[CI] Build the Rust extensions for aarch64 too (#37820)

Co-authored-by: Mohammad Angkad <mohammad.angkad@radixark.ai>

.github/workflows/_pr-test-rust-ext-build.yml

@@ -27,7 +27,16 @@ on:
         type: string
         required: true
       restore_runs_on:
-        description: 'Runner label for the cache-hit path and the collect job. Wants free capacity and nothing else, since neither compiles nor imports the modules.'
+        description: >
+          Runner label for the cache-hit path. Only republishes bytes, so it wants
+          free capacity and is the one job here with no architecture requirement.
+        type: string
+        default: ubuntu-latest
+      collect_runs_on:
+        description: >
+          Runner label for the collect job. Must be runs_on's ARCHITECTURE: its GLIBC
+          gate reads the .so with objdump, and the hosted images ship a single-target
+          binutils.
         type: string
         default: ubuntu-latest
       artifact_name:
@@ -39,7 +48,10 @@ on:
         type: string
         default: rust-ext-x86_64-cp310-cp312
       max_glibc:
-        description: 'Highest GLIBC symbol version the built .so files may require. Set by the oldest test runner image, jammy at glibc 2.35 - the pools are not all on one image.'
+        description: >
+          Highest GLIBC symbol version the built .so files may require. A ceiling set
+          by the oldest image that has to LOAD them, not by whoever builds them; the
+          pools are not all on one image, and x86_64's oldest is jammy at 2.35.
         type: string
         default: '2.35'
       git_ref:
@@ -298,7 +310,7 @@ jobs:
   # entry both cover the whole set, so they join in one job that sees all of it.
   collect:
     needs: compile
-    runs-on: ${{ inputs.restore_runs_on }}
+    runs-on: ${{ inputs.collect_runs_on }}
     timeout-minutes: 15
     name: Collect Rust Ext
     outputs:

.github/workflows/_pr-test-stage.yml

@@ -59,7 +59,10 @@ on:
         type: string
         default: ''
       skip_prebuilt_rust_ext:
-        description: 'Take neither the artifact nor the cache, and compile during install. For a stage no prebuild targets: the cache key is arch-blind, so an aarch64 stage would be handed x86_64 modules.'
+        description: >
+          Take neither the artifact nor the cache, and compile during install. No
+          stage sets it: both architectures have a prebuild. Kept as the opt-out for
+          a pool that gets one later than its stages do.
         type: boolean
         default: false
 

.github/workflows/nightly-test-nvidia.yml

@@ -152,8 +152,6 @@ jobs:
       run_timeout_minutes: '360'
       job_timeout_minutes: '420'
       scheduled: true
-      # aarch64: the rust-ext cache key is x86_64-only, same reason base-c skips it.
-      skip_prebuilt_rust_ext: true
     secrets: inherit
 
   nightly-8-gpu-h200:

.github/workflows/pr-test.yml

@@ -231,6 +231,26 @@ jobs:
       skip_pr_test_health_check: ${{ inputs.skip_pr_test_health_check == true }}
     secrets: inherit
 
+  # gb300 is Grace, and PyO3 modules are portable across neither arch nor
+  # interpreter, so it can use nothing the job above builds.
+  rust-ext-build-aarch64:
+    needs: [check-changes, call-gate]
+    # Same gate as rust-ext-build: every main-ref run has to seed the cache, and a
+    # docs-only PR still installs sglang on gb300.
+    if: |
+      !cancelled() &&
+      needs.check-changes.result == 'success' &&
+      (needs.call-gate.result == 'success' || needs.call-gate.result == 'skipped')
+    uses: ./.github/workflows/_pr-test-rust-ext-build.yml
+    with:
+      runs_on: ubuntu-22.04-arm
+      collect_runs_on: ubuntu-22.04-arm
+      artifact_name: rust-ext-aarch64
+      cache_key_prefix: rust-ext-aarch64-cp310-cp312
+      git_ref: ${{ inputs.git_ref || '' }}
+      skip_pr_test_health_check: ${{ inputs.skip_pr_test_health_check == true }}
+    secrets: inherit
+
   call-sgl-kernel-tests:
     needs: [check-changes, call-gate, sgl-kernel-build-wheels, rust-ext-build]
     if: |
@@ -455,7 +475,7 @@ jobs:
     secrets: inherit
 
   base-c-test-4-gpu-gb300:
-    needs: [check-changes, call-gate, wait-for-base-b, sgl-kernel-build-wheels, rust-ext-build]
+    needs: [check-changes, call-gate, wait-for-base-b, sgl-kernel-build-wheels, rust-ext-build-aarch64]
     if: ${{ !failure() && !cancelled() }}
     uses: ./.github/workflows/_pr-test-stage.yml
     with:
@@ -466,9 +486,8 @@ jobs:
       partitions: ${{ needs.check-changes.outputs.partitions }}
       run_timeout_minutes: '30'
       timeout_per_file: '1800'
-      # aarch64 has no prebuild, so this stage compiles during install - which also
-      # makes it the only stage still covering that path. Keep it that way.
-      skip_prebuilt_rust_ext: true
+      # The one aarch64 stage, so it takes the aarch64 build, not rust-ext-build's.
+      rust_ext_artifact: ${{ needs.rust-ext-build-aarch64.outputs.artifact_name }}
     secrets: inherit
 
   base-c-test-8-gpu-b300:
@@ -499,6 +518,7 @@ jobs:
         call-sgl-kernel-tests,
 
         rust-ext-build,
+        rust-ext-build-aarch64,
 
         wait-for-base-a,
         wait-for-base-b,

.github/workflows/seed-rust-ext-cache.yml

@@ -15,6 +15,10 @@ on:
       # The key's prefix lives in this file, so a bump there also moves the key.
       - '.github/workflows/_pr-test-rust-ext-build.yml'
       - 'scripts/ci/utils/stage_rust_ext_modules.sh'
+      # Not hashed into the key, but they carry the prefix each side uses: the one
+      # seeded per arch here, and the one consumers restore with there.
+      - '.github/workflows/seed-rust-ext-cache.yml'
+      - '.github/actions/download-rust-ext/action.yml'
   workflow_dispatch:
 
 # Only the newest merge needs to seed; earlier ones are already stale.
@@ -35,3 +39,13 @@ jobs:
     with:
       runs_on: ubuntu-22.04
     secrets: inherit
+
+  # One entry per arch, since the arch is part of the key.
+  seed-aarch64:
+    uses: ./.github/workflows/_pr-test-rust-ext-build.yml
+    with:
+      runs_on: ubuntu-22.04-arm
+      collect_runs_on: ubuntu-22.04-arm
+      artifact_name: rust-ext-aarch64
+      cache_key_prefix: rust-ext-aarch64-cp310-cp312
+    secrets: inherit

scripts/ci/utils/stage_rust_ext_modules.sh

@@ -55,12 +55,23 @@ max_allowed="${MAX_GLIBC:-}"
 # Newer glibc than the test runners fails at import: "GLIBC_2.xx not found".
 status=0
 for so in "${built[@]}"; do
-    # grep exits 1 with no match; without `|| true` pipefail kills the script.
-    needed=$(objdump -T "$so" \
+    # Its own invocation, not the head of a pipeline: there its failure is
+    # invisible, and the empty symbol list a wrong-arch objdump leaves behind reads
+    # exactly like "needs no glibc", passing this gate silently.
+    if ! symbols=$(objdump -T "$so" 2>&1); then
+        echo "::error::objdump could not read ${so}: ${symbols}"
+        echo "::error::this gate must run on the same architecture as the modules"
+        status=1
+        continue
+    fi
+    # grep exits 1 with no match; without `|| true` pipefail kills the script. With
+    # objdump known to have succeeded, no match really is no requirement.
+    references=$(printf '%s\n' "${symbols}" | grep -coE 'GLIBC_2\.[0-9]+' || true)
+    needed=$(printf '%s\n' "${symbols}" \
         | grep -oE 'GLIBC_2\.[0-9]+' \
         | sed 's/GLIBC_//' \
         | sort -V | tail -1 || true)
-    echo "${so}: requires glibc <= ${needed:-none}"
+    echo "${so}: requires glibc <= ${needed:-none} (${references} GLIBC references)"
     if [ -n "${needed}" ] \
        && [ "$(printf '%s\n%s\n' "${max_allowed}" "${needed}" | sort -V | tail -1)" != "${max_allowed}" ]; then
         echo "::error::${so} requires glibc ${needed} > ${max_allowed} supported by the test runners; build on an older image"