let users determine the availability of GPUs

client.c

@@ -849,3 +849,31 @@ void c_unset_env() {
     send_msg(server_socket, &m);
     send_bytes(server_socket, command_line.label, m.u.size);
 }
+
+void c_set_free_percentage() {
+    struct Msg m;
+    m.type = SET_FREE_PERC;
+    m.u.size = command_line.gpus;
+    send_msg(server_socket, &m);
+}
+
+void c_get_free_percentage() {
+    struct Msg m;
+    int res;
+
+    m.type = GET_FREE_PERC;
+    send_msg(server_socket, &m);
+
+    /* Receive the answer */
+    res = recv_msg(server_socket, &m);
+    if (res != sizeof(m))
+        error("Error in get_free_percentage");
+
+    switch (m.type) {
+        case GET_FREE_PERC:
+            printf("%d\n", m.u.size);
+            return;
+        default:
+            warning("Wrong internal message in get_free_percentage");
+    }
+}

gpu.c

@@ -10,6 +10,7 @@
 
 #define TS_VISIBLE_DEVICES "TS_VISIBLE_DEVICES"
 
+extern int free_percentage;
 int *used_gpus;
 int num_total_gpus;
 
@@ -86,7 +87,7 @@ int * getGpuList(int *num) {
             goto Error;
         }
 
-        if (mem.free > .9 * mem.total)
+        if (mem.free > free_percentage / 100. * mem.total)
             gpuList[count++] = visible[i];
     }
     *num = count;

jobs.c

@@ -22,6 +22,8 @@ int num_total_gpus;
 int busy_slots = 0;
 int max_slots = 1;
 
+int free_percentage = 90;
+
 struct Notify {
     int socket;
     int jobid;
@@ -1632,3 +1634,17 @@ void s_unset_env(int s, int size) {
     unsetenv(var);
     free(var);
 }
+
+void s_set_free_percentage(int new_percentage) {
+    if (new_percentage > 0)
+        free_percentage = new_percentage;
+    else
+        warning("Received new_percentage=%i", new_percentage);
+}
+
+void s_get_free_percentage(int s) {
+    struct Msg m;
+    m.type = GET_FREE_PERC;
+    m.u.size = free_percentage;
+    send_msg(s, &m);
+}

main.c

@@ -100,6 +100,8 @@ static struct option longOptions[] = {
         {"getenv",        required_argument, NULL, 0},
         {"setenv",        required_argument, NULL, 0},
         {"unsetenv",      required_argument, NULL, 0},
+        {"set_gpu_free_perc", required_argument, NULL, 0},
+        {"get_gpu_free_perc", no_argument,       NULL, 0},
         {NULL, 0,                            NULL, 0}
 };
 
@@ -127,6 +129,11 @@ void parse_opts(int argc, char **argv) {
                 } else if (strcmp(longOptions[optionIdx].name, "unsetenv") == 0) {
                     command_line.request = c_UNSET_ENV;
                     command_line.label = optarg;  /* reuse this var */
+                } else if (strcmp(longOptions[optionIdx].name, "set_gpu_free_perc") == 0) {
+                    command_line.request = c_SET_FREE_PERC;
+                    command_line.gpus = atoi(optarg); /* reuse this var */
+                } else if (strcmp(longOptions[optionIdx].name, "get_gpu_free_perc") == 0) {
+                    command_line.request = c_GET_FREE_PERC;
                 } else
                     error("Wrong option %s.", longOptions[optionIdx].name);
                 break;
@@ -657,6 +664,16 @@ int main(int argc, char **argv) {
                 error("The command %i needs the server", command_line.request);
             c_unset_env();
             break;
+        case c_SET_FREE_PERC:
+            if (!command_line.need_server)
+                error("The command %i needs the server", command_line.request);
+            c_set_free_percentage();
+            break;
+        case c_GET_FREE_PERC:
+            if (!command_line.need_server)
+                error("The command %i needs the server", command_line.request);
+            c_get_free_percentage();
+            break;
     }
 
     if (command_line.need_server) {

main.h

@@ -50,7 +50,9 @@ enum MsgTypes {
     GET_CMD,
     GET_ENV,
     SET_ENV,
-    UNSET_ENV
+    UNSET_ENV,
+    SET_FREE_PERC,
+    GET_FREE_PERC
 };
 
 enum Request {
@@ -80,7 +82,9 @@ enum Request {
     c_SHOW_CMD,
     c_GET_ENV,
     c_SET_ENV,
-    c_UNSET_ENV
+    c_UNSET_ENV,
+    c_SET_FREE_PERC,
+    c_GET_FREE_PERC
 };
 
 struct CommandLine {
@@ -293,6 +297,10 @@ void c_set_env();
 
 void c_unset_env();
 
+void c_set_free_percentage();
+
+void c_get_free_percentage();
+
 /* jobs.c */
 void s_list(int s);
 
@@ -364,6 +372,10 @@ void s_set_env(int s, int size);
 
 void s_unset_env(int s, int size);
 
+void s_set_free_percentage(int new_percentage);
+
+void s_get_free_percentage(int s);
+
 /* server.c */
 void server_main(int notify_fd, char *_path);
 

server.c

@@ -480,6 +480,12 @@ client_read(int index) {
         case UNSET_ENV:
             s_unset_env(s, m.u.size);
             break;
+        case SET_FREE_PERC:
+            s_set_free_percentage(m.u.size);
+            break;
+        case GET_FREE_PERC:
+            s_get_free_percentage(s);
+            break;
         case GET_VERSION:
             s_send_version(s);
             break;