(#3) test: shared E2E lib + configurable matrix

tests/hf_comparator.py

@@ -23,8 +23,16 @@ def main():
     p.add_argument("--ref-dir", required=True)
     p.add_argument("--mon-dir", required=True)
     p.add_argument("--result-file", required=True)
+    p.add_argument("--standard", default="allclose",
+                   help="comparison standard (only 'allclose' is implemented)")
     args = p.parse_args()
 
+    if args.standard != "allclose":
+        raise SystemExit(
+            f"hf_comparator: standard={args.standard!r} is not implemented; "
+            "only 'allclose' is supported"
+        )
+
     from tests.hf_reference import (
         _HFRef,
         _load_hf_refs_from_disk,

tests/hf_monitored_runner.py

@@ -45,7 +45,7 @@ def main():
         model_cls = HookedGPT2LMHeadModel
 
     model = model_cls.from_pretrained(
-        hf_model_id, attn_implementation="eager", torch_dtype=torch.float16
+        hf_model_id, attn_implementation="eager", dtype=torch.float16
     ).to(device).eval()
 
     tokenizer = AutoTokenizer.from_pretrained(hf_model_id)
@@ -99,8 +99,15 @@ def main():
     from monitoring._native_engine import RingConfig
     ring_cfg = RingConfig()
     ring_cfg.task_ring_entries = int(os.environ.get("E2E_RING_TASK_ENTRIES", "16384"))
-    ring_cfg.payload_ring_bytes = int(os.environ.get("E2E_RING_PAYLOAD_BYTES", str(4 * 1024**3)))
-    ring_cfg.pinned_staging_bytes = int(os.environ.get("E2E_RING_PINNED_BYTES", str(4 * 1024**3)))
+    # E2E_RING_PAYLOAD_MB is the matrix-level knob (set by cell_env via --ring-mb).
+    # E2E_RING_PAYLOAD_BYTES is a fine-grained override; if neither is set, default
+    # to 256 MB so the runner does not OOM on GPUs smaller than the ring buffer.
+    _payload_mb = int(os.environ.get("E2E_RING_PAYLOAD_MB", "256"))
+    ring_cfg.payload_ring_bytes = int(os.environ.get("E2E_RING_PAYLOAD_BYTES",
+                                                      str(_payload_mb * 1024 * 1024)))
+    _pinned_mb = int(os.environ.get("E2E_RING_PINNED_MB", "256"))
+    ring_cfg.pinned_staging_bytes = int(os.environ.get("E2E_RING_PINNED_BYTES",
+                                                        str(_pinned_mb * 1024 * 1024)))
     ring_cfg.drain_poll_timeout_us = int(os.environ.get("E2E_DRAIN_POLL_TIMEOUT_US", "100"))
     ring_cfg.clone_slices = int(os.environ.get("E2E_CLONE_SLICES", "0")) != 0
     ring_cfg.insert_queue_max_bytes = int(os.environ.get("E2E_INSERT_QUEUE_MAX_BYTES", str(512 * 1024**2)))

tests/hf_reference_runner.py

@@ -33,7 +33,7 @@ def main():
 
     # Load ORIGINAL model (NOT hooked)
     model = AutoModelForCausalLM.from_pretrained(
-        hf_model_id, attn_implementation="eager", torch_dtype=torch.float16
+        hf_model_id, attn_implementation="eager", dtype=torch.float16
     ).to(device).eval()
 
     tokenizer = AutoTokenizer.from_pretrained(hf_model_id)

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 "<gid>:<row>" 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"]

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 "<group_id>:<row_index>" -- 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 ``"<group_id>:<row_index>"`` 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)

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_<name>`` 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 ``<database>.<table>`` 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