Fix workflow thread issue pytorch#1511

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

Hide file tree

Showing 3 changed files with 8 additions and 3 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
@@ -150,7 +150,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];