forge_orchestration/scheduler/

sim.rs

//! Simulation-native workload model (the Eustress Engine primitive)
//!
//! Forge's differentiating job is running "AI models living inside simulation
//! environments at scale." A simulation environment is not a single workload:
//! it is a *world shard* (the authoritative sim state advancing one tick at a
//! time) plus the *agent policy models* that observe and act inside that world.
//! Those members are useless apart from one another — an agent policy with no
//! world to act in, or a world with no agents, makes no progress — so they must
//! be scheduled together, co-located, and torn down together.
//!
//! This module gives that primitive a name ([`SimCell`]) and expands it into the
//! scheduler's existing [`Workload`] type so the rest of the stack (constraints,
//! affinity, queues, node accounting) works unchanged. The all-or-nothing gang
//! placement of the expanded members is handled by [`crate::scheduler::gang`];
//! the tick-deadline ordering is handled by [`crate::scheduler::deadline`].
//!
//! # Design notes
//!
//! [`Workload`] has no free-form metadata map, so the linkage between a member
//! and its cell/gang is carried two ways:
//!
//! 1. **Deterministically inside the id**: a member's [`Workload::id`] is
//!    `"simcell/{cell_id}/{role}/{ordinal}"`. Given any member workload you can
//!    recover its owning cell with [`parse_member_id`].
//! 2. **Structurally via [`SimMember`]**: [`SimCell::members`] returns the
//!    workloads, while [`SimCell::gang_group`] returns a [`GangGroup`] that keeps
//!    the role/ordinal/co-placement key alongside each workload for the gang
//!    scheduler. No metadata map is required.

use std::time::Duration;

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};

use super::{ResourceRequirements, Workload};

/// Axis-aligned spherical region used for spatial interest management.
///
/// Two cells whose regions do not overlap cannot influence each other this tick,
/// which lets a higher layer decide they need not share interconnect locality.
/// Kept deliberately minimal (sphere = center + radius) so it is cheap to test
/// for overlap; richer AABB/oct-tree schemes can wrap this later.
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub struct Region3D {
    /// World-space X coordinate of the region center.
    pub x: f64,
    /// World-space Y coordinate of the region center.
    pub y: f64,
    /// World-space Z coordinate of the region center.
    pub z: f64,
    /// Interest radius around the center.
    pub radius: f64,
}

impl Region3D {
    /// Create a new spherical region.
    pub fn new(x: f64, y: f64, z: f64, radius: f64) -> Self {
        Self { x, y, z, radius }
    }

    /// Squared distance between two region centers (avoids a `sqrt`).
    pub fn center_distance_sq(&self, other: &Region3D) -> f64 {
        let dx = self.x - other.x;
        let dy = self.y - other.y;
        let dz = self.z - other.z;
        dx * dx + dy * dy + dz * dz
    }

    /// Whether two interest spheres overlap (their regions can interact).
    pub fn overlaps(&self, other: &Region3D) -> bool {
        let r = self.radius + other.radius;
        self.center_distance_sq(other) <= r * r
    }
}

/// How the members of a [`SimCell`] must be placed relative to one another.
///
/// This is the *constraint* the gang scheduler honors; it is independent of the
/// node-scoring algorithm used to pick *which* node.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum CoPlacement {
    /// All members must land on the same node (shared host memory / loopback).
    SameNode,
    /// All members must land on the same node *and* their GPUs should be
    /// interconnect-local (prefer contiguous / peer-capable devices) so the
    /// world↔agent tensor exchange stays on-box (NVLink-style).
    InterconnectLocalGpu,
    /// Members should be spread across distinct nodes (e.g. replicated worlds
    /// for fault isolation). All members must still be placeable, but on
    /// different nodes.
    Spread,
}

impl CoPlacement {
    /// Whether this policy requires every member on a single node.
    pub fn requires_same_node(&self) -> bool {
        matches!(self, CoPlacement::SameNode | CoPlacement::InterconnectLocalGpu)
    }

    /// Whether this policy wants interconnect-local (peer) GPUs.
    pub fn wants_gpu_locality(&self) -> bool {
        matches!(self, CoPlacement::InterconnectLocalGpu)
    }
}

/// Role of a member within a [`SimCell`]. Used to build deterministic ids and to
/// let the gang scheduler reason about world-vs-agent without string parsing.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum MemberRole {
    /// The authoritative simulation shard (advances world state per tick).
    World,
    /// An agent policy model observing/acting inside the world.
    Agent,
}

