Merge 7e03592 into 603e89f

pytorch · Jun 24, 2023 · 8393f47 · 8393f47
2 parents 603e89f + 7e03592
commit 8393f47
Showing 1 changed file with 16 additions and 1 deletion.
diff --git a/frontend/server/src/main/java/org/pytorch/serve/job/GRPCJob.java b/frontend/server/src/main/java/org/pytorch/serve/job/GRPCJob.java
@@ -4,6 +4,7 @@
 
 import com.google.protobuf.ByteString;
 import io.grpc.Status;
+import io.grpc.stub.ServerCallStreamObserver;
 import io.grpc.stub.StreamObserver;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -67,10 +68,16 @@ public void response(
             int statusCode,
             String statusPhrase,
             Map<String, String> responseHeaders) {
-
         ByteString output = ByteString.copyFrom(body);
         if (this.getCmd() == WorkerCommands.PREDICT
                 || this.getCmd() == WorkerCommands.STREAMPREDICT) {
+            if (((ServerCallStreamObserver<PredictionResponse>) predictionResponseObserver)
+                    .isCancelled()) {
+                logger.warn(
+                        "grpc client call already cancelled, not able to send this response for requestId: {}",
+                        getPayload().getRequestId());
+                return;
+            }
             PredictionResponse reply =
                     PredictionResponse.newBuilder().setPrediction(output).build();
             predictionResponseObserver.onNext(reply);
@@ -118,6 +125,14 @@ public void sendError(int status, String error) {
         Status responseStatus = GRPCUtils.getGRPCStatusCode(status);
         if (this.getCmd() == WorkerCommands.PREDICT
                 || this.getCmd() == WorkerCommands.STREAMPREDICT) {
+            if (((ServerCallStreamObserver<PredictionResponse>) predictionResponseObserver)
+                    .isCancelled()) {
+                logger.warn(
+                        "grpc client call already cancelled, not able to send error: {}, for requestId: {}",
+                        error,
+                        getPayload().getRequestId());
+                return;
+            }
             predictionResponseObserver.onError(
                     responseStatus
                             .withDescription(error)