[Vedaant-J]

Stacked on #101 (vedaant/dg-per-seed-recording).

Summary

This PR runs the per-seed deduplication on the sampler's thread pool and resolves seed anchor times once per request. Together with the previous PR, multi-threaded sampling is 2–3.3× faster than main at 3–6 hops, with identical output.

Problem. After the previous PR, a process with 8 sampler threads still spent most of a 5–6-hop request in serial work, and one per-edge lookup remained in the parallel phase:

• With 8 threads at 5 hops, the commit phase took 6.6 s of a 7.3 s run (100M graph, 5,000 seeds), with about 2.3 cores busy.
• Every sampled (node, relation) pair looked up its seed's anchor time by hashing the seed type's name string, about 8% of single-thread CPU.

Cause. The per-seed step only touches one seed's map and edges, but it ran on one thread. The anchor lookup went seed_type_id → type name → HashMap<String, Vec<i64>> on every call to sample_segment_edges.

Change.

• commit_edges splits the seeds into contiguous ranges of roughly equal edge counts (4 per thread) and runs the deduplication and the final map update on the engine's pool. Local id assignment and edge output remain one in-order pass, so output is unchanged. The parallel path runs only when the pool has more than one thread and the layer has at least DiskGraphEngine::parallel_commit_min_edges edges (default 32,768); single-thread processes run the same code inline. The default comes from forcing the threshold to 1 and to infinity at 8 threads: layers of a few thousand to ~25k edges ran 10–50% slower on the pool, layers of 100k+ edges ran faster. Cuts fall on seed boundaries, so ranges are uneven; 4 ranges per thread let the pool balance them.
• The in-order pass keeps each destination type's state at hand instead of looking it up in type_states for every new node.
• sample_with_cutoff resolves each seed type's anchor slice once; sample_segment_edges takes that slice table instead of the request.

Back-compat

Situation	Behaviour
Standalone sampling, any thread count	Same local ids, edges, frontiers and stats as the previous path (tested at 4 threads).
1-thread process	No parallel dispatch; same code inline.
Fused feature-fetching sampling	Anchor change only; output unchanged.
legacy_recording = true	Previous path.
parallel_commit_min_edges	New public engine field; only changes when the parallel dedup runs, never the output.

Measured results

Same indexes, workload, ctxbench-style driver and machine as the previous PR. Sampling seeds/s, better of two interleaved rounds:

graph	threads	hops	main	previous PR	this PR	this PR / main	with feature fetch, main → this PR	cores busy, main → this PR	p99 per batch, main → this PR
100M	1	2	36,496	31,250	35,211	0.96×	18,587 → 17,794	1.0 → 1.0	30 → 29 ms
100M	1	3	7,899	9,960	10,846	1.37×	5,123 → 6,090	1.0 → 1.0	107 → 85 ms
100M	1	4	2,836	3,213	4,340	1.53×	2,094 → 2,769	1.0 → 1.0	250 → 191 ms
100M	1	5	331	578	648	1.96×	270 → 435	1.0 → 1.0	1.95 → 1.23 s
100M	1	6	154	226	269	1.74×	138 → 222	1.0 → 1.0	3.82 → 2.44 s
100M	8	2	51,546	53,191	56,180	1.09×	27,778 → 28,902	3.1 → 3.2	19 → 18 ms
100M	8	3	9,804	12,987	19,685	2.01×	7,728 → 13,021	2.5 → 3.2	67 → 41 ms
100M	8	4	4,444	5,359	8,865	1.99×	3,771 → 6,831	2.6 → 3.5	139 → 82 ms
100M	8	5	421	704	1,268	3.01×	393 → 1,041	1.8 → 3.2	1.34 → 0.52 s
100M	8	6	235	397	630	2.68×	225 → 563	1.9 → 3.1	2.38 → 1.00 s
1B	1	2	30,675	27,473	30,303	0.99×	15,385 → 15,773	1.0 → 1.0	34 → 33 ms
1B	1	3	6,821	8,913	9,980	1.46×	4,355 → 5,556	1.0 → 1.0	125 → 99 ms
1B	1	4	2,900	3,183	3,843	1.33×	2,033 → 2,449	1.0 → 1.0	266 → 224 ms
1B	1	5	327	554	612	1.87×	262 → 418	1.0 → 1.0	2.12 → 1.33 s
1B	1	6	162	236	282	1.74×	142 → 228	1.0 → 1.0	3.85 → 2.39 s
1B	8	2	48,544	48,077	53,191	1.10×	25,253 → 26,455	3.7 → 3.8	21 → 20 ms
1B	8	3	9,634	13,587	19,455	2.02×	7,396 → 12,255	2.8 → 3.6	73 → 46 ms
1B	8	4	4,651	5,525	8,000	1.72×	3,975 → 5,995	2.5 → 3.8	133 → 89 ms
1B	8	5	402	716	1,334	3.32×	374 → 1,066	1.8 → 3.3	1.42 → 0.51 s
1B	8	6	233	407	680	2.91×	223 → 593	1.9 → 3.5	2.39 → 0.92 s

