Fixing stuck training processes

studio/backend/core/training/trainer.py

@@ -2,13 +2,16 @@
 Unsloth Training Backend
 Integrates Unsloth training capabilities with the FastAPI backend
 """
+import os
+# Prevent tokenizer parallelism deadlocks when datasets uses multiprocessing fork
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
 import torch
 from utils.hardware import clear_gpu_cache
 torch._dynamo.config.recompile_limit = 64
 from unsloth import FastLanguageModel, FastVisionModel, is_bfloat16_supported
 from unsloth.chat_templates import get_chat_template
 
-import os
 import json
 import threading
 import math

studio/backend/core/training/training.py

@@ -6,6 +6,7 @@ from typing import Any, Generator, Tuple
 import logging
 
 from .trainer import get_trainer, TrainingProgress
+from utils.hardware import clear_gpu_cache
 
 logger = logging.getLogger(__name__)
 
@@ -101,6 +102,31 @@ class TrainingBackend:
             True if training started successfully, False otherwise.
         """
         try:
+            # Wait for any previous training thread to finish
+            old_thread = getattr(self.trainer, "training_thread", None)
+            if old_thread and old_thread.is_alive():
+                logger.info("Waiting for previous training thread to finish...")
+                old_thread.join(timeout=30)
+
+            # Explicitly free old SFTTrainer and CUDA resources before loading new model.
+            # Without this, forked multiprocessing workers (num_proc tokenization) inherit
+            # stale CUDA state from the previous run, causing extreme slowdowns or crashes.
+            if self.trainer.trainer is not None:
+                logger.info("Cleaning up previous SFTTrainer...")
+                self.trainer.trainer = None
+            if self.trainer.model is not None:
+                self.trainer.model = None
+            if self.trainer.tokenizer is not None:
+                self.trainer.tokenizer = None
+            # Flush all pending async CUDA ops so forked tokenization processes
+            # don't inherit stale async state that causes pool join to hang.
+            import torch as _torch
+            if _torch.cuda.is_available():
+                _torch.cuda.synchronize()
+            import gc
+            gc.collect()
+            clear_gpu_cache()
+
             # Reset stop flag and clear history
             self.trainer.should_stop = False
             self.trainer.save_on_stop = True