add success into GenerationResult

Author: davidmlw

tensorrt_llm/_torch/pyexecutor/llm_request.py

@@ -163,14 +163,16 @@ def __init__(self,
                  return_log_probs: bool = False,
                  return_context_logits: bool = False,
                  return_generation_logits: bool = False,
-                 exclude_last_generation_logits: bool = False):
+                 exclude_last_generation_logits: bool = False,
+                 success: bool = False):
         self._streaming = streaming
         self._context_logits = LogitsStorage(
             prompt_len, use_device_memory) if return_context_logits else None
         self._generation_logits = LogitsStorage(
             max_new_tokens, use_device_memory, exclude_last_generation_logits
         ) if return_generation_logits else None
         self._log_probs = LogProbStorage() if return_log_probs else None
+        self._success = success
 
     def append_context_logits(self, context_logits: torch.Tensor):
         if self._context_logits:

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -34,7 +34,7 @@
 from ..speculative.drafter import Drafter
 from .kv_cache_transceiver import KvCacheTransceiver
 from .llm_request import (ExecutorRequest, LlmRequest, LlmRequestState,
-                          LlmResponse, LlmResult, executor_request_to_llm_request)
+                          LlmResponse, LlmResult, executor_request_to_llm_request, PyResult)
 from .model_engine import ModelEngine
 from .sampler import Sampler, SampleState, SampleStateTensors, TorchSampler
 from .scheduler import RequestScheduler, ScheduledRequests
@@ -1086,11 +1086,11 @@ def _prepare_draft_requests(self):
 
     def _sleep(self, sleep_request):
         self.is_sleep_request = False
-        self._enqueue_responses({sleep_request.id: LlmResponse(request_id=sleep_request.id, result=LlmResult(result=None, py_result=None, is_final=True), client_id=sleep_request.id)})
+        self._enqueue_responses({sleep_request.id: LlmResponse(request_id=sleep_request.id, result=LlmResult(result=None, py_result=PyResult(0, 0, success=True), is_final=True), client_id=sleep_request.id)})
 
     def _wakeup(self, wakeup_request):
         self.is_wakeup_request = False
-        self._enqueue_responses({wakeup_request.id: LlmResponse(request_id=wakeup_request.id, result=LlmResult(result=None, py_result=None, is_final=True), client_id=wakeup_request.id)})
+        self._enqueue_responses({wakeup_request.id: LlmResponse(request_id=wakeup_request.id, result=LlmResult(result=None, py_result=PyResult(0, 0, success=True), is_final=True), client_id=wakeup_request.id)})
 
     def _update_weight(self, update_weight_request):
         self.is_update_weight_request = False
@@ -1119,13 +1119,13 @@ def _update_weight(self, update_weight_request):
             self.model_engine.model.load_weights(weights)
 
             torch.cuda.synchronize()
-            update_weight_response = LlmResponse(request_id=update_weight_request.id, result=LlmResult(result=None, py_result=None, is_final=True),     client_id=update_weight_request.id)
+            update_weight_response = LlmResponse(request_id=update_weight_request.id, result=LlmResult(result=None, py_result=PyResult(0, 0, success=True), is_final=True),     client_id=update_weight_request.id)
             self._enqueue_responses({update_weight_request.id: update_weight_response})
         except Exception as e:
             print(
                 f"Error in VllmInternalWorkerExtension.update_weights_from_ipc_handles: {e}"
             )
-            update_weight_response = LlmResponse(request_id=update_weight_request.id, result=LlmResult(result=None, py_result=None, is_final=True), client_id=update_weight_request.id)
+            update_weight_response = LlmResponse(request_id=update_weight_request.id, result=LlmResult(result=None, py_result=PyResult(0, 0, success=False), is_final=True), client_id=update_weight_request.id)
             self._enqueue_responses({update_weight_request.id: update_weight_response})
 
     def _executor_loop_overlap(self):

tensorrt_llm/executor/result.py

@@ -146,6 +146,7 @@ def __init__(self,
         self.disaggregated_params = None
         self.decoding_iter = 0
         self._done = False
+        self._success = False
 
         if has_event_loop():
             self.aqueue = AsyncQueue()
@@ -303,6 +304,7 @@ def _handle_response(self,
                 response_result.deserialize()
 
             self._done = response_result.is_final
+            self._success = True # TODO: replace with response_result._py_result._success
             context_phase_params = response_result.context_phase_params
             self.decoding_iter = response_result.decoding_iter
             if context_phase_params is not None:
@@ -332,10 +334,13 @@ def _handle_response(self,
                     handler := self._background_error_handler()):
                 handler()
         elif is_update_weights_response(response):
+            self._success = response.result._py_result._success
             self._done = True
         elif is_sleep_response(response):
+            self._success = response.result._py_result._success
             self._done = True
         elif is_wakeup_response(response):
+            self._success = response.result._py_result._success
             self._done = True
         elif isinstance(response, ErrorResponse):
             if self._background_error_handler is not None and (
@@ -463,6 +468,10 @@ def aborted(self) -> bool:
     def finished(self) -> bool:
         return self._done
 
+    @property
+    def success(self) -> bool:
+        return self._success
+
     def clear_logprob_params(self) -> None:
         # Remove temporary attribute used in executor
         # for a cleaner external-facing output.
@@ -533,7 +542,7 @@ def _exception(self, timeout: Optional[float] = None):
 
     def _repr_fields(self):
         return [
-            'request_id', 'prompt_token_ids', 'outputs', 'finished',
+            'request_id', 'prompt_token_ids', 'outputs', 'finished', 'success',
             "context_logits"
         ]