MB-71397: Add a maximal-effort strategy to avoid OOM errors for GPU indexes.

gpu.go

@@ -23,23 +23,30 @@ package faiss
 #include <faiss/c_api/gpu/GpuAutoTune_c.h>
 #include <faiss/c_api/gpu/GpuClonerOptions_c.h>
 #include <faiss/c_api/gpu/DeviceUtils_c.h>
+#include <faiss/c_api/gpu/GpuIndex_c.h>
+#include <faiss/c_api/gpu/GpuIndexIVF_c.h>
 */
 import "C"
 import (
 	"errors"
 	"fmt"
 	"math/rand"
+	"reflect"
 	"sort"
 	"sync"
 	"sync/atomic"
 	"time"
-	"unsafe"
 )
 
 var (
 	errAccessingGPUDevices = errors.New("error accessing GPU devices")
 	errNilIndex            = errors.New("index is nil")
 	errNoGPUDevices        = errors.New("no GPU devices available")
+	errNotEnoughGPUMemory  = errors.New("selected GPU device does not have enough free memory for cloning")
+	errTrainFailed         = errors.New("training failed on GPU index")
+	errAddFailed           = errors.New("adding vectors failed on GPU index")
+	errMemoryReserveFailed = errors.New("failed to reserve memory on GPU index")
+	errSearchFailed        = errors.New("search failed on GPU index")
 )
 
 // memorySpace controls where GPU index data is allocated.
