nuvec_distributed/
lib.rs

1//! Distributed backends for nuvec: MLX's TCP ring, and MLX's JACCL (RDMA over
2//! Thunderbolt 5; experimental, feature `jaccl`).
3//!
4//! [`init`] reads MLX's launch environment, connects a backend and makes it the default
5//! group of `nuvec::distributed`'s collectives, as `mx.distributed.init()`. The ring reads
6//! `MLX_HOSTFILE` (a JSON list with one entry per rank of `"host:port"` addresses) and
7//! `MLX_RANK`; JACCL reads the variables `mlx.launch --backend jaccl` writes
8//! (`jaccl::Config::from_env`). [`init_backend`] names the backend. Every rank runs the
9//! same program; the processes may share one Mac (over loopback) or run on several (over
10//! Ethernet or Thunderbolt IP for the ring, Thunderbolt 5 cables for JACCL).
11//!
12//! ```no_run
13//! use nuvec_core::{array, distributed};
14//!
15//! # fn main() -> nuvec_core::Result<()> {
16//! let group = nuvec_distributed::init(false)?; // size 1 without a hostfile
17//! let x = array![1.0, 2.0];
18//! let total = distributed::all_sum(&x, None); // the sum over every rank
19//! println!("rank {} of {}: {:?}", group.rank(), group.size(), total.to_vec::<f32>()?);
20//! # Ok(())
21//! # }
22//! ```
23//!
24//! **Stability:** no promise when this crate is depended on directly; depend on `nuvec`,
25//! whose re-exports are semver-tracked.
26
27#![deny(unsafe_code)]
28
29mod auth;
30mod backend;
31mod hostfile;
32#[cfg(feature = "jaccl")]
33pub mod jaccl;
34mod json;
35mod keepalive;
36mod pool;
37mod reduce;
38mod ring;
39mod socket;
40
41use std::sync::Arc;
42use std::time::Duration;
43
44use nuvec_core::Result;
45use nuvec_core::distributed::Group;
46
47pub use backend::{Backend, init_backend, is_available};
48
49/// Options of [`init_with`].
50#[derive(Clone, Debug)]
51#[non_exhaustive]
52pub struct Options {
53    /// How long connecting and every transfer may wait for a peer before failing with an
54    /// error naming it. Default 30 s (`NUVEC_DISTRIBUTED_TIMEOUT`, in seconds, for
55    /// [`init`]).
56    pub timeout: Duration,
57    key: auth::Secret,
58}
59
60impl Default for Options {
61    fn default() -> Self {
62        Self {
63            timeout: Duration::from_secs(30),
64            key: auth::Secret::default(),
65        }
66    }
67}
68
69impl Options {
70    /// Sets the timeout.
71    pub fn timeout(mut self, timeout: Duration) -> Self {
72        self.timeout = timeout;
73        self
74    }
75
76    /// Authenticates every connection with this pre-shared secret (at least 16 bytes, the
77    /// same on every rank; `NUVEC_DISTRIBUTED_KEY` for [`init`]): a host without it cannot
78    /// join the group. The traffic itself is not encrypted; on an untrusted network, run the
79    /// group over a VPN.
80    pub fn key(mut self, secret: impl AsRef<[u8]>) -> Self {
81        self.key = auth::Secret(Some(secret.as_ref().to_vec()));
82        self
83    }
84}
85
86/// Initialises a backend from MLX's environment and makes it the default group, as
87/// `mx.distributed.init()`: [`init_backend`] with [`Backend::Any`].
88pub fn init(strict: bool) -> Result<Group> {
89    init_backend(strict, Backend::Any)
90}
91
92/// Connects rank `rank` of the ring over `nodes` (per rank its addresses, one connection
93/// per address, as MLX's hostfile) and returns its group (not made the default). A single
94/// rank is the group of this process alone.
95pub fn init_with(nodes: &[Vec<String>], rank: usize, options: &Options) -> Result<Group> {
96    if nodes.len() <= 1 {
97        return Ok(Group::single());
98    }
99    let key = options.key.key()?;
100    let ring = ring::Ring::connect(nodes, rank, options.timeout, key.as_ref())?;
101    Ok(Group::new(Arc::new(ring)))
102}
nuvec_distributed/lib.rs

nuvec_distributed/
lib.rs
