rlx_distributed/cluster/
stats.rs

1// RLX — versatile ML compiler + runtime.
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4//! **Per-node monitoring** — build/forward timing, resident memory, and
5//! throughput, collected by the coordinator and rendered as a table.
6
7use serde::{Deserialize, Serialize};
8use std::ops::Range;
9
10/// What a worker reports about building + serving its stage (JSON on stdout).
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct NodeReport {
13    pub addr: String,
14    pub layers: Range<usize>,
15    pub device: String,
16    pub precision: String,
17    /// Time to build the stage graph + load weights (ms).
18    pub build_ms: u64,
19    /// Resident bytes after build.
20    pub resident_bytes: u64,
21    pub n_params: usize,
22    pub n_packed: usize,
23}
24
25/// Coordinator-side per-stage timing (network + compute for one forward).
26#[derive(Debug, Clone)]
27pub struct StageTiming {
28    pub addr: String,
29    pub layers: Range<usize>,
30    pub device: String,
31    pub build_ms: u64,
32    pub forward_ms: u64,
33    pub resident_bytes: u64,
34}
35
36/// Full run result: the model output + per-node monitoring.
37#[derive(Debug, Clone)]
38pub struct ClusterRun {
39    pub timings: Vec<StageTiming>,
40    pub total_forward_ms: u64,
41    /// Final-stage output (e.g. logits), flat.
42    pub output: Vec<f32>,
43}
44
45impl ClusterRun {
46    /// Throughput in tokens/s given the sequence length driven.
47    pub fn tokens_per_s(&self, seq: usize) -> f64 {
48        if self.total_forward_ms == 0 {
49            return 0.0;
50        }
51        seq as f64 / (self.total_forward_ms as f64 / 1000.0)
52    }
53
54    /// Aggregate a per-node metric with a reduce op — the SAME [`ReduceKind`](rlx_driver::ReduceKind)
55    /// semantics the collective all-reduce uses for tensor-parallel partial sums.
56    /// (Here over locally-gathered stats; [`crate::cluster::all_reduce`] runs the
57    /// real cross-node collective.)
58    pub fn reduce_resident_gb(&self, kind: rlx_driver::ReduceKind) -> f64 {
59        use rlx_driver::ReduceKind::*;
60        let vals = self.timings.iter().map(|t| t.resident_bytes as f64 / 1e9);
61        match kind {
62            Sum => vals.sum(),
63            Mean => {
64                let n = self.timings.len().max(1) as f64;
65                vals.sum::<f64>() / n
66            }
67            Max => vals.fold(f64::MIN, f64::max),
68            Min => vals.fold(f64::MAX, f64::min),
69        }
70    }
71    /// Critical-path forward time = MAX over stages (a reduce-Max); the pipeline
72    /// bottleneck the planner's throughput policy minimizes.
73    pub fn critical_path_ms(&self) -> u64 {
74        self.timings.iter().map(|t| t.forward_ms).max().unwrap_or(0)
75    }
76
77    /// A monitor table for stdout.
78    pub fn table(&self) -> String {
79        let mut s = String::new();
80        s.push_str(&format!(
81            "{:<22}{:<10}{:<8}{:>10}{:>12}{:>12}\n",
82            "node", "layers", "device", "build", "forward", "resident"
83        ));
84        s.push_str(&"─".repeat(74));
85        s.push('\n');
86        for t in &self.timings {
87            s.push_str(&format!(
88                "{:<22}{:<10}{:<8}{:>9}s{:>11}s{:>10.1}G\n",
89                t.addr,
90                format!("{}..{}", t.layers.start, t.layers.end),
91                t.device,
92                format!("{:.1}", t.build_ms as f64 / 1000.0),
93                format!("{:.1}", t.forward_ms as f64 / 1000.0),
94                t.resident_bytes as f64 / 1e9,
95            ));
96        }
97        s.push_str(&"─".repeat(74));
98        s.push('\n');
99        // Cluster aggregates via reduce ops: Σ resident, max critical-path stage.
100        use rlx_driver::ReduceKind::{Max, Sum};
101        s.push_str(&format!(
102            "reduce: Σ resident {:.1} GB, max resident {:.1} GB | total forward {:.1}s, critical path {:.1}s across {} stages\n",
103            self.reduce_resident_gb(Sum),
104            self.reduce_resident_gb(Max),
105            self.total_forward_ms as f64 / 1000.0,
106            self.critical_path_ms() as f64 / 1000.0,
107            self.timings.len(),
108        ));
109        s
110    }
111}
rlx_distributed/cluster/stats.rs

rlx_distributed/cluster/
stats.rs
