guidellm 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

guidellm/request/__init__.py

@@ -1,13 +1,15 @@
-from .base import GenerationMode, RequestGenerator
-from .emulated import EmulatedConfig, EmulatedRequestGenerator
-from .file import FileRequestGenerator
-from .transformers import TransformersDatasetRequestGenerator
+from .loader import (
+    GenerativeRequestLoader,
+    GenerativeRequestLoaderDescription,
+    RequestLoader,
+    RequestLoaderDescription,
+)
+from .request import GenerationRequest
 
 __all__ = [
-    "EmulatedConfig",
-    "EmulatedRequestGenerator",
-    "FileRequestGenerator",
-    "GenerationMode",
-    "RequestGenerator",
-    "TransformersDatasetRequestGenerator",
+    "RequestLoader",
+    "RequestLoaderDescription",
+    "GenerativeRequestLoaderDescription",
+    "GenerativeRequestLoader",
+    "GenerationRequest",
 ]

guidellm/request/loader.py

@@ -0,0 +1,281 @@
+from abc import abstractmethod
+from collections.abc import Iterable, Iterator
+from pathlib import Path
+from typing import (
+    Any,
+    Literal,
+    Optional,
+    Union,
+)
+
+from datasets import Dataset, DatasetDict, IterableDataset, IterableDatasetDict
+from transformers import PreTrainedTokenizerBase  # type: ignore[import]
+
+from guidellm.dataset import ColumnInputTypes, load_dataset
+from guidellm.objects import StandardBaseModel
+from guidellm.request.request import GenerationRequest
+
+__all__ = [
+    "RequestLoaderDescription",
+    "RequestLoader",
+    "GenerativeRequestLoaderDescription",
+    "GenerativeRequestLoader",
+]
+
+
+class RequestLoaderDescription(StandardBaseModel):
+    type_: Literal["request_loader"] = "request_loader"
+
+
+class RequestLoader(Iterable):
+    @abstractmethod
+    def __iter__(self): ...
+
+    @abstractmethod
+    def __len__(self): ...
+
+    @property
+    @abstractmethod
+    def description(self) -> RequestLoaderDescription: ...
+
+
+class GenerativeRequestLoaderDescription(RequestLoaderDescription):
+    type_: Literal["generative_request_loader"] = "generative_request_loader"  # type: ignore[assignment]
+    data: str
+    data_args: Optional[dict[str, Any]]
+    processor: str
+    processor_args: Optional[dict[str, Any]]
+
+
+class GenerativeRequestLoader(RequestLoader):
+    DEFAULT_PROMPT_COLUMNS = [
+        "prompt",
+        "prompts",
+        "instruction",
+        "instructions",
+        "question",
+        "questions",
+        "input",
+        "inputs",
+        "context",
+        "content",
+        "conversation",
+        "conversations",
+        "text",
+    ]
+
+    def __init__(
+        self,
+        data: Union[
+            str,
+            Path,
+            Iterable[Union[str, dict[str, Any]]],
+            Dataset,
+            DatasetDict,
+            IterableDataset,
+            IterableDatasetDict,
+        ],
+        data_args: Optional[dict[str, Any]],
+        processor: Optional[Union[str, Path, PreTrainedTokenizerBase]],
+        processor_args: Optional[dict[str, Any]],
+        shuffle: bool = True,
+        iter_type: Literal["finite", "infinite"] = "finite",
+        random_seed: int = 42,
+    ):
+        self.data = data
+        self.data_args = data_args
+        dataset, args_column_mappings = load_dataset(
+            data,
+            data_args,
+            processor,
+            processor_args,
+            random_seed,
+        )
+        self.dataset = dataset
+        self.processor = processor
+        self.processor_args = processor_args
+        self.shuffle = shuffle
+        self.iter_type = iter_type
+        self.random_seed = random_seed
+
+        self.column_mappings = self._create_column_mappings(args_column_mappings)
+        self.preserve_iter_state = iter_type == "infinite"  # ensure no caching requests
+        self._preserved_iter = None
+
+    def __iter__(self) -> Iterator[GenerationRequest]:
+        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)
+
+            self._preserved_iter = None
+
+    def __len__(self) -> int:
+        if self.iter_type == "finite":
+            return self.num_unique_items()
+
+        raise ValueError(f"Unable to determine length of dataset: {self.data}")
+
+    @property
+    def description(self) -> GenerativeRequestLoaderDescription:
+        return GenerativeRequestLoaderDescription(
+            data=str(self.data),
+            data_args=self.data_args,
+            processor=str(self.processor),
+            processor_args=self.processor_args,
+        )
+
+    def num_unique_items(self, raise_err: bool = True) -> int:
+        try:
+            return len(self.dataset)
+        except Exception:  # noqa: BLE001, S110
+            pass
+
+        dataset_size = self.dataset.info.dataset_size
+        if dataset_size is not None:
+            return dataset_size
+
+        if raise_err:
+            raise ValueError("Unable to determine number of items in the dataset")
+
+        return -1
+
+    def _create_column_mappings(
+        self,
+        args_column_mappings: dict[ColumnInputTypes, str],
+    ) -> dict[ColumnInputTypes, str]:
+        column_mappings: dict[ColumnInputTypes, str] = {}
+
+        if "text_column" in args_column_mappings:
+            column_mappings["prompt_column"] = args_column_mappings["text_column"]
+        else:
+            column_mappings["prompt_column"] = self._extract_text_column()
+
+        if "prompt_tokens_count_column" in args_column_mappings:
+            column_mappings["prompt_tokens_count_column"] = args_column_mappings[
+                "prompt_tokens_count_column"
+            ]
+        elif prompt_tokens_count_column := self._extract_prompt_tokens_count_column():
+            column_mappings["prompt_tokens_count_column"] = prompt_tokens_count_column
+
+        if "output_tokens_count_column" in args_column_mappings:
+            column_mappings["output_tokens_count_column"] = args_column_mappings[
+                "output_tokens_count_column"
+            ]
+        elif output_tokens_count_column := self._extract_output_tokens_count_column():
+            column_mappings["output_tokens_count_column"] = output_tokens_count_column
+
+        return column_mappings
+
+    def _extract_text_column(self) -> str:
+        column_names = self._dataset_columns(
+            err_msg=(
+                "Unable to determine text column from dataset and it is required. "
+                "To specify the text column, set the 'text_column' key in the "
+                "'data_args' dictionary."
+            )
+        )
+
+        if not column_names:
+            raise ValueError(
+                "Unable to determine text column from dataset and it is required. "
+                "To specify the text column, set the 'text_column' key in the "
+                "'data_args' dictionary."
+            )
+
+        if len(column_names) == 1:
+            return column_names[0]
+
+        for def_column in self.DEFAULT_PROMPT_COLUMNS:
+            if def_column in column_names:
+                return def_column
+
+        raise ValueError(
+            f"Unable to determine text column from dataset columns: {column_names}. "
+            "To specify the text column, set the 'text_column' key in the "
+            "'data_args' dictionary."
+        )
+
+    def _extract_prompt_tokens_count_column(self) -> Optional[str]:
+        column_names = self._dataset_columns()
+
+        if column_names and "prompt_tokens_count" in column_names:
+            return "prompt_tokens_count"
+
+        if column_names and "prompt_tokens" in column_names:
+            return "prompt_tokens"
+
+        return None
+
+    def _extract_output_tokens_count_column(self) -> Optional[str]:
+        column_names = self._dataset_columns()
+
+        if column_names and "output_tokens_count" in column_names:
+            return "output_tokens_count"
+
+        if column_names and "output_tokens" in column_names:
+            return "output_tokens"
+
+        return None
+
+    def _dataset_columns(self, err_msg: Optional[str] = None) -> Optional[list[str]]:
+        try:
+            column_names = self.dataset.column_names
+
+            if not column_names and err_msg:
+                raise ValueError(f"No column names found in dataset: {self.data}")
+        except Exception as err:
+            if err_msg:
+                raise ValueError(err_msg) from err
+
+            column_names = None
+
+        return column_names
+
+    def _get_dataset_iter(
+        self, scope_create_count: int
+    ) -> Optional[Iterator[dict[str, Any]]]:
+        if scope_create_count > 0 and self.iter_type != "infinite":
+            return None
+
+        if self.preserve_iter_state and self._preserved_iter is not None:
+            return self._preserved_iter
+
+        dataset = (
+            self.dataset
+            if not self.shuffle
+            else self.dataset.shuffle(seed=self.random_seed)
+        )
+
+        dataset_iter = iter(dataset)
+
+        if self.preserve_iter_state:
+            self._preserved_iter = dataset_iter
+
+        return dataset_iter
+
+    def _create_request(self, item: dict[str, Any]) -> GenerationRequest:
+        prompt_tokens = (
+            item[self.column_mappings["prompt_tokens_count_column"]]
+            if "prompt_tokens_count_column" in self.column_mappings
+            else None
+        )
+        output_tokens = (
+            item[self.column_mappings["output_tokens_count_column"]]
+            if "output_tokens_count_column" in self.column_mappings
+            else None
+        )
+
+        return GenerationRequest(
+            request_type="text_completions",
+            content=item[self.column_mappings["prompt_column"]],
+            stats=(
+                {"prompt_tokens": prompt_tokens} if prompt_tokens is not None else {}
+            ),
+            constraints=(
+                {"output_tokens": output_tokens} if output_tokens is not None else {}
+            ),
+        )

