name: PR Test - Build Rust Extensions

# Builds the rust/ workspace's PyO3 extension modules once so the CUDA test
# stages install with SGLANG_BUILD_RUST_EXTS=none instead of each rebuilding them.
#
# Why: uv holds the editable sdist lock in the shared ~/.cache/uv for the whole
# build, cargo included, so CUDA jobs on one host serialize on it - and once the
# Rust build passed a few minutes, the queue's tail hit uv's lock timeout.
#
# The win is compiling once per run instead of in all ~25 install steps; the two
# layers below only decide who pays for it.
# - The cache, keyed by a source hash, carries a build across runs - no version to
#   bump, unlike sgl-kernel, which publishes one because it ships to users.
# - The artifact hands it to the stages, which the cache cannot: it is best-effort,
#   evictable, and this repo is at its 10 GB limit. Artifacts are durable per run.
#
# Compiling is split off from restoring and collecting because only it needs the
# build node, and queueing there on a cache hit would put that wait in front of
# every stage. Republishing bytes needs no particular host - the glibc a module
# requires is recorded in the module, not decided by whoever uploads it.

on:
  workflow_call:
    inputs:
      runs_on:
        description: 'Runner label for the compile job. Its glibc has to satisfy max_glibc.'
        type: string
        required: true
      restore_runs_on:
        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:
        description: 'Artifact name. Suffix it per caller: artifacts are immutable per name per run, so two callers sharing a run would collide.'
        type: string
        default: rust-ext-x86_64
      cache_key_prefix:
        description: 'Cache key prefix. Callers share it on purpose to reuse each other''s build; it encodes the arch and the interpreter ABI set, since the modules are portable across neither.'
        type: string
        default: rust-ext-x86_64-cp310-cp312
      max_glibc:
        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:
        type: string
        default: ''
      skip_pr_test_health_check:
        description: 'Forwarded from the caller for the check-maintenance action.'
        type: boolean
        default: false
    outputs:
      artifact_name:
        description: 'Artifact holding the built modules. Empty when neither job published one, whether skipped or failed; consumers that still run compile them during install.'
        value: ${{ jobs.restore.outputs.artifact_name || jobs.collect.outputs.artifact_name }}

# Reusable workflows do not inherit the caller's env; mirror what
# check-maintenance reads.
env:
  SGLANG_IS_IN_CI: true
  SKIP_PR_TEST_HEALTH_CHECK: ${{ inputs.skip_pr_test_health_check && 'true' || 'false' }}
  PR_TEST_BYPASS_MAINTENANCE_ON_MAIN: ${{ github.ref == 'refs/heads/main' && 'true' || 'false' }}

