Convert new types to TypeAliasTypes

Author: sjmonson

src/guidellm/benchmark/__init__.py

@@ -46,9 +46,17 @@
     enable_scenarios,
     get_builtin_scenarios,
 )
+from .types import (
+    AggregatorInputT,
+    DataInputT,
+    OutputFormatT,
+    ProcessorInputT,
+    ProgressInputT,
+)
 
 __all__ = [
     "Aggregator",
+    "AggregatorInputT",
     "AggregatorState",
     "AsyncProfile",
     "Benchmark",
@@ -60,6 +68,7 @@
     "BenchmarkerProgressGroup",
     "CompilableAggregator",
     "ConcurrentProfile",
+    "DataInputT",
     "GenerativeBenchmark",
     "GenerativeBenchmarkerCSV",
     "GenerativeBenchmarkerConsole",
@@ -73,8 +82,11 @@
     "GenerativeStatsProgressAggregator",
     "GenerativeTextScenario",
     "InjectExtrasAggregator",
+    "OutputFormatT",
+    "ProcessorInputT",
     "Profile",
     "ProfileType",
+    "ProgressInputT",
     "Scenario",
     "SchedulerStatsAggregator",
     "SerializableAggregator",

src/guidellm/benchmark/entrypoints.py

@@ -24,11 +24,11 @@
 from guidellm.benchmark.progress import BenchmarkerProgressGroup
 from guidellm.benchmark.scenario import enable_scenarios
 from guidellm.benchmark.types import (
-    OutputFormatType,
-    DataInputType,
-    ProcessorInputType,
-    ProgressInputType,
-    AggregatorInputType
+    AggregatorInputT,
+    DataInputT,
+    OutputFormatT,
+    ProcessorInputT,
+    ProgressInputT,
 )
 from guidellm.request import GenerativeRequestLoader
 from guidellm.scheduler import (
@@ -49,16 +49,15 @@
 
 # Helper functions
 
+
 async def initialize_backend(
     backend: BackendType | Backend,
     target: str,
     model: str | None,
     backend_kwargs: dict[str, Any] | None,
 ) -> Backend:
     backend = (
-        Backend.create(
-            backend, target=target, model=model, **(backend_kwargs or {})
-        )
+        Backend.create(backend, target=target, model=model, **(backend_kwargs or {}))
         if not isinstance(backend, Backend)
         else backend
     )
@@ -95,18 +94,19 @@ async def resolve_profile(
         )
     return profile
 
+
 async def resolve_output_formats(
-    output_formats: OutputFormatType,
+    output_formats: OutputFormatT,
     output_path: str | Path | None,
 ) -> dict[str, GenerativeBenchmarkerOutput]:
-    output_formats = GenerativeBenchmarkerOutput.resolve(
+    return GenerativeBenchmarkerOutput.resolve(
         output_formats=(output_formats or {}), output_path=output_path
     )
-    return output_formats
+
 
 async def finalize_outputs(
     report: GenerativeBenchmarksReport,
-    resolved_output_formats: dict[str, GenerativeBenchmarkerOutput]
+    resolved_output_formats: dict[str, GenerativeBenchmarkerOutput],
 ):
     output_format_results = {}
     for key, output in resolved_output_formats.items():
@@ -122,7 +122,7 @@ async def finalize_outputs(
 @enable_scenarios
 async def benchmark_generative_text(  # noqa: C901
     target: str,
-    data: DataInputType,
+    data: DataInputT,
     profile: StrategyType | ProfileType | Profile,
     rate: list[float] | None = None,
     random_seed: int = 42,
@@ -131,18 +131,18 @@ async def benchmark_generative_text(  # noqa: C901
     backend_kwargs: dict[str, Any] | None = None,
     model: str | None = None,
     # Data configuration
-    processor: ProcessorInputType | None = None,
+    processor: ProcessorInputT | None = None,
     processor_args: dict[str, Any] | None = None,
     data_args: dict[str, Any] | None = None,
     data_sampler: Literal["random"] | None = None,
     # Output configuration
     output_path: str | Path | None = _CURRENT_WORKING_DIR,
-    output_formats: OutputFormatType = ("console", "json", "html", "csv"),
+    output_formats: OutputFormatT = ("console", "json", "html", "csv"),
     # Updates configuration
-    progress: ProgressInputType | None = None,
+    progress: ProgressInputT | None = None,
     print_updates: bool = False,
     # Aggregators configuration
-    add_aggregators: AggregatorInputType | None = None,
+    add_aggregators: AggregatorInputT | None = None,
     warmup: float | None = None,
     cooldown: float | None = None,
     request_samples: int | None = 20,
@@ -259,7 +259,9 @@ async def benchmark_generative_text(  # noqa: C901
         )
 
     with console.print_update_step(title="Resolving output formats") as console_step:
-        resolved_output_formats = await resolve_output_formats(output_formats, output_path)
+        resolved_output_formats = await resolve_output_formats(
+            output_formats, output_path
+        )
         console_step.finish(
             title="Output formats resolved",
             details={key: str(val) for key, val in resolved_output_formats.items()},
@@ -314,7 +316,7 @@ async def benchmark_generative_text(  # noqa: C901
 async def reimport_benchmarks_report(
     file: Path,
     output_path: Path | None,
-    output_formats: OutputFormatType = ("console", "json", "html", "csv"),
+    output_formats: OutputFormatT = ("console", "json", "html", "csv"),
 ) -> tuple[GenerativeBenchmarksReport, dict[str, Any]]:
     """
     The command-line entry point for re-importing and displaying an
@@ -326,10 +328,15 @@ async def reimport_benchmarks_report(
         title=f"Loading benchmarks from {file}"
     ) as console_step:
         report = GenerativeBenchmarksReport.load_file(file)
-        console_step.finish(f"Import of old benchmarks complete; loaded {len(report.benchmarks)} benchmark(s)")
+        console_step.finish(
+            "Import of old benchmarks complete;"
+            f" loaded {len(report.benchmarks)} benchmark(s)"
+        )
 
     with console.print_update_step(title="Resolving output formats") as console_step:
-        resolved_output_formats = await resolve_output_formats(output_formats, output_path)
+        resolved_output_formats = await resolve_output_formats(
+            output_formats, output_path
+        )
         console_step.finish(
             title="Output formats resolved",
             details={key: str(val) for key, val in resolved_output_formats.items()},

src/guidellm/benchmark/scenario.py

@@ -9,21 +9,18 @@
 import yaml
 from loguru import logger
 from pydantic import BeforeValidator, Field, PositiveFloat, PositiveInt, SkipValidation
-from transformers.tokenization_utils_base import (  # type: ignore[import]
-    PreTrainedTokenizerBase,
-)
 
 from guidellm.backends import Backend, BackendType
 from guidellm.benchmark.profile import Profile, ProfileType
-from guidellm.benchmark.types import DataInputType, ProcessorInputType, AggregatorInputType
+from guidellm.benchmark.types import AggregatorInputT, DataInputT, ProcessorInputT
 from guidellm.scheduler import StrategyType
 from guidellm.utils import StandardBaseModel
 
-__ALL__ = [
-    "Scenario",
+__all__ = [
     "GenerativeTextScenario",
-    "get_builtin_scenarios",
+    "Scenario",
     "enable_scenarios",
+    "get_builtin_scenarios",
 ]
 
 SCENARIO_DIR = Path(__file__).parent / "scenarios/"
@@ -111,7 +108,7 @@ class Config:
         arbitrary_types_allowed = True
 
     data: Annotated[
-        DataInputType,
+        DataInputT,
         # BUG: See https://github.com/pydantic/pydantic/issues/9541
         SkipValidation,
     ]
@@ -125,12 +122,12 @@ class Config:
     backend_kwargs: dict[str, Any] | None = None
     model: str | None = None
     # Data configuration
-    processor: ProcessorInputType | None = None
+    processor: ProcessorInputT | None = None
     processor_args: dict[str, Any] | None = None
     data_args: dict[str, Any] | None = None
     data_sampler: Literal["random"] | None = None
     # Aggregators configuration
-    add_aggregators: AggregatorInputType | None = None
+    add_aggregators: AggregatorInputT | None = None
     warmup: Annotated[float | None, Field(gt=0, le=1)] = None
     cooldown: Annotated[float | None, Field(gt=0, le=1)] = None
     request_samples: PositiveInt | None = 20

src/guidellm/benchmark/types.py

@@ -1,45 +1,59 @@
 from __future__ import annotations
+
 from collections.abc import Iterable
-from typing import Any
 from pathlib import Path
-from datasets import Dataset, DatasetDict, IterableDataset, IterableDatasetDict
-
-from guidellm.benchmark.output import (
-    GenerativeBenchmarkerOutput,
-)
+from typing import Any, TypeAliasType
 
+from datasets import Dataset, DatasetDict, IterableDataset, IterableDatasetDict
 from transformers import (  # type: ignore[import]
     PreTrainedTokenizerBase,
 )
 
-from guidellm.benchmark.progress import BenchmarkerProgress
-
 from guidellm.benchmark.aggregator import (
     Aggregator,
     CompilableAggregator,
 )
+from guidellm.benchmark.output import (
+    GenerativeBenchmarkerOutput,
+)
+from guidellm.benchmark.progress import BenchmarkerProgress
+
+__all__ = [
+    "AggregatorInputT",
+    "DataInputT",
+    "OutputFormatT",
+    "ProcessorInputT",
+    "ProgressInputT",
+]
 
 
-DataInputType = (
+DataInputT = TypeAliasType(
+    "DataInputT",
     Iterable[str]
     | Iterable[dict[str, Any]]
     | Dataset
     | DatasetDict
     | IterableDataset
     | IterableDatasetDict
     | str
-    | Path
+    | Path,
 )
 
-OutputFormatType = (
+OutputFormatT = TypeAliasType(
+    "OutputFormatT",
     tuple[str, ...]
     | list[str]
     | dict[str, str | dict[str, Any] | GenerativeBenchmarkerOutput]
-    | None
+    | None,
 )
 
-ProcessorInputType = str | Path | PreTrainedTokenizerBase
+ProcessorInputT = TypeAliasType("ProcessorInputT", str | Path | PreTrainedTokenizerBase)
 
-ProgressInputType = tuple[str, ...] | list[str] | list[BenchmarkerProgress]
+ProgressInputT = TypeAliasType(
+    "ProgressInputT", tuple[str, ...] | list[str] | list[BenchmarkerProgress]
+)
 
-AggregatorInputType = dict[str, str | dict[str, Any] | Aggregator | CompilableAggregator]
+AggregatorInputT = TypeAliasType(
+    "AggregatorInputT",
+    dict[str, str | dict[str, Any] | Aggregator | CompilableAggregator],
+)