ci: build container once and share across downstream tests (#661)

Signed-off-by: Charlie Truong <chtruong@nvidia.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
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