@@ -54,8 +61,8 @@ const (
 )
 
 const (
-	// the minimum amount of free memory that must be available on a GPU to be considered for index cloning.
-	minGPUFreeMemory = 512 * 1024 * 1024 // 512 MiB
+	// the default minimum amount of free memory that must be available on a GPU to be considered for index cloning.
+	defaultGPUMinFreeMemory = 512 * 1024 * 1024 // 512 MiB
 	// the default memory space to use for GPU indices.
 	defaultGPUMemoryMode = memorySpaceUnified
 	// the default amount of pinned memory to allocate for each GPU clone operation.
@@ -65,19 +72,31 @@ const (
 var (
 	gpuCount     int
 	loadBalancer *gpuLoadBalancer
+	serializer   *gpuIndexSerializer
 )
 
+var reflectStaticSizefaissGPUIndex uint64
+
 func init() {
+	var g faissGPUIndex
+	reflectStaticSizefaissGPUIndex = uint64(reflect.TypeOf(g).Size())
+
 	var err error
 	gpuCount, err = numGPUs()
 	if err != nil || gpuCount <= 0 {
 		gpuCount = 0
 	}
-	if gpuCount > 0 {
-		// TODO: verify if 500 milliseconds is a good interval
+	// Only initialize the load balancer if multiple GPUs are available,
+	// as it's not needed for a single GPU and adds overhead from monitoring.
+	if gpuCount > 1 {
 		loadBalancer = newGPULoadBalancer(500 * time.Millisecond)
 		go loadBalancer.monitor()
 	}
+	// Initialize the GPU serializer if at least one GPU is available;
+	// it will be used to serialize access to each GPU device to prevent oversubscription.
+	if gpuCount > 0 {
+		serializer = newGPUIndexSerializer(gpuCount)
+	}
 }
 
 // numGPUs returns the number of available GPU devices.
@@ -90,6 +109,22 @@ func numGPUs() (int, error) {
 	return int(rv), nil
 }
 
+// getFreeGPUMemory returns the amount of free memory in bytes for the specified GPU device.
+// It handles any CGO error by returning 0, which will allow callers to naturally exclude that device from selection.
+func getFreeGPUMemory(device int) uint64 {
+	var freeBytes C.size_t
+	if c := C.faiss_gpu_free_memory(C.int(device), &freeBytes); c == 0 {
+		return uint64(freeBytes)
+	}
+	return 0
+}
+
+// hasEnoughGPUMemory returns true if the GPU has enough free memory for the required allocation plus a buffer.
+func hasEnoughGPUMemory(freeMem, requiredMemory uint64) bool {
+	return freeMem > requiredMemory+defaultGPUMinFreeMemory
+}
+
+// --------------------------------------------------------------------------------
 // gpuLoadBalancer monitors GPU free memory on a fixed interval, keeps a
 // memory-sorted list of devices, and hands them out in round-robin order.
 // At each interval the list is re-sorted and the round-robin counter resets
@@ -137,17 +172,15 @@ func (lb *gpuLoadBalancer) refresh() {
 	for i := 0; i < gpuCount; i++ {
 		go func(device int) {
 			defer wg.Done()
-			var freeBytes C.size_t
-			if C.faiss_gpu_free_memory(C.int(device), &freeBytes) == 0 {
-				lb.freeMemory[device] = uint64(freeBytes)
-			}
+			lb.freeMemory[device] = getFreeGPUMemory(device)
 		}(i)
 	}
 	wg.Wait()
 
-	// Only include devices that reported non-zero free memory, and have at least minGPUFreeMemory free.
+	// Only include devices that reported non-zero free memory and
+	// have at least the minimum required free memory for cloning.
 	for i, mem := range lb.freeMemory {
-		if mem > minGPUFreeMemory {
+		if mem > 0 && mem >= defaultGPUMinFreeMemory {
 			lb.scratchDevs = append(lb.scratchDevs, i)
 		}
 	}
@@ -189,37 +222,139 @@ func (lb *gpuLoadBalancer) nextDevice() (int, error) {
 }
 
 func getBestGPUDevice() (int, error) {
-	if gpuCount == 0 || loadBalancer == nil {
+	if gpuCount == 0 || serializer == nil {
 		return 0, errNoGPUDevices
 	}
+	// If loadBalancer is nil, it means we have only 1 GPU, so just return device 0.
+	if loadBalancer == nil {
+		return 0, nil
+	}
 	return loadBalancer.nextDevice()
 }
 
-// only expose API used by zapx
-type GPUIndexImpl struct {
-	idx         *faissIndex
+// --------------------------------------------------------------------------------
+// gpuIndexSerializer provides a simple mutex-like mechanism to serialize index clones to
+// each GPU device, preventing oversubscription and out-of-memory errors when multiple
+// goroutines attempt to clone to the same GPU at the same time.
+type gpuIndexSerializer struct {
+	deviceMu []sync.Mutex
+}
+
+func newGPUIndexSerializer(numGPUs int) *gpuIndexSerializer {
+	s := &gpuIndexSerializer{
+		deviceMu: make([]sync.Mutex, numGPUs),
+	}
+	return s
+}
+
+func (s *gpuIndexSerializer) acquire(device int) {
+	s.deviceMu[device].Lock()
+}
+
+func (s *gpuIndexSerializer) release(device int) {
+	s.deviceMu[device].Unlock()
+}
+
+// --------------------------------------------------------------------------------
+
+type GPUIndex interface {
+	// D returns the dimension of the indexed vectors.
+	D() int
+	// Add adds vectors to the index.
+	// Resurns two errors:
+	//  - The first error indicates failure to reserve memory for the add operation.
+	//  - The second error indicates failure to add the vectors after successful memory reservation.
+	Add(x []float32) (error, error)
+	// Train trains the index on a representative set of vectors.
+	Train(x []float32) error
+	// Search queries the index with the vectors in x.
+	// Returns the IDs of the k nearest neighbors for each query vector and the
+	// corresponding distances.
+	Search(x []float32, k int64) (distances []float32, labels []int64, err error)
+	// Size estimates the memory footprint of the index assuming in bytes,
+	// if the underlying faiss index is memory-mapped and not fully loaded into memory.
+	Size() uint64
+	// Close frees the memory used by the index.
+	Close()
+	// gPtr returns the underlying C pointer to the FaissGpuIndex.
+	gPtr() *C.FaissGpuIndex
+}
+
+type faissGPUIndex struct {
+	idx         *C.FaissGpuIndex
 	gpuResource *C.FaissStandardGpuResources
+	device      int
+	codeSize    uint64
 }
 
-func (g *GPUIndexImpl) cPtr() *C.FaissIndex {
-	return g.idx.idx
+func (g *faissGPUIndex) gPtr() *C.FaissGpuIndex {
+	return g.idx
 }
 
-func (g *GPUIndexImpl) Train(x []float32) error {
-	return g.idx.Train(x)
+func (g *faissGPUIndex) D() int {
+	return int(C.faiss_GpuIndex_d(g.idx))
 }
 
-func (g *GPUIndexImpl) Add(x []float32) error {
-	return g.idx.Add(x)
+func (g *faissGPUIndex) Add(x []float32) (error, error) {
+	n := len(x) / g.D()
+	// cast to gpu ivf index first
+	var ivfIdx *C.FaissGpuIndexIVF
+	ivfIdx = C.faiss_GpuIndexIVF_cast(g.idx)
+	if ivfIdx != nil {
+		// try to reserve memory for the new vectors first if possible
+		// acquire device lock for memory allocation
+		serializer.acquire(g.device)
+		// early exit path. Use a heuristic to determine the minimum
+		// memory needed to add the vectors before attempting to reserve memory or add to the index,
+		// to avoid unnecessary work and errors when adding a large batch of vectors that exceed GPU memory.
+		requiredMemory := uint64(n) * g.codeSize
+		freeMem := getFreeGPUMemory(g.device)
+		if !hasEnoughGPUMemory(freeMem, requiredMemory) {
+			serializer.release(g.device)
+			return errNotEnoughGPUMemory, nil
+		}
+		// actually reserve the memory on the GPU for the new vectors.
+		if c := C.faiss_GpuIndexIVF_reserve_memory(ivfIdx, C.size_t(n)); c != 0 {
+			serializer.release(g.device)
+			return errMemoryReserveFailed, nil
+		}
+		serializer.release(g.device)
+	}
+	if c := C.faiss_GpuIndex_add(g.idx, C.idx_t(n), (*C.float)(&x[0])); c != 0 {
+		return nil, errAddFailed
+	}
+	return nil, nil
 }
 
-func (g *GPUIndexImpl) Search(x []float32, k int64) ([]float32, []int64, error) {
-	return g.idx.Search(x, k)
+func (g *faissGPUIndex) Train(x []float32) error {
+	n := len(x) / g.D()
+	if c := C.faiss_GpuIndex_train(g.idx, C.idx_t(n), (*C.float)(&x[0])); c != 0 {
+		return errTrainFailed
+	}
+	return nil
 }
 
-func (g *GPUIndexImpl) Close() {
+func (g *faissGPUIndex) Search(x []float32, k int64) ([]float32, []int64, error) {
+	n := len(x) / g.D()
+	distances := make([]float32, int64(n)*k)
+	labels := make([]int64, int64(n)*k)
+	if c := C.faiss_GpuIndex_search(
+		g.idx,
+		C.idx_t(n),
+		(*C.float)(&x[0]),
+		C.idx_t(k),
+		(*C.float)(&distances[0]),
+		(*C.idx_t)(&labels[0]),
+	); c != 0 {
+		return nil, nil, errSearchFailed
+	}
+
+	return distances, labels, nil
+}
+
+func (g *faissGPUIndex) Close() {
 	if g.idx != nil {
-		g.idx.Close()
+		C.faiss_GpuIndex_free(g.idx)
 		g.idx = nil
 	}
 	if g.gpuResource != nil {
@@ -228,6 +363,14 @@ func (g *GPUIndexImpl) Close() {
 	}
 }
 
+func (g *faissGPUIndex) Size() uint64 {
+	return reflectStaticSizefaissGPUIndex
+}
+
+type GPUIndexImpl struct {
+	GPUIndex
+}
+
 // CloneToGPU transfers a CPU index to the best available GPU based on free memory.
 func CloneToGPU(cpuIndex *IndexImpl) (*GPUIndexImpl, error) {
 	if cpuIndex == nil {
@@ -239,12 +382,24 @@ func CloneToGPU(cpuIndex *IndexImpl) (*GPUIndexImpl, error) {
 	if err != nil {
 		return nil, err
 	}
-
+	// Acquire the GPU resource for the duration of the clone operation to prevent oversubscription.
+	serializer.acquire(device)
+	defer serializer.release(device)
+	// amount of GPU memory required for the index clone
+	requiredMemory, err := cpuIndex.IndexSize()
+	if err != nil {
+		return nil, err
+	}
+	// first check if we have enough free memory on the selected GPU before attempting the clone, to avoid unnecessary work and errors.
+	freeMem := getFreeGPUMemory(device)
+	// The GPU must have enough free memory for the index plus a minimum buffer.
+	if !hasEnoughGPUMemory(freeMem, requiredMemory) {
+		return nil, errNotEnoughGPUMemory
+	}
 	var gpuResource *C.FaissStandardGpuResources
 	if code := C.faiss_StandardGpuResources_new(&gpuResource); code != 0 {
 		return nil, fmt.Errorf("failed to initialize GPU resources: error code %d, err: %v", code, getLastError())
 	}
-
 	// Disable the pre-allocated temp memory pool so that all GPU memory is
 	// available for index data; unified memory mode handles intermediate
 	// allocations via cudaMalloc/cudaFree on demand.
@@ -256,16 +411,13 @@ func CloneToGPU(cpuIndex *IndexImpl) (*GPUIndexImpl, error) {
 		C.faiss_StandardGpuResources_free(gpuResource)
 		return nil, fmt.Errorf("failed to disable GPU pinned memory: error code %d, err: %v", code, getLastError())
 	}
-
 	var clonerOpts *C.FaissGpuClonerOptions
 	if code := C.faiss_GpuClonerOptions_new(&clonerOpts); code != 0 {
 		C.faiss_StandardGpuResources_free(gpuResource)
 		return nil, fmt.Errorf("failed to create cloner options: error code %d, err: %v", code, getLastError())
 	}
 	defer C.faiss_GpuClonerOptions_free(clonerOpts)
-
 	C.faiss_GpuClonerOptions_set_memorySpace(clonerOpts, C.int(defaultGPUMemoryMode))
-
 	var gpuIdx *C.FaissGpuIndex
 	code := C.faiss_index_cpu_to_gpu_with_options(
 		gpuResource,
@@ -278,25 +430,27 @@ func CloneToGPU(cpuIndex *IndexImpl) (*GPUIndexImpl, error) {
 		C.faiss_StandardGpuResources_free(gpuResource)
 		return nil, fmt.Errorf("failed to transfer index to GPU device %d: error code %d, err: %v", device, code, getLastError())
 	}
-
-	idx := &faissIndex{
-		idx: (*C.FaissIndex)(unsafe.Pointer(gpuIdx)),
+	codeSize, err := cpuIndex.CodeSize()
+	if err != nil {
+		C.faiss_StandardGpuResources_free(gpuResource)
+		return nil, err
 	}
-
-	return &GPUIndexImpl{
-		idx:         idx,
+	g := &faissGPUIndex{
+		idx:         gpuIdx,
 		gpuResource: gpuResource,
-	}, nil
+		device:      device,
+		codeSize:    codeSize,
+	}
+	return &GPUIndexImpl{g}, nil
 }
 
 func CloneToCPU(gpuIndex *GPUIndexImpl) (*IndexImpl, error) {
 	if gpuIndex == nil {
 		return nil, errNilIndex
 	}
-
 	var cpuIdx *C.FaissIndex
 	code := C.faiss_index_gpu_to_cpu(
-		gpuIndex.cPtr(),
+		gpuIndex.gPtr(),
 		&cpuIdx,
 	)
 	if code != 0 {

gpu_stub.go

@@ -22,12 +22,13 @@ import "errors"
 // GPUIndexImpl is an opaque type when not built with GPU support.
 type GPUIndexImpl struct{}
 
-func (g *GPUIndexImpl) Train(x []float32) error { return errGPUNotBuilt }
-func (g *GPUIndexImpl) Add(x []float32) error   { return errGPUNotBuilt }
+func (g *GPUIndexImpl) Train(x []float32) error        { return errGPUNotBuilt }
+func (g *GPUIndexImpl) Add(x []float32) (error, error) { return nil, errGPUNotBuilt }
 func (g *GPUIndexImpl) Search(x []float32, k int64) ([]float32, []int64, error) {
 	return nil, nil, errGPUNotBuilt
 }
-func (g *GPUIndexImpl) Close() {}
+func (g *GPUIndexImpl) Close()       {}
+func (g *GPUIndexImpl) Size() uint64 { return 0 }
 
 var errGPUNotBuilt = errors.New("not built with GPU support (requires -tags gpu)")