Add debug thread name

pytorch · Apr 7, 2022 · 110abd2 · 110abd2
1 parent 67b66a2
commit 110abd2
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 6 deletions.
diff --git a/frontend/server/src/main/java/org/pytorch/serve/ensemble/DagExecutor.java b/frontend/server/src/main/java/org/pytorch/serve/ensemble/DagExecutor.java
@@ -1,5 +1,6 @@
 package org.pytorch.serve.ensemble;
 
+import com.google.common.util.concurrent.ThreadFactoryBuilder;
 import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.List;
@@ -14,10 +15,9 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
+import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
-import java.util.concurrent.ThreadFactory;
-import com.google.common.util.concurrent.ThreadFactoryBuilder;
 import org.pytorch.serve.archive.model.ModelNotFoundException;
 import org.pytorch.serve.archive.model.ModelVersionNotFoundException;
 import org.pytorch.serve.http.InternalServerException;

diff --git a/frontend/server/src/main/java/org/pytorch/serve/workflow/WorkflowManager.java b/frontend/server/src/main/java/org/pytorch/serve/workflow/WorkflowManager.java
@@ -1,5 +1,6 @@
 package org.pytorch.serve.workflow;
 
+import com.google.common.util.concurrent.ThreadFactoryBuilder;
 import com.google.gson.JsonObject;
 import com.google.gson.JsonParser;
 import io.netty.channel.ChannelHandlerContext;
@@ -23,6 +24,7 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
+import java.util.concurrent.ThreadFactory;
 import org.pytorch.serve.archive.DownloadArchiveException;
 import org.pytorch.serve.archive.model.ModelNotFoundException;
 import org.pytorch.serve.archive.model.ModelVersionNotFoundException;
@@ -50,10 +52,6 @@
 
 public final class WorkflowManager {
     private static final Logger logger = LoggerFactory.getLogger(WorkflowManager.class);
-
-    private final ExecutorService inferenceExecutorService =
-            Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors());
-
     private static WorkflowManager workflowManager;
     private final ConfigManager configManager;
     private final ConcurrentHashMap<String, WorkFlow> workflowMap;
@@ -367,6 +365,12 @@ public WorkFlow getWorkflow(String workflowName) {
 
     public void predict(ChannelHandlerContext ctx, String wfName, RequestInput input)
             throws WorkflowException {
+        ThreadFactory namedThreadFactory =
+                new ThreadFactoryBuilder().setNameFormat("wf-manager-thread-%d").build();
+        ExecutorService inferenceExecutorService =
+                Executors.newFixedThreadPool(
+                        Runtime.getRuntime().availableProcessors(), namedThreadFactory);
+
         WorkFlow wf = workflowMap.get(wfName);
         if (wf != null) {
             DagExecutor dagExecutor = new DagExecutor(wf.getDag());
@@ -422,6 +426,7 @@ public void predict(ChannelHandlerContext ctx, String wfName, RequestInput input
                             });
             try {
                 predictionFuture.get();
+                inferenceExecutorService.shutdown();
             } catch (ExecutionException | InterruptedException e) {
                 throw new WorkflowException("Workflow failed ", e);
             }

diff --git a/...er/src/main/java/org/pytorch/serve/workflow/api/http/WorkflowInferenceRequestHandler.java b/...er/src/main/java/org/pytorch/serve/workflow/api/http/WorkflowInferenceRequestHandler.java
@@ -13,6 +13,7 @@
 import org.pytorch.serve.archive.DownloadArchiveException;
 import org.pytorch.serve.archive.model.ModelException;
 import org.pytorch.serve.archive.workflow.WorkflowException;
+import org.pytorch.serve.archive.workflow.WorkflowNotFoundException;
 import org.pytorch.serve.http.BadRequestException;
 import org.pytorch.serve.http.HttpRequestHandlerChain;
 import org.pytorch.serve.http.ResourceNotFoundException;