Implement session persistence and checkpointing
Release / create-release (push) Successful in 2m32s
Release / build-linux-x86_64 (push) Successful in 8m21s
Release / build-macos-arm64 (push) Successful in 2m1s
ci.yml / build (pull_request) Successful in 3m45s

Introduces session persistence via the `--session` option, allowing sampled sets to be restored. Implements chunked dumps and shared checkpointing for alignment and tally artifacts, ensuring state restoration upon interruption. Refines sampling logic and introduces serialization mechanisms for state management.
This commit is contained in:
Eric Coissac
2026-09-12 07:48:51 +02:00
parent 7861e886e7
commit 2693a6af09
12 changed files with 418 additions and 148 deletions
+96 -47
View File
@@ -2729,9 +2729,10 @@ anywhere in it. `Session::open(dir, params: &[u8], force: bool)` handles
directory lifecycle and the params-conflict check against an opaque byte
blob; `Session::store(name, bytes)`/`Session::restore(name) -> Option<Mmap>`
handle a checksummed artifact cache. `obikphylo`
(`siblings/algorithms/snp_distance.rs`) is the consumer: it serializes its
own `PairwiseTally`/`PartitionDispersion` to bytes (via `rkyv`) and hands
them to `obiksession`, which never sees their type.
(`siblings/algorithms/session_cache.rs`, shared by `snp_distance.rs` and
`sankoff.rs` — see below) is the consumer: it serializes its own
`PairwiseTally`/`PartitionDispersion`/`SnpAlignment` to bytes (via `rkyv`)
and hands them to `obiksession`, which never sees their type.
**Serialization**: `rkyv` 0.8.18, added as planned. One deviation from the
original zero-copy pitch, scoped down deliberately: artifacts are restored
@@ -2763,51 +2764,99 @@ of `obiksession` itself.
`store_tally`) — `ratio_ceiling`/`gamma_shape` remain pure post-processing
over whatever tally is in hand, cached or fresh, exactly as planned.
**Locking**: `obiksession::Session` holds one `obisys::DirLock` for the
*whole session directory*, for the `Session` value's entire lifetime — not
the per-artifact-entry locking originally sketched. Simpler, and still
fully correct/crash-safe (same OS-auto-release guarantee `DirLock` already
provides): a second `obikmer phylo` process pointed at the same
`--session DIR` blocks until the first releases it, rather than each
artifact being independently lockable. Verified end-to-end (throwaway
8-genome index): same params → second run restores in `0ms` vs. the first
run's `~150-200ms`, output matrices identical; different params without
`--session-force` → clean error, process exits; with `--session-force` →
cache discarded, fresh sample computed and stored under the new params.
**Locking, reconsidered and kept as the whole-session `DirLock`.** First
written up as a corner cut ("per-entry locking was sketched, not
implemented"). On actually working through the concrete race, whole-session
locking turned out to be the *correct* choice, not a simplification:
`--session-force` resets a session by deleting its whole `artifacts/`
directory. If the directory lock were only held briefly (during
`Session::open`, then released so a long-running `Session` doesn't block
other processes from touching *different* artifacts), a second process
calling `open(..., force: true)` could acquire the lock, see mismatched
params, and wipe `artifacts/` **while a first process's already-open
`Session` is mid-checkpoint** — genuine corruption risk, not just wasted
work. Holding the lock for the `Session`'s entire lifetime makes that
sequence impossible: a `force` reset must wait for every currently-open
`Session` on that directory to be dropped first. The concurrency this
gives up (two *different* processes reading/writing *different* artifacts
of the same session at once) isn't exercised by anything today — only one
code path (`snp_distance`/`sankoff_bundle`, sequentially per invocation)
uses sessions at all.
**Explicitly not done in this pass** (each a real follow-up, not silently
dropped):
- **Per-layer chunked/resumable dumps — genuinely missing, not a
justified tradeoff.** Each artifact is a single store/restore of the
whole tally; a crash mid-`sample_index` loses the *entire* in-progress
call, however far it had gotten. An earlier pass at this note argued the
gap was fine because "the real cost is the I/O scan, not the tally
serialization, and redoing the tally is cheap" — that's a non sequitur,
caught and corrected on review: the cost that matters for resumability
is however much of the *scan* was already done when it died, not how
cheap the final tally is to serialize. That argument was extrapolated
from an 8-genome throwaway smoke test (~150-200ms total), not verified
against the actual scale this project targets (hundreds of partitions,
potentially hours — see the `merge_partitions` runs discussed earlier in
this project). On a long real run, losing 90% of a completed scan to a
crash is exactly the expensive case per-layer checkpointing exists to
avoid, and nothing here measures whether that's rare or common in
practice. Left undone for the same reason as the `SankoffBundle` gap
below (time, not a technical blocker) — not because the resumability
question was settled in favor of skipping it.
- **`SankoffBundle` integration.** `--session` only wires into the
standalone `snp_distance()` path today. `--sankoff`/`--tnt`/`--phyg`/
`--iqtree` still always resample, even with `--session` given — their
`SnpAlignment`/`BasePairTally`/`CardinalityTally` aren't cached. Same
mechanism would extend to them; not done yet.
- **Exhaustive-sample caching for `--subsample`-varying exploration.** The
design discussion's suggestion — cache the exhaustive (`n = None`)
sample once and subsample *from* it on each run instead of one session
per exact `n` — wasn't implemented; a session today is still scoped to
one exact parameter tuple, `--subsample` included.
- Only the whole-session lock (see above), not the finer per-entry
locking originally sketched for reducing contention between processes
sharing one `--session` for *different* artifacts.
**Per-layer chunked/resumable dumps — implemented.** `subsample::sample_index`
gained a `skip_layers: usize` parameter (skip the first N (partition,
layer) pairs, in its own deterministic iteration order, entirely — no
I/O, no `on_layer` call) and its `on_layer` callback gained a `raw_index`
argument (that pair's 0-based position in the same order), so a caller can
tell exactly how far a run has gotten. `snp_distance`/`sankoff_bundle`
checkpoint their tally/dispersion (and, for `sankoff_bundle`, the
in-progress `SnpAlignment`) to the session every
`session_cache::CHECKPOINT_INTERVAL_LAYERS` (8) layers or
`CHECKPOINT_INTERVAL` (30s) of wall time, whichever comes first — bounding
how much completed scanning a crash can lose, without paying a full
tally-rewrite on every one of what can be thousands of layers. A `progress`
artifact records how many layers are reflected in the checkpoint, with a
`PROGRESS_COMPLETE` (`u64::MAX`) sentinel once every layer has been
processed (distinct from any real layer count, so "finished" is never
confused with "stopped after N layers where N happens to equal the total").
Verified with a real kill: built a 10-genome/64-partition index, ran
`--distance snp-jc --session DIR` under `timeout 0.4`, confirmed a
checkpoint at `progress=8` (exactly `CHECKPOINT_INTERVAL_LAYERS`) was on
disk, then re-ran the same command and confirmed it resumed from layer 8
and completed successfully — 0-effort verification that the mechanism
itself is real, not just plausible-sounding.
**Important correction from that same test — resume is *not* bit-reproducible,
and shouldn't be.** The output matrix from the resumed run differed from
an uninterrupted control run against the same index/params. Root cause:
`sample_layer`'s random draws come from `rand::rng()`, a per-process,
OS-seeded generator with no continuity across a process boundary — the
layers processed before a kill and the layers processed after a resume
come from two independent random streams. An earlier revision of this doc
(and of the code comments) claimed resume "reproduces the exact same
sample an uninterrupted run would have" — that's wrong, and was corrected
after the user pointed out the actual reasoning error: true randomness has
no notion of continuity in the first place, so there was never a
correctness requirement for a resumed run to match an uninterrupted one
bit-for-bit. The sample produced after a resume is exactly as legitimate a
random draw as one from an uninterrupted run — just a *different* one,
same as running the command twice without `--session` already gives two
different samples today. That variance is explicitly wanted (see below),
not a defect to fix. Comments in `subsample.rs`/`session_cache.rs` were
corrected to state this plainly instead of the false reproducibility claim.
**Explicitly declined: deterministic/seeded sampling.** Raised as a
possible fix for the above (seed `sample_layer`'s RNG from a per-session
value plus `(partition, layer)`, making a resume reproduce bit-identical
results, and, if applied everywhere, making *any* two runs with identical
parameters produce identical samples). The user explicitly declined this
for now: repeated runs are relied on to measure sampling variance, and a
`--seed` flag would need its own explicit design if wanted later — this is
not an accidental gap, it's a stated preference to revisit only if asked.
**`SankoffBundle` integration — implemented.** `sankoff_bundle` accepts
the same `session: Option<&Session>` and uses the *same* artifact names
(`session_cache::{TALLY_ARTIFACT, DISPERSION_ARTIFACT, PROGRESS_ARTIFACT}`)
as `snp_distance`, plus its own `ALIGNMENT_ARTIFACT` (`sankoff.rs`) folded
in lockstep with the tally at every checkpoint (same `on_layer` call
produces both, so they can never drift out of sync). Consequence verified
directly: a tally checkpointed by `--sankoff --session DIR` is restored by
a *later, separate* `--distance snp-k2p --session DIR` invocation with no
`--sankoff` at all, and vice versa — cross-consumer reuse, not just
same-command reuse, exactly the synergy this design aimed for.
**Per-artifact caching, as designed**: only the tally/dispersion (plus,
for `sankoff_bundle`, the alignment) are cached — `ratio_ceiling`/
`gamma_shape` remain pure post-processing over whatever tally is in hand,
cached or fresh, exactly as planned.
**Still not done — genuine follow-up, not a justified tradeoff:**
**exhaustive-sample caching for `--subsample`-varying exploration.** The
design discussion's suggestion — cache the exhaustive (`n = None`) sample
once and subsample *from* it on each run instead of one session per exact
`n` — wasn't implemented; a session today is still scoped to one exact
parameter tuple, `--subsample` included.
### Output format: PHYLIP-relaxed by default for the distance matrix