From 667e61734ca9b2dbf3611c3e6dfaf47188cef7f0 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 20 Mar 2026 12:19:56 +0100 Subject: [PATCH] fix: scope event tracking disable to CUDA Graph capture only disable_event_tracking() was called globally in GpuDqnTrainer::new(), polluting the CUDA context for ALL subsequent operations including other models (TFT, PPO). This caused test_tft_adapter_deterministic to fail when run after DQN tests. Fix: disable/enable only inside capture_training_graph(), not at construction. The capture region is the only place where CudaEvents are disallowed. Result: 855/855 ml lib tests pass (was 854/855 with 1 flaky failure). Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 453ff8054..9dcd08375 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -403,9 +403,9 @@ impl GpuDqnTrainer { let b = config.batch_size; let total_params = compute_total_params(&config); - // Disable event tracking — required for CUDA Graph capture (no CudaEvents - // inside captured regions). Safe: single stream, explicit ordering. - unsafe { stream.context().disable_event_tracking(); } + // Event tracking: kept ENABLED during normal ops (buffer allocation, DtoD). + // DISABLED only during CUDA Graph capture (in capture_training_graph). + // Global disable pollutes the context for other tests/models. // Stack size for the training kernel: ~46KB/thread with NUM_ATOMS=51, // ~12KB/thread with NUM_ATOMS=11. DIST_SIZE(SCRATCH1_DIM=256) * 4 = 11KB