diff --git a/src/python/library/tritonclient/grpc/_client.py b/src/python/library/tritonclient/grpc/_client.py index 1c68115c4..9fc50493a 100755 --- a/src/python/library/tritonclient/grpc/_client.py +++ b/src/python/library/tritonclient/grpc/_client.py @@ -1741,7 +1741,8 @@ def wrapped_callback(call_future): raise_error_grpc(rpc_error) def start_stream( - self, callback, stream_timeout=None, headers=None, compression_algorithm=None + self, callback, stream_timeout=None, headers=None, compression_algorithm=None, + daemon=None, ): """Starts a grpc bi-directional stream to send streaming inferences. Note: When using stream, user must ensure the InferenceServerClient.close() @@ -1768,6 +1769,8 @@ def start_stream( Optional grpc compression algorithm to be used on client side. Currently supports "deflate", "gzip" and None. By default, no compression is used. + daemon : bool + Make internal thread daemonic if set True. By default set to None. Raises ------ @@ -1784,7 +1787,7 @@ def start_stream( ) metadata = self._get_metadata(headers) - self._stream = _InferStream(callback, self._verbose) + self._stream = _InferStream(callback, self._verbose, daemon=daemon) try: response_iterator = self._client_stub.ModelStreamInfer( diff --git a/src/python/library/tritonclient/grpc/_infer_stream.py b/src/python/library/tritonclient/grpc/_infer_stream.py index 92ec05890..ca3181c71 100755 --- a/src/python/library/tritonclient/grpc/_infer_stream.py +++ b/src/python/library/tritonclient/grpc/_infer_stream.py @@ -52,9 +52,11 @@ class _InferStream: 'error' would be None for a successful inference. verbose : bool Enables verbose mode if set True. + daemon : bool + Make handler thread daemonic if set True (default None). """ - def __init__(self, callback, verbose): + def __init__(self, callback, verbose, daemon=None): self._callback = callback self._verbose = verbose self._request_queue = queue.Queue() @@ -62,6 +64,7 @@ def __init__(self, callback, verbose): self._cancelled = False self._active = True self._response_iterator = None + self.daemon = daemon def __del__(self): self.close(cancel_requests=True) @@ -100,7 +103,7 @@ def _init_handler(self, response_iterator): if self._handler is not None: raise_error("Attempted to initialize already initialized InferStream") # Create a new thread to handle the gRPC response stream - self._handler = threading.Thread(target=self._process_response) + self._handler = threading.Thread(target=self._process_response, daemon=self.daemon) self._handler.start() if self._verbose: print("stream started...")