nuvec_distributed/
lib.rs

1//! Distributed backends for nuvec: MLX's TCP ring.
2//!
3//! [`init`] reads MLX's launch environment (`MLX_HOSTFILE`, a JSON list with one entry per
4//! rank of `"host:port"` addresses, and `MLX_RANK`), connects the ring and makes it the
5//! default group of `nuvec::distributed`'s collectives. Every rank runs the same program;
6//! the processes may share one Mac (over loopback) or run on several (over Ethernet or
7//! Thunderbolt IP).
8//!
9//! ```no_run
10//! use nuvec_core::{array, distributed};
11//!
12//! # fn main() -> nuvec_core::Result<()> {
13//! let group = nuvec_distributed::init(false)?; // size 1 without a hostfile
14//! let x = array![1.0, 2.0];
15//! let total = distributed::all_sum(&x, None); // the sum over every rank
16//! println!("rank {} of {}: {:?}", group.rank(), group.size(), total.to_vec::<f32>()?);
17//! # Ok(())
18//! # }
19//! ```
20//!
21//! **Stability:** no promise when this crate is depended on directly; depend on `nuvec`,
22//! whose re-exports are semver-tracked.
23
24#![forbid(unsafe_code)]
25
26mod hostfile;
27mod ring;
28
29use std::sync::Arc;
30use std::time::Duration;
31
32use nuvec_core::distributed::{Group, set_default_group};
33use nuvec_core::{Error, Result};
34
35/// Options of [`init_with`].
36#[derive(Clone, Debug)]
37#[non_exhaustive]
38pub struct Options {
39    /// How long connecting and every transfer may wait for a peer before failing with an
40    /// error naming it. Default 30 s (`NUVEC_DISTRIBUTED_TIMEOUT`, in seconds, for
41    /// [`init`]).
42    pub timeout: Duration,
43}
44
45impl Default for Options {
46    fn default() -> Self {
47        Self {
48            timeout: Duration::from_secs(30),
49        }
50    }
51}
52
53impl Options {
54    /// Sets the timeout.
55    pub fn timeout(mut self, timeout: Duration) -> Self {
56        self.timeout = timeout;
57        self
58    }
59}
60
61/// Initialises the ring from MLX's environment (`MLX_HOSTFILE`, `MLX_RANK`) and makes it
62/// the default group, as `mx.distributed.init()`. Without the environment: an error if
63/// `strict`, else the group of this process alone.
64pub fn init(strict: bool) -> Result<Group> {
65    let (hostfile, rank) = (std::env::var("MLX_HOSTFILE"), std::env::var("MLX_RANK"));
66    let (Ok(hostfile), Ok(rank)) = (hostfile, rank) else {
67        if strict {
68            return Err(Error::InvalidArgument(
69                "distributed: MLX_HOSTFILE and MLX_RANK are not set".into(),
70            ));
71        }
72        return Ok(Group::single());
73    };
74    let text = std::fs::read_to_string(&hostfile)
75        .map_err(|e| Error::Io(format!("distributed: reading {hostfile}: {e}")))?;
76    let rank = rank
77        .trim()
78        .parse()
79        .map_err(|_| Error::InvalidArgument(format!("distributed: MLX_RANK {rank:?}")))?;
80    let mut options = Options::default();
81    if let Ok(secs) = std::env::var("NUVEC_DISTRIBUTED_TIMEOUT") {
82        let invalid = || {
83            Error::InvalidArgument(format!(
84                "distributed: NUVEC_DISTRIBUTED_TIMEOUT {secs:?} (seconds, a positive number)"
85            ))
86        };
87        let parsed: f64 = secs.trim().parse().map_err(|_| invalid())?;
88        options.timeout = Duration::try_from_secs_f64(parsed)
89            .ok()
90            .filter(|d| !d.is_zero())
91            .ok_or_else(invalid)?;
92    }
93    let group = init_with(&hostfile::parse(&text)?, rank, &options)?;
94    set_default_group(group.clone());
95    Ok(group)
96}
97
98/// Connects rank `rank` of the ring over `nodes` (per rank its addresses, one connection
99/// per address, as MLX's hostfile) and returns its group (not made the default). A single
100/// rank is the group of this process alone.
101pub fn init_with(nodes: &[Vec<String>], rank: usize, options: &Options) -> Result<Group> {
102    if nodes.len() <= 1 {
103        return Ok(Group::single());
104    }
105    let ring = ring::Ring::connect(nodes, rank, options.timeout)?;
106    Ok(Group::new(Arc::new(ring)))
107}
nuvec_distributed/lib.rs

nuvec_distributed/
lib.rs
