obikmer/docmd/implementation/merge_parallelism.md

Merge parallelism and memory pressure

Problem observed

Running obikmer merge over 109 indexes (108 sources + 1 bootstrap) on a 192-core machine produces a fatal OOM during the merge_partitions stage:

memory allocation of 9126805520 bytes failed

A single allocation of ~8.5 GB fails. This is not an aggregate; it is one malloc call from hashbrown during a HashMap resize.

---

Root cause

The merge pipeline per partition

source unitigs.bin
  → iter_indexed_canonical_kmers()
  → GraphDeBruijn::push()       ← HashSet<u64> + 1 byte flags, all in RAM
  → compute_degrees_and_mark_starts()
  → try_for_each_unitig()
  → unitigs.bin (new layer)
  → Layer::build() → MPHF + evidence

GraphDeBruijn is a FastHashMap<CanonicalKmer, AtomicU8> — a HashSet<u64> with one flag byte per node. Neighbor lookup is implicit: 4 probes into the same map. No edges are stored. The full kmer set of one partition must reside in RAM simultaneously to compute degrees and mark unitig starts.

The matrix builders that follow (pass 2) are mmapped files — they do not consume significant RAM. The pressure is entirely in pass 1.

Unbounded Rayon parallelism

With 192 cores, Rayon ran up to 192 partitions concurrently. Each partition built its own GraphDeBruijn accumulating all kmers absent from the destination. Peak memory = 192 × peak_partition_hashset.

The 8.5 GB single allocation

hashbrown allocates the entire backing array in one call when rehashing. At load factor 7/8: capacity × (sizeof(K,V) + 1 control byte). For (u64, AtomicU8) with alignment: ~16 bytes per slot.

9 127 MB / 16 bytes ≈ 570 M slots → ~380 M new kmers in one partition

Plausible for the largest partition of 108 Salix/Betula sources (~450 Mbp each).

---

Partition size distribution

obikmer utils --partition-stats measures the sum of unitigs.bin file sizes per partition across all source indexes (pure stat() syscalls, negligible cost).

Observed on a 9-genome pilot (256 partitions):

Stat	Value
min	30.5 MB
max	232.1 MB
mean	40.1 MB
median	37.2 MB
p95	47.1 MB
max/median ratio	6.23×

The distribution is bimodal with a heavy tail:

• 238/256 partitions in a narrow 30–50 MB band
• 4 structurally extreme partitions (3–6× the median): 221, 233, 135, 191

These correspond to minimizers over-represented in repetitive regions shared across all sources. They are extreme in every run on this dataset.

With 109 sources, outlier partitions do not scale linearly: only kmers absent from the destination enter the GraphDeBruijn, and inter-source overlap is high for closely related species. Partition 221 is the likely trigger for the 8.5 GB crash.

---

Solution: LFD scheduling + memory budget semaphore

Principle

Pre-sort partitions by decreasing estimated size (First Fit Decreasing — FFD), then schedule them through a continuous memory budget semaphore. Each worker acquires an estimated cost before starting and releases it on completion.

Large partitions run first when the full budget is available; small partitions fill the gaps. No hard outlier threshold is needed.

MemoryBudget (obisys)

pub struct MemoryBudget { … }

impl MemoryBudget {
    pub fn new(total: u64) -> Self;
    pub fn acquire(&self, cost: u64);  // blocks until budget available
    pub fn release(&self, cost: u64);
    pub fn peak_active(&self) -> usize;
}

Non-deadlock guarantee: when active == 0, acquire always succeeds regardless of cost. Without this, a partition whose estimated cost exceeds the total budget would block forever.

Adaptive expansion factor

The expansion factor converts raw unitigs.bin bytes into an estimated GraphDeBruijn RAM footprint. hashbrown stores each kmer as (u64, AtomicU8) ≈ 16 bytes/kmer at 7/8 load factor; unitig files encode ≈ 2 bits/base. The ratio depends on average unitig length (short unitigs: ~2×; long unitigs: up to ~50×).

Phase 1 — sequential pilot (worst partition)

The largest partition runs alone first. Its actual g.len() seeds the expansion factor before any parallel job starts. FALLBACK_EXPANSION = 4× is used only for empty partitions.

let worst_g_len = dst_partition.merge_partition(worst_id, …)?;
//                              ↑ now returns SKResult<usize> (was SKResult<()>)

let seed_expansion = worst_g_len as u64 * 16 * 1000 / worst_bytes;
let max_expansion = AtomicU64::new(seed_expansion);

Phase 2 — parallel with adaptive updates

order[1..].into_par_iter().for_each(|&i| {
    let cost = partition_sizes[i] * max_expansion.load(Relaxed) / 1000;
    budget.acquire(cost);
    let g_len = dst_partition.merge_partition(i, …)?;
    budget.release(cost);          // releases estimated cost, not actual

    let actual = g_len as u64 * 16 * 1000 / partition_sizes[i];
    max_expansion.fetch_max(actual, Relaxed);  // always pessimistic (max)
});

budget.release(cost) uses the estimated cost, not the actual one. The budget tracks reservations, not physical RAM; each partition pays what it promised at acquisition.

On the safety margin

There is no separate multiplier k. It is redundant with budget_fraction: both reduce effective concurrency by the same amount. A single parameter is easier to calibrate. budget_fraction = 0.5 (default) reserves half of available RAM for the OS, MPHF build, pass 2, and estimation error.

--budget-fraction is exposed as a CLI flag — the only escape hatch for pathological cases (extreme repetitive content, unusually long unitigs) that still cause OOM.

RAM source

obisys::available_memory_bytes() — wraps sysinfo::System::available_memory(), falls back to total / 2 on macOS when the memory compressor returns 0.

---

Diagnostic report

After the parallel phase, merge_partition emits a structured report via tracing::info!:

─── merge_partitions memory report ───
  available RAM : 512.0 GB   budget 50% = 256.0 GB
  expansion factor — seed: 4.2×  final max: 6.1×  (mean: 1.8×  median: 1.6×)
  peak concurrent workers: 42
  expansion factor distribution (256 partitions with data):
     0.50× –  1.25× │██████████████████████████████ 148
     1.25× –  2.00× │████████████████████████        82
     …
     5.50× –  6.25× │█                                2
  top partitions by actual expansion factor:
    partition  221 : 6.10×  (232.1 MB unitigs → 48M kmers, reserved at 4.20×)
    partition  135 : 5.82×  (127.3 MB unitigs → 24M kmers, reserved at 4.20×)
    …
──────────────────────────────────────

Fields useful for diagnosis:

Field	Interpretation
seed vs final max expansion	gap indicates partitions with higher expansion than the worst-by-size
reserved at X×	the factor used at acquisition; if much lower than actual, the budget was under-reserved for that partition
peak concurrent workers	effective parallelism achieved under the budget constraint
mean / median expansion	typical dataset characteristic; stable across runs on the same data

---

Parameters

Parameter	Default	CLI flag	Notes
fallback_expansion	4×	—	seed for empty partitions only
budget_fraction	0.5	--budget-fraction	reduce if OOM persists
RAM source	obisys::available_memory_bytes()	—	falls back to total/2 on macOS