Catch exception in python

Signed-off-by: Shunkang <[email protected]>

Author: Shunkang

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -22,6 +22,7 @@
 from tensorrt_llm._torch.pyexecutor.seq_slot_manager import SeqSlotManager
 from tensorrt_llm._utils import (customized_gc_thresholds, global_mpi_rank,
                                  is_trace_enabled, nvtx_range, trace_func)
+from tensorrt_llm.bindings.exceptions import RequestSpecificException
 from tensorrt_llm.bindings.executor import (DisServingRequestStats,
                                             FinishReason, InflightBatchingStats,
                                             IterationStats, KvCacheStats,
@@ -686,8 +687,7 @@ def _executor_loop_pp(self):
                         logger.warning(
                             "num_fitting_reqs=0 and fitting_disagg_gen_init_requests is empty, may not have enough kvCache"
                         )
-                        self.kv_cache_transceiver.check_context_transfer_status(
-                            1)
+                        self._check_disagg_ctx_cache_transfer_status(1)
 
                 self.num_scheduled_requests = scheduled_batch.batch_size
 
@@ -887,7 +887,11 @@ def _prepare_and_schedule_batch(self):
                 logger.warning(
                     "num_fitting_reqs=0 and fitting_disagg_gen_init_requests is empty, may not have enough kvCache"
                 )
-                self.kv_cache_transceiver.check_context_transfer_status(1)
+                self._check_disagg_ctx_cache_transfer_status(1)
+        else:
+            assert scheduled_batch.batch_size > 0, (
+                "fail to schedule any pending request, "
+                "probably run out of resource.")
 
         self.num_scheduled_requests = scheduled_batch.batch_size
         logger.debug(
@@ -1258,7 +1262,7 @@ def _check_disagg_gen_transfer_status(self):
 
         if need_check:
             at_least_num = 1 if need_check_one else 0
-            self.kv_cache_transceiver.check_gen_transfer_status(at_least_num)
+            self._check_disagg_gen_cache_transfer_status(at_least_num)
 
         return
 
@@ -1361,8 +1365,7 @@ def _recv_disagg_gen_cache(self, new_gen_reqs):
             req.is_disagg_generation_transmission_in_progress
             for req in self.active_requests
         ])
-        self.kv_cache_transceiver.check_gen_transfer_status(
-            1 if block_transfer else 0)
+        self._check_disagg_gen_cache_transfer_status(1 if block_transfer else 0)
 
         return
 
@@ -1382,7 +1385,7 @@ def _send_disagg_ctx_cache(self, scheduled_ctx_requests):
                         self.resource_manager.resource_managers[
                             resource_mgr_type].free_resources(req)
 
-        self.kv_cache_transceiver.check_context_transfer_status(0)
+        self._check_disagg_ctx_cache_transfer_status(0)
 
         # Keep track of ctx requests that are in transmission
         ctx_transmission_reqs = [
@@ -1392,6 +1395,38 @@ def _send_disagg_ctx_cache(self, scheduled_ctx_requests):
 
         return ctx_transmission_reqs
 
+    def _check_cache_transfer_status_helper(self,
+                                            method_name: str,
+                                            method_call,
+                                            atLeastNum: int = 0):
+        """Helper method to handle cache transfer status checking with error handling."""
+        try:
+            method_call(atLeastNum)
+        except RequestSpecificException as e:
+            error_msg = str(e)
+            logger.error(
+                f"Encountered a request-specific error in {method_name}: {error_msg}"
+            )
+            request_ids = [e.request_id]
+            self._handle_errors(error_msg, request_ids)
+        except Exception as e:
+            error_msg = str(e)
+            logger.error(
+                f"Encountered a system error in {method_name}: {error_msg}")
+            self._handle_errors(error_msg)
+
+    @nvtx_range("_check_disagg_ctx_cache_transfer_status")
+    def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0):
+        self._check_cache_transfer_status_helper(
+            "checking context transfer status",
+            self.kv_cache_transceiver.check_context_transfer_status, atLeastNum)
+
+    @nvtx_range("_check_disagg_gen_cache_transfer_status")
+    def _check_disagg_gen_cache_transfer_status(self, atLeastNum: int = 0):
+        self._check_cache_transfer_status_helper(
+            "checking generation transfer status",
+            self.kv_cache_transceiver.check_gen_transfer_status, atLeastNum)
+
     def _forward_step(self,
                       scheduled_requests,
                       new_tensors_device: Optional[SampleStateTensors] = None):
@@ -1501,27 +1536,26 @@ def _update_requests(self, sample_state: SampleState):
 
     def _handle_errors(self,
                        error_msg: Optional[str] = None,
-                       *,
-                       requests: Optional[List[LlmRequest]] = None):
-        error_responses: Dict[int, LlmResponse] = {}
+                       request_ids: Optional[List[int]] = None):
+        error_responses = {}
         error_msg = error_msg or "error"
-        failed_requests = requests if requests is not None else self.active_requests
-        for request in failed_requests:
+        for request in self.active_requests:
+            if request_ids is not None and request.py_request_id not in request_ids:
+                continue
             req_id = request.py_request_id
             request.state = LlmRequestState.GENERATION_COMPLETE
             self._terminate_request(request)
             error_responses[req_id] = LlmResponse(
                 request_id=req_id,
                 error_msg=error_msg,
                 client_id=request.py_client_id)
-        if requests is None:
-            self.active_requests.clear()
+
+        if request_ids is not None:
+            for req_id in request_ids:
+                self.active_requests.remove(req_id)
         else:
-            self.active_requests = [
-                request for request in self.active_requests
-                if request not in requests
-            ]
-        self._enqueue_responses(error_responses.items())
+            self.active_requests.clear()
+        self._enqueue_responses(error_responses)
 
     def _terminate_request(self, request: LlmRequest):
         self.resource_manager.free_resources(request)