Restrict GPU access from worker to deviceIds

pytorch · Jul 2, 2024 · 4f184ce · 4f184ce
1 parent e979f41
commit 4f184ce
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/frontend/server/src/main/java/org/pytorch/serve/wlm/WorkerLifeCycle.java b/frontend/server/src/main/java/org/pytorch/serve/wlm/WorkerLifeCycle.java
@@ -140,6 +140,9 @@ private void startWorkerPython(int port, String deviceIds)
                 argl.add(EnvironmentUtils.getPythonRunTime(model));
             }
         } else if (model.getParallelLevel() == 0) {
+            if (deviceIds != null) {
+                envp.add("CUDA_VISIBLE_DEVICES=" + deviceIds);
+            }
             argl.add(EnvironmentUtils.getPythonRunTime(model));
         }