# 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 3x validating a real HTTPS Colab URL; falls back to localhost on failure. """ 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 _short_colab_url(url: str, port: int) -> str: """Truncated display form of a Colab proxy URL; falls back to the full URL.""" try: port_prefix = f"{port}-" idx = url.index(port_prefix) next_dash = url.index("-", idx + len(port_prefix)) return url[: next_dash + 1] + "..." except (ValueError, IndexError): return url def _is_colab_proxy_url(url: str, port: int) -> bool: """True when *url* looks like a real Colab kernel proxy, not a localhost fallback.""" return bool(url and isinstance(url, str) and url.startswith("https://") and str(port) in url) def _is_colab_runtime() -> bool: """True on a hosted Colab notebook kernel. Reuses the backend's main Colab detector (``/content`` + Colab env / ``google.colab``) instead of a single env var, which is not always present on hosted runtimes. """ try: from main import _IS_COLAB return bool(_IS_COLAB) except Exception: return False def _colab_login_credentials_path() -> Path: from auth.storage import DB_PATH return DB_PATH.parent / ".colab_notebook_login" def _store_colab_login_credentials(username: str, password: str) -> None: """Persist Colab admin credentials for notebook re-runs after interrupt.""" path = _colab_login_credentials_path() try: path.parent.mkdir(parents = True, exist_ok = True) path.write_text(f"{username}\n{password}\n", encoding = "utf-8") try: import os os.chmod(path, 0o600) except OSError: pass except OSError as e: logger.info(f"Could not persist Colab login credentials ({e}).") def _load_colab_login_credentials() -> "tuple[str, str] | None": """Return stored Colab admin credentials from a previous ``start()`` run, if any.""" path = _colab_login_credentials_path() try: if not path.is_file(): return None lines = path.read_text(encoding = "utf-8").splitlines() if len(lines) >= 2 and lines[0] and lines[1]: return lines[0], lines[1] except (OSError, UnicodeDecodeError) as e: logger.info(f"Could not load Colab login credentials ({e}).") return None def _clear_colab_login_credentials() -> None: """Drop the cached Colab credentials once they no longer authenticate.""" path = _colab_login_credentials_path() try: path.unlink(missing_ok = True) except OSError as e: logger.info(f"Could not clear Colab login credentials ({e}).") def _colab_credentials_still_valid(username: str, password: str) -> bool: """True when *password* still matches the stored admin hash. Guards against redisplaying a cached first-run password after the user has changed the admin password through the app, which would print credentials that no longer authenticate to the current Cloudflare tunnel. """ try: from auth.storage import get_user_and_secret from auth.hashing import verify_password except Exception as e: logger.info(f"Could not load auth to validate cached Colab credentials ({e}).") return False try: row = get_user_and_secret(username) if not row: return False salt, pwd_hash = row[0], row[1] return bool(verify_password(password, salt, pwd_hash)) except Exception as e: logger.info(f"Could not validate cached Colab credentials ({e}).") return False def _colab_wants_cloudflare(cloudflare: "bool | None") -> bool: """Resolve whether to open a Cloudflare tunnel. ``None`` auto-enables on real Colab (the in-cell proxy embed is often blank); pass ``False`` to opt out. """ if cloudflare is not None: return cloudflare return _is_colab_runtime() def _finalize_colab_admin_password() -> "tuple[str, str] | None": """Clear the bootstrap-password gate on Colab so Cloudflare tunnels can start. Returns ``(username, password)`` for display in the notebook. On first run the random admin password is finalized; on later runs (e.g. after interrupt) the stored credentials are re-displayed so the Cloudflare link stays usable. Anyone who can read this cell already controls the runtime. """ if not _is_colab_runtime(): return None try: from auth.storage import ( DEFAULT_ADMIN_USERNAME, ensure_default_admin, generate_bootstrap_password, get_bootstrap_password, requires_password_change, update_password, ) except Exception as e: logger.warning( f"Could not load auth for Colab setup ({e}); Cloudflare link may be blocked." ) return None try: ensure_default_admin() username = DEFAULT_ADMIN_USERNAME if not requires_password_change(username): creds = _load_colab_login_credentials() if creds is not None and _colab_credentials_still_valid(username, creds[1]): return creds # The admin password was changed through the app after the first run, # so the cached copy is stale; drop it instead of printing dead credentials. _clear_colab_login_credentials() return None password = get_bootstrap_password() or generate_bootstrap_password() if not update_password(username, password): logger.warning( "Could not finalize Colab admin password; Cloudflare link may be blocked." ) return None _store_colab_login_credentials(username, password) return username, password except Exception as e: logger.warning( f"Could not finalize Colab admin password ({e}); Cloudflare link may be blocked." ) return None def _colab_login_html(username: str, password: str) -> str: """Notebook card with Colab admin credentials (shown once after auto-finalize).""" return f"""

