Skip to content

MB-71397: Add on CPU index if Add OOMs on GPU. #418

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
CascadingRadium merged 1 commit into from
May 26, 2026
Merged

MB-71397: Add on CPU index if Add OOMs on GPU. #418

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 16 additions & 2 deletions faiss_vector_index_gpu_float32.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ package zap
import (
"encoding/binary"
"encoding/json"
"errors"
"reflect"
"sync/atomic"

Expand Down Expand Up @@ -177,7 +178,6 @@ func (f *faissGPUFloat32Index) close() {
f.waitGPU()
f.teardownGPU()
f.cpuIdx.Close()
f.idxBytes = nil
}

// teardownGPU stops the batcher first (while gpuIdx is still live so that
Expand Down Expand Up @@ -317,8 +317,21 @@ func (f *faissGPUFloat32Index) trainAndAdd(trainingData *vectorSet, vecsToAdd *v

err = gpuState.idx.Add(vecsToAdd.floatData)
if err != nil {
// Fallback to full CPU train and add if:
// 1. We get a non-OOM error from the GPU index
if !errors.Is(err, faiss.ErrGPUOutOfMemory) {
f.teardownGPU()
return f.trainAndAddCPU(trainingData, vecsToAdd)
}
// 2. We get an OOM error but syncing GPU to CPU fails.
if f.syncGPUToCPU() != nil {
f.teardownGPU()
return f.trainAndAddCPU(trainingData, vecsToAdd)
}
// OOM but CPU sync succeeded: both indexes are trained.
// Teardown GPU to avoid discrepancy, add to CPU instead.
f.teardownGPU()
return f.trainAndAddCPU(trainingData, vecsToAdd)
return f.cpuIdx.Add(vecsToAdd.floatData)
}

err = f.syncGPUToCPU()
Expand Down Expand Up @@ -353,6 +366,7 @@ func (f *faissGPUFloat32Index) mergeFrom(other faissIndex, offset int64) error {
// syncGPUToCPU clones the current GPU index state back to the CPU index,
// replacing the old CPU index.
func (f *faissGPUFloat32Index) syncGPUToCPU() error {
f.waitGPU()
gpuState := f.gpu.Load()
if gpuState == nil {
return nil
Expand Down
2 changes: 1 addition & 1 deletion go.mod
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ go 1.25.0
require (
github.com/RoaringBitmap/roaring/v2 v2.14.5
github.com/blevesearch/bleve_index_api v1.3.11
github.com/blevesearch/go-faiss v1.1.2
github.com/blevesearch/go-faiss v1.1.3-0.20260525132456-c1cb753e04cd
github.com/blevesearch/mmap-go v1.2.0
github.com/blevesearch/scorch_segment_api/v2 v2.4.7
github.com/blevesearch/vellum v1.2.0
Expand Down
4 changes: 2 additions & 2 deletions go.sum
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,8 @@ github.com/bits-and-blooms/bitset v1.24.2 h1:M7/NzVbsytmtfHbumG+K2bremQPMJuqv1JD
github.com/bits-and-blooms/bitset v1.24.2/go.mod h1:7hO7Gc7Pp1vODcmWvKMRA9BNmbv6a/7QIWpPxHddWR8=
github.com/blevesearch/bleve_index_api v1.3.11 h1:x29vbV8OjWfLcrDVd7Lr1q+BkLNS0JWNEig0MCVnKH4=
github.com/blevesearch/bleve_index_api v1.3.11/go.mod h1:xvd48t5XMeeioWQ5/jZvgLrV98flT2rdvEJ3l/ki4Ko=
github.com/blevesearch/go-faiss v1.1.2 h1:ojv2S7ot3orbk8wMfJWryq37G4eIL8Y8PLLZYd8ZLHY=
github.com/blevesearch/go-faiss v1.1.2/go.mod h1:OMGQwOaRRYxrmeNdMrXJPvVx8gBnvE5RYrr0BahNnkk=
github.com/blevesearch/go-faiss v1.1.3-0.20260525132456-c1cb753e04cd h1:ftEpy+Ma4N/O4zgIRn4pKZOhmi8UnvApcYT6hTf0IYQ=
github.com/blevesearch/go-faiss v1.1.3-0.20260525132456-c1cb753e04cd/go.mod h1:w3W9AiWsFRGVaMG+/cmJi7iHEAuGyC6blsgO1EzCK/M=
github.com/blevesearch/mmap-go v1.2.0 h1:l33nNKPFcBjJUMwem6sAYJPUzhUCABoK9FxZDGiFNBI=
github.com/blevesearch/mmap-go v1.2.0/go.mod h1:Vd6+20GBhEdwJnU1Xohgt88XCD/CTWcqbCNxkZpyBo0=
github.com/blevesearch/scorch_segment_api/v2 v2.4.7 h1:GlMzW08hcsM3DnLUxhyF/1PcDal1qtvvIuytuph5djw=
Expand Down
Loading