rebase to v1.0.0rc6 and works

Author: davidmlw

tensorrt_llm/_torch/pyexecutor/executor_request_queue.py

@@ -18,6 +18,9 @@
 from .sampler import Sampler, TorchSampler
 
 SHUTDOWN_REQUEST_ID = -1
+UPDATE_WEIGHT_REQUEST_ID = -2
+SLEEP_REQUEST_ID = -3
+WAKEUP_REQUEST_ID = -4
 
 
 @dataclasses.dataclass
@@ -28,15 +31,33 @@ class RequestQueueItem:
     child_req_ids: Optional[list] = None
     is_canceled_request: bool = False
     query: Optional[list] = None  # only used in `StarAttention`
+    weight_ipc_handles: Optional[dict] = None
+    sleep_level: Optional[int] = None
+    wakeup_level: Optional[int] = None
 
     @property
     def is_shutdown_request(self):
         return self.id == SHUTDOWN_REQUEST_ID
 
     @property
     def is_normal_request(self):
-        return not (self.is_shutdown_request or self.is_canceled_request)
+        return self.id > 0 and not self.is_canceled_request
 
+    @property
+    def is_update_weight_request(self):
+        return self.id == UPDATE_WEIGHT_REQUEST_ID
+
+    @property
+    def is_sleep_request(self):
+        return self.id == SLEEP_REQUEST_ID
+
+    @property
+    def is_wakeup_request(self):
+        return self.id == WAKEUP_REQUEST_ID
+
+    @property
+    def is_control_request(self):
+        return self.is_update_weight_request or self.is_sleep_request or self.is_wakeup_request
 
 class ExecutorRequestQueue:
     """Handles fetching and processing of new requests from the request queue."""
