[Bugfix] Race condition in MultiprocExecutor when concurent access to Scheduler (#1448)

Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

vllm_omni/diffusion/executor/multiproc_executor.py

@@ -160,33 +160,43 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         }
 
         try:
-            # Broadcast RPC request to all workers via unified message queue
-            self.scheduler.mq.enqueue(rpc_request)
+            # Acquire lock with timeout awareness so that a stalled add_req
+            # (holding the lock while blocked on dequeue) does not prevent
+            # this RPC from honouring its own timeout.
+            lock_timeout = None if deadline is None else max(0, deadline - time.monotonic())
+            acquired = self.scheduler._lock.acquire(timeout=lock_timeout if lock_timeout is not None else -1)
+            if not acquired:
+                raise TimeoutError(f"RPC call to {method} timed out waiting for scheduler lock.")
+            try:
+                # Broadcast RPC request to all workers via unified message queue
+                self.scheduler.mq.enqueue(rpc_request)
 
-            # Determine which workers we expect responses from
-            num_responses = 1 if unique_reply_rank is not None else self.od_config.num_gpus
+                # Determine which workers we expect responses from
+                num_responses = 1 if unique_reply_rank is not None else self.od_config.num_gpus
 
-            responses = []
-            for _ in range(num_responses):
-                dequeue_timeout = None if deadline is None else (deadline - time.monotonic())
-                try:
-                    if self.scheduler.result_mq is None:
-                        raise RuntimeError("Result queue not initialized")
+                responses = []
+                for _ in range(num_responses):
+                    dequeue_timeout = None if deadline is None else max(0, deadline - time.monotonic())
+                    try:
+                        if self.scheduler.result_mq is None:
+                            raise RuntimeError("Result queue not initialized")
 
-                    response = self.scheduler.result_mq.dequeue(timeout=dequeue_timeout)
+                        response = self.scheduler.result_mq.dequeue(timeout=dequeue_timeout)
 
-                    # Check if response indicates an error
-                    if isinstance(response, dict) and response.get("status") == "error":
-                        raise RuntimeError(
-                            f"Worker failed with error '{response.get('error')}', "
-                            "please check the stack trace above for the root cause"
-                        )
+                        # Check if response indicates an error
+                        if isinstance(response, dict) and response.get("status") == "error":
+                            raise RuntimeError(
+                                f"Worker failed with error '{response.get('error')}', "
+                                "please check the stack trace above for the root cause"
+                            )
 
-                    responses.append(response)
-                except TimeoutError as e:
-                    raise TimeoutError(f"RPC call to {method} timed out.") from e
+                        responses.append(response)
+                    except TimeoutError as e:
+                        raise TimeoutError(f"RPC call to {method} timed out.") from e
 
-            return responses[0] if unique_reply_rank is not None else responses
+                return responses[0] if unique_reply_rank is not None else responses
+            finally:
+                self.scheduler._lock.release()
 
         except Exception as e:
             logger.error(f"RPC call failed: {e}")