removed hard codes

justanhduc · Nov 4, 2021 · 0cf6dc0 · 0cf6dc0
1 parent 1c1d28c
commit 0cf6dc0
Show file tree

Hide file tree

Showing 5 changed files with 17 additions and 17 deletions.
diff --git a/gpu.c b/gpu.c
@@ -7,7 +7,7 @@
 
 #include "main.h"
 
-int * getFreeGpuList(int *numFree) {
+int * getGpuList(int *num, int unoccupied) {
     int * gpuList;
     unsigned int nDevices;
     int i, j = 0, count = 0;
@@ -31,19 +31,22 @@ int * getFreeGpuList(int *numFree) {
             goto Error;
         }
 
-        result = nvmlDeviceGetMemoryInfo(dev, &mem);
-        if (result != 0) {
-            error("Failed to get GPU memory for GPU %d: %s", i, nvmlErrorString(result));
-            goto Error;
-        }
+        if (unoccupied) {
+            result = nvmlDeviceGetMemoryInfo(dev, &mem);
+            if (result != 0) {
+                error("Failed to get GPU memory for GPU %d: %s", i, nvmlErrorString(result));
+                goto Error;
+            }
 
-        if (mem.free > .1 * mem.total) {
-            gpuList[j] = i;
-            count++;
-            j++;
+            if (mem.free < .1 * mem.total)
+                continue;
         }
+
+        gpuList[j] = i;
+        count++;
+        j++;
     }
-    *numFree = count;
+    *num = count;
     result = nvmlShutdown();
     if (NVML_SUCCESS != result)
         error("Failed to shutdown NVML: %s", nvmlErrorString(result));

diff --git a/jobs.c b/jobs.c
@@ -660,7 +660,7 @@ int next_run_job() {
 
     /* Query GPUs */
     int numFree;
-    int *freeGpuList = getFreeGpuList(&numFree);
+    int *freeGpuList = getGpuList(&numFree, 1);
 
     /* Look for a runnable task */
     p = firstjob;

diff --git a/main.h b/main.h
@@ -457,4 +457,4 @@ char *get_environment();
 int tail_file(const char *fname, int last_lines);
 
 /* gpu.c */
-int *getFreeGpuList(int *numFree);
+int *getGpuList(int *num, int unoccupied);
diff --git a/server.c b/server.c
@@ -74,8 +74,7 @@ int *used_gpus;
 int num_total_gpus;
 
 static void initialize_gpus() {
-    /* TODO: fixed hard code */
-    num_total_gpus = 100;
+    getGpuList(&num_total_gpus, 0);
     used_gpus = (int *) malloc(num_total_gpus * sizeof(int));
     memset(used_gpus, 0, num_total_gpus * sizeof(int));  /* 0 is not in used, 1 is in used */
 }

diff --git a/server_start.c b/server_start.c
@@ -19,8 +19,6 @@
 #include "main.h"
 
 extern int server_socket;
-int *used_gpus;
-int num_total_gpus;
 
 static char *socket_path;
 static int should_check_owner = 0;