From 25b51fad3be828cd519ce47250df59332e4d2d89 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Sat, 7 Mar 2026 04:16:10 +0000 Subject: [PATCH] fix: wait for training shutdown before export load, clear stop flag on reset MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 1. Export route: stop_training() only signals the subprocess — wait up to 30s for it to actually exit before loading the export checkpoint, avoiding a GPU memory race. 2. Training reset: clear _should_stop so /api/train/status returns phase=idle instead of staying stuck on phase=stopped after a user-triggered stop. Co-Authored-By: Claude Opus 4.6 --- studio/backend/routes/export.py | 8 ++++++++ studio/backend/routes/training.py | 1 + 2 files changed, 9 insertions(+) diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index 6b335bc139..7f18f50eaa 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -86,6 +86,14 @@ async def load_checkpoint( if trn.is_training_active(): logger.info("Stopping active training to free GPU memory for export") trn.stop_training() + # Wait for training subprocess to actually exit before proceeding, + # otherwise it may still hold GPU memory when export tries to load. + for _ in range(60): # up to 30s + if not trn.is_training_active(): + break + import time; time.sleep(0.5) + else: + logger.warning("Training subprocess did not exit within 30s, proceeding anyway") except Exception as e: logger.warning("Could not stop training: %s", e) diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index c0ee6deff7..84315291a9 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -304,6 +304,7 @@ async def reset_training( ) logger.info("Reset training state: clearing runtime + metric history") + backend._should_stop = False # Clear stop flag so status returns to idle backend.trainer._update_progress( is_training=False, is_completed=False, error=None, status_message="Ready to train", step=0, loss=0.0, epoch=0,