diff --git a/DevDocMD/theory/evolutionary_distances.md b/DevDocMD/theory/evolutionary_distances.md index 8036cb8d..813fa0fb 100644 --- a/DevDocMD/theory/evolutionary_distances.md +++ b/DevDocMD/theory/evolutionary_distances.md @@ -2715,38 +2715,99 @@ sample)` and the `snp_distance` stage timer reads `0ms` (`formula(kind)` is `O(n²)` post-processing, no I/O), vs. ~150-230ms for a fresh sample on the same tiny test index. -#### Future direction, not implemented: explicit `--session` +#### Explicit `--session` (implemented 2026-09-12) -Raised in discussion, not started. The fix above only covers reuse -*within one process*. A further idea: a `--session DIR` flag naming a -directory (outside the index) that persists the CLI selection parameters -plus every intermediate artifact `sample_index` would otherwise -recompute — the site set itself, `PairwiseTally`/`PartitionDispersion`, -the pseudo-alignment, etc. — across *separate* `obikmer phylo` -invocations, e.g. running `--sankoff` today and `--distance snp-k2p` -tomorrow against the identical sample. Without `--session`, an implicit -*temporary* session would still be created (scoped to the index directory, -matching the tmp-cache design discussed earlier in this file's α-estimation -section) — same mechanism, just not named/kept by the user. +The in-process fix above only covered reuse *within one command*. +`--session DIR` (`cmd/phylo/args.rs`) names a directory (outside the +index) that persists the `snp-*` `--distance` sample — the +`PairwiseTally`/`PartitionDispersion` pair `sample_index` would otherwise +rebuild from scratch — across *separate* `obikmer phylo` invocations. -Two open design points if this is picked up: -- **Cache-key validity**: as established for the tally cache above, the - content depends on `n`/`free_loss`/`no_ambiguity`/excluded-set/ - `entropy_bias` — a session is only reusable for the exact tuple it was - built under. `--subsample` in particular is the parameter most likely to - change between exploratory runs, so an *exhaustive* (`n = None`) session - that later runs subsample *from*, rather than one session per exact `n`, - would likely see far more real reuse. -- **Concurrency-safe cleanup**: `obisys::DirLock` (`lock.rs`) is the - existing pattern to follow — an OS advisory lock (`flock`/`LockFileEx`), - auto-released on process exit *including a crash*, no stale-lock cleanup - logic needed. Applied per cache/session entry (not just once for the - whole index as `DirLock` does today for annex writes): a process - deciding whether to reclaim an old temporary session first tries to - acquire that entry's lock — success means nobody's using it, safe to - delete; failure means another process holds it, leave it alone. Avoids - the hazard of one process deleting another concurrently-running - process's temp session, which a naive "wipe at startup" would risk. +**Crate placement, as planned**: a new **`obiksession`** crate +(`src/obiksession/`), domain-agnostic — no `PairwiseTally`/"site" concept +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` +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. + +**Serialization**: `rkyv` 0.8.18, added as planned. One deviation from the +original zero-copy pitch, scoped down deliberately: artifacts are restored +via `rkyv::from_bytes` (full owned deserialize) rather than `rkyv::access` +(zero-copy over the `mmap`). True zero-copy would require every +`PairwiseTally` query method (`categories`, `base_freq`, `pair`, ...) to +work generically over `Archived` as well as the owned type +— a separate, larger change not needed to get the actual win (skipping +`sample_index`'s expensive re-scan of the sibling annex; deserializing an +already-in-memory-sized count array is comparatively cheap CPU, not I/O). +`#[derive(Archive, Serialize, Deserialize)]` was added to `PairStats`, +`PairwiseTally`, and `PartitionDispersion` (`pairwise.rs`) for this. + +**Startup behavior, as designed**: `Session::open` creates `DIR` if +missing (fresh, params recorded); if `DIR` already holds different saved +params, `cmd/phylo` reports a hard error and exits rather than silently +preferring one side, unless `--session-force` is given (which discards the +directory's cached artifacts and starts over under the new params) — +implements exactly the "explicit escape hatch, not silent override" +decision from the design discussion. `obiksession` itself doesn't decode +either side's bytes to produce a field-by-field diff (true to "opaque +blob," left to whoever needs it); `cmd/phylo`'s conflict message is +currently a single generic sentence listing which flags could be the +cause, not a computed diff — an honest v1 simplification, not a limitation +of `obiksession` itself. + +**Per-artifact caching, as designed**: only `PairwiseTally`/ +`PartitionDispersion` are cached (`snp_distance.rs`'s `restore_tally`/ +`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. + +**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. ### Output format: PHYLIP-relaxed by default for the distance matrix diff --git a/UserDocMD/usage/phylo.md b/UserDocMD/usage/phylo.md index 916362ca..d74af9e6 100644 --- a/UserDocMD/usage/phylo.md +++ b/UserDocMD/usage/phylo.md @@ -202,7 +202,18 @@ A family is eligible for a genome pair $(i,j)$ only if both genomes carry exactl `--subsample`, `--free-loss`, `--no-ambiguity`, `--entropy`/`--entropy-sd` are shared by `--pseudo-alignment`, `--sankoff` (and everything it implies: `--tnt`/`--phyg`/`--iqtree`), and a `snp-*` `--distance` value — one draw feeds all of them in a single invocation. `--subsample` is mandatory for `--pseudo-alignment`/`--sankoff`; for a `snp-*` `--distance` value it is optional (omitted means every non-monomorphic family in the index, not an approximation). -Combining `--sankoff` (or `--tnt`/`--phyg`/`--iqtree`) with a `snp-*` `--distance` value in the same command reuses that one draw for both — the distance and the Sankoff calibration/alignment are guaranteed to be computed from the *identical* set of sampled sites, never two independent samples, so the two outputs are directly comparable. This only holds within a single command; running them as two separate `obikmer phylo` invocations draws two independent samples even with the same flags. +Combining `--sankoff` (or `--tnt`/`--phyg`/`--iqtree`) with a `snp-*` `--distance` value in the same command reuses that one draw for both — the distance and the Sankoff calibration/alignment are guaranteed to be computed from the *identical* set of sampled sites, never two independent samples, so the two outputs are directly comparable. This only holds within a single command; running them as two separate `obikmer phylo` invocations draws two independent samples even with the same flags — unless `--session` is used (below), which currently covers a `snp-*` `--distance` value only, not `--sankoff`/`--tnt`/`--phyg`/`--iqtree`. + +### `--session`: reusing a sample across separate commands + +| Option | Default | Description | +|---|---|---| +| `--session DIR` | none | Persist the `snp-*` `--distance` sample in `DIR` so a later, separate `obikmer phylo` invocation with the exact same selection parameters restores it instead of resampling | +| `--session-force` | off | With `--session DIR`: overwrite its saved parameters and cached sample instead of erroring out when this run's parameters don't match. No effect without `--session` | + +`DIR` is created if it doesn't exist. If it already holds a sample built with different `--subsample`/`--free-loss`/`--no-ambiguity`/`--exclude-genome`/`--min-shared-family`/`--entropy`/`--entropy-sd` values than this run, the command exits with an error rather than silently using either the old or the new values — pass `--session-force` to discard the old sample and rebuild under the new parameters, or point `--session` at a different directory to keep both. + +Only the `snp-*` `--distance` sample (not `--sankoff`/`--tnt`/`--phyg`/`--iqtree`'s calibration/alignment) is persisted today. Without `--session`, nothing changes: every invocation resamples, as before. Without `--subsample`, every variable family (family size ≥ 2) is used. With `--subsample N`, roughly `N` families are kept instead, drawn in proportion to how many candidate families each part of the index actually holds, so the sample stays representative of the whole index. If the index has fewer than `N` candidate families, `--subsample` has no effect. diff --git a/src/Cargo.lock b/src/Cargo.lock index ed71e1af..2c9192a0 100644 --- a/src/Cargo.lock +++ b/src/Cargo.lock @@ -213,6 +213,29 @@ version = "3.20.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5d20789868f4b01b2f2caec9f5c4e0213b41e3e5702a50157d699ae31ced2fcb" +[[package]] +name = "bytecheck" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26333eeac754f0ad8a6bcd0eb0ac012156302e4e16b852b72ee399aea4f12c29" +dependencies = [ + "bytecheck_derive", + "ptr_meta", + "rancor", + "simdutf8", +] + +[[package]] +name = "bytecheck_derive" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "46d07918caa9eeaaf06b7873925c53a61daac173539b4f7715090745e44e4e69" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + [[package]] name = "byteorder" version = "1.5.0" @@ -1236,6 +1259,26 @@ dependencies = [ "windows 0.48.0", ] +[[package]] +name = "munge" +version = "0.4.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e17401f259eba956ca16491461b6e8f72913a0a114e39736ce404410f915a0c" +dependencies = [ + "munge_macro", +] + +[[package]] +name = "munge_macro" +version = "0.4.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4568f25ccbd45ab5d5603dc34318c1ec56b117531781260002151b8530a9f931" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "nanorand" version = "0.6.1" @@ -1499,6 +1542,7 @@ dependencies = [ "obikrope", "obikselect", "obikseq", + "obiksession", "obikstats", "obipipeline", "obiread", @@ -1541,6 +1585,7 @@ dependencies = [ "obikindex", "obikindexer", "obikseq", + "obiksession", "obipipeline", "obiread", "obiskbuilder", @@ -1549,6 +1594,7 @@ dependencies = [ "petgraph", "rand 0.10.2", "rayon", + "rkyv", "speedytree", "tempfile", "tracing", @@ -1615,6 +1661,17 @@ dependencies = [ "xxhash-rust", ] +[[package]] +name = "obiksession" +version = "0.1.0" +dependencies = [ + "memmap2", + "obisys", + "tempfile", + "tracing", + "xxhash-rust", +] + [[package]] name = "obikstats" version = "0.1.0" @@ -1963,6 +2020,26 @@ dependencies = [ "xxhash-rust", ] +[[package]] +name = "ptr_meta" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "743da816b98c921cdbe8628ef7381b76f25ecf4da599fc80aca90eae7ef70cc0" +dependencies = [ + "ptr_meta_derive", +] + +[[package]] +name = "ptr_meta_derive" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c8d9ca532f185d5d4db7a7c9d51420b452168ea1c2b913953281bd6fe1fcbd0" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + [[package]] name = "quote" version = "1.0.45" @@ -1990,6 +2067,15 @@ version = "0.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dc33ff2d4973d518d823d61aa239014831e521c75da58e3df4840d3f47749d09" +[[package]] +name = "rancor" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b534442d0fcdb55d66f373d9cac6d33b6293a2335bc2136dbd06ce0e87d2572" +dependencies = [ + "ptr_meta", +] + [[package]] name = "rand" version = "0.8.6" @@ -2151,6 +2237,15 @@ version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" +[[package]] +name = "rend" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "663ba70707f96e871406fe10d68128412e619b06d1d47cb91c3a4c6501176240" +dependencies = [ + "bytecheck", +] + [[package]] name = "ring" version = "0.17.14" @@ -2165,6 +2260,36 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "rkyv" +version = "0.8.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9776093b7ca170454ab1406954f7b7d97a57c51dc6c0642957fb2ef25c2d399" +dependencies = [ + "bytecheck", + "bytes", + "hashbrown", + "indexmap", + "munge", + "ptr_meta", + "rancor", + "rend", + "rkyv_derive", + "tinyvec", + "uuid", +] + +[[package]] +name = "rkyv_derive" +version = "0.8.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c25ef604ac7dd839d44d64648952ea23c97866f124ff671b0ed2cf3ad9bb06e" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + [[package]] name = "rustc-hash" version = "2.1.2" @@ -2359,6 +2484,12 @@ version = "0.3.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "703d5c7ef118737c72f1af64ad2f6f8c5e1921f818cdcb97b8fe6fc69bf66214" +[[package]] +name = "simdutf8" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" + [[package]] name = "smallvec" version = "1.15.1" @@ -2587,6 +2718,21 @@ dependencies = [ "serde_json", ] +[[package]] +name = "tinyvec" +version = "1.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cf0ded5c4e56918d8f8a339e1bb67d038d3bc6d144ac407904015ba2e4cde9b" +dependencies = [ + "tinyvec_macros", +] + +[[package]] +name = "tinyvec_macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" + [[package]] name = "tracing" version = "0.1.44" diff --git a/src/Cargo.toml b/src/Cargo.toml index b67a5f86..06573eb3 100644 --- a/src/Cargo.toml +++ b/src/Cargo.toml @@ -1,5 +1,5 @@ [workspace] resolver = "3" -members = ["obikseq", "obiread", "obiskbuilder", "obifastwrite", "obikmer","obikrope","obipipeline", "obiskio","obidebruinj", "obicompactvec", "obisys", "obikindex", "obikindexer", "obikquery", "obikdump", "obikfilter", "obikselect", "obikrebuild", "obikmerge", "obikstats", "obikidxcache", "obitaxonomy", "obikentropy", "obikphylo", "obikalgorithm"] +members = ["obikseq", "obiread", "obiskbuilder", "obifastwrite", "obikmer","obikrope","obipipeline", "obiskio","obidebruinj", "obicompactvec", "obisys", "obikindex", "obikindexer", "obikquery", "obikdump", "obikfilter", "obikselect", "obikrebuild", "obikmerge", "obikstats", "obikidxcache", "obitaxonomy", "obikentropy", "obikphylo", "obikalgorithm", "obiksession"] [profile.release] debug = 1 diff --git a/src/obikmer/Cargo.toml b/src/obikmer/Cargo.toml index 5ad6428a..402c7cf1 100644 --- a/src/obikmer/Cargo.toml +++ b/src/obikmer/Cargo.toml @@ -24,6 +24,7 @@ obikstats = { path = "../obikstats" } obikquery = { path = "../obikquery" } obikidxcache = { path = "../obikidxcache" } obikphylo = { path = "../obikphylo" } +obiksession = { path = "../obiksession" } obikrope = { path = "../obikrope" } obifastwrite = { path = "../obifastwrite" } obiskbuilder = { path = "../obiskbuilder" } diff --git a/src/obikmer/src/cmd/phylo/args.rs b/src/obikmer/src/cmd/phylo/args.rs index 10c173ae..305607bb 100644 --- a/src/obikmer/src/cmd/phylo/args.rs +++ b/src/obikmer/src/cmd/phylo/args.rs @@ -224,6 +224,27 @@ pub struct PhyloArgs { #[arg(long)] pub entropy_sd: Option, + /// Persist the `snp-*` `--distance` sample (the pairwise substitution + /// tally) in `DIR` across separate `obikmer phylo` invocations: a later + /// run against the same index with the exact same + /// `--subsample`/`--free-loss`/`--no-ambiguity`/`--exclude-genome`/ + /// `--min-shared-family`/`--entropy`/`--entropy-sd` restores it instead + /// of resampling. `DIR` is created if it doesn't exist. If it already + /// holds a sample built under *different* selection parameters, this + /// run is rejected with an error listing what changed, unless + /// `--session-force` is also given (which discards the old sample and + /// starts fresh under this run's parameters). Only covers a `snp-*` + /// `--distance` today — `--sankoff`/`--tnt`/`--phyg`/`--iqtree` don't + /// read or write a `--session` yet. + #[arg(long, value_name = "DIR")] + pub session: Option, + + /// With `--session DIR`: overwrite its saved selection parameters (and + /// discard its cached sample) instead of erroring out when this run's + /// parameters don't match. No effect without `--session`. + #[arg(long, requires = "session")] + pub session_force: bool, + /// Calibrate a 16-state Sankoff cost matrix (and its matching /// pseudo-alignment) from an already-built sibling annex — requires /// `--subsample `, and shares `--free-loss`/`--no-ambiguity`/ diff --git a/src/obikmer/src/cmd/phylo/mod.rs b/src/obikmer/src/cmd/phylo/mod.rs index 790ea7ba..aaee15d8 100644 --- a/src/obikmer/src/cmd/phylo/mod.rs +++ b/src/obikmer/src/cmd/phylo/mod.rs @@ -24,6 +24,43 @@ use phylip::write_phylip_relaxed; use sankoff::{write_sankoff_alignment_fasta, write_sankoff_matrix_csv, write_sankoff_params}; use tnt::write_sankoff_tnt; +/// Deterministic byte encoding of the `snp-*` `--distance` selection +/// parameters, for `obiksession::Session::open`'s params-conflict check. +/// `obiksession` treats this as an opaque blob (see its own crate docs) — +/// only this module needs to know the layout, and only well enough to +/// produce/compare it, never to pretty-print a per-field diff (v1 keeps +/// the conflict message generic rather than decoding both sides field by +/// field — see its use in `run` below). +fn session_params_blob( + n: Option, + free_loss: bool, + no_ambiguity: bool, + excluded: &[bool], + entropy_bias: Option, +) -> Vec { + let mut buf = Vec::new(); + match n { + Some(v) => { + buf.push(1); + buf.extend_from_slice(&(v as u64).to_le_bytes()); + } + None => buf.push(0), + } + buf.push(free_loss as u8); + buf.push(no_ambiguity as u8); + buf.extend_from_slice(&(excluded.len() as u64).to_le_bytes()); + buf.extend(excluded.iter().map(|&b| b as u8)); + match entropy_bias { + Some(EntropyBias { mu, sigma }) => { + buf.push(1); + buf.extend_from_slice(&mu.to_le_bytes()); + buf.extend_from_slice(&sigma.to_le_bytes()); + } + None => buf.push(0), + } + buf +} + pub use args::PhyloArgs; pub fn run(args: PhyloArgs) { @@ -236,6 +273,47 @@ pub fn run(args: PhyloArgs) { None }; + // ── `--session`: persist the snp-* sample across separate invocations ── + // Only covers `snp-*` `--distance` today (see `--session`'s own docs); + // `None` when `--session` isn't given, matching today's always-resample + // default exactly. + let session = args.session.as_ref().map(|dir| { + let params = session_params_blob( + args.subsample, + args.free_loss, + args.no_ambiguity, + &snp_exclude_mask, + entropy_bias, + ); + match obiksession::Session::open(dir, ¶ms, args.session_force) { + Ok(Ok(session)) => { + match session.outcome() { + obiksession::SessionOutcome::Created => { + info!("--session {}: starting fresh", dir.display()); + } + obiksession::SessionOutcome::Reused => { + info!("--session {}: found, will reuse its cached sample if present", dir.display()); + } + } + session + } + Ok(Err(_conflict)) => { + eprintln!( + "error: --session {} was built with different selection parameters \ + (--subsample/--free-loss/--no-ambiguity/--exclude-genome/--min-shared-family/\ + --entropy/--entropy-sd) than this run — pass --session-force to discard its \ + cached sample and start fresh, or point --session at a different directory", + dir.display() + ); + std::process::exit(1); + } + Err(e) => { + eprintln!("error opening --session {}: {e}", dir.display()); + std::process::exit(1); + } + } + }); + // ── SNP pseudo-alignment (`--pseudo-alignment`) ───────────────────────────── if args.pseudo_alignment { let Some(subsample_n) = args.subsample else { @@ -389,6 +467,7 @@ pub fn run(args: PhyloArgs) { &snp_exclude_mask, entropy_bias, gamma_shape, + session.as_ref(), ) } } diff --git a/src/obikphylo/Cargo.toml b/src/obikphylo/Cargo.toml index 1ebe4797..08a6a0dd 100644 --- a/src/obikphylo/Cargo.toml +++ b/src/obikphylo/Cargo.toml @@ -12,12 +12,14 @@ obicompactvec = { path = "../obicompactvec" } obikidxcache = { path = "../obikidxcache" } obiskbuilder = { path = "../obiskbuilder" } obipipeline = { path = "../obipipeline" } +obiksession = { path = "../obiksession" } kodama = "0.3.0" memmap2 = "0.9" ndarray = "0.17" petgraph = "0.6.4" rand = "0.10" rayon = "1" +rkyv = "0.8.18" speedytree = "0.1" tracing = "0.1.44" diff --git a/src/obikphylo/src/siblings/algorithms/pairwise.rs b/src/obikphylo/src/siblings/algorithms/pairwise.rs index 1691522a..a651428c 100644 --- a/src/obikphylo/src/siblings/algorithms/pairwise.rs +++ b/src/obikphylo/src/siblings/algorithms/pairwise.rs @@ -18,6 +18,7 @@ //! machines), not something to optimise away. use ndarray::Array2; +use rkyv::{Archive, Deserialize, Serialize}; use super::subsample::SurvivingFamily; @@ -30,7 +31,12 @@ use super::subsample::SurvivingFamily; /// [`BasePairTally`] the latter (they are *not* the same population — see /// `BasePairTally::same`'s own docs). `cardinality` is restricted to /// variable families only, matching [`CardinalityTally`]'s own scope. -#[derive(Default, Clone, Copy)] +/// +/// `Archive`/`Serialize`/`Deserialize`: part of [`PairwiseTally`]'s own +/// `--session` persistence (see that struct's docs) — `rkyv` needs every +/// field type to derive these too, hence deriving here on an otherwise +/// purely internal struct. +#[derive(Default, Clone, Copy, Archive, Serialize, Deserialize)] struct PairStats { subst: [[u64; 4]; 4], same_all: [u64; 4], @@ -40,6 +46,13 @@ struct PairStats { /// See the module docs. `pub(crate)`: consumed by /// `algorithms::sankoff`'s orchestration, not yet exposed past this crate. +/// +/// `Archive`/`Serialize`/`Deserialize` (`rkyv`): lets `--session` dump/ +/// restore a tally across separate CLI invocations as opaque bytes handed +/// to `obiksession::Session::store`/`restore` — `obiksession` itself knows +/// nothing about this type, it only stores/checksums/mmaps the bytes this +/// crate produces (see `obiksession`'s own crate docs on that split). +#[derive(Archive, Serialize, Deserialize)] pub(crate) struct PairwiseTally { n_genomes: usize, /// Upper triangle only (`i < j`), flat-indexed via [`Self::flat_index`]. @@ -300,7 +313,7 @@ pub struct CardinalityTally { /// applied here either, matching `reduce_pairwise`'s own raw per-pair fold /// (`PairStats` accumulates over every genome pair unconditionally; /// exclusion is a filter applied only in derived, post-hoc views). -#[derive(Default, Clone)] +#[derive(Default, Clone, Archive, Serialize, Deserialize)] pub(crate) struct PartitionDispersion { /// One entry per partition id seen so far (grows lazily — partitions /// are visited in increasing order in practice, but nothing here diff --git a/src/obikphylo/src/siblings/algorithms/snp_distance.rs b/src/obikphylo/src/siblings/algorithms/snp_distance.rs index 41d06920..4db4f8aa 100644 --- a/src/obikphylo/src/siblings/algorithms/snp_distance.rs +++ b/src/obikphylo/src/siblings/algorithms/snp_distance.rs @@ -23,12 +23,51 @@ use ndarray::Array2; use obikidxcache::index_cache::IndexCache; -use obikindex::OKIResult; +use obikindex::{OKIError, OKIResult}; +use obiksession::Session; use super::pairwise::{PairwiseTally, PartitionDispersion}; use super::sibling_family_size_histogram; use super::subsample::{EntropyBias, sample_index}; +const TALLY_ARTIFACT: &str = "pairwise_tally"; +const DISPERSION_ARTIFACT: &str = "partition_dispersion"; + +/// Restore a `PairwiseTally`/`PartitionDispersion` pair from `session`, if +/// both artifacts are cached there — `None` on any miss (not cached, I/O +/// error, or a deserialization failure), each treated identically: +/// something's missing or unusable, fall back to recomputing rather than +/// surfacing a hard error over what's meant to be a transparent cache. +/// +/// Full owned deserialization (`rkyv::from_bytes`), not zero-copy +/// (`rkyv::access`) — simpler to integrate with `PairwiseTally`'s existing +/// query methods, which all take `&self` by value semantics, not an +/// `Archived`. Still skips the expensive part (`sample_index` +/// re-scanning the sibling annex); the `O(n²)` deserialize of an already- +/// in-memory-sized byte buffer is comparatively cheap. Making this +/// genuinely zero-copy later would mean every `PairwiseTally` accessor +/// (`categories`, `base_freq`, ...) working generically over +/// `Archived` too — a bigger, separate change. +fn restore_tally(session: &Session) -> Option<(PairwiseTally, PartitionDispersion)> { + let tally_bytes = session.restore(TALLY_ARTIFACT).ok().flatten()?; + let dispersion_bytes = session.restore(DISPERSION_ARTIFACT).ok().flatten()?; + let tally = rkyv::from_bytes::(&tally_bytes[..]).ok()?; + let dispersion = + rkyv::from_bytes::(&dispersion_bytes[..]).ok()?; + tracing::info!("--session: restored cached SNP tally — skipping resampling"); + Some((tally, dispersion)) +} + +fn store_tally(session: &Session, tally: &PairwiseTally, dispersion: &PartitionDispersion) -> OKIResult<()> { + let tally_bytes = rkyv::to_bytes::(tally) + .map_err(|e| OKIError::Io(std::io::Error::other(e.to_string())))?; + session.store(TALLY_ARTIFACT, &tally_bytes).map_err(OKIError::Io)?; + let dispersion_bytes = rkyv::to_bytes::(dispersion) + .map_err(|e| OKIError::Io(std::io::Error::other(e.to_string())))?; + session.store(DISPERSION_ARTIFACT, &dispersion_bytes).map_err(OKIError::Io)?; + Ok(()) +} + /// One `snp-*` `--distance` value. `pub`: part of /// [`crate::siblings::extensions::SiblingExt::snp_distance`]'s public /// signature. @@ -231,6 +270,17 @@ fn formula(kind: SnpDistanceKind) -> fn(&PairwiseTally, usize, usize, Option, gamma_shape: GammaShape, + session: Option<&Session>, ) -> OKIResult> { // Fail fast, before paying for `sample_index`, not just inside // `distance_matrix` (which runs after sampling either way — the right @@ -251,33 +302,44 @@ pub(crate) fn snp_distance( )); } - let n_genomes = cache.meta().genomes().len(); - let mut tally = PairwiseTally::new(n_genomes); - let mut dispersion = PartitionDispersion::default(); - - let (target, entropy_bias) = match n { - Some(target) => (target, entropy_bias), + let cached = session.and_then(restore_tally); + let (tally, dispersion) = match cached { + Some(pair) => pair, None => { - let counts = sibling_family_size_histogram(cache)?; - let total_eligible = (counts[1] + counts[2] + counts[3]) as usize; - (total_eligible, None) + let n_genomes = cache.meta().genomes().len(); + let mut tally = PairwiseTally::new(n_genomes); + let mut dispersion = PartitionDispersion::default(); + + let (target, entropy_bias) = match n { + Some(target) => (target, entropy_bias), + None => { + let counts = sibling_family_size_histogram(cache)?; + let total_eligible = (counts[1] + counts[2] + counts[3]) as usize; + (total_eligible, None) + } + }; + + if target > 0 { + sample_index( + cache, + target, + free_loss, + no_ambiguity, + excluded, + entropy_bias, + |partition, _layer, survivors| { + super::pairwise::reduce_pairwise(&survivors, partition, &mut tally, &mut dispersion); + }, + )?; + } + + if let Some(session) = session { + store_tally(session, &tally, &dispersion)?; + } + (tally, dispersion) } }; - if target > 0 { - sample_index( - cache, - target, - free_loss, - no_ambiguity, - excluded, - entropy_bias, - |partition, _layer, survivors| { - super::pairwise::reduce_pairwise(&survivors, partition, &mut tally, &mut dispersion); - }, - )?; - } - distance_matrix(&tally, &dispersion, kind, gamma_shape) } diff --git a/src/obikphylo/src/siblings/extensions/sibling_ext.rs b/src/obikphylo/src/siblings/extensions/sibling_ext.rs index 758a9dfe..a1bbddc1 100644 --- a/src/obikphylo/src/siblings/extensions/sibling_ext.rs +++ b/src/obikphylo/src/siblings/extensions/sibling_ext.rs @@ -147,6 +147,13 @@ pub trait SiblingExt { /// ([`SnpDistanceKind::supports_gamma`]). `Auto` estimates `alpha` from /// the data itself (method-of-moments over per-partition substitution /// rate dispersion) rather than requiring a user-supplied value. + /// + /// `session`: `--session`, `None` disables persistence (always + /// resample). `Some` must already be open under these exact selection + /// parameters (the caller's responsibility, via + /// `obiksession::Session::open`'s params-conflict check) — restores a + /// cached tally when present instead of resampling, and stores a + /// freshly computed one back for next time otherwise. fn snp_distance( &self, kind: SnpDistanceKind, @@ -156,6 +163,7 @@ pub trait SiblingExt { excluded: &[bool], entropy_bias: Option, gamma_shape: GammaShape, + session: Option<&obiksession::Session>, ) -> OKIResult>; /// The Family Overlap annex — number of shared *variable* families per @@ -335,8 +343,9 @@ impl SiblingExt for IndexCache { excluded: &[bool], entropy_bias: Option, gamma_shape: GammaShape, + session: Option<&obiksession::Session>, ) -> OKIResult> { - snp_distance(self, kind, n, free_loss, no_ambiguity, excluded, entropy_bias, gamma_shape) + snp_distance(self, kind, n, free_loss, no_ambiguity, excluded, entropy_bias, gamma_shape, session) } fn family_overlap(&self) -> OKIResult { diff --git a/src/obiksession/Cargo.toml b/src/obiksession/Cargo.toml new file mode 100644 index 00000000..3ee4a5a2 --- /dev/null +++ b/src/obiksession/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "obiksession" +version = "0.1.0" +edition = "2024" + +[dependencies] +obisys = { path = "../obisys", default-features = false } +memmap2 = "0.9.11" +xxhash-rust = { version = "0.8.18", features = ["xxh3"] } +tracing = "0.1" + +[dev-dependencies] +tempfile = "3" diff --git a/src/obiksession/src/lib.rs b/src/obiksession/src/lib.rs new file mode 100644 index 00000000..1c7b8b31 --- /dev/null +++ b/src/obiksession/src/lib.rs @@ -0,0 +1,19 @@ +//! Domain-agnostic session persistence: a named directory that survives +//! across separate CLI invocations, holding an opaque parameter blob (used +//! to detect stale/mismatched reuse) plus a checksummed cache of named +//! binary artifacts. +//! +//! Deliberately knows nothing about what the parameters or artifacts +//! *mean* — no k-mer/genome/tally concept anywhere in this crate. A +//! caller (e.g. `obikphylo`, for its `--session` support around +//! `PairwiseTally`/`PartitionDispersion`/`SnpAlignment`) serializes its own +//! parameter struct and artifacts to bytes however it likes (`rkyv`, +//! `bincode`, ...) and hands raw `&[u8]` to this crate; [`Session`] only +//! deals in byte blobs, directory lifecycle, and locking. Same split +//! already used elsewhere in this project between a generic mechanism +//! crate and the domain crate that plugs into it (`obisys`/`obikindex`, +//! `obicompactvec`/`obikindexer`). + +mod session; + +pub use session::{ParamsConflict, Session, SessionOutcome}; diff --git a/src/obiksession/src/session.rs b/src/obiksession/src/session.rs new file mode 100644 index 00000000..b057b03e --- /dev/null +++ b/src/obiksession/src/session.rs @@ -0,0 +1,249 @@ +use std::fs; +use std::io; +use std::path::{Path, PathBuf}; + +use memmap2::Mmap; +use obisys::DirLock; + +const PARAMS_FILE: &str = "params.bin"; +const ARTIFACTS_DIR: &str = "artifacts"; +const CHECKSUM_EXT: &str = "xxh3"; + +/// Whether [`Session::open`] found a fresh (just-created, or just-reset by +/// `force`) directory or reused one whose saved parameters matched the +/// requested ones exactly. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SessionOutcome { + /// New directory, or an existing one reset because its saved + /// parameters didn't match and `force: true` was passed. No cached + /// artifact in it can be trusted — there either isn't one yet, or the + /// ones that were there have just been deleted. + Created, + /// Existing directory whose saved parameters matched byte-for-byte. + /// Every artifact already cached under it was produced under these + /// same parameters and is safe to [`Session::restore`]. + Reused, +} + +/// Returned by [`Session::open`] when the directory already holds a +/// parameter blob that doesn't match the one just requested, and `force` +/// wasn't set. This crate has no idea what the bytes mean, so it can't +/// produce a useful diagnostic on its own — the caller, who serialized +/// both, is expected to deserialize `saved`/`requested` itself and report +/// exactly which field(s) differ before deciding whether to error out, +/// point at a different directory, or retry [`Session::open`] with +/// `force: true`. +#[derive(Debug)] +pub struct ParamsConflict { + pub saved: Vec, + pub requested: Vec, +} + +/// An open session directory: exclusively locked (via [`obisys::DirLock`]) +/// for as long as this value lives, so no other `obiksession`-using +/// process can read a half-written artifact or reset the directory out +/// from under this one. +#[derive(Debug)] +pub struct Session { + dir: PathBuf, + outcome: SessionOutcome, + _lock: DirLock, +} + +impl Session { + /// Open (or create, or reset) a session directory — see + /// [`SessionOutcome`]/[`ParamsConflict`] for the three possible + /// outcomes. Blocks until the directory's lock is free (see + /// [`obisys::DirLock::acquire`]) — a second process pointed at the + /// same session directory waits rather than racing it. + pub fn open(dir: &Path, params: &[u8], force: bool) -> io::Result> { + let lock = DirLock::acquire(dir)?; + let params_path = dir.join(PARAMS_FILE); + + let saved = read_checked(¶ms_path)?; + match saved { + Some(saved) if saved == params => Ok(Ok(Session { + dir: dir.to_path_buf(), + outcome: SessionOutcome::Reused, + _lock: lock, + })), + Some(saved) if !force => Ok(Err(ParamsConflict { saved, requested: params.to_vec() })), + _ => { + // No saved params yet, or a mismatch with `force: true`: + // start clean. Any artifact cached under different params + // is stale by construction — deleting the whole + // `artifacts/` directory is simpler and safer than trying + // to figure out which entries are still valid. + let artifacts_dir = dir.join(ARTIFACTS_DIR); + if artifacts_dir.exists() { + fs::remove_dir_all(&artifacts_dir)?; + } + fs::create_dir_all(&artifacts_dir)?; + write_checked(¶ms_path, params)?; + Ok(Ok(Session { + dir: dir.to_path_buf(), + outcome: SessionOutcome::Created, + _lock: lock, + })) + } + } + } + + pub fn dir(&self) -> &Path { + &self.dir + } + + pub fn outcome(&self) -> SessionOutcome { + self.outcome + } + + fn artifact_path(&self, name: &str) -> PathBuf { + self.dir.join(ARTIFACTS_DIR).join(name) + } + + /// Restore a previously [`store`](Self::store)d artifact, `mmap`ed + /// read-only for zero-copy access by the caller (e.g. `rkyv::access` + /// straight over the returned bytes). + /// + /// `Ok(None)` if `name` isn't cached yet, *or* if its checksum doesn't + /// match what was recorded at [`store`](Self::store) time — treated + /// the same as "not cached" rather than an error, so a partial dump + /// left behind by a killed process just triggers a recompute on the + /// next run instead of a hard failure. + pub fn restore(&self, name: &str) -> io::Result> { + verify_mmap(&self.artifact_path(name)) + } + + /// Store `bytes` under `name`, plus a checksum sidecar, atomically: + /// both are written to a temporary file and renamed into place, so a + /// crash mid-write can never leave a corrupt file that a later + /// [`restore`](Self::restore) mistakes for valid (the checksum file is + /// only ever renamed into place *after* the data file's rename + /// succeeds, so the two can't disagree about which write completed). + pub fn store(&self, name: &str, bytes: &[u8]) -> io::Result<()> { + write_checked(&self.artifact_path(name), bytes) + } +} + +fn checksum_path_for(path: &Path) -> PathBuf { + let mut os_str = path.as_os_str().to_owned(); + os_str.push("."); + os_str.push(CHECKSUM_EXT); + PathBuf::from(os_str) +} + +/// `mmap` `path` and verify it against its checksum sidecar. `Ok(None)` +/// when `path` doesn't exist, its checksum sidecar doesn't exist/parse, or +/// the checksum doesn't match — every case here means "nothing usable is +/// cached," never an error the caller has to handle specially. +fn verify_mmap(path: &Path) -> io::Result> { + if !path.exists() { + return Ok(None); + } + let Ok(expected_str) = fs::read_to_string(checksum_path_for(path)) else { + return Ok(None); + }; + let Ok(expected) = expected_str.trim().parse::() else { + return Ok(None); + }; + + let file = fs::File::open(path)?; + // SAFETY: the file is only ever written by `write_checked` (atomic + // temp-then-rename) and this whole directory is held under this + // process's exclusive `DirLock` for the `Session`'s entire lifetime — + // no other process can be concurrently modifying it underneath this + // mapping. + let mmap = unsafe { Mmap::map(&file)? }; + let actual = xxhash_rust::xxh3::xxh3_64(&mmap[..]); + if actual != expected { + tracing::warn!( + path = %path.display(), + "session artifact failed checksum verification — treating as absent" + ); + return Ok(None); + } + Ok(Some(mmap)) +} + +fn read_checked(path: &Path) -> io::Result>> { + Ok(verify_mmap(path)?.map(|mmap| mmap[..].to_vec())) +} + +fn write_checked(path: &Path, bytes: &[u8]) -> io::Result<()> { + if let Some(parent) = path.parent() { + fs::create_dir_all(parent)?; + } + let tmp = path.with_extension("tmp"); + fs::write(&tmp, bytes)?; + fs::rename(&tmp, path)?; + + let checksum = xxhash_rust::xxh3::xxh3_64(bytes); + let checksum_path = checksum_path_for(path); + let checksum_tmp = checksum_path.with_extension("tmp"); + fs::write(&checksum_tmp, checksum.to_string())?; + fs::rename(&checksum_tmp, checksum_path)?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn fresh_directory_is_created() { + let tmp = tempfile::tempdir().unwrap(); + let dir = tmp.path().join("session"); + let session = Session::open(&dir, b"params-v1", false).unwrap().unwrap(); + assert_eq!(session.outcome(), SessionOutcome::Created); + assert!(session.restore("missing").unwrap().is_none()); + } + + #[test] + fn matching_params_reuse_and_restore_artifact() { + let tmp = tempfile::tempdir().unwrap(); + let dir = tmp.path().join("session"); + { + let session = Session::open(&dir, b"params-v1", false).unwrap().unwrap(); + session.store("tally", b"some serialized bytes").unwrap(); + } + let session = Session::open(&dir, b"params-v1", false).unwrap().unwrap(); + assert_eq!(session.outcome(), SessionOutcome::Reused); + let restored = session.restore("tally").unwrap().unwrap(); + assert_eq!(&restored[..], b"some serialized bytes"); + } + + #[test] + fn mismatched_params_without_force_conflicts() { + let tmp = tempfile::tempdir().unwrap(); + let dir = tmp.path().join("session"); + Session::open(&dir, b"params-v1", false).unwrap().unwrap(); + let conflict = Session::open(&dir, b"params-v2", false).unwrap().unwrap_err(); + assert_eq!(conflict.saved, b"params-v1"); + assert_eq!(conflict.requested, b"params-v2"); + } + + #[test] + fn mismatched_params_with_force_resets_artifacts() { + let tmp = tempfile::tempdir().unwrap(); + let dir = tmp.path().join("session"); + { + let session = Session::open(&dir, b"params-v1", false).unwrap().unwrap(); + session.store("tally", b"old data").unwrap(); + } + let session = Session::open(&dir, b"params-v2", true).unwrap().unwrap(); + assert_eq!(session.outcome(), SessionOutcome::Created); + assert!(session.restore("tally").unwrap().is_none()); + } + + #[test] + fn corrupted_artifact_is_treated_as_absent() { + let tmp = tempfile::tempdir().unwrap(); + let dir = tmp.path().join("session"); + let session = Session::open(&dir, b"params-v1", false).unwrap().unwrap(); + session.store("tally", b"good data").unwrap(); + // Corrupt the data file after the fact without touching the + // checksum sidecar — simulates a partial/garbled write. + std::fs::write(dir.join("artifacts").join("tally"), b"corrupted!").unwrap(); + assert!(session.restore("tally").unwrap().is_none()); + } +} diff --git a/src/obisys/src/lock.rs b/src/obisys/src/lock.rs index b6a13aeb..dd6304e8 100644 --- a/src/obisys/src/lock.rs +++ b/src/obisys/src/lock.rs @@ -14,6 +14,7 @@ use tracing::info; /// Windows) via `std::fs::File::lock`/`try_lock`, not a hand-rolled PID /// file: the OS releases it automatically on process exit, including a /// crash — no stale-lock cleanup logic needed. +#[derive(Debug)] pub struct DirLock { _file: std::fs::File, }