nuvec_distributed/jaccl/
mod.rs

1//! **Experimental:** MLX's JACCL backend, RDMA over Thunderbolt 5 (macOS 26.2+), behind the
2//! `jaccl` feature. It is unverified on Thunderbolt 5 hardware; its algorithms are
3//! tested on an emulated fabric.
4//!
5//! RDMA must be enabled once per Mac in recoveryOS (`rdma_ctl enable`), and every pair of
6//! ranks cabled (the mesh) or every rank cabled to its neighbours (the ring). The group is
7//! described by MLX's device file and started as MLX's (`mlx.launch --backend jaccl` or
8//! `--backend jaccl-ring`); see [`Config::from_env`]. [`devices`] lists this Mac's RDMA
9//! devices to check a setup.
10
11mod config;
12#[cfg(test)]
13mod emulated;
14mod endpoint;
15mod frame;
16mod group;
17mod mesh;
18mod ring;
19mod side_channel;
20#[cfg(test)]
21mod tests;
22#[cfg(target_os = "macos")]
23mod verbs;
24
25use std::net::Ipv4Addr;
26use std::sync::{Arc, Mutex};
27use std::time::Duration;
28
29use nuvec_core::distributed::Group;
30use nuvec_core::{Error, Result};
31
32pub use config::{Config, Topology, parse_devices};
33#[cfg(target_os = "macos")]
34#[doc(hidden)]
35pub use verbs::layout as verbs_layout;
36
37use side_channel::{AllGather, SideChannel, Star};
38
39/// Options of [`init_with`].
40#[derive(Clone)]
41#[non_exhaustive]
42pub struct Options {
43    /// How long connecting and every transfer may wait for a peer before failing with an
44    /// error naming it. Default 30 s (`NUVEC_DISTRIBUTED_TIMEOUT`, in seconds, for `init`).
45    pub timeout: Duration,
46    side_channel: Option<AllGather>,
47    key: crate::auth::Secret,
48}
49
50impl std::fmt::Debug for Options {
51    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
52        f.debug_struct("Options")
53            .field("timeout", &self.timeout)
54            .field(
55                "side_channel",
56                &self.side_channel.as_ref().map(|_| "custom"),
57            )
58            .field("key", &self.key)
59            .finish()
60    }
61}
62
63impl Default for Options {
64    fn default() -> Self {
65        Self {
66            timeout: Duration::from_secs(30),
67            side_channel: None,
68            key: crate::auth::Secret::default(),
69        }
70    }
71}
72
73impl Options {
74    /// Sets the timeout.
75    pub fn timeout(mut self, timeout: Duration) -> Self {
76        self.timeout = timeout;
77        self
78    }
79
80    /// Authenticates the side channel's connections with this pre-shared secret (at least 16
81    /// bytes, the same on every rank; `NUVEC_DISTRIBUTED_KEY` for `init`).
82    pub fn key(mut self, secret: impl AsRef<[u8]>) -> Self {
83        self.key = crate::auth::Secret(Some(secret.as_ref().to_vec()));
84        self
85    }
86
87    /// Exchanges the connection details through `all_gather` instead of the coordinator's
88    /// TCP star, as MLX's `all_gather_factory`: it takes this rank's bytes and returns every
89    /// rank's (equal lengths), concatenated in rank order.
90    pub fn side_channel(
91        mut self,
92        all_gather: impl Fn(&[u8]) -> Result<Vec<u8>> + Send + Sync + 'static,
93    ) -> Self {
94        self.side_channel = Some(Arc::new(all_gather));
95        self
96    }
97}
98
99/// Connects this rank's JACCL group (not made the default group).
100pub fn init_with(config: &Config, options: &Options) -> Result<Group> {
101    #[cfg(target_os = "macos")]
102    {
103        let fabric = verbs::Verbs::load()?;
104        let comm = connect(config, &fabric, options, frame::large_frames())?;
105        Ok(Group::new(Arc::new(comm)))
106    }
107    #[cfg(not(target_os = "macos"))]
108    {
109        let _ = (config, options);
110        Err(unavailable())
111    }
112}
113
114/// Connects over any fabric (the RDMA library, or the emulated one in tests).
115fn connect(
116    config: &Config,
117    fabric: &dyn endpoint::Fabric,
118    options: &Options,
119    large: bool,
120) -> Result<group::Jaccl> {
121    let side = match &options.side_channel {
122        Some(f) => SideChannel::Custom(f.clone()),
123        None => {
124            let coordinator = config.coordinator.as_deref().ok_or_else(|| {
125                Error::InvalidArgument(
126                    "jaccl: no coordinator (JACCL_COORDINATOR or MLX_JACCL_COORDINATOR) and no \
127                     side channel"
128                        .into(),
129                )
130            })?;
131            SideChannel::Tcp(Mutex::new(Star::connect(
132                config.rank(),
133                config.size(),
134                coordinator,
135                options.timeout,
136                options.key.key()?.as_ref(),
137            )?))
138        }
139    };
140    group::connect(config, fabric, side, large, options.timeout)
141}
142
143#[cfg(not(target_os = "macos"))]
144pub(crate) fn unavailable() -> Error {
145    Error::Backend(
146        "jaccl: RDMA is not available (librdma not found, or no RDMA device: enable RDMA over \
147         Thunderbolt with `rdma_ctl enable` in recoveryOS)"
148            .into(),
149    )
150}
151
152/// Whether RDMA can run here: the system library loads and lists at least one device.
153pub fn is_available() -> bool {
154    devices().is_ok_and(|d| !d.is_empty())
155}
156
157/// A port's state.
158#[derive(Clone, Copy, Debug, PartialEq, Eq)]
159#[non_exhaustive]
160pub enum PortState {
161    /// No link.
162    Down,
163    /// Initialising.
164    Init,
165    /// Armed.
166    Armed,
167    /// Ready.
168    Active,
169    /// Ready, deferred.
170    ActiveDefer,
171    /// A state the library reports beyond these.
172    Unknown(i32),
173}
174
175/// One of this Mac's RDMA devices (a Thunderbolt port).
176#[derive(Clone, Debug, PartialEq, Eq)]
177#[non_exhaustive]
178pub struct RdmaDevice {
179    /// The device name, as in the device file (e.g. `rdma_en3`).
180    pub name: String,
181    /// Port 1's state.
182    pub port_state: PortState,
183    /// The IPv4 address of the port's IPv4-mapped GID; `None` means the Thunderbolt
184    /// interface has no IPv4 address yet, which connecting needs.
185    pub ipv4: Option<Ipv4Addr>,
186}
187
188/// This Mac's RDMA devices (empty when RDMA is not enabled); an error when the system
189/// library cannot be loaded.
190pub fn devices() -> Result<Vec<RdmaDevice>> {
191    #[cfg(target_os = "macos")]
192    {
193        verbs::devices()
194    }
195    #[cfg(not(target_os = "macos"))]
196    {
197        Err(unavailable())
198    }
199}
200
201/// Runs `f` for every item, the first on this thread and the others on one thread each
202/// (one per cable), and returns the first error.
203pub(crate) fn on_wires<T: Send>(
204    items: Vec<T>,
205    f: impl Fn(usize, T) -> Result<()> + Sync,
206) -> Result<()> {
207    let mut items = items.into_iter().enumerate();
208    let Some((_, first)) = items.next() else {
209        return Ok(());
210    };
211    if items.len() == 0 {
212        return f(0, first);
213    }
214    std::thread::scope(|s| {
215        let f = &f;
216        let handles: Vec<_> = items.map(|(w, item)| s.spawn(move || f(w, item))).collect();
217        let first = f(0, first);
218        let rest: Vec<Result<()>> = handles
219            .into_iter()
220            .map(|h| {
221                h.join().unwrap_or_else(|_| {
222                    Err(Error::Backend("jaccl: a cable's thread panicked".into()))
223                })
224            })
225            .collect();
226        first.and(rest.into_iter().collect())
227    })
228}
nuvec_distributed/jaccl/mod.rs

nuvec_distributed/jaccl/
mod.rs