guidellm/request/request.py

@@ -0,0 +1,79 @@
+import uuid
+from typing import Any, Literal, Optional
+
+from pydantic import Field
+
+from guidellm.objects.pydantic import StandardBaseModel
+
+__all__ = ["GenerationRequest"]
+
+
+class GenerationRequest(StandardBaseModel):
+    """
+    A class representing a request for generation.
+    This class is used to encapsulate the details of a generation request,
+    including the request ID, type, content, parameters, statistics, and constraints.
+    It is designed to be used with the BackendRequestsWorker class to handle
+    the generation process.
+
+    :param request_id: The unique identifier for the request.
+    :param request_type: The type of request (e.g., text, chat).
+    :param content: The content for the request to send to the backend.
+        If request_type is 'text', this should be a string or list of strings
+        which will be resolved by backend.text_completions.
+        If request_type is 'chat', this should be a string,
+        a list of (str, Dict[str, Union[str, Dict[str, str]], Path, Image]),
+        or Any raw content which will be resolved by backend.chat_completions.
+        If raw content, raw_content=True must be passed in the params.
+    :param params: Additional parameters for the request passed in as kwargs.
+        For an http backend, these are passed into the body of the request.
+    :param stats: Statistics for the request, such as the number of prompt tokens.
+        Used for tracking and reporting purposes.
+    :param constraints: Constraints for the request, such as the maximum number
+        of output tokens. Used for controlling the behavior of the backend.
+    """
+
+    request_id: Optional[str] = Field(
+        default_factory=lambda: str(uuid.uuid4()),
+        description="The unique identifier for the request.",
+    )
+    request_type: Literal["text_completions", "chat_completions"] = Field(
+        default="text_completions",
+        description=(
+            "The type of request (e.g., text, chat). "
+            "If request_type='text_completions', resolved by backend.text_completions. "
+            "If request_typ='chat_completions', resolved by backend.chat_completions."
+        ),
+    )
+    content: Any = Field(
+        description=(
+            "The content for the request to send to the backend. "
+            "If request_type is 'text', this should be a string or list of strings "
+            "which will be resolved by backend.text_completions. "
+            "If request_type is 'chat', this should be a string, "
+            "a list of (str, Dict[str, Union[str, Dict[str, str]], Path, Image]), "
+            "or Any raw content which will be resolved by backend.chat_completions. "
+            "If raw content, raw_content=True must be passed in the params."
+        )
+    )
+    params: dict[str, Any] = Field(
+        default_factory=dict,
+        description=(
+            "Additional parameters for the request that will be passed in as kwargs. "
+            "For an http backend, these are passed into the body of the request. "
+        ),
+    )
+    stats: dict[Literal["prompt_tokens"], int] = Field(
+        default_factory=dict,
+        description=(
+            "Statistics for the request, such as the number of prompt tokens. "
+            "Used for tracking and reporting purposes."
+        ),
+    )
+    constraints: dict[Literal["output_tokens"], int] = Field(
+        default_factory=dict,
+        description=(
+            "Constraints for the request, such as the maximum number of output tokens. "
+            "Used for controlling the behavior of the backend."
+        ),
+    )

