fix: unload competing subprocesses before load across all routes

studio/backend/routes/export.py

@@ -63,6 +63,32 @@ async def load_checkpoint(
     try:
         # Version switching is handled automatically by the subprocess-based
         # export backend — no need for ensure_transformers_version() here.
+
+        # Free GPU memory: shut down any running inference/training subprocesses
+        # before loading the export checkpoint (they'd compete for VRAM).
+        try:
+            from core.inference import get_inference_backend
+            inf = get_inference_backend()
+            if inf.active_model_name:
+                logger.info(
+                    "Unloading inference model '%s' to free GPU memory for export",
+                    inf.active_model_name,
+                )
+                inf._shutdown_subprocess()
+                inf.active_model_name = None
+                inf.models.clear()
+        except Exception as e:
+            logger.warning("Could not unload inference model: %s", e)
+
+        try:
+            from core.training import get_training_backend
+            trn = get_training_backend()
+            if trn.is_training_active():
+                logger.info("Stopping active training to free GPU memory for export")
+                trn.stop_training()
+        except Exception as e:
+            logger.warning("Could not stop training: %s", e)
+
         backend = get_export_backend()
         success, message = backend.load_checkpoint(
             checkpoint_path=request.checkpoint_path,

studio/backend/routes/inference.py

@@ -163,6 +163,19 @@ async def load_model(
             logger.info("Unloading GGUF model before loading Unsloth model")
             llama_backend.unload_model()
 
+        # Shut down any export subprocess to free VRAM
+        try:
+            from core.export import get_export_backend
+            exp_backend = get_export_backend()
+            if exp_backend.current_checkpoint:
+                logger.info("Shutting down export subprocess to free GPU memory for inference")
+                exp_backend._shutdown_subprocess()
+                exp_backend.current_checkpoint = None
+                exp_backend.is_vision = False
+                exp_backend.is_peft = False
+        except Exception as e:
+            logger.warning("Could not shut down export subprocess: %s", e)
+
         # Auto-detect quantization for LoRA adapters from adapter_config.json
         # The training pipeline patches this file with "unsloth_training_method"
         # which is 'qlora' or 'lora'. Only LoRA (16-bit) needs load_in_4bit=False.

studio/backend/routes/training.py

@@ -192,7 +192,7 @@ async def start_training(
             "tensorboard_dir": request.tensorboard_dir or "",
         }
 
-        # Free GPU memory: shut down any running inference subprocess
+        # Free GPU memory: shut down any running inference/export subprocesses
         # before training starts (they'd compete for VRAM otherwise)
         try:
             from core.inference import get_inference_backend
@@ -208,6 +208,18 @@ async def start_training(
         except Exception as e:
             logger.warning("Could not unload inference model: %s", e)
 
+        try:
+            from core.export import get_export_backend
+            exp_backend = get_export_backend()
+            if exp_backend.current_checkpoint:
+                logger.info("Shutting down export subprocess to free GPU memory for training")
+                exp_backend._shutdown_subprocess()
+                exp_backend.current_checkpoint = None
+                exp_backend.is_vision = False
+                exp_backend.is_peft = False
+        except Exception as e:
+            logger.warning("Could not shut down export subprocess: %s", e)
+
         # start_training now spawns a subprocess (non-blocking)
         success = backend.start_training(**training_kwargs)