* fix(studio/colab): merge iframe+keepalive into start(), add proxy_headers to uvicorn - Move serve_kernel_port_as_iframe and keepalive loop into colab.start() so both run in the same cell execution context, eliminating the race where the proxy URL was shown before the iframe cell had a chance to run - Add a 2s sleep after run_server() before show_link() to give Colab's proxy infrastructure time to register the bound port - Add proxy_headers=True and forwarded_allow_ips="*" to uvicorn Config so X-Forwarded-Proto/Host from Colab's reverse proxy are trusted - Simplify notebook start cell (no more separate iframe cell needed) * fix(studio/colab): fix iframe blocking and server thread crash in Colab Two root causes for the long-standing proxy/iframe breakage: 1. SecurityHeadersMiddleware set X-Frame-Options: DENY and frame-ancestors 'none' unconditionally, blocking serve_kernel_port_as_iframe regardless of server health. Fix: detect Colab via COLAB_BACKEND_URL/COLAB_GPU env vars, relax frame-ancestors to *.prod.colab.dev and omit X-Frame-Options. 2. asyncio.run() in the daemon thread conflicted with nest_asyncio's global patches applied on the main thread, causing the server to crash silently after ready_event fired. Fix: use explicit new_event_loop() + run_until_complete() in the daemon thread to bypass nest_asyncio's asyncio.run patch. Also replace blind time.sleep(2) with a health endpoint poll so the link and iframe are only shown once the server is truly reachable. * fix(studio/colab): use reliable /content + google.colab path for Colab detection COLAB_BACKEND_URL and COLAB_GPU env vars aren't consistently set across all Colab runtime versions. Use /content dir + google.colab package path as a more reliable signal, computed once at module load. * fix(studio/colab): fix port mismatch, health-check silence, and CSP framing Four bugs causing the iframe and URL button to always fail: 1. Port not propagated back: run_server auto-increments when 8888 is taken, but start() kept using the original port for show_link() and serve_kernel_port_as_iframe() — now reads app.state.server_port. 2. Silent health-check failure: the poll loop never checked whether any attempt succeeded; on all-fail it continued and showed a dead link — now exits early with a clear error message. 3. CSP frame-ancestors too narrow: '*.prod.colab.dev' only matches one subdomain level; actual Colab proxy URLs are two levels deep (e.g. foo.region.prod.colab.dev), and the parent frame may also be colab.research.google.com or a sandboxed null-origin output iframe — changed to '*' in Colab mode (single-user sandbox, no security loss). 4. _IS_COLAB detection hardcoded python3.10/3.11 paths: Python 3.12+ Colab runtimes wouldn't match when env vars aren't set — replaced with a glob over python3.*/dist-packages/google/colab. * fix(studio/colab): harden Colab startup against every known failure mode colab.py: - get_colab_url: retry eval_js up to 3x (10s timeout each), validate that result is a real https:// URL containing the port before accepting it; log a clear warning when falling back to localhost - show_link: safe short_url truncation (try/except around str.index so an unexpected URL shape never blocks the link card from rendering); also emit the URL via logger so it's visible in cell text output even if HTML display is suppressed - start: detect "already running" at entry — on cell re-run Studio is still healthy on port 8888; skip re-launch and go straight to show+iframe so the user never ends up with mismatched port state - start: wrap run_server in try/except (SystemExit + Exception) so startup errors surface as readable messages rather than cell crashes - start: check frontend_path/index.html exists, not just the directory - start: remove unused `import sys` - start / keepalive: catch KeyboardInterrupt so interrupting the cell prints a clean "stopped" message instead of a raw traceback - extract _is_studio_healthy() and _show_and_embed() helpers to deduplicate the fast-path and normal-path logic main.py: - _build_csp: in Colab mode, extend script-src to include *.prod.colab.dev and *.googleusercontent.com (Colab injects scripts from these origins into the output iframe scaffolding) - _build_csp: in Colab mode, extend connect-src with blob:, data:, wss://*.prod.colab.dev, and wss://*.googleusercontent.com so WebSocket streams and Colab kernel traffic are not blocked by CSP * fix(studio/colab): fix iframe width responsiveness and height sizing Replace serve_kernel_port_as_iframe with a raw CSS iframe for two reasons: 1. Width responsiveness: serve_kernel_port_as_iframe sets the width as an HTML attribute (width="100%") which Colab's output machinery can bake into a fixed pixel value on first render, causing the Studio to stop following the notebook panel width when it opens/closes or the window resizes. A CSS style property (style="width:100%") participates in normal reflow and always tracks the parent container width. 2. Height sizing: the hardcoded height=1200 was too tall on short monitors (forced outer-page scroll) and wasted space on tall ones. A small JS snippet reads screen.availHeight and sets height to ~82% of the screen, clamped to [600, 1100]px, with a resize listener that re-fits on zoom changes and panel open/close events. Also eliminate the double eval_js call: _show_and_embed now fetches the Colab proxy URL once and passes it to show_link via the new _url kwarg, so google.colab.kernel.proxyPort is only called once per invocation. Falls back to serve_kernel_port_as_iframe if IPython.display.HTML is unavailable for any reason. * fix(studio/colab): fix link button + add fullscreen hover button to iframe Link button: target="_blank" is blocked by Colab's output sandbox. Switch to onclick="window.open(url,'_blank')" which the sandbox allows. Fullscreen: add a small button that appears on hover in the top-right corner of the iframe. Clicking it calls requestFullscreen() on the wrapper div and stretches the iframe to 100vh/100vw. Exits back to normal on fullscreen change. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * revert(studio/colab): remove fullscreen button * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio/colab): address review feedback - Wrap both urlopen calls in with statements to prevent socket/fd leaks - Replace JS resize listener with CSS height:82vh — simpler, responsive, and no risk of leaked window listeners on cell re-runs - Use importlib.util.find_spec("google.colab") instead of a glob path to detect Colab; more robust across Python versions and venv layouts * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio/colab): fall back to href navigation when window.open is blocked window.open from a cross-origin sandboxed Colab output iframe can be silently blocked by the browser (returns null, no exception). The old code returned false unconditionally, so a blocked popup left the button doing nothing. Now: if window.open succeeds the new tab opens and the href is suppressed; if it returns null the browser follows the href, navigating the output cell to Studio — always does something useful. * fix(studio/colab): remove button, give iframe a branded header bar The "Open Unsloth Studio" button was unreliable in Colab's sandboxed output context regardless of how window.open was called. Since the iframe already loads Studio inline, the button added no value and confused users with a URL that 404s outside the output cell. Replace the separate link card + bare iframe with a single block: a slim black header bar (Unsloth logo + truncated URL) flush on top of the full-height responsive iframe. Cleaner and removes the broken button entirely. * studio: gate uvicorn proxy_headers/forwarded_allow_ips behind _IS_COLAB forwarded_allow_ips="*" was applied unconditionally, so every Studio deployment trusted X-Forwarded-* headers from any client. Only Colab needs that, because its reverse proxy fronts the kernel. For a normal local/standalone Studio this is an unwanted relaxation, especially when bound to 0.0.0.0. Now proxy_headers/forwarded_allow_ips are only set when _IS_COLAB. Standalone runs fall back to uvicorn's defaults (proxy_headers honored from loopback only), restoring the prior security posture, while Colab keeps the wide trust its proxy requires. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han <danielhanchen@gmail.com>
937 lines
34 KiB
Python
937 lines
34 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
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
# Suppress annoying C-level dependency warnings globally (e.g. SwigPyPacked)
|
|
os.environ["PYTHONWARNINGS"] = "ignore"
|
|
|
|
# Add the backend directory to Python path early so local modules are importable
|
|
backend_dir = Path(__file__).parent
|
|
if str(backend_dir) not in sys.path:
|
|
sys.path.insert(0, str(backend_dir))
|
|
|
|
from utils.cpu_threads import configure_cpu_threads
|
|
|
|
try:
|
|
configure_cpu_threads()
|
|
except ValueError as exc:
|
|
configured = os.environ.get("UNSLOTH_CPU_THREADS")
|
|
raise SystemExit(
|
|
f"Error: Invalid UNSLOTH_CPU_THREADS value {configured!r}: {exc}"
|
|
) from None
|
|
|
|
# Fix for Anaconda/conda-forge Python: seed platform._sys_version_cache before
|
|
# any library imports that trigger attrs -> rich -> structlog -> platform crash.
|
|
# See: https://github.com/python/cpython/issues/102396
|
|
import _platform_compat # noqa: F401
|
|
|
|
from loggers import get_logger
|
|
from startup_banner import print_studio_access_banner, print_studio_stop_hint
|
|
|
|
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 _install_uvicorn_startup_log_rewrite(bind_host: str, display_host: str) -> None:
|
|
"""Rewrite Uvicorn's startup log line: swap wildcard bind for the
|
|
externally-reachable address, replace the CTRL+C suffix with our Mac-aware
|
|
stop hint, and rename the prefix to "Unsloth Studio running on"."""
|
|
import logging
|
|
import re
|
|
|
|
rewrite_host = (
|
|
bind_host in ("0.0.0.0", "::")
|
|
and bool(display_host)
|
|
and display_host != bind_host
|
|
)
|
|
new_suffix = "(To stop: press Ctrl+C -- on macOS, Control+C not Command+C)"
|
|
old_suffix_re = re.compile(r"\(Press CTRL\+C to quit\)")
|
|
old_prefix = "Uvicorn running on "
|
|
new_prefix = "Unsloth Studio running on "
|
|
|
|
def _rewrite(text: str) -> str:
|
|
if text.startswith(old_prefix):
|
|
text = new_prefix + text[len(old_prefix) :]
|
|
return old_suffix_re.sub(new_suffix, text)
|
|
|
|
class _UvicornStartupRewrite(logging.Filter):
|
|
def filter(self, record: logging.LogRecord) -> bool:
|
|
try:
|
|
msg = record.msg if isinstance(record.msg, str) else ""
|
|
if (
|
|
msg.startswith(old_prefix)
|
|
and isinstance(record.args, tuple)
|
|
and len(record.args) >= 3
|
|
):
|
|
if rewrite_host and record.args[1] == bind_host:
|
|
record.args = (
|
|
record.args[0],
|
|
display_host,
|
|
record.args[2],
|
|
*record.args[3:],
|
|
)
|
|
record.msg = _rewrite(msg)
|
|
cmsg = getattr(record, "color_message", None)
|
|
if isinstance(cmsg, str):
|
|
record.color_message = _rewrite(cmsg)
|
|
except Exception:
|
|
pass
|
|
return True
|
|
|
|
f = _UvicornStartupRewrite()
|
|
for name in ("uvicorn", "uvicorn.error"):
|
|
logging.getLogger(name).addFilter(f)
|
|
|
|
|
|
def _local_port_open(host: str, port: int, timeout: float = 1.0) -> bool:
|
|
"""Return True iff a TCP connection to (host, port) succeeds within timeout."""
|
|
import socket
|
|
|
|
try:
|
|
with socket.create_connection((host, port), timeout = timeout):
|
|
return True
|
|
except OSError:
|
|
return False
|
|
|
|
|
|
def _working_local_url(port: int) -> "str | None":
|
|
"""Return a working loopback URL on this machine, or None if neither
|
|
127.0.0.1 nor ::1 responds. Used as a fallback when external reachability fails."""
|
|
if _local_port_open("127.0.0.1", port):
|
|
return f"http://127.0.0.1:{port}"
|
|
if _local_port_open("::1", port):
|
|
return f"http://[::1]:{port}"
|
|
return None
|
|
|
|
|
|
def _stdout_color_ok() -> bool:
|
|
"""Whether to emit ANSI color codes on stdout. Mirrors startup_banner."""
|
|
if os.environ.get("NO_COLOR", "").strip():
|
|
return False
|
|
if os.environ.get("FORCE_COLOR", "").strip():
|
|
return True
|
|
try:
|
|
return sys.stdout.isatty()
|
|
except (AttributeError, OSError, ValueError):
|
|
return False
|
|
|
|
|
|
def _verify_global_reachability(display_host: str, port: int) -> None:
|
|
"""Probe check-host.net to confirm display_host:port is reachable from the
|
|
public internet. Synchronous so the caller can render output between the
|
|
banner URL section and the trailing stop hint. Bounded at ~15s; failures
|
|
are swallowed (the verifier failing is not Studio failing). Only meaningful
|
|
when bound to a wildcard host."""
|
|
import ipaddress
|
|
import json
|
|
import time
|
|
import urllib.error
|
|
import urllib.parse
|
|
import urllib.request
|
|
|
|
if not display_host or display_host in ("0.0.0.0", "::"):
|
|
return
|
|
|
|
use_color = _stdout_color_ok()
|
|
dim = "\033[38;5;245m" if use_color else ""
|
|
ok_c = "\033[38;5;120;1m" if use_color else ""
|
|
err_c = "\033[38;5;203;1m" if use_color else ""
|
|
warn_c = "\033[38;5;215;1m" if use_color else ""
|
|
local_url_c = "\033[38;5;108;1m" if use_color else "" # matches banner's URL color
|
|
reset = "\033[0m" if use_color else ""
|
|
|
|
url = f"http://{display_host}:{port}"
|
|
|
|
# Private / loopback / link-local addresses are not globally routable.
|
|
try:
|
|
addr = ipaddress.ip_address(display_host)
|
|
if addr.is_loopback or addr.is_private or addr.is_link_local:
|
|
print(
|
|
f"{dim} Note: {display_host} is a private/LAN address -- "
|
|
f"reachable on this network only, not from the public internet."
|
|
f"{reset}",
|
|
flush = True,
|
|
)
|
|
return
|
|
except ValueError:
|
|
# Not an IP literal; probe by hostname.
|
|
pass
|
|
|
|
try:
|
|
qs = urllib.parse.urlencode({"host": f"{display_host}:{port}", "max_nodes": 3})
|
|
req = urllib.request.Request(
|
|
f"https://check-host.net/check-tcp?{qs}",
|
|
headers = {
|
|
"Accept": "application/json",
|
|
"User-Agent": "unsloth-studio-reachability/1",
|
|
},
|
|
)
|
|
with urllib.request.urlopen(req, timeout = 5) as resp:
|
|
init = json.loads(resp.read().decode("utf-8", errors = "replace"))
|
|
req_id = init.get("request_id")
|
|
if not req_id:
|
|
return
|
|
|
|
results = {}
|
|
deadline = time.monotonic() + 15.0
|
|
poll_req = urllib.request.Request(
|
|
f"https://check-host.net/check-result/{req_id}",
|
|
headers = {
|
|
"Accept": "application/json",
|
|
"User-Agent": "unsloth-studio-reachability/1",
|
|
},
|
|
)
|
|
while time.monotonic() < deadline:
|
|
time.sleep(1.5)
|
|
try:
|
|
with urllib.request.urlopen(poll_req, timeout = 5) as resp:
|
|
results = json.loads(resp.read().decode("utf-8", errors = "replace"))
|
|
except Exception:
|
|
continue
|
|
if results and all(v is not None for v in results.values()):
|
|
break
|
|
# Two decisive nodes is enough; stop polling early.
|
|
decisive = [
|
|
v
|
|
for v in results.values()
|
|
if isinstance(v, list)
|
|
and v
|
|
and isinstance(v[0], dict)
|
|
and ("time" in v[0] or "error" in v[0])
|
|
]
|
|
if len(decisive) >= 2:
|
|
break
|
|
|
|
ok_nodes = err_nodes = 0
|
|
for v in results.values():
|
|
if not isinstance(v, list) or not v or not isinstance(v[0], dict):
|
|
continue
|
|
if "time" in v[0]:
|
|
ok_nodes += 1
|
|
elif "error" in v[0]:
|
|
err_nodes += 1
|
|
total = ok_nodes + err_nodes
|
|
|
|
print("", flush = True)
|
|
if ok_nodes:
|
|
print(
|
|
f"{ok_c} Reachability check: {url}/ is reachable from the "
|
|
f"public internet ({ok_nodes}/{total} probe nodes connected).{reset}",
|
|
flush = True,
|
|
)
|
|
elif err_nodes:
|
|
print(
|
|
f"{err_c} Reachability check: {url}/ is NOT reachable from "
|
|
f"the public internet ({err_nodes}/{total} probe nodes failed).{reset}",
|
|
flush = True,
|
|
)
|
|
print(f"{dim} Common causes:{reset}", flush = True)
|
|
print(
|
|
f"{dim} * AWS -- the instance's Security Group doesn't "
|
|
f"allow inbound TCP {port}.{reset}",
|
|
flush = True,
|
|
)
|
|
print(
|
|
f"{dim} * GCP -- no firewall rule allowing TCP {port} "
|
|
f"for the instance's network tag.{reset}",
|
|
flush = True,
|
|
)
|
|
print(
|
|
f"{dim} * Azure / other clouds -- equivalent NSG / "
|
|
f"firewall rule missing.{reset}",
|
|
flush = True,
|
|
)
|
|
print(
|
|
f"{dim} * Home -- your router isn't port-forwarding "
|
|
f"{port} to this machine.{reset}",
|
|
flush = True,
|
|
)
|
|
print(
|
|
f"{dim} Workaround that needs no firewall changes -- "
|
|
f"SSH local-forward from your laptop:{reset}",
|
|
flush = True,
|
|
)
|
|
print(
|
|
f"{dim} ssh -L {port}:localhost:{port} "
|
|
f"<user>@{display_host}{reset}",
|
|
flush = True,
|
|
)
|
|
print(
|
|
f"{dim} then open http://localhost:{port}/ in your browser.{reset}",
|
|
flush = True,
|
|
)
|
|
# Only offer the local URL if loopback actually answers.
|
|
local_url = _working_local_url(port)
|
|
if local_url:
|
|
print(
|
|
f"{local_url_c} You can access Unsloth Studio locally "
|
|
f"in the meantime: {local_url}{reset}",
|
|
flush = True,
|
|
)
|
|
else:
|
|
print(
|
|
f"{warn_c} Reachability check: probe nodes did not respond "
|
|
f"in time -- could not verify {url}/.{reset}",
|
|
flush = True,
|
|
)
|
|
except urllib.error.URLError:
|
|
# Outbound HTTPS blocked; skip silently.
|
|
pass
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def _get_pid_on_port(port: int) -> "tuple[int, str] | None":
|
|
"""Return (pid, process_name) of the process listening on *port*, or None.
|
|
|
|
Uses psutil when available. Falls back gracefully to None so callers
|
|
can still report the port conflict without process details.
|
|
|
|
Works on Windows, macOS, and Linux wherever psutil is installed.
|
|
"""
|
|
try:
|
|
import psutil
|
|
except ImportError:
|
|
return None
|
|
try:
|
|
for conn in psutil.net_connections(kind = "tcp"):
|
|
if conn.status == "LISTEN" and conn.laddr.port == port:
|
|
if conn.pid is None:
|
|
return None
|
|
try:
|
|
proc = psutil.Process(conn.pid)
|
|
return (conn.pid, proc.name())
|
|
except (psutil.NoSuchProcess, psutil.AccessDenied):
|
|
return (conn.pid, "<unknown>")
|
|
except (psutil.AccessDenied, OSError) as e:
|
|
# psutil.net_connections() needs elevated privileges on some platforms
|
|
logger.debug("Failed to scan network connections for port %s: %s", port, e)
|
|
return None
|
|
|
|
|
|
def _is_port_free(host: str, port: int) -> bool:
|
|
"""Check if a port is available for binding.
|
|
|
|
When *host* is ``0.0.0.0`` (wildcard), we also check whether anything
|
|
is already listening on ``127.0.0.1`` (and ``::1`` when IPv6 is
|
|
available). An SSH tunnel or similar process may hold the loopback
|
|
address while our wildcard bind still succeeds, making Unsloth Studio
|
|
unreachable via ``localhost``.
|
|
|
|
Works on Windows, macOS, and Linux.
|
|
"""
|
|
import socket
|
|
|
|
# 1. Can we bind to the requested address?
|
|
# Use getaddrinfo so both IPv4 ("0.0.0.0") and IPv6 ("::") hosts
|
|
# resolve to the correct address family automatically.
|
|
try:
|
|
addr_info = socket.getaddrinfo(host, port, socket.AF_UNSPEC, socket.SOCK_STREAM)
|
|
family, socktype, proto, _, sockaddr = addr_info[0]
|
|
with socket.socket(family, socktype, proto) as s:
|
|
s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
|
s.bind(sockaddr)
|
|
except OSError:
|
|
return False
|
|
|
|
# 2. When binding to all interfaces, verify that localhost is not
|
|
# already claimed by another process (e.g. an SSH -L tunnel).
|
|
# We attempt a TCP connect -- if it succeeds something is listening.
|
|
if host in ("0.0.0.0", "::"):
|
|
for loopback, family in [
|
|
("127.0.0.1", socket.AF_INET),
|
|
("::1", socket.AF_INET6),
|
|
]:
|
|
try:
|
|
with socket.socket(family, socket.SOCK_STREAM) as s:
|
|
s.settimeout(1)
|
|
if s.connect_ex((loopback, port)) == 0:
|
|
# Connection succeeded -- port is taken on loopback
|
|
return False
|
|
except OSError:
|
|
# IPv6 disabled or other OS-level restriction -- skip
|
|
continue
|
|
|
|
return True
|
|
|
|
|
|
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}"
|
|
)
|
|
|
|
|
|
from utils.paths.storage_roots import studio_root as _studio_root
|
|
|
|
_PID_FILE = _studio_root() / "studio.pid"
|
|
|
|
# Direct backend launches bypass the CLI's env re-export; do it here for
|
|
# real custom roots so unsloth-zoo's import-time LLAMA_CPP_DEFAULT_DIR
|
|
# picks up the custom build. Skip for legacy-default to avoid flipping
|
|
# default-mode installs into env-override.
|
|
try:
|
|
_LEGACY_STUDIO_ROOT = (Path.home() / ".unsloth" / "studio").resolve()
|
|
except (OSError, ValueError):
|
|
_LEGACY_STUDIO_ROOT = Path.home() / ".unsloth" / "studio"
|
|
try:
|
|
_STUDIO_ROOT_RESOLVED = _studio_root().resolve()
|
|
except (OSError, ValueError):
|
|
_STUDIO_ROOT_RESOLVED = _studio_root()
|
|
if _STUDIO_ROOT_RESOLVED != _LEGACY_STUDIO_ROOT:
|
|
if not os.environ.get("UNSLOTH_STUDIO_HOME"):
|
|
os.environ["UNSLOTH_STUDIO_HOME"] = str(_STUDIO_ROOT_RESOLVED)
|
|
if not os.environ.get("UNSLOTH_LLAMA_CPP_PATH"):
|
|
os.environ["UNSLOTH_LLAMA_CPP_PATH"] = str(_STUDIO_ROOT_RESOLVED / "llama.cpp")
|
|
|
|
|
|
def _write_pid_file():
|
|
"""Write the current process PID to the studio PID file."""
|
|
try:
|
|
_PID_FILE.parent.mkdir(parents = True, exist_ok = True)
|
|
_PID_FILE.write_text(str(os.getpid()))
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
def _remove_pid_file():
|
|
"""Remove the PID file if it belongs to this process."""
|
|
try:
|
|
if _PID_FILE.is_file():
|
|
stored = _PID_FILE.read_text().strip()
|
|
if stored == str(os.getpid()):
|
|
_PID_FILE.unlink(missing_ok = True)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
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.
|
|
"""
|
|
_remove_pid_file()
|
|
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
|
|
|
|
|
|
_DEFAULT_FRONTEND_PATH = Path(__file__).resolve().parent.parent / "frontend" / "dist"
|
|
|
|
|
|
def _iter_frontend_fallback_candidates() -> "list[Path]":
|
|
"""Yield `studio/frontend/dist` paths to try when the default is missing.
|
|
|
|
Covers PATH-shadowed binaries whose __file__ resolves into a
|
|
site-packages tree that never received a vite build (e.g. plain
|
|
`pip install unsloth` from PyPI).
|
|
"""
|
|
import ast
|
|
import re
|
|
|
|
out: list[Path] = []
|
|
home_str = (
|
|
os.environ.get("UNSLOTH_STUDIO_HOME")
|
|
or os.environ.get("STUDIO_HOME")
|
|
or str(Path.home() / ".unsloth" / "studio")
|
|
)
|
|
venv_dir = Path(home_str).expanduser() / "unsloth_studio"
|
|
# Installer venv site-packages.
|
|
for pattern in (
|
|
"lib/python*/site-packages/studio/frontend/dist",
|
|
"Lib/site-packages/studio/frontend/dist",
|
|
):
|
|
out.extend(venv_dir.glob(pattern))
|
|
# Editable source roots referenced from the installer venv.
|
|
for sp_pattern in ("lib/python*/site-packages", "Lib/site-packages"):
|
|
for sp in venv_dir.glob(sp_pattern):
|
|
for finder in sp.glob("__editable___*_finder.py"):
|
|
try:
|
|
src = finder.read_text(encoding = "utf-8")
|
|
except OSError:
|
|
continue
|
|
# Tolerate single- or multi-line dict literals; [^}]* still
|
|
# rejects nested dicts, which the setuptools template never
|
|
# emits for editable installs.
|
|
m = re.search(
|
|
r"^MAPPING\s*(?::[^=]*)?=\s*(\{[^}]*\})", src, re.M | re.S
|
|
)
|
|
if not m:
|
|
continue
|
|
try:
|
|
mapping = ast.literal_eval(m.group(1))
|
|
except (SyntaxError, ValueError):
|
|
continue
|
|
# Defensive: literal_eval can return a set / list / None if the
|
|
# matched literal is not a dict (regex captures `{...}`).
|
|
if not isinstance(mapping, dict):
|
|
continue
|
|
studio_pkg = mapping.get("studio")
|
|
if studio_pkg:
|
|
out.append(Path(studio_pkg) / "frontend" / "dist")
|
|
return out
|
|
|
|
|
|
def _resolve_frontend_path(frontend_path: Path) -> tuple[Optional[Path], list[Path]]:
|
|
"""Pick a frontend dir that actually contains `index.html`.
|
|
|
|
Returns (chosen, attempted). `chosen` is None if nothing servable was
|
|
found; `attempted` is the full ordered list for diagnostics.
|
|
"""
|
|
attempted: list[Path] = []
|
|
seen: set[Path] = set()
|
|
|
|
def _try(p: Path) -> bool:
|
|
try:
|
|
key = p.resolve()
|
|
except OSError:
|
|
key = p
|
|
if key in seen:
|
|
return False
|
|
seen.add(key)
|
|
attempted.append(p)
|
|
return (p / "index.html").is_file()
|
|
|
|
if _try(Path(frontend_path)):
|
|
return attempted[-1], attempted
|
|
for alt in _iter_frontend_fallback_candidates():
|
|
if _try(alt):
|
|
return attempted[-1], attempted
|
|
return None, attempted
|
|
|
|
|
|
def run_server(
|
|
host: str = "127.0.0.1",
|
|
port: int = 8888,
|
|
frontend_path: Path = _DEFAULT_FRONTEND_PATH,
|
|
silent: bool = False,
|
|
api_only: bool = False,
|
|
llama_parallel_slots: int = 1,
|
|
):
|
|
"""
|
|
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
|
|
api_only: Run API server only, no frontend serving (for Tauri desktop app)
|
|
llama_parallel_slots: Number of parallel slots for llama-server
|
|
|
|
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
|
|
|
|
# On Windows the default console encoding (cp1252) cannot encode emoji.
|
|
# Reconfigure stdout to UTF-8 so startup messages do not crash the server.
|
|
if sys.platform == "win32" and hasattr(sys.stdout, "reconfigure"):
|
|
try:
|
|
sys.stdout.reconfigure(encoding = "utf-8", errors = "replace")
|
|
except Exception:
|
|
pass
|
|
|
|
# Set env var BEFORE importing main so CORS middleware picks it up
|
|
if api_only:
|
|
os.environ["UNSLOTH_API_ONLY"] = "1"
|
|
|
|
import nest_asyncio
|
|
|
|
nest_asyncio.apply()
|
|
|
|
import asyncio
|
|
from threading import Thread, Event
|
|
import uvicorn
|
|
|
|
from main import app, setup_frontend, _IS_COLAB
|
|
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
|
|
blocker = _get_pid_on_port(port)
|
|
port = _find_free_port(host, port + 1)
|
|
if not silent:
|
|
print("")
|
|
print("=" * 50)
|
|
if blocker:
|
|
pid, name = blocker
|
|
print(
|
|
f"Port {original_port} is already in use by " f"{name} (PID {pid})."
|
|
)
|
|
else:
|
|
print(f"Port {original_port} is already in use.")
|
|
print(f"Unsloth Studio will use port {port} instead.")
|
|
print(f"Open http://localhost:{port} in your browser.")
|
|
print("=" * 50)
|
|
print("")
|
|
|
|
# Setup frontend if path provided (skip in api-only mode).
|
|
# Falls back through alternate locations if the default lacks a built
|
|
# dist; errors out loudly rather than silently serving 404 on `/`.
|
|
if frontend_path and not api_only:
|
|
chosen, attempted = _resolve_frontend_path(Path(frontend_path))
|
|
if chosen is not None and setup_frontend(app, chosen):
|
|
if not silent:
|
|
# Resolve so logs always show an absolute path for support.
|
|
try:
|
|
display = chosen.resolve()
|
|
except OSError:
|
|
display = chosen
|
|
print(f"[OK] Frontend loaded from {display}")
|
|
else:
|
|
home_str = (
|
|
os.environ.get("UNSLOTH_STUDIO_HOME")
|
|
or os.environ.get("STUDIO_HOME")
|
|
or str(Path.home() / ".unsloth" / "studio")
|
|
)
|
|
# Windows ships the user-facing shim at $STUDIO_HOME/bin/unsloth.exe
|
|
# (a hardlink to the venv exe); Linux/macOS use the venv binary
|
|
# at $STUDIO_HOME/unsloth_studio/bin/unsloth.
|
|
home = Path(home_str).expanduser()
|
|
if sys.platform == "win32":
|
|
installer_bin = home / "bin" / "unsloth.exe"
|
|
else:
|
|
installer_bin = home / "unsloth_studio" / "bin" / "unsloth"
|
|
tried_lines = "\n".join(f" - {p}" for p in attempted) or " (none)"
|
|
raise SystemExit(
|
|
"[ERROR] Studio frontend build not found.\n"
|
|
f"Tried:\n{tried_lines}\n"
|
|
"\n"
|
|
"Likely cause: another 'unsloth' on PATH is shadowing the "
|
|
"installer's binary and points at a site-packages tree with "
|
|
"no built dist.\n"
|
|
"\n"
|
|
"Fix one of:\n"
|
|
f" - run the installer's binary directly: {installer_bin} studio\n"
|
|
" - pass --frontend <path/to/studio/frontend/dist>\n"
|
|
" - pass --api-only to skip serving the web UI\n"
|
|
" - reinstall: curl -fsSL https://unsloth.ai/install.sh | sh"
|
|
)
|
|
|
|
# Resolve once; shared by the log rewrite and the banner.
|
|
display_host = _resolve_external_ip() if host == "0.0.0.0" else host
|
|
_install_uvicorn_startup_log_rewrite(host, display_host)
|
|
|
|
ready_event = Event()
|
|
startup_failed = Event()
|
|
startup_errors = []
|
|
|
|
class _ReadyServer(uvicorn.Server):
|
|
async def startup(self, *args, **kwargs):
|
|
await super().startup(*args, **kwargs)
|
|
if getattr(self, "started", False) and not self.should_exit:
|
|
ready_event.set()
|
|
|
|
# server_header=False suppresses uvicorn's "Server: uvicorn"; SecurityHeadersMiddleware sets its own.
|
|
config_kwargs = dict(
|
|
host = host,
|
|
port = port,
|
|
log_level = "info",
|
|
access_log = False,
|
|
server_header = False,
|
|
)
|
|
# Only in Colab: trust X-Forwarded-* from Colab's reverse proxy so the app
|
|
# sees the real https origin. forwarded_allow_ips="*" is fine inside Colab's
|
|
# single-user sandbox, but would be an unwanted security relaxation for a
|
|
# normal local/standalone Studio, so leave uvicorn's safe defaults
|
|
# (forwarded headers trusted from loopback only) in place there.
|
|
if _IS_COLAB:
|
|
config_kwargs["proxy_headers"] = True
|
|
config_kwargs["forwarded_allow_ips"] = "*"
|
|
config = uvicorn.Config(app, **config_kwargs)
|
|
_server = _ReadyServer(config)
|
|
_shutdown_event = Event()
|
|
|
|
# Expose the actual bound port so request-handling code can build
|
|
# loopback URLs that point at the real backend, not whatever port a
|
|
# reverse proxy or tunnel exposed in the request URL. Only publish
|
|
# an explicit value when we know the concrete port; for ephemeral
|
|
# binds (port==0) leave it unset and let request handlers fall back
|
|
# to the ASGI request scope or request.base_url.
|
|
app.state.server_port = port if port and port > 0 else None
|
|
app.state.llama_parallel_slots = llama_parallel_slots
|
|
|
|
# Expose a shutdown callable via app.state before the server can accept
|
|
# requests so /api/shutdown is available as soon as readiness is published.
|
|
def _trigger_shutdown():
|
|
_graceful_shutdown(_server)
|
|
if _shutdown_event is not None:
|
|
_shutdown_event.set()
|
|
|
|
app.state.trigger_shutdown = _trigger_shutdown
|
|
|
|
# Run server in a daemon thread.
|
|
# Use an explicit new_event_loop() + run_until_complete() instead of
|
|
# asyncio.run() to avoid nest_asyncio's global patches to asyncio.run
|
|
# interfering when called from a thread while Colab/IPython already has
|
|
# a running loop on the main thread.
|
|
def _run():
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
try:
|
|
loop.run_until_complete(_server.serve())
|
|
except BaseException as exc:
|
|
startup_errors.append(exc)
|
|
startup_failed.set()
|
|
finally:
|
|
loop.close()
|
|
if not ready_event.is_set():
|
|
startup_failed.set()
|
|
|
|
thread = Thread(target = _run, daemon = True)
|
|
thread.start()
|
|
|
|
# Wait until uvicorn has completed lifespan startup and bound sockets, or
|
|
# until the server exits/fails before startup. This intentionally has no
|
|
# correctness deadline: a slow but live startup should remain in progress.
|
|
try:
|
|
while not ready_event.is_set():
|
|
if startup_failed.is_set() or not thread.is_alive():
|
|
if startup_errors:
|
|
raise RuntimeError(
|
|
"Uvicorn server failed before startup completed"
|
|
) from startup_errors[0]
|
|
raise RuntimeError("Uvicorn server exited before startup completed")
|
|
ready_event.wait(timeout = 0.1)
|
|
except KeyboardInterrupt:
|
|
_graceful_shutdown(_server)
|
|
_shutdown_event.set()
|
|
raise
|
|
|
|
_write_pid_file()
|
|
import atexit
|
|
|
|
atexit.register(_remove_pid_file)
|
|
|
|
# Output port for Tauri to parse when in api-only mode. Emit only after
|
|
# uvicorn sockets are bound and FastAPI lifespan/startup has completed.
|
|
if api_only:
|
|
print(f"TAURI_PORT={port}", flush = True)
|
|
|
|
if not silent:
|
|
wildcard_bind = host in ("0.0.0.0", "::")
|
|
# For wildcard binds, run the reachability check between the URL
|
|
# section and the stop hint so the stop hint stays last on screen.
|
|
print_studio_access_banner(
|
|
port = port,
|
|
bind_host = host,
|
|
display_host = display_host,
|
|
include_stop_hint = not wildcard_bind,
|
|
)
|
|
if wildcard_bind:
|
|
_verify_global_reachability(display_host, port)
|
|
print_studio_stop_hint()
|
|
|
|
return app
|
|
|
|
|
|
# For direct execution (also invoked by CLI via os.execvp / subprocess)
|
|
if __name__ == "__main__":
|
|
import argparse
|
|
import signal
|
|
import traceback
|
|
|
|
# Ensure stderr can handle Unicode on Windows (tracebacks with non-ASCII paths)
|
|
if sys.platform == "win32" and hasattr(sys.stderr, "reconfigure"):
|
|
try:
|
|
sys.stderr.reconfigure(encoding = "utf-8", errors = "replace")
|
|
except Exception:
|
|
pass
|
|
|
|
parser = argparse.ArgumentParser(description = "Run Unsloth UI Backend server")
|
|
parser.add_argument(
|
|
"--host",
|
|
default = "127.0.0.1",
|
|
help = "Host to bind to (default: 127.0.0.1; use 0.0.0.0 for network/cloud access)",
|
|
)
|
|
parser.add_argument("--port", type = int, default = 8888, help = "Port to bind to")
|
|
parser.add_argument(
|
|
"--frontend",
|
|
type = str,
|
|
default = _DEFAULT_FRONTEND_PATH,
|
|
help = "Path to frontend build",
|
|
)
|
|
parser.add_argument("--silent", action = "store_true", help = "Suppress output")
|
|
parser.add_argument(
|
|
"--api-only",
|
|
action = "store_true",
|
|
help = "API server only, no frontend (for Tauri)",
|
|
)
|
|
# Mirror unsloth_cli/commands/studio.py's _PARALLEL_*. Default 1
|
|
# applies only to direct backend launches; `unsloth studio run`
|
|
# always passes its own value (4) explicitly.
|
|
_PARALLEL_MIN = 1
|
|
_PARALLEL_MAX = 64
|
|
_PARALLEL_DEFAULT_PLAIN = 1
|
|
parser.add_argument(
|
|
"--parallel",
|
|
"--n-parallel",
|
|
type = int,
|
|
default = _PARALLEL_DEFAULT_PLAIN,
|
|
help = (
|
|
f"llama-server parallel decode slots ({_PARALLEL_MIN}..{_PARALLEL_MAX}). "
|
|
f"Default {_PARALLEL_DEFAULT_PLAIN}; `unsloth studio run` uses 4."
|
|
),
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
if not _PARALLEL_MIN <= args.parallel <= _PARALLEL_MAX:
|
|
parser.error(f"--parallel must be between {_PARALLEL_MIN} and {_PARALLEL_MAX}")
|
|
|
|
kwargs = dict(
|
|
host = args.host,
|
|
port = args.port,
|
|
silent = args.silent,
|
|
api_only = args.api_only,
|
|
llama_parallel_slots = args.parallel,
|
|
)
|
|
if args.frontend is not None:
|
|
kwargs["frontend_path"] = Path(args.frontend)
|
|
|
|
try:
|
|
run_server(**kwargs)
|
|
except Exception:
|
|
sys.stderr.write("\n")
|
|
sys.stderr.write("=" * 60 + "\n")
|
|
sys.stderr.write("ERROR: Unsloth Studio failed to start.\n")
|
|
sys.stderr.write("=" * 60 + "\n")
|
|
traceback.print_exc(file = sys.stderr)
|
|
sys.stderr.write("\n")
|
|
sys.stderr.write(
|
|
"If a package is missing, try re-running: unsloth studio setup\n"
|
|
)
|
|
sys.stderr.flush()
|
|
sys.exit(1)
|
|
|
|
# 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)
|