Prep scheduler for multiturn

Revert "Revert loop logic changes"

This reverts commit bcc2f8c.

Revert "Strip out multiturn features"

This reverts commit a524469.

Author: sjmonson

src/guidellm/request/__init__.py

@@ -5,14 +5,17 @@
     RequestLoaderDescription,
 )
 from .request import GenerationRequest
+from .session import GenerativeRequestSession, RequestSession
 from .types import RequestT, ResponseT
 
 __all__ = [
     "GenerationRequest",
     "GenerativeRequestLoader",
     "GenerativeRequestLoaderDescription",
+    "GenerativeRequestSession",
     "RequestLoader",
     "RequestLoaderDescription",
+    "RequestSession",
     "RequestT",
     "ResponseT",
 ]

src/guidellm/request/loader.py

@@ -15,6 +15,7 @@
 from guidellm.dataset import ColumnInputTypes, load_dataset
 from guidellm.objects import StandardBaseModel
 from guidellm.request.request import GenerationRequest
+from guidellm.request.session import GenerativeRequestSession
 
 __all__ = [
     "GenerativeRequestLoader",
@@ -105,14 +106,14 @@ def __init__(
         self.preserve_iter_state = iter_type == "infinite"  # ensure no caching requests
         self._preserved_iter = None
 
-    def __iter__(self) -> Iterator[GenerationRequest]:
+    def __iter__(self) -> Iterator[GenerativeRequestSession]:
         scope_create_count = 0
 
         while (dataset_iter := self._get_dataset_iter(scope_create_count)) is not None:
             scope_create_count += 1
 
             for item in dataset_iter:
-                yield self._create_request(item)
+                yield GenerativeRequestSession(self._create_request(item))
 
             self._preserved_iter = None
 

src/guidellm/request/session.py

@@ -0,0 +1,55 @@
+from abc import ABC, abstractmethod
+from typing import Generic, TypeVar
+
+from guidellm.backend.response import ResponseSummary
+from guidellm.request.request import GenerationRequest
+
+__all__ = ["GenerativeRequestSession", "RequestSession"]
+
+RequestT = TypeVar("RequestT")
+ResponseT = TypeVar("ResponseT")
+
+
+class RequestSession(ABC, Generic[RequestT, ResponseT]):
+    """
+    A series of requests that build upon each other to
+    form a conversion between the user and the model.
+    """
+
+    @abstractmethod
+    def __len__(self) -> int: ...
+
+    @abstractmethod
+    def get_next_request(self) -> RequestT: ...
+
+    @abstractmethod
+    def get_next_delay(self) -> float: ...
+
+    @abstractmethod
+    def push_response(self, response: ResponseT) -> None: ...
+
+    @property
+    @abstractmethod
+    def complete(self) -> bool: ...
+
+
+class GenerativeRequestSession(RequestSession[GenerationRequest, ResponseSummary]):
+    def __init__(self, request: GenerationRequest) -> None:
+        self.request = request
+        self._complete = False
+
+    def __len__(self) -> int:
+        return 1
+
+    def get_next_request(self) -> GenerationRequest:
+        return self.request
+
+    def get_next_delay(self) -> float:
+        return 0.0
+
+    def push_response(self, response: ResponseSummary) -> None:  # noqa: ARG002
+        self._complete = True
+
+    @property
+    def complete(self) -> bool:
+        return self._complete

src/guidellm/scheduler/result.py

@@ -6,6 +6,7 @@
 )
 
 from guidellm.objects import StandardBaseModel
+from guidellm.request.session import RequestSession
 from guidellm.request.types import RequestT, ResponseT
 from guidellm.scheduler.strategy import SchedulingStrategy
 
@@ -142,7 +143,7 @@ class SchedulerRequestResult(
 
 @dataclass
 class WorkerProcessRequest(Generic[RequestT, ResponseT]):
-    request: RequestT
+    session: RequestSession[RequestT, ResponseT]
     timeout_time: float
     queued_time: float
 

src/guidellm/scheduler/scheduler.py

@@ -130,17 +130,15 @@ async def run(
             futures, queues, stop_event = await self._start_processes(
                 manager, executor, scheduling_strategy
             )
-            run_info, requests_iter, times_iter = self._run_setup(
+            run_info, requests_iter = self._run_setup(
                 futures, scheduling_strategy, max_number, max_duration
             )
 
             # Add some initial requests to the queue
             requests_iter = self._add_requests(
                 requests_iter,
                 queues.requests,
-                times_iter,
                 run_info,
-                loop_limit=run_info.strategy.queued_requests_limit,
             )
             # Wait for the test to start
             await asyncio.sleep(time.time() - scheduling_strategy.start_time)
@@ -171,7 +169,6 @@ async def run(
                     requests_iter = self._add_requests(
                         requests_iter,
                         queues.requests,
-                        times_iter,
                         run_info,
                     )
                     await asyncio.sleep(0)  # enable requests to start
@@ -244,6 +241,7 @@ async def _start_processes(
                     queues,
                     scheduling_strategy,
                     stop_event,
+                    False,  # TODO: Make configurable
                     requests_limit,
                     id_,
                     num_processes,
@@ -260,9 +258,8 @@ def _run_setup(
         scheduling_strategy: SchedulingStrategy,
         max_number: Optional[int],
         max_duration: Optional[float],
-    ) -> tuple[SchedulerRunInfo, Iterator[Any], Iterator[float]]:
+    ) -> tuple[SchedulerRunInfo, Iterator[Any]]:
         requests_iter = iter(self.request_loader)
-        times_iter = iter(scheduling_strategy.request_times())
         end_time = scheduling_strategy.start_time + (max_duration or math.inf)
         end_number = max_number or math.inf
 
@@ -288,42 +285,39 @@ def _run_setup(
             strategy=scheduling_strategy,
         )
 
-        return info, requests_iter, times_iter
+        return info, requests_iter
 
     def _add_requests(
         self,
         requests_iter: Optional[Iterator[Any]],
         requests_queue: Queue[WorkerProcessRequest[RequestT, ResponseT]],
-        times_iter: Iterator[float],
         run_info: SchedulerRunInfo,
-        loop_limit: Optional[int] = None,
     ) -> Optional[Iterator[Any]]:
         if requests_iter is not None:
             try:
                 added_count = 0
 
+                if time.time() >= run_info.end_time:
+                    raise StopIteration
+
                 while not requests_queue.full() and added_count < (
-                    loop_limit or settings.max_add_requests_per_loop
+                    run_info.strategy.queued_requests_limit
+                    or settings.max_add_requests_per_loop
                 ):
                     if run_info.created_requests >= run_info.end_number:
                         raise StopIteration
 
-                    if (
-                        next(times_iter) >= run_info.end_time
-                        or time.time() >= run_info.end_time
-                    ):
-                        raise StopIteration
-
-                    work_req = WorkerProcessRequest[RequestT, ResponseT](
-                        request=next(requests_iter),
+                    session = next(requests_iter)
+                    work_req = WorkerProcessRequest(
+                        session=session,
                         timeout_time=run_info.end_time,
                         queued_time=time.time(),
                     )
                     requests_queue.put(work_req)
 
-                    run_info.created_requests += 1
-                    run_info.queued_requests += 1
-                    added_count += 1
+                    run_info.created_requests += len(session)
+                    run_info.queued_requests += len(session)
+                    added_count += len(session)
             except StopIteration:
                 # we've reached the limit number, limit time, or exhausted the requests
                 # set to None to stop adding more and tell the loop no more requests

src/guidellm/scheduler/worker.py

@@ -122,8 +122,8 @@ async def resolve_scheduler_request(
         start_time: float,
         results_queue: Queue[WorkerProcessResult[RequestT, ResponseT]],
         process_id: int,
-    ):
-        request = process_request.request
+    ) -> WorkerProcessRequest[RequestT, ResponseT]:
+        request = process_request.session.get_next_request()
         timeout_time = process_request.timeout_time
         queued_time = process_request.queued_time
 
@@ -170,17 +170,22 @@ async def resolve_scheduler_request(
         )
         asyncio.create_task(self.send_result(results_queue, result))
 
+        process_request.session.push_response(response)
+        return process_request
+
     def process_loop_asynchronous(
         self,
         queues: MPQueues[RequestT, ResponseT],
         strategy: SchedulingStrategy,
         stop_event: Event,
+        prioritize_sessions: bool,
         max_concurrency: int,
         process_id: int,
         num_processes: int,
     ):
         async def _process_runner():
             lock = asyncio.Semaphore(max_concurrency)
+            pending_requests: list[WorkerProcessRequest[RequestT, ResponseT]] = []
             times_iter = islice(
                 strategy.request_times(),
                 process_id,
@@ -197,18 +202,50 @@ async def _process_runner():
                 await asyncio.sleep(start_time - time.time() - 1)
                 await lock.acquire()
 
+                process_request = None
                 try:
-                    process_request = queues.requests.get_nowait()
+                    process_request = (
+                        pending_requests.pop()
+                        if pending_requests
+                        else queues.requests.get_nowait()
+                    )
                     dequeued_time = time.time()
                 except QueueEmpty:
                     lock.release()
                     continue
 
+                async def wait_then_requeue(
+                    process_request: WorkerProcessRequest[RequestT, ResponseT],
+                ):
+                    # Wait to requeue the request session if it specifies a delay
+                    if delay := process_request.session.get_next_delay():
+                        await asyncio.sleep(delay)
+
+                    # Push session to the stack
+                    process_request.queued_time = time.time()
+                    pending_requests.append(process_request)
+                    if prioritize_sessions:
+                        # Release the lock with the session on top of the stack
+                        lock.release()
+
                 def _request_callback(
-                    _: asyncio.Future[WorkerProcessRequest[RequestT, ResponseT]],
+                    future: asyncio.Future[WorkerProcessRequest[RequestT, ResponseT]],
                 ):
+                    # If we are prioritizing sessions, hold
+                    # the lock until the session is done
                     nonlocal lock
-                    lock.release()
+                    if not prioritize_sessions:
+                        lock.release()
+
+                    try:
+                        process_request = future.result()
+                    except asyncio.CancelledError:
+                        return
+                    if not process_request.session.complete:
+                        asyncio.create_task(wait_then_requeue(process_request))
+                    elif prioritize_sessions:
+                        # no more requests in this session, release the lock
+                        lock.release()
 
                 task = asyncio.create_task(
                     self.resolve_scheduler_request(
@@ -282,6 +319,7 @@ def process_loop_asynchronous(
         queues: MPQueues[GenerationRequest, ResponseSummary],
         strategy: SchedulingStrategy,
         stop_event: Event,
+        prioritize_sessions: bool,
         max_concurrency: int,
         process_id: int,
         num_processes: int,
@@ -291,6 +329,7 @@ def process_loop_asynchronous(
             queues=queues,
             strategy=strategy,
             stop_event=stop_event,
+            prioritize_sessions=prioritize_sessions,
             max_concurrency=max_concurrency,
             process_id=process_id,
             num_processes=num_processes,