gam_solve/

cross_node.rs

//! Cross-node deterministic border-Gram reduction (#987, extending #973).
//!
//! [`crate::streaming_border`] made the in-process accumulation of the
//! Schur border Gram `G = Σ_n x_n x_nᵀ` bit-reproducible by construction: the
//! chunk partition is a pure function of `(n_rows, chunk_size)`, per-chunk
//! partials are deterministic [`chunk_gram_flat`] reductions, and the
//! cross-chunk fold is a fixed pairwise tree keyed by **chunk index** — never by
//! arrival order, thread timing, or device count. This module extends that same
//! fixed-shape-by-construction discipline **one level up**, to a fleet of
//! worker nodes, with three properties the frontier corpus regime
//! (10⁹–10¹¹ tokens, hundreds of TB of activations) demands:
//!
//! 1. **Node count never changes bits.** A node's partials are not leaves of a
//!    *separate* per-node tree that then gets merged (that shape would depend
//!    on the node count). Instead every node computes the **globally indexed**
//!    per-chunk partials it owns and ships `(chunk_index, k·k partial)`
//!    messages; the coordinator folds them through the *single* global cascade
//!    of [`StreamingBorderGram`], which accepts any arrival order and folds in
//!    chunk-index order. The reduction topology is therefore a pure function of
//!    `(n_rows, chunk_size)` alone — running on 1 node, 3 nodes, or 64 nodes
//!    yields the identical bit pattern, because the tree never saw the node
//!    count. (The chunk→node *assignment* is rank-indexed and deterministic,
//!    but it only decides who computes a partial, never how partials combine.)
//! 2. **Checkpoint/resume is the job model, not an afterthought.** Any worker's
//!    death resumes from its serialized [`NodeWorkerCheckpoint`] (a cursor into
//!    its owned chunk sequence); the coordinator's full state — the in-order
//!    fold forest, the pending out-of-order partials, and the per-rank receipt
//!    cursors — serializes to a [`CrossNodeCheckpoint`]. Resume-equals-
//!    straight-through holds at the bit level on both sides because both
//!    cursors are positions in deterministic sequences.
//! 3. **Partials, never rows, cross the wire.** A worker streams its shard rows
//!    locally (object store / mmap — `gam_sae::corpus`) and ships
//!    only `k·k` f64 partials. The coordinator's ingest seam is
//!    `StreamingBorderGram::submit_chunk_gram`; both producers route through
//!    the one [`chunk_gram_flat`] free function, so a shipped partial is
//!    bit-identical to the partial the coordinator would have computed from the
//!    same rows.
//!
//! ## Chunk→rank assignment
//!
//! Round-robin by chunk index: rank `r` of `n_ranks` owns chunks
//! `{j : j ≡ r (mod n_ranks)}`, in increasing order. Round-robin (rather than
//! contiguous ranges) keeps the coordinator's in-order fold frontier advancing
//! steadily while all ranks make progress at similar rates, which bounds the
//! pending out-of-order buffer by O(`n_ranks` × inter-node skew) instead of
//! O(total chunks). The assignment is a pure function of
//! `(chunk_index, n_ranks)`; no scheduler, no work stealing — work stealing
//! would not change bits (the fold is index-keyed) but it *would* break the
//! one-cursor-per-rank resume model, so it is deliberately absent.
//!
//! Pure library: no networking, no flags, no environment variables. The
//! transport (MPI, gRPC, files on a shared filesystem) is the caller's; this
//! module owns the deterministic topology, the cursors, and the validation.

use crate::streaming_border::{BorderGramCheckpoint, StreamingBorderGram, chunk_gram_flat};
use ndarray::{Array2, ArrayView2};
use serde::{Deserialize, Serialize};

/// The deterministic chunk partition + rank-indexed assignment shared by every
/// participant of one cross-node pass. A pure function of its four fields; two
/// participants constructed with the same fields agree on every derived
/// quantity, with no communication.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct CrossNodePartition {
    /// Border dimension `k` (columns of every chunk; partials are `k·k`).
    pub border_dim: usize,
    /// Total row count of the full pass.
    pub n_rows: usize,
    /// Fixed chunk size (rows per chunk; the last chunk may be shorter).
    pub chunk_size: usize,
    /// Number of worker ranks in the fleet.
    pub n_ranks: usize,
}

