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 crate::cluster::resolve_local_hostname;
23use crate::config::{
24    ClusterConfig, ClusterController, ClusterWorker, DEFAULT_CONTROLLER_PORT, LocalDevices,
25};
26
27/// Parsed `--gpus` argument value.
28///
29/// Two forms accepted by [`GpusSpec::parse`]:
30/// - `--gpus all`: resolve to all visible CUDA devices via `nvidia-smi -L`.
31/// - `--gpus 0,1,2`: explicit comma-separated physical device indices.
32#[derive(Debug, Clone, PartialEq, Eq)]
33pub enum GpusSpec {
34    /// Use every visible CUDA device. Resolved against `nvidia-smi -L` at
35    /// [`GpusSpec::resolve`] time.
36    All,
37    /// Explicit list of physical CUDA device indices.
38    List(Vec<u8>),
39}
40
41impl GpusSpec {
42    /// Parse a `--gpus` value. Loud errors on empty, malformed, or duplicate
43    /// device indices.
44    pub fn parse(raw: &str) -> Result<Self, String> {
45        let trimmed = raw.trim();
46        if trimmed.is_empty() {
47            return Err("--gpus requires a value (e.g. `--gpus 0,1` or `--gpus all`)".to_string());
48        }
49        if trimmed.eq_ignore_ascii_case("all") {
50            return Ok(GpusSpec::All);
51        }
52        let mut out = Vec::new();
53        for part in trimmed.split(',') {
54            let p = part.trim();
55            if p.is_empty() {
56                return Err(format!("--gpus: empty entry in {trimmed:?}"));
57            }
58            let idx: u8 = p
59                .parse()
60                .map_err(|e| format!("--gpus: cannot parse {p:?} as device index: {e}"))?;
61            out.push(idx);
62        }
63        let mut sorted = out.clone();
64        sorted.sort_unstable();
65        for win in sorted.windows(2) {
66            if win[0] == win[1] {
67                return Err(format!(
68                    "--gpus: duplicate device index {} in {trimmed:?}",
69                    win[0]
70                ));
71            }
72        }
73        Ok(GpusSpec::List(out))
74    }
75
76    /// Resolve to a concrete list of physical CUDA device indices.
77    ///
78    /// `List` returns its entries verbatim. `All` shells out to
79    /// `nvidia-smi -L` and counts the result -- loud error if nvidia-smi
80    /// is missing or returns 0 GPUs.
81    pub fn resolve(&self) -> Result<Vec<u8>, String> {
82        match self {
83            GpusSpec::List(v) => Ok(v.clone()),
84            GpusSpec::All => {
85                // `require_devices` turns an empty sweep into the best
86                // available explanation: a driver that failed to
87                // enumerate, hardware present without its stack
88                // installed, or genuinely no GPU. An explicit `--gpus
89                // all` must fail loudly rather than resolve to zero.
90                let devices = local_gpu_count().map_err(|e| format!("--gpus all: {e}"))?;
91                if devices > u8::MAX as usize {
92                    return Err(format!(
93                        "--gpus all: {devices} GPUs detected, which exceeds \
94                         the supported device-index range (0..255). Specify \
95                         devices explicitly via --gpus."
96                    ));
97                }
98                Ok((0u8..devices as u8).collect())
99            }
100        }
101    }
102}
103
104/// Number of GPUs on this box, or a caller-facing reason there are none.
105///
106/// One entry point for every "how many GPUs are here" question in fdl,
107/// across every vendor. The error is the point: an empty device list has
108/// several causes (no driver, driver present but its tool broken,
109/// hardware present without its stack installed, genuinely no card) and
110/// a command that was *asked* for GPUs must say which one it hit rather
111/// than silently resolving to zero.
112///
113/// Counts **physical** devices: `--gpus`/`local_devices` select from the
114/// full set, and applying a visibility mask here would make the
115/// selection depend on a mask the selection itself is about to set.
116pub fn local_gpu_count() -> Result<usize, String> {
117    flodl_hw::survey().require_devices().map(|d| d.len())
118}
119
120/// Build a `ClusterConfig` for single-host loopback from a list of physical
121/// CUDA device indices.
122///
123/// Used when `--gpus` is set on a cluster-aware command and no `cluster:`
124/// block is in YAML. Returns a config with one host (this machine), N ranks
125/// (`0..devices.len()`), NCCL loopback transport (`lo`).
126///
127/// `controller.port` defaults to [`DEFAULT_CONTROLLER_PORT`] (1337),
128/// overridable via `FLODL_CONTROLLER_PORT`. Concurrent `fdl` cluster
129/// commands on the same host must use distinct ports to avoid
130/// rendezvous collisions.
131pub fn synthesize_local_cluster(devices: &[u8]) -> Result<ClusterConfig, String> {
132    if devices.is_empty() {
133        return Err("synthesize_local_cluster: device list is empty".to_string());
134    }
135    let hostname = resolve_local_hostname();
136    let path = std::env::current_dir()
137        .map(|p| p.to_string_lossy().into_owned())
138        .map_err(|e| format!("synthesize_local_cluster: cannot read current_dir: {e}"))?;
139    let port = std::env::var("FLODL_CONTROLLER_PORT")
140        .ok()
141        .and_then(|s| s.parse::<u16>().ok())
142        .unwrap_or(DEFAULT_CONTROLLER_PORT);
143
144    Ok(ClusterConfig {
145        controller: ClusterController {
146            host: "127.0.0.1".to_string(),
147            port,
148            path: path.clone(),
149            docker: None,
150            arch: None,
151            data_path: None,
152            join: None,
153        },
154        workers: vec![ClusterWorker {
155            host: hostname,
156            ranks: (0..devices.len()).collect(),
157            local_devices: LocalDevices::Explicit(devices.to_vec()),
158            nccl_socket_ifname: "lo".to_string(),
159            path,
160            ssh: None,
161            tunnel: false,
162            arch: None,
163            data_path: None,
164            gpu_ram_share: None,
165            docker: None,
166            env: std::collections::BTreeMap::new(),
167        }],
168        env: std::collections::BTreeMap::new(),
169        gpu_ram_share: None,
170    })
171}
172
173/// Set `CUDA_VISIBLE_DEVICES` to restrict the spawned process to the given
174/// physical CUDA device indices.
175///
176/// Used on the non-cluster path (`--gpus 0,1` on `fdl test`, `clippy`, etc.)
177/// so the single child process sees only the requested GPUs. NVIDIA Docker
178/// forwards `CUDA_VISIBLE_DEVICES` to containers automatically.
179///
180/// Empty slice removes the var. The caller normally avoids calling with an
181/// empty slice (a loud error earlier in the resolution path).
182///
183/// # Safety
184///
185/// Calls `std::env::set_var` which is unsafe in multi-threaded programs.
186/// Must be called from `main` before any threads are spawned, which is the
187/// case for the fdl-cli dispatch flow.
188pub unsafe fn apply_cuda_visible_devices(devices: &[u8]) {
189    let joined = devices
190        .iter()
191        .map(|d| d.to_string())
192        .collect::<Vec<_>>()
193        .join(",");
194    // Both vendors' spellings: HIP prefers its own variable over
195    // CUDA_VISIBLE_DEVICES (first one set wins), so setting only the
196    // CUDA one would leave an AMD box unmasked whenever HIP_VISIBLE_DEVICES
197    // is already in the environment. Inert where the other vendor's
198    // runtime never looks.
199    for key in ["CUDA_VISIBLE_DEVICES", "HIP_VISIBLE_DEVICES"] {
200        if joined.is_empty() {
201            unsafe { std::env::remove_var(key) };
202        } else {
203            unsafe { std::env::set_var(key, &joined) };
204        }
205    }
206}
207
208#[cfg(test)]
209mod tests {
210    use super::*;
211
212    #[test]
213    fn parse_all_case_insensitive() {
214        assert_eq!(GpusSpec::parse("all").unwrap(), GpusSpec::All);
215        assert_eq!(GpusSpec::parse("ALL").unwrap(), GpusSpec::All);
216        assert_eq!(GpusSpec::parse("All").unwrap(), GpusSpec::All);
217    }
218
219    #[test]
220    fn parse_single_index() {
221        assert_eq!(GpusSpec::parse("0").unwrap(), GpusSpec::List(vec![0]));
222        assert_eq!(GpusSpec::parse("3").unwrap(), GpusSpec::List(vec![3]));
223    }
224
225    #[test]
226    fn parse_multiple_indices() {
227        assert_eq!(
228            GpusSpec::parse("0,1,2").unwrap(),
229            GpusSpec::List(vec![0, 1, 2])
230        );
231        assert_eq!(GpusSpec::parse("3,1").unwrap(), GpusSpec::List(vec![3, 1]));
232    }
233
234    #[test]
235    fn parse_tolerates_whitespace() {
236        assert_eq!(
237            GpusSpec::parse(" 0 , 1 ").unwrap(),
238            GpusSpec::List(vec![0, 1])
239        );
240        assert_eq!(GpusSpec::parse("  all  ").unwrap(), GpusSpec::All);
241    }
242
243    #[test]
244    fn parse_rejects_empty() {
245        let err = GpusSpec::parse("").unwrap_err();
246        assert!(err.contains("--gpus requires a value"), "got: {err}");
247        let err = GpusSpec::parse("   ").unwrap_err();
248        assert!(err.contains("--gpus requires a value"), "got: {err}");
249    }
250
251    #[test]
252    fn parse_rejects_empty_entry() {
253        let err = GpusSpec::parse("0,,1").unwrap_err();
254        assert!(err.contains("empty entry"), "got: {err}");
255        let err = GpusSpec::parse(",0").unwrap_err();
256        assert!(err.contains("empty entry"), "got: {err}");
257    }
258
259    #[test]
260    fn parse_rejects_non_numeric() {
261        let err = GpusSpec::parse("0,abc").unwrap_err();
262        assert!(err.contains("cannot parse"), "got: {err}");
263        assert!(err.contains("abc"), "got: {err}");
264    }
265
266    #[test]
267    fn parse_rejects_duplicates() {
268        let err = GpusSpec::parse("0,1,0").unwrap_err();
269        assert!(err.contains("duplicate"), "got: {err}");
270        assert!(err.contains("0"), "got: {err}");
271    }
272
273    #[test]
274    fn resolve_list_returns_verbatim() {
275        let r = GpusSpec::List(vec![3, 1]).resolve().unwrap();
276        assert_eq!(r, vec![3, 1]);
277    }
278
279    #[test]
280    fn synthesize_local_cluster_basic_shape() {
281        // We don't control hostname/cwd here, so just assert structural invariants.
282        let c = synthesize_local_cluster(&[0, 1]).unwrap();
283        assert_eq!(c.controller.host, "127.0.0.1");
284        assert_eq!(c.workers.len(), 1);
285        let w = &c.workers[0];
286        assert_eq!(w.ranks, vec![0, 1]);
287        assert_eq!(w.local_devices, LocalDevices::Explicit(vec![0, 1]));
288        assert_eq!(w.nccl_socket_ifname, "lo");
289        assert!(w.arch.is_none());
290        assert!(w.ssh.is_none());
291        assert!(!w.host.trim().is_empty(), "hostname must be non-empty");
292        assert!(!w.path.trim().is_empty(), "path must be non-empty");
293    }
294
295    #[test]
296    fn synthesize_local_cluster_validates() {
297        // The synthesized config must pass ClusterConfig::validate (so the
298        // launcher accepts it without special-casing).
299        let c = synthesize_local_cluster(&[0, 1]).unwrap();
300        c.validate()
301            .expect("synthesized cluster must pass validate");
302    }
303
304    #[test]
305    fn synthesize_local_cluster_single_device() {
306        // N=1 is structurally valid (validate enforces 0..world_size with
307        // ranks=[0], devices=[0]). Caller decides whether to use it.
308        let c = synthesize_local_cluster(&[2]).unwrap();
309        c.validate()
310            .expect("single-device synthesized config validates");
311        assert_eq!(c.workers[0].ranks, vec![0]);
312        assert_eq!(c.workers[0].local_devices, LocalDevices::Explicit(vec![2]));
313    }
314
315    #[test]
316    fn synthesize_local_cluster_rejects_empty() {
317        let err = synthesize_local_cluster(&[]).unwrap_err();
318        assert!(err.contains("empty"), "got: {err}");
319    }
320
321    #[test]
322    fn synthesize_local_cluster_respects_controller_port_env() {
323        // SAFETY: cargo test parallelism. Use a unique env var name probe
324        // pattern instead of FLODL_CONTROLLER_PORT to avoid clobbering other
325        // tests -- but here we DO want to test the env reading path, so we
326        // accept the race. Single-threaded mod tests would be cleaner.
327        // For now, accept that this test runs serially-enough.
328        unsafe {
329            std::env::set_var("FLODL_CONTROLLER_PORT", "31415");
330        }
331        let c = synthesize_local_cluster(&[0]).unwrap();
332        unsafe {
333            std::env::remove_var("FLODL_CONTROLLER_PORT");
334        }
335        assert_eq!(c.controller.port, 31415);
336    }
337}
flodl_cli/gpus.rs

flodl_cli/
gpus.rs
