Kernelstep profiler

.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/

.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

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"

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`. |
 

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<std::path::Path>) -> Option<UnixStream> {
+    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<Telemetry>,
@@ -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();

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<String>,
     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<String>,
+    /// Unix stream socket path for live NDJSON latency telemetry.
+    pub latency_sock: Option<String>,
+    pub latency_flush_interval: Duration,
     pub summary_file: Option<String>,
     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));