_run-ci-rocm.yml

name: CI Job (ROCm)

on:
  workflow_call:
    inputs:
      execute_command:
        type: string
        required: true
      runs_on:
        type: string
        required: true
        description: 'JSON-encoded list of runner labels (all required). e.g. ''["amd", "mi350", "4gpu"]'' picks runners that carry ALL the listed labels.'
      container_image:
        # No default: pr-test-rocm.yml's resolve-ci-image is the single
        # source of the image name, so a dated tag can't rot here.
        type: string
        required: true
      skip_dependency_install:
        # Default true: the ROCm image ships SGLang / Megatron-LM baked in.
        # Pass false to fetch and install the refs named by ci-megatron-pr: /
        # ci-sglang-pr:; whichever is unnamed keeps the image's own copy.
        type: boolean
        required: false
        default: true
      ref:
        description: 'Git ref to check out (empty = the triggering commit). Set by release CI to test a release branch.'
        type: string
        required: false
        default: ''
    secrets:
      WANDB_API_KEY:
        required: false

jobs:
  plan:
    runs-on: ubuntu-latest
    permissions:
      contents: read
    outputs:
      has_tests: ${{ steps.select.outputs.has_tests }}
      ref: ${{ steps.select.outputs.ref }}
    steps:
      - uses: actions/checkout@v4
        with:
          ref: ${{ inputs.ref }}
          persist-credentials: false
      - uses: actions/setup-python@v5
        with:
          python-version: '3.12'
      - name: Resolve suite plan
        id: select
        env:
          PYTHONPATH: ${{ github.workspace }}
        run: |
          echo "ref=$(git rev-parse HEAD)" >> "$GITHUB_OUTPUT"
          ${{ inputs.execute_command }} --list-only --github-output "$GITHUB_OUTPUT"

  run:
    needs: plan
    if: needs.plan.result == 'success' && needs.plan.outputs.has_tests == 'true'
    # fromJSON parses the input list so `runs-on` receives an actual array.
    # Without this, GitHub treats the string verbatim as one label literal
    # (e.g. the label name '["amd", "mi350", "4gpu"]' with brackets), which no
    # runner carries, and the job queues forever.
    runs-on: ${{ fromJSON(inputs.runs_on) }}
    container:
      image: ${{ inputs.container_image }}
      options: >
        --device=/dev/kfd
        --device=/dev/dri
        --env HIP_VISIBLE_DEVICES
        --security-opt seccomp=unconfined
        --group-add video
        --group-add render
        --cap-add=SYS_PTRACE
        --ipc=host
        --shm-size=32g
        --ulimit memlock=-1
        --ulimit stack=67108864
        --memory=0
        --memory-swap=0
        -v /data/miles_ci:/data/miles_ci
        -v /data/miles_ci/models:/root/models
        -v /data/miles_ci/datasets:/root/datasets
        -v /data/miles_ci/hf_cache:/root/.cache/huggingface
        --privileged
        --ulimit nofile=65535:65535
    defaults:
      run:
        working-directory: ${{ github.workspace }}
    env:
      PYTHONPATH: ${{ github.workspace }}
      MILES_HARDWARE_PLATFORM: rocm
      GITHUB_COMMIT_NAME: ${{ github.event.pull_request.head.sha || github.sha }}_${{ github.event.pull_request.number || 'non-pr' }}
      WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}

    steps:
      - name: Checkout repository
        uses: actions/checkout@v4
        with:
          ref: ${{ needs.plan.outputs.ref }}
          persist-credentials: false

      - name: Cleanup Ray processes
        shell: bash
        run: |
          pkill -9 -f 'ray::' 2>/dev/null || true
          pkill -9 -f raylet 2>/dev/null || true
          pkill -9 -f gcs_server 2>/dev/null || true
          pkill -9 -f 'ray-dashboard' 2>/dev/null || true
          pkill -9 sglang 2>/dev/null || true
          ray stop --force 2>/dev/null || true
          sleep 3

      - name: Wait for GPU ready
        shell: bash
        run: |
          # When a self-hosted runner accepts a new job seconds after the
          # previous job's container exits, the AMD GPU driver subsystem may
          # still be tearing down its prior context. Block here until
          # rocm-smi succeeds, with a bounded retry budget so a
          # genuinely-broken GPU fails fast.
          for i in $(seq 1 60); do
            if rocm-smi --showproductname >/dev/null 2>&1; then
              echo "GPU ready after $((i-1)) retry(s)"
              rocm-smi --showproductname 2>&1 | head -20 || true
              exit 0
            fi
            sleep 2
          done
          echo "ERROR: rocm-smi did not become ready after 120s" >&2
          rocm-smi 2>&1 | tail -20 || true
          exit 1

      - name: Resolve dependency refs
        if: ${{ !inputs.skip_dependency_install }}
        id: resolve-refs
        shell: bash
        env:
          PR_BODY: ${{ github.event.pull_request.body || '' }}
          INPUT_MEGATRON_PR: ${{ github.event.inputs.ci_megatron_pr || '' }}
          INPUT_SGLANG_PR: ${{ github.event.inputs.ci_sglang_pr || '' }}
        run: |
          MEGATRON_PR="${INPUT_MEGATRON_PR}"
          SGLANG_PR="${INPUT_SGLANG_PR}"
          if [ -n "$PR_BODY" ]; then
            PR_MEGATRON_PR=$(echo "$PR_BODY" | grep -m1 -oP '^ci-megatron-pr:\s+\K\S+' || true)
            PR_SGLANG_PR=$(echo "$PR_BODY" | grep -m1 -oP '^ci-sglang-pr:\s+\K\S+' || true)
            [ -z "$MEGATRON_PR" ] && [ -n "$PR_MEGATRON_PR" ] && MEGATRON_PR="$PR_MEGATRON_PR"
            [ -z "$SGLANG_PR" ] && [ -n "$PR_SGLANG_PR" ] && SGLANG_PR="$PR_SGLANG_PR"
          fi
          resolve_fetch_ref() {
            local ref="$1"
            if [[ "$ref" =~ ^#([0-9]+)$ ]]; then
              echo "refs/pull/${BASH_REMATCH[1]}/head"
            else
              echo "$ref"
            fi
          }
          MEGATRON_FETCH=$(resolve_fetch_ref "$MEGATRON_PR")
          SGLANG_FETCH=$(resolve_fetch_ref "$SGLANG_PR")
          echo "ci_megatron_pr=$MEGATRON_FETCH" >> $GITHUB_OUTPUT
          echo "ci_sglang_pr=$SGLANG_FETCH" >> $GITHUB_OUTPUT
          echo "Resolved: megatron=$MEGATRON_PR -> fetch=$MEGATRON_FETCH, sglang=$SGLANG_PR -> fetch=$SGLANG_FETCH"

      - name: Install dependencies
        if: ${{ !inputs.skip_dependency_install }}
        shell: bash
        env:
          MEGATRON_PR: ${{ steps.resolve-refs.outputs.ci_megatron_pr }}
          SGLANG_PR: ${{ steps.resolve-refs.outputs.ci_sglang_pr }}
        run: |
          if command -v rustup >/dev/null 2>&1; then
            rustup default stable
          fi
          if ! command -v protoc >/dev/null 2>&1; then
            apt-get update
            apt-get install -y protobuf-compiler
          fi
          rustc --version
          protoc --version
          if [ -n "$SGLANG_PR" ]; then
            cd /sgl-workspace/sglang && git reset --hard HEAD && git clean -fd && git fetch origin "$SGLANG_PR" && git checkout -f FETCH_HEAD && git log --oneline -1 && pip install -e python --no-deps --break-system-packages
          fi
          if [ -n "$MEGATRON_PR" ]; then
            cd /root/Megatron-LM
            git reset --hard HEAD
            git clean -fd
            git fetch origin "$MEGATRON_PR"
            git checkout -f FETCH_HEAD
            git log --oneline -1
            pip install -e . --no-deps --break-system-packages
          fi
          cd $GITHUB_WORKSPACE && pip install -e . --no-deps --break-system-packages
          pip install polars zstandard xxhash blake3 "numba==0.65.1" --break-system-packages

      - name: Install (miles only)
        if: ${{ inputs.skip_dependency_install }}
        shell: bash
        run: |
          cd $GITHUB_WORKSPACE && pip install -e . --no-deps --break-system-packages
          pip install polars zstandard xxhash blake3 "numba==0.65.1" --break-system-packages

      - name: Install Megatron Bridge override
        if: ${{ github.event.inputs.ci_megatron_bridge_pr != '' || contains(github.event.pull_request.body, 'ci-megatron-bridge-pr:') }}
        uses: ./.github/actions/install-megatron-bridge
        with:
          ref: ${{ github.event.inputs.ci_megatron_bridge_pr || '' }}
          pr_body: ${{ github.event.pull_request.body || '' }}

      - name: Load one HIP runtime
        shell: bash
        run: |
          torch_lib="$(python3 -c 'import os, torch; print(os.path.join(os.path.dirname(torch.__file__), "lib"))')"
          for lib in libamdhip64.so libhsa-runtime64.so; do
            target="$(readlink -f "/opt/rocm/lib/${lib}")"
            [ -f "$torch_lib/$lib" ] && [ -f "$target" ] && ln -sfn "$target" "$torch_lib/$lib"
            ls -l "$torch_lib/$lib"
          done

      - name: Run tests
        shell: bash
        run: ${{ inputs.execute_command }}