feat(mem): report real container gpu memory and enforce limits accordingly

Author: Hoang Thanh Loi

src/allocator/allocator.c

@@ -42,7 +42,9 @@ int oom_check(const int dev, size_t addon) {
     else
         d=dev;
     uint64_t limit = get_current_device_memory_limit(d);
-    size_t _usage = get_gpu_memory_usage(d);
+    // Use real NVML-reported memory usage instead of internally tracked value
+    // This ensures OOM is triggered based on actual GPU memory consumption
+    size_t _usage = get_gpu_memory_real_usage(d);
 
     if (limit == 0) {
         return 0;

src/cuda/memory.c

@@ -493,17 +493,22 @@ CUresult cuMemGetInfo_v2(size_t* free, size_t* total) {
     LOG_DEBUG("cuMemGetInfo_v2");
     ENSURE_INITIALIZED();
     CHECK_DRV_API(cuCtxGetDevice(&dev));
-    size_t usage = get_current_device_memory_usage(dev);
+    // Use real NVML-reported memory usage for accurate free memory calculation
+    size_t usage = get_gpu_memory_real_usage(dev);
     size_t limit = get_current_device_memory_limit(dev);
     if (limit == 0) {
         CUDA_OVERRIDE_CALL(cuda_library_entry,cuMemGetInfo_v2, free, total);
         LOG_INFO("orig free=%ld total=%ld", *free, *total);
-        *free = *total - usage;
+        *free = (*total > usage) ? (*total - usage) : 0;
         LOG_INFO("after free=%ld total=%ld", *free, *total);
         return CUDA_SUCCESS;
     } else if (limit < usage) {
         LOG_WARN("limit < usage; usage=%ld, limit=%ld", usage, limit);
-        return CUDA_ERROR_INVALID_VALUE;
+        // Return 0 free memory instead of error when over limit
+        CUDA_OVERRIDE_CALL(cuda_library_entry,cuMemGetInfo_v2, free, total);
+        *free = 0;
+        *total = limit;
+        return CUDA_SUCCESS;
     } else {
         CUDA_OVERRIDE_CALL(cuda_library_entry,cuMemGetInfo_v2, free, total);
         LOG_INFO("orig free=%ld total=%ld limit=%ld usage=%ld",

src/multiprocess/multiprocess_memory_limit.c

@@ -52,6 +52,7 @@ int _record_kernel_interval = 1;
 
 void do_init_device_memory_limits(uint64_t*, int);
 void exit_withlock(int exitcode);
+uint64_t nvml_get_device_memory_usage(const int dev);
 
 void set_current_gpu_status(int status){
     int i;
@@ -255,6 +256,20 @@ size_t get_gpu_memory_usage(const int dev) {
     return total;
 }
 
+size_t get_gpu_memory_real_usage(const int dev) {
+    LOG_INFO("get_gpu_memory_real_usage dev=%d",dev);
+    ensure_initialized();
+    // Query NVML directly for real-time memory usage instead of using cached monitor value
+    // This ensures OOM checks use current actual GPU memory, not stale data
+    size_t nvml_usage = nvml_get_device_memory_usage(dev);
+    size_t tracked_usage = get_gpu_memory_usage(dev);
+    // Use the maximum of NVML-reported and tracked value to be conservative
+    // NVML gives real GPU memory, tracked gives what we've accounted for
+    size_t real_usage = (nvml_usage > tracked_usage) ? nvml_usage : tracked_usage;
+    LOG_INFO("get_gpu_memory_real_usage dev=%d nvml_usage=%lu tracked_usage=%lu real_usage=%lu", dev, nvml_usage, tracked_usage, real_usage);
+    return real_usage;
+}
+
 int set_gpu_device_memory_monitor(int32_t pid,int dev,size_t monitor){
     //LOG_WARN("set_gpu_device_memory_monitor:%d %d %lu",pid,dev,monitor);
     int i;
@@ -307,13 +322,15 @@ uint64_t nvml_get_device_memory_usage(const int dev) {
     ret = nvmlDeviceGetHandleByIndex(dev, &ndev);
     if (ret != NVML_SUCCESS) {
         LOG_ERROR("NVML get device %d error, %s", dev, nvmlErrorString(ret));
+        return 0;
     }
     unsigned int pcnt = SHARED_REGION_MAX_PROCESS_NUM;
     nvmlProcessInfo_v1_t infos[SHARED_REGION_MAX_PROCESS_NUM];
     LOG_DEBUG("before nvmlDeviceGetComputeRunningProcesses");
     ret = nvmlDeviceGetComputeRunningProcesses(ndev, &pcnt, infos);
     if (ret != NVML_SUCCESS) {
         LOG_ERROR("NVML get process error, %s", nvmlErrorString(ret));
+        return 0;
     }
     int i = 0;
     uint64_t usage = 0;
@@ -322,9 +339,13 @@ uint64_t nvml_get_device_memory_usage(const int dev) {
     for (; i < pcnt; i++) {
         int slot = 0;
         for (; slot < region->proc_num; slot++) {
-            if (infos[i].pid != region->procs[slot].pid)
+            // NVML returns host PIDs, so we need to compare with hostpid, not pid
+            // pid is the container PID (from getpid()), hostpid is the real host PID
+            if (infos[i].pid != region->procs[slot].hostpid)
                 continue;
             usage += infos[i].usedGpuMemory;
+            LOG_DEBUG("nvml_get_device_memory_usage: matched hostpid=%d, usedGpuMemory=%lu", 
+                infos[i].pid, infos[i].usedGpuMemory);
         }
     }
     unlock_shrreg();
@@ -851,8 +872,8 @@ uint64_t get_current_device_memory_usage(const int dev) {
     if (dev < 0 || dev >= CUDA_DEVICE_MAX_COUNT) {
         LOG_ERROR("Illegal device id: %d", dev);
     }
-    result = get_gpu_memory_usage(dev);
-//    result= nvml_get_device_memory_usage(dev);
+    // Use real NVML-reported memory usage for accurate memory tracking
+    result = get_gpu_memory_real_usage(dev);
     finish=clock();
     LOG_DEBUG("get_current_device_memory_usage:tick=%lu result=%lu\n",finish-start,result);
     return result;

src/multiprocess/multiprocess_memory_limit.h

@@ -133,6 +133,7 @@ int set_host_pid(int hostpid);
 uint64_t get_current_device_memory_monitor(const int dev);
 uint64_t get_current_device_memory_usage(const int dev);
 size_t get_gpu_memory_usage(const int dev);
+size_t get_gpu_memory_real_usage(const int dev);
 
 // Priority-related
 int get_current_priority();

src/nvml/hook.c

@@ -338,10 +338,10 @@ nvmlReturn_t _nvmlDeviceGetMemoryInfo(nvmlDevice_t device,void* memory,int versi
     if (cudadev < 0) {
         return NVML_SUCCESS;
     }
+    // get_current_device_memory_usage now returns real NVML-reported usage
     size_t usage = get_current_device_memory_usage(cudadev);
-    size_t monitor = get_current_device_memory_monitor(cudadev);
     size_t limit = get_current_device_memory_limit(cudadev);
-    LOG_DEBUG("usage=%ld limit=%ld monitor=%ld", usage, limit, monitor);
+    LOG_DEBUG("usage=%ld limit=%ld", usage, limit);
     if (limit == 0) {
         switch (version) {
         case 1:
@@ -354,12 +354,12 @@ nvmlReturn_t _nvmlDeviceGetMemoryInfo(nvmlDevice_t device,void* memory,int versi
     } else {
         switch (version) {
         case 1:
-             ((nvmlMemory_t*)memory)->free = (limit-usage);
+             ((nvmlMemory_t*)memory)->free = (limit > usage) ? (limit - usage) : 0;
              ((nvmlMemory_t*)memory)->total = limit;
              ((nvmlMemory_t*)memory)->used = usage;
             return NVML_SUCCESS;
         case 2:
-            ((nvmlMemory_v2_t *)memory)->free = (limit-usage);
+            ((nvmlMemory_v2_t *)memory)->free = (limit > usage) ? (limit - usage) : 0;
             ((nvmlMemory_v2_t *)memory)->total = limit;
             ((nvmlMemory_v2_t *)memory)->used = usage;
             return NVML_SUCCESS;