[wlm] Sets default maxWorkers based on OMP_NUM_THREADS (deepjavalibra…

…ry#776)
KexinFeng · Aug 16, 2023 · 499eb4f · 499eb4f
1 parent 70df327
commit 499eb4f
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/wlm/src/main/java/ai/djl/serving/wlm/util/WlmConfigManager.java b/wlm/src/main/java/ai/djl/serving/wlm/util/WlmConfigManager.java
@@ -15,6 +15,7 @@
 import ai.djl.Device;
 import ai.djl.Model;
 import ai.djl.ndarray.NDManager;
+import ai.djl.util.Utils;
 
 /** This manages some configurations used by the {@link ai.djl.serving.wlm.WorkLoadManager}. */
 public final class WlmConfigManager {
@@ -199,7 +200,16 @@ public int getDefaultMaxWorkers(Model model) {
             }
             return 2;
         }
-        return Runtime.getRuntime().availableProcessors();
+
+        int cpuCores = Runtime.getRuntime().availableProcessors();
+        int ompThreads = Integer.parseInt(Utils.getenv("OMP_NUM_THREADS", "-1"));
+        if (ompThreads > 0) {
+            if (ompThreads > cpuCores) {
+                ompThreads = cpuCores;
+            }
+            return cpuCores * 2 / ompThreads;
+        }
+        return 2;
     }
 
     private static int getWorkersProperty(Model model, String key, int def) {