From d82938a68ad6fbd68b0b7faf33b8ee88387e19e4 Mon Sep 17 00:00:00 2001 From: sputti-czi Date: Tue, 7 Jul 2026 16:47:49 -0400 Subject: [PATCH 1/2] feat: add gap duration --- Cargo.toml | 6 + benches/otel_utils.rs | 99 ++++++++ docs/otel-support.md | 48 +++- src/cloud_daemon.rs | 197 +++++++++++++-- src/lib.rs | 2 +- src/otel_utils.rs | 576 +++++++++++++++++++++++++++++++++++++++--- src/profiler.rs | 184 ++++++++------ src/step_tracker.rs | 37 ++- 8 files changed, 1019 insertions(+), 130 deletions(-) create mode 100644 benches/otel_utils.rs diff --git a/Cargo.toml b/Cargo.toml index fa83832..a4d2e8c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,6 +20,12 @@ path = "src/integration_benchmark.rs" name = "clock" harness = false +# copybara:strip_begin(otel) +[[bench]] +name = "otel_utils" +harness = false +# copybara:strip_end + [features] default = [] explicit-optin = [] diff --git a/benches/otel_utils.rs b/benches/otel_utils.rs new file mode 100644 index 0000000..945305b --- /dev/null +++ b/benches/otel_utils.rs @@ -0,0 +1,99 @@ +// 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. + +// copybara:strip_begin(otel) +use criterion::{criterion_group, criterion_main, Criterion}; +use nccl_profiler::otel_utils::{ + DurationHistogram, EventStep, GapTracker, HistogramManager, NcclOpKey, +}; +use opentelemetry_sdk::metrics::{ + new_view, Aggregation, Instrument, InstrumentKind, ManualReader, SdkMeterProvider, Stream, +}; + +// install a real meter provider with the same base2 exponential histogram +// aggregation as init_meter_provider() so record costs are representative +fn init_meter_provider() { + let mut builder = SdkMeterProvider::builder().with_reader(ManualReader::builder().build()); + for name in [ + "*latency", + "nccl.collective.duration", + "nccl.collective.gap", + ] { + let mut histogram_instrument = Instrument::new().name(name); + histogram_instrument.kind = Some(InstrumentKind::Histogram); + let mask = Stream::new().aggregation(Aggregation::Base2ExponentialHistogram { + max_size: 160, + max_scale: 20, + record_min_max: true, + }); + if let Ok(view) = new_view(histogram_instrument, mask) { + builder = builder.with_view(view); + } + } + opentelemetry::global::set_meter_provider(builder.build()); +} + +fn criterion_benchmark(c: &mut Criterion) { + init_meter_provider(); + + let step = EventStep { + step: 0, + size: 65536, + start_time: 1234567, + fifo_wait_dur_ns: None, + dur_ns: 512, + }; + + let mut send_manager = HistogramManager::new("nccl.net_send.latency", "ns", 16); + let send_histogram = send_manager.get_histogram(NcclOpKey::NetSend(0x123, 0, 1)); + c.bench_function("net_send latency record", |b| { + b.iter(|| send_histogram.record(&step)) + }); + + let mut recv_manager = HistogramManager::new("nccl.net_recv.latency", "ns", 16); + let recv_histogram = recv_manager.get_histogram(NcclOpKey::NetRecv(0x123, 0, 1)); + c.bench_function("net_recv latency record", |b| { + b.iter(|| recv_histogram.record(&step)) + }); + + let mut duration_histogram = DurationHistogram::new("nccl.collective.duration", "ns", 16); + c.bench_function("collective duration record", |b| { + b.iter(|| duration_histogram.record(4096, 0x123, "ncclAllReduce", 0, 1 << 22)) + }); + + // the clock closure mirrors the shipped default path of + // Profiler::recent_timer_ns (an Instant read per call) + let gap_tracker = GapTracker::new(); + gap_tracker.init_instrument(); + let t0 = std::time::Instant::now(); + let now_ns = move || t0.elapsed().as_nanos() as u64; + c.bench_function("gap idle transition", |b| { + b.iter(|| { + gap_tracker.activity_begin(now_ns); + gap_tracker.activity_end(now_ns()); + }) + }); + + gap_tracker.activity_begin(now_ns); + c.bench_function("gap nested activity", |b| { + b.iter(|| { + gap_tracker.activity_begin(now_ns); + gap_tracker.activity_end(now_ns()); + }) + }); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); +// copybara:strip_end diff --git a/docs/otel-support.md b/docs/otel-support.md index d8adc7a..86df1fc 100644 --- a/docs/otel-support.md +++ b/docs/otel-support.md @@ -46,7 +46,53 @@ To prevent high cardinality issues, CoMMA uses a "Top K" cardinality management - `nccl.metric.aggregated`: Set to `true`. - `nccl.hostname`: Hostname of the node. -The "Top K" list is dynamically updated at the interval defined by `NCCL_PROFILER_OTEL_METRICS_CARDINALITY_GROUPING_INTERVAL`. +The "Top K" list is dynamically updated at the interval defined by `NCCL_PROFILER_OTEL_METRICS_CARDINALITY_GROUPING_INTERVAL`, and keys of closed communicators are evicted immediately. + +Note that the OTel SDK exports with cumulative temporality and keeps a stream for every attribute set it has ever recorded until process exit — including streams for keys that were later demoted from the "Top K" set. CoMMA's cardinality limit bounds how many keys record with high-fidelity attributes at any time, but under communicator churn the SDK-side stream count grows with the total number of distinct attribute sets ever promoted. Keep `NCCL_PROFILER_OTEL_METRICS_MAX_CARDINALITY` modest when communicators churn frequently. This applies to every "Top K" managed metric below. + +### `nccl.net_recv.latency` (Histogram, Unit: `ns`) + +This metric records the latency of network receive operations. Receive steps are only generated when `NCCL_PROFILER_TRACK_RECV_STEPS` is enabled (subject to `NCCL_PROFILER_P2P_RECV_SAMPLE_RATE` for point-to-point operations). + +It uses the same "Top K" cardinality management strategy and attributes as `nccl.net_send.latency`. For receive connections, `nccl.source.rank` is the remote sender and `nccl.destination.rank` is the local rank. + +A receive step spans buffer post to data arrival, so its duration includes any delay before the remote sender was ready — unlike send steps, which measure post-clearance transfer time. The two metrics are therefore not symmetric; receive latencies are most meaningful compared across edges rather than against send latencies. + +### `nccl.collective.duration` (Histogram, Unit: `ns`) + +This metric records, for each completed NCCL operation, the span of its network / kernel activity: from the start of its first proxy op (or kernel channel, when `NCCL_PROFILER_TRACK_KERNEL_CH` is enabled) to the end of its last one. This is the same operation lifetime CoMMA computes for traces and summaries. Four classes of operations are not recorded: + +- Operations that produce no proxy op or kernel channel activity, e.g. single-node NVLink-only collectives with kernel channel tracking disabled. +- Point-to-point operations not selected by sampling (`NCCL_PROFILER_P2P_SAMPLE_RATE`, `NCCL_PROFILER_P2P_RECV_SAMPLE_RATE`). With the default recv sample rate of `0.1`, `ncclRecv` durations are a 10% sample: bucket shapes are unbiased, but counts and sums under-report by 10x, and asymmetrically versus `ncclSend`. +- Operations reclaimed by the hang timeout (`NCCL_PROFILER_NCCLOP_TIMEOUT`) before completing. +- Operations on the small-message fast paths: point-to-point operations at or below `NCCL_PROFILER_SMALL_MSG_THRESHOLD` (default 64 KiB) are never tracked, and with the default `NCCL_PROFILER_SKIP_SMALL_COLLECTIVE=true` collectives other than AllReduce at or below the threshold are not tracked either. The `lt1m` size class therefore only covers operations outside these fast paths. + +Attributes: +- `nccl.comm.hash`: Hexadecimal string identifying the NCCL communicator. +- `nccl.collective.name`: Name of the operation (e.g., `ncclAllReduce`, `ncclSend`). +- `nccl.rank`: Rank of the process within the communicator. +- `nccl.size.class`: Coarse operation size class: `lt1m` (< 1 MiB), `1m_16m` (1-16 MiB), or `gt16m` (> 16 MiB). This preserves the latency-bound vs bandwidth-bound split with bounded cardinality. +- `nccl.hostname`: Hostname of the node. + +Distinct attribute sets are capped by `NCCL_PROFILER_OTEL_METRICS_MAX_CARDINALITY` with the same "Top K" strategy as the latency metrics: keys are re-ranked by activity at every `NCCL_PROFILER_OTEL_METRICS_CARDINALITY_GROUPING_INTERVAL`, keys of closed communicators are evicted, and operations outside the top K are recorded with `nccl.metric.aggregated` set to `true`. + +### `nccl.collective.gap` (Histogram, Unit: `ns`) + +This metric records the duration of intervals where the process has no NCCL activity in flight. Activity is the union of operation enqueue windows, proxy op windows (network transfers), and kernel channel windows (when `NCCL_PROFILER_TRACK_KERNEL_CH` is enabled), so an operation keeps the process busy from its API call until its network / kernel work actually completes, not merely until it is enqueued. Each idle interval is recorded once, when the activity that ends it starts. Overlapping activity never contributes to a gap. + +Attributes: +- `nccl.hostname`: Hostname of the node. +- `nccl.pid`: Process id. Together with the hostname this identifies the process; a comm-local rank would be ambiguous for a metric that spans communicators. + +The gap is a per-process property by definition, so it deliberately carries no communicator dimension. + +Caveats — what counts as activity depends on the tracking configuration: +- The metric assumes the default `NCCL_PROFILER_TRACK_NCCLOP=true`. With operation tracking disabled, only proxy op windows (and kernel channel windows, when enabled) count as activity, and everything else is reported as gap. +- Operations whose execution produces no tracked activity window (see the duration carve-outs above) are only counted while they are enqueued, so their execution time appears as gap. Enable kernel channel tracking for full coverage of NVLink-only collectives. +- Operations on the small-message fast paths (see the duration carve-outs above) are covered only while enqueued — small point-to-point operations not at all — so their transfer time, possibly tens of microseconds each, is reported as gap. Interpret gap totals with care for workloads dominated by messages below `NCCL_PROFILER_SMALL_MSG_THRESHOLD`. +- Point-to-point operations not selected by sampling (`NCCL_PROFILER_P2P_SAMPLE_RATE`, `NCCL_PROFILER_P2P_RECV_SAMPLE_RATE`) produce no activity window at all, so their entire enqueue and transfer time is reported as gap. With the default recv sample rate of `0.1`, 90% of receive transfers count as idle time; set the sample rates to `1.0` before interpreting gap totals for p2p-heavy workloads such as pipeline parallelism. +- Sub-microsecond gaps can appear between an operation's enqueue window and the start of its proxy activity; they land in the lowest buckets and carry negligible weight in gap-time totals. +- If NCCL aborts a plan launch on an error path it may never stop the events it started; the in-flight count then stays above zero and the metric reports no further gaps for the process lifetime. CoMMA logs a one-time warning when it detects a stuck in-flight count. This only occurs after NCCL errors, which the job surfaces on its own. ### `nccl.collective.seq_num` (Gauge) diff --git a/src/cloud_daemon.rs b/src/cloud_daemon.rs index 14f3320..d2eeddc 100644 --- a/src/cloud_daemon.rs +++ b/src/cloud_daemon.rs @@ -337,7 +337,8 @@ async fn build_bufwriter( async fn exporter( profiler: &'static Profiler, mut rx: mpsc::Receiver, - otel_latency_hist_manager: Option>>, + otel_send_latency_hist_manager: Option>>, + otel_recv_latency_hist_manager: Option>>, ) -> std::io::Result<()> { let mut latency_file = if let Some(template) = profiler.config.latency_file.as_ref() { let path = template.replace("%p", &format!("{}", profiler.pid)); @@ -378,6 +379,17 @@ async fn exporter( None }; + let mut otel_duration_histogram: Option = + if profiler.config.otel_enable { + Some(otel_utils::DurationHistogram::new( + "nccl.collective.duration", + "ns", + profiler.config.otel_metrics_max_cardinality, + )) + } else { + None + }; + let mut summary_interval = tokio::time::interval(profiler.config.summary_interval); // the very first tick completes immediately @@ -449,6 +461,32 @@ async fn exporter( let _ = otel_utils::record_ncclop_seqnum(gauge, profiler, op); } } + + if let Some(histogram) = otel_duration_histogram.as_mut() { + match &telemetry { + Telemetry::NcclOp(op) => { + let _ = otel_utils::record_ncclop_duration(histogram, profiler, op); + } + Telemetry::CommClose(comm_hash) => { + histogram.close_comm(*comm_hash); + } + _ => {} + } + } + + if let Telemetry::CommClose(comm_hash) = &telemetry { + for manager in [ + otel_send_latency_hist_manager.as_ref(), + otel_recv_latency_hist_manager.as_ref(), + ] + .into_iter() + .flatten() + { + if let Ok(mut inner) = manager.lock() { + inner.close_comm(*comm_hash); + } + } + } }, None => break, } @@ -470,11 +508,23 @@ async fn exporter( } }, _ = async { otel_metrics_grouping_interval.as_mut().unwrap().tick().await }, if otel_metrics_grouping_interval.is_some() => { - if let Some(manager) = otel_latency_hist_manager.as_ref() { + for manager in [ + otel_send_latency_hist_manager.as_ref(), + otel_recv_latency_hist_manager.as_ref(), + ] + .into_iter() + .flatten() + { if let Ok(mut inner) = manager.lock() { inner.update_priority(); } } + if let Some(histogram) = otel_duration_histogram.as_mut() { + histogram.update_priority(); + } + if let Some(gap_tracker) = profiler.gap_tracker.as_ref() { + gap_tracker.check_stalled(); + } }, } } @@ -500,7 +550,8 @@ async fn exporter( struct Exporter { tx: mpsc::Sender, - otel_latency_hist_manager: Option>>, + otel_send_latency_hist_manager: Option>>, + otel_recv_latency_hist_manager: Option>>, gpuviz: Option>>>, track_step_fifo_wait: bool, } @@ -510,7 +561,8 @@ impl Exporter { fn new(tx: mpsc::Sender) -> Self { Self { tx, - otel_latency_hist_manager: None, + otel_send_latency_hist_manager: None, + otel_recv_latency_hist_manager: None, gpuviz: None, track_step_fifo_wait: false, } @@ -548,7 +600,12 @@ impl Export for Exporter { key: &NcclOpKey, ) -> Option>>> { let mut histograms: Vec>> = Vec::new(); - if let Some(manager) = self.otel_latency_hist_manager.as_ref() { + let otel_manager = match key { + NcclOpKey::NetSend(..) => self.otel_send_latency_hist_manager.as_ref(), + NcclOpKey::NetRecv(..) => self.otel_recv_latency_hist_manager.as_ref(), + _ => None, + }; + if let Some(manager) = otel_manager { if let Ok(mut lg) = manager.lock() { let h = Arc::new(lg.get_histogram(key.clone())); histograms.push(h as _); @@ -590,23 +647,35 @@ async fn main_loop( if otel_utils::init_tracer_provider(&profiler.config).is_none() { log::warn!("failed to init otel tracer provider"); } + if let Some(gap_tracker) = profiler.gap_tracker.as_ref() { + gap_tracker.init_instrument(); + } } const TELEMETRY_CHANNEL_SZ: usize = 4096; let (tx, rx) = mpsc::channel::(TELEMETRY_CHANNEL_SZ); - let otel_latency_hist_manager = if profiler.config.otel_enable { - Some(Arc::new(Mutex::new(otel_utils::HistogramManager::new( - "nccl.net_send.latency", - "ns", - profiler.config.otel_metrics_max_cardinality, - )))) - } else { - None + let new_latency_hist_manager = |name| { + if profiler.config.otel_enable { + Some(Arc::new(Mutex::new(otel_utils::HistogramManager::new( + name, + "ns", + profiler.config.otel_metrics_max_cardinality, + )))) + } else { + None + } }; - let export_worker = - tokio::task::spawn(exporter(profiler, rx, otel_latency_hist_manager.clone())); + let otel_send_latency_hist_manager = new_latency_hist_manager("nccl.net_send.latency"); + let otel_recv_latency_hist_manager = new_latency_hist_manager("nccl.net_recv.latency"); + let export_worker = tokio::task::spawn(exporter( + profiler, + rx, + otel_send_latency_hist_manager.clone(), + otel_recv_latency_hist_manager.clone(), + )); let exporter = Exporter { tx, - otel_latency_hist_manager, + otel_send_latency_hist_manager, + otel_recv_latency_hist_manager, gpuviz: profiler .gpuviz_lib .as_ref() @@ -661,6 +730,102 @@ mod tests { // create a mutex to avoid multiple test cases allocating ncclops concurrently static NCCLOP_TEST_MUTEX: Mutex<()> = Mutex::new(()); + #[test] + fn otel_latency_histogram_routing() { + const MAX_CARDINALITY: usize = 4; + let (tx, _rx) = mpsc::channel::(4); + let send_manager = Arc::new(Mutex::new(otel_utils::HistogramManager::new( + "nccl.net_send.latency", + "ns", + MAX_CARDINALITY, + ))); + let recv_manager = Arc::new(Mutex::new(otel_utils::HistogramManager::new( + "nccl.net_recv.latency", + "ns", + MAX_CARDINALITY, + ))); + let exporter = Exporter { + tx, + otel_send_latency_hist_manager: Some(send_manager.clone()), + otel_recv_latency_hist_manager: Some(recv_manager.clone()), + gpuviz: None, + track_step_fifo_wait: false, + }; + + assert!(exporter + .get_latency_histogram(&NcclOpKey::NetSend(0x123, 0, 1)) + .is_some()); + assert_eq!(send_manager.lock().unwrap().num_active(), 1); + assert_eq!(recv_manager.lock().unwrap().num_active(), 0); + + assert!(exporter + .get_latency_histogram(&NcclOpKey::NetRecv(0x123, 0, 1)) + .is_some()); + assert_eq!(send_manager.lock().unwrap().num_active(), 1); + assert_eq!(recv_manager.lock().unwrap().num_active(), 1); + } + + #[test] + fn otel_gap_spans_op_execution() { + use crate::event_ffi::AsFFI as _; + use crate::profiler::THREAD_STATE; + + let _lg = NCCLOP_TEST_MUTEX.lock().unwrap(); + + let mut profiler = Profiler::new(Version::V2); + profiler.pid = 42; // must match the pid in the mock proxyop descriptor + profiler.config.track_interprocess_proxyop = false; + profiler.config.otel_enable = true; + profiler.gap_tracker = Some(otel_utils::GapTracker::new()); + + let profiler_ptr = Box::into_raw(Box::new(profiler)); + // SAFETY: the box is only reclaimed after the daemon is joined and + // the thread state referencing it is dropped + let profiler: &'static Profiler = unsafe { &*profiler_ptr }; + profiler.spawn_daemon(); + THREAD_STATE.with_borrow_mut(|state| *state = Some(profiler.init_thread_state())); + + { + let gap_tracker = profiler.gap_tracker.as_ref().unwrap(); + let comm = Communicator::new(); + + let mut coll_descr = profiler_shim::tests::dummy_coll_descr(); + // large enough to bypass the small-collective fast paths + coll_descr.0.__bindgen_anon_1.coll.count = 1 << 20; + let op = crate::profiler::start_event_handler(&coll_descr, &comm) + .unwrap() + .unwrap(); + assert!(std::matches!(op, event::Event::NcclOp(_))); + assert_eq!(gap_tracker.in_flight(), 1); + let op_handle = op.into_ffi(); + + // the op's network work starts while it is being enqueued + let mut proxyop_descr = profiler_shim::tests::dummy_proxyop_descr(); + proxyop_descr.0.parentObj = op_handle; + let proxyop = crate::profiler::start_event_handler(&proxyop_descr, &comm) + .unwrap() + .unwrap(); + assert!(std::matches!(proxyop, event::Event::ProxyOp(_))); + assert_eq!(gap_tracker.in_flight(), 2); + + // NCCL stops the op event once it is enqueued, but the op stays + // in flight until its proxy op completes + // SAFETY: op_handle is a valid handle returned by into_ffi() + let op = unsafe { event::Event::from_ffi(op_handle) }.unwrap(); + crate::profiler::stop_event_handler(op).unwrap(); + assert_eq!(gap_tracker.in_flight(), 1); + + crate::profiler::stop_event_handler(proxyop).unwrap(); + assert_eq!(gap_tracker.in_flight(), 0); + } + + profiler.join_daemon(); + THREAD_STATE.with_borrow_mut(|state| *state = None); + // SAFETY: reclaim the profiler leaked above; the daemon and thread + // state that referenced it are gone + let _ = unsafe { Box::from_raw(profiler_ptr) }; + } + #[test] fn e2e_mock() { let _lg = NCCLOP_TEST_MUTEX.lock().unwrap(); diff --git a/src/lib.rs b/src/lib.rs index 5e2c600..fe260ef 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -23,7 +23,7 @@ mod gcp_acs_proto; // copybara:strip(oss_protobuf) mod gpuviz; mod histogram; mod nccl_metadata; -mod otel_utils; // copybara:strip(otel) +pub mod otel_utils; // copybara:strip(otel) mod profiler; pub mod profiler_shim; mod shm_fifo; diff --git a/src/otel_utils.rs b/src/otel_utils.rs index f7736b1..ff061d1 100644 --- a/src/otel_utils.rs +++ b/src/otel_utils.rs @@ -17,9 +17,10 @@ use crate::daemon::AtomicHistogram; use crate::event; use crate::event::ProfilerEvent as _; use crate::nccl_metadata; -use crate::nccl_metadata::NcclOpKey; use crate::profiler::Profiler; -use crate::step_tracker::EventStep; + +pub use crate::nccl_metadata::NcclOpKey; +pub use crate::step_tracker::EventStep; use opentelemetry::context::Context as OtelContext; use opentelemetry::metrics::Gauge; @@ -32,7 +33,7 @@ use opentelemetry_sdk::metrics::{new_view, Aggregation, Instrument, InstrumentKi use opentelemetry_sdk::Resource; use std::collections::HashMap; -use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; use std::sync::{Arc, LazyLock, OnceLock}; use std::time::{Duration, SystemTime}; @@ -49,19 +50,24 @@ pub fn init_meter_provider(config: &config::Config) -> Option<()> { .with_tonic() .build() .ok()?; - let mut histogram_instrument = Instrument::new().name("*latency"); - histogram_instrument.kind = Some(InstrumentKind::Histogram); - let mask = Stream::new().aggregation(Aggregation::Base2ExponentialHistogram { - max_size: config.otel_latency_histogram_max_size, - max_scale: config.otel_latency_histogram_max_scale as _, - record_min_max: true, - }); - let mut meter_provider_builder = opentelemetry_sdk::metrics::SdkMeterProvider::builder() .with_resource(resource.clone()) .with_periodic_exporter(otlp_exporter); - if let Ok(view) = new_view(histogram_instrument, mask) { - meter_provider_builder = meter_provider_builder.with_view(view); + for name in [ + "*latency", + "nccl.collective.duration", + "nccl.collective.gap", + ] { + let mut histogram_instrument = Instrument::new().name(name); + histogram_instrument.kind = Some(InstrumentKind::Histogram); + let mask = Stream::new().aggregation(Aggregation::Base2ExponentialHistogram { + max_size: config.otel_latency_histogram_max_size, + max_scale: config.otel_latency_histogram_max_scale as _, + record_min_max: true, + }); + if let Ok(view) = new_view(histogram_instrument, mask) { + meter_provider_builder = meter_provider_builder.with_view(view); + } } let meter_provider = METER_PROVIDER.get_or_init(|| meter_provider_builder.build()); global::set_meter_provider(meter_provider.clone()); @@ -89,9 +95,7 @@ pub fn init_tracer_provider(_config: &config::Config) -> Option<()> { #[derive(Debug)] pub struct LatencyHistogram { inner: OtelHistogram, - op_key: NcclOpKey, - hostname: String, - high_fidelity: bool, + attributes: Vec, local_counter: AtomicUsize, counter: Arc, } @@ -104,41 +108,61 @@ impl LatencyHistogram { high_fidelity: bool, counter: Arc, ) -> Self { + // attribute sets are built once per connection so the per-step + // record path performs no allocation + let attributes = if !high_fidelity { + vec![ + KeyValue::new("nccl.metric.aggregated", true), + KeyValue::new("nccl.hostname", hostname), + ] + } else { + match &op_key { + NcclOpKey::NetSend(comm_hash, src, dst) => vec![ + KeyValue::new("nccl.communicator.hash", format!("0x{:016x}", comm_hash)), + KeyValue::new("nccl.source.rank", *src as i64), + KeyValue::new("nccl.destination.rank", *dst as i64), + KeyValue::new("nccl.hostname", hostname), + ], + NcclOpKey::NetRecv(comm_hash, local, peer) => vec![ + KeyValue::new("nccl.communicator.hash", format!("0x{:016x}", comm_hash)), + KeyValue::new("nccl.source.rank", *peer as i64), + KeyValue::new("nccl.destination.rank", *local as i64), + KeyValue::new("nccl.hostname", hostname), + ], + _ => Vec::new(), + } + }; Self { inner, - op_key, - hostname, - high_fidelity, + attributes, local_counter: AtomicUsize::new(0), counter, } } + pub fn record(&self, step: &EventStep) { + if !self.attributes.is_empty() { + self.inner.record(step.dur_ns as _, &self.attributes); + } + self.local_counter.fetch_add(1, Ordering::Relaxed); + } + #[cfg(test)] fn high_fidelity(&self) -> bool { - self.high_fidelity + !self + .attributes + .contains(&KeyValue::new("nccl.metric.aggregated", true)) + } + + #[cfg(test)] + fn attributes(&self) -> &[KeyValue] { + &self.attributes } } impl AtomicHistogram for LatencyHistogram { fn record(&self, step: &EventStep) { - if let NcclOpKey::NetSend(comm_hash, src, dst) = &self.op_key { - let attributes: &[_] = if self.high_fidelity { - &[ - KeyValue::new("nccl.communicator.hash", format!("0x{:016x}", comm_hash)), - KeyValue::new("nccl.source.rank", *src as i64), - KeyValue::new("nccl.destination.rank", *dst as i64), - KeyValue::new("nccl.hostname", self.hostname.clone()), - ] - } else { - &[ - KeyValue::new("nccl.metric.aggregated", true), - KeyValue::new("nccl.hostname", self.hostname.clone()), - ] - }; - self.inner.record(step.dur_ns as _, attributes); - } - self.local_counter.fetch_add(1, Ordering::Relaxed); + LatencyHistogram::record(self, step); } } @@ -230,14 +254,324 @@ impl HistogramManager { self.num_active = 0; for (i, e) in info_vec.iter_mut().enumerate() { e.1.high_fidelity = i < self.cardinality_limit; - self.num_active += 1 + if e.1.high_fidelity { + self.num_active += 1 + } + } + } + + pub fn close_comm(&mut self, comm_hash: u64) { + let removed_active = self + .info + .iter() + .filter(|(key, info)| key.get_comm_hash() == comm_hash && info.high_fidelity) + .count(); + self.info.retain(|key, _| key.get_comm_hash() != comm_hash); + self.num_active -= removed_active; + } + + #[cfg(test)] + pub(crate) fn num_active(&self) -> usize { + self.num_active + } +} + +const SIZE_CLASS_SMALL_MAX: usize = 1 << 20; +const SIZE_CLASS_MEDIUM_MAX: usize = 16 << 20; + +fn size_class(byte_count: usize) -> &'static str { + if byte_count < SIZE_CLASS_SMALL_MAX { + "lt1m" + } else if byte_count <= SIZE_CLASS_MEDIUM_MAX { + "1m_16m" + } else { + "gt16m" + } +} + +type DurationKey = ( + /* comm_hash = */ u64, + /* op name = */ &'static str, + /* size class = */ &'static str, + /* rank = */ usize, +); + +#[derive(Debug)] +struct DurationInfo { + attributes: Option>, // Some iff this key is within top-k + count: usize, +} + +#[derive(Debug)] +pub struct DurationHistogram { + inner: OtelHistogram, + hostname: String, + info: HashMap, + aggregated_attributes: Vec, + num_active: usize, + cardinality_limit: usize, +} + +impl DurationHistogram { + pub fn new(name: &str, unit: &str, cardinality_limit: usize) -> Self { + let meter = opentelemetry::global::meter("CoMMA"); + let hostname = get_hostname_libc().unwrap_or_default(); + Self { + inner: meter + .u64_histogram(String::from(name)) + .with_unit(String::from(unit)) + .build(), + aggregated_attributes: vec![ + KeyValue::new("nccl.metric.aggregated", true), + KeyValue::new("nccl.hostname", hostname.clone()), + ], + hostname, + info: HashMap::new(), + num_active: 0, + cardinality_limit, + } + } + + fn build_attributes(hostname: &str, key: &DurationKey) -> Vec { + vec![ + KeyValue::new("nccl.comm.hash", format!("0x{:016x}", key.0)), + KeyValue::new("nccl.collective.name", key.1), + KeyValue::new("nccl.size.class", key.2), + KeyValue::new("nccl.rank", key.3 as i64), + KeyValue::new("nccl.hostname", String::from(hostname)), + ] + } + + pub fn record( + &mut self, + dur_ns: u64, + comm_hash: u64, + op_name: &'static str, + rank: usize, + byte_count: usize, + ) { + let key = (comm_hash, op_name, size_class(byte_count), rank); + let info = match self.info.entry(key) { + std::collections::hash_map::Entry::Occupied(e) => e.into_mut(), + std::collections::hash_map::Entry::Vacant(e) => { + let attributes = (self.num_active < self.cardinality_limit) + .then(|| Self::build_attributes(&self.hostname, e.key())); + if attributes.is_some() { + self.num_active += 1; + } + e.insert(DurationInfo { + attributes, + count: 0, + }) + } + }; + info.count += 1; + match info.attributes.as_ref() { + Some(attributes) => self.inner.record(dur_ns, attributes), + None => self.inner.record(dur_ns, &self.aggregated_attributes), + } + } + + // sort the keys by number of records to get the new "Top K"; counts are + // per-interval so keys of idle communicators age out of the top-k + pub fn update_priority(&mut self) { + let hostname = &self.hostname; + let mut info_vec: Vec<_> = self.info.iter_mut().collect(); + info_vec.sort_by_key(|e| std::cmp::Reverse(e.1.count)); + self.num_active = 0; + for (i, (key, info)) in info_vec.into_iter().enumerate() { + if i < self.cardinality_limit { + if info.attributes.is_none() { + info.attributes = Some(Self::build_attributes(hostname, key)); + } + self.num_active += 1; + } else { + info.attributes = None; + } + info.count = 0; } } + pub fn close_comm(&mut self, comm_hash: u64) { + let removed_active = self + .info + .iter() + .filter(|(key, info)| key.0 == comm_hash && info.attributes.is_some()) + .count(); + self.info.retain(|key, _| key.0 != comm_hash); + self.num_active -= removed_active; + } + #[cfg(test)] fn num_active(&self) -> usize { self.num_active } + + #[cfg(test)] + fn high_fidelity( + &self, + comm_hash: u64, + op_name: &'static str, + rank: usize, + byte_count: usize, + ) -> bool { + self.info + .get(&(comm_hash, op_name, size_class(byte_count), rank)) + .is_some_and(|info| info.attributes.is_some()) + } +} + +pub fn record_ncclop_duration( + histogram: &mut DurationHistogram, + _profiler: &Profiler, + op: &event::NcclOp, +) -> Option<()> { + let duration = op.child_duration()?; + let descr = op.get_descr(); + let name = if let Some(coll) = descr.try_cast_to_coll() { + ncclop_otel_name(coll.op_type()) + } else if let Some(p2p) = descr.try_cast_to_p2p() { + if p2p.is_send() { + "ncclSend" + } else { + "ncclRecv" + } + } else { + return None; + }; + histogram.record( + duration.as_nanos() as _, + op.comm_hash(), + name, + op.basic_info().rank(), + op.byte_count(), + ); + Some(()) +} + +const GAP_IDLE_NEVER: u64 = 0; +const GAP_IN_FLIGHT_MASK: u64 = (1 << 32) - 1; +const GAP_BEGIN_UNIT: u64 = 1 << 32; + +/// Tracks intervals where this process has no NCCL activity in flight. +/// +/// An activity window is an op enqueue (`NcclOp` start to stop), a proxy op +/// or a kernel channel; the union of these windows spans each operation from +/// its API call to the end of its network / kernel work, so the recorded gaps +/// are the complement of NCCL engagement and interval-union-correct under +/// overlap by construction. +/// +/// The in-flight count and idle timestamp are lock-free atomics, so activity +/// begin / end cost two atomic operations each on the caller's thread. +#[derive(Debug)] +pub struct GapTracker { + // low 32 bits: activity windows in flight; high 32 bits: total begins, + // so check_stalled can tell a leaked window apart from a busy process + state: AtomicU64, + idle_since_ns: AtomicU64, + last_state: AtomicU64, + stall_logged: AtomicBool, + histogram: OnceLock>, + attributes: [KeyValue; 2], +} + +impl GapTracker { + pub fn new() -> Self { + Self { + state: AtomicU64::new(0), + idle_since_ns: AtomicU64::new(GAP_IDLE_NEVER), + last_state: AtomicU64::new(0), + stall_logged: AtomicBool::new(false), + histogram: OnceLock::new(), + // the gap is process-wide (activity windows span communicators + // and threads), so it is labeled with a process-stable identity + // rather than a comm-local rank + attributes: [ + KeyValue::new("nccl.hostname", get_hostname_libc().unwrap_or_default()), + // SAFETY: `getpid()` takes no input and does not modify rust-managed state + KeyValue::new("nccl.pid", unsafe { libc::getpid() } as i64), + ], + } + } + + /// Must be called after the meter provider is installed; before that, + /// activity_begin / activity_end only maintain the in-flight count. + pub fn init_instrument(&self) { + let meter = opentelemetry::global::meter("CoMMA"); + let _ = self.histogram.get_or_init(|| { + meter + .u64_histogram("nccl.collective.gap") + .with_unit("ns") + .build() + }); + } + + pub fn activity_begin(&self, now_ns: T) + where + T: FnOnce() -> u64, + { + let Some(gap_ns) = self.begin_transition(now_ns) else { + return; + }; + let Some(histogram) = self.histogram.get() else { + return; + }; + histogram.record(gap_ns, &self.attributes); + } + + pub fn activity_end(&self, now_ns: u64) { + // stamp before decrementing so a concurrent 0 -> 1 observer never + // reads a timestamp from a previous idle period; fetch_max keeps the + // idle start at the latest end under concurrent activity ends + self.idle_since_ns.fetch_max(now_ns, Ordering::AcqRel); + self.state.fetch_sub(1, Ordering::AcqRel); + } + + fn begin_transition(&self, now_ns: T) -> Option + where + T: FnOnce() -> u64, + { + let prev = self.state.fetch_add(GAP_BEGIN_UNIT | 1, Ordering::AcqRel); + if prev & GAP_IN_FLIGHT_MASK != 0 { + return None; + } + let idle_since = self.idle_since_ns.load(Ordering::Acquire); + if idle_since == GAP_IDLE_NEVER { + return None; + } + Some(now_ns().saturating_sub(idle_since)) + } + + /// Called periodically off the hot path. If NCCL abandons a started + /// event on an error path the in-flight count never returns to zero and + /// the metric goes silent; warn once when windows stay in flight with no + /// begin or end for a whole check interval. + pub fn check_stalled(&self) -> bool { + let state = self.state.load(Ordering::Acquire); + let prev = self.last_state.swap(state, Ordering::Relaxed); + let stalled = state & GAP_IN_FLIGHT_MASK != 0 + && state == prev + && !self.stall_logged.swap(true, Ordering::Relaxed); + if stalled { + log::warn!( + "{} NCCL activity window(s) stuck in flight; \ + nccl.collective.gap will report no further gaps", + state & GAP_IN_FLIGHT_MASK + ); + } + stalled + } + + #[cfg(test)] + pub(crate) fn in_flight(&self) -> usize { + (self.state.load(Ordering::Acquire) & GAP_IN_FLIGHT_MASK) as usize + } +} + +impl Default for GapTracker { + fn default() -> Self { + Self::new() + } } fn coll_type_to_num(t: nccl_metadata::NcclOpType) -> u32 { @@ -459,6 +793,7 @@ mod tests { std::mem::drop(histograms); manager.update_priority(); + assert_eq!(manager.num_active(), MAX_CARDINALITY); for (i, (hist_idx, _)) in n_telemetry.iter().enumerate() { let key = NcclOpKey::NetSend(0x123, 42, *hist_idx); @@ -466,4 +801,169 @@ mod tests { assert_eq!(h.high_fidelity(), i < MAX_CARDINALITY); } } + + #[test] + fn latency_histogram_attributes() { + let mut manager = HistogramManager::new("test.latency", "ns", 16); + + let send = manager.get_histogram(NcclOpKey::NetSend(0x123, 4, 8)); + assert!(send + .attributes() + .contains(&KeyValue::new("nccl.source.rank", 4_i64))); + assert!(send + .attributes() + .contains(&KeyValue::new("nccl.destination.rank", 8_i64))); + + // for recv connections the local rank is the destination and the + // peer is the remote sender + let recv = manager.get_histogram(NcclOpKey::NetRecv(0x123, 4, 8)); + assert!(recv + .attributes() + .contains(&KeyValue::new("nccl.source.rank", 8_i64))); + assert!(recv + .attributes() + .contains(&KeyValue::new("nccl.destination.rank", 4_i64))); + + let aggregated = LatencyHistogram::new( + manager.histogram.clone(), + NcclOpKey::NetRecv(0x123, 4, 8), + String::from("host"), + false, + Arc::new(AtomicUsize::new(0)), + ); + assert!(aggregated + .attributes() + .contains(&KeyValue::new("nccl.metric.aggregated", true))); + } + + #[test] + fn latency_histogram_close_comm() { + const MAX_CARDINALITY: usize = 2; + let mut manager = HistogramManager::new("test.latency.close", "ns", MAX_CARDINALITY); + + manager.get_histogram(NcclOpKey::NetSend(0x1, 0, 1)); + manager.get_histogram(NcclOpKey::NetRecv(0x1, 0, 1)); + let low = manager.get_histogram(NcclOpKey::NetSend(0x2, 0, 1)); + assert_eq!(manager.num_active(), MAX_CARDINALITY); + assert!(!low.high_fidelity()); + + // closing a comm frees its slots for later connections + manager.close_comm(0x1); + assert_eq!(manager.num_active(), 0); + let h = manager.get_histogram(NcclOpKey::NetRecv(0x2, 0, 1)); + assert!(h.high_fidelity()); + } + + #[test] + fn size_class_boundaries() { + assert_eq!(size_class(0), "lt1m"); + assert_eq!(size_class((1 << 20) - 1), "lt1m"); + assert_eq!(size_class(1 << 20), "1m_16m"); + assert_eq!(size_class(16 << 20), "1m_16m"); + assert_eq!(size_class((16 << 20) + 1), "gt16m"); + } + + #[test] + fn duration_histogram_cardinality() { + const MAX_CARDINALITY: usize = 4; + let mut histogram = DurationHistogram::new("test.duration", "ns", MAX_CARDINALITY); + + for comm_hash in 0..(MAX_CARDINALITY as u64 * 2) { + histogram.record(1024, comm_hash, "ncclAllReduce", 0, 65536); + } + assert_eq!(histogram.num_active(), MAX_CARDINALITY); + assert!(histogram.high_fidelity(0, "ncclAllReduce", 0, 65536)); + assert!(!histogram.high_fidelity(MAX_CARDINALITY as u64, "ncclAllReduce", 0, 65536)); + + // repeated keys within the cap do not consume more slots, and the + // rank is part of the key + histogram.record(1024, 0, "ncclAllReduce", 0, 65536); + histogram.record(1024, 0, "ncclAllReduce", 1, 65536); + assert_eq!(histogram.num_active(), MAX_CARDINALITY); + assert!(!histogram.high_fidelity(0, "ncclAllReduce", 1, 65536)); + } + + #[test] + fn duration_histogram_update_priority() { + const MAX_CARDINALITY: usize = 2; + let mut histogram = DurationHistogram::new("test.duration.priority", "ns", MAX_CARDINALITY); + + // comms 0 and 1 take the slots first; 2 and 3 arrive later but are + // more active in this interval + for comm_hash in 0..4 { + histogram.record(1024, comm_hash, "ncclAllReduce", 0, 65536); + } + for _ in 0..8 { + histogram.record(1024, 2, "ncclAllReduce", 0, 65536); + histogram.record(1024, 3, "ncclAllReduce", 0, 65536); + } + assert!(histogram.high_fidelity(0, "ncclAllReduce", 0, 65536)); + assert!(!histogram.high_fidelity(2, "ncclAllReduce", 0, 65536)); + + histogram.update_priority(); + assert_eq!(histogram.num_active(), MAX_CARDINALITY); + assert!(!histogram.high_fidelity(0, "ncclAllReduce", 0, 65536)); + assert!(!histogram.high_fidelity(1, "ncclAllReduce", 0, 65536)); + assert!(histogram.high_fidelity(2, "ncclAllReduce", 0, 65536)); + assert!(histogram.high_fidelity(3, "ncclAllReduce", 0, 65536)); + } + + #[test] + fn duration_histogram_close_comm() { + const MAX_CARDINALITY: usize = 2; + let mut histogram = DurationHistogram::new("test.duration.close", "ns", MAX_CARDINALITY); + + histogram.record(1024, 1, "ncclAllReduce", 0, 65536); + histogram.record(1024, 1, "ncclAllGather", 0, 65536); + histogram.record(1024, 2, "ncclAllReduce", 0, 65536); + assert_eq!(histogram.num_active(), MAX_CARDINALITY); + assert!(!histogram.high_fidelity(2, "ncclAllReduce", 0, 65536)); + + // closing a comm frees its slots for later comms + histogram.close_comm(1); + assert_eq!(histogram.num_active(), 0); + histogram.record(1024, 2, "ncclAllReduce", 0, 1 << 24); + assert!(histogram.high_fidelity(2, "ncclAllReduce", 0, 1 << 24)); + } + + #[test] + fn gap_tracker_transitions() { + let tracker = GapTracker::new(); + + // no idle interval exists before the first activity completes + assert_eq!(tracker.begin_transition(|| 100), None); + tracker.activity_end(200); + assert_eq!(tracker.begin_transition(|| 500), Some(300)); + + // overlapping activity: only the 0 -> 1 transition observes a gap, + // and the gap starts when the last in-flight activity ends + assert_eq!(tracker.begin_transition(|| 600), None); + tracker.activity_end(800); + tracker.activity_end(700); // out-of-order end keeps the max stamp + assert_eq!(tracker.begin_transition(|| 1000), Some(200)); + tracker.activity_end(1100); + + // a gap can never be negative + assert_eq!(tracker.begin_transition(|| 1050), Some(0)); + } + + #[test] + fn gap_tracker_stall_detection() { + let tracker = GapTracker::new(); + assert!(!tracker.check_stalled()); + + // a leaked window is reported once, after a whole quiet interval + assert_eq!(tracker.begin_transition(|| 100), None); + assert!(!tracker.check_stalled()); + assert!(tracker.check_stalled()); + assert!(!tracker.check_stalled()); + + // ongoing activity is never mistaken for a stall + let tracker = GapTracker::new(); + assert_eq!(tracker.begin_transition(|| 100), None); + assert!(!tracker.check_stalled()); + tracker.activity_end(200); + assert_eq!(tracker.begin_transition(|| 300), Some(100)); + assert!(!tracker.check_stalled()); + } } diff --git a/src/profiler.rs b/src/profiler.rs index 500f72b..7866003 100644 --- a/src/profiler.rs +++ b/src/profiler.rs @@ -23,6 +23,7 @@ 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::otel_utils; use crate::profiler_shim; use crate::slab; use crate::spsc; @@ -61,6 +62,7 @@ pub struct Profiler { daemon: Mutex>, ncclop_cnt: AtomicU64, + pub gap_tracker: Option, pub remote_net_bytes: Option>, pub free_ncclop: slab::AtomicFreeList, pub free_proxyop: slab::AtomicFreeList, @@ -102,6 +104,11 @@ impl Profiler { daemon: Mutex::new(None), ncclop_cnt: AtomicU64::new(0), + gap_tracker: if config.otel_enable { + Some(otel_utils::GapTracker::new()) + } else { + None + }, remote_net_bytes: if config.heartbeat_collective_progress { Some(Arc::new(AtomicUsize::new(0))) } else { @@ -742,89 +749,122 @@ where } _ => panic!("unknown event type"), }; + if event.as_ref().is_some_and(is_nccl_activity) { + with_thread_state(|thread_state| { + let profiler = thread_state.profiler; + if let Some(gap_tracker) = profiler.gap_tracker.as_ref() { + gap_tracker.activity_begin(|| profiler.recent_timer_ns()); + } + }); + } Ok(event) } +// events that represent NCCL activity for gap tracking: an op being enqueued +// or its network / kernel work in progress; the proxy op and kernel channel +// windows keep an op in flight until its work actually completes, well after +// NCCL stops the op event itself at enqueue time +fn is_nccl_activity(event: &event::Event) -> bool { + std::matches!( + event, + event::Event::NcclOp(_) + | event::Event::NcclOpLite(_) + | event::Event::SmallNcclOp(_) + | event::Event::ProxyOp(_) + | event::Event::ProxyOpLite(_) + | event::Event::KernelCh(_) + ) +} + pub fn stop_event_handler(event: event::Event) -> NcclResult<()> { - with_thread_state(|thread_state| match event { - event::Event::Group(group) => { - thread_state.send_to_daemon(daemon::Message::Group(group), true); - } - event::Event::ProxyOpLite(data) => { - thread_state.fifo.prefetch_next(); - if let Some(ncclop) = data.parent_op { - if let Some(start_time) = thread_state.dec_ncclop_ref(data.info.pid, ncclop) { - let msg = daemon::Message::ProxyOpLite( - start_time, - start_time.elapsed().as_nanos() as u64, - data.info.clone(), - ); - thread_state.send_to_daemon(msg, true); - } - } - thread_state.proxyop_free_list.free(data); - } - event::Event::ProxyOp(mut data) => { - thread_state.fifo.prefetch_next(); - if let Some(step) = data.step_tracker.finalize() { - data.get_steps_mut(thread_state).push(step); - } - if thread_state.profiler.config.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); + let end_nccl_activity = is_nccl_activity(&event); + with_thread_state(|thread_state| { + match event { + event::Event::Group(group) => { + thread_state.send_to_daemon(daemon::Message::Group(group), true); } - if let Some(steps) = data.steps.take() { - let msg = daemon::Message::StepBatch(data.info.clone(), steps, true); - thread_state.send_to_daemon(msg, true); - } else { - let msg = daemon::Message::ProxyOp(data.info.id); - thread_state.send_to_daemon(msg, true); + event::Event::ProxyOpLite(data) => { + thread_state.fifo.prefetch_next(); + if let Some(ncclop) = data.parent_op { + if let Some(start_time) = thread_state.dec_ncclop_ref(data.info.pid, ncclop) { + let msg = daemon::Message::ProxyOpLite( + start_time, + start_time.elapsed().as_nanos() as u64, + data.info.clone(), + ); + thread_state.send_to_daemon(msg, true); + } + } + thread_state.proxyop_free_list.free(data); } - thread_state.proxyop_free_list.free(data); - } - event::Event::ProxyStep(mut data) => { - if data.end_time.is_none() { - data.end_time = Some(thread_state.profiler.recent_timer_instant()); + event::Event::ProxyOp(mut data) => { + thread_state.fifo.prefetch_next(); + if let Some(step) = data.step_tracker.finalize() { + data.get_steps_mut(thread_state).push(step); + } + if thread_state.profiler.config.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); + } + if let Some(steps) = data.steps.take() { + let msg = daemon::Message::StepBatch(data.info.clone(), steps, true); + thread_state.send_to_daemon(msg, true); + } else { + let msg = daemon::Message::ProxyOp(data.info.id); + thread_state.send_to_daemon(msg, true); + } + thread_state.proxyop_free_list.free(data); } - let step = - data.finalize(|t| (*t - thread_state.profiler.init_instant).as_nanos() as u64); - // SAFETY: NCCL guarantees that the proxyop event handle is - // live at this moment and the proxystep is on the same thread - // as the parent event handle. - // Therefore, dereference this pointer is safe as - // 1. the pointer is valid - // 2. there is no other threads accessing it - let parent = unsafe { &mut *data.parent }; - let steps = parent.get_steps_mut(thread_state); - steps.push(step); - if steps.is_full() { - let steps = parent.steps.take().unwrap(); - let msg = daemon::Message::StepBatch(parent.info.clone(), steps, false); - thread_state.send_to_daemon(msg, false); + event::Event::ProxyStep(mut data) => { + if data.end_time.is_none() { + data.end_time = Some(thread_state.profiler.recent_timer_instant()); + } + let step = + data.finalize(|t| (*t - thread_state.profiler.init_instant).as_nanos() as u64); + // SAFETY: NCCL guarantees that the proxyop event handle is + // live at this moment and the proxystep is on the same thread + // as the parent event handle. + // Therefore, dereference this pointer is safe as + // 1. the pointer is valid + // 2. there is no other threads accessing it + let parent = unsafe { &mut *data.parent }; + let steps = parent.get_steps_mut(thread_state); + steps.push(step); + if steps.is_full() { + let steps = parent.steps.take().unwrap(); + let msg = daemon::Message::StepBatch(parent.info.clone(), steps, false); + thread_state.send_to_daemon(msg, false); + } + thread_state.proxystep_free_list.free(data); } - thread_state.proxystep_free_list.free(data); - } - event::Event::KernelCh(kernelch) => { - thread_state.fifo.prefetch_next(); - if let Some(ncclop) = kernelch.parent_op { - if let Some(start_time) = - thread_state.dec_ncclop_ref(thread_state.profiler.pid, ncclop) - { - let msg = daemon::Message::KernelCh( - start_time, - start_time.elapsed().as_nanos() as u64, - ncclop, - ); - thread_state.send_to_daemon(msg, true); + event::Event::KernelCh(kernelch) => { + thread_state.fifo.prefetch_next(); + if let Some(ncclop) = kernelch.parent_op { + if let Some(start_time) = + thread_state.dec_ncclop_ref(thread_state.profiler.pid, ncclop) + { + let msg = daemon::Message::KernelCh( + start_time, + start_time.elapsed().as_nanos() as u64, + ncclop, + ); + thread_state.send_to_daemon(msg, true); + } } + thread_state.kernelch_free_list.free(kernelch); + } + event::Event::Dummy(_) => (), + event::Event::SmallNcclOp(_) => (), + event::Event::NcclOpLite(_) => {} + event::Event::NcclOp(_) => {} + } + if end_nccl_activity { + let profiler = thread_state.profiler; + if let Some(gap_tracker) = profiler.gap_tracker.as_ref() { + gap_tracker.activity_end(profiler.recent_timer_ns()); } - thread_state.kernelch_free_list.free(kernelch); } - event::Event::Dummy(_) => (), - event::Event::SmallNcclOp(_) => (), - event::Event::NcclOpLite(_) => {} - event::Event::NcclOp(_) => {} }); Ok(()) } diff --git a/src/step_tracker.rs b/src/step_tracker.rs index ae7a09a..a855e8e 100644 --- a/src/step_tracker.rs +++ b/src/step_tracker.rs @@ -58,14 +58,16 @@ impl EventStepInProgress { } fn finalize(&self) -> EventStep { - let fifo_wait_dur_ns = self.fifo_ready_time.map(|t| (t - self.start_time) as u32); + let fifo_wait_dur_ns = self + .fifo_ready_time + .map(|t| u32::try_from(t - self.start_time).unwrap_or(u32::MAX)); let net_start_time = self.fifo_ready_time.unwrap_or(self.start_time); EventStep { step: self.step, size: self.size.unwrap(), start_time: self.start_time, fifo_wait_dur_ns, - dur_ns: (self.end_time.unwrap() - net_start_time) as _, + dur_ns: u32::try_from(self.end_time.unwrap() - net_start_time).unwrap_or(u32::MAX), } } } @@ -299,4 +301,35 @@ mod tests { assert_eq!(step.size, SZ); } } + + #[test] + fn steptracker_saturates_dur_over_u32_max() { + const SZ: usize = 65536; + let mut tracker = StepTracker::new(true, true, false); + + let clock = std::cell::Cell::new(0u64); + let get_time = || -> u64 { + let t = clock.get(); + clock.set(t + 6_000_000_000); // 6 s between events, > u32::MAX ns + t + }; + + let args = nccl_metadata::ProxyOpStateV1::new(0, 0); + let first = tracker.update_step( + profiler_shim::proxy_event_state::v1::SEND_TRANSMITTED, + &args, + get_time, + ); + assert!(first.is_none()); + + let args = nccl_metadata::ProxyOpStateV1::new(0, SZ); + let _ = tracker.update_step( + profiler_shim::proxy_event_state::v1::SEND_DONE, + &args, + get_time, + ); + + let step = tracker.finalize().expect("expected a finalized step"); + assert_eq!(step.dur_ns, u32::MAX); + } } From 615269750188aea3468245597f38b9912c467ff6 Mon Sep 17 00:00:00 2001 From: sputti-czi Date: Tue, 7 Jul 2026 16:54:56 -0400 Subject: [PATCH 2/2] chore --- src/otel_utils.rs | 14 +------------- 1 file changed, 1 insertion(+), 13 deletions(-) diff --git a/src/otel_utils.rs b/src/otel_utils.rs index ff061d1..5050b13 100644 --- a/src/otel_utils.rs +++ b/src/otel_utils.rs @@ -1,16 +1,4 @@ -// 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. +// commited by sputti-czi use crate::config; use crate::daemon::AtomicHistogram;