guidellm/scheduler/__init__.py

@@ -1,4 +1,52 @@
-from .base import Scheduler, SchedulerResult
-from .load_generator import LoadGenerationMode, LoadGenerator
+from .result import (
+    SchedulerRequestInfo,
+    SchedulerRequestResult,
+    SchedulerResult,
+    SchedulerRunInfo,
+)
+from .scheduler import Scheduler
+from .strategy import (
+    AsyncConstantStrategy,
+    AsyncPoissonStrategy,
+    ConcurrentStrategy,
+    SchedulingStrategy,
+    StrategyType,
+    SynchronousStrategy,
+    ThroughputStrategy,
+    strategy_display_str,
+)
+from .types import RequestT, ResponseT
+from .worker import (
+    GenerativeRequestsWorker,
+    GenerativeRequestsWorkerDescription,
+    RequestsWorker,
+    ResolveStatus,
+    WorkerDescription,
+    WorkerProcessRequest,
+    WorkerProcessResult,
+)
 
-__all__ = ["LoadGenerationMode", "LoadGenerator", "Scheduler", "SchedulerResult"]
+__all__ = [
+    "SchedulerRequestInfo",
+    "SchedulerRequestResult",
+    "SchedulerResult",
+    "SchedulerRunInfo",
+    "Scheduler",
+    "AsyncConstantStrategy",
+    "AsyncPoissonStrategy",
+    "ConcurrentStrategy",
+    "SchedulingStrategy",
+    "StrategyType",
+    "SynchronousStrategy",
+    "ThroughputStrategy",
+    "strategy_display_str",
+    "RequestT",
+    "ResponseT",
+    "WorkerProcessRequest",
+    "WorkerProcessResult",
+    "ResolveStatus",
+    "WorkerDescription",
+    "RequestsWorker",
+    "GenerativeRequestsWorkerDescription",
+    "GenerativeRequestsWorker",
+]

