Skip to main content

Module launch

Module launch 

Source
Expand description

Local multi-process launcher — run an N-rank cluster on one host over loopback TCP, with no hand-written hosts.json and no N terminals.

A distributed binary has two roles. When launched with --rank R it is a worker: it forms the process group and runs its block. With no --rank, it is the launcher: it picks free loopback ports, generates a hostfile, and re-spawns itself once per rank as a separate OS process. The same shape scales to a real cluster — only the hostfile IPs change.

use rlx_distributed::launch::{worker_args, LocalCluster};

fn main() -> anyhow::Result<()> {
    match worker_args() {
        Some(w) => run_worker(w.rank, &w.hostfile),  // spawned worker
        None => {                                     // top-level launcher
            let out = LocalCluster::new(3).arg("--device").arg("cpu").run()?;
            for line in out {
                println!("{line}");
            }
        }
    }
    Ok(())
}

Structs§

LocalCluster
A self-spawning local cluster: re-runs the current executable once per rank as a separate OS process, each pointed at a generated loopback hostfile, capturing one rank’s stdout.
WorkerArgs
The launch-level arguments every spawned worker receives. Model-specific flags (e.g. --device) stay in argv for the binary to parse itself.

Functions§

free_loopback_ports
Reserve n free loopback TCP ports by binding ephemeral sockets and immediately releasing them. There is an inherent (small) race: the OS may hand a freed port to another process before the worker rebinds it. For local launches this is negligible and matches what mlx-lm / torchrun do for single-host runs.
worker_args
Parse this process’s worker invocation: --rank N --hostfile PATH (rank falls back to RLX_RANK). Returns None when neither --rank nor RLX_RANK is present — meaning this is the top-level launcher, which should build a LocalCluster.