jobs:
  restore:
    runs-on: ${{ inputs.restore_runs_on }}
    timeout-minutes: 15
    name: Restore Rust Ext
    outputs:
      hit: ${{ steps.cache.outputs.cache-hit }}
      # From the last step, so a failed job publishes no name for always() consumers.
      artifact_name: ${{ steps.publish.outputs.name }}
    steps:
      - uses: actions/checkout@v4
        with:
          ref: ${{ inputs.git_ref || github.sha }}
          # Just what the cache key hashes, plus the action and script this job
          # runs: the workspace is cold here and the rest of the tree is mostly
          # docs. Both jobs must hash the same Rust extension inputs, which this
          # preserves.
          # Cone mode off is what allows naming a single file.
          sparse-checkout: |
            rust
            proto
            python/setup.py
            python/pyproject.toml
            python/sglang/srt/rust_extensions/torch_build.py
            .github
            scripts/ci/utils
          sparse-checkout-cone-mode: false

      - uses: ./.github/actions/check-maintenance

      # A no-op on the hosted default, which ships zstd; here for whatever
      # restore_runs_on is pointed at, since a reader without it sees no entry.
      - name: Ensure zstd so the saved entry is readable
        run: bash scripts/ci/utils/ensure_zstd.sh

      # The setup hook and torch helper select and configure what gets built;
      # pyproject.toml pins the libtorch ABI used by sglang-radix-tree.
      - name: Restore built modules
        id: cache
        uses: actions/cache/restore@v4
        with:
          path: |
            python/sglang/srt/rust_extensions/_*.so
            python/sglang/srt/mem_cache/rust_tree_core/mem_cache*.so
          key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'proto/**', 'python/setup.py', 'python/pyproject.toml', 'python/sglang/srt/rust_extensions/torch_build.py', '.github/workflows/_pr-test-rust-ext-build.yml', 'scripts/ci/utils/stage_rust_ext_modules.sh') }}

      # On a miss: different hash = rust/setup.py moved; no entries = evicted.
      - name: Report cache lookup
        env:
          GH_TOKEN: ${{ github.token }}
          PRIMARY_KEY: ${{ steps.cache.outputs.cache-primary-key }}
          MATCHED_KEY: ${{ steps.cache.outputs.cache-matched-key }}
          KEY_PREFIX: ${{ inputs.cache_key_prefix }}
        run: |
          if [ -n "${MATCHED_KEY}" ]; then
              echo "hit: ${MATCHED_KEY}"
              ls -l python/sglang/srt/rust_extensions/_*.so \
                python/sglang/srt/mem_cache/rust_tree_core/mem_cache*.so
          else
              echo "miss: ${PRIMARY_KEY}"
              echo "entries under ${KEY_PREFIX}- (created / ref / size / key):"
              gh cache list --repo "${GITHUB_REPOSITORY}" --key "${KEY_PREFIX}-" \
                  --limit 15 --json createdAt,ref,sizeInBytes,key \
                  --jq '.[] | [.createdAt, .ref, ((.sizeInBytes / 1048576 | floor | tostring) + " MiB"), .key] | @tsv' \
                  || echo "(gh cache list unavailable: token lacks actions:read)"
          fi

      # No MAX_GLIBC: these are the bytes the collect job already checked before
      # saving them under this key. The module count is still worth re-checking,
      # so a truncated entry fails here rather than as a test import error.
      - name: Stage modules for upload
        if: steps.cache.outputs.cache-hit == 'true'
        run: bash scripts/ci/utils/stage_rust_ext_modules.sh

      - name: Upload extension modules
        if: steps.cache.outputs.cache-hit == 'true'
        uses: actions/upload-artifact@v4
        with:
          name: ${{ inputs.artifact_name }}
          # Archive holds package-relative paths, so it unpacks into python/sglang/srt/.
          path: rust-ext-staging/
          if-no-files-found: error
          retention-days: 1

      - name: Publish artifact name
        id: publish
        if: steps.cache.outputs.cache-hit == 'true'
        run: echo "name=${{ inputs.artifact_name }}" >> "$GITHUB_OUTPUT"

  compile:
    needs: restore
    if: needs.restore.outputs.hit != 'true'
    runs-on: ${{ inputs.runs_on }}
    timeout-minutes: 60
    name: Build Rust Ext
    # No crate sets abi3, so the pools' interpreters (h100 ships 3.10, h20 ships
    # 3.12) each need their own module set, and each its own target dir, since
    # PyO3's fingerprint tracks the interpreter. The two passes therefore share
    # nothing, which is why they run as parallel jobs rather than back to back.
    strategy:
      # Both legs run to completion, so a failure in one still leaves the other's
      # log to compare against.
      fail-fast: false
      matrix:
        python-version: ['3.10', '3.12']
    steps:
      - uses: actions/checkout@v4
        with:
          ref: ${{ inputs.git_ref || github.sha }}

      - uses: ./.github/actions/check-maintenance

      - name: Set up Python ${{ matrix.python-version }}
        id: py
        uses: actions/setup-python@v5
        with:
          python-version: ${{ matrix.python-version }}

      - name: Install protoc and Rust toolchain
        run: bash scripts/ci/utils/install_rust_protoc.sh

      - name: Build extension modules
        run: |
          set -euxo pipefail
          export PATH="${CARGO_HOME:-$HOME/.cargo}/bin:${PATH}"
          # A hosted runner is fresh per job, so there is no shared tree to reuse or
          # to guard against a concurrent dropper.
          export CARGO_TARGET_DIR="${RUNNER_TEMP:-/tmp}/sglang-cargo-target-${GITHUB_RUN_ID:-norun}-$$"
          mkdir -p "${CARGO_TARGET_DIR}"
          python3 -m pip install --upgrade pip
          command -v uv >/dev/null 2>&1 || pip install uv
          # Per-job path: a self-hosted runs_on is persistent and shared, so a fixed
          # one both inherits the previous job's venv and races a concurrent build.
          venv="${RUNNER_TEMP:-/tmp}/sglang-ci-rust-ext-${GITHUB_RUN_ID:-norun}-$$"
          # Best-effort, like ci_cleanup_venv.sh: under set -e a failing EXIT trap
          # would fail the step, and nothing here is worth keeping for a postmortem.
          trap 'rm -rf "${venv}" "${CARGO_TARGET_DIR}" || true' EXIT
          # build_rust needs the build backend and torch, not sglang's ~294 other runtime deps.
          uv venv "${venv}" --python "${{ steps.py.outputs.python-path }}" --seed
          # shellcheck disable=SC1091
          source "${venv}/bin/activate"
          uv pip install "setuptools>=61.0" "setuptools-rust>=1.11" "setuptools-scm>=8.0" "torch==2.13.0" wheel
          # torch-sys resolves libtorch from the active interpreter and bakes that
          # path into the cargo cache. venv is per-run and deleted on exit, so a
          # persistent target dir would reuse a "Fresh torch-sys" whose -L points at
          # a gone directory and fail with "unable to find library -ltorch". Pin
          # LIBTORCH and drop the torch-shim units so they rebuild against this venv.
          LIBTORCH="$(python -c 'import pathlib, torch; print(pathlib.Path(torch.__file__).parent)')"
          export LIBTORCH
          export LD_LIBRARY_PATH="${LIBTORCH}/lib:${LD_LIBRARY_PATH:-}"
          cargo clean --release --manifest-path rust/sglang-radix-tree/Cargo.toml \
              -p torch-sys -p sglang-radix-tree 2>/dev/null || true
          (cd python && SGLANG_BUILD_RUST_EXTS=all python setup.py build_rust --inplace)
          python - <<'PY'
          import importlib.util
          import os
          import pathlib
          import runpy
          import shutil
          import subprocess
          import sysconfig

          import torch  # noqa: F401 - preload libtorch before the extension

          root = pathlib.Path.cwd()
          suffix = sysconfig.get_config_var("EXT_SUFFIX")
          production_path = pathlib.Path(
              "python/sglang/srt/mem_cache/rust_tree_core/mem_cache" + suffix
          ).resolve()
          inspection_path = production_path.with_name(
              "mem_cache_inspection" + suffix
          )

          helper = runpy.run_path(
              root / "python/sglang/srt/rust_extensions/torch_build.py"
          )
          build = helper["torch_build_configuration"](
              compat_header=root / "rust/sglang-radix-tree/torch_2_13_compat.h",
              python_module="sglang.srt.mem_cache.rust_tree_core.mem_cache",
              torch_module=torch,
              include_absolute_rpath=False,
          )
          subprocess.run(
              [
                  "cargo",
                  "build",
                  "--release",
                  "--locked",
                  "--manifest-path",
                  "rust/sglang-radix-tree/Cargo.toml",
                  "--features",
                  "python-extension,inspection",
              ],
              env=build.environment,
              check=True,
          )
          release_dir = pathlib.Path(os.environ["CARGO_TARGET_DIR"]) / "release"
          if target := os.environ.get("CARGO_BUILD_TARGET"):
              release_dir = pathlib.Path(os.environ["CARGO_TARGET_DIR"]) / target / "release"
          shutil.copy2(release_dir / "libmem_cache.so", inspection_path)

          def load(name, path):
              spec = importlib.util.spec_from_file_location(name, path)
              assert spec is not None and spec.loader is not None
              module = importlib.util.module_from_spec(spec)
              spec.loader.exec_module(module)
              return module

          production = load(
              "sglang.srt.mem_cache.rust_tree_core.mem_cache", production_path
          )
          inspection = load(
              "sglang.srt.mem_cache.rust_tree_core.mem_cache_inspection",
              inspection_path,
          )
          assert not hasattr(
              production.RustUnifiedTreeCoreBinding, "inspect_contains_node"
          )
          assert hasattr(
              inspection.RustUnifiedTreeCoreBinding, "inspect_contains_node"
          )
          PY

      # Both globs sit under python/sglang/srt/, which upload-artifact strips as the
      # common prefix, so the collect job restores them by unpacking back into it.
      - name: Upload this interpreter's modules
        uses: actions/upload-artifact@v4
        with:
          name: ${{ inputs.artifact_name }}-build-py${{ matrix.python-version }}
          path: |
            python/sglang/srt/rust_extensions/_*.so
            python/sglang/srt/mem_cache/rust_tree_core/mem_cache*.so
          if-no-files-found: error
          retention-days: 1

  # The interpreters compile independently, but the suffix-set check and the cache
  # entry both cover the whole set, so they join in one job that sees all of it.
  collect:
    needs: compile
    runs-on: ${{ inputs.collect_runs_on }}
    timeout-minutes: 15
    name: Collect Rust Ext
    outputs:
      # From the last step, so a failed job publishes no name for always() consumers.
      artifact_name: ${{ steps.publish.outputs.name }}
    steps:
      - uses: actions/checkout@v4
        with:
          ref: ${{ inputs.git_ref || github.sha }}
          # Mirrors the restore job: this one writes the entry that job reads, so
          # both have to check out the same set for hashFiles to agree.
          sparse-checkout: |
            rust
            proto
            python/setup.py
            python/pyproject.toml
            python/sglang/srt/rust_extensions/torch_build.py
            .github
            scripts/ci/utils
          sparse-checkout-cone-mode: false

      - name: Collect the per-interpreter modules
        uses: actions/download-artifact@v4
        with:
          pattern: ${{ inputs.artifact_name }}-build-py*
          merge-multiple: true
          path: python/sglang/srt

      # First point that sees every interpreter, which is what the suffix-set check
      # compares. MAX_GLIBC belongs here rather than in the restore job: these are
      # freshly compiled bytes, not ones that already passed under this key.
      - name: Verify modules and stage for upload
        env:
          MAX_GLIBC: ${{ inputs.max_glibc }}
        run: bash scripts/ci/utils/stage_rust_ext_modules.sh

      # Without it the entry saved below lands under a version no reader with zstd
      # can find. This runner's image ships without it.
      - name: Ensure zstd so the restore job can read what this job saves
        run: bash scripts/ci/utils/ensure_zstd.sh

      # After the verify step, so a rejected build cannot poison this key for every
      # later run.
      - name: Save built modules
        uses: actions/cache/save@v4
        with:
          path: |
            python/sglang/srt/rust_extensions/_*.so
            python/sglang/srt/mem_cache/rust_tree_core/mem_cache*.so
          key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'proto/**', 'python/setup.py', 'python/pyproject.toml', 'python/sglang/srt/rust_extensions/torch_build.py', '.github/workflows/_pr-test-rust-ext-build.yml', 'scripts/ci/utils/stage_rust_ext_modules.sh') }}

      - name: Upload extension modules
        uses: actions/upload-artifact@v4
        with:
          name: ${{ inputs.artifact_name }}
          # Archive holds package-relative paths, so it unpacks into python/sglang/srt/.
          path: rust-ext-staging/
          if-no-files-found: error
          retention-days: 1

      - name: Publish artifact name
        id: publish
        run: echo "name=${{ inputs.artifact_name }}" >> "$GITHUB_OUTPUT"