[python] Sets rolloing batch threads as daemon thread

deepjavalibrary · Dec 8, 2023 · 9e98b2b · 9e98b2b
1 parent 9cd5a7e
commit 9e98b2b
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/engines/python/src/main/java/ai/djl/python/engine/RollingBatch.java b/engines/python/src/main/java/ai/djl/python/engine/RollingBatch.java
@@ -87,6 +87,7 @@ class RollingBatch implements Runnable {
     @Override
     public void run() {
         currentThread = Thread.currentThread();
+        currentThread.setDaemon(true);
         while (!stop) {
             int size;
             Input batch = new Input();

diff --git a/wlm/src/main/java/ai/djl/serving/wlm/WorkLoadManager.java b/wlm/src/main/java/ai/djl/serving/wlm/WorkLoadManager.java
@@ -43,7 +43,13 @@ public class WorkLoadManager {
 
     /** Constructs a {@link WorkLoadManager} instance. */
     public WorkLoadManager() {
-        threadPool = Executors.newCachedThreadPool();
+        threadPool =
+                Executors.newCachedThreadPool(
+                        r -> {
+                            Thread t = Executors.defaultThreadFactory().newThread(r);
+                            t.setDaemon(true);
+                            return t;
+                        });
         workerPools = new ConcurrentHashMap<>();
     }