forge-orchestration 0.5.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 tracing::{debug, info, warn};

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

/// 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)>,
    /// 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(),
            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");
        }

        // ---- Step 1: snapshot the candidate nodes -------------------------
        // Cloning the whole store is acceptable: gangs are small and infrequent
        // relative to single-workload scheduling, and it gives us a trivially
        // correct rollback (drop the snapshot).
        let mut snapshot: HashMap<NodeId, NodeResources> = nodes.clone();

        // ---- Step 2: try to place all members on the snapshot -------------
        let placements = match group.co_placement {
            CoPlacement::SameNode | CoPlacement::InterconnectLocalGpu => {
                self.place_same_node(group, &mut snapshot)
            }
            CoPlacement::Spread => self.place_spread(group, &mut snapshot),
        };

        let placements = match placements {
            Ok(p) => p,
            Err(reason) => {
                debug!(group = %group.id, %reason, "gang placement aborted; nothing reserved");
                return GangDecision::failed(&group.id, reason);
            }
        };

        // Invariant: all-or-nothing groups must have placed every member.
        if group.all_or_nothing && placements.len() != group.members.len() {
            warn!(
                group = %group.id,
                placed = placements.len(),
                expected = group.members.len(),
                "partial gang placement detected; aborting"
            );
            return GangDecision::failed(
                &group.id,
                format!(
                    "partial placement ({}/{} members); all-or-nothing aborts",
                    placements.len(),
                    group.members.len()
                ),
            );
        }

        // ---- Step 3: commit by replaying allocations onto the real nodes --
        // The snapshot already proved every allocation fits; replay on `nodes`.
        for (committed, (member, node_id)) in group
            .members
            .iter()
            .zip(placements.iter().map(|(_, n)| *n))
            .enumerate()
        {
            let ok = nodes
                .get_mut(&node_id)
                .map(|node| node.allocate(&member.workload.resources))
                .unwrap_or(false);
            if !ok {
                // Defensive: a concurrent mutation changed the real store
                // between snapshot and commit. Roll back only the prefix we
                // actually committed (members `0..committed`) and fail cleanly.
                self.rollback(group, &placements[..committed], nodes);
                return GangDecision::failed(
                    &group.id,
                    "node store changed during commit; rolled back",
                );
            }
        }

        info!(
            group = %group.id,
            members = placements.len(),
            policy = ?group.co_placement,
            "gang committed atomically"
        );

        GangDecision {
            group_id: group.id.clone(),
            placements,
            committed: true,
            reason: "all members placed".to_string(),
        }
    }

    /// Place every member on a single node (SameNode / InterconnectLocalGpu).
    ///
    /// Returns the per-member `(workload_id, node_id)` placement on success, or a
    /// failure reason. All allocation happens on `snapshot`.
    fn place_same_node(
        &self,
        group: &GangGroup,
        snapshot: &mut HashMap<NodeId, NodeResources>,
    ) -> Result<Vec<(String, NodeId)>, String> {
        // Rank candidate nodes by the inner algorithm's score for the *world*
        // member (the heaviest anchor), best first. Ties broken by available
        // CPU so the ordering is deterministic for tests.
        let anchor = &group.members[0].workload;
        let mut ranked: Vec<NodeId> = snapshot.keys().copied().collect();
        ranked.sort_by(|a, b| {
            let sa = self.algorithm.score(anchor, &snapshot[a]);
            let sb = self.algorithm.score(anchor, &snapshot[b]);
            sb.partial_cmp(&sa)
                .unwrap_or(std::cmp::Ordering::Equal)
                .then_with(|| snapshot[b].cpu_available().cmp(&snapshot[a].cpu_available()))
        });

        for node_id in ranked {
            // Work on a per-node trial copy so a failed node leaves the outer
            // snapshot untouched for the next candidate.
            let mut trial = snapshot[&node_id].clone();

            if group.co_placement.wants_gpu_locality()
                && !self.has_interconnect_local_gpus(&trial, group)
            {
                continue;
            }

            if Self::try_fit_all_on_node(group, &mut trial) {
                // Commit the trial node back into the snapshot and record
                // placements (all members map to this one node).
                snapshot.insert(node_id, trial);
                let placements = group
                    .members
                    .iter()
                    .map(|m| (m.workload.id.clone(), node_id))
                    .collect();
                return Ok(placements);
            }
        }

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

    /// Place each member on a distinct node (Spread).
    fn place_spread(
        &self,
        group: &GangGroup,
        snapshot: &mut HashMap<NodeId, NodeResources>,
    ) -> Result<Vec<(String, NodeId)>, String> {
        let mut placements = Vec::with_capacity(group.members.len());
        let mut used: Vec<NodeId> = Vec::with_capacity(group.members.len());

        for member in &group.members {
            // Best-scoring unused node that fits this member.
            let mut ranked: Vec<NodeId> = snapshot
                .keys()
                .copied()
                .filter(|n| !used.contains(n))
                .collect();
            ranked.sort_by(|a, b| {
                let sa = self.algorithm.score(&member.workload, &snapshot[a]);
                let sb = self.algorithm.score(&member.workload, &snapshot[b]);
                sb.partial_cmp(&sa).unwrap_or(std::cmp::Ordering::Equal)
            });

            let chosen = ranked.into_iter().find(|n| {
                snapshot
                    .get(n)
                    .map(|node| node.can_fit(&member.workload.resources))
                    .unwrap_or(false)
            });

            match chosen {
                Some(node_id) => {
                    // Allocate on the snapshot so subsequent members see it taken.
                    if let Some(node) = snapshot.get_mut(&node_id) {
                        node.allocate(&member.workload.resources);
                    }
                    used.push(node_id);
                    placements.push((member.workload.id.clone(), node_id));
                }
                None => {
                    return Err(format!(
                        "spread: no distinct node available for member {}",
                        member.workload.id
                    ));
                }
            }
        }

        Ok(placements)
    }

    /// 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
    }

    /// Roll back already-committed member allocations on the real node store.
    fn rollback(
        &self,
        group: &GangGroup,
        placements: &[(String, NodeId)],
        nodes: &mut HashMap<NodeId, NodeResources>,
    ) {
        let by_id: HashMap<&str, &SimMember> = group
            .members
            .iter()
            .map(|m| (m.workload.id.as_str(), m))
            .collect();

        for (workload_id, node_id) in placements {
            if let (Some(member), Some(node)) = (by_id.get(workload_id.as_str()), nodes.get_mut(node_id))
            {
                // Release CPU/mem; gpu ids are tracked at a higher layer, so we
                // release the count by passing the currently-allocated ids that
                // match this member's GPU demand is not knowable here — release
                // CPU/memory and leave GPU ids to the owning Scheduler::release.
                node.release(&member.workload.resources, &[]);
            }
        }
    }
}

#[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");
    }
}