From dbbcdb4f09f62fa62d406c2527cd6310ba1f9039 Mon Sep 17 00:00:00 2001
From: Roland Tannous <rolandtannous@gravityq.ai>
Date: Mon, 23 Feb 2026 07:26:22 +0000
Subject: [PATCH] feat: clear unsloth_compiled_cache on startup, shutdown, and
 between model loads

---
 setup.sh                                   |  5 ++++
 studio/backend/core/inference/inference.py |  4 +++
 studio/backend/core/training/trainer.py    |  4 +++
 studio/backend/main.py                     |  8 +++---
 studio/backend/utils/cache_cleanup.py      | 30 ++++++++++++++++++++++
 5 files changed, 48 insertions(+), 3 deletions(-)
 create mode 100644 studio/backend/utils/cache_cleanup.py

diff --git a/setup.sh b/setup.sh
index b54f782ded..1f55e141fd 100755
--- a/setup.sh
+++ b/setup.sh
@@ -24,6 +24,11 @@ echo "╔═══════════════════════
 echo "║     Unsloth Studio Setup Script      ║"
 echo "╚══════════════════════════════════════╝"
 
+# ── Clean up stale Unsloth compiled caches ──
+rm -rf "$SCRIPT_DIR/unsloth_compiled_cache"
+rm -rf "$SCRIPT_DIR/studio/backend/unsloth_compiled_cache"
+rm -rf "$SCRIPT_DIR/studio/tmp/unsloth_compiled_cache"
+
 # ── Detect Colab (like unsloth does) ──
 IS_COLAB=false
 keynames=$'\n'$(printenv | cut -d= -f1)
diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py
index 6423e7a128..e90e6c0c2a 100644
--- a/studio/backend/core/inference/inference.py
+++ b/studio/backend/core/inference/inference.py
@@ -184,6 +184,10 @@ class InferenceBackend:
                 # Clear GPU memory cache
                 clear_gpu_cache()
 
+                # Remove stale compiled cache so the next model gets a fresh one
+                from utils.cache_cleanup import clear_unsloth_compiled_cache
+                clear_unsloth_compiled_cache()
+
                 logger.info(f"Model '{model_name}' successfully unloaded.")
                 return True
             except Exception as e:
diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py
index d78c429e5a..e996603ae6 100644
--- a/studio/backend/core/training/trainer.py
+++ b/studio/backend/core/training/trainer.py
@@ -126,6 +126,10 @@ class UnslothTrainer:
             print("\nClearing GPU memory before training...")
             clear_gpu_cache()
 
+            # Remove stale compiled cache so the new model gets a fresh one
+            from utils.cache_cleanup import clear_unsloth_compiled_cache
+            clear_unsloth_compiled_cache()
+
             # Detect if this is a vision model AND dataset is multimodal
             # A vision-capable model with a text-only dataset should use FastLanguageModel
             self.is_vlm = is_vision_model(model_name) and is_dataset_multimodal
diff --git a/studio/backend/main.py b/studio/backend/main.py
index e7a33750ff..5ccd6cec12 100644
--- a/studio/backend/main.py
+++ b/studio/backend/main.py
@@ -3,7 +3,6 @@ Main FastAPI application for Unsloth UI Backend
 """
 import os
 import secrets
-import shutil
 from contextlib import asynccontextmanager
 
 from fastapi import FastAPI
@@ -19,12 +18,15 @@ from auth import storage
 from utils.hardware import detect_hardware, get_device, DeviceType
 import utils.hardware.hardware as _hw_module
 
-UNSLOTH_CACHE_DIR = Path(__file__).parent / "unsloth_compiled_cache"
+from utils.cache_cleanup import clear_unsloth_compiled_cache
 
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     """Startup: detect hardware, print setup token if needed. Shutdown: clean up compiled cache."""
+    # Clean up any stale compiled cache from previous runs
+    clear_unsloth_compiled_cache()
+
     # Detect hardware first — sets DEVICE global used everywhere
     detect_hardware()
 
@@ -52,7 +54,7 @@ async def lifespan(app: FastAPI):
     yield
     # Cleanup
     _hw_module.DEVICE = None
-    shutil.rmtree(UNSLOTH_CACHE_DIR, ignore_errors=True)
+    clear_unsloth_compiled_cache()
 
 
 # Create FastAPI app
diff --git a/studio/backend/utils/cache_cleanup.py b/studio/backend/utils/cache_cleanup.py
new file mode 100644
index 0000000000..4673cf995f
--- /dev/null
+++ b/studio/backend/utils/cache_cleanup.py
@@ -0,0 +1,30 @@
+"""
+Utility for cleaning up the Unsloth compiled cache directory.
+
+The unsloth_compiled_cache is created by unsloth_zoo/compiler.py during
+FastModel.from_pretrained() and contains model-type-specific compiled Python
+files. It should be cleared between model loads to avoid stale artefacts.
+"""
+import shutil
+import logging
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+# Possible locations where unsloth_compiled_cache may appear
+_BACKEND_DIR = Path(__file__).resolve().parent.parent          # studio/backend
+_PROJECT_ROOT = _BACKEND_DIR.parent.parent                     # repo root
+
+_CACHE_DIRS = [
+    _BACKEND_DIR / "unsloth_compiled_cache",
+    _PROJECT_ROOT / "unsloth_compiled_cache",
+    _PROJECT_ROOT / "studio" / "tmp" / "unsloth_compiled_cache",
+]
+
+
+def clear_unsloth_compiled_cache() -> None:
+    """Remove every known unsloth_compiled_cache directory (idempotent)."""
+    for cache_dir in _CACHE_DIRS:
+        if cache_dir.exists():
+            logger.info(f"Removing unsloth compiled cache: {cache_dir}")
+            shutil.rmtree(cache_dir, ignore_errors=True)