# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Colab-specific helpers for running Unsloth Studio.
Uses Colab's built-in proxy - no external tunneling needed!
"""
from pathlib import Path
import sys
# 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
_backend_dir = str(Path(__file__).parent)
if _backend_dir not in sys.path:
sys.path.insert(0, _backend_dir)
import _platform_compat # noqa: F401
from loggers import get_logger
logger = get_logger(__name__)
def get_colab_url(port: int = 8888) -> str:
"""
Get the actual Colab proxy URL for a port.
Retries up to 3 times and validates that the result is a real HTTPS Colab
URL before returning. Falls back to http://localhost:{port} only when all
attempts fail.
"""
import time as _time
fallback = f"http://localhost:{port}"
try:
from google.colab.output import eval_js
except ImportError:
return fallback
for attempt in range(3):
try:
url = eval_js(f"google.colab.kernel.proxyPort({port})", timeout_sec = 10)
# A valid Colab proxy URL starts with https:// and embeds the port.
if (
url
and isinstance(url, str)
and url.startswith("https://")
and str(port) in url
):
return url.rstrip("/")
except Exception as e:
logger.info(f"Note: Could not get Colab URL (attempt {attempt + 1}/3: {e})")
if attempt < 2:
_time.sleep(1)
logger.warning(
f"Could not get a valid Colab proxy URL after 3 attempts — using localhost fallback. "
f"The link/iframe may not work from outside the runtime."
)
return fallback
def show_link(port: int = 8888, *, _url: "str | None" = None):
"""Display a styled clickable link to the UI.
*_url* is an optional pre-fetched Colab proxy URL. When omitted,
``get_colab_url(port)`` is called internally. Pass it from
``_show_and_embed`` to avoid a second ``eval_js`` round-trip.
"""
from IPython.display import display, HTML
url = _url if _url is not None else get_colab_url(port)
# Build a truncated display URL. Wrap in try/except so an unexpected URL
# shape never prevents the link from rendering.
try:
port_prefix = f"{port}-"
idx = url.index(port_prefix)
next_dash = url.index("-", idx + len(port_prefix))
short_url = url[: next_dash + 1] + "..."
except (ValueError, IndexError):
short_url = url
# Also emit a plain-text line so the URL is visible even if HTML display
# is suppressed or fails.
logger.info(f"🌐 Unsloth Studio URL: {url}")
html = f"""
If the link doesn't work, you can scroll down to view the UI generated directly in Colab.
{short_url}
"""
display(HTML(html))
def _is_studio_healthy(port: int, timeout: float = 2.0) -> bool:
"""Return True if a Studio backend is already answering health checks on *port*."""
import urllib.request
try:
with urllib.request.urlopen(
f"http://localhost:{port}/api/health", timeout = timeout
):
return True
except Exception:
return False
def _show_and_embed(port: int):
"""Embed the Studio inline for *port* with a branded header bar.
Fetches the Colab proxy URL once (registering the port with Colab's
reverse-proxy at the same time) then renders a header bar + full-height
iframe as a single HTML block.
Falls back to ``serve_kernel_port_as_iframe`` if ``IPython.display.HTML``
is unavailable for any reason.
"""
url = get_colab_url(port)
logger.info(f"🌐 Unsloth Studio URL: {url}")
try:
from IPython.display import HTML, display
iframe_id = f"unsloth-studio-{port}"
# Truncated URL shown in the header — best-effort, falls back to full URL.
try:
port_prefix = f"{port}-"
idx = url.index(port_prefix)
next_dash = url.index("-", idx + len(port_prefix))
short_url = url[: next_dash + 1] + "..."
except (ValueError, IndexError):
short_url = url
display(
HTML(f"""
Unsloth Studio{short_url}
""")
)
except Exception:
# Fallback: Colab's built-in helper (less control, but always works)
try:
from google.colab import output as colab_output
colab_output.serve_kernel_port_as_iframe(port, height = 900, width = "100%")
except ImportError:
pass
def start(port: int = 8888):
"""
Start Unsloth Studio server in Colab and display the URL.
Usage:
from colab import start
start()
"""
import time
logger.info("🦥 Starting Unsloth Studio...")
# --- Fast path: Studio is already running (cell re-run) ---
# Re-launching would either collide on the port or silently shift to a new
# port and confuse the user. Just re-show the link and iframe instead.
if _is_studio_healthy(port):
logger.info(
f" Studio is already running on port {port} — reusing existing server."
)
_show_and_embed(port)
try:
for _ in range(10000):
time.sleep(300)
print("=", end = "", flush = True)
except KeyboardInterrupt:
logger.info("\nUnsloth Studio keepalive stopped.")
return
logger.info(" Loading backend...")
from run import run_server
# Auto-detect frontend path
repo_root = Path(__file__).parent.parent
frontend_path = repo_root / "frontend" / "dist"
if not (frontend_path / "index.html").exists():
logger.info("❌ Frontend not built! Please run the setup cell first.")
return
logger.info(" Starting server...")
try:
app = run_server(
host = "0.0.0.0", port = port, frontend_path = frontend_path, silent = True
)
except SystemExit as exc:
logger.error(f"❌ Unsloth Studio failed to start: {exc}")
return
except Exception as exc:
logger.error(f"❌ Unsloth Studio failed to start: {exc}")
return
# run_server auto-increments the port when the requested one is already in
# use (e.g. Jupyter occupying 8888). Read back the actual bound port so the
# Colab proxy URL and iframe always point at the right place.
actual_port: int = getattr(getattr(app, "state", None), "server_port", None) or port
logger.info(f" Server started on port {actual_port}!")
# Poll health endpoint to confirm the server is truly reachable before
# showing the link and registering the iframe — avoids the race where
# ready_event fires but the process hasn't finished binding.
import urllib.request
server_ready = False
for _ in range(40):
try:
with urllib.request.urlopen(
f"http://localhost:{actual_port}/api/health", timeout = 1
):
server_ready = True
break
except Exception:
time.sleep(0.5)
if not server_ready:
logger.error(
f"❌ Unsloth Studio did not become healthy on port {actual_port}. "
"Check for errors above."
)
return
_show_and_embed(actual_port)
# Keep kernel alive so the daemon server thread stays running.
# Handle KeyboardInterrupt cleanly so the user gets a readable message
# rather than a raw traceback when they interrupt the cell.
try:
for _ in range(10000):
time.sleep(300)
print("=", end = "", flush = True)
except KeyboardInterrupt:
logger.info("\nUnsloth Studio keepalive stopped.")
if __name__ == "__main__":
start()