Enhance probing features with new crash demos and improved error hand…

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:

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()

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()

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"

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<Option<JoinHandle<()>>>,
 }
 
+fn lock_compactor_handle(
+    m: &Mutex<Option<JoinHandle<()>>>,
+) -> MutexGuard<'_, Option<JoinHandle<()>>> {
+    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<ColdCompactor> = Lazy::new(|| ColdCompactor {
@@ -1176,15 +1188,15 @@ 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.
     pub fn stop(&self) {
         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();
         }
     }

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

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),
     }
 }
 

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"
 

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<String>,
 
     /// 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}");