Four concurrent processes of 8 threads each (aggregate sampling seeds/s; a scaled-down version of the 16×8 layout used on large hosts, limited by this shared machine):

graph	hops	main	this PR	×	with feature fetch ×	cores busy, main → this PR
100M	2	197,466	241,907	1.23×	1.14×	12.5 → 12.6
100M	3	39,897	74,313	1.86×	1.60×	9.6 → 12.7
100M	4	17,195	31,178	1.81×	1.62×	9.7 → 14.1
100M	5	1,409	4,261	3.02×	2.62×	6.4 → 10.9
100M	6	877	2,338	2.67×	2.48×	7.0 → 12.5
1B	2	191,409	218,250	1.14×	1.06×	13.7 → 14.0
1B	3	38,287	74,560	1.95×	1.58×	10.3 → 14.2
1B	4	17,512	33,131	1.89×	1.64×	10.1 → 15.2
1B	5	1,349	4,406	3.27×	2.82×	6.7 → 13.1
1B	6	748	2,218	2.96×	2.72×	7.1 → 13.7

Peak RSS per process rises by 0.0–0.3 GB over main.

Training path (sample_runtime_batch, link prediction)

Same setup as the training-path table in #101: sample_runtime_batch with a fresh SamplingState per batch, 1B index, two-tower link-prediction stage table, b=512 with one engine negative per row (three seed groups of 512), features skipped unless noted. Per-batch time, previous path (legacy_recording = true) → this PR, better of three interleaved rounds of 4 batches; identical node counts in every cell.

threads	2 hops	3 hops	4 hops	5 hops	6 hops
2	96 → 87 ms (1.10×)	714 → 470 ms (1.52×)	1.64 → 1.11 s (1.48×)	7.23 → 3.90 s (1.86×)	16.5 → 9.31 s (1.78×)
8	79 → 67 ms (1.17×)	676 → 371 ms (1.82×)	1.46 → 0.88 s (1.66×)	6.74 → 3.01 s (2.24×)	14.9 → 7.02 s (2.12×)

With feature fetch for every sampled node (8 threads): 114 → 101 ms (1.13×) at 2 hops, 1.89 → 1.28 s (1.48×) at 4 hops.

50B index

