diff --git a/docs/set-up/config-reference.mdx b/docs/set-up/config-reference.mdx index 959972a5b7..0d49c872e9 100644 --- a/docs/set-up/config-reference.mdx +++ b/docs/set-up/config-reference.mdx @@ -493,6 +493,8 @@ models: model_puller_max_workers: 1 # Number of retries when the puller fails with a transient error (IncompleteRead, ChunkedEncodingError, connection broken). Same volume is reused so partial downloads can be completed. | default: 3 model_puller_retries: 3 + # Additional labels copied onto model resources managed by the models controller. + model_labels: {} # BusyBox image repository used for helper containers (permissions/find/chown). | default: 'busybox' busybox_image: busybox # BusyBox image tag used for helper containers. | default: 'latest' @@ -577,6 +579,8 @@ models: default_tolerations: # Default Kubernetes node selector for all NIM deployments. Can be overridden per-deployment via k8s_nim_operator_config. Example: {'node-type': 'gpu-node', 'zone': 'us-west1-a'} default_node_selector: + # Additional labels copied onto model resources managed by the models controller. + model_labels: # Default Kubernetes labels applied to NIMService and NIMCache resources and their child resources (e.g. pods). Merged with controller-managed labels; controller labels take precedence on conflict. Example: {'team': 'ml-platform', 'environment': 'prod'} default_labels: # Default Kubernetes annotations applied to NIMService and NIMCache resources and their child resources (e.g. pods, PVCs). Merged with controller-managed annotations; controller annotations take precedence on conflict. Example: {'prometheus.io/scrape': 'true'} diff --git a/packages/nmp_testing/src/nmp/testing/docker.py b/packages/nmp_testing/src/nmp/testing/docker.py index 65803a00f6..c0929d7607 100644 --- a/packages/nmp_testing/src/nmp/testing/docker.py +++ b/packages/nmp_testing/src/nmp/testing/docker.py @@ -304,26 +304,38 @@ def ensure_mock_sidecar_image(docker_client: docker.DockerClient, image_name: st MODELS_CONTROLLER_MANAGED_LABEL = "nmp.nvidia.com/managed-by=models-controller" -def cleanup_model_deployment_containers(docker_client: docker.DockerClient) -> int: - """Stop and remove all containers managed by the models controller. +def _matches_labels(container_labels: dict[str, str], labels: dict[str, str]) -> bool: + """Return True when all expected Docker labels are present.""" + return all(container_labels.get(key) == value for key, value in labels.items()) + + +def cleanup_model_deployment_containers( + docker_client: docker.DockerClient, + labels: dict[str, str] | None = None, +) -> int: + """Stop and remove models-controller containers matching optional owner labels. Finds containers with label nmp.nvidia.com/managed-by=models-controller - (NIM and sidecar containers created by the Docker backend), stops and - removes them. Intended for integration test teardown so failed tests - don't leave stuck containers; use as a pytest fixture teardown. + (NIM, sidecar, and puller containers created by the Docker backend), stops + and removes them. When labels are supplied, only containers matching all + owner labels are removed. Intended for integration test teardown so failed + tests don't leave stuck containers; use as a pytest fixture teardown. Uses the same retry logic as DockerTestContext for DinD compatibility. Args: docker_client: Docker client to use. + labels: Optional labels used to scope cleanup to this test owner. Returns: Number of containers removed. """ + labels = labels or {} try: containers = docker_client.containers.list( all=True, filters={"label": MODELS_CONTROLLER_MANAGED_LABEL}, + ignore_removed=True, ) except Exception: return 0 @@ -333,6 +345,9 @@ def cleanup_model_deployment_containers(docker_client: docker.DockerClient) -> i for container in containers: try: + if labels and not _matches_labels(container.labels or {}, labels): + continue + name = container.name @retry( @@ -500,10 +515,15 @@ def print_diagnostics(self) -> None: # Print relevant containers print("\n--- Docker PS (test containers) ---") try: - for c in self.docker_client.containers.list(all=True): + for c in self.docker_client.containers.list(all=True, ignore_removed=True): labels = c.labels if labels.get("nmp.nvidia.com/managed-by") == "models-controller": - print(f" {c.name}: {c.status}") + owner = { + key: value + for key, value in labels.items() + if key in {"nmp.nvidia.com/test-run", "nmp.nvidia.com/test-worker"} + } + print(f" {c.name}: {c.status} labels={owner}") except Exception as e: print(f"Could not list containers: {e}") diff --git a/packages/nmp_testing/tests/unit/test_docker.py b/packages/nmp_testing/tests/unit/test_docker.py new file mode 100644 index 0000000000..ad8c36d4df --- /dev/null +++ b/packages/nmp_testing/tests/unit/test_docker.py @@ -0,0 +1,54 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for Docker testing helpers.""" + +from unittest.mock import MagicMock + +from nmp.testing.docker import MODELS_CONTROLLER_MANAGED_LABEL, cleanup_model_deployment_containers + + +def _container(name: str, labels: dict[str, str]) -> MagicMock: + container = MagicMock() + container.name = name + container.labels = labels + return container + + +def test_cleanup_model_deployment_containers_filters_by_owner_labels(): + """Cleanup removes only managed containers matching the supplied owner labels.""" + owner_labels = { + "nmp.nvidia.com/test-run": "run-1", + "nmp.nvidia.com/test-worker": "gw0", + } + matching_container = _container( + "matching", + { + MODELS_CONTROLLER_MANAGED_LABEL: "models-controller", + **owner_labels, + }, + ) + other_worker_container = _container( + "other-worker", + { + MODELS_CONTROLLER_MANAGED_LABEL: "models-controller", + "nmp.nvidia.com/test-run": "run-1", + "nmp.nvidia.com/test-worker": "gw1", + }, + ) + + docker_client = MagicMock() + docker_client.containers.list.return_value = [matching_container, other_worker_container] + + removed = cleanup_model_deployment_containers(docker_client, labels=owner_labels) + + assert removed == 1 + docker_client.containers.list.assert_called_once_with( + all=True, + filters={"label": MODELS_CONTROLLER_MANAGED_LABEL}, + ignore_removed=True, + ) + matching_container.stop.assert_called_once() + matching_container.remove.assert_called_once_with(force=True) + other_worker_container.stop.assert_not_called() + other_worker_container.remove.assert_not_called() diff --git a/services/core/inference-gateway/tests/integration/conftest.py b/services/core/inference-gateway/tests/integration/conftest.py index 6b3ec18dd4..7c0008c21c 100644 --- a/services/core/inference-gateway/tests/integration/conftest.py +++ b/services/core/inference-gateway/tests/integration/conftest.py @@ -220,18 +220,34 @@ def docker_test_context( @pytest.fixture -def models_controller_container_cleanup(docker_client: docker.DockerClient) -> Generator[None, None, None]: +def docker_owner_labels(worker_id: str, testrun_uid: str) -> dict[str, str]: + """Labels used to scope Docker resources to this pytest worker/run.""" + return { + "nmp.nvidia.com/test-run": testrun_uid, + "nmp.nvidia.com/test-worker": worker_id, + } + + +@pytest.fixture +def models_controller_container_cleanup( + docker_client: docker.DockerClient, + docker_owner_labels: dict[str, str], +) -> Generator[None, None, None]: """Teardown: remove all containers with label nmp.nvidia.com/managed-by=models-controller. Ensures failed tests (e.g. stuck in PENDING) don't leave NIM/sidecar containers. Request this via controller_with_docker_and_igw; no per-test try/finally needed. """ yield - cleanup_model_deployment_containers(docker_client) + cleanup_model_deployment_containers(docker_client, labels=docker_owner_labels) @pytest.fixture -def docker_backend_config(worker_id: str, mock_sidecar_image: str) -> dict[str, Any]: +def docker_backend_config( + worker_id: str, + mock_sidecar_image: str, + docker_owner_labels: dict[str, str], +) -> dict[str, Any]: """Configuration for Docker backend in tests. Uses worker_id from pytest-xdist to allocate unique port ranges per worker. @@ -244,6 +260,7 @@ def docker_backend_config(worker_id: str, mock_sidecar_image: str) -> dict[str, "models_docker_port_range_end": end_port, "docker_timeout": 60, "models_docker_host_service_name": "localhost", + "model_labels": docker_owner_labels, } diff --git a/services/core/jobs/src/nmp/core/jobs/controllers/backends/kubernetes/kubernetes_job.py b/services/core/jobs/src/nmp/core/jobs/controllers/backends/kubernetes/kubernetes_job.py index a23367760b..dcac06443b 100644 --- a/services/core/jobs/src/nmp/core/jobs/controllers/backends/kubernetes/kubernetes_job.py +++ b/services/core/jobs/src/nmp/core/jobs/controllers/backends/kubernetes/kubernetes_job.py @@ -311,7 +311,26 @@ def create_step_update(self, step: PlatformJobStepWithContext, job: V1Job) -> Jo error_details["message"] = status_details.get("message", "Job encountered an error") status_details["events"] = self.get_kube_job_events(job) task_has_error = update_all_tasks(self._nmp_sdk, self._core_v1, self.namespace, step) - if task_has_error: + teardown_lifecycle_statuses = { + PlatformJobStatus.PAUSING, + PlatformJobStatus.PAUSED, + PlatformJobStatus.CANCELLING, + PlatformJobStatus.CANCELLED, + } + # Task-level errors can appear while Kubernetes is tearing down pods for + # a requested pause or cancel. Preserve those user-requested lifecycle + # states so the dispatcher can finish transitioning to PAUSED/CANCELLED. + if task_has_error and status in teardown_lifecycle_statuses: + logger.debug( + "Task error observed during container teardown", + extra={ + "workspace": step.workspace, + "job": step.job, + "step": step.name, + "status": status, + }, + ) + elif task_has_error: status = PlatformJobStatus.ERROR if "message" not in error_details: error_details["message"] = "One or more tasks are in error state" diff --git a/services/core/jobs/tests/controllers/test_kubernetes_backend.py b/services/core/jobs/tests/controllers/test_kubernetes_backend.py index f2c42721da..c0c1f75391 100644 --- a/services/core/jobs/tests/controllers/test_kubernetes_backend.py +++ b/services/core/jobs/tests/controllers/test_kubernetes_backend.py @@ -1048,6 +1048,45 @@ def test_sync_job_paused_with_errored_pods_from_sigterm(kubernetes_job, test_ste assert job_update.status == PlatformJobStatus.PAUSED +def test_sync_job_paused_ignores_task_errors_from_suspend(kubernetes_job, test_step_pending): + """Task-level pod errors observed during suspension must not override PAUSED.""" + mock_job_spec = MagicMock() + mock_job_spec.suspend = True + + mock_job_status = MagicMock() + mock_job_status.active = None + mock_job_status.succeeded = None + mock_job_status.failed = None + mock_job_status.terminating = None + mock_job_status.completion_time = None + + mock_job = MagicMock() + mock_job.status = mock_job_status + mock_job.spec = mock_job_spec + kubernetes_job._batch_v1.read_namespaced_job.return_value = mock_job + + with ( + patch("nmp.core.jobs.controllers.backends.kubernetes.kubernetes_job.list_pod_status") as mock_list_pod_status, + patch("nmp.core.jobs.controllers.backends.kubernetes.kubernetes_job.update_all_tasks") as mock_update_all_tasks, + ): + mock_list_pod_status.return_value = [ + PodStatus( + task_id="test-task", + name="test-pod", + errors={"test-container": 137}, + completed=set(), + active=set(), + waiting={}, + phase="Failed", + ) + ] + mock_update_all_tasks.return_value = True + + job_update = kubernetes_job.sync(test_step_pending) + + assert job_update.status == PlatformJobStatus.PAUSED + + def test_sync_job_pausing_with_errored_pods_from_sigterm(kubernetes_job, test_step_pending): """Test that a suspended job with both running and errored pods reports PAUSING, not ERROR. @@ -1136,6 +1175,45 @@ def test_sync_job_cancelling_with_errored_pods(kubernetes_job, test_step_cancell assert job_update.status == PlatformJobStatus.CANCELLED +def test_sync_job_cancelled_ignores_task_errors_from_termination(kubernetes_job, test_step_cancelling): + """Task-level pod errors observed during termination must not override CANCELLED.""" + mock_job_spec = MagicMock() + mock_job_spec.suspend = False + + mock_job_status = MagicMock() + mock_job_status.active = None + mock_job_status.succeeded = None + mock_job_status.failed = None + mock_job_status.terminating = None + mock_job_status.completion_time = None + + mock_job = MagicMock() + mock_job.status = mock_job_status + mock_job.spec = mock_job_spec + kubernetes_job._batch_v1.read_namespaced_job.return_value = mock_job + + with ( + patch("nmp.core.jobs.controllers.backends.kubernetes.kubernetes_job.list_pod_status") as mock_list_pod_status, + patch("nmp.core.jobs.controllers.backends.kubernetes.kubernetes_job.update_all_tasks") as mock_update_all_tasks, + ): + mock_list_pod_status.return_value = [ + PodStatus( + task_id="test-task", + name="test-pod", + errors={"test-container": 137}, + completed=set(), + active=set(), + waiting={}, + phase="Failed", + ) + ] + mock_update_all_tasks.return_value = True + + job_update = kubernetes_job.sync(test_step_cancelling) + + assert job_update.status == PlatformJobStatus.CANCELLED + + def test_sync_job_cancelling(kubernetes_job, test_step_cancelling): """Test syncing a cancelling job that isn't ready to be cancelled.""" # Mock cancelling job status diff --git a/services/core/models/src/nmp/core/models/controllers/backends/docker/backend.py b/services/core/models/src/nmp/core/models/controllers/backends/docker/backend.py index 7bee18711f..6cd1d174d8 100644 --- a/services/core/models/src/nmp/core/models/controllers/backends/docker/backend.py +++ b/services/core/models/src/nmp/core/models/controllers/backends/docker/backend.py @@ -45,6 +45,11 @@ logger = getLogger(__name__) +def _matches_labels(container_labels: dict[str, str], expected_labels: dict[str, str]) -> bool: + """Return True when all expected Docker labels are present.""" + return all(container_labels.get(key) == value for key, value in expected_labels.items()) + + class DockerServiceBackend(ServiceBackend): """Docker-based backend for managing model deployments. @@ -439,14 +444,18 @@ async def list_managed_deployment_names(self) -> list[str]: self._reconciler.list_containers, all=True, filters={"label": f"{MODEL_MANAGED_BY_LABEL}={MODEL_MANAGED_BY_MODELS_CONTROLLER}"}, + ignore_removed=True, ) except Exception as e: logger.warning(f"Failed to list managed containers for orphan reconciliation: {e}") return [] seen: set[str] = set() + owner_labels = self._backend_config.model_labels for container in containers: labels = container.labels or {} + if owner_labels and not _matches_labels(labels, owner_labels): + continue ws = labels.get("nmp.nvidia.com/deployment-workspace") n = labels.get("nmp.nvidia.com/deployment-name") if ws and n: diff --git a/services/core/models/src/nmp/core/models/controllers/backends/docker/config.py b/services/core/models/src/nmp/core/models/controllers/backends/docker/config.py index 730561699f..46c86d684a 100644 --- a/services/core/models/src/nmp/core/models/controllers/backends/docker/config.py +++ b/services/core/models/src/nmp/core/models/controllers/backends/docker/config.py @@ -174,6 +174,11 @@ class DockerBackendConfig(BaseModel): "ChunkedEncodingError, connection broken). Same volume is reused so partial downloads can be completed.", ) + model_labels: dict[str, str] = Field( + default_factory=dict, + description=("Additional labels copied onto model resources managed by the models controller. "), + ) + busybox_image: str = Field( default="busybox", description="BusyBox image repository used for helper containers (permissions/find/chown).", diff --git a/services/core/models/src/nmp/core/models/controllers/backends/docker/creation_reconciler.py b/services/core/models/src/nmp/core/models/controllers/backends/docker/creation_reconciler.py index e9273d6d14..59cd43a0bf 100644 --- a/services/core/models/src/nmp/core/models/controllers/backends/docker/creation_reconciler.py +++ b/services/core/models/src/nmp/core/models/controllers/backends/docker/creation_reconciler.py @@ -282,6 +282,20 @@ def get_health_path_from_container(self, container: Container) -> str: engine = str(labels.get(ENGINE_LABEL, ENGINE_NIM)).lower() return ENGINE_HEALTH_PATHS.get(engine, ENGINE_HEALTH_PATHS[ENGINE_NIM]) + def _managed_container_labels( + self, + deployment: ModelDeployment, + extra_labels: dict[str, str] | None = None, + ) -> dict[str, str]: + """Build labels for Docker resources managed by the models controller.""" + return { + **self._backend_config.model_labels, + "nmp.nvidia.com/deployment-workspace": deployment.workspace, + "nmp.nvidia.com/deployment-name": deployment.name, + MODEL_MANAGED_BY_LABEL: MODEL_MANAGED_BY_MODELS_CONTROLLER, + **(extra_labels or {}), + } + # ====================================================================== # Network / URL helpers # ====================================================================== @@ -348,6 +362,7 @@ async def find_available_port(self) -> Optional[int]: self.list_containers, all=True, filters={"label": f"{MODEL_MANAGED_BY_LABEL}={MODEL_MANAGED_BY_MODELS_CONTROLLER}"}, + ignore_removed=True, ) except Exception as e: logger.error(f"Failed to list containers: {e}") @@ -1048,13 +1063,13 @@ async def cleanup_and_error(status_message: str, error_details: dict) -> tuple[D "detach": True, "device_requests": device_requests, "volumes": volumes, - "labels": { - "nmp.nvidia.com/deployment-workspace": deployment.workspace, - "nmp.nvidia.com/deployment-name": deployment.name, - ENGINE_LABEL: engine, - HEALTH_PATH_LABEL: _resolve_health_path(engine, view), - MODEL_MANAGED_BY_LABEL: MODEL_MANAGED_BY_MODELS_CONTROLLER, - }, + "labels": self._managed_container_labels( + deployment, + { + ENGINE_LABEL: engine, + HEALTH_PATH_LABEL: _resolve_health_path(engine, view), + }, + ), "restart_policy": {"Name": "unless-stopped"}, } @@ -1135,11 +1150,7 @@ async def cleanup_and_error(status_message: str, error_details: dict) -> tuple[D state.volume_name: {"bind": "/model-store", "mode": "rw"}, state.scratch_volume_name: {"bind": "/scratch", "mode": "rw"}, }, - "labels": { - "nmp.nvidia.com/deployment-workspace": deployment.workspace, - "nmp.nvidia.com/deployment-name": deployment.name, - MODEL_MANAGED_BY_LABEL: MODEL_MANAGED_BY_MODELS_CONTROLLER, - }, + "labels": self._managed_container_labels(deployment), "restart_policy": {"Name": "unless-stopped"}, "healthcheck": {"test": ["NONE"]}, "ports": {}, @@ -1292,12 +1303,10 @@ async def _start_model_puller_container(self, state: CreationState) -> Container "environment": env_vars, "user": "1000:1000", "volumes": {volume_name: {"bind": "/model-store", "mode": "rw"}}, - "labels": { - "nmp.nvidia.com/deployment-workspace": deployment.workspace, - "nmp.nvidia.com/deployment-name": deployment.name, - "nmp.nvidia.com/managed-by": "models-controller", - "nmp.nvidia.com/container-type": "model-puller", - }, + "labels": self._managed_container_labels( + deployment, + {"nmp.nvidia.com/container-type": "model-puller"}, + ), "detach": True, "remove": False, } @@ -1386,12 +1395,10 @@ async def _run_plugin_puller( "environment": env_vars, "user": "1000:1000", "volumes": {volume_name: {"bind": "/model-store", "mode": "rw"}}, - "labels": { - "nmp.nvidia.com/deployment-workspace": deployment.workspace, - "nmp.nvidia.com/deployment-name": deployment.name, - MODEL_MANAGED_BY_LABEL: MODEL_MANAGED_BY_MODELS_CONTROLLER, - "nmp.nvidia.com/container-type": "plugin-puller", - }, + "labels": self._managed_container_labels( + deployment, + {"nmp.nvidia.com/container-type": "plugin-puller"}, + ), "detach": True, "remove": False, } diff --git a/services/core/models/src/nmp/core/models/controllers/backends/k8s_nim_operator/config.py b/services/core/models/src/nmp/core/models/controllers/backends/k8s_nim_operator/config.py index 6f76449388..d20cb4b4e8 100644 --- a/services/core/models/src/nmp/core/models/controllers/backends/k8s_nim_operator/config.py +++ b/services/core/models/src/nmp/core/models/controllers/backends/k8s_nim_operator/config.py @@ -199,6 +199,11 @@ class K8sNimOperatorConfig(BaseModel): "Example: {'node-type': 'gpu-node', 'zone': 'us-west1-a'}", examples=[{"node-type": "gpu-node", "zone": "us-west1-a"}], ) + model_labels: Optional[Dict[str, str]] = Field( + default=None, + description="Additional labels copied onto model resources managed by the models controller.", + examples=[{"example.com/test-label": "A", "example.com/test-suite": "B"}], + ) default_labels: Optional[Dict[str, str]] = Field( default=None, description="Default Kubernetes labels applied to NIMService and NIMCache resources and their child resources (e.g. pods). " diff --git a/services/core/models/src/nmp/core/models/controllers/backends/k8s_nim_operator/nimservice_compiler.py b/services/core/models/src/nmp/core/models/controllers/backends/k8s_nim_operator/nimservice_compiler.py index 7e105c815f..55d9f09fa3 100644 --- a/services/core/models/src/nmp/core/models/controllers/backends/k8s_nim_operator/nimservice_compiler.py +++ b/services/core/models/src/nmp/core/models/controllers/backends/k8s_nim_operator/nimservice_compiler.py @@ -140,8 +140,8 @@ def compile_nimcache( files_full_url = _get_files_hf_url() - cr_labels = _merge_default_labels( - backend_config.default_labels, + cr_labels = _merge_model_labels( + backend_config, { "app.kubernetes.io/name": resource_name, MODEL_MANAGED_BY_LABEL: MODEL_MANAGED_BY_MODELS_CONTROLLER, @@ -379,8 +379,8 @@ def compile_nimservice( env=sidecar_env_vars, ) ] - spec_labels = _merge_default_labels( - backend_config.default_labels, + spec_labels = _merge_model_labels( + backend_config, { "app.kubernetes.io/name": resource_name, MODEL_MANAGED_BY_LABEL: MODEL_MANAGED_BY_MODELS_CONTROLLER, @@ -436,8 +436,8 @@ def compile_nimservice( nimservice_metadata: dict[str, Any] = { "name": resource_name, "namespace": k8s_namespace, - "labels": _merge_default_labels( - backend_config.default_labels, + "labels": _merge_model_labels( + backend_config, { "app.kubernetes.io/name": resource_name, MODEL_MANAGED_BY_LABEL: MODEL_MANAGED_BY_MODELS_CONTROLLER, @@ -708,9 +708,10 @@ def _deep_merge(base: dict[str, Any], override: dict[str, Any]) -> None: base[key] = value -def _merge_default_labels(default_labels: Optional[dict[str, str]], base: dict[str, str]) -> dict[str, str]: - """Merge default labels with base; base takes precedence on conflict.""" - out = dict(default_labels or {}) +def _merge_model_labels(backend_config: K8sNimOperatorConfig, base: dict[str, str]) -> dict[str, str]: + """Merge configured model labels with base; base takes precedence on conflict.""" + out = dict(backend_config.default_labels or {}) + out.update(backend_config.model_labels or {}) out.update(base) return out diff --git a/services/core/models/src/nmp/core/models/controllers/backends/k8s_nim_operator/reconcilers/k8s.py b/services/core/models/src/nmp/core/models/controllers/backends/k8s_nim_operator/reconcilers/k8s.py index 80de3669a5..1108761348 100644 --- a/services/core/models/src/nmp/core/models/controllers/backends/k8s_nim_operator/reconcilers/k8s.py +++ b/services/core/models/src/nmp/core/models/controllers/backends/k8s_nim_operator/reconcilers/k8s.py @@ -149,6 +149,7 @@ async def create(self, resolved: ResolvedDeployment) -> DeploymentStatusUpdate: model_source=source_tag, namespace=self._k8s_namespace, annotations=self._backend_config.default_annotations, + extra_labels=self._backend_config.model_labels, ) job = vllm_k8s_compiler.compile_puller_job( resource_name=resource_name, @@ -165,6 +166,7 @@ async def create(self, resolved: ResolvedDeployment) -> DeploymentStatusUpdate: user_id=user_id, group_id=group_id, model_source=source_tag, + extra_labels=self._backend_config.model_labels, ) self._create_or_skip(self._core_v1.create_namespaced_persistent_volume_claim, pvc, "PVC") @@ -522,6 +524,7 @@ def _serving_spec( startup_grace_seconds=startup_grace, init_containers=init_containers, sidecar_containers=sidecar_containers, + extra_labels=self._backend_config.model_labels, mount_model_store=mount_model_store, ) svc_obj = vllm_k8s_compiler.compile_service( @@ -530,6 +533,7 @@ def _serving_spec( name=deployment.name, engine=engine, namespace=self._k8s_namespace, + extra_labels=self._backend_config.model_labels, ) return dep_obj, svc_obj diff --git a/services/core/models/src/nmp/core/models/controllers/backends/k8s_nim_operator/vllm_k8s_compiler.py b/services/core/models/src/nmp/core/models/controllers/backends/k8s_nim_operator/vllm_k8s_compiler.py index 32633a45dc..227991ffc8 100644 --- a/services/core/models/src/nmp/core/models/controllers/backends/k8s_nim_operator/vllm_k8s_compiler.py +++ b/services/core/models/src/nmp/core/models/controllers/backends/k8s_nim_operator/vllm_k8s_compiler.py @@ -153,6 +153,7 @@ def compile_pvc( model_source: Optional[str] = None, namespace: Optional[str] = None, annotations: Optional[dict[str, str]] = None, + extra_labels: Optional[dict[str, str]] = None, ) -> k8s_client.V1PersistentVolumeClaim: """Compile the model-weights PVC. @@ -165,7 +166,7 @@ def compile_pvc( metadata=k8s_client.V1ObjectMeta( name=pvc_name(resource_name), namespace=namespace, - labels=common_labels(workspace, name, engine), + labels=common_labels(workspace, name, engine, extra=extra_labels), annotations=_merge_annotations(annotations, model_source), ), spec=k8s_client.V1PersistentVolumeClaimSpec( @@ -195,6 +196,7 @@ def compile_puller_job( backoff_limit: int = DEFAULT_BACKOFF_LIMIT, ttl_seconds_after_finished: int = DEFAULT_TTL_SECONDS_AFTER_FINISHED, annotations: Optional[dict[str, str]] = None, + extra_labels: Optional[dict[str, str]] = None, ) -> k8s_client.V1Job: """Compile the weight-puller Job. @@ -209,7 +211,7 @@ def compile_puller_job( where the server can mount it (correct across any StorageClass ``volumeBindingMode``). """ - labels = common_labels(workspace, name, engine) + labels = common_labels(workspace, name, engine, extra=extra_labels) job_annotations = _merge_annotations(annotations, model_source) env_list = [k8s_client.V1EnvVar(name=k, value=str(v)) for k, v in (env or {}).items()] @@ -392,7 +394,7 @@ def compile_deployment( metadata=k8s_client.V1ObjectMeta( name=resource_name, namespace=namespace, - labels=common_labels(workspace, name, engine), + labels=common_labels(workspace, name, engine, extra=extra_labels), ), spec=k8s_client.V1DeploymentSpec( replicas=1, @@ -413,13 +415,14 @@ def compile_service( engine: str, port: int = 8000, namespace: Optional[str] = None, + extra_labels: Optional[dict[str, str]] = None, ) -> k8s_client.V1Service: """Compile the ClusterIP Service exposing the server port for IGW routing.""" return k8s_client.V1Service( metadata=k8s_client.V1ObjectMeta( name=resource_name, namespace=namespace, - labels=common_labels(workspace, name, engine), + labels=common_labels(workspace, name, engine, extra=extra_labels), ), spec=k8s_client.V1ServiceSpec( type="ClusterIP", diff --git a/services/core/models/tests/integration/conftest.py b/services/core/models/tests/integration/conftest.py index 9a0b06b4ff..be86ea273b 100644 --- a/services/core/models/tests/integration/conftest.py +++ b/services/core/models/tests/integration/conftest.py @@ -309,18 +309,34 @@ def docker_test_context( @pytest.fixture -def models_controller_container_cleanup(docker_client: docker.DockerClient) -> Generator[None, None, None]: +def docker_owner_labels(worker_id: str, testrun_uid: str) -> dict[str, str]: + """Labels used to scope Docker resources to this pytest worker/run.""" + return { + "nmp.nvidia.com/test-run": testrun_uid, + "nmp.nvidia.com/test-worker": worker_id, + } + + +@pytest.fixture +def models_controller_container_cleanup( + docker_client: docker.DockerClient, + docker_owner_labels: dict[str, str], +) -> Generator[None, None, None]: """Teardown: remove all containers with label nmp.nvidia.com/managed-by=models-controller. Ensures failed tests (e.g. stuck in PENDING) don't leave NIM/sidecar containers. Request this via controller_with_docker; no per-test try/finally needed. """ yield - cleanup_model_deployment_containers(docker_client) + cleanup_model_deployment_containers(docker_client, labels=docker_owner_labels) @pytest.fixture -def docker_backend_config(worker_id: str, mock_sidecar_image: str) -> dict[str, Any]: +def docker_backend_config( + worker_id: str, + mock_sidecar_image: str, + docker_owner_labels: dict[str, str], +) -> dict[str, Any]: """Configuration for Docker backend in tests. Uses worker_id from pytest-xdist to allocate unique port ranges per worker. @@ -333,6 +349,7 @@ def docker_backend_config(worker_id: str, mock_sidecar_image: str) -> dict[str, "models_docker_port_range_end": end_port, "docker_timeout": 60, "models_docker_host_service_name": "localhost", + "model_labels": docker_owner_labels, } diff --git a/services/core/models/tests/integration/test_models_controller.py b/services/core/models/tests/integration/test_models_controller.py index c602c955bf..01cc6934ec 100644 --- a/services/core/models/tests/integration/test_models_controller.py +++ b/services/core/models/tests/integration/test_models_controller.py @@ -419,7 +419,7 @@ def _get_worker_port_range(worker_id: str, ports_per_worker: int = 100) -> tuple @pytest.fixture -def docker_backend_config(worker_id): +def docker_backend_config(worker_id, docker_owner_labels): """Configuration for Docker backend in tests. Uses worker_id from pytest-xdist to allocate unique port ranges @@ -431,6 +431,7 @@ def docker_backend_config(worker_id): "models_docker_port_range_end": end_port, "docker_timeout": 60, "models_docker_host_service_name": "localhost", + "model_labels": docker_owner_labels, } diff --git a/services/core/models/tests/unit/controllers/backends/test_vllm_k8s_compiler.py b/services/core/models/tests/unit/controllers/backends/test_vllm_k8s_compiler.py index 6b56d82a7e..88dc6ed14a 100644 --- a/services/core/models/tests/unit/controllers/backends/test_vllm_k8s_compiler.py +++ b/services/core/models/tests/unit/controllers/backends/test_vllm_k8s_compiler.py @@ -45,6 +45,60 @@ def test_compile_pvc_basic(): assert pvc.spec.storage_class_name is None +def test_model_labels_apply_to_raw_k8s_object_metadata(): + labels = {"example.com/test-label": "A", "example.com/test-suite": "B"} + + pvc = c.compile_pvc( + resource_name="md-default-qwen", + workspace="default", + name="qwen", + engine="vllm", + disk_size="50Gi", + extra_labels=labels, + ) + job = c.compile_puller_job( + resource_name="md-default-qwen", + workspace="default", + name="qwen", + engine="vllm", + image="hf-cli:25.10", + container_args=["download", "default/qwen", "--local-dir", "/model-store"], + extra_labels=labels, + ) + deployment = c.compile_deployment( + resource_name="md-default-qwen", + workspace="default", + name="qwen", + engine="vllm", + image="vllm:v1", + args=["default/qwen"], + health_path="/v1/health/ready", + extra_labels=labels, + ) + service = c.compile_service( + resource_name="md-default-qwen", + workspace="default", + name="qwen", + engine="vllm", + extra_labels=labels, + ) + assert pvc.metadata is not None and pvc.metadata.labels is not None + assert job.metadata is not None and job.metadata.labels is not None + assert deployment.metadata is not None and deployment.metadata.labels is not None + assert ( + deployment.spec is not None + and deployment.spec.template is not None + and deployment.spec.template.metadata is not None + and deployment.spec.template.metadata.labels is not None + ) + assert service.metadata is not None and service.metadata.labels is not None + assert pvc.metadata.labels["example.com/test-label"] == "A" + assert job.metadata.labels["example.com/test-suite"] == "B" + assert deployment.metadata.labels["example.com/test-label"] == "A" + assert deployment.spec.template.metadata.labels["example.com/test-suite"] == "B" + assert service.metadata.labels["example.com/test-label"] == "A" + + def test_compile_pvc_storage_class_and_model_source(): pvc = c.compile_pvc( resource_name="md-default-qwen", diff --git a/services/core/models/tests/unit/controllers/test_backend_config_fields.py b/services/core/models/tests/unit/controllers/test_backend_config_fields.py index da07d7d967..96ffbcc6ae 100644 --- a/services/core/models/tests/unit/controllers/test_backend_config_fields.py +++ b/services/core/models/tests/unit/controllers/test_backend_config_fields.py @@ -579,6 +579,31 @@ def test_default_labels_applied_to_nimservice_metadata_and_spec(sample_deploymen assert spec_labels["nmp.nvidia.com/deployment-workspace"] == sample_deployment.workspace +def test_model_labels_applied_to_nimservice_metadata_and_spec(sample_deployment, minimal_nim_config): + """Test that model_labels from backend config are applied to NIMService CR metadata and spec (pods).""" + backend_config = K8sNimOperatorConfig( + model_labels={"nmp.nvidia.com/test-run": "run-1", "nmp.nvidia.com/test-worker": "gw0"}, + ) + + nimservice = compile_nimservice( + backend_config=backend_config, + deployment=sample_deployment, + config=minimal_nim_config, + k8s_namespace="default", + resource_name="test-resource", + ) + + meta_labels = nimservice.metadata["labels"] + assert meta_labels["nmp.nvidia.com/test-run"] == "run-1" + assert meta_labels["nmp.nvidia.com/test-worker"] == "gw0" + assert meta_labels["app.kubernetes.io/name"] == "test-resource" + + spec_labels = nimservice.spec.labels + assert spec_labels["nmp.nvidia.com/test-run"] == "run-1" + assert spec_labels["nmp.nvidia.com/test-worker"] == "gw0" + assert spec_labels["nmp.nvidia.com/deployment-workspace"] == sample_deployment.workspace + + def test_default_annotations_applied_to_nimservice_metadata_and_spec(sample_deployment, minimal_nim_config): """Test that default_annotations from backend config are applied to NIMService CR metadata and spec (pods).""" backend_config = K8sNimOperatorConfig( diff --git a/services/core/models/tests/unit/controllers/test_docker_backend.py b/services/core/models/tests/unit/controllers/test_docker_backend.py index 36abc64745..c327bec787 100644 --- a/services/core/models/tests/unit/controllers/test_docker_backend.py +++ b/services/core/models/tests/unit/controllers/test_docker_backend.py @@ -407,6 +407,58 @@ async def test_docker_backend_create_model_deployment( assert mock_container.start.call_count == 2 +@pytest.mark.asyncio +async def test_docker_backend_adds_configured_labels_to_managed_containers( + mock_nmp_sdk, + mock_docker_client, + reset_shared_resource_manager_base, + sample_deployment, + sample_config, +): + """Configured owner labels are copied to managed NIM and sidecar containers.""" + owner_labels = { + "nmp.nvidia.com/test-run": "run-1", + "nmp.nvidia.com/test-worker": "gw0", + } + platform_config = PlatformConfig( # type: ignore[abstract] + files_url="http://files-service:8000", + ) + with ( + patch("nmp.core.models.controllers.backends.docker.backend.get_platform_config", return_value=platform_config), + patch("nmp.common.resources.manager.get_platform_config") as mock_resource_config, + patch("nmp.common.resources.manager.detect_gpu_device_ids") as mock_detect_gpu_device_ids, + ): + mock_resource_config.return_value.docker = create_mock_docker_config("0,1,2,3") + mock_detect_gpu_device_ids.return_value = [0, 1, 2, 3] + backend = DockerServiceBackend( + nmp_sdk=mock_nmp_sdk, + config={ + "model_labels": owner_labels, + "models_docker_networking_mode": "dond", + }, + ) + + mock_container = MagicMock() + mock_container.id = "1234567890abcdef" + mock_container.start = MagicMock() + mock_docker_client.containers.create.return_value = mock_container + mock_docker_client.images.get.return_value = MagicMock() + mock_docker_client.containers.list.return_value = [] + sample_config.model_spec.lora_enabled = True + + await backend.create_model_deployment( + ModelContext(model_deployment=sample_deployment, model_deployment_config=sample_config) + ) + await drive_creation_to_completion(backend, sample_deployment) + + assert mock_docker_client.containers.create.call_count == 2 + for call in mock_docker_client.containers.create.call_args_list: + labels = call.kwargs["labels"] + assert labels[MODEL_MANAGED_BY_LABEL] == MODEL_MANAGED_BY_MODELS_CONTROLLER + assert labels["nmp.nvidia.com/test-run"] == "run-1" + assert labels["nmp.nvidia.com/test-worker"] == "gw0" + + def _make_vllm_config(*, lora_enabled: bool = False, image_name=None, image_tag=None): config = MagicMock() kwargs = dict( @@ -1129,6 +1181,22 @@ def test_docker_backend_dind_mode_initialization(docker_backend_with_dind_mode): assert docker_backend_with_dind_mode._backend_config.models_docker_port_range_end == 49652 +@pytest.mark.asyncio +async def test_find_available_port_ignores_containers_removed_during_list(docker_backend_with_dind_mode): + """Port allocation tolerates another worker removing a container during Docker list hydration.""" + reconciler = docker_backend_with_dind_mode._reconciler + + with patch.object(reconciler, "list_containers", return_value=[]) as list_containers: + port = await reconciler.find_available_port() + + assert port == 49152 + list_containers.assert_called_once_with( + all=True, + filters={"label": f"{MODEL_MANAGED_BY_LABEL}={MODEL_MANAGED_BY_MODELS_CONTROLLER}"}, + ignore_removed=True, + ) + + def test_docker_backend_local_mode_initialization(mock_nmp_sdk, mock_docker_client): """Test Docker backend with local mode (default) initializes correctly.""" platform_config = PlatformConfig( # type: ignore[abstract] @@ -4235,6 +4303,46 @@ async def test_list_managed_deployment_names_skips_missing_labels(backend_with_m assert names == [] +@pytest.mark.asyncio +async def test_list_managed_deployment_names_filters_by_owner_labels(mock_nmp_sdk, mock_docker_client): + """Owner labels keep orphan reconciliation scoped to this backend owner.""" + platform_config = PlatformConfig( # type: ignore[abstract] + files_url="http://files-service:8000", + ) + owner_labels = { + "nmp.nvidia.com/test-run": "run-1", + "nmp.nvidia.com/test-worker": "gw0", + } + with patch("nmp.core.models.controllers.backends.docker.backend.get_platform_config", return_value=platform_config): + backend = DockerServiceBackend( + nmp_sdk=mock_nmp_sdk, + config={"model_labels": owner_labels}, + ) + + matching_container = MagicMock() + matching_container.labels = { + MODEL_MANAGED_BY_LABEL: MODEL_MANAGED_BY_MODELS_CONTROLLER, + "nmp.nvidia.com/deployment-workspace": "ws-a", + "nmp.nvidia.com/deployment-name": "dep1", + **owner_labels, + } + other_worker_container = MagicMock() + other_worker_container.labels = { + MODEL_MANAGED_BY_LABEL: MODEL_MANAGED_BY_MODELS_CONTROLLER, + "nmp.nvidia.com/deployment-workspace": "ws-b", + "nmp.nvidia.com/deployment-name": "dep2", + "nmp.nvidia.com/test-run": "run-1", + "nmp.nvidia.com/test-worker": "gw1", + } + + with patch.object( + backend._reconciler, "list_containers", return_value=[matching_container, other_worker_container] + ): + names = await backend.list_managed_deployment_names() + + assert names == ["ws-a/dep1"] + + @pytest.mark.asyncio async def test_delete_model_deployment_by_id_calls_delete_by_model_deployment_id(backend_with_mock_client): """delete_model_deployment(workspace, name) delegates to _delete_by_model_deployment_id."""