Implement worker support for multiturn

Signed-off-by: Samuel Monson <[email protected]>

Author: sjmonson

src/guidellm/request/loader.py

@@ -105,7 +105,7 @@ def __init__(
         self.preserve_iter_state = iter_type == "infinite"  # ensure no caching requests
         self._preserved_iter = None
 
-    def __iter__(self) -> Iterator[list[GenerationRequest]]:
+    def __iter__(self) -> Iterator[list[tuple[GenerationRequest, float]]]:
         scope_create_count = 0
 
         while (dataset_iter := self._get_dataset_iter(scope_create_count)) is not None:
@@ -260,7 +260,9 @@ def _get_dataset_iter(
 
         return dataset_iter
 
-    def _create_requests(self, item: dict[str, Any]) -> list[GenerationRequest]:
+    def _create_requests(
+        self, item: dict[str, Any]
+    ) -> list[tuple[GenerationRequest, float]]:
         prompts = list(item[self.column_mappings["prompt_column"]])
         prompts_tokens: list[Optional[int]] = (
             list(item[self.column_mappings["prompt_tokens_count_column"]])
@@ -281,15 +283,24 @@ def _create_requests(self, item: dict[str, Any]) -> list[GenerationRequest]:
             )
 
         return [
-            GenerationRequest(
-                request_type=settings.preferred_route,
-                content=prompt,
-                stats=(
-                {"prompt_tokens": prompt_tokens} if prompt_tokens is not None else {}
-                ),
-                constraints=(
-                    {"output_tokens": output_tokens} if output_tokens is not None else {}
+            (
+                GenerationRequest(
+                    request_type=settings.preferred_route,
+                    content=prompt,
+                    stats=(
+                        {"prompt_tokens": prompt_tokens}
+                        if prompt_tokens is not None
+                        else {}
+                    ),
+                    constraints=(
+                        {"output_tokens": output_tokens}
+                        if output_tokens is not None
+                        else {}
+                    ),
                 ),
+                0.0,  # TODO: delay
+            )
+            for prompt, prompt_tokens, output_tokens in zip(
+                prompts, prompts_tokens, outputs_tokens
             )
-            for prompt, prompt_tokens, output_tokens in zip(prompts, prompts_tokens, outputs_tokens)
         ]

src/guidellm/scheduler/__init__.py

@@ -15,16 +15,20 @@
 from .objects import (
     BackendInterface,
     BackendT,
+    HistoryT,
     MeasuredRequestTimings,
     MultiTurnRequestT,
+    MultiTurnT,
     RequestSchedulerTimings,
     RequestT,
     ResponseT,
+    ScheduledRequestAugmentation,
     ScheduledRequestInfo,
     SchedulerMessagingPydanticRegistry,
     SchedulerState,
     SchedulerUpdateAction,
     SchedulerUpdateActionProgress,
+    TurnT,
 )
 from .scheduler import Scheduler
 from .strategies import (
@@ -56,6 +60,7 @@
     "ConstraintInitializer",
     "ConstraintsInitializerFactory",
     "Environment",
+    "HistoryT",
     "LastCompletionRequestTimings",
     "MaxDurationConstraint",
     "MaxErrorRateConstraint",
@@ -64,13 +69,15 @@
     "MaxNumberConstraint",
     "MeasuredRequestTimings",
     "MultiTurnRequestT",
+    "MultiTurnT",
     "NoDelayRequestTimings",
     "NonDistributedEnvironment",
     "PoissonRateRequestTimings",
     "PydanticConstraintInitializer",
     "RequestSchedulerTimings",
     "RequestT",
     "ResponseT",
+    "ScheduledRequestAugmentation",
     "ScheduledRequestInfo",
     "ScheduledRequestTimings",
     "Scheduler",
@@ -84,6 +91,7 @@
     "StrategyType",
     "SynchronousStrategy",
     "ThroughputStrategy",
+    "TurnT",
     "UnserializableConstraintInitializer",
     "WorkerProcess",
     "WorkerProcessGroup",

src/guidellm/scheduler/objects.py

@@ -19,7 +19,6 @@
     Literal,
     Protocol,
     TypeVar,
-    Union,
 )
 
 from pydantic import Field, computed_field
@@ -35,34 +34,50 @@
 __all__ = [
     "BackendInterface",
     "BackendT",
+    "HistoryT",
     "MeasuredRequestTimings",
     "MultiTurnRequestT",
+    "MultiTurnT",
     "RequestSchedulerTimings",
     "RequestT",
     "ResponseT",
+    "ScheduledRequestAugmentation",
     "ScheduledRequestInfo",
     "SchedulerMessagingPydanticRegistry",
     "SchedulerState",
     "SchedulerUpdateAction",
     "SchedulerUpdateActionProgress",
+    "TurnT",
 ]
 
 RequestT = TypeVar("RequestT")
 """Generic request object type for scheduler processing."""
 
+# TODO: Remove
+MultiTurnRequestT = RequestT
+
 ResponseT = TypeVar("ResponseT")
 """Generic response object type returned by backend processing."""
 
-MultiTurnRequestT = TypeAliasType(
-    "MultiTurnRequestT",
-    Union[
-        list[Union[RequestT, tuple[RequestT, float]]],
-        tuple[Union[RequestT, tuple[RequestT, float]]],
-    ],
+TurnT = TypeAliasType(
+    "TurnT",
+    tuple[RequestT, "ScheduledRequestAugmentation", "ScheduledRequestInfo"],
+    type_params=(RequestT,),
+)
+
+MultiTurnT = TypeAliasType(
+    "MultiTurnT",
+    list[TurnT[RequestT]],
     type_params=(RequestT,),
 )
 """Multi-turn request structure supporting conversation history with optional delays."""
 
+HistoryT = TypeAliasType(
+    "HistoryT",
+    list[tuple[RequestT, ResponseT]],
+    type_params=(RequestT, ResponseT),
+)
+
 
 class SchedulerMessagingPydanticRegistry(RegistryMixin[RegistryObjT]):
     """
@@ -71,6 +86,21 @@ class SchedulerMessagingPydanticRegistry(RegistryMixin[RegistryObjT]):
     """
 
 
+@SchedulerMessagingPydanticRegistry.register()
+class ScheduledRequestAugmentation(StandardBaseModel):
+    """
+    Adjustments to scheduler logic for a paired request.
+    """
+
+    post_requeue_delay: float = Field(
+        description=(
+            "Delay in seconds to wait after a request to "
+            "queue the next request in the conversation."
+        ),
+        default=0.0,
+    )
+
+
 @SchedulerMessagingPydanticRegistry.register()
 class RequestSchedulerTimings(StandardBaseModel):
     """

src/guidellm/scheduler/worker.py

@@ -31,9 +31,12 @@
 
 from guidellm.scheduler.objects import (
     BackendInterface,
+    HistoryT,
     MultiTurnRequestT,
+    MultiTurnT,
     RequestT,
     ResponseT,
+    ScheduledRequestAugmentation,
     ScheduledRequestInfo,
     SchedulerMessagingPydanticRegistry,
 )
@@ -118,6 +121,9 @@ def __init__(
         self.startup_completed = False
         self.backend_started = False
         self.messaging_started = False
+        self.turns_queue: list[
+            tuple[HistoryT[RequestT, ResponseT], MultiTurnT[RequestT]]
+        ] = []
 
     def run(self):
         """
@@ -302,16 +308,19 @@ async def _cancel_requests_loop(self):
             self._send_update("cancelled", None, request, request_info)
 
     async def _process_next_request(self):
-        request: RequestT | MultiTurnRequestT[RequestT] | None = None
+        request: RequestT | None = None
         request_info: ScheduledRequestInfo | None = None
         response: ResponseT | None = None
+        aug: ScheduledRequestAugmentation | None = None
 
         try:
             # Pull request from the queue
-            request, request_info = await self.messaging.get()
-
-            if isinstance(request, (list, tuple)):
-                raise NotImplementedError("Multi-turn requests are not yet supported")
+            history, conversation = (
+                self.turns_queue.pop(0)
+                if self.turns_queue
+                else ([], await self.messaging.get())
+            )
+            request, aug, request_info = conversation.pop(0)
 
             # Calculate targeted start and set pending state for request
             request_info.scheduler_node_id = self.messaging.worker_index
@@ -341,6 +350,12 @@ async def _process_next_request(self):
             request_info.scheduler_timings.resolve_end = time.time()
             self._send_update("completed", response, request, request_info)
 
+            # If multi-turn, queue up next turn(s)
+            # TODO: Move to callback and support delay
+            if conversation:  # more turns to process
+                history.append((request, response))
+                self.turns_queue.append((history, conversation))
+
             response = request = request_info = None
         except asyncio.CancelledError:
             # Handle cancellation

src/guidellm/scheduler/worker_group.py

@@ -26,8 +26,10 @@
 from guidellm.scheduler.objects import (
     BackendInterface,
     MultiTurnRequestT,
+    MultiTurnT,
     RequestT,
     ResponseT,
+    ScheduledRequestAugmentation,
     ScheduledRequestInfo,
     SchedulerMessagingPydanticRegistry,
     SchedulerState,
@@ -471,9 +473,9 @@ def __init__(
 
     def requests_generator(
         self,
-        requests: Iterable[RequestT | MultiTurnRequestT[RequestT]] | None,
-        cycle_requests: Iterable[RequestT | MultiTurnRequestT[RequestT]] | None,
-    ) -> Generator[tuple[RequestT | MultiTurnRequestT[RequestT],], None, None]:
+        requests: Iterable[Iterable[tuple[RequestT, float]]] | None,
+        cycle_requests: Iterable[Iterable[tuple[RequestT, float]]] | None,
+    ) -> Generator[MultiTurnT[RequestT], None, None]:
         """
         Generate request-info pairs for worker processing with constraint evaluation.
 
@@ -494,31 +496,40 @@ def _iter():
                 while True:
                     yield from cycle_requests
 
-        count = 0
-        request_info: ScheduledRequestInfo = None
+        count: int = 0
+        stop_queueing: bool = False
+
+        def _turn_iter(requests_chain: Iterable[tuple[RequestT, float]]):
+            nonlocal count, stop_queueing
+            for request, delay in requests_chain:
+                count += 1
+
+                if hasattr(request, "request_id"):
+                    request_id = request.request_id
+                elif hasattr(request, "id"):
+                    request_id = request.id
+                else:
+                    request_id = str(uuid.uuid4())
+                request_augmentation = ScheduledRequestAugmentation(
+                    post_requeue_delay=delay
+                )
+                request_info: ScheduledRequestInfo = ScheduledRequestInfo(
+                    request_id=request_id,
+                    status="queued",
+                    scheduler_process_id=0,
+                    scheduler_start_time=self.start_time,
+                )
+                state_update = self._locked_update(request_info)
+                yield (request, request_augmentation, request_info)
+
+                if state_update.stop_queueing:
+                    stop_queueing = True
+                    return
+
         for request_chain in _iter():
-            if isinstance(request_chain, (list, tuple)):
-                request = request_chain[0]
-            else:
-                request = request_chain
-            count += 1
-
-            if hasattr(request, "request_id"):
-                request_id = request.request_id
-            elif hasattr(request, "id"):
-                request_id = request.id
-            else:
-                request_id = str(uuid.uuid4())
-            request_info: ScheduledRequestInfo = ScheduledRequestInfo(
-                request_id=request_id,
-                status="queued",
-                scheduler_process_id=0,
-                scheduler_start_time=self.start_time,
-            )
-            state_update = self._locked_update(request_info)
-            yield (request, request_info)
+            yield list(_turn_iter(request_chain))
 
-            if state_update.stop_queueing:
+            if stop_queueing:
                 self.stop_send_requests_event.set()
                 return