Implement session persistence and checkpointing
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:
@@ -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
|
||||
|
||||
|
||||
@@ -202,18 +202,22 @@ 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 — unless `--session` is used (below), which currently covers a `snp-*` `--distance` value only, not `--sankoff`/`--tnt`/`--phyg`/`--iqtree`.
|
||||
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).
|
||||
|
||||
Every invocation, with or without `--session`, draws its own fresh random sample by default — running the same command twice gives two different (but equally valid) samples, which is useful for measuring sampling variance and is kept that way deliberately. `--session` does not change this: it makes a *specific* sample reusable on request, it does not make sampling itself reproducible from one independent run to the next.
|
||||
|
||||
### `--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 DIR` | none | Persist the sample (and, for `--sankoff`, its calibration/alignment) 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.
|
||||
A `--sankoff`-family run and a `snp-*` `--distance` run share the same cached sample when pointed at the same `--session DIR` — build it once with either, reuse it from the other, in either order, across separate commands.
|
||||
|
||||
If a run using `--session` is interrupted (crash, kill, `Ctrl-C`), the next run against the same `--session DIR` resumes from the last automatic checkpoint (roughly every 8 partitions'-worth of sampling progress) instead of starting over. The resumed run's sample is **not** guaranteed to be identical to what an uninterrupted run would have produced past that checkpoint — each process draws its own independent random sequence, same as any two separate invocations do — but nothing already checkpointed is lost, and no work needs redoing beyond that point.
|
||||
|
||||
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.
|
||||
|
||||
|
||||
Generated
+1
-1
@@ -1523,7 +1523,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "obikmer"
|
||||
version = "1.3.4"
|
||||
version = "1.3.5"
|
||||
dependencies = [
|
||||
"clap",
|
||||
"csv",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "obikmer"
|
||||
version = "1.3.4"
|
||||
version = "1.3.5"
|
||||
edition = "2024"
|
||||
|
||||
[[bin]]
|
||||
|
||||
@@ -370,6 +370,7 @@ pub fn run(args: PhyloArgs) {
|
||||
&snp_exclude_mask,
|
||||
entropy_bias,
|
||||
args.sankoff_ratio_ceiling,
|
||||
session.as_ref(),
|
||||
)
|
||||
.unwrap_or_else(|e| {
|
||||
eprintln!("error computing Sankoff calibration bundle: {e}");
|
||||
|
||||
@@ -5,6 +5,7 @@
|
||||
|
||||
use obikidxcache::index_cache::IndexCache;
|
||||
use obikindex::OKIResult;
|
||||
use rkyv::{Archive, Deserialize, Serialize};
|
||||
|
||||
use super::masking::{iupac_code, masked_state};
|
||||
use super::subsample::{EntropyBias, SurvivingFamily, sample_index};
|
||||
@@ -17,6 +18,10 @@ use super::subsample::{EntropyBias, SurvivingFamily, sample_index};
|
||||
/// index simply doesn't appear in `genome_indices`. `pub`, not
|
||||
/// `pub(crate)`: part of the public signature of
|
||||
/// [`crate::siblings::extensions::SiblingExt::snp_pseudo_alignment`].
|
||||
///
|
||||
/// `Archive`/`Serialize`/`Deserialize` (`rkyv`): `--session` persistence
|
||||
/// for `SankoffBundle` — same reasoning as `PairwiseTally`'s own derive.
|
||||
#[derive(Archive, Serialize, Deserialize)]
|
||||
pub struct SnpAlignment {
|
||||
pub sequences: Vec<Vec<u8>>,
|
||||
/// This index's own genome numbering (matching `IndexCache::meta().genomes()`'s
|
||||
@@ -54,7 +59,8 @@ pub(crate) fn snp_pseudo_alignment(
|
||||
no_ambiguity,
|
||||
excluded,
|
||||
entropy_bias,
|
||||
|_partition, _layer, survivors| {
|
||||
0,
|
||||
|_partition, _layer, _raw_index, survivors| {
|
||||
reduce_alignment(&survivors, &genome_indices, free_loss, no_ambiguity, &mut sequences);
|
||||
},
|
||||
)?;
|
||||
|
||||
@@ -22,6 +22,7 @@ mod masking;
|
||||
mod minorant_selection;
|
||||
mod pairwise;
|
||||
mod sankoff;
|
||||
mod session_cache;
|
||||
mod snp_distance;
|
||||
mod stats;
|
||||
mod subsample;
|
||||
|
||||
@@ -12,16 +12,40 @@
|
||||
|
||||
use ndarray::Array2;
|
||||
use obikidxcache::index_cache::IndexCache;
|
||||
use obikindex::OKIResult;
|
||||
use obikindex::{OKIError, OKIResult};
|
||||
use obiksession::Session;
|
||||
|
||||
use super::alignment::{SnpAlignment, reduce_alignment};
|
||||
use super::pairwise::{
|
||||
BasePairTally, CardinalityTally, PairwiseTally, PartitionDispersion, RawSnpDistanceOutput,
|
||||
reduce_pairwise,
|
||||
};
|
||||
use super::session_cache::{
|
||||
CHECKPOINT_INTERVAL, CHECKPOINT_INTERVAL_LAYERS, PROGRESS_COMPLETE, restore_tally_checkpoint,
|
||||
store_tally_checkpoint,
|
||||
};
|
||||
use super::snp_distance::{GammaShape, SnpDistanceKind, distance_matrix};
|
||||
use super::subsample::{EntropyBias, sample_index};
|
||||
|
||||
/// The alignment's own `--session` artifact — kept in lockstep with the
|
||||
/// shared `pairwise_tally`/`partition_dispersion`/`progress` artifacts
|
||||
/// (`session_cache`): both are folded from the *same* `on_layer` call, so
|
||||
/// checkpointing them together, under the same `progress` marker, is what
|
||||
/// keeps a resumed alignment's column count consistent with how many
|
||||
/// layers the resumed tally actually reflects.
|
||||
const ALIGNMENT_ARTIFACT: &str = "snp_alignment";
|
||||
|
||||
fn restore_alignment(session: &Session) -> Option<SnpAlignment> {
|
||||
let bytes = session.restore(ALIGNMENT_ARTIFACT).ok().flatten()?;
|
||||
rkyv::from_bytes::<SnpAlignment, rkyv::rancor::Error>(&bytes[..]).ok()
|
||||
}
|
||||
|
||||
fn store_alignment(session: &Session, alignment: &SnpAlignment) -> OKIResult<()> {
|
||||
let bytes = rkyv::to_bytes::<rkyv::rancor::Error>(alignment)
|
||||
.map_err(|e| OKIError::Io(std::io::Error::other(e.to_string())))?;
|
||||
session.store(ALIGNMENT_ARTIFACT, &bytes).map_err(OKIError::Io)
|
||||
}
|
||||
|
||||
/// Every output the `--sankoff`/`--tnt`/`--phyg`/`--iqtree` pipeline needs,
|
||||
/// computed together from one shared, possibly-subsampled/entropy-biased
|
||||
/// selection — see the module docs. `pub`: part of the public signature of
|
||||
@@ -67,6 +91,14 @@ impl SankoffBundle {
|
||||
/// [`BasePairTally`] only (see [`PairwiseTally::cardinality_tally`]'s own
|
||||
/// docs for why [`CardinalityTally`] uses a different — `excluded`-only —
|
||||
/// inclusion rule).
|
||||
///
|
||||
/// `session`: `--session`, `None` means no persistence (always resample).
|
||||
/// `Some(session)` must already be open under the exact selection
|
||||
/// parameters this call is about to use — the caller's job, same contract
|
||||
/// as [`super::snp_distance::snp_distance`]'s own `session` parameter,
|
||||
/// which this shares an artifact namespace with (see `session_cache`'s
|
||||
/// module docs): a tally checkpointed here is restorable by a later
|
||||
/// `--distance snp-*` call against the same session, and vice versa.
|
||||
pub(crate) fn sankoff_bundle(
|
||||
cache: &IndexCache,
|
||||
n: usize,
|
||||
@@ -75,27 +107,72 @@ pub(crate) fn sankoff_bundle(
|
||||
excluded: &[bool],
|
||||
entropy_bias: Option<EntropyBias>,
|
||||
ratio_ceiling: f64,
|
||||
session: Option<&Session>,
|
||||
) -> OKIResult<SankoffBundle> {
|
||||
let n_genomes = cache.meta().genomes().len();
|
||||
let genome_indices: Vec<usize> = (0..n_genomes)
|
||||
.filter(|&g| !excluded.get(g).copied().unwrap_or(false))
|
||||
.collect();
|
||||
let mut sequences: Vec<Vec<u8>> = vec![Vec::new(); genome_indices.len()];
|
||||
let mut tally = PairwiseTally::new(n_genomes);
|
||||
let mut dispersion = PartitionDispersion::default();
|
||||
|
||||
sample_index(
|
||||
cache,
|
||||
n,
|
||||
free_loss,
|
||||
no_ambiguity,
|
||||
excluded,
|
||||
entropy_bias,
|
||||
|partition, _layer, survivors| {
|
||||
reduce_alignment(&survivors, &genome_indices, free_loss, no_ambiguity, &mut sequences);
|
||||
reduce_pairwise(&survivors, partition, &mut tally, &mut dispersion);
|
||||
},
|
||||
)?;
|
||||
let restored = session.and_then(|s| Some((restore_tally_checkpoint(s)?, restore_alignment(s)?)));
|
||||
let (mut tally, mut dispersion, mut sequences, skip_layers) = match restored {
|
||||
Some(((tally, dispersion, skip_layers), alignment)) => {
|
||||
(tally, dispersion, alignment.sequences, skip_layers)
|
||||
}
|
||||
None => {
|
||||
let sequences = vec![Vec::new(); genome_indices.len()];
|
||||
(PairwiseTally::new(n_genomes), PartitionDispersion::default(), sequences, 0)
|
||||
}
|
||||
};
|
||||
|
||||
if skip_layers != usize::MAX {
|
||||
// Same checkpoint cadence and rationale as `snp_distance`'s own
|
||||
// resumable loop (see `session_cache::CHECKPOINT_INTERVAL*`'s own
|
||||
// docs) — the alignment's `sequences` are folded in lockstep with
|
||||
// the tally so a resumed run's column count always matches how
|
||||
// many layers the restored tally actually reflects.
|
||||
let mut last_checkpoint_layer = skip_layers;
|
||||
let mut last_checkpoint_at = std::time::Instant::now();
|
||||
sample_index(
|
||||
cache,
|
||||
n,
|
||||
free_loss,
|
||||
no_ambiguity,
|
||||
excluded,
|
||||
entropy_bias,
|
||||
skip_layers,
|
||||
|partition, _layer, raw_index, survivors| {
|
||||
reduce_alignment(&survivors, &genome_indices, free_loss, no_ambiguity, &mut sequences);
|
||||
reduce_pairwise(&survivors, partition, &mut tally, &mut dispersion);
|
||||
let Some(session) = session else { return };
|
||||
let layers_done = raw_index + 1;
|
||||
let due_by_count =
|
||||
layers_done.saturating_sub(last_checkpoint_layer) >= CHECKPOINT_INTERVAL_LAYERS;
|
||||
let due_by_time = last_checkpoint_at.elapsed() >= CHECKPOINT_INTERVAL;
|
||||
if !(due_by_count || due_by_time) {
|
||||
return;
|
||||
}
|
||||
let alignment = SnpAlignment { sequences: sequences.clone(), genome_indices: genome_indices.clone() };
|
||||
let stored = store_tally_checkpoint(session, &tally, &dispersion, layers_done as u64)
|
||||
.and_then(|()| store_alignment(session, &alignment));
|
||||
match stored {
|
||||
Ok(()) => {
|
||||
last_checkpoint_layer = layers_done;
|
||||
last_checkpoint_at = std::time::Instant::now();
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("--session: failed to checkpoint at layer {layers_done}: {e}");
|
||||
}
|
||||
}
|
||||
},
|
||||
)?;
|
||||
|
||||
if let Some(session) = session {
|
||||
let alignment = SnpAlignment { sequences: sequences.clone(), genome_indices: genome_indices.clone() };
|
||||
store_tally_checkpoint(session, &tally, &dispersion, PROGRESS_COMPLETE)?;
|
||||
store_alignment(session, &alignment)?;
|
||||
}
|
||||
}
|
||||
|
||||
let raw = tally.raw_snp_distance();
|
||||
let included = tally.included(ratio_ceiling, excluded);
|
||||
|
||||
@@ -0,0 +1,101 @@
|
||||
//! Shared `--session` checkpoint plumbing for [`super::snp_distance`] and
|
||||
//! [`super::sankoff`] — both fold [`super::subsample::sample_index`]'s
|
||||
//! layer-by-layer output into a [`PairwiseTally`]/[`PartitionDispersion`]
|
||||
//! pair, and both persist it under the *same* artifact names on purpose:
|
||||
//! a tally checkpointed by one is restorable by the other. Concretely,
|
||||
//! `--sankoff --session DIR` today followed by `--distance snp-k2p
|
||||
//! --session DIR` tomorrow restores the first run's tally instead of
|
||||
//! resampling, even though only `--sankoff` originally built it —
|
||||
//! `snp_distance`'s own selection-parameter conflict check (via
|
||||
//! `obiksession::Session::open`, the caller's responsibility) is what
|
||||
//! guarantees that's actually safe to do.
|
||||
|
||||
use obikindex::{OKIError, OKIResult};
|
||||
use obiksession::Session;
|
||||
|
||||
use super::pairwise::{PairwiseTally, PartitionDispersion};
|
||||
|
||||
pub(crate) const TALLY_ARTIFACT: &str = "pairwise_tally";
|
||||
pub(crate) const DISPERSION_ARTIFACT: &str = "partition_dispersion";
|
||||
pub(crate) const PROGRESS_ARTIFACT: &str = "progress";
|
||||
/// Sentinel `progress` value meaning "every layer was processed — this
|
||||
/// tally is complete," distinct from any real layer count so a resumed run
|
||||
/// can tell "finished" apart from "stopped after N layers" without relying
|
||||
/// on N happening to equal the index's actual layer count.
|
||||
pub(crate) const PROGRESS_COMPLETE: u64 = u64::MAX;
|
||||
|
||||
/// How much in-progress sampling work a checkpoint can lose: at most
|
||||
/// `CHECKPOINT_INTERVAL_LAYERS` layers *or* `CHECKPOINT_INTERVAL` of wall
|
||||
/// time, whichever comes first — not "checkpoint every layer" (rewriting
|
||||
/// the whole, potentially large, tally/dispersion after every single one of
|
||||
/// what can be thousands of layers would make the checkpointing overhead
|
||||
/// itself dominate), and not "only at the end" either (which would defeat
|
||||
/// the point: on a long run against a real production index, losing 90% of
|
||||
/// an already-completed scan to a crash is exactly the expensive case
|
||||
/// resumability exists to avoid — the actual justification, not "the tally
|
||||
/// is cheap to rebuild," which was an earlier, wrong framing of this
|
||||
/// tradeoff caught on review).
|
||||
pub(crate) const CHECKPOINT_INTERVAL_LAYERS: usize = 8;
|
||||
pub(crate) const CHECKPOINT_INTERVAL: std::time::Duration = std::time::Duration::from_secs(30);
|
||||
|
||||
/// Restore a checkpointed `(PairwiseTally, PartitionDispersion, resume-from
|
||||
/// layer index)` triple from `session` — `None` on any miss (an artifact
|
||||
/// not cached, an I/O error, or a deserialization failure), all treated
|
||||
/// identically: something's missing or unusable, fall back to starting
|
||||
/// from scratch (`skip_layers = 0`) rather than surfacing a hard error over
|
||||
/// what's meant to be a transparent cache. The resume index is
|
||||
/// [`PROGRESS_COMPLETE`] when a prior run finished every layer. Layers
|
||||
/// processed after a resume draw a fresh, independent random sample (see
|
||||
/// `subsample::sample_index`'s own `skip_layers` docs) — deliberately not
|
||||
/// bit-identical to what an uninterrupted run would have produced,
|
||||
/// consistent with wanting run-to-run sampling variance rather than fixed
|
||||
/// reproducibility.
|
||||
///
|
||||
/// Full owned deserialization (`rkyv::from_bytes`) for the tally/dispersion,
|
||||
/// not zero-copy (`rkyv::access`) — simpler to integrate with
|
||||
/// `PairwiseTally`'s existing query methods, which all take `&self` by
|
||||
/// value semantics, not an `Archived<PairwiseTally>`. Still skips the
|
||||
/// expensive part (`sample_index` re-scanning the sibling annex for every
|
||||
/// already-completed layer); the 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<PairwiseTally>`
|
||||
/// too — a bigger, separate change.
|
||||
pub(crate) fn restore_tally_checkpoint(session: &Session) -> Option<(PairwiseTally, PartitionDispersion, usize)> {
|
||||
let tally_bytes = session.restore(TALLY_ARTIFACT).ok().flatten()?;
|
||||
let dispersion_bytes = session.restore(DISPERSION_ARTIFACT).ok().flatten()?;
|
||||
let progress_bytes = session.restore(PROGRESS_ARTIFACT).ok().flatten()?;
|
||||
let tally = rkyv::from_bytes::<PairwiseTally, rkyv::rancor::Error>(&tally_bytes[..]).ok()?;
|
||||
let dispersion =
|
||||
rkyv::from_bytes::<PartitionDispersion, rkyv::rancor::Error>(&dispersion_bytes[..]).ok()?;
|
||||
let progress_raw: [u8; 8] = progress_bytes.get(..8)?.try_into().ok()?;
|
||||
let progress = u64::from_le_bytes(progress_raw);
|
||||
let skip_layers = if progress == PROGRESS_COMPLETE { usize::MAX } else { progress as usize };
|
||||
if skip_layers == usize::MAX {
|
||||
tracing::info!("--session: restored a complete cached SNP tally — skipping resampling");
|
||||
} else {
|
||||
tracing::info!("--session: resuming a cached SNP tally from layer {skip_layers} onward");
|
||||
}
|
||||
Some((tally, dispersion, skip_layers))
|
||||
}
|
||||
|
||||
/// Store `tally`/`dispersion` plus how far sampling has gotten
|
||||
/// (`layers_done`, or [`PROGRESS_COMPLETE`] once every layer has been
|
||||
/// processed) — a checkpoint [`restore_tally_checkpoint`] can resume from
|
||||
/// later, whether that "later" is after a crash or a clean, separate
|
||||
/// invocation.
|
||||
pub(crate) fn store_tally_checkpoint(
|
||||
session: &Session,
|
||||
tally: &PairwiseTally,
|
||||
dispersion: &PartitionDispersion,
|
||||
progress: u64,
|
||||
) -> OKIResult<()> {
|
||||
let tally_bytes = rkyv::to_bytes::<rkyv::rancor::Error>(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::<rkyv::rancor::Error>(dispersion)
|
||||
.map_err(|e| OKIError::Io(std::io::Error::other(e.to_string())))?;
|
||||
session.store(DISPERSION_ARTIFACT, &dispersion_bytes).map_err(OKIError::Io)?;
|
||||
session.store(PROGRESS_ARTIFACT, &progress.to_le_bytes()).map_err(OKIError::Io)?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -23,51 +23,17 @@
|
||||
|
||||
use ndarray::Array2;
|
||||
use obikidxcache::index_cache::IndexCache;
|
||||
use obikindex::{OKIError, OKIResult};
|
||||
use obikindex::OKIResult;
|
||||
use obiksession::Session;
|
||||
|
||||
use super::pairwise::{PairwiseTally, PartitionDispersion};
|
||||
use super::session_cache::{
|
||||
CHECKPOINT_INTERVAL, CHECKPOINT_INTERVAL_LAYERS, PROGRESS_COMPLETE, restore_tally_checkpoint,
|
||||
store_tally_checkpoint,
|
||||
};
|
||||
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<PairwiseTally>`. 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<PairwiseTally>` 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::<PairwiseTally, rkyv::rancor::Error>(&tally_bytes[..]).ok()?;
|
||||
let dispersion =
|
||||
rkyv::from_bytes::<PartitionDispersion, rkyv::rancor::Error>(&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::<rkyv::rancor::Error>(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::<rkyv::rancor::Error>(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.
|
||||
@@ -302,44 +268,68 @@ pub(crate) fn snp_distance(
|
||||
));
|
||||
}
|
||||
|
||||
let cached = session.and_then(restore_tally);
|
||||
let (tally, dispersion) = match cached {
|
||||
Some(pair) => pair,
|
||||
let (mut tally, mut dispersion, skip_layers) = match session.and_then(restore_tally_checkpoint) {
|
||||
Some(triple) => triple,
|
||||
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)
|
||||
(PairwiseTally::new(n_genomes), PartitionDispersion::default(), 0)
|
||||
}
|
||||
};
|
||||
|
||||
if skip_layers != usize::MAX {
|
||||
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 {
|
||||
// Checkpoint every `CHECKPOINT_INTERVAL_LAYERS` layers or
|
||||
// `CHECKPOINT_INTERVAL` of wall time, whichever comes first —
|
||||
// see those constants' own docs for why "just recompute
|
||||
// everything" isn't the right call to make on a long run
|
||||
// against a real production index.
|
||||
let mut last_checkpoint_layer = skip_layers;
|
||||
let mut last_checkpoint_at = std::time::Instant::now();
|
||||
sample_index(
|
||||
cache,
|
||||
target,
|
||||
free_loss,
|
||||
no_ambiguity,
|
||||
excluded,
|
||||
entropy_bias,
|
||||
skip_layers,
|
||||
|partition, _layer, raw_index, survivors| {
|
||||
super::pairwise::reduce_pairwise(&survivors, partition, &mut tally, &mut dispersion);
|
||||
let Some(session) = session else { return };
|
||||
let layers_done = raw_index + 1;
|
||||
let due_by_count =
|
||||
layers_done.saturating_sub(last_checkpoint_layer) >= CHECKPOINT_INTERVAL_LAYERS;
|
||||
let due_by_time = last_checkpoint_at.elapsed() >= CHECKPOINT_INTERVAL;
|
||||
if !(due_by_count || due_by_time) {
|
||||
return;
|
||||
}
|
||||
match store_tally_checkpoint(session, &tally, &dispersion, layers_done as u64) {
|
||||
Ok(()) => {
|
||||
last_checkpoint_layer = layers_done;
|
||||
last_checkpoint_at = std::time::Instant::now();
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("--session: failed to checkpoint at layer {layers_done}: {e}");
|
||||
}
|
||||
}
|
||||
},
|
||||
)?;
|
||||
}
|
||||
|
||||
if let Some(session) = session {
|
||||
store_tally_checkpoint(session, &tally, &dispersion, PROGRESS_COMPLETE)?;
|
||||
}
|
||||
}
|
||||
|
||||
distance_matrix(&tally, &dispersion, kind, gamma_shape)
|
||||
}
|
||||
|
||||
|
||||
@@ -124,8 +124,33 @@ fn circular_entropy_values(
|
||||
/// its own `rayon::join`/`scope` worker, since `Arc` (unlike a plain
|
||||
/// borrowed slice) satisfies the `Send + 'static` those need.
|
||||
///
|
||||
/// Returns the actual number of sites kept, which may be less than `n` (see
|
||||
/// the module docs).
|
||||
/// `skip_layers`: the number of (partition, layer) pairs, in this
|
||||
/// function's own iteration order (`cache.partitions()` then `0..n_layer`,
|
||||
/// both deterministic for an unchanged index), to skip *entirely* —
|
||||
/// neither their (cheap) eligibility-bitset lookup's `sample_layer` call
|
||||
/// nor `on_layer` runs for them. Exists for `--session` resume: quotas are
|
||||
/// computed from `total_eligible` (summed over *every* layer regardless of
|
||||
/// `skip_layers`), so a resumed run reproduces the exact same *per-layer
|
||||
/// quotas* an uninterrupted one would have — i.e. how many sites each
|
||||
/// remaining layer contributes stays correct. It does **not** reproduce
|
||||
/// the exact same *sample*: `sample_layer`'s draws come from `rand::rng()`
|
||||
/// (`Self::sample_layer`'s own docs), a per-process, OS-seeded generator
|
||||
/// with no continuity across a process boundary, so the layers processed
|
||||
/// after a resume draw different (but equally valid) random outcomes than
|
||||
/// an uninterrupted run would have. Confirmed by direct test (kill mid-run,
|
||||
/// resume, compare to an uninterrupted run against the same index/params:
|
||||
/// different distance matrices, both legitimate samples). This is
|
||||
/// deliberate, not a defect: run-to-run sampling variance is wanted (see
|
||||
/// `DevDocMD/theory/evolutionary_distances.md`, "`--session`" section) —
|
||||
/// making resume bit-reproducible would require seeding deterministically,
|
||||
/// which would also remove that variance for ordinary repeated runs
|
||||
/// unless scoped carefully, and was explicitly declined for now. `0` for
|
||||
/// every caller not doing checkpointed resume.
|
||||
///
|
||||
/// Returns the actual number of *newly processed* sites kept (not
|
||||
/// counting whatever `skip_layers` layers already contributed on an
|
||||
/// earlier, interrupted call) — may be less than `n` (see the module
|
||||
/// docs).
|
||||
pub(crate) fn sample_index(
|
||||
cache: &IndexCache,
|
||||
n: usize,
|
||||
@@ -133,7 +158,8 @@ pub(crate) fn sample_index(
|
||||
no_ambiguity: bool,
|
||||
excluded: &[bool],
|
||||
entropy_bias: Option<EntropyBias>,
|
||||
mut on_layer: impl FnMut(usize, usize, Arc<Vec<SurvivingFamily>>),
|
||||
skip_layers: usize,
|
||||
mut on_layer: impl FnMut(usize, usize, usize, Arc<Vec<SurvivingFamily>>),
|
||||
) -> OKIResult<usize> {
|
||||
let n_genomes = cache.meta().genomes().len();
|
||||
let fast_mode = is_fast_mode(cache);
|
||||
@@ -171,7 +197,10 @@ pub(crate) fn sample_index(
|
||||
}
|
||||
|
||||
let mut total_kept = 0usize;
|
||||
for (part, l, mut eligible, n_minorants) in layers {
|
||||
for (raw_index, (part, l, mut eligible, n_minorants)) in layers.into_iter().enumerate() {
|
||||
if raw_index < skip_layers {
|
||||
continue;
|
||||
}
|
||||
let count_layer = eligible.view().count_ones();
|
||||
if count_layer == 0 {
|
||||
continue;
|
||||
@@ -184,6 +213,7 @@ pub(crate) fn sample_index(
|
||||
cache,
|
||||
part,
|
||||
l,
|
||||
raw_index,
|
||||
&mut eligible,
|
||||
n_minorants,
|
||||
quota,
|
||||
@@ -207,6 +237,7 @@ fn sample_layer(
|
||||
cache: &IndexCache,
|
||||
partition: usize,
|
||||
layer_idx: usize,
|
||||
raw_index: usize,
|
||||
eligible: &mut TempBitVecBuilder,
|
||||
n_minorants: usize,
|
||||
quota: usize,
|
||||
@@ -216,7 +247,7 @@ fn sample_layer(
|
||||
no_ambiguity: bool,
|
||||
excluded: &[bool],
|
||||
entropy_bias: Option<EntropyBias>,
|
||||
on_layer: &mut impl FnMut(usize, usize, Arc<Vec<SurvivingFamily>>),
|
||||
on_layer: &mut impl FnMut(usize, usize, usize, Arc<Vec<SurvivingFamily>>),
|
||||
) -> OKIResult<usize> {
|
||||
if n_minorants == 0 {
|
||||
return Ok(0);
|
||||
@@ -330,7 +361,7 @@ fn sample_layer(
|
||||
}
|
||||
|
||||
if !survivors.is_empty() {
|
||||
on_layer(partition, layer_idx, Arc::new(survivors));
|
||||
on_layer(partition, layer_idx, raw_index, Arc::new(survivors));
|
||||
}
|
||||
|
||||
Ok(kept)
|
||||
|
||||
@@ -113,6 +113,14 @@ pub trait SiblingExt {
|
||||
/// `crate::siblings::CardinalityTally`'s own docs for why that
|
||||
/// wouldn't make sense: cardinality reflects each genome's own
|
||||
/// coverage/duplication structure, not the pair's mutual divergence).
|
||||
/// `session`: `--session`, `None` disables persistence (always
|
||||
/// resample). `Some` must already be open under these exact selection
|
||||
/// parameters (the caller's responsibility) — restores a cached tally/
|
||||
/// alignment when present instead of resampling, and checkpoints a
|
||||
/// fresh one back as sampling progresses otherwise. Shares its
|
||||
/// artifact namespace with [`Self::snp_distance`]'s own `session`
|
||||
/// parameter internally: a tally checkpointed by one is restorable by
|
||||
/// the other, given the same session and selection parameters.
|
||||
fn sankoff_bundle(
|
||||
&self,
|
||||
n: usize,
|
||||
@@ -121,6 +129,7 @@ pub trait SiblingExt {
|
||||
excluded: &[bool],
|
||||
entropy_bias: Option<EntropyBias>,
|
||||
ratio_ceiling: f64,
|
||||
session: Option<&obiksession::Session>,
|
||||
) -> OKIResult<SankoffBundle>;
|
||||
|
||||
/// A `snp-*` `--distance` matrix — one of the closed-form corrections
|
||||
@@ -330,8 +339,9 @@ impl SiblingExt for IndexCache {
|
||||
excluded: &[bool],
|
||||
entropy_bias: Option<EntropyBias>,
|
||||
ratio_ceiling: f64,
|
||||
session: Option<&obiksession::Session>,
|
||||
) -> OKIResult<SankoffBundle> {
|
||||
sankoff_bundle(self, n, free_loss, no_ambiguity, excluded, entropy_bias, ratio_ceiling)
|
||||
sankoff_bundle(self, n, free_loss, no_ambiguity, excluded, entropy_bias, ratio_ceiling, session)
|
||||
}
|
||||
|
||||
fn snp_distance(
|
||||
|
||||
Reference in New Issue
Block a user