impl CrossNodePartition {
    pub fn new(
        border_dim: usize,
        n_rows: usize,
        chunk_size: usize,
        n_ranks: usize,
    ) -> Result<Self, String> {
        if border_dim == 0 {
            return Err("CrossNodePartition: border_dim must be positive".to_string());
        }
        if chunk_size == 0 {
            return Err("CrossNodePartition: chunk_size must be positive".to_string());
        }
        if n_ranks == 0 {
            return Err("CrossNodePartition: n_ranks must be positive".to_string());
        }
        Ok(Self {
            border_dim,
            n_rows,
            chunk_size,
            n_ranks,
        })
    }

    /// Total number of chunks of the pass: `ceil(n_rows / chunk_size)`.
    /// Identical to [`StreamingBorderGram::n_chunks`] for the same partition
    /// parameters — the global tree this assignment feeds.
    pub fn n_chunks(&self) -> usize {
        self.n_rows.div_ceil(self.chunk_size)
    }

    /// Row range covered by global chunk `chunk_index` — the same pure function
    /// as [`StreamingBorderGram::chunk_rows`], duplicated here so a worker can
    /// slice its rows without constructing a coordinator-side accumulator.
    pub fn chunk_rows(&self, chunk_index: usize) -> std::ops::Range<usize> {
        let lo = chunk_index * self.chunk_size;
        let hi = ((chunk_index + 1) * self.chunk_size).min(self.n_rows);
        lo..hi
    }

    /// Number of chunks rank `rank` owns.
    pub fn chunks_owned_by(&self, rank: usize) -> usize {
        let n = self.n_chunks();
        if rank >= self.n_ranks || n == 0 {
            return 0;
        }
        // Chunks r, r + n_ranks, r + 2·n_ranks, … below n.
        if rank < n {
            (n - rank - 1) / self.n_ranks + 1
        } else {
            0
        }
    }

    /// The `ordinal`-th (0-based) global chunk index owned by `rank`, or `None`
    /// past the end of the rank's sequence. The worker cursor is an ordinal
    /// into exactly this sequence.
    pub fn owned_chunk(&self, rank: usize, ordinal: usize) -> Option<usize> {
        if rank >= self.n_ranks {
            return None;
        }
        let idx = rank + ordinal * self.n_ranks;
        if idx < self.n_chunks() {
            Some(idx)
        } else {
            None
        }
    }
}

/// One shipped partial: the global chunk index plus the deterministic `k·k`
/// per-chunk Gram. This is the only message that crosses the node boundary —
/// `k·k` f64 values per chunk, never rows.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct NodePartial {
    /// Rank that produced this partial (its assignment is validated on
    /// receipt, so a misconfigured worker is rejected loudly).
    pub rank: usize,
    /// Global chunk index of the partial.
    pub chunk_index: usize,
    /// Flattened `k·k` row-major per-chunk Gram, as produced by
    /// [`chunk_gram_flat`] over the chunk's rows.
    pub gram: Vec<f64>,
}

/// Serialized cursor of one worker: everything needed for a **replacement**
/// process (same rank, any host) to continue the dead worker's deterministic
/// chunk sequence from where receipts stopped. Pure data; the worker's row
/// source re-seeks by row range, which is a pure function of the partition.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct NodeWorkerCheckpoint {
    pub partition: CrossNodePartition,
    pub rank: usize,
    /// Ordinal (into the rank's owned-chunk sequence) of the next chunk to
    /// compute and ship. Everything below it has been durably shipped.
    pub next_ordinal: usize,
}

