From cd0fe8549ce95c6d2734323fa5aa6ebb113ff726 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 21 May 2026 02:08:55 +0200 Subject: [PATCH] fix(gpu-log): use std::thread for drain task (no tokio runtime needed) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The drain task previously used tokio::spawn + tokio::time::sleep, gated by Handle::try_current().is_ok(). alpha_train's main() is synchronous (not #[tokio::main]) so Handle::try_current() returned Err and the drain task was silently skipped — the trace file was never created. Refactor to std::thread::spawn + std::thread::sleep. No runtime required; spawn is unconditional; Drop joins synchronously instead of aborting; BufWriter flushes on every drain cycle so even short runs (< 1 sec) capture their tail. Existing gpu_log_ring_invariants tests migrated from #[tokio::test] back to plain #[test]. Also resolves an exit-1 issue on alpha_train shutdown — likely a tokio task abort panicking through the synchronous Drop path. Local smoke (alpha_train --kernel-step-trace ...): - EXIT=0 - step_trace.jsonl: 1497 lines (~500 steps × 3 smoothness records) - valid JSONL, schema matches smoothness_controller emission contract --- crates/ml-alpha/src/gpu_log.rs | 166 ++++++++++-------- crates/ml-alpha/src/trainer/perception.rs | 51 +++--- .../ml-alpha/tests/gpu_log_ring_invariants.rs | 26 +-- 3 files changed, 125 insertions(+), 118 deletions(-) diff --git a/crates/ml-alpha/src/gpu_log.rs b/crates/ml-alpha/src/gpu_log.rs index 11c4f3bb9..eae9afbc1 100644 --- a/crates/ml-alpha/src/gpu_log.rs +++ b/crates/ml-alpha/src/gpu_log.rs @@ -6,8 +6,10 @@ //! trainer is invoked with `--kernel-step-trace `, the trainer //! allocates a mapped-pinned ring buffer + a 1-element device step //! counter, passes their pointers into logging-capable kernels, and -//! starts a tokio drain task that polls the ring at 500 ms cadence and -//! writes structured JSONL records to ``. +//! starts a std::thread drain that polls the ring at 500 ms cadence and +//! writes structured JSONL records to ``. The drain runs on a +//! plain OS thread (no async runtime needed) so it works whether the +//! caller is `#[tokio::main]` or a synchronous `fn main()`. //! //! Constants below mirror `crates/ml-alpha/cuda/gpu_log_ids.h` — single //! source of truth; manual mirror with on-build verification. @@ -17,12 +19,12 @@ use std::io::{BufWriter, Write}; use std::path::PathBuf; use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; -use std::time::{Duration, Instant}; +use std::thread::JoinHandle; +use std::time::Duration; use anyhow::{Context, Result}; use cudarc::driver::{CudaSlice, CudaStream}; use serde_json::{json, Value}; -use tokio::task::JoinHandle; use tracing::{debug, warn}; use crate::pinned_mem::MappedRecordBuffer; @@ -223,102 +225,110 @@ pub fn decode_to_json(kernel_id: u8, record_type: u8, step: u32, payload: &[f32] /// and serializes to the writer. Drops + warns if the drainer falls /// behind the ring window. /// -/// The writer is flushed every 1024 records or every 500 ms, -/// whichever comes first, and flushed one final time on stop. +/// The writer is flushed at the end of every drain cycle (cheap +/// because `BufWriter` coalesces writes) and once more on stop — +/// even sub-second runs capture their tail. /// -/// Stops cleanly when `stop` is set. +/// Runs on a plain `std::thread` — no async runtime required, so it +/// works equally from `#[tokio::main]` binaries and synchronous +/// `fn main()` binaries (e.g. `alpha_train`). +/// +/// Stops cleanly when `stop` is set; the drain thread observes the +/// flag at the top of its next 500 ms tick. pub fn spawn_drain_task_to_file( ring: Arc, step_counter_host_shadow: Arc>, stop: Arc, output_path: PathBuf, ) -> Result> { + // Open the file eagerly so a permission / path error surfaces to the + // caller instead of being swallowed inside the spawned thread. let file = File::create(&output_path) .with_context(|| format!("kernel-step-trace: create {}", output_path.display()))?; let mut writer = BufWriter::new(file); + let interval = Duration::from_millis(500); - Ok(tokio::spawn(async move { - let mut last_drained_step: u32 = 0; - let interval = Duration::from_millis(500); - let mut records_since_flush: usize = 0; - let mut last_flush = Instant::now(); - const FLUSH_RECORDS: usize = 1024; - const FLUSH_INTERVAL: Duration = Duration::from_millis(500); + let handle = std::thread::Builder::new() + .name("gpu_log_drain".into()) + .spawn(move || { + let mut last_drained_step: u32 = 0; - loop { - if stop.load(Ordering::Relaxed) { - debug!(target: "gpu_log", last_drained_step, "drain task exiting on stop signal"); - if let Err(e) = writer.flush() { - warn!(target: "gpu_log", error = %e, "final flush failed"); + loop { + if stop.load(Ordering::Relaxed) { + if let Err(e) = writer.flush() { + warn!(target: "gpu_log", error = %e, "final flush failed"); + } + debug!(target: "gpu_log", last_drained_step, "drain task exiting on stop signal"); + break; } - break; - } - tokio::time::sleep(interval).await; + std::thread::sleep(interval); - let current_step_i32 = step_counter_host_shadow.read_record(0); - if current_step_i32 < 0 { - // Counter overflow (unlikely at 4G steps but defensive). - warn!(target: "gpu_log", current_step_i32, "step counter went negative; skipping drain cycle"); - continue; - } - let current_step = current_step_i32 as u32; - if current_step <= last_drained_step { - continue; - } + let current_step_i32 = step_counter_host_shadow.read_record(0); + if current_step_i32 < 0 { + // Counter overflow (unlikely at 4G steps but defensive). + warn!(target: "gpu_log", current_step_i32, "step counter went negative; skipping drain cycle"); + continue; + } + let current_step = current_step_i32 as u32; + if current_step <= last_drained_step { + continue; + } - let max_drain_window = (N_SLOTS / N_RECORDS_PER_STEP) as u32; - let oldest_safe_step = current_step.saturating_sub(max_drain_window); - if last_drained_step < oldest_safe_step { - let skipped = oldest_safe_step.saturating_sub(last_drained_step); - warn!(target: "gpu_log", - last_drained_step, current_step, skipped, - "drainer fell behind — records dropped"); - last_drained_step = oldest_safe_step; - } + let max_drain_window = (N_SLOTS / N_RECORDS_PER_STEP) as u32; + let oldest_safe_step = current_step.saturating_sub(max_drain_window); + if last_drained_step < oldest_safe_step { + let skipped = oldest_safe_step.saturating_sub(last_drained_step); + warn!(target: "gpu_log", + last_drained_step, current_step, skipped, + "drainer fell behind — records dropped"); + last_drained_step = oldest_safe_step; + } - for step in (last_drained_step + 1)..=current_step { - for kid in 0..N_KERNELS as u8 { - for rt in 0..N_RT_PER_KERNEL as u8 { - let slot = ((step as usize) * N_RECORDS_PER_STEP - + (kid as usize) * N_RT_PER_KERNEL - + rt as usize) & (N_SLOTS - 1); - let rec = ring.buffer.read_record(slot); - if rec.header.magic != MAGIC { - continue; + for step in (last_drained_step + 1)..=current_step { + for kid in 0..N_KERNELS as u8 { + for rt in 0..N_RT_PER_KERNEL as u8 { + let slot = ((step as usize) * N_RECORDS_PER_STEP + + (kid as usize) * N_RT_PER_KERNEL + + rt as usize) & (N_SLOTS - 1); + let rec = ring.buffer.read_record(slot); + if rec.header.magic != MAGIC { + continue; + } + if rec.header.step != step { + // Slot was written by a later step; this kernel + // didn't emit at THIS step. + continue; + } + let n = (rec.header.payload_words as usize).min(PAYLOAD_F32); + let value = decode_to_json( + rec.header.kernel_id, + rec.header.record_type, + rec.header.step, + &rec.payload[..n], + ); + if let Err(e) = serde_json::to_writer(&mut writer, &value) { + warn!(target: "gpu_log", error = %e, "JSONL serialize failed"); + continue; + } + if let Err(e) = writeln!(&mut writer) { + warn!(target: "gpu_log", error = %e, "JSONL newline write failed"); + continue; + } } - if rec.header.step != step { - // Slot was written by a later step; this kernel - // didn't emit at THIS step. - continue; - } - let n = (rec.header.payload_words as usize).min(PAYLOAD_F32); - let value = decode_to_json( - rec.header.kernel_id, - rec.header.record_type, - rec.header.step, - &rec.payload[..n], - ); - if let Err(e) = serde_json::to_writer(&mut writer, &value) { - warn!(target: "gpu_log", error = %e, "JSONL serialize failed"); - continue; - } - if let Err(e) = writeln!(&mut writer) { - warn!(target: "gpu_log", error = %e, "JSONL newline write failed"); - continue; - } - records_since_flush += 1; } } - } - last_drained_step = current_step; + last_drained_step = current_step; - if records_since_flush >= FLUSH_RECORDS || last_flush.elapsed() >= FLUSH_INTERVAL { + // Flush on every drain cycle. BufWriter coalesces writes + // already, so the syscall cost is amortised over the cycle; + // the upside is that even sub-second runs (a kill-9'd + // smoke, a `cargo test` Ctrl-C) leave a complete tail on + // disk instead of losing it in BufWriter. if let Err(e) = writer.flush() { warn!(target: "gpu_log", error = %e, "periodic flush failed"); } - records_since_flush = 0; - last_flush = Instant::now(); } - } - })) + }) + .context("spawn gpu_log_drain thread")?; + Ok(handle) } diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index f714e0e12..a43440f45 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -350,7 +350,7 @@ pub struct PerceptionTrainer { #[cfg(feature = "kernel-step-trace")] log_drain_stop: Option>, #[cfg(feature = "kernel-step-trace")] - log_drain_handle: Option>, + log_drain_handle: Option>, // X6: LN_b weights (formerly ln_gain_d / ln_bias_d) moved to // `self.trunk.ln_b_gain_d` / `ln_b_bias_d`. LN_b is the LayerNorm // AFTER Mamba2 stack 2 — stabilises the trunk-to-CfC distribution. @@ -1177,26 +1177,16 @@ impl PerceptionTrainer { (f, module) }; let log_drain_stop = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)); - // The drain task uses `tokio::spawn`, which requires a runtime - // handle on the current thread. In production the trainer runs - // inside a tokio runtime (binaries are `#[tokio::main]`). In - // sync `#[test]` units there is no runtime — we skip the drain - // (records still land in the ring; the test simply doesn't read - // them via the JSONL file path). This is the canonical pattern - // for opt-in async telemetry. - let log_drain_handle = if tokio::runtime::Handle::try_current().is_ok() { - Some(crate::gpu_log::spawn_drain_task_to_file( - log_ring.clone(), - log_step_counter_host_shadow.clone(), - log_drain_stop.clone(), - trace_path.clone(), - )?) - } else { - tracing::debug!(target: "gpu_log", - path = %trace_path.display(), - "no tokio runtime available — skipping drain task spawn (records still emitted to ring)"); - None - }; + // The drain runs on a plain `std::thread`, so spawning is + // unconditional — it works the same whether the trainer is + // invoked from `#[tokio::main]` or a synchronous `fn main()` + // (e.g. `alpha_train`). No runtime-handle gate required. + let log_drain_handle = Some(crate::gpu_log::spawn_drain_task_to_file( + log_ring.clone(), + log_step_counter_host_shadow.clone(), + log_drain_stop.clone(), + trace_path.clone(), + )?); ( Some(log_ring), Some(log_step_counter_d), @@ -4230,12 +4220,12 @@ fn download(stream: &Arc, src: &CudaSlice) -> Result> Ok(staging.read_all()) } -// ── Drop: shut down the GPU log drain task (feature-gated) ── +// ── Drop: shut down the GPU log drain thread (feature-gated) ── // Setting the stop flag lets the drainer exit cleanly on its next 500ms -// tick; `.abort()` is the belt-and-braces guard against a dropped -// runtime / hanging shutdown. Drop cannot be async, so we cannot await -// the handle — abort is the canonical pattern for tokio JoinHandle in -// non-async destructors. +// tick; the drain thread observes the flag at the top of each iteration, +// flushes the BufWriter, and exits. We then `join()` synchronously so +// the trace file is fully flushed by the time the trainer is dropped. +// The bounded 500 ms tick guarantees the join completes promptly. #[cfg(feature = "kernel-step-trace")] impl Drop for PerceptionTrainer { fn drop(&mut self) { @@ -4244,7 +4234,14 @@ impl Drop for PerceptionTrainer { stop.store(true, Ordering::Relaxed); } if let Some(h) = self.log_drain_handle.take() { - h.abort(); + // Drain thread observes the stop flag on its next tick (≤500 ms), + // flushes the BufWriter, and exits. `join()` blocks until the + // thread terminates — guarantees the JSONL tail lands on disk. + if let Err(e) = h.join() { + tracing::warn!(target: "gpu_log", + error = ?e, + "gpu_log_drain thread join failed"); + } } } } diff --git a/crates/ml-alpha/tests/gpu_log_ring_invariants.rs b/crates/ml-alpha/tests/gpu_log_ring_invariants.rs index 188b6adea..80ce246e0 100644 --- a/crates/ml-alpha/tests/gpu_log_ring_invariants.rs +++ b/crates/ml-alpha/tests/gpu_log_ring_invariants.rs @@ -117,8 +117,8 @@ fn host_write_round_trip() { } } -#[tokio::test(flavor = "current_thread")] -async fn magic_validation_skips_torn() { +#[test] +fn magic_validation_skips_torn() { let _dev = test_device(); let ring = Arc::new(LogRing::alloc().expect("alloc ring")); let counter = alloc_host_counter(); @@ -153,10 +153,10 @@ async fn magic_validation_skips_torn() { let h = spawn_drain_task_to_file(ring.clone(), counter.clone(), stop.clone(), path.clone()) .expect("spawn drain"); // Give drainer one tick (>500 ms cadence + a bit of slack). - tokio::time::sleep(Duration::from_millis(700)).await; + std::thread::sleep(Duration::from_millis(700)); stop.store(true, Ordering::Relaxed); - // Wait for the next sleep boundary so the drain task observes the stop flag. - let _ = tokio::time::timeout(Duration::from_secs(2), h).await; + // Wait for the drain thread to observe the stop flag and flush. + let _ = h.join(); // The valid record should produce exactly one JSONL line; the torn // (magic=0) and mismatched-step records must NOT appear. @@ -202,8 +202,8 @@ fn ring_wrap_modulo_arithmetic() { } } -#[tokio::test(flavor = "current_thread")] -async fn step_gap_detection_emits_warn() { +#[test] +fn step_gap_detection_emits_warn() { let _dev = test_device(); let ring = Arc::new(LogRing::alloc().expect("alloc ring")); let counter = alloc_host_counter(); @@ -222,17 +222,17 @@ async fn step_gap_detection_emits_warn() { tmp.path().to_path_buf(), ) .expect("spawn drain"); - tokio::time::sleep(Duration::from_millis(700)).await; + std::thread::sleep(Duration::from_millis(700)); stop.store(true, Ordering::Relaxed); - let _ = tokio::time::timeout(Duration::from_secs(2), h).await; + let _ = h.join(); // Pass if no panic — drainer clamped `last_drained_step` to `oldest_safe_step` // and emitted a warn. We don't tap tracing here; the absence of panic + // the drainer's known structure (warns then clamps) is the pass condition. } -#[tokio::test(flavor = "current_thread")] -async fn decoder_dispatch_unknown_falls_through() { +#[test] +fn decoder_dispatch_unknown_falls_through() { let _dev = test_device(); let ring = Arc::new(LogRing::alloc().expect("alloc ring")); let counter = alloc_host_counter(); @@ -256,9 +256,9 @@ async fn decoder_dispatch_unknown_falls_through() { tmp.path().to_path_buf(), ) .expect("spawn drain"); - tokio::time::sleep(Duration::from_millis(700)).await; + std::thread::sleep(Duration::from_millis(700)); stop.store(true, Ordering::Relaxed); - let _ = tokio::time::timeout(Duration::from_secs(2), h).await; + let _ = h.join(); // The drainer's iteration is `for kid in 0..N_KERNELS` — so kid=99 will // never be visited by the drainer's slot walk. But the `decode_to_json()`