Keep deploys terminated during model load

src/orchestration/models/model_deployment/deployment_linker.py

@@ -29,6 +29,7 @@
     InventoryRepository,
 )
 from orchestration.repositories.model_deployment_repo import (
+    ENDED_STATES,
     ModelDeploymentRepository,
 )
 from orchestration.models.model_deployment.model_ref import (
@@ -259,15 +260,9 @@ async def _drive_deploy(self, node_id: UUID, deploy: dict,
             # is VISIBLE in the deploy row + dashboard — set_state writes state
             # only, which previously left load_model failures with an empty
             # error_message and no clue why the container never started.
-            async with self._db.acquire() as conn:
-                async with conn.transaction():
-                    await self._inventory.release_gpu(
-                        node_id, gpu_required, tx=conn,
-                    )
-                    await self._deploys.update_state(
-                        deploy["id"], "FAILED", tx=conn,
-                        error_message=f"load_model error: {e}",
-                    )
+            await self._mark_failed(
+                node_id, deploy["id"], gpu_required, f"load_model error: {e}",
+            )
             return
 
         # The worker can return a CommandResult with status='failed' (readiness
@@ -279,20 +274,33 @@ async def _drive_deploy(self, node_id: UUID, deploy: dict,
                 "linker: load_model returned status=failed for deploy=%s: %s",
                 deploy["id"], detail,
             )
-            async with self._db.acquire() as conn:
-                async with conn.transaction():
-                    await self._inventory.release_gpu(
-                        node_id, gpu_required, tx=conn,
-                    )
-                    await self._deploys.update_state(
-                        deploy["id"], "FAILED", tx=conn,
-                        error_message=f"load_model failed on worker: {detail}",
-                    )
+            await self._mark_failed(
+                node_id, deploy["id"], gpu_required,
+                f"load_model failed on worker: {detail}",
+            )
             return
 
         # Model loaded on the worker. Promote DEPLOYING → RUNNING.
-        async with self._db.acquire() as conn:
-            await self._deploys.set_state(deploy["id"], "RUNNING", tx=conn)
+        promoted = await self._deploys.update_state(
+            deploy["id"], "RUNNING", unless_in=ENDED_STATES,
+        )
+        if not promoted:
+            # Terminated while loading: terminate already released the GPU and
+            # its unload ran before the container existed, so unload it now.
+            logger.info(
+                "linker: deploy=%s ended during load_model; unloading",
+                deploy["id"],
+            )
+            try:
+                await self._controller.unload_model(
+                    node_id=str(node_id), deployment_id=str(deploy["id"]),
+                )
+            except Exception:
+                logger.exception(
+                    "linker: unload after terminate failed for deploy=%s",
+                    deploy["id"],
+                )
+            return
         # Publish the inference endpoint (the node's CP-reachable advertise_url,
         # NOT the worker's 127.0.0.1 loopback) so the data plane can route.
         try:
@@ -313,3 +321,17 @@ async def _drive_deploy(self, node_id: UUID, deploy: dict,
             logger.exception(
                 "linker: failed to set endpoint for deploy=%s", deploy["id"],
             )
+
+    async def _mark_failed(self, node_id: UUID, deploy_id, gpu_required: int,
+                           error_message: str) -> None:
+        # Release the GPU only if the FAILED write lands. A terminated deploy's
+        # GPU was already released by terminate and may now belong to another.
+        async with self._db.acquire() as conn:
+            async with conn.transaction():
+                if await self._deploys.update_state(
+                    deploy_id, "FAILED", tx=conn,
+                    error_message=error_message, unless_in=ENDED_STATES,
+                ):
+                    await self._inventory.release_gpu(
+                        node_id, gpu_required, tx=conn,
+                    )

src/orchestration/models/model_deployment/worker.py

@@ -3,16 +3,13 @@
 
 from orchestration.scheduling.placement_scoring import score_node
 from orchestration.provisioning.engine.registry import get_adapter
+from orchestration.repositories.model_deployment_repo import ENDED_STATES
 
 log = logging.getLogger(__name__)
 
 MAX_PROVISION_RETRIES = 4
 PROVISION_WAIT_SECONDS = 40
 
-# A terminate or stop can land while provisioning is still running. The
-# worker's own state writes must not overwrite it.
-_ENDED_STATES = ("TERMINATING", "TERMINATED", "STOPPED")
-
 
 class DeploymentEndedError(Exception):
     pass
@@ -51,7 +48,7 @@ def __init__(
     # -------------------------------------------------
     async def _set_state(self, deployment_id: UUID, state: str, **kwargs) -> bool:
         return await self.deployments.update_state(
-            deployment_id, state, unless_in=_ENDED_STATES, **kwargs
+            deployment_id, state, unless_in=ENDED_STATES, **kwargs
         )
 
     async def _mark_running(self, deployment_id: UUID) -> None:

src/orchestration/repositories/model_deployment_repo.py

@@ -10,6 +10,10 @@
 from orchestration.config import settings
 from orchestration.repositories.base_repo import BaseRepository
 
+# A deploy in one of these states is being torn down; other writers must not
+# move it back to a live state.
+ENDED_STATES = ("TERMINATING", "TERMINATED", "STOPPED")
+
 _fernet = Fernet(settings.secret_encryption_key.encode()) if settings.secret_encryption_key else None
 
 

src/tests/orchestration/model_deployment/test_deployment_linker.py

@@ -442,3 +442,54 @@ async def test_already_bound_deploy_is_promoted_without_reallocate(pool):
     assert row["state"] == "RUNNING"
     assert node["gpu_allocated"] == 1            # NOT double-allocated
     controller.load_model.assert_awaited_once()  # model load fired
+
+
+def _terminate_during_load(pool, deploy_id, node_id, *, then=None):
+    async def load_model(**_kw):
+        async with pool.acquire() as c:
+            await c.execute(
+                "UPDATE model_deployments SET state='TERMINATED' "
+                "WHERE deployment_id=$1", deploy_id,
+            )
+            # Terminate released the GPU and another deploy took it.
+            await c.execute(
+                "UPDATE compute_inventory SET gpu_allocated=1 WHERE id=$1",
+                node_id,
+            )
+        if then:
+            raise then
+    return load_model
+
+
+@pytest.mark.parametrize("load_error", [None, RuntimeError("pull failed")])
+async def test_terminate_during_load_stays_terminated(pool, load_error):
+    pool_id, node_id = await _seed_pool_and_node(pool, gpu_total=1)
+    deploy_id = await _seed_pending_deploy(pool, pool_id)
+    controller = AsyncMock(spec=WorkerController)
+    controller.load_model.side_effect = _terminate_during_load(
+        pool, deploy_id, node_id, then=load_error,
+    )
+    linker = DeploymentLinker(
+        db_pool=pool,
+        inventory_repo=InventoryRepository(pool),
+        deployment_repo=ModelDeploymentRepository(pool, event_bus=None),
+        worker_controller=controller,
+    )
+
+    await linker.on_worker_ready(node_id)
+
+    async with pool.acquire() as c:
+        row = await c.fetchrow(
+            "SELECT state, endpoint FROM model_deployments "
+            "WHERE deployment_id=$1", deploy_id,
+        )
+        node = await c.fetchrow(
+            "SELECT gpu_allocated FROM compute_inventory WHERE id=$1", node_id,
+        )
+    assert row["state"] == "TERMINATED"
+    assert row["endpoint"] is None
+    assert node["gpu_allocated"] == 1
+    if load_error is None:
+        controller.unload_model.assert_awaited_once_with(
+            node_id=str(node_id), deployment_id=str(deploy_id),
+        )