forge-orchestration 0.6.0

//! Gang / coscheduling: all-or-nothing placement of a [`GangGroup`].
//!
//! A simulation unit ([`crate::scheduler::sim::SimCell`]) only makes progress if
//! *every* member is running: an agent policy with no world deadlocks waiting
//! for observations; a world with no agents produces no actions. Scheduling its
//! members independently risks the classic *partial-deadlock*: some members get
//! slots, the rest queue forever, and the running members burn GPUs producing
//! nothing. Gang scheduling forbids that: either all members are placed, or none
//! are and the reservation is rolled back.
//!
//! # Transaction model
//!
//! [`GangScheduler::schedule_gang`] is transactional against the live node store:
//!
//! 1. **Snapshot** every candidate node ([`NodeResources`] is `Clone`).
//! 2. **Try-place** all members on the snapshot, honoring [`CoPlacement`], using
//!    an inner [`SchedulingAlgorithm`] only to *score/order* candidate nodes
//!    (default [`BinPackScheduler`]). All mutations happen on the snapshot.
//! 3. **Commit or abort**: if and only if all members were placed on the
//!    snapshot, replay the exact same allocations onto the real nodes and return
//!    `committed = true`. Otherwise touch nothing real and return `committed =
//!    false` with a reason.
//!
//! Because step 2 only mutates the snapshot, an aborted gang leaves the cluster
//! byte-for-byte unchanged — there is no window where a half-placed gang holds
//! resources. Placement reuses [`NodeResources::can_fit`]/[`NodeResources::allocate`]
//! exactly, so gang accounting and single-workload accounting stay consistent.

use std::collections::HashMap;

use serde::{Deserialize, Serialize};
use tracing::{debug, info};

use super::algorithms::{BinPackScheduler, SchedulingAlgorithm};
use super::sim::{CoPlacement, GangGroup};
use super::{NodeResources, ResourceRequirements, Workload};
use crate::types::NodeId;

/// A resource reservation a committed gang made on one node, sufficient to
/// release it exactly later (CPU/mem via `resources`, GPUs via `gpu_ids`).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GangReservation {
    /// Node the reservation is on.
    pub node: NodeId,
    /// Aggregate resources reserved on that node by the gang.
    pub resources: ResourceRequirements,
    /// GPU device ids reserved on that node.
    pub gpu_ids: Vec<u32>,
}

/// Outcome of a gang placement attempt.
#[derive(Debug, Clone)]
pub struct GangDecision {
    /// The group that was scheduled.
    pub group_id: String,
    /// `(workload_id, node_id)` for every member, in member order. Empty when
    /// the gang failed.
    pub placements: Vec<(String, NodeId)>,
    /// Per-node reservations made by a committed gang, for exact release later.
    /// Empty when the gang failed.
    pub reservations: Vec<GangReservation>,
    /// Whether the placement was committed to the real node store.
    pub committed: bool,
    /// Human-readable reason (success detail or failure cause).
    pub reason: String,
}

impl GangDecision {
    fn failed(group_id: &str, reason: impl Into<String>) -> Self {
        Self {
            group_id: group_id.to_string(),
            placements: Vec::new(),
            reservations: Vec::new(),
            committed: false,
            reason: reason.into(),
        }
    }
}

/// Schedules [`GangGroup`]s atomically. The inner algorithm only scores nodes;
/// the all-or-nothing transaction is implemented here.
pub struct GangScheduler {
    algorithm: Box<dyn SchedulingAlgorithm>,
}

impl Default for GangScheduler {
    fn default() -> Self {
        Self::new()
    }
}

impl GangScheduler {
    /// Create a gang scheduler with the default [`BinPackScheduler`] for scoring.
    pub fn new() -> Self {
        Self {
            algorithm: Box::new(BinPackScheduler::new()),
        }
    }