Unsloth Studio Login (Colab)

Log in as {username} with this password. This cell is visible only in your notebook session.

Password: {password}

""" def _show_colab_login_credentials(username: str, password: str) -> None: """Display Colab admin credentials in the notebook output.""" from IPython.display import HTML, display logger.info(f"🔐 Unsloth Studio login — user: {username}") display(HTML(_colab_login_html(username, password))) def _ready_card_html( url: str, port: int, *, has_cloudflare_link: bool = False, cloudflare_requested: bool = False, ) -> str: """Branded ready card for the in-notebook Studio view. Colab ``*.prod.colab.dev`` proxy URLs are session-scoped and 404 when opened as a top-level tab or on another device, so never ``window.open`` them. On real Colab the Cloudflare link is the supported entry point because in-cell proxy embeds often stay blank. """ short_url = _short_colab_url(url, port) if _is_colab_runtime() or _is_colab_proxy_url(url, port): if has_cloudflare_link: embed_note = ( "Open Studio with the Cloudflare link above. In-cell proxy previews on " "current Colab often stay blank, so the tunnel link is the supported path." ) elif cloudflare_requested: embed_note = ( "Could not open a Cloudflare tunnel, so Studio may be unreachable on Colab. " "Check the logs above and re-run this cell. Pass " '' "cloudflare=True after fixing any tunnel errors." ) else: embed_note = ( "Colab proxy links cannot be opened in a new tab (they 404 outside this " 'notebook). Re-run with start(cloudflare=True) for a working link.' ) return f"""

Unsloth Studio is Ready!

{embed_note}

{short_url}

""" return f"""

Unsloth Studio is Ready!

Open Unsloth Studio

If the link doesn't work, you can scroll down to view the UI generated directly in Colab.

{short_url}

