wip

Author: jopemachine

src/ai/backend/manager/models/deployment_policy/row.py

@@ -3,7 +3,7 @@
 import logging
 import uuid
 from datetime import datetime
-from typing import TYPE_CHECKING, Union
+from typing import TYPE_CHECKING
 
 import sqlalchemy as sa
 from pydantic import BaseModel, Field, model_validator
@@ -62,7 +62,7 @@ class BlueGreenSpec(BaseModel):
     promote_delay_seconds: int = 0
 
 
-DeploymentStrategySpec = Union[RollingUpdateSpec, BlueGreenSpec]
+DeploymentStrategySpec = RollingUpdateSpec | BlueGreenSpec
 
 
 def _get_endpoint_join_condition() -> sa.ColumnElement[bool]:

src/ai/backend/manager/sokovan/deployment/strategy/rolling_update.py

@@ -80,7 +80,8 @@ def evaluate_cycle(
         sweep handles it by transitioning to ROLLING_BACK when the
         deploying timeout is exceeded.
         """
-        assert isinstance(spec, RollingUpdateSpec)
+        if not isinstance(spec, RollingUpdateSpec):
+            raise TypeError(f"Expected RollingUpdateSpec, got {type(spec).__name__}")
         desired = deployment.replica_spec.target_replica_count
         deploying_revision_id = deployment.deploying_revision_id
         if deploying_revision_id is None:
@@ -136,7 +137,11 @@ def _check_completed(
         desired: int,
     ) -> StrategyCycleResult | None:
         """Return COMPLETED result if all old routes are replaced and enough new are healthy."""
-        if classified.old_active or classified.new_healthy_count < desired:
+        if (
+            classified.old_active
+            or classified.new_provisioning_count
+            or classified.new_healthy_count < desired
+        ):
             return None
         log.info(
             "deployment {}: rolling update complete ({} healthy routes)",