use super::*;
use crate::collectives::Collectives;
use crate::engine::partitioned::topology::{allowed_cpus, Placement};
fn config(partitions: usize) -> PartitionConfig {
PartitionConfig {
placement: Placement::Unpinned { partitions },
bind_memory: false,
threads_per_partition: Some(2),
}
}
#[test]
fn partition_bits_is_log2_of_the_count() {
for (p, bits) in [(1usize, 0u8), (2, 1), (4, 2), (8, 3)] {
let runtime = PartitionRuntime::new(&config(p)).expect("resolve");
assert_eq!(runtime.num_partitions(), p);
assert_eq!(runtime.partition_bits(), bits);
}
}
#[test]
fn threads_per_partition_overrides_the_resolved_width() {
let cpus = allowed_cpus();
let placement = Placement::Explicit(vec![cpus.clone(), cpus.clone()]);
let config = PartitionConfig { placement, bind_memory: false, threads_per_partition: None };
let resolved = PartitionRuntime::new(&config).expect("resolve");
for slot in &resolved.slots {
assert_eq!(slot.threads, cpus.len());
}
let narrowed =
PartitionRuntime::new(&PartitionConfig { threads_per_partition: Some(1), ..config })
.expect("resolve with an explicit width");
for slot in &narrowed.slots {
assert_eq!(slot.threads, 1);
assert_eq!(slot.cpus.as_ref(), Some(&cpus), "the CPU set is untouched");
}
}
#[test]
fn map_partitions_returns_results_in_rank_order() {
let runtime = PartitionRuntime::new(&config(4)).expect("resolve");
let got = runtime.map_partitions(vec![10usize, 20, 30, 40], |rank, item, transport| {
assert_eq!(transport.rank() as usize, rank);
assert_eq!(transport.size(), 4);
assert!(rayon::current_thread_index().is_some());
(rank, item)
});
assert_eq!(got, vec![(0, 10), (1, 20), (2, 30), (3, 40)]);
}
#[test]
fn map_partitions_runs_collectives_between_partitions() {
let runtime = PartitionRuntime::new(&config(4)).expect("resolve");
let got = runtime
.map_partitions(vec![3u8, 7, 1, 5], |_, item, transport| transport.allreduce_max_u8(item));
assert_eq!(got, vec![7, 7, 7, 7]);
}
#[test]
#[should_panic]
fn a_panicking_partition_does_not_hang_the_group() {
let runtime = PartitionRuntime::new(&config(4)).expect("resolve");
runtime.map_partitions(vec![0usize; 4], |rank, _, transport| {
if rank == 2 {
panic!("rank 2 fell over");
}
for _ in 0..4 {
transport.allreduce_max_u8(rank as u8);
}
});
}
#[test]
fn placement_summary_names_every_slot() {
let runtime = PartitionRuntime::new(&config(2)).expect("resolve");
let summary = runtime.placement_summary();
assert!(summary.contains("0:"), "{summary}");
assert!(summary.contains("1:"), "{summary}");
}