diff --git a/studio/Unsloth_Studio_Colab.ipynb b/studio/Unsloth_Studio_Colab.ipynb index 44282b2255..612d739806 100644 --- a/studio/Unsloth_Studio_Colab.ipynb +++ b/studio/Unsloth_Studio_Colab.ipynb @@ -1,134 +1,145 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": { - "id": "view-in-github", - "colab_type": "text" - }, - "source": [ - "\"Open" - ] + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "view-in-github", + "colab_type": "text" + }, + "source": [ + "\"Open" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "6b87de59" + }, + "source": [ + "To run this, press \"*Runtime*\" and press \"*Run all*\" on a **free** Tesla T4 Google Colab instance!\n", + "
\n", + "\n", + "\n", + " Join Discord if you need help + ⭐ Star us on Github ⭐\n", + "
\n", + "\n", + "To install Unsloth Studio on your local device, follow [our guide](https://unsloth.ai/docs/new/unsloth-studio/install). Unsloth Studio is licensed [AGPL-3.0](https://github.com/unslothai/unsloth/blob/main/studio/LICENSE.AGPL-3.0).\n", + "\n", + "### Unsloth Studio\n", + "\n", + "Train and run open models with [**Unsloth Studio**](https://unsloth.ai/docs/new/unsloth-studio/start). NEW! Installation should now only take 2 mins!\n", + "\n", + "\n", + "We are actively working on making Unsloth Studio install on Colab T4 GPUs faster.\n", + "\n", + "[Features](https://unsloth.ai/docs/new/unsloth-studio#features) • [Quickstart](https://unsloth.ai/docs/new/unsloth-studio/start) • [Data Recipes](https://unsloth.ai/docs/new/unsloth-studio/data-recipe) • [Unsloth Chat](https://unsloth.ai/docs/new/unsloth-studio/chat) • [Export](https://unsloth.ai/docs/new/unsloth-studio/export)" + ], + "id": "6b87de59" + }, + { + "cell_type": "markdown", + "metadata": { + "id": "e4206349" + }, + "source": [ + "

" + ], + "id": "e4206349" + }, + { + "cell_type": "markdown", + "metadata": { + "id": "27da2957" + }, + "source": [ + "### Setup: Clone repo and run setup" + ], + "id": "27da2957" + }, + { + "cell_type": "code", + "metadata": { + "id": "27e68f91" + }, + "source": "!git clone --depth 1 --branch main https://github.com/unslothai/unsloth.git\n%cd /content/unsloth\n!chmod +x studio/setup.sh && ./studio/setup.sh --local", + "execution_count": null, + "outputs": [], + "id": "27e68f91" + }, + { + "cell_type": "markdown", + "metadata": { + "id": "3e1771a9" + }, + "source": [ + "### Start Unsloth Studio" + ], + "id": "3e1771a9" + }, + { + "cell_type": "code", + "metadata": { + "id": "277e431e" + }, + "source": [ + "import sys\n", + "sys.path.insert(0, \"/content/unsloth/studio/backend\")\n", + "from colab import start\n", + "\n", + "# On Colab, start() auto-opens a Cloudflare link and prints admin login credentials.\n", + "# Use the Cloudflare link above the ready card to open Studio (in-cell iframes often stay blank).\n", + "start()\n", + "\n", + "# To skip the Cloudflare tunnel and try the in-notebook proxy iframe only:\n", + "# start(cloudflare=False)" + ], + "execution_count": null, + "outputs": [], + "id": "277e431e" + }, + { + "cell_type": "markdown", + "metadata": { + "id": "f2b0c6a1" + }, + "source": [ + "And we're done! If you have any questions on Unsloth, we have a [Discord](https://discord.gg/unsloth) channel! If you find any bugs or want to keep updated with the latest LLM stuff, or need help, join projects etc, feel free to join our Discord!\n", + "\n", + "Some other resources:\n", + "1. Looking to use Unsloth locally? Read our [Installation Guide](https://unsloth.ai/docs/get-started/install) for details on installing Unsloth on Windows, Docker, AMD, Intel GPUs.\n", + "2. Learn how to do Reinforcement Learning with our [RL Guide and notebooks](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide).\n", + "3. Read our guides and notebooks for [Text-to-speech (TTS)](https://unsloth.ai/docs/basics/text-to-speech-tts-fine-tuning) and [vision](https://unsloth.ai/docs/basics/vision-fine-tuning) model support.\n", + "4. Explore our [LLM Tutorials Directory](https://unsloth.ai/docs/models/tutorials-how-to-fine-tune-and-run-llms) to find dedicated guides for each model.\n", + "5. Need help with Inference? Read our [Inference & Deployment page](https://unsloth.ai/docs/basics/inference-and-deployment) for details on using vLLM, llama.cpp, Ollama etc.\n", + "\n", + "
\n", + " \n", + " \n", + " \n", + "\n", + " Join Discord if you need help + ⭐️ Star us on Github ⭐️\n", + "\n", + " This notebook is licensed AGPL-3.0\n", + "
" + ], + "id": "f2b0c6a1" + } + ], + "metadata": { + "accelerator": "GPU", + "colab": { + "gpuType": "T4", + "provenance": [], + "include_colab_link": true + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "name": "python" + } }, - { - "cell_type": "markdown", - "id": "6b87de59", - "metadata": { - "id": "6b87de59" - }, - "source": [ - "To run this, press \"*Runtime*\" and press \"*Run all*\" on a **free** Tesla T4 Google Colab instance!\n", - "
\n", - "\n", - "\n", - " Join Discord if you need help + ⭐ Star us on Github ⭐\n", - "
\n", - "\n", - "To install Unsloth Studio on your local device, follow [our guide](https://unsloth.ai/docs/new/unsloth-studio/install). Unsloth Studio is licensed [AGPL-3.0](https://github.com/unslothai/unsloth/blob/main/studio/LICENSE.AGPL-3.0).\n", - "\n", - "### Unsloth Studio\n", - "\n", - "Train and run open models with [**Unsloth Studio**](https://unsloth.ai/docs/new/unsloth-studio/start). NEW! Installation should now only take 2 mins!\n", - "\n", - "\n", - "We are actively working on making Unsloth Studio install on Colab T4 GPUs faster.\n", - "\n", - "[Features](https://unsloth.ai/docs/new/unsloth-studio#features) • [Quickstart](https://unsloth.ai/docs/new/unsloth-studio/start) • [Data Recipes](https://unsloth.ai/docs/new/unsloth-studio/data-recipe) • [Unsloth Chat](https://unsloth.ai/docs/new/unsloth-studio/chat) • [Export](https://unsloth.ai/docs/new/unsloth-studio/export)" - ] - }, - { - "cell_type": "markdown", - "id": "e4206349", - "metadata": { - "id": "e4206349" - }, - "source": [ - "

" - ] - }, - { - "cell_type": "markdown", - "id": "27da2957", - "metadata": { - "id": "27da2957" - }, - "source": [ - "### Setup: Clone repo and run setup" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "27e68f91", - "metadata": { - "id": "27e68f91" - }, - "outputs": [], - "source": "!git clone --depth 1 --branch main https://github.com/unslothai/unsloth.git\n%cd /content/unsloth\n!chmod +x studio/setup.sh && ./studio/setup.sh --local" - }, - { - "cell_type": "markdown", - "id": "3e1771a9", - "metadata": { - "id": "3e1771a9" - }, - "source": [ - "### Start Unsloth Studio" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "277e431e", - "metadata": { - "id": "277e431e" - }, - "outputs": [], - "source": "import sys\nsys.path.insert(0, \"/content/unsloth/studio/backend\")\nfrom colab import start\n\n# Default: in-tab iframe only. start() blocks to keep the kernel alive.\nstart()\n\n# For a shareable Cloudflare link, replace start() above with:\n# start(cloudflare=True)" - }, - { - "cell_type": "markdown", - "id": "f2b0c6a1", - "metadata": { - "id": "f2b0c6a1" - }, - "source": [ - "And we're done! If you have any questions on Unsloth, we have a [Discord](https://discord.gg/unsloth) channel! If you find any bugs or want to keep updated with the latest LLM stuff, or need help, join projects etc, feel free to join our Discord!\n", - "\n", - "Some other resources:\n", - "1. Looking to use Unsloth locally? Read our [Installation Guide](https://unsloth.ai/docs/get-started/install) for details on installing Unsloth on Windows, Docker, AMD, Intel GPUs.\n", - "2. Learn how to do Reinforcement Learning with our [RL Guide and notebooks](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide).\n", - "3. Read our guides and notebooks for [Text-to-speech (TTS)](https://unsloth.ai/docs/basics/text-to-speech-tts-fine-tuning) and [vision](https://unsloth.ai/docs/basics/vision-fine-tuning) model support.\n", - "4. Explore our [LLM Tutorials Directory](https://unsloth.ai/docs/models/tutorials-how-to-fine-tune-and-run-llms) to find dedicated guides for each model.\n", - "5. Need help with Inference? Read our [Inference & Deployment page](https://unsloth.ai/docs/basics/inference-and-deployment) for details on using vLLM, llama.cpp, Ollama etc.\n", - "\n", - "
\n", - " \n", - " \n", - " \n", - "\n", - " Join Discord if you need help + ⭐️ Star us on Github ⭐️\n", - "\n", - " This notebook is licensed AGPL-3.0\n", - "
" - ] - } - ], - "metadata": { - "accelerator": "GPU", - "colab": { - "gpuType": "T4", - "provenance": [], - "include_colab_link": true - }, - "kernelspec": { - "display_name": "Python 3", - "name": "python3" - }, - "language_info": { - "name": "python" - } - }, - "nbformat": 4, - "nbformat_minor": 5 + "nbformat": 4, + "nbformat_minor": 5 } \ No newline at end of file diff --git a/studio/backend/colab.py b/studio/backend/colab.py index 1762469bcf..baa18a2fec 100644 --- a/studio/backend/colab.py +++ b/studio/backend/colab.py @@ -1,9 +1,7 @@ # 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. -""" +"""Colab helpers for Unsloth Studio. Uses Colab's built-in proxy.""" from pathlib import Path import sys @@ -22,11 +20,9 @@ logger = get_logger(__name__) def get_colab_url(port: int = 8888) -> str: - """ - Get the Colab proxy URL for a port. + """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. + Retries 3x validating a real HTTPS Colab URL; falls back to localhost on failure. """ import time as _time @@ -55,28 +51,244 @@ def get_colab_url(port: int = 8888) -> str: 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. +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)) - short_url = url[: next_dash + 1] + "..." + return url[: next_dash + 1] + "..." except (ValueError, IndexError): - short_url = url + return url - # Plain-text line so the URL shows even if HTML display fails. - logger.info(f"🌐 Unsloth Studio URL: {url}") - html = f""" +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") + 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().splitlines() + if len(lines) >= 2 and lines[0] and lines[1]: + return lines[0], lines[1] + except OSError 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 to Studio with the Cloudflare link above using these credentials. This cell + is visible only in your notebook session. +

+

+ Username: {username}
+ 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"""

""" - display(HTML(html)) + + +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, main.py injects that password into same-origin GETs, and a public - tunnel GET (no Origin) reads as same-origin, so sharing the link would leak admin - access. Fails safe to pending if the state cannot be read. + 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 @@ -121,9 +370,8 @@ def _bootstrap_password_pending() -> bool: 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 by design, so we start it directly. - Refused while the bootstrap password is pending; any failure collapses to None - and the Colab proxy still works. + 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( @@ -152,9 +400,9 @@ def start_cloudflare_tunnel(port: int) -> "str | None": 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 only sets this when it opens the tunnel itself, which it skips on Colab, - so we set it here. Otherwise the frontend's API examples fall back to an - unreachable server_url. Best-effort. + 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 @@ -183,8 +431,7 @@ def _stop_cloudflare_tunnel() -> None: 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 that merely serves /api/health. + The service-marker check stops the reuse path reusing or tunneling a foreign process. """ import json, urllib.request try: @@ -222,31 +469,45 @@ def _shareable_link_html(cloudflare_url: str) -> str: """ -def _show_and_embed(port: int, *, cloudflare_url: "str | None" = None): - """Render the Unsloth header + iframe for *port*, with a shareable-link card above - when *cloudflare_url* is set. Falls back to serve_kernel_port_as_iframe.""" - url = get_colab_url(port) - logger.info(f"🌐 Unsloth Studio URL: {url}") - if cloudflare_url: - logger.info(f"🔗 Shareable Cloudflare link: {cloudflare_url}") +# 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 - 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 - - if cloudflare_url: - display(HTML(_shareable_link_html(cloudflare_url))) - + short_url = _short_colab_url(url, port) + iframe_id = f"unsloth-studio-{port}" + try: display( HTML(f"""
""") ) - except Exception: - # Fallback: Colab's built-in helper. + 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, + ) + + if cloudflare_url: try: - from google.colab import output as colab_output - colab_output.serve_kernel_port_as_iframe(port, height = 900, width = "100%") - except ImportError: - pass + from IPython.display import HTML, display + display(HTML(_shareable_link_html(cloudflare_url))) + except Exception as e: + logger.info(f"Could not render Cloudflare link card ({e}).") + + if colab_login: + try: + _show_colab_login_credentials(*colab_login) + except Exception as e: + logger.info(f"Could not render Colab login card ({e}).") + + 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 = False): +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: Opt in to a shareable Cloudflare HTTPS link reachable from any - device (default OFF). It exposes Unsloth's login page beyond Colab, so it - stays an explicit opt-in; the default shows only the in-tab proxy iframe. + 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() # Colab-proxy iframe only (default) - start(cloudflare=True) # also open a shareable Cloudflare link + 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: Unsloth already running (cell re-run). Re-launching would collide on - # the port, so just re-show the link and iframe. + # 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: - cf_url = start_cloudflare_tunnel(port) if cloudflare else None + 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) + _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) @@ -313,7 +634,6 @@ def start(port: int = 8888, *, cloudflare: bool = False): 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" @@ -323,8 +643,7 @@ def start(port: int = 8888, *, cloudflare: bool = False): logger.info(" Starting server...") try: - # cloudflare=False: this helper owns the tunnel (Colab's own - # start(cloudflare=...) drives it), so pin it off explicitly. + # cloudflare=False: this helper owns the tunnel (via start(cloudflare=...)), so pin it off. app = run_server( host = "0.0.0.0", port = port, @@ -339,14 +658,12 @@ def start(port: int = 8888, *, cloudflare: bool = False): 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. + # 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 endpoint before showing the link — avoids the race where ready_event - # fires but the process hasn't finished binding. + # Poll health before showing the link: avoids the race where ready_event fires pre-bind. import urllib.request server_ready = False @@ -365,12 +682,17 @@ def start(port: int = 8888, *, cloudflare: bool = False): ) return - # Open the tunnel now the server is healthy, publish its URL for /api/health, and - # tear it down on interrupt (try/finally) rather than orphan the process. + # Server healthy: finalize Colab auth, open the tunnel, publish URL, tear down on interrupt. try: - cf_url = start_cloudflare_tunnel(actual_port) if cloudflare else None + 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) + _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): diff --git a/studio/backend/tests/test_colab_embed.py b/studio/backend/tests/test_colab_embed.py new file mode 100644 index 0000000000..dae0c7dae0 --- /dev/null +++ b/studio/backend/tests/test_colab_embed.py @@ -0,0 +1,479 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression coverage for Colab iframe embedding (#7344).""" + +import sys +import types +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import colab + + +def _mock_google_colab_modules(colab_mod): + """Mock ``google`` and ``google.colab`` for environments without Google packages.""" + google_mod = types.ModuleType("google") + google_mod.colab = colab_mod + return {"google": google_mod, "google.colab": colab_mod} + + +def test_short_colab_url_truncates_proxy_host(): + url = "https://8888-gpu-a100-s-kkb-usc1f0-9hzedjcxrlu8-f.us-central1-0.prod.colab.dev/" + assert colab._short_colab_url(url, 8888) == "https://8888-gpu-..." + + +def test_short_colab_url_falls_back_on_unexpected_shape(): + assert colab._short_colab_url("https://example.com", 8888) == "https://example.com" + + +def test_is_colab_proxy_url_requires_https_proxy(): + assert colab._is_colab_proxy_url("https://8888-test.prod.colab.dev/", 8888) is True + assert colab._is_colab_proxy_url("http://localhost:8888", 8888) is False + assert colab._is_colab_proxy_url("http://127.0.0.1:8888", 8888) is False + + +def test_ready_card_html_does_not_open_colab_proxy_in_new_tab(): + """Colab proxy hosts 404 as top-level tabs (#7349 reporter); never window.open them.""" + html = colab._ready_card_html("https://8888-test.prod.colab.dev/", 8888) + assert "window.open" not in html + assert 'href="https://8888-test.prod.colab.dev/"' not in html + assert "start(cloudflare=True)" in html + + +def test_ready_card_html_points_to_cloudflare_when_link_ready(monkeypatch): + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + html = colab._ready_card_html( + "https://8888-test.prod.colab.dev/", + 8888, + has_cloudflare_link = True, + ) + assert "Cloudflare link above" in html + + +def test_ready_card_html_warns_when_cloudflare_tunnel_missing(monkeypatch): + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + html = colab._ready_card_html( + "https://8888-test.prod.colab.dev/", + 8888, + cloudflare_requested = True, + ) + assert "Could not open a Cloudflare tunnel" in html + + +def test_warn_colab_cloudflare_missing_logs_on_colab_without_tunnel(monkeypatch): + warnings: list[str] = [] + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr(colab.logger, "warning", lambda msg, **kwargs: warnings.append(msg)) + colab._warn_colab_cloudflare_missing(use_cloudflare = True, cloudflare_url = None) + assert warnings + assert "Cloudflare tunnel unavailable" in warnings[0] + + +def test_warn_colab_cloudflare_missing_skips_when_tunnel_ready(monkeypatch, caplog): + import logging + + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + with caplog.at_level(logging.WARNING): + colab._warn_colab_cloudflare_missing( + use_cloudflare = True, + cloudflare_url = "https://share.trycloudflare.com", + ) + assert "Cloudflare tunnel unavailable" not in caplog.text + + +def test_is_colab_runtime_uses_backend_colab_detector(monkeypatch): + fake_main = types.ModuleType("main") + fake_main._IS_COLAB = True + monkeypatch.setitem(sys.modules, "main", fake_main) + assert colab._is_colab_runtime() is True + fake_main._IS_COLAB = False + assert colab._is_colab_runtime() is False + + +def test_ready_card_html_uses_cloudflare_hint_on_colab_runtime_localhost(monkeypatch): + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + html = colab._ready_card_html("http://localhost:8888", 8888) + assert "window.open" not in html + assert "start(cloudflare=True)" in html + + +def test_ready_card_html_keeps_open_button_for_localhost_outside_colab(monkeypatch): + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: False) + html = colab._ready_card_html("http://localhost:8888", 8888) + assert "window.open" in html + assert 'href="http://localhost:8888"' in html + assert "Open Unsloth Studio" in html + + +def test_embed_kernel_port_iframe_uses_colab_helper(monkeypatch): + colab_output = MagicMock() + google_colab = SimpleNamespace(output = colab_output) + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + with patch.dict("sys.modules", _mock_google_colab_modules(google_colab)): + assert colab._embed_kernel_port_iframe(8888) is True + colab_output.serve_kernel_port_as_iframe.assert_called_once_with( + 8888, + height = colab._COLAB_IFRAME_HEIGHT, + width = "100%", + ) + + +def test_embed_kernel_port_iframe_returns_false_without_colab(): + with patch.dict("sys.modules", _mock_google_colab_modules(None)): + assert colab._embed_kernel_port_iframe(8888) is False + + +def test_embed_kernel_port_iframe_skips_colabtools_without_runtime(monkeypatch): + """colabtools can queue JS without appending an iframe; only trust the helper on Colab.""" + colab_output = MagicMock() + google_colab = SimpleNamespace(output = colab_output) + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: False) + with patch.dict("sys.modules", _mock_google_colab_modules(google_colab)): + assert colab._embed_kernel_port_iframe(8888) is False + colab_output.serve_kernel_port_as_iframe.assert_not_called() + + +def test_show_and_embed_prefers_kernel_port_iframe(monkeypatch): + calls: list[str] = [] + + monkeypatch.setattr(colab, "get_colab_url", lambda port: f"https://{port}-test.prod.colab.dev/") + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr( + colab, + "show_link", + lambda port, + *, + _url = None, + has_cloudflare_link = False, + cloudflare_requested = False: calls.append("show_link"), + ) + monkeypatch.setattr( + colab, + "_embed_kernel_port_iframe", + lambda port: calls.append("kernel_iframe") or True, + ) + monkeypatch.setattr( + colab, + "_embed_html_iframe", + lambda url, port: calls.append("html_iframe") or True, + ) + + colab._show_and_embed(8888) + + assert calls == ["show_link", "kernel_iframe"] + + +def test_show_and_embed_falls_back_to_html_iframe(monkeypatch): + calls: list[str] = [] + + monkeypatch.setattr(colab, "get_colab_url", lambda port: f"https://{port}-test.prod.colab.dev/") + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: False) + monkeypatch.setattr( + colab, + "show_link", + lambda port, *, _url = None, has_cloudflare_link = False: None, + ) + monkeypatch.setattr(colab, "_embed_kernel_port_iframe", lambda port: False) + monkeypatch.setattr( + colab, + "_embed_html_iframe", + lambda url, port: calls.append((url, port)) or True, + ) + + colab._show_and_embed(8888) + + assert calls == [("https://8888-test.prod.colab.dev/", 8888)] + + +def test_colab_wants_cloudflare_auto_enables_on_runtime(monkeypatch): + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + assert colab._colab_wants_cloudflare(None) is True + assert colab._colab_wants_cloudflare(True) is True + assert colab._colab_wants_cloudflare(False) is False + + +def test_colab_wants_cloudflare_defaults_off_outside_runtime(monkeypatch): + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: False) + assert colab._colab_wants_cloudflare(None) is False + assert colab._colab_wants_cloudflare(True) is True + + +def test_finalize_colab_admin_password_skips_outside_runtime(monkeypatch): + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: False) + assert colab._finalize_colab_admin_password() is None + + +def test_finalize_colab_admin_password_clears_bootstrap_gate(monkeypatch): + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr(colab, "_load_colab_login_credentials", lambda: None) + stored: list[tuple[str, str]] = [] + monkeypatch.setattr( + colab, + "_store_colab_login_credentials", + lambda username, password: stored.append((username, password)), + ) + + storage = SimpleNamespace( + DEFAULT_ADMIN_USERNAME = "unsloth", + ensure_default_admin = MagicMock(), + get_bootstrap_password = MagicMock(return_value = "alpha-beta-gamma"), + generate_bootstrap_password = MagicMock(return_value = "alpha-beta-gamma"), + requires_password_change = MagicMock(return_value = True), + update_password = MagicMock(return_value = True), + ) + auth_pkg = types.ModuleType("auth") + auth_pkg.storage = storage + with patch.dict("sys.modules", {"auth": auth_pkg, "auth.storage": storage}): + result = colab._finalize_colab_admin_password() + + assert result == ("unsloth", "alpha-beta-gamma") + storage.ensure_default_admin.assert_called_once() + storage.update_password.assert_called_once_with("unsloth", "alpha-beta-gamma") + assert stored == [("unsloth", "alpha-beta-gamma")] + + +def test_start_skips_finalize_when_cloudflare_disabled(monkeypatch): + import time + + finalize_calls: list[str] = [] + monkeypatch.setattr(colab, "_is_studio_healthy", lambda port: True) + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr( + colab, + "_finalize_colab_admin_password", + lambda: finalize_calls.append("finalize") or ("unsloth", "secret"), + ) + monkeypatch.setattr( + colab, "start_cloudflare_tunnel", lambda port: "https://share.trycloudflare.com" + ) + monkeypatch.setattr(colab, "_publish_cloudflare_url", lambda url: None) + monkeypatch.setattr(colab, "_show_and_embed", lambda port, **kwargs: None) + monkeypatch.setattr(colab, "_stop_cloudflare_tunnel", lambda: None) + monkeypatch.setattr(time, "sleep", lambda _: (_ for _ in ()).throw(KeyboardInterrupt)) + + colab.start(cloudflare = False) + + assert finalize_calls == [] + + +def test_finalize_colab_admin_password_redisplay_on_rerun(monkeypatch): + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr( + colab, + "_load_colab_login_credentials", + lambda: ("unsloth", "saved-pass"), + ) + monkeypatch.setattr(colab, "_colab_credentials_still_valid", lambda username, password: True) + + storage = SimpleNamespace( + DEFAULT_ADMIN_USERNAME = "unsloth", + ensure_default_admin = MagicMock(), + get_bootstrap_password = MagicMock(), + generate_bootstrap_password = MagicMock(), + requires_password_change = MagicMock(return_value = False), + update_password = MagicMock(), + ) + auth_pkg = types.ModuleType("auth") + auth_pkg.storage = storage + with patch.dict("sys.modules", {"auth": auth_pkg, "auth.storage": storage}): + result = colab._finalize_colab_admin_password() + + assert result == ("unsloth", "saved-pass") + storage.update_password.assert_not_called() + + +def test_finalize_colab_admin_password_drops_stale_cached_credentials(monkeypatch): + """After an in-app password change the cached first-run password no longer + authenticates, so it must not be redisplayed (#7349 Codex review).""" + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr( + colab, + "_load_colab_login_credentials", + lambda: ("unsloth", "stale-pass"), + ) + monkeypatch.setattr(colab, "_colab_credentials_still_valid", lambda username, password: False) + cleared: list[bool] = [] + monkeypatch.setattr(colab, "_clear_colab_login_credentials", lambda: cleared.append(True)) + + storage = SimpleNamespace( + DEFAULT_ADMIN_USERNAME = "unsloth", + ensure_default_admin = MagicMock(), + get_bootstrap_password = MagicMock(), + generate_bootstrap_password = MagicMock(), + requires_password_change = MagicMock(return_value = False), + update_password = MagicMock(), + ) + auth_pkg = types.ModuleType("auth") + auth_pkg.storage = storage + with patch.dict("sys.modules", {"auth": auth_pkg, "auth.storage": storage}): + result = colab._finalize_colab_admin_password() + + assert result is None + assert cleared == [True] + storage.update_password.assert_not_called() + + +def test_colab_credentials_still_valid_matches_stored_hash(monkeypatch): + from auth.hashing import hash_password + + salt, pwd_hash = hash_password("right-pass") + storage = SimpleNamespace( + get_user_and_secret = MagicMock(return_value = (salt, pwd_hash, "jwt", False)), + ) + with patch.dict("sys.modules", {"auth.storage": storage}): + assert colab._colab_credentials_still_valid("unsloth", "right-pass") is True + assert colab._colab_credentials_still_valid("unsloth", "wrong-pass") is False + + +def test_colab_credentials_still_valid_false_when_user_missing(monkeypatch): + storage = SimpleNamespace(get_user_and_secret = MagicMock(return_value = None)) + with patch.dict("sys.modules", {"auth.storage": storage}): + assert colab._colab_credentials_still_valid("unsloth", "any") is False + + +def test_colab_login_html_includes_credentials(): + html = colab._colab_login_html("unsloth", "alpha-beta-gamma-delta") + assert "unsloth" in html + assert "alpha-beta-gamma-delta" in html + + +def test_show_and_embed_renders_cloudflare_before_colab_login(monkeypatch): + displayed: list[str] = [] + ipython_display = SimpleNamespace( + HTML = lambda html: SimpleNamespace(html = html), + display = lambda html: displayed.append(html.html), + ) + + monkeypatch.setattr(colab, "get_colab_url", lambda port: "https://8888-test.prod.colab.dev/") + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr( + colab, + "show_link", + lambda port, *, _url = None, has_cloudflare_link = False, cloudflare_requested = False: None, + ) + monkeypatch.setattr(colab, "_embed_kernel_port_iframe", lambda port: True) + with patch.dict("sys.modules", {"IPython.display": ipython_display}): + colab._show_and_embed( + 8888, + cloudflare_url = "https://share.trycloudflare.com", + colab_login = ("unsloth", "secret-pass"), + ) + + assert len(displayed) == 2 + assert "share.trycloudflare.com" in displayed[0] + assert "secret-pass" in displayed[1] + + +def test_show_and_embed_skips_iframe_on_colab_when_cloudflare_ready(monkeypatch): + calls: list[str] = [] + + monkeypatch.setattr(colab, "get_colab_url", lambda port: f"https://{port}-test.prod.colab.dev/") + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr( + colab, + "show_link", + lambda port, *, _url = None, has_cloudflare_link = False, cloudflare_requested = False: None, + ) + monkeypatch.setattr( + colab, + "_embed_kernel_port_iframe", + lambda port: calls.append("kernel_iframe") or True, + ) + monkeypatch.setattr( + colab, + "_embed_html_iframe", + lambda url, port: calls.append("html_iframe") or True, + ) + + colab._show_and_embed(8888, cloudflare_url = "https://share.trycloudflare.com") + + assert calls == [] + + +def test_show_and_embed_uses_kernel_helper_on_colab_runtime_despite_localhost(monkeypatch): + calls: list[str] = [] + + monkeypatch.setattr(colab, "get_colab_url", lambda port: f"http://localhost:{port}") + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + + monkeypatch.setattr( + colab, + "show_link", + lambda port, + *, + _url = None, + has_cloudflare_link = False, + cloudflare_requested = False: calls.append("show_link"), + ) + monkeypatch.setattr( + colab, + "_embed_kernel_port_iframe", + lambda port: calls.append("kernel_iframe") or True, + ) + monkeypatch.setattr( + colab, + "_embed_html_iframe", + lambda url, port: calls.append("html_iframe") or True, + ) + + colab._show_and_embed(8888) + + assert calls == ["show_link", "kernel_iframe"] + + +def test_show_and_embed_skips_kernel_helper_for_localhost_outside_colab(monkeypatch): + calls: list[str] = [] + + monkeypatch.setattr(colab, "get_colab_url", lambda port: f"http://localhost:{port}") + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: False) + + monkeypatch.setattr( + colab, + "show_link", + lambda port, + *, + _url = None, + has_cloudflare_link = False, + cloudflare_requested = False: calls.append("show_link"), + ) + monkeypatch.setattr( + colab, + "_embed_kernel_port_iframe", + lambda port: calls.append("kernel_iframe") or True, + ) + monkeypatch.setattr( + colab, + "_embed_html_iframe", + lambda url, port: calls.append("html_iframe") or True, + ) + + colab._show_and_embed(8888) + + assert calls == ["show_link", "html_iframe"] + + +def test_show_and_embed_still_embeds_when_show_link_fails(monkeypatch): + calls: list[str] = [] + + monkeypatch.setattr(colab, "get_colab_url", lambda port: f"https://{port}-test.prod.colab.dev/") + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr( + colab, + "show_link", + lambda port, *, _url = None: (_ for _ in ()).throw(RuntimeError("no display")), + ) + monkeypatch.setattr( + colab, + "_embed_kernel_port_iframe", + lambda port: calls.append("kernel_iframe") or True, + ) + monkeypatch.setattr( + colab, + "_embed_html_iframe", + lambda url, port: calls.append("html_iframe") or True, + ) + + colab._show_and_embed(8888) + + assert calls == ["kernel_iframe"]