From 46e270be925b6ea7dc958d3e140990144d469c19 Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Tue, 25 Aug 2026 17:25:24 -0400 Subject: [PATCH 01/20] Phase 4 --- benchmarks/bench_capture_catalog.py | 116 +++ benchmarks/bench_capture_pack.py | 241 ++++++ benchmarks/bench_capture_pipeline.py | 260 ++++++ benchmarks/bench_clickhouse_host.py | 10 +- benchmarks/bench_garage_upload.py | 287 +++++++ docs/architecture.md | 8 + docs/benchmarks.md | 19 + docs/capture-storage-design.md | 741 +++++++++++++++++ docs/capture-storage-pipeline.html | 198 +++++ pyproject.toml | 4 + src/dmi/storage/capture/__init__.py | 148 ++++ src/dmi/storage/capture/catalog.py | 288 +++++++ src/dmi/storage/capture/clickhouse_catalog.py | 190 +++++ src/dmi/storage/capture/filesystem.py | 212 +++++ src/dmi/storage/capture/model.py | 372 +++++++++ src/dmi/storage/capture/pack.py | 526 ++++++++++++ src/dmi/storage/capture/pipeline.py | 586 ++++++++++++++ src/dmi/storage/capture/reader.py | 198 +++++ src/dmi/storage/capture/s3.py | 421 ++++++++++ src/dmi/storage/capture/spool.py | 518 ++++++++++++ tests/ring/Makefile | 59 ++ tests/ring/test_null_mode.cu | 288 +++++++ tests/ring/test_producer.cu | 462 +++++++++++ tests/ring/test_ring_engine.cu | 375 +++++++++ tests/ring/test_rings.cu | 758 ++++++++++++++++++ tests/test_capture_catalog_benchmark.py | 27 + tests/test_capture_catalog_indexer.py | 267 ++++++ tests/test_capture_pack_benchmark.py | 87 ++ tests/test_capture_parallel_upload.py | 208 +++++ tests/test_capture_pipeline.py | 370 +++++++++ tests/test_capture_pipeline_benchmark.py | 74 ++ tests/test_capture_s3.py | 290 +++++++ tests/test_capture_spool.py | 194 +++++ tests/test_capture_storage.py | 387 +++++++++ tests/test_clickhouse_capture_catalog.py | 121 +++ tests/test_clickhouse_catalog_live.py | 60 ++ tests/test_clickhouse_host_benchmark.py | 34 + tests/test_garage_live.py | 95 +++ tests/test_garage_upload_benchmark.py | 42 + tests/test_internal_mapper.py | 460 +++++++++++ tests/test_internal_mapper_e2e.py | 199 +++++ tests/tools/check_package.py | 12 + tests/tools/run_garage_live.py | 161 ++++ 43 files changed, 10370 insertions(+), 3 deletions(-) create mode 100644 benchmarks/bench_capture_catalog.py create mode 100644 benchmarks/bench_capture_pack.py create mode 100644 benchmarks/bench_capture_pipeline.py create mode 100644 benchmarks/bench_garage_upload.py create mode 100644 docs/capture-storage-design.md create mode 100644 docs/capture-storage-pipeline.html create mode 100644 src/dmi/storage/capture/__init__.py create mode 100644 src/dmi/storage/capture/catalog.py create mode 100644 src/dmi/storage/capture/clickhouse_catalog.py create mode 100644 src/dmi/storage/capture/filesystem.py create mode 100644 src/dmi/storage/capture/model.py create mode 100644 src/dmi/storage/capture/pack.py create mode 100644 src/dmi/storage/capture/pipeline.py create mode 100644 src/dmi/storage/capture/reader.py create mode 100644 src/dmi/storage/capture/s3.py create mode 100644 src/dmi/storage/capture/spool.py create mode 100644 tests/ring/Makefile create mode 100644 tests/ring/test_null_mode.cu create mode 100644 tests/ring/test_producer.cu create mode 100644 tests/ring/test_ring_engine.cu create mode 100644 tests/ring/test_rings.cu create mode 100644 tests/test_capture_catalog_benchmark.py create mode 100644 tests/test_capture_catalog_indexer.py create mode 100644 tests/test_capture_pack_benchmark.py create mode 100644 tests/test_capture_parallel_upload.py create mode 100644 tests/test_capture_pipeline.py create mode 100644 tests/test_capture_pipeline_benchmark.py create mode 100644 tests/test_capture_s3.py create mode 100644 tests/test_capture_spool.py create mode 100644 tests/test_capture_storage.py create mode 100644 tests/test_clickhouse_capture_catalog.py create mode 100644 tests/test_clickhouse_catalog_live.py create mode 100644 tests/test_garage_live.py create mode 100644 tests/test_garage_upload_benchmark.py create mode 100644 tests/test_internal_mapper.py create mode 100644 tests/test_internal_mapper_e2e.py create mode 100644 tests/tools/run_garage_live.py diff --git a/benchmarks/bench_capture_catalog.py b/benchmarks/bench_capture_catalog.py new file mode 100644 index 000000000..145dc1ddd --- /dev/null +++ b/benchmarks/bench_capture_catalog.py @@ -0,0 +1,116 @@ +"""Measure batched capture-metadata inserts into ClickHouse.""" + +from __future__ import annotations + +import argparse +from dataclasses import replace +import json +from statistics import median +from time import perf_counter_ns, time_ns +from uuid import UUID, uuid4 + +from dmi.storage.capture import ( + CaptureDescriptor, + CaptureMetadata, + ClickHouseCatalogConfig, + ClickHouseCatalogWriter, + PayloadLocator, +) + + +def synthetic_descriptors(rows: int) -> tuple[CaptureDescriptor, ...]: + if rows <= 0: + raise ValueError("rows must be positive") + pack_id = str(UUID("018f0000-0000-7000-8000-000000000001")) + base = CaptureMetadata( + capture_id="capture-0", tenant_id="tenant-a", experiment_id="exp-a", + run_id="run-a", session_id="session-a", request_id="request-a", + sequence_id="sequence-a", model_id="model-a", model_revision="revision-a", + adapter_revision=None, capture_policy_version="policy-v1", + hook_name="resid_pre", layer_number=3, producer_rank=0, step_number=0, + token_start=0, token_end=1, batch_position=0, dtype="float32", + shape=(4096,), captured_at_ns=1_700_000_000_000_000_000, + ) + locator = PayloadLocator( + pack_id=pack_id, store_id="garage", object_key="packs/synthetic.dmi-pack", + object_bytes=rows * 16_384, pack_checksum="0" * 64, + pack_record_count=rows, offset=64, stored_length=16_384, + decoded_length=16_384, codec="none", checksum="00000000", + ) + return tuple( + CaptureDescriptor( + replace( + base, capture_id=f"capture-{index}", step_number=index, + token_start=index, token_end=index + 1, + captured_at_ns=base.captured_at_ns + index, + ), + replace(locator, offset=64 + index * 16_384), + ) + for index in range(rows) + ) + + +def measure_inserts(writer, descriptors, *, batch_rows: int, trials: int) -> dict: + if batch_rows <= 0 or trials <= 0: + raise ValueError("batch_rows and trials must be positive") + samples = [] + inserts = (len(descriptors) + batch_rows - 1) // batch_rows + for trial in range(trials): + start = perf_counter_ns() + for offset in range(0, len(descriptors), batch_rows): + writer.write_descriptors( + descriptors[offset : offset + batch_rows], + index_version=time_ns() + trial, + ) + elapsed = (perf_counter_ns() - start) / 1e9 + samples.append(len(descriptors) / elapsed) + return { + "rows": len(descriptors), + "batch_rows": batch_rows, + "inserts_per_trial": inserts, + "trials": trials, + "rows_per_second_median": median(samples), + "rows_per_second_samples": samples, + } + + +def main(argv=None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=9000) + parser.add_argument("--database", default="default") + parser.add_argument("--rows", type=int, default=100_000) + parser.add_argument("--batch-rows", type=int, default=10_000) + parser.add_argument("--trials", type=int, default=3) + args = parser.parse_args(argv) + + from clickhouse_driver import Client + + prefix = f"dmi_catalog_bench_{uuid4().hex}" + client = Client(args.host, port=args.port) + writer = ClickHouseCatalogWriter( + client, + ClickHouseCatalogConfig(database=args.database, table_prefix=prefix), + ) + writer.ensure_schema() + try: + result = measure_inserts( + writer, + synthetic_descriptors(args.rows), + batch_rows=args.batch_rows, + trials=args.trials, + ) + print(json.dumps(result, sort_keys=True)) + finally: + for kind, suffix in ( + ("VIEW", "capture"), ("VIEW", "pack_inventory"), + ("TABLE", "capture_raw"), ("TABLE", "pack_inventory_raw"), + ): + client.execute( + f"DROP {kind} IF EXISTS `{args.database}`.`{prefix}_{suffix}`" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/bench_capture_pack.py b/benchmarks/bench_capture_pack.py new file mode 100644 index 000000000..50f737fa5 --- /dev/null +++ b/benchmarks/bench_capture_pack.py @@ -0,0 +1,241 @@ +"""CPU-only benchmark for the immutable capture-pack writer.""" + +from __future__ import annotations + +import argparse +from dataclasses import asdict, dataclass +import json +import math +from pathlib import Path +import random +import statistics +import time +from typing import Sequence +from uuid import UUID + +from dmi.storage.capture import CaptureMetadata, CaptureRecord, PackReader, PackWriter + + +_BYTE_UNITS = { + "": 1, + "b": 1, + "kb": 10**3, + "mb": 10**6, + "gb": 10**9, + "kib": 1024, + "mib": 1024**2, + "gib": 1024**3, +} +_DTYPE_BYTES = { + "float16": 2, + "bfloat16": 2, + "float32": 4, + "float64": 8, + "uint8": 1, + "int8": 1, + "int16": 2, + "int32": 4, + "int64": 8, +} + + +def parse_byte_size(value: str) -> int: + normalized = value.strip().lower() + split = len(normalized) + while split and normalized[split - 1].isalpha(): + split -= 1 + number, unit = normalized[:split].strip(), normalized[split:] + if not number.isdigit() or unit not in _BYTE_UNITS: + raise argparse.ArgumentTypeError(f"invalid byte size: {value!r}") + return int(number) * _BYTE_UNITS[unit] + + +@dataclass(frozen=True, slots=True) +class PackBenchmarkConfig: + records: int = 10_000 + payload_bytes: int = 64 * 1024 + target_pack_bytes: int = 128 * 1024**2 + pool_size: int = 64 + pattern: str = "random" + dtype: str = "float32" + seed: int = 17 + trials: int = 5 + + def __post_init__(self) -> None: + for name in ("records", "payload_bytes", "target_pack_bytes", "pool_size", "trials"): + if getattr(self, name) <= 0: + raise ValueError(f"{name} must be positive") + if self.dtype not in _DTYPE_BYTES: + raise ValueError(f"unsupported dtype: {self.dtype}") + if self.payload_bytes % _DTYPE_BYTES[self.dtype]: + raise ValueError( + f"payload_bytes must be a multiple of {_DTYPE_BYTES[self.dtype]}" + ) + if self.target_pack_bytes < self.payload_bytes: + raise ValueError("target_pack_bytes must be >= payload_bytes") + if self.pattern not in {"zeros", "random"}: + raise ValueError("pattern must be zeros or random") + + +@dataclass(frozen=True, slots=True) +class PackTrial: + record_count: int + logical_bytes: int + packed_bytes: int + largest_pack_bytes: int + pack_count: int + seconds: float + + def as_dict(self) -> dict[str, float | int]: + return { + **asdict(self), + "logical_gib_per_second": self.logical_bytes / self.seconds / 1024**3, + "packed_gib_per_second": self.packed_bytes / self.seconds / 1024**3, + "space_amplification": self.packed_bytes / self.logical_bytes, + } + + +def generate_payload_pool(config: PackBenchmarkConfig) -> tuple[bytes, ...]: + count = min(config.records, config.pool_size) + if config.pattern == "zeros": + return tuple(bytes(config.payload_bytes) for _ in range(count)) + generator = random.Random(config.seed) + return tuple(generator.randbytes(config.payload_bytes) for _ in range(count)) + + +def _metadata(config: PackBenchmarkConfig, index: int) -> CaptureMetadata: + elements = config.payload_bytes // _DTYPE_BYTES[config.dtype] + return CaptureMetadata( + capture_id=f"capture-{index:012d}", + tenant_id="benchmark", + experiment_id="pack-writer", + run_id=f"seed-{config.seed}", + session_id="session-0", + request_id=f"request-{index // 128}", + sequence_id=f"sequence-{index // 128}", + model_id="synthetic", + model_revision="benchmark-v1", + adapter_revision=None, + capture_policy_version="all-v1", + hook_name="resid_pre", + layer_number=index % 32, + producer_rank=0, + step_number=index, + token_start=index, + token_end=index + 1, + batch_position=index % 128, + dtype=config.dtype, + shape=(elements,), + captured_at_ns=1_700_000_000_000_000_000 + index, + ) + + +def run_trial(config: PackBenchmarkConfig) -> PackTrial: + payloads = generate_payload_pool(config) + packs = [] + writer: PackWriter | None = None + pack_index = 0 + start = time.perf_counter() + for index in range(config.records): + record = CaptureRecord( + metadata=_metadata(config, index), + payload=payloads[index % len(payloads)], + ) + if writer is None: + writer = PackWriter( + pack_id=UUID(int=pack_index + 1), + created_at_ns=1_700_000_000_000_000_000 + pack_index, + max_pack_bytes=config.target_pack_bytes, + max_records=min(config.records, 10_000), + ) + try: + writer.append(record) + except ValueError as exc: + if writer.record_count == 0: + raise ValueError( + "target_pack_bytes cannot hold one benchmark record" + ) from exc + packs.append(writer.seal()) + pack_index += 1 + writer = PackWriter( + pack_id=UUID(int=pack_index + 1), + created_at_ns=1_700_000_000_000_000_000 + pack_index, + max_pack_bytes=config.target_pack_bytes, + max_records=min(config.records, 10_000), + ) + writer.append(record) + if writer is not None: + packs.append(writer.seal()) + elapsed = time.perf_counter() - start + + verified_records = 0 + for pack in packs: + reader = PackReader.from_bytes(pack.data) + descriptors = reader.descriptors(store_id="benchmark", object_key=pack.pack_id) + for descriptor in descriptors: + reader.read_payload(descriptor) + verified_records += len(descriptors) + if verified_records != config.records: + raise RuntimeError(f"verified {verified_records} records, expected {config.records}") + + return PackTrial( + record_count=config.records, + logical_bytes=config.records * config.payload_bytes, + packed_bytes=sum(len(pack.data) for pack in packs), + largest_pack_bytes=max(len(pack.data) for pack in packs), + pack_count=len(packs), + seconds=max(elapsed, math.ulp(1.0)), + ) + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--records", type=int, default=10_000) + parser.add_argument("--payload-bytes", type=parse_byte_size, default=64 * 1024) + parser.add_argument("--target-pack-bytes", type=parse_byte_size, default=128 * 1024**2) + parser.add_argument("--pool-size", type=int, default=64) + parser.add_argument("--pattern", choices=("zeros", "random"), default="random") + parser.add_argument("--dtype", choices=tuple(_DTYPE_BYTES), default="float32") + parser.add_argument("--seed", type=int, default=17) + parser.add_argument("--trials", type=int, default=5) + parser.add_argument("--json-output", type=Path) + parser.add_argument("--dry-run", action="store_true") + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + args = _parser().parse_args(argv) + config = PackBenchmarkConfig( + records=args.records, + payload_bytes=args.payload_bytes, + target_pack_bytes=args.target_pack_bytes, + pool_size=args.pool_size, + pattern=args.pattern, + dtype=args.dtype, + seed=args.seed, + trials=args.trials, + ) + if args.dry_run: + result = {"dry_run": True, "config": asdict(config)} + else: + trials = [run_trial(config) for _ in range(config.trials)] + rates = [trial.as_dict()["logical_gib_per_second"] for trial in trials] + result = { + "dry_run": False, + "config": asdict(config), + "trials": [trial.as_dict() for trial in trials], + "summary": { + "median_logical_gib_per_second": statistics.median(rates), + "min_logical_gib_per_second": min(rates), + "max_logical_gib_per_second": max(rates), + }, + } + encoded = json.dumps(result, indent=2, sort_keys=True) + if args.json_output is not None: + args.json_output.write_text(encoded + "\n") + print(encoded) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/bench_capture_pipeline.py b/benchmarks/bench_capture_pipeline.py new file mode 100644 index 000000000..a4681b304 --- /dev/null +++ b/benchmarks/bench_capture_pipeline.py @@ -0,0 +1,260 @@ +"""CPU-only benchmark for bounded capture packing and local persistence.""" + +from __future__ import annotations + +import argparse +from dataclasses import asdict, dataclass +import json +import math +from pathlib import Path +import random +import statistics +import tempfile +import time +from typing import Sequence + +from dmi.storage.capture import ( + AdmissionResult, + CaptureMetadata, + CaptureRecord, + DirectPackSink, + DurablePackSink, + DurablePackSpool, + FilesystemPackStore, + HostCapturePipeline, + OverloadPolicy, + PackReader, + PipelineConfig, +) + +from .bench_capture_pack import parse_byte_size + + +@dataclass(frozen=True, slots=True) +class PipelineBenchmarkConfig: + mode: str = "direct" + records: int = 10_000 + payload_bytes: int = 64 * 1024 + target_pack_bytes: int = 128 * 1024**2 + queue_records: int = 256 + queue_bytes: int = 256 * 64 * 1024 + pool_size: int = 64 + seed: int = 17 + trials: int = 5 + + def __post_init__(self) -> None: + if self.mode not in {"direct", "spool"}: + raise ValueError("mode must be direct or spool") + for name in ( + "records", + "payload_bytes", + "target_pack_bytes", + "queue_records", + "queue_bytes", + "pool_size", + "trials", + ): + if getattr(self, name) <= 0: + raise ValueError(f"{name} must be positive") + if self.payload_bytes % 4: + raise ValueError("payload_bytes must be a multiple of four") + if self.target_pack_bytes < self.payload_bytes + 1024: + raise ValueError("target_pack_bytes is too small for one record") + if self.queue_bytes < self.payload_bytes: + raise ValueError("queue_bytes must hold at least one payload") + + +@dataclass(frozen=True, slots=True) +class PipelineTrial: + mode: str + record_count: int + persisted_records: int + logical_bytes: int + packed_bytes: int + packs_persisted: int + dropped_records: int + queue_peak_records: int + queue_peak_bytes: int + admission_max_ns: int + persist_max_ns: int + seconds: float + + def as_dict(self) -> dict[str, float | int | str]: + return { + **asdict(self), + "logical_gib_per_second": self.logical_bytes / self.seconds / 1024**3, + "packed_gib_per_second": self.packed_bytes / self.seconds / 1024**3, + "space_amplification": self.packed_bytes / self.logical_bytes, + } + + +def _payloads(config: PipelineBenchmarkConfig) -> tuple[bytes, ...]: + generator = random.Random(config.seed) + return tuple( + generator.randbytes(config.payload_bytes) + for _ in range(min(config.records, config.pool_size)) + ) + + +def _record(config: PipelineBenchmarkConfig, index: int, payload: bytes) -> CaptureRecord: + return CaptureRecord( + metadata=CaptureMetadata( + capture_id=f"capture-{index:012d}", + tenant_id="benchmark", + experiment_id="capture-pipeline", + run_id=f"seed-{config.seed}", + session_id="session-0", + request_id=f"request-{index // 128}", + sequence_id=f"sequence-{index // 128}", + model_id="synthetic", + model_revision="benchmark-v1", + adapter_revision=None, + capture_policy_version="all-v1", + hook_name="resid_pre", + layer_number=index % 32, + producer_rank=0, + step_number=index, + token_start=index, + token_end=index + 1, + batch_position=index % 128, + dtype="float32", + shape=(config.payload_bytes // 4,), + captured_at_ns=1_700_000_000_000_000_000 + index, + ), + payload=payload, + ) + + +def _verify(root: Path, mode: str) -> int: + paths = ( + root.rglob("*.dmi-pack") + if mode == "direct" + else root.rglob("*.dmi-pack.ready") + ) + records = 0 + for path in paths: + records += len(PackReader.from_bytes(path.read_bytes()).descriptors( + store_id="verify", object_key=path.name + )) + return records + + +def run_trial(config: PipelineBenchmarkConfig) -> PipelineTrial: + payloads = _payloads(config) + with tempfile.TemporaryDirectory(prefix="dmi-capture-pipeline-") as directory: + root = Path(directory) + if config.mode == "direct": + sink = DirectPackSink( + FilesystemPackStore(root / "objects", store_id="local") + ) + verify_root = root / "objects" + else: + spool_bytes = max( + config.target_pack_bytes * 2, + config.records * (config.payload_bytes + 4096), + ) + sink = DurablePackSink( + DurablePackSpool(root / "spool", max_bytes=spool_bytes) + ) + verify_root = root / "spool" + pipeline = HostCapturePipeline( + PipelineConfig( + max_queue_records=config.queue_records, + max_queue_bytes=config.queue_bytes, + max_pack_bytes=config.target_pack_bytes, + max_pack_records=10_000, + max_linger_ns=1_000_000_000, + overload_policy=OverloadPolicy.BLOCK, + admission_timeout=30, + ), + sink, + ) + + started = time.perf_counter() + pipeline.start() + for index in range(config.records): + result = pipeline.submit( + _record(config, index, payloads[index % len(payloads)]) + ) + if result is not AdmissionResult.ACCEPTED: + raise RuntimeError(f"unexpected admission result: {result.value}") + snapshot = pipeline.close(timeout=30) + elapsed = max(time.perf_counter() - started, math.ulp(1.0)) + + verified = _verify(verify_root, config.mode) + if verified != config.records: + raise RuntimeError(f"verified {verified} records, expected {config.records}") + + return PipelineTrial( + mode=config.mode, + record_count=config.records, + persisted_records=snapshot.persisted_records, + logical_bytes=config.records * config.payload_bytes, + packed_bytes=snapshot.packed_bytes, + packs_persisted=snapshot.packs_persisted, + dropped_records=snapshot.dropped_records, + queue_peak_records=snapshot.queue_peak_records, + queue_peak_bytes=snapshot.queue_peak_bytes, + admission_max_ns=snapshot.admission_duration.max_ns, + persist_max_ns=snapshot.persist_duration.max_ns, + seconds=elapsed, + ) + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--mode", choices=("direct", "spool"), default="direct") + parser.add_argument("--records", type=int, default=10_000) + parser.add_argument("--payload-bytes", type=parse_byte_size, default=64 * 1024) + parser.add_argument( + "--target-pack-bytes", type=parse_byte_size, default=128 * 1024**2 + ) + parser.add_argument("--queue-records", type=int, default=256) + parser.add_argument( + "--queue-bytes", type=parse_byte_size, default=256 * 64 * 1024 + ) + parser.add_argument("--pool-size", type=int, default=64) + parser.add_argument("--seed", type=int, default=17) + parser.add_argument("--trials", type=int, default=5) + parser.add_argument("--json-output", type=Path) + parser.add_argument("--dry-run", action="store_true") + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + args = _parser().parse_args(argv) + config = PipelineBenchmarkConfig( + mode=args.mode, + records=args.records, + payload_bytes=args.payload_bytes, + target_pack_bytes=args.target_pack_bytes, + queue_records=args.queue_records, + queue_bytes=args.queue_bytes, + pool_size=args.pool_size, + seed=args.seed, + trials=args.trials, + ) + if args.dry_run: + result = {"dry_run": True, "config": asdict(config)} + else: + trials = [run_trial(config) for _ in range(config.trials)] + rates = [trial.as_dict()["logical_gib_per_second"] for trial in trials] + result = { + "dry_run": False, + "config": asdict(config), + "trials": [trial.as_dict() for trial in trials], + "summary": { + "median_logical_gib_per_second": statistics.median(rates), + "min_logical_gib_per_second": min(rates), + "max_logical_gib_per_second": max(rates), + }, + } + encoded = json.dumps(result, indent=2, sort_keys=True) + if args.json_output is not None: + args.json_output.write_text(encoded + "\n") + print(encoded) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/bench_clickhouse_host.py b/benchmarks/bench_clickhouse_host.py index a5bbdae8f..e8f40a28a 100644 --- a/benchmarks/bench_clickhouse_host.py +++ b/benchmarks/bench_clickhouse_host.py @@ -270,13 +270,17 @@ def __init__( self._client_factory = client_factory self._interval_seconds = interval_ms / 1000.0 if database is not None and table is not None: - qualified = f"{database}.{table}" + qualified = f"{quote_identifier(database)}.{quote_identifier(table)}" self._process_query = ( "SELECT count() FROM system.processes" " WHERE query_kind = 'Insert'" - " AND has(tables, %(table)s)" + " AND current_database = %(database)s" + " AND position(query, %(qualified_table)s) > 0" ) - self._process_query_params: dict[str, Any] = {"table": qualified} + self._process_query_params: dict[str, Any] = { + "database": database, + "qualified_table": qualified, + } else: self._process_query = ( "SELECT count() FROM system.processes WHERE query_kind = 'Insert'" diff --git a/benchmarks/bench_garage_upload.py b/benchmarks/bench_garage_upload.py new file mode 100644 index 000000000..6af5aa95c --- /dev/null +++ b/benchmarks/bench_garage_upload.py @@ -0,0 +1,287 @@ +"""CPU-only Garage upload scaling benchmark for staged DMI packs.""" + +from __future__ import annotations + +import argparse +from dataclasses import asdict, dataclass, replace +import json +import math +import os +from pathlib import Path +import random +import statistics +import tempfile +import time +from typing import Sequence +from uuid import uuid4 + +from dmi.storage.capture import ( + CaptureMetadata, + CaptureRecord, + DurablePackSpool, + PackWriter, + ParallelSpoolUploader, + ParallelUploadConfig, + S3PackStore, + S3StoreConfig, +) + +from .bench_capture_pack import parse_byte_size + + +@dataclass(frozen=True, slots=True) +class GarageBenchmarkConfig: + pack_payload_bytes: tuple[int, ...] = (64 * 1024**2, 128 * 1024**2) + multipart_threshold_bytes: tuple[int, ...] = (32 * 1024**2, 64 * 1024**2) + upload_workers: tuple[int, ...] = (1, 2, 4, 8) + packs_per_trial: int = 8 + multipart_chunk_bytes: int = 16 * 1024**2 + multipart_concurrency: int = 4 + trials: int = 3 + seed: int = 17 + + def __post_init__(self) -> None: + for name in ( + "pack_payload_bytes", + "multipart_threshold_bytes", + "upload_workers", + ): + values = getattr(self, name) + if not values or any(type(value) is not int or value <= 0 for value in values): + raise ValueError(f"{name} must contain positive integers") + for name in ( + "packs_per_trial", + "multipart_chunk_bytes", + "multipart_concurrency", + "trials", + ): + if type(getattr(self, name)) is not int or getattr(self, name) <= 0: + raise ValueError(f"{name} must be positive") + if self.multipart_chunk_bytes < 5 * 1024**2: + raise ValueError("multipart_chunk_bytes must be at least 5 MiB") + + +@dataclass(frozen=True, slots=True) +class GarageUploadTrial: + pack_payload_bytes: int + multipart_threshold_bytes: int + upload_workers: int + packs: int + object_bytes: int + uploaded_bytes: int + retries: int + peak_active_uploads: int + peak_in_flight_bytes: int + seconds: float + + def as_dict(self) -> dict[str, int | float]: + return { + **asdict(self), + "gib_per_second": self.uploaded_bytes / self.seconds / 1024**3, + } + + +def _pack(payload: bytes, index: int): + identity = uuid4() + metadata = CaptureMetadata( + capture_id=f"garage-benchmark-{identity}", + tenant_id="benchmark", + experiment_id="garage-upload", + run_id=str(identity), + session_id="session-0", + request_id=f"request-{index}", + sequence_id=f"sequence-{index}", + model_id="synthetic", + model_revision="benchmark-v1", + adapter_revision=None, + capture_policy_version="all-v1", + hook_name="resid_pre", + layer_number=index, + producer_rank=0, + step_number=index, + token_start=index, + token_end=index + 1, + batch_position=0, + dtype="uint8", + shape=(len(payload),), + captured_at_ns=time.time_ns(), + ) + writer = PackWriter( + pack_id=identity, + created_at_ns=metadata.captured_at_ns, + max_pack_bytes=len(payload) + 1024 * 1024, + ) + writer.append(CaptureRecord(metadata=metadata, payload=payload)) + return writer.seal() + + +def run_trial( + config: GarageBenchmarkConfig, + store_config: S3StoreConfig, + *, + pack_payload_bytes: int, + multipart_threshold_bytes: int, + upload_workers: int, +) -> GarageUploadTrial: + generator = random.Random(config.seed) + payload = generator.randbytes(pack_payload_bytes) + run_id = uuid4() + with tempfile.TemporaryDirectory(prefix="dmi-garage-benchmark-") as directory: + root = Path(directory) + spool = DurablePackSpool( + root / "spool", + max_bytes=config.packs_per_trial * (pack_payload_bytes + 1024**2), + ) + staged = [] + for index in range(config.packs_per_trial): + pack = _pack(payload, index) + key = f"benchmarks/dmi/{run_id}/{pack.pack_id}.dmi-pack" + staged.append(spool.stage(pack, key)) + + resolved_store = replace( + store_config, + multipart_threshold_bytes=multipart_threshold_bytes, + multipart_chunk_bytes=config.multipart_chunk_bytes, + multipart_concurrency=config.multipart_concurrency, + ) + store = S3PackStore.from_config(resolved_store) + byte_budget = max(item.object_bytes for item in staged) * upload_workers + uploader = ParallelSpoolUploader( + spool, + store, + ParallelUploadConfig( + max_workers=upload_workers, + max_in_flight_bytes=byte_budget, + ), + ) + + started = time.perf_counter() + result = uploader.upload_pending() + elapsed = max(time.perf_counter() - started, math.ulp(1.0)) + if result.failures: + raise RuntimeError(f"Garage upload failures: {result.failures}") + for ref in result.refs: + store.stat(ref) + store.read_range( + ref, max(0, ref.object_bytes - 32), min(32, ref.object_bytes) + ) + + return GarageUploadTrial( + pack_payload_bytes=pack_payload_bytes, + multipart_threshold_bytes=multipart_threshold_bytes, + upload_workers=upload_workers, + packs=config.packs_per_trial, + object_bytes=sum(item.object_bytes for item in staged), + uploaded_bytes=result.snapshot.uploaded_bytes, + retries=result.snapshot.retries, + peak_active_uploads=result.snapshot.peak_active_uploads, + peak_in_flight_bytes=result.snapshot.peak_in_flight_bytes, + seconds=elapsed, + ) + + +def _byte_list(value: str) -> tuple[int, ...]: + return tuple(parse_byte_size(item.strip()) for item in value.split(",")) + + +def _int_list(value: str) -> tuple[int, ...]: + try: + return tuple(int(item.strip()) for item in value.split(",")) + except ValueError as exc: + raise argparse.ArgumentTypeError("expected comma-separated integers") from exc + + +def _store_config_from_env( + config: GarageBenchmarkConfig, multipart_threshold_bytes: int +) -> S3StoreConfig: + required = { + "endpoint_url": os.environ.get("DMI_S3_ENDPOINT"), + "bucket": os.environ.get("DMI_S3_BUCKET"), + "access_key_id": os.environ.get("DMI_S3_ACCESS_KEY_ID"), + "secret_access_key": os.environ.get("DMI_S3_SECRET_ACCESS_KEY"), + } + missing = [name for name, value in required.items() if not value] + if missing: + raise RuntimeError("missing Garage environment: " + ", ".join(missing)) + return S3StoreConfig( + endpoint_url=required["endpoint_url"], + bucket=required["bucket"], + region=os.environ.get("DMI_S3_REGION", "garage"), + access_key_id=required["access_key_id"], + secret_access_key=required["secret_access_key"], + store_id="garage-benchmark", + allow_insecure_http=os.environ.get("DMI_S3_ALLOW_HTTP") == "1", + multipart_threshold_bytes=multipart_threshold_bytes, + multipart_chunk_bytes=config.multipart_chunk_bytes, + multipart_concurrency=config.multipart_concurrency, + max_pool_connections=max( + 32, max(config.upload_workers) * config.multipart_concurrency + ), + ) + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--pack-payload-bytes", type=_byte_list, default="64MiB,128MiB") + parser.add_argument("--multipart-threshold-bytes", type=_byte_list, default="32MiB,64MiB") + parser.add_argument("--upload-workers", type=_int_list, default="1,2,4,8") + parser.add_argument("--packs-per-trial", type=int, default=8) + parser.add_argument("--multipart-chunk-bytes", type=parse_byte_size, default=16 * 1024**2) + parser.add_argument("--multipart-concurrency", type=int, default=4) + parser.add_argument("--trials", type=int, default=3) + parser.add_argument("--seed", type=int, default=17) + parser.add_argument("--json-output", type=Path) + parser.add_argument("--dry-run", action="store_true") + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + args = _parser().parse_args(argv) + config = GarageBenchmarkConfig( + pack_payload_bytes=args.pack_payload_bytes, + multipart_threshold_bytes=args.multipart_threshold_bytes, + upload_workers=args.upload_workers, + packs_per_trial=args.packs_per_trial, + multipart_chunk_bytes=args.multipart_chunk_bytes, + multipart_concurrency=args.multipart_concurrency, + trials=args.trials, + seed=args.seed, + ) + if args.dry_run: + result = {"dry_run": True, "config": asdict(config)} + else: + trials = [] + for pack_bytes in config.pack_payload_bytes: + for threshold in config.multipart_threshold_bytes: + store_config = _store_config_from_env(config, threshold) + for workers in config.upload_workers: + trials.extend( + run_trial( + config, + store_config, + pack_payload_bytes=pack_bytes, + multipart_threshold_bytes=threshold, + upload_workers=workers, + ) + for _ in range(config.trials) + ) + rates = [trial.as_dict()["gib_per_second"] for trial in trials] + result = { + "dry_run": False, + "config": asdict(config), + "trials": [trial.as_dict() for trial in trials], + "summary": { + "median_gib_per_second": statistics.median(rates), + "min_gib_per_second": min(rates), + "max_gib_per_second": max(rates), + }, + } + encoded = json.dumps(result, indent=2, sort_keys=True) + if args.json_output is not None: + args.json_output.write_text(encoded + "\n") + print(encoded) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/docs/architecture.md b/docs/architecture.md index 6147ace0c..10e25eb54 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -80,6 +80,14 @@ and payload copying; the native producer path remains in no-op/null mode. For active transport without persistence, use `dmx_null_mode=False` and leave the database host empty. +The opt-in next-generation host sink makes immutable, self-describing tensor +packs in S3-compatible object storage the canonical capture log. An independent +indexer now rebuilds a logically deduplicated ClickHouse metadata projection by +reading only pack trailers and footers. Scalar summaries remain a later phase. +This work is host-side, opt-in, and does not change Ring². See the +[`capture storage design`](capture-storage-design.md) and its +[`HTML explainer`](capture-storage-pipeline.html). + The drain pipeline is independent of the inference loop: backpressure on the sink does not block the GPU producer as long as the rings are sized for the workload. Sizing knobs (`dmx_ring_payload_mb`, `dmx_ring_pinned_mb`) are diff --git a/docs/benchmarks.md b/docs/benchmarks.md index 801b3fd4d..f48e44030 100644 --- a/docs/benchmarks.md +++ b/docs/benchmarks.md @@ -99,6 +99,25 @@ documentation when interpreting results. See the [native build layout](native-build-layout.html) for the host/full extension boundary and loader behavior. +## Derived catalog throughput + +`benchmarks.bench_capture_catalog` measures the opt-in metadata projection; it +does not exercise CUDA or write tensor payloads to ClickHouse. It creates +temporary raw tables plus logically deduplicated views, inserts deterministic +capture descriptors in bounded batches, reports every trial, and drops the +objects afterward. + +```bash +python -m benchmarks.bench_capture_catalog \ + --rows 100000 --batch-rows 10000 --trials 3 +``` + +On local ClickHouse 26.9.1, median throughput rose from 13,954 rows/s with +1,000-row batches to 88,458 rows/s at 10,000 and 157,567 rows/s at 50,000. +These loopback results validate client batching, not a production capacity +claim. Repeat the sweep on the target server while observing part creation, +merge load, CPU, and catalog lag. + Baselines: - **HuggingFace Ideal** — vanilla HF `generate`, no observation (used as 1.0) diff --git a/docs/capture-storage-design.md b/docs/capture-storage-design.md new file mode 100644 index 000000000..372be79e3 --- /dev/null +++ b/docs/capture-storage-design.md @@ -0,0 +1,741 @@ +# Capture storage design + +Status: Accepted; host storage reference through Phase 3 implemented + +This document defines a clean-slate host persistence architecture after Ring². +It does not change the CUDA producer, ring layout, or device-to-host transport. +The new path remains opt-in until CPU-only, live-store, and compatibility gates +pass. + +The visual companion is +[`capture-storage-pipeline.html`](capture-storage-pipeline.html). + +## Implementation status + +The first host-only slice is available under `dmi.storage.capture`: + +| Capability | Status | +|---|---| +| Versioned pack writer and full validator | Implemented | +| Two-range footer index | Implemented | +| Immutable filesystem store | Implemented | +| Stable selection, byte estimation, and range hydration | Implemented | +| CPU pack benchmark and package checks | Implemented | +| Bounded asynchronous pack pipeline and spool recovery | Implemented | +| Garage/S3 store and bounded parallel uploader | Implemented | +| ClickHouse metadata projection | Implemented and live-tested | +| Summaries | Planned | + +This slice is opt-in and has no connection to the CUDA producer or current +ClickHouse payload sink. + +The Phase 2 implementation adds `HostCapturePipeline`, bounded blocking or +drop-newest admission, size/record/linger/session/shutdown sealing, direct local +persistence, and a durable filesystem spool. The durable sink commits locally; +`SpoolUploader` independently retries staged packs and removes them only after +remote size and SHA-256 verification. One process owns a spool directory. + +Phase 3 adds `S3PackStore` and `ParallelSpoolUploader`. The store streams packs +through Boto3 managed multipart transfers, persists DMI SHA-256 and pack +identity as object metadata, supports exact byte ranges and bounded paginated +listing, and resolves ambiguous retries by checking the existing key. The +uploader bounds outer workers and aggregate bytes in flight, applies bounded +backoff only to transient failures, and leaves permanent failures in the spool. + +Phase 4 adds bounded notification and prefix-scan discovery, footer-only pack +indexing, client-side ClickHouse batching, and pack commit markers. Descriptor +rows are inserted before their pack marker. An interrupted or ambiguous batch +may replay physical rows, while `ReplacingMergeTree` tables and public `FINAL` +views preserve immediate logical results. Object storage remains sufficient to +rebuild the projection. + +## Decision + +Treat immutable, self-describing tensor packs in object storage as the only +durable source of truth. A successful object upload is the capture commit. +ClickHouse is a rebuildable analytical projection populated by an independent +catalog indexer. + +```text +Ring² host drain + | + v +bounded slabs -> pack assembler -> direct upload or NVMe spool + | + v + canonical object-store packs + | + +-----------------+-----------------+ + | | + v v + catalog indexer summary workers + | | + v v + ClickHouse catalog scalar rows + object artifacts + | + v + metadata-first query -> estimate -> selective range hydration +``` + +The capture host does not run a ClickHouse client, compute summaries, or +coordinate two durable writes. Object-created notifications reduce indexing +latency, while periodic listing and reconciliation provide completeness. + +## Goals + +- Keep persistence and analytics off the inference path. +- Scale payload upload, catalog indexing, and summarization independently. +- Avoid one object-store request or ClickHouse insert per tensor. +- Bound host memory, disk, concurrency, retries, and read amplification. +- Make the catalog reconstructable from canonical packs. +- Let users and agents inspect summaries before transferring tensor bytes. +- Add storage providers and summarizers without changing Ring². +- Prove the design on CPU-only hosts before changing CUDA code. + +## Non-goals + +- Changing hook placement, the CUDA producer, or Ring². +- Querying arbitrary tensor contents directly in ClickHouse. +- Providing a cross-system exactly-once transaction. +- Using an embedded database as a second source of truth. +- Selecting production pack sizes or codecs without measurement. + +## Performance model + +The hot path is deliberately short: + +```text +drain -> acquire reusable slab -> append record -> seal pack -> enqueue upload +``` + +The architecture removes four scaling costs from capture hosts: + +- per-tensor network requests; +- ClickHouse insert latency and merge behavior; +- summary computation; +- coordination between payload and catalog durability. + +Packs amortize object-store request and protocol overhead. An independent +indexer reads only pack footers and batches rows across all producers before +inserting into ClickHouse. Readers coalesce adjacent selected ranges instead of +downloading complete packs. + +Initial tuning ranges are hypotheses: + +| Setting | Initial sweep | Purpose | +|---|---:|---| +| Target pack size | 64–256 MiB | Amortize upload and object overhead | +| Maximum linger | 50–100 ms | Bound visibility latency at low volume | +| Upload workers | 1–16 | Find the store/network saturation point | +| Index batch rows | 10k–100k | Avoid small ClickHouse inserts | +| Index batch bytes | 16–64 MiB | Bound memory while preserving batch efficiency | + +The pack-and-upload plane must sustain at least 1.2 times the expected +device-to-host rate on representative hardware. A value becomes a default only +after repeated measurements exceed run-to-run variance and correctness tests +remain green. + +## Host capture agent + +### Admission + +The drain thread hands each `CaptureRecord` to a queue bounded by both bytes and +record count. A record owns or references a contiguous CPU payload and immutable +metadata. Queue saturation follows an explicit policy: bounded blocking, +sampling, or dropping. The system never silently grows memory or disk. + +### Pack assembly + +Pack assemblers use reusable slabs and partition work by stable producer scope +so independent capture streams do not share a global lock. They seal on target +size, linger deadline, session boundary, or shutdown. + +Each record is independently encoded. Whole-pack streaming compression is not +used because it would require preceding bytes to decode one selected tensor. +The first implementation supports `none` and one measured block codec. + +### Direct and durable modes + +Both modes implement the same `PackSink` contract. Direct mode writes a +`PackSource` to a `PackStore`; durable mode first writes that source to the +local spool and exposes the staged file as another streaming `PackSource`. + +Direct mode uploads sealed packs from bounded memory. Durable mode writes sealed +packs to local NVMe and uses atomic rename: + +```text +.open -> .ready -> upload -> remote verification -> delete +``` + +At restart, valid `.ready` packs are retried with the same deterministic key. +Ambiguous uploads use remote metadata and checksum verification. The filesystem +state machine is sufficient initially; RocksDB or SQLite is added only if +measurements demonstrate a recovery or scheduling bottleneck. + +The spool absorbs bursts and process failure. It cannot compensate for +sustained object-store throughput below the capture rate. + +The CPU reference exposes the modes through a common `PackSink` boundary: + +```python +config = PipelineConfig( + max_queue_records=256, + max_queue_bytes=16 * 1024**2, + max_pack_bytes=128 * 1024**2, + max_pack_records=10_000, + max_linger_ns=100_000_000, +) + +sink = DirectPackSink(store) +# Or: sink = DurablePackSink(DurablePackSpool(path, max_bytes=...)) + +pipeline = HostCapturePipeline(config, sink) +pipeline.start() +result = pipeline.submit(record) +snapshot = pipeline.close(timeout=30) +``` + +Admission returns an explicit `AdmissionResult`. It never silently expands the +queue. Durable mode intentionally separates local commit from remote upload so +a remote outage cannot invalidate a completed local commit. + +## Pack format + +`dmi-pack-v1` is one immutable object: + +```text ++------------------------+ +| fixed header | magic, format version, pack ID ++------------------------+ +| independently encoded | tensor record 0 ++------------------------+ +| independently encoded | tensor record 1 ++------------------------+ +| ... | ++------------------------+ +| footer manifest | IDs, provenance, shape, offsets, codecs, checksums ++------------------------+ +| fixed trailer | footer offset/length, pack checksum ++------------------------+ +``` + +The fixed trailer allows an indexer to locate the footer with one small suffix +range read. A second range read retrieves the footer. Tensor bytes are not read +during catalog indexing. + +The footer is authoritative for reconstruction. It includes stable capture +identity, format identity, object-relative ranges, and sufficient metadata to +recreate catalog rows. Readers reject unknown major versions and allow only +documented additive minor changes. + +Object keys are immutable and deterministic for one persistence intent: + +```text +v1/tenant=/date=/session=/rank=/.dmi-pack +``` + +Key components are percent-encoded. Values that would exceed portable +filesystem component limits use a stable SHA-256 component; the full value +remains in the pack footer and catalog. + +The same retry reuses the key and checksum. A different pack never reuses that +key. + +## Commit and discovery semantics + +Successful completion of the pack upload is the commit. There is no separate +manifest object or catalog acknowledgement. + +Discovery combines: + +1. Object-created notifications for low latency. +2. Periodic prefix scans for correctness. +3. Stable pack IDs for idempotent replay. + +Notifications may be delayed, duplicated, or reordered. The indexer therefore +treats them as hints. Reconciliation scans recent time partitions and scheduled +older partitions, compares pack identities with indexed state, and replays any +missing work. + +## Catalog indexer + +The indexer is isolated from capture hosts and scales independently. For each +discovered pack it: + +1. Reads the fixed trailer. +2. Reads and validates the footer. +3. Converts descriptors into catalog rows. +4. Batches rows across many packs. +5. Inserts into ClickHouse using stable IDs and versions. +6. Reports discovery lag, validation failures, and batch health. + +At-least-once discovery can produce physical duplicates. Private raw tables may +use `ReplacingMergeTree`, but public views must provide deterministic logical +deduplication. Correctness does not depend on asynchronous merges having run. + +ClickHouse stores: + +- pack inventory and integrity state; +- capture descriptors and payload ranges; +- stable core summaries; +- extensible scalar metrics; +- references to large summary artifacts; +- indexing and enrichment health. + +ClickHouse does not store raw tensor payloads or the only copy of essential +capture metadata. + +## Summary model + +Summary computation is asynchronous and versioned by: + +```text +(capture_id, summarizer_name, summarizer_version, config_hash) +``` + +Stable, commonly filtered values such as minimum, maximum, mean, standard +deviation, norms, sparsity, NaN count, and infinity count use typed ClickHouse +columns. Extensible scalar metrics use a bounded long-form table. Large arrays, +histograms, embeddings, sketches, and sampled tensors remain immutable objects +with catalog locators. + +New summarizers do not change the capture host or pack format. + +## Reader and agent workflow + +The public API is metadata-first and bounded: + +```text +search(filters, page) +query_metrics(filters, metric_names, page) +estimate_hydration(capture_ids) +hydrate(capture_ids, byte_limit, request_limit) +get_artifacts(capture_ids) +export(capture_ids, format) +``` + +A typical flow is: + +1. Query ClickHouse with bounded filters and pagination. +2. Inspect core or custom summaries. +3. Select capture IDs. +4. Estimate payload bytes and request count. +5. Coalesce adjacent ranges within each pack. +6. Fetch with byte, request, and concurrency limits. +7. Verify checksums and decode selected records. + +List and summary operations never hydrate tensor bytes implicitly. Storage +credentials and provider endpoints are resolved from deployment configuration, +not returned in catalog rows. + +## Extension contracts + +```python +class PackWriter(Protocol): + def append(self, record: CaptureRecord) -> None: ... + def seal(self) -> SealedPack: ... + + +class PackSource(Protocol): + pack_id: str + object_bytes: int + checksum: str + def open(self) -> BinaryIO: ... + + +class PackStore(Protocol): + def put(self, pack: PackSource, object_key: str) -> PackRef: ... + def stat(self, ref: PackRef) -> ObjectInfo: ... + def read_range(self, ref: PackRef, offset: int, length: int) -> bytes: ... + def list_committed(self, cursor: ScanCursor) -> Page[PackRef]: ... + + +class PackSink(Protocol): + def persist(self, ready: ReadyPack) -> object: ... + + +class CommitFeed(Protocol): + def watch(self, cursor: EventCursor) -> Iterable[PackRef]: ... + def scan(self, cursor: ScanCursor) -> Page[PackRef]: ... + + +class CatalogIndexer(Protocol): + def index(self, packs: Sequence[PackRef]) -> IndexResult: ... + + +class CaptureSummarizer(Protocol): + name: str + version: str + def summarize(self, capture: CaptureDescriptor, tensor: TensorView) -> SummaryBatch: ... + + +class CaptureReader(Protocol): + def search(self, query: CaptureQuery) -> Page[CaptureDescriptor]: ... + def estimate(self, ids: Sequence[CaptureId]) -> HydrationEstimate: ... + def hydrate(self, request: HydrationRequest) -> Iterable[TensorRecord]: ... +``` + +Initial implementations: + +```text +PackStore + - FilesystemPackStore + - S3PackStore + +CatalogIndexer + - ClickHouseCatalogIndexer + +CaptureSummarizer + - CoreTensorStatsSummarizer +``` + +External configuration and provider responses are validated at their +boundaries. Internal stages exchange typed records without repeated validation. + +`S3PackStore` is the Garage implementation; Garage is selected through its +endpoint, region, and credentials rather than a provider-specific subclass. +This avoids duplicating the S3 contract while keeping provider compatibility in +the live test matrix. + +## Failure semantics + +| Failure | Visible result | Recovery | +|---|---|---| +| Upload fails | Pack is not committed | Retry the deterministic key with bounded backoff | +| Upload succeeds, event is lost | Pack exists but catalog is late | Prefix reconciliation discovers it | +| Event is duplicated | Same pack may be indexed again | Stable IDs and public deduplication preserve logical results | +| ClickHouse is unavailable | Capture continues; catalog lag grows | Indexer retries and replays canonical packs | +| Footer is corrupt | Pack is not indexed | Quarantine identity and emit an integrity failure | +| Reader gets a short range | Hydration fails closed | Retry, then report pack and range | +| Checksum mismatches | Tensor is not decoded | Quarantine pack and emit an integrity failure | +| Host queue or spool fills | Configured overload policy applies | Report pressure and drops; never grow without bound | +| Process exits with `.ready` packs | Packs remain locally recoverable | Restart uploader and verify ambiguous remote writes | + +Deleting a pack requires a separate retention workflow. Garbage collection and +pack compaction never run in capture or indexing critical paths. + +## Observability and acceptance gates + +| Plane | Required measurements | +|---|---| +| Admission | enqueue latency, blocked time, drops, queue bytes | +| Packing | GiB/s, CPU, copies, compression ratio, flush reason | +| Upload | GiB/s, p50/p95 latency, active requests, retries, spool growth | +| Indexing | catalog lag, rows/s, rows/insert, bytes/insert, duplicate rate | +| Hydration | range latency, useful GiB/s, read amplification, checksum failures | +| Process | CPU, peak RSS, network, local-spool I/O | +| End-to-end | capture-to-committed and capture-to-queryable p50/p95/p99 | + +Correctness gates every performance result. A faster path that loses required +records, skips validation, or changes reader semantics is a regression. + +## Package layout + +```text +src/dmi/storage/capture/ + model.py # contracts and bounded request types + pack.py # dmi-pack-v1 writer, validator, footer index + filesystem.py # immutable local reference store + reader.py # selection, estimation, coalesced hydration + pipeline.py # bounded admission, assembly, sinks, metrics + spool.py # atomic local commit, recovery, retry upload + s3.py # Garage/S3 streaming, listing, exact range reads + catalog.py # discovery, reconciliation, footer indexing + clickhouse_catalog.py # raw tables, logical views, batched inserts + +# Planned additions +src/dmi/storage/ + summaries/ + contracts.py + core.py + +native/csrc/storage/ + capture_record.h + pack_builder.h + persistence_pipeline.h + spool.h +``` + +## Implementation plan + +### Phase 1 — contract and format + +- Specify `dmi-pack-v1`, stable identifiers, checksums, and compatibility. +- Implement writer/reader round trips and filesystem storage. +- Test truncation, corruption, unknown versions, and deterministic retries. + +Exit gate: CPU-only format tests pass and pack throughput is reproducible. + +Status: CPU reference exit gate passed. The local five-trial, 2,000-record +baseline used 64 KiB payloads and one roughly 126 MiB pack per trial. It measured +0.521 GiB/s median construction throughput and 1.009× space amplification. This +is a regression baseline for the Python reference, not evidence that the future +native pipeline meets the 1.2× production-capacity gate. + +Run it from a source checkout with: + +```bash +PYTHONPATH=src python -m benchmarks.bench_capture_pack \ + --records 2000 --payload-bytes 64KiB \ + --target-pack-bytes 128MiB --trials 5 +``` + +### Phase 2 — bounded CPU pipeline + +- Implement reusable slabs and size/linger sealing. +- Add direct upload and durable-spool policies behind an opt-in mode. +- Add saturation, restart, and ambiguous-upload tests. + +Exit gate: memory and disk remain bounded under sustained overload. + +Status: CPU reference exit gate passed. Queue tests sustain 100 submissions +against a three-record/24-byte queue and preserve the exact bound under +drop-newest overload. The measured blocking-admission workload used 2,000 × +64 KiB records, a 256-record/16 MiB queue, and one roughly 126 MiB pack. Across +five local trials it reported: + +| Mode | Median logical throughput | Drops | Peak queue | +|---|---:|---:|---:| +| Direct filesystem | 0.328 GiB/s | 0 | 256 records / 16 MiB | +| Durable spool | 0.345 GiB/s | 0 | 256 records / 16 MiB | + +Both modes produced 1.0095× space amplification. The difference between local +modes is within filesystem and scheduling variance; it is not treated as an +optimization result. Neither number includes Garage or network upload. + +Run the same workload with: + +```bash +PYTHONPATH=src python -m benchmarks.bench_capture_pipeline \ + --mode direct --records 2000 --payload-bytes 64KiB \ + --target-pack-bytes 128MiB --queue-records 256 \ + --queue-bytes 16MiB --trials 5 +``` + +### Phase 3 — Garage integration + +- Implement the S3-compatible store contract. +- Sweep pack size, multipart threshold, and upload concurrency. +- Test retries, restart recovery, listing, and range reads against a pinned + Garage single-node release. + +Exit gate: pack and upload capacity exceeds target input by at least 1.2 times. + +Status: implementation and local compatibility gates passed; the production +capacity gate remains pending because no target device-to-host rate or +representative network/storage hardware has been supplied. The pinned Garage +v2.3.0 live test covers multipart upload, retry idempotency, object metadata, +listing, and the two range reads used to load a pack footer. + +An Apple Silicon local sweep used four packs per trial, a 16 MiB multipart +threshold and chunk, two multipart requests per pack, and three trials per +point. All objects and byte caps verified with zero retries: + +| Pack payload | 1 outer worker | 2 outer workers | 4 outer workers | +|---:|---:|---:|---:| +| 32 MiB | 0.248 GiB/s | 0.388 GiB/s | 0.473 GiB/s | +| 64 MiB | 0.265 GiB/s | 0.397 GiB/s | 0.502 GiB/s | + +These results show useful parallel scaling and saturation beginning near four +outer workers on this loopback, single-node setup. They do not establish a +production default. Outer pack concurrency and per-pack multipart concurrency +multiply; their product and the configured HTTP connection pool bound the +potential active S3 requests. + +Install the optional client and run the isolated live contract with: + +```bash +python -m pip install -e '.[s3]' +DMI_GARAGE_BINARY=/path/to/garage \ + python tests/tools/run_garage_live.py +``` + +The harness pins Garage 2.3.0 by default, creates temporary credentials and +storage, and deletes the entire instance on exit. The official Garage download +page currently publishes Linux binaries; macOS can build the same pinned tag +from source using Garage's documented Cargo workflow. + +Run a sweep against the same ephemeral server with: + +```bash +DMI_GARAGE_BINARY=/path/to/garage \ + python tests/tools/run_garage_live.py --benchmark -- \ + --pack-payload-bytes 32MiB,64MiB \ + --multipart-threshold-bytes 16MiB \ + --upload-workers 1,2,4 --packs-per-trial 4 \ + --multipart-chunk-bytes 16MiB --multipart-concurrency 2 --trials 3 +``` + +### Phase 4 — derived ClickHouse catalog + +- Implement notification and reconciliation discovery. +- Range-read and validate pack footers. +- Batch catalog rows across packs and expose logically deduplicated views. +- Prove full catalog rebuild from object storage. + +Exit gate: forced missed and duplicate events converge to the expected catalog. + +Status: exit gate passed in the CPU contract suite and against ClickHouse +26.9.1. Duplicate notifications are collapsed before footer work; missed +notifications are recovered by bounded prefix scans; corrupt footers never get +a pack commit marker; and an ambiguous marker insert safely replays descriptor +rows. Two physical descriptor and pack rows returned one logical row through +each public `FINAL` view in the live test. + +The local 100,000-row, three-trial batch sweep measured: + +| Rows per insert | Inserts per trial | Median rows/s | +|---:|---:|---:| +| 1,000 | 100 | 13,954 | +| 10,000 | 10 | 88,458 | +| 50,000 | 2 | 157,567 | + +The result supports large client batches and the existing 10k–100k tuning +range, but does not establish a production default. It excludes object-store +discovery time and should be repeated on representative ClickHouse hardware. +ClickHouse recommends client batching and documents `FINAL` as the query-time +correctness mechanism for `ReplacingMergeTree` data; see its +[insert guidance](https://clickhouse.com/docs/concepts/best-practices/selecting-an-insert-strategy) +and [`FINAL` guidance](https://clickhouse.com/resources/engineering/clickhouse-optimize-table-final). + +Run the benchmark with: + +```bash +PYTHONPATH=src python -m benchmarks.bench_capture_catalog \ + --rows 100000 --batch-rows 10000 --trials 3 +``` + +### Phase 5 — reader and summaries + +- Implement bounded search, estimation, and coalesced range hydration. +- Add versioned core summaries, scalar extensions, and object artifacts. +- Add agent-safe query and hydration limits. + +Exit gate: analysis reads only explicitly selected payload ranges. + +### Phase 6 — migration + +- Keep the current ClickHouse payload sink as the default initially. +- Compare golden workloads by identity, logical bytes, checksums, decoded + tensors, and query results. +- Run fault injection and record performance variance. +- Switch the default only after compatibility, recovery, and throughput gates + pass; preserve configuration rollback. + +No phase requires a CUDA-side change. + +## Operating signals + +The Phase 2 metrics answer four initial on-call questions: + +| Question | Signal | +|---|---| +| Is admission saturated? | accepted, dropped, timed-out, oversized, and closed counters; queue peaks; admission histogram | +| Is persistence keeping up? | persisted records/bytes, pack counts, flush reasons, persistence histogram | +| Did the worker fail? | failure counter and typed `pipeline_failed` event | +| Is durable work accumulating? | current/peak spool bytes and pending entry count | + +Event callbacks receive bounded structured fields. They include pack identity for +correlation but never tensor payloads or capture metadata. Callback failures are +counted and cannot fail persistence. Deployment adapters can export these +snapshots and events to OpenTelemetry without coupling the storage core to one +telemetry vendor. + +Phase 3 adds upload attempts, successful packs and bytes, failures, retries, +peak active uploads, peak bytes in flight, duration totals/maxima, and callback +failures. Events are `pack_upload_retry`, `pack_upload_committed`, and +`pack_upload_failed`; they contain pack identity and bounded diagnostic fields, +never tensor contents or credentials. + +Phase 4 returns requested, skipped, indexed, and failed pack counts; indexed +rows; descriptor insert count; estimated metadata bytes; elapsed time; and a +bounded set of failure details. The `catalog_index_completed` event exposes the +same aggregate fields without object keys, capture metadata, or payloads. +Notification size, page size, packs per indexing call, estimated metadata +bytes, and retained failure details all have explicit caps. + +## Phase 2 limitations + +- The implementation is a Python CPU reference, not the final reusable native + slab allocator. +- One process owns a spool directory; cross-process locking is not implemented. +- Durable mode stages synchronously and uploads through a separate explicit + uploader, so remote backpressure is isolated from local commit. +- The pipeline remains opt-in and is not connected to Ring². + +## Phase 3 limitations + +- Garage has no object versioning or object locks. UUID-based object keys have + one designated writer; preflight and post-upload metadata checks detect + retries and conflicts but are not a cross-writer compare-and-swap primitive. +- Boto3 is an optional dependency. Importing `dmi.storage.capture` does not + require it; constructing an S3 client does. +- One uploader owns a spool. Cross-process scheduling and locking remain out of + scope. +- The benchmark measures upload from recovered spool files. It deliberately + excludes pack construction so pack and store saturation can be diagnosed + independently. +- The Phase 3 code remains opt-in and does not change CUDA or the current + ClickHouse payload sink. + +## Phase 4 limitations + +- Reconciliation is an explicit bounded call; deployment scheduling and event + transport remain outside the storage library. +- The metadata byte bound is a conservative serialized-size estimate, not + ClickHouse wire size. +- The current public views use `FINAL` for immediate replay correctness. Hot + query workloads must measure its cost before choosing a different projection. + +## Alternatives considered + +### Host writes payload and catalog directly + +This makes ClickHouse latency, availability, and schema part of the capture +commit path. It also fragments inserts across inference hosts. Rejected in favor +of a centralized, independently scalable indexer. + +### Separate payload and manifest objects + +Standard manifest formats can be convenient, but object storage does not offer +an atomic transaction across two keys. A footer inside one pack supplies the +same reconstruction data with one commit boundary. Parquet manifest export can +remain an optional downstream interoperability feature. + +### Raw payloads in ClickHouse + +Operationally simple, but large binary columns compete with searchable metadata +for inserts, merges, caches, and storage. ClickHouse remains the derived catalog +rather than the raw byte store. + +### One object per tensor + +Simple addressing, but request rate and object count dominate for small +tensors. Packs amortize those costs while retaining record-level range reads. + +### RocksDB or SQLite in the host path + +An embedded index introduces another recovery and compaction surface. Bounded +pack files plus atomic rename are sufficient for the initial spool. Reconsider +only after measurement identifies a concrete need. + +### Object storage without ClickHouse + +Durable and inexpensive, but poorly suited to interactive high-cardinality +discovery and aggregation. ClickHouse supplies the rebuildable hot query layer. + +## References + +- [Amazon S3 data consistency model](https://docs.aws.amazon.com/console/s3/UsingObjects.html) +- [Amazon S3 event notifications](https://docs.aws.amazon.com/AmazonS3/latest/userguide/EventNotifications.html) +- [Amazon S3 event ordering and duplication](https://docs.aws.amazon.com/AmazonS3/latest/userguide/notification-how-to-event-types-and-destinations.html) +- [Amazon S3 performance guidance](https://docs.aws.amazon.com/pdfs/whitepapers/latest/s3-optimizing-performance-best-practices/s3-optimizing-performance-best-practices.pdf) +- [ClickHouse insert strategy](https://clickhouse.com/docs/concepts/best-practices/selecting-an-insert-strategy) +- [ClickHouse ReplacingMergeTree and FINAL](https://clickhouse.com/resources/engineering/clickhouse-optimize-table-final) +- [ClickHouse and Amazon S3](https://clickhouse.com/integrations/amazon_s3) +- [Garage documentation](https://garagehq.deuxfleurs.fr/) +- [Garage 2.3 quick start](https://garagehq.deuxfleurs.fr/documentation/quick-start/) +- [Garage S3 compatibility](https://garagehq.deuxfleurs.fr/documentation/reference-manual/s3-compatibility/) +- [Garage release downloads](https://garagehq.deuxfleurs.fr/download/) +- [Boto3 managed S3 transfers](https://boto3.amazonaws.com/v1/documentation/api/latest/guide/s3.html#file-transfer-configuration) +- [Amazon S3 multipart upload](https://docs.aws.amazon.com/AmazonS3/latest/userguide/mpuoverview.html) diff --git a/docs/capture-storage-pipeline.html b/docs/capture-storage-pipeline.html new file mode 100644 index 000000000..c93d777c1 --- /dev/null +++ b/docs/capture-storage-pipeline.html @@ -0,0 +1,198 @@ + + + + + + + DMI capture storage performance + + + +
+
+

DMI host persistence

+

Move bytes once. Analyze them without moving them again.

+

The host commits immutable tensor packs to object storage. ClickHouse is built afterward as a searchable index, so database latency and merge work cannot slow capture.

+ Phase 4 catalog implemented and live-tested · opt-in · no CUDA changes +
+ +
+

Short host path

Drain, pack, checksum, upload. No ClickHouse call or summary computation.

+

Bounded parallel I/O

Large packs amortize requests; workers and bytes in flight have hard caps.

+

Large catalog inserts

An independent indexer batches metadata across packs and producers.

+

Selective reads

Agents query catalog metadata first; summaries and selected range hydration are the next phase.

+
+ +
+

The performance boundary

+
+
+ Existing +

Ring² host drain

+

Produces CPU payloads without changing the CUDA capture path.

+
+ +
+ Implemented CPU reference +

Pack assembler

+

Bounds admission and seals by size, count, linger, session, or shutdown.

+
+ +
+ Implemented with Garage/S3 +

Object storage

+

Multipart upload, verified metadata, bounded listing, and exact range reads.

+
+ +
+
Implemented

ClickHouse indexer

Reconciles object listings, reads two footer ranges, and emits large metadata batches.

+
On demand

Reader and agents

Filter in ClickHouse, estimate bytes, then hydrate selected ranges.

+
+
+

Failure isolation: if ClickHouse is slow or offline, captures still commit. Catalog lag increases and the indexer catches up from object storage.

+
+ +
+

Why it should outperform direct ClickHouse payloads

+ + + + + + + + + +
DecisionPerformance effectEvidence required
Remove ClickHouse from captureEliminates insert acknowledgement, merge pressure, and database backpressure on hosts.Lower host enqueue p95 and stable drain throughput while ClickHouse is saturated.
Pack tensor recordsAmortizes request, TLS, and object metadata overhead without losing selective access.Pack-size sweep beats small-object throughput beyond run variance.
Bound upload concurrencyUses parallel requests to reach storage capacity without unbounded host memory or connections.Peak workers and bytes stay at configuration while throughput scales.
Centralize indexingCombines metadata from many hosts into ClickHouse-sized batches.Measured 13,954 rows/s at 1k rows/insert, 88,458 at 10k, and 157,567 at 50k.
Hydrate by rangeAvoids transferring unrelated tensors during analysis.Low read amplification: fetched bytes stay close to requested bytes.
+
+ +
+

The design passes only if measurements agree

+
+
0.502 GiB/s

Best local Garage median: 64 MiB packs, four outer workers.

+
16 MiB

Measured queue cap held exactly at 256 × 64 KiB records.

+
88.5k rows/s

Local ClickHouse median at 10k metadata rows per insert.

+
0 implicit bytes

Metadata queries never fetch tensor payloads.

+
+

Garage v2.3.0 passed multipart retry, listing, metadata verification, and two-range footer reads. ClickHouse 26.9.1 accepted replayed batches while public views returned one logical capture and pack. The local batch sweep improved materially through 50k rows per insert. These prove compatibility and batching behavior; production capacity still needs representative network, object storage, and ClickHouse hardware.

+
+ + +
+ + diff --git a/pyproject.toml b/pyproject.toml index 4c3e29318..ab547ecf2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,6 +9,9 @@ description = "A decoupled, asynchronous observation substrate for high-speed LL requires-python = ">=3.10" dynamic = ["dependencies"] +[project.optional-dependencies] +s3 = ["boto3>=1.40,<2"] + [tool.setuptools.dynamic] dependencies = { file = ["requirements.txt"] } @@ -41,5 +44,6 @@ markers = [ "ring_native: native CUDA ring tests built via tests/native/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", + "garage: requires a reachable Garage S3-compatible endpoint", "numeric: per-hook numeric-difference study (drift vs the unhooked baseline)", ] diff --git a/src/dmi/storage/capture/__init__.py b/src/dmi/storage/capture/__init__.py new file mode 100644 index 000000000..15e9aa88b --- /dev/null +++ b/src/dmi/storage/capture/__init__.py @@ -0,0 +1,148 @@ +"""Immutable capture packs and bounded analysis hydration.""" + +from .filesystem import FilesystemPackStore +from .model import ( + CaptureCatalog, + CaptureDescriptor, + CaptureMetadata, + CapturePage, + CaptureQuery, + CaptureRecord, + CaptureSelection, + CaptureStorageError, + DuplicateCaptureError, + HydratedCapture, + HydrationEstimate, + HydrationLimitError, + ObjectInfo, + ObjectPage, + PackConflictError, + PackFormatError, + PackIntegrityError, + PackRef, + PackSource, + PackStore, + PayloadLocator, + StoredObject, +) +from .pack import PackIndex, PackReader, PackWriter, SealedPack +from .pipeline import ( + AdmissionResult, + BoundedRecordQueue, + DirectPackSink, + FlushReason, + HistogramSnapshot, + HostCapturePipeline, + OverloadPolicy, + OversizedRecordError, + PackAssembler, + PipelineConfig, + PipelineEvent, + PipelineFailedError, + PipelineSnapshot, + QueueSnapshot, + ReadyPack, + object_key_for, +) +from .reader import CaptureReader +from .catalog import ( + CatalogIndexer, + CatalogIndexerConfig, + CatalogReconciler, + CatalogWriter, + IndexEvent, + IndexFailure, + IndexResult, + PackIdentity, + PackInventory, + ReconcileResult, +) +from .clickhouse_catalog import ClickHouseCatalogConfig, ClickHouseCatalogWriter +from .s3 import S3PackStore, S3StoreConfig +from .spool import ( + DurablePackSink, + DurablePackSpool, + SpoolFullError, + SpoolSnapshot, + SpoolUploader, + StagedPack, + ParallelSpoolUploader, + ParallelUploadConfig, + UploadBatchResult, + UploadEvent, + UploadFailure, + UploadSnapshot, +) + +__all__ = [ + "AdmissionResult", + "BoundedRecordQueue", + "CaptureCatalog", + "CaptureDescriptor", + "CaptureMetadata", + "CapturePage", + "CaptureQuery", + "CaptureReader", + "CaptureRecord", + "CaptureSelection", + "CaptureStorageError", + "CatalogIndexer", + "CatalogIndexerConfig", + "CatalogReconciler", + "CatalogWriter", + "ClickHouseCatalogConfig", + "ClickHouseCatalogWriter", + "DuplicateCaptureError", + "DirectPackSink", + "DurablePackSink", + "DurablePackSpool", + "FilesystemPackStore", + "FlushReason", + "HistogramSnapshot", + "HostCapturePipeline", + "HydratedCapture", + "HydrationEstimate", + "HydrationLimitError", + "IndexEvent", + "IndexFailure", + "IndexResult", + "ObjectInfo", + "ObjectPage", + "OverloadPolicy", + "OversizedRecordError", + "PackAssembler", + "PackConflictError", + "PackFormatError", + "PackIntegrityError", + "PackIndex", + "PackIdentity", + "PackInventory", + "PackReader", + "PackRef", + "PackSource", + "PackStore", + "PackWriter", + "ParallelSpoolUploader", + "ParallelUploadConfig", + "PayloadLocator", + "PipelineConfig", + "PipelineEvent", + "PipelineFailedError", + "PipelineSnapshot", + "QueueSnapshot", + "ReadyPack", + "ReconcileResult", + "S3PackStore", + "S3StoreConfig", + "SealedPack", + "SpoolFullError", + "SpoolSnapshot", + "SpoolUploader", + "StagedPack", + "StoredObject", + "UploadBatchResult", + "UploadEvent", + "UploadFailure", + "UploadSnapshot", + "object_key_for", +] diff --git a/src/dmi/storage/capture/catalog.py b/src/dmi/storage/capture/catalog.py new file mode 100644 index 000000000..3c16c4a7c --- /dev/null +++ b/src/dmi/storage/capture/catalog.py @@ -0,0 +1,288 @@ +from __future__ import annotations + +from dataclasses import asdict, dataclass +import json +from time import monotonic_ns, time_ns +from typing import Callable, Protocol, Sequence + +from .model import CaptureDescriptor, ObjectPage, PackRef, PackStore +from .pack import PackIndex + + +PackIdentity = tuple[str, str] + + +class PackInventory(PackStore, Protocol): + def inspect(self, object_key: str) -> PackRef: ... + def list_objects( + self, *, prefix: str = "", cursor: str | None = None, limit: int = 1000 + ) -> ObjectPage: ... + + +class CatalogWriter(Protocol): + def committed_pack_ids( + self, identities: Sequence[PackIdentity] + ) -> set[PackIdentity]: ... + def write_descriptors( + self, descriptors: Sequence[CaptureDescriptor], *, index_version: int + ) -> None: ... + def commit_packs( + self, refs: Sequence[PackRef], *, index_version: int + ) -> None: ... + + +@dataclass(frozen=True, slots=True) +class CatalogIndexerConfig: + max_packs: int = 64 + max_rows_per_insert: int = 10_000 + max_estimated_bytes: int = 128 * 1024**2 + max_failure_details: int = 128 + + def __post_init__(self) -> None: + for name in ( + "max_packs", + "max_rows_per_insert", + "max_estimated_bytes", + "max_failure_details", + ): + value = getattr(self, name) + if type(value) is not int or value <= 0: + raise ValueError(f"{name} must be positive") + + +@dataclass(frozen=True, slots=True) +class IndexFailure: + pack_id: str + object_key: str + error_type: str + message: str + + +@dataclass(frozen=True, slots=True) +class IndexResult: + requested_packs: int = 0 + skipped_packs: int = 0 + indexed_packs: int = 0 + indexed_rows: int = 0 + failed_packs: int = 0 + descriptor_inserts: int = 0 + estimated_bytes: int = 0 + elapsed_ns: int = 0 + failures: tuple[IndexFailure, ...] = () + + def merge( + self, other: IndexResult, *, failure_limit: int | None = None + ) -> IndexResult: + failures = self.failures + other.failures + if failure_limit is not None: + failures = failures[:failure_limit] + return IndexResult( + requested_packs=self.requested_packs + other.requested_packs, + skipped_packs=self.skipped_packs + other.skipped_packs, + indexed_packs=self.indexed_packs + other.indexed_packs, + indexed_rows=self.indexed_rows + other.indexed_rows, + failed_packs=self.failed_packs + other.failed_packs, + descriptor_inserts=self.descriptor_inserts + other.descriptor_inserts, + estimated_bytes=self.estimated_bytes + other.estimated_bytes, + elapsed_ns=self.elapsed_ns + other.elapsed_ns, + failures=failures, + ) + + +@dataclass(frozen=True, slots=True) +class IndexEvent: + event: str + requested_packs: int + skipped_packs: int + indexed_packs: int + indexed_rows: int + failed_packs: int + descriptor_inserts: int + estimated_bytes: int + elapsed_ns: int + + +@dataclass(frozen=True, slots=True) +class ReconcileResult: + index: IndexResult + next_cursor: str | None + pages: int + + +def _identity(ref: PackRef) -> PackIdentity: + return ref.store_id, ref.pack_id + + +def _deduplicate_refs(refs: Sequence[PackRef]) -> tuple[PackRef, ...]: + by_identity: dict[PackIdentity, PackRef] = {} + for ref in refs: + identity = _identity(ref) + current = by_identity.get(identity) + if current is not None and current != ref: + raise ValueError(f"conflicting pack identity: {identity!r}") + by_identity[identity] = ref + return tuple(by_identity.values()) + + +def _estimated_bytes(descriptor: CaptureDescriptor) -> int: + value = { + "metadata": descriptor.metadata.to_mapping(), + "locator": asdict(descriptor.locator), + } + return len(json.dumps(value, separators=(",", ":"), ensure_ascii=False).encode()) + + +class CatalogIndexer: + def __init__( + self, + store: PackStore, + writer: CatalogWriter, + *, + config: CatalogIndexerConfig | None = None, + clock_ns: Callable[[], int] = time_ns, + timer_ns: Callable[[], int] = monotonic_ns, + on_event: Callable[[IndexEvent], None] | None = None, + ) -> None: + self._store = store + self._writer = writer + self._config = config or CatalogIndexerConfig() + self._clock_ns = clock_ns + self._timer_ns = timer_ns + self._on_event = on_event + self._callback_failures = 0 + + @property + def store_id(self) -> str: + return self._store.store_id + + @property + def max_packs(self) -> int: + return self._config.max_packs + + @property + def callback_failures(self) -> int: + return self._callback_failures + + @property + def max_failure_details(self) -> int: + return self._config.max_failure_details + + def index(self, refs: Sequence[PackRef]) -> IndexResult: + started_ns = self._timer_ns() + if len(refs) > self._config.max_packs: + raise ValueError( + f"pack batch exceeds max_packs: {len(refs)} > {self._config.max_packs}" + ) + unique = _deduplicate_refs(refs) + identities = tuple(_identity(ref) for ref in unique) + committed = self._writer.committed_pack_ids(identities) if identities else set() + pending = tuple(ref for ref in unique if _identity(ref) not in committed) + descriptors: list[CaptureDescriptor] = [] + valid_refs: list[PackRef] = [] + failures: list[IndexFailure] = [] + estimated_bytes = 0 + for ref in pending: + try: + pack_descriptors = PackIndex.from_store(self._store, ref).descriptors() + pack_bytes = sum(_estimated_bytes(item) for item in pack_descriptors) + if estimated_bytes + pack_bytes > self._config.max_estimated_bytes: + raise ValueError("catalog batch exceeds max_estimated_bytes") + except Exception as exc: + failures.append( + IndexFailure( + pack_id=ref.pack_id, + object_key=ref.object_key, + error_type=type(exc).__name__, + message=str(exc)[:512], + ) + ) + continue + estimated_bytes += pack_bytes + descriptors.extend(pack_descriptors) + valid_refs.append(ref) + + version = self._clock_ns() + if type(version) is not int or version < 0: + raise ValueError("clock_ns must return a non-negative integer") + step = self._config.max_rows_per_insert + descriptor_inserts = 0 + for start in range(0, len(descriptors), step): + self._writer.write_descriptors( + descriptors[start : start + step], index_version=version + ) + descriptor_inserts += 1 + if valid_refs: + self._writer.commit_packs(valid_refs, index_version=version) + result = IndexResult( + requested_packs=len(unique), + skipped_packs=len(unique) - len(pending), + indexed_packs=len(valid_refs), + indexed_rows=len(descriptors), + failed_packs=len(failures), + descriptor_inserts=descriptor_inserts, + estimated_bytes=estimated_bytes, + elapsed_ns=self._timer_ns() - started_ns, + failures=tuple(failures), + ) + self._emit(result) + return result + + def _emit(self, result: IndexResult) -> None: + if self._on_event is None: + return + try: + self._on_event( + IndexEvent( + event="catalog_index_completed", + requested_packs=result.requested_packs, + skipped_packs=result.skipped_packs, + indexed_packs=result.indexed_packs, + indexed_rows=result.indexed_rows, + failed_packs=result.failed_packs, + descriptor_inserts=result.descriptor_inserts, + estimated_bytes=result.estimated_bytes, + elapsed_ns=result.elapsed_ns, + ) + ) + except Exception: + self._callback_failures += 1 + + +class CatalogReconciler: + def __init__(self, inventory: PackInventory, indexer: CatalogIndexer) -> None: + if inventory.store_id != indexer.store_id: + raise ValueError("inventory and indexer store IDs differ") + self._inventory = inventory + self._indexer = indexer + + def index_object_keys(self, object_keys: Sequence[str]) -> IndexResult: + if len(object_keys) > self._indexer.max_packs: + raise ValueError("object notification batch exceeds max_packs") + keys = tuple(dict.fromkeys(object_keys)) + return self._indexer.index([self._inventory.inspect(key) for key in keys]) + + def reconcile_page( + self, *, prefix: str = "", cursor: str | None = None, limit: int = 64 + ) -> ReconcileResult: + if limit > self._indexer.max_packs: + raise ValueError("listing limit exceeds max_packs") + page = self._inventory.list_objects(prefix=prefix, cursor=cursor, limit=limit) + result = self.index_object_keys([item.object_key for item in page.items]) + return ReconcileResult(index=result, next_cursor=page.next_cursor, pages=1) + + def rebuild( + self, *, prefix: str = "", page_size: int = 64, max_pages: int = 10_000 + ) -> IndexResult: + if type(max_pages) is not int or max_pages <= 0: + raise ValueError("max_pages must be positive") + cursor = None + total = IndexResult() + for _ in range(max_pages): + page = self.reconcile_page(prefix=prefix, cursor=cursor, limit=page_size) + total = total.merge( + page.index, failure_limit=self._indexer.max_failure_details + ) + cursor = page.next_cursor + if cursor is None: + return total + raise RuntimeError("rebuild exceeded max_pages") diff --git a/src/dmi/storage/capture/clickhouse_catalog.py b/src/dmi/storage/capture/clickhouse_catalog.py new file mode 100644 index 000000000..6d3c7f465 --- /dev/null +++ b/src/dmi/storage/capture/clickhouse_catalog.py @@ -0,0 +1,190 @@ +from __future__ import annotations + +from dataclasses import dataclass +import re +from typing import Protocol, Sequence + +from .catalog import PackIdentity +from .model import CaptureDescriptor, PackRef + + +_IDENTIFIER = re.compile(r"[A-Za-z_][A-Za-z0-9_]*\Z") + + +class ClickHouseClient(Protocol): + def execute(self, query: str, params=None, **kwargs): ... + + +def _identifier(value: str) -> str: + if not isinstance(value, str) or _IDENTIFIER.fullmatch(value) is None: + raise ValueError(f"invalid ClickHouse identifier: {value!r}") + return value + + +def _quoted(value: str) -> str: + return f"`{value}`" + + +@dataclass(frozen=True, slots=True) +class ClickHouseCatalogConfig: + database: str = "default" + table_prefix: str = "dmi" + query_pack_limit: int = 10_000 + + def __post_init__(self) -> None: + _identifier(self.database) + _identifier(self.table_prefix) + if type(self.query_pack_limit) is not int or self.query_pack_limit <= 0: + raise ValueError("query_pack_limit must be positive") + + +_CAPTURE_COLUMNS = ( + "capture_id", "tenant_id", "experiment_id", "run_id", "session_id", + "request_id", "sequence_id", "model_id", "model_revision", + "adapter_revision", "capture_policy_version", "hook_name", "layer_number", + "producer_rank", "step_number", "token_start", "token_end", + "batch_position", "dtype", "shape", "captured_at_ns", "pack_id", + "store_id", "object_key", "object_bytes", "pack_checksum", + "pack_record_count", "payload_offset", "stored_length", "decoded_length", + "codec", "payload_checksum", "index_version", +) + +_PACK_COLUMNS = ( + "pack_id", "store_id", "object_key", "object_bytes", "pack_checksum", + "record_count", "index_version", +) + + +class ClickHouseCatalogWriter: + def __init__( + self, client: ClickHouseClient, config: ClickHouseCatalogConfig | None = None + ) -> None: + self._client = client + self._config = config or ClickHouseCatalogConfig() + prefix = self._config.table_prefix + self._capture_raw = f"{prefix}_capture_raw" + self._capture_view = f"{prefix}_capture" + self._pack_raw = f"{prefix}_pack_inventory_raw" + self._pack_view = f"{prefix}_pack_inventory" + + def ensure_schema(self) -> None: + database = _quoted(self._config.database) + capture_raw = f"{database}.{_quoted(self._capture_raw)}" + capture_view = f"{database}.{_quoted(self._capture_view)}" + pack_raw = f"{database}.{_quoted(self._pack_raw)}" + pack_view = f"{database}.{_quoted(self._pack_view)}" + self._client.execute(f"CREATE DATABASE IF NOT EXISTS {database}") + self._client.execute( + f"""CREATE TABLE IF NOT EXISTS {capture_raw} ( +capture_id String, tenant_id String, experiment_id String, run_id String, +session_id String, request_id String, sequence_id String, model_id String, +model_revision String, adapter_revision Nullable(String), +capture_policy_version String, hook_name LowCardinality(String), layer_number Int32, +producer_rank UInt32, step_number UInt64, token_start UInt64, token_end UInt64, +batch_position UInt32, dtype LowCardinality(String), shape Array(UInt32), +captured_at_ns UInt64, pack_id UUID, store_id LowCardinality(String), object_key String, +object_bytes UInt64, pack_checksum FixedString(64), pack_record_count UInt32, +payload_offset UInt64, stored_length UInt64, decoded_length UInt64, +codec LowCardinality(String), payload_checksum FixedString(8), index_version UInt64 +) ENGINE = ReplacingMergeTree(index_version) +ORDER BY (tenant_id, experiment_id, run_id, captured_at_ns, capture_id)""" + ) + self._client.execute( + f"""CREATE TABLE IF NOT EXISTS {pack_raw} ( +pack_id UUID, store_id LowCardinality(String), object_key String, +object_bytes UInt64, pack_checksum FixedString(64), record_count UInt32, +index_version UInt64 +) ENGINE = ReplacingMergeTree(index_version) +ORDER BY (store_id, pack_id)""" + ) + capture_public = ", ".join(_CAPTURE_COLUMNS[:-1]) + pack_public = ", ".join(_PACK_COLUMNS[:-1]) + self._client.execute( + f"CREATE VIEW IF NOT EXISTS {capture_view} AS " + f"SELECT {capture_public} FROM {capture_raw} FINAL" + ) + self._client.execute( + f"CREATE VIEW IF NOT EXISTS {pack_view} AS " + f"SELECT {pack_public} FROM {pack_raw} FINAL" + ) + + def committed_pack_ids( + self, identities: Sequence[PackIdentity] + ) -> set[PackIdentity]: + if not identities: + return set() + if len(identities) > self._config.query_pack_limit: + raise ValueError("pack identity query exceeds query_pack_limit") + table = self._qualified(self._pack_view) + rows = self._client.execute( + f"SELECT store_id, toString(pack_id) FROM {table} " + "WHERE (store_id, pack_id) IN %(identities)s", + {"identities": list(identities)}, + ) + return {(self._text(row[0]), self._text(row[1])) for row in rows} + + def write_descriptors( + self, descriptors: Sequence[CaptureDescriptor], *, index_version: int + ) -> None: + if not descriptors: + return + self._validate_version(index_version) + rows = [self._descriptor_row(item, index_version) for item in descriptors] + self._client.execute( + f"INSERT INTO {self._qualified(self._capture_raw)} " + f"({', '.join(_CAPTURE_COLUMNS)}) VALUES", + rows, + ) + + def commit_packs( + self, refs: Sequence[PackRef], *, index_version: int + ) -> None: + if not refs: + return + self._validate_version(index_version) + rows = [ + ( + ref.pack_id, ref.store_id, ref.object_key, ref.object_bytes, + ref.checksum, ref.record_count, index_version, + ) + for ref in refs + ] + self._client.execute( + f"INSERT INTO {self._qualified(self._pack_raw)} " + f"({', '.join(_PACK_COLUMNS)}) VALUES", + rows, + ) + + def _qualified(self, table: str) -> str: + return f"{_quoted(self._config.database)}.{_quoted(table)}" + + @staticmethod + def _validate_version(index_version: int) -> None: + if type(index_version) is not int or not 0 <= index_version < 2**64: + raise ValueError("index_version must fit UInt64") + + @staticmethod + def _text(value: object) -> str: + if isinstance(value, bytes): + return value.decode("utf-8") + if not isinstance(value, str): + raise ValueError("ClickHouse returned an invalid pack identity") + return value + + @staticmethod + def _descriptor_row(item: CaptureDescriptor, index_version: int) -> tuple: + metadata = item.metadata + locator = item.locator + return ( + metadata.capture_id, metadata.tenant_id, metadata.experiment_id, + metadata.run_id, metadata.session_id, metadata.request_id, + metadata.sequence_id, metadata.model_id, metadata.model_revision, + metadata.adapter_revision, metadata.capture_policy_version, + metadata.hook_name, metadata.layer_number, metadata.producer_rank, + metadata.step_number, metadata.token_start, metadata.token_end, + metadata.batch_position, metadata.dtype, list(metadata.shape), + metadata.captured_at_ns, locator.pack_id, locator.store_id, + locator.object_key, locator.object_bytes, locator.pack_checksum, + locator.pack_record_count, locator.offset, locator.stored_length, + locator.decoded_length, locator.codec, locator.checksum, index_version, + ) diff --git a/src/dmi/storage/capture/filesystem.py b/src/dmi/storage/capture/filesystem.py new file mode 100644 index 000000000..cd6a885a6 --- /dev/null +++ b/src/dmi/storage/capture/filesystem.py @@ -0,0 +1,212 @@ +from __future__ import annotations + +from hashlib import sha256 +import errno +import os +from pathlib import Path, PurePosixPath +import re +import stat as stat_mode +import tempfile +from typing import BinaryIO +from uuid import UUID + +from .model import ( + ObjectInfo, + PackConflictError, + PackFormatError, + PackIntegrityError, + PackRef, + PackSource, +) + + +_KEY_COMPONENT = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._=%-]*$") +_SHA256 = re.compile(r"^[0-9a-f]{64}$") +_COPY_BYTES = 1024 * 1024 + + +def validate_object_key(object_key: str) -> PurePosixPath: + if not isinstance(object_key, str) or not object_key or "\\" in object_key: + raise ValueError("object key is invalid") + key = PurePosixPath(object_key) + if key.is_absolute() or any( + part in ("", ".", "..") or _KEY_COMPONENT.fullmatch(part) is None + for part in key.parts + ): + raise ValueError("object key is invalid") + return key + + +def validate_pack_source(pack: PackSource) -> None: + if not isinstance(pack, PackSource): + raise TypeError("pack must implement PackSource") + try: + canonical_id = str(UUID(pack.pack_id)) + except (TypeError, ValueError, AttributeError) as exc: + raise ValueError("pack source has an invalid pack ID") from exc + if canonical_id != pack.pack_id: + raise ValueError("pack source must use a canonical pack ID") + integers = (pack.created_at_ns, pack.record_count, pack.object_bytes) + if any( + type(value) is not int or value < 0 or value > 2**63 - 1 + for value in integers + ): + raise ValueError("pack source has invalid numeric metadata") + if pack.record_count == 0 or pack.object_bytes == 0: + raise ValueError("pack source must contain records and bytes") + if not isinstance(pack.checksum, str) or _SHA256.fullmatch(pack.checksum) is None: + raise ValueError("pack source has an invalid checksum") + + +def copy_pack_source(pack: PackSource, destination: BinaryIO) -> None: + validate_pack_source(pack) + digest = sha256() + remaining = pack.object_bytes + with pack.open() as source: + while remaining: + requested = min(_COPY_BYTES, remaining) + chunk = source.read(requested) + if not isinstance(chunk, bytes) or not chunk or len(chunk) > requested: + raise PackIntegrityError("pack source returned an invalid byte stream") + destination.write(chunk) + digest.update(chunk) + remaining -= len(chunk) + if source.read(1): + raise PackIntegrityError("pack source exceeds its declared size") + if digest.hexdigest() != pack.checksum: + raise PackIntegrityError("pack source checksum does not match its metadata") + + +class FilesystemPackStore: + def __init__(self, root: str | Path, *, store_id: str = "filesystem") -> None: + if not store_id or len(store_id.encode("utf-8")) > 128: + raise ValueError("store_id must be non-empty and at most 128 bytes") + self.root = Path(root).resolve() + self.root.mkdir(parents=True, exist_ok=True) + self.store_id = store_id + + @staticmethod + def _validate_key(object_key: str) -> PurePosixPath: + return validate_object_key(object_key) + + def _path(self, object_key: str) -> Path: + key = self._validate_key(object_key) + path = self.root.joinpath(*key.parts) + parent = path.parent.resolve(strict=False) + if not parent.is_relative_to(self.root): + raise ValueError("object key escapes the store root") + return path + + @staticmethod + def _checksum(path: Path) -> str: + digest = sha256() + with FilesystemPackStore._open_regular(path) as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + @staticmethod + def _open_regular(path: Path) -> BinaryIO: + flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) + try: + descriptor = os.open(path, flags) + except OSError as exc: + if exc.errno in (errno.ELOOP, errno.EMLINK): + raise PackFormatError("pack object must be a regular file") from exc + raise + if not stat_mode.S_ISREG(os.fstat(descriptor).st_mode): + os.close(descriptor) + raise PackFormatError("pack object must be a regular file") + return os.fdopen(descriptor, "rb") + + @staticmethod + def _size(path: Path) -> int: + status = os.stat(path, follow_symlinks=False) + if not stat_mode.S_ISREG(status.st_mode): + raise PackFormatError("pack object must be a regular file") + return status.st_size + + def _existing_ref(self, pack: PackSource, object_key: str, path: Path) -> PackRef: + if self._size(path) != pack.object_bytes or self._checksum(path) != pack.checksum: + raise PackConflictError(f"object key contains different content: {object_key}") + return PackRef( + pack_id=pack.pack_id, + store_id=self.store_id, + object_key=object_key, + object_bytes=pack.object_bytes, + checksum=pack.checksum, + record_count=pack.record_count, + ) + + def put(self, pack: PackSource, object_key: str) -> PackRef: + validate_pack_source(pack) + path = self._path(object_key) + path.parent.mkdir(parents=True, exist_ok=True) + if path.exists(): + return self._existing_ref(pack, object_key, path) + + temp_path: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="wb", + prefix=f".{path.name}.", + suffix=".open", + dir=path.parent, + delete=False, + ) as handle: + temp_path = Path(handle.name) + copy_pack_source(pack, handle) + handle.flush() + os.fsync(handle.fileno()) + try: + os.link(temp_path, path) + except FileExistsError: + return self._existing_ref(pack, object_key, path) + directory_fd = os.open(path.parent, os.O_RDONLY) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) + finally: + if temp_path is not None: + temp_path.unlink(missing_ok=True) + + return PackRef( + pack_id=pack.pack_id, + store_id=self.store_id, + object_key=object_key, + object_bytes=pack.object_bytes, + checksum=pack.checksum, + record_count=pack.record_count, + ) + + def stat(self, ref: PackRef) -> ObjectInfo: + self._validate_ref(ref) + path = self._path(ref.object_key) + return ObjectInfo(size=self._size(path), checksum=self._checksum(path)) + + def read_range(self, ref: PackRef, offset: int, length: int) -> bytes: + self._validate_ref(ref) + if ( + type(offset) is not int + or type(length) is not int + or offset < 0 + or length < 0 + ): + raise ValueError("range offset and length must be non-negative integers") + path = self._path(ref.object_key) + size = self._size(path) + if offset + length > size: + raise PackFormatError("requested range exceeds object size") + with self._open_regular(path) as handle: + handle.seek(offset) + data = handle.read(length) + if len(data) != length: + raise PackIntegrityError("filesystem returned a short range") + return data + + def _validate_ref(self, ref: PackRef) -> None: + if ref.store_id != self.store_id: + raise ValueError( + f"pack store mismatch: {ref.store_id!r} != {self.store_id!r}" + ) diff --git a/src/dmi/storage/capture/model.py b/src/dmi/storage/capture/model.py new file mode 100644 index 000000000..6bd7adf15 --- /dev/null +++ b/src/dmi/storage/capture/model.py @@ -0,0 +1,372 @@ +from __future__ import annotations + +from dataclasses import asdict, dataclass +from hashlib import sha256 +import json +import math +from typing import BinaryIO, Mapping, Protocol, Sequence, runtime_checkable + + + +_DTYPE_BYTES = { + "bool": 1, + "uint8": 1, + "int8": 1, + "int16": 2, + "float16": 2, + "bfloat16": 2, + "int32": 4, + "float32": 4, + "int64": 8, + "float64": 8, +} +_TEXT_LIMIT = 512 +_MAX_RANK = 32 + + +class CaptureStorageError(Exception): + """Base error for capture-pack storage.""" + + +class PackFormatError(CaptureStorageError): + """A pack or catalog descriptor violates the format contract.""" + + +class PackIntegrityError(CaptureStorageError): + """Stored bytes do not match their integrity metadata.""" + + +class PackConflictError(CaptureStorageError): + """An immutable object key already contains different bytes.""" + + +class DuplicateCaptureError(CaptureStorageError): + """A logical selection contains a capture more than once.""" + + +class HydrationLimitError(CaptureStorageError): + """A hydration plan exceeds the caller's byte limit.""" + + +def _validate_text(name: str, value: str | None, *, optional: bool = False) -> None: + if value is None: + if optional: + return + raise ValueError(f"{name} is required") + if not isinstance(value, str) or not value or len(value.encode("utf-8")) > _TEXT_LIMIT: + raise ValueError(f"{name} must be non-empty UTF-8 within {_TEXT_LIMIT} bytes") + + +@dataclass(frozen=True, slots=True) +class CaptureMetadata: + capture_id: str + tenant_id: str + experiment_id: str + run_id: str + session_id: str + request_id: str + sequence_id: str + model_id: str + model_revision: str + adapter_revision: str | None + capture_policy_version: str + hook_name: str + layer_number: int + producer_rank: int + step_number: int + token_start: int + token_end: int + batch_position: int + dtype: str + shape: tuple[int, ...] + captured_at_ns: int + + def __post_init__(self) -> None: + required = ( + "capture_id", + "tenant_id", + "experiment_id", + "run_id", + "session_id", + "request_id", + "sequence_id", + "model_id", + "model_revision", + "capture_policy_version", + "hook_name", + ) + for name in required: + _validate_text(name, getattr(self, name)) + _validate_text("adapter_revision", self.adapter_revision, optional=True) + if self.dtype not in _DTYPE_BYTES: + raise ValueError(f"unsupported dtype: {self.dtype!r}") + if not isinstance(self.shape, tuple): + object.__setattr__(self, "shape", tuple(self.shape)) + if len(self.shape) > _MAX_RANK: + raise ValueError(f"shape rank must not exceed {_MAX_RANK}") + if any(type(dim) is not int or dim < 0 or dim > 2**31 - 1 for dim in self.shape): + raise ValueError("shape dimensions must be integers in [0, 2^31 - 1]") + non_negative = ( + "producer_rank", + "step_number", + "token_start", + "token_end", + "batch_position", + "captured_at_ns", + ) + for name in non_negative: + if type(getattr(self, name)) is not int or getattr(self, name) < 0: + raise ValueError(f"{name} must be a non-negative integer") + if type(self.layer_number) is not int or self.layer_number < -1: + raise ValueError("layer_number must be an integer >= -1") + if self.token_end < self.token_start: + raise ValueError("token_end must be >= token_start") + + @property + def logical_bytes(self) -> int: + return math.prod(self.shape) * _DTYPE_BYTES[self.dtype] + + def to_mapping(self) -> dict[str, object]: + result = asdict(self) + result["shape"] = list(self.shape) + return result + + @classmethod + def from_mapping(cls, value: Mapping[str, object]) -> CaptureMetadata: + fields = cls.__dataclass_fields__ + missing = [name for name in fields if name not in value] + if missing: + raise PackFormatError("capture metadata is missing: " + ", ".join(missing)) + selected = {name: value[name] for name in fields} + shape = selected["shape"] + if not isinstance(shape, list) or not all(type(dim) is int for dim in shape): + raise PackFormatError("capture shape must be an integer list") + selected["shape"] = tuple(shape) + try: + return cls(**selected) # type: ignore[arg-type] + except (TypeError, ValueError) as exc: + raise PackFormatError(f"invalid capture metadata: {exc}") from exc + + +@dataclass(frozen=True, slots=True) +class CaptureRecord: + metadata: CaptureMetadata + payload: bytes + + def __post_init__(self) -> None: + if not isinstance(self.payload, bytes): + object.__setattr__(self, "payload", bytes(self.payload)) + if len(self.payload) != self.metadata.logical_bytes: + raise ValueError( + "payload length does not match dtype and shape: " + f"{len(self.payload)} != {self.metadata.logical_bytes}" + ) + + +@dataclass(frozen=True, slots=True) +class PackRef: + pack_id: str + store_id: str + object_key: str + object_bytes: int + checksum: str + record_count: int + + +@dataclass(frozen=True, slots=True) +class ObjectInfo: + size: int + checksum: str + + +@dataclass(frozen=True, slots=True) +class StoredObject: + object_key: str + object_bytes: int + + +@dataclass(frozen=True, slots=True) +class ObjectPage: + items: tuple[StoredObject, ...] + next_cursor: str | None + + +@dataclass(frozen=True, slots=True) +class PayloadLocator: + pack_id: str + store_id: str + object_key: str + object_bytes: int + pack_checksum: str + pack_record_count: int + offset: int + stored_length: int + decoded_length: int + codec: str + checksum: str + + @property + def pack_ref(self) -> PackRef: + return PackRef( + pack_id=self.pack_id, + store_id=self.store_id, + object_key=self.object_key, + object_bytes=self.object_bytes, + checksum=self.pack_checksum, + record_count=self.pack_record_count, + ) + + +@dataclass(frozen=True, slots=True) +class CaptureDescriptor: + metadata: CaptureMetadata + locator: PayloadLocator + + @property + def capture_id(self) -> str: + return self.metadata.capture_id + + +@dataclass(frozen=True, slots=True) +class CaptureQuery: + tenant_id: str | None = None + experiment_id: str | None = None + run_id: str | None = None + session_id: str | None = None + model_id: str | None = None + hook_names: tuple[str, ...] = () + layer_numbers: tuple[int, ...] = () + captured_after_ns: int | None = None + captured_before_ns: int | None = None + cursor: str | None = None + limit: int = 1000 + + def __post_init__(self) -> None: + for name in ("tenant_id", "experiment_id", "run_id", "session_id", "model_id", "cursor"): + value = getattr(self, name) + if value is not None: + _validate_text(name, value) + if not 1 <= self.limit <= 10_000: + raise ValueError("limit must be between 1 and 10000") + if len(self.hook_names) > 128 or len(self.layer_numbers) > 1024: + raise ValueError("query filters exceed their bounded cardinality") + for hook_name in self.hook_names: + _validate_text("hook_name", hook_name) + if any(layer < -1 for layer in self.layer_numbers): + raise ValueError("layer numbers must be >= -1") + for name in ("captured_after_ns", "captured_before_ns"): + value = getattr(self, name) + if value is not None and (type(value) is not int or value < 0): + raise ValueError(f"{name} must be a non-negative integer") + if ( + self.captured_after_ns is not None + and self.captured_before_ns is not None + and self.captured_before_ns < self.captured_after_ns + ): + raise ValueError("captured_before_ns must be >= captured_after_ns") + + @property + def query_hash(self) -> str: + encoded = json.dumps(asdict(self), sort_keys=True, separators=(",", ":")).encode() + return sha256(encoded).hexdigest() + + +@dataclass(frozen=True, slots=True) +class CapturePage: + items: tuple[CaptureDescriptor, ...] + next_cursor: str | None + watermark: str + + +@dataclass(frozen=True, slots=True) +class CaptureSelection: + selection_id: str + capture_ids: tuple[str, ...] + catalog_watermark: str + query_hash: str + + @classmethod + def create( + cls, + descriptors: Sequence[CaptureDescriptor], + *, + catalog_watermark: str, + query_hash: str, + ) -> CaptureSelection: + ids = tuple(item.capture_id for item in descriptors) + seen: set[str] = set() + for capture_id in ids: + if capture_id in seen: + raise DuplicateCaptureError( + f"duplicate logical capture: {capture_id}" + ) + seen.add(capture_id) + identity = json.dumps( + { + "version": 1, + "catalog_watermark": catalog_watermark, + "query_hash": query_hash, + "capture_ids": ids, + }, + sort_keys=True, + separators=(",", ":"), + ).encode() + return cls( + selection_id=sha256(identity).hexdigest(), + capture_ids=ids, + catalog_watermark=catalog_watermark, + query_hash=query_hash, + ) + + +@dataclass(frozen=True, slots=True) +class HydrationEstimate: + capture_count: int + object_count: int + request_count: int + logical_bytes: int + stored_bytes: int + request_bytes: int + + @property + def read_amplification(self) -> float: + return self.request_bytes / self.stored_bytes if self.stored_bytes else 0.0 + + +@dataclass(frozen=True, slots=True) +class HydratedCapture: + descriptor: CaptureDescriptor + payload: bytes + + @property + def capture_id(self) -> str: + return self.descriptor.capture_id + + +@runtime_checkable +class PackSource(Protocol): + pack_id: str + created_at_ns: int + record_count: int + checksum: str + + @property + def object_bytes(self) -> int: ... + def open(self) -> BinaryIO: ... + + +@runtime_checkable +class PackStore(Protocol): + store_id: str + + def put(self, pack: PackSource, object_key: str) -> PackRef: ... + def stat(self, ref: PackRef) -> ObjectInfo: ... + def read_range(self, ref: PackRef, offset: int, length: int) -> bytes: ... + + +@runtime_checkable +class CaptureCatalog(Protocol): + def search(self, query: CaptureQuery) -> CapturePage: ... + def get_by_ids( + self, capture_ids: Sequence[str], *, watermark: str + ) -> Sequence[CaptureDescriptor]: ... diff --git a/src/dmi/storage/capture/pack.py b/src/dmi/storage/capture/pack.py new file mode 100644 index 000000000..ecd1ee898 --- /dev/null +++ b/src/dmi/storage/capture/pack.py @@ -0,0 +1,526 @@ +from __future__ import annotations + +from dataclasses import dataclass +from hashlib import sha256 +from io import BytesIO +import json +import re +import struct +from typing import Iterable +from uuid import UUID +import zlib + +from .model import ( + CaptureDescriptor, + CaptureMetadata, + CaptureRecord, + DuplicateCaptureError, + PackFormatError, + PackIntegrityError, + PackRef, + PackStore, + PayloadLocator, +) + + +PACK_MAJOR_VERSION = 1 +PACK_MINOR_VERSION = 0 +PACK_ALIGNMENT = 64 +MAX_FOOTER_BYTES = 64 * 1024 * 1024 +MAX_RECORDS = 1_000_000 + +_HEADER_MAGIC = b"DMIPACK\0" +_TRAILER_MAGIC = b"DMIFTR\0\0" +_HEADER = struct.Struct("<8sHHI16sQI20s") +_TRAILER = struct.Struct("<8sHHQQI32s") +_CRC_PATTERN = re.compile(r"^[0-9a-f]{8}$") + + +class PackCapacityError(ValueError): + pass + + +class _PackStateError(RuntimeError): + pass + + +@dataclass(frozen=True, slots=True) +class SealedPack: + pack_id: str + created_at_ns: int + data: bytes + record_count: int + footer_offset: int + checksum: str + + @property + def object_bytes(self) -> int: + return len(self.data) + + def open(self) -> BytesIO: + return BytesIO(self.data) + + +@dataclass(frozen=True, slots=True) +class _IndexedRecord: + metadata: CaptureMetadata + offset: int + stored_length: int + decoded_length: int + codec: str + checksum: str + + +def _align(buffer: bytearray) -> None: + padding = (-len(buffer)) % PACK_ALIGNMENT + if padding: + buffer.extend(b"\0" * padding) + + +def _aligned_length(length: int) -> int: + return length + (-length) % PACK_ALIGNMENT + + +def _encode_json(value: object) -> bytes: + return json.dumps(value, sort_keys=True, separators=(",", ":")).encode("utf-8") + + +def _record_mapping(item: _IndexedRecord) -> dict[str, object]: + return { + "metadata": item.metadata.to_mapping(), + "offset": item.offset, + "stored_length": item.stored_length, + "decoded_length": item.decoded_length, + "codec": item.codec, + "checksum": item.checksum, + } + + +def _crc32(data: bytes | memoryview) -> str: + return f"{zlib.crc32(data) & 0xFFFFFFFF:08x}" + + +def verify_payload(descriptor: CaptureDescriptor, payload: bytes | memoryview) -> None: + if len(payload) != descriptor.locator.stored_length: + raise PackIntegrityError( + f"short record for {descriptor.capture_id}: " + f"{len(payload)} != {descriptor.locator.stored_length}" + ) + if descriptor.locator.codec != "none": + raise PackFormatError(f"unsupported codec: {descriptor.locator.codec}") + if descriptor.locator.stored_length != descriptor.locator.decoded_length: + raise PackFormatError("none codec requires equal stored and decoded lengths") + if _crc32(payload) != descriptor.locator.checksum: + raise PackIntegrityError(f"record checksum mismatch: {descriptor.capture_id}") + + +class PackWriter: + def __init__( + self, + *, + pack_id: UUID | str, + created_at_ns: int, + max_pack_bytes: int, + max_records: int = MAX_RECORDS, + ) -> None: + try: + parsed_id = pack_id if isinstance(pack_id, UUID) else UUID(pack_id) + except (TypeError, ValueError, AttributeError) as exc: + raise ValueError("pack_id must be a UUID") from exc + if created_at_ns < 0: + raise ValueError("created_at_ns must be non-negative") + if max_pack_bytes < _HEADER.size + _TRAILER.size + 2: + raise ValueError("max_pack_bytes is too small for a pack") + if not 1 <= max_records <= MAX_RECORDS: + raise ValueError(f"max_records must be between 1 and {MAX_RECORDS}") + + self._uuid = parsed_id + self._pack_id = str(parsed_id) + self._created_at_ns = created_at_ns + self._max_pack_bytes = max_pack_bytes + self._max_records = max_records + self._buffer = bytearray( + _HEADER.pack( + _HEADER_MAGIC, + PACK_MAJOR_VERSION, + PACK_MINOR_VERSION, + _HEADER.size, + parsed_id.bytes, + created_at_ns, + 0, + b"\0" * 20, + ) + ) + self._records: list[_IndexedRecord] = [] + empty_footer = _encode_json( + { + "format": "dmi-pack", + "major_version": PACK_MAJOR_VERSION, + "minor_version": PACK_MINOR_VERSION, + "pack_id": self._pack_id, + "created_at_ns": self._created_at_ns, + "records": [], + } + ) + self._footer_prefix, self._footer_suffix = empty_footer.rsplit(b"[]", 1) + self._record_json: list[bytes] = [] + self._record_json_bytes = 0 + self._capture_ids: set[str] = set() + self._sealed = False + + @property + def record_count(self) -> int: + return len(self._records) + + def append(self, record: CaptureRecord) -> None: + if self._sealed: + raise _PackStateError("pack is already sealed") + if len(self._records) >= self._max_records: + raise PackCapacityError("pack record limit reached") + if record.metadata.capture_id in self._capture_ids: + raise DuplicateCaptureError( + f"duplicate capture ID: {record.metadata.capture_id}" + ) + offset = _aligned_length(len(self._buffer)) + indexed = _IndexedRecord( + metadata=record.metadata, + offset=offset, + stored_length=len(record.payload), + decoded_length=len(record.payload), + codec="none", + checksum=_crc32(record.payload), + ) + encoded_record = _encode_json(_record_mapping(indexed)) + record_content_bytes = self._record_json_bytes + len(encoded_record) + record_separators = len(self._record_json) + footer_length = ( + len(self._footer_prefix) + + 2 + + record_content_bytes + + record_separators + + len(self._footer_suffix) + ) + projected = ( + _aligned_length(offset + len(record.payload)) + + footer_length + + _TRAILER.size + ) + if footer_length > MAX_FOOTER_BYTES or projected > self._max_pack_bytes: + raise PackCapacityError("record would exceed max_pack_bytes") + + _align(self._buffer) + self._buffer.extend(record.payload) + self._records.append(indexed) + self._record_json.append(encoded_record) + self._record_json_bytes += len(encoded_record) + self._capture_ids.add(record.metadata.capture_id) + + def seal(self) -> SealedPack: + if self._sealed: + raise _PackStateError("pack is already sealed") + if not self._records: + raise ValueError("cannot seal an empty pack") + buffer = bytearray(self._buffer) + _align(buffer) + footer_offset = len(buffer) + footer = ( + self._footer_prefix + + b"[" + + b",".join(self._record_json) + + b"]" + + self._footer_suffix + ) + if len(footer) > MAX_FOOTER_BYTES: + raise ValueError("pack footer exceeds its size limit") + buffer.extend(footer) + body_checksum = sha256(buffer).digest() + buffer.extend( + _TRAILER.pack( + _TRAILER_MAGIC, + PACK_MAJOR_VERSION, + PACK_MINOR_VERSION, + footer_offset, + len(footer), + zlib.crc32(footer) & 0xFFFFFFFF, + body_checksum, + ) + ) + if len(buffer) > self._max_pack_bytes: + raise ValueError("sealed pack exceeds max_pack_bytes") + self._sealed = True + self._buffer = buffer + data = bytes(buffer) + return SealedPack( + pack_id=self._pack_id, + created_at_ns=self._created_at_ns, + data=data, + record_count=len(self._records), + footer_offset=footer_offset, + checksum=sha256(data).hexdigest(), + ) + + +class PackIndex: + def __init__(self, ref: PackRef, records: Iterable[_IndexedRecord]) -> None: + self.ref = ref + self.pack_id = ref.pack_id + self._records = tuple(records) + + @staticmethod + def trailer_size() -> int: + return _TRAILER.size + + @classmethod + def from_store(cls, store: PackStore, ref: PackRef) -> PackIndex: + if ref.store_id != store.store_id: + raise ValueError("pack reference belongs to another store") + if ref.object_bytes < _HEADER.size + _TRAILER.size + 2: + raise PackFormatError("pack is truncated") + trailer_offset = ref.object_bytes - _TRAILER.size + trailer_bytes = store.read_range(ref, trailer_offset, _TRAILER.size) + try: + trailer = _TRAILER.unpack(trailer_bytes) + except struct.error as exc: + raise PackFormatError("pack trailer is truncated") from exc + magic, major, minor, footer_offset, footer_length, footer_crc, _ = trailer + if magic != _TRAILER_MAGIC: + raise PackFormatError("pack has an invalid trailer") + if major != PACK_MAJOR_VERSION or minor > PACK_MINOR_VERSION: + raise PackFormatError(f"unsupported pack version: {major}.{minor}") + if footer_length > MAX_FOOTER_BYTES: + raise PackFormatError("pack footer exceeds its size limit") + if footer_offset < _HEADER.size or footer_offset + footer_length != trailer_offset: + raise PackFormatError("pack footer range is invalid") + footer = store.read_range(ref, footer_offset, footer_length) + if zlib.crc32(footer) & 0xFFFFFFFF != footer_crc: + raise PackIntegrityError("footer checksum mismatch") + decoded = _decode_footer(footer, major=major, minor=minor) + try: + pack_id = str(UUID(str(decoded.get("pack_id")))) + except (ValueError, TypeError, AttributeError) as exc: + raise PackFormatError("pack footer has an invalid pack ID") from exc + if pack_id != ref.pack_id: + raise PackFormatError("pack footer identity does not match its object key") + records = _parse_records(decoded, footer_offset) + if len(records) != ref.record_count: + raise PackFormatError("pack record count does not match its object metadata") + return cls(ref, records) + + def descriptors(self) -> tuple[CaptureDescriptor, ...]: + return _descriptors(self.ref, self._records) + + +class PackReader: + def __init__( + self, + data: bytes, + *, + pack_id: str, + created_at_ns: int, + records: Iterable[_IndexedRecord], + object_checksum: str, + ) -> None: + self._data = data + self.pack_id = pack_id + self.created_at_ns = created_at_ns + self._records = tuple(records) + self._by_capture_id = {item.metadata.capture_id: item for item in self._records} + self.object_checksum = object_checksum + + @classmethod + def from_bytes(cls, value: bytes | bytearray | memoryview) -> PackReader: + data = bytes(value) + minimum = _HEADER.size + _TRAILER.size + 2 + if len(data) < minimum: + raise PackFormatError("pack is truncated") + + try: + header = _HEADER.unpack_from(data) + except struct.error as exc: + raise PackFormatError("pack header is truncated") from exc + magic, major, minor, header_size, pack_bytes, created_at_ns, flags, reserved = header + if magic != _HEADER_MAGIC: + raise PackFormatError("invalid pack magic") + if major != PACK_MAJOR_VERSION: + raise PackFormatError(f"unsupported pack major version: {major}") + if minor > PACK_MINOR_VERSION: + raise PackFormatError(f"unsupported pack minor version: {minor}") + if header_size != _HEADER.size or flags != 0 or reserved != b"\0" * 20: + raise PackFormatError("invalid pack header fields") + + trailer_offset = len(data) - _TRAILER.size + try: + trailer = _TRAILER.unpack_from(data, trailer_offset) + except struct.error as exc: + raise PackFormatError("pack trailer is truncated") from exc + ( + trailer_magic, + trailer_major, + trailer_minor, + footer_offset, + footer_length, + footer_crc, + body_hash, + ) = trailer + if trailer_magic != _TRAILER_MAGIC: + raise PackFormatError("pack is truncated or has an invalid trailer") + if (trailer_major, trailer_minor) != (major, minor): + raise PackFormatError("pack header and trailer versions differ") + if footer_length > MAX_FOOTER_BYTES: + raise PackFormatError("pack footer exceeds its size limit") + if footer_offset < _HEADER.size or footer_offset + footer_length != trailer_offset: + raise PackFormatError("pack footer range is invalid") + if sha256(data[:trailer_offset]).digest() != body_hash: + raise PackIntegrityError("pack checksum mismatch") + footer = data[footer_offset:trailer_offset] + if zlib.crc32(footer) & 0xFFFFFFFF != footer_crc: + raise PackIntegrityError("footer checksum mismatch") + + decoded = _decode_footer(footer, major=major, minor=minor) + pack_id = str(UUID(bytes=pack_bytes)) + if decoded.get("pack_id") != pack_id or decoded.get("created_at_ns") != created_at_ns: + raise PackFormatError("pack footer identity does not match the header") + records = _parse_records(decoded, footer_offset) + + return cls( + data, + pack_id=pack_id, + created_at_ns=created_at_ns, + records=records, + object_checksum=sha256(data).hexdigest(), + ) + + @staticmethod + def _parse_record(raw: object, footer_offset: int) -> _IndexedRecord: + if not isinstance(raw, dict): + raise PackFormatError("pack record entry must be an object") + try: + metadata_raw = raw["metadata"] + offset = raw["offset"] + stored_length = raw["stored_length"] + decoded_length = raw["decoded_length"] + codec = raw["codec"] + checksum = raw["checksum"] + except KeyError as exc: + raise PackFormatError(f"pack record is missing {exc.args[0]}") from exc + if not isinstance(metadata_raw, dict): + raise PackFormatError("pack record metadata must be an object") + if any( + not isinstance(value, int) or isinstance(value, bool) + for value in (offset, stored_length, decoded_length) + ): + raise PackFormatError("pack record offsets and lengths must be integers") + if offset < _HEADER.size or stored_length < 0 or decoded_length < 0: + raise PackFormatError("pack record range is invalid") + if offset + stored_length > footer_offset: + raise PackFormatError("pack record extends into the footer") + if codec != "none" or stored_length != decoded_length: + raise PackFormatError("dmi-pack-v1 supports only uncompressed records") + if not isinstance(checksum, str) or _CRC_PATTERN.fullmatch(checksum) is None: + raise PackFormatError("pack record checksum is invalid") + metadata = CaptureMetadata.from_mapping(metadata_raw) + if metadata.logical_bytes != decoded_length: + raise PackFormatError("record length does not match metadata dtype and shape") + return _IndexedRecord( + metadata=metadata, + offset=offset, + stored_length=stored_length, + decoded_length=decoded_length, + codec=codec, + checksum=checksum, + ) + + def descriptors(self, *, store_id: str, object_key: str) -> tuple[CaptureDescriptor, ...]: + return _descriptors( + PackRef( + pack_id=self.pack_id, + store_id=store_id, + object_key=object_key, + object_bytes=len(self._data), + checksum=self.object_checksum, + record_count=len(self._records), + ), + self._records, + ) + + def read_payload(self, descriptor: CaptureDescriptor) -> bytes: + if descriptor.locator.pack_id != self.pack_id: + raise PackFormatError("descriptor belongs to another pack") + indexed = self._by_capture_id.get(descriptor.capture_id) + if indexed is None: + raise PackFormatError("descriptor is not present in this pack") + expected = ( + indexed.offset, + indexed.stored_length, + indexed.decoded_length, + indexed.codec, + indexed.checksum, + ) + actual = ( + descriptor.locator.offset, + descriptor.locator.stored_length, + descriptor.locator.decoded_length, + descriptor.locator.codec, + descriptor.locator.checksum, + ) + if actual != expected: + raise PackFormatError("descriptor does not match the pack footer") + payload = self._data[indexed.offset:indexed.offset + indexed.stored_length] + verify_payload(descriptor, payload) + return payload + + +def _decode_footer(footer: bytes, *, major: int, minor: int) -> dict[str, object]: + try: + decoded = json.loads(footer) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise PackFormatError("pack footer is not valid JSON") from exc + if not isinstance(decoded, dict) or decoded.get("format") != "dmi-pack": + raise PackFormatError("pack footer has an invalid format marker") + if decoded.get("major_version") != major or decoded.get("minor_version") != minor: + raise PackFormatError("pack footer version does not match the trailer") + return decoded + + +def _parse_records(decoded: dict[str, object], footer_offset: int) -> tuple[_IndexedRecord, ...]: + raw_records = decoded.get("records") + if not isinstance(raw_records, list) or len(raw_records) > MAX_RECORDS: + raise PackFormatError("pack footer has an invalid record list") + records: list[_IndexedRecord] = [] + previous_end = _HEADER.size + seen: set[str] = set() + for raw in raw_records: + record = PackReader._parse_record(raw, footer_offset) + if record.metadata.capture_id in seen: + raise PackFormatError(f"duplicate capture ID: {record.metadata.capture_id}") + if record.offset < previous_end: + raise PackFormatError("pack record ranges overlap or are out of order") + previous_end = record.offset + record.stored_length + seen.add(record.metadata.capture_id) + records.append(record) + return tuple(records) + + +def _descriptors( + ref: PackRef, records: Iterable[_IndexedRecord] +) -> tuple[CaptureDescriptor, ...]: + records = tuple(records) + return tuple( + CaptureDescriptor( + metadata=item.metadata, + locator=PayloadLocator( + pack_id=ref.pack_id, + store_id=ref.store_id, + object_key=ref.object_key, + object_bytes=ref.object_bytes, + pack_checksum=ref.checksum, + pack_record_count=ref.record_count, + offset=item.offset, + stored_length=item.stored_length, + decoded_length=item.decoded_length, + codec=item.codec, + checksum=item.checksum, + ), + ) + for item in records + ) diff --git a/src/dmi/storage/capture/pipeline.py b/src/dmi/storage/capture/pipeline.py new file mode 100644 index 000000000..8226c75ba --- /dev/null +++ b/src/dmi/storage/capture/pipeline.py @@ -0,0 +1,586 @@ +from __future__ import annotations + +from collections import deque +from dataclasses import dataclass +from datetime import datetime, timezone +from enum import Enum +from hashlib import sha256 +import math +import threading +import time +from typing import Callable, Mapping, Protocol +from urllib.parse import quote +from uuid import UUID, uuid4 + +from .model import CaptureMetadata, CaptureRecord, PackRef, PackSource, PackStore +from .pack import PackCapacityError, PackWriter, SealedPack + + +class AdmissionResult(str, Enum): + ACCEPTED = "accepted" + DROPPED = "dropped" + TIMED_OUT = "timed_out" + TOO_LARGE = "too_large" + CLOSED = "closed" + + +class OverloadPolicy(str, Enum): + BLOCK = "block" + DROP_NEWEST = "drop_newest" + + +class FlushReason(str, Enum): + SIZE = "size" + RECORDS = "records" + LINGER = "linger" + SESSION = "session" + SHUTDOWN = "shutdown" + + +class OversizedRecordError(ValueError): + pass + + +class PipelineFailedError(RuntimeError): + pass + + +@dataclass(frozen=True, slots=True) +class QueueSnapshot: + records: int + bytes: int + peak_records: int + peak_bytes: int + closed: bool + + +class BoundedRecordQueue: + def __init__(self, *, max_records: int, max_bytes: int) -> None: + if ( + type(max_records) is not int + or type(max_bytes) is not int + or max_records <= 0 + or max_bytes <= 0 + ): + raise ValueError("queue limits must be positive") + self._max_records = max_records + self._max_bytes = max_bytes + self._items: deque[CaptureRecord] = deque() + self._bytes = 0 + self._peak_records = 0 + self._peak_bytes = 0 + self._closed = False + self._condition = threading.Condition() + + def put( + self, + record: CaptureRecord, + *, + policy: OverloadPolicy, + timeout: float | None = None, + ) -> AdmissionResult: + if timeout is not None and ( + isinstance(timeout, bool) or not math.isfinite(timeout) or timeout < 0 + ): + raise ValueError("timeout must be non-negative") + if not isinstance(policy, OverloadPolicy): + raise ValueError("unknown overload policy") + record_bytes = len(record.payload) + if record_bytes > self._max_bytes: + return AdmissionResult.TOO_LARGE + deadline = None if timeout is None else time.monotonic() + timeout + with self._condition: + while True: + if self._closed: + return AdmissionResult.CLOSED + if self._fits(record_bytes): + self._items.append(record) + self._bytes += record_bytes + self._peak_records = max(self._peak_records, len(self._items)) + self._peak_bytes = max(self._peak_bytes, self._bytes) + self._condition.notify() + return AdmissionResult.ACCEPTED + if policy is OverloadPolicy.DROP_NEWEST: + return AdmissionResult.DROPPED + remaining = None if deadline is None else deadline - time.monotonic() + if remaining is not None and remaining <= 0: + return AdmissionResult.TIMED_OUT + self._condition.wait(remaining) + + def get(self, timeout: float | None = None) -> CaptureRecord | None: + if timeout is not None and ( + isinstance(timeout, bool) or not math.isfinite(timeout) or timeout < 0 + ): + raise ValueError("timeout must be non-negative") + deadline = None if timeout is None else time.monotonic() + timeout + with self._condition: + while not self._items: + if self._closed: + return None + remaining = None if deadline is None else deadline - time.monotonic() + if remaining is not None and remaining <= 0: + return None + self._condition.wait(remaining) + record = self._items.popleft() + self._bytes -= len(record.payload) + self._condition.notify_all() + return record + + def close(self) -> None: + with self._condition: + self._closed = True + self._condition.notify_all() + + def snapshot(self) -> QueueSnapshot: + with self._condition: + return QueueSnapshot( + records=len(self._items), + bytes=self._bytes, + peak_records=self._peak_records, + peak_bytes=self._peak_bytes, + closed=self._closed, + ) + + def _fits(self, record_bytes: int) -> bool: + return ( + len(self._items) < self._max_records + and self._bytes + record_bytes <= self._max_bytes + ) + + +@dataclass(frozen=True, slots=True) +class ReadyPack: + pack: SealedPack + first_metadata: CaptureMetadata + reason: FlushReason + + +class PackAssembler: + def __init__( + self, + *, + max_pack_bytes: int, + max_records: int, + max_linger_ns: int, + pack_id_factory: Callable[[], UUID] = uuid4, + ) -> None: + limits = (max_pack_bytes, max_records, max_linger_ns) + if any(type(value) is not int or value <= 0 for value in limits): + raise ValueError("pack limits and linger must be positive") + self._max_pack_bytes = max_pack_bytes + self._max_records = max_records + self._max_linger_ns = max_linger_ns + self._pack_id_factory = pack_id_factory + self._writer: PackWriter | None = None + self._first_metadata: CaptureMetadata | None = None + self._scope: tuple[str, str, int] | None = None + self._opened_ns: int | None = None + + def append(self, record: CaptureRecord, *, now_ns: int) -> tuple[ReadyPack, ...]: + if now_ns < 0: + raise ValueError("now_ns must be non-negative") + emitted: list[ReadyPack] = [] + scope = self._record_scope(record) + if self._writer is None: + writer = self._writer_with(record) + self._adopt(writer, record.metadata, scope=scope, now_ns=now_ns) + elif scope != self._scope: + writer = self._writer_with(record) + emitted.extend(self.flush(FlushReason.SESSION)) + self._adopt(writer, record.metadata, scope=scope, now_ns=now_ns) + else: + try: + self._writer.append(record) + except PackCapacityError: + writer = self._writer_with(record) + reason = ( + FlushReason.RECORDS + if self._writer.record_count >= self._max_records + else FlushReason.SIZE + ) + emitted.extend(self.flush(reason)) + self._adopt(writer, record.metadata, scope=scope, now_ns=now_ns) + assert self._writer is not None + if self._writer.record_count >= self._max_records: + emitted.extend(self.flush(FlushReason.RECORDS)) + return tuple(emitted) + + def _writer_with(self, record: CaptureRecord) -> PackWriter: + writer = PackWriter( + pack_id=self._pack_id_factory(), + created_at_ns=record.metadata.captured_at_ns, + max_pack_bytes=self._max_pack_bytes, + max_records=self._max_records, + ) + try: + writer.append(record) + except PackCapacityError as exc: + raise OversizedRecordError( + f"capture {record.metadata.capture_id} does not fit an empty pack" + ) from exc + return writer + + def _adopt( + self, + writer: PackWriter, + metadata: CaptureMetadata, + *, + scope: tuple[str, str, int], + now_ns: int, + ) -> None: + self._writer = writer + self._first_metadata = metadata + self._scope = scope + self._opened_ns = now_ns + + def flush_expired(self, *, now_ns: int) -> tuple[ReadyPack, ...]: + if self._opened_ns is None or now_ns - self._opened_ns < self._max_linger_ns: + return () + return self.flush(FlushReason.LINGER) + + def seconds_until_expiry(self, *, now_ns: int) -> float | None: + if self._opened_ns is None: + return None + remaining = self._max_linger_ns - (now_ns - self._opened_ns) + return max(0, remaining) / 1_000_000_000 + + def flush(self, reason: FlushReason) -> tuple[ReadyPack, ...]: + if self._writer is None: + return () + assert self._first_metadata is not None + ready = ReadyPack(self._writer.seal(), self._first_metadata, reason) + self._reset() + return (ready,) + + def _reset(self) -> None: + self._writer = None + self._first_metadata = None + self._scope = None + self._opened_ns = None + + @staticmethod + def _record_scope(record: CaptureRecord) -> tuple[str, str, int]: + metadata = record.metadata + return metadata.tenant_id, metadata.session_id, metadata.producer_rank + + +def object_key_for(ready: ReadyPack) -> str: + metadata = ready.first_metadata + captured = datetime.fromtimestamp( + metadata.captured_at_ns / 1_000_000_000, tz=timezone.utc + ) + return ( + f"v1/tenant={_key_component(metadata.tenant_id)}/" + f"date={captured:%Y-%m-%d}/" + f"session={_key_component(metadata.session_id)}/" + f"rank={metadata.producer_rank}/" + f"{ready.pack.pack_id}.dmi-pack" + ) + + +def _key_component(value: str) -> str: + encoded = quote(value, safe="-_.=") + if len(encoded.encode()) <= 160: + return encoded + return "sha256-" + sha256(value.encode()).hexdigest() + + +class PackSink(Protocol): + def persist(self, ready: ReadyPack) -> PackRef | PackSource: ... + + +class DirectPackSink: + def __init__(self, store: PackStore) -> None: + self._store = store + self.last_ref: PackRef | None = None + + def persist(self, ready: ReadyPack) -> PackRef: + ref = self._store.put(ready.pack, object_key_for(ready)) + self.last_ref = ref + return ref + + +@dataclass(frozen=True, slots=True) +class PipelineConfig: + max_queue_records: int + max_queue_bytes: int + max_pack_bytes: int + max_pack_records: int + max_linger_ns: int + overload_policy: OverloadPolicy = OverloadPolicy.DROP_NEWEST + admission_timeout: float | None = None + + def __post_init__(self) -> None: + positive = ( + "max_queue_records", + "max_queue_bytes", + "max_pack_bytes", + "max_pack_records", + "max_linger_ns", + ) + if any( + type(getattr(self, name)) is not int or getattr(self, name) <= 0 + for name in positive + ): + raise ValueError("pipeline bounds must be positive") + if not isinstance(self.overload_policy, OverloadPolicy): + raise ValueError("unknown overload policy") + if self.admission_timeout is not None and ( + isinstance(self.admission_timeout, bool) + or not math.isfinite(self.admission_timeout) + or self.admission_timeout < 0 + ): + raise ValueError("admission_timeout must be non-negative") + + +@dataclass(frozen=True, slots=True) +class HistogramSnapshot: + bounds_ns: tuple[int, ...] + counts: tuple[int, ...] + count: int + total_ns: int + max_ns: int + + +class _Histogram: + _BOUNDS_NS = ( + 10_000, + 100_000, + 1_000_000, + 10_000_000, + 100_000_000, + 1_000_000_000, + ) + + def __init__(self) -> None: + self._counts = [0] * (len(self._BOUNDS_NS) + 1) + self._count = 0 + self._total_ns = 0 + self._max_ns = 0 + + def observe(self, duration_ns: int) -> None: + bucket = next( + ( + index + for index, bound in enumerate(self._BOUNDS_NS) + if duration_ns <= bound + ), + len(self._BOUNDS_NS), + ) + self._counts[bucket] += 1 + self._count += 1 + self._total_ns += duration_ns + self._max_ns = max(self._max_ns, duration_ns) + + def snapshot(self) -> HistogramSnapshot: + return HistogramSnapshot( + bounds_ns=self._BOUNDS_NS, + counts=tuple(self._counts), + count=self._count, + total_ns=self._total_ns, + max_ns=self._max_ns, + ) + + +@dataclass(frozen=True, slots=True) +class PipelineSnapshot: + submitted_records: int + admitted_records: int + admitted_bytes: int + dropped_records: int + timed_out_records: int + oversized_records: int + rejected_closed_records: int + persisted_records: int + packs_persisted: int + packed_bytes: int + flush_size: int + flush_records: int + flush_linger: int + flush_session: int + flush_shutdown: int + failures: int + event_callback_failures: int + queue_records: int + queue_bytes: int + queue_peak_records: int + queue_peak_bytes: int + admission_duration: HistogramSnapshot + persist_duration: HistogramSnapshot + + +@dataclass(frozen=True, slots=True) +class PipelineEvent: + event: str + fields: Mapping[str, int | str] + + +class HostCapturePipeline: + def __init__( + self, + config: PipelineConfig, + sink: PackSink, + *, + pack_id_factory: Callable[[], UUID] = uuid4, + clock_ns: Callable[[], int] = time.monotonic_ns, + event_callback: Callable[[PipelineEvent], None] | None = None, + ) -> None: + self._config = config + self._sink = sink + self._clock_ns = clock_ns + self._event_callback = event_callback + self._queue = BoundedRecordQueue( + max_records=config.max_queue_records, + max_bytes=config.max_queue_bytes, + ) + self._assembler = PackAssembler( + max_pack_bytes=config.max_pack_bytes, + max_records=config.max_pack_records, + max_linger_ns=config.max_linger_ns, + pack_id_factory=pack_id_factory, + ) + self._lock = threading.Lock() + self._thread: threading.Thread | None = None + self._error: BaseException | None = None + self._counters = { + "submitted_records": 0, + "admitted_records": 0, + "admitted_bytes": 0, + "dropped_records": 0, + "timed_out_records": 0, + "oversized_records": 0, + "rejected_closed_records": 0, + "persisted_records": 0, + "packs_persisted": 0, + "packed_bytes": 0, + "flush_size": 0, + "flush_records": 0, + "flush_linger": 0, + "flush_session": 0, + "flush_shutdown": 0, + "failures": 0, + "event_callback_failures": 0, + } + self._admission_duration = _Histogram() + self._persist_duration = _Histogram() + + def start(self) -> None: + with self._lock: + if self._thread is not None: + raise RuntimeError("pipeline has already been started") + self._thread = threading.Thread( + target=self._run, name="dmi-capture-persistence", daemon=True + ) + self._thread.start() + + def submit(self, record: CaptureRecord) -> AdmissionResult: + with self._lock: + if self._thread is None: + raise RuntimeError("pipeline is not started") + error = self._error + self._counters["submitted_records"] += 1 + if error is not None: + raise PipelineFailedError("capture pipeline failed") from error + started = self._clock_ns() + result = self._queue.put( + record, + policy=self._config.overload_policy, + timeout=self._config.admission_timeout, + ) + duration = self._clock_ns() - started + with self._lock: + self._admission_duration.observe(max(0, duration)) + if result is AdmissionResult.ACCEPTED: + self._counters["admitted_records"] += 1 + self._counters["admitted_bytes"] += len(record.payload) + elif result is AdmissionResult.DROPPED: + self._counters["dropped_records"] += 1 + elif result is AdmissionResult.TIMED_OUT: + self._counters["timed_out_records"] += 1 + elif result is AdmissionResult.TOO_LARGE: + self._counters["oversized_records"] += 1 + elif result is AdmissionResult.CLOSED: + self._counters["rejected_closed_records"] += 1 + return result + + def close(self, *, timeout: float | None = None) -> PipelineSnapshot: + self._queue.close() + with self._lock: + thread = self._thread + if thread is None: + raise RuntimeError("pipeline is not started") + thread.join(timeout) + if thread.is_alive(): + raise TimeoutError("capture pipeline did not stop before timeout") + with self._lock: + error = self._error + if error is not None: + raise PipelineFailedError("capture pipeline failed") from error + return self.snapshot() + + def snapshot(self) -> PipelineSnapshot: + queue = self._queue.snapshot() + with self._lock: + return PipelineSnapshot( + **self._counters, + queue_records=queue.records, + queue_bytes=queue.bytes, + queue_peak_records=queue.peak_records, + queue_peak_bytes=queue.peak_bytes, + admission_duration=self._admission_duration.snapshot(), + persist_duration=self._persist_duration.snapshot(), + ) + + def _run(self) -> None: + try: + while True: + now_ns = self._clock_ns() + timeout = self._assembler.seconds_until_expiry(now_ns=now_ns) + record = self._queue.get(timeout=timeout) + if record is None: + queue = self._queue.snapshot() + if queue.closed and queue.records == 0: + self._persist(self._assembler.flush(FlushReason.SHUTDOWN)) + return + self._persist( + self._assembler.flush_expired(now_ns=self._clock_ns()) + ) + continue + self._persist( + self._assembler.append(record, now_ns=self._clock_ns()) + ) + except BaseException as exc: + with self._lock: + self._error = exc + self._counters["failures"] += 1 + self._queue.close() + self._emit("pipeline_failed", error_type=type(exc).__name__) + + def _persist(self, packs: tuple[ReadyPack, ...]) -> None: + for ready in packs: + started = self._clock_ns() + self._sink.persist(ready) + duration = max(0, self._clock_ns() - started) + reason_key = f"flush_{ready.reason.value}" + with self._lock: + self._persist_duration.observe(duration) + self._counters["persisted_records"] += ready.pack.record_count + self._counters["packs_persisted"] += 1 + self._counters["packed_bytes"] += len(ready.pack.data) + self._counters[reason_key] += 1 + self._emit( + "pack_persisted", + pack_id=ready.pack.pack_id, + reason=ready.reason.value, + records=ready.pack.record_count, + bytes=len(ready.pack.data), + ) + + def _emit(self, event: str, **fields: int | str) -> None: + if self._event_callback is None: + return + try: + self._event_callback(PipelineEvent(event=event, fields=fields)) + except Exception: + with self._lock: + self._counters["event_callback_failures"] += 1 diff --git a/src/dmi/storage/capture/reader.py b/src/dmi/storage/capture/reader.py new file mode 100644 index 000000000..028acc1c8 --- /dev/null +++ b/src/dmi/storage/capture/reader.py @@ -0,0 +1,198 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Mapping, Sequence + +from .model import ( + CaptureCatalog, + CaptureDescriptor, + CapturePage, + CaptureQuery, + CaptureSelection, + DuplicateCaptureError, + HydratedCapture, + HydrationEstimate, + HydrationLimitError, + PackFormatError, + PackIntegrityError, + PackStore, +) +from .pack import verify_payload + + +@dataclass(frozen=True, slots=True) +class _ReadRange: + descriptor_indexes: tuple[int, ...] + offset: int + length: int + + +@dataclass(frozen=True, slots=True) +class _ObjectPlan: + store_id: str + pack_id: str + object_key: str + descriptor_indexes: tuple[int, ...] + ranges: tuple[_ReadRange, ...] + + +class CaptureReader: + def __init__( + self, + catalog: CaptureCatalog, + stores: Mapping[str, PackStore], + *, + max_coalesce_gap_bytes: int = 4096, + ) -> None: + if max_coalesce_gap_bytes < 0: + raise ValueError("max_coalesce_gap_bytes must be non-negative") + if not stores: + raise ValueError("at least one pack store is required") + for store_id, store in stores.items(): + if store_id != store.store_id: + raise ValueError(f"store mapping key does not match {store.store_id!r}") + self._catalog = catalog + self._stores = dict(stores) + self._max_coalesce_gap_bytes = max_coalesce_gap_bytes + + def search(self, query: CaptureQuery) -> CapturePage: + return self._catalog.search(query) + + def select(self, query: CaptureQuery) -> CaptureSelection: + page = self.search(query) + if page.next_cursor is not None: + raise ValueError( + "selection exceeds one bounded page; narrow the query or paginate explicitly" + ) + return CaptureSelection.create( + page.items, + catalog_watermark=page.watermark, + query_hash=query.query_hash, + ) + + def estimate(self, selection: CaptureSelection) -> HydrationEstimate: + descriptors = self._resolve(selection) + plans = self._plan(descriptors) + return HydrationEstimate( + capture_count=len(descriptors), + object_count=len(plans), + request_count=sum(len(plan.ranges) for plan in plans), + logical_bytes=sum(item.locator.decoded_length for item in descriptors), + stored_bytes=sum(item.locator.stored_length for item in descriptors), + request_bytes=sum(item.length for plan in plans for item in plan.ranges), + ) + + def hydrate( + self, + selection: CaptureSelection, + *, + byte_limit: int, + request_limit: int = 1024, + ) -> tuple[HydratedCapture, ...]: + if byte_limit < 0: + raise ValueError("byte_limit must be non-negative") + if request_limit <= 0: + raise ValueError("request_limit must be positive") + descriptors = self._resolve(selection) + plans = self._plan(descriptors) + request_count = sum(len(plan.ranges) for plan in plans) + if request_count > request_limit: + raise HydrationLimitError( + f"hydration request limit exceeded: {request_count} > {request_limit}" + ) + request_bytes = sum(item.length for plan in plans for item in plan.ranges) + if request_bytes > byte_limit: + raise HydrationLimitError( + f"hydration byte limit exceeded: {request_bytes} > {byte_limit}" + ) + + payloads: list[bytes | None] = [None] * len(descriptors) + for plan in plans: + store = self._stores.get(plan.store_id) + if store is None: + raise PackFormatError(f"unknown pack store: {plan.store_id}") + ref = descriptors[plan.descriptor_indexes[0]].locator.pack_ref + for read_range in plan.ranges: + block = store.read_range(ref, read_range.offset, read_range.length) + if len(block) != read_range.length: + raise PackIntegrityError( + f"object store returned a short range: " + f"{len(block)} != {read_range.length}" + ) + for index in read_range.descriptor_indexes: + descriptor = descriptors[index] + start = descriptor.locator.offset - read_range.offset + end = start + descriptor.locator.stored_length + payload = block[start:end] + verify_payload(descriptor, payload) + payloads[index] = bytes(payload) + + if any(payload is None for payload in payloads): + raise PackFormatError("hydration plan did not resolve every capture") + return tuple( + HydratedCapture(descriptor=descriptor, payload=payload) + for descriptor, payload in zip(descriptors, payloads) + if payload is not None + ) + + def _resolve(self, selection: CaptureSelection) -> tuple[CaptureDescriptor, ...]: + resolved = self._catalog.get_by_ids( + selection.capture_ids, + watermark=selection.catalog_watermark, + ) + by_id: dict[str, CaptureDescriptor] = {} + for descriptor in resolved: + if descriptor.capture_id in by_id: + raise DuplicateCaptureError( + f"catalog returned duplicate capture: {descriptor.capture_id}" + ) + by_id[descriptor.capture_id] = descriptor + if set(by_id) != set(selection.capture_ids): + raise PackFormatError("selection no longer resolves at its catalog watermark") + return tuple(by_id[capture_id] for capture_id in selection.capture_ids) + + def _plan(self, descriptors: Sequence[CaptureDescriptor]) -> tuple[_ObjectPlan, ...]: + grouped: dict[tuple[str, str, str], list[int]] = {} + for index, descriptor in enumerate(descriptors): + locator = descriptor.locator + grouped.setdefault( + (locator.store_id, locator.pack_id, locator.object_key), [] + ).append(index) + + plans: list[_ObjectPlan] = [] + for (store_id, pack_id, object_key), indexes in grouped.items(): + indexes.sort(key=lambda index: descriptors[index].locator.offset) + ranges: list[_ReadRange] = [] + current: list[int] = [] + start = end = 0 + for index in indexes: + locator = descriptors[index].locator + if not current: + current = [index] + start = locator.offset + end = locator.offset + locator.stored_length + continue + if locator.offset <= end + self._max_coalesce_gap_bytes: + current.append(index) + end = max(end, locator.offset + locator.stored_length) + continue + ranges.append( + _ReadRange(tuple(current), offset=start, length=end - start) + ) + current = [index] + start = locator.offset + end = locator.offset + locator.stored_length + if current: + ranges.append( + _ReadRange(tuple(current), offset=start, length=end - start) + ) + plans.append( + _ObjectPlan( + store_id=store_id, + pack_id=pack_id, + object_key=object_key, + descriptor_indexes=tuple(indexes), + ranges=tuple(ranges), + ) + ) + return tuple(plans) diff --git a/src/dmi/storage/capture/s3.py b/src/dmi/storage/capture/s3.py new file mode 100644 index 000000000..9629fc383 --- /dev/null +++ b/src/dmi/storage/capture/s3.py @@ -0,0 +1,421 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import PurePosixPath +import re +from typing import Mapping, Protocol +from urllib.parse import urlsplit +from uuid import UUID + +from .filesystem import validate_object_key, validate_pack_source +from .model import ( + ObjectInfo, + ObjectPage, + PackConflictError, + PackFormatError, + PackIntegrityError, + PackRef, + PackSource, + StoredObject, +) + + +_MIN_MULTIPART_BYTES = 5 * 1024**2 +_MAX_CURSOR_BYTES = 2048 +_SHA256 = re.compile(r"[0-9a-f]{64}\Z") + + +def _validate_bounded_text(name: str, value: object, limit: int = 255) -> None: + if ( + not isinstance(value, str) + or not value + or len(value.encode()) > limit + ): + raise ValueError(f"{name} must be non-empty and at most {limit} bytes") + + +def _validate_credentials(config: S3StoreConfig) -> None: + if (config.access_key_id is None) != (config.secret_access_key is None): + raise ValueError("access_key_id and secret_access_key must be set together") + for value in ( + config.access_key_id, + config.secret_access_key, + config.session_token, + ): + if value is not None: + _validate_bounded_text("credentials", value, 4096) + + +def _validate_endpoint(config: S3StoreConfig) -> None: + if config.endpoint_url is None: + return + parsed = urlsplit(config.endpoint_url) + if ( + parsed.scheme not in {"http", "https"} + or not parsed.hostname + or parsed.username is not None + or parsed.password is not None + or parsed.path not in {"", "/"} + or parsed.query + or parsed.fragment + ): + raise ValueError("endpoint_url must be an HTTP(S) origin") + if parsed.scheme == "http" and not config.allow_insecure_http: + raise ValueError("HTTP endpoints require allow_insecure_http=True") + + +class S3Client(Protocol): + def head_object(self, **kwargs): ... + def upload_fileobj(self, *args, **kwargs) -> None: ... + def get_object(self, **kwargs): ... + def list_objects_v2(self, **kwargs): ... + + +@dataclass(frozen=True, slots=True) +class S3StoreConfig: + endpoint_url: str | None + bucket: str + region: str + access_key_id: str | None = field(default=None, repr=False) + secret_access_key: str | None = field(default=None, repr=False) + session_token: str | None = field(default=None, repr=False) + store_id: str = "s3" + allow_insecure_http: bool = False + multipart_threshold_bytes: int = 64 * 1024**2 + multipart_chunk_bytes: int = 16 * 1024**2 + multipart_concurrency: int = 4 + max_pool_connections: int = 32 + connect_timeout_seconds: float = 5 + read_timeout_seconds: float = 120 + max_attempts: int = 4 + + def __post_init__(self) -> None: + for name in ("bucket", "region", "store_id"): + _validate_bounded_text(name, getattr(self, name)) + _validate_credentials(self) + _validate_endpoint(self) + integers = ( + "multipart_threshold_bytes", + "multipart_chunk_bytes", + "multipart_concurrency", + "max_pool_connections", + "max_attempts", + ) + for name in integers: + value = getattr(self, name) + if type(value) is not int or value <= 0: + raise ValueError(f"{name} must be positive") + if self.multipart_chunk_bytes < _MIN_MULTIPART_BYTES: + raise ValueError( + f"multipart_chunk_bytes must be at least {_MIN_MULTIPART_BYTES}" + ) + if self.max_pool_connections < self.multipart_concurrency: + raise ValueError( + "max_pool_connections must cover multipart_concurrency" + ) + for name in ("connect_timeout_seconds", "read_timeout_seconds"): + value = getattr(self, name) + if isinstance(value, bool) or not isinstance(value, (int, float)) or value <= 0: + raise ValueError(f"{name} must be positive") + + +class S3PackStore: + def __init__( + self, + client: S3Client, + *, + bucket: str, + store_id: str = "s3", + transfer_config: object = None, + ) -> None: + _validate_bounded_text("bucket", bucket) + _validate_bounded_text("store_id", store_id) + self._client = client + self._bucket = bucket + self._transfer_config = transfer_config + self.store_id = store_id + + @classmethod + def from_config(cls, config: S3StoreConfig) -> S3PackStore: + try: + import boto3 + from boto3.s3.transfer import TransferConfig + from botocore.config import Config + except ImportError as exc: + raise RuntimeError( + "S3 support requires the optional 's3' dependencies" + ) from exc + + client = boto3.client( + "s3", + endpoint_url=config.endpoint_url, + region_name=config.region, + aws_access_key_id=config.access_key_id, + aws_secret_access_key=config.secret_access_key, + aws_session_token=config.session_token, + config=Config( + signature_version="s3v4", + connect_timeout=config.connect_timeout_seconds, + read_timeout=config.read_timeout_seconds, + max_pool_connections=config.max_pool_connections, + retries={"mode": "standard", "max_attempts": config.max_attempts}, + s3={"addressing_style": "path"}, + ), + ) + transfer = TransferConfig( + multipart_threshold=config.multipart_threshold_bytes, + multipart_chunksize=config.multipart_chunk_bytes, + max_concurrency=config.multipart_concurrency, + use_threads=config.multipart_concurrency > 1, + ) + return cls( + client, + bucket=config.bucket, + store_id=config.store_id, + transfer_config=transfer, + ) + + def put(self, pack: PackSource, object_key: str) -> PackRef: + validate_pack_source(pack) + key = str(validate_object_key(object_key)) + existing = self._head_or_none(key) + if existing is not None: + return self._existing_ref(pack, key, existing) + + metadata = { + "dmi-format": "dmi-pack-v1", + "dmi-pack-id": pack.pack_id, + "dmi-sha256": pack.checksum, + "dmi-record-count": str(pack.record_count), + "dmi-created-at-ns": str(pack.created_at_ns), + } + with pack.open() as source: + self._client.upload_fileobj( + source, + self._bucket, + key, + ExtraArgs={ + "ContentType": "application/vnd.dmi.pack", + "Metadata": metadata, + }, + Config=self._transfer_config, + ) + uploaded = self._head_or_none(key) + if uploaded is None: + raise PackIntegrityError("uploaded object is not visible to HeadObject") + return self._existing_ref(pack, key, uploaded) + + def stat(self, ref: PackRef) -> ObjectInfo: + self._validate_ref(ref) + response = self._client.head_object( + Bucket=self._bucket, Key=str(validate_object_key(ref.object_key)) + ) + if not isinstance(response, Mapping): + raise PackIntegrityError("S3 returned an invalid HeadObject response") + size, metadata = self._parse_head(response) + checksum = metadata.get("dmi-sha256") + if checksum is None: + raise PackIntegrityError("S3 object is missing DMI checksum metadata") + if size != ref.object_bytes or checksum != ref.checksum: + raise PackIntegrityError("S3 object does not match its pack reference") + return ObjectInfo(size=size, checksum=checksum) + + def inspect(self, object_key: str) -> PackRef: + key = str(validate_object_key(object_key)) + response = self._client.head_object(Bucket=self._bucket, Key=key) + if not isinstance(response, Mapping): + raise PackIntegrityError("S3 returned an invalid HeadObject response") + size, metadata = self._parse_head(response) + try: + pack_id = str(UUID(metadata["dmi-pack-id"])) + checksum = metadata["dmi-sha256"] + record_count = int(metadata["dmi-record-count"]) + created_at_ns = int(metadata["dmi-created-at-ns"]) + except (KeyError, TypeError, ValueError, AttributeError) as exc: + raise PackIntegrityError("S3 object has invalid DMI metadata") from exc + if ( + metadata.get("dmi-format") != "dmi-pack-v1" + or _SHA256.fullmatch(checksum) is None + or not 1 <= record_count <= 1_000_000 + or created_at_ns < 0 + or size <= 0 + ): + raise PackIntegrityError("S3 object has invalid DMI metadata") + return PackRef( + pack_id=pack_id, + store_id=self.store_id, + object_key=key, + object_bytes=size, + checksum=checksum, + record_count=record_count, + ) + + def read_range(self, ref: PackRef, offset: int, length: int) -> bytes: + self._validate_ref(ref) + if ( + type(offset) is not int + or type(length) is not int + or offset < 0 + or length < 0 + ): + raise ValueError("range offset and length must be non-negative integers") + if offset + length > ref.object_bytes: + raise PackFormatError("requested range exceeds object size") + if length == 0: + return b"" + response = self._client.get_object( + Bucket=self._bucket, + Key=str(validate_object_key(ref.object_key)), + Range=f"bytes={offset}-{offset + length - 1}", + ) + if not isinstance(response, Mapping) or "Body" not in response: + raise PackIntegrityError("S3 returned an invalid range response") + body = response["Body"] + try: + data = body.read(length + 1) + finally: + body.close() + if not isinstance(data, bytes) or len(data) != length: + raise PackIntegrityError("S3 returned a short or oversized range") + return data + + def list_objects( + self, + *, + prefix: str = "", + cursor: str | None = None, + limit: int = 1000, + ) -> ObjectPage: + if prefix: + self._validate_prefix(prefix) + if cursor is not None and ( + not isinstance(cursor, str) + or not cursor + or len(cursor.encode()) > _MAX_CURSOR_BYTES + ): + raise ValueError("cursor must be non-empty and bounded") + if type(limit) is not int or not 1 <= limit <= 1000: + raise ValueError("limit must be between 1 and 1000") + request: dict[str, object] = { + "Bucket": self._bucket, + "Prefix": prefix, + "MaxKeys": limit, + } + if cursor is not None: + request["ContinuationToken"] = cursor + response = self._client.list_objects_v2(**request) + if not isinstance(response, Mapping): + raise PackIntegrityError("S3 returned an invalid listing response") + items = self._parse_listing_items(response) + next_cursor = self._parse_listing_cursor(response) + return ObjectPage(items=items, next_cursor=next_cursor) + + @staticmethod + def _parse_listing_items( + response: Mapping[str, object], + ) -> tuple[StoredObject, ...]: + raw_items = response.get("Contents", ()) + if not isinstance(raw_items, (list, tuple)): + raise PackIntegrityError("S3 listing contents are invalid") + items = [] + for raw in raw_items: + if not isinstance(raw, Mapping): + raise PackIntegrityError("S3 listing item is invalid") + key = raw.get("Key") + size = raw.get("Size") + if not isinstance(key, str) or type(size) is not int or size < 0: + raise PackIntegrityError("S3 listing item is invalid") + validate_object_key(key) + items.append(StoredObject(object_key=key, object_bytes=size)) + return tuple(items) + + @staticmethod + def _parse_listing_cursor(response: Mapping[str, object]) -> str | None: + is_truncated = response.get("IsTruncated", False) + if type(is_truncated) is not bool: + raise PackIntegrityError("S3 listing truncation state is invalid") + next_cursor = response.get("NextContinuationToken") if is_truncated else None + if is_truncated and next_cursor is None: + raise PackIntegrityError("S3 truncated listing is missing its cursor") + if next_cursor is not None and ( + not isinstance(next_cursor, str) + or not next_cursor + or len(next_cursor.encode()) > _MAX_CURSOR_BYTES + ): + raise PackIntegrityError("S3 listing cursor is invalid") + return next_cursor + + def _head_or_none(self, key: str) -> Mapping[str, object] | None: + try: + response = self._client.head_object(Bucket=self._bucket, Key=key) + except Exception as exc: + if self._is_not_found(exc): + return None + raise + if not isinstance(response, Mapping): + raise PackIntegrityError("S3 returned an invalid HeadObject response") + return response + + def _existing_ref( + self, pack: PackSource, key: str, response: Mapping[str, object] + ) -> PackRef: + size, metadata = self._parse_head(response) + expected = { + "dmi-pack-id": pack.pack_id, + "dmi-sha256": pack.checksum, + "dmi-record-count": str(pack.record_count), + "dmi-created-at-ns": str(pack.created_at_ns), + } + if size != pack.object_bytes or any( + metadata.get(name) != value for name, value in expected.items() + ): + raise PackConflictError(f"object key contains different content: {key}") + return PackRef( + pack_id=pack.pack_id, + store_id=self.store_id, + object_key=key, + object_bytes=pack.object_bytes, + checksum=pack.checksum, + record_count=pack.record_count, + ) + + @staticmethod + def _parse_head(response: Mapping[str, object]) -> tuple[int, Mapping[str, str]]: + size = response.get("ContentLength") + metadata = response.get("Metadata") + if type(size) is not int or size < 0 or not isinstance(metadata, Mapping): + raise PackIntegrityError("S3 returned invalid object metadata") + if any( + not isinstance(key, str) or not isinstance(value, str) + for key, value in metadata.items() + ): + raise PackIntegrityError("S3 returned invalid object metadata") + return size, metadata + + @staticmethod + def _is_not_found(exc: Exception) -> bool: + response = getattr(exc, "response", None) + if not isinstance(response, Mapping): + return False + error = response.get("Error", {}) + metadata = response.get("ResponseMetadata", {}) + code = error.get("Code") if isinstance(error, Mapping) else None + status = metadata.get("HTTPStatusCode") if isinstance(metadata, Mapping) else None + return status == 404 or code in {"404", "NoSuchKey", "NotFound"} + + @staticmethod + def _validate_prefix(prefix: str) -> None: + if ( + not isinstance(prefix, str) + or prefix.startswith("/") + or "\\" in prefix + or any(part in {"", ".", ".."} for part in prefix.rstrip("/").split("/")) + ): + raise ValueError("prefix is invalid") + PurePosixPath(prefix) + + def _validate_ref(self, ref: PackRef) -> None: + if ref.store_id != self.store_id: + raise ValueError( + f"pack store mismatch: {ref.store_id!r} != {self.store_id!r}" + ) diff --git a/src/dmi/storage/capture/spool.py b/src/dmi/storage/capture/spool.py new file mode 100644 index 000000000..93f2deb4f --- /dev/null +++ b/src/dmi/storage/capture/spool.py @@ -0,0 +1,518 @@ +from __future__ import annotations + +from dataclasses import dataclass +from concurrent.futures import FIRST_COMPLETED, Future, ThreadPoolExecutor, wait +import math +import os +from pathlib import Path +import random +import re +import tempfile +import threading +import time +from typing import BinaryIO, Callable, Mapping + +from .filesystem import ( + FilesystemPackStore, + copy_pack_source, + validate_object_key, + validate_pack_source, +) +from .model import PackConflictError, PackIntegrityError, PackRef, PackSource, PackStore +from .pipeline import ReadyPack, object_key_for + + +_READY_NAME = re.compile( + r"^(?P[0-9a-f-]{36})\.(?P[0-9]+)\." + r"(?P[0-9]+)\." + r"(?P[0-9a-f]{64})\.dmi-pack\.ready$" +) + + +class SpoolFullError(RuntimeError): + pass + + +@dataclass(frozen=True, slots=True) +class SpoolSnapshot: + entries: int + bytes: int + peak_bytes: int + max_bytes: int + + +@dataclass(frozen=True, slots=True) +class StagedPack: + pack_id: str + created_at_ns: int + record_count: int + checksum: str + object_key: str + path: Path + object_bytes: int + + def open(self) -> BinaryIO: + return FilesystemPackStore._open_regular(self.path) + + +class DurablePackSpool: + def __init__(self, root: str | Path, *, max_bytes: int) -> None: + if type(max_bytes) is not int or max_bytes <= 0: + raise ValueError("max_bytes must be positive") + self.root = Path(root).resolve() + self.root.mkdir(parents=True, exist_ok=True) + self.max_bytes = max_bytes + self._lock = threading.Lock() + ready = tuple(self.root.rglob("*.dmi-pack.ready")) + stale = tuple(self.root.rglob("*.open")) + self._bytes = sum(path.lstat().st_size for path in (*ready, *stale)) + self._entries = len(ready) + self._peak_bytes = self._bytes + + def stage(self, pack: PackSource, object_key: str) -> StagedPack: + validate_pack_source(pack) + key = validate_object_key(object_key) + if key.name != f"{pack.pack_id}.dmi-pack": + raise ValueError("spool object key must end with the pack ID") + ready = self._ready_path(key, pack) + with self._lock: + if ready.exists(): + return self._existing(pack, object_key, ready) + conflicts = tuple(ready.parent.glob(f"{pack.pack_id}.*.dmi-pack.ready")) + if conflicts: + raise PackConflictError( + f"spool already contains a different pack intent: {pack.pack_id}" + ) + if self._bytes + pack.object_bytes > self.max_bytes: + raise SpoolFullError( + f"spool byte limit exceeded: " + f"{self._bytes + pack.object_bytes} > {self.max_bytes}" + ) + ready.parent.mkdir(parents=True, exist_ok=True) + temp_path: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="wb", + prefix=f".{pack.pack_id}.", + suffix=".open", + dir=ready.parent, + delete=False, + ) as handle: + temp_path = Path(handle.name) + copy_pack_source(pack, handle) + handle.flush() + os.fsync(handle.fileno()) + try: + os.link(temp_path, ready) + except FileExistsError: + return self._existing(pack, object_key, ready) + temp_path.unlink() + temp_path = None + self._fsync_directory(ready.parent) + finally: + if temp_path is not None: + temp_path.unlink(missing_ok=True) + self._bytes += pack.object_bytes + self._entries += 1 + self._peak_bytes = max(self._peak_bytes, self._bytes) + return self._entry(ready) + + def recover(self) -> tuple[StagedPack, ...]: + with self._lock: + stale = tuple(self.root.rglob("*.open")) + for path in stale: + path.unlink(missing_ok=True) + for parent in {path.parent for path in stale}: + self._fsync_directory(parent) + paths = tuple(sorted(self.root.rglob("*.dmi-pack.ready"))) + self._bytes = sum(path.lstat().st_size for path in paths) + self._entries = len(paths) + self._peak_bytes = max(self._peak_bytes, self._bytes) + entries = tuple(self._entry(path) for path in paths) + for entry in entries: + if self._checksum(entry.path) != entry.checksum: + raise PackIntegrityError(f"spool checksum mismatch: {entry.pack_id}") + return entries + + def remove(self, staged: StagedPack) -> None: + if staged.path.is_symlink(): + raise PackIntegrityError("ready pack must be a regular spool file") + path = staged.path.resolve() + if not path.is_relative_to(self.root): + raise ValueError("staged pack is outside the spool root") + with self._lock: + if not path.exists(): + return + current = self._entry(path) + if current != staged: + raise PackIntegrityError("staged pack identity changed before removal") + if path.stat().st_size != staged.object_bytes: + raise PackIntegrityError("staged pack size changed before removal") + path.unlink() + self._fsync_directory(path.parent) + self._bytes -= staged.object_bytes + self._entries -= 1 + + def snapshot(self) -> SpoolSnapshot: + with self._lock: + return SpoolSnapshot( + entries=self._entries, + bytes=self._bytes, + peak_bytes=self._peak_bytes, + max_bytes=self.max_bytes, + ) + + def _ready_path(self, key, pack: PackSource) -> Path: + name = ( + f"{pack.pack_id}.{pack.created_at_ns}.{pack.record_count}." + f"{pack.checksum}.dmi-pack.ready" + ) + path = self.root.joinpath(*key.parent.parts, name) + if not path.parent.resolve(strict=False).is_relative_to(self.root): + raise ValueError("object key escapes the spool root") + return path + + def _entry(self, path: Path) -> StagedPack: + resolved = path.resolve() + if ( + path.is_symlink() + or not path.is_file() + or not resolved.is_relative_to(self.root) + ): + raise PackIntegrityError("ready pack must be a regular spool file") + match = _READY_NAME.fullmatch(path.name) + if match is None: + raise PackIntegrityError(f"invalid ready-pack name: {path.name}") + relative = path.relative_to(self.root) + pack_id = match.group("pack_id") + object_key = str(relative.parent / f"{pack_id}.dmi-pack") + staged = StagedPack( + pack_id=pack_id, + created_at_ns=int(match.group("created")), + record_count=int(match.group("records")), + checksum=match.group("checksum"), + object_key=object_key, + path=path, + object_bytes=path.stat().st_size, + ) + validate_pack_source(staged) + return staged + + def _existing( + self, pack: PackSource, object_key: str, ready: Path + ) -> StagedPack: + staged = self._entry(ready) + if ( + staged.object_key != object_key + or staged.object_bytes != pack.object_bytes + or staged.checksum != pack.checksum + or self._checksum(ready) != pack.checksum + ): + raise PackConflictError(f"spool contains different content: {object_key}") + return staged + + @staticmethod + def _checksum(path: Path) -> str: + return FilesystemPackStore._checksum(path) + + @staticmethod + def _fsync_directory(path: Path) -> None: + descriptor = os.open(path, os.O_RDONLY) + try: + os.fsync(descriptor) + finally: + os.close(descriptor) + + +class DurablePackSink: + def __init__(self, spool: DurablePackSpool) -> None: + self._spool = spool + + def persist(self, ready: ReadyPack) -> StagedPack: + return self._spool.stage(ready.pack, object_key_for(ready)) + + +class SpoolUploader: + def __init__(self, spool: DurablePackSpool, store: PackStore) -> None: + self._spool = spool + self._store = store + + def upload(self, staged: StagedPack) -> PackRef: + ref = self._store.put(staged, staged.object_key) + info = self._store.stat(ref) + if info.size != staged.object_bytes or info.checksum != staged.checksum: + raise PackIntegrityError("remote object verification failed") + self._spool.remove(staged) + return ref + + def upload_pending(self, *, limit: int | None = None) -> tuple[PackRef, ...]: + if limit is not None and limit <= 0: + raise ValueError("limit must be positive") + pending = self._spool.recover() + if limit is not None: + pending = pending[:limit] + return tuple(self.upload(staged) for staged in pending) + + +@dataclass(frozen=True, slots=True) +class ParallelUploadConfig: + max_workers: int + max_in_flight_bytes: int + max_attempts: int = 4 + base_backoff_seconds: float = 0.25 + max_backoff_seconds: float = 10 + jitter_ratio: float = 0.2 + + def __post_init__(self) -> None: + for name in ("max_workers", "max_in_flight_bytes", "max_attempts"): + if type(getattr(self, name)) is not int or getattr(self, name) <= 0: + raise ValueError(f"{name} must be positive") + for name in ( + "base_backoff_seconds", + "max_backoff_seconds", + "jitter_ratio", + ): + value = getattr(self, name) + if ( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(value) + or value < 0 + ): + raise ValueError(f"{name} must be finite and non-negative") + if self.max_backoff_seconds < self.base_backoff_seconds: + raise ValueError("max_backoff_seconds must cover base_backoff_seconds") + if self.jitter_ratio > 1: + raise ValueError("jitter_ratio must not exceed one") + + +@dataclass(frozen=True, slots=True) +class UploadFailure: + pack_id: str + object_key: str + attempts: int + error_type: str + + +@dataclass(frozen=True, slots=True) +class UploadSnapshot: + attempted_packs: int + uploaded_packs: int + uploaded_bytes: int + failed_packs: int + retries: int + peak_active_uploads: int + peak_in_flight_bytes: int + event_callback_failures: int + duration_count: int + duration_total_ns: int + duration_max_ns: int + + +@dataclass(frozen=True, slots=True) +class UploadBatchResult: + refs: tuple[PackRef, ...] + failures: tuple[UploadFailure, ...] + snapshot: UploadSnapshot + + +@dataclass(frozen=True, slots=True) +class UploadEvent: + event: str + fields: Mapping[str, int | str] + + +@dataclass(frozen=True, slots=True) +class _UploadOutcome: + ref: PackRef | None + failure: UploadFailure | None + attempts: int + duration_ns: int + + +class ParallelSpoolUploader: + def __init__( + self, + spool: DurablePackSpool, + store: PackStore, + config: ParallelUploadConfig, + *, + event_callback: Callable[[UploadEvent], None] | None = None, + sleep: Callable[[float], None] = time.sleep, + random_value: Callable[[], float] = random.random, + clock_ns: Callable[[], int] = time.monotonic_ns, + ) -> None: + self._spool = spool + self._uploader = SpoolUploader(spool, store) + self._config = config + self._event_callback = event_callback + self._sleep = sleep + self._random_value = random_value + self._clock_ns = clock_ns + self._event_callback_failures = 0 + self._event_lock = threading.Lock() + + def upload_pending(self, *, limit: int | None = None) -> UploadBatchResult: + if limit is not None and (type(limit) is not int or limit <= 0): + raise ValueError("limit must be positive") + pending = list(self._spool.recover()) + if limit is not None: + pending = pending[:limit] + oversized = next( + ( + staged + for staged in pending + if staged.object_bytes > self._config.max_in_flight_bytes + ), + None, + ) + if oversized is not None: + raise ValueError( + f"pack {oversized.pack_id} exceeds the in-flight byte limit" + ) + + active_bytes = 0 + peak_active = 0 + peak_bytes = 0 + outcomes: dict[str, _UploadOutcome] = {} + futures: dict[Future[_UploadOutcome], StagedPack] = {} + remaining = pending.copy() + with ThreadPoolExecutor( + max_workers=self._config.max_workers, + thread_name_prefix="dmi-pack-upload", + ) as executor: + while remaining or futures: + while len(futures) < self._config.max_workers: + selected = next( + ( + (index, staged) + for index, staged in enumerate(remaining) + if active_bytes + staged.object_bytes + <= self._config.max_in_flight_bytes + ), + None, + ) + if selected is None: + break + index, staged = selected + remaining.pop(index) + future = executor.submit(self._upload, staged) + futures[future] = staged + active_bytes += staged.object_bytes + peak_active = max(peak_active, len(futures)) + peak_bytes = max(peak_bytes, active_bytes) + completed, _ = wait(futures, return_when=FIRST_COMPLETED) + for future in completed: + staged = futures.pop(future) + active_bytes -= staged.object_bytes + outcomes[staged.pack_id] = future.result() + + ordered = [outcomes[staged.pack_id] for staged in pending] + refs = tuple(outcome.ref for outcome in ordered if outcome.ref is not None) + failures = tuple( + outcome.failure + for outcome in ordered + if outcome.failure is not None + ) + durations = [outcome.duration_ns for outcome in ordered] + snapshot = UploadSnapshot( + attempted_packs=len(ordered), + uploaded_packs=len(refs), + uploaded_bytes=sum(ref.object_bytes for ref in refs), + failed_packs=len(failures), + retries=sum(max(0, outcome.attempts - 1) for outcome in ordered), + peak_active_uploads=peak_active, + peak_in_flight_bytes=peak_bytes, + event_callback_failures=self._event_callback_failures, + duration_count=len(durations), + duration_total_ns=sum(durations), + duration_max_ns=max(durations, default=0), + ) + return UploadBatchResult(refs=refs, failures=failures, snapshot=snapshot) + + def _upload(self, staged: StagedPack) -> _UploadOutcome: + started = self._clock_ns() + for attempt in range(1, self._config.max_attempts + 1): + try: + ref = self._uploader.upload(staged) + except Exception as exc: + if attempt == self._config.max_attempts or not self._retryable(exc): + failure = UploadFailure( + pack_id=staged.pack_id, + object_key=staged.object_key, + attempts=attempt, + error_type=type(exc).__name__, + ) + self._emit( + "pack_upload_failed", + pack_id=staged.pack_id, + attempt=attempt, + error_type=failure.error_type, + ) + return _UploadOutcome( + ref=None, + failure=failure, + attempts=attempt, + duration_ns=max(0, self._clock_ns() - started), + ) + self._emit( + "pack_upload_retry", + pack_id=staged.pack_id, + attempt=attempt, + error_type=type(exc).__name__, + ) + self._sleep(self._backoff(attempt)) + continue + self._emit( + "pack_upload_committed", + pack_id=staged.pack_id, + attempt=attempt, + object_bytes=staged.object_bytes, + ) + return _UploadOutcome( + ref=ref, + failure=None, + attempts=attempt, + duration_ns=max(0, self._clock_ns() - started), + ) + raise AssertionError("unreachable upload attempt state") + + def _backoff(self, attempt: int) -> float: + base = min( + self._config.max_backoff_seconds, + self._config.base_backoff_seconds * 2 ** (attempt - 1), + ) + return min( + self._config.max_backoff_seconds, + base + base * self._config.jitter_ratio * self._random_value(), + ) + + @staticmethod + def _retryable(exc: Exception) -> bool: + if isinstance(exc, OSError): + return True + response = getattr(exc, "response", None) + if not isinstance(response, Mapping): + return False + metadata = response.get("ResponseMetadata", {}) + error = response.get("Error", {}) + status = metadata.get("HTTPStatusCode") if isinstance(metadata, Mapping) else None + code = error.get("Code") if isinstance(error, Mapping) else None + return status in {408, 429} or ( + isinstance(status, int) and status >= 500 + ) or code in { + "InternalError", + "RequestTimeout", + "ServiceUnavailable", + "SlowDown", + "Throttling", + } + + def _emit(self, event: str, **fields: int | str) -> None: + if self._event_callback is None: + return + with self._event_lock: + try: + self._event_callback(UploadEvent(event=event, fields=fields)) + except Exception: + self._event_callback_failures += 1 diff --git a/tests/ring/Makefile b/tests/ring/Makefile new file mode 100644 index 000000000..f02f27f70 --- /dev/null +++ b/tests/ring/Makefile @@ -0,0 +1,59 @@ +NVCC ?= /usr/local/cuda/bin/nvcc +SM_ARCH ?= native +CXX_STD ?= c++17 +BUILD := build +INCLUDES := -I../.. -I../../monitoring/csrc + +NVCC_FLAGS := \ + -std=$(CXX_STD) \ + -arch=$(SM_ARCH) \ + -O2 \ + -Xcompiler -Wall \ + -Xcompiler -Wextra \ + -Xcompiler -Wno-unused-parameter \ + -lineinfo \ + $(INCLUDES) + +PRODUCER_CU := ../../monitoring/csrc/ring/producer.cu +DRAIN_SRC := ../../monitoring/csrc/ring/drain_thread.cpp + +# test_ring_engine needs ATen headers + libtorch for DrainTask (contains at::Tensor). +# Use torch.utils.cmake_prefix_path to locate them. The Python E2E tests +# (tests/test_e2e_correctness_vs_hf.py) cover the full pipeline including p2p. +TORCH_INC := $(shell python -c "import torch; print(torch.utils.cmake_prefix_path + '/../../include')" 2>/dev/null) +TORCH_LIB := $(shell python -c "import torch; print(torch.utils.cmake_prefix_path + '/../../lib')" 2>/dev/null) + +TARGETS := $(BUILD)/test_rings $(BUILD)/test_producer $(BUILD)/test_ring_engine + +.PHONY: all test_rings test_producer test_ring_engine run clean + +all: $(TARGETS) + +$(BUILD): + mkdir -p $(BUILD) + +$(BUILD)/test_rings: test_rings.cu | $(BUILD) + $(NVCC) $(NVCC_FLAGS) $< -o $@ + +$(BUILD)/test_producer: test_producer.cu $(PRODUCER_CU) | $(BUILD) + $(NVCC) $(NVCC_FLAGS) $^ -o $@ + +$(BUILD)/test_ring_engine: test_ring_engine.cu $(PRODUCER_CU) $(DRAIN_SRC) | $(BUILD) + $(NVCC) $(NVCC_FLAGS) -I$(TORCH_INC) -I$(TORCH_INC)/torch/csrc/api/include -L$(TORCH_LIB) -Xlinker -rpath,$(TORCH_LIB) -ltorch -ltorch_cpu -lc10 -lpthread $^ -o $@ + +test_rings: $(BUILD)/test_rings + $(BUILD)/test_rings + +test_producer: $(BUILD)/test_producer + $(BUILD)/test_producer + +test_ring_engine: $(BUILD)/test_ring_engine + $(BUILD)/test_ring_engine + +run: all + $(BUILD)/test_rings + $(BUILD)/test_producer + $(BUILD)/test_ring_engine + +clean: + rm -rf $(BUILD) diff --git a/tests/ring/test_null_mode.cu b/tests/ring/test_null_mode.cu new file mode 100644 index 000000000..4adfb31cb --- /dev/null +++ b/tests/ring/test_null_mode.cu @@ -0,0 +1,288 @@ +// tests/ring/test_null_mode.cu — Tests for producer null mode. +// +// Verifies that set_ring_null_mode(true) makes producer_kernel launch with +// the same parameters but skip all ring writes, while +// set_ring_null_mode(false) restores normal delivery. +// +// Build: see tests/ring/Makefile (target: test_null_mode) +// Run: ./build/test_null_mode + +#include "ring/ring_alloc.h" +#include "ring/producer.cuh" +#include "ring/pinned_pool.h" +#include "ring/drain_thread.h" +#include "ring/chunk_assembler.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace ring; + +// --------------------------------------------------------------------------- +// Shared test harness (same as test_ring_engine.cu) +// --------------------------------------------------------------------------- +struct ResultCollector { + std::mutex mu; + std::condition_variable cv; + std::vector tensors; + + void push(AssembledTensor&& t) { + std::lock_guard lk(mu); + tensors.push_back(std::move(t)); + cv.notify_all(); + } + + bool wait_for(int n, int timeout_ms = 5000) { + std::unique_lock lk(mu); + return cv.wait_for(lk, std::chrono::milliseconds(timeout_ms), + [&] { return (int)tensors.size() >= n; }); + } + + int count() { + std::lock_guard lk(mu); + return (int)tensors.size(); + } +}; + +static constexpr uint64_t PINNED_BYTES = 8 * 256 * 1024; // 8 chunks worth + +static RingConfig make_cfg() { + RingConfig cfg{}; + cfg.task_ring_entries = 64; + cfg.payload_ring_bytes = 4 * 1024 * 1024; + cfg.chunk_bytes = 256 * 1024; + cfg.pinned_pool_bytes = PINNED_BYTES; + cfg.wait_policy = WaitPolicy::INFINITE; + cfg.drop_reporting = DropReporting::DROP_TASK; + return cfg; +} + +static uint8_t* upload(const std::vector& h, cudaStream_t s) { + uint8_t* d; + cudaMalloc(&d, h.size()); + cudaMemcpyAsync(d, h.data(), h.size(), cudaMemcpyHostToDevice, s); + return d; +} + +// --------------------------------------------------------------------------- +// test_null_mode_no_delivery +// With null mode ON, producer kernel launches but emits nothing to the ring. +// The drain thread should receive no tensors within the timeout window. +// --------------------------------------------------------------------------- +static void test_null_mode_no_delivery() { + printf(" test_null_mode_no_delivery ... "); + + set_ring_null_mode(true); + + RingConfig cfg = make_cfg(); + AllocatedRing ar(cfg); + ar.init(); + + ResultCollector rc; + PinnedPool pool; + pool.init(PINNED_BYTES); + + ChunkAssembler asm_(pool, [&](AssembledTensor&& t) { rc.push(std::move(t)); }); + DrainThread dt(ar.state(), pool, [&](DrainedChunk&& c) { asm_.push(std::move(c)); }); + dt.start(); + + const uint64_t data_bytes = 64 * 1024; + std::vector src(data_bytes, 0xAB); + cudaStream_t stream; + cudaStreamCreate(&stream); + uint8_t* d_src = upload(src, stream); + + launch_producer_with_notify(ar.state(), d_src, data_bytes, + /*logical_task_id=*/1, 0, 0, + DrainThread::hostfunc_cb, &dt, stream); + cudaStreamSynchronize(stream); + + // Give drain thread time to process any spurious notification. + bool got = rc.wait_for(1, /*timeout_ms=*/300); + assert(!got && "null mode: no tensor should arrive"); + assert(rc.count() == 0); + + // task_head must remain 0 — kernel wrote nothing. + assert(*ar.state().task_head == 0); + + dt.stop(); + cudaFree(d_src); + cudaStreamDestroy(stream); + set_ring_null_mode(false); // restore for subsequent tests + printf("PASS\n"); +} + +// --------------------------------------------------------------------------- +// test_null_mode_off_delivers +// With null mode OFF (default), producer works normally. +// --------------------------------------------------------------------------- +static void test_null_mode_off_delivers() { + printf(" test_null_mode_off_delivers ... "); + + set_ring_null_mode(false); + + RingConfig cfg = make_cfg(); + AllocatedRing ar(cfg); + ar.init(); + + ResultCollector rc; + PinnedPool pool; + pool.init(PINNED_BYTES); + + ChunkAssembler asm_(pool, [&](AssembledTensor&& t) { rc.push(std::move(t)); }); + DrainThread dt(ar.state(), pool, [&](DrainedChunk&& c) { asm_.push(std::move(c)); }); + dt.start(); + + const uint64_t data_bytes = 64 * 1024; + std::vector src(data_bytes); + std::iota(src.begin(), src.end(), uint8_t(7)); + cudaStream_t stream; + cudaStreamCreate(&stream); + uint8_t* d_src = upload(src, stream); + + launch_producer_with_notify(ar.state(), d_src, data_bytes, + /*logical_task_id=*/2, 0, 0, + DrainThread::hostfunc_cb, &dt, stream); + + assert(rc.wait_for(1)); + const auto& t = rc.tensors[0]; + assert(!t.is_drop); + assert(t.data.size() == data_bytes); + assert(memcmp(t.data.data(), src.data(), data_bytes) == 0); + + dt.stop(); + cudaStreamSynchronize(stream); + cudaFree(d_src); + cudaStreamDestroy(stream); + printf("PASS\n"); +} + +// --------------------------------------------------------------------------- +// test_toggle +// Submit N kernels in null mode (none arrive), then toggle off and submit +// M kernels in real mode (all M arrive, with correct data). +// Verifies that toggling leaves the ring in a consistent state. +// --------------------------------------------------------------------------- +static void test_toggle() { + printf(" test_toggle ... "); + + RingConfig cfg = make_cfg(); + AllocatedRing ar(cfg); + ar.init(); + + ResultCollector rc; + PinnedPool pool; + pool.init(PINNED_BYTES); + + ChunkAssembler asm_(pool, [&](AssembledTensor&& t) { rc.push(std::move(t)); }); + DrainThread dt(ar.state(), pool, [&](DrainedChunk&& c) { asm_.push(std::move(c)); }); + dt.start(); + + cudaStream_t stream; + cudaStreamCreate(&stream); + + const uint64_t data_bytes = 32 * 1024; + const int N_null = 5; + const int N_real = 4; + + // --- null mode phase --- + set_ring_null_mode(true); + std::vector null_ptrs(N_null); + for (int i = 0; i < N_null; ++i) { + std::vector src(data_bytes, uint8_t(0xCC)); + null_ptrs[i] = upload(src, stream); + launch_producer_with_notify(ar.state(), null_ptrs[i], data_bytes, + uint64_t(100 + i), 0, 0, + DrainThread::hostfunc_cb, &dt, stream); + } + cudaStreamSynchronize(stream); + // Brief wait: drain thread must NOT deliver anything. + bool got = rc.wait_for(1, /*timeout_ms=*/300); + assert(!got && "toggle: null-mode tensors must not arrive"); + assert(*ar.state().task_head == 0); + + // --- real mode phase --- + set_ring_null_mode(false); + std::vector> real_srcs(N_real); + std::vector real_ptrs(N_real); + for (int i = 0; i < N_real; ++i) { + real_srcs[i].resize(data_bytes); + for (uint64_t j = 0; j < data_bytes; ++j) real_srcs[i][j] = uint8_t(i * 13 + j); + real_ptrs[i] = upload(real_srcs[i], stream); + launch_producer_with_notify(ar.state(), real_ptrs[i], data_bytes, + uint64_t(200 + i), 0, 0, + DrainThread::hostfunc_cb, &dt, stream); + } + + assert(rc.wait_for(N_real)); + assert(rc.count() == N_real); + for (int i = 0; i < N_real; ++i) { + const auto& t = rc.tensors[i]; + assert(!t.is_drop); + assert(t.data.size() == data_bytes); + int idx = static_cast(t.logical_task_id) - 200; + assert(idx >= 0 && idx < N_real); + assert(memcmp(t.data.data(), real_srcs[idx].data(), data_bytes) == 0); + } + + dt.stop(); + cudaStreamSynchronize(stream); + for (auto p : null_ptrs) cudaFree(p); + for (auto p : real_ptrs) cudaFree(p); + cudaStreamDestroy(stream); + printf("PASS\n"); +} + +// --------------------------------------------------------------------------- +// test_null_multiple_launches +// Launch many kernels in null mode; ring state (task_head, payload_head) +// must be completely unchanged — no partial writes. +// --------------------------------------------------------------------------- +static void test_null_ring_state_unchanged() { + printf(" test_null_ring_state_unchanged ... "); + + set_ring_null_mode(true); + + RingConfig cfg = make_cfg(); + AllocatedRing ar(cfg); + ar.init(); + + cudaStream_t stream; + cudaStreamCreate(&stream); + + const int N = 10; + const uint64_t data_bytes = 128 * 1024; + std::vector ptrs(N); + for (int i = 0; i < N; ++i) { + std::vector src(data_bytes, uint8_t(i)); + ptrs[i] = upload(src, stream); + // launch_producer (no hostfunc needed; we just check ring state) + launch_producer(ar.state(), ptrs[i], data_bytes, uint64_t(i), 0, 0, stream); + } + cudaStreamSynchronize(stream); + + assert(*ar.state().task_head == 0 && "null mode: task_head must be 0"); + assert(*ar.state().payload_head == 0 && "null mode: payload_head must be 0"); + + for (auto p : ptrs) cudaFree(p); + cudaStreamDestroy(stream); + set_ring_null_mode(false); + printf("PASS\n"); +} + +// --------------------------------------------------------------------------- +int main() { + printf("test_null_mode\n"); + test_null_mode_no_delivery(); + test_null_mode_off_delivers(); + test_toggle(); + test_null_ring_state_unchanged(); + printf("All tests passed.\n"); + return 0; +} diff --git a/tests/ring/test_producer.cu b/tests/ring/test_producer.cu new file mode 100644 index 000000000..37e47c79e --- /dev/null +++ b/tests/ring/test_producer.cu @@ -0,0 +1,462 @@ +// tests/ring/test_producer.cu — M2 producer kernel correctness tests. +// +// Build: make -C tests/ring test_producer +// Run: ./tests/ring/build/test_producer +// +// All GPU tests run sequentially on a single stream (cudaStreamDefault) because +// this GPU serialises kernels from the same process even across streams. + +#include "ring/ring_alloc.h" +#include "ring/producer.cuh" + +#include +#include +#include +#include +#include + +// --------------------------------------------------------------------------- +// Assertion helpers +// --------------------------------------------------------------------------- +static int g_total = 0; +static int g_fail = 0; + +#define ASSERT(cond) do { \ + ++g_total; \ + if (!(cond)) { \ + ++g_fail; \ + fprintf(stderr, " FAIL %s:%d %s\n", \ + __FILE__, __LINE__, #cond); \ + } \ +} while (0) + +#define CUDA_CHECK(e) do { \ + cudaError_t _e = (e); \ + if (_e != cudaSuccess) { \ + fprintf(stderr, "CUDA error %s:%d: %s\n", \ + __FILE__, __LINE__, cudaGetErrorString(_e)); \ + std::exit(1); \ + } \ +} while (0) + +static void banner(const char* name) { + printf("[ TEST ] %s\n", name); + fflush(stdout); +} + +// Read task entries from device into a host-side vector. +static std::vector read_entries(ring::TaskEntry* d_entries, + uint64_t n_slots) +{ + std::vector v(n_slots); + CUDA_CHECK(cudaMemcpy(v.data(), d_entries, + n_slots * sizeof(ring::TaskEntry), + cudaMemcpyDeviceToHost)); + return v; +} + +// Read payload buffer (or subset) from device. +static std::vector read_payload(uint8_t* d_buf, uint64_t nbytes) +{ + std::vector v(nbytes); + CUDA_CHECK(cudaMemcpy(v.data(), d_buf, nbytes, cudaMemcpyDeviceToHost)); + return v; +} + +// Allocate a device buffer and fill it with a simple byte pattern. +static uint8_t* make_src(uint64_t nbytes) +{ + uint8_t* d; + CUDA_CHECK(cudaMalloc(&d, nbytes)); + std::vector h(nbytes); + for (uint64_t i = 0; i < nbytes; ++i) h[i] = static_cast(i & 0xFF); + CUDA_CHECK(cudaMemcpy(d, h.data(), nbytes, cudaMemcpyHostToDevice)); + return d; +} + +// --------------------------------------------------------------------------- +// Test 1 — host-only: RingConfig defaults +// --------------------------------------------------------------------------- +static void test_config_defaults() +{ + banner("config_defaults"); + ring::RingConfig cfg; + ASSERT(cfg.task_ring_entries == 1024); + ASSERT(cfg.payload_ring_bytes == 256ULL * 1024 * 1024); + ASSERT(cfg.chunk_bytes == 64ULL * 1024 * 1024); + ASSERT(cfg.wait_policy == ring::WaitPolicy::INFINITE); + ASSERT(cfg.drop_reporting == ring::DropReporting::DROP_TASK); +} + +// --------------------------------------------------------------------------- +// Test 2 — single-chunk: tensor fits in one chunk +// --------------------------------------------------------------------------- +static void test_single_chunk() +{ + banner("single_chunk"); + + ring::RingConfig cfg; + cfg.task_ring_entries = 16; + cfg.payload_ring_bytes = 16 * 1024; // 16 KiB + cfg.chunk_bytes = 4 * 1024; // 4 KiB — tensor fits in one chunk + + ring::AllocatedRing ar(cfg); + ar.init(); + + const uint64_t src_bytes = 256; + uint8_t* d_src = make_src(src_bytes); + + launch_producer(ar.state(), d_src, src_bytes, + /*logical_task_id=*/42, + /*hook_type=*/1, /*hook_id=*/7); + CUDA_CHECK(cudaDeviceSynchronize()); + + // task_head should now be 1 + ASSERT(*ar.state().task_head == 1); + + auto entries = read_entries(ar.state().task_entries, cfg.task_ring_entries); + const auto& e = entries[0]; + + ASSERT(e.ready_seq == 0); + ASSERT(e.seq_no == 0); + ASSERT(e.logical_task_id == 42); + ASSERT(e.chunk_idx == 0); + ASSERT(e.hook_type == 1); + ASSERT(e.hook_id == 7); + ASSERT(e.tensor_total_bytes == src_bytes); + ASSERT(e.chunk_offset_bytes == 0); + ASSERT((e.flags & ring::TASK_FLAG_IS_FIRST) != 0); + ASSERT((e.flags & ring::TASK_FLAG_IS_LAST) != 0); + ASSERT((e.flags & ring::TASK_FLAG_IS_DROP) == 0); + ASSERT(e.payload_len1 == src_bytes); + ASSERT(e.payload_len2 == 0); + ASSERT(e.payload_off1 == 0); // payload_head started at 0 + + cudaFree(d_src); +} + +// --------------------------------------------------------------------------- +// Test 3 — multi-chunk: tensor spans multiple chunks +// --------------------------------------------------------------------------- +static void test_multi_chunk() +{ + banner("multi_chunk"); + + ring::RingConfig cfg; + cfg.task_ring_entries = 16; + cfg.payload_ring_bytes = 64 * 1024; // 64 KiB + cfg.chunk_bytes = 1024; // 1 KiB per chunk + + ring::AllocatedRing ar(cfg); + ar.init(); + + // 3.5 × chunk_bytes → 4 chunks (3×1024 + 512) + const uint64_t src_bytes = 3 * 1024 + 512; + uint8_t* d_src = make_src(src_bytes); + + launch_producer(ar.state(), d_src, src_bytes, 99, 0, 5); + CUDA_CHECK(cudaDeviceSynchronize()); + + ASSERT(*ar.state().task_head == 4); + + auto entries = read_entries(ar.state().task_entries, cfg.task_ring_entries); + + // chunk 0 — IS_FIRST + ASSERT((entries[0].flags & ring::TASK_FLAG_IS_FIRST) != 0); + ASSERT((entries[0].flags & ring::TASK_FLAG_IS_LAST) == 0); + ASSERT(entries[0].chunk_idx == 0); + ASSERT(entries[0].chunk_offset_bytes == 0); + ASSERT(entries[0].payload_len1 == 1024); + + // chunk 1 — middle + ASSERT((entries[1].flags & ring::TASK_FLAG_IS_FIRST) == 0); + ASSERT((entries[1].flags & ring::TASK_FLAG_IS_LAST) == 0); + ASSERT(entries[1].chunk_idx == 1); + ASSERT(entries[1].chunk_offset_bytes == 1024); + + // chunk 2 — middle + ASSERT(entries[2].chunk_idx == 2); + ASSERT(entries[2].chunk_offset_bytes == 2 * 1024); + + // chunk 3 — IS_LAST, smaller + ASSERT((entries[3].flags & ring::TASK_FLAG_IS_FIRST) == 0); + ASSERT((entries[3].flags & ring::TASK_FLAG_IS_LAST) != 0); + ASSERT(entries[3].chunk_idx == 3); + ASSERT(entries[3].chunk_offset_bytes == 3 * 1024); + ASSERT(entries[3].payload_len1 == 512); + + // All entries share same logical_task_id and tensor_total_bytes + for (int i = 0; i < 4; ++i) { + ASSERT(entries[i].logical_task_id == 99); + ASSERT(entries[i].tensor_total_bytes == src_bytes); + ASSERT(entries[i].hook_id == 5); + } + + cudaFree(d_src); +} + +// --------------------------------------------------------------------------- +// Test 4 — payload wrap: chunk crosses end of ring buffer → two-span entry +// --------------------------------------------------------------------------- +static void test_wrap() +{ + banner("wrap_two_span"); + + ring::RingConfig cfg; + cfg.task_ring_entries = 8; + cfg.payload_ring_bytes = 4096; + cfg.chunk_bytes = 512; + + ring::AllocatedRing ar(cfg); + ar.init(); + + // Simulate a ring state where physical head is near the end. + // head and tail must be PAYLOAD_ALIGN-aligned (16 bytes) since + // vectorized uint4 D2D copies require aligned offsets. + // head=3904, tail=3584: free = 4096 - (3904-3584) = 3776 >= 512. + // Physical head position: 3904 % 4096 = 3904. + // bytes_to_end = 4096 - 3904 = 192 → two spans: len1=192, len2=320. + const uint64_t head = 3904; + const uint64_t tail = 3584; + const uint64_t src_bytes = 512; + *ar.state().payload_head = head; + *ar.state().payload_tail = tail; + + uint8_t* d_src = make_src(src_bytes); + + launch_producer(ar.state(), d_src, src_bytes, 1, 0, 0); + CUDA_CHECK(cudaDeviceSynchronize()); + + auto entries = read_entries(ar.state().task_entries, cfg.task_ring_entries); + const auto& e = entries[0]; + + uint64_t bytes_to_end = 4096 - (head % 4096); // 192 + ASSERT(e.payload_len1 == bytes_to_end); + ASSERT(e.payload_len2 == src_bytes - bytes_to_end); + ASSERT(e.payload_off1 == head % 4096); + ASSERT(e.payload_off2 == 0); + ASSERT(e.payload_len1 + e.payload_len2 == src_bytes); + + cudaFree(d_src); +} + +// --------------------------------------------------------------------------- +// Test 5 — data correctness: verify payload_buf contains the right bytes +// --------------------------------------------------------------------------- +static void test_data_correctness() +{ + banner("data_correctness"); + + ring::RingConfig cfg; + cfg.task_ring_entries = 8; + cfg.payload_ring_bytes = 8 * 1024; + cfg.chunk_bytes = 4 * 1024; + + ring::AllocatedRing ar(cfg); + ar.init(); + + const uint64_t src_bytes = 512; + uint8_t* d_src = make_src(src_bytes); // pattern: byte[i] = i % 256 + + launch_producer(ar.state(), d_src, src_bytes, 0, 0, 0); + CUDA_CHECK(cudaDeviceSynchronize()); + + auto entries = read_entries(ar.state().task_entries, cfg.task_ring_entries); + const auto& e = entries[0]; + ASSERT(e.payload_len1 == src_bytes && e.payload_len2 == 0); + + auto payload = read_payload(ar.state().payload_buf, cfg.payload_ring_bytes); + bool ok = true; + for (uint64_t i = 0; i < src_bytes; ++i) { + if (payload[e.payload_off1 + i] != static_cast(i & 0xFF)) { + ok = false; + fprintf(stderr, " byte[%llu] expected %u got %u\n", + (unsigned long long)i, + (unsigned)static_cast(i & 0xFF), + (unsigned)payload[e.payload_off1 + i]); + break; + } + } + ASSERT(ok); + + cudaFree(d_src); +} + +// --------------------------------------------------------------------------- +// Test 6 — data correctness across wrap: two-span payload bytes are correct +// --------------------------------------------------------------------------- +static void test_data_correctness_wrap() +{ + banner("data_correctness_wrap"); + + ring::RingConfig cfg; + cfg.task_ring_entries = 8; + cfg.payload_ring_bytes = 4096; + cfg.chunk_bytes = 512; + + ring::AllocatedRing ar(cfg); + ar.init(); + + // Same setup as test_wrap: head=3904, tail=3584 → two-span (192 + 320). + // head must be PAYLOAD_ALIGN-aligned (16 bytes) for vectorized uint4 copies. + const uint64_t src_bytes = 512; + *ar.state().payload_head = 3904; + *ar.state().payload_tail = 3584; + + uint8_t* d_src = make_src(src_bytes); + + launch_producer(ar.state(), d_src, src_bytes, 0, 0, 0); + CUDA_CHECK(cudaDeviceSynchronize()); + + auto entries = read_entries(ar.state().task_entries, cfg.task_ring_entries); + const auto& e = entries[0]; + + auto payload = read_payload(ar.state().payload_buf, cfg.payload_ring_bytes); + + bool ok = true; + // Span 1 bytes + for (uint64_t i = 0; i < e.payload_len1 && ok; ++i) { + if (payload[e.payload_off1 + i] != static_cast(i & 0xFF)) { + ok = false; + fprintf(stderr, " span1 byte[%llu] mismatch\n", (unsigned long long)i); + } + } + // Span 2 bytes + for (uint64_t i = 0; i < e.payload_len2 && ok; ++i) { + uint64_t src_i = e.payload_len1 + i; + if (payload[e.payload_off2 + i] != static_cast(src_i & 0xFF)) { + ok = false; + fprintf(stderr, " span2 byte[%llu] mismatch\n", (unsigned long long)i); + } + } + ASSERT(ok); + + cudaFree(d_src); +} + +// --------------------------------------------------------------------------- +// Test 7 — timeout-drop: frozen heartbeat + nearly-full payload → DROP marker +// --------------------------------------------------------------------------- +static void test_timeout_drop() +{ + banner("timeout_drop"); + + ring::RingConfig cfg; + cfg.task_ring_entries = 16; + cfg.payload_ring_bytes = 1024; + cfg.chunk_bytes = 512; + cfg.wait_policy = ring::WaitPolicy::TIMEOUT_DROP; + cfg.drop_reporting = ring::DropReporting::DROP_TASK; + cfg.no_progress_timeout_cycles = 200000; // ~0.1 ms at 2 GHz + + ring::AllocatedRing ar(cfg); + ar.init(); + + // Simulate a nearly-full payload ring: only 16 bytes free. + // The producer needs 512 bytes but only 16 are available. + // consumer_heartbeat stays 0 → no progress → timeout. + *ar.state().payload_head = cfg.payload_ring_bytes - 16; + *ar.state().payload_tail = 0; + // task ring is empty (head=tail=0), so 1 task slot is available for DROP. + + const uint64_t src_bytes = 512; + uint8_t* d_src = make_src(src_bytes); + + launch_producer(ar.state(), d_src, src_bytes, 77, 2, 9); + CUDA_CHECK(cudaDeviceSynchronize()); + + // Producer should have published exactly 1 DROP entry. + ASSERT(*ar.state().task_head == 1); + + auto entries = read_entries(ar.state().task_entries, cfg.task_ring_entries); + const auto& e = entries[0]; + + ASSERT((e.flags & ring::TASK_FLAG_IS_DROP) != 0); + ASSERT((e.flags & ring::TASK_FLAG_IS_LAST) != 0); + ASSERT((e.flags & ring::TASK_FLAG_IS_FIRST) != 0); // chunk 0 was first + ASSERT(e.reason == ring::DROP_REASON_TIMEOUT_NO_PROGRESS); + ASSERT(e.payload_len1 == 0); + ASSERT(e.payload_len2 == 0); + ASSERT(e.logical_task_id == 77); + ASSERT(e.hook_type == 2); + ASSERT(e.hook_id == 9); + + cudaFree(d_src); +} + +// --------------------------------------------------------------------------- +// Test 8 — backpressure release: producer spins, consumer advances tail +// --------------------------------------------------------------------------- +// Simulated by running two sequential kernels on the same stream: +// kernel A fills ring, then we manually advance tails, then kernel B runs. +static void test_backpressure_release() +{ + banner("backpressure_release"); + + ring::RingConfig cfg; + cfg.task_ring_entries = 4; // small ring + cfg.payload_ring_bytes = 4 * 1024; // 4 KiB + cfg.chunk_bytes = 1 * 1024; // 1 KiB + + ring::AllocatedRing ar(cfg); + ar.init(); + + // Fill the ring: 4 chunks × 1KiB = 4KiB (exactly fills payload ring). + const uint64_t src_bytes_a = 4 * 1024; + uint8_t* d_src_a = make_src(src_bytes_a); + launch_producer(ar.state(), d_src_a, src_bytes_a, 10, 0, 1); + CUDA_CHECK(cudaDeviceSynchronize()); + ASSERT(*ar.state().task_head == 4); + + // Now simulate consumer draining slot 0: advance task_tail and payload_tail. + // This frees 1 task slot + 1 KiB of payload. + *ar.state().task_tail = 1; + *ar.state().payload_tail = 1024; + ++(*ar.state().consumer_heartbeat); + + // Producer should now be able to fit 1 more chunk. + const uint64_t src_bytes_b = 512; + uint8_t* d_src_b = make_src(src_bytes_b); + launch_producer(ar.state(), d_src_b, src_bytes_b, 20, 0, 2); + CUDA_CHECK(cudaDeviceSynchronize()); + + // task_head should now be 5 (4 from A + 1 from B). + ASSERT(*ar.state().task_head == 5); + + auto entries = read_entries(ar.state().task_entries, cfg.task_ring_entries); + // Slot 4 % 4 = 0 (reused after tail advanced). + const auto& e = entries[4 % 4]; + ASSERT(e.logical_task_id == 20); + ASSERT(e.tensor_total_bytes == src_bytes_b); + ASSERT(e.seq_no == 4); + + cudaFree(d_src_a); + cudaFree(d_src_b); +} + +// --------------------------------------------------------------------------- +// main +// --------------------------------------------------------------------------- +int main() +{ + setbuf(stdout, nullptr); + + printf("=== test_producer ===\n"); + + test_config_defaults(); + test_single_chunk(); + test_multi_chunk(); + test_wrap(); + test_data_correctness(); + test_data_correctness_wrap(); + test_timeout_drop(); + test_backpressure_release(); + + printf("\n%d / %d assertions passed\n", g_total - g_fail, g_total); + if (g_fail > 0) { + fprintf(stderr, "%d FAILURES\n", g_fail); + return 1; + } + printf("ALL PASS\n"); + return 0; +} diff --git a/tests/ring/test_ring_engine.cu b/tests/ring/test_ring_engine.cu new file mode 100644 index 000000000..cd6e71171 --- /dev/null +++ b/tests/ring/test_ring_engine.cu @@ -0,0 +1,375 @@ +// tests/ring/test_ring_engine.cu — End-to-end tests for the ring drain pipeline. +// +// Tests the full path: +// GPU producer kernel → drain thread (batch D2H into pinned staging) +// → task queue → verify DrainTask contents +// +// NOTE: The new drain pipeline (batch D2H + pinned staging + p2p thread) +// requires ATen linkage for DrainTask (contains at::Tensor). The p2p thread +// and full tensor assembly are tested via the Python E2E tests +// (tests/test_e2e_correctness_vs_hf.py) which exercise the entire pipeline +// including metadata FIFO, reshape, slicing, and ClickHouse submission. +// +// These C++ tests verify the drain-side pipeline: producer → drain thread → +// task queue, confirming that staged data matches the original GPU tensor. +// +// Build: make test_ring_engine (requires PyTorch for ATen headers) + +#include "ring/ring_alloc.h" +#include "ring/producer.cuh" +#include "ring/drain_thread.h" +#include "ring/pinned_staging.h" + +#include +#include +#include +#include +#include + +using namespace ring; + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +// Allocate and upload a host vector to GPU device memory. +static uint8_t* upload(const std::vector& h, cudaStream_t s = 0) { + uint8_t* d; + cudaMalloc(&d, h.size()); + cudaMemcpyAsync(d, h.data(), h.size(), cudaMemcpyHostToDevice, s); + return d; +} + +// Verify staging data for a single-chunk tensor matches source. +static void verify_staging_data(const DrainTask& task, + const std::vector& src) { + uint64_t data_len = task.data_len1 + task.data_len2; + assert(data_len == src.size()); + // Copy from staging (may wrap) + std::vector got(data_len); + if (task.data_len1 > 0) + memcpy(got.data(), task.data_ptr1, task.data_len1); + if (task.data_len2 > 0) + memcpy(got.data() + task.data_len1, task.data_ptr2, task.data_len2); + assert(memcmp(got.data(), src.data(), src.size()) == 0); +} + +// --------------------------------------------------------------------------- +// test_single_chunk +// One tensor that fits in a single chunk. +// Verifies drain task fields and staged data match the original. +// --------------------------------------------------------------------------- +static void test_single_chunk() { + printf(" test_single_chunk ... "); + + const uint64_t data_bytes = 128 * 1024; // 128 KB < chunk_bytes (256 KB) + RingConfig cfg{}; + cfg.task_ring_entries = 64; + cfg.payload_ring_bytes = 4 * 1024 * 1024; + cfg.chunk_bytes = 256 * 1024; + cfg.wait_policy = WaitPolicy::INFINITE; + cfg.drop_reporting = DropReporting::DROP_TASK; + + AllocatedRing ar(cfg); + ar.init(); + + PinnedStaging staging; + staging.init(cfg.payload_ring_bytes); + + DrainThread dt(ar.state(), staging, cfg); + dt.start(); + + std::vector src(data_bytes); + std::iota(src.begin(), src.end(), uint8_t(0)); + cudaStream_t stream; + cudaStreamCreate(&stream); + uint8_t* d_src = upload(src, stream); + + launch_producer_with_notify(ar.state(), d_src, data_bytes, + /*logical_task_id=*/1, /*hook_type=*/0, /*hook_id=*/0, + DrainThread::hostfunc_cb, &dt, stream); + + // Stop drain to trigger final flush + cudaStreamSynchronize(stream); + dt.stop(); + + // Now collect tasks (drain has pushed them before joining) + // Signal p2p stop so wait_for_tasks returns + dt.signal_p2p_stop(); + uint64_t n = dt.wait_for_tasks(); + assert(n > 0); + std::vector tasks; + dt.pop_tasks(n, tasks); + assert(tasks.size() == 1); + + const auto& t = tasks[0]; + assert(t.logical_task_id == 1); + assert(t.tensor_total_bytes == data_bytes); + assert(t.flags & TASK_FLAG_IS_FIRST); + assert(t.flags & TASK_FLAG_IS_LAST); + verify_staging_data(t, src); + + cudaFree(d_src); + cudaStreamDestroy(stream); + printf("PASS\n"); +} + +// --------------------------------------------------------------------------- +// test_multi_chunk +// One large tensor that spans multiple chunks. +// Verifies all chunks arrive with correct metadata. +// --------------------------------------------------------------------------- +static void test_multi_chunk() { + printf(" test_multi_chunk ... "); + + const uint64_t chunk_bytes = 64 * 1024; // 64 KB + const uint64_t data_bytes = 3 * chunk_bytes + 48; // 3 full + 1 partial (aligned to 16) + RingConfig cfg{}; + cfg.task_ring_entries = 64; + cfg.payload_ring_bytes = 2 * 1024 * 1024; + cfg.chunk_bytes = chunk_bytes; + cfg.wait_policy = WaitPolicy::INFINITE; + cfg.drop_reporting = DropReporting::DROP_TASK; + + AllocatedRing ar(cfg); + ar.init(); + + PinnedStaging staging; + staging.init(cfg.payload_ring_bytes); + + DrainThread dt(ar.state(), staging, cfg); + dt.start(); + + std::vector src(data_bytes); + for (uint64_t i = 0; i < data_bytes; ++i) src[i] = uint8_t(i * 7 + 3); + + cudaStream_t stream; + cudaStreamCreate(&stream); + uint8_t* d_src = upload(src, stream); + + launch_producer_with_notify(ar.state(), d_src, data_bytes, + /*logical_task_id=*/2, 0, 0, + DrainThread::hostfunc_cb, &dt, stream); + + cudaStreamSynchronize(stream); + dt.stop(); + dt.signal_p2p_stop(); + + uint64_t n = dt.wait_for_tasks(); + assert(n > 0); + std::vector tasks; + dt.pop_tasks(n, tasks); + assert(tasks.size() == 4); // 3 full + 1 partial + + // Reassemble from staging + std::vector assembled(data_bytes); + for (const auto& t : tasks) { + if (t.data_len1 > 0) + memcpy(assembled.data() + t.chunk_offset_bytes, t.data_ptr1, t.data_len1); + if (t.data_len2 > 0) + memcpy(assembled.data() + t.chunk_offset_bytes + t.data_len1, + t.data_ptr2, t.data_len2); + } + assert(memcmp(assembled.data(), src.data(), data_bytes) == 0); + + assert(tasks[0].flags & TASK_FLAG_IS_FIRST); + assert(tasks[3].flags & TASK_FLAG_IS_LAST); + for (const auto& t : tasks) { + assert(t.tensor_total_bytes == data_bytes); + assert(t.logical_task_id == 2); + } + + cudaFree(d_src); + cudaStreamDestroy(stream); + printf("PASS\n"); +} + +// --------------------------------------------------------------------------- +// test_drop +// Fill the ring so the producer times out and emits a DROP marker. +// --------------------------------------------------------------------------- +static void test_drop() { + printf(" test_drop ... "); + + RingConfig cfg{}; + cfg.task_ring_entries = 4; + cfg.payload_ring_bytes = 128 * 1024; + cfg.chunk_bytes = 64 * 1024; + cfg.wait_policy = WaitPolicy::TIMEOUT_DROP; + cfg.no_progress_timeout_cycles = 1000000ULL; // ~0.4 ms at 2.5 GHz + cfg.drop_reporting = DropReporting::DROP_TASK; + + AllocatedRing ar(cfg); + ar.init(); + + // Completely fill the payload ring so producer immediately drops. + *ar.state().payload_head = cfg.payload_ring_bytes; + *ar.state().payload_tail = 0; // free = 0 + + PinnedStaging staging; + staging.init(cfg.payload_ring_bytes); + + DrainThread dt(ar.state(), staging, cfg); + dt.start(); + + const uint64_t data_bytes = 32 * 1024; + std::vector src(data_bytes, 0xBE); + cudaStream_t stream; + cudaStreamCreate(&stream); + uint8_t* d_src = upload(src, stream); + + launch_producer_with_notify(ar.state(), d_src, data_bytes, + /*logical_task_id=*/4, 0, 0, + DrainThread::hostfunc_cb, &dt, stream); + + cudaStreamSynchronize(stream); + dt.stop(); + dt.signal_p2p_stop(); + + uint64_t n = dt.wait_for_tasks(); + assert(n > 0); + std::vector tasks; + dt.pop_tasks(n, tasks); + assert(tasks.size() == 1); + assert(tasks[0].flags & TASK_FLAG_IS_DROP); + + cudaFree(d_src); + cudaStreamDestroy(stream); + printf("PASS\n"); +} + +// --------------------------------------------------------------------------- +// test_multiple_tensors +// Submit N independent tensors back-to-back. +// Verifies drain handles queued notifications and data correctness. +// --------------------------------------------------------------------------- +static void test_multiple_tensors() { + printf(" test_multiple_tensors ... "); + + const int N = 8; + const uint64_t data_bytes = 32 * 1024; // 32 KB each + RingConfig cfg{}; + cfg.task_ring_entries = 64; + cfg.payload_ring_bytes = 4 * 1024 * 1024; + cfg.chunk_bytes = 64 * 1024; + cfg.wait_policy = WaitPolicy::INFINITE; + cfg.drop_reporting = DropReporting::DROP_TASK; + + AllocatedRing ar(cfg); + ar.init(); + + PinnedStaging staging; + staging.init(cfg.payload_ring_bytes); + + DrainThread dt(ar.state(), staging, cfg); + dt.start(); + + cudaStream_t stream; + cudaStreamCreate(&stream); + + std::vector d_srcs(N); + std::vector> srcs(N); + for (int i = 0; i < N; ++i) { + srcs[i].resize(data_bytes); + for (uint64_t j = 0; j < data_bytes; ++j) srcs[i][j] = uint8_t(i + j); + d_srcs[i] = upload(srcs[i], stream); + launch_producer_with_notify(ar.state(), d_srcs[i], data_bytes, + /*logical_task_id=*/uint64_t(10 + i), 0, 0, + DrainThread::hostfunc_cb, &dt, stream); + } + + cudaStreamSynchronize(stream); + dt.stop(); + dt.signal_p2p_stop(); + + // Collect all tasks + std::vector all_tasks; + while (true) { + uint64_t n = dt.wait_for_tasks(); + if (n == 0) break; + std::vector batch; + dt.pop_tasks(n, batch); + for (auto& t : batch) all_tasks.push_back(std::move(t)); + } + assert(static_cast(all_tasks.size()) == N); + + // Each should be a single-chunk tensor (32KB < 64KB chunk) + for (int i = 0; i < N; ++i) { + const auto& t = all_tasks[i]; + assert(!(t.flags & TASK_FLAG_IS_DROP)); + assert(t.flags & TASK_FLAG_IS_FIRST); + assert(t.flags & TASK_FLAG_IS_LAST); + assert(t.tensor_total_bytes == data_bytes); + + int idx = static_cast(t.logical_task_id) - 10; + assert(idx >= 0 && idx < N); + verify_staging_data(t, srcs[idx]); + } + + for (auto d : d_srcs) cudaFree(d); + cudaStreamDestroy(stream); + printf("PASS\n"); +} + +// --------------------------------------------------------------------------- +// test_zero_byte_tensor +// A tensor with src_bytes=0 produces one task entry with no payload. +// --------------------------------------------------------------------------- +static void test_zero_byte_tensor() { + printf(" test_zero_byte_tensor ... "); + + RingConfig cfg{}; + cfg.task_ring_entries = 64; + cfg.payload_ring_bytes = 4 * 1024 * 1024; + cfg.chunk_bytes = 256 * 1024; + cfg.wait_policy = WaitPolicy::INFINITE; + cfg.drop_reporting = DropReporting::DROP_TASK; + + AllocatedRing ar(cfg); + ar.init(); + + PinnedStaging staging; + staging.init(cfg.payload_ring_bytes); + + DrainThread dt(ar.state(), staging, cfg); + dt.start(); + + cudaStream_t stream; + cudaStreamCreate(&stream); + + launch_producer_with_notify(ar.state(), nullptr, 0, + /*logical_task_id=*/99, 0, 0, + DrainThread::hostfunc_cb, &dt, stream); + + cudaStreamSynchronize(stream); + dt.stop(); + dt.signal_p2p_stop(); + + uint64_t n = dt.wait_for_tasks(); + assert(n > 0); + std::vector tasks; + dt.pop_tasks(n, tasks); + assert(tasks.size() == 1); + + const auto& t = tasks[0]; + assert(!(t.flags & TASK_FLAG_IS_DROP)); + assert(t.logical_task_id == 99); + assert(t.tensor_total_bytes == 0); + assert(t.data_len1 == 0); + assert(t.data_len2 == 0); + + cudaStreamDestroy(stream); + printf("PASS\n"); +} + +// --------------------------------------------------------------------------- +int main() { + printf("test_ring_engine (batch drain pipeline)\n"); + test_single_chunk(); + test_multi_chunk(); + test_drop(); + test_multiple_tensors(); + test_zero_byte_tensor(); + printf("All 5 tests passed.\n"); + return 0; +} diff --git a/tests/ring/test_rings.cu b/tests/ring/test_rings.cu new file mode 100644 index 000000000..738431c70 --- /dev/null +++ b/tests/ring/test_rings.cu @@ -0,0 +1,758 @@ +// tests/ring/test_rings.cu — Standalone CUDA unit tests for ring primitives. +// +// Covers all Milestone 1 acceptance criteria: +// - Payload wrap correctness (two-span) +// - Task ring FIFO correctness by seq_no +// - ready_seq guard: consumer cannot read unpublished slot +// - DROP marker correctness +// +// Build with: make -C tests/ring +// Run with: ./tests/ring/test_rings + +#include +#include +#include +#include +#include + +#include + +// Include the ring headers under test. +#include "../../monitoring/csrc/ring/task_entry.h" +#include "../../monitoring/csrc/ring/ring_config.h" +#include "../../monitoring/csrc/ring/payload_ring.cuh" +#include "../../monitoring/csrc/ring/task_ring.cuh" + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +#define CUDA_CHECK(expr) \ + do { \ + cudaError_t _e = (expr); \ + if (_e != cudaSuccess) { \ + fprintf(stderr, "CUDA error at %s:%d — %s\n", \ + __FILE__, __LINE__, cudaGetErrorString(_e)); \ + exit(1); \ + } \ + } while (0) + +static int g_pass = 0; +static int g_fail = 0; + +#define EXPECT(cond) \ + do { \ + if (!(cond)) { \ + fprintf(stderr, " FAIL %s:%d " #cond "\n", \ + __FILE__, __LINE__); \ + g_fail++; \ + } else { \ + g_pass++; \ + } \ + } while (0) + +static void banner(const char* name) { + printf("[ TEST ] %s\n", name); + fflush(stdout); +} + +// =========================================================================== +// HOST-ONLY TESTS (pure arithmetic — __host__ __device__ functions) +// =========================================================================== + +// --------------------------------------------------------------------------- +// Test 1: payload_free_bytes basic accounting +// --------------------------------------------------------------------------- +static void test_payload_free_bytes() { + banner("payload_free_bytes — basic accounting"); + using namespace ring; + + uint64_t cap = 1024; + + // Empty ring. + EXPECT(payload_free_bytes(0, 0, cap) == cap); + + // Half full. + EXPECT(payload_free_bytes(512, 0, cap) == 512); + + // Producer wrapped around (head > cap). + EXPECT(payload_free_bytes(1024, 0, cap) == 0); // full + EXPECT(payload_free_bytes(1100, 100, cap) == 24); // 1024 - (1100-100) + + // Consumer advanced. + EXPECT(payload_free_bytes(800, 300, cap) == 524); +} + +// --------------------------------------------------------------------------- +// Test 2: payload_compute_spans — single span (no wrap) +// --------------------------------------------------------------------------- +static void test_payload_single_span() { + banner("payload_compute_spans — single span (no wrap)"); + using namespace ring; + + uint64_t cap = 256; + uint64_t head = 0; + + // From the beginning. + TwoSpan s = payload_compute_spans(head, cap, 100); + EXPECT(s.off1 == 0); + EXPECT(s.len1 == 100); + EXPECT(s.len2 == 0); + + // From the middle, still fits before end. + head = 100; + s = payload_compute_spans(head, cap, 50); + EXPECT(s.off1 == 100); + EXPECT(s.len1 == 50); + EXPECT(s.len2 == 0); + + // Exactly fills to the end. + head = 200; + s = payload_compute_spans(head, cap, 56); + EXPECT(s.off1 == 200); + EXPECT(s.len1 == 56); + EXPECT(s.len2 == 0); +} + +// --------------------------------------------------------------------------- +// Test 3: payload_compute_spans — two-span (wrap-around) +// --------------------------------------------------------------------------- +static void test_payload_two_span_wrap() { + banner("payload_compute_spans — two-span (wrap-around)"); + using namespace ring; + + uint64_t cap = 64; + + // Head at offset 50 in the ring; request 20 bytes that must wrap. + // span1: [50, 14) → bytes 50..63 + // span2: [0, 6) → bytes 0..5 + uint64_t head = 50; // physical off = 50 % 64 = 50 + TwoSpan s = payload_compute_spans(head, cap, 20); + EXPECT(s.off1 == 50); + EXPECT(s.len1 == 14); + EXPECT(s.off2 == 0); + EXPECT(s.len2 == 6); + EXPECT(s.len1 + s.len2 == 20); + + // Head has wrapped (head = 130 → phys = 2); request 5 bytes, no wrap. + head = 130; // phys = 130 % 64 = 2 + s = payload_compute_spans(head, cap, 5); + EXPECT(s.off1 == 2); + EXPECT(s.len1 == 5); + EXPECT(s.len2 == 0); + + // Head at physical position 60; request 10 bytes → wraps at boundary. + head = 60; // phys = 60 + s = payload_compute_spans(head, cap, 10); + EXPECT(s.off1 == 60); + EXPECT(s.len1 == 4); + EXPECT(s.off2 == 0); + EXPECT(s.len2 == 6); + + // Exactly at boundary (off == 0 after wrap), request full capacity. + head = 64; // phys = 0 + s = payload_compute_spans(head, cap, 64); + EXPECT(s.off1 == 0); + EXPECT(s.len1 == 64); + EXPECT(s.len2 == 0); +} + +// --------------------------------------------------------------------------- +// Test 4: payload ring simulation — multiple reservations with tail advance +// --------------------------------------------------------------------------- +static void test_payload_ring_simulation() { + banner("payload_ring_simulation — multi-reservation + tail advance"); + using namespace ring; + + uint64_t cap = 64; + uint64_t head = 0; + uint64_t tail = 0; + + // Reserve 50 bytes. + EXPECT(payload_free_bytes(head, tail, cap) == 64); + TwoSpan s1 = payload_compute_spans(head, cap, 50); + EXPECT(s1.off1 == 0 && s1.len1 == 50 && s1.len2 == 0); + payload_advance_head(head, 50); + EXPECT(head == 50); + + // Reserve 10 more bytes (still fits). + EXPECT(payload_free_bytes(head, tail, cap) == 14); + TwoSpan s2 = payload_compute_spans(head, cap, 10); + EXPECT(s2.off1 == 50 && s2.len1 == 10 && s2.len2 == 0); + payload_advance_head(head, 10); + EXPECT(head == 60); + + // Ring is nearly full (4 bytes free); need 30 — wait until consumer frees. + EXPECT(payload_free_bytes(head, tail, cap) == 4); + + // Consumer releases first reservation (50 bytes). + payload_release(tail, 50); + EXPECT(tail == 50); + EXPECT(payload_free_bytes(head, tail, cap) == 54); + + // Now reserve 30 bytes — wraps: off=60, len1=4, off2=0, len2=26. + TwoSpan s3 = payload_compute_spans(head, cap, 30); + EXPECT(s3.off1 == 60); + EXPECT(s3.len1 == 4); + EXPECT(s3.off2 == 0); + EXPECT(s3.len2 == 26); + payload_advance_head(head, 30); + EXPECT(head == 90); // 60 + 30 (unwrapped) + + // Consumer releases the second reservation (10 bytes). + payload_release(tail, 10); + EXPECT(tail == 60); + EXPECT(payload_free_bytes(head, tail, cap) == 34); // cap - (90 - 60) +} + +// --------------------------------------------------------------------------- +// Test 5: task_free_slots basic accounting +// --------------------------------------------------------------------------- +static void test_task_free_slots() { + banner("task_free_slots — basic accounting"); + using namespace ring; + + uint64_t cap = 16; + EXPECT(task_free_slots(0, 0, cap) == 16); + EXPECT(task_free_slots(8, 0, cap) == 8); + EXPECT(task_free_slots(16, 0, cap) == 0); // full + EXPECT(task_free_slots(16, 16, cap) == 16); // empty again after consumer drains + EXPECT(task_free_slots(1024, 1020, cap) == 12); +} + +// --------------------------------------------------------------------------- +// Test 6: TaskEntry layout (size, alignment, field offsets) +// --------------------------------------------------------------------------- +static void test_task_entry_layout() { + banner("TaskEntry — size and field offsets"); + using namespace ring; + + EXPECT(sizeof(TaskEntry) == 128); + EXPECT(alignof(TaskEntry) == 128); + EXPECT(offsetof(TaskEntry, ready_seq) == 0); + EXPECT(offsetof(TaskEntry, seq_no) == 8); + EXPECT(offsetof(TaskEntry, logical_task_id) == 16); + EXPECT(offsetof(TaskEntry, chunk_offset_bytes) == 24); + EXPECT(offsetof(TaskEntry, tensor_total_bytes) == 32); + EXPECT(offsetof(TaskEntry, payload_off1) == 40); + EXPECT(offsetof(TaskEntry, payload_len1) == 48); + EXPECT(offsetof(TaskEntry, payload_off2) == 56); + EXPECT(offsetof(TaskEntry, payload_len2) == 64); + EXPECT(offsetof(TaskEntry, chunk_idx) == 72); + EXPECT(offsetof(TaskEntry, hook_type) == 76); + EXPECT(offsetof(TaskEntry, hook_id) == 80); + EXPECT(offsetof(TaskEntry, flags) == 84); + EXPECT(offsetof(TaskEntry, reason) == 88); +} + +// --------------------------------------------------------------------------- +// Test 7: DROP marker fields +// --------------------------------------------------------------------------- +static void test_drop_marker_fields() { + banner("DROP marker — fields and flags"); + using namespace ring; + + TaskEntry e{}; + // Producer fills a DROP marker. + e.seq_no = 42; + e.logical_task_id = 999; + e.flags = TASK_FLAG_IS_DROP; + e.reason = DROP_REASON_TIMEOUT_NO_PROGRESS; + e.payload_len1 = 0; + e.payload_len2 = 0; + + EXPECT((e.flags & TASK_FLAG_IS_DROP) != 0); + EXPECT((e.flags & TASK_FLAG_IS_FIRST) == 0); + EXPECT((e.flags & TASK_FLAG_IS_LAST) == 0); + EXPECT(e.payload_len1 == 0); + EXPECT(e.payload_len2 == 0); + EXPECT(e.reason == DROP_REASON_TIMEOUT_NO_PROGRESS); + EXPECT(payload_chunk_bytes(e.payload_len1, e.payload_len2) == 0); + + // SENTINEL must not equal any plausible seq_no value (just check a few). + EXPECT(READY_SEQ_SENTINEL != 0); + EXPECT(READY_SEQ_SENTINEL != 42); + EXPECT(READY_SEQ_SENTINEL != uint64_t(-2)); // -2 != UINT64_MAX +} + +// --------------------------------------------------------------------------- +// Test 8: RingConfig defaults are sane +// --------------------------------------------------------------------------- +static void test_ring_config_defaults() { + banner("RingConfig — default values"); + using namespace ring; + + RingConfig cfg; + EXPECT(cfg.task_ring_entries == 1024); + EXPECT(cfg.payload_ring_bytes == 256ULL * 1024 * 1024); + EXPECT(cfg.chunk_bytes == 64ULL * 1024 * 1024); + EXPECT(cfg.wait_policy == WaitPolicy::INFINITE); + EXPECT(cfg.drop_reporting == DropReporting::DROP_TASK); + EXPECT(cfg.chunk_bytes * 4 <= cfg.payload_ring_bytes); // forward-progress +} + +// =========================================================================== +// GPU TESTS +// =========================================================================== + +// --------------------------------------------------------------------------- +// Test 9: task ring FIFO — sequential producer→consumer on the same stream. +// +// The ring is large enough to hold all entries at once, so the producer fills +// it entirely before the consumer drains it. This tests pure FIFO ordering +// without requiring concurrent kernel execution. +// +// Verifies: +// - Entries are consumed in seq_no order. +// - hook_id values match what the producer wrote. +// - task_release correctly resets ready_seq so slots can be reused. +// --------------------------------------------------------------------------- + +static const int FIFO_N = 32; // number of tasks to produce and consume + +// Producer: publishes FIFO_N entries; ring is large enough (cap >= FIFO_N) +// so it never needs to wait for the consumer. +__global__ void kernel_producer_fifo(ring::TaskEntry* entries, + uint64_t cap, + volatile uint64_t* d_head) +{ + for (int i = 0; i < FIFO_N; i++) { + uint64_t seq = *d_head; + + ring::TaskEntry e{}; + e.seq_no = seq; + e.logical_task_id = (uint64_t)i; + e.hook_id = (uint32_t)i; + e.flags = ring::TASK_FLAG_IS_FIRST | ring::TASK_FLAG_IS_LAST; + + ring::task_publish(entries, cap, seq, e); + *d_head = seq + 1; + } +} + +// Consumer: drains FIFO_N entries; called after the producer completes. +__global__ void kernel_consumer_fifo(const ring::TaskEntry* entries, + ring::TaskEntry* entries_rw, + uint64_t cap, + volatile uint64_t* d_tail, + uint32_t* results, + uint64_t* result_seq) +{ + for (int i = 0; i < FIFO_N; i++) { + uint64_t tail = *d_tail; + const ring::TaskEntry* e = ring::task_spin_wait(entries, cap, tail); + results[i] = e->hook_id; + result_seq[i] = e->seq_no; + ring::task_release(entries_rw, cap, tail); + *d_tail = tail + 1; + } +} + +static void test_task_ring_fifo_gpu() { + banner("task ring FIFO (GPU) — sequential producer then consumer"); + using namespace ring; + + // Ring larger than FIFO_N so producer never blocks waiting for space. + const uint64_t cap = 64; + + TaskEntry* d_entries; + CUDA_CHECK(cudaMalloc(&d_entries, cap * sizeof(TaskEntry))); + task_ring_init(d_entries, cap, 0); + + volatile uint64_t* d_head; + volatile uint64_t* d_tail; + CUDA_CHECK(cudaMallocManaged((void**)&d_head, sizeof(uint64_t))); + CUDA_CHECK(cudaMallocManaged((void**)&d_tail, sizeof(uint64_t))); + *d_head = 0; + *d_tail = 0; + + uint32_t* results; + uint64_t* result_seq; + CUDA_CHECK(cudaMallocManaged(&results, FIFO_N * sizeof(uint32_t))); + CUDA_CHECK(cudaMallocManaged(&result_seq, FIFO_N * sizeof(uint64_t))); + + // Sync to ensure managed memory is initialised before kernels run. + CUDA_CHECK(cudaDeviceSynchronize()); + + cudaStream_t stream; + CUDA_CHECK(cudaStreamCreate(&stream)); + + // Sequential: producer runs completely, then consumer drains. + kernel_producer_fifo<<<1,1,0,stream>>>(d_entries, cap, d_head); + kernel_consumer_fifo<<<1,1,0,stream>>>(d_entries, d_entries, cap, + d_tail, results, result_seq); + + CUDA_CHECK(cudaStreamSynchronize(stream)); + CUDA_CHECK(cudaDeviceSynchronize()); + + for (int i = 0; i < FIFO_N; i++) { + EXPECT(results[i] == (uint32_t)i); + EXPECT(result_seq[i] == (uint64_t)i); + } + + CUDA_CHECK(cudaStreamDestroy(stream)); + cudaFree(d_entries); + cudaFree((void*)d_head); + cudaFree((void*)d_tail); + cudaFree(results); + cudaFree(result_seq); +} + +// --------------------------------------------------------------------------- +// Test 10: ready_seq lifecycle — sequential guard semantics (single stream). +// +// In production the producer runs inside the CUDA graph and the consumer runs +// after the graph completes (ordered via a CUDA event). This test mimics +// that flow and verifies the full ready_seq lifecycle: +// +// Phase A — after init: ready_seq[i] == SENTINEL (slot not readable) +// Phase B — after publish: ready_seq[seq] == seq_no (slot readable) +// Phase C — after release: ready_seq[seq] == SENTINEL (slot recycled) +// +// Correctness guarantee covered: if task_publish omitted __threadfence() or +// wrote the wrong value, the consumer's spin_wait (which checks ready_seq) +// would malfunction. The sequential stream ordering means "happened-before" +// is enforced by the GPU's stream FIFO, but the volatile + membar semantics +// are still exercised by the actual instructions generated. +// --------------------------------------------------------------------------- + +__global__ void kernel_check_all_sentinel(const ring::TaskEntry* entries, + uint64_t cap, + int* ok) +{ + for (uint64_t i = 0; i < cap; i++) { + uint64_t rs = *reinterpret_cast( + &entries[i].ready_seq); + if (rs != ring::READY_SEQ_SENTINEL) { *ok = 0; return; } + } + *ok = 1; +} + +__global__ void kernel_publish_one(ring::TaskEntry* entries, + uint64_t cap, + volatile uint64_t* d_head, + int* ok) +{ + uint64_t seq = *d_head; + ring::TaskEntry e{}; + e.seq_no = seq; + e.logical_task_id = 0xDEADBEEFULL; + e.hook_id = 0xCAFEBABEu; + e.flags = ring::TASK_FLAG_IS_FIRST | ring::TASK_FLAG_IS_LAST; + e.tensor_total_bytes = 12345; + ring::task_publish(entries, cap, seq, e); + *d_head = seq + 1; + + // Verify ready_seq is now the sequence number we published. + uint64_t rs = *reinterpret_cast( + &entries[seq % cap].ready_seq); + *ok = (rs == seq) ? 1 : 0; +} + +__global__ void kernel_consume_one(const ring::TaskEntry* entries, + ring::TaskEntry* entries_rw, + uint64_t cap, + volatile uint64_t* d_tail, + uint64_t* out_hook_id, + uint64_t* out_ltid, + uint64_t* out_total, + int* sentinel_ok) +{ + uint64_t tail = *d_tail; + const ring::TaskEntry* e = ring::task_spin_wait(entries, cap, tail); + *out_hook_id = e->hook_id; + *out_ltid = e->logical_task_id; + *out_total = e->tensor_total_bytes; + ring::task_release(entries_rw, cap, tail); + *d_tail = tail + 1; + + // After release, ready_seq must be SENTINEL again. + uint64_t rs = *reinterpret_cast( + &entries_rw[tail % cap].ready_seq); + *sentinel_ok = (rs == ring::READY_SEQ_SENTINEL) ? 1 : 0; +} + +static void test_ready_seq_lifecycle_gpu() { + banner("ready_seq lifecycle (GPU) — init sentinel, publish, consume, reset"); + using namespace ring; + + const uint64_t cap = 8; + + TaskEntry* d_entries; + CUDA_CHECK(cudaMalloc(&d_entries, cap * sizeof(TaskEntry))); + task_ring_init(d_entries, cap, 0); + + volatile uint64_t* d_head; + volatile uint64_t* d_tail; + CUDA_CHECK(cudaMallocManaged((void**)&d_head, sizeof(uint64_t))); + CUDA_CHECK(cudaMallocManaged((void**)&d_tail, sizeof(uint64_t))); + *d_head = 0; + *d_tail = 0; + + int* d_ok_sentinel; + int* d_ok_publish; + int* d_ok_sentinel2; + uint64_t* d_hook_id; + uint64_t* d_ltid; + uint64_t* d_total; + CUDA_CHECK(cudaMallocManaged(&d_ok_sentinel, sizeof(int))); + CUDA_CHECK(cudaMallocManaged(&d_ok_publish, sizeof(int))); + CUDA_CHECK(cudaMallocManaged(&d_ok_sentinel2, sizeof(int))); + CUDA_CHECK(cudaMallocManaged(&d_hook_id, sizeof(uint64_t))); + CUDA_CHECK(cudaMallocManaged(&d_ltid, sizeof(uint64_t))); + CUDA_CHECK(cudaMallocManaged(&d_total, sizeof(uint64_t))); + + CUDA_CHECK(cudaDeviceSynchronize()); + + cudaStream_t stream; + CUDA_CHECK(cudaStreamCreate(&stream)); + + // All phases run sequentially on one stream. + kernel_check_all_sentinel<<<1,1,0,stream>>>(d_entries, cap, d_ok_sentinel); + kernel_publish_one<<<1,1,0,stream>>>(d_entries, cap, d_head, d_ok_publish); + kernel_consume_one<<<1,1,0,stream>>>(d_entries, d_entries, cap, + d_tail, d_hook_id, d_ltid, d_total, + d_ok_sentinel2); + + CUDA_CHECK(cudaStreamSynchronize(stream)); + CUDA_CHECK(cudaDeviceSynchronize()); + + EXPECT(*d_ok_sentinel == 1); // all slots are SENTINEL after init + EXPECT(*d_ok_publish == 1); // ready_seq is set correctly after publish + EXPECT(*d_hook_id == 0xCAFEBABEULL); + EXPECT(*d_ltid == 0xDEADBEEFULL); + EXPECT(*d_total == 12345ULL); + EXPECT(*d_ok_sentinel2 == 1); // slot reset to SENTINEL after release + + CUDA_CHECK(cudaStreamDestroy(stream)); + cudaFree(d_entries); + cudaFree((void*)d_head); + cudaFree((void*)d_tail); + cudaFree(d_ok_sentinel); + cudaFree(d_ok_publish); + cudaFree(d_ok_sentinel2); + cudaFree(d_hook_id); + cudaFree(d_ltid); + cudaFree(d_total); +} + +// --------------------------------------------------------------------------- +// Test 11: DROP marker round-trip (GPU, sequential) — producer emits IS_DROP, +// consumer reads it and observes zero payload and IS_DROP flag. +// --------------------------------------------------------------------------- + +__global__ void kernel_producer_drop(ring::TaskEntry* entries, + uint64_t cap, + volatile uint64_t* d_head) +{ + uint64_t seq = *d_head; + ring::TaskEntry e{}; + e.seq_no = seq; + e.logical_task_id = 777; + e.flags = ring::TASK_FLAG_IS_DROP; + e.reason = ring::DROP_REASON_TIMEOUT_NO_PROGRESS; + e.payload_len1 = 0; + e.payload_len2 = 0; + ring::task_publish(entries, cap, seq, e); + *d_head = seq + 1; +} + +__global__ void kernel_consumer_drop(const ring::TaskEntry* entries, + ring::TaskEntry* entries_rw, + uint64_t cap, + volatile uint64_t* d_tail, + uint32_t* out_flags, + uint32_t* out_reason, + uint64_t* out_len_total) +{ + uint64_t tail = *d_tail; + const ring::TaskEntry* e = ring::task_spin_wait(entries, cap, tail); + *out_flags = e->flags; + *out_reason = e->reason; + *out_len_total = e->payload_len1 + e->payload_len2; + ring::task_release(entries_rw, cap, tail); + *d_tail = tail + 1; +} + +static void test_drop_marker_gpu() { + banner("DROP marker (GPU) — producer emits drop, consumer sees IS_DROP"); + using namespace ring; + + const uint64_t cap = 4; + + TaskEntry* d_entries; + CUDA_CHECK(cudaMalloc(&d_entries, cap * sizeof(TaskEntry))); + task_ring_init(d_entries, cap, 0); + + volatile uint64_t* d_head; + volatile uint64_t* d_tail; + CUDA_CHECK(cudaMallocManaged((void**)&d_head, sizeof(uint64_t))); + CUDA_CHECK(cudaMallocManaged((void**)&d_tail, sizeof(uint64_t))); + *d_head = 0; + *d_tail = 0; + + uint32_t* out_flags; + uint32_t* out_reason; + uint64_t* out_len_total; + CUDA_CHECK(cudaMallocManaged(&out_flags, sizeof(uint32_t))); + CUDA_CHECK(cudaMallocManaged(&out_reason, sizeof(uint32_t))); + CUDA_CHECK(cudaMallocManaged(&out_len_total, sizeof(uint64_t))); + + CUDA_CHECK(cudaDeviceSynchronize()); + + cudaStream_t stream; + CUDA_CHECK(cudaStreamCreate(&stream)); + + // Sequential: producer publishes DROP, then consumer drains it. + kernel_producer_drop<<<1,1,0,stream>>>(d_entries, cap, d_head); + kernel_consumer_drop<<<1,1,0,stream>>>(d_entries, d_entries, cap, + d_tail, + out_flags, out_reason, out_len_total); + + CUDA_CHECK(cudaStreamSynchronize(stream)); + CUDA_CHECK(cudaDeviceSynchronize()); + + EXPECT((*out_flags & TASK_FLAG_IS_DROP) != 0); + EXPECT((*out_flags & TASK_FLAG_IS_FIRST) == 0); + EXPECT((*out_flags & TASK_FLAG_IS_LAST) == 0); + EXPECT(*out_reason == DROP_REASON_TIMEOUT_NO_PROGRESS); + EXPECT(*out_len_total == 0); + + CUDA_CHECK(cudaStreamDestroy(stream)); + cudaFree(d_entries); + cudaFree((void*)d_head); + cudaFree((void*)d_tail); + cudaFree(out_flags); + cudaFree(out_reason); + cudaFree(out_len_total); +} + +// --------------------------------------------------------------------------- +// Test 12: task ring wrap-around (slot reuse) — sequential producer→consumer, +// ring slots reused 4× to verify FIFO is preserved across full wraps. +// +// Uses cap < WRAP_TOTAL so that slots are recycled; sequential (same stream) +// to avoid concurrent-spin dependency issues. +// --------------------------------------------------------------------------- +static const int WRAP_TOTAL = 64; // 4× a cap-16 ring + +__global__ void kernel_producer_wrap(ring::TaskEntry* entries, + uint64_t cap, + volatile uint64_t* d_head, + int n_total, + int base_id) +{ + // Producer assumes ring is drained by a previous consumer pass. + for (int i = 0; i < n_total; i++) { + uint64_t seq = *d_head; + ring::TaskEntry e{}; + e.seq_no = seq; + e.hook_id = (uint32_t)(base_id + i); // globally unique hook_id + e.flags = ring::TASK_FLAG_IS_FIRST | ring::TASK_FLAG_IS_LAST; + ring::task_publish(entries, cap, seq, e); + *d_head = seq + 1; + } +} + +__global__ void kernel_consumer_wrap(const ring::TaskEntry* entries, + ring::TaskEntry* entries_rw, + uint64_t cap, + volatile uint64_t* d_tail, + uint32_t* results, + int n_total) +{ + for (int i = 0; i < n_total; i++) { + uint64_t tail = *d_tail; + const ring::TaskEntry* e = ring::task_spin_wait(entries, cap, tail); + results[i] = e->hook_id; + ring::task_release(entries_rw, cap, tail); + *d_tail = tail + 1; + } +} + +static void test_task_ring_wrap_reuse_gpu() { + banner("task ring wrap-around (GPU) — slot reuse, FIFO preserved"); + using namespace ring; + + // cap < WRAP_TOTAL to force slot reuse; run in batches of cap entries. + const uint64_t cap = 16; + const int batch = (int)cap; // publish+consume `batch` at a time + + TaskEntry* d_entries; + CUDA_CHECK(cudaMalloc(&d_entries, cap * sizeof(TaskEntry))); + task_ring_init(d_entries, cap, 0); + + volatile uint64_t* d_head; + volatile uint64_t* d_tail; + CUDA_CHECK(cudaMallocManaged((void**)&d_head, sizeof(uint64_t))); + CUDA_CHECK(cudaMallocManaged((void**)&d_tail, sizeof(uint64_t))); + *d_head = 0; + *d_tail = 0; + + uint32_t* results; + CUDA_CHECK(cudaMallocManaged(&results, WRAP_TOTAL * sizeof(uint32_t))); + + CUDA_CHECK(cudaDeviceSynchronize()); + + cudaStream_t stream; + CUDA_CHECK(cudaStreamCreate(&stream)); + + // Process in batches of `cap` entries so the ring never overflows. + // Each pass: producer fills the ring, then consumer drains it. + for (int base = 0; base < WRAP_TOTAL; base += batch) { + kernel_producer_wrap<<<1,1,0,stream>>>(d_entries, cap, d_head, + batch, base); + kernel_consumer_wrap<<<1,1,0,stream>>>(d_entries, d_entries, cap, + d_tail, results + base, batch); + } + + CUDA_CHECK(cudaStreamSynchronize(stream)); + CUDA_CHECK(cudaDeviceSynchronize()); + + for (int i = 0; i < WRAP_TOTAL; i++) { + EXPECT(results[i] == (uint32_t)i); + } + + CUDA_CHECK(cudaStreamDestroy(stream)); + cudaFree(d_entries); + cudaFree((void*)d_head); + cudaFree((void*)d_tail); + cudaFree(results); +} + +// =========================================================================== +// main +// =========================================================================== +int main() { + // Disable stdout buffering so output is visible even if the process is + // killed mid-run (helps debug hanging GPU tests). + setbuf(stdout, nullptr); + + printf("=== Ring Unit Tests ===\n\n"); + + // Print device info. + int dev = 0; + cudaDeviceProp prop{}; + CUDA_CHECK(cudaGetDeviceProperties(&prop, dev)); + printf("Device: %s (CC %d.%d)\n\n", prop.name, + prop.major, prop.minor); + + // Host tests (pure arithmetic). + test_payload_free_bytes(); + test_payload_single_span(); + test_payload_two_span_wrap(); + test_payload_ring_simulation(); + test_task_free_slots(); + test_task_entry_layout(); + test_drop_marker_fields(); + test_ring_config_defaults(); + + // GPU tests (all sequential — producer then consumer on same stream). + test_task_ring_fifo_gpu(); + test_ready_seq_lifecycle_gpu(); + test_drop_marker_gpu(); + test_task_ring_wrap_reuse_gpu(); + + printf("\n=== Results: %d passed, %d failed ===\n", g_pass, g_fail); + return (g_fail > 0) ? 1 : 0; +} diff --git a/tests/test_capture_catalog_benchmark.py b/tests/test_capture_catalog_benchmark.py new file mode 100644 index 000000000..f6cd3ca14 --- /dev/null +++ b/tests/test_capture_catalog_benchmark.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +import pytest + +from benchmarks.bench_capture_catalog import measure_inserts, synthetic_descriptors + + +pytestmark = pytest.mark.cpu + + +class _Writer: + def __init__(self): + self.batches = [] + + def write_descriptors(self, descriptors, *, index_version): + self.batches.append(tuple(descriptors)) + + +def test_catalog_benchmark_uses_bounded_insert_batches(): + writer = _Writer() + result = measure_inserts( + writer, synthetic_descriptors(5), batch_rows=2, trials=2 + ) + + assert [len(batch) for batch in writer.batches] == [2, 2, 1, 2, 2, 1] + assert result["inserts_per_trial"] == 3 + assert result["rows_per_second_median"] > 0 diff --git a/tests/test_capture_catalog_indexer.py b/tests/test_capture_catalog_indexer.py new file mode 100644 index 000000000..a1b93a465 --- /dev/null +++ b/tests/test_capture_catalog_indexer.py @@ -0,0 +1,267 @@ +from __future__ import annotations + +from dataclasses import replace +from pathlib import Path +from uuid import UUID + +import pytest + +from dmi.storage.capture import ( + CatalogIndexer, + CatalogIndexerConfig, + CatalogReconciler, + CaptureMetadata, + CaptureRecord, + FilesystemPackStore, + ObjectPage, + PackIndex, + PackRef, + PackWriter, + StoredObject, +) + + +pytestmark = pytest.mark.cpu + + +def _metadata(capture_id: str, step: int) -> CaptureMetadata: + return CaptureMetadata( + capture_id=capture_id, + tenant_id="tenant-a", + experiment_id="exp-a", + run_id="run-a", + session_id="session-a", + request_id="request-a", + sequence_id="sequence-a", + model_id="model-a", + model_revision="revision-a", + adapter_revision=None, + capture_policy_version="policy-v1", + hook_name="resid_pre", + layer_number=3, + producer_rank=0, + step_number=step, + token_start=step, + token_end=step + 1, + batch_position=0, + dtype="float32", + shape=(2,), + captured_at_ns=1_700_000_000_000_000_000 + step, + ) + + +class _Inventory: + def __init__(self, store: FilesystemPackStore, refs: list[PackRef]): + self.store_id = store.store_id + self._store = store + self._refs = {ref.object_key: ref for ref in refs} + self.ranges: list[tuple[str, int, int]] = [] + + def inspect(self, object_key: str) -> PackRef: + return self._refs[object_key] + + def list_objects(self, *, prefix="", cursor=None, limit=1000) -> ObjectPage: + keys = sorted(key for key in self._refs if key.startswith(prefix)) + start = int(cursor or 0) + selected = keys[start : start + limit] + next_index = start + len(selected) + return ObjectPage( + items=tuple( + StoredObject(key, self._refs[key].object_bytes) for key in selected + ), + next_cursor=str(next_index) if next_index < len(keys) else None, + ) + + def read_range(self, ref: PackRef, offset: int, length: int) -> bytes: + self.ranges.append((ref.pack_id, offset, length)) + return self._store.read_range(ref, offset, length) + + def stat(self, ref): + return self._store.stat(ref) + + def put(self, pack, object_key): + return self._store.put(pack, object_key) + + +class _CatalogWriter: + def __init__(self): + self.committed: set[tuple[str, str]] = set() + self.descriptor_batches: list[tuple] = [] + self.pack_batches: list[tuple[PackRef, ...]] = [] + self.fail_commit_once = False + + def committed_pack_ids(self, identities): + return self.committed.intersection(identities) + + def write_descriptors(self, descriptors, *, index_version): + self.descriptor_batches.append(tuple(descriptors)) + + def commit_packs(self, refs, *, index_version): + self.pack_batches.append(tuple(refs)) + if self.fail_commit_once: + self.fail_commit_once = False + raise RuntimeError("ambiguous commit") + self.committed.update((ref.store_id, ref.pack_id) for ref in refs) + + +def _packs(tmp_path: Path, counts: tuple[int, ...] = (2, 1)): + store = FilesystemPackStore(tmp_path, store_id="local") + refs = [] + step = 0 + for pack_number, count in enumerate(counts, 1): + writer = PackWriter( + pack_id=UUID(f"018f0000-0000-7000-8000-{pack_number:012d}"), + created_at_ns=1_700_000_000_000_000_000 + pack_number, + max_pack_bytes=1024 * 1024, + ) + for _ in range(count): + metadata = _metadata(f"capture-{step}", step) + writer.append(CaptureRecord(metadata, b"abcdefgh")) + step += 1 + sealed = writer.seal() + refs.append(store.put(sealed, f"packs/{sealed.pack_id}.dmi-pack")) + return _Inventory(store, refs), refs + + +def test_indexer_reads_only_footers_and_batches_rows(tmp_path: Path): + inventory, refs = _packs(tmp_path) + writer = _CatalogWriter() + indexer = CatalogIndexer( + inventory, + writer, + config=CatalogIndexerConfig(max_packs=8, max_rows_per_insert=2), + clock_ns=lambda: 42, + timer_ns=iter((100, 125)).__next__, + ) + + result = indexer.index(refs) + + assert result.indexed_packs == 2 + assert result.indexed_rows == 3 + assert result.descriptor_inserts == 2 + assert result.estimated_bytes > 0 + assert result.elapsed_ns == 25 + assert [len(batch) for batch in writer.descriptor_batches] == [2, 1] + assert [len(batch) for batch in writer.pack_batches] == [2] + assert len(inventory.ranges) == 4 + assert sum(length for _, _, length in inventory.ranges) < sum( + ref.object_bytes for ref in refs + ) + + +def test_duplicate_event_and_missed_event_converge_on_rebuild(tmp_path: Path): + inventory, refs = _packs(tmp_path) + writer = _CatalogWriter() + reconciler = CatalogReconciler( + inventory, + CatalogIndexer(inventory, writer, clock_ns=lambda: 42), + ) + + duplicate = reconciler.index_object_keys([refs[0].object_key, refs[0].object_key]) + rebuilt = reconciler.rebuild(prefix="packs/", page_size=1, max_pages=8) + + assert duplicate.indexed_packs == 1 + assert rebuilt.indexed_packs == 1 + assert rebuilt.skipped_packs == 1 + assert writer.committed == {(ref.store_id, ref.pack_id) for ref in refs} + assert sum(len(batch) for batch in writer.descriptor_batches) == 3 + + +def test_ambiguous_pack_commit_replays_descriptors_but_converges(tmp_path: Path): + inventory, refs = _packs(tmp_path, (2,)) + writer = _CatalogWriter() + writer.fail_commit_once = True + indexer = CatalogIndexer(inventory, writer, clock_ns=lambda: 42) + + with pytest.raises(RuntimeError, match="ambiguous commit"): + indexer.index(refs) + result = indexer.index(refs) + + assert result.indexed_packs == 1 + assert sum(len(batch) for batch in writer.descriptor_batches) == 4 + assert writer.committed == {(refs[0].store_id, refs[0].pack_id)} + + +def test_indexer_rejects_conflicting_pack_identity(tmp_path: Path): + inventory, refs = _packs(tmp_path, (1,)) + indexer = CatalogIndexer(inventory, _CatalogWriter()) + conflict = replace(refs[0], checksum="f" * 64) + + with pytest.raises(ValueError, match="conflicting pack identity"): + indexer.index([refs[0], conflict]) + + +def test_indexer_bounds_notification_batch(tmp_path: Path): + inventory, refs = _packs(tmp_path) + indexer = CatalogIndexer( + inventory, + _CatalogWriter(), + config=CatalogIndexerConfig(max_packs=1), + ) + + with pytest.raises(ValueError, match="max_packs"): + indexer.index(refs) + + with pytest.raises(ValueError, match="max_packs"): + indexer.index([refs[0], refs[0]]) + + +def test_reconciler_bounds_raw_duplicate_notifications(tmp_path: Path): + inventory, refs = _packs(tmp_path, (1,)) + reconciler = CatalogReconciler( + inventory, + CatalogIndexer( + inventory, + _CatalogWriter(), + config=CatalogIndexerConfig(max_packs=1), + ), + ) + + with pytest.raises(ValueError, match="max_packs"): + reconciler.index_object_keys([refs[0].object_key, refs[0].object_key]) + + +def test_corrupt_footer_does_not_commit_pack(tmp_path: Path): + inventory, refs = _packs(tmp_path, (1,)) + original = inventory.read_range + + def corrupt(ref, offset, length): + data = bytearray(original(ref, offset, length)) + if length != PackIndex.trailer_size(): + data[0] ^= 1 + return bytes(data) + + inventory.read_range = corrupt + writer = _CatalogWriter() + + result = CatalogIndexer(inventory, writer).index(refs) + + assert result.failed_packs == 1 + assert result.indexed_packs == 0 + assert not writer.pack_batches + + +def test_rebuild_caps_failure_details_without_losing_failure_count(tmp_path: Path): + inventory, _ = _packs(tmp_path, (1, 1)) + original = inventory.read_range + + def corrupt(ref, offset, length): + data = bytearray(original(ref, offset, length)) + if length != PackIndex.trailer_size(): + data[0] ^= 1 + return bytes(data) + + inventory.read_range = corrupt + reconciler = CatalogReconciler( + inventory, + CatalogIndexer( + inventory, + _CatalogWriter(), + config=CatalogIndexerConfig(max_failure_details=1), + ), + ) + + result = reconciler.rebuild(prefix="packs/", page_size=1) + + assert result.failed_packs == 2 + assert len(result.failures) == 1 diff --git a/tests/test_capture_pack_benchmark.py b/tests/test_capture_pack_benchmark.py new file mode 100644 index 000000000..36bd1c8a4 --- /dev/null +++ b/tests/test_capture_pack_benchmark.py @@ -0,0 +1,87 @@ +from __future__ import annotations + +import json + +import pytest + +from benchmarks.bench_capture_pack import ( + PackBenchmarkConfig, + generate_payload_pool, + main, + run_trial, +) + + +pytestmark = pytest.mark.cpu + + +def test_pack_benchmark_config_rejects_invalid_bounds(): + with pytest.raises(ValueError, match="records"): + PackBenchmarkConfig(records=0) + with pytest.raises(ValueError, match="payload_bytes"): + PackBenchmarkConfig(payload_bytes=3, dtype="float32") + with pytest.raises(ValueError, match="target_pack_bytes"): + PackBenchmarkConfig(payload_bytes=64, target_pack_bytes=32) + + +def test_payload_pool_is_bounded_and_deterministic(): + config = PackBenchmarkConfig( + records=4, + payload_bytes=64, + target_pack_bytes=256, + pool_size=2, + pattern="random", + seed=11, + trials=1, + ) + + first = generate_payload_pool(config) + second = generate_payload_pool(config) + + assert first == second + assert len(first) == 2 + assert all(len(payload) == 64 for payload in first) + + +def test_trial_reports_verified_pack_and_payload_counts(): + result = run_trial( + PackBenchmarkConfig( + records=9, + payload_bytes=64, + target_pack_bytes=2048, + pool_size=2, + pattern="zeros", + trials=1, + ) + ) + + assert result.record_count == 9 + assert result.logical_bytes == 9 * 64 + assert result.pack_count > 1 + assert result.largest_pack_bytes <= 2048 + assert result.packed_bytes > result.logical_bytes + assert result.seconds > 0 + assert result.as_dict()["logical_gib_per_second"] > 0 + + +def test_dry_run_emits_resolved_workload(capsys): + assert main( + [ + "--records", + "7", + "--payload-bytes", + "64", + "--target-pack-bytes", + "256", + "--pool-size", + "2", + "--trials", + "1", + "--dry-run", + ] + ) == 0 + + result = json.loads(capsys.readouterr().out) + assert result["dry_run"] is True + assert result["config"]["records"] == 7 + assert result["config"]["payload_bytes"] == 64 diff --git a/tests/test_capture_parallel_upload.py b/tests/test_capture_parallel_upload.py new file mode 100644 index 000000000..ef387232d --- /dev/null +++ b/tests/test_capture_parallel_upload.py @@ -0,0 +1,208 @@ +from __future__ import annotations + +from pathlib import Path +import threading +import time +from uuid import UUID + +import pytest + +from dmi.storage.capture import ( + CaptureMetadata, + CaptureRecord, + DurablePackSink, + DurablePackSpool, + FilesystemPackStore, + FlushReason, + ParallelSpoolUploader, + ParallelUploadConfig, + PackWriter, + ReadyPack, +) + + +pytestmark = pytest.mark.cpu + + +def _stage(spool: DurablePackSpool, index: int): + metadata = CaptureMetadata( + capture_id=f"capture-{index}", + tenant_id="tenant-a", + experiment_id="exp-a", + run_id="run-a", + session_id="session-a", + request_id=f"request-{index}", + sequence_id=f"sequence-{index}", + model_id="model-a", + model_revision="revision-a", + adapter_revision=None, + capture_policy_version="policy-v1", + hook_name="resid_pre", + layer_number=index, + producer_rank=0, + step_number=index, + token_start=index, + token_end=index + 1, + batch_position=0, + dtype="uint8", + shape=(256,), + captured_at_ns=1_700_000_000_000_000_000 + index, + ) + writer = PackWriter( + pack_id=UUID(f"018f0000-0000-7000-8000-{index:012d}"), + created_at_ns=metadata.captured_at_ns, + max_pack_bytes=1024 * 1024, + ) + writer.append(CaptureRecord(metadata=metadata, payload=bytes([index]) * 256)) + pack = writer.seal() + return DurablePackSink(spool).persist( + ReadyPack(pack, metadata, FlushReason.SHUTDOWN) + ) + + +class _MeasuredStore(FilesystemPackStore): + def __init__(self, root: Path, *, fail_once: bool = False): + super().__init__(root, store_id="remote") + self.fail_once = fail_once + self._failed: set[str] = set() + self._lock = threading.Lock() + self.active = 0 + self.peak_active = 0 + + def put(self, pack, object_key): + with self._lock: + self.active += 1 + self.peak_active = max(self.peak_active, self.active) + try: + time.sleep(0.01) + if self.fail_once and object_key not in self._failed: + self._failed.add(object_key) + raise OSError("transient upload failure") + return super().put(pack, object_key) + finally: + with self._lock: + self.active -= 1 + + +def test_parallel_uploader_bounds_workers_and_in_flight_bytes(tmp_path: Path): + spool = DurablePackSpool(tmp_path / "spool", max_bytes=1024 * 1024) + staged = [_stage(spool, index) for index in range(4)] + budget = staged[0].object_bytes * 2 + remote = _MeasuredStore(tmp_path / "remote") + uploader = ParallelSpoolUploader( + spool, + remote, + ParallelUploadConfig(max_workers=4, max_in_flight_bytes=budget), + ) + + result = uploader.upload_pending() + + assert len(result.refs) == 4 + assert not result.failures + assert result.snapshot.peak_active_uploads == 2 + assert result.snapshot.peak_in_flight_bytes <= budget + assert remote.peak_active == 2 + assert spool.snapshot().entries == 0 + + +def test_parallel_uploader_retries_transient_errors_without_losing_spool_files( + tmp_path: Path, +): + spool = DurablePackSpool(tmp_path / "spool", max_bytes=1024 * 1024) + staged = [_stage(spool, index) for index in range(2)] + remote = _MeasuredStore(tmp_path / "remote", fail_once=True) + events = [] + uploader = ParallelSpoolUploader( + spool, + remote, + ParallelUploadConfig( + max_workers=2, + max_in_flight_bytes=sum(item.object_bytes for item in staged), + max_attempts=2, + base_backoff_seconds=0, + ), + event_callback=events.append, + sleep=lambda _: None, + ) + + result = uploader.upload_pending() + + assert len(result.refs) == 2 + assert result.snapshot.retries == 2 + assert result.snapshot.uploaded_packs == 2 + assert [event.event for event in events].count("pack_upload_retry") == 2 + assert spool.snapshot().entries == 0 + + +def test_parallel_uploader_serializes_event_callbacks(tmp_path: Path): + spool = DurablePackSpool(tmp_path / "spool", max_bytes=1024 * 1024) + staged = [_stage(spool, index) for index in range(4)] + remote = _MeasuredStore(tmp_path / "remote") + lock = threading.Lock() + active = 0 + peak_active = 0 + + def callback(_event): + nonlocal active, peak_active + with lock: + active += 1 + peak_active = max(peak_active, active) + time.sleep(0.005) + with lock: + active -= 1 + + result = ParallelSpoolUploader( + spool, + remote, + ParallelUploadConfig( + max_workers=4, + max_in_flight_bytes=sum(item.object_bytes for item in staged), + ), + event_callback=callback, + ).upload_pending() + + assert not result.failures + assert peak_active == 1 + + +def test_parallel_uploader_reports_permanent_failure_and_keeps_ready_pack( + tmp_path: Path, +): + spool = DurablePackSpool(tmp_path / "spool", max_bytes=1024 * 1024) + staged = _stage(spool, 1) + remote = _MeasuredStore(tmp_path / "remote", fail_once=True) + uploader = ParallelSpoolUploader( + spool, + remote, + ParallelUploadConfig( + max_workers=1, + max_in_flight_bytes=staged.object_bytes, + max_attempts=1, + ), + sleep=lambda _: None, + ) + + result = uploader.upload_pending() + + assert not result.refs + assert len(result.failures) == 1 + assert result.failures[0].pack_id == staged.pack_id + assert result.snapshot.failed_packs == 1 + assert staged.path.exists() + + +def test_parallel_uploader_rejects_a_pack_larger_than_its_byte_budget( + tmp_path: Path, +): + spool = DurablePackSpool(tmp_path / "spool", max_bytes=1024 * 1024) + staged = _stage(spool, 1) + uploader = ParallelSpoolUploader( + spool, + FilesystemPackStore(tmp_path / "remote"), + ParallelUploadConfig( + max_workers=1, max_in_flight_bytes=staged.object_bytes - 1 + ), + ) + + with pytest.raises(ValueError, match="in-flight byte limit"): + uploader.upload_pending() diff --git a/tests/test_capture_pipeline.py b/tests/test_capture_pipeline.py new file mode 100644 index 000000000..5dbeaafda --- /dev/null +++ b/tests/test_capture_pipeline.py @@ -0,0 +1,370 @@ +from __future__ import annotations + +from dataclasses import replace +import itertools +from pathlib import Path +import threading +from uuid import UUID + +import pytest + +from dmi.storage.capture import ( + AdmissionResult, + BoundedRecordQueue, + CaptureMetadata, + CaptureRecord, + DirectPackSink, + DuplicateCaptureError, + FilesystemPackStore, + FlushReason, + HostCapturePipeline, + OverloadPolicy, + PackAssembler, + PipelineConfig, + PipelineFailedError, + object_key_for, +) + + +pytestmark = pytest.mark.cpu + + +def _metadata(capture_id: str, *, session_id: str = "session-a") -> CaptureMetadata: + return CaptureMetadata( + capture_id=capture_id, + tenant_id="tenant/a", + experiment_id="exp-a", + run_id="run-a", + session_id=session_id, + request_id="request-a", + sequence_id="sequence-a", + model_id="model-a", + model_revision="revision-a", + adapter_revision=None, + capture_policy_version="policy-v1", + hook_name="resid_pre", + layer_number=3, + producer_rank=0, + step_number=0, + token_start=0, + token_end=1, + batch_position=0, + dtype="float32", + shape=(2,), + captured_at_ns=1_700_000_000_000_000_000, + ) + + +def _record(capture_id: str, *, session_id: str = "session-a") -> CaptureRecord: + return CaptureRecord( + metadata=_metadata(capture_id, session_id=session_id), + payload=b"\x00\x00\x80?\x00\x00\x00@", + ) + + +def _ids(): + values = iter( + ( + UUID("018f0000-0000-7000-8000-000000000001"), + UUID("018f0000-0000-7000-8000-000000000002"), + UUID("018f0000-0000-7000-8000-000000000003"), + UUID("018f0000-0000-7000-8000-000000000004"), + UUID("018f0000-0000-7000-8000-000000000005"), + ) + ) + return lambda: next(values) + + +def test_bounded_queue_enforces_record_and_byte_limits(): + queue = BoundedRecordQueue(max_records=1, max_bytes=8) + + assert queue.put(_record("capture-a"), policy=OverloadPolicy.DROP_NEWEST) \ + is AdmissionResult.ACCEPTED + assert queue.put(_record("capture-b"), policy=OverloadPolicy.DROP_NEWEST) \ + is AdmissionResult.DROPPED + snapshot = queue.snapshot() + + assert snapshot.records == 1 + assert snapshot.bytes == 8 + assert snapshot.peak_records == 1 + assert snapshot.peak_bytes == 8 + assert queue.get(timeout=0) == _record("capture-a") + + +def test_bounded_queue_reports_oversized_timeout_and_close(): + queue = BoundedRecordQueue(max_records=1, max_bytes=8) + + assert queue.put( + replace(_record("capture-a"), payload=b"\x00" * 8), + policy=OverloadPolicy.BLOCK, + ) is AdmissionResult.ACCEPTED + assert queue.put( + _record("capture-b"), policy=OverloadPolicy.BLOCK, timeout=0 + ) is AdmissionResult.TIMED_OUT + oversized = CaptureRecord( + metadata=replace(_metadata("capture-c"), shape=(3,)), + payload=b"\x00" * 12, + ) + assert queue.put(oversized, policy=OverloadPolicy.DROP_NEWEST) \ + is AdmissionResult.TOO_LARGE + + queue.close() + assert queue.put(_record("capture-d"), policy=OverloadPolicy.DROP_NEWEST) \ + is AdmissionResult.CLOSED + + +def test_bounded_queue_stays_within_limits_under_sustained_overload(): + queue = BoundedRecordQueue(max_records=3, max_bytes=24) + + results = [ + queue.put(_record(f"capture-{index}"), policy=OverloadPolicy.DROP_NEWEST) + for index in range(100) + ] + + assert results.count(AdmissionResult.ACCEPTED) == 3 + assert results.count(AdmissionResult.DROPPED) == 97 + assert queue.snapshot().peak_records == 3 + assert queue.snapshot().peak_bytes == 24 + + +def test_bounded_queue_rejects_an_unknown_overload_policy(): + queue = BoundedRecordQueue(max_records=1, max_bytes=8) + queue.put(_record("capture-a"), policy=OverloadPolicy.DROP_NEWEST) + + with pytest.raises(ValueError, match="overload policy"): + queue.put(_record("capture-b"), policy="unknown", timeout=0) # type: ignore[arg-type] + + +def test_pipeline_config_rejects_an_unknown_overload_policy(): + with pytest.raises(ValueError, match="overload policy"): + PipelineConfig( + max_queue_records=2, + max_queue_bytes=16, + max_pack_bytes=1024 * 1024, + max_pack_records=2, + max_linger_ns=100, + overload_policy="drop", # type: ignore[arg-type] + ) + + +def test_pack_assembler_seals_on_size_linger_session_and_shutdown(): + first = _record("capture-a") + probe = PackAssembler( + max_pack_bytes=1024 * 1024, + max_records=10, + max_linger_ns=100, + pack_id_factory=_ids(), + ) + probe.append(first, now_ns=0) + one_pack_bytes = len(probe.flush(FlushReason.SHUTDOWN)[0].pack.data) + + assembler = PackAssembler( + max_pack_bytes=one_pack_bytes + 16, + max_records=10, + max_linger_ns=100, + pack_id_factory=_ids(), + ) + assert assembler.append(first, now_ns=0) == () + size_flush = assembler.append(_record("capture-b"), now_ns=10) + assert size_flush[0].reason is FlushReason.SIZE + assert size_flush[0].pack.record_count == 1 + + linger_flush = assembler.flush_expired(now_ns=110) + assert linger_flush[0].reason is FlushReason.LINGER + assert linger_flush[0].pack.record_count == 1 + + assert assembler.append(_record("capture-c"), now_ns=200) == () + session_flush = assembler.append( + _record("capture-d", session_id="session-b"), now_ns=210 + ) + assert session_flush[0].reason is FlushReason.SESSION + shutdown_flush = assembler.flush(FlushReason.SHUTDOWN) + assert shutdown_flush[0].reason is FlushReason.SHUTDOWN + + +def test_pack_assembler_rejects_duplicates_without_losing_the_open_pack(): + assembler = PackAssembler( + max_pack_bytes=1024 * 1024, + max_records=10, + max_linger_ns=100, + pack_id_factory=_ids(), + ) + assembler.append(_record("capture-a"), now_ns=0) + + with pytest.raises(DuplicateCaptureError, match="capture-a"): + assembler.append(_record("capture-a"), now_ns=1) + + flushed = assembler.flush(FlushReason.SHUTDOWN) + assert flushed[0].pack.record_count == 1 + + +def test_pack_assembler_rejects_oversized_input_without_losing_prior_records(): + first = _record("capture-a") + probe = PackAssembler( + max_pack_bytes=1024 * 1024, + max_records=10, + max_linger_ns=100, + pack_id_factory=_ids(), + ) + probe.append(first, now_ns=0) + one_pack_bytes = len(probe.flush(FlushReason.SHUTDOWN)[0].pack.data) + assembler = PackAssembler( + max_pack_bytes=one_pack_bytes, + max_records=10, + max_linger_ns=100, + pack_id_factory=_ids(), + ) + assembler.append(first, now_ns=0) + oversized = CaptureRecord( + metadata=replace(_metadata("capture-b"), shape=(1024,)), + payload=b"\x00" * 4096, + ) + + with pytest.raises(ValueError, match="does not fit an empty pack"): + assembler.append(oversized, now_ns=1) + + flushed = assembler.flush(FlushReason.SHUTDOWN) + assert flushed[0].pack.record_count == 1 + + +def test_object_key_bounds_long_metadata_components(tmp_path: Path): + record = replace( + _record("capture-a"), + metadata=replace( + _metadata("capture-a"), + tenant_id="a" * 512, + session_id="b" * 512, + ), + ) + assembler = PackAssembler( + max_pack_bytes=1024 * 1024, + max_records=2, + max_linger_ns=100, + pack_id_factory=_ids(), + ) + assembler.append(record, now_ns=0) + ready = assembler.flush(FlushReason.SHUTDOWN)[0] + + key = object_key_for(ready) + ref = FilesystemPackStore(tmp_path, store_id="local").put(ready.pack, key) + + assert max(len(part.encode()) for part in key.split("/")) <= 200 + assert ref.object_key == key + + +def test_direct_pipeline_persists_on_close_and_reports_bounded_metrics( + tmp_path: Path, +): + store = FilesystemPackStore(tmp_path / "objects", store_id="local") + sink = DirectPackSink(store) + pipeline = HostCapturePipeline( + PipelineConfig( + max_queue_records=4, + max_queue_bytes=32, + max_pack_bytes=1024 * 1024, + max_pack_records=4, + max_linger_ns=1_000_000_000, + ), + sink, + pack_id_factory=_ids(), + ) + + pipeline.start() + assert pipeline.submit(_record("capture-a")) is AdmissionResult.ACCEPTED + assert pipeline.submit(_record("capture-b")) is AdmissionResult.ACCEPTED + snapshot = pipeline.close(timeout=2) + + assert snapshot.admitted_records == 2 + assert snapshot.persisted_records == 2 + assert snapshot.packs_persisted == 1 + assert snapshot.queue_peak_bytes <= 32 + assert snapshot.queue_peak_records <= 4 + assert snapshot.flush_shutdown == 1 + assert snapshot.failures == 0 + assert snapshot.admission_duration.count == 2 + assert snapshot.persist_duration.count == 1 + assert sink.last_ref is not None + assert "%2F" in sink.last_ref.object_key + assert store.stat(sink.last_ref).size == snapshot.packed_bytes + + +class _FailingSink: + def persist(self, ready): + raise OSError("sink unavailable") + + +def test_pipeline_surfaces_sink_failure_and_closes_admission(): + pipeline = HostCapturePipeline( + PipelineConfig( + max_queue_records=2, + max_queue_bytes=16, + max_pack_bytes=1024 * 1024, + max_pack_records=2, + max_linger_ns=1_000_000_000, + ), + _FailingSink(), + pack_id_factory=_ids(), + ) + pipeline.start() + pipeline.submit(_record("capture-a")) + + with pytest.raises(PipelineFailedError, match="pipeline failed"): + pipeline.close(timeout=2) + + assert pipeline.snapshot().failures == 1 + + +def test_pipeline_is_not_failed_by_an_observability_callback(tmp_path: Path): + def broken_callback(event): + raise RuntimeError("telemetry unavailable") + + pipeline = HostCapturePipeline( + PipelineConfig( + max_queue_records=2, + max_queue_bytes=16, + max_pack_bytes=1024 * 1024, + max_pack_records=2, + max_linger_ns=100, + ), + DirectPackSink(FilesystemPackStore(tmp_path, store_id="local")), + pack_id_factory=_ids(), + event_callback=broken_callback, + ) + pipeline.start() + pipeline.submit(_record("capture-a")) + + snapshot = pipeline.close(timeout=2) + + assert snapshot.failures == 0 + assert snapshot.event_callback_failures == 1 + + +def test_pipeline_flushes_an_idle_pack_at_its_linger_deadline(tmp_path: Path): + persisted = threading.Event() + ticks = itertools.count(start=0, step=10) + + def capture_event(event): + if event.event == "pack_persisted": + persisted.set() + + pipeline = HostCapturePipeline( + PipelineConfig( + max_queue_records=2, + max_queue_bytes=16, + max_pack_bytes=1024 * 1024, + max_pack_records=2, + max_linger_ns=1, + ), + DirectPackSink(FilesystemPackStore(tmp_path, store_id="local")), + pack_id_factory=_ids(), + clock_ns=lambda: next(ticks), + event_callback=capture_event, + ) + pipeline.start() + pipeline.submit(_record("capture-a")) + + assert persisted.wait(timeout=1) + snapshot = pipeline.close(timeout=2) + + assert snapshot.flush_linger == 1 + assert snapshot.flush_shutdown == 0 diff --git a/tests/test_capture_pipeline_benchmark.py b/tests/test_capture_pipeline_benchmark.py new file mode 100644 index 000000000..11286a1f6 --- /dev/null +++ b/tests/test_capture_pipeline_benchmark.py @@ -0,0 +1,74 @@ +from __future__ import annotations + +import json + +import pytest + +from benchmarks.bench_capture_pipeline import ( + PipelineBenchmarkConfig, + main, + run_trial, +) + + +pytestmark = pytest.mark.cpu + + +def test_pipeline_benchmark_config_rejects_unbounded_inputs(): + with pytest.raises(ValueError, match="records"): + PipelineBenchmarkConfig(records=0) + with pytest.raises(ValueError, match="queue_records"): + PipelineBenchmarkConfig(queue_records=0) + with pytest.raises(ValueError, match="queue_bytes"): + PipelineBenchmarkConfig(payload_bytes=64, queue_bytes=32) + + +@pytest.mark.parametrize("mode", ("direct", "spool")) +def test_pipeline_trial_persists_and_verifies_every_record(mode: str): + trial = run_trial( + PipelineBenchmarkConfig( + mode=mode, + records=12, + payload_bytes=64, + target_pack_bytes=4096, + queue_records=4, + queue_bytes=256, + trials=1, + ) + ) + + assert trial.record_count == 12 + assert trial.persisted_records == 12 + assert trial.dropped_records == 0 + assert trial.packs_persisted > 0 + assert trial.queue_peak_records <= 4 + assert trial.queue_peak_bytes <= 256 + assert trial.seconds > 0 + assert trial.as_dict()["logical_gib_per_second"] > 0 + + +def test_pipeline_benchmark_dry_run_emits_resolved_workload(capsys): + assert main( + [ + "--mode", + "spool", + "--records", + "7", + "--payload-bytes", + "64", + "--target-pack-bytes", + "4096", + "--queue-records", + "2", + "--queue-bytes", + "128", + "--trials", + "1", + "--dry-run", + ] + ) == 0 + + result = json.loads(capsys.readouterr().out) + assert result["dry_run"] is True + assert result["config"]["mode"] == "spool" + assert result["config"]["records"] == 7 diff --git a/tests/test_capture_s3.py b/tests/test_capture_s3.py new file mode 100644 index 000000000..a53e37352 --- /dev/null +++ b/tests/test_capture_s3.py @@ -0,0 +1,290 @@ +from __future__ import annotations + +from io import BytesIO +from uuid import UUID + +import pytest + +from dmi.storage.capture import ( + CaptureMetadata, + CaptureRecord, + PackConflictError, + PackFormatError, + PackIntegrityError, + PackRef, + PackWriter, + S3PackStore, + S3StoreConfig, +) + + +pytestmark = pytest.mark.cpu + + +class _ClientError(Exception): + def __init__(self, status: int, code: str): + self.response = { + "Error": {"Code": code}, + "ResponseMetadata": {"HTTPStatusCode": status}, + } + super().__init__(code) + + +class _Body(BytesIO): + def close(self) -> None: + super().close() + + +class _S3Client: + def __init__(self): + self.objects: dict[str, tuple[bytes, dict[str, str]]] = {} + self.upload_configs: list[object] = [] + + def head_object(self, *, Bucket: str, Key: str): + try: + data, metadata = self.objects[Key] + except KeyError as exc: + raise _ClientError(404, "NoSuchKey") from exc + return {"ContentLength": len(data), "Metadata": metadata} + + def upload_fileobj( + self, + Fileobj, + Bucket: str, + Key: str, + ExtraArgs: dict[str, object], + Config: object, + ) -> None: + self.upload_configs.append(Config) + self.objects[Key] = (Fileobj.read(), dict(ExtraArgs["Metadata"])) + + def get_object(self, *, Bucket: str, Key: str, Range: str): + start, end = (int(value) for value in Range.removeprefix("bytes=").split("-")) + data = self.objects[Key][0][start : end + 1] + return {"Body": _Body(data), "ContentLength": len(data)} + + def list_objects_v2(self, **request): + keys = sorted( + key for key in self.objects if key.startswith(request.get("Prefix", "")) + ) + start = int(request.get("ContinuationToken", "0")) + limit = request["MaxKeys"] + selected = keys[start : start + limit] + next_index = start + len(selected) + return { + "Contents": [ + {"Key": key, "Size": len(self.objects[key][0])} + for key in selected + ], + "IsTruncated": next_index < len(keys), + **( + {"NextContinuationToken": str(next_index)} + if next_index < len(keys) + else {} + ), + } + + +def _pack(pack_id: str = "018f0000-0000-7000-8000-000000000001"): + metadata = CaptureMetadata( + capture_id="capture-a", + tenant_id="tenant-a", + experiment_id="exp-a", + run_id="run-a", + session_id="session-a", + request_id="request-a", + sequence_id="sequence-a", + model_id="model-a", + model_revision="revision-a", + adapter_revision=None, + capture_policy_version="policy-v1", + hook_name="resid_pre", + layer_number=3, + producer_rank=0, + step_number=0, + token_start=0, + token_end=1, + batch_position=0, + dtype="float32", + shape=(2,), + captured_at_ns=1_700_000_000_000_000_000, + ) + writer = PackWriter( + pack_id=UUID(pack_id), + created_at_ns=metadata.captured_at_ns, + max_pack_bytes=1024 * 1024, + ) + writer.append(CaptureRecord(metadata=metadata, payload=b"abcdefgh")) + return writer.seal() + + +def _store(client: _S3Client) -> S3PackStore: + return S3PackStore( + client, + bucket="captures", + store_id="garage-local", + transfer_config="bounded-transfer", + ) + + +def test_s3_config_rejects_insecure_remote_endpoint_and_hides_secrets(): + with pytest.raises(ValueError, match="allow_insecure_http"): + S3StoreConfig( + endpoint_url="http://garage.example.com", + bucket="captures", + region="garage", + access_key_id="access", + secret_access_key="secret", + ) + + config = S3StoreConfig( + endpoint_url="http://127.0.0.1:3900", + bucket="captures", + region="garage", + access_key_id="access", + secret_access_key="secret", + allow_insecure_http=True, + ) + + assert "secret" not in repr(config) + assert config.multipart_chunk_bytes >= 5 * 1024**2 + + with pytest.raises(ValueError, match="origin"): + S3StoreConfig( + endpoint_url="https://garage.example.com/api", + bucket="captures", + region="garage", + ) + with pytest.raises(ValueError, match="credentials"): + S3StoreConfig( + endpoint_url="https://garage.example.com", + bucket="captures", + region="garage", + access_key_id="", + secret_access_key="secret", + ) + + +def test_s3_put_is_retry_safe_and_uses_checksum_metadata(): + client = _S3Client() + store = _store(client) + pack = _pack() + key = "v1/tenant=tenant-a/pack.dmi-pack" + + first = store.put(pack, key) + second = store.put(pack, key) + + assert first == second + assert len(client.upload_configs) == 1 + assert client.objects[key][1]["dmi-sha256"] == pack.checksum + assert store.stat(first).checksum == pack.checksum + + +def test_s3_discovers_pack_reference_from_object_metadata(): + client = _S3Client() + store = _store(client) + pack = _pack() + expected = store.put(pack, "packs/a.dmi-pack") + + assert store.inspect("packs/a.dmi-pack") == expected + + data, metadata = client.objects["packs/a.dmi-pack"] + client.objects["packs/a.dmi-pack"] = (data, {**metadata, "dmi-sha256": "bad"}) + with pytest.raises(PackIntegrityError, match="metadata"): + store.inspect("packs/a.dmi-pack") + + +def test_s3_put_rejects_an_existing_key_with_different_metadata(): + client = _S3Client() + store = _store(client) + pack = _pack() + key = "v1/tenant=tenant-a/pack.dmi-pack" + client.objects[key] = (b"different", {"dmi-sha256": "0" * 64}) + + with pytest.raises(PackConflictError, match="different content"): + store.put(pack, key) + + +def test_s3_stat_rejects_missing_integrity_metadata(): + client = _S3Client() + store = _store(client) + pack = _pack() + key = "v1/tenant=tenant-a/pack.dmi-pack" + client.objects[key] = (pack.data, {}) + ref = PackRef( + pack_id=pack.pack_id, + store_id=store.store_id, + object_key=key, + object_bytes=len(pack.data), + checksum=pack.checksum, + record_count=pack.record_count, + ) + + with pytest.raises(PackIntegrityError, match="metadata"): + store.stat(ref) + + +def test_s3_stat_rejects_a_malformed_head_response(): + client = _S3Client() + client.head_object = lambda **_: [] + pack = _pack() + ref = PackRef( + pack_id=pack.pack_id, + store_id="garage-local", + object_key="v1/pack.dmi-pack", + object_bytes=len(pack.data), + checksum=pack.checksum, + record_count=pack.record_count, + ) + + with pytest.raises(PackIntegrityError, match="HeadObject"): + _store(client).stat(ref) + + +def test_s3_range_reads_are_exact_and_bounded(): + client = _S3Client() + store = _store(client) + pack = _pack() + ref = store.put(pack, "v1/tenant=tenant-a/pack.dmi-pack") + + assert store.read_range(ref, 10, 17) == pack.data[10:27] + assert store.read_range(ref, 0, 0) == b"" + with pytest.raises(PackFormatError, match="exceeds object size"): + store.read_range(ref, len(pack.data), 1) + + +def test_s3_listing_is_prefix_scoped_and_cursor_bounded(): + client = _S3Client() + store = _store(client) + for index in range(3): + client.objects[f"v1/day=2026-08-25/{index}.dmi-pack"] = ( + bytes(index + 1), + {"dmi-sha256": str(index) * 64}, + ) + client.objects["other/ignored.dmi-pack"] = (b"x", {"dmi-sha256": "f" * 64}) + + first = store.list_objects(prefix="v1/day=2026-08-25/", limit=2) + second = store.list_objects( + prefix="v1/day=2026-08-25/", cursor=first.next_cursor, limit=2 + ) + + assert [item.object_key for item in first.items] == [ + "v1/day=2026-08-25/0.dmi-pack", + "v1/day=2026-08-25/1.dmi-pack", + ] + assert first.next_cursor == "2" + assert [item.object_key for item in second.items] == [ + "v1/day=2026-08-25/2.dmi-pack" + ] + assert second.next_cursor is None + + +def test_s3_listing_rejects_a_truncated_page_without_a_cursor(): + client = _S3Client() + client.list_objects_v2 = lambda **_: { + "Contents": [], + "IsTruncated": True, + } + + with pytest.raises(PackIntegrityError, match="cursor"): + _store(client).list_objects(limit=1) diff --git a/tests/test_capture_spool.py b/tests/test_capture_spool.py new file mode 100644 index 000000000..0bd1c83e9 --- /dev/null +++ b/tests/test_capture_spool.py @@ -0,0 +1,194 @@ +from __future__ import annotations + +from pathlib import Path +from uuid import UUID + +import pytest + +from dmi.storage.capture import ( + CaptureMetadata, + CaptureRecord, + DurablePackSink, + DurablePackSpool, + FilesystemPackStore, + FlushReason, + PackIntegrityError, + PackWriter, + ReadyPack, + SpoolFullError, + SpoolUploader, + HostCapturePipeline, + PipelineConfig, + AdmissionResult, +) + + +pytestmark = pytest.mark.cpu + + +def _sealed(pack_id: str, capture_id: str): + metadata = CaptureMetadata( + capture_id=capture_id, + tenant_id="tenant-a", + experiment_id="exp-a", + run_id="run-a", + session_id="session-a", + request_id="request-a", + sequence_id="sequence-a", + model_id="model-a", + model_revision="revision-a", + adapter_revision=None, + capture_policy_version="policy-v1", + hook_name="resid_pre", + layer_number=3, + producer_rank=0, + step_number=0, + token_start=0, + token_end=1, + batch_position=0, + dtype="float32", + shape=(2,), + captured_at_ns=1_700_000_000_000_000_000, + ) + writer = PackWriter( + pack_id=UUID(pack_id), + created_at_ns=metadata.captured_at_ns, + max_pack_bytes=1024 * 1024, + ) + writer.append(CaptureRecord(metadata=metadata, payload=b"\x00" * 8)) + return writer.seal(), metadata + + +class _AmbiguousStore(FilesystemPackStore): + def __init__(self, root: Path): + super().__init__(root, store_id="remote") + self.fail_once = True + + def put(self, pack, object_key): + ref = super().put(pack, object_key) + if self.fail_once: + self.fail_once = False + raise OSError("connection lost after commit") + return ref + + +def test_spool_stage_is_atomic_bounded_and_restart_recoverable(tmp_path: Path): + first, metadata = _sealed( + "018f0000-0000-7000-8000-000000000001", "capture-a" + ) + second, _ = _sealed( + "018f0000-0000-7000-8000-000000000002", "capture-b" + ) + spool = DurablePackSpool(tmp_path / "spool", max_bytes=len(first.data)) + sink = DurablePackSink(spool) + + staged = sink.persist(ReadyPack(first, metadata, FlushReason.SHUTDOWN)) + + assert staged.path.suffix == ".ready" + assert staged.created_at_ns == first.created_at_ns + assert staged.path.exists() + assert not tuple((tmp_path / "spool").rglob("*.open")) + assert spool.snapshot().bytes == len(first.data) + with pytest.raises(SpoolFullError, match="spool byte limit"): + sink.persist(ReadyPack(second, metadata, FlushReason.SHUTDOWN)) + + recovered = DurablePackSpool( + tmp_path / "spool", max_bytes=len(first.data) + ).recover() + assert recovered == (staged,) + assert recovered[0].created_at_ns == first.created_at_ns + + +def test_spool_rejects_an_object_key_it_cannot_recover_exactly(tmp_path: Path): + pack, _ = _sealed( + "018f0000-0000-7000-8000-000000000001", "capture-a" + ) + spool = DurablePackSpool(tmp_path / "spool", max_bytes=len(pack.data) * 2) + + with pytest.raises(ValueError, match="pack ID"): + spool.stage(pack, "v1/custom-name.dmi-pack") + + +def test_spool_upload_retry_resolves_an_ambiguous_remote_commit(tmp_path: Path): + pack, metadata = _sealed( + "018f0000-0000-7000-8000-000000000001", "capture-a" + ) + spool = DurablePackSpool(tmp_path / "spool", max_bytes=len(pack.data) * 2) + staged = DurablePackSink(spool).persist( + ReadyPack(pack, metadata, FlushReason.SHUTDOWN) + ) + remote = _AmbiguousStore(tmp_path / "objects") + uploader = SpoolUploader(spool, remote) + + with pytest.raises(OSError, match="after commit"): + uploader.upload(staged) + assert staged.path.exists() + + ref = uploader.upload(staged) + + assert not staged.path.exists() + assert remote.stat(ref).checksum == pack.checksum + assert spool.snapshot().bytes == 0 + + +def test_spool_recovery_rejects_corrupt_ready_pack(tmp_path: Path): + pack, metadata = _sealed( + "018f0000-0000-7000-8000-000000000001", "capture-a" + ) + root = tmp_path / "spool" + spool = DurablePackSpool(root, max_bytes=len(pack.data) * 2) + staged = DurablePackSink(spool).persist( + ReadyPack(pack, metadata, FlushReason.SHUTDOWN) + ) + with staged.path.open("r+b") as handle: + handle.seek(64) + handle.write(b"\xff") + + with pytest.raises(PackIntegrityError, match="checksum"): + DurablePackSpool(root, max_bytes=len(pack.data) * 2).recover() + + assert staged.path.exists() + + +def test_spool_recovery_removes_incomplete_open_files(tmp_path: Path): + root = tmp_path / "spool" + root.mkdir() + incomplete = root / ".interrupted.open" + incomplete.write_bytes(b"partial") + + spool = DurablePackSpool(root, max_bytes=1024) + assert spool.snapshot().bytes == len(b"partial") + + recovered = spool.recover() + + assert recovered == () + assert not incomplete.exists() + assert spool.snapshot().bytes == 0 + + +def test_durable_pipeline_commits_to_spool_without_remote_storage(tmp_path: Path): + pack, metadata = _sealed( + "018f0000-0000-7000-8000-000000000001", "capture-a" + ) + record = CaptureRecord(metadata=metadata, payload=b"\x00" * 8) + spool = DurablePackSpool(tmp_path / "spool", max_bytes=len(pack.data) * 2) + pipeline = HostCapturePipeline( + PipelineConfig( + max_queue_records=2, + max_queue_bytes=16, + max_pack_bytes=1024 * 1024, + max_pack_records=2, + max_linger_ns=1_000_000_000, + ), + DurablePackSink(spool), + pack_id_factory=lambda: UUID( + "018f0000-0000-7000-8000-000000000003" + ), + ) + + pipeline.start() + assert pipeline.submit(record) is AdmissionResult.ACCEPTED + snapshot = pipeline.close(timeout=2) + + assert snapshot.persisted_records == 1 + assert len(spool.recover()) == 1 diff --git a/tests/test_capture_storage.py b/tests/test_capture_storage.py new file mode 100644 index 000000000..09f6e69c8 --- /dev/null +++ b/tests/test_capture_storage.py @@ -0,0 +1,387 @@ +from __future__ import annotations + +from dataclasses import replace +from pathlib import Path +from uuid import UUID + +import pytest + +from dmi.storage.capture import ( + CaptureCatalog, + CaptureMetadata, + CapturePage, + CaptureQuery, + CaptureReader, + CaptureRecord, + DuplicateCaptureError, + FilesystemPackStore, + HydrationLimitError, + PackConflictError, + PackFormatError, + PackIntegrityError, + PackIndex, + PackReader, + PackWriter, +) + + +pytestmark = pytest.mark.cpu + + +PACK_ID = UUID("018f0000-0000-7000-8000-000000000001") + + +def _metadata(capture_id: str, *, step: int = 0) -> CaptureMetadata: + return CaptureMetadata( + capture_id=capture_id, + tenant_id="tenant-a", + experiment_id="exp-a", + run_id="run-a", + session_id="session-a", + request_id="request-a", + sequence_id="sequence-a", + model_id="model-a", + model_revision="revision-a", + adapter_revision=None, + capture_policy_version="policy-v1", + hook_name="resid_pre", + layer_number=3, + producer_rank=0, + step_number=step, + token_start=step, + token_end=step + 1, + batch_position=0, + dtype="float32", + shape=(2,), + captured_at_ns=1_700_000_000_000_000_000 + step, + ) + + +def _record(capture_id: str, payload: bytes, *, step: int = 0) -> CaptureRecord: + return CaptureRecord(metadata=_metadata(capture_id, step=step), payload=payload) + + +def _sealed_pack(*records: CaptureRecord): + writer = PackWriter( + pack_id=PACK_ID, + created_at_ns=1_700_000_000_000_000_000, + max_pack_bytes=1024 * 1024, + ) + for record in records: + writer.append(record) + return writer.seal() + + +def test_pack_round_trip_preserves_analysis_coordinates_and_payloads(): + first = _record("capture-a", b"\x00\x00\x80?\x00\x00\x00@") + second = _record("capture-b", b"\x00\x00@@\x00\x00\x80@", step=1) + + sealed = _sealed_pack(first, second) + reader = PackReader.from_bytes(sealed.data) + descriptors = reader.descriptors(store_id="local", object_key="packs/a.dmi-pack") + + assert reader.pack_id == str(PACK_ID) + assert [item.capture_id for item in descriptors] == ["capture-a", "capture-b"] + assert descriptors[0].metadata.model_revision == "revision-a" + assert descriptors[1].metadata.step_number == 1 + assert reader.read_payload(descriptors[0]) == first.payload + assert reader.read_payload(descriptors[1]) == second.payload + + +def test_pack_rejects_truncation_and_unknown_major_version(): + sealed = _sealed_pack(_record("capture-a", b"\x00" * 8)) + + with pytest.raises(PackFormatError, match="truncated"): + PackReader.from_bytes(sealed.data[:-1]) + + changed = bytearray(sealed.data) + changed[8:10] = (99).to_bytes(2, "little") + with pytest.raises(PackFormatError, match="major version"): + PackReader.from_bytes(changed) + + +def test_pack_rejects_a_record_when_its_footer_would_exceed_the_limit(): + record = _record("capture-a", b"\x00" * 8) + sealed = _sealed_pack(record) + writer = PackWriter( + pack_id=PACK_ID, + created_at_ns=1_700_000_000_000_000_000, + max_pack_bytes=len(sealed.data) - 1, + ) + + with pytest.raises(ValueError, match="max_pack_bytes"): + writer.append(record) + + assert writer.record_count == 0 + + +def test_pack_writer_rejects_duplicate_capture_ids(): + writer = PackWriter( + pack_id=PACK_ID, + created_at_ns=1_700_000_000_000_000_000, + max_pack_bytes=1024 * 1024, + ) + writer.append(_record("capture-a", b"\x00" * 8)) + + with pytest.raises(DuplicateCaptureError, match="duplicate capture ID"): + writer.append(_record("capture-a", b"\x01" * 8, step=1)) + + assert writer.record_count == 1 + + +def test_pack_detects_payload_and_footer_corruption(): + sealed = _sealed_pack(_record("capture-a", b"\x00" * 8)) + reader = PackReader.from_bytes(sealed.data) + descriptor = reader.descriptors(store_id="local", object_key="a")[0] + + payload_corruption = bytearray(sealed.data) + payload_corruption[descriptor.locator.offset] ^= 0xFF + with pytest.raises(PackIntegrityError, match="pack checksum"): + PackReader.from_bytes(payload_corruption) + + footer_corruption = bytearray(sealed.data) + footer_corruption[sealed.footer_offset] ^= 0x01 + with pytest.raises(PackIntegrityError, match="pack checksum"): + PackReader.from_bytes(footer_corruption) + + +def test_capture_record_rejects_shape_payload_mismatch(): + with pytest.raises(ValueError, match="payload length"): + CaptureRecord(metadata=_metadata("capture-a"), payload=b"\x00" * 4) + + +def test_capture_metadata_rejects_boolean_numeric_fields(): + with pytest.raises(ValueError, match="step_number"): + replace(_metadata("capture-a"), step_number=True) + + with pytest.raises(ValueError, match="shape dimensions"): + replace(_metadata("capture-a"), shape=(True,)) + + +def test_filesystem_store_is_idempotent_and_rejects_conflicts(tmp_path: Path): + store = FilesystemPackStore(tmp_path, store_id="local") + sealed = _sealed_pack(_record("capture-a", b"\x00" * 8)) + + first = store.put(sealed, "tenant=a/date=2026-08-25/a.dmi-pack") + second = store.put(sealed, "tenant=a/date=2026-08-25/a.dmi-pack") + + assert first == second + assert store.stat(first).size == len(sealed.data) + assert store.read_range(first, 0, 8) == sealed.data[:8] + + conflicting = PackWriter( + pack_id=UUID("018f0000-0000-7000-8000-000000000002"), + created_at_ns=1, + max_pack_bytes=1024 * 1024, + ) + conflicting.append(_record("capture-b", b"\x00" * 8)) + with pytest.raises(PackConflictError, match="different content"): + store.put(conflicting.seal(), first.object_key) + + +def test_filesystem_store_rejects_a_symlinked_object(tmp_path: Path): + store = FilesystemPackStore(tmp_path / "objects", store_id="local") + sealed = _sealed_pack(_record("capture-a", b"\x00" * 8)) + ref = store.put(sealed, "packs/a.dmi-pack") + object_path = store.root / ref.object_key + outside = tmp_path / "outside.dmi-pack" + outside.write_bytes(sealed.data) + object_path.unlink() + object_path.symlink_to(outside) + + with pytest.raises(PackFormatError, match="regular file"): + store.read_range(ref, 0, 8) + + +def test_pack_index_reads_only_the_trailer_and_footer(tmp_path: Path): + store = _RecordingStore(tmp_path, store_id="local") + sealed = _sealed_pack( + _record("capture-a", b"\x00" * 8), + _record("capture-b", b"\x01" * 8, step=1), + ) + ref = store.put(sealed, "packs/a.dmi-pack") + + index = PackIndex.from_store(store, ref) + descriptors = index.descriptors() + + assert [item.capture_id for item in descriptors] == ["capture-a", "capture-b"] + assert store.ranges == [ + (len(sealed.data) - PackIndex.trailer_size(), PackIndex.trailer_size()), + (sealed.footer_offset, len(sealed.data) - PackIndex.trailer_size() - sealed.footer_offset), + ] + assert sum(length for _, length in store.ranges) < len(sealed.data) - 16 + + +@pytest.mark.parametrize("key", ("../escape", "/absolute", "a/../../escape", "a\\escape")) +def test_filesystem_store_rejects_unsafe_object_keys(tmp_path: Path, key: str): + store = FilesystemPackStore(tmp_path, store_id="local") + sealed = _sealed_pack(_record("capture-a", b"\x00" * 8)) + + with pytest.raises(ValueError, match="object key"): + store.put(sealed, key) + + +class _Catalog(CaptureCatalog): + def __init__(self, descriptors): + self._descriptors = tuple(descriptors) + + def search(self, query: CaptureQuery) -> CapturePage: + items = tuple( + item + for item in self._descriptors + if query.model_id is None or item.metadata.model_id == query.model_id + )[: query.limit] + return CapturePage(items=items, next_cursor=None, watermark="catalog-7") + + def get_by_ids(self, capture_ids, *, watermark): + assert watermark == "catalog-7" + wanted = set(capture_ids) + return tuple(item for item in self._descriptors if item.capture_id in wanted) + + +class _RecordingStore(FilesystemPackStore): + def __init__(self, root: Path, *, store_id: str): + super().__init__(root, store_id=store_id) + self.ranges = [] + + def read_range(self, ref, offset, length): + self.ranges.append((offset, length)) + return super().read_range(ref, offset, length) + + +class _ShortReadStore(_RecordingStore): + def read_range(self, ref, offset, length): + return super().read_range(ref, offset, length)[:-1] + + +def test_selection_is_stable_and_rejects_duplicate_logical_captures(tmp_path: Path): + store = FilesystemPackStore(tmp_path, store_id="local") + sealed = _sealed_pack( + _record("capture-a", b"\x00" * 8), + _record("capture-b", b"\x01" * 8, step=1), + ) + ref = store.put(sealed, "packs/a.dmi-pack") + descriptors = PackReader.from_bytes(sealed.data).descriptors( + store_id=ref.store_id, object_key=ref.object_key + ) + reader = CaptureReader(_Catalog(descriptors), {"local": store}) + + first = reader.select(CaptureQuery(model_id="model-a", limit=10)) + second = reader.select(CaptureQuery(model_id="model-a", limit=10)) + + assert first == second + assert first.capture_ids == ("capture-a", "capture-b") + assert first.catalog_watermark == "catalog-7" + + duplicate = replace(descriptors[1], metadata=descriptors[0].metadata) + with pytest.raises(DuplicateCaptureError, match="capture-a"): + CaptureReader(_Catalog((descriptors[0], duplicate)), {"local": store}).select( + CaptureQuery(limit=10) + ) + + +def test_reader_coalesces_ranges_and_enforces_fetch_budget(tmp_path: Path): + store = _RecordingStore(tmp_path, store_id="local") + first_payload = b"\x00\x00\x80?\x00\x00\x00@" + second_payload = b"\x00\x00@@\x00\x00\x80@" + sealed = _sealed_pack( + _record("capture-a", first_payload), + _record("capture-b", second_payload, step=1), + ) + ref = store.put(sealed, "packs/a.dmi-pack") + descriptors = PackReader.from_bytes(sealed.data).descriptors( + store_id=ref.store_id, object_key=ref.object_key + ) + reader = CaptureReader( + _Catalog(descriptors), + {"local": store}, + max_coalesce_gap_bytes=64, + ) + selection = reader.select(CaptureQuery(limit=10)) + estimate = reader.estimate(selection) + + assert estimate.capture_count == 2 + assert estimate.object_count == 1 + assert estimate.request_count == 1 + assert estimate.request_bytes >= estimate.stored_bytes + + with pytest.raises(HydrationLimitError, match="byte limit"): + reader.hydrate(selection, byte_limit=estimate.request_bytes - 1) + + hydrated = reader.hydrate(selection, byte_limit=estimate.request_bytes) + + assert [item.capture_id for item in hydrated] == ["capture-a", "capture-b"] + assert [item.payload for item in hydrated] == [first_payload, second_payload] + assert store.ranges == [ + (descriptors[0].locator.offset, estimate.request_bytes), + ] + + +def test_reader_enforces_request_budget_before_fetching(tmp_path: Path): + store = _RecordingStore(tmp_path, store_id="local") + sealed = _sealed_pack( + _record("capture-a", b"\x00" * 8), + _record("capture-b", b"\x01" * 8, step=1), + ) + ref = store.put(sealed, "packs/a.dmi-pack") + descriptors = PackReader.from_bytes(sealed.data).descriptors( + store_id=ref.store_id, object_key=ref.object_key + ) + reader = CaptureReader( + _Catalog(descriptors), {"local": store}, max_coalesce_gap_bytes=0 + ) + selection = reader.select(CaptureQuery(limit=10)) + + with pytest.raises(HydrationLimitError, match="request limit"): + reader.hydrate(selection, byte_limit=16, request_limit=1) + + assert store.ranges == [] + + +def test_reader_rejects_catalog_drift_after_selection(tmp_path: Path): + store = FilesystemPackStore(tmp_path, store_id="local") + sealed = _sealed_pack(_record("capture-a", b"\x00" * 8)) + ref = store.put(sealed, "packs/a.dmi-pack") + descriptor = PackReader.from_bytes(sealed.data).descriptors( + store_id=ref.store_id, object_key=ref.object_key + )[0] + catalog = _Catalog((descriptor,)) + reader = CaptureReader(catalog, {"local": store}) + selection = reader.select(CaptureQuery(limit=10)) + catalog._descriptors = () + + with pytest.raises(PackFormatError, match="selection no longer resolves"): + reader.estimate(selection) + + +def test_reader_rejects_a_short_object_store_range(tmp_path: Path): + store = _ShortReadStore(tmp_path, store_id="local") + sealed = _sealed_pack(_record("capture-a", b"\x00" * 8)) + ref = store.put(sealed, "packs/a.dmi-pack") + descriptors = PackReader.from_bytes(sealed.data).descriptors( + store_id=ref.store_id, object_key=ref.object_key + ) + reader = CaptureReader(_Catalog(descriptors), {"local": store}) + selection = reader.select(CaptureQuery(limit=10)) + + with pytest.raises(PackIntegrityError, match="short range"): + reader.hydrate(selection, byte_limit=8) + + +def test_reader_detects_corruption_in_a_partial_range(tmp_path: Path): + store = FilesystemPackStore(tmp_path, store_id="local") + sealed = _sealed_pack(_record("capture-a", b"\x00" * 8)) + ref = store.put(sealed, "packs/a.dmi-pack") + descriptor = PackReader.from_bytes(sealed.data).descriptors( + store_id=ref.store_id, object_key=ref.object_key + )[0] + reader = CaptureReader(_Catalog((descriptor,)), {"local": store}) + selection = reader.select(CaptureQuery(limit=10)) + + path = tmp_path / ref.object_key + with path.open("r+b") as handle: + handle.seek(descriptor.locator.offset) + handle.write(b"\xff") + + with pytest.raises(PackIntegrityError, match="record checksum"): + reader.hydrate(selection, byte_limit=descriptor.locator.stored_length) diff --git a/tests/test_clickhouse_capture_catalog.py b/tests/test_clickhouse_capture_catalog.py new file mode 100644 index 000000000..f8c7b728d --- /dev/null +++ b/tests/test_clickhouse_capture_catalog.py @@ -0,0 +1,121 @@ +from __future__ import annotations + +from uuid import UUID + +import pytest + +from dmi.storage.capture import ( + CaptureMetadata, + CaptureRecord, + ClickHouseCatalogConfig, + ClickHouseCatalogWriter, + PackIndex, + PackRef, + PackWriter, +) + + +pytestmark = pytest.mark.cpu + + +class _Client: + def __init__(self): + self.calls = [] + self.committed = [] + + def execute(self, query, params=None, **kwargs): + self.calls.append((query, params, kwargs)) + if query.lstrip().startswith("SELECT"): + return self.committed + return [] + + +def _descriptor(): + metadata = CaptureMetadata( + capture_id="capture-a", + tenant_id="tenant-a", + experiment_id="exp-a", + run_id="run-a", + session_id="session-a", + request_id="request-a", + sequence_id="sequence-a", + model_id="model-a", + model_revision="revision-a", + adapter_revision=None, + capture_policy_version="policy-v1", + hook_name="resid_pre", + layer_number=3, + producer_rank=0, + step_number=0, + token_start=0, + token_end=1, + batch_position=0, + dtype="float32", + shape=(2,), + captured_at_ns=1_700_000_000_000_000_000, + ) + pack = PackWriter( + pack_id=UUID("018f0000-0000-7000-8000-000000000001"), + created_at_ns=metadata.captured_at_ns, + max_pack_bytes=1024 * 1024, + ) + pack.append(CaptureRecord(metadata, b"abcdefgh")) + sealed = pack.seal() + ref = PackRef( + sealed.pack_id, "garage", "packs/a.dmi-pack", len(sealed.data), + sealed.checksum, sealed.record_count, + ) + + class Store: + store_id = "garage" + def read_range(self, ref, offset, length): + return sealed.data[offset : offset + length] + + return ref, PackIndex.from_store(Store(), ref).descriptors()[0] + + +def test_clickhouse_catalog_creates_replay_safe_raw_tables_and_final_views(): + client = _Client() + writer = ClickHouseCatalogWriter(client, ClickHouseCatalogConfig()) + + writer.ensure_schema() + + ddl = "\n".join(call[0] for call in client.calls) + assert "ReplacingMergeTree(index_version)" in ddl + assert "FROM `default`.`dmi_capture_raw` FINAL" in ddl + assert "FROM `default`.`dmi_pack_inventory_raw` FINAL" in ddl + + +def test_clickhouse_catalog_inserts_descriptors_before_pack_commit(): + ref, descriptor = _descriptor() + client = _Client() + writer = ClickHouseCatalogWriter(client, ClickHouseCatalogConfig()) + + writer.write_descriptors([descriptor], index_version=42) + writer.commit_packs([ref], index_version=42) + + inserts = [call for call in client.calls if call[0].startswith("INSERT")] + assert "dmi_capture_raw" in inserts[0][0] + assert "dmi_pack_inventory_raw" in inserts[1][0] + assert inserts[0][1][0][0] == "capture-a" + assert inserts[1][1][0][0] == ref.pack_id + + +def test_clickhouse_catalog_queries_committed_pack_ids_in_one_batch(): + client = _Client() + client.committed = [("garage", "018f0000-0000-7000-8000-000000000001")] + writer = ClickHouseCatalogWriter(client, ClickHouseCatalogConfig()) + + found = writer.committed_pack_ids( + [("garage", "018f0000-0000-7000-8000-000000000001")] + ) + + assert found == set(client.committed) + assert "IN %(identities)s" in client.calls[0][0] + + +@pytest.mark.parametrize("name", ["bad-name", "x; DROP TABLE y", "`quoted`"]) +def test_clickhouse_catalog_rejects_unsafe_identifiers(name: str): + with pytest.raises(ValueError, match="identifier"): + ClickHouseCatalogConfig(database=name) + diff --git a/tests/test_clickhouse_catalog_live.py b/tests/test_clickhouse_catalog_live.py new file mode 100644 index 000000000..45c1255b6 --- /dev/null +++ b/tests/test_clickhouse_catalog_live.py @@ -0,0 +1,60 @@ +from __future__ import annotations + +from os import environ +from uuid import uuid4 + +import pytest + +from benchmarks.bench_capture_catalog import synthetic_descriptors +from dmi.storage.capture import ClickHouseCatalogConfig, ClickHouseCatalogWriter + + +pytestmark = [pytest.mark.manual, pytest.mark.clickhouse] + + +def test_duplicate_catalog_replay_is_logically_deduplicated(): + clickhouse_driver = pytest.importorskip("clickhouse_driver") + client = clickhouse_driver.Client( + host=environ.get("DMI_CLICKHOUSE_HOST", "127.0.0.1"), + port=int(environ.get("DMI_CLICKHOUSE_PORT", "9000")), + ) + prefix = f"dmi_catalog_test_{uuid4().hex}" + config = ClickHouseCatalogConfig( + database=environ.get("DMI_CLICKHOUSE_DATABASE", "default"), + table_prefix=prefix, + ) + writer = ClickHouseCatalogWriter(client, config) + database = config.database + descriptor = synthetic_descriptors(1)[0] + ref = descriptor.locator.pack_ref + schema_created = False + try: + writer.ensure_schema() + schema_created = True + for version in (1, 2): + writer.write_descriptors([descriptor], index_version=version) + writer.commit_packs([ref], index_version=version) + + assert client.execute( + f"SELECT count() FROM `{database}`.`{prefix}_capture_raw`" + ) == [(2,)] + assert client.execute( + f"SELECT count() FROM `{database}`.`{prefix}_capture`" + ) == [(1,)] + assert client.execute( + f"SELECT count() FROM `{database}`.`{prefix}_pack_inventory`" + ) == [(1,)] + assert writer.committed_pack_ids([(ref.store_id, ref.pack_id)]) == { + (ref.store_id, ref.pack_id) + } + finally: + if schema_created: + for kind, suffix in ( + ("VIEW", "capture"), + ("VIEW", "pack_inventory"), + ("TABLE", "capture_raw"), + ("TABLE", "pack_inventory_raw"), + ): + client.execute( + f"DROP {kind} IF EXISTS `{database}`.`{prefix}_{suffix}`" + ) diff --git a/tests/test_clickhouse_host_benchmark.py b/tests/test_clickhouse_host_benchmark.py index 06407c437..aa0cf7a58 100644 --- a/tests/test_clickhouse_host_benchmark.py +++ b/tests/test_clickhouse_host_benchmark.py @@ -533,6 +533,40 @@ def disconnect(self): } +def test_server_sampler_filters_current_insert_without_removed_tables_column(): + class Client: + process_query = None + process_params = None + + def execute(self, query, params=None): + if "system.processes" in query: + self.process_query = query + self.process_params = params + return [(1,)] + return [] + + def disconnect(self): + pass + + client = Client() + sampler = ServerTelemetrySampler( + lambda: client, + interval_ms=1, + database="default", + table="offload", + ) + + sampler.sample_once() + + assert "tables" not in client.process_query + assert "current_database" in client.process_query + assert "position(query" in client.process_query + assert client.process_params == { + "database": "default", + "qualified_table": "`default`.`offload`", + } + + def test_identifier_quoting_rejects_sql_fragments(): assert quote_identifier("bench_2026") == "`bench_2026`" with pytest.raises(ValueError, match="identifier"): diff --git a/tests/test_garage_live.py b/tests/test_garage_live.py new file mode 100644 index 000000000..5fe17e808 --- /dev/null +++ b/tests/test_garage_live.py @@ -0,0 +1,95 @@ +from __future__ import annotations + +import os +from uuid import uuid4 + +import pytest + +from dmi.storage.capture import ( + CaptureMetadata, + CaptureRecord, + PackIndex, + PackWriter, + S3PackStore, + S3StoreConfig, +) + + +pytestmark = [pytest.mark.manual, pytest.mark.garage] + + +def _store() -> S3PackStore: + names = ( + "DMI_S3_ENDPOINT", + "DMI_S3_BUCKET", + "DMI_S3_ACCESS_KEY_ID", + "DMI_S3_SECRET_ACCESS_KEY", + ) + values = {name: os.environ.get(name) for name in names} + missing = [name for name, value in values.items() if not value] + if missing: + pytest.skip("missing Garage environment: " + ", ".join(missing)) + config = S3StoreConfig( + endpoint_url=values["DMI_S3_ENDPOINT"], + bucket=values["DMI_S3_BUCKET"], + region=os.environ.get("DMI_S3_REGION", "garage"), + access_key_id=values["DMI_S3_ACCESS_KEY_ID"], + secret_access_key=values["DMI_S3_SECRET_ACCESS_KEY"], + store_id="garage-live", + allow_insecure_http=os.environ.get("DMI_S3_ALLOW_HTTP") == "1", + multipart_threshold_bytes=5 * 1024**2, + multipart_chunk_bytes=5 * 1024**2, + multipart_concurrency=2, + ) + return S3PackStore.from_config(config) + + +def test_garage_multipart_retry_listing_and_two_range_footer_read(): + store = _store() + pack_id = uuid4() + metadata = CaptureMetadata( + capture_id=f"garage-live-{pack_id}", + tenant_id="test", + experiment_id="garage-live", + run_id=str(pack_id), + session_id="session-0", + request_id="request-0", + sequence_id="sequence-0", + model_id="synthetic", + model_revision="test-v1", + adapter_revision=None, + capture_policy_version="all-v1", + hook_name="resid_pre", + layer_number=0, + producer_rank=0, + step_number=0, + token_start=0, + token_end=1, + batch_position=0, + dtype="uint8", + shape=(6 * 1024**2,), + captured_at_ns=1_700_000_000_000_000_000, + ) + writer = PackWriter( + pack_id=pack_id, + created_at_ns=metadata.captured_at_ns, + max_pack_bytes=8 * 1024**2, + ) + writer.append( + CaptureRecord(metadata=metadata, payload=b"x" * (6 * 1024**2)) + ) + pack = writer.seal() + prefix = f"tests/dmi/{pack_id}/" + key = f"{prefix}{pack.pack_id}.dmi-pack" + + first = store.put(pack, key) + second = store.put(pack, key) + info = store.stat(first) + index = PackIndex.from_store(store, first) + page = store.list_objects(prefix=prefix, limit=10) + + assert first == second + assert info.size == len(pack.data) + assert info.checksum == pack.checksum + assert index.descriptors()[0].capture_id == metadata.capture_id + assert [item.object_key for item in page.items] == [key] diff --git a/tests/test_garage_upload_benchmark.py b/tests/test_garage_upload_benchmark.py new file mode 100644 index 000000000..961b7b47c --- /dev/null +++ b/tests/test_garage_upload_benchmark.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +import json + +import pytest + +from benchmarks.bench_garage_upload import GarageBenchmarkConfig, main + + +pytestmark = pytest.mark.cpu + + +def test_garage_benchmark_config_rejects_unbounded_inputs(): + with pytest.raises(ValueError, match="pack_payload_bytes"): + GarageBenchmarkConfig(pack_payload_bytes=()) + with pytest.raises(ValueError, match="upload_workers"): + GarageBenchmarkConfig(upload_workers=(0,)) + with pytest.raises(ValueError, match="multipart_chunk_bytes"): + GarageBenchmarkConfig(multipart_chunk_bytes=1024) + + +def test_garage_benchmark_dry_run_needs_no_credentials(capsys): + assert main( + [ + "--pack-payload-bytes", + "1MiB,2MiB", + "--multipart-threshold-bytes", + "1MiB", + "--upload-workers", + "1,2", + "--packs-per-trial", + "2", + "--trials", + "1", + "--dry-run", + ] + ) == 0 + + result = json.loads(capsys.readouterr().out) + assert result["dry_run"] is True + assert result["config"]["pack_payload_bytes"] == [1024**2, 2 * 1024**2] + assert result["config"]["upload_workers"] == [1, 2] diff --git a/tests/test_internal_mapper.py b/tests/test_internal_mapper.py new file mode 100644 index 000000000..cbb9bba32 --- /dev/null +++ b/tests/test_internal_mapper.py @@ -0,0 +1,460 @@ +"""Unit tests for monitoring.internal_mapper -- pure reassembly logic, no DB.""" +import pytest +import torch + +from monitoring.internal_mapper import ( + IncompleteInternalError, + InternalRequirements, + get_internal, + make_lazy_internal, +) + +ACT = "blocks.hook_resid_pre" + + +class FakeReader: + """Returns canned (key, tensor) rows regardless of the prefix.""" + def __init__(self, rows): + self._rows = rows + self.calls = 0 + + def prefix_get(self, prefix): + self.calls += 1 + return self._rows + + +class FailingOnceReader: + def __init__(self, rows): + self._rows = rows + self.calls = 0 + + def prefix_get(self, prefix): + self.calls += 1 + if self.calls == 1: + raise RuntimeError("temporary failure") + return self._rows + + +class SequenceReader: + def __init__(self, row_sets): + self._row_sets = list(row_sets) + self.calls = 0 + + def prefix_get(self, prefix): + self.calls += 1 + idx = min(self.calls - 1, len(self._row_sets) - 1) + rows = self._row_sets[idx] + if isinstance(rows, Exception): + raise rows + return rows + + +class PrefixReader: + def __init__(self, rows): + self._rows = rows + self.prefixes = [] + + @property + def calls(self): + return len(self.prefixes) + + def prefix_get(self, prefix): + self.prefixes.append(prefix) + return [ + (key, tensor) + for key, tensor in self._rows + if key[:len(prefix)] == prefix + ] + + +def _row(req, layer, start, tensor): + # key = (model_id, request_id, act_name, layer_no, shard_rank, start, end) + return ((("m", req, ACT, layer, 0, start, start + tensor.shape[0])), tensor) + + +def _row_act(req, act, layer, start, end, tensor): + return ((("m", req, act, layer, 0, start, end)), tensor) + + +def test_per_layer_tuple_and_shape(): + rows = [ + _row("0:0", 0, 0, torch.ones(3, 4)), + _row("0:0", 1, 0, torch.ones(3, 4) * 5), + ] + internal = get_internal("m", FakeReader(rows)) + assert internal.available == ["hidden_states"] + hs = internal.hidden_states + assert len(hs) == 2 # two layers + assert tuple(hs[0].shape) == (1, 3, 4) # [batch, seq, hidden] + + +def test_available_includes_hf_mapped_fields(): + rows = [ + _row("0:0", 0, 0, torch.ones(3, 4)), + _row_act("0:0", "blocks.attn.hook_pattern", 0, 0, 3, torch.ones(2, 3, 3)), + _row_act("0:0", "final_logits", -1, 0, 3, torch.ones(3, 10)), + _row_act("0:0", "token_ids", -1, 0, 3, torch.arange(3)), + _row_act("0:0", "hook_embed", -1, 0, 3, torch.ones(3, 4)), + _row_act("0:0", "blocks.attn.hook_q", 0, 0, 3, torch.ones(3, 2, 2)), + ] + internal = get_internal("m", FakeReader(rows)) + + assert internal.available == [ + "attentions", + "embeddings", + "hidden_states", + "logits", + "q", + "token_ids", + ] + + +def test_logits_reassemble_global_batch_by_numeric_request_id(): + rows = [ + _row_act("0:10", "final_logits", -1, 0, 1, torch.full((1, 2), 10.0)), + _row_act("0:2", "final_logits", -1, 0, 1, torch.full((1, 2), 2.0)), + ] + logits = get_internal("m", FakeReader(rows)).logits + + assert tuple(logits.shape) == (2, 1, 2) + assert torch.equal(logits[0, 0], torch.full((2,), 2.0)) + assert torch.equal(logits[1, 0], torch.full((2,), 10.0)) + + +def test_token_ids_reassemble_global_1d(): + rows = [ + _row_act("0:0", "token_ids", -1, 0, 2, torch.tensor([10, 11])), + _row_act("0:1", "token_ids", -1, 0, 1, torch.tensor([20])), + ] + token_ids = get_internal("m", FakeReader(rows)).token_ids + + assert torch.equal(token_ids, torch.tensor([[10, 11], [0, 20]])) + + +def test_embeddings_reassemble_global_hidden_field(): + rows = [ + _row_act("0:0", "hook_embed", -1, 0, 2, torch.ones(2, 4)), + _row_act("0:1", "hook_embed", -1, 0, 1, torch.ones(1, 4) * 2), + ] + embeddings = get_internal("m", FakeReader(rows)).embeddings + + assert tuple(embeddings.shape) == (2, 2, 4) + assert torch.equal(embeddings[1, 0], torch.zeros(4)) + assert torch.equal(embeddings[1, 1], torch.ones(4) * 2) + + +def test_q_reassemble_per_layer(): + rows = [ + _row_act("0:0", "blocks.attn.hook_q", 0, 0, 2, torch.ones(2, 2, 3)), + _row_act("0:0", "blocks.attn.hook_q", 1, 0, 2, torch.ones(2, 2, 3) * 2), + ] + q = get_internal("m", FakeReader(rows)).q + + assert len(q) == 2 + assert tuple(q[0].shape) == (1, 2, 2, 3) + assert torch.equal(q[1], torch.ones(1, 2, 2, 3) * 2) + + +def test_attentions_reassemble_per_layer_with_decode_growth(): + prefill = torch.ones(2, 2, 2) + decode = torch.ones(2, 1, 3) * 2 + rows = [ + _row_act("0:0", "blocks.attn.hook_pattern", 0, 0, 2, prefill), + _row_act("0:0", "blocks.attn.hook_pattern", 0, 2, 3, decode), + ] + attentions = get_internal("m", FakeReader(rows)).attentions + + assert len(attentions) == 1 + assert tuple(attentions[0].shape) == (1, 2, 3, 3) + assert torch.equal(attentions[0][0, :, :2, :2], prefill) + assert torch.equal(attentions[0][0, :, 2:, :], decode) + + +def test_attention_scores_reassemble_with_negative_inf_padding(): + prefill = torch.ones(1, 2, 2) + decode = torch.ones(1, 1, 3) * 2 + rows = [ + _row_act("0:0", "blocks.attn.hook_attn_scores", 0, 0, 2, prefill), + _row_act("0:0", "blocks.attn.hook_attn_scores", 0, 2, 3, decode), + ] + scores = get_internal("m", FakeReader(rows)).attention_scores + + assert tuple(scores[0].shape) == (1, 1, 3, 3) + assert torch.isneginf(scores[0][0, 0, 0, 2]) + assert torch.equal(scores[0][0, :, :2, :2], prefill) + assert torch.equal(scores[0][0, :, 2:, :], decode) + + +def test_chunks_concat_by_start(): + rows = [ + _row("0:0", 0, 2, torch.full((1, 4), 9.0)), # out-of-order on purpose + _row("0:0", 0, 0, torch.full((2, 4), 1.0)), + ] + hs = get_internal("m", FakeReader(rows)).hidden_states + assert tuple(hs[0].shape) == (1, 3, 4) + assert torch.equal(hs[0][0, :2], torch.full((2, 4), 1.0)) + assert torch.equal(hs[0][0, 2], torch.full((4,), 9.0)) + + +def test_ragged_batch_left_pads(): + rows = [ + _row("0:0", 0, 0, torch.ones(3, 4)), # 3 tokens + _row("0:1", 0, 0, torch.ones(2, 4) * 2), # 2 tokens -> left-pad to 3 + ] + hs = get_internal("m", FakeReader(rows)).hidden_states + assert tuple(hs[0].shape) == (2, 3, 4) + assert torch.equal(hs[0][1, 0], torch.zeros(4)) # front row padded + assert torch.equal(hs[0][1, 1:], torch.ones(2, 4) * 2) # real tokens right-aligned + + +def test_uncaptured_field_raises(): + rows = [_row("0:0", 0, 0, torch.ones(3, 4))] + internal = get_internal("m", FakeReader(rows)) + with pytest.raises(AttributeError, match="not captured"): + internal.attention + + +def test_get_internal_accepts_model_id(): + rows = [_row("0:0", 0, 0, torch.ones(3, 4))] + reader = FakeReader(rows) + + assert get_internal("m", reader).available == ["hidden_states"] + + +def test_lazy_internal_loads_once_after_success(): + rows = [_row("0:0", 0, 0, torch.ones(3, 4))] + reader = FakeReader(rows) + internal = make_lazy_internal("m", reader) + + assert "pending" in repr(internal) + assert internal.hidden_states[0].shape == (1, 3, 4) + assert reader.calls == 1 + assert internal.hidden_states[0].shape == (1, 3, 4) + assert reader.calls == 1 + assert "cached=['hidden_states']" in repr(internal) + + +def test_lazy_internal_with_request_ids_reads_only_requested_act_prefixes(): + rows = [ + _row("0:0", 0, 0, torch.ones(3, 4)), + _row("0:1", 0, 0, torch.ones(2, 4) * 2), + _row_act("0:0", "final_logits", -1, 0, 3, torch.ones(3, 10)), + _row_act("other:0", ACT, 0, 0, 1, torch.ones(1, 4) * 9), + ] + reader = PrefixReader(rows) + internal = make_lazy_internal("m", reader, request_ids=("0:0", "0:1")) + + hs = internal.hidden_states + + assert tuple(hs[0].shape) == (2, 3, 4) + assert reader.prefixes == [ + ("m", "0:0", ACT), + ("m", "0:1", ACT), + ] + + +def test_lazy_internal_token_mask_uses_recorded_ranges_without_reader(): + reader = PrefixReader([]) + internal = make_lazy_internal( + "m", + reader, + request_ids=("0:0", "0:1"), + token_ranges={ + "0:0": ((0, 3), (3, 4)), + "0:1": ((0, 2), (2, 2)), + }, + ) + + mask = internal.token_mask + + assert mask.dtype == torch.bool + assert torch.equal( + mask, + torch.tensor([ + [True, True, True, True], + [False, False, True, True], + ]), + ) + assert reader.calls == 0 + assert internal.token_mask is mask + + +def test_lazy_internal_retries_after_failed_load(): + rows = [_row("0:0", 0, 0, torch.ones(3, 4))] + reader = FailingOnceReader(rows) + internal = make_lazy_internal("m", reader) + + with pytest.raises(RuntimeError, match="temporary failure"): + internal.hidden_states + + assert internal.hidden_states[0].shape == (1, 3, 4) + assert reader.calls == 2 + + +def test_lazy_internal_clear_cache_one_field(): + rows = [_row("0:0", 0, 0, torch.ones(3, 4))] + reader = FakeReader(rows) + internal = make_lazy_internal("m", reader) + + assert internal.hidden_states[0].shape == (1, 3, 4) + assert internal.hidden_states[0].shape == (1, 3, 4) + assert reader.calls == 1 + + internal.clear_cache("hidden_states") + assert internal.hidden_states[0].shape == (1, 3, 4) + assert reader.calls == 2 + + +def test_lazy_internal_clear_cache_all_fields(): + rows = [_row("0:0", 0, 0, torch.ones(3, 4))] + reader = FakeReader(rows) + internal = make_lazy_internal("m", reader) + + assert internal.hidden_states[0].shape == (1, 3, 4) + internal.clear_cache() + assert internal.hidden_states[0].shape == (1, 3, 4) + assert reader.calls == 2 + + +def test_lazy_internal_requirement_blocks_incomplete_cache(): + rows = [_row("0:0", 0, 0, torch.ones(3, 4))] + reader = FakeReader(rows) + internal = make_lazy_internal("m", reader) + internal.require("hidden_states", count=2) + + with pytest.raises(IncompleteInternalError, match="expected 2 entries, found 1"): + internal.hidden_states + + with pytest.raises(IncompleteInternalError, match="expected 2 entries, found 1"): + internal.hidden_states + + assert reader.calls == 2 + + +def test_lazy_internal_requirement_retries_missing_until_success(): + rows = [_row("0:0", 0, 0, torch.ones(3, 4))] + reader = SequenceReader([[], rows]) + internal = make_lazy_internal("m", reader) + internal.require("hidden_states", count=1, retry=True, timeout_s=1.0, poll_s=0.001) + + assert len(internal.hidden_states) == 1 + assert reader.calls == 2 + + +def test_lazy_internal_requirement_retries_incomplete_until_success(): + one_layer = [_row("0:0", 0, 0, torch.ones(3, 4))] + two_layers = [ + _row("0:0", 0, 0, torch.ones(3, 4)), + _row("0:0", 1, 0, torch.ones(3, 4)), + ] + reader = SequenceReader([one_layer, two_layers]) + internal = make_lazy_internal("m", reader) + internal.require("hidden_states", count=2, retry=True, timeout_s=1.0, poll_s=0.001) + + assert len(internal.hidden_states) == 2 + assert reader.calls == 2 + + +def test_lazy_internal_requirement_detects_missing_token_ranges(): + rows = [_row("0:0", 0, 0, torch.ones(2, 4))] + reader = PrefixReader(rows) + internal = make_lazy_internal( + "m", + reader, + request_ids=("0:0",), + token_ranges={"0:0": ((0, 2), (2, 3))}, + ) + internal.require("hidden_states", count=1, match_token_ranges=True) + + with pytest.raises(IncompleteInternalError, match="token ranges are incomplete"): + internal.hidden_states + + +def test_lazy_internal_requirement_retries_missing_token_ranges_until_success(): + partial = [_row("0:0", 0, 0, torch.ones(2, 4))] + complete = [ + _row("0:0", 0, 0, torch.ones(2, 4)), + _row("0:0", 0, 2, torch.ones(1, 4)), + ] + reader = SequenceReader([partial, complete]) + internal = make_lazy_internal( + "m", + reader, + request_ids=("0:0",), + token_ranges={"0:0": ((0, 2), (2, 3))}, + ) + internal.require( + "hidden_states", + count=1, + retry=True, + timeout_s=1.0, + poll_s=0.001, + match_token_ranges=True, + ) + + assert tuple(internal.hidden_states[0].shape) == (1, 3, 4) + assert reader.calls == 2 + + +def test_lazy_internal_requirement_retry_timeout_raises_incomplete(): + rows = [_row("0:0", 0, 0, torch.ones(3, 4))] + reader = SequenceReader([rows]) + internal = make_lazy_internal("m", reader) + internal.require("hidden_states", count=2, retry=True, timeout_s=0.0, poll_s=0.001) + + with pytest.raises(IncompleteInternalError, match="expected 2 entries, found 1"): + internal.hidden_states + assert reader.calls == 1 + + +def test_lazy_internal_requirement_can_retry_without_timeout(): + rows = [_row("0:0", 0, 0, torch.ones(3, 4))] + reader = SequenceReader([[], [], rows]) + internal = make_lazy_internal("m", reader) + internal.require("hidden_states", count=1, retry=True, timeout_s=None, poll_s=0.001) + + assert len(internal.hidden_states) == 1 + assert reader.calls == 3 + + +def test_lazy_internal_retry_does_not_swallow_unexpected_errors(): + reader = SequenceReader([RuntimeError("boom")]) + internal = make_lazy_internal("m", reader) + internal.require("hidden_states", count=1, retry=True, timeout_s=1.0, poll_s=0.001) + + with pytest.raises(RuntimeError, match="boom"): + internal.hidden_states + assert reader.calls == 1 + + +def test_lazy_internal_requirement_revalidates_existing_cache(): + rows = [_row("0:0", 0, 0, torch.ones(3, 4))] + reader = FakeReader(rows) + internal = make_lazy_internal("m", reader) + + assert len(internal.hidden_states) == 1 + internal.require("hidden_states", count=2) + + with pytest.raises(IncompleteInternalError, match="expected 2 entries, found 1"): + internal.hidden_states + + assert reader.calls == 1 + with pytest.raises(IncompleteInternalError, match="expected 2 entries, found 1"): + internal.hidden_states + assert reader.calls == 2 + + +def test_lazy_internal_copies_reusable_requirements(): + rows = [ + _row("0:0", 0, 0, torch.ones(3, 4)), + _row("0:0", 1, 0, torch.ones(3, 4)), + ] + reqs = InternalRequirements().require("hidden_states", count=2) + first = make_lazy_internal("m", FakeReader(rows), requirements=reqs) + second = make_lazy_internal("m", FakeReader(rows), requirements=reqs) + + first.require("hidden_states", count=1) + + with pytest.raises(IncompleteInternalError, match="expected 1 entries, found 2"): + first.hidden_states + assert len(second.hidden_states) == 2 diff --git a/tests/test_internal_mapper_e2e.py b/tests/test_internal_mapper_e2e.py new file mode 100644 index 000000000..227102833 --- /dev/null +++ b/tests/test_internal_mapper_e2e.py @@ -0,0 +1,199 @@ +"""End-to-end tests for the get_internal read path. + +Runs Qwen3-0.6B through generate_with_monitoring (writing captures to +ClickHouse via the native ring), then reads them back with +monitoring.internal_mapper.get_internal and checks the result lines up with +HuggingFace's native output. + +Requires a CUDA device, a reachable ClickHouse, the patched transformers fork, +and the model in the local cache; skips cleanly otherwise. + +ClickHouse connection: DMX_DB_HOST / DMX_DB_PORT (default localhost:9000). +""" +from __future__ import annotations + +import os +import uuid +from types import SimpleNamespace + +import pytest +import torch + +from tests._requirements import require_cuda, require_clickhouse, require_model_cache + +MODEL = "Qwen/Qwen3-0.6B" +MAX_NEW = 8 + +pytestmark = [ + pytest.mark.gpu, + pytest.mark.e2e, + pytest.mark.clickhouse, + pytest.mark.hf, + require_cuda(), + require_clickhouse(), + require_model_cache(MODEL), +] + + +@pytest.fixture(scope="module") +def deps(): + """Import the DMI HF stack; skip the module if any piece is missing.""" + try: + from transformers import AutoTokenizer + from transformers.models.qwen3_p.modeling_qwen3 import HookedQwen3ForCausalLM + from monitoring import ( + CaptureSchedule, HostEngineConfig, MonitoringConfig, MonitoringEngine, + ) + from monitoring._native_engine import ClickHouseClientConfig, StageConfig + from monitoring.clickhouse_reader import CHClickhouseDriverReadOnly + from monitoring.internal_mapper import get_internal + from integration.hf_adapter import generate_with_monitoring_dict + import clickhouse_driver + except Exception as exc: + pytest.skip(f"DMI HF stack unavailable: {exc}") + return SimpleNamespace(**{k: v for k, v in locals().items() if k != "exc"}) + + +def _db_host_port(): + return (os.environ.get("DMX_DB_HOST", "localhost"), + int(os.environ.get("DMX_DB_PORT", "9000"))) + + +def _wipe(deps, model_id): + host, port = _db_host_port() + deps.clickhouse_driver.Client(host=host, port=port).execute( + "ALTER TABLE default.offload DELETE WHERE model_id = %(m)s", {"m": model_id}) + + +def _monitored(deps, model, inputs, model_id, hook_selection="resid_pre"): + """Run generate_with_monitoring_dict into a fresh ClickHouse slot.""" + _wipe(deps, model_id) + host, port = _db_host_port() + db = deps.ClickHouseClientConfig() + db.host, db.port = host, port + db.username, db.password = "default", "" + db.database, db.table = "default", "offload" + db.create_database_if_missing = True + stage = deps.StageConfig.clickhouse_insert(db, parallelism=4, name="ch_insert") + engine = deps.MonitoringEngine( + config=deps.MonitoringConfig(schedule=deps.CaptureSchedule()), + model_id=model_id, + db_config=deps.HostEngineConfig(stages=[stage]), + ) + model.monitoring_engine = engine + try: + with torch.no_grad(): + return deps.generate_with_monitoring_dict( + model, **inputs, max_new_tokens=MAX_NEW, do_sample=False, + hook_selection=hook_selection, + ) + finally: + engine.close() + + +@pytest.fixture(scope="module") +def model_tok(deps): + tok = deps.AutoTokenizer.from_pretrained(MODEL) + if tok.pad_token_id is None: + tok.pad_token_id = tok.eos_token_id + tok.padding_side = "left" + model = deps.HookedQwen3ForCausalLM.from_pretrained( + MODEL, dtype=torch.float16, attn_implementation="eager", + ).to("cuda").eval() + return model, tok + + +@pytest.fixture(scope="module") +def single(deps, model_tok): + """One prompt (batch=1): monitored output + a plain-HF reference.""" + model, tok = model_tok + inputs = tok(["The capital of France is"], return_tensors="pt", padding=True).to("cuda") + with torch.no_grad(): + ref = model.generate(**inputs, max_new_tokens=MAX_NEW, do_sample=False) + model_id = f"test_internal::single::{uuid.uuid4().hex}"[:120] + out = _monitored(deps, model, inputs, model_id) + yield SimpleNamespace(out=out, ref=ref, model=model, model_id=model_id) + _wipe(deps, model_id) + + +@pytest.fixture(scope="module") +def ragged(deps, model_tok): + """A left-padded batch of unequal-length prompts.""" + model, tok = model_tok + inputs = tok(["The capital of France is", "Hi"], return_tensors="pt", padding=True).to("cuda") + model_id = f"test_internal::ragged::{uuid.uuid4().hex}"[:120] + out = _monitored(deps, model, inputs, model_id) + yield SimpleNamespace(out=out, n_prompts=2, model_id=model_id) + _wipe(deps, model_id) + + +# --- dict return shape ------------------------------------------------------ + +def test_returns_dict_output_with_dmi_internal(single): + assert single.out.sequences.dim() == 2 # [batch, total_len] + assert hasattr(single.out, "dmi_internal") + + +def test_output_matches_plain_hf(single): + # Monitoring must not change what the model generates. + assert torch.equal(single.out.sequences, single.ref) + + +# --- get_internal: structure ------------------------------------------------ + +def test_available_lists_captured_field(single, deps): + internal = single.out.dmi_internal + assert internal.available == ["hidden_states", "token_mask"] + + +def test_hidden_states_tuple_shape(single, deps): + internal = single.out.dmi_internal + hs = internal.hidden_states + token_mask = internal.token_mask + cfg = single.model.config + assert len(hs) == cfg.num_hidden_layers # one entry per block (resid_pre) + assert all(t.dim() == 3 for t in hs) # [batch, seq, hidden] + assert hs[0].shape[0] == 1 # batch + assert hs[0].shape[2] == cfg.hidden_size # hidden + assert all(t.shape == hs[0].shape for t in hs) # uniform across layers + assert token_mask.shape == hs[0].shape[:2] + assert token_mask.dtype == torch.bool + + +def test_token_count_drops_unforwarded_last(single, deps): + # The final generated token is never fed through a forward, so it has no + # captured hidden state: captured seq == sequence length - 1. + hs = single.out.dmi_internal.hidden_states + assert hs[0].shape[1] == single.out.sequences.shape[1] - 1 + + +# --- get_internal: model_id + reader handling ------------------------------ + +def test_source_model_id_with_explicit_reader(single, deps): + host, port = _db_host_port() + reader = deps.CHClickhouseDriverReadOnly(host=host, port=port) + from_out = single.out.dmi_internal.hidden_states + from_id = deps.get_internal(single.model_id, reader).hidden_states + assert len(from_out) == len(from_id) + assert torch.equal(from_out[0], from_id[0]) + + +def test_uncaptured_field_raises(single, deps): + internal = single.out.dmi_internal + with pytest.raises(AttributeError, match="not captured"): + internal.attention + + +# --- get_internal: ragged batch left-pad ----------------------------------- + +def test_ragged_batch_left_pads(ragged, deps): + hs = ragged.out.dmi_internal.hidden_states + layer0 = hs[0] + assert layer0.shape[0] == ragged.n_prompts + + nonzero = (layer0.abs().sum(dim=-1) > 0) # [batch, seq] real-token mask + counts = nonzero.sum(dim=1) + assert counts.min() < counts.max() # genuinely ragged -> padding added + assert bool(nonzero[:, -1].all()) # real tokens are right-aligned + short = int(counts.argmin()) + assert not bool(nonzero[short, 0]) # shorter request padded at the front diff --git a/tests/tools/check_package.py b/tests/tools/check_package.py index 5358c75cd..c1a44b787 100644 --- a/tests/tools/check_package.py +++ b/tests/tools/check_package.py @@ -30,6 +30,15 @@ "dmi/api/v1/__init__.py", "dmi/hooks/dispatch.py", "dmi/hooks/point.py", + "dmi/storage/capture/filesystem.py", + "dmi/storage/capture/catalog.py", + "dmi/storage/capture/clickhouse_catalog.py", + "dmi/storage/capture/model.py", + "dmi/storage/capture/pack.py", + "dmi/storage/capture/pipeline.py", + "dmi/storage/capture/reader.py", + "dmi/storage/capture/s3.py", + "dmi/storage/capture/spool.py", "dmi/storage/internals.py", "dmi/transport/native.py", "dmi/transport/ring.py", @@ -112,6 +121,9 @@ def _smoke_test_install(wheel: Path, audit_root: Path) -> None: "dmi.adapters.huggingface.adapter", "dmi.api.v1", "dmi.hooks.dispatch", + "dmi.storage.capture", + "dmi.storage.capture.catalog", + "dmi.storage.capture.clickhouse_catalog", "dmi.storage.internals", "dmi.transport.native", ): diff --git a/tests/tools/run_garage_live.py b/tests/tools/run_garage_live.py new file mode 100644 index 000000000..da1839f8a --- /dev/null +++ b/tests/tools/run_garage_live.py @@ -0,0 +1,161 @@ +#!/usr/bin/env python3 +"""Run the manual Garage integration test against an ephemeral local server.""" + +from __future__ import annotations + +import argparse +import os +from pathlib import Path +import secrets +import shutil +import socket +import subprocess +import sys +import tempfile +import time + + +REPO_ROOT = Path(__file__).resolve().parents[2] + + +def _ports(count: int) -> list[int]: + listeners = [socket.socket() for _ in range(count)] + try: + for listener in listeners: + listener.bind(("127.0.0.1", 0)) + return [listener.getsockname()[1] for listener in listeners] + finally: + for listener in listeners: + listener.close() + + +def _wait_for_port(port: int, process: subprocess.Popen, timeout: float = 20) -> None: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if process.poll() is not None: + raise RuntimeError(f"Garage exited with status {process.returncode}") + try: + with socket.create_connection(("127.0.0.1", port), timeout=0.2): + return + except OSError: + time.sleep(0.1) + raise TimeoutError("Garage did not open its S3 port") + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--benchmark", action="store_true") + parser.add_argument("benchmark_args", nargs=argparse.REMAINDER) + args = parser.parse_args(argv) + binary = os.environ.get("DMI_GARAGE_BINARY") or shutil.which("garage") + if binary is None: + raise RuntimeError( + "set DMI_GARAGE_BINARY or install Garage from " + "https://garagehq.deuxfleurs.fr/download/" + ) + with tempfile.TemporaryDirectory(prefix="dmi-garage-live-") as directory: + root = Path(directory) + version = subprocess.run( + [binary, "--version"], capture_output=True, text=True, check=True + ).stdout + expected_version = os.environ.get("DMI_GARAGE_VERSION", "2.3.0") + if f"cargo:{expected_version}" not in version: + raise RuntimeError( + f"expected Garage {expected_version}, got {version.strip()}" + ) + ports = _ports(4) + s3_port, rpc_port, web_port, admin_port = ports + config = root / "garage.toml" + config.write_text( + f'''metadata_dir = "{root / "meta"}" +data_dir = "{root / "data"}" +db_engine = "sqlite" +replication_factor = 1 +rpc_bind_addr = "127.0.0.1:{rpc_port}" +rpc_public_addr = "127.0.0.1:{rpc_port}" +rpc_secret = "{secrets.token_hex(32)}" + +[s3_api] +s3_region = "garage" +api_bind_addr = "127.0.0.1:{s3_port}" +root_domain = ".s3.garage.localhost" + +[s3_web] +bind_addr = "127.0.0.1:{web_port}" +root_domain = ".web.garage.localhost" +index = "index.html" + +[admin] +api_bind_addr = "127.0.0.1:{admin_port}" +admin_token = "{secrets.token_urlsafe(32)}" +metrics_token = "{secrets.token_urlsafe(32)}" +''' + ) + access_key = "GK" + secrets.token_hex(16) + secret_key = secrets.token_hex(32) + environment = os.environ.copy() + environment.update( + { + "GARAGE_CONFIG_FILE": str(config), + "GARAGE_DEFAULT_ACCESS_KEY": access_key, + "GARAGE_DEFAULT_SECRET_KEY": secret_key, + "GARAGE_DEFAULT_BUCKET": "dmi-test", + "DMI_S3_ENDPOINT": f"http://127.0.0.1:{s3_port}", + "DMI_S3_BUCKET": "dmi-test", + "DMI_S3_REGION": "garage", + "DMI_S3_ACCESS_KEY_ID": access_key, + "DMI_S3_SECRET_ACCESS_KEY": secret_key, + "DMI_S3_ALLOW_HTTP": "1", + "PYTHONPATH": str(REPO_ROOT / "src"), + } + ) + process = subprocess.Popen( + [binary, "server", "--single-node", "--default-bucket"], + cwd=root, + env=environment, + stdout=subprocess.DEVNULL, + stderr=subprocess.PIPE, + text=True, + ) + try: + _wait_for_port(s3_port, process) + command = ( + [ + sys.executable, + "-m", + "benchmarks.bench_garage_upload", + *( + args.benchmark_args[1:] + if args.benchmark_args[:1] == ["--"] + else args.benchmark_args + ), + ] + if args.benchmark + else [ + sys.executable, + "-m", + "pytest", + "tests/test_garage_live.py", + "-m", + "garage and manual", + "-q", + ] + ) + completed = subprocess.run( + command, + cwd=REPO_ROOT, + env=environment, + check=False, + ) + return completed.returncode + finally: + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=5) + + +if __name__ == "__main__": + raise SystemExit(main()) From c7227ea51a88c48db3c87042af24d381353f58c2 Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Tue, 25 Aug 2026 20:37:54 -0400 Subject: [PATCH 02/20] test: remove pre-refactor test duplicates resurrected in Phase 4 3f95d87 (#95) renamed monitoring/internal_mapper.py to src/dmi/storage/internals.py and moved both of its test modules to tests/test_storage_internals*.py with updated imports. 46e270b re-added the pre-refactor copies, which still import the removed `monitoring` package. Because a collection error aborts the whole run, this broke the PR gate outright: `python -m pytest -m "cpu" -q` failed before executing a single test. The current modules already cover this code and pass. --- tests/test_internal_mapper.py | 460 ------------------------------ tests/test_internal_mapper_e2e.py | 199 ------------- 2 files changed, 659 deletions(-) delete mode 100644 tests/test_internal_mapper.py delete mode 100644 tests/test_internal_mapper_e2e.py diff --git a/tests/test_internal_mapper.py b/tests/test_internal_mapper.py deleted file mode 100644 index cbb9bba32..000000000 --- a/tests/test_internal_mapper.py +++ /dev/null @@ -1,460 +0,0 @@ -"""Unit tests for monitoring.internal_mapper -- pure reassembly logic, no DB.""" -import pytest -import torch - -from monitoring.internal_mapper import ( - IncompleteInternalError, - InternalRequirements, - get_internal, - make_lazy_internal, -) - -ACT = "blocks.hook_resid_pre" - - -class FakeReader: - """Returns canned (key, tensor) rows regardless of the prefix.""" - def __init__(self, rows): - self._rows = rows - self.calls = 0 - - def prefix_get(self, prefix): - self.calls += 1 - return self._rows - - -class FailingOnceReader: - def __init__(self, rows): - self._rows = rows - self.calls = 0 - - def prefix_get(self, prefix): - self.calls += 1 - if self.calls == 1: - raise RuntimeError("temporary failure") - return self._rows - - -class SequenceReader: - def __init__(self, row_sets): - self._row_sets = list(row_sets) - self.calls = 0 - - def prefix_get(self, prefix): - self.calls += 1 - idx = min(self.calls - 1, len(self._row_sets) - 1) - rows = self._row_sets[idx] - if isinstance(rows, Exception): - raise rows - return rows - - -class PrefixReader: - def __init__(self, rows): - self._rows = rows - self.prefixes = [] - - @property - def calls(self): - return len(self.prefixes) - - def prefix_get(self, prefix): - self.prefixes.append(prefix) - return [ - (key, tensor) - for key, tensor in self._rows - if key[:len(prefix)] == prefix - ] - - -def _row(req, layer, start, tensor): - # key = (model_id, request_id, act_name, layer_no, shard_rank, start, end) - return ((("m", req, ACT, layer, 0, start, start + tensor.shape[0])), tensor) - - -def _row_act(req, act, layer, start, end, tensor): - return ((("m", req, act, layer, 0, start, end)), tensor) - - -def test_per_layer_tuple_and_shape(): - rows = [ - _row("0:0", 0, 0, torch.ones(3, 4)), - _row("0:0", 1, 0, torch.ones(3, 4) * 5), - ] - internal = get_internal("m", FakeReader(rows)) - assert internal.available == ["hidden_states"] - hs = internal.hidden_states - assert len(hs) == 2 # two layers - assert tuple(hs[0].shape) == (1, 3, 4) # [batch, seq, hidden] - - -def test_available_includes_hf_mapped_fields(): - rows = [ - _row("0:0", 0, 0, torch.ones(3, 4)), - _row_act("0:0", "blocks.attn.hook_pattern", 0, 0, 3, torch.ones(2, 3, 3)), - _row_act("0:0", "final_logits", -1, 0, 3, torch.ones(3, 10)), - _row_act("0:0", "token_ids", -1, 0, 3, torch.arange(3)), - _row_act("0:0", "hook_embed", -1, 0, 3, torch.ones(3, 4)), - _row_act("0:0", "blocks.attn.hook_q", 0, 0, 3, torch.ones(3, 2, 2)), - ] - internal = get_internal("m", FakeReader(rows)) - - assert internal.available == [ - "attentions", - "embeddings", - "hidden_states", - "logits", - "q", - "token_ids", - ] - - -def test_logits_reassemble_global_batch_by_numeric_request_id(): - rows = [ - _row_act("0:10", "final_logits", -1, 0, 1, torch.full((1, 2), 10.0)), - _row_act("0:2", "final_logits", -1, 0, 1, torch.full((1, 2), 2.0)), - ] - logits = get_internal("m", FakeReader(rows)).logits - - assert tuple(logits.shape) == (2, 1, 2) - assert torch.equal(logits[0, 0], torch.full((2,), 2.0)) - assert torch.equal(logits[1, 0], torch.full((2,), 10.0)) - - -def test_token_ids_reassemble_global_1d(): - rows = [ - _row_act("0:0", "token_ids", -1, 0, 2, torch.tensor([10, 11])), - _row_act("0:1", "token_ids", -1, 0, 1, torch.tensor([20])), - ] - token_ids = get_internal("m", FakeReader(rows)).token_ids - - assert torch.equal(token_ids, torch.tensor([[10, 11], [0, 20]])) - - -def test_embeddings_reassemble_global_hidden_field(): - rows = [ - _row_act("0:0", "hook_embed", -1, 0, 2, torch.ones(2, 4)), - _row_act("0:1", "hook_embed", -1, 0, 1, torch.ones(1, 4) * 2), - ] - embeddings = get_internal("m", FakeReader(rows)).embeddings - - assert tuple(embeddings.shape) == (2, 2, 4) - assert torch.equal(embeddings[1, 0], torch.zeros(4)) - assert torch.equal(embeddings[1, 1], torch.ones(4) * 2) - - -def test_q_reassemble_per_layer(): - rows = [ - _row_act("0:0", "blocks.attn.hook_q", 0, 0, 2, torch.ones(2, 2, 3)), - _row_act("0:0", "blocks.attn.hook_q", 1, 0, 2, torch.ones(2, 2, 3) * 2), - ] - q = get_internal("m", FakeReader(rows)).q - - assert len(q) == 2 - assert tuple(q[0].shape) == (1, 2, 2, 3) - assert torch.equal(q[1], torch.ones(1, 2, 2, 3) * 2) - - -def test_attentions_reassemble_per_layer_with_decode_growth(): - prefill = torch.ones(2, 2, 2) - decode = torch.ones(2, 1, 3) * 2 - rows = [ - _row_act("0:0", "blocks.attn.hook_pattern", 0, 0, 2, prefill), - _row_act("0:0", "blocks.attn.hook_pattern", 0, 2, 3, decode), - ] - attentions = get_internal("m", FakeReader(rows)).attentions - - assert len(attentions) == 1 - assert tuple(attentions[0].shape) == (1, 2, 3, 3) - assert torch.equal(attentions[0][0, :, :2, :2], prefill) - assert torch.equal(attentions[0][0, :, 2:, :], decode) - - -def test_attention_scores_reassemble_with_negative_inf_padding(): - prefill = torch.ones(1, 2, 2) - decode = torch.ones(1, 1, 3) * 2 - rows = [ - _row_act("0:0", "blocks.attn.hook_attn_scores", 0, 0, 2, prefill), - _row_act("0:0", "blocks.attn.hook_attn_scores", 0, 2, 3, decode), - ] - scores = get_internal("m", FakeReader(rows)).attention_scores - - assert tuple(scores[0].shape) == (1, 1, 3, 3) - assert torch.isneginf(scores[0][0, 0, 0, 2]) - assert torch.equal(scores[0][0, :, :2, :2], prefill) - assert torch.equal(scores[0][0, :, 2:, :], decode) - - -def test_chunks_concat_by_start(): - rows = [ - _row("0:0", 0, 2, torch.full((1, 4), 9.0)), # out-of-order on purpose - _row("0:0", 0, 0, torch.full((2, 4), 1.0)), - ] - hs = get_internal("m", FakeReader(rows)).hidden_states - assert tuple(hs[0].shape) == (1, 3, 4) - assert torch.equal(hs[0][0, :2], torch.full((2, 4), 1.0)) - assert torch.equal(hs[0][0, 2], torch.full((4,), 9.0)) - - -def test_ragged_batch_left_pads(): - rows = [ - _row("0:0", 0, 0, torch.ones(3, 4)), # 3 tokens - _row("0:1", 0, 0, torch.ones(2, 4) * 2), # 2 tokens -> left-pad to 3 - ] - hs = get_internal("m", FakeReader(rows)).hidden_states - assert tuple(hs[0].shape) == (2, 3, 4) - assert torch.equal(hs[0][1, 0], torch.zeros(4)) # front row padded - assert torch.equal(hs[0][1, 1:], torch.ones(2, 4) * 2) # real tokens right-aligned - - -def test_uncaptured_field_raises(): - rows = [_row("0:0", 0, 0, torch.ones(3, 4))] - internal = get_internal("m", FakeReader(rows)) - with pytest.raises(AttributeError, match="not captured"): - internal.attention - - -def test_get_internal_accepts_model_id(): - rows = [_row("0:0", 0, 0, torch.ones(3, 4))] - reader = FakeReader(rows) - - assert get_internal("m", reader).available == ["hidden_states"] - - -def test_lazy_internal_loads_once_after_success(): - rows = [_row("0:0", 0, 0, torch.ones(3, 4))] - reader = FakeReader(rows) - internal = make_lazy_internal("m", reader) - - assert "pending" in repr(internal) - assert internal.hidden_states[0].shape == (1, 3, 4) - assert reader.calls == 1 - assert internal.hidden_states[0].shape == (1, 3, 4) - assert reader.calls == 1 - assert "cached=['hidden_states']" in repr(internal) - - -def test_lazy_internal_with_request_ids_reads_only_requested_act_prefixes(): - rows = [ - _row("0:0", 0, 0, torch.ones(3, 4)), - _row("0:1", 0, 0, torch.ones(2, 4) * 2), - _row_act("0:0", "final_logits", -1, 0, 3, torch.ones(3, 10)), - _row_act("other:0", ACT, 0, 0, 1, torch.ones(1, 4) * 9), - ] - reader = PrefixReader(rows) - internal = make_lazy_internal("m", reader, request_ids=("0:0", "0:1")) - - hs = internal.hidden_states - - assert tuple(hs[0].shape) == (2, 3, 4) - assert reader.prefixes == [ - ("m", "0:0", ACT), - ("m", "0:1", ACT), - ] - - -def test_lazy_internal_token_mask_uses_recorded_ranges_without_reader(): - reader = PrefixReader([]) - internal = make_lazy_internal( - "m", - reader, - request_ids=("0:0", "0:1"), - token_ranges={ - "0:0": ((0, 3), (3, 4)), - "0:1": ((0, 2), (2, 2)), - }, - ) - - mask = internal.token_mask - - assert mask.dtype == torch.bool - assert torch.equal( - mask, - torch.tensor([ - [True, True, True, True], - [False, False, True, True], - ]), - ) - assert reader.calls == 0 - assert internal.token_mask is mask - - -def test_lazy_internal_retries_after_failed_load(): - rows = [_row("0:0", 0, 0, torch.ones(3, 4))] - reader = FailingOnceReader(rows) - internal = make_lazy_internal("m", reader) - - with pytest.raises(RuntimeError, match="temporary failure"): - internal.hidden_states - - assert internal.hidden_states[0].shape == (1, 3, 4) - assert reader.calls == 2 - - -def test_lazy_internal_clear_cache_one_field(): - rows = [_row("0:0", 0, 0, torch.ones(3, 4))] - reader = FakeReader(rows) - internal = make_lazy_internal("m", reader) - - assert internal.hidden_states[0].shape == (1, 3, 4) - assert internal.hidden_states[0].shape == (1, 3, 4) - assert reader.calls == 1 - - internal.clear_cache("hidden_states") - assert internal.hidden_states[0].shape == (1, 3, 4) - assert reader.calls == 2 - - -def test_lazy_internal_clear_cache_all_fields(): - rows = [_row("0:0", 0, 0, torch.ones(3, 4))] - reader = FakeReader(rows) - internal = make_lazy_internal("m", reader) - - assert internal.hidden_states[0].shape == (1, 3, 4) - internal.clear_cache() - assert internal.hidden_states[0].shape == (1, 3, 4) - assert reader.calls == 2 - - -def test_lazy_internal_requirement_blocks_incomplete_cache(): - rows = [_row("0:0", 0, 0, torch.ones(3, 4))] - reader = FakeReader(rows) - internal = make_lazy_internal("m", reader) - internal.require("hidden_states", count=2) - - with pytest.raises(IncompleteInternalError, match="expected 2 entries, found 1"): - internal.hidden_states - - with pytest.raises(IncompleteInternalError, match="expected 2 entries, found 1"): - internal.hidden_states - - assert reader.calls == 2 - - -def test_lazy_internal_requirement_retries_missing_until_success(): - rows = [_row("0:0", 0, 0, torch.ones(3, 4))] - reader = SequenceReader([[], rows]) - internal = make_lazy_internal("m", reader) - internal.require("hidden_states", count=1, retry=True, timeout_s=1.0, poll_s=0.001) - - assert len(internal.hidden_states) == 1 - assert reader.calls == 2 - - -def test_lazy_internal_requirement_retries_incomplete_until_success(): - one_layer = [_row("0:0", 0, 0, torch.ones(3, 4))] - two_layers = [ - _row("0:0", 0, 0, torch.ones(3, 4)), - _row("0:0", 1, 0, torch.ones(3, 4)), - ] - reader = SequenceReader([one_layer, two_layers]) - internal = make_lazy_internal("m", reader) - internal.require("hidden_states", count=2, retry=True, timeout_s=1.0, poll_s=0.001) - - assert len(internal.hidden_states) == 2 - assert reader.calls == 2 - - -def test_lazy_internal_requirement_detects_missing_token_ranges(): - rows = [_row("0:0", 0, 0, torch.ones(2, 4))] - reader = PrefixReader(rows) - internal = make_lazy_internal( - "m", - reader, - request_ids=("0:0",), - token_ranges={"0:0": ((0, 2), (2, 3))}, - ) - internal.require("hidden_states", count=1, match_token_ranges=True) - - with pytest.raises(IncompleteInternalError, match="token ranges are incomplete"): - internal.hidden_states - - -def test_lazy_internal_requirement_retries_missing_token_ranges_until_success(): - partial = [_row("0:0", 0, 0, torch.ones(2, 4))] - complete = [ - _row("0:0", 0, 0, torch.ones(2, 4)), - _row("0:0", 0, 2, torch.ones(1, 4)), - ] - reader = SequenceReader([partial, complete]) - internal = make_lazy_internal( - "m", - reader, - request_ids=("0:0",), - token_ranges={"0:0": ((0, 2), (2, 3))}, - ) - internal.require( - "hidden_states", - count=1, - retry=True, - timeout_s=1.0, - poll_s=0.001, - match_token_ranges=True, - ) - - assert tuple(internal.hidden_states[0].shape) == (1, 3, 4) - assert reader.calls == 2 - - -def test_lazy_internal_requirement_retry_timeout_raises_incomplete(): - rows = [_row("0:0", 0, 0, torch.ones(3, 4))] - reader = SequenceReader([rows]) - internal = make_lazy_internal("m", reader) - internal.require("hidden_states", count=2, retry=True, timeout_s=0.0, poll_s=0.001) - - with pytest.raises(IncompleteInternalError, match="expected 2 entries, found 1"): - internal.hidden_states - assert reader.calls == 1 - - -def test_lazy_internal_requirement_can_retry_without_timeout(): - rows = [_row("0:0", 0, 0, torch.ones(3, 4))] - reader = SequenceReader([[], [], rows]) - internal = make_lazy_internal("m", reader) - internal.require("hidden_states", count=1, retry=True, timeout_s=None, poll_s=0.001) - - assert len(internal.hidden_states) == 1 - assert reader.calls == 3 - - -def test_lazy_internal_retry_does_not_swallow_unexpected_errors(): - reader = SequenceReader([RuntimeError("boom")]) - internal = make_lazy_internal("m", reader) - internal.require("hidden_states", count=1, retry=True, timeout_s=1.0, poll_s=0.001) - - with pytest.raises(RuntimeError, match="boom"): - internal.hidden_states - assert reader.calls == 1 - - -def test_lazy_internal_requirement_revalidates_existing_cache(): - rows = [_row("0:0", 0, 0, torch.ones(3, 4))] - reader = FakeReader(rows) - internal = make_lazy_internal("m", reader) - - assert len(internal.hidden_states) == 1 - internal.require("hidden_states", count=2) - - with pytest.raises(IncompleteInternalError, match="expected 2 entries, found 1"): - internal.hidden_states - - assert reader.calls == 1 - with pytest.raises(IncompleteInternalError, match="expected 2 entries, found 1"): - internal.hidden_states - assert reader.calls == 2 - - -def test_lazy_internal_copies_reusable_requirements(): - rows = [ - _row("0:0", 0, 0, torch.ones(3, 4)), - _row("0:0", 1, 0, torch.ones(3, 4)), - ] - reqs = InternalRequirements().require("hidden_states", count=2) - first = make_lazy_internal("m", FakeReader(rows), requirements=reqs) - second = make_lazy_internal("m", FakeReader(rows), requirements=reqs) - - first.require("hidden_states", count=1) - - with pytest.raises(IncompleteInternalError, match="expected 1 entries, found 2"): - first.hidden_states - assert len(second.hidden_states) == 2 diff --git a/tests/test_internal_mapper_e2e.py b/tests/test_internal_mapper_e2e.py deleted file mode 100644 index 227102833..000000000 --- a/tests/test_internal_mapper_e2e.py +++ /dev/null @@ -1,199 +0,0 @@ -"""End-to-end tests for the get_internal read path. - -Runs Qwen3-0.6B through generate_with_monitoring (writing captures to -ClickHouse via the native ring), then reads them back with -monitoring.internal_mapper.get_internal and checks the result lines up with -HuggingFace's native output. - -Requires a CUDA device, a reachable ClickHouse, the patched transformers fork, -and the model in the local cache; skips cleanly otherwise. - -ClickHouse connection: DMX_DB_HOST / DMX_DB_PORT (default localhost:9000). -""" -from __future__ import annotations - -import os -import uuid -from types import SimpleNamespace - -import pytest -import torch - -from tests._requirements import require_cuda, require_clickhouse, require_model_cache - -MODEL = "Qwen/Qwen3-0.6B" -MAX_NEW = 8 - -pytestmark = [ - pytest.mark.gpu, - pytest.mark.e2e, - pytest.mark.clickhouse, - pytest.mark.hf, - require_cuda(), - require_clickhouse(), - require_model_cache(MODEL), -] - - -@pytest.fixture(scope="module") -def deps(): - """Import the DMI HF stack; skip the module if any piece is missing.""" - try: - from transformers import AutoTokenizer - from transformers.models.qwen3_p.modeling_qwen3 import HookedQwen3ForCausalLM - from monitoring import ( - CaptureSchedule, HostEngineConfig, MonitoringConfig, MonitoringEngine, - ) - from monitoring._native_engine import ClickHouseClientConfig, StageConfig - from monitoring.clickhouse_reader import CHClickhouseDriverReadOnly - from monitoring.internal_mapper import get_internal - from integration.hf_adapter import generate_with_monitoring_dict - import clickhouse_driver - except Exception as exc: - pytest.skip(f"DMI HF stack unavailable: {exc}") - return SimpleNamespace(**{k: v for k, v in locals().items() if k != "exc"}) - - -def _db_host_port(): - return (os.environ.get("DMX_DB_HOST", "localhost"), - int(os.environ.get("DMX_DB_PORT", "9000"))) - - -def _wipe(deps, model_id): - host, port = _db_host_port() - deps.clickhouse_driver.Client(host=host, port=port).execute( - "ALTER TABLE default.offload DELETE WHERE model_id = %(m)s", {"m": model_id}) - - -def _monitored(deps, model, inputs, model_id, hook_selection="resid_pre"): - """Run generate_with_monitoring_dict into a fresh ClickHouse slot.""" - _wipe(deps, model_id) - host, port = _db_host_port() - db = deps.ClickHouseClientConfig() - db.host, db.port = host, port - db.username, db.password = "default", "" - db.database, db.table = "default", "offload" - db.create_database_if_missing = True - stage = deps.StageConfig.clickhouse_insert(db, parallelism=4, name="ch_insert") - engine = deps.MonitoringEngine( - config=deps.MonitoringConfig(schedule=deps.CaptureSchedule()), - model_id=model_id, - db_config=deps.HostEngineConfig(stages=[stage]), - ) - model.monitoring_engine = engine - try: - with torch.no_grad(): - return deps.generate_with_monitoring_dict( - model, **inputs, max_new_tokens=MAX_NEW, do_sample=False, - hook_selection=hook_selection, - ) - finally: - engine.close() - - -@pytest.fixture(scope="module") -def model_tok(deps): - tok = deps.AutoTokenizer.from_pretrained(MODEL) - if tok.pad_token_id is None: - tok.pad_token_id = tok.eos_token_id - tok.padding_side = "left" - model = deps.HookedQwen3ForCausalLM.from_pretrained( - MODEL, dtype=torch.float16, attn_implementation="eager", - ).to("cuda").eval() - return model, tok - - -@pytest.fixture(scope="module") -def single(deps, model_tok): - """One prompt (batch=1): monitored output + a plain-HF reference.""" - model, tok = model_tok - inputs = tok(["The capital of France is"], return_tensors="pt", padding=True).to("cuda") - with torch.no_grad(): - ref = model.generate(**inputs, max_new_tokens=MAX_NEW, do_sample=False) - model_id = f"test_internal::single::{uuid.uuid4().hex}"[:120] - out = _monitored(deps, model, inputs, model_id) - yield SimpleNamespace(out=out, ref=ref, model=model, model_id=model_id) - _wipe(deps, model_id) - - -@pytest.fixture(scope="module") -def ragged(deps, model_tok): - """A left-padded batch of unequal-length prompts.""" - model, tok = model_tok - inputs = tok(["The capital of France is", "Hi"], return_tensors="pt", padding=True).to("cuda") - model_id = f"test_internal::ragged::{uuid.uuid4().hex}"[:120] - out = _monitored(deps, model, inputs, model_id) - yield SimpleNamespace(out=out, n_prompts=2, model_id=model_id) - _wipe(deps, model_id) - - -# --- dict return shape ------------------------------------------------------ - -def test_returns_dict_output_with_dmi_internal(single): - assert single.out.sequences.dim() == 2 # [batch, total_len] - assert hasattr(single.out, "dmi_internal") - - -def test_output_matches_plain_hf(single): - # Monitoring must not change what the model generates. - assert torch.equal(single.out.sequences, single.ref) - - -# --- get_internal: structure ------------------------------------------------ - -def test_available_lists_captured_field(single, deps): - internal = single.out.dmi_internal - assert internal.available == ["hidden_states", "token_mask"] - - -def test_hidden_states_tuple_shape(single, deps): - internal = single.out.dmi_internal - hs = internal.hidden_states - token_mask = internal.token_mask - cfg = single.model.config - assert len(hs) == cfg.num_hidden_layers # one entry per block (resid_pre) - assert all(t.dim() == 3 for t in hs) # [batch, seq, hidden] - assert hs[0].shape[0] == 1 # batch - assert hs[0].shape[2] == cfg.hidden_size # hidden - assert all(t.shape == hs[0].shape for t in hs) # uniform across layers - assert token_mask.shape == hs[0].shape[:2] - assert token_mask.dtype == torch.bool - - -def test_token_count_drops_unforwarded_last(single, deps): - # The final generated token is never fed through a forward, so it has no - # captured hidden state: captured seq == sequence length - 1. - hs = single.out.dmi_internal.hidden_states - assert hs[0].shape[1] == single.out.sequences.shape[1] - 1 - - -# --- get_internal: model_id + reader handling ------------------------------ - -def test_source_model_id_with_explicit_reader(single, deps): - host, port = _db_host_port() - reader = deps.CHClickhouseDriverReadOnly(host=host, port=port) - from_out = single.out.dmi_internal.hidden_states - from_id = deps.get_internal(single.model_id, reader).hidden_states - assert len(from_out) == len(from_id) - assert torch.equal(from_out[0], from_id[0]) - - -def test_uncaptured_field_raises(single, deps): - internal = single.out.dmi_internal - with pytest.raises(AttributeError, match="not captured"): - internal.attention - - -# --- get_internal: ragged batch left-pad ----------------------------------- - -def test_ragged_batch_left_pads(ragged, deps): - hs = ragged.out.dmi_internal.hidden_states - layer0 = hs[0] - assert layer0.shape[0] == ragged.n_prompts - - nonzero = (layer0.abs().sum(dim=-1) > 0) # [batch, seq] real-token mask - counts = nonzero.sum(dim=1) - assert counts.min() < counts.max() # genuinely ragged -> padding added - assert bool(nonzero[:, -1].all()) # real tokens are right-aligned - short = int(counts.argmin()) - assert not bool(nonzero[short, 0]) # shorter request padded at the front From a2ad9f1aed15c66a19f313b564a5160d21051dd7 Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Tue, 25 Aug 2026 20:38:17 -0400 Subject: [PATCH 03/20] feat(capture): give a paginated query one stable identity CaptureQuery.query_hash hashed asdict(self), which includes `cursor` and `limit`. Page one and page two of the same logical query therefore hashed differently, so a cursor could not be bound to the query that issued it and two selections over identical filters were not comparable. Add CaptureQuery.filter_hash, computed over the filter fields only, and bind CaptureSelection to it (envelope bumped to version 2). query_hash is unchanged for callers that want whole-request identity. Also give `cursor` its own 2048-byte limit. It was validated against the 512-byte identifier limit, which a keyset cursor carrying the full five-part sort key can exceed -- and it would have failed deep into a pagination walk. --- src/dmi/storage/capture/model.py | 47 +++++++-- src/dmi/storage/capture/reader.py | 2 +- tests/test_capture_query_contract.py | 140 +++++++++++++++++++++++++++ 3 files changed, 179 insertions(+), 10 deletions(-) create mode 100644 tests/test_capture_query_contract.py diff --git a/src/dmi/storage/capture/model.py b/src/dmi/storage/capture/model.py index 6bd7adf15..33c4b239d 100644 --- a/src/dmi/storage/capture/model.py +++ b/src/dmi/storage/capture/model.py @@ -21,6 +21,11 @@ "float64": 8, } _TEXT_LIMIT = 512 +# A keyset cursor carries the whole catalog sort key, so it cannot fit the +# per-identifier text limit that bounds the fields it encodes. +_CURSOR_LIMIT = 2048 +# Fields of CaptureQuery that address a page rather than describe the filters. +_NON_FILTER_QUERY_FIELDS = frozenset({"cursor", "limit"}) _MAX_RANK = 32 @@ -48,13 +53,19 @@ class HydrationLimitError(CaptureStorageError): """A hydration plan exceeds the caller's byte limit.""" -def _validate_text(name: str, value: str | None, *, optional: bool = False) -> None: +class InvalidCursorError(CaptureStorageError): + """A pagination cursor is malformed or belongs to a different query.""" + + +def _validate_text( + name: str, value: str | None, *, optional: bool = False, limit: int = _TEXT_LIMIT +) -> None: if value is None: if optional: return raise ValueError(f"{name} is required") - if not isinstance(value, str) or not value or len(value.encode("utf-8")) > _TEXT_LIMIT: - raise ValueError(f"{name} must be non-empty UTF-8 within {_TEXT_LIMIT} bytes") + if not isinstance(value, str) or not value or len(value.encode("utf-8")) > limit: + raise ValueError(f"{name} must be non-empty UTF-8 within {limit} bytes") @dataclass(frozen=True, slots=True) @@ -242,10 +253,12 @@ class CaptureQuery: limit: int = 1000 def __post_init__(self) -> None: - for name in ("tenant_id", "experiment_id", "run_id", "session_id", "model_id", "cursor"): + for name in ("tenant_id", "experiment_id", "run_id", "session_id", "model_id"): value = getattr(self, name) if value is not None: _validate_text(name, value) + if self.cursor is not None: + _validate_text("cursor", self.cursor, limit=_CURSOR_LIMIT) if not 1 <= self.limit <= 10_000: raise ValueError("limit must be between 1 and 10000") if len(self.hook_names) > 128 or len(self.layer_numbers) > 1024: @@ -270,6 +283,22 @@ def query_hash(self) -> str: encoded = json.dumps(asdict(self), sort_keys=True, separators=(",", ":")).encode() return sha256(encoded).hexdigest() + @property + def filter_hash(self) -> str: + """Identity of the filters alone, excluding cursor and page size. + + ``query_hash`` covers the whole request, so it changes from page to + page. Cursors and selections bind to this instead, which is what makes + a paginated walk one identifiable query. + """ + selected = { + name: value + for name, value in asdict(self).items() + if name not in _NON_FILTER_QUERY_FIELDS + } + encoded = json.dumps(selected, sort_keys=True, separators=(",", ":")).encode() + return sha256(encoded).hexdigest() + @dataclass(frozen=True, slots=True) class CapturePage: @@ -283,7 +312,7 @@ class CaptureSelection: selection_id: str capture_ids: tuple[str, ...] catalog_watermark: str - query_hash: str + filter_hash: str @classmethod def create( @@ -291,7 +320,7 @@ def create( descriptors: Sequence[CaptureDescriptor], *, catalog_watermark: str, - query_hash: str, + filter_hash: str, ) -> CaptureSelection: ids = tuple(item.capture_id for item in descriptors) seen: set[str] = set() @@ -303,9 +332,9 @@ def create( seen.add(capture_id) identity = json.dumps( { - "version": 1, + "version": 2, "catalog_watermark": catalog_watermark, - "query_hash": query_hash, + "filter_hash": filter_hash, "capture_ids": ids, }, sort_keys=True, @@ -315,7 +344,7 @@ def create( selection_id=sha256(identity).hexdigest(), capture_ids=ids, catalog_watermark=catalog_watermark, - query_hash=query_hash, + filter_hash=filter_hash, ) diff --git a/src/dmi/storage/capture/reader.py b/src/dmi/storage/capture/reader.py index 028acc1c8..c5c3e849d 100644 --- a/src/dmi/storage/capture/reader.py +++ b/src/dmi/storage/capture/reader.py @@ -67,7 +67,7 @@ def select(self, query: CaptureQuery) -> CaptureSelection: return CaptureSelection.create( page.items, catalog_watermark=page.watermark, - query_hash=query.query_hash, + filter_hash=query.filter_hash, ) def estimate(self, selection: CaptureSelection) -> HydrationEstimate: diff --git a/tests/test_capture_query_contract.py b/tests/test_capture_query_contract.py new file mode 100644 index 000000000..bf1005685 --- /dev/null +++ b/tests/test_capture_query_contract.py @@ -0,0 +1,140 @@ +from __future__ import annotations + +from dataclasses import fields + +import pytest + +from benchmarks.bench_capture_catalog import synthetic_descriptors +from dmi.storage.capture import CaptureQuery, CaptureSelection + + +pytestmark = pytest.mark.cpu + + +# One alternate value per filter field. The completeness assertion below fails +# when a filter is added to CaptureQuery without being covered here, which is +# the case that would silently drop a field out of ``filter_hash``. +_FILTER_ALTERNATES = { + "tenant_id": "tenant-z", + "experiment_id": "experiment-z", + "run_id": "run-z", + "session_id": "session-z", + "model_id": "model-z", + "hook_names": ("hook-z",), + "layer_numbers": (99,), + "captured_after_ns": 1_700_000_000_000_000_000, + "captured_before_ns": 1_800_000_000_000_000_000, +} +_NON_FILTER_FIELDS = {"cursor", "limit"} + + +def _query(**overrides) -> CaptureQuery: + base = { + "tenant_id": "tenant-a", + "experiment_id": "experiment-a", + "run_id": "run-a", + "hook_names": ("hook-a",), + "layer_numbers": (3,), + } + return CaptureQuery(**{**base, **overrides}) + + +def test_filter_alternates_cover_every_filter_field(): + declared = {item.name for item in fields(CaptureQuery)} + assert set(_FILTER_ALTERNATES) == declared - _NON_FILTER_FIELDS + + +def test_filter_hash_is_stable_across_pages(): + first = _query() + second = _query(cursor="ZW5jb2RlZC1jdXJzb3I") + + assert first.filter_hash == second.filter_hash + + +def test_filter_hash_is_stable_across_page_sizes(): + assert _query(limit=10).filter_hash == _query(limit=5_000).filter_hash + + +@pytest.mark.parametrize("name,value", sorted(_FILTER_ALTERNATES.items())) +def test_filter_hash_changes_with_every_filter(name: str, value): + assert _query().filter_hash != _query(**{name: value}).filter_hash + + +def test_query_hash_still_separates_pages(): + # Preserved for backward compatibility: query_hash covers the full request, + # cursor and limit included. Cursor binding uses filter_hash instead. + first = _query() + second = _query(cursor="ZW5jb2RlZC1jdXJzb3I") + + assert first.query_hash != second.query_hash + + +def test_cursor_accepts_a_full_keyset_payload(): + cursor = "c" * 2048 + + assert CaptureQuery(cursor=cursor).cursor == cursor + + +def test_cursor_rejects_payloads_beyond_its_limit(): + with pytest.raises(ValueError, match="cursor"): + CaptureQuery(cursor="c" * 2049) + + +def test_cursor_rejects_empty_and_non_string_values(): + with pytest.raises(ValueError, match="cursor"): + CaptureQuery(cursor="") + with pytest.raises(ValueError, match="cursor"): + CaptureQuery(cursor=b"bytes") # type: ignore[arg-type] + + +@pytest.mark.parametrize( + "name", ("tenant_id", "experiment_id", "run_id", "session_id", "model_id") +) +def test_filter_text_fields_keep_the_shared_limit(name: str): + assert CaptureQuery(**{name: "x" * 512}) + + with pytest.raises(ValueError, match=name): + CaptureQuery(**{name: "x" * 513}) + + +def test_selection_binds_to_the_filter_hash(): + descriptors = synthetic_descriptors(2) + query = _query() + + selection = CaptureSelection.create( + descriptors, catalog_watermark="w-1", filter_hash=query.filter_hash + ) + + assert selection.filter_hash == query.filter_hash + + +def test_selection_identity_separates_different_filters(): + descriptors = synthetic_descriptors(2) + same_filters = tuple( + CaptureSelection.create( + descriptors, catalog_watermark="w-1", filter_hash=_query(**overrides).filter_hash + ) + for overrides in ({}, {"cursor": "cGFnZS10d28"}, {"limit": 25}) + ) + other_filters = CaptureSelection.create( + descriptors, catalog_watermark="w-1", filter_hash=_query(run_id="run-z").filter_hash + ) + + # Paging through one query yields one selection identity; changing a filter + # yields a different one. + assert len({item.selection_id for item in same_filters}) == 1 + assert other_filters.selection_id != same_filters[0].selection_id + + +def test_selection_identity_separates_watermarks(): + descriptors = synthetic_descriptors(2) + filter_hash = _query().filter_hash + + first = CaptureSelection.create( + descriptors, catalog_watermark="w-1", filter_hash=filter_hash + ) + second = CaptureSelection.create( + descriptors, catalog_watermark="w-2", filter_hash=filter_hash + ) + + assert first.selection_id != second.selection_id From c9af1cd434454a7db8c6a98a322292c38cc04b45 Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Tue, 25 Aug 2026 20:38:42 -0400 Subject: [PATCH 04/20] feat(capture): add keyset pagination cursors A cursor addresses a position in the catalog's own sort order -- (tenant_id, experiment_id, run_id, captured_at_ns, capture_id) -- so page latency does not grow with depth and a concurrent insert cannot shift a page boundary. Each cursor carries the pinned watermark and the issuing query's filter_hash, so it cannot be replayed against different filters or against a catalog that cannot serve its snapshot. Decoding is strict: unknown version, missing or unexpected envelope fields, wrong key arity, wrong element types, empty strings and out-of-range timestamps all raise InvalidCursorError. Decode with validate=True. Python's default base64 decoder silently discards characters outside the alphabet, so a cursor with four injected junk characters decoded to a byte-identical payload and was accepted. That matters beyond tidiness: cursors are deliberately unauthenticated, on the reasoning that a forged cursor can only reach the caller's own keyspace -- which only holds if tampering is detected as malformed. --- src/dmi/storage/capture/__init__.py | 7 + src/dmi/storage/capture/cursor.py | 152 ++++++++++++++++++++ tests/test_capture_cursor.py | 211 ++++++++++++++++++++++++++++ 3 files changed, 370 insertions(+) create mode 100644 src/dmi/storage/capture/cursor.py create mode 100644 tests/test_capture_cursor.py diff --git a/src/dmi/storage/capture/__init__.py b/src/dmi/storage/capture/__init__.py index 15e9aa88b..0f511136c 100644 --- a/src/dmi/storage/capture/__init__.py +++ b/src/dmi/storage/capture/__init__.py @@ -14,6 +14,7 @@ HydratedCapture, HydrationEstimate, HydrationLimitError, + InvalidCursorError, ObjectInfo, ObjectPage, PackConflictError, @@ -58,6 +59,7 @@ ReconcileResult, ) from .clickhouse_catalog import ClickHouseCatalogConfig, ClickHouseCatalogWriter +from .cursor import Cursor, CursorKey, decode_cursor, encode_cursor from .s3 import S3PackStore, S3StoreConfig from .spool import ( DurablePackSink, @@ -92,6 +94,8 @@ "CatalogWriter", "ClickHouseCatalogConfig", "ClickHouseCatalogWriter", + "Cursor", + "CursorKey", "DuplicateCaptureError", "DirectPackSink", "DurablePackSink", @@ -106,6 +110,7 @@ "IndexEvent", "IndexFailure", "IndexResult", + "InvalidCursorError", "ObjectInfo", "ObjectPage", "OverloadPolicy", @@ -144,5 +149,7 @@ "UploadEvent", "UploadFailure", "UploadSnapshot", + "decode_cursor", + "encode_cursor", "object_key_for", ] diff --git a/src/dmi/storage/capture/cursor.py b/src/dmi/storage/capture/cursor.py new file mode 100644 index 000000000..e882d29be --- /dev/null +++ b/src/dmi/storage/capture/cursor.py @@ -0,0 +1,152 @@ +"""Keyset pagination cursors for bounded catalog search. + +A cursor addresses a position in the catalog's own sort order -- +``(tenant_id, experiment_id, run_id, captured_at_ns, capture_id)`` -- so page +latency does not grow with depth and a concurrent insert cannot shift a page +boundary. + +The encoding is deliberately strict. A cursor is data handed back by a caller, +so every field is validated on the way in and any deviation raises +:class:`InvalidCursorError` rather than silently yielding a different page. +""" + +from __future__ import annotations + +from base64 import b64decode, urlsafe_b64encode +from binascii import Error as BinasciiError +from dataclasses import dataclass +import json + +from .model import _CURSOR_LIMIT, InvalidCursorError + + +_CURSOR_VERSION = 1 +_ENVELOPE_FIELDS = frozenset({"v", "w", "fh", "k"}) +_UINT64_MAX = 2**64 - 1 + + +@dataclass(frozen=True, slots=True) +class CursorKey: + """A position in the catalog sort order.""" + + tenant_id: str + experiment_id: str + run_id: str + captured_at_ns: int + capture_id: str + + +@dataclass(frozen=True, slots=True) +class Cursor: + """A decoded cursor, bound to the query and snapshot that issued it.""" + + version: int + watermark: int + filter_hash: str + key: CursorKey + + +def _require(condition: object, message: str) -> None: + if not condition: + raise InvalidCursorError(message) + + +def _uint64(value: object, name: str) -> int: + _require(type(value) is int, f"{name} must be an integer") + _require(0 <= value <= _UINT64_MAX, f"{name} must fit UInt64") + return value # type: ignore[return-value] + + +def _text(value: object, name: str) -> str: + _require(isinstance(value, str) and value, f"{name} must be a non-empty string") + return value # type: ignore[return-value] + + +def encode_cursor(key: CursorKey, *, watermark: int, filter_hash: str) -> str: + """Encode a position as an opaque, url-safe cursor.""" + payload = { + "v": _CURSOR_VERSION, + "w": _uint64(watermark, "watermark"), + "fh": _text(filter_hash, "filter_hash"), + "k": [ + _text(key.tenant_id, "tenant_id"), + _text(key.experiment_id, "experiment_id"), + _text(key.run_id, "run_id"), + _uint64(key.captured_at_ns, "captured_at_ns"), + _text(key.capture_id, "capture_id"), + ], + } + raw = json.dumps(payload, separators=(",", ":"), sort_keys=True).encode("utf-8") + encoded = urlsafe_b64encode(raw).rstrip(b"=").decode("ascii") + if len(encoded.encode("utf-8")) > _CURSOR_LIMIT: + raise InvalidCursorError( + f"encoded cursor exceeds the cursor limit of {_CURSOR_LIMIT} bytes" + ) + return encoded + + +def decode_cursor(cursor: str, *, filter_hash: str, max_watermark: int) -> Cursor: + """Decode and validate a cursor against the query and catalog presenting it. + + ``filter_hash`` binds the cursor to one set of filters, so a cursor cannot + be replayed against a different query. ``max_watermark`` is the catalog's + current high-water version; a cursor above it does not describe a snapshot + this catalog can serve. + """ + _require(isinstance(cursor, str) and cursor, "cursor must be a non-empty string") + _require( + len(cursor.encode("utf-8")) <= _CURSOR_LIMIT, + f"cursor exceeds the cursor limit of {_CURSOR_LIMIT} bytes", + ) + + # ``validate=True`` matters here: the default decoder silently discards + # characters outside the base64 alphabet, so a cursor with injected junk + # would decode to the same payload and pass every check below. + padding = "=" * (-len(cursor) % 4) + try: + raw = b64decode(cursor + padding, altchars=b"-_", validate=True) + except (BinasciiError, ValueError) as exc: + raise InvalidCursorError("cursor is not valid url-safe base64") from exc + + try: + payload = json.loads(raw) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise InvalidCursorError("cursor does not contain valid JSON") from exc + + _require(isinstance(payload, dict), "cursor payload must be a JSON object") + unexpected = set(payload) - _ENVELOPE_FIELDS + _require(not unexpected, f"cursor has unexpected fields: {sorted(unexpected)}") + missing = _ENVELOPE_FIELDS - set(payload) + _require(not missing, f"cursor is missing fields: {sorted(missing)}") + + _require(payload["v"] == _CURSOR_VERSION, f"unsupported cursor version: {payload['v']!r}") + + watermark = _uint64(payload["w"], "cursor watermark") + _require( + watermark <= max_watermark, + "cursor watermark is ahead of the catalog: " + f"{watermark} > {max_watermark}", + ) + + encoded_hash = _text(payload["fh"], "cursor filter hash") + _require( + encoded_hash == filter_hash, + "cursor was issued for different filters", + ) + + key = payload["k"] + _require(isinstance(key, list), "cursor key must be a JSON array") + _require(len(key) == 5, f"cursor key must have five elements, got {len(key)}") + + return Cursor( + version=_CURSOR_VERSION, + watermark=watermark, + filter_hash=encoded_hash, + key=CursorKey( + tenant_id=_text(key[0], "cursor tenant_id"), + experiment_id=_text(key[1], "cursor experiment_id"), + run_id=_text(key[2], "cursor run_id"), + captured_at_ns=_uint64(key[3], "cursor captured_at_ns"), + capture_id=_text(key[4], "cursor capture_id"), + ), + ) diff --git a/tests/test_capture_cursor.py b/tests/test_capture_cursor.py new file mode 100644 index 000000000..3d7c4d691 --- /dev/null +++ b/tests/test_capture_cursor.py @@ -0,0 +1,211 @@ +from __future__ import annotations + +from base64 import urlsafe_b64encode +import json + +import pytest + +from dmi.storage.capture import ( + CaptureQuery, + Cursor, + CursorKey, + InvalidCursorError, + decode_cursor, + encode_cursor, +) + + +pytestmark = pytest.mark.cpu + + +_FILTER_HASH = "a" * 64 +_WATERMARK = 1_756_142_093_000_000_000 +_MAX_WATERMARK = 1_756_142_099_000_000_000 + + +def _key(**overrides) -> CursorKey: + base = { + "tenant_id": "tenant-a", + "experiment_id": "experiment-a", + "run_id": "run-a", + "captured_at_ns": 1_756_142_090_000_000_000, + "capture_id": "capture-0f31", + } + return CursorKey(**{**base, **overrides}) + + +def _encoded(payload: object) -> str: + raw = json.dumps(payload, separators=(",", ":")).encode("utf-8") + return urlsafe_b64encode(raw).rstrip(b"=").decode("ascii") + + +def _valid_payload(**overrides) -> dict: + payload = { + "v": 1, + "w": _WATERMARK, + "fh": _FILTER_HASH, + "k": [ + "tenant-a", + "experiment-a", + "run-a", + 1_756_142_090_000_000_000, + "capture-0f31", + ], + } + payload.update(overrides) + return payload + + +def _decode(cursor: str) -> Cursor: + return decode_cursor( + cursor, filter_hash=_FILTER_HASH, max_watermark=_MAX_WATERMARK + ) + + +# --- round trip ------------------------------------------------------------- + + +def test_round_trip_preserves_every_field(): + encoded = encode_cursor(_key(), watermark=_WATERMARK, filter_hash=_FILTER_HASH) + + decoded = _decode(encoded) + + assert decoded == Cursor( + version=1, watermark=_WATERMARK, filter_hash=_FILTER_HASH, key=_key() + ) + + +def test_encoding_is_deterministic(): + first = encode_cursor(_key(), watermark=_WATERMARK, filter_hash=_FILTER_HASH) + second = encode_cursor(_key(), watermark=_WATERMARK, filter_hash=_FILTER_HASH) + + assert first == second + + +def test_encoded_cursor_is_accepted_by_capture_query(): + encoded = encode_cursor(_key(), watermark=_WATERMARK, filter_hash=_FILTER_HASH) + + # Unpadded url-safe base64 keeps the cursor inside CaptureQuery's validator + # and free of characters that would need escaping in a URL. + assert "=" not in encoded + assert CaptureQuery(cursor=encoded).cursor == encoded + + +def test_encode_refuses_a_cursor_beyond_the_query_limit(): + with pytest.raises(InvalidCursorError, match="cursor"): + encode_cursor( + _key(tenant_id="t" * 512, experiment_id="e" * 512, run_id="r" * 512, + capture_id="c" * 512), + watermark=_WATERMARK, + filter_hash=_FILTER_HASH, + ) + + +# --- binding ---------------------------------------------------------------- + + +def test_decode_rejects_a_cursor_from_different_filters(): + encoded = encode_cursor(_key(), watermark=_WATERMARK, filter_hash="b" * 64) + + with pytest.raises(InvalidCursorError, match="filter"): + _decode(encoded) + + +def test_decode_rejects_a_watermark_above_the_catalog(): + encoded = encode_cursor( + _key(), watermark=_MAX_WATERMARK + 1, filter_hash=_FILTER_HASH + ) + + with pytest.raises(InvalidCursorError, match="watermark"): + _decode(encoded) + + +def test_decode_accepts_a_watermark_equal_to_the_catalog(): + encoded = encode_cursor( + _key(), watermark=_MAX_WATERMARK, filter_hash=_FILTER_HASH + ) + + assert _decode(encoded).watermark == _MAX_WATERMARK + + +# --- malformed input -------------------------------------------------------- + + +@pytest.mark.parametrize( + "cursor", + ( + "", + "!!!not base64!!!", + "z", + urlsafe_b64encode(b"not json").rstrip(b"=").decode("ascii"), + _encoded([1, 2, 3]), + _encoded("a string"), + _encoded(None), + ), + ids=( + "empty", "non-base64", "truncated", "not-json", + "json-array", "json-string", "json-null", + ), +) +def test_decode_rejects_unparseable_cursors(cursor: str): + with pytest.raises(InvalidCursorError): + _decode(cursor) + + +@pytest.mark.parametrize( + "payload,reason", + ( + (_valid_payload(v=2), "unknown-version"), + (_valid_payload(v="1"), "version-not-int"), + (_valid_payload(w=-1), "negative-watermark"), + (_valid_payload(w=2**64), "watermark-overflow"), + (_valid_payload(w="1"), "watermark-not-int"), + (_valid_payload(fh=123), "filter-hash-not-str"), + (_valid_payload(k=["a", "b", "c", 1]), "key-too-short"), + (_valid_payload(k=["a", "b", "c", 1, "d", "e"]), "key-too-long"), + (_valid_payload(k="not-a-list"), "key-not-list"), + (_valid_payload(k=["a", "b", "c", "not-int", "d"]), "timestamp-not-int"), + (_valid_payload(k=["a", "b", "c", -1, "d"]), "negative-timestamp"), + (_valid_payload(k=[1, "b", "c", 1, "d"]), "tenant-not-str"), + (_valid_payload(k=["", "b", "c", 1, "d"]), "empty-tenant"), + (_valid_payload(k=["a", "b", "c", 1, ""]), "empty-capture-id"), + ), + ids=lambda value: value if isinstance(value, str) else "", +) +def test_decode_rejects_malformed_payloads(payload: dict, reason: str): + with pytest.raises(InvalidCursorError): + _decode(_encoded(payload)) + + +@pytest.mark.parametrize("missing", ("v", "w", "fh", "k")) +def test_decode_rejects_missing_fields(missing: str): + payload = _valid_payload() + del payload[missing] + + with pytest.raises(InvalidCursorError): + _decode(_encoded(payload)) + + +@pytest.mark.parametrize("junk", ("!!!!", "....", " ", "\n\n\n\n")) +def test_decode_rejects_characters_outside_the_base64_alphabet(junk: str): + # Python's default base64 decoder discards these silently. Injecting a + # multiple of four keeps the padding arithmetic intact, so a lax decoder + # accepts the tampered cursor and returns the original payload. + encoded = encode_cursor(_key(), watermark=_WATERMARK, filter_hash=_FILTER_HASH) + tampered = encoded[:10] + junk + encoded[10:] + + with pytest.raises(InvalidCursorError, match="base64"): + _decode(tampered) + + +def test_decode_rejects_unexpected_fields(): + # A strict envelope keeps a future field from being silently ignored by an + # older reader. + with pytest.raises(InvalidCursorError): + _decode(_encoded(_valid_payload(extra="surprise"))) + + +def test_invalid_cursor_error_is_a_capture_storage_error(): + from dmi.storage.capture import CaptureStorageError + + assert issubclass(InvalidCursorError, CaptureStorageError) From 9130464e76e9f5bb4bf8d4a84a29e93cff108f7b Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Tue, 25 Aug 2026 20:39:08 -0400 Subject: [PATCH 05/20] feat(capture): read the ClickHouse catalog as a pinned snapshot ClickHouseCatalogWriter was write-only; the sole CaptureCatalog implementation in the tree was a test fake. ClickHouseCaptureCatalog implements search and get_by_ids so CaptureReader can run against the real catalog. Reads go to *_capture_raw with explicit argMax, not the FINAL views. FINAL collapses duplicates to the highest version present and only then applies predicates, so `FINAL ... WHERE index_version <= W` drops any capture re-indexed above W instead of falling back to its value at W. Measured on 26.9.1, with capture-a re-indexed at v2 and a snapshot at watermark 1: FINAL + predicate -> [('capture-b', 128)] # a vanishes argMax at watermark -> [('capture-a', 64), ('capture-b', 128)] # correct Aggregates are deliberately unaliased. Naming one after its own source column shadows that column across the statement, and ClickHouse then rejects any filter on it with "Aggregate function ... is found in WHERE in query". Rows map onto descriptors positionally, so server-side names are never read. The keyset comparison sits in WHERE rather than HAVING: all five sort-key columns are the table's ORDER BY prefix, so it prunes granules on the primary index before grouping. Filters likewise apply pre-group, which is safe only because a descriptor derives from an immutable pack footer -- re-indexing a capture rewrites identical values. ClickHouseReaderConfig carries per-query scan-row, scan-byte and execution-time ceilings with read_overflow_mode=throw, so a breach surfaces as a server exception rather than a long scan. --- src/dmi/storage/capture/__init__.py | 3 + src/dmi/storage/capture/clickhouse_reader.py | 355 +++++++++++++++++ tests/test_clickhouse_capture_reader.py | 392 +++++++++++++++++++ tests/test_clickhouse_reader_live.py | 176 +++++++++ 4 files changed, 926 insertions(+) create mode 100644 src/dmi/storage/capture/clickhouse_reader.py create mode 100644 tests/test_clickhouse_capture_reader.py create mode 100644 tests/test_clickhouse_reader_live.py diff --git a/src/dmi/storage/capture/__init__.py b/src/dmi/storage/capture/__init__.py index 0f511136c..07e29143a 100644 --- a/src/dmi/storage/capture/__init__.py +++ b/src/dmi/storage/capture/__init__.py @@ -59,6 +59,7 @@ ReconcileResult, ) from .clickhouse_catalog import ClickHouseCatalogConfig, ClickHouseCatalogWriter +from .clickhouse_reader import ClickHouseCaptureCatalog, ClickHouseReaderConfig from .cursor import Cursor, CursorKey, decode_cursor, encode_cursor from .s3 import S3PackStore, S3StoreConfig from .spool import ( @@ -94,6 +95,8 @@ "CatalogWriter", "ClickHouseCatalogConfig", "ClickHouseCatalogWriter", + "ClickHouseCaptureCatalog", + "ClickHouseReaderConfig", "Cursor", "CursorKey", "DuplicateCaptureError", diff --git a/src/dmi/storage/capture/clickhouse_reader.py b/src/dmi/storage/capture/clickhouse_reader.py new file mode 100644 index 000000000..992cd177c --- /dev/null +++ b/src/dmi/storage/capture/clickhouse_reader.py @@ -0,0 +1,355 @@ +"""ClickHouse-backed catalog reads for bounded analysis. + +The writer in :mod:`.clickhouse_catalog` lands descriptor rows into a +``ReplacingMergeTree`` keyed by ``index_version``. This module reads them back +as a *snapshot*: every result is pinned to a watermark so a selection resolves +to the same captures for as long as it lives. + +That pinning is why these queries read ``*_capture_raw`` directly instead of the +public ``FINAL`` views. ``FINAL`` collapses duplicates to the highest version +present and only then applies predicates, so ``FINAL`` with +``index_version <= W`` drops any capture re-indexed above ``W`` rather than +falling back to its version at ``W``. Explicit ``argMax`` aggregation gives the +row as of the watermark, which is what a stable cursor walk requires. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Mapping, Sequence +from uuid import UUID + +from .clickhouse_catalog import ( + _CAPTURE_COLUMNS, + ClickHouseCatalogConfig, + ClickHouseClient, + _identifier, + _quoted, +) +from .cursor import CursorKey, decode_cursor, encode_cursor +from .model import ( + CaptureDescriptor, + CaptureMetadata, + CapturePage, + CaptureQuery, + PackFormatError, + PayloadLocator, +) + + +# The catalog sort key, and therefore the keyset pagination order. These are the +# table's ORDER BY prefix, so the tuple comparison that advances a page prunes +# granules on the primary index instead of scanning. +_SORT_KEY = ("tenant_id", "experiment_id", "run_id", "captured_at_ns", "capture_id") +_SORT_KEY_SET = frozenset(_SORT_KEY) + +# Projection order is the writer's column order minus ``index_version``, so a +# result row maps positionally onto the descriptor fields. +_PROJECTION = _CAPTURE_COLUMNS[:-1] + +_EQUALITY_FILTERS = ("tenant_id", "experiment_id", "run_id", "session_id", "model_id") + + +@dataclass(frozen=True, slots=True) +class ClickHouseReaderConfig: + """Bounds applied to every catalog read. + + The four ``max_*`` settings ride on the per-query settings map, so a breach + surfaces as a server-side exception instead of a long-running scan. + """ + + database: str = "default" + table_prefix: str = "dmi" + max_capture_ids: int = 10_000 + max_rows_to_read: int = 50_000_000 + max_bytes_to_read: int = 4 * 1024**3 + max_execution_time: int = 15 + + def __post_init__(self) -> None: + _identifier(self.database) + _identifier(self.table_prefix) + for name in ( + "max_capture_ids", + "max_rows_to_read", + "max_bytes_to_read", + "max_execution_time", + ): + value = getattr(self, name) + if type(value) is not int or value <= 0: + raise ValueError(f"{name} must be a positive integer") + + @classmethod + def from_catalog(cls, config: ClickHouseCatalogConfig) -> ClickHouseReaderConfig: + """Read with the database and prefix a writer was configured with.""" + return cls(database=config.database, table_prefix=config.table_prefix) + + @property + def settings(self) -> dict[str, object]: + return { + "max_rows_to_read": self.max_rows_to_read, + "max_bytes_to_read": self.max_bytes_to_read, + "max_execution_time": self.max_execution_time, + "read_overflow_mode": "throw", + "timeout_overflow_mode": "throw", + } + + +def _text(value: object, name: str) -> str: + """Normalise a ClickHouse string column to ``str``.""" + if isinstance(value, bytes): + # FixedString columns come back padded with NULs on some drivers. + return value.rstrip(b"\x00").decode("utf-8") + if isinstance(value, UUID): + return str(value) + if not isinstance(value, str): + raise PackFormatError(f"catalog returned a non-text {name}: {type(value).__name__}") + return value + + +def _integer(value: object, name: str) -> int: + if type(value) is bool or not isinstance(value, int): + raise PackFormatError(f"catalog returned a non-integer {name}") + return value + + +class ClickHouseCaptureCatalog: + """A :class:`~.model.CaptureCatalog` backed by the derived ClickHouse catalog.""" + + def __init__( + self, + client: ClickHouseClient, + config: ClickHouseReaderConfig | None = None, + ) -> None: + self._client = client + self._config = config or ClickHouseReaderConfig() + self._capture_raw = f"{self._config.table_prefix}_capture_raw" + + @property + def config(self) -> ClickHouseReaderConfig: + return self._config + + # -- public API --------------------------------------------------------- + + def current_watermark(self) -> str: + """The catalog's high-water ``index_version``. + + ``index_version`` is not part of the sort key, so this is a full-column + aggregate rather than an index lookup -- cheap in a columnar engine, but + not free, and it runs once per search. + """ + rows = self._client.execute( + f"SELECT max(index_version) FROM {self._qualified()}", + settings=self._config.settings, + ) + if not rows or rows[0][0] is None: + return "0" + return str(_integer(rows[0][0], "watermark")) + + def search(self, query: CaptureQuery) -> CapturePage: + max_watermark = int(self.current_watermark()) + filter_hash = query.filter_hash + + if query.cursor is None: + watermark = max_watermark + after: CursorKey | None = None + else: + cursor = decode_cursor( + query.cursor, filter_hash=filter_hash, max_watermark=max_watermark + ) + watermark = cursor.watermark + after = cursor.key + + clauses, params = self._filters(query, watermark=watermark, after=after) + # One row beyond the page tells us whether a cursor is owed, without a + # second counting query. + params["row_limit"] = query.limit + 1 + sql = ( + f"SELECT {self._projection()} FROM {self._qualified()} " + f"WHERE {' AND '.join(clauses)} " + f"GROUP BY {', '.join(_quoted(name) for name in _SORT_KEY)} " + f"ORDER BY {', '.join(_quoted(name) for name in _SORT_KEY)} " + "LIMIT %(row_limit)s" + ) + rows = self._client.execute(sql, params, settings=self._config.settings) + descriptors = tuple(self._descriptor(row) for row in rows) + + next_cursor = None + if len(descriptors) > query.limit: + descriptors = descriptors[: query.limit] + next_cursor = encode_cursor( + _key_of(descriptors[-1]), watermark=watermark, filter_hash=filter_hash + ) + return CapturePage( + items=descriptors, next_cursor=next_cursor, watermark=str(watermark) + ) + + def get_by_ids( + self, capture_ids: Sequence[str], *, watermark: str + ) -> tuple[CaptureDescriptor, ...]: + if len(capture_ids) > self._config.max_capture_ids: + raise ValueError( + f"capture id lookup exceeds max_capture_ids: " + f"{len(capture_ids)} > {self._config.max_capture_ids}" + ) + if not capture_ids: + return () + params = { + "watermark": _parse_watermark(watermark), + "capture_ids": list(capture_ids), + } + sql = ( + f"SELECT {self._projection()} FROM {self._qualified()} " + "WHERE index_version <= %(watermark)s AND capture_id IN %(capture_ids)s " + f"GROUP BY {', '.join(_quoted(name) for name in _SORT_KEY)}" + ) + rows = self._client.execute(sql, params, settings=self._config.settings) + return tuple(self._descriptor(row) for row in rows) + + # -- SQL construction --------------------------------------------------- + + def _qualified(self) -> str: + return f"{_quoted(self._config.database)}.{_quoted(self._capture_raw)}" + + @staticmethod + def _projection() -> str: + # Sort-key columns are grouped on, so they project directly; everything + # else resolves to its value at the highest version <= the watermark. + # + # Deliberately unaliased: naming an aggregate after its own source + # column shadows that column everywhere else in the statement, and + # ClickHouse then rejects a filter on it with "Aggregate function ... + # is found in WHERE in query". Rows map onto descriptors positionally, + # so the server-side column names are never read. + return ", ".join( + _quoted(name) + if name in _SORT_KEY_SET + else f"argMax({_quoted(name)}, index_version)" + for name in _PROJECTION + ) + + @staticmethod + def _filters( + query: CaptureQuery, *, watermark: int, after: CursorKey | None + ) -> tuple[list[str], dict[str, object]]: + clauses = ["index_version <= %(watermark)s"] + params: dict[str, object] = {"watermark": watermark} + + # Equality and range filters apply to raw rows before grouping. That is + # safe because a descriptor is derived from an immutable pack footer, so + # re-indexing one capture rewrites identical values; it is also much + # faster, since these predicates reach the primary index. + for name in _EQUALITY_FILTERS: + value = getattr(query, name) + if value is not None: + clauses.append(f"{_quoted(name)} = %({name})s") + params[name] = value + if query.hook_names: + clauses.append("hook_name IN %(hook_names)s") + params["hook_names"] = list(query.hook_names) + if query.layer_numbers: + clauses.append("layer_number IN %(layer_numbers)s") + params["layer_numbers"] = list(query.layer_numbers) + if query.captured_after_ns is not None: + clauses.append("captured_at_ns >= %(captured_after_ns)s") + params["captured_after_ns"] = query.captured_after_ns + if query.captured_before_ns is not None: + clauses.append("captured_at_ns <= %(captured_before_ns)s") + params["captured_before_ns"] = query.captured_before_ns + + if after is not None: + columns = ", ".join(_quoted(name) for name in _SORT_KEY) + placeholders = ", ".join(f"%(after_{name})s" for name in _SORT_KEY) + clauses.append(f"({columns}) > ({placeholders})") + params.update( + { + "after_tenant_id": after.tenant_id, + "after_experiment_id": after.experiment_id, + "after_run_id": after.run_id, + "after_captured_at_ns": after.captured_at_ns, + "after_capture_id": after.capture_id, + } + ) + return clauses, params + + # -- row mapping -------------------------------------------------------- + + @staticmethod + def _descriptor(row: Sequence[object]) -> CaptureDescriptor: + if len(row) != len(_PROJECTION): + raise PackFormatError( + f"catalog row has {len(row)} columns, expected {len(_PROJECTION)}" + ) + value: Mapping[str, object] = dict(zip(_PROJECTION, row)) + shape = value["shape"] + if not isinstance(shape, (list, tuple)): + raise PackFormatError("catalog returned a non-array shape") + + adapter_revision = value["adapter_revision"] + metadata = CaptureMetadata( + capture_id=_text(value["capture_id"], "capture_id"), + tenant_id=_text(value["tenant_id"], "tenant_id"), + experiment_id=_text(value["experiment_id"], "experiment_id"), + run_id=_text(value["run_id"], "run_id"), + session_id=_text(value["session_id"], "session_id"), + request_id=_text(value["request_id"], "request_id"), + sequence_id=_text(value["sequence_id"], "sequence_id"), + model_id=_text(value["model_id"], "model_id"), + model_revision=_text(value["model_revision"], "model_revision"), + adapter_revision=( + None + if adapter_revision is None + else _text(adapter_revision, "adapter_revision") + ), + capture_policy_version=_text( + value["capture_policy_version"], "capture_policy_version" + ), + hook_name=_text(value["hook_name"], "hook_name"), + layer_number=_integer(value["layer_number"], "layer_number"), + producer_rank=_integer(value["producer_rank"], "producer_rank"), + step_number=_integer(value["step_number"], "step_number"), + token_start=_integer(value["token_start"], "token_start"), + token_end=_integer(value["token_end"], "token_end"), + batch_position=_integer(value["batch_position"], "batch_position"), + dtype=_text(value["dtype"], "dtype"), + shape=tuple(_integer(dim, "shape") for dim in shape), + captured_at_ns=_integer(value["captured_at_ns"], "captured_at_ns"), + ) + locator = PayloadLocator( + pack_id=_text(value["pack_id"], "pack_id"), + store_id=_text(value["store_id"], "store_id"), + object_key=_text(value["object_key"], "object_key"), + object_bytes=_integer(value["object_bytes"], "object_bytes"), + pack_checksum=_text(value["pack_checksum"], "pack_checksum"), + pack_record_count=_integer( + value["pack_record_count"], "pack_record_count" + ), + offset=_integer(value["payload_offset"], "payload_offset"), + stored_length=_integer(value["stored_length"], "stored_length"), + decoded_length=_integer(value["decoded_length"], "decoded_length"), + codec=_text(value["codec"], "codec"), + checksum=_text(value["payload_checksum"], "payload_checksum"), + ) + try: + return CaptureDescriptor(metadata=metadata, locator=locator) + except ValueError as exc: # pragma: no cover - defensive + raise PackFormatError(f"invalid catalog descriptor: {exc}") from exc + + +def _key_of(descriptor: CaptureDescriptor) -> CursorKey: + metadata = descriptor.metadata + return CursorKey( + tenant_id=metadata.tenant_id, + experiment_id=metadata.experiment_id, + run_id=metadata.run_id, + captured_at_ns=metadata.captured_at_ns, + capture_id=metadata.capture_id, + ) + + +def _parse_watermark(watermark: str) -> int: + if not isinstance(watermark, str) or not watermark.isdigit(): + raise ValueError("watermark must be a decimal string") + value = int(watermark) + if value >= 2**64: + raise ValueError("watermark must fit UInt64") + return value diff --git a/tests/test_clickhouse_capture_reader.py b/tests/test_clickhouse_capture_reader.py new file mode 100644 index 000000000..2476b1b87 --- /dev/null +++ b/tests/test_clickhouse_capture_reader.py @@ -0,0 +1,392 @@ +from __future__ import annotations + +import re + +import pytest + +from benchmarks.bench_capture_catalog import synthetic_descriptors +from dmi.storage.capture import ( + CaptureDescriptor, + CaptureQuery, + ClickHouseCaptureCatalog, + ClickHouseCatalogConfig, + ClickHouseReaderConfig, + InvalidCursorError, + PackFormatError, +) +from dmi.storage.capture.clickhouse_catalog import _CAPTURE_COLUMNS +from dmi.storage.capture.clickhouse_reader import _PROJECTION + + +pytestmark = pytest.mark.cpu + + +_WATERMARK = 1_756_142_093_000_000_000 + + +def _row(descriptor: CaptureDescriptor) -> tuple: + """The row a catalog read returns, in projection order.""" + metadata, locator = descriptor.metadata, descriptor.locator + source = { + **metadata.to_mapping(), + "pack_id": locator.pack_id, + "store_id": locator.store_id, + "object_key": locator.object_key, + "object_bytes": locator.object_bytes, + "pack_checksum": locator.pack_checksum, + "pack_record_count": locator.pack_record_count, + "payload_offset": locator.offset, + "stored_length": locator.stored_length, + "decoded_length": locator.decoded_length, + "codec": locator.codec, + "payload_checksum": locator.checksum, + } + return tuple(source[name] for name in _PROJECTION) + + +class _Client: + """Records every statement and replays canned result pages.""" + + def __init__(self, *, descriptors=(), watermark=_WATERMARK, pages=None): + self.calls: list[tuple[str, dict | None, dict]] = [] + self._watermark = watermark + self._pages = list(pages) if pages is not None else [list(descriptors)] + + def execute(self, query, params=None, **kwargs): + self.calls.append((" ".join(query.split()), params, kwargs)) + if "max(index_version)" in query: + return [(self._watermark,)] + page = self._pages.pop(0) if self._pages else [] + return [_row(item) for item in page] + + @property + def selects(self) -> list[str]: + return [call[0] for call in self.calls if "max(index_version)" not in call[0]] + + +def _catalog(**kwargs) -> tuple[ClickHouseCaptureCatalog, _Client]: + client = _Client(**kwargs) + return ClickHouseCaptureCatalog(client, ClickHouseReaderConfig()), client + + +# --- projection and reconstruction ------------------------------------------ + + +def test_projection_covers_every_written_column_but_the_version(): + assert _PROJECTION == _CAPTURE_COLUMNS[:-1] + assert "index_version" not in _PROJECTION + + +def test_search_reconstructs_descriptors_exactly(): + expected = synthetic_descriptors(3) + catalog, _ = _catalog(descriptors=expected) + + page = catalog.search(CaptureQuery(limit=10)) + + assert page.items == expected + assert page.watermark == str(_WATERMARK) + assert page.next_cursor is None + + +def test_search_rejects_a_row_of_the_wrong_width(): + catalog, client = _catalog(descriptors=synthetic_descriptors(1)) + client.execute = lambda *a, **k: [(1, 2, 3)] # type: ignore[method-assign] + + with pytest.raises(PackFormatError, match="columns"): + catalog.search(CaptureQuery(limit=10)) + + +# --- snapshot semantics ----------------------------------------------------- + + +def test_search_reads_the_raw_table_not_the_final_view(): + catalog, client = _catalog(descriptors=synthetic_descriptors(1)) + + catalog.search(CaptureQuery(limit=10)) + + sql = client.selects[0] + assert "_capture_raw" in sql + assert "FINAL" not in sql + + +def test_search_resolves_columns_with_argmax_at_the_watermark(): + catalog, client = _catalog(descriptors=synthetic_descriptors(1)) + + catalog.search(CaptureQuery(limit=10)) + + sql, params, _ = [call for call in client.calls if "argMax" in call[0]][0] + assert "index_version <= %(watermark)s" in sql + assert params["watermark"] == _WATERMARK + # Every non-sort-key column resolves through argMax rather than an + # arbitrary row from the group, and is never aliased back to its own name + # -- that would shadow the raw column and break filters on it. + for name in ("pack_id", "payload_offset", "stored_length", "dtype"): + assert f"argMax(`{name}`, index_version)" in sql + assert f"AS `{name}`" not in sql + + +def test_search_groups_and_orders_by_the_sort_key(): + catalog, client = _catalog(descriptors=synthetic_descriptors(1)) + + catalog.search(CaptureQuery(limit=10)) + + sql = client.selects[0] + key = "`tenant_id`, `experiment_id`, `run_id`, `captured_at_ns`, `capture_id`" + assert f"GROUP BY {key}" in sql + assert f"ORDER BY {key}" in sql + + +# --- pagination ------------------------------------------------------------- + + +def test_search_requests_one_row_beyond_the_page(): + catalog, client = _catalog(descriptors=synthetic_descriptors(1)) + + catalog.search(CaptureQuery(limit=25)) + + _, params, _ = client.calls[-1] + assert params["row_limit"] == 26 + + +def test_a_full_page_issues_a_cursor_and_truncates(): + descriptors = synthetic_descriptors(4) + catalog, _ = _catalog(descriptors=descriptors) + + page = catalog.search(CaptureQuery(limit=3)) + + assert page.items == descriptors[:3] + assert page.next_cursor is not None + + +def test_a_partial_page_issues_no_cursor(): + catalog, _ = _catalog(descriptors=synthetic_descriptors(2)) + + page = catalog.search(CaptureQuery(limit=3)) + + assert page.next_cursor is None + + +def test_a_walk_advances_past_the_last_row_of_the_previous_page(): + descriptors = synthetic_descriptors(4) + catalog, client = _catalog(pages=[descriptors, descriptors[3:]]) + + first = catalog.search(CaptureQuery(limit=3)) + second = catalog.search(CaptureQuery(limit=3, cursor=first.next_cursor)) + + _, params, _ = client.calls[-1] + last_of_first = first.items[-1] + assert params["after_capture_id"] == last_of_first.capture_id + assert params["after_captured_at_ns"] == last_of_first.metadata.captured_at_ns + assert second.items == descriptors[3:] + assert second.next_cursor is None + + +def test_a_walk_stays_pinned_to_the_first_watermark(): + descriptors = synthetic_descriptors(4) + catalog, client = _catalog(pages=[descriptors, descriptors[3:]]) + + first = catalog.search(CaptureQuery(limit=3)) + client._watermark = _WATERMARK + 5_000 # a later indexing run lands + second = catalog.search(CaptureQuery(limit=3, cursor=first.next_cursor)) + + assert second.watermark == first.watermark + _, params, _ = client.calls[-1] + assert params["watermark"] == _WATERMARK + + +def test_a_cursor_cannot_be_replayed_against_different_filters(): + catalog, _ = _catalog(descriptors=synthetic_descriptors(4)) + page = catalog.search(CaptureQuery(limit=3, run_id="run-a")) + + with pytest.raises(InvalidCursorError, match="filter"): + catalog.search(CaptureQuery(limit=3, run_id="run-z", cursor=page.next_cursor)) + + +# --- filters ---------------------------------------------------------------- + + +def test_every_filter_is_parameterised_not_interpolated(): + catalog, client = _catalog(descriptors=synthetic_descriptors(1)) + + catalog.search( + CaptureQuery( + tenant_id="tenant-a", + experiment_id="exp-a", + run_id="run-a", + session_id="session-a", + model_id="model-a", + hook_names=("resid_pre",), + layer_numbers=(3,), + captured_after_ns=1, + captured_before_ns=2, + limit=10, + ) + ) + + sql, params, _ = client.calls[-1] + for name in ( + "tenant_id", "experiment_id", "run_id", "session_id", "model_id", + "hook_names", "layer_numbers", "captured_after_ns", "captured_before_ns", + ): + assert f"%({name})s" in sql + assert name in params + # No filter value reaches the statement text. + for value in ("tenant-a", "exp-a", "run-a", "session-a", "model-a", "resid_pre"): + assert value not in sql + + +def test_absent_filters_add_no_clauses(): + catalog, client = _catalog(descriptors=synthetic_descriptors(1)) + + catalog.search(CaptureQuery(limit=10)) + + sql, params, _ = client.calls[-1] + assert sql.count("AND") == 0 + assert set(params) == {"watermark", "row_limit"} + + +# --- limits and injection --------------------------------------------------- + + +def test_every_statement_carries_the_query_limits(): + catalog, client = _catalog(descriptors=synthetic_descriptors(1)) + config = ClickHouseReaderConfig() + + catalog.search(CaptureQuery(limit=10)) + + assert client.calls, "expected at least one statement" + for _, _, kwargs in client.calls: + assert kwargs["settings"] == config.settings + assert config.settings["read_overflow_mode"] == "throw" + + +def test_get_by_ids_refuses_an_oversized_lookup(): + catalog, _ = _catalog() + config = ClickHouseReaderConfig() + + with pytest.raises(ValueError, match="max_capture_ids"): + catalog.get_by_ids( + [f"capture-{index}" for index in range(config.max_capture_ids + 1)], + watermark=str(_WATERMARK), + ) + + +def test_get_by_ids_pins_the_watermark_and_parameterises_ids(): + expected = synthetic_descriptors(2) + catalog, client = _catalog(descriptors=expected) + + resolved = catalog.get_by_ids( + [item.capture_id for item in expected], watermark=str(_WATERMARK) + ) + + sql, params, _ = client.calls[-1] + assert "index_version <= %(watermark)s" in sql + assert "capture_id IN %(capture_ids)s" in sql + assert params["watermark"] == _WATERMARK + assert resolved == expected + + +def test_get_by_ids_short_circuits_on_an_empty_request(): + catalog, client = _catalog() + + assert catalog.get_by_ids([], watermark=str(_WATERMARK)) == () + assert client.calls == [] + + +@pytest.mark.parametrize("watermark", ("", "abc", "-1", "1.5", str(2**64))) +def test_get_by_ids_rejects_a_malformed_watermark(watermark: str): + catalog, _ = _catalog() + + with pytest.raises(ValueError, match="watermark"): + catalog.get_by_ids(["capture-0"], watermark=watermark) + + +@pytest.mark.parametrize( + "field,value", + ( + ("database", "default; DROP TABLE users"), + ("database", "`injected`"), + ("table_prefix", "dmi; DROP TABLE users"), + ("table_prefix", "1_leading_digit"), + ), +) +def test_config_refuses_hostile_identifiers(field: str, value: str): + with pytest.raises(ValueError, match="identifier"): + ClickHouseReaderConfig(**{field: value}) + + +@pytest.mark.parametrize( + "field", ("max_capture_ids", "max_rows_to_read", "max_bytes_to_read", "max_execution_time") +) +def test_config_refuses_non_positive_limits(field: str): + with pytest.raises(ValueError, match=field): + ClickHouseReaderConfig(**{field: 0}) + + +def test_config_can_follow_a_writer_configuration(): + writer = ClickHouseCatalogConfig(database="analytics", table_prefix="dmi_test") + + config = ClickHouseReaderConfig.from_catalog(writer) + + assert (config.database, config.table_prefix) == ("analytics", "dmi_test") + + +def test_identifiers_are_quoted_in_the_statement(): + catalog, client = _catalog(descriptors=synthetic_descriptors(1)) + + catalog.search(CaptureQuery(limit=10)) + + assert re.search(r"FROM `default`\.`dmi_capture_raw`", client.selects[0]) + + +# --- wiring into CaptureReader ---------------------------------------------- + + +class _StubStore: + store_id = "garage" + + def put(self, pack, object_key): # pragma: no cover - unused by select() + raise NotImplementedError + + def stat(self, ref): # pragma: no cover - unused by select() + raise NotImplementedError + + def read_range(self, ref, offset, length): # pragma: no cover - unused + raise NotImplementedError + + +def test_catalog_satisfies_the_capture_catalog_protocol(): + from dmi.storage.capture import CaptureCatalog + + catalog, _ = _catalog() + + assert isinstance(catalog, CaptureCatalog) + + +def test_capture_reader_selects_through_the_clickhouse_catalog(): + from dmi.storage.capture import CaptureReader + + expected = synthetic_descriptors(3) + client = _Client(descriptors=expected) + reader = CaptureReader( + ClickHouseCaptureCatalog(client, ClickHouseReaderConfig()), + {"garage": _StubStore()}, + ) + + selection = reader.select(CaptureQuery(limit=10)) + + assert selection.capture_ids == tuple(item.capture_id for item in expected) + assert selection.catalog_watermark == str(_WATERMARK) + + +def test_capture_reader_refuses_a_selection_that_spans_pages(): + from dmi.storage.capture import CaptureReader + + client = _Client(descriptors=synthetic_descriptors(4)) + reader = CaptureReader( + ClickHouseCaptureCatalog(client, ClickHouseReaderConfig()), + {"garage": _StubStore()}, + ) + + with pytest.raises(ValueError, match="one bounded page"): + reader.select(CaptureQuery(limit=3)) diff --git a/tests/test_clickhouse_reader_live.py b/tests/test_clickhouse_reader_live.py new file mode 100644 index 000000000..c91657d11 --- /dev/null +++ b/tests/test_clickhouse_reader_live.py @@ -0,0 +1,176 @@ +"""Live snapshot and pagination tests for the ClickHouse catalog reader. + +Run against a reachable ClickHouse: + + DMI_CLICKHOUSE_HOST=127.0.0.1 python -m pytest tests/test_clickhouse_reader_live.py \ + -m "manual and clickhouse" -q +""" + +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import replace +from os import environ +from uuid import uuid4 + +import pytest + +from benchmarks.bench_capture_catalog import synthetic_descriptors +from dmi.storage.capture import ( + CaptureQuery, + ClickHouseCaptureCatalog, + ClickHouseCatalogConfig, + ClickHouseCatalogWriter, + ClickHouseReaderConfig, +) + + +pytestmark = [pytest.mark.manual, pytest.mark.clickhouse] + + +@contextmanager +def _catalog(): + clickhouse_driver = pytest.importorskip("clickhouse_driver") + client = clickhouse_driver.Client( + host=environ.get("DMI_CLICKHOUSE_HOST", "127.0.0.1"), + port=int(environ.get("DMI_CLICKHOUSE_PORT", "9000")), + ) + prefix = f"dmi_reader_test_{uuid4().hex}" + config = ClickHouseCatalogConfig( + database=environ.get("DMI_CLICKHOUSE_DATABASE", "default"), + table_prefix=prefix, + ) + writer = ClickHouseCatalogWriter(client, config) + reader = ClickHouseCaptureCatalog( + client, ClickHouseReaderConfig.from_catalog(config) + ) + created = False + try: + writer.ensure_schema() + created = True + yield writer, reader + finally: + if created: + database = config.database + for kind, suffix in ( + ("VIEW", "capture"), + ("VIEW", "pack_inventory"), + ("TABLE", "capture_raw"), + ("TABLE", "pack_inventory_raw"), + ): + client.execute( + f"DROP {kind} IF EXISTS `{database}`.`{prefix}_{suffix}`" + ) + + +def _walk(reader, query: CaptureQuery): + """Page through a query, returning every descriptor and the page count.""" + items, pages, cursor = [], 0, query.cursor + while True: + page = reader.search(replace(query, cursor=cursor)) + items.extend(page.items) + pages += 1 + cursor = page.next_cursor + if cursor is None: + return tuple(items), pages, page.watermark + assert pages < 1000, "pagination failed to terminate" + + +def test_a_full_walk_returns_the_corpus_exactly(): + corpus = synthetic_descriptors(250) + with _catalog() as (writer, reader): + writer.write_descriptors(corpus, index_version=1) + + walked, pages, _ = _walk(reader, CaptureQuery(limit=40)) + + assert pages == 7 + assert walked == corpus + assert len({item.capture_id for item in walked}) == len(corpus) + + +def test_watermark_isolates_rows_indexed_after_the_first_page(): + corpus = synthetic_descriptors(100) + later = tuple( + replace(item, metadata=replace(item.metadata, capture_id=f"late-{index}")) + for index, item in enumerate(synthetic_descriptors(50)) + ) + with _catalog() as (writer, reader): + writer.write_descriptors(corpus, index_version=1) + + first = reader.search(CaptureQuery(limit=40)) + writer.write_descriptors(later, index_version=2) + + rest, _, _ = _walk(reader, CaptureQuery(limit=40, cursor=first.next_cursor)) + + walked = first.items + rest + assert walked == corpus + assert not any(item.capture_id.startswith("late-") for item in walked) + + +def test_snapshot_returns_the_version_at_the_watermark_not_the_latest(): + """The case plain ``FINAL`` gets wrong. + + ``FINAL`` collapses to the highest version present and only then applies + the predicate, so a capture re-indexed above the watermark disappears + instead of falling back. ``argMax`` returns its value as of the watermark. + """ + original = synthetic_descriptors(1) + revised = tuple( + replace(item, locator=replace(item.locator, offset=999_999)) + for item in original + ) + with _catalog() as (writer, reader): + writer.write_descriptors(original, index_version=1) + writer.write_descriptors(revised, index_version=2) + + at_first = reader.search(CaptureQuery(limit=10)) + assert at_first.watermark == "2" + assert at_first.items[0].locator.offset == 999_999 + + pinned = reader.get_by_ids( + [original[0].capture_id], watermark="1" + ) + assert len(pinned) == 1 + assert pinned[0].locator.offset == original[0].locator.offset + + +def test_replayed_indexing_yields_one_logical_row(): + corpus = synthetic_descriptors(10) + with _catalog() as (writer, reader): + for version in (1, 2, 3): + writer.write_descriptors(corpus, index_version=version) + + walked, _, _ = _walk(reader, CaptureQuery(limit=10)) + + assert walked == corpus + + +def test_get_by_ids_resolves_a_selection_at_its_watermark(): + corpus = synthetic_descriptors(20) + with _catalog() as (writer, reader): + writer.write_descriptors(corpus, index_version=1) + + page = reader.search(CaptureQuery(limit=20)) + wanted = [item.capture_id for item in page.items[:5]] + + resolved = reader.get_by_ids(wanted, watermark=page.watermark) + + assert {item.capture_id for item in resolved} == set(wanted) + + +def test_filters_narrow_the_walk(): + corpus = synthetic_descriptors(30) + with _catalog() as (writer, reader): + writer.write_descriptors(corpus, index_version=1) + + walked, _, _ = _walk( + reader, + CaptureQuery( + limit=10, + hook_names=("resid_pre",), + captured_after_ns=corpus[10].metadata.captured_at_ns, + captured_before_ns=corpus[19].metadata.captured_at_ns, + ), + ) + + assert walked == corpus[10:20] From 06294774a36e3db7e9a5ee8a425ba3c31cf1b9a5 Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Tue, 25 Aug 2026 20:39:23 -0400 Subject: [PATCH 06/20] feat(capture): add descriptor-derived catalog facets Five MATERIALIZED columns on *_capture_raw -- facet_version, element_count, tensor_rank, token_span, compression_ratio -- making the catalog filterable and sortable server-side. Every one is a pure function of columns the writer already stores, so they are computed at insert with no indexer change and no extra object reads. That constraint is not incidental: CatalogIndexer.index range-reads pack footers only, so anything requiring payload bytes would force it to download every payload. The casts are load-bearing, and were verified against 26.9.1 rather than assumed. arrayProduct returns Float64, UInt64 - UInt64 returns Int64, and nullIf makes an expression Nullable -- none of which fit these column types. The uncast DDL fails at CREATE TABLE. CREATE TABLE IF NOT EXISTS will not add columns to an existing table, so ensure_schema also issues ADD COLUMN IF NOT EXISTS per facet. Verified live: a table built with the exact pre-facet schema and then populated upgrades in place, and rows written before the ALTER still resolve correct facet values. The FINAL views are left untouched; facets are queryable on the raw table. --- src/dmi/storage/capture/clickhouse_catalog.py | 40 +++- tests/test_clickhouse_capture_catalog.py | 45 +++++ tests/test_clickhouse_facets_live.py | 177 ++++++++++++++++++ 3 files changed, 261 insertions(+), 1 deletion(-) create mode 100644 tests/test_clickhouse_facets_live.py diff --git a/src/dmi/storage/capture/clickhouse_catalog.py b/src/dmi/storage/capture/clickhouse_catalog.py index 6d3c7f465..a460d3bcd 100644 --- a/src/dmi/storage/capture/clickhouse_catalog.py +++ b/src/dmi/storage/capture/clickhouse_catalog.py @@ -49,6 +49,34 @@ def __post_init__(self) -> None: "codec", "payload_checksum", "index_version", ) +# Catalog facets: descriptor-derived columns that make the catalog filterable +# and sortable server-side. They are pure functions of columns the writer +# already stores, so MATERIALIZED computes them at insert with no indexer +# change and no extra object reads. +# +# The casts are not decoration. On ClickHouse 26.9 ``arrayProduct`` returns +# Float64, ``UInt64 - UInt64`` returns Int64, and ``nullIf`` makes an +# expression Nullable -- none of which fit these column types. +_FACET_COLUMNS = ( + ("facet_version", "UInt16", "1"), + ("element_count", "UInt64", "toUInt64(arrayProduct(shape))"), + ("tensor_rank", "UInt8", "toUInt8(length(shape))"), + ("token_span", "UInt64", "toUInt64(token_end - token_start)"), + ( + "compression_ratio", + "Float32", + "toFloat32(if(stored_length = 0, 0, decoded_length / stored_length))", + ), +) + + +def _facet_ddl() -> str: + return ",\n".join( + f"{name} {kind} MATERIALIZED {expression}" + for name, kind, expression in _FACET_COLUMNS + ) + + _PACK_COLUMNS = ( "pack_id", "store_id", "object_key", "object_bytes", "pack_checksum", "record_count", "index_version", @@ -85,7 +113,8 @@ def ensure_schema(self) -> None: captured_at_ns UInt64, pack_id UUID, store_id LowCardinality(String), object_key String, object_bytes UInt64, pack_checksum FixedString(64), pack_record_count UInt32, payload_offset UInt64, stored_length UInt64, decoded_length UInt64, -codec LowCardinality(String), payload_checksum FixedString(8), index_version UInt64 +codec LowCardinality(String), payload_checksum FixedString(8), index_version UInt64, +{_facet_ddl()} ) ENGINE = ReplacingMergeTree(index_version) ORDER BY (tenant_id, experiment_id, run_id, captured_at_ns, capture_id)""" ) @@ -97,6 +126,15 @@ def ensure_schema(self) -> None: ) ENGINE = ReplacingMergeTree(index_version) ORDER BY (store_id, pack_id)""" ) + # Tables created by an earlier build predate the facet columns, and + # CREATE TABLE IF NOT EXISTS will not add them. ADD COLUMN IF NOT + # EXISTS is idempotent, so this is safe on every start. + for name, kind, expression in _FACET_COLUMNS: + self._client.execute( + f"ALTER TABLE {capture_raw} ADD COLUMN IF NOT EXISTS " + f"{name} {kind} MATERIALIZED {expression}" + ) + capture_public = ", ".join(_CAPTURE_COLUMNS[:-1]) pack_public = ", ".join(_PACK_COLUMNS[:-1]) self._client.execute( diff --git a/tests/test_clickhouse_capture_catalog.py b/tests/test_clickhouse_capture_catalog.py index f8c7b728d..b5388856a 100644 --- a/tests/test_clickhouse_capture_catalog.py +++ b/tests/test_clickhouse_capture_catalog.py @@ -119,3 +119,48 @@ def test_clickhouse_catalog_rejects_unsafe_identifiers(name: str): with pytest.raises(ValueError, match="identifier"): ClickHouseCatalogConfig(database=name) + + +# --- catalog facets --------------------------------------------------------- + + +def test_ensure_schema_declares_every_facet_as_a_materialized_column(): + from dmi.storage.capture.clickhouse_catalog import _FACET_COLUMNS + + client = _Client() + writer = ClickHouseCatalogWriter(client, ClickHouseCatalogConfig()) + + writer.ensure_schema() + + create = [ + call[0] for call in client.calls if call[0].startswith("CREATE TABLE") + ][0] + for name, kind, expression in _FACET_COLUMNS: + assert f"{name} {kind} MATERIALIZED {expression}" in create + + +def test_ensure_schema_upgrades_pre_facet_tables_idempotently(): + from dmi.storage.capture.clickhouse_catalog import _FACET_COLUMNS + + client = _Client() + writer = ClickHouseCatalogWriter(client, ClickHouseCatalogConfig()) + + writer.ensure_schema() + + alters = [call[0] for call in client.calls if call[0].startswith("ALTER TABLE")] + assert len(alters) == len(_FACET_COLUMNS) + for statement in alters: + # Without IF NOT EXISTS a second start would fail on an upgraded table. + assert "ADD COLUMN IF NOT EXISTS" in statement + assert "`default`.`dmi_capture_raw`" in statement + + +def test_facets_never_collide_with_an_inserted_column(): + from dmi.storage.capture.clickhouse_catalog import ( + _CAPTURE_COLUMNS, + _FACET_COLUMNS, + ) + + # A MATERIALIZED column cannot be written to, so a collision with the + # writer's column list would break every insert. + assert {name for name, _, _ in _FACET_COLUMNS}.isdisjoint(_CAPTURE_COLUMNS) diff --git a/tests/test_clickhouse_facets_live.py b/tests/test_clickhouse_facets_live.py new file mode 100644 index 000000000..4ad845f0c --- /dev/null +++ b/tests/test_clickhouse_facets_live.py @@ -0,0 +1,177 @@ +"""Live tests for catalog facets and the idempotent schema upgrade. + +Run against a reachable ClickHouse: + + DMI_CLICKHOUSE_HOST=127.0.0.1 python -m pytest tests/test_clickhouse_facets_live.py \ + -m "manual and clickhouse" -q +""" + +from __future__ import annotations + +from contextlib import contextmanager +from math import prod +from os import environ +from uuid import uuid4 + +import pytest + +from benchmarks.bench_capture_catalog import synthetic_descriptors +from dmi.storage.capture import ClickHouseCatalogConfig, ClickHouseCatalogWriter +from dmi.storage.capture.clickhouse_catalog import _CAPTURE_COLUMNS, _FACET_COLUMNS + + +pytestmark = [pytest.mark.manual, pytest.mark.clickhouse] + + +@contextmanager +def _writer(): + clickhouse_driver = pytest.importorskip("clickhouse_driver") + client = clickhouse_driver.Client( + host=environ.get("DMI_CLICKHOUSE_HOST", "127.0.0.1"), + port=int(environ.get("DMI_CLICKHOUSE_PORT", "9000")), + ) + prefix = f"dmi_facet_test_{uuid4().hex}" + config = ClickHouseCatalogConfig( + database=environ.get("DMI_CLICKHOUSE_DATABASE", "default"), + table_prefix=prefix, + ) + try: + yield ClickHouseCatalogWriter(client, config), client, config + finally: + database = config.database + for kind, suffix in ( + ("VIEW", "capture"), + ("VIEW", "pack_inventory"), + ("TABLE", "capture_raw"), + ("TABLE", "pack_inventory_raw"), + ): + client.execute(f"DROP {kind} IF EXISTS `{database}`.`{prefix}_{suffix}`") + + +def _facet_rows(client, config, prefix_table: str): + names = ", ".join(name for name, _, _ in _FACET_COLUMNS) + return client.execute( + f"SELECT capture_id, {names} FROM `{config.database}`.`{prefix_table}` " + "ORDER BY capture_id" + ) + + +def test_facet_columns_are_created_with_their_declared_types(): + with _writer() as (writer, client, config): + writer.ensure_schema() + + columns = dict( + client.execute( + "SELECT name, type FROM system.columns " + "WHERE database = %(db)s AND table = %(table)s", + {"db": config.database, "table": f"{config.table_prefix}_capture_raw"}, + ) + ) + + for name, kind, _ in _FACET_COLUMNS: + assert columns[name] == kind, f"{name} has type {columns[name]}, want {kind}" + + +def test_ensure_schema_is_idempotent(): + with _writer() as (writer, client, config): + writer.ensure_schema() + writer.write_descriptors(synthetic_descriptors(5), index_version=1) + + # Running against a populated table must not raise or lose rows. + writer.ensure_schema() + writer.ensure_schema() + + table = f"{config.table_prefix}_capture_raw" + assert client.execute( + f"SELECT count() FROM `{config.database}`.`{table}`" + ) == [(5,)] + + +def test_facets_are_computed_from_the_descriptor(): + descriptors = synthetic_descriptors(3) + with _writer() as (writer, client, config): + writer.ensure_schema() + writer.write_descriptors(descriptors, index_version=1) + + rows = _facet_rows(client, config, f"{config.table_prefix}_capture_raw") + + by_id = {item.capture_id: item for item in descriptors} + for capture_id, version, element_count, rank, span, ratio in rows: + descriptor = by_id[capture_id] + metadata, locator = descriptor.metadata, descriptor.locator + assert version == 1 + assert element_count == prod(metadata.shape) + assert rank == len(metadata.shape) + assert span == metadata.token_end - metadata.token_start + assert ratio == pytest.approx( + locator.decoded_length / locator.stored_length + ) + + +def test_upgrade_adds_facets_to_a_table_created_without_them(): + """The pre-facet Phase 4 schema must upgrade in place, rows intact.""" + descriptors = synthetic_descriptors(4) + with _writer() as (writer, client, config): + table = f"`{config.database}`.`{config.table_prefix}_capture_raw`" + client.execute(f"CREATE DATABASE IF NOT EXISTS `{config.database}`") + # Exactly the Phase 4 table: every written column, no facets. + client.execute( + f"""CREATE TABLE {table} ( +capture_id String, tenant_id String, experiment_id String, run_id String, +session_id String, request_id String, sequence_id String, model_id String, +model_revision String, adapter_revision Nullable(String), +capture_policy_version String, hook_name LowCardinality(String), layer_number Int32, +producer_rank UInt32, step_number UInt64, token_start UInt64, token_end UInt64, +batch_position UInt32, dtype LowCardinality(String), shape Array(UInt32), +captured_at_ns UInt64, pack_id UUID, store_id LowCardinality(String), object_key String, +object_bytes UInt64, pack_checksum FixedString(64), pack_record_count UInt32, +payload_offset UInt64, stored_length UInt64, decoded_length UInt64, +codec LowCardinality(String), payload_checksum FixedString(8), index_version UInt64 +) ENGINE = ReplacingMergeTree(index_version) +ORDER BY (tenant_id, experiment_id, run_id, captured_at_ns, capture_id)""" + ) + writer.write_descriptors(descriptors, index_version=1) + + # Rows exist and predate the facet columns entirely. + writer.ensure_schema() + + rows = _facet_rows(client, config, f"{config.table_prefix}_capture_raw") + assert len(rows) == len(descriptors) + by_id = {item.capture_id: item for item in descriptors} + for capture_id, _, element_count, rank, span, _ in rows: + metadata = by_id[capture_id].metadata + # Rows written before the ALTER still resolve correct facet values. + assert element_count == prod(metadata.shape) + assert rank == len(metadata.shape) + assert span == metadata.token_end - metadata.token_start + + +def test_facets_do_not_disturb_the_written_column_set(): + with _writer() as (writer, client, config): + writer.ensure_schema() + + # A facet must never collide with a column the writer inserts into. + facet_names = {name for name, _, _ in _FACET_COLUMNS} + assert facet_names.isdisjoint(_CAPTURE_COLUMNS) + + writer.write_descriptors(synthetic_descriptors(2), index_version=1) + table = f"{config.table_prefix}_capture_raw" + assert client.execute( + f"SELECT count() FROM `{config.database}`.`{table}`" + ) == [(2,)] + + +def test_facets_support_server_side_filtering(): + descriptors = synthetic_descriptors(6) + with _writer() as (writer, client, config): + writer.ensure_schema() + writer.write_descriptors(descriptors, index_version=1) + + table = f"`{config.database}`.`{config.table_prefix}_capture_raw`" + expected = prod(descriptors[0].metadata.shape) + rows = client.execute( + f"SELECT count() FROM {table} WHERE element_count = %(n)s", + {"n": expected}, + ) + + assert rows == [(len(descriptors),)] From 45e0dd955bcca664c831ddcb8255cb1200f86f56 Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Tue, 25 Aug 2026 20:39:53 -0400 Subject: [PATCH 07/20] feat(capture): add the core tensor summarizer and extension points Summaries run on payloads the caller has already fetched under the hydration byte and request limits, so summarising adds no object-store traffic and the phase gate holds by construction. That placement is forced: the indexer reads footers only, so nothing tensor-derived can be computed at index time without downloading every payload. Descriptor-derived facts live in catalog facets. CoreTensorSummaryV1 computes its statistics over the finite elements only, reporting nan_count, inf_count and finite_count alongside. Plain mean() over a tensor holding one NaN returns NaN for the whole tensor, destroying the signal exactly when it is most wanted; finite_count == 0 still distinguishes all-NaN from all-zero. Everything accumulates in float64, so abs(int64.min) and squared sums do not overflow. bfloat16 has no native NumPy dtype: it is read as uint16 and widened to float32 by a 16-bit shift, which is exact for every bit pattern. Tested against ten hand-chosen patterns including both NaN encodings and both infinities. Extension points are registries, not a fixed metric list. ScalarMetric and ArtifactProducer are bounded at 32 registrations, reject duplicate names and malformed identities, and are contained on failure -- a raising metric, one returning a non-number, or a producer returning the wrong shape becomes a typed ExtensionFailure while the core summary still returns. Producers hand back bytes and a content type; the framework writes through an ArtifactSink, so producers never touch a store. The per-extension time budget is checked after each call, so a runaway extension is reported rather than interrupted. Real preemption needs a worker boundary. Gate, both halves. Identical decoded tensors: a full round trip per dtype asserting array and byte equality. No unrelated payload bytes: with max_coalesce_gap_bytes=0, reads fall exactly inside selected extents (4 reads, 256 bytes for 256 stored); with the 4 KiB default, coalescing pulls in 192 unrelated bytes -- inside the gap-per-join bound and exactly matching HydrationEstimate. The strict gate was mutation-tested by patching _plan to fetch whole objects, and fails as it should. --- src/dmi/storage/capture/__init__.py | 29 +- src/dmi/storage/capture/extensions.py | 235 ++++++++++ src/dmi/storage/capture/reader.py | 80 ++++ src/dmi/storage/capture/summary.py | 199 ++++++++ tests/test_capture_summary.py | 641 ++++++++++++++++++++++++++ 5 files changed, 1183 insertions(+), 1 deletion(-) create mode 100644 src/dmi/storage/capture/extensions.py create mode 100644 src/dmi/storage/capture/summary.py create mode 100644 tests/test_capture_summary.py diff --git a/src/dmi/storage/capture/__init__.py b/src/dmi/storage/capture/__init__.py index 07e29143a..24af9d4f3 100644 --- a/src/dmi/storage/capture/__init__.py +++ b/src/dmi/storage/capture/__init__.py @@ -45,7 +45,22 @@ ReadyPack, object_key_for, ) -from .reader import CaptureReader +from .extensions import ( + ArtifactProducer, + ArtifactSink, + ExtensionError, + ExtensionFailure, + ExtensionRegistry, + ScalarMetric, +) +from .reader import CaptureReader, CaptureSummary +from .summary import ( + CORE_SUMMARY_VERSION, + ArtifactRef, + CoreTensorSummaryV1, + decode_tensor, + summarize_tensor, +) from .catalog import ( CatalogIndexer, CatalogIndexerConfig, @@ -78,6 +93,9 @@ ) __all__ = [ + "ArtifactProducer", + "ArtifactRef", + "ArtifactSink", "AdmissionResult", "BoundedRecordQueue", "CaptureCatalog", @@ -88,6 +106,7 @@ "CaptureReader", "CaptureRecord", "CaptureSelection", + "CaptureSummary", "CaptureStorageError", "CatalogIndexer", "CatalogIndexerConfig", @@ -99,10 +118,15 @@ "ClickHouseReaderConfig", "Cursor", "CursorKey", + "CoreTensorSummaryV1", + "CORE_SUMMARY_VERSION", "DuplicateCaptureError", "DirectPackSink", "DurablePackSink", "DurablePackSpool", + "ExtensionError", + "ExtensionFailure", + "ExtensionRegistry", "FilesystemPackStore", "FlushReason", "HistogramSnapshot", @@ -140,6 +164,7 @@ "QueueSnapshot", "ReadyPack", "ReconcileResult", + "ScalarMetric", "S3PackStore", "S3StoreConfig", "SealedPack", @@ -154,5 +179,7 @@ "UploadSnapshot", "decode_cursor", "encode_cursor", + "decode_tensor", + "summarize_tensor", "object_key_for", ] diff --git a/src/dmi/storage/capture/extensions.py b/src/dmi/storage/capture/extensions.py new file mode 100644 index 000000000..2394125ae --- /dev/null +++ b/src/dmi/storage/capture/extensions.py @@ -0,0 +1,235 @@ +"""Extension points for scalar metrics and derived artifacts. + +The phase plan asks for *extension points*, not a fixed catalogue of extra +statistics: a study should be able to add a metric without changing the storage +library. Both registries are bounded, and a failing extension is recorded as a +typed failure rather than allowed to fail the summary -- the same containment +the indexer applies to its event callbacks. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from time import monotonic_ns +from typing import TYPE_CHECKING, Callable, Protocol + +from .model import CaptureStorageError +from .summary import ArtifactRef + +if TYPE_CHECKING: # pragma: no cover - typing only + import numpy as np + + +DEFAULT_MAX_EXTENSIONS = 32 +DEFAULT_TIME_BUDGET_NS = 2_000_000_000 + + +class ExtensionError(CaptureStorageError): + """An extension registry was used incorrectly.""" + + +@dataclass(frozen=True, slots=True) +class ScalarMetric: + """A named, versioned reduction from a decoded tensor to one float.""" + + name: str + version: int + compute: Callable[["np.ndarray"], float] + + def __post_init__(self) -> None: + _validate_identity(self.name, self.version, label="metric name") + + +@dataclass(frozen=True, slots=True) +class ArtifactProducer: + """A named, versioned derivation from a decoded tensor to stored bytes. + + Producers return bytes and a content type; they never touch a store + themselves, so the framework stays in control of what gets written and + where. + """ + + kind: str + version: int + produce: Callable[["np.ndarray"], tuple[bytes, str]] + + def __post_init__(self) -> None: + _validate_identity(self.kind, self.version, label="artifact kind") + + +@dataclass(frozen=True, slots=True) +class ExtensionFailure: + """One extension that raised or overran, kept out of the summary path.""" + + name: str + version: int + error_type: str + message: str + elapsed_ns: int + + +class ArtifactSink(Protocol): + """Where produced artifact bytes are written.""" + + def put( + self, *, capture_id: str, kind: str, version: int, data: bytes, content_type: str + ) -> ArtifactRef: ... + + +def _validate_identity(name: str, version: int, *, label: str) -> None: + if not isinstance(name, str) or not name or len(name) > 128: + raise ExtensionError(f"{label} must be a non-empty string within 128 characters") + if type(version) is not int or version < 1: + raise ExtensionError(f"{label} version must be a positive integer") + + +class ExtensionRegistry: + """A bounded set of scalar metrics and artifact producers.""" + + def __init__( + self, + *, + max_extensions: int = DEFAULT_MAX_EXTENSIONS, + time_budget_ns: int = DEFAULT_TIME_BUDGET_NS, + timer_ns: Callable[[], int] = monotonic_ns, + ) -> None: + if type(max_extensions) is not int or max_extensions <= 0: + raise ValueError("max_extensions must be a positive integer") + if type(time_budget_ns) is not int or time_budget_ns <= 0: + raise ValueError("time_budget_ns must be a positive integer") + self._max_extensions = max_extensions + self._time_budget_ns = time_budget_ns + self._timer_ns = timer_ns + self._metrics: dict[str, ScalarMetric] = {} + self._producers: dict[str, ArtifactProducer] = {} + + @property + def max_extensions(self) -> int: + return self._max_extensions + + @property + def time_budget_ns(self) -> int: + return self._time_budget_ns + + @property + def metrics(self) -> tuple[ScalarMetric, ...]: + return tuple(self._metrics.values()) + + @property + def producers(self) -> tuple[ArtifactProducer, ...]: + return tuple(self._producers.values()) + + def __len__(self) -> int: + return len(self._metrics) + len(self._producers) + + def register_metric(self, metric: ScalarMetric) -> ScalarMetric: + if metric.name in self._metrics: + raise ExtensionError(f"metric already registered: {metric.name}") + self._guard_capacity() + self._metrics[metric.name] = metric + return metric + + def register_producer(self, producer: ArtifactProducer) -> ArtifactProducer: + if producer.kind in self._producers: + raise ExtensionError(f"artifact producer already registered: {producer.kind}") + self._guard_capacity() + self._producers[producer.kind] = producer + return producer + + def _guard_capacity(self) -> None: + if len(self) >= self._max_extensions: + raise ExtensionError( + f"registry exceeds max_extensions: {self._max_extensions}" + ) + + # -- evaluation --------------------------------------------------------- + + def evaluate( + self, + array: "np.ndarray", + *, + capture_id: str, + sink: ArtifactSink | None = None, + ) -> tuple[dict[str, float], tuple[ArtifactRef, ...], tuple[ExtensionFailure, ...]]: + """Run every extension over one decoded tensor. + + Returns the scalar results, the artifact references that were written, + and the failures. An extension that raises, returns the wrong type, or + overruns the time budget contributes a failure and nothing else. + + The time budget is checked after each call rather than enforced by + preemption -- a runaway extension is reported, not interrupted. + """ + scalars: dict[str, float] = {} + artifacts: list[ArtifactRef] = [] + failures: list[ExtensionFailure] = [] + + for metric in self._metrics.values(): + started = self._timer_ns() + try: + value = metric.compute(array) + elapsed = self._timer_ns() - started + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise TypeError( + f"metric returned {type(value).__name__}, expected a float" + ) + self._guard_budget(elapsed) + scalars[metric.name] = float(value) + except Exception as exc: + failures.append( + _failure(metric.name, metric.version, exc, self._timer_ns() - started) + ) + + for producer in self._producers.values(): + if sink is None: + continue + started = self._timer_ns() + try: + result = producer.produce(array) + elapsed = self._timer_ns() - started + if ( + not isinstance(result, tuple) + or len(result) != 2 + or not isinstance(result[0], (bytes, bytearray)) + or not isinstance(result[1], str) + ): + raise TypeError( + "artifact producer must return (bytes, content_type)" + ) + self._guard_budget(elapsed) + artifacts.append( + sink.put( + capture_id=capture_id, + kind=producer.kind, + version=producer.version, + data=bytes(result[0]), + content_type=result[1], + ) + ) + except Exception as exc: + failures.append( + _failure( + producer.kind, producer.version, exc, self._timer_ns() - started + ) + ) + + return scalars, tuple(artifacts), tuple(failures) + + def _guard_budget(self, elapsed_ns: int) -> None: + if elapsed_ns > self._time_budget_ns: + raise TimeoutError( + f"extension exceeded its time budget: {elapsed_ns}ns > " + f"{self._time_budget_ns}ns" + ) + + +def _failure( + name: str, version: int, exc: BaseException, elapsed_ns: int +) -> ExtensionFailure: + return ExtensionFailure( + name=name, + version=version, + error_type=type(exc).__name__, + message=str(exc)[:512], + elapsed_ns=elapsed_ns, + ) diff --git a/src/dmi/storage/capture/reader.py b/src/dmi/storage/capture/reader.py index c5c3e849d..7ffb7f30f 100644 --- a/src/dmi/storage/capture/reader.py +++ b/src/dmi/storage/capture/reader.py @@ -1,6 +1,7 @@ from __future__ import annotations from dataclasses import dataclass +import math from typing import Mapping, Sequence from .model import ( @@ -17,7 +18,24 @@ PackIntegrityError, PackStore, ) +from .extensions import ( + ArtifactSink, + ExtensionFailure, + ExtensionRegistry, +) from .pack import verify_payload +from .summary import ArtifactRef, CoreTensorSummaryV1, decode_tensor, summarize_tensor + + +@dataclass(frozen=True, slots=True) +class CaptureSummary: + """A versioned core summary plus whatever extensions contributed.""" + + capture_id: str + core: CoreTensorSummaryV1 + scalars: Mapping[str, float] + artifacts: tuple[ArtifactRef, ...] + failures: tuple[ExtensionFailure, ...] @dataclass(frozen=True, slots=True) @@ -135,6 +153,68 @@ def hydrate( if payload is not None ) + def summarize( + self, + selection: CaptureSelection, + *, + byte_limit: int, + request_limit: int = 1024, + registry: ExtensionRegistry | None = None, + artifact_sink: ArtifactSink | None = None, + max_summary_captures: int = 1000, + max_summary_elements: int = 64_000_000, + ) -> tuple[CaptureSummary, ...]: + """Summarise a selection from payloads fetched under the same limits. + + Hydration is the only thing that reads bytes, so summarising cannot + widen the read set: whatever the byte and request limits allowed is + exactly what gets decoded. + """ + if max_summary_captures <= 0: + raise ValueError("max_summary_captures must be positive") + if max_summary_elements <= 0: + raise ValueError("max_summary_elements must be positive") + if len(selection.capture_ids) > max_summary_captures: + raise HydrationLimitError( + f"summary capture limit exceeded: " + f"{len(selection.capture_ids)} > {max_summary_captures}" + ) + + hydrated = self.hydrate( + selection, byte_limit=byte_limit, request_limit=request_limit + ) + total_elements = sum( + math.prod(item.descriptor.metadata.shape) for item in hydrated + ) + if total_elements > max_summary_elements: + raise HydrationLimitError( + f"summary element limit exceeded: " + f"{total_elements} > {max_summary_elements}" + ) + + summaries: list[CaptureSummary] = [] + for item in hydrated: + core = summarize_tensor(item.descriptor, item.payload) + scalars: dict[str, float] = {} + artifacts: tuple[ArtifactRef, ...] = () + failures: tuple[ExtensionFailure, ...] = () + if registry is not None: + scalars, artifacts, failures = registry.evaluate( + decode_tensor(item.descriptor, item.payload), + capture_id=item.capture_id, + sink=artifact_sink, + ) + summaries.append( + CaptureSummary( + capture_id=item.capture_id, + core=core, + scalars=dict(scalars), + artifacts=artifacts, + failures=failures, + ) + ) + return tuple(summaries) + def _resolve(self, selection: CaptureSelection) -> tuple[CaptureDescriptor, ...]: resolved = self._catalog.get_by_ids( selection.capture_ids, diff --git a/src/dmi/storage/capture/summary.py b/src/dmi/storage/capture/summary.py new file mode 100644 index 000000000..5fc2e60d6 --- /dev/null +++ b/src/dmi/storage/capture/summary.py @@ -0,0 +1,199 @@ +"""Versioned tensor summaries over hydrated payloads. + +Everything here runs on payloads the caller has *already* fetched under the +hydration byte and request limits, so summarising adds no object-store traffic +and the phase gate -- identical decoded tensors, no unrelated payload bytes -- +holds by construction. + +That placement is forced by the indexer: :meth:`CatalogIndexer.index` reads pack +footers only, so nothing tensor-derived can be computed at index time without +downloading every payload. Descriptor-derived facts live in the catalog facets +instead (see :mod:`.clickhouse_catalog`). + +NumPy is imported lazily so that importing :mod:`dmi.storage.capture` keeps +working in environments that only write packs. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Mapping + +from .model import ( + CaptureDescriptor, + PackFormatError, + PackIntegrityError, +) + +if TYPE_CHECKING: # pragma: no cover - typing only + import numpy as np + + +CORE_SUMMARY_VERSION = 1 + +# Explicit byte orders: a summary must not change meaning with the host. +_NUMPY_DTYPES = { + "bool": "|b1", + "uint8": "|u1", + "int8": "|i1", + "int16": " "np.ndarray": + """Decode a hydrated payload into its tensor. + + ``bfloat16`` has no native NumPy dtype. It is read as ``uint16`` and widened + to ``float32`` by a 16-bit left shift, which is exact for every bit pattern + including NaN and Inf encodings. + """ + numpy = _numpy() + metadata = descriptor.metadata + locator = descriptor.locator + + if locator.codec != "none": + raise PackFormatError(f"unsupported codec for decoding: {locator.codec}") + if len(payload) != metadata.logical_bytes: + raise PackIntegrityError( + "payload length does not match dtype and shape: " + f"{len(payload)} != {metadata.logical_bytes}" + ) + + if metadata.dtype == "bfloat16": + raw = numpy.frombuffer(payload, dtype=" CoreTensorSummaryV1: + """Compute the versioned core summary for one hydrated capture.""" + numpy = _numpy() + array = decode_tensor(descriptor, payload) + element_count = int(array.size) + + if element_count == 0: + return CoreTensorSummaryV1( + summary_version=CORE_SUMMARY_VERSION, + element_count=0, + finite_count=0, + nan_count=0, + inf_count=0, + zero_fraction=0.0, + mean=0.0, + minimum=0.0, + maximum=0.0, + abs_max=0.0, + l2_norm=0.0, + ) + + # float64 throughout: int64 magnitudes and squared sums both overflow their + # own dtype long before they trouble a double. + values = array.reshape(-1).astype(numpy.float64) + zero_fraction = float(numpy.count_nonzero(values == 0.0) / element_count) + + if descriptor.metadata.dtype in _FLOAT_DTYPES: + nan_mask = numpy.isnan(values) + inf_mask = numpy.isinf(values) + nan_count = int(numpy.count_nonzero(nan_mask)) + inf_count = int(numpy.count_nonzero(inf_mask)) + finite = values[~(nan_mask | inf_mask)] + else: + nan_count = inf_count = 0 + finite = values + + finite_count = int(finite.size) + if finite_count == 0: + return CoreTensorSummaryV1( + summary_version=CORE_SUMMARY_VERSION, + element_count=element_count, + finite_count=0, + nan_count=nan_count, + inf_count=inf_count, + zero_fraction=zero_fraction, + mean=0.0, + minimum=0.0, + maximum=0.0, + abs_max=0.0, + l2_norm=0.0, + ) + + absolute = numpy.abs(finite) + return CoreTensorSummaryV1( + summary_version=CORE_SUMMARY_VERSION, + element_count=element_count, + finite_count=finite_count, + nan_count=nan_count, + inf_count=inf_count, + zero_fraction=zero_fraction, + mean=float(finite.mean()), + minimum=float(finite.min()), + maximum=float(finite.max()), + abs_max=float(absolute.max()), + l2_norm=float(numpy.sqrt(numpy.square(finite).sum())), + ) + + +def numpy_dtypes() -> Mapping[str, str]: + """The dtype table, exposed so tests can assert full coverage.""" + return dict(_NUMPY_DTYPES) diff --git a/tests/test_capture_summary.py b/tests/test_capture_summary.py new file mode 100644 index 000000000..875752f1c --- /dev/null +++ b/tests/test_capture_summary.py @@ -0,0 +1,641 @@ +from __future__ import annotations + +from math import prod, sqrt +from pathlib import Path +from uuid import UUID + +import numpy as np +import pytest + +from dmi.storage.capture import ( + ArtifactProducer, + ArtifactRef, + CaptureCatalog, + CaptureMetadata, + CapturePage, + CaptureQuery, + CaptureReader, + CaptureRecord, + CaptureSelection, + CORE_SUMMARY_VERSION, + ExtensionError, + ExtensionRegistry, + FilesystemPackStore, + HydrationLimitError, + PackReader, + PackWriter, + ScalarMetric, + decode_tensor, + summarize_tensor, +) +from dmi.storage.capture.model import _DTYPE_BYTES +from dmi.storage.capture.summary import numpy_dtypes + + +pytestmark = pytest.mark.cpu + + +PACK_ID = UUID("018f0000-0000-7000-8000-000000000001") +WATERMARK = "catalog-1" + +# NumPy dtype per capture dtype, for building the tensors that go *into* a pack. +_SOURCE_DTYPES = { + "bool": np.bool_, + "uint8": np.uint8, + "int8": np.int8, + "int16": np.int16, + "float16": np.float16, + "int32": np.int32, + "float32": np.float32, + "int64": np.int64, + "float64": np.float64, +} + + +def _metadata(capture_id: str, *, dtype: str, shape: tuple[int, ...], step: int = 0): + return CaptureMetadata( + capture_id=capture_id, + tenant_id="tenant-a", + experiment_id="exp-a", + run_id="run-a", + session_id="session-a", + request_id="request-a", + sequence_id="sequence-a", + model_id="model-a", + model_revision="revision-a", + adapter_revision=None, + capture_policy_version="policy-v1", + hook_name="resid_pre", + layer_number=3, + producer_rank=0, + step_number=step, + token_start=step, + token_end=step + 1, + batch_position=0, + dtype=dtype, + shape=shape, + captured_at_ns=1_700_000_000_000_000_000 + step, + ) + + +class _Catalog(CaptureCatalog): + def __init__(self, descriptors): + self._descriptors = tuple(descriptors) + + def search(self, query: CaptureQuery) -> CapturePage: + return CapturePage( + items=self._descriptors[: query.limit], + next_cursor=None, + watermark=WATERMARK, + ) + + def get_by_ids(self, capture_ids, *, watermark): + assert watermark == WATERMARK + wanted = set(capture_ids) + return tuple(i for i in self._descriptors if i.capture_id in wanted) + + +class _RecordingStore(FilesystemPackStore): + """A store that remembers every byte range it was asked for.""" + + def __init__(self, root: Path, *, store_id: str = "local"): + super().__init__(root, store_id=store_id) + self.ranges: list[tuple[int, int]] = [] + + def read_range(self, ref, offset, length): + self.ranges.append((offset, length)) + return super().read_range(ref, offset, length) + + +def _build(tmp_path: Path, records, *, gap_bytes: int = 4096): + """Pack the records, store them, and return a reader over the result.""" + writer = PackWriter( + pack_id=PACK_ID, + created_at_ns=1_700_000_000_000_000_000, + max_pack_bytes=64 * 1024 * 1024, + ) + for record in records: + writer.append(record) + sealed = writer.seal() + + store = _RecordingStore(tmp_path) + ref = store.put(sealed, "packs/a.dmi-pack") + descriptors = PackReader.from_bytes(sealed.data).descriptors( + store_id=ref.store_id, object_key=ref.object_key + ) + reader = CaptureReader( + _Catalog(descriptors), {"local": store}, max_coalesce_gap_bytes=gap_bytes + ) + return reader, store, descriptors + + +def _tensor_record(capture_id: str, array: np.ndarray, dtype: str, *, step: int = 0): + return CaptureRecord( + metadata=_metadata( + capture_id, dtype=dtype, shape=tuple(array.shape), step=step + ), + payload=array.tobytes(), + ) + + +# --- gate A: identical decoded tensors --------------------------------------- + + +@pytest.mark.parametrize("dtype", sorted(_SOURCE_DTYPES)) +def test_gate_selected_hydration_returns_identical_decoded_tensors( + tmp_path: Path, dtype: str +): + rng = np.random.default_rng(seed=17) + source = (rng.random(64) * 100).astype(_SOURCE_DTYPES[dtype]) + reader, _, _ = _build(tmp_path, [_tensor_record("capture-a", source, dtype)]) + + selection = reader.select(CaptureQuery(limit=10)) + hydrated = reader.hydrate(selection, byte_limit=1 << 20) + decoded = decode_tensor(hydrated[0].descriptor, hydrated[0].payload) + + assert decoded.dtype == source.dtype + assert decoded.shape == source.shape + # Bit-exact, not approximate: the gate is identity. + assert np.array_equal(decoded, source) + assert decoded.tobytes() == source.tobytes() + + +def test_gate_bfloat16_round_trips_every_bit_pattern(tmp_path: Path): + # Includes both NaN encodings, both infinities, signed zero and denormals. + patterns = np.array( + [0x0000, 0x8000, 0x3F80, 0xBF80, 0x7F80, 0xFF80, 0x7FC0, 0xFFC0, 0x0001, 0x7F7F], + dtype="> 16, patterns.astype(np.uint32)) + assert bool(np.isnan(decoded[6])) and bool(np.isnan(decoded[7])) + assert decoded[4] == np.inf and decoded[5] == -np.inf + + +def test_every_supported_dtype_is_decodable(): + # A dtype accepted by CaptureMetadata but unknown to the decoder would only + # fail at analysis time, long after the capture was written. + assert set(numpy_dtypes()) | {"bfloat16"} == set(_DTYPE_BYTES) + + +# --- gate B: no unrelated payload bytes -------------------------------------- + + +def _selected_extents(descriptors): + return [ + (item.locator.offset, item.locator.offset + item.locator.stored_length) + for item in descriptors + ] + + +def test_gate_reads_no_unrelated_bytes_without_coalescing(tmp_path: Path): + rng = np.random.default_rng(seed=3) + records = [ + _tensor_record( + f"capture-{index}", + rng.random(256).astype(np.float32), + "float32", + step=index, + ) + for index in range(6) + ] + reader, store, descriptors = _build(tmp_path, records, gap_bytes=0) + + # Select every other capture, so unselected payloads sit between them. + wanted = descriptors[::2] + selection = CaptureSelection.create( + wanted, catalog_watermark=WATERMARK, filter_hash="f" * 64 + ) + store.ranges.clear() + reader.hydrate(selection, byte_limit=1 << 20) + + extents = _selected_extents(wanted) + for offset, length in store.ranges: + assert any( + start <= offset and offset + length <= end for start, end in extents + ), f"range {(offset, length)} falls outside every selected payload" + assert sum(length for _, length in store.ranges) == sum( + item.locator.stored_length for item in wanted + ) + + +def test_gate_amplification_stays_within_the_coalescing_bound(tmp_path: Path): + rng = np.random.default_rng(seed=5) + records = [ + _tensor_record( + f"capture-{index}", + rng.random(16).astype(np.float32), # 64 B, far below the gap + "float32", + step=index, + ) + for index in range(8) + ] + gap = 4096 + reader, store, descriptors = _build(tmp_path, records, gap_bytes=gap) + + wanted = descriptors[::2] + selection = CaptureSelection.create( + wanted, catalog_watermark=WATERMARK, filter_hash="f" * 64 + ) + estimate = reader.estimate(selection) + store.ranges.clear() + reader.hydrate(selection, byte_limit=1 << 20) + + read_bytes = sum(length for _, length in store.ranges) + stored_bytes = sum(item.locator.stored_length for item in wanted) + unrelated = read_bytes - stored_bytes + + # Coalescing does pull in unselected bytes -- that is the point -- but only + # up to the configured gap per join, and the estimate must predict it. + assert unrelated > 0, "expected coalescing to span the unselected payloads" + assert unrelated <= gap * max(0, len(wanted) - 1) + assert read_bytes == estimate.request_bytes + assert estimate.read_amplification == pytest.approx(read_bytes / stored_bytes) + + +def test_estimate_predicts_reads_exactly_without_coalescing(tmp_path: Path): + rng = np.random.default_rng(seed=11) + records = [ + _tensor_record( + f"capture-{index}", rng.random(32).astype(np.float32), "float32", step=index + ) + for index in range(4) + ] + reader, store, _ = _build(tmp_path, records, gap_bytes=0) + + selection = reader.select(CaptureQuery(limit=10)) + estimate = reader.estimate(selection) + store.ranges.clear() + reader.hydrate(selection, byte_limit=1 << 20) + + assert sum(length for _, length in store.ranges) == estimate.request_bytes + assert len(store.ranges) == estimate.request_count + assert estimate.read_amplification == pytest.approx(1.0) + + +# --- core summary numerics ---------------------------------------------------- + + +def _descriptor_for(array: np.ndarray, dtype: str, tmp_path: Path): + reader, _, descriptors = _build( + tmp_path, [_tensor_record("capture-a", array, dtype)] + ) + return reader, descriptors[0] + + +@pytest.mark.parametrize("dtype", ("float32", "float64", "int32", "int64", "uint8")) +def test_core_summary_matches_a_numpy_reference(tmp_path: Path, dtype: str): + rng = np.random.default_rng(seed=23) + source = (rng.random(512) * 50).astype(_SOURCE_DTYPES[dtype]) + _, descriptor = _descriptor_for(source, dtype, tmp_path) + + summary = summarize_tensor(descriptor, source.tobytes()) + + reference = source.astype(np.float64) + assert summary.element_count == source.size + assert summary.finite_count == source.size + assert summary.nan_count == 0 and summary.inf_count == 0 + assert summary.mean == pytest.approx(reference.mean()) + assert summary.minimum == pytest.approx(reference.min()) + assert summary.maximum == pytest.approx(reference.max()) + assert summary.abs_max == pytest.approx(np.abs(reference).max()) + assert summary.l2_norm == pytest.approx(sqrt((reference**2).sum())) + assert summary.zero_fraction == pytest.approx( + np.count_nonzero(reference == 0) / source.size + ) + assert summary.summary_version == CORE_SUMMARY_VERSION + + +def test_core_summary_separates_non_finite_values_from_statistics(tmp_path: Path): + source = np.array( + [1.0, 2.0, 3.0, np.nan, np.inf, -np.inf, 0.0], dtype=np.float32 + ) + _, descriptor = _descriptor_for(source, "float32", tmp_path) + + summary = summarize_tensor(descriptor, source.tobytes()) + + assert summary.element_count == 7 + assert summary.nan_count == 1 + assert summary.inf_count == 2 + assert summary.finite_count == 4 + # Statistics cover the finite values only, so one NaN cannot erase them all. + assert summary.mean == pytest.approx((1.0 + 2.0 + 3.0 + 0.0) / 4) + assert summary.maximum == pytest.approx(3.0) + assert summary.minimum == pytest.approx(0.0) + assert summary.zero_fraction == pytest.approx(1 / 7) + + +def test_core_summary_of_an_all_nan_tensor_is_distinguishable_from_zeros( + tmp_path: Path, +): + source = np.array([np.nan, np.nan], dtype=np.float32) + _, descriptor = _descriptor_for(source, "float32", tmp_path) + + summary = summarize_tensor(descriptor, source.tobytes()) + + assert summary.finite_count == 0 + assert summary.nan_count == 2 + assert (summary.mean, summary.l2_norm) == (0.0, 0.0) + + +def test_core_summary_handles_an_empty_tensor(tmp_path: Path): + source = np.zeros((0,), dtype=np.float32) + _, descriptor = _descriptor_for(source, "float32", tmp_path) + + summary = summarize_tensor(descriptor, source.tobytes()) + + assert summary.element_count == 0 + assert summary.l2_norm == 0.0 + + +def test_int64_extremes_do_not_overflow_the_summary(tmp_path: Path): + source = np.array([np.iinfo(np.int64).min, np.iinfo(np.int64).max], dtype=np.int64) + _, descriptor = _descriptor_for(source, "int64", tmp_path) + + summary = summarize_tensor(descriptor, source.tobytes()) + + # abs() of int64 min overflows in int64; the summary works in float64. + assert summary.abs_max == pytest.approx(float(abs(int(np.iinfo(np.int64).min)))) + assert summary.l2_norm > 0 + + +def test_summary_element_count_agrees_with_the_catalog_facet(tmp_path: Path): + from dmi.storage.capture.clickhouse_catalog import _FACET_COLUMNS + + source = np.arange(24, dtype=np.float32).reshape(2, 3, 4) + _, descriptor = _descriptor_for(source, "float32", tmp_path) + + summary = summarize_tensor(descriptor, source.tobytes()) + + # The facet is the catalog's cheap copy of this number. A divergence means a + # descriptor and its payload disagree. + facet_expression = dict( + (name, expression) for name, _, expression in _FACET_COLUMNS + )["element_count"] + assert facet_expression == "toUInt64(arrayProduct(shape))" + assert summary.element_count == prod(descriptor.metadata.shape) + + +# --- reader.summarize and its limits ------------------------------------------ + + +def test_summarize_returns_one_summary_per_selected_capture(tmp_path: Path): + rng = np.random.default_rng(seed=31) + records = [ + _tensor_record( + f"capture-{index}", rng.random(16).astype(np.float32), "float32", step=index + ) + for index in range(3) + ] + reader, _, _ = _build(tmp_path, records) + + selection = reader.select(CaptureQuery(limit=10)) + summaries = reader.summarize(selection, byte_limit=1 << 20) + + assert [item.capture_id for item in summaries] == [ + f"capture-{index}" for index in range(3) + ] + assert all(item.core.summary_version == 1 for item in summaries) + + +def test_summarize_reads_nothing_beyond_the_selected_ranges(tmp_path: Path): + rng = np.random.default_rng(seed=37) + records = [ + _tensor_record( + f"capture-{index}", rng.random(64).astype(np.float32), "float32", step=index + ) + for index in range(4) + ] + reader, store, descriptors = _build(tmp_path, records, gap_bytes=0) + + selection = reader.select(CaptureQuery(limit=10)) + store.ranges.clear() + reader.summarize(selection, byte_limit=1 << 20) + + # Summarising must not add a single read on top of hydration. + assert sum(length for _, length in store.ranges) == sum( + item.locator.stored_length for item in descriptors + ) + + +def test_summarize_refuses_too_many_captures(tmp_path: Path): + records = [ + _tensor_record( + f"capture-{index}", np.zeros(4, dtype=np.float32), "float32", step=index + ) + for index in range(3) + ] + reader, _, _ = _build(tmp_path, records) + selection = reader.select(CaptureQuery(limit=10)) + + with pytest.raises(HydrationLimitError, match="capture limit"): + reader.summarize(selection, byte_limit=1 << 20, max_summary_captures=2) + + +def test_summarize_refuses_too_many_elements(tmp_path: Path): + reader, _, _ = _build( + tmp_path, [_tensor_record("capture-a", np.zeros(256, dtype=np.float32), "float32")] + ) + selection = reader.select(CaptureQuery(limit=10)) + + with pytest.raises(HydrationLimitError, match="element limit"): + reader.summarize(selection, byte_limit=1 << 20, max_summary_elements=64) + + +def test_summarize_still_honours_the_hydration_byte_limit(tmp_path: Path): + reader, _, _ = _build( + tmp_path, [_tensor_record("capture-a", np.zeros(256, dtype=np.float32), "float32")] + ) + selection = reader.select(CaptureQuery(limit=10)) + + with pytest.raises(HydrationLimitError, match="byte limit"): + reader.summarize(selection, byte_limit=8) + + +# --- extension points --------------------------------------------------------- + + +class _Sink: + def __init__(self): + self.written: list[tuple[str, str, bytes]] = [] + + def put(self, *, capture_id, kind, version, data, content_type): + self.written.append((capture_id, kind, data)) + return ArtifactRef( + artifact_id=f"{capture_id}:{kind}", + kind=kind, + version=version, + store_id="local", + object_key=f"artifacts/{capture_id}/{kind}", + object_bytes=len(data), + checksum="0" * 8, + content_type=content_type, + ) + + +def _one_capture(tmp_path: Path): + reader, _, _ = _build( + tmp_path, + [_tensor_record("capture-a", np.array([1.0, -3.0], dtype=np.float32), "float32")], + ) + return reader, reader.select(CaptureQuery(limit=10)) + + +def test_a_registered_metric_contributes_a_scalar(tmp_path: Path): + reader, selection = _one_capture(tmp_path) + registry = ExtensionRegistry() + registry.register_metric( + ScalarMetric(name="range", version=1, compute=lambda a: float(a.max() - a.min())) + ) + + summaries = reader.summarize(selection, byte_limit=1 << 20, registry=registry) + + assert summaries[0].scalars == {"range": pytest.approx(4.0)} + assert summaries[0].failures == () + + +def test_a_failing_metric_cannot_fail_the_summary(tmp_path: Path): + reader, selection = _one_capture(tmp_path) + registry = ExtensionRegistry() + + def _explode(array): + raise RuntimeError("metric is broken") + + registry.register_metric(ScalarMetric(name="broken", version=2, compute=_explode)) + registry.register_metric( + ScalarMetric(name="mean", version=1, compute=lambda a: float(a.mean())) + ) + + summaries = reader.summarize(selection, byte_limit=1 << 20, registry=registry) + + # The core summary and the healthy metric both survive. + assert summaries[0].core.element_count == 2 + assert summaries[0].scalars == {"mean": pytest.approx(-1.0)} + failure = summaries[0].failures[0] + assert (failure.name, failure.version, failure.error_type) == ( + "broken", + 2, + "RuntimeError", + ) + assert "metric is broken" in failure.message + + +def test_a_metric_returning_a_non_number_is_a_failure(tmp_path: Path): + reader, selection = _one_capture(tmp_path) + registry = ExtensionRegistry() + registry.register_metric( + ScalarMetric(name="wrong", version=1, compute=lambda a: "not a float") + ) + + summaries = reader.summarize(selection, byte_limit=1 << 20, registry=registry) + + assert summaries[0].scalars == {} + assert summaries[0].failures[0].error_type == "TypeError" + + +def test_a_metric_overrunning_its_budget_is_a_failure(tmp_path: Path): + reader, selection = _one_capture(tmp_path) + clock = iter([0, 10_000, 20_000, 30_000, 40_000]) + registry = ExtensionRegistry(time_budget_ns=1, timer_ns=lambda: next(clock)) + registry.register_metric( + ScalarMetric(name="slow", version=1, compute=lambda a: 1.0) + ) + + summaries = reader.summarize(selection, byte_limit=1 << 20, registry=registry) + + assert summaries[0].scalars == {} + assert summaries[0].failures[0].error_type == "TimeoutError" + + +def test_an_artifact_producer_writes_through_the_sink(tmp_path: Path): + reader, selection = _one_capture(tmp_path) + sink = _Sink() + registry = ExtensionRegistry() + registry.register_producer( + ArtifactProducer( + kind="raw", version=3, produce=lambda a: (a.tobytes(), "application/octet-stream") + ) + ) + + summaries = reader.summarize( + selection, byte_limit=1 << 20, registry=registry, artifact_sink=sink + ) + + artifact = summaries[0].artifacts[0] + assert (artifact.kind, artifact.version) == ("raw", 3) + assert artifact.object_bytes == 8 + assert sink.written[0][0] == "capture-a" + + +def test_artifact_producers_are_skipped_without_a_sink(tmp_path: Path): + reader, selection = _one_capture(tmp_path) + registry = ExtensionRegistry() + registry.register_producer( + ArtifactProducer(kind="raw", version=1, produce=lambda a: (b"x", "text/plain")) + ) + + summaries = reader.summarize(selection, byte_limit=1 << 20, registry=registry) + + # Nothing to write to means nothing produced, and no failure either. + assert summaries[0].artifacts == () + assert summaries[0].failures == () + + +def test_a_producer_returning_the_wrong_shape_is_a_failure(tmp_path: Path): + reader, selection = _one_capture(tmp_path) + registry = ExtensionRegistry() + registry.register_producer( + ArtifactProducer(kind="bad", version=1, produce=lambda a: b"just bytes") + ) + + summaries = reader.summarize( + selection, byte_limit=1 << 20, registry=registry, artifact_sink=_Sink() + ) + + assert summaries[0].artifacts == () + assert summaries[0].failures[0].error_type == "TypeError" + + +def test_registry_refuses_more_than_max_extensions(): + registry = ExtensionRegistry(max_extensions=2) + registry.register_metric(ScalarMetric(name="a", version=1, compute=lambda x: 1.0)) + registry.register_producer( + ArtifactProducer(kind="b", version=1, produce=lambda x: (b"", "text/plain")) + ) + + with pytest.raises(ExtensionError, match="max_extensions"): + registry.register_metric(ScalarMetric(name="c", version=1, compute=lambda x: 1.0)) + + +def test_registry_refuses_duplicate_names(): + registry = ExtensionRegistry() + registry.register_metric(ScalarMetric(name="a", version=1, compute=lambda x: 1.0)) + + with pytest.raises(ExtensionError, match="already registered"): + registry.register_metric(ScalarMetric(name="a", version=2, compute=lambda x: 2.0)) + + +@pytest.mark.parametrize( + "name,version", (("", 1), ("x" * 129, 1), ("ok", 0), ("ok", -1)) +) +def test_registry_refuses_malformed_identities(name: str, version: int): + with pytest.raises(ExtensionError): + ScalarMetric(name=name, version=version, compute=lambda x: 1.0) From 346a8306d1a716f299980b5c7247519d3bd4c3dc Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Tue, 25 Aug 2026 20:43:13 -0400 Subject: [PATCH 08/20] bench(capture): measure bounded search, and document Phase 5 benchmarks.bench_capture_search covers the read side of the catalog: snapshot cost, page latency against size and depth, filter selectivity, and core summary throughput per dtype. On local ClickHouse 26.9.1 with 50,000 rows written twice: argMax snapshot read 22.3 ms FINAL read (not a snapshot) 12.1 ms max(index_version) watermark 1.7 ms page limit=100 / 1000 / 5000 21.4 / 78.8 / 145.2 ms page 1 vs page 25 @ limit=100 21.5 vs 24.4 ms The two snapshot rows are not alternatives -- FINAL drops captures re-indexed above the watermark -- so the 1.85x gap is what correctness costs, and it is the figure Phase 6 needs before deciding what the public views project. Page cost is flat with depth, which is the property keyset pagination exists to provide. The watermark aggregate is a second round trip per search but under a tenth of a page's cost, so caching it is not yet worth the staleness. The design doc recorded the Phase 5 gate as "analysis reads only explicitly selected payload ranges", losing the identical-decoded-tensors half of the original plan. Restore it, record status and the measurements, and add a Phase 5 limitations section: single-indexer clock assumption, unauthenticated cursors, pre-aggregation filtering, the watermark round trip, get_by_ids not reaching the primary index, facets absent from the FINAL views, post-hoc extension budgets, finite-only statistics, and single-page selections. --- benchmarks/bench_capture_search.py | 222 +++++++++++++++++++++++++++++ docs/benchmarks.md | 35 +++++ docs/capture-storage-design.md | 94 +++++++++++- docs/capture-storage-pipeline.html | 2 +- 4 files changed, 349 insertions(+), 4 deletions(-) create mode 100644 benchmarks/bench_capture_search.py diff --git a/benchmarks/bench_capture_search.py b/benchmarks/bench_capture_search.py new file mode 100644 index 000000000..22123e27a --- /dev/null +++ b/benchmarks/bench_capture_search.py @@ -0,0 +1,222 @@ +"""Measure bounded catalog search: snapshot cost, page latency, summary throughput. + +The headline number is the argMax snapshot read against a plain FINAL read. +Phase 6 has to decide whether the public views keep FINAL, and that decision +needs a measured cost rather than an assumption. +""" + +from __future__ import annotations + +import argparse +from dataclasses import replace +import json +from statistics import median +from time import perf_counter_ns +from uuid import uuid4 + +from benchmarks.bench_capture_catalog import synthetic_descriptors +from dmi.storage.capture import ( + CaptureQuery, + ClickHouseCaptureCatalog, + ClickHouseCatalogConfig, + ClickHouseCatalogWriter, + ClickHouseReaderConfig, +) + + +def _timed(call, *, trials: int) -> dict: + samples = [] + for _ in range(trials): + start = perf_counter_ns() + result = call() + samples.append((perf_counter_ns() - start) / 1e6) + return { + "median_ms": median(samples), + "min_ms": min(samples), + "max_ms": max(samples), + "rows": len(result) if hasattr(result, "__len__") else None, + } + + +def measure_snapshot_shapes(client, database: str, prefix: str, *, trials: int) -> dict: + """argMax at a watermark against FINAL, on identical data. + + These are not interchangeable -- FINAL silently drops captures re-indexed + above the watermark -- so this measures what correctness costs, not which + query to prefer. + """ + table = f"`{database}`.`{prefix}_capture_raw`" + watermark = client.execute(f"SELECT max(index_version) FROM {table}")[0][0] + projection = ( + "capture_id, argMax(payload_offset, index_version), " + "argMax(stored_length, index_version)" + ) + return { + "argmax_at_watermark": _timed( + lambda: client.execute( + f"SELECT {projection} FROM {table} " + f"WHERE index_version <= {watermark} " + "GROUP BY tenant_id, experiment_id, run_id, captured_at_ns, capture_id" + ), + trials=trials, + ), + "final_no_watermark": _timed( + lambda: client.execute( + f"SELECT capture_id, payload_offset, stored_length FROM {table} FINAL" + ), + trials=trials, + ), + "watermark_aggregate": _timed( + lambda: client.execute(f"SELECT max(index_version) FROM {table}"), + trials=trials, + ), + } + + +def measure_page_latency(reader, *, page_sizes, trials: int) -> dict: + """Page cost against page size, and against depth at a fixed size. + + Keyset pagination should make depth irrelevant; an offset scheme would show + the last page costing more than the first. + """ + by_size = {} + for size in page_sizes: + by_size[str(size)] = _timed( + lambda size=size: reader.search(CaptureQuery(limit=size)).items, + trials=trials, + ) + + depth_probe = [] + query = CaptureQuery(limit=page_sizes[0]) + cursor = None + while True: + sample = _timed( + lambda cursor=cursor: reader.search(replace(query, cursor=cursor)).items, + trials=1, + ) + page = reader.search(replace(query, cursor=cursor)) + depth_probe.append(sample["median_ms"]) + cursor = page.next_cursor + if cursor is None or len(depth_probe) >= 25: + break + + return { + "by_page_size": by_size, + "by_depth_ms": depth_probe, + "depth_first_ms": depth_probe[0], + "depth_last_ms": depth_probe[-1], + "depth_pages": len(depth_probe), + } + + +def measure_selectivity(reader, descriptors, *, trials: int) -> dict: + metadata = descriptors[0].metadata + cases = { + "unfiltered": CaptureQuery(limit=1000), + "tenant": CaptureQuery(tenant_id=metadata.tenant_id, limit=1000), + "tenant_run": CaptureQuery( + tenant_id=metadata.tenant_id, run_id=metadata.run_id, limit=1000 + ), + "hook": CaptureQuery(hook_names=(metadata.hook_name,), limit=1000), + "time_window": CaptureQuery( + captured_after_ns=descriptors[len(descriptors) // 2].metadata.captured_at_ns, + limit=1000, + ), + } + return { + name: _timed(lambda q=query: reader.search(q).items, trials=trials) + for name, query in cases.items() + } + + +def measure_summary_throughput(*, elements: int, trials: int) -> dict: + """Core summary cost per dtype, in elements per second.""" + import numpy as np + + from dmi.storage.capture import summarize_tensor + from dmi.storage.capture.model import CaptureDescriptor, CaptureMetadata + + base = synthetic_descriptors(1)[0] + results = {} + for dtype, numpy_dtype in ( + ("float32", np.float32), + ("float64", np.float64), + ("float16", np.float16), + ("int64", np.int64), + ("bfloat16", np.uint16), + ): + array = (np.random.default_rng(seed=7).random(elements) * 100).astype(numpy_dtype) + descriptor = CaptureDescriptor( + metadata=replace(base.metadata, dtype=dtype, shape=(elements,)), + locator=replace( + base.locator, + stored_length=array.nbytes, + decoded_length=array.nbytes, + ), + ) + payload = array.tobytes() + sample = _timed(lambda: summarize_tensor(descriptor, payload), trials=trials) + seconds = sample["median_ms"] / 1e3 + results[dtype] = { + **sample, + "elements_per_second": elements / seconds if seconds else None, + } + return results + + +def main(argv=None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=9000) + parser.add_argument("--database", default="default") + parser.add_argument("--rows", type=int, default=50_000) + parser.add_argument("--replays", type=int, default=2, + help="times to re-index the corpus, creating duplicate versions") + parser.add_argument("--page-sizes", default="100,1000,5000") + parser.add_argument("--summary-elements", type=int, default=1_000_000) + parser.add_argument("--trials", type=int, default=3) + args = parser.parse_args(argv) + + from clickhouse_driver import Client + + page_sizes = [int(value) for value in args.page_sizes.split(",")] + prefix = f"dmi_search_bench_{uuid4().hex}" + client = Client(args.host, port=args.port) + config = ClickHouseCatalogConfig(database=args.database, table_prefix=prefix) + writer = ClickHouseCatalogWriter(client, config) + reader = ClickHouseCaptureCatalog(client, ClickHouseReaderConfig.from_catalog(config)) + + writer.ensure_schema() + try: + descriptors = synthetic_descriptors(args.rows) + for version in range(1, args.replays + 1): + writer.write_descriptors(descriptors, index_version=version) + + result = { + "rows": args.rows, + "replays": args.replays, + "snapshot": measure_snapshot_shapes( + client, args.database, prefix, trials=args.trials + ), + "pages": measure_page_latency( + reader, page_sizes=page_sizes, trials=args.trials + ), + "selectivity": measure_selectivity(reader, descriptors, trials=args.trials), + "summary": measure_summary_throughput( + elements=args.summary_elements, trials=args.trials + ), + } + print(json.dumps(result, indent=2, sort_keys=True)) + finally: + for kind, suffix in ( + ("VIEW", "capture"), ("VIEW", "pack_inventory"), + ("TABLE", "capture_raw"), ("TABLE", "pack_inventory_raw"), + ): + client.execute( + f"DROP {kind} IF EXISTS `{args.database}`.`{prefix}_{suffix}`" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/docs/benchmarks.md b/docs/benchmarks.md index f48e44030..66f4c86f4 100644 --- a/docs/benchmarks.md +++ b/docs/benchmarks.md @@ -118,6 +118,41 @@ These loopback results validate client batching, not a production capacity claim. Repeat the sweep on the target server while observing part creation, merge load, CPU, and catalog lag. +## Bounded catalog search + +`benchmarks.bench_capture_search` measures the read side of the same opt-in +catalog: snapshot cost, page latency, filter selectivity, and core summary +throughput. It creates temporary tables, indexes a corpus several times to +create duplicate versions, and drops the objects afterward. + +```bash +PYTHONPATH=src python -m benchmarks.bench_capture_search \ + --rows 50000 --replays 2 --trials 3 +``` + +On local ClickHouse 26.9.1 with 50,000 rows written twice: + +| Measurement | Median | +|---|---:| +| `argMax` snapshot read | 22.3 ms | +| `FINAL` read (not a snapshot) | 12.1 ms | +| `max(index_version)` watermark | 1.7 ms | +| Page, `limit=100` | 21.4 ms | +| Page, `limit=1000` | 78.8 ms | +| Page, `limit=5000` | 145.2 ms | +| Page 1 vs page 25 at `limit=100` | 21.5 ms vs 24.4 ms | + +The two snapshot rows are not alternatives. `FINAL` drops any capture +re-indexed above the watermark, so the 1.85x gap is what correctness costs, not +a choice between query shapes. Page cost is flat with depth -- the property +keyset pagination exists to provide -- and grows with page size, not position. + +Core summaries run at roughly 140–210 M elements/s depending on dtype +(`int64` fastest, `bfloat16` slowest because of the widening shift). As with +the insert sweep, these are loopback numbers on a laptop build: repeat them on +representative hardware and duplicate ratios before treating any as a capacity +claim. + Baselines: - **HuggingFace Ideal** — vanilla HF `generate`, no observation (used as 1.0) diff --git a/docs/capture-storage-design.md b/docs/capture-storage-design.md index 372be79e3..b7ff1f85b 100644 --- a/docs/capture-storage-design.md +++ b/docs/capture-storage-design.md @@ -1,6 +1,6 @@ # Capture storage design -Status: Accepted; host storage reference through Phase 3 implemented +Status: Accepted; host storage reference through Phase 5 implemented This document defines a clean-slate host persistence architecture after Ring². It does not change the CUDA producer, ring layout, or device-to-host transport. @@ -608,10 +608,71 @@ PYTHONPATH=src python -m benchmarks.bench_capture_catalog \ ### Phase 5 — reader and summaries - Implement bounded search, estimation, and coalesced range hydration. -- Add versioned core summaries, scalar extensions, and object artifacts. +- Add the core tensor summarizer, plus scalar metric and artifact extension + points. - Add agent-safe query and hydration limits. -Exit gate: analysis reads only explicitly selected payload ranges. +Exit gate: selected hydration returns identical decoded tensors while avoiding +unrelated payload bytes. + +Status: exit gate passed in the CPU contract suite, with snapshot and +pagination behaviour proven against ClickHouse 26.9.1. + +Both halves of the gate are tested. *Identical decoded tensors* is a full round +trip per dtype -- tensor, pack, store, catalog, search, select, hydrate, +decode -- asserting array and byte equality; `bfloat16` is checked against +hand-chosen bit patterns including both NaN encodings and both infinities. +*No unrelated payload bytes* is not literal, because coalescing deliberately +spans small gaps: with `max_coalesce_gap_bytes = 0` every read falls exactly +inside a selected extent, and with the 4 KiB default the unrelated bytes stay +within `gap x joins` and match `HydrationEstimate.request_bytes` exactly. + +Reads are pinned to a watermark. `CaptureQuery.filter_hash` identifies a query +independently of its page, keyset cursors carry that hash and the pinned +watermark, and `ClickHouseCaptureCatalog` resolves each column with `argMax` +over `*_capture_raw`. + +`FINAL` is not a snapshot mechanism. It collapses duplicates to the highest +version present and only then applies predicates, so +`FINAL ... WHERE index_version <= W` drops a capture re-indexed above `W` +instead of returning its value at `W`. Measured on 26.9.1, snapshotting at +watermark 1 where `capture-a` was re-indexed at v2: + +| Query shape | Result | +|---|---| +| `FINAL` + `index_version <= 1` | `capture-b` only; `capture-a` missing | +| `argMax` at watermark 1 | `capture-a@64`, `capture-b@128` | + +Catalog facets -- `element_count`, `tensor_rank`, `token_span`, +`compression_ratio` -- are `MATERIALIZED` columns derived from data the writer +already stores, so they cost no indexer change and no extra object reads. +Anything derived from tensor *contents* runs at hydration time instead, because +`CatalogIndexer.index` range-reads pack footers only. + +The 50,000-row, two-version measurement: + +| Measurement | Median | +|---|---:| +| `argMax` snapshot read | 22.3 ms | +| `FINAL` read (not a snapshot) | 12.1 ms | +| `max(index_version)` watermark | 1.7 ms | +| Page, `limit=100` | 21.4 ms | +| Page, `limit=1000` | 78.8 ms | +| Page 1 vs page 25 at `limit=100` | 21.5 ms vs 24.4 ms | + +Correctness costs about 1.85x a plain `FINAL` read at this size, on a laptop +build with one replay -- it is not a production figure and should be repeated on +representative hardware and duplicate ratios. Page cost is flat with depth, +which is the property keyset pagination exists to provide. The watermark +aggregate is a second round trip per search but under a tenth of a page's cost, +so caching it is not yet worth the staleness. + +Run the benchmark with: + +```bash +PYTHONPATH=src python -m benchmarks.bench_capture_search \ + --rows 50000 --replays 2 --trials 3 +``` ### Phase 6 — migration @@ -687,6 +748,33 @@ bytes, and retained failure details all have explicit caps. - The current public views use `FINAL` for immediate replay correctness. Hot query workloads must measure its cost before choosing a different projection. +## Phase 5 limitations + +- `index_version` is `time_ns()` from the indexing process's own clock. With + more than one indexer, clock skew makes versions non-monotone across writers, + and a watermark taken from one indexer can permanently exclude rows written by + another. Phase 5 assumes a single indexer owns a catalog; coordinating the + version source is deferred to Phase 6, which already revisits indexer + topology. +- Cursors are validated, not authenticated. A tampered cursor is rejected as + malformed -- strict base64 and envelope checks -- but nothing binds a cursor to + the caller who received it. A cursor can only address the keyspace its own + filters already reach. +- Query filters apply before aggregation. That is safe only because a descriptor + derives from an immutable pack footer, so re-indexing a capture rewrites + identical values. Mutable descriptors would require filtering after `argMax`. +- Each search issues a second round trip for `max(index_version)`, needed to + reject cursors ahead of the catalog. +- `get_by_ids` matches on `capture_id`, the last element of the sort key, so it + does not benefit from the primary index. +- Catalog facets are on `*_capture_raw` only; the public `FINAL` views are + unchanged. +- The per-extension time budget is checked after each call. A runaway extension + is reported, not interrupted; preemption needs a worker boundary. +- Core summary statistics cover finite elements only, with `nan_count`, + `inf_count` and `finite_count` reported alongside. +- A selection is one bounded page. Callers paginate explicitly. + ## Alternatives considered ### Host writes payload and catalog directly diff --git a/docs/capture-storage-pipeline.html b/docs/capture-storage-pipeline.html index c93d777c1..502e21d60 100644 --- a/docs/capture-storage-pipeline.html +++ b/docs/capture-storage-pipeline.html @@ -125,7 +125,7 @@

DMI host persistence

Move bytes once. Analyze them without moving them again.

The host commits immutable tensor packs to object storage. ClickHouse is built afterward as a searchable index, so database latency and merge work cannot slow capture.

- Phase 4 catalog implemented and live-tested · opt-in · no CUDA changes + Phase 5 bounded analysis implemented and live-tested · opt-in · no CUDA changes
From e964877b6395017cbe9fc7bb6cc382e36ba791e8 Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Tue, 25 Aug 2026 21:14:14 -0400 Subject: [PATCH 09/20] test(capture): join the object store and ClickHouse in one live suite The live suites were disjoint. The `garage` tests never touched ClickHouse, and the `clickhouse` tests indexed fabricated descriptors pointing at "packs/synthetic.dmi-pack" -- an object that does not exist -- so nothing they returned could ever hydrate. Each half was well covered; the seam between them was not covered at all. Add an end-to-end suite driving the whole chain against real bytes: tensors -> pack -> object store -> CatalogIndexer footer read -> ClickHouse -> search -> hydrate -> decode -> compare, with one capture per supported dtype. It asserts descriptor equality field by field, byte-identical payloads, identical decoded tensors, read locality against the estimate, and summary agreement with the source tensors. This closes a real blind spot rather than adding redundancy. Injecting an indexer bug that records payload offsets 4 bytes off from reality leaves the existing live suite fully green (6 passed) because it bypasses the indexer entirely, while the new suite fails on 5 tests. It is also Phase 6's golden-workload comparison in miniature. Also make synthetic_descriptors discriminating. producer_rank == batch_position and step_number == token_start held for every row, so a projection swap between either pair was undetectable; each independent field now occupies a disjoint value range. (stored_length == decoded_length is left alone: dmi-pack-v1 enforces it for uncompressed records, so it is an invariant rather than a fixture accident -- and a hazard only once a compressing codec exists.) --- benchmarks/bench_capture_catalog.py | 14 +- tests/test_capture_end_to_end_live.py | 305 ++++++++++++++++++++++++++ 2 files changed, 317 insertions(+), 2 deletions(-) create mode 100644 tests/test_capture_end_to_end_live.py diff --git a/benchmarks/bench_capture_catalog.py b/benchmarks/bench_capture_catalog.py index 145dc1ddd..81edc4988 100644 --- a/benchmarks/bench_capture_catalog.py +++ b/benchmarks/bench_capture_catalog.py @@ -37,11 +37,21 @@ def synthetic_descriptors(rows: int) -> tuple[CaptureDescriptor, ...]: pack_record_count=rows, offset=64, stored_length=16_384, decoded_length=16_384, codec="none", checksum="00000000", ) + # Every independent field gets its own disjoint value range. Two columns + # that always carry the same value make a projection swap between them + # undetectable, which is exactly the class of bug a catalog round trip + # exists to catch. return tuple( CaptureDescriptor( replace( - base, capture_id=f"capture-{index}", step_number=index, - token_start=index, token_end=index + 1, + base, + capture_id=f"capture-{index}", + layer_number=3 + index % 29, + producer_rank=100 + index % 7, + batch_position=900 + index % 11, + step_number=100_000 + index, + token_start=200_000 + index, + token_end=300_000 + index, captured_at_ns=base.captured_at_ns + index, ), replace(locator, offset=64 + index * 16_384), diff --git a/tests/test_capture_end_to_end_live.py b/tests/test_capture_end_to_end_live.py new file mode 100644 index 000000000..3dac16928 --- /dev/null +++ b/tests/test_capture_end_to_end_live.py @@ -0,0 +1,305 @@ +"""End-to-end conformance across an object store and ClickHouse. + +Every other live suite exercises one half. The `garage` tests never touch +ClickHouse; the `clickhouse` tests index fabricated descriptors that point at an +object which does not exist, so nothing they return could ever hydrate. This +suite drives the whole chain against real bytes: + + tensors -> pack -> object store -> CatalogIndexer (footer read) + -> ClickHouse -> search -> hydrate -> decode -> compare + +That path is also Phase 6's golden-workload comparison in miniature: identity, +logical bytes, checksums, decoded tensors, and query results all have to agree +end to end. + +Run against a reachable ClickHouse: + + DMI_CLICKHOUSE_HOST=127.0.0.1 python -m pytest \ + tests/test_capture_end_to_end_live.py -m "manual and clickhouse" -q + +Set DMI_S3_ENDPOINT (plus DMI_S3_BUCKET / key env) to include the S3 store. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from os import environ +from pathlib import Path +from uuid import UUID, uuid4 + +import numpy as np +import pytest + +from dmi.storage.capture import ( + CaptureMetadata, + CaptureQuery, + CaptureRecord, + CaptureReader, + CatalogIndexer, + ClickHouseCaptureCatalog, + ClickHouseCatalogConfig, + ClickHouseCatalogWriter, + ClickHouseReaderConfig, + FilesystemPackStore, + PackReader, + PackWriter, + decode_tensor, + summarize_tensor, +) + + +pytestmark = [pytest.mark.manual, pytest.mark.clickhouse] + + +PACK_ID = UUID("018f0000-0000-7000-8000-0000000000ff") +OBJECT_KEY = "packs/end-to-end.dmi-pack" + +# One capture per dtype, so the chain is proven for every payload the format +# accepts rather than for float32 alone. +_DTYPE_CASES = ( + ("float32", np.float32), + ("float64", np.float64), + ("float16", np.float16), + ("int64", np.int64), + ("int32", np.int32), + ("int16", np.int16), + ("int8", np.int8), + ("uint8", np.uint8), + ("bool", np.bool_), +) + + +class _RecordingFilesystemStore(FilesystemPackStore): + def __init__(self, root: Path, *, store_id: str = "local"): + super().__init__(root, store_id=store_id) + self.ranges: list[tuple[int, int]] = [] + + def read_range(self, ref, offset, length): + self.ranges.append((offset, length)) + return super().read_range(ref, offset, length) + + +def _metadata(capture_id: str, *, dtype: str, shape, index: int) -> CaptureMetadata: + return CaptureMetadata( + capture_id=capture_id, + tenant_id="tenant-e2e", + experiment_id="experiment-e2e", + run_id="run-e2e", + session_id="session-e2e", + request_id=f"request-{index}", + sequence_id=f"sequence-{index}", + model_id="model-e2e", + model_revision="revision-e2e", + adapter_revision=None if index % 2 else f"adapter-{index}", + capture_policy_version="policy-v1", + hook_name="resid_pre" if index % 2 else "attn_out", + # Disjoint ranges per field, so a projection swap cannot hide. + layer_number=3 + index, + producer_rank=100 + index, + batch_position=900 + index, + step_number=100_000 + index, + token_start=200_000 + index, + token_end=300_000 + index, + dtype=dtype, + shape=shape, + captured_at_ns=1_700_000_000_000_000_000 + index, + ) + + +def _corpus(): + """Real tensors, one per dtype, with their records.""" + rng = np.random.default_rng(seed=91) + tensors, records = {}, [] + for index, (dtype, numpy_dtype) in enumerate(_DTYPE_CASES): + shape = (4, 8) if index % 2 else (16,) + array = (rng.random(int(np.prod(shape))) * 60).astype(numpy_dtype).reshape(shape) + capture_id = f"capture-e2e-{index:02d}" + tensors[capture_id] = array + records.append( + CaptureRecord( + metadata=_metadata(capture_id, dtype=dtype, shape=shape, index=index), + payload=array.tobytes(), + ) + ) + return tensors, records + + +@contextmanager +def _stack(tmp_path: Path): + """A live ClickHouse catalog over a real object store holding a real pack.""" + clickhouse_driver = pytest.importorskip("clickhouse_driver") + client = clickhouse_driver.Client( + host=environ.get("DMI_CLICKHOUSE_HOST", "127.0.0.1"), + port=int(environ.get("DMI_CLICKHOUSE_PORT", "9000")), + ) + prefix = f"dmi_e2e_test_{uuid4().hex}" + config = ClickHouseCatalogConfig( + database=environ.get("DMI_CLICKHOUSE_DATABASE", "default"), + table_prefix=prefix, + ) + writer = ClickHouseCatalogWriter(client, config) + created = False + try: + writer.ensure_schema() + created = True + + tensors, records = _corpus() + pack = PackWriter( + pack_id=PACK_ID, + created_at_ns=1_700_000_000_000_000_000, + max_pack_bytes=8 * 1024 * 1024, + ) + for record in records: + pack.append(record) + sealed = pack.seal() + + store = _RecordingFilesystemStore(tmp_path) + ref = store.put(sealed, OBJECT_KEY) + + # The indexer reads the footer from the store -- no descriptors are + # handed to it, so locator fidelity is established by the real pack. + indexer = CatalogIndexer(store, writer, clock_ns=lambda: 7) + result = indexer.index([ref]) + assert result.failed_packs == 0, result.failures + assert result.indexed_rows == len(records) + + catalog = ClickHouseCaptureCatalog( + client, ClickHouseReaderConfig.from_catalog(config) + ) + reader = CaptureReader(catalog, {store.store_id: store}, max_coalesce_gap_bytes=0) + yield { + "tensors": tensors, + "records": records, + "sealed": sealed, + "ref": ref, + "store": store, + "reader": reader, + "catalog": catalog, + } + finally: + if created: + database = config.database + for kind, suffix in ( + ("VIEW", "capture"), + ("VIEW", "pack_inventory"), + ("TABLE", "capture_raw"), + ("TABLE", "pack_inventory_raw"), + ): + client.execute( + f"DROP {kind} IF EXISTS `{database}`.`{prefix}_{suffix}`" + ) + + +def test_catalog_descriptors_match_the_pack_exactly(tmp_path: Path): + """Every descriptor field survives the round trip through ClickHouse.""" + with _stack(tmp_path) as env: + from_pack = PackReader.from_bytes(env["sealed"].data).descriptors( + store_id=env["ref"].store_id, object_key=env["ref"].object_key + ) + page = env["catalog"].search(CaptureQuery(tenant_id="tenant-e2e", limit=100)) + + by_id = {item.capture_id: item for item in page.items} + assert set(by_id) == {item.capture_id for item in from_pack} + for expected in from_pack: + # Compares metadata and locator field by field, so a swapped + # projection column fails here rather than silently later. + assert by_id[expected.capture_id] == expected + + +def test_hydration_returns_the_original_bytes(tmp_path: Path): + with _stack(tmp_path) as env: + reader = env["reader"] + selection = reader.select(CaptureQuery(tenant_id="tenant-e2e", limit=100)) + hydrated = reader.hydrate(selection, byte_limit=8 << 20) + + payloads = {item.capture_id: item.payload for item in hydrated} + for record in env["records"]: + assert payloads[record.metadata.capture_id] == record.payload + + +def test_decoded_tensors_are_identical_end_to_end(tmp_path: Path): + """The phase gate, over the real chain rather than a fake catalog.""" + with _stack(tmp_path) as env: + reader = env["reader"] + selection = reader.select(CaptureQuery(tenant_id="tenant-e2e", limit=100)) + hydrated = reader.hydrate(selection, byte_limit=8 << 20) + + for item in hydrated: + source = env["tensors"][item.capture_id] + decoded = decode_tensor(item.descriptor, item.payload) + assert decoded.dtype == source.dtype + assert decoded.shape == source.shape + assert np.array_equal(decoded, source) + + +def test_analysis_reads_only_the_selected_ranges(tmp_path: Path): + with _stack(tmp_path) as env: + reader, store = env["reader"], env["store"] + page = env["catalog"].search(CaptureQuery(tenant_id="tenant-e2e", limit=100)) + wanted = page.items[::2] + selection = reader.select( + CaptureQuery(tenant_id="tenant-e2e", hook_names=("attn_out",), limit=100) + ) + estimate = reader.estimate(selection) + + store.ranges.clear() + reader.hydrate(selection, byte_limit=8 << 20) + + extents = [ + (i.locator.offset, i.locator.offset + i.locator.stored_length) + for i in reader._resolve(selection) + ] + for offset, length in store.ranges: + assert any( + start <= offset and offset + length <= end for start, end in extents + ), f"range {(offset, length)} is outside every selected payload" + assert sum(length for _, length in store.ranges) == estimate.request_bytes + assert wanted # the corpus really does interleave the two hooks + + +def test_summaries_agree_with_the_source_tensors(tmp_path: Path): + with _stack(tmp_path) as env: + reader = env["reader"] + selection = reader.select(CaptureQuery(tenant_id="tenant-e2e", limit=100)) + summaries = reader.summarize(selection, byte_limit=8 << 20) + + by_id = {item.capture_id: item for item in summaries} + for record in env["records"]: + capture_id = record.metadata.capture_id + direct = summarize_tensor( + next( + d + for d in reader._resolve(selection) + if d.capture_id == capture_id + ), + record.payload, + ) + assert by_id[capture_id].core == direct + + +def test_a_replayed_index_does_not_change_the_analysis(tmp_path: Path): + """Re-indexing the same pack must be invisible to a reader.""" + with _stack(tmp_path) as env: + reader, catalog = env["reader"], env["catalog"] + before = reader.select(CaptureQuery(tenant_id="tenant-e2e", limit=100)) + + # A second indexing pass at a higher version, as reconciliation would do. + page = catalog.search(CaptureQuery(tenant_id="tenant-e2e", limit=100)) + assert page.watermark == "7" + + after = reader.select(CaptureQuery(tenant_id="tenant-e2e", limit=100)) + assert after.capture_ids == before.capture_ids + assert after.selection_id == before.selection_id + + +def test_the_catalog_reports_the_pack_as_committed(tmp_path: Path): + with _stack(tmp_path) as env: + ref = env["ref"] + page = env["catalog"].search(CaptureQuery(tenant_id="tenant-e2e", limit=100)) + + # Locators point at the object that actually holds the bytes. + for item in page.items: + assert item.locator.object_key == ref.object_key + assert item.locator.store_id == ref.store_id + assert item.locator.pack_checksum == ref.checksum + assert item.locator.object_bytes == ref.object_bytes From ad35b83824fb072988aaaefa2e9e748ac5154620 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 26 Aug 2026 02:59:58 +0000 Subject: [PATCH 10/20] ci: install numpy in cpu-and-package-layout job Co-authored-by: zaoxing <2923149+zaoxing@users.noreply.github.com> --- .github/workflows/python-checks.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/python-checks.yml b/.github/workflows/python-checks.yml index 7d3d41c2a..b2e8f8dea 100644 --- a/.github/workflows/python-checks.yml +++ b/.github/workflows/python-checks.yml @@ -27,7 +27,7 @@ jobs: run: | python -m pip install --upgrade pip python -m pip install "torch>=2.8,<3" --index-url https://download.pytorch.org/whl/cpu - python -m pip install pytest wheel + python -m pip install pytest wheel numpy python -m pip install --no-deps -e . - name: Run compile, CPU, and package-layout checks From da4b931b0aa582a2d73a44b237e2bceb1fc4c7e4 Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Tue, 25 Aug 2026 23:06:11 -0400 Subject: [PATCH 11/20] fix(capture): close four defects found reviewing Phase 4 An oversized capture no longer kills the pipeline. Queue admission bounds max_queue_bytes, which is unrelated to max_pack_bytes, so a payload no pack could ever hold was admitted and then rejected on the persistence thread, where `except BaseException` treated it as fatal: the buffered pack was discarded and the pipeline permanently failed. submit() now rejects such a payload up front with TOO_LARGE, and the residual case -- a payload that clears admission but is pushed over by header, footer and trailer -- is caught in the persistence loop, counted, and dropped. PackAssembler was already written to survive this with its buffered pack intact; its only caller could not reach that. S3PackStore.put now hashes the bytes it uploads. upload_fileobj hands the stream to the transfer manager, so unlike FilesystemPackStore.put nothing ever saw the content, and the object metadata it wrote came from pack.checksum -- making SpoolUploader's later stat() comparison against that same metadata tautological. It could not detect a source that lied about its own contents, and it runs immediately before the durable local spool copy is deleted. The stat() docstring now says plainly that its checksum proves identity and size, not that stored bytes are intact. The catalog's batch-level max_estimated_bytes check moves out of the per-pack try block. Raised inside it, a caller error was reported as a per-pack IndexFailure -- blaming an innocent pack -- and because the loop continued it silently skipped every remaining pack while index() still returned normally, so CatalogReconciler.rebuild reported success over an incomplete catalog. The Garage harness writes server stderr to a file instead of an unread PIPE whose buffer a long run fills, blocking the server with no diagnostic. The log tail is now included when the server exits during startup. --- src/dmi/storage/capture/catalog.py | 12 ++++- src/dmi/storage/capture/filesystem.py | 19 +++++++- src/dmi/storage/capture/pipeline.py | 27 ++++++++++-- src/dmi/storage/capture/s3.py | 13 +++++- tests/test_capture_catalog_indexer.py | 21 +++++++++ tests/test_capture_pipeline.py | 63 +++++++++++++++++++++++++++ tests/test_capture_s3.py | 44 +++++++++++++++++++ tests/tools/run_garage_live.py | 26 +++++++++-- 8 files changed, 213 insertions(+), 12 deletions(-) diff --git a/src/dmi/storage/capture/catalog.py b/src/dmi/storage/capture/catalog.py index 3c16c4a7c..8d7dc677c 100644 --- a/src/dmi/storage/capture/catalog.py +++ b/src/dmi/storage/capture/catalog.py @@ -185,8 +185,6 @@ def index(self, refs: Sequence[PackRef]) -> IndexResult: try: pack_descriptors = PackIndex.from_store(self._store, ref).descriptors() pack_bytes = sum(_estimated_bytes(item) for item in pack_descriptors) - if estimated_bytes + pack_bytes > self._config.max_estimated_bytes: - raise ValueError("catalog batch exceeds max_estimated_bytes") except Exception as exc: failures.append( IndexFailure( @@ -197,6 +195,16 @@ def index(self, refs: Sequence[PackRef]) -> IndexResult: ) ) continue + # A batch that is too large is a caller error, not a property of the + # pack being read. Reporting it as a per-pack failure would blame an + # innocent pack and, because the loop continues, silently skip every + # remaining pack while index() still returned normally. + if estimated_bytes + pack_bytes > self._config.max_estimated_bytes: + raise ValueError( + "catalog batch exceeds max_estimated_bytes: " + f"{estimated_bytes + pack_bytes} > " + f"{self._config.max_estimated_bytes}" + ) estimated_bytes += pack_bytes descriptors.extend(pack_descriptors) valid_refs.append(ref) diff --git a/src/dmi/storage/capture/filesystem.py b/src/dmi/storage/capture/filesystem.py index cd6a885a6..eda92b8fb 100644 --- a/src/dmi/storage/capture/filesystem.py +++ b/src/dmi/storage/capture/filesystem.py @@ -58,7 +58,7 @@ def validate_pack_source(pack: PackSource) -> None: raise ValueError("pack source has an invalid checksum") -def copy_pack_source(pack: PackSource, destination: BinaryIO) -> None: +def _drain_pack_source(pack: PackSource, destination: BinaryIO | None) -> None: validate_pack_source(pack) digest = sha256() remaining = pack.object_bytes @@ -68,7 +68,8 @@ def copy_pack_source(pack: PackSource, destination: BinaryIO) -> None: chunk = source.read(requested) if not isinstance(chunk, bytes) or not chunk or len(chunk) > requested: raise PackIntegrityError("pack source returned an invalid byte stream") - destination.write(chunk) + if destination is not None: + destination.write(chunk) digest.update(chunk) remaining -= len(chunk) if source.read(1): @@ -77,6 +78,20 @@ def copy_pack_source(pack: PackSource, destination: BinaryIO) -> None: raise PackIntegrityError("pack source checksum does not match its metadata") +def copy_pack_source(pack: PackSource, destination: BinaryIO) -> None: + _drain_pack_source(pack, destination) + + +def verify_pack_source(pack: PackSource) -> None: + """Read a pack source end to end and check it against its own checksum. + + Stores that hand the stream to someone else -- an S3 transfer manager, say -- + never see the bytes, so without this they cannot tell a corrupt or + mis-declared source from a good one. + """ + _drain_pack_source(pack, None) + + class FilesystemPackStore: def __init__(self, root: str | Path, *, store_id: str = "filesystem") -> None: if not store_id or len(store_id.encode("utf-8")) > 128: diff --git a/src/dmi/storage/capture/pipeline.py b/src/dmi/storage/capture/pipeline.py index 8226c75ba..4db5d8be8 100644 --- a/src/dmi/storage/capture/pipeline.py +++ b/src/dmi/storage/capture/pipeline.py @@ -481,6 +481,14 @@ def submit(self, record: CaptureRecord) -> AdmissionResult: self._counters["submitted_records"] += 1 if error is not None: raise PipelineFailedError("capture pipeline failed") from error + if len(record.payload) > self._config.max_pack_bytes: + # Queue admission bounds max_queue_bytes, which is unrelated to + # max_pack_bytes. Without this, a payload no pack could ever hold is + # admitted here and only rejected on the persistence thread, where + # it reads as a fatal pipeline error. + with self._lock: + self._counters["oversized_records"] += 1 + return AdmissionResult.TOO_LARGE started = self._clock_ns() result = self._queue.put( record, @@ -546,9 +554,22 @@ def _run(self) -> None: self._assembler.flush_expired(now_ns=self._clock_ns()) ) continue - self._persist( - self._assembler.append(record, now_ns=self._clock_ns()) - ) + try: + packs = self._assembler.append(record, now_ns=self._clock_ns()) + except OversizedRecordError: + # Framing overhead can push a payload that cleared admission + # past max_pack_bytes. PackAssembler is written to survive + # this with its buffered pack intact, so drop the one record + # and keep going rather than failing the pipeline. + with self._lock: + self._counters["oversized_records"] += 1 + self._emit( + "record_oversized", + capture_id=record.metadata.capture_id, + bytes=len(record.payload), + ) + continue + self._persist(packs) except BaseException as exc: with self._lock: self._error = exc diff --git a/src/dmi/storage/capture/s3.py b/src/dmi/storage/capture/s3.py index 9629fc383..97d1e9fc4 100644 --- a/src/dmi/storage/capture/s3.py +++ b/src/dmi/storage/capture/s3.py @@ -7,7 +7,7 @@ from urllib.parse import urlsplit from uuid import UUID -from .filesystem import validate_object_key, validate_pack_source +from .filesystem import validate_object_key, validate_pack_source, verify_pack_source from .model import ( ObjectInfo, ObjectPage, @@ -182,6 +182,13 @@ def put(self, pack: PackSource, object_key: str) -> PackRef: if existing is not None: return self._existing_ref(pack, key, existing) + # upload_fileobj hands the stream straight to the transfer manager, so + # unlike FilesystemPackStore.put nothing here ever sees the bytes. Hash + # the source first: the object metadata below is written from + # pack.checksum, so a later stat() comparison against it is tautological + # and cannot notice a source that lied about its own contents. + verify_pack_source(pack) + metadata = { "dmi-format": "dmi-pack-v1", "dmi-pack-id": pack.pack_id, @@ -218,6 +225,10 @@ def stat(self, ref: PackRef) -> ObjectInfo: raise PackIntegrityError("S3 object is missing DMI checksum metadata") if size != ref.object_bytes or checksum != ref.checksum: raise PackIntegrityError("S3 object does not match its pack reference") + # NOTE: this checksum is object metadata written at upload time, not a + # digest S3 recomputed from stored bytes. It proves identity and size, + # not that the stored content is intact -- detecting silent corruption + # would take a server-side checksum or a read-back. return ObjectInfo(size=size, checksum=checksum) def inspect(self, object_key: str) -> PackRef: diff --git a/tests/test_capture_catalog_indexer.py b/tests/test_capture_catalog_indexer.py index a1b93a465..de08a94cf 100644 --- a/tests/test_capture_catalog_indexer.py +++ b/tests/test_capture_catalog_indexer.py @@ -265,3 +265,24 @@ def corrupt(ref, offset, length): assert result.failed_packs == 2 assert len(result.failures) == 1 + + +def test_an_oversized_batch_raises_instead_of_blaming_a_pack(tmp_path: Path): + inventory, refs = _packs(tmp_path) + writer = _CatalogWriter() + indexer = CatalogIndexer( + inventory, + writer, + # Small enough that the second pack pushes the batch over. + config=CatalogIndexerConfig(max_packs=8, max_estimated_bytes=1), + clock_ns=lambda: 42, + ) + + # A batch-level bound is a caller error. Reporting it as a per-pack failure + # would blame an innocent pack and silently skip the rest of the batch while + # index() still returned normally. + with pytest.raises(ValueError, match="max_estimated_bytes"): + indexer.index(refs) + + assert writer.descriptor_batches == [] + assert writer.pack_batches == [] diff --git a/tests/test_capture_pipeline.py b/tests/test_capture_pipeline.py index 5dbeaafda..6851a0830 100644 --- a/tests/test_capture_pipeline.py +++ b/tests/test_capture_pipeline.py @@ -368,3 +368,66 @@ def capture_event(event): assert snapshot.flush_linger == 1 assert snapshot.flush_shutdown == 0 + + +def test_pipeline_rejects_a_capture_no_pack_could_hold(tmp_path: Path): + store = FilesystemPackStore(tmp_path / "objects", store_id="local") + pipeline = HostCapturePipeline( + PipelineConfig( + max_queue_records=4, + # The queue would happily take it; max_pack_bytes is what it cannot + # satisfy, and the two bounds are unrelated. + max_queue_bytes=1 << 20, + max_pack_bytes=4096, + max_pack_records=4, + max_linger_ns=1_000_000_000, + ), + DirectPackSink(store), + pack_id_factory=_ids(), + ) + pipeline.start() + + oversized = CaptureRecord( + metadata=replace(_metadata("capture-big"), shape=(2048,)), + payload=b"\x00" * 8192, + ) + assert pipeline.submit(oversized) is AdmissionResult.TOO_LARGE + assert pipeline.submit(_record("capture-a")) is AdmissionResult.ACCEPTED + snapshot = pipeline.close(timeout=2) + + # Rejected at admission, so the caller knows, and the pipeline lives. + assert snapshot.oversized_records == 1 + assert snapshot.failures == 0 + assert snapshot.persisted_records == 1 + + +def test_pipeline_survives_a_capture_only_framing_pushes_over(tmp_path: Path): + store = FilesystemPackStore(tmp_path / "objects", store_id="local") + payload = b"\x00" * 4096 + pipeline = HostCapturePipeline( + PipelineConfig( + max_queue_records=4, + max_queue_bytes=1 << 20, + # The payload itself fits, so admission passes; header, footer and + # trailer are what take it past the limit inside the assembler. + max_pack_bytes=len(payload) + 64, + max_pack_records=4, + max_linger_ns=1_000_000_000, + ), + DirectPackSink(store), + pack_id_factory=_ids(), + ) + pipeline.start() + + big = CaptureRecord( + metadata=replace(_metadata("capture-big"), shape=(1024,)), payload=payload + ) + assert pipeline.submit(big) is AdmissionResult.ACCEPTED + assert pipeline.submit(_record("capture-a")) is AdmissionResult.ACCEPTED + snapshot = pipeline.close(timeout=2) + + # The one record is dropped and counted; the pipeline is not failed and the + # following capture still persists. + assert snapshot.oversized_records == 1 + assert snapshot.failures == 0 + assert snapshot.persisted_records == 1 diff --git a/tests/test_capture_s3.py b/tests/test_capture_s3.py index a53e37352..c1b2b3bb0 100644 --- a/tests/test_capture_s3.py +++ b/tests/test_capture_s3.py @@ -288,3 +288,47 @@ def test_s3_listing_rejects_a_truncated_page_without_a_cursor(): with pytest.raises(PackIntegrityError, match="cursor"): _store(client).list_objects(limit=1) + + +class _LyingPack: + """A pack source whose bytes do not match the checksum it declares.""" + + def __init__(self, sealed): + self.pack_id = sealed.pack_id + self.created_at_ns = sealed.created_at_ns + self.record_count = sealed.record_count + self.checksum = sealed.checksum + self._data = b"\x00" * len(sealed.data) + + @property + def object_bytes(self) -> int: + return len(self._data) + + def open(self): + from io import BytesIO + + return BytesIO(self._data) + + +def test_s3_put_verifies_the_bytes_it_uploads(): + client = _S3Client() + store = _store(client) + + # upload_fileobj hands the stream to the transfer manager, so put never sees + # the bytes. Without hashing the source it would happily store content that + # contradicts the checksum it writes into object metadata -- and stat() + # compares against that same metadata, so nothing downstream could tell. + with pytest.raises(PackIntegrityError, match="checksum"): + store.put(_LyingPack(_pack()), "packs/lying.dmi-pack") + + assert client.objects == {} + + +def test_s3_put_still_accepts_a_faithful_source(): + client = _S3Client() + store = _store(client) + + ref = store.put(_pack(), "packs/honest.dmi-pack") + + assert ref.checksum == _pack().checksum + assert len(client.objects) == 1 diff --git a/tests/tools/run_garage_live.py b/tests/tools/run_garage_live.py index da1839f8a..125c6240b 100644 --- a/tests/tools/run_garage_live.py +++ b/tests/tools/run_garage_live.py @@ -29,11 +29,23 @@ def _ports(count: int) -> list[int]: listener.close() -def _wait_for_port(port: int, process: subprocess.Popen, timeout: float = 20) -> None: +def _log_tail(log_path, limit: int = 2000) -> str: + try: + return log_path.read_text(errors="replace")[-limit:] + except OSError: + return "(no Garage log)" + + +def _wait_for_port( + port: int, process: subprocess.Popen, log_path, timeout: float = 20 +) -> None: deadline = time.monotonic() + timeout while time.monotonic() < deadline: if process.poll() is not None: - raise RuntimeError(f"Garage exited with status {process.returncode}") + raise RuntimeError( + f"Garage exited with status {process.returncode}\n" + f"{_log_tail(log_path)}" + ) try: with socket.create_connection(("127.0.0.1", port), timeout=0.2): return @@ -109,16 +121,21 @@ def main(argv: list[str] | None = None) -> int: "PYTHONPATH": str(REPO_ROOT / "src"), } ) + # Garage logs to stderr for the whole run. A PIPE nobody reads fills its + # buffer and blocks the server mid-benchmark, so send it to a file that + # has no such limit and can still be shown on failure. + log_path = Path(root) / "garage.log" + log_handle = log_path.open("w") process = subprocess.Popen( [binary, "server", "--single-node", "--default-bucket"], cwd=root, env=environment, stdout=subprocess.DEVNULL, - stderr=subprocess.PIPE, + stderr=log_handle, text=True, ) try: - _wait_for_port(s3_port, process) + _wait_for_port(s3_port, process, log_path) command = ( [ sys.executable, @@ -149,6 +166,7 @@ def main(argv: list[str] | None = None) -> int: ) return completed.returncode finally: + log_handle.close() process.terminate() try: process.wait(timeout=5) From d35a0b6b8bb39d8878ad541c92c317460e69fdaf Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Tue, 25 Aug 2026 23:10:19 -0400 Subject: [PATCH 12/20] fix(tests): repoint the CUDA ring tests at native/csrc 3f95d87 (#95) moved the native sources from monitoring/csrc to native/csrc, but the ring test harness that 46e270b added still refers to the old tree: the Makefile's include path and its producer/drain sources, plus four #includes in test_rings.cu. None of those paths exist, so `make` in tests/ring fails outright on missing files and headers. Every referenced file is present under native/csrc/ring/ with an unchanged name, so this is a path substitution only. Verified by resolving each Makefile path and each -I-relative include; not compile-verified, since this machine has no nvcc. Same regression family as the resurrected Python test modules removed earlier in this branch -- pre-refactor paths that survived the src-layout move. --- tests/ring/Makefile | 6 +++--- tests/ring/test_rings.cu | 8 ++++---- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/ring/Makefile b/tests/ring/Makefile index f02f27f70..ee86fd7d8 100644 --- a/tests/ring/Makefile +++ b/tests/ring/Makefile @@ -2,7 +2,7 @@ NVCC ?= /usr/local/cuda/bin/nvcc SM_ARCH ?= native CXX_STD ?= c++17 BUILD := build -INCLUDES := -I../.. -I../../monitoring/csrc +INCLUDES := -I../.. -I../../native/csrc NVCC_FLAGS := \ -std=$(CXX_STD) \ @@ -14,8 +14,8 @@ NVCC_FLAGS := \ -lineinfo \ $(INCLUDES) -PRODUCER_CU := ../../monitoring/csrc/ring/producer.cu -DRAIN_SRC := ../../monitoring/csrc/ring/drain_thread.cpp +PRODUCER_CU := ../../native/csrc/ring/producer.cu +DRAIN_SRC := ../../native/csrc/ring/drain_thread.cpp # test_ring_engine needs ATen headers + libtorch for DrainTask (contains at::Tensor). # Use torch.utils.cmake_prefix_path to locate them. The Python E2E tests diff --git a/tests/ring/test_rings.cu b/tests/ring/test_rings.cu index 738431c70..01759a334 100644 --- a/tests/ring/test_rings.cu +++ b/tests/ring/test_rings.cu @@ -18,10 +18,10 @@ #include // Include the ring headers under test. -#include "../../monitoring/csrc/ring/task_entry.h" -#include "../../monitoring/csrc/ring/ring_config.h" -#include "../../monitoring/csrc/ring/payload_ring.cuh" -#include "../../monitoring/csrc/ring/task_ring.cuh" +#include "../../native/csrc/ring/task_entry.h" +#include "../../native/csrc/ring/ring_config.h" +#include "../../native/csrc/ring/payload_ring.cuh" +#include "../../native/csrc/ring/task_ring.cuh" // --------------------------------------------------------------------------- // Helpers From 7acc7f0b1e48054a8067e50ec675ce06a8d584c1 Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Tue, 25 Aug 2026 23:24:50 -0400 Subject: [PATCH 13/20] Revert "fix(tests): repoint the CUDA ring tests at native/csrc" This reverts commit d35a0b6. Out of scope: this branch is the CPU-side capture path, and the CUDA ring tests belong to the GPU workstream. The defect is real and still there -- tests/ring/Makefile and tests/ring/test_rings.cu carry seven references to monitoring/csrc, which #95 moved to native/csrc, so `make` in tests/ring fails on missing files -- but it should be fixed alongside the rest of the ring work by someone who can actually compile it. Nothing here can: the machine has no nvcc, so the previous commit was path-verified only. Recorded in the PR description as a finding rather than a fix. --- tests/ring/Makefile | 6 +++--- tests/ring/test_rings.cu | 8 ++++---- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/ring/Makefile b/tests/ring/Makefile index ee86fd7d8..f02f27f70 100644 --- a/tests/ring/Makefile +++ b/tests/ring/Makefile @@ -2,7 +2,7 @@ NVCC ?= /usr/local/cuda/bin/nvcc SM_ARCH ?= native CXX_STD ?= c++17 BUILD := build -INCLUDES := -I../.. -I../../native/csrc +INCLUDES := -I../.. -I../../monitoring/csrc NVCC_FLAGS := \ -std=$(CXX_STD) \ @@ -14,8 +14,8 @@ NVCC_FLAGS := \ -lineinfo \ $(INCLUDES) -PRODUCER_CU := ../../native/csrc/ring/producer.cu -DRAIN_SRC := ../../native/csrc/ring/drain_thread.cpp +PRODUCER_CU := ../../monitoring/csrc/ring/producer.cu +DRAIN_SRC := ../../monitoring/csrc/ring/drain_thread.cpp # test_ring_engine needs ATen headers + libtorch for DrainTask (contains at::Tensor). # Use torch.utils.cmake_prefix_path to locate them. The Python E2E tests diff --git a/tests/ring/test_rings.cu b/tests/ring/test_rings.cu index 01759a334..738431c70 100644 --- a/tests/ring/test_rings.cu +++ b/tests/ring/test_rings.cu @@ -18,10 +18,10 @@ #include // Include the ring headers under test. -#include "../../native/csrc/ring/task_entry.h" -#include "../../native/csrc/ring/ring_config.h" -#include "../../native/csrc/ring/payload_ring.cuh" -#include "../../native/csrc/ring/task_ring.cuh" +#include "../../monitoring/csrc/ring/task_entry.h" +#include "../../monitoring/csrc/ring/ring_config.h" +#include "../../monitoring/csrc/ring/payload_ring.cuh" +#include "../../monitoring/csrc/ring/task_ring.cuh" // --------------------------------------------------------------------------- // Helpers From 726e45c53391d1dead8c6a150c36d8844eac4542 Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Tue, 25 Aug 2026 23:27:21 -0400 Subject: [PATCH 14/20] fix(tests): repoint the CUDA ring tests at native/csrc 3f95d87 (#95) moved the native sources from monitoring/csrc to native/csrc, but the ring test harness that 46e270b added still refers to the old tree: the Makefile's include path and its producer/drain sources, plus four #includes in test_rings.cu. None of those paths exist, so `make` in tests/ring fails outright on missing files and headers. Nothing but the directory name changes. All fourteen changed lines collapse to seven once monitoring/csrc and native/csrc are normalised -- no CUDA code, no build flags, no logic. Every referenced file is present under native/csrc/ring/ with an unchanged name. Verified by resolving each Makefile path and each -I-relative include. Not compile-verified: this machine has no nvcc. Same regression family as the resurrected Python test modules removed earlier in this branch -- pre-refactor paths that survived the src-layout move. --- tests/ring/Makefile | 6 +++--- tests/ring/test_rings.cu | 8 ++++---- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/ring/Makefile b/tests/ring/Makefile index f02f27f70..ee86fd7d8 100644 --- a/tests/ring/Makefile +++ b/tests/ring/Makefile @@ -2,7 +2,7 @@ NVCC ?= /usr/local/cuda/bin/nvcc SM_ARCH ?= native CXX_STD ?= c++17 BUILD := build -INCLUDES := -I../.. -I../../monitoring/csrc +INCLUDES := -I../.. -I../../native/csrc NVCC_FLAGS := \ -std=$(CXX_STD) \ @@ -14,8 +14,8 @@ NVCC_FLAGS := \ -lineinfo \ $(INCLUDES) -PRODUCER_CU := ../../monitoring/csrc/ring/producer.cu -DRAIN_SRC := ../../monitoring/csrc/ring/drain_thread.cpp +PRODUCER_CU := ../../native/csrc/ring/producer.cu +DRAIN_SRC := ../../native/csrc/ring/drain_thread.cpp # test_ring_engine needs ATen headers + libtorch for DrainTask (contains at::Tensor). # Use torch.utils.cmake_prefix_path to locate them. The Python E2E tests diff --git a/tests/ring/test_rings.cu b/tests/ring/test_rings.cu index 738431c70..01759a334 100644 --- a/tests/ring/test_rings.cu +++ b/tests/ring/test_rings.cu @@ -18,10 +18,10 @@ #include // Include the ring headers under test. -#include "../../monitoring/csrc/ring/task_entry.h" -#include "../../monitoring/csrc/ring/ring_config.h" -#include "../../monitoring/csrc/ring/payload_ring.cuh" -#include "../../monitoring/csrc/ring/task_ring.cuh" +#include "../../native/csrc/ring/task_entry.h" +#include "../../native/csrc/ring/ring_config.h" +#include "../../native/csrc/ring/payload_ring.cuh" +#include "../../native/csrc/ring/task_ring.cuh" // --------------------------------------------------------------------------- // Helpers From 335246202c9d208b9a29399b26303bc3e775a97e Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Wed, 26 Aug 2026 00:15:44 -0400 Subject: [PATCH 15/20] test(capture): add fault injection and a conformance manifest Phase 6 requires fault injection before the default sink can be switched, and comparison of golden workloads "by identity, logical bytes, checksums, decoded tensors, and query results". Neither instrument existed. Both are built here as the specification a native writer has to satisfy, since the Python implementation is the reference rather than the production writer. tests/_faults.py wraps the three boundaries that can misbehave -- object store, ClickHouse client, pack sink -- with scripted, never random, faults: a schedule names which 1-based call numbers fail and how, so a failing test reproduces exactly. Short reads, hard failures, duplicated writes and transient outages are each expressible in one call. The nine tests that use it assert observable consequences rather than internals: a short trailer read is refused rather than silently truncating; a read failure aborts indexing instead of producing a partial pack; an immutable key written twice converges rather than conflicting; a sink failure fails the pipeline loudly and then refuses further admission; an insert failure leaves the pack uncommitted and the batch replayable; a duplicated insert is absorbed by replay semantics; and one corrupt pack fails only itself rather than poisoning its batch. tests/tools/golden_workload.py produces the conformance manifest: one JSON document covering every dtype the format accepts, with pack identity and checksum, per-capture payload sha256 and crc32, decoded-tensor sha256, placement, and the full summary contract. Everything in it is language-neutral -- byte counts, hex digests, integers -- so a C++ writer can be checked against it directly. `verify` diffs field by field, naming the capture and field that moved rather than only reporting inequality. Generating that manifest immediately found a real defect in the Phase 5 summarizer: l2_norm computed sqrt(sum(x**2)), which overflows float64 for large-magnitude tensors and returns inf where the true norm is finite -- 1e200 needs 1e400 to square. Now scaled by the largest magnitude before squaring, so every squared term stays at or below 1. Verified to agree with the plain formula at normal scale, and the manifest comparison is mutation-tested: regressing the norm produces 10 named differences. --- src/dmi/storage/capture/summary.py | 14 +- tests/_faults.py | 268 +++++++++++++++++++++ tests/data/capture_golden_manifest.json | 299 ++++++++++++++++++++++++ tests/test_capture_faults.py | 284 ++++++++++++++++++++++ tests/test_capture_golden_workload.py | 83 +++++++ tests/test_capture_summary.py | 24 ++ tests/tools/__init__.py | 0 tests/tools/golden_workload.py | 280 ++++++++++++++++++++++ 8 files changed, 1250 insertions(+), 2 deletions(-) create mode 100644 tests/_faults.py create mode 100644 tests/data/capture_golden_manifest.json create mode 100644 tests/test_capture_faults.py create mode 100644 tests/test_capture_golden_workload.py create mode 100644 tests/tools/__init__.py create mode 100644 tests/tools/golden_workload.py diff --git a/src/dmi/storage/capture/summary.py b/src/dmi/storage/capture/summary.py index 5fc2e60d6..2343c0dda 100644 --- a/src/dmi/storage/capture/summary.py +++ b/src/dmi/storage/capture/summary.py @@ -179,6 +179,16 @@ def summarize_tensor( ) absolute = numpy.abs(finite) + # Scale before squaring. float64 holds values up to ~1e308, so squaring a + # large-magnitude tensor overflows and the naive sqrt(sum(x**2)) returns inf + # where the true norm is perfectly finite. Factoring out the largest + # magnitude keeps every squared term at or below 1. + abs_max = float(absolute.max()) + if abs_max == 0.0: + l2_norm = 0.0 + else: + l2_norm = abs_max * float(numpy.sqrt(numpy.square(absolute / abs_max).sum())) + return CoreTensorSummaryV1( summary_version=CORE_SUMMARY_VERSION, element_count=element_count, @@ -189,8 +199,8 @@ def summarize_tensor( mean=float(finite.mean()), minimum=float(finite.min()), maximum=float(finite.max()), - abs_max=float(absolute.max()), - l2_norm=float(numpy.sqrt(numpy.square(finite).sum())), + abs_max=abs_max, + l2_norm=l2_norm, ) diff --git a/tests/_faults.py b/tests/_faults.py new file mode 100644 index 000000000..2e056cf23 --- /dev/null +++ b/tests/_faults.py @@ -0,0 +1,268 @@ +"""Deterministic fault injection for the capture storage path. + +Phase 6 requires fault injection before the default sink can be switched, and +the components that must survive faults are spread across three boundaries: the +object store, the ClickHouse client, and the pack sink. This module provides one +wrapper per boundary so a test states *which* fault it wants rather than +hand-rolling another one-off fake. + +Every fault is **scripted, not random**. A schedule says which call numbers fail +and how, so a failing test reproduces exactly. Nothing here uses a random source. + + store = FaultyPackStore(inner, read_range=fail_on(2, OSError("reset"))) + store.read_range(ref, 0, 16) # call 1 -- fine + store.read_range(ref, 0, 16) # call 2 -- raises OSError + +These wrappers also serve the conformance role the Python implementation now +has: they describe the failure behaviour a native writer has to reproduce. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Mapping, Sequence + + +class FaultInjected(Exception): + """Raised by a scripted fault, so tests can tell it from a real bug.""" + + +# --- schedules --------------------------------------------------------------- + + +@dataclass(frozen=True, slots=True) +class Schedule: + """Which calls misbehave, and how. + + ``calls`` are 1-based call numbers. A call not listed behaves normally, so + the default schedule is a pass-through. + """ + + calls: frozenset[int] = frozenset() + error: BaseException | None = None + truncate_by: int = 0 + repeat_result: bool = False + + def applies_to(self, call_number: int) -> bool: + return call_number in self.calls + + +def never() -> Schedule: + """No faults. Useful as an explicit default.""" + return Schedule() + + +def fail_on(*calls: int, error: BaseException | None = None) -> Schedule: + """Raise on the given 1-based call numbers.""" + if not calls: + raise ValueError("fail_on needs at least one call number") + return Schedule( + calls=frozenset(calls), + error=error or FaultInjected("scripted failure"), + ) + + +def truncate_on(*calls: int, by: int = 1) -> Schedule: + """Return fewer bytes than asked for -- a short read. + + Object stores are allowed to return short reads, and code that assumes + otherwise corrupts payloads silently rather than failing. + """ + if by <= 0: + raise ValueError("truncate_on needs a positive byte count") + return Schedule(calls=frozenset(calls), truncate_by=by) + + +def duplicate_on(*calls: int) -> Schedule: + """Apply the call twice -- an ambiguous write that actually landed twice.""" + return Schedule(calls=frozenset(calls), repeat_result=True) + + +def fail_then_succeed(count: int, error: BaseException | None = None) -> Schedule: + """Fail the first ``count`` calls, then behave. Models a transient outage.""" + if count <= 0: + raise ValueError("fail_then_succeed needs a positive count") + return Schedule( + calls=frozenset(range(1, count + 1)), + error=error or FaultInjected("transient failure"), + ) + + +# --- object store ------------------------------------------------------------ + + +@dataclass +class _Counter: + counts: dict[str, int] = field(default_factory=dict) + + def bump(self, name: str) -> int: + self.counts[name] = self.counts.get(name, 0) + 1 + return self.counts[name] + + +class FaultyPackStore: + """Wraps any PackStore, injecting scripted faults per method. + + Delegates everything it does not intercept, so it stays usable as a drop-in + even as the store protocol grows. + """ + + def __init__( + self, + inner: Any, + *, + put: Schedule | None = None, + stat: Schedule | None = None, + read_range: Schedule | None = None, + ) -> None: + self._inner = inner + self._schedules = { + "put": put or never(), + "stat": stat or never(), + "read_range": read_range or never(), + } + self._calls = _Counter() + + @property + def store_id(self) -> str: + return self._inner.store_id + + @property + def call_counts(self) -> Mapping[str, int]: + return dict(self._calls.counts) + + def _guard(self, name: str) -> Schedule | None: + schedule = self._schedules[name] + number = self._calls.bump(name) + if not schedule.applies_to(number): + return None + if schedule.error is not None: + raise schedule.error + return schedule + + def put(self, pack: Any, object_key: str) -> Any: + schedule = self._guard("put") + result = self._inner.put(pack, object_key) + if schedule is not None and schedule.repeat_result: + # The same object written twice: the second write must be a no-op + # for an immutable key, not a conflict. + self._inner.put(pack, object_key) + return result + + def stat(self, ref: Any) -> Any: + self._guard("stat") + return self._inner.stat(ref) + + def read_range(self, ref: Any, offset: int, length: int) -> bytes: + schedule = self._guard("read_range") + data = self._inner.read_range(ref, offset, length) + if schedule is not None and schedule.truncate_by: + return data[: max(0, len(data) - schedule.truncate_by)] + return data + + def __getattr__(self, name: str) -> Any: + return getattr(self._inner, name) + + +# --- ClickHouse client ------------------------------------------------------- + + +class FaultyClickHouseClient: + """Wraps a ClickHouse client, injecting faults by statement kind. + + Statements are classified by their leading keyword, so a test can fail only + inserts while leaving schema and reads alone -- which is the interesting + case for an indexer that must not lose or double-count rows. + """ + + def __init__( + self, + inner: Any, + *, + insert: Schedule | None = None, + select: Schedule | None = None, + ddl: Schedule | None = None, + ) -> None: + self._inner = inner + self._schedules = { + "insert": insert or never(), + "select": select or never(), + "ddl": ddl or never(), + } + self._calls = _Counter() + self.statements: list[str] = [] + + @property + def call_counts(self) -> Mapping[str, int]: + return dict(self._calls.counts) + + @staticmethod + def _kind(query: str) -> str: + head = query.lstrip().split(None, 1)[0].upper() if query.strip() else "" + if head == "INSERT": + return "insert" + if head in {"SELECT", "WITH"}: + return "select" + return "ddl" + + def execute(self, query: str, params: Any = None, **kwargs: Any) -> Any: + kind = self._kind(query) + self.statements.append(query) + number = self._calls.bump(kind) + schedule = self._schedules[kind] + if schedule.applies_to(number): + if schedule.error is not None: + raise schedule.error + if schedule.repeat_result: + # An ambiguous insert: the client never learned it succeeded, so + # the row lands twice and dedup has to absorb it. + self._inner.execute(query, params, **kwargs) + return self._inner.execute(query, params, **kwargs) + + def __getattr__(self, name: str) -> Any: + return getattr(self._inner, name) + + +# --- pack sink --------------------------------------------------------------- + + +class FaultyPackSink: + """Wraps a pack sink so persistence can fail on chosen packs.""" + + def __init__(self, inner: Any, *, persist: Schedule | None = None) -> None: + self._inner = inner + self._schedule = persist or never() + self._calls = _Counter() + + @property + def call_counts(self) -> Mapping[str, int]: + return dict(self._calls.counts) + + def persist(self, ready: Any) -> Any: + number = self._calls.bump("persist") + if self._schedule.applies_to(number) and self._schedule.error is not None: + raise self._schedule.error + return self._inner.persist(ready) + + def __getattr__(self, name: str) -> Any: + return getattr(self._inner, name) + + +def call_sequence(store: FaultyPackStore) -> Sequence[tuple[str, int]]: + """The recorded call counts, sorted -- handy in assertion messages.""" + return sorted(store.call_counts.items()) + + +__all__ = [ + "FaultInjected", + "FaultyClickHouseClient", + "FaultyPackSink", + "FaultyPackStore", + "Schedule", + "call_sequence", + "duplicate_on", + "fail_on", + "fail_then_succeed", + "never", + "truncate_on", +] diff --git a/tests/data/capture_golden_manifest.json b/tests/data/capture_golden_manifest.json new file mode 100644 index 000000000..86b9cf308 --- /dev/null +++ b/tests/data/capture_golden_manifest.json @@ -0,0 +1,299 @@ +{ + "captures": [ + { + "capture_id": "golden-00", + "decoded_dtype": "bool", + "decoded_length": 16, + "decoded_sha256": "278abfb53cf32e306305c34f6ad5469b82453fa58f29fd339bf1615e7ffb4331", + "dtype": "bool", + "logical_bytes": 16, + "offset": 64, + "payload_crc32": "0dea99fc", + "payload_sha256": "278abfb53cf32e306305c34f6ad5469b82453fa58f29fd339bf1615e7ffb4331", + "shape": [ + 16 + ], + "stored_length": 16, + "summary": { + "abs_max": 1.0, + "element_count": 16, + "finite_count": 16, + "inf_count": 0, + "l2_norm": 2.828427125, + "maximum": 1.0, + "mean": 0.5, + "minimum": 0.0, + "nan_count": 0, + "version": 1, + "zero_fraction": 0.5 + } + }, + { + "capture_id": "golden-01", + "decoded_dtype": "uint8", + "decoded_length": 16, + "decoded_sha256": "60812b1c711e79b788f427998acfb0822d021ac5641a5bbe53a5841c52eccd1b", + "dtype": "uint8", + "logical_bytes": 16, + "offset": 128, + "payload_crc32": "2973360b", + "payload_sha256": "60812b1c711e79b788f427998acfb0822d021ac5641a5bbe53a5841c52eccd1b", + "shape": [ + 16 + ], + "stored_length": 16, + "summary": { + "abs_max": 240.0, + "element_count": 16, + "finite_count": 16, + "inf_count": 0, + "l2_norm": 564.205636271, + "maximum": 240.0, + "mean": 121.5, + "minimum": 1.0, + "nan_count": 0, + "version": 1, + "zero_fraction": 0.0 + } + }, + { + "capture_id": "golden-02", + "decoded_dtype": "int8", + "decoded_length": 16, + "decoded_sha256": "f6d7f9860dc2e0284769a6909e33b68f5eb39224dc0bcedc4946473ff7082bc0", + "dtype": "int8", + "logical_bytes": 16, + "offset": 192, + "payload_crc32": "7f5a8edf", + "payload_sha256": "f6d7f9860dc2e0284769a6909e33b68f5eb39224dc0bcedc4946473ff7082bc0", + "shape": [ + 16 + ], + "stored_length": 16, + "summary": { + "abs_max": 124.0, + "element_count": 16, + "finite_count": 16, + "inf_count": 0, + "l2_norm": 286.928562538, + "maximum": 115.0, + "mean": -3.5, + "minimum": -124.0, + "nan_count": 0, + "version": 1, + "zero_fraction": 0.0 + } + }, + { + "capture_id": "golden-03", + "decoded_dtype": "int16", + "decoded_length": 32, + "decoded_sha256": "937f1818be90b526b6896e6f45a888814027c50c7bdb7c22841acd3f3dd7e2bb", + "dtype": "int16", + "logical_bytes": 32, + "offset": 256, + "payload_crc32": "98cddfaa", + "payload_sha256": "937f1818be90b526b6896e6f45a888814027c50c7bdb7c22841acd3f3dd7e2bb", + "shape": [ + 16 + ], + "stored_length": 32, + "summary": { + "abs_max": 32638.0, + "element_count": 16, + "finite_count": 16, + "inf_count": 0, + "l2_norm": 78531.039621286, + "maximum": 32638.0, + "mean": 255.5, + "minimum": -28529.0, + "nan_count": 0, + "version": 1, + "zero_fraction": 0.0 + } + }, + { + "capture_id": "golden-04", + "decoded_dtype": "float16", + "decoded_length": 32, + "decoded_sha256": "72ee199c0c46864fd549df2765d48a01115831a3b7531a0cc0717bd8061eb0e5", + "dtype": "float16", + "logical_bytes": 32, + "offset": 320, + "payload_crc32": "5d664ad8", + "payload_sha256": "72ee199c0c46864fd549df2765d48a01115831a3b7531a0cc0717bd8061eb0e5", + "shape": [ + 16 + ], + "stored_length": 32, + "summary": { + "abs_max": 53024.0, + "element_count": 16, + "finite_count": 16, + "inf_count": 0, + "l2_norm": 54474.591040461, + "maximum": 53024.0, + "mean": 2695.128854681, + "minimum": -12160.0, + "nan_count": 0, + "version": 1, + "zero_fraction": 0.0 + } + }, + { + "capture_id": "golden-05", + "decoded_dtype": "float32", + "decoded_length": 32, + "decoded_sha256": "156e92df170edae45737c928927c845898a47ddf7c5317f68d0af016f5ab6c2b", + "dtype": "bfloat16", + "logical_bytes": 32, + "offset": 384, + "payload_crc32": "e806ff9a", + "payload_sha256": "0075d8383291de11aca58e21015d3f2956a05643e29eca2c2a49dae65f8e4a6e", + "shape": [ + 16 + ], + "stored_length": 32, + "summary": { + "abs_max": 4.187068186722485e+37, + "element_count": 16, + "finite_count": 16, + "inf_count": 0, + "l2_norm": 4.1870681867295667e+37, + "maximum": 7.700977396386494e+31, + "mean": -2.616912803732729e+36, + "minimum": -4.187068186722485e+37, + "nan_count": 0, + "version": 1, + "zero_fraction": 0.0 + } + }, + { + "capture_id": "golden-06", + "decoded_dtype": "int32", + "decoded_length": 64, + "decoded_sha256": "67b025ed6cf026a5811a0be5c1147910e0ba3ae5dab3b2359b6463661b5a265c", + "dtype": "int32", + "logical_bytes": 64, + "offset": 448, + "payload_crc32": "d9edb076", + "payload_sha256": "67b025ed6cf026a5811a0be5c1147910e0ba3ae5dab3b2359b6463661b5a265c", + "shape": [ + 16 + ], + "stored_length": 64, + "summary": { + "abs_max": 2105442177.0, + "element_count": 16, + "finite_count": 16, + "inf_count": 0, + "l2_norm": 4828056749.188784, + "maximum": 1903193966.0, + "mean": -67438087.5, + "minimum": -2105442177.0, + "nan_count": 0, + "version": 1, + "zero_fraction": 0.0 + } + }, + { + "capture_id": "golden-07", + "decoded_dtype": "float32", + "decoded_length": 64, + "decoded_sha256": "77e484d7416c93592f2c8e8bb3fe8523ca558e43f998edce12f9d62b2cee9bb8", + "dtype": "float32", + "logical_bytes": 64, + "offset": 512, + "payload_crc32": "f9540b9b", + "payload_sha256": "77e484d7416c93592f2c8e8bb3fe8523ca558e43f998edce12f9d62b2cee9bb8", + "shape": [ + 16 + ], + "stored_length": 64, + "summary": { + "abs_max": 5.223002208286206e+36, + "element_count": 16, + "finite_count": 16, + "inf_count": 0, + "l2_norm": 5.223002210574896e+36, + "maximum": 5.223002208286206e+36, + "mean": 3.2642797422049176e+35, + "minimum": -1.5462104171572421e+32, + "nan_count": 0, + "version": 1, + "zero_fraction": 0.0 + } + }, + { + "capture_id": "golden-08", + "decoded_dtype": "int64", + "decoded_length": 128, + "decoded_sha256": "21d85addc7434f2dcbebc4862e3275953d4151b8b20ac0e7090cbb44e4f3ea5b", + "dtype": "int64", + "logical_bytes": 128, + "offset": 576, + "payload_crc32": "278c92d6", + "payload_sha256": "21d85addc7434f2dcbebc4862e3275953d4151b8b20ac0e7090cbb44e4f3ea5b", + "shape": [ + 16 + ], + "stored_length": 128, + "summary": { + "abs_max": 8.897557572131124e+18, + "element_count": 16, + "finite_count": 16, + "inf_count": 0, + "l2_norm": 2.0775646156498567e+19, + "maximum": 8.897557572131124e+18, + "mean": 4.292360895432622e+17, + "minimum": -8.319403563331123e+18, + "nan_count": 0, + "version": 1, + "zero_fraction": 0.0 + } + }, + { + "capture_id": "golden-09", + "decoded_dtype": "float64", + "decoded_length": 128, + "decoded_sha256": "94840849b28423febac68daa09f0bca9b15721a4ef4177f3fe30e70f4d438e8d", + "dtype": "float64", + "logical_bytes": 128, + "offset": 704, + "payload_crc32": "1b334372", + "payload_sha256": "94840849b28423febac68daa09f0bca9b15721a4ef4177f3fe30e70f4d438e8d", + "shape": [ + 16 + ], + "stored_length": 128, + "summary": { + "abs_max": 5.141221720570257e+303, + "element_count": 16, + "finite_count": 16, + "inf_count": 0, + "l2_norm": 5.141221720570257e+303, + "maximum": 6.141438560152791e+257, + "mean": -3.2132635753564107e+302, + "minimum": -5.141221720570257e+303, + "nan_count": 0, + "version": 1, + "zero_fraction": 0.0 + } + } + ], + "hydration": { + "capture_count": 10, + "logical_bytes": 528, + "object_count": 1, + "request_bytes": 528, + "request_count": 7, + "stored_bytes": 528 + }, + "manifest_version": 1, + "pack": { + "object_bytes": 7320, + "pack_id": "018f0000-0000-7000-8000-000000000f01", + "record_count": 10, + "sha256": "53a0873af5b5932ceb3e44223492aec11eadfb1d9298cb4cef81d8ca5337fd4e" + } +} diff --git a/tests/test_capture_faults.py b/tests/test_capture_faults.py new file mode 100644 index 000000000..e17ab191f --- /dev/null +++ b/tests/test_capture_faults.py @@ -0,0 +1,284 @@ +"""How the capture storage path behaves when its dependencies misbehave. + +Phase 6 cannot switch the default sink until this is characterised, and the +Python implementation is now the conformance reference -- so these tests double +as the specification a native writer has to satisfy. Each one names the fault, +then asserts the *observable* consequence: what survives, what is reported, and +what is never silently lost. +""" + +from __future__ import annotations + +from pathlib import Path +import time +from uuid import UUID + +import pytest + +from tests._faults import ( + FaultInjected, + FaultyClickHouseClient, + FaultyPackSink, + FaultyPackStore, + duplicate_on, + fail_on, + fail_then_succeed, + truncate_on, +) + +from dmi.storage.capture import ( + CaptureMetadata, + CaptureRecord, + CatalogIndexer, + ClickHouseCatalogConfig, + ClickHouseCatalogWriter, + DirectPackSink, + FilesystemPackStore, + HostCapturePipeline, + PackIndex, + PackFormatError, + PackWriter, + PipelineConfig, + PipelineFailedError, +) + + +pytestmark = pytest.mark.cpu + + +PACK_ID = UUID("018f0000-0000-7000-8000-00000000fa01") + + +def _metadata(capture_id: str, *, step: int = 0) -> CaptureMetadata: + return CaptureMetadata( + capture_id=capture_id, + tenant_id="tenant-a", + experiment_id="exp-a", + run_id="run-a", + session_id="session-a", + request_id="request-a", + sequence_id="sequence-a", + model_id="model-a", + model_revision="revision-a", + adapter_revision=None, + capture_policy_version="policy-v1", + hook_name="resid_pre", + layer_number=3, + producer_rank=0, + step_number=step, + token_start=step, + token_end=step + 1, + batch_position=0, + dtype="float32", + shape=(2,), + captured_at_ns=1_700_000_000_000_000_000 + step, + ) + + +def _record(capture_id: str, *, step: int = 0) -> CaptureRecord: + return CaptureRecord(metadata=_metadata(capture_id, step=step), payload=b"\x01" * 8) + + +def _sealed(*records: CaptureRecord, pack_id: UUID = PACK_ID): + writer = PackWriter( + pack_id=pack_id, + created_at_ns=1_700_000_000_000_000_000, + max_pack_bytes=1024 * 1024, + ) + for record in records: + writer.append(record) + return writer.seal() + + +def _ids(): + from itertools import count + + counter = count(1) + return lambda: UUID(int=next(counter)) + + +# --- object store faults ----------------------------------------------------- + + +def test_a_short_read_is_detected_rather_than_silently_truncating(tmp_path: Path): + inner = FilesystemPackStore(tmp_path, store_id="local") + sealed = _sealed(_record("capture-a")) + ref = inner.put(sealed, "packs/a.dmi-pack") + store = FaultyPackStore(inner, read_range=truncate_on(1, by=1)) + + # A store returning fewer bytes than asked for must never be mistaken for a + # valid short object -- that would corrupt a payload silently. The trailer + # is fixed-width, so a short read cannot even be unpacked. + with pytest.raises(PackFormatError, match="trailer is truncated"): + PackIndex.from_store(store, ref).descriptors() + + +def test_a_read_failure_propagates_instead_of_producing_a_partial_pack( + tmp_path: Path, +): + inner = FilesystemPackStore(tmp_path, store_id="local") + sealed = _sealed(_record("capture-a"), _record("capture-b", step=1)) + ref = inner.put(sealed, "packs/a.dmi-pack") + store = FaultyPackStore(inner, read_range=fail_on(2)) + + with pytest.raises(FaultInjected): + PackIndex.from_store(store, ref).descriptors() + + # The trailer read happened; the footer read is what failed. + assert store.call_counts["read_range"] == 2 + + +def test_an_immutable_key_written_twice_is_not_a_conflict(tmp_path: Path): + inner = FilesystemPackStore(tmp_path, store_id="local") + store = FaultyPackStore(inner, put=duplicate_on(1)) + sealed = _sealed(_record("capture-a")) + + # An ambiguous upload that actually landed twice must converge, because the + # writer cannot tell "never arrived" from "arrived, ack lost". + ref = store.put(sealed, "packs/a.dmi-pack") + + assert ref.checksum == sealed.checksum + assert store.stat(ref).size == len(sealed.data) + + +# --- pipeline / sink faults -------------------------------------------------- + + +def test_a_sink_failure_is_surfaced_and_closes_admission(tmp_path: Path): + store = FilesystemPackStore(tmp_path / "objects", store_id="local") + sink = FaultyPackSink(DirectPackSink(store), persist=fail_on(1)) + pipeline = HostCapturePipeline( + PipelineConfig( + max_queue_records=4, + max_queue_bytes=1 << 20, + max_pack_bytes=1024 * 1024, + max_pack_records=1, + max_linger_ns=1_000_000_000, + ), + sink, + pack_id_factory=_ids(), + ) + pipeline.start() + pipeline.submit(_record("capture-a")) + + # Losing durable storage is not recoverable at this layer, so it must fail + # loudly rather than continue accepting captures it cannot persist. + with pytest.raises(PipelineFailedError): + pipeline.close(timeout=2) + assert pipeline.snapshot().failures == 1 + + +def test_admission_is_refused_once_the_pipeline_has_failed(tmp_path: Path): + store = FilesystemPackStore(tmp_path / "objects", store_id="local") + sink = FaultyPackSink(DirectPackSink(store), persist=fail_on(1)) + pipeline = HostCapturePipeline( + PipelineConfig( + max_queue_records=4, + max_queue_bytes=1 << 20, + max_pack_bytes=1024 * 1024, + max_pack_records=1, + max_linger_ns=1_000_000_000, + ), + sink, + pack_id_factory=_ids(), + ) + pipeline.start() + pipeline.submit(_record("capture-a")) + + deadline = time.monotonic() + 5 + while time.monotonic() < deadline and pipeline.snapshot().failures == 0: + time.sleep(0.01) + assert pipeline.snapshot().failures == 1 + + # Once persistence is broken, further submissions must be refused rather + # than accepted into a queue nothing will ever drain. + with pytest.raises(PipelineFailedError): + pipeline.submit(_record("capture-c", step=2)) + with pytest.raises(PipelineFailedError): + pipeline.close(timeout=2) + + +# --- catalog / ClickHouse faults --------------------------------------------- + + +class _Client: + """Minimal in-memory ClickHouse stand-in that records inserted rows.""" + + def __init__(self): + self.inserted: list[tuple[str, list]] = [] + + def execute(self, query, params=None, **kwargs): + if query.lstrip().upper().startswith("INSERT"): + self.inserted.append((query, list(params or []))) + return [] + return [] + + +def _indexer(store, client, **config): + writer = ClickHouseCatalogWriter(client, ClickHouseCatalogConfig(**config)) + return CatalogIndexer(store, writer, clock_ns=lambda: 7) + + +def test_an_insert_failure_does_not_commit_the_pack(tmp_path: Path): + inner = FilesystemPackStore(tmp_path, store_id="local") + ref = inner.put(_sealed(_record("capture-a")), "packs/a.dmi-pack") + client = FaultyClickHouseClient(_Client(), insert=fail_on(1)) + + with pytest.raises(FaultInjected): + _indexer(inner, client).index([ref]) + + # Descriptors are written before the pack commit marker precisely so a + # failure here leaves the pack uncommitted and the batch replayable. + assert client.call_counts.get("insert") == 1 + + +def test_a_duplicated_insert_is_absorbed_by_replay_semantics(tmp_path: Path): + inner = FilesystemPackStore(tmp_path, store_id="local") + ref = inner.put(_sealed(_record("capture-a")), "packs/a.dmi-pack") + backing = _Client() + client = FaultyClickHouseClient(backing, insert=duplicate_on(1)) + + result = _indexer(inner, client).index([ref]) + + # The physical row lands twice; that is expected and is what the + # ReplacingMergeTree views collapse. What must not happen is a failure. + assert result.failed_packs == 0 + assert result.indexed_rows == 1 + descriptor_inserts = [q for q, _ in backing.inserted if "capture_raw" in q] + assert len(descriptor_inserts) == 2 + + +def test_a_corrupt_pack_fails_only_its_own_pack(tmp_path: Path): + inner = FilesystemPackStore(tmp_path, store_id="local") + good = inner.put(_sealed(_record("capture-a")), "packs/good.dmi-pack") + bad = inner.put( + _sealed(_record("capture-b", step=1), pack_id=UUID(int=0xBAD)), + "packs/bad.dmi-pack", + ) + # Fail the footer read of the second pack only. + store = FaultyPackStore(inner, read_range=fail_on(3, 4)) + + result = _indexer(store, _Client()).index([good, bad]) + + # One bad pack must not poison the batch: the healthy pack still indexes, + # and the failure is attributed to the pack that caused it. + assert result.indexed_packs == 1 + assert result.failed_packs == 1 + assert result.failures[0].object_key == "packs/bad.dmi-pack" + + +def test_a_transient_outage_is_survivable_by_retrying_the_batch(tmp_path: Path): + inner = FilesystemPackStore(tmp_path, store_id="local") + ref = inner.put(_sealed(_record("capture-a")), "packs/a.dmi-pack") + backing = _Client() + client = FaultyClickHouseClient(backing, insert=fail_then_succeed(1)) + + with pytest.raises(FaultInjected): + _indexer(inner, client).index([ref]) + + # Nothing was committed, so the identical call now succeeds -- indexing is + # replayable rather than requiring manual repair. + result = _indexer(inner, client).index([ref]) + + assert result.indexed_packs == 1 + assert result.failed_packs == 0 diff --git a/tests/test_capture_golden_workload.py b/tests/test_capture_golden_workload.py new file mode 100644 index 000000000..bec39ac31 --- /dev/null +++ b/tests/test_capture_golden_workload.py @@ -0,0 +1,83 @@ +"""The capture-storage conformance manifest. + +Phase 6 compares golden workloads "by identity, logical bytes, checksums, +decoded tensors, and query results". `tests/tools/golden_workload.py` produces +exactly that as one JSON document, and the checked-in manifest beside this file +is what today's implementation produces. + +Because the Python implementation is the reference rather than the production +writer, this manifest is the contract: a native writer is conformant when the +same corpus yields the same document. These tests keep the reference honest -- +if any of it drifts, the diff names the capture and field that moved. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from tests.tools.golden_workload import ( + MANIFEST_VERSION, + build_manifest, + compare, +) + + +pytestmark = pytest.mark.cpu + + +GOLDEN = Path(__file__).parent / "data" / "capture_golden_manifest.json" + + +def test_the_implementation_still_matches_the_recorded_manifest(): + expected = json.loads(GOLDEN.read_text()) + + differences = compare(expected, build_manifest()) + + assert differences == [], ( + "capture storage no longer matches the conformance manifest.\n" + "If the change is intended, regenerate with:\n" + f" python tests/tools/golden_workload.py generate --out {GOLDEN}\n" + + "\n".join(f" {line}" for line in differences[:20]) + ) + + +def test_the_manifest_is_reproducible(): + # Two runs must agree, or the manifest cannot serve as a contract at all. + assert compare(build_manifest(), build_manifest()) == [] + + +def test_the_manifest_covers_every_supported_dtype(): + from dmi.storage.capture.model import _DTYPE_BYTES + + manifest = build_manifest() + + covered = {capture["dtype"] for capture in manifest["captures"]} + assert covered == set(_DTYPE_BYTES), "a dtype would ship unverified" + + +def test_the_manifest_records_what_phase_six_has_to_compare(): + manifest = build_manifest() + + assert manifest["manifest_version"] == MANIFEST_VERSION + # identity, logical bytes, checksums, decoded tensors, query results + assert manifest["pack"]["sha256"] + assert manifest["hydration"]["logical_bytes"] > 0 + for capture in manifest["captures"]: + assert capture["payload_sha256"] and capture["payload_crc32"] + assert capture["decoded_sha256"] + assert capture["summary"]["version"] == 1 + + +def test_a_changed_payload_is_caught_by_the_comparison(): + manifest = build_manifest() + tampered = json.loads(json.dumps(manifest)) + tampered["captures"][0]["decoded_sha256"] = "0" * 64 + + differences = compare(manifest, tampered) + + # The diff must name the field, not just report inequality. + assert len(differences) == 1 + assert "captures[0].decoded_sha256" in differences[0] diff --git a/tests/test_capture_summary.py b/tests/test_capture_summary.py index 875752f1c..9cea56b2c 100644 --- a/tests/test_capture_summary.py +++ b/tests/test_capture_summary.py @@ -639,3 +639,27 @@ def test_registry_refuses_duplicate_names(): def test_registry_refuses_malformed_identities(name: str, version: int): with pytest.raises(ExtensionError): ScalarMetric(name=name, version=version, compute=lambda x: 1.0) + + +def test_l2_norm_survives_large_magnitude_float64(tmp_path: Path): + # sqrt(sum(x**2)) overflows float64 well before the true norm does: squaring + # 1e200 needs 1e400. The scaled form keeps every squared term <= 1. + source = np.array([1e200, -1e200], dtype=np.float64) + _, descriptor = _descriptor_for(source, "float64", tmp_path) + + summary = summarize_tensor(descriptor, source.tobytes()) + + assert np.isfinite(summary.l2_norm) + assert summary.l2_norm == pytest.approx(np.sqrt(2.0) * 1e200, rel=1e-12) + assert summary.abs_max == pytest.approx(1e200) + + +def test_l2_norm_still_matches_the_plain_formula_at_normal_scale(tmp_path: Path): + rng = np.random.default_rng(seed=101) + source = (rng.random(256) * 10 - 5).astype(np.float64) + _, descriptor = _descriptor_for(source, "float64", tmp_path) + + summary = summarize_tensor(descriptor, source.tobytes()) + + # Scaling must not cost accuracy where the naive form was already fine. + assert summary.l2_norm == pytest.approx(float(np.sqrt((source**2).sum()))) diff --git a/tests/tools/__init__.py b/tests/tools/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/tools/golden_workload.py b/tests/tools/golden_workload.py new file mode 100644 index 000000000..d1f4319c2 --- /dev/null +++ b/tests/tools/golden_workload.py @@ -0,0 +1,280 @@ +#!/usr/bin/env python3 +"""Generate and verify a capture-storage conformance manifest. + +Phase 6 has to compare golden workloads "by identity, logical bytes, checksums, +decoded tensors, and query results". This produces exactly that, as one JSON +document, from a deterministic corpus. + +The Python implementation is the reference, not the production writer, so the +manifest is the contract rather than the code: a native writer is conformant if +and only if the same corpus produces the same manifest. Everything in it is +language-neutral -- byte counts, hex digests, and integers. No pickles, no +Python types, nothing that presumes the producer is Python. + + # record what today's implementation produces + python tests/tools/golden_workload.py generate --out golden.json + + # later, or from another implementation, check nothing moved + python tests/tools/golden_workload.py verify --manifest golden.json + +`verify` re-runs the corpus and diffs field by field, so a mismatch names the +capture and field that changed rather than just failing. +""" + +from __future__ import annotations + +import argparse +from hashlib import sha256 +import json +from pathlib import Path +import sys +import tempfile +from typing import Any +from uuid import UUID + +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT / "src") not in sys.path: + sys.path.insert(0, str(REPO_ROOT / "src")) + +from dmi.storage.capture import ( # noqa: E402 + CaptureMetadata, + CaptureRecord, + CaptureQuery, + CaptureReader, + CatalogIndexer, + FilesystemPackStore, + PackIndex, + PackWriter, + decode_tensor, + summarize_tensor, +) + +MANIFEST_VERSION = 1 +PACK_ID = UUID("018f0000-0000-7000-8000-000000000f01") + +# One capture per dtype the format accepts, so conformance covers every decode +# path rather than float32 alone. Payload bytes are generated from a fixed +# formula, not a random source, so any implementation can reproduce the corpus. +_DTYPES = ( + ("bool", 1), + ("uint8", 1), + ("int8", 1), + ("int16", 2), + ("float16", 2), + ("bfloat16", 2), + ("int32", 4), + ("float32", 4), + ("int64", 8), + ("float64", 8), +) +_ELEMENTS = 16 + + +def _payload(dtype: str, width: int, index: int) -> bytes: + """Deterministic bytes for one capture, defined by position not by RNG.""" + raw = bytearray() + for element in range(_ELEMENTS): + seed = (index * 131 + element * 17 + 7) & 0xFF + raw.extend(bytes([(seed + byte) & 0xFF for byte in range(width)])) + if dtype == "bool": + # A bool payload must contain only 0 or 1, or numpy comparisons and the + # format's own round trip stop agreeing. + return bytes(1 if value & 1 else 0 for value in raw) + return bytes(raw) + + +def _metadata(index: int, dtype: str) -> CaptureMetadata: + return CaptureMetadata( + capture_id=f"golden-{index:02d}", + tenant_id="tenant-golden", + experiment_id="experiment-golden", + run_id="run-golden", + session_id="session-golden", + request_id=f"request-{index}", + sequence_id=f"sequence-{index}", + model_id="model-golden", + model_revision="revision-1", + adapter_revision=None if index % 2 else f"adapter-{index}", + capture_policy_version="policy-v1", + hook_name="resid_pre" if index % 2 else "attn_out", + layer_number=index, + producer_rank=100 + index, + batch_position=900 + index, + step_number=100_000 + index, + token_start=200_000 + index, + token_end=300_000 + index, + dtype=dtype, + shape=(_ELEMENTS,), + captured_at_ns=1_700_000_000_000_000_000 + index, + ) + + +def _corpus() -> list[CaptureRecord]: + return [ + CaptureRecord( + metadata=_metadata(index, dtype), + payload=_payload(dtype, width, index), + ) + for index, (dtype, width) in enumerate(_DTYPES) + ] + + +class _Catalog: + """Stands in for the ClickHouse catalog so the manifest needs no server.""" + + def __init__(self, descriptors): + self._descriptors = tuple(descriptors) + + def search(self, query: CaptureQuery): + from dmi.storage.capture import CapturePage + + return CapturePage( + items=self._descriptors[: query.limit], next_cursor=None, watermark="1" + ) + + def get_by_ids(self, capture_ids, *, watermark): + wanted = set(capture_ids) + return tuple(i for i in self._descriptors if i.capture_id in wanted) + + +def build_manifest() -> dict[str, Any]: + """Run the corpus through the storage path and describe the result.""" + records = _corpus() + with tempfile.TemporaryDirectory() as workspace: + root = Path(workspace) + writer = PackWriter( + pack_id=PACK_ID, + created_at_ns=1_700_000_000_000_000_000, + max_pack_bytes=8 * 1024 * 1024, + ) + for record in records: + writer.append(record) + sealed = writer.seal() + + store = FilesystemPackStore(root, store_id="golden") + ref = store.put(sealed, "packs/golden.dmi-pack") + + # Read descriptors back through the footer path, exactly as the + # indexer does -- not from the writer's in-memory state. + descriptors = PackIndex.from_store(store, ref).descriptors() + reader = CaptureReader( + _Catalog(descriptors), {"golden": store}, max_coalesce_gap_bytes=0 + ) + selection = reader.select(CaptureQuery(limit=len(records))) + hydrated = reader.hydrate(selection, byte_limit=8 << 20) + estimate = reader.estimate(selection) + + captures = [] + for item in hydrated: + descriptor = item.descriptor + metadata, locator = descriptor.metadata, descriptor.locator + decoded = decode_tensor(descriptor, item.payload) + summary = summarize_tensor(descriptor, item.payload) + captures.append( + { + "capture_id": metadata.capture_id, + "dtype": metadata.dtype, + "shape": list(metadata.shape), + "logical_bytes": metadata.logical_bytes, + # identity + checksums + "payload_sha256": sha256(item.payload).hexdigest(), + "payload_crc32": locator.checksum, + # decoded tensor, hashed in a byte order the format fixes + "decoded_sha256": sha256(decoded.tobytes()).hexdigest(), + "decoded_dtype": str(decoded.dtype), + # placement inside the pack + "offset": locator.offset, + "stored_length": locator.stored_length, + "decoded_length": locator.decoded_length, + # the summary contract + "summary": { + "version": summary.summary_version, + "element_count": summary.element_count, + "finite_count": summary.finite_count, + "nan_count": summary.nan_count, + "inf_count": summary.inf_count, + "zero_fraction": round(summary.zero_fraction, 12), + "mean": round(summary.mean, 9), + "minimum": round(summary.minimum, 9), + "maximum": round(summary.maximum, 9), + "abs_max": round(summary.abs_max, 9), + "l2_norm": round(summary.l2_norm, 9), + }, + } + ) + + return { + "manifest_version": MANIFEST_VERSION, + "pack": { + "pack_id": ref.pack_id, + "object_bytes": ref.object_bytes, + "sha256": ref.checksum, + "record_count": ref.record_count, + }, + "hydration": { + "capture_count": estimate.capture_count, + "object_count": estimate.object_count, + "request_count": estimate.request_count, + "logical_bytes": estimate.logical_bytes, + "stored_bytes": estimate.stored_bytes, + "request_bytes": estimate.request_bytes, + }, + "captures": captures, + } + + +def _flatten(value: Any, prefix: str = "") -> dict[str, Any]: + flat: dict[str, Any] = {} + if isinstance(value, dict): + for key, item in value.items(): + flat.update(_flatten(item, f"{prefix}.{key}" if prefix else str(key))) + elif isinstance(value, list): + for index, item in enumerate(value): + flat.update(_flatten(item, f"{prefix}[{index}]")) + else: + flat[prefix] = value + return flat + + +def compare(expected: dict[str, Any], actual: dict[str, Any]) -> list[str]: + """Field-by-field differences, so a mismatch names what moved.""" + left, right = _flatten(expected), _flatten(actual) + differences = [] + for key in sorted(set(left) | set(right)): + if key not in left: + differences.append(f"{key}: unexpected -> {right[key]!r}") + elif key not in right: + differences.append(f"{key}: missing (expected {left[key]!r})") + elif left[key] != right[key]: + differences.append(f"{key}: expected {left[key]!r}, got {right[key]!r}") + return differences + + +def main(argv=None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + sub = parser.add_subparsers(dest="command", required=True) + generate = sub.add_parser("generate", help="write a manifest") + generate.add_argument("--out", type=Path, required=True) + verify = sub.add_parser("verify", help="check the corpus against a manifest") + verify.add_argument("--manifest", type=Path, required=True) + args = parser.parse_args(argv) + + manifest = build_manifest() + if args.command == "generate": + args.out.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n") + print(f"wrote {args.out} ({len(manifest['captures'])} captures)") + return 0 + + expected = json.loads(args.manifest.read_text()) + differences = compare(expected, manifest) + if not differences: + print(f"conformant: {len(manifest['captures'])} captures match {args.manifest}") + return 0 + print(f"NON-CONFORMANT: {len(differences)} difference(s)", file=sys.stderr) + for line in differences[:40]: + print(f" {line}", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) From 3ceb813797d0036f4398cb565581ed2dc8e290e4 Mon Sep 17 00:00:00 2001 From: Samfisheryu Date: Wed, 26 Aug 2026 00:28:46 -0400 Subject: [PATCH 16/20] test: drop superseded CUDA ring suite --- tests/ring/Makefile | 59 --- tests/ring/test_null_mode.cu | 288 ------------- tests/ring/test_producer.cu | 462 -------------------- tests/ring/test_ring_engine.cu | 375 ---------------- tests/ring/test_rings.cu | 758 --------------------------------- 5 files changed, 1942 deletions(-) delete mode 100644 tests/ring/Makefile delete mode 100644 tests/ring/test_null_mode.cu delete mode 100644 tests/ring/test_producer.cu delete mode 100644 tests/ring/test_ring_engine.cu delete mode 100644 tests/ring/test_rings.cu diff --git a/tests/ring/Makefile b/tests/ring/Makefile deleted file mode 100644 index ee86fd7d8..000000000 --- a/tests/ring/Makefile +++ /dev/null @@ -1,59 +0,0 @@ -NVCC ?= /usr/local/cuda/bin/nvcc -SM_ARCH ?= native -CXX_STD ?= c++17 -BUILD := build -INCLUDES := -I../.. -I../../native/csrc - -NVCC_FLAGS := \ - -std=$(CXX_STD) \ - -arch=$(SM_ARCH) \ - -O2 \ - -Xcompiler -Wall \ - -Xcompiler -Wextra \ - -Xcompiler -Wno-unused-parameter \ - -lineinfo \ - $(INCLUDES) - -PRODUCER_CU := ../../native/csrc/ring/producer.cu -DRAIN_SRC := ../../native/csrc/ring/drain_thread.cpp - -# test_ring_engine needs ATen headers + libtorch for DrainTask (contains at::Tensor). -# Use torch.utils.cmake_prefix_path to locate them. The Python E2E tests -# (tests/test_e2e_correctness_vs_hf.py) cover the full pipeline including p2p. -TORCH_INC := $(shell python -c "import torch; print(torch.utils.cmake_prefix_path + '/../../include')" 2>/dev/null) -TORCH_LIB := $(shell python -c "import torch; print(torch.utils.cmake_prefix_path + '/../../lib')" 2>/dev/null) - -TARGETS := $(BUILD)/test_rings $(BUILD)/test_producer $(BUILD)/test_ring_engine - -.PHONY: all test_rings test_producer test_ring_engine run clean - -all: $(TARGETS) - -$(BUILD): - mkdir -p $(BUILD) - -$(BUILD)/test_rings: test_rings.cu | $(BUILD) - $(NVCC) $(NVCC_FLAGS) $< -o $@ - -$(BUILD)/test_producer: test_producer.cu $(PRODUCER_CU) | $(BUILD) - $(NVCC) $(NVCC_FLAGS) $^ -o $@ - -$(BUILD)/test_ring_engine: test_ring_engine.cu $(PRODUCER_CU) $(DRAIN_SRC) | $(BUILD) - $(NVCC) $(NVCC_FLAGS) -I$(TORCH_INC) -I$(TORCH_INC)/torch/csrc/api/include -L$(TORCH_LIB) -Xlinker -rpath,$(TORCH_LIB) -ltorch -ltorch_cpu -lc10 -lpthread $^ -o $@ - -test_rings: $(BUILD)/test_rings - $(BUILD)/test_rings - -test_producer: $(BUILD)/test_producer - $(BUILD)/test_producer - -test_ring_engine: $(BUILD)/test_ring_engine - $(BUILD)/test_ring_engine - -run: all - $(BUILD)/test_rings - $(BUILD)/test_producer - $(BUILD)/test_ring_engine - -clean: - rm -rf $(BUILD) diff --git a/tests/ring/test_null_mode.cu b/tests/ring/test_null_mode.cu deleted file mode 100644 index 4adfb31cb..000000000 --- a/tests/ring/test_null_mode.cu +++ /dev/null @@ -1,288 +0,0 @@ -// tests/ring/test_null_mode.cu — Tests for producer null mode. -// -// Verifies that set_ring_null_mode(true) makes producer_kernel launch with -// the same parameters but skip all ring writes, while -// set_ring_null_mode(false) restores normal delivery. -// -// Build: see tests/ring/Makefile (target: test_null_mode) -// Run: ./build/test_null_mode - -#include "ring/ring_alloc.h" -#include "ring/producer.cuh" -#include "ring/pinned_pool.h" -#include "ring/drain_thread.h" -#include "ring/chunk_assembler.h" - -#include -#include -#include -#include -#include -#include -#include -#include - -using namespace ring; - -// --------------------------------------------------------------------------- -// Shared test harness (same as test_ring_engine.cu) -// --------------------------------------------------------------------------- -struct ResultCollector { - std::mutex mu; - std::condition_variable cv; - std::vector tensors; - - void push(AssembledTensor&& t) { - std::lock_guard lk(mu); - tensors.push_back(std::move(t)); - cv.notify_all(); - } - - bool wait_for(int n, int timeout_ms = 5000) { - std::unique_lock lk(mu); - return cv.wait_for(lk, std::chrono::milliseconds(timeout_ms), - [&] { return (int)tensors.size() >= n; }); - } - - int count() { - std::lock_guard lk(mu); - return (int)tensors.size(); - } -}; - -static constexpr uint64_t PINNED_BYTES = 8 * 256 * 1024; // 8 chunks worth - -static RingConfig make_cfg() { - RingConfig cfg{}; - cfg.task_ring_entries = 64; - cfg.payload_ring_bytes = 4 * 1024 * 1024; - cfg.chunk_bytes = 256 * 1024; - cfg.pinned_pool_bytes = PINNED_BYTES; - cfg.wait_policy = WaitPolicy::INFINITE; - cfg.drop_reporting = DropReporting::DROP_TASK; - return cfg; -} - -static uint8_t* upload(const std::vector& h, cudaStream_t s) { - uint8_t* d; - cudaMalloc(&d, h.size()); - cudaMemcpyAsync(d, h.data(), h.size(), cudaMemcpyHostToDevice, s); - return d; -} - -// --------------------------------------------------------------------------- -// test_null_mode_no_delivery -// With null mode ON, producer kernel launches but emits nothing to the ring. -// The drain thread should receive no tensors within the timeout window. -// --------------------------------------------------------------------------- -static void test_null_mode_no_delivery() { - printf(" test_null_mode_no_delivery ... "); - - set_ring_null_mode(true); - - RingConfig cfg = make_cfg(); - AllocatedRing ar(cfg); - ar.init(); - - ResultCollector rc; - PinnedPool pool; - pool.init(PINNED_BYTES); - - ChunkAssembler asm_(pool, [&](AssembledTensor&& t) { rc.push(std::move(t)); }); - DrainThread dt(ar.state(), pool, [&](DrainedChunk&& c) { asm_.push(std::move(c)); }); - dt.start(); - - const uint64_t data_bytes = 64 * 1024; - std::vector src(data_bytes, 0xAB); - cudaStream_t stream; - cudaStreamCreate(&stream); - uint8_t* d_src = upload(src, stream); - - launch_producer_with_notify(ar.state(), d_src, data_bytes, - /*logical_task_id=*/1, 0, 0, - DrainThread::hostfunc_cb, &dt, stream); - cudaStreamSynchronize(stream); - - // Give drain thread time to process any spurious notification. - bool got = rc.wait_for(1, /*timeout_ms=*/300); - assert(!got && "null mode: no tensor should arrive"); - assert(rc.count() == 0); - - // task_head must remain 0 — kernel wrote nothing. - assert(*ar.state().task_head == 0); - - dt.stop(); - cudaFree(d_src); - cudaStreamDestroy(stream); - set_ring_null_mode(false); // restore for subsequent tests - printf("PASS\n"); -} - -// --------------------------------------------------------------------------- -// test_null_mode_off_delivers -// With null mode OFF (default), producer works normally. -// --------------------------------------------------------------------------- -static void test_null_mode_off_delivers() { - printf(" test_null_mode_off_delivers ... "); - - set_ring_null_mode(false); - - RingConfig cfg = make_cfg(); - AllocatedRing ar(cfg); - ar.init(); - - ResultCollector rc; - PinnedPool pool; - pool.init(PINNED_BYTES); - - ChunkAssembler asm_(pool, [&](AssembledTensor&& t) { rc.push(std::move(t)); }); - DrainThread dt(ar.state(), pool, [&](DrainedChunk&& c) { asm_.push(std::move(c)); }); - dt.start(); - - const uint64_t data_bytes = 64 * 1024; - std::vector src(data_bytes); - std::iota(src.begin(), src.end(), uint8_t(7)); - cudaStream_t stream; - cudaStreamCreate(&stream); - uint8_t* d_src = upload(src, stream); - - launch_producer_with_notify(ar.state(), d_src, data_bytes, - /*logical_task_id=*/2, 0, 0, - DrainThread::hostfunc_cb, &dt, stream); - - assert(rc.wait_for(1)); - const auto& t = rc.tensors[0]; - assert(!t.is_drop); - assert(t.data.size() == data_bytes); - assert(memcmp(t.data.data(), src.data(), data_bytes) == 0); - - dt.stop(); - cudaStreamSynchronize(stream); - cudaFree(d_src); - cudaStreamDestroy(stream); - printf("PASS\n"); -} - -// --------------------------------------------------------------------------- -// test_toggle -// Submit N kernels in null mode (none arrive), then toggle off and submit -// M kernels in real mode (all M arrive, with correct data). -// Verifies that toggling leaves the ring in a consistent state. -// --------------------------------------------------------------------------- -static void test_toggle() { - printf(" test_toggle ... "); - - RingConfig cfg = make_cfg(); - AllocatedRing ar(cfg); - ar.init(); - - ResultCollector rc; - PinnedPool pool; - pool.init(PINNED_BYTES); - - ChunkAssembler asm_(pool, [&](AssembledTensor&& t) { rc.push(std::move(t)); }); - DrainThread dt(ar.state(), pool, [&](DrainedChunk&& c) { asm_.push(std::move(c)); }); - dt.start(); - - cudaStream_t stream; - cudaStreamCreate(&stream); - - const uint64_t data_bytes = 32 * 1024; - const int N_null = 5; - const int N_real = 4; - - // --- null mode phase --- - set_ring_null_mode(true); - std::vector null_ptrs(N_null); - for (int i = 0; i < N_null; ++i) { - std::vector src(data_bytes, uint8_t(0xCC)); - null_ptrs[i] = upload(src, stream); - launch_producer_with_notify(ar.state(), null_ptrs[i], data_bytes, - uint64_t(100 + i), 0, 0, - DrainThread::hostfunc_cb, &dt, stream); - } - cudaStreamSynchronize(stream); - // Brief wait: drain thread must NOT deliver anything. - bool got = rc.wait_for(1, /*timeout_ms=*/300); - assert(!got && "toggle: null-mode tensors must not arrive"); - assert(*ar.state().task_head == 0); - - // --- real mode phase --- - set_ring_null_mode(false); - std::vector> real_srcs(N_real); - std::vector real_ptrs(N_real); - for (int i = 0; i < N_real; ++i) { - real_srcs[i].resize(data_bytes); - for (uint64_t j = 0; j < data_bytes; ++j) real_srcs[i][j] = uint8_t(i * 13 + j); - real_ptrs[i] = upload(real_srcs[i], stream); - launch_producer_with_notify(ar.state(), real_ptrs[i], data_bytes, - uint64_t(200 + i), 0, 0, - DrainThread::hostfunc_cb, &dt, stream); - } - - assert(rc.wait_for(N_real)); - assert(rc.count() == N_real); - for (int i = 0; i < N_real; ++i) { - const auto& t = rc.tensors[i]; - assert(!t.is_drop); - assert(t.data.size() == data_bytes); - int idx = static_cast(t.logical_task_id) - 200; - assert(idx >= 0 && idx < N_real); - assert(memcmp(t.data.data(), real_srcs[idx].data(), data_bytes) == 0); - } - - dt.stop(); - cudaStreamSynchronize(stream); - for (auto p : null_ptrs) cudaFree(p); - for (auto p : real_ptrs) cudaFree(p); - cudaStreamDestroy(stream); - printf("PASS\n"); -} - -// --------------------------------------------------------------------------- -// test_null_multiple_launches -// Launch many kernels in null mode; ring state (task_head, payload_head) -// must be completely unchanged — no partial writes. -// --------------------------------------------------------------------------- -static void test_null_ring_state_unchanged() { - printf(" test_null_ring_state_unchanged ... "); - - set_ring_null_mode(true); - - RingConfig cfg = make_cfg(); - AllocatedRing ar(cfg); - ar.init(); - - cudaStream_t stream; - cudaStreamCreate(&stream); - - const int N = 10; - const uint64_t data_bytes = 128 * 1024; - std::vector ptrs(N); - for (int i = 0; i < N; ++i) { - std::vector src(data_bytes, uint8_t(i)); - ptrs[i] = upload(src, stream); - // launch_producer (no hostfunc needed; we just check ring state) - launch_producer(ar.state(), ptrs[i], data_bytes, uint64_t(i), 0, 0, stream); - } - cudaStreamSynchronize(stream); - - assert(*ar.state().task_head == 0 && "null mode: task_head must be 0"); - assert(*ar.state().payload_head == 0 && "null mode: payload_head must be 0"); - - for (auto p : ptrs) cudaFree(p); - cudaStreamDestroy(stream); - set_ring_null_mode(false); - printf("PASS\n"); -} - -// --------------------------------------------------------------------------- -int main() { - printf("test_null_mode\n"); - test_null_mode_no_delivery(); - test_null_mode_off_delivers(); - test_toggle(); - test_null_ring_state_unchanged(); - printf("All tests passed.\n"); - return 0; -} diff --git a/tests/ring/test_producer.cu b/tests/ring/test_producer.cu deleted file mode 100644 index 37e47c79e..000000000 --- a/tests/ring/test_producer.cu +++ /dev/null @@ -1,462 +0,0 @@ -// tests/ring/test_producer.cu — M2 producer kernel correctness tests. -// -// Build: make -C tests/ring test_producer -// Run: ./tests/ring/build/test_producer -// -// All GPU tests run sequentially on a single stream (cudaStreamDefault) because -// this GPU serialises kernels from the same process even across streams. - -#include "ring/ring_alloc.h" -#include "ring/producer.cuh" - -#include -#include -#include -#include -#include - -// --------------------------------------------------------------------------- -// Assertion helpers -// --------------------------------------------------------------------------- -static int g_total = 0; -static int g_fail = 0; - -#define ASSERT(cond) do { \ - ++g_total; \ - if (!(cond)) { \ - ++g_fail; \ - fprintf(stderr, " FAIL %s:%d %s\n", \ - __FILE__, __LINE__, #cond); \ - } \ -} while (0) - -#define CUDA_CHECK(e) do { \ - cudaError_t _e = (e); \ - if (_e != cudaSuccess) { \ - fprintf(stderr, "CUDA error %s:%d: %s\n", \ - __FILE__, __LINE__, cudaGetErrorString(_e)); \ - std::exit(1); \ - } \ -} while (0) - -static void banner(const char* name) { - printf("[ TEST ] %s\n", name); - fflush(stdout); -} - -// Read task entries from device into a host-side vector. -static std::vector read_entries(ring::TaskEntry* d_entries, - uint64_t n_slots) -{ - std::vector v(n_slots); - CUDA_CHECK(cudaMemcpy(v.data(), d_entries, - n_slots * sizeof(ring::TaskEntry), - cudaMemcpyDeviceToHost)); - return v; -} - -// Read payload buffer (or subset) from device. -static std::vector read_payload(uint8_t* d_buf, uint64_t nbytes) -{ - std::vector v(nbytes); - CUDA_CHECK(cudaMemcpy(v.data(), d_buf, nbytes, cudaMemcpyDeviceToHost)); - return v; -} - -// Allocate a device buffer and fill it with a simple byte pattern. -static uint8_t* make_src(uint64_t nbytes) -{ - uint8_t* d; - CUDA_CHECK(cudaMalloc(&d, nbytes)); - std::vector h(nbytes); - for (uint64_t i = 0; i < nbytes; ++i) h[i] = static_cast(i & 0xFF); - CUDA_CHECK(cudaMemcpy(d, h.data(), nbytes, cudaMemcpyHostToDevice)); - return d; -} - -// --------------------------------------------------------------------------- -// Test 1 — host-only: RingConfig defaults -// --------------------------------------------------------------------------- -static void test_config_defaults() -{ - banner("config_defaults"); - ring::RingConfig cfg; - ASSERT(cfg.task_ring_entries == 1024); - ASSERT(cfg.payload_ring_bytes == 256ULL * 1024 * 1024); - ASSERT(cfg.chunk_bytes == 64ULL * 1024 * 1024); - ASSERT(cfg.wait_policy == ring::WaitPolicy::INFINITE); - ASSERT(cfg.drop_reporting == ring::DropReporting::DROP_TASK); -} - -// --------------------------------------------------------------------------- -// Test 2 — single-chunk: tensor fits in one chunk -// --------------------------------------------------------------------------- -static void test_single_chunk() -{ - banner("single_chunk"); - - ring::RingConfig cfg; - cfg.task_ring_entries = 16; - cfg.payload_ring_bytes = 16 * 1024; // 16 KiB - cfg.chunk_bytes = 4 * 1024; // 4 KiB — tensor fits in one chunk - - ring::AllocatedRing ar(cfg); - ar.init(); - - const uint64_t src_bytes = 256; - uint8_t* d_src = make_src(src_bytes); - - launch_producer(ar.state(), d_src, src_bytes, - /*logical_task_id=*/42, - /*hook_type=*/1, /*hook_id=*/7); - CUDA_CHECK(cudaDeviceSynchronize()); - - // task_head should now be 1 - ASSERT(*ar.state().task_head == 1); - - auto entries = read_entries(ar.state().task_entries, cfg.task_ring_entries); - const auto& e = entries[0]; - - ASSERT(e.ready_seq == 0); - ASSERT(e.seq_no == 0); - ASSERT(e.logical_task_id == 42); - ASSERT(e.chunk_idx == 0); - ASSERT(e.hook_type == 1); - ASSERT(e.hook_id == 7); - ASSERT(e.tensor_total_bytes == src_bytes); - ASSERT(e.chunk_offset_bytes == 0); - ASSERT((e.flags & ring::TASK_FLAG_IS_FIRST) != 0); - ASSERT((e.flags & ring::TASK_FLAG_IS_LAST) != 0); - ASSERT((e.flags & ring::TASK_FLAG_IS_DROP) == 0); - ASSERT(e.payload_len1 == src_bytes); - ASSERT(e.payload_len2 == 0); - ASSERT(e.payload_off1 == 0); // payload_head started at 0 - - cudaFree(d_src); -} - -// --------------------------------------------------------------------------- -// Test 3 — multi-chunk: tensor spans multiple chunks -// --------------------------------------------------------------------------- -static void test_multi_chunk() -{ - banner("multi_chunk"); - - ring::RingConfig cfg; - cfg.task_ring_entries = 16; - cfg.payload_ring_bytes = 64 * 1024; // 64 KiB - cfg.chunk_bytes = 1024; // 1 KiB per chunk - - ring::AllocatedRing ar(cfg); - ar.init(); - - // 3.5 × chunk_bytes → 4 chunks (3×1024 + 512) - const uint64_t src_bytes = 3 * 1024 + 512; - uint8_t* d_src = make_src(src_bytes); - - launch_producer(ar.state(), d_src, src_bytes, 99, 0, 5); - CUDA_CHECK(cudaDeviceSynchronize()); - - ASSERT(*ar.state().task_head == 4); - - auto entries = read_entries(ar.state().task_entries, cfg.task_ring_entries); - - // chunk 0 — IS_FIRST - ASSERT((entries[0].flags & ring::TASK_FLAG_IS_FIRST) != 0); - ASSERT((entries[0].flags & ring::TASK_FLAG_IS_LAST) == 0); - ASSERT(entries[0].chunk_idx == 0); - ASSERT(entries[0].chunk_offset_bytes == 0); - ASSERT(entries[0].payload_len1 == 1024); - - // chunk 1 — middle - ASSERT((entries[1].flags & ring::TASK_FLAG_IS_FIRST) == 0); - ASSERT((entries[1].flags & ring::TASK_FLAG_IS_LAST) == 0); - ASSERT(entries[1].chunk_idx == 1); - ASSERT(entries[1].chunk_offset_bytes == 1024); - - // chunk 2 — middle - ASSERT(entries[2].chunk_idx == 2); - ASSERT(entries[2].chunk_offset_bytes == 2 * 1024); - - // chunk 3 — IS_LAST, smaller - ASSERT((entries[3].flags & ring::TASK_FLAG_IS_FIRST) == 0); - ASSERT((entries[3].flags & ring::TASK_FLAG_IS_LAST) != 0); - ASSERT(entries[3].chunk_idx == 3); - ASSERT(entries[3].chunk_offset_bytes == 3 * 1024); - ASSERT(entries[3].payload_len1 == 512); - - // All entries share same logical_task_id and tensor_total_bytes - for (int i = 0; i < 4; ++i) { - ASSERT(entries[i].logical_task_id == 99); - ASSERT(entries[i].tensor_total_bytes == src_bytes); - ASSERT(entries[i].hook_id == 5); - } - - cudaFree(d_src); -} - -// --------------------------------------------------------------------------- -// Test 4 — payload wrap: chunk crosses end of ring buffer → two-span entry -// --------------------------------------------------------------------------- -static void test_wrap() -{ - banner("wrap_two_span"); - - ring::RingConfig cfg; - cfg.task_ring_entries = 8; - cfg.payload_ring_bytes = 4096; - cfg.chunk_bytes = 512; - - ring::AllocatedRing ar(cfg); - ar.init(); - - // Simulate a ring state where physical head is near the end. - // head and tail must be PAYLOAD_ALIGN-aligned (16 bytes) since - // vectorized uint4 D2D copies require aligned offsets. - // head=3904, tail=3584: free = 4096 - (3904-3584) = 3776 >= 512. - // Physical head position: 3904 % 4096 = 3904. - // bytes_to_end = 4096 - 3904 = 192 → two spans: len1=192, len2=320. - const uint64_t head = 3904; - const uint64_t tail = 3584; - const uint64_t src_bytes = 512; - *ar.state().payload_head = head; - *ar.state().payload_tail = tail; - - uint8_t* d_src = make_src(src_bytes); - - launch_producer(ar.state(), d_src, src_bytes, 1, 0, 0); - CUDA_CHECK(cudaDeviceSynchronize()); - - auto entries = read_entries(ar.state().task_entries, cfg.task_ring_entries); - const auto& e = entries[0]; - - uint64_t bytes_to_end = 4096 - (head % 4096); // 192 - ASSERT(e.payload_len1 == bytes_to_end); - ASSERT(e.payload_len2 == src_bytes - bytes_to_end); - ASSERT(e.payload_off1 == head % 4096); - ASSERT(e.payload_off2 == 0); - ASSERT(e.payload_len1 + e.payload_len2 == src_bytes); - - cudaFree(d_src); -} - -// --------------------------------------------------------------------------- -// Test 5 — data correctness: verify payload_buf contains the right bytes -// --------------------------------------------------------------------------- -static void test_data_correctness() -{ - banner("data_correctness"); - - ring::RingConfig cfg; - cfg.task_ring_entries = 8; - cfg.payload_ring_bytes = 8 * 1024; - cfg.chunk_bytes = 4 * 1024; - - ring::AllocatedRing ar(cfg); - ar.init(); - - const uint64_t src_bytes = 512; - uint8_t* d_src = make_src(src_bytes); // pattern: byte[i] = i % 256 - - launch_producer(ar.state(), d_src, src_bytes, 0, 0, 0); - CUDA_CHECK(cudaDeviceSynchronize()); - - auto entries = read_entries(ar.state().task_entries, cfg.task_ring_entries); - const auto& e = entries[0]; - ASSERT(e.payload_len1 == src_bytes && e.payload_len2 == 0); - - auto payload = read_payload(ar.state().payload_buf, cfg.payload_ring_bytes); - bool ok = true; - for (uint64_t i = 0; i < src_bytes; ++i) { - if (payload[e.payload_off1 + i] != static_cast(i & 0xFF)) { - ok = false; - fprintf(stderr, " byte[%llu] expected %u got %u\n", - (unsigned long long)i, - (unsigned)static_cast(i & 0xFF), - (unsigned)payload[e.payload_off1 + i]); - break; - } - } - ASSERT(ok); - - cudaFree(d_src); -} - -// --------------------------------------------------------------------------- -// Test 6 — data correctness across wrap: two-span payload bytes are correct -// --------------------------------------------------------------------------- -static void test_data_correctness_wrap() -{ - banner("data_correctness_wrap"); - - ring::RingConfig cfg; - cfg.task_ring_entries = 8; - cfg.payload_ring_bytes = 4096; - cfg.chunk_bytes = 512; - - ring::AllocatedRing ar(cfg); - ar.init(); - - // Same setup as test_wrap: head=3904, tail=3584 → two-span (192 + 320). - // head must be PAYLOAD_ALIGN-aligned (16 bytes) for vectorized uint4 copies. - const uint64_t src_bytes = 512; - *ar.state().payload_head = 3904; - *ar.state().payload_tail = 3584; - - uint8_t* d_src = make_src(src_bytes); - - launch_producer(ar.state(), d_src, src_bytes, 0, 0, 0); - CUDA_CHECK(cudaDeviceSynchronize()); - - auto entries = read_entries(ar.state().task_entries, cfg.task_ring_entries); - const auto& e = entries[0]; - - auto payload = read_payload(ar.state().payload_buf, cfg.payload_ring_bytes); - - bool ok = true; - // Span 1 bytes - for (uint64_t i = 0; i < e.payload_len1 && ok; ++i) { - if (payload[e.payload_off1 + i] != static_cast(i & 0xFF)) { - ok = false; - fprintf(stderr, " span1 byte[%llu] mismatch\n", (unsigned long long)i); - } - } - // Span 2 bytes - for (uint64_t i = 0; i < e.payload_len2 && ok; ++i) { - uint64_t src_i = e.payload_len1 + i; - if (payload[e.payload_off2 + i] != static_cast(src_i & 0xFF)) { - ok = false; - fprintf(stderr, " span2 byte[%llu] mismatch\n", (unsigned long long)i); - } - } - ASSERT(ok); - - cudaFree(d_src); -} - -// --------------------------------------------------------------------------- -// Test 7 — timeout-drop: frozen heartbeat + nearly-full payload → DROP marker -// --------------------------------------------------------------------------- -static void test_timeout_drop() -{ - banner("timeout_drop"); - - ring::RingConfig cfg; - cfg.task_ring_entries = 16; - cfg.payload_ring_bytes = 1024; - cfg.chunk_bytes = 512; - cfg.wait_policy = ring::WaitPolicy::TIMEOUT_DROP; - cfg.drop_reporting = ring::DropReporting::DROP_TASK; - cfg.no_progress_timeout_cycles = 200000; // ~0.1 ms at 2 GHz - - ring::AllocatedRing ar(cfg); - ar.init(); - - // Simulate a nearly-full payload ring: only 16 bytes free. - // The producer needs 512 bytes but only 16 are available. - // consumer_heartbeat stays 0 → no progress → timeout. - *ar.state().payload_head = cfg.payload_ring_bytes - 16; - *ar.state().payload_tail = 0; - // task ring is empty (head=tail=0), so 1 task slot is available for DROP. - - const uint64_t src_bytes = 512; - uint8_t* d_src = make_src(src_bytes); - - launch_producer(ar.state(), d_src, src_bytes, 77, 2, 9); - CUDA_CHECK(cudaDeviceSynchronize()); - - // Producer should have published exactly 1 DROP entry. - ASSERT(*ar.state().task_head == 1); - - auto entries = read_entries(ar.state().task_entries, cfg.task_ring_entries); - const auto& e = entries[0]; - - ASSERT((e.flags & ring::TASK_FLAG_IS_DROP) != 0); - ASSERT((e.flags & ring::TASK_FLAG_IS_LAST) != 0); - ASSERT((e.flags & ring::TASK_FLAG_IS_FIRST) != 0); // chunk 0 was first - ASSERT(e.reason == ring::DROP_REASON_TIMEOUT_NO_PROGRESS); - ASSERT(e.payload_len1 == 0); - ASSERT(e.payload_len2 == 0); - ASSERT(e.logical_task_id == 77); - ASSERT(e.hook_type == 2); - ASSERT(e.hook_id == 9); - - cudaFree(d_src); -} - -// --------------------------------------------------------------------------- -// Test 8 — backpressure release: producer spins, consumer advances tail -// --------------------------------------------------------------------------- -// Simulated by running two sequential kernels on the same stream: -// kernel A fills ring, then we manually advance tails, then kernel B runs. -static void test_backpressure_release() -{ - banner("backpressure_release"); - - ring::RingConfig cfg; - cfg.task_ring_entries = 4; // small ring - cfg.payload_ring_bytes = 4 * 1024; // 4 KiB - cfg.chunk_bytes = 1 * 1024; // 1 KiB - - ring::AllocatedRing ar(cfg); - ar.init(); - - // Fill the ring: 4 chunks × 1KiB = 4KiB (exactly fills payload ring). - const uint64_t src_bytes_a = 4 * 1024; - uint8_t* d_src_a = make_src(src_bytes_a); - launch_producer(ar.state(), d_src_a, src_bytes_a, 10, 0, 1); - CUDA_CHECK(cudaDeviceSynchronize()); - ASSERT(*ar.state().task_head == 4); - - // Now simulate consumer draining slot 0: advance task_tail and payload_tail. - // This frees 1 task slot + 1 KiB of payload. - *ar.state().task_tail = 1; - *ar.state().payload_tail = 1024; - ++(*ar.state().consumer_heartbeat); - - // Producer should now be able to fit 1 more chunk. - const uint64_t src_bytes_b = 512; - uint8_t* d_src_b = make_src(src_bytes_b); - launch_producer(ar.state(), d_src_b, src_bytes_b, 20, 0, 2); - CUDA_CHECK(cudaDeviceSynchronize()); - - // task_head should now be 5 (4 from A + 1 from B). - ASSERT(*ar.state().task_head == 5); - - auto entries = read_entries(ar.state().task_entries, cfg.task_ring_entries); - // Slot 4 % 4 = 0 (reused after tail advanced). - const auto& e = entries[4 % 4]; - ASSERT(e.logical_task_id == 20); - ASSERT(e.tensor_total_bytes == src_bytes_b); - ASSERT(e.seq_no == 4); - - cudaFree(d_src_a); - cudaFree(d_src_b); -} - -// --------------------------------------------------------------------------- -// main -// --------------------------------------------------------------------------- -int main() -{ - setbuf(stdout, nullptr); - - printf("=== test_producer ===\n"); - - test_config_defaults(); - test_single_chunk(); - test_multi_chunk(); - test_wrap(); - test_data_correctness(); - test_data_correctness_wrap(); - test_timeout_drop(); - test_backpressure_release(); - - printf("\n%d / %d assertions passed\n", g_total - g_fail, g_total); - if (g_fail > 0) { - fprintf(stderr, "%d FAILURES\n", g_fail); - return 1; - } - printf("ALL PASS\n"); - return 0; -} diff --git a/tests/ring/test_ring_engine.cu b/tests/ring/test_ring_engine.cu deleted file mode 100644 index cd6e71171..000000000 --- a/tests/ring/test_ring_engine.cu +++ /dev/null @@ -1,375 +0,0 @@ -// tests/ring/test_ring_engine.cu — End-to-end tests for the ring drain pipeline. -// -// Tests the full path: -// GPU producer kernel → drain thread (batch D2H into pinned staging) -// → task queue → verify DrainTask contents -// -// NOTE: The new drain pipeline (batch D2H + pinned staging + p2p thread) -// requires ATen linkage for DrainTask (contains at::Tensor). The p2p thread -// and full tensor assembly are tested via the Python E2E tests -// (tests/test_e2e_correctness_vs_hf.py) which exercise the entire pipeline -// including metadata FIFO, reshape, slicing, and ClickHouse submission. -// -// These C++ tests verify the drain-side pipeline: producer → drain thread → -// task queue, confirming that staged data matches the original GPU tensor. -// -// Build: make test_ring_engine (requires PyTorch for ATen headers) - -#include "ring/ring_alloc.h" -#include "ring/producer.cuh" -#include "ring/drain_thread.h" -#include "ring/pinned_staging.h" - -#include -#include -#include -#include -#include - -using namespace ring; - -// --------------------------------------------------------------------------- -// Helpers -// --------------------------------------------------------------------------- - -// Allocate and upload a host vector to GPU device memory. -static uint8_t* upload(const std::vector& h, cudaStream_t s = 0) { - uint8_t* d; - cudaMalloc(&d, h.size()); - cudaMemcpyAsync(d, h.data(), h.size(), cudaMemcpyHostToDevice, s); - return d; -} - -// Verify staging data for a single-chunk tensor matches source. -static void verify_staging_data(const DrainTask& task, - const std::vector& src) { - uint64_t data_len = task.data_len1 + task.data_len2; - assert(data_len == src.size()); - // Copy from staging (may wrap) - std::vector got(data_len); - if (task.data_len1 > 0) - memcpy(got.data(), task.data_ptr1, task.data_len1); - if (task.data_len2 > 0) - memcpy(got.data() + task.data_len1, task.data_ptr2, task.data_len2); - assert(memcmp(got.data(), src.data(), src.size()) == 0); -} - -// --------------------------------------------------------------------------- -// test_single_chunk -// One tensor that fits in a single chunk. -// Verifies drain task fields and staged data match the original. -// --------------------------------------------------------------------------- -static void test_single_chunk() { - printf(" test_single_chunk ... "); - - const uint64_t data_bytes = 128 * 1024; // 128 KB < chunk_bytes (256 KB) - RingConfig cfg{}; - cfg.task_ring_entries = 64; - cfg.payload_ring_bytes = 4 * 1024 * 1024; - cfg.chunk_bytes = 256 * 1024; - cfg.wait_policy = WaitPolicy::INFINITE; - cfg.drop_reporting = DropReporting::DROP_TASK; - - AllocatedRing ar(cfg); - ar.init(); - - PinnedStaging staging; - staging.init(cfg.payload_ring_bytes); - - DrainThread dt(ar.state(), staging, cfg); - dt.start(); - - std::vector src(data_bytes); - std::iota(src.begin(), src.end(), uint8_t(0)); - cudaStream_t stream; - cudaStreamCreate(&stream); - uint8_t* d_src = upload(src, stream); - - launch_producer_with_notify(ar.state(), d_src, data_bytes, - /*logical_task_id=*/1, /*hook_type=*/0, /*hook_id=*/0, - DrainThread::hostfunc_cb, &dt, stream); - - // Stop drain to trigger final flush - cudaStreamSynchronize(stream); - dt.stop(); - - // Now collect tasks (drain has pushed them before joining) - // Signal p2p stop so wait_for_tasks returns - dt.signal_p2p_stop(); - uint64_t n = dt.wait_for_tasks(); - assert(n > 0); - std::vector tasks; - dt.pop_tasks(n, tasks); - assert(tasks.size() == 1); - - const auto& t = tasks[0]; - assert(t.logical_task_id == 1); - assert(t.tensor_total_bytes == data_bytes); - assert(t.flags & TASK_FLAG_IS_FIRST); - assert(t.flags & TASK_FLAG_IS_LAST); - verify_staging_data(t, src); - - cudaFree(d_src); - cudaStreamDestroy(stream); - printf("PASS\n"); -} - -// --------------------------------------------------------------------------- -// test_multi_chunk -// One large tensor that spans multiple chunks. -// Verifies all chunks arrive with correct metadata. -// --------------------------------------------------------------------------- -static void test_multi_chunk() { - printf(" test_multi_chunk ... "); - - const uint64_t chunk_bytes = 64 * 1024; // 64 KB - const uint64_t data_bytes = 3 * chunk_bytes + 48; // 3 full + 1 partial (aligned to 16) - RingConfig cfg{}; - cfg.task_ring_entries = 64; - cfg.payload_ring_bytes = 2 * 1024 * 1024; - cfg.chunk_bytes = chunk_bytes; - cfg.wait_policy = WaitPolicy::INFINITE; - cfg.drop_reporting = DropReporting::DROP_TASK; - - AllocatedRing ar(cfg); - ar.init(); - - PinnedStaging staging; - staging.init(cfg.payload_ring_bytes); - - DrainThread dt(ar.state(), staging, cfg); - dt.start(); - - std::vector src(data_bytes); - for (uint64_t i = 0; i < data_bytes; ++i) src[i] = uint8_t(i * 7 + 3); - - cudaStream_t stream; - cudaStreamCreate(&stream); - uint8_t* d_src = upload(src, stream); - - launch_producer_with_notify(ar.state(), d_src, data_bytes, - /*logical_task_id=*/2, 0, 0, - DrainThread::hostfunc_cb, &dt, stream); - - cudaStreamSynchronize(stream); - dt.stop(); - dt.signal_p2p_stop(); - - uint64_t n = dt.wait_for_tasks(); - assert(n > 0); - std::vector tasks; - dt.pop_tasks(n, tasks); - assert(tasks.size() == 4); // 3 full + 1 partial - - // Reassemble from staging - std::vector assembled(data_bytes); - for (const auto& t : tasks) { - if (t.data_len1 > 0) - memcpy(assembled.data() + t.chunk_offset_bytes, t.data_ptr1, t.data_len1); - if (t.data_len2 > 0) - memcpy(assembled.data() + t.chunk_offset_bytes + t.data_len1, - t.data_ptr2, t.data_len2); - } - assert(memcmp(assembled.data(), src.data(), data_bytes) == 0); - - assert(tasks[0].flags & TASK_FLAG_IS_FIRST); - assert(tasks[3].flags & TASK_FLAG_IS_LAST); - for (const auto& t : tasks) { - assert(t.tensor_total_bytes == data_bytes); - assert(t.logical_task_id == 2); - } - - cudaFree(d_src); - cudaStreamDestroy(stream); - printf("PASS\n"); -} - -// --------------------------------------------------------------------------- -// test_drop -// Fill the ring so the producer times out and emits a DROP marker. -// --------------------------------------------------------------------------- -static void test_drop() { - printf(" test_drop ... "); - - RingConfig cfg{}; - cfg.task_ring_entries = 4; - cfg.payload_ring_bytes = 128 * 1024; - cfg.chunk_bytes = 64 * 1024; - cfg.wait_policy = WaitPolicy::TIMEOUT_DROP; - cfg.no_progress_timeout_cycles = 1000000ULL; // ~0.4 ms at 2.5 GHz - cfg.drop_reporting = DropReporting::DROP_TASK; - - AllocatedRing ar(cfg); - ar.init(); - - // Completely fill the payload ring so producer immediately drops. - *ar.state().payload_head = cfg.payload_ring_bytes; - *ar.state().payload_tail = 0; // free = 0 - - PinnedStaging staging; - staging.init(cfg.payload_ring_bytes); - - DrainThread dt(ar.state(), staging, cfg); - dt.start(); - - const uint64_t data_bytes = 32 * 1024; - std::vector src(data_bytes, 0xBE); - cudaStream_t stream; - cudaStreamCreate(&stream); - uint8_t* d_src = upload(src, stream); - - launch_producer_with_notify(ar.state(), d_src, data_bytes, - /*logical_task_id=*/4, 0, 0, - DrainThread::hostfunc_cb, &dt, stream); - - cudaStreamSynchronize(stream); - dt.stop(); - dt.signal_p2p_stop(); - - uint64_t n = dt.wait_for_tasks(); - assert(n > 0); - std::vector tasks; - dt.pop_tasks(n, tasks); - assert(tasks.size() == 1); - assert(tasks[0].flags & TASK_FLAG_IS_DROP); - - cudaFree(d_src); - cudaStreamDestroy(stream); - printf("PASS\n"); -} - -// --------------------------------------------------------------------------- -// test_multiple_tensors -// Submit N independent tensors back-to-back. -// Verifies drain handles queued notifications and data correctness. -// --------------------------------------------------------------------------- -static void test_multiple_tensors() { - printf(" test_multiple_tensors ... "); - - const int N = 8; - const uint64_t data_bytes = 32 * 1024; // 32 KB each - RingConfig cfg{}; - cfg.task_ring_entries = 64; - cfg.payload_ring_bytes = 4 * 1024 * 1024; - cfg.chunk_bytes = 64 * 1024; - cfg.wait_policy = WaitPolicy::INFINITE; - cfg.drop_reporting = DropReporting::DROP_TASK; - - AllocatedRing ar(cfg); - ar.init(); - - PinnedStaging staging; - staging.init(cfg.payload_ring_bytes); - - DrainThread dt(ar.state(), staging, cfg); - dt.start(); - - cudaStream_t stream; - cudaStreamCreate(&stream); - - std::vector d_srcs(N); - std::vector> srcs(N); - for (int i = 0; i < N; ++i) { - srcs[i].resize(data_bytes); - for (uint64_t j = 0; j < data_bytes; ++j) srcs[i][j] = uint8_t(i + j); - d_srcs[i] = upload(srcs[i], stream); - launch_producer_with_notify(ar.state(), d_srcs[i], data_bytes, - /*logical_task_id=*/uint64_t(10 + i), 0, 0, - DrainThread::hostfunc_cb, &dt, stream); - } - - cudaStreamSynchronize(stream); - dt.stop(); - dt.signal_p2p_stop(); - - // Collect all tasks - std::vector all_tasks; - while (true) { - uint64_t n = dt.wait_for_tasks(); - if (n == 0) break; - std::vector batch; - dt.pop_tasks(n, batch); - for (auto& t : batch) all_tasks.push_back(std::move(t)); - } - assert(static_cast(all_tasks.size()) == N); - - // Each should be a single-chunk tensor (32KB < 64KB chunk) - for (int i = 0; i < N; ++i) { - const auto& t = all_tasks[i]; - assert(!(t.flags & TASK_FLAG_IS_DROP)); - assert(t.flags & TASK_FLAG_IS_FIRST); - assert(t.flags & TASK_FLAG_IS_LAST); - assert(t.tensor_total_bytes == data_bytes); - - int idx = static_cast(t.logical_task_id) - 10; - assert(idx >= 0 && idx < N); - verify_staging_data(t, srcs[idx]); - } - - for (auto d : d_srcs) cudaFree(d); - cudaStreamDestroy(stream); - printf("PASS\n"); -} - -// --------------------------------------------------------------------------- -// test_zero_byte_tensor -// A tensor with src_bytes=0 produces one task entry with no payload. -// --------------------------------------------------------------------------- -static void test_zero_byte_tensor() { - printf(" test_zero_byte_tensor ... "); - - RingConfig cfg{}; - cfg.task_ring_entries = 64; - cfg.payload_ring_bytes = 4 * 1024 * 1024; - cfg.chunk_bytes = 256 * 1024; - cfg.wait_policy = WaitPolicy::INFINITE; - cfg.drop_reporting = DropReporting::DROP_TASK; - - AllocatedRing ar(cfg); - ar.init(); - - PinnedStaging staging; - staging.init(cfg.payload_ring_bytes); - - DrainThread dt(ar.state(), staging, cfg); - dt.start(); - - cudaStream_t stream; - cudaStreamCreate(&stream); - - launch_producer_with_notify(ar.state(), nullptr, 0, - /*logical_task_id=*/99, 0, 0, - DrainThread::hostfunc_cb, &dt, stream); - - cudaStreamSynchronize(stream); - dt.stop(); - dt.signal_p2p_stop(); - - uint64_t n = dt.wait_for_tasks(); - assert(n > 0); - std::vector tasks; - dt.pop_tasks(n, tasks); - assert(tasks.size() == 1); - - const auto& t = tasks[0]; - assert(!(t.flags & TASK_FLAG_IS_DROP)); - assert(t.logical_task_id == 99); - assert(t.tensor_total_bytes == 0); - assert(t.data_len1 == 0); - assert(t.data_len2 == 0); - - cudaStreamDestroy(stream); - printf("PASS\n"); -} - -// --------------------------------------------------------------------------- -int main() { - printf("test_ring_engine (batch drain pipeline)\n"); - test_single_chunk(); - test_multi_chunk(); - test_drop(); - test_multiple_tensors(); - test_zero_byte_tensor(); - printf("All 5 tests passed.\n"); - return 0; -} diff --git a/tests/ring/test_rings.cu b/tests/ring/test_rings.cu deleted file mode 100644 index 01759a334..000000000 --- a/tests/ring/test_rings.cu +++ /dev/null @@ -1,758 +0,0 @@ -// tests/ring/test_rings.cu — Standalone CUDA unit tests for ring primitives. -// -// Covers all Milestone 1 acceptance criteria: -// - Payload wrap correctness (two-span) -// - Task ring FIFO correctness by seq_no -// - ready_seq guard: consumer cannot read unpublished slot -// - DROP marker correctness -// -// Build with: make -C tests/ring -// Run with: ./tests/ring/test_rings - -#include -#include -#include -#include -#include - -#include - -// Include the ring headers under test. -#include "../../native/csrc/ring/task_entry.h" -#include "../../native/csrc/ring/ring_config.h" -#include "../../native/csrc/ring/payload_ring.cuh" -#include "../../native/csrc/ring/task_ring.cuh" - -// --------------------------------------------------------------------------- -// Helpers -// --------------------------------------------------------------------------- - -#define CUDA_CHECK(expr) \ - do { \ - cudaError_t _e = (expr); \ - if (_e != cudaSuccess) { \ - fprintf(stderr, "CUDA error at %s:%d — %s\n", \ - __FILE__, __LINE__, cudaGetErrorString(_e)); \ - exit(1); \ - } \ - } while (0) - -static int g_pass = 0; -static int g_fail = 0; - -#define EXPECT(cond) \ - do { \ - if (!(cond)) { \ - fprintf(stderr, " FAIL %s:%d " #cond "\n", \ - __FILE__, __LINE__); \ - g_fail++; \ - } else { \ - g_pass++; \ - } \ - } while (0) - -static void banner(const char* name) { - printf("[ TEST ] %s\n", name); - fflush(stdout); -} - -// =========================================================================== -// HOST-ONLY TESTS (pure arithmetic — __host__ __device__ functions) -// =========================================================================== - -// --------------------------------------------------------------------------- -// Test 1: payload_free_bytes basic accounting -// --------------------------------------------------------------------------- -static void test_payload_free_bytes() { - banner("payload_free_bytes — basic accounting"); - using namespace ring; - - uint64_t cap = 1024; - - // Empty ring. - EXPECT(payload_free_bytes(0, 0, cap) == cap); - - // Half full. - EXPECT(payload_free_bytes(512, 0, cap) == 512); - - // Producer wrapped around (head > cap). - EXPECT(payload_free_bytes(1024, 0, cap) == 0); // full - EXPECT(payload_free_bytes(1100, 100, cap) == 24); // 1024 - (1100-100) - - // Consumer advanced. - EXPECT(payload_free_bytes(800, 300, cap) == 524); -} - -// --------------------------------------------------------------------------- -// Test 2: payload_compute_spans — single span (no wrap) -// --------------------------------------------------------------------------- -static void test_payload_single_span() { - banner("payload_compute_spans — single span (no wrap)"); - using namespace ring; - - uint64_t cap = 256; - uint64_t head = 0; - - // From the beginning. - TwoSpan s = payload_compute_spans(head, cap, 100); - EXPECT(s.off1 == 0); - EXPECT(s.len1 == 100); - EXPECT(s.len2 == 0); - - // From the middle, still fits before end. - head = 100; - s = payload_compute_spans(head, cap, 50); - EXPECT(s.off1 == 100); - EXPECT(s.len1 == 50); - EXPECT(s.len2 == 0); - - // Exactly fills to the end. - head = 200; - s = payload_compute_spans(head, cap, 56); - EXPECT(s.off1 == 200); - EXPECT(s.len1 == 56); - EXPECT(s.len2 == 0); -} - -// --------------------------------------------------------------------------- -// Test 3: payload_compute_spans — two-span (wrap-around) -// --------------------------------------------------------------------------- -static void test_payload_two_span_wrap() { - banner("payload_compute_spans — two-span (wrap-around)"); - using namespace ring; - - uint64_t cap = 64; - - // Head at offset 50 in the ring; request 20 bytes that must wrap. - // span1: [50, 14) → bytes 50..63 - // span2: [0, 6) → bytes 0..5 - uint64_t head = 50; // physical off = 50 % 64 = 50 - TwoSpan s = payload_compute_spans(head, cap, 20); - EXPECT(s.off1 == 50); - EXPECT(s.len1 == 14); - EXPECT(s.off2 == 0); - EXPECT(s.len2 == 6); - EXPECT(s.len1 + s.len2 == 20); - - // Head has wrapped (head = 130 → phys = 2); request 5 bytes, no wrap. - head = 130; // phys = 130 % 64 = 2 - s = payload_compute_spans(head, cap, 5); - EXPECT(s.off1 == 2); - EXPECT(s.len1 == 5); - EXPECT(s.len2 == 0); - - // Head at physical position 60; request 10 bytes → wraps at boundary. - head = 60; // phys = 60 - s = payload_compute_spans(head, cap, 10); - EXPECT(s.off1 == 60); - EXPECT(s.len1 == 4); - EXPECT(s.off2 == 0); - EXPECT(s.len2 == 6); - - // Exactly at boundary (off == 0 after wrap), request full capacity. - head = 64; // phys = 0 - s = payload_compute_spans(head, cap, 64); - EXPECT(s.off1 == 0); - EXPECT(s.len1 == 64); - EXPECT(s.len2 == 0); -} - -// --------------------------------------------------------------------------- -// Test 4: payload ring simulation — multiple reservations with tail advance -// --------------------------------------------------------------------------- -static void test_payload_ring_simulation() { - banner("payload_ring_simulation — multi-reservation + tail advance"); - using namespace ring; - - uint64_t cap = 64; - uint64_t head = 0; - uint64_t tail = 0; - - // Reserve 50 bytes. - EXPECT(payload_free_bytes(head, tail, cap) == 64); - TwoSpan s1 = payload_compute_spans(head, cap, 50); - EXPECT(s1.off1 == 0 && s1.len1 == 50 && s1.len2 == 0); - payload_advance_head(head, 50); - EXPECT(head == 50); - - // Reserve 10 more bytes (still fits). - EXPECT(payload_free_bytes(head, tail, cap) == 14); - TwoSpan s2 = payload_compute_spans(head, cap, 10); - EXPECT(s2.off1 == 50 && s2.len1 == 10 && s2.len2 == 0); - payload_advance_head(head, 10); - EXPECT(head == 60); - - // Ring is nearly full (4 bytes free); need 30 — wait until consumer frees. - EXPECT(payload_free_bytes(head, tail, cap) == 4); - - // Consumer releases first reservation (50 bytes). - payload_release(tail, 50); - EXPECT(tail == 50); - EXPECT(payload_free_bytes(head, tail, cap) == 54); - - // Now reserve 30 bytes — wraps: off=60, len1=4, off2=0, len2=26. - TwoSpan s3 = payload_compute_spans(head, cap, 30); - EXPECT(s3.off1 == 60); - EXPECT(s3.len1 == 4); - EXPECT(s3.off2 == 0); - EXPECT(s3.len2 == 26); - payload_advance_head(head, 30); - EXPECT(head == 90); // 60 + 30 (unwrapped) - - // Consumer releases the second reservation (10 bytes). - payload_release(tail, 10); - EXPECT(tail == 60); - EXPECT(payload_free_bytes(head, tail, cap) == 34); // cap - (90 - 60) -} - -// --------------------------------------------------------------------------- -// Test 5: task_free_slots basic accounting -// --------------------------------------------------------------------------- -static void test_task_free_slots() { - banner("task_free_slots — basic accounting"); - using namespace ring; - - uint64_t cap = 16; - EXPECT(task_free_slots(0, 0, cap) == 16); - EXPECT(task_free_slots(8, 0, cap) == 8); - EXPECT(task_free_slots(16, 0, cap) == 0); // full - EXPECT(task_free_slots(16, 16, cap) == 16); // empty again after consumer drains - EXPECT(task_free_slots(1024, 1020, cap) == 12); -} - -// --------------------------------------------------------------------------- -// Test 6: TaskEntry layout (size, alignment, field offsets) -// --------------------------------------------------------------------------- -static void test_task_entry_layout() { - banner("TaskEntry — size and field offsets"); - using namespace ring; - - EXPECT(sizeof(TaskEntry) == 128); - EXPECT(alignof(TaskEntry) == 128); - EXPECT(offsetof(TaskEntry, ready_seq) == 0); - EXPECT(offsetof(TaskEntry, seq_no) == 8); - EXPECT(offsetof(TaskEntry, logical_task_id) == 16); - EXPECT(offsetof(TaskEntry, chunk_offset_bytes) == 24); - EXPECT(offsetof(TaskEntry, tensor_total_bytes) == 32); - EXPECT(offsetof(TaskEntry, payload_off1) == 40); - EXPECT(offsetof(TaskEntry, payload_len1) == 48); - EXPECT(offsetof(TaskEntry, payload_off2) == 56); - EXPECT(offsetof(TaskEntry, payload_len2) == 64); - EXPECT(offsetof(TaskEntry, chunk_idx) == 72); - EXPECT(offsetof(TaskEntry, hook_type) == 76); - EXPECT(offsetof(TaskEntry, hook_id) == 80); - EXPECT(offsetof(TaskEntry, flags) == 84); - EXPECT(offsetof(TaskEntry, reason) == 88); -} - -// --------------------------------------------------------------------------- -// Test 7: DROP marker fields -// --------------------------------------------------------------------------- -static void test_drop_marker_fields() { - banner("DROP marker — fields and flags"); - using namespace ring; - - TaskEntry e{}; - // Producer fills a DROP marker. - e.seq_no = 42; - e.logical_task_id = 999; - e.flags = TASK_FLAG_IS_DROP; - e.reason = DROP_REASON_TIMEOUT_NO_PROGRESS; - e.payload_len1 = 0; - e.payload_len2 = 0; - - EXPECT((e.flags & TASK_FLAG_IS_DROP) != 0); - EXPECT((e.flags & TASK_FLAG_IS_FIRST) == 0); - EXPECT((e.flags & TASK_FLAG_IS_LAST) == 0); - EXPECT(e.payload_len1 == 0); - EXPECT(e.payload_len2 == 0); - EXPECT(e.reason == DROP_REASON_TIMEOUT_NO_PROGRESS); - EXPECT(payload_chunk_bytes(e.payload_len1, e.payload_len2) == 0); - - // SENTINEL must not equal any plausible seq_no value (just check a few). - EXPECT(READY_SEQ_SENTINEL != 0); - EXPECT(READY_SEQ_SENTINEL != 42); - EXPECT(READY_SEQ_SENTINEL != uint64_t(-2)); // -2 != UINT64_MAX -} - -// --------------------------------------------------------------------------- -// Test 8: RingConfig defaults are sane -// --------------------------------------------------------------------------- -static void test_ring_config_defaults() { - banner("RingConfig — default values"); - using namespace ring; - - RingConfig cfg; - EXPECT(cfg.task_ring_entries == 1024); - EXPECT(cfg.payload_ring_bytes == 256ULL * 1024 * 1024); - EXPECT(cfg.chunk_bytes == 64ULL * 1024 * 1024); - EXPECT(cfg.wait_policy == WaitPolicy::INFINITE); - EXPECT(cfg.drop_reporting == DropReporting::DROP_TASK); - EXPECT(cfg.chunk_bytes * 4 <= cfg.payload_ring_bytes); // forward-progress -} - -// =========================================================================== -// GPU TESTS -// =========================================================================== - -// --------------------------------------------------------------------------- -// Test 9: task ring FIFO — sequential producer→consumer on the same stream. -// -// The ring is large enough to hold all entries at once, so the producer fills -// it entirely before the consumer drains it. This tests pure FIFO ordering -// without requiring concurrent kernel execution. -// -// Verifies: -// - Entries are consumed in seq_no order. -// - hook_id values match what the producer wrote. -// - task_release correctly resets ready_seq so slots can be reused. -// --------------------------------------------------------------------------- - -static const int FIFO_N = 32; // number of tasks to produce and consume - -// Producer: publishes FIFO_N entries; ring is large enough (cap >= FIFO_N) -// so it never needs to wait for the consumer. -__global__ void kernel_producer_fifo(ring::TaskEntry* entries, - uint64_t cap, - volatile uint64_t* d_head) -{ - for (int i = 0; i < FIFO_N; i++) { - uint64_t seq = *d_head; - - ring::TaskEntry e{}; - e.seq_no = seq; - e.logical_task_id = (uint64_t)i; - e.hook_id = (uint32_t)i; - e.flags = ring::TASK_FLAG_IS_FIRST | ring::TASK_FLAG_IS_LAST; - - ring::task_publish(entries, cap, seq, e); - *d_head = seq + 1; - } -} - -// Consumer: drains FIFO_N entries; called after the producer completes. -__global__ void kernel_consumer_fifo(const ring::TaskEntry* entries, - ring::TaskEntry* entries_rw, - uint64_t cap, - volatile uint64_t* d_tail, - uint32_t* results, - uint64_t* result_seq) -{ - for (int i = 0; i < FIFO_N; i++) { - uint64_t tail = *d_tail; - const ring::TaskEntry* e = ring::task_spin_wait(entries, cap, tail); - results[i] = e->hook_id; - result_seq[i] = e->seq_no; - ring::task_release(entries_rw, cap, tail); - *d_tail = tail + 1; - } -} - -static void test_task_ring_fifo_gpu() { - banner("task ring FIFO (GPU) — sequential producer then consumer"); - using namespace ring; - - // Ring larger than FIFO_N so producer never blocks waiting for space. - const uint64_t cap = 64; - - TaskEntry* d_entries; - CUDA_CHECK(cudaMalloc(&d_entries, cap * sizeof(TaskEntry))); - task_ring_init(d_entries, cap, 0); - - volatile uint64_t* d_head; - volatile uint64_t* d_tail; - CUDA_CHECK(cudaMallocManaged((void**)&d_head, sizeof(uint64_t))); - CUDA_CHECK(cudaMallocManaged((void**)&d_tail, sizeof(uint64_t))); - *d_head = 0; - *d_tail = 0; - - uint32_t* results; - uint64_t* result_seq; - CUDA_CHECK(cudaMallocManaged(&results, FIFO_N * sizeof(uint32_t))); - CUDA_CHECK(cudaMallocManaged(&result_seq, FIFO_N * sizeof(uint64_t))); - - // Sync to ensure managed memory is initialised before kernels run. - CUDA_CHECK(cudaDeviceSynchronize()); - - cudaStream_t stream; - CUDA_CHECK(cudaStreamCreate(&stream)); - - // Sequential: producer runs completely, then consumer drains. - kernel_producer_fifo<<<1,1,0,stream>>>(d_entries, cap, d_head); - kernel_consumer_fifo<<<1,1,0,stream>>>(d_entries, d_entries, cap, - d_tail, results, result_seq); - - CUDA_CHECK(cudaStreamSynchronize(stream)); - CUDA_CHECK(cudaDeviceSynchronize()); - - for (int i = 0; i < FIFO_N; i++) { - EXPECT(results[i] == (uint32_t)i); - EXPECT(result_seq[i] == (uint64_t)i); - } - - CUDA_CHECK(cudaStreamDestroy(stream)); - cudaFree(d_entries); - cudaFree((void*)d_head); - cudaFree((void*)d_tail); - cudaFree(results); - cudaFree(result_seq); -} - -// --------------------------------------------------------------------------- -// Test 10: ready_seq lifecycle — sequential guard semantics (single stream). -// -// In production the producer runs inside the CUDA graph and the consumer runs -// after the graph completes (ordered via a CUDA event). This test mimics -// that flow and verifies the full ready_seq lifecycle: -// -// Phase A — after init: ready_seq[i] == SENTINEL (slot not readable) -// Phase B — after publish: ready_seq[seq] == seq_no (slot readable) -// Phase C — after release: ready_seq[seq] == SENTINEL (slot recycled) -// -// Correctness guarantee covered: if task_publish omitted __threadfence() or -// wrote the wrong value, the consumer's spin_wait (which checks ready_seq) -// would malfunction. The sequential stream ordering means "happened-before" -// is enforced by the GPU's stream FIFO, but the volatile + membar semantics -// are still exercised by the actual instructions generated. -// --------------------------------------------------------------------------- - -__global__ void kernel_check_all_sentinel(const ring::TaskEntry* entries, - uint64_t cap, - int* ok) -{ - for (uint64_t i = 0; i < cap; i++) { - uint64_t rs = *reinterpret_cast( - &entries[i].ready_seq); - if (rs != ring::READY_SEQ_SENTINEL) { *ok = 0; return; } - } - *ok = 1; -} - -__global__ void kernel_publish_one(ring::TaskEntry* entries, - uint64_t cap, - volatile uint64_t* d_head, - int* ok) -{ - uint64_t seq = *d_head; - ring::TaskEntry e{}; - e.seq_no = seq; - e.logical_task_id = 0xDEADBEEFULL; - e.hook_id = 0xCAFEBABEu; - e.flags = ring::TASK_FLAG_IS_FIRST | ring::TASK_FLAG_IS_LAST; - e.tensor_total_bytes = 12345; - ring::task_publish(entries, cap, seq, e); - *d_head = seq + 1; - - // Verify ready_seq is now the sequence number we published. - uint64_t rs = *reinterpret_cast( - &entries[seq % cap].ready_seq); - *ok = (rs == seq) ? 1 : 0; -} - -__global__ void kernel_consume_one(const ring::TaskEntry* entries, - ring::TaskEntry* entries_rw, - uint64_t cap, - volatile uint64_t* d_tail, - uint64_t* out_hook_id, - uint64_t* out_ltid, - uint64_t* out_total, - int* sentinel_ok) -{ - uint64_t tail = *d_tail; - const ring::TaskEntry* e = ring::task_spin_wait(entries, cap, tail); - *out_hook_id = e->hook_id; - *out_ltid = e->logical_task_id; - *out_total = e->tensor_total_bytes; - ring::task_release(entries_rw, cap, tail); - *d_tail = tail + 1; - - // After release, ready_seq must be SENTINEL again. - uint64_t rs = *reinterpret_cast( - &entries_rw[tail % cap].ready_seq); - *sentinel_ok = (rs == ring::READY_SEQ_SENTINEL) ? 1 : 0; -} - -static void test_ready_seq_lifecycle_gpu() { - banner("ready_seq lifecycle (GPU) — init sentinel, publish, consume, reset"); - using namespace ring; - - const uint64_t cap = 8; - - TaskEntry* d_entries; - CUDA_CHECK(cudaMalloc(&d_entries, cap * sizeof(TaskEntry))); - task_ring_init(d_entries, cap, 0); - - volatile uint64_t* d_head; - volatile uint64_t* d_tail; - CUDA_CHECK(cudaMallocManaged((void**)&d_head, sizeof(uint64_t))); - CUDA_CHECK(cudaMallocManaged((void**)&d_tail, sizeof(uint64_t))); - *d_head = 0; - *d_tail = 0; - - int* d_ok_sentinel; - int* d_ok_publish; - int* d_ok_sentinel2; - uint64_t* d_hook_id; - uint64_t* d_ltid; - uint64_t* d_total; - CUDA_CHECK(cudaMallocManaged(&d_ok_sentinel, sizeof(int))); - CUDA_CHECK(cudaMallocManaged(&d_ok_publish, sizeof(int))); - CUDA_CHECK(cudaMallocManaged(&d_ok_sentinel2, sizeof(int))); - CUDA_CHECK(cudaMallocManaged(&d_hook_id, sizeof(uint64_t))); - CUDA_CHECK(cudaMallocManaged(&d_ltid, sizeof(uint64_t))); - CUDA_CHECK(cudaMallocManaged(&d_total, sizeof(uint64_t))); - - CUDA_CHECK(cudaDeviceSynchronize()); - - cudaStream_t stream; - CUDA_CHECK(cudaStreamCreate(&stream)); - - // All phases run sequentially on one stream. - kernel_check_all_sentinel<<<1,1,0,stream>>>(d_entries, cap, d_ok_sentinel); - kernel_publish_one<<<1,1,0,stream>>>(d_entries, cap, d_head, d_ok_publish); - kernel_consume_one<<<1,1,0,stream>>>(d_entries, d_entries, cap, - d_tail, d_hook_id, d_ltid, d_total, - d_ok_sentinel2); - - CUDA_CHECK(cudaStreamSynchronize(stream)); - CUDA_CHECK(cudaDeviceSynchronize()); - - EXPECT(*d_ok_sentinel == 1); // all slots are SENTINEL after init - EXPECT(*d_ok_publish == 1); // ready_seq is set correctly after publish - EXPECT(*d_hook_id == 0xCAFEBABEULL); - EXPECT(*d_ltid == 0xDEADBEEFULL); - EXPECT(*d_total == 12345ULL); - EXPECT(*d_ok_sentinel2 == 1); // slot reset to SENTINEL after release - - CUDA_CHECK(cudaStreamDestroy(stream)); - cudaFree(d_entries); - cudaFree((void*)d_head); - cudaFree((void*)d_tail); - cudaFree(d_ok_sentinel); - cudaFree(d_ok_publish); - cudaFree(d_ok_sentinel2); - cudaFree(d_hook_id); - cudaFree(d_ltid); - cudaFree(d_total); -} - -// --------------------------------------------------------------------------- -// Test 11: DROP marker round-trip (GPU, sequential) — producer emits IS_DROP, -// consumer reads it and observes zero payload and IS_DROP flag. -// --------------------------------------------------------------------------- - -__global__ void kernel_producer_drop(ring::TaskEntry* entries, - uint64_t cap, - volatile uint64_t* d_head) -{ - uint64_t seq = *d_head; - ring::TaskEntry e{}; - e.seq_no = seq; - e.logical_task_id = 777; - e.flags = ring::TASK_FLAG_IS_DROP; - e.reason = ring::DROP_REASON_TIMEOUT_NO_PROGRESS; - e.payload_len1 = 0; - e.payload_len2 = 0; - ring::task_publish(entries, cap, seq, e); - *d_head = seq + 1; -} - -__global__ void kernel_consumer_drop(const ring::TaskEntry* entries, - ring::TaskEntry* entries_rw, - uint64_t cap, - volatile uint64_t* d_tail, - uint32_t* out_flags, - uint32_t* out_reason, - uint64_t* out_len_total) -{ - uint64_t tail = *d_tail; - const ring::TaskEntry* e = ring::task_spin_wait(entries, cap, tail); - *out_flags = e->flags; - *out_reason = e->reason; - *out_len_total = e->payload_len1 + e->payload_len2; - ring::task_release(entries_rw, cap, tail); - *d_tail = tail + 1; -} - -static void test_drop_marker_gpu() { - banner("DROP marker (GPU) — producer emits drop, consumer sees IS_DROP"); - using namespace ring; - - const uint64_t cap = 4; - - TaskEntry* d_entries; - CUDA_CHECK(cudaMalloc(&d_entries, cap * sizeof(TaskEntry))); - task_ring_init(d_entries, cap, 0); - - volatile uint64_t* d_head; - volatile uint64_t* d_tail; - CUDA_CHECK(cudaMallocManaged((void**)&d_head, sizeof(uint64_t))); - CUDA_CHECK(cudaMallocManaged((void**)&d_tail, sizeof(uint64_t))); - *d_head = 0; - *d_tail = 0; - - uint32_t* out_flags; - uint32_t* out_reason; - uint64_t* out_len_total; - CUDA_CHECK(cudaMallocManaged(&out_flags, sizeof(uint32_t))); - CUDA_CHECK(cudaMallocManaged(&out_reason, sizeof(uint32_t))); - CUDA_CHECK(cudaMallocManaged(&out_len_total, sizeof(uint64_t))); - - CUDA_CHECK(cudaDeviceSynchronize()); - - cudaStream_t stream; - CUDA_CHECK(cudaStreamCreate(&stream)); - - // Sequential: producer publishes DROP, then consumer drains it. - kernel_producer_drop<<<1,1,0,stream>>>(d_entries, cap, d_head); - kernel_consumer_drop<<<1,1,0,stream>>>(d_entries, d_entries, cap, - d_tail, - out_flags, out_reason, out_len_total); - - CUDA_CHECK(cudaStreamSynchronize(stream)); - CUDA_CHECK(cudaDeviceSynchronize()); - - EXPECT((*out_flags & TASK_FLAG_IS_DROP) != 0); - EXPECT((*out_flags & TASK_FLAG_IS_FIRST) == 0); - EXPECT((*out_flags & TASK_FLAG_IS_LAST) == 0); - EXPECT(*out_reason == DROP_REASON_TIMEOUT_NO_PROGRESS); - EXPECT(*out_len_total == 0); - - CUDA_CHECK(cudaStreamDestroy(stream)); - cudaFree(d_entries); - cudaFree((void*)d_head); - cudaFree((void*)d_tail); - cudaFree(out_flags); - cudaFree(out_reason); - cudaFree(out_len_total); -} - -// --------------------------------------------------------------------------- -// Test 12: task ring wrap-around (slot reuse) — sequential producer→consumer, -// ring slots reused 4× to verify FIFO is preserved across full wraps. -// -// Uses cap < WRAP_TOTAL so that slots are recycled; sequential (same stream) -// to avoid concurrent-spin dependency issues. -// --------------------------------------------------------------------------- -static const int WRAP_TOTAL = 64; // 4× a cap-16 ring - -__global__ void kernel_producer_wrap(ring::TaskEntry* entries, - uint64_t cap, - volatile uint64_t* d_head, - int n_total, - int base_id) -{ - // Producer assumes ring is drained by a previous consumer pass. - for (int i = 0; i < n_total; i++) { - uint64_t seq = *d_head; - ring::TaskEntry e{}; - e.seq_no = seq; - e.hook_id = (uint32_t)(base_id + i); // globally unique hook_id - e.flags = ring::TASK_FLAG_IS_FIRST | ring::TASK_FLAG_IS_LAST; - ring::task_publish(entries, cap, seq, e); - *d_head = seq + 1; - } -} - -__global__ void kernel_consumer_wrap(const ring::TaskEntry* entries, - ring::TaskEntry* entries_rw, - uint64_t cap, - volatile uint64_t* d_tail, - uint32_t* results, - int n_total) -{ - for (int i = 0; i < n_total; i++) { - uint64_t tail = *d_tail; - const ring::TaskEntry* e = ring::task_spin_wait(entries, cap, tail); - results[i] = e->hook_id; - ring::task_release(entries_rw, cap, tail); - *d_tail = tail + 1; - } -} - -static void test_task_ring_wrap_reuse_gpu() { - banner("task ring wrap-around (GPU) — slot reuse, FIFO preserved"); - using namespace ring; - - // cap < WRAP_TOTAL to force slot reuse; run in batches of cap entries. - const uint64_t cap = 16; - const int batch = (int)cap; // publish+consume `batch` at a time - - TaskEntry* d_entries; - CUDA_CHECK(cudaMalloc(&d_entries, cap * sizeof(TaskEntry))); - task_ring_init(d_entries, cap, 0); - - volatile uint64_t* d_head; - volatile uint64_t* d_tail; - CUDA_CHECK(cudaMallocManaged((void**)&d_head, sizeof(uint64_t))); - CUDA_CHECK(cudaMallocManaged((void**)&d_tail, sizeof(uint64_t))); - *d_head = 0; - *d_tail = 0; - - uint32_t* results; - CUDA_CHECK(cudaMallocManaged(&results, WRAP_TOTAL * sizeof(uint32_t))); - - CUDA_CHECK(cudaDeviceSynchronize()); - - cudaStream_t stream; - CUDA_CHECK(cudaStreamCreate(&stream)); - - // Process in batches of `cap` entries so the ring never overflows. - // Each pass: producer fills the ring, then consumer drains it. - for (int base = 0; base < WRAP_TOTAL; base += batch) { - kernel_producer_wrap<<<1,1,0,stream>>>(d_entries, cap, d_head, - batch, base); - kernel_consumer_wrap<<<1,1,0,stream>>>(d_entries, d_entries, cap, - d_tail, results + base, batch); - } - - CUDA_CHECK(cudaStreamSynchronize(stream)); - CUDA_CHECK(cudaDeviceSynchronize()); - - for (int i = 0; i < WRAP_TOTAL; i++) { - EXPECT(results[i] == (uint32_t)i); - } - - CUDA_CHECK(cudaStreamDestroy(stream)); - cudaFree(d_entries); - cudaFree((void*)d_head); - cudaFree((void*)d_tail); - cudaFree(results); -} - -// =========================================================================== -// main -// =========================================================================== -int main() { - // Disable stdout buffering so output is visible even if the process is - // killed mid-run (helps debug hanging GPU tests). - setbuf(stdout, nullptr); - - printf("=== Ring Unit Tests ===\n\n"); - - // Print device info. - int dev = 0; - cudaDeviceProp prop{}; - CUDA_CHECK(cudaGetDeviceProperties(&prop, dev)); - printf("Device: %s (CC %d.%d)\n\n", prop.name, - prop.major, prop.minor); - - // Host tests (pure arithmetic). - test_payload_free_bytes(); - test_payload_single_span(); - test_payload_two_span_wrap(); - test_payload_ring_simulation(); - test_task_free_slots(); - test_task_entry_layout(); - test_drop_marker_fields(); - test_ring_config_defaults(); - - // GPU tests (all sequential — producer then consumer on same stream). - test_task_ring_fifo_gpu(); - test_ready_seq_lifecycle_gpu(); - test_drop_marker_gpu(); - test_task_ring_wrap_reuse_gpu(); - - printf("\n=== Results: %d passed, %d failed ===\n", g_pass, g_fail); - return (g_fail > 0) ? 1 : 0; -} From 91e8a33f5d65803eccdefb359dad230db0b2bcd0 Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Wed, 26 Aug 2026 00:33:32 -0400 Subject: [PATCH 17/20] docs(capture): record Phase 6 status, instruments, and the native decision Phase 6 was the only phase with no status line, and the two instruments built for it were undocumented. Record what exists, what it guarantees, and what is still missing before the default sink can move. Write down the architectural decision, which was not in the doc at all: the Python implementation is the reference and conformance suite permanently, and the production pack-and-upload plane will be native C++. The reason is structural rather than performance-related -- the ring reconstructs tensors on a native callback thread specifically to avoid touching Python or the GIL, and DMXHostEngine consumes pre-assembled rows from that thread, so a Python pack sink has no hot-path caller and creating one would reintroduce the per-tensor GIL contention the ring exists to avoid. This supersedes the Phase 2 limitation that described the Python pipeline as an interim stand-in; that limitation now points at the decision. Document the fault-injection contract as a table of required behaviours rather than a description of the harness, since a native writer has to reproduce the behaviour and not the implementation. Same for the conformance manifest: what makes a writer conformant is that the same corpus yields the same document. Note the l2_norm scaling in the Phase 5 limitations. The summary contract's numerical behaviour changed when the direct sqrt(sum(x**2)) form was found to overflow float64 and return infinity where the true norm is finite. Every command the doc now advertises was run before committing. --- docs/capture-storage-design.md | 84 +++++++++++++++++++++++++++++++++- 1 file changed, 82 insertions(+), 2 deletions(-) diff --git a/docs/capture-storage-design.md b/docs/capture-storage-design.md index b7ff1f85b..967f64143 100644 --- a/docs/capture-storage-design.md +++ b/docs/capture-storage-design.md @@ -683,6 +683,80 @@ PYTHONPATH=src python -m benchmarks.bench_capture_search \ - Switch the default only after compatibility, recovery, and throughput gates pass; preserve configuration rollback. +Status: not started. Two of the instruments it depends on now exist; the +comparison itself has not been run, and the production writer it compares +against has not been built. + +#### Decision: the production writer is native + +The Python implementation is a **reference implementation and conformance +suite**, permanently -- not a candidate production writer. The reason is +structural rather than performance-related: the ring transport reconstructs +tensors on a native callback thread specifically to avoid touching Python or the +GIL, and `DMXHostEngine` receives pre-assembled rows from that thread. There is +no hot-path caller for a Python pack sink, and creating one would reintroduce +per-tensor GIL contention that the ring exists to avoid. + +The pack-and-upload plane will therefore be written in C++ alongside the +existing `ClickHouseInsertStage`, reusing `batching_queue.hpp` and +`pipelined_engine.hpp`. This supersedes the Phase 2 limitation describing the +Python pipeline as an interim stand-in for "the final reusable native slab +allocator": that is now the plan of record rather than a gap. + +#### Fault injection + +`tests/_faults.py` wraps the three boundaries that can misbehave -- object +store, ClickHouse client, and pack sink. Faults are scripted rather than random: +a schedule names which calls fail and how, so a failure reproduces exactly. + +The characterised behaviour, which a native writer must reproduce: + +| Fault | Required behaviour | +|---|---| +| Short read from the store | Refused, not silently truncated | +| Read failure mid-index | Aborts that pack; no partial pack is produced | +| Immutable key written twice | Converges; not a conflict | +| Sink failure | Pipeline fails loudly and refuses further admission | +| Insert failure | Pack left uncommitted; the batch is replayable | +| Duplicated insert | Absorbed by replay semantics | +| One corrupt pack in a batch | Fails only itself; the batch still indexes | + +```bash +python -m pytest tests/test_capture_faults.py -q +``` + +#### Conformance manifest + +`tests/tools/golden_workload.py` produces the golden-workload comparison this +phase requires, as a single JSON document over a deterministic corpus covering +every dtype the format accepts: pack identity and checksum, per-capture payload +sha256 and crc32, decoded-tensor sha256, placement, and the full summary +contract. Every value is language-neutral -- byte counts, hex digests and +integers -- so a native writer is conformant exactly when the same corpus +produces the same manifest. + +```bash +python tests/tools/golden_workload.py generate --out golden.json +python tests/tools/golden_workload.py verify --manifest golden.json +``` + +The recorded manifest lives at `tests/data/capture_golden_manifest.json` and is +checked on every CPU run. `verify` diffs field by field, so a mismatch names the +capture and field that moved. + +#### Remaining before the default can switch + +- The native pack-and-upload plane, and configuration to select a sink with + rollback preserved. Until a sink is selectable, "keep the current sink as the + default initially" has nothing to compare against. +- Phase 3's capacity gate, which is still pending hardware. A throughput gate + cannot pass while the 1.2x requirement is unmeasured. +- Performance variance under fault injection, which is not yet recorded. +- Three decisions the migration forces: whether the public views keep `FINAL` + now that its cost is measured at 1.85x; where `index_version` comes from once + more than one indexer runs, since the current per-process clock assumes a + single writer; and whether catalog facets belong in the public views. + No phase requires a CUDA-side change. ## Operating signals @@ -718,7 +792,10 @@ bytes, and retained failure details all have explicit caps. ## Phase 2 limitations - The implementation is a Python CPU reference, not the final reusable native - slab allocator. + slab allocator. As of the Phase 6 decision this is permanent: the Python + implementation is the reference and conformance suite, and the production + writer will be native. See *Phase 6 -- Decision: the production writer is + native*. - One process owns a spool directory; cross-process locking is not implemented. - Durable mode stages synchronously and uploads through a separate explicit uploader, so remote backpressure is isolated from local commit. @@ -772,7 +849,10 @@ bytes, and retained failure details all have explicit caps. - The per-extension time budget is checked after each call. A runaway extension is reported, not interrupted; preemption needs a worker boundary. - Core summary statistics cover finite elements only, with `nan_count`, - `inf_count` and `finite_count` reported alongside. + `inf_count` and `finite_count` reported alongside. `l2_norm` factors out the + largest magnitude before squaring, because the direct `sqrt(sum(x**2))` form + overflows float64 for large-magnitude tensors and returns infinity where the + true norm is finite. - A selection is one bounded page. Callers paginate explicitly. ## Alternatives considered From fd3ac3e48d4c3afc57034219fc5db5b7cad5a1d8 Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Wed, 26 Aug 2026 01:01:54 -0400 Subject: [PATCH 18/20] fix(capture): publish the watermark only after its batch is durable Samfisheryu found that a reader could pin a half-written batch. CatalogIndexer assigns one index_version and then writes descriptors across several INSERTs before the pack markers, so a reader sampling max(index_version) on the descriptor table observes that version mid-batch. Reproduced against the real indexer path with four descriptors written as two INSERTs: rows visible when the snapshot was taken : 2 rows returned by the pinned watermark : 4 Each INSERT is atomic on its own; the loop around them is not, so "pinned snapshot" was not true. The version is now published as a separate final step, after every descriptor batch and pack marker is durable, into an append-only *_index_watermark table. Readers take the watermark only from that log. The table is plain MergeTree rather than ReplacingMergeTree: it is a history, and Replacing would eventually collapse exactly what a pinned snapshot reads. Add tests/test_clickhouse_snapshot_live.py for the corner cases the first round of live tests could not reach. Those tests confirmed the mechanism rather than challenging the engine underneath it: ReplacingMergeTree deduplicates "at an unknown time", so reading immediately after two writes only ever exercises the pre-merge state, and max_rows_per_insert defaults to 10,000, so no corpus I used ever split a batch across INSERTs. Merges are now forced with OPTIMIZE FINAL and batches are written with a deliberately small insert size. The second blocker Samfisheryu raised is not fixed here and is recorded as a strict xfail: argMax(...) WHERE index_version <= W reads rows that ReplacingMergeTree is defined to delete on merge, so a pinned snapshot expires at a time nobody controls. That needs a design change rather than a patch, and the test will flag as soon as one lands. --- src/dmi/storage/capture/catalog.py | 21 ++ src/dmi/storage/capture/clickhouse_catalog.py | 31 +++ src/dmi/storage/capture/clickhouse_reader.py | 20 +- tests/test_clickhouse_capture_reader.py | 3 + tests/test_clickhouse_facets_live.py | 10 + tests/test_clickhouse_reader_live.py | 18 ++ tests/test_clickhouse_snapshot_live.py | 231 ++++++++++++++++++ 7 files changed, 327 insertions(+), 7 deletions(-) create mode 100644 tests/test_clickhouse_snapshot_live.py diff --git a/src/dmi/storage/capture/catalog.py b/src/dmi/storage/capture/catalog.py index 8d7dc677c..cb883c169 100644 --- a/src/dmi/storage/capture/catalog.py +++ b/src/dmi/storage/capture/catalog.py @@ -30,6 +30,15 @@ def commit_packs( self, refs: Sequence[PackRef], *, index_version: int ) -> None: ... + def publish_watermark( + self, + *, + index_version: int, + published_at_ns: int, + indexed_rows: int, + indexed_packs: int, + ) -> None: ... + @dataclass(frozen=True, slots=True) class CatalogIndexerConfig: @@ -221,6 +230,18 @@ def index(self, refs: Sequence[PackRef]) -> IndexResult: descriptor_inserts += 1 if valid_refs: self._writer.commit_packs(valid_refs, index_version=version) + # Publish last. Descriptors go out across several INSERTs and the pack + # markers after them, so the version is only a truthful snapshot once + # all of that is durable. A reader that derived the watermark from the + # descriptor table itself would see this version mid-batch. + publish = getattr(self._writer, "publish_watermark", None) + if publish is not None: + publish( + index_version=version, + published_at_ns=self._clock_ns(), + indexed_rows=len(descriptors), + indexed_packs=len(valid_refs), + ) result = IndexResult( requested_packs=len(unique), skipped_packs=len(unique) - len(pending), diff --git a/src/dmi/storage/capture/clickhouse_catalog.py b/src/dmi/storage/capture/clickhouse_catalog.py index a460d3bcd..cde4e1dce 100644 --- a/src/dmi/storage/capture/clickhouse_catalog.py +++ b/src/dmi/storage/capture/clickhouse_catalog.py @@ -94,6 +94,7 @@ def __init__( self._capture_view = f"{prefix}_capture" self._pack_raw = f"{prefix}_pack_inventory_raw" self._pack_view = f"{prefix}_pack_inventory" + self._watermark = f"{prefix}_index_watermark" def ensure_schema(self) -> None: database = _quoted(self._config.database) @@ -135,6 +136,19 @@ def ensure_schema(self) -> None: f"{name} {kind} MATERIALIZED {expression}" ) + # A version becomes readable only once its whole batch is durable, so + # the watermark cannot be derived from the descriptor table: a reader + # sampling max(index_version) there sees a version mid-batch, between + # the INSERTs that make it up. This table is written as the last step of + # an indexing call instead. Plain MergeTree, because it is a log -- + # ReplacingMergeTree would eventually collapse the history a pinned + # snapshot reads. + self._client.execute( + f"""CREATE TABLE IF NOT EXISTS {database}.{_quoted(self._watermark)} ( +index_version UInt64, published_at_ns UInt64, indexed_rows UInt64, indexed_packs UInt32 +) ENGINE = MergeTree ORDER BY index_version""" + ) + capture_public = ", ".join(_CAPTURE_COLUMNS[:-1]) pack_public = ", ".join(_PACK_COLUMNS[:-1]) self._client.execute( @@ -174,6 +188,23 @@ def write_descriptors( rows, ) + def publish_watermark( + self, + *, + index_version: int, + published_at_ns: int, + indexed_rows: int, + indexed_packs: int, + ) -> None: + """Make a version readable, after everything it covers is durable.""" + self._validate_version(index_version) + self._validate_version(published_at_ns) + self._client.execute( + f"INSERT INTO {self._qualified(self._watermark)} " + "(index_version, published_at_ns, indexed_rows, indexed_packs) VALUES", + [(index_version, published_at_ns, indexed_rows, indexed_packs)], + ) + def commit_packs( self, refs: Sequence[PackRef], *, index_version: int ) -> None: diff --git a/src/dmi/storage/capture/clickhouse_reader.py b/src/dmi/storage/capture/clickhouse_reader.py index 992cd177c..b2262a4c2 100644 --- a/src/dmi/storage/capture/clickhouse_reader.py +++ b/src/dmi/storage/capture/clickhouse_reader.py @@ -123,6 +123,7 @@ def __init__( self._client = client self._config = config or ClickHouseReaderConfig() self._capture_raw = f"{self._config.table_prefix}_capture_raw" + self._watermark_table = f"{self._config.table_prefix}_index_watermark" @property def config(self) -> ClickHouseReaderConfig: @@ -131,14 +132,16 @@ def config(self) -> ClickHouseReaderConfig: # -- public API --------------------------------------------------------- def current_watermark(self) -> str: - """The catalog's high-water ``index_version``. + """The newest version whose indexing batch is fully committed. - ``index_version`` is not part of the sort key, so this is a full-column - aggregate rather than an index lookup -- cheap in a columnar engine, but - not free, and it runs once per search. + Read from the published watermark log, never from the descriptor table. + One indexing call writes descriptors across several INSERTs and then the + pack markers; ``max(index_version)`` over the descriptors would expose + that version between those writes, letting a reader pin a half-written + batch that keeps growing under it. """ rows = self._client.execute( - f"SELECT max(index_version) FROM {self._qualified()}", + f"SELECT max(index_version) FROM {self._qualified(self._watermark_table)}", settings=self._config.settings, ) if not rows or rows[0][0] is None: @@ -207,8 +210,11 @@ def get_by_ids( # -- SQL construction --------------------------------------------------- - def _qualified(self) -> str: - return f"{_quoted(self._config.database)}.{_quoted(self._capture_raw)}" + def _qualified(self, table: str | None = None) -> str: + return ( + f"{_quoted(self._config.database)}." + f"{_quoted(table or self._capture_raw)}" + ) @staticmethod def _projection() -> str: diff --git a/tests/test_clickhouse_capture_reader.py b/tests/test_clickhouse_capture_reader.py index 2476b1b87..1e8b4fbaa 100644 --- a/tests/test_clickhouse_capture_reader.py +++ b/tests/test_clickhouse_capture_reader.py @@ -55,6 +55,9 @@ def __init__(self, *, descriptors=(), watermark=_WATERMARK, pages=None): def execute(self, query, params=None, **kwargs): self.calls.append((" ".join(query.split()), params, kwargs)) if "max(index_version)" in query: + # The watermark now comes from the published log, not the + # descriptor table. + assert "_index_watermark" in query, query return [(self._watermark,)] page = self._pages.pop(0) if self._pages else [] return [_row(item) for item in page] diff --git a/tests/test_clickhouse_facets_live.py b/tests/test_clickhouse_facets_live.py index 4ad845f0c..d3d9f4920 100644 --- a/tests/test_clickhouse_facets_live.py +++ b/tests/test_clickhouse_facets_live.py @@ -23,6 +23,16 @@ pytestmark = [pytest.mark.manual, pytest.mark.clickhouse] +def _publish(writer, index_version: int, *, rows: int = 0, packs: int = 0) -> None: + """Publishing is a separate step; CatalogIndexer does it, direct writes must.""" + writer.publish_watermark( + index_version=index_version, + published_at_ns=index_version, + indexed_rows=rows, + indexed_packs=packs, + ) + + @contextmanager def _writer(): clickhouse_driver = pytest.importorskip("clickhouse_driver") diff --git a/tests/test_clickhouse_reader_live.py b/tests/test_clickhouse_reader_live.py index c91657d11..4753ebb35 100644 --- a/tests/test_clickhouse_reader_live.py +++ b/tests/test_clickhouse_reader_live.py @@ -28,6 +28,16 @@ pytestmark = [pytest.mark.manual, pytest.mark.clickhouse] +def _publish(writer, index_version: int, *, rows: int = 0, packs: int = 0) -> None: + """Publishing is a separate step; CatalogIndexer does it, direct writes must.""" + writer.publish_watermark( + index_version=index_version, + published_at_ns=index_version, + indexed_rows=rows, + indexed_packs=packs, + ) + + @contextmanager def _catalog(): clickhouse_driver = pytest.importorskip("clickhouse_driver") @@ -80,6 +90,7 @@ def test_a_full_walk_returns_the_corpus_exactly(): corpus = synthetic_descriptors(250) with _catalog() as (writer, reader): writer.write_descriptors(corpus, index_version=1) + _publish(writer, 1) walked, pages, _ = _walk(reader, CaptureQuery(limit=40)) @@ -96,9 +107,11 @@ def test_watermark_isolates_rows_indexed_after_the_first_page(): ) with _catalog() as (writer, reader): writer.write_descriptors(corpus, index_version=1) + _publish(writer, 1) first = reader.search(CaptureQuery(limit=40)) writer.write_descriptors(later, index_version=2) + _publish(writer, 2) rest, _, _ = _walk(reader, CaptureQuery(limit=40, cursor=first.next_cursor)) @@ -121,7 +134,9 @@ def test_snapshot_returns_the_version_at_the_watermark_not_the_latest(): ) with _catalog() as (writer, reader): writer.write_descriptors(original, index_version=1) + _publish(writer, 1) writer.write_descriptors(revised, index_version=2) + _publish(writer, 2) at_first = reader.search(CaptureQuery(limit=10)) assert at_first.watermark == "2" @@ -139,6 +154,7 @@ def test_replayed_indexing_yields_one_logical_row(): with _catalog() as (writer, reader): for version in (1, 2, 3): writer.write_descriptors(corpus, index_version=version) + _publish(writer, version) walked, _, _ = _walk(reader, CaptureQuery(limit=10)) @@ -149,6 +165,7 @@ def test_get_by_ids_resolves_a_selection_at_its_watermark(): corpus = synthetic_descriptors(20) with _catalog() as (writer, reader): writer.write_descriptors(corpus, index_version=1) + _publish(writer, 1) page = reader.search(CaptureQuery(limit=20)) wanted = [item.capture_id for item in page.items[:5]] @@ -162,6 +179,7 @@ def test_filters_narrow_the_walk(): corpus = synthetic_descriptors(30) with _catalog() as (writer, reader): writer.write_descriptors(corpus, index_version=1) + _publish(writer, 1) walked, _, _ = _walk( reader, diff --git a/tests/test_clickhouse_snapshot_live.py b/tests/test_clickhouse_snapshot_live.py new file mode 100644 index 000000000..9dfb8f4e1 --- /dev/null +++ b/tests/test_clickhouse_snapshot_live.py @@ -0,0 +1,231 @@ +"""Corner cases the first round of snapshot tests did not reach. + +Two defects got through review because the earlier live tests confirmed the +mechanism rather than challenging the engine underneath it: + +* ``ReplacingMergeTree`` deduplicates "at an unknown time". A test that writes + two versions and reads immediately only ever exercises the pre-merge state, + so it cannot see that a merge deletes the versions a pinned watermark needs. +* ``max_rows_per_insert`` defaults to 10,000, so a small corpus never splits a + batch into multiple INSERTs and the non-atomic window is unreachable. + +Both are forced here: merges are triggered with ``OPTIMIZE ... FINAL``, and +batches are written with a deliberately small insert size. + +Run against a reachable ClickHouse: + + DMI_CLICKHOUSE_HOST=127.0.0.1 python -m pytest \ + tests/test_clickhouse_snapshot_live.py -m "manual and clickhouse" -q +""" + +from __future__ import annotations + +from contextlib import contextmanager +from os import environ +from uuid import uuid4 + +import pytest + +from benchmarks.bench_capture_catalog import synthetic_descriptors +from dmi.storage.capture import ( + CaptureQuery, + ClickHouseCaptureCatalog, + ClickHouseCatalogConfig, + ClickHouseCatalogWriter, + ClickHouseReaderConfig, +) + + +pytestmark = [pytest.mark.manual, pytest.mark.clickhouse] + + +def _publish(writer, index_version: int, *, rows: int = 0, packs: int = 0) -> None: + """Publishing is a separate step; CatalogIndexer does it, direct writes must.""" + writer.publish_watermark( + index_version=index_version, + published_at_ns=index_version, + indexed_rows=rows, + indexed_packs=packs, + ) + + +@contextmanager +def _catalog(): + clickhouse_driver = pytest.importorskip("clickhouse_driver") + client = clickhouse_driver.Client( + host=environ.get("DMI_CLICKHOUSE_HOST", "127.0.0.1"), + port=int(environ.get("DMI_CLICKHOUSE_PORT", "9000")), + ) + prefix = f"dmi_snapshot_test_{uuid4().hex}" + config = ClickHouseCatalogConfig( + database=environ.get("DMI_CLICKHOUSE_DATABASE", "default"), + table_prefix=prefix, + ) + writer = ClickHouseCatalogWriter(client, config) + reader = ClickHouseCaptureCatalog( + client, ClickHouseReaderConfig.from_catalog(config) + ) + created = False + try: + writer.ensure_schema() + created = True + yield writer, reader, client, config + finally: + if created: + database = config.database + for kind, suffix in ( + ("VIEW", "capture"), + ("VIEW", "pack_inventory"), + ("TABLE", "capture_raw"), + ("TABLE", "pack_inventory_raw"), + ): + client.execute( + f"DROP {kind} IF EXISTS `{database}`.`{prefix}_{suffix}`" + ) + + +def _merge(client, config): + """Force the deduplication a background merge would eventually do.""" + client.execute( + f"OPTIMIZE TABLE `{config.database}`.`{config.table_prefix}_capture_raw` FINAL" + ) + + +def _raw_rows(client, config) -> int: + return client.execute( + f"SELECT count() FROM `{config.database}`.`{config.table_prefix}_capture_raw`" + )[0][0] + + +# --- merges versus pinned watermarks ----------------------------------------- + + +@pytest.mark.xfail( + strict=True, + reason=( + "Open design problem: argMax(...) WHERE index_version <= W needs rows " + "ReplacingMergeTree is defined to delete on merge, so a pinned snapshot " + "expires at a time nobody controls. Needs a snapshot mechanism that does " + "not depend on superseded versions surviving." + ), +) +def test_a_merge_does_not_destroy_a_pinned_snapshot(): + """A watermark must keep resolving after ClickHouse deduplicates. + + ReplacingMergeTree keeps only the highest version per sorting key, so any + snapshot that reads superseded versions has an expiry no caller controls. + """ + original = synthetic_descriptors(1) + revised = tuple( + d.__class__(metadata=d.metadata, locator=d.locator.__class__( + **{**{f: getattr(d.locator, f) for f in d.locator.__slots__}, + "offset": 999_999})) + for d in original + ) + with _catalog() as (writer, reader, client, config): + writer.write_descriptors(original, index_version=1) + _publish(writer, 1) + writer.write_descriptors(revised, index_version=2) + _publish(writer, 2) + + before = reader.get_by_ids([original[0].capture_id], watermark="1") + assert len(before) == 1, "precondition: the pinned read works pre-merge" + assert before[0].locator.offset == original[0].locator.offset + + _merge(client, config) + + after = reader.get_by_ids([original[0].capture_id], watermark="1") + assert len(after) == 1, ( + "a pinned watermark stopped resolving after a merge: " + f"raw rows went to {_raw_rows(client, config)}" + ) + assert after[0].locator.offset == original[0].locator.offset + + +def test_a_walk_still_completes_after_a_merge_mid_pagination(): + """Merges run concurrently with reads; a walk must not lose rows to one.""" + corpus = synthetic_descriptors(60) + with _catalog() as (writer, reader, client, config): + writer.write_descriptors(corpus, index_version=1) + _publish(writer, 1) + # Re-index everything, so every row has a superseded version to lose. + writer.write_descriptors(corpus, index_version=2) + _publish(writer, 2) + + first = reader.search(CaptureQuery(limit=20)) + _merge(client, config) + + items = list(first.items) + cursor = first.next_cursor + while cursor is not None: + page = reader.search(CaptureQuery(limit=20, cursor=cursor)) + items.extend(page.items) + cursor = page.next_cursor + + assert len(items) == len(corpus), "a merge mid-walk dropped rows" + + +# --- batch atomicity --------------------------------------------------------- + + +def test_a_watermark_is_not_published_before_its_batch_completes(): + """One logical batch must become visible all at once. + + CatalogIndexer assigns one index_version and then writes descriptors across + several INSERTs. A reader that samples between them pins a half-written + batch, and the same watermark then returns more rows later. + """ + corpus = synthetic_descriptors(4) + with _catalog() as (writer, reader, client, config): + # Two INSERTs at one version, as max_rows_per_insert would produce for + # any batch larger than the insert size. + writer.write_descriptors(corpus[:2], index_version=123) + mid_batch = reader.current_watermark() + rows_visible = _raw_rows(client, config) + writer.write_descriptors(corpus[2:], index_version=123) + _publish(writer, 123, rows=len(corpus)) + + # Mid-batch the version must not be readable at all: two of four rows + # were durable, so publishing 123 there would pin a snapshot that keeps + # growing under the caller. + assert int(mid_batch) < 123, ( + f"watermark {mid_batch} was readable with only {rows_visible} of " + f"{len(corpus)} rows durable" + ) + + resolved = reader.search( + CaptureQuery(limit=10, tenant_id=corpus[0].metadata.tenant_id) + ) + assert len(resolved.items) == len(corpus) + + +def test_an_indexed_batch_is_visible_all_at_once(): + """A watermark taken after indexing must see the whole batch.""" + corpus = synthetic_descriptors(30) + with _catalog() as (writer, reader, client, config): + for start in range(0, len(corpus), 7): + writer.write_descriptors(corpus[start : start + 7], index_version=9) + _publish(writer, 9, rows=len(corpus)) + + page = reader.search(CaptureQuery(limit=100)) + + assert len(page.items) == len(corpus) + assert page.watermark == "9" + + +# --- ordering and boundary conditions ---------------------------------------- + + +def test_a_watermark_below_every_row_returns_nothing(): + corpus = synthetic_descriptors(5) + with _catalog() as (writer, reader, _, _): + writer.write_descriptors(corpus, index_version=10) + _publish(writer, 10) + + assert reader.get_by_ids([corpus[0].capture_id], watermark="9") == () + + +def test_an_empty_catalog_reports_a_zero_watermark(): + with _catalog() as (_, reader, _, _): + assert reader.current_watermark() == "0" + assert reader.search(CaptureQuery(limit=10)).items == () From 5bb57dbec467151e5e35d80eb11e7c4b7f139eb4 Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Wed, 26 Aug 2026 01:08:26 -0400 Subject: [PATCH 19/20] fix(capture): bound snapshots by committed packs, not descriptor versions Second blocker from Samfisheryu: a pinned watermark stopped resolving after a merge. Reproduced on 26.9.1 -- BEFORE merge: raw rows = 2, watermark 1 -> [(capture-a, 64)] AFTER merge: raw rows = 1, watermark 1 -> [] argMax(...) WHERE index_version <= W reads rows that ReplacingMergeTree is defined to delete: it keeps only the highest version per sorting key. So the snapshot had an expiry nobody controlled, and a saved cursor could stop resolving at an unpredictable time. The snapshot boundary moves to the packs committed at or before the watermark, recorded in an append-only *_pack_commit_log that no merge rewrites. A capture belongs to exactly one immutable pack, so this is the same question asked of durable data instead of collapsible data. What makes it sound is that a descriptor is derived from an immutable pack footer, so re-indexing rewrites byte-identical rows. There is no content to choose between, which is why argMax is now deduplication rather than version selection and why it no longer matters which duplicate survives a merge. Two existing live tests encoded the old semantics and were updated rather than the code bent to fit them: - test_watermark_isolates_rows_indexed_after_the_first_page gave its "later" captures the same pack_id as the corpus, so it simulated adding captures to an already-committed pack. Packs are sealed before commit, so that cannot happen; the later captures now get their own pack. - test_snapshot_returns_the_version_at_the_watermark_not_the_latest mutated a descriptor's offset between versions to prove version selection worked. That scenario is outside the contract now, so it is replaced by test_replay_is_invisible_because_it_rewrites_identical_descriptors, which asserts the invariant the design depends on and fails loudly if a future change ever makes a re-indexed descriptor differ. The limitation this trades into: the snapshot can no longer distinguish content across versions, and relies on that immutability invariant instead. The new test is the guard. --- src/dmi/storage/capture/clickhouse_catalog.py | 21 +++++ src/dmi/storage/capture/clickhouse_reader.py | 57 ++++++++---- tests/test_clickhouse_reader_live.py | 67 +++++++++----- tests/test_clickhouse_snapshot_live.py | 90 ++++++++++++++----- 4 files changed, 173 insertions(+), 62 deletions(-) diff --git a/src/dmi/storage/capture/clickhouse_catalog.py b/src/dmi/storage/capture/clickhouse_catalog.py index cde4e1dce..057544658 100644 --- a/src/dmi/storage/capture/clickhouse_catalog.py +++ b/src/dmi/storage/capture/clickhouse_catalog.py @@ -95,6 +95,7 @@ def __init__( self._pack_raw = f"{prefix}_pack_inventory_raw" self._pack_view = f"{prefix}_pack_inventory" self._watermark = f"{prefix}_index_watermark" + self._commit_log = f"{prefix}_pack_commit_log" def ensure_schema(self) -> None: database = _quoted(self._config.database) @@ -149,6 +150,18 @@ def ensure_schema(self) -> None: ) ENGINE = MergeTree ORDER BY index_version""" ) + # The snapshot boundary. A capture belongs to exactly one immutable + # pack, so "the catalog as of W" is "the packs committed at or before + # W" -- a fact about packs, not about descriptor rows. Keeping it here, + # append-only, is what lets the descriptor table stay a + # ReplacingMergeTree: every version of a descriptor row is byte + # identical, so it does not matter which one a merge keeps. + self._client.execute( + f"""CREATE TABLE IF NOT EXISTS {database}.{_quoted(self._commit_log)} ( +pack_id UUID, store_id LowCardinality(String), index_version UInt64 +) ENGINE = MergeTree ORDER BY (index_version, store_id, pack_id)""" + ) + capture_public = ", ".join(_CAPTURE_COLUMNS[:-1]) pack_public = ", ".join(_PACK_COLUMNS[:-1]) self._client.execute( @@ -223,6 +236,14 @@ def commit_packs( f"({', '.join(_PACK_COLUMNS)}) VALUES", rows, ) + # Append the same commit to the log the snapshot reads. A pack + # committed twice appends twice; the reader treats it as a set, so + # replay is absorbed without needing deduplication here. + self._client.execute( + f"INSERT INTO {self._qualified(self._commit_log)} " + "(pack_id, store_id, index_version) VALUES", + [(ref.pack_id, ref.store_id, index_version) for ref in refs], + ) def _qualified(self, table: str) -> str: return f"{_quoted(self._config.database)}.{_quoted(table)}" diff --git a/src/dmi/storage/capture/clickhouse_reader.py b/src/dmi/storage/capture/clickhouse_reader.py index b2262a4c2..d5eb9558a 100644 --- a/src/dmi/storage/capture/clickhouse_reader.py +++ b/src/dmi/storage/capture/clickhouse_reader.py @@ -1,16 +1,28 @@ """ClickHouse-backed catalog reads for bounded analysis. -The writer in :mod:`.clickhouse_catalog` lands descriptor rows into a -``ReplacingMergeTree`` keyed by ``index_version``. This module reads them back -as a *snapshot*: every result is pinned to a watermark so a selection resolves -to the same captures for as long as it lives. - -That pinning is why these queries read ``*_capture_raw`` directly instead of the -public ``FINAL`` views. ``FINAL`` collapses duplicates to the highest version -present and only then applies predicates, so ``FINAL`` with -``index_version <= W`` drops any capture re-indexed above ``W`` rather than -falling back to its version at ``W``. Explicit ``argMax`` aggregation gives the -row as of the watermark, which is what a stable cursor walk requires. +Reads are pinned to a watermark, so a selection resolves to the same captures +for as long as it lives. The snapshot boundary is **the set of packs committed +at or before that watermark**, read from an append-only commit log -- not a +range of descriptor versions. + +That distinction is load-bearing. Descriptor rows live in a +``ReplacingMergeTree``, which is defined to keep only the highest version per +sorting key and delete the rest during a merge. Any snapshot phrased as +``index_version <= W`` over those rows therefore expires at a time nobody +controls: after a merge, the version it wanted is simply gone. Bounding on pack +commits instead depends only on an append-only log, which no merge rewrites. + +The reason this works is that a descriptor is derived from an immutable pack +footer, so re-indexing a pack rewrites byte-identical rows. There is no content +to choose between, which is why ``argMax`` here is deduplication rather than +version selection, and why it does not matter which duplicate a merge keeps. +``test_replay_is_invisible_because_it_rewrites_identical_descriptors`` guards +that invariant; if it ever breaks, this design has to be revisited. + +The watermark itself comes from a published log rather than from +``max(index_version)`` over the descriptors, because one indexing call writes +descriptors across several INSERTs before the pack markers -- sampling the +descriptor table mid-call pins a batch that then keeps growing. """ from __future__ import annotations @@ -124,6 +136,7 @@ def __init__( self._config = config or ClickHouseReaderConfig() self._capture_raw = f"{self._config.table_prefix}_capture_raw" self._watermark_table = f"{self._config.table_prefix}_index_watermark" + self._commit_log = f"{self._config.table_prefix}_pack_commit_log" @property def config(self) -> ClickHouseReaderConfig: @@ -202,7 +215,9 @@ def get_by_ids( } sql = ( f"SELECT {self._projection()} FROM {self._qualified()} " - "WHERE index_version <= %(watermark)s AND capture_id IN %(capture_ids)s " + "WHERE capture_id IN %(capture_ids)s AND pack_id IN " + f"(SELECT pack_id FROM {self._qualified(self._commit_log)} " + "WHERE index_version <= %(watermark)s) " f"GROUP BY {', '.join(_quoted(name) for name in _SORT_KEY)}" ) rows = self._client.execute(sql, params, settings=self._config.settings) @@ -219,7 +234,9 @@ def _qualified(self, table: str | None = None) -> str: @staticmethod def _projection() -> str: # Sort-key columns are grouped on, so they project directly; everything - # else resolves to its value at the highest version <= the watermark. + # else collapses whatever duplicate rows survive. Any version of a + # descriptor is byte identical to any other, so argMax here is + # deduplication rather than version selection. # # Deliberately unaliased: naming an aggregate after its own source # column shadows that column everywhere else in the statement, and @@ -233,11 +250,19 @@ def _projection() -> str: for name in _PROJECTION ) - @staticmethod def _filters( - query: CaptureQuery, *, watermark: int, after: CursorKey | None + self, query: CaptureQuery, *, watermark: int, after: CursorKey | None ) -> tuple[list[str], dict[str, object]]: - clauses = ["index_version <= %(watermark)s"] + # The snapshot is the set of packs committed at or before the + # watermark, not a range of descriptor versions. Descriptor rows are + # byte identical across re-indexing, so which version survives a merge + # is irrelevant -- but a version *range* over them is not durable, + # because ReplacingMergeTree deletes superseded rows. + clauses = [ + "pack_id IN (SELECT pack_id FROM " + f"{self._qualified(self._commit_log)} " + "WHERE index_version <= %(watermark)s)" + ] params: dict[str, object] = {"watermark": watermark} # Equality and range filters apply to raw rows before grouping. That is diff --git a/tests/test_clickhouse_reader_live.py b/tests/test_clickhouse_reader_live.py index 4753ebb35..de43947db 100644 --- a/tests/test_clickhouse_reader_live.py +++ b/tests/test_clickhouse_reader_live.py @@ -28,6 +28,17 @@ pytestmark = [pytest.mark.manual, pytest.mark.clickhouse] +def _commit(writer, descriptors, index_version: int) -> None: + """Snapshots are bounded by committed packs, so a direct write must commit.""" + seen, refs = set(), [] + for item in descriptors: + ref = item.locator.pack_ref + if (ref.store_id, ref.pack_id) not in seen: + seen.add((ref.store_id, ref.pack_id)) + refs.append(ref) + writer.commit_packs(refs, index_version=index_version) + + def _publish(writer, index_version: int, *, rows: int = 0, packs: int = 0) -> None: """Publishing is a separate step; CatalogIndexer does it, direct writes must.""" writer.publish_watermark( @@ -90,6 +101,7 @@ def test_a_full_walk_returns_the_corpus_exactly(): corpus = synthetic_descriptors(250) with _catalog() as (writer, reader): writer.write_descriptors(corpus, index_version=1) + _commit(writer, corpus, 1) _publish(writer, 1) walked, pages, _ = _walk(reader, CaptureQuery(limit=40)) @@ -101,16 +113,25 @@ def test_a_full_walk_returns_the_corpus_exactly(): def test_watermark_isolates_rows_indexed_after_the_first_page(): corpus = synthetic_descriptors(100) + # Its own pack: a pack is sealed before it is committed, so captures never + # appear inside one that the catalog already knows about. + late_pack = str(uuid4()) later = tuple( - replace(item, metadata=replace(item.metadata, capture_id=f"late-{index}")) + replace( + item, + metadata=replace(item.metadata, capture_id=f"late-{index}"), + locator=replace(item.locator, pack_id=late_pack), + ) for index, item in enumerate(synthetic_descriptors(50)) ) with _catalog() as (writer, reader): writer.write_descriptors(corpus, index_version=1) + _commit(writer, corpus, 1) _publish(writer, 1) first = reader.search(CaptureQuery(limit=40)) writer.write_descriptors(later, index_version=2) + _commit(writer, later, 2) _publish(writer, 2) rest, _, _ = _walk(reader, CaptureQuery(limit=40, cursor=first.next_cursor)) @@ -120,33 +141,32 @@ def test_watermark_isolates_rows_indexed_after_the_first_page(): assert not any(item.capture_id.startswith("late-") for item in walked) -def test_snapshot_returns_the_version_at_the_watermark_not_the_latest(): - """The case plain ``FINAL`` gets wrong. +def test_replay_is_invisible_because_it_rewrites_identical_descriptors(): + """Why the snapshot needs no version selection. - ``FINAL`` collapses to the highest version present and only then applies - the predicate, so a capture re-indexed above the watermark disappears - instead of falling back. ``argMax`` returns its value as of the watermark. + Descriptors are derived from an immutable pack footer, so re-indexing a + pack writes byte-identical rows. That invariant is what lets the snapshot + be "packs committed at or before W" and lets a merge collapse duplicate + descriptor rows freely -- there is no content to choose between. + + If a future change ever makes a re-indexed descriptor differ from the + original, this test fails, and the snapshot design has to be revisited. """ - original = synthetic_descriptors(1) - revised = tuple( - replace(item, locator=replace(item.locator, offset=999_999)) - for item in original - ) + corpus = synthetic_descriptors(5) with _catalog() as (writer, reader): - writer.write_descriptors(original, index_version=1) + writer.write_descriptors(corpus, index_version=1) + _commit(writer, corpus, 1) _publish(writer, 1) - writer.write_descriptors(revised, index_version=2) - _publish(writer, 2) - - at_first = reader.search(CaptureQuery(limit=10)) - assert at_first.watermark == "2" - assert at_first.items[0].locator.offset == 999_999 + at_first = reader.get_by_ids([corpus[0].capture_id], watermark="1") - pinned = reader.get_by_ids( - [original[0].capture_id], watermark="1" + writer.write_descriptors(corpus, index_version=2) + _commit(writer, corpus, 2) + _publish(writer, 2) + at_second = reader.get_by_ids( + [corpus[0].capture_id], watermark=reader.current_watermark() ) - assert len(pinned) == 1 - assert pinned[0].locator.offset == original[0].locator.offset + + assert at_first == at_second, "a replay changed a descriptor" def test_replayed_indexing_yields_one_logical_row(): @@ -154,6 +174,7 @@ def test_replayed_indexing_yields_one_logical_row(): with _catalog() as (writer, reader): for version in (1, 2, 3): writer.write_descriptors(corpus, index_version=version) + _commit(writer, corpus, version) _publish(writer, version) walked, _, _ = _walk(reader, CaptureQuery(limit=10)) @@ -165,6 +186,7 @@ def test_get_by_ids_resolves_a_selection_at_its_watermark(): corpus = synthetic_descriptors(20) with _catalog() as (writer, reader): writer.write_descriptors(corpus, index_version=1) + _commit(writer, corpus, 1) _publish(writer, 1) page = reader.search(CaptureQuery(limit=20)) @@ -179,6 +201,7 @@ def test_filters_narrow_the_walk(): corpus = synthetic_descriptors(30) with _catalog() as (writer, reader): writer.write_descriptors(corpus, index_version=1) + _commit(writer, corpus, 1) _publish(writer, 1) walked, _, _ = _walk( diff --git a/tests/test_clickhouse_snapshot_live.py b/tests/test_clickhouse_snapshot_live.py index 9dfb8f4e1..7ca6444de 100644 --- a/tests/test_clickhouse_snapshot_live.py +++ b/tests/test_clickhouse_snapshot_live.py @@ -22,6 +22,7 @@ from contextlib import contextmanager from os import environ +from dataclasses import replace from uuid import uuid4 import pytest @@ -39,6 +40,17 @@ pytestmark = [pytest.mark.manual, pytest.mark.clickhouse] +def _commit(writer, descriptors, index_version: int) -> None: + """Snapshots are bounded by committed packs, so a direct write must commit.""" + seen, refs = set(), [] + for item in descriptors: + ref = item.locator.pack_ref + if (ref.store_id, ref.pack_id) not in seen: + seen.add((ref.store_id, ref.pack_id)) + refs.append(ref) + writer.commit_packs(refs, index_version=index_version) + + def _publish(writer, index_version: int, *, rows: int = 0, packs: int = 0) -> None: """Publishing is a separate step; CatalogIndexer does it, direct writes must.""" writer.publish_watermark( @@ -100,46 +112,71 @@ def _raw_rows(client, config) -> int: # --- merges versus pinned watermarks ----------------------------------------- -@pytest.mark.xfail( - strict=True, - reason=( - "Open design problem: argMax(...) WHERE index_version <= W needs rows " - "ReplacingMergeTree is defined to delete on merge, so a pinned snapshot " - "expires at a time nobody controls. Needs a snapshot mechanism that does " - "not depend on superseded versions surviving." - ), -) def test_a_merge_does_not_destroy_a_pinned_snapshot(): """A watermark must keep resolving after ClickHouse deduplicates. - ReplacingMergeTree keeps only the highest version per sorting key, so any - snapshot that reads superseded versions has an expiry no caller controls. + The snapshot is the set of packs committed at or before the watermark, so + it survives a merge collapsing duplicate descriptor rows. Replay writes + byte-identical descriptors, which is why it does not matter which copy the + merge keeps. """ - original = synthetic_descriptors(1) - revised = tuple( - d.__class__(metadata=d.metadata, locator=d.locator.__class__( - **{**{f: getattr(d.locator, f) for f in d.locator.__slots__}, - "offset": 999_999})) - for d in original - ) + corpus = synthetic_descriptors(3) with _catalog() as (writer, reader, client, config): - writer.write_descriptors(original, index_version=1) + writer.write_descriptors(corpus, index_version=1) + _commit(writer, corpus, 1) _publish(writer, 1) - writer.write_descriptors(revised, index_version=2) + + # Replay the identical batch, as an ambiguous commit would. + writer.write_descriptors(corpus, index_version=2) + _commit(writer, corpus, 2) _publish(writer, 2) - before = reader.get_by_ids([original[0].capture_id], watermark="1") + before = reader.get_by_ids([corpus[0].capture_id], watermark="1") assert len(before) == 1, "precondition: the pinned read works pre-merge" - assert before[0].locator.offset == original[0].locator.offset _merge(client, config) - after = reader.get_by_ids([original[0].capture_id], watermark="1") + after = reader.get_by_ids([corpus[0].capture_id], watermark="1") assert len(after) == 1, ( "a pinned watermark stopped resolving after a merge: " f"raw rows went to {_raw_rows(client, config)}" ) - assert after[0].locator.offset == original[0].locator.offset + assert after[0] == before[0], "the resolved descriptor changed under a merge" + + +def test_a_pack_committed_after_the_watermark_is_not_visible(): + """The snapshot boundary has to actually exclude later work.""" + early = synthetic_descriptors(2) + # A genuinely later pack: its own pack id, and its own captures. A capture + # id belongs to exactly one pack, so reusing ids across packs would be an + # invalid corpus rather than a harder test. + late_pack = str(uuid4()) + later = tuple( + replace( + item, + metadata=replace(item.metadata, capture_id=f"late-{index}"), + locator=replace(item.locator, pack_id=late_pack), + ) + for index, item in enumerate(synthetic_descriptors(2)) + ) + with _catalog() as (writer, reader, client, config): + writer.write_descriptors(early, index_version=1) + _commit(writer, early, 1) + _publish(writer, 1) + pinned = reader.current_watermark() + + writer.write_descriptors(later, index_version=2) + _commit(writer, later, 2) + _publish(writer, 2) + _merge(client, config) + + # The pinned snapshot sees the early pack and nothing after it. + assert len(reader.get_by_ids([early[0].capture_id], watermark=pinned)) == 1 + assert reader.get_by_ids([later[0].capture_id], watermark=pinned) == () + + # And the later watermark sees both. + current = reader.current_watermark() + assert len(reader.get_by_ids([later[0].capture_id], watermark=current)) == 1 def test_a_walk_still_completes_after_a_merge_mid_pagination(): @@ -147,9 +184,11 @@ def test_a_walk_still_completes_after_a_merge_mid_pagination(): corpus = synthetic_descriptors(60) with _catalog() as (writer, reader, client, config): writer.write_descriptors(corpus, index_version=1) + _commit(writer, corpus, 1) _publish(writer, 1) # Re-index everything, so every row has a superseded version to lose. writer.write_descriptors(corpus, index_version=2) + _commit(writer, corpus, 2) _publish(writer, 2) first = reader.search(CaptureQuery(limit=20)) @@ -183,6 +222,7 @@ def test_a_watermark_is_not_published_before_its_batch_completes(): mid_batch = reader.current_watermark() rows_visible = _raw_rows(client, config) writer.write_descriptors(corpus[2:], index_version=123) + _commit(writer, corpus[2:], 123) _publish(writer, 123, rows=len(corpus)) # Mid-batch the version must not be readable at all: two of four rows @@ -205,6 +245,7 @@ def test_an_indexed_batch_is_visible_all_at_once(): with _catalog() as (writer, reader, client, config): for start in range(0, len(corpus), 7): writer.write_descriptors(corpus[start : start + 7], index_version=9) + _commit(writer, corpus[start : start + 7], 9) _publish(writer, 9, rows=len(corpus)) page = reader.search(CaptureQuery(limit=100)) @@ -220,6 +261,7 @@ def test_a_watermark_below_every_row_returns_nothing(): corpus = synthetic_descriptors(5) with _catalog() as (writer, reader, _, _): writer.write_descriptors(corpus, index_version=10) + _commit(writer, corpus, 10) _publish(writer, 10) assert reader.get_by_ids([corpus[0].capture_id], watermark="9") == () From fd881045e6fe3a3ad14049e8173ad5916ca4192a Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Wed, 26 Aug 2026 01:45:24 -0400 Subject: [PATCH 20/20] fix(capture): close five defects found by review, with the coverage that missed them Each fix comes with a test that names the gap in the existing suites, because the pattern matters more than the individual bug: all five sit just outside a boundary the tests already exercised. Object keys. quote() treats "~" as always-safe per RFC 3986 and ignores its `safe` argument for it, but the object-key pattern rejects "~". A tenant or session id containing one produced a key every store refuses, the sink raised, and the persistence thread treated that as fatal -- so one tenant name took down capture for every tenant. The gap: key generation and key validation were tested separately and never against each other. Now parametrised over every character the encoder passes through unchanged. Catalog commit ordering. commit_packs wrote the inventory before the commit log, but committed_pack_ids reads the inventory to skip replays while readers bound the snapshot by the log. A crash between the two writes left a pack skipped forever and never visible -- silent, permanent loss. The log is now written first, so the same crash only costs redundant work. The gap: the fault tests failed whole operations, never the window between two writes that must agree. Reconciliation. index_object_keys inspected every key outside CatalogIndexer's per-pack failure handling, so one foreign object in the bucket aborted a whole rebuild instead of being recorded as one failure. A bucket holds whatever anyone put in it. The gap: reconciliation was only ever tested over buckets containing valid packs. Version monotonicity. index_version comes from a wall clock, which steps backwards on NTP correction and can repeat on a coarse clock, letting a batch land under a watermark a reader already pinned -- the same defect the published watermark fixed for the mid-batch case. The indexer now advances past its last published version rather than trusting the clock. It advances rather than raising: a version is an ordering token, not a timestamp, and refusing would stop capture over a clock correction. published_at_ns still records real time, and cross-process skew remains documented under Phase 5 limitations. Search benchmark. It wrote descriptors but never committed packs or published a watermark, so after the snapshot change every search returned zero rows and the page-latency table in docs/benchmarks.md was no longer reproducible. It also leaked the two new tables on cleanup. Both fixed; verified it returns rows again and drops everything it creates. Two existing tests changed with the behaviour rather than around it: the commit-order assertion now expects the log first and says why, and the backwards-clock test asserts published versions stay strictly increasing instead of expecting a refusal. --- benchmarks/bench_capture_search.py | 17 + src/dmi/storage/capture/catalog.py | 47 ++- src/dmi/storage/capture/clickhouse_catalog.py | 18 +- src/dmi/storage/capture/pipeline.py | 6 +- tests/test_capture_review_findings.py | 296 ++++++++++++++++++ tests/test_clickhouse_capture_catalog.py | 6 +- 6 files changed, 379 insertions(+), 11 deletions(-) create mode 100644 tests/test_capture_review_findings.py diff --git a/benchmarks/bench_capture_search.py b/benchmarks/bench_capture_search.py index 22123e27a..387b59e5d 100644 --- a/benchmarks/bench_capture_search.py +++ b/benchmarks/bench_capture_search.py @@ -189,8 +189,24 @@ def main(argv=None) -> int: writer.ensure_schema() try: descriptors = synthetic_descriptors(args.rows) + # Descriptors alone are not readable. A snapshot is bounded by committed + # packs and the watermark comes from the published log, so a benchmark + # that only writes descriptors measures empty result sets. + refs, seen = [], set() + for item in descriptors: + ref = item.locator.pack_ref + if (ref.store_id, ref.pack_id) not in seen: + seen.add((ref.store_id, ref.pack_id)) + refs.append(ref) for version in range(1, args.replays + 1): writer.write_descriptors(descriptors, index_version=version) + writer.commit_packs(refs, index_version=version) + writer.publish_watermark( + index_version=version, + published_at_ns=version, + indexed_rows=len(descriptors), + indexed_packs=len(refs), + ) result = { "rows": args.rows, @@ -211,6 +227,7 @@ def main(argv=None) -> int: for kind, suffix in ( ("VIEW", "capture"), ("VIEW", "pack_inventory"), ("TABLE", "capture_raw"), ("TABLE", "pack_inventory_raw"), + ("TABLE", "index_watermark"), ("TABLE", "pack_commit_log"), ): client.execute( f"DROP {kind} IF EXISTS `{args.database}`.`{prefix}_{suffix}`" diff --git a/src/dmi/storage/capture/catalog.py b/src/dmi/storage/capture/catalog.py index cb883c169..548838f34 100644 --- a/src/dmi/storage/capture/catalog.py +++ b/src/dmi/storage/capture/catalog.py @@ -159,6 +159,7 @@ def __init__( self._timer_ns = timer_ns self._on_event = on_event self._callback_failures = 0 + self._published_version: int | None = None @property def store_id(self) -> str: @@ -221,6 +222,16 @@ def index(self, refs: Sequence[PackRef]) -> IndexResult: version = self._clock_ns() if type(version) is not int or version < 0: raise ValueError("clock_ns must return a non-negative integer") + # index_version has to increase strictly, or a batch lands underneath a + # watermark a reader already pinned and that snapshot grows after the + # fact. A wall clock does not guarantee that: NTP steps backwards, and + # a coarse clock can return the same value twice in a row. Advance past + # the last published version rather than failing -- the version is an + # ordering token, not a timestamp, and published_at_ns still records the + # real time. Cross-process skew is a separate problem, documented under + # the Phase 5 limitations. + if self._published_version is not None and version <= self._published_version: + version = self._published_version + 1 step = self._config.max_rows_per_insert descriptor_inserts = 0 for start in range(0, len(descriptors), step): @@ -234,6 +245,7 @@ def index(self, refs: Sequence[PackRef]) -> IndexResult: # markers after them, so the version is only a truthful snapshot once # all of that is durable. A reader that derived the watermark from the # descriptor table itself would see this version mid-batch. + self._published_version = version publish = getattr(self._writer, "publish_watermark", None) if publish is not None: publish( @@ -288,7 +300,40 @@ def index_object_keys(self, object_keys: Sequence[str]) -> IndexResult: if len(object_keys) > self._indexer.max_packs: raise ValueError("object notification batch exceeds max_packs") keys = tuple(dict.fromkeys(object_keys)) - return self._indexer.index([self._inventory.inspect(key) for key in keys]) + # A bucket holds whatever anyone put in it. Inspection runs outside + # CatalogIndexer.index's per-pack handling, so without this one foreign + # object would abort an entire rebuild instead of being one failure. + refs: list[PackRef] = [] + failures: list[IndexFailure] = [] + for key in keys: + try: + refs.append(self._inventory.inspect(key)) + except Exception as exc: + failures.append( + IndexFailure( + pack_id="", + object_key=key, + error_type=type(exc).__name__, + message=str(exc)[:512], + ) + ) + result = self._indexer.index(refs) + if not failures: + return result + rejected = IndexResult( + requested_packs=len(failures), + skipped_packs=0, + indexed_packs=0, + indexed_rows=0, + failed_packs=len(failures), + descriptor_inserts=0, + estimated_bytes=0, + elapsed_ns=0, + failures=tuple(failures), + ) + return result.merge( + rejected, failure_limit=self._indexer.max_failure_details + ) def reconcile_page( self, *, prefix: str = "", cursor: str | None = None, limit: int = 64 diff --git a/src/dmi/storage/capture/clickhouse_catalog.py b/src/dmi/storage/capture/clickhouse_catalog.py index 057544658..93ccef5bc 100644 --- a/src/dmi/storage/capture/clickhouse_catalog.py +++ b/src/dmi/storage/capture/clickhouse_catalog.py @@ -231,19 +231,21 @@ def commit_packs( ) for ref in refs ] - self._client.execute( - f"INSERT INTO {self._qualified(self._pack_raw)} " - f"({', '.join(_PACK_COLUMNS)}) VALUES", - rows, - ) - # Append the same commit to the log the snapshot reads. A pack - # committed twice appends twice; the reader treats it as a set, so - # replay is absorbed without needing deduplication here. + # Order matters. committed_pack_ids reads the inventory to skip replays, + # and readers bound the snapshot by the commit log. Writing the + # inventory first would mean a crash in between leaves the pack skipped + # forever *and* never visible -- silent, permanent loss. This way round + # the same crash only costs redundant work on the next pass. self._client.execute( f"INSERT INTO {self._qualified(self._commit_log)} " "(pack_id, store_id, index_version) VALUES", [(ref.pack_id, ref.store_id, index_version) for ref in refs], ) + self._client.execute( + f"INSERT INTO {self._qualified(self._pack_raw)} " + f"({', '.join(_PACK_COLUMNS)}) VALUES", + rows, + ) def _qualified(self, table: str) -> str: return f"{_quoted(self._config.database)}.{_quoted(table)}" diff --git a/src/dmi/storage/capture/pipeline.py b/src/dmi/storage/capture/pipeline.py index 4db5d8be8..b5f081484 100644 --- a/src/dmi/storage/capture/pipeline.py +++ b/src/dmi/storage/capture/pipeline.py @@ -279,7 +279,11 @@ def object_key_for(ready: ReadyPack) -> str: def _key_component(value: str) -> str: - encoded = quote(value, safe="-_.=") + # quote() treats "~" as always-safe per RFC 3986 and ignores `safe` for it, + # but the object-key pattern does not allow it. Left alone, an identifier + # containing "~" produces a key every store rejects, and the sink failure + # is fatal to the whole pipeline. + encoded = quote(value, safe="-_.=").replace("~", "%7E") if len(encoded.encode()) <= 160: return encoded return "sha256-" + sha256(value.encode()).hexdigest() diff --git a/tests/test_capture_review_findings.py b/tests/test_capture_review_findings.py new file mode 100644 index 000000000..ea398f72c --- /dev/null +++ b/tests/test_capture_review_findings.py @@ -0,0 +1,296 @@ +"""Regressions found by review that the existing suites could not reach. + +Each test names the gap in coverage that let the bug through, because the +pattern matters more than the individual defect: every one of these sits just +outside a boundary the tests already exercised. +""" + +from __future__ import annotations + +from pathlib import Path +from urllib.parse import quote +from uuid import UUID + +import pytest + +from tests._faults import FaultInjected, FaultyClickHouseClient, fail_on + +from dmi.storage.capture import ( + AdmissionResult, + CaptureMetadata, + CaptureRecord, + CatalogIndexer, + CatalogIndexerConfig, + CatalogReconciler, + ClickHouseCatalogConfig, + ClickHouseCatalogWriter, + DirectPackSink, + FilesystemPackStore, + HostCapturePipeline, + PackRef, + PackWriter, + PipelineConfig, + object_key_for, +) +from dmi.storage.capture.filesystem import validate_object_key + + +pytestmark = pytest.mark.cpu + + +PACK_ID = UUID("018f0000-0000-7000-8000-00000000ab01") + + +def _metadata(**overrides) -> CaptureMetadata: + base = dict( + capture_id="capture-a", + tenant_id="tenant-a", + experiment_id="exp-a", + run_id="run-a", + session_id="session-a", + request_id="request-a", + sequence_id="sequence-a", + model_id="model-a", + model_revision="revision-a", + adapter_revision=None, + capture_policy_version="policy-v1", + hook_name="resid_pre", + layer_number=3, + producer_rank=0, + step_number=0, + token_start=0, + token_end=1, + batch_position=0, + dtype="float32", + shape=(2,), + captured_at_ns=1_700_000_000_000_000_000, + ) + base.update(overrides) + return CaptureMetadata(**base) + + +def _record(metadata: CaptureMetadata) -> CaptureRecord: + return CaptureRecord(metadata=metadata, payload=b"\x01" * 8) + + +def _sealed(*records: CaptureRecord, pack_id: UUID = PACK_ID): + writer = PackWriter( + pack_id=pack_id, + created_at_ns=1_700_000_000_000_000_000, + max_pack_bytes=1024 * 1024, + ) + for record in records: + writer.append(record) + return writer.seal() + + +def _ids(): + from itertools import count + + counter = count(1) + return lambda: UUID(int=next(counter)) + + +# --- object keys ------------------------------------------------------------- +# +# Gap: key generation and key validation were tested separately, never against +# each other, so a character one produced and the other refused went unnoticed. + + +@pytest.mark.parametrize( + "character", + [c for c in map(chr, range(32, 127)) if quote(c, safe="-_.=") == c], + ids=lambda c: f"U+{ord(c):04X}", +) +def test_every_character_key_encoding_passes_through_is_accepted(character: str): + """Whatever the encoder leaves alone, the validator must accept. + + `quote` treats `~` as always-safe per RFC 3986 regardless of its `safe` + argument, so it survives encoding -- but the key pattern rejects it. + """ + metadata = _metadata(tenant_id=f"acme{character}lab") + + from dmi.storage.capture.pipeline import _key_component + + encoded = _key_component(metadata.tenant_id) + validate_object_key(f"tenant={encoded}/x.dmi-pack") + + +def test_an_identifier_needing_no_escape_still_yields_a_usable_key(tmp_path: Path): + """The end-to-end consequence: an unencodable id kills the pipeline. + + The key is rejected by the store, the sink raises, and the persistence + thread treats that as fatal -- so one tenant name takes down capture for + every tenant. + """ + store = FilesystemPackStore(tmp_path / "objects", store_id="local") + pipeline = HostCapturePipeline( + PipelineConfig( + max_queue_records=4, + max_queue_bytes=1 << 20, + max_pack_bytes=1024 * 1024, + max_pack_records=1, + max_linger_ns=1_000_000_000, + ), + DirectPackSink(store), + pack_id_factory=_ids(), + ) + pipeline.start() + + assert ( + pipeline.submit(_record(_metadata(tenant_id="acme~lab"))) + is AdmissionResult.ACCEPTED + ) + snapshot = pipeline.close(timeout=2) + + assert snapshot.failures == 0, "an unencodable identifier failed the pipeline" + assert snapshot.persisted_records == 1 + + +# --- catalog commit ordering ------------------------------------------------- +# +# Gap: the fault tests failed whole operations, never the gap *between* two +# writes that must agree. + + +class _Client: + def __init__(self): + self.statements: list[str] = [] + self.published: list[tuple[str, list]] = [] + + def execute(self, query, params=None, **kwargs): + self.statements.append(query) + if query.lstrip().upper().startswith("INSERT"): + self.published.append((query, list(params or []))) + return [] + + +def test_a_pack_is_never_both_skipped_on_replay_and_invisible_to_readers(): + """The durability window between the two commit writes. + + `committed_pack_ids` reads the inventory to skip replays; readers bound the + snapshot by the commit log. If the inventory is written first and the + process dies before the log, the pack is skipped forever *and* never + visible -- silent, permanent data loss. Writing the log first makes the + same crash merely redundant work. + """ + backing = _Client() + # Fail the second of the two INSERTs that commit_packs performs. + client = FaultyClickHouseClient(backing, insert=fail_on(2)) + writer = ClickHouseCatalogWriter(client, ClickHouseCatalogConfig()) + ref = PackRef( + pack_id=str(PACK_ID), + store_id="local", + object_key="packs/a.dmi-pack", + object_bytes=1024, + checksum="0" * 64, + record_count=1, + ) + + with pytest.raises(FaultInjected): + writer.commit_packs([ref], index_version=1) + + written = [s for s in backing.statements if s.startswith("INSERT")] + assert len(written) == 1, "expected exactly one of the two writes to land" + assert "pack_commit_log" in written[0], ( + "the surviving write was the inventory, so this pack is now skipped on " + "replay and invisible to readers; the commit log must be written first" + ) + + +# --- indexer robustness ------------------------------------------------------ +# +# Gap: reconciliation was tested over buckets containing only valid packs. + + +class _Inventory: + """A store whose bucket contains one object that is not a pack.""" + + store_id = "local" + + def __init__(self, store, refs, bad_key: str): + self._store = store + self._refs = {ref.object_key: ref for ref in refs} + self._bad_key = bad_key + + def inspect(self, object_key: str) -> PackRef: + if object_key == self._bad_key: + raise ValueError(f"not a dmi pack: {object_key}") + return self._refs[object_key] + + def read_range(self, ref, offset, length): + return self._store.read_range(ref, offset, length) + + def put(self, pack, object_key): + return self._store.put(pack, object_key) + + def stat(self, ref): + return self._store.stat(ref) + + +def test_one_foreign_object_in_the_bucket_does_not_abort_reconciliation( + tmp_path: Path, +): + """A bucket holds whatever anyone put there. + + `index_object_keys` inspects every key before handing the batch to the + indexer, outside its per-pack failure handling, so a single unreadable + object aborts the whole rebuild instead of being recorded as one failure. + """ + store = FilesystemPackStore(tmp_path, store_id="local") + good = store.put(_sealed(_record(_metadata())), "packs/good.dmi-pack") + inventory = _Inventory(store, [good], bad_key="packs/README.txt") + indexer = CatalogIndexer(inventory, ClickHouseCatalogWriter(_Client(), + ClickHouseCatalogConfig()), + config=CatalogIndexerConfig(max_packs=8), + clock_ns=lambda: 7) + reconciler = CatalogReconciler(inventory, indexer) + + result = reconciler.index_object_keys(["packs/good.dmi-pack", "packs/README.txt"]) + + assert result.indexed_packs == 1, "the valid pack should still be indexed" + assert result.failed_packs == 1 + assert "README" in result.failures[0].object_key + + +# --- version monotonicity ---------------------------------------------------- +# +# Gap: every indexer test used a fixed or increasing clock. + + +def test_a_clock_that_steps_backwards_cannot_publish_under_a_pinned_watermark( + tmp_path: Path, +): + """index_version is wall-clock, and wall clocks move backwards. + + An NTP correction (or a second indexer with a skewed clock) lets a newer + batch take a version below one a reader already pinned, so rows appear + inside a snapshot that was taken before they existed -- the same defect the + published-watermark change fixed for the mid-batch case. + """ + store = FilesystemPackStore(tmp_path, store_id="local") + first = store.put(_sealed(_record(_metadata())), "packs/first.dmi-pack") + second = store.put( + _sealed(_record(_metadata(capture_id="capture-b")), pack_id=UUID(int=2)), + "packs/second.dmi-pack", + ) + + client = _Client() + writer = ClickHouseCatalogWriter(client, ClickHouseCatalogConfig()) + clock = iter([2_000, 2_001, 1_000, 1_001]) # second batch stamped *earlier* + indexer = CatalogIndexer(store, writer, clock_ns=lambda: next(clock)) + + first_result = indexer.index([first]) + second_result = indexer.index([second]) + + # Both batches must still be indexed -- refusing would stop capture over a + # clock correction -- but the second must not reuse or undercut a version a + # reader may already have pinned. + assert first_result.indexed_packs == 1 and second_result.indexed_packs == 1 + published = [ + params[0][0] + for query, params in client.published + if "index_watermark" in query + ] + assert published == sorted(published), f"versions went backwards: {published}" + assert len(set(published)) == len(published), "a version was reused" diff --git a/tests/test_clickhouse_capture_catalog.py b/tests/test_clickhouse_capture_catalog.py index b5388856a..ca1a0a724 100644 --- a/tests/test_clickhouse_capture_catalog.py +++ b/tests/test_clickhouse_capture_catalog.py @@ -96,7 +96,11 @@ def test_clickhouse_catalog_inserts_descriptors_before_pack_commit(): inserts = [call for call in client.calls if call[0].startswith("INSERT")] assert "dmi_capture_raw" in inserts[0][0] - assert "dmi_pack_inventory_raw" in inserts[1][0] + # Commit log before inventory: the inventory is the replay guard, so a crash + # between the two writes must not leave a pack skipped forever *and* never + # visible to readers. + assert "dmi_pack_commit_log" in inserts[1][0] + assert "dmi_pack_inventory_raw" in inserts[2][0] assert inserts[0][1][0][0] == "capture-a" assert inserts[1][1][0][0] == ref.pack_id