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

scx_cosmos/
gpu.rs
