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:
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