fix: mark review comments

Signed-off-by: Alon Kellner <[email protected]>

Author: AlonKellner-RedHat

src/guidellm/__main__.py

@@ -25,6 +25,7 @@
 
 import asyncio
 import codecs
+import json
 from pathlib import Path
 
 import click
@@ -383,19 +384,40 @@ def benchmark():
     help="Maximum global error rate across all benchmarks.",
 )
 @click.option(
-    "--stop-over-saturated",
-    "--stop-osd",  # alias
-    default=BenchmarkGenerativeTextArgs.get_default("stop_over_saturated"),
+    "--over-saturation",
+    "--detect-saturation",  # alias
+    default=None,
     help=(
-        "Set this flag to stop the benchmark if the model is over-saturated. "
-        "Defaults to False."
+        "Enable over-saturation detection. Can be a flag (bool) or a JSON dict with "
+        'configuration (e.g., \'{"enabled": true, "min_seconds": 30}\'). '
+        "Defaults to None (disabled)."
     ),
-    is_flag=True,
+    type=click.UNPROCESSED,
 )
-def run(**kwargs):
+def run(**kwargs):  # noqa: C901
     # Only set CLI args that differ from click defaults
     kwargs = cli_tools.set_if_not_default(click.get_current_context(), **kwargs)
 
+    # Handle over_saturation parsing (can be bool flag or JSON dict string)
+    if "over_saturation" in kwargs and kwargs["over_saturation"] is not None:
+        over_sat = kwargs["over_saturation"]
+        if isinstance(over_sat, str):
+            try:
+                # Try parsing as JSON dict
+                kwargs["over_saturation"] = json.loads(over_sat)
+            except (json.JSONDecodeError, ValueError):
+                # If not valid JSON, treat as bool flag
+                kwargs["over_saturation"] = over_sat.lower() in (
+                    "true",
+                    "1",
+                    "yes",
+                    "on",
+                )
+        elif isinstance(over_sat, bool):
+            # Already a bool, keep as is
+            pass
+        # If it's already a dict, keep as is
+
     # Handle remapping for request params
     request_type = kwargs.pop("request_type", None)
     request_formatter_kwargs = kwargs.pop("request_formatter_kwargs", None)

src/guidellm/benchmark/entrypoints.py

@@ -323,7 +323,7 @@ async def resolve_profile(
     max_errors: int | None,
     max_error_rate: float | None,
     max_global_error_rate: float | None,
-    stop_over_saturated: bool | None = None,
+    over_saturation: bool | dict[str, Any] | None = None,
     console: Console | None = None,
 ) -> Profile:
     """
@@ -344,7 +344,7 @@ async def resolve_profile(
     :param max_errors: Maximum number of errors before stopping
     :param max_error_rate: Maximum error rate threshold before stopping
     :param max_global_error_rate: Maximum global error rate threshold before stopping
-    :param stop_over_saturated: Whether to stop if over-saturation is detected
+    :param over_saturation: Over-saturation detection configuration (bool or dict)
     :param console: Console instance for progress reporting, or None
     :return: Configured Profile instance ready for benchmarking
     :raises ValueError: If constraints are provided with a pre-configured Profile
@@ -361,7 +361,7 @@ async def resolve_profile(
         "max_errors": max_errors,
         "max_error_rate": max_error_rate,
         "max_global_error_rate": max_global_error_rate,
-        "stop_over_saturated": stop_over_saturated,
+        "over_saturation": over_saturation,
     }.items():
         if val is not None:
             constraints[key] = val
@@ -503,7 +503,7 @@ async def benchmark_generative_text(
         max_errors=args.max_errors,
         max_error_rate=args.max_error_rate,
         max_global_error_rate=args.max_global_error_rate,
-        stop_over_saturated=args.stop_over_saturated,
+        over_saturation=args.over_saturation,
         console=console,
     )
     output_formats = await resolve_output_formats(

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

@@ -283,9 +283,13 @@ def get_default(cls: type[BenchmarkGenerativeTextArgs], field: str) -> Any:
     max_global_error_rate: float | None = Field(
         default=None, description="Maximum global error rate (0-1) before stopping"
     )
-    stop_over_saturated: bool = Field(
-        default=False,
-        description="Whether to stop the benchmark if over-saturation is detected",
+    over_saturation: bool | dict[str, Any] | None = Field(
+        default=None,
+        description=(
+            "Over-saturation detection configuration. Can be a bool to enable/disable "
+            "with defaults, or a dict with configuration parameters (enabled, "
+            "min_seconds, max_window_seconds, moe_threshold, etc.)."
+        ),
     )
 
     @field_validator("data", "data_args", "rate", mode="wrap")

src/guidellm/scheduler/__init__.py

@@ -21,7 +21,6 @@
     MaxNumberConstraint,
     OverSaturationConstraint,
     OverSaturationConstraintInitializer,
-    OverSaturationDetector,
     PydanticConstraintInitializer,
     SerializableConstraintInitializer,
     UnserializableConstraintInitializer,
@@ -71,7 +70,6 @@
     "NonDistributedEnvironment",
     "OverSaturationConstraint",
     "OverSaturationConstraintInitializer",
-    "OverSaturationDetector",
     "PydanticConstraintInitializer",
     "RequestT",
     "ResponseT",

src/guidellm/scheduler/constraints/__init__.py

@@ -13,16 +13,15 @@
     UnserializableConstraintInitializer,
 )
 from .factory import ConstraintsInitializerFactory
-from .over_saturation import (
-    OverSaturationConstraint,
-    OverSaturationConstraintInitializer,
-    OverSaturationDetector,
-)
 from .protocols import (
     Constraint,
     ConstraintInitializer,
     SerializableConstraintInitializer,
 )
+from .saturation import (
+    OverSaturationConstraint,
+    OverSaturationConstraintInitializer,
+)
 from .standard import (
     MaxDurationConstraint,
     MaxErrorRateConstraint,
@@ -43,7 +42,6 @@
     "MaxNumberConstraint",
     "OverSaturationConstraint",
     "OverSaturationConstraintInitializer",
-    "OverSaturationDetector",
     "PydanticConstraintInitializer",
     "RequestsExhaustedConstraint",
     "SerializableConstraintInitializer",