    /// Create a gang scheduler with a custom node-scoring algorithm.
    pub fn with_algorithm<A: SchedulingAlgorithm + 'static>(algorithm: A) -> Self {
        Self {
            algorithm: Box::new(algorithm),
        }
    }

    /// Attempt to place every member of `group` atomically.
    ///
    /// On success the real `nodes` map is mutated (resources allocated) and the
    /// returned [`GangDecision`] has `committed = true`. On failure `nodes` is
    /// left exactly as it was found.
    pub fn schedule_gang(
        &self,
        group: &GangGroup,
        nodes: &mut HashMap<NodeId, NodeResources>,
    ) -> GangDecision {
        if group.is_empty() {
            return GangDecision::failed(&group.id, "gang has no members");
        }

        let result = match group.co_placement {
            CoPlacement::SameNode | CoPlacement::InterconnectLocalGpu => {
                self.commit_same_node(group, nodes)
            }
            CoPlacement::Spread => self.commit_spread(group, nodes),
        };

        match result {
            Ok((placements, reservations)) => {
                info!(
                    group = %group.id,
                    members = placements.len(),
                    policy = ?group.co_placement,
                    "gang committed atomically"
                );
                GangDecision {
                    group_id: group.id.clone(),
                    placements,
                    reservations,
                    committed: true,
                    reason: "all members placed".to_string(),
                }
            }
            Err(reason) => {
                debug!(group = %group.id, %reason, "gang placement aborted; nothing reserved");
                GangDecision::failed(&group.id, reason)
            }
        }
    }

    /// Rank node ids best-first by the inner algorithm's score for `workload`,
    /// breaking ties by available CPU for deterministic ordering. `exclude`
    /// skips already-used nodes (Spread). Reads `nodes`; clones no node.
    fn ranked_nodes(
        &self,
        workload: &Workload,
        nodes: &HashMap<NodeId, NodeResources>,
        exclude: &[NodeId],
    ) -> Vec<NodeId> {
        let mut ranked: Vec<NodeId> = nodes
            .keys()
            .copied()
            .filter(|n| !exclude.contains(n))
            .collect();
        ranked.sort_by(|a, b| {
            let sa = self.algorithm.score(workload, &nodes[a]);
            let sb = self.algorithm.score(workload, &nodes[b]);
            sb.partial_cmp(&sa)
                .unwrap_or(std::cmp::Ordering::Equal)
                .then_with(|| nodes[b].cpu_available().cmp(&nodes[a].cpu_available()))
        });
        ranked
    }

    /// Commit all members onto a single node (SameNode / InterconnectLocalGpu).
    ///
    /// Clones at most one candidate node at a time (the trial), never the whole
    /// store. The winning trial — which already holds the exact allocations,
    /// including chosen GPU device ids — is committed verbatim.
    fn commit_same_node(
        &self,
        group: &GangGroup,
        nodes: &mut HashMap<NodeId, NodeResources>,
    ) -> Result<(Vec<(String, NodeId)>, Vec<GangReservation>), String> {
        // The world member (member 0) is the heaviest anchor for ranking.
        let anchor = &group.members[0].workload;
        for node_id in self.ranked_nodes(anchor, nodes, &[]) {
            if group.co_placement.wants_gpu_locality()
                && !self.has_interconnect_local_gpus(&nodes[&node_id], group)
            {
                continue;
            }
            let before_gpus: Vec<u32> = nodes[&node_id].gpus_allocated.clone();
            let mut trial = nodes[&node_id].clone();
            if Self::try_fit_all_on_node(group, &mut trial) {
                // The GPU device ids the gang took on this node (diff vs before).
                let gpu_ids: Vec<u32> = trial
                    .gpus_allocated
                    .iter()
                    .filter(|d| !before_gpus.contains(d))
                    .copied()
                    .collect();
                // Trial proved every member fits; commit it verbatim (it carries
                // the exact CPU/mem/GPU allocations). Single mutation, atomic.
                nodes.insert(node_id, trial);
                let placements: Vec<(String, NodeId)> = group
                    .members
                    .iter()
                    .map(|m| (m.workload.id.clone(), node_id))
                    .collect();
                let reservations = vec![GangReservation {
                    node: node_id,
                    resources: Self::sum_resources(group),
                    gpu_ids,
                }];
                return Ok((placements, reservations));
            }
        }

        Err(format!(
            "no single node fits all {} members under {:?}",
            group.members.len(),
            group.co_placement
        ))
    }

    /// Sum all members' resource requirements (single-node reservation total).
    fn sum_resources(group: &GangGroup) -> ResourceRequirements {
        let mut cpu = 0u64;
        let mut mem = 0u64;
        let mut gpu = 0u32;
        for m in &group.members {
            cpu += m.workload.resources.cpu_millis;
            mem += m.workload.resources.memory_mb;
            gpu += m.workload.resources.gpu_count;
        }
        ResourceRequirements {
            cpu_millis: cpu,
            memory_mb: mem,
            gpu_count: gpu,
            gpu_memory_mb: 0,
            storage_mb: 0,
            network_mbps: 0,
        }
    }

    /// Commit each member on a distinct node (Spread), directly on the real
    /// store with exact rollback (CPU/mem + the GPU device ids actually taken).
    /// Touches at most one node per member; clones no node.
    fn commit_spread(
        &self,
        group: &GangGroup,
        nodes: &mut HashMap<NodeId, NodeResources>,
    ) -> Result<(Vec<(String, NodeId)>, Vec<GangReservation>), String> {
        let mut placements: Vec<(String, NodeId)> = Vec::with_capacity(group.members.len());
        // Per committed member: (node, resources, gpu_ids) for precise rollback.
        let mut committed: Vec<(NodeId, ResourceRequirements, Vec<u32>)> = Vec::new();
        let mut used: Vec<NodeId> = Vec::new();

        for member in &group.members {
            let req = &member.workload.resources;
            let chosen = self
                .ranked_nodes(&member.workload, nodes, &used)
                .into_iter()
                .find(|n| nodes.get(n).map(|node| node.can_fit(req)).unwrap_or(false));

            match chosen {
                Some(node_id) => {
                    let node = nodes.get_mut(&node_id).expect("ranked node exists");
                    // Capture the GPU device ids `allocate` reserves via a
                    // snapshot diff, so rollback frees exactly those.
                    let before: Vec<u32> = node.gpus_allocated.clone();
                    if !node.allocate(req) {
                        // can_fit just passed, so this should not happen; be safe.
                        Self::release_committed(&committed, nodes);
                        return Err(format!(
                            "spread: allocation raced for member {}",
                            member.workload.id
                        ));
                    }
                    let gpu_ids: Vec<u32> = node
                        .gpus_allocated
                        .iter()
                        .filter(|d| !before.contains(d))
                        .copied()
                        .collect();
                    committed.push((node_id, req.clone(), gpu_ids));
                    used.push(node_id);
                    placements.push((member.workload.id.clone(), node_id));
                }
                None => {
                    // All-or-nothing: undo every member committed so far.
                    Self::release_committed(&committed, nodes);
                    return Err(format!(
                        "spread: no distinct node available for member {}",
                        member.workload.id
                    ));
                }
            }
        }

        let reservations = committed
            .into_iter()
            .map(|(node, resources, gpu_ids)| GangReservation {
                node,
                resources,
                gpu_ids,
            })
            .collect();
        Ok((placements, reservations))
    }

    /// Release a set of committed allocations (exact CPU/mem + GPU device ids).
    fn release_committed(
        committed: &[(NodeId, ResourceRequirements, Vec<u32>)],
        nodes: &mut HashMap<NodeId, NodeResources>,
    ) {
        for (node_id, req, gpu_ids) in committed {
            if let Some(node) = nodes.get_mut(node_id) {
                node.release(req, gpu_ids);
            }
        }
    }

    /// Try to fit *all* members on a single node copy. Mutates `node` only; the
    /// caller decides whether to keep the result. Returns whether all fit.
    fn try_fit_all_on_node(group: &GangGroup, node: &mut NodeResources) -> bool {
        for member in &group.members {
            if !node.allocate(&member.workload.resources) {
                return false;
            }
        }
        true
    }

    /// Whether `node` has enough mutually interconnect-local (peer) GPUs for the
    /// group's total GPU demand. We approximate peer-locality the same way the
    /// existing [`super::algorithms::GpuLocalityScheduler`] does: contiguous
    /// device ids among the currently-free GPUs that satisfy the per-member
    /// memory requirement.
    fn has_interconnect_local_gpus(&self, node: &NodeResources, group: &GangGroup) -> bool {
        let needed = group.total_gpu_count() as usize;
        if needed == 0 {
            return true;
        }

        // Largest per-GPU memory requirement among members.
        let max_per_gpu_mem = group
            .members
            .iter()
            .map(|m| m.workload.resources.gpu_memory_mb)
            .max()
            .unwrap_or(0);

        let mut free: Vec<u32> = node
            .gpus
            .iter()
            .filter(|g| !node.gpus_allocated.contains(&g.device_id))
            .filter(|g| g.available_memory_mb() >= max_per_gpu_mem)
            .map(|g| g.device_id)
            .collect();
        free.sort_unstable();

        if free.len() < needed {
            return false;
        }

        // Look for a contiguous run of `needed` device ids (peer-capable group).
        let mut run = 1usize;
        for w in free.windows(2) {
            if w[1] == w[0] + 1 {
                run += 1;
                if run >= needed {
                    return true;
                }
            } else {
                run = 1;
            }
        }
        // No contiguous run, but enough free GPUs: still placeable, just not
        // peer-optimal. Treat as acceptable so the gang doesn't deadlock; the
        // scoring step already preferred better nodes first.
        free.len() >= needed
    }}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::scheduler::sim::{AgentPolicy, SimCell, SimWorld};
    use crate::scheduler::ResourceRequirements;
    use crate::types::{GpuResources, NodeId};
    use std::time::Duration;

    fn gpu_node(device_ids: &[u32], cpu: u64, mem: u64, gpu_mem: u64) -> NodeResources {
        let mut node = NodeResources::new(NodeId::new(), cpu, mem);
        for &id in device_ids {
            node = node.with_gpu(GpuResources::new(id, "A100", gpu_mem).with_tensor_cores(true));
        }
        node
    }

    fn store(nodes: Vec<NodeResources>) -> HashMap<NodeId, NodeResources> {
        nodes.into_iter().map(|n| (n.node_id, n)).collect()
    }

    /// A cell with a CPU world + 2 GPU agents, same-node.
    fn cell() -> SimCell {
        SimCell::new("c1", SimWorld::cpu(1000, 2048), Duration::from_millis(50))
            .with_agent(AgentPolicy::gpu("a0", 200, 256, 4096))
            .with_agent(AgentPolicy::gpu("a1", 200, 256, 4096))
    }

    #[test]
    fn gang_fully_places_when_one_node_fits_all() {
        let group = cell().gang_group();
        // Node with 2 contiguous GPUs and ample CPU/mem.
        let node = gpu_node(&[0, 1], 8000, 16384, 8192);
        let nid = node.node_id;
        let mut nodes = store(vec![node]);

        let decision = GangScheduler::new().schedule_gang(&group, &mut nodes);

        assert!(decision.committed, "reason: {}", decision.reason);
        assert_eq!(decision.placements.len(), 3);
        // All members on the same node.
        assert!(decision.placements.iter().all(|(_, n)| *n == nid));
        // Real node store reflects the allocation: 2 GPUs taken.
        assert_eq!(nodes[&nid].gpus_allocated.len(), 2);
        assert_eq!(nodes[&nid].cpu_allocated, 1000 + 200 + 200);
    }

    #[test]
    fn gang_fully_fails_and_reserves_nothing_when_no_node_fits() {
        let group = cell().gang_group();
        // Only ONE GPU available; the gang needs two => must fail atomically.
        let node = gpu_node(&[0], 8000, 16384, 8192);
        let nid = node.node_id;
        let before_cpu = node.cpu_allocated;
        let mut nodes = store(vec![node]);

        let decision = GangScheduler::new().schedule_gang(&group, &mut nodes);

        assert!(!decision.committed);
        assert!(decision.placements.is_empty());
        // Nothing reserved: node untouched.
        assert_eq!(nodes[&nid].cpu_allocated, before_cpu);
        assert_eq!(nodes[&nid].gpus_allocated.len(), 0);
    }

    #[test]
    fn gang_does_not_split_across_two_partial_nodes() {
        let group = cell().gang_group();
        // Two nodes each with exactly ONE GPU. Total = 2 GPUs, enough in
        // aggregate, but SameNode forbids splitting => must fail and reserve
        // nothing on either node.
        let n1 = gpu_node(&[0], 8000, 16384, 8192);
        let n2 = gpu_node(&[0], 8000, 16384, 8192);
        let id1 = n1.node_id;
        let id2 = n2.node_id;
        let mut nodes = store(vec![n1, n2]);

        let decision = GangScheduler::new().schedule_gang(&group, &mut nodes);

        assert!(!decision.committed);
        assert_eq!(nodes[&id1].gpus_allocated.len(), 0);
        assert_eq!(nodes[&id2].gpus_allocated.len(), 0);
    }

    #[test]
    fn spread_places_members_on_distinct_nodes() {
        // Three CPU-only members, Spread across three nodes.
        let cell = SimCell::new("s1", SimWorld::cpu(500, 512), Duration::from_millis(50))
            .with_agent(AgentPolicy::new(
                "a0",
                ResourceRequirements::new().cpu(500).memory(512),
            ))
            .with_agent(AgentPolicy::new(
                "a1",
                ResourceRequirements::new().cpu(500).memory(512),
            ))
            .with_co_placement(CoPlacement::Spread);
        let group = cell.gang_group();

        let nodes_vec: Vec<NodeResources> = (0..3)
            .map(|_| NodeResources::new(NodeId::new(), 4000, 4096))
            .collect();
        let mut nodes = store(nodes_vec);

        let decision = GangScheduler::new().schedule_gang(&group, &mut nodes);

        assert!(decision.committed, "reason: {}", decision.reason);
        assert_eq!(decision.placements.len(), 3);
        // All placement nodes are distinct.
        let mut placed: Vec<NodeId> = decision.placements.iter().map(|(_, n)| *n).collect();
        placed.sort_by_key(|n| *n.as_uuid());
        placed.dedup();
        assert_eq!(placed.len(), 3, "spread must use distinct nodes");
    }

    #[test]
    fn spread_fails_when_too_few_nodes() {
        let cell = SimCell::new("s2", SimWorld::cpu(500, 512), Duration::from_millis(50))
            .with_agent(AgentPolicy::new(
                "a0",
                ResourceRequirements::new().cpu(500).memory(512),
            ))
            .with_co_placement(CoPlacement::Spread);
        let group = cell.gang_group(); // 2 members

        // Only one node => cannot spread 2 members.
        let node = NodeResources::new(NodeId::new(), 4000, 4096);
        let nid = node.node_id;
        let mut nodes = store(vec![node]);

        let decision = GangScheduler::new().schedule_gang(&group, &mut nodes);
        assert!(!decision.committed);
        assert_eq!(nodes[&nid].cpu_allocated, 0, "spread abort reserves nothing");
    }

    #[test]
    fn spread_rollback_frees_committed_gpu_devices() {
        // world + 2 GPU agents, Spread, but only 2 nodes: the 3rd member can't
        // get a distinct node, so the gang aborts. A GPU agent committed before
        // the failure must have its GPU *device* released, not just CPU/mem.
        // (The previous rollback passed empty gpu_ids and leaked the device.)
        let cell = SimCell::new("g", SimWorld::cpu(1000, 1024), Duration::from_millis(50))
            .with_agent(AgentPolicy::gpu("a0", 200, 256, 4096))
            .with_agent(AgentPolicy::gpu("a1", 200, 256, 4096))
            .with_co_placement(CoPlacement::Spread);
        let group = cell.gang_group(); // 3 members

        let n1 = gpu_node(&[0], 8000, 16384, 8192);
        let n2 = gpu_node(&[0], 8000, 16384, 8192);
        let id1 = n1.node_id;
        let id2 = n2.node_id;
        let mut nodes = store(vec![n1, n2]);

        let decision = GangScheduler::new().schedule_gang(&group, &mut nodes);

        assert!(!decision.committed);
        assert!(decision.placements.is_empty());
        // Full rollback: neither CPU nor GPU devices left reserved on any node.
        assert_eq!(nodes[&id1].cpu_allocated, 0);
        assert_eq!(nodes[&id2].cpu_allocated, 0);
        assert_eq!(nodes[&id1].gpus_allocated.len(), 0, "GPU device must be freed on rollback");
        assert_eq!(nodes[&id2].gpus_allocated.len(), 0, "GPU device must be freed on rollback");
    }
}