Attempts to fix stranded messages

Author: markurtz

src/guidellm/scheduler/worker.py

@@ -78,7 +78,7 @@ def __init__(
         startup_barrier: ProcessingBarrier,
         shutdown_event: ProcessingEvent,
         error_event: ProcessingEvent,
-        completed_event: ProcessingEvent,
+        requests_completed_event: ProcessingEvent,
         backend: BackendInterface[RequestT, MeasuredRequestTimingsT, ResponseT],
         request_timings: ScheduledRequestTimings,
     ):
@@ -90,7 +90,8 @@ def __init__(
         :param startup_barrier: Multiprocessing barrier for coordinated startup
         :param shutdown_event: Event for signaling graceful shutdown
         :param error_event: Event for signaling error conditions across processes
-        :param completed_event: Event for signaling when this worker has completed
+        :param requests_completed_event: Event for signaling when the main process
+            has stopped sending requests / all requests are added to the queue
         :param backend: Backend instance for processing requests
         :param request_timings: Timing strategy for request scheduling
         """
@@ -99,7 +100,7 @@ def __init__(
         self.startup_barrier = startup_barrier
         self.shutdown_event = shutdown_event
         self.error_event = error_event
-        self.completed_event = completed_event
+        self.requests_completed_event = requests_completed_event
         self.backend = backend
         self.request_timings = request_timings
         self.startup_completed = False
@@ -126,8 +127,6 @@ def run(self):
                 f"Worker process {self.messaging.worker_index} encountered an "
                 f"error: {err}"
             ) from err
-        finally:
-            self.completed_event.set()
 
     async def run_async(self):
         """
@@ -212,11 +211,10 @@ async def _run_async_requests_processing(self):
             await self.backend.validate()
 
             # Get messaging system ready
-            processing_cancelled = ThreadingEvent()
             all_requests_processed = ThreadingEvent()
             await self.messaging.start(
                 send_stop_criteria=[all_requests_processed],
-                receive_stop_criteria=[processing_cancelled],
+                receive_stop_criteria=[self.requests_completed_event, self.error_event],
                 pydantic_models=list(
                     SchedulerMessagingPydanticRegistry.registry.values()
                 ),
@@ -255,7 +253,6 @@ def _task_done(task):
                 pending_tasks.add(request_task)
                 request_task.add_done_callback(_task_done)
         except (asyncio.CancelledError, Exception) as err:
-            processing_cancelled.set()
             await self._cancel_remaining_requests(pending_tasks, all_requests_processed)
             await self.messaging.stop()
             await self.backend.process_shutdown()
@@ -323,27 +320,17 @@ def _send_update(
         prev_status = request_info.status
 
         try:
-            if (new_status == "in_progress" and prev_status != "in_progress") or (
-                new_status != "in_progress" and prev_status == "pending"
-            ):
-                request_info.status = "in_progress"
-                self.messaging.put_sync(
-                    (None, request, request_info.model_copy()),
-                    timeout=-1,
-                )
-                prev_status = new_status
-
-            if prev_status == "in_progress" and new_status in {
-                "completed",
-                "errored",
-                "cancelled",
-            }:
-                request_info.status = new_status
-                self.messaging.put_sync(
-                    (response, request, request_info),  # last update, no copy
-                    timeout=-1,
-                )
-                prev_status = new_status
+            request_info.status = new_status
+            request_info = (
+                request_info.model_copy()
+                if new_status not in {"completed", "errored", "cancelled"}
+                else request_info  # last update, don't need to copy
+            )
+            self.messaging.put_sync(
+                (response, request, request_info),
+                timeout=-1,
+            )
+            prev_status = new_status
         except Exception as exc:
             # Reset status to last one that succeeded or started function with
             # Calling logic can retry after handling error, if possible

src/guidellm/scheduler/worker_group.py

@@ -120,7 +120,7 @@ def __init__(
         self.mp_context = None
         self.mp_manager = None
         self.processes: list[BaseProcess] = None
-        self.processes_completed_events: list[Event] = None
+        self.requests_completed_event: Event = None
         self.startup_barrier: Barrier = None
         self.shutdown_event: Event = None
         self.error_event: Event = None
@@ -176,8 +176,11 @@ async def create_processes(self):
             raise RuntimeError("num_processes resolved to 0; increase limits/config")
 
         per_proc_max_conc = max_conc // num_processes
-        per_proc_max_receive_buffer = max(
-            1, math.floor(per_proc_max_conc * settings.mp_proc_receive_buffer_per)
+        max_pending_size = max(
+            1, math.floor(max_conc * settings.mp_max_pending_buffer_percent)
+        )
+        per_proc_max_buffer_size = max(
+            1, math.floor(per_proc_max_conc * settings.mp_max_worker_buffer_percent)
         )
 
         # Initialize multiprocessing components
@@ -186,12 +189,13 @@ async def create_processes(self):
         self.startup_barrier = self.mp_context.Barrier(num_processes + 1)
         self.shutdown_event = self.mp_context.Event()
         self.error_event = self.mp_context.Event()
+        self.requests_completed_event = self.mp_context.Event()
 
         if settings.mp_messaging_object == "queue":
             self.messaging = InterProcessMessagingQueue(
                 serialization=settings.mp_serialization,
                 encoding=settings.mp_encoding,
-                max_send_size=max_conc,
+                max_pending_size=max_pending_size,
                 max_buffer_send_size=settings.mp_requests_send_buffer_size,
                 poll_interval=settings.mp_poll_interval,
             )
@@ -200,7 +204,7 @@ async def create_processes(self):
                 manager=self.mp_manager,
                 serialization=settings.mp_serialization,
                 encoding=settings.mp_encoding,
-                max_send_size=max_conc,
+                max_pending_size=max_pending_size,
                 max_buffer_send_size=settings.mp_requests_send_buffer_size,
                 poll_interval=settings.mp_poll_interval,
             )
@@ -209,32 +213,30 @@ async def create_processes(self):
                 num_workers=num_processes,
                 serialization=settings.mp_serialization,
                 encoding=settings.mp_encoding,
-                max_send_size=max_conc,
+                max_pending_size=max_pending_size,
                 max_buffer_send_size=settings.mp_requests_send_buffer_size,
                 poll_interval=settings.mp_poll_interval,
             )
 
         # Initialize worker processes
         self.processes = []
-        self.processes_completed_events = []
         for rank in range(num_processes):
             # Distribute any remainder across the first N ranks
             async_limit = per_proc_max_conc + (
                 1 if rank < (max_conc % num_processes) else 0
             )
 
-            worker_completed_event = self.mp_context.Event()
             worker = WorkerProcess[RequestT, MeasuredRequestTimingsT, ResponseT](
                 messaging=self.messaging.create_worker_copy(
                     worker_index=rank,
                     max_buffer_send_size=None,
-                    max_buffer_receive_size=per_proc_max_receive_buffer,
+                    max_buffer_receive_size=per_proc_max_buffer_size,
                 ),
                 async_limit=async_limit,
                 startup_barrier=self.startup_barrier,
                 shutdown_event=self.shutdown_event,
                 error_event=self.error_event,
-                completed_event=worker_completed_event,
+                requests_completed_event=self.requests_completed_event,
                 backend=self.backend,
                 request_timings=self.strategy.create_request_timings(
                     local_rank=rank,
@@ -245,7 +247,6 @@ async def create_processes(self):
             proc = self.mp_context.Process(target=worker.run, daemon=False)
             proc.start()
             self.processes.append(proc)
-            self.processes_completed_events.append(worker_completed_event)
 
         reason, _ = await synchronous_to_exitable_async(
             synchronous=None,
@@ -279,7 +280,7 @@ async def start(self, start_time: float):
         self._state = _WorkerGroupState[RequestT, MeasuredRequestTimingsT, ResponseT](
             start_time=start_time,
             num_processes=len(self.processes),
-            processes_completed_events=self.processes_completed_events,
+            processes=self.processes,
             constraints=self.constraints,
             shutdown_event=self.shutdown_event,
         )
@@ -289,6 +290,7 @@ async def start(self, start_time: float):
             ),
             receive_callback=self._state.update_callback_receive,
             send_stop_criteria=[self.shutdown_event, self.error_event],
+            send_stopped_event=self.requests_completed_event,
             receive_stop_criteria=[self.error_event, self._state.stop_callback_receive],
             pydantic_models=list(SchedulerMessagingPydanticRegistry.registry.values()),
         )
@@ -408,7 +410,7 @@ def __init__(
         self,
         start_time: float,
         num_processes: int,
-        processes_completed_events: list[Event],
+        processes: list[BaseProcess],
         constraints: dict[str, Constraint],
         shutdown_event: Event,
     ):
@@ -419,7 +421,7 @@ def __init__(
             num_processes=num_processes,
             start_time=start_time,
         )
-        self.processes_completed_events = processes_completed_events
+        self.processes = processes
         self._constraints = constraints
         self._internal_constraints: dict[str, Constraint] = {}
         self._shutdown_event = shutdown_event
@@ -544,7 +546,7 @@ def stop_callback_receive(
             and messaging.send_stopped_event.is_set()  # No more requests will be added
             and self._shutdown_event.is_set()  # processing should stop
             and all(
-                event.is_set() for event in self.processes_completed_events
+                not proc.is_alive() for proc in self.processes
             )  # no more updates will be added by workers
         )
 
@@ -601,21 +603,19 @@ def _update_new_request(self):
         self._state.queued_requests += 1
 
     def _update_new_response(self, info: ScheduledRequestInfo[MeasuredRequestTimingsT]):
-        if info.status == "in_progress":
+        if info.status == "in_progress" or (
+            info.status == "cancelled" and info.scheduler_timings.resolve_start is None
+            # Cancelled request that never sent a progress update
+        ):
             self._state.queued_requests -= 1
             self._state.processing_requests += 1
-        elif info.status in ("completed", "errored", "cancelled"):
+
+        if info.status in ("completed", "errored", "cancelled"):
             self._state.processing_requests -= 1
             self._state.processed_requests += 1
             self._state.successful_requests += 1 if info.status == "completed" else 0
             self._state.errored_requests += 1 if info.status == "errored" else 0
             self._state.cancelled_requests += 1 if info.status == "cancelled" else 0
-        else:
-            raise ValueError(
-                f"Unknown request status: {info.status}. "
-                "Supported statuses are: queued, pending, in_progress, "
-                "completed, errored, cancelled."
-            )
 
     def _update_with_constraints(
         self, info: ScheduledRequestInfo[MeasuredRequestTimingsT]

src/guidellm/settings.py

@@ -140,7 +140,8 @@ class Settings(BaseSettings):
     mp_messaging_object: Literal["queue", "manager_queue", "pipe"] = "queue"
     mp_requests_send_buffer_size: int = 1
     mp_poll_interval: float = 0.1
-    mp_proc_receive_buffer_per: float = 0.1
+    mp_max_pending_buffer_percent: float = 0.5
+    mp_max_worker_buffer_percent: float = 0.2
     max_concurrency: int = 512
     max_worker_processes: int = 10
     scheduler_start_delay_non_distributed: float = 0.1