nuvec_distributed/jaccl/
config.rs

1//! MLX's JACCL configuration: the rank, the device file (per rank and peer, the RDMA
2//! devices of the cables between them), the coordinator and the ring preference.
3
4use nuvec_core::{Error, Result};
5
6use super::mesh::MAX_PEERS;
7use super::ring::MAX_CABLES;
8use crate::json::{self, Value};
9
10/// A JACCL group's configuration, as MLX's `jaccl::Config`.
11#[derive(Clone, Debug, PartialEq, Eq)]
12pub struct Config {
13    pub(crate) rank: usize,
14    pub(crate) devices: Vec<Vec<Vec<String>>>,
15    pub(crate) coordinator: Option<String>,
16    pub(crate) prefer_ring: bool,
17}
18
19/// The topology a configuration selects.
20#[derive(Clone, Copy, Debug, PartialEq, Eq)]
21#[non_exhaustive]
22pub enum Topology {
23    /// Every pair of ranks is cabled; `wires` cables per pair.
24    Mesh {
25        /// Cables per pair (the smallest count over all pairs).
26        wires: usize,
27    },
28    /// Every rank is cabled to its two neighbours.
29    Ring {
30        /// Cables per direction.
31        cables: usize,
32    },
33}
34
35/// MLX's device file: a JSON array with one entry per rank, each an array with one entry
36/// per peer: `null` (no cable; the rank itself), a device name, or a list of names (one per
37/// cable; the `k`-th cable of a pair is the `k`-th name on both sides).
38pub fn parse_devices(text: &str) -> Result<Vec<Vec<Vec<String>>>> {
39    let invalid = |what: String| Error::InvalidArgument(format!("jaccl: device file: {what}"));
40    let Value::Array(ranks) = json::parse(text, "jaccl: device file")? else {
41        return Err(invalid("expected an array of arrays".into()));
42    };
43    let size = ranks.len();
44    ranks
45        .into_iter()
46        .enumerate()
47        .map(|(rank, row)| {
48            let Value::Array(peers) = row else {
49                return Err(invalid(format!("rank {rank}: expected an array")));
50            };
51            if peers.len() != size {
52                return Err(invalid(format!(
53                    "rank {rank} lists {} peers, the file has {size} ranks",
54                    peers.len()
55                )));
56            }
57            peers
58                .into_iter()
59                .enumerate()
60                .map(|(peer, entry)| match entry {
61                    Value::Null => Ok(Vec::new()),
62                    Value::String(name) => Ok(vec![name]),
63                    Value::Array(names) => names
64                        .into_iter()
65                        .map(|n| match n {
66                            Value::String(n) => Ok(n),
67                            _ => Err(invalid(format!(
68                                "rank {rank} → {peer}: device names are strings"
69                            ))),
70                        })
71                        .collect(),
72                    _ => Err(invalid(format!(
73                        "rank {rank} → {peer}: expected null, a device name or a list of names"
74                    ))),
75                })
76                .collect()
77        })
78        .collect()
79}
80
81impl Config {
82    /// Rank `rank` of the group the device matrix `devices` describes (`size` rows of `size`
83    /// entries).
84    pub fn new(rank: usize, devices: Vec<Vec<Vec<String>>>) -> Result<Self> {
85        let size = devices.len();
86        for (r, row) in devices.iter().enumerate() {
87            if row.len() != size {
88                return Err(Error::InvalidArgument(format!(
89                    "jaccl: the device matrix has {size} rows but row {r} has {} entries",
90                    row.len()
91                )));
92            }
93        }
94        if rank >= size {
95            return Err(Error::InvalidArgument(format!(
96                "jaccl: rank {rank} is not a rank of a group of {size}"
97            )));
98        }
99        Ok(Self {
100            rank,
101            devices,
102            coordinator: None,
103            prefer_ring: false,
104        })
105    }
106
107    /// The coordinator's `host:port` (rank 0 listens on it during initialisation).
108    pub fn coordinator(mut self, address: impl Into<String>) -> Self {
109        self.coordinator = Some(address.into());
110        self
111    }
112
113    /// Prefers the ring topology when the cabling allows both.
114    pub fn prefer_ring(mut self, prefer: bool) -> Self {
115        self.prefer_ring = prefer;
116        self
117    }
118
119    /// This process's rank.
120    pub fn rank(&self) -> usize {
121        self.rank
122    }
123
124    /// The number of ranks.
125    pub fn size(&self) -> usize {
126        self.devices.len()
127    }
128
129    /// The configuration from MLX's environment, `None` when it is not set: `JACCL_RANK` or
130    /// `MLX_RANK`, `JACCL_IBV_DEVICES` or `MLX_IBV_DEVICES` (the device file),
131    /// `JACCL_COORDINATOR` or `MLX_JACCL_COORDINATOR`, and `JACCL_RING` or `MLX_JACCL_RING`
132    /// (set: prefer the ring), as `mlx.launch --backend jaccl` writes them.
133    pub fn from_env() -> Result<Option<Self>> {
134        Self::from_vars(|name| std::env::var(name).ok())
135    }
136
137    fn from_vars(var: impl Fn(&str) -> Option<String>) -> Result<Option<Self>> {
138        let env = |names: &[&str]| names.iter().find_map(|n| var(n));
139        let (Some(file), Some(rank)) = (
140            env(&["JACCL_IBV_DEVICES", "MLX_IBV_DEVICES"]),
141            env(&["JACCL_RANK", "MLX_RANK"]),
142        ) else {
143            return Ok(None);
144        };
145        let text = std::fs::read_to_string(&file)
146            .map_err(|e| Error::Io(format!("jaccl: reading the device file {file}: {e}")))?;
147        let rank = rank
148            .trim()
149            .parse()
150            .map_err(|_| Error::InvalidArgument(format!("jaccl: rank {rank:?}")))?;
151        let mut config = Self::new(rank, parse_devices(&text)?)?
152            .prefer_ring(env(&["JACCL_RING", "MLX_JACCL_RING"]).is_some());
153        config.coordinator = env(&["JACCL_COORDINATOR", "MLX_JACCL_COORDINATOR"]);
154        Ok(Some(config))
155    }
156
157    /// The first rule a mesh breaks, or `None` for a valid mesh (MLX's `is_valid_mesh`).
158    fn mesh_error(&self) -> Option<String> {
159        let size = self.size();
160        if size < 2 {
161            return Some("a group needs at least 2 ranks".into());
162        }
163        for (src, row) in self.devices.iter().enumerate() {
164            for (dst, cables) in row.iter().enumerate() {
165                if src == dst && !cables.is_empty() {
166                    return Some(format!("rank {src} lists a cable to itself"));
167                }
168                if src != dst && cables.is_empty() {
169                    return Some(format!("rank {src} → {dst}: no cable in the device file"));
170                }
171            }
172        }
173        None
174    }
175
176    /// The first rule a ring breaks, or `None` for a valid ring (MLX's `is_valid_ring`).
177    fn ring_error(&self) -> Option<String> {
178        let size = self.size();
179        if size < 2 {
180            return Some("a group needs at least 2 ranks".into());
181        }
182        let cables = self.devices[0][1].len();
183        if cables == 0 {
184            return Some("rank 0 → 1: no cable in the device file".into());
185        }
186        if cables > MAX_CABLES {
187            return Some(format!(
188                "the ring supports up to {MAX_CABLES} cables per direction, rank 0 → 1 has {cables}"
189            ));
190        }
191        for src in 0..size {
192            for dst in [(src + size - 1) % size, (src + 1) % size] {
193                let n = self.devices[src][dst].len();
194                if n != cables {
195                    return Some(format!(
196                        "rank {src} → {dst}: {n} cables, rank 0 → 1 has {cables}"
197                    ));
198                }
199            }
200        }
201        None
202    }
203
204    /// The topology MLX's `init` selects: the ring when preferred and valid, else the mesh
205    /// when valid, else the ring when valid.
206    pub fn topology(&self) -> Result<Topology> {
207        let (mesh, ring) = (self.mesh_error(), self.ring_error());
208        let as_ring = || Topology::Ring {
209            cables: self.devices[0][1].len(),
210        };
211        match (mesh, ring) {
212            (_, None) if self.prefer_ring => Ok(as_ring()),
213            // MLX selects a valid mesh of any size and then rejects more than 8 ranks.
214            (None, _) if self.size() > MAX_PEERS => Err(Error::InvalidArgument(format!(
215                "jaccl: the mesh supports up to {MAX_PEERS} ranks, the file has {}",
216                self.size()
217            ))),
218            (None, _) => Ok(Topology::Mesh {
219                wires: (0..self.size())
220                    .flat_map(|s| {
221                        (0..self.size())
222                            .filter(move |&d| d != s)
223                            .map(move |d| (s, d))
224                    })
225                    .map(|(s, d)| self.devices[s][d].len())
226                    .min()
227                    .unwrap_or(1),
228            }),
229            (Some(_), None) => Ok(as_ring()),
230            (Some(mesh), Some(ring)) => Err(Error::InvalidArgument(format!(
231                "jaccl: the device file defines neither a mesh ({mesh}) nor a ring ({ring})"
232            ))),
233        }
234    }
235
236    /// The mesh's devices of this rank: per cable, per peer (empty for itself).
237    pub(crate) fn mesh_devices(&self, wires: usize) -> Vec<Vec<String>> {
238        (0..wires)
239            .map(|w| {
240                (0..self.size())
241                    .map(|p| {
242                        if p == self.rank {
243                            String::new()
244                        } else {
245                            self.devices[self.rank][p][w].clone()
246                        }
247                    })
248                    .collect()
249            })
250            .collect()
251    }
252
253    /// The ring's devices of this rank: to the left and to the right neighbour.
254    pub(crate) fn ring_devices(&self) -> (Vec<String>, Vec<String>) {
255        let size = self.size();
256        let (left, right) = ((self.rank + size - 1) % size, (self.rank + 1) % size);
257        (
258            self.devices[self.rank][left].clone(),
259            self.devices[self.rank][right].clone(),
260        )
261    }
262}
263
264#[cfg(test)]
265mod tests {
266    use super::*;
267
268    const MESH4: &str = r#"[
269        [null, "rdma_en5", "rdma_en4", "rdma_en3"],
270        ["rdma_en5", null, "rdma_en3", "rdma_en4"],
271        ["rdma_en4", "rdma_en3", null, "rdma_en5"],
272        ["rdma_en3", "rdma_en4", "rdma_en5", null]
273    ]"#;
274
275    #[test]
276    fn reads_mlx_device_files() {
277        let d = parse_devices(MESH4).unwrap();
278        assert_eq!(d[1][2], ["rdma_en3"]);
279        assert!(d[2][2].is_empty());
280        let c = Config::new(2, d).unwrap();
281        assert_eq!(c.topology().unwrap(), Topology::Mesh { wires: 1 });
282        assert_eq!(
283            c.clone().prefer_ring(true).topology().unwrap(),
284            Topology::Ring { cables: 1 }
285        );
286        assert_eq!(
287            c.mesh_devices(1),
288            [vec!["rdma_en4", "rdma_en3", "", "rdma_en5"]]
289        );
290        assert_eq!(
291            c.ring_devices(),
292            (vec!["rdma_en3".into()], vec!["rdma_en5".into()])
293        );
294        // Two cables between two Macs.
295        let two = r#"[[null, ["rdma_en2", "rdma_en3"]], [["rdma_en2", "rdma_en3"], null]]"#;
296        let c = Config::new(0, parse_devices(two).unwrap()).unwrap();
297        assert_eq!(c.topology().unwrap(), Topology::Mesh { wires: 2 });
298        // A ring of four: neighbours only.
299        let ring = r#"[[null, "a", null, "b"], ["a", null, "c", null],
300                       [null, "c", null, "d"], ["b", null, "d", null]]"#;
301        let c = Config::new(1, parse_devices(ring).unwrap()).unwrap();
302        assert_eq!(c.topology().unwrap(), Topology::Ring { cables: 1 });
303    }
304
305    #[test]
306    fn reads_mlx_and_jaccl_variables() {
307        let file = std::env::temp_dir().join(format!("nuvec-jaccl-{}.json", std::process::id()));
308        std::fs::write(&file, MESH4).unwrap();
309        let path = file.to_str().unwrap().to_string();
310        let vars = |pairs: &[(&str, &str)]| {
311            let pairs: Vec<(String, String)> = pairs
312                .iter()
313                .map(|(a, b)| (a.to_string(), b.to_string()))
314                .collect();
315            move |n: &str| pairs.iter().find(|(k, _)| k == n).map(|(_, v)| v.clone())
316        };
317        assert_eq!(Config::from_vars(vars(&[])).unwrap(), None);
318        let c = Config::from_vars(vars(&[
319            ("MLX_IBV_DEVICES", &path),
320            ("MLX_RANK", "1"),
321            ("JACCL_RANK", "3"),
322            ("MLX_JACCL_COORDINATOR", "10.0.0.1:5000"),
323            ("MLX_JACCL_RING", "1"),
324        ]))
325        .unwrap()
326        .unwrap();
327        assert_eq!((c.rank(), c.size(), c.prefer_ring), (3, 4, true));
328        assert_eq!(c.coordinator.as_deref(), Some("10.0.0.1:5000"));
329        assert!(Config::from_vars(vars(&[("MLX_IBV_DEVICES", &path), ("MLX_RANK", "x")])).is_err());
330        std::fs::remove_file(file).unwrap();
331    }
332
333    #[test]
334    fn invalid_files_name_the_entry() {
335        for (text, want) in [
336            ("[[null, \"a\"]]", "rank 0 lists 2 peers"),
337            ("[[null, 3], [\"a\", null]]", "rank 0 → 1"),
338            ("{}", "array of arrays"),
339        ] {
340            let e = parse_devices(text).unwrap_err().to_string();
341            assert!(e.contains(want), "{text}: {e}");
342        }
343        let self_cable = parse_devices(r#"[["x", "a"], ["a", null]]"#).unwrap();
344        // Not a mesh; MLX's ring rule reads only the neighbours' entries.
345        let c = Config::new(0, self_cable).unwrap();
346        assert!(
347            c.mesh_error()
348                .unwrap()
349                .contains("rank 0 lists a cable to itself")
350        );
351        assert_eq!(c.topology().unwrap(), Topology::Ring { cables: 1 });
352        let missing =
353            parse_devices(r#"[[null, "a", null], ["a", null, null], [null, null, null]]"#).unwrap();
354        let e = Config::new(0, missing)
355            .unwrap()
356            .topology()
357            .unwrap_err()
358            .to_string();
359        assert!(
360            e.contains("rank 0 → 2: no cable") && e.contains("rank 0 → 2: 0 cables"),
361            "{e}"
362        );
363        let nine: Vec<Vec<Vec<String>>> = (0..9)
364            .map(|s| {
365                (0..9)
366                    .map(|d| if s == d { vec![] } else { vec!["x".into()] })
367                    .collect()
368            })
369            .collect();
370        let e = Config::new(0, nine)
371            .unwrap()
372            .topology()
373            .unwrap_err()
374            .to_string();
375        assert!(e.contains("up to 8 ranks"), "{e}");
376        let five = r#"[[null, ["a","b","c","d","e"]], [["a","b","c","d","e"], null]]"#;
377        let c = Config::new(0, parse_devices(five).unwrap()).unwrap();
378        assert_eq!(c.topology().unwrap(), Topology::Mesh { wires: 5 });
379        let e = c.prefer_ring(true).ring_error().unwrap();
380        assert!(e.contains("up to 4 cables"), "{e}");
381        assert!(Config::new(3, parse_devices(MESH4).unwrap()).is_ok());
382        assert!(Config::new(4, parse_devices(MESH4).unwrap()).is_err());
383    }
384}
nuvec_distributed/jaccl/config.rs

nuvec_distributed/jaccl/
config.rs
