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(