# 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"""

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}

""" 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()