Skip to content

ci: bump base image/runner to v26.4 and drop v26.3-only workarounds #4178

ci: bump base image/runner to v26.4 and drop v26.3-only workarounds

ci: bump base image/runner to v26.4 and drop v26.3-only workarounds #4178

Workflow file for this run

name: Primus-CI-TAS
on:
workflow_dispatch:
push:
branches:
- main
tags:
- "v*"
pull_request:
# Re-list the implicit defaults (opened/synchronize/reopened) and add
# ready_for_review so flipping a Draft PR to Ready triggers its first run.
types: [opened, synchronize, reopened, ready_for_review]
concurrency:
group: ${{ github.workflow }}-${{ github.head_ref || github.event.merge_group.head_ref || github.ref }}
cancel-in-progress: true
# Default to read-only. No job uses GITHUB_TOKEN for writes (Docker Hub push
# uses a dedicated PAT secret), so this is behavior-neutral. Jobs that need
# more must opt in at the job level.
permissions:
contents: read
env:
PRIMUS_TURBO_COMMIT: 56c789e58f72aaf733b7715f1536be1ed33b69a1 # dev/kyle/flydsl_attn_deepseekv4 (dsv4 sparse-MLA attention + cr=4 fixes + flydsl kernels packaging __init__.py fix)
PRIMUS_TURBO_AITER_COMMIT: 0f3c58e6edb6754940bcf9fd5f09ccb6f389f52e # AITER v0.1.14.post1 (tag commit) — required by Primus-Turbo main aiter_utils.py
ROCSHMEM_COMMIT: 17ff985c026f9f97f85068647e863ab541dd5645 # Update version to 3.2.0 for 7.2.0 rocm release (#351) (#355)
UCCL_COMMIT: 5afb4117893c58cc0c8557d9286336141a301053 # [EP]: fix fp8 error of internode_ll on amd gfx950 arch. (#710)
TRITON_COMMIT: 88b227e23f0445f3f695bad05bbf1a363b4f50e0
BASE_IMAGE: docker.io/rocm/primus:v26.4
MAXTEXT_BASE_IMAGE: docker.io/rocm/jax-training:maxtext-v26.2
jobs:
code-lint:
runs-on: ubuntu-latest
strategy:
matrix:
python-version: ["3.12"]
steps:
- run: echo "🎉 Begin Primus Python Lint."
- run: echo "🐧 This job is now running on a ${{ runner.os }} server hosted by GitHub!"
- uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
- run: git config --global --add safe.directory /github/workspace
- name: Set up Python
uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0
with:
python-version: ${{ matrix.python-version }}
- name: Display Python version
run: python -c "import sys; print(sys.version)"
# Mirrors `.pre-commit-config.yaml` exactly (hooks + pinned revs/args).
- name: Install pre-commit
run: pip install pre-commit
- name: Run pre-commit (all hooks, all files)
run: pre-commit run --all-files --show-diff-on-failure
- name: Check version/commit pin consistency
run: python tools/ci/check_version_consistency.py
# Flag PRs that introduce known-vulnerable dependencies. Warn-only during
# burn-in; drop warn-only to turn it into a hard gate once the team is ready.
dependency-review:
if: github.event_name == 'pull_request'
runs-on: ubuntu-latest
steps:
- name: Checkout repository
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
- name: Dependency Review
uses: actions/dependency-review-action@a1d282b36b6f3519aa1f3fc636f609c47dddb294 # v5.0.0
with:
fail-on-severity: high
warn-only: true
build-docker:
needs: [code-lint]
# Skip on Draft PRs (keep code-lint ungated); the event_name guard preserves
# push/tag and workflow_dispatch runs where github.event.pull_request is absent.
if: ${{ github.event_name != 'pull_request' || !github.event.pull_request.draft }}
runs-on: build-docker
strategy:
matrix:
python-version: ["3.12"]
steps:
- run: echo "🎉 Begin Build Primus Docker Image."
- uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
submodules: recursive
- name: Show Environment Info
run: |
echo "Hostname: $(hostname)"
echo "PWD: $(pwd)"
echo "HOME: $HOME"
echo "GITHUB_WORKSPACE: $GITHUB_WORKSPACE"
echo "Runner Temp Dir: $RUNNER_TEMP"
echo "Runner Tool Cache: $RUNNER_TOOL_CACHE"
- name: Parse Commit Info
run: |
if [[ "${{ github.event_name }}" == "pull_request" ]]; then
echo "IMAGE_TAG=pr-${{ github.event.pull_request.number }}" >> $GITHUB_ENV
elif [[ "${{ github.event_name }}" == "push" && "${{ github.ref }}" == "refs/heads/main" ]]; then
echo "IMAGE_TAG=latest" >> $GITHUB_ENV
elif [[ "${{ github.event_name }}" == "release" ]]; then
TAG_NAME="${{ github.ref }}"
TAG="${TAG_NAME#refs/tags/}"
echo "IMAGE_TAG=$TAG" >> $GITHUB_ENV
else
echo "IMAGE_TAG=others" >> $GITHUB_ENV
fi
- name: Build and Push Docker Image
run: |
# echo "> Login to ROCm Docker Hub"
# docker login -u rocmshared -p ${{ secrets.ROCM_DOCKER_HUB_TOKEN }}
echo "> Build Docker Image with tag: ${{ env.IMAGE_TAG }}"
echo "> Build dependencies"
start_time=$(date +%s)
docker build -f $GITHUB_WORKSPACE/.github/workflows/docker/Dockerfile \
--network=host \
-t tasimage/primus:${{env.IMAGE_TAG}} \
--build-arg BASE_IMAGE=${BASE_IMAGE} \
--build-arg PRIMUS_TURBO_COMMIT=${PRIMUS_TURBO_COMMIT} \
--build-arg PRIMUS_TURBO_AITER_COMMIT=${PRIMUS_TURBO_AITER_COMMIT} \
--build-arg ROCSHMEM_COMMIT=${ROCSHMEM_COMMIT} \
--build-arg PRIMUS_TURBO_FRAMEWORK=PYTORCH \
--build-arg UCCL_COMMIT=${UCCL_COMMIT} \
--build-arg TRITON_COMMIT=${TRITON_COMMIT} \
$GITHUB_WORKSPACE/.github/workflows/docker
end_time=$(date +%s)
elapsed=$((end_time - start_time))
echo "⏱️ [build primus docker] Total elapsed time: ${elapsed} seconds"
docker tag tasimage/primus:${{env.IMAGE_TAG}} docker.io/tasimage/primus:${{env.IMAGE_TAG}}
docker login -u tasimage -p ${{ secrets.PRIMUS_DOCKER_HUB_TOKEN }}
docker push docker.io/tasimage/primus:${{env.IMAGE_TAG}}
# docker login -u rocmshared -p ${{ secrets.ROCM_DOCKER_HUB_TOKEN }}
# Primus v26.4 already includes AINIC under /workspace. Re-enable this
# Dockerfile.ainic build only when we need to refresh the tasimage -ainic image.
echo "> Build Docker Image with tag: ${{ env.IMAGE_TAG }}-ainic"
start_time=$(date +%s)
mkdir -p $GITHUB_WORKSPACE/.github/workflows/docker/ainic
cp /apps/tas/0_public/primus_docker_ci/ainic/ainic_bundle_1.117.5-a-77.tar.gz $GITHUB_WORKSPACE/.github/workflows/docker/ainic/ || { echo "Error: Failed to copy ainic bundle"; exit 1; }
docker build -f $GITHUB_WORKSPACE/.github/workflows/docker/Dockerfile.ainic \
--network=host \
-t tasimage/primus:${{env.IMAGE_TAG}}-ainic \
--build-arg BASE_IMAGE=docker.io/tasimage/primus:${{env.IMAGE_TAG}} \
--build-arg AINIC_BUNDLE_PATH=ainic \
$GITHUB_WORKSPACE/.github/workflows/docker
end_time=$(date +%s)
elapsed=$((end_time - start_time))
echo "⏱️ [build primus docker-ainic] Total elapsed time: ${elapsed} seconds"
docker tag tasimage/primus:${{env.IMAGE_TAG}}-ainic docker.io/tasimage/primus:${{env.IMAGE_TAG}}-ainic
docker login -u tasimage -p ${{ secrets.PRIMUS_DOCKER_HUB_TOKEN }}
docker push docker.io/tasimage/primus:${{env.IMAGE_TAG}}-ainic
# docker login -u rocmshared -p ${{ secrets.ROCM_DOCKER_HUB_TOKEN }}
echo "> Build Docker Image with tag: ${{ env.IMAGE_TAG }}-jax"
start_time=$(date +%s)
docker build -f $GITHUB_WORKSPACE/.github/workflows/docker/Dockerfile \
--network=host \
-t tasimage/primus:${{env.IMAGE_TAG}}-jax \
--build-arg BASE_IMAGE=${MAXTEXT_BASE_IMAGE} \
--build-arg PRIMUS_TURBO_COMMIT=${PRIMUS_TURBO_COMMIT} \
--build-arg PRIMUS_TURBO_AITER_COMMIT=${PRIMUS_TURBO_AITER_COMMIT} \
--build-arg PRIMUS_TURBO_FRAMEWORK=JAX \
--build-arg UCCL_COMMIT=${UCCL_COMMIT} \
--build-arg ROCSHMEM_COMMIT=${ROCSHMEM_COMMIT} \
--build-arg TRITON_COMMIT=${TRITON_COMMIT} .
end_time=$(date +%s)
elapsed=$((end_time - start_time))
echo "⏱️ [build primus docker-jax] Total elapsed time: ${elapsed} seconds"
echo "> Docker tag image for Docker Hub"
docker tag tasimage/primus:${{env.IMAGE_TAG}}-jax docker.io/tasimage/primus:${{env.IMAGE_TAG}}-jax
docker login -u tasimage -p ${{ secrets.PRIMUS_DOCKER_HUB_TOKEN }}
docker push docker.io/tasimage/primus:${{env.IMAGE_TAG}}-jax
# docker login -u rocmshared -p ${{ secrets.ROCM_DOCKER_HUB_TOKEN }}
echo "> Build Docker Image with tag: ${{ env.IMAGE_TAG }}-jax-ainic"
start_time=$(date +%s)
mkdir -p $GITHUB_WORKSPACE/.github/workflows/docker/ainic
cp /apps/tas/0_public/primus_docker_ci/ainic/ainic_bundle_1.117.5-a-56.tar.gz $GITHUB_WORKSPACE/.github/workflows/docker/ainic/ || { echo "Error: Failed to copy ainic bundle"; exit 1; }
docker build -f $GITHUB_WORKSPACE/.github/workflows/docker/Dockerfile_jax.ainic \
--network=host \
-t tasimage/primus:${{env.IMAGE_TAG}}-jax-ainic \
--build-arg BASE_IMAGE=${MAXTEXT_BASE_IMAGE} \
--build-arg AINIC_BUNDLE_PATH=ainic \
$GITHUB_WORKSPACE/.github/workflows/docker
end_time=$(date +%s)
elapsed=$((end_time - start_time))
echo "⏱️ [build primus docker-jax-ainic] Total elapsed time: ${elapsed} seconds"
docker tag tasimage/primus:${{env.IMAGE_TAG}}-jax-ainic docker.io/tasimage/primus:${{env.IMAGE_TAG}}-jax-ainic
docker login -u tasimage -p ${{ secrets.PRIMUS_DOCKER_HUB_TOKEN }}
docker push docker.io/tasimage/primus:${{env.IMAGE_TAG}}-jax-ainic
# docker login -u rocmshared -p ${{ secrets.ROCM_DOCKER_HUB_TOKEN }}
# echo "> Docker cleanup local images"
# docker rmi tasimage/primus:${{env.IMAGE_TAG}}
# docker rmi tasimage/primus:${{env.IMAGE_TAG}}-jax
echo "> build-docker success"
run-unittest-torch:
permissions:
contents: read
actions: read # let "Write runtime summary" list this job's own steps for auto-discovered timing
env:
PRIMUS_WORKDIR: /mnt/apps_proxy/tas/0_public/primus_ci/actions-runner-torch
# PRIMUS_WORKDIR: /wekafs/primus-data/primus_safe_ci/torch
PRIMUS_TURBO_ATTN_V3_ATOMIC_FP32: 1
needs: [code-lint]
if: ${{ github.event_name != 'pull_request' || !github.event.pull_request.draft }}
# runs-on: [primus-lm-cicd-torch-j8knc]
runs-on: [primus-lm-cicd-v26.4-tas8n-a16-40]
steps:
- run: echo "🎉 Begin Primus-Turbo Checkout."
- name: Clean stale Primus-Turbo checkout
run: rm -rf "${GITHUB_WORKSPACE}/Primus-Turbo"
- name: Set commit hash to env
run: echo "PRIMUS_TURBO_COMMIT=${PRIMUS_TURBO_COMMIT}" >> $GITHUB_ENV
- name: Checkout Repo Primus-Turbo
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
repository: AMD-AGI/Primus-Turbo
submodules: "recursive"
path: Primus-Turbo
ref: ${{ env.PRIMUS_TURBO_COMMIT }}
- run: echo "Begin AITER + Primus-Turbo Install."
- name: Install AITER
run: |
: > "$RUNNER_TEMP/runtime.tsv" # reset the CI runtime log for this job
echo "✅ [Uninstall old aiter] started at: $(date)"
pip3 uninstall aiter amd-aiter -y || true
rm -rf /tmp/aiter || true
cd /tmp
git clone https://github.com/ROCm/aiter.git
cd aiter
git checkout -f ${PRIMUS_TURBO_AITER_COMMIT}
git submodule sync
git submodule update --init --recursive
start_time=$(date +%s)
echo "✅ [Build aiter] started at: $(date)"
PREBUILD_KERNELS=3 pip install --no-cache-dir --use-pep517 .
end_time=$(date +%s)
elapsed=$((end_time - start_time))
echo "✅ [Build aiter] ended at: $(date)"
echo "⏱️ [Build aiter] Total elapsed time: ${elapsed} seconds"
echo -e "Build aiter\t${elapsed}" >> "$RUNNER_TEMP/runtime.tsv"
- name: Install Primus-Turbo
run: |
rm -rf /tmp/Primus-Turbo || true
mv Primus-Turbo /tmp/
echo "Primus-Turbo dir: /tmp/Primus-Turbo"
git config --global --add safe.directory /tmp/Primus-Turbo || true
cd /tmp/Primus-Turbo || true
start_time=$(date +%s)
echo "✅ [Pip install requirements] started at: $(date)"
mkdir -p ${PRIMUS_WORKDIR}/primus-cache
MAX_JOBS=128 pip install --cache-dir=${PRIMUS_WORKDIR}/primus-cache --no-build-isolation --no-clean -r requirements.txt
end_time=$(date +%s)
elapsed=$((end_time - start_time))
echo "✅ [Pip install requirements] ended at: $(date)"
echo "⏱️ [Pip install requirements] Total elapsed time: ${elapsed} seconds"
echo -e "primus-turbo: pip install requirements\t${elapsed}" >> "$RUNNER_TEMP/runtime.tsv"
start_time=$(date +%s)
echo "✅ [build primus-turbo] started at: $(date)"
pip3 install --no-build-isolation -e . -v
end_time=$(date +%s)
elapsed=$((end_time - start_time))
echo "✅ [build primus-turbo] ended at: $(date)"
echo "⏱️ [build primus-turbo] Total elapsed time: ${elapsed} seconds"
echo -e "primus-turbo: build/install\t${elapsed}" >> "$RUNNER_TEMP/runtime.tsv"
- run: echo "🎉 Begin Primus Unit Test."
- uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
submodules: recursive
- name: Show Environment Info
run: |
echo "Hostname: $(hostname)"
echo "PWD: $(pwd)"
echo "HOME: $HOME"
echo "GITHUB_WORKSPACE: $GITHUB_WORKSPACE"
echo "Runner Temp Dir: $RUNNER_TEMP"
echo "Runner Tool Cache: $RUNNER_TOOL_CACHE"
- name: Install Primus
run: |
pip install -r requirements.txt
# origami is intentionally NOT installed here, matching
# .github/workflows/docker/Dockerfile: no measurable perf gain on
# gfx942, and turbo falls back to heuristic kernel selection without it.
- name: Set UT_LOG_PATH
run: |
ts="$(date +%Y%m%d-%H%M%S)"
commit_id="${GITHUB_SHA::7}"
if [[ "${{ github.event_name }}" == "pull_request" ]]; then
echo "UT_LOG_PATH=${PRIMUS_WORKDIR}/ut_out/pr-${{ github.event.pull_request.number }}-${ts}-${commit_id}" >> $GITHUB_ENV
elif [[ "${{ github.event_name }}" == "push" && "${{ github.ref }}" == "refs/heads/main" ]]; then
echo "UT_LOG_PATH=${PRIMUS_WORKDIR}/ut_out/main-${ts}-${commit_id}" >> $GITHUB_ENV
elif [[ "${{ github.event_name }}" == "release" ]]; then
TAG_NAME="${{ github.ref }}"
TAG="${TAG_NAME#refs/tags/}"
echo "UT_LOG_PATH=${PRIMUS_WORKDIR}/ut_out/${TAG}-${ts}-${commit_id}" >> $GITHUB_ENV
else
echo "UT_LOG_PATH=${PRIMUS_WORKDIR}/ut_out/others-${ts}-${commit_id}" >> $GITHUB_ENV
fi
- name: Run CLI Shell Tests
run: |
echo "Running Primus CLI shell tests..."
bash ./tests/runner/run_all_tests.sh
- name: Run Primus Core Tests
run: |
echo "Running Primus Core tests..."
# Note: The tests `test_fp8_te_linear` and `test_te_linear` are temporarily skipped due to intermittent failures.
# Note: `test_limit_layers_moe_with_dense_keeps_two` is temporarily disabled pending a fix to
# `_limit_layers_for_projection` dense-layer handling (num_layers collapses to 1 instead of 2).
# Note HSA_NO_SCRATCH_RECLAIM=1 must be set to avoid RCCL perf hit (TAS-8N Node), rocm ver:70125424
export HSA_NO_SCRATCH_RECLAIM=1
mkdir -p "${GITHUB_WORKSPACE}/test-reports"
pytest --maxfail=1 -s ./tests/unit_tests/ \
--cov=primus --cov-report=term-missing:skip-covered \
--junitxml="${GITHUB_WORKSPACE}/test-reports/core-unit.xml" \
--deselect=tests/unit_tests/megatron/cco/test_tp_overlap.py::TPOverlapTestCase::test_fp8_te_linear \
--deselect=tests/unit_tests/megatron/cco/test_tp_overlap.py::TPOverlapTestCase::test_te_linear \
--deselect=tests/unit_tests/megatron/transformer/moe/test_token_dispatcher.py::TestFlexDispatcher::test_forward_backward \
--deselect=tests/unit_tests/megatron/transformer/moe/test_token_dispatcher.py::TestFlexDispatcher::test_capacity_forward_backward \
--deselect=tests/unit_tests/core/projection/test_performance_projection.py::test_limit_layers_moe_with_dense_keeps_two
- name: Snapshot unit-test coverage
if: always()
continue-on-error: true
run: |
# Keep the unit data for the final unit-vs-E2E comparison table
# (.coverage.unit is fed to `coverage combine` after the E2E steps).
python -m coverage json -o coverage_unit.json 2>/dev/null || true
cp .coverage "$GITHUB_WORKSPACE/.coverage.unit" 2>/dev/null || true
- name: Setup E2E training coverage
if: always()
continue-on-error: true
run: |
# Record coverage of the training subprocesses (primus-cli -> torchrun
# -> python). A .pth makes every new interpreter call
# coverage.process_startup(); COVERAGE_PROCESS_START points it at the rc
# and COVERAGE_FILE keeps E2E data separate from the unit .coverage.
# parallel=true -> one data file per rank; sigterm=true captures ranks
# killed on teardown. These env vars are inherited by the E2E steps.
SP=$(python -c "import site; print(site.getsitepackages()[0])")
echo "import coverage; coverage.process_startup()" > "$SP/primus_e2e_coverage.pth"
printf '[run]\nparallel = true\nsource = primus\nsigterm = true\n' > "$GITHUB_WORKSPACE/.coveragerc_e2e"
echo "COVERAGE_PROCESS_START=$GITHUB_WORKSPACE/.coveragerc_e2e" >> "$GITHUB_ENV"
echo "COVERAGE_FILE=$GITHUB_WORKSPACE/.coverage_e2e" >> "$GITHUB_ENV"
- name: Decide E2E scope (torch)
run: |
# Fail-safe: default to running all E2E (push/release/dispatch, or any
# diff failure). On PRs, narrow to the suites the changed files affect.
M=1; T=1
if [[ "${{ github.event_name }}" == "pull_request" ]]; then
base="${{ github.event.pull_request.base.sha }}"
git fetch --no-tags --depth=200 origin "${{ github.base_ref }}" 2>/dev/null || true
changed="$(git diff --name-only "${base}" HEAD 2>/dev/null || true)"
if [[ -n "${changed}" ]]; then
e2e="$(printf '%s\n' "${changed}" | python tools/ci/select_tests.py)"
echo "Selected torch E2E scope: ${e2e:-<none>}"
if [[ "${e2e}" != "all" ]]; then
echo "${e2e}" | grep -qw megatron || M=0
echo "${e2e}" | grep -qw torchtitan || T=0
fi
fi
fi
echo "RUN_MEGATRON_E2E=${M}" >> "$GITHUB_ENV"
echo "RUN_TORCHTITAN_E2E=${T}" >> "$GITHUB_ENV"
- name: Run Primus Model Tests -- Megatron-LM
env:
HF_TOKEN: ${{secrets.HF_TOKEN}}
run: |
if [[ "${RUN_MEGATRON_E2E:-1}" != "1" ]]; then echo "Skipping Megatron-LM E2E: no relevant changes."; exit 0; fi
echo "Set UT_LOG_PATH: ${{ env.UT_LOG_PATH }}"
rm -rf "${{ env.UT_LOG_PATH }}"
mkdir -p "${{ env.UT_LOG_PATH }}"
# MASTER_PORT=10009 DATA_PATH=/wekafs/primus-data \
MASTER_PORT=10009 DATA_PATH=/mnt/apps_proxy/tas/0_public/data HSA_NO_SCRATCH_RECLAIM=1 \
GPU_ARCHS=gfx942 \
pytest --maxfail=1 -s ./tests/trainer/test_megatron_trainer.py \
--junitxml="${GITHUB_WORKSPACE}/test-reports/megatron-e2e.xml"
- name: Run Primus Model Tests -- TorchTitan
env:
HF_TOKEN: ${{secrets.HF_TOKEN}}
run: |
if [[ "${RUN_TORCHTITAN_E2E:-1}" != "1" ]]; then echo "Skipping TorchTitan E2E: no relevant changes."; exit 0; fi
echo "Set UT_LOG_PATH: ${{ env.UT_LOG_PATH }}"
rm -rf "${{ env.UT_LOG_PATH }}"
mkdir -p "${{ env.UT_LOG_PATH }}"
# MASTER_PORT=10009 DATA_PATH=/wekafs/primus-data \
MASTER_PORT=10009 DATA_PATH=/mnt/apps_proxy/tas/0_public/data HSA_NO_SCRATCH_RECLAIM=1 \
pytest --maxfail=1 -s ./tests/trainer/test_torchtitan_trainer.py \
--junitxml="${GITHUB_WORKSPACE}/test-reports/torchtitan-e2e.xml"
- name: Run Primus CLI tool smoke (benchmark / preflight, E2E coverage)
if: always()
continue-on-error: true
run: |
# Best-effort CLI smoke so primus/tools (GEMM / attention / RCCL
# benches + preflight probes) is exercised under the inherited E2E
# coverage injection. Never fails CI: coverage is kept up to any exit.
export HSA_NO_SCRATCH_RECLAIM=1
OUT="${{ env.UT_LOG_PATH }}/cli_smoke"; mkdir -p "$OUT"
cli() { python -m primus.cli.main "$@" || true; }
# single-GPU benches + preflight info
cli benchmark gemm --M 512 --N 512 --K 512 --duration 2 --output-file "$OUT/gemm.md"
cli benchmark gemm-dense --seqlen 256 --hidden-size 512 --intermediate-size 512 --duration 2 --output-file "$OUT/gemm_dense.md"
cli benchmark gemm-deepseek --seqlen 256 --hidden-size 512 --intermediate-size 512 --moe-intermediate-size 512 --duration 2 --output-file "$OUT/gemm_deepseek.md"
cli benchmark attention --backend flash --mbs-list 1,2 --report-csv-path "$OUT/attn.csv"
cli preflight --host --gpu --dump-path "$OUT/preflight"
# multi-GPU collectives (best-effort; strided needs >=8 GPUs)
NGPU=$(python -c "import torch; print(torch.cuda.device_count())" 2>/dev/null || echo 1)
[ "$NGPU" -ge 2 ] && torchrun --nproc_per_node="$NGPU" -m primus.cli.main benchmark rccl --op all_reduce all_gather --min-bytes 1K --max-bytes 256K --num-sizes 2 --iters 3 --warmup 1 --output-file "$OUT/rccl.md" || true
[ "$NGPU" -ge 8 ] && torchrun --nproc_per_node=8 -m primus.cli.main benchmark strided-allgather --sizes-mb 8 --iters 3 --warmup 1 || true
- name: Upload test reports
if: always()
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: torch-test-reports
path: test-reports/
if-no-files-found: warn
retention-days: 14
- name: Write test summary
if: always()
run: |
python tools/ci/junit_summary.py --title torch test-reports/*.xml >> "$GITHUB_STEP_SUMMARY" || true
- name: Write runtime summary
if: always()
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
run: |
python tools/ci/runtime_summary.py --title torch >> "$GITHUB_STEP_SUMMARY" || true
- name: Build torch coverage json (unit + E2E)
if: always()
continue-on-error: true
run: |
# Don't instrument coverage's own helper processes here.
unset COVERAGE_PROCESS_START
# Merge per-rank E2E data, then line-merge unit + E2E into one dataset.
# The JSON is consumed by the coverage-summary job (no per-job summary).
python -m coverage combine 2>/dev/null || true
COVERAGE_FILE="$GITHUB_WORKSPACE/.coverage_all" python -m coverage combine --keep \
"$GITHUB_WORKSPACE/.coverage.unit" "$GITHUB_WORKSPACE/.coverage_e2e" 2>/dev/null || true
COVERAGE_FILE="$GITHUB_WORKSPACE/.coverage_all" python -m coverage json -o coverage_combined.json 2>/dev/null || true
- name: Upload torch coverage json
if: always()
continue-on-error: true
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: coverage-torch
path: |
coverage_unit.json
coverage_combined.json
if-no-files-found: ignore
- name: Clean
if: always()
run: |
rm -rf ${PRIMUS_WORKDIR}/Primus-Turbo
rm -rf ${PRIMUS_WORKDIR}/Primus
# Remove E2E coverage artifacts and, importantly, the subprocess
# injection .pth so it never leaks into later jobs on this persistent runner.
rm -f "${GITHUB_WORKSPACE}/.coverage_e2e"* "${GITHUB_WORKSPACE}/.coverage.unit" "${GITHUB_WORKSPACE}/.coverage_all" "${GITHUB_WORKSPACE}/.coveragerc_e2e" coverage_unit.json coverage_combined.json || true
SP=$(python -c "import site; print(site.getsitepackages()[0])" 2>/dev/null) && rm -f "$SP/primus_e2e_coverage.pth" || true
run-unittest-jax:
permissions:
contents: read
actions: read # let "Write runtime summary" list this job's own steps for auto-discovered timing
env:
# PRIMUS_WORKDIR: /wekafs/primus-data/primus_safe_ci/jax
PRIMUS_WORKDIR: /mnt/apps_proxy/tas/0_public/primus_docker_jax_ci/actions-runner
needs: [code-lint]
if: ${{ github.event_name != 'pull_request' || !github.event.pull_request.draft }}
runs-on: [primus-jax-tas-runner] # docker container primus_jax_github_runner on tas a16-31
steps:
- run: echo "🎉 Begin Primus-Turbo Checkout."
- name: Clean stale Primus-Turbo checkout
run: rm -rf "${GITHUB_WORKSPACE}/Primus-Turbo"
- name: Set commit hash to env
run: echo "PRIMUS_TURBO_COMMIT=${PRIMUS_TURBO_COMMIT}" >> $GITHUB_ENV
- name: Checkout Repo Primus-Turbo
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
repository: AMD-AGI/Primus-Turbo
path: Primus-Turbo
ref: ${{ env.PRIMUS_TURBO_COMMIT }}
- name: Init Primus-Turbo submodules (full clone)
working-directory: Primus-Turbo
run: |
rm -rf 3rdparty/composable_kernel .git/modules/3rdparty/composable_kernel
git submodule sync --recursive
git submodule update --init --recursive
- run: echo "Begin Primus-Turbo Install."
- name: Install Primus-Turbo
run: |
rm -rf /tmp/Primus-Turbo || true
mv Primus-Turbo /tmp/
echo "Primus-Turbo dir: /tmp/Primus-Turbo"
git config --global --add safe.directory /tmp/Primus-Turbo
cd /tmp/Primus-Turbo
: > "$RUNNER_TEMP/runtime.tsv" # reset the CI runtime log for this job
start_time=$(date +%s)
echo "✅ [Pip install requirements] started at: $(date)"
mkdir -p ${PRIMUS_WORKDIR}/primus-cache
python3 -m pip install --upgrade pip setuptools
end_time=$(date +%s)
elapsed=$((end_time - start_time))
echo "✅ [Pip install requirements] ended at: $(date)"
echo "⏱️ [Pip install requirements] Total elapsed time: ${elapsed} seconds"
echo -e "pip install/upgrade\t${elapsed}" >> "$RUNNER_TEMP/runtime.tsv"
start_time=$(date +%s)
echo "✅ [build primus-turbo] started at: $(date)"
end_time=$(date +%s)
elapsed=$((end_time - start_time))
echo "✅ [build primus-turbo] ended at: $(date)"
echo "⏱️ [build primus-turbo] Torch installation causes segfault, so we skip it and actually not install turbo. Total elapsed time: ${elapsed} seconds"
- run: echo "🎉 Begin Primus Unit Test."
- uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
submodules: recursive
- name: Show Environment Info
run: |
echo "Hostname: $(hostname)"
echo "PWD: $(pwd)"
echo "HOME: $HOME"
echo "GITHUB_WORKSPACE: $GITHUB_WORKSPACE"
echo "Runner Temp Dir: $RUNNER_TEMP"
echo "Runner Tool Cache: $RUNNER_TOOL_CACHE"
- name: Install Primus
run: |
pip install -r requirements-jax.txt
- name: Set UT_LOG_PATH
run: |
ts="$(date +%Y%m%d-%H%M%S)"
commit_id="${GITHUB_SHA::7}"
if [[ "${{ github.event_name }}" == "pull_request" ]]; then
echo "UT_LOG_PATH=${PRIMUS_WORKDIR}/ut_out/pr-${{ github.event.pull_request.number }}-${ts}-${commit_id}" >> $GITHUB_ENV
elif [[ "${{ github.event_name }}" == "push" && "${{ github.ref }}" == "refs/heads/main" ]]; then
echo "UT_LOG_PATH=${PRIMUS_WORKDIR}/ut_out/main-${ts}-${commit_id}" >> $GITHUB_ENV
elif [[ "${{ github.event_name }}" == "release" ]]; then
TAG_NAME="${{ github.ref }}"
TAG="${TAG_NAME#refs/tags/}"
echo "UT_LOG_PATH=${PRIMUS_WORKDIR}/ut_out/${TAG}-${ts}-${commit_id}" >> $GITHUB_ENV
else
echo "UT_LOG_PATH=${PRIMUS_WORKDIR}/ut_out/others-${ts}-${commit_id}" >> $GITHUB_ENV
fi
- name: Run Shell Tests
run: |
echo "Running Primus CLI shell tests..."
bash ./tests/runner/run_all_tests.sh
- name: Setup MaxText E2E coverage
if: always()
continue-on-error: true
run: |
# Instrument the MaxText (JAX) training subprocesses so their coverage
# of primus/ (notably backends/maxtext) is recorded. Same mechanism as
# the torch job: a .pth runs coverage.process_startup() in every child.
pip install coverage >/dev/null 2>&1 || true
SP=$(python -c "import site; print(site.getsitepackages()[0])")
echo "import coverage; coverage.process_startup()" > "$SP/primus_e2e_coverage.pth"
printf '[run]\nparallel = true\nsource = primus\nsigterm = true\n' > "$GITHUB_WORKSPACE/.coveragerc_e2e"
echo "COVERAGE_PROCESS_START=$GITHUB_WORKSPACE/.coveragerc_e2e" >> "$GITHUB_ENV"
echo "COVERAGE_FILE=$GITHUB_WORKSPACE/.coverage_e2e" >> "$GITHUB_ENV"
- name: Decide E2E scope (jax)
run: |
# Fail-safe: default to running MaxText E2E; on PRs skip it when no
# relevant (maxtext / shared) changes are detected.
X=1
if [[ "${{ github.event_name }}" == "pull_request" ]]; then
base="${{ github.event.pull_request.base.sha }}"
git fetch --no-tags --depth=200 origin "${{ github.base_ref }}" 2>/dev/null || true
changed="$(git diff --name-only "${base}" HEAD 2>/dev/null || true)"
if [[ -n "${changed}" ]]; then
e2e="$(printf '%s\n' "${changed}" | python tools/ci/select_tests.py)"
echo "Selected jax E2E scope: ${e2e:-<none>}"
if [[ "${e2e}" != "all" ]]; then
echo "${e2e}" | grep -qw maxtext || X=0
fi
fi
fi
echo "RUN_MAXTEXT_E2E=${X}" >> "$GITHUB_ENV"
- name: Run Unit Tests
env:
HF_TOKEN: ${{secrets.HF_TOKEN}}
run: |
if [[ "${RUN_MAXTEXT_E2E:-1}" != "1" ]]; then echo "Skipping MaxText E2E: no relevant changes."; exit 0; fi
echo "Set UT_LOG_PATH: ${{ env.UT_LOG_PATH }}"
rm -rf "${{ env.UT_LOG_PATH }}"
mkdir -p "${{ env.UT_LOG_PATH }}"
mkdir -p "${GITHUB_WORKSPACE}/test-reports"
# run_unit_tests.py shells out to pytest; PYTEST_ADDOPTS injects the
# JUnit report without having to touch the wrapper script.
export PYTEST_ADDOPTS="--junitxml=${GITHUB_WORKSPACE}/test-reports/maxtext-e2e.xml"
# MASTER_PORT=10009 DATA_PATH=/wekafs/primus-data \
MASTER_PORT=10009 DATA_PATH=/mnt/apps_proxy/tas/0_public/data \
JAX_SKIP_UT=1 python ./tests/run_unit_tests.py --jax
- name: Upload test reports
if: always()
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: jax-test-reports
path: test-reports/
if-no-files-found: warn
retention-days: 14
- name: Write test summary
if: always()
run: |
python tools/ci/junit_summary.py --title jax test-reports/*.xml >> "$GITHUB_STEP_SUMMARY" || true
- name: Write runtime summary
if: always()
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
run: |
python tools/ci/runtime_summary.py --title jax >> "$GITHUB_STEP_SUMMARY" || true
- name: Build jax coverage json (MaxText E2E)
if: always()
continue-on-error: true
run: |
unset COVERAGE_PROCESS_START
# The JSON is consumed by the coverage-summary job (no per-job summary).
python -m coverage combine 2>/dev/null || true
python -m coverage json -o coverage_maxtext_e2e.json 2>/dev/null || true
- name: Upload jax coverage json
if: always()
continue-on-error: true
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: coverage-jax
path: coverage_maxtext_e2e.json
if-no-files-found: ignore
- name: Clean
if: always()
run: |
rm -rf ${PRIMUS_WORKDIR}/Primus-Turbo
rm -rf ${PRIMUS_WORKDIR}/Primus
# Remove E2E coverage artifacts + the injection .pth (persistent runner).
rm -f "${GITHUB_WORKSPACE}/.coverage_e2e"* "${GITHUB_WORKSPACE}/.coveragerc_e2e" coverage_maxtext_e2e.json || true
SP=$(python -c "import site; print(site.getsitepackages()[0])" 2>/dev/null) && rm -f "$SP/primus_e2e_coverage.pth" || true
coverage-summary:
# Combine the torch and jax coverage JSON artifacts into one unit-vs-E2E
# table. Uses module-level JSON (not coverage data), so no cross-runner path
# mapping is needed: torch (megatron/torchtitan) and jax (maxtext) E2E cover
# near-disjoint modules and are merged per module by max covered lines.
needs: [run-unittest-torch, run-unittest-jax]
if: always()
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
- name: Download coverage artifacts
uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
continue-on-error: true
with:
path: cov-artifacts
- name: Render combined coverage summary
continue-on-error: true
run: |
unit=$(find cov-artifacts -name coverage_unit.json | head -1)
torch=$(find cov-artifacts -name coverage_combined.json | head -1)
jax=$(find cov-artifacts -name coverage_maxtext_e2e.json | head -1)
if [ -n "$unit" ]; then
python tools/ci/coverage_summary.py "$unit" "Unit vs Unit+E2E (torch + jax)" $torch $jax >> "$GITHUB_STEP_SUMMARY"
else
echo "No unit coverage artifact found; skipping combined summary." >> "$GITHUB_STEP_SUMMARY"
fi