nuvec_distributed/
backend.rs

1//! MLX's backend selection: `init(strict, backend)` with a group cached per backend, and
2//! `is_available(backend)`.
3
4use std::sync::{Mutex, PoisonError};
5use std::time::Duration;
6
7use nuvec_core::distributed::{Group, set_default_group};
8use nuvec_core::{Error, Result};
9
10use crate::{Options, hostfile, init_with};
11
12/// A distributed backend (MLX's `backend` argument).
13#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
14#[non_exhaustive]
15pub enum Backend {
16    /// The first backend whose environment is set: the ring, then JACCL.
17    Any,
18    /// MLX's TCP ring.
19    Ring,
20    /// MLX's JACCL, RDMA over Thunderbolt 5 (experimental, feature `jaccl`).
21    Jaccl,
22}
23
24impl std::str::FromStr for Backend {
25    type Err = Error;
26
27    /// MLX's names: `any`, `ring`, `jaccl`.
28    fn from_str(s: &str) -> Result<Self> {
29        match s {
30            "any" => Ok(Self::Any),
31            "ring" => Ok(Self::Ring),
32            "jaccl" => Ok(Self::Jaccl),
33            _ => Err(Error::InvalidArgument(format!(
34                "distributed: no backend {s:?} (any, ring or jaccl)"
35            ))),
36        }
37    }
38}
39
40/// Initialised groups per backend; `Any` maps to the first one (MLX's `get_backends`).
41static GROUPS: Mutex<Vec<(Backend, Group)>> = Mutex::new(Vec::new());
42
43/// `NUVEC_DISTRIBUTED_TIMEOUT` in seconds, or 30 s.
44fn timeout() -> Result<Duration> {
45    let Ok(secs) = std::env::var("NUVEC_DISTRIBUTED_TIMEOUT") else {
46        return Ok(Duration::from_secs(30));
47    };
48    let invalid = || {
49        Error::InvalidArgument(format!(
50            "distributed: NUVEC_DISTRIBUTED_TIMEOUT {secs:?} (seconds, a positive number)"
51        ))
52    };
53    let parsed: f64 = secs.trim().parse().map_err(|_| invalid())?;
54    Duration::try_from_secs_f64(parsed)
55        .ok()
56        .filter(|d| !d.is_zero())
57        .ok_or_else(invalid)
58}
59
60/// `NUVEC_DISTRIBUTED_KEY`, the pre-shared secret of every connection.
61fn secret() -> Option<String> {
62    std::env::var("NUVEC_DISTRIBUTED_KEY").ok()
63}
64
65/// The ring from `MLX_HOSTFILE` and `MLX_RANK`; `None` when they are not set.
66fn ring() -> Result<Option<Group>> {
67    let (Ok(hostfile), Ok(rank)) = (std::env::var("MLX_HOSTFILE"), std::env::var("MLX_RANK"))
68    else {
69        return Ok(None);
70    };
71    let text = std::fs::read_to_string(&hostfile)
72        .map_err(|e| Error::Io(format!("distributed: reading {hostfile}: {e}")))?;
73    let rank = rank
74        .trim()
75        .parse()
76        .map_err(|_| Error::InvalidArgument(format!("distributed: MLX_RANK {rank:?}")))?;
77    let mut options = Options::default().timeout(timeout()?);
78    if let Some(secret) = secret() {
79        options = options.key(secret);
80    }
81    init_with(&hostfile::parse(&text)?, rank, &options).map(Some)
82}
83
84/// JACCL from its environment; `None` when it is not set.
85#[cfg(feature = "jaccl")]
86fn jaccl() -> Result<Option<Group>> {
87    let Some(config) = crate::jaccl::Config::from_env()? else {
88        return Ok(None);
89    };
90    let mut options = crate::jaccl::Options::default().timeout(timeout()?);
91    if let Some(secret) = secret() {
92        options = options.key(secret);
93    }
94    crate::jaccl::init_with(&config, &options).map(Some)
95}
96
97#[cfg(not(feature = "jaccl"))]
98fn jaccl() -> Result<Option<Group>> {
99    Ok(None)
100}
101
102const RING_VARS: &str = "the ring needs MLX_HOSTFILE and MLX_RANK";
103const JACCL_VARS: &str = "JACCL needs MLX_RANK (or JACCL_RANK), MLX_IBV_DEVICES (or \
104                          JACCL_IBV_DEVICES) and MLX_JACCL_COORDINATOR (or JACCL_COORDINATOR)";
105
106/// Initialises `backend` from MLX's environment and makes it the default group, as MLX's
107/// `init(strict, backend)`. A backend initialised before returns its group (`Any`: the first
108/// group initialised). Without the environment: an error naming the missing variables if
109/// `strict`, else the group of this process alone.
110pub fn init_backend(strict: bool, backend: Backend) -> Result<Group> {
111    if let Some(g) = cached(backend) {
112        return Ok(g);
113    }
114    let (group, missing) = match backend {
115        Backend::Ring => (ring()?.map(|g| (Backend::Ring, g)), RING_VARS.to_string()),
116        Backend::Jaccl => {
117            if cfg!(not(feature = "jaccl")) {
118                return Err(Error::InvalidArgument(
119                    "distributed: the JACCL backend is experimental and not built in: enable \
120                     the `jaccl` feature"
121                        .into(),
122                ));
123            }
124            (
125                jaccl()?.map(|g| (Backend::Jaccl, g)),
126                JACCL_VARS.to_string(),
127            )
128        }
129        Backend::Any => {
130            let group = match ring()? {
131                Some(g) => Some((Backend::Ring, g)),
132                None => jaccl()?.map(|g| (Backend::Jaccl, g)),
133            };
134            let tried = if cfg!(feature = "jaccl") {
135                format!("{RING_VARS}; {JACCL_VARS}")
136            } else {
137                RING_VARS.to_string()
138            };
139            (group, tried)
140        }
141    };
142    match group {
143        Some((b, g)) => {
144            let mut groups = GROUPS.lock().unwrap_or_else(PoisonError::into_inner);
145            for key in [b, Backend::Any] {
146                if !groups.iter().any(|(k, _)| *k == key) {
147                    groups.push((key, g.clone()));
148                }
149            }
150            set_default_group(g.clone());
151            Ok(g)
152        }
153        None if strict => Err(Error::InvalidArgument(format!(
154            "distributed: no backend could be initialised ({missing})"
155        ))),
156        None => Ok(Group::single()),
157    }
158}
159
160/// The group initialised for `backend`, unless it failed or was closed: then every cached
161/// entry of it is dropped (and the default group reset), so the next `init` forms a new one.
162fn cached(backend: Backend) -> Option<Group> {
163    let mut groups = GROUPS.lock().unwrap_or_else(PoisonError::into_inner);
164    let g = groups
165        .iter()
166        .find(|(k, _)| *k == backend)
167        .map(|(_, g)| g.clone())?;
168    if g.health().is_ok() {
169        return Some(g);
170    }
171    groups.retain(|(_, c)| c.health().is_ok());
172    nuvec_core::distributed::set_default_group(Group::single());
173    None
174}
175
176/// Whether `backend` can run here, as MLX's `is_available(backend)`: the ring always;
177/// JACCL when built in and the system's RDMA library lists a device.
178pub fn is_available(backend: Backend) -> bool {
179    match backend {
180        Backend::Any | Backend::Ring => true,
181        #[cfg(feature = "jaccl")]
182        Backend::Jaccl => crate::jaccl::is_available(),
183        #[cfg(not(feature = "jaccl"))]
184        Backend::Jaccl => false,
185    }
186}
187
188#[cfg(test)]
189mod tests {
190    use super::*;
191
192    /// The tests run without the distributed environment (the multi-process tests set it in
193    /// their child processes only).
194    fn clean() -> bool {
195        let set = ["MLX_HOSTFILE", "MLX_IBV_DEVICES", "JACCL_IBV_DEVICES"]
196            .iter()
197            .any(|v| std::env::var(v).is_ok());
198        if set {
199            eprintln!("skipped: a distributed environment is set");
200        }
201        !set
202    }
203
204    #[test]
205    fn selects_backends_as_mlx() {
206        assert_eq!("jaccl".parse::<Backend>().unwrap(), Backend::Jaccl);
207        assert!("mpi".parse::<Backend>().is_err());
208        assert!(is_available(Backend::Ring) && is_available(Backend::Any));
209        if !clean() {
210            return;
211        }
212        assert_eq!(init_backend(false, Backend::Any).unwrap().size(), 1);
213        assert_eq!(init_backend(false, Backend::Ring).unwrap().size(), 1);
214        let e = init_backend(true, Backend::Ring).unwrap_err().to_string();
215        assert!(e.contains("MLX_HOSTFILE and MLX_RANK"), "{e}");
216        let e = init_backend(true, Backend::Jaccl).unwrap_err().to_string();
217        if cfg!(feature = "jaccl") {
218            assert!(e.contains("MLX_IBV_DEVICES"), "{e}");
219            assert_eq!(init_backend(false, Backend::Jaccl).unwrap().size(), 1);
220        } else {
221            assert!(e.contains("enable the `jaccl` feature"), "{e}");
222            assert!(!is_available(Backend::Jaccl));
223        }
224        let e = init_backend(true, Backend::Any).unwrap_err().to_string();
225        assert!(e.contains("MLX_HOSTFILE"), "{e}");
226    }
227}
nuvec_distributed/backend.rs

nuvec_distributed/
backend.rs
