Merge pull request #136 from Project-HAMi/fix_vllm

Fix vllm issues: Project-HAMi/HAMi#1381, Project-HAMi/HAMi#1461

Author: archlitchi

src/cuda/hook.c

@@ -134,7 +134,7 @@ cuda_entry_t cuda_library_entry[] = {
     {.name = "cuMemCreate"},
     {.name = "cuMemRelease"},
     {.name = "cuMemMap"},
-    {.name = "cucuMemImportFromShareableHandle"},
+    {.name = "cuMemImportFromShareableHandle"},
     {.name = "cuMemAllocAsync"},
     {.name = "cuMemFreeAsync"},
     /* cuda11.7 new api memory part */

src/cuda/memory.c

@@ -493,8 +493,8 @@ 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);
-    size_t limit = get_current_device_memory_limit(dev);
+    size_t usage = get_current_device_memory_usage(cuda_to_nvml_map(dev));
+    size_t limit = get_current_device_memory_limit(cuda_to_nvml_map(dev));
     if (limit == 0) {
         CUDA_OVERRIDE_CALL(cuda_library_entry,cuMemGetInfo_v2, free, total);
         LOG_INFO("orig free=%ld total=%ld", *free, *total);

src/libvgpu.c

@@ -848,7 +848,7 @@ void preInit(){
 
 void postInit(){
     allocator_init();
-
+    map_cuda_visible_devices();
     try_lock_unified_lock();
     nvmlReturn_t res = set_task_pid();
     try_unlock_unified_lock();
@@ -860,7 +860,6 @@ void postInit(){
         pidfound=1;
     }
 
-    map_cuda_visible_devices();
     //add_gpu_device_memory_usage(getpid(),0,context_size,0);
     env_utilization_switch = set_env_utilization_switch();
     init_utilization_watcher();

src/utils.c

@@ -193,7 +193,7 @@ int parse_cuda_visible_env() {
     for (int i = 0; i < CUDA_DEVICE_MAX_COUNT; i++) {
         LOG_INFO("device %d -> %d",i,cuda_to_nvml_map(i));
     }
-    LOG_DEBUG("get default cuda from %s",getenv("CUDA_VISIBLE_DEVICES"));
+    LOG_INFO("get default cuda from %s", getenv("CUDA_VISIBLE_DEVICES"));
     return count;
 }