rlx_distributed/cluster/
mod.rs

1// RLX — versatile ML compiler + runtime.
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4//! **Cluster orchestration** — the high-level DX over the pipeline primitives.
5//! Turn a declarative [`ClusterConfig`] into a probed, HW-planned, monitored run:
6//!
7//! ```no_run
8//! # use rlx_distributed::cluster::*;
9//! # fn cost() -> ModelCost { unimplemented!() }
10//! let mut cx = Cluster::from_path("cluster.toml")?;
11//! cx.probe("/home/user/rlx-models/target/release/examples/dsv4_cluster")?; // ssh-probe every node
12//! for a in cx.plan(cost())? { println!("{} -> layers {:?} on {}", a.addr, a.layers, a.device); }
13//! // The model crate builds each stage; the coordinator drives + monitors:
14//! let children = cx.launch("/home/user/.../dsv4_cluster")?;            // spawn workers
15//! let run = cx.drive(vec![/* input NamedTensor */])?;                   // relay + time
16//! println!("{}", run.table());
17//! # Ok::<(), anyhow::Error>(())
18//! ```
19//!
20//! The model-specific bits (build a `Stage` for a layer range, load its weights)
21//! stay in the model crate behind [`crate::Stage`] / [`crate::ParamSource`]; this
22//! module only orchestrates.
23
24pub mod caps;
25pub mod config;
26pub mod placement;
27pub mod stats;
28
29pub use caps::{DeviceInfo, NodeCaps, probe_local, probe_remote};
30pub use config::{ClusterConfig, KvPolicy, NodeConfig, PlacementPolicy};
31pub use placement::{Assignment, ModelCost, plan_placement};
32pub use rlx_driver::{ProcessGroup, ReduceKind};
33pub use stats::{ClusterRun, NodeReport, StageTiming};
34
35use crate::graph::transport::run_pipeline_tcp_timed;
36use crate::graph::{NamedTensor, Stage};
37use anyhow::{Context, Result};
38use std::path::Path;
39use std::process::Child;
40
41/// Cross-node **all-reduce** over the cluster's [`ProcessGroup`] — the collective
42/// complementary to the pipeline's sequential hidden-state relay. Model crates use
43/// it for TENSOR parallelism (shard one layer across nodes, then `Sum` the partial
44/// matmul outputs) or to aggregate per-node values (`Mean`/`Max`/`Min`). The
45/// ProcessGroup rides the same rlx-driver transports (TCP / Thunderbolt / MLX).
46pub fn all_reduce(pg: &ProcessGroup, buf: &mut [f32], kind: ReduceKind) -> Result<()> {
47    pg.all_reduce(buf, kind)
48        .map_err(|e| anyhow::anyhow!("all_reduce: {e:?}"))
49}
50
51/// A configured, probed, planned cluster.
52pub struct Cluster {
53    pub cfg: ClusterConfig,
54    /// Probed hardware, one per node (same order as `cfg.nodes`). Empty until [`Cluster::probe`].
55    pub caps: Vec<NodeCaps>,
56    /// Layer→node plan. Empty until [`Cluster::plan`].
57    pub plan: Vec<Assignment>,
58}
59
60impl Cluster {
61    pub fn from_config(cfg: ClusterConfig) -> Self {
62        Self {
63            cfg,
64            caps: Vec::new(),
65            plan: Vec::new(),
66        }
67    }
68    pub fn from_path(p: impl AsRef<Path>) -> Result<Self> {
69        Ok(Self::from_config(ClusterConfig::from_path(p)?))
70    }
71
72    /// Probe every node's hardware. `remote_bin` is the worker binary path on the
73    /// remote hosts (its `--probe` mode self-reports); nodes without `ssh` are
74    /// probed locally.
75    pub fn probe(&mut self, remote_bin: &str) -> Result<&[NodeCaps]> {
76        self.caps.clear();
77        for n in &self.cfg.nodes {
78            let caps = match &n.ssh {
79                Some(host) => probe_remote(host, remote_bin, &n.addr, &n.ckpt_dir)
80                    .with_context(|| format!("probe {host}"))?,
81                None => probe_local(&n.addr, &n.ckpt_dir, true),
82            };
83            self.caps.push(caps);
84        }
85        Ok(&self.caps)
86    }
87
88    /// Plan placement from the model cost model + probed caps (probe first).
89    pub fn plan(&mut self, model: ModelCost) -> Result<&[Assignment]> {
90        anyhow::ensure!(!self.caps.is_empty(), "probe() before plan()");
91        let nodes: Vec<_> = self
92            .caps
93            .iter()
94            .cloned()
95            .zip(self.cfg.nodes.iter().cloned())
96            .collect();
97        let reserve = (self.cfg.reserve_ram_gb * 1e9) as u64;
98        self.plan = plan_placement(&model, &nodes, self.cfg.placement.policy, reserve)?;
99        Ok(&self.plan)
100    }
101
102    /// The worker command-line for node `i` from its plan + config — so the model
103    /// crate's worker binary honours device / precision / rng / kv / layers.
104    pub fn worker_argv(&self, i: usize) -> Vec<String> {
105        let a = &self.plan[i];
106        let n = &self.cfg.nodes[i];
107        let mut v = vec![
108            "--role".into(),
109            "worker".into(),
110            "--index".into(),
111            i.to_string(),
112            "--layers".into(),
113            format!("{}:{}", a.layers.start, a.layers.end),
114            "--ckpt".into(),
115            n.ckpt_dir.clone(),
116            "--addr".into(),
117            n.addr.clone(),
118            "--seq".into(),
119            self.cfg.seq.to_string(),
120            "--device".into(),
121            n.device.clone(),
122            "--precision".into(),
123            n.precision.clone(),
124            "--kv".into(),
125            format!("{:?}", n.kv_cache).to_lowercase(),
126            "--rng".into(),
127            self.cfg.seed_for(i).to_string(),
128        ];
129        if a.first {
130            v.push("--first".into());
131        }
132        if a.last {
133            v.push("--last".into());
134        }
135        v
136    }
137
138    /// Spawn a worker process per node (ssh for remote, local otherwise). The
139    /// caller drives them via [`Cluster::drive`], then reaps the children.
140    pub fn launch(&self, remote_bin: &str) -> Result<Vec<Child>> {
141        anyhow::ensure!(!self.plan.is_empty(), "plan() before launch()");
142        use std::process::Stdio;
143        let mut kids = Vec::new();
144        for (i, n) in self.cfg.nodes.iter().enumerate() {
145            let argv = self.worker_argv(i);
146            // Pipe stdout so the coordinator can await "serving on" + collect the
147            // node's build report; leave stderr inherited for live error output.
148            let child = match &n.ssh {
149                Some(host) => std::process::Command::new("ssh")
150                    .arg(host)
151                    .arg(format!("{remote_bin} {}", argv.join(" ")))
152                    .stdout(Stdio::piped())
153                    .spawn()?,
154                // Local node: this very binary, not the remote path. Its stderr is
155                // otherwise a black hole (the coordinator's own stderr races it);
156                // capture to a file when RLX_WORKER_ERR_DIR is set, for diagnosis.
157                None => {
158                    let mut c = std::process::Command::new(std::env::current_exe()?);
159                    c.args(&argv).stdout(Stdio::piped());
160                    if let Ok(dir) = std::env::var("RLX_WORKER_ERR_DIR")
161                        && let Ok(f) = std::fs::File::create(format!("{dir}/local_worker_{i}.err"))
162                    {
163                        c.stderr(Stdio::from(f));
164                    }
165                    c.spawn()?
166                }
167            };
168            kids.push(child);
169        }
170        Ok(kids)
171    }
172
173    /// Drive one forward through the planned stages, timing each, and collect the
174    /// output. `inputs` seeds the first stage (e.g. `input_ids`). Boundary tensor
175    /// name is `hidden`/`hidden_in` per the model's stage I/O convention.
176    pub fn drive(&self, inputs: Vec<NamedTensor>) -> Result<ClusterRun> {
177        let stages: Vec<Stage> = (0..self.plan.len()).map(|i| self.meta_stage(i)).collect();
178        let addrs: Vec<String> = self.cfg.nodes.iter().map(|n| n.addr.clone()).collect();
179        let (out, per_stage_ms) = run_pipeline_tcp_timed(&stages, &addrs, inputs)?;
180        let timings = self
181            .plan
182            .iter()
183            .zip(per_stage_ms)
184            .map(|(a, ms)| StageTiming {
185                addr: a.addr.clone(),
186                layers: a.layers.clone(),
187                device: a.device.clone(),
188                build_ms: 0,
189                forward_ms: ms,
190                resident_bytes: a.est_bytes,
191            })
192            .collect::<Vec<_>>();
193        let total: u64 = timings.iter().map(|t| t.forward_ms).sum();
194        Ok(ClusterRun {
195            timings,
196            total_forward_ms: total,
197            output: out.into_iter().next().map(|t| t.data).unwrap_or_default(),
198        })
199    }
200
201    /// Coordinator-side boundary metadata for stage `i` (no weights, no graph):
202    /// first stage takes `input_ids`, others take `hidden_in`; last emits
203    /// `logits`, others `hidden_in`.
204    fn meta_stage(&self, i: usize) -> Stage {
205        let (first, last) = (self.plan[i].first, self.plan[i].last);
206        Stage {
207            index: i,
208            graph: rlx_ir::Graph::new("meta"),
209            inputs: vec![if first { "input_ids" } else { "hidden_in" }.into()],
210            outputs: vec![if last { "logits" } else { "hidden_in" }.into()],
211            output_shapes: vec![],
212            params: vec![],
213        }
214    }
215}
rlx_distributed/cluster/mod.rs

rlx_distributed/cluster/
mod.rs
