[serving][post 7/24] Fixes tensor_parallel_degree detection on CPU (#…

…2229)
deepjavalibrary · Jul 25, 2024 · b752c31 · b752c31
1 parent a9c32a1
commit b752c31
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 8 deletions.
diff --git a/engines/python/src/main/java/ai/djl/python/engine/Connection.java b/engines/python/src/main/java/ai/djl/python/engine/Connection.java
@@ -131,9 +131,7 @@ static String[] getPythonStartCmd(
         String recommendedEntryPoint = pyEnv.getRecommendedEntryPoint();
 
         if (PyEnv.isMultiNode()) {
-
             int worldSize = tensorParallelDegree * pipelineParallelDegree;
-
             if (tensorParallelDegree * pipelineParallelDegree % clusterSize != 0) {
                 throw new IllegalArgumentException(
                         "Error: Cannot use cluster size: "
@@ -211,9 +209,7 @@ static String[] getPythonStartCmd(
             args[46] = "--recommended-entry-point";
             args[47] = recommendedEntryPoint == null ? "" : recommendedEntryPoint;
             return args;
-        }
-
-        if (pyEnv.isMpiMode()) {
+        } else if (pyEnv.isMpiMode()) {
             String cudaDevices = getVisibleDevices(workerId, tensorParallelDegree);
             logger.info("Set CUDA_VISIBLE_DEVICES={}", cudaDevices);
             String[] args = new String[42];

diff --git a/wlm/src/main/java/ai/djl/serving/wlm/LmiConfigRecommender.java b/wlm/src/main/java/ai/djl/serving/wlm/LmiConfigRecommender.java
@@ -145,9 +145,16 @@ private static void setTensorParallelDegree(Properties lmiProperties) {
         }
         String tpDegree = Utils.getenv("TENSOR_PARALLEL_DEGREE", "max");
         if ("max".equals(tpDegree)) {
-            tpDegree = String.valueOf(CudaUtils.getGpuCount());
+            int numGpus = CudaUtils.getGpuCount();
+            if (numGpus > 0) {
+                tpDegree = String.valueOf(numGpus);
+            } else {
+                tpDegree = null;
+            }
+        }
+        if (tpDegree != null) {
+            lmiProperties.setProperty("option.tensor_parallel_degree", tpDegree);
         }
-        lmiProperties.setProperty("option.tensor_parallel_degree", tpDegree);
     }
 
     private static void setPipelineParallelDegree(Properties lmiProperties) {

diff --git a/wlm/src/main/java/ai/djl/serving/wlm/LmiUtils.java b/wlm/src/main/java/ai/djl/serving/wlm/LmiUtils.java
@@ -68,7 +68,7 @@ static void configureLmiModel(ModelInfo<?, ?> modelInfo) throws ModelException {
 
         LmiConfigRecommender.configure(modelInfo, prop, modelConfig);
         logger.info(
-                "Detected mpi_mode: {}, rolling_batch: {}, tensor_parallel_degree {}, for"
+                "Detected mpi_mode: {}, rolling_batch: {}, tensor_parallel_degree: {}, for"
                         + " modelType: {}",
                 prop.getProperty("option.mpi_mode"),
                 prop.getProperty("option.rolling_batch"),

diff --git a/wlm/src/main/java/ai/djl/serving/wlm/ModelInfo.java b/wlm/src/main/java/ai/djl/serving/wlm/ModelInfo.java
@@ -36,6 +36,7 @@
 import ai.djl.serving.wlm.util.EventManager;
 import ai.djl.serving.wlm.util.WlmConfigManager;
 import ai.djl.serving.wlm.util.WlmOutOfMemoryException;
+import ai.djl.translate.NoopServingTranslatorFactory;
 import ai.djl.translate.TranslateException;
 import ai.djl.util.NeuronUtils;
 import ai.djl.util.Utils;
@@ -262,6 +263,9 @@ public void load(Device device) throws ModelException, IOException {
                 // override model_id
                 builder.optOption("model_id", downloadDir.toAbsolutePath().toString());
             }
+            if (translator == null && translatorFactory == null && "Python".equals(engineName)) {
+                builder.optTranslatorFactory(new NoopServingTranslatorFactory());
+            }
             ZooModel<I, O> m = builder.build().loadModel();
             m.setProperty("metric_dimension", id);