inferencelayer 0.2.3

//! Fleet REPLICATION planning — how to allocate a heterogeneous worker pool between pipeline
//! **depth** (split one model instance across stages, to fit a bigger model) and **replication**
//! (run several independent pipelines, for more throughput).
//!
//! The two axes have different ceilings (see `ARCHITECTURE.md` §14):
//! - **Depth** is capped by the model's layer count (a stage needs ≥ 1 layer) and, in practice, by
//!   per-token latency (each stage is a network hop).
//! - **Replication** is capped by total VRAM — each replica must independently hold the whole model.
//!
//! This module is the *smart, model-aware* decision: given the model (layers + bytes), the worker
//! pool (per-device usable VRAM), and a selectable [`ReplicaPolicy`], it produces a concrete
//! replica→worker assignment, or a clear error when the pool can't satisfy the request. It is a
//! PURE planner — no sockets, no GPU — so it is trivially testable and the serving layer just
//! consumes its output.
use anyhow::{Result, ensure};

/// How to trade pipeline depth against replica count. Selectable from the CLI (`--replicas`).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ReplicaPolicy {
    /// One replica over as many workers as the model allows — maximum depth, fits the biggest
    /// model, lowest throughput. This is the classic single-pipeline behaviour.
    Depth,
    /// As many replicas as the pool can independently hold — maximum throughput. The model-aware
    /// default when the pool has spare capacity: a small model packs into many replicas, a huge
    /// one collapses to a single deep pipeline automatically.
    Throughput,
    /// Exactly `n` replicas — fails if the pool can't form that many.
    Fixed(usize),
}

impl ReplicaPolicy {
    /// Parse the `--replicas` argument: `1`/`depth` → [`Self::Depth`], `max`/`auto` →
    /// [`Self::Throughput`], a number `N` → [`Self::Fixed`]`(N)`.
    pub fn parse(s: &str) -> Result<Self> {
        match s.trim().to_ascii_lowercase().as_str() {
            "1" | "depth" => Ok(Self::Depth),
            "max" | "auto" | "throughput" => Ok(Self::Throughput),
            other => {
                let n: usize = other.parse().map_err(|_| {
                    anyhow::anyhow!("--replicas expects 1|depth|max|auto|<N>, got {s:?}")
                })?;
                ensure!(n >= 1, "--replicas must be ≥ 1");
                Ok(Self::Fixed(n))
            }
        }
    }
}

/// One pipeline replica: the worker indices (into the input VRAM slice) that form it, in descending
/// VRAM order, plus their combined usable VRAM.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Replica {
    pub workers: Vec<usize>,
    pub vram: u64,
}

/// The result of [`plan_replicas`]: the replicas to build, the leftover workers that couldn't
/// complete a replica, and a human-readable explanation of the allocation.
#[derive(Debug, Clone)]
pub struct FleetPlan {
    pub replicas: Vec<Replica>,
    pub spares: Vec<usize>,
    pub note: String,
}

/// Upper bound on replicas the pool could form for this model (VRAM-limited): `⌊ΣVRAM / model⌋`.
/// The device-count/layer cap can lower the achievable count below this; [`plan_replicas`] with
/// [`ReplicaPolicy::Throughput`] returns the actual achievable number.
pub fn max_replicas(model_bytes: u64, worker_vram: &[u64]) -> usize {
    if model_bytes == 0 {
        return 0;
    }
    (worker_vram.iter().sum::<u64>() / model_bytes) as usize
}

/// Indices of `worker_vram` sorted by descending VRAM (ties broken by index for determinism).
fn by_vram_desc(worker_vram: &[u64]) -> Vec<usize> {
    let mut idx: Vec<usize> = (0..worker_vram.len()).collect();
    idx.sort_by(|&a, &b| worker_vram[b].cmp(&worker_vram[a]).then(a.cmp(&b)));
    idx
}

