From 8c493bbd20cacd98166d84e19582a91567e7e70c Mon Sep 17 00:00:00 2001 From: Manan17 Date: Mon, 9 Mar 2026 07:00:06 +0000 Subject: [PATCH 1/3] CLI fix for backend changes --- cli/commands/export.py | 7 +++++-- cli/commands/train.py | 10 ++++++---- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/cli/commands/export.py b/cli/commands/export.py index ba487aeff3..71aa95b7f0 100644 --- a/cli/commands/export.py +++ b/cli/commands/export.py @@ -25,8 +25,11 @@ def list_checkpoints( typer.echo("No checkpoints found.") raise typer.Exit() - for display, path in checkpoints: - typer.echo(f"{display}: {path}") + for model_name, ckpt_list, metadata in checkpoints: + typer.echo(f"\n{model_name}:") + for display, path, loss in ckpt_list: + loss_str = f" (loss: {loss:.4f})" if loss is not None else "" + typer.echo(f" {display}{loss_str}: {path}") def export( diff --git a/cli/commands/train.py b/cli/commands/train.py index 2e78db7375..be124407b3 100644 --- a/cli/commands/train.py +++ b/cli/commands/train.py @@ -81,7 +81,7 @@ def train( ) raise typer.Exit(code=2) - from studio.backend.core.training import UnslothTrainer + from studio.backend.core.training.trainer import UnslothTrainer trainer = UnslothTrainer() @@ -101,18 +101,20 @@ def train( typer.echo("Model preparation failed", err=True) raise typer.Exit(code=1) - ds = trainer.load_and_format_dataset( + result = trainer.load_and_format_dataset( dataset_source=cfg.data.dataset or "", format_type=cfg.data.format_type, local_datasets=cfg.data.local_dataset, ) - if ds is None: + if result is None: typer.echo("Dataset load failed", err=True) raise typer.Exit(code=1) + ds, eval_ds = result + training_kwargs = cfg.training_kwargs() training_kwargs["wandb_token"] = wandb_token # CLI/env takes precedence - started = trainer.start_training(dataset=ds, **training_kwargs) + started = trainer.start_training(dataset=ds, eval_dataset=eval_ds, **training_kwargs) if not started: typer.echo("Training failed to start", err=True) From 9be55f0c1bb6c7c8e51df303982bdefb4948e9cd Mon Sep 17 00:00:00 2001 From: Manan17 Date: Tue, 10 Mar 2026 02:20:56 +0000 Subject: [PATCH 2/3] fixing cancel training --- studio/backend/core/training/training.py | 14 ++++++++++++++ studio/backend/routes/training.py | 15 ++++++++++----- .../features/studio/training-start-overlay.tsx | 8 ++++++-- .../training/hooks/use-training-actions.ts | 6 ++++++ 4 files changed, 36 insertions(+), 7 deletions(-) diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 101ba13621..47ed0c88cf 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -226,6 +226,20 @@ class TrainingBackend: ) return True + def force_terminate(self) -> None: + """Force-kill the training subprocess so state can be reset immediately.""" + with self._lock: + if self._proc is not None and self._proc.is_alive(): + logger.info("Force-terminating training subprocess (pid=%s)", self._proc.pid) + self._proc.terminate() + proc = self._proc + + if proc is not None: + proc.join(timeout=5.0) + if proc.is_alive(): + proc.kill() + proc.join(timeout=2.0) + def is_training_active(self) -> bool: """Check if training is currently active.""" with self._lock: diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 45fbff0172..1590c0864d 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -309,11 +309,16 @@ async def reset_training( is_active = backend.is_training_active() if is_active: - logger.warning("Rejected reset while training active: is_active=%s", is_active) - raise HTTPException( - status_code=409, - detail="Training is still running. Stop training and wait for it to finish before resetting.", - ) + if backend._should_stop: + # Stop was already requested — force-terminate so we can reset immediately + logger.info("Force-terminating subprocess for immediate reset") + backend.force_terminate() + else: + logger.warning("Rejected reset while training active: is_active=%s", is_active) + raise HTTPException( + status_code=409, + detail="Training is still running. Stop training and wait for it to finish before resetting.", + ) logger.info("Reset training state: clearing runtime + metric history") backend._should_stop = False # Clear stop flag so status returns to idle diff --git a/studio/frontend/src/features/studio/training-start-overlay.tsx b/studio/frontend/src/features/studio/training-start-overlay.tsx index e2d660343a..898e47297a 100644 --- a/studio/frontend/src/features/studio/training-start-overlay.tsx +++ b/studio/frontend/src/features/studio/training-start-overlay.tsx @@ -31,7 +31,7 @@ export function TrainingStartOverlay({ message, currentStep, }: TrainingStartOverlayProps): ReactElement { - const { stopTrainingRun } = useTrainingActions(); + const { stopTrainingRun, dismissTrainingRun } = useTrainingActions(); const isStarting = useTrainingRuntimeStore((s) => s.isStarting); const [cancelDialogOpen, setCancelDialogOpen] = useState(false); const [cancelRequested, setCancelRequested] = useState(false); @@ -77,7 +77,11 @@ export function TrainingStartOverlay({ setCancelDialogOpen(false); useTrainingRuntimeStore.getState().setStopRequested(true); void stopTrainingRun(false).then((ok) => { - if (!ok) setCancelRequested(false); + if (ok) { + void dismissTrainingRun(); + } else { + setCancelRequested(false); + } }); }} > diff --git a/studio/frontend/src/features/training/hooks/use-training-actions.ts b/studio/frontend/src/features/training/hooks/use-training-actions.ts index b748daa907..773c8f267b 100644 --- a/studio/frontend/src/features/training/hooks/use-training-actions.ts +++ b/studio/frontend/src/features/training/hooks/use-training-actions.ts @@ -105,6 +105,12 @@ export function useTrainingActions() { } } + // Abort if cancel was requested during dataset check + if (useTrainingRuntimeStore.getState().stopRequested) { + runtimeStore.setStarting(false); + return false; + } + // Re-read config after potential store updates from dataset check const payload = buildTrainingStartPayload(useTrainingConfigStore.getState()); const response = await startTraining(payload); From fd7ca8bda8116805b2de18273f30adbb11dfcb84 Mon Sep 17 00:00:00 2001 From: Manan17 Date: Tue, 10 Mar 2026 02:35:32 +0000 Subject: [PATCH 3/3] distinguish cancel and stop for force terminate --- studio/backend/core/training/training.py | 4 ++++ studio/backend/routes/training.py | 6 +++--- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 47ed0c88cf..14ff57bbd7 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -71,6 +71,7 @@ class TrainingBackend: # Progress state (updated by pump thread from subprocess events) self._progress = TrainingProgress() self._should_stop = False + self._cancel_requested = False # True only for stop(save=False) # Training Metrics (consumed by routes for SSE and /metrics) self.loss_history: list = [] @@ -114,6 +115,7 @@ class TrainingBackend: # Reset state self._should_stop = False + self._cancel_requested = False self._progress = TrainingProgress(is_training=True, status_message="Initializing training...") self.loss_history.clear() self.lr_history.clear() @@ -213,6 +215,8 @@ class TrainingBackend: def stop_training(self, save: bool = True) -> bool: """Send stop signal to the training subprocess.""" self._should_stop = True + if not save: + self._cancel_requested = True with self._lock: if self._stop_queue is not None: try: diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 1590c0864d..6de29c4882 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -309,9 +309,9 @@ async def reset_training( is_active = backend.is_training_active() if is_active: - if backend._should_stop: - # Stop was already requested — force-terminate so we can reset immediately - logger.info("Force-terminating subprocess for immediate reset") + if backend._cancel_requested: + # Cancel (save=False) was requested — force-terminate so we can reset immediately + logger.info("Force-terminating subprocess for immediate reset (cancel path)") backend.force_terminate() else: logger.warning("Rejected reset while training active: is_active=%s", is_active)