Merge remote-tracking branch 'origin/main' into onur/remove-trt-llm

Signed-off-by: Charlie Truong <chtruong@nvidia.com>

Author: chtruong814

.github/actions/test-template/action.yml

@@ -34,57 +34,25 @@ inputs:
     description: "Run tests on CPU only"
     required: false
     default: "false"
-  azure-client-id:
-    description: "Azure Client ID"
-    required: true
-  azure-tenant-id:
-    description: "Azure Tenant ID"
-    required: true
-  azure-subscription-id:
-    description: "Azure Subscription ID"
-    required: true
-  has-azure-credentials:
-    description: "Has Azure credentials"
-    required: false
-    default: "false"
   PAT:
     description: "GitHub Personal Access Token"
     required: true
+  container-image:
+    description: "Pre-built container image to use for test"
+    required: true
   inference-framework:
-    description: "Inference Framework"
+    description: "Inference framework extra to install at runtime (trtllm, vllm, trt-onnx, inframework)"
     required: true
   test-data-path:
     description: "Test data path"
     required: true
   runner:
     description: "The GHA runner to use"
     required: true
-  ngc-api-user:
-    description: "NGC API User"
-    required: false
-  ngc-api-key:
-    description: "NGC API Key"
-    required: false
 
 runs:
   using: "composite"
   steps:
-    - name: Install Azure CLI
-      if: ${{ inputs.has-azure-credentials == 'true' }}
-      shell: bash
-      run: |
-        echo ::group::Install Azure CLI
-        curl -sL https://aka.ms/InstallAzureCLIDeb | sudo bash
-        echo ::endgroup::
-
-    - name: Azure Login
-      if: ${{ inputs.has-azure-credentials == 'true' }}
-      uses: azure/login@v2
-      with:
-        client-id: ${{ inputs.azure-client-id }}
-        tenant-id: ${{ inputs.azure-tenant-id }}
-        subscription-id: ${{ inputs.azure-subscription-id }}
-
     - name: Install uuidgen
       shell: bash -x -e -u -o pipefail {0}
       if: ${{ contains(inputs.runner, 'aws') }}
@@ -98,37 +66,6 @@ runs:
         path: Export-Deploy
 
 
-    - name: Cache uv
-      uses: actions/cache@v4
-      id: cache
-      with:
-        path: cache-mount
-        key: ${{ runner.os }}-uv-${{ hashFiles('**/uv.lock') }}
-
-    - name: Login to NGC
-      if: ${{ inputs.ngc-api-user != '' && inputs.ngc-api-key != '' }}
-      shell: bash
-      run: |
-        echo ::group::Use credentials
-        echo ${{ inputs.ngc-api-key }} | docker login -u '${{ inputs.ngc-api-user }}' --password-stdin nvcr.io
-        echo ::endgroup::
-
-    - name: Restore Docker cache mounts
-      uses: reproducible-containers/buildkit-cache-dance@5b81f4d29dc8397a7d341dba3aeecc7ec54d6361
-      with:
-        cache-dir: cache-mount
-        dockerfile: docker/Dockerfile.pytorch
-        skip-extraction: ${{ steps.cache.outputs.cache-hit }}
-
-    - name: Build container
-      shell: bash
-      env:
-        GH_TOKEN: ${{ inputs.PAT }}
-      run: |
-        echo ::group::Build test container
-        docker build -f docker/Dockerfile.pytorch --build-arg INFERENCE_FRAMEWORK=${{ inputs.inference-framework }} --secret id=GH_TOKEN -t export-deploy .
-        echo ::endgroup::
-
     - name: Start container
       shell: bash
       run: |
@@ -154,7 +91,7 @@ runs:
           --workdir /workspace \
           --volume $(pwd)/Export-Deploy:/workspace \
           --volume ${{ inputs.test-data-path }}:/home/TestData \
-          export-deploy \
+          ${{ inputs.container-image }} \
           bash -c "sleep $(( ${{ inputs.timeout }} * 60 + 60 ))"
         RUN_TEST_EOF
         )
@@ -163,6 +100,45 @@ runs:
         bash retry_job.sh
         echo ::endgroup::
 