/// Worker-side driver for one rank: walks the rank's deterministic owned-chunk
/// sequence, turning row slices into shippable [`NodePartial`]s.
///
/// The worker does **not** do I/O: the caller streams rows (from its shards /
/// object store) for the row range [`NodeWorker::next_chunk_rows`] names, hands
/// them to [`NodeWorker::emit`], and ships the returned partial. The cursor
/// advances only on `emit`, so "ship durably, then checkpoint" gives exactly-
/// once production under crash-resume (re-shipping an already-folded chunk is
/// rejected by the coordinator as a duplicate, which is the safe failure).
#[derive(Clone, Debug)]
pub struct NodeWorker {
    partition: CrossNodePartition,
    rank: usize,
    next_ordinal: usize,
}

impl NodeWorker {
    /// Fresh worker for `rank`, starting at the beginning of its sequence.
    pub fn new(partition: CrossNodePartition, rank: usize) -> Result<Self, String> {
        if rank >= partition.n_ranks {
            return Err(format!(
                "NodeWorker: rank {rank} out of range (n_ranks = {})",
                partition.n_ranks
            ));
        }
        Ok(Self {
            partition,
            rank,
            next_ordinal: 0,
        })
    }

    /// Resume a (replacement) worker from a serialized cursor. Validates the
    /// cursor against the partition so a checkpoint from a different pass is
    /// rejected loudly.
    pub fn resume(state: NodeWorkerCheckpoint) -> Result<Self, String> {
        if state.rank >= state.partition.n_ranks {
            return Err(format!(
                "NodeWorkerCheckpoint: rank {} out of range (n_ranks = {})",
                state.rank, state.partition.n_ranks
            ));
        }
        let owned = state.partition.chunks_owned_by(state.rank);
        if state.next_ordinal > owned {
            return Err(format!(
                "NodeWorkerCheckpoint: next_ordinal {} exceeds owned chunk count {owned}",
                state.next_ordinal
            ));
        }
        Ok(Self {
            partition: state.partition,
            rank: state.rank,
            next_ordinal: state.next_ordinal,
        })
    }

    /// Serialize the cursor. Write this (durably) after each successful ship.
    pub fn checkpoint(&self) -> NodeWorkerCheckpoint {
        NodeWorkerCheckpoint {
            partition: self.partition,
            rank: self.rank,
            next_ordinal: self.next_ordinal,
        }
    }

    /// Global chunk index and row range of the next chunk to compute, or
    /// `None` when done. The caller fetches exactly these rows.
    pub fn next_chunk_rows(&self) -> Option<(usize, std::ops::Range<usize>)> {
        let idx = self.partition.owned_chunk(self.rank, self.next_ordinal)?;
        Some((idx, self.partition.chunk_rows(idx)))
    }

    /// Compute the next chunk's deterministic partial from its rows and advance
    /// the cursor. `rows` must be exactly the rows of
    /// [`NodeWorker::next_chunk_rows`] (shape-validated here; content is the
    /// caller's contract, same as the in-process path).
    pub fn emit(&mut self, rows: ArrayView2<'_, f64>) -> Result<NodePartial, String> {
        let (chunk_index, range) = self
            .next_chunk_rows()
            .ok_or_else(|| format!("NodeWorker rank {}: sequence exhausted", self.rank))?;
        if rows.nrows() != range.len() || rows.ncols() != self.partition.border_dim {
            return Err(format!(
                "NodeWorker rank {}: chunk {chunk_index} has shape ({}, {}) but expected ({}, {})",
                self.rank,
                rows.nrows(),
                rows.ncols(),
                range.len(),
                self.partition.border_dim
            ));
        }
        let gram = chunk_gram_flat(rows);
        self.next_ordinal += 1;
        Ok(NodePartial {
            rank: self.rank,
            chunk_index,
            gram,
        })
    }
}

/// Serializable coordinator state: the inner accumulation state plus the
/// per-rank receipt cursors.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct CrossNodeCheckpoint {
    pub partition: CrossNodePartition,
    /// The wrapped [`StreamingBorderGram`] state (fold forest + pending +
    /// chunk frontier).
    pub inner: BorderGramCheckpoint,
    /// Per-rank count of partials received so far — each is an ordinal cursor
    /// into that rank's deterministic owned-chunk sequence, used both for
    /// receipt validation (in-sequence, no gaps per rank) and to tell a
    /// restarted fleet where each rank should resume.
    pub received_per_rank: Vec<usize>,
}

