Fix workflow thread issue pytorch#1511

maaquib · Apr 6, 2022 · 3327be2 · 3327be2
1 parent 375f004
commit 3327be2
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 4 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
@@ -41,8 +41,9 @@ public DagExecutor(Dag dag) {
     public ArrayList<NodeOutput> execute(RequestInput input, ArrayList<String> topoSortedList) {
 
         CompletionService<NodeOutput> executorCompletionService = null;
+        ExecutorService executorService = null;
         if (topoSortedList == null) {
-            ExecutorService executorService = Executors.newFixedThreadPool(4);
+            executorService = Executors.newFixedThreadPool(4);
             executorCompletionService = new ExecutorCompletionService<>(executorService);
         }
 
@@ -90,6 +91,8 @@ public ArrayList<NodeOutput> execute(RequestInput input, ArrayList<String> topoS
                     logger.error(e.getMessage());
                     String[] error = e.getMessage().split(":");
                     throw new InternalServerException(error[error.length - 1]); // NOPMD
+                } finally {
+                    executorService.shutdownNow();
                 }
             } else {
                 for (String name : readyToExecute) {
@@ -150,7 +153,7 @@ private NodeOutput invokeModel(
                     InterruptedException {
         try {
 
-            logger.error(String.format("Invoking -  %s for attempt %d", nodeName, retryAttempt));
+            logger.info(String.format("Invoking -  %s for attempt %d", nodeName, retryAttempt));
             CompletableFuture<byte[]> respFuture = new CompletableFuture<>();
             RestJob job = ApiUtils.addRESTInferenceJob(null, workflowModel.getName(), null, input);
             job.setResponsePromise(respFuture);

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
@@ -366,7 +366,7 @@ public WorkFlow getWorkflow(String workflowName) {
     }
 
     public void predict(ChannelHandlerContext ctx, String wfName, RequestInput input)
-            throws WorkflowNotFoundException {
+            throws WorkflowNotFoundException, WorkflowException {
         WorkFlow wf = workflowMap.get(wfName);
         if (wf != null) {
             DagExecutor dagExecutor = new DagExecutor(wf.getDag());
@@ -420,6 +420,11 @@ public void predict(ChannelHandlerContext ctx, String wfName, RequestInput input
                                                 error[error.length - 1].strip()));
                                 return null;
                             });
+            try {
+                predictionFuture.get();
+            } catch (ExecutionException | InterruptedException e) {
+                throw new WorkflowException("Workflow failed ", e);
+            }
         } else {
             throw new WorkflowNotFoundException("Workflow not found: " + wfName);
         }

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