fix(deployments): handle Docker restart races (#868)

Update Docker deployment status reconciliation to observe Docker inspect state for exited and missing containers.

Stopped one-shot containers remain observable, retryable OnFailure exits stay non-terminal until the backoff limit is exhausted, and Docker GPU allocations are released only when status handling reaches a terminal deployment result.

Signed-off-by: Ryan S <267728323+ironcommit@users.noreply.github.com>

Author: ironcommit

packages/nmp_customization_common/src/nmp/customization_common/config.py

@@ -7,7 +7,7 @@
 from pydantic import Field
 
 
-class CustomizationCommonConfig(create_service_config_class("customizer")):  # type: ignore[misc]
+class CustomizationCommonConfig(create_service_config_class("customizer")):  # ty: ignore[unsupported-base]
     """Environment variables use the ``NMP_CUSTOMIZER_`` prefix."""
 
     tasks_image: str | None = Field(

plugins/nemo-agents/examples/calculator-agent/src/calculator_agent/register.py

@@ -8,10 +8,10 @@
 
 from collections.abc import AsyncGenerator
 
-from nat.builder.builder import Builder  # type: ignore
-from nat.builder.function import FunctionGroup  # type: ignore
-from nat.cli.register_workflow import register_function_group  # type: ignore
-from nat.data_models.function import FunctionGroupBaseConfig  # type: ignore
+from nat.builder.builder import Builder
+from nat.builder.function import FunctionGroup
+from nat.cli.register_workflow import register_function_group
+from nat.data_models.function import FunctionGroupBaseConfig
 from pydantic import Field
 
 

plugins/nemo-deployments/openapi/openapi.yaml

@@ -78,7 +78,7 @@
     SecretResolutionError,
     resolve_deployment_config_secrets,
 )
-from nemo_deployments_plugin.types import Endpoint, RestartPolicy
+from nemo_deployments_plugin.types import NON_TERMINAL_DEPLOYMENT_STATUSES, Endpoint, RestartPolicy
 from nemo_platform_plugin.capabilities import docker_from_env_kwargs, probe_docker
 from nemo_platform_plugin.client.adapter import client_from_platform
 from nemo_platform_plugin.config import LOOPBACK_ADDRESSES
@@ -139,6 +139,21 @@ def _config_files_tar(config_files: list[ConfigFile]) -> bytes:
     return buf.getvalue()
 
 
+def _docker_inspect_attrs(container: DockerContainer) -> dict[str, Any]:
+    return container.attrs or {}
+
+
+def _docker_inspect_exit_code(container: DockerContainer) -> int:
+    state = _docker_inspect_attrs(container).get("State") or {}
+    if not isinstance(state, dict):
+        return 1
+    return int(state.get("ExitCode", 1))
+
+
+def _docker_inspect_restart_count(container: DockerContainer) -> int:
+    return int(_docker_inspect_attrs(container).get("RestartCount", 0))
+
+
 class DockerDeploymentBackend(DeploymentBackend):
     """Manage deployments and volumes as Docker containers and volumes."""
 
@@ -676,12 +691,15 @@ async def read_status(self, *, workspace: str, name: str) -> BackendStatusUpdate
             container = await asyncio.to_thread(self._client.containers.get, c_name)
             if not self._container_matches_deployment_group(container, workspace, name):
                 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 status_update.status not in NON_TERMINAL_DEPLOYMENT_STATUSES and self._gpu_pool is not None:
+                    self._gpu_pool.release_gpu(dep_key)
+                return status_update
             await asyncio.to_thread(container.reload)
         except self._docker_errors.NotFound:
             restart_policy = await self._resolve_restart_policy(workspace, 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:
+            if status_update.status not in NON_TERMINAL_DEPLOYMENT_STATUSES and self._gpu_pool is not None:
                 self._gpu_pool.release_gpu(dep_key)
             return status_update
         except (
@@ -751,8 +769,8 @@ 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))
-            restart_count = int(container.attrs.get("RestartCount", 0))
+            exit_code = _docker_inspect_exit_code(container)
+            restart_count = _docker_inspect_restart_count(container)
             return self._status_from_exited_container(
                 exit_code=exit_code,
                 restart_policy=restart_policy,
@@ -805,6 +823,13 @@ def _status_from_exited_container(
             )
         if restart_policy == "OnFailure":
             backoff_limit = int(labels.get(BACKOFF_LIMIT_LABEL, "6"))
+            if backoff_limit == 0:
+                return BackendStatusUpdate(
+                    status="STARTING",
+                    status_message=(f"Container exited (code {exit_code}); retry {restart_count}/unlimited"),
+                    exit_code=exit_code,
+                    endpoints=resolved_endpoints,
+                )
             if restart_count < backoff_limit:
                 return BackendStatusUpdate(
                     status="STARTING",
@@ -887,9 +912,8 @@ def _wait_for_exit() -> int:
                 endpoints=endpoints,
             )
         if container.status in _EXITED_CONTAINER_STATES:
-            attrs = container.attrs or {}
-            exit_code = int(attrs.get("State", {}).get("ExitCode", 1))
-            restart_count = int(attrs.get("RestartCount", 0))
+            exit_code = _docker_inspect_exit_code(container)
+            restart_count = _docker_inspect_restart_count(container)
             return self._status_from_exited_container(
                 exit_code=exit_code,
                 restart_policy=restart_policy,

plugins/nemo-deployments/src/nemo_deployments_plugin/backends/docker/backend.py

@@ -95,6 +95,8 @@ def restart_policy_kwargs(restart_policy: RestartPolicy, backoff_limit: int) ->
     if restart_policy == "Always":
         return {"restart_policy": {"Name": "always"}}
     if restart_policy == "OnFailure":
+        if backoff_limit < 1:
+            raise DeploymentConfigError("backoff_limit must be at least 1 for OnFailure restart policy")
         return {"restart_policy": {"Name": "on-failure", "MaximumRetryCount": backoff_limit}}
     return {}
 

plugins/nemo-deployments/src/nemo_deployments_plugin/backends/docker/containers.py

@@ -285,7 +285,12 @@ class DeploymentConfig(NemoEntity, entity_type=ENTITY_TYPE_DEPLOYMENT_CONFIG):
     volume_mounts: list[VolumeMount] = Field(default_factory=list, alias="volumeMounts")
     config_files: list[ConfigFile] = Field(default_factory=list, alias="configFiles")
     restart_policy: RestartPolicy = Field(default="Always", alias="restartPolicy")
-    backoff_limit: int = Field(default=6, alias="backoffLimit")
+    backoff_limit: int = Field(
+        default=6,
+        ge=1,
+        alias="backoffLimit",
+        description="Retry limit for OnFailure deployments; must be positive. Never deployments disable retries internally.",
+    )
     drift_recovery: DriftRecoveryPolicy = Field(default_factory=DriftRecoveryPolicy, alias="driftRecovery")
     labels: dict[str, str] = Field(default_factory=dict)
     backend_config: DeploymentBackendConfig = Field(default_factory=DeploymentBackendConfig, alias="backendConfig")

plugins/nemo-deployments/src/nemo_deployments_plugin/entities.py

@@ -102,7 +102,11 @@ class CreateDeploymentConfigRequest(BaseModel):
     volume_mounts: list[VolumeMount] = Field(default_factory=list)
     config_files: list[ConfigFile] = Field(default_factory=list)
     restart_policy: RestartPolicy = "Always"
-    backoff_limit: int = 6
+    backoff_limit: int = Field(
+        default=6,
+        ge=1,
+        description="Retry limit for OnFailure deployments; must be positive. Never deployments disable retries internally.",
+    )
     drift_recovery: DriftRecoveryPolicy | None = None
     labels: dict[str, str] = Field(default_factory=dict)
     backend_config: DeploymentBackendConfig = Field(default_factory=DeploymentBackendConfig)

plugins/nemo-deployments/src/nemo_deployments_plugin/schema.py

@@ -49,8 +49,19 @@ def _worker_port_base(worker_id: str) -> int:
     return TEST_PORT_RANGE_START + worker_index * TEST_PORT_RANGE_SIZE
 
 
-def _build_docker_backend(worker_id: str = "master", **config_overrides: Any) -> DockerDeploymentBackend:
-    mock_entities = AsyncMock()
+@pytest.fixture
+def mock_entities() -> AsyncMock:
+    return AsyncMock()
+
+
+def _build_docker_backend(
+    worker_id: str = "master",
+    *,
+    mock_entities: AsyncMock | None = None,
+    **config_overrides: Any,
+) -> DockerDeploymentBackend:
+    if mock_entities is None:
+        mock_entities = AsyncMock()
     mock_sdk = MagicMock()
     port_base = _worker_port_base(worker_id)
     executor_config: dict[str, Any] = {
@@ -70,8 +81,8 @@ def _build_docker_backend(worker_id: str = "master", **config_overrides: Any) ->
 
 
 @pytest.fixture
-def docker_backend(worker_id: str) -> DockerDeploymentBackend:
-    return _build_docker_backend(worker_id)
+def docker_backend(worker_id: str, mock_entities: AsyncMock) -> DockerDeploymentBackend:
+    return _build_docker_backend(worker_id, mock_entities=mock_entities)
 
 
 def _never_config() -> DeploymentConfig:
@@ -148,9 +159,12 @@ async def test_volume_lifecycle(docker_backend: DockerDeploymentBackend) -> None
 
 
 @pytest.mark.asyncio
-async def test_never_deployment_succeeds(docker_backend: DockerDeploymentBackend) -> None:
+async def test_never_deployment_succeeds(
+    docker_backend: DockerDeploymentBackend,
+    mock_entities: AsyncMock,
+) -> None:
     config = _never_config()
-    docker_backend._entities.get.return_value = config  # ty: ignore[unresolved-attribute]
+    mock_entities.get.return_value = config
     c_name = container_name("itest", "echo-job")
     client = docker.from_env()
 
@@ -189,7 +203,14 @@ async def test_never_deployment_outlives_observe_wait_then_succeeds(worker_id: s
         oneshot_observe_timeout_seconds=observe_timeout_seconds,
     )
     config = _never_sleep_config(sleep_seconds=job_sleep_seconds)
-    docker_backend._entities.get.return_value = config  # ty: ignore[unresolved-attribute]
+    deployment = Deployment(name="sleep-job", workspace="itest", deployment_config="sleep-cfg")
+
+    async def get_side_effect(entity_type, name, workspace=None):
+        if entity_type is Deployment:
+            return deployment
+        return config
+
+    docker_backend._entities.get.side_effect = get_side_effect  # ty: ignore[unresolved-attribute]
     c_name = container_name("itest", "sleep-job")
     client = docker.from_env()
 
@@ -233,7 +254,55 @@ async def test_never_deployment_outlives_observe_wait_then_succeeds(worker_id: s
 
 
 @pytest.mark.asyncio
-async def test_lost_detection_for_always(docker_backend: DockerDeploymentBackend) -> None:
+async def test_removed_never_deployment_reports_missing(
+    docker_backend: DockerDeploymentBackend,
+    mock_entities: AsyncMock,
+) -> None:
+    config = _never_config()
+    deployment = Deployment(name="removed-job", workspace="itest", deployment_config="echo-cfg")
+
+    async def get_side_effect(entity_type, name, workspace=None):
+        if entity_type is Deployment:
+            return deployment
+        return config
+
+    mock_entities.get.side_effect = get_side_effect
+    c_name = container_name("itest", "removed-job")
+    client = docker.from_env()
+
+    try:
+        created = await docker_backend.create_deployment(
+            workspace="itest",
+            name="removed-job",
+            config_name="echo-cfg",
+            labels={"managed-by": MANAGED_BY_LABEL},
+            backend_config={},
+        )
+        assert created.status == "SUCCEEDED"
+        assert created.exit_code == 0
+
+        container = client.containers.get(c_name)
+        result = container.wait(timeout=20)
+        assert result["StatusCode"] == 0
+        container.remove(force=True)
+
+        status = await docker_backend.read_status(workspace="itest", name="removed-job")
+
+        assert status.status == "FAILED"
+        assert status.exit_code is None
+        assert status.error_details == {
+            "expected_container_name": c_name,
+        }
+    finally:
+        await docker_backend.delete_deployment("itest", "removed-job")
+        force_remove_container(client, c_name)
+
+
+@pytest.mark.asyncio
+async def test_lost_detection_for_always(
+    docker_backend: DockerDeploymentBackend,
+    mock_entities: AsyncMock,
+) -> None:
     deployment = Deployment(name="lost-srv", workspace="itest", deployment_config="http-cfg")
     config = _always_http_config()
 
@@ -242,7 +311,7 @@ async def get_side_effect(entity_type, name, workspace=None):
             return deployment
         return config
 
-    docker_backend._entities.get.side_effect = get_side_effect  # ty: ignore[unresolved-attribute]
+    mock_entities.get.side_effect = get_side_effect
     c_name = container_name("itest", "lost-srv")
     client = docker.from_env()
 

plugins/nemo-deployments/tests/integration/backends/docker/test_docker_backend.py

@@ -119,9 +119,9 @@ def lora_config(*, restart_policy: RestartPolicy = "Always") -> DeploymentConfig
     )
 
 
-def container_attrs(*, status: str = "running", exit_code: int = 0) -> dict[str, Any]:
+def container_attrs(*, status: str = "running", exit_code: int = 0, restart_count: int = 0) -> dict[str, Any]:
     del status
     return {
         "State": {"ExitCode": exit_code, "StartedAt": "2026-01-01T00:00:00Z"},
-        "RestartCount": 0,
+        "RestartCount": restart_count,
     }