impl MemberRole {
    /// Stable lowercase token used inside workload ids.
    pub fn token(&self) -> &'static str {
        match self {
            MemberRole::World => "world",
            MemberRole::Agent => "agent",
        }
    }
}

/// An agent policy model co-resident with a sim world.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentPolicy {
    /// Human-readable policy name (e.g. `"ppo-forager"`).
    pub name: String,
    /// Resources this policy needs. Agents typically need a GPU slice; set
    /// `gpu_count`/`gpu_memory_mb` via [`ResourceRequirements::gpu`].
    pub resources: ResourceRequirements,
}

impl AgentPolicy {
    /// Create an agent policy with explicit resource requirements.
    pub fn new(name: impl Into<String>, resources: ResourceRequirements) -> Self {
        Self {
            name: name.into(),
            resources,
        }
    }

    /// Convenience constructor for a GPU-backed policy.
    ///
    /// `gpu_memory_mb` is per-GPU; `cpu_millis`/`memory_mb` are the host-side
    /// sidecar footprint of the policy server.
    pub fn gpu(
        name: impl Into<String>,
        cpu_millis: u64,
        memory_mb: u64,
        gpu_memory_mb: u64,
    ) -> Self {
        Self::new(
            name,
            ResourceRequirements::new()
                .cpu(cpu_millis)
                .memory(memory_mb)
                .gpu(1, gpu_memory_mb),
        )
    }
}

/// The simulation shard (authoritative world) of a [`SimCell`].
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SimWorld {
    /// Resources the world simulator needs (often CPU-heavy, GPU-optional).
    pub resources: ResourceRequirements,
}

impl SimWorld {
    /// Create a world shard with explicit resource requirements.
    pub fn new(resources: ResourceRequirements) -> Self {
        Self { resources }
    }

    /// Convenience constructor for a CPU-only world simulator.
    pub fn cpu(cpu_millis: u64, memory_mb: u64) -> Self {
        Self::new(ResourceRequirements::new().cpu(cpu_millis).memory(memory_mb))
    }
}

/// A gang-schedulable, checkpointable simulation unit: one world shard plus its
/// co-resident agent policy models, advancing on a fixed tick cadence.
///
/// This is *the* sim-native workload type. It is also called a "habitat" in
/// product docs; [`SimCell`] is the canonical name here.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SimCell {
    /// Stable identifier for the cell (used to build member ids and the gang id).
    pub id: String,
    /// Optional spatial footprint for interest management.
    pub region: Option<Region3D>,
    /// The authoritative simulation world shard.
    pub world: SimWorld,
    /// Co-resident agent policy models.
    pub agents: Vec<AgentPolicy>,
    /// Tick cadence (wall-clock budget per simulation step).
    pub tick: Duration,
    /// Absolute time the next tick is due — the EEVDF virtual deadline used by
    /// [`crate::scheduler::deadline`].
    pub next_deadline: DateTime<Utc>,
    /// How members must be placed relative to one another.
    pub co_placement: CoPlacement,
    /// Base scheduling priority applied to every expanded member.
    pub priority: i32,
}

/// A single expanded member of a [`SimCell`], pairing the scheduler-native
/// [`Workload`] with the structural metadata the gang scheduler needs.
#[derive(Debug, Clone)]
pub struct SimMember {
    /// The scheduler-native workload (carries [`ResourceRequirements`]).
    pub workload: Workload,
    /// Role within the cell.
    pub role: MemberRole,
    /// Ordinal within the role (world is always `0`; agents are `0..n`).
    pub ordinal: usize,
}

/// Identity recovered from a member workload id.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MemberIdentity {
    /// Owning cell id.
    pub cell_id: String,
    /// Member role.
    pub role: MemberRole,
    /// Ordinal within the role.
    pub ordinal: usize,
}

/// Prefix for all sim member workload ids.
pub const MEMBER_ID_PREFIX: &str = "simcell";

/// Build the deterministic workload id for a member.
pub fn member_id(cell_id: &str, role: MemberRole, ordinal: usize) -> String {
    format!("{MEMBER_ID_PREFIX}/{cell_id}/{}/{ordinal}", role.token())
}

/// Co-placement group key shared by all members of a cell. Members carrying the
/// same key must satisfy the cell's [`CoPlacement`] policy together.
pub fn co_placement_key(cell_id: &str) -> String {
    format!("{MEMBER_ID_PREFIX}/{cell_id}")
}

