diff --git a/.gitignore b/.gitignore index d7715a391..cf2859f68 100644 --- a/.gitignore +++ b/.gitignore @@ -14,6 +14,7 @@ build/ venv/ env/ ENV/ +.venv/ # IDE .vscode/ diff --git a/pyproject.toml b/pyproject.toml index 3b1539913..e495d7def 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,6 +32,21 @@ monitoring = [ include = ["monitoring*", "benchmark*", "integration"] [tool.pytest.ini_options] +# `manual` tools and native-ring sources are not collected by default. The +# default suite is the CPU contract/unit tests: +# python -m pytest -m "not gpu and not e2e and not manual" -q +addopts = "-ra -m 'not manual'" +norecursedirs = ["tests/tools", "tests/ring", ".venv", "integration", "build"] markers = [ + "cpu: pure-CPU contract/unit test; the default suite (no CUDA / ClickHouse / vLLM runtime / model weights / native extension build required)", + "gpu: requires a CUDA device", + "multi_gpu: requires >= 2 CUDA devices (TP / EP / routing)", + "e2e: end-to-end pipeline through the native backend + host engine", + "clickhouse: requires a reachable ClickHouse instance", + "vllm: requires the vLLM runtime importable", + "hf: requires HuggingFace weights / model cache", + "ring_native: native CUDA ring tests built via tests/ring/Makefile (needs nvcc)", "slow: tests that take more than ~30 s (per-hook isolation full sweep, large E2E sweeps); skipped by default unless `-m slow` is passed", + "manual: investigation / tooling, not a regression gate; not collected by default", + "numeric: per-hook numeric-difference study (drift vs the unhooked baseline)", ] diff --git a/tests/_requirements.py b/tests/_requirements.py new file mode 100644 index 000000000..bb40407ab --- /dev/null +++ b/tests/_requirements.py @@ -0,0 +1,128 @@ +"""Resource skip-guards for the test suite. + +Every GPU / E2E / native test should fail *closed with a precise reason* when a +prerequisite is missing, never with an import error or an opaque crash. Each +helper here returns a ``pytest.mark.skipif`` marker, so it can be used either as +a decorator:: + + from tests._requirements import require_cuda + + @require_cuda() + def test_kernel(): + ... + +or composed into a module-level mark list alongside a category marker:: + + pytestmark = [pytest.mark.gpu, require_cuda()] + +This module must stay importable on a CPU-only box with no CUDA, ClickHouse, +vLLM, model weights, or native build toolchain present. All heavy imports +(``torch`` in particular) are deferred into the helper bodies so that merely +importing this file costs nothing. +""" +from __future__ import annotations + +import importlib.util +import os +import shutil +import socket + +import pytest + +__all__ = [ + "require_cuda", + "require_gpus", + "require_clickhouse", + "require_vllm", + "require_model_cache", + "require_nvcc", +] + + +def _cuda_device_count() -> int: + """CUDA device count, or 0 if torch/CUDA is unavailable. Never raises.""" + try: + import torch + except Exception: + return 0 + try: + return torch.cuda.device_count() if torch.cuda.is_available() else 0 + except Exception: + return 0 + + +def require_cuda(): + """Skip unless at least one CUDA device is visible.""" + return pytest.mark.skipif( + _cuda_device_count() < 1, reason="no CUDA device available" + ) + + +def require_gpus(n: int): + """Skip unless at least ``n`` CUDA devices are visible.""" + have = _cuda_device_count() + return pytest.mark.skipif( + have < n, reason=f"needs >= {n} CUDA device(s), found {have}" + ) + + +def require_clickhouse(host: str | None = None, port: int | None = None): + """Skip unless a ClickHouse TCP port is reachable. + + Host/port default to the ``DMX_DB_HOST`` / ``DMX_DB_PORT`` env vars (and + finally ``127.0.0.1:9000``), matching the runners' connection defaults. + """ + host = host or os.environ.get("DMX_DB_HOST", "127.0.0.1") + port = int(port if port is not None else os.environ.get("DMX_DB_PORT", "9000")) + reachable = False + try: + with socket.create_connection((host, port), timeout=1.0): + reachable = True + except OSError: + reachable = False + return pytest.mark.skipif( + not reachable, reason=f"ClickHouse unreachable at {host}:{port}" + ) + + +def require_vllm(): + """Skip unless the vLLM runtime is importable.""" + available = importlib.util.find_spec("vllm") is not None + return pytest.mark.skipif(not available, reason="vLLM not importable") + + +def _model_in_cache(model: str) -> bool: + """Best-effort check that ``model`` is a local path or a cached HF repo.""" + # Explicit local path (a checkpoint dir). + if os.path.sep in model and os.path.exists(model): + return True + # HuggingFace hub cache layout: ``models----``. + cache_root = os.environ.get( + "HF_HUB_CACHE", + os.path.join( + os.environ.get( + "HF_HOME", os.path.expanduser("~/.cache/huggingface") + ), + "hub", + ), + ) + repo_dir = "models--" + model.replace("/", "--") + return os.path.isdir(os.path.join(cache_root, repo_dir)) + + +def require_model_cache(model: str): + """Skip unless ``model`` (HF repo id or local path) is already on disk. + + Keeps the default/offline suites from triggering a network download. + """ + return pytest.mark.skipif( + not _model_in_cache(model), + reason=f"model {model!r} not found in local cache", + ) + + +def require_nvcc(): + """Skip unless the CUDA compiler ``nvcc`` is on PATH (native ring tests).""" + return pytest.mark.skipif( + shutil.which("nvcc") is None, reason="nvcc not found on PATH" + ) diff --git a/tests/e2e_matrix.py b/tests/e2e_matrix.py new file mode 100644 index 000000000..6098abc92 --- /dev/null +++ b/tests/e2e_matrix.py @@ -0,0 +1,398 @@ +"""Configurable E2E matrix harness (plan §8). + +One matrix-driven entry point replacing the hardcoded shell sweeps. Each +axis is a comma-separated multi-value flag; the harness takes the Cartesian +product, runs every cell as subprocesses (reusing the existing runners and +comparators -- no inference logic is reimplemented), and writes one JSON +record per cell. + + python -m tests.e2e_matrix \ + --backend hf,vllm --model gpt2,qwen3 \ + --mode eager,cuda_graph --standard transport_bitwise \ + --hooks vllm-full --tp 1 --out results/e2e.jsonl + +Cell dispatch +------------- +- ``vllm`` + ``bitwise`` / ``transport_bitwise`` + vllm_ref_runner (RefDiskWorker, D2D->disk) + vllm_monitored_runner + (ring->ClickHouse) -> vllm_identical_comparator. +- ``vllm`` + ``row_count`` / ``allclose`` + vllm_monitored_runner -> vllm_rowcnt_comparator. +- ``hf`` (any standard) + hf_reference_runner + hf_monitored_runner -> hf_comparator. + +The public ``E2E_HOOK_SELECTION`` input is translated to the internal +``DMX_HOOK_SELECTION`` runtime contract in each subprocess env (plan §2). + +``--dry-run`` prints the planned cells + dispatch commands without touching +CUDA / ClickHouse, so the expansion and env translation are unit-testable on +a CPU-only box. +""" +from __future__ import annotations + +import argparse +import itertools +import json +import os +import shutil +import subprocess +import sys +import tempfile +from dataclasses import dataclass, field +from typing import List, Optional + +from tests.lib.report import CellResult, checks_from_legacy_result, write_jsonl, human_table + +PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + +# model key -> vLLM ref model source filename (under model_executor/models). +_VLLM_REF_FILES = { + "gpt2": "gpt2_ref.py", + "qwen2_moe": "qwen2_moe_ref.py", + "qwen3": "qwen3_ref.py", + "llama": "llama_ref.py", +} + +# Standards that compare reference D2D buffers against ring/ClickHouse output. +_VLLM_IDENTICAL_STANDARDS = {"bitwise", "transport_bitwise"} + + +@dataclass(frozen=True) +class Cell: + """One point in the matrix.""" + + backend: str + model: str + mode: str + standard: str + hooks: str + tp: int = 1 + ring_mb: int = 4096 + dtype: str = "bfloat16" + prompt_set: str = "smoke" + + @property + def enforce_eager(self) -> str: + return "1" if self.mode == "eager" else "0" + + +# --------------------------------------------------------------------------- +# Axis expansion +# --------------------------------------------------------------------------- + + +def _split(s: str) -> List[str]: + return [x.strip() for x in s.split(",") if x.strip()] + + +def build_cells(args: argparse.Namespace) -> List[Cell]: + """Cartesian product over every multi-value axis.""" + cells: List[Cell] = [] + for backend, model, mode, standard, hooks, tp, ring_mb, dtype, pset in itertools.product( + _split(args.backend), _split(args.model), _split(args.mode), + _split(args.standard), _split(args.hooks), _split(str(args.tp)), + _split(str(args.ring_mb)), _split(args.dtype), _split(args.prompt_set), + ): + cells.append(Cell( + backend=backend, model=model, mode=mode, standard=standard, + hooks=hooks, tp=int(tp), ring_mb=int(ring_mb), dtype=dtype, + prompt_set=pset, + )) + return cells + + +def cell_env(cell: Cell, args: argparse.Namespace, base: Optional[dict] = None) -> dict: + """Build the subprocess env for a cell. + + Sets the public ``E2E_*`` knobs *and* the translated internal + ``DMX_HOOK_SELECTION`` (plan §2) so the runners see one consistent + configuration. + """ + env = dict(base if base is not None else os.environ) + env["E2E_MODEL"] = cell.model + env["E2E_ENFORCE_EAGER"] = cell.enforce_eager + env["E2E_CUDA_GRAPHS"] = "0" if cell.mode == "eager" else "1" + env["E2E_DTYPE"] = cell.dtype + env["E2E_TP_SIZE"] = str(cell.tp) + env["E2E_RING_PAYLOAD_MB"] = str(cell.ring_mb) + env["E2E_RING_PINNED_MB"] = str(cell.ring_mb) + env["E2E_PROMPT_SET"] = cell.prompt_set + # Public hook-selection input + internal runtime contract translation. + env["E2E_HOOK_SELECTION"] = cell.hooks + env["DMX_HOOK_SELECTION"] = cell.hooks + env["E2E_NUM_PROMPTS"] = str(args.num_prompts) + env["E2E_MAX_NEW_TOKENS"] = str(args.max_new_tokens) + env["E2E_MAX_MODEL_LEN"] = str(args.max_model_len) + env["E2E_MAX_NUM_BATCHED_TOKENS"] = str(args.max_batched_tokens) + env["E2E_GPU_MEM_UTIL"] = str(args.gpu_mem_util) + env["E2E_TOLERANCE"] = str(args.tolerance) + env["DMX_DB_HOST"] = args.db_host + env["DMX_DB_PORT"] = str(args.db_port) + env["VLLM_DISABLE_COMPILE_CACHE"] = "1" + return env + + +# --------------------------------------------------------------------------- +# Dispatch planning (no side effects -- the dry-run surface) +# --------------------------------------------------------------------------- + + +@dataclass +class Step: + """One planned subprocess: a label + argv (env applied at run time).""" + + label: str + argv: List[str] + + +def _runner(mod: str, *flags: str) -> List[str]: + return [sys.executable, "-m", mod, *flags] + + +def plan_cell(cell: Cell, run_dir: str) -> tuple[List[Step], str, str]: + """Return (steps, comparator_module, result_file) for a cell. + + Pure planning: builds the subprocess argv list without executing, so the + same code path feeds both ``--dry-run`` and the real runner. + """ + ref_dir = os.path.join(run_dir, "ref") + mon_dir = os.path.join(run_dir, "mon") + result_file = os.path.join(run_dir, "result.json") + steps: List[Step] = [] + + if cell.backend == "vllm": + if cell.standard in _VLLM_IDENTICAL_STANDARDS: + config_file = os.path.join(ref_dir, "ref_config.json") + steps.append(Step("enable_ref_hooks", ["", "enable_ref_hooks"])) + steps.append(Step("vllm_ref", _runner("tests.vllm_ref_runner", "--output-dir", ref_dir))) + steps.append(Step("vllm_monitored", _runner("tests.vllm_monitored_runner", "--output-dir", mon_dir))) + steps.append(Step("compare", _runner( + "tests.vllm_identical_comparator", + "--ref-config", config_file, "--mon-dir", mon_dir, + "--result-file", result_file))) + return steps, "tests.vllm_identical_comparator", result_file + # row_count / allclose -> monitored-only + rowcnt comparator + steps.append(Step("vllm_monitored", _runner("tests.vllm_monitored_runner", "--output-dir", mon_dir))) + steps.append(Step("compare", _runner( + "tests.vllm_rowcnt_comparator", + "--ref-dir", ref_dir, "--mon-dir", mon_dir, + "--result-file", result_file))) + return steps, "tests.vllm_rowcnt_comparator", result_file + + if cell.backend == "hf": + steps.append(Step("hf_ref", _runner("tests.hf_reference_runner", "--output-dir", ref_dir))) + steps.append(Step("hf_monitored", _runner("tests.hf_monitored_runner", "--output-dir", mon_dir))) + steps.append(Step("compare", _runner( + "tests.hf_comparator", + "--ref-dir", ref_dir, "--mon-dir", mon_dir, + "--result-file", result_file))) + return steps, "tests.hf_comparator", result_file + + raise ValueError(f"unknown backend {cell.backend!r}") + + +# --------------------------------------------------------------------------- +# Execution +# --------------------------------------------------------------------------- + + +def _enable_vllm_ref_hooks(cell: Cell, run_dir: str, env: dict) -> tuple[str, str, str]: + """Run the in-process enable_ref_hooks preprocessor for a vLLM cell. + + Returns (model_file, backup_file, config_file). Mirrors the flow in + test_vllm_identical: back up the ref model source, generate the hooked + ref + config, leaving restore to the caller's finally. + """ + ref_dir = os.path.join(run_dir, "ref") + os.makedirs(ref_dir, exist_ok=True) + models_dir = os.path.join( + PROJECT_ROOT, "integration", "vllm", "vllm", + "model_executor", "models") + ref_filename = _VLLM_REF_FILES.get(cell.model) + if ref_filename is None: + raise ValueError(f"no vLLM ref model registered for {cell.model!r}") + model_file = os.path.join(models_dir, ref_filename) + backup_file = os.path.join(run_dir, f"{ref_filename}.bak") + config_file = os.path.join(ref_dir, "ref_config.json") + max_len = int(os.environ.get("E2E_REF_MAX_LEN", "8192")) + + shutil.copy2(model_file, backup_file) + sys.path.insert(0, models_dir) + from enable_ref_hooks import enable_ref_hooks # type: ignore + enable_ref_hooks( + model_file=model_file, hooks=cell.hooks, max_len=max_len, + output_dir=ref_dir, config_out=config_file, + ) + return model_file, backup_file, config_file + + +def _run_steps(steps: List[Step], env: dict, run_dir: str, + *, restore=None, timeout: Optional[float] = None) -> Optional[str]: + """Execute the runner/comparator steps in order. + + ``restore`` is a zero-arg callback run after the reference step (for the + vLLM identical flow, which restores the ref model source before the + monitored run). ``timeout`` bounds each subprocess so a hung runner + fails the cell instead of hanging the whole matrix. Returns an error + string on first failure, else None. + """ + for step in steps: + if step.argv and step.argv[0] == "": + continue # enable_ref_hooks handled by the caller + # vLLM identical: restore the patched ref source after the ref run, + # before the monitored run starts. + if restore is not None and step.label in ("vllm_monitored", "hf_monitored"): + restore() + restore = None + env_step = dict(env) + if step.label in ("vllm_ref",): + env_step["REF_CONFIG"] = os.path.join(run_dir, "ref", "ref_config.json") + try: + proc = subprocess.run( + step.argv, env=env_step, cwd=PROJECT_ROOT, + capture_output=True, text=True, timeout=timeout) + except subprocess.TimeoutExpired: + return f"step {step.label} timed out after {timeout}s" + if proc.returncode != 0: + tail = (proc.stderr or "")[-2000:] + return f"step {step.label} failed (rc={proc.returncode}): {tail}" + return None + + +def run_cell(cell: Cell, args: argparse.Namespace) -> CellResult: + """Run one cell end-to-end and return its :class:`CellResult`.""" + cr = CellResult( + backend=cell.backend, model=cell.model, mode=cell.mode, + standard=cell.standard, hook_selection=cell.hooks, tp=cell.tp, + extra={"ring_mb": cell.ring_mb, "dtype": cell.dtype, + "prompt_set": cell.prompt_set}, + ) + run_dir = tempfile.mkdtemp(prefix="e2e_matrix_") + os.makedirs(os.path.join(run_dir, "ref"), exist_ok=True) + os.makedirs(os.path.join(run_dir, "mon"), exist_ok=True) + env = cell_env(cell, args) + backup_file = model_file = None + try: + steps, _comparator, result_file = plan_cell(cell, run_dir) + restore = None + + if cell.backend == "vllm" and cell.standard in _VLLM_IDENTICAL_STANDARDS: + model_file, backup_file, _config = _enable_vllm_ref_hooks(cell, run_dir, env) + + def restore(): # noqa: E306 -- restore ref source pre-monitored run + shutil.copy2(backup_file, model_file) + + elif cell.backend == "vllm": + # rowcnt comparator still expects a ref meta.json (skipped marker). + with open(os.path.join(run_dir, "ref", "meta.json"), "w") as f: + json.dump({"skipped": True}, f) + + err = _run_steps(steps, env, run_dir, restore=restore, + timeout=args.cell_timeout) + if err is not None: + cr.error = err + return cr.finalize() + + with open(result_file) as f: + legacy = json.load(f) + cr.checks = checks_from_legacy_result(legacy) + return cr.finalize() + + except Exception as exc: # noqa: BLE001 -- one bad cell must not abort the matrix + cr.error = f"{type(exc).__name__}: {exc}" + return cr.finalize() + finally: + if backup_file and model_file and os.path.exists(backup_file): + shutil.copy2(backup_file, model_file) + if not args.keep_artifacts: + shutil.rmtree(run_dir, ignore_errors=True) + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def build_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser( + prog="tests.e2e_matrix", + description="Configurable E2E matrix over backend/model/mode/standard/hooks/tp.") + p.add_argument("--backend", default="vllm", help="comma list: hf,vllm") + p.add_argument("--model", default="gpt2", help="comma list: gpt2,qwen3,llama,qwen2_moe,") + p.add_argument("--mode", default="eager", help="comma list: eager,cuda_graph") + p.add_argument("--standard", default="transport_bitwise", + help="comma list: bitwise,allclose,row_count,transport_bitwise") + p.add_argument("--hooks", default="vllm-full", + help="comma list: preset (vllm-full,hidden-states) or single hook") + p.add_argument("--tp", default="1", help="comma list of tensor-parallel sizes") + p.add_argument("--ring-mb", dest="ring_mb", default="4096", + help="comma list of ring payload/pinned sizes in MB") + p.add_argument("--dtype", default="bfloat16", help="comma list: bfloat16,float16,float32") + p.add_argument("--prompt-set", dest="prompt_set", default="smoke", + help="comma list: smoke,math,chat,random") + p.add_argument("--num-prompts", type=int, default=8) + p.add_argument("--max-new-tokens", type=int, default=20) + p.add_argument("--max-model-len", type=int, default=512) + p.add_argument("--max-batched-tokens", type=int, default=512) + p.add_argument("--gpu-mem-util", type=float, default=0.5) + p.add_argument("--tolerance", type=float, default=0.01, + help="abs tolerance forwarded to comparators (E2E_TOLERANCE)") + p.add_argument("--db-host", default="localhost") + p.add_argument("--db-port", type=int, default=9000) + p.add_argument("--out", default=None, help="JSONL output path (one record per cell)") + p.add_argument("--cell-timeout", type=float, default=1800.0, + help="per-subprocess timeout in seconds (hung runner fails the cell)") + p.add_argument("--keep-artifacts", action="store_true", + help="keep per-cell temp run dirs") + p.add_argument("--dry-run", action="store_true", + help="print planned cells + dispatch commands; no CUDA/ClickHouse") + return p + + +def _dry_run(cells: List[Cell], args: argparse.Namespace) -> int: + print(f"# {len(cells)} cell(s) planned\n") + for i, cell in enumerate(cells): + steps, comparator, _result = plan_cell(cell, run_dir="") + env = cell_env(cell, args, base={}) + print(f"[{i}] backend={cell.backend} model={cell.model} mode={cell.mode} " + f"standard={cell.standard} hooks={cell.hooks} tp={cell.tp} " + f"ring_mb={cell.ring_mb} dtype={cell.dtype} prompt_set={cell.prompt_set}") + print(f" env: E2E_HOOK_SELECTION={env['E2E_HOOK_SELECTION']} -> " + f"DMX_HOOK_SELECTION={env['DMX_HOOK_SELECTION']} " + f"E2E_ENFORCE_EAGER={env['E2E_ENFORCE_EAGER']}") + for step in steps: + print(f" - {step.label}: {' '.join(step.argv)}") + print(f" comparator: {comparator}\n") + return 0 + + +def main(argv: Optional[List[str]] = None) -> int: + args = build_parser().parse_args(argv) + cells = build_cells(args) + if not cells: + print("no cells to run (check axis values)", file=sys.stderr) + return 2 + + if args.dry_run: + return _dry_run(cells, args) + + results: List[CellResult] = [] + for i, cell in enumerate(cells): + print(f"\n=== cell {i + 1}/{len(cells)}: {cell.backend}/{cell.model}/{cell.mode}/" + f"{cell.standard}/{cell.hooks}/tp{cell.tp} ===", flush=True) + cr = run_cell(cell, args) + verdict = "ERROR" if cr.error else ("PASS" if cr.passed else "FAIL") + print(f" -> {verdict}" + (f": {cr.error}" if cr.error else ""), flush=True) + results.append(cr) + + print("\n" + human_table(results)) + if args.out: + write_jsonl(results, args.out) + print(f"\nwrote {len(results)} record(s) to {args.out}") + + # Exit non-zero if any cell failed or errored, so CI can gate on it. + return 0 if all(r.passed for r in results) else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/isolate_hook.py b/tests/isolate_hook.py index fca14337a..3a4868f28 100644 --- a/tests/isolate_hook.py +++ b/tests/isolate_hook.py @@ -13,8 +13,11 @@ forward pass also writes to other buffers, which keeps the un-hooked tensors live and could mask hook-perturbation regressions. -This module owns the source-level patching. It is callable both as a -library (``patch_compare_model``) and as a CLI for ad-hoc isolation runs. +This module owns the source-level patching. The canonical entry point is +the ``isolated_hook`` context manager, which snapshots the file bytes, +patches, and on exit restores *and asserts byte-identical restoration* so +the vendored submodule is never left dirty. It is also callable as a CLI +for ad-hoc isolation runs. The patch is applied **in-place** to the source file, with the original saved to a sibling ``.copy_isolate_backup`` so a Ctrl-C / crash doesn't @@ -30,6 +33,7 @@ import argparse import contextlib +import hashlib import os import re import shutil @@ -173,24 +177,64 @@ def unpatch(framework: str, model_key: str) -> Path: @contextlib.contextmanager -def patch_compare_model( - framework: str, model_key: str, hook: str +def isolated_hook( + framework: str, model_key: str, hook: str, *, invalidate: bool = True, ) -> Iterator[tuple[Path, list[str]]]: - """Context-manager form of patch / unpatch. - - Yields ``(model_path, commented_bufs)``. Always restores on exit, - even on exception, even on Ctrl-C (best effort -- the backup file - persists across crashes). + """Single hardened isolation contract (plan §6). + + Snapshots the vendored ``_compare`` source bytes, patches the file so + only ``hook``'s ``.copy_()`` line fires, yields + ``(model_path, commented_bufs)``, and on exit -- always, even on + exception or Ctrl-C -- restores from backup and **asserts the file is + byte-identical to the snapshot** via a SHA-256 compare. A non-identical + restore raises ``RuntimeError`` loudly so a dirty vendored submodule can + never escape unnoticed. + + With ``invalidate=True`` (default) the module's cached bytecode is + dropped after patching and after restoring, so a subprocess import + always picks up the on-disk source rather than a stale ``.pyc``. The + numeric-difference study (plan §9) consumes this context manager rather + than the raw ``patch`` / ``unpatch`` functions. """ - p, commented = patch(framework, model_key, hook) + p = compare_model_path(framework, model_key) + original_digest = hashlib.sha256(p.read_bytes()).hexdigest() + _, commented = patch(framework, model_key, hook) + if invalidate: + invalidate_bytecode(framework, model_key) try: yield p, commented finally: try: unpatch(framework, model_key) except FileNotFoundError: - # Already unpatched (e.g. if the body called unpatch directly). + # Already unpatched (e.g. if the body called unpatch directly); + # fall through to the hash check, which confirms the on-disk + # source matches the snapshot regardless. pass + if invalidate: + invalidate_bytecode(framework, model_key) + restored_digest = hashlib.sha256(p.read_bytes()).hexdigest() + if restored_digest != original_digest: + raise RuntimeError( + f"isolated_hook failed to restore {p} byte-identically " + f"(framework={framework!r}, model_key={model_key!r}, " + f"hook={hook!r}); the vendored submodule may be left dirty. " + f"Restore it manually, e.g. `git -C {REPO_ROOT} checkout -- {p}`." + ) + + +@contextlib.contextmanager +def patch_compare_model( + framework: str, model_key: str, hook: str +) -> Iterator[tuple[Path, list[str]]]: + """Backward-compatible alias for :func:`isolated_hook`. + + Kept for existing callers. Delegates to the hardened contract but + leaves bytecode invalidation to the caller (historical behavior). + Prefer :func:`isolated_hook` for new code. + """ + with isolated_hook(framework, model_key, hook, invalidate=False) as y: + yield y def _bytecode_paths_for(p: Path) -> list[Path]: diff --git a/tests/lib/__init__.py b/tests/lib/__init__.py new file mode 100644 index 000000000..bda07b33a --- /dev/null +++ b/tests/lib/__init__.py @@ -0,0 +1,25 @@ +"""Shared E2E test library (plan §7). + +Consolidates the read / merge / align / compare / report logic that the +configurable matrix (:mod:`tests.e2e_matrix`), the pytest wrappers, and the +numeric-difference study all build on, so each rule lives in exactly one +place. + +Submodules: + align -- left-pad strip, EOS trim, request_id ":" parsing + compare -- Check + bitwise / allclose / row_count / transport_bitwise + report -- Check / CellResult dataclasses -> JSON(L) + human table + clickhouse_io -- read offload rows by request_id, dtype/hook maps, row counts + segments -- merge chunked segments -> dense tensors (segment_merger) + disk_ref -- load .pt / structured reference tensors written by ref workers + hf_reference -- ROL + GEN HF rollouts (re-export of tests.hf_reference) + +Only ``align``, ``compare``, and ``report`` are imported eagerly here; they +are pure-CPU (torch-only). The IO-heavy submodules are imported on demand +to keep ``import tests.lib`` cheap and offline-friendly. +""" +from __future__ import annotations + +from tests.lib import align, compare, report # noqa: F401 + +__all__ = ["align", "compare", "report"] diff --git a/tests/lib/align.py b/tests/lib/align.py new file mode 100644 index 000000000..ca31eca17 --- /dev/null +++ b/tests/lib/align.py @@ -0,0 +1,88 @@ +"""Alignment helpers shared by the comparators / matrix (plan §7). + +Consolidates the left-pad strip, EOS trim, and ``request_id`` parsing that +several comparators reimplement. Pure-CPU, ``torch``-only, unit-tested +without CUDA. +""" +from __future__ import annotations + +import re +from typing import Optional, Tuple + +import torch + +# request_id canonical form is ":" -- the group is the +# batched generate() call, the row the position within that batch. +_REQUEST_ID_RE = re.compile(r"^(\d+):(\d+)$") + +# vLLM appends a "-<8 hex>" UUID suffix to request ids; the ref/disk workers +# strip it so monitored and reference rows key the same. +_VLLM_SUFFIX_RE = re.compile(r"-[0-9a-f]{8}$") + + +def parse_request_id(req_id: str) -> Tuple[int, int]: + """Parse ``":"`` into ``(group_id, row_index)``. + + Raises ``ValueError`` on an unexpected format so a malformed id surfaces + loudly rather than silently sorting wrong. + """ + m = _REQUEST_ID_RE.match(req_id) + if not m: + raise ValueError(f"unexpected request_id format: {req_id!r}") + return int(m.group(1)), int(m.group(2)) + + +def normalize_request_id(req_id: str) -> str: + """Strip a trailing vLLM ``-<8hex>`` UUID suffix, if present.""" + return _VLLM_SUFFIX_RE.sub("", req_id) + + +def strip_left_pad(ids_row: torch.Tensor, attn_row: torch.Tensor) -> torch.Tensor: + """Drop left-padding from a single sequence using its attention mask. + + Returns the last ``attn_row.sum()`` ids (HF left-pads, so the real + tokens are the trailing run). Empty (all-pad) rows return an empty + slice. + """ + true_len = int(attn_row.sum().item()) + if true_len <= 0: + return ids_row[:0] + return ids_row[-true_len:] + + +def trim_eos(ids: torch.Tensor, eos_id: int, + *, keep_eos: bool = False) -> torch.Tensor: + """Trim a 1-D id sequence at the first EOS token. + + With ``keep_eos=False`` (default) the EOS itself is dropped; with + ``keep_eos=True`` it is retained. If no EOS is present the sequence is + returned unchanged. + """ + flat = ids.reshape(-1) + hits = torch.nonzero(flat == eos_id, as_tuple=False) + if hits.numel() == 0: + return flat + first = int(hits[0].item()) + return flat[: first + 1] if keep_eos else flat[:first] + + +def align_to_min_len(a: torch.Tensor, b: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Trim two tensors to a common length along dim 0. + + Used before a value comparison when reference and monitored captures + cover slightly different token spans (e.g. generate() drops the final + never-forwarded token). + """ + n = min(a.shape[0], b.shape[0]) + return a[:n], b[:n] + + +def logits_align_skip(db_len: int, ref_len: int) -> int: + """Rows to skip at the head of a DB ``final_logits`` block to align to ref. + + The DB (``logits_to_keep=0``) stores every position + ``[prompt_0..prompt_{N-1}, decode_0..decode_{G-1}]``; ``generate()``'s + ``output_logits`` yields ``[prefill_last, decode_0..decode_{G-2}]``. + The prefill-last DB row sits at ``prompt_len - 1 = db_len - ref_len - 1``. + """ + return max(0, db_len - ref_len - 1) diff --git a/tests/lib/clickhouse_io.py b/tests/lib/clickhouse_io.py new file mode 100644 index 000000000..a3151409e --- /dev/null +++ b/tests/lib/clickhouse_io.py @@ -0,0 +1,124 @@ +"""ClickHouse read helpers for the comparators / matrix (plan §7). + +Consolidates the row-decode logic, dtype table, and short-hook -> CH +``act_name`` map that ``vllm_identical_comparator``, ``compare_disk_vs_ch``, +and ``vllm_rowcnt_comparator`` each reimplement. + +``clickhouse_driver`` is imported lazily inside the functions so importing +this module stays CPU/offline-friendly (the unit suite never reaches a DB). +""" +from __future__ import annotations + +from typing import Dict, List, Tuple + +import torch + +# CH stores the torch dtype as its ``str(dtype)``; map back on read. +DTYPE_MAP: Dict[str, torch.dtype] = { + "torch.bfloat16": torch.bfloat16, "torch.float": torch.float32, + "torch.float32": torch.float32, "torch.half": torch.float16, + "torch.float16": torch.float16, "torch.int": torch.int32, + "torch.int32": torch.int32, "torch.long": torch.int64, + "torch.int64": torch.int64, "torch.uint8": torch.uint8, + "torch.int8": torch.int8, "torch.short": torch.int16, + "torch.double": torch.float64, "torch.bool": torch.bool, +} + +# Short hook name (the ``_buf_`` suffix / disk filename stem) -> the +# ClickHouse ``act_name``. Must match tensor_meta.h hook_type_name() and the +# p2p make_act_name() convention. +HOOK_TO_CH_ACT: Dict[str, str] = { + "resid_pre": "blocks.hook_resid_pre", + "ln1": "blocks.hook_ln1", + "q": "blocks.attn.hook_q", + "k": "blocks.attn.hook_k", + "v": "blocks.attn.hook_v", + "z": "blocks.attn.hook_z", + "attn_scores": "blocks.attn.hook_attn_scores", + "pattern": "blocks.attn.hook_pattern", + "attn_out": "blocks.hook_attn_out", + "resid_mid": "blocks.hook_resid_mid", + "ln2": "blocks.hook_ln2", + "mlp_in": "blocks.hook_mlp_in", + "mlp_out": "blocks.hook_mlp_out", + "mlp_post": "blocks.hook_mlp_post", + "embed": "hook_embed", + "pos_embed": "hook_pos_embed", + "resid_final": "hook_resid_final", + "final_ln": "hook_final_ln", + "final_logits": "final_logits", + "token_ids": "token_ids", + "router_logits": "blocks.mlp.hook_router_logits", + "topk_ids": "blocks.mlp.hook_topk_ids", + "topk_weights": "blocks.mlp.hook_topk_weights", +} + +# A CH row key: (req_id, act_name, layer_no, shard_rank, start_token, end_token). +RowKey = Tuple[str, str, int, int, int, int] + + +def _decode(v) -> str: + return v.decode() if isinstance(v, bytes) else v + + +def read_offload_rows( + db_host: str, db_port: int, *, + database: str = "default", table: str = "offload", +) -> Tuple[Dict[RowKey, torch.Tensor], int]: + """Read every row from ``.`` into a keyed dict. + + Returns ``(rows_by_key, num_rows)`` where the key is :data:`RowKey` and + the value the decoded CPU tensor. Raises whatever ``clickhouse_driver`` + raises on a connection / query error -- callers decide whether that is a + soft "db unreachable" skip or a hard failure. + """ + import clickhouse_driver + + client = clickhouse_driver.Client(db_host, port=db_port) + raw_rows = client.execute( + "SELECT model_id, request_id, act_name, layer_no, shard_rank, " + "start_token_idx, end_token_idx, dtype, shape, bytes " + f"FROM {database}.{table}", + settings={"strings_as_bytes": True}, + ) + + out: Dict[RowKey, torch.Tensor] = {} + for row in raw_rows: + _, req_id, act_name, layer_no, shard_rank, s, e, dtype_str, shape, payload = row + dt = DTYPE_MAP.get(_decode(dtype_str), torch.float32) + t = torch.frombuffer(bytearray(payload), dtype=dt).reshape(list(shape)) + out[(_decode(req_id), _decode(act_name), int(layer_no), + int(shard_rank), int(s), int(e))] = t + return out, len(raw_rows) + + +def per_hook_counts(rows_by_key: Dict[RowKey, torch.Tensor]) -> Dict[str, int]: + """Count rows per ``act_name`` (input to the ``row_count`` standard).""" + counts: Dict[str, int] = {} + for key in rows_by_key: + act = key[1] + counts[act] = counts.get(act, 0) + 1 + return counts + + +def group_by_request( + rows_by_key: Dict[RowKey, torch.Tensor], +) -> Dict[str, Dict[Tuple[int, str], List[Tuple[int, int, torch.Tensor]]]]: + """Regroup CH rows as ``req_id -> (layer_no, act_name) -> [(s, e, t)]``. + + ``act_name`` is canonicalised by stripping a leading ``"blocks."`` so a + per-layer hook keys as ``(layer_no, "hook_resid_pre")`` and a global hook + as ``(-1, "final_logits")``. Segments are left unsorted; merge callers + sort by start token. + """ + grouped: Dict[str, Dict[Tuple[int, str], List[Tuple[int, int, torch.Tensor]]]] = {} + for (req_id, act_name, layer_no, _shard, s, e), t in rows_by_key.items(): + if act_name.startswith("blocks."): + canon = act_name[len("blocks."):] + lno = layer_no + else: + canon = act_name + lno = -1 + grouped.setdefault(req_id, {}).setdefault((lno, canon), []).append( + (s, e, t.detach().cpu())) + return grouped diff --git a/tests/lib/compare.py b/tests/lib/compare.py new file mode 100644 index 000000000..8ee02a58c --- /dev/null +++ b/tests/lib/compare.py @@ -0,0 +1,211 @@ +"""Comparison standards for the E2E matrix and pytest wrappers (plan §7). + +Four standards, one common interface. Each returns a :class:`Check` +recording pass/fail plus the numeric drift (``max_abs`` / ``mean_abs`` / +``first_diff_pos``) **even on a pass**, so a green cell still surfaces a +"barely passing" trend: + +- ``bitwise`` -- exact equality (raw bytes / ``torch.equal``). +- ``allclose`` -- ``torch.allclose(atol, rtol)`` with a named, + reported threshold. +- ``row_count`` -- schema + segment-count validation only. +- ``transport_bitwise`` -- ``.copy_()`` reference buffers vs ClickHouse ring + output; exact, same engine as ``bitwise`` but a + distinct name so the gating policy (§8) can treat + transport separately from model-output transparency. + +This module is pure-CPU and only depends on ``torch``; it is unit-tested +without CUDA / ClickHouse / vLLM. +""" +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Callable, Optional + +import torch + + +@dataclass +class Check: + """One comparison result. + + ``max_abs`` / ``mean_abs`` / ``first_diff_pos`` are recorded whenever + they can be computed (even on pass) so trends stay visible. ``detail`` + is a short human string for tables / assertion messages. + """ + + name: str + passed: bool + max_abs: Optional[float] = None + mean_abs: Optional[float] = None + first_diff_pos: Optional[int] = None + detail: str = "" + + def to_dict(self) -> dict: + d: dict = {"name": self.name, "passed": self.passed} + if self.max_abs is not None: + d["max_abs"] = self.max_abs + if self.mean_abs is not None: + d["mean_abs"] = self.mean_abs + if self.first_diff_pos is not None: + d["first_diff_pos"] = self.first_diff_pos + if self.detail: + d["detail"] = self.detail + return d + + +# --------------------------------------------------------------------------- +# Low-level helpers +# --------------------------------------------------------------------------- + + +_INT_VIEW = {1: torch.uint8, 2: torch.int16, 4: torch.int32, 8: torch.int64} + + +def bytes_identical(a: torch.Tensor, b: torch.Tensor) -> bool: + """True iff ``a`` and ``b`` have identical raw bytes. + + Reinterprets as a same-width integer dtype so the compare is a true + bitwise check (no NaN / signed-zero surprises). Falls back to a + storage-bytes compare for exotic element sizes. + """ + a_c = a.contiguous() + b_c = b.contiguous() + if a_c.shape != b_c.shape or a_c.element_size() != b_c.element_size(): + return False + dt = _INT_VIEW.get(a_c.element_size()) + if dt is None: + return bytes(a_c.untyped_storage()) == bytes(b_c.untyped_storage()) + return torch.equal(a_c.view(dt), b_c.view(dt)) + + +def _abs_diff_stats(a: torch.Tensor, b: torch.Tensor) -> tuple[float, float, Optional[int]]: + """Return (max_abs, mean_abs, first_diff_pos) over the flattened tensors. + + ``first_diff_pos`` is the index of the first element that differs in the + flattened view, or ``None`` if the tensors are elementwise equal. + """ + af = a.detach().float().reshape(-1) + bf = b.detach().float().reshape(-1) + n = min(af.numel(), bf.numel()) + af = af[:n] + bf = bf[:n] + diff = (af - bf).abs() + max_abs = float(diff.max().item()) if n else 0.0 + mean_abs = float(diff.mean().item()) if n else 0.0 + ne = torch.nonzero(af != bf, as_tuple=False) + first = int(ne[0].item()) if ne.numel() else None + return max_abs, mean_abs, first + + +def _precheck(a: torch.Tensor, b: torch.Tensor, name: str) -> Optional[Check]: + """Shape/dtype gate shared by the tensor standards. + + Returns a failing :class:`Check` on mismatch, else ``None``. + """ + if a.shape != b.shape: + return Check(name, False, detail=f"shape mismatch: {list(a.shape)} vs {list(b.shape)}") + if a.dtype != b.dtype: + return Check(name, False, detail=f"dtype mismatch: {a.dtype} vs {b.dtype}") + return None + + +# --------------------------------------------------------------------------- +# The four standards +# --------------------------------------------------------------------------- + + +def bitwise(a: torch.Tensor, b: torch.Tensor, name: str = "bitwise") -> Check: + """Exact equality. Records drift stats when it fails.""" + pre = _precheck(a, b, name) + if pre is not None: + return pre + if bytes_identical(a, b): + return Check(name, True, max_abs=0.0, mean_abs=0.0, detail="bitwise equal") + max_abs, mean_abs, first = _abs_diff_stats(a, b) + return Check(name, False, max_abs=max_abs, mean_abs=mean_abs, + first_diff_pos=first, detail=f"max_abs={max_abs:.6e}") + + +def transport_bitwise(a: torch.Tensor, b: torch.Tensor, + name: str = "transport_bitwise") -> Check: + """``.copy_()`` reference buffer vs ring/ClickHouse output -- exact. + + Identical engine to :func:`bitwise`; a distinct standard name so the + §8 gating policy can keep transport bitwise even when model-output + transparency is allowed to use ``allclose`` under CUDA graphs. + """ + return bitwise(a, b, name) + + +def allclose(a: torch.Tensor, b: torch.Tensor, name: str = "allclose", + *, atol: float = 1e-3, rtol: float = 0.0) -> Check: + """``torch.allclose`` with a named, reported threshold. + + Always records max/mean abs diff -- even on a pass -- so the threshold + headroom stays visible. + """ + pre = _precheck(a, b, name) + if pre is not None: + return pre + max_abs, mean_abs, first = _abs_diff_stats(a, b) + passed = bool(torch.allclose(a.float(), b.float(), atol=atol, rtol=rtol)) + return Check( + name, passed, max_abs=max_abs, mean_abs=mean_abs, + first_diff_pos=None if passed else first, + detail=f"max_abs={max_abs:.6e} (atol={atol:g}, rtol={rtol:g})", + ) + + +def row_count(per_hook_counts: dict, name: str = "row_count", *, + min_per_layer_types: int = 10, + require_final_logits: bool = True) -> Check: + """Schema + segment-count validation only (no value comparison). + + ``per_hook_counts`` maps a ClickHouse ``act_name`` to the number of rows + captured for it. Validates that: + + - there are enough per-layer hook types (``blocks.*``), + - every per-layer hook captured the same number of rows, and + - the global ``final_logits`` hook is present (when required). + """ + if not per_hook_counts: + return Check(name, False, detail="no rows") + per_layer = {k: v for k, v in per_hook_counts.items() if k.startswith("blocks.")} + problems: list[str] = [] + if len(per_layer) < min_per_layer_types: + problems.append(f"only {len(per_layer)} per-layer types (<{min_per_layer_types})") + counts = set(per_layer.values()) + if len(counts) > 1: + problems.append(f"uneven per-layer counts: {sorted(counts)}") + if require_final_logits and "final_logits" not in per_hook_counts: + problems.append("final_logits missing") + passed = not problems + detail = "ok" if passed else "; ".join(problems) + return Check(name, passed, detail=detail) + + +# Registry of the tensor-pair standards for matrix dispatch by name. +TENSOR_STANDARDS: dict[str, Callable[..., Check]] = { + "bitwise": bitwise, + "transport_bitwise": transport_bitwise, + "allclose": allclose, +} + +ALL_STANDARDS = tuple(TENSOR_STANDARDS) + ("row_count",) + + +def compare_tensors(a: torch.Tensor, b: torch.Tensor, standard: str, + name: Optional[str] = None, **kwargs) -> Check: + """Dispatch a tensor-pair comparison by standard name. + + ``standard`` must be one of :data:`TENSOR_STANDARDS` (``row_count`` is + not a tensor-pair standard -- call :func:`row_count` directly). + """ + fn = TENSOR_STANDARDS.get(standard) + if fn is None: + raise ValueError( + f"unknown tensor standard {standard!r}; " + f"expected one of {sorted(TENSOR_STANDARDS)}" + ) + return fn(a, b, name or standard, **kwargs) diff --git a/tests/lib/disk_ref.py b/tests/lib/disk_ref.py new file mode 100644 index 000000000..0e988a2b4 --- /dev/null +++ b/tests/lib/disk_ref.py @@ -0,0 +1,84 @@ +"""Load reference tensors written to disk by the ref workers (plan §7). + +Two on-disk reference shapes exist: + +- The vLLM ``RefDiskWorker`` writes per-request ``.pt`` files named + ``{hook}_L{layer}_T{start}_{end}[_SR{rank}].pt`` under + ``//``. :func:`scan_pt_ref_files` parses those. +- The HF reference runner writes a structured dump consumed via + ``tests.hf_reference._load_hf_refs_from_disk``; :func:`load_hf_refs` + re-exports it (lazily, to avoid importing the heavy HF module on CPU). +""" +from __future__ import annotations + +import re +from dataclasses import dataclass +from pathlib import Path +from typing import List, Optional + +# {hook}[_L{layer}]_T{start}_{end}[_SR{rank}].pt +PT_RE = re.compile( + r"^(?P\w+?)(?:_L(?P\d+))?_T(?P\d+)_(?P\d+)" + r"(?:_SR(?P\d+))?\.pt$" +) + + +@dataclass(frozen=True) +class RefFile: + """One parsed reference ``.pt`` file.""" + + req_id: str + hook: str + layer: int # -1 for global hooks + shard: int # 0 when TP == 1 + start: int + end: int + path: str + + @property + def label(self) -> str: + lbl = f"{self.req_id}/{self.hook}" + if self.layer >= 0: + lbl += f"_L{self.layer}" + return lbl + f"_T{self.start}_{self.end}" + + +def parse_pt_name(name: str) -> Optional[dict]: + """Parse a ref ``.pt`` filename into its fields, or ``None`` if unmatched.""" + m = PT_RE.match(name) + if not m: + return None + return { + "hook": m.group("hook"), + "layer": int(m.group("layer")) if m.group("layer") is not None else -1, + "shard": int(m.group("shard")) if m.group("shard") is not None else 0, + "start": int(m.group("start")), + "end": int(m.group("end")), + } + + +def scan_pt_ref_files(ref_dir: str) -> List[RefFile]: + """Walk ``//*.pt`` and return parsed :class:`RefFile`s.""" + root = Path(ref_dir) + out: List[RefFile] = [] + for req_dir in sorted(root.iterdir()): + if not req_dir.is_dir(): + continue + for pt_file in sorted(req_dir.iterdir()): + parsed = parse_pt_name(pt_file.name) + if parsed is None: + continue + out.append(RefFile(req_id=req_dir.name, path=str(pt_file), **parsed)) + return out + + +def load_pt(path: str): + """Load a single reference tensor (CPU, weights-only).""" + import torch + return torch.load(path, weights_only=True, map_location="cpu") + + +def load_hf_refs(ref_dir: str): + """Load the HF structured reference dump (delegates to tests.hf_reference).""" + from tests.hf_reference import _load_hf_refs_from_disk + return _load_hf_refs_from_disk(ref_dir) diff --git a/tests/lib/hf_reference.py b/tests/lib/hf_reference.py new file mode 100644 index 000000000..4bdaf7eb0 --- /dev/null +++ b/tests/lib/hf_reference.py @@ -0,0 +1,28 @@ +"""HF reference rollouts for the matrix / wrappers (plan §7). + +Re-exports the ROL (manual KV-cache rollout: full logits + hidden states + +attn patterns) and GEN (``generate()`` token_ids + decode scores) reference +helpers under stable public names. The canonical implementation still +lives in :mod:`tests.hf_reference`; this shim gives the shared ``tests.lib`` +namespace a single import point without moving the 700-line module (that +relocation is deferred to the legacy-removal PR so this one stays additive). +""" +from __future__ import annotations + +from tests.hf_reference import ( # noqa: F401 (re-exported) + _HFRef as HFRef, + _HFGenRef as HFGenRef, + _hf_greedy_rollout_collect_all_batched as rollout_collect_all, + _hf_generate_collect_scores_batched as generate_collect_scores, + _hf_generate_collect_hidden_states_batched as generate_collect_hidden_states, + _load_hf_refs_from_disk as load_refs_from_disk, +) + +__all__ = [ + "HFRef", + "HFGenRef", + "rollout_collect_all", + "generate_collect_scores", + "generate_collect_hidden_states", + "load_refs_from_disk", +] diff --git a/tests/lib/report.py b/tests/lib/report.py new file mode 100644 index 000000000..7923d9e10 --- /dev/null +++ b/tests/lib/report.py @@ -0,0 +1,126 @@ +"""Machine- and human-readable matrix output (plan §7, §8). + +A :class:`CellResult` is one cell of the E2E matrix -- one +``(backend, model, mode, standard, hook_selection, tp, ...)`` point. It +serialises to a single JSON record (one per line in the JSONL artifact, §8) +and renders into a compact human table. + +Pure-CPU; depends only on the stdlib and :mod:`tests.lib.compare`. +""" +from __future__ import annotations + +import json +from dataclasses import dataclass, field, asdict +from typing import Any, List, Optional + +from tests.lib.compare import Check + + +@dataclass +class CellResult: + """Result of one matrix cell. + + ``checks`` holds the per-tensor / per-hook :class:`Check` objects. + ``passed`` is the cell verdict (all checks passed and no error). + ``error`` carries a setup/dispatch failure message when the cell could + not run at all (distinct from a cell that ran and failed a check). + ``extra`` stashes axis values that don't have a first-class field + (ring sizes, dtype, prompt-set, ...). + """ + + backend: str + model: str + mode: str + standard: str + hook_selection: str + tp: int = 1 + passed: bool = False + checks: List[Check] = field(default_factory=list) + error: Optional[str] = None + extra: dict = field(default_factory=dict) + + def finalize(self) -> "CellResult": + """Set ``passed`` from the checks (no error + every check passed).""" + self.passed = self.error is None and bool(self.checks) and all( + c.passed for c in self.checks) + return self + + def to_record(self) -> dict: + rec: dict[str, Any] = { + "backend": self.backend, + "model": self.model, + "mode": self.mode, + "standard": self.standard, + "hook_selection": self.hook_selection, + "tp": self.tp, + "passed": self.passed, + "checks": [c.to_dict() for c in self.checks], + } + if self.error is not None: + rec["error"] = self.error + if self.extra: + rec["extra"] = self.extra + return rec + + +def checks_from_legacy_result(result: dict) -> List[Check]: + """Adapt a legacy comparator ``result.json`` into :class:`Check` objects. + + The existing comparators emit ``{"tests": [{"name", "passed", "detail"}]}``; + this lets the matrix dispatch to them unchanged and still produce the + new record shape. + """ + out: List[Check] = [] + for t in result.get("tests", []): + out.append(Check( + name=t.get("name", "?"), + passed=bool(t.get("passed", False)), + detail=t.get("detail", ""), + )) + return out + + +def write_jsonl(results: List[CellResult], path: str) -> None: + """Write one JSON record per line (creates parent dirs).""" + import os + os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True) + with open(path, "w") as f: + for r in results: + f.write(json.dumps(r.to_record()) + "\n") + + +def read_jsonl(path: str) -> List[dict]: + """Read a JSONL artifact back into a list of records.""" + out: List[dict] = [] + with open(path) as f: + for line in f: + line = line.strip() + if line: + out.append(json.loads(line)) + return out + + +def human_table(results: List[CellResult]) -> str: + """Render the matrix as a compact fixed-width table.""" + header = ("backend", "model", "mode", "standard", "hooks", "tp", "result", "n_fail") + rows: List[tuple] = [] + for r in results: + n_fail = sum(1 for c in r.checks if not c.passed) + verdict = "ERROR" if r.error is not None else ("PASS" if r.passed else "FAIL") + rows.append(( + r.backend, r.model, r.mode, r.standard, r.hook_selection, + str(r.tp), verdict, str(n_fail), + )) + widths = [len(h) for h in header] + for row in rows: + for i, cell in enumerate(row): + widths[i] = max(widths[i], len(cell)) + fmt = " ".join(f"{{:<{w}}}" for w in widths) + lines = [fmt.format(*header), fmt.format(*("-" * w for w in widths))] + lines += [fmt.format(*row) for row in rows] + n_pass = sum(1 for r in results if r.passed) + n_err = sum(1 for r in results if r.error is not None) + lines.append("") + lines.append(f"{n_pass}/{len(results)} cells passed" + + (f", {n_err} errored" if n_err else "")) + return "\n".join(lines) diff --git a/tests/lib/segments.py b/tests/lib/segments.py new file mode 100644 index 000000000..1fc9e066b --- /dev/null +++ b/tests/lib/segments.py @@ -0,0 +1,33 @@ +"""Chunked-segment merge helpers (plan §7). + +Thin wrapper over :mod:`monitoring.segment_merger` so the matrix, the +pytest wrappers, and the numeric study all merge chunked ring/CH segments +through one entry point. Re-exports the canonical implementation rather +than reimplementing the per-act-name merge rules. +""" +from __future__ import annotations + +from typing import List, Optional, Tuple + +import torch + +from monitoring.segment_merger import ( # noqa: F401 (re-exported) + merge_segments, + segment_manager, + parse_internal_id, + get_delta_token_len, +) + + +def merge_request_chunks( + chunks: List[Tuple[int, int, torch.Tensor]], act_name: str, + *, drop_token_cnt_to: Optional[int] = None, +) -> torch.Tensor: + """Sort ``(start, end, tensor)`` chunks by start token and merge them. + + Convenience over :func:`merge_segments` for the + ``group_by_request`` output shape (lists of ``(s, e, t)`` triples). + """ + ordered = sorted(chunks, key=lambda c: c[0]) + return merge_segments( + [t for _, _, t in ordered], act_name, drop_token_cnt_to=drop_token_cnt_to) diff --git a/tests/numeric_study.py b/tests/numeric_study.py new file mode 100644 index 000000000..c8c11cf09 --- /dev/null +++ b/tests/numeric_study.py @@ -0,0 +1,742 @@ +"""Per-hook numeric-difference study (plan §9 / Phase 5). + +Enables **one hook at a time** and reports the drift its monitoring path +introduces versus the **unhooked** baseline model. The non-goal carried from +the issue holds: this does not *fix* numeric drift, it makes it *visible, +categorized, and reproducible*. + +Algorithm (plan §9): + + 1. Run the baseline **unhooked** model once; capture token ids + full logits. + 2. For each hook ``H`` in the selection, enable **only** ``H`` and run the + monitored model: + - ``--variant p`` (default): the production ``_p`` Hooked variant driven + with ``hook_selection=H`` -- ``hook_selection`` already isolates a + single hook, so no source patching is needed. + - ``--variant compare``: the ``_compare`` variant under the hardened + :func:`tests.isolate_hook.isolated_hook` context manager (plan §6), + which patches the vendored source so only ``H``'s ``.copy_()`` line + fires and asserts byte-identical restoration on exit. + 3. Compare against the baseline and record, per hook: + - bitwise pass/fail (eager) or allclose-within-threshold (cuda graph), + via the shared standards in :mod:`tests.lib.compare`, + - max abs diff, mean abs diff, first differing token position, + - top-k vocab diffs at the first differing position, + - whether greedy (argmax) token ids diverged. + 4. Emit a machine-readable JSON artifact **and** a human-readable table. + +Alert policy (plan §9), expressed through the §8 standards-by-mode choice: + + - **Eager** -> ``bitwise`` standard: *any* non-bitwise logits diff alerts. + - **CUDA graph** -> ``allclose`` standard with a per-model/per-dtype + threshold: a max abs diff over the threshold alerts. + - **Always** alert on greedy token-id divergence (a hook flipped the + argmax), on a runner error, on a shape mismatch (a possible hook identity + swap), or on an empty capture. + +The capture side reuses the proven subprocess-rollout pattern from +``tests/test_per_hook_isolation.py`` (a clean CUDA context / ring-transport +instance per cell), but saves the **full** ``[N, vocab]`` logits so the study +can report top-k vocab drift at the first divergence. + +The comparison / report / alert core (``compute_drift``, ``format_table``, the +``*_to_dict`` helpers) is pure and CPU-testable; ``torch`` is imported lazily +inside the helpers that need it, so this module loads on a torch-less box. + +CLI:: + + python -m tests.numeric_study \\ + --framework hf --model qwen3 --mode eager \\ + --hooks q,k,resid_pre,final_logits \\ + --out results/numeric_qwen3_eager.json +""" +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys +from dataclasses import dataclass, field +from pathlib import Path +from textwrap import dedent +from typing import Any, Dict, List, Optional + +from tests.lib.compare import Check, allclose, bitwise + +REPO_ROOT = Path(__file__).resolve().parent.parent + +# Per-(model, dtype) max-abs-diff thresholds for the CUDA-graph allclose gate. +# Eager ignores the threshold (it gates strictly on bitwise equality); the +# CUDA-graph path tolerates inductor's per-class fusion noise up to this bound. +_DEFAULT_THRESHOLD = 0.15 +_CUDA_GRAPH_THRESHOLDS: Dict[tuple, float] = { + ("gpt2", "float16"): 0.15, + ("qwen3", "float16"): 0.15, + ("qwen2_moe", "float16"): 0.25, +} + +# Default hooks when ``--hooks`` is not given: an attention projection, a +# residual-stream read, and the final-logits read. +_DEFAULT_HOOKS = ["q", "k", "resid_pre", "final_logits"] + + +def cuda_graph_threshold(model_key: str, dtype: str) -> float: + """Resolve the CUDA-graph max-abs-diff alert threshold for a cell.""" + return _CUDA_GRAPH_THRESHOLDS.get((model_key, dtype), _DEFAULT_THRESHOLD) + + +def standard_for_mode(mode: str) -> str: + """The §8 comparison standard for a mode: bitwise (eager) / allclose (cg).""" + return "bitwise" if mode == "eager" else "allclose" + + +# --------------------------------------------------------------------------- +# Result dataclasses (JSON-serializable; no torch types stored) +# --------------------------------------------------------------------------- + + +@dataclass +class VocabDiff: + """One vocab entry's logit drift at the first differing token position.""" + + token_id: int + baseline_logit: float + monitored_logit: float + abs_diff: float + + def to_dict(self) -> Dict[str, Any]: + return { + "token_id": self.token_id, + "baseline_logit": self.baseline_logit, + "monitored_logit": self.monitored_logit, + "abs_diff": self.abs_diff, + } + + +@dataclass +class HookDrift: + """Per-hook drift record vs the unhooked baseline. + + ``check`` is the shared-lib :class:`~tests.lib.compare.Check` for the + logits comparison (carries passed / max_abs / mean_abs / first_diff_pos / + detail). The extra fields capture what the study adds on top. + """ + + hook: str + check: Optional[Check] = None + n_positions: int = 0 + vocab_size: int = 0 + first_diff_pos: int = -1 # token position; -1 == none + token_ids_diverged: bool = False + n_token_diff: int = 0 + topk_vocab_diffs: List[VocabDiff] = field(default_factory=list) + shape_mismatch: bool = False + error: Optional[str] = None + alert: bool = False + alert_reasons: List[str] = field(default_factory=list) + + def to_dict(self) -> Dict[str, Any]: + return { + "hook": self.hook, + "check": self.check.to_dict() if self.check is not None else None, + "n_positions": self.n_positions, + "vocab_size": self.vocab_size, + "first_diff_pos": self.first_diff_pos, + "token_ids_diverged": self.token_ids_diverged, + "n_token_diff": self.n_token_diff, + "topk_vocab_diffs": [v.to_dict() for v in self.topk_vocab_diffs], + "shape_mismatch": self.shape_mismatch, + "error": self.error, + "alert": self.alert, + "alert_reasons": self.alert_reasons, + } + + +@dataclass +class StudyResult: + framework: str + model: str + mode: str + variant: str + dtype: str + standard: str + threshold: float + topk: int + hooks: List[HookDrift] = field(default_factory=list) + + @property + def any_alert(self) -> bool: + return any(h.alert for h in self.hooks) + + def to_dict(self) -> Dict[str, Any]: + return { + "framework": self.framework, + "model": self.model, + "mode": self.mode, + "variant": self.variant, + "dtype": self.dtype, + "standard": self.standard, + "threshold": self.threshold, + "topk": self.topk, + "any_alert": self.any_alert, + "hooks": [h.to_dict() for h in self.hooks], + } + + +# --------------------------------------------------------------------------- +# Pure comparison core (CPU-testable; torch imported lazily) +# --------------------------------------------------------------------------- + + +def _first_diff_row(abs_diff: "Any") -> int: + """Index of the first token position holding any nonzero diff, or -1. + + ``abs_diff`` is a ``[N, vocab]`` tensor of absolute differences. + """ + if abs_diff.numel() == 0: + return -1 + row_has_diff = (abs_diff > 0).any(dim=-1) + nz = row_has_diff.nonzero(as_tuple=False) + return int(nz[0].item()) if nz.numel() > 0 else -1 + + +def compute_drift( + hook: str, + baseline: Dict[str, "Any"], + monitored: Dict[str, "Any"], + *, + mode: str, + threshold: float, + topk: int = 5, +) -> HookDrift: + """Compare one monitored capture against the unhooked baseline. + + ``baseline`` / ``monitored`` are dicts with ``token_ids`` (int64 ``[N]``) + and ``logits`` (float ``[N, vocab]``). Uses the shared-lib standard for + ``mode`` (``bitwise`` eager / ``allclose`` cuda graph) for the verdict and + augments it with token-divergence + top-k vocab drift. The §9 alert + policy is applied here. Never raises on bad input: records ``error`` / + ``shape_mismatch`` so the caller can alert. + """ + import torch + + drift = HookDrift(hook=hook) + + b_logits = baseline.get("logits") + m_logits = monitored.get("logits") + if b_logits is None or m_logits is None: + drift.error = "missing logits in capture" + return _finalize_alert(drift, mode=mode) + + b_logits = b_logits.float() + m_logits = m_logits.float() + drift.n_positions = int(b_logits.shape[0]) if b_logits.ndim >= 1 else 0 + drift.vocab_size = int(b_logits.shape[-1]) if b_logits.ndim >= 1 else 0 + + if drift.n_positions == 0 or m_logits.shape[0] == 0: + drift.error = "empty capture (no positions)" + return _finalize_alert(drift, mode=mode) + + if tuple(b_logits.shape) != tuple(m_logits.shape): + drift.shape_mismatch = True + drift.error = ( + f"logits shape mismatch: baseline {tuple(b_logits.shape)} " + f"vs monitored {tuple(m_logits.shape)}" + ) + # Still run the shared standard so the Check records the mismatch. + drift.check = _run_standard(b_logits, m_logits, mode=mode, threshold=threshold) + return _finalize_alert(drift, mode=mode) + + # Shared-lib verdict (bitwise / allclose) -- carries max/mean/first stats. + drift.check = _run_standard(b_logits, m_logits, mode=mode, threshold=threshold) + + abs_diff = (b_logits - m_logits).abs() + drift.first_diff_pos = _first_diff_row(abs_diff) + + # Greedy (argmax) token divergence, recomputed from logits and + # cross-checked against the stored ids. + b_arg = b_logits.argmax(dim=-1) + m_arg = m_logits.argmax(dim=-1) + drift.n_token_diff = int((b_arg != m_arg).sum().item()) + drift.token_ids_diverged = drift.n_token_diff > 0 + b_ids = baseline.get("token_ids") + m_ids = monitored.get("token_ids") + if b_ids is not None and m_ids is not None and b_ids.shape == m_ids.shape: + if not torch.equal(b_ids, m_ids): + drift.token_ids_diverged = True + drift.n_token_diff = max(drift.n_token_diff, int((b_ids != m_ids).sum().item())) + + # Top-k vocab diffs at the first differing position. + pos = drift.first_diff_pos + if pos >= 0 and topk > 0: + row = abs_diff[pos] + k = min(topk, int(row.numel())) + top = torch.topk(row, k) + for rank in range(k): + tid = int(top.indices[rank].item()) + drift.topk_vocab_diffs.append( + VocabDiff( + token_id=tid, + baseline_logit=float(b_logits[pos, tid].item()), + monitored_logit=float(m_logits[pos, tid].item()), + abs_diff=float(top.values[rank].item()), + ) + ) + + return _finalize_alert(drift, mode=mode) + + +def _run_standard(b_logits: "Any", m_logits: "Any", *, mode: str, threshold: float) -> Check: + """Run the §8 standard for ``mode`` and return its :class:`Check`.""" + name = standard_for_mode(mode) + if name == "bitwise": + return bitwise(b_logits, m_logits, name="logits_bitwise") + return allclose(b_logits, m_logits, name="logits_allclose", atol=threshold, rtol=0.0) + + +def _finalize_alert(drift: HookDrift, *, mode: str) -> HookDrift: + """Apply the §9 alert policy to a populated :class:`HookDrift` in place.""" + reasons: List[str] = [] + if drift.error is not None: + reasons.append(f"capture error: {drift.error}") + if drift.shape_mismatch: + reasons.append("logits shape mismatch (possible hook identity swap)") + if drift.token_ids_diverged: + reasons.append(f"greedy token ids diverged at {drift.n_token_diff} position(s)") + if drift.check is not None and not drift.check.passed and not drift.shape_mismatch: + if mode == "eager": + reasons.append(f"non-bitwise drift in eager mode ({drift.check.detail})") + else: + reasons.append(f"exceeds cuda-graph threshold ({drift.check.detail})") + drift.alert_reasons = reasons + drift.alert = bool(reasons) + return drift + + +# --------------------------------------------------------------------------- +# Human-readable table +# --------------------------------------------------------------------------- + + +def format_table(result: StudyResult) -> str: + """Render a fixed-width human-readable table for the study result.""" + header = ( + f"numeric-difference study: {result.framework}/{result.model} " + f"mode={result.mode} variant={result.variant} dtype={result.dtype} " + f"standard={result.standard}" + + (f" (threshold={result.threshold:g})" if result.standard == "allclose" else "") + ) + cols = ("hook", "verdict", "max_abs", "mean_abs", "first_diff", "tok_div", "alert") + widths = (16, 8, 12, 12, 10, 8, 6) + sep = " " + + def _row(values: tuple) -> str: + return sep.join(str(v).ljust(w) for v, w in zip(values, widths)) + + lines = [header, "-" * len(header), _row(cols)] + for h in result.hooks: + if h.error is not None and h.check is None: + verdict, mx, mn, fd = "ERR", "-", "-", "-" + else: + chk = h.check + verdict = "pass" if (chk and chk.passed) else "fail" + mx = f"{chk.max_abs:.4g}" if chk and chk.max_abs is not None else "-" + mn = f"{chk.mean_abs:.4g}" if chk and chk.mean_abs is not None else "-" + fd = str(h.first_diff_pos) if h.first_diff_pos >= 0 else "-" + lines.append( + _row(( + h.hook, verdict, mx, mn, fd, + str(h.n_token_diff) if h.token_ids_diverged else "0", + "ALERT" if h.alert else "ok", + )) + ) + for vd in h.topk_vocab_diffs: + lines.append( + sep + f" tok {vd.token_id}: base={vd.baseline_logit:.4g} " + f"mon={vd.monitored_logit:.4g} |Δ|={vd.abs_diff:.4g}" + ) + if h.alert: + for reason in h.alert_reasons: + lines.append(sep + f" ! {reason}") + n_alert = sum(1 for h in result.hooks if h.alert) + lines.append("") + lines.append(f"{len(result.hooks) - n_alert}/{len(result.hooks)} hooks clean" + + (f", {n_alert} alerting" if n_alert else "")) + return "\n".join(lines) + + +# --------------------------------------------------------------------------- +# Subprocess rollout runners (full-logits capture) +# --------------------------------------------------------------------------- + +# Each rollout subprocess saves {token_ids: int64[N], logits: float32[N, vocab]}. +# Same on-disk shape for HF and vLLM so the comparison code is framework-agnostic. + +_HF_RUNNER = dedent(""" + import argparse, os + import torch + + ap = argparse.ArgumentParser() + ap.add_argument('--model-key', required=True) + ap.add_argument('--hook', required=True) + ap.add_argument('--mode', required=True) + ap.add_argument('--variant', required=True, choices=['p', 'compare']) + ap.add_argument('--rollout', required=True, choices=['baseline', 'hooked']) + ap.add_argument('--max-new-tokens', type=int, default=4) + ap.add_argument('--prompt', default='Hello') + ap.add_argument('--out', required=True) + args = ap.parse_args() + + MODEL_ALIASES = { + 'gpt2': 'gpt2', + 'qwen3': 'Qwen/Qwen3-0.6B', + 'qwen2_moe': 'Qwen/Qwen1.5-MoE-A2.7B', + 'llama': 'meta-llama/Llama-3.1-8B', + } + hf_id = MODEL_ALIASES[args.model_key] + device = torch.device('cuda') + dtype = torch.float16 + + from transformers import AutoTokenizer + tok = AutoTokenizer.from_pretrained(hf_id) + if tok.pad_token_id is None: + tok.pad_token_id = tok.eos_token_id + + if args.rollout == 'baseline': + from transformers import AutoModelForCausalLM + model = AutoModelForCausalLM.from_pretrained( + hf_id, torch_dtype=dtype, attn_implementation='eager' + ).to(device).eval() + elif args.variant == 'compare': + # Patched _compare model (driver wraps this subprocess in isolated_hook). + if args.model_key == 'qwen3': + from transformers.models.qwen3_compare.modeling_qwen3 import CompareQwen3ForCausalLM as cls + elif args.model_key == 'gpt2': + from transformers.models.gpt2_compare.modeling_gpt2 import CompareGPT2LMHeadModel as cls + elif args.model_key == 'qwen2_moe': + from transformers.models.qwen2_moe_compare.modeling_qwen2_moe import CompareQwen2MoeForCausalLM as cls + elif args.model_key == 'llama': + from transformers.models.llama_compare.modeling_llama import CompareLlamaForCausalLM as cls + else: + raise ValueError(f'unsupported model_key={args.model_key!r}') + model = cls.from_pretrained( + hf_id, torch_dtype=dtype, attn_implementation='eager' + ).to(device).eval() + model.allocate_compare_buffers(1, 32, dtype=dtype, tp_size=1) + else: # variant p: production _p Hooked variant + hook_selection=H + if args.model_key == 'qwen3': + from transformers.models.qwen3_p.modeling_qwen3 import HookedQwen3ForCausalLM as cls + elif args.model_key == 'gpt2': + from transformers.models.gpt2_p.modeling_gpt2 import HookedGPT2LMHeadModel as cls + elif args.model_key == 'qwen2_moe': + from transformers.models.qwen2_moe_p.modeling_qwen2_moe import HookedQwen2MoeForCausalLM as cls + else: + raise ValueError(f'unsupported model_key={args.model_key!r}') + model = cls.from_pretrained( + hf_id, torch_dtype=dtype, attn_implementation='eager' + ).to(device).eval() + + inputs = tok([args.prompt], return_tensors='pt', padding=True).to(device) + gen_kwargs = dict( + **inputs, + max_new_tokens=args.max_new_tokens, + do_sample=False, + pad_token_id=tok.pad_token_id, + return_dict_in_generate=True, + output_scores=True, + ) + if args.mode == 'cuda_graph': + from transformers import CompileConfig + gen_kwargs['cache_implementation'] = 'static' + gen_kwargs['compile_config'] = CompileConfig(mode='reduce-overhead', fullgraph=False) + + if args.rollout == 'hooked' and args.variant == 'p': + from monitoring import MonitoringEngine, MonitoringConfig + from monitoring.config import CaptureSchedule + from monitoring._native_engine import RingConfig + from integration.hf_adapter import generate_with_monitoring + cfg = MonitoringConfig(schedule=CaptureSchedule(capture_prefill=True, capture_decode=True)) + engine = MonitoringEngine(config=cfg, model_id='numeric_study') + ring_cfg = RingConfig() + ring_cfg.task_ring_entries = 1024 + ring_cfg.payload_ring_bytes = 64 * 1024 * 1024 + ring_cfg.pinned_staging_bytes = 64 * 1024 * 1024 + engine.enable_ring_transport(ring_cfg) + model.monitoring_engine = engine + try: + out = generate_with_monitoring(model, hook_selection=args.hook, **gen_kwargs) + finally: + engine.close() + else: + with torch.no_grad(): + out = model.generate(**gen_kwargs) + + scores = torch.stack(out.scores, dim=0) # [N, 1, vocab] + logits = scores.squeeze(1).float().cpu() # [N, vocab] + token_ids = logits.argmax(dim=-1).to(torch.int64) # [N] + torch.save({'token_ids': token_ids, 'logits': logits}, args.out) + print(f'OK {args.rollout}/{args.variant} N={logits.shape[0]} V={logits.shape[1]} -> {args.out}') +""") + + +_VLLM_RUNNER = dedent(""" + import argparse, os + os.environ.setdefault('VLLM_DISABLE_COMPILE_CACHE', '1') + import torch + from vllm import LLM, SamplingParams + + ap = argparse.ArgumentParser() + ap.add_argument('--model-key', required=True) + ap.add_argument('--hook', required=True) + ap.add_argument('--mode', required=True) + ap.add_argument('--variant', required=True, choices=['p', 'compare']) + ap.add_argument('--rollout', required=True, choices=['baseline', 'hooked']) + ap.add_argument('--max-new-tokens', type=int, default=4) + ap.add_argument('--prompt', default='Hello') + ap.add_argument('--out', required=True) + args = ap.parse_args() + + MODEL_ALIASES = { + 'gpt2': 'gpt2', + 'qwen3': 'Qwen/Qwen3-0.6B', + 'qwen2_moe': 'Qwen/Qwen1.5-MoE-A2.7B', + } + model_name = MODEL_ALIASES[args.model_key] + + llm_kwargs = dict( + model=model_name, + max_model_len=128, + gpu_memory_utilization=0.5, + enforce_eager=(args.mode == 'eager'), + ) + if args.rollout == 'hooked': + additional_config = {'dmx_hook_selection': args.hook, 'dmx_db_host': ''} + if args.variant == 'compare': + llm_kwargs['worker_cls'] = 'tests.compare_worker.CompareWorker' + else: + llm_kwargs['worker_cls'] = 'integration.vllm_adapter.DMXGPUWorker' + llm_kwargs['additional_config'] = additional_config + + llm = LLM(**llm_kwargs) + params = SamplingParams(temperature=0.0, max_tokens=args.max_new_tokens, logprobs=20) + outputs = llm.generate([args.prompt], params) + + completion = outputs[0].outputs[0] + ids = list(completion.token_ids) + step_logprobs = completion.logprobs or [] + # Dense [N, V] from the sparse top-k logprob dicts; unreported entries stay + # at a large negative floor (sufficient for chosen-token + top-k drift). + vocab = int(getattr(llm.llm_engine.model_config.hf_config, 'vocab_size', 0)) or 1 + N = len(ids) + logits = torch.full((N, vocab), -1e30, dtype=torch.float32) + for i in range(N): + if i < len(step_logprobs) and step_logprobs[i] is not None: + for tid, lp in step_logprobs[i].items(): + logits[i, int(tid)] = float(lp.logprob) + token_ids = torch.tensor(ids, dtype=torch.int64) + torch.save({'token_ids': token_ids, 'logits': logits}, args.out) + print(f'OK {args.rollout}/{args.variant} N={N} V={vocab} -> {args.out}') + + try: + llm.collective_rpc('stop_monitoring') + except Exception: + pass +""") + + +def _build_subprocess_env() -> dict: + """Pin CUDA_VISIBLE_DEVICES=0 and put conda lib on LD_LIBRARY_PATH so vLLM + imports resolve (mirrors test_per_hook_isolation / test_no_graph_breaks).""" + env = os.environ.copy() + conda_prefix = env.get("CONDA_PREFIX") + if conda_prefix: + ld = env.get("LD_LIBRARY_PATH", "") + env["LD_LIBRARY_PATH"] = f"{conda_prefix}/lib:{ld}" if ld else f"{conda_prefix}/lib" + env.setdefault("CUDA_VISIBLE_DEVICES", "0") + return env + + +def _run_rollout( + *, + framework: str, + model_key: str, + hook: str, + mode: str, + variant: str, + rollout: str, + out_path: Path, + max_new_tokens: int, + prompt: str, + env: dict, + timeout: int = 600, +) -> None: + """Spawn one rollout subprocess; raise with captured output on failure.""" + runner = _HF_RUNNER if framework == "hf" else _VLLM_RUNNER + cmd = [ + sys.executable, "-c", runner, + "--model-key", model_key, + "--hook", hook, + "--mode", mode, + "--variant", variant, + "--rollout", rollout, + "--max-new-tokens", str(max_new_tokens), + "--prompt", prompt, + "--out", str(out_path), + ] + proc = subprocess.run( + cmd, capture_output=True, text=True, timeout=timeout, env=env, cwd=REPO_ROOT, + ) + if proc.returncode != 0: + raise RuntimeError( + f"rollout={rollout} variant={variant} hook={hook} failed " + f"(rc={proc.returncode})\n" + f"--- stdout ---\n{proc.stdout}\n" + f"--- stderr (tail) ---\n{proc.stderr[-3000:]}" + ) + + +# --------------------------------------------------------------------------- +# Orchestration +# --------------------------------------------------------------------------- + + +def run_study( + *, + framework: str, + model: str, + mode: str, + hooks: List[str], + out_dir: Path, + variant: str = "p", + dtype: str = "float16", + topk: int = 5, + max_new_tokens: int = 4, + prompt: str = "Hello", + threshold: Optional[float] = None, +) -> StudyResult: + """Run the full per-hook study and return a populated :class:`StudyResult`. + + Captures the unhooked baseline once, then one monitored rollout per hook, + comparing each against the baseline and applying the §9 alert policy. For + ``variant='compare'`` each hooked rollout runs inside the hardened + :func:`tests.isolate_hook.isolated_hook` context manager (plan §6). + """ + import torch + + out_dir = Path(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + env = _build_subprocess_env() + thr = threshold if threshold is not None else cuda_graph_threshold(model, dtype) + + result = StudyResult( + framework=framework, model=model, mode=mode, variant=variant, + dtype=dtype, standard=standard_for_mode(mode), threshold=thr, topk=topk, + ) + + # 1. Baseline (unhooked) captured once; hook arg is unused by the runner. + baseline_path = out_dir / "baseline.pt" + _run_rollout( + framework=framework, model_key=model, hook=hooks[0] if hooks else "q", + mode=mode, variant="p", rollout="baseline", out_path=baseline_path, + max_new_tokens=max_new_tokens, prompt=prompt, env=env, + ) + baseline = torch.load(baseline_path, map_location="cpu") + + # 2-3. One hook at a time; compare; apply alert policy. + for hook in hooks: + drift = HookDrift(hook=hook) + try: + hooked_path = out_dir / f"hooked_{hook}.pt" + if variant == "compare": + from tests.isolate_hook import isolated_hook + + with isolated_hook(framework, model, hook): + _run_rollout( + framework=framework, model_key=model, hook=hook, mode=mode, + variant="compare", rollout="hooked", out_path=hooked_path, + max_new_tokens=max_new_tokens, prompt=prompt, env=env, + ) + else: + _run_rollout( + framework=framework, model_key=model, hook=hook, mode=mode, + variant="p", rollout="hooked", out_path=hooked_path, + max_new_tokens=max_new_tokens, prompt=prompt, env=env, + ) + monitored = torch.load(hooked_path, map_location="cpu") + drift = compute_drift( + hook, baseline, monitored, mode=mode, threshold=thr, topk=topk, + ) + except Exception as exc: # subprocess crash / OOM / dirty restore -> alert + drift.error = f"{type(exc).__name__}: {exc}" + _finalize_alert(drift, mode=mode) + result.hooks.append(drift) + + return result + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def main(argv: Optional[List[str]] = None) -> int: + ap = argparse.ArgumentParser( + description="Per-hook numeric-difference study vs the unhooked baseline." + ) + ap.add_argument("--framework", choices=["hf", "vllm"], default="hf") + ap.add_argument("--model", default="qwen3") + ap.add_argument("--mode", choices=["eager", "cuda_graph"], default="eager") + ap.add_argument("--variant", choices=["p", "compare"], default="p", + help="p: production _p+hook_selection; compare: _compare under isolated_hook") + ap.add_argument("--hooks", default=",".join(_DEFAULT_HOOKS), + help="Comma-separated hook short-names (e.g. q,k,resid_pre,final_logits).") + ap.add_argument("--dtype", default="float16") + ap.add_argument("--topk", type=int, default=5) + ap.add_argument("--max-new-tokens", type=int, default=4) + ap.add_argument("--prompt", default="Hello") + ap.add_argument("--threshold", type=float, default=None, + help="Override the CUDA-graph max-abs-diff alert threshold.") + ap.add_argument("--work-dir", default=None, + help="Scratch dir for per-rollout .pt files (default: temp dir).") + ap.add_argument("--out", default=None, + help="Write the JSON artifact here (default: stdout table only).") + args = ap.parse_args(argv) + + hooks = [h.strip() for h in args.hooks.split(",") if h.strip()] + if not hooks: + ap.error("no hooks selected") + + import tempfile + + tmp_ctx = None + if args.work_dir: + work_dir = Path(args.work_dir) + else: + tmp_ctx = tempfile.TemporaryDirectory(prefix="numeric_study_") + work_dir = Path(tmp_ctx.name) + + try: + result = run_study( + framework=args.framework, model=args.model, mode=args.mode, + hooks=hooks, out_dir=work_dir, variant=args.variant, dtype=args.dtype, + topk=args.topk, max_new_tokens=args.max_new_tokens, prompt=args.prompt, + threshold=args.threshold, + ) + finally: + if tmp_ctx is not None: + tmp_ctx.cleanup() + + print(format_table(result)) + if args.out: + out_path = Path(args.out) + out_path.parent.mkdir(parents=True, exist_ok=True) + out_path.write_text(json.dumps(result.to_dict(), indent=2), encoding="utf-8") + print(f"\nwrote JSON artifact -> {out_path}") + + # Non-zero exit if any hook alerted, so CI / wrappers can gate on it. + return 1 if result.any_alert else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_adapter_protocol.py b/tests/test_adapter_protocol.py index 11151beb4..e9d41414f 100644 --- a/tests/test_adapter_protocol.py +++ b/tests/test_adapter_protocol.py @@ -16,9 +16,13 @@ import dataclasses +import pytest + from monitoring.adaptor_base import BackendAdaptor from monitoring.step_context import StepContext +pytestmark = pytest.mark.cpu + # --------------------------------------------------------------------------- # Test doubles diff --git a/tests/test_config.py b/tests/test_config.py index d5bd177a7..3c7ab4da1 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -15,6 +15,8 @@ from monitoring.config import CaptureSchedule, MonitoringConfig +pytestmark = pytest.mark.cpu + # ============================================================================= # CaptureSchedule diff --git a/tests/test_e2e_correctness_vs_hf.py b/tests/test_e2e_correctness_vs_hf.py index 31cb562f1..8f2080940 100644 --- a/tests/test_e2e_correctness_vs_hf.py +++ b/tests/test_e2e_correctness_vs_hf.py @@ -76,6 +76,13 @@ import pytest import torch +pytestmark = [ + pytest.mark.gpu, + pytest.mark.e2e, + pytest.mark.clickhouse, + pytest.mark.hf, +] + from monitoring.clickhouse_reader import CHClickhouseDriverReadOnly from monitoring.segment_merger import merge_segments, parse_internal_id diff --git a/tests/test_e2e_lib.py b/tests/test_e2e_lib.py new file mode 100644 index 000000000..0f03017be --- /dev/null +++ b/tests/test_e2e_lib.py @@ -0,0 +1,309 @@ +"""CPU unit coverage for the shared E2E lib + matrix expansion (plan §7, §8). + +These are pure-CPU (torch-only) — no CUDA / ClickHouse / vLLM / weights — and +guard the de-duplicated comparison/align/report logic plus the matrix's cell +expansion and env translation (exercised via ``--dry-run`` internals). +""" +from __future__ import annotations + +import json + +import pytest +import torch + +from tests.lib import align, compare +from tests.lib.compare import Check +from tests.lib.report import ( + CellResult, + checks_from_legacy_result, + human_table, + read_jsonl, + write_jsonl, +) + +pytestmark = pytest.mark.cpu + + +# --------------------------------------------------------------------------- +# compare.py — the four standards +# --------------------------------------------------------------------------- + + +class TestCompareStandards: + def test_bitwise_equal_records_zero_drift(self): + a = torch.arange(12, dtype=torch.float32).reshape(3, 4) + c = compare.bitwise(a, a.clone(), "x") + assert c.passed and c.max_abs == 0.0 and c.mean_abs == 0.0 + assert c.first_diff_pos is None + + def test_bitwise_detects_diff_and_reports_first_pos(self): + a = torch.arange(6, dtype=torch.float32) + b = a.clone() + b[2] = 99.0 + c = compare.bitwise(a, b, "x") + assert not c.passed + assert c.first_diff_pos == 2 + assert c.max_abs == pytest.approx(99.0 - 2.0) + + def test_bitwise_shape_mismatch_fails_closed(self): + c = compare.bitwise(torch.zeros(4), torch.zeros(5), "x") + assert not c.passed and "shape mismatch" in c.detail + + def test_bitwise_dtype_mismatch_fails_closed(self): + c = compare.bitwise(torch.zeros(4, dtype=torch.float32), + torch.zeros(4, dtype=torch.float16), "x") + assert not c.passed and "dtype mismatch" in c.detail + + def test_allclose_passes_within_tol_but_records_drift(self): + a = torch.zeros(8) + b = a + 1e-4 + c = compare.allclose(a, b, "x", atol=1e-3) + assert c.passed + # drift recorded even on a pass + assert c.max_abs == pytest.approx(1e-4, abs=1e-9) + assert c.mean_abs == pytest.approx(1e-4, abs=1e-9) + + def test_allclose_fails_outside_tol(self): + a = torch.zeros(8) + b = a + 1.0 + c = compare.allclose(a, b, "x", atol=1e-3) + assert not c.passed and c.first_diff_pos == 0 + + def test_transport_bitwise_is_exact(self): + a = torch.randn(4, 4) + assert compare.transport_bitwise(a, a.clone()).passed + assert not compare.transport_bitwise(a, a + 1e-6).passed + + def test_row_count_ok(self): + counts = {f"blocks.hook_{i}": 5 for i in range(12)} + counts["final_logits"] = 3 + c = compare.row_count(counts) + assert c.passed, c.detail + + def test_row_count_uneven_fails(self): + counts = {f"blocks.hook_{i}": 5 for i in range(12)} + counts["blocks.hook_0"] = 4 + counts["final_logits"] = 3 + c = compare.row_count(counts) + assert not c.passed and "uneven" in c.detail + + def test_row_count_missing_final_logits_fails(self): + counts = {f"blocks.hook_{i}": 5 for i in range(12)} + c = compare.row_count(counts) + assert not c.passed and "final_logits" in c.detail + + def test_row_count_too_few_types_fails(self): + c = compare.row_count({"blocks.hook_0": 1, "final_logits": 1}) + assert not c.passed + + def test_compare_tensors_dispatch_and_unknown(self): + a = torch.zeros(3) + assert compare.compare_tensors(a, a.clone(), "bitwise").passed + with pytest.raises(ValueError): + compare.compare_tensors(a, a, "row_count") + + def test_bytes_identical_neg_zero(self): + # -0.0 and 0.0 are equal numerically but differ bitwise. + a = torch.tensor([0.0]) + b = torch.tensor([-0.0]) + assert torch.equal(a, b) # numerically equal + assert not compare.bytes_identical(a, b) # but not byte-identical + + +# --------------------------------------------------------------------------- +# align.py +# --------------------------------------------------------------------------- + + +class TestAlign: + def test_parse_request_id(self): + assert align.parse_request_id("3:7") == (3, 7) + with pytest.raises(ValueError): + align.parse_request_id("nope") + + def test_normalize_request_id_strips_vllm_suffix(self): + assert align.normalize_request_id("12:0-deadbeef") == "12:0" + assert align.normalize_request_id("12:0") == "12:0" + + def test_strip_left_pad(self): + ids = torch.tensor([0, 0, 5, 6, 7]) + attn = torch.tensor([0, 0, 1, 1, 1]) + assert torch.equal(align.strip_left_pad(ids, attn), torch.tensor([5, 6, 7])) + + def test_strip_left_pad_all_padding(self): + ids = torch.tensor([0, 0]) + attn = torch.tensor([0, 0]) + assert align.strip_left_pad(ids, attn).numel() == 0 + + def test_trim_eos_drops_and_keeps(self): + ids = torch.tensor([1, 2, 9, 3]) + assert torch.equal(align.trim_eos(ids, 9), torch.tensor([1, 2])) + assert torch.equal(align.trim_eos(ids, 9, keep_eos=True), torch.tensor([1, 2, 9])) + + def test_trim_eos_absent_returns_all(self): + ids = torch.tensor([1, 2, 3]) + assert torch.equal(align.trim_eos(ids, 9), ids) + + def test_align_to_min_len(self): + a = torch.arange(5) + b = torch.arange(3) + ra, rb = align.align_to_min_len(a, b) + assert ra.numel() == rb.numel() == 3 + + def test_logits_align_skip(self): + # skip = max(0, db_len - ref_len - 1): the DB keeps every position + # while generate() yields gen-1 rows, so the head offset is the + # prefill span. Never negative when ref is longer than db. + assert align.logits_align_skip(db_len=9, ref_len=4) == 4 + assert align.logits_align_skip(db_len=2, ref_len=5) == 0 + + +# --------------------------------------------------------------------------- +# report.py +# --------------------------------------------------------------------------- + + +class TestReport: + def test_cellresult_finalize_and_record(self): + cr = CellResult("vllm", "qwen3", "eager", "transport_bitwise", "vllm-full", tp=1) + cr.checks = [Check("token_ids", True), Check("layer.0.q", True, max_abs=0.0)] + cr.finalize() + rec = cr.to_record() + assert rec["passed"] is True + assert rec["backend"] == "vllm" and rec["hook_selection"] == "vllm-full" + assert len(rec["checks"]) == 2 + assert rec["checks"][1]["max_abs"] == 0.0 + + def test_cellresult_fails_when_any_check_fails(self): + cr = CellResult("hf", "gpt2", "eager", "bitwise", "vllm-full") + cr.checks = [Check("a", True), Check("b", False, detail="boom")] + cr.finalize() + assert cr.passed is False + + def test_cellresult_no_checks_is_not_passed(self): + cr = CellResult("hf", "gpt2", "eager", "bitwise", "vllm-full").finalize() + assert cr.passed is False + + def test_cellresult_error_sets_failed(self): + cr = CellResult("hf", "gpt2", "eager", "bitwise", "vllm-full") + cr.error = "runner crashed" + cr.checks = [Check("a", True)] + cr.finalize() + assert cr.passed is False + assert cr.to_record()["error"] == "runner crashed" + + def test_legacy_result_adapter(self): + legacy = {"tests": [ + {"name": "rows_found", "passed": True, "detail": "10 rows"}, + {"name": "x", "passed": False, "detail": "max_abs=1e-3"}, + ]} + checks = checks_from_legacy_result(legacy) + assert [c.name for c in checks] == ["rows_found", "x"] + assert checks[0].passed and not checks[1].passed + + def test_jsonl_roundtrip(self, tmp_path): + crs = [] + for passed in (True, False): + cr = CellResult("vllm", "gpt2", "eager", "row_count", "vllm-full") + cr.checks = [Check("c", passed)] + crs.append(cr.finalize()) + out = tmp_path / "e2e.jsonl" + write_jsonl(crs, str(out)) + recs = read_jsonl(str(out)) + assert len(recs) == 2 + assert recs[0]["passed"] is True and recs[1]["passed"] is False + # each line is valid standalone JSON + for line in out.read_text().splitlines(): + json.loads(line) + + def test_human_table_renders_counts(self): + crs = [ + CellResult("vllm", "gpt2", "eager", "bitwise", "vllm-full").finalize(), + ] + crs[0].checks = [Check("c", True)] + crs[0].finalize() + table = human_table(crs) + assert "backend" in table and "1/1 cells passed" in table + + +# --------------------------------------------------------------------------- +# e2e_matrix — cell expansion + env translation (the dry-run surface) +# --------------------------------------------------------------------------- + + +class TestMatrixExpansion: + def _args(self, **over): + from tests.e2e_matrix import build_parser + argv = [] + for k, v in over.items(): + argv += [f"--{k.replace('_', '-')}", str(v)] + return build_parser().parse_args(argv) + + def test_cartesian_product_count(self): + from tests.e2e_matrix import build_cells + args = self._args(backend="hf,vllm", model="gpt2,qwen3", + mode="eager,cuda_graph", standard="row_count") + cells = build_cells(args) + assert len(cells) == 2 * 2 * 2 * 1 + + def test_env_translates_hook_selection(self): + from tests.e2e_matrix import build_cells, cell_env + args = self._args(backend="vllm", model="qwen3", mode="cuda_graph", + standard="transport_bitwise", hooks="q") + cell = build_cells(args)[0] + env = cell_env(cell, args, base={}) + # public -> internal contract (plan §2) + assert env["E2E_HOOK_SELECTION"] == "q" + assert env["DMX_HOOK_SELECTION"] == "q" + # cuda_graph -> not eager + assert env["E2E_ENFORCE_EAGER"] == "0" + assert env["E2E_CUDA_GRAPHS"] == "1" + assert env["E2E_MODEL"] == "qwen3" + + def test_eager_sets_enforce_eager(self): + from tests.e2e_matrix import build_cells, cell_env + args = self._args(backend="hf", model="gpt2", mode="eager", standard="bitwise") + env = cell_env(build_cells(args)[0], args, base={}) + assert env["E2E_ENFORCE_EAGER"] == "1" + assert env["E2E_CUDA_GRAPHS"] == "0" + + def test_plan_vllm_identical_dispatch(self): + from tests.e2e_matrix import build_cells, plan_cell + args = self._args(backend="vllm", model="qwen3", mode="eager", + standard="transport_bitwise", hooks="vllm-full") + steps, comparator, _result = plan_cell(build_cells(args)[0], "/run") + labels = [s.label for s in steps] + assert labels == ["enable_ref_hooks", "vllm_ref", "vllm_monitored", "compare"] + assert comparator == "tests.vllm_identical_comparator" + + def test_plan_vllm_rowcount_dispatch(self): + from tests.e2e_matrix import build_cells, plan_cell + args = self._args(backend="vllm", model="gpt2", mode="eager", + standard="row_count") + steps, comparator, _ = plan_cell(build_cells(args)[0], "/run") + assert [s.label for s in steps] == ["vllm_monitored", "compare"] + assert comparator == "tests.vllm_rowcnt_comparator" + + def test_plan_hf_dispatch(self): + from tests.e2e_matrix import build_cells, plan_cell + args = self._args(backend="hf", model="gpt2", mode="eager", standard="allclose") + steps, comparator, _ = plan_cell(build_cells(args)[0], "/run") + assert [s.label for s in steps] == ["hf_ref", "hf_monitored", "compare"] + assert comparator == "tests.hf_comparator" + + def test_unknown_backend_raises(self): + from tests.e2e_matrix import Cell, plan_cell + with pytest.raises(ValueError): + plan_cell(Cell("nope", "gpt2", "eager", "bitwise", "vllm-full"), "/run") + + def test_main_dry_run_no_side_effects(self, capsys): + from tests.e2e_matrix import main + rc = main(["--backend", "hf,vllm", "--model", "gpt2", + "--standard", "row_count", "--dry-run"]) + assert rc == 0 + out = capsys.readouterr().out + assert "2 cell(s) planned" in out + + def test_main_empty_axis_returns_2(self): + from tests.e2e_matrix import main + assert main(["--backend", "", "--dry-run"]) == 2 diff --git a/tests/test_gpt2_parity.py b/tests/test_gpt2_parity.py index 08ad887e5..2cac0705a 100644 --- a/tests/test_gpt2_parity.py +++ b/tests/test_gpt2_parity.py @@ -5,6 +5,11 @@ from transformers.models.gpt2_p.modeling_gpt2 import HookedGPT2Model from transformers.models.gpt2_p.modeling_gpt2 import GPT2LMHeadModel as HFModifiedGPT2 +from tests._requirements import require_model_cache + +# Runs on CPU but pulls real gpt2 weights via from_pretrained -> `hf`, not `cpu`. +pytestmark = [pytest.mark.hf, require_model_cache("gpt2")] + @pytest.fixture(scope="module") def gpt2_tokenizer(): diff --git a/tests/test_hf_eos_strip.py b/tests/test_hf_eos_strip.py index 85034d8b2..ee58af598 100644 --- a/tests/test_hf_eos_strip.py +++ b/tests/test_hf_eos_strip.py @@ -21,6 +21,8 @@ import pytest import torch +pytestmark = pytest.mark.cpu + from integration.hf_adapter import HFAdaptor diff --git a/tests/test_hook_spec_flags.py b/tests/test_hook_spec_flags.py index c081f12ac..f95510a2a 100644 --- a/tests/test_hook_spec_flags.py +++ b/tests/test_hook_spec_flags.py @@ -11,11 +11,14 @@ from unittest.mock import MagicMock +import pytest import torch import torch.nn as nn from monitoring.ring_transport import HookSpec, ModelShapeConfig, RingTransport +pytestmark = pytest.mark.cpu + def test_hook_spec_flag_defaults_false(): spec = HookSpec(hook_type=0, module=nn.Identity()) diff --git a/tests/test_moe_v1_routing_hooks.py b/tests/test_moe_v1_routing_hooks.py index 271092f59..1959c94fd 100644 --- a/tests/test_moe_v1_routing_hooks.py +++ b/tests/test_moe_v1_routing_hooks.py @@ -2,6 +2,8 @@ import json +import pytest + from transformers import Qwen2MoeConfig from transformers.models.qwen2_moe_compare.modeling_qwen2_moe import CompareQwen2MoeForCausalLM from transformers.models.qwen2_moe_p.modeling_qwen2_moe import HookedQwen2MoeForCausalLM @@ -19,6 +21,8 @@ ) from tests.ref_disk_worker import _ARCH_REMAP as _REF_ARCH_REMAP +pytestmark = pytest.mark.cpu + def test_moe_v1_routing_hook_types_registered() -> None: assert _id_by_short["router_logits"] == HOOK_TYPE_ROUTER_LOGITS diff --git a/tests/test_no_graph_breaks.py b/tests/test_no_graph_breaks.py index 0b9b4b6f4..e640cccb0 100644 --- a/tests/test_no_graph_breaks.py +++ b/tests/test_no_graph_breaks.py @@ -34,6 +34,8 @@ import pytest +pytestmark = [pytest.mark.gpu, pytest.mark.vllm, pytest.mark.hf] + REPO_ROOT = Path(__file__).resolve().parent.parent HF_GPT2_MODEL = "gpt2" diff --git a/tests/test_numeric_study.py b/tests/test_numeric_study.py new file mode 100644 index 000000000..0a9c3ff07 --- /dev/null +++ b/tests/test_numeric_study.py @@ -0,0 +1,217 @@ +"""Tests for the per-hook numeric-difference study (plan §9 / Phase 5). + +Two layers, mirroring ``test_per_hook_isolation.py``: + +1. **CPU unit tests** (always run): exercise the pure comparison / alert / + report core (``compute_drift``, ``format_table``, serialization, threshold + + standard selection) on small hand-built CPU tensors. These guard the + study's verdict logic -- including its use of the shared ``tests.lib.compare`` + standards -- independently of any GPU rollout. + +2. **GPU smoke** (marked ``numeric`` + ``gpu`` + ``slow``, opt-in): run the + real study for a couple of observational hooks and assert the machinery + produces serializable per-hook records and a self-consistent eager verdict. +""" +from __future__ import annotations + +import json + +import pytest + +from tests._requirements import require_cuda +from tests.lib.compare import Check +from tests.numeric_study import ( + HookDrift, + StudyResult, + VocabDiff, + compute_drift, + cuda_graph_threshold, + format_table, + run_study, + standard_for_mode, +) + + +# --------------------------------------------------------------------------- +# CPU unit tests for the comparison core +# --------------------------------------------------------------------------- + + +def _cap(logits): + import torch + + t = torch.tensor(logits, dtype=torch.float32) + return {"logits": t, "token_ids": t.argmax(dim=-1).to(torch.int64)} + + +@pytest.mark.cpu +class TestComputeDrift: + def test_identical_eager_passes_bitwise(self): + base = _cap([[1.0, 2.0, 0.5], [0.1, 0.2, 3.0]]) + d = compute_drift("q", base, base, mode="eager", threshold=0.15, topk=3) + assert d.check is not None and d.check.name == "logits_bitwise" + assert d.check.passed is True + assert d.check.max_abs == 0.0 + assert d.first_diff_pos == -1 + assert d.token_ids_diverged is False + assert d.topk_vocab_diffs == [] + assert d.alert is False + assert d.n_positions == 2 and d.vocab_size == 3 + + def test_eager_any_drift_alerts(self): + base = _cap([[5.0, 1.0, 0.0], [0.0, 4.0, 1.0]]) + mon = _cap([[5.0, 1.0, 0.0], [0.0, 4.0, 1.25]]) # tiny perturb, no argmax flip + d = compute_drift("k", base, mon, mode="eager", threshold=0.15, topk=2) + assert d.check.passed is False + assert d.first_diff_pos == 1 + assert d.check.max_abs == pytest.approx(0.25, abs=1e-6) + assert d.token_ids_diverged is False + assert d.alert is True + assert any("eager" in r for r in d.alert_reasons) + # Top-k drift at the first differing position points at the perturbed id. + assert d.topk_vocab_diffs[0].token_id == 2 + assert d.topk_vocab_diffs[0].abs_diff == pytest.approx(0.25, abs=1e-6) + + def test_cuda_graph_below_threshold_no_alert(self): + base = _cap([[5.0, 1.0, 0.0]]) + mon = _cap([[5.0, 1.05, 0.0]]) # max_abs 0.05 < 0.15 + d = compute_drift("q", base, mon, mode="cuda_graph", threshold=0.15, topk=1) + assert d.check.name == "logits_allclose" + assert d.check.passed is True + assert d.alert is False + + def test_cuda_graph_above_threshold_alerts(self): + base = _cap([[5.0, 1.0, 0.0]]) + mon = _cap([[5.0, 1.5, 0.0]]) # max_abs 0.5 > 0.15, no argmax flip + d = compute_drift("q", base, mon, mode="cuda_graph", threshold=0.15, topk=1) + assert d.check.passed is False + assert d.alert is True + assert any("threshold" in r for r in d.alert_reasons) + + def test_argmax_flip_always_alerts_even_under_cuda_graph(self): + base = _cap([[5.0, 1.0, 0.0]]) + mon = _cap([[1.0, 9.0, 0.0]]) # argmax 0 -> 1 + d = compute_drift("resid_pre", base, mon, mode="cuda_graph", threshold=100.0, topk=1) + assert d.token_ids_diverged is True + assert d.n_token_diff == 1 + assert d.alert is True + assert any("token ids diverged" in r for r in d.alert_reasons) + + def test_shape_mismatch_alerts(self): + import torch + + base = _cap([[1.0, 2.0, 3.0]]) + mon = { + "logits": torch.zeros((1, 4), dtype=torch.float32), + "token_ids": torch.zeros((1,), dtype=torch.int64), + } + d = compute_drift("q", base, mon, mode="eager", threshold=0.15, topk=1) + assert d.shape_mismatch is True + assert d.alert is True + assert any("identity swap" in r for r in d.alert_reasons) + + def test_empty_capture_alerts(self): + import torch + + empty = { + "logits": torch.zeros((0, 3), dtype=torch.float32), + "token_ids": torch.zeros((0,), dtype=torch.int64), + } + d = compute_drift("q", empty, empty, mode="eager", threshold=0.15, topk=1) + assert d.error is not None + assert d.alert is True + + def test_missing_logits_alerts(self): + d = compute_drift("q", {"token_ids": None}, {"token_ids": None}, + mode="eager", threshold=0.15, topk=1) + assert d.error is not None + assert d.alert is True + + +@pytest.mark.cpu +class TestSelectorsAndReport: + def test_standard_for_mode(self): + assert standard_for_mode("eager") == "bitwise" + assert standard_for_mode("cuda_graph") == "allclose" + + def test_threshold_lookup_and_fallback(self): + assert cuda_graph_threshold("qwen2_moe", "float16") == 0.25 + assert cuda_graph_threshold("nonesuch", "float64") == pytest.approx(0.15) + + def test_format_table_contains_hooks_and_verdict(self): + result = StudyResult( + framework="hf", model="qwen3", mode="eager", variant="p", + dtype="float16", standard="bitwise", threshold=0.15, topk=2, + ) + ok = HookDrift(hook="resid_pre", check=Check("logits_bitwise", True, max_abs=0.0)) + bad = HookDrift( + hook="q", + check=Check("logits_bitwise", False, max_abs=0.3, mean_abs=0.01, detail="max_abs=3e-01"), + first_diff_pos=1, alert=True, alert_reasons=["non-bitwise drift in eager mode"], + topk_vocab_diffs=[VocabDiff(7, 1.0, 1.3, 0.3)], + ) + result.hooks.extend([ok, bad]) + table = format_table(result) + assert "resid_pre" in table + assert "ALERT" in table + assert "tok 7" in table + assert "1/2 hooks clean" in table + + def test_result_to_dict_is_json_serializable(self): + result = StudyResult( + framework="hf", model="qwen3", mode="cuda_graph", variant="compare", + dtype="float16", standard="allclose", threshold=0.15, topk=1, + ) + result.hooks.append( + HookDrift( + hook="q", + check=Check("logits_allclose", False, max_abs=0.2, mean_abs=0.01), + topk_vocab_diffs=[VocabDiff(1, 0.0, 0.2, 0.2)], + alert=True, alert_reasons=["exceeds cuda-graph threshold"], + ) + ) + payload = result.to_dict() + s = json.dumps(payload) # must not raise + rt = json.loads(s) + assert rt["any_alert"] is True + assert rt["variant"] == "compare" + assert rt["hooks"][0]["check"]["name"] == "logits_allclose" + assert rt["hooks"][0]["topk_vocab_diffs"][0]["token_id"] == 1 + + +# --------------------------------------------------------------------------- +# GPU smoke +# --------------------------------------------------------------------------- + + +@pytest.mark.numeric +@pytest.mark.gpu +@pytest.mark.slow +@require_cuda() +def test_numeric_study_hf_qwen3_eager_smoke(tmp_path): + """Eager observational hooks must not drift the logits vs the unhooked + baseline; the study reports per-hook records and the eager verdict is + self-consistent (alert iff the bitwise check failed).""" + hooks = ["resid_pre", "final_logits"] + result = run_study( + framework="hf", model="qwen3", mode="eager", + hooks=hooks, out_dir=tmp_path, max_new_tokens=4, + ) + + # Always surface the table, even on pass. + print("\n" + format_table(result)) + + assert [h.hook for h in result.hooks] == hooks + json.dumps(result.to_dict()) # serializable + + for h in result.hooks: + assert h.error is None, f"{h.hook} rollout errored: {h.error}" + assert h.check is not None + # Eager verdict must be self-consistent with the bitwise outcome. + assert h.alert == (not h.check.passed), ( + f"{h.hook}: alert={h.alert} but check.passed={h.check.passed}" + ) + assert not result.any_alert, ( + "eager observational hooks drifted vs the unhooked baseline:\n" + + format_table(result) + ) diff --git a/tests/test_per_hook_isolation.py b/tests/test_per_hook_isolation.py index 97b77735b..e718d224b 100644 --- a/tests/test_per_hook_isolation.py +++ b/tests/test_per_hook_isolation.py @@ -61,8 +61,7 @@ _COPY_LINE_RE, _patched_source, compare_model_path, - invalidate_bytecode, - patch_compare_model, + isolated_hook, ) REPO_ROOT = Path(__file__).resolve().parent.parent @@ -89,6 +88,7 @@ def allocate(self): """).strip("\n") +@pytest.mark.cpu class TestPatcherCorrectness: """Verify the regex + line-by-line patching logic.""" @@ -126,6 +126,7 @@ def test_indentation_preserved_on_commented_lines(self): assert indent == 8, f"unexpected indent on: {line!r}" +@pytest.mark.cpu class TestPatcherRoundTrip: """Verify the on-disk patch / unpatch context manager.""" @@ -134,10 +135,10 @@ class TestPatcherRoundTrip: ("vllm", "gpt2"), ("vllm", "qwen3"), ("vllm", "llama"), ]) def test_round_trip_byte_identical(self, framework, model_key): - """File contents before and after patch_compare_model must match.""" + """File contents before and after isolated_hook must match.""" p = compare_model_path(framework, model_key) original = p.read_bytes() - with patch_compare_model(framework, model_key, "q") as (model_path, commented): + with isolated_hook(framework, model_key, "q") as (model_path, commented): assert p.read_bytes() != original, "patch did not modify the file" assert len(commented) > 0, "no _buf_* capture lines were commented" assert "q" not in commented, "q itself should not be in commented list" @@ -147,6 +148,28 @@ def test_round_trip_byte_identical(self, framework, model_key): backup = p.with_suffix(p.suffix + ".copy_isolate_backup") assert not backup.exists() + def test_dirty_restore_raises_loudly(self, tmp_path, monkeypatch): + """If the file can't be restored byte-identically, exit must raise. + + Simulates a body that clobbers the source *and* removes the backup + (the failure mode §6 hardens against): the context manager's + hash compare on exit must surface it as a loud RuntimeError rather + than silently leaving the vendored submodule dirty. + """ + from tests import isolate_hook + target = tmp_path / "fake_compare.py" + target.write_text(SAMPLE_COMPARE_SOURCE) + fake_paths = {("test", "fake"): target} + monkeypatch.setattr(isolate_hook, "_COMPARE_MODEL_PATHS", fake_paths) + + with pytest.raises(RuntimeError, match="byte-identically"): + with isolate_hook.isolated_hook("test", "fake", "q"): + # Corrupt the file and delete the backup so neither unpatch + # nor the snapshot can restore the original bytes. + target.write_text("corrupted contents that won't be restored") + backup = target.with_suffix(target.suffix + ".copy_isolate_backup") + backup.unlink() + def test_stale_backup_raises(self, tmp_path, monkeypatch): """If a previous run crashed mid-patch leaving a backup, refuse.""" from tests import isolate_hook @@ -163,6 +186,7 @@ def test_stale_backup_raises(self, tmp_path, monkeypatch): isolate_hook.patch("test", "fake", "q") +@pytest.mark.cpu class TestRealCompareModelsContainAllExpectedHooks: """The smoke cells assume specific hooks have a `.copy_()` line in the real _compare files. If a hook is missing the smoke fails opaquely, @@ -488,6 +512,7 @@ def _run_rollout( ) +@pytest.mark.gpu @pytest.mark.slow @pytest.mark.parametrize( "framework,model_key,hook,mode", SMOKE_CELLS, @@ -514,11 +539,10 @@ def test_per_hook_isolation_smoke( _run_rollout(tmp_path, "ours", framework, model_key, hook, mode, env) # Ref: _compare variant patched to isolate H. Patch persists for - # the lifetime of the with-block; the ``invalidate_bytecode`` call - # ensures the subprocess imports the patched source rather than a - # stale .pyc. - with patch_compare_model(framework, model_key, hook): - invalidate_bytecode(framework, model_key) + # the lifetime of the with-block; ``isolated_hook`` invalidates the + # cached bytecode so the subprocess imports the patched source rather + # than a stale .pyc, and asserts byte-identical restoration on exit. + with isolated_hook(framework, model_key, hook): _run_rollout(tmp_path, "ref", framework, model_key, hook, mode, env) L_orig = torch.load(tmp_path / "orig.pt", map_location="cpu") diff --git a/tests/test_producer_chunked_schema.py b/tests/test_producer_chunked_schema.py index bcc77f843..0508eb364 100644 --- a/tests/test_producer_chunked_schema.py +++ b/tests/test_producer_chunked_schema.py @@ -21,18 +21,35 @@ def setup_module(module): # noqa: D401 -- pytest hook _load_extension() # ensure .so loaded -> three ring ops registered +# --- Registration / wiring: no CUDA device required (CPU default suite) ------ + +@pytest.mark.cpu def test_producer_op_registered(): assert torch.ops.ring.producer.default is not None +@pytest.mark.cpu def test_producer_prefix_op_registered(): assert torch.ops.ring.producer_prefix.default is not None +@pytest.mark.cpu def test_producer_chunked_op_registered(): assert torch.ops.ring.producer_chunked.default is not None +@pytest.mark.cpu +def test_hook_point_strip_attrs_default_to_static(): + """HookPoint instances default to the static path.""" + from monitoring.hook_points import HookPoint + hp = HookPoint() + assert hp._strip_tensor is None + assert hp._strip_row_bytes == 0 + + +# --- Device smoke: allocate CUDA tensors and dispatch the op (GPU only) ------- + +@pytest.mark.gpu def test_producer_static_smoke(): """Static op accepts (Tensor(a!), Tensor, int, int). C++ impl early-returns when no engine is active, so this is a pure schema @@ -44,6 +61,7 @@ def test_producer_static_smoke(): torch.ops.ring.producer(ring_payload, x, 0, 0) +@pytest.mark.gpu def test_producer_prefix_smoke(): """Prefix op accepts (Tensor(a!), Tensor, Tensor, int, int, int).""" if not torch.cuda.is_available(): @@ -54,6 +72,7 @@ def test_producer_prefix_smoke(): torch.ops.ring.producer_prefix(ring_payload, x, row_count, 8, 0, 0) +@pytest.mark.gpu def test_producer_chunked_smoke(): """Chunked op accepts (Tensor(a!), Tensor, Tensor, int, int).""" if not torch.cuda.is_available(): @@ -64,14 +83,7 @@ def test_producer_chunked_smoke(): torch.ops.ring.producer_chunked(ring_payload, x, chunk_bytes, 0, 0) -def test_hook_point_strip_attrs_default_to_static(): - """HookPoint instances default to the static path.""" - from monitoring.hook_points import HookPoint - hp = HookPoint() - assert hp._strip_tensor is None - assert hp._strip_row_bytes == 0 - - +@pytest.mark.gpu def test_hook_point_strip_attrs_settable_for_prefix_mode(): """Setting _strip_tensor + _strip_row_bytes > 0 selects prefix mode.""" if not torch.cuda.is_available(): @@ -85,6 +97,7 @@ def test_hook_point_strip_attrs_settable_for_prefix_mode(): assert hp._strip_row_bytes == 8 +@pytest.mark.gpu def test_hook_point_strip_attrs_settable_for_chunked_mode(): """Setting _strip_tensor + _strip_row_bytes == 0 selects chunked mode.""" if not torch.cuda.is_available(): diff --git a/tests/test_ring_transport_no_framework_strings.py b/tests/test_ring_transport_no_framework_strings.py index cb5acf064..9bc8c7d82 100644 --- a/tests/test_ring_transport_no_framework_strings.py +++ b/tests/test_ring_transport_no_framework_strings.py @@ -25,6 +25,8 @@ import pytest +pytestmark = pytest.mark.cpu + RING_TRANSPORT = ( Path(__file__).resolve().parent.parent / "monitoring" / "ring_transport.py" ) diff --git a/tests/test_tp_shapes.py b/tests/test_tp_shapes.py index d56f6889c..c6a3ea441 100644 --- a/tests/test_tp_shapes.py +++ b/tests/test_tp_shapes.py @@ -6,6 +6,8 @@ import pytest import torch +pytestmark = pytest.mark.cpu + from monitoring.ring_transport import ( HOOK_TYPE_RESID_PRE, HOOK_TYPE_LN1, diff --git a/tests/test_vllm_identical.py b/tests/test_vllm_identical.py index 92f153ed6..41087a516 100644 --- a/tests/test_vllm_identical.py +++ b/tests/test_vllm_identical.py @@ -13,12 +13,12 @@ E2E_MAX_NEW_TOKENS Tokens to generate per prompt (default 20) E2E_ENFORCE_EAGER "1" to disable CUDA graphs (default "1") E2E_DTYPE Model dtype, e.g. "bfloat16", "float16", "auto" (default "bfloat16") - E2E_HOOKS Hook selection (default "vllm-full") E2E_REF_MAX_LEN Max first-dim for buffers (default 8192) E2E_MAX_NUM_BATCHED_TOKENS vLLM scheduler max_num_batched_tokens (default 512) E2E_RING_PAYLOAD_MB Ring payload size (default 4096) E2E_RING_PINNED_MB Pinned staging size (default 4096) - E2E_HOOK_SELECTION Hook selection for monitored run (default "vllm-full") + E2E_HOOK_SELECTION Public hook selection input (default "vllm-full"); + translated to DMX_HOOK_SELECTION for subprocesses DMX_DB_HOST ClickHouse host (default "localhost") DMX_DB_PORT ClickHouse port (default 9000) @@ -40,6 +40,13 @@ import pytest import torch +pytestmark = [ + pytest.mark.gpu, + pytest.mark.vllm, + pytest.mark.clickhouse, + pytest.mark.e2e, +] + _MODEL_REF_FILES = { "gpt2": "gpt2_ref.py", "qwen2_moe": "qwen2_moe_ref.py", @@ -54,7 +61,7 @@ def test_vllm_identical(subtests): """Bitwise comparison: ref model (disk) vs monitored model (ClickHouse).""" model_key = os.environ.get("E2E_MODEL", "gpt2") - hooks = os.environ.get("E2E_HOOKS", "vllm-full") + hooks = os.environ.get("E2E_HOOK_SELECTION", "vllm-full") max_len = int(os.environ.get("E2E_REF_MAX_LEN", "8192")) enforce_eager = os.environ.get("E2E_ENFORCE_EAGER", "1") diff --git a/tests/test_vllm_rowcnt.py b/tests/test_vllm_rowcnt.py index 13f1b1ceb..da0d466be 100644 --- a/tests/test_vllm_rowcnt.py +++ b/tests/test_vllm_rowcnt.py @@ -12,7 +12,8 @@ E2E_ENFORCE_EAGER "1" to disable torch.compile + CUDA graphs (default "0") E2E_RING_PAYLOAD_MB Ring payload size in MB (default 4096) E2E_RING_PINNED_MB Pinned staging size in MB (default 4096) - E2E_HOOK_SELECTION Hook selection preset (default "vllm-full") + E2E_HOOK_SELECTION Public hook selection preset (default "vllm-full"); + translated to DMX_HOOK_SELECTION for subprocesses E2E_COMPARE_LAYERS "all" or comma-separated layer IDs for value comparison. Requires model supported by extract_hidden_states. GPT-2 not supported -- value comparison skipped with warning. @@ -40,6 +41,13 @@ import pytest import torch +pytestmark = [ + pytest.mark.gpu, + pytest.mark.vllm, + pytest.mark.clickhouse, + pytest.mark.e2e, +] + _MODEL_ALIASES = { "gpt2": "gpt2", "qwen2_moe": "Qwen/Qwen1.5-MoE-A2.7B", @@ -54,6 +62,12 @@ def test_vllm_rowcnt(subtests): model_key = os.environ.get("E2E_MODEL", "gpt2") model_id = _MODEL_ALIASES.get(model_key, model_key) + # Translate the public E2E_HOOK_SELECTION input into the internal + # DMX_HOOK_SELECTION runtime contract that the runner actually reads. + sub_env = dict(os.environ) + sub_env["DMX_HOOK_SELECTION"] = os.environ.get( + "E2E_HOOK_SELECTION", "vllm-full") + run_dir = tempfile.mkdtemp(prefix="vllm_rowcnt_") ref_dir = os.path.join(run_dir, "ref") mon_dir = os.path.join(run_dir, "mon") @@ -76,7 +90,7 @@ def test_vllm_rowcnt(subtests): r2 = subprocess.run( [sys.executable, "-m", "tests.vllm_monitored_runner", "--output-dir", mon_dir], - env=os.environ, capture_output=True, text=True, cwd=project_root, + env=sub_env, capture_output=True, text=True, cwd=project_root, ) if r2.returncode != 0: pytest.fail(f"Monitored runner failed:\n{r2.stderr[-2000:]}") @@ -88,7 +102,7 @@ def test_vllm_rowcnt(subtests): "--ref-dir", ref_dir, "--mon-dir", mon_dir, "--result-file", result_file], - env=os.environ, capture_output=True, text=True, cwd=project_root, + env=sub_env, capture_output=True, text=True, cwd=project_root, ) if r3.returncode != 0: pytest.fail(f"Comparator failed:\n{r3.stderr[-2000:]}") diff --git a/tests/tools/README.md b/tests/tools/README.md new file mode 100644 index 000000000..e6225cd19 --- /dev/null +++ b/tests/tools/README.md @@ -0,0 +1,35 @@ +# tests/tools — manual analysis & release-sweep scripts + +These are **manual** entry points: debugging aids, transport/correctness sweeps, +and release-candidate regression wrappers. They are intentionally **not** part of +the pytest regression gates — `pyproject.toml` lists `tests/tools` under +`norecursedirs`, so pytest never discovers anything here. + +Run them by hand, from the **repository root**, when you want a full sweep or are +debugging a specific backend. They require a GPU (most also need ClickHouse and +the vLLM runtime); they are not CPU-safe. + +| Script | What it does | +|---|---| +| `run_regression.sh` | Full release sweep: CPU unit tests + HF/vLLM transport correctness across models/modes/TP. Calls the `run_tp_compare_*` wrappers. | +| `run_tp_compare_hf.sh` | Single HF transport-correctness run (`.copy_()` buffers vs ClickHouse) for one model/mode/TP. | +| `run_tp_compare_vllm.sh` | Single vLLM transport-correctness run for one model/mode/TP. | +| `run_qwen2_moe_vllm_pipeline.sh` | Qwen2-MoE / EP vLLM ref → monitored → compare pipeline. | +| `identical_vllm.sh` | Wrapper around the vLLM bitwise-identical pytest check. | +| `verify_vllm.sh` | vLLM row-count + identical verification sweep across ring sizes. | +| `verify_hf.sh` | HF E2E correctness sweep across ring sizes. | + +Example: + +```bash +# from the repo root +LD_PRELOAD=/path/to/libstdc++.so.6 CUDA_VISIBLE_DEVICES=0,1 \ + bash tests/tools/run_regression.sh +``` + +> Native CUDA ring tests live separately under `tests/ring/` (built via its +> `Makefile`, marker `ring_native`, needs `nvcc`) and are likewise excluded from +> default pytest discovery. + +As the configurable E2E matrix (`tests/e2e_matrix`) lands, these hardcoded +wrappers are expected to be superseded by matrix invocations. diff --git a/tests/identical_vllm.sh b/tests/tools/identical_vllm.sh similarity index 93% rename from tests/identical_vllm.sh rename to tests/tools/identical_vllm.sh index 444d6839a..2d9add082 100755 --- a/tests/identical_vllm.sh +++ b/tests/tools/identical_vllm.sh @@ -11,7 +11,7 @@ # - LD_PRELOAD for libstdc++ (caller's responsibility) # # Usage: -# LD_PRELOAD=/path/to/libstdc++.so.6 bash tests/identical_vllm.sh +# LD_PRELOAD=/path/to/libstdc++.so.6 bash tests/tools/identical_vllm.sh set -e export VLLM_DISABLE_COMPILE_CACHE=1 @@ -34,7 +34,7 @@ run_identical_test() { E2E_MODEL=$model_key \ E2E_ENFORCE_EAGER=$eager \ E2E_DTYPE=bfloat16 \ - E2E_HOOKS=vllm-full \ + E2E_HOOK_SELECTION=vllm-full \ E2E_REF_MAX_LEN=8192 \ E2E_RING_PAYLOAD_MB=$ring_mb \ E2E_RING_PINNED_MB=$ring_mb \ diff --git a/tests/run_qwen2_moe_vllm_pipeline.sh b/tests/tools/run_qwen2_moe_vllm_pipeline.sh similarity index 96% rename from tests/run_qwen2_moe_vllm_pipeline.sh rename to tests/tools/run_qwen2_moe_vllm_pipeline.sh index 7cc8437ff..8f3617d34 100755 --- a/tests/run_qwen2_moe_vllm_pipeline.sh +++ b/tests/tools/run_qwen2_moe_vllm_pipeline.sh @@ -1,7 +1,7 @@ #!/bin/bash set -euo pipefail -ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" CONDA_SH="${CONDA_SH:-/home/sixian/miniforge3/etc/profile.d/conda.sh}" CONDA_ENV_NAME="${CONDA_ENV_NAME:-ring_offload}" @@ -15,7 +15,7 @@ source "$CONDA_SH" conda activate "$CONDA_ENV_NAME" MODEL="${E2E_MODEL:-qwen2_moe}" -HOOKS="${E2E_HOOKS:-vllm-full}" +HOOKS="${E2E_HOOK_SELECTION:-vllm-full}" TP="${E2E_TP_SIZE:-2}" MAX_MODEL_LEN="${E2E_MAX_MODEL_LEN:-128}" REF_MAX_LEN="${E2E_REF_MAX_LEN:-512}" diff --git a/tests/run_regression.sh b/tests/tools/run_regression.sh similarity index 77% rename from tests/run_regression.sh rename to tests/tools/run_regression.sh index d179b2384..02c3ab69a 100755 --- a/tests/run_regression.sh +++ b/tests/tools/run_regression.sh @@ -1,6 +1,6 @@ #!/bin/bash # Full regression test suite. -# Usage: LD_PRELOAD=... CUDA_VISIBLE_DEVICES=0,1 bash tests/run_regression.sh +# Usage: LD_PRELOAD=... CUDA_VISIBLE_DEVICES=0,1 bash tests/tools/run_regression.sh set -e PASS=0 @@ -23,7 +23,7 @@ run_test() { fi } -PROJECT_ROOT=$(cd "$(dirname "$0")/.." && pwd) +PROJECT_ROOT=$(cd "$(dirname "$0")/../.." && pwd) cd "$PROJECT_ROOT" # --- Unit tests --- @@ -34,9 +34,9 @@ run_test "unit: test_tp_shapes + test_config" \ for model in qwen3 gpt2; do for mode in eager cudagraph; do run_test "vllm: $model $mode tp=1" \ - bash tests/run_tp_compare_vllm.sh "$model" "$mode" 1 + bash tests/tools/run_tp_compare_vllm.sh "$model" "$mode" 1 run_test "vllm: $model $mode tp=2" \ - bash tests/run_tp_compare_vllm.sh "$model" "$mode" 2 + bash tests/tools/run_tp_compare_vllm.sh "$model" "$mode" 2 done done @@ -44,13 +44,13 @@ done for model in gpt2 qwen3; do for mode in eager cudagraph; do run_test "hf: $model $mode tp=1" \ - bash tests/run_tp_compare_hf.sh "$model" "$mode" 1 + bash tests/tools/run_tp_compare_hf.sh "$model" "$mode" 1 done done # HF TP=2 (qwen3 only — gpt2 lacks tp_plan support) for mode in eager cudagraph; do run_test "hf: qwen3 $mode tp=2" \ - bash tests/run_tp_compare_hf.sh qwen3 "$mode" 2 + bash tests/tools/run_tp_compare_hf.sh qwen3 "$mode" 2 done # --- Summary --- diff --git a/tests/run_tp_compare_hf.sh b/tests/tools/run_tp_compare_hf.sh similarity index 95% rename from tests/run_tp_compare_hf.sh rename to tests/tools/run_tp_compare_hf.sh index df4f13ed0..745130895 100755 --- a/tests/run_tp_compare_hf.sh +++ b/tests/tools/run_tp_compare_hf.sh @@ -1,6 +1,6 @@ #!/bin/bash # HF transport correctness test: single run with compare model. -# Usage: bash tests/run_tp_compare_hf.sh [model] [mode] [tp] +# Usage: bash tests/tools/run_tp_compare_hf.sh [model] [mode] [tp] # model: gpt2 (default) or qwen3 # mode: eager (default) or cudagraph # tp: 1 (default) or 2 diff --git a/tests/run_tp_compare_vllm.sh b/tests/tools/run_tp_compare_vllm.sh similarity index 96% rename from tests/run_tp_compare_vllm.sh rename to tests/tools/run_tp_compare_vllm.sh index 06a27801e..f23143921 100755 --- a/tests/run_tp_compare_vllm.sh +++ b/tests/tools/run_tp_compare_vllm.sh @@ -1,6 +1,6 @@ #!/bin/bash # Transport correctness test: single run, compare .copy_() buffers vs ClickHouse. -# Usage: bash tests/run_tp_compare_vllm.sh [model] [mode] [tp] +# Usage: bash tests/tools/run_tp_compare_vllm.sh [model] [mode] [tp] # model: qwen3 (default) or gpt2 # mode: eager (default) or cudagraph # tp: 1 (default) or 2 diff --git a/tests/verify_hf.sh b/tests/tools/verify_hf.sh similarity index 100% rename from tests/verify_hf.sh rename to tests/tools/verify_hf.sh diff --git a/tests/verify_vllm.sh b/tests/tools/verify_vllm.sh similarity index 96% rename from tests/verify_vllm.sh rename to tests/tools/verify_vllm.sh index 877831985..6ae30bd92 100755 --- a/tests/verify_vllm.sh +++ b/tests/tools/verify_vllm.sh @@ -17,7 +17,7 @@ # - LD_PRELOAD for libstdc++ (caller's responsibility) # # Usage: -# LD_PRELOAD=/path/to/libstdc++.so.6 bash tests/verify_vllm.sh +# LD_PRELOAD=/path/to/libstdc++.so.6 bash tests/tools/verify_vllm.sh set -e export VLLM_DISABLE_COMPILE_CACHE=1 @@ -83,7 +83,7 @@ run_identical_test() { E2E_MODEL=$model_key \ E2E_ENFORCE_EAGER=$eager \ E2E_DTYPE=bfloat16 \ - E2E_HOOKS=vllm-full \ + E2E_HOOK_SELECTION=vllm-full \ E2E_REF_MAX_LEN=8192 \ E2E_RING_PAYLOAD_MB=$ring_mb \ E2E_RING_PINNED_MB=$ring_mb \