/// Coordinator-side reduction: receives [`NodePartial`]s from the fleet and
/// folds them into the single global fixed-tree accumulator.
///
/// Receipt validation is per rank and in-sequence: rank `r`'s `i`-th accepted
/// partial must be its `i`-th owned chunk. This makes the per-rank cursor in
/// [`CrossNodeCheckpoint::received_per_rank`] a complete description of what
/// has been received, which is what lets a dead rank resume from a bare
/// ordinal. Cross-rank arrival order is unconstrained (the inner accumulator
/// buffers out-of-order chunks), so slow nodes never block fast ones.
pub struct CrossNodeGramReduction {
    partition: CrossNodePartition,
    inner: StreamingBorderGram,
    received_per_rank: Vec<usize>,
}

impl CrossNodeGramReduction {
    /// Fresh coordinator for the given partition.
    pub fn new(partition: CrossNodePartition) -> Result<Self, String> {
        let inner =
            StreamingBorderGram::new(partition.border_dim, partition.n_rows, partition.chunk_size)?;
        Ok(Self {
            received_per_rank: vec![0; partition.n_ranks],
            partition,
            inner,
        })
    }

    /// The shared partition (workers must be constructed with an equal one).
    pub fn partition(&self) -> CrossNodePartition {
        self.partition
    }

    /// Serialize the full coordinator state. Resume-equals-straight-through is
    /// inherited bit-for-bit from the inner accumulator; the per-rank cursors
    /// resume receipt validation exactly where it stopped.
    pub fn checkpoint(&self) -> CrossNodeCheckpoint {
        CrossNodeCheckpoint {
            partition: self.partition,
            inner: self.inner.checkpoint(),
            received_per_rank: self.received_per_rank.clone(),
        }
    }

    /// Reconstruct a coordinator from a checkpoint, validating the cursor
    /// structure against the partition so corruption is rejected loudly.
    pub fn resume(state: CrossNodeCheckpoint) -> Result<Self, String> {
        if state.received_per_rank.len() != state.partition.n_ranks {
            return Err(format!(
                "CrossNodeCheckpoint: {} rank cursors for n_ranks = {}",
                state.received_per_rank.len(),
                state.partition.n_ranks
            ));
        }
        if state.inner.border_dim != state.partition.border_dim
            || state.inner.n_rows != state.partition.n_rows
            || state.inner.chunk_size != state.partition.chunk_size
        {
            return Err(
                "CrossNodeCheckpoint: inner accumulator partition disagrees with the cross-node \
                 partition"
                    .to_string(),
            );
        }
        for (rank, &cursor) in state.received_per_rank.iter().enumerate() {
            if cursor > state.partition.chunks_owned_by(rank) {
                return Err(format!(
                    "CrossNodeCheckpoint: rank {rank} cursor {cursor} exceeds its owned chunk \
                     count {}",
                    state.partition.chunks_owned_by(rank)
                ));
            }
        }
        let inner = StreamingBorderGram::resume(state.inner)?;
        Ok(Self {
            partition: state.partition,
            inner,
            received_per_rank: state.received_per_rank,
        })
    }

    /// Finish the pass, returning the `k×k` border Gram. Errors if any rank's
    /// sequence is incomplete. The result is a pure function of the row content
    /// and `(n_rows, chunk_size)` — identical bits for any node count, any
    /// arrival interleaving, and any checkpoint/resume history on either side.
    pub fn finish(self) -> Result<Array2<f64>, String> {
        for (rank, &cursor) in self.received_per_rank.iter().enumerate() {
            let owned = self.partition.chunks_owned_by(rank);
            if cursor != owned {
                return Err(format!(
                    "CrossNodeGramReduction: finish() with rank {rank} at ordinal {cursor} of \
                     {owned} owned chunks"
                ));
            }
        }
        self.inner.finish()
    }
}