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 diff --git a/benchmarks/bench_capture_catalog.py b/benchmarks/bench_capture_catalog.py new file mode 100644 index 000000000..81edc4988 --- /dev/null +++ b/benchmarks/bench_capture_catalog.py @@ -0,0 +1,126 @@ +"""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", + ) + # 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}", + 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), + ) + 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_capture_search.py b/benchmarks/bench_capture_search.py new file mode 100644 index 000000000..387b59e5d --- /dev/null +++ b/benchmarks/bench_capture_search.py @@ -0,0 +1,239 @@ +"""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) + # 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, + "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"), + ("TABLE", "index_watermark"), ("TABLE", "pack_commit_log"), + ): + client.execute( + f"DROP {kind} IF EXISTS `{args.database}`.`{prefix}_{suffix}`" + ) + 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..66f4c86f4 100644 --- a/docs/benchmarks.md +++ b/docs/benchmarks.md @@ -99,6 +99,60 @@ 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. + +## 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 new file mode 100644 index 000000000..967f64143 --- /dev/null +++ b/docs/capture-storage-design.md @@ -0,0 +1,909 @@ +# Capture storage design + +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. +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 the core tensor summarizer, plus scalar metric and artifact extension + points. +- Add agent-safe query and hydration limits. + +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 + +- 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. + +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 + +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. 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. +- 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. + +## 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. `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 + +### 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..502e21d60 --- /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 5 bounded analysis 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..24af9d4f3 --- /dev/null +++ b/src/dmi/storage/capture/__init__.py @@ -0,0 +1,185 @@ +"""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, + InvalidCursorError, + 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 .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, + CatalogReconciler, + CatalogWriter, + IndexEvent, + IndexFailure, + IndexResult, + PackIdentity, + PackInventory, + 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 ( + DurablePackSink, + DurablePackSpool, + SpoolFullError, + SpoolSnapshot, + SpoolUploader, + StagedPack, + ParallelSpoolUploader, + ParallelUploadConfig, + UploadBatchResult, + UploadEvent, + UploadFailure, + UploadSnapshot, +) + +__all__ = [ + "ArtifactProducer", + "ArtifactRef", + "ArtifactSink", + "AdmissionResult", + "BoundedRecordQueue", + "CaptureCatalog", + "CaptureDescriptor", + "CaptureMetadata", + "CapturePage", + "CaptureQuery", + "CaptureReader", + "CaptureRecord", + "CaptureSelection", + "CaptureSummary", + "CaptureStorageError", + "CatalogIndexer", + "CatalogIndexerConfig", + "CatalogReconciler", + "CatalogWriter", + "ClickHouseCatalogConfig", + "ClickHouseCatalogWriter", + "ClickHouseCaptureCatalog", + "ClickHouseReaderConfig", + "Cursor", + "CursorKey", + "CoreTensorSummaryV1", + "CORE_SUMMARY_VERSION", + "DuplicateCaptureError", + "DirectPackSink", + "DurablePackSink", + "DurablePackSpool", + "ExtensionError", + "ExtensionFailure", + "ExtensionRegistry", + "FilesystemPackStore", + "FlushReason", + "HistogramSnapshot", + "HostCapturePipeline", + "HydratedCapture", + "HydrationEstimate", + "HydrationLimitError", + "IndexEvent", + "IndexFailure", + "IndexResult", + "InvalidCursorError", + "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", + "ScalarMetric", + "S3PackStore", + "S3StoreConfig", + "SealedPack", + "SpoolFullError", + "SpoolSnapshot", + "SpoolUploader", + "StagedPack", + "StoredObject", + "UploadBatchResult", + "UploadEvent", + "UploadFailure", + "UploadSnapshot", + "decode_cursor", + "encode_cursor", + "decode_tensor", + "summarize_tensor", + "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..548838f34 --- /dev/null +++ b/src/dmi/storage/capture/catalog.py @@ -0,0 +1,362 @@ +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: ... + + 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: + 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 + self._published_version: int | None = None + + @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) + 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 + # 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) + + 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): + 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) + # 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. + self._published_version = version + 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), + 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)) + # 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 + ) -> 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..93ccef5bc --- /dev/null +++ b/src/dmi/storage/capture/clickhouse_catalog.py @@ -0,0 +1,282 @@ +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", +) + +# 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", +) + + +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" + self._watermark = f"{prefix}_index_watermark" + self._commit_log = f"{prefix}_pack_commit_log" + + 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, +{_facet_ddl()} +) 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)""" + ) + # 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}" + ) + + # 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""" + ) + + # 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( + 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 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: + 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 + ] + # 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)}" + + @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/clickhouse_reader.py b/src/dmi/storage/capture/clickhouse_reader.py new file mode 100644 index 000000000..d5eb9558a --- /dev/null +++ b/src/dmi/storage/capture/clickhouse_reader.py @@ -0,0 +1,386 @@ +"""ClickHouse-backed catalog reads for bounded analysis. + +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 + +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" + 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: + return self._config + + # -- public API --------------------------------------------------------- + + def current_watermark(self) -> str: + """The newest version whose indexing batch is fully committed. + + 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(self._watermark_table)}", + 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 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) + return tuple(self._descriptor(row) for row in rows) + + # -- SQL construction --------------------------------------------------- + + 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: + # Sort-key columns are grouped on, so they project directly; everything + # 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 + # 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 + ) + + def _filters( + self, query: CaptureQuery, *, watermark: int, after: CursorKey | None + ) -> tuple[list[str], dict[str, object]]: + # 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 + # 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/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/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/filesystem.py b/src/dmi/storage/capture/filesystem.py new file mode 100644 index 000000000..eda92b8fb --- /dev/null +++ b/src/dmi/storage/capture/filesystem.py @@ -0,0 +1,227 @@ +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 _drain_pack_source(pack: PackSource, destination: BinaryIO | None) -> 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") + if destination is not None: + 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") + + +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: + 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..33c4b239d --- /dev/null +++ b/src/dmi/storage/capture/model.py @@ -0,0 +1,401 @@ +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 +# 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 + + +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.""" + + +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")) > limit: + raise ValueError(f"{name} must be non-empty UTF-8 within {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"): + 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: + 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() + + @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: + 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 + filter_hash: str + + @classmethod + def create( + cls, + descriptors: Sequence[CaptureDescriptor], + *, + catalog_watermark: str, + filter_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": 2, + "catalog_watermark": catalog_watermark, + "filter_hash": filter_hash, + "capture_ids": ids, + }, + sort_keys=True, + separators=(",", ":"), + ).encode() + return cls( + selection_id=sha256(identity).hexdigest(), + capture_ids=ids, + catalog_watermark=catalog_watermark, + filter_hash=filter_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..b5f081484 --- /dev/null +++ b/src/dmi/storage/capture/pipeline.py @@ -0,0 +1,611 @@ +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: + # 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() + + +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 + 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, + 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 + 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 + 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..7ffb7f30f --- /dev/null +++ b/src/dmi/storage/capture/reader.py @@ -0,0 +1,278 @@ +from __future__ import annotations + +from dataclasses import dataclass +import math +from typing import Mapping, Sequence + +from .model import ( + CaptureCatalog, + CaptureDescriptor, + CapturePage, + CaptureQuery, + CaptureSelection, + DuplicateCaptureError, + HydratedCapture, + HydrationEstimate, + HydrationLimitError, + PackFormatError, + 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) +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, + filter_hash=query.filter_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 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, + 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..97d1e9fc4 --- /dev/null +++ b/src/dmi/storage/capture/s3.py @@ -0,0 +1,432 @@ +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, verify_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) + + # 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, + "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") + # 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: + 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/src/dmi/storage/capture/summary.py b/src/dmi/storage/capture/summary.py new file mode 100644 index 000000000..2343c0dda --- /dev/null +++ b/src/dmi/storage/capture/summary.py @@ -0,0 +1,209 @@ +"""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) + # 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, + 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=abs_max, + l2_norm=l2_norm, + ) + + +def numpy_dtypes() -> Mapping[str, str]: + """The dtype table, exposed so tests can assert full coverage.""" + return dict(_NUMPY_DTYPES) 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_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..de08a94cf --- /dev/null +++ b/tests/test_capture_catalog_indexer.py @@ -0,0 +1,288 @@ +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 + + +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_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) 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 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_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..6851a0830 --- /dev/null +++ b/tests/test_capture_pipeline.py @@ -0,0 +1,433 @@ +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 + + +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_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_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 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_capture_s3.py b/tests/test_capture_s3.py new file mode 100644 index 000000000..c1b2b3bb0 --- /dev/null +++ b/tests/test_capture_s3.py @@ -0,0 +1,334 @@ +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) + + +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/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_capture_summary.py b/tests/test_capture_summary.py new file mode 100644 index 000000000..9cea56b2c --- /dev/null +++ b/tests/test_capture_summary.py @@ -0,0 +1,665 @@ +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) + + +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/test_clickhouse_capture_catalog.py b/tests/test_clickhouse_capture_catalog.py new file mode 100644 index 000000000..ca1a0a724 --- /dev/null +++ b/tests/test_clickhouse_capture_catalog.py @@ -0,0 +1,170 @@ +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] + # 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 + + +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) + + + +# --- 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_capture_reader.py b/tests/test_clickhouse_capture_reader.py new file mode 100644 index 000000000..1e8b4fbaa --- /dev/null +++ b/tests/test_clickhouse_capture_reader.py @@ -0,0 +1,395 @@ +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: + # 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] + + @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_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_facets_live.py b/tests/test_clickhouse_facets_live.py new file mode 100644 index 000000000..d3d9f4920 --- /dev/null +++ b/tests/test_clickhouse_facets_live.py @@ -0,0 +1,187 @@ +"""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] + + +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") + 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),)] 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_clickhouse_reader_live.py b/tests/test_clickhouse_reader_live.py new file mode 100644 index 000000000..de43947db --- /dev/null +++ b/tests/test_clickhouse_reader_live.py @@ -0,0 +1,217 @@ +"""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] + + +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( + 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_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) + _commit(writer, corpus, 1) + _publish(writer, 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) + # 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}"), + 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)) + + walked = first.items + rest + assert walked == corpus + assert not any(item.capture_id.startswith("late-") for item in walked) + + +def test_replay_is_invisible_because_it_rewrites_identical_descriptors(): + """Why the snapshot needs no version selection. + + 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. + """ + corpus = synthetic_descriptors(5) + with _catalog() as (writer, reader): + writer.write_descriptors(corpus, index_version=1) + _commit(writer, corpus, 1) + _publish(writer, 1) + at_first = reader.get_by_ids([corpus[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 at_first == at_second, "a replay changed a descriptor" + + +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) + _commit(writer, corpus, version) + _publish(writer, 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) + _commit(writer, corpus, 1) + _publish(writer, 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) + _commit(writer, corpus, 1) + _publish(writer, 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] diff --git a/tests/test_clickhouse_snapshot_live.py b/tests/test_clickhouse_snapshot_live.py new file mode 100644 index 000000000..7ca6444de --- /dev/null +++ b/tests/test_clickhouse_snapshot_live.py @@ -0,0 +1,273 @@ +"""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 dataclasses import replace +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 _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( + 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 ----------------------------------------- + + +def test_a_merge_does_not_destroy_a_pinned_snapshot(): + """A watermark must keep resolving after ClickHouse deduplicates. + + 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. + """ + corpus = synthetic_descriptors(3) + with _catalog() as (writer, reader, client, config): + writer.write_descriptors(corpus, index_version=1) + _commit(writer, corpus, 1) + _publish(writer, 1) + + # 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([corpus[0].capture_id], watermark="1") + assert len(before) == 1, "precondition: the pinned read works pre-merge" + + _merge(client, config) + + 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] == 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(): + """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) + _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)) + _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) + _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 + # 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) + _commit(writer, corpus[start : start + 7], 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) + _commit(writer, corpus, 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 == () 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/tools/__init__.py b/tests/tools/__init__.py new file mode 100644 index 000000000..e69de29bb 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/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()) diff --git a/tests/tools/run_garage_live.py b/tests/tools/run_garage_live.py new file mode 100644 index 000000000..125c6240b --- /dev/null +++ b/tests/tools/run_garage_live.py @@ -0,0 +1,179 @@ +#!/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 _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}\n" + f"{_log_tail(log_path)}" + ) + 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"), + } + ) + # 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=log_handle, + text=True, + ) + try: + _wait_for_port(s3_port, process, log_path) + 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: + log_handle.close() + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=5) + + +if __name__ == "__main__": + raise SystemExit(main())