paulistrings/engine/partitioned/
runtime.rs

1//! [`PartitionRuntime`]: the resolved placement, one pinned pool per partition, and the scoped fan-out every in-process partitioned call runs inside.
2//! A layer runs inside `ThreadPool::install`, so transport calls come from a pool worker, which is why the MPI transport needs `MPI_THREAD_SERIALIZED`.
3
4use std::sync::Arc;
5
6use super::in_process::{InProcessTransport, DEFAULT_WAIT_TIMEOUT};
7use super::topology::{
8    bind_current_thread_memory, build_pool, pin_current_thread, PartitionConfig, PartitionSlot,
9    TopologyError,
10};
11use super::PLACEMENT_LOG_TARGET as LOG_TARGET;
12
13/// The resolved placement and its pinned pools, built once and shared behind an [`Arc`] by every [`PartitionedSum`](crate::PartitionedSum) that runs on it.
14pub struct PartitionRuntime {
15    /// In rank order.
16    slots: Vec<PartitionSlot>,
17    /// In rank order.
18    pools: Vec<rayon::ThreadPool>,
19    bind_memory: bool,
20    /// How long a partition waits for a partner's collective before declaring it dead.
21    wait_timeout: std::time::Duration,
22}
23
24impl PartitionRuntime {
25    /// Resolves `config` against the machine and builds one pinned pool per partition.
26    ///
27    /// # Errors
28    ///
29    /// Whatever [`PartitionConfig::num_partitions`] reports (a partition count that is not a power of two, an empty or out-of-mask explicit CPU set), or [`TopologyError::Io`] if a pool cannot be built.
30    pub fn new(config: &PartitionConfig) -> Result<Arc<Self>, TopologyError> {
31        let slots = config.resolve()?;
32        let mut pools = Vec::with_capacity(slots.len());
33        for (rank, slot) in slots.iter().enumerate() {
34            pools.push(build_pool(slot, config.bind_memory, &format!("ps-part{rank}"))?);
35        }
36        // A device group waits for a partner's collective as long as for its wire, `PAULISTRINGS_WIRE_TIMEOUT_S`.
37        #[cfg(feature = "cuda")]
38        let wait_timeout = if slots.iter().any(|s| s.device.is_some()) {
39            crate::engine::gpu::wire_timeout()
40        } else {
41            DEFAULT_WAIT_TIMEOUT
42        };
43        #[cfg(not(feature = "cuda"))]
44        let wait_timeout = DEFAULT_WAIT_TIMEOUT;
45        Ok(Arc::new(Self { slots, pools, bind_memory: config.bind_memory, wait_timeout }))
46    }
47
48    /// Partitions in the group — always a power of two.
49    pub fn num_partitions(&self) -> usize {
50        self.slots.len()
51    }
52
53    /// The resolved slots, in rank order.
54    #[cfg(feature = "cuda")]
55    pub(crate) fn slots(&self) -> &[PartitionSlot] {
56        &self.slots
57    }
58
59    /// `log2` of the partition count.
60    pub(crate) fn partition_bits(&self) -> u8 {
61        debug_assert!(self.slots.len().is_power_of_two());
62        self.slots.len().trailing_zeros() as u8
63    }
64
65    /// A one-line summary of the placement, for the entry log line.
66    pub(crate) fn placement_summary(&self) -> String {
67        let mut summary = String::new();
68        for (rank, slot) in self.slots.iter().enumerate() {
69            if rank > 0 {
70                summary.push_str(", ");
71            }
72            match (&slot.cpus, slot.device) {
73                (_, Some(device)) => {
74                    summary.push_str(&format!("{rank}:gpu{device} x{}", slot.threads))
75                }
76                (Some(cpus), None) => summary.push_str(&format!("{rank}:{cpus}x{}", slot.threads)),
77                (None, None) => summary.push_str(&format!("{rank}:unpinned x{}", slot.threads)),
78            }
79        }
80        summary
81    }
82
83    /// Runs `f` on partition 0's pool: the distributed shape of [`map_partitions`](Self::map_partitions).
84    pub(crate) fn install<R: Send>(&self, f: impl FnOnce() -> R + Send) -> R {
85        self.pools[0].install(f)
86    }
87
88    /// Runs `f(rank, item, transport)` once per partition inside its own pool, concurrently, and returns the results in rank order; a partition's panic is re-raised here.
89    pub(crate) fn map_partitions<I, O, F>(&self, items: Vec<I>, f: F) -> Vec<O>
90    where
91        I: Send,
92        O: Send,
93        F: Fn(usize, I, &InProcessTransport) -> O + Send + Sync,
94    {
95        let size = self.num_partitions();
96        assert_eq!(items.len(), size, "map_partitions: one item per partition");
97        // A fresh group per call, each partition owning its endpoint, so a partner's panic drops its senders and surfaces as a panic instead of a hang.
98        let transports = InProcessTransport::group_with_timeout(size as u32, self.wait_timeout);
99        let f = &f;
100
101        let mut items = items.into_iter();
102        let mut transports = transports.into_iter();
103        let item0 = items.next().expect("at least one partition");
104        let transport0 = transports.next().expect("at least one transport");
105
106        std::thread::scope(|scope| {
107            let handles: Vec<_> = items
108                .zip(transports)
109                .enumerate()
110                .map(|(i, (item, transport))| {
111                    let rank = i + 1;
112                    let slot = &self.slots[rank];
113                    let pool = &self.pools[rank];
114                    let bind_memory = self.bind_memory;
115                    scope.spawn(move || {
116                        place_current_thread(slot, bind_memory);
117                        pool.install(move || f(rank, item, &transport))
118                    })
119                })
120                .collect();
121
122            // The calling thread is not ours to pin; the work runs on pool 0's pinned workers.
123            let mut out = Vec::with_capacity(size);
124            out.push(self.pools[0].install(|| f(0, item0, &transport0)));
125            for handle in handles {
126                match handle.join() {
127                    Ok(o) => out.push(o),
128                    Err(payload) => std::panic::resume_unwind(payload),
129                }
130            }
131            out
132        })
133    }
134}
135
136/// Pins a partition's driving thread to its slot, warning rather than failing when the platform declines.
137fn place_current_thread(slot: &PartitionSlot, bind_memory: bool) {
138    #[cfg(feature = "cuda")]
139    if let Some(device) = slot.device {
140        crate::engine::cuda_context::bind_device_context(device);
141    }
142    if let Some(cpus) = &slot.cpus {
143        if let Err(err) = pin_current_thread(cpus) {
144            log::warn!(target: LOG_TARGET, "failed to pin partition driver to {cpus}: {err}");
145        }
146    }
147    if bind_memory {
148        if let Err(err) = bind_current_thread_memory(slot.node) {
149            log::warn!(
150                target: LOG_TARGET,
151                "failed to bind partition driver memory to {:?}: {err}",
152                slot.node,
153            );
154        }
155    }
156}
157
158#[cfg(test)]
159mod tests;
paulistrings/engine/partitioned/runtime.rs

paulistrings/engine/partitioned/
runtime.rs
