resolve review comments

deepjavalibrary · Mar 28, 2024 · 69ab513 · 69ab513
1 parent 11f96b5
commit 69ab513
Showing 1 changed file with 10 additions and 6 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
@@ -103,9 +103,7 @@ private static void setEngine(
             lmiProperties.setProperty("option.mpi_mode", "true");
         }
         //  TODO TrtLLM python backend: Change it once TrtLLM supports T5 with inflight batching.
-        if (isTrtLLMEnabled(features)
-                && "disable".equals(rollingBatch)
-                && "t5".equals(modelConfig.getModelType())) {
+        if (isT5TrtLLM(modelConfig, features, rollingBatch)) {
             engine = "MPI";
             lmiProperties.setProperty("option.mpi_mode", "true");
         }
@@ -129,9 +127,8 @@ private static void setDynamicBatch(
             ModelInfo<?, ?> modelInfo,
             String features) {
         // TODO TrtLLM python backend: Change it once TrtLLM supports T5 with inflight batching.
-        if ("t5".equals(modelConfig.getModelType())
-                && "disable".equals(lmiProperties.getProperty("option.rolling_batch"))
-                && isTrtLLMEnabled(features)) {
+        String rollingBatch = lmiProperties.getProperty("option.rolling_batch");
+        if (isT5TrtLLM(modelConfig, features, rollingBatch)) {
 
             // To do runtime compilation for TensorRT-LLM T5 model.
             lmiProperties.setProperty("trtllm_python_backend", String.valueOf(true));
@@ -156,6 +153,13 @@ private static boolean isTrtLLMEnabled(String features) {
         return features != null && features.contains("trtllm");
     }
 
+    private static boolean isT5TrtLLM(
+            LmiUtils.HuggingFaceModelConfig modelConfig, String features, String rollingBatch) {
+        return isTrtLLMEnabled(features)
+                && "disable".equals(rollingBatch)
+                && "t5".equals(modelConfig.getModelType());
+    }
+
     private static boolean isTextGenerationModel(LmiUtils.HuggingFaceModelConfig modelConfig) {
         for (String arch : modelConfig.getArchitectures()) {
             boolean isTextGenerationModel =