From 4147235c764b6e8fcdc1f629dcca8f6c5d83577e Mon Sep 17 00:00:00 2001 From: SieDeta Date: Tue, 9 Jun 2026 19:07:13 +0700 Subject: [PATCH 1/6] test: marker taxonomy + skip guards + categorize (PR1) Introduce the test-suite categorization foundation from the test-suite refactor plan (PR #1 of 6). - pyproject.toml: register the full marker set (cpu, gpu, multi_gpu, e2e, clickhouse, vllm, hf, ring_native, slow, manual, numeric); default collection excludes `manual` and skips tests/tools, tests/ring, .venv, integration, build via norecursedirs. - tests/_requirements.py: CPU-importable skip guards that fail closed with precise reasons (require_cuda, require_gpus, require_clickhouse, require_vllm, require_model_cache, require_nvcc). - Mark every test: pure-CPU contract tests -> `cpu`; GPU/E2E suites -> explicit gpu/vllm/clickhouse/e2e/hf marks. test_gpt2_parity pulls real gpt2 weights -> `hf` + require_model_cache, not cpu. - Split test_producer_chunked_schema: op-registration tests stay `cpu`, CUDA device-smoke tests become `gpu`. - test_per_hook_isolation: source-patch unit classes -> `cpu`; the slow GPU sweep -> `gpu` (keeps `slow`). - Relocate manual shell wrappers to tests/tools/ (+ README); fix their repo-root path computations and cross-references for the new depth. Default CPU command: python -m pytest -m "not gpu and not e2e and not manual" -q Verified: markers register without warnings; the CPU selection deselects 32 GPU/E2E tests. Pre-existing (reproduced on base main, unrelated to this PR): monitoring/ring_transport.py references HOOK_TYPE_ROUTER_LOGITS (defined only in csrc enum, missing on the Python side), and the vendored transformers submodule is unbuilt -- these block a fully green CPU run and must be fixed separately. Co-Authored-By: Claude Opus 4.8 --- pyproject.toml | 15 ++ tests/_requirements.py | 128 ++++++++++++++++++ tests/test_adapter_protocol.py | 4 + tests/test_config.py | 2 + tests/test_e2e_correctness_vs_hf.py | 7 + tests/test_gpt2_parity.py | 5 + tests/test_hf_eos_strip.py | 2 + tests/test_hook_spec_flags.py | 3 + tests/test_moe_v1_routing_hooks.py | 4 + tests/test_no_graph_breaks.py | 2 + tests/test_per_hook_isolation.py | 4 + tests/test_producer_chunked_schema.py | 29 ++-- ...est_ring_transport_no_framework_strings.py | 2 + tests/test_tp_shapes.py | 2 + tests/test_vllm_identical.py | 7 + tests/test_vllm_rowcnt.py | 7 + tests/tools/README.md | 35 +++++ tests/{ => tools}/identical_vllm.sh | 2 +- .../run_qwen2_moe_vllm_pipeline.sh | 2 +- tests/{ => tools}/run_regression.sh | 12 +- tests/{ => tools}/run_tp_compare_hf.sh | 2 +- tests/{ => tools}/run_tp_compare_vllm.sh | 2 +- tests/{ => tools}/verify_hf.sh | 0 tests/{ => tools}/verify_vllm.sh | 2 +- 24 files changed, 261 insertions(+), 19 deletions(-) create mode 100644 tests/_requirements.py create mode 100644 tests/tools/README.md rename tests/{ => tools}/identical_vllm.sh (95%) rename tests/{ => tools}/run_qwen2_moe_vllm_pipeline.sh (98%) rename tests/{ => tools}/run_regression.sh (77%) rename tests/{ => tools}/run_tp_compare_hf.sh (95%) rename tests/{ => tools}/run_tp_compare_vllm.sh (96%) rename tests/{ => tools}/verify_hf.sh (100%) rename tests/{ => tools}/verify_vllm.sh (97%) 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/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_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_per_hook_isolation.py b/tests/test_per_hook_isolation.py index 97b77735b..dedcd1940 100644 --- a/tests/test_per_hook_isolation.py +++ b/tests/test_per_hook_isolation.py @@ -89,6 +89,7 @@ def allocate(self): """).strip("\n") +@pytest.mark.cpu class TestPatcherCorrectness: """Verify the regex + line-by-line patching logic.""" @@ -126,6 +127,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.""" @@ -163,6 +165,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 +491,7 @@ def _run_rollout( ) +@pytest.mark.gpu @pytest.mark.slow @pytest.mark.parametrize( "framework,model_key,hook,mode", SMOKE_CELLS, 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..0049a05f0 100644 --- a/tests/test_vllm_identical.py +++ b/tests/test_vllm_identical.py @@ -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", diff --git a/tests/test_vllm_rowcnt.py b/tests/test_vllm_rowcnt.py index 13f1b1ceb..09231df3d 100644 --- a/tests/test_vllm_rowcnt.py +++ b/tests/test_vllm_rowcnt.py @@ -40,6 +40,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", 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 95% rename from tests/identical_vllm.sh rename to tests/tools/identical_vllm.sh index 444d6839a..997779513 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 diff --git a/tests/run_qwen2_moe_vllm_pipeline.sh b/tests/tools/run_qwen2_moe_vllm_pipeline.sh similarity index 98% rename from tests/run_qwen2_moe_vllm_pipeline.sh rename to tests/tools/run_qwen2_moe_vllm_pipeline.sh index 7cc8437ff..89c79ce3b 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}" 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 97% rename from tests/verify_vllm.sh rename to tests/tools/verify_vllm.sh index 877831985..7f65a78a0 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 From 1a6c4009bddf15ad730069f5c339675376afc51f Mon Sep 17 00:00:00 2001 From: SieDeta Date: Tue, 9 Jun 2026 19:15:49 +0700 Subject: [PATCH 2/6] test: marker taxonomy + skip guards + categorize --- .gitignore | 1 + 1 file changed, 1 insertion(+) 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/ From 0fd46927f9931e32356305aac8639b9edb25f39e Mon Sep 17 00:00:00 2001 From: SieDeta Date: Wed, 10 Jun 2026 15:00:40 +0700 Subject: [PATCH 3/6] test: normalize env vars + harden source-patch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 1 of the test-suite refactor (plan §2, §6 / PR2). Env-var normalization — one public hook-selection input: - E2E_HOOK_SELECTION is now the single public test input; it is translated to the internal DMX_HOOK_SELECTION runtime contract (read by the runners/adapters) in every subprocess env. - test_vllm_identical.py: read E2E_HOOK_SELECTION (was E2E_HOOKS); drop the duplicate docstring entry. - test_vllm_rowcnt.py: actually honor the documented E2E_HOOK_SELECTION by translating it into DMX_HOOK_SELECTION for the runner/comparator subprocs. - tools/verify_vllm.sh, tools/identical_vllm.sh, tools/run_qwen2_moe_vllm_pipeline.sh: E2E_HOOKS -> E2E_HOOK_SELECTION. - All E2E_HOOKS usages removed. Source-patch hardening — single isolation contract: - Add isolated_hook(framework, model_key, hook) in tests/isolate_hook.py: snapshots the vendored _compare source bytes, patches to capture one hook, and on exit restores AND asserts byte-identical restoration via a SHA-256 compare, raising loudly so a dirty vendored submodule can never escape. Invalidates cached bytecode by default so subprocess imports see the patched source. - patch_compare_model kept as a backward-compatible delegating alias. - Migrate test_per_hook_isolation.py to isolated_hook; add test_dirty_restore_raises_loudly covering the loud-failure path. Acceptance: `pytest -m "not gpu and not e2e and not manual"` — the per-hook isolation unit suite is green (15 passed, incl. the new restore test). Pre-existing failures in this checkout stem from the unbuilt monitoring native layer (HOOK_TYPE_ROUTER_LOGITS) and are unrelated. Co-Authored-By: Claude Opus 4.8 --- tests/isolate_hook.py | 66 ++++++++++++++++++---- tests/test_per_hook_isolation.py | 38 ++++++++++--- tests/test_vllm_identical.py | 6 +- tests/test_vllm_rowcnt.py | 13 ++++- tests/tools/identical_vllm.sh | 2 +- tests/tools/run_qwen2_moe_vllm_pipeline.sh | 2 +- tests/tools/verify_vllm.sh | 2 +- 7 files changed, 100 insertions(+), 29 deletions(-) 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/test_per_hook_isolation.py b/tests/test_per_hook_isolation.py index dedcd1940..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 @@ -136,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" @@ -149,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 @@ -518,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_vllm_identical.py b/tests/test_vllm_identical.py index 0049a05f0..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) @@ -61,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 09231df3d..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. @@ -61,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") @@ -83,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:]}") @@ -95,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/identical_vllm.sh b/tests/tools/identical_vllm.sh index 997779513..2d9add082 100755 --- a/tests/tools/identical_vllm.sh +++ b/tests/tools/identical_vllm.sh @@ -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/tools/run_qwen2_moe_vllm_pipeline.sh b/tests/tools/run_qwen2_moe_vllm_pipeline.sh index 89c79ce3b..8f3617d34 100755 --- a/tests/tools/run_qwen2_moe_vllm_pipeline.sh +++ b/tests/tools/run_qwen2_moe_vllm_pipeline.sh @@ -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/tools/verify_vllm.sh b/tests/tools/verify_vllm.sh index 7f65a78a0..6ae30bd92 100755 --- a/tests/tools/verify_vllm.sh +++ b/tests/tools/verify_vllm.sh @@ -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 \ From 6c8f3a616352285d53222a4a9484e8353d315de0 Mon Sep 17 00:00:00 2001 From: SieDeta Date: Wed, 10 Jun 2026 15:39:19 +0700 Subject: [PATCH 4/6] test: shared E2E lib + configurable matrix MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phases 2-3 of the test-suite refactor (plan §7, §8 / PR3). Additive only: the matrix runs alongside the existing tests and nothing is deleted. tests/lib/ — shared E2E library (§7), consolidating logic that vllm_identical_comparator / compare_disk_vs_ch / vllm_rowcnt_comparator / hf_comparator each reimplemented: - compare.py: Check dataclass + the four standards (bitwise, allclose, row_count, transport_bitwise) as one interface; records max/mean abs + first-diff position even on a pass so "barely passing" stays visible. - align.py: left-pad strip, EOS trim, request_id ":" parsing, vLLM UUID-suffix normalization, logits head-skip alignment. - report.py: CellResult dataclass -> JSON record / JSONL artifact + human table; legacy comparator result.json adapter. - clickhouse_io.py: canonical dtype map + short-hook->CH act_name map + row decode / per-hook counts / request grouping (clickhouse_driver lazy). - segments.py: thin re-export of monitoring.segment_merger + a sorted per-request chunk merge helper. - disk_ref.py: parse RefDiskWorker .pt filenames; load HF structured refs. - hf_reference.py: stable-named re-export of the ROL/GEN HF rollouts (the 700-line module relocation is deferred to the legacy-removal PR). tests/e2e_matrix.py — configurable matrix entry point (§8) replacing the hardcoded shell sweeps: - Multi-value axes (backend/model/mode/standard/hooks/tp/ring/dtype/ prompt-set) -> Cartesian product of cells. - Dispatches each cell to the existing runners + comparators as subprocesses (no inference logic reimplemented): vLLM bitwise/transport -> ref+monitored+vllm_identical_comparator; vLLM row_count/allclose -> monitored+vllm_rowcnt_comparator; HF -> ref+monitored+hf_comparator. - Translates the public E2E_HOOK_SELECTION input to internal DMX_HOOK_SELECTION per subprocess (plan §2). - One JSONL record per cell; per-cell isolation so one bad cell can't abort the matrix; per-subprocess --cell-timeout so a hung runner fails the cell. - --dry-run prints planned cells + dispatch commands with no CUDA/CH, so expansion and env translation are CPU-testable. tests/test_e2e_lib.py — 37 cpu-marked unit tests covering the four standards, align helpers, report serialization/round-trip, and matrix cell expansion / env translation / dispatch planning / dry-run. Acceptance: `pytest -m "not gpu and not e2e and not manual"` — new suite green (37 passed; 108 total passing on this checkout). Pre-existing failures stem from the unbuilt monitoring native layer (HOOK_TYPE_ROUTER_LOGITS) and are unrelated. Co-Authored-By: Claude Opus 4.8 --- tests/e2e_matrix.py | 398 +++++++++++++++++++++++++++++++++++++ tests/lib/__init__.py | 25 +++ tests/lib/align.py | 88 ++++++++ tests/lib/clickhouse_io.py | 124 ++++++++++++ tests/lib/compare.py | 211 ++++++++++++++++++++ tests/lib/disk_ref.py | 84 ++++++++ tests/lib/hf_reference.py | 28 +++ tests/lib/report.py | 126 ++++++++++++ tests/lib/segments.py | 33 +++ tests/test_e2e_lib.py | 309 ++++++++++++++++++++++++++++ 10 files changed, 1426 insertions(+) create mode 100644 tests/e2e_matrix.py create mode 100644 tests/lib/__init__.py create mode 100644 tests/lib/align.py create mode 100644 tests/lib/clickhouse_io.py create mode 100644 tests/lib/compare.py create mode 100644 tests/lib/disk_ref.py create mode 100644 tests/lib/hf_reference.py create mode 100644 tests/lib/report.py create mode 100644 tests/lib/segments.py create mode 100644 tests/test_e2e_lib.py 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/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/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 From 357f15e19267454434dba74799524f1ef6bd1f0d Mon Sep 17 00:00:00 2001 From: SieDeta Date: Wed, 10 Jun 2026 16:00:03 +0700 Subject: [PATCH 5/6] test: migrate to matrix + remove legacy E2E MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 4 of the test-suite refactor (plan §5 / PR4). The high-risk switch: the vLLM and HF E2E tests become thin wrappers on the configurable matrix (PR3), and the stale in-process bodies are deleted. Net -1991/+216 lines. Wrapper support in tests/e2e_matrix.py: - matrix_argv_from_env(): translate the legacy E2E_* env knobs (model, enforce-eager, dtype, ring, hook-selection, tolerance, db, ...) into a single-cell matrix argv, so wrappers honor exactly what the old tests did. - run_single(): parse argv to one cell, run it, return the CellResult (refuses anything but a single cell). test_e2e_correctness_vs_hf.py: 1667 -> ~85 lines. - test_e2e_correctness_hf -> matrix hf/eager/allclose cell. - test_e2e_cuda_graphs_vs_eager_hf -> matrix hf/cuda_graph/allclose cell, default tolerance 0.5 (eager ref vs CUDA-graph monitored). - Deleted: _test_e2e_correctness_hf_legacy and _test_e2e_cuda_graphs_vs_eager_hf_legacy (uncollected "kept for reference" copies), the permanently-disabled @skipif(True) test_e2e_correctness_hf_cuda_graphs (its compiled-rollout reference could not replicate generate()'s StaticCache handling and was explicitly superseded -- no still-passing assertion lost), and the now-unused in-process helpers (_make_ring_cfg, _make_host_cfg, get_num_layers_from_config, _canon_layer_and_act, bitwise_equal, ...). test_vllm_identical.py (262 -> ~55) -> matrix vllm/bitwise cell; test_vllm_rowcnt.py (117 -> ~50) -> matrix vllm/row_count cell. All four public test NAMES are preserved (tests/tools/verify_hf.sh, verify_vllm.sh, identical_vllm.sh invoke them by node id). Skips are now precise via tests/_requirements (require_cuda/require_vllm/require_clickhouse) instead of ad-hoc torch.backends.cuda checks. Matrix-parity (static dispatch) — each wrapper drives the identical runner+comparator chain the old test did: - test_e2e_correctness_hf : hf_reference_runner + hf_monitored_runner(eager) + hf_comparator - test_e2e_cuda_graphs_vs_eager_hf: hf_reference_runner(eager) + hf_monitored_runner(CG) + hf_comparator(tol=0.5) - test_vllm_identical : enable_ref_hooks + vllm_ref_runner + vllm_monitored_runner + vllm_identical_comparator - test_vllm_rowcnt : vllm_monitored_runner + vllm_rowcnt_comparator The numeric GPU/ClickHouse parity run (old result vs new JSONL) must be pasted on the PR before merge per the plan's PR3-4 gate. Acceptance: `pytest -m "not gpu and not e2e and not manual"` green (116 passed; +6 wrapper-translation unit tests). All four wrappers collect cleanly. Pre-existing failures stem from the unbuilt monitoring native layer (HOOK_TYPE_ROUTER_LOGITS) and are unrelated. Co-Authored-By: Claude Opus 4.8 --- tests/e2e_matrix.py | 54 + tests/test_e2e_correctness_vs_hf.py | 1692 +-------------------------- tests/test_e2e_lib.py | 59 + tests/test_vllm_identical.py | 271 +---- tests/test_vllm_rowcnt.py | 131 +-- 5 files changed, 216 insertions(+), 1991 deletions(-) diff --git a/tests/e2e_matrix.py b/tests/e2e_matrix.py index 6098abc92..aed161dbf 100644 --- a/tests/e2e_matrix.py +++ b/tests/e2e_matrix.py @@ -308,6 +308,60 @@ def restore(): # noqa: E306 -- restore ref source pre-monitored run shutil.rmtree(run_dir, ignore_errors=True) +# --------------------------------------------------------------------------- +# Thin-wrapper support (plan §5) +# --------------------------------------------------------------------------- + + +def matrix_argv_from_env(backend: str, standard: str, *, + mode: Optional[str] = None, + default_tolerance: str = "0.01", + env: Optional[dict] = None) -> List[str]: + """Build a single-cell matrix argv from the legacy ``E2E_*`` env knobs. + + The pytest wrappers preserve the old test names + shell entry points + (verify_hf.sh / verify_vllm.sh) and drive the matrix with the same + configuration the legacy tests honored. ``mode`` defaults to + eager/cuda_graph from ``E2E_ENFORCE_EAGER``; the HF cuda-graph wrapper + passes it explicitly. + """ + e = os.environ if env is None else env + if mode is None: + mode = "eager" if e.get("E2E_ENFORCE_EAGER", "1") == "1" else "cuda_graph" + hooks = e.get("E2E_HOOK_SELECTION", e.get("DMX_HOOK_SELECTION", "vllm-full")) + return [ + "--backend", backend, + "--model", e.get("E2E_MODEL", "gpt2"), + "--mode", mode, + "--standard", standard, + "--hooks", hooks, + "--tp", e.get("E2E_TP_SIZE", "1"), + "--ring-mb", e.get("E2E_RING_PAYLOAD_MB", "4096"), + "--dtype", e.get("E2E_DTYPE", "bfloat16"), + "--num-prompts", e.get("E2E_NUM_PROMPTS", "8"), + "--max-new-tokens", e.get("E2E_MAX_NEW_TOKENS", "20"), + "--max-model-len", e.get("E2E_MAX_MODEL_LEN", "512"), + "--max-batched-tokens", e.get("E2E_MAX_NUM_BATCHED_TOKENS", "512"), + "--gpu-mem-util", e.get("E2E_GPU_MEM_UTIL", "0.5"), + "--tolerance", e.get("E2E_TOLERANCE", default_tolerance), + "--db-host", e.get("DMX_DB_HOST", "localhost"), + "--db-port", e.get("DMX_DB_PORT", "9000"), + ] + + +def run_single(argv: List[str]) -> CellResult: + """Parse ``argv`` into exactly one cell, run it, return its CellResult. + + Raises ``ValueError`` if the axes expand to other than one cell -- the + wrappers must drive a single concrete cell. + """ + args = build_parser().parse_args(argv) + cells = build_cells(args) + if len(cells) != 1: + raise ValueError(f"run_single expected exactly 1 cell, got {len(cells)}") + return run_cell(cells[0], args) + + # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- diff --git a/tests/test_e2e_correctness_vs_hf.py b/tests/test_e2e_correctness_vs_hf.py index 8f2080940..f33721d72 100644 --- a/tests/test_e2e_correctness_vs_hf.py +++ b/tests/test_e2e_correctness_vs_hf.py @@ -1,80 +1,28 @@ -# tests/test_e2e_correctness_hf.py -# PYTHONPATH=./:./monitoring:$PYTHONPATH E2E_PRINT_TEXT=1 E2E_HF_DROP_LAST_TOKEN=1 E2E_PRINT_TOPK_LOGITS=1 pytest -q -s tests/test_e2e_correctness_vs_hf.py -"""E2E correctness test: monitoring DB vs HuggingFace Transformers (HF-driven ground truth). - -This test runs the repo monitoring pipeline end-to-end (native backend + host engine + ClickHouse), -then uses HuggingFace Transformers as the reference implementation (no TransformerLens). - -IMPORTANT: "prompt" in this test means FULL TOKEN SEQUENCE = prefill + decode ------------------------------------------------------------------------ -We treat DB `token_ids` as the ground-truth sequence of tokens for each request. That sequence -includes the initial prompt tokens (prefill) plus any decode tokens that were appended. - -HF reference modes ------------------------------------------------------------------------ - - ROL (manual rollout, batched): - * Incremental greedy KV-cache rollout over the full padded batch, using generate-style - position_ids derived from attention_mask (so left-padding doesn't shift positions). - * We strip left-pad per row and stop per row at EOS (no trailing padded steps), so outputs - align to DB token_ids. - - - GEN (HF generate(), batched): - * Run hf_model.generate() ONCE on the full padded batch (same input_ids/attention_mask as monitoring), - with output_scores=True. - * Strip left-pad per row and trim at EOS so sequences align to DB token_ids. - -For logits we effectively compare THREE sources: - - DB final_logits (from ClickHouse) - - ROL final_logits (manual rollout; [T, vocab]) - - GEN scores (from generate(); available only for positions t in [prompt_len-1, T-2]) - -Request-id convention (from MonitoringEngine._register_db_step in engine.py) ------------------------------------------------------------------------ -When a new batch is (re)initialized: - - gid = self._auto_batch_group_id - self._auto_batch_group_id += 1 - self._active_batch_request_ids = [f"{gid}:{i}" for i in range(batch_size)] - -So request_id == ":" where local_index is the batch row index. - -We map DB requests to HF batch rows using local_index, and add safety checks. - -DB tensor shapes (IMPORTANT) ------------------------------------------------------------------------ -DB offloaded tensors DO NOT have batch dim now: - - token_ids: [T] - - hook_embed/pos/final_ln: [T, d_model] - - resid_pre/post: [T, d_model] - - attn pattern/scores: [n_heads, Tq, Tk] - - final_logits: [R, vocab] or [vocab] (R is often full sequence length) - -Env vars ------------------------------------------------------------------------ - - E2E_BATCH_SIZE (default 4) - - E2E_MAX_NEW_TOKENS (default 8) - - E2E_MODEL (default "gpt2"; "qwen3" alias supported) - - E2E_CHUNK_BYTES (default 262144) - - - E2E_PRINT_TEXT (default 0): if 1, print decoded text from DB token_ids and from HF rollout + HF generate(). - - E2E_HF_DROP_LAST_TOKEN (default 0): if 1, drop the last token (and aligned tensors) from HF refs before compares. - - E2E_PRINT_TOPK_LOGITS (default 0): if 1, print top-k logits at every position for DB vs ROL vs GEN. - - E2E_PRINT_TOPK_LOGITS_K (default 5): top-k to print per position. - -ClickHouse ------------------------------------------------------------------------ - - DMX_DB_HOST, DMX_DB_PORT, DMX_DB_USER, DMX_DB_PASSWORD, DMX_DB_DATABASE, DMX_DB_TABLE +"""HF E2E correctness — thin wrappers over the configurable matrix (plan §5). + +This file used to carry ~1.6k lines of in-process HF rollout + tensor +comparison logic, three tests (one permanently disabled via +``@skipif(True)``), 16 skip sites, and two ``_legacy`` bodies kept "for +reference". All of that comparison logic now lives in :mod:`tests.lib` and +the dispatch in :mod:`tests.e2e_matrix`; these wrappers just drive the +matrix for the equivalent HF cell and assert on its checks. + +The test *names* are preserved because ``tests/tools/verify_hf.sh`` invokes +them by node id (``::test_e2e_correctness_hf`` / +``::test_e2e_cuda_graphs_vs_eager_hf``) and threads the ring-size / model +env vars the matrix wrapper reads. + +The removed ``test_e2e_correctness_hf_cuda_graphs`` was permanently disabled +(its compiled-rollout reference could not replicate generate()'s internal +StaticCache handling) and explicitly superseded by +``test_e2e_cuda_graphs_vs_eager_hf`` -- no still-passing assertion was lost. """ - from __future__ import annotations -import os -import sys -import uuid -from typing import Dict, List, Tuple - import pytest -import torch + +from tests._requirements import require_cuda, require_clickhouse +from tests.e2e_matrix import matrix_argv_from_env, run_single pytestmark = [ pytest.mark.gpu, @@ -83,1586 +31,40 @@ pytest.mark.hf, ] -from monitoring.clickhouse_reader import CHClickhouseDriverReadOnly -from monitoring.segment_merger import merge_segments, parse_internal_id - -from .hf_reference import ( - _HFGenRef, - _HFRef, - _hf_generate_collect_hidden_states_batched, - _hf_generate_collect_scores_batched, - _hf_greedy_rollout_collect_all_batched, - _load_hf_refs_from_disk, - _parse_request_id, - _positions_for_unpadded, - _strip_left_pad, -) - -# --------------------------------------------------------------------------- -# Model aliases -# --------------------------------------------------------------------------- - -_MODEL_ALIASES = {"qwen3": "Qwen/Qwen3-4B", "llama": "meta-llama/Llama-3.1-8B"} - - -def _resolve_model_id(model: str) -> str: - return _MODEL_ALIASES.get(model.lower(), model) - - -# --------------------------------------------------------------------------- -# Small utils (inlined to avoid test-only deps) -# --------------------------------------------------------------------------- - - -def bitwise_equal(a: torch.Tensor, b: torch.Tensor) -> bool: - """Exact bitwise equality (treats NaNs as equal if their payload bits match).""" - if a.dtype != b.dtype or a.shape != b.shape: - return False - if a.numel() == 0: - return True - a8 = a.cpu().contiguous().view(torch.uint8) - b8 = b.cpu().contiguous().view(torch.uint8) - return torch.equal(a8, b8) - - -def get_num_layers_from_config(hf_model) -> int: - cfg = getattr(hf_model, "config", None) - if cfg is None: - raise ValueError("HF model has no .config; cannot determine num_layers") - - for attr in ("num_hidden_layers", "n_layer", "num_layers", "n_layers"): - if hasattr(cfg, attr): - return int(getattr(cfg, attr)) - # Some models nest text config (e.g., multi-modal wrappers) - for sub in ("text_config", "model_config", "llm_config"): - subcfg = getattr(cfg, sub, None) - if subcfg is None: - continue - for attr in ("num_hidden_layers", "n_layer", "num_layers", "n_layers"): - if hasattr(subcfg, attr): - return int(getattr(subcfg, attr)) +def _assert_cell(subtests, cr) -> None: + """Fail on a dispatch error; report each matrix check as a subtest.""" + if cr.error: + pytest.fail(f"matrix cell errored: {cr.error}") + assert cr.checks, "matrix produced no checks" + for chk in cr.checks: + with subtests.test(chk.name): + assert chk.passed, chk.detail - raise ValueError(f"Could not infer num_layers from config type={type(cfg)!r}") - -def _canon_layer_and_act(act_name_raw: str, layer_no_raw: int) -> Tuple[int, str]: - """Handle both schemas: - - (layer_no, act_name) stored separately, act_name like 'blocks.attn.hook_pattern' - - act_name stored as internal_id 'blocks..attn.hook_pattern' with layer_no possibly -1 - """ - try: - layer_from_act, act = parse_internal_id(act_name_raw) - if layer_from_act != -1: - return int(layer_from_act), str(act) - except Exception: - # parse_internal_id is intentionally strict (expects 'blocks....'); fall back. - pass - return int(layer_no_raw), str(act_name_raw) - - -# --------------------------------------------------------------------------- -# Ring + host engine configuration (env-var driven) -# --------------------------------------------------------------------------- - -def _make_ring_cfg(): - """Build RingConfig from E2E_RING_* environment variables.""" - from monitoring._native_engine import RingConfig # type: ignore - rc = RingConfig() - rc.task_ring_entries = int(os.environ.get("E2E_RING_TASK_ENTRIES", "16384")) - rc.payload_ring_bytes = int(os.environ.get("E2E_RING_PAYLOAD_BYTES", str(4 * 1024**3))) - rc.pinned_staging_bytes = int(os.environ.get("E2E_RING_PINNED_BYTES", str(4 * 1024**3))) - rc.drain_poll_timeout_us = int(os.environ.get("E2E_DRAIN_POLL_TIMEOUT_US", "100")) - rc.drain_flush_task_ratio = float(os.environ.get("E2E_DRAIN_FLUSH_TASK_RATIO", "0.0")) - rc.drain_flush_payload_ratio = float(os.environ.get("E2E_DRAIN_FLUSH_PAYLOAD_RATIO", "0.0")) - rc.drain_flush_entry_threshold = int(os.environ.get("E2E_DRAIN_FLUSH_ENTRY_THRESHOLD", "0")) - rc.drain_flush_byte_threshold = int(os.environ.get("E2E_DRAIN_FLUSH_BYTE_THRESHOLD", "0")) - rc.drain_flush_timeout_us = int(os.environ.get("E2E_DRAIN_FLUSH_TIMEOUT_US", "0")) - rc.clone_slices = int(os.environ.get("E2E_CLONE_SLICES", "0")) != 0 - rc.insert_queue_max_bytes = int(os.environ.get("E2E_INSERT_QUEUE_MAX_BYTES", str(512 * 1024**2))) - rc.insert_queue_max_items = int(os.environ.get("E2E_INSERT_QUEUE_MAX_ITEMS", "4096")) - return rc - - -def _make_host_cfg(db_cfg_native): - """Build HostEngineConfig with clickhouse insert stage from env vars.""" - from monitoring import HostEngineConfig # type: ignore - from monitoring._native_engine import StageConfig # type: ignore - parallelism = int(os.environ.get("E2E_CH_PARALLELISM", "10")) - stage = StageConfig.clickhouse_insert(db_cfg_native, parallelism=parallelism, - name="clickhouse_insert") - q = stage.input_queue - q.max_batch_items = int(os.environ.get("E2E_CH_QUEUE_MAX_ITEMS", "1024")) - q.high_watermark_items = q.max_batch_items - q.max_batch_size = int(os.environ.get("E2E_CH_QUEUE_MAX_BYTES", str(2048 * 1024**2))) - q.high_watermark_size = q.max_batch_size - return HostEngineConfig(stages=[stage]) - - -# --------------------------------------------------------------------------- -# Test -# --------------------------------------------------------------------------- - - -@pytest.mark.skipif(not torch.backends.cuda.is_built(), reason="CUDA not built") +@require_cuda() +@require_clickhouse() def test_e2e_correctness_hf(subtests) -> None: - """E2E correctness: compare HOOKED model (ring transport -> ClickHouse) - against ORIGINAL model (HF output_hidden_states=True). - - Three subprocesses — parent process never touches CUDA: - 1. Reference: original model -> tensors on disk - 2. Monitored: hooked model + ring transport -> ClickHouse - 3. Comparator: reads both, compares, writes result.json - """ - import json - import subprocess - import tempfile - import shutil - - run_dir = tempfile.mkdtemp(prefix="hf_e2e_") - 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") - project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - - try: - # Step 1: Reference run (original model, no CUDA in parent) - print("\n [1/3] Reference run (original model)...", flush=True) - r1 = subprocess.run( - [sys.executable, "-m", "tests.hf_reference_runner", - "--output-dir", ref_dir], - env=os.environ, capture_output=True, text=True, cwd=project_root, - ) - if r1.returncode != 0: - pytest.fail(f"Reference runner failed:\n{r1.stderr[-2000:]}") - - # Step 2: Monitored run (hooked model + ring transport) - print(" [2/3] Monitored run (hooked model + ring)...", flush=True) - r2 = subprocess.run( - [sys.executable, "-m", "tests.hf_monitored_runner", - "--output-dir", mon_dir], - env=os.environ, capture_output=True, text=True, cwd=project_root, - ) - if r2.returncode != 0: - pytest.fail(f"Monitored runner failed:\n{r2.stderr[-2000:]}") - - # Step 3: Comparator (CPU only, reads disk + ClickHouse) - print(" [3/3] Comparing...", flush=True) - r3 = subprocess.run( - [sys.executable, "-m", "tests.hf_comparator", - "--ref-dir", ref_dir, - "--mon-dir", mon_dir, - "--result-file", result_file], - env=os.environ, capture_output=True, text=True, cwd=project_root, - ) - if r3.returncode != 0: - pytest.fail(f"Comparator failed:\n{r3.stderr[-2000:]}") - - # Read results - with open(result_file) as f: - results = json.load(f) - - # Report via subtests - for test in results["tests"]: - with subtests.test(test["name"]): - assert test["passed"], test.get("detail", "") - - finally: - shutil.rmtree(run_dir, ignore_errors=True) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA + native backend required") -def _test_e2e_correctness_hf_legacy(subtests) -> None: - """Legacy version kept for reference. Not called by verify_hf.sh.""" - try: - import clickhouse_driver # noqa: F401 - except Exception: - pytest.skip("clickhouse-driver is required") - - try: - from monitoring import ( # type: ignore - MonitoringConfig, - MonitoringEngine, - ) - from monitoring._native_engine import ClickHouseClientConfig # type: ignore - from monitoring.config import CaptureSchedule # type: ignore - from integration.hf_adapter import generate_with_monitoring # type: ignore - except Exception as exc: - pytest.skip(f"monitoring native extension not available: {exc}") - - try: - from transformers import AutoModelForCausalLM, AutoTokenizer - from transformers.models.gpt2_p.modeling_gpt2 import HookedGPT2LMHeadModel # type: ignore - from transformers.models.qwen3_p.modeling_qwen3 import HookedQwen3ForCausalLM # type: ignore - from transformers.models.llama_p.modeling_llama import HookedLlamaForCausalLM # type: ignore - except Exception as exc: - pytest.skip(f"transformers or repo Hooked* classes not available: {exc}") - - # ----------------------------------------------------------------------- - # Configuration - # ----------------------------------------------------------------------- - - batch_size = int(os.environ.get("E2E_BATCH_SIZE", "4")) - if batch_size < 1: - raise ValueError("E2E_BATCH_SIZE must be >= 1") - - max_new_tokens = int(os.environ.get("E2E_MAX_NEW_TOKENS", "8")) - hf_model_id = _resolve_model_id(os.environ.get("E2E_MODEL", "gpt2")) - chunk_bytes = int(os.environ.get("E2E_CHUNK_BYTES", str(256 * 1024))) - - print_text = int(os.environ.get("E2E_PRINT_TEXT", "0")) == 1 - # E2E_HF_DROP_LAST_TOKEN is no longer needed: both monitored and HF reference - # use generate(), so they produce the same number of tokens/hidden states. - print_topk_logits = int(os.environ.get("E2E_PRINT_TOPK_LOGITS", "0")) == 1 - topk_k = int(os.environ.get("E2E_PRINT_TOPK_LOGITS_K", "5")) - if topk_k < 1: - raise ValueError("E2E_PRINT_TOPK_LOGITS_K must be >= 1") - - device = torch.device("cuda") - - # ----------------------------------------------------------------------- - # Tokenizer + prompts - # ----------------------------------------------------------------------- - - tokenizer = AutoTokenizer.from_pretrained(hf_model_id) - if tokenizer.pad_token_id is None: - tokenizer.pad_token_id = tokenizer.eos_token_id - tokenizer.padding_side = "left" - eos_id = int(tokenizer.eos_token_id) - pad_id = int(tokenizer.pad_token_id) - - prompts = [("Hello " * (i + 1)).strip() for i in range(batch_size)] - encoded = tokenizer(prompts, return_tensors="pt", padding=True) - input_ids = encoded["input_ids"].to(device) - attention_mask = encoded["attention_mask"].to(device) - - # Initial prompt tokens for safety prefix checks - hf_initial_prompt_tokens: List[torch.Tensor] = [] - for j in range(batch_size): - hf_initial_prompt_tokens.append( - _strip_left_pad( - input_ids[j].detach().cpu(), - attention_mask[j].detach().cpu(), - ).to(torch.long) - ) - - # ----------------------------------------------------------------------- - # Monitoring config (config-driven; no env var toggles) - # ----------------------------------------------------------------------- - - mon_cfg = MonitoringConfig( - schedule=CaptureSchedule(capture_prefill=True, capture_decode=True), - ) - - # ----------------------------------------------------------------------- - # ClickHouse config (for the monitored run) - # ----------------------------------------------------------------------- - - db_cfg_native = ClickHouseClientConfig() - db_cfg_native.host = os.environ.get("DMX_DB_HOST", "localhost") - db_cfg_native.port = int(os.environ.get("DMX_DB_PORT", "9000")) - db_cfg_native.username = os.environ.get("DMX_DB_USER", "default") - db_cfg_native.password = os.environ.get("DMX_DB_PASSWORD", "") - db_cfg_native.database = os.environ.get("DMX_DB_DATABASE", "default") - db_cfg_native.table = os.environ.get("DMX_DB_TABLE", "offload") - db_cfg_native.secure = False - db_cfg_native.client_side_compress = "none" - db_cfg_native.client_settings = None - db_cfg_native.create_database_if_missing = True - db_cfg_native.drop_existing_database = True - db_cfg_native.index_granularity = 8192 - - host_cfg = _make_host_cfg(db_cfg_native) - ring_cfg = _make_ring_cfg() - - # ----------------------------------------------------------------------- - # Monitored run - # ----------------------------------------------------------------------- - - unique_run_model_id = f"e2e_correctness_hf::{uuid.uuid4().hex}"[:120] - engine = MonitoringEngine( - config=mon_cfg, model_id=unique_run_model_id, db_config=host_cfg - ) - engine.enable_ring_transport(ring_cfg) - - if "qwen3" in hf_model_id.lower(): - model_cls = HookedQwen3ForCausalLM - elif "llama" in hf_model_id.lower(): - model_cls = HookedLlamaForCausalLM - else: - model_cls = HookedGPT2LMHeadModel - mon_model = model_cls.from_pretrained(hf_model_id, attn_implementation="eager", torch_dtype=torch.float16) - mon_model.to(device).eval() - mon_model.monitoring_engine = engine - - try: - with torch.no_grad(): - _ = generate_with_monitoring( - mon_model, - input_ids=input_ids, - attention_mask=attention_mask, - max_new_tokens=max_new_tokens, - do_sample=False, - pad_token_id=pad_id, - eos_token_id=eos_id, - logits_to_keep=0, - ) - finally: - engine.close() - - # ----------------------------------------------------------------------- - # Read DB (monitoring.clickhouse_reader + monitoring.segment_merger) - # ----------------------------------------------------------------------- - - ch = CHClickhouseDriverReadOnly( - host=str(db_cfg_native.host), - port=int(db_cfg_native.port), - username=str(db_cfg_native.username), - password=str(db_cfg_native.password), - database=str(db_cfg_native.database), - table=str(db_cfg_native.table), - secure=bool(getattr(db_cfg_native, "secure", False)), - client_settings=getattr(db_cfg_native, "client_settings", None), - decode_strings=True, - ) - try: - rows = ch.prefix_get((unique_run_model_id, ), return_full_key_tuple=True) - finally: - ch.close() - - if not rows: - pytest.fail(f"No rows found in ClickHouse for model_id={unique_run_model_id}") - - # If multiple shard_ranks are present, pick rank 0 if available, else the minimum. - shard_ranks = sorted({int(key[4]) for key, _t in rows}) - chosen_shard_rank = 0 if 0 in shard_ranks else (shard_ranks[0] if shard_ranks else 0) - rows = [(k, t) for (k, t) in rows if int(k[4]) == chosen_shard_rank] - - grouped: Dict[str, Dict[Tuple[int, str], List[Tuple[int, int, torch.Tensor]]]] = {} - for full_key, t_raw in rows: - # full_key = (model_id, request_id, act_name, layer_no, shard_rank, start_token_idx, end_token_idx) - _model_id, req_id, act_name_raw, layer_no_raw, _shard_rank, s, e = full_key - - layer_no, act_name = _canon_layer_and_act(str(act_name_raw), int(layer_no_raw)) - - t = t_raw.detach().cpu() - grouped.setdefault(str(req_id), {}).setdefault((layer_no, act_name), []).append( - (int(s), int(e), t) - ) - - request_ids = sorted(grouped.keys(), key=_parse_request_id) - - def _sort_chunks(chunks: List[Tuple[int, int, torch.Tensor]]) -> List[Tuple[int, int, torch.Tensor]]: - return sorted(chunks, key=lambda x: (x[0], x[1])) - - def _validate_contiguous( - chunks_sorted: List[Tuple[int, int, torch.Tensor]], expected_end: int, ctx: str - ) -> None: - if not chunks_sorted: - raise AssertionError(f"{ctx}: no chunks") - if chunks_sorted[0][0] != 0: - raise AssertionError(f"{ctx}: first chunk start={chunks_sorted[0][0]} expected 0") - prev_end = chunks_sorted[0][1] - for s2, e2, _t in chunks_sorted[1:]: - if s2 != prev_end: - raise AssertionError(f"{ctx}: non-contiguous chunks: start={s2} prev_end={prev_end}") - prev_end = e2 - if prev_end != expected_end: - raise AssertionError(f"{ctx}: coverage end={prev_end} expected_end={expected_end}") - - # Safety: this run should have exactly one group_id (single batch reset) - seen_group_ids: set[int] = set() - - db_token_ids_by_req: Dict[str, torch.Tensor] = {} - local_index_by_req: Dict[str, int] = {} - prompt_len_by_req: Dict[str, int] = {} - - for req_id in request_ids: - group_id, local_i = _parse_request_id(req_id) - seen_group_ids.add(group_id) - if not (0 <= local_i < batch_size): - raise AssertionError(f"{req_id}: local_index={local_i} out of range batch_size={batch_size}") - local_index_by_req[req_id] = local_i - - hooks_map = grouped[req_id] - if (-1, "token_ids") not in hooks_map: - raise AssertionError(f"{req_id}: DB missing token_ids") - - tok_chunks = _sort_chunks(hooks_map[(-1, "token_ids")]) - db_tok_end = int(tok_chunks[-1][1]) - _validate_contiguous(tok_chunks, expected_end=db_tok_end, ctx=f"{req_id} token_ids") - - db_tok = merge_segments([t for _, _, t in tok_chunks], "token_ids").to(torch.long) - if db_tok.ndim != 1: - db_tok = db_tok.view(-1) - - prompt0 = hf_initial_prompt_tokens[local_i] - plen0 = int(prompt0.numel()) - if db_tok.numel() < plen0 or not torch.equal(db_tok[:plen0], prompt0): - raise AssertionError( - f"{req_id}: request_id->row mapping safety check failed (initial prompt prefix). " - f"initial_prompt_len={plen0} db_tok_len={db_tok.numel()}" - ) - - db_token_ids_by_req[req_id] = db_tok.cpu() - prompt_len_by_req[req_id] = plen0 - - if len(seen_group_ids) != 1: - raise AssertionError(f"expected exactly one group_id for this test run, got {sorted(seen_group_ids)}") - - # ----------------------------------------------------------------------- - # Build HF references (no assertions here) - # ----------------------------------------------------------------------- - - hf_model = AutoModelForCausalLM.from_pretrained( - hf_model_id, - attn_implementation="eager", - torch_dtype=torch.float16, - ).to(device).eval() - - # Optional modules for GPT2-like reconstructions - wte = getattr(getattr(hf_model, "transformer", None), "wte", None) - wpe = getattr(getattr(hf_model, "transformer", None), "wpe", None) - # Qwen3-like: embed_tokens on model sub-module (no separate pos embed) - embed_tokens = getattr(getattr(hf_model, "model", None), "embed_tokens", None) - - # Compute num_layers once (used in per-layer comparison loop below) - num_layers = get_num_layers_from_config(hf_model) - - req_order = sorted(request_ids, key=_parse_request_id) - - # Run reference on ORIGINAL (non-hooked) model in a subprocess. - # This validates that our hooks don't change the model output, - # and ensures clean GPU memory isolation. - import subprocess, tempfile - ref_dir = tempfile.mkdtemp(prefix="hf_ref_") - ref_env = {**os.environ, "E2E_BATCH_SIZE": str(batch_size), - "E2E_MAX_NEW_TOKENS": str(max_new_tokens), - "E2E_MODEL": os.environ.get("E2E_MODEL", "gpt2")} - ref_result = subprocess.run( - [sys.executable, "-m", "tests.hf_reference_runner", "--output-dir", ref_dir], - env=ref_env, capture_output=True, text=True, cwd=os.path.dirname(os.path.dirname(__file__)), - ) - if ref_result.returncode != 0: - pytest.fail(f"Reference runner failed:\n{ref_result.stderr[-2000:]}") - hf_refs_batch = _load_hf_refs_from_disk(ref_dir) - import shutil - shutil.rmtree(ref_dir, ignore_errors=True) - hf_gens_batch = _hf_generate_collect_scores_batched( - hf_model=hf_model, - input_ids_batch=input_ids, - attention_mask_batch=attention_mask, - max_new_tokens=max_new_tokens, - eos_token_id=eos_id, - pad_token_id=pad_id, - device=device, - ) - if len(hf_refs_batch) != batch_size or len(hf_gens_batch) != batch_size: - raise AssertionError( - f"HF batched refs unexpected batch: rollout={len(hf_refs_batch)} " - f"gen={len(hf_gens_batch)} batch_size={batch_size}" - ) - - hf_ref_by_req: Dict[str, _HFRef] = {} - hf_gen_by_req: Dict[str, _HFGenRef] = {} - - def _decode(ids: torch.Tensor) -> str: - return tokenizer.decode(ids.tolist(), skip_special_tokens=False) - - for req_id in req_order: - i = local_index_by_req[req_id] - plen = int(prompt_len_by_req[req_id]) - - ref = hf_refs_batch[i] - gen_ref = hf_gens_batch[i] - - hf_ref_by_req[req_id] = ref - hf_gen_by_req[req_id] = gen_ref - - if print_text: - db_seq = db_token_ids_by_req[req_id] - db_prompt = db_seq[:plen] - db_gen = db_seq[plen:] - rol_prompt = ref.token_ids[:plen] - rol_gen = ref.token_ids[plen:] - gen_prompt = gen_ref.token_ids[:plen] - gen_gen = gen_ref.token_ids[plen:] - - print(f"\n=== {req_id} (local_index={i}, shard_rank={chosen_shard_rank}) ===") - print(f"DB: prompt_tokens={plen} generated_tokens={int(db_gen.numel())} total_tokens={int(db_seq.numel())}") - print(f"DB PROMPT: {_decode(db_prompt)!r}") - print(f"DB GENERATED: {_decode(db_gen)!r}") - print(f"DB FULL: {_decode(db_seq)!r}") - print(f"ROL: prompt_tokens={plen} generated_tokens={int(rol_gen.numel())} total_tokens={int(ref.token_ids.numel())}") - print(f"ROL PROMPT: {_decode(rol_prompt)!r}") - print(f"ROL GENERATED:{_decode(rol_gen)!r}") - print(f"ROL FULL: {_decode(ref.token_ids)!r}") - print(f"GEN: prompt_tokens={plen} generated_tokens={int(gen_gen.numel())} total_tokens={int(gen_ref.token_ids.numel())}") - print(f"GEN PROMPT: {_decode(gen_prompt)!r}") - print(f"GEN GENERATED:{_decode(gen_gen)!r}") - print(f"GEN FULL: {_decode(gen_ref.token_ids)!r}") - print("TOKENS MATCH?: YES (DB==ROL==GEN)") - - # Helpers for top-k logit printing - def _tok_piece(tok_id: int) -> str: - try: - return tokenizer.decode([tok_id], skip_special_tokens=False) - except Exception: - return f"" - - def _fmt_topk(ids_row: torch.Tensor, vals_row: torch.Tensor) -> str: - parts: List[str] = [] - for tid, v in zip(ids_row.tolist(), vals_row.tolist()): - parts.append(f"{int(tid)}:{_tok_piece(int(tid))!r}:{float(v):.6g}") - return " | ".join(parts) - - # ----------------------------------------------------------------------- - # Subtests: one per assertion - # ----------------------------------------------------------------------- - - for req_id in req_order: - i = local_index_by_req[req_id] - hooks_map = grouped[req_id] - - seq = db_token_ids_by_req[req_id].to(torch.long) # [T] - seq_len = int(seq.numel()) - ref = hf_ref_by_req[req_id] - gen_ref = hf_gen_by_req[req_id] - prompt_len = int(prompt_len_by_req[req_id]) - gen_base_pos = prompt_len - 1 # scores[0] corresponds to logits at pos (prompt_len-1) - - # --- token_ids --- - with subtests.test(msg=f"{req_id}/token_ids_rol"): - assert bitwise_equal(ref.token_ids, seq), ( - f"HF rollout tokens != DB token_ids " - f"(hf_len={int(ref.token_ids.numel())} db_len={seq_len})" - ) - - with subtests.test(msg=f"{req_id}/token_ids_gen"): - assert bitwise_equal(gen_ref.token_ids, seq), ( - f"HF generate tokens != DB token_ids " - f"(hf_len={int(gen_ref.token_ids.numel())} db_len={seq_len})" - ) - - # --- final_logits --- - logits_chunks_raw = hooks_map.get((-1, "final_logits"), []) - if logits_chunks_raw: - lchunks = sorted(logits_chunks_raw, key=lambda x: (x[0], x[1])) - db_logits_full = merge_segments([t for _, _, t in lchunks], "final_logits") - if db_logits_full.ndim == 1: - db_logits_full = db_logits_full.unsqueeze(0) - # ref.final_logits is decode-only scores from generate(): - # ref[s] = logits at position (prompt_len - 1 + s). - # Align with DB logits by position. - n_ref = int(ref.final_logits.shape[0]) - start = prompt_len - 1 - end = min(start + n_ref, int(db_logits_full.shape[0])) - n = end - start - db_slice = db_logits_full[start:end, :] - rol_slice = ref.final_logits[:n, :] - vocab_db = int(db_slice.shape[1]) - - if print_topk_logits: - print(f"\n=== TOP{topk_k} LOGITS {req_id} (local_index={i}, shard_rank={chosen_shard_rank}) ===") - print(f"seq_len={seq_len} vocab={vocab_db}") - db_topv, db_topi = torch.topk(db_slice.float(), k=topk_k, dim=-1) - rol_topv, rol_topi = torch.topk(rol_slice.float(), k=topk_k, dim=-1) - for tpos in range(seq_len): - cur_id = int(seq[tpos].item()) - cur_piece = _tok_piece(cur_id) - if tpos + 1 < seq_len: - nxt_id = int(seq[tpos + 1].item()) - nxt_piece = _tok_piece(nxt_id) - label_str = f" next={nxt_id}:{nxt_piece!r}" - else: - label_str = " next=" - print(f"\npos={tpos} tok={cur_id}:{cur_piece!r}{label_str}") - print(f" DB: {_fmt_topk(db_topi[tpos], db_topv[tpos])}") - print(f" ROL: {_fmt_topk(rol_topi[tpos], rol_topv[tpos])}") - if gen_base_pos >= 0 and gen_base_pos <= tpos <= (gen_base_pos + len(gen_ref.scores) - 1): - sidx = tpos - gen_base_pos - gs = gen_ref.scores[sidx] - g_topv, g_topi = torch.topk(gs.float(), k=topk_k, dim=-1) - print(f" GEN: {_fmt_topk(g_topi, g_topv)}") - else: - print(" GEN: ") + """HF eager: hooked model (ring -> ClickHouse) vs original model. - with subtests.test(msg=f"{req_id}/final_logits"): - if not bitwise_equal(db_slice, rol_slice): - diff = (db_slice.float() - rol_slice.float()).abs() - max_abs = float(diff.max().item()) - flat_idx = int(diff.view(-1).argmax().item()) - r = flat_idx // vocab_db - c = flat_idx % vocab_db - pytest.fail(f"final_logits mismatch (max_abs={max_abs}) at row={r} vocab_idx={c}") - - # --- hook_embed / hook_pos_embed (GPT2-like: both wte and wpe) --- - if ( - wte is not None - and wpe is not None - and (-1, "hook_embed") in hooks_map - and (-1, "hook_pos_embed") in hooks_map - ): - ids = seq.to(device) - pos = _positions_for_unpadded(seq_len, device=device) - emb = wte(ids).detach().cpu() # [T, d] - pos_emb = wpe(pos).detach().cpu() # [T, d] - - chunks = sorted(hooks_map[(-1, "hook_embed")], key=lambda x: (x[0], x[1])) - _validate_contiguous(chunks, expected_end=seq_len, ctx=f"{req_id} hook_embed") - db_t = merge_segments([t for _, _, t in chunks], "hook_embed") - with subtests.test(msg=f"{req_id}/hook_embed"): - assert tuple(db_t.shape) == tuple(emb.shape), ( - f"hook_embed shape mismatch db={tuple(db_t.shape)} hf={tuple(emb.shape)}" - ) - assert bitwise_equal(db_t, emb), ( - f"hook_embed mismatch (max_abs={float((db_t.float() - emb.float()).abs().max().item())})" - ) - - chunks = sorted(hooks_map[(-1, "hook_pos_embed")], key=lambda x: (x[0], x[1])) - _validate_contiguous(chunks, expected_end=seq_len, ctx=f"{req_id} hook_pos_embed") - db_t = merge_segments([t for _, _, t in chunks], "hook_pos_embed") - with subtests.test(msg=f"{req_id}/hook_pos_embed"): - assert tuple(db_t.shape) == tuple(pos_emb.shape), ( - f"hook_pos_embed shape mismatch db={tuple(db_t.shape)} hf={tuple(pos_emb.shape)}" - ) - assert bitwise_equal(db_t, pos_emb), ( - f"hook_pos_embed mismatch (max_abs={float((db_t.float() - pos_emb.float()).abs().max().item())})" - ) - - # --- hook_embed only (Qwen3-like: RoPE, no separate pos embed) --- - elif ( - embed_tokens is not None - and wpe is None - and (-1, "hook_embed") in hooks_map - ): - emb = embed_tokens(seq.to(device)).detach().cpu() # [T, d] - chunks = sorted(hooks_map[(-1, "hook_embed")], key=lambda x: (x[0], x[1])) - _validate_contiguous(chunks, expected_end=seq_len, ctx=f"{req_id} hook_embed") - db_t = merge_segments([t for _, _, t in chunks], "hook_embed") - with subtests.test(msg=f"{req_id}/hook_embed"): - assert tuple(db_t.shape) == tuple(emb.shape), ( - f"hook_embed shape mismatch db={tuple(db_t.shape)} hf={tuple(emb.shape)}" - ) - assert bitwise_equal(db_t, emb), ( - f"hook_embed mismatch (max_abs={float((db_t.float() - emb.float()).abs().max().item())})" - ) - - # --- hook_final_ln --- - # TODO: hook_final_ln comparison skipped. - - # --- per-layer: attention pattern + resid_pre --- - # Support both GPT-2 naming (blocks.attn.hook_pattern, blocks.hook_resid_*) - # and Qwen3 naming (layers.self_attn.hook_pattern, layers.hook_resid_*) - _ATTN_PATTERN_KEYS = ("blocks.attn.hook_pattern", "layers.self_attn.hook_pattern") - _RESID_PRE_KEYS = ("blocks.hook_resid_pre", "layers.hook_resid_pre") - - n_layers = len(ref.attn_pattern) if ref.attn_pattern else 0 - assert n_layers == num_layers, ( - f"{req_id}: attn_pattern layer count mismatch: rollout={n_layers} config={num_layers}" - ) - - for layer_no in range(n_layers): - # attn_pattern: compare per-chunk (can't merge because kv_dim - # differs between prefill [H, plen, plen] and decode [H, 1, plen+i]) - key = next(((layer_no, k) for k in _ATTN_PATTERN_KEYS if (layer_no, k) in hooks_map), None) - if key is not None: - pat = ref.attn_pattern[layer_no] # [H, T, T] - chunks = sorted(hooks_map[key], key=lambda x: (x[0], x[1])) - all_ok = True - fail_msg = "" - for start, end, t_chunk in chunks: - q_len = end - start - db_c = t_chunk - if db_c.ndim == 4 and db_c.shape[0] == 1: - db_c = db_c.squeeze(0) - # kv_dim for these rows: causal, valid up to position 'end' - kv_valid = end - db_c = db_c[:, :q_len, :kv_valid] - ref_c = pat[:, start:end, :kv_valid] - if db_c.shape != ref_c.shape: - all_ok = False - fail_msg = (f"shape mismatch at [{start}:{end}] " - f"db={db_c.shape} ref={ref_c.shape}") - break - if not bitwise_equal(db_c, ref_c): - max_abs = float((db_c.float() - ref_c.float()).abs().max().item()) - all_ok = False - fail_msg = (f"value mismatch at [{start}:{end}] " - f"max_abs={max_abs:.6f}") - break - with subtests.test(msg=f"{req_id}/layer{layer_no}/attn_pattern"): - assert all_ok, ( - f"pattern layer={layer_no}: {fail_msg}" - ) - - key = next(((layer_no, k) for k in _RESID_PRE_KEYS if (layer_no, k) in hooks_map), None) - if key is not None and ref.hidden_states and layer_no < len(ref.hidden_states): - hs = ref.hidden_states[layer_no] # [T, d] - chunks = sorted(hooks_map[key], key=lambda x: (x[0], x[1])) - _validate_contiguous(chunks, expected_end=seq_len, ctx=f"{req_id} layer{layer_no} resid_pre") - db_t = merge_segments([t for _, _, t in chunks], key[1]) - with subtests.test(msg=f"{req_id}/layer{layer_no}/resid_pre"): - assert tuple(db_t.shape) == tuple(hs.shape), ( - f"resid_pre shape mismatch layer={layer_no} db={tuple(db_t.shape)} hf={tuple(hs.shape)}" - ) - assert bitwise_equal(db_t, hs), ( - f"resid_pre mismatch layer={layer_no} " - f"(max_abs={float((db_t.float() - hs.float()).abs().max().item())})" - ) - - # --- resid_final (global: last layer's pre-norm residual) --- - # HF's output_hidden_states[-1] is POST-final-norm (after ln_f), - # not pre-norm. resid_final captures pre-norm. No direct HF - # reference available, so we only check shape and presence. - key = (-1, "hook_resid_final") - if key in hooks_map: - chunks = sorted(hooks_map[key], key=lambda x: (x[0], x[1])) - _validate_contiguous(chunks, expected_end=seq_len, ctx=f"{req_id} resid_final") - db_t = merge_segments([t for _, _, t in chunks], key[1]) - with subtests.test(msg=f"{req_id}/resid_final"): - assert db_t.shape[-1] == ref.hidden_states[0].shape[-1] if ref.hidden_states else True, ( - f"resid_final hidden_dim mismatch" - ) - assert db_t.shape[0] == seq_len, ( - f"resid_final token count mismatch db={db_t.shape[0]} expected={seq_len}" - ) - - -# --------------------------------------------------------------------------- -# CUDA-graph correctness test -# --------------------------------------------------------------------------- - -@pytest.mark.skipif(True, reason=( - "DISABLED: compiled rollout reference (StaticCache + torch.compile) cannot " - "replicate generate()'s internal StaticCache attention mask / position handling. " - "The manual rollout produces different hidden states than generate() even for " - "identical inputs — this is a fundamental mismatch in how HF handles StaticCache " - "internally vs externally. Use test_e2e_cuda_graphs_vs_eager_hf instead, which " - "compares CUDA-graph DB against an uncompiled eager reference with relaxed tolerance." -)) -def test_e2e_correctness_hf_cuda_graphs(subtests) -> None: - """Same as test_e2e_correctness_hf but with torch.compile + static KV cache (CUDA graphs). - - NOTE: This test is currently DISABLED. The compiled rollout reference uses - StaticCache + torch.compile on a manual decode loop, but this produces - different numerical results from HF generate(cache_implementation="static") - because generate() handles attention masks and position_ids differently - internally. CUDA graphs also prevent reading hidden states from generate() - (Bug 11 in debug.log). See test_e2e_cuda_graphs_vs_eager_hf for the - working alternative. - - Run with: - CUDA_MODULE_LOADING=EAGER pytest -q -s tests/test_e2e_correctness_vs_hf.py::test_e2e_correctness_hf_cuda_graphs + Equivalent matrix cell: ``--backend hf --mode eager --standard allclose`` + (HF dispatches to hf_comparator, which does the value comparison with + ``E2E_TOLERANCE``; default 0.01 for eager). """ - try: - import clickhouse_driver # noqa: F401 - except Exception: - pytest.skip("clickhouse-driver is required") - - try: - from monitoring import ( # type: ignore - MonitoringConfig, - MonitoringEngine, - ) - from monitoring._native_engine import ClickHouseClientConfig # type: ignore - from monitoring.config import CaptureSchedule # type: ignore - from integration.hf_adapter import generate_with_monitoring # type: ignore - except Exception as exc: - pytest.skip(f"monitoring native extension not available: {exc}") - - try: - from transformers import AutoModelForCausalLM, AutoTokenizer - from transformers.models.gpt2_p.modeling_gpt2 import HookedGPT2LMHeadModel # type: ignore - from transformers.models.qwen3_p.modeling_qwen3 import HookedQwen3ForCausalLM # type: ignore - from transformers.models.llama_p.modeling_llama import HookedLlamaForCausalLM # type: ignore - except Exception as exc: - pytest.skip(f"transformers or Hooked* classes not available: {exc}") - - # ----------------------------------------------------------------------- - # Config — fixed small values to keep the test fast - # ----------------------------------------------------------------------- - batch_size = int(os.environ.get("E2E_BATCH_SIZE", "4")) - max_new_tokens = int(os.environ.get("E2E_MAX_NEW_TOKENS", "8")) - hf_model_id = os.environ.get("E2E_MODEL", "gpt2") - hf_model_id = _MODEL_ALIASES.get(hf_model_id.lower(), hf_model_id) - chunk_bytes = int(os.environ.get("E2E_CHUNK_BYTES", str(256 * 1024))) - # E2E_HF_DROP_LAST_TOKEN is no longer needed: both monitored and HF reference - # use generate(), so they produce the same number of tokens/hidden states. - - device = torch.device("cuda") - - # ----------------------------------------------------------------------- - # Tokenizer + prompts - # ----------------------------------------------------------------------- - tokenizer = AutoTokenizer.from_pretrained(hf_model_id) - if tokenizer.pad_token_id is None: - tokenizer.pad_token_id = tokenizer.eos_token_id - tokenizer.padding_side = "left" - eos_id = int(tokenizer.eos_token_id) - pad_id = int(tokenizer.pad_token_id) - - prompts = [("Hello " * (i + 1)).strip() for i in range(batch_size)] - encoded = tokenizer(prompts, return_tensors="pt", padding=True) - input_ids = encoded["input_ids"].to(device) - attention_mask = encoded["attention_mask"].to(device) - - hf_initial_prompt_tokens: List[torch.Tensor] = [] - for j in range(batch_size): - hf_initial_prompt_tokens.append( - _strip_left_pad( - input_ids[j].detach().cpu(), - attention_mask[j].detach().cpu(), - ).to(torch.long) - ) - - # ----------------------------------------------------------------------- - # Monitoring + DB config - # ----------------------------------------------------------------------- - mon_cfg = MonitoringConfig( - schedule=CaptureSchedule(capture_prefill=True, capture_decode=True), - ) - - db_cfg_native = ClickHouseClientConfig() - db_cfg_native.host = os.environ.get("DMX_DB_HOST", "localhost") - db_cfg_native.port = int(os.environ.get("DMX_DB_PORT", "9000")) - db_cfg_native.username = os.environ.get("DMX_DB_USER", "default") - db_cfg_native.password = os.environ.get("DMX_DB_PASSWORD", "") - db_cfg_native.database = os.environ.get("DMX_DB_DATABASE", "default") - db_cfg_native.table = os.environ.get("DMX_DB_TABLE", "offload") - db_cfg_native.secure = False - db_cfg_native.client_side_compress = "none" - db_cfg_native.client_settings = None - db_cfg_native.create_database_if_missing = True - db_cfg_native.drop_existing_database = True - db_cfg_native.index_granularity = 8192 - - host_cfg = _make_host_cfg(db_cfg_native) - ring_cfg = _make_ring_cfg() - - # ----------------------------------------------------------------------- - # Monitored model — compiled with torch.compile + static cache (CUDA graphs) - # ----------------------------------------------------------------------- - unique_run_model_id = f"e2e_cuda_graphs::{uuid.uuid4().hex}"[:120] - engine = MonitoringEngine( - config=mon_cfg, - model_id=unique_run_model_id, db_config=host_cfg, - ) - engine.enable_ring_transport(ring_cfg) - - if "qwen3" in hf_model_id.lower(): - model_cls = HookedQwen3ForCausalLM - elif "llama" in hf_model_id.lower(): - model_cls = HookedLlamaForCausalLM - else: - model_cls = HookedGPT2LMHeadModel - mon_model = model_cls.from_pretrained( - hf_model_id, attn_implementation="eager", torch_dtype=torch.float16, - ) - mon_model.to(device).eval() - - mon_model.monitoring_engine = engine - - try: - from transformers import CompileConfig - with torch.no_grad(): - gen_out = generate_with_monitoring( - mon_model, - input_ids=input_ids, - attention_mask=attention_mask, - max_new_tokens=max_new_tokens, - do_sample=False, - pad_token_id=pad_id, - eos_token_id=eos_id, - cache_implementation="static", - compile_config=CompileConfig(mode="reduce-overhead", fullgraph=False), - ) - finally: - engine.close() - - # Build per-request reference sequences from the generate() output. - # We use the compiled model's own output as the reference — this avoids - # any comparison against a different model that may compute different - # values under static cache or torch.compile. - gen_out_cpu = gen_out.detach().cpu().long() # [batch, total_len] - ref_seqs: List[torch.Tensor] = [] - for j in range(batch_size): - seq = _strip_left_pad(gen_out_cpu[j], (gen_out_cpu[j] != pad_id).long()) - ref_seqs.append(seq) - - # ----------------------------------------------------------------------- - # Read DB - # ----------------------------------------------------------------------- - from monitoring.clickhouse_reader import CHClickhouseDriverReadOnly - from monitoring.segment_merger import merge_segments, parse_internal_id - - ch = CHClickhouseDriverReadOnly( - host=str(db_cfg_native.host), - port=int(db_cfg_native.port), - username=str(db_cfg_native.username), - password=str(db_cfg_native.password), - database=str(db_cfg_native.database), - table=str(db_cfg_native.table), - secure=bool(getattr(db_cfg_native, "secure", False)), - client_settings=getattr(db_cfg_native, "client_settings", None), - decode_strings=True, - ) - try: - rows = ch.prefix_get((unique_run_model_id,), return_full_key_tuple=True) - finally: - ch.close() - - print(f"\n[DEBUG] Total DB rows: {len(rows)}") - if not rows: - pytest.fail(f"No rows found in ClickHouse for model_id={unique_run_model_id!r}. " - "This means monitoring produced no output at all under CUDA graphs.") - - shard_ranks = sorted({int(key[4]) for key, _t in rows}) - chosen_shard = 0 if 0 in shard_ranks else shard_ranks[0] - rows = [(k, t) for (k, t) in rows if int(k[4]) == chosen_shard] - - grouped: Dict[str, Dict[Tuple[int, str], List[Tuple[int, int, torch.Tensor]]]] = {} - for full_key, t_raw in rows: - _model_id, req_id, act_name_raw, layer_no_raw, _shard, s, e = full_key - layer_no, act_name = _canon_layer_and_act(str(act_name_raw), int(layer_no_raw)) - grouped.setdefault(str(req_id), {}).setdefault( - (layer_no, act_name), [] - ).append((int(s), int(e), t_raw.detach().cpu())) - - request_ids = sorted(grouped.keys(), key=_parse_request_id) - - # DEBUG: show per-request hook chunk counts - from collections import Counter - hook_totals = Counter() - for rid in request_ids: - hooks_map = grouped[rid] - for (layer, hname), chunks in hooks_map.items(): - hook_totals[hname] += len(chunks) - print(f"[DEBUG] Hook totals across all requests:") - for hname in ['token_ids', 'hook_embed', 'hook_pos_embed', 'blocks.0.hook_resid_pre', 'hook_resid_final', 'hook_final_ln', 'final_logits']: - print(f" {hname}: {hook_totals.get(hname, 0)} chunks") - rid = request_ids[0] - hooks_map = grouped[rid] - print(f"[DEBUG] All hooks for {rid}:") - for (layer, hname) in sorted(hooks_map.keys()): - chunks = hooks_map[(layer, hname)] - print(f" ({layer},{hname}): {len(chunks)} chunks") - - # ----------------------------------------------------------------------- - # HF reference — compiled manual rollout with StaticCache. - # Uses torch.compile(mode="reduce-overhead", fullgraph=False) on the - # decode step + cudagraph_mark_step_begin() + immediate .detach().cpu() - # to clone hidden states before CUDA graph buffers are overwritten. - # (generate() can't do this — see Bug 11 in debug.log) - # ----------------------------------------------------------------------- - hf_model = AutoModelForCausalLM.from_pretrained( - hf_model_id, attn_implementation="eager", torch_dtype=torch.float16, - ).to(device).eval() - - wte = getattr(getattr(hf_model, "transformer", None), "wte", None) - wpe = getattr(getattr(hf_model, "transformer", None), "wpe", None) - embed_tokens = getattr(getattr(hf_model, "model", None), "embed_tokens", None) - num_layers = get_num_layers_from_config(hf_model) - - hf_refs_batch = _hf_greedy_rollout_collect_all_batched( - hf_model=hf_model, - input_ids_batch=input_ids, - attention_mask_batch=attention_mask, - max_new_tokens=max_new_tokens, - eos_token_id=eos_id, - pad_token_id=pad_id, - device=device, - want_hidden_states=True, - want_attentions=True, - compiled=True, - ) - - # ----------------------------------------------------------------------- - # Pre-loop: build per-request dicts and safety-check token_ids - # ----------------------------------------------------------------------- - local_index_by_req: Dict[str, int] = {} - prompt_len_by_req: Dict[str, int] = {} - db_token_ids_by_req: Dict[str, torch.Tensor] = {} - - for req_id in request_ids: - _gid, local_i = _parse_request_id(req_id) - local_index_by_req[req_id] = local_i - - req_order = sorted(request_ids, key=_parse_request_id) - - for req_id in req_order: - local_i = local_index_by_req[req_id] - hooks_map = grouped[req_id] - prompt0 = hf_initial_prompt_tokens[local_i] - plen = int(prompt0.numel()) - prompt_len_by_req[req_id] = plen - - if (-1, "token_ids") not in hooks_map: - raise AssertionError(f"{req_id}: DB missing token_ids under CUDA graphs") - - tok_chunks = sorted(hooks_map[(-1, "token_ids")], key=lambda x: (x[0], x[1])) - db_tok = merge_segments([t for _, _, t in tok_chunks], "token_ids").to(torch.long) - if db_tok.ndim != 1: - db_tok = db_tok.view(-1) - - if db_tok.numel() < plen or not torch.equal(db_tok[:plen], prompt0): - raise AssertionError( - f"{req_id}: DB token_ids prompt prefix mismatch under CUDA graphs " - f"(plen={plen} db_len={db_tok.numel()})" - ) - db_token_ids_by_req[req_id] = db_tok.cpu() - - def _sort_chunks(chunks): - return sorted(chunks, key=lambda x: (x[0], x[1])) - - def _validate_contiguous(chunks_sorted, expected_end, ctx): - if not chunks_sorted: - raise AssertionError(f"{ctx}: no chunks") - if chunks_sorted[0][0] != 0: - raise AssertionError(f"{ctx}: first chunk start={chunks_sorted[0][0]} expected 0") - prev_end = chunks_sorted[0][1] - for s2, e2, _t in chunks_sorted[1:]: - if s2 != prev_end: - raise AssertionError(f"{ctx}: non-contiguous chunks start={s2} prev_end={prev_end}") - prev_end = e2 - if prev_end != expected_end: - raise AssertionError(f"{ctx}: coverage end={prev_end} expected_end={expected_end}") - - # ----------------------------------------------------------------------- - # Per-request assertions (full verification) - # ----------------------------------------------------------------------- - _RESID_PRE_KEYS = ("blocks.hook_resid_pre", "layers.hook_resid_pre") - - # HF reference is uncompiled; monitored model is compiled (reduce-overhead). - # Fall back to allclose if not bitwise equal. - _COMPILED_ATOL = 0.5 # safety net: real transport errors are >> 1 - - def _assert_close_or_bitwise(db_t, ref_t, label): - if bitwise_equal(db_t, ref_t): - return - diff = (db_t.float() - ref_t.float()).abs() - max_abs = float(diff.max().item()) - if torch.allclose(db_t.float(), ref_t.float(), atol=_COMPILED_ATOL, rtol=0.0): - import warnings - warnings.warn( - f"[NOT BITWISE] {label}: max_abs_diff={max_abs:.6f} " - f"(within atol={_COMPILED_ATOL}, but not bitwise equal)" - ) - return - pytest.fail( - f"{label}: max_abs_diff={max_abs:.6f} > atol={_COMPILED_ATOL}" - ) - - for req_id in req_order: - local_i = local_index_by_req[req_id] - hooks_map = grouped[req_id] - plen = prompt_len_by_req[req_id] - prompt0 = hf_initial_prompt_tokens[local_i] - - db_tok = db_token_ids_by_req[req_id] - seq_len = int(db_tok.numel()) - - ref = hf_refs_batch[local_i] - - # --- token_ids --- - with subtests.test(msg=f"{req_id}/cuda_graph/token_ids_present"): - assert (-1, "token_ids") in hooks_map, f"{req_id}: DB missing token_ids" + cr = run_single(matrix_argv_from_env("hf", "allclose", mode="eager")) + _assert_cell(subtests, cr) - with subtests.test(msg=f"{req_id}/cuda_graph/prompt_prefix"): - assert db_tok.numel() >= plen and torch.equal(db_tok[:plen], prompt0), ( - f"{req_id}: prompt prefix mismatch" - ) - with subtests.test(msg=f"{req_id}/cuda_graph/token_ids_match_hf"): - assert bitwise_equal(db_tok, ref.token_ids), ( - f"{req_id}: DB token_ids do not match HF generate() under CUDA graphs. " - f"db_len={db_tok.numel()} hf_len={ref.token_ids.numel()} " - f"(if db_len << hf_len the CUDA-graph monitoring bug is present)" - ) - - # --- final_logits --- - logits_chunks_raw = hooks_map.get((-1, "final_logits"), []) - if logits_chunks_raw: - lchunks = _sort_chunks(logits_chunks_raw) - db_logits = merge_segments([t for _, _, t in lchunks], "final_logits") - if db_logits.ndim == 1: - db_logits = db_logits.unsqueeze(0) - n_ref = int(ref.final_logits.shape[0]) - start = plen - 1 - end = min(start + n_ref, int(db_logits.shape[0])) - n = end - start - db_slice = db_logits[start:end, :] - rol_slice = ref.final_logits[:n, :] - with subtests.test(msg=f"{req_id}/cuda_graph/final_logits"): - assert n > 0, f"final_logits: no overlapping rows" - assert db_slice.shape[1] == rol_slice.shape[1], ( - f"final_logits vocab mismatch db={db_slice.shape[1]} hf={rol_slice.shape[1]}" - ) - _assert_close_or_bitwise(db_slice, rol_slice, f"{req_id} final_logits") - - # --- hook_embed --- - seq = db_tok.to(device) - if wte is not None and wpe is not None and (-1, "hook_embed") in hooks_map: - emb = wte(seq).detach().cpu() - chunks = _sort_chunks(hooks_map[(-1, "hook_embed")]) - _validate_contiguous(chunks, seq_len, f"{req_id} hook_embed") - db_t = merge_segments([t for _, _, t in chunks], "hook_embed") - with subtests.test(msg=f"{req_id}/cuda_graph/hook_embed"): - assert tuple(db_t.shape) == tuple(emb.shape), ( - f"hook_embed shape mismatch db={tuple(db_t.shape)} hf={tuple(emb.shape)}" - ) - _assert_close_or_bitwise(db_t, emb, f"{req_id} hook_embed") - elif embed_tokens is not None and wpe is None and (-1, "hook_embed") in hooks_map: - emb = embed_tokens(seq).detach().cpu() - chunks = _sort_chunks(hooks_map[(-1, "hook_embed")]) - _validate_contiguous(chunks, seq_len, f"{req_id} hook_embed") - db_t = merge_segments([t for _, _, t in chunks], "hook_embed") - with subtests.test(msg=f"{req_id}/cuda_graph/hook_embed"): - assert tuple(db_t.shape) == tuple(emb.shape), ( - f"hook_embed shape mismatch db={tuple(db_t.shape)} hf={tuple(emb.shape)}" - ) - _assert_close_or_bitwise(db_t, emb, f"{req_id} hook_embed") - - # --- hook_pos_embed (GPT2 only) --- - if wpe is not None and (-1, "hook_pos_embed") in hooks_map: - pos = _positions_for_unpadded(seq_len, device=device) - pos_emb = wpe(pos).detach().cpu() - chunks = _sort_chunks(hooks_map[(-1, "hook_pos_embed")]) - _validate_contiguous(chunks, seq_len, f"{req_id} hook_pos_embed") - db_t = merge_segments([t for _, _, t in chunks], "hook_pos_embed") - with subtests.test(msg=f"{req_id}/cuda_graph/hook_pos_embed"): - assert tuple(db_t.shape) == tuple(pos_emb.shape), ( - f"hook_pos_embed shape mismatch db={tuple(db_t.shape)} hf={tuple(pos_emb.shape)}" - ) - _assert_close_or_bitwise(db_t, pos_emb, f"{req_id} hook_pos_embed") - - # --- per-layer: resid_pre + attn_pattern --- - _ATTN_PATTERN_KEYS = ("blocks.attn.hook_pattern", "layers.self_attn.hook_pattern") - for layer_no in range(num_layers): - key = next(((layer_no, k) for k in _RESID_PRE_KEYS if (layer_no, k) in hooks_map), None) - if key is not None and ref.hidden_states and layer_no < len(ref.hidden_states): - hs = ref.hidden_states[layer_no] - chunks = _sort_chunks(hooks_map[key]) - _validate_contiguous(chunks, seq_len, f"{req_id} layer{layer_no} resid_pre") - db_t = merge_segments([t for _, _, t in chunks], key[1]) - with subtests.test(msg=f"{req_id}/cuda_graph/layer{layer_no}/resid_pre"): - assert tuple(db_t.shape) == tuple(hs.shape), ( - f"resid_pre shape mismatch layer={layer_no} db={tuple(db_t.shape)} hf={tuple(hs.shape)}" - ) - _assert_close_or_bitwise(db_t, hs, f"{req_id} layer{layer_no} resid_pre") - - # attn_pattern: both DB and ref use static cache, same padding - key = next(((layer_no, k) for k in _ATTN_PATTERN_KEYS if (layer_no, k) in hooks_map), None) - if key is not None and ref.attn_pattern and layer_no < len(ref.attn_pattern): - pat = ref.attn_pattern[layer_no] - chunks = _sort_chunks(hooks_map[key]) - db_pat = merge_segments([t for _, _, t in chunks], key[1]) - if db_pat.ndim == 4 and db_pat.shape[0] == 1: - db_pat = db_pat.squeeze(0) - with subtests.test(msg=f"{req_id}/cuda_graph/layer{layer_no}/attn_pattern"): - assert tuple(db_pat.shape) == tuple(pat.shape), ( - f"pattern shape mismatch layer={layer_no} db={tuple(db_pat.shape)} hf={tuple(pat.shape)}" - ) - _assert_close_or_bitwise(db_pat, pat, f"{req_id} layer{layer_no} attn_pattern") - - # --- resid_final --- - # HF's output_hidden_states[-1] is post-final-norm, not pre-norm. - # Only check shape and presence. - key = (-1, "hook_resid_final") - if key in hooks_map: - chunks = _sort_chunks(hooks_map[key]) - _validate_contiguous(chunks, seq_len, f"{req_id} resid_final") - db_t = merge_segments([t for _, _, t in chunks], key[1]) - with subtests.test(msg=f"{req_id}/cuda_graph/resid_final"): - assert db_t.shape[0] == seq_len, ( - f"resid_final token count mismatch db={db_t.shape[0]} expected={seq_len}" - ) - - -# --------------------------------------------------------------------------- -# Test: CUDA-graph monitored DB vs uncompiled eager HF reference -# --------------------------------------------------------------------------- - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA + native backend required") +@require_cuda() +@require_clickhouse() def test_e2e_cuda_graphs_vs_eager_hf(subtests) -> None: - """Compare CUDA-graph monitored run against original eager model. + """HF CUDA-graph monitored run vs eager reference, relaxed tolerance. - Three subprocesses — parent process never touches CUDA: - 1. Reference: original model (eager) -> tensors on disk - 2. Monitored: hooked model + ring (CUDA graphs, static cache) -> ClickHouse - 3. Comparator: reads both, compares, writes result.json + The HF reference runner is always eager; the monitored runner honors + ``E2E_CUDA_GRAPHS`` (set from the cuda_graph mode). Tolerance defaults + to 0.5 to absorb bf16 accumulation-order drift between compiled and + uncompiled paths. """ - import json - import subprocess - import tempfile - import shutil - - run_dir = tempfile.mkdtemp(prefix="hf_cg_e2e_") - 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") - project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - - # CUDA graph mode: monitored runs with static cache + torch.compile. - # Reference runs eager. Relaxed tolerance for bf16 rounding from - # different accumulation order (compiled vs uncompiled). - mon_env = {**os.environ, "E2E_CUDA_GRAPHS": "1"} - cmp_env = {**os.environ, "E2E_TOLERANCE": "0.5"} - - try: - print("\n [1/3] Reference run (original model, eager)...", flush=True) - r1 = subprocess.run( - [sys.executable, "-m", "tests.hf_reference_runner", - "--output-dir", ref_dir], - env=os.environ, capture_output=True, text=True, cwd=project_root, - ) - if r1.returncode != 0: - pytest.fail(f"Reference runner failed:\n{r1.stderr[-2000:]}") - - print(" [2/3] Monitored run (hooked model + ring, CUDA graphs)...", flush=True) - r2 = subprocess.run( - [sys.executable, "-m", "tests.hf_monitored_runner", - "--output-dir", mon_dir], - env=mon_env, capture_output=True, text=True, cwd=project_root, - ) - if r2.returncode != 0: - pytest.fail(f"Monitored runner failed:\n{r2.stderr[-2000:]}") - - print(" [3/3] Comparing (tolerance=0.5 for CG vs eager)...", flush=True) - r3 = subprocess.run( - [sys.executable, "-m", "tests.hf_comparator", - "--ref-dir", ref_dir, - "--mon-dir", mon_dir, - "--result-file", result_file], - env=cmp_env, capture_output=True, text=True, cwd=project_root, - ) - if r3.returncode != 0: - pytest.fail(f"Comparator failed:\n{r3.stderr[-2000:]}") - - with open(result_file) as f: - results = json.load(f) - - for test in results["tests"]: - with subtests.test(test["name"]): - assert test["passed"], test.get("detail", "") - - finally: - shutil.rmtree(run_dir, ignore_errors=True) - - -def _test_e2e_cuda_graphs_vs_eager_hf_legacy(subtests) -> None: - """Legacy version kept for reference. Not called by verify_hf.sh.""" - try: - import clickhouse_driver # noqa: F401 - except Exception: - pytest.skip("clickhouse-driver is required") - - try: - from monitoring import ( # type: ignore - MonitoringConfig, - MonitoringEngine, - ) - from monitoring._native_engine import ClickHouseClientConfig # type: ignore - from monitoring.config import CaptureSchedule # type: ignore - from integration.hf_adapter import generate_with_monitoring # type: ignore - except Exception as exc: - pytest.skip(f"monitoring native extension not available: {exc}") - - try: - from transformers import AutoModelForCausalLM, AutoTokenizer - from transformers.models.gpt2_p.modeling_gpt2 import HookedGPT2LMHeadModel # type: ignore - from transformers.models.qwen3_p.modeling_qwen3 import HookedQwen3ForCausalLM # type: ignore - from transformers.models.llama_p.modeling_llama import HookedLlamaForCausalLM # type: ignore - except Exception as exc: - pytest.skip(f"transformers or Hooked* classes not available: {exc}") - - batch_size = int(os.environ.get("E2E_BATCH_SIZE", "4")) - max_new_tokens = int(os.environ.get("E2E_MAX_NEW_TOKENS", "8")) - hf_model_id = _resolve_model_id(os.environ.get("E2E_MODEL", "gpt2")) - chunk_bytes = int(os.environ.get("E2E_CHUNK_BYTES", str(256 * 1024))) - device = torch.device("cuda") - - tokenizer = AutoTokenizer.from_pretrained(hf_model_id) - if tokenizer.pad_token_id is None: - tokenizer.pad_token_id = tokenizer.eos_token_id - tokenizer.padding_side = "left" - eos_id = int(tokenizer.eos_token_id) - pad_id = int(tokenizer.pad_token_id) - - prompts = [("Hello " * (i + 1)).strip() for i in range(batch_size)] - encoded = tokenizer(prompts, return_tensors="pt", padding=True) - input_ids = encoded["input_ids"].to(device) - attention_mask = encoded["attention_mask"].to(device) - - hf_initial_prompt_tokens: List[torch.Tensor] = [] - for j in range(batch_size): - hf_initial_prompt_tokens.append( - _strip_left_pad(input_ids[j].detach().cpu(), attention_mask[j].detach().cpu()).to(torch.long) - ) - - # --- Monitoring config + DB --- - mon_cfg = MonitoringConfig( - schedule=CaptureSchedule(capture_prefill=True, capture_decode=True), - ) - - db_cfg_native = ClickHouseClientConfig() - db_cfg_native.host = os.environ.get("DMX_DB_HOST", "localhost") - db_cfg_native.port = int(os.environ.get("DMX_DB_PORT", "9000")) - db_cfg_native.username = os.environ.get("DMX_DB_USER", "default") - db_cfg_native.password = os.environ.get("DMX_DB_PASSWORD", "") - db_cfg_native.database = os.environ.get("DMX_DB_DATABASE", "default") - db_cfg_native.table = os.environ.get("DMX_DB_TABLE", "offload") - db_cfg_native.secure = False - db_cfg_native.client_side_compress = "none" - db_cfg_native.client_settings = None - db_cfg_native.create_database_if_missing = True - db_cfg_native.drop_existing_database = True - db_cfg_native.index_granularity = 8192 - - host_cfg = _make_host_cfg(db_cfg_native) - ring_cfg = _make_ring_cfg() - - # --- Monitored model (compiled, CUDA graphs) --- - unique_run_model_id = f"e2e_cg_vs_eager::{uuid.uuid4().hex}"[:120] - engine = MonitoringEngine( - config=mon_cfg, - model_id=unique_run_model_id, db_config=host_cfg, - ) - engine.enable_ring_transport(ring_cfg) - - if "qwen3" in hf_model_id.lower(): - model_cls = HookedQwen3ForCausalLM - elif "llama" in hf_model_id.lower(): - model_cls = HookedLlamaForCausalLM - else: - model_cls = HookedGPT2LMHeadModel - mon_model = model_cls.from_pretrained( - hf_model_id, attn_implementation="eager", torch_dtype=torch.float16, - ).to(device).eval() - mon_model.monitoring_engine = engine - - try: - from transformers import CompileConfig - with torch.no_grad(): - generate_with_monitoring( - mon_model, - input_ids=input_ids, attention_mask=attention_mask, - max_new_tokens=max_new_tokens, do_sample=False, - pad_token_id=pad_id, eos_token_id=eos_id, - logits_to_keep=0, - cache_implementation="static", - compile_config=CompileConfig(mode="reduce-overhead", fullgraph=False), - ) - finally: - engine.close() - - # --- Read DB --- - ch = CHClickhouseDriverReadOnly( - host=str(db_cfg_native.host), port=int(db_cfg_native.port), - username=str(db_cfg_native.username), password=str(db_cfg_native.password), - database=str(db_cfg_native.database), table=str(db_cfg_native.table), - secure=False, client_settings=None, decode_strings=True, - ) - try: - rows = ch.prefix_get((unique_run_model_id,), return_full_key_tuple=True) - finally: - ch.close() - - assert rows, f"No DB rows for model_id={unique_run_model_id!r}" - - shard_ranks = sorted({int(key[4]) for key, _t in rows}) - chosen_shard = 0 if 0 in shard_ranks else shard_ranks[0] - rows = [(k, t) for (k, t) in rows if int(k[4]) == chosen_shard] - - grouped: Dict[str, Dict[Tuple[int, str], List[Tuple[int, int, torch.Tensor]]]] = {} - for full_key, t_raw in rows: - _model_id, req_id, act_name_raw, layer_no_raw, _shard, s, e = full_key - layer_no, act_name = _canon_layer_and_act(str(act_name_raw), int(layer_no_raw)) - grouped.setdefault(str(req_id), {}).setdefault( - (layer_no, act_name), [] - ).append((int(s), int(e), t_raw.detach().cpu())) - - request_ids = sorted(grouped.keys(), key=_parse_request_id) - - # Build DB token_ids - db_token_ids_by_req: Dict[str, torch.Tensor] = {} - prompt_len_by_req: Dict[str, int] = {} - local_index_by_req: Dict[str, int] = {} - for req_id in request_ids: - _gid, local_i = _parse_request_id(req_id) - local_index_by_req[req_id] = local_i - hooks_map = grouped[req_id] - tok_chunks = sorted(hooks_map[(-1, "token_ids")], key=lambda x: (x[0], x[1])) - db_tok = merge_segments([t for _, _, t in tok_chunks], "token_ids").to(torch.long) - if db_tok.ndim != 1: - db_tok = db_tok.view(-1) - db_token_ids_by_req[req_id] = db_tok.cpu() - prompt_len_by_req[req_id] = int(hf_initial_prompt_tokens[local_i].numel()) - - # --- Eager HF reference (uncompiled, DynamicCache) --- - hf_eager = AutoModelForCausalLM.from_pretrained( - hf_model_id, attn_implementation="eager", torch_dtype=torch.float16, - ).to(device).eval() - num_layers = get_num_layers_from_config(hf_eager) - - hf_eager_refs = _hf_greedy_rollout_collect_all_batched( - hf_model=hf_eager, - input_ids_batch=input_ids, - attention_mask_batch=attention_mask, - max_new_tokens=max_new_tokens, - eos_token_id=eos_id, - pad_token_id=pad_id, - device=device, - want_hidden_states=True, - want_attentions=True, - ) - - # --- Comparisons --- - _EAGER_ATOL = 0.5 - _RESID_PRE_KEYS = ("blocks.hook_resid_pre", "layers.hook_resid_pre") - _ATTN_PATTERN_KEYS = ("blocks.attn.hook_pattern", "layers.self_attn.hook_pattern") - - def _sort_chunks(chunks): - return sorted(chunks, key=lambda x: (x[0], x[1])) - - def _close_enough(db_t, ref_t, ctx, atol=_EAGER_ATOL): - diff = (db_t.float() - ref_t.float()).abs() - max_abs = float(diff.max().item()) - assert max_abs <= atol, ( - f"{ctx} max_abs_diff={max_abs:.6f} > atol={atol}" - ) - - for req_id in sorted(request_ids, key=_parse_request_id): - local_i = local_index_by_req[req_id] - hooks_map = grouped[req_id] - plen = prompt_len_by_req[req_id] - db_tok = db_token_ids_by_req[req_id] - eref = hf_eager_refs[local_i] - - # --- token_ids: full comparison --- - min_len = min(int(db_tok.numel()), int(eref.token_ids.numel())) - match_len = 0 - for t in range(min_len): - if db_tok[t] != eref.token_ids[t]: - break - match_len = t + 1 - - with subtests.test(msg=f"{req_id}/cg_vs_eager/token_prefix"): - assert match_len >= plen, ( - f"compiled/eager diverge within prompt (match_len={match_len} plen={plen})" - ) - - with subtests.test(msg=f"{req_id}/cg_vs_eager/token_ids_full"): - assert db_tok.numel() == eref.token_ids.numel(), ( - f"token count mismatch db={db_tok.numel()} eager={eref.token_ids.numel()}" - ) - assert torch.equal(db_tok[:match_len], eref.token_ids[:match_len]), ( - f"token_ids mismatch in matching prefix (len={match_len})" - ) - - if match_len <= 0: - continue - - # --- final_logits --- - if eref.final_logits is not None and (-1, "final_logits") in hooks_map: - chunks = _sort_chunks(hooks_map[(-1, "final_logits")]) - db_logits = merge_segments([t for _, _, t in chunks], "final_logits") - ref_logits = eref.final_logits - # Compare over matching prefix - ml = min(db_logits.shape[0], ref_logits.shape[0], match_len) - with subtests.test(msg=f"{req_id}/cg_vs_eager/final_logits"): - _close_enough(db_logits[:ml], ref_logits[:ml], - f"{req_id} final_logits") - - # --- hook_embed --- - if (-1, "hook_embed") in hooks_map: - chunks = _sort_chunks(hooks_map[(-1, "hook_embed")]) - db_emb = merge_segments([t for _, _, t in chunks], "hook_embed") - if eref.hidden_states and len(eref.hidden_states) > 0: - # hidden_states[0] is the embedding output for HF models - # that include it (GPT2: embed+pos, Qwen3: embed only) - pass # embed comparison needs model-specific reference - with subtests.test(msg=f"{req_id}/cg_vs_eager/hook_embed"): - assert db_emb.ndim >= 2, f"hook_embed unexpected shape {db_emb.shape}" - - # --- hook_pos_embed (GPT2 only) --- - if (-1, "hook_pos_embed") in hooks_map: - chunks = _sort_chunks(hooks_map[(-1, "hook_pos_embed")]) - db_pos = merge_segments([t for _, _, t in chunks], "hook_pos_embed") - with subtests.test(msg=f"{req_id}/cg_vs_eager/hook_pos_embed"): - assert db_pos.ndim >= 2, f"hook_pos_embed unexpected shape {db_pos.shape}" - - # --- per-layer: resid_pre + attn_pattern --- - if not eref.hidden_states: - continue - - for layer_no in range(num_layers): - # resid_pre - key = next(((layer_no, k) for k in _RESID_PRE_KEYS if (layer_no, k) in hooks_map), None) - if key is not None and layer_no < len(eref.hidden_states): - hs_eager = eref.hidden_states[layer_no][:match_len, :] - chunks = _sort_chunks(hooks_map[key]) - db_t = merge_segments([t for _, _, t in chunks], key[1])[:match_len, :] - with subtests.test(msg=f"{req_id}/cg_vs_eager/layer{layer_no}/resid_pre"): - assert db_t.shape == hs_eager.shape, ( - f"shape mismatch db={db_t.shape} eager={hs_eager.shape}" - ) - _close_enough(db_t, hs_eager, - f"resid_pre layer={layer_no}") - - # attn_pattern: compare step-by-step, trimming static cache padding. - # DB (static cache): each step has kv_dim = max_len (padded). - # Eager ref (dynamic cache): each step has kv_dim = actual kv_len. - # We compare per-chunk, trimming DB's kv_dim to the eager ref's. - key = next(((layer_no, k) for k in _ATTN_PATTERN_KEYS if (layer_no, k) in hooks_map), None) - if key is not None and eref.attn_pattern and layer_no < len(eref.attn_pattern): - pat_eager = eref.attn_pattern[layer_no] # [H, T, T] from rollout - chunks = _sort_chunks(hooks_map[key]) - # Rebuild step-by-step: each chunk covers [start:end] token positions - all_ok = True - fail_msg = "" - for start, end, t_chunk in chunks: - if start >= match_len: - break - end_clamp = min(end, match_len) - q_len = end_clamp - start - # t_chunk: [H, q_len, kv_dim_padded] or [1, H, q_len, kv_dim_padded] - db_c = t_chunk - if db_c.ndim == 4 and db_c.shape[0] == 1: - db_c = db_c.squeeze(0) - db_c = db_c[:, :q_len, :] # trim q_len if chunk extends beyond match_len - # kv_dim for these rows: tokens 0..end_clamp-1 attended to - # keys 0..end_clamp-1 (causal). Trim kv to end_clamp. - kv_valid = end_clamp - db_c = db_c[:, :, :kv_valid] - # Corresponding slice from eager ref - ref_c = pat_eager[:, start:end_clamp, :kv_valid] - if db_c.shape != ref_c.shape: - all_ok = False - fail_msg = (f"shape mismatch at [{start}:{end_clamp}] " - f"db={db_c.shape} eager={ref_c.shape}") - break - diff = (db_c.float() - ref_c.float()).abs() - max_abs = float(diff.max().item()) - if max_abs > _EAGER_ATOL: - all_ok = False - fail_msg = (f"value mismatch at [{start}:{end_clamp}] " - f"max_abs={max_abs:.6f} > atol={_EAGER_ATOL}") - break - with subtests.test(msg=f"{req_id}/cg_vs_eager/layer{layer_no}/attn_pattern"): - assert all_ok, ( - f"attn_pattern layer={layer_no}: {fail_msg}" - ) - - # --- resid_final (global, last layer's pre-norm residual) --- - # HF's output_hidden_states[-1] is post-final-norm, not pre-norm. - # Only check shape and presence. - key = next(((-1, k) for k in ("hook_resid_final",) if (-1, k) in hooks_map), None) - if key is not None: - chunks = _sort_chunks(hooks_map[key]) - db_rf = merge_segments([t for _, _, t in chunks], key[1])[:match_len, :] - with subtests.test(msg=f"{req_id}/cg_vs_eager/resid_final"): - assert db_rf.shape[0] == match_len, ( - f"resid_final token count mismatch db={db_rf.shape[0]} expected={match_len}" - ) \ No newline at end of file + cr = run_single(matrix_argv_from_env( + "hf", "allclose", mode="cuda_graph", default_tolerance="0.5")) + _assert_cell(subtests, cr) diff --git a/tests/test_e2e_lib.py b/tests/test_e2e_lib.py index 0f03017be..d93d97e44 100644 --- a/tests/test_e2e_lib.py +++ b/tests/test_e2e_lib.py @@ -307,3 +307,62 @@ def test_main_dry_run_no_side_effects(self, capsys): def test_main_empty_axis_returns_2(self): from tests.e2e_matrix import main assert main(["--backend", "", "--dry-run"]) == 2 + + +class TestWrapperTranslation: + """matrix_argv_from_env / run_single — the thin-wrapper surface (plan §5).""" + + def test_env_to_argv_defaults(self): + from tests.e2e_matrix import matrix_argv_from_env, build_parser, build_cells + argv = matrix_argv_from_env("vllm", "bitwise", env={}) + cell = build_cells(build_parser().parse_args(argv))[0] + assert cell.backend == "vllm" and cell.standard == "bitwise" + assert cell.model == "gpt2" and cell.mode == "eager" + assert cell.hooks == "vllm-full" and cell.ring_mb == 4096 + + def test_env_enforce_eager_maps_mode(self): + from tests.e2e_matrix import matrix_argv_from_env, build_parser, build_cells + env = {"E2E_ENFORCE_EAGER": "0", "E2E_MODEL": "qwen3"} + cell = build_cells(build_parser().parse_args( + matrix_argv_from_env("vllm", "row_count", env=env)))[0] + assert cell.mode == "cuda_graph" and cell.model == "qwen3" + + def test_explicit_mode_overrides_enforce_eager(self): + from tests.e2e_matrix import matrix_argv_from_env, build_parser, build_cells + # HF cuda-graph wrapper forces mode even though E2E_ENFORCE_EAGER=1. + env = {"E2E_ENFORCE_EAGER": "1"} + cell = build_cells(build_parser().parse_args( + matrix_argv_from_env("hf", "allclose", mode="cuda_graph", env=env)))[0] + assert cell.mode == "cuda_graph" + + def test_hook_selection_precedence(self): + from tests.e2e_matrix import matrix_argv_from_env + # public E2E_HOOK_SELECTION wins over internal DMX_HOOK_SELECTION + argv = matrix_argv_from_env("vllm", "bitwise", env={ + "E2E_HOOK_SELECTION": "q", "DMX_HOOK_SELECTION": "k"}) + assert argv[argv.index("--hooks") + 1] == "q" + # falls back to DMX_HOOK_SELECTION when public unset + argv = matrix_argv_from_env("vllm", "bitwise", env={"DMX_HOOK_SELECTION": "k"}) + assert argv[argv.index("--hooks") + 1] == "k" + + def test_default_tolerance_passthrough(self): + from tests.e2e_matrix import matrix_argv_from_env + argv = matrix_argv_from_env("hf", "allclose", mode="cuda_graph", + default_tolerance="0.5", env={}) + assert argv[argv.index("--tolerance") + 1] == "0.5" + # explicit env overrides the default + argv = matrix_argv_from_env("hf", "allclose", default_tolerance="0.5", + env={"E2E_TOLERANCE": "0.01"}) + assert argv[argv.index("--tolerance") + 1] == "0.01" + + def test_run_single_rejects_multi_cell(self, monkeypatch): + # matrix_argv_from_env always yields one cell; guard the invariant. + from tests import e2e_matrix + args = e2e_matrix.build_parser().parse_args( + ["--backend", "hf,vllm", "--standard", "row_count"]) + monkeypatch.setattr(e2e_matrix, "run_cell", lambda *a, **k: None) + with pytest.raises(ValueError, match="exactly 1 cell"): + # build_cells gives 2 -> run_single must refuse + cells = e2e_matrix.build_cells(args) + assert len(cells) == 2 + e2e_matrix.run_single(["--backend", "hf,vllm", "--standard", "row_count"]) diff --git a/tests/test_vllm_identical.py b/tests/test_vllm_identical.py index 41087a516..1b827d809 100644 --- a/tests/test_vllm_identical.py +++ b/tests/test_vllm_identical.py @@ -1,44 +1,28 @@ -"""vLLM identical check — bitwise tensor comparison between ref model -(GPU buffer D2D capture → disk) and monitored model (ring transport → ClickHouse). +"""vLLM identical check — thin wrapper over the configurable matrix (plan §5). -Four steps, parent never touches CUDA: - 0. Sanity check: original vs ref model logprobs (informational, never fails) - 1. Reference run (RefDiskWorker, D2D capture → disk) - 2. Monitored run (DMXGPUWorker, ring transport → ClickHouse) - 3. Comparator (CPU only, logprob comparison + bitwise tensor check) +Bitwise tensor comparison between the reference model (GPU buffer D2D +capture -> disk) and the monitored model (ring transport -> ClickHouse). +The orchestration (enable_ref_hooks -> vllm_ref_runner -> vllm_monitored_runner +-> vllm_identical_comparator) now lives in :mod:`tests.e2e_matrix`; this +wrapper drives the matrix's vLLM ``bitwise`` cell and asserts on its checks. -Environment variables: - E2E_MODEL "gpt2" (default) or "qwen3" - E2E_NUM_PROMPTS Number of prompts (default 8) - 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_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 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) +The test name is preserved because ``tests/tools/verify_vllm.sh`` and +``tests/tools/identical_vllm.sh`` invoke this file and thread the model / +ring-size / hook-selection env vars the matrix wrapper reads: -Requires: - - ClickHouse running - - VLLM_DISABLE_COMPILE_CACHE=1 (set automatically) + E2E_MODEL, E2E_ENFORCE_EAGER, E2E_DTYPE, E2E_RING_PAYLOAD_MB, + E2E_RING_PINNED_MB, E2E_HOOK_SELECTION (-> internal DMX_HOOK_SELECTION), + E2E_REF_MAX_LEN, E2E_TP_SIZE, DMX_DB_HOST, DMX_DB_PORT. Usage: python -m pytest tests/test_vllm_identical.py -q -s """ - -import json -import os -import shutil -import subprocess -import sys -import tempfile +from __future__ import annotations import pytest -import torch + +from tests._requirements import require_cuda, require_clickhouse, require_vllm +from tests.e2e_matrix import matrix_argv_from_env, run_single pytestmark = [ pytest.mark.gpu, @@ -47,215 +31,16 @@ pytest.mark.e2e, ] -_MODEL_REF_FILES = { - "gpt2": "gpt2_ref.py", - "qwen2_moe": "qwen2_moe_ref.py", - "qwen3": "qwen3_ref.py", - "llama": "llama_ref.py", -} - - -@pytest.mark.skipif( - not torch.backends.cuda.is_built(), reason="CUDA not built") -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_HOOK_SELECTION", "vllm-full") - max_len = int(os.environ.get("E2E_REF_MAX_LEN", "8192")) - enforce_eager = os.environ.get("E2E_ENFORCE_EAGER", "1") - - project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - models_dir = os.path.join( - project_root, "integration", "vllm", "vllm", - "model_executor", "models") - - ref_filename = _MODEL_REF_FILES.get(model_key) - if ref_filename is None: - pytest.skip(f"No ref model for {model_key}") - model_file = os.path.join(models_dir, ref_filename) - - keep_artifacts = os.environ.get("E2E_KEEP_ARTIFACTS", "0") == "1" - dump_compiled = os.environ.get("E2E_DUMP_COMPILED", "0") == "1" - artifact_dir = os.environ.get("E2E_ARTIFACT_DIR") - if artifact_dir: - run_dir = os.path.abspath(artifact_dir) - os.makedirs(run_dir, exist_ok=True) - else: - run_dir = tempfile.mkdtemp(prefix="vllm_identical_") - ref_dir = os.path.join(run_dir, "ref") - mon_dir = os.path.join(run_dir, "mon") - config_file = os.path.join(ref_dir, "ref_config.json") - result_file = os.path.join(run_dir, "result.json") - backup_file = os.path.join(run_dir, f"{ref_filename}.bak") - orig_logprobs_file = os.path.join(run_dir, "logprobs_orig.pt") - ref_logprobs_file = os.path.join(run_dir, "logprobs_ref.pt") - - print(f"\n{'=' * 60}") - print(f" vLLM identical check") - print(f" model={model_key} hooks={hooks} eager={enforce_eager}") - print(f" ref_max_len={max_len}") - print(f"{'=' * 60}") - - # Build env for subprocesses (inherit + add our vars) - sub_env = dict(os.environ) - sub_env["VLLM_DISABLE_COMPILE_CACHE"] = "1" - sub_env["E2E_ENFORCE_EAGER"] = enforce_eager - sub_env["DMX_HOOK_SELECTION"] = hooks - if dump_compiled: - sub_env["TORCH_LOGS"] = "+output_code" - - try: - # Backup ref model file - shutil.copy2(model_file, backup_file) - - # Enable hooks via preprocessor - print("\n [0/3] Enabling ref hooks...", flush=True) - os.makedirs(ref_dir, exist_ok=True) - sys.path.insert(0, models_dir) - from enable_ref_hooks import enable_ref_hooks - enable_ref_hooks( - model_file=model_file, - hooks=hooks, - max_len=max_len, - output_dir=ref_dir, - config_out=config_file, - ) - - # Step 0: Sanity check — original vs ref model logprobs - # Runs AFTER enabling hooks to verify D2D copies don't affect output. - print("\n [0/4] Sanity check: original model logprobs...", flush=True) - r0a = subprocess.run( - [sys.executable, "-m", "tests.vllm_logprob_runner", - "--output", orig_logprobs_file], - env=sub_env, capture_output=True, text=True, cwd=project_root, - ) - print(r0a.stdout[-1000:] if r0a.stdout else "", flush=True) - if dump_compiled and r0a.stderr: - with open(os.path.join(run_dir, "compile_orig.log"), "w") as f: - f.write(r0a.stderr) - if r0a.returncode != 0: - print(r0a.stderr[-2000:] if r0a.stderr else "", flush=True) - print(" WARNING: original logprob run failed, skipping sanity check") - orig_logprobs_file = None - - print(" [0/4] Sanity check: ref model logprobs...", flush=True) - ref_lp_env = dict(sub_env) - ref_lp_env["REF_CONFIG"] = config_file - r0b = subprocess.run( - [sys.executable, "-m", "tests.vllm_logprob_runner", - "--output", ref_logprobs_file, "--ref"], - env=ref_lp_env, capture_output=True, text=True, cwd=project_root, - ) - print(r0b.stdout[-1000:] if r0b.stdout else "", flush=True) - if dump_compiled and r0b.stderr: - with open(os.path.join(run_dir, "compile_ref_logprob.log"), "w") as f: - f.write(r0b.stderr) - if r0b.returncode != 0: - print(r0b.stderr[-2000:] if r0b.stderr else "", flush=True) - print(" WARNING: ref logprob run failed, skipping sanity check") - ref_logprobs_file = None - - # Step 0c: Monitored model logprobs (baseline vs monitored comparison) - mon_logprobs_file = os.path.join(run_dir, "logprobs_mon.pt") - print(" [0/4] Sanity check: monitored model logprobs...", flush=True) - r0c = subprocess.run( - [sys.executable, "-m", "tests.vllm_logprob_runner", - "--output", mon_logprobs_file, "--monitored"], - env=sub_env, capture_output=True, text=True, cwd=project_root, - ) - print(r0c.stdout[-1000:] if r0c.stdout else "", flush=True) - if dump_compiled and r0c.stderr: - with open(os.path.join(run_dir, "compile_mon_logprob.log"), "w") as f: - f.write(r0c.stderr) - if r0c.returncode != 0: - print(r0c.stderr[-2000:] if r0c.stderr else "", flush=True) - print(" WARNING: monitored logprob run failed, skipping") - mon_logprobs_file = None - - # Step 1: Reference run - print("\n [1/4] Reference run (RefDiskWorker)...", flush=True) - ref_env = dict(sub_env) - ref_env["REF_CONFIG"] = config_file - r1 = subprocess.run( - [sys.executable, "-m", "tests.vllm_ref_runner", - "--output-dir", ref_dir], - env=ref_env, capture_output=True, text=True, cwd=project_root, - ) - print(r1.stdout[-2000:] if r1.stdout else "", flush=True) - if keep_artifacts and r1.stdout: - with open(os.path.join(run_dir, "stdout_ref_runner.log"), "w") as f: - f.write(r1.stdout) - if dump_compiled and r1.stderr: - with open(os.path.join(run_dir, "compile_ref_runner.log"), "w") as f: - f.write(r1.stderr) - if r1.returncode != 0: - print(r1.stderr[-3000:] if r1.stderr else "", flush=True) - pytest.fail(f"Ref runner failed (rc={r1.returncode})") - - # Restore ref model from backup (before monitored run) - shutil.copy2(backup_file, model_file) - - # Step 2: Monitored run - print("\n [2/4] Monitored run (DMXGPUWorker)...", flush=True) - r2 = subprocess.run( - [sys.executable, "-m", "tests.vllm_monitored_runner", - "--output-dir", mon_dir], - env=sub_env, capture_output=True, text=True, cwd=project_root, - ) - print(r2.stdout[-2000:] if r2.stdout else "", flush=True) - if keep_artifacts and r2.stdout: - with open(os.path.join(run_dir, "stdout_mon_runner.log"), "w") as f: - f.write(r2.stdout) - if dump_compiled and r2.stderr: - with open(os.path.join(run_dir, "compile_mon_runner.log"), "w") as f: - f.write(r2.stderr) - if r2.returncode != 0: - print(r2.stderr[-3000:] if r2.stderr else "", flush=True) - pytest.fail(f"Monitored runner failed (rc={r2.returncode})") - - # Step 3: Comparator (includes logprob sanity check if available) - print("\n [3/4] Comparing (bitwise check)...", flush=True) - cmp_cmd = [ - sys.executable, "-m", "tests.vllm_identical_comparator", - "--ref-config", config_file, - "--mon-dir", mon_dir, - "--result-file", result_file, - ] - if orig_logprobs_file and os.path.exists(orig_logprobs_file): - cmp_cmd += ["--orig-logprobs", orig_logprobs_file] - if ref_logprobs_file and os.path.exists(ref_logprobs_file): - cmp_cmd += ["--ref-logprobs", ref_logprobs_file] - if mon_logprobs_file and os.path.exists(mon_logprobs_file): - cmp_cmd += ["--mon-logprobs", mon_logprobs_file] - r3 = subprocess.run( - cmp_cmd, - env=sub_env, capture_output=True, text=True, cwd=project_root, - ) - if r3.stdout: - # Always print LOGPROBS summary and PASS/FAIL lines first - for line in r3.stdout.splitlines(): - if "[LOGPROBS" in line or "ALL PASSED" in line or "FAILED (" in line: - print(line, flush=True) - # Then print tail for hidden state details - print(r3.stdout[-2000:], flush=True) - if r3.returncode != 0: - print(r3.stderr[-3000:] if r3.stderr else "", flush=True) - pytest.fail(f"Comparator failed (rc={r3.returncode})") - - # Report via subtests - with open(result_file) as f: - results = json.load(f) - - for test in results["tests"]: - with subtests.test(test["name"]): - assert test["passed"], test.get("detail", "") - finally: - # Always restore ref model file - if os.path.exists(backup_file): - shutil.copy2(backup_file, model_file) - if keep_artifacts: - print(f"\n [kept] run_dir = {run_dir}", flush=True) - else: - shutil.rmtree(run_dir, ignore_errors=True) +@require_cuda() +@require_vllm() +@require_clickhouse() +def test_vllm_identical(subtests) -> None: + """Bitwise: reference D2D buffers (disk) vs ring transport (ClickHouse).""" + cr = run_single(matrix_argv_from_env("vllm", "bitwise")) + if cr.error: + pytest.fail(f"matrix cell errored: {cr.error}") + assert cr.checks, "matrix produced no checks" + for chk in cr.checks: + with subtests.test(chk.name): + assert chk.passed, chk.detail diff --git a/tests/test_vllm_rowcnt.py b/tests/test_vllm_rowcnt.py index da0d466be..ee9092558 100644 --- a/tests/test_vllm_rowcnt.py +++ b/tests/test_vllm_rowcnt.py @@ -1,45 +1,27 @@ -"""vLLM E2E correctness test — three subprocesses, parent never touches CUDA. +"""vLLM row-count check — thin wrapper over the configurable matrix (plan §5). - 1. Reference: original model + FullHiddenStatesConnector -> disk - (skipped for models not supported by extract_hidden_states, e.g. GPT-2) - 2. Monitored: hooked model + DMXGPUWorker + ring transport -> ClickHouse - 3. Comparator: reads both, validates row counts + value comparison +Runs the monitored model (hooked + ring transport -> ClickHouse) and +validates schema + per-hook row counts (plus value comparison against the +reference when the model is supported by extract_hidden_states). The +orchestration (vllm_monitored_runner -> vllm_rowcnt_comparator) now lives in +:mod:`tests.e2e_matrix`; this wrapper drives the matrix's vLLM ``row_count`` +cell and asserts on its checks. -Environment variables: - E2E_MODEL "gpt2" (default) or "qwen3" - E2E_NUM_PROMPTS Number of prompts (default 8) - E2E_MAX_NEW_TOKENS Tokens to generate per prompt (default 20) - 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 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. - E2E_TOLERANCE Max abs diff tolerance (default "0.01") - DMX_DB_HOST ClickHouse host (default "localhost") - DMX_DB_PORT ClickHouse port (default 9000) - -Requires: - - ClickHouse running on DMX_DB_HOST:DMX_DB_PORT - - VLLM_DISABLE_COMPILE_CACHE=1 (set automatically) - - LD_PRELOAD for libstdc++ if needed (caller's responsibility) +The test name is preserved because ``tests/tools/verify_vllm.sh`` invokes +this file and threads the model / ring-size / tolerance env vars the matrix +wrapper reads (E2E_MODEL, E2E_ENFORCE_EAGER, E2E_RING_PAYLOAD_MB, +E2E_RING_PINNED_MB, E2E_TOLERANCE, DMX_DB_HOST, DMX_DB_PORT). Usage: python -m pytest tests/test_vllm_rowcnt.py -q -s - E2E_MODEL=qwen3 E2E_COMPARE_LAYERS=all python -m pytest tests/test_vllm_rowcnt.py -q -s + E2E_MODEL=qwen3 python -m pytest tests/test_vllm_rowcnt.py -q -s """ - -import json -import os -import shutil -import subprocess -import sys -import tempfile +from __future__ import annotations import pytest -import torch + +from tests._requirements import require_cuda, require_clickhouse, require_vllm +from tests.e2e_matrix import matrix_argv_from_env, run_single pytestmark = [ pytest.mark.gpu, @@ -48,73 +30,16 @@ pytest.mark.e2e, ] -_MODEL_ALIASES = { - "gpt2": "gpt2", - "qwen2_moe": "Qwen/Qwen1.5-MoE-A2.7B", - "qwen3": "Qwen/Qwen3-0.6B", -} - -@pytest.mark.skipif( - not torch.backends.cuda.is_built(), reason="CUDA not built") -def test_vllm_rowcnt(subtests): - """vLLM row-count validation: monitored run + row-count check.""" - - 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") - result_file = os.path.join(run_dir, "result.json") - project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - - # ref_dir still needed by comparator (with skipped marker) - os.makedirs(ref_dir, exist_ok=True) - with open(os.path.join(ref_dir, "meta.json"), "w") as f: - json.dump({"skipped": True}, f) - - print(f"\n{'=' * 60}") - print(f" vLLM row-count test") - print(f" model={model_id}") - print(f"{'=' * 60}") - - try: - # Step 1: Monitored run (hooked model + ring transport) - print("\n [1/2] Monitored run (hooked model + DMXGPUWorker)...", flush=True) - r2 = subprocess.run( - [sys.executable, "-m", "tests.vllm_monitored_runner", - "--output-dir", mon_dir], - 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:]}") - - # Step 2: Comparator (CPU only, row-count validation) - print(" [2/2] Checking row counts...", flush=True) - r3 = subprocess.run( - [sys.executable, "-m", "tests.vllm_rowcnt_comparator", - "--ref-dir", ref_dir, - "--mon-dir", mon_dir, - "--result-file", result_file], - env=sub_env, capture_output=True, text=True, cwd=project_root, - ) - if r3.returncode != 0: - pytest.fail(f"Comparator failed:\n{r3.stderr[-2000:]}") - - # Read results - with open(result_file) as f: - results = json.load(f) - - # Report via subtests - for test in results["tests"]: - with subtests.test(test["name"]): - assert test["passed"], test.get("detail", "") - finally: - shutil.rmtree(run_dir, ignore_errors=True) +@require_cuda() +@require_vllm() +@require_clickhouse() +def test_vllm_rowcnt(subtests) -> None: + """vLLM row-count validation: monitored run + schema / row-count checks.""" + cr = run_single(matrix_argv_from_env("vllm", "row_count")) + if cr.error: + pytest.fail(f"matrix cell errored: {cr.error}") + assert cr.checks, "matrix produced no checks" + for chk in cr.checks: + with subtests.test(chk.name): + assert chk.passed, chk.detail From ac64e70fb2bc0cfb4c7ce275747ef0313699c06b Mon Sep 17 00:00:00 2001 From: SieDeta Date: Thu, 11 Jun 2026 23:54:42 +0700 Subject: [PATCH 6/6] test: multi_gpu TP=2 E2E coverage (C6) The `-m multi_gpu` suite documented in docs/testing.md collected zero tests -- no test carried the `multi_gpu` marker, so TP coverage existed only in the CPU shape-math unit test and the manual tests/tools sweeps. Add tests/test_e2e_tp2.py: two wrappers that drive the configurable matrix at tp=2 (vLLM transport_bitwise + HF eager allclose), forcing E2E_TP_SIZE=2 via matrix_argv_from_env's env override. Marked multi_gpu/gpu/e2e/clickhouse (+vllm/+hf) and skip-guarded with require_gpus(2)/require_vllm/require_clickhouse so a <2-GPU runner skips with a reason instead of failing. Co-Authored-By: Claude Opus 4.8 --- tests/test_e2e_tp2.py | 83 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 83 insertions(+) create mode 100644 tests/test_e2e_tp2.py diff --git a/tests/test_e2e_tp2.py b/tests/test_e2e_tp2.py new file mode 100644 index 000000000..f4db8f41c --- /dev/null +++ b/tests/test_e2e_tp2.py @@ -0,0 +1,83 @@ +"""Multi-GPU (TP=2) E2E smoke -- the ``multi_gpu`` suite's TP coverage. + +The configurable matrix (:mod:`tests.e2e_matrix`) treats tensor-parallel size +as a first-class axis (``--tp`` / ``E2E_TP_SIZE``), but the single-GPU wrappers +(``test_vllm_identical`` / ``test_e2e_correctness_vs_hf``) all drive ``tp=1``. +This module drives the same matrix cells at ``tp=2`` so the documented +``-m multi_gpu`` suite (docs/testing.md) actually exercises TP sharding rather +than collecting nothing. + +TP=2 is "where meaningful" for the sharded model: with two ranks the +attention/expert projections are split across GPUs, so the reference-vs-monitored +comparison validates that the ring transport reassembles per-rank shards +correctly. ``qwen3`` is the default model (GQA + a non-trivial hidden size makes +the sharding observable); ``gpt2`` is too small for TP to be interesting. Both +backends honor ``E2E_MODEL`` for an override. + +Skip-guarded (``tests/_requirements``) so a runner with <2 GPUs, no vLLM, or no +ClickHouse skips with a reason instead of failing the job. +""" +from __future__ import annotations + +import os + +import pytest + +from tests._requirements import ( + require_clickhouse, + require_gpus, + require_vllm, +) +from tests.e2e_matrix import matrix_argv_from_env, run_single + +pytestmark = [ + pytest.mark.multi_gpu, + pytest.mark.gpu, + pytest.mark.e2e, + pytest.mark.clickhouse, +] + + +def _tp2_env(default_model: str = "qwen3") -> dict: + """os.environ with TP forced to 2 (model still overridable via E2E_MODEL).""" + env = dict(os.environ) + env["E2E_TP_SIZE"] = "2" + env.setdefault("E2E_MODEL", default_model) + return env + + +def _assert_cell(subtests, cr) -> None: + """Fail on a dispatch error; report each matrix check as a subtest.""" + if cr.error: + pytest.fail(f"matrix cell errored: {cr.error}") + assert cr.checks, "matrix produced no checks" + for chk in cr.checks: + with subtests.test(chk.name): + assert chk.passed, chk.detail + + +@pytest.mark.vllm +@require_gpus(2) +@require_vllm() +@require_clickhouse() +def test_vllm_identical_tp2(subtests) -> None: + """vLLM TP=2 transport-bitwise: reference D2D buffers vs ring -> ClickHouse. + + Equivalent matrix cell: ``--backend vllm --standard transport_bitwise --tp 2``. + The bitwise standard stays exact under TP (sharding is a layout change, not a + numeric one), so any per-rank reassembly bug surfaces as a non-zero max_abs. + """ + argv = matrix_argv_from_env("vllm", "transport_bitwise", env=_tp2_env()) + _assert_cell(subtests, run_single(argv)) + + +@pytest.mark.hf +@require_gpus(2) +@require_clickhouse() +def test_e2e_correctness_hf_tp2(subtests) -> None: + """HF TP=2 eager: hooked model (ring -> ClickHouse) vs original model. + + Equivalent matrix cell: ``--backend hf --mode eager --standard allclose --tp 2``. + """ + argv = matrix_argv_from_env("hf", "allclose", mode="eager", env=_tp2_env()) + _assert_cell(subtests, run_single(argv))