[python] Adjuests mpi workers based CUDA_VISIBLE_DEVICES (deepjavalib…

…rary#782)
KexinFeng · Aug 16, 2023 · 64d412d · 64d412d
1 parent 3e23857
commit 64d412d
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/engines/python/src/main/java/ai/djl/python/engine/PyEnv.java b/engines/python/src/main/java/ai/djl/python/engine/PyEnv.java
@@ -261,6 +261,14 @@ public void setTensorParallelDegree(int tensorParallelDegree) {
 
     int getMpiWorkers() {
         int gpuCount = CudaUtils.getGpuCount();
+        String visibleDevices = Utils.getenv("CUDA_VISIBLE_DEVICES");
+        if (gpuCount > 0 && visibleDevices != null) {
+            int visibleCount = visibleDevices.split(",").length;
+            if (visibleCount > gpuCount || visibleCount < 1) {
+                throw new AssertionError("Invalid CUDA_VISIBLE_DEVICES: " + visibleDevices);
+            }
+            gpuCount = visibleCount;
+        }
         return gpuCount / getTensorParallelDegree();
     }