use crate::adapter::net::behavior::fold::{
CapabilityFold, Fold, IslandId, IslandTopologyFold, JobId, NodeId, ReservationFold,
};
use crate::adapter::net::behavior::gang::{
commit_active, match_islands, release_island, single_island_claim, ActiveCommitOutcome,
ClaimError, ClaimOutcome, Claimant, Epoch, MatchCriteria, ReplicaCohort, ReplicaSet,
};
use crate::adapter::net::current_timestamp_micros;
use crate::adapter::net::identity::EntityKeypair;
use super::adapter::WorkflowAdapter;
use super::types::TaskId;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ActiveClaim {
pub island: IslandId,
}
pub struct CapabilityRequirement {
pub criteria: MatchCriteria,
pub reserve_ttl_us: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ClaimResult {
Active(ActiveClaim),
Rejected,
}
pub trait ClaimPipeline {
type Error;
fn claim(&mut self, req: &CapabilityRequirement) -> Result<ClaimResult, Self::Error>;
fn release(&mut self, claim: &ActiveClaim) -> Result<(), Self::Error>;
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StepGate {
Running(ActiveClaim),
Waiting,
}
#[derive(Debug)]
pub enum StepError<E> {
Pipeline(E),
Workflow(super::super::error::CortexAdapterError),
}
pub fn drive_capability_step<P: ClaimPipeline>(
wf: &WorkflowAdapter,
pipeline: &mut P,
task: TaskId,
req: &CapabilityRequirement,
) -> Result<StepGate, StepError<P::Error>> {
match pipeline.claim(req).map_err(StepError::Pipeline)? {
ClaimResult::Active(claim) => {
wf.start(task).map_err(StepError::Workflow)?;
Ok(StepGate::Running(claim))
}
ClaimResult::Rejected => {
wf.wait(task).map_err(StepError::Workflow)?;
Ok(StepGate::Waiting)
}
}
}
pub struct GangClaimContext<'a> {
pub capability: &'a Fold<CapabilityFold>,
pub topology: &'a Fold<IslandTopologyFold>,
pub reservations: &'a Fold<ReservationFold>,
pub keypair: &'a EntityKeypair,
pub node_id: NodeId,
}
pub struct GangClaimPipeline<'a> {
ctx: GangClaimContext<'a>,
claimant: Claimant<'a>,
cohort: ReplicaCohort,
replica_set: ReplicaSet,
reachable: Vec<NodeId>,
job: JobId,
}
impl<'a> GangClaimPipeline<'a> {
pub fn new(
ctx: GangClaimContext<'a>,
replica_set: ReplicaSet,
reachable: Vec<NodeId>,
job: JobId,
) -> Self {
Self::with_generation(ctx, replica_set, reachable, job, 1)
}
pub fn with_generation(
ctx: GangClaimContext<'a>,
replica_set: ReplicaSet,
reachable: Vec<NodeId>,
job: JobId,
start_generation: u64,
) -> Self {
let cohort = ReplicaCohort::new(replica_set.members());
let claimant =
Claimant::with_generation(ctx.reservations, ctx.keypair, ctx.node_id, start_generation);
Self {
ctx,
claimant,
cohort,
replica_set,
reachable,
job,
}
}
fn next_gen(&mut self) -> u64 {
self.claimant.next_gen()
}
}
impl ClaimPipeline for GangClaimPipeline<'_> {
type Error = ClaimError;
fn claim(&mut self, req: &CapabilityRequirement) -> Result<ClaimResult, ClaimError> {
let islands = match_islands(
self.ctx.capability,
self.ctx.topology,
&req.criteria,
&std::collections::HashSet::new(),
);
if islands.is_empty() {
return Ok(ClaimResult::Rejected);
}
let until = current_timestamp_micros().saturating_add(req.reserve_ttl_us);
let mut reserved = None;
for island in islands {
let gen = self.next_gen();
if single_island_claim(
self.ctx.reservations,
self.ctx.keypair,
self.ctx.node_id,
gen,
island,
until,
)? == ClaimOutcome::Won
{
reserved = Some(island);
break;
}
}
let Some(island) = reserved else {
return Ok(ClaimResult::Rejected);
};
let epoch: Epoch = self.next_gen();
match commit_active(
&self.claimant,
&mut self.cohort,
&self.replica_set,
&self.reachable,
island,
self.job,
epoch,
)? {
ActiveCommitOutcome::Committed => Ok(ClaimResult::Active(ActiveClaim { island })),
ActiveCommitOutcome::NoQuorum { .. } | ActiveCommitOutcome::LostReservation => {
let gen = self.next_gen();
let _ = release_island(
self.ctx.reservations,
self.ctx.keypair,
self.ctx.node_id,
gen,
island,
);
Ok(ClaimResult::Rejected)
}
}
}
fn release(&mut self, claim: &ActiveClaim) -> Result<(), ClaimError> {
let gen = self.next_gen();
release_island(
self.ctx.reservations,
self.ctx.keypair,
self.ctx.node_id,
gen,
claim.island,
)?;
Ok(())
}
}
pub fn release_step<P: ClaimPipeline>(
pipeline: &mut P,
claim: &ActiveClaim,
) -> Result<(), P::Error> {
pipeline.release(claim)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum StepKind {
#[default]
Pure,
Idempotent,
SideEffecting,
}
impl StepKind {
pub fn may_reexecute(self, already_completed: bool) -> bool {
match self {
StepKind::Pure | StepKind::Idempotent => true,
StepKind::SideEffecting => !already_completed,
}
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use std::time::Duration;
use super::*;
use crate::adapter::net::behavior::fold::{
CapabilityFilter, CapabilityMembership, CapabilityQuery, EnvelopeMeta, FoldKind,
IslandRecord, NodeState, ReservationQuery, ReservationState, SignedAnnouncement, UnitSet,
};
use crate::adapter::net::behavior::gang::{NumericFilter, SelectionPolicy};
use crate::adapter::net::cortex::workflow::TaskStatus;
use crate::adapter::net::redex::Redex;
struct ForcedPipeline {
result: ClaimResult,
calls: u32,
releases: u32,
}
impl ClaimPipeline for ForcedPipeline {
type Error = std::convert::Infallible;
fn claim(&mut self, _req: &CapabilityRequirement) -> Result<ClaimResult, Self::Error> {
self.calls += 1;
Ok(self.result)
}
fn release(&mut self, _claim: &ActiveClaim) -> Result<(), Self::Error> {
self.releases += 1;
Ok(())
}
}
fn requirement() -> CapabilityRequirement {
CapabilityRequirement {
criteria: MatchCriteria {
capability: CapabilityQuery::Composite(CapabilityFilter {
tags_all: vec!["gpu:h100".into()],
..Default::default()
}),
numeric: NumericFilter {
min_units: 8,
..Default::default()
},
selection: SelectionPolicy::LeastLoaded,
prefer_capability: None,
},
reserve_ttl_us: 60_000_000,
}
}
async fn submitted_task(wf: &WorkflowAdapter, id: TaskId) {
let seq = wf.submit(id).unwrap();
wf.wait_for_seq(seq).await.unwrap();
}
#[tokio::test]
async fn forced_reject_leaves_the_step_waiting_never_running() {
let redex = Redex::new();
let wf = WorkflowAdapter::open(&redex, 0x0F10_00D1).await.unwrap();
submitted_task(&wf, 1).await;
let mut pipeline = ForcedPipeline {
result: ClaimResult::Rejected,
calls: 0,
releases: 0,
};
let gate = drive_capability_step(&wf, &mut pipeline, 1, &requirement()).unwrap();
let seq = wf.wait(1).unwrap(); wf.wait_for_seq(seq).await.unwrap();
assert_eq!(gate, StepGate::Waiting);
assert_eq!(
pipeline.calls, 1,
"the requirement is handed to the pipeline"
);
assert_eq!(wf.get(1).unwrap().status, TaskStatus::Waiting);
}
#[tokio::test]
async fn forced_active_runs_the_step() {
let redex = Redex::new();
let wf = WorkflowAdapter::open(&redex, 0x0F10_00D2).await.unwrap();
submitted_task(&wf, 1).await;
let mut pipeline = ForcedPipeline {
result: ClaimResult::Active(ActiveClaim { island: 0xA0 }),
calls: 0,
releases: 0,
};
let gate = drive_capability_step(&wf, &mut pipeline, 1, &requirement()).unwrap();
let seq = wf.start(1).unwrap();
wf.wait_for_seq(seq).await.unwrap();
assert_eq!(gate, StepGate::Running(ActiveClaim { island: 0xA0 }));
assert_eq!(wf.get(1).unwrap().status, TaskStatus::Running);
if let StepGate::Running(claim) = gate {
release_step(&mut pipeline, &claim).unwrap();
wf.fail(1).unwrap();
}
assert_eq!(
pipeline.releases, 1,
"the held claim is released on abnormal exit"
);
}
#[test]
fn step_kind_reexecute_is_advisory_and_blocks_completed_side_effects() {
assert!(StepKind::Pure.may_reexecute(true));
assert!(StepKind::Idempotent.may_reexecute(true));
assert!(StepKind::SideEffecting.may_reexecute(false));
assert!(!StepKind::SideEffecting.may_reexecute(true));
assert_eq!(StepKind::default(), StepKind::Pure);
}
fn new_fold<K: FoldKind>() -> Fold<K> {
Fold::with_sweep_interval(Duration::ZERO)
}
fn announce_capability(fold: &Fold<CapabilityFold>, kp: &EntityKeypair, node: u64) {
let m = CapabilityMembership {
class_hash: 0x67_70_75,
tags: vec!["gpu:h100".into()],
hardware: None,
state: NodeState::Idle,
region: None,
price_quote: None,
reflex_addr: None,
allowed_nodes: Vec::new(),
allowed_subnets: Vec::new(),
allowed_groups: Vec::new(),
metadata: BTreeMap::new(),
owner: None,
};
fold.apply(
SignedAnnouncement::sign(
kp,
CapabilityFold::KIND_ID,
m.class_hash,
node,
1,
EnvelopeMeta::default(),
m,
)
.unwrap(),
)
.unwrap();
}
fn announce_island(fold: &Fold<IslandTopologyFold>, kp: &EntityKeypair, node: u64, id: u64) {
let record = IslandRecord {
id,
units: UnitSet::new((0..8).collect()),
host: node,
capabilities: vec!["model:a1".into()],
load: 0.2,
p50_latency_us: 1_000,
};
fold.apply(
SignedAnnouncement::sign(
kp,
IslandTopologyFold::KIND_ID,
0,
node,
1,
EnvelopeMeta::default(),
record,
)
.unwrap(),
)
.unwrap();
}
#[tokio::test]
async fn gang_pipeline_claims_active_and_runs_when_capacity_exists() {
let caps = new_fold::<CapabilityFold>();
let topo = new_fold::<IslandTopologyFold>();
let res = new_fold::<ReservationFold>();
let gpu = EntityKeypair::generate();
let gn = gpu.entity_id().node_id();
announce_capability(&caps, &gpu, gn);
announce_island(&topo, &gpu, gn, 0xA0);
let leader = EntityKeypair::generate();
let ln = leader.entity_id().node_id();
let mut pipeline = GangClaimPipeline::new(
GangClaimContext {
capability: &caps,
topology: &topo,
reservations: &res,
keypair: &leader,
node_id: ln,
},
ReplicaSet::new([1, 2, 3]),
vec![1, 2, 3], 42,
);
let redex = Redex::new();
let wf = WorkflowAdapter::open(&redex, 0x0F10_00D3).await.unwrap();
submitted_task(&wf, 1).await;
let gate = drive_capability_step(&wf, &mut pipeline, 1, &requirement()).unwrap();
let seq = wf.start(1).unwrap();
wf.wait_for_seq(seq).await.unwrap();
assert_eq!(gate, StepGate::Running(ActiveClaim { island: 0xA0 }));
assert_eq!(wf.get(1).unwrap().status, TaskStatus::Running);
assert!(matches!(
res.query(ReservationQuery::State(0xA0))[0].1,
ReservationState::Active { holder, .. } if holder == ln
));
}
#[tokio::test]
async fn gang_pipeline_release_returns_the_island_to_free() {
let caps = new_fold::<CapabilityFold>();
let topo = new_fold::<IslandTopologyFold>();
let res = new_fold::<ReservationFold>();
let gpu = EntityKeypair::generate();
let gn = gpu.entity_id().node_id();
announce_capability(&caps, &gpu, gn);
announce_island(&topo, &gpu, gn, 0xA0);
let leader = EntityKeypair::generate();
let ln = leader.entity_id().node_id();
let mut pipeline = GangClaimPipeline::new(
GangClaimContext {
capability: &caps,
topology: &topo,
reservations: &res,
keypair: &leader,
node_id: ln,
},
ReplicaSet::new([1, 2, 3]),
vec![1, 2, 3],
42,
);
let redex = Redex::new();
let wf = WorkflowAdapter::open(&redex, 0x0F10_00D6).await.unwrap();
submitted_task(&wf, 1).await;
let gate = drive_capability_step(&wf, &mut pipeline, 1, &requirement()).unwrap();
let claim = match gate {
StepGate::Running(c) => c,
StepGate::Waiting => panic!("expected the claim to commit Active"),
};
assert!(matches!(
res.query(ReservationQuery::State(0xA0))[0].1,
ReservationState::Active { .. }
));
release_step(&mut pipeline, &claim).unwrap();
assert_eq!(
res.query(ReservationQuery::State(0xA0))[0].1,
ReservationState::Free,
"released island returns to the pool",
);
}
#[tokio::test]
async fn gang_pipeline_rejects_and_waits_with_no_capacity_leaving_nothing_reserved() {
let caps = new_fold::<CapabilityFold>();
let topo = new_fold::<IslandTopologyFold>();
let res = new_fold::<ReservationFold>();
let gpu = EntityKeypair::generate();
let gn = gpu.entity_id().node_id();
announce_capability(&caps, &gpu, gn);
let leader = EntityKeypair::generate();
let ln = leader.entity_id().node_id();
let mut pipeline = GangClaimPipeline::new(
GangClaimContext {
capability: &caps,
topology: &topo,
reservations: &res,
keypair: &leader,
node_id: ln,
},
ReplicaSet::new([1, 2, 3]),
vec![1, 2, 3],
42,
);
let redex = Redex::new();
let wf = WorkflowAdapter::open(&redex, 0x0F10_00D4).await.unwrap();
submitted_task(&wf, 1).await;
let gate = drive_capability_step(&wf, &mut pipeline, 1, &requirement()).unwrap();
let seq = wf.wait(1).unwrap();
wf.wait_for_seq(seq).await.unwrap();
assert_eq!(gate, StepGate::Waiting);
assert_eq!(wf.get(1).unwrap().status, TaskStatus::Waiting);
assert!(res.query(ReservationQuery::State(0xA0)).is_empty());
}
#[tokio::test]
async fn gang_pipeline_minority_partition_cannot_run_the_step() {
let caps = new_fold::<CapabilityFold>();
let topo = new_fold::<IslandTopologyFold>();
let res = new_fold::<ReservationFold>();
let gpu = EntityKeypair::generate();
let gn = gpu.entity_id().node_id();
announce_capability(&caps, &gpu, gn);
announce_island(&topo, &gpu, gn, 0xA0);
let leader = EntityKeypair::generate();
let ln = leader.entity_id().node_id();
let mut pipeline = GangClaimPipeline::new(
GangClaimContext {
capability: &caps,
topology: &topo,
reservations: &res,
keypair: &leader,
node_id: ln,
},
ReplicaSet::new([1, 2, 3, 4, 5]),
vec![1, 2], 42,
);
let redex = Redex::new();
let wf = WorkflowAdapter::open(&redex, 0x0F10_00D5).await.unwrap();
submitted_task(&wf, 1).await;
let gate = drive_capability_step(&wf, &mut pipeline, 1, &requirement()).unwrap();
assert_eq!(
gate,
StepGate::Waiting,
"minority side can't reach Active → step waits"
);
let state = res.query(ReservationQuery::State(0xA0));
assert!(
state.is_empty() || matches!(state[0].1, ReservationState::Free),
"minority reserve released (Free), never left Reserved or leaked Active: {state:?}",
);
}
}