+    - name: Install framework deps
+      if: ${{ inputs.inference-framework != 'inframework' }}
+      shell: bash
+      run: |
+        echo ::group::Install ${{ inputs.inference-framework }} dependencies
+        docker exec nemo_container_${{ github.run_id }} bash -c '
+          uv sync --extra ${{ inputs.inference-framework }} --locked --all-groups \
+            --no-install-package torch \
+            --no-install-package torchvision \
+            --no-install-package triton \
+            --no-install-package nvidia-cublas-cu12 \
+            --no-install-package nvidia-cuda-cupti-cu12 \
+            --no-install-package nvidia-cuda-nvrtc-cu12 \
+            --no-install-package nvidia-cuda-runtime-cu12 \
+            --no-install-package nvidia-cudnn-cu12 \
+            --no-install-package nvidia-cufft-cu12 \
+            --no-install-package nvidia-cufile-cu12 \
+            --no-install-package nvidia-curand-cu12 \
+            --no-install-package nvidia-cusolver-cu12 \
+            --no-install-package nvidia-cusparse-cu12 \
+            --no-install-package nvidia-cusparselt-cu12 \
+            --no-install-package nvidia-nccl-cu12 \
+            --no-install-package nvidia-cublas \
+            --no-install-package nvidia-cuda-cupti \
+            --no-install-package nvidia-cuda-nvrtc \
+            --no-install-package nvidia-cuda-runtime \
+            --no-install-package nvidia-cudnn-cu13 \
+            --no-install-package nvidia-cufft \
+            --no-install-package nvidia-cufile \
+            --no-install-package nvidia-curand \
+            --no-install-package nvidia-cusolver \
+            --no-install-package nvidia-cusparse \
+            --no-install-package nvidia-cusparselt-cu13 \
+            --no-install-package nvidia-nccl-cu13 \
+            --no-install-package tensorrt-llm \
+            --no-install-package vllm
+        '
+        echo ::endgroup::
+
     - name: Create run-script
       id: create
       shell: bash

.github/workflows/cicd-main.yml

@@ -33,10 +33,12 @@ jobs:
   pre-flight:
     uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_cicd_preflight.yml@v0.80.1
     with:
-      default_runner_prefix: ${{ vars.DEFAULT_RUNNER_PREFIX }}
-      non_nvidia_runner_prefix: ${{ vars.NON_NVIDIA_RUNNER_PREFIX }}
-      default_test_data_path: ${{ vars.DEFAULT_TEST_DATA_PATH }}
-      non_nvidia_test_data_path: ${{ vars.NON_NVIDIA_TEST_DATA_PATH }}
+      default_runner_prefix: nemo-ci-aws-gpu-x2
+      non_nvidia_runner_prefix: nemo-ci-aws-gpu-x2-ephemeral
+      default_test_data_path: /mnt/datadrive/TestData/nemo-fw/TestData
+      non_nvidia_test_data_path: /mnt/datadrive/TestData/nemo-fw/TestData
+      default_registry: 766267172432.dkr.ecr.us-east-1.amazonaws.com
+      non_nvidia_registry: 766267172432.dkr.ecr.us-east-1.amazonaws.com
       sso_users_filename: ${{ vars.SSO_USERS_FILENAME }}
     secrets:
       NVIDIA_MANAGEMENT_ORG_PAT: ${{ secrets.NVIDIA_MANAGEMENT_ORG_PAT }}
@@ -82,17 +84,77 @@ jobs:
         run: |
           echo "Running CI tests"
 
