diff --git a/plugins/nemo-deployments/src/nemo_deployments_plugin/backends/docker/backend.py b/plugins/nemo-deployments/src/nemo_deployments_plugin/backends/docker/backend.py index 43c410457d..b27f62b900 100644 --- a/plugins/nemo-deployments/src/nemo_deployments_plugin/backends/docker/backend.py +++ b/plugins/nemo-deployments/src/nemo_deployments_plugin/backends/docker/backend.py @@ -446,12 +446,16 @@ def _run_and_wait() -> int: async def read_status(self, *, workspace: str, name: str) -> BackendStatusUpdate: c_name = container_name(workspace, name) + dep_key = deployment_key(workspace, name) try: container = await asyncio.to_thread(self._client.containers.get, c_name) await asyncio.to_thread(container.reload) except self._docker_errors.NotFound: restart_policy = await self._resolve_restart_policy(workspace, name) - return missing_container_status(restart_policy, container_name=c_name) + status_update = missing_container_status(restart_policy, container_name=c_name) + if restart_policy in _ONE_SHOT_RESTART_POLICIES and self._gpu_pool is not None: + self._gpu_pool.release_gpu(dep_key) + return status_update except ( self._docker_errors.APIError, ReadTimeout, @@ -517,7 +521,6 @@ async def read_status(self, *, workspace: str, name: str) -> BackendStatusUpdate if state in ("exited", "dead"): exit_code = int(container.attrs.get("State", {}).get("ExitCode", 1)) if exit_code == 0 and restart_policy in ("Never", "OnFailure"): - dep_key = deployment_key(workspace, name) if self._gpu_pool is not None: self._gpu_pool.release_gpu(dep_key) return BackendStatusUpdate( @@ -543,7 +546,6 @@ async def read_status(self, *, workspace: str, name: str) -> BackendStatusUpdate exit_code=exit_code, endpoints=endpoints, ) - dep_key = deployment_key(workspace, name) if self._gpu_pool is not None: self._gpu_pool.release_gpu(dep_key) status = map_exited_status(exit_code, restart_policy) diff --git a/plugins/nemo-deployments/tests/unit/backends/docker/test_backend_mocked.py b/plugins/nemo-deployments/tests/unit/backends/docker/test_backend_mocked.py index 9be1747f83..3051cc2194 100644 --- a/plugins/nemo-deployments/tests/unit/backends/docker/test_backend_mocked.py +++ b/plugins/nemo-deployments/tests/unit/backends/docker/test_backend_mocked.py @@ -19,9 +19,11 @@ RESTART_POLICY_LABEL, companion_container_name, container_name, + deployment_key, ) from nemo_deployments_plugin.constants import MANAGED_BY_LABEL from nemo_deployments_plugin.entities import Deployment +from nemo_deployments_plugin.types import RestartPolicy @pytest.mark.asyncio @@ -455,10 +457,46 @@ async def get_side_effect(entity_type, name, workspace=None): return sample_config(restart_policy="Always") mock_entities.get.side_effect = get_side_effect + gpu_pool = MagicMock() + docker_backend._gpu_pool = gpu_pool update = await docker_backend.read_status(workspace="default", name="srv") assert update.status == "LOST" + gpu_pool.release_gpu.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("restart_policy", ["Never", "OnFailure"]) +async def test_read_status_failed_and_releases_gpu_when_missing_one_shot( + docker_backend: DockerDeploymentBackend, + mock_entities: AsyncMock, + mock_docker_client: MagicMock, + restart_policy: RestartPolicy, +) -> None: + expected_container_name = container_name("default", "job") + mock_docker_client.containers.get.side_effect = NotFound("missing") + + deployment_entity = MagicMock() + deployment_entity.deployment_config = "cfg1" + + async def get_side_effect(entity_type, name, workspace=None): + if entity_type is Deployment: + return deployment_entity + return sample_config(restart_policy=restart_policy) + + mock_entities.get.side_effect = get_side_effect + gpu_pool = MagicMock() + docker_backend._gpu_pool = gpu_pool + + update = await docker_backend.read_status(workspace="default", name="job") + + assert update.status == "FAILED" + assert update.exit_code is None + assert update.error_details == { + "expected_container_name": expected_container_name, + } + gpu_pool.release_gpu.assert_called_once_with(deployment_key("default", "job")) @pytest.mark.asyncio