ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: add cancel or save and stop training

Browse source
This commit is contained in:
Manan17 2026-02-15 00:00:22 +00:00
commit 97c6a09b84
3 changed files with 79 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

42
studio/backend/core/training/trainer.py
View file

@ -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="")

14
studio/backend/core/training/training.py
View file

@ -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)

39
studio/backend/routes/training.py
View file

@ -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: