From fc3c920463ccdd1369a7abec0f7691d2343d768a Mon Sep 17 00:00:00 2001 From: Rahul Rampure Date: Wed, 3 Jun 2026 16:27:43 +0530 Subject: [PATCH 1/2] Probe GPU devices for health --- gpu.go | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/gpu.go b/gpu.go index e472531..e8689f8 100644 --- a/gpu.go +++ b/gpu.go @@ -214,14 +214,18 @@ type gpuSnapshotStore struct { func newGPUSnapshotStore() *gpuSnapshotStore { snapshots := make([]*gpuSnapshot, gpuCount) for device := 0; device < gpuCount; device++ { - // get total free memory for the GPU device. - totMemory := uint64(0) - var freeBytes C.size_t - if c := C.faiss_gpu_free_memory( - C.int(device), - &freeBytes, - ); c == 0 { - totMemory = uint64(freeBytes) + cDev := C.int(device) + // first probe if the device is healthy and can be used + if c := C.faiss_probe_gpu(cDev); c == 0 { + // get total free memory for the device + totMemory := uint64(0) + var freeBytes C.size_t + if c := C.faiss_gpu_free_memory( + cDev, + &freeBytes, + ); c == 0 { + totMemory = uint64(freeBytes) + } } // if we fail to get the free memory for the GPU, // we still create a snapshot with 0 total and free memory, From 11119bb6e61c1d75baa0d73ed81f312fae0227e6 Mon Sep 17 00:00:00 2001 From: Rahul Rampure Date: Wed, 3 Jun 2026 17:43:46 +0530 Subject: [PATCH 2/2] fixup API --- gpu.go | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/gpu.go b/gpu.go index e8689f8..5f0b020 100644 --- a/gpu.go +++ b/gpu.go @@ -215,10 +215,10 @@ func newGPUSnapshotStore() *gpuSnapshotStore { snapshots := make([]*gpuSnapshot, gpuCount) for device := 0; device < gpuCount; device++ { cDev := C.int(device) + totMemory := uint64(0) // first probe if the device is healthy and can be used - if c := C.faiss_probe_gpu(cDev); c == 0 { - // get total free memory for the device - totMemory := uint64(0) + var probeResult C.int + if c := C.faiss_probe_gpu(cDev, &probeResult); c == 0 && probeResult == 0 { var freeBytes C.size_t if c := C.faiss_gpu_free_memory( cDev,