atom-test.yaml

name: Atom Test

on:
  push:
    branches: [main]
  pull_request:
    types: [opened, synchronize, reopened, ready_for_review, labeled]
    branches: [main]
    paths-ignore:
      - '**/*.md'
      - 'docs/**'
      - 'LICENSE'
      - '.gitignore'
  workflow_dispatch:

concurrency:
  group: ${{ github.workflow }}-${{ github.ref }}
  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}

env:
  ATOM_BRANCH: "main"
  ATOM_REPOSITORY_URL: "ROCm/ATOM"
  BASE_IMAGE: "rocm/atom-dev:latest"
  DOCKER_PULL_MAX_ATTEMPTS: "3"
  DOCKER_PULL_RETRY_DELAY_SECONDS: "10"
  GITHUB_REPO_URL: ${{ github.event.pull_request.head.repo.clone_url || 'https://github.com/ROCm/Aiter.git' }}
  GITHUB_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.event.head_commit.id || github.sha }}

jobs:
  check-signal:
    if: >-
      github.event_name != 'pull_request' ||
      (github.event.pull_request.draft == false &&
      (contains(github.event.pull_request.labels.*.name, 'ci:atom') ||
      contains(github.event.pull_request.labels.*.name, 'ci:atom_full') ||
      contains(github.event.pull_request.labels.*.name, 'ci:all')))
    runs-on: ubuntu-latest
    steps:
      - name: Checkout code
        uses: actions/checkout@v4

      - name: Wait for Checks workflow
        run: ./.github/scripts/check_signal.sh
        env:
          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
          GITHUB_SHA: ${{ github.sha }}

  load-atom-models:
    if: >-
      github.event_name != 'pull_request' ||
      (github.event.pull_request.draft == false &&
      (contains(github.event.pull_request.labels.*.name, 'ci:atom') ||
      contains(github.event.pull_request.labels.*.name, 'ci:atom_full') ||
      contains(github.event.pull_request.labels.*.name, 'ci:all')))
    runs-on: ubuntu-latest
    outputs:
      models_json: ${{ steps.models.outputs.models_json }}
    steps:
      - name: Checkout ATOM model configs
        uses: actions/checkout@v4
        with:
          repository: ${{ env.ATOM_REPOSITORY_URL }}
          ref: ${{ env.ATOM_BRANCH }}

      - name: Load ATOM accuracy model matrix
        id: models
        env:
          EVENT_NAME: ${{ github.event_name }}
          LABELS_JSON: ${{ toJSON(github.event.pull_request.labels.*.name) }}
        run: |
          python3 <<'PY'
          import json
          import os

          # Always-on PR gate set (every ci:atom / ci:all). Kept small so the
          # do-mi350x pool isn't overloaded (overload was flaking the Kimi +
          # perf gates). The InferenceX frontier models (DeepSeek-V4-Pro,
          # Qwen3.5-397B-A17B-FP8, GLM-5.1-FP8 — all validated PASS on do-mi350x)
          # are covered on-demand via the ci:atom_full label, which runs every
          # pr/main model in ATOM models_accuracy.json with the runner pins below.
          # MiniMax-M2.7 needs HSA_NO_SCRATCH_RECLAIM=1 (scratch reclaim corrupts
          # its MoE scratch buffers on gfx950); injected via model_env_overrides
          # below. It runs under ci:atom_full, not on every PR.
          default_models = [
              {"model_name": "DeepSeek-R1-0528"},
              {"model_name": "gpt-oss-120b"},
              # Kimi-K2.5 ATOM (in-tree) downstream gate. Config (model_path,
              # extraArgs -tp4, threshold) is pulled from ATOM models_accuracy.json.
              {"model_name": "Kimi-K2.5-MXFP4"},
          ]

          labels_raw = os.environ.get("LABELS_JSON") or "[]"
          labels = json.loads(labels_raw)
          if labels is None:
              labels = []

          with open(".github/benchmark/models_accuracy.json", encoding="utf-8") as f:
              atom_models = json.load(f)

          runner_overrides = {
              "atom-mi355-8gpu.predownload": "linux-aiter-do-mi350x-8",
              "linux-atom-do-mi350x-8": "linux-aiter-do-mi350x-8",
              "linux-atom-mi35x-4": "linux-aiter-do-mi350x-4",
              "linux-atom-mi35x-1": "linux-aiter-do-mi350x-1",
          }

          # Per-model runner pin (overrides the generic mapping above). Kimi-K2.5
          # runs the downstream gate on the AITER MI350X pool (TP4 -> 4 GPUs).
          model_runner_overrides = {
              "Kimi-K2.5-MXFP4": "linux-aiter-do-mi350x-4",
              "DeepSeek-V4-Pro": "linux-aiter-do-mi350x-8",       # -tp 8
              "Qwen3.5-397B-A17B-FP8": "linux-aiter-do-mi350x-4",  # -tp 4
              "MiniMax-M2.7-MXFP4": "linux-aiter-do-mi350x-2",     # -tp 2 (MXFP4)
              "GLM-5.1-FP8": "linux-aiter-do-mi350x-8",            # -tp 8
          }

          # Per-model extra env injected into the test container (one KEY=VALUE per
          # line, merged with any env_vars already declared in the ATOM model
          # config). MiniMax MoE needs scratch reclaim disabled or its scratch
          # buffers get corrupted on gfx950, failing the accuracy gate.
          model_env_overrides = {
              "MiniMax-M2.7-MXFP4": "HSA_NO_SCRATCH_RECLAIM=1",
          }

          # Extra CLI args appended (after the ATOM config's extraArgs) for
          # specific models. MiniMax's trust_remote_code modeling breaks ATOM's
          # torch.compile/inductor path during warmup (BackendCompilerFailed ->
          # FileNotFoundError '<frozen posixpath>'); --enforce-eager only
          # disables cudagraph capture in ATOM, so --level 0 (NO_COMPILATION) is
          # what bypasses the failing backend. The MXFP4 ATOM entry has no -tp,
          # so pin TP2 to match the do-mi350x-2 runner.
          model_extra_args_append = {
              "MiniMax-M2.7-MXFP4": "-tp 2 --enforce-eager --level 0",
          }

          def prepare_model(model):
              model = dict(model)
              model["runner"] = runner_overrides.get(model["runner"], model["runner"])
              model["runner"] = model_runner_overrides.get(model["model_name"], model["runner"])
              extra_env = model_env_overrides.get(model["model_name"])
              if extra_env:
                  existing_env = str(model.get("env_vars", "")).strip()
                  model["env_vars"] = (
                      f"{existing_env}\n{extra_env}" if existing_env else extra_env
                  )
              extra_args = model_extra_args_append.get(model["model_name"])
              if extra_args:
                  base_args = str(model.get("extraArgs", "")).strip()
                  model["extraArgs"] = f"{base_args} {extra_args}".strip()
              model["label"] = model.get("runner", "")
              model["accuracy_test_threshold"] = str(model.get("accuracy_threshold", 0))
              model["run_on_pr"] = True
              model["job_name"] = (
                  f"Accuracy ({model['model_name']}, {model['model_path']}, "
                  f"{model.get('extraArgs', '')}, {model['runner']})"
              )
              return model

          # Models force-included in ci:atom_full regardless of their ATOM
          # test_level. The MiniMax MXFP4 variant is tagged "nightly" in the ATOM
          # config, so the pr/main filter below would otherwise skip it.
          full_atom_extra_models = {"MiniMax-M2.7-MXFP4"}

          full_atom = os.environ["EVENT_NAME"] == "pull_request" and "ci:atom_full" in labels
          if not full_atom:
              default_model_names = [model["model_name"] for model in default_models]
              atom_models_by_name = {model["model_name"]: model for model in atom_models}
              missing_models = [
                  name for name in default_model_names if name not in atom_models_by_name
              ]
              if missing_models:
                  raise SystemExit(
                      "Missing default model(s) in ATOM accuracy config: "
                      + ", ".join(missing_models)
                  )
              selected = [
                  prepare_model(atom_models_by_name[name]) for name in default_model_names
              ]
          else:
              selected = [
                  prepare_model(model)
                  for model in atom_models
                  if str(model.get("test_level", "")).lower() in {"pr", "main"}
                  or model["model_name"] in full_atom_extra_models
              ]

          with open(os.environ["GITHUB_OUTPUT"], "a", encoding="utf-8") as out:
              out.write(f"models_json={json.dumps(selected)}\n")

          mode = "full" if full_atom else "default"
          print(f"Selected {len(selected)} ATOM accuracy model(s) for {mode} mode")
          PY

  atom_benchmark:
    if: >-
      github.event_name != 'pull_request' ||
      (github.event.pull_request.draft == false &&
      (contains(github.event.pull_request.labels.*.name, 'ci:atom') ||
      contains(github.event.pull_request.labels.*.name, 'ci:atom_full') ||
      contains(github.event.pull_request.labels.*.name, 'ci:all')))
    needs: [check-signal, load-atom-models]
    name: ${{ matrix.job_name }}
    strategy:
      fail-fast: false
      max-parallel: 1
      matrix:
        include: ${{ fromJson(needs.load-atom-models.outputs.models_json) }}
    runs-on: ${{ matrix.runner }}
    steps:
      - name: Checkout code
        if: matrix.run_on_pr == true || github.event_name != 'pull_request'
        uses: actions/checkout@v4
        with:
          repository: ${{ env.ATOM_REPOSITORY_URL }}
          ref: ${{ env.ATOM_BRANCH }}

      - name: Docker login
        run: docker login -u rocmshared -p ${{ secrets.DOCKER_PASSWORD }} || true

      - name: Download the ATOM base image
        if: matrix.run_on_pr == true || github.event_name != 'pull_request'
        run: |
          for attempt in $(seq 1 "${DOCKER_PULL_MAX_ATTEMPTS}"); do
            if docker pull "${{ env.BASE_IMAGE }}"; then
              echo "Docker pull succeeded on attempt ${attempt}"
              exit 0
            fi
            echo "Docker pull attempt ${attempt} failed"
            if [ "${attempt}" -lt "${DOCKER_PULL_MAX_ATTEMPTS}" ]; then
              sleep "${DOCKER_PULL_RETRY_DELAY_SECONDS}"
            fi
          done
          echo "Docker pull failed after ${DOCKER_PULL_MAX_ATTEMPTS} attempts"
          exit 1

      - name: Generate Dockerfile
        if: matrix.run_on_pr == true || github.event_name != 'pull_request'
        run: |
          cat <<EOF > Dockerfile.mod
          FROM ${{ env.BASE_IMAGE }}

          RUN echo "=== Aiter version BEFORE uninstall ===" && pip show amd-aiter || true
          RUN pip uninstall -y amd-aiter
          RUN pip install --upgrade "pybind11>=3.0.1"
          RUN pip show pybind11
          RUN rm -rf /app/aiter-test
          RUN git clone --no-checkout ${{ env.GITHUB_REPO_URL }} /app/aiter-test && \\
              cd /app/aiter-test && \\
              git checkout --force ${{ env.GITHUB_COMMIT_SHA }} && \\
              git submodule sync && git submodule update --init --recursive && \\
              MAX_JOBS=64 PREBUILD_KERNELS=0 GPU_ARCHS=gfx950 pip install -e . && \\
              ./.github/scripts/install_triton.sh
          RUN echo "=== triton version ===" && pip show triton || true
          RUN echo "=== Aiter version AFTER installation ===" && pip show amd-aiter || true
          EOF

      - name: Build the ATOM test image
        if: matrix.run_on_pr == true || github.event_name != 'pull_request'
        timeout-minutes: 30
        run: |
          docker build --network=host \
            --no-cache \
            -t rocm/aiter-ci:atom-test \
            -f Dockerfile.mod .

      - name: Start CI container
        if: matrix.run_on_pr == true || github.event_name != 'pull_request'
        run: |
          echo "Clean up containers..."
          docker ps -aq -f name=atom_test | xargs -r docker stop | xargs -r docker rm

          if [ -f "/etc/podinfo/gha-render-devices" ]; then
            DEVICE_FLAG=$(cat /etc/podinfo/gha-render-devices)
          else
            DEVICE_FLAG="--device /dev/dri"
          fi

          MODEL_MOUNT=""
          if [ -d "/models" ]; then
            MODEL_MOUNT="-v /models:/models"
            echo "Using host model cache mount: /models"
          else
            echo "Warning: /models directory not found on runner; model cache will stay container-local."
          fi

          cat > /tmp/env_file.txt << 'EOF'
          ${{ matrix.env_vars }}
          EOF

          echo "Starting container: rocm/aiter-ci:atom-test"
          echo "Model-specific environment variables for ${{ matrix.model_name }}:"
          cat /tmp/env_file.txt

          docker run -dt --device=/dev/kfd $DEVICE_FLAG \
          $MODEL_MOUNT \
          --ipc=host --group-add video \
          --shm-size=16G \
          --privileged \
          --cap-add=SYS_PTRACE \
          --env-file /tmp/env_file.txt \
          --security-opt seccomp=unconfined \
          --ulimit memlock=-1 \
          --ulimit stack=67108864 \
          -e ATOM_DISABLE_MMAP=true \
          -e HF_TOKEN="${HF_TOKEN:-${{ secrets.HF_TOKEN_TEST }}}" \
          -v "${{ github.workspace }}:/workspace" \
          -w /workspace \
          --name atom_aiter_test \
          rocm/aiter-ci:atom-test
        env:
          GITHUB_WORKSPACE: ${{ github.workspace }}

      - name: Check shm size
        if: matrix.run_on_pr == true || github.event_name != 'pull_request'
        run: |
          df -h # for testing
          df -h /dev/shm
          docker exec atom_aiter_test df -h /dev/shm
    
      - name: Check version
        if: matrix.run_on_pr == true || github.event_name != 'pull_request'
        run: |
          docker exec atom_aiter_test bash -lc "pip show amd-aiter atom && pip list"

      - name: Download Models
        if: matrix.run_on_pr == true || github.event_name != 'pull_request'
        timeout-minutes: 150
        run: |
          set -euo pipefail
          if docker exec atom_aiter_test bash -lc "[ -d /models ]"; then
            model_cache_dir="/models"
          else
            model_cache_dir="/run"
          fi

          target_dir="${model_cache_dir}/${{ matrix.model_path }}"
          echo "Downloading model to ${target_dir}"
          if ! docker exec \
            -e HF_TOKEN="${HF_TOKEN:-${{ secrets.HF_TOKEN_TEST }}}" \
            -e MODEL_ID="${{ matrix.model_path }}" \
            -e TARGET_DIR="${target_dir}" \
            -e MODEL_DOWNLOAD_TIMEOUT="${MODEL_DOWNLOAD_TIMEOUT}" \
            -e MODEL_LOCK_WAIT_SECONDS="${MODEL_LOCK_WAIT_SECONDS}" \
            -e MODEL_LOCK_POLL_INTERVAL="${MODEL_LOCK_POLL_INTERVAL}" \
            -e MODEL_PROGRESS_INTERVAL="${MODEL_PROGRESS_INTERVAL}" \
            atom_aiter_test bash -lc 'bash /workspace/.github/scripts/download_model_with_lock.sh "$MODEL_ID" "$TARGET_DIR"'; then
            echo "Model download failed for '${{ matrix.model_path }}'. Aborting."
            exit 1
          fi
        env:
          MODEL_DOWNLOAD_TIMEOUT: "3h"
          MODEL_LOCK_WAIT_SECONDS: "7200"
          MODEL_LOCK_POLL_INTERVAL: "30"
          MODEL_PROGRESS_INTERVAL: "60"

      - name: Run ATOM accuracy test
        if: matrix.run_on_pr == true || github.event_name != 'pull_request'
        timeout-minutes: 90
        env:
          MODEL_EXTRA_ARGS: ${{ matrix.extraArgs }}
          CLIENT_COMMAND: ${{ matrix.client_command || '' }}
        run: |
          set -euo pipefail
          echo ""
          echo "========== Launching ATOM server =========="
          if docker exec atom_aiter_test bash -lc "[ -d /models ]"; then
            model_cache_dir="/models"
          else
            model_cache_dir="/run"
          fi

          if docker exec atom_aiter_test bash -lc "[ -d '${model_cache_dir}/${{ matrix.model_path }}' ]"; then
            model_path="${model_cache_dir}/${{ matrix.model_path }}"
          else
            model_path="${{ matrix.model_path }}"
          fi
          printf '.github/scripts/atom_test.sh launch %q %s\n' "$model_path" "$MODEL_EXTRA_ARGS" | \
            docker exec -i atom_aiter_test bash -l
          echo ""
          echo "========== Running accuracy test =========="
          docker exec \
            -e CLIENT_COMMAND="${CLIENT_COMMAND}" \
            atom_aiter_test bash -lc "
            .github/scripts/atom_test.sh accuracy $model_path
          " 2>&1 | tee atom_accuracy_output.txt

      - name: Check accuracy test results
        if: (matrix.run_on_pr == true || github.event_name != 'pull_request') && success()
        run: |
          result_file=$(ls -1t accuracy_test_results/*.json 2>/dev/null | head -n 1)
          if [ -z "$result_file" ] || [ ! -f "$result_file" ]; then
              echo "ERROR: No results JSON file found in accuracy_test_results/"
              exit 2
          else
              echo "RESULT_FILE: $result_file"
          fi
          flexible_extract_value=$(jq '.results.gsm8k["exact_match,flexible-extract"]' "$result_file")
          echo "Flexible extract value: $flexible_extract_value"
          echo "Accuracy test threshold: ${{ matrix.accuracy_test_threshold }}"

          result=$(awk -v val="$flexible_extract_value" -v threshold="${{ matrix.accuracy_test_threshold }}" 'BEGIN {print (val < threshold) ? 1 : 0}')
          if [ "$result" -eq 1 ]; then
            echo "Accuracy test failed: Flexible extract value $flexible_extract_value is less than the threshold ${{ matrix.accuracy_test_threshold }}."
            exit 1
          else
            echo "Accuracy test passed: Flexible extract value $flexible_extract_value is greater than or equal to the threshold ${{ matrix.accuracy_test_threshold }}."
            exit 0
          fi

      - name: Collect Test Summary
        if: (matrix.run_on_pr == true || github.event_name != 'pull_request') && success()
        run: |
          echo "Accuracy Test Summary for ${{ matrix.model_name }}:" >> $GITHUB_STEP_SUMMARY
          awk '/\|Tasks\|Version\|/,/^$/ { if (NF > 0) print }' atom_accuracy_output.txt >> $GITHUB_STEP_SUMMARY

      - name: Upload output
        if: (matrix.run_on_pr == true || github.event_name != 'pull_request') && always()
        uses: actions/upload-artifact@v4
        with:
          name: ${{ matrix.model_name }}_${{ matrix.label }}_${{ matrix.runner }}_atom_accuracy_output.txt
          path: atom_accuracy_output.txt
          retention-days: 7

      - name: Clean Up
        if: (matrix.run_on_pr == true || github.event_name != 'pull_request') && always()
        run: |
          docker stop atom_aiter_test || true
          docker rm atom_aiter_test || true