Make stream input optional (#58)

triton-inference-server · Oct 9, 2023 · 2e4a4e7 · 2e4a4e7
1 parent 2e29214
commit 2e4a4e7
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/Quick_Deploy/vLLM/model_repository/vllm/1/model.py b/Quick_Deploy/vLLM/model_repository/vllm/1/model.py
@@ -166,15 +166,20 @@ async def generate(self, request):
         self.ongoing_request_count += 1
         try:
             request_id = random_uuid()
+
             prompt = pb_utils.get_input_tensor_by_name(request, "PROMPT").as_numpy()[0]
             if isinstance(prompt, bytes):
                 prompt = prompt.decode("utf-8")
-            stream = pb_utils.get_input_tensor_by_name(request, "STREAM").as_numpy()[0]
+
+            # stream is an optional input
+            stream = False
+            stream_input_tensor = pb_utils.get_input_tensor_by_name(request, "STREAM")
+            if stream_input_tensor:
+                stream = stream_input_tensor.as_numpy()[0]
 
             # Request parameters are not yet supported via
             # BLS. Provide an optional mechanism to receive serialized
             # parameters as an input tensor until support is added
-
             parameters_input_tensor = pb_utils.get_input_tensor_by_name(request, "SAMPLING_PARAMETERS")
             if parameters_input_tensor:
                 parameters = parameters_input_tensor.as_numpy()[0].decode("utf-8")

diff --git a/Quick_Deploy/vLLM/model_repository/vllm/config.pbtxt b/Quick_Deploy/vLLM/model_repository/vllm/config.pbtxt
@@ -49,6 +49,7 @@ input [
     name: "STREAM"
     data_type: TYPE_BOOL
     dims: [ 1 ]
+    optional: true
   },
   {
     name: "SAMPLING_PARAMETERS"