Gate run_server re-exec behind allow_reexec so embedders are not restarted

run_server is a library/embed entry point (colab.start calls it directly), so the
unconditional os.execv for the torch CUDA LD_LIBRARY_PATH fix would replace the live
Colab/Jupyter kernel and drop in-memory state. Gate the re-exec behind a new
allow_reexec flag defaulting to False; the run.py __main__ path already re-execs
before calling run_server, and the unsloth_cli studio/ui entrypoints opt in with
allow_reexec=True so the CLI keeps the CUDA LD fix.

studio/backend/run.py

@@ -12,6 +12,80 @@ import sys
 # Suppress annoying C-level dependency warnings globally (e.g. SwigPyPacked)
 os.environ["PYTHONWARNINGS"] = "ignore"
 
+
+def _fix_torch_cuda_ld_path():
+    """Prepend torch's bundled CUDA libs to LD_LIBRARY_PATH.
+
+    PyTorch wheels ship their own CUDA runtime (libcudart, libcublas, etc.)
+    inside ``site-packages/nvidia/*/lib/``.  On Linux the dynamic linker
+    checks LD_LIBRARY_PATH **before** the RUNPATH baked into the .so files,
+    so a user's pre-existing LD_LIBRARY_PATH pointing at a different system
+    CUDA (e.g. /usr/local/cuda-13/lib64) will shadow torch's libs and cause
+    symbol-version errors at import time.
+
+    Fix: detect torch's lib dirs (without importing torch) and prepend them
+    so they take priority.  Returns True if LD_LIBRARY_PATH was changed.
+    """
+    if sys.platform != "linux":
+        return False
+    ld_path = os.environ.get("LD_LIBRARY_PATH", "")
+    if not ld_path:
+        return False
+    try:
+        import importlib.util
+
+        spec = importlib.util.find_spec("torch")
+        if not spec or not spec.origin:
+            return False
+        torch_dir = os.path.dirname(spec.origin)
+        site_pkgs = os.path.dirname(torch_dir)
+        nvidia_dir = os.path.join(site_pkgs, "nvidia")
+
+        lib_dirs = []
+        torch_lib = os.path.join(torch_dir, "lib")
+        if os.path.isdir(torch_lib):
+            lib_dirs.append(torch_lib)
+        if os.path.isdir(nvidia_dir):
+            for sub in sorted(os.listdir(nvidia_dir)):
+                lib = os.path.join(nvidia_dir, sub, "lib")
+                if os.path.isdir(lib):
+                    lib_dirs.append(lib)
+        if not lib_dirs:
+            return False
+
+        # Already at the front -- nothing to do
+        existing = ld_path.split(":")
+        if existing[: len(lib_dirs)] == lib_dirs:
+            return False
+
+        # Prepend torch dirs, deduplicate
+        torch_set = set(lib_dirs)
+        cleaned = [p for p in existing if p not in torch_set]
+        os.environ["LD_LIBRARY_PATH"] = ":".join(lib_dirs + cleaned)
+        return True
+    except Exception:
+        return False
+
+
+_LD_FIXED_SENTINEL = "_UNSLOTH_STUDIO_LD_FIXED"
+
+
+def _maybe_reexec_for_cuda_ld_path():
+    """Re-exec once so the dynamic linker sees corrected LD_LIBRARY_PATH.
+
+    Must only be called from a true entry point (``if __name__ == "__main__"``
+    or an explicit startup function), never at module import time, because
+    os.execv replaces the entire process.
+    """
+    if _LD_FIXED_SENTINEL in os.environ:
+        return
+    if not _fix_torch_cuda_ld_path():
+        return
+    os.environ[_LD_FIXED_SENTINEL] = "1"
+    argv = getattr(sys, "orig_argv", None) or [sys.executable] + sys.argv
+    os.execv(sys.executable, argv)
+
+
 from pathlib import Path
 
 # Add the backend directory to Python path
@@ -158,6 +232,7 @@ def run_server(
     port: int = 8000,
     frontend_path: Path = Path(__file__).resolve().parent.parent / "frontend" / "dist",
     silent: bool = False,
+    allow_reexec: bool = False,
 ):
     """
     Start the FastAPI server.
@@ -167,6 +242,11 @@ def run_server(
         port: Port to bind to (auto-increments if in use)
         frontend_path: Path to frontend build directory (optional)
         silent: Suppress startup messages
+        allow_reexec: Re-exec the process (os.execv) to apply the torch CUDA
+            LD_LIBRARY_PATH fix. Enable ONLY from true CLI/process entrypoints.
+            Embedders (e.g. Colab via colab.start) must leave this False,
+            otherwise the re-exec replaces the notebook kernel and drops
+            in-memory state.
 
     Note:
         Signal handlers are NOT registered here so that embedders
@@ -175,6 +255,14 @@ def run_server(
     """
     global _server, _shutdown_event
 
+    # Only re-exec when invoked as a real CLI/process entrypoint. Embedders
+    # (e.g. Colab via colab.start -> run_server) keep allow_reexec=False,
+    # because os.execv would replace the live notebook kernel and drop
+    # in-memory state. CLI entrypoints (run.py __main__, unsloth_cli in-venv
+    # fallback) opt in explicitly.
+    if allow_reexec:
+        _maybe_reexec_for_cuda_ld_path()
+
     import nest_asyncio
 
     nest_asyncio.apply()
@@ -238,6 +326,8 @@ def run_server(
 
 # For direct execution (also invoked by CLI via os.execvp / subprocess)
 if __name__ == "__main__":
+    _maybe_reexec_for_cuda_ld_path()
+
     import argparse
     import signal
 

unsloth_cli/commands/studio.py

@@ -138,6 +138,8 @@ def studio_default(
         port = port,
         frontend_path = frontend,
         silent = silent,
+        # CLI entrypoint: apply the torch CUDA LD_LIBRARY_PATH fix via re-exec.
+        allow_reexec = True,
     )
 
     from studio.backend.run import _shutdown_event

unsloth_cli/commands/ui.py

@@ -85,6 +85,8 @@ def ui(
         port = port,
         frontend_path = frontend,
         silent = silent,
+        # CLI entrypoint: apply the torch CUDA LD_LIBRARY_PATH fix via re-exec.
+        allow_reexec = True,
     )
 
     from studio.backend.run import _shutdown_event