paulistrings/engine/partitioned/
topology.rs

1//! CPU sets, NUMA node discovery, thread and memory pinning, and the pinned Rayon pool one partition runs on.
2//! Discovery is sysfs plus `libc`; on non-Linux targets every entry point reports one node and pins nothing.
3
4use std::fmt;
5use std::io;
6
7use super::PLACEMENT_LOG_TARGET as LOG_TARGET;
8
9/// A set of logical CPU indices; an unsorted literal is tolerated, since every consumer normalizes.
10#[derive(Clone, Debug, PartialEq, Eq)]
11pub struct CpuSet(
12    /// The CPU indices, ascending and unique.
13    pub Vec<usize>,
14);
15
16impl CpuSet {
17    /// Parses a Linux cpulist, the `"0-7,16-23"` form written by sysfs, skipping whitespace and empty tokens.
18    ///
19    /// # Errors
20    ///
21    /// [`TopologyError::EmptySet`] if the list names no CPU, and [`TopologyError::InvalidCpuList`] if a token is not a decimal index or an ascending `lo-hi` range.
22    pub fn parse(s: &str) -> Result<Self, TopologyError> {
23        let invalid = || TopologyError::InvalidCpuList(s.trim().to_string());
24        let index = |token: &str| token.trim().parse::<usize>().map_err(|_| invalid());
25
26        let mut cpus = Vec::new();
27        for token in s.split(',') {
28            let token = token.trim();
29            if token.is_empty() {
30                continue;
31            }
32            match token.split_once('-') {
33                None => cpus.push(index(token)?),
34                Some((lo, hi)) => {
35                    let (lo, hi) = (index(lo)?, index(hi)?);
36                    if hi < lo {
37                        return Err(invalid());
38                    }
39                    cpus.extend(lo..=hi);
40                }
41            }
42        }
43        if cpus.is_empty() {
44            return Err(TopologyError::EmptySet);
45        }
46        cpus.sort_unstable();
47        cpus.dedup();
48        Ok(Self(cpus))
49    }
50
51    /// The number of distinct CPUs in the set.
52    #[must_use]
53    pub fn len(&self) -> usize {
54        self.normalized().len()
55    }
56
57    /// Whether the set names no CPU.
58    #[must_use]
59    pub fn is_empty(&self) -> bool {
60        self.0.is_empty()
61    }
62
63    /// The CPUs present in both sets.
64    #[must_use]
65    pub fn intersect(&self, other: &Self) -> Self {
66        let other = other.normalized();
67        Self(self.normalized().into_iter().filter(|cpu| other.binary_search(cpu).is_ok()).collect())
68    }
69
70    fn is_subset_of(&self, other: &Self) -> bool {
71        let other = other.normalized();
72        self.0.iter().all(|cpu| other.binary_search(cpu).is_ok())
73    }
74
75    fn normalized(&self) -> Vec<usize> {
76        let mut cpus = self.0.clone();
77        cpus.sort_unstable();
78        cpus.dedup();
79        cpus
80    }
81
82    fn union(&self, other: &Self) -> Self {
83        let mut cpus = self.normalized();
84        cpus.extend(other.normalized());
85        cpus.sort_unstable();
86        cpus.dedup();
87        Self(cpus)
88    }
89}
90
91impl fmt::Display for CpuSet {
92    /// Writes the sysfs cpulist form, compressing runs into `lo-hi` ranges.
93    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
94        let cpus = self.normalized();
95        let mut first = true;
96        let mut i = 0;
97        while i < cpus.len() {
98            let mut j = i;
99            while j + 1 < cpus.len() && cpus[j + 1] == cpus[j] + 1 {
100                j += 1;
101            }
102            if !first {
103                write!(f, ",")?;
104            }
105            first = false;
106            if i == j {
107                write!(f, "{}", cpus[i])?;
108            } else {
109                write!(f, "{}-{}", cpus[i], cpus[j])?;
110            }
111            i = j + 1;
112        }
113        Ok(())
114    }
115}
116
117/// The current affinity mask, or `0..available_parallelism` where it cannot be read.
118#[must_use]
119pub(crate) fn allowed_cpus() -> CpuSet {
120    #[cfg(target_os = "linux")]
121    {
122        // SAFETY: `set` is a valid, correctly sized `cpu_set_t`; pid 0 is the calling thread.
123        unsafe {
124            let mut set: libc::cpu_set_t = std::mem::zeroed();
125            if libc::sched_getaffinity(0, std::mem::size_of::<libc::cpu_set_t>(), &mut set) == 0 {
126                let cpus: Vec<usize> = (0..libc::CPU_SETSIZE as usize)
127                    .filter(|&cpu| libc::CPU_ISSET(cpu, &set))
128                    .collect();
129                if !cpus.is_empty() {
130                    return CpuSet(cpus);
131                }
132            }
133        }
134    }
135    CpuSet((0..available_parallelism()).collect())
136}
137
138/// The NUMA nodes visible in the current affinity mask, ascending by node id, each intersected with the mask.
139///
140/// Where sysfs exposes no usable node, the whole affinity mask is reported as node 0.
141#[must_use]
142pub fn numa_nodes() -> Vec<(usize, CpuSet)> {
143    let allowed = allowed_cpus();
144    let mut nodes = sysfs_numa_nodes(&allowed);
145    nodes.sort_unstable_by_key(|(id, _)| *id);
146    if nodes.is_empty() {
147        vec![(0, allowed)]
148    } else {
149        nodes
150    }
151}
152
153#[cfg(target_os = "linux")]
154pub(crate) fn sysfs_numa_nodes(allowed: &CpuSet) -> Vec<(usize, CpuSet)> {
155    const ROOT: &str = "/sys/devices/system/node";
156    let Ok(entries) = std::fs::read_dir(ROOT) else {
157        return Vec::new();
158    };
159    let mut nodes = Vec::new();
160    for entry in entries.flatten() {
161        let name = entry.file_name();
162        let Some(id) = name.to_str().and_then(|n| n.strip_prefix("node")) else {
163            continue;
164        };
165        let Ok(id) = id.parse::<usize>() else {
166            continue;
167        };
168        let Ok(list) = std::fs::read_to_string(entry.path().join("cpulist")) else {
169            continue;
170        };
171        let Ok(cpus) = CpuSet::parse(&list) else {
172            continue;
173        };
174        let cpus = cpus.intersect(allowed);
175        if !cpus.is_empty() {
176            nodes.push((id, cpus));
177        }
178    }
179    nodes
180}
181
182#[cfg(not(target_os = "linux"))]
183pub(crate) fn sysfs_numa_nodes(_allowed: &CpuSet) -> Vec<(usize, CpuSet)> {
184    Vec::new()
185}
186
187/// Pins the calling thread to `set`; a no-op on non-Linux targets.
188pub(super) fn pin_current_thread(set: &CpuSet) -> io::Result<()> {
189    #[cfg(target_os = "linux")]
190    {
191        if set.is_empty() {
192            return Err(io::Error::new(
193                io::ErrorKind::InvalidInput,
194                "cannot pin a thread to an empty CPU set",
195            ));
196        }
197        // SAFETY: `mask` is a valid, zeroed `cpu_set_t`; every index is checked against `CPU_SETSIZE` before `CPU_SET` touches it.
198        unsafe {
199            let mut mask: libc::cpu_set_t = std::mem::zeroed();
200            for &cpu in &set.0 {
201                if cpu >= libc::CPU_SETSIZE as usize {
202                    return Err(io::Error::new(
203                        io::ErrorKind::InvalidInput,
204                        format!("CPU {cpu} is beyond CPU_SETSIZE"),
205                    ));
206                }
207                libc::CPU_SET(cpu, &mut mask);
208            }
209            if libc::sched_setaffinity(0, std::mem::size_of::<libc::cpu_set_t>(), &mask) != 0 {
210                return Err(io::Error::last_os_error());
211            }
212        }
213    }
214    #[cfg(not(target_os = "linux"))]
215    let _ = set;
216    Ok(())
217}
218
219/// Binds the calling thread's allocations to one NUMA node (`MPOL_BIND`), or restores `MPOL_DEFAULT` for `None`; a no-op on non-Linux targets.
220pub(super) fn bind_current_thread_memory(node: Option<usize>) -> io::Result<()> {
221    #[cfg(target_os = "linux")]
222    {
223        const BITS: usize = 8 * std::mem::size_of::<libc::c_ulong>();
224        let rc = match node {
225            Some(node) => {
226                let words = node / BITS + 1;
227                let mut mask = vec![0 as libc::c_ulong; words];
228                mask[node / BITS] = 1 << (node % BITS);
229                // SAFETY: `mask` outlives the call and holds `words` words, exactly the `maxnode` bits the kernel reads.
230                unsafe {
231                    libc::syscall(
232                        libc::SYS_set_mempolicy,
233                        libc::MPOL_BIND,
234                        mask.as_ptr(),
235                        (words * BITS) as libc::c_ulong,
236                    )
237                }
238            }
239            // SAFETY: `MPOL_DEFAULT` takes a null mask of zero length.
240            None => unsafe {
241                libc::syscall(
242                    libc::SYS_set_mempolicy,
243                    libc::MPOL_DEFAULT,
244                    std::ptr::null::<libc::c_ulong>(),
245                    0 as libc::c_ulong,
246                )
247            },
248        };
249        if rc != 0 {
250            return Err(io::Error::last_os_error());
251        }
252    }
253    #[cfg(not(target_os = "linux"))]
254    let _ = node;
255    Ok(())
256}
257
258/// One resolved partition: where its pool's workers run and how many there are.
259#[derive(Clone, Debug)]
260pub(crate) struct PartitionSlot {
261    /// CPUs to pin the pool's workers to, or `None` to leave them unpinned.
262    pub(crate) cpus: Option<CpuSet>,
263    /// NUMA node to bind the workers' allocations to, when the slot sits in exactly one node.
264    pub(crate) node: Option<usize>,
265    /// Worker count for the slot's Rayon pool.
266    pub(crate) threads: usize,
267    /// The CUDA device a device partition runs on, `None` for a host partition.
268    pub(crate) device: Option<u32>,
269}
270
271/// Builds one partition's Rayon pool, each worker pinning itself before Rayon's main loop and only warning when pinning fails.
272pub(super) fn build_pool(
273    slot: &PartitionSlot,
274    bind_memory: bool,
275    name: &str,
276) -> Result<rayon::ThreadPool, TopologyError> {
277    let cpus = slot.cpus.clone();
278    let node = slot.node;
279    #[cfg(feature = "cuda")]
280    let device = slot.device;
281    let pool_name = name.to_string();
282    rayon::ThreadPoolBuilder::new()
283        .num_threads(slot.threads)
284        .thread_name(move |index| format!("{pool_name}-{index}"))
285        .spawn_handler(move |thread| {
286            let cpus = cpus.clone();
287            let mut builder = std::thread::Builder::new();
288            if let Some(name) = thread.name() {
289                builder = builder.name(name.to_string());
290            }
291            if let Some(size) = thread.stack_size() {
292                builder = builder.stack_size(size);
293            }
294            builder.spawn(move || {
295                #[cfg(feature = "cuda")]
296                if let Some(device) = device {
297                    crate::engine::cuda_context::bind_device_context(device);
298                }
299                if let Some(cpus) = &cpus {
300                    if let Err(err) = pin_current_thread(cpus) {
301                        log::warn!(target: LOG_TARGET, "failed to pin worker to {cpus}: {err}");
302                    }
303                }
304                if bind_memory {
305                    if let Err(err) = bind_current_thread_memory(node) {
306                        log::warn!(target: LOG_TARGET, "failed to bind worker memory to {node:?}: {err}");
307                    }
308                }
309                thread.run();
310            })?;
311            Ok(())
312        })
313        .build()
314        .map_err(|err| TopologyError::Io(io::Error::other(err.to_string())))
315}
316
317/// How partitions map onto the machine.
318#[derive(Clone, Debug, PartialEq, Eq)]
319pub enum Placement {
320    /// One partition per NUMA node in the affinity mask.
321    ///
322    /// The partition count is the node count rounded **down** to a power of two (at least one); when that is fewer than the node count, adjacent nodes are merged into a partition rather than dropped, so every allowed CPU stays in play.
323    /// Each partition takes as many threads as it has CPUs.
324    Auto {
325        /// Upper bound on the partition count, itself rounded down to a power of two.
326        max_partitions: Option<usize>,
327    },
328    /// One partition per listed CPU set, exactly as given.
329    ///
330    /// Sets may overlap (the resolve logs a warning and proceeds), but each must be non-empty and a subset of the process's affinity mask.
331    Explicit(
332        /// The CPU sets, one per partition; the count must be a power of two.
333        Vec<CpuSet>,
334    ),
335    /// Partitions with no pinning at all, for tests and CI, each taking an even share of `available_parallelism` (at least one thread).
336    Unpinned {
337        /// Number of partitions; must be a power of two.
338        partitions: usize,
339    },
340    /// `per_device` partitions on each listed CUDA device, in rank order, for the `cuda` backend's [`GpuPartitionedSum`](crate::gpu::GpuPartitionedSum).
341    ///
342    /// `devices.len() × per_device` must be a power of two; each slot is unpinned with a small host pool for the export and receive plumbing, and every worker of that pool binds the device's context when it starts.
343    #[cfg(feature = "cuda")]
344    Devices {
345        /// Device ordinals, one entry per device.
346        devices: Vec<u32>,
347        /// Partitions sharing each device; `1` in production, more to test the exchange on one device.
348        per_device: usize,
349    },
350}
351
352/// Host workers of a device partition's pool.
353#[cfg(feature = "cuda")]
354const DEVICE_PARTITION_THREADS: usize = 4;
355
356/// Placement plus the knobs the partitioned engine reads alongside it.
357#[derive(Clone, Debug, PartialEq, Eq)]
358pub struct PartitionConfig {
359    /// How partitions map onto the machine.
360    pub placement: Placement,
361    /// Whether each pool's workers bind their allocations to the partition's NUMA node (Python's `pin_memory=`).
362    pub bind_memory: bool,
363    /// Workers in every partition's pool (at least one) on the CPU set the placement resolved, or `None` for the placement's own count.
364    ///
365    /// A measurement knob for holding the total thread count fixed across partition counts; it overrides a device partition's small host pool too.
366    pub threads_per_partition: Option<usize>,
367}
368
369impl Default for PartitionConfig {
370    /// One partition per NUMA node, memory bound to it.
371    fn default() -> Self {
372        Self {
373            placement: Placement::Auto { max_partitions: None },
374            bind_memory: true,
375            threads_per_partition: None,
376        }
377    }
378}
379
380impl PartitionConfig {
381    /// Resolves the placement against the machine into one slot per partition, a power of two of them.
382    pub(crate) fn resolve(&self) -> Result<Vec<PartitionSlot>, TopologyError> {
383        let mut slots = match &self.placement {
384            Placement::Auto { max_partitions } => resolve_auto(*max_partitions),
385            Placement::Explicit(sets) => resolve_explicit(sets)?,
386            Placement::Unpinned { partitions } => resolve_unpinned(*partitions)?,
387            #[cfg(feature = "cuda")]
388            Placement::Devices { devices, per_device } => resolve_devices(devices, *per_device)?,
389        };
390        if let Some(threads) = self.threads_per_partition {
391            for slot in &mut slots {
392                slot.threads = threads.max(1);
393            }
394        }
395        Ok(slots)
396    }
397
398    /// The partition count the placement resolves to on this machine, without building any pool.
399    ///
400    /// # Errors
401    ///
402    /// [`TopologyError::NotPowerOfTwo`] if an explicit or unpinned partition count is not a power of two, [`TopologyError::EmptySet`] if an explicit set is empty, and [`TopologyError::CpuNotAllowed`] if an explicit set names a CPU outside the affinity mask.
403    pub fn num_partitions(&self) -> Result<usize, TopologyError> {
404        Ok(self.resolve()?.len())
405    }
406}
407
408#[cfg(feature = "cuda")]
409fn resolve_devices(
410    devices: &[u32],
411    per_device: usize,
412) -> Result<Vec<PartitionSlot>, TopologyError> {
413    let count = devices.len() * per_device;
414    if count == 0 || !count.is_power_of_two() {
415        return Err(TopologyError::NotPowerOfTwo(count));
416    }
417    Ok(devices
418        .iter()
419        .flat_map(|&d| {
420            std::iter::repeat_n(
421                PartitionSlot {
422                    cpus: None,
423                    node: None,
424                    threads: DEVICE_PARTITION_THREADS,
425                    device: Some(d),
426                },
427                per_device,
428            )
429        })
430        .collect())
431}
432
433fn resolve_auto(max_partitions: Option<usize>) -> Vec<PartitionSlot> {
434    let nodes = numa_nodes();
435    let mut count = prev_power_of_two(nodes.len());
436    if let Some(max) = max_partitions {
437        count = prev_power_of_two(count.min(max));
438    }
439
440    // Merge adjacent nodes when the count was rounded down, the first `nodes.len() % count` partitions taking one node more.
441    let (base, remainder) = (nodes.len() / count, nodes.len() % count);
442    let mut slots = Vec::with_capacity(count);
443    let mut rest = nodes.as_slice();
444    for partition in 0..count {
445        let take = base + usize::from(partition < remainder);
446        let (group, tail) = rest.split_at(take);
447        rest = tail;
448        let cpus =
449            group.iter().fold(CpuSet(Vec::new()), |accumulator, (_, set)| accumulator.union(set));
450        slots.push(PartitionSlot {
451            node: (group.len() == 1).then(|| group[0].0),
452            threads: cpus.len(),
453            cpus: Some(cpus),
454            device: None,
455        });
456    }
457    slots
458}
459
460fn resolve_explicit(sets: &[CpuSet]) -> Result<Vec<PartitionSlot>, TopologyError> {
461    if !sets.len().is_power_of_two() {
462        return Err(TopologyError::NotPowerOfTwo(sets.len()));
463    }
464    let allowed = allowed_cpus();
465    let nodes = numa_nodes();
466
467    let mut slots = Vec::with_capacity(sets.len());
468    for set in sets {
469        if set.is_empty() {
470            return Err(TopologyError::EmptySet);
471        }
472        let cpus = CpuSet(set.normalized());
473        if let Some(&cpu) = cpus.0.iter().find(|cpu| !allowed.0.contains(cpu)) {
474            return Err(TopologyError::CpuNotAllowed(cpu));
475        }
476        let mut containing =
477            nodes.iter().filter(|(_, node)| cpus.is_subset_of(node)).map(|(id, _)| *id);
478        let node = match (containing.next(), containing.next()) {
479            (Some(id), None) => Some(id),
480            _ => None,
481        };
482        slots.push(PartitionSlot { node, threads: cpus.len(), cpus: Some(cpus), device: None });
483    }
484
485    for (i, slot) in slots.iter().enumerate() {
486        for other in slots.iter().skip(i + 1) {
487            let (a, b) = (slot.cpus.as_ref(), other.cpus.as_ref());
488            if let (Some(a), Some(b)) = (a, b) {
489                let shared = a.intersect(b);
490                if !shared.is_empty() {
491                    log::warn!(
492                        target: LOG_TARGET,
493                        "explicit partitions {a} and {b} overlap on {shared}; \
494                         their pools will contend for those CPUs",
495                    );
496                }
497            }
498        }
499    }
500    Ok(slots)
501}
502
503fn resolve_unpinned(partitions: usize) -> Result<Vec<PartitionSlot>, TopologyError> {
504    if !partitions.is_power_of_two() {
505        return Err(TopologyError::NotPowerOfTwo(partitions));
506    }
507    let threads = (available_parallelism() / partitions).max(1);
508    Ok(vec![PartitionSlot { cpus: None, node: None, threads, device: None }; partitions])
509}
510
511/// The largest power of two `<= n`, and 1 for `n == 0`.
512fn prev_power_of_two(n: usize) -> usize {
513    1 << n.max(1).ilog2()
514}
515
516fn available_parallelism() -> usize {
517    std::thread::available_parallelism().map_or(1, std::num::NonZeroUsize::get)
518}
519
520/// What can go wrong discovering topology or resolving a [`PartitionConfig`].
521#[derive(Debug)]
522pub enum TopologyError {
523    /// A cpulist string was not the `"0-7,16-23"` form.
524    InvalidCpuList(
525        /// The offending list, trimmed.
526        String,
527    ),
528    /// A partition count was not a power of two.
529    NotPowerOfTwo(
530        /// The rejected count.
531        usize,
532    ),
533    /// A CPU set was empty where a partition needs at least one CPU.
534    EmptySet,
535    /// A syscall or sysfs read failed, or a pool could not be built.
536    Io(
537        /// The underlying error.
538        io::Error,
539    ),
540    /// An explicit set named a CPU outside the process's affinity mask.
541    CpuNotAllowed(
542        /// The CPU index that is outside the process's affinity mask.
543        usize,
544    ),
545}
546
547impl fmt::Display for TopologyError {
548    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
549        match self {
550            Self::InvalidCpuList(list) => {
551                write!(f, "{list:?} is not a cpulist of the form \"0-7,16-23\"")
552            }
553            Self::NotPowerOfTwo(count) => write!(
554                f,
555                "partition count {count} is not a power of two; a partition index is a fixed \
556                 set of GF(2) hash rows",
557            ),
558            Self::EmptySet => write!(f, "a partition needs a non-empty CPU set"),
559            Self::Io(err) => write!(f, "topology I/O failed: {err}"),
560            Self::CpuNotAllowed(cpu) => {
561                write!(f, "CPU {cpu} is outside this process's affinity mask",)
562            }
563        }
564    }
565}
566
567impl std::error::Error for TopologyError {
568    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
569        match self {
570            Self::Io(err) => Some(err),
571            _ => None,
572        }
573    }
574}
575
576impl From<io::Error> for TopologyError {
577    fn from(err: io::Error) -> Self {
578        Self::Io(err)
579    }
580}
581
582#[cfg(test)]
583mod tests;
paulistrings/engine/partitioned/topology.rs

paulistrings/engine/partitioned/
topology.rs
