From 8efff49ad69d170e9c7bb904cc5095a1e4cdca4e Mon Sep 17 00:00:00 2001
From: maverick123123 <yuming.wu@dynamia.ai>
Date: Thu, 23 Jul 2026 16:12:48 +0800
Subject: [PATCH] fix: probe CUDA context on correct device in set_task_pid()

set_task_pid() polls NVML processes on the first device where
nvml_to_cuda_map(i) >= 0, but the probe context was always created
on hardcoded CUDA device 0 via cuDevicePrimaryCtxRetain(&pctx, 0).

When CUDA_VISIBLE_DEVICES does not start with the pod's device 0
(e.g., CVD=1 or CVD=1,0), the polled NVML device and the probe
context device are different physical GPUs. getextrapid() never
finds the new PID, causing a spurious 'host pid is error!' and
subsequent Device 0 OOM due to incorrect memory accounting.

Fix: use the cudaDev found in the NVML enumeration loop for both
cuDevicePrimaryCtxRetain and cuDevicePrimaryCtxRelease, ensuring
the probe context runs on the same device being polled.

Fixes #225
---
 src/utils.c | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/src/utils.c b/src/utils.c
index e96b43e3..b525daad 100755
--- a/src/utils.c
+++ b/src/utils.c
@@ -111,6 +111,7 @@ nvmlReturn_t set_task_pid() {
     CHECK_NVML_API(nvmlDeviceGetCount(&nvmlCounts));
     
     int cudaDev;
+    int probeDev = 0;
     for (i=0;i<nvmlCounts;i++){
         cudaDev=nvml_to_cuda_map(i);
         if (cudaDev<0) {
@@ -123,14 +124,15 @@ nvmlReturn_t set_task_pid() {
                 LOG_ERROR("Device2GetComputeRunningProcesses failed %d,%d\n",res,i);
                 return res;
             }
-        }while(res==NVML_ERROR_INSUFFICIENT_SIZE); 
+        }while(res==NVML_ERROR_INSUFFICIENT_SIZE);
         mergepid(&previous,&merged_num,(nvmlProcessInfo_t1 *)tmp_pids_on_device,pre_pids_on_device);
+        probeDev = cudaDev;
         break;
     }
     previous = merged_num;
     merged_num = 0;
     memset(tmp_pids_on_device,0,sizeof(nvmlProcessInfo_v1_t)*SHARED_REGION_MAX_PROCESS_NUM);
-    CHECK_CU_RESULT(cuDevicePrimaryCtxRetain(&pctx,0));
+    CHECK_CU_RESULT(cuDevicePrimaryCtxRetain(&pctx,probeDev));
     for (i=0;i<nvmlCounts;i++) {
         cudaDev=nvml_to_cuda_map(i);
         if (cudaDev<0) {
@@ -168,7 +170,7 @@ nvmlReturn_t set_task_pid() {
             }
         }
     }
-    CHECK_CU_RESULT(cuDevicePrimaryCtxRelease(0));
+    CHECK_CU_RESULT(cuDevicePrimaryCtxRelease(probeDev));
     return NVML_SUCCESS; 
 }