Reenablement of flows and fixes

Author: markurtz

src/guidellm/__main__.py

@@ -286,7 +286,7 @@ def benchmark():
 )
 @click.option(
     "--data-num-workers",
-    default=1,
+    default=None,
     type=int,
     help="The number of worker processes to use for data loading.",
 )
@@ -505,11 +505,9 @@ def run(
             output_formats=output_formats,
             # Updates configuration
             progress=(
-                [
-                    GenerativeConsoleBenchmarkerProgress(
-                        display_scheduler_stats=display_scheduler_stats
-                    )
-                ]
+                GenerativeConsoleBenchmarkerProgress(
+                    display_scheduler_stats=display_scheduler_stats
+                )
                 if not disable_progress
                 else None
             ),

src/guidellm/benchmark/__init__.py

@@ -15,11 +15,7 @@
     SynchronousProfile,
     ThroughputProfile,
 )
-from .progress import (
-    BenchmarkerProgress,
-    BenchmarkerProgressGroup,
-    GenerativeConsoleBenchmarkerProgress,
-)
+from .progress import BenchmarkerProgress, GenerativeConsoleBenchmarkerProgress
 from .schemas import (
     Benchmark,
     BenchmarkArgs,
@@ -44,7 +40,6 @@
     "Benchmarker",
     "BenchmarkerDict",
     "BenchmarkerProgress",
-    "BenchmarkerProgressGroup",
     "ConcurrentProfile",
     "EstimatedBenchmarkState",
     "GenerativeAudioMetricsSummary",

src/guidellm/benchmark/benchmarker.py

@@ -23,20 +23,19 @@
 from typing import Generic
 
 from guidellm.benchmark.profile import Profile
+from guidellm.benchmark.progress import BenchmarkerProgress
 from guidellm.benchmark.schemas import (
     BenchmarkArgs,
     BenchmarkT,
     EstimatedBenchmarkState,
 )
+from guidellm.logger import logger
 from guidellm.scheduler import (
     BackendInterface,
     Environment,
-    NonDistributedEnvironment,
     RequestT,
     ResponseT,
     Scheduler,
-    SchedulerState,
-    SchedulingStrategy,
 )
 from guidellm.utils import ThreadSafeSingletonMixin
 
@@ -65,19 +64,13 @@ async def run(
         requests: Iterable[RequestT | Iterable[RequestT | tuple[RequestT, float]]],
         backend: BackendInterface[RequestT, ResponseT],
         profile: Profile,
-        environment: Environment | None = None,
+        environment: Environment,
+        progress: BenchmarkerProgress[BenchmarkT] | None = None,
         sample_requests: int | None = 20,
         warmup: float | None = None,
         cooldown: float | None = None,
         prefer_response_metrics: bool = True,
-    ) -> AsyncIterator[
-        tuple[
-            EstimatedBenchmarkState | None,
-            BenchmarkT | None,
-            SchedulingStrategy,
-            SchedulerState | None,
-        ]
-    ]:
+    ) -> AsyncIterator[BenchmarkT]:
         """
         Execute benchmark runs across multiple scheduling strategies.
 
@@ -95,15 +88,17 @@ async def run(
         :raises Exception: If benchmark execution or compilation fails.
         """
         with self.thread_lock:
-            if environment is None:
-                environment = NonDistributedEnvironment()
+            if progress:
+                await progress.on_initialize(profile)
 
             run_id = str(uuid.uuid4())
             strategies_generator = profile.strategies_generator()
             strategy, constraints = next(strategies_generator)
 
             while strategy is not None:
