pub mod active;
pub mod claim;
pub mod contention;
pub mod filter;
pub mod multi;
#[cfg(feature = "redex")]
pub mod placement;
pub mod quorum;
pub mod schedule;
#[cfg(test)]
mod proptest;
pub use active::{commit_active, ActiveCommitOutcome, ReplicaCohort};
pub use claim::{
activate_announcement, activate_island, release_announcement, release_island,
reserve_announcement, single_island_claim, ClaimError, ClaimOutcome, Claimant,
};
pub use contention::claim_first_available;
pub use filter::{
candidate_hosts, numeric_filter, select_islands, select_with_affinity, NumericFilter,
SelectionPolicy,
};
pub use multi::{acquire_gang, try_acquire_gang, AcquireAttempt, GangClaim, GangOutcome};
#[cfg(feature = "redex")]
pub use placement::{colocated_island_config, pinned_island_replicas, COLOCATE_WITH_STRICT_KEY};
pub use quorum::{Epoch, FenceLedger, QuorumWitness, ReplicaSet};
pub use schedule::{
schedule_gang, schedule_single, GangRequest, GangScheduler, ScheduleError, Scheduled,
};
use std::collections::HashSet;
use crate::adapter::net::behavior::fold::{
CapabilityFold, CapabilityQuery, Fold, IslandId, IslandQuery, IslandRecord, IslandTopologyFold,
NodeId,
};
#[derive(Debug, Clone)]
pub struct MatchCriteria {
pub capability: CapabilityQuery,
pub numeric: NumericFilter,
pub selection: SelectionPolicy,
pub prefer_capability: Option<String>,
}
pub fn match_islands(
capability_fold: &Fold<CapabilityFold>,
topology_fold: &Fold<IslandTopologyFold>,
criteria: &MatchCriteria,
down_nodes: &HashSet<NodeId>,
) -> Vec<IslandId> {
let matches = capability_fold.query(criteria.capability.clone());
let mut hosts = candidate_hosts(&matches);
if !down_nodes.is_empty() {
hosts.retain(|host| !down_nodes.contains(host));
}
if hosts.is_empty() {
return Vec::new();
}
let candidates: Vec<IslandRecord> = topology_fold
.query(IslandQuery::HostedByAny(hosts))
.into_iter()
.map(|(_, record)| record)
.filter(|record| criteria.numeric.accepts(record))
.collect();
select_with_affinity(
candidates,
criteria.selection,
criteria.prefer_capability.clone(),
)
}
pub fn match_islands_sensed(
capability_fold: &Fold<CapabilityFold>,
topology_fold: &Fold<IslandTopologyFold>,
criteria: &MatchCriteria,
down_nodes: &HashSet<NodeId>,
sensed_non_viable: &HashSet<NodeId>,
sensed_viable_order: &[NodeId],
) -> Vec<IslandId> {
let pruned: HashSet<NodeId> = if sensed_non_viable.is_empty() {
down_nodes.clone()
} else {
down_nodes.union(sensed_non_viable).copied().collect()
};
let mut ordered = match_islands(capability_fold, topology_fold, criteria, &pruned);
if sensed_viable_order.is_empty() || ordered.len() < 2 {
return ordered;
}
let snapshot: std::collections::HashMap<IslandId, NodeId> = topology_fold
.query(IslandQuery::All)
.into_iter()
.map(|(island, record)| (island, record.host))
.collect();
let bands: std::collections::HashMap<IslandId, usize> = ordered
.iter()
.map(|island| {
let band = snapshot
.get(island)
.and_then(|host| {
sensed_viable_order
.iter()
.position(|provider| provider == host)
})
.unwrap_or(usize::MAX);
(*island, band)
})
.collect();
ordered.sort_by_key(|island| bands.get(island).copied().unwrap_or(usize::MAX));
ordered
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use std::time::Duration;
use super::*;
use crate::adapter::net::behavior::fold::{
CapabilityFilter, CapabilityMembership, EnvelopeMeta, Fold, FoldKind, IslandRecord,
IslandTopologyFold, NodeState, ReservationFold, ReservationQuery, ReservationState,
SignedAnnouncement, UnitSet,
};
use crate::adapter::net::current_timestamp_micros;
use crate::adapter::net::identity::EntityKeypair;
fn announce_capability(
fold: &Fold<CapabilityFold>,
kp: &EntityKeypair,
node: u64,
tags: Vec<String>,
) {
announce_capability_in(fold, kp, node, tags, None);
}
fn announce_capability_in(
fold: &Fold<CapabilityFold>,
kp: &EntityKeypair,
node: u64,
tags: Vec<String>,
region: Option<String>,
) {
let membership = CapabilityMembership {
class_hash: 0x67_70_75, tags,
hardware: None,
state: NodeState::Idle,
region,
price_quote: None,
reflex_addr: None,
allowed_nodes: Vec::new(),
allowed_subnets: Vec::new(),
allowed_groups: Vec::new(),
metadata: BTreeMap::new(),
};
let ann = SignedAnnouncement::sign(
kp,
CapabilityFold::KIND_ID,
membership.class_hash,
node,
1,
EnvelopeMeta::default(),
membership,
)
.expect("sign cap");
fold.apply(ann).expect("apply cap");
}
fn announce_island(
fold: &Fold<IslandTopologyFold>,
kp: &EntityKeypair,
node: u64,
id: IslandId,
units: usize,
load: f32,
) {
let record = IslandRecord {
id,
units: UnitSet::new((0..units as u32).collect()),
host: node,
capabilities: vec!["model:a1".into()],
load,
p50_latency_us: 1_500,
};
let ann = SignedAnnouncement::sign(
kp,
IslandTopologyFold::KIND_ID,
0,
node,
1,
EnvelopeMeta::default(),
record,
)
.expect("sign island");
fold.apply(ann).expect("apply island");
}
fn new_fold<K: crate::adapter::net::behavior::fold::FoldKind>() -> Fold<K> {
Fold::with_sweep_interval(Duration::ZERO)
}
#[test]
fn match_islands_narrows_by_capability_then_numeric_then_orders() {
let caps: Fold<CapabilityFold> = new_fold();
let topo: Fold<IslandTopologyFold> = new_fold();
let kp_a = EntityKeypair::generate();
let kp_b = EntityKeypair::generate();
let kp_c = EntityKeypair::generate();
let (na, nb, nc) = (
kp_a.entity_id().node_id(),
kp_b.entity_id().node_id(),
kp_c.entity_id().node_id(),
);
announce_capability(&caps, &kp_a, na, vec!["gpu:h100".into()]);
announce_capability(&caps, &kp_b, nb, vec!["gpu:h100".into()]);
announce_capability(&caps, &kp_c, nc, vec!["gpu:a10".into()]);
announce_island(&topo, &kp_a, na, 0xA0, 8, 0.6);
announce_island(&topo, &kp_a, na, 0xA5, 8, 0.2);
announce_island(&topo, &kp_b, nb, 0xB0, 8, 0.4);
announce_island(&topo, &kp_c, nc, 0xC0, 8, 0.0);
let criteria = MatchCriteria {
capability: CapabilityQuery::Composite(CapabilityFilter {
tags_all: vec!["gpu:h100".into()],
..Default::default()
}),
numeric: NumericFilter {
min_units: 8,
max_load: Some(0.5),
..Default::default()
},
selection: SelectionPolicy::LeastLoaded,
prefer_capability: None,
};
let order = match_islands(&caps, &topo, &criteria, &HashSet::new());
assert_eq!(order, vec![0xA5, 0xB0]);
}
#[test]
fn sensed_match_prunes_non_viable_and_ranks_viable_first() {
let caps: Fold<CapabilityFold> = new_fold();
let topo: Fold<IslandTopologyFold> = new_fold();
let kp_a = EntityKeypair::generate();
let kp_b = EntityKeypair::generate();
let kp_c = EntityKeypair::generate();
let (na, nb, nc) = (
kp_a.entity_id().node_id(),
kp_b.entity_id().node_id(),
kp_c.entity_id().node_id(),
);
for (kp, node) in [(&kp_a, na), (&kp_b, nb), (&kp_c, nc)] {
announce_capability(&caps, kp, node, vec!["gpu:h100".into()]);
}
announce_island(&topo, &kp_a, na, 0xA0, 8, 0.1);
announce_island(&topo, &kp_b, nb, 0xB0, 8, 0.2);
announce_island(&topo, &kp_c, nc, 0xC0, 8, 0.3);
let 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,
};
assert_eq!(
match_islands(&caps, &topo, &criteria, &HashSet::new()),
vec![0xA0, 0xB0, 0xC0],
);
let non_viable: HashSet<NodeId> = [nb].into_iter().collect();
let order = match_islands_sensed(
&caps,
&topo,
&criteria,
&HashSet::new(),
&non_viable,
&[nc, na],
);
assert_eq!(
order,
vec![0xC0, 0xA0],
"NotReady host pruned; sensed rank leads the claim order",
);
assert_eq!(
match_islands(&caps, &topo, &criteria, &HashSet::new()),
vec![0xA0, 0xB0, 0xC0],
"one interest's NotReady never suspends the entry",
);
}
#[test]
fn sensed_match_with_empty_delta_is_identical_and_potential_is_never_pruned() {
let caps: Fold<CapabilityFold> = new_fold();
let topo: Fold<IslandTopologyFold> = new_fold();
let kp_a = EntityKeypair::generate();
let kp_b = EntityKeypair::generate();
let (na, nb) = (kp_a.entity_id().node_id(), kp_b.entity_id().node_id());
announce_capability(&caps, &kp_a, na, vec!["gpu:h100".into()]);
announce_capability(&caps, &kp_b, nb, vec!["gpu:h100".into()]);
announce_island(&topo, &kp_a, na, 0xA0, 8, 0.1);
announce_island(&topo, &kp_b, nb, 0xB0, 8, 0.2);
let 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,
};
let plain = match_islands(&caps, &topo, &criteria, &HashSet::new());
assert_eq!(
match_islands_sensed(
&caps,
&topo,
&criteria,
&HashSet::new(),
&HashSet::new(),
&[],
),
plain,
"empty sensed delta ⇒ byte-identical to match_islands",
);
assert_eq!(
match_islands_sensed(
&caps,
&topo,
&criteria,
&HashSet::new(),
&HashSet::new(),
&[nb],
),
vec![0xB0, 0xA0],
"potential hosts trail the viable band but are never pruned",
);
}
#[test]
fn match_islands_empty_when_no_capability_match() {
let caps: Fold<CapabilityFold> = new_fold();
let topo: Fold<IslandTopologyFold> = new_fold();
let kp = EntityKeypair::generate();
let n = kp.entity_id().node_id();
announce_capability(&caps, &kp, n, vec!["gpu:a10".into()]);
announce_island(&topo, &kp, n, 0xA0, 8, 0.1);
let criteria = MatchCriteria {
capability: CapabilityQuery::Composite(CapabilityFilter {
tags_all: vec!["gpu:h100".into()],
..Default::default()
}),
numeric: NumericFilter::default(),
selection: SelectionPolicy::LeastLoaded,
prefer_capability: None,
};
assert!(match_islands(&caps, &topo, &criteria, &HashSet::new()).is_empty());
}
#[test]
fn dead_host_islands_are_pruned_from_matching() {
let caps: Fold<CapabilityFold> = new_fold();
let topo: Fold<IslandTopologyFold> = new_fold();
let kp_a = EntityKeypair::generate();
let kp_b = EntityKeypair::generate();
let na = kp_a.entity_id().node_id();
let nb = kp_b.entity_id().node_id();
announce_capability(&caps, &kp_a, na, vec!["gpu:h100".into()]);
announce_capability(&caps, &kp_b, nb, vec!["gpu:h100".into()]);
announce_island(&topo, &kp_a, na, 0xA0, 8, 0.1);
announce_island(&topo, &kp_b, nb, 0xB0, 8, 0.2);
let criteria = MatchCriteria {
capability: CapabilityQuery::Composite(CapabilityFilter {
tags_all: vec!["gpu:h100".into()],
..Default::default()
}),
numeric: NumericFilter::default(),
selection: SelectionPolicy::LeastLoaded,
prefer_capability: None,
};
assert_eq!(
match_islands(&caps, &topo, &criteria, &HashSet::new()),
vec![0xA0, 0xB0],
);
let a_down: HashSet<NodeId> = [na].into_iter().collect();
assert_eq!(match_islands(&caps, &topo, &criteria, &a_down), vec![0xB0]);
let both_down: HashSet<NodeId> = [na, nb].into_iter().collect();
assert!(match_islands(&caps, &topo, &criteria, &both_down).is_empty());
}
#[test]
fn region_filters_at_the_host_stage_not_the_island() {
let caps: Fold<CapabilityFold> = new_fold();
let topo: Fold<IslandTopologyFold> = new_fold();
let kp_east = EntityKeypair::generate();
let kp_west = EntityKeypair::generate();
let ne = kp_east.entity_id().node_id();
let nw = kp_west.entity_id().node_id();
announce_capability_in(
&caps,
&kp_east,
ne,
vec!["gpu:h100".into()],
Some("us-east".into()),
);
announce_capability_in(
&caps,
&kp_west,
nw,
vec!["gpu:h100".into()],
Some("us-west".into()),
);
announce_island(&topo, &kp_east, ne, 0xE0, 8, 0.1);
announce_island(&topo, &kp_west, nw, 0xF0, 8, 0.1);
let east_only = MatchCriteria {
capability: CapabilityQuery::Composite(CapabilityFilter {
tags_all: vec!["gpu:h100".into()],
region: Some("us-east".into()),
..Default::default()
}),
numeric: NumericFilter::default(),
selection: SelectionPolicy::LeastLoaded,
prefer_capability: None,
};
assert_eq!(
match_islands(&caps, &topo, &east_only, &HashSet::new()),
vec![0xE0]
);
let any_region = MatchCriteria {
capability: CapabilityQuery::Composite(CapabilityFilter {
tags_all: vec!["gpu:h100".into()],
..Default::default()
}),
..east_only.clone()
};
let mut both = match_islands(&caps, &topo, &any_region, &HashSet::new());
both.sort_unstable();
assert_eq!(both, vec![0xE0, 0xF0]);
let nowhere = MatchCriteria {
capability: CapabilityQuery::Composite(CapabilityFilter {
tags_all: vec!["gpu:h100".into()],
region: Some("ap-south".into()),
..Default::default()
}),
..east_only.clone()
};
assert!(match_islands(&caps, &topo, &nowhere, &HashSet::new()).is_empty());
}
#[test]
fn pipeline_then_claim_run_release() {
let caps: Fold<CapabilityFold> = new_fold();
let topo: Fold<IslandTopologyFold> = new_fold();
let reservations: Fold<ReservationFold> = new_fold();
let kp = EntityKeypair::generate();
let node = kp.entity_id().node_id();
announce_capability(&caps, &kp, node, vec!["gpu:h100".into()]);
announce_island(&topo, &kp, node, 0xA0, 8, 0.3);
let 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,
};
let order = match_islands(&caps, &topo, &criteria, &HashSet::new());
let island = *order.first().expect("a candidate island");
assert_eq!(island, 0xA0);
let deadline = current_timestamp_micros() + 60_000_000;
assert_eq!(
single_island_claim(&reservations, &kp, node, 1, island, deadline).unwrap(),
ClaimOutcome::Won,
);
assert_eq!(
activate_island(&reservations, &kp, node, 2, island, 0x42).unwrap(),
ClaimOutcome::Won,
);
assert!(matches!(
reservations.query(ReservationQuery::State(island))[0].1,
ReservationState::Active { job_id: 0x42, .. }
));
assert_eq!(
release_island(&reservations, &kp, node, 3, island).unwrap(),
ClaimOutcome::Won,
);
assert_eq!(
reservations.query(ReservationQuery::State(island))[0].1,
ReservationState::Free,
);
}
}