* fix(studio/colab): restore iframe embed via serve_kernel_port_as_iframe Colab's output sanitizer often strips custom <iframe> tags from IPython.display.HTML without raising, leaving a blank cell even though display() succeeded. The kernel-port helper is the supported embedding path and registers the proxy correctly. - Prefer serve_kernel_port_as_iframe; keep raw HTML iframe as fallback - Always show the clickable link card via show_link() so the proxy URL is visible even when iframe embedding fails - Add regression tests for embed ordering and URL truncation Fixes #7344 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio/colab): harden iframe embed fallbacks per Codex review Guard show_link so a display failure cannot skip embedding, and only use serve_kernel_port_as_iframe when get_colab_url returned a real Colab proxy URL so localhost/colabtools environments still get the HTML iframe path. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio/colab): stop opening Colab proxy URLs in a new tab (#7349) Colab *.prod.colab.dev proxy hosts are session-scoped and return HTTP 404 when opened as a top-level tab or from another device. Replace the clickable Open button for those URLs with an in-notebook ready card, keep serve_kernel_port_as_iframe for the UI, and point users at start(cloudflare=True) for a real shareable / new-window link. * fix(studio/colab): use kernel iframe on real Colab when eval_js fails (#7349) Gate serve_kernel_port_as_iframe on COLAB_RELEASE_TAG + google.colab import instead of a successful proxyPort URL. When eval_js fails and get_colab_url falls back to localhost, real Colab notebooks still embed via the kernel helper (port-only). colabtools without COLAB_RELEASE_TAG keeps the HTML iframe path. Thanks @mfielding92 for the runtime diagnosis. * Mock top-level google package in Colab embed tests * test(studio/colab): mock top-level google package in Colab tests Patching only sys.modules["google.colab"] fails when no google namespace is installed: import google.colab resolves the parent first and returns False in _is_colab_runtime(). Add a shared helper that mocks both google and google.colab for deterministic tests across environments. * Tighten comments in Colab embed helpers and tests * fix(studio/colab): default Cloudflare on Colab with durable login credentials Colab proxy iframes often load an empty document even when the kernel helper appends the frame, leaving users unable to reach Studio to change the bootstrap password and blocking start(cloudflare=True). On real Colab runtime: - Default cloudflare to True (pass cloudflare=False to opt out) - Finalize the random admin password and print credentials in the notebook - Persist credentials across cell re-runs after interrupt - Show Cloudflare link before login credentials; skip blank proxy iframe when ready - Reuse main._IS_COLAB for runtime detection (not COLAB_RELEASE_TAG alone) - Only trust serve_kernel_port_as_iframe on real Colab; colabtools falls back to HTML - Keep embedding when the link card display fails Addresses Codex review feedback on #7349 and @mfielding92's catch-22 report. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio/colab): skip credential finalize when cloudflare=False Only call _finalize_colab_admin_password() when opening a Cloudflare tunnel. start(cloudflare=False) should not clear the bootstrap-password gate or show a login card that references a missing tunnel link. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio/colab): drop stale cached Colab credentials after password change On a Colab rerun the finalize path redisplayed the cached first-run password whenever the bootstrap gate was already cleared. If the admin changed the password through the app, that cached copy no longer authenticates, so the notebook printed dead credentials. Validate the cached password against the current stored hash before redisplaying and drop the cache when it no longer matches. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han <danielhanchen@gmail.com>
This commit is contained in:
Souravrajvi0
GitHub
parent
330586de7c
commit
d17567af3e
3 changed files with 1024 additions and 212 deletions
|
|
@ -12,7 +12,6 @@
|
|||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "6b87de59",
|
||||
"metadata": {
|
||||
"id": "6b87de59"
|
||||
},
|
||||
|
|
@ -34,61 +33,72 @@
|
|||
"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",
|
||||
"id": "e4206349",
|
||||
"metadata": {
|
||||
"id": "e4206349"
|
||||
},
|
||||
"source": [
|
||||
"<p align=\"left\"><img src=\"https://github.com/unslothai/unsloth/raw/main/studio/frontend/public/studio%20github%20landscape%20colab%20display.png\" width=\"600\"></p>"
|
||||
]
|
||||
],
|
||||
"id": "e4206349"
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "27da2957",
|
||||
"metadata": {
|
||||
"id": "27da2957"
|
||||
},
|
||||
"source": [
|
||||
"### Setup: Clone repo and run setup"
|
||||
]
|
||||
],
|
||||
"id": "27da2957"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "27e68f91",
|
||||
"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": [],
|
||||
"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"
|
||||
"id": "27e68f91"
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "3e1771a9",
|
||||
"metadata": {
|
||||
"id": "3e1771a9"
|
||||
},
|
||||
"source": [
|
||||
"### Start Unsloth Studio"
|
||||
]
|
||||
],
|
||||
"id": "3e1771a9"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "277e431e",
|
||||
"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": [],
|
||||
"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)"
|
||||
"id": "277e431e"
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "f2b0c6a1",
|
||||
"metadata": {
|
||||
"id": "f2b0c6a1"
|
||||
},
|
||||
|
|
@ -111,7 +121,8 @@
|
|||
"\n",
|
||||
" <b>This notebook is licensed <a href=\"https://github.com/unslothai/unsloth/blob/main/studio/LICENSE.AGPL-3.0\">AGPL-3.0</a></b>\n",
|
||||
"</div>"
|
||||
]
|
||||
],
|
||||
"id": "f2b0c6a1"
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
|
|
|
|||
|
|
@ -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"""
|
||||
<div style="display: inline-block; padding: 20px; background: #ffffff; border: 2px solid #000000;
|
||||
border-radius: 12px; margin: 10px 0; font-family: system-ui, -apple-system, sans-serif;">
|
||||
<h2 style="color: #000000; margin: 0 0 12px 0; font-size: 22px; font-weight: 800;">
|
||||
Unsloth Studio Login (Colab)
|
||||
</h2>
|
||||
<p style="color: #333333; margin: 0 0 12px 0; font-size: 14px; font-weight: bold;">
|
||||
Log in to Studio with the Cloudflare link above using these credentials. This cell
|
||||
is visible only in your notebook session.
|
||||
</p>
|
||||
<p style="color: #333333; margin: 0; font-size: 14px; font-family: monospace; font-weight: bold;">
|
||||
Username: <code>{username}</code><br>
|
||||
Password: <code>{password}</code>
|
||||
</p>
|
||||
</div>
|
||||
"""
|
||||
|
||||
|
||||
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 "
|
||||
'<code style="background:#f3f3f3;padding:2px 6px;border-radius:4px;">'
|
||||
"cloudflare=True</code> 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 <code style="background:#f3f3f3;padding:2px 6px;'
|
||||
'border-radius:4px;">start(cloudflare=True)</code> for a working link.'
|
||||
)
|
||||
return f"""
|
||||
<div style="display: inline-block; padding: 20px; background: #ffffff; border: 2px solid #000000;
|
||||
border-radius: 12px; margin: 10px 0; font-family: system-ui, -apple-system, sans-serif;">
|
||||
<h2 style="color: #000000; margin: 0 0 12px 0; font-size: 26px; font-weight: 800;
|
||||
display: flex; align-items: center; gap: 12px;">
|
||||
<img src="https://github.com/unslothai/unsloth/raw/main/studio/frontend/public/unsloth-gem.png"
|
||||
height="48" style="display:block;">
|
||||
Unsloth Studio is Ready!
|
||||
</h2>
|
||||
<p style="color: #333333; margin: 0 0 8px 0; font-size: 15px; font-weight: bold;">
|
||||
{embed_note}
|
||||
</p>
|
||||
<p style="color: #666666; margin: 16px 0 0 0; font-size: 13px; font-family: monospace; font-weight: bold;">
|
||||
{short_url}
|
||||
</p>
|
||||
</div>
|
||||
"""
|
||||
|
||||
return f"""
|
||||
<div style="display: inline-block; padding: 20px; background: #ffffff; border: 2px solid #000000;
|
||||
border-radius: 12px; margin: 10px 0; font-family: system-ui, -apple-system, sans-serif;">
|
||||
<h2 style="color: #000000; margin: 0 0 12px 0; font-size: 26px; font-weight: 800;
|
||||
|
|
@ -100,15 +312,52 @@ def show_link(port: int = 8888, *, _url: "str | None" = None):
|
|||
</p>
|
||||
</div>
|
||||
"""
|
||||
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
|
||||
|
||||
short_url = _short_colab_url(url, port)
|
||||
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)))
|
||||
|
||||
display(
|
||||
HTML(f"""
|
||||
<div style="font-family:system-ui,-apple-system,sans-serif;margin:8px 0;
|
||||
|
|
@ -266,41 +527,101 @@ def _show_and_embed(port: int, *, cloudflare_url: "str | None" = None):
|
|||
</div>
|
||||
""")
|
||||
)
|
||||
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):
|
||||
|
|
|
|||
479
studio/backend/tests/test_colab_embed.py
Normal file
479
studio/backend/tests/test_colab_embed.py
Normal file
|
|
@ -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"]
|
||||
Loading…
Add table
Add a link
Reference in a new issue