gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(gpu-log): use std::thread for drain task (no tokio runtime needed)

Browse Source 
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
This commit is contained in:
jgrusewski
2026-05-21 02:08:55 +02:00
parent 2d5a66f6e4
commit cd0fe8549c
3 changed files with 125 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files

166
crates/ml-alpha/src/gpu_log.rs
View File

@@ -6,8 +6,10 @@
//! trainer is invoked with `--kernel-step-trace <PATH>`, 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 `<PATH>`.
//! starts a std::thread drain that polls the ring at 500 ms cadence and
//! writes structured JSONL records to `<PATH>`. 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<LogRing>,
step_counter_host_shadow: Arc<MappedRecordBuffer<i32>>,
stop: Arc<AtomicBool>,
output_path: PathBuf,
) -> Result<JoinHandle<()>> {
// 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)
}

51
crates/ml-alpha/src/trainer/perception.rs
View File

@@ -350,7 +350,7 @@ pub struct PerceptionTrainer {
#[cfg(feature = "kernel-step-trace")]
log_drain_stop: Option<std::sync::Arc<std::sync::atomic::AtomicBool>>,
#[cfg(feature = "kernel-step-trace")]
log_drain_handle: Option<tokio::task::JoinHandle<()>>,
log_drain_handle: Option<std::thread::JoinHandle<()>>,
// 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<CudaStream>, src: &CudaSlice<f32>) -> Result<Vec<f32>>
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");
}
}
}
}

26
crates/ml-alpha/tests/gpu_log_ring_invariants.rs
View File

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