rlx_distributed/
launch.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//! Local multi-process launcher — run an N-rank cluster on one host over
6//! loopback TCP, with **no hand-written `hosts.json` and no N terminals**.
7//!
8//! A distributed binary has two roles. When launched with `--rank R` it is a
9//! **worker**: it forms the process group and runs its block. With no
10//! `--rank`, it is the **launcher**: it picks free loopback ports, generates
11//! a hostfile, and re-spawns itself once per rank as a separate OS process.
12//! The same shape scales to a real cluster — only the hostfile IPs change.
13//!
14//! ```no_run
15//! use rlx_distributed::launch::{worker_args, LocalCluster};
16//!
17//! # fn run_worker(_rank: u32, _hostfile: &str) {}
18//! fn main() -> anyhow::Result<()> {
19//!     match worker_args() {
20//!         Some(w) => run_worker(w.rank, &w.hostfile),  // spawned worker
21//!         None => {                                     // top-level launcher
22//!             let out = LocalCluster::new(3).arg("--device").arg("cpu").run()?;
23//!             for line in out {
24//!                 println!("{line}");
25//!             }
26//!         }
27//!     }
28//!     Ok(())
29//! }
30//! ```
31
32use crate::config::{Hostfile, TransportBackend};
33use anyhow::{Context, Result};
34use std::io::{BufRead, BufReader};
35use std::net::{Ipv4Addr, TcpListener};
36use std::process::{Command, Stdio};
37
38/// Reserve `n` free loopback TCP ports by binding ephemeral sockets and
39/// immediately releasing them. There is an inherent (small) race: the OS
40/// may hand a freed port to another process before the worker rebinds it.
41/// For local launches this is negligible and matches what mlx-lm / torchrun
42/// do for single-host runs.
43pub fn free_loopback_ports(n: u32) -> Result<Vec<u16>> {
44    (0..n)
45        .map(|_| {
46            let l = TcpListener::bind((Ipv4Addr::LOCALHOST, 0))
47                .context("binding an ephemeral loopback port")?;
48            // The listener drops at the end of this closure, freeing the port
49            // for the worker process to bind.
50            Ok(l.local_addr().context("reading local_addr")?.port())
51        })
52        .collect()
53}
54
55/// The launch-level arguments every spawned worker receives. Model-specific
56/// flags (e.g. `--device`) stay in argv for the binary to parse itself.
57#[derive(Debug, Clone)]
58pub struct WorkerArgs {
59    pub rank: u32,
60    pub hostfile: String,
61}
62
63/// Parse this process's worker invocation: `--rank N --hostfile PATH` (rank
64/// falls back to `RLX_RANK`). Returns `None` when neither `--rank` nor
65/// `RLX_RANK` is present — meaning this is the top-level launcher, which
66/// should build a [`LocalCluster`].
67pub fn worker_args() -> Option<WorkerArgs> {
68    let mut rank: Option<u32> = None;
69    let mut hostfile: Option<String> = None;
70    let mut it = std::env::args().skip(1);
71    while let Some(a) = it.next() {
72        match a.as_str() {
73            "--rank" => rank = it.next().and_then(|v| v.parse().ok()),
74            "--hostfile" => hostfile = it.next(),
75            _ => {}
76        }
77    }
78    let rank = rank.or_else(|| std::env::var("RLX_RANK").ok().and_then(|v| v.parse().ok()))?;
79    Some(WorkerArgs {
80        rank,
81        hostfile: hostfile.unwrap_or_default(),
82    })
83}
84
85/// A self-spawning local cluster: re-runs the current executable once per
86/// rank as a separate OS process, each pointed at a generated loopback
87/// hostfile, capturing one rank's stdout.
88///
89/// This exercises the real `DistConfig::load` / `connect` path over real TCP
90/// sockets in separate address spaces — the genuine deployment shape minus
91/// the physical wire.
92pub struct LocalCluster {
93    world: u32,
94    backend: TransportBackend,
95    extra_args: Vec<String>,
96    capture_rank: u32,
97    quiet: bool,
98}
99
100impl LocalCluster {
101    /// A `world`-rank cluster over loopback TCP, capturing rank 0's stdout.
102    pub fn new(world: u32) -> Self {
103        Self {
104            world,
105            backend: TransportBackend::Tcp,
106            extra_args: Vec::new(),
107            capture_rank: 0,
108            quiet: false,
109        }
110    }
111
112    /// Override the transport backend written into the hostfile.
113    pub fn backend(mut self, backend: TransportBackend) -> Self {
114        self.backend = backend;
115        self
116    }
117
118    /// Forward an extra CLI argument to every spawned worker (e.g.
119    /// `--device mlx`). Workers also receive `--rank i --hostfile <path>`.
120    pub fn arg(mut self, a: impl Into<String>) -> Self {
121        self.extra_args.push(a.into());
122        self
123    }
124
125    /// Forward several extra arguments at once.
126    pub fn args<I, S>(mut self, it: I) -> Self
127    where
128        I: IntoIterator<Item = S>,
129        S: Into<String>,
130    {
131        self.extra_args.extend(it.into_iter().map(Into::into));
132        self
133    }
134
135    /// Which rank's stdout to capture and return from [`run`](Self::run).
136    pub fn capture_rank(mut self, rank: u32) -> Self {
137        self.capture_rank = rank;
138        self
139    }
140
141    /// Suppress the one-line "spawning N ranks" banner on stderr.
142    pub fn quiet(mut self, quiet: bool) -> Self {
143        self.quiet = quiet;
144        self
145    }
146
147    /// Spawn `world` workers, block until all exit, and return the captured
148    /// rank's stdout line by line. Errors if any worker exits non-zero.
149    pub fn run(self) -> Result<Vec<String>> {
150        let exe = std::env::current_exe().context("resolving current_exe")?;
151        let ports = free_loopback_ports(self.world)?;
152        let hostfile = Hostfile::loopback(&ports, self.backend).write_temp("cluster")?;
153
154        if !self.quiet {
155            eprintln!(
156                "LocalCluster: spawning {} ranks ({} backend), hostfile {}",
157                self.world,
158                self.backend.as_str(),
159                hostfile.display()
160            );
161        }
162
163        let mut children = Vec::with_capacity(self.world as usize);
164        for rank in 0..self.world {
165            let capture = rank == self.capture_rank;
166            let child = Command::new(&exe)
167                .arg("--rank")
168                .arg(rank.to_string())
169                .arg("--hostfile")
170                .arg(&hostfile)
171                .args(&self.extra_args)
172                .stdout(if capture {
173                    Stdio::piped()
174                } else {
175                    Stdio::null()
176                })
177                .stderr(Stdio::inherit())
178                .spawn()
179                .with_context(|| format!("spawning worker rank {rank}"))?;
180            children.push((rank, child));
181        }
182
183        // Drain the captured rank's stdout to EOF (it closes when that worker
184        // exits); the other ranks run concurrently meanwhile.
185        let mut lines = Vec::new();
186        if let Some((_, child)) = children.iter_mut().find(|(r, _)| *r == self.capture_rank)
187            && let Some(out) = child.stdout.take()
188        {
189            lines.extend(BufReader::new(out).lines().map_while(Result::ok));
190        }
191
192        let mut failures = Vec::new();
193        for (rank, mut child) in children {
194            match child.wait() {
195                Ok(status) if status.success() => {}
196                Ok(status) => failures.push(format!("rank {rank} exited with {status}")),
197                Err(e) => failures.push(format!("rank {rank} wait() failed: {e}")),
198            }
199        }
200        let _ = std::fs::remove_file(&hostfile);
201
202        if !failures.is_empty() {
203            anyhow::bail!("LocalCluster worker failure: {}", failures.join("; "));
204        }
205        Ok(lines)
206    }
207}
208
209#[cfg(test)]
210mod tests {
211    use super::*;
212
213    #[test]
214    fn allocates_distinct_free_ports() {
215        let ports = free_loopback_ports(4).unwrap();
216        assert_eq!(ports.len(), 4);
217        // Distinct and non-zero.
218        let mut sorted = ports.clone();
219        sorted.sort_unstable();
220        sorted.dedup();
221        assert_eq!(sorted.len(), 4, "ports should be distinct: {ports:?}");
222        assert!(ports.iter().all(|&p| p != 0));
223    }
224
225    #[test]
226    fn worker_args_none_without_flag() {
227        // The test harness runs with no `--rank` and no RLX_RANK.
228        assert!(std::env::var("RLX_RANK").is_err());
229        assert!(worker_args().is_none());
230    }
231
232    #[test]
233    fn builder_collects_args() {
234        let c = LocalCluster::new(2)
235            .backend(TransportBackend::Thunderbolt)
236            .arg("--device")
237            .arg("mlx")
238            .args(["--decode", "--max-tokens", "8"]);
239        assert_eq!(c.world, 2);
240        assert_eq!(c.backend, TransportBackend::Thunderbolt);
241        assert_eq!(
242            c.extra_args,
243            ["--device", "mlx", "--decode", "--max-tokens", "8"]
244        );
245    }
246}
rlx_distributed/launch.rs

rlx_distributed/
launch.rs
