fix(nemo-deployments): release docker gpu for missing one-shot containers

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)

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