Add pid and per-pid GPU memory usage to get_gpu_info_c

Author: pseudotensor

src/gpu/utils.cu

@@ -52,7 +52,16 @@ int get_compute_capability(int d_idx, int *major, int *minor, int *ratioperf) {
 }
 
 
-void get_gpu_info_c(unsigned int *n_gpus, int *gpu_percent_usage, unsigned long long *gpu_total_memory, unsigned long long *gpu_free_memory, char **gpu_name, int *majors, int *minors) {
+void get_gpu_info_c(unsigned int *n_gpus, int *gpu_percent_usage, unsigned long long *gpu_total_memory,
+ unsigned long long *gpu_free_memory,
+ char **gpu_name,
+ int *majors, int *minors,
+ unsigned int *num_pids, unsigned int *pids, unsigned long long *usedGpuMemorys) {
+
+  bool verbose=false;
+  if(verbose){
+    std::cerr << "inside get_gpu_info_c c function" << std::endl;
+  }
 
   nvmlReturn_t rv;
   rv = nvmlInit();
@@ -69,17 +78,25 @@ void get_gpu_info_c(unsigned int *n_gpus, int *gpu_percent_usage, unsigned long
     nvmlReturn_t rv;
     rv = nvmlDeviceGetHandleByIndex(i, &device);
     assert(rv == NVML_SUCCESS);
+
     nvmlUtilization_t utilization;
     rv = nvmlDeviceGetUtilizationRates(device, &utilization);
     assert(rv == NVML_SUCCESS);
     gpu_percent_usage[i] = utilization.gpu;
+
+    if(verbose){
+      std::cerr << "i=" << i << " usage=" << gpu_percent_usage[i] << std::endl;
+    }
+
     nvmlMemory_t memory;
     rv = nvmlDeviceGetMemoryInfo(device, &memory);
     assert(rv == NVML_SUCCESS);
     gpu_total_memory[i] = memory.total;
     gpu_free_memory[i] = memory.free;
+
     rv = nvmlDeviceGetName(device, gpu_name[i], 30);
     assert(rv == NVML_SUCCESS);
+
 #if (CUDART_VERSION >= 9000)
     rv = nvmlDeviceGetCudaComputeCapability(device, &majors[i], &minors[i]);
     assert(rv == NVML_SUCCESS);
@@ -91,6 +108,34 @@ void get_gpu_info_c(unsigned int *n_gpus, int *gpu_percent_usage, unsigned long
     majors[i] = -1;
     minors[i] = -1;
 #endif
+
+    if(verbose){
+      std::cerr << "i=" << i << " majors=" << majors[i] << " minors=" << minors[i] << std::endl;
+    }
+
+    unsigned int max_pids=2000;
+    unsigned int infoCount;
+    //rv = nvmlDeviceGetComputeRunningProcesses(device, &infoCount, NULL);
+    //assert(rv == NVML_SUCCESS);
+    infoCount = max_pids;
+    nvmlProcessInfo_t infos[infoCount];
+    unsigned int num_pid_local;
+    num_pids[i] = infoCount;
+    rv = nvmlDeviceGetComputeRunningProcesses(device, &num_pids[i], infos);
+    assert(rv == NVML_SUCCESS);
+    if(num_pids[i] > max_pids){
+       std::cerr << "Too many pids: " << num_pids[i] << " .  Increase max_pids: " << max_pids << std::endl;
+       assert(num_pids[i] <= max_pids);
+    }
+    for (unsigned int pidi=0; pidi < num_pids[i]; pidi++) {
+        pids[pidi + i * max_pids] = infos[pidi].pid;
+        usedGpuMemorys[pidi + i * max_pids] = infos[pidi].usedGpuMemory;
+
+        if(verbose){
+          std::cerr << "i=" << i << " pidi=" << pidi << " pids=" << pids[pidi + i * max_pids] << " gpumemory=" << usedGpuMemorys[pidi + i * max_pids] << std::endl;
+        }
+
+    }
   }
 
 }

src/interface_py/h2o4gpu/util/gpu.py

@@ -148,33 +148,42 @@ def get_gpu_info_subprocess(return_usage=False):
     return (total_gpus, total_mem, gpu_type)
 
 
-def get_gpu_info_c(return_usage=False, return_capability=False,
-                   return_all=False, verbose=False):
+def get_gpu_info_c(return_usage=False,
+                   return_free_memory=False,
+                   return_capability=False,
+                   return_memory_by_pid=False,
+                   return_all=False,
+                   verbose=False):
     """Gets the GPU info from C call
 
     :return:
         Total number of GPUs and total available memory
          (and optionally GPU usage)
     """
-    max_gpus = 128
+    max_gpus = 16
     total_gpus = 0
     total_gpus_actual = 0
     which_gpus = []
-    usages_tmp = np.empty(max_gpus, dtype=np.int32)
-    total_mems_tmp = np.empty(max_gpus, dtype=np.uint64)
-    free_mems_tmp = np.empty(max_gpus, dtype=np.uint64)
+    usages_tmp = np.zeros(max_gpus, dtype=np.int32)
+    total_mems_tmp = np.zeros(max_gpus, dtype=np.uint64)
+    free_mems_tmp = np.zeros(max_gpus, dtype=np.uint64)
     # This 30 should be same as the gpu type in get_gpu_info_c
     gpu_types_tmp = [' ' * 30 for _ in range(64)]
-    majors_tmp = np.empty(max_gpus, dtype=np.int32)
-    minors_tmp = np.empty(max_gpus, dtype=np.int32)
+    majors_tmp = np.zeros(max_gpus, dtype=np.int32)
+    minors_tmp = np.zeros(max_gpus, dtype=np.int32)
+    max_pids = 2000
+    num_pids_tmp = np.zeros(max_pids, dtype=np.uint32)
+    pids_tmp = np.zeros(max_pids * max_gpus, dtype=np.uint32)
+    usedGpuMemorys_tmp = np.zeros(max_pids * max_gpus, dtype=np.uint64)
 
     try:
         from ..libs.lib_utils import GPUlib
         lib = GPUlib().get()
 
         total_gpus_actual = \
             lib.get_gpu_info_c(usages_tmp, total_mems_tmp, free_mems_tmp,
-                               gpu_types_tmp, majors_tmp, minors_tmp)
+                               gpu_types_tmp, majors_tmp, minors_tmp,
+                               num_pids_tmp, pids_tmp, usedGpuMemorys_tmp)
 
         # This will drop the GPU count, but the returned usage
         total_gpus, which_gpus = cuda_vis_check(total_gpus_actual)
@@ -185,7 +194,10 @@ def get_gpu_info_c(return_usage=False, return_capability=False,
     # pylint: disable=broad-except
     except Exception as e:
         if verbose:
+            import sys
+            sys.stderr.write("Exception: %s" % str(e))
             print(e)
+            sys.stdout.flush()
 
     if return_capability or return_all:
         if list(minors_tmp)[0] == -1:
@@ -198,32 +210,49 @@ def get_gpu_info_c(return_usage=False, return_capability=False,
     usages_actual = np.resize(usages_tmp, total_gpus_actual)
     majors_actual = np.resize(majors_tmp, total_gpus_actual)
     minors_actual = np.resize(minors_tmp, total_gpus_actual)
+    num_pids_actual = np.resize(num_pids_tmp, total_gpus_actual)
+    pids_actual = np.resize(pids_tmp, total_gpus_actual * max_pids)
+    usedGpuMemorys_actual = np.resize(usedGpuMemorys_tmp,
+                                      total_gpus_actual * max_pids)
 
     total_mems = np.resize(np.copy(total_mems_actual), total_gpus)
     free_mems = np.resize(np.copy(free_mems_actual), total_gpus)
     gpu_types = np.resize(np.copy(gpu_types_actual), total_gpus)
     usages = np.resize(np.copy(usages_actual), total_gpus)
     majors = np.resize(np.copy(majors_actual), total_gpus)
     minors = np.resize(np.copy(minors_actual), total_gpus)
+    num_pids = np.resize(np.copy(num_pids_actual), total_gpus)
+    pids = np.resize(np.copy(pids_actual), total_gpus * max_pids)
+    usedGpuMemorys = np.resize(np.copy(usedGpuMemorys_actual),
+                               total_gpus * max_pids)
+
     gpu_i = 0
     for j in range(total_gpus_actual):
         if j in which_gpus:
             total_mems[gpu_i] = total_mems_actual[j]
             free_mems[gpu_i] = free_mems_actual[j]
             gpu_types[gpu_i] = gpu_types_actual[j]
             usages[gpu_i] = usages_actual[j]
+            minors[gpu_i] = minors_actual[j]
+            majors[gpu_i] = majors_actual[j]
+            num_pids[gpu_i] = num_pids_actual[j]
+            pids[gpu_i] = pids_actual[j]
+            usedGpuMemorys[gpu_i] = usedGpuMemorys_actual[j]
             gpu_i += 1
-
-    if return_all:
-        return (total_gpus, total_mems,
-                free_mems, gpu_types, usages, majors, minors)
-    if return_usage and return_capability:
-        return (total_gpus, total_mems, gpu_types, usages, majors, minors)
-    if return_usage:
-        return (total_gpus, total_mems, gpu_types, usages)
-    if return_capability:
-        return (total_gpus, total_mems, gpu_types, majors, minors)
-    return (total_gpus, total_mems, gpu_types)
+    pids = np.reshape(pids, (total_gpus, max_pids))
+    usedGpuMemorys = np.reshape(usedGpuMemorys, (total_gpus, max_pids))
+
+    to_return = [total_gpus, total_mems, gpu_types]
+    if return_all or return_usage:
+        to_return.append(usages)
+    if return_all or return_free_memory:
+        to_return.append(free_mems)
+    if return_all or return_capability:
+        to_return.extend([majors, minors])
+    if return_all or return_memory_by_pid:
+        to_return.extend([num_pids, pids, usedGpuMemorys])
+
+    return tuple(to_return)
 
 
 def cudaresetdevice(gpu_id, n_gpus):

src/swig/util/gpu.i

@@ -2,7 +2,11 @@
 %{
 extern int cudaresetdevice(int wDev, int nDev);
 extern int get_compute_capability(int d_idx, int *major, int *minor, int *ratioperf);
-extern void get_gpu_info_c(unsigned int *n_gpus, int *gpu_percent_usage, unsigned long long *gpu_total_memory, unsigned long long *gpu_free_memory, char **gpu_name, int *majors, int *minors);
+extern void get_gpu_info_c(unsigned int *n_gpus, int *gpu_percent_usage,
+ unsigned long long *gpu_total_memory, unsigned long long *gpu_free_memory,
+ char **gpu_name,
+ int *majors, int *minors,
+ unsigned int *num_pids, unsigned int *pids, unsigned long long *usedGpuMemorys);
 extern int cudaresetdevice_bare(void);
 %}
 
@@ -38,15 +42,23 @@ extern int cudaresetdevice_bare(void);
   free((char *) $1);
 }
 
-%apply int *OUTPUT {int *major, int *minor, int *ratioperf}
-%apply int *OUTPUT {unsigned int *n_gpus}
+%apply int *OUTPUT {int *major, int *minor, int *ratioperf};
+
+%apply int *OUTPUT {unsigned int *n_gpus};
 %apply (int *INPLACE_ARRAY1) {int *gpu_percent_usage};
 %apply (int *INPLACE_ARRAY1) {int *majors};
 %apply (int *INPLACE_ARRAY1) {int *minors};
 %apply (unsigned long long *INPLACE_ARRAY1) {unsigned long long *gpu_total_memory};
 %apply (unsigned long long *INPLACE_ARRAY1) {unsigned long long *gpu_free_memory};
+%apply (unsigned int *INPLACE_ARRAY1) {unsigned int *num_pids};
+%apply (unsigned int *INPLACE_ARRAY1) {unsigned int *pids};
+%apply (unsigned long long *INPLACE_ARRAY1) {unsigned long long *usedGpuMemorys};
 
 extern int cudaresetdevice(int wDev, int nDev);
 extern int get_compute_capability(int d_idx, int *major, int *minor, int *ratioperf);
-extern void get_gpu_info_c(unsigned int *n_gpus, int *gpu_percent_usage, unsigned long long *gpu_total_memory, unsigned long long *gpu_free_memory, char **gpu_name, int *majors, int *minors);
+extern void get_gpu_info_c(unsigned int *n_gpus, int *gpu_percent_usage,
+ unsigned long long *gpu_total_memory, unsigned long long *gpu_free_memory,
+ char **gpu_name,
+ int *majors, int *minors,
+ unsigned int *num_pids, unsigned int *pids, unsigned long long *usedGpuMemorys);
 extern int cudaresetdevice_bare(void);