@@ -66,6 +87,7 @@ def __init__(self, dist: Distributed, enable_attention_dp: bool,
         self.new_active_requests_queue_latency_ms = 0
         self.is_shutdown = False
         self.should_exclude_last_generation_logits = False
+        self.control_requests: List[RequestQueueItem] = []
 
     def _get_from_request_queue(
             self,
@@ -226,6 +248,20 @@ def enqueue_cancel_request(self, req_id: int):
             self.request_queue.put(
                 RequestQueueItem(req_id, is_canceled_request=True))
 
+    def enqueue_sleep_request(self, req_id: int, sleep_level: int):
+        with self.enqueue_lock:
+            print(f"enqueue_sleep_request: {req_id} {sleep_level}")
+            self.request_queue.put(
+                RequestQueueItem(req_id, sleep_level=sleep_level))
+
+    def enqueue_wakeup_request(self, req_id: int, wakeup_level: int):
+        with self.enqueue_lock:
+            self.request_queue.put(RequestQueueItem(req_id, wakeup_level=wakeup_level))
+
+    def enqueue_update_weight_request(self, req_id: int, weight_ipc_handles: dict):
+        with self.enqueue_lock:
+            self.request_queue.put(RequestQueueItem(req_id, weight_ipc_handles=weight_ipc_handles))
+
     def enqueue_shutdown_request(self):
         with self.enqueue_lock:
             self.request_queue.put(RequestQueueItem(SHUTDOWN_REQUEST_ID))
@@ -431,6 +467,8 @@ def _validate_and_filter_requests(
                 break
             elif req_item.is_canceled_request:
                 self.canceled_req_ids.append(req_item.id)
+            elif req_item.is_control_request:
+                self.control_requests.append(req_item)
             else:
                 valid_new_requests.append(req_item)
 

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -334,6 +334,32 @@ def cancel_request(self, id: int):
         """
         self.executor_request_queue.enqueue_cancel_request(id)
 
+    def enqueue_sleep_request(self, id: int, sleep_level: int):
+        """
+        Enqueue a sleep request with provided request id and sleep level
+        Args:
+            id (int): The request id for which to sleep
+            sleep_level (int): The sleep level to apply to the request
+        """
+        self.executor_request_queue.enqueue_sleep_request(id, sleep_level)
+
+    def enqueue_wakeup_request(self, id: int, wakeup_level: int):
+        """
+        Enqueue a wakeup request with provided request id
+        Args:
+            id (int): The request id for which to wakeup
+        """
+        self.executor_request_queue.enqueue_wakeup_request(id, wakeup_level)
+
+    def enqueue_update_weight_request(self, id: int, weight_ipc_handles: dict):
+        """
+        Enqueue a update weight request with provided request id and weight ipc handles
+        Args:
+            id (int): The request id for which to update weight
+            weight_ipc_handles (dict): The weight ipc handles to update
+        """
+        self.executor_request_queue.enqueue_update_weight_request(id, weight_ipc_handles)
+
     def shutdown(self):
         """
         Signals the server to shutdown.
@@ -1080,19 +1106,19 @@ def update_weight_from_ipc_handles(self, handles):
 
     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=PyResult(0, 0, success=True), 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=PyResult(0, 0, success=True), 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
 
         try:
             self.update_weight_from_ipc_handles(update_weight_request.weight_ipc_handles)
             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})
+            self._enqueue_responses([(update_weight_request.id, update_weight_response)])
         except Exception as e:
             print(
                 f"Error in update_weights_from_ipc_handles: {e}"
@@ -1101,6 +1127,20 @@ def _update_weight(self, update_weight_request):
             #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 _handle_control_request(self):
+        if len(self.executor_request_queue.control_requests) > 0:
+            assert len(self.executor_request_queue.control_requests) == 1, f"control request should be the only request in the list, but got {len(self.executor_request_queue.control_requests)}"
+            control_request = self.executor_request_queue.control_requests.pop()
+            if (control_request.is_update_weight_request):
+                self._update_weight(control_request)
+            elif (control_request.is_sleep_request):
+                self._sleep(control_request)
+            elif (control_request.is_wakeup_request):
+                self._wakeup(control_request)
+            else:
+                assert False, "Invalid control request"
+
+
     def _executor_loop_overlap(self):
         torch.cuda.set_device(self.device_id)
         # ensure the context is created, otherwise, some MPI calls will fail.
@@ -1122,20 +1162,10 @@ def _executor_loop_overlap(self):
                     iter_start_time = time.time()
 
                 scheduled_batch, iter_stats = self._prepare_and_schedule_batch()
+                self._handle_control_request()
+
                 if scheduled_batch is None:
                     break
-                if self.is_control_request:
-                    self.is_control_request = False
-                    assert len(new_requests) == 1, f"control request should be the only request in the list, but got {len(new_requests)}"
-                    if (new_requests[0].is_update_weight_request()):
-                        self._update_weight(new_requests[0])
-                    elif (new_requests[0].is_sleep_request()):
-                        self._sleep(new_requests[0])
-                    elif (new_requests[0].is_wakeup_request()):
-                        self._wakeup(new_requests[0])
-                    else:
-                        assert False, "Invalid control request"
-                    continue
 
                 self._pause_requests(scheduled_batch.paused_requests)
 

tensorrt_llm/_torch/utils.py

@@ -284,6 +284,7 @@ def nvml_context() -> Generator[None, None, None]:
 
 def device_id_to_physical_device_id(device_id: int) -> int:
     """Convert a logical device ID to a physical device ID considering CUDA_VISIBLE_DEVICES."""
+    import os
     if "CUDA_VISIBLE_DEVICES" in os.environ:
         device_ids = os.environ["CUDA_VISIBLE_DEVICES"].split(",")
         try:

tensorrt_llm/executor/worker.py

@@ -460,13 +460,13 @@ def _deduce_max_tokens(request: GenerationRequest,
                     f"is larger than max_seq_len {executor_config.max_seq_len}")
             return default_max_tokens
         if request.is_weight_update_request():
-            req_id = self.engine.enqueue_request(request, weight_ipc_handles=request.weight_ipc_handles)
+            req_id = self.engine.enqueue_update_weight_request(request.id, weight_ipc_handles=request.weight_ipc_handles)
             return req_id
         elif request.is_sleep_request():
-            req_id = self.engine.enqueue_request(request, sleep_level=request.sleep_level)
+            req_id = self.engine.enqueue_sleep_request(request.id, sleep_level=request.sleep_level)
             return req_id
         elif request.is_wakeup_request():
-            req_id = self.engine.enqueue_request(request, wakeup_level=request.wakeup_level)
+            req_id = self.engine.enqueue_wakeup_request(request.id, wakeup_level=request.wakeup_level)
             return req_id
         try:
             executor_request = tllm.Request(