ci(docker): PR-side image build and run the GPU matrix inside the freshly built image

.github/workflows/docker-build-pr.yml

@@ -0,0 +1,65 @@
+name: Docker Build (PR check)
+
+on:
+  pull_request:
+    paths:
+      - 'docker/Dockerfile'
+      - 'docker/build.py'
+      - 'docker/patch/**'
+      - 'requirements.txt'
+      - '.github/workflows/docker-build-pr.yml'
+
+concurrency:
+  group: docker-build-pr-${{ github.event.pull_request.number }}
+  cancel-in-progress: true
+
+jobs:
+  build-only:
+    runs-on: ["h200", "2gpu"]
+    timeout-minutes: 180
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+        with:
+          driver-opts: |
+            image=moby/buildkit:latest
+            network=host
+
+      - name: Install Python + dependencies
+        run: |
+          apt-get update && apt-get install -y python3 python3-pip
+          pip3 install --break-system-packages typer
+
+      - name: Login to Docker Hub
+        # Fork PRs have no secrets: build without pushing, skip the test stage.
+        if: github.event.pull_request.head.repo.full_name == github.repository
+        uses: docker/login-action@v3
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_TOKEN }}
+
+      - name: Build (and push a PR-scoped tag for the test stage)
+        run: |
+          if [ "${{ github.event.pull_request.head.repo.full_name }}" = "${{ github.repository }}" ]; then
+            python3 docker/build.py --variant cu13-x86 --image-tag custom \
+              --custom-tag pr-${{ github.event.pull_request.number }} --push
+          else
+            python3 docker/build.py --variant cu13-x86 --image-tag dev
+          fi
+
+  # Smoke the freshly built image with a real GPU suite so a Dockerfile change
+  # is validated by tests, not only by "the build finished".
+  test-in-fresh-image:
+    needs: [build-only]
+    if: github.event.pull_request.head.repo.full_name == github.repository
+    uses: ./.github/workflows/_run-ci.yml
+    with:
+      container_image: radixark/miles:pr-${{ github.event.pull_request.number }}
+      runs_on: '["h200", "2gpu"]'
+      execute_command: >-
+        python tests/ci/run_suite.py --hw h200 --suite stage-b-2-gpu-h200
+        --cadence pr --labels ''
+    secrets: inherit

docker/Dockerfile

@@ -120,7 +120,7 @@ RUN if [ "${ENABLE_CUDA_13}" = "1" ]; then \
 # https://github.com/NVIDIA/TransformerEngine/pull/3141; drop it after TE v2.18.
 COPY docker/patch/ /tmp/patches/
 RUN if [ "${ENABLE_CUDA_13}" = "1" ] && [ -d /tmp/patches/cu13 ]; then \
-      TE_DIR=$(python -c 'import transformer_engine; print(transformer_engine.__path__[0])') && \
+      TE_DIR=$(python -c 'import importlib.util; print(importlib.util.find_spec("transformer_engine").submodule_search_locations[0])') && \
       for p in /tmp/patches/cu13/*.patch; do \
         echo "Applying $(basename $p) to $TE_DIR" && \
         patch -d "$TE_DIR" -p1 < "$p"; \