diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 2426f4a0ce9b..e96f30e1f89a 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -1243,6 +1243,9 @@ jobs: dockerhub_password: ${{ secrets.DOCKERHUB_ACCESS_TOKEN }} checkpoint_enabled: 'true' checkpoint_storage_size: 64Gi + # Same shared model cache as regular deploy tests; empty vars -> HF download. + model_cache_server: ${{ vars.AZURE_MODEL_CACHE_SERVER }} + model_cache_path: ${{ vars.AZURE_MODEL_CACHE_PATH }} deploy-operator-checkpoint-sglang: name: SGLang DynamoCheckpoint Operator Setup @@ -1274,6 +1277,9 @@ jobs: dockerhub_password: ${{ secrets.DOCKERHUB_ACCESS_TOKEN }} checkpoint_enabled: 'true' checkpoint_storage_size: 64Gi + # Same shared model cache as regular deploy tests; empty vars -> HF download. + model_cache_server: ${{ vars.AZURE_MODEL_CACHE_SERVER }} + model_cache_path: ${{ vars.AZURE_MODEL_CACHE_PATH }} deploy-operator-checkpoint-trtllm: name: TRTLLM DynamoCheckpoint Operator Setup @@ -1305,6 +1311,9 @@ jobs: dockerhub_password: ${{ secrets.DOCKERHUB_ACCESS_TOKEN }} checkpoint_enabled: 'true' checkpoint_storage_size: 64Gi + # Same shared model cache as regular deploy tests; empty vars -> HF download. + model_cache_server: ${{ vars.AZURE_MODEL_CACHE_SERVER }} + model_cache_path: ${{ vars.AZURE_MODEL_CACHE_PATH }} deploy-snapshot-agent-checkpoint-vllm: name: vLLM DynamoCheckpoint Snapshot Agent Setup @@ -1341,6 +1350,9 @@ jobs: dockerhub_password: ${{ secrets.DOCKERHUB_ACCESS_TOKEN }} checkpoint_enabled: 'true' checkpoint_storage_size: 64Gi + # Same shared model cache as regular deploy tests; empty vars -> HF download. + model_cache_server: ${{ vars.AZURE_MODEL_CACHE_SERVER }} + model_cache_path: ${{ vars.AZURE_MODEL_CACHE_PATH }} - name: Connect to vCluster id: connect-vcluster uses: ./.github/actions/connect-vcluster @@ -1389,6 +1401,9 @@ jobs: dockerhub_password: ${{ secrets.DOCKERHUB_ACCESS_TOKEN }} checkpoint_enabled: 'true' checkpoint_storage_size: 64Gi + # Same shared model cache as regular deploy tests; empty vars -> HF download. + model_cache_server: ${{ vars.AZURE_MODEL_CACHE_SERVER }} + model_cache_path: ${{ vars.AZURE_MODEL_CACHE_PATH }} - name: Connect to vCluster id: connect-vcluster uses: ./.github/actions/connect-vcluster @@ -1437,6 +1452,9 @@ jobs: dockerhub_password: ${{ secrets.DOCKERHUB_ACCESS_TOKEN }} checkpoint_enabled: 'true' checkpoint_storage_size: 64Gi + # Same shared model cache as regular deploy tests; empty vars -> HF download. + model_cache_server: ${{ vars.AZURE_MODEL_CACHE_SERVER }} + model_cache_path: ${{ vars.AZURE_MODEL_CACHE_PATH }} - name: Connect to vCluster id: connect-vcluster uses: ./.github/actions/connect-vcluster @@ -1486,6 +1504,9 @@ jobs: dockerhub_password: ${{ secrets.DOCKERHUB_ACCESS_TOKEN }} checkpoint_enabled: 'true' checkpoint_storage_size: 64Gi + # Same shared model cache as regular deploy tests; empty vars -> HF download. + model_cache_server: ${{ vars.AZURE_MODEL_CACHE_SERVER }} + model_cache_path: ${{ vars.AZURE_MODEL_CACHE_PATH }} - name: Connect to vCluster id: connect-vcluster uses: ./.github/actions/connect-vcluster @@ -1511,6 +1532,9 @@ jobs: image: ${{ secrets.AZURE_ACR_HOSTNAME }}/ai-dynamo/dynamo:${{ needs.snapshot-placeholder-vllm.outputs.image_tag }} test_name: checkpoint_dgd_restore_vllm test_file: tests/deploy/test_dynamocheckpoint.py + # Mount the shared model cache only when both endpoint vars are set + # (matches the PV/PVC creation gate); otherwise workers download from HF. + model_cache_pvc: ${{ vars.AZURE_MODEL_CACHE_SERVER != '' && vars.AZURE_MODEL_CACHE_PATH != '' && 'model-cache' || '' }} extra_pytest_args: >- -m dynamocheckpoint --checkpoint-backend=vllm @@ -1552,6 +1576,9 @@ jobs: dockerhub_password: ${{ secrets.DOCKERHUB_ACCESS_TOKEN }} checkpoint_enabled: 'true' checkpoint_storage_size: 64Gi + # Same shared model cache as regular deploy tests; empty vars -> HF download. + model_cache_server: ${{ vars.AZURE_MODEL_CACHE_SERVER }} + model_cache_path: ${{ vars.AZURE_MODEL_CACHE_PATH }} - name: Connect to vCluster id: connect-vcluster uses: ./.github/actions/connect-vcluster @@ -1577,6 +1604,9 @@ jobs: image: ${{ secrets.AZURE_ACR_HOSTNAME }}/ai-dynamo/dynamo:${{ needs.snapshot-placeholder-sglang.outputs.image_tag }} test_name: checkpoint_dgd_restore_sglang test_file: tests/deploy/test_dynamocheckpoint.py + # Mount the shared model cache only when both endpoint vars are set + # (matches the PV/PVC creation gate); otherwise workers download from HF. + model_cache_pvc: ${{ vars.AZURE_MODEL_CACHE_SERVER != '' && vars.AZURE_MODEL_CACHE_PATH != '' && 'model-cache' || '' }} extra_pytest_args: >- -m dynamocheckpoint --checkpoint-backend=sglang @@ -1618,6 +1648,9 @@ jobs: dockerhub_password: ${{ secrets.DOCKERHUB_ACCESS_TOKEN }} checkpoint_enabled: 'true' checkpoint_storage_size: 64Gi + # Same shared model cache as regular deploy tests; empty vars -> HF download. + model_cache_server: ${{ vars.AZURE_MODEL_CACHE_SERVER }} + model_cache_path: ${{ vars.AZURE_MODEL_CACHE_PATH }} - name: Connect to vCluster id: connect-vcluster uses: ./.github/actions/connect-vcluster @@ -1643,6 +1676,9 @@ jobs: image: ${{ secrets.AZURE_ACR_HOSTNAME }}/ai-dynamo/dynamo:${{ needs.snapshot-placeholder-trtllm.outputs.image_tag }} test_name: checkpoint_dgd_restore_trtllm test_file: tests/deploy/test_dynamocheckpoint.py + # Mount the shared model cache only when both endpoint vars are set + # (matches the PV/PVC creation gate); otherwise workers download from HF. + model_cache_pvc: ${{ vars.AZURE_MODEL_CACHE_SERVER != '' && vars.AZURE_MODEL_CACHE_PATH != '' && 'model-cache' || '' }} extra_pytest_args: >- -m dynamocheckpoint --checkpoint-backend=trtllm diff --git a/tests/deploy/test_dynamocheckpoint.py b/tests/deploy/test_dynamocheckpoint.py index 26bc50f6556f..a8cb8e3013dc 100644 --- a/tests/deploy/test_dynamocheckpoint.py +++ b/tests/deploy/test_dynamocheckpoint.py @@ -168,9 +168,10 @@ class CheckpointBackendConfig: "--free-gpu-memory-fraction", "0.10", ), - # Keep the raw DGD PVC-free: the checkpoint operator mounts - # snapshot-pvc at /checkpoints for checkpoint/restore pods, so HF_HOME - # there preserves model files across restore without a model-cache PVC. + # UCX_TLS is always set. HF_HOME defaults to the snapshot PVC so restore + # pods keep weights without a model-cache PVC; when CI passes + # --model-cache-pvc, _new_checkpoint_spec skips this HF_HOME so the + # shared cache mount can own it (same as regular deploy tests). env=(("UCX_TLS", "tcp,self"), ("HF_HOME", TRTLLM_HF_HOME)), # Match the base TRTLLM snapshot recipe and avoid cold-worker/restore # rollout overlap during initial DGD startup. @@ -233,6 +234,9 @@ def _new_checkpoint_spec( namespace: str, image: str, frontend_image: str, + *, + model_cache_pvc: str | None = None, + model_cache_mount: str | None = None, ) -> DeploymentSpec: spec_path = Path(_get_workspace_dir()).joinpath(*backend.manifest) deployment_spec = DeploymentSpec(str(spec_path)) @@ -269,6 +273,10 @@ def _new_checkpoint_spec( if backend.env: env = container.setdefault("env", []) for name, value in backend.env: + # Container HF_HOME would shadow the deployment-level value that + # mount_model_cache_pvc sets; skip it when the shared cache is used. + if name == "HF_HOME" and model_cache_pvc: + continue for item in env: if item.get("name") == name: item["value"] = value @@ -281,6 +289,11 @@ def _new_checkpoint_spec( checkpoint["targetContainerName"] = backend.target_container if backend.checkpoint_startup_policy is not None: checkpoint["startupPolicy"] = backend.checkpoint_startup_policy + + if model_cache_pvc: + mount = model_cache_mount or "/models" + deployment_spec.mount_model_cache_pvc(model_cache_pvc, mount) + return deployment_spec @@ -594,6 +607,8 @@ async def test_dgd_checkpoint_restore_deploy( namespace=namespace, image=image, frontend_image=frontend_image, + model_cache_pvc=request.config.getoption("--model-cache-pvc") or None, + model_cache_mount=request.config.getoption("--model-cache-mount") or None, ) async with ManagedDeployment( diff --git a/tests/utils/managed_deployment.py b/tests/utils/managed_deployment.py index 6ecd08bf5a82..0cdf478704b0 100644 --- a/tests/utils/managed_deployment.py +++ b/tests/utils/managed_deployment.py @@ -117,6 +117,20 @@ def _sidecar_container_name(self) -> str: def _sidecar_container(self, create: bool = False) -> Optional[dict]: return self._find_container(self._sidecar_container_name(), create=create) + def _non_main_checkpoint_target_container(self) -> Optional[dict]: + checkpoint = self._spec.get("experimental", {}).get("checkpoint", {}) + target_name = checkpoint.get("targetContainerName") + if not target_name or target_name == "main": + return None + + target = self._find_container(target_name) + if target is None: + raise ValueError( + f"v1beta1 component {self.name!r} checkpoint targetContainerName " + f"{target_name!r} does not name a podTemplate container" + ) + return target + # ----- Image ----- @property def image(self) -> Optional[str]: @@ -188,7 +202,34 @@ def envs(self, value: list[dict[str, str]]): self._spec["envs"] = value def add_pvc_mount(self, pvc_name: str, mount_point: str) -> None: - """Add a service-level volumeMount for a PVC declared in ``spec.pvcs``. Idempotent.""" + """Mount a PVC on this service/component. Idempotent.""" + if self._schema == SCHEMA_V1BETA1: + checkpoint_target = self._non_main_checkpoint_target_container() + pod_spec = self._spec.setdefault("podTemplate", {}).setdefault("spec", {}) + volumes = pod_spec.setdefault("volumes", []) + if not any(volume.get("name") == pvc_name for volume in volumes): + volumes.append( + { + "name": pvc_name, + "persistentVolumeClaim": {"claimName": pvc_name}, + } + ) + + container = self._main_container(create=True) + assert container is not None + cache_containers = [container] + if checkpoint_target is not None: + cache_containers.append(checkpoint_target) + env = checkpoint_target.setdefault("env", []) + if not any(item.get("name") == "HF_HOME" for item in env): + env.append({"name": "HF_HOME", "value": mount_point}) + + for cache_container in cache_containers: + mounts = cache_container.setdefault("volumeMounts", []) + if not any(mount.get("name") == pvc_name for mount in mounts): + mounts.append({"name": pvc_name, "mountPath": mount_point}) + return + mounts = self._spec.setdefault("volumeMounts", []) if not any(m.get("name") == pvc_name for m in mounts): mounts.append({"name": pvc_name, "mountPoint": mount_point}) @@ -495,10 +536,23 @@ def mount_model_cache_pvc( self, pvc_name: str, mount_point: str = "/models" ) -> None: """Reference a pre-existing PVC and mount it at ``mount_point`` on every - service, with ``HF_HOME`` pointed at it so models come from the shared cache - instead of HuggingFace. Idempotent; used by CI via --model-cache-pvc. + service/component, with ``HF_HOME`` pointed at it so models come from the + shared cache instead of HuggingFace. Idempotent; used by CI via + --model-cache-pvc. """ spec = self._deployment_spec["spec"] + if self._schema == SCHEMA_V1BETA1: + services = self.services + for service in services: + service._non_main_checkpoint_target_container() + + env = spec.setdefault("env", []) + if not any(item.get("name") == "HF_HOME" for item in env): + env.append({"name": "HF_HOME", "value": mount_point}) + for service in services: + service.add_pvc_mount(pvc_name, mount_point) + return + pvcs = spec.setdefault("pvcs", []) if not any(p.get("name") == pvc_name for p in pvcs): pvcs.append({"name": pvc_name, "create": False}) diff --git a/tests/utils/test_managed_deployment.py b/tests/utils/test_managed_deployment.py new file mode 100644 index 000000000000..ffecbea26978 --- /dev/null +++ b/tests/utils/test_managed_deployment.py @@ -0,0 +1,225 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import copy +from pathlib import Path +from typing import Any + +import pytest +import yaml + +from tests.utils.managed_deployment import DeploymentSpec + +pytestmark = [ + pytest.mark.pre_merge, + pytest.mark.gpu_0, + pytest.mark.unit, +] + + +def _deployment_spec(tmp_path: Path, raw_spec: dict[str, Any]) -> DeploymentSpec: + manifest = tmp_path / "deployment.yaml" + manifest.write_text(yaml.safe_dump(raw_spec)) + return DeploymentSpec(str(manifest)) + + +def test_mount_model_cache_pvc_uses_v1beta1_native_fields( + tmp_path: Path, +) -> None: + deployment = _deployment_spec( + tmp_path, + { + "apiVersion": "nvidia.com/v1beta1", + "kind": "DynamoGraphDeployment", + "metadata": {"name": "test"}, + "spec": { + "env": [{"name": "EXISTING", "value": "value"}], + "components": [ + { + "name": "Frontend", + "podTemplate": { + "spec": { + "containers": [{"name": "main", "image": "frontend"}] + } + }, + }, + { + "name": "Worker", + "experimental": { + "checkpoint": { + "targetContainerName": "snapshot-me", + } + }, + "podTemplate": { + "spec": { + "containers": [ + {"name": "main", "image": "worker"}, + { + "name": "snapshot-me", + "image": "snapshot-worker", + "env": [ + { + "name": "EXISTING_TARGET", + "value": "value", + } + ], + }, + {"name": "unrelated-sidecar", "image": "sidecar"}, + ] + } + }, + }, + ], + }, + }, + ) + + deployment.mount_model_cache_pvc("model-cache", "/models") + expected_after_first_mount = copy.deepcopy(deployment.spec()) + deployment.mount_model_cache_pvc("model-cache", "/models") + + raw_spec = deployment.spec() + assert raw_spec == expected_after_first_mount + assert raw_spec["spec"]["env"] == [ + {"name": "EXISTING", "value": "value"}, + {"name": "HF_HOME", "value": "/models"}, + ] + assert "pvcs" not in raw_spec["spec"] + assert "envs" not in raw_spec["spec"] + + for component in raw_spec["spec"]["components"]: + assert "volumeMounts" not in component + pod_spec = component["podTemplate"]["spec"] + assert pod_spec["volumes"] == [ + { + "name": "model-cache", + "persistentVolumeClaim": {"claimName": "model-cache"}, + } + ] + main = next( + container + for container in pod_spec["containers"] + if container["name"] == "main" + ) + assert main["volumeMounts"] == [{"name": "model-cache", "mountPath": "/models"}] + assert "env" not in main + + if component["name"] == "Worker": + checkpoint_target = next( + container + for container in pod_spec["containers"] + if container["name"] == "snapshot-me" + ) + assert checkpoint_target["env"] == [ + {"name": "EXISTING_TARGET", "value": "value"}, + {"name": "HF_HOME", "value": "/models"}, + ] + assert checkpoint_target["volumeMounts"] == [ + {"name": "model-cache", "mountPath": "/models"} + ] + + unrelated_sidecar = next( + container + for container in pod_spec["containers"] + if container["name"] == "unrelated-sidecar" + ) + assert "env" not in unrelated_sidecar + assert "volumeMounts" not in unrelated_sidecar + + +def test_mount_model_cache_pvc_rejects_missing_checkpoint_target( + tmp_path: Path, +) -> None: + deployment = _deployment_spec( + tmp_path, + { + "apiVersion": "nvidia.com/v1beta1", + "kind": "DynamoGraphDeployment", + "metadata": {"name": "test"}, + "spec": { + "components": [ + { + "name": "Worker", + "experimental": { + "checkpoint": { + "targetContainerName": "missing", + } + }, + "podTemplate": { + "spec": { + "containers": [{"name": "main", "image": "worker"}] + } + }, + } + ] + }, + }, + ) + + with pytest.raises( + ValueError, + match=( + "v1beta1 component 'Worker' checkpoint targetContainerName 'missing' " + "does not name a podTemplate container" + ), + ): + deployment.mount_model_cache_pvc("model-cache", "/models") + + assert deployment.spec()["spec"] == { + "components": [ + { + "name": "Worker", + "experimental": { + "checkpoint": { + "targetContainerName": "missing", + } + }, + "podTemplate": { + "spec": { + "containers": [{"name": "main", "image": "worker"}], + } + }, + } + ] + } + + +def test_mount_model_cache_pvc_preserves_v1alpha1_fields( + tmp_path: Path, +) -> None: + deployment = _deployment_spec( + tmp_path, + { + "apiVersion": "nvidia.com/v1alpha1", + "kind": "DynamoGraphDeployment", + "metadata": {"name": "test"}, + "spec": { + "services": { + "Frontend": {"componentType": "frontend"}, + "Worker": {"componentType": "worker"}, + } + }, + }, + ) + + deployment.mount_model_cache_pvc("model-cache", "/models") + + assert deployment.spec() == { + "apiVersion": "nvidia.com/v1alpha1", + "kind": "DynamoGraphDeployment", + "metadata": {"name": "test"}, + "spec": { + "services": { + "Frontend": { + "componentType": "frontend", + "volumeMounts": [{"name": "model-cache", "mountPoint": "/models"}], + }, + "Worker": { + "componentType": "worker", + "volumeMounts": [{"name": "model-cache", "mountPoint": "/models"}], + }, + }, + "pvcs": [{"name": "model-cache", "create": False}], + "envs": [{"name": "HF_HOME", "value": "/models"}], + }, + }