diff --git a/frontier/scheduler/replica_stage_scheduler/stage_execution_context.py b/frontier/scheduler/replica_stage_scheduler/stage_execution_context.py index e92a4edc..1dcda539 100644 --- a/frontier/scheduler/replica_stage_scheduler/stage_execution_context.py +++ b/frontier/scheduler/replica_stage_scheduler/stage_execution_context.py @@ -35,9 +35,14 @@ class StageExecutionContext: The context is intentionally independent of event timing and child lane queues. A complete operation first enters the ready FIFO, then the owner - admits it atomically. EP child schedulers may start only after their - wave's ticket has been acquired, and the ticket remains active through the - wave-level combine/cleanup boundary. + admits it atomically. An EP wave is admitted only from the FIFO head, so + it waits for every operation queued before it. A full-stage operation may + be admitted ahead of earlier queued full-stage operations, but never ahead + of an EP wave queued before it. Full-stage operations are therefore + admitted in the order their lane stage schedulers present them, and + ``admission_seq`` records enqueue order only. EP child schedulers may + start only after their wave's ticket has been acquired, and the ticket + remains active through the wave-level combine/cleanup boundary. """ def __init__( @@ -320,24 +325,39 @@ def _validate_ticket(self, ticket: StageAdmissionTicket) -> None: ) def try_acquire(self, ticket: StageAdmissionTicket) -> bool: - """Acquire the FIFO-head ticket if this stage is currently idle.""" + """Acquire ``ticket`` if the stage can admit it now. + + An EP wave must be the FIFO head of an idle stage. A full-stage ticket + must have no EP wave queued ahead of it. A ticket that is already + active is not queued and is refused. + """ self._validate_ticket(ticket) if ticket.scope == EP_WAVE: - if self._active_ep_ticket is not None or self._active_full_stage_tickets: + if ( + self._active_ep_ticket is not None + or self._active_full_stage_tickets + or not self._ready_fifo + or self._ready_fifo[0] != ticket + ): return False - elif self._active_ep_ticket is not None: - return False - elif len(self._active_full_stage_tickets) >= self._full_stage_capacity: - return False - elif self._forward_group_sealed: - return False - if not self._ready_fifo or self._ready_fifo[0] != ticket: - return False - self._ready_fifo.popleft() - if ticket.scope == EP_WAVE: + self._ready_fifo.popleft() self._active_ep_ticket = ticket else: + if ( + self._active_ep_ticket is not None + or len(self._active_full_stage_tickets) >= self._full_stage_capacity + or self._forward_group_sealed + ): + return False + for position, queued in enumerate(self._ready_fifo): + if queued == ticket: + break + if queued.scope == EP_WAVE: + return False + else: + return False + del self._ready_fifo[position] self._active_full_stage_tickets.add(ticket) self._refresh_active_ticket_view() return True diff --git a/tests/comparison/stage_admission_pp/compare_lanes.py b/tests/comparison/stage_admission_pp/compare_lanes.py new file mode 100644 index 00000000..1d16ecc3 --- /dev/null +++ b/tests/comparison/stage_admission_pp/compare_lanes.py @@ -0,0 +1,272 @@ +#!/usr/bin/env python3 +"""Compare DP-lane stage admission between vLLM traces and Frontier ledgers. + +Reads the vLLM ground truth written by ``run_vllm_worker.sh`` and the G7 cases +of ``tests.e2e.stage_admission_matrix`` for a set run before the change and one +run after it, computes the per-forward lane metrics of the stage-admission plan +(§4.7, M1–M5) on both sides with the same definitions, and writes the +workflow-gap table, summary and status for the calibration case. + +A vLLM forward row is one ``pp_boundary`` record: its lane is the DP rank the +driver pinned its requests to, and its stage is ``pp_rank``. Stage-0 +intervals end at ``send_start_ts``, taken after the post-forward synchronize; +last-stage intervals end at the record's wall-clock ``timestamp`` converted to +the monotonic clock with the offset the driver sampled around the round. A +Frontier forward row is one ``ATTN_DP_LANE`` ledger row. +""" + +from __future__ import annotations + +import argparse +import csv +import json +import statistics +from collections import defaultdict +from pathlib import Path + +from tests.e2e.stage_admission_matrix import ( + ADMISSION_DEADLOCK, + ATTN_DP_LANE, + SUCCESS, + interval_overlap, + matrix_root, + read_ledger, +) + +MODELS = ("moe", "dense") +BURSTS = (8, 16) +CO_START_BOUND = 0.5 +CO_EXECUTION_BOUND = 0.10 +# Dense DP ranks meet once per forward, in the DP metadata all-reduce that +# vLLM runs after ``forward_start_ts``. A rank's recorded interval therefore +# includes its wait for the other rank (start offsets) as well as its own +# duration variation (end offsets); the dummy predictor models neither, so +# dense co-execution is reported, not gated (plan §4.7 V5, D-9). MoE ranks +# stay aligned by EP collectives. +CO_EXECUTION_GATED = {"moe": True, "dense": False} +INFORMATIONAL = "INFORMATIONAL" +HOLDS, LOST = "HOLDS", "LOST" +FRONTIER_OWNER = "frontier/scheduler/replica_stage_scheduler/stage_execution_context.py" + + +def _forward(lane: int, stage: int, start: float, end: float, indices) -> dict: + return {"lane": lane, "stage": stage, "start": start, "end": end, + "indices": tuple(sorted(indices))} + + +def vllm_forwards(scenario_dir: Path) -> dict[tuple[int, int], dict]: + """Formal vLLM forwards keyed by ``(burst size, round)``.""" + requests = {row["request_id"]: row for row in map(json.loads, (scenario_dir / "requests.jsonl").read_text().splitlines())} + summary = json.loads((scenario_dir / "summary.json").read_text()) + offsets = { + entry["label"]: (entry["wall_minus_monotonic_before"] + entry["wall_minus_monotonic_after"]) / 2 + for entry in summary["rounds"] + } + runs: dict[tuple[int, int], dict] = {} + for line in (scenario_dir / "pp_boundary.jsonl").read_text().splitlines(): + record = json.loads(line) + members = [requests[request_id] for request_id in record["request_ids"]] + labels = {member["burst"] for member in members} + ranks = {member["rank"] for member in members} + if len(labels) != 1 or len(ranks) != 1: + raise ValueError(f"forward mixes bursts or ranks: {record['request_ids']}") + label = labels.pop() + if label == "warmup": + continue + if record["is_last_rank"]: + end = record["timestamp"] - offsets[label] + else: + end = record["send_start_ts"] + member = members[0] + run = runs.setdefault( + (int(label.split("-")[0][1:]), member["round"]), + {"forwards": [], "requests": [row for row in requests.values() if row["burst"] == label]}, + ) + run["forwards"].append(_forward(ranks.pop(), record["pp_rank"], record["forward_start_ts"], end, + (m["index"] for m in members))) + for run in runs.values(): + rows = run.pop("requests") + run["submitted"] = len(rows) + run["completed"] = sum(1 for row in rows if row["num_output_tokens"] == 1) + return runs + + +def vllm_placement(scenario_dir: Path) -> dict: + """Check from engine iterations that each formal request ran on its pinned rank.""" + pinned = {row["request_id"]: row["rank"] for row in map(json.loads, (scenario_dir / "requests.jsonl").read_text().splitlines())} + scheduled_by = defaultdict(set) + for path in sorted((scenario_dir / "dp_placement").glob("*.jsonl")): + for record in map(json.loads, path.read_text().splitlines()): + if record["kind"] == "engine_iteration": + for request_id in record["scheduled_new_req_ids"]: + scheduled_by[request_id].add(record["engine"]) + misplaced = sorted(rid for rid, rank in pinned.items() if scheduled_by.get(rid, {rank}) != {rank}) + unseen = sorted(rid for rid in pinned if rid not in scheduled_by) + return {"requests": len(pinned), "misplaced": misplaced, "unseen": unseen, + "ok": not misplaced and not unseen} + + +def frontier_run(set_dir: Path, case_id: str) -> dict: + case_dir = set_dir / case_id + run = json.loads((case_dir / "run.json").read_text()) + case = json.loads((case_dir / "case.json").read_text()) + result = {"outcome": run["outcome"], "submitted": case["num_requests"], "forwards": []} + if run["outcome"] != SUCCESS: + result["completed"] = None + return result + result["completed"] = case["num_requests"] + for row in read_ledger(case_dir / "metrics"): + if row["execution_scope"] != ATTN_DP_LANE: + continue + result["forwards"].append(_forward(row["replica_local_id"], row["stage_id"], row["stage_start_ts"], + row["stage_end_ts"], (int(rid) for rid in row["request_ids"]))) + lanes_match_index = all(index % 2 == forward["lane"] for forward in result["forwards"] + for index in forward["indices"]) + result["placement_ok"] = lanes_match_index + return result + + +def lane_metrics(forwards: list[dict]) -> dict: + """M2–M5 of plan §4.7 from one run's forwards.""" + by_lane_stage = defaultdict(list) + for forward in forwards: + by_lane_stage[(forward["lane"], forward["stage"])].append(forward) + sequences = { + f"lane{lane}/stage{stage}": [list(f["indices"]) for f in sorted(rows, key=lambda f: f["start"])] + for (lane, stage), rows in sorted(by_lane_stage.items()) + } + metrics = {"M2_sequences": sequences} + for stage in sorted({forward["stage"] for forward in forwards}): + lane0 = sorted(by_lane_stage[(0, stage)], key=lambda f: f["start"]) + lane1 = sorted(by_lane_stage[(1, stage)], key=lambda f: f["start"]) + pairing = [] + for forward in lane0: + overlaps = [(min(forward["end"], other["end"]) - max(forward["start"], other["start"]), other) + for other in lane1] + overlap, partner = max(overlaps, key=lambda item: item[0], default=(0.0, None)) + pairing.append([list(forward["indices"]), list(partner["indices"]) if partner and overlap > 0 else None]) + durations = [f["end"] - f["start"] for f in lane0 + lane1] + skew = abs(lane0[0]["start"] - lane1[0]["start"]) / statistics.median(durations) if lane0 and lane1 else None + overlap = interval_overlap([(f["start"], f["end"], f["lane"]) for f in lane0 + lane1]) + metrics[f"stage{stage}"] = { + "M3_pairing": pairing, + "M4_co_start": skew, + "M5_co_execution": overlap["multi_lane_busy_time"] / overlap["busy_time"] if overlap["busy_time"] else None, + "median_forward_duration": statistics.median(durations) if durations else None, + "self_overlap": overlap["self_overlap"], + } + return metrics + + +def _row(check, model, burst, round_index, metric, groundtruth, after, base, status, note=""): + return {"check": check, "model": model, "burst": burst, "round": round_index, "metric": metric, + "groundtruth": json.dumps(groundtruth), "frontier_after": json.dumps(after), + "frontier_base": json.dumps(base), "status": status, + "frontier_owner": FRONTIER_OWNER if status == "MISMATCH" else "", "note": note} + + +def compare(vllm_run: Path, frontier_root: Path, before: str, after: str) -> tuple[list[dict], dict]: + rows, details = [], {"placement": {}, "runs": {}} + for model in MODELS: + scenario_dir = vllm_run / "runs" / model + vllm_runs = vllm_forwards(scenario_dir) + details["placement"][model] = vllm_placement(scenario_dir) + for burst in BURSTS: + case_id = f"G7-{model}-dp2-pp2-n{burst}" + base = frontier_run(frontier_root / before, case_id) + new = frontier_run(frontier_root / after, case_id) + base_metrics = lane_metrics(base["forwards"]) if base["outcome"] == SUCCESS else None + new_metrics = lane_metrics(new["forwards"]) if new["outcome"] == SUCCESS else None + rounds = sorted(r for (b, r) in vllm_runs if b == burst) + vllm_metrics = {r: lane_metrics(vllm_runs[(burst, r)]["forwards"]) for r in rounds} + details["runs"][case_id] = {"frontier_base": base_metrics, "frontier_after": new_metrics, + "frontier_base_outcome": base["outcome"], + "frontier_after_outcome": new["outcome"], + "frontier_after_placement_ok": new.get("placement_ok"), + "vllm": vllm_metrics} + base_m4 = base_metrics["stage0"]["M4_co_start"] if base_metrics else None + # Negative controls: the base rule deadlocks MoE and starts dense + # lanes one forward apart. They describe the base, not vLLM. + if model == "moe": + rows.append(_row("N1", model, burst, "base", "base outcome", None, new["outcome"], + base["outcome"], HOLDS if base["outcome"] == ADMISSION_DEADLOCK else LOST, + note=f"expected {ADMISSION_DEADLOCK}")) + else: + rows.append(_row("N4", model, burst, "base", "M4 stage-0 co-start", None, + new_metrics["stage0"]["M4_co_start"] if new_metrics else None, base_m4, + HOLDS if base_m4 is not None and base_m4 >= CO_START_BOUND else LOST, + note=f"expected >= {CO_START_BOUND}")) + for r in rounds: + run = vllm_runs[(burst, r)] + completed = run["completed"] == run["submitted"] == burst + status = "MATCH" if completed and new["outcome"] == SUCCESS else "MISMATCH" + rows.append(_row("V1", model, burst, r, "M1 completion", + f"{run['completed']}/{run['submitted']}", new["outcome"], base["outcome"], status)) + gt = vllm_metrics[r] + after_m2 = new_metrics["M2_sequences"] if new_metrics else None + rows.append(_row("V2", model, burst, r, "M2 lane sequences", gt["M2_sequences"], after_m2, + base_metrics["M2_sequences"] if base_metrics else None, + "MATCH" if after_m2 == gt["M2_sequences"] else "MISMATCH")) + after_m3 = new_metrics["stage0"]["M3_pairing"] if new_metrics else None + rows.append(_row("V3", model, burst, r, "M3 stage-0 pairing", gt["stage0"]["M3_pairing"], after_m3, + base_metrics["stage0"]["M3_pairing"] if base_metrics else None, + "MATCH" if after_m3 == gt["stage0"]["M3_pairing"] else "MISMATCH")) + after_m4 = new_metrics["stage0"]["M4_co_start"] if new_metrics else None + m4_ok = (gt["stage0"]["M4_co_start"] < CO_START_BOUND and after_m4 is not None + and after_m4 < CO_START_BOUND) + rows.append(_row("V4", model, burst, r, "M4 stage-0 co-start", gt["stage0"]["M4_co_start"], + after_m4, base_m4, "MATCH" if m4_ok else "MISMATCH")) + gt_m5 = statistics.mean(vllm_metrics[r]["stage0"]["M5_co_execution"] for r in rounds) + after_m5 = new_metrics["stage0"]["M5_co_execution"] if new_metrics else None + base_m5 = base_metrics["stage0"]["M5_co_execution"] if base_metrics else None + if CO_EXECUTION_GATED[model]: + m5_ok = after_m5 is not None and abs(after_m5 - gt_m5) <= CO_EXECUTION_BOUND + m5_status = "MATCH" if m5_ok else "MISMATCH" + else: + m5_status = INFORMATIONAL + rows.append(_row("V5", model, burst, "mean", "M5 stage-0 co-execution", gt_m5, after_m5, base_m5, + m5_status)) + return rows, details + + +def main(argv=None) -> int: + parser = argparse.ArgumentParser(description=__doc__.split("\n", 1)[0]) + parser.add_argument("--vllm-run", type=Path, required=True, help="evidence directory of one worker run") + parser.add_argument("--before", default="base") + parser.add_argument("--after", default="after") + parser.add_argument("--output", type=Path, required=True, help="case analysis directory") + args = parser.parse_args(argv) + + rows, details = compare(args.vllm_run, matrix_root(), args.before, args.after) + args.output.mkdir(parents=True, exist_ok=True) + with (args.output / "workflow_gap_table.csv").open("w", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=list(rows[0])) + writer.writeheader() + writer.writerows(rows) + (args.output / "lane_metrics.json").write_text(json.dumps(details, indent=1, sort_keys=True)) + mismatches = [row for row in rows if row["status"] == "MISMATCH"] + placement_ok = all(p["ok"] for p in details["placement"].values()) + placement_unseen = sum(len(p["unseen"]) for p in details["placement"].values()) + controls = [row for row in rows if row["status"] in (HOLDS, LOST)] + status = { + "analysis_state": "COMPLETE", + "status": "PASS" if not mismatches and placement_ok else "FAIL", + "correction_state": "not_applicable", + "rows": len(rows), + "mismatches": len(mismatches), + "vllm_placement_ok": placement_ok, + "vllm_placement_unseen_requests": placement_unseen, + "negative_control_holds": all(row["status"] == HOLDS for row in controls), + "next_action": ("record C7 in the test report" if not mismatches and placement_ok + else "report each MISMATCH row with its cause before P4; adjust nothing"), + } + (args.output / "workflow_gap_status.json").write_text(json.dumps(status, indent=1)) + for row in rows: + print(f"{row['check']} {row['model']:<5} n{row['burst']:<3} r{row['round']!s:<5} {row['status']:<9} " + f"gt={row['groundtruth'][:40]} after={row['frontier_after'][:40]} base={row['frontier_base'][:40]}") + print(json.dumps(status)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/comparison/stage_admission_pp/run_vllm_worker.sh b/tests/comparison/stage_admission_pp/run_vllm_worker.sh new file mode 100644 index 00000000..e7dea8ff --- /dev/null +++ b/tests/comparison/stage_admission_pp/run_vllm_worker.sh @@ -0,0 +1,104 @@ +#!/usr/bin/env bash +# Worker entry point for the stage-admission vLLM ground truth (4 GPUs, +# vllm/vllm-openai:v0.10.2). Builds the instrumented overlay, runs the MoE and +# dense bursts, and publishes the traces to the cloud-volume archive and to the +# calibration case directory on the mounted workspace. +# +# Required environment: +# RUN_TAG identifier of this run +# FRONTIER_TREE Frontier worktree on the mounted workspace +# GROUNDTRUTH vLLM-BS checkout on the mounted workspace +# CASE_DIR calibration case directory on the mounted workspace; reads +# inputs/, writes runs/vllm-instrumented// +# ARCHIVE_DIR cloud-volume directory for this run +# Optional: +# OVERLAY_PATCH recorded unified diff applied to the accepted overlay +set -euo pipefail +set +x +: "${RUN_TAG:?}" "${FRONTIER_TREE:?}" "${GROUNDTRUTH:?}" "${CASE_DIR:?}" "${ARCHIVE_DIR:?}" +EVIDENCE_DIR="$CASE_DIR/runs/vllm-instrumented/$RUN_TAG" + +for d in /usr/local/nvidia/lib64 /usr/local/nvidia/lib /usr/lib/x86_64-linux-gnu; do + if [ -e "$d/libcuda.so.1" ]; then + export LD_LIBRARY_PATH="$d${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" + break + fi +done + +PY=python3 +SCRIPT_DIR="$FRONTIER_TREE/tests/comparison/stage_admission_pp" +WORK=/tmp/stage_admission_pp/$RUN_TAG +mkdir -p "$WORK/runs" +export VLLM_CACHE_ROOT="$WORK/vllm_cache" HF_HOME="$WORK/hf" HF_HUB_OFFLINE=1 \ + TRANSFORMERS_OFFLINE=1 VLLM_NO_USAGE_STATS=1 DO_NOT_TRACK=1 + +publish() { + local target="$1" + mkdir -p "$target" + for item in runs overlay_report.json worker_env.json vllm_import.txt; do + if [ -e "$WORK/$item" ]; then cp -r "$WORK/$item" "$target/"; fi + done + echo "status=$status" > "$target/COMPLETE" +} +# The worker writes the mounted workspace as root; hand the evidence back to +# the owner of the case directory. +publish_evidence() { + publish "$EVIDENCE_DIR" + chown -R "$(stat -c %u:%g "$CASE_DIR")" "$EVIDENCE_DIR" +} +status=0 +"$PY" - <<'PY' | tee "$WORK/worker_env.json" +import json, platform, torch +print(json.dumps({ + "python": platform.python_version(), + "torch": torch.__version__, + "cuda_available": torch.cuda.is_available(), + "device_count": torch.cuda.device_count(), + "devices": [torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())], +})) +PY + +SITE_VLLM=$("$PY" -c 'import importlib.util, os; print(os.path.dirname(importlib.util.find_spec("vllm").origin))') +"$PY" "$SCRIPT_DIR/vllm_burst_driver.py" overlay \ + --site-vllm "$SITE_VLLM" --checkout "$GROUNDTRUTH" --destination "$WORK/overlay" \ + --expected-changes "$CASE_DIR/inputs/fork_changed_files.txt" \ + --report "$WORK/overlay_report.json" ${OVERLAY_PATCH:+--patch "$OVERLAY_PATCH"} || status=3 +if [ "$status" -ne 0 ]; then + publish "$ARCHIVE_DIR"; publish_evidence + echo "WORKER_STATUS=$status overlay rejected" + exit "$status" +fi +export PYTHONPATH="$WORK/overlay" +"$PY" -c 'import vllm, vllm.v1.frontier_trace as t; print("VLLM_IMPORT", vllm.__version__, vllm.__file__, t.__file__)' \ + | tee "$WORK/vllm_import.txt" + +run_scenario() { + local name="$1"; shift + local out="$WORK/runs/$name" + mkdir -p "$out" + # Engine cores write their placement records at interpreter exit, which a + # forked multiprocessing child skips; spawned children run it. + if VLLM_FRONTIER_INSTRUMENTATION=1 VLLM_WORKER_MULTIPROC_METHOD=spawn \ + VLLM_FRONTIER_PP_BOUNDARY_LOG_PATH="$out/pp_boundary.jsonl" \ + VLLM_FRONTIER_DP_PLACEMENT_LOG_DIR="$out/dp_placement" \ + timeout 1500 "$PY" "$SCRIPT_DIR/vllm_burst_driver.py" run --output-dir "$out" "$@" \ + > "$out/driver.log" 2>&1; then + echo "SCENARIO_PASS $name" + else + echo "SCENARIO_FAIL $name exit=$?" + tail -n 60 "$out/driver.log" + status=1 + fi +} +run_scenario moe --model-config "$FRONTIER_TREE/data/config/models/Qwen3-30B-A3B-tiny.json" --enable-expert-parallel +run_scenario dense --model-config "$FRONTIER_TREE/data/config/models/Llama-3.2-1B-Instruct.json" + +publish "$ARCHIVE_DIR" +publish_evidence + +for name in moe dense; do + grep -h "DRIVER_DONE" "$WORK/runs/$name/driver.log" | sed "s/^/$name /" || true + wc -l "$WORK/runs/$name/pp_boundary.jsonl" 2>/dev/null || true +done +echo "WORKER_STATUS=$status RUN_TAG=$RUN_TAG" +exit "$status" diff --git a/tests/comparison/stage_admission_pp/vllm_burst_driver.py b/tests/comparison/stage_admission_pp/vllm_burst_driver.py new file mode 100644 index 00000000..21a04109 --- /dev/null +++ b/tests/comparison/stage_admission_pp/vllm_burst_driver.py @@ -0,0 +1,281 @@ +#!/usr/bin/env python3 +"""vLLM ground truth for stage admission of DP ranks under pipeline parallelism. + +Runs inside the ``vllm/vllm-openai:v0.10.2`` image on one 4-GPU worker. Two +subcommands: + +``overlay`` + Build the instrumented vLLM package: copy the image's installed ``vllm`` + package (which carries the compiled extensions) and copy every + ``vllm/**/*.py`` of the ground-truth checkout over it. The overlay is + accepted only when the files where the image and the checkout differ are + exactly the checkout's own changes over its upstream base, listed in + ``--expected-changes``. An accepted overlay may then take one recorded + ``--patch`` (a unified diff), whose SHA-256 and files enter the report. + +``run`` + Start one ``AsyncLLM`` with DP=2, PP=2, TP=1 (EP for the MoE model), run + warmup requests, then bursts of prefill-only requests pinned to rank + ``i mod 2``. Every request of a burst is added before any output is + awaited. The engines are idle between rounds. Per-forward traces come + from the checkout's own instrumentation; this script records each + request's rank, submit and finish times, and the wall/monotonic clock + offset around each round. + +The script imports vLLM only inside ``run`` so that ``overlay`` never loads +the package it is building. +""" + +from __future__ import annotations + +import argparse +import asyncio +import filecmp +import hashlib +import json +import os +import random +import re +import shutil +import sys +import time +from pathlib import Path + + +def build_overlay(site_vllm: Path, checkout: Path, destination: Path, expected_changes: Path) -> dict: + target = destination / "vllm" + if target.exists(): + shutil.rmtree(target) + shutil.copytree(site_vllm, target, symlinks=True) + differing = [] + for source in sorted((checkout / "vllm").rglob("*.py")): + relative = source.relative_to(checkout) + installed = site_vllm.parent / relative + if not installed.exists() or not filecmp.cmp(source, installed, shallow=False): + differing.append(str(relative)) + copy_target = destination / relative + copy_target.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(source, copy_target) + expected = sorted( + line.strip() for line in expected_changes.read_text().splitlines() + if line.strip().endswith(".py") + ) + return { + "site_vllm": str(site_vllm), + "checkout": str(checkout), + "overlay": str(target), + "differing_py_files": differing, + "expected_py_changes": expected, + "unexpected": sorted(set(differing) - set(expected)), + "missing": sorted(set(expected) - set(differing)), + "accepted": set(differing) == set(expected), + } + + +def apply_patch(patch: Path, root: Path) -> list[str]: + """Apply the unified diff ``patch`` to files under ``root``. + + The worker image need not carry ``patch`` or ``git``, so hunks are applied + here as text replacements; each hunk must match its file exactly once. + An empty hunk line is a context line whose leading space was trimmed. + """ + lines = patch.read_text().splitlines(keepends=True) + hunks: dict[str, list[tuple[str, str]]] = {} + index = 0 + while index < len(lines): + line = lines[index] + index += 1 + if line.startswith("+++ "): + target = line[4:].strip().removeprefix("b/") + hunks.setdefault(target, []) + elif line.startswith("@@ "): + header = re.match(r"@@ -\d+(?:,(\d+))? \+\d+(?:,(\d+))? @@", line) + old_count, new_count = (int(count or 1) for count in header.groups()) + old, new = [], [] + while old_count or new_count: + tag, text = lines[index][0], lines[index][1:] + index += 1 + if tag == "\n": + tag, text = " ", "\n" + if tag not in " -+": + raise ValueError(f"{patch}: unexpected hunk line {lines[index - 1]!r}") + if tag in " -": + old.append(text) + old_count -= 1 + if tag in " +": + new.append(text) + new_count -= 1 + hunks[target].append(("".join(old), "".join(new))) + for relative, edits in hunks.items(): + path = root / relative + text = path.read_text() + for old, new in edits: + if text.count(old) != 1: + raise ValueError(f"{patch}: a hunk does not match {relative} exactly once") + text = text.replace(old, new) + path.write_text(text) + return sorted(hunks) + + +def write_model_dir(model_config: Path, model_dir: Path) -> dict: + config = json.loads(model_config.read_text()) + model_dir.mkdir(parents=True, exist_ok=True) + (model_dir / "config.json").write_text(json.dumps(config, indent=1)) + return config + + +def prompt_token_ids(request_id: str, length: int, vocab_size: int) -> list[int]: + generator = random.Random(request_id) + return [generator.randrange(100, vocab_size - 100) for _ in range(length)] + + +async def run_bursts(args: argparse.Namespace) -> dict: + from vllm import SamplingParams + from vllm.engine.arg_utils import AsyncEngineArgs + from vllm.inputs import TokensPrompt + from vllm.sampling_params import RequestOutputKind + from vllm.v1.engine.async_llm import AsyncLLM + + output_dir = Path(args.output_dir) + model_config = write_model_dir(Path(args.model_config), output_dir / "model") + engine_args = AsyncEngineArgs( + model=str(output_dir / "model"), + load_format="dummy", + skip_tokenizer_init=True, + dtype="bfloat16", + tensor_parallel_size=1, + pipeline_parallel_size=args.pipeline_parallel_size, + data_parallel_size=args.data_parallel_size, + enable_expert_parallel=args.enable_expert_parallel, + enforce_eager=True, + enable_prefix_caching=False, + enable_chunked_prefill=True, + max_num_batched_tokens=args.prompt_tokens, + max_num_seqs=args.max_num_seqs, + block_size=16, + max_model_len=args.max_model_len, + gpu_memory_utilization=args.gpu_memory_utilization, + seed=0, + disable_log_stats=True, + ) + engine = AsyncLLM.from_engine_args(engine_args) + sampling = SamplingParams( + max_tokens=1, ignore_eos=True, temperature=0.0, detokenize=False, + output_kind=RequestOutputKind.FINAL_ONLY, + ) + vocab_size = int(model_config["vocab_size"]) + records: list[dict] = [] + + async def wait_until_idle() -> float: + started = time.monotonic() + while engine.engine_core.dp_engines_running(): + if time.monotonic() - started > args.idle_timeout_s: + raise RuntimeError("DP engines did not pause between rounds") + await asyncio.sleep(0.05) + await asyncio.sleep(args.idle_gap_s) + return time.monotonic() - started + + async def burst(label: str, round_index: int, num_requests: int) -> dict: + offset_before = time.time() - time.monotonic() + queues = [] + for index in range(num_requests): + request_id = f"{label}-q{index}" + rank = index % args.data_parallel_size + prompt = TokensPrompt( + prompt_token_ids=prompt_token_ids(request_id, args.prompt_tokens, vocab_size) + ) + submitted = time.monotonic() + queue = await engine.add_request(request_id, prompt, sampling, data_parallel_rank=rank) + queues.append((request_id, index, rank, submitted, queue)) + for request_id, index, rank, submitted, queue in queues: + output = await queue.get() + while not output.finished: + output = await queue.get() + records.append({ + "request_id": request_id, "burst": label, "round": round_index, + "index": index, "rank": rank, "submit_monotonic": submitted, + "finish_monotonic": time.monotonic(), + "num_prompt_tokens": len(output.prompt_token_ids), + "num_output_tokens": len(output.outputs[0].token_ids), + "finish_reason": output.outputs[0].finish_reason, + }) + offset_after = time.time() - time.monotonic() + return {"label": label, "round": round_index, "num_requests": num_requests, + "wall_minus_monotonic_before": offset_before, + "wall_minus_monotonic_after": offset_after, + "idle_wait_s": await wait_until_idle()} + + rounds = [] + try: + rounds.append(await burst("warmup", 0, args.warmups)) + for num_requests in args.bursts: + for round_index in range(args.rounds): + rounds.append(await burst(f"b{num_requests}-r{round_index}", round_index, num_requests)) + cache_config = engine.vllm_config.cache_config + summary = { + "model_config": args.model_config, + "num_gpu_blocks": cache_config.num_gpu_blocks, + "block_size": cache_config.block_size, + "engine_args": {key: value for key, value in vars(engine_args).items() + if isinstance(value, (bool, int, float, str, type(None)))}, + "rounds": rounds, + } + finally: + engine.shutdown() + (output_dir / "requests.jsonl").write_text("".join(json.dumps(row) + "\n" for row in records)) + return summary + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__.split("\n", 1)[0]) + commands = parser.add_subparsers(dest="command", required=True) + overlay = commands.add_parser("overlay") + overlay.add_argument("--site-vllm", type=Path, required=True) + overlay.add_argument("--checkout", type=Path, required=True) + overlay.add_argument("--destination", type=Path, required=True) + overlay.add_argument("--expected-changes", type=Path, required=True) + overlay.add_argument("--report", type=Path, required=True) + overlay.add_argument("--patch", type=Path) + run = commands.add_parser("run") + run.add_argument("--model-config", required=True) + run.add_argument("--output-dir", required=True) + run.add_argument("--enable-expert-parallel", action="store_true") + run.add_argument("--data-parallel-size", type=int, default=2) + run.add_argument("--pipeline-parallel-size", type=int, default=2) + run.add_argument("--prompt-tokens", type=int, default=256) + run.add_argument("--max-num-seqs", type=int, default=4) + run.add_argument("--max-model-len", type=int, default=512) + run.add_argument("--gpu-memory-utilization", type=float, default=0.5) + run.add_argument("--bursts", type=int, nargs="+", default=[8, 16]) + run.add_argument("--rounds", type=int, default=3) + run.add_argument("--warmups", type=int, default=4) + run.add_argument("--idle-gap-s", type=float, default=1.0) + run.add_argument("--idle-timeout-s", type=float, default=60.0) + args = parser.parse_args(argv) + + if args.command == "overlay": + report = build_overlay(args.site_vllm, args.checkout, args.destination, args.expected_changes) + if report["accepted"] and args.patch is not None: + files = apply_patch(args.patch, args.destination) + report["patch"] = { + "path": str(args.patch), + "sha256": hashlib.sha256(args.patch.read_bytes()).hexdigest(), + "files": files, + "equal_to_image_after_patch": { + name: filecmp.cmp(args.destination / name, args.site_vllm.parent / name, shallow=False) + for name in files + }, + } + args.report.write_text(json.dumps(report, indent=1)) + print(json.dumps({key: report[key] for key in ("accepted", "unexpected", "missing")})) + return 0 if report["accepted"] else 3 + Path(args.output_dir).mkdir(parents=True, exist_ok=True) + summary = asyncio.run(run_bursts(args)) + Path(args.output_dir, "summary.json").write_text(json.dumps(summary, indent=1)) + print("DRIVER_DONE", json.dumps({"rounds": len(summary["rounds"]), + "num_gpu_blocks": summary["num_gpu_blocks"]})) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/e2e/stage_admission_matrix.py b/tests/e2e/stage_admission_matrix.py new file mode 100644 index 00000000..c872b011 --- /dev/null +++ b/tests/e2e/stage_admission_matrix.py @@ -0,0 +1,833 @@ +#!/usr/bin/env python3 +"""Case matrix for stage admission of attention-DP lanes under pipeline parallelism. + +Runs the case list of ``build_cases`` on the current source tree and writes, +for each case, its inputs (``case.json``), its run provenance (``run.json``) +and one outcome artifact: + +* ``success``: ``sha256sums.txt`` over the copied metrics tree; +* ``admission_deadlock``: ``state_report.json`` read from the live scheduler + objects after the sequential run ends with work left; +* ``configuration_rejection`` / ``other_failure``: ``error.txt``. + +Each case runs in its own child process because ``IS_MOE`` is process-global, +and each child runs in its own session so a case that exceeds +``--case-timeout`` is killed with everything it started. Every child writes +its simulator output under ``/work/``, a path shared by all +sets, so that files embedding the output path compare byte for byte between a +set run before a change and one run after it. Sets therefore run one at a +time: ``run`` holds an exclusive lock on the matrix root. + +Usage:: + + python -m tests.e2e.stage_admission_matrix run --set base [--group G3a] + python -m tests.e2e.stage_admission_matrix compare --before base --after after +""" + +from __future__ import annotations + +import argparse +import csv +import hashlib +import json +import fcntl +import os +import shutil +import signal +import subprocess +import sys +import time +import traceback +from collections import defaultdict +from concurrent.futures import ThreadPoolExecutor +from dataclasses import asdict, dataclass +from importlib import metadata as importlib_metadata +from pathlib import Path +from typing import Iterable, Sequence + +from tests.scratch_root import resolve_scratch_root + + +REPO_ROOT = Path(__file__).resolve().parents[2] +MATRIX_DIR_NAME = "stage_admission_ordering" +DRAIN_MESSAGE = "Sequential simulation ended with non-empty scheduler state" +ATTN_DP_LANE = "ATTN_DP_LANE" + +SUCCESS = "success" +ADMISSION_DEADLOCK = "admission_deadlock" +CONFIGURATION_REJECTION = "configuration_rejection" +OTHER_FAILURE = "other_failure" + +# Plan §4.1 synthetic fixture and §4.7 vLLM-aligned fixture. +SYNTHETIC = "synthetic" +VLLM_ALIGNED = "vllm_aligned" +VLLM_ALIGNED_MODELS = {True: "Qwen3-30B-A3B-tiny", False: "Llama-3.2-1B-Instruct"} + +PREFILL_ONLY = (16, 1) +PREFILL_DECODE = (16, 3) +VLLM_ALIGNED_PREFILL_ONLY = (256, 1) +ONLINE_QPS = 20.0 +ONLINE_QPS_SWEEP = (5.0, 80.0) +DEFAULT_CASE_TIMEOUT_S = 600 + + +@dataclass(frozen=True) +class Case: + case_id: str + group: str + fixture: str = SYNTHETIC + is_moe: bool = False + attn_dp: int = 1 + stages: int = 1 + num_requests: int = 0 + prefill_tokens: int = 0 + decode_tokens: int = 0 + arrival: str = "static" + qps: float = 1e6 + sys_arch: str = "co-location" + simulation_mode: str = "offline" + cc_backend: str = "analytical" + recipe: str | None = None + recipe_env: tuple[tuple[str, str], ...] = () + contention_witness: bool = False + + @property + def path(self) -> str: + """Plan §4.2 path before P0 classification refines L/T.""" + if self.recipe is not None or self.stages == 1 or self.attn_dp == 1: + return "U" + return "LT" + + +def _shape_id(group: str, is_moe: bool, attn_dp: int, stages: int, num_requests: int) -> str: + kind = "moe" if is_moe else "dense" + return f"{group}-{kind}-dp{attn_dp}-pp{stages}-n{num_requests}" + + +def _synthetic(group: str, is_moe: bool, attn_dp: int, stages: int, num_requests: int, + lengths: tuple[int, int], suffix: str = "", **fields) -> Case: + return Case( + case_id=_shape_id(group, is_moe, attn_dp, stages, num_requests) + suffix, + group=group, + is_moe=is_moe, + attn_dp=attn_dp, + stages=stages, + num_requests=num_requests, + prefill_tokens=lengths[0], + decode_tokens=lengths[1], + **fields, + ) + + +def _release_recipes() -> list[Case]: + architecture_root = REPO_ROOT / "examples" / "architecture" + recipes = [] + for architecture in ("co-location", "pdd", "pd-af-disagg"): + for mode in ("offline", "online"): + for script in sorted((architecture_root / architecture / mode).glob("*.sh")): + recipes.append( + Case( + case_id=f"G1-{architecture}-{mode}-{script.stem}", + group="G1", + recipe=str(script.relative_to(REPO_ROOT)), + ) + ) + return recipes + + +def build_cases() -> list[Case]: + """Return the plan §4.2 case list in a fixed order.""" + cases: list[Case] = [] + # R0 reproduces the author-reported shapes of design.md with their inputs: + # Poisson arrivals, the default CC backend and prefill 16 / decode 3. + r0_shapes = [ + (True, 2, 2, 3), (True, 2, 2, 4), (True, 2, 2, 6), (True, 4, 2, 8), + (True, 2, 1, 6), (True, 2, 1, 12), (True, 4, 1, 8), (True, 4, 1, 12), + (True, 1, 2, 6), (True, 1, 3, 6), + (False, 2, 2, 6), (False, 4, 2, 8), (False, 2, 1, 6), (False, 4, 1, 8), + (False, 1, 2, 6), (True, 2, 3, 6), + ] + for is_moe, attn_dp, stages, num_requests in r0_shapes: + cases.append( + _synthetic("R0", is_moe, attn_dp, stages, num_requests, PREFILL_DECODE, + arrival="poisson", cc_backend="default") + ) + cases.extend(_release_recipes()) + for attn_dp in (2, 4): + for stages in (1, 2, 3): + for num_requests in (4, 8, 12): + cases.append(_synthetic("G3a", True, attn_dp, stages, num_requests, PREFILL_ONLY)) + for attn_dp in (2, 4): + for stages in (1, 2, 3): + for num_requests in (4, 8): + cases.append(_synthetic("G3b", True, attn_dp, stages, num_requests, PREFILL_DECODE)) + for attn_dp in (2, 4): + for stages in (1, 2, 3): + for num_requests in (4, 8): + cases.append( + _synthetic("G4", False, attn_dp, stages, num_requests, PREFILL_DECODE, + contention_witness=stages > 1 and num_requests == 8) + ) + for is_moe in (True, False): + for stages in (1, 2, 3): + cases.append(_synthetic("G5", is_moe, 1, stages, 6, PREFILL_DECODE)) + for is_moe in (True, False): + for num_requests in (8, 16): + cases.append( + _synthetic("G7", is_moe, 2, 2, num_requests, VLLM_ALIGNED_PREFILL_ONLY, + fixture=VLLM_ALIGNED) + ) + # Plan §7 (R2-03). PDD admits multi-lane contexts only for MoE: dense PDD + # requires attn_dp == 1. Online cells use Poisson arrivals at ONLINE_QPS. + # On main, MONOLITHIC and PREFILL place every request of one scheduling + # call from lane 0 on, so incremental online arrivals all land on lane 0; + # the "-burst" cells deliver all requests at t=0 to reach several lanes. + pdd_shapes = [(True, attn_dp, stages) for attn_dp in (2, 4) for stages in (1, 2, 3)] + pdd_shapes.append((False, 1, 2)) + online = dict(simulation_mode="online", arrival="poisson", qps=ONLINE_QPS) + burst = dict(simulation_mode="online", suffix="-burst") + for group, timing in (("G8", {}), ("G9", online)): + for is_moe, attn_dp, stages in pdd_shapes: + cases.append( + _synthetic(group, is_moe, attn_dp, stages, 8, PREFILL_DECODE, + sys_arch="pd-disaggregation", **timing) + ) + for attn_dp in (2, 4): + for stages in (2, 3): + cases.append( + _synthetic("G9", True, attn_dp, stages, 8, PREFILL_DECODE, + sys_arch="pd-disaggregation", **burst) + ) + for is_moe, lengths in ((True, PREFILL_ONLY), (False, PREFILL_DECODE)): + for attn_dp in (2, 4): + for stages in (1, 2, 3): + cases.append(_synthetic("G10", is_moe, attn_dp, stages, 8, lengths, **online)) + cases.append(_synthetic("G10", is_moe, attn_dp, stages, 8, lengths, **burst)) + for qps in ONLINE_QPS_SWEEP: + cases.append( + _synthetic("G10", is_moe, 2, 2, 8, lengths, suffix=f"-q{qps:g}", + **dict(online, qps=qps)) + ) + # PD-AF contexts have capacity 1 (DECODE_ATTN requires attn_dp == 1), so + # these are unchanged controls for PREFILL pipeline stages. + for mode in ("offline", "online"): + suffix = "_online" if mode == "online" else "" + for stem in ("dense_model_basic", "moe_model_basic"): + script = f"examples/architecture/pd-af-disagg/{mode}/{stem}{suffix}.sh" + cases.append( + Case(case_id=f"G11-pd-af-disagg-{mode}-{stem}-pp2", group="G11", + recipe=script, recipe_env=(("PREFILL_PP", "2"),)) + ) + return cases + + +# --------------------------------------------------------------------------- +# Fixture (runs inside the child process) +# --------------------------------------------------------------------------- + + +def _synthetic_model(is_moe: bool): + from frontier.config import BaseModelConfig + from frontier.types import ActivationType, NormType + + model = BaseModelConfig( + num_layers=6, num_q_heads=4, num_kv_heads=2, embedding_dim=256, + mlp_hidden_dim=64, max_position_embeddings=4096, use_gated_mlp=True, + use_bias=False, use_qkv_bias=False, activation=ActivationType.SILU, + norm=NormType.RMS_NORM, post_attn_norm=True, vocab_size=1024, + is_moe=is_moe, num_experts=8 if is_moe else 0, + num_experts_per_tok=2 if is_moe else 0, torch_dtype="bfloat16", + ) + model._model_name = f"stage_admission_{'moe' if is_moe else 'dense'}" + registered = BaseModelConfig.create_from_name.__func__ + BaseModelConfig.create_from_name = classmethod( + lambda cls, name: model if name == model._model_name else registered(cls, name) + ) + return model._model_name + + +def build_config(case: Case, output_dir: Path, cache_dir: Path): + """Build the SimulationConfig of one synthetic or vLLM-aligned case.""" + from frontier.cc_backend.cc_backend_config import AnalyticalCCBackendConfig + from frontier.config import ( + ClusterConfig, FixedRequestLengthGeneratorConfig, MetricsConfig, + PoissonRequestIntervalGeneratorConfig, RandomForrestExecutionTimePredictorConfig, + ReplicaConfig, RoundRobinClusterSchedulerConfig, SimulationConfig, + StaticRequestIntervalGeneratorConfig, SyntheticRequestGeneratorConfig, + VllmV1SchedulerConfig, + ) + + if case.fixture == SYNTHETIC: + model_name = _synthetic_model(case.is_moe) + device, network_device = "a100", "a100_pairwise_nvlink" + scheduler = VllmV1SchedulerConfig( + num_blocks=128, block_size=16, batch_size_cap=4, + max_tokens_in_batch=16, enable_chunked_prefill=True, + ) + elif case.fixture == VLLM_ALIGNED: + model_name = VLLM_ALIGNED_MODELS[case.is_moe] + device, network_device = "h800", "h800_dgx" + scheduler = VllmV1SchedulerConfig( + num_blocks=1024, block_size=16, batch_size_cap=4, + max_tokens_in_batch=case.prefill_tokens, enable_chunked_prefill=True, + ) + else: + raise ValueError(f"unknown fixture {case.fixture!r}") + + moe_fields = ( + dict(moe_tensor_parallel_size=1, moe_expert_parallel_size=case.attn_dp) + if case.is_moe else {} + ) + replica = ReplicaConfig( + model_name=model_name, device=device, network_device=network_device, + num_pipeline_stages=case.stages, attn_tensor_parallel_size=1, + attn_dp=case.attn_dp, memory_margin_fraction=0.1, **moe_fields, + ) + cluster_fields = {} + if case.cc_backend == "analytical": + cluster_fields["cc_backend_config"] = AnalyticalCCBackendConfig() + elif case.cc_backend != "default": + raise ValueError(f"unknown CC backend selector {case.cc_backend!r}") + if case.sys_arch == "pd-disaggregation": + # One Replica per role; both roles take the fixture's replica config. + cluster_fields.update(prefill_cluster_num_replicas=1, decode_cluster_num_replicas=1) + cluster = ClusterConfig( + replica_config=replica, + replica_scheduler_config=scheduler, + cluster_scheduler_config=RoundRobinClusterSchedulerConfig(), + execution_time_predictor_config=RandomForrestExecutionTimePredictorConfig( + enable_dummy_mode=True + ), + **cluster_fields, + ) + if case.arrival == "static": + interval = StaticRequestIntervalGeneratorConfig() + elif case.arrival == "poisson": + interval = PoissonRequestIntervalGeneratorConfig(qps=case.qps) + else: + raise ValueError(f"unknown arrival process {case.arrival!r}") + return SimulationConfig( + simulation_mode=case.simulation_mode, sys_arch=case.sys_arch, + enable_parallel_clusters=False, decode_cuda_graph_mode="none", + cluster_config=cluster, + metrics_config=MetricsConfig( + output_dir=str(output_dir), cache_dir=str(cache_dir), + run_id=case.case_id, write_metrics=True, store_request_metrics=True, + store_plots=False, enable_chrome_trace=False, write_json_trace=False, + ), + request_generator_config=SyntheticRequestGeneratorConfig( + num_requests=case.num_requests, + length_generator_config=FixedRequestLengthGeneratorConfig( + prefill_tokens=case.prefill_tokens, decode_tokens=case.decode_tokens, + ), + interval_generator_config=interval, + ), + ) + + +# --------------------------------------------------------------------------- +# Drain state and outcome classification (child process) +# --------------------------------------------------------------------------- + + +def _ticket_view(ticket) -> dict: + return { + "admission_seq": ticket.admission_seq, + "operation_id": str(ticket.operation_id), + "scope": ticket.scope, + } + + +def build_state_report(simulator) -> dict: + """Read stage contexts, lane queues and sync rooms of every cluster after a drain.""" + lanes = {} + contexts = [] + rooms = [] + for cluster_type, cluster_scheduler in simulator.scheduler._cluster_schedulers.items(): + _read_cluster_state(cluster_type.name, cluster_scheduler, lanes, contexts, rooms) + return { + "simulation_time": simulator._time, + "contexts": contexts, + "lanes": lanes, + "sync_rooms": rooms, + } + + +def _read_cluster_state(cluster: str, cluster_scheduler, lanes: dict, + contexts: list, rooms: list) -> None: + queued_owner = {} + for (replica_id, lane_id), replica_scheduler in sorted( + cluster_scheduler._replica_schedulers.items(), key=lambda item: str(item[0]) + ): + stage_views = [] + for stage_id in range(replica_scheduler._num_stages): + stage = replica_scheduler.get_replica_stage_scheduler(stage_id) + heap = [] + for batch in stage.get_queue_batches(): + ticket = batch._stage_admission_ticket + queued_owner[(replica_id, stage_id, ticket.admission_seq)] = lane_id + heap.append({"batch_id": batch.id, "global_id": batch.global_id, + **_ticket_view(ticket)}) + stage_views.append({"busy": stage.is_busy, "heap": heap}) + lanes[f"{cluster}/{replica_id}/{lane_id}"] = { + "cluster": cluster, "replica_id": replica_id, "lane": lane_id, + "stages": stage_views, + } + + for (replica_id, stage_id), context in sorted(cluster_scheduler._stage_execution_contexts.items()): + contexts.append({ + "cluster": cluster, + "replica_id": replica_id, + "stage_id": stage_id, + "capacity": context.full_stage_capacity, + "sealed": context.forward_group_sealed, + "bound_group": context._forward_group_id, + "ep_wave_active": context._active_ep_ticket is not None, + "active_full_stage": sorted( + (_ticket_view(ticket) for ticket in context._active_full_stage_tickets), + key=lambda view: view["admission_seq"], + ), + "fifo": [ + {**_ticket_view(ticket), + "lane": queued_owner.get((replica_id, stage_id, ticket.admission_seq))} + for ticket in context.queued_tickets + ], + }) + + for room_name in ("_prefill_sync_waiting_room", "_decode_sync_waiting_room"): + by_replica = getattr(cluster_scheduler, room_name) or {} + for replica_id, by_stage in by_replica.items(): + for stage_id, by_step in by_stage.items(): + for step, by_layer in by_step.items(): + for layer, by_sync in by_layer.items(): + for sync_stage, room in by_sync.items(): + if not room["batches"]: + continue + rooms.append({ + "cluster": cluster, "room": room_name.strip("_"), + "replica_id": replica_id, "stage_id": stage_id, + "step": step, "layer": layer, "sync_stage": str(sync_stage), + "lanes_present": sorted(room["batches"]), + }) + + +def has_admission_deadlock_signature(report: dict) -> bool: + """Plan §4.3: a busy lane's queued ticket heads the FIFO while an idle lane + with queued work, needed by that lane's sync room, is refused behind it.""" + lanes = report["lanes"] + for context in report["contexts"]: + if (context["ep_wave_active"] or context["sealed"] or not context["fifo"] + or len(context["active_full_stage"]) >= context["capacity"]): + continue + head = context["fifo"][0] + if head["scope"] != "FULL_STAGE_WORLD" or head["lane"] is None: + continue + cluster, replica_id, stage_id = context["cluster"], context["replica_id"], context["stage_id"] + head_stage = lanes[f"{cluster}/{replica_id}/{head['lane']}"]["stages"][stage_id] + if not head_stage["busy"]: + continue + for lane in lanes.values(): + if ((lane["cluster"], lane["replica_id"]) != (cluster, replica_id) + or lane["lane"] == head["lane"]): + continue + stage = lane["stages"][stage_id] + if stage["busy"] or not stage["heap"]: + continue + if stage["heap"][0]["admission_seq"] <= head["admission_seq"]: + continue + for room in report["sync_rooms"]: + if ((room["cluster"], room["replica_id"], room["stage_id"]) + == (cluster, replica_id, stage_id) + and head["lane"] in room["lanes_present"] + and lane["lane"] not in room["lanes_present"]): + return True + return False + + +def _run_simulator_case(case: Case, work_dir: Path, case_dir: Path) -> dict: + output_root = work_dir / "metrics" + try: + config = build_config(case, output_root, work_dir / "cache") + from frontier.simulator import Simulator + + simulator = Simulator(config) + except ValueError as exc: + (case_dir / "error.txt").write_text(traceback.format_exc()) + return {"outcome": CONFIGURATION_REJECTION, "exception": repr(exc)} + resolved_config = json.loads( + (Path(config.metrics_config.output_dir) / "config.json").read_text() + ) + (case_dir / "resolved_config.json").write_text(json.dumps(resolved_config, indent=1, sort_keys=True)) + try: + simulator.run() + except RuntimeError as exc: + if not str(exc).startswith(DRAIN_MESSAGE): + (case_dir / "error.txt").write_text(traceback.format_exc()) + return {"outcome": OTHER_FAILURE, "exception": repr(exc)[:2000]} + report = build_state_report(simulator) + (case_dir / "state_report.json").write_text(json.dumps(report, indent=1, sort_keys=True)) + outcome = ADMISSION_DEADLOCK if has_admission_deadlock_signature(report) else OTHER_FAILURE + return {"outcome": outcome, "exception": DRAIN_MESSAGE, + "simulation_time": report["simulation_time"]} + except Exception as exc: # classified, never converted into success + (case_dir / "error.txt").write_text(traceback.format_exc()) + return {"outcome": OTHER_FAILURE, "exception": repr(exc)[:2000]} + requests = list(simulator._all_requests) + completed = sum(1 for request in requests if request.completed) + if completed != len(requests): + (case_dir / "error.txt").write_text( + f"run returned with {completed} of {len(requests)} requests completed\n" + ) + return {"outcome": OTHER_FAILURE, "exception": "incomplete requests"} + return {"outcome": SUCCESS, "completed_requests": completed, + "metrics_run_dir": str(Path(config.metrics_config.output_dir).relative_to(work_dir))} + + +def _run_recipe_case(case: Case, work_dir: Path, case_dir: Path) -> dict: + env = dict(os.environ) + env.update({ + "PYTHON_BIN": sys.executable, + "METRICS_OUTPUT_DIR": str(work_dir / "metrics"), + "RUN_ID": case.case_id, + **dict(case.recipe_env), + }) + result = subprocess.run( + ["bash", str(REPO_ROOT / case.recipe)], cwd=REPO_ROOT, env=env, + capture_output=True, text=True, + ) + (case_dir / "stdout.log").write_text(result.stdout[-200_000:] + result.stderr[-200_000:]) + if result.returncode != 0: + (case_dir / "error.txt").write_text(result.stderr[-50_000:]) + return {"outcome": OTHER_FAILURE, "exception": f"exit code {result.returncode}"} + return {"outcome": SUCCESS} + + +def run_case_in_child(case: Case, root: Path, set_name: str) -> None: + """Child entry point: run one case, then publish its artifacts.""" + case_dir = root / set_name / case.case_id + work_dir = root / "work" / case.case_id + for directory in (case_dir, work_dir): + if directory.exists(): + shutil.rmtree(directory) + directory.mkdir(parents=True) + started = time.time() + if case.recipe is None: + result = _run_simulator_case(case, work_dir, case_dir) + else: + result = _run_recipe_case(case, work_dir, case_dir) + result["wall_start"] = started + result["wall_end"] = time.time() + if result["outcome"] == SUCCESS: + shutil.copytree(work_dir / "metrics", case_dir / "metrics") + (case_dir / "sha256sums.txt").write_text(sha256_lines(case_dir / "metrics")) + shutil.rmtree(work_dir) + (case_dir / "outcome.json").write_text(json.dumps(result, indent=1, sort_keys=True)) + + +def sha256_lines(directory: Path) -> str: + lines = [] + for path in sorted(p for p in directory.rglob("*") if p.is_file()): + digest = hashlib.sha256(path.read_bytes()).hexdigest() + lines.append(f"{digest} {path.relative_to(directory)}") + return "\n".join(lines) + "\n" + + +# --------------------------------------------------------------------------- +# Lane-overlap metric (plan §4.5) +# --------------------------------------------------------------------------- + + +def interval_overlap(intervals: Iterable[tuple[float, float, object]]) -> dict: + """Half-open ``[start, end)`` intervals keyed by lane. + + Returns the time with at least one open lane, the time with at least two + distinct open lanes, the peak number of distinct open lanes, the latest + end, and whether any lane overlaps itself. + """ + intervals = list(intervals) + events = [] + for start, end, lane in intervals: + if end > start: + events.append((start, 1, lane)) + events.append((end, -1, lane)) + # Ends sort before starts at the same instant: touching intervals do not overlap. + events.sort(key=lambda event: (event[0], event[1])) + open_by_lane: dict[object, int] = defaultdict(int) + busy_time = multi_lane_time = 0.0 + peak_lanes = 0 + self_overlap = False + previous_time = None + for event_time, delta, lane in events: + if previous_time is not None: + open_lanes = sum(1 for count in open_by_lane.values() if count > 0) + if open_lanes >= 1: + busy_time += event_time - previous_time + if open_lanes >= 2: + multi_lane_time += event_time - previous_time + open_by_lane[lane] += delta + if open_by_lane[lane] > 1: + self_overlap = True + peak_lanes = max(peak_lanes, sum(1 for count in open_by_lane.values() if count > 0)) + previous_time = event_time + return { + "busy_time": busy_time, + "multi_lane_busy_time": multi_lane_time, + "peak_lanes": peak_lanes, + "makespan": max((end for _, end, _ in intervals), default=0.0), + "self_overlap": self_overlap, + } + + +def read_ledger(metrics_dir: Path) -> list[dict]: + paths = sorted(metrics_dir.rglob("frontier_stage_batch_ledger.jsonl")) + if len(paths) != 1: + raise ValueError(f"expected one stage ledger under {metrics_dir}, found {len(paths)}") + return [json.loads(line) for line in paths[0].read_text().splitlines() if line.strip()] + + +def lane_intervals(rows: Sequence[dict]) -> dict[tuple, list[tuple[float, float, int]]]: + """``ATTN_DP_LANE`` ledger intervals per physical stage.""" + by_stage: dict[tuple, list[tuple[float, float, int]]] = defaultdict(list) + for row in rows: + if row["execution_scope"] != ATTN_DP_LANE: + continue + key = (row["cluster_type"], row["replica_id"], row["stage_id"]) + by_stage[key].append((row["stage_start_ts"], row["stage_end_ts"], row["replica_local_id"])) + return dict(by_stage) + + +def ledger_lane_metric(metrics_dir: Path) -> dict: + """Plan §4.5 metric per physical stage, keyed ``cluster/replica/stage``.""" + stages = {} + for (cluster_type, replica_id, stage_id), intervals in sorted(lane_intervals(read_ledger(metrics_dir)).items()): + metric = interval_overlap(intervals) + metric["lanes"] = sorted({lane for _, _, lane in intervals}) + stages[f"{cluster_type}/{replica_id}/{stage_id}"] = metric + return stages + + +# --------------------------------------------------------------------------- +# Parent: set runner and provenance +# --------------------------------------------------------------------------- + + +def _git(*args: str) -> str: + return subprocess.run(["git", "-C", str(REPO_ROOT), *args], check=True, + capture_output=True, text=True).stdout.strip() + + +def set_provenance() -> dict: + distributions = sorted( + f"{dist.metadata['Name']}=={dist.version}" for dist in importlib_metadata.distributions() + ) + status = _git("status", "--porcelain", "--", ".", ":!task_memory") + return { + "interpreter": sys.executable, + "python_vv": subprocess.run([sys.executable, "-VV"], check=True, + capture_output=True, text=True).stdout.strip(), + "distributions_sha256": hashlib.sha256("\n".join(distributions).encode()).hexdigest(), + "git_head": _git("rev-parse", "HEAD"), + "clean_outside_task_memory": status == "", + "status_outside_task_memory": status, + } + + +def matrix_root() -> Path: + return resolve_scratch_root() / MATRIX_DIR_NAME + + +def _run_one(case: Case, root: Path, set_name: str, provenance: dict, + case_timeout: float) -> dict: + command = [sys.executable, "-m", "tests.e2e.stage_admission_matrix", "child", + "--set", set_name, "--case", case.case_id] + env = dict(os.environ, PYTHONPATH=str(REPO_ROOT), WANDB_DISABLED="true", + VIDUR_DISABLE_WANDB="1") + child = subprocess.Popen(command, cwd=REPO_ROOT, env=env, stdout=subprocess.PIPE, + stderr=subprocess.PIPE, text=True, start_new_session=True) + try: + stdout, stderr = child.communicate(timeout=case_timeout) + failure = f"child exit code {child.returncode}" + except subprocess.TimeoutExpired: + os.killpg(child.pid, signal.SIGKILL) + stdout, stderr = child.communicate() + failure = f"case timeout after {case_timeout:g} s" + case_dir = root / set_name / case.case_id + case_dir.mkdir(parents=True, exist_ok=True) + outcome_path = case_dir / "outcome.json" + if child.returncode != 0 or not outcome_path.exists(): + (case_dir / "error.txt").write_text(stdout[-50_000:] + stderr[-50_000:]) + outcome = {"outcome": OTHER_FAILURE, "exception": failure} + else: + outcome = json.loads(outcome_path.read_text()) + (case_dir / "case.json").write_text(json.dumps(asdict(case), indent=1, sort_keys=True)) + run = {"command": command, **provenance, **outcome} + (case_dir / "run.json").write_text(json.dumps(run, indent=1, sort_keys=True)) + return {"case_id": case.case_id, "group": case.group, "outcome": outcome["outcome"], + "exception": outcome.get("exception")} + + +def run_set(set_name: str, cases: Sequence[Case], jobs: int, case_timeout: float) -> list[dict]: + root = matrix_root() + (root / set_name).mkdir(parents=True, exist_ok=True) + with (root / "run.lock").open("w") as lock: + try: + fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError: + raise RuntimeError( + f"another set is running under {root}; sets share work/ and run one at a time" + ) from None + return _run_locked_set(root, set_name, cases, jobs, case_timeout) + + +def _run_locked_set(root: Path, set_name: str, cases: Sequence[Case], jobs: int, + case_timeout: float) -> list[dict]: + provenance = set_provenance() + with ThreadPoolExecutor(max_workers=jobs) as pool: + rows = list(pool.map( + lambda case: _run_one(case, root, set_name, provenance, case_timeout), cases + )) + if set_provenance()["git_head"] != provenance["git_head"]: + raise RuntimeError("git HEAD changed while the set was running") + index = root / set_name / "cases.jsonl" + existing = {} + if index.exists(): + existing = {row["case_id"]: row for row in map(json.loads, index.read_text().splitlines())} + existing.update({row["case_id"]: row for row in rows}) + order = [case.case_id for case in build_cases()] + index.write_text("".join(json.dumps(existing[case_id]) + "\n" + for case_id in order if case_id in existing)) + return rows + + +# --------------------------------------------------------------------------- +# Parent: before/after comparison (plan §4.4) +# --------------------------------------------------------------------------- + + +def _case_state(root: Path, set_name: str, case_id: str) -> dict: + case_dir = root / set_name / case_id + run = json.loads((case_dir / "run.json").read_text()) + state = {"outcome": run["outcome"]} + if run["outcome"] == SUCCESS: + state["sha256sums"] = (case_dir / "sha256sums.txt").read_text() + if (case_dir / "metrics").exists(): + state["metrics_dir"] = case_dir / "metrics" + return state + + +def _conservation(case: Case, metrics_dir: Path) -> dict: + paths = sorted(metrics_dir.rglob("request_metrics.csv")) + if len(paths) != 1: + return {"ok": False, "reason": f"{len(paths)} request_metrics.csv files"} + with paths[0].open() as handle: + rows = list(csv.DictReader(handle)) + prefill = sum(int(float(row["request_num_prefill_tokens"])) for row in rows) + decode = sum(int(float(row["request_num_decode_tokens"])) for row in rows) + expected = (case.num_requests, case.num_requests * case.prefill_tokens, + case.num_requests * case.decode_tokens) + observed = (len(rows), prefill, decode) + return {"ok": observed == expected, "observed": observed, "expected": expected} + + +def _differing_files(before: str, after: str) -> list[str]: + def parse(text): + return dict(reversed(line.split(" ", 1)) for line in text.splitlines() if line) + left, right = parse(before), parse(after) + return sorted(name for name in set(left) | set(right) if left.get(name) != right.get(name)) + + +def compare_sets(before: str, after: str) -> list[dict]: + root = matrix_root() + rows = [] + for case in build_cases(): + if not (root / before / case.case_id / "run.json").exists(): + continue + base = _case_state(root, before, case.case_id) + new = _case_state(root, after, case.case_id) + row = {"case_id": case.case_id, "group": case.group, + "before": base["outcome"], "after": new["outcome"]} + if case.group == "R0": + row.update(path="R0", verdict="informational") + elif case.path == "U": + row["path"] = "U" + identical = base["outcome"] == new["outcome"] == SUCCESS and base["sha256sums"] == new["sha256sums"] + row["verdict"] = "PASS" if identical else "STOP" + if base["outcome"] == new["outcome"] == SUCCESS and not identical: + row["differing_files"] = _differing_files(base["sha256sums"], new["sha256sums"]) + elif base["outcome"] == ADMISSION_DEADLOCK: + row["path"] = "L" + if new["outcome"] == SUCCESS: + row["conservation"] = _conservation(case, new["metrics_dir"]) + row["verdict"] = "PASS" if row["conservation"]["ok"] else "STOP" + else: + row["verdict"] = "STOP" + elif base["outcome"] == SUCCESS: + row["path"] = "T" + if new["outcome"] != SUCCESS: + row["verdict"] = "STOP" + else: + before_metric = ledger_lane_metric(base["metrics_dir"]) + after_metric = ledger_lane_metric(new["metrics_dir"]) + row["lane_metric_before"] = before_metric + row["lane_metric_after"] = after_metric + checks_ok = all( + not stage["self_overlap"] and stage["peak_lanes"] <= case.attn_dp + for stage in after_metric.values() + ) + identical = base["sha256sums"] == new["sha256sums"] + if not identical: + row["differing_files"] = _differing_files(base["sha256sums"], new["sha256sums"]) + if case.contention_witness: + # The fraction, not the absolute overlap: admitting more lanes + # together also shortens the busy period. + fraction = lambda metric: (sum(stage["multi_lane_busy_time"] for stage in metric.values()) + / sum(stage["busy_time"] for stage in metric.values())) + row["witness_increase"] = fraction(after_metric) > fraction(before_metric) + checks_ok = checks_ok and row["witness_increase"] + row["verdict"] = ("PASS" if identical and checks_ok + else "EXPLAIN" if checks_ok else "STOP") + else: + row.update(path="class", verdict="STOP") + rows.append(row) + return rows + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__.split("\n", 1)[0]) + commands = parser.add_subparsers(dest="command", required=True) + run_parser = commands.add_parser("run", help="run cases into one named set") + run_parser.add_argument("--set", required=True) + run_parser.add_argument("--group", action="append", default=[]) + run_parser.add_argument("--case", action="append", default=[]) + run_parser.add_argument("--jobs", type=int, default=8) + run_parser.add_argument("--case-timeout", type=float, default=DEFAULT_CASE_TIMEOUT_S, + help="seconds before a case's child session is killed") + child_parser = commands.add_parser("child", help=argparse.SUPPRESS) + child_parser.add_argument("--set", required=True) + child_parser.add_argument("--case", required=True) + compare_parser = commands.add_parser("compare", help="apply the plan §4.4 paths") + compare_parser.add_argument("--before", required=True) + compare_parser.add_argument("--after", required=True) + compare_parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args(argv) + + cases_by_id = {case.case_id: case for case in build_cases()} + if args.command == "child": + run_case_in_child(cases_by_id[args.case], matrix_root(), args.set) + return 0 + if args.command == "compare": + rows = compare_sets(args.before, args.after) + args.output.write_text(json.dumps(rows, indent=1, sort_keys=True, default=str)) + for row in rows: + print(f"{row['case_id']:<48} {row['path']:<5} {row['before']:<24} {row['after']:<24} {row['verdict']}") + return 0 + selected = [case for case in cases_by_id.values() + if (not args.group or case.group in args.group) + and (not args.case or case.case_id in args.case)] + for row in run_set(args.set, selected, args.jobs, args.case_timeout): + print(f"{row['case_id']:<48} {row['outcome']}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/integration/test_stage_admission_pipeline_lanes.py b/tests/integration/test_stage_admission_pipeline_lanes.py new file mode 100644 index 00000000..11970c18 --- /dev/null +++ b/tests/integration/test_stage_admission_pipeline_lanes.py @@ -0,0 +1,65 @@ +"""Simulator regression for attention-DP lanes sharing pipeline stages.""" + +import csv +import json +import os +from pathlib import Path +import subprocess +import sys + +import pytest + +from tests.e2e.stage_admission_matrix import ATTN_DP_LANE, SUCCESS, build_cases, read_ledger, run_case_in_child + +REPO_ROOT = Path(__file__).resolve().parents[2] +SET_NAME = "test" + + +def run_case(root, case_id): + """Run one matrix case in its own process, because ``IS_MOE`` is process-global.""" + result = subprocess.run( + [sys.executable, str(Path(__file__).resolve()), str(root), case_id], + cwd=REPO_ROOT, + env={**os.environ, "PYTHONPATH": str(REPO_ROOT), "WANDB_DISABLED": "true", + "VIDUR_DISABLE_WANDB": "1", "OMP_NUM_THREADS": "1", "OPENBLAS_NUM_THREADS": "1"}, + text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, timeout=300, + ) + (root / "run.log").write_text(result.stdout) + assert result.returncode == 0, result.stdout[-15000:] + case_dir = root / SET_NAME / case_id + outcome = json.loads((case_dir / "outcome.json").read_text()) + assert outcome["outcome"] == SUCCESS, outcome + return case_dir / "metrics" + + +@pytest.mark.parametrize("case_id, expected", [ + # (requests, prefill tokens, decode tokens): 16 prompt tokens and one output token each. + ("G3a-moe-dp2-pp2-n4", (4, 64, 4)), + ("G3a-moe-dp4-pp2-n8", (8, 128, 8)), +]) +def test_moe_lanes_complete_every_request(tmp_path, case_id, expected): + metrics_dir = run_case(tmp_path, case_id) + with next(metrics_dir.rglob("request_metrics.csv")).open() as handle: + rows = list(csv.DictReader(handle)) + observed = ( + len(rows), + sum(int(float(row["request_num_prefill_tokens"])) for row in rows), + sum(int(float(row["request_num_decode_tokens"])) for row in rows), + ) + assert observed == expected + + +def test_dense_lanes_start_in_the_same_first_forward(tmp_path): + """Every request arrives at t=0 and the stage has room for both lanes.""" + rows = read_ledger(run_case(tmp_path, "G4-dense-dp2-pp2-n8")) + first_start = {} + for row in sorted(rows, key=lambda row: row["stage_start_ts"]): + if row["execution_scope"] == ATTN_DP_LANE and row["stage_id"] == 0: + first_start.setdefault(row["replica_local_id"], row["stage_start_ts"]) + assert sorted(first_start) == [0, 1] + assert first_start[0] == first_start[1] + + +if __name__ == "__main__": + cases = {case.case_id: case for case in build_cases()} + run_case_in_child(cases[sys.argv[2]], Path(sys.argv[1]), SET_NAME) diff --git a/tests/unit/test_mixed_layer_decode_ffn_scheduling.py b/tests/unit/test_mixed_layer_decode_ffn_scheduling.py index 7f3db6fe..9c959682 100644 --- a/tests/unit/test_mixed_layer_decode_ffn_scheduling.py +++ b/tests/unit/test_mixed_layer_decode_ffn_scheduling.py @@ -799,6 +799,75 @@ def test_decode_ffn_wave_materialization_attaches_one_parent_ticket( assert context.queued_tickets == (tickets[0],) +def test_decode_ffn_dense_groups_keep_counter_order_around_a_queued_ep_wave( + mixed_model_config, +) -> None: + """Dense groups enter in group-counter order and never pass an EP wave queued ahead.""" + + scheduler, _, _, lane_sinks = _atomicity_scheduler( + mixed_model_config, + layer_id=4, + ep_size=2, + ) + full_stage_sink = _QueuedBatchSink() + scheduler._full_stage_replica_schedulers = {0: full_stage_sink} + scheduler.get_full_stage_replica_scheduler = Mock(return_value=full_stage_sink) + scheduler._m2n_ready_groups = deque( + [ + [(_source_batch(layer_id=layer_id), _transfer_info(layer_id=layer_id))] + for layer_id in (3, 4, 3) + ] + ) + context = scheduler.get_stage_execution_context(0, 2) + for _ in range(3): + scheduler.schedule_ffn_with_m2n_immediate() + + first_dense, second_dense = full_stage_sink._m2n_immediate_batch_queue + lane_batches = [lane_sinks[ep_id]._m2n_immediate_batch_queue[0] for ep_id in (0, 1)] + wave = lane_batches[0]._stage_admission_ticket + assert [first_dense.global_id, lane_batches[0].global_id, second_dense.global_id] == [0, 1, 2] + assert context.queued_tickets == ( + first_dense._stage_admission_ticket, + wave, + second_dense._stage_admission_ticket, + ) + + def stage_scheduler(replica_local_id): + return ReplicaStageScheduler( + replica_id=0, + stage_id=2, + is_last_stage=True, + is_moe=True, + execution_time_predictor=object(), + cluster_type=ClusterType.DECODE_FFN, + replica_local_id=replica_local_id, + stage_execution_context=context, + ) + + full_stage = stage_scheduler(None) + ep_lanes = [stage_scheduler(ep_id) for ep_id in (0, 1)] + full_stage.add_batch(first_dense) + full_stage.add_batch(second_dense) + for lane, batch in zip(ep_lanes, lane_batches): + lane.add_batch(batch) + assert full_stage.get_queue_batches() == [first_dense, second_dense] + + assert full_stage.pop_batch_if_not_busy() is first_dense + full_stage.on_stage_end() + context.release(first_dense._stage_admission_ticket) + assert full_stage.pop_batch_if_not_busy() is None + + assert ep_lanes[0].pop_batch_if_not_busy() is lane_batches[0] + assert ep_lanes[1].pop_batch_if_not_busy() is lane_batches[1] + assert full_stage.pop_batch_if_not_busy() is None + for lane in ep_lanes: + lane.on_stage_end() + context.release(wave) + + assert full_stage.pop_batch_if_not_busy() is second_dense + assert context.queued_tickets == () + + def _atomicity_snapshot( scheduler, source_batch, queue_sinks, *, include_entity_ids: bool = True ): diff --git a/tests/unit/test_shared_forward_group_admission.py b/tests/unit/test_shared_forward_group_admission.py index 7b5284f8..5d97ebc7 100644 --- a/tests/unit/test_shared_forward_group_admission.py +++ b/tests/unit/test_shared_forward_group_admission.py @@ -51,6 +51,43 @@ def test_admitted_lanes_share_identity_after_unequal_batch_histories(cluster_typ assert [batch.global_id for batch in batches] == [2, 1] +@pytest.mark.parametrize("first_lane", [0, 1]) +def test_idle_lane_is_admitted_behind_a_busy_lane_queued_ticket(first_lane): + """Under PP a busy lane can hold the FIFO head; the other lane must still join.""" + + other_lane = 1 - first_lane + context = StageExecutionContext(replica_id=0, stage_id=0, ep_size=2, full_stage_capacity=2) + stages = [make_stage(context, lane, ClusterType.MONOLITHIC) for lane in range(2)] + first_now, first_next = make_batch(first_lane, 0), make_batch(first_lane, 1) + # Distinct provisional ids, so that sharing a bound group is observable. + other_now, other_next = make_batch(other_lane, 5), make_batch(other_lane, 6) + stages[first_lane].add_batch(first_now) + assert stages[first_lane].pop_batch_if_not_busy() is first_now + stages[first_lane].add_batch(first_next) + stages[other_lane].add_batch(other_now) + stages[other_lane].add_batch(other_next) + assert context.queued_tickets[0] == first_next._stage_admission_ticket + + assert stages[other_lane].pop_batch_if_not_busy() is other_now + assert other_now._forward_cohort_provisional_id == first_now._forward_cohort_provisional_id + + wave = context.replace_full_stage_owners_with_ep_wave( + (first_now._stage_admission_ticket, other_now._stage_admission_ticket), + operation_id="wave", participant_ep_ids=(0, 1), + ) + owners = context.replace_ep_wave_with_full_stage_owners(wave, operation_ids=("restored0", "restored1")) + for owner in owners: + context.release(owner) + for stage in stages: + stage.on_stage_end() + + assert stages[first_lane].pop_batch_if_not_busy() is first_next + assert stages[other_lane].pop_batch_if_not_busy() is other_next + assert first_next._forward_cohort_provisional_id == other_next._forward_cohort_provisional_id + assert first_next._forward_cohort_provisional_id > first_now._forward_cohort_provisional_id + assert context.queued_tickets == () + + def test_started_group_blocks_new_lane_through_ep_restore_and_partial_release(): context = StageExecutionContext(replica_id=0, stage_id=0, ep_size=2, full_stage_capacity=4) first = context.enqueue_full_stage(operation_id="first") diff --git a/tests/unit/test_stage_admission_pp_tools.py b/tests/unit/test_stage_admission_pp_tools.py new file mode 100644 index 00000000..ea88f196 --- /dev/null +++ b/tests/unit/test_stage_admission_pp_tools.py @@ -0,0 +1,237 @@ +"""Synthetic checks of the stage-admission vLLM comparison tools. + +``compare_lanes`` is driven end to end on hand-built vLLM traces and Frontier +sets: ideal lane pairing on both sides, and the base rule's two negative +controls (a MoE admission deadlock, dense lanes one forward apart). +``vllm_burst_driver`` is checked for overlay acceptance and patch parsing. +""" + +from __future__ import annotations + +import csv +import json +from pathlib import Path + +import pytest + +from tests.comparison.stage_admission_pp import compare_lanes +from tests.comparison.stage_admission_pp.vllm_burst_driver import apply_patch, build_overlay +from tests.e2e.stage_admission_matrix import ADMISSION_DEADLOCK, MATRIX_DIR_NAME, SUCCESS + + +FORWARD = 0.12 +ROUNDS = 3 + + +def forwards(num_requests: int, late_lane: bool = False) -> list[tuple[int, int, float, float, int]]: + """Two-stage forwards: pair ``k`` runs stage 0 in ``[k*F, (k+1)*F)``. + + With ``late_lane`` lane 1 starts one forward after lane 0, as under the + base rule. + """ + rows = [] + for index in range(num_requests): + lane, slot = index % 2, index // 2 + start = (slot + (1 if late_lane and lane == 1 else 0)) * FORWARD + rows.append((lane, 0, start, start + FORWARD, index)) + rows.append((lane, 1, start + FORWARD, start + 2 * FORWARD, index)) + return rows + + +def write_jsonl(path: Path, records) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("".join(json.dumps(record) + "\n" for record in records)) + + +def write_vllm_scenario(run_dir: Path, model: str, late_round: int | None = None) -> None: + scenario = run_dir / "runs" / model + requests, boundaries, rounds = [], [], [] + placement = {0: [], 1: []} + wall_offset, origin = 1.7e9, 100.0 + labels = [("warmup", 0, 4)] + [ + (f"b{burst}-r{r}", r, burst) for burst in compare_lanes.BURSTS for r in range(ROUNDS) + ] + for label, round_index, num_requests in labels: + late = label != "warmup" and round_index == late_round + for lane, stage, start, end, index in forwards(num_requests, late): + request_id = f"{label}-q{index}" + boundaries.append({ + "request_ids": [request_id], "pp_rank": stage, "is_last_rank": stage == 1, + "forward_start_ts": origin + start, + "send_start_ts": None if stage else origin + end, + "timestamp": wall_offset + origin + end, + }) + if stage == 0: + placement[lane].append({"kind": "engine_iteration", "engine": lane, + "scheduled_new_req_ids": [request_id]}) + requests.extend( + {"request_id": f"{label}-q{index}", "burst": label, "round": round_index, + "index": index, "rank": index % 2, "num_output_tokens": 1} + for index in range(num_requests) + ) + rounds.append({"label": label, "wall_minus_monotonic_before": wall_offset, + "wall_minus_monotonic_after": wall_offset}) + origin += 10.0 + write_jsonl(scenario / "requests.jsonl", requests) + write_jsonl(scenario / "pp_boundary.jsonl", boundaries) + (scenario / "summary.json").write_text(json.dumps({"rounds": rounds})) + for engine, records in placement.items(): + write_jsonl(scenario / "dp_placement" / f"dp_placement_{engine}.jsonl", records) + + +def write_frontier_case(set_dir: Path, model: str, burst: int, outcome: str, + late_lane: bool = False) -> None: + case_dir = set_dir / f"G7-{model}-dp2-pp2-n{burst}" + case_dir.mkdir(parents=True) + (case_dir / "run.json").write_text(json.dumps({"outcome": outcome})) + (case_dir / "case.json").write_text(json.dumps({"num_requests": burst})) + if outcome == SUCCESS: + write_jsonl( + case_dir / "metrics" / "run" / "frontier_stage_batch_ledger.jsonl", + ({"execution_scope": "ATTN_DP_LANE", "replica_local_id": lane, "stage_id": stage, + "stage_start_ts": start, "stage_end_ts": end, "request_ids": [str(index)]} + for lane, stage, start, end, index in forwards(burst, late_lane)), + ) + + +@pytest.fixture +def workspace(tmp_path, monkeypatch): + monkeypatch.setenv("FRONTIER_TMP_ROOT", str(tmp_path / "scratch")) + return tmp_path + + +def run_compare(workspace: Path, *, fixed_base: bool = False, late_round: int | None = None, + drop_placement_engine: int | None = None) -> tuple[dict, list[dict]]: + vllm_run = workspace / "vllm" + for model in compare_lanes.MODELS: + write_vllm_scenario(vllm_run, model, late_round if model == "dense" else None) + if drop_placement_engine is not None: + (vllm_run / "runs" / "moe" / "dp_placement" + / f"dp_placement_{drop_placement_engine}.jsonl").unlink() + frontier = workspace / "scratch" / MATRIX_DIR_NAME + for burst in compare_lanes.BURSTS: + for model in compare_lanes.MODELS: + write_frontier_case(frontier / "after", model, burst, SUCCESS) + if fixed_base: + write_frontier_case(frontier / "base", model, burst, SUCCESS) + elif model == "moe": + write_frontier_case(frontier / "base", model, burst, ADMISSION_DEADLOCK) + else: + write_frontier_case(frontier / "base", model, burst, SUCCESS, late_lane=True) + output = workspace / "analysis" + compare_lanes.main(["--vllm-run", str(vllm_run), "--output", str(output)]) + status = json.loads((output / "workflow_gap_status.json").read_text()) + with (output / "workflow_gap_table.csv").open() as handle: + rows = list(csv.DictReader(handle)) + return status, rows + + +def statuses(rows: list[dict], checks: tuple[str, ...]) -> set[str]: + return {row["status"] for row in rows if row["check"] in checks} + + +def test_ideal_after_revision_matches_and_base_controls_hold(workspace) -> None: + status, rows = run_compare(workspace) + + assert status["status"] == "PASS" + assert status["mismatches"] == 0 + assert status["negative_control_holds"] is True + assert statuses(rows, ("V1", "V2", "V3", "V4")) == {"MATCH"} + assert {(row["model"], row["status"]) for row in rows if row["check"] == "V5"} == { + ("moe", "MATCH"), ("dense", compare_lanes.INFORMATIONAL) + } + assert sorted((row["check"], row["model"], row["status"]) for row in rows + if row["check"] in ("N1", "N4")) == [ + ("N1", "moe", "HOLDS"), ("N1", "moe", "HOLDS"), + ("N4", "dense", "HOLDS"), ("N4", "dense", "HOLDS"), + ] + + +def test_fixed_base_loses_the_controls_without_a_mismatch(workspace) -> None: + status, rows = run_compare(workspace, fixed_base=True) + + assert status["status"] == "PASS" + assert status["mismatches"] == 0 + assert status["negative_control_holds"] is False + assert statuses(rows, ("V1", "V2", "V3", "V4")) == {"MATCH"} + assert statuses(rows, ("N1", "N4")) == {"LOST"} + + +def test_vllm_round_with_a_late_lane_is_reported(workspace) -> None: + status, rows = run_compare(workspace, late_round=1) + + mismatched = sorted((row["check"], row["model"], row["burst"], row["round"]) + for row in rows if row["status"] == "MISMATCH") + assert mismatched == [ + ("V3", "dense", "16", "1"), ("V3", "dense", "8", "1"), + ("V4", "dense", "16", "1"), ("V4", "dense", "8", "1"), + ] + assert status["status"] == "FAIL" + + +def test_missing_placement_log_fails_the_comparison(workspace) -> None: + status, _ = run_compare(workspace, drop_placement_engine=1) + + assert status["vllm_placement_ok"] is False + assert status["vllm_placement_unseen_requests"] > 0 + assert status["status"] == "FAIL" + + +def write_tree(root: Path, files: dict[str, str]) -> None: + for relative, text in files.items(): + path = root / relative + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(text) + + +@pytest.mark.parametrize("extra_change, accepted", [(False, True), (True, False)]) +def test_overlay_acceptance_compares_file_sets(tmp_path, extra_change, accepted) -> None: + # Path order puts vllm/a/c.py first; string order puts vllm/a-b.py first. + changed = ["vllm/a-b.py", "vllm/a/c.py"] + write_tree(tmp_path / "site", {"vllm/a-b.py": "old\n", "vllm/a/c.py": "old\n", "vllm/d.py": "same\n"}) + checkout_files = {"vllm/a-b.py": "new\n", "vllm/a/c.py": "new\n", "vllm/d.py": "same\n"} + if extra_change: + checkout_files["vllm/d.py"] = "changed\n" + write_tree(tmp_path / "checkout", checkout_files) + expected = tmp_path / "expected_changes.txt" + expected.write_text("".join(f"{name}\n" for name in sorted(changed))) + + report = build_overlay(tmp_path / "site" / "vllm", tmp_path / "checkout", + tmp_path / "overlay", expected) + + assert report["accepted"] is accepted + assert report["unexpected"] == ([] if accepted else ["vllm/d.py"]) + + +def test_apply_patch_keeps_every_section_of_a_repeated_file(tmp_path) -> None: + write_tree(tmp_path, {"pkg/mod.py": "a = 1\nb = 2\nc = 3\n"}) + patch = tmp_path / "change.patch" + patch.write_text( + "--- a/pkg/mod.py\n+++ b/pkg/mod.py\n@@ -1,1 +1,1 @@\n-a = 1\n+a = 10\n" + "--- a/pkg/mod.py\n+++ b/pkg/mod.py\n@@ -3,1 +3,1 @@\n-c = 3\n+c = 30\n" + ) + + assert apply_patch(patch, tmp_path) == ["pkg/mod.py"] + assert (tmp_path / "pkg" / "mod.py").read_text() == "a = 10\nb = 2\nc = 30\n" + + +def test_apply_patch_reads_a_trimmed_context_line(tmp_path) -> None: + write_tree(tmp_path, {"pkg/mod.py": "a = 1\n\nb = 2\n"}) + patch = tmp_path / "change.patch" + patch.write_text("--- a/pkg/mod.py\n+++ b/pkg/mod.py\n@@ -1,3 +1,3 @@\n a = 1\n\n-b = 2\n+b = 20\n") + + apply_patch(patch, tmp_path) + + assert (tmp_path / "pkg" / "mod.py").read_text() == "a = 1\n\nb = 20\n" + + +def test_apply_patch_rejects_an_unknown_hunk_line(tmp_path) -> None: + write_tree(tmp_path, {"pkg/mod.py": "a = 1\n"}) + patch = tmp_path / "change.patch" + patch.write_text( + "--- a/pkg/mod.py\n+++ b/pkg/mod.py\n@@ -1,1 +1,1 @@\n-a = 1\n" + "\\ No newline at end of file\n+a = 2\n" + ) + + with pytest.raises(ValueError, match="unexpected hunk line"): + apply_patch(patch, tmp_path) diff --git a/tests/unit/test_stage_execution_context.py b/tests/unit/test_stage_execution_context.py index 4eb05de0..0e471f5c 100644 --- a/tests/unit/test_stage_execution_context.py +++ b/tests/unit/test_stage_execution_context.py @@ -91,6 +91,74 @@ def test_admission_fifo_cannot_skip_an_earlier_ready_wave() -> None: context.release(second) +def _context_with_wave_between_full_stage_tickets(): + context = StageExecutionContext( + replica_id=0, + stage_id=0, + ep_size=2, + full_stage_capacity=2, + ) + full0 = context.enqueue_full_stage(operation_id="full0") + full1 = context.enqueue_full_stage(operation_id="full1") + wave0 = context.enqueue_ep_wave(operation_id="wave0", participant_ep_ids=(0, 1)) + full2 = context.enqueue_full_stage(operation_id="full2") + return context, full0, full1, wave0, full2 + + +def test_full_stage_ticket_passes_queued_full_stage_work_but_not_a_queued_wave() -> None: + context, full0, full1, wave0, full2 = _context_with_wave_between_full_stage_tickets() + + assert context.try_acquire(full1) is True + assert context.queued_tickets == (full0, wave0, full2) + # Capacity remains, but wave0 is queued ahead of full2. + assert context.try_acquire(full2) is False + assert context.try_acquire(full0) is True + + +def test_queued_ep_wave_orders_full_stage_work_on_both_sides() -> None: + context, full0, full1, wave0, full2 = _context_with_wave_between_full_stage_tickets() + + assert context.try_acquire(full0) is True + assert context.try_acquire(full1) is True + context.release(full1) + assert context.try_acquire(full2) is False + assert context.try_acquire(wave0) is False + context.release(full0) + assert context.try_acquire(wave0) is True + assert context.try_acquire(full2) is False + context.release(wave0) + assert context.try_acquire(full2) is True + context.release(full2) + assert context.is_idle + assert context.queued_tickets == () + + +def test_idle_single_owner_stage_admits_a_later_queued_full_stage_ticket() -> None: + context = StageExecutionContext(replica_id=0, stage_id=0, ep_size=1) + full0 = context.enqueue_full_stage(operation_id="full0") + full1 = context.enqueue_full_stage(operation_id="full1") + + assert context.try_acquire(full1) is True + assert context.queued_tickets == (full0,) + + +def test_active_full_stage_ticket_is_refused_without_changing_the_stage() -> None: + context = StageExecutionContext( + replica_id=0, + stage_id=0, + ep_size=1, + full_stage_capacity=2, + ) + first = context.enqueue_full_stage(operation_id=("lane", 0)) + second = context.enqueue_full_stage(operation_id=("lane", 1)) + assert context.try_acquire(first) is True + + assert context.try_acquire(first) is False + assert context.is_active(first) + assert context.queued_tickets == (second,) + assert context.try_acquire(second) is True + + def test_release_requires_the_active_operation_ticket() -> None: context = StageExecutionContext(replica_id=0, stage_id=0, ep_size=1) wave = context.enqueue_ep_wave(operation_id=30, participant_ep_ids=(0,))