rlx_driver/
node.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! One-call node bring-up.
6//!
7//! Joining a distributed job used to mean hand-wiring a transport (mesh vs.
8//! star), resolving peer addresses, and wrapping the result in an
9//! `Arc<ProcessGroup>`. [`Node`] collapses that to a builder:
10//!
11//! ```no_run
12//! # use rlx_driver::{Node, Topology};
13//! # fn main() -> Result<(), Box<dyn std::error::Error>> {
14//! // From env (RANK/WORLD/PEERS or DISCOVER/TOPOLOGY):
15//! let group = Node::from_env()?.connect()?;
16//! // Or explicitly:
17//! let group = Node::new(/*rank*/ 1, /*world*/ 2)
18//!     .topology(Topology::Star)          // worker dials the coordinator
19//!     .peers(["10.0.0.1:29500"])?        // just the coordinator address
20//!     .connect()?;
21//! let _ = group;
22//! # Ok(()) }
23//! ```
24//!
25//! The returned [`ProcessGroup`] carries every collective (`all_reduce`,
26//! `all_reduce_typed`, `federated_average`, `broadcast`, …). This is the
27//! entry point a model runner or an edge worker calls to get on the mesh.
28
29use crate::net::{DEFAULT_HEAP_BYTES, NetTransport, TcpTransport};
30use crate::transport::ProcessGroup;
31use std::collections::BTreeMap;
32use std::io;
33use std::net::{IpAddr, Ipv4Addr, SocketAddr, TcpListener, ToSocketAddrs, UdpSocket};
34use std::sync::Arc;
35use std::sync::atomic::{AtomicBool, Ordering};
36use std::time::Duration;
37
38/// Wire topology.
39#[derive(Clone, Copy, Debug, PartialEq, Eq)]
40pub enum Topology {
41    /// Full-mesh TCP — every rank connects to higher ranks and accepts from
42    /// lower. Needed for peer-to-peer collectives (ring all-reduce, etc.).
43    Mesh,
44    /// Star — the coordinator (rank 0) listens; workers dial in and need **no
45    /// inbound port**. NAT/mobile/Pi-behind-a-router friendly; the shape for
46    /// the coordinator/worker ship-graph model and federated averaging.
47    Star,
48}
49
50#[derive(Clone, Debug)]
51enum PeerSpec {
52    Static(Vec<SocketAddr>),
53    Discover { disc_port: u16, data_base: u16 },
54}
55
56/// Builder that turns minimal config into a connected [`ProcessGroup`].
57pub struct Node {
58    rank: u32,
59    world: u32,
60    peers: PeerSpec,
61    topology: Topology,
62    heap_bytes: usize,
63    /// For star discovery from behind NAT (e.g. a Docker/QEMU node): the host
64    /// to *unicast* the discovery query to, since broadcast can't cross the
65    /// bridge. `None` = LAN broadcast discovery.
66    discover_host: Option<String>,
67}
68
69impl Node {
70    /// A node at `rank` of `world`, full-mesh, default heap. Set peers with
71    /// [`Self::peers`] or [`Self::discover`] before [`Self::connect`].
72    pub fn new(rank: u32, world: u32) -> Self {
73        Self {
74            rank,
75            world,
76            peers: PeerSpec::Static(Vec::new()),
77            topology: Topology::Mesh,
78            heap_bytes: DEFAULT_HEAP_BYTES,
79            discover_host: None,
80        }
81    }
82
83    /// Unicast the discovery query to `host` instead of broadcasting — for a
84    /// worker behind NAT (Docker/QEMU) that can't hear the coordinator's LAN
85    /// broadcast. It learns the coordinator's port from the reply and dials
86    /// `host:port`. Only used with a discovering [`Topology::Star`] worker.
87    pub fn discover_via(mut self, host: impl Into<String>) -> Self {
88        self.discover_host = Some(host.into());
89        self
90    }
91
92    /// Static peer addresses, one per rank (mesh) or just the coordinator
93    /// (star, `peers[0]`). Accepts anything `ToSocketAddrs` (host:port or ip:port).
94    pub fn peers<A: ToSocketAddrs>(
95        mut self,
96        addrs: impl IntoIterator<Item = A>,
97    ) -> io::Result<Self> {
98        let mut v = Vec::new();
99        for a in addrs {
100            let sa = a.to_socket_addrs()?.next().ok_or_else(|| {
101                io::Error::new(io::ErrorKind::InvalidInput, "peer resolved to no address")
102            })?;
103            v.push(sa);
104        }
105        self.peers = PeerSpec::Static(v);
106        Ok(self)
107    }
108
109    /// Discover peers by UDP broadcast (zero-config on a shared LAN): each node
110    /// announces `rank → ip:(data_base+rank)` on `disc_port` until all `world`
111    /// addresses are known.
112    pub fn discover(mut self, disc_port: u16, data_base: u16) -> Self {
113        self.peers = PeerSpec::Discover {
114            disc_port,
115            data_base,
116        };
117        self
118    }
119
120    pub fn topology(mut self, t: Topology) -> Self {
121        self.topology = t;
122        self
123    }
124
125    /// Per-rank symmetric-heap size in bytes (default [`DEFAULT_HEAP_BYTES`]).
126    pub fn heap_bytes(mut self, n: usize) -> Self {
127        self.heap_bytes = n;
128        self
129    }
130
131    pub fn rank(&self) -> u32 {
132        self.rank
133    }
134    pub fn world(&self) -> u32 {
135        self.world
136    }
137
138    /// Read node config from the environment:
139    ///   `RANK`, `WORLD`                       — this rank / total ranks
140    ///   `PEERS=host:port,host:port,…`         — static addresses, or
141    ///   `DISCOVER=1` (+ `DISC_PORT`/`DATA_PORT`) — UDP auto-discovery
142    ///   `TOPOLOGY=mesh|star` (`DIAL_OUT=1` = star)
143    ///   `HEAP_MB`                             — per-rank heap (optional)
144    pub fn from_env() -> Result<Self, String> {
145        let var = |k: &str| std::env::var(k).ok();
146        let rank: u32 = var("RANK")
147            .as_deref()
148            .unwrap_or("0")
149            .parse()
150            .map_err(|_| "RANK must be an integer".to_string())?;
151        let world: u32 = var("WORLD")
152            .as_deref()
153            .unwrap_or("1")
154            .parse()
155            .map_err(|_| "WORLD must be an integer".to_string())?;
156        let mut node = Node::new(rank, world);
157
158        let star =
159            var("TOPOLOGY").as_deref() == Some("star") || var("DIAL_OUT").is_some_and(|v| v != "0");
160        node = node.topology(if star { Topology::Star } else { Topology::Mesh });
161
162        if let Some(mb) = var("HEAP_MB").and_then(|v| v.parse::<usize>().ok()) {
163            node = node.heap_bytes(mb << 20);
164        }
165
166        if var("DISCOVER").is_some_and(|v| v != "0") {
167            let dp = var("DISC_PORT")
168                .and_then(|v| v.parse().ok())
169                .unwrap_or(29600);
170            let db = var("DATA_PORT")
171                .and_then(|v| v.parse().ok())
172                .unwrap_or(29500);
173            node = node.discover(dp, db);
174            if let Some(h) = var("DISCOVER_HOST") {
175                node = node.discover_via(h);
176            }
177        } else {
178            let peers = var("PEERS").unwrap_or_else(|| "127.0.0.1:29500,127.0.0.1:29501".into());
179            let addrs: Vec<String> = peers.split(',').map(|s| s.trim().to_string()).collect();
180            node = node
181                .peers(addrs.iter().map(String::as_str))
182                .map_err(|e| format!("PEERS: {e}"))?;
183        }
184        Ok(node)
185    }
186
187    /// Establish the transport and return a ready [`ProcessGroup`].
188    pub fn connect(self) -> io::Result<Arc<ProcessGroup>> {
189        // Star + discovery: the coordinator *announces* itself and workers
190        // *find* it. Unlike mesh discovery (every rank must hear every other),
191        // the coordinator never waits to hear the workers — so a NAT'd node
192        // that can't broadcast (Docker/QEMU) doesn't stall the group; it just
193        // unicasts its query via `discover_via`.
194        if self.topology == Topology::Star
195            && let PeerSpec::Discover {
196                disc_port,
197                data_base,
198            } = &self.peers
199        {
200            let (disc_port, data_base) = (*disc_port, *data_base);
201            let transport = if self.rank == 0 {
202                let listener = TcpListener::bind(("0.0.0.0", data_base))?;
203                let stop = Arc::new(AtomicBool::new(false));
204                let ann = stop.clone();
205                std::thread::spawn(move || announce_coordinator(data_base, disc_port, &ann));
206                #[cfg(feature = "mdns")]
207                let _mdns = mdns_advertise(data_base); // advertise while we listen
208                let t = NetTransport::coordinator_listen(self.world, listener, self.heap_bytes);
209                stop.store(true, Ordering::SeqCst);
210                t?
211            } else {
212                let coord = discover_coordinator(disc_port, self.discover_host.as_deref())?;
213                NetTransport::worker_dial(self.rank, self.world, coord, self.heap_bytes)?
214            };
215            return Ok(Arc::new(ProcessGroup::new(Arc::new(transport))));
216        }
217
218        // Otherwise resolve a concrete peer list (mesh discovery collects all).
219        let peers = match self.peers {
220            PeerSpec::Static(v) => v,
221            PeerSpec::Discover {
222                disc_port,
223                data_base,
224            } => discover_peers(self.rank, self.world, disc_port, data_base),
225        };
226        let transport = match self.topology {
227            Topology::Mesh => {
228                if peers.len() != self.world as usize {
229                    return Err(io::Error::new(
230                        io::ErrorKind::InvalidInput,
231                        format!(
232                            "mesh needs WORLD={} peer addresses, got {}",
233                            self.world,
234                            peers.len()
235                        ),
236                    ));
237                }
238                TcpTransport::bind(self.rank, self.world, peers, self.heap_bytes)?
239            }
240            Topology::Star => {
241                let coord = *peers.first().ok_or_else(|| {
242                    io::Error::new(
243                        io::ErrorKind::InvalidInput,
244                        "star needs the coordinator address (peers[0])",
245                    )
246                })?;
247                if self.rank == 0 {
248                    let listener = TcpListener::bind(coord)?;
249                    NetTransport::coordinator_listen(self.world, listener, self.heap_bytes)?
250                } else {
251                    NetTransport::worker_dial(self.rank, self.world, coord, self.heap_bytes)?
252                }
253            }
254        };
255        Ok(Arc::new(ProcessGroup::new(Arc::new(transport))))
256    }
257}
258
259/// The default-route local IP (picks the outbound interface); loopback if the
260/// probe fails.
261pub fn local_ip() -> IpAddr {
262    UdpSocket::bind("0.0.0.0:0")
263        .and_then(|s| {
264            s.connect("8.8.8.8:80")?; // no packet sent — just selects the iface
265            s.local_addr()
266        })
267        .map(|a| a.ip())
268        .unwrap_or(IpAddr::V4(Ipv4Addr::LOCALHOST))
269}
270
271/// UDP-broadcast rendezvous: announce `rank → ip:(data_base+rank)` on
272/// `disc_port` and collect until all `world` addresses are known; return them
273/// sorted by rank. Zero-config peer discovery on a shared LAN.
274pub fn discover_peers(rank: u32, world: u32, disc_port: u16, data_base: u16) -> Vec<SocketAddr> {
275    let my_addr = SocketAddr::new(local_ip(), data_base + rank as u16);
276    let sock = UdpSocket::bind(("0.0.0.0", disc_port)).expect("discovery bind");
277    sock.set_broadcast(true).ok();
278    sock.set_read_timeout(Some(Duration::from_millis(150))).ok();
279    let bcast = SocketAddr::new(IpAddr::V4(Ipv4Addr::BROADCAST), disc_port);
280    let msg = format!("RLXDISC {rank} {my_addr}");
281
282    let mut peers: BTreeMap<u32, SocketAddr> = BTreeMap::new();
283    peers.insert(rank, my_addr);
284    let mut buf = [0u8; 256];
285    while (peers.len() as u32) < world {
286        let _ = sock.send_to(msg.as_bytes(), bcast);
287        if let Ok((n, _)) = sock.recv_from(&mut buf)
288            && let Ok(s) = std::str::from_utf8(&buf[..n])
289        {
290            let mut it = s.split_whitespace();
291            if it.next() == Some("RLXDISC")
292                && let (Some(r), Some(a)) = (it.next(), it.next())
293                && let (Ok(r), Ok(a)) = (r.parse::<u32>(), a.parse::<SocketAddr>())
294            {
295                peers.insert(r, a);
296            }
297        }
298    }
299    // Brief drain so late joiners still hear our announcement.
300    for _ in 0..5 {
301        let _ = sock.send_to(msg.as_bytes(), bcast);
302    }
303    peers.into_values().collect()
304}
305
306fn parse_rlxport(b: &[u8]) -> Option<u16> {
307    let s = std::str::from_utf8(b).ok()?;
308    let mut it = s.split_whitespace();
309    if it.next() != Some("RLXPORT") {
310        return None;
311    }
312    it.next()?.parse().ok()
313}
314
315/// Coordinator side of star discovery: a **pure unicast responder** on
316/// `disc_port` — a worker sends `RLXQ`, we reply `RLXPORT <data_port>`. Runs
317/// until `stop`. Deliberately does NOT broadcast on this socket: a coordinator
318/// that both beacons and receives on one socket starves incoming queries
319/// (its own zero-latency loopback beacons always win the recv race). LAN
320/// zero-config is handled by mDNS instead ([`mdns_advertise`]); this responder
321/// covers a worker that reaches us by name/IP (LAN, `host.docker.internal`, or
322/// a Tailscale MagicDNS name — anything routable).
323pub fn announce_coordinator(data_port: u16, disc_port: u16, stop: &AtomicBool) {
324    let Ok(sock) = UdpSocket::bind(("0.0.0.0", disc_port)) else {
325        return;
326    };
327    sock.set_read_timeout(Some(Duration::from_millis(200))).ok();
328    let msg = format!("RLXPORT {data_port}");
329    let mut buf = [0u8; 64];
330    while !stop.load(Ordering::SeqCst) {
331        if let Ok((n, from)) = sock.recv_from(&mut buf)
332            && buf[..n].starts_with(b"RLXQ")
333        {
334            let _ = sock.send_to(msg.as_bytes(), from); // reply to the querier
335        }
336    }
337}
338
339/// Worker side of star discovery: return the coordinator's dial address.
340/// `host = Some(h)` (NAT/Docker) unicasts a query to `h:disc_port` and dials
341/// `h:port` from the reply; `None` listens for the LAN broadcast and dials the
342/// announcing host's own IP.
343pub fn discover_coordinator(disc_port: u16, host: Option<&str>) -> io::Result<SocketAddr> {
344    let mut buf = [0u8; 64];
345    let no_addr = || {
346        io::Error::new(
347            io::ErrorKind::InvalidInput,
348            "discover host resolved to no address",
349        )
350    };
351    match host {
352        Some(h) => {
353            let target = (h, disc_port)
354                .to_socket_addrs()?
355                .next()
356                .ok_or_else(no_addr)?;
357            let sock = UdpSocket::bind(("0.0.0.0", 0))?;
358            sock.set_read_timeout(Some(Duration::from_millis(300))).ok();
359            loop {
360                sock.send_to(b"RLXQ", target)?;
361                if let Ok((n, _)) = sock.recv_from(&mut buf)
362                    && let Some(port) = parse_rlxport(&buf[..n])
363                {
364                    return (h, port).to_socket_addrs()?.next().ok_or_else(no_addr);
365                }
366            }
367        }
368        None => {
369            let _ = disc_port;
370            #[cfg(feature = "mdns")]
371            {
372                mdns_discover(Duration::from_secs(15)).ok_or_else(|| {
373                    io::Error::new(
374                        io::ErrorKind::NotFound,
375                        "mDNS: no _rlx-coord._udp coordinator found on the LAN",
376                    )
377                })
378            }
379            #[cfg(not(feature = "mdns"))]
380            Err(io::Error::new(
381                io::ErrorKind::InvalidInput,
382                "no coordinator address — set DISCOVER_HOST (LAN IP / host.docker.internal / \
383                 tailnet name) for unicast discovery, or build with the `mdns` feature for \
384                 zero-config LAN discovery",
385            ))
386        }
387    }
388}
389
390/// Advertise the coordinator as `_rlx-coord._udp.local.` so LAN workers can
391/// browse for it — zero-config, Bonjour/Avahi-compatible, and multicast-based
392/// (`224.0.0.251`), which switches/APs forward far more reliably than raw
393/// broadcast. Hold the returned daemon alive while listening.
394#[cfg(feature = "mdns")]
395pub fn mdns_advertise(data_port: u16) -> Option<mdns_sd::ServiceDaemon> {
396    use mdns_sd::{ServiceDaemon, ServiceInfo};
397    let mdns = ServiceDaemon::new().ok()?;
398    let ip = local_ip().to_string();
399    let props: &[(&str, &str)] = &[];
400    let info = ServiceInfo::new(
401        "_rlx-coord._udp.local.",
402        "coord",
403        "rlx-coord.local.",
404        ip.as_str(),
405        data_port,
406        props,
407    )
408    .ok()?;
409    mdns.register(info).ok()?;
410    Some(mdns)
411}
412
413/// Browse for a `_rlx-coord._udp.local.` coordinator and return its address.
414#[cfg(feature = "mdns")]
415pub fn mdns_discover(timeout: Duration) -> Option<SocketAddr> {
416    use mdns_sd::{ServiceDaemon, ServiceEvent};
417    use std::time::Instant;
418    let mdns = ServiceDaemon::new().ok()?;
419    let rx = mdns.browse("_rlx-coord._udp.local.").ok()?;
420    let deadline = Instant::now() + timeout;
421    while let Some(left) = deadline.checked_duration_since(Instant::now())
422        && let Ok(ev) = rx.recv_timeout(left)
423    {
424        if let ServiceEvent::ServiceResolved(info) = ev
425            && let Some(ip) = info.get_addresses().iter().next()
426        {
427            return Some(SocketAddr::new(*ip, info.get_port()));
428        }
429    }
430    None
431}
432
433#[cfg(test)]
434mod tests {
435    use super::*;
436    use crate::collective::ReduceKind;
437
438    #[test]
439    fn node_mesh_connect_and_all_reduce() {
440        // Grab two free loopback ports, then bring up a 2-rank mesh via the
441        // builder and check a collective round-trips.
442        let free = || {
443            let l = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).unwrap();
444            l.local_addr().unwrap()
445        };
446        let (a0, a1) = (free(), free());
447        let peers = vec![a0.to_string(), a1.to_string()];
448
449        let p1 = peers.clone();
450        let h = std::thread::spawn(move || {
451            let g = Node::new(1, 2)
452                .peers(p1.iter().map(String::as_str))
453                .unwrap()
454                .connect()
455                .unwrap();
456            let mut d = vec![2.0f32; 3];
457            g.all_reduce(&mut d, ReduceKind::Sum).unwrap();
458            assert_eq!(d, vec![3.0; 3]);
459            g.barrier().unwrap();
460        });
461        let g = Node::new(0, 2)
462            .peers(peers.iter().map(String::as_str))
463            .unwrap()
464            .connect()
465            .unwrap();
466        let mut d = vec![1.0f32; 3];
467        g.all_reduce(&mut d, ReduceKind::Sum).unwrap();
468        assert_eq!(d, vec![3.0; 3]); // 1 + 2
469        g.barrier().unwrap();
470        h.join().unwrap();
471    }
472}
rlx_driver/node.rs

rlx_driver/
node.rs
