diff --git a/src/guidellm/scheduler/schemas.py b/src/guidellm/scheduler/schemas.py index ae16c763d..446ada4cf 100644 --- a/src/guidellm/scheduler/schemas.py +++ b/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,), ) """ diff --git a/src/guidellm/scheduler/worker_group.py b/src/guidellm/scheduler/worker_group.py index 297bff08a..23b29d504 100644 --- a/src/guidellm/scheduler/worker_group.py +++ b/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,7 +31,6 @@ BackendInterface, ConversationT, DatasetIterT, - RequestDataT, RequestT, ResponseT, SchedulerState, @@ -39,7 +38,7 @@ ) 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() diff --git a/src/guidellm/schemas/__init__.py b/src/guidellm/schemas/__init__.py index 890b2f9da..cf9e3ccbb 100644 --- a/src/guidellm/schemas/__init__.py +++ b/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", diff --git a/src/guidellm/schemas/dag.py b/src/guidellm/schemas/dag.py new file mode 100644 index 000000000..748dd5b7f --- /dev/null +++ b/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." + )