diff --git a/Cargo.lock b/Cargo.lock index 651b58ac..e0309640 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3431,6 +3431,7 @@ name = "probing-memtable" version = "0.2.5" dependencies = [ "libc", + "log", "memmap2", "pco", "xxhash-rust", @@ -3500,10 +3501,13 @@ dependencies = [ "pyo3-build-config", "regex", "rustc-demangle", + "serde", "serde_json", + "sha2", "signal-hook-registry", "tempfile", "tokio", + "uuid", ] [[package]] @@ -4061,6 +4065,17 @@ dependencies = [ "digest", ] +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + [[package]] name = "shlex" version = "2.0.1" diff --git a/Makefile b/Makefile index dcfb840c..283cc0f1 100644 --- a/Makefile +++ b/Makefile @@ -47,6 +47,9 @@ CLIPPY_WORKSPACE := cargo clippy --workspace --all-targets --no-default-features CLIPPY_CORE := cargo clippy -p probing-core --all-targets --no-default-features $(CLIPPY_DENY) CLIPPY_WEB := cd web && cargo clippy --all-targets $(CLIPPY_DENY) +FMT_WORKSPACE := cargo fmt --all +FMT_WEB := cd web && cargo fmt --all + # ============================================================================== .PHONY: help help: @@ -60,6 +63,7 @@ help: @echo " install-wheel pip install dist/probing-*.whl" @echo " venv create/refresh project .venv (used by develop and CI)" @echo " test / lint Full test and lint suites" + @echo " fmt rustfmt workspace + web (same paths as CI)" @echo " check-dev Quick env sanity check" @echo " clean Remove build artifacts" @echo "" @@ -203,7 +207,7 @@ PYTEST_WHEEL_FLAGS := --import-mode=importlib -o pythonpath= -o "addopts=--verbo PYTEST_WHEEL_EXTRA ?= .PHONY: test test-rust test-rust-unit test-rust-regression test-python test-python-unit test-python-regression test-doctest test-python-wheel coverage-python-wheel -.PHONY: lint lint-python lint-rust lint-core clippy clippy-fix coverage coverage-rust coverage-python bootstrap clean docs-install docs docs-serve docs-clean +.PHONY: fmt fmt-check lint lint-python lint-rust lint-core clippy clippy-fix coverage coverage-rust coverage-python bootstrap clean docs-install docs docs-serve docs-clean test: test-rust test-python test-rust: test-rust-unit test-rust-regression @@ -241,6 +245,12 @@ coverage-python-wheel: $(MAKE) test-python-wheel PYTEST_WHEEL_EXTRA="--cov=probing --cov=tests --cov-report=xml:coverage.xml" lint: lint-python lint-rust +fmt: + $(FMT_WORKSPACE) + $(FMT_WEB) +fmt-check: + $(FMT_WORKSPACE) -- --check + $(FMT_WEB) -- --check lint-core: $(CLIPPY_CORE) lint-python: diff --git a/examples/crash_demo.py b/examples/crash_demo.py new file mode 100644 index 00000000..8d111cdf --- /dev/null +++ b/examples/crash_demo.py @@ -0,0 +1,42 @@ +#!/usr/bin/env python3 +"""单进程 crash demo — ``PROBING=1 python examples/crash_demo.py [--mode exception]``. + +Probing 由 site hook 自动加载；本脚本不 import probing。 +``record``：后台线程异常，写入 memtable 后正常退出；``exception``：主线程崩溃。 +""" + +from __future__ import annotations + +import argparse +import os +import threading + + +def _crash(*, record_only: bool) -> None: + def boom() -> None: + def inner() -> None: + raise ValueError("demo crash from crash_demo.py") + + inner() + + if record_only: + t = threading.Thread(target=boom, name="crash-demo", daemon=True) + t.start() + t.join() + else: + boom() + + +def main() -> None: + p = argparse.ArgumentParser() + p.add_argument("--mode", choices=("record", "exception"), default="record") + args = p.parse_args() + + os.environ.setdefault("RANK", "0") + os.environ.setdefault("LOCAL_RANK", "0") + + _crash(record_only=args.mode == "record") + + +if __name__ == "__main__": + main() diff --git a/examples/crash_torchrun_demo.py b/examples/crash_torchrun_demo.py new file mode 100644 index 00000000..57e953e2 --- /dev/null +++ b/examples/crash_torchrun_demo.py @@ -0,0 +1,51 @@ +#!/usr/bin/env python3 +"""torchrun crash demo — run via ``./examples/run_crash_torchrun.sh``. + +Probing loads from ``PROBING=…`` (site hook); crash capture needs no imports here. +Modes: ``record`` (thread exception, process stays up), ``exception`` / ``all`` (fatal). +""" + +from __future__ import annotations + +import argparse +import os +import threading + +import torch.distributed as dist + + +def _crash(rank: int, *, record_only: bool) -> None: + def boom() -> None: + raise RuntimeError(f"demo crash on rank {rank}") + + if record_only: + t = threading.Thread(target=boom, name=f"crash-r{rank}", daemon=True) + t.start() + t.join() + else: + boom() + + +def main() -> None: + p = argparse.ArgumentParser() + p.add_argument("--mode", choices=("record", "exception", "all"), default="record") + p.add_argument("--crash-rank", type=int, default=1) + p.add_argument("--backend", default=os.environ.get("PROBING_TORCH_BACKEND", "gloo")) + args = p.parse_args() + + dist.init_process_group(backend=args.backend) + rank, world = dist.get_rank(), dist.get_world_size() + print(f"[rank {rank}/{world}] ready", flush=True) + + dist.barrier() + if args.mode == "all" or rank == args.crash_rank: + _crash(rank, record_only=args.mode == "record") + + if args.mode == "record": + dist.barrier() + + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/examples/run_crash_torchrun.sh b/examples/run_crash_torchrun.sh new file mode 100755 index 00000000..1d4dd167 --- /dev/null +++ b/examples/run_crash_torchrun.sh @@ -0,0 +1,26 @@ +#!/usr/bin/env bash +# torchrun crash demo — run from repo root: ./examples/run_crash_torchrun.sh [mode] [crash-rank] +# mode: record (default) | exception | all +set -euo pipefail + +cd "$(dirname "$0")/.." +[[ -f .venv/bin/activate ]] && source .venv/bin/activate + +MODE="${1:-record}" +CRASH_RANK="${2:-1}" +NPROC="${NPROC:-4}" +MASTER_PORT="${MASTER_PORT:-29571}" + +export PROBING="${PROBING:-regex:crash_torchrun_demo.py}" +export MASTER_ADDR=127.0.0.1 +export MASTER_PORT +[[ "$(uname -s)" == Darwin ]] && export GLOO_SOCKET_IFNAME="${GLOO_SOCKET_IFNAME:-lo0}" + +python -c "import torch" 2>/dev/null || { + echo "need torch: uv pip install torch" >&2 + exit 1 +} + +exec torchrun --standalone --nnodes=1 --nproc_per_node="$NPROC" \ + --master_addr=127.0.0.1 --master_port="$MASTER_PORT" --local_addr=127.0.0.1 \ + examples/crash_torchrun_demo.py --mode "$MODE" --crash-rank "$CRASH_RANK" diff --git a/probing/core/src/core/memtable_sql.rs b/probing/core/src/core/memtable_sql.rs index db3d2a58..73821b5f 100644 --- a/probing/core/src/core/memtable_sql.rs +++ b/probing/core/src/core/memtable_sql.rs @@ -29,7 +29,7 @@ use std::collections::{BTreeSet, HashSet}; use std::panic::AssertUnwindSafe; use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::{Arc, Mutex}; +use std::sync::{Arc, Mutex, MutexGuard}; use std::thread::JoinHandle; use std::time::Duration; @@ -1104,6 +1104,18 @@ pub struct ColdCompactor { handle: Mutex>>, } +fn lock_compactor_handle( + m: &Mutex>>, +) -> MutexGuard<'_, Option>> { + match m.lock() { + Ok(guard) => guard, + Err(e) => { + log::debug!("ColdCompactor handle mutex poisoned; recovering"); + e.into_inner() + } + } +} + impl ColdCompactor { pub fn instance() -> &'static Self { static INSTANCE: Lazy = Lazy::new(|| ColdCompactor { @@ -1176,7 +1188,7 @@ impl ColdCompactor { } }) .expect("spawn memc-compactor thread"); - *self.handle.lock().unwrap() = Some(handle); + *lock_compactor_handle(&self.handle) = Some(handle); } /// Signal the thread to flush and exit, then join it. @@ -1184,7 +1196,7 @@ impl ColdCompactor { if !self.running.swap(false, Ordering::SeqCst) { return; } - if let Some(h) = self.handle.lock().unwrap().take() { + if let Some(h) = lock_compactor_handle(&self.handle).take() { let _ = h.join(); } } diff --git a/probing/core/src/trace/mod.rs b/probing/core/src/trace/mod.rs index fd26b574..461523e9 100644 --- a/probing/core/src/trace/mod.rs +++ b/probing/core/src/trace/mod.rs @@ -3,8 +3,8 @@ mod step; pub use span::{attr, Attribute, Ele, Event, Location, Span, SpanStatus, Timestamp}; pub use step::{ - advance_micro_step, current_micro_step, set_micro_batches, step_snapshot, sync_micro_step, - StepSnapshot, + advance_micro_step, crash_atomic_step, crash_step_snapshot, current_micro_step, + set_micro_batches, step_snapshot, sync_micro_step, StepSnapshot, }; // --- Custom Error Type --- diff --git a/probing/core/src/trace/step.rs b/probing/core/src/trace/step.rs index 64bae4f9..748bb166 100644 --- a/probing/core/src/trace/step.rs +++ b/probing/core/src/trace/step.rs @@ -1,4 +1,15 @@ use std::cell::RefCell; +use std::sync::atomic::{AtomicI64, AtomicU64, Ordering}; + +static GLOBAL_MICRO_STEP: AtomicU64 = AtomicU64::new(0); +static GLOBAL_MICRO_BATCHES: AtomicU64 = AtomicU64::new(1); +static CACHED_RANK: AtomicI64 = AtomicI64::new(0); +static CACHED_WORLD_SIZE: AtomicI64 = AtomicI64::new(1); + +fn refresh_cached_dist() { + CACHED_RANK.store(read_rank(), Ordering::Relaxed); + CACHED_WORLD_SIZE.store(read_world_size(), Ordering::Relaxed); +} /// Training step coordinates. /// @@ -25,6 +36,7 @@ struct StepContext { impl Default for StepContext { fn default() -> Self { + refresh_cached_dist(); Self { micro_step: 0, micro_batches: read_micro_batches(), @@ -49,16 +61,53 @@ impl StepContext { fn sync_micro_step(&mut self, step: u64) -> StepSnapshot { self.micro_step = step; + publish_global(self.micro_step, self.micro_batches); self.snapshot() } fn advance_micro_step(&mut self) -> StepSnapshot { self.micro_step = self.micro_step.saturating_add(1); + publish_global(self.micro_step, self.micro_batches); self.snapshot() } fn set_micro_batches(&mut self, micro_batches: u64) { self.micro_batches = micro_batches.max(1); + GLOBAL_MICRO_BATCHES.store(self.micro_batches, Ordering::Relaxed); + } +} + +fn publish_global(micro_step: u64, micro_batches: u64) { + GLOBAL_MICRO_STEP.fetch_max(micro_step, Ordering::Relaxed); + GLOBAL_MICRO_BATCHES.store(micro_batches.max(1), Ordering::Relaxed); +} + +/// Best-effort step coordinates for crash reporting (prefers thread-local, falls +/// back to process-wide high-water mark when the crashing thread never advanced step). +pub fn crash_step_snapshot() -> StepSnapshot { + let local = step_snapshot(); + if local.micro_step > 0 { + return local; + } + atomic_step_snapshot() +} + +/// Async-signal-safe: global high-water step only (no thread-local / env reads). +pub fn crash_atomic_step() -> StepSnapshot { + atomic_step_snapshot() +} + +fn atomic_step_snapshot() -> StepSnapshot { + let micro = GLOBAL_MICRO_STEP.load(Ordering::Relaxed); + let batches = GLOBAL_MICRO_BATCHES.load(Ordering::Relaxed).max(1); + let training = training_step_for(micro, batches); + StepSnapshot { + micro_step: micro, + local_step: training, + global_step: training, + micro_batches: batches, + rank: CACHED_RANK.load(Ordering::Relaxed), + world_size: CACHED_WORLD_SIZE.load(Ordering::Relaxed), } } diff --git a/probing/extensions/python/Cargo.toml b/probing/extensions/python/Cargo.toml index 3d1b8e57..cc8b8789 100644 --- a/probing/extensions/python/Cargo.toml +++ b/probing/extensions/python/Cargo.toml @@ -25,7 +25,10 @@ ctor = { workspace = true } log = { workspace = true } nix = { workspace = true } once_cell = { workspace = true } +serde = { workspace = true } serde_json = { workspace = true } +sha2 = "0.10" +uuid = { version = "1.0", features = ["v4"] } tokio = { workspace = true } html-escape = "0.2" diff --git a/probing/extensions/python/src/extensions/python.rs b/probing/extensions/python/src/extensions/python.rs index f3dd9d18..d8312371 100644 --- a/probing/extensions/python/src/extensions/python.rs +++ b/probing/extensions/python/src/extensions/python.rs @@ -22,7 +22,6 @@ pub use tbls::PythonProbeDataSource; use crate::features::stack_tracer::{SignalTracer, StackTracer}; use crate::python::enable_crash_handler; use crate::python::enable_monitoring; -use crate::python::CRASH_HANDLER; mod exttbls; mod tbls; @@ -51,8 +50,8 @@ impl Display for PyExtList { /// Python integration with the probing system #[derive(Debug, Default, ProbeExtension)] pub struct PythonExt { - /// Path to Python crash handler script (executed when interpreter crashes) - #[option(aliases = ["crash.handler"])] + /// Enable the crash handler module (`probing.crash`). Aliases: `crash.handler`. + #[option(aliases = ["crash.handler", "crash.enabled"])] crash_handler: Maybe, /// Path to Python monitoring handler script @@ -84,6 +83,10 @@ impl ProbeExtensionCall for PythonExt { ); let normalized_path = path.trim_start_matches('/'); + if normalized_path.starts_with("crash/") { + return crate::features::crash::handle_http(normalized_path, params, body) + .map_err(EngineError::PluginError); + } call_python_handler(normalized_path, params, body) } } @@ -102,7 +105,11 @@ impl PythonExt { )), Maybe::Just(handler) => { self.crash_handler = crash_handler.clone(); - CRASH_HANDLER.lock().unwrap().replace(handler.to_string()); + let lowered = handler.trim().to_ascii_lowercase(); + if lowered == "off" || lowered == "0" || lowered == "false" { + log::info!("Python crash handler disabled"); + return Ok(()); + } match enable_crash_handler() { Ok(_) => { log::info!("Python crash handler enabled: {handler}"); diff --git a/probing/extensions/python/src/extensions/python/exttbls.rs b/probing/extensions/python/src/extensions/python/exttbls.rs index 3b55867e..72bca5cf 100644 --- a/probing/extensions/python/src/extensions/python/exttbls.rs +++ b/probing/extensions/python/src/extensions/python/exttbls.rs @@ -13,7 +13,7 @@ //! epoch, `I64`) is always present, matching the previous TimeSeries layout. use std::collections::HashMap; -use std::sync::{Arc, Mutex}; +use std::sync::{Arc, Mutex, MutexGuard}; use crate::features::native_bridge::with_detached_native; use once_cell::sync::Lazy; @@ -404,6 +404,26 @@ impl std::fmt::Debug for ExternBacking { pub static EXTERN_TABLES: Lazy>>>> = Lazy::new(|| Mutex::new(Default::default())); +fn lock_extern_tables() -> MutexGuard<'static, HashMap>>> { + match EXTERN_TABLES.lock() { + Ok(guard) => guard, + Err(e) => { + log::warn!("EXTERN_TABLES mutex poisoned; recovering"); + e.into_inner() + } + } +} + +fn lock_backing(backing: &Mutex) -> MutexGuard<'_, ExternBacking> { + match backing.lock() { + Ok(guard) => guard, + Err(e) => { + log::warn!("ExternBacking mutex poisoned; recovering"); + e.into_inner() + } + } +} + #[pyclass(from_py_object)] #[derive(Clone, Debug)] pub struct ExternalTable(Arc>, usize); @@ -437,9 +457,7 @@ impl ExternalTable { table_doc, column_docs, ))); - backing - .lock() - .expect("extern table lock") + lock_backing(backing.as_ref()) .ensure_registered() .expect("failed to register extern table for SQL catalog"); backing @@ -471,7 +489,7 @@ impl ExternalTable { table_doc, column_docs.unwrap_or_default(), ); - EXTERN_TABLES.lock().unwrap().insert(name, backing.clone()); + lock_extern_tables().insert(name, backing.clone()); ExternalTable(backing, ncolumn) }) } @@ -480,9 +498,9 @@ impl ExternalTable { fn get(_cls: &Bound<'_, PyType>, name: &str) -> PyResult { let name = name.to_string(); with_detached_native(move || { - let binding = EXTERN_TABLES.lock().unwrap(); + let binding = lock_extern_tables(); if let Some(backing) = binding.get(&name) { - let ncolumn = backing.lock().unwrap().columns.len(); + let ncolumn = lock_backing(backing.as_ref()).columns.len(); Ok(ExternalTable(backing.clone(), ncolumn)) } else { Err(pyo3::exceptions::PyValueError::new_err(format!( @@ -508,9 +526,9 @@ impl ExternalTable { let _ = chunk_size; let name = name.to_string(); with_detached_native(move || { - let mut binding = EXTERN_TABLES.lock().unwrap(); + let mut binding = lock_extern_tables(); if let Some(backing) = binding.get(&name) { - let ncolumn = backing.lock().unwrap().columns.len(); + let ncolumn = lock_backing(backing.as_ref()).columns.len(); Ok(ExternalTable(backing.clone(), ncolumn)) } else { let ncolumn = columns.len(); @@ -532,14 +550,14 @@ impl ExternalTable { fn drop(_cls: &Bound<'_, PyType>, name: &str) -> PyResult<()> { let name = name.to_string(); with_detached_native(move || { - let _ = EXTERN_TABLES.lock().unwrap().remove(&name); + let _ = lock_extern_tables().remove(&name); Ok(()) }) } fn names(&self) -> Vec { let backing = self.0.clone(); - with_detached_native(move || backing.lock().unwrap().columns.clone()) + with_detached_native(move || lock_backing(backing.as_ref()).columns.clone()) } fn append(&mut self, values: Vec>) -> PyResult<()> { @@ -551,9 +569,7 @@ impl ExternalTable { let eles = Self::extract_eles(values); let backing = self.0.clone(); with_detached_native(move || { - backing - .lock() - .unwrap() + lock_backing(backing.as_ref()) .append(now_micros(), &eles) .map_err(pyo3::exceptions::PyValueError::new_err) }) @@ -568,9 +584,7 @@ impl ExternalTable { let eles = Self::extract_eles(values); let backing = self.0.clone(); with_detached_native(move || { - backing - .lock() - .unwrap() + lock_backing(backing.as_ref()) .append(t, &eles) .map_err(pyo3::exceptions::PyValueError::new_err) }) @@ -587,7 +601,7 @@ impl ExternalTable { fn take(&self, limit: Option) -> PyResult> { let backing = self.0.clone(); with_detached_native(move || { - let rows = backing.lock().unwrap().take(limit); + let rows = lock_backing(backing.as_ref()).take(limit); let result = rows .iter() .map(|(t, vals)| { diff --git a/probing/extensions/python/src/features/crash.rs b/probing/extensions/python/src/features/crash.rs deleted file mode 100644 index e0852c3f..00000000 --- a/probing/extensions/python/src/features/crash.rs +++ /dev/null @@ -1,203 +0,0 @@ -//! Fatal-signal backtrace dumper. -//! -//! On `SIGSEGV` / `SIGBUS` / `SIGABRT` / `SIGILL` / `SIGFPE` we print the -//! crashing thread's native backtrace to stderr, then restore the default -//! disposition and re-raise the signal so the process still produces a core -//! dump / exits with the usual status. This is purely a debugging aid for -//! diagnosing the profiler/native crashes. -//! -//! The handler runs on a dedicated `sigaltstack` (so a stack-overflow crash can -//! still be reported) and avoids heap allocation on its hot path: integers are -//! formatted into stack buffers and symbol names are written as their raw bytes -//! via `write(2)`. Symbolization (`backtrace::*_unsynchronized`) is not strictly -//! async-signal-safe, but this only ever runs once, while the process is already -//! dying; a recursive fault is caught by a guard that immediately re-raises with -//! the default handler. - -use core::ffi::{c_int, c_void}; -use std::sync::atomic::{AtomicBool, Ordering}; - -use nix::libc; - -static INSTALLED: AtomicBool = AtomicBool::new(false); -static IN_HANDLER: AtomicBool = AtomicBool::new(false); - -const FATAL_SIGNALS: [c_int; 5] = [ - libc::SIGSEGV, - libc::SIGBUS, - libc::SIGABRT, - libc::SIGILL, - libc::SIGFPE, -]; - -const MAX_FRAMES: usize = 256; -const ALT_STACK_SIZE: usize = 256 * 1024; - -static mut ALT_STACK: [u8; ALT_STACK_SIZE] = [0u8; ALT_STACK_SIZE]; - -fn sig_name(sig: c_int) -> &'static str { - match sig { - libc::SIGSEGV => "SIGSEGV", - libc::SIGBUS => "SIGBUS", - libc::SIGABRT => "SIGABRT", - libc::SIGILL => "SIGILL", - libc::SIGFPE => "SIGFPE", - _ => "SIGNAL", - } -} - -#[inline] -unsafe fn write_str(s: &str) { - let _ = libc::write(2, s.as_ptr() as *const c_void, s.len()); -} - -#[inline] -unsafe fn write_bytes(b: &[u8]) { - let _ = libc::write(2, b.as_ptr() as *const c_void, b.len()); -} - -/// Write `v` as lowercase hex (no `0x` prefix), no allocation. -unsafe fn write_hex(v: usize) { - if v == 0 { - write_str("0"); - return; - } - let mut buf = [0u8; 16]; - let mut i = buf.len(); - let mut val = v; - while val > 0 { - i -= 1; - let d = (val & 0xf) as u8; - buf[i] = if d < 10 { b'0' + d } else { b'a' + (d - 10) }; - val >>= 4; - } - write_bytes(&buf[i..]); -} - -/// Write `v` as decimal, no allocation. -unsafe fn write_dec(v: usize) { - if v == 0 { - write_str("0"); - return; - } - let mut buf = [0u8; 20]; - let mut i = buf.len(); - let mut val = v; - while val > 0 { - i -= 1; - buf[i] = b'0' + (val % 10) as u8; - val /= 10; - } - write_bytes(&buf[i..]); -} - -unsafe fn current_tid() -> u64 { - #[cfg(target_os = "linux")] - { - libc::syscall(libc::SYS_gettid) as u64 - } - #[cfg(target_os = "macos")] - { - let mut t: u64 = 0; - libc::pthread_threadid_np(0, &mut t); - t - } - #[cfg(not(any(target_os = "linux", target_os = "macos")))] - { - 0 - } -} - -unsafe fn restore_and_reraise(sig: c_int) { - let mut sa: libc::sigaction = std::mem::zeroed(); - sa.sa_sigaction = libc::SIG_DFL; - libc::sigemptyset(&mut sa.sa_mask); - sa.sa_flags = 0; - libc::sigaction(sig, &sa, std::ptr::null_mut()); - libc::raise(sig); -} - -unsafe extern "C" fn crash_handler(sig: c_int, info: *mut libc::siginfo_t, _uctx: *mut c_void) { - // A fault *inside* the dump (e.g. while symbolizing) must not loop forever. - if IN_HANDLER.swap(true, Ordering::SeqCst) { - restore_and_reraise(sig); - return; - } - - write_str("\n==== probing: fatal signal "); - write_str(sig_name(sig)); - write_str(" ("); - write_dec(sig as usize); - write_str(") on thread "); - write_dec(current_tid() as usize); - write_str(" ====\n"); - - if !info.is_null() { - let addr = (*info).si_addr() as usize; - write_str("fault address: 0x"); - write_hex(addr); - write_str("\n"); - } - - write_str("native backtrace (crashing thread):\n"); - - let mut idx = 0usize; - backtrace::trace_unsynchronized(|frame| { - let ip = frame.ip() as usize; - write_str(" #"); - write_dec(idx); - write_str(" 0x"); - write_hex(ip); - write_str(" "); - - let mut wrote_name = false; - backtrace::resolve_frame_unsynchronized(frame, |symbol| { - if !wrote_name { - if let Some(name) = symbol.name() { - if let Some(s) = name.as_str() { - write_bytes(s.as_bytes()); - wrote_name = true; - } - } - } - }); - if !wrote_name { - write_str(""); - } - write_str("\n"); - - idx += 1; - idx < MAX_FRAMES - }); - - write_str("==== end probing backtrace; re-raising "); - write_str(sig_name(sig)); - write_str(" ====\n"); - - restore_and_reraise(sig); -} - -/// Install backtrace-on-crash handlers for the common fatal signals. Idempotent. -pub fn install_crash_handler() { - if INSTALLED.swap(true, Ordering::SeqCst) { - return; - } - - unsafe { - let mut ss: libc::stack_t = std::mem::zeroed(); - ss.ss_sp = core::ptr::addr_of_mut!(ALT_STACK) as *mut c_void; - ss.ss_size = ALT_STACK_SIZE; - ss.ss_flags = 0; - libc::sigaltstack(&ss, std::ptr::null_mut()); - - for &sig in FATAL_SIGNALS.iter() { - let mut sa: libc::sigaction = std::mem::zeroed(); - sa.sa_sigaction = crash_handler as *const () as usize; - sa.sa_flags = libc::SA_SIGINFO | libc::SA_ONSTACK; - libc::sigemptyset(&mut sa.sa_mask); - libc::sigaction(sig, &sa, std::ptr::null_mut()); - } - } - - log::info!("probing: crash backtrace handler installed"); -} diff --git a/probing/extensions/python/src/features/crash/api.rs b/probing/extensions/python/src/features/crash/api.rs new file mode 100644 index 00000000..436384a1 --- /dev/null +++ b/probing/extensions/python/src/features/crash/api.rs @@ -0,0 +1,56 @@ +//! Python-facing crash API (``probing._core``). + +use pyo3::prelude::*; + +use super::config; +use super::grace; +use super::handler::{self, CrashInput}; + +#[pyfunction] +pub fn crash_enabled() -> bool { + config::enabled() +} + +#[pyfunction] +#[allow(clippy::too_many_arguments)] // PyO3 signature mirrors Python excepthook payload +#[pyo3(signature = (kind, exception_type, message, traceback, top_frame, native_backtrace="", finalize=true, thread_stacks="", crash_thread=""))] +pub fn record_crash( + kind: &str, + exception_type: &str, + message: &str, + traceback: &str, + top_frame: &str, + native_backtrace: &str, + finalize: bool, + thread_stacks: &str, + crash_thread: &str, +) -> PyResult { + Ok(handler::record(CrashInput { + kind: kind.to_string(), + exception_type: exception_type.to_string(), + message: message.to_string(), + top_frame: top_frame.to_string(), + traceback: traceback.to_string(), + native_backtrace: native_backtrace.to_string(), + crash_thread: crash_thread.to_string(), + thread_stacks: thread_stacks.to_string(), + finalize, + })) +} + +#[pyfunction] +#[pyo3(signature = (op, group_size=0, bytes=0, global_step=-1))] +pub fn note_last_comm(op: &str, group_size: i32, bytes: i64, global_step: i64) { + handler::note_last_comm(op, group_size, bytes, global_step); +} + +#[pyfunction] +pub fn request_crash_hold() { + grace::request_hold(); +} + +#[pyfunction] +#[pyo3(signature = (exit_after=true))] +pub fn request_crash_release(exit_after: bool) { + grace::request_release(exit_after); +} diff --git a/probing/extensions/python/src/features/crash/grace.rs b/probing/extensions/python/src/features/crash/grace.rs new file mode 100644 index 00000000..5b8f6665 --- /dev/null +++ b/probing/extensions/python/src/features/crash/grace.rs @@ -0,0 +1,134 @@ +//! Grace period and interactive hold before process exit (non-signal path). + +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Mutex; +use std::time::{Duration, Instant}; + +use super::config; +use super::context; + +static HELD: AtomicBool = AtomicBool::new(false); +static RELEASED: AtomicBool = AtomicBool::new(false); +static EXIT_AFTER_RELEASE: Mutex = Mutex::new(true); +static SIGNALS_REGISTERED: AtomicBool = AtomicBool::new(false); + +pub fn is_held() -> bool { + HELD.load(Ordering::SeqCst) +} + +pub fn request_hold() { + HELD.store(true, Ordering::SeqCst); +} + +pub fn request_release(exit_after: bool) { + if let Ok(mut guard) = EXIT_AFTER_RELEASE.lock() { + *guard = exit_after; + } + RELEASED.store(true, Ordering::SeqCst); + HELD.store(false, Ordering::SeqCst); +} + +pub fn should_run_grace() -> bool { + if config::force_hold() { + return true; + } + if config::grace_sec() == 0 { + return false; + } + if config::grace_all_ranks() { + return true; + } + let ctx = context::snapshot(); + ctx.rank == 0 || ctx.local_rank == 0 +} + +pub fn grace_and_maybe_hold(exit_code: i32) -> i32 { + register_signals(); + let ctx = context::snapshot(); + + if config::force_hold() { + request_hold(); + } + + if !should_run_grace() && !is_held() { + return exit_code; + } + + let deadline = Instant::now() + Duration::from_secs(config::grace_sec()); + loop { + if RELEASED.load(Ordering::SeqCst) { + let exit_after = EXIT_AFTER_RELEASE.lock().map(|g| *g).unwrap_or(true); + return if exit_after { exit_code } else { 0 }; + } + if hold_triggered(ctx.pid) { + request_hold(); + print_hold_banner(ctx.pid, ctx.rank); + while !RELEASED.load(Ordering::SeqCst) { + std::thread::sleep(Duration::from_millis(200)); + } + let exit_after = EXIT_AFTER_RELEASE.lock().map(|g| *g).unwrap_or(true); + return if exit_after { exit_code } else { 0 }; + } + if !is_held() && config::grace_sec() > 0 && Instant::now() >= deadline { + return exit_code; + } + if !is_held() && config::grace_sec() == 0 && !config::force_hold() { + return exit_code; + } + if !is_held() && config::grace_sec() > 0 { + let remaining = deadline.saturating_duration_since(Instant::now()).as_secs(); + eprint!( + "\r[probing crash] exiting in {remaining}s — ENTER/touch/kill -USR1 to hold..." + ); + } + std::thread::sleep(Duration::from_millis(250)); + } +} + +fn register_signals() { + if SIGNALS_REGISTERED.swap(true, Ordering::SeqCst) { + return; + } + #[cfg(unix)] + { + use nix::sys::signal::{self, SigHandler, Signal}; + unsafe { + let _ = signal::signal(Signal::SIGUSR1, SigHandler::Handler(on_hold_signal)); + let _ = signal::signal(Signal::SIGUSR2, SigHandler::Handler(on_release_signal)); + } + } +} + +#[cfg(unix)] +extern "C" fn on_hold_signal(_: nix::libc::c_int) { + request_hold(); +} + +#[cfg(unix)] +extern "C" fn on_release_signal(_: nix::libc::c_int) { + request_release(true); +} + +fn hold_triggered(pid: i32) -> bool { + if config::force_hold() { + return true; + } + if std::path::Path::new(&context::hold_file_path(pid)).exists() { + return true; + } + if is_held() { + return true; + } + false +} + +fn print_hold_banner(pid: i32, rank: i32) { + eprintln!( + "\nProcess HELD for debugging (pid={pid}, rank={rank}).\n\ + attach : gdb -p {pid}\n\ + python : py-spy dump --pid {pid}\n\ + release: curl -X POST http://127.0.0.1:/apis/pythonext/crash/release\n\ + rm {} && kill -USR2 {pid}\n", + context::hold_file_path(pid) + ); +} diff --git a/probing/extensions/python/src/features/crash/handler.rs b/probing/extensions/python/src/features/crash/handler.rs new file mode 100644 index 00000000..8cd7110a --- /dev/null +++ b/probing/extensions/python/src/features/crash/handler.rs @@ -0,0 +1,259 @@ +//! Python exception crash path: build event → spill → report → grace. + +use std::fs; +use std::io::Write; +use std::path::{Path, PathBuf}; +use std::time::{SystemTime, UNIX_EPOCH}; + +use once_cell::sync::Lazy; +use probing_core::trace::crash_step_snapshot; +use probing_memtable::discover::default_dir; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +use std::sync::Mutex; +use uuid::Uuid; + +use crate::features::tracing::crash_span_snapshot; + +use super::grace; +use super::report; +use super::{config, context}; + +static LAST_COMM: Lazy>> = Lazy::new(|| Mutex::new(None)); + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct CrashEvent { + pub event_id: String, + pub timestamp_ns: i64, + pub kind: String, + pub rank: i32, + pub local_rank: i32, + pub world_size: i32, + pub host: String, + pub pid: i32, + pub exception_type: String, + pub message: String, + pub top_frame: String, + pub traceback: String, + pub native_backtrace: String, + pub crash_thread: String, + pub thread_stacks: String, + pub fingerprint: String, + pub global_step: i64, + pub local_step: i64, + pub micro_step: i64, + pub training_phase: String, + pub active_span: String, + pub last_comm_op: String, +} + +pub struct CrashInput { + pub kind: String, + pub exception_type: String, + pub message: String, + pub top_frame: String, + pub traceback: String, + pub native_backtrace: String, + pub crash_thread: String, + pub thread_stacks: String, + pub finalize: bool, +} + +pub fn note_last_comm(op: &str, _group_size: i32, _bytes: i64, _global_step: i64) { + if op.is_empty() { + return; + } + if let Ok(mut guard) = LAST_COMM.lock() { + *guard = Some(op.to_string()); + } +} + +pub fn signal_spill_path(pid: i32) -> PathBuf { + crash_dir().join(pid.to_string()).join("signal-latest.json") +} + +pub fn record(input: CrashInput) -> i32 { + if !config::enabled() { + return if input.finalize { 1 } else { 0 }; + } + + let ctx = context::snapshot(); + let training = training_context(); + let message = truncate(&input.message, 2000); + let fp = fingerprint(&input.exception_type, &message, &input.top_frame); + + let event = CrashEvent { + event_id: new_event_id(), + timestamp_ns: SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_nanos() as i64) + .unwrap_or(0), + kind: input.kind, + rank: ctx.rank, + local_rank: ctx.local_rank, + world_size: ctx.world_size, + host: ctx.host, + pid: ctx.pid, + exception_type: input.exception_type, + message, + top_frame: input.top_frame, + traceback: truncate(&input.traceback, 64000), + native_backtrace: truncate(&input.native_backtrace, 64000), + crash_thread: input.crash_thread, + thread_stacks: truncate(&input.thread_stacks, 64000), + fingerprint: fp, + global_step: training.global_step, + local_step: training.local_step, + micro_step: training.micro_step, + training_phase: training.training_phase, + active_span: training.active_span, + last_comm_op: training.last_comm_op, + }; + + let spill_path = if config::spill_enabled() { + spill_event(&event) + } else { + None + }; + let grace_sec = if input.finalize && grace::should_run_grace() { + config::grace_sec() + } else { + 0 + }; + + if input.finalize { + report::print_report( + &event, + grace_sec, + config::force_hold(), + spill_path.as_deref(), + ); + } else { + report::print_summary(&event, spill_path.as_deref()); + } + + if !input.finalize { + return 1; + } + grace::grace_and_maybe_hold(1) +} + +struct TrainingContext { + global_step: i64, + local_step: i64, + micro_step: i64, + training_phase: String, + active_span: String, + last_comm_op: String, +} + +fn training_context() -> TrainingContext { + let step = crash_step_snapshot(); + let spans = crash_span_snapshot(); + let comm = LAST_COMM.lock().ok().and_then(|g| g.clone()); + TrainingContext { + global_step: i64::try_from(step.global_step).unwrap_or(-1), + local_step: i64::try_from(step.local_step).unwrap_or(-1), + micro_step: i64::try_from(step.micro_step).unwrap_or(-1), + training_phase: spans.training_phase.clone(), + active_span: spans.active_name.clone(), + last_comm_op: comm.unwrap_or_default(), + } +} + +fn crash_dir() -> PathBuf { + default_dir().join("crash") +} + +fn spill_event(event: &CrashEvent) -> Option { + let json = serde_json::to_string(event).ok()?; + let path = crash_dir() + .join(event.pid.to_string()) + .join(format!("{}.json", event.event_id)); + write_atomic(&path, json.as_bytes()).ok()?; + let _ = write_atomic( + &crash_dir().join(event.pid.to_string()).join("latest.json"), + json.as_bytes(), + ); + Some(path) +} + +fn write_atomic(path: &Path, data: &[u8]) -> std::io::Result<()> { + if let Some(parent) = path.parent() { + fs::create_dir_all(parent)?; + } + let tmp = path.with_extension("json.tmp"); + { + let mut file = fs::File::create(&tmp)?; + file.write_all(data)?; + file.sync_all()?; + } + fs::rename(tmp, path) +} + +fn new_event_id() -> String { + Uuid::new_v4().simple().to_string() +} + +fn truncate(text: &str, limit: usize) -> String { + if text.len() <= limit { + text.to_string() + } else { + format!("{}...", &text[..limit.saturating_sub(3)]) + } +} + +fn fingerprint(exception_type: &str, message: &str, top_frame: &str) -> String { + let msg = message.trim().replace('\n', " "); + let msg = if msg.len() <= 120 { + msg + } else { + msg[..120].to_string() + }; + let payload = format!("{exception_type}|{top_frame}|{msg}"); + let digest = Sha256::digest(payload.as_bytes()); + digest[..8].iter().map(|b| format!("{:02x}", b)).collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn fingerprint_stable() { + let a = fingerprint("ValueError", "bad", "train.py:10 in step"); + let b = fingerprint("ValueError", "bad", "train.py:10 in step"); + assert_eq!(a, b); + assert_eq!(a.len(), 16); + } + + #[test] + fn spill_writes_files() { + let event = CrashEvent { + event_id: new_event_id(), + timestamp_ns: 1, + kind: "python_exception".into(), + rank: 0, + local_rank: 0, + world_size: 1, + host: "localhost".into(), + pid: 4242, + exception_type: "RuntimeError".into(), + message: "boom".into(), + top_frame: "t.py:1".into(), + traceback: String::new(), + native_backtrace: String::new(), + crash_thread: String::new(), + thread_stacks: String::new(), + fingerprint: "abc".into(), + global_step: 0, + local_step: 0, + micro_step: 0, + training_phase: String::new(), + active_span: String::new(), + last_comm_op: String::new(), + }; + let path = spill_event(&event).expect("spill"); + assert!(path.exists()); + } +} diff --git a/probing/extensions/python/src/features/crash/mod.rs b/probing/extensions/python/src/features/crash/mod.rs new file mode 100644 index 00000000..55b466c3 --- /dev/null +++ b/probing/extensions/python/src/features/crash/mod.rs @@ -0,0 +1,122 @@ +//! Crash capture — stderr + spill + grace. +//! +//! Python ``excepthook`` only. Post-mortem: logs and ``default_dir()/crash//latest.json``. + +mod api; +mod grace; +mod handler; +mod report; +mod signal; + +pub use api::{ + crash_enabled, note_last_comm, record_crash, request_crash_hold, request_crash_release, +}; +pub use signal::install_crash_handler; + +pub(crate) mod config { + pub fn enabled() -> bool { + match std::env::var("PROBING_CRASH") { + Ok(val) => { + let lower = val.trim().to_ascii_lowercase(); + !matches!(lower.as_str(), "0" | "false" | "no" | "off") + } + Err(_) => true, + } + } + + pub fn grace_sec() -> u64 { + if env_truthy("PROBING_CRASH_NO_GRACE", false) { + return 0; + } + std::env::var("PROBING_CRASH_GRACE_SEC") + .ok() + .and_then(|v| v.trim().parse().ok()) + .unwrap_or(20) + } + + pub fn force_hold() -> bool { + env_truthy("PROBING_CRASH_HOLD", false) + } + + pub fn grace_all_ranks() -> bool { + env_truthy("PROBING_CRASH_GRACE_ALL_RANKS", false) + } + + pub fn spill_enabled() -> bool { + match std::env::var("PROBING_CRASH_SPILL") { + Ok(val) => { + let lower = val.trim().to_ascii_lowercase(); + !matches!(lower.as_str(), "0" | "false" | "no" | "off") + } + Err(_) => true, + } + } + + fn env_truthy(name: &str, default: bool) -> bool { + match std::env::var(name) { + Ok(val) => matches!( + val.trim().to_ascii_lowercase().as_str(), + "1" | "true" | "yes" | "on" + ), + Err(_) => default, + } + } +} + +pub(crate) mod context { + pub struct Snapshot { + pub rank: i32, + pub local_rank: i32, + pub world_size: i32, + pub host: String, + pub pid: i32, + } + + pub fn snapshot() -> Snapshot { + Snapshot { + rank: env_i32("RANK"), + local_rank: env_i32("LOCAL_RANK"), + world_size: env_i32("WORLD_SIZE"), + host: std::env::var("POD_IP") + .ok() + .filter(|s| !s.is_empty()) + .unwrap_or_else(|| { + std::env::var("HOSTNAME").unwrap_or_else(|_| "localhost".into()) + }), + pid: std::process::id() as i32, + } + } + + pub fn hold_file_path(pid: i32) -> String { + probing_memtable::discover::default_dir() + .join("crash") + .join(format!("hold-{pid}")) + .to_string_lossy() + .into() + } + + fn env_i32(key: &str) -> i32 { + std::env::var(key) + .ok() + .and_then(|v| v.trim().parse().ok()) + .unwrap_or(-1) + } +} + +pub fn handle_http( + path: &str, + _params: &std::collections::HashMap, + _body: &[u8], +) -> Result, String> { + match path { + "crash/hold" => { + grace::request_hold(); + Ok(r#"{"ok":true,"held":true}"#.into()) + } + "crash/release" => { + grace::request_release(true); + Ok(r#"{"ok":true,"released":true}"#.into()) + } + _ => Err(format!("unknown crash path: {path}")), + } +} diff --git a/probing/extensions/python/src/features/crash/report.rs b/probing/extensions/python/src/features/crash/report.rs new file mode 100644 index 00000000..0861d8b8 --- /dev/null +++ b/probing/extensions/python/src/features/crash/report.rs @@ -0,0 +1,215 @@ +//! Human-readable crash report rendering. + +use super::handler::CrashEvent; + +const WIDTH: usize = 72; + +fn rule(title: &str) -> String { + let head = format!("── {title} "); + let fill = WIDTH.saturating_sub(head.len()); + format!("{head}{}", "─".repeat(fill)) +} + +fn rank_label(event: &CrashEvent) -> String { + if event.world_size > 0 { + format!("rank {}/{}", event.rank, event.world_size) + } else if event.rank >= 0 { + format!("rank {}", event.rank) + } else { + "rank ?".into() + } +} + +fn step_line(event: &CrashEvent) -> String { + if event.global_step > 0 || event.local_step > 0 || event.micro_step > 0 { + let phase = if event.training_phase.is_empty() { + String::new() + } else { + format!(" phase={}", event.training_phase) + }; + format!( + " step global={} local={} micro={}{phase}", + event.global_step, event.local_step, event.micro_step + ) + } else { + " step (not in training loop)".into() + } +} + +fn context_lines(event: &CrashEvent) -> Vec { + let mut parts = vec![rank_label(event)]; + if event.local_rank >= 0 { + parts.push(format!("local_rank={}", event.local_rank)); + } + parts.push(format!("host={}", event.host)); + parts.push(format!("pid={}", event.pid)); + + let mut lines = vec![format!(" {}", parts.join(" "))]; + lines.push(step_line(event)); + if !event.active_span.is_empty() { + lines.push(format!(" span {}", event.active_span)); + } + if !event.last_comm_op.is_empty() { + lines.push(format!(" comm {}", event.last_comm_op)); + } + lines +} + +fn indent_block(text: &str, max_lines: usize) -> Vec { + let lines: Vec<&str> = text.trim_end().lines().collect(); + if lines.is_empty() { + return vec![]; + } + if lines.len() > max_lines { + let skip = lines.len() - max_lines; + let mut out = vec![format!(" … {skip} earlier frames omitted …")]; + for line in &lines[skip..] { + out.push(format!(" {line}")); + } + return out; + } + lines.iter().map(|line| format!(" {line}")).collect() +} + +fn stack_sections(event: &CrashEvent, all_threads_max: usize) -> Vec { + let mut lines = Vec::new(); + let python_hook = matches!(event.kind.as_str(), "python_exception" | "thread_exception"); + + if !python_hook && !event.traceback.is_empty() { + let crash_thread = if event.crash_thread.is_empty() { + "MainThread" + } else { + &event.crash_thread + }; + lines.push(format!( + " crash thread ({crash_thread}) — exception traceback" + )); + lines.extend(indent_block(&event.traceback, usize::MAX)); + } + + if !event.thread_stacks.is_empty() { + lines.push(" all threads — faulthandler snapshot".into()); + lines.extend(indent_block(&event.thread_stacks, all_threads_max)); + } else if !python_hook && !event.native_backtrace.is_empty() { + lines.push(" native backtrace".into()); + lines.extend(indent_block(&event.native_backtrace, 16)); + } + + lines +} + +fn format_event( + event: &CrashEvent, + fatal: bool, + grace_sec: u64, + hold_active: bool, + spill_path: Option<&std::path::Path>, +) -> String { + let title = if fatal { + format!("PROBING CRASH · {} · FATAL", event.kind) + } else { + format!("PROBING CRASH · {} · NON-FATAL", event.kind) + }; + let stacks_max = if fatal { 64 } else { 48 }; + let short_id = if event.event_id.len() > 16 { + &event.event_id[..16] + } else { + &event.event_id + }; + + let mut lines = vec![ + rule(&title), + format!(" exception {}: {}", event.exception_type, event.message), + format!(" at {}", event.top_frame), + ]; + lines.extend(context_lines(event)); + lines.push(format!( + " ids event={short_id} fingerprint={}", + event.fingerprint + )); + lines.extend(stack_sections(event, stacks_max)); + + if fatal { + if let Some(path) = spill_path { + lines.push(format!(" spill {}", path.display())); + } + if grace_sec > 0 { + lines.push(format!( + " hold {grace_sec}s grace — gdb -p {} · touch {} · kill -USR1", + event.pid, + super::context::hold_file_path(event.pid) + )); + } else if hold_active { + lines.push( + " hold active — kill -USR2 or POST /apis/pythonext/crash/release".into(), + ); + } + } else { + lines.push( + " status process continues — main thread did NOT exit; training may stall".into(), + ); + if let Some(path) = spill_path { + lines.push(format!(" spill {}", path.display())); + } + } + + lines.push("─".repeat(WIDTH)); + lines.join("\n") +} + +pub fn print_report( + event: &CrashEvent, + grace_sec: u64, + hold_active: bool, + spill_path: Option<&std::path::Path>, +) { + eprintln!( + "{}", + format_event(event, true, grace_sec, hold_active, spill_path) + ); +} + +pub fn print_summary(event: &CrashEvent, spill_path: Option<&std::path::Path>) { + eprintln!("{}", format_event(event, false, 0, false, spill_path)); +} + +#[cfg(test)] +mod tests { + use super::*; + + fn sample_event() -> CrashEvent { + CrashEvent { + event_id: "test-event-id".into(), + timestamp_ns: 1, + kind: "thread_exception".into(), + rank: 1, + local_rank: 1, + world_size: 4, + host: "node-a".into(), + pid: 99, + exception_type: "RuntimeError".into(), + message: "demo crash".into(), + top_frame: "demo.py:19 in boom".into(), + traceback: String::new(), + native_backtrace: String::new(), + crash_thread: "crash-r1".into(), + thread_stacks: "Thread 0x1 (MainThread)\n".into(), + fingerprint: "abc123".into(), + global_step: 42, + local_step: 42, + micro_step: 42, + training_phase: "backward".into(), + active_span: "train.step".into(), + last_comm_op: "all_reduce".into(), + } + } + + #[test] + fn summary_includes_training_context() { + let text = format_event(&sample_event(), false, 0, false, None); + assert!(text.contains("NON-FATAL")); + assert!(text.contains("global=42")); + assert!(text.contains("train.step")); + assert!(!text.contains("memtable")); + } +} diff --git a/probing/extensions/python/src/features/crash/signal.rs b/probing/extensions/python/src/features/crash/signal.rs new file mode 100644 index 00000000..d7757852 --- /dev/null +++ b/probing/extensions/python/src/features/crash/signal.rs @@ -0,0 +1,425 @@ +//! Fatal-signal backtrace dumper. +//! +//! On `SIGSEGV` / `SIGBUS` / `SIGABRT` / `SIGILL` / `SIGFPE` we print the +//! crashing thread's native backtrace to stderr, then restore the default +//! disposition and re-raise the signal so the process still produces a core +//! dump / exits with the usual status. This is purely a debugging aid for +//! diagnosing the profiler/native crashes. +//! +//! The handler runs on a dedicated `sigaltstack` (so a stack-overflow crash can +//! still be reported) and avoids heap allocation on its hot path: integers are +//! formatted into stack buffers and symbol names are written as their raw bytes +//! via `write(2)`. Symbolization (`backtrace::*_unsynchronized`) is not strictly +//! async-signal-safe, but this only ever runs once, while the process is already +//! dying; a recursive fault is caught by a guard that immediately re-raises with +//! the default handler. + +use core::ffi::{c_int, c_void}; +use std::sync::atomic::{AtomicBool, AtomicI32, AtomicUsize, Ordering}; + +use nix::libc; +use probing_core::trace::crash_atomic_step; + +use super::context; +use super::handler; + +static INSTALLED: AtomicBool = AtomicBool::new(false); +static IN_HANDLER: AtomicBool = AtomicBool::new(false); +static SIGNAL_SPILL_PATH_LEN: AtomicUsize = AtomicUsize::new(0); +static CACHED_RANK: AtomicI32 = AtomicI32::new(-1); +static CACHED_LOCAL_RANK: AtomicI32 = AtomicI32::new(-1); + +const SIGNAL_SPILL_PATH_CAP: usize = 384; +static mut SIGNAL_SPILL_PATH: [u8; SIGNAL_SPILL_PATH_CAP] = [0u8; SIGNAL_SPILL_PATH_CAP]; + +const FATAL_SIGNALS: [c_int; 5] = [ + libc::SIGSEGV, + libc::SIGBUS, + libc::SIGABRT, + libc::SIGILL, + libc::SIGFPE, +]; + +const MAX_FRAMES: usize = 256; +const ALT_STACK_SIZE: usize = 256 * 1024; + +static mut ALT_STACK: [u8; ALT_STACK_SIZE] = [0u8; ALT_STACK_SIZE]; + +fn sig_name(sig: c_int) -> &'static str { + match sig { + libc::SIGSEGV => "SIGSEGV", + libc::SIGBUS => "SIGBUS", + libc::SIGABRT => "SIGABRT", + libc::SIGILL => "SIGILL", + libc::SIGFPE => "SIGFPE", + _ => "SIGNAL", + } +} + +#[inline] +unsafe fn write_str(s: &str) { + let _ = libc::write(2, s.as_ptr() as *const c_void, s.len()); +} + +#[inline] +unsafe fn write_bytes(b: &[u8]) { + let _ = libc::write(2, b.as_ptr() as *const c_void, b.len()); +} + +/// Write `v` as lowercase hex (no `0x` prefix), no allocation. +unsafe fn write_hex(v: usize) { + if v == 0 { + write_str("0"); + return; + } + let mut buf = [0u8; 16]; + let mut i = buf.len(); + let mut val = v; + while val > 0 { + i -= 1; + let d = (val & 0xf) as u8; + buf[i] = if d < 10 { b'0' + d } else { b'a' + (d - 10) }; + val >>= 4; + } + write_bytes(&buf[i..]); +} + +/// Write `v` as decimal, no allocation. +unsafe fn write_dec(v: usize) { + if v == 0 { + write_str("0"); + return; + } + let mut buf = [0u8; 20]; + let mut i = buf.len(); + let mut val = v; + while val > 0 { + i -= 1; + buf[i] = b'0' + (val % 10) as u8; + val /= 10; + } + write_bytes(&buf[i..]); +} + +unsafe fn current_tid() -> u64 { + #[cfg(target_os = "linux")] + { + libc::syscall(libc::SYS_gettid) as u64 + } + #[cfg(target_os = "macos")] + { + let mut t: u64 = 0; + libc::pthread_threadid_np(0, &mut t); + t + } + #[cfg(not(any(target_os = "linux", target_os = "macos")))] + { + 0 + } +} + +unsafe fn restore_and_reraise(sig: c_int) { + let mut sa: libc::sigaction = std::mem::zeroed(); + sa.sa_sigaction = libc::SIG_DFL; + libc::sigemptyset(&mut sa.sa_mask); + sa.sa_flags = 0; + libc::sigaction(sig, &sa, std::ptr::null_mut()); + libc::raise(sig); +} + +unsafe extern "C" fn crash_handler(sig: c_int, info: *mut libc::siginfo_t, _uctx: *mut c_void) { + // A fault *inside* the dump (e.g. while symbolizing) must not loop forever. + if IN_HANDLER.swap(true, Ordering::SeqCst) { + restore_and_reraise(sig); + return; + } + + write_str("\n==== probing: fatal signal "); + write_str(sig_name(sig)); + write_str(" ("); + write_dec(sig as usize); + write_str(") on thread "); + write_dec(current_tid() as usize); + write_str(" ====\n"); + + if !info.is_null() { + let addr = (*info).si_addr() as usize; + write_str("fault address: 0x"); + write_hex(addr); + write_str("\n"); + } + + write_str("native backtrace (crashing thread):\n"); + + let mut idx = 0usize; + backtrace::trace_unsynchronized(|frame| { + let ip = frame.ip() as usize; + write_str(" #"); + write_dec(idx); + write_str(" 0x"); + write_hex(ip); + write_str(" "); + + let mut wrote_name = false; + backtrace::resolve_frame_unsynchronized(frame, |symbol| { + if !wrote_name { + if let Some(name) = symbol.name() { + if let Some(s) = name.as_str() { + write_bytes(s.as_bytes()); + wrote_name = true; + } + } + } + }); + if !wrote_name { + write_str(""); + } + write_str("\n"); + + idx += 1; + idx < MAX_FRAMES + }); + + write_str("==== end probing backtrace; re-raising "); + write_str(sig_name(sig)); + write_str(" ====\n"); + + spill_fatal_signal( + sig, + current_tid(), + if info.is_null() { + 0 + } else { + (*info).si_addr() as usize + }, + ); + + crash_grace_wait(); + + restore_and_reraise(sig); +} + +unsafe fn env_is_truthy(name: &str) -> bool { + let Ok(val) = std::env::var(name) else { + return false; + }; + matches!( + val.trim().to_ascii_lowercase().as_str(), + "1" | "true" | "yes" | "on" + ) +} + +unsafe fn env_grace_sec() -> u64 { + if env_is_truthy("PROBING_CRASH_NO_GRACE") { + return 0; + } + match std::env::var("PROBING_CRASH_GRACE_SEC") { + Ok(val) => val.trim().parse().unwrap_or(20), + Err(_) => 20, + } +} + +unsafe fn env_rank(name: &str) -> i32 { + std::env::var(name) + .ok() + .and_then(|v| v.trim().parse().ok()) + .unwrap_or(-1) +} + +unsafe fn should_signal_grace() -> bool { + if env_is_truthy("PROBING_CRASH_HOLD") { + return true; + } + let grace = env_grace_sec(); + if grace == 0 { + return false; + } + if env_is_truthy("PROBING_CRASH_GRACE_ALL_RANKS") { + return true; + } + env_rank("RANK") == 0 || env_rank("LOCAL_RANK") == 0 +} + +unsafe fn crash_grace_wait() { + if !should_signal_grace() { + return; + } + let pid = libc::getpid() as usize; + write_str("\n[probing crash] fatal signal grace — attach: gdb -p "); + write_dec(pid); + write_str("\n"); + + let grace = if env_is_truthy("PROBING_CRASH_HOLD") { + 0 + } else { + env_grace_sec() + }; + + let deadline = if grace == 0 { + i64::MAX + } else { + libc::time(std::ptr::null_mut()) + grace as i64 + }; + loop { + if env_is_truthy("PROBING_CRASH_HOLD") { + libc::usleep(200_000); + continue; + } + let now = libc::time(std::ptr::null_mut()); + if now >= deadline { + return; + } + libc::usleep(250_000); + } +} + +fn prepare_signal_spill() { + let ctx = context::snapshot(); + CACHED_RANK.store(ctx.rank, Ordering::Relaxed); + CACHED_LOCAL_RANK.store(ctx.local_rank, Ordering::Relaxed); + + if let Some(parent) = handler::signal_spill_path(ctx.pid).parent() { + let _ = std::fs::create_dir_all(parent); + } + let path = handler::signal_spill_path(ctx.pid); + let bytes = path.as_os_str().as_encoded_bytes(); + unsafe { + if bytes.len() + 1 >= SIGNAL_SPILL_PATH_CAP { + return; + } + let p = std::ptr::addr_of_mut!(SIGNAL_SPILL_PATH).cast::(); + std::ptr::copy_nonoverlapping(bytes.as_ptr(), p, bytes.len()); + p.add(bytes.len()).write(0); + SIGNAL_SPILL_PATH_LEN.store(bytes.len(), Ordering::Relaxed); + } +} + +unsafe fn spill_fatal_signal(sig: c_int, tid: u64, fault_addr: usize) { + let path_len = SIGNAL_SPILL_PATH_LEN.load(Ordering::Relaxed); + if path_len == 0 { + return; + } + + let step = crash_atomic_step(); + let rank = CACHED_RANK.load(Ordering::Relaxed); + let local_rank = CACHED_LOCAL_RANK.load(Ordering::Relaxed); + let pid = libc::getpid(); + + let mut buf = [0u8; 512]; + let mut n = 0usize; + macro_rules! append { + ($s:expr) => {{ + let b = $s.as_bytes(); + if n + b.len() < buf.len() { + buf[n..n + b.len()].copy_from_slice(b); + n += b.len(); + } + }}; + } + + append!("{\"kind\":\"fatal_signal\",\"signal\":\""); + append!(sig_name(sig)); + append!("\",\"pid\":"); + n = append_dec(&mut buf, n, pid as usize); + append!(",\"rank\":"); + n = append_dec(&mut buf, n, rank.max(0) as usize); + append!(",\"local_rank\":"); + n = append_dec(&mut buf, n, local_rank.max(0) as usize); + append!(",\"global_step\":"); + n = append_dec(&mut buf, n, step.global_step as usize); + append!(",\"micro_step\":"); + n = append_dec(&mut buf, n, step.micro_step as usize); + append!(",\"tid\":"); + n = append_dec(&mut buf, n, tid as usize); + append!(",\"fault_addr\":\"0x"); + n = append_hex(&mut buf, n, fault_addr); + append!("\"}\n"); + + let path_ptr = std::ptr::addr_of!(SIGNAL_SPILL_PATH).cast::(); + let fd = libc::open( + path_ptr, + libc::O_CREAT | libc::O_TRUNC | libc::O_WRONLY, + 0o644, + ); + if fd >= 0 { + let _ = libc::write(fd, buf.as_ptr() as *const c_void, n); + let _ = libc::fsync(fd); + let _ = libc::close(fd); + } +} + +unsafe fn append_dec(buf: &mut [u8], mut n: usize, val: usize) -> usize { + if val == 0 { + if n < buf.len() { + buf[n] = b'0'; + n += 1; + } + return n; + } + let mut tmp = [0u8; 20]; + let mut i = tmp.len(); + let mut v = val; + while v > 0 { + i -= 1; + tmp[i] = b'0' + (v % 10) as u8; + v /= 10; + } + let digits = &tmp[i..]; + if n + digits.len() < buf.len() { + buf[n..n + digits.len()].copy_from_slice(digits); + n += digits.len(); + } + n +} + +unsafe fn append_hex(buf: &mut [u8], mut n: usize, val: usize) -> usize { + if val == 0 { + if n < buf.len() { + buf[n] = b'0'; + n += 1; + } + return n; + } + let mut tmp = [0u8; 16]; + let mut i = tmp.len(); + let mut v = val; + while v > 0 { + i -= 1; + let d = (v & 0xf) as u8; + tmp[i] = if d < 10 { b'0' + d } else { b'a' + (d - 10) }; + v >>= 4; + } + let digits = &tmp[i..]; + if n + digits.len() < buf.len() { + buf[n..n + digits.len()].copy_from_slice(digits); + n += digits.len(); + } + n +} + +/// Install backtrace-on-crash handlers for the common fatal signals. Idempotent. +pub fn install_crash_handler() { + if INSTALLED.swap(true, Ordering::SeqCst) { + return; + } + + prepare_signal_spill(); + + unsafe { + let mut ss: libc::stack_t = std::mem::zeroed(); + ss.ss_sp = core::ptr::addr_of_mut!(ALT_STACK) as *mut c_void; + ss.ss_size = ALT_STACK_SIZE; + ss.ss_flags = 0; + libc::sigaltstack(&ss, std::ptr::null_mut()); + + for &sig in FATAL_SIGNALS.iter() { + let mut sa: libc::sigaction = std::mem::zeroed(); + sa.sa_sigaction = crash_handler as *const () as usize; + sa.sa_flags = libc::SA_SIGINFO | libc::SA_ONSTACK; + libc::sigemptyset(&mut sa.sa_mask); + libc::sigaction(sig, &sa, std::ptr::null_mut()); + } + } + + log::info!("probing: crash backtrace handler installed"); +} diff --git a/probing/extensions/python/src/features/tracing.rs b/probing/extensions/python/src/features/tracing.rs index feb36397..1c055642 100644 --- a/probing/extensions/python/src/features/tracing.rs +++ b/probing/extensions/python/src/features/tracing.rs @@ -2,7 +2,7 @@ use pyo3::prelude::*; use pyo3::types::{PyDict, PyList, PyModule}; use pyo3::IntoPyObjectExt; use std::cell::RefCell; -use std::sync::{Arc, Mutex}; +use std::sync::{Arc, Mutex, MutexGuard}; use probing_core::trace::Span as RawSpan; use probing_core::trace::{ @@ -12,7 +12,15 @@ use probing_core::trace::{ use crate::features::convert::{ele_to_python, python_to_ele}; -const SPAN_LOCK_POISONED: &str = "Failed to acquire lock on span (lock poisoned)"; +fn lock_span(m: &Mutex) -> MutexGuard<'_, RawSpan> { + match m.lock() { + Ok(guard) => guard, + Err(e) => { + log::debug!("span mutex poisoned; recovering"); + e.into_inner() + } + } +} fn parse_py_attributes(py: Python, attributes: Vec>) -> PyResult> { let mut converted = Vec::new(); @@ -57,6 +65,83 @@ where // Thread-local storage for span context thread_local! { static SPAN_STACK: RefCell>> = const { RefCell::new(Vec::new()) }; + static SPAN_HINT_STACK: RefCell> = const { RefCell::new(Vec::new()) }; + static CRASH_SPAN_CAPTURE: RefCell> = const { RefCell::new(None) }; +} + +#[derive(Clone, Debug, Default)] +struct SpanHint { + trace_id: u64, + span_id: u64, + name: String, + phase: Option, +} + +#[derive(Clone, Debug, Default)] +pub struct CrashSpanSnapshot { + pub active_name: String, + pub trace_id: u64, + pub span_id: u64, + pub training_phase: String, + pub stack: Vec, +} + +pub fn crash_span_snapshot() -> CrashSpanSnapshot { + if let Some(captured) = CRASH_SPAN_CAPTURE.with(|cap| cap.borrow_mut().take()) { + return captured; + } + snapshot_from_hint_stack() +} + +fn snapshot_from_hint_stack() -> CrashSpanSnapshot { + SPAN_HINT_STACK.with(|stack| { + let stack = stack.borrow(); + if stack.is_empty() { + return CrashSpanSnapshot::default(); + } + let names: Vec = stack.iter().map(|h| h.name.clone()).collect(); + let active = stack.last().cloned().unwrap_or_default(); + let training_phase = stack + .iter() + .rev() + .filter_map(|h| h.phase.as_deref()) + .find(|p| matches!(*p, "forward" | "backward" | "optimizer")) + .unwrap_or("") + .to_string(); + CrashSpanSnapshot { + active_name: active.name, + trace_id: active.trace_id, + span_id: active.span_id, + training_phase, + stack: names, + } + }) +} + +fn capture_span_snapshot_for_crash() { + let snap = snapshot_from_hint_stack(); + if !snap.active_name.is_empty() { + CRASH_SPAN_CAPTURE.with(|cap| *cap.borrow_mut() = Some(snap)); + } +} + +fn push_span_hint(span: &Span) { + let hint = SpanHint { + trace_id: span.with_inner(|s| s.trace_id), + span_id: span.with_inner(|s| s.span_id), + name: span.with_inner(|s| s.name.clone()), + phase: span.with_inner(|s| s.phase.clone()), + }; + SPAN_HINT_STACK.with(|stack| stack.borrow_mut().push(hint)); +} + +fn pop_span_hint(span_id: u64) { + SPAN_HINT_STACK.with(|stack| { + let mut stack = stack.borrow_mut(); + if let Some(pos) = stack.iter().rposition(|h| h.span_id == span_id) { + stack.remove(pos); + } + }); } /// Python binding for Span @@ -68,11 +153,11 @@ pub struct Span { impl Span { fn with_inner(&self, f: impl FnOnce(&RawSpan) -> R) -> R { - f(&self.inner.lock().expect(SPAN_LOCK_POISONED)) + f(&lock_span(&self.inner)) } fn with_inner_mut(&mut self, f: impl FnOnce(&mut RawSpan) -> R) -> R { - f(&mut self.inner.lock().expect(SPAN_LOCK_POISONED)) + f(&mut lock_span(&self.inner)) } } @@ -238,7 +323,7 @@ impl Span { /// Gets all events as a list. fn get_events(&self, py: Python) -> PyResult> { let list = PyList::empty(py); - let inner = self.inner.lock().expect(SPAN_LOCK_POISONED); + let inner = lock_span(&self.inner); for event in &inner.events { let event_dict = PyDict::new(py); event_dict.set_item("name", &event.name)?; @@ -287,6 +372,7 @@ impl Span { /// Context manager entry (for `with` statement support). fn __enter__(slf: PyRef) -> PyResult> { let py = slf.py(); + push_span_hint(&slf); let span_obj: Py = Py::new(py, slf.clone())?.into(); SPAN_STACK.with(|stack| { stack.borrow_mut().push(span_obj); @@ -297,12 +383,16 @@ impl Span { /// Context manager exit (for `with` statement support). fn __exit__( slf: PyRef, - _exc_type: Option<&Bound<'_, PyAny>>, + exc_type: Option<&Bound<'_, PyAny>>, _exc_val: Option<&Bound<'_, PyAny>>, _exc_tb: Option<&Bound<'_, PyAny>>, ) -> PyResult { + if exc_type.is_some() && !exc_type.unwrap().is_none() { + capture_span_snapshot_for_crash(); + } let self_id = slf.span_id(); - slf.inner.lock().expect(SPAN_LOCK_POISONED).end(); + lock_span(&slf.inner).end(); + pop_span_hint(self_id); SPAN_STACK.with(|stack| { let mut stack = stack.borrow_mut(); diff --git a/probing/extensions/python/src/python.rs b/probing/extensions/python/src/python.rs index 10a558cb..4fc233d8 100644 --- a/probing/extensions/python/src/python.rs +++ b/probing/extensions/python/src/python.rs @@ -1,17 +1,9 @@ use std::ffi::CString; use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::Mutex; -use anyhow::Result; -use once_cell::sync::Lazy; use pyo3::prelude::*; -use pyo3::types::PyTuple; use pyo3::{types::PyDict, Bound, Python}; -use crate::pycode::get_code; - -pub static CRASH_HANDLER: Lazy>> = Lazy::new(|| Mutex::new(None)); -pub static OLD_HANDLER: Lazy>> = Lazy::new(|| None); pub static PROBING_ENABLED: AtomicBool = AtomicBool::new(false); fn run_embedded( @@ -26,45 +18,6 @@ fn run_embedded( py.run(&code, globals, locals) } -fn call_default_handler(typ: Py, value: Py, traceback: Py) -> Result<()> { - let code = get_code("crash_handler.py").unwrap_or_default(); - Python::attach(|py| -> Result<()> { - let global = PyDict::new(py); - run_embedded(py, &code, Some(&global), None)?; - if let Some(handler) = global.get_item("crash_handler")? { - let args = PyTuple::new(py, [typ, value, traceback])?; - handler.call(args, None)?; - } - Ok(()) - }) -} - -fn call_custom_handler( - handler: &str, - typ: Py, - value: Py, - traceback: Py, -) -> Result<()> { - Python::attach(|py| -> Result<()> { - let locals = PyDict::new(py); - if handler.contains('.') { - let parts: Vec<&str> = handler.split('.').collect(); - let pkg = py.import(parts[0])?; - locals.set_item(parts[0], pkg)?; - } - locals.set_item("type", typ)?; - locals.set_item("value", value)?; - locals.set_item("traceback", traceback)?; - let ret = (|| { - let expr = CString::new(handler)?; - py.eval(&expr, None, Some(&locals)) - })(); - - println!("crash handler: {ret:?}"); - Ok(()) - }) -} - fn script_basename(py: Python) -> Option { let sys = py.import("sys").ok()?; let argv = sys.getattr("argv").ok()?; @@ -125,27 +78,12 @@ pub fn set_enabled(enabled: bool) { PROBING_ENABLED.store(enabled, Ordering::SeqCst); } -#[pyfunction] -pub fn crash_handler(typ: Py, value: Py, traceback: Py) { - let handler = CRASH_HANDLER.lock().unwrap().clone(); - log::debug!("call crash handler: {handler:?}"); - if let Some(handler) = handler { - let ret = match handler.as_str() { - "default" => call_default_handler(typ, value, traceback), - handler => call_custom_handler(handler, typ, value, traceback), - }; - if let Err(err) = ret { - log::error!("error calling crash handler: {err}"); - } - } -} - +/// Install the Python crash handler module (``probing.crash.install``). pub fn enable_crash_handler() -> anyhow::Result<()> { Python::attach(|py| -> anyhow::Result<()> { - log::debug!("enable crash handler"); - let sys = py.import("sys")?; - let func = wrap_pyfunction!(crash_handler, &sys)?; - sys.setattr("excepthook", func)?; + log::debug!("enable crash handler via probing.crash"); + let crash = py.import("probing.crash")?; + crash.call_method0("install")?; Ok(()) })?; Ok(()) @@ -164,7 +102,7 @@ pub fn enable_monitoring(filename: &str) -> anyhow::Result<()> { filename }; - let code = get_code(filename).ok_or_else(|| { + let code = crate::pycode::get_code(filename).ok_or_else(|| { anyhow::anyhow!( "monitoring script not found: {filename} (embed under pycode/ or set PROBING_CODE_ROOT)" ) diff --git a/probing/memtable/Cargo.toml b/probing/memtable/Cargo.toml index 7a1801fa..d8f874cc 100644 --- a/probing/memtable/Cargo.toml +++ b/probing/memtable/Cargo.toml @@ -8,6 +8,7 @@ description = "Self-describing columnar memory table with ring buffer" [dependencies] +log = "0.4" xxhash-rust = { version = "0.8", features = ["xxh3"] } memmap2 = "0.9" libc = "0.2" diff --git a/probing/memtable/src/memtable.rs b/probing/memtable/src/memtable.rs index d79077f8..d24c51b3 100644 --- a/probing/memtable/src/memtable.rs +++ b/probing/memtable/src/memtable.rs @@ -641,10 +641,10 @@ impl MemTable { /// MEMT is single-writer: the `&mut self` borrow guarantees exclusive /// access, so no lock is taken. pub fn push_row(&mut self, values: &[Value]) { - assert!( - validate_row_schema(self.backing.bytes(), values), - "value types do not match schema" - ); + if !validate_row_schema(self.backing.bytes(), values) { + log::warn!("push_row: value types do not match schema"); + return; + } self.push_row_unchecked(values); } pub fn push_row_unchecked(&mut self, values: &[Value]) { @@ -1683,19 +1683,19 @@ mod tests { } #[test] - #[should_panic(expected = "value types do not match schema")] fn push_row_rejects_wrong_column_count() { let schema = Schema::new().col("a", DType::U32).col("b", DType::I64); let mut t = MemTable::new(&schema, 256, 2); t.push_row(&[Value::U32(1)]); // only 1 value for 2 columns + assert_eq!(t.num_rows(0), 0); } #[test] - #[should_panic(expected = "value types do not match schema")] fn push_row_rejects_wrong_dtype() { let schema = Schema::new().col("a", DType::U32); let mut t = MemTable::new(&schema, 256, 2); t.push_row(&[Value::Str("oops")]); // Str instead of U32 + assert_eq!(t.num_rows(0), 0); } // ── MemTableWriter tests ────────────────────────── diff --git a/probing/memtable/src/row.rs b/probing/memtable/src/row.rs index afc4a0f8..4e542f44 100644 --- a/probing/memtable/src/row.rs +++ b/probing/memtable/src/row.rs @@ -10,11 +10,77 @@ pub(crate) fn panic_stale(context: &str) -> ! { panic!("stale read: chunk recycled ({context})") } -fn var_field_size(buf: &[u8], off: usize) -> usize { - if off + 4 > buf.len() { - panic_stale("var_field_size"); +fn read_i32(data: &[u8], off: usize) -> i32 { + if off + 4 > data.len() { + panic_stale("read_i32"); + } + i32::from_le_bytes([data[off], data[off + 1], data[off + 2], data[off + 3]]) +} + +fn read_u32(data: &[u8], off: usize) -> u32 { + if off + 4 > data.len() { + panic_stale("read_u32"); } - let raw = i32::from_le_bytes(buf[off..off + 4].try_into().unwrap()); + u32::from_le_bytes([data[off], data[off + 1], data[off + 2], data[off + 3]]) +} + +fn read_i64(data: &[u8], off: usize) -> i64 { + if off + 8 > data.len() { + panic_stale("read_i64"); + } + i64::from_le_bytes([ + data[off], + data[off + 1], + data[off + 2], + data[off + 3], + data[off + 4], + data[off + 5], + data[off + 6], + data[off + 7], + ]) +} + +fn read_u64(data: &[u8], off: usize) -> u64 { + if off + 8 > data.len() { + panic_stale("read_u64"); + } + u64::from_le_bytes([ + data[off], + data[off + 1], + data[off + 2], + data[off + 3], + data[off + 4], + data[off + 5], + data[off + 6], + data[off + 7], + ]) +} + +fn read_f32(data: &[u8], off: usize) -> f32 { + if off + 4 > data.len() { + panic_stale("read_f32"); + } + f32::from_le_bytes([data[off], data[off + 1], data[off + 2], data[off + 3]]) +} + +fn read_f64(data: &[u8], off: usize) -> f64 { + if off + 8 > data.len() { + panic_stale("read_f64"); + } + f64::from_le_bytes([ + data[off], + data[off + 1], + data[off + 2], + data[off + 3], + data[off + 4], + data[off + 5], + data[off + 6], + data[off + 7], + ]) +} + +fn var_field_size(buf: &[u8], off: usize) -> usize { + let raw = read_i32(buf, off); if raw < 0 { 4 } else { @@ -23,10 +89,7 @@ fn var_field_size(buf: &[u8], off: usize) -> usize { } fn resolve_var(buf: &[u8], off: usize, chunk_start: usize) -> &[u8] { - if off + 4 > buf.len() { - panic_stale("resolve_var offset"); - } - let raw = i32::from_le_bytes(buf[off..off + 4].try_into().unwrap()); + let raw = read_i32(buf, off); if raw < 0 { let ref_off = chunk_start + (-raw) as usize; if ref_off + 4 > buf.len() { @@ -100,31 +163,29 @@ impl<'a> Row<'a> { } pub fn col_u8(&self, col: usize) -> u8 { - self.data[self.col_offset(col)] + let off = self.col_offset(col); + if off >= self.data.len() { + panic_stale("col_u8"); + } + self.data[off] } pub fn col_u32(&self, col: usize) -> u32 { - let off = self.col_offset(col); - u32::from_le_bytes(self.data[off..off + 4].try_into().unwrap()) + read_u32(self.data, self.col_offset(col)) } pub fn col_i32(&self, col: usize) -> i32 { - let off = self.col_offset(col); - i32::from_le_bytes(self.data[off..off + 4].try_into().unwrap()) + read_i32(self.data, self.col_offset(col)) } pub fn col_i64(&self, col: usize) -> i64 { - let off = self.col_offset(col); - i64::from_le_bytes(self.data[off..off + 8].try_into().unwrap()) + read_i64(self.data, self.col_offset(col)) } pub fn col_f32(&self, col: usize) -> f32 { - let off = self.col_offset(col); - f32::from_le_bytes(self.data[off..off + 4].try_into().unwrap()) + read_f32(self.data, self.col_offset(col)) } pub fn col_f64(&self, col: usize) -> f64 { - let off = self.col_offset(col); - f64::from_le_bytes(self.data[off..off + 8].try_into().unwrap()) + read_f64(self.data, self.col_offset(col)) } pub fn col_u64(&self, col: usize) -> u64 { - let off = self.col_offset(col); - u64::from_le_bytes(self.data[off..off + 8].try_into().unwrap()) + read_u64(self.data, self.col_offset(col)) } pub fn col_str(&self, col: usize) -> &str { let b = self.resolve_var_col(col); @@ -175,7 +236,12 @@ impl<'a> RowCursor<'a> { } fn read_fixed(&mut self) -> [u8; N] { - let v: [u8; N] = self.data[self.pos..self.pos + N].try_into().unwrap(); + let end = self.pos + N; + if end > self.data.len() { + panic_stale("read_fixed"); + } + let mut v = [0u8; N]; + v.copy_from_slice(&self.data[self.pos..end]); self.pos += N; v } diff --git a/probing/memtable/src/writer.rs b/probing/memtable/src/writer.rs index 468979d0..d2f2b5f0 100644 --- a/probing/memtable/src/writer.rs +++ b/probing/memtable/src/writer.rs @@ -57,18 +57,19 @@ impl<'a> RowWriter<'a> { fn write_str_dedup(&mut self, data: &[u8]) { if !self.overflow { - if let Some(off) = self.dedup.as_ref().unwrap().lookup(self.col_idx, data) { - self.write_raw(&(-(off as i32)).to_le_bytes()); - return; + if let Some(dedup) = self.dedup.as_ref() { + if let Some(off) = dedup.lookup(self.col_idx, data) { + self.write_raw(&(-(off as i32)).to_le_bytes()); + return; + } } } let chunk_off = self.pos - self.chunk_start; self.write_lp(data); if !self.overflow { - self.dedup - .as_mut() - .unwrap() - .insert(self.col_idx, data, chunk_off); + if let Some(dedup) = self.dedup.as_mut() { + dedup.insert(self.col_idx, data, chunk_off); + } } } diff --git a/python/probing/__init__.py b/python/probing/__init__.py index d5d3db26..455602bc 100644 --- a/python/probing/__init__.py +++ b/python/probing/__init__.py @@ -79,6 +79,16 @@ def is_enabled(): except Exception: pass + try: + from probing._entrypoint import should_activate_probing + + if should_activate_probing(): + from probing.crash import install + + install() + except Exception: + pass + __all__ = [ "VERSION", "ExternalTable", diff --git a/python/probing/_entrypoint.py b/python/probing/_entrypoint.py index cd0789b7..afff7a2c 100644 --- a/python/probing/_entrypoint.py +++ b/python/probing/_entrypoint.py @@ -34,6 +34,33 @@ def is_probing_cli() -> bool: _LIGHTWEIGHT_MODULES = ("probing.nccl", "probing.skills", "probing.dev_pth") +def should_activate_probing() -> bool: + """True when ``PROBING`` / ``PROBING_ORIGINAL`` targets this process.""" + raw = os.environ.get("PROBING_ORIGINAL") or os.environ.get("PROBING", "0") + token = raw.strip().lower() + if token in ("0", "", "false", "no", "off"): + return False + if token in ("1", "followed", "2", "nested"): + return True + if raw.lower().startswith("regex:"): + import re + + pattern = raw.split(":", 1)[1] + candidates: list[str] = [c for c in sys.argv if c] + candidates.append(current_script_name()) + try: + import __main__ + + main_file = getattr(__main__, "__file__", None) + if main_file: + candidates.append(main_file) + candidates.append(os.path.basename(main_file)) + except Exception: + pass + return any(re.search(pattern, c) for c in candidates if c) + return raw == current_script_name() + + def is_lightweight_module() -> bool: """``python -m probing.nccl|skills|dev_pth`` must not start the engine.""" helper_suffixes = ( diff --git a/python/probing/crash.py b/python/probing/crash.py new file mode 100644 index 00000000..7a005a37 --- /dev/null +++ b/python/probing/crash.py @@ -0,0 +1,169 @@ +"""Python crash hooks — exception capture only; Rust handles spill and reporting.""" + +from __future__ import annotations + +import faulthandler +import os +import sys +import tempfile +import threading +import traceback +from types import TracebackType +from typing import Callable, Optional, Type + +from probing._core import crash_enabled, record_crash + +_INSTALLED = False +_PREV_EXCEPTHOOK: Optional[Callable] = None +_PREV_THREAD_EXCEPTHOOK: Optional[Callable] = None + + +def _call_original_excepthook( + exc_type: Type[BaseException], + exc_value: BaseException, + exc_tb: Optional[TracebackType], +) -> None: + hook = _PREV_EXCEPTHOOK + if hook is not None and hook is not _probing_excepthook: + hook(exc_type, exc_value, exc_tb) + else: + sys.__excepthook__(exc_type, exc_value, exc_tb) + + +def _call_original_thread_excepthook(args: threading.ExceptHookArgs) -> None: + hook = _PREV_THREAD_EXCEPTHOOK + if hook is not None and hook is not _probing_thread_excepthook: + hook(args) + elif hasattr(threading, "__excepthook__"): + threading.__excepthook__(args) + + +def _top_frame( + exc_type: Type[BaseException], + exc_value: BaseException, + exc_tb: Optional[TracebackType], +) -> str: + if exc_tb is not None: + frames = traceback.extract_tb(exc_tb) + if frames: + last = frames[-1] + return f"{last.filename}:{last.lineno} in {last.name}" + return f"{getattr(exc_type, '__name__', type(exc_value).__name__)}:" + + +def _capture_thread_stacks() -> str: + fd, path = tempfile.mkstemp(suffix=".probing-tb") + try: + with os.fdopen(fd, "w") as out: + faulthandler.dump_traceback(file=out, all_threads=True) + with open(path, encoding="utf-8") as inp: + return inp.read() + except Exception: + return "" + finally: + try: + os.unlink(path) + except OSError: + pass + + +def _dispatch( + *, + kind: str, + exc_type: Type[BaseException], + exc_value: BaseException, + exc_tb: Optional[TracebackType], + native_backtrace: str = "", + exit_code: int = 1, + finalize: bool = True, + crash_thread: str = "", +) -> int: + exc_name = getattr(exc_type, "__name__", type(exc_value).__name__) + tb_text = "".join(traceback.format_exception(exc_type, exc_value, exc_tb)) + top = _top_frame(exc_type, exc_value, exc_tb) + if not crash_thread: + try: + crash_thread = threading.current_thread().name or "MainThread" + except Exception: + crash_thread = "" + thread_stacks = _capture_thread_stacks() + + code = record_crash( + kind, + exc_name, + str(exc_value), + tb_text, + top, + native_backtrace, + finalize, + thread_stacks, + crash_thread, + ) + + if not finalize: + _ensure_hooks() + + return code if finalize else exit_code + + +def _probing_excepthook( + exc_type: Type[BaseException], + exc_value: BaseException, + exc_tb: Optional[TracebackType], +) -> None: + if exc_type is KeyboardInterrupt: + _call_original_excepthook(exc_type, exc_value, exc_tb) + return + _call_original_excepthook(exc_type, exc_value, exc_tb) + code = _dispatch( + kind="python_exception", + exc_type=exc_type, + exc_value=exc_value, + exc_tb=exc_tb, + exit_code=1, + finalize=True, + ) + sys.exit(code) + + +def _probing_thread_excepthook(args: threading.ExceptHookArgs) -> None: + _call_original_thread_excepthook(args) + thread = args.thread + _dispatch( + kind="thread_exception", + exc_type=args.exc_type, + exc_value=args.exc_value, + exc_tb=args.exc_traceback, + exit_code=1, + finalize=False, + crash_thread=thread.name if thread is not None else "", + ) + + +def install() -> None: + global _INSTALLED, _PREV_EXCEPTHOOK, _PREV_THREAD_EXCEPTHOOK + if not crash_enabled(): + return + if not _INSTALLED: + _INSTALLED = True + try: + faulthandler.enable(all_threads=True, file=sys.stderr) + except Exception: + pass + _PREV_EXCEPTHOOK = sys.excepthook + if hasattr(threading, "excepthook"): + _PREV_THREAD_EXCEPTHOOK = threading.excepthook + _ensure_hooks() + + +def _ensure_hooks() -> None: + if sys.excepthook is not _probing_excepthook: + sys.excepthook = _probing_excepthook + if ( + hasattr(threading, "excepthook") + and threading.excepthook is not _probing_thread_excepthook + ): + threading.excepthook = _probing_thread_excepthook + + +__all__ = ["install"] diff --git a/python/probing/ext/torch.py b/python/probing/ext/torch.py index 445c08bf..9ca0509e 100644 --- a/python/probing/ext/torch.py +++ b/python/probing/ext/torch.py @@ -99,6 +99,12 @@ def init(): _patch_dist_init_process_group() collective_hook() + try: + from probing.crash import install + + install() + except Exception: + pass def deinit(): diff --git a/python/probing/handlers/pythonext.py b/python/probing/handlers/pythonext.py index 9cbcb1f5..49a27d24 100644 --- a/python/probing/handlers/pythonext.py +++ b/python/probing/handlers/pythonext.py @@ -550,6 +550,24 @@ def get_magics_list() -> str: return json.dumps({"error": str(e), "traceback": traceback.format_exc()}) +@ext_handler("pythonext", "crash/hold") +def crash_hold() -> str: + """Extend grace hold for post-crash debugging.""" + import probing._core as core + + core.request_crash_hold() + return '{"ok":true,"held":true}' + + +@ext_handler("pythonext", "crash/release") +def crash_release() -> str: + """Release grace hold and exit the process.""" + import probing._core as core + + core.request_crash_release(True) + return '{"ok":true,"released":true}' + + @ext_handler("pythonext", "trace/variables") def get_trace_variables(function: Optional[str] = None, limit: int = 100) -> str: """Get trace variables from database. diff --git a/python/probing/profiling/collective/record.py b/python/probing/profiling/collective/record.py index 68bf7c69..60b04b2d 100644 --- a/python/probing/profiling/collective/record.py +++ b/python/probing/profiling/collective/record.py @@ -109,6 +109,17 @@ def record_comm_lite( async_op=async_op, ) CommCollective(duration_ms=duration_ms, **fields).save() + try: + from probing._core import note_last_comm + + note_last_comm( + op, + group_size, + nbytes, + int(fields.get("global_step", -1)), + ) + except Exception: + pass if write_trace_event: record_span( op, diff --git a/python/probing/site_hook.py b/python/probing/site_hook.py index ec6ef3ca..a8ff0d01 100644 --- a/python/probing/site_hook.py +++ b/python/probing/site_hook.py @@ -13,6 +13,7 @@ current_script_name, is_lightweight_module, is_probing_cli, + should_activate_probing, ) _RAN = False @@ -66,6 +67,7 @@ def _init_probing() -> None: import probing # noqa: F401 _execute_init_script(script_init) + _install_crash_handler() elif probe_value.lower() in ("2", "nested"): print( @@ -76,13 +78,12 @@ def _init_probing() -> None: os.environ["PROBING"] = probe_value _execute_init_script(script_init) + _install_crash_handler() elif probe_value.lower().startswith("regex:"): pattern = probe_value.split(":", 1)[1] try: - import re - - if re.search(pattern, current_script) is not None: + if should_activate_probing(): print( f"Activating probing for script matching '{pattern}'", file=sys.stderr, @@ -90,6 +91,7 @@ def _init_probing() -> None: import probing # noqa: F401 _execute_init_script(script_init) + _install_crash_handler() os.environ["PROBING"] = probe_value except Exception as exc: print(f"Error in regex pattern '{pattern}': {exc}", file=sys.stderr) @@ -103,6 +105,7 @@ def _init_probing() -> None: import probing # noqa: F401 _execute_init_script(script_init) + _install_crash_handler() os.environ["PROBING"] = probe_value except ImportError as exc: @@ -111,6 +114,15 @@ def _init_probing() -> None: print(f"Unexpected error in probing site hook: {exc}", file=sys.stderr) +def _install_crash_handler() -> None: + try: + from probing.crash import install + + install() + except Exception as exc: + print(f"probing crash handler install skipped: {exc}", file=sys.stderr) + + def _execute_init_script(script_init: str | None) -> None: if script_init is None: return diff --git a/python/probing/torchrun_cluster.py b/python/probing/torchrun_cluster.py index 4bf32ba6..0409ffae 100644 --- a/python/probing/torchrun_cluster.py +++ b/python/probing/torchrun_cluster.py @@ -88,7 +88,11 @@ def _reachable_addr(bound: str) -> str: host, port = bound.rsplit(":", 1) host = host.strip().strip("[]") if host in ("0.0.0.0", "::", "", "*"): - host = os.environ.get("MASTER_ADDR") or socket.gethostname() + master = os.environ.get("MASTER_ADDR", "").strip() + if master in ("127.0.0.1", "localhost"): + host = master + else: + host = master or socket.gethostname() return f"{host}:{port}" @@ -295,6 +299,11 @@ def refresh_node_role() -> bool: return False +def master_info() -> dict[str, str] | None: + """Return cluster master HTTP info after ``setup_torchrun_cluster()``.""" + return _MASTER_INFO + + def maybe_setup_torchrun_cluster() -> dict[str, str] | None: """Idempotent entry point; safe to call after ``init_process_group``.""" return setup_torchrun_cluster() diff --git a/skills/catalog.yaml b/skills/catalog.yaml index 0bdf9035..7ee7cf93 100644 --- a/skills/catalog.yaml +++ b/skills/catalog.yaml @@ -36,7 +36,20 @@ skills: - module_bottleneck - training_hang - memory_leak + - crash_triage path: health_overview/steps.yaml +- id: crash_triage + category: triage + priority: 1 + entry: true + description: Crash triage — grep PROBING CRASH in logs; read spill JSON files + tables: [] + pages: + - agent + related: + - health_overview + - training_hang + path: crash_triage/steps.yaml - id: training_hang category: reliability priority: 1 diff --git a/skills/crash_triage/SKILL.md b/skills/crash_triage/SKILL.md new file mode 100644 index 00000000..afa252ec --- /dev/null +++ b/skills/crash_triage/SKILL.md @@ -0,0 +1,25 @@ +--- +name: crash_triage +description: > + Crash triage: grep PROBING CRASH in rank logs; read spill JSON under + {PROBING_DATA_DIR}/crash//latest.json +category: triage +tags: [crash, error, exception, startup, 报错, 崩溃] +keywords: + en: ['crash', 'error', 'exception', 'failed', 'startup failure'] + zh: ['崩溃', '报错', '异常', '启动失败', '挂了'] +--- + +# Crash triage + +训练崩溃或启动失败时： + +1. **日志**：在各 rank stderr / launcher 日志中搜索 ``PROBING CRASH`` +2. **Spill**：``{PROBING_DATA_DIR}/crash//latest.json``（与 memtable 同根目录，默认 Linux ``/dev/shm/probing``，macOS ``$TMPDIR/probing``） +3. **Signal crash**：同目录下 ``signal-latest.json`` + +torchrun 级联退出后 memtable 不可用——以日志 + spill 为准。 + +## Related skills + +- 进程仍存活但无进展 → skill: training_hang diff --git a/skills/crash_triage/steps.yaml b/skills/crash_triage/steps.yaml new file mode 100644 index 00000000..c7eb35fa --- /dev/null +++ b/skills/crash_triage/steps.yaml @@ -0,0 +1,28 @@ +apiVersion: probing.dev/v1 +kind: Skill + +metadata: + id: crash_triage + title: "崩溃 / 启动失败诊断" + title_en: "Crash and startup failure triage" + category: triage + tags: [crash, error, exception, startup, 报错, 崩溃] + triggers: + keywords: + zh: ["崩溃", "报错", "异常", "启动失败", "挂了", "exit code"] + en: ["crash", "error", "exception", "failed", "startup failure"] + docs: | + 查 rank 日志中的 ``PROBING CRASH``，或 spill： + ``{PROBING_DATA_DIR}/crash//latest.json`` + +spec: + steps: [] + + summary_template: | + ## Crash 诊断 + 1. grep 各 rank 日志 ``PROBING CRASH`` + 2. 读 spill ``$PROBING_DATA_DIR/crash//latest.json`` + + next_steps: + - "全 rank 同一 fingerprint → 检查代码/环境" + - "无 exception 但卡住 → skill: training_hang" diff --git a/src/lib.rs b/src/lib.rs index 13e695b3..1a0d0fd5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -174,6 +174,7 @@ fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> { // Register all classes m.add_class::()?; m.add_function(wrap_pyfunction!(register_table_docs, m)?)?; + m.add_function(wrap_pyfunction!(start_local, m)?)?; m.add_class::()?; // Register all functions @@ -191,7 +192,14 @@ fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> { use probing_python::features::python_api::{is_enabled, should_enable_probing}; m.add_function(wrap_pyfunction!(is_enabled, m)?)?; m.add_function(wrap_pyfunction!(should_enable_probing, m)?)?; - m.add_function(wrap_pyfunction!(start_local, m)?)?; + use probing_python::features::crash::{ + crash_enabled, note_last_comm, record_crash, request_crash_hold, request_crash_release, + }; + m.add_function(wrap_pyfunction!(record_crash, m)?)?; + m.add_function(wrap_pyfunction!(crash_enabled, m)?)?; + m.add_function(wrap_pyfunction!(note_last_comm, m)?)?; + m.add_function(wrap_pyfunction!(request_crash_hold, m)?)?; + m.add_function(wrap_pyfunction!(request_crash_release, m)?)?; // Register config functions directly to the module (flattened) config::register_config_functions(m)?; diff --git a/tests/regression/skills/test_install_tools.py b/tests/regression/skills/test_install_tools.py index ff0c36da..d916adcf 100644 --- a/tests/regression/skills/test_install_tools.py +++ b/tests/regression/skills/test_install_tools.py @@ -41,14 +41,14 @@ def test_skill_roots_include_bundled(): assert "bundled" in labels or "repo" in labels -def test_merged_catalog_has_eight_skills(): +def test_merged_catalog_has_nine_skills(): catalog = load_catalog() - assert len(catalog.skills) == 8 + assert len(catalog.skills) == 9 def test_list_skills_tool(): skills = list_skills(limit=10) - assert len(skills) >= 8 + assert len(skills) >= 9 assert any(s.id == "health_overview" for s in skills) diff --git a/tests/regression/spec/api_spec.json b/tests/regression/spec/api_spec.json index fdd14395..97c68a40 100644 --- a/tests/regression/spec/api_spec.json +++ b/tests/regression/spec/api_spec.json @@ -103,6 +103,18 @@ "method": "GET", "uses_body": false, "response": { "content_type": "application/json", "cors": false } + }, + { + "local_path": "crash/hold", + "method": "GET", + "uses_body": false, + "response": { "content_type": "application/json", "cors": false } + }, + { + "local_path": "crash/release", + "method": "GET", + "uses_body": false, + "response": { "content_type": "application/json", "cors": false } } ], "other_extensions": [ diff --git a/tests/unit/probing/skills/test_loader.py b/tests/unit/probing/skills/test_loader.py index 147c09b3..ba4d2400 100644 --- a/tests/unit/probing/skills/test_loader.py +++ b/tests/unit/probing/skills/test_loader.py @@ -26,10 +26,11 @@ def test_skills_root_exists(): assert root.is_dir() -def test_catalog_loads_eight_skills(): +def test_catalog_loads_nine_skills(): catalog = load_catalog() - assert len(catalog.skills) == 8 + assert len(catalog.skills) == 9 ids = {p.id for p in catalog.skills} + assert "crash_triage" in ids assert "slow_rank" in ids assert "nccl_culprit_victim" in ids assert "health_overview" in ids diff --git a/web/assets/tailwind.css b/web/assets/tailwind.css index f2c15b1d..c01b64df 100644 --- a/web/assets/tailwind.css +++ b/web/assets/tailwind.css @@ -818,6 +818,16 @@ video { inset: 0px; } +.inset-x-0 { + left: 0px; + right: 0px; +} + +.inset-y-\[7px\] { + top: 7px; + bottom: 7px; +} + .-right-3 { right: -0.75rem; } @@ -826,6 +836,10 @@ video { bottom: 0px; } +.bottom-1\/2 { + bottom: 50%; +} + .bottom-3 { bottom: 0.75rem; } @@ -906,6 +920,30 @@ video { top: 1.125rem; } +.top-\[10px\] { + top: 10px; +} + +.top-\[11px\] { + top: 11px; +} + +.top-\[14px\] { + top: 14px; +} + +.top-\[52px\] { + top: 52px; +} + +.top-\[5px\] { + top: 5px; +} + +.top-\[9px\] { + top: 9px; +} + .top-full { top: 100%; } @@ -1090,6 +1128,10 @@ video { display: none; } +.h-0\.5 { + height: 0.125rem; +} + .h-1\.5 { height: 0.375rem; } @@ -1154,18 +1196,34 @@ video { height: 2.25rem; } +.h-\[12px\] { + height: 12px; +} + .h-\[20px\] { height: 20px; } +.h-\[22px\] { + height: 22px; +} + .h-\[28px\] { height: 28px; } +.h-\[7px\] { + height: 7px; +} + .h-full { height: 100%; } +.h-px { + height: 1px; +} + .h-screen { height: 100vh; } @@ -1487,6 +1545,10 @@ video { flex-shrink: 0; } +.grow { + flex-grow: 1; +} + .table-auto { table-layout: auto; } @@ -1629,6 +1691,10 @@ video { align-items: baseline; } +.items-stretch { + align-items: stretch; +} + .justify-end { justify-content: flex-end; } @@ -3337,6 +3403,18 @@ video { background-color: rgb(253 230 138 / var(--tw-bg-opacity, 1)); } +.bg-amber-200\/70 { + background-color: rgb(253 230 138 / 0.7); +} + +.bg-amber-200\/80 { + background-color: rgb(253 230 138 / 0.8); +} + +.bg-amber-400\/60 { + background-color: rgb(251 191 36 / 0.6); +} + .bg-amber-50 { --tw-bg-opacity: 1; background-color: rgb(255 251 235 / var(--tw-bg-opacity, 1)); @@ -3378,6 +3456,10 @@ video { background-color: rgb(191 219 254 / var(--tw-bg-opacity, 1)); } +.bg-blue-200\/70 { + background-color: rgb(191 219 254 / 0.7); +} + .bg-blue-400 { --tw-bg-opacity: 1; background-color: rgb(96 165 250 / var(--tw-bg-opacity, 1)); @@ -3437,6 +3519,10 @@ video { background-color: rgb(167 243 208 / var(--tw-bg-opacity, 1)); } +.bg-emerald-200\/70 { + background-color: rgb(167 243 208 / 0.7); +} + .bg-emerald-50 { --tw-bg-opacity: 1; background-color: rgb(236 253 245 / var(--tw-bg-opacity, 1)); @@ -3466,6 +3552,10 @@ video { background-color: rgb(229 231 235 / var(--tw-bg-opacity, 1)); } +.bg-gray-200\/70 { + background-color: rgb(229 231 235 / 0.7); +} + .bg-gray-300 { --tw-bg-opacity: 1; background-color: rgb(209 213 219 / var(--tw-bg-opacity, 1)); @@ -3574,6 +3664,15 @@ video { background-color: rgb(243 232 255 / var(--tw-bg-opacity, 1)); } +.bg-purple-200\/70 { + background-color: rgb(233 213 255 / 0.7); +} + +.bg-purple-500 { + --tw-bg-opacity: 1; + background-color: rgb(168 85 247 / var(--tw-bg-opacity, 1)); +} + .bg-red-200 { --tw-bg-opacity: 1; background-color: rgb(254 202 202 / var(--tw-bg-opacity, 1)); @@ -4790,6 +4889,10 @@ video { padding-bottom: 1px; } +.pb-0\.5 { + padding-bottom: 0.125rem; +} + .pb-1 { padding-bottom: 0.25rem; } @@ -5759,6 +5862,11 @@ video { --tw-ring-color: rgb(91 33 182 / var(--tw-ring-opacity, 1)); } +.ring-white { + --tw-ring-opacity: 1; + --tw-ring-color: rgb(255 255 255 / var(--tw-ring-opacity, 1)); +} + .ring-yellow-50 { --tw-ring-opacity: 1; --tw-ring-color: rgb(254 252 232 / var(--tw-ring-opacity, 1)); @@ -6075,6 +6183,10 @@ video { background-color: rgb(249 250 251 / var(--tw-bg-opacity, 1)); } +.hover\:bg-gray-50\/50:hover { + background-color: rgb(249 250 251 / 0.5); +} + .hover\:bg-gray-50\/60:hover { background-color: rgb(249 250 251 / 0.6); } diff --git a/web/src/components/mod.rs b/web/src/components/mod.rs index 16114173..b5bcdc14 100644 --- a/web/src/components/mod.rs +++ b/web/src/components/mod.rs @@ -38,6 +38,7 @@ pub mod profiling; pub mod profiling_sidebar_hint; pub mod sidebar; pub mod source_viewer; +pub mod span_timeline; pub mod stat_card; pub mod table_view; pub mod timeline_viewer; diff --git a/web/src/components/span_timeline.rs b/web/src/components/span_timeline.rs new file mode 100644 index 00000000..9f6191e1 --- /dev/null +++ b/web/src/components/span_timeline.rs @@ -0,0 +1,280 @@ +//! Vertical timeline lane paired with the span tree on the Spans page. +//! +//! Each visible tree row gets a timeline row on the left; expanding the tree +//! grows the timeline stack so hierarchy and timing stay aligned. + +use dioxus::prelude::*; + +use crate::api::SpanInfo; + +const TIMELINE_LANE_PX: f64 = 148.0; +const MIN_BAR_PX: f64 = 3.0; + +/// Nanosecond window covering all spans in the current tree. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct TraceTimeWindow { + pub start_ns: i64, + pub end_ns: i64, +} + +impl TraceTimeWindow { + pub fn from_spans(spans: &[SpanInfo]) -> Self { + let mut start = i64::MAX; + let mut end = i64::MIN; + + fn walk(span: &SpanInfo, start: &mut i64, end: &mut i64) { + *start = (*start).min(span.start_timestamp); + let span_end = span.end_timestamp.unwrap_or(span.start_timestamp); + *end = (*end).max(span_end); + for child in &span.children { + walk(child, start, end); + } + } + + for span in spans { + walk(span, &mut start, &mut end); + } + + if start == i64::MAX { + return Self { + start_ns: 0, + end_ns: 1, + }; + } + if end <= start { + end = start + 1; + } + Self { + start_ns: start, + end_ns: end, + } + } + + pub fn range_ns(&self) -> i64 { + (self.end_ns - self.start_ns).max(1) + } + + pub fn offset_px(&self, timestamp_ns: i64) -> f64 { + let pct = (timestamp_ns - self.start_ns) as f64 / self.range_ns() as f64; + (pct.clamp(0.0, 1.0) * TIMELINE_LANE_PX).max(0.0) + } + + pub fn width_px(&self, start_ns: i64, end_ns: Option) -> f64 { + let end = end_ns.unwrap_or(self.end_ns); + let dur = (end - start_ns).max(0) as f64; + let px = dur / self.range_ns() as f64 * TIMELINE_LANE_PX; + px.max(MIN_BAR_PX) + } +} + +pub fn format_axis_label(duration_ns: f64) -> String { + if duration_ns >= 1_000_000_000.0 { + format!("{:.2}s", duration_ns / 1_000_000_000.0) + } else if duration_ns >= 1_000_000.0 { + format!("{:.1}ms", duration_ns / 1_000_000.0) + } else if duration_ns >= 1_000.0 { + format!("{:.0}µs", duration_ns / 1_000.0) + } else { + format!("{:.0}ns", duration_ns) + } +} + +fn span_bar_style(phase: Option<&str>, active: bool) -> (&'static str, &'static str) { + if active { + return ("bg-amber-200/80", "bg-amber-500"); + } + match phase { + Some("forward") => ("bg-blue-200/70", "bg-blue-500"), + Some("backward") => ("bg-purple-200/70", "bg-purple-500"), + Some("optimizer") | Some("step") => ("bg-amber-200/70", "bg-amber-500"), + Some("idle") => ("bg-gray-200/70", "bg-gray-400"), + _ => ("bg-emerald-200/70", "bg-emerald-500"), + } +} + +fn span_tooltip(span: &SpanInfo, window: TraceTimeWindow) -> String { + let start = format_axis_label((span.start_timestamp - window.start_ns) as f64); + let end = span + .end_timestamp + .map(|t| format_axis_label((t - window.start_ns) as f64)) + .unwrap_or_else(|| "active".to_string()); + let dur = span + .end_timestamp + .map(|t| format_axis_label((t - span.start_timestamp) as f64)) + .unwrap_or_else(|| "active".to_string()); + format!( + "{}\nphase: {}\noffset: {} · end: {}\nduration: {}", + span.name, + span.phase.as_deref().unwrap_or("—"), + start, + end, + dur, + ) +} + +#[component] +pub fn SpanTimelineHeader(window: TraceTimeWindow) -> Element { + let total = format_axis_label(window.range_ns() as f64); + let mid = format_axis_label(window.range_ns() as f64 / 2.0); + rsx! { + div { + class: "flex shrink-0 border-b border-gray-200 bg-gray-50/90 sticky top-0 z-10", + div { + class: "shrink-0 border-r border-gray-200 px-2 py-1.5", + style: "width: {TIMELINE_LANE_PX}px", + div { class: "text-[10px] font-semibold uppercase tracking-wide text-gray-500 mb-1", + "Timeline" + } + div { class: "relative h-4 text-[9px] text-gray-400 font-mono tabular-nums", + span { class: "absolute left-0 top-0", "0" } + span { class: "absolute left-1/2 -translate-x-1/2 top-0", "{mid}" } + span { class: "absolute right-0 top-0", "{total}" } + div { class: "absolute inset-x-0 top-[14px] h-px bg-gray-200" } + div { class: "absolute left-0 top-[11px] w-px h-[7px] bg-gray-300" } + div { class: "absolute left-1/2 top-[11px] w-px h-[7px] bg-gray-300" } + div { class: "absolute right-0 top-[11px] w-px h-[7px] bg-gray-300" } + } + } + div { class: "flex-1 px-3 py-1.5 text-[10px] font-semibold uppercase tracking-wide text-gray-500", + "Span tree" + } + } + } +} + +#[component] +pub fn SpanTimelineLegend() -> Element { + rsx! { + div { + class: "flex flex-wrap items-center gap-x-4 gap-y-1 px-2 py-1.5 border-b border-gray-100 bg-white text-[10px] text-gray-500 sticky top-[52px] z-10", + span { class: "font-medium text-gray-600", "Lane" } + div { class: "inline-flex items-center gap-1", + span { class: "w-3 h-2 rounded-sm bg-blue-500" } + span { "forward" } + } + div { class: "inline-flex items-center gap-1", + span { class: "w-3 h-2 rounded-sm bg-purple-500" } + span { "backward" } + } + div { class: "inline-flex items-center gap-1", + span { class: "w-3 h-2 rounded-sm bg-amber-500" } + span { "optimizer" } + } + div { class: "inline-flex items-center gap-1", + span { class: "w-3 h-2 rounded-sm bg-emerald-500" } + span { "other" } + } + div { class: "inline-flex items-center gap-1", + span { class: "w-3 h-2 rounded-sm bg-amber-500 animate-pulse" } + span { "active" } + } + } + } +} + +#[component] +pub fn SpanTimelineBar(span: SpanInfo, window: TraceTimeWindow, depth: usize) -> Element { + let active = span.end_timestamp.is_none(); + let left = window.offset_px(span.start_timestamp); + let width = window.width_px(span.start_timestamp, span.end_timestamp); + let (track_bg, bar_bg) = span_bar_style(span.phase.as_deref(), active); + let indent = depth * 10; + let tooltip = span_tooltip(&span, window); + let lane_inner = TIMELINE_LANE_PX - indent as f64; + let bar_left = left.min(lane_inner - MIN_BAR_PX).max(0.0); + let bar_width = width.min(lane_inner - bar_left); + let guide_left = indent.saturating_sub(6); + + rsx! { + div { + class: "shrink-0 border-r border-gray-100 flex items-center py-0.5 relative", + style: "width: {TIMELINE_LANE_PX}px; padding-left: {indent}px", + title: "{tooltip}", + if depth > 0 { + div { + class: "absolute top-0 bottom-1/2 w-px bg-gray-200", + style: "left: {guide_left}px", + } + div { + class: "absolute top-1/2 w-2 h-px bg-gray-200", + style: "left: {guide_left}px", + } + } + div { class: "relative h-[22px] flex-1 min-w-0 pr-1", + div { class: "absolute inset-y-[7px] inset-x-0 rounded-full {track_bg}" } + div { + class: "absolute top-[5px] h-[12px] rounded-sm {bar_bg} shadow-sm", + style: "left: {bar_left:.2}px; width: {bar_width:.2}px;", + } + div { + class: "absolute top-[9px] w-1.5 h-1.5 rounded-full {bar_bg} ring-2 ring-white -translate-x-1/2", + style: "left: {bar_left:.2}px;", + } + if active { + div { + class: "absolute top-[10px] h-0.5 bg-amber-400/60 animate-pulse", + style: "left: calc({bar_left:.2}px + {bar_width:.2}px); right: 0;", + } + } + } + } + } +} + +/// Empty lane cell for detail rows (attributes / events) that have no bar. +#[component] +pub fn SpanTimelineSpacer() -> Element { + rsx! { + div { + class: "shrink-0 border-r border-gray-100 bg-gray-50/30", + style: "width: {TIMELINE_LANE_PX}px", + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn span(start: i64, end: Option) -> SpanInfo { + SpanInfo { + span_id: 1, + trace_id: 1, + parent_id: None, + name: "test".into(), + start_timestamp: start, + end_timestamp: end, + thread_id: 0, + phase: None, + location: None, + attributes: None, + children: vec![], + events: vec![], + } + } + + #[test] + fn window_from_spans_uses_min_max() { + let roots = vec![ + span(100, Some(500)), + SpanInfo { + children: vec![span(200, Some(800))], + ..span(50, Some(300)) + }, + ]; + let w = TraceTimeWindow::from_spans(&roots); + assert_eq!(w.start_ns, 50); + assert_eq!(w.end_ns, 800); + } + + #[test] + fn offset_and_width_px() { + let w = TraceTimeWindow { + start_ns: 0, + end_ns: 1000, + }; + assert!((w.offset_px(500) - TIMELINE_LANE_PX / 2.0).abs() < 0.01); + assert!((w.width_px(0, Some(1000)) - TIMELINE_LANE_PX).abs() < 0.01); + assert!(w.width_px(0, Some(1)) >= MIN_BAR_PX); + } +} diff --git a/web/src/pages/traces.rs b/web/src/pages/traces.rs index 517ce710..d2155df9 100644 --- a/web/src/pages/traces.rs +++ b/web/src/pages/traces.rs @@ -9,6 +9,9 @@ use crate::components::common::{query_result, AsyncBoundary, EmptyState}; use crate::components::icon::Icon; use crate::components::page::{PageContainer, PageTitle}; use crate::components::poll_status::{ManualRefreshStatus, RefreshButton}; +use crate::components::span_timeline::{ + SpanTimelineBar, SpanTimelineHeader, SpanTimelineLegend, SpanTimelineSpacer, TraceTimeWindow, +}; use crate::hooks::use_app_resource; use crate::state::investigation::{ clear_spans_investigation_filters, investigation_context_key, set_trace_context, @@ -64,7 +67,7 @@ pub fn Traces() -> Element { PageTitle { title: "Spans".to_string(), subtitle: Some( - "Hierarchical tracing spans from python.trace_event. For chrome trace timelines, use Profiling → Chrome trace.".to_string(), + "Span tree with an aligned timeline lane — expand nodes to grow both views. For kernel-level chrome traces, use Profiling → Chrome trace.".to_string(), ), icon: Some(&icondata::AiApiOutlined), header_right: Some(rsx! { @@ -353,15 +356,25 @@ fn TraceTreePanel( } } } else { - div { class: "px-2 py-2 max-h-[calc(100vh-14rem)] overflow-y-auto font-mono text-xs leading-5", - for span in filtered { - SpanView { - key: "{span.span_id}", - span: span.clone(), - depth: 0, - highlight: highlight.clone(), - expand_all, - collapse_all, + { + let time_window = TraceTimeWindow::from_spans(&filtered); + rsx! { + div { class: "max-h-[calc(100vh-14rem)] overflow-y-auto font-mono text-xs leading-5", + SpanTimelineHeader { window: time_window } + SpanTimelineLegend {} + div { class: "px-0 py-1", + for span in filtered { + SpanView { + key: "{span.span_id}", + span: span.clone(), + depth: 0, + highlight: highlight.clone(), + expand_all, + collapse_all, + time_window, + } + } + } } } } @@ -560,6 +573,7 @@ fn SpanView( highlight: SpanHighlight, expand_all: Signal, collapse_all: Signal, + time_window: TraceTimeWindow, ) -> Element { let mut expanded = use_signal(|| depth < 2); let has_children = !span.children.is_empty(); @@ -590,86 +604,108 @@ fn SpanView( rsx! { div { class: "min-w-0", div { - class: "{row_class}", - style: if indent > 0 { format!("padding-left: {indent}px") } else { String::new() }, - onclick: move |_| { - set_trace_context(trace_id, Some(&span_name), Some(thread_id)); - }, - if has_details { - button { - class: "shrink-0 w-4 h-4 flex items-center justify-center text-gray-400 hover:text-gray-700", - onclick: move |e| { - e.stop_propagation(); - expanded.set(!expanded()); + class: "flex items-stretch min-w-0 hover:bg-gray-50/50", + div { + class: "flex flex-1 min-w-0 items-stretch", + SpanTimelineBar { + span: span.clone(), + window: time_window, + depth, + } + div { + class: "{row_class} flex-1 min-w-0 border-b border-gray-50", + style: if indent > 0 { format!("padding-left: {indent}px") } else { String::new() }, + onclick: move |_| { + set_trace_context(trace_id, Some(&span_name), Some(thread_id)); }, - if expanded() { - Icon { icon: &icondata::AiCaretDownOutlined, class: "w-3 h-3" } + if has_details { + button { + class: "shrink-0 w-4 h-4 flex items-center justify-center text-gray-400 hover:text-gray-700", + onclick: move |e| { + e.stop_propagation(); + expanded.set(!expanded()); + }, + if expanded() { + Icon { icon: &icondata::AiCaretDownOutlined, class: "w-3 h-3" } + } else { + Icon { icon: &icondata::AiCaretRightOutlined, class: "w-3 h-3" } + } + } } else { - Icon { icon: &icondata::AiCaretRightOutlined, class: "w-3 h-3" } + span { class: "w-4 shrink-0" } + } + span { class: "font-semibold text-gray-900 shrink-0", "{span.name}" } + if let Some(ref phase) = span.phase { + span { + class: format!( + "shrink-0 px-1.5 py-px rounded text-[10px] font-sans font-medium bg-{} text-{}", + colors::CONTENT_ACCENT_BG, + colors::CONTENT_ACCENT_TEXT, + ), + "{phase}" + } + } + if let Some(ref location) = span.location { + if !location.is_empty() { + span { class: "text-gray-500 truncate max-w-[14rem]", "{location}" } + } + } + span { class: "text-gray-400 shrink-0", "id:{span.span_id}" } + if let Some(parent) = span.parent_id { + span { class: "text-gray-400 shrink-0", "↑{parent}" } + } + span { class: "text-gray-400 shrink-0", "t:{span.thread_id}" } + if let Some(dur) = duration { + span { class: "text-emerald-700 font-medium shrink-0", "{duration_label(dur)}" } + } else { + span { class: "text-amber-600 shrink-0", "active" } + } + if has_events { + span { class: "text-gray-400 shrink-0", "{span.events.len()}evt" } + } + if has_children { + span { class: "text-gray-400 shrink-0", "{span.children.len()}↓" } } } - } else { - span { class: "w-4 shrink-0" } - } - span { class: "font-semibold text-gray-900 shrink-0", "{span.name}" } - if let Some(ref phase) = span.phase { - span { - class: format!( - "shrink-0 px-1.5 py-px rounded text-[10px] font-sans font-medium bg-{} text-{}", - colors::CONTENT_ACCENT_BG, - colors::CONTENT_ACCENT_TEXT, - ), - "{phase}" - } - } - if let Some(ref location) = span.location { - if !location.is_empty() { - span { class: "text-gray-500 truncate max-w-[14rem]", "{location}" } - } - } - span { class: "text-gray-400 shrink-0", "id:{span.span_id}" } - if let Some(parent) = span.parent_id { - span { class: "text-gray-400 shrink-0", "↑{parent}" } - } - span { class: "text-gray-400 shrink-0", "t:{span.thread_id}" } - if let Some(dur) = duration { - span { class: "text-emerald-700 font-medium shrink-0", "{duration_label(dur)}" } - } else { - span { class: "text-amber-600 shrink-0", "active" } - } - if has_events { - span { class: "text-gray-400 shrink-0", "{span.events.len()}evt" } - } - if has_children { - span { class: "text-gray-400 shrink-0", "{span.children.len()}↓" } } } if expanded() && has_details { - div { - class: "space-y-0.5 pb-1", - style: format!("padding-left: {}px", indent + 20), - if has_attrs { - AttributesInline { raw: span.attributes.clone().unwrap_or_default() } - } - if has_events { - for event in span.events.iter() { - EventView { event: event.clone() } + if has_attrs { + div { class: "flex items-stretch min-w-0", + SpanTimelineSpacer {} + div { + class: "flex-1 min-w-0 pb-0.5", + style: format!("padding-left: {}px", indent + 20), + AttributesInline { raw: span.attributes.clone().unwrap_or_default() } } } - if has_children { - for child in span.children.iter() { - SpanView { - key: "{child.span_id}", - span: child.clone(), - depth: depth + 1, - highlight: highlight.clone(), - expand_all, - collapse_all, + } + if has_events { + for event in span.events.iter() { + div { class: "flex items-stretch min-w-0", + SpanTimelineSpacer {} + div { + class: "flex-1 min-w-0", + style: format!("padding-left: {}px", indent + 20), + EventView { event: event.clone(), time_window, span_start: span.start_timestamp } } } } } + if has_children { + for child in span.children.iter() { + SpanView { + key: "{child.span_id}", + span: child.clone(), + depth: depth + 1, + highlight: highlight.clone(), + expand_all, + collapse_all, + time_window, + } + } + } } } } @@ -715,11 +751,19 @@ fn attribute_value(val: &serde_json::Value) -> String { } #[component] -fn EventView(event: EventInfo) -> Element { +fn EventView(event: EventInfo, time_window: TraceTimeWindow, span_start: i64) -> Element { + let offset = crate::components::span_timeline::format_axis_label( + (event.timestamp - time_window.start_ns) as f64, + ); + let rel = + crate::components::span_timeline::format_axis_label((event.timestamp - span_start) as f64); rsx! { - div { class: "flex flex-wrap items-baseline gap-x-2 gap-y-0 py-0.5 text-gray-600", + div { + class: "flex flex-wrap items-baseline gap-x-2 gap-y-0 py-0.5 text-gray-600", + title: "t+{rel} from span start · {offset} from trace origin", span { class: "text-blue-500 shrink-0", "●" } span { class: "text-gray-800", "{event.name}" } + span { class: "text-gray-400 shrink-0 font-mono text-[10px]", "+{rel}" } if let Some(ref attrs) = event.attributes { if !attrs.is_empty() { span { class: "text-gray-500 break-all", "{attrs}" }