Skip to content

Commit f3d9d06

Browse files
committed
Add builtin scenarios
1 parent aa6dc13 commit f3d9d06

File tree

3 files changed

+25
-6
lines changed

3 files changed

+25
-6
lines changed

src/guidellm/__main__.py

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
from guidellm.backend import BackendType
1010
from guidellm.benchmark import ProfileType
1111
from guidellm.benchmark.entrypoints import benchmark_with_scenario
12-
from guidellm.benchmark.scenario import GenerativeTextScenario
12+
from guidellm.benchmark.scenario import GenerativeTextScenario, get_builtin_scenarios
1313
from guidellm.config import print_config
1414
from guidellm.preprocess.dataset import ShortPromptStrategy, process_dataset
1515
from guidellm.scheduler import StrategyType
@@ -38,7 +38,7 @@ def cli():
3838
dir_okay=False,
3939
path_type=Path, # type: ignore[type-var]
4040
),
41-
click.STRING
41+
click.Choice(get_builtin_scenarios()),
4242
),
4343
default=None,
4444
help=("TODO: A scenario or path to config"),
@@ -280,8 +280,7 @@ def benchmark(
280280
elif isinstance(scenario, Path):
281281
_scenario = GenerativeTextScenario.from_file(scenario, overrides)
282282
else:
283-
# TODO: Add support for builtin scenarios
284-
raise NotImplementedError
283+
_scenario = GenerativeTextScenario.from_builtin(scenario, overrides)
285284
except ValidationError as e:
286285
errs = e.errors(include_url=False, include_context=True, include_input=True)
287286
param_name = "--" + str(errs[0]["loc"][0]).replace("_", "-")

src/guidellm/benchmark/scenario.py

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
from collections.abc import Iterable
2+
from functools import cache
23
from pathlib import Path
3-
from typing import Annotated, Any, Literal, Optional, Union
4+
from typing import Annotated, Any, Literal, Optional, TypeVar, Union
45

56
from datasets import Dataset, DatasetDict, IterableDataset, IterableDatasetDict
67
from pydantic import BeforeValidator, Field, NonNegativeInt, PositiveFloat, PositiveInt
@@ -13,7 +14,14 @@
1314
from guidellm.objects.pydantic import StandardBaseModel
1415
from guidellm.scheduler.strategy import StrategyType
1516

16-
__ALL__ = ["Scenario", "GenerativeTextScenario"]
17+
__ALL__ = ["Scenario", "GenerativeTextScenario", "get_builtin_scenarios"]
18+
19+
SCENARIO_DIR = Path(__file__).parent / "scenarios/"
20+
21+
22+
@cache
23+
def get_builtin_scenarios() -> list[str]:
24+
return [p.stem for p in SCENARIO_DIR.glob("*.json")]
1725

1826

1927
def parse_float_list(value: Union[str, float, list[float]]) -> list[float]:
@@ -32,9 +40,21 @@ def parse_float_list(value: Union[str, float, list[float]]) -> list[float]:
3240
) from err
3341

3442

43+
T = TypeVar("T", bound="Scenario")
44+
45+
3546
class Scenario(StandardBaseModel):
3647
target: str
3748

49+
@classmethod
50+
def from_builtin(cls: type[T], name: str, overrides: Optional[dict] = None) -> T:
51+
filename = SCENARIO_DIR / f"{name}.json"
52+
53+
if not filename.is_file():
54+
raise ValueError(f"{name} is not a vaild builtin scenario")
55+
56+
return cls.from_file(filename, overrides)
57+
3858

3959
class GenerativeTextScenario(Scenario):
4060
# FIXME: This solves an issue with Pydantic and class types

src/guidellm/benchmark/scenarios/__init__.py

Whitespace-only changes.

0 commit comments

Comments
 (0)