diff --git a/benchmarks/DESTINATION_MEMORY_REVIEW.md b/benchmarks/DESTINATION_MEMORY_REVIEW.md new file mode 100644 index 0000000..435ca0f --- /dev/null +++ b/benchmarks/DESTINATION_MEMORY_REVIEW.md @@ -0,0 +1,114 @@ +# BED-9372 destination memory & part-size review + +Acceptance criteria: peak memory and part sizes stay bounded independently of +table cardinality, against the destination's 1,000-item batch, without +unsupported upload artifacts. + +## Revision notes + +- An earlier revision called peak RSS "orthogonal to the batching fix." Wrong: + growth was linear (~1.2 KB per relationship) because each wrapper now carries + up to `batch_size` relationships. Root cause and fix below. +- Causal chain: the batching fix is the feature; the memory fix was required to + keep it viable at scale (RSS scaled linearly with cardinality); and the memory + fix exposed a latent DLT 1.26.0 data-corruption defect, which required an + in-process correction — without it, delivery silently duplicates items. +- The first mitigation attempt (buffer=333, DLT untouched) demonstrated that + defect: the faker BHE scheduler test ingested 2,875 nodes where 1,000 were + expected. Corrected in-process; all numbers below re-measured after. + +## Root cause + +DLT's writer buffer flushes on **item count** (`data_writer.buffer_max_items`, +default 5,000), and each edge wrapper counts as one item regardless of its +content. Post-batching that is up to + + 5,000 items x 150 edges = 750,000 relationships in RAM per flush. + +Measured: peak RSS 255 MB at 100k rows -> 1,257 MB at 1M rows (~1.2 KB per +relationship); the `batch_size=1` baseline stayed flat at ~130 MB. + +## DLT 1.26.0 jsonl batching defect (corrected in-process) + +`DestinationJsonlLoadJob.get_batches` (`dlt/destinations/job_impl.py:240`) +yields the accumulated batch at the end of every load-file line without +resetting it, so multi-line files re-deliver earlier items as a growing prefix: + + buffer=333, 1,200 rows -> 1,208 wrappers delivered as + 333 + 666 + 999 + 1000 + 208 = 3,206 items + +Stock settings are only accidentally safe: 5,000 divides evenly by the +destinations' `batch_size=1000`, so partial batches stay empty until each +file's last line. Any other buffer value triggers duplicates. Upstream +refactored this code in release 1.27.0 (`JsonlFileBatchIterator`, verified) +with correct semantics; `openhound.core.dlt_jsonl_batching` installs those +semantics process-wide for dlt 1.26.x, making any buffer value safe. + +## Shipped mitigation + +`writer_buffer_max_items()` scales the buffer so buffered *relationships* stay +near a fixed budget; applied via `DATA_WRITER__BUFFER_MAX_ITEMS` (`setdefault`, +user overrides win) in `Converter.pipeline` and the benchmark harness: + + buffer_max_items = min(5000, max(1, 50_000 // batch_size)) # 333 @ 150 + +More frequent flushes write to the same open file; wall time is unchanged +within noise. The jsonl batching module above is what makes 333 safe. + +## Method + +`benchmarks/opengraph_batching_benchmark.py`, one-edge shape, 4 files, DLT +1.26.0 with the batching module active, single load worker. Table-wide = +`batch_size=150`; baseline = pre-fix `batch_size=1`. Untuned rows set +`DATA_WRITER__BUFFER_MAX_ITEMS=5000` explicitly (delivery identical with and +without the correction, since 5,000 aligns with the destination batch size). + +## Results + +Peak RSS (sampled at 50 ms): + +| Scale | Mode | buffer | Wrappers | Peak RSS | Wall | +|------:|------|-------:|---------:|---------:|-----:| +| 100k | baseline | 5,000 | 100,000 | 130 MB | 18.2s | +| 100k | table-wide untuned | 5,000 | 667 | 255 MB | 6.5s | +| 100k | table-wide tuned | 333 | 667 | 240 MB | 6.6s | +| 1M | baseline | 5,000 | 1,000,000 | 132 MB | 111.8s | +| 1M | table-wide untuned | 5,000 | 6,667 | 1,257 MB | 36.2s | +| 1M | table-wide tuned | 333 | 6,667 | 499 MB | 38.9s | +| 4M | table-wide tuned | 333 | 26,667 | 498 MB | 161.5s | + +Tuned RSS is **flat across cardinality** (499 MB @ 1M vs 498 MB @ 4M); the +untuned trend extrapolates to multiple GB at customer scale (12.5M+ rows). The +benchmark reports `peak_rss_per_edge` and enforces a guard band (300 MiB floor ++ 512 B/edge) — fires on untuned runs, silent on tuned ones. + +Destination callbacks / parts (one JSON part per callback): + +| Scale | Mode | Callbacks / Parts | Max rel/callback | Max bytes/callback | +|------:|------|------------------:|-----------------:|-------------------:| +| 100k | baseline | 100 / 100 | 1,000 | 178 KB | +| 1M | baseline | 1,000 / 1,000 | 1,000 | 180 KB | +| 1M | table-wide untuned | 7 / 7 | 150,000 | 27.0 MB | +| 1M | table-wide tuned | 7 / 7 | 150,000 | 27.0 MB | +| 4M | table-wide tuned | 27 / 27 | 150,000 | 27.3 MB | + +(An earlier revision listed 27 parts / ~9 MB at 1M and 107 parts at 4M — +artifacts of the defect's fragmented callbacks.) All runs deliver exact item +totals, `inner_relationships` = row count, and 0 warnings when tuned. + +## Bounds + +- Destination callback/part: `1,000 x 150 = 150,000` relationships (~27 MB), + independent of cardinality and of the extract buffer. +- Process peak RSS: flat 1M -> 4M post-mitigation; guarded by the benchmark band. + +## Conclusion + +Both acceptance bounds hold, and no unsupported upload artifacts are created. +Verified against release tags: only the 1.26 series carries the defect (fixed +in 1.27.0; `buffer_max_items` semantics unchanged through 1.30.0, so the +coordination applies on all versions). If the pin moves past 1.26, +`ensure_dlt_jsonl_batching` becomes a no-op and upstream's corrected iterator +takes over — worth reporting upstream with the minimal reproduction above. +Future wrapper-width or buffering changes will surface via `peak_rss_per_edge` +and the warnings list. diff --git a/benchmarks/_bench_assets.py b/benchmarks/_bench_assets.py new file mode 100644 index 0000000..65f4bee --- /dev/null +++ b/benchmarks/_bench_assets.py @@ -0,0 +1,148 @@ +"""Generic synthetic assets and input generation for the BED-9372 benchmark. + +These assets are deliberately extension-agnostic (not Okta/SAML-specific) so the +benchmark exercises the shared opengraph source, per the ticket requirement to +use generic synthetic assets plus at least one non-Okta high-cardinality shape. +""" + +from __future__ import annotations + +import gzip +import json +import math +from dataclasses import dataclass, field +from pathlib import Path + +from openhound.core.asset import BaseAsset +from openhound.core.models.entries_dataclass import ( + Edge, + EdgePath, + Node as DNode, + NodeProperties as DNodeProperties, +) + + +def _edge(idx: int, k: int = 0) -> Edge: + return Edge( + kind="BENCH_Relationship", + start=EdgePath(match_by="id", value=f"start-{idx}-{k}"), + end=EdgePath(match_by="id", value=f"end-{idx}-{k}"), + ) + + +class OneEdgeAsset(BaseAsset): + """High-cardinality zero-or-one-edge shape (e.g. membership/grant rows).""" + + idx: int + + @property + def as_node(self): + return None + + @property + def edges(self): + return [_edge(self.idx)] + + +class MultiEdgeAsset(BaseAsset): + """A row emitting several edges, to stress inner-relationship growth.""" + + idx: int + n: int + + @property + def as_node(self): + return None + + @property + def edges(self): + return [_edge(self.idx, k) for k in range(self.n)] + + +@dataclass +class _BenchNode(DNode): + id: str = field(default="") + + def __post_init__(self): + self.id = f"node-{self.properties.name}" + + +class NodeAndEdgeAsset(BaseAsset): + """Node-bearing row that also emits one containment/ownership edge.""" + + idx: int + + @property + def as_node(self): + return _BenchNode( + kinds=["BENCH_Node"], + properties=DNodeProperties( + name=f"n{self.idx}", + displayname=f"Node {self.idx}", + environmentid="bench-env", + ), + ) + + @property + def edges(self): + return [ + Edge( + kind="BENCH_Relationship", + start=EdgePath(match_by="id", value=f"node-n{self.idx}"), + end=EdgePath(match_by="id", value=f"end-{self.idx}-0"), + ) + ] + + +def _single_edge_row(idx: int, epr: int) -> dict: + """Row builder for shapes that emit exactly one edge per row. + + These shapes cannot vary the edge count, so reject any edges_per_row other + than 1 rather than silently discarding it. + """ + if epr != 1: + raise ValueError( + f"edges_per_row={epr} is unsupported for single-edge shapes; use 1" + ) + return {"idx": idx} + + +# Maps a shape name to (asset model, row builder). The row builder returns the +# raw dict that read_jsonl will feed back into the model. +ASSET_SHAPES: dict[str, tuple[type[BaseAsset], object]] = { + "one_edge": (OneEdgeAsset, _single_edge_row), + "multi_edge": (MultiEdgeAsset, lambda idx, epr: {"idx": idx, "n": epr}), + "node_and_edge": (NodeAndEdgeAsset, _single_edge_row), +} + + +def model_for_shape(shape: str) -> type[BaseAsset]: + return ASSET_SHAPES[shape][0] + + +def _write_gz(path: Path, rows: list[dict]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with gzip.open(path, "wt", encoding="utf-8") as fh: + for row in rows: + fh.write(json.dumps(row) + "\n") + + +def write_synthetic_input( + input_dir: Path, shape: str, rows: int, edges_per_row: int, files: int +) -> str: + """Write `rows` synthetic rows for `shape` across `files` .jsonl.gz files. + + Returns the table (subdirectory) name used by the opengraph file_glob. + """ + model, row_builder = ASSET_SHAPES[shape] + table = model.__name__.lower() + per_file = math.ceil(rows / files) + written = 0 + for f in range(files): + count = min(per_file, rows - written) + if count <= 0: + break + batch = [row_builder(written + i, edges_per_row) for i in range(count)] + written += count + _write_gz(input_dir / table / f"part-{f:04d}.jsonl.gz", batch) + return table diff --git a/benchmarks/_bench_run.py b/benchmarks/_bench_run.py new file mode 100644 index 0000000..094975f --- /dev/null +++ b/benchmarks/_bench_run.py @@ -0,0 +1,259 @@ +"""Pipeline execution and metric collection for the BED-9372 benchmark. + +Runs a real dlt.pipeline with the real opengraph source and an instrumented +destination (mirroring opengraph_file), recording per-callback item/ +relationship/byte counts plus wall time, process CPU, and peak RSS. +""" + +from __future__ import annotations + +import gzip +import os +import time +from dataclasses import dataclass, field +from pathlib import Path + +# Single load worker reduces the Windows WinError 5 race on DLT's atomic +# os.replace. Must be set before dlt import. +os.environ.setdefault("LOAD__WORKERS", "1") + +import dlt # noqa: E402 +import psutil # noqa: E402 +from dlt.common import json as dlt_json # noqa: E402 + +from openhound.core.lookup import LookupManager # noqa: E402 +from openhound.sources.opengraph.source import ( # noqa: E402 + GraphResource, + opengraph, + writer_buffer_max_items, +) + +from _bench_assets import model_for_shape # noqa: E402 +from _peak_rss import PeakRSSSampler # noqa: E402 +from _win_atomic_retry import install as _install_win_atomic_retry # noqa: E402 + +# Retry DLT's atomic state-file replace so high-part-count runs survive the +# Windows WinError 5 transient. No-op when the replace succeeds first try. +_install_win_atomic_retry() + + +@dataclass +class BenchMetrics: + wall_seconds: float = 0.0 + process_cpu_seconds: float = 0.0 + peak_rss_bytes: int = 0 + edge_wrappers: int = 0 + node_wrappers: int = 0 + inner_relationships: int = 0 + normalized_dlt_items: int = 0 + destination_callbacks: int = 0 + destination_parts: int = 0 + dlt_package_files: int = 0 + part_uncompressed_bytes: int = 0 + part_compressed_bytes: int = 0 + max_relationships_per_callback: int = 0 + max_bytes_per_callback: int = 0 + max_relationships_per_part: int = 0 + max_bytes_per_part: int = 0 + writer_buffer_max_items: int = 0 + peak_rss_per_edge: float = 0.0 + warnings: list[str] = field(default_factory=list) + + +# RSS guard band: fail if peak RSS exceeds floor + per-edge allowance. +RSS_FLOOR_BYTES = 300 * 1024 * 1024 +RSS_PER_EDGE_ALLOWANCE = 512 + + +def _instrumented_destination(real_output_dir: str, metrics: BenchMetrics): + """Return a dlt destination mirroring opengraph_file's write path (1000-item + batch, nodes/edges grouping, one JSON part per callback) while recording the + per-callback metrics the ticket requires. + """ + part_counter: dict[str, int] = {} + + @dlt.destination(skip_dlt_columns_and_tables=True, batch_size=1000) + def instrumented( + items, table, output_path=real_output_dir, source_kind="benchmark" + ): + table_name = table.get("name") or "opengraph" + part_counter[table_name] = part_counter.get(table_name, 0) + 1 + metrics.destination_callbacks += 1 + + nodes = [] + edges = [] + for item in items: + g = item["graph"] + if g["entity_type"] == "node": + nodes.append(g["content"]) + if g["entity_type"] == "edge": + edges.extend(g["content"]) + + payload = dlt_json.dumps( + { + "graph": {"nodes": nodes, "edges": edges}, + "metadata": {"source_kind": source_kind}, + } + ) + raw = payload.encode("utf-8") if isinstance(payload, str) else payload + metrics.max_relationships_per_callback = max( + metrics.max_relationships_per_callback, len(edges) + ) + metrics.max_bytes_per_callback = max(metrics.max_bytes_per_callback, len(raw)) + + file_path = Path(output_path) / f"{table_name}-{part_counter[table_name]}.json" + file_path.write_bytes(raw) + + return instrumented(output_path=real_output_dir, source_kind="benchmark") + + +def _measure_output_parts(output_dir: Path, metrics: BenchMetrics) -> None: + for part in sorted(output_dir.glob("*.json")): + raw = part.read_bytes() + metrics.destination_parts += 1 + metrics.part_uncompressed_bytes += len(raw) + compressed = len(gzip.compress(raw)) + metrics.part_compressed_bytes += compressed + try: + doc = dlt_json.loadb(raw) + edges = len(doc["graph"]["edges"]) + except Exception: + edges = 0 + metrics.max_relationships_per_part = max( + metrics.max_relationships_per_part, edges + ) + metrics.max_bytes_per_part = max(metrics.max_bytes_per_part, len(raw)) + + +def _count_source_wrappers( + input_dir: Path, table: str, shape: str, batch_size: int, lookup: LookupManager +) -> tuple[int, int, int]: + """Iterate the real source once to count edge/node wrappers and inner edges. + + This is a separate pass so the destination timing is not polluted; it uses + the same source configuration the pipeline run uses. + """ + model = model_for_shape(shape) + source = opengraph( + [GraphResource(table=table, model=model)], + bucket_url=input_dir.as_uri(), + lookup=lookup, + extras={}, + batch_size=batch_size, + ) + edge_wrappers = node_wrappers = inner = 0 + for item in source.resources[f"{model.__name__.lower()}_fs"]: + g = item["graph"] + if g["entity_type"] == "edge": + edge_wrappers += 1 + inner += len(g["content"]) + else: + node_wrappers += 1 + return edge_wrappers, node_wrappers, inner + + +def run_pipeline( + input_dir: Path, + output_dir: Path, + table: str, + shape: str, + batch_size: int, + lookup: LookupManager, + work_dir: Path, +) -> BenchMetrics: + output_dir.mkdir(parents=True, exist_ok=True) + work_dir.mkdir(parents=True, exist_ok=True) + metrics = BenchMetrics() + + # Mirror the Converter's writer-buffer coordination: an explicit override + # wins while running; the prior environment is restored afterwards so + # repeated in-process runs never inherit a stale buffer. + buffer_env = "DATA_WRITER__BUFFER_MAX_ITEMS" + prior_buffer = os.environ.get(buffer_env) + if prior_buffer is None: + os.environ[buffer_env] = str(writer_buffer_max_items(batch_size)) + metrics.writer_buffer_max_items = int(os.environ[buffer_env]) + try: + metrics.edge_wrappers, metrics.node_wrappers, metrics.inner_relationships = ( + _count_source_wrappers(input_dir, table, shape, batch_size, lookup) + ) + + model = model_for_shape(shape) + source = opengraph( + [GraphResource(table=table, model=model)], + bucket_url=input_dir.as_uri(), + lookup=lookup, + extras={}, + batch_size=batch_size, + ) + dest = _instrumented_destination(str(output_dir), metrics) + # Isolate DLT's working dir per run so runs never share load packages/state. + pipeline = dlt.pipeline( + pipeline_name="bench_opengraph_convert", + dataset_name="bench", + destination=dest, + pipelines_dir=str(work_dir), + ) + + proc = psutil.Process() + cpu_before = proc.cpu_times() + sampler = PeakRSSSampler(proc) + sampler.start() + t0 = time.perf_counter() + load_info = pipeline.run(source) + metrics.wall_seconds = time.perf_counter() - t0 + sampler.stop() + cpu_after = proc.cpu_times() + + metrics.process_cpu_seconds = (cpu_after.user - cpu_before.user) + ( + cpu_after.system - cpu_before.system + ) + metrics.peak_rss_bytes = sampler.peak_rss + if metrics.inner_relationships > 0: + metrics.peak_rss_per_edge = metrics.peak_rss_bytes / metrics.inner_relationships + rss_band = RSS_FLOOR_BYTES + RSS_PER_EDGE_ALLOWANCE * metrics.inner_relationships + if metrics.peak_rss_bytes > rss_band: + metrics.warnings.append( + f"peak RSS {metrics.peak_rss_bytes} exceeds guard band " + f"{rss_band} (floor {RSS_FLOOR_BYTES} + " + f"{RSS_PER_EDGE_ALLOWANCE} B/edge); staging memory may be " + "scaling with table cardinality" + ) + _measure_output_parts(output_dir, metrics) + _collect_dlt_metrics(pipeline, load_info, metrics) + finally: + if prior_buffer is None: + del os.environ[buffer_env] + else: + os.environ[buffer_env] = prior_buffer + return metrics + + +def _collect_dlt_metrics(pipeline, load_info, metrics: BenchMetrics) -> None: + # Normalized DLT items come from the normalize step's per-job writer metrics + # (items_count), excluding DLT's internal pipeline-state table. + try: + trace = pipeline.last_trace + for step in trace.steps: + if step.step != "normalize": + continue + for _load_id, runs in step.step_info.metrics.items(): + for run in runs: + for fname, writer in run["job_metrics"].items(): + if fname.startswith("_dlt_pipeline_state"): + continue + metrics.normalized_dlt_items += int( + getattr(writer, "items_count", 0) + ) + except Exception as exc: # pragma: no cover - defensive: trace shape drift + metrics.warnings.append(f"normalize metrics unavailable: {exc}") + + # DLT package files: completed load jobs excluding the internal state table. + try: + for package in load_info.load_packages: + for job in package.jobs["completed_jobs"]: + if job.job_file_info.table_name.startswith("_dlt_pipeline_state"): + continue + metrics.dlt_package_files += 1 + except Exception as exc: # pragma: no cover - defensive: LoadInfo shape drift + metrics.warnings.append(f"load package metrics unavailable: {exc}") diff --git a/benchmarks/_peak_rss.py b/benchmarks/_peak_rss.py new file mode 100644 index 0000000..9f53297 --- /dev/null +++ b/benchmarks/_peak_rss.py @@ -0,0 +1,44 @@ +"""Cross-platform peak RSS sampler for the BED-9372 benchmark. + +Windows has no resource.getrusage peak-RSS equivalent, so a background thread +polls psutil RSS during the run and keeps the maximum observed value. +""" + +from __future__ import annotations + +import threading + +import psutil + + +class PeakRSSSampler: + def __init__(self, process: psutil.Process, interval: float = 0.05): + self._process = process + self._interval = interval + self._stop = threading.Event() + self._thread: threading.Thread | None = None + self.peak_rss = 0 + + def _sample_once(self) -> None: + try: + rss = self._process.memory_info().rss + if rss > self.peak_rss: + self.peak_rss = rss + except psutil.Error: + pass + + def _run(self) -> None: + while not self._stop.is_set(): + self._sample_once() + self._stop.wait(self._interval) + + def start(self) -> None: + self._sample_once() + self._thread = threading.Thread(target=self._run, daemon=True) + self._thread.start() + + def stop(self) -> None: + self._stop.set() + if self._thread is not None: + self._thread.join(timeout=2.0) + self._sample_once() diff --git a/benchmarks/_win_atomic_retry.py b/benchmarks/_win_atomic_retry.py new file mode 100644 index 0000000..b0a7c63 --- /dev/null +++ b/benchmarks/_win_atomic_retry.py @@ -0,0 +1,44 @@ +"""Benchmark-local retry shim for DLT's atomic file replace on Windows. + +``os.replace`` in ``FileStorage.save_atomic`` intermittently raises WinError 5 +when a handle on the just-written file lingers (NTFS releases rename locks late; +see CPython gh-90161 and dlt-hub/dlt PR #3853). At the batch_size=1 baseline the +state file is committed ~1000x, hitting the race often enough to exhaust DLT's +retries. Applied here (not in shipped code) since it only affects this synthetic +worst case; a genuine permission error still surfaces once the budget is spent. +""" + +from __future__ import annotations + +import time + +_MAX_ATTEMPTS = 20 +_BACKOFF_SECONDS = 0.02 + + +def install() -> None: + """Idempotently wrap FileStorage.save_atomic with a WinError 5 retry loop.""" + from dlt.common.storages import file_storage as fs + + if getattr(fs.FileStorage.save_atomic, "_win_atomic_retry", False): + return + + original = fs.FileStorage.save_atomic + + def save_atomic_with_retry( + storage_path: str, relative_path: str, data, file_type: str = "t" + ) -> str: + last_exc: Exception | None = None + for attempt in range(_MAX_ATTEMPTS): + try: + return original(storage_path, relative_path, data, file_type=file_type) + except PermissionError as exc: # WinError 5 on the os.replace + last_exc = exc + if attempt == _MAX_ATTEMPTS - 1: + break + time.sleep(_BACKOFF_SECONDS * (attempt + 1)) + assert last_exc is not None + raise last_exc + + save_atomic_with_retry._win_atomic_retry = True # type: ignore[attr-defined] + fs.FileStorage.save_atomic = staticmethod(save_atomic_with_retry) diff --git a/benchmarks/opengraph_batching_benchmark.py b/benchmarks/opengraph_batching_benchmark.py new file mode 100644 index 0000000..bb6b1e9 --- /dev/null +++ b/benchmarks/opengraph_batching_benchmark.py @@ -0,0 +1,194 @@ +"""End-to-end benchmark for BED-9372 table-wide edge batching. + +Runs the real opengraph source through a real dlt.pipeline and instrumented +destination at the ticket's 100k/1M row scales, recording the required metrics +(wall time, CPU, peak RSS, wrappers, inner relationships, normalized DLT items, +destination callbacks/parts, package files, bytes, and per-callback/part maxima). +LookupManager and the synthetic asset are held constant so gains reflect batching. + +Usage: + python benchmarks/opengraph_batching_benchmark.py --rows 100000 --batch-size 150 + python benchmarks/opengraph_batching_benchmark.py --rows 1000000 --edges-per-row 1 + +Pass --baseline to force per-row wrapping (batch_size=1) for a before/after run. +""" + +from __future__ import annotations + +import argparse +import json +import math +import sys +from dataclasses import asdict, dataclass +from pathlib import Path +from uuid import uuid4 + +# Support running as a plain script (python benchmarks/opengraph_batching_benchmark.py) +# as well as a module (python -m benchmarks.opengraph_batching_benchmark). +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +import duckdb # noqa: E402 + +from openhound.core.lookup import LookupManager # noqa: E402 + +from _bench_assets import ASSET_SHAPES, write_synthetic_input # noqa: E402 +from _bench_run import BenchMetrics, run_pipeline # noqa: E402 + + +@dataclass +class BenchConfig: + rows: int + batch_size: int + edges_per_row: int + shape: str + files: int + baseline: bool + keep_output: bool + output_root: Path + owns_output_root: bool + quiet: bool = False + + +def _parse_args(argv: list[str] | None = None) -> BenchConfig: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--rows", type=int, default=100_000, help="Number of source rows.") + p.add_argument( + "--batch-size", type=int, default=150, help="Source edge batch size." + ) + p.add_argument( + "--edges-per-row", + type=int, + default=1, + help="Edges emitted per source row (1 = high-cardinality one-edge shape).", + ) + p.add_argument( + "--shape", + choices=sorted(ASSET_SHAPES), + default="one_edge", + help="Synthetic asset shape (one_edge, multi_edge, node_and_edge).", + ) + p.add_argument( + "--files", + type=int, + default=4, + help="Number of input .jsonl.gz files to spread rows across.", + ) + p.add_argument( + "--baseline", + action="store_true", + help="Force per-row wrapping (batch_size=1) for a before/after comparison.", + ) + p.add_argument( + "--keep-output", + action="store_true", + help="Keep the generated input/output instead of using a temp dir.", + ) + p.add_argument( + "--output-root", + type=Path, + default=None, + help="Directory for input/output (default: a temp dir under the system tmp).", + ) + p.add_argument( + "--quiet", + action="store_true", + help="Silence DLT INFO logging so the JSON report is the only output.", + ) + ns = p.parse_args(argv) + if ns.rows < 1: + p.error("--rows must be >= 1") + if ns.batch_size < 1: + p.error("--batch-size must be >= 1") + if ns.edges_per_row < 0: + p.error("--edges-per-row must be >= 0") + if ns.files < 1: + p.error("--files must be >= 1") + + owns_output_root = ns.output_root is None + root = ns.output_root or Path( + __import__("tempfile").mkdtemp(prefix="openhound-bench-") + ) + return BenchConfig( + rows=ns.rows, + batch_size=1 if ns.baseline else ns.batch_size, + edges_per_row=ns.edges_per_row, + shape=ns.shape, + files=ns.files, + baseline=ns.baseline, + keep_output=ns.keep_output, + output_root=root, + owns_output_root=owns_output_root, + quiet=ns.quiet, + ) + + +def _in_memory_lookup() -> LookupManager: + # Hold DuckDB/model work constant: a real LookupManager over an empty + # in-memory schema so lookup cost is fixed and negligible across runs. + conn = duckdb.connect(":memory:") + return LookupManager(conn, "main") + + +def _print_report(cfg: BenchConfig, metrics: BenchMetrics, run_dir: Path) -> None: + expected_wrappers = ( + cfg.rows * cfg.edges_per_row + if cfg.batch_size == 1 + else math.ceil((cfg.rows * cfg.edges_per_row) / cfg.batch_size) + ) + report = { + "config": asdict(cfg) | {"output_root": str(cfg.output_root)}, + "run_dir": str(run_dir), + "expected_edge_wrappers_ceiling": expected_wrappers, + "metrics": asdict(metrics), + } + print(json.dumps(report, indent=2, default=str)) + + +def _silence_dlt_logging() -> None: + import logging + import os + + # DLT's log level is env-driven and resolved when the pipeline runs, so the + # env var is what takes effect; setLevel covers an already-created logger. + # CRITICAL suppresses DLT's ERROR log for the recovered WinError 5 transient; + # a real terminal failure still raises regardless of log level. + os.environ["RUNTIME__LOG_LEVEL"] = "CRITICAL" + logging.getLogger("dlt").setLevel(logging.CRITICAL) + + +def main(argv: list[str] | None = None) -> None: + cfg = _parse_args(argv) + if cfg.quiet: + _silence_dlt_logging() + cfg.output_root.mkdir(parents=True, exist_ok=True) + # Fresh run dir per invocation: preserves the caller-supplied root while + # isolating table/output/dlt_work children from any stale prior run. + run_dir = cfg.output_root / f"run-{uuid4().hex[:8]}" + input_dir = run_dir / "input" + output_dir = run_dir / "output" + work_dir = run_dir / "dlt_work" + + try: + table = write_synthetic_input( + input_dir, cfg.shape, cfg.rows, cfg.edges_per_row, cfg.files + ) + + metrics = run_pipeline( + input_dir=input_dir, + output_dir=output_dir, + table=table, + shape=cfg.shape, + batch_size=cfg.batch_size, + lookup=_in_memory_lookup(), + work_dir=work_dir, + ) + _print_report(cfg, metrics, run_dir) + finally: + if not cfg.keep_output and cfg.owns_output_root: + import shutil + + shutil.rmtree(cfg.output_root, ignore_errors=True) + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index e87b588..1f0fe20 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,6 +51,11 @@ dotenv = [".env"] [tool.ruff] exclude = ["src/openhound/cookiecutter-templates/*"] +[tool.ruff.lint.per-file-ignores] +# Harness sets env vars / sys.path before its imports and uses broad guards +# around DLT's trace/LoadInfo shapes. +"benchmarks/*" = ["E402", "I001", "BLE001", "PERF102", "PLR1730", "RUF100"] + [tool.mypy] exclude = ["src/openhound/cookiecutter-templates/*"] disable_error_code = ["import-untyped"] diff --git a/src/openhound/__init__.py b/src/openhound/__init__.py index 26bb1fd..5071c8f 100644 --- a/src/openhound/__init__.py +++ b/src/openhound/__init__.py @@ -1,5 +1,9 @@ from importlib.metadata import PackageNotFoundError, version +from openhound.core.dlt_jsonl_batching import ensure_dlt_jsonl_batching + +ensure_dlt_jsonl_batching() + try: __version__ = version("openhound") except PackageNotFoundError: diff --git a/src/openhound/core/convert.py b/src/openhound/core/convert.py index 5af52e1..51c183a 100644 --- a/src/openhound/core/convert.py +++ b/src/openhound/core/convert.py @@ -1,4 +1,5 @@ import logging +import os from dataclasses import dataclass from enum import Enum from pathlib import Path @@ -16,7 +17,12 @@ from openhound.core.progress import Progress from openhound.destinations.bloodhound_enterprise.destination import ingest from openhound.destinations.opengraph.destination import opengraph_file -from openhound.sources.opengraph.source import GraphResource, opengraph +from openhound.sources.opengraph.source import ( + DEFAULT_EDGE_BATCH_SIZE, + GraphResource, + opengraph, + writer_buffer_max_items, +) logger = logging.getLogger(__name__) @@ -47,6 +53,7 @@ def __init__( @property def pipeline(self) -> Pipeline: + self._coordinate_writer_buffer() if self.method == Method.ingest: logger.debug( "Initializing BloodHound Enterprise client for converter ingest method" @@ -67,6 +74,13 @@ def pipeline(self) -> Pipeline: ) return pipeline + @staticmethod + def _coordinate_writer_buffer() -> None: + # setdefault keeps any explicit DATA_WRITER__BUFFER_MAX_ITEMS override. + os.environ.setdefault( + "DATA_WRITER__BUFFER_MAX_ITEMS", str(writer_buffer_max_items()) + ) + def run( self, source_object: DltSource, @@ -99,6 +113,7 @@ def run( lookup=self.lookup, bucket_url=str(self.input_path), extras=extra_context, + batch_size=DEFAULT_EDGE_BATCH_SIZE, ) ) diff --git a/src/openhound/core/dlt_jsonl_batching.py b/src/openhound/core/dlt_jsonl_batching.py new file mode 100644 index 0000000..2ad3321 --- /dev/null +++ b/src/openhound/core/dlt_jsonl_batching.py @@ -0,0 +1,67 @@ +"""Correct jsonl batching for dlt callable-destination load jobs. + +dlt 1.26.x re-yields accumulated items per load-file line (duplicated data) +unless buffer_max_items divides evenly by batch size. This module installs +dlt's later corrected semantics: full batches yielded with reset, one +trailing partial per file. +""" + +import logging +from collections.abc import Iterable +from typing import Any + +import dlt +from dlt.common import json +from dlt.common.storages import FileStorage +from dlt.common.typing import TDataItems + +logger = logging.getLogger(__name__) + +_AFFECTED_DLT_SERIES = (1, 26) +_INSTALLED_ATTR = "_openhound_jsonl_batching_installed" + + +def _jsonl_get_batches(self: Any, start_index: int) -> Iterable[TDataItems]: + current_batch: TDataItems = [] + + with FileStorage.open_zipsafe_ro(self._file_path) as f: + for line in f: + encoded_json = json.typed_loads(line) + if isinstance(encoded_json, dict): + encoded_json = [encoded_json] + + for item in encoded_json: + if start_index > 0: + start_index -= 1 + continue + for column in self._skipped_columns: + item.pop(column, None) + current_batch.append(item) + if len(current_batch) == self._config.batch_size: + yield current_batch + current_batch = [] + + if current_batch: + yield current_batch + + +def ensure_dlt_jsonl_batching() -> bool: + """Install correct get_batches on affected dlt versions.""" + if getattr(dlt, _INSTALLED_ATTR, False): + return True + + from dlt.destinations.job_impl import DestinationJsonlLoadJob + + version = tuple(int(part) for part in dlt.__version__.split(".")[:2]) + if not hasattr(DestinationJsonlLoadJob, "get_batches") or version != _AFFECTED_DLT_SERIES: + logger.debug("dlt %s needs no jsonl batching correction", dlt.__version__) + return False + + DestinationJsonlLoadJob.get_batches = _jsonl_get_batches # type: ignore[method-assign] + setattr(dlt, _INSTALLED_ATTR, True) + logger.info( + "Installed OpenHound jsonl batching for dlt %s DestinationJsonlLoadJob " + "(upstream yielded cumulative partial batches across load-file lines)", + dlt.__version__, + ) + return True diff --git a/src/openhound/sources/opengraph/source.py b/src/openhound/sources/opengraph/source.py index eb0c56f..0183d6b 100644 --- a/src/openhound/sources/opengraph/source.py +++ b/src/openhound/sources/opengraph/source.py @@ -17,14 +17,33 @@ class GraphResource: model: BaseAsset +DEFAULT_EDGE_BATCH_SIZE = 150 + +# Each wrapper holds up to batch_size edges; scale DLT's item-count writer +# buffer so buffered edges stay bounded regardless of batch size. +DLT_BUFFERED_EDGE_BUDGET = 50_000 +DLT_DEFAULT_BUFFER_MAX_ITEMS = 5_000 + + +def writer_buffer_max_items(batch_size: int = DEFAULT_EDGE_BATCH_SIZE) -> int: + if batch_size < 1: + raise ValueError(f"batch_size must be >= 1, got {batch_size}") + return min( + DLT_DEFAULT_BUFFER_MAX_ITEMS, + max(1, DLT_BUFFERED_EDGE_BUDGET // batch_size), + ) + + @dlt.source(name="opengraph", max_table_nesting=0) def opengraph( graph_resources: list[GraphResource], bucket_url: str, lookup: LookupManager, extras: dict | None = None, - batch_size: int = 150, + batch_size: int = DEFAULT_EDGE_BATCH_SIZE, ): + if batch_size < 1: + raise ValueError(f"batch_size must be >= 1, got {batch_size}") def apply_context(obj): obj._lookup = lookup @@ -40,9 +59,19 @@ def apply_context(obj): | read_jsonl() ) - @dlt.transformer(parallelized=False, name=table_name, columns=GraphContent) - def generate_graph(resources, model, apply_context: Callable | None = None): - for resource in resources: + @dlt.resource(name=table_name, columns=GraphContent) + def generate_graph( + reader=reader, + model=graph_resource.model, + apply_context: Callable | None = apply_context, + ): + # One generator consumes the whole reader stream, so the edge + # accumulator spans the entire table (across chunks and files) and + # the final partial batch flushes once, giving exactly + # ceil(total_edges / batch_size) wrappers in encounter order. Tradeoff: + # a failure re-extracts the whole table rather than resuming mid-chunk. + edge_parts = [] + for resource in reader: parsed_resource = model(**resource) if apply_context: apply_context(parsed_resource) @@ -56,16 +85,13 @@ def generate_graph(resources, model, apply_context: Callable | None = None): }, } - edge_parts = [] for edge in parsed_resource.edges: edge_parts.append(asdict(edge)) if len(edge_parts) >= batch_size: yield {"graph": {"content": edge_parts, "entity_type": "edge"}} edge_parts = [] - if edge_parts: - yield {"graph": {"content": edge_parts, "entity_type": "edge"}} + if edge_parts: + yield {"graph": {"content": edge_parts, "entity_type": "edge"}} - yield reader | generate_graph( - model=graph_resource.model, apply_context=apply_context - ) + yield generate_graph() diff --git a/tests/test_dlt_jsonl_batching.py b/tests/test_dlt_jsonl_batching.py new file mode 100644 index 0000000..ed9a3f5 --- /dev/null +++ b/tests/test_dlt_jsonl_batching.py @@ -0,0 +1,73 @@ +"""Tests for OpenHound's dlt jsonl batching module.""" + +import json as jsonlib +from types import SimpleNamespace + +from dlt.destinations.job_impl import DestinationJsonlLoadJob + +import openhound # noqa: F401 ensures batching is installed process-wide +from openhound.core.dlt_jsonl_batching import ensure_dlt_jsonl_batching + + +def _make_job(tmp_path, lines, batch_size, skipped_columns=()): + path = tmp_path / "load.jsonl" + path.write_text("\n".join(jsonlib.dumps(line) for line in lines) + "\n", encoding="utf-8") + job = DestinationJsonlLoadJob.__new__(DestinationJsonlLoadJob) + job._file_path = str(path) + job._config = SimpleNamespace(batch_size=batch_size) + job._skipped_columns = list(skipped_columns) + return job + + +def _collect(job, start_index=0): + return [list(batch) for batch in job.get_batches(start_index)] + + +def test_batching_is_installed_for_pinned_dlt(): + assert ensure_dlt_jsonl_batching() is True + + +def test_multi_line_files_deliver_each_item_exactly_once(tmp_path): + lines = [list(range(offset, offset + 10)) for offset in range(0, 40, 10)] + job = _make_job(tmp_path, lines, batch_size=7) + + batches = _collect(job) + + delivered = [item for batch in batches for item in batch] + expected = list(range(40)) + assert delivered == expected + assert [len(batch) for batch in batches] == [7, 7, 7, 7, 7, 5] + + +def test_start_index_resumes_mid_file(tmp_path): + lines = [[1, 2, 3], [4, 5, 6], [7, 8, 9]] + job = _make_job(tmp_path, lines, batch_size=2) + + batches = _collect(job, start_index=4) + + delivered = [item for batch in batches for item in batch] + assert delivered == [5, 6, 7, 8, 9] + assert [len(batch) for batch in batches] == [2, 2, 1] + + +def test_skipped_columns_are_removed(tmp_path): + records = [ + [{"_dlt_id": "a", "value": 1}, {"_dlt_id": "b", "value": 2}], + [{"_dlt_id": "c", "value": 3}], + ] + job = _make_job(tmp_path, records, batch_size=10, skipped_columns=("_dlt_id",)) + + batches = _collect(job) + + delivered = [item for batch in batches for item in batch] + assert delivered == [{"value": 1}, {"value": 2}, {"value": 3}] + + +def test_single_dict_lines_are_wrapped(tmp_path): + job = _make_job(tmp_path, [{"n": 1}, {"n": 2}], batch_size=1) + + batches = _collect(job) + + delivered = [item for batch in batches for item in batch] + assert delivered == [{"n": 1}, {"n": 2}] + assert len(batches) == 2 diff --git a/tests/test_opengraph_batching.py b/tests/test_opengraph_batching.py new file mode 100644 index 0000000..a544572 --- /dev/null +++ b/tests/test_opengraph_batching.py @@ -0,0 +1,351 @@ +"""Regression tests for table-wide edge batching in the opengraph source. + +The edge accumulator must span the whole table (across DLT chunks and input +files), flush one final partial batch, and preserve the flattened edge sequence +exactly (order + duplicates, no dedup). +""" + +import gzip +import json +import math +from dataclasses import dataclass, field + +import pytest +from dlt.extract.exceptions import ResourceExtractionError + +from openhound.core.asset import BaseAsset +from openhound.core.models.entries_dataclass import ( + Edge, + EdgePath, +) +from openhound.core.models.entries_dataclass import ( + Node as DNode, +) +from openhound.core.models.entries_dataclass import ( + NodeProperties as DNodeProperties, +) +from openhound.sources.opengraph.source import GraphResource, opengraph + +# read_jsonl yields chunks of this many rows; table-wide batching spans them. +CHUNK = 1000 + + +def _edge(idx, k=0): + return Edge( + kind="TEST_Edge", + start=EdgePath(match_by="id", value=f"s{idx}-{k}"), + end=EdgePath(match_by="id", value=f"e{idx}-{k}"), + ) + + +class OneEdgeAsset(BaseAsset): + idx: int + + @property + def as_node(self): + return None + + @property + def edges(self): + return [_edge(self.idx)] + + +class MultiEdgeAsset(BaseAsset): + idx: int + n: int + + @property + def as_node(self): + return None + + @property + def edges(self): + return [_edge(self.idx, k) for k in range(self.n)] + + +@dataclass +class _TestNode(DNode): + id: str = field(default="") + + def __post_init__(self): + self.id = f"node-{self.properties.name}" + + +class NodeAndEdgeAsset(BaseAsset): + idx: int + + @property + def as_node(self): + return _TestNode( + kinds=["TestKind"], + properties=DNodeProperties( + name=f"n{self.idx}", displayname="d", environmentid="env" + ), + ) + + @property + def edges(self): + return [_edge(self.idx)] + + +class EmptyAsset(BaseAsset): + idx: int + + @property + def as_node(self): + return None + + @property + def edges(self): + return [] + + +class FailAtAsset(BaseAsset): + # Emits one edge per row until idx == extras["fail_at"], where it raises. + # Used to assert accumulator state is not leaked into a re-extraction. + idx: int + + @property + def as_node(self): + return None + + @property + def edges(self): + if self.idx == self._extras.get("fail_at"): + raise RuntimeError("boom") + return [_edge(self.idx)] + + +def _write_gz(path, rows): + path.parent.mkdir(parents=True, exist_ok=True) + with gzip.open(path, "wt", encoding="utf-8") as fh: + for row in rows: + fh.write(json.dumps(row) + "\n") + + +def _collect(bucket, table, model, batch_size=150): + source = opengraph( + [GraphResource(table=table, model=model)], + bucket_url=bucket.as_uri(), + lookup=None, + extras={}, + batch_size=batch_size, + ) + return list(source.resources[f"{model.__name__.lower()}_fs"]) + + +def _source(bucket, resources, batch_size=150): + return opengraph( + [GraphResource(table=t, model=m) for t, m in resources], + bucket_url=bucket.as_uri(), + lookup=None, + extras={}, + batch_size=batch_size, + ) + + +def _split(items): + nodes = [i for i in items if i["graph"]["entity_type"] == "node"] + edges = [i for i in items if i["graph"]["entity_type"] == "edge"] + flat = [e for i in edges for e in i["graph"]["content"]] + return nodes, edges, flat + + +def _one_edge_rows(bucket, table, count, files=1): + per = math.ceil(count / files) + written = 0 + for f in range(files): + rows = [{"idx": i} for i in range(written, min(written + per, count))] + written += len(rows) + _write_gz(bucket / table / f"part{f}.jsonl.gz", rows) + + +def test_batching_spans_rows_within_chunk(tmp_path): + # 150 one-edge rows fit in a single chunk -> one wrapper, not 150. + _one_edge_rows(tmp_path, "oneedgeasset", 150) + _, edges, flat = _split(_collect(tmp_path, "oneedgeasset", OneEdgeAsset)) + assert len(edges) == 1 + assert len(edges[0]["graph"]["content"]) == 150 + assert len(flat) == 150 + + +def test_batching_crosses_1000_row_chunk_boundary(tmp_path): + # 1001 rows cross the read_jsonl chunk boundary (1000 + 1). Table-wide, the + # accumulator spans chunks -> exactly ceil(1001/150) == 7 wrappers, and the + # final remainder is the only partial wrapper. + _one_edge_rows(tmp_path, "oneedgeasset", 1001) + _, edges, flat = _split(_collect(tmp_path, "oneedgeasset", OneEdgeAsset)) + expected = math.ceil(1001 / 150) + assert len(edges) == expected == 7 + assert [len(e["graph"]["content"]) for e in edges] == [150] * 6 + [101] + assert len(flat) == 1001 + + +def test_batching_across_multiple_input_files(tmp_path): + # The whole table is one stream, so batches span input files too. + _one_edge_rows(tmp_path, "oneedgeasset", 400, files=2) + _, edges, flat = _split(_collect(tmp_path, "oneedgeasset", OneEdgeAsset)) + assert len(edges) == math.ceil(400 / 150) == 3 + assert [len(e["graph"]["content"]) for e in edges] == [150, 150, 100] + assert len(flat) == 400 + + +def test_batching_across_files_and_chunks(tmp_path): + # Multiple files each spanning several chunks still yield one table-wide + # count with a single trailing partial wrapper. + _one_edge_rows(tmp_path, "oneedgeasset", 1500, files=3) + _, edges, flat = _split(_collect(tmp_path, "oneedgeasset", OneEdgeAsset)) + assert len(edges) == math.ceil(1500 / 150) == 10 + assert all(len(e["graph"]["content"]) == 150 for e in edges) + assert len(flat) == 1500 + + +def test_flattened_sequence_preserves_order_and_duplicates(tmp_path): + _one_edge_rows(tmp_path, "oneedgeasset", 20) + _, _, flat = _split(_collect(tmp_path, "oneedgeasset", OneEdgeAsset, batch_size=7)) + expected = [_edge(i) for i in range(20)] + got = [(e["kind"], e["start"]["value"], e["end"]["value"]) for e in flat] + want = [(e.kind, e.start.value, e.end.value) for e in expected] + assert got == want # order preserved + + +def test_duplicate_edges_are_not_deduplicated(tmp_path): + # Two rows emitting identical edges must both survive. + _write_gz(tmp_path / "oneedgeasset" / "p.jsonl.gz", [{"idx": 5}, {"idx": 5}]) + _, edges, flat = _split(_collect(tmp_path, "oneedgeasset", OneEdgeAsset)) + assert len(flat) == 2 + assert flat[0] == flat[1] + assert len(edges) == 1 # both batched into one wrapper + + +def test_batch_size_one_preserved(tmp_path): + _one_edge_rows(tmp_path, "oneedgeasset", 5) + _, edges, flat = _split( + _collect(tmp_path, "oneedgeasset", OneEdgeAsset, batch_size=1) + ) + assert len(edges) == 5 + assert all(len(e["graph"]["content"]) == 1 for e in edges) + assert len(flat) == 5 + + +@pytest.mark.parametrize("bad", [0, -1, -150]) +def test_batch_size_below_one_rejected(tmp_path, bad): + _one_edge_rows(tmp_path, "oneedgeasset", 1) + with pytest.raises(ValueError): + _collect(tmp_path, "oneedgeasset", OneEdgeAsset, batch_size=bad) + + +def test_exact_batch_size_single_row(tmp_path): + # One row emitting exactly batch_size edges -> exactly one full wrapper. + _write_gz(tmp_path / "multiedgeasset" / "p.jsonl.gz", [{"idx": 0, "n": 150}]) + _, edges, flat = _split(_collect(tmp_path, "multiedgeasset", MultiEdgeAsset)) + assert len(edges) == 1 + assert len(edges[0]["graph"]["content"]) == 150 + assert len(flat) == 150 + + +def test_batch_size_plus_one_single_row(tmp_path): + # batch_size + 1 edges -> a full wrapper plus a one-edge remainder. + _write_gz(tmp_path / "multiedgeasset" / "p.jsonl.gz", [{"idx": 0, "n": 151}]) + _, edges, flat = _split(_collect(tmp_path, "multiedgeasset", MultiEdgeAsset)) + assert [len(e["graph"]["content"]) for e in edges] == [150, 1] + assert len(flat) == 151 + + +def test_nodes_unchanged_and_never_mixed_with_edges(tmp_path): + _one_edge_rows(tmp_path, "nodeandedgeasset", 300) + nodes, edges, flat = _split( + _collect(tmp_path, "nodeandedgeasset", NodeAndEdgeAsset) + ) + # One node wrapper per row, node content unchanged (single dict, not a list). + assert len(nodes) == 300 + assert all(isinstance(n["graph"]["content"], dict) for n in nodes) + # Edges are still batched across rows and never mixed into a node wrapper. + assert len(edges) == math.ceil(300 / 150) == 2 + assert all(isinstance(e["graph"]["content"], list) for e in edges) + assert len(flat) == 300 + + +def test_empty_input_file_yields_nothing(tmp_path): + _write_gz(tmp_path / "oneedgeasset" / "p.jsonl.gz", []) + assert _collect(tmp_path, "oneedgeasset", OneEdgeAsset) == [] + + +def test_zero_output_rows_yield_nothing(tmp_path): + _one_edge_rows(tmp_path, "emptyasset", 10) + assert _collect(tmp_path, "emptyasset", EmptyAsset) == [] + + +def test_minimal_wrapper_count_matches_ceiling(tmp_path): + # Table-wide batching is strictly minimal: ceil(N / batch_size) wrappers for + # a range of counts and batch sizes, crossing chunk and file boundaries. + for count, batch, files in [(999, 150, 1), (1001, 150, 1), (2345, 200, 4)]: + _one_edge_rows(tmp_path, "oneedgeasset", count, files=files) + _, edges, flat = _split( + _collect(tmp_path, "oneedgeasset", OneEdgeAsset, batch_size=batch) + ) + assert len(edges) == math.ceil(count / batch) + assert len(flat) == count + for f in range(files): + (tmp_path / "oneedgeasset" / f"part{f}.jsonl.gz").unlink() + + +def test_flattened_sequence_matches_per_row_baseline(tmp_path): + # Parity with the pre-fix per-row wrapping: identical flattened edge sequence + # (order + duplicates), just regrouped into fewer wrappers. Crosses chunks. + _one_edge_rows(tmp_path, "oneedgeasset", 1500) + _, batched, flat_after = _split(_collect(tmp_path, "oneedgeasset", OneEdgeAsset)) + _, per_row, flat_before = _split( + _collect(tmp_path, "oneedgeasset", OneEdgeAsset, batch_size=1) + ) + + def key(e): + return (e["kind"], e["start"]["value"], e["end"]["value"]) + + assert [key(e) for e in flat_after] == [key(e) for e in flat_before] + assert len(per_row) == 1500 + assert len(batched) == math.ceil(1500 / 150) == 10 + + +def test_multiple_graph_resources_isolated(tmp_path): + # Two resources with different row counts each batch table-wide with their + # own accumulator; neither leaks edges into the other. + _one_edge_rows(tmp_path, "oneedgeasset", 200) + _write_gz(tmp_path / "multiedgeasset" / "p.jsonl.gz", [{"idx": 0, "n": 151}]) + source = _source( + tmp_path, + [("oneedgeasset", OneEdgeAsset), ("multiedgeasset", MultiEdgeAsset)], + ) + + _, one_edges, one_flat = _split(list(source.resources["oneedgeasset_fs"])) + _, multi_edges, multi_flat = _split(list(source.resources["multiedgeasset_fs"])) + + assert len(one_edges) == math.ceil(200 / 150) == 2 + assert [len(e["graph"]["content"]) for e in one_edges] == [150, 50] + assert len(one_flat) == 200 + assert [len(e["graph"]["content"]) for e in multi_edges] == [150, 1] + assert len(multi_flat) == 151 + + +def test_failure_does_not_leak_accumulator_into_retry(tmp_path): + # A failure mid-stream aborts extraction; re-iterating builds a fresh + # generator with a fresh accumulator, so the successful retry has no + # duplicated or leaked pending edges. + _one_edge_rows(tmp_path, "failatasset", 300) + + def build(fail_at): + source = opengraph( + [GraphResource(table="failatasset", model=FailAtAsset)], + bucket_url=tmp_path.as_uri(), + lookup=None, + extras={"fail_at": fail_at}, + batch_size=150, + ) + return source.resources["failatasset_fs"] + + with pytest.raises(ResourceExtractionError): + list(build(fail_at=170)) + + _, edges, flat = _split(list(build(fail_at=-1))) + assert len(edges) == math.ceil(300 / 150) == 2 + assert len(flat) == 300 diff --git a/tests/test_writer_buffer.py b/tests/test_writer_buffer.py new file mode 100644 index 0000000..34e5d61 --- /dev/null +++ b/tests/test_writer_buffer.py @@ -0,0 +1,74 @@ +"""Tests for BED-9372 writer-buffer coordination.""" + +import inspect +import math + +import pytest + +from openhound.core.convert import Converter, Method +from openhound.sources.opengraph.source import ( + DEFAULT_EDGE_BATCH_SIZE, + DLT_BUFFERED_EDGE_BUDGET, + DLT_DEFAULT_BUFFER_MAX_ITEMS, + opengraph, + writer_buffer_max_items, +) + +ENV_VAR = "DATA_WRITER__BUFFER_MAX_ITEMS" + + +def test_source_signature_default_matches_constant(): + sig = inspect.signature(opengraph) + assert sig.parameters["batch_size"].default == DEFAULT_EDGE_BATCH_SIZE + + +def test_buffer_formula_bounds_buffered_edges(): + items = writer_buffer_max_items(DEFAULT_EDGE_BATCH_SIZE) + assert items == math.floor(DLT_BUFFERED_EDGE_BUDGET / DEFAULT_EDGE_BATCH_SIZE) + assert items * DEFAULT_EDGE_BATCH_SIZE <= DLT_BUFFERED_EDGE_BUDGET + + assert writer_buffer_max_items(1) == DLT_DEFAULT_BUFFER_MAX_ITEMS + assert writer_buffer_max_items(10**9) == 1 + for batch_size in (1, 2, 7, 149, 150, 151, 1_000, 100_000): + assert ( + writer_buffer_max_items(batch_size) * batch_size + <= DLT_BUFFERED_EDGE_BUDGET + batch_size + ) + + +def test_buffer_formula_rejects_invalid_batch_size(): + for bad in (0, -1): + with pytest.raises(ValueError, match="batch_size must be >= 1"): + writer_buffer_max_items(bad) + + +def test_converter_applies_env_default(monkeypatch): + monkeypatch.delenv(ENV_VAR, raising=False) + converter = Converter( + name="t", + input_path=None, + lookup=None, + output_path=None, + source_kind="test", + method=Method.write, + ) + converter._coordinate_writer_buffer() + import os + + assert os.environ[ENV_VAR] == str(writer_buffer_max_items()) + + +def test_converter_respects_explicit_override(monkeypatch): + monkeypatch.setenv(ENV_VAR, "777") + converter = Converter( + name="t", + input_path=None, + lookup=None, + output_path=None, + source_kind="test", + method=Method.write, + ) + converter._coordinate_writer_buffer() + import os + + assert os.environ[ENV_VAR] == "777"