fix(test): mount model-cache PVC on v1beta1 pod templates

.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

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(

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})