Merge branch 'main' of github.com:triton-inference-server/vllm_backend into yinggeh-DLIS-7061-add-vllm-metrics

Author: yinggeh

src/model.py

@@ -25,8 +25,10 @@
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 import asyncio
+import gc
 import json
 import os
+import queue
 import threading
 from typing import Dict, List
 
@@ -115,13 +117,19 @@ def initialize(self, args):
         # Counter to keep track of ongoing request counts
         self.ongoing_request_count = 0
 
+        # Starting the response thread. It allows vLLM to keep making progress while
+        # response sender(s) are sending responses to server frontend.
+        self._response_queue = queue.Queue()
+        self._response_thread = threading.Thread(target=self.response_loop)
+        self._response_thread.start()
+
         # Starting asyncio event loop to process the received requests asynchronously.
         self._loop = asyncio.get_event_loop()
-        self._loop_thread = threading.Thread(
+        self._event_thread = threading.Thread(
             target=self.engine_loop, args=(self._loop,)
         )
         self._shutdown_event = asyncio.Event()
-        self._loop_thread.start()
+        self._event_thread.start()
 
     def init_engine(self):
         # Currently, Triton needs to use decoupled policy for asynchronously
@@ -290,6 +298,27 @@ def get_sampling_params_dict(self, params_json):
 
         return params_dict
 
+    def response_loop(self):
+        while True:
+            item = self._response_queue.get()
+            # To signal shutdown a None item will be added to the queue.
+            if item is None:
+                break
+            response_state, response, response_flag = item
+            response_sender = response_state["response_sender"]
+            try:
+                response_sender.send(response, response_flag)
+                # Stop checking for cancellation if the last response is generated.
+                if not response_state["last_response_generated"]:
+                    response_state["is_cancelled"] = response_sender.is_cancelled()
+            except Exception as e:
+                self.logger.log_error(
+                    f"An error occurred while sending a response: {e}"
+                )
+            finally:
+                if response_flag == pb_utils.TRITONSERVER_RESPONSE_COMPLETE_FINAL:
+                    self.ongoing_request_count -= 1
+
     def create_response(self, vllm_output, prepend_input):
         """
         Parses the output from the vLLM engine into Triton
@@ -330,7 +359,13 @@ async def generate(self, request):
         Forwards single request to LLM engine and returns responses.
         """
         response_sender = request.get_response_sender()
+        response_state = {
+            "response_sender": response_sender,
+            "is_cancelled": False,
+            "last_response_generated": False,  # last response ready but not yet sent
+        }
         self.ongoing_request_count += 1
+        decrement_ongoing_request_count = True
         try:
             request_id = random_uuid()
             prompt = pb_utils.get_input_tensor_by_name(
@@ -385,13 +420,31 @@ async def generate(self, request):
                 lora_local_path = self.lora_repository[lora_name]
                 lora_request = LoRARequest(lora_id, lora_int_id, lora_local_path)
 
-            async for output in self.llm_engine.generate(
-                prompt, sampling_params, request_id, lora_request=lora_request
-            ):
-                if response_sender.is_cancelled():
+            response_iterator = await self.llm_engine.add_request(
+                request_id, prompt, sampling_params, lora_request=lora_request
+            )
+
+            async for output in response_iterator:
+                is_cancelled = response_state["is_cancelled"]
+                if not stream:
+                    is_cancelled = response_sender.is_cancelled()
+                if is_cancelled:
                     self.logger.log_info("[vllm] Cancelling the request")
                     await self.llm_engine.abort(request_id)
                     self.logger.log_info("[vllm] Successfully cancelled the request")
+                    if stream:
+                        response_state["last_response_generated"] = True
+                        response = pb_utils.InferenceResponse(
+                            error=pb_utils.TritonError(
+                                message="Request was cancelled",
+                                code=pb_utils.TritonError.CANCELLED,
+                            )
+                        )
+                        flags = pb_utils.TRITONSERVER_RESPONSE_COMPLETE_FINAL
+                        decrement_ongoing_request_count = False
+                        self._response_queue.put_nowait(
+                            (response_state, response, flags)
+                        )
                     break
                 if stream:
                     prev_outputs_lengths = None
@@ -400,15 +453,13 @@ async def generate(self, request):
                             len(prev_output.text)
                             for prev_output in prev_outputs.outputs
                         ]
+                    response = self.create_stream_response(output, prev_outputs_lengths)
+                    flags = 0
                     if output.finished:
-                        response_sender.send(
-                            self.create_stream_response(output, prev_outputs_lengths),
-                            flags=pb_utils.TRITONSERVER_RESPONSE_COMPLETE_FINAL,
-                        )
-                    else:
-                        response_sender.send(
-                            self.create_stream_response(output, prev_outputs_lengths)
-                        )
+                        response_state["last_response_generated"] = True
+                        flags = pb_utils.TRITONSERVER_RESPONSE_COMPLETE_FINAL
+                        decrement_ongoing_request_count = False
+                    self._response_queue.put_nowait((response_state, response, flags))
                 prev_outputs = output
 
             last_output = output
@@ -420,7 +471,7 @@ async def generate(self, request):
                 )
 
         except Exception as e:
-            self.logger.log_info(f"[vllm] Error generating stream: {e}")
+            self.logger.log_error(f"[vllm] Error generating stream: {e}")
             error = pb_utils.TritonError(f"Error generating stream: {e}")
             triton_output_tensor = pb_utils.Tensor(
                 "text_output", np.asarray(["N/A"], dtype=self.output_dtype)
@@ -433,7 +484,8 @@ async def generate(self, request):
             )
             raise e
         finally:
-            self.ongoing_request_count -= 1
+            if decrement_ongoing_request_count:
+                self.ongoing_request_count -= 1
 
     def verify_loras(self, request):
         # We will check if the requested lora exists here, if not we will send a
@@ -500,6 +552,20 @@ def finalize(self):
         """
         self.logger.log_info("[vllm] Issuing finalize to vllm backend")
         self._shutdown_event.set()
-        if self._loop_thread is not None:
-            self._loop_thread.join()
-            self._loop_thread = None
+
+        # Shutdown the event thread.
+        if self._event_thread is not None:
+            self._event_thread.join()
+            self._event_thread = None
+
+        # Shutdown the response thread.
+        self._response_queue.put(None)
+        if self._response_thread is not None:
+            self._response_thread.join()
+            self._response_thread = None
+
+        # When using parallel tensors, the stub process may not shutdown due to
+        # unreleased references, so manually run the garbage collector once.
+        self.logger.log_info("[vllm] Running Garbage Collector on finalize...")
+        gc.collect()
+        self.logger.log_info("[vllm] Garbage Collector on finalize... done")