[tnx] update default neuron rolling batch for correct mpi mode config (…

…#1789)
deepjavalibrary · Apr 20, 2024 · 8bf6106 · 8bf6106
1 parent 812c164
commit 8bf6106
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/wlm/src/main/java/ai/djl/serving/wlm/LmiConfigRecommender.java b/wlm/src/main/java/ai/djl/serving/wlm/LmiConfigRecommender.java
@@ -87,6 +87,8 @@ private static void setRollingBatch(
         } else if (!isTextGenerationModel(modelConfig)) {
             // Non text-generation use-cases are not compatible with rolling batch
             rollingBatch = "disable";
+        } else if (isTnxEnabled(features)) {
+            rollingBatch = "tnx";
         } else if (isLmiDistEnabled(features)
                 && "lmi-dist".equals(MODEL_TO_ROLLING_BATCH.get(modelType))) {
             rollingBatch = "lmi-dist";
@@ -175,6 +177,10 @@ private static boolean isTrtLlmEnabled(String features) {
         return features != null && features.contains("trtllm");
     }
 
+    private static boolean isTnxEnabled(String features) {
+        return features != null && features.contains("tnx");
+    }
+
     private static boolean isT5TrtLlm(
             LmiUtils.HuggingFaceModelConfig modelConfig, String features) {
         return isTrtLlmEnabled(features) && "t5".equals(modelConfig.getModelType());