fix(deployments): address CodeRabbit review on Docker backend (AIRCORE-756)

- Exclude already-assigned host ports during multi-port allocation
- Use label constants; fix HTTP/TCP probe host/port resolution
- Remove TYPE_CHECKING-only docker imports
- Deduplicate Docker availability check via docker_availability module
- Harden port unit tests; add type hints and cleanup improvements

Signed-off-by: Tyler Bray <tbray@nvidia.com>

Author: tylersbray

plugins/nemo-deployments/pyproject.toml

@@ -38,6 +38,6 @@ dev = ["pytest>=8.3.4", "pytest-asyncio>=0.25.3", "httpx>=0.27", "fastapi>=0.115
 [tool.pytest.ini_options]
 testpaths = ["tests"]
 asyncio_mode = "auto"
-pythonpath = ["src", "tests/unit"]
+pythonpath = ["src", "tests/unit", "tests/integration"]
 
 [tool.nemo.openapi]

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

@@ -32,9 +32,11 @@
 )
 from nemo_deployments_plugin.backends.docker.gpu import GPUAllocationError, get_shared_gpu_pool
 from nemo_deployments_plugin.backends.docker.labels import (
+    BACKOFF_LIMIT_LABEL,
     CONFIG_NAME_LABEL,
     DEPLOYMENT_NAME_LABEL,
     DEPLOYMENT_WORKSPACE_LABEL,
+    MANAGED_BY_KEY,
     RESTART_POLICY_LABEL,
     container_name,
     deployment_identity_labels,
@@ -147,7 +149,11 @@ async def create_deployment(
 
         host_ports: dict[int, int] = {}
         for port_spec in container_spec.ports:
-            host_port = await find_available_port(self._client, docker_cfg)
+            host_port = await find_available_port(
+                self._client,
+                docker_cfg,
+                exclude_ports=set(host_ports.values()),
+            )
             if host_port is None:
                 if gpu_ids:
                     self._gpu_pool.release_gpu(dep_key)  # type: ignore[union-attr]
@@ -289,7 +295,7 @@ async def read_status(self, *, workspace: str, name: str) -> BackendStatusUpdate
                 )
             if restart_policy == "OnFailure":
                 restart_count = int(container.attrs.get("RestartCount", 0))
-                backoff_limit = int(labels.get("nmp.nvidia.com/backoff-limit", "6"))
+                backoff_limit = int(labels.get(BACKOFF_LIMIT_LABEL, "6"))
                 if restart_count < backoff_limit:
                     return BackendStatusUpdate(
                         status="STARTING",
@@ -350,7 +356,7 @@ async def list_managed_deployment_names(self) -> list[str]:
         seen: set[str] = set()
         for container in containers:
             container_labels = container.labels or {}
-            if container_labels.get("managed-by") != MANAGED_BY_LABEL:
+            if container_labels.get(MANAGED_BY_KEY) != MANAGED_BY_LABEL:
                 continue
             ws = container_labels.get(DEPLOYMENT_WORKSPACE_LABEL)
             dep_name = container_labels.get(DEPLOYMENT_NAME_LABEL)
@@ -417,7 +423,7 @@ def _container_matches_deployment(
             labels.get(DEPLOYMENT_WORKSPACE_LABEL) == workspace
             and labels.get(DEPLOYMENT_NAME_LABEL) == name
             and labels.get(CONFIG_NAME_LABEL) == config_name
-            and labels.get("managed-by") == MANAGED_BY_LABEL
+            and labels.get(MANAGED_BY_KEY) == MANAGED_BY_LABEL
         )
 
     async def _resolve_restart_policy(self, workspace: str, name: str) -> RestartPolicy:

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

@@ -9,13 +9,12 @@
 import logging
 import os
 import socket
-from typing import TYPE_CHECKING, Any
+from typing import Any
 
 from nemo_deployments_plugin.backends.docker.labels import managed_by_filter
 from nemo_deployments_plugin.entities import DockerDeploymentConfig
 
-if TYPE_CHECKING:
-    import docker
+import docker
 
 logger = logging.getLogger(__name__)
 
@@ -53,7 +52,12 @@ def collect_used_host_ports(containers: list[Any]) -> set[int]:
     return used
 
 
-async def find_available_port(client: docker.DockerClient, docker_cfg: DockerDeploymentConfig) -> int | None:
+async def find_available_port(
+    client: docker.DockerClient,
+    docker_cfg: DockerDeploymentConfig,
+    *,
+    exclude_ports: set[int] | None = None,
+) -> int | None:
     try:
         containers = await asyncio.to_thread(
             client.containers.list,
@@ -65,6 +69,8 @@ async def find_available_port(client: docker.DockerClient, docker_cfg: DockerDep
         return None
 
     used_ports = collect_used_host_ports(containers)
+    if exclude_ports:
+        used_ports = used_ports | exclude_ports
     for port in range(docker_cfg.port_range_start, docker_cfg.port_range_end + 1):
         if port not in used_ports and is_port_free(port):
             return port

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

@@ -8,15 +8,12 @@
 import asyncio
 import logging
 import socket
-from typing import TYPE_CHECKING
-from urllib.parse import urljoin
+from urllib.parse import urljoin, urlparse
 
 import httpx
+from docker.models.containers import Container as DockerContainer
 from nemo_deployments_plugin.entities import Probe
 
-if TYPE_CHECKING:
-    from docker.models.containers import Container as DockerContainer
-
 logger = logging.getLogger(__name__)
 
 
@@ -39,7 +36,7 @@ async def check_readiness_probe(
         return await _check_http_probe(host_url, probe, host_ports=host_ports, named_ports=named_ports)
 
     if probe.tcp_socket is not None and host_url is not None:
-        return await _check_tcp_probe(host_url, probe)
+        return await _check_tcp_probe(host_url, probe, host_ports=host_ports, named_ports=named_ports)
 
     return True, "probe type not implemented; treating as ready"
 
@@ -75,6 +72,27 @@ def _run() -> tuple[int, str]:
     return False, f"exec probe exit {exit_code}: {output[:200]}"
 
 
+def _probe_host(host_url: str) -> str:
+    return urlparse(host_url).hostname or "127.0.0.1"
+
+
+def _resolve_probe_port(
+    port: int | str,
+    *,
+    host_ports: dict[int, int] | None,
+    named_ports: dict[str, int] | None,
+) -> int | None:
+    if isinstance(port, int):
+        return host_ports.get(port, port) if host_ports else port
+    if named_ports and port in named_ports:
+        return named_ports[port]
+    if host_ports:
+        for container_port, host_port in host_ports.items():
+            if str(container_port) == port:
+                return host_port
+    return None
+
+
 async def _check_http_probe(
     host_url: str,
     probe: Probe,
@@ -86,16 +104,11 @@ async def _check_http_probe(
     port = probe.http_get.port
     path = probe.http_get.path
     scheme = probe.http_get.scheme.lower()
-    base = host_url
-    if isinstance(port, int):
-        base = f"{scheme}://127.0.0.1:{port}"
-    elif isinstance(port, str) and named_ports and port in named_ports:
-        base = f"{scheme}://127.0.0.1:{named_ports[port]}"
-    elif isinstance(port, str) and host_ports:
-        for container_port, host_port in host_ports.items():
-            if str(container_port) == port:
-                base = f"{scheme}://127.0.0.1:{host_port}"
-                break
+    host = _probe_host(host_url)
+    resolved_port = _resolve_probe_port(port, host_ports=host_ports, named_ports=named_ports)
+    if resolved_port is None:
+        return False, f"http probe port not mapped: {port}"
+    base = f"{scheme}://{host}:{resolved_port}"
     url = urljoin(f"{base.rstrip('/')}/", path.lstrip("/"))
     timeout = probe.timeout_seconds
     try:
@@ -108,16 +121,23 @@ async def _check_http_probe(
         return False, f"http probe failed: {exc}"
 
 
-async def _check_tcp_probe(host_url: str, probe: Probe) -> tuple[bool, str]:
+async def _check_tcp_probe(
+    host_url: str,
+    probe: Probe,
+    *,
+    host_ports: dict[int, int] | None = None,
+    named_ports: dict[str, int] | None = None,
+) -> tuple[bool, str]:
     assert probe.tcp_socket is not None
     port_value = probe.tcp_socket.port
-    if not isinstance(port_value, int):
-        return False, "tcp probe requires numeric port"
-    host = "127.0.0.1"
+    host = _probe_host(host_url)
+    target_port = _resolve_probe_port(port_value, host_ports=host_ports, named_ports=named_ports)
+    if target_port is None:
+        return False, f"tcp probe port not mapped: {port_value}"
     timeout = probe.timeout_seconds
 
     def _connect() -> None:
-        with socket.create_connection((host, port_value), timeout=timeout):
+        with socket.create_connection((host, target_port), timeout=timeout):
             return
 
     try:

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

@@ -26,10 +26,14 @@ def format_duration(seconds: float) -> str:
 
 
 def map_exited_status(exit_code: int, restart_policy: RestartPolicy) -> DeploymentStatus:
+    """Map a stopped container's exit code to a terminal deployment status.
+
+    Restart-policy-specific retry handling happens in the backend before this
+    helper is called for non-zero exits.
+    """
+    del restart_policy
     if exit_code in SUCCESSFUL_EXIT_CODES:
         return "SUCCEEDED"
-    if restart_policy == "Always":
-        return "FAILED"
     return "FAILED"
 
 

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

@@ -7,13 +7,12 @@
 
 import asyncio
 import logging
-from typing import TYPE_CHECKING, Any
+from typing import Any
 
 from nemo_deployments_plugin.backends.base import VolumeStatusUpdate
 from nemo_deployments_plugin.backends.docker.labels import docker_volume_name, volume_identity_labels
 
-if TYPE_CHECKING:
-    import docker
+import docker
 
 logger = logging.getLogger(__name__)
 

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

@@ -10,8 +10,9 @@
 
 import pytest
 from docker.errors import NotFound
+from docker_availability import skip_without_docker
 from nemo_deployments_plugin.backends.docker.backend import DockerDeploymentBackend
-from nemo_deployments_plugin.backends.docker.labels import container_name, docker_volume_name
+from nemo_deployments_plugin.backends.docker.labels import container_name
 from nemo_deployments_plugin.backends.registry import BACKEND_CLASSES
 from nemo_deployments_plugin.constants import MANAGED_BY_LABEL
 from nemo_deployments_plugin.entities import (
@@ -23,17 +24,11 @@
     DockerDeploymentConfig,
 )
 
-try:
-    import docker
-
-    docker.from_env().ping()
-    _DOCKER_AVAILABLE = True
-except Exception:
-    _DOCKER_AVAILABLE = False
+import docker
 
 pytestmark = [
     pytest.mark.skipif("docker" not in BACKEND_CLASSES, reason="Docker backend not registered"),
-    pytest.mark.skipif(not _DOCKER_AVAILABLE, reason="Docker daemon not available"),
+    skip_without_docker,
 ]
 
 
@@ -101,6 +96,7 @@ async def test_never_deployment_succeeds(docker_backend: DockerDeploymentBackend
     config = _never_config()
     docker_backend._entities.get.return_value = config  # type: ignore[attr-defined]
     c_name = container_name("itest", "echo-job")
+    client = docker.from_env()
 
     try:
         created = await docker_backend.create_deployment(
@@ -123,7 +119,7 @@ async def test_never_deployment_succeeds(docker_backend: DockerDeploymentBackend
     finally:
         await docker_backend.delete_deployment("itest", "echo-job")
         try:
-            docker_backend._client.containers.get(c_name).remove(force=True)
+            client.containers.get(c_name).remove(force=True)
         except NotFound:
             pass
 
@@ -140,6 +136,7 @@ async def get_side_effect(entity_type, name, workspace=None):
 
     docker_backend._entities.get.side_effect = get_side_effect  # type: ignore[attr-defined]
     c_name = container_name("itest", "lost-srv")
+    client = docker.from_env()
 
     try:
         created = await docker_backend.create_deployment(
@@ -151,14 +148,14 @@ async def get_side_effect(entity_type, name, workspace=None):
         )
         assert created.status == "STARTING"
 
-        container = docker_backend._client.containers.get(c_name)
+        container = client.containers.get(c_name)
         container.remove(force=True)
 
         status = await docker_backend.read_status(workspace="itest", name="lost-srv")
         assert status.status == "LOST"
     finally:
         await docker_backend.delete_deployment("itest", "lost-srv")
         try:
-            docker_backend._client.volumes.remove(docker_volume_name("itest", "data"))
-        except Exception:
+            client.containers.get(c_name).remove(force=True)
+        except NotFound:
             pass

plugins/nemo-deployments/tests/integration/conftest.py

@@ -5,14 +5,6 @@
 
 from __future__ import annotations
 
-import pytest
+from docker_availability import DOCKER_AVAILABLE, skip_without_docker
 
-try:
-    import docker
-
-    docker.from_env().ping()
-    DOCKER_AVAILABLE = True
-except Exception:
-    DOCKER_AVAILABLE = False
-
-skip_without_docker = pytest.mark.skipif(not DOCKER_AVAILABLE, reason="Docker daemon not available")
+__all__ = ["DOCKER_AVAILABLE", "skip_without_docker"]

plugins/nemo-deployments/tests/integration/docker_availability.py

@@ -0,0 +1,20 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Shared Docker daemon availability check for integration tests."""
+
+from __future__ import annotations
+
+import pytest
+
+try:
+    import docker
+
+    docker.from_env().ping()
+    DOCKER_AVAILABLE: bool = True
+except Exception:
+    DOCKER_AVAILABLE = False
+
+skip_without_docker: pytest.MarkDecorator = pytest.mark.skipif(
+    not DOCKER_AVAILABLE, reason="Docker daemon not available"
+)