Merge branch 'main' into fix-muliple-rates

Author: markurtz

src/guidellm/benchmark/benchmarker.py

@@ -57,6 +57,7 @@ async def run(
         backend: BackendInterface[RequestT, ResponseT],
         profile: Profile,
         environment: Environment,
+        data: list[Any],
         progress: BenchmarkerProgress[BenchmarkT] | None = None,
         sample_requests: int | None = 20,
         warmup: float | None = None,
@@ -149,6 +150,7 @@ async def run(
                     environment=environment,
                     strategy=strategy,
                     constraints=constraints,
+                    data=data,
                 )
                 if progress:
                     await progress.on_benchmark_complete(benchmark)

src/guidellm/benchmark/entrypoints.py

@@ -436,6 +436,7 @@ async def benchmark_generative_text(
         backend=backend,
         profile=profile,
         environment=NonDistributedEnvironment(),
+        data=args.data,
         progress=progress,
         sample_requests=args.sample_requests,
         warmup=args.warmup,

src/guidellm/benchmark/output.py

@@ -649,6 +649,8 @@ def _get_benchmark_status_metrics_stats(
         status_dist_summary: StatusDistributionSummary = getattr(
             benchmark.metrics, metric
         )
+        if not hasattr(status_dist_summary, status):
+            return [], []
         dist_summary: DistributionSummary = getattr(status_dist_summary, status)
 
         headers = [
@@ -688,7 +690,7 @@ def _get_benchmark_extras_headers_and_values(
         values: list[str] = [
             benchmark.benchmarker.profile.model_dump_json(),
             json.dumps(benchmark.benchmarker.backend),
-            json.dumps(benchmark.benchmarker.requests["attributes"]["data"]),
+            json.dumps(benchmark.benchmarker.requests["data"]),
         ]
 
         if len(headers) != len(values):

src/guidellm/benchmark/schemas.py

@@ -1674,6 +1674,7 @@ def compile(
         environment: Environment,
         strategy: SchedulingStrategy,
         constraints: dict[str, dict[str, Any]],
+        data: list[Any],
     ) -> GenerativeBenchmark:
         """
         Compile final generative benchmark from accumulated state.
@@ -1702,7 +1703,7 @@ def compile(
             ),
             benchmarker=BenchmarkerDict(
                 profile=profile,
-                requests=InfoMixin.extract_from_obj(requests),
+                requests={"data": data},
                 backend=backend.info,
                 environment=environment.info,
             ),

src/guidellm/data/deserializers/deserializer.py

@@ -1,10 +1,9 @@
 from __future__ import annotations
 
-import contextlib
 from collections.abc import Callable
 from typing import Any, Protocol, Union, runtime_checkable
 
-from datasets import Dataset, IterableDataset
+from datasets import Dataset, DatasetDict, IterableDataset, IterableDatasetDict
 from transformers import PreTrainedTokenizerBase
 
 from guidellm.data.utils import resolve_dataset_split
@@ -29,7 +28,7 @@ def __call__(
         processor_factory: Callable[[], PreTrainedTokenizerBase],
         random_seed: int,
         **data_kwargs: dict[str, Any],
-    ) -> dict[str, list]: ...
+    ) -> Dataset | IterableDataset | DatasetDict | IterableDatasetDict: ...
 
 
 class DatasetDeserializerFactory(
@@ -47,51 +46,16 @@ def deserialize(
         remove_columns: list[str] | None = None,
         **data_kwargs: dict[str, Any],
     ) -> Dataset | IterableDataset:
-        dataset = None
+        dataset: Dataset
 
         if type_ is None:
-            errors = []
-            # Note: There is no priority order for the deserializers, so all deserializers
-            #  must be mutually exclusive to ensure deterministic behavior.
-            for name, deserializer in cls.registry.items():
-                deserializer_fn: DatasetDeserializer = (
-                    deserializer() if isinstance(deserializer, type) else deserializer
-                )
-
-                try:
-                    with contextlib.suppress(DataNotSupportedError):
-                        dataset = deserializer_fn(
-                            data=data,
-                            processor_factory=processor_factory,
-                            random_seed=random_seed,
-                            **data_kwargs,
-                        )
-                except Exception as e:
-                    errors.append(e)
-
-                if dataset is not None:
-                    break # Found one that works. Continuing could overwrite it.
-
-            if dataset is None and len(errors) > 0:
-                raise DataNotSupportedError(f"data deserialization failed; {len(errors)} errors occurred while "
-                                            f"attempting to deserialize data {data}: {errors}")
-
-        elif deserializer := cls.get_registered_object(type_) is not None:
-            deserializer_fn: DatasetDeserializer = (
-                deserializer() if isinstance(deserializer, type) else deserializer
+            dataset = cls._deserialize_with_registered_deserializers(
+                data, processor_factory, random_seed, **data_kwargs
             )
 
-            dataset = deserializer_fn(
-                data=data,
-                processor_factory=processor_factory,
-                random_seed=random_seed,
-                **data_kwargs,
-            )
-
-        if dataset is None:
-            raise DataNotSupportedError(
-                f"No suitable deserializer found for data {data} "
-                f"with kwargs {data_kwargs} and deserializer type {type_}."
+        else:
+            dataset = cls._deserialize_with_specified_deserializer(
+                data, type_, processor_factory, random_seed, **data_kwargs
             )
 
         if resolve_split:
@@ -107,3 +71,74 @@ def deserialize(
             dataset = dataset.remove_columns(remove_columns)
 
         return dataset
+
+    @classmethod
+    def _deserialize_with_registered_deserializers(
+        cls,
+        data: Any,
+        processor_factory: Callable[[], PreTrainedTokenizerBase],
+        random_seed: int = 42,
+        **data_kwargs: dict[str, Any],
+    ) -> Dataset:
+        if cls.registry is None:
+            raise RuntimeError("registry is None; cannot deserialize dataset")
+        dataset: Dataset | None = None
+
+        errors: dict[str, Exception] = {}
+        # Note: There is no priority order for the deserializers, so all deserializers
+        #  must be mutually exclusive to ensure deterministic behavior.
+        for _name, deserializer in cls.registry.items():
+            deserializer_fn: DatasetDeserializer = (
+                deserializer() if isinstance(deserializer, type) else deserializer
+            )
+
+            try:
+                dataset = deserializer_fn(
+                    data=data,
+                    processor_factory=processor_factory,
+                    random_seed=random_seed,
+                    **data_kwargs,
+                )
+            except Exception as e:  # noqa: BLE001 # The exceptions are saved.
+                errors[_name] = e
+
+            if dataset is not None:
+                return dataset  # Success
+
+        if len(errors) > 0:
+            err_msgs = ""
+            def sort_key(item):
+                return (isinstance(item[1], DataNotSupportedError), item[0])
+            for key, err in sorted(errors.items(), key=sort_key):
+                err_msgs += f"\n  - Deserializer '{key}': ({type(err).__name__}) {err}"
+            raise ValueError(
+                "Data deserialization failed, likely because the input doesn't "
+                f"match any of the input formats. See the {len(errors)} error(s) that "
+                f"occurred while attempting to deserialize the data {data}:{err_msgs}"
+            )
+        return dataset
+
+    @classmethod
+    def _deserialize_with_specified_deserializer(
+        cls,
+        data: Any,
+        type_: str,
+        processor_factory: Callable[[], PreTrainedTokenizerBase],
+        random_seed: int = 42,
+        **data_kwargs: dict[str, Any],
+    ) -> Dataset:
+        deserializer_from_type = cls.get_registered_object(type_)
+        if deserializer_from_type is None:
+            raise ValueError(f"Deserializer type '{type_}' is not registered.")
+        if isinstance(deserializer_from_type, type):
+            deserializer_fn = deserializer_from_type()
+        else:
+            deserializer_fn = deserializer_from_type
+
+        return deserializer_fn(
+            data=data,
+            processor_factory=processor_factory,
+            random_seed=random_seed,
+            **data_kwargs,
+        )
+

src/guidellm/data/deserializers/file.py

@@ -34,11 +34,11 @@ def __call__(
         processor_factory: Callable[[], PreTrainedTokenizerBase],
         random_seed: int,
         **data_kwargs: dict[str, Any],
-    ) -> dict[str, list]:
+    ) -> Dataset:
         _ = (processor_factory, random_seed)  # Ignore unused args format errors
 
         if (
-            not isinstance(data, (str, Path))
+            not isinstance(data, str | Path)
             or not (path := Path(data)).exists()
             or not path.is_file()
             or path.suffix.lower() not in {".txt", ".text"}
@@ -62,10 +62,10 @@ def __call__(
         processor_factory: Callable[[], PreTrainedTokenizerBase],
         random_seed: int,
         **data_kwargs: dict[str, Any],
-    ) -> dict[str, list]:
+    ) -> Dataset:
         _ = (processor_factory, random_seed)
         if (
-            not isinstance(data, (str, Path))
+            not isinstance(data, str | Path)
             or not (path := Path(data)).exists()
             or not path.is_file()
             or path.suffix.lower() != ".csv"
@@ -86,10 +86,10 @@ def __call__(
         processor_factory: Callable[[], PreTrainedTokenizerBase],
         random_seed: int,
         **data_kwargs: dict[str, Any],
-    ) -> dict[str, list]:
+    ) -> Dataset:
         _ = (processor_factory, random_seed)
         if (
-            not isinstance(data, (str, Path))
+            not isinstance(data, str | Path)
             or not (path := Path(data)).exists()
             or not path.is_file()
             or path.suffix.lower() not in {".json", ".jsonl"}
@@ -110,10 +110,10 @@ def __call__(
         processor_factory: Callable[[], PreTrainedTokenizerBase],
         random_seed: int,
         **data_kwargs: dict[str, Any],
-    ) -> dict[str, list]:
+    ) -> Dataset:
         _ = (processor_factory, random_seed)
         if (
-            not isinstance(data, (str, Path))
+            not isinstance(data, str | Path)
             or not (path := Path(data)).exists()
             or not path.is_file()
             or path.suffix.lower() != ".parquet"
@@ -134,10 +134,10 @@ def __call__(
         processor_factory: Callable[[], PreTrainedTokenizerBase],
         random_seed: int,
         **data_kwargs: dict[str, Any],
-    ) -> dict[str, list]:
+    ) -> Dataset:
         _ = (processor_factory, random_seed)
         if (
-            not isinstance(data, (str, Path))
+            not isinstance(data, str | Path)
             or not (path := Path(data)).exists()
             or not path.is_file()
             or path.suffix.lower() != ".arrow"
@@ -158,10 +158,10 @@ def __call__(
         processor_factory: Callable[[], PreTrainedTokenizerBase],
         random_seed: int,
         **data_kwargs: dict[str, Any],
-    ) -> dict[str, list]:
+    ) -> Dataset:
         _ = (processor_factory, random_seed)
         if (
-            not isinstance(data, (str, Path))
+            not isinstance(data, str | Path)
             or not (path := Path(data)).exists()
             or not path.is_file()
             or path.suffix.lower() not in {".hdf5", ".h5"}
@@ -185,7 +185,7 @@ def __call__(
     ) -> dict[str, list]:
         _ = (processor_factory, random_seed)
         if (
-            not isinstance(data, (str, Path))
+            not isinstance(data, str | Path)
             or not (path := Path(data)).exists()
             or not path.is_file()
             or path.suffix.lower() != ".db"
@@ -209,7 +209,7 @@ def __call__(
     ) -> dict[str, list]:
         _ = (processor_factory, random_seed)
         if (
-            not isinstance(data, (str, Path))
+            not isinstance(data, str | Path)
             or not (path := Path(data)).exists()
             or not path.is_file()
             or path.suffix.lower() != ".tar"

src/guidellm/data/deserializers/huggingface.py

@@ -36,7 +36,7 @@ def __call__(
         processor_factory: Callable[[], PreTrainedTokenizerBase],
         random_seed: int,
         **data_kwargs: dict[str, Any],
-    ) -> dict[str, list]:
+    ) -> Dataset | IterableDataset | DatasetDict | IterableDatasetDict:
         _ = (processor_factory, random_seed)
 
         if isinstance(