-
Notifications
You must be signed in to change notification settings - Fork 4
fix: Table-wide edge batching in the OpenGraph convert source - BED-9372 #70
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
ktstrader
wants to merge
6
commits into
main
Choose a base branch
from
fix/BED-9372-table-wide-edge-batching
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Open
Changes from all commits
Commits
Show all changes
6 commits
Select commit
Hold shift + click to select a range
6cae801
fix: implemented table wide edge batching and created benchmark outsi…
ktstrader 6022cb9
fix: issues found by coderabbit
ktstrader fbc82bb
fix: a couple more nits from coderabbit
ktstrader d4d96c5
fix: bound writer buffer memory and correct dlt jsonl batching
ktstrader dd88b09
fix: Isolate each benchmark run from stale workspace data
ktstrader 5da71fa
fix: a couple of minor issues identified by coderabbit
ktstrader File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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), | ||
| } | ||
|
coderabbitai[bot] marked this conversation as resolved.
|
||
|
|
||
|
|
||
| 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 | ||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.