Use process build for GPU command (#2632)

* Add GPU ID validation * update * updates * Update WorkerThread.java
pytorch · Sep 29, 2023 · 7b1c359 · 7b1c359
1 parent bfb3d42
commit 7b1c359
Showing 1 changed file with 10 additions and 6 deletions.
diff --git a/frontend/server/src/main/java/org/pytorch/serve/wlm/WorkerThread.java b/frontend/server/src/main/java/org/pytorch/serve/wlm/WorkerThread.java
@@ -120,12 +120,16 @@ public String getGpuUsage() {
             try {
                 // TODO : add a generic code to capture gpu details for different devices instead of
                 // just NVIDIA
-                process =
-                        Runtime.getRuntime()
-                                .exec(
-                                        "nvidia-smi -i "
-                                                + gpuId
-                                                + " --query-gpu=utilization.gpu,utilization.memory,memory.used --format=csv");
+                ProcessBuilder pb =
+                        new ProcessBuilder(
+                                "nvidia-smi",
+                                "-i",
+                                String.valueOf(gpuId),
+                                "--query-gpu=utilization.gpu,utilization.memory,memory.used",
+                                "--format=csv");
+
+                // Start the process
+                process = pb.start();
                 process.waitFor();
                 int exitCode = process.exitValue();
                 if (exitCode != 0) {