avoid kv cache manager race conditions

Signed-off-by: Qiang Xu <qiangx@nvidia.com>

Author: qiangxu1996

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -214,9 +214,9 @@ def __init__(self,
         self.responses = {}
         self.result_wait_queues = {}
 
-        self.sm_disagg_request_lock = threading.Lock()
-        self.ctx_request_cv = threading.Condition(self.sm_disagg_request_lock)
-        self.gen_request_cv = threading.Condition(self.sm_disagg_request_lock)
+        self.sm_disagg_lock = threading.Lock()
+        self.ctx_request_cv = threading.Condition(self.sm_disagg_lock)
+        self.gen_request_cv = threading.Condition(self.sm_disagg_lock)
 
         # kv cache events
         self.kv_cache_manager = self.resource_manager.resource_managers.get(
@@ -1536,7 +1536,7 @@ def _executor_loop_sm_disagg_ctx(self, stream):
                         self.executor_request_queue.
                         get_new_active_requests_queue_latency())
 
-                with self.sm_disagg_request_lock:
+                with self.sm_disagg_lock:
                     ctx_requests = get_context_requests(self.active_requests)
                     if self.is_shutdown and len(ctx_requests) == 0 \
                             and self.executor_request_queue.get_waiting_queue_size() == 0:
@@ -1559,7 +1559,8 @@ def _executor_loop_sm_disagg_ctx(self, stream):
 
                 if scheduled_batch.batch_size > 0 or (
                         self.enable_attention_dp and self.dist.tp_size > 1):
-                    self.resource_manager.prepare_resources(scheduled_batch)
+                    with self.sm_disagg_lock:
+                        self.resource_manager.prepare_resources(scheduled_batch)
 
                     with torch.cuda.stream(stream):
                         batch_outputs = self._forward_step(
@@ -1569,23 +1570,25 @@ def _executor_loop_sm_disagg_ctx(self, stream):
                         # To avoid long sync time in critical section below
                         sample_state.sampler_event.synchronize()
 
-                    with self.sm_disagg_request_lock:
+                    with self.sm_disagg_lock:
                         self._update_request_states(scheduled_batch)
                         self._update_requests(sample_state,
                                               self.resource_manager)
                         self._handle_canceled_requests()
                         finished_requests = self._handle_responses()
-                        self.ctx_request_cv.notify()
 
-                    attn_metadata = getattr(self.ctx_model_engine,
-                                            'attn_metadata', None)
-                    kv_cache_dtype_byte_size = getattr(
-                        self.ctx_model_engine, 'kv_cache_dtype_byte_size', None)
-                    self.resource_manager.update_resources(
-                        scheduled_batch, attn_metadata,
-                        kv_cache_dtype_byte_size)
-                    if self.enable_kv_cache_events:
-                        self._add_kv_cache_events()
+                        attn_metadata = getattr(self.ctx_model_engine,
+                                                'attn_metadata', None)
+                        kv_cache_dtype_byte_size = getattr(
+                            self.ctx_model_engine, 'kv_cache_dtype_byte_size',
+                            None)
+                        self.resource_manager.update_resources(
+                            scheduled_batch, attn_metadata,
+                            kv_cache_dtype_byte_size)
+                        if self.enable_kv_cache_events:
+                            self._add_kv_cache_events()
+
+                        self.ctx_request_cv.notify()
 
                 if self.enable_iter_perf_stats and sample_state is not None:
                     iter_stats.iter_counter = self.ctx_model_engine.iter_counter
@@ -1621,7 +1624,7 @@ def _executor_loop_sm_disagg_gen_overlap(self, stream):
                         num_new_active_requests=0,
                         new_active_requests_queue_latency_ms=0)
 
-                with self.sm_disagg_request_lock:
+                with self.sm_disagg_lock:
                     self._pad_attention_dp_dummy_request()
 
                     gen_requests = get_generation_requests(self.active_requests)
@@ -1641,7 +1644,8 @@ def _executor_loop_sm_disagg_gen_overlap(self, stream):
                     self._pause_requests(scheduled_batch.paused_requests)
 
                 if scheduled_batch.batch_size > 0:
-                    self.resource_manager.prepare_resources(scheduled_batch)
+                    with self.sm_disagg_lock:
+                        self.resource_manager.prepare_resources(scheduled_batch)
 
                     # The generation requests that just finished context phase
                     # needs to be in front of the batch due to the assumptions
@@ -1662,7 +1666,7 @@ def _executor_loop_sm_disagg_gen_overlap(self, stream):
                             self.previous_batch.sample_state.sampler_event.synchronize(
                             )
 
-                    with self.sm_disagg_request_lock:
+                    with self.sm_disagg_lock:
                         if self.previous_batch is not None:
                             self._update_requests(
                                 self.previous_batch.sample_state)
@@ -1672,7 +1676,7 @@ def _executor_loop_sm_disagg_gen_overlap(self, stream):
                             scheduled_batch, batch_outputs)
                         assert sample_state is not None, "Sampling failed"
 
-                    with self.sm_disagg_request_lock:
+                    with self.sm_disagg_lock:
                         self._update_request_states(scheduled_batch)
 
                         if self.previous_batch is not None: