diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..71e2eb6 --- /dev/null +++ b/.gitignore @@ -0,0 +1,21 @@ +# Rust build +/target/ +Cargo.lock + +# Built profiler plugin +*.so +libnccl-profiler.so + +# Local KernelStep / setup build stamps +.kernelstep-* + +# Nested NCCL build tree +third_party/nccl/build/ + +# Editor / OS +.DS_Store +*.swp +*~ +# Local tool caches +.cargo-cache/ +.rustup-cache/ diff --git a/.gitmodules b/.gitmodules index 162279d..d628fae 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,3 +1,4 @@ [submodule "third_party/nccl"] path = third_party/nccl - url = https://github.com/NVIDIA/nccl.git + url = https://github.com/alizmhdi/nccl.git + branch = kernelstep-profiler diff --git a/Cargo.toml b/Cargo.toml index fa83832..9a54f22 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,6 +3,7 @@ name = "nccl-profiler" version = "0.3.0" edition = "2021" + [profile.release] debug = 1 opt-level = 3 @@ -47,6 +48,7 @@ opentelemetry-otlp = { version = "0.29", features = ["http-proto", "reqwest-bloc prost = "0.13.3" prost-types = "0.13.3" rand = "0.9.0" +serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" static_assertions = "1.1.0" tempfile = "3.19.0" diff --git a/docs/configuration.md b/docs/configuration.md index e2ff59c..cb2aef7 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -55,6 +55,7 @@ CoMMA can export raw telemetry data to local files in JSON format. | Environment Variable | Type | Default | Description | | :--- | :--- | :--- | :--- | | `NCCL_PROFILER_LATENCY_FILE` | String | *None* | Path to the local file where raw event telemetry will be written (e.g., `/tmp/latency-%p.json`). Supports `%p` for PID. | +| `NCCL_PROFILER_LATENCY_SOCK` | String | *None* | Unix stream socket for live raw event telemetry. Sends the same newline-delimited JSON as `LATENCY_FILE`; supports `%p` for PID but the monitor usually uses one shared socket per host. | | `NCCL_PROFILER_SUMMARY_FILE` | String | *None* | Path to the local file where periodic summaries will be written. | | `NCCL_PROFILER_SUMMARY_INTERVAL` | Duration | `60s` | Interval at which periodic summaries are calculated and written to `SUMMARY_FILE`. | diff --git a/src/cloud_daemon.rs b/src/cloud_daemon.rs index 14f3320..143a7f5 100644 --- a/src/cloud_daemon.rs +++ b/src/cloud_daemon.rs @@ -38,6 +38,7 @@ use std::sync::{Arc, Mutex, Once}; use std::time::{Duration, Instant}; use tokio::fs::OpenOptions; use tokio::io::AsyncWriteExt; +use tokio::net::UnixStream; use tokio::sync::mpsc::error::TrySendError; use tokio::sync::{mpsc, oneshot}; @@ -88,6 +89,9 @@ impl Telemetry { Telemetry::NcclOp(ncclop) => { writeln!(buf, "{}", ncclop.trace_record(&mut time_to_num))?; } + Telemetry::P2pParent(parent) => { + writeln!(buf, "{}", parent.trace_record(&mut time_to_num))?; + } Telemetry::ProxyOp(proxyop) => { writeln!(buf, "{}", proxyop.trace_record(&mut time_to_num))?; } @@ -334,6 +338,20 @@ async fn build_bufwriter( } } +async fn connect_latency_sock(path: impl AsRef) -> Option { + match UnixStream::connect(&path).await { + Ok(stream) => Some(stream), + Err(e) => { + error!( + "Failed to connect latency telemetry socket {:?}: {}.", + path.as_ref().as_os_str(), + e + ); + None + } + } +} + async fn exporter( profiler: &'static Profiler, mut rx: mpsc::Receiver, @@ -346,6 +364,18 @@ async fn exporter( None }; + let latency_sock_path = profiler + .config + .latency_sock + .as_ref() + .map(|template| template.replace("%p", &format!("{}", profiler.pid))); + let mut latency_sock = if let Some(path) = latency_sock_path.as_ref() { + connect_latency_sock(path).await + } else { + None + }; + let mut next_latency_sock_retry = Instant::now(); + let mut summary_file = if let Some(template) = profiler.config.summary_file.as_ref() { let path = template.replace("%p", &format!("{}", profiler.pid)); build_bufwriter(path).await @@ -386,6 +416,16 @@ async fn exporter( let mut uploader_interval = tokio::time::interval(profiler.config.heartbeat_upload_interval); uploader_interval.tick().await; + // Periodic latency-file flush for online consumers (straggler monitor). + let mut latency_flush_interval = + if profiler.config.latency_flush_interval > Duration::from_secs(0) { + let mut i = tokio::time::interval(profiler.config.latency_flush_interval); + i.tick().await; + Some(i) + } else { + None + }; + let mut otel_metrics_grouping_interval = if profiler.config.otel_enable { let mut i = tokio::time::interval(profiler.config.otel_metrics_cardinality_grouping_interval); @@ -412,6 +452,30 @@ async fn exporter( } } + if latency_sock.is_none() { + if let Some(path) = latency_sock_path.as_ref() { + if Instant::now() >= next_latency_sock_retry { + latency_sock = connect_latency_sock(path).await; + if latency_sock.is_none() { + next_latency_sock_retry = Instant::now() + Duration::from_secs(1); + } + } + } + } + + if let Some(stream) = latency_sock.as_mut() { + let r = telemetry + .write_to_file( + stream, + |t| profiler.instant_to_timestamp(t).as_micros() as _) + .await; + if let Err(e) = r { + error!("Failed to stream latency telemetry to socket: {}. Will retry.", e); + latency_sock = None; + next_latency_sock_retry = Instant::now() + Duration::from_secs(1); + } + } + if let Some(summary) = summary.as_mut() { match &telemetry { Telemetry::NcclOp(op) => { @@ -453,6 +517,15 @@ async fn exporter( None => break, } }, + _ = async { latency_flush_interval.as_mut().unwrap().tick().await }, + if latency_flush_interval.is_some() && latency_file.is_some() => { + if let Some(file) = latency_file.as_mut() { + if let Err(e) = file.flush().await { + error!("Failed to flush latency telemetry file: {}. Stop logging.", e); + latency_file = None; + } + } + }, _ = summary_interval.tick(), if summary.is_some() => { if let Some(file) = summary_file.as_mut() { let s = summary.as_mut().unwrap(); diff --git a/src/config.rs b/src/config.rs index a071f68..a120e1d 100644 --- a/src/config.rs +++ b/src/config.rs @@ -72,6 +72,10 @@ pub struct Config { pub track_step_fifo_wait: bool, pub aggregate_steps: bool, pub track_kernel_ch: bool, + pub track_kernel_step: bool, + /// Unix socket path template for mid-flight gate RPC (`%p` = pid). + /// Default derived from `latency_file` dir as `control-%p.sock`. + pub control_sock: Option, pub ncclop_completion_delay: Duration, pub comm_hash_ipc_timeout: Duration, @@ -89,6 +93,9 @@ pub struct Config { // Export method & config pub latency_file: Option, + /// Unix stream socket path for live NDJSON latency telemetry. + pub latency_sock: Option, + pub latency_flush_interval: Duration, pub summary_file: Option, pub summary_interval: Duration, @@ -124,8 +131,11 @@ impl Config { field_from_env!(s, track_steps, false); field_from_env!(s, track_recv_steps, false); field_from_env!(s, track_step_fifo_wait, true); - field_from_env!(s, aggregate_steps, true); - field_from_env!(s, track_kernel_ch, false); + // Coarse default: Coll/P2P start+end only (no ProxyStep / KernelStep). + field_from_env!(s, aggregate_steps, false); + field_from_env!(s, track_kernel_ch, true); + field_from_env!(s, track_kernel_step, false); + field_from_env!(s, control_sock); field_from_env!(s, ncclop_completion_delay, Duration::from_secs(2)); field_from_env!(s, comm_hash_ipc_timeout, Duration::from_secs(1)); @@ -152,6 +162,9 @@ impl Config { field_from_env!(s, use_cached_clock, false); field_from_env!(s, latency_file); + field_from_env!(s, latency_sock); + // 0 disables periodic flushing (stock behavior: flush on shutdown only). + field_from_env!(s, latency_flush_interval, Duration::from_secs(0)); field_from_env!(s, summary_file); field_from_env!(s, summary_interval, Duration::from_secs(60)); diff --git a/src/control_rpc.rs b/src/control_rpc.rs new file mode 100644 index 0000000..5d9015b --- /dev/null +++ b/src/control_rpc.rs @@ -0,0 +1,227 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Unix-socket RPC for mid-flight CoMMA metric gates (used by comma-monitor). +//! +//! Socket path: `{dir}/control-{pid}.sock` (same directory as `latency-{pid}.txt`). +//! Protocol: one JSON object per line. +//! +//! Requests: +//! ```json +//! {"cmd":"ping"} +//! {"cmd":"get"} +//! {"cmd":"set","track_kernel_step":false,"track_steps":true} +//! ``` +//! +//! Responses: +//! ```json +//! {"ok":true,"pid":1234,"gates":{...}} +//! {"ok":false,"error":"..."} +//! ``` + +use crate::profiler::Profiler; +use crate::runtime_gates::{GateSnapshot, GateUpdate}; + +use log::{info, warn}; +use serde::{Deserialize, Serialize}; + +use std::fs; +use std::io::{BufRead, BufReader, Write}; +use std::os::unix::fs::PermissionsExt; +use std::os::unix::net::{UnixListener, UnixStream}; +use std::path::{Path, PathBuf}; +use std::time::Duration; + +#[derive(Debug, Deserialize)] +struct Request { + cmd: String, + #[serde(flatten)] + update: GateUpdate, +} + +#[derive(Debug, Serialize)] +struct Response<'a> { + ok: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pid: Option, + #[serde(skip_serializing_if = "Option::is_none")] + changed: Option, + #[serde(skip_serializing_if = "Option::is_none")] + gates: Option<&'a GateSnapshot>, + #[serde(skip_serializing_if = "Option::is_none")] + error: Option, +} + +/// Resolve control socket path from config / latency file / env. +pub fn resolve_sock_path( + config_sock: Option<&str>, + latency_file: Option<&str>, + pid: libc::pid_t, +) -> Option { + let template = config_sock + .map(str::to_string) + .or_else(|| std::env::var("NCCL_PROFILER_CONTROL_SOCK").ok()) + .or_else(|| { + latency_file.and_then(|lf| { + let p = Path::new(lf); + let dir = p.parent().unwrap_or_else(|| Path::new(".")); + Some(dir.join("control-%p.sock").to_string_lossy().into_owned()) + }) + })?; + Some(PathBuf::from(template.replace("%p", &pid.to_string()))) +} + +pub fn spawn_control_server(profiler: &'static Profiler) { + let Some(path) = resolve_sock_path( + profiler.config.control_sock.as_deref(), + profiler.config.latency_file.as_deref(), + profiler.pid, + ) else { + info!("CoMMA control RPC disabled (no control sock / latency_file path)"); + return; + }; + if let Err(e) = std::thread::Builder::new() + .name("comma-control-rpc".into()) + .spawn(move || run_server(path, profiler)) + { + warn!("failed to spawn CoMMA control RPC thread: {e}"); + } +} + +fn run_server(path: PathBuf, profiler: &'static Profiler) { + if let Some(parent) = path.parent() { + let _ = fs::create_dir_all(parent); + // World-writable dir so the host-side monitor (non-root) can connect + // when CoMMA runs as root inside the training container. + let _ = fs::set_permissions(parent, fs::Permissions::from_mode(0o777)); + } + let _ = fs::remove_file(&path); + // Linux sockaddr_un.sun_path is ~108 bytes; refuse long paths early. + let path_bytes = path.as_os_str().as_encoded_bytes(); + if path_bytes.len() >= 108 { + warn!( + "CoMMA control RPC path too long ({} >= 108): {}", + path_bytes.len(), + path.display() + ); + return; + } + let listener = match UnixListener::bind(&path) { + Ok(l) => l, + Err(e) => { + warn!("CoMMA control RPC bind {}: {e}", path.display()); + return; + } + }; + // Connect requires write on the socket inode; container often runs as root. + let _ = fs::set_permissions(&path, fs::Permissions::from_mode(0o666)); + // Make accept interruptible-ish for process teardown. + let _ = listener.set_nonblocking(false); + info!( + "CoMMA control RPC listening on {} (pid={})", + path.display(), + profiler.pid + ); + loop { + match listener.accept() { + Ok((stream, _)) => { + if let Err(e) = handle_client(stream, profiler) { + warn!("CoMMA control RPC client error: {e}"); + } + } + Err(e) => { + warn!("CoMMA control RPC accept error: {e}"); + std::thread::sleep(Duration::from_millis(50)); + } + } + } +} + +fn handle_client(stream: UnixStream, profiler: &'static Profiler) -> std::io::Result<()> { + let _ = stream.set_read_timeout(Some(Duration::from_secs(5))); + let _ = stream.set_write_timeout(Some(Duration::from_secs(5))); + let mut reader = BufReader::new(stream.try_clone()?); + let mut writer = stream; + let mut line = String::new(); + loop { + line.clear(); + let n = reader.read_line(&mut line)?; + if n == 0 { + break; + } + let resp = dispatch_line(line.trim(), profiler); + writeln!(writer, "{resp}")?; + writer.flush()?; + } + Ok(()) +} + +fn dispatch_line(line: &str, profiler: &'static Profiler) -> String { + if line.is_empty() { + return err_json("empty request"); + } + let req: Request = match serde_json::from_str(line) { + Ok(r) => r, + Err(e) => return err_json(&format!("invalid json: {e}")), + }; + match req.cmd.as_str() { + "ping" => { + let snap = profiler.gates.snapshot(); + ok_json(profiler.pid, None, Some(&snap)) + } + "get" => { + let snap = profiler.gates.snapshot(); + ok_json(profiler.pid, None, Some(&snap)) + } + "set" => { + let changed = profiler.gates.apply_update(&req.update); + let snap = profiler.gates.snapshot(); + ok_json(profiler.pid, Some(changed), Some(&snap)) + } + other => err_json(&format!("unknown cmd {other:?} (want ping|get|set)")), + } +} + +fn ok_json(pid: libc::pid_t, changed: Option, gates: Option<&GateSnapshot>) -> String { + serde_json::to_string(&Response { + ok: true, + pid: Some(pid as i32), + changed, + gates, + error: None, + }) + .unwrap_or_else(|_| "{\"ok\":false,\"error\":\"serialize\"}".into()) +} + +fn err_json(msg: &str) -> String { + serde_json::to_string(&Response { + ok: false, + pid: None, + changed: None, + gates: None, + error: Some(msg.to_string()), + }) + .unwrap_or_else(|_| "{\"ok\":false,\"error\":\"serialize\"}".into()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn resolve_from_latency_template() { + let p = resolve_sock_path(None, Some("/tmp/comma/latency-%p.txt"), 42).unwrap(); + assert_eq!(p, PathBuf::from("/tmp/comma/control-42.sock")); + } +} diff --git a/src/daemon.rs b/src/daemon.rs index 8961594..8746f7f 100644 --- a/src/daemon.rs +++ b/src/daemon.rs @@ -16,6 +16,7 @@ use crate::event; use crate::event::ProfilerEvent as _; use crate::fixed_batch; use crate::nccl_metadata; +use crate::nccl_metadata::P2p as _; use crate::profiler; use crate::profiler::Communicator; use crate::profiler::Profiler; @@ -26,7 +27,7 @@ use crate::step_tracker::EventStep; use log::error; -use std::collections::{BTreeMap, HashMap, VecDeque}; +use std::collections::{BTreeMap, HashMap, HashSet, VecDeque}; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -118,6 +119,8 @@ impl FifoReceiver { #[derive(Debug)] pub enum ControlMessage { NewThread(ThreadControl), + /// Apply runtime gate update (also used by tests / future IPC). + SetGates(crate::runtime_gates::GateUpdate), } #[derive(Debug)] @@ -141,6 +144,7 @@ pub enum Message { /* duration ns */ u64, /* parent handle */ usize, ), + KernelStep(event::KernelEventStep, /* parent handle */ usize), CommOpen(Communicator), CommClose(/* comm_hash = */ u64), } @@ -170,11 +174,33 @@ pub enum Telemetry { Group(Box), NcclOpIssued(Box), // copybara:strip(hang detection) NcclOp(Box), + P2pParent(Box), ProxyOp(Box), CommOpen(Communicator), CommClose(/* comm_hash = */ u64), } +/// Accumulator for P2Ps (and optional COLLs) that share one NCCL Group. +#[derive(Default)] +struct P2pGroupAcc { + expected: usize, + completed: usize, + ended: bool, + has_coll: bool, + min_start: Option, + max_end: Option, + send_bytes: usize, + peers: HashSet, + comm_hash: u64, + rank: usize, +} + +impl P2pGroupAcc { + fn will_emit_parent(&self) -> bool { + self.ended && !self.has_coll && self.expected >= 2 + } +} + pub struct PollingContext<'a> { profiler: &'a Profiler, pub ncclops: BTreeMap>, @@ -187,6 +213,7 @@ pub struct PollingContext<'a> { peer_rank_fifo: HashMap>, pending_ipc_msg: HashMap>, + p2p_groups: HashMap, } impl<'a> PollingContext<'a> { @@ -201,10 +228,104 @@ impl<'a> PollingContext<'a> { free_step_batch: slab::FreeList::default(), peer_rank_fifo: HashMap::new(), pending_ipc_msg: HashMap::new(), + p2p_groups: HashMap::new(), + } + } + + fn note_ncclop_start(&mut self, op: &event::NcclOp) { + let Some(gid) = op.parent_group_id() else { + return; + }; + let g = self.p2p_groups.entry(gid).or_default(); + if op.is_p2p() { + g.expected += 1; + } else { + g.has_coll = true; } } - fn reclaim_ncclop(&mut self, op: event::NcclOp) { + fn record_p2p_completion(&mut self, op: &event::NcclOp) { + let Some(gid) = op.parent_group_id() else { + return; + }; + let g = self.p2p_groups.entry(gid).or_default(); + g.completed += 1; + let start = op.basic_info().start_time(); + g.min_start = Some(g.min_start.map_or(start, |t| t.min(start))); + if let Some(end) = op.basic_info().end_time() { + g.max_end = Some(g.max_end.map_or(end, |t| t.max(end))); + } + g.comm_hash = op.comm_hash(); + g.rank = op.basic_info().rank(); + if let Some(p2p) = op.get_descr().try_cast_to_p2p() { + g.peers.insert(p2p.peer()); + if p2p.is_send() { + g.send_bytes = g.send_bytes.saturating_add(op.byte_count()); + } + } + } + + fn will_emit_parent(&self, gid: u64) -> bool { + self.p2p_groups + .get(&gid) + .map(P2pGroupAcc::will_emit_parent) + .unwrap_or(false) + } + + fn try_emit_p2p_parent(&mut self, gid: u64) { + let Some(g) = self.p2p_groups.get(&gid) else { + return; + }; + if !(g.will_emit_parent() && g.completed == g.expected) { + return; + } + let g = self.p2p_groups.remove(&gid).unwrap(); + let start = g.min_start.unwrap_or_else(Instant::now); + let end = g.max_end.unwrap_or(start); + let name = if g.peers.len() >= 2 { + "all_to_all" + } else { + "sendrecv" + }; + self.pending_telemetry + .push_back(Telemetry::P2pParent(Box::new(event::P2pParent { + group_id: gid, + rank: g.rank, + comm_hash: g.comm_hash, + start_time: start, + end_time: if end > start { end } else { start }, + size: g.send_bytes, + n_children: g.completed, + n_peers: g.peers.len(), + name, + }))); + } + + fn flush_p2p_parents(&mut self) { + let gids: Vec = self.p2p_groups.keys().copied().collect(); + for gid in gids { + if let Some(g) = self.p2p_groups.get_mut(&gid) { + if g.ended && !g.has_coll && g.completed >= 2 { + g.expected = g.completed; + } + } + self.try_emit_p2p_parent(gid); + } + } + + fn reclaim_ncclop(&mut self, mut op: event::NcclOp) { + if let Some(gid) = op.parent_group_id() { + if op.is_p2p() { + self.record_p2p_completion(&op); + if !self.will_emit_parent(gid) { + op.clear_parent_group_id(); + } + self.pending_telemetry + .push_back(Telemetry::NcclOp(Box::new(op))); + self.try_emit_p2p_parent(gid); + return; + } + } self.pending_telemetry .push_back(Telemetry::NcclOp(Box::new(op))); } @@ -337,12 +458,17 @@ impl<'a> PollingContext<'a> { // 1. we know the parent (and therefore comm hash) // 2. this proxyop is originated from current process OR // we are tracking interprocess proxyop + // Prefer runtime gates (comma-monitor mid-flight enable) over the + // process-start Config snapshot. + let gates = &self.profiler.gates; + let track_steps = gates.track_steps(); + let aggregate_cfg = gates.aggregate_steps(); + let track_interprocess = gates.track_interprocess_proxyop(); if info.parent().is_some() { - aggregate_steps = self.profiler.config.aggregate_steps - && (info.pid == self.profiler.pid - || self.profiler.config.track_interprocess_proxyop); + aggregate_steps = + aggregate_cfg && (info.pid == self.profiler.pid || track_interprocess); } - op.init_step_tracking(self.profiler.config.track_steps, aggregate_steps); + op.init_step_tracking(track_steps, aggregate_steps); self.handle_proxyop_start(thread_state, info, op.basic_info().start_time()); thread_state.proxyops.insert(id, op); @@ -388,8 +514,11 @@ impl<'a> PollingContext<'a> { proxyops: Vec>, send_ipc: bool, ) { - let config = &self.profiler.config; - let record_proxyop = config.track_proxyop || config.track_steps; + // Runtime gates: when the monitor escalates after an anomaly, newly + // completed ProxyOps must be attached to the parent COLL JSON even + // though process-start Config still has track_proxyop/steps=false. + let gates = &self.profiler.gates; + let record_proxyop = gates.track_proxyop() || gates.track_steps(); if let Some(parent_handle) = info.parent() { if info.pid == self.profiler.pid { if let Some(ncclop) = self.get_ncclop(parent_handle) { @@ -432,6 +561,9 @@ impl<'a> PollingContext<'a> { { match msg { Message::Group(group) => { + let gid = group.id(); + self.p2p_groups.entry(gid).or_default().ended = true; + self.try_emit_p2p_parent(gid); if self.profiler.config.track_group { self.pending_telemetry.push_back(Telemetry::Group(group)); } @@ -439,6 +571,7 @@ impl<'a> PollingContext<'a> { Message::NcclOp(op) => { let id = op.id(); let op = self.free_ncclop.take_and_free(op); + self.note_ncclop_start(&op); let _ = self.ncclops.insert(id, Box::new(op.clone())); // copybara:strip_begin(hang detection) self.pending_telemetry @@ -549,6 +682,11 @@ impl<'a> PollingContext<'a> { ); } } + Message::KernelStep(step, parent) => { + if let Some(ncclop) = self.get_ncclop(parent) { + ncclop.add_kernel_step(step); + } + } Message::CommOpen(comm) => { self.pending_telemetry.push_back(Telemetry::CommOpen(comm)); } @@ -731,6 +869,9 @@ where ctrl.daemon_state.idx = threads.len(); threads.push(ctrl); } + ControlMessage::SetGates(update) => { + let _ = ctx.profiler.gates.apply_update(&update); + } } } @@ -867,6 +1008,7 @@ where } ctx.reclaim_all_ncclops_in_map(); + ctx.flush_p2p_parents(); exporter.export(ctx, Some(RETRY_MS)); } diff --git a/src/event.rs b/src/event.rs index fc6e34b..90abc08 100644 --- a/src/event.rs +++ b/src/event.rs @@ -24,9 +24,12 @@ use crate::step_tracker::EventStep; use serde_json::json; +use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; use std::time::Instant; +static NEXT_GROUP_ID: AtomicU64 = AtomicU64::new(1); + #[derive(Debug)] pub enum Event { Group(Box), @@ -34,6 +37,7 @@ pub enum Event { NcclOp(usize), ProxyOpLite(slab::AllocatedNode), // Lite == no step tracking KernelCh(slab::AllocatedNode), + KernelStep(slab::AllocatedNode), ProxyOp(slab::AllocatedNode), Dummy(usize), SmallNcclOp(usize), @@ -109,6 +113,7 @@ pub trait ProfilerEvent { #[repr(align(16))] pub struct Group { basic_info: BasicInfo, + id: u64, } impl Group { @@ -118,8 +123,52 @@ impl Group { { Self { basic_info: BasicInfo::from_descr(descr, time), + id: NEXT_GROUP_ID.fetch_add(1, Ordering::Relaxed), } } + + pub fn id(&self) -> u64 { + self.id + } +} + +/// Parent covering every P2P send/recv in one NCCL group. +/// Emitted as `cat: P2P_GROUP` (not COLL) when a group contains 2+ P2Ps +/// and no native collective. `name` is `all_to_all` or `sendrecv`. +#[derive(Debug, Clone)] +pub struct P2pParent { + pub group_id: u64, + pub rank: usize, + pub comm_hash: u64, + pub start_time: Instant, + pub end_time: Instant, + pub size: usize, + pub n_children: usize, + pub n_peers: usize, + pub name: &'static str, +} + +impl P2pParent { + pub fn trace_record(&self, mut time_to_num: F) -> serde_json::Value + where + F: FnMut(Instant) -> u64, + { + json!({ + "ph": "X", + "ts": time_to_num(self.start_time), + "dur": (self.end_time.saturating_duration_since(self.start_time)).as_micros(), + "cat": "P2P_GROUP", + "name": self.name, + "rank": self.rank, + "comm_hash": format!("0x{:016x}", self.comm_hash), + "group_id": self.group_id, + "n_children": self.n_children, + "n_peers": self.n_peers, + "args": { + "size": self.size, + }, + }) + } } // NcclOp: collective and P2p @@ -132,6 +181,25 @@ pub struct NcclOp { comm_hash: Option, // starting from v4 comm_hash is no longer part of the event descriptor descr: nccl_metadata::EventMetadata, proxyops: Option>, + kernel_steps: Option>, + /// Stable NCCL Group id when this op was launched inside ncclGroupStart/End. + parent_group_id: Option, +} + +/// Per-slice Simple-prims KernelStep attached to a Coll/P2p NcclOp. +#[derive(Debug, Clone)] +pub struct KernelEventStep { + pub channel_id: u8, + pub is_send: bool, + pub peer: u8, + pub step: u32, + pub size: u32, + /// NCCL `startTs`: wait/step begin; 0 if none/recv. + pub start_ts: u64, + /// NCCL `readyTs`: transfer/comm begin. + pub ready_ts: u64, + /// Stop-ring end time (GPU globaltimer). + pub end_ts: u64, } #[derive(Debug, Clone)] @@ -265,9 +333,19 @@ impl NcclOp { comm_hash: comm_hash_override, descr: descr.clone_to_metadata(), proxyops: None, + kernel_steps: None, + parent_group_id: event_ffi::peek_group_id(descr.parent_obj()), } } + pub fn parent_group_id(&self) -> Option { + self.parent_group_id + } + + pub fn clear_parent_group_id(&mut self) { + self.parent_group_id = None; + } + pub fn id(&self) -> usize { self.id } @@ -330,6 +408,15 @@ impl NcclOp { } self.proxyops.as_mut().unwrap().push(proxyop); } + + pub fn add_kernel_step(&mut self, step: KernelEventStep) { + // Nest only. Parent Coll/P2p timing is refined by KernelCh / ProxyOp, + // not by KernelStep GPU spans. + if self.kernel_steps.is_none() { + self.kernel_steps = Some(Vec::new()); + } + self.kernel_steps.as_mut().unwrap().push(step); + } } impl ProfilerEvent for NcclOp { @@ -372,6 +459,9 @@ impl ProfilerEvent for NcclOp { json["name"] = json!(if p2p.is_send() { "send" } else { "recv" }); json["peer"] = json!(p2p.peer()); } + if let Some(gid) = self.parent_group_id { + json["parent"] = json!(gid); + } if let Some(proxyops) = self.proxyops.as_ref() { json["proxyops"] = proxyops @@ -380,6 +470,33 @@ impl ProfilerEvent for NcclOp { .collect(); } + if let Some(steps) = self.kernel_steps.as_ref() { + json["kernel_steps"] = steps + .iter() + .map(|s| { + // Same JSON keys as ProxyStep: start_time → fifo_ready_time → end_time. + let start_time = if s.start_ts != 0 { + s.start_ts + } else { + s.ready_ts + }; + let mut r = json!({ + "channel": s.channel_id, + "is_send": s.is_send, + "peer": s.peer, + "step": s.step, + "size": s.size, + "start_time": start_time, + "end_time": s.end_ts, + }); + if s.start_ts != 0 { + r["fifo_ready_time"] = json!(s.ready_ts); + } + r + }) + .collect(); + } + if let Some(child_start_time) = self.child_start_time { json["child_start_ts"] = json!(time_to_num(child_start_time)); if !is_reclaimed { @@ -633,4 +750,88 @@ mod tests { basic_info.update_end_time(t0); assert_eq!(basic_info.end_time(), Some(t1)); } + + #[test] + fn kernel_steps_nested_without_refining_parent_duration() { + let start = Instant::now(); + let mut descr: profiler_shim::ncclProfilerEventDescr_v4_t = unsafe { std::mem::zeroed() }; + descr.type_ = profiler_shim::ncclProfileColl as _; + let mut op = NcclOp { + basic_info: BasicInfo { + rank: 0, + start_time: start, + end_time: None, + }, + id: 0, + is_p2p: false, + child_start_time: None, + comm_hash: Some(0x123), + descr: nccl_metadata::EventMetadata::V4(profiler_shim::EventDescrV4(descr)), + proxyops: None, + kernel_steps: None, + parent_group_id: None, + }; + + op.add_kernel_step(KernelEventStep { + channel_id: 0, + is_send: true, + peer: 0, + step: 1, + size: 4, + start_ts: 1_000_000 - 2_000, // wait begin + ready_ts: 1_000_000, // transfer begin + end_ts: 1_000_000 + 5_000, // 5 us transfer + }); + op.add_kernel_step(KernelEventStep { + channel_id: 0, + is_send: false, + peer: 0, + step: 1, + size: 4, + start_ts: 0, + ready_ts: 1_000_000 + 1_000, + end_ts: 1_000_000 + 12_000, + }); + + assert!( + op.basic_info.end_time().is_none(), + "KernelSteps must not refine parent end_time (KernelCh/ProxyOp do)" + ); + assert_eq!(op.child_start_time(), None); + + let rec = op.trace_record(|_| 0); + let ks = rec["kernel_steps"].as_array().expect("kernel_steps"); + assert_eq!(ks.len(), 2); + assert_eq!(ks[0]["start_time"], 1_000_000 - 2_000); + assert_eq!(ks[0]["fifo_ready_time"], 1_000_000); + assert_eq!(ks[0]["end_time"], 1_000_000 + 5_000); + assert!(ks[1].get("fifo_ready_time").is_none()); + assert_eq!(ks[1]["start_time"], 1_000_000 + 1_000); + assert_eq!(ks[1]["end_time"], 1_000_000 + 12_000); + } + + #[test] + fn p2p_parent_emits_p2p_group() { + let start = Instant::now(); + let parent = P2pParent { + group_id: 7, + rank: 1, + comm_hash: 0xabc, + start_time: start, + end_time: start + Duration::from_micros(250), + size: 4096, + n_children: 6, + n_peers: 3, + name: "all_to_all", + }; + let rec = parent.trace_record(|_| 42); + assert_eq!(rec["cat"], "P2P_GROUP"); + assert_eq!(rec["name"], "all_to_all"); + assert_eq!(rec["group_id"], 7); + assert_eq!(rec["n_children"], 6); + assert_eq!(rec["n_peers"], 3); + assert_eq!(rec["comm_hash"], "0x0000000000000abc"); + assert_eq!(rec["args"]["size"], 4096); + assert_eq!(rec["dur"], 250); + } } diff --git a/src/event_ffi.rs b/src/event_ffi.rs index 1b81c71..b0335d0 100644 --- a/src/event_ffi.rs +++ b/src/event_ffi.rs @@ -13,7 +13,7 @@ // limitations under the License. use crate::event::{Event, Group, ProxyStep}; -use crate::profiler::{KernelCh, ProxyOpLocalData}; +use crate::profiler::{KernelCh, KernelStepLocal, ProxyOpLocalData}; use crate::slab; use static_assertions::const_assert; @@ -34,6 +34,7 @@ pub enum Type { SmallNcclOp, ProxyStep, KernelCh, + KernelStep, } impl Type { @@ -48,6 +49,7 @@ impl Type { Type::SmallNcclOp => 0b110, Type::ProxyStep => 0b111, Type::KernelCh => 0b1000, + Type::KernelStep => 0b1001, } } @@ -62,6 +64,7 @@ impl Type { 0b110 => Type::SmallNcclOp, 0b111 => Type::ProxyStep, 0b1000 => Type::KernelCh, + 0b1001 => Type::KernelStep, _ => panic!("unknown bit pattern"), } } @@ -85,6 +88,22 @@ fn to_ncclop_handle(handle: Handle) -> usize { get_handle_inner(handle) >> N_TYPE_BITS } +/// Read the monotonic Group id from a live NCCL Group handle without taking +/// ownership. Returns None when `handle` is not a Group. +pub fn peek_group_id(handle: Handle) -> Option { + if handle.is_null() || get_handle_type(handle) != Type::Group { + return None; + } + let ptr = get_handle_inner(handle) as *const Group; + if ptr.is_null() { + None + } else { + // SAFETY: NCCL holds the Group event until Group stop. P2P/COLL start + // runs while that parent Group handle is still live. + Some(unsafe { (*ptr).id() }) + } +} + pub trait AsFFI: Sized { fn into_ffi(self) -> Handle; @@ -97,6 +116,7 @@ pub trait AsFFI: Sized { const_assert!(std::mem::align_of::() >= (1 << N_TYPE_BITS)); const_assert!(std::mem::align_of::() >= (1 << N_TYPE_BITS)); const_assert!(std::mem::align_of::() >= (1 << N_TYPE_BITS)); +const_assert!(std::mem::align_of::() >= (1 << N_TYPE_BITS)); const_assert!(std::mem::align_of::() >= (1 << N_TYPE_BITS)); impl AsFFI for Event { @@ -126,6 +146,10 @@ impl AsFFI for Event { let ptr = slab::AllocatedNode::into_raw(op); handle(ptr as _, Type::KernelCh) } + Event::KernelStep(op) => { + let ptr = slab::AllocatedNode::into_raw(op); + handle(ptr as _, Type::KernelStep) + } /* Event::ProxyOp(id) => { let v = (id as usize) << N_TYPE_BITS; @@ -186,6 +210,11 @@ impl AsFFI for Event { let proxyop = slab::AllocatedNode::from_raw(handle); Some(Event::KernelCh(proxyop)) } + Type::KernelStep => { + let handle = get_handle_inner(handle) as _; + let step = slab::AllocatedNode::from_raw(handle); + Some(Event::KernelStep(step)) + } /* Type::ProxyOp => { let handle = get_handle_inner(handle); @@ -295,6 +324,20 @@ mod tests { } } + #[test] + fn peek_group_id_matches_live_handle() { + let group_descr = dummy_group_descr(); + let event = Event::new_group(&group_descr, Instant::now()); + let Event::Group(ref group) = event else { + panic!("expected Group"); + }; + let want = group.id(); + let handle = Event::into_ffi(event); + assert_eq!(peek_group_id(handle), Some(want)); + // SAFETY: handle came from into_ffi in this test. + let _ = unsafe { Event::from_ffi(handle) }; + } + #[test] fn proxyop_event_mock() { use crate::nccl_metadata::Version as _; diff --git a/src/lib.rs b/src/lib.rs index 5e2c600..d062fcf 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -15,6 +15,7 @@ pub mod clock; mod cloud_daemon; mod config; +mod control_rpc; mod daemon; mod event; mod event_ffi; @@ -26,6 +27,7 @@ mod nccl_metadata; mod otel_utils; // copybara:strip(otel) mod profiler; pub mod profiler_shim; +mod runtime_gates; mod shm_fifo; mod slab; mod spsc; @@ -34,7 +36,9 @@ mod step_tracker; use std::sync::OnceLock; use event_ffi::AsFFI as _; -use profiler_shim::{ncclResult_t, EventDescrV1, EventDescrV2, EventDescrV3, EventDescrV4}; +use profiler_shim::{ + ncclResult_t, EventDescrV1, EventDescrV2, EventDescrV3, EventDescrV4, EventDescrV6, +}; /// # Safety /// @@ -194,6 +198,43 @@ unsafe extern "C" fn profiler_init_v4( } } +#[allow(clippy::missing_safety_doc)] +unsafe extern "C" fn profiler_init_v6( + context: *mut *mut libc::c_void, + comm_id: u64, + e_activation_mask: *mut i32, + comm_name: *const libc::c_char, + n_nodes: i32, + n_ranks: i32, + rank: i32, + log_fn: profiler_shim::ncclDebugLogger_t, +) -> ncclResult_t { + { + let mut lg = LOGGER_LOCK.lock().unwrap(); + if !*lg { + let logger = LOGGER.get_or_init(|| NcclLogger(log_fn)); + log::set_logger(logger) + .map(|()| log::set_max_level(log::LevelFilter::Trace)) + .unwrap(); + *lg = true; + } + } + match profiler::init_handler_v6( + &mut *e_activation_mask, + comm_name, + comm_id, + n_nodes, + n_ranks, + rank, + ) { + Ok(comm) => { + *context = Box::into_raw(comm) as _; + profiler_shim::ncclResult_t_ncclSuccess + } + Err(e) => e, + } +} + #[allow(clippy::missing_safety_doc)] unsafe extern "C" fn profiler_start_event_v1( context: *mut libc::c_void, @@ -258,6 +299,22 @@ unsafe extern "C" fn profiler_start_event_v4( } } +#[allow(clippy::missing_safety_doc)] +unsafe extern "C" fn profiler_start_event_v6( + context: *mut libc::c_void, + e_handle: *mut *mut libc::c_void, + e_descr: *mut profiler_shim::ncclProfilerEventDescr_v6_t, +) -> ncclResult_t { + let descr = &*(e_descr as *const EventDescrV6); + match profiler::start_event_handler(descr, context as _) { + Ok(event) => { + *e_handle = event.map_or(std::ptr::null_mut(), event::Event::into_ffi); + profiler_shim::ncclResult_t_ncclSuccess + } + Err(e) => e, + } +} + #[allow(clippy::missing_safety_doc)] unsafe extern "C" fn profiler_stop_event(e_handle: *mut libc::c_void) -> ncclResult_t { if e_handle.is_null() { @@ -367,6 +424,55 @@ unsafe extern "C" fn profiler_record_event_state_v4( profiler_shim::ncclResult_t_ncclSuccess } +#[allow(clippy::missing_safety_doc)] +unsafe extern "C" fn profiler_record_event_state_v6( + e_handle: *mut libc::c_void, + e_state: profiler_shim::ncclProfilerEventState_v6_t, + e_state_args: *mut profiler_shim::ncclProfilerEventStateArgs_v6_t, +) -> ncclResult_t { + let handle_type = event_ffi::get_handle_type(e_handle); + match handle_type { + event_ffi::Type::ProxyStep => { + if let Some(mut event) = event::Event::from_ffi(e_handle) { + let step_state = nccl_metadata::ProxyStepStateV6::cast_from_union(&*e_state_args); + if let Err(e) = + profiler::record_proxystep_event_state_handler(&mut event, e_state, step_state) + { + return e; + } + let _ = event::Event::into_ffi(event); + } + } + + event_ffi::Type::ProxyOpLite => { + if let Some(mut event) = event::Event::from_ffi(e_handle) { + if let Err(e) = profiler::record_proxyop_event_state_handler_v4(&mut event, e_state) + { + return e; + } + let _ = event::Event::into_ffi(event); + } + } + + event_ffi::Type::KernelStep => { + if let Some(mut event) = event::Event::from_ffi(e_handle) { + let step_state = nccl_metadata::ProxyStepStateV6::cast_from_union(&*e_state_args); + if let Err(e) = profiler::record_kernelstep_event_state_handler( + &mut event, + e_state, + step_state.kernel_step_ptimer(), + ) { + return e; + } + let _ = event::Event::into_ffi(event); + } + } + + _ => (), + } + profiler_shim::ncclResult_t_ncclSuccess +} + #[allow(clippy::missing_safety_doc)] unsafe extern "C" fn profiler_finalize(context: *mut libc::c_void) -> ncclResult_t { if let Err(e) = profiler::finalize_handler(Box::from_raw(context.cast())) { @@ -383,6 +489,7 @@ unsafe impl Sync for profiler_shim::ncclProfiler_v1_t {} unsafe impl Sync for profiler_shim::ncclProfiler_v2_t {} unsafe impl Sync for profiler_shim::ncclProfiler_v3_t {} unsafe impl Sync for profiler_shim::ncclProfiler_v4_t {} +unsafe impl Sync for profiler_shim::ncclProfiler_v6_t {} #[allow(non_upper_case_globals)] #[no_mangle] @@ -432,6 +539,18 @@ pub static ncclProfiler_v4: profiler_shim::ncclProfiler_v4_t = profiler_shim::nc finalize: Some(profiler_finalize), }; +#[allow(non_upper_case_globals)] +#[no_mangle] +pub static ncclProfiler_v6: profiler_shim::ncclProfiler_v6_t = profiler_shim::ncclProfiler_v6_t { + // SAFETY: string has no interior NUL bytes + name: unsafe { static_cstr!("GCP_NCCL_PROFILER_V6").as_ptr() }, + init: Some(profiler_init_v6), + startEvent: Some(profiler_start_event_v6), + stopEvent: Some(profiler_stop_event), + recordEventState: Some(profiler_record_event_state_v6), + finalize: Some(profiler_finalize), +}; + // Helper for testing the Profiler type #[cfg(test)] fn scoped_profiler_test(profiler: profiler::Profiler, f: F) diff --git a/src/nccl_metadata.rs b/src/nccl_metadata.rs index e1dcdf3..3c02a2b 100644 --- a/src/nccl_metadata.rs +++ b/src/nccl_metadata.rs @@ -99,6 +99,7 @@ pub enum NcclOpType { AllGather, ReduceScatter, AllReduce, + AlltoAll, Send, Recv, Unknown, @@ -111,6 +112,7 @@ static NCCLOP_NAME_LOOKUP: LazyLock> = LazyLo (c"AllGather", NcclOpType::AllGather), (c"ReduceScatter", NcclOpType::ReduceScatter), (c"AllReduce", NcclOpType::AllReduce), + (c"AlltoAll", NcclOpType::AlltoAll), (c"Send", NcclOpType::Send), (c"Recv", NcclOpType::Recv), ] @@ -129,6 +131,7 @@ impl NcclOpType { 4 => NcclOpType::AllReduce, 6 => NcclOpType::Send, 7 => NcclOpType::Recv, + 8 => NcclOpType::AlltoAll, _ => NcclOpType::Unknown, } } @@ -147,6 +150,7 @@ impl NcclOpType { b'A' => match bytes[last_idx] { b'r' => NcclOpType::AllGather, b'e' => NcclOpType::AllReduce, + b'l' => NcclOpType::AlltoAll, _ => NcclOpType::Unknown, }, b'S' => NcclOpType::Send, @@ -175,6 +179,7 @@ impl NcclOpType { NcclOpType::AllGather => "all_gather", NcclOpType::ReduceScatter => "reduce_scatter", NcclOpType::AllReduce => "all_reduce", + NcclOpType::AlltoAll => "all_to_all", NcclOpType::Send => "send", NcclOpType::Recv => "recv", _ => "unknown", @@ -188,6 +193,7 @@ impl NcclOpType { NcclOpType::AllGather => c"all_gather", NcclOpType::ReduceScatter => c"reduce_scatter", NcclOpType::AllReduce => c"all_reduce", + NcclOpType::AlltoAll => c"all_to_all", NcclOpType::Send => c"send", NcclOpType::Recv => c"recv", _ => c"unknown", @@ -539,6 +545,7 @@ pub enum EventMetadata { V2(profiler_shim::EventDescrV2), V3(profiler_shim::EventDescrV3), V4(profiler_shim::EventDescrV4), + V6(profiler_shim::EventDescrV6), } /// # Safety @@ -555,6 +562,7 @@ impl EventMetadata { Self::V2(descr) => NcclOpKey::from_descr(descr, alt_comm_hash), Self::V3(descr) => NcclOpKey::from_descr(descr, alt_comm_hash), Self::V4(descr) => NcclOpKey::from_descr(descr, alt_comm_hash), + Self::V6(descr) => NcclOpKey::from_descr(descr, alt_comm_hash), } } @@ -568,6 +576,7 @@ impl EventMetadata { Self::V2(descr) => event_byte_count(descr), Self::V3(descr) => event_byte_count(descr), Self::V4(descr) => event_byte_count(descr), + Self::V6(descr) => event_byte_count(descr), } } @@ -577,6 +586,7 @@ impl EventMetadata { Self::V2(descr) => descr.try_cast_to_coll().map(|x| x as _), Self::V3(descr) => descr.try_cast_to_coll().map(|x| x as _), Self::V4(descr) => descr.try_cast_to_coll().map(|x| x as _), + Self::V6(descr) => descr.try_cast_to_coll().map(|x| x as _), } } @@ -586,13 +596,14 @@ impl EventMetadata { Self::V2(descr) => descr.try_cast_to_p2p().map(|x| x as _), Self::V3(descr) => descr.try_cast_to_p2p().map(|x| x as _), Self::V4(descr) => descr.try_cast_to_p2p().map(|x| x as _), + Self::V6(descr) => descr.try_cast_to_p2p().map(|x| x as _), } } } pub trait Event { fn rank(&self) -> i32; - fn type_(&self) -> u8; + fn type_(&self) -> u64; fn parent_obj(&self) -> *mut libc::c_void; fn clone_to_metadata(&self) -> EventMetadata; } @@ -602,6 +613,7 @@ pub trait Version: Event { type P2p: P2p; type ProxyOp: ProxyOp; type ProxyStep: ProxyStep; + type KernelStep: KernelStep; fn version() -> profiler::Version; @@ -641,6 +653,15 @@ pub trait Version: Event { &*ptr.cast() } + /// # Safety + /// + /// type of this descriptor must be kernelstep (API v6+) + #[inline(always)] + unsafe fn cast_to_kernelstep(&self) -> &Self::KernelStep { + let ptr = self as *const Self; + &*ptr.cast() + } + #[inline(always)] fn try_cast_to_coll(&self) -> Option<&Self::Coll> { if self.type_() as u32 == profiler_shim::ncclProfileColl { @@ -725,6 +746,37 @@ pub trait ProxyStep: Event { fn step(&self) -> i32; } +/// Per-slice Simple-prims kernel step (profiler API v6+). +pub trait KernelStep: Event { + fn channel_id(&self) -> u8; + fn is_send(&self) -> bool; + fn peer(&self) -> u8; + fn step(&self) -> u32; + fn size(&self) -> u32; + /// Wait/step begin (GPU globaltimer); 0 if none/recv. + fn start_ts(&self) -> u64; + /// Transfer/comm begin (GPU globaltimer). + fn ready_ts(&self) -> u64; + /// CoMMA ProxyStep-style wait: `ready_ts - start_ts` when `start_ts != 0`. + fn fifo_wait_dur_ns(&self) -> u32 { + derive_fifo_wait_dur_ns(self.start_ts(), self.ready_ts()) + } +} + +/// Derive ProxyStep-compatible `fifo_wait_dur_ns` from NCCL `startTs`/`readyTs`. +#[inline(always)] +pub fn derive_fifo_wait_dur_ns(start_ts: u64, ready_ts: u64) -> u32 { + if start_ts == 0 || ready_ts <= start_ts { + return 0; + } + let d = ready_ts - start_ts; + if d > u32::MAX as u64 { + u32::MAX + } else { + d as u32 + } +} + pub trait ProxyOpState { fn version() -> profiler::Version; fn steps(&self) -> i32; @@ -744,8 +796,8 @@ macro_rules! impl_event { } #[inline(always)] - fn type_(&self) -> u8 { - self.0.type_() + fn type_(&self) -> u64 { + self.0.type_() as u64 } #[inline(always)] @@ -770,8 +822,8 @@ macro_rules! descr_impl_event { } #[inline(always)] - fn type_(&self) -> u8 { - self.0.type_ + fn type_(&self) -> u64 { + self.0.type_ as u64 } #[inline(always)] @@ -959,6 +1011,84 @@ macro_rules! def_proxystep { }; } +/// Stub KernelStep for profiler API versions that lack the union field. +/// Must never be cast to outside of dead code paths (mask gated on V6). +macro_rules! def_kernelstep_unsupported { + ($t:tt, $d:tt) => { + #[allow(unused_parens)] + #[repr(transparent)] + pub struct $t($d); + impl super::KernelStep for $t { + fn channel_id(&self) -> u8 { + panic!("KernelStep unsupported on this profiler API version") + } + fn is_send(&self) -> bool { + panic!("KernelStep unsupported on this profiler API version") + } + fn peer(&self) -> u8 { + panic!("KernelStep unsupported on this profiler API version") + } + fn step(&self) -> u32 { + panic!("KernelStep unsupported on this profiler API version") + } + fn size(&self) -> u32 { + panic!("KernelStep unsupported on this profiler API version") + } + fn start_ts(&self) -> u64 { + panic!("KernelStep unsupported on this profiler API version") + } + fn ready_ts(&self) -> u64 { + panic!("KernelStep unsupported on this profiler API version") + } + } + }; +} + +macro_rules! def_kernelstep { + ($t:tt, $d:tt) => { + #[allow(unused_parens)] + #[repr(transparent)] + pub struct $t($d); + impl super::KernelStep for $t { + #[inline(always)] + fn channel_id(&self) -> u8 { + let s = unsafe { &self.0 .0.__bindgen_anon_1.kernelStep }; + s.channelId + } + #[inline(always)] + fn is_send(&self) -> bool { + let s = unsafe { &self.0 .0.__bindgen_anon_1.kernelStep }; + s.isSend != 0 + } + #[inline(always)] + fn peer(&self) -> u8 { + let s = unsafe { &self.0 .0.__bindgen_anon_1.kernelStep }; + s.peer + } + #[inline(always)] + fn step(&self) -> u32 { + let s = unsafe { &self.0 .0.__bindgen_anon_1.kernelStep }; + s.step + } + #[inline(always)] + fn size(&self) -> u32 { + let s = unsafe { &self.0 .0.__bindgen_anon_1.kernelStep }; + s.size + } + #[inline(always)] + fn start_ts(&self) -> u64 { + let s = unsafe { &self.0 .0.__bindgen_anon_1.kernelStep }; + s.startTs + } + #[inline(always)] + fn ready_ts(&self) -> u64 { + let s = unsafe { &self.0 .0.__bindgen_anon_1.kernelStep }; + s.readyTs + } + } + }; +} + mod v1 { use super::algo::IntoAlgo; use super::proto::IntoProto; @@ -986,6 +1116,9 @@ mod v1 { def_proxystep!(ProxyStep, (profiler_shim::EventDescrV1)); impl_event!(ProxyStep, EventMetadata::V1); + + def_kernelstep_unsupported!(KernelStep, (profiler_shim::EventDescrV1)); + impl_event!(KernelStep, EventMetadata::V1); } mod v2 { @@ -1015,6 +1148,9 @@ mod v2 { def_proxystep!(ProxyStep, (profiler_shim::EventDescrV2)); impl_event!(ProxyStep, EventMetadata::V2); + + def_kernelstep_unsupported!(KernelStep, (profiler_shim::EventDescrV2)); + impl_event!(KernelStep, EventMetadata::V2); } mod v3 { @@ -1044,6 +1180,9 @@ mod v3 { def_proxystep!(ProxyStep, (profiler_shim::EventDescrV3)); impl_event!(ProxyStep, EventMetadata::V3); + + def_kernelstep_unsupported!(KernelStep, (profiler_shim::EventDescrV3)); + impl_event!(KernelStep, EventMetadata::V3); } mod v4 { @@ -1109,6 +1248,77 @@ mod v4 { def_proxystep!(ProxyStep, (profiler_shim::EventDescrV4)); impl_event!(ProxyStep, EventMetadata::V4); + + def_kernelstep_unsupported!(KernelStep, (profiler_shim::EventDescrV4)); + impl_event!(KernelStep, EventMetadata::V4); +} + +mod v6 { + use super::algo::IntoAlgo; + use super::proto::IntoProto; + use super::*; + + #[repr(transparent)] + pub struct Coll(profiler_shim::EventDescrV6); + + impl_event!(Coll, EventMetadata::V6); + + impl NcclOp for Coll { + // SAFETY: this type could only be constructed via cast_*(), + // which must be called with type_ == ncclProfileColl. + // Therefore accessing the corresponding union field is safe. + + #[inline(always)] + fn comm_hash(&self) -> Option { + None + } + + #[inline(always)] + fn byte_count(&self) -> usize { + let coll = unsafe { &self.0 .0.__bindgen_anon_1.coll }; + let dt_bytes = unsafe { datatype_c_str_ptr_to_nbytes(coll.datatype) }; + (coll.count * dt_bytes) as _ + } + } + + impl_coll!(Coll, nChannels); + + #[repr(transparent)] + pub struct P2p(profiler_shim::EventDescrV6); + + impl_event!(P2p, EventMetadata::V6); + + impl NcclOp for P2p { + // SAFETY: this type could only be constructed via cast_*(), + // which must be called with type_ == ncclProfileP2p. + // Therefore accessing the corresponding union field is safe. + + #[inline(always)] + fn comm_hash(&self) -> Option { + None + } + + #[inline(always)] + fn byte_count(&self) -> usize { + let p2p = unsafe { &self.0 .0.__bindgen_anon_1.p2p }; + let dt_bytes = unsafe { datatype_c_str_ptr_to_nbytes(p2p.datatype) }; + (p2p.count * dt_bytes) as _ + } + } + + impl_p2p!(P2p); + + #[repr(transparent)] + pub struct ProxyOp(profiler_shim::EventDescrV6); + + impl_event!(ProxyOp, EventMetadata::V6); + impl_proxyop!(ProxyOp); + + def_proxystep!(ProxyStep, (profiler_shim::EventDescrV6)); + impl_event!(ProxyStep, EventMetadata::V6); + + def_kernelstep!(KernelStep, (profiler_shim::EventDescrV6)); + impl_event!(KernelStep, EventMetadata::V6); } descr_impl_event!(profiler_shim::EventDescrV1, EventMetadata::V1); @@ -1118,6 +1328,7 @@ impl Version for profiler_shim::EventDescrV1 { type P2p = v1::P2p; type ProxyOp = v1::ProxyOp; type ProxyStep = v1::ProxyStep; + type KernelStep = v1::KernelStep; fn version() -> profiler::Version { profiler::Version::V1 @@ -1131,6 +1342,7 @@ impl Version for profiler_shim::EventDescrV2 { type P2p = v2::P2p; type ProxyOp = v2::ProxyOp; type ProxyStep = v2::ProxyStep; + type KernelStep = v2::KernelStep; fn version() -> profiler::Version { profiler::Version::V2 @@ -1144,6 +1356,7 @@ impl Version for profiler_shim::EventDescrV3 { type P2p = v3::P2p; type ProxyOp = v3::ProxyOp; type ProxyStep = v3::ProxyStep; + type KernelStep = v3::KernelStep; fn version() -> profiler::Version { profiler::Version::V3 @@ -1157,12 +1370,27 @@ impl Version for profiler_shim::EventDescrV4 { type P2p = v4::P2p; type ProxyOp = v4::ProxyOp; type ProxyStep = v4::ProxyStep; + type KernelStep = v4::KernelStep; fn version() -> profiler::Version { profiler::Version::V4 } } +descr_impl_event!(profiler_shim::EventDescrV6, EventMetadata::V6); + +impl Version for profiler_shim::EventDescrV6 { + type Coll = v6::Coll; + type P2p = v6::P2p; + type ProxyOp = v6::ProxyOp; + type ProxyStep = v6::ProxyStep; + type KernelStep = v6::KernelStep; + + fn version() -> profiler::Version { + profiler::Version::V6 + } +} + #[repr(transparent)] pub struct ProxyOpStateV1(profiler_shim::ncclProfilerEventStateArgs_v1_t); @@ -1240,10 +1468,47 @@ impl ProxyStepState for ProxyStepStateV4 { } } +/// V6 state args are typedef'd to v5 (includes kernelStep.pTimer). +#[repr(transparent)] +pub struct ProxyStepStateV6(profiler_shim::ncclProfilerEventStateArgs_v6_t); + +impl ProxyStepStateV6 { + /// # Safety + /// + /// input must be known to be a proxystep variant. + pub unsafe fn cast_from_union(u: &profiler_shim::ncclProfilerEventStateArgs_v6_t) -> &Self { + let ptr = u as *const profiler_shim::ncclProfilerEventStateArgs_v6_t; + &*ptr.cast() + } + + #[inline(always)] + pub fn kernel_step_ptimer(&self) -> u64 { + unsafe { self.0.kernelStep.pTimer } + } +} + +impl ProxyStepState for ProxyStepStateV6 { + fn trans_size(&self) -> usize { + unsafe { self.0.proxyStep.transSize } + } +} + #[cfg(test)] mod tests { use super::*; + #[test] + fn derive_fifo_wait_from_timestamps() { + assert_eq!(derive_fifo_wait_dur_ns(0, 1_000_000), 0); + assert_eq!(derive_fifo_wait_dur_ns(1_000_000, 1_000_000), 0); + assert_eq!(derive_fifo_wait_dur_ns(1_000_000, 999_999), 0); + assert_eq!(derive_fifo_wait_dur_ns(1_000_000, 1_002_000), 2_000); + assert_eq!( + derive_fifo_wait_dur_ns(1, 1 + (u32::MAX as u64) + 10), + u32::MAX + ); + } + #[test] fn to_json_lossless() { let op = NcclOpKey::Collective(0x123, 42, NcclOpType::AllGather, algo::RING, proto::LL128); diff --git a/src/otel_utils.rs b/src/otel_utils.rs index f7736b1..9470c1d 100644 --- a/src/otel_utils.rs +++ b/src/otel_utils.rs @@ -248,6 +248,7 @@ fn coll_type_to_num(t: nccl_metadata::NcclOpType) -> u32 { T::AllGather => 3, T::ReduceScatter => 4, T::AllReduce => 5, + T::AlltoAll => 6, _ => 0xabcd, // we don't use zero as zero span ID is invalid } } @@ -261,6 +262,7 @@ fn ncclop_otel_name(op_type: nccl_metadata::NcclOpType) -> &'static str { NcclOpType::AllGather => "ncclAllGather", NcclOpType::ReduceScatter => "ncclReduceScatter", NcclOpType::AllReduce => "ncclAllReduce", + NcclOpType::AlltoAll => "ncclAlltoAll", NcclOpType::Send => "ncclSend", NcclOpType::Recv => "ncclRecv", _ => "unknown nccl op", diff --git a/src/profiler.rs b/src/profiler.rs index 500f72b..ab8ced8 100644 --- a/src/profiler.rs +++ b/src/profiler.rs @@ -22,8 +22,9 @@ use crate::event_ffi::AsFFI as _; use crate::gpuviz; use crate::nccl_metadata; use crate::nccl_metadata::ProxyOp; -use crate::nccl_metadata::{Coll as _, Event as _, NcclOp as _, P2p as _, ProxyStep as _}; +use crate::nccl_metadata::{Coll as _, Event as _, KernelStep as _, NcclOp as _, P2p as _, ProxyStep as _}; use crate::profiler_shim; +use crate::runtime_gates::RuntimeGates; use crate::slab; use crate::spsc; use crate::step_tracker::StepTracker; @@ -46,6 +47,7 @@ pub enum Version { V2, V3, V4, + V6, } /// struct that holds global states for profiler @@ -58,6 +60,8 @@ pub struct Profiler { pub init_instant: Instant, pub gpuviz_lib: Option>, // copybara:strip(gpuviz) pub ctrl_fifo: ArrayQueue, + /// Runtime start-gate + dynamic NCCL activation mask. + pub gates: RuntimeGates, daemon: Mutex>, ncclop_cnt: AtomicU64, @@ -99,6 +103,7 @@ impl Profiler { }, // copybara:strip_end ctrl_fifo: ArrayQueue::new(CTRL_FIFO_SZ), + gates: RuntimeGates::new(config, version), daemon: Mutex::new(None), ncclop_cnt: AtomicU64::new(0), @@ -144,6 +149,7 @@ impl Profiler { let mut lg = self.daemon.lock().unwrap(); let daemon = daemon::Daemon::new(self); *lg = Some(daemon); + crate::control_rpc::spawn_control_server(self); } pub fn join_daemon(&'static self) { @@ -174,7 +180,7 @@ pub fn thread_local_state(profiler: &'_ Profiler) -> (ThreadLocalState<'_>, daem ncclop_tx.set_batch(std::cmp::max(1, profiler.config.fifo_batch_size)); fifo_tx.set_batch(std::cmp::max(1, profiler.config.fifo_batch_size)); - let n_free_proxystep = if profiler.version == Version::V4 { + let n_free_proxystep = if matches!(profiler.version, Version::V4 | Version::V6) { slab::FREELIST_BATCH } else { 0 @@ -190,6 +196,7 @@ pub fn thread_local_state(profiler: &'_ Profiler) -> (ThreadLocalState<'_>, daem proxystep_free_list: slab::FreeList::new_list(n_free_proxystep), steps_free_list: slab::FreeList::new_list(slab::FREELIST_BATCH * 4), kernelch_free_list: slab::FreeList::new_list(slab::FREELIST_BATCH), + kernelstep_free_list: slab::FreeList::new_list(slab::FREELIST_BATCH * 4), proxyop_id: 0, rng: SmallRng::from_rng(&mut rand::rng()), }; @@ -230,6 +237,7 @@ pub struct ThreadLocalState<'a> { pub proxystep_free_list: slab::FreeList, pub steps_free_list: slab::FreeList, pub kernelch_free_list: slab::FreeList, + pub kernelstep_free_list: slab::FreeList, proxyop_id: u32, rng: SmallRng, @@ -347,6 +355,20 @@ pub struct KernelCh { pub parent_op: Option, } +#[derive(Debug)] +#[repr(align(16))] +pub struct KernelStepLocal { + pub parent_op: Option, + pub channel_id: u8, + pub is_send: bool, + pub peer: u8, + pub step: u32, + pub size: u32, + pub start_ts: u64, + pub ready_ts: u64, + pub end_ts: Option, +} + #[derive(Debug)] #[repr(align(16))] pub struct ProxyOpLocalData { @@ -461,20 +483,14 @@ pub fn init_handler_v4( } *lg += 1; - mask |= profiler_shim::ncclProfileGroup; - mask |= profiler_shim::ncclProfileColl; - mask |= profiler_shim::ncclProfileP2p; - mask |= profiler_shim::ncclProfileProxyOp; - - if config::CONFIG.track_kernel_ch { - mask |= profiler_shim::ncclProfileKernelCh; - } - - if config::CONFIG.track_steps | config::CONFIG.aggregate_steps { - mask |= profiler_shim::ncclProfileProxyStep; - } + let profiler = PROFILER.get().unwrap(); + // Retain NCCL's process-global mask pointer so the daemon can rewrite it. + profiler + .gates + .register_mask_ptr(e_activation_mask as *mut i32); + mask = profiler.gates.publish_mask(); } - *e_activation_mask = mask as i32; + *e_activation_mask = mask; let mut comm = Box::new(Communicator::new()); comm.comm_hash = Some(comm_hash); if config::CONFIG.heartbeat_collective_progress { @@ -485,6 +501,25 @@ pub fn init_handler_v4( Ok(comm) } +pub fn init_handler_v6( + e_activation_mask: &mut i32, + comm_name: *const libc::c_char, + comm_hash: u64, + n_nodes: i32, + n_ranks: i32, + rank: i32, +) -> NcclResult> { + init_handler_v4( + e_activation_mask, + comm_name, + comm_hash, + n_nodes, + n_ranks, + rank, + Version::V6, + ) +} + pub fn start_event_handler( descr: &E, comm: *const Communicator, @@ -493,11 +528,15 @@ where E: nccl_metadata::Version + nccl_metadata::Event, { let config: &config::Config = &config::CONFIG; + let gates = PROFILER.get().map(|p| &p.gates); let event = match descr.type_() as u32 { profiler_shim::ncclProfileGroup => { - if config.track_ncclop { + let track_ncclop = gates.map(|g| g.track_ncclop()).unwrap_or(config.track_ncclop); + let track_group = gates.map(|g| g.track_group()).unwrap_or(config.track_group); + let track_proxyop = gates.map(|g| g.track_proxyop()).unwrap_or(config.track_proxyop); + if track_ncclop || track_group { Some(event::Event::new_group(descr, Instant::now())) - } else if config.track_proxyop { + } else if track_proxyop { Some(event::Event::new_dummyop(42)) } else { None @@ -510,9 +549,17 @@ where let comm = unsafe { &*comm }; let byte_count = descr.byte_count(); let op_type = descr.op_type(); - if config.track_ncclop { + let track_ncclop = gates.map(|g| g.track_ncclop()).unwrap_or(config.track_ncclop); + let track_proxyop = gates.map(|g| g.track_proxyop()).unwrap_or(config.track_proxyop); + if track_ncclop { use nccl_metadata::NcclOpType; - let skip_step_tracking = !(config.track_steps || config.aggregate_steps); + let track_steps = gates + .map(|g| g.proxy_step_enabled()) + .unwrap_or(config.track_steps || config.aggregate_steps); + let track_kernel_step = gates + .map(|g| g.track_kernel_step()) + .unwrap_or(config.track_kernel_step); + let skip_step_tracking = !(track_steps || track_kernel_step); let skip_small_msg = config.skip_small_collective || config.skip_small_collective_steps; if skip_small_msg @@ -565,7 +612,7 @@ where .unwrap_or_else(|| event::Event::new_dummyop(42)) })) } - } else if config.track_proxyop { + } else if track_proxyop { let comm_hash = comm .comm_hash .unwrap_or_else(|| descr.comm_hash().unwrap_or(42)); @@ -579,6 +626,11 @@ where let descr = unsafe { descr.cast_to_p2p() }; // SAFETY: for coll and p2p, NCCL guarantees the comm pointer is valid let comm = unsafe { &*comm }; + // AllToAll/self-copy emits peer==rank P2P with no GPU KernelStep path. + // Drop those no-ops so strict coverage validators see only real send/recv. + if descr.peer() == descr.rank() { + return Ok(None); + } let should_sample = with_thread_state(|thread_state| { if thread_state.rnd_decision(thread_state.profiler.config.p2p_sample_rate) { if descr.is_send() { @@ -592,7 +644,13 @@ where }); if should_sample { - if config.track_ncclop { + let track_ncclop = gates + .map(|g| g.track_ncclop()) + .unwrap_or(config.track_ncclop); + let track_proxyop = gates + .map(|g| g.track_proxyop()) + .unwrap_or(config.track_proxyop); + if track_ncclop { let byte_count = descr.byte_count(); if byte_count > config.small_msg_threshold { with_thread_state(|thread_state| { @@ -610,7 +668,7 @@ where } else { None } - } else if config.track_proxyop { + } else if track_proxyop { let comm_hash = comm .comm_hash .unwrap_or_else(|| descr.comm_hash().unwrap_or(42)); @@ -637,6 +695,12 @@ where with_thread_state(|thread_state| { let profiler = thread_state.profiler; let proxyop_id = thread_state.next_proxyop_id(); + let track_fifo = gates + .map(|g| g.track_step_fifo_wait()) + .unwrap_or(profiler.config.track_step_fifo_wait); + let track_recv = gates + .map(|g| g.track_recv_steps()) + .unwrap_or(profiler.config.track_recv_steps); let comm: Option = // SAFETY: comm is valid when pids match if config.heartbeat_collective_progress && profiler.pid == descr.pid() { @@ -648,13 +712,13 @@ where proxyop_id, descr, E::version() == Version::V1, - profiler.config.track_step_fifo_wait, + track_fifo, comm, ); let is_send = local_data.info.is_send; let mut skip_step_tracking = parent_type == event_ffi::Type::NcclOpLite; - if !skip_step_tracking && !is_send && !profiler.config.track_recv_steps { + if !skip_step_tracking && !is_send && !track_recv { skip_step_tracking = true; } @@ -682,65 +746,124 @@ where } } profiler_shim::ncclProfileKernelCh => { - let parent_ffi = descr.parent_obj(); - let parent_type = event_ffi::get_handle_type(parent_ffi); - if parent_ffi.is_null() || parent_type == event_ffi::Type::SmallNcclOp { + let allow_ch = gates + .map(|g| g.track_kernel_ch()) + .unwrap_or(config.track_kernel_ch); + if !allow_ch { None } else { - let parent = event_ffi::ProxyParent::from_ffi(parent_ffi); - if let event_ffi::ProxyParent::NcclOp(ncclop) = parent { - with_thread_state(|thread_state| { - thread_state.inc_ncclop_ref(thread_state.profiler.pid, ncclop); - let kernelch = thread_state - .kernelch_free_list - .alloc_new( - KernelCh { - parent_op: Some(ncclop), - }, - None, - true, - ) - .unwrap(); - Some(event::Event::KernelCh(kernelch)) - }) - } else { + let parent_ffi = descr.parent_obj(); + let parent_type = event_ffi::get_handle_type(parent_ffi); + if parent_ffi.is_null() || parent_type == event_ffi::Type::SmallNcclOp { None + } else { + let parent = event_ffi::ProxyParent::from_ffi(parent_ffi); + if let event_ffi::ProxyParent::NcclOp(ncclop) = parent { + with_thread_state(|thread_state| { + thread_state.inc_ncclop_ref(thread_state.profiler.pid, ncclop); + let kernelch = thread_state + .kernelch_free_list + .alloc_new( + KernelCh { + parent_op: Some(ncclop), + }, + None, + true, + ) + .unwrap(); + Some(event::Event::KernelCh(kernelch)) + }) + } else { + None + } } } } profiler_shim::ncclProfileProxyStep => { - // SAFETY: just checked that event type is proxystep - let descr = unsafe { descr.cast_to_proxystep() }; - let parent_ffi = descr.parent_obj(); - if parent_ffi.is_null() { + // Start-gate: refuse NEW ProxyStep events when steps are disabled. + // In-flight ProxySteps still receive state/stop via their existing handles. + let allow_steps = gates + .map(|g| g.proxy_step_enabled()) + .unwrap_or_else(|| config.track_steps || config.aggregate_steps); + if !allow_steps { None } else { - // SAFETY: When not set to null, - // NCCL always sets the parent_obj to a valid handle returned by profiler - // API. So `event::Event::from_ffi()` would always be called on a valid handle - let parent = unsafe { event::Event::from_ffi(parent_ffi) }; - if let Some(event::Event::ProxyOp(op)) = parent { - with_thread_state(|thread_state| { - let data = thread_state - .proxystep_free_list - .alloc_new( - event::ProxyStep::new( - descr.step(), - slab::AllocatedNode::into_raw(op), - ), - None, - true, - ) - .unwrap(); - Some(event::Event::ProxyStep(data)) - }) + // SAFETY: just checked that event type is proxystep + let descr = unsafe { descr.cast_to_proxystep() }; + let parent_ffi = descr.parent_obj(); + if parent_ffi.is_null() { + None } else { - let _ = parent.map(event::Event::into_ffi); + // SAFETY: When not set to null, + // NCCL always sets the parent_obj to a valid handle returned by profiler + // API. So `event::Event::from_ffi()` would always be called on a valid handle + let parent = unsafe { event::Event::from_ffi(parent_ffi) }; + if let Some(event::Event::ProxyOp(op)) = parent { + with_thread_state(|thread_state| { + let data = thread_state + .proxystep_free_list + .alloc_new( + event::ProxyStep::new( + descr.step(), + slab::AllocatedNode::into_raw(op), + ), + None, + true, + ) + .unwrap(); + Some(event::Event::ProxyStep(data)) + }) + } else { + let _ = parent.map(event::Event::into_ffi); + None + } + } + } + } + profiler_shim::ncclProfileKernelStep => { + let allow_ks = gates + .map(|g| g.track_kernel_step()) + .unwrap_or(config.track_kernel_step); + if !allow_ks || E::version() != Version::V6 { + None + } else { + let descr = unsafe { descr.cast_to_kernelstep() }; + let parent_ffi = descr.parent_obj(); + let parent_type = event_ffi::get_handle_type(parent_ffi); + if parent_ffi.is_null() || parent_type == event_ffi::Type::SmallNcclOp { None + } else { + let parent = event_ffi::ProxyParent::from_ffi(parent_ffi); + if let event_ffi::ProxyParent::NcclOp(ncclop) = parent { + with_thread_state(|thread_state| { + thread_state.inc_ncclop_ref(thread_state.profiler.pid, ncclop); + let step = thread_state + .kernelstep_free_list + .alloc_new( + KernelStepLocal { + parent_op: Some(ncclop), + channel_id: descr.channel_id(), + is_send: descr.is_send(), + peer: descr.peer(), + step: descr.step(), + size: descr.size(), + start_ts: descr.start_ts(), + ready_ts: descr.ready_ts(), + end_ts: None, + }, + None, + true, + ) + .unwrap(); + Some(event::Event::KernelStep(step)) + }) + } else { + None + } } } } - _ => panic!("unknown event type"), + _ => None, }; Ok(event) } @@ -769,7 +892,7 @@ pub fn stop_event_handler(event: event::Event) -> NcclResult<()> { if let Some(step) = data.step_tracker.finalize() { data.get_steps_mut(thread_state).push(step); } - if thread_state.profiler.config.track_proxyop { + if thread_state.profiler.gates.track_proxyop() { // if we are tracking proxyop also send the extra info about this proxyop let msg = daemon::Message::ProxyOpExtra(data.extra.clone()); thread_state.send_to_daemon(msg, false); @@ -821,6 +944,28 @@ pub fn stop_event_handler(event: event::Event) -> NcclResult<()> { } thread_state.kernelch_free_list.free(kernelch); } + event::Event::KernelStep(step) => { + thread_state.fifo.prefetch_next(); + if let Some(ncclop) = step.parent_op { + let end = step.end_ts.unwrap_or(step.ready_ts); + let msg = daemon::Message::KernelStep( + event::KernelEventStep { + channel_id: step.channel_id, + is_send: step.is_send, + peer: step.peer, + step: step.step, + size: step.size, + start_ts: step.start_ts, + ready_ts: step.ready_ts, + end_ts: end, + }, + ncclop, + ); + thread_state.send_to_daemon(msg, true); + let _ = thread_state.dec_ncclop_ref(thread_state.profiler.pid, ncclop); + } + thread_state.kernelstep_free_list.free(step); + } event::Event::Dummy(_) => (), event::Event::SmallNcclOp(_) => (), event::Event::NcclOpLite(_) => {} @@ -924,6 +1069,20 @@ pub fn record_proxyop_event_state_handler_v4( Ok(()) } +pub fn record_kernelstep_event_state_handler( + event: &mut event::Event, + e_state: profiler_shim::ncclProfilerEventState_v1_t, + p_timer: u64, +) -> NcclResult<()> { + if e_state != profiler_shim::proxy_event_state::v6::KERNEL_STEP_STOP { + return Ok(()); + } + if let event::Event::KernelStep(step) = event { + step.end_ts = Some(p_timer); + } + Ok(()) +} + #[allow(clippy::boxed_local)] pub fn finalize_handler(comm: Box) -> NcclResult<()> { if config::CONFIG.telemetry_mode > 0 { diff --git a/src/profiler_shim.rs b/src/profiler_shim.rs index 17b9fa2..c804d56 100644 --- a/src/profiler_shim.rs +++ b/src/profiler_shim.rs @@ -67,6 +67,17 @@ impl AsRef for EventDescrV4 { } } +#[repr(transparent)] +#[derive(Debug, Clone)] +pub struct EventDescrV6(pub ncclProfilerEventDescr_v6_t); + +impl AsRef for EventDescrV6 { + #[inline(always)] + fn as_ref(&self) -> &ncclProfilerEventDescr_v6_t { + &self.0 + } +} + pub type EventDescr = EventDescrV2; // alias of `ncclProfilerEventState_vX_t` to make the name shorter @@ -112,6 +123,11 @@ pub mod proxy_event_state { pub const RECV_FLUSH_WAIT: u32 = ncclProfilerEventState_t_ncclProfilerProxyStepRecvFlushWait; } + + pub mod v6 { + use super::*; + pub const KERNEL_STEP_STOP: u32 = ncclProfilerEventState_t_ncclProfilerKernelStepStop; + } } #[cfg(test)] diff --git a/src/runtime_gates.rs b/src/runtime_gates.rs new file mode 100644 index 0000000..90c1279 --- /dev/null +++ b/src/runtime_gates.rs @@ -0,0 +1,261 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Runtime start-gate for CoMMA logging metrics. +//! +//! Toggle bits without touching in-flight events: +//! - Update atomics that `start_event` consults (refuse new starts only). +//! - Rewrite NCCL's process-global `ncclProfilerEventMask` so *new* tasks +//! pick up enablement at enqueue. +//! +//! Live control is via Unix-socket RPC (`control_rpc`); see monitor `control` module. + +use crate::config::Config; +use crate::profiler::Version; +use crate::profiler_shim; + +use log::{info, warn}; + +use serde::{Deserialize, Serialize}; + +use std::sync::atomic::{AtomicBool, AtomicI32, AtomicPtr, Ordering}; + +/// Desired metric enablement + retained pointer to NCCL's activation mask. +#[derive(Debug)] +pub struct RuntimeGates { + track_group: AtomicBool, + track_ncclop: AtomicBool, + track_proxyop: AtomicBool, + track_interprocess_proxyop: AtomicBool, + track_steps: AtomicBool, + track_recv_steps: AtomicBool, + track_step_fifo_wait: AtomicBool, + aggregate_steps: AtomicBool, + track_kernel_ch: AtomicBool, + track_kernel_step: AtomicBool, + api_v6: AtomicBool, + /// Points at NCCL `ncclProfilerEventMask` (same address every `init`). + event_mask: AtomicPtr, +} + +impl RuntimeGates { + pub fn new(config: &Config, version: Version) -> Self { + Self { + track_group: AtomicBool::new(config.track_group), + track_ncclop: AtomicBool::new(config.track_ncclop), + track_proxyop: AtomicBool::new(config.track_proxyop), + track_interprocess_proxyop: AtomicBool::new(config.track_interprocess_proxyop), + track_steps: AtomicBool::new(config.track_steps), + track_recv_steps: AtomicBool::new(config.track_recv_steps), + track_step_fifo_wait: AtomicBool::new(config.track_step_fifo_wait), + aggregate_steps: AtomicBool::new(config.aggregate_steps), + track_kernel_ch: AtomicBool::new(config.track_kernel_ch), + track_kernel_step: AtomicBool::new(config.track_kernel_step), + api_v6: AtomicBool::new(matches!(version, Version::V6)), + event_mask: AtomicPtr::new(std::ptr::null_mut()), + } + } + + pub fn track_group(&self) -> bool { + self.track_group.load(Ordering::Acquire) + } + pub fn track_ncclop(&self) -> bool { + self.track_ncclop.load(Ordering::Acquire) + } + pub fn track_proxyop(&self) -> bool { + self.track_proxyop.load(Ordering::Acquire) + } + pub fn track_interprocess_proxyop(&self) -> bool { + self.track_interprocess_proxyop.load(Ordering::Acquire) + } + pub fn track_steps(&self) -> bool { + self.track_steps.load(Ordering::Acquire) + } + pub fn track_recv_steps(&self) -> bool { + self.track_recv_steps.load(Ordering::Acquire) + } + pub fn track_step_fifo_wait(&self) -> bool { + self.track_step_fifo_wait.load(Ordering::Acquire) + } + pub fn aggregate_steps(&self) -> bool { + self.aggregate_steps.load(Ordering::Acquire) + } + pub fn track_kernel_ch(&self) -> bool { + self.track_kernel_ch.load(Ordering::Acquire) + } + pub fn track_kernel_step(&self) -> bool { + self.track_kernel_step.load(Ordering::Acquire) + } + + pub fn proxy_step_enabled(&self) -> bool { + self.track_steps() || self.aggregate_steps() + } + + pub fn snapshot(&self) -> GateSnapshot { + GateSnapshot { + track_group: self.track_group(), + track_ncclop: self.track_ncclop(), + track_proxyop: self.track_proxyop(), + track_interprocess_proxyop: self.track_interprocess_proxyop(), + track_steps: self.track_steps(), + track_recv_steps: self.track_recv_steps(), + track_step_fifo_wait: self.track_step_fifo_wait(), + aggregate_steps: self.aggregate_steps(), + track_kernel_ch: self.track_kernel_ch(), + track_kernel_step: self.track_kernel_step(), + mask: self.compute_mask(), + } + } + + /// Remember NCCL's global mask pointer (idempotent). + pub fn register_mask_ptr(&self, ptr: *mut i32) { + if ptr.is_null() { + return; + } + let prev = self.event_mask.swap(ptr, Ordering::AcqRel); + if !prev.is_null() && prev != ptr { + warn!("CoMMA activation mask pointer changed unexpectedly"); + } + } + + /// Compute mask from current gates and publish to NCCL + return value for `init`. + pub fn publish_mask(&self) -> i32 { + let mask = self.compute_mask(); + self.store_mask(mask); + mask + } + + pub fn compute_mask(&self) -> i32 { + let mut mask: i32 = 0; + mask |= profiler_shim::ncclProfileGroup as i32; + mask |= profiler_shim::ncclProfileColl as i32; + mask |= profiler_shim::ncclProfileP2p as i32; + // ProxyOp bit stays on when proxyop tracking is enabled (MoE/P2P paths). + if self.track_proxyop() { + mask |= profiler_shim::ncclProfileProxyOp as i32; + } else { + // Keep ProxyOp bit for completion even if deep tracking is off — + // historical CoMMA always set this bit. Prefer always-on ProxyOp mask. + mask |= profiler_shim::ncclProfileProxyOp as i32; + } + if self.track_kernel_ch() { + mask |= profiler_shim::ncclProfileKernelCh as i32; + } + if self.track_kernel_step() && self.api_v6.load(Ordering::Acquire) { + mask |= profiler_shim::ncclProfileKernelStep as i32; + } + if self.proxy_step_enabled() { + mask |= profiler_shim::ncclProfileProxyStep as i32; + } + mask + } + + fn store_mask(&self, mask: i32) { + let ptr = self.event_mask.load(Ordering::Acquire); + if ptr.is_null() { + return; + } + // SAFETY: ptr is &ncclProfilerEventMask for the process lifetime. + unsafe { + (*(ptr as *const AtomicI32)).store(mask, Ordering::Release); + } + } + + /// Apply a parsed update; republish mask. Returns true if anything changed. + pub fn apply_update(&self, update: &GateUpdate) -> bool { + let mut changed = false; + changed |= swap_bool(&self.track_group, update.track_group); + changed |= swap_bool(&self.track_ncclop, update.track_ncclop); + changed |= swap_bool(&self.track_proxyop, update.track_proxyop); + changed |= swap_bool(&self.track_interprocess_proxyop, update.track_interprocess_proxyop); + changed |= swap_bool(&self.track_steps, update.track_steps); + changed |= swap_bool(&self.track_recv_steps, update.track_recv_steps); + changed |= swap_bool(&self.track_step_fifo_wait, update.track_step_fifo_wait); + changed |= swap_bool(&self.aggregate_steps, update.aggregate_steps); + changed |= swap_bool(&self.track_kernel_ch, update.track_kernel_ch); + changed |= swap_bool(&self.track_kernel_step, update.track_kernel_step); + if changed { + let mask = self.publish_mask(); + info!( + "CoMMA runtime gates updated: {:?} mask=0x{:x}", + self.snapshot(), + mask + ); + } + changed + } +} + +fn swap_bool(slot: &AtomicBool, next: Option) -> bool { + let Some(v) = next else { + return false; + }; + slot.swap(v, Ordering::AcqRel) != v +} + +#[derive(Debug, Default, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct GateUpdate { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub track_group: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub track_ncclop: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub track_proxyop: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub track_interprocess_proxyop: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub track_steps: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub track_recv_steps: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub track_step_fifo_wait: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub aggregate_steps: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub track_kernel_ch: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub track_kernel_step: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct GateSnapshot { + pub track_group: bool, + pub track_ncclop: bool, + pub track_proxyop: bool, + pub track_interprocess_proxyop: bool, + pub track_steps: bool, + pub track_recv_steps: bool, + pub track_step_fifo_wait: bool, + pub aggregate_steps: bool, + pub track_kernel_ch: bool, + pub track_kernel_step: bool, + pub mask: i32, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn gate_update_json_roundtrip() { + let u = GateUpdate { + track_kernel_step: Some(false), + track_steps: Some(true), + ..Default::default() + }; + let s = serde_json::to_string(&u).unwrap(); + let back: GateUpdate = serde_json::from_str(&s).unwrap(); + assert_eq!(back, u); + } +} diff --git a/third_party/nccl b/third_party/nccl index 5067397..52eb456 160000 --- a/third_party/nccl +++ b/third_party/nccl @@ -1 +1 @@ -Subproject commit 5067397c2676d5aed50042fc39e5c8ee96eb0027 +Subproject commit 52eb456805c2757f0ef0d8e5be2f83e7f75dd377