+  cicd-container-build:
+    needs: [pre-flight, cicd-wait-in-queue]
+    runs-on: ${{ needs.pre-flight.outputs.runner_prefix }}
+    environment: nemo-ci
+    if: |
+      (
+        success()
+        || needs.pre-flight.outputs.is_ci_workload == 'true'
+        || needs.pre-flight.outputs.force_run_all == 'true'
+      )
+      && !cancelled()
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v6
+
+      - name: Login to NGC
+        uses: docker/login-action@v3
+        with:
+          registry: nvcr.io
+          username: ${{ secrets.NGC_API_USER }}
+          password: ${{ secrets.NGC_API_KEY }}
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+
+      - name: Compute cache keys
+        id: cache_keys
+        shell: bash
+        run: |
+          if [[ "${{ github.ref }}" == "refs/heads/main" ]]; then
+            KEY="main"
+          else
+            PR_NUMBER=$(echo "${{ github.ref_name }}" | grep -oP '^pull-request/\K[0-9]+$' || true)
+            if [[ -n "$PR_NUMBER" ]]; then
+              KEY="$PR_NUMBER"
+            else
+              KEY=$(echo "${{ github.ref_name }}" | tr '/' '-' | tr -cd '[:alnum:]._-')
+            fi
+          fi
+          REGISTRY="${{ needs.pre-flight.outputs.registry }}"
+          echo "key=$KEY" | tee -a "$GITHUB_OUTPUT"
+          echo "cache-to=type=registry,ref=${REGISTRY}/export-deploy:${KEY}-buildcache,mode=max" | tee -a "$GITHUB_OUTPUT"
+
+      - name: Build and push
+        uses: docker/build-push-action@v5
+        with:
+          file: ./docker/Dockerfile.pytorch
+          push: true
+          context: .
+          build-args: |
+            INFERENCE_FRAMEWORK=inframework
+          cache-from: |
+            type=registry,ref=${{ needs.pre-flight.outputs.registry }}/export-deploy:${{ steps.cache_keys.outputs.key }}-buildcache
+            type=registry,ref=${{ needs.pre-flight.outputs.registry }}/export-deploy:main-buildcache
+          cache-to: ${{ steps.cache_keys.outputs.cache-to }}
+          tags: |
+            ${{ needs.pre-flight.outputs.registry }}/export-deploy:${{ github.sha }}
+          secrets: |
+            GH_TOKEN=${{ secrets.PAT }}
+
   cicd-unit-tests-vllm:
     strategy:
       fail-fast: false
       matrix:
         include:
           - script: L0_Unit_Tests_GPU
-            runner: ${{ needs.pre-flight.outputs.runner_prefix }}-gpu-x2
+            runner: ${{ needs.pre-flight.outputs.runner_prefix }}
           - script: L0_Unit_Tests_CPU
-            runner: linux-amd64-cpu16
+            runner: ${{ needs.pre-flight.outputs.runner_prefix }}
             cpu-only: true