""" def show_link( port: int = 8888, *, _url: "str | None" = None, has_cloudflare_link: bool = False, cloudflare_requested: bool = False, ): """Display a styled ready card for the UI. Colab proxy URLs are informational only (no new-tab open; they 404 outside the cell); non-proxy URLs keep a clickable open button. *_url* is an optional pre-fetched proxy URL 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) logger.info(f"🌐 Unsloth Studio URL: {url}") display( HTML( _ready_card_html( url, port, has_cloudflare_link = has_cloudflare_link, cloudflare_requested = cloudflare_requested, ) ) ) def _warn_colab_cloudflare_missing(*, use_cloudflare: bool, cloudflare_url: "str | None") -> None: """Log a prominent warning when Colab expected a tunnel but none was opened.""" if not use_cloudflare or cloudflare_url or not _is_colab_runtime(): return logger.warning( "Colab Cloudflare tunnel unavailable — Studio is unlikely to be reachable in this " "notebook. Check the logs above for tunnel or auth errors, then re-run start()." ) def _bootstrap_password_pending() -> bool: """True while the default admin still owes a bootstrap-password change. While pending, a public tunnel GET (no Origin) reads as same-origin and gets the injected password, so sharing the link would leak admin access. Fails safe to pending. """ try: from auth.storage import requires_password_change, DEFAULT_ADMIN_USERNAME return bool(requires_password_change(DEFAULT_ADMIN_USERNAME)) except Exception as e: logger.info(f"Could not check admin password state ({e}); refusing tunnel to be safe.") return True def start_cloudflare_tunnel(port: int) -> "str | None": """Open a shareable Cloudflare quick tunnel to localhost:*port*, or None. run_server suppresses the tunnel on Colab, so we start it directly. Refused while the bootstrap password is pending; any failure collapses to None (Colab proxy still works). """ if _bootstrap_password_pending(): logger.warning( "Cloudflare link not started: the admin account still has its temporary " "bootstrap password, which is exposed to anyone who can load the page. " "Open Unsloth in this tab, log in and change the admin password, then re-run " "start(cloudflare=True) to get the shareable link." ) return None try: from cloudflare_tunnel import start_studio_tunnel except Exception as e: logger.info(f"Cloudflare tunnel unavailable ({e}); using Colab proxy only.") return None try: url = start_studio_tunnel(port) except Exception as e: logger.info(f"Cloudflare tunnel failed to start ({e}); using Colab proxy only.") return None # Success is logged by _show_and_embed; note only misses here. if not url: logger.info("Cloudflare tunnel did not produce a URL; using Colab proxy only.") return url def _publish_cloudflare_url(cloudflare_url: "str | None") -> None: """Publish a directly-started tunnel URL onto app.state so /api/health advertises it. run_server sets this only when it opens the tunnel itself (skipped on Colab), so we set it here; otherwise the frontend's API examples fall back to an unreachable server_url. Best-effort. """ if not cloudflare_url: return try: from main import app as _studio_app _studio_app.state.cloudflare_url = cloudflare_url except Exception as e: logger.info(f"Could not publish Cloudflare URL to /api/health ({e}).") def _stop_cloudflare_tunnel() -> None: """Best-effort teardown of the Cloudflare tunnel started by start_cloudflare_tunnel.""" try: from cloudflare_tunnel import stop_studio_tunnel stop_studio_tunnel() except Exception: pass # Stop /api/health advertising a dead tunnel. try: from main import app as _studio_app _studio_app.state.cloudflare_url = None except Exception: pass def _is_studio_healthy(port: int, timeout: float = 2.0) -> bool: """True only if Unsloth Studio (not some other app) answers /api/health on *port*. The service-marker check stops the reuse path reusing or tunneling a foreign process. """ import json, urllib.request try: with urllib.request.urlopen(f"http://localhost:{port}/api/health", timeout = timeout) as r: return json.loads(r.read()).get("service") == "Unsloth UI Backend" except Exception: return False def _shareable_link_html( cloudflare_url: str, password: "str | None" = None, username: "str | None" = None, ) -> str: """Branded card for the shareable Cloudflare link, styled like the show_link banner. *password* renders under the link so the credential sits in the card with the button it unlocks. The username is always the default admin, so it reads inline. """ login_block = "" if password: login_block = f"""

Password

{password}

Log in as {username} with this password. Shown only in your notebook session, and never included in the shared link.

""" return f"""

Shareable Unsloth Link is Ready!

Open Unsloth Studio

This Cloudflare HTTPS link works from any device, so you can share it with anyone.

🔗 {cloudflare_url}

