Test response_sender.send() after closed

triton-inference-server · Oct 4, 2023 · 430b16d · 430b16d
1 parent ce3142d
commit 430b16d
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/qa/L0_backend_python/decoupled/decoupled_test.py b/qa/L0_backend_python/decoupled/decoupled_test.py
@@ -288,9 +288,6 @@ def test_decoupled_execute_cancel(self):
             log_text = f.read()
         self.assertIn("[execute_cancel] Request not cancelled at 1.0 s", log_text)
         self.assertIn("[execute_cancel] Request cancelled at ", log_text)
-        self.assertIn(
-            "add response output: output: DUMMY_OUT, type: FP32, shape: [1,1]", log_text
-        )
 
 
 if __name__ == "__main__":

diff --git a/qa/python_models/execute_cancel/model.py b/qa/python_models/execute_cancel/model.py
@@ -102,11 +102,26 @@ def response_thread(execute_processed_requests, processed_requests):
                         flags=pb_utils.TRITONSERVER_RESPONSE_COMPLETE_FINAL
                     )
                 else:
+                    # test response_sender.send() after it is closed upon checking for
+                    # confirmed request cancelled
                     dummy_tensor = pb_utils.Tensor(
                         "DUMMY_OUT", np.array([[1]], np.single)
                     )
                     dummy_response = pb_utils.InferenceResponse([dummy_tensor])
-                    response_sender.send(dummy_response)
+                    raised_exception = False
+                    try:
+                        response_sender.send(dummy_response)
+                    except pb_utils.TritonModelException as e:
+                        if (
+                            "Unable to send response. Response sender has been closed."
+                            not in str(e)
+                        ):
+                            raise e
+                        raised_exception = True
+                    if not raised_exception:
+                        raise pb_utils.TritonModelException(
+                            "response_sender did not raise an exception on send after closed"
+                        )
 
         thread = threading.Thread(
             target=response_thread,