wip

Author: jopemachine

src/ai/backend/manager/data/deployment/types.py

@@ -380,24 +380,12 @@ class DeploymentInfo:
     deploying_revision_id: UUID | None = None
     sub_step: DeploymentSubStep | None = None
 
-    def resolve_revision_spec(self, revision_id: UUID | None = None) -> ModelRevisionSpec | None:
-        """Resolve a revision spec with fallback: revision_id → deploying → current."""
-        for candidate_id in (revision_id, self.deploying_revision_id, self.current_revision_id):
-            if candidate_id is None:
-                continue
-            found = next(
-                (r for r in self.model_revisions if r.revision_id == candidate_id),
-                None,
-            )
-            if found is not None:
-                return found
-        return None
-
-    def current_revision_spec(self) -> ModelRevisionSpec | None:
-        return self.resolve_revision_spec(self.current_revision_id)
-
-    def deploying_revision_spec(self) -> ModelRevisionSpec | None:
-        return self.resolve_revision_spec(self.deploying_revision_id)
+    def resolve_revision_spec(self, revision_id: UUID) -> ModelRevisionSpec | None:
+        """Find a ModelRevisionSpec by revision_id from model_revisions."""
+        return next(
+            (r for r in self.model_revisions if r.revision_id == revision_id),
+            None,
+        )
 
 
 @dataclass

src/ai/backend/manager/repositories/scheduler/types/session_creation.py

@@ -117,9 +117,18 @@ def from_deployment_info(
     ) -> Self:
         session_creation_id = secrets.token_urlsafe(16)
         # Use the specified revision, falling back to deploying then current.
-        target_revision = deployment_info.resolve_revision_spec(revision_id)
-        if target_revision is None:
+        target_revision_id = (
+            revision_id
+            or deployment_info.deploying_revision_id
+            or deployment_info.current_revision_id
+        )
+        if target_revision_id is None:
             raise DeploymentHasNoTargetRevision("Deployment has no revision for session creation")
+        target_revision = deployment_info.resolve_revision_spec(target_revision_id)
+        if target_revision is None:
+            raise DeploymentHasNoTargetRevision(
+                f"Revision {target_revision_id} not found in model_revisions"
+            )
 
         # Prepare mount spec
         mount_spec = target_revision.mounts.to_mount_spec()

src/ai/backend/manager/sokovan/deployment/deployment_controller.py

@@ -147,11 +147,14 @@ async def update_deployment(
         modified_endpoint = await self._deployment_repository.get_modified_endpoint(
             endpoint_id=endpoint_id, updater=updater
         )
-        current_revision = modified_endpoint.current_revision_spec()
-        if current_revision:
-            await self._scheduling_controller.validate_session_spec(
-                SessionValidationSpec.from_revision(model_revision=current_revision)
+        if modified_endpoint.current_revision_id is not None:
+            current_revision = modified_endpoint.resolve_revision_spec(
+                modified_endpoint.current_revision_id
             )
+            if current_revision:
+                await self._scheduling_controller.validate_session_spec(
+                    SessionValidationSpec.from_revision(model_revision=current_revision)
+                )
         res = await self._deployment_repository.update_endpoint_with_spec(updater)
         try:
             await self.mark_lifecycle_needed(DeploymentLifecycleType.CHECK_REPLICA)

src/ai/backend/manager/sokovan/deployment/executor.py

@@ -127,13 +127,20 @@ async def check_pending_deployments(
         valid_deployments: list[DeploymentWithHistory] = []
         for deployment in deployments:
             info = deployment.deployment_info
-            current_revision = info.current_revision_spec()
-            if not current_revision:
+            if info.current_revision_id is None:
                 log.warning(
                     "Deployment {} has no current revision, skipping",
                     info.id,
                 )
                 continue
+            current_revision = info.resolve_revision_spec(info.current_revision_id)
+            if not current_revision:
+                log.warning(
+                    "Deployment {} current revision {} not found in model_revisions, skipping",
+                    info.id,
+                    info.current_revision_id,
+                )
+                continue
             targets = scaling_group_targets[info.metadata.resource_group]
             if not targets:
                 log.warning(
@@ -443,11 +450,15 @@ async def _register_endpoint(
 
         with recorder.phase("register_endpoint"):
             with recorder.step("check_target_revision"):
-                current_revision = deployment.current_revision_spec()
-                if not current_revision:
+                if deployment.current_revision_id is None:
                     raise ModelDefinitionNotFound(
                         f"No current revision for deployment {deployment.id}"
                     )
+                current_revision = deployment.resolve_revision_spec(deployment.current_revision_id)
+                if not current_revision:
+                    raise ModelDefinitionNotFound(
+                        f"Current revision {deployment.current_revision_id} not found for deployment {deployment.id}"
+                    )
 
             with recorder.step("generate_model_definition"):
                 model_definition = (