Initial agentic relationship support in scheduler

src/guidellm/scheduler/schemas.py

@@ -16,7 +16,7 @@
 from pydantic import Field
 from typing_extensions import TypeAliasType
 
-from guidellm.schemas import RequestInfo, RequestSettings, StandardBaseModel
+from guidellm.schemas import RequestInfo, StandardBaseModel, StartNode
 from guidellm.utils.registry import RegistryMixin, RegistryObjT
 
 __all__ = [
@@ -63,7 +63,7 @@
 # NOTE: This is the interface between data and scheduler.
 DatasetIterT = TypeAliasType(
     "DatasetIterT",
-    Iterable[Iterable[tuple[RequestT, RequestSettings]]],
+    Iterable[StartNode[RequestT]],
     type_params=(RequestT,),
 )
 """

src/guidellm/scheduler/worker_group.py

@@ -15,7 +15,7 @@
 import threading
 import time
 import uuid
-from collections.abc import AsyncIterator, Generator, Iterable
+from collections.abc import AsyncIterator, Generator
 from multiprocessing import get_context
 from multiprocessing.context import BaseContext
 from multiprocessing.managers import BaseManager
@@ -31,15 +31,14 @@
     BackendInterface,
     ConversationT,
     DatasetIterT,
-    RequestDataT,
     RequestT,
     ResponseT,
     SchedulerState,
     SchedulerUpdateAction,
 )
 from guidellm.scheduler.strategies import SchedulingStrategy
 from guidellm.scheduler.worker import WorkerProcess
-from guidellm.schemas import RequestInfo, RequestSettings
+from guidellm.schemas import DAGNode, RequestInfo, RequestNode
 from guidellm.settings import settings
 from guidellm.utils.messaging import (
     InterProcessMessaging,
@@ -565,41 +564,45 @@ def requests_generator(
             count = 0
             stop_queueing: bool = False
 
-            def _turn_iter(
-                requests_chain: Iterable[tuple[RequestT, RequestSettings]],
-            ) -> Generator[RequestDataT[RequestT], None, None]:
+            def apply_request_info(
+                node: DAGNode[RequestT], conv_id: str, turn_idx: int
+            ):
                 nonlocal count, stop_queueing
-                # NOTE: This allows users to correlate requests in post-processing
-                conv_id = str(uuid.uuid4())
-                for i, (request, setting) in enumerate(requests_chain):
-                    count += 1
-
-                    request_id = self._find_request_id(request)
-                    request_info: RequestInfo = RequestInfo(
-                        request_id=request_id,
-                        conversation_id=conv_id,
-                        turn_index=i,
-                        status="queued",
-                        scheduler_process_id=0,
-                        scheduler_start_time=self.start_time,
-                        settings=setting,
-                    )
-                    state_update = self._locked_update(request_info)
-                    request_info.timings.queued = time.time()
-                    if self.messaging.buffer_receive_queue is None:
-                        raise RuntimeError("buffer receive queue is None")
-                    self.messaging.buffer_receive_queue.sync_put(
-                        (None, request, request_info, state_update.state)
-                    )
+                if not isinstance(node, RequestNode):
+                    return
 
-                    yield request, request_info
+                count += 1
+
+                request_id = self._find_request_id(node.request)
+                request_info: RequestInfo = RequestInfo(
+                    request_id=request_id,
+                    conversation_id=conv_id,
+                    turn_index=turn_idx,
+                    status="queued",
+                    scheduler_process_id=0,
+                    scheduler_start_time=self.start_time,
+                    settings=node.settings,
+                )
+                state_update = self._locked_update(request_info)
+                request_info.timings.queued = time.time()
+                if self.messaging.buffer_receive_queue is None:
+                    raise RuntimeError("buffer receive queue is None")
+                self.messaging.buffer_receive_queue.sync_put(
+                    (None, node.request, request_info, state_update.state)
+                )
 
-                    if state_update.stop_queueing:
-                        stop_queueing = True
-                        return
+                if state_update.stop_queueing:
+                    stop_queueing = True
 
             for request_chain in requests:
-                yield list(_turn_iter(request_chain))
+                # TODO Spawn should probably rotate conversation ID
+                conv_id = str(uuid.uuid4())
+                node_iter = request_chain.apply_to_all(apply_request_info)
+                # FIXME this is wrong node_idx != turn_idx
+                for node_idx in node_iter:
+                    node_iter.send({"conv_id": conv_id, "turn_idx": node_idx})
+
+                yield request_chain
 
                 if stop_queueing:
                     self.stop_send_requests_event.set()

src/guidellm/schemas/__init__.py

@@ -23,6 +23,7 @@
     TotalT,
     standard_model_config,
 )
+from .dag import DAGNode, ForkNode, JoinNode, RequestNode, SpawnNode, StartNode
 from .info import RequestInfo, RequestSettings, RequestTimings
 from .request import (
     GenerationRequest,
@@ -42,23 +43,29 @@
 
 __all__ = [
     "BaseModelT",
+    "DAGNode",
     "DistributionSummary",
     "ErroredT",
+    "ForkNode",
     "FunctionObjT",
     "GenerationRequest",
     "GenerationRequestArguments",
     "GenerationResponse",
     "GenerativeRequestStats",
     "IncompleteT",
+    "JoinNode",
     "Percentiles",
     "PydanticClassRegistryMixin",
     "RegisterClassT",
     "ReloadableBaseModel",
     "RequestInfo",
+    "RequestNode",
     "RequestSettings",
     "RequestTimings",
+    "SpawnNode",
     "StandardBaseDict",
     "StandardBaseModel",
+    "StartNode",
     "StatusBreakdown",
     "StatusDistributionSummary",
     "SuccessfulT",

src/guidellm/schemas/dag.py

@@ -0,0 +1,120 @@
+from __future__ import annotations
+
+from collections.abc import Generator
+from typing import Generic, Protocol, TypeVar
+
+from pydantic import Field
+
+from guidellm.schemas.base import StandardBaseModel
+from guidellm.schemas.info import RequestInfo, RequestSettings
+
+RequestT = TypeVar("RequestT")
+
+
+class DAGMutator(Protocol[RequestT]):
+    def __call__(self, node: DAGNode[RequestT], **kwargs) -> None:
+        """
+        A callable that mutates a DAGNode.
+        """
+        ...
+
+
+class DAGNode(StandardBaseModel, Generic[RequestT]):
+    """
+    A node in a directed acyclic graph (DAG).
+    """
+
+    next: list[DAGNode[RequestT]] = Field(
+        default_factory=list, description="The next nodes in the DAG."
+    )
+
+    def _apply_to_all(
+        self, func: DAGMutator[RequestT], idx: int
+    ) -> Generator[int, dict, None]:
+        sent = yield idx
+        func(self, **sent)
+        for child in self.next:
+            child._apply_to_all(func, idx + 1)  # noqa: SLF001
+
+    def apply_to_all(self, func: DAGMutator[RequestT]) -> Generator[int, dict, None]:
+        """
+        Apply a function to this node and all its descendants.
+        """
+        yield from self._apply_to_all(func, 0)
+
+
+class StartNode(DAGNode[RequestT], Generic[RequestT]):
+    """
+    A node that represents the start of the DAG.
+    """
+
+    next: list[DAGNode[RequestT]] = Field(
+        default_factory=list,
+        min_length=1,
+        max_length=1,
+        description="The next node in the DAG.",
+    )
+
+
+class ForkNode(DAGNode[RequestT], Generic[RequestT]):
+    """
+    A node that represents a fork in the DAG.
+    """
+
+    next: list[DAGNode[RequestT]] = Field(
+        default_factory=list, min_length=1, description="The next nodes in the DAG."
+    )
+
+
+class SpawnNode(DAGNode[RequestT], Generic[RequestT]):
+    """
+    A node that represents a spawn point in the DAG.
+    """
+
+    next: list[DAGNode[RequestT]] = Field(
+        default_factory=list, min_length=1, description="The next nodes in the DAG."
+    )
+
+
+class JoinNode(DAGNode[RequestT], Generic[RequestT]):
+    """
+    A node that represents a join point in the DAG.
+    """
+
+    next: list[DAGNode[RequestT]] = Field(
+        default_factory=list,
+        min_length=1,
+        max_length=1,
+        description="The next nodes in the DAG.",
+    )
+
+
+class EndNode(DAGNode[RequestT], Generic[RequestT]):
+    """
+    A node that represents the end of the DAG.
+    """
+
+    next: list[DAGNode[RequestT]] = Field(
+        default_factory=list, max_length=0, description="The next nodes in the DAG."
+    )
+
+
+class RequestNode(DAGNode[RequestT], Generic[RequestT]):
+    """
+    A node that represents a request in the DAG.
+    """
+
+    next: list[DAGNode[RequestT]] = Field(
+        default_factory=list,
+        min_length=1,
+        max_length=1,
+        description="The next nodes in the DAG.",
+    )
+    request: RequestT = Field(..., description="The request associated with this node.")
+    settings: RequestSettings = Field(
+        default_factory=RequestSettings,
+        description="The settings associated with this node.",
+    )
+    info: RequestInfo | None = Field(
+        default_factory=None, description="The info associated with this node."
+    )