flodl_cli/
gpus.rs

1//! `--gpus` flag parsing + single-host cluster envelope synthesis.
2//!
3//! The `--gpus` flag has uniform semantics ("use these GPUs") but the
4//! mechanism depends on the command kind:
5//!
6//! - **Cluster-aware commands** (`cluster: true`): N >= 2 GPUs trigger
7//!   synthesis of a single-host cluster envelope (master=127.0.0.1, lo
8//!   transport, one host with N ranks) and spawn-per-rank via the existing
9//!   launcher (see [`crate::cluster::prepare_cluster_env`]). The library
10//!   inside each spawned process reads the envelope from `FLODL_INTERNAL_CLUSTER_JSON`
11//!   and uses the same code path as multi-host. N = 1 is degenerate — no
12//!   synthesis, just runs single-process on that device.
13//!
14//! - **Non-cluster commands** (`test`, `clippy`, etc.): `--gpus` sets
15//!   `CUDA_VISIBLE_DEVICES` on the single child process. No envelope, no
16//!   spawning. Tests internally manage their own multi-rank coordination
17//!   (typically via the threaded `NcclRankComm` pattern in unit tests).
18//!
19//! Caller (`main.rs`) decides which mechanism applies based on whether the
20//! resolved command's `cluster:` chain enables dispatch.
21
22use std::process::Command;
23
24use crate::cluster::resolve_local_hostname;
25use crate::config::{
26    ClusterConfig, ClusterController, ClusterWorker, LocalDevices,
27    DEFAULT_CONTROLLER_PORT,
28};
29
30/// Parsed `--gpus` argument value.
31///
32/// Two forms accepted by [`GpusSpec::parse`]:
33/// - `--gpus all`: resolve to all visible CUDA devices via `nvidia-smi -L`.
34/// - `--gpus 0,1,2`: explicit comma-separated physical device indices.
35#[derive(Debug, Clone, PartialEq, Eq)]
36pub enum GpusSpec {
37    /// Use every visible CUDA device. Resolved against `nvidia-smi -L` at
38    /// [`GpusSpec::resolve`] time.
39    All,
40    /// Explicit list of physical CUDA device indices.
41    List(Vec<u8>),
42}
43
44impl GpusSpec {
45    /// Parse a `--gpus` value. Loud errors on empty, malformed, or duplicate
46    /// device indices.
47    pub fn parse(raw: &str) -> Result<Self, String> {
48        let trimmed = raw.trim();
49        if trimmed.is_empty() {
50            return Err(
51                "--gpus requires a value (e.g. `--gpus 0,1` or `--gpus all`)".to_string(),
52            );
53        }
54        if trimmed.eq_ignore_ascii_case("all") {
55            return Ok(GpusSpec::All);
56        }
57        let mut out = Vec::new();
58        for part in trimmed.split(',') {
59            let p = part.trim();
60            if p.is_empty() {
61                return Err(format!("--gpus: empty entry in {trimmed:?}"));
62            }
63            let idx: u8 = p.parse().map_err(|e| {
64                format!("--gpus: cannot parse {p:?} as device index: {e}")
65            })?;
66            out.push(idx);
67        }
68        let mut sorted = out.clone();
69        sorted.sort_unstable();
70        for win in sorted.windows(2) {
71            if win[0] == win[1] {
72                return Err(format!(
73                    "--gpus: duplicate device index {} in {trimmed:?}",
74                    win[0]
75                ));
76            }
77        }
78        Ok(GpusSpec::List(out))
79    }
80
81    /// Resolve to a concrete list of physical CUDA device indices.
82    ///
83    /// `List` returns its entries verbatim. `All` shells out to
84    /// `nvidia-smi -L` and counts the result -- loud error if nvidia-smi
85    /// is missing or returns 0 GPUs.
86    pub fn resolve(&self) -> Result<Vec<u8>, String> {
87        match self {
88            GpusSpec::List(v) => Ok(v.clone()),
89            GpusSpec::All => {
90                let count = count_visible_gpus_via_nvidia_smi()?;
91                if count == 0 {
92                    return Err(
93                        "--gpus all: nvidia-smi reports 0 GPUs visible. Install \
94                         NVIDIA drivers or specify devices explicitly (e.g. \
95                         --gpus 0)."
96                            .to_string(),
97                    );
98                }
99                if count > u8::MAX as usize {
100                    return Err(format!(
101                        "--gpus all: nvidia-smi reports {count} GPUs which \
102                         exceeds the supported device-index range (0..255). \
103                         Specify devices explicitly via --gpus."
104                    ));
105                }
106                Ok((0u8..count as u8).collect())
107            }
108        }
109    }
110}
111
112/// Count visible CUDA devices via `nvidia-smi -L`.
113///
114/// Each GPU is one line starting with `GPU <idx>:`. Returns the number of
115/// such lines. Loud error if nvidia-smi is missing or exits non-zero.
116pub fn count_visible_gpus_via_nvidia_smi() -> Result<usize, String> {
117    let out = Command::new("nvidia-smi")
118        .arg("-L")
119        .output()
120        .map_err(|e| {
121            format!(
122                "failed to run `nvidia-smi -L`: {e}. Install NVIDIA drivers \
123                 or specify devices explicitly (e.g. --gpus 0)."
124            )
125        })?;
126    if !out.status.success() {
127        let stderr = String::from_utf8_lossy(&out.stderr);
128        return Err(format!(
129            "`nvidia-smi -L` exited non-zero: {}",
130            stderr.trim()
131        ));
132    }
133    let stdout = String::from_utf8_lossy(&out.stdout);
134    Ok(stdout.lines().filter(|l| l.starts_with("GPU ")).count())
135}
136
137/// Build a `ClusterConfig` for single-host loopback from a list of physical
138/// CUDA device indices.
139///
140/// Used when `--gpus` is set on a cluster-aware command and no `cluster:`
141/// block is in YAML. Returns a config with one host (this machine), N ranks
142/// (`0..devices.len()`), NCCL loopback transport (`lo`).
143///
144/// `controller.port` defaults to [`DEFAULT_CONTROLLER_PORT`] (1337),
145/// overridable via `FLODL_CONTROLLER_PORT`. Concurrent `fdl` cluster
146/// commands on the same host must use distinct ports to avoid
147/// rendezvous collisions.
148pub fn synthesize_local_cluster(devices: &[u8]) -> Result<ClusterConfig, String> {
149    if devices.is_empty() {
150        return Err("synthesize_local_cluster: device list is empty".to_string());
151    }
152    let hostname = resolve_local_hostname();
153    let path = std::env::current_dir()
154        .map(|p| p.to_string_lossy().into_owned())
155        .map_err(|e| {
156            format!("synthesize_local_cluster: cannot read current_dir: {e}")
157        })?;
158    let port = std::env::var("FLODL_CONTROLLER_PORT")
159        .ok()
160        .and_then(|s| s.parse::<u16>().ok())
161        .unwrap_or(DEFAULT_CONTROLLER_PORT);
162
163    Ok(ClusterConfig {
164        controller: ClusterController {
165            host: "127.0.0.1".to_string(),
166            port,
167            path: path.clone(),
168            docker: None,
169            arch: None,
170            data_path: None,
171            join: None,
172        },
173        workers: vec![ClusterWorker {
174            host: hostname,
175            ranks: (0..devices.len()).collect(),
176            local_devices: LocalDevices::Explicit(devices.to_vec()),
177            nccl_socket_ifname: "lo".to_string(),
178            path,
179            ssh: None,
180            tunnel: false,
181            arch: None,
182            data_path: None,
183            docker: None,
184            env: std::collections::BTreeMap::new(),
185        }],
186        env: std::collections::BTreeMap::new(),
187    })
188}
189
190/// Set `CUDA_VISIBLE_DEVICES` to restrict the spawned process to the given
191/// physical CUDA device indices.
192///
193/// Used on the non-cluster path (`--gpus 0,1` on `fdl test`, `clippy`, etc.)
194/// so the single child process sees only the requested GPUs. NVIDIA Docker
195/// forwards `CUDA_VISIBLE_DEVICES` to containers automatically.
196///
197/// Empty slice removes the var. The caller normally avoids calling with an
198/// empty slice (a loud error earlier in the resolution path).
199///
200/// # Safety
201///
202/// Calls `std::env::set_var` which is unsafe in multi-threaded programs.
203/// Must be called from `main` before any threads are spawned, which is the
204/// case for the fdl-cli dispatch flow.
205pub unsafe fn apply_cuda_visible_devices(devices: &[u8]) {
206    let joined = devices
207        .iter()
208        .map(|d| d.to_string())
209        .collect::<Vec<_>>()
210        .join(",");
211    if joined.is_empty() {
212        unsafe { std::env::remove_var("CUDA_VISIBLE_DEVICES") };
213    } else {
214        unsafe { std::env::set_var("CUDA_VISIBLE_DEVICES", &joined) };
215    }
216}
217
218#[cfg(test)]
219mod tests {
220    use super::*;
221
222    #[test]
223    fn parse_all_case_insensitive() {
224        assert_eq!(GpusSpec::parse("all").unwrap(), GpusSpec::All);
225        assert_eq!(GpusSpec::parse("ALL").unwrap(), GpusSpec::All);
226        assert_eq!(GpusSpec::parse("All").unwrap(), GpusSpec::All);
227    }
228
229    #[test]
230    fn parse_single_index() {
231        assert_eq!(GpusSpec::parse("0").unwrap(), GpusSpec::List(vec![0]));
232        assert_eq!(GpusSpec::parse("3").unwrap(), GpusSpec::List(vec![3]));
233    }
234
235    #[test]
236    fn parse_multiple_indices() {
237        assert_eq!(
238            GpusSpec::parse("0,1,2").unwrap(),
239            GpusSpec::List(vec![0, 1, 2])
240        );
241        assert_eq!(GpusSpec::parse("3,1").unwrap(), GpusSpec::List(vec![3, 1]));
242    }
243
244    #[test]
245    fn parse_tolerates_whitespace() {
246        assert_eq!(
247            GpusSpec::parse(" 0 , 1 ").unwrap(),
248            GpusSpec::List(vec![0, 1])
249        );
250        assert_eq!(GpusSpec::parse("  all  ").unwrap(), GpusSpec::All);
251    }
252
253    #[test]
254    fn parse_rejects_empty() {
255        let err = GpusSpec::parse("").unwrap_err();
256        assert!(err.contains("--gpus requires a value"), "got: {err}");
257        let err = GpusSpec::parse("   ").unwrap_err();
258        assert!(err.contains("--gpus requires a value"), "got: {err}");
259    }
260
261    #[test]
262    fn parse_rejects_empty_entry() {
263        let err = GpusSpec::parse("0,,1").unwrap_err();
264        assert!(err.contains("empty entry"), "got: {err}");
265        let err = GpusSpec::parse(",0").unwrap_err();
266        assert!(err.contains("empty entry"), "got: {err}");
267    }
268
269    #[test]
270    fn parse_rejects_non_numeric() {
271        let err = GpusSpec::parse("0,abc").unwrap_err();
272        assert!(err.contains("cannot parse"), "got: {err}");
273        assert!(err.contains("abc"), "got: {err}");
274    }
275
276    #[test]
277    fn parse_rejects_duplicates() {
278        let err = GpusSpec::parse("0,1,0").unwrap_err();
279        assert!(err.contains("duplicate"), "got: {err}");
280        assert!(err.contains("0"), "got: {err}");
281    }
282
283    #[test]
284    fn resolve_list_returns_verbatim() {
285        let r = GpusSpec::List(vec![3, 1]).resolve().unwrap();
286        assert_eq!(r, vec![3, 1]);
287    }
288
289    #[test]
290    fn synthesize_local_cluster_basic_shape() {
291        // We don't control hostname/cwd here, so just assert structural invariants.
292        let c = synthesize_local_cluster(&[0, 1]).unwrap();
293        assert_eq!(c.controller.host, "127.0.0.1");
294        assert_eq!(c.workers.len(), 1);
295        let w = &c.workers[0];
296        assert_eq!(w.ranks, vec![0, 1]);
297        assert_eq!(w.local_devices, LocalDevices::Explicit(vec![0, 1]));
298        assert_eq!(w.nccl_socket_ifname, "lo");
299        assert!(w.arch.is_none());
300        assert!(w.ssh.is_none());
301        assert!(!w.host.trim().is_empty(), "hostname must be non-empty");
302        assert!(!w.path.trim().is_empty(), "path must be non-empty");
303    }
304
305    #[test]
306    fn synthesize_local_cluster_validates() {
307        // The synthesized config must pass ClusterConfig::validate (so the
308        // launcher accepts it without special-casing).
309        let c = synthesize_local_cluster(&[0, 1]).unwrap();
310        c.validate().expect("synthesized cluster must pass validate");
311    }
312
313    #[test]
314    fn synthesize_local_cluster_single_device() {
315        // N=1 is structurally valid (validate enforces 0..world_size with
316        // ranks=[0], devices=[0]). Caller decides whether to use it.
317        let c = synthesize_local_cluster(&[2]).unwrap();
318        c.validate().expect("single-device synthesized config validates");
319        assert_eq!(c.workers[0].ranks, vec![0]);
320        assert_eq!(c.workers[0].local_devices, LocalDevices::Explicit(vec![2]));
321    }
322
323    #[test]
324    fn synthesize_local_cluster_rejects_empty() {
325        let err = synthesize_local_cluster(&[]).unwrap_err();
326        assert!(err.contains("empty"), "got: {err}");
327    }
328
329    #[test]
330    fn synthesize_local_cluster_respects_controller_port_env() {
331        // SAFETY: cargo test parallelism. Use a unique env var name probe
332        // pattern instead of FLODL_CONTROLLER_PORT to avoid clobbering other
333        // tests -- but here we DO want to test the env reading path, so we
334        // accept the race. Single-threaded mod tests would be cleaner.
335        // For now, accept that this test runs serially-enough.
336        unsafe {
337            std::env::set_var("FLODL_CONTROLLER_PORT", "31415");
338        }
339        let c = synthesize_local_cluster(&[0]).unwrap();
340        unsafe {
341            std::env::remove_var("FLODL_CONTROLLER_PORT");
342        }
343        assert_eq!(c.controller.port, 31415);
344    }
345}
flodl_cli/gpus.rs

flodl_cli/
gpus.rs
