chore: device count

api.h

@@ -7,4 +7,6 @@
 
 #define RPC_nvmlDeviceGetName 3
 
+#define RPC_nvmlDeviceGetCount_v2 4
+
 #endif

client.cu

@@ -145,6 +145,12 @@ nvmlReturn_t nvmlInit_v2()
     return send_rpc_message(RPC_nvmlInit_v2);
 }
 
+nvmlReturn_t nvmlDeviceGetCountHandler_v2()
+{
+    open_rpc_client();
+    return send_rpc_message(RPC_nvmlDeviceGetCount_v2);
+}
+
 nvmlReturn_t nvmlShutdown()
 {
     open_rpc_client();
@@ -163,12 +169,14 @@ void *dlsym(void *handle, const char *name) __THROW
 
     if (!strcmp(name, "nvmlInitWithFlags"))
         return (void *)nvmlInitWithFlags;
-    // if (!strcmp(name, "nvmlInit_v2"))
-    //     return (void *)nvmlInit_v2;
-    // if (!strcmp(name, "nvmlShutdown"))
-    //     return (void *)nvmlShutdown;
-    // if (!strcmp(name, "nvmlDeviceGetName"))
-    //     return (void *)nvmlDeviceGetName;
+    if (!strcmp(name, "nvmlInit_v2"))
+        return (void *)nvmlInit_v2;
+    if (!strcmp(name, "nvmlShutdown"))
+        return (void *)nvmlShutdown;
+    if (!strcmp(name, "nvmlDeviceGetName"))
+        return (void *)nvmlDeviceGetName;
+    if (!strcmp(name, "nvmlDeviceGetCount_v2"))
+        return (void *)nvmlDeviceGetCountHandler_v2;
 
     static void *(*real_dlsym)(void *, const char *) = NULL;
     if (real_dlsym == NULL)

local.sh

@@ -3,13 +3,13 @@
 libscuda_path="$(pwd)/libscuda.so"
 client_path="$(pwd)/client.cu"
 server_path="$(pwd)/server.cu"
-server_out_path="$(pwd)/server"
+server_out_path="$(pwd)/server.so"
 
 build() {
   echo "building client..."
 
   if [[ "$(uname)" == "Linux" ]]; then
-    nvcc -shared -Xcompiler -fPIC -o $libscuda_path $client_path -lcudart
+    nvcc -Xcompiler -fPIC -shared -o $libscuda_path $client_path 
   else
     echo "No compiler options set for os "$(uname)""
   fi
@@ -24,12 +24,12 @@ server() {
   echo "building server..." 
 
   if [[ "$(uname)" == "Linux" ]]; then
-    nvcc -shared -Xcompiler -fPIC -o $server_out_path $server_path -lcudart
+    nvcc -o $server_out_path $server_path -lnvidia-ml
   else
     echo "No compiler options set for os "$(uname)""
   fi
 
-  echo "starting server..."
+  echo "starting server... $server_out_path"
 
   "$server_out_path"
 }

server.cu

@@ -69,8 +69,22 @@ void *client_handler(void *arg)
             printf("Sending nvmlInitWithFlags response %d %d\n", request_id, result);
             if (write(connfd, &result, sizeof(nvmlReturn_t)) < 0)
                 goto exit;
-            break;
         }
+        case RPC_nvmlDeviceGetCount_v2:
+        {
+            unsigned int dcount;
+            if (read(connfd, &dcount, sizeof(unsigned int)) < 0)
+                goto exit;
+            printf("Received device count request %d %d\n", request_id, dcount);
+            nvmlReturn_t result = nvmlDeviceGetCount_v2(&dcount);
+
+            if (result == NVML_SUCCESS) {
+                printf("Number of devices: %u\n", dcount);
+            }
+
+            if (write(connfd, &result, sizeof(nvmlReturn_t)) < 0)
+                goto exit;
+            }
         }
     }