diff --git a/crates/ml-alpha/examples/alpha_train.rs b/crates/ml-alpha/examples/alpha_train.rs index 2afd8ba84..231f696a3 100644 --- a/crates/ml-alpha/examples/alpha_train.rs +++ b/crates/ml-alpha/examples/alpha_train.rs @@ -161,6 +161,17 @@ struct Cli { /// `docs/superpowers/specs/2026-05-20-crt-train-output-smoothness-design.md`. #[arg(long, default_value_t = 0.0)] smoothness_base_lambda: f32, + + /// Enable per-step kernel-state JSONL trace to the given file. + /// + /// When set, the trainer allocates the GPU log ring + spawns a drain + /// task that writes one JSONL record per kernel emission to this path. + /// Requires the `kernel-step-trace` Cargo feature at compile time. + /// When unset (default), no ring allocated, zero overhead. + /// + /// Example: `--kernel-step-trace /feature-cache/.../step_trace.jsonl` + #[arg(long)] + kernel_step_trace: Option, } #[derive(Serialize, serde::Deserialize, Default)] @@ -275,6 +286,7 @@ fn main() -> Result<()> { horizon_weights, n_batch: cli.batch_size, smoothness_base_lambda: cli.smoothness_base_lambda, + kernel_step_trace_path: cli.kernel_step_trace.clone(), }; let mut trainer = PerceptionTrainer::new(&dev, &trainer_cfg).context("trainer init")?; diff --git a/crates/ml-alpha/src/gpu_log.rs b/crates/ml-alpha/src/gpu_log.rs index d4d88bf3e..11c4f3bb9 100644 --- a/crates/ml-alpha/src/gpu_log.rs +++ b/crates/ml-alpha/src/gpu_log.rs @@ -2,23 +2,28 @@ //! //! See `docs/superpowers/specs/2026-05-21-gpu-log-ring-design.md`. //! -//! Behaviour: when the `kernel-step-trace` feature is enabled, 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 emits structured -//! tracing events. +//! Behaviour: when the `kernel-step-trace` feature is enabled AND the +//! 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 ``. //! //! Constants below mirror `crates/ml-alpha/cuda/gpu_log_ids.h` — single //! source of truth; manual mirror with on-build verification. +use std::fs::File; +use std::io::{BufWriter, Write}; +use std::path::PathBuf; use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; -use std::time::Duration; +use std::time::{Duration, Instant}; use anyhow::{Context, Result}; use cudarc::driver::{CudaSlice, CudaStream}; +use serde_json::{json, Value}; use tokio::task::JoinHandle; -use tracing::{debug, info, warn}; +use tracing::{debug, warn}; use crate::pinned_mem::MappedRecordBuffer; @@ -126,82 +131,126 @@ pub fn alloc_step_counter(stream: &Arc) -> Result> { stream.alloc_zeros::(1).context("step_counter alloc") } -// ─── Decoder dispatch ──────────────────────────────────────────────── +// ─── JSONL serializer ──────────────────────────────────────────────── +// +// Each (kernel_id, record_type) pair maps to a fixed schema; the +// `payload` field names below come from the kernel-side emission +// sites (mirror of `smoothness_lambda_controller.cu`'s `log_record()` +// calls). Adding a new logging-capable kernel = adding a new arm. -fn default_decoder(step: u32, kernel_id: u8, record_type: u8, payload: &[f32]) { - debug!(target: "gpu_log", step, kernel_id, record_type, ?payload, "unknown record"); -} - -fn smoothness_input(step: u32, payload: &[f32]) { - if payload.len() >= 10 { - info!(target: "gpu_log", - step, - raw_h30 = payload[0], raw_h100 = payload[1], raw_h300 = payload[2], - raw_h1000 = payload[3], raw_h6000 = payload[4], - jitter_in_h30 = payload[5], jitter_in_h100 = payload[6], - jitter_in_h300 = payload[7], jitter_in_h1000 = payload[8], - jitter_in_h6000 = payload[9], - "smoothness_ctrl RT_INPUT"); +/// Kernel-name mirror for the `kname` JSONL field. +fn kernel_name(kernel_id: u8) -> &'static str { + match kernel_id { + KID_SMOOTHNESS_CONTROLLER => "smoothness_controller", + KID_OUTPUT_SMOOTHNESS => "output_smoothness", + KID_BCE_MULTI_HORIZON => "bce_multi_horizon", + KID_HORIZON_LAMBDA => "horizon_lambda", + KID_HEADS_GRN_FWD => "heads_grn_fwd", + KID_HEADS_GRN_BWD => "heads_grn_bwd", + KID_CFC_STEP => "cfc_step", + KID_RESERVED_7 => "reserved_7", + _ => "unknown", } } -fn smoothness_state(step: u32, payload: &[f32]) { - if payload.len() >= 10 { - info!(target: "gpu_log", - step, - ema_h30 = payload[0], ema_h100 = payload[1], ema_h300 = payload[2], - ema_h1000 = payload[3], ema_h6000 = payload[4], - target_h30 = payload[5], target_h100 = payload[6], target_h300 = payload[7], - target_h1000 = payload[8], target_h6000 = payload[9], - "smoothness_ctrl RT_STATE"); +/// Record-type-name mirror for the `rt_name` JSONL field. +fn rt_name(record_type: u8) -> &'static str { + match record_type { + RT_INPUT => "input", + RT_STATE => "state", + RT_OUTPUT => "output", + RT_DEBUG => "debug", + _ => "unknown", } } -fn smoothness_output(step: u32, payload: &[f32]) { - if payload.len() >= 10 { - info!(target: "gpu_log", - step, - excess_h30 = payload[0], excess_h100 = payload[1], excess_h300 = payload[2], - excess_h1000 = payload[3], excess_h6000 = payload[4], - lambda_h30 = payload[5], lambda_h100 = payload[6], lambda_h300 = payload[7], - lambda_h1000 = payload[8], lambda_h6000 = payload[9], - "smoothness_ctrl RT_OUTPUT"); - } -} - -/// Static dispatch table indexed by (kernel_id, record_type). -/// Add per-(kid, rt) decoders here as new kernels opt into the ring. -pub fn decode(kernel_id: u8, record_type: u8, step: u32, payload: &[f32]) { +/// Build the per-(kid, rt) payload map. Returns a JSON object with +/// field names matching the kernel-side emission contract. Unknown +/// (kid, rt) pairs fall through to an array-of-floats payload so no +/// records are silently dropped. +fn payload_json(kernel_id: u8, record_type: u8, payload: &[f32]) -> Value { match (kernel_id, record_type) { - (KID_SMOOTHNESS_CONTROLLER, RT_INPUT) => smoothness_input(step, payload), - (KID_SMOOTHNESS_CONTROLLER, RT_STATE) => smoothness_state(step, payload), - (KID_SMOOTHNESS_CONTROLLER, RT_OUTPUT) => smoothness_output(step, payload), - _ => default_decoder(step, kernel_id, record_type, payload), + (KID_SMOOTHNESS_CONTROLLER, RT_INPUT) if payload.len() >= 10 => json!({ + "raw_h30": payload[0], "raw_h100": payload[1], "raw_h300": payload[2], + "raw_h1000": payload[3], "raw_h6000": payload[4], + "jitter_in_h30": payload[5], "jitter_in_h100": payload[6], + "jitter_in_h300": payload[7], "jitter_in_h1000": payload[8], + "jitter_in_h6000": payload[9], + }), + (KID_SMOOTHNESS_CONTROLLER, RT_STATE) if payload.len() >= 10 => json!({ + "ema_h30": payload[0], "ema_h100": payload[1], "ema_h300": payload[2], + "ema_h1000": payload[3], "ema_h6000": payload[4], + "target_h30": payload[5], "target_h100": payload[6], "target_h300": payload[7], + "target_h1000": payload[8], "target_h6000": payload[9], + }), + (KID_SMOOTHNESS_CONTROLLER, RT_OUTPUT) if payload.len() >= 10 => json!({ + "excess_h30": payload[0], "excess_h100": payload[1], "excess_h300": payload[2], + "excess_h1000": payload[3], "excess_h6000": payload[4], + "lambda_h30": payload[5], "lambda_h100": payload[6], "lambda_h300": payload[7], + "lambda_h1000": payload[8], "lambda_h6000": payload[9], + }), + _ => { + // Unknown (kid, rt) — preserve the raw payload as a "values" + // array so the JSONL line is still well-formed and the + // record isn't silently lost. + let values: Vec = payload.iter().map(|v| json!(*v)).collect(); + json!({ "values": values }) + } } } +/// Build the full JSONL record for one (kid, rt, step, payload) tuple. +pub fn decode_to_json(kernel_id: u8, record_type: u8, step: u32, payload: &[f32]) -> Value { + json!({ + "step": step, + "kid": kernel_id, + "kname": kernel_name(kernel_id), + "rt": record_type, + "rt_name": rt_name(record_type), + "payload": payload_json(kernel_id, record_type, payload), + }) +} + // ─── Drain task ────────────────────────────────────────────────────── -/// Spawn the drain task. Polls the step counter (via a mapped-pinned host -/// shadow that some other component updates each tick — typically a -/// secondary memcpy_dtod_async in `step_batched`, or a separate background -/// reader) at 500 ms cadence; for each new step in -/// `[last_drained..current]`, walks `N_RECORDS_PER_STEP` slots, validates -/// `magic == MAGIC` and `header.step == expected_step`, dispatches to -/// `decode()`. Drops + warns if the drainer falls behind the ring window. +/// Spawn the drain task with a JSONL file sink. Opens `output_path` +/// (truncating any existing file), wraps in a `BufWriter`, and writes +/// one JSON object per line for every valid record observed. Polls +/// the step counter (via a mapped-pinned host shadow updated each +/// captured-graph replay) at 500 ms cadence; for each new step in +/// `[last_drained..current]`, walks `N_RECORDS_PER_STEP` slots, +/// validates `magic == MAGIC` and `header.step == expected_step`, +/// 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. /// /// Stops cleanly when `stop` is set. -pub fn spawn_drain_task( +pub fn spawn_drain_task_to_file( ring: Arc, step_counter_host_shadow: Arc>, stop: Arc, -) -> JoinHandle<()> { - tokio::spawn(async move { + output_path: PathBuf, +) -> Result> { + let file = File::create(&output_path) + .with_context(|| format!("kernel-step-trace: create {}", output_path.display()))?; + let mut writer = BufWriter::new(file); + + 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); + 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"); + } break; } tokio::time::sleep(interval).await; @@ -243,12 +292,33 @@ pub fn spawn_drain_task( continue; } let n = (rec.header.payload_words as usize).min(PAYLOAD_F32); - decode(rec.header.kernel_id, rec.header.record_type, - rec.header.step, &rec.payload[..n]); + 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; + + if records_since_flush >= FLUSH_RECORDS || last_flush.elapsed() >= FLUSH_INTERVAL { + if let Err(e) = writer.flush() { + warn!(target: "gpu_log", error = %e, "periodic flush failed"); + } + records_since_flush = 0; + last_flush = Instant::now(); + } } - }) + })) } diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index 4ea3ea787..f714e0e12 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -94,6 +94,13 @@ pub struct PerceptionTrainerConfig { /// regularization is negligible). Set to e.g. 0.01 to enable. /// See `docs/superpowers/specs/2026-05-20-crt-train-output-smoothness-design.md`. pub smoothness_base_lambda: f32, + + /// Optional path for the kernel-step-trace JSONL drain. None → + /// ring not allocated, drain task not spawned. Some(path) → ring + /// allocated + drain writes JSONL records to `path` (truncated). + /// Per `feedback_no_feature_flags`: gated by the compile-time + /// `kernel-step-trace` feature; specific name justifies the gate. + pub kernel_step_trace_path: Option, } impl Default for PerceptionTrainerConfig { @@ -107,6 +114,7 @@ impl Default for PerceptionTrainerConfig { horizon_weights: [1.0; N_HORIZONS], n_batch: 1, smoothness_base_lambda: 0.0, + kernel_step_trace_path: None, } } } @@ -317,26 +325,30 @@ pub struct PerceptionTrainer { smoothness_controller_fn: CudaFunction, _smoothness_controller_module: Arc, // ── GPU log ring (feature: kernel-step-trace) ── - // When enabled, the trainer allocates a mapped-pinned ring + a device + // When the feature is compiled AND `cfg.kernel_step_trace_path` is + // Some(path), the trainer allocates a mapped-pinned ring + a device // step counter, passes their pointers into logging-capable kernels // (currently `smoothness_lambda_controller`), and spawns a tokio drain - // task that polls the ring at 500 ms cadence and emits structured - // tracing events. See `docs/superpowers/specs/2026-05-21-gpu-log-ring-design.md`. + // task that polls the ring at 500 ms cadence and writes JSONL records + // to `path`. When the path is None, every Option below stays None — + // no allocations, no tick kernel, no drain task, zero overhead even + // with the feature compiled in. + // See `docs/superpowers/specs/2026-05-21-gpu-log-ring-design.md`. #[cfg(feature = "kernel-step-trace")] - log_ring: std::sync::Arc, + log_ring: Option>, #[cfg(feature = "kernel-step-trace")] - log_step_counter_d: CudaSlice, + log_step_counter_d: Option>, /// Mapped-pinned host shadow of `log_step_counter_d`. Updated via /// memcpy_dtod_async inside the captured region; read by the drain /// task each tick. #[cfg(feature = "kernel-step-trace")] - log_step_counter_host_shadow: std::sync::Arc>, + log_step_counter_host_shadow: Option>>, #[cfg(feature = "kernel-step-trace")] - gpu_log_tick_fn: CudaFunction, + gpu_log_tick_fn: Option, #[cfg(feature = "kernel-step-trace")] - _gpu_log_module: Arc, + _gpu_log_module: Option>, #[cfg(feature = "kernel-step-trace")] - log_drain_stop: std::sync::Arc, + log_drain_stop: Option>, #[cfg(feature = "kernel-step-trace")] log_drain_handle: Option>, // X6: LN_b weights (formerly ln_gain_d / ln_bias_d) moved to @@ -1135,45 +1147,67 @@ impl PerceptionTrainer { // DtoD memcpy of the device step counter into the host shadow // is captured in the graph, so the drainer sees the latest // counter without an extra sync. + // + // Runtime activation: only allocate + spawn when + // `cfg.kernel_step_trace_path` is Some(path). Otherwise every + // field below stays None and the trainer carries zero overhead. #[cfg(feature = "kernel-step-trace")] - let log_ring = std::sync::Arc::new(crate::gpu_log::LogRing::alloc()?); - #[cfg(feature = "kernel-step-trace")] - let log_step_counter_d = crate::gpu_log::alloc_step_counter(&stream)?; - #[cfg(feature = "kernel-step-trace")] - let log_step_counter_host_shadow = { - let buf = unsafe { crate::pinned_mem::MappedRecordBuffer::::new(1) } - .map_err(|e| anyhow::anyhow!("log step counter shadow: {e}"))?; - buf.write_record(0, 0); - std::sync::Arc::new(buf) - }; - #[cfg(feature = "kernel-step-trace")] - let (gpu_log_tick_fn, gpu_log_module) = { - const GPU_LOG_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/gpu_log_ring.cubin")); - let module = ctx.load_cubin(GPU_LOG_CUBIN.to_vec()) - .context("load gpu_log_ring cubin")?; - let f = module.load_function("gpu_log_tick").context("load gpu_log_tick")?; - (f, module) - }; - #[cfg(feature = "kernel-step-trace")] - 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 through the structured tracing path). This is the - // canonical pattern for opt-in async telemetry. - #[cfg(feature = "kernel-step-trace")] - let log_drain_handle = if tokio::runtime::Handle::try_current().is_ok() { - Some(crate::gpu_log::spawn_drain_task( - log_ring.clone(), - log_step_counter_host_shadow.clone(), - log_drain_stop.clone(), - )) + let ( + log_ring, + log_step_counter_d, + log_step_counter_host_shadow, + gpu_log_tick_fn, + gpu_log_module, + log_drain_stop, + log_drain_handle, + ) = if let Some(trace_path) = cfg.kernel_step_trace_path.clone() { + let log_ring = std::sync::Arc::new(crate::gpu_log::LogRing::alloc()?); + let log_step_counter_d = crate::gpu_log::alloc_step_counter(&stream)?; + let log_step_counter_host_shadow = { + let buf = unsafe { crate::pinned_mem::MappedRecordBuffer::::new(1) } + .map_err(|e| anyhow::anyhow!("log step counter shadow: {e}"))?; + buf.write_record(0, 0); + std::sync::Arc::new(buf) + }; + let (gpu_log_tick_fn, gpu_log_module) = { + const GPU_LOG_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/gpu_log_ring.cubin")); + let module = ctx.load_cubin(GPU_LOG_CUBIN.to_vec()) + .context("load gpu_log_ring cubin")?; + let f = module.load_function("gpu_log_tick").context("load gpu_log_tick")?; + (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 + }; + ( + Some(log_ring), + Some(log_step_counter_d), + Some(log_step_counter_host_shadow), + Some(gpu_log_tick_fn), + Some(gpu_log_module), + Some(log_drain_stop), + log_drain_handle, + ) } else { - tracing::debug!(target: "gpu_log", - "no tokio runtime available — skipping drain task spawn (records still emitted to ring)"); - None + (None, None, None, None, None, None, None) }; let k = cfg.seq_len; @@ -1211,7 +1245,7 @@ impl PerceptionTrainer { #[cfg(feature = "kernel-step-trace")] log_drain_stop, #[cfg(feature = "kernel-step-trace")] - log_drain_handle, + log_drain_handle, // already Option<_>; field type unchanged // LayerNorm state. `ln_gain_d` / `ln_bias_d` constructed // above via `upload` (gain=1, bias=0). LN_HIDDEN matches // HIDDEN_DIM (128) by construction — checked at kernel @@ -1674,15 +1708,23 @@ impl PerceptionTrainer { // thread/single-block kernel: no host malloc, no host branch // (per `pearl_cudarc_disable_event_tracking_for_graph_capture` // and `pearl_no_host_branches_in_captured_graph`). + // + // Activated only when the runtime trace path is Some — i.e. + // when the trainer was constructed with + // `cfg.kernel_step_trace_path = Some(_)`. With path = None all + // the trace fields are None and this block is skipped, so the + // tick kernel never launches and zero overhead is incurred. #[cfg(feature = "kernel-step-trace")] + if let (Some(tick_fn), Some(step_counter_d)) = + (self.gpu_log_tick_fn.as_ref(), self.log_step_counter_d.as_mut()) { let tick_cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1, 1, 1), shared_mem_bytes: 0, }; - let mut launch = self.stream.launch_builder(&self.gpu_log_tick_fn); - launch.arg(&mut self.log_step_counter_d); + let mut launch = self.stream.launch_builder(tick_fn); + launch.arg(step_counter_d); unsafe { launch.launch(tick_cfg).context("gpu_log_tick launch")?; } } @@ -2154,17 +2196,27 @@ impl PerceptionTrainer { block_dim: (N_HORIZONS as u32, 1, 1), shared_mem_bytes: 0, }; - // GPU log ring pointers. With feature off, both are 0 (null) and - // the kernel's `log_record()` short-circuits at the nullptr guard - // — behaviorally identical to the pre-log kernel. + // GPU log ring pointers. When the feature is OFF, both are 0 + // (null). When the feature is ON but the runtime trace path is + // None (no `--kernel-step-trace`), the ring + step counter + // Options are also None — pass null and the kernel's + // `log_record()` short-circuits at the nullptr guard, behaviorally + // identical to the pre-log kernel. #[cfg(feature = "kernel-step-trace")] - let log_ring_ptr: cudarc::driver::sys::CUdeviceptr = self.log_ring.dev_ptr(); + let log_ring_ptr: cudarc::driver::sys::CUdeviceptr = self + .log_ring + .as_ref() + .map(|r| r.dev_ptr()) + .unwrap_or(0); #[cfg(not(feature = "kernel-step-trace"))] let log_ring_ptr: cudarc::driver::sys::CUdeviceptr = 0; #[cfg(feature = "kernel-step-trace")] - let log_step_ptr: cudarc::driver::sys::CUdeviceptr = { - let (p, _g) = self.log_step_counter_d.device_ptr(&self.stream); - p + let log_step_ptr: cudarc::driver::sys::CUdeviceptr = match self.log_step_counter_d.as_ref() { + Some(slice) => { + let (p, _g) = slice.device_ptr(&self.stream); + p + } + None => 0, }; #[cfg(not(feature = "kernel-step-trace"))] let log_step_ptr: cudarc::driver::sys::CUdeviceptr = 0; @@ -2754,16 +2806,21 @@ impl PerceptionTrainer { // The drain task reads the host shadow at 500ms cadence to know // how many steps have completed since its last drain. Captured // inside the same graph as the other shadow memcpys, so a single - // end-of-step sync makes all of them visible. + // end-of-step sync makes all of them visible. Skipped when the + // ring isn't allocated (runtime trace path is None). #[cfg(feature = "kernel-step-trace")] - unsafe { - let (src_ptr, _g) = self.log_step_counter_d.device_ptr(&self.stream); - cudarc::driver::result::memcpy_dtod_async( - self.log_step_counter_host_shadow.dev_ptr, - src_ptr, - std::mem::size_of::(), - self.stream.cu_stream(), - ).context("log_step_counter_host_shadow shadow")?; + if let (Some(step_counter_d), Some(shadow)) = + (self.log_step_counter_d.as_ref(), self.log_step_counter_host_shadow.as_ref()) + { + unsafe { + let (src_ptr, _g) = step_counter_d.device_ptr(&self.stream); + cudarc::driver::result::memcpy_dtod_async( + shadow.dev_ptr, + src_ptr, + std::mem::size_of::(), + self.stream.cu_stream(), + ).context("log_step_counter_host_shadow shadow")?; + } } Ok(()) } @@ -4183,7 +4240,9 @@ fn download(stream: &Arc, src: &CudaSlice) -> Result> impl Drop for PerceptionTrainer { fn drop(&mut self) { use std::sync::atomic::Ordering; - self.log_drain_stop.store(true, Ordering::Relaxed); + if let Some(stop) = self.log_drain_stop.as_ref() { + stop.store(true, Ordering::Relaxed); + } if let Some(h) = self.log_drain_handle.take() { h.abort(); } diff --git a/crates/ml-alpha/tests/gpu_log_ring_invariants.rs b/crates/ml-alpha/tests/gpu_log_ring_invariants.rs index 49eb195e3..188b6adea 100644 --- a/crates/ml-alpha/tests/gpu_log_ring_invariants.rs +++ b/crates/ml-alpha/tests/gpu_log_ring_invariants.rs @@ -17,11 +17,12 @@ use std::sync::atomic::{AtomicBool, Ordering}; use std::time::Duration; use ml_alpha::gpu_log::{ - spawn_drain_task, LogHeader, LogRecord, LogRing, KID_SMOOTHNESS_CONTROLLER, MAGIC, + spawn_drain_task_to_file, LogHeader, LogRecord, LogRing, KID_SMOOTHNESS_CONTROLLER, MAGIC, N_RECORDS_PER_STEP, N_RT_PER_KERNEL, N_SLOTS, PAYLOAD_F32, RT_INPUT, RT_OUTPUT, RT_STATE, }; use ml_alpha::pinned_mem::MappedRecordBuffer; use ml_core::device::MlDevice; +use tempfile::NamedTempFile; // ── helpers ────────────────────────────────────────────────────────── @@ -147,15 +148,29 @@ async fn magic_validation_skips_torn() { counter.write_record(0, 1); - let h = spawn_drain_task(ring.clone(), counter.clone(), stop.clone()); + let tmp = NamedTempFile::new().expect("tempfile"); + let path = tmp.path().to_path_buf(); + 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; 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; - // Pass if no panic — torn slot was skipped, mismatch slot was skipped. - // (We don't intercept tracing output here; the absence of panic is the - // pass condition — drainer must defensively handle both cases.) + + // The valid record should produce exactly one JSONL line; the torn + // (magic=0) and mismatched-step records must NOT appear. + let contents = std::fs::read_to_string(&path).expect("read trace file"); + let lines: Vec<&str> = contents.lines().filter(|l| !l.is_empty()).collect(); + assert_eq!(lines.len(), 1, "expected exactly 1 valid record, got {}", lines.len()); + let v: serde_json::Value = serde_json::from_str(lines[0]).expect("parse JSONL"); + assert_eq!(v["step"], 1); + assert_eq!(v["kid"], KID_SMOOTHNESS_CONTROLLER); + assert_eq!(v["rt"], RT_INPUT); + assert_eq!(v["kname"], "smoothness_controller"); + assert_eq!(v["rt_name"], "input"); + // Payload field comes from the smoothness_input schema. + assert!((v["payload"]["raw_h30"].as_f64().unwrap() - 100.0).abs() < 1e-6); } #[test] @@ -199,7 +214,14 @@ async fn step_gap_detection_emits_warn() { let big_jump = ((N_SLOTS / N_RECORDS_PER_STEP) + 10) as i32; counter.write_record(0, big_jump); - let h = spawn_drain_task(ring.clone(), counter.clone(), stop.clone()); + let tmp = NamedTempFile::new().expect("tempfile"); + let h = spawn_drain_task_to_file( + ring.clone(), + counter.clone(), + stop.clone(), + tmp.path().to_path_buf(), + ) + .expect("spawn drain"); tokio::time::sleep(Duration::from_millis(700)).await; stop.store(true, Ordering::Relaxed); let _ = tokio::time::timeout(Duration::from_secs(2), h).await; @@ -226,17 +248,27 @@ async fn decoder_dispatch_unknown_falls_through() { ring.buffer.write_record(slot, rec); counter.write_record(0, 1); - let h = spawn_drain_task(ring.clone(), counter.clone(), stop.clone()); + let tmp = NamedTempFile::new().expect("tempfile"); + let h = spawn_drain_task_to_file( + ring.clone(), + counter.clone(), + stop.clone(), + tmp.path().to_path_buf(), + ) + .expect("spawn drain"); tokio::time::sleep(Duration::from_millis(700)).await; stop.store(true, Ordering::Relaxed); let _ = tokio::time::timeout(Duration::from_secs(2), h).await; // 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()` function - // itself MUST handle unknown kernel_id without panic. Exercise it directly: - ml_alpha::gpu_log::decode(unknown_kid, RT_INPUT, 1, &[7.0; 10]); - ml_alpha::gpu_log::decode(0, 99, 1, &[7.0; 10]); // unknown record_type - // Pass if no panic. + // never be visited by the drainer's slot walk. But the `decode_to_json()` + // function itself MUST handle unknown kernel_id without panic. Exercise + // it directly — unknown (kid, rt) pairs fall through to a "values" array. + let v1 = ml_alpha::gpu_log::decode_to_json(unknown_kid, RT_INPUT, 1, &[7.0; 10]); + assert_eq!(v1["kname"], "unknown"); + assert!(v1["payload"]["values"].is_array()); + let v2 = ml_alpha::gpu_log::decode_to_json(0, 99, 1, &[7.0; 10]); // unknown record_type + assert_eq!(v2["rt_name"], "unknown"); } // ── CUDA integration test #7 ───────────────────────────────────────── diff --git a/crates/ml-alpha/tests/perception_overfit.rs b/crates/ml-alpha/tests/perception_overfit.rs index a79b98fcd..f72cae6dc 100644 --- a/crates/ml-alpha/tests/perception_overfit.rs +++ b/crates/ml-alpha/tests/perception_overfit.rs @@ -73,6 +73,7 @@ fn stacked_trainer_loss_shrinks_on_constant_signal() { horizon_weights: [1.0; 5], n_batch: 1, smoothness_base_lambda: 0.0, + kernel_step_trace_path: None, }; let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init"); @@ -147,6 +148,7 @@ fn stacked_trainer_loss_shrinks_with_stride_4() { horizon_weights: [1.0; 5], n_batch: 1, smoothness_base_lambda: 0.0, + kernel_step_trace_path: None, }; let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init"); @@ -201,6 +203,7 @@ fn stacked_trainer_loss_shrinks_at_batch_32() { horizon_weights: [1.0; 5], n_batch: 32, smoothness_base_lambda: 0.0, + kernel_step_trace_path: None, }; let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init"); @@ -299,6 +302,7 @@ fn evaluate_alone_succeeds() { horizon_weights: [1.0; 5], n_batch: 1, smoothness_base_lambda: 0.0, + kernel_step_trace_path: None, }; let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init"); let ts = 1_000_000u64; @@ -327,6 +331,7 @@ fn evaluate_works_after_captured_training_step() { horizon_weights: [1.0; 5], n_batch: 1, smoothness_base_lambda: 0.0, + kernel_step_trace_path: None, }; let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init"); @@ -362,6 +367,7 @@ fn evaluate_works_after_capture_no_replay() { horizon_weights: [1.0; 5], n_batch: 1, smoothness_base_lambda: 0.0, + kernel_step_trace_path: None, }; let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init"); let ts = 1_000_000u64; @@ -392,6 +398,7 @@ fn horizon_ema_and_lambda_track_after_training() { horizon_weights: [1.0; 5], n_batch: 1, smoothness_base_lambda: 0.0, + kernel_step_trace_path: None, }; let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init"); @@ -448,6 +455,7 @@ fn evaluate_works_after_warmup_only() { horizon_weights: [1.0; 5], n_batch: 1, smoothness_base_lambda: 0.0, + kernel_step_trace_path: None, }; let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init"); let ts = 1_000_000u64; diff --git a/infra/k8s/argo/alpha-perception-template.yaml b/infra/k8s/argo/alpha-perception-template.yaml index 1cf9f0ddd..a4221af0a 100644 --- a/infra/k8s/argo/alpha-perception-template.yaml +++ b/infra/k8s/argo/alpha-perception-template.yaml @@ -80,6 +80,10 @@ spec: value: "1" - name: smoothness-base-lambda value: "0.0" + - name: kernel-step-trace-enable + value: "false" + - name: kernel-step-trace-path + value: "" volumes: - name: git-ssh-key @@ -299,8 +303,20 @@ spec: fi export PATH="${CARGO_HOME}/bin:${PATH}" + + FEATURES_FLAG="" + if [ "{{workflow.parameters.kernel-step-trace-enable}}" = "true" ]; then + FEATURES_FLAG="--features kernel-step-trace" + echo "Building with --features kernel-step-trace" + # Bust cache: feature-on binary differs from feature-off at same SHA. + if [ -x "$BIN_DIR/alpha_train" ]; then + echo "Removing cached binary (feature change requires rebuild)" + rm -f "$BIN_DIR/alpha_train" + fi + fi + echo "Building alpha_train (ml-alpha)..." - cargo build --release -p ml-alpha --example alpha_train + cargo build --release $FEATURES_FLAG -p ml-alpha --example alpha_train mkdir -p "$BIN_DIR" cp "$CARGO_TARGET_DIR/release/examples/alpha_train" "$BIN_DIR/" @@ -432,6 +448,11 @@ spec: EXTRA_FLAGS="$EXTRA_FLAGS --auto-horizon-weights" fi + TRACE_FLAG="" + if [ -n "{{workflow.parameters.kernel-step-trace-path}}" ]; then + TRACE_FLAG="--kernel-step-trace {{workflow.parameters.kernel-step-trace-path}}" + fi + "$BIN" \ --mbp10-data-dir "$MBP10_DIR" \ --predecoded-dir "$PREDECODED_DIR" \ @@ -452,6 +473,7 @@ spec: --cv-train-window {{workflow.parameters.cv-train-window}} \ --decision-stride {{workflow.parameters.decision-stride}} \ --smoothness-base-lambda {{workflow.parameters.smoothness-base-lambda}} \ + ${TRACE_FLAG} \ $EXTRA_FLAGS echo "=== Training complete ==="