diff --git a/.github/workflows/collectivex-sweep.yml b/.github/workflows/collectivex-sweep.yml index b56bc285aa..0523837709 100644 --- a/.github/workflows/collectivex-sweep.yml +++ b/.github/workflows/collectivex-sweep.yml @@ -1,40 +1,34 @@ -# CollectiveX Sweep — one structured run instead of thousands of dispatches. -# -# Shape (mirrors the InferenceX CI tracker): setup -> sweep (a MATRIX job = "a job with other jobs -# in it") -> aggregate (the collector "at the end"). The matrix unit is a SHARD = one allocation that -# sweeps many cases sharing (sku, backend, mode, resource) — generate_matrix's own grouping, chunked -# so no cell exceeds the job budget. Each cell emits a handful of per-case JSONs; the aggregate job -# collects every shard into ONE line-delimited file (results/aggregate/*.ndjson) so there aren't -# thousands of individual result files. Run once per backend (deepep / uccl / flashinfer / -# deepep-hybrid / nccl-ep, + deepep_v2) for full parity. +# Generate a shard matrix, run each shard on its GPU pool, and upload raw results. name: CollectiveX Sweep +permissions: + actions: read + contents: read on: workflow_dispatch: inputs: backend: - description: EP library to sweep (deepep matrix is remapped onto the others, capability-filtered) + description: "EP library to sweep — 'all' runs every EP backend in one matrix" type: choice - default: deepep - options: [deepep, uccl, flashinfer, deepep-hybrid, nccl-ep] - deepep_v2: - description: DeepEP V2 from-source kernels (kernel_gen=v2; deepep backend only) - type: boolean - default: false + default: all + options: [all, deepep-v2, mori] suites: description: "'all' or comma-list of suite names" type: string default: all only_sku: - description: Restrict to one SKU (h100-dgxc|h200|b300|b200-dgxc|gb200|gb300|mi355x); blank = all + description: Restrict to one GHA runner pool (h100-dgxc|h200-dgxc|b300|b200-dgxc|gb200|gb300|mi300x|mi355x); blank = all type: string default: '' - max_cases: - description: Max cases per shard cell (chunk larger shards) + exclude_skus: + description: Comma-list of runner pools to drop from the matrix (partial run, e.g. b300); disjoint from only_sku; blank = none type: string - default: '14' - + default: '' + ep_sizes: + description: Keep only shards whose expert-parallel degree is in this comma-list (8 keeps EP8 only, dropping EP16 so GB SKUs co-schedule with 8-GPU SKUs; blank = all) + type: string + default: '' concurrency: - group: cx-sweep-${{ github.ref }}-${{ inputs.backend }}-${{ inputs.deepep_v2 }}-${{ inputs.only_sku }} + group: cx-${{ github.ref }}-${{ inputs.backend }}-${{ inputs.only_sku }} cancel-in-progress: false jobs: @@ -46,26 +40,52 @@ jobs: n: ${{ steps.gen.outputs.n }} steps: - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v5.0.0 - with: { clean: true } - - run: pip install --quiet pyyaml + with: { clean: true, persist-credentials: false } + # PyYAML is needed only here to resolve suites.yaml/workloads.yaml; shard extraction on GPU runners stays stdlib-only. + - name: Install matrix dependencies + run: python3 -m pip install --quiet PyYAML==6.0.2 - id: gen working-directory: experimental/CollectiveX + env: + INPUT_BACKEND: ${{ inputs.backend }} + INPUT_SUITES: ${{ inputs.suites }} + INPUT_ONLY_SKU: ${{ inputs.only_sku }} + INPUT_EXCLUDE_SKUS: ${{ inputs.exclude_skus }} + INPUT_EP_SIZES: ${{ inputs.ep_sizes }} run: | set -euo pipefail - ov=""; [ "${{ inputs.backend }}" != "deepep" ] && ov="--backend ${{ inputs.backend }}" - v2=""; [ "${{ inputs.deepep_v2 }}" = "true" ] && v2="--deepep-v2" - os=""; [ -n "${{ inputs.only_sku }}" ] && os="--only-sku ${{ inputs.only_sku }}" - # full matrix (with cases) -> artifact for the cells; slim (no cases) -> the strategy output. - python3 sweep_matrix.py --suites "${{ inputs.suites }}" --max-cases "${{ inputs.max_cases }}" $ov $v2 $os --out matrix_full.json >/dev/null - SLIM=$(python3 -c "import json;m=json.load(open('matrix_full.json'));print(json.dumps({'include':[{k:v for k,v in x.items() if k!='cases'} for x in m['include']]}))") - echo "matrix=$SLIM" >> "$GITHUB_OUTPUT" - echo "n=$(python3 -c "import json;print(len(json.load(open('matrix_full.json'))['include']))")" >> "$GITHUB_OUTPUT" - python3 -c "import json;m=json.load(open('matrix_full.json'));print('shard-cells:',len(m['include']),'cases:',sum(x['n'] for x in m['include']))" + args=(--suites "$INPUT_SUITES") + case "$INPUT_BACKEND" in + all) args+=(--backends all) ;; + *) args+=(--backend "$INPUT_BACKEND") ;; + esac + [ -n "$INPUT_ONLY_SKU" ] && args+=(--only-sku "$INPUT_ONLY_SKU") + [ -n "$INPUT_EXCLUDE_SKUS" ] && args+=(--exclude-skus "$INPUT_EXCLUDE_SKUS") + [ -n "$INPUT_EP_SIZES" ] && args+=(--ep-sizes "$INPUT_EP_SIZES") + python3 sweep_matrix.py "${args[@]}" --out matrix_full.json >/dev/null + python3 - "$GITHUB_OUTPUT" <<'PY' + import json + import pathlib + import sys + + matrix = json.loads(pathlib.Path("matrix_full.json").read_text()) + cells = matrix["include"] + slim = {"include": [ + {key: value for key, value in cell.items() if key != "case_ids"} + for cell in cells + ]} + with open(sys.argv[1], "a", encoding="utf-8") as output: + output.write(f"matrix={json.dumps(slim, separators=(',', ':'))}\n") + output.write(f"n={len(cells)}\n") + print(f"execution-cells: {len(cells)}") + PY - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: cxsweep-matrix-${{ github.run_id }} path: experimental/CollectiveX/matrix_full.json if-no-files-found: error + overwrite: true + retention-days: 14 # ---- sweep: ONE matrix cell per shard (the parent job with child jobs) ---- sweep: @@ -73,82 +93,170 @@ jobs: if: ${{ fromJSON(needs.setup.outputs.n) > 0 }} strategy: fail-fast: false - max-parallel: 10 # don't saturate the ~20-runner fleet; cells queue as slots free + # Fixed global cap; real throttling is each SKU pool's runner count and its + # cluster's Slurm partition. setup interleaves shards across SKUs, so the + # first jobs under this cap spread over pools instead of queuing on one. + max-parallel: 10 matrix: ${{ fromJSON(needs.setup.outputs.matrix) }} - # h200 label spans two clusters; pin to the validated dgxc pool (mirrors collectivex-experimental). - runs-on: ${{ matrix.sku == 'h200' && 'h200-dgxc' || matrix.sku }} + runs-on: ${{ matrix.sku }} timeout-minutes: 350 env: - CX_BENCH: ${{ matrix.backend }} - CX_DEEPEP_V2: ${{ matrix.deepep_v2 && '1' || '' }} - CX_NODES: ${{ matrix.nodes }} - CX_SHARD_FILE: results/.shard_${{ matrix.id }}.json + COLLX_BENCH: ${{ matrix.backend }} + COLLX_NODES: ${{ matrix.nodes }} + COLLX_GPUS_PER_NODE: ${{ matrix.gpus_per_node }} + COLLX_SCALE_UP_DOMAIN: ${{ matrix.scale_up_domain }} + COLLX_SHARD_FILE: .shards/${{ matrix.id }}.json + COLLX_SHARD_SKU: ${{ matrix.sku }} + COLLECTIVEX_CANONICAL_GHA: '1' COLLECTIVEX_SOURCE_SHA: ${{ github.sha }} - CX_NODELIST: ${{ matrix.sku == 'mi355x' && 'mia1-p01-g10,mia1-p01-g15' || '' }} - CX_STAGE_DIR: ${{ matrix.sku == 'gb200' && '/mnt/lustre01/users-public/sa-shared/cx-stage' || '' }} + COLLECTIVEX_ARTIFACT_NAME: ${{ format('cxshard-{0}-{1}-{2}', matrix.id, github.run_id, github.run_attempt) }} + # Consolidated shards run one bounded build-group in one Slurm allocation. + # MI300X's compute partition has a 180-minute ceiling; the other production + # pools accept 300 minutes. Allocations release as soon as the shard finishes. + COLLX_TIME: ${{ matrix.sku == 'mi300x' && '180' || '300' }} + COLLECTIVEX_EXECUTION_ID: ${{ github.run_id }}_${{ github.run_attempt }}_${{ matrix.id }} + COLLX_JOB_PARENT: ${{ matrix.launcher == 'mi-amds' && format('/tmp/inferencex-collectivex-parent-{0}-{1}-{2}', github.run_id, github.run_attempt, matrix.id) || '/tmp' }} + COLLX_JOB_ROOT: ${{ matrix.launcher == 'mi-amds' && format('/tmp/inferencex-collectivex-parent-{0}-{1}-{2}/inferencex-collectivex-{0}-{1}-{2}', github.run_id, github.run_attempt, matrix.id) || format('/tmp/inferencex-collectivex-{0}-{1}-{2}', github.run_id, github.run_attempt, matrix.id) }} + COLLX_SOURCE_ROOT: ${{ matrix.launcher == 'mi-amds' && format('/tmp/inferencex-collectivex-parent-{0}-{1}-{2}/inferencex-collectivex-{0}-{1}-{2}/source', github.run_id, github.run_attempt, matrix.id) || format('/tmp/inferencex-collectivex-{0}-{1}-{2}/source', github.run_id, github.run_attempt, matrix.id) }} + HOME: ${{ matrix.launcher == 'mi-amds' && format('/tmp/inferencex-collectivex-parent-{0}-{1}-{2}/inferencex-collectivex-{0}-{1}-{2}/home', github.run_id, github.run_attempt, matrix.id) || format('/tmp/inferencex-collectivex-{0}-{1}-{2}/home', github.run_id, github.run_attempt, matrix.id) }} steps: - - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v5.0.0 - with: { clean: true } + # Create the private mode-0700 /tmp job root and git-fetch the exact COLLECTIVEX_SOURCE_SHA into it; shards never execute from the runner's shared checkout. + - name: Prepare isolated source + id: source + env: + COLLECTIVEX_REPOSITORY: ${{ github.repository }} + run: | + set -euo pipefail + [ "${COLLX_JOB_ROOT%/*}" = "$COLLX_JOB_PARENT" ] \ + && [[ "${COLLX_JOB_ROOT##*/}" =~ ^inferencex-collectivex-[0-9]+-[0-9]+-[A-Za-z0-9._-]+$ ]] \ + || { echo "CollectiveX isolated root is invalid" >&2; exit 1; } + [ "$COLLX_SOURCE_ROOT" = "$COLLX_JOB_ROOT/source" ] \ + || { echo "CollectiveX source root is invalid" >&2; exit 1; } + if [ "$COLLX_JOB_PARENT" != /tmp ]; then + [[ "$COLLX_JOB_PARENT" =~ ^/tmp/inferencex-collectivex-parent-[0-9]+-[0-9]+-[A-Za-z0-9._-]+$ ]] \ + || { echo "CollectiveX isolated parent is invalid" >&2; exit 1; } + shared_parent="$GITHUB_WORKSPACE/.collectivex-jobs" + install -d -m 700 "$shared_parent" + [ "$(stat -c '%u:%a' "$shared_parent")" = "$(id -u):700" ] \ + || { echo "CollectiveX shared parent is unsafe" >&2; exit 1; } + [ ! -e "$COLLX_JOB_PARENT" ] && [ ! -L "$COLLX_JOB_PARENT" ] \ + || { echo "CollectiveX isolated parent already exists" >&2; exit 1; } + ln -s "$shared_parent" "$COLLX_JOB_PARENT" + fi + umask 077 + [ ! -e "$COLLX_JOB_ROOT" ] && [ ! -L "$COLLX_JOB_ROOT" ] \ + || { echo "CollectiveX isolated root already exists" >&2; exit 1; } + mkdir -m 700 "$COLLX_JOB_ROOT" + trap 'rc=$?; if [ "$rc" != 0 ]; then rm -rf -- "$COLLX_JOB_ROOT"; [ "$COLLX_JOB_PARENT" = /tmp ] || rm -f -- "$COLLX_JOB_PARENT"; fi; exit "$rc"' EXIT + mkdir -m 700 "$HOME" "$COLLX_JOB_ROOT/control" "$COLLX_JOB_ROOT/artifact" "$COLLX_SOURCE_ROOT" + export GIT_CONFIG_NOSYSTEM=1 GIT_CONFIG_GLOBAL=/dev/null + if ! git init -q "$COLLX_SOURCE_ROOT" \ + || ! git -C "$COLLX_SOURCE_ROOT" remote add origin \ + "https://github.com/${COLLECTIVEX_REPOSITORY}.git" \ + || ! git -C "$COLLX_SOURCE_ROOT" -c credential.helper= -c protocol.version=2 \ + fetch -q --no-tags --depth=1 origin "$COLLECTIVEX_SOURCE_SHA" \ + || ! git -C "$COLLX_SOURCE_ROOT" -c advice.detachedHead=false \ + checkout -q --detach FETCH_HEAD \ + || [ "$(git -C "$COLLX_SOURCE_ROOT" rev-parse HEAD)" != "$COLLECTIVEX_SOURCE_SHA" ]; then + echo "CollectiveX source preparation failed" >&2 + exit 1 + fi + echo 'prepared=true' >> "$GITHUB_OUTPUT" + trap - EXIT - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 with: name: cxsweep-matrix-${{ github.run_id }} - path: experimental/CollectiveX - - name: Extract this shard's cases (stdlib only — no runner deps) - working-directory: experimental/CollectiveX + path: ${{ env.COLLX_JOB_ROOT }}/control + # Materialize this shard's .shards/.json case control from the run-scoped matrix artifact downloaded above. + - name: Extract execution control run: | set -euo pipefail - python3 -c " - import json - m=json.load(open('matrix_full.json')) - s=[x for x in m['include'] if x['id']=='${{ matrix.id }}'] - assert s, 'shard ${{ matrix.id }} not in matrix' - s=s[0] - json.dump({'id':s['id'],'sku':s['sku'],'backend':s['backend'],'nodes':s['nodes'],'deepep_v2':s['deepep_v2'],'cases':s['cases']}, open('results/.shard_${{ matrix.id }}.json','w')) - print('shard ${{ matrix.id }}:', len(s['cases']), 'cases') - " - - name: Sweep shard ${{ matrix.id }} (${{ matrix.n }} cases, one allocation) + cd "$COLLX_SOURCE_ROOT/experimental/CollectiveX" 2>/dev/null \ + || { echo "CollectiveX source is unavailable" >&2; exit 1; } + python3 sweep_matrix.py \ + --extract-from "$COLLX_JOB_ROOT/control/matrix_full.json" \ + --shard-id '${{ matrix.id }}' \ + --out '${{ env.COLLX_SHARD_FILE }}' >/dev/null + # Merge the base + per-SKU operator-config secrets into one validated mode-0600 file, then run the SKU launcher (allocate -> stage -> build -> run cases). + - name: Execute sweep cell ${{ matrix.id }} + id: sweep_shard env: - RUNNER_NAME: ${{ runner.name }} - run: bash "experimental/CollectiveX/launchers/launch_${RUNNER_NAME%%_*}.sh" - - name: Shard summary - if: always() - run: python3 experimental/CollectiveX/summarize.py --results-dir experimental/CollectiveX/results --markdown >> "$GITHUB_STEP_SUMMARY" || true + COLLECTIVEX_OPERATOR_CONFIG_CONTENT: ${{ secrets.COLLECTIVEX_OPERATOR_CONFIG_V1 }} + COLLECTIVEX_NETWORK_CONFIG_CONTENT: ${{ secrets.COLLECTIVEX_NETWORK_CONFIG_V1 }} + COLLECTIVEX_H100_CONFIG_CONTENT: ${{ secrets.COLLECTIVEX_H100_CONFIG_V1 }} + COLLECTIVEX_B300_CONFIG_CONTENT: ${{ secrets.COLLECTIVEX_B300_CONFIG_V1 }} + COLLECTIVEX_B200_CONFIG_CONTENT: ${{ secrets.COLLECTIVEX_B200_CONFIG_V1 }} + COLLECTIVEX_MI300_CONFIG_CONTENT: ${{ secrets.COLLECTIVEX_MI300_CONFIG_V1 }} + COLLECTIVEX_MI355_CONFIG_CONTENT: ${{ secrets.COLLECTIVEX_MI355_CONFIG_V1 }} + COLLECTIVEX_OPERATOR_CONFIG_REQUIRED: '1' + run: | + set -euo pipefail + umask 077 + operator_config="$COLLX_JOB_ROOT/operator-config.json" + python3 "$COLLX_SOURCE_ROOT/experimental/CollectiveX/runtime/config.py" \ + merge-operator-config "$operator_config" + # The merged config has been written to $operator_config; hand the + # launcher the path only. The raw secret *_CONTENT vars must be cleared + # here, otherwise collx_load_operator_config sees them still set and tries + # to re-create operator-config.json (O_EXCL) that this step already + # wrote, failing with "cannot create ephemeral runner configuration". + unset COLLECTIVEX_OPERATOR_CONFIG_CONTENT COLLECTIVEX_NETWORK_CONFIG_CONTENT + unset COLLECTIVEX_H100_CONFIG_CONTENT COLLECTIVEX_B300_CONFIG_CONTENT + unset COLLECTIVEX_B200_CONFIG_CONTENT + unset COLLECTIVEX_MI300_CONFIG_CONTENT COLLECTIVEX_MI355_CONFIG_CONTENT + unset COLLECTIVEX_OPERATOR_CONFIG_REQUIRED + export COLLECTIVEX_OPERATOR_CONFIG="$operator_config" + export COLLECTIVEX_OPERATOR_CONFIG_EPHEMERAL=1 + cd "$COLLX_SOURCE_ROOT" 2>/dev/null \ + || { echo "CollectiveX source is unavailable" >&2; exit 1; } + bash "experimental/CollectiveX/launchers/launch_${{ matrix.launcher }}.sh" + # always(): cancel any Slurm allocation a killed launcher left recorded, append the summary table, and stage result JSONs so a red or partial leg still uploads. + - name: Summarize and stage shard results + id: stage_artifact + if: ${{ always() && steps.source.outputs.prepared == 'true' }} + run: | + set -euo pipefail + cd "$COLLX_SOURCE_ROOT" 2>/dev/null \ + || { echo "CollectiveX source is unavailable" >&2; exit 1; } + source experimental/CollectiveX/runtime/common.sh + collx_cleanup_allocation "$COLLX_JOB_ROOT" \ + || { echo "CollectiveX allocation cleanup failed; results may still be changing" >&2; exit 1; } + python3 experimental/CollectiveX/summarize.py \ + --results-dir experimental/CollectiveX/results >> "$GITHUB_STEP_SUMMARY" || true + shopt -s nullglob + results=(experimental/CollectiveX/results/*.json) + if [ "${#results[@]}" -eq 0 ]; then + echo "staged=false" >> "$GITHUB_OUTPUT" + echo "No result JSON to stage; leg produced none." + exit 0 + fi + cp -- "${results[@]}" "$COLLX_JOB_ROOT/artifact/" + echo "staged=true" >> "$GITHUB_OUTPUT" + # The neutral per-case JSONs are the workflow's only output; a consumer downloads them and decides what to display. - name: Upload shard results - if: always() + id: upload_artifact + if: always() && steps.stage_artifact.outputs.staged == 'true' uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: - name: cxshard-${{ matrix.id }}-${{ github.run_id }} - path: experimental/CollectiveX/results/*.json # glob skips the hidden .shard_*.json - if-no-files-found: warn - - # ---- aggregate: collect every shard into ONE ndjson (the "result aggregator at the end") ---- - aggregate: - needs: sweep - if: always() - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v5.0.0 - with: { clean: true } - - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 - with: - pattern: cxshard-*-${{ github.run_id }} - path: _shards - merge-multiple: true - - name: Aggregate shards -> one ndjson - working-directory: experimental/CollectiveX + name: ${{ format('cxshard-{0}-{1}-{2}', matrix.id, github.run_id, github.run_attempt) }} + path: | + ${{ env.COLLX_JOB_ROOT }}/artifact/*.json + if-no-files-found: error + retention-days: 14 + # Delete the private /tmp job root after the preceding step releases any allocation. + - name: Cleanup isolated workspace + if: ${{ always() && steps.source.outputs.prepared == 'true' }} run: | set -euo pipefail - tag="${{ inputs.backend }}${{ inputs.deepep_v2 && '-v2' || '' }}" - python3 aggregate_results.py --in-dir ../../_shards --out "results/aggregate/collectivex_${tag}_${{ github.run_id }}.ndjson" - { - echo "## CollectiveX sweep aggregate (${tag})" - echo '```' - wc -l results/aggregate/*.ndjson 2>/dev/null || echo "no ndjson" - echo '```' - } >> "$GITHUB_STEP_SUMMARY" - - name: Upload aggregate - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 - with: - name: cxsweep-aggregate-${{ inputs.backend }}${{ inputs.deepep_v2 && '-v2' || '' }}-${{ github.run_id }} - path: experimental/CollectiveX/results/aggregate/*.ndjson - if-no-files-found: warn + [ "$COLLX_JOB_PARENT" = /tmp ] \ + || [[ "$COLLX_JOB_PARENT" =~ ^/tmp/inferencex-collectivex-parent-[0-9]+-[0-9]+-[A-Za-z0-9._-]+$ ]] \ + || { echo "CollectiveX cleanup parent is invalid" >&2; exit 1; } + [ "${COLLX_JOB_ROOT%/*}" = "$COLLX_JOB_PARENT" ] \ + && [[ "${COLLX_JOB_ROOT##*/}" =~ ^inferencex-collectivex-[0-9]+-[0-9]+-[A-Za-z0-9._-]+$ ]] \ + || { echo "CollectiveX cleanup root is invalid" >&2; exit 1; } + [ "$COLLX_SOURCE_ROOT" = "$COLLX_JOB_ROOT/source" ] \ + || { echo "CollectiveX cleanup source is invalid" >&2; exit 1; } + cd "$COLLX_SOURCE_ROOT" 2>/dev/null \ + || { echo "CollectiveX source is unavailable" >&2; exit 1; } + rm -rf -- "$COLLX_JOB_ROOT" + [ "$COLLX_JOB_PARENT" = /tmp ] || rm -f -- "$COLLX_JOB_PARENT" diff --git a/AGENTS.md b/AGENTS.md index 6a385a5c16..f1eac75137 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -2,6 +2,8 @@ Guidance for AI agents working with InferenceX. +> **Temporary language override (highest repository priority):** Until the user explicitly removes this override, ignore any instruction elsewhere in this repository that requires Simplified Chinese or bilingual output, including instructions in later sections, nested agent files, contribution guides, checklists, and templates. Work in English only: do not create or update `_zh.md` files, Chinese documentation, translations, or Chinese text in commits, pull requests, issues, reviews, or comments. Leave existing Chinese content untouched. + > **Mandatory reading: [`CONTRIBUTING.md`](CONTRIBUTING.md)** — read it before opening or reviewing any PR. It covers the full PR review flow, the CODEOWNER sign-off process, the `/reuse-sweep-run` merge path, post-merge responsibilities, and critical cluster rules (e.g. never leaving root-owned files on AMD runners). > **PR and GitHub-issue titles & descriptions must be bilingual — include a Simplified Chinese version in addition to English.** Title format: ` / <中文标题>`. In the PR/issue body, follow the English content with its Chinese translation (e.g. a `## 中文说明` section mirroring the summary; don't translate code blocks, logs, or stack traces — summarize around them). **PR comments must include a Chinese translation too** — conversation comments, review summaries, and inline review comments alike: short comments as a single ` / <中文>` line, longer ones with the Chinese translation as a trailing paragraph (`中文：...`). Exception: the CODEOWNER sign-off template stays English-verbatim (the sign-off verifier triggers on its exact phrase); bot-generated comments follow their own workflow templates. This applies to every PR and every issue, matching the bilingual docs rule in Code Conventions. diff --git a/experimental/CollectiveX/.gitignore b/experimental/CollectiveX/.gitignore new file mode 100644 index 0000000000..f34ac59bad --- /dev/null +++ b/experimental/CollectiveX/.gitignore @@ -0,0 +1,14 @@ +__pycache__/ +*.pyc +results/ +unsupported/ +.shards/ +.collx_workloads/ +.collx_backend/ +/matrix_full.json +gpucore.* + +# Local plans and infrastructure inventory. +goal.md +notes.md +private-infra.md diff --git a/experimental/CollectiveX/README.md b/experimental/CollectiveX/README.md new file mode 100644 index 0000000000..192eebf142 --- /dev/null +++ b/experimental/CollectiveX/README.md @@ -0,0 +1,149 @@ +# CollectiveX + +CollectiveX is an experimental MoE expert-parallel communication benchmark. It measures dispatch, +combine, and paired roundtrip latency across EP libraries and accelerator systems, then uploads +neutral result artifacts. + +CollectiveX schedules benchmarks, executes them on real allocations, and uploads the neutral +artifacts each run emits. It does not validate those artifacts, promote, rank, recommend, select, or +decide what a consumer displays. Any downstream display or comparison is the consumer's +responsibility. The full measurement methodology is in [docs/methodology.md](docs/methodology.md). + +## Execution Profile + +The workload uses packed placement and one pinned `fixed-profile` resource configuration per +backend/topology; there is no tuning sweep. Dispatch and combine are fixed BF16 on every backend; +precision is not a swept dimension. Every case runs the single normal-mode contract: + +- Normal mode uses `layout-and-dispatch-v1`, rank-deduplicated token payloads, and activation-only + combine. Coverage is uniform routing only. + +Cases use a fixed timing profile from `configs/suites.yaml`: 128 trials x 8 timed iterations (1024 samples per component) +with 32 synchronized full roundtrip warmups before each measured component at every trial/point. Component measurement order rotates each +trial so every timed component occupies every position in the sequence; each iteration takes the +cross-rank maximum before nearest-rank p50/p90/p95/p99, and roundtrip p99 is the headline latency. A +keyed BLAKE2b counter produces byte-identical routing and gate weights on every runtime. + +Correctness is checked against the reference activation. The combine gate is `rtol=0.05, atol=0.02` +for the BF16 communication path. Any failed rank or point makes the case ineligible in the result +it writes. + +The matrix covers H100, H200, B200, B300, GB200, GB300, MI300X, MI325X, and MI355X. `sweep_matrix.py` materializes +the requested SKUs, backends, EP sizes, and token ladders, then extracts strict per-shard controls +and rejects missing, stale, malformed, or altered shard controls. `--only-sku`, `--exclude-skus`, and +`--ep-sizes` select a subset; the matrix is generated per dispatch, with no frozen digest or locked +case count. + +| Systems | EP8 | EP16 | +|---|---|---| +| H100/H200/B200/B300 | 1x8 NVLink, scale-up | 2x8 NVLink + RDMA, scale-out | +| MI300X/MI325X/MI355X | 1x8 XGMI, scale-up | 2x8 XGMI + RDMA, scale-out | +| GB200/GB300 | 2x4 MNNVL, scale-up | 4x4 MNNVL, scale-up | + +Physical host count does not determine scope: both GB topologies stay inside one 72-GPU MNNVL +scale-up domain. + +| Backend | Current scope | +|---|---| +| DeepEP V2 | PR #605 `ElasticBuffer` plus exact upstream #630 and #640 fixes: LSA for scale-up and GIN for x86 EP16 scale-out | +| MoRI | AMD EP8 uses IntraNode-family kernels (MI355X IntraNode, MI300X/MI325X asyncLL); EP16 pins InterNodeV1 over 2x8 XGMI + RDMA | + +DeepEP V2 means the `ElasticBuffer` implementation introduced by +[DeepEP PR #605](https://github.com/deepseek-ai/DeepEP/pull/605), not a newer legacy `Buffer` build. +The pinned source is the [PR #630](https://github.com/deepseek-ai/DeepEP/pull/630) head, whose parent +is the #605 merge tree, plus the exact one-line library matcher from upstream +[PR #640](https://github.com/deepseek-ai/DeepEP/pull/640). The first fixes pure scale-up +initialization when GIN is unavailable; the second prevents NCCL shared-memory mappings from being +misclassified as duplicate NCCL libraries. Scale-up cases request NCCL Device API LSA and fail closed +unless the realized LSA team covers the full EP world. x86 EP16 scale-out cases instead require the +hybrid path with GIN, two logical scale-out domains represented by two physical RDMA ranks, and eight +scale-up ranks per domain; GB EP16 remains MNNVL scale-up and therefore uses LSA. Whether a given +SKU/backend/EP cell is attempted is a capability fact; whether it succeeded is decided by the +benchmark's return code. + +## Workflow And Artifacts + +`.github/workflows/collectivex-sweep.yml` has two jobs. `setup` generates a public-SKU matrix +(`backend`, `suites`, `only_sku`, `exclude_skus`, `ep_sizes` inputs) and uploads the matrix. +`sweep` extracts a strict ignored `.shards/.json` control per matrix entry, executes one +allocation per shard, fetches pinned DeepEP source before allocation when required, and uploads the +result artifacts with `always()` so a red or partial run still uploads. + +Each shard emits per-case result JSON and a small mechanical summary. A case counts as successful on +the benchmark's own return code; there is no completeness or privacy validation step, and failed or +unsupported cells produce no synthetic record. No step promotes a run, +builds a dataset, or advances a channel; the neutral artifacts are the output. A consumer downloads +them and decides what to display. + +Credentials stay in encrypted config and are never uploaded. Per-step runner logs are kept on the +runner for postmortem; result artifacts carry only the fields listed in the methodology. + +## Runner Configuration + +Runner-local Slurm and storage values use a strict per-SKU JSON document at +`$XDG_CONFIG_HOME/inferencex/collectivex.json` or `COLLECTIVEX_OPERATOR_CONFIG`. Unknown runners, +fields, duplicate keys, and non-JSON input fail closed; configuration is never evaluated as shell. +GHA passes encrypted `COLLECTIVEX_OPERATOR_CONFIG_V1` content to the launcher, which validates it +and overlays it, per field, onto the tracked baseline. The secret is optional: when it is absent a +SKU runs entirely from its tracked baseline, and any field the baseline already supplies needs no +secret. + +All public per-SKU platform data lives in the tracked `configs/platform_config.json` registry: +identity (vendor/arch/machine/product), fixed placement, launcher, the operator-config fields each +SKU must supply (`operator_fields`), the tracked baseline values for those fields (`operator` block, +overridden per field by an operator document/secret when present), and the SKU's scale-out RDMA +selectors (`network` block, overlaid the same way). A SKU whose `operator` block already covers its +`operator_fields` needs no secret; SKUs without a complete baseline stay secret-fed and fail closed +with a `validation-missing-required-*` marker if the secret is gone. `capability.py` derives EP +topologies from it, and a suite's `platforms: all` in `configs/suites.yaml` resolves to every SKU +registered there. + +Every selected non-MNNVL EP16 placement additionally requires `socket_ifname` and `rdma_devices` for +its operator-approved fabric; optional `ib_gid_index`, `rdma_service_level`, and `rdma_traffic_class` +are also allowlisted. Service level and traffic class are mapped into MoRI's RDMA/IO QoS environment. +CollectiveX does not heuristically select a management route or HCA. After allocation, every +non-MNNVL scale-out node must prove that all configured interfaces and active HCA ports exist before +backend setup. Scale-up and MNNVL jobs clear these overrides. Scale-out NCCL/RCCL is pinned to `IB` +with exact-match HCA selectors so a socket fallback fails instead of being mislabeled as RDMA. + +`ib_gid_index` is applied only when every selected HCA port reports an Ethernet link layer, where it +selects the operator-approved RoCE GID. Native InfiniBand profiles retain explicit HCA and service +level pinning but leave the RoCE-only GID override unset so NVSHMEM/NCCL can use the native LID path. +Mixed Ethernet and InfiniBand HCA lists are rejected. + +`stage_dir` is a pre-existing, runner-owned, non-symlinked base outside the checkout and workflow +workspace. It is not group- or world-writable and is visible at the same path on the runner and every +allocated node. Jobs create only a marked mode-0700 execution child, prove cross-node read/write +visibility, and remove that exact child after allocation teardown; they never mount the runner +checkout or create a stage beneath image storage on AMD. When an AMD operator row omits `stage_dir`, +the runner derives a private base beside its standard `_work` directory on the shared runner +filesystem; the root-owned squash cache is never used as a repository stage. + +H200, B200, and B300 runners may omit `stage_dir`; their isolated execution child is created under a +runner-owned mode-0700 base in the validated operating-system account home, independent of the +workflow's temporary `HOME`. H100 may also omit `stage_dir`; its private base is created beside, never +beneath, the configured shared container directory so it is compute-visible. Canonical B300 execution +ignores any legacy configured `stage_dir` and always uses the validated compute-visible account-home +base; an execution-ID suffix isolates parallel B300 workers. Canonical GB300 execution likewise +ignores its legacy group-writable `stage_dir` and derives an execution-specific private base beneath the +validated compute-visible account home. Backend preparation runs from that staged tree on every node. + +Enroot imports the configured image tag into a per-run-scoped squash keyed by image tag and image +platform, so one run never reuses another run's imported filesystem. Backend source pins and image +references live in `configs/backends.json`; the DeepEP V2 build is fetched at the pinned commit, +verified by the wheel's commit-derived local-version tag and `ElasticBuffer` presence, and cached in +a cluster-local build cache keyed by architecture, image, and commit. Only the fixed +`/cx-cache` mount reaches the container. + +## Local Checks + +```bash +uv run --with-requirements experimental/CollectiveX/requirements.txt \ + python -m unittest discover experimental/CollectiveX/tests -p 'test_*.py' +uv run --with-requirements experimental/CollectiveX/requirements.txt \ + python experimental/CollectiveX/sweep_matrix.py --backends all --out /tmp/cx-matrix.json >/dev/null +bash -n experimental/CollectiveX/runtime/*.sh experimental/CollectiveX/launchers/*.sh +``` + +Core paths are `capability.py`, `configs/`, `sweep_matrix.py`, `summarize.py`, +`bench/`, `runtime/`, `launchers/`, and `tests/`. diff --git a/experimental/CollectiveX/bench/ep_backend.py b/experimental/CollectiveX/bench/ep_backend.py new file mode 100644 index 0000000000..fba516858f --- /dev/null +++ b/experimental/CollectiveX/bench/ep_backend.py @@ -0,0 +1,305 @@ +#!/usr/bin/env python3 +"""Shared lifecycle and input generation for EP backends.""" +from __future__ import annotations + +import abc +import types +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import torch + +from ep_harness import ( + time_us, + token_ladder, +) + + +@dataclass +class RankInputs: + """Inputs for one token-ladder shape at tokens_per_rank tokens on this rank. + + topk_idx/topk_weights are this rank's contiguous slice of the global routing trace + (host tensors; moved to device at make_problem time); activations are the rank's token + activations (already on device). The global trace is retained so Pass 1 can compute + routing statistics and input snapshots. + """ + + tokens_per_rank: int + topk_idx: "torch.Tensor" + topk_weights: "torch.Tensor" + activations: "torch.Tensor" + global_idx: "torch.Tensor | None" = None + global_weights: "torch.Tensor | None" = None + + +@dataclass +class WorkloadSpec: + """Numeric shape + materialised inputs for one fully-specified sweep line. + + Fully default-constructible so make_inputs can early-return a tensor-free + spec (ok=False + rc) on an empty ladder; the driver prints message and + returns rc + """ + + ok: bool = True + rc: int = 0 + message: str = "" + ep_size: int = 0 + experts_per_rank: int = 0 + cap: "int | None" = None + dropped: list = field(default_factory=list) + max_tokens_per_rank: int = 0 + ladder: list = field(default_factory=list) + points: dict = field(default_factory=dict) + + +class EPBackend(abc.ABC): + """One expert-parallel dispatch/combine transport under a fixed benchmark contract. + + Subclasses implement the transport (create_buffer, dispatch, stage, + combine, recv_tokens, inspect_dispatch, combine_transformed); + everything the driver and the oracles need beyond that is provided here. + Dispatch and combine are fixed BF16, so no adapter selects a precision codec. + """ + + name: str = "" + SUPPORTED_MODES: tuple = ("normal",) + stage_device_work = False + combine_needs_redispatch = False + dispatch_needs_combine_cleanup = False + # Adapters that reduce activations and top-k weights independently must carry + # the complete local weighted expert sum in the activation tensor. + combine_weight_semantics = "unweighted-rank-sum" + roundtrip_only = False + + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + if not getattr(cls, "name", ""): + raise TypeError( + f"{cls.__name__} must declare a non-empty class-level `name`" + ) + + def __init__(self, args, rank, world_size, local_rank, device): + self.args = args + self.rank = rank + self.world_size = world_size + self.local_rank = local_rank + self.device = device + self.mode = args.mode + if self.mode not in self.SUPPORTED_MODES: + raise ValueError(f"{self.name} does not support mode {self.mode!r}") + + # ---- Abstract transport contract ------------------------------------------------- + + @abc.abstractmethod + def create_buffer(self, spec: WorkloadSpec): + """Size the communicator from spec before the first dispatch.""" + + @abc.abstractmethod + def dispatch(self, problem): + """Scatter tokens to their experts; return an opaque per-call handle.""" + + @abc.abstractmethod + def stage(self, problem, handle): + """Prepare the combine input on handle (copy into place).""" + + @abc.abstractmethod + def combine(self, problem, handle): + """Gather the staged tokens back to their source rank; return combined activations.""" + + @abc.abstractmethod + def recv_tokens(self, handle): + """Number of tokens this rank received in dispatch (stable for a fixed trace).""" + + @abc.abstractmethod + def inspect_dispatch(self, problem, handle): + """Normalized post-dispatch view for the token-rank correctness oracle.""" + + @abc.abstractmethod + def combine_transformed(self, problem, handle, transformed): + """Combine an oracle-transformed payload in place of the staged input.""" + + # ---- Input generation (shared) --------------------------------------------------- + + def buffer_cap(self, args): + """Max tokens/rank the communicator can serve, or None when unbounded.""" + return None + + def make_inputs(self, args) -> WorkloadSpec: + """Resolve the token ladder and materialise per-rank inputs for the sweep. + + Buffer sizing needs the ladder *numbers* (not the input tensors), so this + runs before create_buffer. Returns a tensor-free spec with ok=False + when the ladder is empty. + """ + ep_size = self.world_size + experts_per_rank = args.experts // ep_size + cap = self.buffer_cap(args) + ladder, dropped = token_ladder(args.tokens_ladder, cap) + if not ladder: + return WorkloadSpec( + ok=False, rc=2, + message=f"empty token ladder (phase={args.phase}, cap={cap})", + ) + spec = WorkloadSpec( + ep_size=ep_size, + experts_per_rank=experts_per_rank, + cap=cap, + dropped=list(dropped), + max_tokens_per_rank=max(ladder), + ladder=list(ladder), + ) + for tokens_per_rank in ladder: + spec.points[tokens_per_rank] = self._build_rank_inputs(args, tokens_per_rank) + return spec + + def _build_rank_inputs(self, args, tokens_per_rank) -> RankInputs: + """Build one rank's deterministic inputs for a tokens-per-rank shape.""" + import torch + import routing + + ep_size = self.world_size + num_logical = getattr(args, "num_logical_experts", args.experts) + global_tokens = tokens_per_rank * ep_size + idx_g, w_g = routing.build_global_routing( + global_tokens, num_logical, args.topk, args.routing, args.seed + ) + idx_s, w_s = routing.rank_slice(idx_g, w_g, self.rank, tokens_per_rank) + activations = routing.rank_activations( + tokens_per_rank, args.hidden, args.seed, self.rank, self.device, torch.bfloat16 + ) + return RankInputs( + tokens_per_rank=tokens_per_rank, + topk_idx=idx_s.contiguous(), + topk_weights=w_s.contiguous(), + activations=activations, + global_idx=idx_g, + global_weights=w_g, + ) + + def make_problem(self, T, idx, weights, x): + """Assemble the per-shape problem namespace (BF16 dispatch sends x directly).""" + import torch + + return types.SimpleNamespace( + T=T, + x=x, + dispatch_x=x, + topk_idx=idx.to(self._topk_idx_dtype()), + topk_weights=weights.to(torch.float32), + ) + + def _topk_idx_dtype(self): + """Integer dtype the backend's kernels expect for top-k routing indices.""" + import torch + return torch.int64 + + # ---- Timing template methods ----------------------------------------------------- + + def timed_components(self): + """Components measured for this backend: roundtrip always; the rest unless + the backend exposes only a stateful paired round trip.""" + components = ["roundtrip"] + if not self.roundtrip_only: + components.extend(["dispatch", "combine"]) + if self.stage_device_work: + components.append("stage") + return components + + def warm(self, problem, count): + """Untimed synchronized full round trips (fabric/clock warm-up; cold-jump-safe). + + Caches the dynamic receive cardinality once so adapters never read a device + scalar during a timed trial (the count is stable for a fixed routing trace). + """ + import torch + + for _ in range(count): + handle = self.dispatch(problem) + if not hasattr(problem, "recv_tokens"): + problem.recv_tokens = self.recv_tokens(handle) + self.stage(problem, handle) + self.combine(problem, handle) + torch.cuda.synchronize() + + def run_roundtrip(self, problem): + """One full dispatch -> stage -> combine round trip; returns combined activations.""" + handle = self.dispatch(problem) + self.stage(problem, handle) + return self.combine(problem, handle) + + def benchmark_component(self, component, problem, warmup, iters): + """Measure one named component; every component gets the same warm-up first.""" + if component == "roundtrip": + return self.benchmark_roundtrip(problem, warmup, iters) + if component == "dispatch": + return self.benchmark_dispatch(problem, warmup, iters) + if component == "stage": + return self.benchmark_stage(problem, warmup, iters) + if component == "combine": + return self.benchmark_combine(problem, warmup, iters) + raise RuntimeError(f"unknown timed component {component!r}") + + def benchmark_roundtrip(self, problem, warmup, iters): + import torch + + self.warm(problem, warmup) + return time_us(torch, lambda p=problem: self.run_roundtrip(p), 0, iters) + + def benchmark_dispatch(self, problem, warmup, iters): + import torch + + self.warm(problem, warmup) + + def finish_dispatch(hh, p=problem): + self.stage(p, hh) + self.combine(p, hh) + + dispatch_needs_cleanup = self.dispatch_needs_combine_cleanup + return time_us( + torch, lambda p=problem: self.dispatch(p), 0, iters, + post=finish_dispatch if dispatch_needs_cleanup else None, + ) + + def benchmark_stage(self, problem, warmup, iters): + import torch + + self.warm(problem, warmup) + + def prep_stage(p=problem): + return self.dispatch(p) + + return time_us( + torch, lambda hh, p=problem: self.stage(p, hh), 0, iters, pre=prep_stage, + ) + + def benchmark_combine(self, problem, warmup, iters): + import torch + + self.warm(problem, warmup) + + def prep_combine(p=problem): + hh = self.dispatch(p) + self.stage(p, hh) + return hh + + if self.combine_needs_redispatch: + return time_us( + torch, lambda hh, p=problem: self.combine(p, hh), 0, iters, pre=prep_combine, + ) + hh = prep_combine() + torch.cuda.synchronize() + return time_us(torch, lambda p=problem, hx=hh: self.combine(p, hx), 0, iters) + + def finalize(self, rc): + """Barrier and tear down the process group; returns rc.""" + import torch.distributed as dist + + try: + dist.barrier() + dist.destroy_process_group() + except Exception: + pass + return rc diff --git a/experimental/CollectiveX/bench/ep_deepep_v2.py b/experimental/CollectiveX/bench/ep_deepep_v2.py new file mode 100644 index 0000000000..ed29516549 --- /dev/null +++ b/experimental/CollectiveX/bench/ep_deepep_v2.py @@ -0,0 +1,282 @@ +#!/usr/bin/env python3 +"""DeepEP PR #605 adapter with the exact upstream PR #630 and #640 fixes.""" + +from __future__ import annotations + +import inspect +import json +import os +import sys +import types +from pathlib import Path + +import torch +import torch.distributed as dist +from ep_backend import EPBackend + +try: + import deep_ep + from deep_ep import ElasticBuffer # type: ignore +except Exception as exc: # pragma: no cover - requires the benchmark image + print(f"ERROR: DeepEP V2 import failed: {exc!r}", file=sys.stderr) + raise + + +# Source pins (PR #605 head + #630/#640 fixes) live in configs/backends.json; +# the launcher fetches and builds them from that checkout. This adapter no longer +# verifies the wheel's commit tag against the pin — it checks only that the loaded +# deep_ep exposes ElasticBuffer (the from-source PR #605 capability). + + +def _jit_cache_directory( + args, + world_size: int, + max_tokens: int, + allow_hybrid_mode: bool, + realized: dict[str, int | bool], +) -> str: + values = ( + args.runner, world_size, args.hidden, args.topk, args.experts, + getattr(args, "num_logical_experts", args.experts), max_tokens, + int(allow_hybrid_mode), realized["allocated_qps"], realized["num_sms"], + ) + return "jit-" + "-".join(str(value) for value in values) + + +def _require_cross_rank_equal(value, label: str) -> None: + gathered = [None] * dist.get_world_size() + dist.all_gather_object(gathered, value) + canonical = {json.dumps(item, sort_keys=True, separators=(",", ":")) for item in gathered} + if len(canonical) != 1: + raise RuntimeError(f"DeepEP V2 {label} differs across ranks") + + +# GIN/GDAKI allocates num_allocated_qps device QPs per peer rank on the local NIC +# (contexts x world_size QPs, before NCCL's own connection QPs). Upstream's hybrid +# default (129, or 65 with fast RDMA atomics) exhausts the per-NIC QP budget at +# EP16: construction dies in ncclDevCommCreate with ibv_create_qp ENOMEM once +# NCCL's regular QPs land on top (identical on H200 bare-metal and B200 pods; on +# CX-7 the budget sits between 784 and 1040 QPs — 49x16 initializes, 65x16 does +# not). Spending a fixed ~512-QP budget keeps every EP size inside that limit +# with headroom: EP8 resolves to 65 (the allocation CX-8 racks already run +# successfully), EP16 to 33 and EP32 to 17 (33 and 49 verified on the failing +# H200 pair). An explicit value also skips upstream's rank-local ibstat probe, +# which is not guaranteed to resolve identically across ranks. +_GIN_QP_BUDGET = 512 + + +def _hybrid_num_allocated_qps(world_size: int) -> int: + return max(9, 1 + _GIN_QP_BUDGET // world_size) + + +def _configure_gin_mode(args, world_size: int) -> bool: + scale_up_domain = int(args.scale_up_domain) + allow_hybrid_mode = world_size > scale_up_domain + if allow_hybrid_mode: + os.environ.pop("EP_DISABLE_GIN", None) + else: + os.environ["EP_DISABLE_GIN"] = "1" + return allow_hybrid_mode + + +def _lsa_topology_is_valid( + gin_enabled: bool, + world_size: int, + scale_up_domain: int, + config: dict[str, int | bool], +) -> bool: + if gin_enabled: + domains = world_size // scale_up_domain + return ( + world_size % scale_up_domain == 0 + and domains > 1 + and config["physical_rdma_ranks"] == domains + and config["physical_nvlink_ranks"] == scale_up_domain + and config["logical_scaleout_ranks"] == domains + and config["logical_scaleup_ranks"] == scale_up_domain + and config["is_scaleup_nvlink"] is True + ) + return ( + config["physical_rdma_ranks"] == 1 + and config["physical_nvlink_ranks"] == world_size + and config["logical_scaleout_ranks"] == 1 + and config["logical_scaleup_ranks"] == world_size + and config["is_scaleup_nvlink"] is True + ) + + +def _require_runtime() -> None: + """Capability check only: the loaded deep_ep must expose ElasticBuffer (still + catches the b300 image-bundled deep_ep 1.2.1 shadowing the from-source build, + which lacks the class).""" + if not inspect.isclass(ElasticBuffer) or ElasticBuffer.__name__ != "ElasticBuffer": + raise RuntimeError("invalid DeepEP V2 runtime: deep_ep.ElasticBuffer is absent") + + +class DeepEPV2Backend(EPBackend): + name = "deepep-v2" + # Invariant by identity contract: this backend IS the PR #605 ElasticBuffer + # implementation; LSA vs hybrid GIN are transport paths, not kernel families. + kernel_generation = "v2-elastic-buffer" + stage_device_work = False + combine_needs_redispatch = False + combine_weight_semantics = "unweighted-rank-sum" + + def __init__(self, args, rank, world_size, local_rank, device): + # deepep-v2 is normal-mode only; base SUPPORTED_MODES=("normal",) enforces it. + super().__init__(args, rank, world_size, local_rank, device) + self.group = dist.group.WORLD + + def create_buffer(self, spec): + # max_tokens is the measured-ladder maximum; the historical values (which + # also folded in the conditioning ramp) are identical because the ramp + # never exceeded the measured maximum, so the JIT directory stays stable. + args, world_size, device = self.args, self.world_size, self.device + self.max_tokens = spec.max_tokens_per_rank + _require_runtime() + jit_root = Path(os.environ["EP_JIT_CACHE_DIR"]) + scale_up_domain = int(args.scale_up_domain) + allow_hybrid_mode = _configure_gin_mode(args, world_size) + gin_enabled = allow_hybrid_mode + self.buffer = ElasticBuffer( + self.group, + num_max_tokens_per_rank=self.max_tokens, + hidden=args.hidden, + num_topk=args.topk, + use_fp8_dispatch=False, # BF16 communication path. + deterministic=False, + allow_hybrid_mode=allow_hybrid_mode, + allow_multiple_reduction=True, + prefer_overlap_with_compute=True, + num_gpu_timeout_secs=100, + explicitly_destroy=True, + # 0 is upstream's use-the-default sentinel; only hybrid (GIN) mode + # needs the explicit budget-derived allocation. + num_allocated_qps=( + _hybrid_num_allocated_qps(world_size) if allow_hybrid_mode else 0 + ), + ) + tuning_num_experts = int(getattr(args, "num_logical_experts", args.experts)) + self.num_sms = int( + self.buffer.get_theoretical_num_sms(tuning_num_experts, args.topk) + ) + self.num_qps = int(self.buffer.get_theoretical_num_qps(self.num_sms)) + properties = torch.cuda.get_device_properties(device) + device_sms = int(properties.multi_processor_count) + jit_config = { + "num_sms": self.num_sms, + "num_qps": self.num_qps, + "allocated_qps": int(self.buffer.num_allocated_qps), + "logical_scaleout_ranks": int(self.buffer.num_scaleout_ranks), + "logical_scaleup_ranks": int(self.buffer.num_scaleup_ranks), + "physical_rdma_ranks": int(self.buffer.num_rdma_ranks), + "physical_nvlink_ranks": int(self.buffer.num_nvlink_ranks), + "is_scaleup_nvlink": self.buffer.num_scaleup_ranks == self.buffer.num_nvlink_ranks, + "device_arch_major": int(properties.major), + "device_arch_minor": int(properties.minor), + "device_sms": device_sms, + "device_smem_bytes": int(properties.shared_memory_per_block_optin), + "gpu_timeout_cycles": 100 * int(properties.clock_rate) * 1000, + } + _require_cross_rank_equal(jit_config, "JIT configuration") + if not _lsa_topology_is_valid( + gin_enabled, world_size, scale_up_domain, jit_config + ): + raise RuntimeError("DeepEP V2 realized communication domains differ from topology") + jit_cache_directory = _jit_cache_directory( + args, + world_size, + self.max_tokens, + allow_hybrid_mode, + jit_config, + ) + os.environ["EP_JIT_CACHE_DIR"] = str(jit_root / jit_cache_directory) + realized_config = { + "num_max_tokens_per_rank": self.max_tokens, + **jit_config, + } + _require_cross_rank_equal(realized_config, "realized tuning/topology") + + def _topk_idx_dtype(self): + # DeepEP V2's kernels key routing indices on deep_ep.topk_idx_t, not int64. + return deep_ep.topk_idx_t + + def dispatch(self, p): + recv_x, recv_topk_idx, recv_topk_weights, handle, _ = self.buffer.dispatch( + p.dispatch_x, + topk_idx=p.topk_idx, + topk_weights=p.topk_weights, + num_experts=self.args.experts, + num_max_tokens_per_rank=self.max_tokens, + expert_alignment=1, + num_sms=self.num_sms, + num_qps=self.num_qps, + async_with_compute_stream=False, + do_handle_copy=True, + do_cpu_sync=True, + do_expand=False, + ) + return types.SimpleNamespace( + recv_x=recv_x, + recv_topk_idx=recv_topk_idx, + recv_topk_weights=recv_topk_weights, + handle=handle, + ) + + def stage(self, p, h): + # BF16: the received buffer is already the semantic payload to combine. + h.combine_input = h.recv_x + + def combine(self, p, h): + combined_x, _, _ = self.buffer.combine( + h.combine_input, + handle=h.handle, + num_sms=self.num_sms, + num_qps=self.num_qps, + async_with_compute_stream=False, + ) + return combined_x + + def inspect_dispatch(self, p, h): + count = self.recv_tokens(h) + local_idx = h.recv_topk_idx[:count] + valid = local_idx >= 0 + expert_ids = torch.where( + valid, + local_idx + self.rank * (self.args.experts // self.world_size), + local_idx, + ) + local = local_idx[valid].to(torch.int64) + return types.SimpleNamespace( + payload=h.recv_x[:count], + expert_ids=expert_ids, + weights=h.recv_topk_weights[:count].masked_fill(~valid, 0), + local_expert_counts=torch.bincount( + local, minlength=self.args.experts // self.world_size + ), + ) + + def combine_transformed(self, p, h, transformed): + combine_input = torch.zeros_like(h.recv_x) + combine_input[: transformed.shape[0]].copy_(transformed.to(combine_input.dtype)) + combined, _, _ = self.buffer.combine( + combine_input, + handle=h.handle, + num_sms=self.num_sms, + num_qps=self.num_qps, + async_with_compute_stream=False, + ) + return combined + + def recv_tokens(self, h): + return int(h.handle.psum_num_recv_tokens_per_scaleup_rank[-1].item()) + + def finalize(self, rc): + try: + dist.barrier() + self.buffer.destroy() + dist.barrier() + dist.destroy_process_group() + except Exception: + return 1 + return rc diff --git a/experimental/CollectiveX/bench/ep_harness.py b/experimental/CollectiveX/bench/ep_harness.py new file mode 100644 index 0000000000..76e7e526dd --- /dev/null +++ b/experimental/CollectiveX/bench/ep_harness.py @@ -0,0 +1,830 @@ +#!/usr/bin/env python3 +"""Shared EP timing, correctness, and result generation.""" +from __future__ import annotations + +import argparse +import datetime as _dt +import json +import math +import os +import re + + +_CASE_ID = re.compile(r"^[a-z0-9][a-z0-9.-]*$") +_NON_SLUG = re.compile(r"[^a-z0-9]+") + + +def is_case_id(value) -> bool: + return bool(isinstance(value, str) and _CASE_ID.fullmatch(value)) + + +def case_id(sku: str, case: dict) -> str: + parts = ( + sku, + case["backend"], + case["workload"], + case["mode"], + case["phase"], + f"ep{int(case['ep'])}", + case["routing"], + ) + values = [_NON_SLUG.sub("-", str(part).lower()).strip("-") for part in parts] + if not all(values): + raise ValueError("case ID contains an empty factor") + return "-".join(values) + + +# Workload and timing values arrive from configs/suites.yaml through the matrix. +CONDITIONING_ROUNDS_PER_SHAPE = 8 +# Dispatch and combine are fixed BF16, so the combine oracle uses one frozen gate. +ORACLE_RTOL = 5e-2 +ORACLE_ATOL = 2e-2 + +def logical_byte_provenance(logical_copies: int, hidden: int) -> dict[str, int]: + """Return comparable logical BF16 activation bytes for one direction. + + Dispatch and combine both move BF16 (2 bytes/value) with no separate scale + payload, so ``scale_bytes`` is always zero. + """ + if logical_copies < 0 or hidden < 0: + raise ValueError("logical byte dimensions must be non-negative") + activation_data_bytes = logical_copies * hidden * 2 + return { + "activation_data_bytes": activation_data_bytes, + "scale_bytes": 0, + "total_logical_bytes": activation_data_bytes, + } + +def format_collective_version(raw) -> str: + """Normalize PyTorch's tuple or packed NCCL/RCCL version representation.""" + if isinstance(raw, int): + if raw < 10_000: + return f"{raw // 1000}.{raw // 100 % 10}.{raw % 100}" + return f"{raw // 10_000}.{raw // 100 % 100}.{raw % 100}" + if isinstance(raw, (tuple, list)): + return ".".join(map(str, raw)) + return str(raw) if raw not in (None, "") else "unknown" + + +def add_common_args(ap: argparse.ArgumentParser) -> None: + """Add the varying v1 inputs; fixed profile values are not CLI axes.""" + ap.add_argument("--mode", required=True, choices=["normal"]) + ap.add_argument("--phase", required=True, choices=["decode", "prefill"], + help="token-size regime label: decode (small T) / prefill (large T)") + ap.add_argument("--tokens-ladder", required=True, + help="space/comma-separated source-tokens-per-rank sweep; the matrix " + "supplies the workload's phase ladder from configs/suites.yaml") + ap.add_argument("--hidden", type=int, required=True) + ap.add_argument("--topk", type=int, required=True) + ap.add_argument("--experts", type=int, required=True, + help="TOTAL experts (fixed across EP degrees)") + ap.add_argument("--routing", required=True, choices=["uniform"]) + ap.add_argument("--case-id", required=True) + ap.add_argument("--suite", required=True) + ap.add_argument("--workload-name", required=True) + ap.add_argument("--seed", type=int, required=True, + help="routing-trace seed; part of the workload identity in configs/suites.yaml") + ap.add_argument( + "--version", + type=int, + required=True, + help="iterable benchmark version copied verbatim into the emitted result", + ) + # The single cross-SKU profile (and its rationale) lives in configs/suites.yaml + # `timing:`; the matrix bakes it into every scheduled case. + ap.add_argument("--warmup", type=int, required=True, + help="untimed full roundtrips before each trial/point") + ap.add_argument("--iters", type=int, required=True, + help="timed iterations per trial") + ap.add_argument("--trials", type=int, required=True, + help="timed trials") + # provenance / output + ap.add_argument("--runner", required=True) + ap.add_argument("--topology-class", required=True) + ap.add_argument("--transport", default="") + ap.add_argument("--scope", required=True, choices=["scale-up", "scale-out"]) + ap.add_argument("--scale-up-transport", required=True) + ap.add_argument("--scale-out-transport", default="") + ap.add_argument("--gpus-per-node", type=int, required=True) + ap.add_argument("--scale-up-domain", type=int, required=True) + ap.add_argument("--out", required=True) + + +def token_ladder(spec: str, cap: int | None) -> tuple[list[int], list[int]]: + """Return (ladder, dropped) from an explicit spec (there is no default — the + model-specific ladders live in configs/suites.yaml); positive ints; clamped to + `cap` with dropped points reported (never silently truncated).""" + want = sorted({t for t in (int(t) for t in spec.replace(",", " ").split() if t) if t > 0}) + if cap is not None: + return [t for t in want if t <= cap], [t for t in want if t > cap] + return want, [] + + +def trial_order(values: list, trial_index: int) -> list: + """Rotate and reverse values so each occupies every timing position.""" + if not values or len(values) != len(set(values)): + raise ValueError("trial order requires non-empty unique values") + if type(trial_index) is not int or trial_index < 0: + raise ValueError("trial_index must be a non-negative integer") + cycle, offset = divmod(trial_index, len(values)) + base = list(values) if cycle % 2 == 0 else list(reversed(values)) + return base[offset:] + base[:offset] + + +def percentile(xs: list[float], q: float) -> float: + if not xs: + return float("nan") + s = sorted(xs) + i = max(0, min(len(s) - 1, math.ceil(q / 100.0 * len(s)) - 1)) + return s[i] + + +def _pcts(xs): + return ({"p50": percentile(xs, 50), "p90": percentile(xs, 90), + "p95": percentile(xs, 95), "p99": percentile(xs, 99)} if xs else None) + + +def _component(percentiles, count, *, derived=False): + if percentiles is None: + return {"availability": "unavailable", "origin": None, + "percentiles_us": None, "sample_count": 0} + return { + "availability": "derived" if derived else "measured", + "origin": "derived-percentile-sum" if derived else "measured", + "percentiles_us": percentiles, + "sample_count": 0 if derived else count, + } + + +# The exact routing fields each row publishes — a whitelist so a new stat in +# routing.routing_stats never leaks into the artifact unreviewed. +_ROUTING_FIELDS = ( + "empty_expert_count", "empty_rank_count", "expert_assignment_rank_cv", + "expert_assignments_per_rank", "expert_load_cv", "expert_load_max", + "expert_load_mean", "expert_load_min", "fanout_histogram", "fanout_max", + "fanout_mean", "fanout_min", "hotspot_ratio", "locality", + "payload_copies_per_rank", "payload_rank_cv", "routed_copies", +) + + +def _write_bytes_atomic(path: str, payload: bytes) -> None: + os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True) + temporary = f"{path}.tmp-{os.getpid()}" + try: + with open(temporary, "wb") as handle: + handle.write(payload) + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, path) + finally: + try: + os.unlink(temporary) + except FileNotFoundError: + pass + + +def _write_json_atomic(path: str, value) -> None: + payload = ( + json.dumps(value, allow_nan=False, ensure_ascii=False, indent=2) + "\n" + ).encode() + return _write_bytes_atomic(path, payload) + + +def time_us(torch, fn, warmup: int, iters: int, pre=None, post=None) -> list[float]: + """Per-iteration CUDA-event latencies (µs) for THIS rank. + + Without `pre`: times `fn()`. With `pre`: runs `pre()` UNTIMED each iteration (sync + before the start event so its GPU work can't bleed in), then times `fn(pre_result)`. + `post(result)` runs after the end event and synchronization, so stateful backends can + consume/reset a timed operation without charging that cleanup to its latency. Returns + the raw per-iteration series; the caller reduces across ranks per iteration before + percentiling. + """ + def sample(): + arg = pre() if pre is not None else None + if pre is not None: + torch.cuda.synchronize() + s = torch.cuda.Event(enable_timing=True) + e = torch.cuda.Event(enable_timing=True) + s.record() + result = fn(arg) if pre is not None else fn() + e.record() + torch.cuda.synchronize() + elapsed = s.elapsed_time(e) * 1000.0 # ms -> us + if post is not None: + post(result) + torch.cuda.synchronize() + return elapsed + + for _ in range(max(0, warmup)): + if pre is not None: + a = pre() + torch.cuda.synchronize() + fn(a) + else: + fn() + # sync EACH warmup iteration, not just once after the loop: the measured-roundtrip fn + # interleaves dispatch+combine on a backend's persistent comm buffer, so back-to-back + # un-synced warmup iterations let iter N+1's dispatch race iter N's combine (CUDA abort + # on a rank -> NCCL-watchdog SIGABRT). Cheap (warmup is small); timed samples already sync. + torch.cuda.synchronize() + return [sample() for _ in range(iters)] + + +def kernel_generation(backend) -> str: + """Return the adapter's declared kernel family.""" + return getattr(backend, "kernel_generation", None) or "n-a" + + +def _reduce_vec(torch, dist, device, vals, op): + t = torch.tensor(vals, device=device, dtype=torch.float64) + dist.all_reduce(t, op=op) + return [float(x) for x in t.tolist()] + + +def _reduce_int(torch, dist, device, v: int, op) -> int: + t = torch.tensor([int(v)], device=device, dtype=torch.int64) + dist.all_reduce(t, op=op) + return int(t.item()) + + +def _same_tensors_across_ranks(torch, dist, device, *tensors) -> bool: + matches = True + for tensor in tensors: + observed = tensor.to(device=device, non_blocking=False) + reference = observed.clone() if dist.get_rank() == 0 else torch.empty_like(observed) + dist.broadcast(reference, src=0) + matches = matches and bool(torch.equal(observed, reference)) + result = torch.tensor([int(matches)], device=device, dtype=torch.int64) + dist.all_reduce(result, op=dist.ReduceOp.MIN) + return bool(result.item()) + + +def _normalized_expert_metadata(torch, expert_ids, weights): + """Sort each row by global expert ID while keeping -1 sentinels last.""" + valid = expert_ids >= 0 + keys = torch.where(valid, expert_ids.to(torch.int64), torch.full_like(expert_ids, 1 << 30)) + order = torch.argsort(keys, dim=1, stable=True) + sorted_ids = torch.gather(expert_ids.to(torch.int64), 1, order) + sorted_weights = torch.gather(weights.to(torch.float32), 1, order) + sorted_valid = sorted_ids >= 0 + return ( + torch.where(sorted_valid, sorted_ids, torch.full_like(sorted_ids, -1)), + sorted_weights.masked_fill(~sorted_valid, 0), + ) + + +def _expert_coefficients(torch, expert): + """Per-expert affine coefficients — the transform and its independently derived + expectation must use these exact formulas, so they exist only here.""" + scale = ((expert * 17 + 5) % 31 + 1).to(torch.float32) / 32 + offset_a = (((expert * 29 + 7) % 37) - 18).to(torch.float32) / 64 + offset_b = (((expert * 43 + 11) % 41) - 20).to(torch.float32) / 128 + return scale, offset_a, offset_b + + +def _column_pattern(torch, ncols, device): + columns = torch.arange(ncols, device=device, dtype=torch.int64) + return (((columns * 13) % 17) - 8).to(torch.float32) / 8 + + +def _expert_transform(torch, payload, expert_ids, weights, combine_weight_semantics): + """Build one local expert aggregate for the v1 unweighted combine contract.""" + if combine_weight_semantics != "unweighted-rank-sum": + raise ValueError("benchmark requires unweighted rank-sum combine") + valid = expert_ids >= 0 + expert = expert_ids.clamp(min=0).to(torch.int64) + gate = weights.to(torch.float32).masked_fill(~valid, 0) + scale, offset_a, offset_b = _expert_coefficients(torch, expert) + scale_sum = (gate * scale).sum(dim=1, keepdim=True) + offset_a_sum = (gate * offset_a).sum(dim=1, keepdim=True) + offset_b_sum = (gate * offset_b).sum(dim=1, keepdim=True) + pattern = _column_pattern(torch, payload.shape[1], payload.device) + transformed = ( + payload.float() * scale_sum + offset_a_sum + offset_b_sum * pattern.unsqueeze(0) + ) + return transformed.to(payload.dtype) + + +def _expected_transformed_combine(torch, problem): + """Independently derive sum_i gate_i * expert_i(x) for each source token.""" + semantic_x = getattr(problem, "oracle_x", problem.x) + expected = torch.zeros_like(semantic_x, dtype=torch.float32) + expert_ids = problem.topk_idx.to(torch.int64) + weights = problem.topk_weights.to(torch.float32) + pattern = _column_pattern(torch, semantic_x.shape[1], semantic_x.device) + for slot in range(expert_ids.shape[1]): + gate = weights[:, slot].unsqueeze(1) + scale, offset_a, offset_b = ( + c.unsqueeze(1) for c in _expert_coefficients(torch, expert_ids[:, slot]) + ) + expert_output = semantic_x.float() * scale + offset_a + offset_b * pattern.unsqueeze(0) + expected.add_(gate * expert_output) + return expected + + +_ORACLE_CHECKS = ( + "combine_values", "counts", "metadata", "multiplicity", "payload", + "source_set", "weights", +) + + +def _oracle_report(**fields): + """One report shape for both the fail-soft and full oracle paths, so every + emitted correctness dict carries identical keys.""" + report = { + "passed": False, + "atol": ORACLE_ATOL, + "rtol": ORACLE_RTOL, + "combine_weight_semantics": "undeclared", + "receive_count": 0, + "max_absolute_error": None, + "max_elementwise_relative_error": None, + "max_relative_error": None, + "max_weight_error": None, + "checks": dict.fromkeys(_ORACLE_CHECKS, False), + } + assert set(fields) <= set(report), sorted(set(fields) - set(report)) + report.update(fields) + assert set(report["checks"]) == set(_ORACLE_CHECKS) + return report + + +def _run_expert_oracle( + torch, + routing, + backend, + problem, + global_idx, + global_weights, + rank: int, + experts_per_rank: int, + seed: int, +): + """Verify one real dispatch/transform/combine without entering a timed region.""" + handle = backend.dispatch(problem) + torch.cuda.synchronize() + try: + view = backend.inspect_dispatch(problem, handle) + source_ids = routing.decode_source_ids(view.payload, seed) + except Exception as inspection_error: + # Drain the in-flight dispatch before reporting: an abandoned handle + # would deadlock the other ranks. + try: + problem.recv_tokens = backend.recv_tokens(handle) + backend.stage(problem, handle) + backend.combine(problem, handle) + torch.cuda.synchronize() + except Exception as cleanup_error: + raise inspection_error from cleanup_error + return _oracle_report( + combine_weight_semantics=getattr( + backend, "combine_weight_semantics", "undeclared" + ), + ) + + receive_count = int(view.payload.shape[0]) + shape_ok = ( + view.payload.ndim == 2 + and view.expert_ids.shape == (receive_count, problem.topk_idx.shape[1]) + and view.weights.shape == view.expert_ids.shape + ) + source_range = bool( + receive_count == 0 + or ((source_ids >= 0) & (source_ids < global_idx.shape[0])).all().item() + ) + if source_range: + expected_idx = global_idx.to(problem.x.device).index_select(0, source_ids) + expected_weights = global_weights.to(problem.x.device).index_select(0, source_ids) + local = (expected_idx // experts_per_rank) == rank + expected_ids = torch.where(local, expected_idx, torch.full_like(expected_idx, -1)) + expected_weights = expected_weights.masked_fill(~local, 0) + expected_payload = routing.activations_for_source_ids( + source_ids, problem.x.shape[1], seed, problem.x.dtype + ) + else: + expected_ids = torch.full_like(view.expert_ids, -1) + expected_weights = torch.zeros_like(view.weights) + expected_payload = torch.empty_like(view.payload) + actual_ids, actual_weights = _normalized_expert_metadata( + torch, view.expert_ids, view.weights + ) + expected_ids, expected_weights = _normalized_expert_metadata( + torch, expected_ids, expected_weights + ) + expected_sources = ( + ((global_idx // experts_per_rank) == rank).any(dim=1).nonzero(as_tuple=True)[0] + ).to(problem.x.device) + source_set_ok = ( + source_range + and source_ids.numel() == torch.unique(source_ids).numel() + and torch.equal(torch.sort(source_ids).values, expected_sources) + ) + payload_ok = source_range and torch.equal(view.payload, expected_payload) + metadata_ok = shape_ok and torch.equal(actual_ids, expected_ids) + max_weight_error = ( + float((actual_weights - expected_weights).abs().max().item()) + if actual_weights.numel() + else 0.0 + ) + weights_ok = max_weight_error == 0.0 + valid_expected = expected_ids >= 0 + expected_local = expected_ids[valid_expected] - rank * experts_per_rank + expected_counts = torch.bincount(expected_local, minlength=experts_per_rank) + counts_ok = torch.equal( + view.local_expert_counts.to(torch.int64), expected_counts.to(torch.int64) + ) + multiplicity_ok = torch.equal( + (actual_ids >= 0).sum(dim=1), (expected_ids >= 0).sum(dim=1) + ) + problem.recv_tokens = receive_count + combine_weight_semantics = backend.combine_weight_semantics + transformed = _expert_transform( + torch, view.payload, actual_ids, actual_weights, combine_weight_semantics + ) + view.combine_input = transformed + combined = backend.combine_transformed(problem, handle, transformed) + torch.cuda.synchronize() + expected_combined = _expected_transformed_combine(torch, problem) + if combined.shape == expected_combined.shape: + # Zero errors stand when the rank legitimately combined nothing. + max_absolute_error = max_elementwise_relative_error = max_relative_error = 0.0 + combine_values_ok = True + if combined.numel(): + absolute_error = (combined.float() - expected_combined).abs() + max_absolute_error = float(absolute_error.max().item()) + max_relative_error = max_absolute_error / ( + float(expected_combined.abs().max().item()) + 1e-6 + ) + max_elementwise_relative_error = float( + (absolute_error / expected_combined.abs().clamp_min(ORACLE_ATOL)).max().item() + ) + combine_values_ok = bool(torch.allclose( + combined.float(), expected_combined, rtol=ORACLE_RTOL, atol=ORACLE_ATOL + )) + else: + max_absolute_error = max_elementwise_relative_error = max_relative_error = None + combine_values_ok = False + checks = { + "combine_values": combine_values_ok, + "counts": counts_ok, + "metadata": metadata_ok, + "multiplicity": multiplicity_ok, + "payload": payload_ok, + "source_set": source_set_ok, + "weights": weights_ok, + } + return _oracle_report( + passed=bool( + all(checks.values()) + and max_relative_error is not None + and max_relative_error < ORACLE_RTOL + ), + combine_weight_semantics=combine_weight_semantics, + receive_count=receive_count, + max_absolute_error=max_absolute_error, + max_elementwise_relative_error=max_elementwise_relative_error, + max_relative_error=max_relative_error, + max_weight_error=max_weight_error, + checks=checks, + ) + + +def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> int: + """Drive the source-tokens-per-rank sweep for one fully-specified line.""" + mode = args.mode + if mode != "normal": + if rank == 0: + print(f"ERROR: unknown CollectiveX case mode {mode!r}") + return 2 + if min(args.iters, args.trials, args.warmup) <= 0: + if rank == 0: + print(f"ERROR: iters/trials/warmup must be positive; got " + f"{args.iters}:{args.trials}:{args.warmup}") + return 2 + import routing # torch-based; imported lazily so the module byte-compiles without torch + + ep_size = world_size + num_logical = getattr(args, "num_logical_experts", args.experts) + if args.experts % ep_size != 0: + if rank == 0: + print(f"ERROR: experts ({args.experts}) must divide ep_size ({ep_size})") + return 2 + experts_per_rank = args.experts // ep_size + gpn = args.gpus_per_node + scale_up_domain = args.scale_up_domain + suite = args.suite + workload_name = args.workload_name + if getattr(backend, "mode", None) != mode: + if rank == 0: + print(f"ERROR: backend mode {getattr(backend, 'mode', None)!r} != {mode!r}") + return 2 + expected_weight_semantics = "unweighted-rank-sum" + if getattr(backend, "combine_weight_semantics", None) != expected_weight_semantics: + if rank == 0: + print( + f"ERROR: {mode} requires combine semantics {expected_weight_semantics}" + ) + return 2 + + spec = backend.make_inputs(args) + if not spec.ok: + if rank == 0: + print(f"ERROR: {spec.message}") + return spec.rc + cap = spec.cap + ladder, dropped = spec.ladder, spec.dropped + if rank == 0 and dropped: + print(f"NOTE: dropped tokens/rank {dropped} — exceed {backend.name} buffer cap {cap} " + f"(hidden={args.hidden}); not silently truncated.") + MAX, MIN, SUM = dist.ReduceOp.MAX, dist.ReduceOp.MIN, dist.ReduceOp.SUM + + # Inputs determine the communicator capacity. + backend.create_buffer(spec) + + # ---- Pass 1: per shape, ascending (a cold-jump-safe ramp): warm untimed, + # then prove workload identity and run the expert oracle. The untimed warm + # rounds settle clocks/fabric BEFORE anything gate-bearing runs at that shape + # and are never measured or emitted. ---- + problems, gate, gts, global_traces, input_snapshots = {}, {}, {}, {}, {} + routing_consistent = True + for T in ladder: + gt = T * ep_size + gts[T] = gt + point = spec.points[T] + problem = backend.make_problem( + T, point.topk_idx.to(device), point.topk_weights.to(device), point.activations + ) + backend.warm(problem, CONDITIONING_ROUNDS_PER_SHAPE) + torch.cuda.synchronize() + problems[T] = problem + idx_g, w_g = point.global_idx, point.global_weights + rstats = routing.routing_stats(idx_g, args.experts, experts_per_rank) + rstats["locality"] = routing.routing_locality( + idx_g, experts_per_rank, ep_size, max(1, T), gpn, scale_up_domain + ) + point_routing_consistent = _same_tensors_across_ranks( + torch, dist, device, idx_g, w_g + ) + routing_consistent = routing_consistent and point_routing_consistent + input_snapshots[T] = ( + problem.x.clone(), problem.topk_idx.clone(), problem.topk_weights.clone() + ) + oracle = _run_expert_oracle( + torch, routing, backend, problem, idx_g, w_g, rank, experts_per_rank, + args.seed, + ) + before_x, before_idx, before_weights = input_snapshots[T] + pre_input_unchanged = ( + torch.equal(problem.x, before_x) + and torch.equal(problem.topk_idx, before_idx) + and torch.equal(problem.topk_weights, before_weights) + ) + global_traces[T] = (idx_g, w_g) + gate[T] = { + "rstats": rstats, + "recv_local": oracle["receive_count"], + "max_rel": oracle["max_relative_error"] or 0.0, + "local_ok": int(oracle["passed"]), + "oracle_pre": oracle, + "pre_input_unchanged": pre_input_unchanged, + } + + # ---- Pass 2: every backend uses the same rotated point order. + # Per-iteration cross-rank MAX samples are pooled across trials. ---- + disp_pool = {T: [] for T in ladder} # pooled per-iteration cross-rank MAX (dispatch) + stage_pool = {T: [] for T in ladder} # measured only when stage launches device work + comb_pool = {T: [] for T in ladder} # ... combine + rt_pool = {T: [] for T in ladder} # independently measured round trip + for trial_index in range(args.trials): + order = trial_order(list(ladder), trial_index) + for T in order: + problem = problems[T] + # timed_components() encodes the roundtrip-only vs full-component contract + # (and whether stage launches device work) once, in the base class. + component_order = trial_order(backend.timed_components(), trial_index) + measured = {name: [] for name in ("dispatch", "stage", "combine", "roundtrip")} + for component_name in component_order: + # The base template gives every component the same synchronized + # full-roundtrip warm-up before its timed trial and encodes the two + # branch rules (dispatch cleanup, combine re-dispatch) internally. + measured[component_name] = backend.benchmark_component( + component_name, problem, args.warmup, args.iters + ) + # per-iteration cross-rank MAX (the distributed-op latency per iter), pooled. + if measured["dispatch"]: + disp_pool[T] += _reduce_vec(torch, dist, device, measured["dispatch"], MAX) + comb_pool[T] += _reduce_vec(torch, dist, device, measured["combine"], MAX) + if measured["stage"]: + stage_pool[T] += _reduce_vec(torch, dist, device, measured["stage"], MAX) + rt_pool[T] += _reduce_vec(torch, dist, device, measured["roundtrip"], MAX) + + # ---- Pass 3: prove timed inputs were immutable and repeat the full oracle. ---- + for T in ladder: + problem = problems[T] + before_x, before_idx, before_weights = input_snapshots[T] + input_unchanged = gate[T]["pre_input_unchanged"] and ( + torch.equal(problem.x, before_x) + and torch.equal(problem.topk_idx, before_idx) + and torch.equal(problem.topk_weights, before_weights) + ) + idx_g, w_g = global_traces[T] + post = _run_expert_oracle( + torch, routing, backend, problem, idx_g, w_g, rank, experts_per_rank, + args.seed, + ) + pre = gate[T]["oracle_pre"] + gate[T].update({ + "input_unchanged": input_unchanged, + "local_ok": int(pre["passed"] and post["passed"] and input_unchanged), + "max_rel": max(pre["max_relative_error"] or 0.0, post["max_relative_error"] or 0.0), + "oracle_post": post, + }) + + # ---- Pass 4: percentiles (p50/p90/p95/p99, nearest-rank) from pooled samples + bytes + row ---- + rows = [] + for T in ladder: + gt = gts[T] + g = gate[T] + rstats = g["rstats"] + d, s, c, rt = disp_pool[T], stage_pool[T], comb_pool[T], rt_pool[T] + dp, sp, cp, rtp = _pcts(d), _pcts(s), _pcts(c), _pcts(rt) + # isolated_sum = SUM of the isolated dispatch+stage+combine percentiles. Stage contributes + # zero when it is explicitly not applicable. This is NOT a measured chained operation + # (can't reveal shared sync / launch amortization / overlap) — do NOT use for throughput + # or SLO capacity. The MEASURED round trip (rtp) is the real chained latency. + isum = ( + {key: dp[key] + (sp[key] if sp is not None else 0.0) + cp[key] for key in dp} + if dp and cp else None + ) + recv_total = _reduce_int(torch, dist, device, g["recv_local"], SUM) + recv_max = _reduce_int(torch, dist, device, g["recv_local"], MAX) + recv_min = _reduce_int(torch, dist, device, g["recv_local"], MIN) + global_ok = _reduce_int(torch, dist, device, g["local_ok"], MIN) + max_rel = _reduce_vec(torch, dist, device, [g["max_rel"]], MAX)[0] + point_ok = bool(global_ok) and recv_total > 0 + throughput = { + percentile_name: gt / (latency_us * 1e-6) + for percentile_name, latency_us in rtp.items() + } + # Canonical LOGICAL payload bytes come from the routing trace (NOT backend recv + # tensors): one copy per unique (token, dest-rank) pair. Dispatch and combine + # move the same logical payload; the roundtrip is their sum and stage moves nothing. + one_way_bytes = logical_byte_provenance(rstats["routed_copies"], args.hidden) + roundtrip_bytes = {field: 2 * value for field, value in one_way_bytes.items()} + stage_bytes = dict.fromkeys(one_way_bytes, 0) + rows.append({ + "components": { + "combine": _component(cp, len(c)), + "dispatch": _component(dp, len(d)), + "isolated_sum": _component(isum, 0, derived=True), + "roundtrip": _component(rtp, len(rt)), + "stage": _component(sp, len(s)), + }, + "correctness": { + "max_relative_error": max_rel, + "passed": point_ok, + }, + "global_tokens": gt, + "byte_provenance": { + "combine": one_way_bytes, + "dispatch": one_way_bytes, + "roundtrip": roundtrip_bytes, + "stage": stage_bytes, + }, + "receive": { + "max": recv_max, + "mean": recv_total / world_size, + "min": recv_min, + "total": recv_total, + }, + "routing": {key: rstats[key] for key in _ROUTING_FIELDS}, + "token_rate_at_latency_percentile": throughput, + "tokens_per_rank": T, + }) + if rank == 0: + component_log = (f"disp p50/p99={dp['p50']:7.1f}/{dp['p99']:7.1f} " + f"comb {cp['p50']:6.1f}/{cp['p99']:6.1f} " if dp and cp + else "components=unavailable ") + print(f" T={T:<5} {component_log}" + f"RT p50/p99={rtp['p50']:7.1f}/{rtp['p99']:7.1f}us n={len(rt)} fanout={rstats['fanout_mean']:.2f} " + f"recv[min/mean/max]={recv_min}/{recv_total // world_size}/{recv_max} " + f"correct={point_ok}") + + # status=valid requires correctness AND a proven-identical routing trace across ranks. + all_ok = bool(rows) and all(r["correctness"]["passed"] for r in rows) and routing_consistent + + generated_at = _dt.datetime.now().astimezone().isoformat() + nodes = int(os.environ.get("SLURM_NNODES", "1")) + scheduled_case = { + "backend": backend.name, + "ep": ep_size, + "experts": num_logical, + "gpus_per_node": gpn, + "hidden": args.hidden, + "ladder": " ".join(map(str, ladder)), + "mode": mode, + "nodes": nodes, + "phase": args.phase, + "routing": args.routing, + "scale_up_domain": scale_up_domain, + "scale_up_transport": args.scale_up_transport, + "scale_out_transport": args.scale_out_transport or None, + "scope": args.scope, + "suite": suite, + "topk": args.topk, + "topology_class": args.topology_class, + "transport": args.transport, + "workload": workload_name, + } + case_factors = {"case": scheduled_case, "sku": args.runner} + computed_case_id = case_id(args.runner, scheduled_case) + if args.case_id != computed_case_id: + raise ValueError( + f"scheduled case ID does not match realized factors: {args.case_id} != {computed_case_id}" + ) + git_run = getattr(args, "git_run", None) or {} + allocation_factors = { + "run_attempt": git_run.get("run_attempt"), + "run_id": git_run.get("run_id"), + "source_sha": git_run.get("source_sha"), + } + try: + attempt_ordinal = int(os.environ.get("COLLX_ATTEMPT_ID", "1")) + except ValueError: + attempt_ordinal = 0 + if attempt_ordinal <= 0: + raise ValueError("COLLX_ATTEMPT_ID must be a positive integer") + headline = next((r for r in rows if r["tokens_per_rank"] == 64), rows[len(rows) // 2]) + doc = { + "version": args.version, + "record_type": "case-attempt", + "generated_at": generated_at, + "identity": { + "allocation_factors": allocation_factors, + "attempt_ordinal": attempt_ordinal, + "case_factors": case_factors, + "case_id": args.case_id, + }, + "workload": { + "cross_rank_consistent": routing_consistent, + }, + "measurement": { + "combine_dtype": "bf16", + "combine_semantics": "activation-only", + "dispatch_dtype": "bf16", + "payload_unit": "token-rank", + "rows": rows, + "sampling": { + "iterations_per_trial": args.iters, + "samples_per_component": args.iters * args.trials, + "trials": args.trials, + "warmup_iterations": args.warmup, + }, + }, + "implementation": { + "kernel_generation": kernel_generation(backend), + "name": backend.name, + }, + "topology": { + "device_product": getattr(args, "runtime_device_product", None), + "gpus_per_node": gpn, + "nodes": nodes, + "placement": "packed", + "scale_up_domain": scale_up_domain, + "scale_up_transport": args.scale_up_transport, + "scale_out_transport": args.scale_out_transport or None, + "scope": args.scope, + "topology_class": args.topology_class, + "transport": args.transport, + "world_size": world_size, + }, + "runtime": getattr(args, "runtime", {}), + "provenance": { + "image": getattr(args, "image", "") or None, + "source_sha": git_run.get("source_sha"), + }, + "outcome": { + "reasons": [] if all_ok else ["semantic correctness or routing identity failed"], + "status": "success" if all_ok else "invalid", + }, + } + if rank == 0: + _write_json_atomic(args.out, doc) + dispatch_percentiles = headline["components"]["dispatch"]["percentiles_us"] + dispatch_p99 = dispatch_percentiles["p99"] if dispatch_percentiles else None + component_summary = (f"disp_p99={dispatch_p99:.1f}us " + if dispatch_p99 is not None + else "components=unavailable ") + print(f"{backend.name} ep-dispatch-combine [{args.phase}/{mode}]: " + f"status={doc['outcome']['status']} {len(rows)} pts, routing_consistent={routing_consistent}, " + f"headline T={headline['tokens_per_rank']} {component_summary}" + f"-> {args.out}") + # CI honesty: run_sweep's return code is the only success signal collx_run_shard (and thus CI) + # reads — the doc is uploaded regardless, via the launcher's always() stage step. A captured + # `invalid` outcome (semantic correctness or cross-rank routing identity failed) must therefore + # fail the leg, not ride as a green success; otherwise a persistent oracle failure is invisible + # in CI and could autopublish an invalid doc. Agree the verdict across ranks (MIN) so every + # rank exits identically and the distributed case fails as one. + outcome_ok = bool(_reduce_int(torch, dist, device, int(all_ok), dist.ReduceOp.MIN)) + return 0 if outcome_ok else 3 diff --git a/experimental/CollectiveX/bench/ep_mori.py b/experimental/CollectiveX/bench/ep_mori.py new file mode 100644 index 0000000000..188d0ee443 --- /dev/null +++ b/experimental/CollectiveX/bench/ep_mori.py @@ -0,0 +1,319 @@ +#!/usr/bin/env python3 +"""CollectiveX MoRI adapter: native BF16 dispatch/combine over mori.ops.""" +from __future__ import annotations + +import os +import re +import sys +import types + +# MoRI registers the whole symmetric heap at import time. The pinned upstream +# inter-node benchmark uses 6 GiB for its InterNodeV1 staging and signal buffers. +os.environ["MORI_SHMEM_HEAP_SIZE"] = "6G" + +import torch +import torch.distributed as dist + +import capability +from ep_backend import EPBackend + +try: + import mori # type: ignore +except Exception as exc: # pragma: no cover - requires the benchmark image + print(f"ERROR: mori import failed: {exc!r}", file=sys.stderr) + raise + + +def _project_local_metadata(torch_module, raw_expert_ids, raw_weights, rank, experts_per_rank): + local_start = rank * experts_per_rank + local = (raw_expert_ids >= local_start) & ( + raw_expert_ids < local_start + experts_per_rank + ) + expert_ids = torch_module.where( + local, raw_expert_ids, torch_module.full_like(raw_expert_ids, -1) + ) + weights = torch_module.where(local, raw_weights, torch_module.zeros_like(raw_weights)) + return expert_ids, weights, raw_expert_ids[local] - local_start + + +class MoRIBackend(EPBackend): + name = "mori" + combine_needs_redispatch = True + dispatch_needs_combine_cleanup = True + + def __init__(self, args, rank, world_size, local_rank, device): + super().__init__(args, rank, world_size, local_rank, device) + # The runner pins the AMD GPU architecture MoRI is built against + # (configs/platform_config.json). This is a hardware-lineage contract, + # independent of the (fixed BF16) communication path. + runner = str(getattr(args, "runner", "")) + platform_entry = capability.PLATFORMS.get(runner, {}) + expected_arch = ( + platform_entry.get("arch") if platform_entry.get("vendor") == "amd" else None + ) + if not expected_arch: + raise RuntimeError( + f"MoRI has no pinned GPU architecture for runner {runner!r}" + ) + + self.ep_size = world_size + self.experts_per_rank = args.experts // self.ep_size + device_properties = torch.cuda.get_device_properties(device) + realized_arch = str(getattr(device_properties, "gcnArchName", "")).split(":", 1)[0] + if realized_arch != expected_arch: + raise RuntimeError( + f"MoRI runner {runner!r} realized architecture {realized_arch!r}, " + f"expected {expected_arch!r}" + ) + # Realized placement must match the registered topology for this SKU/EP + # cell (configs/platform_config.json via capability) field-for-field. + topology = capability.topology_for(runner, world_size) + if topology is None: + raise RuntimeError(f"{runner!r} does not register an EP{world_size} topology") + gpus_per_node = int(args.gpus_per_node) + observed = { + "gpus_per_node": gpus_per_node, + "scale_up_domain": int(args.scale_up_domain), + "scope": args.scope, + "scale_up_transport": args.scale_up_transport, + "scale_out_transport": args.scale_out_transport or None, + "transport": args.transport, + } + for field, value in observed.items(): + if value != topology[field]: + raise RuntimeError( + f"MoRI {field}={value!r} differs from the registered " + f"EP{world_size} topology ({topology[field]!r})" + ) + scale_out = topology["scope"] == "scale-out" + + # The kernel is a pinned function of the cell, not an operator choice: + # scale-out EP16 runs InterNodeV1; scale-up runs the direct intranode + # kernel on gfx950 (MI355X) and the split AsyncLL send/receive kernel on + # gfx942 (MI300X/MI325X). IntraNode is mori's default (`kernel_type` + # kwarg omitted); the named kernels require their enum member — an + # image-lineage check. + # (kernel, generation label, (block_num, rdma_block_num, dispatch_warps, combine_warps)) + kernel_name, self.kernel_generation, blocks = ( + ("InterNodeV1", "inter-node-v1", (96, 64, 8, 8)) if scale_out + else ("IntraNode", "intranode", (80, 0, 16, 8)) if expected_arch == "gfx950" + else ("AsyncLL", "async-ll", (64, 0, 8, 8)) + ) + requested = os.environ.get("COLLX_MORI_KERNEL_TYPE", "") + if requested and re.sub(r"[-_]", "", requested.strip().lower()) != kernel_name.lower(): + raise RuntimeError( + f"COLLX_MORI_KERNEL_TYPE={requested!r} differs from the pinned " + f"{kernel_name} kernel for this cell" + ) + self._kernel_type = None + if kernel_name != "IntraNode": + kernel_enum = getattr(mori.ops, "EpDispatchCombineKernelType", None) + if kernel_enum is None or not hasattr(kernel_enum, kernel_name): + raise RuntimeError( + f"this MoRI image lacks EpDispatchCombineKernelType.{kernel_name}" + ) + self._kernel_type = getattr(kernel_enum, kernel_name) + self._inter_node = kernel_name == "InterNodeV1" + self._async_ll = kernel_name == "AsyncLL" + self.num_qps = 1 + self.block_num, self.rdma_block_num, self.dispatch_warps, self.combine_warps = blocks + self._external_input = self._async_ll or self._inter_node + # Registered-input MoRI copies expert output into a device-side symmetric buffer. External + # input kernels consume the dispatch output directly, so their stage is not applicable. + self.stage_device_work = not self._external_input + # Stash the __init__-only locals the moved create_buffer body reads back. + self._gpus_per_node = gpus_per_node + + def create_buffer(self, spec): + args, world_size, rank = self.args, self.world_size, self.rank + gpus_per_node = self._gpus_per_node + + world_group = torch.distributed.group.WORLD + torch._C._distributed_c10d._register_process_group("default", world_group) + mori.shmem.shmem_torch_process_group_init("default") + realized_qps = int(mori.shmem.shmem_num_qp_per_pe()) + if realized_qps < self.num_qps: + raise RuntimeError( + f"MoRI realized {realized_qps} QPs per PE; {self.num_qps} required" + ) + + self._cap = self.buffer_cap(args) + assert spec.max_tokens_per_rank <= self._cap + # BF16 communication path: no FP8 dispatch dtype, no scale payload, no combine quantization. + dispatch_dtype = torch.bfloat16 + config_kwargs = { + "data_type": dispatch_dtype, + "rank": rank, + "world_size": world_size, + "hidden_dim": args.hidden, + "scale_dim": 0, + "scale_type_size": 1, + "max_token_type_size": ( + torch.tensor([], dtype=torch.bfloat16).element_size() + if self._inter_node + else torch.tensor([], dtype=torch.float32).element_size() + ), + "max_num_inp_token_per_rank": max(512, self._cap), + "num_experts_per_rank": self.experts_per_rank, + "num_experts_per_token": args.topk, + "use_external_inp_buf": self._external_input, + "quant_type": "none", + } + if self._kernel_type is not None: + config_kwargs["kernel_type"] = self._kernel_type + if self._async_ll: + config_kwargs["max_total_recv_tokens"] = 0 + if self._async_ll or self._inter_node: + config_kwargs["block_num"] = self.block_num + config_kwargs["warp_num_per_block"] = self.dispatch_warps + if self._inter_node: + config_kwargs.update({ + "gpu_per_node": gpus_per_node, + "rdma_block_num": self.rdma_block_num, + "num_qp_per_pe": self.num_qps, + }) + self.config = mori.ops.EpDispatchCombineConfig(**config_kwargs) + expected_config = { + "data_type": dispatch_dtype, + "scale_dim": 0, + "scale_type_size": 1, + "use_external_inp_buf": self._external_input, + "quant_type": config_kwargs["quant_type"], + } + if self._async_ll or self._inter_node: + expected_config.update({ + "block_num": self.block_num, + "warp_num_per_block": self.dispatch_warps, + }) + if self._inter_node: + expected_config.update({ + "gpu_per_node": 8, + "rdma_block_num": 64, + "num_qp_per_pe": 1, + }) + if any(getattr(self.config, key, None) != value for key, value in expected_config.items()): + raise RuntimeError("MoRI requested launch/topology configuration was not realized") + # The newer pinned MoRI revision can otherwise replace explicit values + # with token-dependent tuning rules from the image. + os.environ["MORI_EP_LAUNCH_CONFIG_MODE"] = "MANUAL" + self.op = mori.ops.EpDispatchCombineOp(self.config) + if getattr(self.op, "launch_config_mode", None) != "MANUAL": + raise RuntimeError("MoRI explicit launch configuration was not applied") + + def buffer_cap(self, args): + # Rows reserved per rank for a uniformly routed dispatch. + return 512 + + def make_problem(self, T, idx, weights, x): + indices = idx.to(torch.int32) + gate_weights = weights.to(torch.float32) + return types.SimpleNamespace( + T=T, + x=x, + dispatch_x=x, + topk_idx=indices, + topk_weights=gate_weights, + indices=indices, + weights=gate_weights, + scales=torch.empty((T, 0), dtype=torch.uint8, device=self.device), + ) + + def dispatch(self, p): + dispatch_output, dispatch_weights, _scales, dispatch_indices, recv_num = ( + self.op.dispatch( + p.dispatch_x, + p.weights, + p.scales, + p.indices, + block_num=self.block_num, + rdma_block_num=self.rdma_block_num, + warp_per_block=self.dispatch_warps, + ) + ) + if self._async_ll: + self.op.dispatch_recv(warp_per_block=self.dispatch_warps) + return types.SimpleNamespace( + dispatch_output=dispatch_output, + dispatch_weights=dispatch_weights, + dispatch_indices=dispatch_indices, + recv_num=recv_num[0], + combine_input=None, + ) + + def stage(self, p, h): + rows = getattr(p, "recv_tokens", None) + if not isinstance(rows, int) or rows < 0 or rows > h.dispatch_output.size(0): + raise RuntimeError("MoRI receive count was not validated before staging") + h.combine_input = h.dispatch_output + if self._external_input: + return None + buffer = self.op.get_registered_combine_input_buffer( + torch.bfloat16, hidden_dim=h.combine_input.size(1) + ) + buffer[:rows, :].copy_(h.combine_input[:rows, :]) + h.combine_input = buffer + + def combine(self, p, h): + combine_indices = p.indices if self._async_ll else h.dispatch_indices + combined, _weights = self.op.combine( + h.combine_input, + None, + combine_indices, + block_num=self.block_num, + rdma_block_num=self.rdma_block_num, + warp_per_block=self.combine_warps, + ) + if self._async_ll: + self.op.combine_recv(warp_per_block=self.combine_warps) + return combined[:p.T] + + def inspect_dispatch(self, p, h): + count = self.recv_tokens(h) + if h.dispatch_weights is None: + raise RuntimeError("MoRI dispatch did not expose gate weights") + if count < 0 or any( + tensor.ndim == 0 or count > tensor.size(0) + for tensor in (h.dispatch_output, h.dispatch_indices, h.dispatch_weights) + ): + raise RuntimeError("MoRI receive count exceeds dispatch metadata") + raw_expert_ids = h.dispatch_indices[:count].to(torch.int64) + expert_ids, weights, local_expert_ids = _project_local_metadata( + torch, + raw_expert_ids, + h.dispatch_weights[:count].to(torch.float32), + self.rank, + self.experts_per_rank, + ) + return types.SimpleNamespace( + payload=h.dispatch_output[:count], + expert_ids=expert_ids, + weights=weights, + local_expert_counts=torch.bincount( + local_expert_ids, minlength=self.experts_per_rank + ), + ) + + def combine_transformed(self, p, h, transformed): + h.combine_input = transformed.to(torch.bfloat16) + rows = getattr(p, "recv_tokens", None) + if not isinstance(rows, int) or rows < 0 or rows > h.combine_input.size(0): + raise RuntimeError("MoRI receive count was not validated before transformed combine") + if not self._external_input: + buffer = self.op.get_registered_combine_input_buffer( + torch.bfloat16, hidden_dim=h.combine_input.size(1) + ) + buffer[:rows, :].copy_(h.combine_input[:rows, :]) + h.combine_input = buffer + return self.combine(p, h) + + def recv_tokens(self, h): + return int(h.recv_num.item()) + + def finalize(self, rc): + try: + dist.barrier() + except Exception: + pass + sys.stdout.flush() + sys.stderr.flush() + os._exit(rc if 0 <= rc <= 255 else 1) diff --git a/experimental/CollectiveX/bench/routing.py b/experimental/CollectiveX/bench/routing.py new file mode 100644 index 0000000000..b61d2be160 --- /dev/null +++ b/experimental/CollectiveX/bench/routing.py @@ -0,0 +1,173 @@ +#!/usr/bin/env python3 +"""Deterministic platform-independent MoE routing and activations.""" +from __future__ import annotations + +import hashlib +import struct + +import torch + +_MASK64 = (1 << 64) - 1 + +SOURCE_ID_BITS = 32 +SOURCE_ID_COLUMNS = SOURCE_ID_BITS + + +def build_global_routing(global_tokens: int, experts: int, topk: int, routing: str, seed: int): + """Return one byte-stable counter-generated routing window on CPU.""" + if routing != "uniform": + raise ValueError(f"unknown routing {routing!r} (uniform)") + if global_tokens <= 0 or experts <= 0 or topk <= 0 or topk > experts: + raise ValueError("global_tokens/experts/topk must be positive and topk <= experts") + + # Keyed BLAKE2b as a deterministic random oracle over the coordinate tuple: + # stdlib, byte-stable on every runtime, and uniform by construction. + key = (int(seed) & _MASK64).to_bytes(8, "little") + + def counter(token: int, slot: int, attempt: int, stream: int) -> int: + message = struct.pack("<4Q", token, slot, attempt, stream) + return int.from_bytes( + hashlib.blake2b(message, key=key, digest_size=8).digest(), "little" + ) + + indices, weights = [], [] + for token in range(int(global_tokens)): + selected, used = [], set() + for slot in range(int(topk)): + attempt = 0 + while True: + expert = counter(token, slot, attempt, 0) % int(experts) + if expert not in used: + used.add(expert) + selected.append(expert) + break + attempt += 1 + raw = [1 + counter(token, slot, 0, 1) % 65535 for slot in range(int(topk))] + denominator = float(sum(raw)) + indices.append(selected) + weights.append([value / denominator for value in raw]) + return ( + torch.tensor(indices, dtype=torch.int64), + torch.tensor(weights, dtype=torch.float32), + ) + + +def rank_slice(idx, weights, rank: int, tokens_per_rank: int): + lo = rank * tokens_per_rank + return idx[lo:lo + tokens_per_rank].contiguous(), weights[lo:lo + tokens_per_rank].contiguous() + + +def rank_activations(tokens: int, hidden: int, seed: int, rank: int, device, + dtype=torch.bfloat16): + """Exact counter-derived inputs with a quantization-safe source-token prefix.""" + source = torch.arange(tokens, device=device, dtype=torch.int64) + rank * tokens + return activations_for_source_ids(source, hidden, seed, dtype) + + +def activations_for_source_ids(source, hidden: int, seed: int, dtype=torch.bfloat16): + """Materialize canonical activations for arbitrary global source-token IDs.""" + if hidden < SOURCE_ID_COLUMNS: + raise ValueError(f"hidden must be at least {SOURCE_ID_COLUMNS}") + source = source.to(torch.int64) + column = torch.arange(hidden, device=source.device, dtype=torch.int64) + # Integer lattice, frozen (the oracle regenerates these exact bytes): + # 131/17/19 are odd multipliers coprime to the prime modulus 257, so distinct (source, column, seed) + # land on distinct residues and corruption cannot alias into a correct-looking pattern; %257-128 yields k in + # [-128, 128] and k/64 is exactly representable in bfloat16, so the oracle's bit-exact payload compare sees + # transport corruption, never representation error. + values = (source[:, None] * 131 + column[None, :] * 17 + int(seed) * 19) % 257 - 128 + output = values.to(dtype).mul_(1 / 64) + if bool((source < 0).any().item()) or bool((source >= (1 << SOURCE_ID_BITS)).any().item()): + raise ValueError("source token ID is outside the bounded identity contract") + source_columns = torch.arange(SOURCE_ID_BITS, device=source.device, dtype=torch.int64) + source_bits = ((source[:, None] >> source_columns[None, :]) & 1) * 2 - 1 + # Magnitude one sits inside the ordinary [-2, 2] activation range, so the identity cannot set + # an FP8 block scale. Decode depends only on sign and remains stable after dequantization. + output[:, :SOURCE_ID_BITS] = source_bits.to(dtype) + return output + + +def decode_source_ids(payload, seed: int): + """Decode and validate source IDs carried by rank_activations.""" + if payload.ndim != 2 or payload.shape[1] < SOURCE_ID_COLUMNS: + raise ValueError("received payload cannot carry the source-token prefix") + prefix = payload[:, :SOURCE_ID_COLUMNS].float() + if not bool(torch.isfinite(prefix).all().item()) or bool((prefix.abs() < 0.25).any().item()): + raise ValueError("received source-token prefix is not quantization-stable") + bits = prefix >= 0 + powers = 1 << torch.arange(SOURCE_ID_BITS, device=payload.device, dtype=torch.int64) + source = (bits[:, :SOURCE_ID_BITS].to(torch.int64) * powers).sum(dim=1) + return source + + +def routing_locality(idx, experts_per_rank: int, ep_size: int, tokens_per_rank: int, + gpus_per_node: int, scale_up_domain: int = None) -> dict: + """Locality of rank-deduplicated payload copies under packed placement.""" + gt = idx.shape[0] + assignments = (idx // experts_per_rank).clamp(max=ep_size - 1) + destinations = torch.zeros((gt, ep_size), dtype=torch.bool) + destinations.scatter_(1, assignments, True) + token, dest = destinations.nonzero(as_tuple=True) + src = (token // max(1, tokens_per_rank)).clamp(max=ep_size - 1) + sud = scale_up_domain or (gpus_per_node * ep_size) # default: all one domain + phys = torch.arange(ep_size, dtype=torch.int64) + pd, ps = phys[dest], phys[src] + local = (dest == src) + same_node = (pd // gpus_per_node) == (ps // gpus_per_node) + same_dom = (pd // sud) == (ps // sud) + n = dest.numel() + return { + "placement": "packed", + "local_rank_fraction": float(local.float().mean()), + "same_node_fraction": float(same_node.float().mean()), + "same_scaleup_domain_fraction": float(same_dom.float().mean()), + "cross_node_fraction": float((~same_node).float().mean()), + "cross_domain_fraction": float((~same_dom).float().mean()), + "gpus_per_node": gpus_per_node, "scale_up_domain": sud, "copies": int(n), + } + + +def routing_stats(idx, experts: int, experts_per_rank: int) -> dict: + """Realized routing properties for the GLOBAL trace — published per point so the + fan-out / load can never be silently misread. idx is the global [gt, topk] tensor; + """ + ep = max(1, experts // max(1, experts_per_rank)) + ranks = (idx // experts_per_rank) # [gt, topk] destination rank per assignment + # unique destination ranks per token (fan-out) + onehot = torch.zeros(idx.shape[0], ep, dtype=torch.bool) + onehot.scatter_(1, ranks.clamp(max=ep - 1), True) + fanout = onehot.sum(dim=1) # [gt] + hist = torch.bincount(fanout, minlength=ep + 1)[1:ep + 1].tolist() # counts for fan-out 1..ep + load = torch.bincount(idx.reshape(-1), minlength=experts).float() + # Keep expert assignments (compute load) separate from rank-deduplicated payload copies + # (network load). Conflating them overstates traffic when two experts share a rank. + assignment_load = torch.bincount( + ranks.reshape(-1).clamp(max=ep - 1), minlength=ep + ).float() + payload_load = onehot.sum(dim=0).float() + # One-number imbalance summaries so a row is self-describing for the distribution-sensitivity + # suite (no need to read the full histograms): CV = std/mean of the load; hotspot_ratio = + # worst expert load over the mean. + def _cv(t): + m = float(t.mean()) + return float(t.std(unbiased=False) / m) if m > 0 else 0.0 + expert_load_cv = _cv(load) + assignment_rank_cv = _cv(assignment_load) + payload_rank_cv = _cv(payload_load) + hotspot_ratio = float(load.max() / load.mean()) if float(load.mean()) > 0 else 0.0 + # Empty experts capture compute skew; empty destination ranks capture network skew. + empty_expert_count = int((load == 0).sum()) + empty_rank_count = int((payload_load == 0).sum()) + return { + "fanout_mean": float(fanout.float().mean()), + "fanout_min": int(fanout.min()), "fanout_max": int(fanout.max()), + "fanout_histogram": hist, # index k-1 = #tokens with fan-out k + "expert_assignments_per_rank": [int(x) for x in assignment_load.tolist()], + "payload_copies_per_rank": [int(x) for x in payload_load.tolist()], + "routed_copies": int(fanout.sum()), # total (token, dest-rank) pairs + "expert_load_min": int(load.min()), "expert_load_max": int(load.max()), + "expert_load_mean": float(load.mean()), "expert_load_cv": expert_load_cv, + "expert_assignment_rank_cv": assignment_rank_cv, + "payload_rank_cv": payload_rank_cv, "hotspot_ratio": hotspot_ratio, + "empty_expert_count": empty_expert_count, "empty_rank_count": empty_rank_count, + } diff --git a/experimental/CollectiveX/bench/run_ep.py b/experimental/CollectiveX/bench/run_ep.py new file mode 100644 index 0000000000..8672f36523 --- /dev/null +++ b/experimental/CollectiveX/bench/run_ep.py @@ -0,0 +1,203 @@ +#!/usr/bin/env python3 +"""CollectiveX v1 EP benchmark entrypoint for torchrun or rank environments.""" + +from __future__ import annotations + +import argparse +import ctypes +import os +import platform +import sys + +# Make the sibling bench/ modules importable when run as `bench/run_ep.py` under +# torchrun (it executes the file as __main__, not as a package). +HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path[:0] = [HERE, os.path.dirname(HERE)] + +import ep_harness # noqa: E402 (stdlib-only; safe before torch) + + +def _loaded_collective_version() -> str | None: + try: + with open("/proc/self/maps", encoding="utf-8") as handle: + paths = { + os.path.realpath(line.rstrip().split()[-1]) + for line in handle + if any(name in line for name in ("libnccl.so", "librccl.so")) + and os.path.isfile(line.rstrip().split()[-1]) + } + if len(paths) != 1: + return None + version = ctypes.c_int() + library = ctypes.CDLL(paths.pop()) + if library.ncclGetVersion(ctypes.byref(version)) != 0: + return None + return ep_harness.format_collective_version(version.value) + except (AttributeError, OSError): + return None + + +def _runtime_info(torch, *, vendor: str) -> dict: + """Return the runtime versions needed to compare and debug results.""" + runtime_kind = "cuda" if vendor == "nvidia" else "hip" + collective_kind = "nccl" if vendor == "nvidia" else "rccl" + return { + "accelerator_runtime": getattr(torch.version, runtime_kind), + "collective_library": {"kind": collective_kind, "version": _loaded_collective_version()}, + "framework": str(torch.__version__), + "vendor": vendor, + } + + +def main() -> int: + ap = argparse.ArgumentParser(description="CollectiveX EP dispatch/combine sweep") + ap.add_argument("--backend", required=True, choices=["deepep-v2", "mori"]) + ep_harness.add_common_args(ap) + args = ap.parse_args() + + if not ep_harness.is_case_id(args.case_id): + print(f"ERROR: invalid native case ID {args.case_id!r}", file=sys.stderr) + return 2 + # Seed and timing arrive baked into the case argv from the single + # configs/suites.yaml source; there is no separate canonical constant to + # cross-check against. + + try: + import torch + import torch.distributed as dist + except Exception as exc: # pragma: no cover + print(f"ERROR: torch unavailable: {exc!r}", file=sys.stderr) + return 3 + + rank = int(os.environ.get("RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + torch.cuda.set_device(local_rank) + device = torch.device(f"cuda:{local_rank}") + + import capability + + machine = {"x86_64": "amd64", "aarch64": "arm64"}.get( + platform.machine(), platform.machine() + ) + props = torch.cuda.get_device_properties(device) + if torch.version.hip: + vendor = "amd" + accelerator = str(getattr(props, "gcnArchName", "")).split(":", 1)[0] + else: + vendor = "nvidia" + major, minor = torch.cuda.get_device_capability(device) + accelerator = f"sm{major}{minor}" + device_name = torch.cuda.get_device_name(device) + device_count = torch.cuda.device_count() + identity_issues = capability.runtime_identity_issues( + args.runner, + vendor=vendor, + arch=accelerator, + machine=machine, + device_name=device_name, + device_count=device_count, + world_size=world_size, + ) + if identity_issues: + print( + f"ERROR: runtime identity does not match {args.runner}: " + + "; ".join(identity_issues), + file=sys.stderr, + ) + return 5 + observed_gpus_per_node = args.gpus_per_node + observed_nodes = world_size // observed_gpus_per_node + topology = capability.topology_for(args.runner, world_size) + observed_topology = { + "nodes": observed_nodes, + "gpus_per_node": observed_gpus_per_node, + "scale_up_domain": args.scale_up_domain, + "scope": args.scope, + "scale_up_transport": args.scale_up_transport, + "scale_out_transport": args.scale_out_transport or None, + "transport": args.transport, + "topology_class": args.topology_class, + } + if topology is None or any( + observed_topology[field] != topology[field] for field in observed_topology + ): + print( + f"ERROR: runtime topology does not match {args.runner} EP{world_size}", + file=sys.stderr, + ) + return 5 + schedulable, reason = capability.resolve( + args.runner, + args.backend, + ep=world_size, + nodes=observed_nodes, + routing=args.routing, + mode=args.mode, + ) + if not schedulable: + print(f"ERROR: scheduled case is unsupported: {reason}", file=sys.stderr) + return 5 + args.runtime_device_product = device_name + args.image = os.environ.get("COLLECTIVEX_IMAGE", "") + _run = { + "run_id": os.environ.get("GITHUB_RUN_ID"), + "run_attempt": os.environ.get("GITHUB_RUN_ATTEMPT"), + "source_sha": os.environ.get("COLLECTIVEX_SOURCE_SHA") + or os.environ.get("GITHUB_SHA"), + } + args.git_run = _run if any(_run.values()) else None + + # Import the backend class only after torch initializes. The selected mode is an + # explicit case dimension; adapters do not infer it from the token ladder. + if args.backend == "mori": + from ep_mori import MoRIBackend as Backend + else: + from ep_deepep_v2 import DeepEPV2Backend as Backend + + # MoRI registers the default GPU process group with its SHMEM runtime. Keep that + # group device-only so scale-out does not also depend on a host Gloo fabric. + if not dist.is_initialized(): + if args.backend == "mori": + dist.init_process_group( + backend="nccl", + rank=rank, + world_size=world_size, + device_id=device, + ) + else: + # PR #605 reuses PyTorch's NCCL communicator through ``_comm_ptr``. Supplying + # device_id eagerly forms it before ElasticBuffer construction. + dist.init_process_group("nccl", device_id=device) + + args.runtime = _runtime_info(torch, vendor=vendor) + + # Construct + run inside a try so a backend exception (esp. a new adapter on GPU) prints its + # FULL traceback to STDOUT — torchrun captures per-rank stdout but only summarizes stderr, so an + # uncaught exception is otherwise invisible in CI. Print on every rank (prefixed) then re-raise. + try: + backend = Backend(args, rank, world_size, local_rank, device) + if rank == 0: + print( + f"[run_ep] backend={args.backend} phase={args.phase} mode={args.mode} " + f"world={world_size} ep_size={world_size} hidden={args.hidden} " + f"topk={args.topk} experts={args.experts} dtype=bf16 " + f"routing={args.routing} seed={args.seed}" + ) + rc = ep_harness.run_sweep(args, backend, torch, dist, device, rank, world_size) + except Exception: + import traceback + + print( + f"[run_ep][rank{rank}] backend={args.backend} FAILED:\n" + + traceback.format_exc(), + flush=True, + ) + raise + # finalize() handles backend-specific teardown: DeepEP returns rc cleanly; + # MoRI hard-exits past its post-shmem_finalize teardown assertion. + return backend.finalize(rc) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/experimental/CollectiveX/capability.py b/experimental/CollectiveX/capability.py new file mode 100644 index 0000000000..fc90cbcdfa --- /dev/null +++ b/experimental/CollectiveX/capability.py @@ -0,0 +1,248 @@ +#!/usr/bin/env python3 +"""Public runner and backend capability registry for CollectiveX. + +Per-SKU platform identity lives in configs/platform_config.json; this module +derives the EP topologies from it and holds backend capability.""" + +from __future__ import annotations + +import json +import re +from pathlib import Path +from typing import Any + + +DEEPEP_V2_SKU_CAPABILITIES = { + "h100-dgxc": { + "schedulable": True, + "basis": "current-runner-ep8-lsa-elasticbuffer-validated", + }, + "h200-dgxc": {"schedulable": True, "basis": "upstream-sm90-requirement"}, + "b200-dgxc": {"schedulable": True, "basis": "upstream-sm100-result"}, + "gb200": {"schedulable": True, "basis": "upstream-sm100-result"}, + "b300": {"schedulable": True, "basis": "pinned-pr605-pr630-sm103-maps-sm100f"}, + "gb300": {"schedulable": True, "basis": "pinned-pr605-pr630-sm103-maps-sm100f"}, + "mi300x": {"schedulable": False, "basis": "nvidia-only"}, + "mi325x": {"schedulable": False, "basis": "nvidia-only"}, + "mi355x": {"schedulable": False, "basis": "nvidia-only"}, +} + + +def _topologies( + product: str, *, gpus_per_node: int, scale_up_domain: int, scale_up_transport: str +) -> dict[int, dict[str, Any]]: + scale_up_class = ( + f"{product}-nvl72-mnnvl" + if scale_up_transport == "mnnvl" + else f"{product}-xgmi" + if scale_up_transport == "xgmi" + else f"{product}-{scale_up_transport}-island" + ) + return { + 8: { + "nodes": 8 // gpus_per_node, + "gpus_per_node": gpus_per_node, + "scale_up_domain": scale_up_domain, + "scope": "scale-up", + "scale_up_transport": scale_up_transport, + "scale_out_transport": None, + "transport": scale_up_transport, + "topology_class": scale_up_class, + }, + 16: { + "nodes": 16 // gpus_per_node, + "gpus_per_node": gpus_per_node, + "scale_up_domain": scale_up_domain, + "scope": "scale-up" if scale_up_domain >= 16 else "scale-out", + "scale_up_transport": scale_up_transport, + "scale_out_transport": None if scale_up_domain >= 16 else "rdma", + "transport": ( + scale_up_transport + if scale_up_domain >= 16 + else f"{scale_up_transport}-rdma" + ), + "topology_class": ( + scale_up_class + if scale_up_domain >= 16 + else f"{product}-{scale_up_transport}-rdma" + ), + }, + } + + +def _platform( + *, vendor: str, arch: str, machine: str, product: str, gpus_per_node: int, + scale_up_domain: int, scale_up_transport: str, launcher: str, +) -> dict[str, Any]: + topologies = _topologies( + product, + gpus_per_node=gpus_per_node, + scale_up_domain=scale_up_domain, + scale_up_transport=scale_up_transport, + ) + return { + "vendor": vendor, + "arch": arch, + "machine": machine, + "product": product, + "gpus_per_node": gpus_per_node, + "scale_up_domain": scale_up_domain, + "ep_degrees": tuple(topologies), + "topologies": topologies, + "launcher": launcher, + } + + +def _load_platforms() -> dict[str, dict[str, Any]]: + """Build the registry from configs/platform_config.json — fails loudly at + import on a missing file, missing field, or wrong-typed value.""" + path = Path(__file__).resolve().parent / "configs" / "platform_config.json" + with path.open(encoding="utf-8") as stream: + document = json.load(stream) + platforms: dict[str, dict[str, Any]] = {} + for name, entry in document["platforms"].items(): + identity = { + field: entry[field] + for field in ( + "vendor", "arch", "machine", "product", "scale_up_transport", "launcher" + ) + } + placement = { + field: entry[field] + for field in ("gpus_per_node", "scale_up_domain") + } + if not all(isinstance(value, str) and value for value in identity.values()): + raise ValueError(f"platform {name!r} has a non-string identity field") + if not all(isinstance(value, int) and value > 0 for value in placement.values()): + raise ValueError(f"platform {name!r} has a non-positive placement field") + platforms[name] = _platform(**identity, **placement) + return platforms + + +PLATFORMS = _load_platforms() + +# Source pins and images live in configs/backends.json; this registry holds +# only what scheduling reads: vendor and per-SKU capability. +BACKENDS = { + "deepep-v2": { + "vendors": {"nvidia"}, + "sku_capabilities": DEEPEP_V2_SKU_CAPABILITIES, + }, + "mori": {"vendors": {"amd"}}, +} +SWEEP_BACKENDS = tuple(BACKENDS) + +# Backend-specific topology limits require repeated native execution evidence. +# Keep these narrower than platform overrides so working reference paths remain +# measurable on the same fabric. +BACKEND_TOPOLOGY_CELL_OVERRIDES: dict[tuple[str, str, int], str] = { + ("h100-dgxc", "deepep-v2", 16): ( + "DeepEP V2 EP16 requires NCCL GIN over the H100 scale-out fabric, " + "unverified on current runners; EP8 (LSA) validated on-node 2026-07-11 — " + "scheduled EP8 only for now" + ), + ("b300", "deepep-v2", 16): ( + "DeepEP V2 EP16 requires GDRCopy /dev/gdrdrv for NVSHMEM-IBGDA, " + "unprovisioned on B300 hosts" + ), + ("mi300x", "mori", 16): ( + "Pinned MoRI distributed initialization does not complete on MI300X EP16" + ), + ("mi325x", "mori", 16): ( + "MoRI InterNodeV1 EP16 unvalidated on MI325X (gfx942) — pending a 2-node " + "internode run plus the MI325X internode RDMA selectors in the network " + "config; scheduled EP8 only for now" + ), +} + + +def runtime_identity_issues( + sku: str, *, vendor: str, arch: str, machine: str, device_name: str, + device_count: int, world_size: int, +) -> list[str]: + """Validate public product identity on every rank without private device identifiers.""" + platform = PLATFORMS.get(sku) + if platform is None: + return [f"unknown runner identity {sku!r}"] + issues = [] + for field, observed in (("vendor", vendor), ("arch", arch), ("machine", machine)): + if observed != platform[field]: + issues.append(f"{field}={observed!r}, expected {platform[field]!r}") + products = set(re.findall(r"[a-z]+\d+[a-z]*", device_name.lower())) + if platform["product"] not in products: + issues.append(f"device product {device_name!r} does not identify {platform['product']}") + if device_count != platform["gpus_per_node"]: + issues.append( + f"visible GPUs={device_count}, expected {platform['gpus_per_node']} per node" + ) + if world_size not in platform["ep_degrees"]: + issues.append(f"EP{world_size} is not registered for {sku}") + return issues + + +def topology_for(sku: str, ep: int) -> dict[str, Any] | None: + """Return the exact public topology registered for one SKU/EP cell.""" + platform = PLATFORMS.get(sku) + if platform is None: + return None + return platform["topologies"].get(ep) + + +def resolve( + sku: str, + backend: str, + *, + ep: int | None = None, + nodes: int | None = None, + routing: str = "uniform", + mode: str = "normal", +) -> tuple[bool, str]: + """Return whether one fixed-profile case can run on a public GHA runner label.""" + platform, implementation = PLATFORMS.get(sku), BACKENDS.get(backend) + if platform is None: + return False, f"unknown GHA runner label {sku!r}" + if implementation is None: + return False, f"unknown backend {backend!r}" + if mode not in {"normal"}: + return False, f"unknown benchmark mode {mode!r}" + if ep is None: + if nodes is None: + ep = platform["ep_degrees"][0] + else: + matches = [ + degree for degree, topology in platform["topologies"].items() + if topology["nodes"] == nodes + ] + if len(matches) != 1: + return False, f"{sku} does not register a unique {nodes}-node EP degree" + ep = matches[0] + topology = topology_for(sku, ep) + if topology is None or (nodes is not None and nodes != topology["nodes"]): + return False, f"{sku} does not register EP{ep} on {nodes} nodes" + if routing != "uniform": + return False, "core routing is uniform" + if platform["vendor"] not in implementation["vendors"]: + return False, f"{backend} does not support {platform['vendor']}" + sku_capability = implementation.get("sku_capabilities", {}).get(sku) + if sku_capability is not None and not sku_capability["schedulable"]: + return False, f"{backend} is unsupported on {sku}: {sku_capability['basis']}" + backend_topology_override = BACKEND_TOPOLOGY_CELL_OVERRIDES.get( + (sku, backend, ep) + ) + if backend_topology_override is not None: + return False, backend_topology_override + return True, "ok" + + +def resolve_disposition( + sku: str, + backend: str, + *, + ep: int | None = None, + nodes: int | None = None, + routing: str = "uniform", + mode: str = "normal", +) -> tuple[str, str]: + """Resolve a BF16 cell to its capability disposition.""" + ok, detail = resolve(sku, backend, ep=ep, nodes=nodes, routing=routing, mode=mode) + return ("supported", "ok") if ok else ("unsupported", detail) diff --git a/experimental/CollectiveX/configs/backends.json b/experimental/CollectiveX/configs/backends.json new file mode 100644 index 0000000000..7a500a8321 --- /dev/null +++ b/experimental/CollectiveX/configs/backends.json @@ -0,0 +1,16 @@ +{ + "images": { + "multiarch": { + "ref": "lmsysorg/sglang:v0.5.11-cu130" + }, + "amd-mori": { + "ref": "rocm/sgl-dev:sglang-0.5.14-rocm720-mi35x-mori-0701" + } + }, + "backends": { + "deepep-v2": { + "repo": "https://github.com/deepseek-ai/DeepEP", + "commit": "fa8a9b16898204afd347c663b89e65ef87dc6ce6" + } + } +} diff --git a/experimental/CollectiveX/configs/platform_config.json b/experimental/CollectiveX/configs/platform_config.json new file mode 100644 index 0000000000..7fb40ab98b --- /dev/null +++ b/experimental/CollectiveX/configs/platform_config.json @@ -0,0 +1,193 @@ +{ + "platforms": { + "h100-dgxc": { + "vendor": "nvidia", + "arch": "sm90", + "machine": "amd64", + "product": "h100", + "gpus_per_node": 8, + "scale_up_domain": 8, + "scale_up_transport": "nvlink", + "launcher": "single-slurm", + "operator_fields": [ + "partition", + "account", + "squash_dir" + ], + "operator": { + "partition": "hpc-gpu-1", + "account": "customer", + "squash_dir": "/mnt/nfs/sa-shared/cx-squash", + "exclude_nodes": "hpc-gpu-1-7" + } + }, + "h200-dgxc": { + "vendor": "nvidia", + "arch": "sm90", + "machine": "amd64", + "product": "h200", + "gpus_per_node": 8, + "scale_up_domain": 8, + "scale_up_transport": "nvlink", + "launcher": "single-slurm", + "operator_fields": [ + "partition", + "squash_dir" + ], + "operator": { + "partition": "main", + "squash_dir": "/home/sa-shared/containers" + }, + "network": { + "rdma_devices": "mlx5_0,mlx5_1,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_6,mlx5_7" + } + }, + "b200-dgxc": { + "vendor": "nvidia", + "arch": "sm100", + "machine": "amd64", + "product": "b200", + "gpus_per_node": 8, + "scale_up_domain": 8, + "scale_up_transport": "nvlink", + "launcher": "single-slurm", + "operator_fields": [ + "partition", + "account", + "squash_dir" + ], + "operator": { + "partition": "gpu-2", + "account": "benchmark", + "qos": "gpu-2_qos", + "squash_dir": "/home/sa-shared/containers" + }, + "network": { + "rdma_devices": "mlx5_0,mlx5_1,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_6,mlx5_7", + "ib_gid_index": "3" + } + }, + "b300": { + "vendor": "nvidia", + "arch": "sm103", + "machine": "amd64", + "product": "b300", + "gpus_per_node": 8, + "scale_up_domain": 8, + "scale_up_transport": "nvlink", + "launcher": "single-slurm", + "operator_fields": [ + "partition", + "account", + "squash_dir" + ], + "operator": { + "partition": "batch_1", + "account": "benchmark", + "qos": "batch_1_qos", + "squash_dir": "/data/home/sa-shared/sqsh" + }, + "network": { + "socket_ifname": "bond0", + "rdma_devices": "mlx5_0,mlx5_1,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_8,mlx5_9,mlx5_10,mlx5_11,mlx5_16,mlx5_17,mlx5_20,mlx5_21,mlx5_22,mlx5_23", + "ib_gid_index": "3" + } + }, + "gb200": { + "vendor": "nvidia", + "arch": "sm100", + "machine": "arm64", + "product": "gb200", + "gpus_per_node": 4, + "scale_up_domain": 72, + "scale_up_transport": "mnnvl", + "launcher": "gb-nv", + "operator_fields": [ + "partition", + "account", + "storage_roots" + ], + "operator": { + "partition": "batch", + "account": "benchmark", + "storage_roots": [ + "/mnt/lustre01/users-public/sa-shared" + ] + } + }, + "gb300": { + "vendor": "nvidia", + "arch": "sm103", + "machine": "arm64", + "product": "gb300", + "gpus_per_node": 4, + "scale_up_domain": 72, + "scale_up_transport": "mnnvl", + "launcher": "gb-nv", + "operator_fields": [ + "partition", + "account", + "squash_dir", + "enroot_cache_path" + ], + "operator": { + "partition": "batch_1", + "account": "benchmark", + "qos": "batch_1_qos", + "squash_dir": "/data/home/sa-shared/collectivex/containers", + "enroot_cache_path": "/data/home/sa-shared/collectivex/enroot-cache" + } + }, + "mi300x": { + "vendor": "amd", + "arch": "gfx942", + "machine": "amd64", + "product": "mi300x", + "gpus_per_node": 8, + "scale_up_domain": 8, + "scale_up_transport": "xgmi", + "launcher": "mi-amds", + "operator_fields": [ + "partition", + "squash_dir" + ] + }, + "mi325x": { + "vendor": "amd", + "arch": "gfx942", + "machine": "amd64", + "product": "mi325x", + "gpus_per_node": 8, + "scale_up_domain": 8, + "scale_up_transport": "xgmi", + "launcher": "mi-amds", + "operator_fields": [ + "partition", + "squash_dir" + ] + }, + "mi355x": { + "vendor": "amd", + "arch": "gfx950", + "machine": "amd64", + "product": "mi355x", + "gpus_per_node": 8, + "scale_up_domain": 8, + "scale_up_transport": "xgmi", + "launcher": "mi-amds", + "network": { + "socket_ifname": "eno0", + "rdma_devices": "rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7", + "ib_gid_index": 1 + }, + "operator_fields": [ + "partition", + "squash_dir" + ], + "operator": { + "partition": "compute", + "squash_dir": "/var/lib/squash" + } + } + } +} diff --git a/experimental/CollectiveX/configs/suites.yaml b/experimental/CollectiveX/configs/suites.yaml new file mode 100644 index 0000000000..6cd3167c2f --- /dev/null +++ b/experimental/CollectiveX/configs/suites.yaml @@ -0,0 +1,44 @@ +# CollectiveX EP comparison suites. +# Iterable benchmark version — the single authoritative CollectiveX version. Bump +# when a methodology change makes new results incomparable with prior ones; it is +# copied verbatim into the matrix, every emitted case result, every extracted shard +# control, and the run summary so the frontend can group and pick runs by it. This +# is the only place it is set. +version: 1 + +# One literal timing profile for every comparison-grade point on every SKU/backend; +# the matrix bakes it into each case, so changing it changes comparability — bump +# `version` with it. Eight timed iterations keep each MoRI burst well below its +# sustained-iteration wedge; 128 trials give 1024 observations per component +# (tighter p99 tail — the estimate that drives the chart); 32 warmups meet +# Blackwell's measured clock-ramp floor (B300 dispatch reads ~1787us cold at +# warmup=8, settles to ~85us at warmup>=30; H100/MI355X settle sooner — the extra +# iters are harmless). Trials re-warm every time, so trial count dominates wall-clock. +timing: + iters_per_trial: 8 + trials_per_point: 128 + warmup_iters_per_trial: 32 + +# Canonical workload shapes, verified against the upstream model configuration. +workloads: + deepseek-v3: + hidden: 7168 + topk: 8 + routed_experts: 256 + seed: 67 + # Measured source-tokens-per-rank per phase. Powers of two for a clean log + # x-axis; decode runs to 512 tok/rank. Per-backend buffer caps clamp over-cap + # points (dropped, not silently truncated — see ep_harness.token_ladder). + token_ladders: + decode: [1, 2, 4, 8, 16, 32, 64, 128, 256, 512] + prefill: [256, 512, 1024, 2048] + +suites: + ep-core: + mode: normal + workloads: [deepseek-v3] + # "all" = every SKU registered in configs/platform_config.json; a suite may list a subset. + platforms: all + ep_degrees: [8, 16] + routings: [uniform] + phases: [decode, prefill] diff --git a/experimental/CollectiveX/docs/methodology.md b/experimental/CollectiveX/docs/methodology.md new file mode 100644 index 0000000000..ac18477a17 --- /dev/null +++ b/experimental/CollectiveX/docs/methodology.md @@ -0,0 +1,218 @@ +# CollectiveX EP Benchmark Methodology + +CollectiveX schedules expert-parallel (EP) communication benchmarks, executes them on real +accelerator allocations, and uploads the neutral artifacts each run emits. It does **not** validate +those artifacts, promote, rank, recommend, select, hide, or decide what any consumer displays. The +frontend reads the neutral matrix, result, and summary artifacts and makes its own coverage +and display decisions. This document describes how a case is scheduled, measured, checked, and +recorded — not a publication or qualification contract. + +## Product Boundary + +CollectiveX is a communication microbenchmark for: + +- comparing EP libraries on one chip/topology; +- comparing EP latency and logical payload bandwidth across systems under the same workload; and +- surfacing unsupported, failed, invalid, and unstable cases rather than hiding them. + +It does not predict serving throughput without a separate correlation study. + +## Matrix + +The implemented workload is `deepseek-v3`: hidden 7168, top-k 8, 256 routed experts, packed +placement, and one pinned fixed resource profile per backend/topology. Dispatch and combine are +fixed BF16 on every backend; precision is not a swept dimension. Every case uses the normal +`layout-and-dispatch-v1` semantics. + +- `ep-core`: uniform routing over the workload's token ladders — for `deepseek-v3`, decode + T=1..512 powers of two and prefill T=256..2048 powers of two. Ladders are model-specific and + live with the workload in `configs/suites.yaml`. + +`sweep_matrix.py` materializes the requested SKUs, backends, EP sizes, and token ladders into a +matrix document, then extracts strict per-shard controls. `--only-sku`, `--exclude-skus`, and +`--ep-sizes` select a subset; a subset produces a smaller matrix, not a different contract. The +matrix is generated per dispatch; there is no frozen matrix digest or locked case count. + +| Systems | EP8 | EP16 | +|---|---|---| +| H100/H200/B200/B300 | 1x8 NVLink, scale-up | 2x8 NVLink + RDMA, scale-out | +| MI300X/MI325X/MI355X | 1x8 XGMI, scale-up | 2x8 XGMI + RDMA, scale-out | +| GB200/GB300 | 2x4 MNNVL, scale-up | 4x4 MNNVL, scale-up | + +Physical host count does not define scope. Both GB cells remain inside one 72-GPU MNNVL scale-up +domain. + +Unsupported combinations are explicitly classified in the matrix, not silently skipped coverage. DeepEP V2 is the +`ElasticBuffer` introduced by PR #605, pinned with upstream PR #630's minimal pure-scale-up fix and +the exact upstream PR #640 library matcher that excludes NCCL shared-memory mappings. Scale-up cases +request NCCL Device API LSA and fail closed unless the realized LSA team covers the full EP world. +x86 EP16 scale-out uses the hybrid path with GIN and requires two logical scale-out domains +represented by two physical RDMA ranks, with eight scale-up ranks per domain. GB EP16 remains MNNVL +scale-up and uses LSA. MoRI EP8 uses IntraNode-family kernels (MI355X IntraNode, MI300X/MI325X +AsyncLL); EP16 uses pinned InterNodeV1 over 2x8 XGMI + RDMA with 96 blocks, 64 RDMA blocks, 8 warps, +one QP per PE, and external input. +Whether a given SKU/backend/EP cell is attempted is a capability fact; whether it succeeded is +decided only by the emitted artifact. + +## Workload Identity + +One deterministic workload is generated over the global token batch from the workload's seed in +`configs/suites.yaml` (part of the workload identity, baked into every scheduled case) and sliced by +source rank; a keyed BLAKE2b counter over the (token, slot, attempt, stream) coordinates produces +byte-identical expert indices and gate weights on every runtime, and the harness proves the +realized routing trace identical across ranks before a case can succeed. + +Routing traffic distinguishes: + +- token-expert assignments, which determine expert compute load; and +- rank-deduplicated token payload copies, which determine EP activation traffic. + +Adapters may not generate routing or reinterpret one quantity as the other. + +## Measurement + +Normal mode uses `layout-and-dispatch-v1`: dispatch timing includes layout plus communication, and +combine returns activation payload through an unweighted rank-sum path. Expert-output staging is +outside isolated combine timing and inside the measured paired roundtrip. Each component declares +availability, origin, and sample count. A paired-only API reports null isolated components. +`isolated_sum` is derived. The artifact records the mode so a reader can keep distinct measurement +contracts separate. + +Every measured component uses one fixed timing profile, defined once in `configs/suites.yaml` +and baked into every scheduled case: + +- 128 trials x 8 timed iterations = 1024 observations; +- 32 synchronized full dispatch-stage-combine warmups before each available measured component at + every trial/point; +- component measurement order rotates each trial (`trial_order`) so every timed component occupies + every position in the sequence, over a per-trial-rotated token ladder; and +- per-iteration maximum latency across ranks before nearest-rank p50/p90/p95/p99. + +Measured roundtrip p99 is the headline latency. Decode and prefill identify the serving regime +represented by one MoE-layer collective; they do not change the timed primitive at an otherwise +identical shape. Ascending through the ladder, each measured shape is conditioned with 8 untimed +full roundtrips — settling clocks, fabric, and buffer state — before it is correctness-checked; +all timing happens after every shape is warmed and checked. Conditioning rounds are never +measured or emitted. + +Logical payload bandwidth is: + +`logical_payload_bytes / measured_latency_seconds` + +Normal-mode payload bytes use rank-deduplicated token-rank BF16 activations (2 bytes per value, no +scale payload) and exclude expert metadata, padding, and backend buffer capacity. Algorithm bandwidth, bus bandwidth, wire +utilization, and physical-link utilization are not emitted without a defined primitive model or +transport counters. Logical bandwidth must never be labeled physical bandwidth. Payload and token +rates are named `rate_at_latency_percentile`: bytes or tokens divided by the matching latency +percentile. They are lower-tail service rates at p99 latency, not p99 percentiles of an inverted +rate distribution. + +## Correctness + +An implementation-independent oracle uses an expert-specific deterministic transform so wrong expert +routing cannot pass an identity roundtrip. For every rank and point it verifies: + +1. destination rank/expert, source token, multiplicity, gate weight, and receive counts; +2. dispatched payload and metadata before timing; +3. combined output before timing; +4. unchanged semantic inputs through all timed samples; and +5. dispatched payload/metadata and combined output again after timing. + +Normal-mode adapters use activation-only, unweighted rank-sum combine. The oracle builds each rank's +gate-weighted expert aggregate before combine, independently derives `sum(gate * expert(token))`, and +checks the dispatch metadata and transformed output. It compares against the reference activation. +The combine gate is `rtol=0.05, atol=0.02` for the BF16 communication path. This threshold is a correctness gate, not an estimate of transport +error. Any failed rank or point makes the case ineligible in the result it writes. +Pre/post dispatch behavior is checked against canonical source-token metadata and expected output. +Native receive slots may be assigned nondeterministically, so physical receive order is not treated +as a correctness property. + +## Result Artifact + +One raw case document carries `record_type: "case-attempt"` and the single `version`, and contains: + +- `identity`: `case_id`, `attempt_ordinal`, `case_factors` (SKU and the scheduled case — backend, + EP size, mode, phase, suite, workload, and the topology coordinate), and `allocation_factors` + (run id, run attempt, source SHA); +- `workload`: `cross_rank_consistent`, whether the routing trace was proven identical across ranks; +- `measurement`: dispatch/combine dtype and semantics, `sampling`, and the per-point `rows`; +- `implementation`: backend name and kernel generation; +- `topology`: requested SKU/product, placement, nodes, scale-up domain, transport, and world size; +- `provenance`: the mounted image tag and source SHA; and +- `outcome`: `status` (`success` or `invalid`) and `reasons`. + +Each `rows` entry carries point latency, byte accounting, token rate, correctness, load, and fanout; +per-point statistics are summarized in place, not emitted as separate documents. Each dispatched +case writes exactly this one raw result document; unsupported or never-run cells produce no +synthetic record. + +## Identity + +Identifiers are readable factor strings: + +- `case_id`: `{sku}-{backend}-{workload}-{mode}-{phase}-ep{ep}-{routing}`, each factor + slug-normalized; and +- `attempt_ordinal`: a positive integer distinguishing repeat executions of one `case_id`. + +Backend source pins live in `configs/backends.json` and are enforced at fetch by commit; the DeepEP +V2 build is verified at run time by its commit-derived wheel tag. + +These IDs let a consumer group matched configurations and separate distinct ones. The backend does +not itself compute cohorts, controlled comparisons, sensitivity pairs, eligibility, or +recommendations — a reader decides which cases to surface and how to compare them. + +## Execution Isolation + +Every non-MNNVL scale-out case uses operator-pinned socket and RDMA selectors. The launcher rejects +missing or partial profiles, then probes every allocated node for the configured interface, active +HCA port, and configured GID before backend initialization. It never substitutes a default route, +inherited runner environment, or transport fallback. Scale-up and MNNVL cases clear the profile; +scale-out NVIDIA forces `NCCL_NET=IB`, while AMD leaves plugin selection to RCCL. Both use exact HCA +matching. Selector values remain in encrypted config and mode-0600 private logs. + +Repository staging uses a pre-existing, runner-owned, group/world non-writable shared base outside +the checkout and workflow workspace. The parent process resolves the exact execution child before +copying; backend preparation then runs from that tree on every allocated node. Cleanup waits for +confirmed allocation teardown and removes only that child. DeepEP V2 source is fetched before allocation at an +exact pinned revision, initializes its pinned `fmt` submodule, and applies the required local patch. + +H200, B200, and B300 may derive that private base beneath the validated operating-system account home +when it is compute-visible. H100 instead derives a sibling of its shared container directory, never a +child of image storage. +Canonical B300 execution ignores the legacy operator `stage_dir` field and always derives the base +from the validated shared account home. Its UID-mapped Actions shell may accept that exact base when +its owner matches the private parent owner; explicit stages and all other runners retain the strict +effective-UID ownership rule. An execution-ID suffix isolates parallel B300 workers. The current +NFS export may realize a newly created base as +UID 0; only that creation path is accepted, while a pre-existing root-owned base is rejected. +Canonical GB300 execution likewise ignores its legacy group-writable `stage_dir` and derives an +execution-specific private base beneath the validated compute-visible account home. + +## Image Pinning And Build Isolation + +Enroot imports configured container tags into a per-run-scoped squash keyed by the image tag and +image platform, so one run never reuses another run's imported filesystem. Image-provided DeepEP is +also checked against exact package versions and its expected API. Source-built DeepEP V2 uses +a separate mode-0700 cluster-local cache mounted only as `/cx-cache`. Its path binds CPU/GPU +architecture, image, and upstream commit. The cache is never an artifact; per-execution +source/results stages remain isolated and disposable, and runtime probes fail closed before reuse. The runner UID is +inside the trusted cluster boundary: this cache guards against stale or accidental mutation, not +hostile same-UID jobs. Only an unpublished partial build may be reset automatically; a cache that +fails integrity or runtime checks is left intact and rejected so a concurrent allocation cannot lose +files it is using. + +## Neutral Artifact Delivery + +There is no results server, attached store, or managed object store. Each shard runs one allocation, +emits per-case result JSON and a small mechanical summary, and uploads them as GitHub artifacts with +`always()` so a red or partial run still uploads. A case counts as successful on the benchmark's own +return code; there is no completeness or privacy validation step before upload, and failed or +unsupported cells produce no synthetic record. + +No step promotes a run, builds a dataset, or advances a channel; the artifacts are the output. Any +downstream display or comparison is the consumer's responsibility. + +## Legacy Data + +Historical numeric schemas 3-5 are outside this benchmark's artifacts. They remain historical +diagnostic evidence and are not produced or consumed by the current sweep. diff --git a/experimental/CollectiveX/launchers/launch_gb-nv.sh b/experimental/CollectiveX/launchers/launch_gb-nv.sh new file mode 100644 index 0000000000..937d5e3ac9 --- /dev/null +++ b/experimental/CollectiveX/launchers/launch_gb-nv.sh @@ -0,0 +1,81 @@ +#!/usr/bin/env bash +# CollectiveX shared GB200/GB300 NVL72 (aarch64) launcher. +# shellcheck disable=SC2034 +# +# EP8/EP16 use one Slurm task per GPU across two or four trays in the same +# MNNVL scale-up domain. +# +# Flow: +# identity -> setup -> repository-stage -> backend-setup -> scheduler-allocation +# -> container-import -> container-launch -> artifact-collection +set -euo pipefail +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +COLLX_DIR="$(cd "$HERE/.." && pwd)"; REPO_ROOT="$(cd "$COLLX_DIR/../.." && pwd)" +# shellcheck source=../runtime/common.sh +source "$HERE/../runtime/common.sh" + +# ---- identity: resolve SKU, backend, platform ------------------------------- +PRODUCT="${COLLX_SHARD_SKU:-}" +case "$PRODUCT" in + gb200|gb300) ;; + *) collx_die "COLLX_SHARD_SKU must be gb200 or gb300" ;; +esac +RUNNER="$PRODUCT" +export COLLX_RUNNER="$RUNNER" COLLX_BENCH="${COLLX_BENCH:-deepep-v2}" +export COLLX_IMAGE_PLATFORM=linux/arm64 COLLX_VENDOR=nvidia +# ---- setup: operator config, canonical env, topology, network profile ------- +collx_launcher_prologue "$RUNNER" +NODES="${COLLX_NODES:-2}"; GPN="${COLLX_GPUS_PER_NODE:-4}" +SCALE_UP_DOMAIN="${COLLX_SCALE_UP_DOMAIN:-72}" +NGPUS="${COLLX_NGPUS:-$((NODES * GPN))}" +if [ "$PRODUCT" = gb200 ]; then default_time=30; else default_time=90; fi +TIME_MIN="${COLLX_TIME:-$default_time}" +IMAGE="${COLLX_IMAGE:-$COLLX_IMAGE_MULTIARCH}" +TS="$(date -u +%Y-%m-%dT%H-%M-%SZ)" +export COLLX_TS="$TS" COLLX_TOPO="${PRODUCT}-nvl72-mnnvl" +export COLLX_SCOPE=scale-up COLLX_TRANSPORT=mnnvl COLLX_SCALE_UP_TRANSPORT=mnnvl +export COLLX_NODES="$NODES" COLLX_GPUS_PER_NODE="$GPN" COLLX_SCALE_UP_DOMAIN="$SCALE_UP_DOMAIN" +export COLLX_NGPUS="$NGPUS" +unset COLLX_SCALE_OUT_TRANSPORT +case "$COLLX_BENCH" in + deepep-v2) ;; + *) collx_die "unsupported $PRODUCT EP backend: $COLLX_BENCH" ;; +esac +collx_require_vars COLLX_PARTITION COLLX_ACCOUNT COLLX_SQUASH_DIR COLLX_STAGE_DIR +[ "$PRODUCT" != gb300 ] || collx_require_vars COLLX_ENROOT_CACHE_PATH +PARTITION="$COLLX_PARTITION"; ACCOUNT="$COLLX_ACCOUNT"; SQUASH_DIR="$COLLX_SQUASH_DIR" +[ -z "${COLLX_ENROOT_CACHE_PATH:-}" ] || export ENROOT_CACHE_PATH="$COLLX_ENROOT_CACHE_PATH" +export NCCL_CUMEM_ENABLE=1 NCCL_MNNVL_ENABLE=1 MC_FORCE_MNNVL=1 +collx_apply_network_profile "$NODES" "$COLLX_TRANSPORT" + +collx_log "$PRODUCT nodes=$NODES x ${GPN}gpu world=$NGPUS bench=$COLLX_BENCH" +collx_select_image "$IMAGE" + +# ---- repository-stage: compute-visible copy of the checkout ----------------- +MOUNT_SRC="$(collx_stage_path "$REPO_ROOT" "$COLLX_STAGE_DIR")" +collx_stage_repo "$REPO_ROOT" "$MOUNT_SRC" +CONTAINER_MOUNTS="$MOUNT_SRC:/ix" +# ---- backend-setup: pinned DeepEP source + isolated build cache ------------- +# The backend case above admits only deepep-v2, so its staging is unconditional. +collx_prepare_deepep_source "$MOUNT_SRC" \ + || collx_die "cannot stage the pinned backend source" +export COLLX_BACKEND_SOURCE_ROOT=/ix/experimental/CollectiveX/.collx_sources +collx_prepare_backend_cache "$COLLX_SQUASH_DIR" \ + || collx_die "cannot prepare the isolated backend cache" +CONTAINER_MOUNTS="$CONTAINER_MOUNTS,$COLLX_PREPARED_BACKEND_CACHE:/cx-cache" +export COLLX_BACKEND_CACHE_ROOT=/cx-cache + +# ---- scheduler-allocation: salloc the trays --------------------------------- +command -v salloc >/dev/null || collx_die "salloc not found" +collx_salloc_jobid --partition="$PARTITION" --account="$ACCOUNT" --nodes="$NODES" \ + --gres=gpu:"$GPN" --ntasks-per-node="$GPN" --exclusive --mem=0 --cpus-per-task=35 \ + --time="$TIME_MIN" +[ -n "$JOB_ID" ] || collx_die "no JOB_ID from salloc" + +# ---- container-import: squash file resolved on the allocation --------------- +SQUASH_FILE="$(collx_ensure_squash_on_job "$JOB_ID" "$SQUASH_DIR" "$IMAGE")" + +# ---- container-launch -> artifact-collection (shared tail) ------------------ +COLLX_DISTRIBUTED_CONTAINER_ARGS=(--container-writable --container-remap-root) +collx_execute_and_collect "$MOUNT_SRC" "$REPO_ROOT" +exit "$FINAL_RC" diff --git a/experimental/CollectiveX/launchers/launch_mi-amds.sh b/experimental/CollectiveX/launchers/launch_mi-amds.sh new file mode 100644 index 0000000000..b3b551ce9f --- /dev/null +++ b/experimental/CollectiveX/launchers/launch_mi-amds.sh @@ -0,0 +1,130 @@ +#!/usr/bin/env bash +# CollectiveX shared AMD Slurm launcher (one or two nodes). +# shellcheck disable=SC2034 +# +# Flow (container import runs inside the allocation retry loop): +# identity -> setup -> repository-stage -> scheduler-allocation + container-import +# -> container-launch -> artifact-collection +set -euo pipefail + +HERE="$(cd -P -- "$(dirname "${BASH_SOURCE[0]}")" && pwd -P)" +COLLX_DIR="$(cd "$HERE/.." && pwd)" +REPO_ROOT="$(cd "$COLLX_DIR/../.." && pwd)" +# shellcheck source=../runtime/common.sh +source "$HERE/../runtime/common.sh" + +# ---- identity: resolve SKU, backend, platform ------------------------------- +RUNNER="${COLLX_SHARD_SKU:-}" +case "$RUNNER" in + mi300x|mi325x) CPUS_PER_NODE=256; DEVICE_MOUNTS=",/dev/kfd:/dev/kfd,/dev/dri:/dev/dri" ;; + mi355x) CPUS_PER_NODE=128; DEVICE_MOUNTS="" ;; + *) collx_die "COLLX_SHARD_SKU is not a registered AMD SKU" ;; +esac +export COLLX_RUNNER="$RUNNER" COLLX_BENCH="${COLLX_BENCH:-mori}" +export COLLX_IMAGE_PLATFORM=linux/amd64 COLLX_VENDOR=amd +# ---- setup: operator config, canonical env, topology, network profile ------- +collx_launcher_prologue "$RUNNER" + +NODES="${COLLX_NODES:-1}"; GPN="${COLLX_GPUS_PER_NODE:-8}" +SCALE_UP_DOMAIN="${COLLX_SCALE_UP_DOMAIN:-8}" +NGPUS="${COLLX_NGPUS:-$((NODES * GPN))}" +TIME_MIN="${COLLX_TIME:-60}" +EXCLUDE_NODES="${COLLX_EXCLUDE_NODES:-}" +NODELIST="${COLLX_NODELIST:-}" +MOUNT_DIR=/ix +TS="$(date -u +%Y-%m-%dT%H-%M-%SZ)" +case "$COLLX_BENCH" in + mori) ;; + *) collx_die "unsupported AMD EP backend: $COLLX_BENCH" ;; +esac + +export MORI_DISABLE_AUTO_XGMI="${MORI_DISABLE_AUTO_XGMI:-0}" +export MORI_ENABLE_SDMA="${MORI_ENABLE_SDMA:-1}" +export MORI_APP_LOG_LEVEL="${MORI_APP_LOG_LEVEL:-info}" +export MORI_SHMEM_LOG_LEVEL="${MORI_SHMEM_LOG_LEVEL:-info}" +export MORI_IO_LOG_LEVEL="${MORI_IO_LOG_LEVEL:-info}" +IMAGE="${COLLX_IMAGE:-$COLLX_IMAGE_AMD_MORI}" +export COLLX_NGPUS="$NGPUS" COLLX_NODES="$NODES" +export COLLX_GPUS_PER_NODE="$GPN" COLLX_SCALE_UP_DOMAIN="$SCALE_UP_DOMAIN" COLLX_TS="$TS" +export COLLX_SCALE_UP_TRANSPORT=xgmi +if [ "$NODES" -gt 1 ]; then + export COLLX_SCOPE=scale-out COLLX_SCALE_OUT_TRANSPORT=rdma + export COLLX_TRANSPORT=xgmi-rdma COLLX_TOPO="${RUNNER}-xgmi-rdma" +else + export COLLX_SCOPE=scale-up COLLX_TRANSPORT=xgmi COLLX_TOPO="${RUNNER}-xgmi" + unset COLLX_SCALE_OUT_TRANSPORT +fi +export COLLX_RUN_TIMEOUT="${COLLX_RUN_TIMEOUT:-1800}" +collx_apply_network_profile "$NODES" "$COLLX_TRANSPORT" +collx_require_vars COLLX_PARTITION COLLX_SQUASH_DIR COLLX_STAGE_DIR +PARTITION="$COLLX_PARTITION"; SQUASH_DIR="$COLLX_SQUASH_DIR" + +collx_log "runner=$RUNNER nodes=$NODES x ${GPN}gpu world=$NGPUS bench=$COLLX_BENCH" + +# ---- repository-stage: compute-visible copy of the checkout ----------------- +MOUNT_SRC="$(collx_stage_path "$REPO_ROOT" "$COLLX_STAGE_DIR")" +collx_stage_repo "$REPO_ROOT" "$MOUNT_SRC" +collx_select_image "$IMAGE" + +# ---- scheduler-allocation + container-import: retry until nodes validate ---- +# Each attempt must pass the network profile AND import the squash; a rejected +# allocation is cancelled and its nodes excluded from the next attempt. +command -v salloc >/dev/null || collx_die "salloc not found on this runner" + +allocation=(--partition="$PARTITION" --nodes="$NODES" --gres=gpu:"$GPN" + --time="$TIME_MIN" --ntasks-per-node="$GPN" + --cpus-per-task="$((CPUS_PER_NODE / GPN))") +if [ "$RUNNER" = mi355x ]; then + allocation+=(--exclusive) +fi +excluded_nodes="$EXCLUDE_NODES" +for allocation_attempt in 1 2 3; do + attempt_allocation=("${allocation[@]}") + if [ -n "$NODELIST" ]; then + collx_log "using configured node pin" + attempt_allocation+=(--nodelist="$NODELIST") + elif [ -n "$excluded_nodes" ]; then + attempt_allocation+=(--exclude="$excluded_nodes") + fi + export COLLX_SALLOC_ATTEMPT="$allocation_attempt" + export COLLX_NETWORK_VALIDATION_ATTEMPT="$allocation_attempt" + collx_salloc_jobid "${attempt_allocation[@]}" + [ -n "$JOB_ID" ] || collx_die "could not resolve allocated JOB_ID from salloc" + reject_reason="" + if ! collx_validate_network_profile_on_job "$JOB_ID" "$NODES" "$COLLX_TRANSPORT"; then + # A node whose RoCE devices do not match the pinned selector (e.g. an + # outlier still using default rocepXXXs0 names instead of the rdmaN udev + # names the rest of the fleet exposes) must be rejected and retried + # elsewhere, not treated as a hard failure. + reject_reason=network + else + if SQUASH_FILE="$(collx_ensure_squash_on_job \ + "$JOB_ID" "$SQUASH_DIR" "$IMAGE" "${COLLX_LOCK_DIR:-}")"; then + break + fi + reject_reason=container-import + fi + if [ -n "$NODELIST" ] || [ "$allocation_attempt" = 3 ]; then + if [ "$reject_reason" = network ]; then + collx_log_tail "${COLLX_NETWORK_PROFILE_LOG:-}" + collx_die "allocated nodes failed the network profile" + fi + collx_die "allocated nodes failed container import" + fi + rejected_nodes="$(collx_allocation_nodes_csv "$JOB_ID")" \ + || collx_die "cannot identify nodes from a rejected allocation" + collx_log "allocated nodes failed $reject_reason validation; retrying elsewhere" + collx_cleanup_allocation || collx_die "cannot release a rejected allocation" + JOB_ID="" + [ -z "$excluded_nodes" ] || excluded_nodes+=, + excluded_nodes+="$rejected_nodes" +done +unset COLLX_SALLOC_ATTEMPT COLLX_NETWORK_VALIDATION_ATTEMPT +CONTAINER_MOUNTS="$MOUNT_SRC:$MOUNT_DIR$DEVICE_MOUNTS" + +# ---- container-launch -> artifact-collection (shared tail) ------------------ +COLLX_DISTRIBUTED_CONTAINER_ARGS=(--container-writable --container-remap-root) +collx_execute_and_collect "$MOUNT_SRC" "$REPO_ROOT" +rm -f "$MOUNT_SRC"/experimental/CollectiveX/gpucore.* 2>/dev/null || true +collx_log "done - result artifacts collected" +exit "$FINAL_RC" diff --git a/experimental/CollectiveX/launchers/launch_single-slurm.sh b/experimental/CollectiveX/launchers/launch_single-slurm.sh new file mode 100644 index 0000000000..65c289ebe4 --- /dev/null +++ b/experimental/CollectiveX/launchers/launch_single-slurm.sh @@ -0,0 +1,151 @@ +#!/usr/bin/env bash +# CollectiveX shared standard NVIDIA Slurm launcher (one or two nodes). +# shellcheck disable=SC2034 +# +# Flow: +# identity -> setup -> repository-stage -> backend-setup -> scheduler-allocation +# -> container-import -> container-launch -> artifact-collection +set -euo pipefail + +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +COLLX_DIR="$(cd "$HERE/.." && pwd)" +REPO_ROOT="$(cd "$COLLX_DIR/../.." && pwd)" +# shellcheck source=../runtime/common.sh +source "$HERE/../runtime/common.sh" + +# ---- identity: resolve SKU, backend, platform ------------------------------- +RUNNER="${COLLX_SHARD_SKU:-}" +ALLOC_EXTRA=(); SRUN_EXTRA=(); LOCAL_IMPORT=0 +case "$RUNNER" in + h100-dgxc) PRODUCT=h100; DEFAULT_TIME=45; REQUIRE_ACCOUNT=1 ;; + h200-dgxc) + PRODUCT=h200; DEFAULT_TIME=45; REQUIRE_ACCOUNT=0 + SRUN_EXTRA=(--container-remap-root) + ;; + b200-dgxc) + PRODUCT=b200; DEFAULT_TIME=30; REQUIRE_ACCOUNT=1 + ALLOC_EXTRA=(--mem=0) + ;; + b300) + PRODUCT=b300; DEFAULT_TIME=45; REQUIRE_ACCOUNT=1 + # Do not restore ALLOC_EXTRA=(-N 1 --mem=0); it blocks two-node B300 jobs. + ALLOC_EXTRA=(--mem=0) + SRUN_EXTRA=(--mpi=none --container-remap-root) + LOCAL_IMPORT=1 + ;; + *) collx_die "COLLX_SHARD_SKU is not a registered NVIDIA SKU" ;; +esac +TOPO="${PRODUCT}-nvlink-island" +export COLLX_RUNNER="$RUNNER" COLLX_BENCH="${COLLX_BENCH:-deepep-v2}" +export COLLX_IMAGE_PLATFORM=linux/amd64 COLLX_VENDOR=nvidia +# ---- setup: operator config, canonical env, topology, network profile ------- +collx_launcher_prologue "$RUNNER" + +NODES="${COLLX_NODES:-1}"; GPN="${COLLX_GPUS_PER_NODE:-8}" +SCALE_UP_DOMAIN="${COLLX_SCALE_UP_DOMAIN:-8}" +NGPUS="${COLLX_NGPUS:-$((NODES * GPN))}" +TIME_MIN="${COLLX_TIME:-$DEFAULT_TIME}" +IMAGE="${COLLX_IMAGE:-$COLLX_IMAGE_MULTIARCH}" +TS="$(date -u +%Y-%m-%dT%H-%M-%SZ)" +case "$COLLX_BENCH" in + deepep-v2) ;; + *) collx_die "unsupported $RUNNER EP backend: $COLLX_BENCH" ;; +esac + +export COLLX_NGPUS="$NGPUS" COLLX_NODES="$NODES" +export COLLX_GPUS_PER_NODE="$GPN" COLLX_SCALE_UP_DOMAIN="$SCALE_UP_DOMAIN" +export COLLX_TS="$TS" COLLX_SCALE_UP_TRANSPORT=nvlink +if [ "$NODES" -gt 1 ]; then + export COLLX_SCOPE=scale-out COLLX_SCALE_OUT_TRANSPORT=rdma + export COLLX_TRANSPORT=nvlink-rdma COLLX_TOPO="${PRODUCT}-nvlink-rdma" +else + export COLLX_SCOPE=scale-up COLLX_TRANSPORT=nvlink COLLX_TOPO="$TOPO" + unset COLLX_SCALE_OUT_TRANSPORT +fi +export NCCL_CUMEM_ENABLE=1 +collx_apply_network_profile "$NODES" "$COLLX_TRANSPORT" +collx_require_vars COLLX_PARTITION COLLX_SQUASH_DIR +[ "$REQUIRE_ACCOUNT" = 0 ] || collx_require_vars COLLX_ACCOUNT +# b300 /home and /scratch are node-local; h100-dgxc /home is login-local (absent on +# compute nodes) — on both, the implicit passwd-home stage base is compute-invisible, +# so the operator config must pin an explicit compute-visible stage_dir. +case "$RUNNER" in + b300|h100-dgxc) collx_require_vars COLLX_STAGE_DIR ;; +esac + +collx_log "runner=$RUNNER nodes=$NODES x ${GPN}gpu world=$NGPUS bench=$COLLX_BENCH" +collx_select_image "$IMAGE" + +# ---- repository-stage: compute-visible copy of the checkout ----------------- +MOUNT_SRC="$(collx_stage_path "$REPO_ROOT" "${COLLX_STAGE_DIR:-}")" +collx_stage_repo "$REPO_ROOT" "$MOUNT_SRC" +CONTAINER_MOUNTS="$MOUNT_SRC:/ix" +# ---- backend-setup: pinned DeepEP source + isolated build cache ------------- +# The backend case above admits only deepep-v2, so its staging is unconditional. +collx_prepare_deepep_source "$MOUNT_SRC" \ + || collx_die "cannot stage the pinned backend source" +export COLLX_BACKEND_SOURCE_ROOT=/ix/experimental/CollectiveX/.collx_sources +collx_prepare_backend_cache "$COLLX_SQUASH_DIR" \ + || collx_die "cannot prepare the isolated backend cache" +CONTAINER_MOUNTS="$CONTAINER_MOUNTS,$COLLX_PREPARED_BACKEND_CACHE:/cx-cache" +export COLLX_BACKEND_CACHE_ROOT=/cx-cache + +# ---- scheduler-allocation: salloc, retry until nodes validate --------------- +# Each attempt must pass the network profile (and accelerator-context on b300); +# a rejected allocation is cancelled and its nodes excluded from the next attempt. +command -v salloc >/dev/null || collx_die "salloc not found on this runner" +allocation=(--partition="$COLLX_PARTITION" --nodes="$NODES" --gres=gpu:"$GPN" + --ntasks-per-node="$GPN" --exclusive --time="$TIME_MIN" "${ALLOC_EXTRA[@]}") +[ -z "${COLLX_ACCOUNT:-}" ] || allocation+=(--account="$COLLX_ACCOUNT") +[ -z "${COLLX_QOS:-}" ] || allocation+=(--qos="$COLLX_QOS") +[ -z "${COLLX_NODELIST:-}" ] || allocation+=(--nodelist="$COLLX_NODELIST") +excluded_nodes="${COLLX_EXCLUDE_NODES:-}" +for allocation_attempt in 1 2 3; do + validation_failure="" + attempt_allocation=("${allocation[@]}") + [ -z "$excluded_nodes" ] || attempt_allocation+=(--exclude="$excluded_nodes") + export COLLX_SALLOC_ATTEMPT="$allocation_attempt" + export COLLX_NETWORK_VALIDATION_ATTEMPT="$allocation_attempt" + collx_salloc_jobid "${attempt_allocation[@]}" + [ -n "$JOB_ID" ] || collx_die "could not resolve allocated JOB_ID from salloc" + if ! collx_validate_network_profile_on_job "$JOB_ID" "$NODES" "$COLLX_TRANSPORT"; then + validation_failure=network + elif [ "$RUNNER" = b300 ] \ + && ! collx_validate_cuda_context_on_job "$JOB_ID" "$NODES" "$GPN"; then + validation_failure=cuda-context + else + break + fi + retryable=0 + [ "$RUNNER:$validation_failure" != h100-dgxc:network ] || retryable=1 + [ "$RUNNER:$validation_failure" != b300:cuda-context ] || retryable=1 + if [ "$retryable" = 0 ] || [ "$allocation_attempt" = 3 ]; then + if [ "$validation_failure" = network ]; then + collx_log_tail "${COLLX_NETWORK_PROFILE_LOG:-}" + collx_die "allocated nodes failed the network profile" + fi + collx_log_tail "$COLLX_CUDA_CONTEXT_LOG" + collx_die "allocated nodes failed accelerator context validation" + fi + rejected_nodes="$(collx_allocation_nodes_csv "$JOB_ID")" \ + || collx_die "cannot identify nodes from a rejected allocation" + collx_log "allocated nodes failed $validation_failure validation; retrying elsewhere" + collx_cleanup_allocation || collx_die "cannot release a rejected allocation" + JOB_ID="" + [ -z "$excluded_nodes" ] || excluded_nodes+=, + excluded_nodes+="$rejected_nodes" +done +unset COLLX_SALLOC_ATTEMPT COLLX_NETWORK_VALIDATION_ATTEMPT + +# ---- container-import: squash file (login-local on b300, else on the job) ---- +if [ "$LOCAL_IMPORT" = 1 ]; then + SQUASH_FILE="$(COLLX_ENROOT_LOCAL_IMPORT=1 collx_ensure_squash "$COLLX_SQUASH_DIR" "$IMAGE")" +else + SQUASH_FILE="$(collx_ensure_squash_on_job "$JOB_ID" "$COLLX_SQUASH_DIR" "$IMAGE")" +fi + +# ---- container-launch -> artifact-collection (shared tail) ------------------ +COLLX_DISTRIBUTED_CONTAINER_ARGS=(--container-writable "${SRUN_EXTRA[@]}") +collx_execute_and_collect "$MOUNT_SRC" "$REPO_ROOT" +collx_log "done - result artifacts collected" +exit "$FINAL_RC" diff --git a/experimental/CollectiveX/requirements.txt b/experimental/CollectiveX/requirements.txt new file mode 100644 index 0000000000..838ee056b1 --- /dev/null +++ b/experimental/CollectiveX/requirements.txt @@ -0,0 +1,2 @@ +# Host-side matrix generation. GPU libraries are supplied by benchmark images. +PyYAML==6.0.2 diff --git a/experimental/CollectiveX/runtime/common.sh b/experimental/CollectiveX/runtime/common.sh new file mode 100644 index 0000000000..5d85087b0e --- /dev/null +++ b/experimental/CollectiveX/runtime/common.sh @@ -0,0 +1,930 @@ +# shellcheck shell=bash +# CollectiveX — shared launcher helpers (sourced, not executed). +# +# Cluster-generic scaffolding only (Slurm/container/build/staging); no +# model-serving. Logging goes to stderr so functions can `echo` a single +# result on stdout. + +unset COLLECTIVEX_OPERATOR_CONFIG_LOADED +COLLX_RUNTIME_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" + +collx_log() { printf '[collectivex] %s\n' "$*" >&2; } +collx_die() { printf '[collectivex] FATAL: %s\n' "$*" >&2; exit 1; } + +# Public benchmark identity (backend source pins + container images) is data, +# not launcher state: configs/backends.json. config.py validates the registry +# and emits the COLLX_* names consumed here, by prepare_backend.sh (which +# re-sources this file in-container), and by the launchers. +collx_load_backend_registry() { + local registry key value name + registry="$(mktemp /tmp/inferencex-collectivex-registry.XXXXXX)" \ + || collx_die "cannot stage the backend registry" + if ! python3 "$COLLX_RUNTIME_DIR/config.py" backend-registry > "$registry"; then + rm -f -- "$registry" + collx_die "invalid backend registry (configs/backends.json)" + fi + while IFS= read -r -d '' key && IFS= read -r -d '' value; do + printf -v "$key" '%s' "$value" + done < "$registry" + rm -f -- "$registry" + for name in COLLX_IMAGE_MULTIARCH COLLX_IMAGE_AMD_MORI \ + COLLX_DEEPEP_V2_REPO COLLX_DEEPEP_V2_COMMIT; do + [ -n "${!name:-}" ] || collx_die "backend registry omits $name" + done +} +collx_load_backend_registry + +# Print bounded command output without maintaining a parallel failure taxonomy. +collx_log_tail() { + local log_path="$1" + if [ -s "$log_path" ]; then + collx_log "--- command log tail ---" + tail -n 100 -- "$log_path" >&2 || true + collx_log "--- end command log tail ---" + fi +} + +# Shared launcher skeleton: the identity-stage boilerplate every launcher runs +# before any SKU-specific work. +collx_launcher_prologue() { + JOB_ID="" + collx_install_launcher_fail_safe + [ -n "${COLLX_SHARD_FILE:-}" ] || collx_die "COLLX_SHARD_FILE is required" + collx_load_operator_config + collx_prepare_stage_dir "$1" +} + +# Shared launcher tail: run the shard, collect artifacts, and fold both return +# codes into FINAL_RC (run failures win; collection failures surface otherwise). +collx_execute_and_collect() { + local mount_src="$1" repo_root="$2" run_rc=0 collect_rc=0 + collx_run_shard || run_rc=$? + collx_collect_results "$mount_src" "$repo_root" || collect_rc=$? + FINAL_RC="$run_rc" + [ "$FINAL_RC" != 0 ] || FINAL_RC="$collect_rc" +} + +collx_job_root_is_safe() { + local root="$1" + if [[ "$root" =~ ^/tmp/inferencex-collectivex-[0-9]+-[0-9]+-[A-Za-z0-9._-]+$ ]]; then + : + elif [[ "$root" =~ ^/tmp/inferencex-collectivex-parent-([0-9]+)-([0-9]+)-([A-Za-z0-9._-]+)/inferencex-collectivex-([0-9]+)-([0-9]+)-([A-Za-z0-9._-]+)$ ]]; then + [ "${BASH_REMATCH[1]}" = "${BASH_REMATCH[4]}" ] \ + && [ "${BASH_REMATCH[2]}" = "${BASH_REMATCH[5]}" ] \ + && [ "${BASH_REMATCH[3]}" = "${BASH_REMATCH[6]}" ] || return 1 + else + return 1 + fi + [ -d "$root" ] && [ ! -L "$root" ] \ + && [ "$(stat -c '%u:%a' "$root" 2>/dev/null)" = "$(id -u):700" ] +} + +# Runner-local deployment settings are strict JSON kept outside the checkout. +# Only the selected runner's allowlisted values are exported; the document is +# never sourced or evaluated as shell. +collx_load_operator_config() { + [ -n "${COLLECTIVEX_OPERATOR_CONFIG_LOADED:-}" ] \ + && [ "$COLLECTIVEX_OPERATOR_CONFIG_LOADED" = "$$" ] && return 0 + local config_path generated=0 parsed_path key value + unset COLLX_PARTITION COLLX_ACCOUNT COLLX_QOS COLLX_SQUASH_DIR COLLX_STAGE_DIR COLLX_ENROOT_CACHE_PATH + unset ENROOT_CACHE_PATH + unset COLLX_EXCLUDE_NODES COLLX_NODELIST COLLX_LOCK_DIR COLLX_MASTER_PORT + unset COLLX_SOCKET_IFNAME COLLX_RDMA_DEVICES COLLX_IB_GID_INDEX COLLX_RDMA_SERVICE_LEVEL + unset COLLX_RDMA_TRAFFIC_CLASS + unset MASTER_ADDR MASTER_PORT RANK WORLD_SIZE LOCAL_RANK LOCAL_WORLD_SIZE + config_path="${COLLECTIVEX_OPERATOR_CONFIG:-${XDG_CONFIG_HOME:-${HOME}/.config}/inferencex/collectivex.json}" + if [ -n "${COLLECTIVEX_OPERATOR_CONFIG_CONTENT:-}" ]; then + umask 077 + if collx_job_root_is_safe "${COLLX_JOB_ROOT:-}"; then + config_path="$COLLX_JOB_ROOT/operator-config.json" + (set -C; : > "$config_path") 2>/dev/null \ + || collx_die "cannot create ephemeral runner configuration" + else + config_path="$(mktemp /tmp/inferencex-collectivex-config.XXXXXX)" \ + || collx_die "cannot create ephemeral runner configuration" + fi + generated=1 + if ! printf '%s' "$COLLECTIVEX_OPERATOR_CONFIG_CONTENT" > "$config_path"; then + unset COLLECTIVEX_OPERATOR_CONFIG_CONTENT + rm -f -- "$config_path" + collx_die "cannot materialize runner configuration" + fi + elif [ "${COLLECTIVEX_OPERATOR_CONFIG_REQUIRED:-0}" = 1 ]; then + unset COLLECTIVEX_OPERATOR_CONFIG_CONTENT + collx_die "runner configuration is unavailable" + fi + unset COLLECTIVEX_OPERATOR_CONFIG_CONTENT COLLECTIVEX_OPERATOR_CONFIG_REQUIRED + if [ ! -e "$config_path" ]; then + [ "${COLLECTIVEX_CANONICAL_GHA:-0}" != 1 ] \ + || collx_die "runner configuration is unavailable" + if [ -z "${COLLX_RUNNER:-${COLLX_SHARD_SKU:-}}" ]; then + COLLECTIVEX_OPERATOR_CONFIG_LOADED="$$" + return 0 + fi + # No operator document, but the SKU is known: the tracked registry's + # per-SKU `operator` block still supplies baseline fields, so emit + # registry-only values ("-" sentinel; a no-op for SKUs without a block). + config_path="-" + fi + umask 077 + parsed_path="$(mktemp /tmp/inferencex-collectivex-parsed.XXXXXX)" || { + [ "$generated" = 0 ] || rm -f -- "$config_path" + collx_die "cannot parse runner configuration" + } + if ! python3 "$COLLX_RUNTIME_DIR/config.py" operator-config "$config_path" \ + "${COLLX_RUNNER:-${COLLX_SHARD_SKU:-}}" \ + > "$parsed_path" + then + rm -f -- "$parsed_path" + [ "$generated" = 0 ] || rm -f -- "$config_path" + unset COLLECTIVEX_OPERATOR_CONFIG COLLECTIVEX_OPERATOR_CONFIG_EPHEMERAL + collx_die "runner-local configuration failed" + fi + while IFS= read -r -d '' key && IFS= read -r -d '' value; do + printf -v "$key" '%s' "$value" + export "${key?}" + done < "$parsed_path" + rm -f -- "$parsed_path" + if [ "$generated" = 1 ] || [ "${COLLECTIVEX_OPERATOR_CONFIG_EPHEMERAL:-0}" = 1 ]; then + rm -f -- "$config_path" || collx_die "cannot remove ephemeral runner configuration" + fi + unset COLLECTIVEX_OPERATOR_CONFIG COLLECTIVEX_OPERATOR_CONFIG_EPHEMERAL + COLLECTIVEX_OPERATOR_CONFIG_LOADED="$$" +} + +# Per-step log files: several callers parse these for markers (salloc grant, +# stage copy-error, per-node network selectors), so they are a data channel, +# not just failure display. Logs persist after the run for postmortem. +collx_private_log_path() { + local path="${COLLX_JOB_ROOT:-/tmp/inferencex-collectivex-$(id -u)}/logs/$1.log" + mkdir -p "${path%/*}" || collx_die "cannot create log directory" + : > "$path" || collx_die "cannot create runtime log" + printf '%s' "$path" +} + +# Host-side utility steps need only the basic login paths. They never receive +# the complete Actions or runner environment. +collx_host_exports() { + printf '%s' 'HOME,PATH,USER,XDG_CACHE_HOME,ENROOT_CACHE_PATH' +} + +collx_require_vars() { + local name + local -a missing=() + for name in "$@"; do + [ -n "${!name:-}" ] || missing+=("$name") + done + [ "${#missing[@]}" -eq 0 ] || collx_die \ + "missing runner-local configuration: ${missing[*]} (set them in COLLECTIVEX_OPERATOR_CONFIG)" +} + +collx_export_gid_index_for_link_layer() { + local link_layer="$1" + unset NVSHMEM_IB_GID_INDEX NCCL_IB_GID_INDEX + [ -n "${COLLX_IB_GID_INDEX:-}" ] || return 0 + case "$link_layer" in + roce) + export NVSHMEM_IB_GID_INDEX="$COLLX_IB_GID_INDEX" + export NCCL_IB_GID_INDEX="$COLLX_IB_GID_INDEX" + ;; + infiniband) ;; + *) collx_die "unsupported RDMA link layer" ;; + esac +} + +# Convert private, runner-local network selectors into the public library +# variables needed inside the container. Values are interface/HCA identifiers, +# never addresses; the rendezvous hostname is derived from the allocation. +collx_apply_network_profile() { + local nodes="$1" transport="$2" + local selector rdma_name rdma_names="" ep_nic="" + local -a selectors + [[ "$nodes" =~ ^[1-9][0-9]*$ ]] || collx_die "invalid network placement" + unset NCCL_NET NCCL_SOCKET_IFNAME GLOO_SOCKET_IFNAME NCCL_IB_HCA + unset NCCL_IB_GID_INDEX NCCL_IB_SL + unset NVSHMEM_ENABLE_NIC_PE_MAPPING + unset NVSHMEM_HCA_LIST NVSHMEM_IB_GID_INDEX NVSHMEM_IB_SL + unset NVSHMEM_IB_ENABLE_IBGDA NVSHMEM_IBGDA_NIC_HANDLER + unset EP_NIC_NAME EP_OVERRIDE_RDMA_SL + unset MORI_RDMA_DEVICES + unset MORI_RDMA_TC MORI_IO_TC MORI_RDMA_SL MORI_IO_SL + # Single-node and MNNVL runs need only the scrub above; everything past this + # point is the scale-out path, so no per-branch scale-out guards remain. + { [ "$nodes" -gt 1 ] && [ "$transport" != mnnvl ]; } || return 0 + [ -n "${COLLX_RDMA_DEVICES:-}" ] \ + || collx_die "RDMA execution requires a private device selector" + if [ -n "${COLLX_SOCKET_IFNAME:-}" ]; then + [[ "$COLLX_SOCKET_IFNAME" =~ ^[A-Za-z][A-Za-z0-9_.-]{0,31}$ ]] \ + || collx_die "invalid private socket interface selector" + export NCCL_SOCKET_IFNAME="$COLLX_SOCKET_IFNAME" GLOO_SOCKET_IFNAME="$COLLX_SOCKET_IFNAME" + fi + [[ "$COLLX_RDMA_DEVICES" =~ ^[A-Za-z][A-Za-z0-9_.-]{0,31}(:[1-9][0-9]*)?(,[A-Za-z][A-Za-z0-9_.-]{0,31}(:[1-9][0-9]*)?)*$ ]] \ + || collx_die "invalid private RDMA device selector" + IFS=, read -r -a selectors <<< "$COLLX_RDMA_DEVICES" + for selector in "${selectors[@]}"; do + rdma_name="${selector%%:*}" + rdma_names="${rdma_names}${rdma_names:+,}${rdma_name}" + [ -n "$ep_nic" ] || ep_nic="$rdma_name" + done + export NVSHMEM_HCA_LIST="$COLLX_RDMA_DEVICES" + export NVSHMEM_ENABLE_NIC_PE_MAPPING=1 + # RCCL selects its own net plugin; NCCL_NET=IB breaks AMD SKUs. + if [ "${COLLX_VENDOR:-nvidia}" = amd ]; then + unset NCCL_NET + else + export NCCL_NET=IB + fi + export NCCL_IB_HCA="=$COLLX_RDMA_DEVICES" + export MORI_RDMA_DEVICES="$rdma_names" EP_NIC_NAME="$ep_nic" + if [ -n "${COLLX_IB_GID_INDEX:-}" ]; then + [[ "$COLLX_IB_GID_INDEX" =~ ^[0-9]+$ ]] && [ "$COLLX_IB_GID_INDEX" -le 255 ] \ + || collx_die "invalid private IB GID index" + fi + if [ -n "${COLLX_RDMA_SERVICE_LEVEL:-}" ]; then + [[ "$COLLX_RDMA_SERVICE_LEVEL" =~ ^[0-9]+$ ]] && [ "$COLLX_RDMA_SERVICE_LEVEL" -le 15 ] \ + || collx_die "invalid private RDMA service level" + export NVSHMEM_IB_SL="$COLLX_RDMA_SERVICE_LEVEL" + export NCCL_IB_SL="$COLLX_RDMA_SERVICE_LEVEL" + export EP_OVERRIDE_RDMA_SL="$COLLX_RDMA_SERVICE_LEVEL" + export MORI_RDMA_SL="$COLLX_RDMA_SERVICE_LEVEL" MORI_IO_SL="$COLLX_RDMA_SERVICE_LEVEL" + fi + if [ -n "${COLLX_RDMA_TRAFFIC_CLASS:-}" ]; then + [[ "$COLLX_RDMA_TRAFFIC_CLASS" =~ ^[0-9]+$ ]] && [ "$COLLX_RDMA_TRAFFIC_CLASS" -le 255 ] \ + || collx_die "invalid private RDMA traffic class" + export MORI_RDMA_TC="$COLLX_RDMA_TRAFFIC_CLASS" MORI_IO_TC="$COLLX_RDMA_TRAFFIC_CLASS" + fi + local nic_handler=gpu + export NVSHMEM_IB_ENABLE_IBGDA=1 NVSHMEM_IBGDA_NIC_HANDLER="$nic_handler" + if [ -n "${COLLX_RDMA_LINK_LAYER:-}" ]; then + case "$COLLX_RDMA_LINK_LAYER" in + roce|infiniband) ;; + *) collx_die "invalid validated RDMA link layer" ;; + esac + collx_export_gid_index_for_link_layer "$COLLX_RDMA_LINK_LAYER" + fi +} + +# Slurm may remove NCCL's leading exact-match marker while propagating an +# inherited environment. Reconstruct it from the validated private selector at +# the container boundary instead of accepting a prefix-matched HCA list. +collx_restore_exact_hca_selector() { + if [ "${COLLX_NODES:-1}" -le 1 ] || [ "${COLLX_TRANSPORT:-}" = mnnvl ]; then + return 0 + fi + [ -n "${COLLX_RDMA_DEVICES:-}" ] \ + || { collx_log "ERROR: scale-out RDMA selector is unavailable"; return 1; } + [[ "$COLLX_RDMA_DEVICES" =~ ^[A-Za-z][A-Za-z0-9_.-]{0,31}(:[1-9][0-9]*)?(,[A-Za-z][A-Za-z0-9_.-]{0,31}(:[1-9][0-9]*)?)*$ ]] \ + || { collx_log "ERROR: invalid scale-out RDMA selector"; return 1; } + export NCCL_IB_HCA="=$COLLX_RDMA_DEVICES" +} + +collx_default_route_interface() { + python3 "$COLLX_RUNTIME_DIR/probe.py" default-route-interface +} + +# Prove that the operator-pinned scale-out fabric exists on every allocated +# node before image import or backend initialization. Selector values and node +# diagnostics stay in the runner-private log. +collx_validate_network_profile_on_job() { + local job_id="$1" nodes="$2" transport="$3" + local log_label=network-profile log rc=0 marker_count link_layer + { [ "$nodes" -gt 1 ] && [ "$transport" != mnnvl ]; } || return 0 + [[ "$job_id" =~ ^[1-9][0-9]*$ && "$nodes" =~ ^[1-9][0-9]*$ ]] \ + || return 1 + [ -n "${COLLX_RDMA_DEVICES:-}" ] || return 1 + case "${COLLX_NETWORK_VALIDATION_ATTEMPT:-1}" in + 1) ;; + 2|3) log_label+="-a${COLLX_NETWORK_VALIDATION_ATTEMPT}" ;; + *) return 1 ;; + esac + log="$(collx_private_log_path "$log_label")" || return 1 + export COLLX_NETWORK_PROFILE_LOG="$log" + srun --jobid="$job_id" --nodes="$nodes" --ntasks="$nodes" --ntasks-per-node=1 \ + --chdir=/tmp --input=all --export="$(collx_host_exports)" \ + python3 /dev/stdin network-profile "${COLLX_SOCKET_IFNAME:-}" \ + "$COLLX_RDMA_DEVICES" "${COLLX_IB_GID_INDEX:-}" \ + < "$COLLX_RUNTIME_DIR/probe.py" > "$log" 2>&1 || rc=$? + if [ "$rc" != 0 ]; then + marker="$(grep -aoE '(socket-interface|rdma-(device|port))-[0-9]+=(missing|down|inactive|default-route-missing|gid-missing|gid-empty|link-layer-missing|link-layer-invalid|link-layer-mixed)' "$log" \ + | tail -n 1 || true)" + [ -z "$marker" ] || collx_log "ERROR: network-profile-$marker" + return "$rc" + fi + socket_ifname="$( + sed -nE 's/^\[collectivex-private\] socket-interface-selected=([A-Za-z][A-Za-z0-9_.-]{0,31})$/\1/p' "$log" \ + | sort -u + )" + marker_count="$(grep -Ec '^\[collectivex-private\] socket-interface-selected=' "$log")" + socket_unique_count="$(printf '%s\n' "$socket_ifname" | sed '/^$/d' | wc -l | tr -d ' ')" + if [ "$socket_unique_count" -lt 1 ] || [ "$marker_count" != "$nodes" ]; then + collx_log "ERROR: network-profile-socket-markers=$marker_count/$nodes unique=$socket_unique_count" + return 1 + fi + if [ "$socket_unique_count" = 1 ]; then + export COLLX_SOCKET_IFNAME="$socket_ifname" + else + unset COLLX_SOCKET_IFNAME + fi + link_layer="$( + sed -nE 's/^\[collectivex-private\] rdma-link-layer=(roce|infiniband)$/\1/p' "$log" \ + | sort -u + )" + marker_count="$(grep -Ec '^\[collectivex-private\] rdma-link-layer=(roce|infiniband)$' "$log")" + case "$marker_count:$link_layer" in + "$nodes":roce|"$nodes":infiniband) ;; + *) return 1 ;; + esac + export COLLX_RDMA_LINK_LAYER="$link_layer" + collx_export_gid_index_for_link_layer "$link_layer" +} + +collx_allocation_nodes_csv() { + local job_id="$1" nodelist node output="" + [[ "$job_id" =~ ^[1-9][0-9]*$ ]] || return 1 + nodelist="$(squeue -h -j "$job_id" -o %N 2>/dev/null)" || return 1 + [[ "$nodelist" =~ ^[][A-Za-z0-9._,-]+$ ]] || return 1 + while IFS= read -r node; do + [ -n "$node" ] || continue + [[ "$node" =~ ^[A-Za-z0-9][A-Za-z0-9._-]*$ ]] || return 1 + [ -z "$output" ] || output+=, + output+="$node" + done < <(scontrol show hostnames "$nodelist" 2>/dev/null) + [ -n "$output" ] || return 1 + printf '%s' "$output" +} + +collx_resolve_slurm_rendezvous() { + local job_id="$1" master_addr master_port socket_ifname="${COLLX_SOCKET_IFNAME:-}" + [[ "$job_id" =~ ^[1-9][0-9]*$ ]] || collx_die "invalid rendezvous allocation" + # Query relative node zero directly so MASTER_ADDR always hosts global rank 0. + # Prefer the address on the already validated cross-node socket interface; + # a short hostname may resolve onto a management network that ranks cannot use. + if [[ "$socket_ifname" =~ ^[A-Za-z][A-Za-z0-9_.-]{0,31}$ ]]; then + master_addr="$(srun --jobid="$job_id" --nodes=1 --ntasks=1 --relative=0 \ + --chdir=/tmp --export="$(collx_host_exports)" bash -s -- "$socket_ifname" \ + 2>/dev/null <<'BASH' | head -n1 +set -euo pipefail +ip -o -4 address show dev "$1" scope global \ + | awk 'NR == 1 {split($4, address, "/"); print address[1]}' +BASH +)" + [[ "$master_addr" =~ ^([0-9]{1,3}\.){3}[0-9]{1,3}$ ]] \ + || collx_die "could not resolve the allocated primary interface" + else + master_addr="$(srun --jobid="$job_id" --nodes=1 --ntasks=1 --relative=0 \ + --chdir=/tmp --export="$(collx_host_exports)" hostname -s 2>/dev/null | head -n1)" + [[ "$master_addr" =~ ^[A-Za-z0-9][A-Za-z0-9._-]*$ ]] \ + || collx_die "could not resolve the allocated primary node" + fi + master_port="${COLLX_MASTER_PORT:-29551}" + [[ "$master_port" =~ ^[1-9][0-9]*$ ]] && [ "$master_port" -le 65535 ] \ + || collx_die "invalid distributed rendezvous port" + export MASTER_ADDR="$master_addr" MASTER_PORT="$master_port" +} + +# Printed into `bash -c` ahead of the rank wrapper. Sources the per-node +# loader/import environment persisted by collx_persist_backend_env, refusing +# a file whose directory shape, ownership, or mode differs from what that step wrote. +collx_source_backend_env() { + cat <<'BASH' +case "${SLURM_NODEID:-}" in ""|*[!0-9]*) exit 66;; esac +. "/ix/experimental/CollectiveX/.collx_backend/env/node-${SLURM_NODEID}.sh" || exit 66 +BASH +} + +# Printed into `bash -c` for one Slurm task per GPU. Every rank derives its +# identity from Slurm rather than accepting caller-supplied rank values. +collx_slurm_rank_wrapper() { + cat <<'BASH' +case "${SLURM_PROCID:-}:${SLURM_NTASKS:-}:${SLURM_LOCALID:-}:${SLURM_NODEID:-}" in + *[!0-9:]*|:*|*::*|*:) exit 67 ;; +esac +[ "$SLURM_NTASKS" = "$COLLX_NGPUS" ] || exit 67 +[ "$SLURM_LOCALID" -lt "$COLLX_GPUS_PER_NODE" ] || exit 67 +. /ix/experimental/CollectiveX/runtime/common.sh || exit 68 +if [ "${COLLX_NODES:-1}" -gt 1 ] && [ "${COLLX_TRANSPORT:-}" != mnnvl ]; then + if [ -z "${COLLX_SOCKET_IFNAME:-}" ]; then + COLLX_SOCKET_IFNAME="$(collx_default_route_interface)" || exit 68 + [[ "$COLLX_SOCKET_IFNAME" =~ ^[A-Za-z][A-Za-z0-9_.-]{0,31}$ ]] || exit 68 + export COLLX_SOCKET_IFNAME + fi + collx_apply_network_profile "$COLLX_NODES" "$COLLX_TRANSPORT" || exit 68 +fi +export RANK="$SLURM_PROCID" WORLD_SIZE="$SLURM_NTASKS" +export LOCAL_RANK="$SLURM_LOCALID" LOCAL_WORLD_SIZE="$COLLX_GPUS_PER_NODE" +exec python3 bench/run_ep.py "$@" +BASH +} + +# Allocate via salloc's stable grant message and assign JOB_ID in this shell. +# Record it so workflow cleanup can release a launcher interrupted by Actions. +collx_salloc_jobid() { + local log_label=scheduler-allocation log job_id root="${COLLX_JOB_ROOT:-}" + case "${COLLX_SALLOC_ATTEMPT:-1}" in + 1) ;; + 2|3) log_label+="-a${COLLX_SALLOC_ATTEMPT}" ;; + *) return 1 ;; + esac + if ! log="$(collx_private_log_path "$log_label")"; then + collx_log "ERROR: scheduler log is unavailable" + return 1 + fi + collx_log "scheduler-request=submit" + if ! (salloc "$@" --no-shell) > "$log" 2>&1; then + collx_log "ERROR: scheduler allocation failed" + collx_log_tail "$log" + return 1 + fi + job_id="$(sed -nE \ + 's/.*Granted job allocation ([1-9][0-9]*).*/\1/p' "$log" | head -n1)" + [[ "$job_id" =~ ^[1-9][0-9]*$ ]] || return 1 + JOB_ID="$job_id" + if [ -n "$root" ]; then + collx_job_root_is_safe "$root" || return 1 + (umask 077; printf '%s\n' "$JOB_ID" > "$root/jobid") || return 1 + fi +} + +# Idempotent cleanup for launcher traps, allocation retries, and workflow recovery. +collx_cleanup_allocation() { + local root="${1:-${COLLX_JOB_ROOT:-}}" path="" job_id="${JOB_ID:-}" active + if [ -n "$root" ]; then + collx_job_root_is_safe "$root" || return 1 + path="$root/jobid" + if [ -z "$job_id" ] && [ -f "$path" ]; then + IFS= read -r job_id < "$path" || return 1 + fi + fi + [ -n "$job_id" ] || return 0 + [[ "$job_id" =~ ^[1-9][0-9]*$ ]] || return 1 + scancel "$job_id" >/dev/null 2>&1 || true + for _ in {1..30}; do + active="$(squeue -h -j "$job_id" -o %A 2>/dev/null)" || active=unknown + if [ -z "$active" ]; then + [ -z "$path" ] || rm -f -- "$path" + return + fi + sleep 1 + done + collx_log "ERROR: scheduled allocation did not terminate during cleanup" + return 1 +} + +# Image references come from configs/backends.json (collx_load_backend_registry). +# Import uses the configured tag because Enroot cannot reliably import a +# digest-qualified Docker Hub reference non-interactively. +collx_select_image() { + local image="$1" + [[ "$image" =~ ^[A-Za-z0-9._/-]+:[A-Za-z0-9._-]+$ ]] \ + || collx_die "configured image reference is malformed" + export COLLECTIVEX_IMAGE="$image" +} + +# Create a per-UID cache under validated cluster-local storage. Only the fixed +# /cx-cache mount enters the container; the operator host path does not. +collx_prepare_backend_cache() { + local cache + unset COLLX_PREPARED_BACKEND_CACHE + cache="$(python3 "$COLLX_RUNTIME_DIR/probe.py" prepare-cache "$1")" || return 1 + [[ "$cache" = /* ]] || return 1 + export COLLX_PREPARED_BACKEND_CACHE="$cache" +} + +# Fetch the pinned DeepEP tree before allocating GPUs. +collx_prepare_deepep_source() { + local mount_src="$1" root source temporary log + root="$mount_src/experimental/CollectiveX/.collx_sources" + source="$root/deepep-v2-$COLLX_DEEPEP_V2_COMMIT" + [ ! -d "$source" ] || return 0 + mkdir -p -- "$root" && chmod 700 "$root" || return 1 + temporary="$(mktemp -d "$root/.deepep-v2.XXXXXX")" || return 1 + log="$(collx_private_log_path backend-source-deepep-v2)" || return 1 + # On b300 the NFS export can realize a newly created stage dir as UID 0 while + # git runs as the UID-mapped Actions user, tripping git's "dubious ownership" + # guard on the source tree and its fmt submodule. HOME is this job's ephemeral + # dir and the runner UID is inside the trusted cluster boundary, so scope the + # exemption globally (also reaches the submodule child git). + git config --global --add safe.directory '*' >> "$log" 2>&1 || true + if GIT_TERMINAL_PROMPT=0 git init -q "$temporary" > "$log" 2>&1 \ + && git -C "$temporary" remote add origin "$COLLX_DEEPEP_V2_REPO" >> "$log" 2>&1 \ + && GIT_TERMINAL_PROMPT=0 git -C "$temporary" fetch -q --no-tags --depth 1 \ + origin "$COLLX_DEEPEP_V2_COMMIT" >> "$log" 2>&1 \ + && git -C "$temporary" -c advice.detachedHead=false checkout -q --detach FETCH_HEAD \ + >> "$log" 2>&1 \ + && [ "$(git -C "$temporary" rev-parse HEAD)" = "$COLLX_DEEPEP_V2_COMMIT" ] \ + && GIT_TERMINAL_PROMPT=0 git -C "$temporary" submodule update -q --init --depth 1 \ + third-party/fmt >> "$log" 2>&1 \ + && python3 "$COLLX_RUNTIME_DIR/stage.py" rewrite-deepep-v2 \ + "$temporary/deep_ep/__init__.py" >> "$log" 2>&1 \ + && mv -- "$temporary" "$source" >> "$log" 2>&1; then + return 0 + fi + rm -rf -- "$temporary" + collx_log "ERROR: DeepEP source preparation failed" + collx_log_tail "$log" + return 1 +} + +collx_materialize_deepep_source() { + local destination="$1" source + [ -n "${COLLX_BACKEND_SOURCE_ROOT:-}" ] || return 1 + source="$COLLX_BACKEND_SOURCE_ROOT/deepep-v2-$COLLX_DEEPEP_V2_COMMIT" + [ -d "$source" ] || return 1 + rm -rf -- "$destination" && cp -R -- "$source" "$destination" +} + +collx_prepare_implicit_stage_base() { + python3 "$COLLX_RUNTIME_DIR/stage.py" implicit-stage-base "${1:-}" "${2:-}" +} + +collx_prepare_runner_shared_stage_base() { + local runner_temp="${RUNNER_TEMP:-}" runner_root + case "$runner_temp" in + /*/_work/_temp) runner_root="${runner_temp%/_work/_temp}" ;; + *) collx_die "canonical AMD execution requires a standard shared runner temp" ;; + esac + [ -n "$runner_root" ] && [ "$runner_root" != "$runner_temp" ] \ + || collx_die "canonical AMD execution requires a shared runner root" + collx_prepare_implicit_stage_base "$runner_root" +} + +collx_prepare_stage_dir() { + local runner="$1" + [ "${COLLECTIVEX_CANONICAL_GHA:-0}" = 1 ] || return 0 + [ -n "${COLLX_SQUASH_DIR:-}" ] \ + || collx_die "canonical CollectiveX execution requires shared container storage" + case "$runner" in b300|gb300) COLLX_STAGE_DIR="" ;; esac + if [ -z "${COLLX_STAGE_DIR:-}" ]; then + case "$runner" in + h100-dgxc) + COLLX_STAGE_DIR="$(collx_prepare_implicit_stage_base "${COLLX_SQUASH_DIR%/*}")" \ + || collx_die "canonical CollectiveX execution cannot create an isolated shared stage directory" + ;; + b300|gb300) + COLLX_STAGE_DIR="$(collx_prepare_implicit_stage_base "" \ + "${COLLECTIVEX_EXECUTION_ID:-${GITHUB_RUN_ID:-}}")" \ + || collx_die "canonical CollectiveX execution cannot create an isolated stage directory" + ;; + h200-dgxc|b200-dgxc) + COLLX_STAGE_DIR="$(collx_prepare_implicit_stage_base)" \ + || collx_die "canonical CollectiveX execution cannot create an isolated stage directory" + ;; + mi300x|mi325x|mi355x) + COLLX_STAGE_DIR="$(collx_prepare_runner_shared_stage_base)" \ + || collx_die "canonical AMD execution cannot create an isolated shared stage directory" + ;; + *) collx_die "canonical CollectiveX execution requires a configured shared stage directory" ;; + esac + elif [ "$runner" = mi300x ]; then + COLLX_STAGE_DIR="$(python3 "$COLLX_RUNTIME_DIR/stage.py" resolve-directory \ + "$COLLX_STAGE_DIR")" \ + || collx_die "canonical MI300X execution cannot resolve the shared stage directory" + fi + export COLLX_STAGE_DIR +} + +collx_squash_path() { + local squash_dir="$1" image="$2" key platform run_scope + case "${COLLX_IMAGE_PLATFORM:-}" in + linux/amd64) platform="" ;; + linux/arm64) platform="_linux_arm64" ;; + *) return 1 ;; + esac + run_scope="${GITHUB_RUN_ID:-${COLLECTIVEX_EXECUTION_ID:-manual}}-${GITHUB_RUN_ATTEMPT:-1}" + run_scope="$(printf '%s' "$run_scope" | tr -cs 'A-Za-z0-9_.-' '-')" || return 1 + run_scope="${run_scope#-}"; run_scope="${run_scope%-}" + [ -n "$run_scope" ] || return 1 + key="${platform}_${run_scope}_$( + printf '%s' "$image" | sed 's#[/:@#]#_#g' + )" + printf '%s' "$squash_dir/${key}.sqsh" +} + +# collx_ensure_squash -> echoes the squash file path. +# Imports via Enroot only if a valid squash is not already present, under a lock. +collx_ensure_squash() { + local squash_dir="$1" image="$2" key sq locks lock_fd log + local enroot_local="" import_rc=0 machine + log="$(collx_private_log_path container-import)" + machine="$(uname -m)" + case "${COLLX_IMAGE_PLATFORM:-}:$machine" in + linux/amd64:x86_64|linux/amd64:amd64|linux/arm64:aarch64|linux/arm64:arm64) ;; + *) collx_log_tail "$log"; return 1 ;; + esac + mkdir -p "$squash_dir" 2>> "$log" \ + || { collx_log_tail "$log"; return 1; } + sq="$(collx_squash_path "$squash_dir" "$image")" \ + || { collx_log_tail "$log"; return 1; } + key="${sq##*/}" + key="${key%.sqsh}" + locks="$squash_dir/.locks" + mkdir -p "$locks" 2>> "$log" \ + || { collx_log_tail "$log"; return 1; } + { exec {lock_fd}>"$locks/${key}.lock"; } 2>> "$log" \ + || { collx_log_tail "$log"; return 1; } + flock -w 900 "$lock_fd" 2>> "$log" \ + || { collx_log_tail "$log"; return 1; } + if unsquashfs -l "$sq" >/dev/null 2>&1; then + collx_log "container squash ready" + else + collx_log "importing configured container image" + rm -f "$sq" 2>> "$log" \ + || { collx_log_tail "$log"; return 1; } + # > "$log" 2>&1 || import_rc=$? + rm -rf -- "$enroot_local" >/dev/null 2>&1 || true + [ "$import_rc" = 0 ] \ + || { collx_log_tail "$log"; return 1; } + else + enroot import -o "$sq" "docker://$image" > "$log" 2>&1 \ + || { collx_log_tail "$log"; return 1; } + fi + unsquashfs -l "$sq" >> "$log" 2>&1 \ + || { collx_log_tail "$log"; return 1; } + fi + flock -u "$lock_fd" + exec {lock_fd}>&- + echo "$sq" +} + +# Import on an allocated compute node so multiarch tags resolve for the target +# architecture. The squash directory must be shared with the submit host. +collx_ensure_squash_on_job() { + local job_id="$1" squash_dir="$2" image="$3" lock_dir="${4:-}" sq key lock + local log_label=container-import log + [[ "$job_id" =~ ^[0-9]+$ ]] || return 1 + case "${COLLX_SALLOC_ATTEMPT:-1}" in + 1) ;; + 2|3) log_label+="-a${COLLX_SALLOC_ATTEMPT}" ;; + *) return 1 ;; + esac + sq="$(collx_squash_path "$squash_dir" "$image")" || return 1 + key="${sq##*/}" + key="${key%.sqsh}" + [ -n "$lock_dir" ] || lock_dir="$squash_dir/.locks" + lock="$lock_dir/${key}.lock" + log="$(collx_private_log_path "$log_label")" + # Run once per node because some clusters use node-local squash storage. + if ! srun --jobid="$job_id" --nodes="${COLLX_NODES:-1}" --ntasks="${COLLX_NODES:-1}" \ + --ntasks-per-node=1 --chdir=/tmp \ + --export="$(collx_host_exports)" \ + bash -s -- "$sq" "$lock" "$image" "$COLLX_IMAGE_PLATFORM" \ + > "$log" 2>&1 <<'BASH' +set -euo pipefail +sq="$1"; lock="$2"; image="$3"; platform="$4" +machine="$(uname -m)" +case "$platform:$machine" in + linux/amd64:x86_64|linux/amd64:amd64|linux/arm64:aarch64|linux/arm64:arm64) ;; + *) exit 13 ;; +esac +compute_home="$(mktemp -d /tmp/inferencex-collectivex-home.XXXXXX)" +trap 'rm -rf -- "$compute_home"' EXIT +export HOME="$compute_home" XDG_CACHE_HOME="$compute_home/.cache" +export ENROOT_TEMP_PATH="$compute_home/enroot-tmp" +export ENROOT_CACHE_PATH="$compute_home/enroot-cache" +export ENROOT_DATA_PATH="$compute_home/enroot-data" +export ENROOT_RUNTIME_PATH="$compute_home/enroot-run" +mkdir -p "$(dirname "$sq")" "$(dirname "$lock")" \ + "$ENROOT_TEMP_PATH" "$ENROOT_CACHE_PATH" "$ENROOT_DATA_PATH" "$ENROOT_RUNTIME_PATH" +exec 9>"$lock" +# Shared storage serializes the import; node-local storage imports in parallel. +flock 9 +if unsquashfs -l "$sq" >/dev/null 2>&1; then + echo 'container squash ready' +else + rm -f -- "$sq" + enroot import -o "$sq" "docker://$image" /dev/null 2>&1 +fi +BASH + then + collx_log "ERROR: container import failed" + collx_log_tail "$log" + return 1 + fi + printf '%s' "$sq" +} + +# A clean nvidia-smi inventory does not prove that a prior cancelled workload +# released every CUDA context. Retaining each primary context catches poisoned +# allocations before a full shard spends time failing every case. +collx_validate_cuda_context_on_job() { + local job_id="$1" nodes="$2" gpus_per_node="$3" log_label=cuda-context log + case "${COLLX_SALLOC_ATTEMPT:-1}" in + 1) ;; + 2|3) log_label+="-a${COLLX_SALLOC_ATTEMPT}" ;; + *) return 1 ;; + esac + log="$(collx_private_log_path "$log_label")" + export COLLX_CUDA_CONTEXT_LOG="$log" + srun --jobid="$job_id" --nodes="$nodes" --ntasks="$nodes" --ntasks-per-node=1 \ + --gres=gpu:"$gpus_per_node" --chdir=/tmp --input=all \ + --export="$(collx_host_exports)" python3 /dev/stdin cuda-context "$gpus_per_node" \ + < "$COLLX_RUNTIME_DIR/probe.py" >"$log" 2>&1 +} + +# Resolve the exact per-execution child before any copy starts, so the parent +# EXIT trap can remove an interrupted partial stage. The configured base must +# already exist on compute-visible storage and must not traverse symlinks. +collx_stage_path() { + local repo_root="$1" stage_base="${2:-}" tag stage_path + tag="${COLLECTIVEX_EXECUTION_ID:-${GITHUB_RUN_ID:-manual-$$}}" + [[ "$tag" =~ ^[A-Za-z0-9][A-Za-z0-9._-]*$ ]] \ + || collx_die "invalid staging execution identity" + if [ -z "$stage_base" ] || [ "$stage_base" = "$repo_root" ]; then + [ -n "${COLLX_SQUASH_DIR:-}" ] \ + || collx_die "CollectiveX staging requires COLLX_STAGE_DIR or COLLX_SQUASH_DIR" + stage_base="$COLLX_SQUASH_DIR" + stage_path="${stage_base%/}/.collectivex-stage-$tag" + else + stage_path="${stage_base%/}/job_$tag" + fi + python3 "$COLLX_RUNTIME_DIR/stage.py" validate-stage-path "$repo_root" "$stage_base" \ + "$stage_path" "${COLLX_JOB_ROOT:-}" "${GITHUB_WORKSPACE:-}" +} + +# Stage only the public benchmark tree into the private execution child. +collx_stage_repo() { + local repo_root="$1" stage_dir="$2" log + python3 "$COLLX_RUNTIME_DIR/stage.py" create-stage "$stage_dir" \ + || collx_die "cannot create the configured stage directory" + collx_log "staging CollectiveX on compute-visible storage" + log="$(collx_private_log_path repository-stage)" + if ! python3 "$COLLX_RUNTIME_DIR/stage.py" copy-repository \ + "$repo_root/experimental/CollectiveX" \ + "$stage_dir/experimental/CollectiveX" > "$log" 2>&1; then + rm -rf -- "$stage_dir" >/dev/null 2>&1 \ + || collx_log "ERROR: cannot remove the incomplete execution stage" + collx_log "ERROR: repository staging failed" + collx_log_tail "$log" + return 1 + fi +} + +# collx_collect_results +# When the run used a staged (compute-visible) mount, copy result JSONs back to +# the original checkout's results/ so the workflow's upload-artifact (which reads +# the checkout, not the stage dir) finds them. No-op when no staging was used. +collx_collect_results() { + local mount_src="$1" repo_root="$2" dst log + local -a files + [ "$mount_src" = "$repo_root" ] && return 0 + log="$(collx_private_log_path "artifact-collection-$$-${RANDOM}")" + dst="$repo_root/experimental/CollectiveX/results" + mkdir -p "$dst" 2>> "$log" \ + || { collx_log "ERROR: cannot create checkout result directory"; return 1; } + shopt -s nullglob + files=("$mount_src/experimental/CollectiveX/results/"*.json) + shopt -u nullglob + [ "${#files[@]}" -gt 0 ] || { collx_log "ERROR: staged run produced no result JSON"; return 1; } + cp -- "${files[@]}" "$dst/" >> "$log" 2>&1 \ + || { collx_log "ERROR: staged result collection failed"; return 1; } + collx_log "collected staged results for artifact validation" +} + +collx_cleanup_stage() { + local mount_src="$1" repo_root="$2" + [ "$mount_src" != "$repo_root" ] || return 0 + if ! python3 "$COLLX_RUNTIME_DIR/stage.py" validate-cleanup "$mount_src"; then + collx_log "ERROR: refusing to remove an invalid stage directory" + return 1 + fi + rm -rf -- "$mount_src" >/dev/null 2>&1 || { + collx_log "ERROR: cannot remove generated stage directory" + return 1 + } + collx_log "removed generated per-execution stage directory" +} + +# Run one shard with one Slurm task per GPU on one or more nodes. +# Launchers provide only allocation/container policy through globals and +# COLLX_DISTRIBUTED_CONTAINER_ARGS; per-case benchmark inputs travel as run_ep.py +# argv decoded from the shard control (config.py case-args), never as env. +# shellcheck disable=SC2153 +collx_run_shard() { + local build_log expected_cases ci=0 failed_cases=0 + local runtime_log argv_file shard wrap + local -a container_args ep_args + [ "${NODES:-0}" -ge 1 ] && [ "${NGPUS:-0}" = "$((NODES * GPN))" ] \ + || collx_die "invalid shard launcher placement" + [ -n "${JOB_ID:-}" ] && [ -n "${SQUASH_FILE:-}" ] \ + && [ -n "${CONTAINER_MOUNTS:-}" ] || collx_die "shard launcher is incomplete" + wrap="$(collx_source_backend_env)"$'\n'"$(collx_slurm_rank_wrapper)" + + collx_resolve_slurm_rendezvous "$JOB_ID" + collx_apply_network_profile "$NODES" "${COLLX_TRANSPORT:-}" + mkdir -p "$MOUNT_SRC/experimental/CollectiveX/results" + container_args=(--container-mounts="$CONTAINER_MOUNTS" --no-container-mount-home + --container-workdir=/ix/experimental/CollectiveX --no-container-entrypoint) + if declare -p COLLX_DISTRIBUTED_CONTAINER_ARGS >/dev/null 2>&1; then + container_args+=("${COLLX_DISTRIBUTED_CONTAINER_ARGS[@]}") + fi + local container_name="cxep_${JOB_ID}" + + shard="${COLLX_SHARD_FILE:-}" + [ -f "$shard" ] || shard="$COLLX_DIR/$shard" + [ -f "$shard" ] || collx_die "shard control is unavailable" + expected_cases="$(python3 "$COLLX_RUNTIME_DIR/config.py" case-count "$shard")" \ + && [[ "$expected_cases" =~ ^[1-9][0-9]*$ ]] \ + || collx_die "could not enumerate shard cases" + + collx_log "shard backend preparation: bench=$COLLX_BENCH nodes=$NODES" + build_log="$(collx_private_log_path backend-prepare)" + if ! srun --jobid="$JOB_ID" --nodes="$NODES" --ntasks-per-node=1 --chdir=/tmp \ + --container-name="$container_name" --container-image="$SQUASH_FILE" \ + "${container_args[@]}" --export=ALL \ + bash /ix/experimental/CollectiveX/runtime/prepare_backend.sh \ + "$build_log" 2>&1; then + collx_log "ERROR: backend preparation failed" + collx_log_tail "$build_log" + return 1 + fi + + argv_file="$(mktemp)" || return 1 + while [ "$ci" -lt "$expected_cases" ]; do + python3 "$COLLX_RUNTIME_DIR/config.py" case-args "$shard" "$ci" \ + "$RUNNER" "$TS" \ + "$NGPUS" "$NODES" "$GPN" "$SCALE_UP_DOMAIN" > "$argv_file" \ + || { rm -f "$argv_file"; collx_die "shard case $ci does not decode against this allocation"; } + mapfile -d '' -t ep_args < "$argv_file" + [ "${#ep_args[@]}" -gt 0 ] \ + || { rm -f "$argv_file"; collx_die "case $ci produced no benchmark arguments"; } + collx_log "EP${NGPUS}[$((ci + 1))/$expected_cases] $COLLX_BENCH" + runtime_log="$(collx_private_log_path "runtime-c$(printf '%03d' "$ci")")" + if ! timeout -k 30 "${COLLX_RUN_TIMEOUT:-900}" \ + srun --jobid="$JOB_ID" --nodes="$NODES" \ + --ntasks="$NGPUS" --ntasks-per-node="$GPN" --chdir=/tmp \ + --container-name="$container_name" --container-image="$SQUASH_FILE" \ + "${container_args[@]}" \ + --export=ALL \ + bash -c "$wrap" _ "${ep_args[@]}" \ + "$runtime_log" 2>&1; then + collx_log "ERROR: case $ci failed" + collx_log_tail "$runtime_log" + failed_cases=$((failed_cases + 1)) + fi + ci=$((ci + 1)) + done + rm -f "$argv_file" + [ "$failed_cases" = 0 ] || { + collx_log "ERROR: $failed_cases/$expected_cases case(s) failed" + return 1 + } +} + +# Remove this allocation's persistent pyxis container before the allocation is +# released. Clusters may run pyxis with container_scope=global, where the named +# --container-writable container every shard uses (cxep_) survives job +# teardown and its unpacked rootfs — tens of GB per node — would otherwise +# accumulate on every allocated node's local image store until it fills and the +# next writable extraction fails with ENOSPC. Best-effort and bounded: teardown +# must never hang or fail on this. +collx_remove_distributed_container() { + local job_id="$1" nodes="${2:-1}" + [ -n "$job_id" ] || return 0 + [ "$nodes" -ge 1 ] 2>/dev/null || return 0 + timeout 120 srun --jobid="$job_id" --nodes="$nodes" --ntasks-per-node=1 \ + --chdir=/tmp enroot remove -f "pyxis_cxep_${job_id}" \ + /dev/null 2>&1 || true +} + +collx_launcher_cleanup() { + local rc="$1" stage_root="${MOUNT_SRC:-}" + trap - EXIT HUP INT TERM + if [ -n "${JOB_ID:-}" ]; then + collx_remove_distributed_container "$JOB_ID" "${NODES:-1}" + if ! collx_cleanup_allocation; then + [ "$rc" != 0 ] || rc=1 + exit "$rc" + fi + fi + if [ -n "${REPO_ROOT:-}" ] && [ -n "$stage_root" ] \ + && [ "$stage_root" != "$REPO_ROOT" ]; then + if [ "$rc" != 0 ] && [ -d "$stage_root/experimental/CollectiveX" ]; then + collx_collect_results "$stage_root" "$REPO_ROOT" || true + fi + if ! collx_cleanup_stage "$stage_root" "$REPO_ROOT"; then + [ "$rc" != 0 ] || rc=1 + fi + fi + exit "$rc" +} + +collx_install_launcher_fail_safe() { + trap 'collx_launcher_cleanup "$?"' EXIT + trap 'collx_launcher_cleanup 129' HUP + trap 'collx_launcher_cleanup 130' INT + trap 'collx_launcher_cleanup 143' TERM +} diff --git a/experimental/CollectiveX/runtime/config.py b/experimental/CollectiveX/runtime/config.py new file mode 100644 index 0000000000..1366d98775 --- /dev/null +++ b/experimental/CollectiveX/runtime/config.py @@ -0,0 +1,284 @@ +#!/usr/bin/env python3 +"""Load private runner settings, the public backend registry, and shard controls.""" + +from __future__ import annotations + +import argparse +import json +import os +import re +import sys + + +FIELDS = { + "partition": "COLLX_PARTITION", "account": "COLLX_ACCOUNT", "qos": "COLLX_QOS", + "squash_dir": "COLLX_SQUASH_DIR", "stage_dir": "COLLX_STAGE_DIR", + "enroot_cache_path": "COLLX_ENROOT_CACHE_PATH", "exclude_nodes": "COLLX_EXCLUDE_NODES", + "nodelist": "COLLX_NODELIST", "lock_dir": "COLLX_LOCK_DIR", + "socket_ifname": "COLLX_SOCKET_IFNAME", "rdma_devices": "COLLX_RDMA_DEVICES", + "ib_gid_index": "COLLX_IB_GID_INDEX", "rdma_service_level": "COLLX_RDMA_SERVICE_LEVEL", + "rdma_traffic_class": "COLLX_RDMA_TRAFFIC_CLASS", +} + + +def _platforms() -> dict: + """The per-SKU platform registry (configs/platform_config.json). Callers + fail closed on a missing file, unknown SKU, or missing field.""" + path = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + "configs", "platform_config.json", + ) + with open(path, encoding="utf-8") as stream: + return json.load(stream)["platforms"] + + +OPERATOR_ENV = ( + "COLLECTIVEX_OPERATOR_CONFIG_CONTENT", "COLLECTIVEX_NETWORK_CONFIG_CONTENT", + "COLLECTIVEX_H100_CONFIG_CONTENT", "COLLECTIVEX_B300_CONFIG_CONTENT", + "COLLECTIVEX_B200_CONFIG_CONTENT", "COLLECTIVEX_MI300_CONFIG_CONTENT", + "COLLECTIVEX_MI355_CONFIG_CONTENT", +) + + +def emit(values: dict[str, object]) -> None: + for field, value in values.items(): + name = FIELDS.get(field, field) + sys.stdout.buffer.write(name.encode() + b"\0" + str(value).encode() + b"\0") + + +def _network_overlay(runner: str) -> dict[str, object]: + """Repo-tracked per-SKU scale-out RDMA selectors — the `network` block of the + SKU's configs/platform_config.json entry — overlaid onto the base operator + config. Only network FIELDS are taken, so identity keys and notes are ignored; + a missing/invalid file is a no-op fallback to the base/secret network fields.""" + try: + block = _platforms().get(runner, {}).get("network", {}) + except (KeyError, OSError, TypeError, json.JSONDecodeError): + return {} + return {key: value for key, value in block.items() if key in FIELDS} + + +def operator_config(path: str, runner: str) -> None: + try: + # The registry's tracked per-SKU `operator` block is the baseline + # (de-secreted by operator decision); an operator config document, when + # provided, overrides it per field. Path "-" means registry-only — and + # for SKUs with no tracked block it preserves the no-config behavior + # (emit nothing). + selected = dict(_platforms()[runner].get("operator", {})) + if path == "-": + if not selected: + return + else: + with open(path, encoding="utf-8") as stream: + document = json.load(stream) + selected.update(document["runners"].get(runner, {})) + # Overlay repo-tracked scale-out RDMA selectors onto the base runner config; + # SKUs without a platform_config.json network block keep their base/secret + # network fields. + selected.update(_network_overlay(runner)) + missing = set(_platforms()[runner]["operator_fields"]) - set(selected) + if missing: + print("validation-missing-required-" + "-".join(sorted(missing)), file=sys.stderr) + raise SystemExit(1) + allowed = set(FIELDS) | {"storage_roots"} + if set(selected) - allowed: + raise ValueError + roots = selected.pop("storage_roots", None) + if roots: + for root in roots: + squash = os.path.join(root, "collectivex", "containers") + stage = os.path.join(root, "collectivex", "stage") + try: + os.makedirs(squash, mode=0o700, exist_ok=True) + os.makedirs(stage, mode=0o700, exist_ok=True) + selected.update(squash_dir=squash, stage_dir=stage) + break + except OSError: + continue + else: + raise ValueError + if any(not isinstance(value, (str, int)) or "\0" in str(value) for value in selected.values()): + raise ValueError + emit(selected) + except (KeyError, OSError, TypeError, ValueError, json.JSONDecodeError): + print("validation-invalid-config", file=sys.stderr) + raise SystemExit(1) + + +def merge_operator_config(path: str) -> None: + """Merge the base and per-cluster encrypted operator documents.""" + def pairs(items): + result = {} + for key, value in items: + if key in result: + raise ValueError("duplicate configuration key") + result[key] = value + return result + + def load_env(name: str) -> dict: + return json.loads( + os.environ[name], object_pairs_hook=pairs, + parse_constant=lambda _: (_ for _ in ()).throw(ValueError()), + ) + + # An absent/blank base secret (all operator config de-secreted into the + # tracked platform_config.json) yields an empty base; the per-SKU registry + # `operator` block then supplies every field downstream. A present base is + # still parsed strictly. Overlay secrets are already skipped when empty. + base = (load_env(OPERATOR_ENV[0]) + if os.environ.get(OPERATOR_ENV[0], "").strip() else {"runners": {}}) + if not isinstance(base.get("runners"), dict): + raise ValueError("invalid operator runners") + base = {"runners": base["runners"]} + for name in OPERATOR_ENV[1:]: + if not os.environ.get(name): + continue + overlay = load_env(name) + if not isinstance(overlay.get("runners"), dict): + raise ValueError("invalid overlay runners") + for runner, fields in overlay["runners"].items(): + if not isinstance(fields, dict) or not fields: + raise ValueError("invalid overlay runner") + base["runners"].setdefault(runner, {}).update(fields) + payload = json.dumps(base, sort_keys=True, separators=(",", ":")) + "\n" + if len(payload.encode()) > 65536: + raise ValueError("merged operator configuration is too large") + descriptor = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) + with os.fdopen(descriptor, "w", encoding="utf-8") as stream: + stream.write(payload) + + +# Same acceptance as common.sh collx_select_image; keep the two in sync. +_IMAGE_REF = re.compile(r"^[A-Za-z0-9._/-]+:[A-Za-z0-9._-]+$") +_GIT_SHA = re.compile(r"^[0-9a-f]{40}$") + + +def backend_registry(path: str | None = None) -> None: + """Validate the public backend registry (configs/backends.json) and emit the + COLLX_* source-pin and image names consumed by common.sh at source time. + Unlike the network overlay this file is required: a missing or malformed + registry fails closed rather than falling back.""" + if path is None: + path = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + "configs", "backends.json", + ) + try: + with open(path, encoding="utf-8") as stream: + document = json.load(stream) + def sha(value: object) -> str: + if not isinstance(value, str) or not _GIT_SHA.fullmatch(value): + raise ValueError + return value + + def image(name: str) -> str: + ref = document["images"][name]["ref"] + if not isinstance(ref, str) or not _IMAGE_REF.fullmatch(ref): + raise ValueError + return ref + + def repo(backend: dict) -> str: + url = backend["repo"] + if not isinstance(url, str) or not url.startswith("https://"): + raise ValueError + return url + + v2 = document["backends"]["deepep-v2"] + emit({ + "COLLX_IMAGE_MULTIARCH": image("multiarch"), + "COLLX_IMAGE_AMD_MORI": image("amd-mori"), + "COLLX_DEEPEP_V2_REPO": repo(v2), + "COLLX_DEEPEP_V2_COMMIT": sha(v2["commit"]), + }) + except (IndexError, KeyError, OSError, TypeError, ValueError, json.JSONDecodeError): + print("validation-invalid-backend-registry", file=sys.stderr) + raise SystemExit(1) + + +def load(path: str) -> dict: + with open(path, encoding="utf-8") as stream: + return json.load(stream) + + +def case_count(path: str) -> None: + print(len(load(path)["cases"]), end="") + + +def _emit_argv(case: dict, version: object, runner: str, ts: str, index: int) -> None: + """Emit one null-delimited run_ep.py argv — the only case-to-invocation codec.""" + get = lambda key, default="": str(case.get(key) or default) + argv = [ + "--backend", get("backend"), + "--mode", str(case["mode"]), + "--phase", str(case["phase"]), + "--routing", str(case["routing"]), + "--gpus-per-node", get("gpus_per_node", "0"), + "--scale-up-domain", get("scale_up_domain", "0"), + "--scope", get("scope", "scale-up"), + "--scale-up-transport", get("scale_up_transport", "unknown"), + "--scale-out-transport", get("scale_out_transport"), + "--tokens-ladder", get("ladder"), + "--hidden", get("hidden"), + "--topk", get("topk"), + "--experts", get("experts"), + "--seed", get("seed"), + "--runner", runner, + "--topology-class", get("topology_class"), + "--transport", get("transport", "unknown"), + "--case-id", get("case_id"), + "--suite", get("suite"), + "--workload-name", get("workload"), + "--version", str(version), + ] + iters, trials, warmup = (get("timing").split(":") + ["", "", ""])[:3] + for flag, value in (("--iters", iters), ("--trials", trials), ("--warmup", warmup)): + if value: + argv += [flag, value] + out = f"results/{runner}_{get('backend')}_{case['phase']}_{ts}-c{index:03d}.json" + argv += ["--out", out] + sys.stdout.buffer.write(b"\0".join(part.encode() for part in argv) + b"\0") + + +def case_args( + path: str, index: int, runner: str, ts: str, + ngpus: str, nodes: str, gpus_per_node: str, scale_up_domain: str, +) -> None: + document = load(path) + cases = document["cases"] + if not 0 <= index < len(cases): + raise SystemExit(1) + case = cases[index] + placement = tuple( + str(case.get(field, "")) + for field in ("ep", "nodes", "gpus_per_node", "scale_up_domain") + ) + if placement != (ngpus, nodes, gpus_per_node, scale_up_domain): + print(f"case placement {placement} differs from the allocation", file=sys.stderr) + raise SystemExit(1) + _emit_argv(case, document["version"], runner, ts, index) + + +def main() -> None: + parser = argparse.ArgumentParser() + commands = parser.add_subparsers(dest="command", required=True) + for name, names in { + "operator-config": ("path", "runner"), "merge-operator-config": ("path",), + "backend-registry": (), "case-count": ("path",), + "case-args": ("path", "index", "runner", "ts", + "ngpus", "nodes", "gpus_per_node", "scale_up_domain"), + }.items(): + command = commands.add_parser(name) + for arg in names: command.add_argument(arg) + args = parser.parse_args() + if args.command == "operator-config": operator_config(args.path, args.runner) + elif args.command == "merge-operator-config": merge_operator_config(args.path) + elif args.command == "backend-registry": backend_registry() + elif args.command == "case-count": case_count(args.path) + elif args.command == "case-args": + case_args(args.path, int(args.index), args.runner, args.ts, + args.ngpus, args.nodes, args.gpus_per_node, args.scale_up_domain) + + +if __name__ == "__main__": + main() diff --git a/experimental/CollectiveX/runtime/prepare_backend.sh b/experimental/CollectiveX/runtime/prepare_backend.sh new file mode 100644 index 0000000000..a8514b5429 --- /dev/null +++ b/experimental/CollectiveX/runtime/prepare_backend.sh @@ -0,0 +1,332 @@ +#!/usr/bin/env bash +# Prepare one backend per allocated node and persist its rank environment. +set -euo pipefail + +cd /ix/experimental/CollectiveX +# shellcheck source=../runtime/common.sh +source runtime/common.sh + +: "${COLLX_RUNNER:?COLLX_RUNNER not set}" +: "${COLLX_BENCH:?COLLX_BENCH not set}" + +collx_log "backend preparation: runner=$COLLX_RUNNER bench=$COLLX_BENCH nodes=${COLLX_NODES:-1}" + +# Match the CUDA target to the registered SKU before compiling. +collx_cuda_arch() { + local expected detected + expected="$(python3 - "$COLLX_RUNNER" <<'PY' +import json, sys +arch = json.load(open("configs/platform_config.json"))["platforms"][sys.argv[1]]["arch"] +digits = arch.removeprefix("sm") +print(f"{digits[:-1]}.{digits[-1]}" if arch.startswith("sm") and digits.isdigit() else "") +PY +)" || { collx_log "ERROR: no platform registry entry for $COLLX_RUNNER"; return 1; } + [ -n "$expected" ] || { + collx_log "ERROR: no CUDA target registered for $COLLX_RUNNER"; return 1 + } + detected="$(python3 - <<'PY' +import torch + +major, minor = torch.cuda.get_device_capability() +print(f"{major}.{minor}") +PY +)" || return 1 + [ "$detected" = "$expected" ] || { + collx_log "ERROR: $COLLX_RUNNER expected CUDA target $expected, detected $detected" + return 1 + } + printf '%s' "$detected" +} + +collx_nvidia_package_root() { + local package="$1" component="$2" + python3 - "$package" "$component" <<'PY' +from importlib import metadata +from pathlib import Path, PurePosixPath +import sys + +package, component = sys.argv[1:] +try: + distribution = metadata.distribution(package) + prefix = f"nvidia/{component}/" + entries = [str(entry).replace("\\", "/") for entry in distribution.files or ()] + if not any(entry.startswith(prefix) for entry in entries): + raise ValueError + root = Path(distribution.locate_file(PurePosixPath("nvidia") / component)).resolve() + if not root.is_dir(): + raise ValueError +except (metadata.PackageNotFoundError, OSError, TypeError, ValueError): + raise SystemExit(1) +print(root, end="") +PY +} + +collx_prepare_cuda_cccl() { + local cccl="" candidate cuda_home nvcc + nvcc="$(command -v nvcc)" \ + || { collx_log "ERROR: CUDA nvcc is unavailable"; return 1; } + nvcc="$(readlink -f -- "$nvcc")" \ + || { collx_log "ERROR: CUDA nvcc cannot be resolved"; return 1; } + case "$nvcc" in + */bin/nvcc) cuda_home="${nvcc%/bin/nvcc}" ;; + *) collx_log "ERROR: CUDA nvcc has an unexpected path"; return 1 ;; + esac + [ -x "$cuda_home/bin/nvcc" ] && [ -d "$cuda_home/include" ] \ + && [ -d "$cuda_home/lib64" ] \ + || { collx_log "ERROR: CUDA toolkit root is incomplete"; return 1; } + for candidate in "$cuda_home"/targets/*/include/cccl; do + if [ -d "$candidate" ]; then + cccl="$candidate" + break + fi + done + [ -n "$cccl" ] || { collx_log "ERROR: CUDA CCCL headers are unavailable"; return 1; } + export CUDA_HOME="$cuda_home" COLLX_CUDA_CCCL="$cccl" + export CPATH="$cccl:${CPATH:-}" + export NVCC_PREPEND_FLAGS="-I$cccl ${NVCC_PREPEND_FLAGS:-}" +} + +collx_prepare_deepep_toolchain() { + local packaged overlay path root temporary + packaged="$(collx_nvidia_package_root nvidia-nvshmem-cu12 nvshmem)" \ + || { collx_log "ERROR: nvidia.nvshmem is unavailable"; return 1; } + root="$(collx_deepep_v2_root)" || return 1 + overlay="$root/nvshmem-overlay" + if ! ( + umask 077 + exec 8>"$root/nvshmem-overlay.lock" || exit 1 + flock 8 || exit 1 + if [ ! -d "$overlay" ]; then + temporary="$root/.nvshmem-overlay.$$" + rm -rf "$temporary" || exit 1 + mkdir -p "$temporary/lib" || exit 1 + ln -s "$packaged/include" "$temporary/include" || exit 1 + for path in "$packaged"/lib/*; do + ln -s "$path" "$temporary/lib/${path##*/}" || exit 1 + done + [ ! -e "$packaged/lib/libnvshmem_host.so.3" ] \ + || ln -sf "$packaged/lib/libnvshmem_host.so.3" \ + "$temporary/lib/libnvshmem_host.so" || exit 1 + mv "$temporary" "$overlay" || exit 1 + fi + [ ! -L "$overlay" ] \ + && [ "$(readlink -f "$overlay/include")" = "$(readlink -f "$packaged/include")" ] \ + && [ -e "$overlay/lib/libnvshmem_host.so" ] \ + && [ -e "$overlay/lib/libnvshmem_device.a" ] + ); then + collx_log "ERROR: DeepEP V2 NVSHMEM overlay is invalid" + return 1 + fi + NVSHMEM_DIR="$overlay" + export NVSHMEM_DIR + collx_prepare_cuda_cccl || return 1 + export LD_LIBRARY_PATH="$NVSHMEM_DIR/lib:${LD_LIBRARY_PATH:-}" +} + +collx_deepep_v2_root() { + local arch cpu base image + arch="$(collx_cuda_arch)" || return 1 + cpu="$(uname -m)" + [[ "$cpu" =~ ^[A-Za-z0-9._-]+$ ]] || return 1 + base="${COLLX_BACKEND_CACHE_ROOT:-}" + [[ "$base" = /* ]] || return 1 + image="$(printf '%s' "${COLLECTIVEX_IMAGE:-manual}" | tr -cs 'A-Za-z0-9_.-' '-')" + printf '%s/deepep-v2-%s-sm%s-%s-%s' \ + "$base" "$cpu" "${arch/./}" "${image#-}" "${COLLX_DEEPEP_V2_COMMIT:0:12}" +} + +collx_activate_deepep_v2() { + local root venv venv_site execution_id + root="$(collx_deepep_v2_root)" || return 1 + venv="$root/venv" + [ -x "$venv/bin/python" ] \ + || { collx_log "ERROR: DeepEP V2 venv interpreter is unavailable"; return 1; } + export VIRTUAL_ENV="$venv" + export PATH="$venv/bin:${PATH#"$venv/bin:"}" + # Ensure the source build wins over any deep_ep wheel bundled in the image. + for venv_site in "$venv"/lib/python*/site-packages; do break; done + [ -d "$venv_site" ] \ + || { collx_log "ERROR: DeepEP V2 venv site-packages is unavailable"; return 1; } + export PYTHONPATH="$venv_site${PYTHONPATH:+:$PYTHONPATH}" + EP_NCCL_ROOT_DIR="$(collx_nvidia_package_root nvidia-nccl-cu13 nccl)" \ + || { collx_log "ERROR: DeepEP V2 NCCL package root is unavailable"; return 1; } + EP_NVSHMEM_ROOT_DIR="$(collx_nvidia_package_root nvidia-nvshmem-cu12 nvshmem)" \ + || { collx_log "ERROR: DeepEP V2 NVSHMEM package root is unavailable"; return 1; } + export EP_NCCL_ROOT_DIR EP_NVSHMEM_ROOT_DIR + export LD_LIBRARY_PATH="$EP_NCCL_ROOT_DIR/lib:$EP_NVSHMEM_ROOT_DIR/lib:${LD_LIBRARY_PATH:-}" + execution_id="${COLLECTIVEX_EXECUTION_ID:-manual}" + [[ "$execution_id" =~ ^[A-Za-z0-9._-]+$ ]] \ + || { collx_log "ERROR: DeepEP V2 execution identity is invalid"; return 1; } + # Shared JIT caches race across nodes; keep this cache node-local. + export EP_JIT_CACHE_DIR="/tmp/collectivex-deepep-v2-jit-$execution_id" + export EP_REUSE_NCCL_COMM=1 + # ElasticBuffer requires NCCL symmetric memory, which needs cuMem. Some + # provider image variants bake NCCL_CUMEM_ENABLE=0 (h100-dgxc's does), and + # pyxis lets image env override srun-passed env, so a launcher-side export + # never reaches the ranks. Exporting it here puts it in the persisted node + # environment, which ranks source inside the container after image env. + export NCCL_CUMEM_ENABLE=1 + [ ! -L "$EP_JIT_CACHE_DIR" ] \ + || { collx_log "ERROR: DeepEP V2 JIT cache path is unsafe"; return 1; } + if ! mkdir -p "$EP_JIT_CACHE_DIR" || ! chmod 700 "$EP_JIT_CACHE_DIR"; then + collx_log "ERROR: DeepEP V2 JIT cache is unavailable" + return 1 + fi + unset EP_SUPPRESS_NCCL_CHECK +} + +# Catch an image-bundled DeepEP shadowing the from-source build: the shadow +# (deep_ep 1.2.1) lacks ElasticBuffer, so its presence is the capability check. +collx_probe_deepep_v2() { + python3 - <<'PY' +import inspect + +import deep_ep + +assert inspect.isclass(deep_ep.ElasticBuffer) +PY +} + +collx_build_deepep_v2() { + local root venv source ready lock_path arch cache_ready + local revision="$COLLX_DEEPEP_V2_COMMIT" + arch="$(collx_cuda_arch)" || return 1 + root="$(collx_deepep_v2_root)" || return 1 + venv="$root/venv"; source="$root/source"; ready="$root/.ready" + lock_path="${root}.lock" + command -v flock >/dev/null || { collx_log "ERROR: flock is required for DeepEP V2"; return 1; } + mkdir -p "${root%/*}" || return 1 + collx_log "DeepEP V2: preparing PR #605 with upstream PR #630 and #640 fixes ($revision)" + if ! ( + [ ! -L "$lock_path" ] \ + || { collx_log "ERROR: DeepEP V2 cache lock is unsafe"; exit 1; } + (umask 077; : >> "$lock_path") && chmod 600 "$lock_path" \ + || { collx_log "ERROR: DeepEP V2 cache-lock-create failed"; exit 1; } + exec 9<>"$lock_path" \ + || { collx_log "ERROR: DeepEP V2 cache-lock-open failed"; exit 1; } + flock 9 \ + || { collx_log "ERROR: DeepEP V2 cache-lock-acquire failed"; exit 1; } + cache_ready=0 + [ -f "$ready" ] && [ -x "$venv/bin/python" ] && [ -d "$source" ] && cache_ready=1 + if [ "$cache_ready" != 1 ]; then + if [ -e "$root" ] || [ -L "$root" ]; then + rm -rf "$root" \ + || { collx_log "ERROR: incomplete DeepEP V2 cache-reset failed"; exit 1; } + fi + mkdir -m 700 "$root" \ + || { collx_log "ERROR: DeepEP V2 cache-create failed"; exit 1; } + python3 -m venv "$venv" \ + || { collx_log "ERROR: DeepEP V2 venv creation failed"; exit 1; } + "$venv/bin/python" -m pip install -q --disable-pip-version-check --no-input \ + "pip==26.1.2" "setuptools==82.0.1" "wheel==0.47.0" "ninja==1.13.0" \ + "numpy==2.2.6" "nvidia-nvshmem-cu12==3.3.9" >&2 2>&1 \ + || { collx_log "ERROR: DeepEP V2 build-tool installation failed"; exit 1; } + "$venv/bin/python" -m pip install -q --disable-pip-version-check --no-input \ + --index-url https://download.pytorch.org/whl/cu130 \ + --extra-index-url https://pypi.org/simple "torch==2.10.0" >&2 2>&1 \ + || { collx_log "ERROR: torch 2.10.0+cu130 installation failed"; exit 1; } + # Torch pins NCCL 2.28.9; the PR #605 ElasticBuffer implementation requires 2.30.4. + "$venv/bin/python" -m pip install -q --disable-pip-version-check --no-input \ + --force-reinstall --no-deps "nvidia-nccl-cu13==2.30.4" >&2 2>&1 \ + || { collx_log "ERROR: NCCL 2.30.4 installation failed"; exit 1; } + collx_activate_deepep_v2 \ + || { collx_log "ERROR: DeepEP V2 environment activation failed"; exit 1; } + collx_prepare_deepep_toolchain \ + || { collx_log "ERROR: DeepEP V2 toolchain preparation failed"; exit 1; } + EP_NVSHMEM_ROOT_DIR="$NVSHMEM_DIR" + export EP_NVSHMEM_ROOT_DIR + collx_materialize_deepep_source "$source" \ + || { collx_log "ERROR: DeepEP V2 staged source is invalid"; exit 1; } + (cd "$source" && TORCH_CUDA_ARCH_LIST="$arch" MAX_JOBS=16 \ + python3 -m pip install -q --no-build-isolation --no-deps --force-reinstall .) >&2 2>&1 \ + || { collx_log "ERROR: DeepEP V2 build failed"; exit 1; } + collx_probe_deepep_v2 \ + || { collx_log "ERROR: DeepEP V2 import probe failed"; exit 1; } + : > "$ready" + fi + ); then + collx_log "ERROR: shared DeepEP V2 environment is incomplete" + return 1 + fi + collx_activate_deepep_v2 || return 1 + collx_prepare_deepep_toolchain || return 1 + EP_NVSHMEM_ROOT_DIR="$NVSHMEM_DIR" + export EP_NVSHMEM_ROOT_DIR + collx_probe_deepep_v2 || { collx_log "ERROR: DeepEP V2 shared runtime probe failed"; return 1; } + collx_log "DeepEP V2 ready ($COLLX_DEEPEP_V2_COMMIT, ElasticBuffer, NCCL Device API; LSA/Gin selected by adapter)" +} + +# Rank steps enter fresh container tasks, so persist each node's environment. +collx_persist_backend_env() { + local root="$PWD/.collx_backend/env" node_id="${SLURM_NODEID:-0}" path temporary name + local -a names=(PATH VIRTUAL_ENV LD_LIBRARY_PATH PYTHONPATH CUDA_HOME CPATH NVCC_PREPEND_FLAGS + NVSHMEM_DIR + EP_NCCL_ROOT_DIR EP_NVSHMEM_ROOT_DIR EP_JIT_CACHE_DIR EP_REUSE_NCCL_COMM + NCCL_CUMEM_ENABLE) + [[ "$node_id" =~ ^[0-9]+$ ]] || return 1 + mkdir -p "$root" || return 1 + chmod 700 "$root" || return 1 + temporary="$(mktemp "$root/.node-${node_id}.XXXXXX")" || return 1 + chmod 600 "$temporary" || { rm -f "$temporary"; return 1; } + for name in "${names[@]}"; do + if declare -p "$name" >/dev/null 2>&1; then + printf 'export %s=%q\n' "$name" "${!name}" >> "$temporary" \ + || { rm -f "$temporary"; return 1; } + fi + done + path="$root/node-${node_id}.sh" + mv -f -- "$temporary" "$path" || { rm -f "$temporary"; return 1; } +} + +# Validate private scale-out selectors on every allocated compute node before a +# backend can initialize or build transport code. +collx_probe_scaleout_network() { + local interface device rdma_name + local -a interfaces devices + if [ "${COLLX_NODES:-1}" -le 1 ] || [ "${COLLX_TRANSPORT:-}" = mnnvl ]; then + return 0 + fi + collx_restore_exact_hca_selector || return 1 + [ -n "${GLOO_SOCKET_IFNAME:-}" ] && [ -n "${NCCL_IB_HCA:-}" ] \ + || { collx_log "ERROR: scale-out network selectors are unavailable"; return 1; } + IFS=, read -r -a interfaces <<< "$GLOO_SOCKET_IFNAME" + for interface in "${interfaces[@]}"; do + [ -d "/sys/class/net/$interface" ] \ + || { collx_log "ERROR: configured scale-out socket interface is absent"; return 1; } + done + IFS=, read -r -a devices <<< "$NCCL_IB_HCA" + for device in "${devices[@]}"; do + device="${device#=}" + rdma_name="${device%%:*}" + [ -d "/sys/class/infiniband/$rdma_name" ] \ + || { collx_log "ERROR: configured scale-out RDMA device is absent"; return 1; } + done +} + +# Prepare and probe one backend without running a benchmark. +collx_prepare_backend() { + local backend="${1:-}" + case "$backend" in + deepep-v2) + collx_build_deepep_v2 || return 1 + ;; + mori) + python3 -c "import mori" \ + || { collx_log "ERROR: MoRI backend import failed"; return 1; } + ;; + *) + collx_log "ERROR: unknown backend preparation request" + return 1 + ;; + esac +} + +rc=0 +collx_apply_network_profile "${COLLX_NODES:-1}" "${COLLX_TRANSPORT:-}" +if collx_probe_scaleout_network && collx_prepare_backend "$COLLX_BENCH"; then + collx_persist_backend_env || rc=1 +else + rc=1 +fi +collx_log "backend preparation: bench=$COLLX_BENCH rc=$rc" +exit "$rc" diff --git a/experimental/CollectiveX/runtime/probe.py b/experimental/CollectiveX/runtime/probe.py new file mode 100644 index 0000000000..922a99682e --- /dev/null +++ b/experimental/CollectiveX/runtime/probe.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python3 +"""Allocation and network checks used by CollectiveX launchers.""" + +from __future__ import annotations + +import argparse +import ctypes +import os +from pathlib import Path + + +def default_route_interface(route_path: Path = Path("/proc/net/route")) -> str: + for line in route_path.read_text().splitlines()[1:]: + fields = line.split() + if len(fields) >= 4 and fields[1] == "00000000" and int(fields[3], 16) & 1: + return fields[0] + return "" + + +def prepare_cache(parent_path: str) -> str: + path = Path(parent_path).resolve() / f".collectivex-backend-cache-{os.getuid()}" + path.mkdir(mode=0o700, exist_ok=True) + os.chmod(path, 0o700) + return str(path) + + +def validate_cuda_context(expected: int) -> None: + cuda = ctypes.CDLL("libcuda.so.1") + count = ctypes.c_int() + if cuda.cuInit(0) != 0 or cuda.cuDeviceGetCount(ctypes.byref(count)) != 0 or count.value != expected: + raise SystemExit(1) + + +def _emit(marker: str) -> None: + # collx_validate_network_profile_on_job (runtime/common.sh) greps these exact strings + # out of the per-node probe log to derive COLLX_SOCKET_IFNAME / COLLX_RDMA_LINK_LAYER and to + # diagnose failures. The marker vocabulary is a string contract with that function — + # keep the two halves in lockstep (see tests/test_runtime.py::NetworkProfileContract). + print(f"[collectivex-private] {marker}") + + +def _check_port(port_path: Path, ordinal: int, gid_index: str, profile: str): + # Return the port's link layer ("roce"/"infiniband") when it is active, carries a + # non-empty GID at the pinned index, and agrees with any already-seen link layer; + # otherwise emit the matching rdma-port-= marker and return None. + if not port_path.is_dir(): + _emit(f"rdma-port-{ordinal}=missing"); return None + state = port_path / "state" + if not state.is_file() or state.read_text().split()[:1] != ["4:"]: + _emit(f"rdma-port-{ordinal}=inactive"); return None + if gid_index: + gid = port_path / "gids" / gid_index + if not gid.is_file(): + _emit(f"rdma-port-{ordinal}=gid-missing"); return None + if not "".join(c for c in gid.read_text() if c not in ":0" and not c.isspace()): + _emit(f"rdma-port-{ordinal}=gid-empty"); return None + link = port_path / "link_layer" + if not link.is_file(): + _emit(f"rdma-port-{ordinal}=link-layer-missing"); return None + layer = {"Ethernet": "roce", "InfiniBand": "infiniband"}.get(link.read_text().strip()) + if layer is None: + _emit(f"rdma-port-{ordinal}=link-layer-invalid"); return None + if profile and profile != layer: + _emit(f"rdma-port-{ordinal}=link-layer-mixed"); return None + return layer + + +def validate_network_profile(socket_names: str, rdma_devices: str, gid_index: str, + sys_root: Path = Path("/sys"), + route_path: Path = Path("/proc/net/route")) -> None: + # Prove the operator-pinned scale-out fabric on this node: resolve the cross-node socket + # interface (operator selector, else this node's default route), confirm it is live, and + # confirm every pinned RDMA port is active with a consistent link layer. On success emit + # the socket-interface-selected and rdma-link-layer markers the launcher consumes; on any + # failure emit the diagnostic marker and exit non-zero. The seam carries exactly ONE socket + # interface: the launcher's marker-extraction regex has no comma, so a multi-interface + # selector could never survive past this probe — fail it loudly here instead. + interface = socket_names or default_route_interface(route_path) + if not interface: + _emit("socket-interface-1=default-route-missing") + raise SystemExit(1) + _emit(f"socket-interface-selected={interface}") + net = sys_root / "class" / "net" / interface + if not net.is_dir(): + _emit("socket-interface-1=missing"); raise SystemExit(1) + operstate = net / "operstate" + state = operstate.read_text().strip() if operstate.is_file() else "" + if state not in ("up", "unknown"): + _emit("socket-interface-1=down"); raise SystemExit(1) + profile = "" + for ordinal, selector in enumerate((s for s in rdma_devices.split(",") if s), start=1): + device, _, configured_port = selector.partition(":") + ports = sys_root / "class" / "infiniband" / device / "ports" + if not ports.is_dir(): + _emit(f"rdma-device-{ordinal}=missing"); raise SystemExit(1) + if configured_port: + layer = _check_port(ports / configured_port, ordinal, gid_index, profile) + if layer is None: raise SystemExit(1) + profile = layer + else: + active = False + for port_path in sorted(p for p in ports.iterdir() if p.is_dir()): + layer = _check_port(port_path, ordinal, gid_index, profile) + if layer is not None: + profile, active = layer, True + if not active: raise SystemExit(1) + if not profile: raise SystemExit(1) + _emit(f"rdma-link-layer={profile}") + + +def main() -> None: + parser = argparse.ArgumentParser(); commands = parser.add_subparsers(dest="command", required=True) + commands.add_parser("default-route-interface") + command = commands.add_parser("prepare-cache"); command.add_argument("parent") + command = commands.add_parser("cuda-context"); command.add_argument("expected", type=int) + command = commands.add_parser("network-profile"); command.add_argument("socket_names"); command.add_argument("rdma_devices"); command.add_argument("gid_index") + args = parser.parse_args() + if args.command == "default-route-interface": print(default_route_interface(), end="") + elif args.command == "prepare-cache": print(prepare_cache(args.parent), end="") + elif args.command == "cuda-context": validate_cuda_context(args.expected) + else: validate_network_profile(args.socket_names, args.rdma_devices, args.gid_index) + + +if __name__ == "__main__": main() diff --git a/experimental/CollectiveX/runtime/stage.py b/experimental/CollectiveX/runtime/stage.py new file mode 100644 index 0000000000..119b77496f --- /dev/null +++ b/experimental/CollectiveX/runtime/stage.py @@ -0,0 +1,103 @@ +#!/usr/bin/env python3 +"""Create, copy, and clean isolated CollectiveX workspaces.""" + +from __future__ import annotations + +import argparse +import os +import pwd +from pathlib import Path +import shutil + + +EXCLUDES = {"__pycache__", "results", ".shards", ".collx_workloads", ".collx_backend", + ".collx_sources", ".venv", ".pytest_cache", "private-infra.md", "goal.md", + "notes.md"} + + +def implicit_stage_base(args) -> None: + # Resolve the account home from /etc/passwd, not $HOME. The GHA launcher deliberately + # points $HOME at a runner-local /tmp sandbox. The passwd home is compute-visible. + base = args.home or pwd.getpwuid(os.getuid()).pw_dir + home = Path(base).resolve() + suffix = "" + if args.isolation_key: + if not all(char.isalnum() or char in "._-" for char in args.isolation_key): + raise SystemExit(1) + suffix = "-" + args.isolation_key + path = home / f".inferencex-collectivex-stage{suffix}" + path.mkdir(mode=0o700, exist_ok=True) + print(path, end="") + + +def resolve_directory(args) -> None: + path = Path(args.path).resolve() + if not path.is_dir(): raise SystemExit(1) + print(path, end="") + + +def validate_stage_path(args) -> None: + base, child = Path(args.base).resolve(), Path(args.child) + if child.parent.resolve() != base or child.exists() or base == Path("/"): + raise SystemExit(1) + for excluded in (args.repo, args.job_root, args.workspace): + if excluded and base == Path(excluded).resolve(): raise SystemExit(1) + print(child, end="") + + +def create_stage(args) -> None: + Path(args.stage).mkdir(mode=0o700) + + +def copy_repository(args) -> None: + source, target = Path(args.source), Path(args.target) + shutil.copytree(source, target, ignore=shutil.ignore_patterns(*EXCLUDES), dirs_exist_ok=False) + + +def validate_cleanup(args) -> None: + root = Path(args.root) + if not root.is_dir() or root.is_symlink() or root == Path("/"): + raise SystemExit(1) + + +def rewrite_deepep_v2(args) -> None: + path = Path(args.path) + old = "for so in [line.strip().split(' ')[-1] for line in f if 'nccl' in line]:" + new = "for so in [line.strip().split(' ')[-1] for line in f if 'libnccl' in line]:" + text = path.read_text() + if text.count(old) != 1: raise SystemExit(1) + path.write_text(text.replace(old, new)) + + +# The runtime/common.sh launcher shells out to these subcommands by literal name and +# positional argv; there are no optional flags. That argv shape is a string contract with +# common.sh — a subcommand or flag common.sh passes but this parser does not declare fails +# with "unrecognized arguments" and aborts the leg at repository-stage. Keep the two halves in +# lockstep (see tests/test_runtime.py::StageContract), which is exactly the contract that broke +# when the --allow-* flags were dropped here but left on the common.sh callers. +SPECS = { + "implicit-stage-base": (("home", "?"), ("isolation_key", "?")), + "resolve-directory": (("path",),), + "validate-stage-path": (("repo",), ("base",), ("child",), ("job_root", "?"), ("workspace", "?")), + "create-stage": (("stage",),), "copy-repository": (("source",), ("target",)), + "validate-cleanup": (("root",),), "rewrite-deepep-v2": (("path",),), +} + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser() + commands = parser.add_subparsers(dest="command", required=True) + handlers = globals() + for name, arguments in SPECS.items(): + command = commands.add_parser(name) + for item in arguments: + command.add_argument(item[0], nargs=item[1] if len(item) > 1 else None, default="") + command.set_defaults(handler=handlers[name.replace("-", "_")]) + return parser + + +def main() -> None: + args = build_parser().parse_args(); args.handler(args) + + +if __name__ == "__main__": main() diff --git a/experimental/CollectiveX/summarize.py b/experimental/CollectiveX/summarize.py new file mode 100644 index 0000000000..85c7048e29 --- /dev/null +++ b/experimental/CollectiveX/summarize.py @@ -0,0 +1,92 @@ +#!/usr/bin/env python3 +"""Render a small shard summary; benchmark gating remains in the harness.""" +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +# Emitted case-attempt documents this summary reads, discriminated by record_type. +# This is a best-effort renderer over whatever raw attempts a shard produced; it +# validates nothing. +CASE_RECORD_TYPE = "case-attempt" + + +def load_results(directory: str, runner: str | None, timestamp: str | None) -> list[dict]: + documents: list[dict] = [] + for path in sorted(Path(directory).glob("*.json")): + if runner and not path.name.startswith(f"{runner}_"): + continue + if timestamp and timestamp not in path.name: + continue + try: + with path.open() as handle: + document = json.load(handle) + except (OSError, ValueError): + continue + if isinstance(document, dict) and document.get("record_type") == CASE_RECORD_TYPE: + documents.append(document) + return documents + + +def _identity(document: dict) -> tuple[str, str, str, str, int]: + factors = document["identity"]["case_factors"] + case = factors["case"] + return ( + factors["sku"], case["suite"], case["routing"], case["phase"], case["ep"], + ) + + +def _headline(document: dict) -> tuple[int | str, float | str, float | str]: + rows = document["measurement"]["rows"] + row = next((item for item in rows if item["tokens_per_rank"] == 64), rows[len(rows) // 2]) + latency = row["components"]["roundtrip"]["percentiles_us"] + return row["tokens_per_rank"], latency["p50"], latency["p99"] + + +def render(documents: list[dict]) -> str: + documents = sorted(documents, key=_identity) + invalid = [d for d in documents if d["outcome"]["status"] != "success"] + lines = ["## CollectiveX EP results", ""] + if invalid: + # The leg is already red (ep_harness.run_sweep returns nonzero on a non-success + # outcome); call the count out loudly so it is not lost in the per-row table. + lines.append( + f"> **{len(invalid)} of {len(documents)} outcome(s) INVALID** — " + "the leg fails; see the outcome column below." + ) + lines.append("") + lines += [ + "| ver | sku | backend | suite | phase | routing | ep | outcome | T* | p50 us | p99 us |", + "|--:|---|---|---|---|---|--:|---|--:|--:|--:|", + ] + for document in documents: + sku, suite, routing, phase, ep = _identity(document) + backend = document["identity"]["case_factors"]["case"]["backend"] + token, p50, p99 = _headline(document) + lines.append( + f"| {document['version']} | {sku} | `{backend}` | {suite} | {phase} | " + f"{routing} | {ep} | " + f"{document['outcome']['status']} | {token} | {p50} | {p99} |" + ) + if not documents: + lines.append("\n> No valid native outcome documents found.") + return "\n".join(lines) + + +def main() -> int: + parser = argparse.ArgumentParser(description="Summarize CollectiveX native v1 outcomes") + parser.add_argument("--results-dir", default="results") + parser.add_argument("--runner") + parser.add_argument("--ts") + args = parser.parse_args() + documents = load_results(args.results_dir, args.runner, args.ts) + print(render(documents)) + # Pure renderer — never gates CI. The per-case leg gate lives in ep_harness.run_sweep: a + # non-success outcome returns nonzero and fails the shard (see collx_run_shard). A dead + # "exit 1 when no success doc" gate here lost its only caller in 41caeaa0. + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/experimental/CollectiveX/sweep_matrix.py b/experimental/CollectiveX/sweep_matrix.py new file mode 100644 index 0000000000..87ffc7d5ec --- /dev/null +++ b/experimental/CollectiveX/sweep_matrix.py @@ -0,0 +1,373 @@ +#!/usr/bin/env python3 +"""Resolve CollectiveX suites and extract execution shards. + +Mode changes measurement semantics and therefore participates in case identity. +Dispatch and combine are fixed BF16 benchmark facts, so a case's coordinates are +suite/workload/backend/topology only; the matrix schedules, it never ranks. +""" +from __future__ import annotations + +import argparse +import itertools +import json +from pathlib import Path +import sys +from typing import Any + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) +sys.path.insert(0, str(HERE / "bench")) + +try: # Shard extraction on GPU runners is intentionally stdlib-only. + import yaml # type: ignore +except ModuleNotFoundError: # pragma: no cover - exercised by the workflow environment + yaml = None + +import capability as cap # noqa: E402 +import ep_harness # noqa: E402 + + +TOPOLOGY_FIELDS = ( + "nodes", "gpus_per_node", "scale_up_domain", "scope", "scale_up_transport", + "scale_out_transport", "transport", "topology_class", +) + + +if yaml is not None: + class _UniqueKeyLoader(yaml.SafeLoader): + pass + + def _unique_mapping(loader: Any, node: Any, deep: bool = False) -> dict[Any, Any]: + result: dict[Any, Any] = {} + for key_node, value_node in node.value: + key = loader.construct_object(key_node, deep=deep) + if key in result: + raise SystemExit(f"duplicate YAML key {key!r} at line {key_node.start_mark.line + 1}") + result[key] = loader.construct_object(value_node, deep=deep) + return result + + _UniqueKeyLoader.add_constructor( + yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, _unique_mapping + ) + + +def _load(name: str) -> dict[str, Any]: + if yaml is None: + raise SystemExit("matrix generation requires PyYAML; shard extraction does not") + try: + with (HERE / "configs" / name).open() as fh: + document = yaml.load(fh, Loader=_UniqueKeyLoader) + except yaml.YAMLError as exc: + raise SystemExit(f"configs/{name} is not valid YAML: {exc}") from exc + if not isinstance(document, dict): + raise SystemExit(f"configs/{name} must contain a YAML object") + return document + + +def _positive_int(value: Any) -> bool: + return not isinstance(value, bool) and isinstance(value, int) and value > 0 + + +def _dims(workloads: dict[str, Any], name: str) -> tuple[int, int, int, int]: + config = workloads[name] + seed = config["seed"] + # Positive (not just non-negative): the case-argv codec treats falsy fields + # as absent, so a seed of 0 would silently defer to the invocation seed. + if not _positive_int(seed): + raise SystemExit(f"workload {name!r} has no valid routing seed: {seed!r}") + return config["hidden"], config["topk"], config["routed_experts"], seed + + +def _timing(document: dict[str, Any]) -> tuple[int, int, int]: + keys = ("iters_per_trial", "trials_per_point", "warmup_iters_per_trial") + timing = document.get("timing") + if (not isinstance(timing, dict) or set(timing) != set(keys) + or not all(_positive_int(timing[key]) for key in keys)): + raise SystemExit(f"suites.yaml timing must define positive ints {keys}: {timing!r}") + return timing[keys[0]], timing[keys[1]], timing[keys[2]] + + +def _ladder(workloads: dict[str, Any], workload: str, phase: str) -> str: + points = workloads[workload].get("token_ladders", {}).get(phase) + if (not isinstance(points, list) or not points + or not all(_positive_int(point) for point in points) + or points != sorted(set(points))): + raise SystemExit( + f"workload {workload!r} has no valid {phase} token ladder: {points!r}" + ) + return " ".join(map(str, points)) + + +def _select_backends(backend: str, backends: str) -> list[str]: + available = list(cap.SWEEP_BACKENDS) + if backend and backends: + raise SystemExit("--backend and --backends are mutually exclusive") + if backends: + names = available if backends == "all" else [ + value.strip() for value in backends.split(",") if value.strip() + ] + else: + names = [backend or "deepep-v2"] + unknown = sorted(set(names) - set(available)) + if unknown: + raise SystemExit(f"unknown backend values {unknown}; have {available}") + if len(names) != len(set(names)): + raise SystemExit("backend selection contains duplicates") + return names + + +def resolve_matrix( + suites: str = "all", + backend: str = "", + backends: str = "", + only_sku: str = "", + exclude_skus: str = "", + ep_sizes: str = "", + max_cases: int = 128, +) -> dict[str, Any]: + """Resolve suite configuration into allocation-sized workflow shards.""" + if max_cases <= 0: + raise SystemExit("--max-cases must be positive") + # --ep-sizes narrows the matrix to specific expert-parallel degrees at dispatch + # time: "8" keeps every EP8 shard and drops EP16, so a comprehensive run can + # co-schedule the 8-GPU SKUs' single-node EP8 with the GB SKUs' two-node EP8 + # without dispatching any EP16 leg. Blank keeps every degree. The resulting + # matrix is a partial subset that only omits cases; it never reclassifies them. + selected_eps: set[int] = set() + for value in (part.strip() for part in ep_sizes.split(",")): + if not value: + continue + if not value.isdigit() or int(value) <= 0: + raise SystemExit(f"invalid --ep-sizes {ep_sizes!r}; expected positive integers") + selected_eps.add(int(value)) + if only_sku and only_sku not in cap.PLATFORMS: + raise SystemExit(f"unknown --only-sku {only_sku!r}; have {sorted(cap.PLATFORMS)}") + # --exclude-skus narrows the matrix to a subset by dropping whole runner pools + # — e.g. exclude a SKU whose cluster is unavailable. It only omits cases. + excluded = {value.strip() for value in exclude_skus.split(",") if value.strip()} + unknown_excluded = sorted(excluded - set(cap.PLATFORMS)) + if unknown_excluded: + raise SystemExit( + f"unknown --exclude-skus {unknown_excluded}; have {sorted(cap.PLATFORMS)}" + ) + if only_sku and only_sku in excluded: + raise SystemExit("--only-sku and --exclude-skus select disjoint pools") + + suites_document = _load("suites.yaml") + iters, trials, warmup = _timing(suites_document) + timing_profile = f"{iters}:{trials}:{warmup}" + workloads = suites_document["workloads"] + registry = suites_document["suites"] + select_all = suites == "all" + names = ( + list(registry) + if select_all + else [value.strip() for value in suites.split(",") if value.strip()] + ) + if not names or len(names) != len(set(names)): + raise SystemExit("suite selection must be non-empty and unique") + unknown = sorted(set(names) - set(registry)) + if unknown: + raise SystemExit(f"unknown suites {unknown}; have {sorted(registry)}") + targets = _select_backends(backend, backends) + + shards: dict[tuple[str, str, int], list[dict[str, Any]]] = {} + requested_cases: list[dict[str, Any]] = [] + for suite_name in names: + suite = registry[suite_name] + mode = suite["mode"] + phases = suite["phases"] + routings = suite["routings"] + suite_backends = set(suite.get("backends", cap.SWEEP_BACKENDS)) + suite_targets = [target for target in targets if target in suite_backends] + if not suite_targets: + continue + # "all" = every SKU in configs/platform_config.json (via capability). + suite_platforms = suite["platforms"] + if suite_platforms == "all": + suite_platforms = sorted(cap.PLATFORMS) + for platform_name in suite_platforms: + if only_sku and platform_name != only_sku: + continue + if platform_name in excluded: + continue + ep_degrees = suite["ep_degrees"] + for workload, ep, phase, routing, target in itertools.product( + suite["workloads"], ep_degrees, phases, routings, + suite_targets, + ): + if selected_eps and ep not in selected_eps: + continue + topology = cap.topology_for(platform_name, ep) + if topology is None: + raise SystemExit( + f"suite {suite_name}: {platform_name} EP{ep} is not registered" + ) + nodes = int(topology["nodes"]) + capability_disposition, capability_detail = cap.resolve_disposition( + platform_name, + target, + ep=ep, + nodes=nodes, + routing=routing, + mode=mode, + ) + hidden, topk, experts, seed = _dims(workloads, workload) + + def add_case( + case_ladder: str, + disposition: str, + reason: str | None, + detail: str | None, + ) -> None: + case: dict[str, Any] = { + "suite": suite_name, + "workload": workload, + "backend": target, + "routing": routing, + "phase": phase, + "ep": ep, + "hidden": hidden, + "topk": topk, + "experts": experts, + "seed": seed, + "ladder": case_ladder, + "mode": mode, + "timing": timing_profile, + **{field: topology[field] for field in TOPOLOGY_FIELDS}, + } + # Same function the harness recomputes at run time — a scheduled + # case ID can never drift from its realized factors. + case["case_id"] = ep_harness.case_id(platform_name, case) + requested_cases.append( + { + "sku": platform_name, + "case": case, + "disposition": disposition, + "reason": reason, + "detail": detail, + } + ) + if disposition == "runnable": + shards.setdefault((platform_name, target, nodes), []).append(case) + + requested_ladder = _ladder(workloads, workload, phase) + if capability_disposition == "unsupported": + add_case( + requested_ladder, + "unsupported", + "backend-platform-unsupported", + capability_detail, + ) + continue + add_case(requested_ladder, "runnable", None, None) + + shards_by_sku: dict[str, list[dict[str, Any]]] = {} + for (sku, target, nodes), cases in sorted(shards.items()): + chunk_size = max_cases + for offset in range(0, len(cases), chunk_size): + chunk = cases[offset:offset + chunk_size] + part = offset // chunk_size + shard_id = f"{sku}-{target}-n{nodes}" + if len(cases) > chunk_size: + shard_id += f"-p{part}" + shards_by_sku.setdefault(sku, []).append({ + "id": shard_id, + "sku": sku, + "backend": target, + "launcher": cap.PLATFORMS[sku]["launcher"], + **{field: chunk[0][field] for field in TOPOLOGY_FIELDS}, + "n": len(chunk), + "case_ids": [case["case_id"] for case in chunk], + "cases": chunk, + }) + include = [ + shards_by_sku[sku][round_index] + for round_index in range(max(map(len, shards_by_sku.values()), default=0)) + for sku in sorted(shards_by_sku) + if round_index < len(shards_by_sku[sku]) + ] + return { + "version": suites_document["version"], + "requested_cases": requested_cases, + "include": include, + } + + +def extract_shard(matrix_path: str, shard_id: str, output_path: str) -> dict[str, Any]: + """Select one generator-produced shard and write its execution document.""" + with open(matrix_path) as fh: + document = json.load(fh) + matches = [item for item in document["include"] if item["id"] == shard_id] + if len(matches) != 1: + raise SystemExit(f"expected one shard {shard_id!r}, found {len(matches)}") + source = matches[0] + control = { + key: source[key] + for key in ("id", "sku", "backend", "nodes", "n", "cases") + } + control["version"] = document["version"] + output = Path(output_path) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(control, sort_keys=True, separators=(",", ":")) + "\n") + return control + + +def main() -> int: + parser = argparse.ArgumentParser(description="CollectiveX matrix resolver") + parser.add_argument("--suites", default="all", help="'all' or comma-list of suites") + parser.add_argument("--backend", default="", help="select one EP backend") + parser.add_argument("--backends", default="", help="'all' or comma-list of EP backends") + parser.add_argument("--only-sku", default="") + parser.add_argument( + "--exclude-skus", + default="", + help="comma-list of runner pools to drop (partial matrix); disjoint from --only-sku", + ) + parser.add_argument( + "--ep-sizes", + default="", + help="comma-list of expert-parallel degrees to keep (e.g. 8 drops EP16); blank = all", + ) + parser.add_argument("--max-cases", type=int, default=128) + parser.add_argument("--extract-from", default="", metavar="MATRIX") + parser.add_argument("--shard-id", default="") + parser.add_argument("--out", default="") + args = parser.parse_args() + + if args.extract_from: + if not all((args.shard_id, args.out)): + parser.error("shard extraction requires --shard-id and --out") + control = extract_shard(args.extract_from, args.shard_id, args.out) + print(f"extracted {control['id']}: {control['n']} cases", file=sys.stderr) + print(json.dumps(control, separators=(",", ":"))) + return 0 + + matrix = resolve_matrix( + suites=args.suites, + backend=args.backend, + backends=args.backends, + only_sku=args.only_sku, + exclude_skus=args.exclude_skus, + ep_sizes=args.ep_sizes, + max_cases=args.max_cases, + ) + if args.out: + with open(args.out, "w") as fh: + json.dump(matrix, fh, sort_keys=True, separators=(",", ":")) + fh.write("\n") + runnable = sum( + item["disposition"] == "runnable" for item in matrix["requested_cases"] + ) + unsupported = len(matrix["requested_cases"]) - runnable + print( + f"resolved {len(matrix['include'])} shard-cells, " + f"{runnable} runnable and {unsupported} unsupported cases", + file=sys.stderr, + ) + print(json.dumps(matrix)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/experimental/CollectiveX/tests/test_ep_backend.py b/experimental/CollectiveX/tests/test_ep_backend.py new file mode 100644 index 0000000000..330d4dc4fd --- /dev/null +++ b/experimental/CollectiveX/tests/test_ep_backend.py @@ -0,0 +1,121 @@ +#!/usr/bin/env python3 +"""Small torch-free smoke tests for the shared EP backend lifecycle.""" +from __future__ import annotations + +import sys +import types +import unittest +from pathlib import Path +from unittest import mock + +ROOT = Path(__file__).resolve().parents[1] +sys.path[:0] = [str(ROOT), str(ROOT / "bench")] + +import ep_backend # noqa: E402 +from ep_backend import EPBackend, RankInputs # noqa: E402 + + +def args(**updates): + values = dict( + experts=8, phase="decode", tokens_ladder="", routing="uniform", seed=0, + hidden=16, topk=2, mode="normal", + ) + values.update(updates) + return types.SimpleNamespace(**values) + + +class FakeBackend(EPBackend): + name = "fake" + + def __init__(self, options, *, cap=None, world_size=1): + super().__init__(options, 0, world_size, 0, "cpu") + self.cap = cap + self.calls: list[str] = [] + + def create_buffer(self, spec): + return None + + def dispatch(self, problem): + self.calls.append("dispatch") + return object() + + def stage(self, problem, handle): + self.calls.append("stage") + + def combine(self, problem, handle): + self.calls.append("combine") + + def recv_tokens(self, handle): + return 0 + + def inspect_dispatch(self, problem, handle): + return None + + def combine_transformed(self, problem, handle, transformed): + return None + + def buffer_cap(self, options): + return self.cap + + def _build_rank_inputs(self, options, tokens): + return RankInputs( + tokens_per_rank=tokens, topk_idx=None, topk_weights=None, + activations=None, + ) + + +class BackendTests(unittest.TestCase): + def test_input_plan_sizes_for_the_measured_ladder(self): + backend = FakeBackend(args(tokens_ladder="8 16"), world_size=2) + spec = backend.make_inputs(backend.args) + self.assertTrue(spec.ok) + self.assertEqual(spec.ladder, [8, 16]) + self.assertEqual(spec.max_tokens_per_rank, 16) + self.assertEqual((spec.ep_size, spec.experts_per_rank), (2, 4)) + self.assertEqual(sorted(spec.points), [8, 16]) + + def test_invalid_or_fully_clamped_ladder_fails_before_execution(self): + for backend, message in ( + (FakeBackend(args(tokens_ladder="0")), "empty token ladder"), + (FakeBackend(args(tokens_ladder="128"), cap=64), "cap=64"), + ): + with self.subTest(message=message): + spec = backend.make_inputs(backend.args) + self.assertEqual(spec.rc, 2) + self.assertIn(message, spec.message) + + def test_timed_components_follow_backend_contract(self): + backend = FakeBackend(args()) + self.assertEqual(backend.timed_components(), ["roundtrip", "dispatch", "combine"]) + backend.stage_device_work = True + self.assertEqual( + backend.timed_components(), ["roundtrip", "dispatch", "combine", "stage"] + ) + backend.roundtrip_only = True + self.assertEqual(backend.timed_components(), ["roundtrip"]) + + def test_dispatch_cleanup_is_outside_timed_call(self): + backend = FakeBackend(args()) + backend.dispatch_needs_combine_cleanup = True + captured = {} + + def fake_time(_torch, operation, _warmup, _iters, **kwargs): + handle = operation() + kwargs["post"](handle) + captured.update(kwargs) + return [1.0] + + with mock.patch.dict(sys.modules, {"torch": types.SimpleNamespace()}), mock.patch.object( + ep_backend, "time_us", side_effect=fake_time + ): + backend.benchmark_dispatch(object(), 0, 1) + self.assertIn("post", captured) + self.assertEqual(backend.calls, ["dispatch", "stage", "combine"]) + + def test_mode_is_fail_closed(self): + with self.assertRaises(ValueError): + FakeBackend(args(mode="unsupported")) + + +if __name__ == "__main__": + unittest.main() diff --git a/experimental/CollectiveX/tests/test_matrix.py b/experimental/CollectiveX/tests/test_matrix.py new file mode 100644 index 0000000000..c2fb3ea74a --- /dev/null +++ b/experimental/CollectiveX/tests/test_matrix.py @@ -0,0 +1,63 @@ +#!/usr/bin/env python3 +"""Matrix, subset, and shard-extraction tests.""" +from __future__ import annotations + +import json +import sys +import tempfile +import unittest +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +import sweep_matrix # noqa: E402 + + +def matrix(**options): + return sweep_matrix.resolve_matrix(max_cases=128, **options) + + +class MatrixTests(unittest.TestCase): + def test_shard_extraction_is_deterministic_and_preserves_cases(self): + document = matrix(backend="deepep-v2", only_sku="h200-dgxc") + cell = document["include"][0] + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + source = root / "matrix.json" + source.write_text(json.dumps(document, sort_keys=True)) + outputs = [ + sweep_matrix.extract_shard( + source, cell["id"], root / f"shard-{index}.json", + ) + for index in range(2) + ] + self.assertEqual(outputs[0], outputs[1]) + self.assertEqual([case["case_id"] for case in outputs[0]["cases"]], cell["case_ids"]) + + def test_sku_and_ep_filters_only_remove_cases(self): + full = matrix(suites="all", backends="all") + for options, keep in ( + ({"exclude_skus": "b300"}, lambda item: item["sku"] != "b300"), + ({"ep_sizes": "8"}, lambda item: item["case"]["ep"] == 8), + ): + partial = matrix(suites="all", backends="all", **options) + expected = { + item["case"]["case_id"]: item for item in full["requested_cases"] if keep(item) + } + actual = {item["case"]["case_id"]: item for item in partial["requested_cases"]} + self.assertEqual(actual, expected) + + def test_invalid_filters_fail_closed(self): + for options in ( + {"exclude_skus": "unknown"}, + {"only_sku": "b300", "exclude_skus": "b300"}, + {"ep_sizes": "0"}, + {"ep_sizes": "eight"}, + ): + with self.subTest(options=options), self.assertRaises(SystemExit): + sweep_matrix.resolve_matrix(**options) + + +if __name__ == "__main__": + unittest.main() diff --git a/experimental/CollectiveX/tests/test_runtime.py b/experimental/CollectiveX/tests/test_runtime.py new file mode 100644 index 0000000000..fb679765fc --- /dev/null +++ b/experimental/CollectiveX/tests/test_runtime.py @@ -0,0 +1,474 @@ +#!/usr/bin/env python3 +"""Focused tests for the standalone runtime helpers.""" + +from __future__ import annotations + +import argparse +import contextlib +import io +import json +import os +from pathlib import Path +import re +import subprocess +import sys +import tempfile +import unittest +from unittest import mock + + +RUNTIME = Path(__file__).resolve().parents[1] / "runtime" +BENCH = Path(__file__).resolve().parents[1] / "bench" +sys.path.insert(0, str(RUNTIME)) +sys.path.insert(0, str(BENCH)) + +import probe # noqa: E402 +import config # noqa: E402 +import stage # noqa: E402 +import ep_harness # noqa: E402 (stdlib-only at module top) + + +# configs/platform_config.json is the single per-SKU registry shared by +# capability.py (topologies), runtime/config.py (operator schema and network +# overlay), runtime/prepare_backend.sh (CUDA target), and +# bench/ep_mori.py (arch pin). Validate its shape once so a malformed entry +# fails here instead of on an allocation. +class PlatformRegistryTests(unittest.TestCase): + REGISTRY = RUNTIME.parent / "configs" / "platform_config.json" + NETWORK_FIELDS = { + "socket_ifname", "rdma_devices", "ib_gid_index", + "rdma_service_level", "rdma_traffic_class", + } + + def test_every_platform_entry_is_complete_and_typed(self) -> None: + platforms = json.loads(self.REGISTRY.read_text())["platforms"] + self.assertTrue(platforms) + for name, entry in platforms.items(): + with self.subTest(sku=name): + for field in ("vendor", "arch", "machine", "product", + "scale_up_transport", "launcher"): + self.assertIsInstance(entry[field], str) + self.assertTrue(entry[field]) + for field in ("gpus_per_node", "scale_up_domain"): + self.assertIsInstance(entry[field], int) + self.assertGreater(entry[field], 0) + self.assertTrue(entry["operator_fields"]) + self.assertLessEqual( + set(entry.get("network", {})), self.NETWORK_FIELDS + ) + if entry["vendor"] == "nvidia": + self.assertRegex(entry["arch"], r"^sm\d+$") + else: + self.assertRegex(entry["arch"], r"^gfx\d+$") + + def test_capability_loads_exactly_the_registry_skus(self) -> None: + sys.path.insert(0, str(RUNTIME.parent)) + try: + import capability + finally: + sys.path.pop(0) + platforms = json.loads(self.REGISTRY.read_text())["platforms"] + self.assertEqual(set(capability.PLATFORMS), set(platforms)) + + +class ProbeTests(unittest.TestCase): + def test_default_route_interface(self) -> None: + with tempfile.TemporaryDirectory() as directory: + route = Path(directory) / "route" + route.write_text( + "Iface Destination Gateway Flags RefCnt Use Metric Mask MTU Window IRTT\n" + "eth9 00000000 00000000 0003 0 0 0 00000000 0 0 0\n" + ) + self.assertEqual(probe.default_route_interface(route), "eth9") + + def test_prepare_cache_is_private_and_reusable(self) -> None: + with tempfile.TemporaryDirectory() as directory: + first = Path(probe.prepare_cache(directory)) + second = Path(probe.prepare_cache(directory)) + self.assertEqual(first, second) + self.assertEqual(first.stat().st_mode & 0o777, 0o700) + + +class ConfigTests(unittest.TestCase): + def test_merge_operator_config_applies_overlays(self) -> None: + base = {"runners": {"h100-dgxc": {"partition": "gpu"}}} + overlay = {"runners": {"h100-dgxc": {"account": "bench"}}} + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "operator.json" + with mock.patch.dict(os.environ, { + "COLLECTIVEX_OPERATOR_CONFIG_CONTENT": json.dumps(base), + "COLLECTIVEX_H100_CONFIG_CONTENT": json.dumps(overlay), + }, clear=True): + config.merge_operator_config(str(path)) + self.assertEqual( + json.loads(path.read_text())["runners"]["h100-dgxc"], + {"partition": "gpu", "account": "bench"}, + ) + self.assertEqual(path.stat().st_mode & 0o777, 0o600) + + def test_merge_operator_config_tolerates_absent_base(self) -> None: + # De-secreted operator config: the base secret is deleted (empty CONTENT). + # Merge must still write {"runners":{}} so the launcher finds a config file + # and falls back to the tracked platform_config.json baseline. + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "operator.json" + with mock.patch.dict( + os.environ, {"COLLECTIVEX_OPERATOR_CONFIG_CONTENT": ""}, clear=True + ): + config.merge_operator_config(str(path)) + self.assertEqual(json.loads(path.read_text()), {"runners": {}}) + + def test_operator_config_emits_allowlisted_values(self) -> None: + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "operator.json" + path.write_text(json.dumps({ + "runners": { + "h100-dgxc": { + "partition": "gpu", + "account": "bench", + "squash_dir": directory, + } + }, + })) + path.chmod(0o600) + read_fd, write_fd = os.pipe() + stdout = sys.stdout + try: + sys.stdout = os.fdopen(write_fd, "w") + config.operator_config(str(path), "h100-dgxc") + sys.stdout.flush() + finally: + sys.stdout.close() + sys.stdout = stdout + payload = os.read(read_fd, 4096) + os.close(read_fd) + self.assertIn(b"COLLX_PARTITION\0gpu\0", payload) + self.assertIn(b"COLLX_SQUASH_DIR\0" + directory.encode() + b"\0", payload) + + def _emit_registry_only(self, runner: str) -> bytes: + read_fd, write_fd = os.pipe() + stdout = sys.stdout + try: + sys.stdout = os.fdopen(write_fd, "w") + config.operator_config("-", runner) + sys.stdout.flush() + finally: + sys.stdout.close() + sys.stdout = stdout + payload = os.read(read_fd, 4096) + os.close(read_fd) + return payload + + def test_operator_config_registry_only_emits_tracked_baseline(self) -> None: + # "-" = no operator document: the registry's per-SKU operator block is + # the tracked baseline (plus its network overlay where present). + payload = self._emit_registry_only("h200-dgxc") + self.assertIn(b"COLLX_PARTITION\0main\0", payload) + self.assertIn(b"COLLX_SQUASH_DIR\0/home/sa-shared/containers\0", payload) + self.assertIn(b"COLLX_RDMA_DEVICES\0", payload) + + def test_operator_config_registry_only_skips_secret_fed_sku(self) -> None: + # SKUs without a tracked operator block keep the historical no-config + # behavior: nothing emitted, no validation failure. mi325x is the last + # still-secret-fed SKU (no SSH access target to derive a baseline). + self.assertEqual(self._emit_registry_only("mi325x"), b"") + +# configs/backends.json is the single definition of backend source pins and +# container images; config.py backend-registry validates it and emits the +# COLLX_* names common.sh consumes at source time. These tests pin both sides +# of that seam so a rename or malformed registry fails here, not on a runner. +class BackendRegistryTests(unittest.TestCase): + EMITTED = { + "COLLX_IMAGE_MULTIARCH", "COLLX_IMAGE_AMD_MORI", + "COLLX_DEEPEP_V2_REPO", "COLLX_DEEPEP_V2_COMMIT", + } + + @staticmethod + def _emitted_pairs() -> dict[str, str]: + result = subprocess.run( + [sys.executable, str(RUNTIME / "config.py"), "backend-registry"], + capture_output=True, check=True, + ) + parts = result.stdout.split(b"\0") + assert parts[-1] == b"" + values = [part.decode() for part in parts[:-1]] + return dict(zip(values[::2], values[1::2], strict=True)) + + def test_registry_emits_exactly_the_declared_names(self) -> None: + pairs = self._emitted_pairs() + self.assertEqual(set(pairs), self.EMITTED) + for name, value in pairs.items(): + if "_COMMIT" in name: + self.assertRegex(value, r"^[0-9a-f]{40}$", name) + elif name.endswith("_REPO"): + self.assertTrue(value.startswith("https://"), name) + else: + self.assertRegex(value, r"^[A-Za-z0-9._/-]+:[A-Za-z0-9._-]+$", name) + + def test_common_sh_requires_every_emitted_name(self) -> None: + # The loader's fail-closed loop and the registry emit must not drift. + common = (RUNTIME / "common.sh").read_text() + for name in self.EMITTED: + self.assertIn(name, common) + + def test_registry_values_come_from_the_tracked_file(self) -> None: + document = json.loads( + (RUNTIME.parent / "configs" / "backends.json").read_text() + ) + pairs = self._emitted_pairs() + self.assertEqual( + pairs["COLLX_DEEPEP_V2_COMMIT"], + document["backends"]["deepep-v2"]["commit"], + ) + self.assertEqual(pairs["COLLX_IMAGE_MULTIARCH"], document["images"]["multiarch"]["ref"]) + + def test_registry_fails_closed(self) -> None: + with tempfile.TemporaryDirectory() as directory: + missing = Path(directory) / "absent.json" + with self.assertRaises(SystemExit): + config.backend_registry(str(missing)) + malformed = Path(directory) / "backends.json" + document = json.loads( + (RUNTIME.parent / "configs" / "backends.json").read_text() + ) + document["backends"]["deepep-v2"]["commit"] = "not-a-sha" + malformed.write_text(json.dumps(document)) + with self.assertRaises(SystemExit): + config.backend_registry(str(malformed)) + + +class StageTests(unittest.TestCase): + def test_create_copy_and_validate_cleanup(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + source = root / "source" + target = root / "stage" + (source / "runtime").mkdir(parents=True) + (source / "runtime" / "common.sh").write_text("test") + (source / "goal.md").write_text("private") + (source / ".shards").mkdir() + (source / ".shards" / "leg.json").write_text("{}") + args = type("Args", (), {"stage": str(target)}) + stage.create_stage(args) + copy_args = type( + "Args", (), {"source": str(source), "target": str(target / "experimental" / "CollectiveX")} + ) + stage.copy_repository(copy_args) + staged = target / "experimental" / "CollectiveX" + self.assertTrue((staged / "runtime" / "common.sh").is_file()) + self.assertFalse((staged / ".shards").exists()) + self.assertFalse((staged / "goal.md").exists()) + cleanup_args = type("Args", (), {"root": str(target)}) + stage.validate_cleanup(cleanup_args) + +# The per-node probe (runtime/probe.py) and the launcher gate +# (runtime/common.sh: collx_validate_network_profile_on_job) share an implicit string contract: +# the probe prints these markers, the launcher greps them back out to derive COLLX_SOCKET_IFNAME +# and COLLX_RDMA_LINK_LAYER. The patterns are duplicated here on purpose — the test fails if +# either side drifts, which is exactly the failure that slipped through when 5506c623 moved the +# probe into Python but left the emit statements behind, silently zeroing the marker count for +# every non-MNNVL multi-node leg. +SOCKET_MARKER = r"^\[collectivex-private\] socket-interface-selected=([A-Za-z][A-Za-z0-9_.-]{0,31})$" +LINK_MARKER = r"^\[collectivex-private\] rdma-link-layer=(roce|infiniband)$" +FAILURE_MARKER = ( + r"(socket-interface|rdma-(device|port))-[0-9]+=" + r"(missing|down|inactive|default-route-missing|gid-missing|gid-empty|" + r"link-layer-missing|link-layer-invalid|link-layer-mixed)" +) + + +class NetworkProfileContract(unittest.TestCase): + def _fabric(self, root: Path, *, state: str = "4: ACTIVE", + link_layer: str = "Ethernet", gid: str = "fe80::1") -> None: + net = root / "class" / "net" / "eth0" + net.mkdir(parents=True) + (net / "operstate").write_text("up\n") + port = root / "class" / "infiniband" / "mlx5_0" / "ports" / "1" + (port / "gids").mkdir(parents=True) + (port / "state").write_text(state + "\n") + (port / "link_layer").write_text(link_layer + "\n") + (port / "gids" / "3").write_text(gid + "\n") + + def _run(self, root: Path, route: Path, socket_names: str = "eth0"): + buffer = io.StringIO() + rc = 0 + try: + with contextlib.redirect_stdout(buffer): + probe.validate_network_profile(socket_names, "mlx5_0:1", "3", + sys_root=root, route_path=route) + except SystemExit: + rc = 1 + return rc, buffer.getvalue().splitlines() + + @staticmethod + def _captures(pattern: str, lines: list) -> list: + return [match.group(1) for line in lines + for match in [re.match(pattern, line)] if match] + + def test_launcher_still_declares_the_marker_patterns(self) -> None: + common = (RUNTIME / "common.sh").read_text() + self.assertIn(SOCKET_MARKER, common) + self.assertIn(LINK_MARKER, common) + self.assertIn(FAILURE_MARKER, common) + + def test_healthy_fabric_emits_the_success_markers_the_launcher_extracts(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + self._fabric(root) + rc, lines = self._run(root, root / "route") + self.assertEqual(rc, 0) + self.assertEqual(self._captures(SOCKET_MARKER, lines), ["eth0"]) + self.assertEqual(self._captures(LINK_MARKER, lines), ["roce"]) + + def test_infiniband_link_layer_maps_to_the_launcher_token(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + self._fabric(root, link_layer="InfiniBand") + rc, lines = self._run(root, root / "route") + self.assertEqual(rc, 0) + self.assertEqual(self._captures(LINK_MARKER, lines), ["infiniband"]) + + def test_socket_interface_resolves_from_default_route(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + self._fabric(root) + route = root / "route" + route.write_text( + "Iface Destination Gateway Flags RefCnt Use Metric Mask MTU Window IRTT\n" + "eth0 00000000 00000000 0003 0 0 0 00000000 0 0 0\n" + ) + rc, lines = self._run(root, route, socket_names="") + self.assertEqual(rc, 0) + self.assertEqual(self._captures(SOCKET_MARKER, lines), ["eth0"]) + + def test_inactive_port_emits_a_launcher_recognized_failure_marker(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + self._fabric(root, state="1: DOWN") + rc, lines = self._run(root, root / "route") + self.assertEqual(rc, 1) + failures = [line for line in lines if re.search(FAILURE_MARKER, line)] + self.assertTrue(any("rdma-port-1=inactive" in line for line in failures), failures) + + def test_all_zero_gid_emits_gid_empty(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + self._fabric(root, gid="0000:0000:0000:0000:0000:0000:0000:0000") + rc, lines = self._run(root, root / "route") + self.assertEqual(rc, 1) + self.assertTrue(any("rdma-port-1=gid-empty" in line for line in lines), lines) + + +class StageContract(unittest.TestCase): + # runtime/common.sh drives runtime/stage.py purely by literal subcommand name and positional + # argv — there are no optional flags. That argv shape is a string contract: a subcommand or + # flag the launcher passes but stage.py does not declare fails with "unrecognized arguments" + # and aborts the leg at repository-stage. This extracts every stage.py call out of common.sh + # and proves stage.py's parser accepts it — the guard that would have caught the --allow-* + # flags surviving on the callers after they were dropped from stage.py's argparse. + @staticmethod + def _invocations(text: str) -> list: + calls = [] + for line in text.splitlines(): + if "stage.py" not in line or line.lstrip().startswith("#"): + continue + subcommand, flags = None, [] + for raw in line.split("stage.py", 1)[1].split(): + token = raw.strip('"').strip("'") + if token.startswith("--"): + flags.append(token.split("=", 1)[0]) + elif subcommand is None and token and not token.startswith(("$", "${")): + subcommand = token + if subcommand: + calls.append((subcommand, flags)) + return calls + + def test_launcher_only_invokes_declared_subcommands_and_flags(self) -> None: + invocations = self._invocations((RUNTIME / "common.sh").read_text()) + self.assertGreaterEqual(len(invocations), len(stage.SPECS), invocations) + parser = stage.build_parser() + for subcommand, flags in invocations: + self.assertIn(subcommand, stage.SPECS, subcommand) + argv = [subcommand] + ["x"] * len(stage.SPECS[subcommand]) + flags + with contextlib.redirect_stderr(io.StringIO()): + try: + parser.parse_args(argv) + except SystemExit: + self.fail(f"common.sh invokes stage.py with an argv shape it rejects: {argv}") + + def test_contract_test_has_teeth(self) -> None: + # A flag common.sh must never pass has to be rejected by the parser — this is the exact + # failure (unrecognized arguments: --allow-parent-owner) the reconcile removed. + parser = stage.build_parser() + with contextlib.redirect_stderr(io.StringIO()): + with self.assertRaises(SystemExit): + parser.parse_args(["validate-stage-path", "x", "x", "x", "--allow-parent-owner"]) + + +# config.py case-args is the single case→invocation codec: collx_run_shard decodes one +# null-delimited argv per case and hands it verbatim to bench/run_ep.py. Parse the +# emitted argv with the same parser shape run_ep builds so the two sides cannot +# drift — a flag the codec emits but run_ep does not declare (or vice versa) fails +# here instead of on a GPU allocation. +class CaseArgvContract(unittest.TestCase): + CASE = { + "backend": "deepep-v2", "mode": "normal", "phase": "decode", + "routing": "uniform", "ep": 16, "nodes": 2, "gpus_per_node": 8, + "scale_up_domain": 8, "scope": "scale-out", + "scale_up_transport": "nvlink", "scale_out_transport": "rdma", + "transport": "nvlink-rdma", "topology_class": "h200-nvlink-rdma", + "hidden": 7168, "topk": 8, "experts": 256, "seed": 67, + "ladder": "1 2 4", "timing": "8:128:32", + "case_id": "h200-dgxc-deepep-v2-deepseek-v3-normal-decode-ep16-uniform", + "suite": "ep-core", "workload": "deepseek-v3", + } + + @staticmethod + def _run_ep_parser() -> argparse.ArgumentParser: + # Mirror of the parser bench/run_ep.py builds in main(). + parser = argparse.ArgumentParser() + parser.add_argument( + "--backend", required=True, choices=["deepep-v2", "mori"] + ) + ep_harness.add_common_args(parser) + return parser + + def _decode(self, stdout: bytes) -> list: + parts = stdout.split(b"\0") + self.assertEqual(parts[-1], b"") + return [part.decode() for part in parts[:-1]] + + def _case_argv(self, placement: list) -> list: + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "shard.json" + path.write_text(json.dumps({"version": 1, "cases": [self.CASE]})) + result = subprocess.run( + [sys.executable, str(RUNTIME / "config.py"), "case-args", + str(path), "0", "h200-dgxc", "TS", *placement], + capture_output=True, check=True, + ) + return self._decode(result.stdout) + + def test_case_args_round_trips_through_the_run_ep_parser(self) -> None: + argv = self._case_argv(["16", "2", "8", "8"]) + args = self._run_ep_parser().parse_args(argv) + self.assertEqual( + (args.backend, args.mode, args.phase, args.routing, args.scope), + ("deepep-v2", "normal", "decode", "uniform", "scale-out"), + ) + self.assertEqual((args.hidden, args.topk, args.experts), (7168, 8, 256)) + self.assertEqual((args.gpus_per_node, args.scale_up_domain), (8, 8)) + self.assertEqual(args.tokens_ladder, "1 2 4") + self.assertEqual(args.scale_out_transport, "rdma") + self.assertEqual(args.case_id, self.CASE["case_id"]) + self.assertEqual(args.version, 1) + self.assertEqual(args.seed, self.CASE["seed"]) + self.assertEqual((args.iters, args.trials, args.warmup), (8, 128, 32)) + self.assertEqual(args.out, "results/h200-dgxc_deepep-v2_decode_TS-c000.json") + + def test_case_args_fails_closed_on_placement_mismatch(self) -> None: + with self.assertRaises(subprocess.CalledProcessError): + self._case_argv(["8", "1", "8", "8"]) + +if __name__ == "__main__": + unittest.main()