-    needs: [pre-flight, cicd-wait-in-queue]
+    needs: [pre-flight, cicd-wait-in-queue, cicd-container-build]
     if: |
       (
         success()
@@ -113,16 +175,11 @@ jobs:
           timeout: 60
           is_unit_test: "true"
           cpu-only: ${{ matrix.cpu-only || false }}
-          has-azure-credentials: "true"
-          azure-client-id: ${{ secrets.AZURE_CLIENT_ID }}
-          azure-tenant-id: ${{ secrets.AZURE_TENANT_ID }}
-          azure-subscription-id: ${{ secrets.AZURE_SUBSCRIPTION_ID }}
           PAT: ${{ secrets.PAT }}
+          container-image: ${{ needs.pre-flight.outputs.registry }}/export-deploy:${{ github.sha }}
           inference-framework: vllm
           test-data-path: ${{ needs.pre-flight.outputs.test_data_path }}
           runner: ${{ matrix.runner }}
-          ngc-api-user: ${{ secrets.NGC_API_USER }}
-          ngc-api-key: ${{ secrets.NGC_API_KEY }}
 
   cicd-e2e-tests-trt-onnx:
     needs: [cicd-wait-in-queue, pre-flight]
@@ -146,22 +203,17 @@ jobs:
         with:
           script: L2_ONNX_TRT
           is_optional: ${{ matrix.is_optional || false }}
-          azure-client-id: ${{ secrets.AZURE_CLIENT_ID }}
-          azure-tenant-id: ${{ secrets.AZURE_TENANT_ID }}
-          azure-subscription-id: ${{ secrets.AZURE_SUBSCRIPTION_ID }}
-          has-azure-credentials: true
           is_unit_test: "false"
           timeout: 60
           PAT: ${{ secrets.PAT }}
+          container-image: ${{ needs.pre-flight.outputs.registry }}/export-deploy:${{ github.sha }}
           inference-framework: trt-onnx
           test-data-path: ${{ needs.pre-flight.outputs.test_data_path }}
-          runner: ${{ needs.pre-flight.outputs.runner_prefix }}-gpu-x2
-          ngc-api-user: ${{ secrets.NGC_API_USER }}
-          ngc-api-key: ${{ secrets.NGC_API_KEY }}
+          runner: ${{ needs.pre-flight.outputs.runner_prefix }}
 
   cicd-e2e-tests-vllm:
     needs: [cicd-unit-tests-vllm, pre-flight]
-    runs-on: ${{ needs.pre-flight.outputs.runner_prefix }}-gpu-x2
+    runs-on: ${{ needs.pre-flight.outputs.runner_prefix }}
     name: ${{ matrix.is_optional && 'PLEASEFIXME_' || '' }}${{ matrix.script }}
     environment: nemo-ci
     if: |
@@ -179,18 +231,13 @@ jobs:
         with:
           script: L2_Launch_vLLM
           is_optional: ${{ matrix.is_optional || false }}
-          azure-client-id: ${{ secrets.AZURE_CLIENT_ID }}
-          azure-tenant-id: ${{ secrets.AZURE_TENANT_ID }}
-          azure-subscription-id: ${{ secrets.AZURE_SUBSCRIPTION_ID }}
-          has-azure-credentials: true
           is_unit_test: "false"
           timeout: 60
           PAT: ${{ secrets.PAT }}
+          container-image: ${{ needs.pre-flight.outputs.registry }}/export-deploy:${{ github.sha }}
           inference-framework: vllm
           test-data-path: ${{ needs.pre-flight.outputs.test_data_path }}
-          runner: ${{ needs.pre-flight.outputs.runner_prefix }}-gpu-x2
-          ngc-api-user: ${{ secrets.NGC_API_USER }}
-          ngc-api-key: ${{ secrets.NGC_API_KEY }}
+          runner: ${{ needs.pre-flight.outputs.runner_prefix }}
 
   cicd-e2e-tests-inframework:
     needs: [pre-flight, cicd-unit-tests-vllm]
@@ -212,18 +259,13 @@ jobs:
         with:
           script: L2_Launch_InFramework
           is_optional: ${{ matrix.is_optional || false }}
-          azure-client-id: ${{ secrets.AZURE_CLIENT_ID }}
-          azure-tenant-id: ${{ secrets.AZURE_TENANT_ID }}
-          azure-subscription-id: ${{ secrets.AZURE_SUBSCRIPTION_ID }}
-          has-azure-credentials: true
           is_unit_test: "false"
           timeout: 60
           PAT: ${{ secrets.PAT }}
+          container-image: ${{ needs.pre-flight.outputs.registry }}/export-deploy:${{ github.sha }}
           inference-framework: inframework
           test-data-path: ${{ needs.pre-flight.outputs.test_data_path }}
-          runner: ${{ needs.pre-flight.outputs.runner_prefix }}-gpu-x2
-          ngc-api-user: ${{ secrets.NGC_API_USER }}
-          ngc-api-key: ${{ secrets.NGC_API_KEY }}
+          runner: ${{ needs.pre-flight.outputs.runner_prefix }}
 
   Nemo_CICD_Test:
     needs:

.github/workflows/config/.secrets.baseline

@@ -117,6 +117,7 @@ vllm = [
     { index = "pypi", marker = "platform_machine == 'aarch64'" },
 ]
 megatron-bridge = { git = "https://github.com/NVIDIA-NeMo/Megatron-Bridge.git", rev = "4cb379fade7c6b85610db0f7412b3601ae15dea8" }
+nvidia-resiliency-ext = { git = "https://github.com/NVIDIA/nvidia-resiliency-ext.git", rev = "b2bb3d728a18795807d9f76c535e005a609a1b01" }
 # nemo-toolkit = { git = "https://github.com/NVIDIA/NeMo.git", rev = "main" }
 
 [tool.uv]
@@ -139,6 +140,7 @@ default-groups = ["linting", "build", "test"]
 link-mode = "copy"
 conflicts = [[{ extra = "vllm" }, { extra = "trt-onnx" }]]
 override-dependencies = [
+    "cachetools>=5.0.0",
     "torch; sys_platform == 'never'",
     "torchvision; sys_platform == 'never'",
     "triton; sys_platform == 'never'",

pyproject.toml

@@ -39,9 +39,12 @@ def teardown_method(self):
             self.deploy_proc = None
 
     def test_deploy_ray(self):
-        mbridge_checkpoint_path = "/home/TestData/megatron_bridge/checkpoints/llama3_145m-mbridge_saved-distckpt"
-
         try:
+            # Checkpoint created via AutoBridge.import_ckpt("meta-llama/Llama-3.2-1B", "/path/to/checkpoint")
+            mbridge_checkpoint_path = (
+                "/home/TestData/export-deploy/checkpoints/mbridge/hf_llama32_1B_mbridge/iter_0000000"
+            )
+
             # Run Ray deployment
             self.deploy_proc = subprocess.Popen(
                 [