Fix host PID detection when CUDA_VISIBLE_DEVICES doesn't start at device 0

src/utils.c

@@ -1,5 +1,6 @@
 #include <stdio.h>
 #include <string.h>
+#include <strings.h>
 #include <dirent.h>
 #include <ctype.h>
 #include <time.h>
@@ -76,43 +77,69 @@ nvmlReturn_t set_task_pid() {
     unsigned int nvmlCounts;
     CHECK_NVML_API(nvmlDeviceGetCount(&nvmlCounts));
 
-    int cudaDev;
-    for (i=0;i<nvmlCounts;i++){
-        cudaDev=nvml_to_cuda_map(i);
-        if (cudaDev<0) {
-            continue;
-        }
-        CHECK_NVML_API(nvmlDeviceGetHandleByIndex(i, &device));
-        do{
-            res = nvmlDeviceGetComputeRunningProcesses(device, &previous, tmp_pids_on_device);
-            if ((res != NVML_SUCCESS) && (res != NVML_ERROR_INSUFFICIENT_SIZE)) {
-                LOG_ERROR("Device2GetComputeRunningProcesses failed %d,%d\n",res,i);
-                return res;
+    /*
+     * The probe context is created on CUDA device 0, so both process
+     * snapshots must be taken on the NVML device that is the same physical
+     * GPU, otherwise the probe PID never appears in the diff. The NVML
+     * device is resolved by UUID: cuda_to_nvml_map() cannot be used here,
+     * because nvml_preInit() - triggered by the nvmlInit() call above -
+     * resets the map to identity, which is only correct when
+     * CUDA_VISIBLE_DEVICES starts with the container's device 0.
+     */
+    int probeDev = 0;
+    int probeNvmlIndex = -1;
+    CUuuid cu_uuid;
+    char want_uuid[48];
+    char have_uuid[96];
+    if (CUDA_OVERRIDE_CALL(cuda_library_entry, cuDeviceGetUuid, &cu_uuid, probeDev) == CUDA_SUCCESS) {
+        const unsigned char *b = (const unsigned char *)cu_uuid.bytes;
+        snprintf(want_uuid, sizeof(want_uuid),
+            "GPU-%02x%02x%02x%02x-%02x%02x-%02x%02x-%02x%02x-%02x%02x%02x%02x%02x%02x",
+            b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7],
+            b[8], b[9], b[10], b[11], b[12], b[13], b[14], b[15]);
+        for (i = 0; i < nvmlCounts; i++) {
+            if (nvmlDeviceGetHandleByIndex(i, &device) != NVML_SUCCESS) {
+                continue;
+            }
+            if (nvmlDeviceGetUUID(device, have_uuid, sizeof(have_uuid)) != NVML_SUCCESS) {
+                continue;
             }
-        }while(res==NVML_ERROR_INSUFFICIENT_SIZE); 
-        mergepid(&previous,&merged_num,(nvmlProcessInfo_t1 *)tmp_pids_on_device,pre_pids_on_device);
-        break;
+            if (strcasecmp(want_uuid, have_uuid) == 0) {
+                probeNvmlIndex = i;
+                break;
+            }
+        }
     }
+    if (probeNvmlIndex < 0) {
+        LOG_WARN("could not match CUDA device %d to an NVML device by uuid, polling NVML device 0", probeDev);
+        probeNvmlIndex = 0;
+    }
+    CHECK_NVML_API(nvmlDeviceGetHandleByIndex(probeNvmlIndex, &device));
+    do {
+        res = nvmlDeviceGetComputeRunningProcesses(device, &previous, tmp_pids_on_device);
+        if ((res != NVML_SUCCESS) && (res != NVML_ERROR_INSUFFICIENT_SIZE)) {
+            LOG_ERROR("Device2GetComputeRunningProcesses failed %d,%d\n", res, probeNvmlIndex);
+            return res;
+        }
+    } while (res == NVML_ERROR_INSUFFICIENT_SIZE);
+    mergepid(&previous, &merged_num, (nvmlProcessInfo_t1 *)tmp_pids_on_device, pre_pids_on_device);
     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));
-    for (i=0;i<nvmlCounts;i++) {
-        cudaDev=nvml_to_cuda_map(i);
-        if (cudaDev<0) {
-            continue;
+    CHECK_CU_RESULT(cuDevicePrimaryCtxRetain(&pctx, probeDev));
+    res = nvmlDeviceGetHandleByIndex(probeNvmlIndex, &device);
+    if (res != NVML_SUCCESS) {
+        LOG_WARN("NVML error at line %d: %d", __LINE__, res);
+        goto cleanup;
+    }
+    do {
+        res = nvmlDeviceGetComputeRunningProcesses(device, &running_processes, tmp_pids_on_device);
+        if ((res != NVML_SUCCESS) && (res != NVML_ERROR_INSUFFICIENT_SIZE)) {
+            LOG_ERROR("Device2GetComputeRunningProcesses failed %d\n", res);
+            goto cleanup;
         }
-        CHECK_NVML_API(nvmlDeviceGetHandleByIndex (i, &device)); 
-        do{
-            res = nvmlDeviceGetComputeRunningProcesses(device, &running_processes, tmp_pids_on_device);
-            if ((res != NVML_SUCCESS) && (res != NVML_ERROR_INSUFFICIENT_SIZE)) {
-                LOG_ERROR("Device2GetComputeRunningProcesses failed %d\n",res);
-                return res;
-            }
-        }while(res == NVML_ERROR_INSUFFICIENT_SIZE);
-        mergepid(&running_processes,&merged_num,(nvmlProcessInfo_t1 *)tmp_pids_on_device,pids_on_device);
-        break;
-    }
+    } while (res == NVML_ERROR_INSUFFICIENT_SIZE);
+    mergepid(&running_processes, &merged_num, (nvmlProcessInfo_t1 *)tmp_pids_on_device, pids_on_device);
     running_processes = merged_num;
     LOG_INFO("current processes num = %u %u",previous,running_processes);
     for (i=0;i<merged_num;i++){
@@ -122,7 +149,8 @@ nvmlReturn_t set_task_pid() {
     unsigned int hostpid = getextrapid(previous,running_processes,pre_pids_on_device,pids_on_device); 
     if (hostpid==0) {
         LOG_ERROR("host pid is error!");
-        return NVML_ERROR_DRIVER_NOT_LOADED;
+        res = NVML_ERROR_DRIVER_NOT_LOADED;
+        goto cleanup;
     }
     LOG_INFO("hostPid=%d",hostpid);
     if (set_host_pid(hostpid)==0) {
@@ -134,8 +162,12 @@ nvmlReturn_t set_task_pid() {
             }
         }
     }
-    CHECK_CU_RESULT(cuDevicePrimaryCtxRelease(0));
-    return NVML_SUCCESS; 
+    res = NVML_SUCCESS;
+cleanup:
+    if (cuDevicePrimaryCtxRelease(probeDev) != CUDA_SUCCESS) {
+        LOG_WARN("failed to release primary context on device %d", probeDev);
+    }
+    return res;
 }
 
 int parse_cuda_visible_env() {