/// Allocate the worker pool per `policy`. See [`ReplicaPolicy`] for the axes.
///
/// - `n_layers`: model depth (max devices per replica — a stage needs ≥ 1 layer).
/// - `model_bytes`: weights a replica must collectively hold (sum of its workers' VRAM ≥ this).
/// - `worker_vram`: usable VRAM per worker, in bytes.
pub fn plan_replicas(
    n_layers: usize,
    model_bytes: u64,
    worker_vram: &[u64],
    policy: ReplicaPolicy,
) -> Result<FleetPlan> {
    ensure!(!worker_vram.is_empty(), "no workers to plan");
    ensure!(n_layers >= 1, "model has no layers");
    ensure!(model_bytes > 0, "model has zero size");
    let sorted = by_vram_desc(worker_vram);
    let total_vram: u64 = worker_vram.iter().sum();

    match policy {
        ReplicaPolicy::Depth => {
            // One replica; a stage per worker, capped at n_layers (extra workers are spares — they
            // can't be placed without giving some stage 0 layers). Pick the highest-VRAM workers.
            let take = worker_vram.len().min(n_layers);
            let workers: Vec<usize> = sorted[..take].to_vec();
            let vram: u64 = workers.iter().map(|&i| worker_vram[i]).sum();
            ensure!(
                vram >= model_bytes,
                "model needs {:.2} GB but the {take} usable worker(s) hold only {:.2} GB — add VRAM or fewer/bigger devices",
                model_bytes as f64 / 1e9,
                vram as f64 / 1e9
            );
            let spares: Vec<usize> = sorted[take..].to_vec();
            let note = format!(
                "depth: 1 replica over {take} stage(s){}",
                if spares.is_empty() {
                    String::new()
                } else {
                    format!(
                        " ({} worker(s) idle — a {n_layers}-layer model can't use more than {n_layers} stages)",
                        spares.len()
                    )
                }
            );
            Ok(FleetPlan {
                replicas: vec![Replica { workers, vram }],
                spares,
                note,
            })
        }
        ReplicaPolicy::Throughput => {
            // Greedy, biggest-first: each replica grabs the largest remaining workers until it holds
            // the model (fewest devices ⇒ most replicas), capped at n_layers devices. Stop when the
            // remaining pool can't fill another replica.
            let mut replicas = Vec::new();
            let mut spares = Vec::new();
            let mut i = 0;
            while i < sorted.len() {
                let mut workers = Vec::new();
                let mut vram = 0u64;
                while i < sorted.len() && vram < model_bytes && workers.len() < n_layers {
                    let w = sorted[i];
                    workers.push(w);
                    vram += worker_vram[w];
                    i += 1;
                }
                if vram >= model_bytes {
                    replicas.push(Replica { workers, vram });
                } else {
                    // Couldn't complete this replica from the (smallest) remaining workers — with
                    // biggest-first, if these don't fit, nothing further will. They're spares.
                    spares.extend(workers);
                    spares.extend(sorted[i..].iter().copied());
                    break;
                }
            }
            ensure!(
                !replicas.is_empty(),
                "model needs {:.2} GB per replica but no {n_layers}-device group in the pool reaches it (ΣVRAM {:.2} GB)",
                model_bytes as f64 / 1e9,
                total_vram as f64 / 1e9
            );
            let note = format!(
                "throughput: {} replica(s){}",
                replicas.len(),
                if spares.is_empty() {
                    String::new()
                } else {
                    format!(" ({} worker(s) spare)", spares.len())
                }
            );
            Ok(FleetPlan {
                replicas,
                spares,
                note,
            })
        }
        ReplicaPolicy::Fixed(r) => {
            // Balance workers across exactly r bins (each new worker to the currently-lightest bin),
            // then require every bin to hold the model within the layer cap.
            ensure!(
                worker_vram.len() >= r,
                "asked for {r} replicas but only {} worker(s) — need ≥ 1 per replica",
                worker_vram.len()
            );
            let mut bins: Vec<Replica> = (0..r)
                .map(|_| Replica {
                    workers: Vec::new(),
                    vram: 0,
                })
                .collect();
            for &w in &sorted {
                // lightest bin that still has a free stage slot
                if let Some(b) = bins
                    .iter_mut()
                    .filter(|b| b.workers.len() < n_layers)
                    .min_by_key(|b| b.vram)
                {
                    b.workers.push(w);
                    b.vram += worker_vram[w];
                }
                // else: every bin is at the layer cap — remaining workers become spares below.
            }
            let placed: usize = bins.iter().map(|b| b.workers.len()).sum();
            let spares: Vec<usize> = sorted[placed..].to_vec();
            for (k, b) in bins.iter().enumerate() {
                ensure!(
                    !b.workers.is_empty(),
                    "cannot form {r} replicas: replica {} got no workers (pool too small)",
                    k + 1
                );
                ensure!(
                    b.vram >= model_bytes,
                    "cannot form {r} replicas: replica {} holds {:.2} GB < model {:.2} GB — fewer replicas or more VRAM (max feasible ≈ {})",
                    k + 1,
                    b.vram as f64 / 1e9,
                    model_bytes as f64 / 1e9,
                    max_replicas(model_bytes, worker_vram)
                );
            }
            let note = format!(
                "fixed: {r} replica(s){}",
                if spares.is_empty() {
                    String::new()
                } else {
                    format!(" ({} worker(s) spare)", spares.len())
                }
            );
            Ok(FleetPlan {
                replicas: bins,
                spares,
                note,
            })
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    const GB: u64 = 1_000_000_000;

    #[test]
    fn should_use_every_worker_in_one_replica_for_depth() {
        // 4 workers, 28-layer model that needs 6 GB — all four form one deep pipeline.
        let plan = plan_replicas(
            28,
            6 * GB,
            &[20 * GB, 2 * GB, 2 * GB, 2 * GB],
            ReplicaPolicy::Depth,
        )
        .unwrap();
        assert_eq!(plan.replicas.len(), 1);
        assert_eq!(plan.replicas[0].workers.len(), 4);
        assert!(plan.spares.is_empty());
    }

    #[test]
    fn should_pack_many_replicas_for_a_small_model() {
        // A 1.5 GB model, four 2 GB workers → each worker independently fits ⇒ 4 replicas.
        let plan = plan_replicas(
            28,
            1_500_000_000,
            &[2 * GB, 2 * GB, 2 * GB, 2 * GB],
            ReplicaPolicy::Throughput,
        )
        .unwrap();
        assert_eq!(plan.replicas.len(), 4);
        assert!(plan.replicas.iter().all(|r| r.workers.len() == 1));
    }

    #[test]
    fn should_group_workers_when_the_model_needs_several_per_replica() {
        // 3.5 GB model, six 2 GB workers → each replica needs 2 workers ⇒ 3 replicas.
        let plan =
            plan_replicas(28, 3_500_000_000, &[2 * GB; 6], ReplicaPolicy::Throughput).unwrap();
        assert_eq!(plan.replicas.len(), 3);
        assert!(plan.replicas.iter().all(|r| r.vram >= 3_500_000_000));
    }

    #[test]
    fn should_reject_more_replicas_than_the_pool_can_hold() {
        // Two 2 GB workers, 3 GB model → at most 1 replica; asking for 2 must error.
        let err =
            plan_replicas(28, 3 * GB, &[2 * GB, 2 * GB], ReplicaPolicy::Fixed(2)).unwrap_err();
        assert!(err.to_string().contains("cannot form 2 replicas"));
    }

    #[test]
    fn should_error_when_the_model_does_not_fit_the_pool_at_all() {
        // Three 1 GB workers can't hold a 5 GB model even pooled.
        assert!(plan_replicas(28, 5 * GB, &[GB, GB, GB], ReplicaPolicy::Depth).is_err());
        assert!(plan_replicas(28, 5 * GB, &[GB, GB, GB], ReplicaPolicy::Throughput).is_err());
    }

    #[test]
    fn should_cap_depth_at_the_layer_count_and_spare_the_rest() {
        // 2-layer model, 4 workers → one replica can use at most 2 stages; 2 workers spare.
        let plan = plan_replicas(2, GB, &[GB, GB, GB, GB], ReplicaPolicy::Depth).unwrap();
        assert_eq!(plan.replicas[0].workers.len(), 2);
        assert_eq!(plan.spares.len(), 2);
    }

    #[test]
    fn should_report_the_vram_limited_replica_ceiling() {
        assert_eq!(max_replicas(2 * GB, &[2 * GB, 2 * GB, 2 * GB, 2 * GB]), 4);
        assert_eq!(max_replicas(3 * GB, &[2 * GB, 2 * GB]), 1);
    }
}