diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 52cf60af10..013c7817b2 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -58,6 +58,7 @@ class UnslothTrainer: self.progress_callbacks = [] self.is_training = False self.should_stop = False + self.save_on_stop = True # Model state tracking self.is_vlm = False @@ -756,16 +757,32 @@ class UnslothTrainer: self.trainer.train() # ========== SAVE MODEL ========== - self.trainer.save_model() - self.tokenizer.save_pretrained(output_dir) - print(f"\nTraining completed! Model saved to {output_dir}\n") - - self._update_progress( - is_training=False, - is_completed=True, - #status_message=status_msg - status_message=f"Training completed! Model saved to {output_dir}", - ) + if self.should_stop and self.save_on_stop: + # Stopped by user — save model at current checkpoint + self.trainer.save_model() + self.tokenizer.save_pretrained(output_dir) + print(f"\nTraining stopped. Model saved to {output_dir}\n") + self._update_progress( + is_training=False, + status_message=f"Training stopped. Model saved to {output_dir}", + ) + elif self.should_stop: + # Cancelled by user — don't save + print("\nTraining cancelled.\n") + self._update_progress( + is_training=False, + status_message="Training cancelled.", + ) + else: + # Normal completion + self.trainer.save_model() + self.tokenizer.save_pretrained(output_dir) + print(f"\nTraining completed! Model saved to {output_dir}\n") + self._update_progress( + is_training=False, + is_completed=True, + status_message=f"Training completed! Model saved to {output_dir}", + ) except Exception as e: logger.error(f"Training error: {e}") @@ -774,10 +791,11 @@ class UnslothTrainer: finally: self.is_training = False - def stop_training(self): + def stop_training(self, save: bool = True): """Stop ongoing training""" - print("\nStopping training...") + print(f"\nStopping training (save={save})...") self.should_stop = True + self.save_on_stop = save self.is_training = False # Clear the status message so timer doesn't show stale status self._update_progress(is_training=False, status_message="") diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index fa8b3b5daf..62febadc13 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -103,6 +103,7 @@ class TrainingBackend: try: # Reset stop flag and clear history self.trainer.should_stop = False + self.trainer.save_on_stop = True self.loss_history = [] self.lr_history = [] self.step_history = [] @@ -224,16 +225,19 @@ class TrainingBackend: ) return False - def stop_training(self) -> bool: + def stop_training(self, save: bool = True) -> bool: """ Stop ongoing training. + Args: + save: If True, save the model at the current checkpoint. + Returns: True if training was successfully stopped. """ try: - logger.info("Stopping training...") - self.trainer.stop_training() + logger.info(f"Stopping training (save={save})...") + self.trainer.stop_training(save=save) return True except Exception as e: logger.error(f"Error stopping training: {e}") @@ -293,6 +297,10 @@ class TrainingBackend: True if training is in progress, False otherwise """ try: + # If user requested stop, training is no longer considered active + if self.trainer.should_stop: + return False + progress = self.trainer.get_training_progress() # Training is active if is_training is True # Also check if we're in loading/preparation phase (status_message indicates activity) diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 20a5e1e254..1a904cd371 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -37,6 +37,11 @@ from models import ( TrainingProgress, ) from models.responses import TrainingStopResponse, TrainingMetricsResponse +from pydantic import BaseModel as PydanticBaseModel + + +class TrainingStopRequest(PydanticBaseModel): + save: bool = True router = APIRouter() logger = logging.getLogger(__name__) @@ -251,10 +256,14 @@ async def start_training( @router.post("/stop", response_model=TrainingStopResponse) async def stop_training( + body: TrainingStopRequest = TrainingStopRequest(), current_subject: str = Depends(get_current_subject), ): """ Stop the currently running training job. + + Body: + save (bool): If True (default), save the model at the current checkpoint. """ try: backend = get_training_backend() @@ -266,7 +275,7 @@ async def stop_training( ) # Call backend stop method - backend.stop_training() + backend.stop_training(save=body.save) return TrainingStopResponse( status="stopped", @@ -281,6 +290,29 @@ async def stop_training( ) +@router.post("/reset") +async def reset_training( + current_subject: str = Depends(get_current_subject), +): + """ + Reset training state so the user can return to configuration. + """ + try: + backend = get_training_backend() + backend.trainer.should_stop = False + backend.trainer.training_progress = backend.trainer.training_progress.__class__() + backend.loss_history = [] + backend.lr_history = [] + backend.step_history = [] + return {"status": "ok"} + except Exception as e: + logger.error(f"Error resetting training: {e}", exc_info=True) + raise HTTPException( + status_code=500, + detail=f"Failed to reset training: {str(e)}", + ) + + @router.get("/status") async def get_training_status( current_subject: str = Depends(get_current_subject), @@ -313,6 +345,9 @@ async def get_training_status( ) or "Ready to train" error_message = getattr(progress, "error", None) if progress else None + # Check if training was stopped by user + trainer_stopped = getattr(backend.trainer, "should_stop", False) + # Derive high-level phase if error_message: phase = "error" @@ -326,6 +361,8 @@ async def get_training_status( phase = "configuring" else: phase = "training" + elif trainer_stopped: + phase = "stopped" elif progress and getattr(progress, "is_completed", False): phase = "completed" elif has_thread: