# 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 helpers for Unsloth Studio. Uses Colab's built-in proxy.
"""
from pathlib import Path
import sys
# Seed platform._sys_version_cache before attrs->rich->structlog->platform crash on conda Python.
# 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 Colab proxy URL for a port.
Retries up to 3 times, validating the result is a real HTTPS Colab URL.
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)
# Valid proxy URL is 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 proxy URL; pass it 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)
# Truncated display URL; try/except so an odd URL shape still renders the link.
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
# Plain-text line so the URL shows even if HTML display 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 proxy URL once (registering the port), then renders header bar +
iframe. Falls back to serve_kernel_port_as_iframe if IPython HTML is unavailable.
"""
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 header URL — 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.
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 already running (cell re-run). Re-launching would collide on
# the port, so just re-show the link and iframe.
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 if in use; read back the bound port so the
# proxy URL and iframe 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 before showing the link — 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 runs; handle KeyboardInterrupt
# cleanly so interrupting the cell gives a readable message.
try:
for _ in range(10000):
time.sleep(300)
print("=", end = "", flush = True)
except KeyboardInterrupt:
logger.info("\nUnsloth Studio keepalive stopped.")
if __name__ == "__main__":
start()
Unsloth Studio is Ready!