diff --git a/plugins/nemo-deployments/src/nemo_deployments_plugin/backends/base.py b/plugins/nemo-deployments/src/nemo_deployments_plugin/backends/base.py index e87e7aac64..4956ac31c1 100644 --- a/plugins/nemo-deployments/src/nemo_deployments_plugin/backends/base.py +++ b/plugins/nemo-deployments/src/nemo_deployments_plugin/backends/base.py @@ -100,9 +100,21 @@ async def create_volume( raise NotImplementedError @abstractmethod - async def read_volume_status(self, *, workspace: str, name: str) -> VolumeStatusUpdate: + async def read_volume_status( + self, + *, + workspace: str, + name: str, + backend_config: dict[str, Any] | None = None, + ) -> VolumeStatusUpdate: raise NotImplementedError @abstractmethod - async def delete_volume(self, workspace: str, name: str) -> VolumeStatusUpdate: + async def delete_volume( + self, + workspace: str, + name: str, + *, + backend_config: dict[str, Any] | None = None, + ) -> VolumeStatusUpdate: raise NotImplementedError 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 d67994a70b..d6f60a1bc1 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 @@ -430,10 +430,24 @@ async def create_volume( driver=driver, ) - async def read_volume_status(self, *, workspace: str, name: str) -> VolumeStatusUpdate: + async def read_volume_status( + self, + *, + workspace: str, + name: str, + backend_config: dict[str, Any] | None = None, + ) -> VolumeStatusUpdate: + # backend_config is part of the DeploymentBackend ABC so the reconciler can pass + # K8s namespace overrides; Docker volume names are global to the daemon. return await volume_ops.read_volume_status(self._client, workspace=workspace, name=name) - async def delete_volume(self, workspace: str, name: str) -> VolumeStatusUpdate: + async def delete_volume( + self, + workspace: str, + name: str, + *, + backend_config: dict[str, Any] | None = None, + ) -> VolumeStatusUpdate: return await volume_ops.delete_volume(self._client, workspace=workspace, name=name) def _container_matches_deployment( diff --git a/plugins/nemo-deployments/src/nemo_deployments_plugin/backends/k8s/backend.py b/plugins/nemo-deployments/src/nemo_deployments_plugin/backends/k8s/backend.py index acafd58274..63d5018745 100644 --- a/plugins/nemo-deployments/src/nemo_deployments_plugin/backends/k8s/backend.py +++ b/plugins/nemo-deployments/src/nemo_deployments_plugin/backends/k8s/backend.py @@ -14,6 +14,7 @@ LogResult, VolumeStatusUpdate, ) +from nemo_deployments_plugin.backends.k8s import volumes as volume_ops from nemo_deployments_plugin.backends.k8s.client import KubernetesClients from nemo_deployments_plugin.backends.k8s.config import K8sExecutorConfig @@ -101,10 +102,42 @@ async def create_volume( access_modes: list[str], backend_config: dict[str, Any], ) -> VolumeStatusUpdate: - raise NotImplementedError("K8s create_volume is implemented in a later phase.") + return await volume_ops.create_volume( + self._clients, + default_namespace=self._executor_config.default_namespace, + workspace=workspace, + name=name, + size=size, + access_modes=access_modes, + backend_config=backend_config, + ) - async def read_volume_status(self, *, workspace: str, name: str) -> VolumeStatusUpdate: - raise NotImplementedError("K8s read_volume_status is implemented in a later phase.") + async def read_volume_status( + self, + *, + workspace: str, + name: str, + backend_config: dict[str, Any] | None = None, + ) -> VolumeStatusUpdate: + return await volume_ops.read_volume_status( + self._clients, + default_namespace=self._executor_config.default_namespace, + workspace=workspace, + name=name, + backend_config=backend_config, + ) - async def delete_volume(self, workspace: str, name: str) -> VolumeStatusUpdate: - raise NotImplementedError("K8s delete_volume is implemented in a later phase.") + async def delete_volume( + self, + workspace: str, + name: str, + *, + backend_config: dict[str, Any] | None = None, + ) -> VolumeStatusUpdate: + return await volume_ops.delete_volume( + self._clients, + default_namespace=self._executor_config.default_namespace, + workspace=workspace, + name=name, + backend_config=backend_config, + ) diff --git a/plugins/nemo-deployments/src/nemo_deployments_plugin/backends/k8s/volumes.py b/plugins/nemo-deployments/src/nemo_deployments_plugin/backends/k8s/volumes.py new file mode 100644 index 0000000000..0753200cb4 --- /dev/null +++ b/plugins/nemo-deployments/src/nemo_deployments_plugin/backends/k8s/volumes.py @@ -0,0 +1,211 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Kubernetes PVC lifecycle helpers for the deployments plugin.""" + +from __future__ import annotations + +import asyncio +import logging +from typing import Any + +from kubernetes.client.rest import ApiException +from nemo_deployments_plugin.backends.base import VolumeStatusUpdate +from nemo_deployments_plugin.backends.k8s.client import KubernetesClients, _kubernetes_modules +from nemo_deployments_plugin.backends.labels import k8s_volume_resource_name, volume_identity_labels +from nemo_deployments_plugin.entities import K8sVolumeConfig + +logger = logging.getLogger(__name__) + + +def resolve_volume_config(backend_config: dict[str, Any]) -> K8sVolumeConfig | None: + """Parse and validate the k8s section of entity backend_config.""" + k8s_section = backend_config.get("k8s") + if not k8s_section: + return None + return K8sVolumeConfig.model_validate(k8s_section) + + +def resolve_volume_namespace(*, default_namespace: str, volume_config: K8sVolumeConfig | None) -> str: + """Resolve target namespace from parsed volume config with executor default fallback.""" + if volume_config is not None and volume_config.namespace: + return volume_config.namespace + return default_namespace + + +def build_pvc_body( + *, + pvc_name: str, + labels: dict[str, str], + size: str, + access_modes: list[str], + storage_class: str | None, +) -> Any: + """Build a ``V1PersistentVolumeClaim`` for create.""" + client, _ = _kubernetes_modules() + spec_kwargs: dict[str, Any] = { + "access_modes": list(access_modes), + "resources": client.V1VolumeResourceRequirements(requests={"storage": size}), + } + if storage_class is not None: + spec_kwargs["storage_class_name"] = storage_class + return client.V1PersistentVolumeClaim( + api_version="v1", + kind="PersistentVolumeClaim", + metadata=client.V1ObjectMeta(name=pvc_name, labels=labels), + spec=client.V1PersistentVolumeClaimSpec(**spec_kwargs), + ) + + +def map_pvc_phase_to_status(*, pvc_name: str, phase: str | None) -> VolumeStatusUpdate: + """Map Kubernetes PVC phase to plugin ``VolumeStatus``.""" + if phase == "Bound": + return VolumeStatusUpdate(status="BOUND", status_message=f"PVC {pvc_name} is bound") + if phase == "Lost": + return VolumeStatusUpdate(status="FAILED", status_message=f"PVC {pvc_name} is lost") + return VolumeStatusUpdate(status="PENDING", status_message=f"PVC {pvc_name} is pending") + + +def _phase_from_pvc(pvc: Any) -> str | None: + if pvc.status is None: + return None + return pvc.status.phase + + +def _pvc_is_deleting(pvc: Any) -> bool: + metadata = getattr(pvc, "metadata", None) + return bool(metadata and getattr(metadata, "deletion_timestamp", None)) + + +def _pvc_labels_match(pvc: Any, expected_labels: dict[str, str]) -> bool: + metadata = getattr(pvc, "metadata", None) + if metadata is None or not metadata.labels: + return False + return all(metadata.labels.get(key) == value for key, value in expected_labels.items()) + + +def status_from_pvc(*, pvc: Any, pvc_name: str, expected_labels: dict[str, str]) -> VolumeStatusUpdate: + """Map a PVC object to plugin status, enforcing identity labels and delete propagation.""" + if not _pvc_labels_match(pvc, expected_labels): + return VolumeStatusUpdate( + status="FAILED", + status_message=f"PVC {pvc_name} exists but is not managed by this plugin", + ) + if _pvc_is_deleting(pvc): + return VolumeStatusUpdate(status="DELETING", status_message=f"PVC {pvc_name} is terminating") + return map_pvc_phase_to_status(pvc_name=pvc_name, phase=_phase_from_pvc(pvc)) + + +async def create_volume( + clients: KubernetesClients, + *, + default_namespace: str, + workspace: str, + name: str, + size: str, + access_modes: list[str], + backend_config: dict[str, Any], +) -> VolumeStatusUpdate: + pvc_name = k8s_volume_resource_name(workspace, name) + try: + volume_config = resolve_volume_config(backend_config) + namespace = resolve_volume_namespace(default_namespace=default_namespace, volume_config=volume_config) + storage_class = volume_config.storage_class if volume_config else None + labels = volume_identity_labels(workspace, name) + body = build_pvc_body( + pvc_name=pvc_name, + labels=labels, + size=size, + access_modes=access_modes, + storage_class=storage_class, + ) + timeout = clients.request_timeout + core_v1 = clients.core_v1 + + def _create() -> Any: + try: + return core_v1.create_namespaced_persistent_volume_claim( + namespace=namespace, + body=body, + _request_timeout=timeout, + ) + except ApiException as exc: + if exc.status == 409: + return core_v1.read_namespaced_persistent_volume_claim( + name=pvc_name, + namespace=namespace, + _request_timeout=timeout, + ) + raise + + pvc = await asyncio.to_thread(_create) + return status_from_pvc(pvc=pvc, pvc_name=pvc_name, expected_labels=labels) + except Exception as exc: + logger.exception("Failed to create PVC %s", pvc_name) + return VolumeStatusUpdate(status="FAILED", status_message=f"Failed to create PVC: {exc}") + + +async def read_volume_status( + clients: KubernetesClients, + *, + default_namespace: str, + workspace: str, + name: str, + backend_config: dict[str, Any] | None = None, +) -> VolumeStatusUpdate: + pvc_name = k8s_volume_resource_name(workspace, name) + expected_labels = volume_identity_labels(workspace, name) + try: + volume_config = resolve_volume_config(backend_config or {}) + namespace = resolve_volume_namespace(default_namespace=default_namespace, volume_config=volume_config) + timeout = clients.request_timeout + core_v1 = clients.core_v1 + + def _read() -> Any: + return core_v1.read_namespaced_persistent_volume_claim( + name=pvc_name, + namespace=namespace, + _request_timeout=timeout, + ) + + pvc = await asyncio.to_thread(_read) + return status_from_pvc(pvc=pvc, pvc_name=pvc_name, expected_labels=expected_labels) + except ApiException as exc: + if exc.status == 404: + return VolumeStatusUpdate(status="FAILED", status_message=f"PVC {pvc_name} not found") + return VolumeStatusUpdate(status="FAILED", status_message=f"Failed to read PVC: {exc}") + except Exception as exc: + return VolumeStatusUpdate(status="FAILED", status_message=f"Failed to read PVC: {exc}") + + +async def delete_volume( + clients: KubernetesClients, + *, + default_namespace: str, + workspace: str, + name: str, + backend_config: dict[str, Any] | None = None, +) -> VolumeStatusUpdate: + pvc_name = k8s_volume_resource_name(workspace, name) + try: + volume_config = resolve_volume_config(backend_config or {}) + namespace = resolve_volume_namespace(default_namespace=default_namespace, volume_config=volume_config) + timeout = clients.request_timeout + core_v1 = clients.core_v1 + + def _delete() -> None: + try: + core_v1.delete_namespaced_persistent_volume_claim( + name=pvc_name, + namespace=namespace, + _request_timeout=timeout, + ) + except ApiException as exc: + if exc.status == 404: + return + raise + + await asyncio.to_thread(_delete) + return VolumeStatusUpdate(status="RELEASED", status_message=f"PVC {pvc_name} released") + except Exception as exc: + return VolumeStatusUpdate(status="FAILED", status_message=f"Failed to delete PVC: {exc}") diff --git a/plugins/nemo-deployments/src/nemo_deployments_plugin/reconciler/volume_reconciler.py b/plugins/nemo-deployments/src/nemo_deployments_plugin/reconciler/volume_reconciler.py index ea0338b9ea..9e64df6885 100644 --- a/plugins/nemo-deployments/src/nemo_deployments_plugin/reconciler/volume_reconciler.py +++ b/plugins/nemo-deployments/src/nemo_deployments_plugin/reconciler/volume_reconciler.py @@ -49,8 +49,9 @@ async def _reconcile_delete(self, volume: Volume) -> None: logger.warning("No executor for volume delete of %s — will retry", volume_id, exc_info=True) return + backend_config = volume.backend_config.model_dump(by_alias=True, exclude_none=True) try: - await backend.delete_volume(volume.workspace, volume.name) + await backend.delete_volume(volume.workspace, volume.name, backend_config=backend_config) except Exception: logger.warning("Backend delete failed for volume %s — will retry", volume_id, exc_info=True) return @@ -87,8 +88,13 @@ async def _reconcile_create(self, volume: Volume, backend: DeploymentBackend) -> ) async def _reconcile_read(self, volume: Volume, backend: DeploymentBackend) -> None: + backend_config = volume.backend_config.model_dump(by_alias=True, exclude_none=True) try: - update = await backend.read_volume_status(workspace=volume.workspace, name=volume.name) + update = await backend.read_volume_status( + workspace=volume.workspace, + name=volume.name, + backend_config=backend_config, + ) await self._update_volume_status(volume, update) except NemoEntityConflictError: raise diff --git a/plugins/nemo-deployments/tests/unit/backends/k8s/conftest.py b/plugins/nemo-deployments/tests/unit/backends/k8s/conftest.py index fab6513555..04c8088ee7 100644 --- a/plugins/nemo-deployments/tests/unit/backends/k8s/conftest.py +++ b/plugins/nemo-deployments/tests/unit/backends/k8s/conftest.py @@ -23,6 +23,7 @@ def mock_k8s_clients() -> MagicMock: clients.core_v1 = MagicMock() clients.apps_v1 = MagicMock() clients.batch_v1 = MagicMock() + clients.request_timeout = 30 return clients diff --git a/plugins/nemo-deployments/tests/unit/backends/k8s/test_volumes.py b/plugins/nemo-deployments/tests/unit/backends/k8s/test_volumes.py new file mode 100644 index 0000000000..a6e3613f59 --- /dev/null +++ b/plugins/nemo-deployments/tests/unit/backends/k8s/test_volumes.py @@ -0,0 +1,195 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from unittest.mock import MagicMock + +import pytest +from kubernetes.client.rest import ApiException +from nemo_deployments_plugin.backends.k8s import volumes as volume_ops +from nemo_deployments_plugin.backends.k8s.client import KubernetesClients +from nemo_deployments_plugin.backends.k8s.volumes import map_pvc_phase_to_status, status_from_pvc +from nemo_deployments_plugin.backends.labels import ( + MANAGED_BY_KEY, + VOLUME_NAME_LABEL, + VOLUME_WORKSPACE_LABEL, + volume_identity_labels, +) + + +@pytest.mark.parametrize( + ("phase", "expected_status"), + [ + ("Pending", "PENDING"), + ("Bound", "BOUND"), + ("Lost", "FAILED"), + (None, "PENDING"), + ], +) +def test_map_pvc_phase_to_status(phase: str | None, expected_status: str) -> None: + update = map_pvc_phase_to_status(pvc_name="dep-vol-default-data-abc12345", phase=phase) + assert update.status == expected_status + + +@pytest.mark.asyncio +async def test_create_volume_emits_pvc_with_storage_class_and_size(k8s_backend, mock_k8s_clients: MagicMock) -> None: + mock_pvc = MagicMock() + mock_pvc.status.phase = "Pending" + mock_pvc.metadata.labels = volume_identity_labels("default", "weights") + mock_pvc.metadata.deletion_timestamp = None + mock_k8s_clients.core_v1.create_namespaced_persistent_volume_claim.return_value = mock_pvc + mock_k8s_clients.request_timeout = 30 + + update = await k8s_backend.create_volume( + workspace="default", + name="weights", + size="10Gi", + access_modes=["ReadWriteOnce"], + backend_config={"k8s": {"storageClass": "fast-ssd", "namespace": "models"}}, + ) + + assert update.status == "PENDING" + call = mock_k8s_clients.core_v1.create_namespaced_persistent_volume_claim + call.assert_called_once() + assert call.call_args.kwargs["namespace"] == "models" + assert call.call_args.kwargs["_request_timeout"] == 30 + body = call.call_args.kwargs["body"] + assert body.spec.resources.requests["storage"] == "10Gi" + assert body.spec.storage_class_name == "fast-ssd" + assert body.spec.access_modes == ["ReadWriteOnce"] + assert body.metadata.labels[VOLUME_WORKSPACE_LABEL] == "default" + assert body.metadata.labels[VOLUME_NAME_LABEL] == "weights" + + +@pytest.mark.asyncio +async def test_create_volume_conflict_reads_existing_pvc(mock_k8s_clients: MagicMock) -> None: + existing = MagicMock() + existing.status.phase = "Bound" + existing.metadata.labels = volume_identity_labels("default", "data") + existing.metadata.deletion_timestamp = None + mock_k8s_clients.core_v1.create_namespaced_persistent_volume_claim.side_effect = ApiException(status=409) + mock_k8s_clients.core_v1.read_namespaced_persistent_volume_claim.return_value = existing + clients = MagicMock(spec=KubernetesClients) + clients.core_v1 = mock_k8s_clients.core_v1 + clients.request_timeout = 60 + + update = await volume_ops.create_volume( + clients, + default_namespace="default", + workspace="default", + name="data", + size="1Gi", + access_modes=["ReadWriteOnce"], + backend_config={}, + ) + + assert update.status == "BOUND" + mock_k8s_clients.core_v1.read_namespaced_persistent_volume_claim.assert_called_once() + + +@pytest.mark.asyncio +async def test_create_volume_conflict_rejects_foreign_pvc(mock_k8s_clients: MagicMock) -> None: + foreign = MagicMock() + foreign.metadata.labels = {MANAGED_BY_KEY: "other-plugin"} + foreign.status.phase = "Bound" + mock_k8s_clients.core_v1.create_namespaced_persistent_volume_claim.side_effect = ApiException(status=409) + mock_k8s_clients.core_v1.read_namespaced_persistent_volume_claim.return_value = foreign + clients = MagicMock(spec=KubernetesClients) + clients.core_v1 = mock_k8s_clients.core_v1 + clients.request_timeout = 60 + + update = await volume_ops.create_volume( + clients, + default_namespace="default", + workspace="default", + name="data", + size="1Gi", + access_modes=["ReadWriteOnce"], + backend_config={}, + ) + + assert update.status == "FAILED" + assert "not managed" in update.status_message + + +def test_status_from_pvc_deleting_reports_deleting() -> None: + labels = volume_identity_labels("default", "data") + pvc = MagicMock() + pvc.metadata.labels = labels + pvc.metadata.deletion_timestamp = "2026-01-01T00:00:00Z" + pvc.status.phase = "Bound" + + update = status_from_pvc(pvc=pvc, pvc_name="dep-vol-default-data-abc12345", expected_labels=labels) + + assert update.status == "DELETING" + + +@pytest.mark.asyncio +async def test_read_volume_status_uses_entity_namespace(k8s_backend, mock_k8s_clients: MagicMock) -> None: + mock_pvc = MagicMock() + mock_pvc.status.phase = "Bound" + mock_pvc.metadata.labels = volume_identity_labels("default", "weights") + mock_pvc.metadata.deletion_timestamp = None + mock_k8s_clients.core_v1.read_namespaced_persistent_volume_claim.return_value = mock_pvc + + update = await k8s_backend.read_volume_status( + workspace="default", + name="weights", + backend_config={"k8s": {"namespace": "models"}}, + ) + + assert update.status == "BOUND" + call = mock_k8s_clients.core_v1.read_namespaced_persistent_volume_claim + assert call.call_args.kwargs["namespace"] == "models" + + +@pytest.mark.asyncio +async def test_create_volume_malformed_backend_config_returns_failed(mock_k8s_clients: MagicMock) -> None: + clients = MagicMock(spec=KubernetesClients) + clients.core_v1 = mock_k8s_clients.core_v1 + clients.request_timeout = 60 + + update = await volume_ops.create_volume( + clients, + default_namespace="default", + workspace="default", + name="data", + size="1Gi", + access_modes=["ReadWriteOnce"], + backend_config={"k8s": {"storageClass": 123}}, + ) + + assert update.status == "FAILED" + mock_k8s_clients.core_v1.create_namespaced_persistent_volume_claim.assert_not_called() + + +@pytest.mark.asyncio +async def test_read_volume_status_malformed_backend_config_returns_failed(k8s_backend) -> None: + update = await k8s_backend.read_volume_status( + workspace="default", + name="weights", + backend_config={"k8s": {"namespace": 42}}, + ) + + assert update.status == "FAILED" + + +@pytest.mark.asyncio +async def test_read_volume_status_not_found(k8s_backend, mock_k8s_clients: MagicMock) -> None: + mock_k8s_clients.core_v1.read_namespaced_persistent_volume_claim.side_effect = ApiException(status=404) + + update = await k8s_backend.read_volume_status(workspace="default", name="missing") + + assert update.status == "FAILED" + assert "not found" in update.status_message + + +@pytest.mark.asyncio +async def test_delete_volume_missing_is_released(k8s_backend, mock_k8s_clients: MagicMock) -> None: + mock_k8s_clients.request_timeout = 30 + mock_k8s_clients.core_v1.delete_namespaced_persistent_volume_claim.side_effect = ApiException(status=404) + + update = await k8s_backend.delete_volume("default", "gone") + + assert update.status == "RELEASED" diff --git a/plugins/nemo-deployments/tests/unit/reconciler/conftest.py b/plugins/nemo-deployments/tests/unit/reconciler/conftest.py index f08e4f8c0c..a28cefeca7 100644 --- a/plugins/nemo-deployments/tests/unit/reconciler/conftest.py +++ b/plugins/nemo-deployments/tests/unit/reconciler/conftest.py @@ -66,10 +66,22 @@ async def get_logs(self, **kwargs: Any) -> LogResult: async def create_volume(self, **kwargs: Any) -> VolumeStatusUpdate: return self.volume_create_status - async def read_volume_status(self, *, workspace: str, name: str) -> VolumeStatusUpdate: + async def read_volume_status( + self, + *, + workspace: str, + name: str, + backend_config: dict[str, Any] | None = None, + ) -> VolumeStatusUpdate: return VolumeStatusUpdate(status="BOUND") - async def delete_volume(self, workspace: str, name: str) -> VolumeStatusUpdate: + async def delete_volume( + self, + workspace: str, + name: str, + *, + backend_config: dict[str, Any] | None = None, + ) -> VolumeStatusUpdate: self.volume_delete_calls.append((workspace, name)) return VolumeStatusUpdate(status="RELEASED")