Rework entrypoint to match new preprocess and finalizer flow

Signed-off-by: Samuel Monson <[email protected]>

Author: sjmonson

src/guidellm/__main__.py

@@ -178,19 +178,13 @@ def benchmark():
 # Data configuration
 @click.option(
     "--request-type",
-    default=BenchmarkGenerativeTextArgs.get_default("data_request_formatter"),
+    default=BenchmarkGenerativeTextArgs.get_default("request_type"),
     type=click.Choice(list(get_literal_vals(GenerativeRequestType))),
     help=(
         f"Request type to create for each data sample. "
         f"Options: {', '.join(get_literal_vals(GenerativeRequestType))}."
     ),
 )
-@click.option(
-    "--request-formatter-kwargs",
-    default=None,
-    callback=cli_tools.parse_json,
-    help="JSON string of arguments to pass to the request formatter.",
-)
 @click.option(
     "--processor",
     default=BenchmarkGenerativeTextArgs.get_default("processor"),
@@ -223,10 +217,17 @@ def benchmark():
     ),
 )
 @click.option(
-    "--data-column-mapper",
-    default=BenchmarkGenerativeTextArgs.get_default("data_column_mapper"),
+    "--data-preprocessors",
+    default=BenchmarkGenerativeTextArgs.get_default("data_preprocessors"),
+    callback=cli_tools.parse_json,
+    multiple=True,
+    help="JSON string of preprocessors to apply to the dataset.",
+)
+@click.option(
+    "--data-finalizer",
+    default=BenchmarkGenerativeTextArgs.get_default("data_finalizer"),
     callback=cli_tools.parse_json,
-    help="JSON string of column mappings to apply to the dataset.",
+    help="JSON string of finalizer to convert dataset rows to requests.",
 )
 @click.option(
     "--data-sampler",
@@ -386,18 +387,6 @@ def run(**kwargs):
     # Only set CLI args that differ from click defaults
     kwargs = cli_tools.set_if_not_default(click.get_current_context(), **kwargs)
 
-    # Handle remapping for request params
-    request_type = kwargs.pop("request_type", None)
-    request_formatter_kwargs = kwargs.pop("request_formatter_kwargs", None)
-    if request_type is not None:
-        kwargs["data_request_formatter"] = (
-            request_type
-            if not request_formatter_kwargs
-            else {"request_type": request_type, **request_formatter_kwargs}
-        )
-    elif request_formatter_kwargs is not None:
-        kwargs["data_request_formatter"] = request_formatter_kwargs
-
     # Handle output path remapping
     if (output_path := kwargs.pop("output_path", None)) is not None:
         path = Path(output_path)

src/guidellm/benchmark/entrypoints.py

@@ -13,7 +13,7 @@
 
 from collections.abc import Callable, Mapping, MutableMapping
 from pathlib import Path
-from typing import Any, Literal
+from typing import Any, Literal, TypeVar
 
 from torch.utils.data import Sampler
 from transformers import PreTrainedTokenizerBase
@@ -36,21 +36,21 @@
 from guidellm.benchmark.schemas.base import TransientPhaseConfig
 from guidellm.data import (
     DataLoader,
+    DatasetFinalizer,
     DatasetPreprocessor,
+    FinalizerRegistry,
     GenerativeRequestCollator,
     PreprocessorRegistry,
     ProcessorFactory,
-    RequestFormatter,
 )
-from guidellm.data.preprocessors import GenerativeColumnMapper
 from guidellm.scheduler import (
     ConstraintInitializer,
     NonDistributedEnvironment,
     StrategyType,
 )
 from guidellm.schemas import GenerationRequest, GenerationResponse
-from guidellm.settings import settings
 from guidellm.utils import Console, InfoMixin
+from guidellm.utils.registry import RegistryMixin
 
 __all__ = [
     "benchmark_generative_text",
@@ -178,19 +178,66 @@ async def resolve_processor(
     return processor
 
 
+BaseTypeT = TypeVar("BaseTypeT")
+
+
+def resolve_item_from_registry(
+    base_type: type[BaseTypeT],
+    registry: type[RegistryMixin],
+    item: Any,
+    extras: dict[str, Any] | None = None,
+) -> BaseTypeT:
+    """
+    Resolve an item from a registry, instantiating it if necessary.
+
+    :param base_type: The expected base type of the item
+    :param item: The item to resolve, either an instance or a string identifier
+    :param registry: The registry to use for resolving string identifiers
+    :return: The resolved item as an instance of the base type
+    :raises ValueError: If the item cannot be resolved from the registry
+    :raises TypeError: If the resolved item is not of the expected base type
+    """
+    if isinstance(item, base_type):
+        return item
+    else:
+        if isinstance(item, str):
+            item_type = item
+            kwargs = {}
+        else:
+            item_dict = dict(item)
+            item_type = item_dict.pop("type", None)
+            if item_type is None:
+                raise ValueError(
+                    f"Item dictionary must contain a 'type' key to resolve from "
+                    f"{registry.__class__.__name__}."
+                )
+            kwargs = item_dict
+
+        if (item_class := registry.get_registered_object(item_type)) is None:
+            raise ValueError(
+                f"Item type '{item_type}' is not registered in the "
+                f"{registry.__class__.__name__}."
+            )
+        if not issubclass(item_class, base_type):
+            raise TypeError(
+                f"Resolved item type '{item_type}' is not a subclass of "
+                f"{base_type.__name__}."
+            )
+        if extras:
+            kwargs.update(extras)
+        return item_class(**kwargs)
+
+
 async def resolve_request_loader(
     data: list[Any],
     model: str,
+    request_type: str,
     data_args: list[dict[str, Any]] | None,
     data_samples: int,
     processor: ProcessorInputT | None,
     processor_args: dict[str, Any] | None,
-    data_column_mapper: (
-        DatasetPreprocessor
-        | dict[str, str | list[str]]
-        | Literal["generative_column_mapper"]
-    ),
-    data_request_formatter: (RequestFormatter | dict[str, str] | str),
+    data_preprocessors: list[DatasetPreprocessor | dict[str, str | list[str]] | str],
+    data_finalizer: (DatasetFinalizer | dict[str, Any] | str),
     data_collator: Callable | Literal["generative"] | None,
     data_sampler: Sampler[int] | Literal["shuffle"] | None,
     data_num_workers: int | None,
@@ -232,54 +279,22 @@ async def resolve_request_loader(
         else None
     )
 
-    data_column_mapper_instance: DatasetPreprocessor
-    if isinstance(data_column_mapper, DatasetPreprocessor):
-        data_column_mapper_instance = data_column_mapper
-    else:
-        column_mappings = (
-            data_column_mapper if isinstance(data_column_mapper, dict) else None
-        )
-        data_column_mapper_instance = GenerativeColumnMapper(
-            column_mappings=column_mappings  # type: ignore[arg-type]
-        )
-
-    data_request_formatter_instance: RequestFormatter
-    if isinstance(data_request_formatter, RequestFormatter):
-        data_request_formatter_instance = data_request_formatter
-    else:
-        if isinstance(data_request_formatter, str):
-            request_type = data_request_formatter
-            formatter_kwargs: dict[str, Any] = {}
-        else:
-            # Extract request_type from formatter dictionary
-            formatter_dict = dict(data_request_formatter)
-            request_type = formatter_dict.pop("request_type", settings.preferred_route)
-            formatter_kwargs = formatter_dict
-
-        if (
-            formatter_class := PreprocessorRegistry.get_registered_object(request_type)
-        ) is None:
-            raise ValueError(
-                f"Request formatter '{request_type}' is not registered in the "
-                f"PreprocessorRegistry."
-            )
-        if not issubclass(formatter_class, RequestFormatter):
-            raise TypeError(
-                f"Request formatter '{request_type}' is not a subclass of "
-                f"RequestFormatter."
-            )
-
-        data_request_formatter_instance = formatter_class(
-            model=model,
-            **formatter_kwargs,
-        )
-
-    # Cast to proper types for the DataLoader preprocessors list
     preprocessors_list: list[DatasetPreprocessor] = [
-        data_column_mapper_instance,
-        data_request_formatter_instance,
+        resolve_item_from_registry(
+            DatasetPreprocessor,  # type: ignore [type-abstract]
+            PreprocessorRegistry,
+            preprocessor,
+        )
+        for preprocessor in data_preprocessors
     ]
 
+    finalizer_instance = resolve_item_from_registry(
+        DatasetFinalizer,  # type: ignore [type-abstract]
+        FinalizerRegistry,
+        data_finalizer,
+        extras={"request_type": request_type},
+    )
+
     request_loader: DataLoader[GenerationRequest] = DataLoader(
         data=data,
         data_args=data_args,
@@ -289,6 +304,7 @@ async def resolve_request_loader(
             processor_args=processor_args,
         ),
         preprocessors=preprocessors_list,
+        finalizer=finalizer_instance,
         collator=(
             data_collator if callable(data_collator) else GenerativeRequestCollator()
         ),
@@ -460,12 +476,13 @@ async def benchmark_generative_text(
     request_loader = await resolve_request_loader(
         data=args.data,
         model=model,
+        request_type=args.request_type,
         data_args=args.data_args,
         data_samples=args.data_samples,
         processor=processor,
         processor_args=args.processor_args,
-        data_column_mapper=args.data_column_mapper,
-        data_request_formatter=args.data_request_formatter,
+        data_preprocessors=args.data_preprocessors,
+        data_finalizer=args.data_finalizer,
         data_collator=args.data_collator,
         data_sampler=args.data_sampler,
         data_num_workers=args.data_num_workers,

src/guidellm/benchmark/schemas/generative/entrypoints.py

@@ -35,9 +35,10 @@
 from guidellm.benchmark.profiles import Profile, ProfileType
 from guidellm.benchmark.scenarios import get_builtin_scenarios
 from guidellm.benchmark.schemas.base import TransientPhaseConfig
-from guidellm.data import DatasetPreprocessor, RequestFormatter
+from guidellm.data import DatasetFinalizer, DatasetPreprocessor
 from guidellm.scheduler import StrategyType
 from guidellm.schemas import StandardBaseModel
+from guidellm.settings import settings
 
 __all__ = ["BenchmarkGenerativeTextArgs"]
 
@@ -179,6 +180,13 @@ def get_default(cls: type[BenchmarkGenerativeTextArgs], field: str) -> Any:
     backend_kwargs: dict[str, Any] | None = Field(
         default=None, description="Additional backend configuration arguments"
     )
+    request_type: str = Field(
+        default_factory=lambda: settings.preferred_route,
+        description=(
+            "Request type for backend operations;"
+            " shorthand for backend_kwargs['request_type']"
+        ),
+    )
     model: str | None = Field(default=None, description="Model identifier for backend")
     # Data configuration
     processor: str | Path | PreTrainedTokenizerBase | None = Field(
@@ -194,23 +202,21 @@ def get_default(cls: type[BenchmarkGenerativeTextArgs], field: str) -> Any:
     data_samples: int = Field(
         default=-1, description="Number of samples to use from datasets (-1 for all)"
     )
-    data_column_mapper: (
-        DatasetPreprocessor
-        | dict[str, str | list[str]]
-        | Literal["generative_column_mapper"]
-    ) = Field(
-        default="generative_column_mapper",
-        description="Column mapping preprocessor for dataset fields",
+    # TODO: Make it easy to cutomize preprocessors without editing the full list
+    data_preprocessors: list[DatasetPreprocessor | dict[str, str | list[str]] | str] = (
+        Field(
+            default_factory=lambda: [  # type: ignore [arg-type]
+                "generative_column_mapper",
+                "encode_audio",
+                "encode_image",
+                "encode_video",
+            ],
+            description="List of dataset preprocessors to apply in order",
+        )
     )
-    data_request_formatter: RequestFormatter | dict[str, Any] | str = Field(
-        default="chat_completions",
-        description="Request formatting preprocessor or template name",
-        validation_alias=AliasChoices(
-            "data_request_formatter",
-            "data-request-formatter",
-            "request_type",
-            "request-type",
-        ),
+    data_finalizer: DatasetFinalizer | str | dict[str, Any] = Field(
+        default="generative_text_finalizer",
+        description="Finalizer for preparing data samples into requests",
     )
     data_collator: Callable | Literal["generative"] | None = Field(
         default="generative", description="Data collator for batch processing"
@@ -284,7 +290,7 @@ def get_default(cls: type[BenchmarkGenerativeTextArgs], field: str) -> Any:
         default=None, description="Maximum global error rate (0-1) before stopping"
     )
 
-    @field_validator("data", "data_args", "rate", mode="wrap")
+    @field_validator("data", "data_args", "rate", "data_preprocessors", mode="wrap")
     @classmethod
     def single_to_list(
         cls, value: Any, handler: ValidatorFunctionWrapHandler
@@ -323,28 +329,25 @@ def serialize_data_collator(
         """Serialize data_collator to string or None."""
         return data_collator if isinstance(data_collator, str) else None
 
-    @field_serializer("data_column_mapper")
+    @field_serializer("data_preprocessors")
     def serialize_data_column_mapper(
         self,
-        data_column_mapper: (
-            DatasetPreprocessor
-            | dict[str, str | list[str]]
-            | Literal["generative_column_mapper"]
-        ),
-    ) -> dict | str:
+        data_preprocessors: list[
+            DatasetPreprocessor | dict[str, str | list[str]] | str
+        ],
+    ) -> list[dict | str]:
         """Serialize data_column_mapper to dict or string."""
-        return data_column_mapper if isinstance(data_column_mapper, dict | str) else {}
+        return [
+            (preprocessor if isinstance(preprocessor, dict | str) else {})
+            for preprocessor in data_preprocessors
+        ]
 
-    @field_serializer("data_request_formatter")
+    @field_serializer("data_finalizer")
     def serialize_data_request_formatter(
-        self, data_request_formatter: RequestFormatter | dict[str, Any] | str
+        self, data_finalizer: DatasetFinalizer | dict[str, Any] | str
     ) -> dict | str:
         """Serialize data_request_formatter to dict or string."""
-        return (
-            data_request_formatter
-            if isinstance(data_request_formatter, dict | str)
-            else {}
-        )
+        return data_finalizer if isinstance(data_finalizer, dict | str) else {}
 
     @field_serializer("data_sampler")
     def serialize_data_sampler(

src/guidellm/data/loaders.py

@@ -142,6 +142,7 @@ def __init__(
         data_samples: int,
         processor_factory: Callable[[], PreTrainedTokenizerBase],
         preprocessors: list[DatasetPreprocessor | DataDependentPreprocessor],
+        finalizer: DatasetFinalizer[DataT],
         collator: Callable,
         sampler: Sampler[int] | Literal["shuffle"] | None = None,
         num_workers: int | None = 1,
@@ -154,6 +155,7 @@ def __init__(
             data_samples=data_samples,
             processor_factory=processor_factory,
             preprocessors=preprocessors,
+            finalizer=finalizer,
             random_seed=random_seed,
         )
         self._info: dict[str, Any] = {
@@ -163,6 +165,7 @@ def __init__(
             "preprocessors": [
                 preprocessor.__class__.__name__ for preprocessor in preprocessors
             ],
+            "finalizer": finalizer.__class__.__name__,
             "collator": collator.__class__.__name__,
             "sampler": str(sampler),
             "num_workers": num_workers,