From 499eb4f2f43943c4a1a55a20dbf8b9f3af91ea7a Mon Sep 17 00:00:00 2001
From: Frank Liu <frankfliu2000@gmail.com>
Date: Tue, 30 May 2023 08:46:01 -0700
Subject: [PATCH] [wlm] Sets default maxWorkers based on OMP_NUM_THREADS (#776)

---
 .../ai/djl/serving/wlm/util/WlmConfigManager.java    | 12 +++++++++++-
 1 file changed, 11 insertions(+), 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
index beade27a8..0cf114b6b 100644
--- 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) {