From 1eeb1067d46c038eed46f2d3473582cb73bae73d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Jul 2026 14:21:42 +0000 Subject: [PATCH] diffusion trainer: emit learning_rate in progress events (Studio pump compatibility) The Studio training pump reads 'learning_rate' from progress events; the diffusion trainer emitted 'lr'. Rename the field (and the CLI reader) so the trainer's events are directly consumable by the existing training status/SSE machinery when it is wired into the worker, without a translation shim. --- studio/backend/core/training/diffusion_lora_trainer.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index 5c52f75483..28e0e5e65e 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -384,6 +384,9 @@ def run_diffusion_lora_training( running_loss += step_loss done = opt_step + 1 if done % cfg.log_every == 0 or done == cfg.train_steps: + # ``learning_rate`` (not ``lr``) is the field the Studio training pump reads, so + # these progress events are directly consumable by the existing training + # status/SSE machinery when the diffusion trainer is wired into the worker. _emit( on_event, "progress", @@ -391,7 +394,7 @@ def run_diffusion_lora_training( total_steps=cfg.train_steps, loss=round(step_loss, 5), avg_loss=round(running_loss / done, 5), - lr=lr_sched.get_last_lr()[0], + learning_rate=lr_sched.get_last_lr()[0], ) if should_stop is not None and should_stop(): @@ -502,7 +505,7 @@ def main(argv: Optional[list[str]] = None) -> int: if t == "progress": print( f"step {ev['step']}/{ev['total_steps']} loss={ev['loss']} " - f"avg={ev['avg_loss']} lr={ev['lr']:.2e}", + f"avg={ev['avg_loss']} lr={ev['learning_rate']:.2e}", flush=True, ) elif t == "complete":