diff --git a/studio/backend/run.py b/studio/backend/run.py index 2f064ab92e..0ae1b4c697 100644 --- a/studio/backend/run.py +++ b/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 diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index 829f086b7d..11d8a1946b 100644 --- a/unsloth_cli/commands/studio.py +++ b/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 diff --git a/unsloth_cli/commands/ui.py b/unsloth_cli/commands/ui.py index eb6a8a69e8..a35c0d3a68 100644 --- a/unsloth_cli/commands/ui.py +++ b/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