/// Recover the cell/role/ordinal from a member workload id produced by
/// [`member_id`]. Returns `None` if the id is not a sim member id.
pub fn parse_member_id(id: &str) -> Option<MemberIdentity> {
    let rest = id.strip_prefix(MEMBER_ID_PREFIX)?.strip_prefix('/')?;
    // rest == "{cell_id}/{role}/{ordinal}". cell_id itself must not contain '/'.
    let mut parts = rest.splitn(3, '/');
    let cell_id = parts.next()?.to_string();
    let role = match parts.next()? {
        "world" => MemberRole::World,
        "agent" => MemberRole::Agent,
        _ => return None,
    };
    let ordinal: usize = parts.next()?.parse().ok()?;
    if cell_id.is_empty() {
        return None;
    }
    Some(MemberIdentity {
        cell_id,
        role,
        ordinal,
    })
}

impl SimCell {
    /// Create a new sim cell. `next_deadline` defaults to `now + tick`.
    pub fn new(id: impl Into<String>, world: SimWorld, tick: Duration) -> Self {
        let tick_chrono = chrono::Duration::from_std(tick)
            .unwrap_or_else(|_| chrono::Duration::milliseconds(0));
        Self {
            id: id.into(),
            region: None,
            world,
            agents: Vec::new(),
            tick,
            next_deadline: Utc::now() + tick_chrono,
            co_placement: CoPlacement::SameNode,
            priority: 0,
        }
    }

    /// Attach a spatial region for interest management.
    pub fn with_region(mut self, region: Region3D) -> Self {
        self.region = Some(region);
        self
    }

    /// Add an agent policy member.
    pub fn with_agent(mut self, agent: AgentPolicy) -> Self {
        self.agents.push(agent);
        self
    }

    /// Set the co-placement policy.
    pub fn with_co_placement(mut self, policy: CoPlacement) -> Self {
        self.co_placement = policy;
        self
    }

    /// Set the base priority applied to every member.
    pub fn with_priority(mut self, priority: i32) -> Self {
        self.priority = priority;
        self
    }

    /// Set an explicit next-tick deadline (overrides the `now + tick` default).
    pub fn with_next_deadline(mut self, deadline: DateTime<Utc>) -> Self {
        self.next_deadline = deadline;
        self
    }

    /// Total number of members (world + agents).
    pub fn member_count(&self) -> usize {
        1 + self.agents.len()
    }

    /// Advance the deadline by one tick (used after a tick completes).
    pub fn advance_deadline(&mut self) {
        let tick_chrono = chrono::Duration::from_std(self.tick)
            .unwrap_or_else(|_| chrono::Duration::milliseconds(0));
        self.next_deadline += tick_chrono;
    }

    /// Expand this cell into the scheduler's [`SimMember`] list. The world is
    /// member 0; agents follow in declaration order. Each member's workload id
    /// is deterministic ([`member_id`]) and its name embeds the cell id.
    pub fn expand(&self) -> Vec<SimMember> {
        let mut members = Vec::with_capacity(self.member_count());

        let world_workload = Workload::new(
            member_id(&self.id, MemberRole::World, 0),
            format!("{}-world", self.id),
        )
        .with_resources(self.world.resources.clone())
        .with_priority(self.priority);
        members.push(SimMember {
            workload: world_workload,
            role: MemberRole::World,
            ordinal: 0,
        });

        for (i, agent) in self.agents.iter().enumerate() {
            let agent_workload = Workload::new(
                member_id(&self.id, MemberRole::Agent, i),
                format!("{}-{}", self.id, agent.name),
            )
            .with_resources(agent.resources.clone())
            .with_priority(self.priority);
            members.push(SimMember {
                workload: agent_workload,
                role: MemberRole::Agent,
                ordinal: i,
            });
        }

        members
    }

    /// Expand into bare [`Workload`]s, as required by the task signature
    /// `fn members(&self) -> Vec<Workload>`. Prefer [`SimCell::expand`] when you
    /// need the role/ordinal structure.
    pub fn members(&self) -> Vec<Workload> {
        self.expand().into_iter().map(|m| m.workload).collect()
    }

    /// Build the all-or-nothing [`GangGroup`] for this cell.
    pub fn gang_group(&self) -> GangGroup {
        GangGroup {
            id: co_placement_key(&self.id),
            members: self.expand(),
            all_or_nothing: true,
            co_placement: self.co_placement,
        }
    }
}

