Add debug thread name

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

Hide file tree

Showing 3 changed files with 13 additions and 7 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;
@@ -57,7 +58,7 @@ public void handleRequest(
 
     private void handlePredictions(
             ChannelHandlerContext ctx, FullHttpRequest req, String[] segments)
-            throws WorkflowException {
+            throws WorkflowNotFoundException {
         RequestInput input = parseRequest(ctx, req);
         logger.info(input.toString());
         String wfName = segments[2];