-                yield None, None, strategy, None
+                if progress:
+                    await progress.on_benchmark_start(strategy)
+
                 args = BenchmarkArgs(
                     run_id=run_id,
                     run_index=len(profile.completed_strategies),
@@ -127,18 +122,23 @@ async def run(
                     env=environment,
                     **constraints or {},
                 ):
-                    benchmark_class.update_estimate(
-                        args,
-                        estimated_state,
-                        response,
-                        request,
-                        request_info,
-                        scheduler_state,
-                    )
-                    yield estimated_state, None, strategy, scheduler_state
-
-                if scheduler_state is None:
-                    raise RuntimeError("Scheduler state is None after execution.")
+                    try:
+                        benchmark_class.update_estimate(
+                            args,
+                            estimated_state,
+                            response,
+                            request,
+                            request_info,
+                            scheduler_state,
+                        )
+                        if progress:
+                            await progress.on_benchmark_update(
+                                estimated_state, scheduler_state
+                            )
+                    except Exception as err:
+                        logger.error(
+                            f"Error updating benchmark estimate/progress: {err}"
+                        )
 
                 benchmark = benchmark_class.compile(
                     args=args,
@@ -151,10 +151,16 @@ async def run(
                     strategy=strategy,
                     constraints=constraints,
                 )
-                yield None, benchmark, strategy, None
+                if progress:
+                    await progress.on_benchmark_complete(benchmark)
+
+                yield benchmark
 
                 try:
                     strategy, constraints = strategies_generator.send(benchmark)
                 except StopIteration:
                     strategy = None
                     constraints = None
+
+            if progress:
+                await progress.on_finalize()

src/guidellm/benchmark/entrypoints.py

@@ -10,9 +10,9 @@
 from guidellm.benchmark.benchmarker import Benchmarker
 from guidellm.benchmark.output import GenerativeBenchmarkerOutput
 from guidellm.benchmark.profile import Profile, ProfileType
-from guidellm.benchmark.progress import BenchmarkerProgressGroup
+from guidellm.benchmark.progress import BenchmarkerProgress
 from guidellm.benchmark.schemas import GenerativeBenchmark, GenerativeBenchmarksReport
-from guidellm.benchmark.types import OutputFormatT, ProcessorInputT, ProgressInputT
+from guidellm.benchmark.types import OutputFormatT, ProcessorInputT
 from guidellm.data import (
     DataLoader,
     DatasetPreprocessor,
@@ -271,7 +271,6 @@ async def resolve_output_formats(
     return resolved
 
 
-# @validate_call(config={"arbitrary_types_allowed": True})
 async def benchmark_generative_text(  # noqa: C901, PLR0915, PLR0912
     # Required
     target: str,
@@ -296,7 +295,7 @@ async def benchmark_generative_text(  # noqa: C901, PLR0915, PLR0912
     ) = "chat_completions",
     data_collator: Callable | Literal["generative"] | None = "generative",
     data_sampler: Sampler[int] | Literal["shuffle"] | None = None,
-    data_num_workers: int | None = 1,
+    data_num_workers: int | None = None,
     dataloader_kwargs: dict[str, Any] | None = None,
     random_seed: int = 42,
     # Output configuration
@@ -308,7 +307,7 @@ async def benchmark_generative_text(  # noqa: C901, PLR0915, PLR0912
         | None
     ) = ("console", "json", "html", "csv"),
     # Updates configuration
-    progress: ProgressInputT | None = None,
+    progress: BenchmarkerProgress | None = None,
     print_updates: bool = False,
     # Benchmarker configuration
     benchmark_cls: type[GenerativeBenchmark] = GenerativeBenchmark,
@@ -366,37 +365,26 @@ async def benchmark_generative_text(  # noqa: C901, PLR0915, PLR0912
         output_formats=output_formats, output_path=output_path, console=console
     )
 
-    progress_group = BenchmarkerProgressGroup(
-        instances=progress or [], enabled=bool(progress)
-    )
     report = GenerativeBenchmarksReport()
     console.print_update(
         title="Setup complete, starting benchmarks...", status="success"
     )
     console.print("\n\n")
 
-    async for (
-        _aggregator_update,
-        benchmark,
-        _strategy,
-        _scheduler_state,
-    ) in progress_group(
-        profile,
-        Benchmarker[
-            GenerativeBenchmark,
-            GenerationRequest,
-            GenerationResponse,
-        ]().run(
-            benchmark_class=benchmark_cls,
-            requests=request_loader,
-            backend=backend,
-            profile=profile,
-            environment=NonDistributedEnvironment(),
-            sample_requests=sample_requests,
-            warmup=warmup,
-            cooldown=cooldown,
-            prefer_response_metrics=True,
-        ),
+    benchmarker: Benchmarker[
+        GenerativeBenchmark, GenerationRequest, GenerationResponse
+    ] = Benchmarker()
+    async for benchmark in benchmarker.run(
+        benchmark_class=benchmark_cls,
+        requests=request_loader,
+        backend=backend,
+        profile=profile,
+        environment=NonDistributedEnvironment(),
+        progress=progress,
+        sample_requests=sample_requests,
+        warmup=warmup,
+        cooldown=cooldown,
+        prefer_response_metrics=True,
     ):
         if benchmark:
             report.benchmarks.append(benchmark)