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