/// An all-or-nothing, co-placed set of member workloads.
///
/// Consumed by [`crate::scheduler::gang::GangScheduler`]. The `id` is the cell's
/// co-placement key; placement either commits *every* member or reserves nothing.
#[derive(Debug, Clone)]
pub struct GangGroup {
    /// Group identifier (the cell co-placement key).
    pub id: String,
    /// The member workloads to place together.
    pub members: Vec<SimMember>,
    /// If true, partial placement is forbidden (the sim-native default).
    pub all_or_nothing: bool,
    /// Co-placement constraint for the whole group.
    pub co_placement: CoPlacement,
}

impl GangGroup {
    /// Number of members in the group.
    pub fn len(&self) -> usize {
        self.members.len()
    }

    /// Whether the group is empty.
    pub fn is_empty(&self) -> bool {
        self.members.is_empty()
    }

    /// Total GPU count required across all members.
    pub fn total_gpu_count(&self) -> u32 {
        self.members
            .iter()
            .map(|m| m.workload.resources.gpu_count)
            .sum()
    }

    /// Iterator over the member workloads.
    pub fn workloads(&self) -> impl Iterator<Item = &Workload> {
        self.members.iter().map(|m| &m.workload)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn sample_cell() -> SimCell {
        SimCell::new("alpha", SimWorld::cpu(2000, 4096), Duration::from_millis(50))
            .with_agent(AgentPolicy::gpu("policy-a", 250, 512, 8192))
            .with_agent(AgentPolicy::gpu("policy-b", 250, 512, 8192))
            .with_co_placement(CoPlacement::InterconnectLocalGpu)
            .with_priority(100)
    }

    #[test]
    fn expands_to_world_plus_agents() {
        let cell = sample_cell();
        let members = cell.expand();

        // 1 world + 2 agents.
        assert_eq!(members.len(), 3);
        assert_eq!(cell.member_count(), 3);

        // Member 0 is the world, CPU-only.
        assert_eq!(members[0].role, MemberRole::World);
        assert_eq!(members[0].workload.resources.gpu_count, 0);
        assert_eq!(members[0].workload.resources.cpu_millis, 2000);

        // Agents are GPU-backed and carry the cell priority.
        assert_eq!(members[1].role, MemberRole::Agent);
        assert_eq!(members[1].workload.resources.gpu_count, 1);
        assert_eq!(members[1].workload.resources.gpu_memory_mb, 8192);
        assert_eq!(members[1].workload.priority, 100);
        assert_eq!(members[2].ordinal, 1);
    }

    #[test]
    fn member_ids_are_deterministic_and_parseable() {
        let cell = sample_cell();
        let members = cell.members();

        assert_eq!(members[0].id, "simcell/alpha/world/0");
        assert_eq!(members[1].id, "simcell/alpha/agent/0");
        assert_eq!(members[2].id, "simcell/alpha/agent/1");

        let id = parse_member_id(&members[2].id).unwrap();
        assert_eq!(id.cell_id, "alpha");
        assert_eq!(id.role, MemberRole::Agent);
        assert_eq!(id.ordinal, 1);

        // Non-member ids are rejected.
        assert!(parse_member_id("some-other-workload").is_none());
        assert!(parse_member_id("simcell//world/0").is_none());
    }

    #[test]
    fn gang_group_carries_co_placement_and_gpu_totals() {
        let cell = sample_cell();
        let group = cell.gang_group();

        assert_eq!(group.id, "simcell/alpha");
        assert!(group.all_or_nothing);
        assert_eq!(group.co_placement, CoPlacement::InterconnectLocalGpu);
        assert_eq!(group.len(), 3);
        // Two GPU agents => 2 GPUs total.
        assert_eq!(group.total_gpu_count(), 2);
    }

    #[test]
    fn region_overlap_is_symmetric_and_correct() {
        let a = Region3D::new(0.0, 0.0, 0.0, 5.0);
        let b = Region3D::new(8.0, 0.0, 0.0, 5.0); // distance 8 < 10 => overlap
        let c = Region3D::new(20.0, 0.0, 0.0, 1.0); // distance 20 > 6 => no overlap

        assert!(a.overlaps(&b));
        assert!(b.overlaps(&a));
        assert!(!a.overlaps(&c));
    }

    #[test]
    fn deadline_defaults_and_advances_by_tick() {
        let mut cell = SimCell::new("t", SimWorld::cpu(100, 128), Duration::from_millis(100));
        let first = cell.next_deadline;
        cell.advance_deadline();
        let delta = cell.next_deadline - first;
        assert_eq!(delta.num_milliseconds(), 100);
    }
}