diff --git a/src/orchestration/models/model_deployment/deployment_linker.py b/src/orchestration/models/model_deployment/deployment_linker.py index 2b3095d0..5c5ce735 100644 --- a/src/orchestration/models/model_deployment/deployment_linker.py +++ b/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, + ) diff --git a/src/orchestration/models/model_deployment/worker.py b/src/orchestration/models/model_deployment/worker.py index e3933426..87d2c37a 100644 --- a/src/orchestration/models/model_deployment/worker.py +++ b/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: diff --git a/src/orchestration/repositories/model_deployment_repo.py b/src/orchestration/repositories/model_deployment_repo.py index a5cad60f..7e961b9f 100644 --- a/src/orchestration/repositories/model_deployment_repo.py +++ b/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 diff --git a/src/tests/orchestration/model_deployment/test_deployment_linker.py b/src/tests/orchestration/model_deployment/test_deployment_linker.py index 1783e539..cb27b77f 100644 --- a/src/tests/orchestration/model_deployment/test_deployment_linker.py +++ b/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), + )