2773f98357
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.
369 lines
12 KiB
Python
369 lines
12 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""
|
|
Run script for Unsloth UI Backend.
|
|
Works independently and can be moved to any directory.
|
|
"""
|
|
|
|
import os
|
|
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
|
|
backend_dir = Path(__file__).parent
|
|
if str(backend_dir) not in sys.path:
|
|
sys.path.insert(0, str(backend_dir))
|
|
|
|
from loggers import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
def _resolve_external_ip() -> str:
|
|
"""
|
|
Resolve the machine's external IP address.
|
|
|
|
Tries (in order):
|
|
1. GCE metadata server (instant, works on Google Cloud VMs)
|
|
2. ifconfig.me (works anywhere with internet)
|
|
3. LAN IP via UDP socket trick (fallback)
|
|
"""
|
|
import urllib.request
|
|
import socket
|
|
|
|
# 1. Try GCE metadata server (responds in <10ms on GCE, times out fast elsewhere)
|
|
try:
|
|
req = urllib.request.Request(
|
|
"http://metadata.google.internal/computeMetadata/v1/instance/network-interfaces/0/access-configs/0/external-ip",
|
|
headers = {"Metadata-Flavor": "Google"},
|
|
)
|
|
with urllib.request.urlopen(req, timeout = 1) as resp:
|
|
ip = resp.read().decode().strip()
|
|
if ip:
|
|
return ip
|
|
except Exception:
|
|
pass
|
|
|
|
# 2. Try public IP service
|
|
try:
|
|
with urllib.request.urlopen("https://ifconfig.me", timeout = 3) as resp:
|
|
ip = resp.read().decode().strip()
|
|
if ip:
|
|
return ip
|
|
except Exception:
|
|
pass
|
|
|
|
# 3. Fallback: LAN IP via UDP socket trick
|
|
try:
|
|
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
|
s.connect(("8.8.8.8", 80))
|
|
ip = s.getsockname()[0]
|
|
s.close()
|
|
return ip
|
|
except Exception:
|
|
return "0.0.0.0"
|
|
|
|
|
|
def _is_port_free(host: str, port: int) -> bool:
|
|
"""Check if a port is available for binding."""
|
|
import socket
|
|
|
|
try:
|
|
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
|
s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
|
s.bind((host, port))
|
|
return True
|
|
except OSError:
|
|
return False
|
|
|
|
|
|
def _find_free_port(host: str, start: int, max_attempts: int = 20) -> int:
|
|
"""Find a free port starting from `start`, trying up to max_attempts ports."""
|
|
for offset in range(max_attempts):
|
|
candidate = start + offset
|
|
if _is_port_free(host, candidate):
|
|
return candidate
|
|
raise RuntimeError(
|
|
f"Could not find a free port in range {start}-{start + max_attempts - 1}"
|
|
)
|
|
|
|
|
|
def _graceful_shutdown(server = None):
|
|
"""Explicitly shut down all subprocess backends and the uvicorn server.
|
|
|
|
Called from signal handlers to ensure child processes are cleaned up
|
|
before the parent exits. This is critical on Windows where atexit
|
|
handlers are unreliable after Ctrl+C.
|
|
"""
|
|
logger.info("Graceful shutdown initiated — cleaning up subprocesses...")
|
|
|
|
# 1. Shut down uvicorn server (releases the listening socket)
|
|
if server is not None:
|
|
server.should_exit = True
|
|
|
|
# 2. Clean up inference subprocess (if instantiated)
|
|
try:
|
|
from core.inference.orchestrator import _inference_backend
|
|
|
|
if _inference_backend is not None:
|
|
_inference_backend._shutdown_subprocess(timeout = 5.0)
|
|
except Exception as e:
|
|
logger.warning("Error shutting down inference subprocess: %s", e)
|
|
|
|
# 3. Clean up export subprocess (if instantiated)
|
|
try:
|
|
from core.export.orchestrator import _export_backend
|
|
|
|
if _export_backend is not None:
|
|
_export_backend._shutdown_subprocess(timeout = 5.0)
|
|
except Exception as e:
|
|
logger.warning("Error shutting down export subprocess: %s", e)
|
|
|
|
# 4. Clean up training subprocess (if active)
|
|
try:
|
|
from core.training.training import _training_backend
|
|
|
|
if _training_backend is not None:
|
|
_training_backend.force_terminate()
|
|
except Exception as e:
|
|
logger.warning("Error shutting down training subprocess: %s", e)
|
|
|
|
# 5. Kill llama-server subprocess (if loaded)
|
|
try:
|
|
from routes.inference import _llama_cpp_backend
|
|
|
|
if _llama_cpp_backend is not None:
|
|
_llama_cpp_backend._kill_process()
|
|
except Exception as e:
|
|
logger.warning("Error shutting down llama-server: %s", e)
|
|
|
|
logger.info("All subprocesses cleaned up")
|
|
|
|
|
|
# The uvicorn server instance — set by run_server(), used by callers
|
|
# that need to tell the server to exit (e.g. signal handlers).
|
|
_server = None
|
|
|
|
# Shutdown event — used to wake the main loop on signal
|
|
_shutdown_event = None
|
|
|
|
|
|
def run_server(
|
|
host: str = "0.0.0.0",
|
|
port: int = 8000,
|
|
frontend_path: Path = Path(__file__).resolve().parent.parent / "frontend" / "dist",
|
|
silent: bool = False,
|
|
allow_reexec: bool = False,
|
|
):
|
|
"""
|
|
Start the FastAPI server.
|
|
|
|
Args:
|
|
host: Host to bind to
|
|
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
|
|
(e.g. Colab notebooks) keep their own interrupt semantics.
|
|
Standalone callers should register handlers after calling this.
|
|
"""
|
|
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()
|
|
|
|
import asyncio
|
|
from threading import Thread, Event
|
|
import time
|
|
import uvicorn
|
|
|
|
from main import app, setup_frontend
|
|
from utils.paths import ensure_studio_directories
|
|
|
|
# Create all standard directories on startup
|
|
ensure_studio_directories()
|
|
|
|
# Auto-find free port if requested port is in use
|
|
if not _is_port_free(host, port):
|
|
original_port = port
|
|
port = _find_free_port(host, port)
|
|
if not silent:
|
|
print(f"Port {original_port} is in use, using port {port} instead")
|
|
|
|
# Setup frontend if path provided
|
|
if frontend_path:
|
|
if setup_frontend(app, frontend_path):
|
|
if not silent:
|
|
print(f"✅ Frontend loaded from {frontend_path}")
|
|
else:
|
|
if not silent:
|
|
print(f"⚠️ Frontend not found at {frontend_path}")
|
|
|
|
# Create the uvicorn server and expose it for signal handlers
|
|
config = uvicorn.Config(
|
|
app, host = host, port = port, log_level = "info", access_log = False
|
|
)
|
|
_server = uvicorn.Server(config)
|
|
_shutdown_event = Event()
|
|
|
|
# Run server in a daemon thread
|
|
def _run():
|
|
asyncio.run(_server.serve())
|
|
|
|
thread = Thread(target = _run, daemon = True)
|
|
thread.start()
|
|
time.sleep(3)
|
|
|
|
if not silent:
|
|
display_host = _resolve_external_ip() if host == "0.0.0.0" else host
|
|
|
|
print("")
|
|
print("=" * 50)
|
|
print(f"🦥 Unsloth Studio is running on port {port}")
|
|
print(f" Local: http://localhost:{port}")
|
|
print(f" External: http://{display_host}:{port}")
|
|
print(f" API: http://{display_host}:{port}/api")
|
|
print(f" Health: http://{display_host}:{port}/api/health")
|
|
print("=" * 50)
|
|
|
|
return app
|
|
|
|
|
|
# For direct execution (also invoked by CLI via os.execvp / subprocess)
|
|
if __name__ == "__main__":
|
|
_maybe_reexec_for_cuda_ld_path()
|
|
|
|
import argparse
|
|
import signal
|
|
|
|
parser = argparse.ArgumentParser(description = "Run Unsloth UI Backend server")
|
|
parser.add_argument("--host", default = "0.0.0.0", help = "Host to bind to")
|
|
parser.add_argument("--port", type = int, default = 8000, help = "Port to bind to")
|
|
parser.add_argument(
|
|
"--frontend",
|
|
type = str,
|
|
default = Path(__file__).resolve().parent.parent / "frontend" / "dist",
|
|
help = "Path to frontend build",
|
|
)
|
|
parser.add_argument("--silent", action = "store_true", help = "Suppress output")
|
|
|
|
args = parser.parse_args()
|
|
|
|
kwargs = dict(host = args.host, port = args.port, silent = args.silent)
|
|
if args.frontend is not None:
|
|
kwargs["frontend_path"] = Path(args.frontend)
|
|
run_server(**kwargs)
|
|
|
|
# ── Signal handler — ensures subprocess cleanup on Ctrl+C ────
|
|
def _signal_handler(signum, frame):
|
|
_graceful_shutdown(_server)
|
|
_shutdown_event.set()
|
|
|
|
signal.signal(signal.SIGINT, _signal_handler)
|
|
signal.signal(signal.SIGTERM, _signal_handler)
|
|
|
|
# On Windows, some terminals send SIGBREAK for Ctrl+C / Ctrl+Break
|
|
if hasattr(signal, "SIGBREAK"):
|
|
signal.signal(signal.SIGBREAK, _signal_handler)
|
|
|
|
# Keep running until shutdown signal.
|
|
# NOTE: Event.wait() without a timeout blocks at the C level on Linux,
|
|
# which prevents Python from delivering SIGINT (Ctrl+C). Using a
|
|
# short timeout in a loop lets the interpreter process pending signals.
|
|
while not _shutdown_event.is_set():
|
|
_shutdown_event.wait(timeout = 1)
|