feat: account for CUDA context memory

src/cuda/CMakeLists.txt

@@ -1,3 +1,3 @@
-add_library(cuda_mod OBJECT context.c device.c hook.c event.c hook.c memory.c stream.c graph.c)
+add_library(cuda_mod OBJECT context.c context_accounting.c device.c hook.c event.c hook.c memory.c stream.c graph.c)
 target_compile_options(cuda_mod PUBLIC ${LIBRARY_COMPILE_FLAGS})
 target_link_libraries(cuda_mod PUBLIC nvidia-ml -lcuda)

src/cuda/context.c

@@ -1,8 +1,52 @@
+#include <errno.h>
+
 #include "include/libcuda_hook.h"
+#include "include/libvgpu.h"
+#include "cuda/context_accounting.h"
 #include "multiprocess/multiprocess_memory_limit.h"
 
 extern size_t context_size;
-extern int ctx_activate[16];
+extern int ctx_activate[CUDA_DEVICE_MAX_COUNT];
+extern int pidfound;
+
+static size_t device_context_size[CUDA_DEVICE_MAX_COUNT];
+static primary_context_accounting_t
+    context_accounting[CUDA_DEVICE_MAX_COUNT];
+static pthread_mutex_t context_accounting_lock = PTHREAD_MUTEX_INITIALIZER;
+static pthread_mutex_t context_device_locks[CUDA_DEVICE_MAX_COUNT] = {
+    PTHREAD_MUTEX_INITIALIZER, PTHREAD_MUTEX_INITIALIZER,
+    PTHREAD_MUTEX_INITIALIZER, PTHREAD_MUTEX_INITIALIZER,
+    PTHREAD_MUTEX_INITIALIZER, PTHREAD_MUTEX_INITIALIZER,
+    PTHREAD_MUTEX_INITIALIZER, PTHREAD_MUTEX_INITIALIZER,
+    PTHREAD_MUTEX_INITIALIZER, PTHREAD_MUTEX_INITIALIZER,
+    PTHREAD_MUTEX_INITIALIZER, PTHREAD_MUTEX_INITIALIZER,
+    PTHREAD_MUTEX_INITIALIZER, PTHREAD_MUTEX_INITIALIZER,
+    PTHREAD_MUTEX_INITIALIZER, PTHREAD_MUTEX_INITIALIZER,
+};
+
+void context_accounting_fork_prepare() {
+    pthread_mutex_lock(&context_accounting_lock);
+}
+
+void context_accounting_fork_parent() {
+    pthread_mutex_unlock(&context_accounting_lock);
+}
+
+void context_accounting_fork_child() {
+    int dev;
+
+    /* context_size is kept.  The parent's probe is a fair estimate for a
+     * child on the same GPU, and the child's own probe overwrites it. */
+    primary_context_accounting_reset(context_accounting,
+                                     CUDA_DEVICE_MAX_COUNT);
+    for (dev = 0; dev < CUDA_DEVICE_MAX_COUNT; dev++) {
+        device_context_size[dev] = 0;
+        ctx_activate[dev] = 0;
+        context_device_locks[dev] =
+            (pthread_mutex_t)PTHREAD_MUTEX_INITIALIZER;
+    }
+    pthread_mutex_unlock(&context_accounting_lock);
+}
 
 
 CUresult cuDevicePrimaryCtxGetState( CUdevice dev, unsigned int* flags, int* active ){
@@ -11,16 +55,144 @@ CUresult cuDevicePrimaryCtxGetState( CUdevice dev, unsigned int* flags, int* act
     return res;
 }
 
+static size_t context_charge_for_device(CUdevice dev) {
+    if (dev >= 0 && dev < CUDA_DEVICE_MAX_COUNT &&
+        device_context_size[dev] > 0) {
+        return device_context_size[dev];
+    }
+    /* Device 0 trusts the host PID probe, which measured it.  Other devices
+     * are measured on their first retain; if that fails, the retain path
+     * falls back to the probed size rather than charging nothing. */
+    return dev == 0 ? context_size : 0;
+}
+
+static CUresult rollback_unaccounted_retain(CUdevice dev,
+                                            int retain_recorded,
+                                            size_t bytes_to_add) {
+    CUresult release_result;
+
+    if (retain_recorded &&
+        primary_context_rollback_retain(&context_accounting[dev],
+                                        bytes_to_add) != 0) {
+        LOG_ERROR("Failed to roll back local context accounting on device %d",
+                  dev);
+    }
+    release_result = CUDA_OVERRIDE_CALL(cuda_library_entry,
+                                        cuDevicePrimaryCtxRelease_v2, dev);
+    if (release_result != CUDA_SUCCESS) {
+        LOG_ERROR("Failed to release an unaccounted primary context on "
+                  "device %d: %d", dev, release_result);
+    }
+    ctx_activate[dev] = (int)context_accounting[dev].retain_count;
+    return CUDA_ERROR_OUT_OF_MEMORY;
+}
+
 CUresult cuDevicePrimaryCtxRetain(CUcontext *pctx, CUdevice dev){
-    LOG_INFO("dev=%d context_size=%ld",dev,context_size);
-    //for Initialization only
-    CUresult res = CUDA_OVERRIDE_CALL(cuda_library_entry,cuDevicePrimaryCtxRetain,pctx,dev);
-    if (ctx_activate[dev] == 0) {
-        add_gpu_device_memory_usage(getpid(),dev,context_size,0); 
+    uint64_t before = 0;
+    uint64_t after = 0;
+    int hostpid;
+    int measure_context = 0;
+    size_t charge;
+    size_t measured_charge = 0;
+    size_t bytes_to_add = 0;
+
+    if (dev < 0 || dev >= CUDA_DEVICE_MAX_COUNT) {
+        return CUDA_OVERRIDE_CALL(cuda_library_entry,
+                                  cuDevicePrimaryCtxRetain, pctx, dev);
+    }
+
+    pthread_mutex_lock(&context_device_locks[dev]);
+    pthread_mutex_lock(&context_accounting_lock);
+    charge = context_charge_for_device(dev);
+    hostpid = get_current_host_pid();
+    if (charge == 0 && pidfound == 1 && hostpid > 0 &&
+        context_accounting[dev].charged_bytes == 0) {
+        measure_context = 1;
+    }
+    pthread_mutex_unlock(&context_accounting_lock);
+
+    if (measure_context) {
+        nvmlReturn_t result = get_used_gpu_memory_by_pid(
+            (unsigned int)hostpid, dev, &before);
+        if (result == NVML_SUCCESS || result == NVML_ERROR_NOT_FOUND) {
+            if (result == NVML_ERROR_NOT_FOUND) {
+                before = 0;
+            }
+        } else {
+            measure_context = 0;
+        }
+    }
+
+    CUresult res = CUDA_OVERRIDE_CALL(cuda_library_entry,
+                                      cuDevicePrimaryCtxRetain, pctx, dev);
+    if (res != CUDA_SUCCESS) {
+        pthread_mutex_unlock(&context_device_locks[dev]);
+        return res;
+    }
+
+    if (measure_context) {
+        int attempt;
+
+        for (attempt = 0; attempt < 10; attempt++) {
+            if (get_used_gpu_memory_by_pid((unsigned int)hostpid, dev,
+                                           &after) == NVML_SUCCESS &&
+                after > before) {
+                measured_charge = after - before;
+                break;
+            }
+            usleep(1000);
+        }
+    }
+    pthread_mutex_lock(&context_accounting_lock);
+    if (measured_charge > 0) {
+        device_context_size[dev] = measured_charge;
+        charge = measured_charge;
+        LOG_INFO("Measured primary context size lazily: "
+                 "dev=%d size=%lu", dev, charge);
+    } else if (charge == 0 && context_size > 0) {
+        /* Measurement failed or was skipped.  Charge the probed size, as main
+         * does on every device.  Not cached in device_context_size, so a later
+         * fresh retain can still measure this device for real. */
+        charge = context_size;
+        LOG_INFO("Primary context size unmeasured on device %d; charging the "
+                 "probed size %lu", dev, charge);
+    }
+    errno = 0;
+    int record_result =
+        (pidfound == 1)
+            ? primary_context_record_accounted_retain(
+                  &context_accounting[dev], charge, &bytes_to_add)
+            : primary_context_record_retain(&context_accounting[dev], charge,
+                                            &bytes_to_add);
+    /* The driver retain already succeeded.  An unknown context size must not
+     * fail the caller, so defer the charge to a later retain that knows it. */
+    if (record_result != 0 && errno == ENODATA) {
+        LOG_WARN("Primary context size unknown on device %d; charge is "
+                 "deferred to a later retain", dev);
+        record_result = primary_context_record_retain(
+            &context_accounting[dev], charge, &bytes_to_add);
     }
-    if (context_size>0) {
-        ctx_activate[dev] = 1;
+    if (record_result != 0) {
+        LOG_ERROR("Cannot account primary context retain on device %d",
+                  dev);
+        res = rollback_unaccounted_retain(dev, 0, 0);
+        pthread_mutex_unlock(&context_accounting_lock);
+        pthread_mutex_unlock(&context_device_locks[dev]);
+        return res;
     }
+    if (bytes_to_add > 0) {
+        if (add_gpu_device_memory_usage(getpid(), dev, bytes_to_add, 0) != 0) {
+            LOG_ERROR("Failed to charge primary context memory on device %d",
+                      dev);
+            res = rollback_unaccounted_retain(dev, 1, bytes_to_add);
+            pthread_mutex_unlock(&context_accounting_lock);
+            pthread_mutex_unlock(&context_device_locks[dev]);
+            return res;
+        }
+    }
+    ctx_activate[dev] = (int)context_accounting[dev].retain_count;
+    pthread_mutex_unlock(&context_accounting_lock);
+    pthread_mutex_unlock(&context_device_locks[dev]);
     return res;
 }
 
@@ -31,11 +203,59 @@ CUresult cuDevicePrimaryCtxSetFlags_v2( CUdevice dev, unsigned int  flags ){
 }
 
 CUresult cuDevicePrimaryCtxRelease_v2( CUdevice dev ){
-    if (ctx_activate[dev] == 1) {
-        rm_gpu_device_memory_usage(getpid(),dev,context_size,0);
+    size_t bytes_to_remove = 0;
+
+    if (dev < 0 || dev >= CUDA_DEVICE_MAX_COUNT) {
+        return CUDA_OVERRIDE_CALL(cuda_library_entry,
+                                  cuDevicePrimaryCtxRelease_v2, dev);
     }
-    ctx_activate[dev] = 0;
+    pthread_mutex_lock(&context_device_locks[dev]);
     CUresult res = CUDA_OVERRIDE_CALL(cuda_library_entry,cuDevicePrimaryCtxRelease_v2,dev);
+    if (res == CUDA_SUCCESS) {
+        pthread_mutex_lock(&context_accounting_lock);
+        if (primary_context_record_release(&context_accounting[dev],
+                                           &bytes_to_remove) != 0) {
+            LOG_WARN("Unbalanced primary context release on device %d", dev);
+        } else if (bytes_to_remove > 0) {
+            if (rm_gpu_device_memory_usage(getpid(), dev, bytes_to_remove,
+                                           0) != 0) {
+                primary_context_restore_charge(&context_accounting[dev],
+                                               bytes_to_remove);
+            }
+        }
+        ctx_activate[dev] = (int)context_accounting[dev].retain_count;
+        pthread_mutex_unlock(&context_accounting_lock);
+    }
+    pthread_mutex_unlock(&context_device_locks[dev]);
+    return res;
+}
+
+CUresult cuDevicePrimaryCtxReset_v2(CUdevice dev) {
+    size_t bytes_to_remove = 0;
+
+    if (dev < 0 || dev >= CUDA_DEVICE_MAX_COUNT) {
+        return CUDA_OVERRIDE_CALL(cuda_library_entry,
+                                  cuDevicePrimaryCtxReset_v2, dev);
+    }
+    pthread_mutex_lock(&context_device_locks[dev]);
+    CUresult res = CUDA_OVERRIDE_CALL(cuda_library_entry,
+                                      cuDevicePrimaryCtxReset_v2, dev);
+    if (res == CUDA_SUCCESS) {
+        pthread_mutex_lock(&context_accounting_lock);
+        /* The driver destroys the context whatever the retain count, so
+         * every outstanding retain and the charge go with it. */
+        bytes_to_remove = context_accounting[dev].charged_bytes;
+        primary_context_accounting_reset(&context_accounting[dev], 1);
+        if (bytes_to_remove > 0 &&
+            rm_gpu_device_memory_usage(getpid(), dev, bytes_to_remove,
+                                       0) != 0) {
+            primary_context_restore_charge(&context_accounting[dev],
+                                           bytes_to_remove);
+        }
+        ctx_activate[dev] = 0;
+        pthread_mutex_unlock(&context_accounting_lock);
+    }
+    pthread_mutex_unlock(&context_device_locks[dev]);
     return res;
 }
 
@@ -149,4 +369,3 @@ CUresult cuCtxSynchronize ( void ){
     CUresult res = CUDA_OVERRIDE_CALL(cuda_library_entry,cuCtxSynchronize);
     return res;
 }
-

src/cuda/context_accounting.c

@@ -0,0 +1,105 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Copyright (c) 2026 The HAMi Authors.
+ */
+
+#include "cuda/context_accounting.h"
+
+#include <errno.h>
+#include <limits.h>
+
+static int record_retain(primary_context_accounting_t *state,
+                         size_t context_bytes, size_t *bytes_to_add,
+                         int require_charge) {
+    if (state == NULL || bytes_to_add == NULL) {
+        errno = EINVAL;
+        return -1;
+    }
+    if (state->retain_count == UINT_MAX) {
+        errno = EOVERFLOW;
+        return -1;
+    }
+    if (require_charge && state->charged_bytes == 0 && context_bytes == 0) {
+        errno = ENODATA;
+        return -1;
+    }
+
+    *bytes_to_add = 0;
+    if (state->charged_bytes == 0 && context_bytes > 0) {
+        state->charged_bytes = context_bytes;
+        *bytes_to_add = context_bytes;
+    }
+    state->retain_count++;
+    return 0;
+}
+
+int primary_context_record_retain(primary_context_accounting_t *state,
+                                  size_t context_bytes,
+                                  size_t *bytes_to_add) {
+    return record_retain(state, context_bytes, bytes_to_add, 0);
+}
+
+int primary_context_record_accounted_retain(
+    primary_context_accounting_t *state, size_t context_bytes,
+    size_t *bytes_to_add) {
+    return record_retain(state, context_bytes, bytes_to_add, 1);
+}
+
+int primary_context_record_release(primary_context_accounting_t *state,
+                                   size_t *bytes_to_remove) {
+    if (state == NULL || bytes_to_remove == NULL) {
+        errno = EINVAL;
+        return -1;
+    }
+    if (state->retain_count == 0) {
+        errno = EINVAL;
+        return -1;
+    }
+
+    *bytes_to_remove = 0;
+    state->retain_count--;
+    if (state->retain_count == 0) {
+        *bytes_to_remove = state->charged_bytes;
+        state->charged_bytes = 0;
+    }
+    return 0;
+}
+
+int primary_context_rollback_retain(primary_context_accounting_t *state,
+                                    size_t bytes_to_add) {
+    if (state == NULL || state->retain_count == 0) {
+        errno = EINVAL;
+        return -1;
+    }
+    if (bytes_to_add > 0 && state->charged_bytes != bytes_to_add) {
+        errno = EINVAL;
+        return -1;
+    }
+
+    state->retain_count--;
+    if (bytes_to_add > 0) {
+        state->charged_bytes = 0;
+    }
+    return 0;
+}
+
+void primary_context_restore_charge(primary_context_accounting_t *state,
+                                    size_t context_bytes) {
+    if (state != NULL && state->retain_count == 0 && context_bytes > 0) {
+        state->charged_bytes = context_bytes;
+    }
+}
+
+void primary_context_accounting_reset(primary_context_accounting_t *states,
+                                      size_t state_count) {
+    size_t i;
+
+    if (states == NULL) {
+        return;
+    }
+    for (i = 0; i < state_count; i++) {
+        states[i].retain_count = 0;
+        states[i].charged_bytes = 0;
+    }
+}