diff --git a/.github/workflows/test-matrix-logic.yml b/.github/workflows/test-matrix-logic.yml index 67910686f4..7851364f18 100644 --- a/.github/workflows/test-matrix-logic.yml +++ b/.github/workflows/test-matrix-logic.yml @@ -5,6 +5,9 @@ on: pull_request: paths: - 'utils/matrix_logic/**' + - 'runners/launch_mi300x-amds*.sh' + - 'runners/mi300x_native_node_preflight.sh' + - 'benchmarks/multi_node/agentic/kimik3_fp4_mi300x_vllm.sh' permissions: contents: read @@ -39,3 +42,8 @@ jobs: run: | cd utils/matrix_logic pytest test_validation.py -v + + - name: test_kimik3_mi300x_native tests + run: | + cd utils/matrix_logic + pytest test_kimik3_mi300x_native.py -v diff --git a/benchmarks/multi_node/agentic/kimik3_fp4_mi300x_vllm.sh b/benchmarks/multi_node/agentic/kimik3_fp4_mi300x_vllm.sh new file mode 100755 index 0000000000..e726886ba6 --- /dev/null +++ b/benchmarks/multi_node/agentic/kimik3_fp4_mi300x_vllm.sh @@ -0,0 +1,114 @@ +#!/usr/bin/env bash +set -euo pipefail +set -x + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "$SCRIPT_DIR/../../benchmark_lib.sh" + +check_env_vars \ + MODEL \ + MODEL_PATH \ + PORT \ + CONC_LIST \ + PREFILL_NUM_WORKERS \ + PREFILL_TP \ + PREFILL_PP_SIZE \ + PREFILL_EP \ + PREFILL_DP_ATTN \ + DECODE_NUM_WORKERS \ + MULTINODE_NODE_COUNT \ + MULTINODE_GPUS_PER_NODE \ + MULTINODE_NODE_RANK \ + MULTINODE_MASTER_ADDR + +require_agentic_kv_offload_none + +fail() { + echo "ERROR: $*" >&2 + exit 1 +} + +if [[ "$MULTINODE_NODE_COUNT" != "2" ]]; then + fail "this entrypoint serves exactly 2 nodes, got MULTINODE_NODE_COUNT=$MULTINODE_NODE_COUNT" +fi +if [[ "$MULTINODE_GPUS_PER_NODE" != "8" ]]; then + fail "this entrypoint requires 8 GPUs per node, got MULTINODE_GPUS_PER_NODE=$MULTINODE_GPUS_PER_NODE" +fi +if [[ "$MULTINODE_NODE_RANK" != "0" && "$MULTINODE_NODE_RANK" != "1" ]]; then + fail "MULTINODE_NODE_RANK must be 0 or 1, got '$MULTINODE_NODE_RANK'" +fi +if [[ "$PREFILL_TP" != "8" || "$PREFILL_PP_SIZE" != "2" ]]; then + fail "this entrypoint serves only TP8 x PP2, got TP$PREFILL_TP x PP$PREFILL_PP_SIZE" +fi +if [[ "$PREFILL_EP" != "1" ]]; then + fail "this entrypoint serves only EP1, got PREFILL_EP=$PREFILL_EP" +fi +if [[ "$PREFILL_DP_ATTN" != "false" ]]; then + fail "this entrypoint does not enable DP attention, got PREFILL_DP_ATTN=$PREFILL_DP_ATTN" +fi +if [[ "$PREFILL_NUM_WORKERS" != "1" || "$DECODE_NUM_WORKERS" != "0" ]]; then + fail "this entrypoint is aggregated: it needs 1 prefill worker and 0 decode workers, got ${PREFILL_NUM_WORKERS}P/${DECODE_NUM_WORKERS}D" +fi + +read -r -a CONCURRENCIES <<< "$CONC_LIST" +if [[ "${#CONCURRENCIES[@]}" -ne 1 ]]; then + fail "one concurrency per allocation is required, got CONC_LIST='$CONC_LIST'" +fi +case "${CONCURRENCIES[0]}" in + 1|2|4|8) ;; + *) fail "concurrency must be 1, 2, 4, or 8, got '${CONCURRENCIES[0]}'" ;; +esac + +if [[ -n "${AITER_SITUV2_A8W4+set}" ]]; then + if [[ "$AITER_SITUV2_A8W4" != "0" && "$AITER_SITUV2_A8W4" != "1" ]]; then + fail "AITER_SITUV2_A8W4 must be 0 or 1 when set, got '$AITER_SITUV2_A8W4'" + fi +fi + +export VLLM_ROCM_USE_AITER=1 +export SAFETENSORS_FAST_GPU=1 +export AITER_BF16_FP8_MOE_BOUND=0 +export VLLM_USE_BREAKABLE_CUDAGRAPH=0 +export VLLM_USE_V2_MODEL_RUNNER=0 +export VLLM_ENGINE_READY_TIMEOUT_S="${VLLM_ENGINE_READY_TIMEOUT_S:-7200}" +export PYTHONNOUSERSITE=1 + +{ set +x; } 2>/dev/null +VLLM_CMD=( + vllm serve "$MODEL_PATH" + --served-model-name "$MODEL" + --host 0.0.0.0 + --port "$PORT" + --tensor-parallel-size 8 + --pipeline-parallel-size 2 + --nnodes 2 + --node-rank "$MULTINODE_NODE_RANK" + --master-addr "$MULTINODE_MASTER_ADDR" + --trust-remote-code + --load-format auto + --moe-backend auto + --gpu-memory-utilization "${VLLM_GPU_MEMORY_UTILIZATION:-0.95}" + --max-model-len 1048576 + --max-num-seqs "$CONC_LIST" + --max-num-batched-tokens "${VLLM_MAX_NUM_BATCHED_TOKENS:-4096}" + --mm-encoder-tp-mode data + --enable-auto-tool-choice + --tool-call-parser kimi_k3 + --reasoning-parser kimi_k3 + --language-model-only +) +if [[ "$MULTINODE_NODE_RANK" == "1" ]]; then + VLLM_CMD+=(--headless) +fi + +echo "AITER_SITUV2_A8W4=${AITER_SITUV2_A8W4-unset}" +printf 'vLLM command:' +printf ' %q' "${VLLM_CMD[@]}" +printf '\n' + +if [[ "${KIMIK3_VLLM_DRY_RUN:-0}" == "1" ]]; then + echo "KIMIK3_VLLM_DRY_RUN=1 set; not starting the server" + exit 0 +fi + +exec "${VLLM_CMD[@]}" diff --git a/benchmarks/multi_node/agentic_srt.sh b/benchmarks/multi_node/agentic_srt.sh index 79a36da524..77dac7cbac 100644 --- a/benchmarks/multi_node/agentic_srt.sh +++ b/benchmarks/multi_node/agentic_srt.sh @@ -2,8 +2,8 @@ set -euo pipefail set -x -# Client-only agentic trace replay for srt-slurm multinode jobs. -# srt-slurm owns server startup; this script runs as benchmark.type=custom +# Client-only agentic trace replay for externally managed multi-node jobs. +# The caller owns server startup; this script runs as benchmark.type=custom # against the already-ready frontend on the head node. INFMAX_CONTAINER_WORKSPACE="${INFMAX_CONTAINER_WORKSPACE:-/infmax-workspace}" diff --git a/configs/amd-master.yaml b/configs/amd-master.yaml index b15ae5bccd..6612ccecc5 100644 --- a/configs/amd-master.yaml +++ b/configs/amd-master.yaml @@ -2244,6 +2244,36 @@ minimaxm3-fp4-mi355x-vllm-agentic: search-space: - { tp: 4, kv-offloading: dram, kv-offload-backend: { name: vllm-simple }, conc-list: [1, 4, 8, 16] } +kimik3-fp4-mi300x-vllm-agentic: + image: vllm/vllm-openai-rocm:kimi-k3 + model: moonshotai/Kimi-K3 + model-prefix: kimik3 + runner: cluster:mi300x-amds + precision: fp4 + framework: vllm + multinode: true + disagg: false + scenarios: + agentic-coding: + - search-space: + - spec-decoding: none + kv-offloading: none + conc-list: [1, 2, 4, 8] + prefill: + num-worker: 1 + tp: 8 + pp: 2 + ep: 1 + dp-attn: false + additional-settings: + - "NATIVE_MULTINODE=1" + decode: + num-worker: 0 + tp: 8 + pp: 2 + ep: 1 + dp-attn: false + dsv4-fp4-mi355x-sglang-disagg-agentic-hicache: image: lmsysorg/sglang-rocm:v0.5.14-rocm720-mi35x-20260710 model: deepseek-ai/DeepSeek-V4-Pro diff --git a/perf-changelog.yaml b/perf-changelog.yaml index 8f5831507e..dffb079ac7 100644 --- a/perf-changelog.yaml +++ b/perf-changelog.yaml @@ -5339,3 +5339,13 @@ - "Pass --use-chat-template to the benchmark, as required for EAGLE-style speculative decoding" - "Validated on the node through the real launcher, every request successful: TEP4 concurrency 64 640/640 at 6895 tok/s (11% above the non-MTP arm's 6198 tok/s, mean TTFT 27.2s to 9.6s), and TEP4 concurrency 1 10/10 at 1233 tok/s (83% above the non-MTP arm's 672 tok/s, mean TPOT 12.6ms to 6.5ms)" pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2312 + +- config-keys: + - kimik3-fp4-mi300x-vllm-agentic + description: + - "Add an opt-in native Slurm path for aggregated Kimi K3 AgentX on two 8xMI300X nodes with TP8 x PP2, EP1, and concurrency 1/2/4/8" + - "Image: vllm/vllm-openai-rocm:kimi-k3; model moonshotai/Kimi-K3" + - "Fail closed on the exact aggregate topology, complete and revision-matched node-local model snapshots, eight gfx942 GPUs per node, and independently validated node-local Enroot squash images" + - "Keep AITER_SITUV2_A8W4 caller-configurable pending exact-shape gfx942 validation, and never download the approximately 1.5 TB target checkpoint inside a benchmark job" + - "Preserve the existing single-node MI300X launcher unless NATIVE_MULTINODE=1, with bounded host-owned artifact handoff and signal/failure cleanup" + pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2401 diff --git a/runners/launch_mi300x-amds-native-multinode.sh b/runners/launch_mi300x-amds-native-multinode.sh new file mode 100755 index 0000000000..51c0816bf5 --- /dev/null +++ b/runners/launch_mi300x-amds-native-multinode.sh @@ -0,0 +1,413 @@ +#!/usr/bin/env bash +set -euo pipefail + +source "$(dirname "${BASH_SOURCE[0]}")/slurm_utils.sh" + +KIMIK3_MODEL_CACHE_ROOT="${KIMIK3_MODEL_CACHE_ROOT:-/raid/hf-hub-cache/models--moonshotai--Kimi-K3}" +KIMIK3_SQUASH_DIR="${KIMIK3_SQUASH_DIR:-/raid/hf-hub-cache/inferencex/squash}" +KIMIK3_SLURM_TIME_MINUTES="${KIMIK3_SLURM_TIME_MINUTES:-480}" +KIMIK3_STARTUP_TIMEOUT_SECONDS="${KIMIK3_STARTUP_TIMEOUT_SECONDS:-7200}" +KIMIK3_HEALTH_POLL_SECONDS="${KIMIK3_HEALTH_POLL_SECONDS:-10}" +KIMIK3_CLEANUP_TIMEOUT_SECONDS="${KIMIK3_CLEANUP_TIMEOUT_SECONDS:-120}" +KIMIK3_CLEANUP_POLL_SECONDS="${KIMIK3_CLEANUP_POLL_SECONDS:-2}" +KIMIK3_PRESCANCEL_TIMEOUT_SECONDS="${KIMIK3_PRESCANCEL_TIMEOUT_SECONDS:-15}" +export KIMIK3_MODEL_CACHE_ROOT KIMIK3_SQUASH_DIR +export PORT="${PORT:-8888}" + +HF_HUB_CACHE_MOUNT="${HF_HUB_CACHE_MOUNT:-/raid/hf-hub-cache/inferencex/agentx-hub}" +HF_HUB_CACHE_CONTAINER="${HF_HUB_CACHE:-/hf-hub-cache}" + +fail_with() { + local rc="$1" + shift + echo "ERROR: $*" >&2 + exit "$rc" +} + +fail() { + fail_with 1 "$@" +} + +require_exact() { + local name="$1" expected="$2" + if [[ "${!name:-}" != "$expected" ]]; then + fail "the native MI300X path requires $name=$expected, got '${!name:-}'" + fi +} + +require_set() { + local name + for name in "$@"; do + if [[ -z "${!name:-}" ]]; then + fail "$name must be set" + fi + done +} + +require_set GITHUB_WORKSPACE RUNNER_NAME IMAGE MODEL RESULT_FILENAME CONC_LIST + +require_exact IS_MULTINODE true +require_exact IS_AGENTIC 1 +require_exact SCENARIO_TYPE agentic-coding +require_exact FRAMEWORK vllm +require_exact MODEL_PREFIX kimik3 +require_exact PRECISION fp4 +require_exact SPEC_DECODING none +require_exact DISAGG false +require_exact PREFILL_EP 1 +require_exact PREFILL_DP_ATTN false +require_exact DECODE_EP 1 +require_exact DECODE_DP_ATTN false + +if [[ "${PREFILL_TP:-}" != "8" || "${PREFILL_PP_SIZE:-}" != "2" || + "${DECODE_TP:-}" != "8" || "${DECODE_PP_SIZE:-}" != "2" ]]; then + fail "the native MI300X path serves only TP8 x PP2, got prefill TP${PREFILL_TP:-?} x PP${PREFILL_PP_SIZE:-?} and decode TP${DECODE_TP:-?} x PP${DECODE_PP_SIZE:-?}" +fi +if [[ "${PREFILL_NUM_WORKERS:-}" != "1" || "${DECODE_NUM_WORKERS:-}" != "0" ]]; then + fail "the native MI300X path is aggregated: it needs 1 prefill worker and 0 decode workers, got ${PREFILL_NUM_WORKERS:-?}P/${DECODE_NUM_WORKERS:-?}D" +fi + +read -r -a CONCURRENCIES <<< "$CONC_LIST" +if [[ "${#CONCURRENCIES[@]}" -ne 1 ]]; then + fail "one concurrency per allocation is required, got CONC_LIST='$CONC_LIST'" +fi +CONC_VALUE="${CONCURRENCIES[0]}" +case "$CONC_VALUE" in + 1|2|4|8) ;; + *) fail "concurrency must be 1, 2, 4, or 8, got '$CONC_VALUE'" ;; +esac + +for knob in KIMIK3_STARTUP_TIMEOUT_SECONDS KIMIK3_HEALTH_POLL_SECONDS \ + KIMIK3_CLEANUP_TIMEOUT_SECONDS KIMIK3_CLEANUP_POLL_SECONDS \ + KIMIK3_PRESCANCEL_TIMEOUT_SECONDS KIMIK3_SLURM_TIME_MINUTES; do + if [[ ! "${!knob}" =~ ^[1-9][0-9]*$ ]]; then + fail "$knob must be a positive integer, got '${!knob}'" + fi +done + +if [[ -n "${AITER_SITUV2_A8W4+set}" ]]; then + if [[ "$AITER_SITUV2_A8W4" != "0" && "$AITER_SITUV2_A8W4" != "1" ]]; then + fail "AITER_SITUV2_A8W4 must be 0 or 1 when set, got '$AITER_SITUV2_A8W4'" + fi +fi + +JOB_ID="" +HEAD_NODE="" +SERVER_PID="" +CLIENT_PID="" +SERVER_LOG_DIR="" +SERVER_LOG="" +SERVER_RC_FILE="" +SERVER_SRUN_PID_FILE="" +EXTRACT_DIR="" +HANDOFF_HOST="" +SCRATCH_HOST="" +CLEANUP_DONE=0 + +dump_server_log() { + if [[ -n "$SERVER_LOG" && -s "$SERVER_LOG" ]]; then + echo "=== last 200 lines of the vLLM server log ===" + tail -n 200 "$SERVER_LOG" + echo "=============================================" + fi +} + +package_server_logs() { + if [[ -n "$SERVER_LOG_DIR" ]]; then + echo "[cleanup] packaging server logs" + bundle_server_logs "$SERVER_LOG_DIR" \ + "$GITHUB_WORKSPACE/multinode_server_logs.tar.gz" + if [[ ! -s "$GITHUB_WORKSPACE/multinode_server_logs.tar.gz" ]]; then + echo "[cleanup] WARNING: the server produced no log to package" + fi + fi +} + +run_bounded() { + local timeout_seconds="$1" + shift + local marker waited pid + marker=$(mktemp) + ( "$@" >/dev/null 2>&1; printf 'done' > "$marker" ) & + pid=$! + waited=0 + while [[ ! -s "$marker" ]] && (( waited < timeout_seconds )); do + sleep "$KIMIK3_CLEANUP_POLL_SECONDS" + waited=$(( waited + KIMIK3_CLEANUP_POLL_SECONDS )) + done + kill "$pid" 2>/dev/null || true + wait "$pid" 2>/dev/null || true + rm -f "$marker" +} + +cleanup() { + if [[ "$CLEANUP_DONE" == "1" ]]; then + return 0 + fi + CLEANUP_DONE=1 + set +e + + if [[ -n "$CLIENT_PID" ]]; then + echo "[cleanup] stopping AgentX client" + kill "$CLIENT_PID" 2>/dev/null + wait "$CLIENT_PID" 2>/dev/null + CLIENT_PID="" + fi + + if [[ -n "$SERVER_PID" ]]; then + echo "[cleanup] stopping server step" + if [[ -n "$SERVER_SRUN_PID_FILE" && -s "$SERVER_SRUN_PID_FILE" ]]; then + kill "$(cat "$SERVER_SRUN_PID_FILE")" 2>/dev/null + fi + kill "$SERVER_PID" 2>/dev/null + wait "$SERVER_PID" 2>/dev/null + SERVER_PID="" + fi + + if [[ -n "$JOB_ID" && -n "$HEAD_NODE" && -n "$SCRATCH_HOST" ]]; then + echo "[cleanup] removing node-local scratch $SCRATCH_HOST" + run_bounded "$KIMIK3_PRESCANCEL_TIMEOUT_SECONDS" \ + srun --overlap --jobid="$JOB_ID" --nodes=1 --ntasks=1 \ + --nodelist="$HEAD_NODE" rm -rf "$SCRATCH_HOST" + fi + + if [[ -n "$JOB_ID" ]]; then + echo "[cleanup] cancelling allocation $JOB_ID" + scancel "$JOB_ID" 2>/dev/null + local waited=0 + while slurm_job_is_active "$JOB_ID"; do + if (( waited >= KIMIK3_CLEANUP_TIMEOUT_SECONDS )); then + echo "[cleanup] WARNING: job $JOB_ID still present after ${waited}s" + break + fi + sleep "$KIMIK3_CLEANUP_POLL_SECONDS" + waited=$(( waited + KIMIK3_CLEANUP_POLL_SECONDS )) + done + fi + + package_server_logs + + [[ -n "$HANDOFF_HOST" ]] && rm -f "$HANDOFF_HOST" + [[ -n "$EXTRACT_DIR" ]] && rm -rf "$EXTRACT_DIR" + [[ -n "$SERVER_LOG_DIR" ]] && rm -rf "$SERVER_LOG_DIR" + return 0 +} + +trap 'rc=$?; cleanup; exit "$rc"' EXIT +trap 'cleanup; exit 130' INT +trap 'cleanup; exit 143' TERM +trap 'cleanup; exit 129' HUP + +set -x + +salloc_stdout=$( + salloc \ + --parsable \ + --partition=compute \ + --exclude=chi-mi300x-049,chi-mi300x-121 \ + --nodes=2 \ + --ntasks-per-node=1 \ + --gres=gpu:8 \ + --cpus-per-task=256 \ + --exclusive \ + --time="$KIMIK3_SLURM_TIME_MINUTES" \ + --no-shell \ + --job-name="$RUNNER_NAME" +) +JOB_ID=$(printf '%s' "$salloc_stdout" | tr -d '[:space:]' | sed -n 's/^\([0-9][0-9]*\).*/\1/p') +if [[ -z "$JOB_ID" ]]; then + fail "salloc did not return a job ID (stdout: '$salloc_stdout')" +fi +echo "Allocated Slurm job $JOB_ID" + +head_node_output=$( + srun --jobid="$JOB_ID" --nodes=2 --ntasks=2 --ntasks-per-node=1 \ + bash -c 'if [[ "$SLURM_PROCID" == "0" ]]; then hostname; fi' +) +HEAD_NODE=$(printf '%s\n' "$head_node_output" | awk 'NF {print $1; exit}') +if [[ -z "$HEAD_NODE" ]]; then + fail "could not resolve the rank-0 hostname for job $JOB_ID" +fi +echo "Rank 0 runs on $HEAD_NODE" + +preflight_output=$( + srun --jobid="$JOB_ID" --nodes=2 --ntasks=2 --ntasks-per-node=1 --export=ALL \ + bash runners/mi300x_native_node_preflight.sh +) +REVISION=$(printf '%s\n' "$preflight_output" | python3 -c ' +import sys + +records = [] +for line in sys.stdin: + if not line.startswith("INFERENCEX_KIMIK3_PREFLIGHT "): + continue + records.append( + dict(item.split("=", 1) for item in line.split()[1:] if "=" in item) + ) + +if len(records) != 2: + sys.exit(f"ERROR: expected one preflight record per node, got {len(records)}") + +hostnames = {record.get("hostname") for record in records} +if len(hostnames) != 2: + sys.exit(f"ERROR: preflight covered {len(hostnames)} distinct node(s): {sorted(hostnames)}") + +revisions = {record.get("revision") for record in records} +if len(revisions) != 1: + sys.exit(f"ERROR: nodes hold different model revisions: {sorted(revisions)}") + +for record in records: + host = record.get("hostname") + if record.get("gpu_count") != "8" or record.get("gpu_arch") != "gfx942": + sys.exit(f"ERROR: node {host} is not 8x gfx942: {record}") + if int(record.get("squash_size_bytes") or 0) <= 0: + sys.exit(f"ERROR: node {host} has no valid container image") + +sizes = {record.get("squash_size_bytes") for record in records} +if len(sizes) != 1: + sys.exit(f"ERROR: nodes imported different builds of the image: {sorted(sizes)} bytes") + +print(revisions.pop()) +') +echo "Both nodes verified at model revision $REVISION" + +IMAGE_PATH="$KIMIK3_SQUASH_DIR/$(printf '%s' "$IMAGE" | sed 's/[\/:@#]/_/g').sqsh" +MODEL_SNAPSHOT="$KIMIK3_MODEL_CACHE_ROOT/snapshots/$REVISION" +MODEL_CONTAINER_PATH="/models/Kimi-K3" + +export MULTINODE_NODE_COUNT=2 +export MULTINODE_GPUS_PER_NODE=8 +export MULTINODE_MASTER_ADDR="$HEAD_NODE" +export MODEL_PATH="$MODEL_CONTAINER_PATH" + +SERVER_LOG_DIR=$(mktemp -d) +SERVER_LOG="$SERVER_LOG_DIR/vllm_server.log" +SERVER_RC_FILE="$SERVER_LOG_DIR/server.rc" +SERVER_SRUN_PID_FILE="$SERVER_LOG_DIR/server.srun.pid" + +{ + set +e + srun --jobid="$JOB_ID" \ + --nodes=2 \ + --ntasks=2 \ + --ntasks-per-node=1 \ + --kill-on-bad-exit=1 \ + --container-image="$IMAGE_PATH" \ + --container-remap-root \ + --container-writable \ + --no-container-mount-home \ + --no-container-entrypoint \ + --container-workdir=/workspace \ + --container-mounts="$GITHUB_WORKSPACE:/workspace,$MODEL_SNAPSHOT:$MODEL_CONTAINER_PATH:ro,/dev/kfd:/dev/kfd,/dev/dri:/dev/dri" \ + --export=ALL \ + bash -c 'export MULTINODE_NODE_RANK="$SLURM_PROCID"; exec bash /workspace/benchmarks/multi_node/agentic/kimik3_fp4_mi300x_vllm.sh' \ + > "$SERVER_LOG" 2>&1 & + srun_pid=$! + printf '%s\n' "$srun_pid" > "$SERVER_SRUN_PID_FILE" + wait "$srun_pid" + printf '%s\n' "$?" > "$SERVER_RC_FILE" +} & +SERVER_PID=$! +echo "Started both server ranks; logging to $SERVER_LOG" + +HEALTH_URL="http://${HEAD_NODE}:${PORT}/health" +startup_deadline=$(( $(date +%s) + KIMIK3_STARTUP_TIMEOUT_SECONDS )) +set +x +while true; do + if [[ -s "$SERVER_RC_FILE" ]]; then + server_rc=$(tr -d '[:space:]' < "$SERVER_RC_FILE") + wait "$SERVER_PID" 2>/dev/null || true + SERVER_PID="" + dump_server_log + fail "the vLLM server step exited with code ${server_rc:-unknown} before becoming healthy" + fi + if curl -sf --max-time 10 "$HEALTH_URL" > /dev/null 2>&1; then + echo "Server is healthy at $HEALTH_URL" + break + fi + if (( $(date +%s) >= startup_deadline )); then + dump_server_log + fail "the vLLM server did not become healthy within ${KIMIK3_STARTUP_TIMEOUT_SECONDS}s" + fi + sleep "$KIMIK3_HEALTH_POLL_SECONDS" +done +set -x + +SCRATCH_HOST="$KIMIK3_SQUASH_DIR/.runs/${JOB_ID}-conc${CONC_VALUE}" +srun --overlap --jobid="$JOB_ID" --nodes=1 --ntasks=1 --nodelist="$HEAD_NODE" \ + mkdir -p "$SCRATCH_HOST/output" "$SCRATCH_HOST/agentic" "$HF_HUB_CACHE_MOUNT" + +HANDOFF_NAME="multinode_agentic_handoff.tar.gz" +HANDOFF_HOST="$GITHUB_WORKSPACE/$HANDOFF_NAME" +: > "$HANDOFF_HOST" + +export KIMIK3_HANDOFF_PATH="/workspace/$HANDOFF_NAME" +export INFMAX_CONTAINER_WORKSPACE=/workspace +export RESULT_DIR=/results/agentic +export AGENTIC_OUTPUT_DIR=/results/output +export AIPERF_SERVER_METRICS_URLS="http://${HEAD_NODE}:${PORT}/metrics" + +CLIENT_WORKER_SCRIPT='set -uo pipefail +mkdir -p /results/output /results/agentic +bash /workspace/benchmarks/multi_node/agentic_srt.sh +client_rc=$? +shopt -s nullglob +cd /results +aggregates=(output/*.json) +tar czf "$KIMIK3_HANDOFF_PATH" ${aggregates[@]+"${aggregates[@]}"} agentic +exit "$client_rc"' + +HF_HUB_CACHE="$HF_HUB_CACHE_CONTAINER" \ +srun --overlap --jobid="$JOB_ID" --nodes=1 --ntasks=1 --nodelist="$HEAD_NODE" \ + --container-image="$IMAGE_PATH" \ + --container-remap-root \ + --container-writable \ + --no-container-mount-home \ + --no-container-entrypoint \ + --container-workdir=/workspace \ + --container-mounts="$GITHUB_WORKSPACE:/workspace,$SCRATCH_HOST:/results,$HF_HUB_CACHE_MOUNT:$HF_HUB_CACHE_CONTAINER,$MODEL_SNAPSHOT:$MODEL_CONTAINER_PATH:ro,/dev/kfd:/dev/kfd,/dev/dri:/dev/dri" \ + --export=ALL \ + bash -c "$CLIENT_WORKER_SCRIPT" & +CLIENT_PID=$! +client_rc=0 +wait "$CLIENT_PID" || client_rc=$? +CLIENT_PID="" + +if [[ ! -s "$HANDOFF_HOST" ]]; then + fail_with "$(( client_rc == 0 ? 1 : client_rc ))" \ + "the AgentX client produced no handoff archive (client exit code $client_rc)" +fi + +archive_entries=$(tar tzf "$HANDOFF_HOST") +if printf '%s\n' "$archive_entries" | grep -q '^/'; then + fail "the handoff archive contains absolute paths" +fi +if printf '%s\n' "$archive_entries" | grep -Eq '(^|/)\.\.(/|$)'; then + fail "the handoff archive contains parent-directory components" +fi + +EXTRACT_DIR=$(mktemp -d) +tar xzf "$HANDOFF_HOST" -C "$EXTRACT_DIR" --no-same-owner --no-same-permissions + +AGGREGATE="$EXTRACT_DIR/output/${RESULT_FILENAME}_conc${CONC_VALUE}.json" +if [[ ! -f "$AGGREGATE" ]]; then + fail_with "$(( client_rc == 0 ? 1 : client_rc ))" \ + "the handoff archive is missing ${RESULT_FILENAME}_conc${CONC_VALUE}.json (client exit code $client_rc)" +fi +copy_to_workspace "$AGGREGATE" "$GITHUB_WORKSPACE/$(basename "$AGGREGATE")" + +RAW_SOURCE="$EXTRACT_DIR/agentic/conc_${CONC_VALUE}" +RAW_DEST="$GITHUB_WORKSPACE/LOGS/agentic/conc_${CONC_VALUE}" +mkdir -p "$RAW_DEST" +if [[ -d "$RAW_SOURCE" ]]; then + cp -R "$RAW_SOURCE/." "$RAW_DEST/" +fi + +rm -f "$HANDOFF_HOST" +HANDOFF_HOST="" + +if (( client_rc != 0 )); then + fail_with "$client_rc" "the AgentX client exited with code $client_rc" +fi + +echo "Native MI300X Kimi K3 AgentX run complete at concurrency ${CONC_VALUE}" diff --git a/runners/launch_mi300x-amds.sh b/runners/launch_mi300x-amds.sh index fdd03889a0..d66ca501f2 100644 --- a/runners/launch_mi300x-amds.sh +++ b/runners/launch_mi300x-amds.sh @@ -1,6 +1,10 @@ #!/usr/bin/env bash set -eo pipefail +if [[ "${NATIVE_MULTINODE:-0}" == "1" ]]; then + exec bash runners/launch_mi300x-amds-native-multinode.sh +fi + export HF_HUB_CACHE_MOUNT="/raid/hf-hub-cache/" export PORT=8888 diff --git a/runners/mi300x_native_node_preflight.sh b/runners/mi300x_native_node_preflight.sh new file mode 100755 index 0000000000..0b06f9fcc1 --- /dev/null +++ b/runners/mi300x_native_node_preflight.sh @@ -0,0 +1,115 @@ +#!/usr/bin/env bash +set -euo pipefail + +set -x + +KIMIK3_MODEL_CACHE_ROOT="${KIMIK3_MODEL_CACHE_ROOT:-/raid/hf-hub-cache/models--moonshotai--Kimi-K3}" +KIMIK3_SQUASH_DIR="${KIMIK3_SQUASH_DIR:-/raid/hf-hub-cache/inferencex/squash}" +KIMIK3_IMAGE="${KIMIK3_IMAGE:-${IMAGE:?IMAGE must be set}}" + +fail() { + echo "ERROR: [$(hostname)] $*" >&2 + exit 1 +} + +TMP_SQUASH="" +remove_tmp_squash() { + if [[ -n "$TMP_SQUASH" ]]; then + rm -f "$TMP_SQUASH" + fi + return 0 +} +trap remove_tmp_squash EXIT +trap 'remove_tmp_squash; exit 130' INT +trap 'remove_tmp_squash; exit 143' TERM +trap 'remove_tmp_squash; exit 129' HUP + +gpu_arch_lines=$( + rocminfo 2>/dev/null | + sed -n 's/^[[:space:]]*Name:[[:space:]]*\(gfx[0-9a-z]*\)[[:space:]]*$/\1/p' || true +) +gpu_count=$(printf '%s\n' "$gpu_arch_lines" | grep -c '^gfx' || true) +gfx942_count=$(printf '%s\n' "$gpu_arch_lines" | grep -c '^gfx942$' || true) +if [[ "$gpu_count" != "8" || "$gfx942_count" != "8" ]]; then + fail "this node must expose exactly 8 gfx942 GPUs; found ${gfx942_count} gfx942 among ${gpu_count} GPU agents" +fi + +refs_main="$KIMIK3_MODEL_CACHE_ROOT/refs/main" +if [[ ! -f "$refs_main" ]]; then + fail "missing model revision pointer $refs_main; stage the snapshot on this node first" +fi +revision=$(tr -d '[:space:]' < "$refs_main") +if ! [[ "$revision" =~ ^[0-9a-f]{40}$ ]]; then + fail "$refs_main must hold a 40-character revision, found '$revision'" +fi + +snapshot_dir="$KIMIK3_MODEL_CACHE_ROOT/snapshots/$revision" +if [[ ! -d "$snapshot_dir" ]]; then + fail "missing model snapshot directory $snapshot_dir" +fi +if [[ ! -f "$snapshot_dir/config.json" ]]; then + fail "missing $snapshot_dir/config.json" +fi + +weight_index="$snapshot_dir/model.safetensors.index.json" +if [[ ! -f "$weight_index" ]]; then + fail "missing model.safetensors.index.json in $snapshot_dir" +fi +python3 - "$weight_index" "$snapshot_dir" <<'PY' +import json +import os +import sys + +index_path, snapshot_dir = sys.argv[1], sys.argv[2] +with open(index_path) as handle: + weight_map = (json.load(handle) or {}).get("weight_map") or {} +if not weight_map: + sys.exit(f"ERROR: {index_path} declares no weight_map entries") + +incomplete = sorted( + shard + for shard in set(weight_map.values()) + if not os.path.isfile(os.path.join(snapshot_dir, shard)) + or os.path.getsize(os.path.join(snapshot_dir, shard)) == 0 +) +if incomplete: + sys.exit( + f"ERROR: missing weight shard(s) in {snapshot_dir}: " + ", ".join(incomplete) + ) +PY + +mkdir -p "$KIMIK3_SQUASH_DIR" +squash_file="$KIMIK3_SQUASH_DIR/$(printf '%s' "$KIMIK3_IMAGE" | sed 's/[\/:@#]/_/g').sqsh" + +export ENROOT_CACHE_PATH="$KIMIK3_SQUASH_DIR/.enroot-cache" +export ENROOT_TEMP_PATH="$KIMIK3_SQUASH_DIR/.enroot-temp" +mkdir -p "$ENROOT_CACHE_PATH" "$ENROOT_TEMP_PATH" + +lock_file="${squash_file}.lock" +exec 9>"$lock_file" +if ! flock -w "${KIMIK3_IMAGE_LOCK_TIMEOUT_SECONDS:-3600}" 9; then + fail "timed out waiting for the image lock $lock_file" +fi + +if [[ -s "$squash_file" ]] && unsquashfs -s "$squash_file" >/dev/null 2>&1; then + echo "[$(hostname)] reusing validated image $squash_file" +else + rm -f "$squash_file" + TMP_SQUASH=$(mktemp "${squash_file}.XXXXXX") + rm -f "$TMP_SQUASH" + echo "[$(hostname)] importing $KIMIK3_IMAGE into $squash_file" + enroot import -o "$TMP_SQUASH" "docker://$KIMIK3_IMAGE" + if ! unsquashfs -s "$TMP_SQUASH" >/dev/null 2>&1; then + fail "imported image $KIMIK3_IMAGE failed unsquashfs validation" + fi + mv "$TMP_SQUASH" "$squash_file" + TMP_SQUASH="" +fi + +squash_size_bytes=$(wc -c < "$squash_file" | tr -d '[:space:]') +if [[ "$squash_size_bytes" -le 0 ]]; then + fail "validated image $squash_file is empty" +fi + +printf 'INFERENCEX_KIMIK3_PREFLIGHT hostname=%s revision=%s gpu_count=%s gpu_arch=gfx942 squash_size_bytes=%s\n' \ + "$(hostname)" "$revision" "$gpu_count" "$squash_size_bytes" diff --git a/utils/matrix_logic/test_kimik3_mi300x_native.py b/utils/matrix_logic/test_kimik3_mi300x_native.py new file mode 100644 index 0000000000..36328322ed --- /dev/null +++ b/utils/matrix_logic/test_kimik3_mi300x_native.py @@ -0,0 +1,765 @@ +"""CPU-only gates for the Kimi K3 MI300X native multi-node AgentX path.""" + +from __future__ import annotations + +import argparse +import json +import os +import shutil +import signal +import subprocess +import sys +import time +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(REPO_ROOT / "utils" / "matrix_logic")) + +from generate_sweep_configs import generate_test_config_sweep # noqa: E402 +from validation import load_config_files, load_runner_file # noqa: E402 + +CONFIG_KEY = "kimik3-fp4-mi300x-vllm-agentic" +SERVER_SCRIPT = ( + REPO_ROOT / "benchmarks" / "multi_node" / "agentic" / "kimik3_fp4_mi300x_vllm.sh" +) +PREFLIGHT_SCRIPT = REPO_ROOT / "runners" / "mi300x_native_node_preflight.sh" +DEFAULT_LAUNCHER = REPO_ROOT / "runners" / "launch_mi300x-amds.sh" +NATIVE_LAUNCHER = REPO_ROOT / "runners" / "launch_mi300x-amds-native-multinode.sh" +IMAGE = "vllm/vllm-openai-rocm:kimi-k3" +REVISION = "0123456789abcdef0123456789abcdef01234567" +OTHER_REVISION = "89abcdef0123456789abcdef0123456789abcdef" +STAGING_COMMANDS = ("hf download", "huggingface-cli download", "wget", "curl") +JOB_ID = "4242" +RESULT_FILENAME = "kimik3_agentic_prefill-tp8-pp2_conc4_mi300x-amds_00" + + +def generate_kimik3_matrix() -> list[dict]: + """Generate the sweep rows the workflow would dispatch for this config key.""" + configs = load_config_files([str(REPO_ROOT / "configs" / "amd-master.yaml")]) + runners = load_runner_file(str(REPO_ROOT / "configs" / "runners.yaml")) + args = argparse.Namespace( + config_keys=[CONFIG_KEY], + seq_lens=None, + conc=None, + scenario_type=["agentic-coding"], + runner_node_filter=None, + ) + return generate_test_config_sweep(args, configs, runners) + + +def test_kimik3_matrix_is_exactly_four_tp8_pp2_aggregate_jobs() -> None: + """One aggregate job per concurrency, all TP8 x PP2, with no AITER default.""" + rows = generate_kimik3_matrix() + + assert [row["conc"] for row in rows] == [[1], [2], [4], [8]] + assert {row["runner"] for row in rows} == {"cluster:mi300x-amds"} + assert {row["framework"] for row in rows} == {"vllm"} + assert {row["disagg"] for row in rows} == {False} + assert { + ( + row["prefill"]["num-worker"], + row["prefill"]["tp"], + row["prefill"]["pp"], + row["prefill"]["ep"], + row["prefill"]["dp-attn"], + row["decode"]["num-worker"], + row["decode"]["tp"], + row["decode"]["pp"], + row["decode"]["ep"], + row["decode"]["dp-attn"], + ) + for row in rows + } == {(1, 8, 2, 1, False, 0, 8, 2, 1, False)} + settings = rows[0]["prefill"]["additional-settings"] + assert settings == ["NATIVE_MULTINODE=1"] + assert all("AITER_SITUV2_A8W4" not in setting for setting in settings) + + +def server_env(rank: int = 0) -> dict[str, str]: + """Build the environment the workflow exports for one vLLM rank.""" + env = { + **os.environ, + "MODEL": "moonshotai/Kimi-K3", + "MODEL_PATH": "/models/Kimi-K3", + "PORT": "8888", + "CONC_LIST": "4", + "KV_OFFLOADING": "none", + "PREFILL_NUM_WORKERS": "1", + "PREFILL_TP": "8", + "PREFILL_PP_SIZE": "2", + "PREFILL_EP": "1", + "PREFILL_DP_ATTN": "false", + "DECODE_NUM_WORKERS": "0", + "MULTINODE_NODE_COUNT": "2", + "MULTINODE_GPUS_PER_NODE": "8", + "MULTINODE_NODE_RANK": str(rank), + "MULTINODE_MASTER_ADDR": "node-a", + "KIMIK3_VLLM_DRY_RUN": "1", + } + env.pop("AITER_SITUV2_A8W4", None) + return env + + +def run_server(env: dict[str, str]) -> subprocess.CompletedProcess: + """Run the rank entrypoint in dry-run mode and capture the rendered command.""" + return subprocess.run( + ["bash", str(SERVER_SCRIPT)], + cwd=str(REPO_ROOT), + env=env, + capture_output=True, + text=True, + ) + + +@pytest.mark.parametrize("rank", [0, 1]) +def test_each_rank_serves_tp8_pp2_and_only_rank_zero_owns_the_endpoint( + rank: int, +) -> None: + """Both ranks share one topology; --headless is the only difference.""" + result = run_server(server_env(rank)) + assert result.returncode == 0, result.stderr + assert "--tensor-parallel-size 8" in result.stdout + assert "--pipeline-parallel-size 2" in result.stdout + assert "--nnodes 2" in result.stdout + assert f"--node-rank {rank}" in result.stdout + assert "--master-addr node-a" in result.stdout + assert ("--headless" in result.stdout) is (rank == 1) + assert "FLASHMLA" not in result.stdout + assert "FLASHINFER" not in result.stdout + + +@pytest.mark.parametrize( + ("name", "value", "message"), + [ + ("PREFILL_TP", "16", "TP8 x PP2"), + ("PREFILL_PP_SIZE", "1", "TP8 x PP2"), + ("PREFILL_EP", "8", "EP1"), + ("DECODE_NUM_WORKERS", "1", "aggregated"), + ("CONC_LIST", "4 8", "one concurrency"), + ("CONC_LIST", "16", "1, 2, 4, or 8"), + ("AITER_SITUV2_A8W4", "auto", "0 or 1"), + ], +) +def test_server_rejects_out_of_contract_values( + name: str, value: str, message: str +) -> None: + """Reject any topology this bring-up has not been validated for.""" + env = server_env() + env[name] = value + result = run_server(env) + assert result.returncode != 0 + assert message in result.stderr + + +def test_aiter_mode_is_not_defaulted_and_accepts_both_modes() -> None: + """gfx942 has no tuned MoE tables, so a8w4 stays an operator input.""" + unset_result = run_server(server_env()) + assert "AITER_SITUV2_A8W4=unset" in unset_result.stdout + for value in ("0", "1"): + env = server_env() + env["AITER_SITUV2_A8W4"] = value + result = run_server(env) + assert result.returncode == 0 + assert f"AITER_SITUV2_A8W4={value}" in result.stdout + + +def write_executable(path: Path, body: str) -> None: + """Write a fake executable onto the test PATH.""" + path.write_text(body) + path.chmod(0o755) + + +def rocminfo_output(gpu_count: int, gpu_arch: str) -> str: + """Mimic rocminfo closely enough to catch sloppy agent counting.""" + blocks = [ + "HSA Agents\n==========\n", + "Agent 1\n*******\n" + " Name: AMD EPYC 9654 96-Core Processor\n" + " Marketing Name: AMD EPYC 9654 96-Core Processor\n" + " Device Type: CPU\n", + ] + for index in range(gpu_count): + blocks.append( + f"Agent {index + 2}\n*******\n" + f" Name: {gpu_arch}\n" + " Marketing Name: AMD Instinct MI300X\n" + " Device Type: GPU\n" + ) + blocks.append( + "*** Done ***\nISA Info:\n ISA 1\n" + f" Name: amdgcn-amd-amdhsa--{gpu_arch}:sramecc+:xnack-\n" + ) + return "".join(blocks) + + +def make_node( + tmp_path: Path, + *, + gpu_count: int = 8, + gpu_arch: str = "gfx942", + with_main_ref: bool = True, + with_weight_index: bool = True, + with_indexed_shard: bool = True, +) -> dict[str, object]: + """Build one node's model cache, squash tree, and external-command fakes.""" + cache_root = tmp_path / "raid" / "hf-hub-cache" / "models--moonshotai--Kimi-K3" + snapshot_dir = cache_root / "snapshots" / REVISION + snapshot_dir.mkdir(parents=True) + (snapshot_dir / "config.json").write_text('{"model_type": "kimi_k3"}\n') + if with_main_ref: + (cache_root / "refs").mkdir(parents=True) + (cache_root / "refs" / "main").write_text(f"{REVISION}\n") + if with_weight_index: + (snapshot_dir / "model.safetensors.index.json").write_text( + json.dumps( + { + "metadata": {}, + "weight_map": { + "model.layers.0.weight": "model-00001-of-00001.safetensors" + }, + } + ) + ) + if with_indexed_shard: + (snapshot_dir / "model-00001-of-00001.safetensors").write_text("weights\n") + + squash_dir = tmp_path / "raid" / "hf-hub-cache" / "inferencex" / "squash" + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + cmd_log = tmp_path / "commands.log" + cmd_log.write_text("") + + write_executable( + bin_dir / "rocminfo", + "#!/usr/bin/env bash\ncat <<'ROCMINFO_EOF'\n" + + rocminfo_output(gpu_count, gpu_arch) + + "ROCMINFO_EOF\n", + ) + write_executable( + bin_dir / "unsquashfs", + "#!/usr/bin/env bash\n" + 'echo "unsquashfs $*" >> "$FAKE_CMD_LOG"\n' + "for arg in \"$@\"; do target=\"$arg\"; done\n" + '[[ -s "$target" ]] || exit 1\n' + "exit 0\n", + ) + write_executable( + bin_dir / "enroot", + "#!/usr/bin/env bash\n" + 'echo "enroot $*" >> "$FAKE_CMD_LOG"\n' + '[[ "${1:-}" == "import" ]] || exit 1\n' + "out=\"\"\n" + "while [[ $# -gt 0 ]]; do\n" + ' case "$1" in\n' + ' -o) out="$2"; shift 2 ;;\n' + " *) shift ;;\n" + " esac\n" + "done\n" + '[[ -n "$out" ]] || exit 1\n' + "printf 'fake-squash-image\\n' > \"$out\"\n", + ) + if shutil.which("flock") is None: + write_executable( + bin_dir / "flock", + "#!/usr/bin/env bash\n" + 'echo "flock $*" >> "$FAKE_CMD_LOG"\n' + "exit 0\n", + ) + + env = { + **os.environ, + "PATH": f"{bin_dir}{os.pathsep}{os.environ['PATH']}", + "KIMIK3_MODEL_CACHE_ROOT": str(cache_root), + "KIMIK3_SQUASH_DIR": str(squash_dir), + "IMAGE": IMAGE, + "FAKE_CMD_LOG": str(cmd_log), + "KIMIK3_IMAGE_LOCK_TIMEOUT_SECONDS": "5", + } + return {"squash_dir": squash_dir, "cmd_log": cmd_log, "env": env} + + +def run_preflight(node: dict[str, object]) -> subprocess.CompletedProcess: + """Run the real preflight script against a staged node.""" + return subprocess.run( + ["bash", str(PREFLIGHT_SCRIPT)], + cwd=str(REPO_ROOT), + env=node["env"], + capture_output=True, + text=True, + ) + + +def preflight_records(stdout: str) -> list[str]: + """Extract the machine-readable preflight records from stdout.""" + return [ + line + for line in stdout.splitlines() + if line.startswith("INFERENCEX_KIMIK3_PREFLIGHT ") + ] + + +def test_preflight_imports_and_validates_image_in_node_local_tree( + tmp_path: Path, +) -> None: + """Import the image node-locally and never download the checkpoint.""" + node = make_node(tmp_path) + result = run_preflight(node) + + assert result.returncode == 0, result.stderr + records = preflight_records(result.stdout) + assert len(records) == 1 + record = records[0] + assert f"revision={REVISION}" in record + assert "gpu_count=8" in record + assert "gpu_arch=gfx942" in record + + squashes = sorted(node["squash_dir"].glob("*.sqsh")) + assert len(squashes) == 1 + assert squashes[0].stat().st_size > 0 + assert f"squash_size_bytes={squashes[0].stat().st_size}" in record + + command_log = node["cmd_log"].read_text() + assert f"docker://{IMAGE}" in command_log + assert not any(command in command_log for command in STAGING_COMMANDS) + script_source = PREFLIGHT_SCRIPT.read_text() + assert not any(command in script_source for command in STAGING_COMMANDS) + + +def test_preflight_reuses_a_valid_squash_without_import(tmp_path: Path) -> None: + """A later job on the same node reuses the already validated image.""" + node = make_node(tmp_path) + assert run_preflight(node).returncode == 0 + + node["cmd_log"].write_text("") + result = run_preflight(node) + + assert result.returncode == 0, result.stderr + assert len(preflight_records(result.stdout)) == 1 + assert "enroot import" not in node["cmd_log"].read_text() + + +@pytest.mark.parametrize( + ("node_kwargs", "message"), + [ + ({"gpu_count": 7}, "exactly 8 gfx942"), + ({"gpu_arch": "gfx950"}, "exactly 8 gfx942"), + ({"with_main_ref": False}, "refs/main"), + ({"with_weight_index": False}, "model.safetensors.index.json"), + ({"with_indexed_shard": False}, "missing weight shard"), + ], +) +def test_preflight_rejects_unusable_node_state( + tmp_path: Path, node_kwargs: dict, message: str +) -> None: + """Reject a node before the job starts loading 1.5 TB of weights onto it.""" + result = run_preflight(make_node(tmp_path, **node_kwargs)) + assert result.returncode != 0 + assert message in result.stderr + + +FAKE_SALLOC = """#!/usr/bin/env bash +echo "salloc $*" >> "$FAKE_CMD_LOG" +echo "salloc: Granted job allocation {job_id}" >&2 +if [[ "$*" == *--parsable* ]]; then + echo "{job_id}" +fi +""" + +FAKE_SCANCEL = """#!/usr/bin/env bash +echo "scancel $*" >> "$FAKE_CMD_LOG" +""" + +FAKE_SQUEUE = """#!/usr/bin/env bash +echo "squeue $*" >> "$FAKE_CMD_LOG" +""" + +FAKE_CURL = """#!/usr/bin/env bash +echo "curl $*" >> "$FAKE_CMD_LOG" +[[ "${FAKE_HEALTH:-ok}" == "ok" ]] +""" + +FAKE_SRUN = """#!/usr/bin/env bash +args="$*" +printf 'srun %s\\n' "$(printf '%s' "$args" | tr '\\n' ' ')" >> "$FAKE_CMD_LOG" + +if [[ "$args" == *mi300x_native_node_preflight.sh* ]]; then + cat "$FAKE_PREFLIGHT_OUTPUT" + exit 0 +fi + +if [[ "$args" == *--kill-on-bad-exit=1* ]]; then + mode="${FAKE_SERVER_MODE:-alive}" + if [[ "$mode" == exit:* ]]; then + exit "${mode#exit:}" + fi + printf '%s\\n' "$$" > "$FAKE_SERVER_PID_FILE" + sleep 600 + exit 0 +fi + +if [[ "$args" == *agentic_srt.sh* ]]; then + handoff="$GITHUB_WORKSPACE/$(basename "$KIMIK3_HANDOFF_PATH")" + staging=$(mktemp -d) + mkdir -p "$staging/output" "$staging/agentic/conc_${CONC_LIST}/aiperf_artifacts" + printf '{"num_requests_successful": 12}\\n' \ + > "$staging/output/${RESULT_FILENAME}_conc${CONC_LIST}.json" + printf '{}\\n' \ + > "$staging/agentic/conc_${CONC_LIST}/aiperf_artifacts/profile_export.json" + ( cd "$staging" && tar czf - output agentic ) > "$handoff" + rm -rf "$staging" + sleep "${FAKE_CLIENT_SLEEP_SECONDS:-0}" + exit "${FAKE_CLIENT_EXIT_CODE:-0}" +fi + +if [[ "$args" == *hostname* ]]; then + echo "node-a" + exit 0 +fi + +exit 0 +""" + + +def preflight_record(hostname: str, revision: str) -> str: + """Build one preflight record line as a node would emit it.""" + return ( + f"INFERENCEX_KIMIK3_PREFLIGHT hostname={hostname} revision={revision} " + "gpu_count=8 gpu_arch=gfx942 squash_size_bytes=33076838400" + ) + + +def make_cluster( + tmp_path: Path, + *, + preflight_node_count: int = 2, + mismatched_revisions: bool = False, +) -> dict[str, object]: + """Stage a two-node cluster with fake Slurm, curl and srun on PATH.""" + workspace = tmp_path / "workspace" + workspace.mkdir() + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + cmd_log = tmp_path / "commands.log" + cmd_log.write_text("") + + write_executable(bin_dir / "salloc", FAKE_SALLOC.format(job_id=JOB_ID)) + write_executable(bin_dir / "scancel", FAKE_SCANCEL) + write_executable(bin_dir / "squeue", FAKE_SQUEUE) + write_executable(bin_dir / "curl", FAKE_CURL) + write_executable(bin_dir / "srun", FAKE_SRUN) + + hostnames = ["node-a", "node-b"][:preflight_node_count] + revisions = [REVISION, OTHER_REVISION if mismatched_revisions else REVISION] + preflight_output = tmp_path / "preflight.txt" + preflight_output.write_text( + "".join( + f"{preflight_record(host, revisions[index])}\n" + for index, host in enumerate(hostnames) + ) + ) + + env = { + **os.environ, + "PATH": f"{bin_dir}{os.pathsep}{os.environ['PATH']}", + "GITHUB_WORKSPACE": str(workspace), + "FAKE_CMD_LOG": str(cmd_log), + "FAKE_PREFLIGHT_OUTPUT": str(preflight_output), + "FAKE_SERVER_PID_FILE": str(tmp_path / "server.pid"), + "RUNNER_NAME": "mi300x-amds_00", + "IMAGE": IMAGE, + "MODEL": "moonshotai/Kimi-K3", + "MODEL_PREFIX": "kimik3", + "PRECISION": "fp4", + "FRAMEWORK": "vllm", + "EXP_NAME": "kimik3_p1x8_d0x8_conc4", + "RESULT_FILENAME": RESULT_FILENAME, + "SCENARIO_TYPE": "agentic-coding", + "SCENARIO_SUBDIR": "agentic/", + "IS_AGENTIC": "1", + "IS_MULTINODE": "true", + "DISAGG": "false", + "SPEC_DECODING": "none", + "KV_OFFLOADING": "none", + "CONC_LIST": "4", + "PORT": "8888", + "PREFILL_NUM_WORKERS": "1", + "PREFILL_TP": "8", + "PREFILL_PP_SIZE": "2", + "PREFILL_EP": "1", + "PREFILL_DP_ATTN": "false", + "DECODE_NUM_WORKERS": "0", + "DECODE_TP": "8", + "DECODE_PP_SIZE": "2", + "DECODE_EP": "1", + "DECODE_DP_ATTN": "false", + "KIMIK3_MODEL_CACHE_ROOT": str(tmp_path / "raid" / "models--moonshotai--Kimi-K3"), + "KIMIK3_SQUASH_DIR": str(tmp_path / "raid" / "squash"), + "KIMIK3_STARTUP_TIMEOUT_SECONDS": "60", + "KIMIK3_HEALTH_POLL_SECONDS": "1", + "KIMIK3_CLEANUP_TIMEOUT_SECONDS": "5", + "KIMIK3_CLEANUP_POLL_SECONDS": "1", + "FAKE_HEALTH": "ok", + } + env.pop("NATIVE_MULTINODE", None) + return { + "workspace": workspace, + "cmd_log": cmd_log, + "env": env, + "preflight_output": preflight_output, + "server_pid_file": tmp_path / "server.pid", + } + + +def run_launcher( + cluster: dict[str, object], script: Path = NATIVE_LAUNCHER, timeout: int = 120 +) -> subprocess.CompletedProcess: + """Run a launcher against the staged cluster.""" + return subprocess.run( + ["bash", str(script)], + cwd=str(REPO_ROOT), + env=cluster["env"], + capture_output=True, + text=True, + timeout=timeout, + ) + + +def grep_supports_perl_regex() -> bool: + """Report whether the local grep supports -P; BSD grep does not.""" + return ( + subprocess.run( + ["grep", "-oP", "x"], input="x", capture_output=True, text=True + ).returncode + == 0 + ) + + +@pytest.mark.skipif( + not grep_supports_perl_regex(), + reason="the pre-existing single-node launcher parses salloc with GNU grep -P", +) +def test_default_launcher_keeps_existing_single_node_path(tmp_path: Path) -> None: + """Without NATIVE_MULTINODE the MI300X launcher behaves exactly as before.""" + cluster = make_cluster(tmp_path) + cluster["env"].update( + { + "TP": "8", + "HF_HUB_CACHE": "/hf-hub-cache", + "SCENARIO_SUBDIR": "fixed_seq_len/", + "IS_MULTINODE": "false", + "IS_AGENTIC": "0", + } + ) + result = run_launcher(cluster, script=DEFAULT_LAUNCHER) + + assert result.returncode == 0, result.stderr + command_log = cluster["cmd_log"].read_text() + assert "--nodes=2" not in command_log + assert "mi300x_native_node_preflight.sh" not in command_log + + +def test_native_launcher_uses_two_full_nodes_and_all_node_preflight( + tmp_path: Path, +) -> None: + """Allocate both nodes exclusively and verify each one before serving.""" + cluster = make_cluster(tmp_path) + cluster["env"]["NATIVE_MULTINODE"] = "1" + result = run_launcher(cluster, script=DEFAULT_LAUNCHER) + + assert result.returncode == 0, result.stdout + result.stderr + lines = cluster["cmd_log"].read_text().splitlines() + + allocation = next(line for line in lines if line.startswith("salloc ")) + assert "--nodes=2" in allocation + assert "--gres=gpu:8" in allocation + + preflight = next( + line for line in lines if "mi300x_native_node_preflight.sh" in line + ) + assert "--ntasks=2" in preflight + + server = next(line for line in lines if "--kill-on-bad-exit=1" in line) + assert "kimik3_fp4_mi300x_vllm.sh" in server + + client = next(line for line in lines if "agentic_srt.sh" in line) + assert "--overlap" in client + assert "--nodelist=node-a" in client + assert "/raid/hf-hub-cache/inferencex/agentx-hub:/hf-hub-cache" in client + assert "/raid/hf-hub-cache:/hf-hub-cache" not in client + assert "--container-writable" in client + assert "--container-writable" in server + + assert f"scancel {JOB_ID}" in cluster["cmd_log"].read_text() + + +def test_native_launcher_rejects_topology_before_salloc(tmp_path: Path) -> None: + """Fail the contract check before holding any allocation.""" + cluster = make_cluster(tmp_path) + cluster["env"]["PREFILL_PP_SIZE"] = "1" + result = run_launcher(cluster) + + assert result.returncode != 0 + assert "TP8 x PP2" in result.stderr + assert "salloc " not in cluster["cmd_log"].read_text() + + +@pytest.mark.parametrize( + "cluster_kwargs", + [{"preflight_node_count": 1}, {"mismatched_revisions": True}], + ids=["one_record", "mismatched_revisions"], +) +def test_native_launcher_rejects_inconsistent_preflight_records( + tmp_path: Path, cluster_kwargs: dict +) -> None: + """Never start a rank until both nodes report the same usable state.""" + cluster = make_cluster(tmp_path, **cluster_kwargs) + result = run_launcher(cluster) + + assert result.returncode != 0 + assert "--kill-on-bad-exit=1" not in cluster["cmd_log"].read_text() + + +def test_server_failure_preserves_failure_and_cancels_allocation( + tmp_path: Path, +) -> None: + """A rank that dies early surfaces as a failure and frees the nodes.""" + cluster = make_cluster(tmp_path) + cluster["env"]["FAKE_SERVER_MODE"] = "exit:23" + cluster["env"]["FAKE_HEALTH"] = "down" + result = run_launcher(cluster) + + assert result.returncode != 0 + assert "exited" in (result.stdout + result.stderr) + assert "agentic_srt.sh" not in cluster["cmd_log"].read_text() + assert f"scancel {JOB_ID}" in cluster["cmd_log"].read_text() + + +def test_native_launcher_cancels_the_allocation_before_packaging_logs( + tmp_path: Path, +) -> None: + """Nothing that does not need the allocation may delay scancel.""" + cluster = make_cluster(tmp_path) + result = run_launcher(cluster) + + assert result.returncode == 0, result.stdout + result.stderr + output = result.stdout + result.stderr + assert output.index("[cleanup] cancelling allocation") < output.index( + "[cleanup] packaging server logs" + ) + + +@pytest.mark.parametrize( + "knob,value", + [ + ("KIMIK3_CLEANUP_POLL_SECONDS", "0"), + ("KIMIK3_CLEANUP_POLL_SECONDS", "0.5"), + ("KIMIK3_HEALTH_POLL_SECONDS", "abc"), + ], +) +def test_native_launcher_rejects_non_positive_timing_knobs( + tmp_path: Path, knob: str, value: str +) -> None: + """A zero or fractional poll would spin forever, once before scancel.""" + cluster = make_cluster(tmp_path) + cluster["env"][knob] = value + result = run_launcher(cluster) + + assert result.returncode != 0 + assert "positive integer" in result.stderr + assert "salloc " not in cluster["cmd_log"].read_text() + + +def test_native_launcher_rejects_mismatched_image_builds(tmp_path: Path) -> None: + """Equal squash size is the only signal both nodes hold the same build.""" + cluster = make_cluster(tmp_path) + cluster["preflight_output"].write_text( + preflight_record("node-a", REVISION) + + preflight_record("node-b", REVISION).replace( + "squash_size_bytes=33076838400", "squash_size_bytes=33076838401" + ) + ) + result = run_launcher(cluster) + + assert result.returncode != 0 + assert "different builds" in result.stdout + result.stderr + assert "--kill-on-bad-exit=1" not in cluster["cmd_log"].read_text() + + +def test_client_failure_propagates_the_client_exit_code(tmp_path: Path) -> None: + """Publish the artifacts, then exit with the replay's own status.""" + cluster = make_cluster(tmp_path) + cluster["env"]["FAKE_CLIENT_EXIT_CODE"] = "23" + result = run_launcher(cluster) + + assert result.returncode == 23, result.stdout + result.stderr + assert (cluster["workspace"] / f"{RESULT_FILENAME}_conc4.json").is_file() + assert f"scancel {JOB_ID}" in cluster["cmd_log"].read_text() + + +def test_sigterm_returns_143_and_reaps_server_and_allocation(tmp_path: Path) -> None: + """A signal must reap the server step and cancel the allocation.""" + cluster = make_cluster(tmp_path) + cluster["env"]["FAKE_HEALTH"] = "down" + cluster["env"]["KIMIK3_STARTUP_TIMEOUT_SECONDS"] = "300" + + process = subprocess.Popen( + ["bash", str(NATIVE_LAUNCHER)], + cwd=str(REPO_ROOT), + env=cluster["env"], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + try: + deadline = time.monotonic() + 30 + while time.monotonic() < deadline: + if "curl " in cluster["cmd_log"].read_text(): + break + time.sleep(0.1) + else: + raise AssertionError("launcher never reached server health polling") + process.terminate() + output = process.communicate(timeout=10)[0] + except BaseException: + process.kill() + raise + + assert process.returncode == 143, output + assert "stopping server step" in output + assert f"scancel {JOB_ID}" in cluster["cmd_log"].read_text() + + server_pid = int(cluster["server_pid_file"].read_text().strip()) + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + try: + os.kill(server_pid, 0) + except ProcessLookupError: + break + time.sleep(0.1) + else: + os.kill(server_pid, signal.SIGKILL) + raise AssertionError(f"server step {server_pid} outlived cleanup") + + +def test_success_extracts_only_host_owned_bounded_artifacts(tmp_path: Path) -> None: + """Only host-owned files cross into the workspace; the handoff is removed.""" + cluster = make_cluster(tmp_path) + workspace = cluster["workspace"] + result = run_launcher(cluster) + + assert result.returncode == 0, result.stdout + result.stderr + + aggregate = workspace / f"{RESULT_FILENAME}_conc4.json" + raw_artifact = ( + workspace / "LOGS" / "agentic" / "conc_4" / "aiperf_artifacts" + / "profile_export.json" + ) + server_logs = workspace / "multinode_server_logs.tar.gz" + for path in (aggregate, raw_artifact, server_logs): + assert path.is_file(), f"missing {path}" + assert path.stat().st_uid == os.getuid() + + handoffs = sorted(workspace.glob("*handoff*")) + assert handoffs == []