fix: resolve type checking errors in over_saturation constraint

- Fix first_iteration -> first_token_iteration attribute name
- Add type ignore for OverSaturationConstraint return type
- Fix validated_kwargs type handling for stop_over_saturated parameter

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

Author: AlonKellner-RedHat

src/guidellm/scheduler/constraints/base.py

@@ -15,7 +15,10 @@
 from guidellm.scheduler.schemas import SchedulerState, SchedulerUpdateAction
 from guidellm.schemas import RequestInfo, StandardBaseModel
 from guidellm.utils import InfoMixin
-from .protocols import Constraint, ConstraintInitializer, SerializableConstraintInitializer
+
+from .protocols import (
+    Constraint,
+)
 
 __all__ = [
     "PydanticConstraintInitializer",
@@ -130,4 +133,3 @@ def __call__(
             "Cannot invoke unserializable constraint instance. "
             "This constraint was not properly serialized and cannot be executed."
         )
-

src/guidellm/scheduler/constraints/factory.py

@@ -11,6 +11,7 @@
 from typing import Any
 
 from guidellm.utils import InfoMixin, RegistryMixin
+
 from .base import UnserializableConstraintInitializer
 from .protocols import (
     Constraint,
@@ -180,4 +181,3 @@ def resolve_constraints(
                 resolved_constraints[key] = cls.create_constraint(key, val)
 
         return resolved_constraints
-

src/guidellm/scheduler/constraints/over_saturation.py

@@ -14,9 +14,13 @@
 
 from pydantic import Field
 
-from guidellm.scheduler.schemas import SchedulerState, SchedulerUpdateAction
+from guidellm.scheduler.schemas import (
+    SchedulerState,
+    SchedulerUpdateAction,
+)
 from guidellm.schemas import RequestInfo
 from guidellm.settings import settings
+
 from .base import PydanticConstraintInitializer
 from .factory import ConstraintsInitializerFactory
 from .protocols import Constraint
@@ -362,10 +366,10 @@ def __call__(
         elif (
             request_info.status == "completed"
             and request_info.timings
-            and request_info.timings.first_iteration
+            and request_info.timings.first_token_iteration
         ):
             ttft = (
-                request_info.timings.first_iteration
+                request_info.timings.first_token_iteration
                 - request_info.timings.request_start
             )
             self.over_saturation_detector.add_finished(
@@ -436,13 +440,13 @@ def create_constraint(self, **_kwargs) -> Constraint:
             minimum_duration=self.min_seconds,
             maximum_window_seconds=self.max_window_seconds,
         )
-        return OverSaturationConstraint(
+        return OverSaturationConstraint(  # type: ignore[return-value]
             over_saturation_detector=over_saturation_detector,
             stop_over_saturated=self.stop_over_saturated,
         )
 
     @classmethod
-    def validated_kwargs(cls, stop_over_saturated: bool, **kwargs) -> dict[str, Any]:
+    def validated_kwargs(cls, stop_over_saturated: bool | None = None, **kwargs) -> dict[str, Any]:
         """
         Validate and process arguments for OverSaturationConstraint creation.
 
@@ -451,8 +455,10 @@ def validated_kwargs(cls, stop_over_saturated: bool, **kwargs) -> dict[str, Any]
         :return: Validated dictionary with stop_over_saturated field
         """
         aliases = ["stop_over_saturated", "stop_over_sat", "stop_osd"]
+        result = stop_over_saturated if stop_over_saturated is not None else False
         for alias in aliases:
-            stop_over_saturated = stop_over_saturated or kwargs.get(alias)
-
-        return {"stop_over_saturated": stop_over_saturated}
+            alias_value = kwargs.get(alias)
+            if alias_value is not None:
+                result = bool(alias_value) or result
 
+        return {"stop_over_saturated": result}

src/guidellm/scheduler/constraints/protocols.py

@@ -85,4 +85,3 @@ def create_constraint(self, **kwargs) -> Constraint:
         :param kwargs: Additional configuration parameters
         :return: Configured constraint evaluation function
         """
-

src/guidellm/scheduler/constraints/standard.py

@@ -20,6 +20,7 @@
 from guidellm.schemas import RequestInfo, StandardBaseModel
 from guidellm.settings import settings
 from guidellm.utils import InfoMixin
+
 from .base import PydanticConstraintInitializer
 from .factory import ConstraintsInitializerFactory
 from .protocols import Constraint
@@ -690,4 +691,3 @@ def __call__(
                 stop_time=stop_time,
             ),
         )
-

tests/unit/scheduler/OVER_SATURATION_TEST_COVERAGE.md

@@ -142,7 +142,7 @@ This document outlines the comprehensive unit test coverage for the over-saturat
 - **Focus**: Basic initialization, core algorithms, critical paths
 
 ### Sanity Tests (`@pytest.mark.sanity`)
-- **Count**: 21 tests  
+- **Count**: 21 tests
 - **Purpose**: Comprehensive validation of feature behavior
 - **Runtime**: 1-3 minutes total
 - **Focus**: Realistic scenarios, robustness, edge cases
@@ -156,22 +156,22 @@ This document outlines the comprehensive unit test coverage for the over-saturat
 - ✅ **Threshold detection**: TTFT violations and concurrent request tracking
 - ✅ **Statistical significance**: Margin of error and confidence testing
 
-### Integration Coverage  
+### Integration Coverage
 - ✅ **Detector ↔ Constraint**: Proper data flow and decision making
 - ✅ **Constraint ↔ Scheduler**: State integration and action generation
 - ✅ **Factory ↔ Initializer**: Proper constraint creation and configuration
 - ✅ **Timing ↔ Detection**: Accurate duration and timing calculations
 
 ### Robustness Coverage
 - ✅ **Empty data**: No crashes or false positives
-- ✅ **Malformed data**: Proper validation and error handling  
+- ✅ **Malformed data**: Proper validation and error handling
 - ✅ **Extreme values**: Numerical stability maintained
 - ✅ **Memory management**: Bounded growth under stress
 - ✅ **Performance**: Efficiency maintained at scale
 
 ### Scenario Coverage
 - ✅ **Gradual degradation**: Detected correctly
-- ✅ **Sudden spikes**: Detected correctly  
+- ✅ **Sudden spikes**: Detected correctly
 - ✅ **Stable performance**: No false positives
 - ✅ **Recovery patterns**: Proper handling
 - ✅ **Variable workloads**: Robust detection

tests/unit/scheduler/test_over_saturation.py

@@ -154,9 +154,7 @@ def test_check_alert_requires_minimum_duration(self):
     @pytest.mark.sanity
     def test_check_alert_requires_minimum_window_size(self):
         """Test that check_alert requires minimum window size."""
-        detector = OverSaturationDetector(
-            minimum_duration=0.0, minimum_window_size=10
-        )
+        detector = OverSaturationDetector(minimum_duration=0.0, minimum_window_size=10)
 
         # Add few requests
         for i in range(5):
@@ -291,15 +289,11 @@ def test_constraint_stops_when_over_saturated(self, detector):
         # Simulate over-saturation by creating positive slopes
         # Add many started requests with increasing concurrent count
         for i in range(20):
-            detector.add_started(
-                {"concurrent_requests": i * 2, "duration": float(i)}
-            )
+            detector.add_started({"concurrent_requests": i * 2, "duration": float(i)})
 
         # Add finished requests with increasing TTFT
         for i in range(20):
-            detector.add_finished(
-                {"ttft": 1.0 + i * 0.1, "duration": float(i) + 10.0}
-            )
+            detector.add_finished({"ttft": 1.0 + i * 0.1, "duration": float(i) + 10.0})
 
         detector.update_duration(30.0)
         detector.check_alert()  # Prime the slope checkers
@@ -363,7 +357,11 @@ class TestOverSaturationConstraintInitializer:
         params=[
             {"stop_over_saturated": True},
             {"stop_over_saturated": False},
-            {"stop_over_saturated": True, "min_seconds": 10.0, "max_window_seconds": 60.0},
+            {
+                "stop_over_saturated": True,
+                "min_seconds": 10.0,
+                "max_window_seconds": 60.0,
+            },
         ]
     )
     def valid_instances(self, request):
@@ -619,4 +617,3 @@ def test_check_slope_requires_minimum_samples(self, slope_checker):
         slope_checker.add_data_point(3.0, 6.0)
         result = slope_checker.check_slope(3.0)
         # Might be True or False depending on confidence intervals
-