guidellm/scheduler/result.py

@@ -0,0 +1,137 @@
+from typing import (
+    Generic,
+    Literal,
+    Optional,
+)
+
+from guidellm.objects import StandardBaseModel
+from guidellm.scheduler.strategy import SchedulingStrategy
+from guidellm.scheduler.types import RequestT, ResponseT
+
+__all__ = [
+    "SchedulerResult",
+    "SchedulerRequestResult",
+    "SchedulerRunInfo",
+    "SchedulerRequestInfo",
+]
+
+
+class SchedulerRunInfo(StandardBaseModel):
+    """
+    Information about the current run of the scheduler.
+    This class holds metadata about the scheduling run,
+    including the start and end times, the number of processes,
+    and the scheduling strategy used.
+    It also tracks the number of requests created, queued, pending,
+    and completed during the run.
+
+    :param start_time: The start time of the scheduling run.
+    :param end_time: The end time of the scheduling run;
+        if None, then this will be math.inf.
+    :param end_number: The maximum number of requests to be processed;
+        if None, then this will be math.inf.
+    :param processes: The number of processes used in the scheduling run.
+    :param strategy: The scheduling strategy used in the run.
+        This should be an instance of SchedulingStrategy.
+    :param created_requests: The number of requests created during the run.
+    :param queued_requests: The number of requests queued during the run.
+    :param scheduled_requests: The number of requests scheduled during the run.
+        (requests pending being sent to the worker but recieved by a process)
+    :param processing_requests: The number of requests actively being run.
+    :param completed_requests: The number of requests completed during the run.
+    """
+
+    start_time: float
+    end_time: float
+    end_number: float
+    processes: int
+    strategy: SchedulingStrategy
+
+    created_requests: int = 0
+    queued_requests: int = 0
+    scheduled_requests: int = 0
+    processing_requests: int = 0
+    completed_requests: int = 0
+
+
+class SchedulerRequestInfo(StandardBaseModel):
+    """
+    Information about a specific request run through the scheduler.
+    This class holds metadata about the request, including
+    the targeted start time, queued time, start time, end time,
+    and the process ID that handled the request.
+
+    :param targeted_start_time: The targeted start time for the request (time.time()).
+    :param queued_time: The time the request was queued (time.time()).
+    :param scheduled_time: The time the request was scheduled (time.time())
+        (any sleep time before the request was sent to the worker).
+    :param worker_start: The time the worker started processing request (time.time()).
+    :param worker_end: The time the worker finished processing request. (time.time()).
+    :param process_id: The ID of the underlying process that handled the request.
+    """
+
+    requested: bool = False
+    completed: bool = False
+    errored: bool = False
+    canceled: bool = False
+
+    targeted_start_time: float = -1
+    queued_time: float = -1
+    dequeued_time: float = -1
+    scheduled_time: float = -1
+    worker_start: float = -1
+    request_start: float = -1
+    request_end: float = -1
+    worker_end: float = -1
+    process_id: int = -1
+
+
+class SchedulerResult(StandardBaseModel):
+    """
+    The yielded, iterative result for a scheduler run.
+    These are triggered on the start and end of the run,
+    as well as on the start and end of each request.
+    Depending on the type, it will hold the request and response
+    along with information and statistics about the request and general run.
+
+    :param type_: The type of the result, which can be one of:
+        - "run_start": Indicates the start of the run.
+        - "run_complete": Indicates the completion of the run (teardown happens after).
+        - "request_start": Indicates the start of a request.
+        - "request_complete": Indicates the completion of a request.
+    :param request: The request that was processed.
+    :param response: The response from the worker for the request.
+    :param request_info: Information about the request, including
+        the targeted start time, queued time, start time, end time,
+        and the process ID that handled the request.
+    :param run_info: Information about the current run of the scheduler,
+        including the start and end times, the number of processes,
+        and the scheduling strategy used.
+        It also tracks the number of requests created, queued, pending,
+        and completed during the run.
+    """
+
+    pydantic_type: Literal["scheduler_result"] = "scheduler_result"
+    type_: Literal[
+        "run_start",
+        "run_complete",
+        "request_scheduled",
+        "request_start",
+        "request_complete",
+    ]
+    run_info: SchedulerRunInfo
+
+
+class SchedulerRequestResult(
+    SchedulerResult,
+    Generic[RequestT, ResponseT],
+):
+    pydantic_type: Literal["scheduler_request_result"] = "scheduler_request_result"  # type: ignore[assignment]
+    type_: Literal[
+        "request_scheduled",
+        "request_start",
+        "request_complete",
+    ]
+    request: RequestT
+    request_info: SchedulerRequestInfo
+    response: Optional[ResponseT] = None