{login_block}
""" # Height for serve_kernel_port_as_iframe (~82vh on a 1080p screen, clamped). _COLAB_IFRAME_HEIGHT = 900 def _embed_kernel_port_iframe(port: int) -> bool: """Embed Studio via Colab's native kernel-port iframe helper. Only trusted on a real Colab runtime: colabtools can import ``google.colab`` and queue browser-side JS without appending an iframe, so callers outside Colab must use the HTML iframe path instead. """ if not _is_colab_runtime(): return False try: from google.colab import output as colab_output except ImportError: return False try: colab_output.serve_kernel_port_as_iframe( port, height = _COLAB_IFRAME_HEIGHT, width = "100%", ) return True except Exception as e: logger.info(f"serve_kernel_port_as_iframe failed ({e}); trying HTML iframe.") return False def _embed_html_iframe(url: str, port: int) -> bool: """Fallback embed: raw HTML iframe when the Colab helper is unavailable.""" try: from IPython.display import HTML, display except ImportError: return False short_url = _short_colab_url(url, port) iframe_id = f"unsloth-studio-{port}" try: display( HTML(f"""
Unsloth Studio {short_url}
""") ) return True except Exception as e: logger.info(f"HTML iframe embed failed ({e}).") return False def _show_and_embed( port: int, *, cloudflare_url: "str | None" = None, colab_login: "tuple[str, str] | None" = None, cloudflare_requested: bool = False, ): """Render the Unsloth ready card + iframe for *port*. Prefer Colab's ``serve_kernel_port_as_iframe`` on real Colab; raw HTML iframe is the fallback. Cloudflare cards stay clickable. """ url = get_colab_url(port) logger.info(f"🌐 Unsloth Studio URL: {url}") if cloudflare_url: logger.info(f"🔗 Shareable Cloudflare link: {cloudflare_url}") _warn_colab_cloudflare_missing( use_cloudflare = cloudflare_requested, cloudflare_url = cloudflare_url, ) # Fold the credentials into the link card rather than a second card below it. credentials_shown = False if cloudflare_url: try: from IPython.display import HTML, display username, password = colab_login if colab_login else (None, None) display(HTML(_shareable_link_html(cloudflare_url, password, username))) credentials_shown = bool(colab_login) except Exception as e: logger.info(f"Could not render Cloudflare link card ({e}).") if colab_login and not credentials_shown: try: _show_colab_login_credentials(*colab_login) except Exception as e: logger.info(f"Could not render Colab login card ({e}).") # With a tunnel up the embed below is skipped, so the ready card would only restate # the link card and print a proxy URL that 404s outside this tab. skip_ready_card = _is_colab_runtime() and bool(cloudflare_url) if not skip_ready_card: try: show_link( port, _url = url, has_cloudflare_link = bool(cloudflare_url), cloudflare_requested = cloudflare_requested, ) except Exception as e: logger.info(f"Could not render Unsloth link card ({e}).") # On Colab with a working tunnel, skip the in-cell proxy embed (often blank). if _is_colab_runtime() and cloudflare_url: return # Real Colab: kernel helper needs only the port (works when eval_js failed). if _is_colab_runtime(): if _embed_kernel_port_iframe(port): return _embed_html_iframe(url, port) def start(port: int = 8888, *, cloudflare: "bool | None" = None): """Start Unsloth Studio in Colab and display the URL. Args: port: Port to bind/serve on. cloudflare: Shareable Cloudflare HTTPS link. ``None`` (default) auto-enables on real Colab because the in-cell proxy embed is often blank; pass ``False`` to skip the tunnel or ``True`` to force it on other runtimes. Usage: start() # Cloudflare link on Colab (auto); proxy iframe elsewhere start(cloudflare=False) # Colab proxy iframe only (often blank on current Colab) start(cloudflare=True) # force Cloudflare link on any runtime """ import time logger.info("🦥 Starting Unsloth Studio...") use_cloudflare = _colab_wants_cloudflare(cloudflare) # Fast path: already running (cell re-run); re-show link/iframe instead of rebinding the port. if _is_studio_healthy(port): logger.info(f" Unsloth is already running on port {port} — reusing existing server.") # try/finally: tear the tunnel down even if interrupted mid-start/render. try: colab_login = _finalize_colab_admin_password() if use_cloudflare else None cf_url = start_cloudflare_tunnel(port) if use_cloudflare else None _publish_cloudflare_url(cf_url) _show_and_embed( port, cloudflare_url = cf_url, colab_login = colab_login, cloudflare_requested = use_cloudflare, ) for _ in range(10000): time.sleep(300) print("=", end = "", flush = True) except KeyboardInterrupt: logger.info("\nUnsloth Studio keepalive stopped.") finally: _stop_cloudflare_tunnel() return logger.info(" Loading backend...") from run import run_server 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: # cloudflare=False: this helper owns the tunnel (via start(cloudflare=...)), so pin it off. app = run_server( host = "0.0.0.0", port = port, frontend_path = frontend_path, silent = True, cloudflare = False, ) 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 may auto-increment the port; read back the bound port for the proxy URL/iframe. actual_port: int = getattr(getattr(app, "state", None), "server_port", None) or port logger.info(f" Server started on port {actual_port}!") # Poll health before showing the link: avoids the race where ready_event fires pre-bind. 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 # Server healthy: finalize Colab auth, open the tunnel, publish URL, tear down on interrupt. try: colab_login = _finalize_colab_admin_password() if use_cloudflare else None cf_url = start_cloudflare_tunnel(actual_port) if use_cloudflare else None _publish_cloudflare_url(cf_url) _show_and_embed( actual_port, cloudflare_url = cf_url, colab_login = colab_login, cloudflare_requested = use_cloudflare, ) # Keep kernel alive so the daemon server thread runs. for _ in range(10000): time.sleep(300) print("=", end = "", flush = True) except KeyboardInterrupt: logger.info("\nUnsloth Studio keepalive stopped.") finally: _stop_cloudflare_tunnel() if __name__ == "__main__": start()