From c3cdb39081fdf40c51fadf3e8df25f7687eb5aa5 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 18 Apr 2026 13:15:25 +0200 Subject: [PATCH] fix: acquire write lock once per epoch, not per training step MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The async write lock (self.agent.write().await) was acquired inside the per-step training loop, causing 696s/epoch overhead from async yield on every iteration — 178 steps × ~4s per lock acquisition. Actual GPU compute was only 3.6s (graph replay at 2ms/step). Fix: acquire the lock once before the loop. The training loop is the only writer during this phase — no contention, no need to release between steps. Expected: epoch time drops from 700s to ~5s. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml/src/trainers/dqn/trainer/training_loop.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 16a65aaa2..e38189078 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -1473,10 +1473,10 @@ impl DQNTrainer { } } + // Acquire write lock ONCE for the entire training loop — not per-step. + // Per-step lock acquisition was burning 696s/epoch due to async yield overhead. + let mut agent = self.agent.write().await; for _step in 0..num_training_steps { - // Single write lock for both sample + fused step — eliminates - // read→write lock cycling overhead (~43ms per step from contention). - let mut agent = self.agent.write().await; let sample_start = std::time::Instant::now(); {