Restrict GPU access from worker to deviceIds

pytorch · Jul 2, 2024 · 2282882 · 2282882
1 parent e979f41
commit 2282882
Showing 1 changed file with 7 additions and 13 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
@@ -130,16 +130,13 @@ private void startWorkerPython(int port, String deviceIds)
                                 modelPath.getAbsolutePath(),
                                 model.getModelArchive().getManifest().getModel().getHandler())));
 
-        if (model.getParallelLevel() > 0) {
-            if (model.getParallelType() != ParallelType.CUSTOM) {
-                attachRunner(argl, envp, port, deviceIds);
-            } else {
-                if (deviceIds != null) {
-                    envp.add("CUDA_VISIBLE_DEVICES=" + deviceIds);
-                }
-                argl.add(EnvironmentUtils.getPythonRunTime(model));
-            }
-        } else if (model.getParallelLevel() == 0) {
+        if (deviceIds != null) {
+            envp.add("CUDA_VISIBLE_DEVICES=" + deviceIds);
+        }
+
+        if (model.getParallelLevel() > 0 && model.getParallelType() != ParallelType.CUSTOM) {
+            attachRunner(argl, envp, port, deviceIds);
+        } else {
             argl.add(EnvironmentUtils.getPythonRunTime(model));
         }
 
@@ -291,9 +288,6 @@ private void startWorkerCPP(int port, String runtimeType, String deviceIds)
     private void attachRunner(
             ArrayList<String> argl, List<String> envp, int port, String deviceIds) {
         envp.add("LOGLEVEL=INFO");
-        if (deviceIds != null) {
-            envp.add("CUDA_VISIBLE_DEVICES=" + deviceIds);
-        }
         ModelConfig.TorchRun torchRun = model.getModelArchive().getModelConfig().getTorchRun();
         envp.add(String.format("OMP_NUM_THREADS=%d", torchRun.getOmpNumberThreads()));
         argl.add("torchrun");