Initial state for datasets rework to enable multimodal and more complicated combinations

Author: markurtz

src/guidellm/__init__.py

@@ -7,6 +7,8 @@
 import logging
 import os
 
+from datasets.utils.logging import disable_progress_bar
+
 with (
     open(os.devnull, "w") as devnull,  # noqa: PTH123
     contextlib.redirect_stderr(devnull),
@@ -19,6 +21,7 @@
     os.environ["TOKENIZERS_PARALLELISM"] = "false"  # Silence warnings for tokenizers
     hf_logging.set_verbosity_error()
     logging.getLogger("transformers").setLevel(logging.ERROR)
+    disable_progress_bar()
 
 from .logger import configure_logger, logger
 from .settings import (

src/guidellm/__main__.py

@@ -56,6 +56,11 @@
 from guidellm.benchmark.scenario import (
     GenerativeTextScenario,
 )
+from guidellm.data import (
+    GenerativeDatasetArgs,
+    GenerativeRequestFormatter,
+    GenerativeRequestType,
+)
 from guidellm.mock_server import MockServer, MockServerConfig
 from guidellm.preprocess.dataset import ShortPromptStrategy, process_dataset
 from guidellm.scheduler import StrategyType
@@ -143,6 +148,7 @@ def benchmark():
 @click.option(
     "--data",
     type=str,
+    multiple=True,
     help=(
         "The HuggingFace dataset ID, a path to a HuggingFace dataset, "
         "a path to a data file csv, json, jsonl, or txt, "
@@ -197,9 +203,7 @@ def benchmark():
     default=None,
     help=(
         "A JSON string containing any arguments to pass to the backend as a "
-        "dict with **kwargs. Headers can be removed by setting their value to "
-        "null. For example: "
-        """'{"headers": {"Authorization": null, "Custom-Header": "Custom-Value"}}'"""
+        "dict with **kwargs."
     ),
 )
 @click.option(
@@ -234,19 +238,72 @@ def benchmark():
 @click.option(
     "--data-args",
     default=None,
-    callback=cli_tools.parse_json,
+    callback=(
+        lambda _ctx, _param, value: [
+            GenerativeDatasetArgs.model_validate_json(val)
+            if val
+            else GenerativeDatasetArgs()
+            for val in value
+        ]
+        if value
+        else None
+    ),
     help=(
         "A JSON string containing any arguments to pass to the dataset creation "
         "as a dict with **kwargs."
     ),
 )
+@click.option(
+    "--data-samples",
+    default=-1,
+    type=int,
+    help=(
+        "The number of samples to use from the dataset. If -1 (default), will use all "
+        "samples in the dataset."
+    ),
+)
 @click.option(
     "--data-sampler",
     default=None,
-    type=click.Choice(["random"]),
+    type=click.Choice(["shuffle"]),
+    help="The data sampler type to use.",
+)
+@click.option(
+    "--data-request-type",
+    default="text_completions",
+    type=str,
+    help=(
+        "The type of request to create for each data sample. "
+        f"For example, {list(get_literal_vals(GenerativeRequestType))}."
+    ),
+)
+@click.option(
+    "--data-request-template",
+    default=None,
+    help=(
+        "A Jinja2 template string or path to a Jinja2 template file to use for "
+        "creating requests from the data samples. If not provided, will use a "
+        "default template based on the request type."
+    ),
+)
+@click.option(
+    "--data-request-extras",
+    default=None,
+    callback=cli_tools.parse_json,
+    help=("A JSON string of extra data to include with each data request."),
+)
+@click.option(
+    "--data-request-nonstreaming",
+    is_flag=True,
+    help="Set this flag to disable streaming for the data requests.",
+)
+@click.option(
+    "--dataloader_kwargs",
+    default=None,
+    callback=cli_tools.parse_json,
     help=(
-        "The data sampler type to use. 'random' will add a random shuffle on the data. "
-        "Defaults to None"
+        "A JSON string containing any arguments to pass to the dataloader constructor "
+        "as a dict with **kwargs."
     ),
 )
 # Output configuration
@@ -387,7 +444,13 @@ def run(
     processor,
     processor_args,
     data_args,
+    data_samples,
     data_sampler,
+    data_request_type,
+    data_request_template,
+    data_request_extras,
+    data_request_nonstreaming,
+    dataloader_kwargs,
     # Output configuration
     output_path,
     output_formats,
@@ -420,7 +483,8 @@ def run(
     asyncio.run(
         benchmark_generative_text(
             target=target,
-            data=data,
+            data=list(data),
+            # Benchmark configuration
             profile=profile,
             rate=rate,
             random_seed=random_seed,
@@ -432,7 +496,22 @@ def run(
             processor=processor,
             processor_args=processor_args,
             data_args=data_args,
-            data_sampler=data_sampler,
+            data_samples=data_samples,
+            data_column_mapper=None,  # use default
+            data_request_formatter=GenerativeRequestFormatter(
+                request_type=data_request_type,
+                request_template=data_request_template,
+                request_extras=data_request_extras,
+                request_defaults=(
+                    {}  # disable defaults if non-streaming
+                    if data_request_nonstreaming
+                    else None
+                ),
+            ),
+            data_preprocessors=None,  # no preprocessors through CLI for now
+            dataloader_sampler=data_sampler,
+            dataloader_collate_fn=None,  # use default
+            dataloader_kwargs=dataloader_kwargs,
             # Output configuration
             output_path=output_path,
             output_formats=[

src/guidellm/backends/__init__.py

@@ -13,6 +13,7 @@
     GenerationRequest,
     GenerationRequestTimings,
     GenerationResponse,
+    GenerationTokenStats,
 )
 from .openai import OpenAIHTTPBackend
 
@@ -22,5 +23,6 @@
     "GenerationRequest",
     "GenerationRequestTimings",
     "GenerationResponse",
+    "GenerationTokenStats",
     "OpenAIHTTPBackend",
 ]

src/guidellm/backends/backend.py

@@ -115,5 +115,6 @@ def requests_limit(self) -> int | None:
     async def default_model(self) -> str | None:
         """
         :return: The default model name or identifier for generation requests.
+            None if no default model is available.
         """
         ...

src/guidellm/backends/objects.py

@@ -6,62 +6,51 @@
 implementations.
 """
 
-import uuid
-from typing import Any, Literal, Optional
+from __future__ import annotations
+
+from typing import Literal
 
 from pydantic import Field
 
+from guidellm.data import (
+    GenerationRequest,
+    GenerationRequestArguments,
+    GenerationRequestTimings,
+)
 from guidellm.scheduler import (
-    MeasuredRequestTimings,
     SchedulerMessagingPydanticRegistry,
 )
 from guidellm.utils import StandardBaseModel
 
 __all__ = [
     "GenerationRequest",
+    "GenerationRequestArguments",
     "GenerationRequestTimings",
     "GenerationResponse",
+    "GenerationTokenStats",
 ]
 
 
 @SchedulerMessagingPydanticRegistry.register()
-class GenerationRequest(StandardBaseModel):
-    """Request model for backend generation operations."""
+class GenerationTokenStats(StandardBaseModel):
+    """Token statistics for generation requests and responses."""
 
-    request_id: str = Field(
-        default_factory=lambda: str(uuid.uuid4()),
-        description="Unique identifier for the request.",
-    )
-    request_type: Literal["text_completions", "chat_completions"] = Field(
-        default="text_completions",
-        description=(
-            "Type of request. 'text_completions' uses backend.text_completions(), "
-            "'chat_completions' uses backend.chat_completions()."
-        ),
-    )
-    content: Any = Field(
-        description=(
-            "Request content. For text_completions: string or list of strings. "
-            "For chat_completions: string, list of messages, or raw content "
-            "(set raw_content=True in params)."
-        )
-    )
-    params: dict[str, Any] = Field(
-        default_factory=dict,
-        description=(
-            "Additional parameters passed to backend methods. "
-            "Common: max_tokens, temperature, stream."
-        ),
+    request: int | None = Field(
+        default=None, description="Number of tokens in the original request."
     )
-    stats: dict[Literal["prompt_tokens"], int] = Field(
-        default_factory=dict,
-        description="Request statistics including prompt token count.",
-    )
-    constraints: dict[Literal["output_tokens"], int] = Field(
-        default_factory=dict,
-        description="Request constraints such as maximum output tokens.",
+    response: int | None = Field(
+        default=None, description="Number of tokens in the generated response."
     )
 
+    def value(
+        self, preference: Literal["request", "response"] | None = None
+    ) -> int | None:
+        if preference == "request":
+            return self.request
+        if preference == "response":
+            return self.response
+        return self.response if self.response is not None else self.request
+
 
 @SchedulerMessagingPydanticRegistry.register()
 class GenerationResponse(StandardBaseModel):
@@ -70,87 +59,32 @@ class GenerationResponse(StandardBaseModel):
     request_id: str = Field(
         description="Unique identifier matching the original GenerationRequest."
     )
-    request_args: dict[str, Any] = Field(
+    request_args: GenerationRequestArguments = Field(
         description="Arguments passed to the backend for this request."
     )
-    value: Optional[str] = Field(
+    text: str | None = Field(
         default=None,
-        description="Complete generated text content. None for streaming responses.",
-    )
-    delta: Optional[str] = Field(
-        default=None, description="Incremental text content for streaming responses."
+        description="The generated response text.",
     )
     iterations: int = Field(
         default=0, description="Number of generation iterations completed."
     )
-    request_prompt_tokens: Optional[int] = Field(
-        default=None, description="Token count from the original request prompt."
-    )
-    request_output_tokens: Optional[int] = Field(
-        default=None,
-        description="Expected output token count from the original request.",
-    )
-    response_prompt_tokens: Optional[int] = Field(
-        default=None, description="Actual prompt token count reported by the backend."
+
+    prompt_stats: GenerationTokenStats = Field(
+        default_factory=GenerationTokenStats,
+        description="Token statistics from the prompt.",
     )
-    response_output_tokens: Optional[int] = Field(
-        default=None, description="Actual output token count reported by the backend."
+    output_stats: GenerationTokenStats = Field(
+        default_factory=GenerationTokenStats,
+        description="Token statistics from the generated output.",
     )
 
-    @property
-    def prompt_tokens(self) -> Optional[int]:
-        """
-        :return: The number of prompt tokens used in the request
-            (response_prompt_tokens if available, otherwise request_prompt_tokens).
-        """
-        return self.response_prompt_tokens or self.request_prompt_tokens
-
-    @property
-    def output_tokens(self) -> Optional[int]:
-        """
-        :return: The number of output tokens generated in the response
-            (response_output_tokens if available, otherwise request_output_tokens).
-        """
-        return self.response_output_tokens or self.request_output_tokens
-
-    @property
-    def total_tokens(self) -> Optional[int]:
-        """
-        :return: The total number of tokens used in the request and response.
-            Sum of prompt_tokens and output_tokens.
-        """
-        if self.prompt_tokens is None or self.output_tokens is None:
-            return None
-        return self.prompt_tokens + self.output_tokens
-
-    def preferred_prompt_tokens(
-        self, preferred_source: Literal["request", "response"]
-    ) -> Optional[int]:
-        if preferred_source == "request":
-            return self.request_prompt_tokens or self.response_prompt_tokens
-        else:
-            return self.response_prompt_tokens or self.request_prompt_tokens
-
-    def preferred_output_tokens(
-        self, preferred_source: Literal["request", "response"]
-    ) -> Optional[int]:
-        if preferred_source == "request":
-            return self.request_output_tokens or self.response_output_tokens
-        else:
-            return self.response_output_tokens or self.request_output_tokens
-
-
-@SchedulerMessagingPydanticRegistry.register()
-@MeasuredRequestTimings.register("generation_request_timings")
-class GenerationRequestTimings(MeasuredRequestTimings):
-    """Timing model for tracking generation request lifecycle events."""
+    def total_tokens(
+        self, preference: Literal["request", "response"] | None = None
+    ) -> int | None:
+        prompt_tokens = self.prompt_stats.value(preference=preference)
+        output_tokens = self.output_stats.value(preference=preference)
 
-    timings_type: Literal["generation_request_timings"] = "generation_request_timings"
-    first_iteration: Optional[float] = Field(
-        default=None,
-        description="Unix timestamp when the first generation iteration began.",
-    )
-    last_iteration: Optional[float] = Field(
-        default=None,
-        description="Unix timestamp when the last generation iteration completed.",
-    )
+        if prompt_tokens is None and output_tokens is None:
+            return None
+        return (prompt_tokens or 0) + (output_tokens or 0)