scx_cosmos/
gpu.rs

1// SPDX-License-Identifier: GPL-2.0
2
3use std::collections::{HashMap, HashSet};
4use std::path::{Path, PathBuf};
5
6use anyhow::Result;
7use log::debug;
8
9use crate::cgroup::CgroupReader;
10
11/// Keep only NVML processes whose GPUs resolve to one NUMA node.
12pub(crate) fn direct_gpu_processes(
13    process_nodes: &HashMap<u32, HashSet<u32>>,
14) -> HashMap<u32, u32> {
15    process_nodes
16        .iter()
17        .filter(|&(&_tgid, nodes)| nodes.len() == 1)
18        .map(|(&tgid, nodes)| (tgid, *nodes.iter().next().unwrap()))
19        .collect()
20}
21
22/// Expand NVML processes to peer processes in the same exact cgroup.
23///
24/// NVML remains the source of GPU usage and NUMA locality. Cgroups are used
25/// only to discover the other processes that belong to the same workload.
26pub(crate) fn expand_gpu_processes(
27    process_nodes: &HashMap<u32, HashSet<u32>>,
28    cgroups: &CgroupReader,
29) -> HashMap<u32, u32> {
30    expand_gpu_processes_with(
31        process_nodes,
32        |tgid| cgroups.process_cgroup(tgid),
33        |cgroup| cgroups.processes(cgroup),
34    )
35}
36
37/// Fit workload hints into the BPF map without partially expanding a cgroup.
38pub(crate) fn fit_gpu_processes(
39    expanded: HashMap<u32, u32>,
40    direct: &HashMap<u32, u32>,
41    max_entries: usize,
42) -> HashMap<u32, u32> {
43    if expanded.len() <= max_entries {
44        expanded
45    } else if direct.len() <= max_entries {
46        direct.clone()
47    } else {
48        HashMap::new()
49    }
50}
51
52fn expand_gpu_processes_with<C, P>(
53    process_nodes: &HashMap<u32, HashSet<u32>>,
54    mut process_cgroup: C,
55    mut cgroup_processes: P,
56) -> HashMap<u32, u32>
57where
58    C: FnMut(u32) -> Result<Option<PathBuf>>,
59    P: FnMut(&Path) -> Result<HashSet<u32>>,
60{
61    #[derive(Default)]
62    struct Observation {
63        nodes: HashSet<u32>,
64        seeds: HashSet<u32>,
65    }
66
67    // Every NVML observation is authoritative, including processes whose
68    // GPUs span multiple nodes and therefore intentionally have no hint.
69    // Peer discovery must not override them when membership changes.
70    let direct = direct_gpu_processes(process_nodes);
71    let mut peer_nodes: HashMap<u32, HashSet<u32>> = HashMap::new();
72    let mut cgroups: HashMap<PathBuf, Observation> = HashMap::new();
73
74    for (&tgid, nodes) in process_nodes {
75        match process_cgroup(tgid) {
76            Ok(Some(cgroup)) => {
77                let observation = cgroups.entry(cgroup).or_default();
78                observation.nodes.extend(nodes.iter().copied());
79                observation.seeds.insert(tgid);
80            }
81            Ok(None) => debug!("GPU process {tgid} is in the root cgroup; not expanding"),
82            Err(error) => debug!("GPU process {tgid} cgroup discovery failed: {error:#}"),
83        }
84    }
85
86    for (cgroup, observation) in cgroups {
87        if observation.nodes.len() != 1 {
88            debug!(
89                "GPU cgroup {} spans multiple NUMA nodes; not expanding",
90                cgroup.display()
91            );
92            continue;
93        }
94        let node = *observation.nodes.iter().next().unwrap();
95        match cgroup_processes(&cgroup) {
96            Ok(processes) => {
97                // Membership is mutable. Do not expand an old path after all
98                // of the NVML seed processes have moved away or exited.
99                if observation.seeds.is_disjoint(&processes) {
100                    debug!(
101                        "GPU cgroup {} no longer contains an observed GPU process; not expanding",
102                        cgroup.display()
103                    );
104                    continue;
105                }
106                debug!(
107                    "GPU cgroup {} expands {} NVML processes to {} workload processes",
108                    cgroup.display(),
109                    observation.seeds.len(),
110                    processes.len()
111                );
112                for tgid in processes {
113                    if !process_nodes.contains_key(&tgid) {
114                        peer_nodes.entry(tgid).or_default().insert(node);
115                    }
116                }
117            }
118            Err(error) => debug!(
119                "GPU cgroup {} process discovery failed: {error:#}",
120                cgroup.display()
121            ),
122        }
123    }
124
125    let mut expanded = direct;
126    for (tgid, nodes) in peer_nodes {
127        // A peer discovered in conflicting workloads is safer without a hint
128        // than with a node chosen according to iteration order.
129        if nodes.len() == 1 {
130            expanded.insert(tgid, *nodes.iter().next().unwrap());
131        }
132    }
133    expanded
134}
135
136#[cfg(test)]
137mod tests {
138    use anyhow::bail;
139
140    use super::*;
141
142    fn nodes(values: &[u32]) -> HashSet<u32> {
143        values.iter().copied().collect()
144    }
145
146    #[test]
147    fn expands_exact_cgroup_processes() {
148        let observed = HashMap::from([(100, nodes(&[0]))]);
149        let expanded = expand_gpu_processes_with(
150            &observed,
151            |tgid| {
152                assert_eq!(tgid, 100);
153                Ok(Some(PathBuf::from("/cgroup/a")))
154            },
155            |path| {
156                assert_eq!(path, Path::new("/cgroup/a"));
157                Ok(nodes(&[100, 200, 300]))
158            },
159        );
160
161        assert_eq!(expanded, HashMap::from([(100, 0), (200, 0), (300, 0)]));
162    }
163
164    #[test]
165    fn falls_back_to_direct_process_on_discovery_failure() {
166        let observed = HashMap::from([(100, nodes(&[0]))]);
167        let expanded = expand_gpu_processes_with(
168            &observed,
169            |_| bail!("unavailable"),
170            |_| bail!("must not read cgroup.procs"),
171        );
172        assert_eq!(expanded, HashMap::from([(100, 0)]));
173    }
174
175    #[test]
176    fn does_not_expand_after_seed_leaves_cgroup() {
177        let observed = HashMap::from([(100, nodes(&[0]))]);
178        let expanded = expand_gpu_processes_with(
179            &observed,
180            |_| Ok(Some(PathBuf::from("/cgroup/a"))),
181            |_| Ok(nodes(&[200, 300])),
182        );
183        assert_eq!(expanded, HashMap::from([(100, 0)]));
184    }
185
186    #[test]
187    fn does_not_expand_cgroup_across_nodes() {
188        let observed = HashMap::from([(100, nodes(&[0])), (101, nodes(&[1]))]);
189        let expanded = expand_gpu_processes_with(
190            &observed,
191            |_| Ok(Some(PathBuf::from("/cgroup/a"))),
192            |_| Ok(nodes(&[100, 101, 200])),
193        );
194
195        assert_eq!(expanded, HashMap::from([(100, 0), (101, 1)]));
196    }
197
198    #[test]
199    fn removes_ambiguous_process_hint() {
200        let observed = HashMap::from([(100, nodes(&[0, 1]))]);
201        let expanded = expand_gpu_processes_with(
202            &observed,
203            |_| Ok(Some(PathBuf::from("/cgroup/a"))),
204            |_| Ok(nodes(&[100, 200])),
205        );
206        assert!(expanded.is_empty());
207    }
208
209    #[test]
210    fn removes_peer_seen_in_conflicting_cgroups() {
211        let observed = HashMap::from([(100, nodes(&[0])), (101, nodes(&[1]))]);
212        let expanded = expand_gpu_processes_with(
213            &observed,
214            |tgid| Ok(Some(PathBuf::from(format!("/cgroup/{tgid}")))),
215            |_| Ok(nodes(&[200])),
216        );
217
218        assert_eq!(expanded, HashMap::from([(100, 0), (101, 1)]));
219        assert!(!expanded.contains_key(&200));
220    }
221
222    #[test]
223    fn direct_nvml_mapping_wins_membership_race() {
224        let observed = HashMap::from([(100, nodes(&[0])), (101, nodes(&[1]))]);
225        let expanded = expand_gpu_processes_with(
226            &observed,
227            |tgid| Ok(Some(PathBuf::from(format!("/cgroup/{tgid}")))),
228            |path| {
229                if path == Path::new("/cgroup/100") {
230                    Ok(nodes(&[100]))
231                } else {
232                    Ok(nodes(&[100, 101]))
233                }
234            },
235        );
236
237        assert_eq!(expanded, HashMap::from([(100, 0), (101, 1)]));
238    }
239
240    #[test]
241    fn ambiguous_nvml_observation_wins_membership_race() {
242        let observed = HashMap::from([(100, nodes(&[0])), (101, nodes(&[0, 1]))]);
243        let expanded = expand_gpu_processes_with(
244            &observed,
245            |tgid| {
246                if tgid == 100 {
247                    Ok(Some(PathBuf::from("/cgroup/a")))
248                } else {
249                    bail!("cgroup lookup raced")
250                }
251            },
252            |_| Ok(nodes(&[100, 101])),
253        );
254
255        assert_eq!(expanded, HashMap::from([(100, 0)]));
256    }
257
258    #[test]
259    fn capacity_falls_back_without_partial_expansion() {
260        let direct = HashMap::from([(100, 0)]);
261        let expanded = HashMap::from([(100, 0), (200, 0), (300, 0)]);
262
263        assert_eq!(fit_gpu_processes(expanded.clone(), &direct, 3), expanded);
264        assert_eq!(fit_gpu_processes(expanded, &direct, 2), direct);
265        assert!(fit_gpu_processes(HashMap::from([(100, 0)]), &direct, 0).is_empty());
266    }
267}
scx_cosmos/gpu.rs

scx_cosmos/
gpu.rs