Hops	2	3	4	5	6
seeds/s (best config, #102)	392k	169k	98k	13.8k	9.3k
vs main (best config)	1.10×	1.16×	1.19×	1.6×	1.4×

50B commerce, i7ie.48xlarge, warm, b=512, numa_balancing=0, identical row counts.

Same published synthetic graph as #93 (49,997M rows, 7 node types × 171 shards, 2.6 TB index). Machine: 192 vCPU, 1.5 TiB RAM, local NVMe, kernel.numa_balancing=0. Harness: P processes × T threads, sampling plus feature fetch for every sampled node, b=512, warm (a warm-up pass, then the same seeds measured). main and this PR return identical row counts in every cell.

Seeds/s, main → this PR:

hops	64×3	16×8	96×2
2	273,570 → 374,560 (1.37×)¹	200,983 → 219,036 (1.09×)	356,736 → 392,243 (1.10×)
3	118,877 → 161,784 (1.36×)	61,169 → 84,639 (1.38×)	146,194 → 169,197 (1.16×)
4	70,050 → 92,251 (1.32×)	29,632 → 44,783 (1.51×)	81,754 → 97,583 (1.19×)
5	7,850 → 13,750 (1.75×)	3,619 → 6,713 (1.86×)	8,560 → 11,687 (1.37×)
6	5,729 → 8,740 (1.53×)	2,114 → 4,146 (1.96×)	6,777 → 9,319 (1.38×)

Best config per hop:

hops	main	this PR	#93 run, b=512	#93 run, best config	DuckDB 2.0 native 8×16 / 8×24
2	356,736	392,243	387,683 (64×3)	421,042 (64×3)	37,332 / 42,651
3	146,194	169,197	—²	143,935 (96×2)	24,549 / 27,605
4	81,754	97,583	77,214 (96×2)	77,316 (96×2)	15,867 / 15,045
5	8,560	13,750	—²	10,919 (96×2)	4,184 / 5,002
6	6,777	9,319	6,774 (96×2)	6,582 (96×2)	3,020 / 3,138

¹ main's measured pass still read 6.7 GB from disk in this cell, so part of the gain is warm-up; 96×2 is the cleaner 2-hop comparison.
² The #93 run measured b=512 only at 2, 4 and 6 hops. Its best-config column is the 192-vCPU row of #93's table, at batch sizes tuned per hop.

At 6 hops, p99 per batch at 16×8 drops from 4.2 s to 2.1 s.

Each cell above was measured once; repeated runs on this machine type have shown about 15% run-to-run spread, so the 1.10–1.19× best-config gains at 2–4 hops are within noise; the 1.4–1.6× gains at 5–6 hops are clearly real. The gains are smaller than on the 64-vCPU machine because the best 50B configs run many processes with 2–3 threads each: there the commit phase is a smaller share and page faults dominate. The 16×8 layout, where commit work is concentrated in fewer processes, gains 1.4–2.0× at 3–6 hops.

Known limitations

• One seed that holds most of a layer's edges (a hub seed) lands in a single range, so its dedup runs on one thread and caps the parallel speedup for that request.
• With 8 threads at 5–6 hops, the remaining serial work is the in-order id assignment and edge output (about 60% of commit time) and the counting sort (about 20%).

Testing Done

• Both equivalence tests from diskgraph: record sampled edges through per-seed maps #101 also run with parallel_commit_min_edges = 1, so at 4 threads every multi-seed layer takes the parallel branch, and compare that run with the previous path.
• tests/per_seed_recording_real_graph.rs (4 threads) passes on both indexes at 2–6 hops.
• cargo test -p diskgraph-engine, -p diskgraph-ingest --test seed_table_ingest --test resolution_runtime_batch and -p diskgraph-pybind pass; cargo clippy --workspace --all-targets -- -D warnings is clean.

Checklist

• No breaking changes.
• Unit tests added.

🤖 Generated with Claude Code

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

No actionable comments were generated in the recent review.

Recent review info

Run configuration

• Source mode: Auto
• Repository: NVIDIA/kumo-engine
• Branch: vedaant/dg-parallel-per-seed
• File: .coderabbit.yaml
• Commit: 9bfbf06
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: 0c07bd4e-1a6d-40d2-8cb8-0c80b0363205

Commits

Reviewing files that changed from the base of the PR and between a15177e and 9bfbf06.

Files selected for processing: 3

• packages/diskgraph/diskgraph-engine/src/engine.rs
• packages/diskgraph/diskgraph-engine/src/state.rs
• packages/diskgraph/diskgraph-engine/src/tests_unit.rs

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.

Walkthrough

The sampling path precomputes seed timestamp slices by node-type ID and passes them through layer processing for temporal anchor lookup. Edge commitment partitions destination deduplication when the engine has multiple threads and the edge count meets the configured threshold of 32,768. Commit scratch uses atomic result storage and per-partition buffers. Tests cover the updated timestamp argument and compare parallel-commit results with legacy recording.

Sequence Diagram(s)

sequenceDiagram
  participant DiskGraphEngine
  participant commit_edges
  participant for_each_part
  participant ThreadPool
  participant dedup_seeds
  participant SeedMaps
  participant TypeState
  DiskGraphEngine->>commit_edges: Commit edges
  commit_edges->>for_each_part: Partition seed-map ranges
  for_each_part->>ThreadPool: Run partition callbacks when multiple partitions apply
  ThreadPool->>dedup_seeds: Deduplicate destinations per partition
  dedup_seeds->>SeedMaps: Record pending or existing results
  commit_edges->>TypeState: Assign IDs to new nodes
  commit_edges->>SeedMaps: Resolve pending entries

Loading

Priority: Normal

Merge Risk: Minimal · up to 9bfbf

No identified issue needs resolution before merge; proceed with normal checks.

Pre-merge checks: 2 · 1

Failed checks: 1 inconclusive

Check name	Status	Explanation	Resolution
Change Map Design Review		The required authored PR conversation comment beginning <!-- kumo-pr-change-map --> is unavailable. Repository search found only the template, checker, and review instructions, not the comment. The …	Provide the complete PR conversation comment beginning <!-- kumo-pr-change-map -->, including its Change map and verification details. Then assess each proposed change against the linked problem and the reviewed diff.

Passed checks: 2 passed

Check name	Status	Explanation
Pr Size		The pull request changes 4 files. This count does not exceed 40.
Public Resource Access		No added or modified external resource reference was found. The diff changes Rust code only. The new gxhash import uses the dependency already declared unchanged in `packages/diskgraph/diskgraph-eng…

Full details: Public Resource Access

Explanation

No added or modified external resource reference was found. The diff changes Rust code only. The new gxhash import uses the dependency already declared unchanged in packages/diskgraph/diskgraph-engine/Cargo.toml; no URL, repository location, or package source location was added or changed.

Full details: Change Map Design Review

Explanation

The required authored PR conversation comment beginning <!-- kumo-pr-change-map --> is unavailable. Repository search found only the template, checker, and review instructions, not the comment. The supplied description is explicitly truncated, and the summary does not provide the map's necessity, abstraction-layer rationale, focused verification, or stopping condition. The diff confirms four changed files and a substantial engine.rs refactor, but it cannot be assessed against the missing proposed design.

Resolution

Provide the complete PR conversation comment beginning <!-- kumo-pr-change-map -->, including its Change map and verification details. Then assess each proposed change against the linked problem and the reviewed diff.

Finishing Touches: 💡 1

🛠️ Fix failing CI checks: 💡

• Commit to this branch
• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

[Vedaant-J]

Change map

File	Function / scope	What changed	Why
packages/diskgraph/diskgraph-engine/src/engine.rs	commit_edges, new for_each_part, dedup_seeds, new field DiskGraphEngine::parallel_commit_min_edges; sample_with_cutoff, process_layer, process_node_relations, sample_segment_edges	Per-seed dedup and map update run on the engine pool over contiguous seed ranges (4 per thread) when threads > 1 and the layer has at least parallel_commit_min_edges edges (default 32,768, with a comment on where it comes from). Destination type states are held across the in-order pass. Seed anchors are resolved once per request and passed as a slice table.	Per-seed work was the serial bottleneck with multiple threads; the anchor lookup hashed a type-name string per sample. The threshold is a field so tests can force the parallel path.
packages/diskgraph/diskgraph-engine/src/state.rs	CommitScratch, new CommitPart, reset and its test	The result array is atomic; per-range new-node lists, capped on reset.	Shared by parallel ranges without unsafe code.
packages/diskgraph/diskgraph-engine/src/pseudo_time_tests.rs	sample_segment_edges call	Passes an empty anchor table.	Signature change.
packages/diskgraph/diskgraph-engine/src/tests_unit.rs	assert_per_seed_matches_legacy	Runs at 4 threads with up to 700 seeds, and also with parallel_commit_min_edges = 1 so every multi-seed layer takes the parallel branch.	Coverage of the parallel path no longer depends on the RNG or the plan.

[aaron-kumo]

The parallel split and the anchor change look correct: threads write to separate AtomicU32 indexes, and rayon's join orders those writes before they're read. Posted two small notes inline.