diff --git a/studio/Unsloth_Studio_Colab.ipynb b/studio/Unsloth_Studio_Colab.ipynb index c3aec04820..00eecfe51d 100644 --- a/studio/Unsloth_Studio_Colab.ipynb +++ b/studio/Unsloth_Studio_Colab.ipynb @@ -84,26 +84,7 @@ "id": "277e431e" }, "outputs": [], - "source": [ - "import sys, time\n", - "sys.path.insert(0, \"/content/unsloth/studio/backend\")\n", - "from colab import start\n", - "start()" - ] - }, - { - "cell_type": "code", - "source": [ - "from google.colab import output\n", - "output.serve_kernel_port_as_iframe(8888, height = 1200, width = \"100%\")\n", - "for _ in range(10000): time.sleep(300), print(\"=\", end = \"\")" - ], - "metadata": { - "id": "wb9UELh--XzX" - }, - "id": "wb9UELh--XzX", - "execution_count": null, - "outputs": [] + "source": "import sys\nsys.path.insert(0, \"/content/unsloth/studio/backend\")\nfrom colab import start\nstart()" }, { "cell_type": "markdown", @@ -150,4 +131,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} +} \ No newline at end of file diff --git a/studio/backend/colab.py b/studio/backend/colab.py index 7336f8a532..1bca16359b 100644 --- a/studio/backend/colab.py +++ b/studio/backend/colab.py @@ -26,30 +26,68 @@ logger = get_logger(__name__) def get_colab_url(port: int = 8888) -> str: """ Get the actual Colab proxy URL for a port. + + Retries up to 3 times and validates that the result is a real HTTPS Colab + URL before returning. Falls back to http://localhost:{port} only when all + attempts fail. """ + import time as _time + + fallback = f"http://localhost:{port}" + try: from google.colab.output import eval_js + except ImportError: + return fallback - # Use Colab's proxy mechanism - url = eval_js(f"google.colab.kernel.proxyPort({port})", timeout_sec = 5) - return url if url else f"http://localhost:{port}" - except Exception as e: - logger.info(f"Note: Could not get Colab URL ({e})") - return f"http://localhost:{port}" + for attempt in range(3): + try: + url = eval_js(f"google.colab.kernel.proxyPort({port})", timeout_sec = 10) + # A valid Colab proxy URL starts with https:// and embeds the port. + if ( + url + and isinstance(url, str) + and url.startswith("https://") + and str(port) in url + ): + return url.rstrip("/") + except Exception as e: + logger.info(f"Note: Could not get Colab URL (attempt {attempt + 1}/3: {e})") + if attempt < 2: + _time.sleep(1) + + logger.warning( + f"Could not get a valid Colab proxy URL after 3 attempts — using localhost fallback. " + f"The link/iframe may not work from outside the runtime." + ) + return fallback -def show_link(port: int = 8888): - """Display a styled clickable link to the UI.""" +def show_link(port: int = 8888, *, _url: "str | None" = None): + """Display a styled clickable link to the UI. + + *_url* is an optional pre-fetched Colab proxy URL. When omitted, + ``get_colab_url(port)`` is called internally. Pass it from + ``_show_and_embed`` to avoid a second ``eval_js`` round-trip. + """ from IPython.display import display, HTML - # Get real Colab proxy URL - url = get_colab_url(port) + url = _url if _url is not None else get_colab_url(port) + + # Build a truncated display URL. Wrap in try/except so an unexpected URL + # shape never prevents the link from rendering. + try: + port_prefix = f"{port}-" + idx = url.index(port_prefix) + next_dash = url.index("-", idx + len(port_prefix)) + short_url = url[: next_dash + 1] + "..." + except (ValueError, IndexError): + short_url = url + + # Also emit a plain-text line so the URL is visible even if HTML display + # is suppressed or fails. + logger.info(f"🌐 Unsloth Studio URL: {url}") - short_url = ( - url[: url.index("-", url.index(f"{port}-") + len(str(port)) + 1) + 1] + "..." - if f"{port}-" in url - else url - ) html = f"""
@@ -59,10 +97,10 @@ def show_link(port: int = 8888): height="48" style="display:block;"> Unsloth Studio is Ready! - + font-weight: 800; font-size: 16px; cursor: pointer;"> Open Unsloth Studio @@ -77,6 +115,75 @@ def show_link(port: int = 8888): display(HTML(html)) +def _is_studio_healthy(port: int, timeout: float = 2.0) -> bool: + """Return True if a Studio backend is already answering health checks on *port*.""" + import urllib.request + + try: + with urllib.request.urlopen( + f"http://localhost:{port}/api/health", timeout = timeout + ): + return True + except Exception: + return False + + +def _show_and_embed(port: int): + """Embed the Studio inline for *port* with a branded header bar. + + Fetches the Colab proxy URL once (registering the port with Colab's + reverse-proxy at the same time) then renders a header bar + full-height + iframe as a single HTML block. + + Falls back to ``serve_kernel_port_as_iframe`` if ``IPython.display.HTML`` + is unavailable for any reason. + """ + url = get_colab_url(port) + logger.info(f"🌐 Unsloth Studio URL: {url}") + + try: + from IPython.display import HTML, display + + iframe_id = f"unsloth-studio-{port}" + + # Truncated URL shown in the header — best-effort, falls back to full URL. + try: + port_prefix = f"{port}-" + idx = url.index(port_prefix) + next_dash = url.index("-", idx + len(port_prefix)) + short_url = url[: next_dash + 1] + "..." + except (ValueError, IndexError): + short_url = url + + display( + HTML(f""" +
+
+ + Unsloth Studio + {short_url} +
+ +
+""") + ) + except Exception: + # Fallback: Colab's built-in helper (less control, but always works) + try: + from google.colab import output as colab_output + + colab_output.serve_kernel_port_as_iframe(port, height = 900, width = "100%") + except ImportError: + pass + + def start(port: int = 8888): """ Start Unsloth Studio server in Colab and display the URL. @@ -85,10 +192,26 @@ def start(port: int = 8888): from colab import start start() """ - import sys + import time logger.info("🦥 Starting Unsloth Studio...") + # --- Fast path: Studio is already running (cell re-run) --- + # Re-launching would either collide on the port or silently shift to a new + # port and confuse the user. Just re-show the link and iframe instead. + if _is_studio_healthy(port): + logger.info( + f" Studio is already running on port {port} — reusing existing server." + ) + _show_and_embed(port) + try: + for _ in range(10000): + time.sleep(300) + print("=", end = "", flush = True) + except KeyboardInterrupt: + logger.info("\nUnsloth Studio keepalive stopped.") + return + logger.info(" Loading backend...") from run import run_server @@ -96,18 +219,63 @@ def start(port: int = 8888): repo_root = Path(__file__).parent.parent frontend_path = repo_root / "frontend" / "dist" - if not frontend_path.exists(): + if not (frontend_path / "index.html").exists(): logger.info("❌ Frontend not built! Please run the setup cell first.") return logger.info(" Starting server...") - # Start server silently - run_server(host = "0.0.0.0", port = port, frontend_path = frontend_path, silent = True) + try: + app = run_server( + host = "0.0.0.0", port = port, frontend_path = frontend_path, silent = True + ) + except SystemExit as exc: + logger.error(f"❌ Unsloth Studio failed to start: {exc}") + return + except Exception as exc: + logger.error(f"❌ Unsloth Studio failed to start: {exc}") + return - logger.info(" Server started!") + # run_server auto-increments the port when the requested one is already in + # use (e.g. Jupyter occupying 8888). Read back the actual bound port so the + # Colab proxy URL and iframe always point at the right place. + actual_port: int = getattr(getattr(app, "state", None), "server_port", None) or port - # Show the clickable link with real URL - show_link(port) + logger.info(f" Server started on port {actual_port}!") + + # Poll health endpoint to confirm the server is truly reachable before + # showing the link and registering the iframe — avoids the race where + # ready_event fires but the process hasn't finished binding. + import urllib.request + + server_ready = False + for _ in range(40): + try: + with urllib.request.urlopen( + f"http://localhost:{actual_port}/api/health", timeout = 1 + ): + server_ready = True + break + except Exception: + time.sleep(0.5) + + if not server_ready: + logger.error( + f"❌ Unsloth Studio did not become healthy on port {actual_port}. " + "Check for errors above." + ) + return + + _show_and_embed(actual_port) + + # Keep kernel alive so the daemon server thread stays running. + # Handle KeyboardInterrupt cleanly so the user gets a readable message + # rather than a raw traceback when they interrupt the cell. + try: + for _ in range(10000): + time.sleep(300) + print("=", end = "", flush = True) + except KeyboardInterrupt: + logger.info("\nUnsloth Studio keepalive stopped.") if __name__ == "__main__": diff --git a/studio/backend/main.py b/studio/backend/main.py index 8457da0b2d..fbf1c31de3 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -328,20 +328,58 @@ from starlette.requests import Request as _StarletteRequest # noqa: E402 _CSP_SCRIPT_NONCE_HEADER = "x-internal-script-nonce" +# /content is Colab's working directory — more reliable than env vars which +# aren't always set depending on Colab runtime version. +import importlib.util as _importlib_util + +_IS_COLAB = os.path.isdir("/content") and ( + bool(os.environ.get("COLAB_BACKEND_URL")) + or bool(os.environ.get("COLAB_JUPYTER_IP")) + or _importlib_util.find_spec("google.colab") is not None +) + + def _build_csp(script_nonce: "str | None" = None) -> str: script_src = "script-src 'self'" if script_nonce: script_src += f" 'nonce-{script_nonce}'" + # In Colab the parent frame can be colab.research.google.com, a multi-level + # *.prod.colab.dev subdomain (e.g. foo.region.prod.colab.dev — note: CSP + # wildcards only match one level, so *.prod.colab.dev misses these), or a + # sandboxed null-origin output iframe. Use '*' so any ancestor is allowed; + # Colab is already a sandboxed single-user environment. + frame_ancestors = "*" if _IS_COLAB else "'none'" + + # In Colab the frontend is served over the Colab reverse-proxy at an HTTPS + # *.prod.colab.dev URL. Colab's kernel communication layer and the output + # iframe scaffolding inject scripts from *.prod.colab.dev and + # *.googleusercontent.com, and make fetch/WebSocket connections to those + # same origins. Widen script-src and connect-src in Colab mode so those + # requests are not blocked. 'unsafe-inline' for scripts is still omitted; + # our own inline script uses a nonce. + if _IS_COLAB: + script_src += " https://*.prod.colab.dev https://*.googleusercontent.com" + connect_src = ( + "'self' blob: data: " + "https://huggingface.co https://datasets-server.huggingface.co " + "https://*.prod.colab.dev wss://*.prod.colab.dev " + "https://*.googleusercontent.com wss://*.googleusercontent.com" + ) + else: + connect_src = ( + "'self' https://huggingface.co https://datasets-server.huggingface.co" + ) + return ( "default-src 'self'; " "img-src 'self' data: blob: https://t0.gstatic.com " "https://t1.gstatic.com https://t2.gstatic.com " "https://t3.gstatic.com https://www.google.com; " - "connect-src 'self' https://huggingface.co https://datasets-server.huggingface.co; " + f"connect-src {connect_src}; " "style-src 'self' 'unsafe-inline'; " f"{script_src}; " "font-src 'self' data:; " - "frame-ancestors 'none'; " + f"frame-ancestors {frame_ancestors}; " "form-action 'self'; " "base-uri 'self'" ) @@ -357,7 +395,10 @@ class SecurityHeadersMiddleware(BaseHTTPMiddleware): if nonce is not None: del response.headers[_CSP_SCRIPT_NONCE_HEADER] response.headers.setdefault("Content-Security-Policy", _build_csp(nonce)) - response.headers.setdefault("X-Frame-Options", "DENY") + # Omit X-Frame-Options in Colab — CSP frame-ancestors handles it, and + # DENY would block serve_kernel_port_as_iframe regardless of CSP. + if not _IS_COLAB: + response.headers.setdefault("X-Frame-Options", "DENY") response.headers.setdefault("X-Content-Type-Options", "nosniff") response.headers.setdefault("Referrer-Policy", "no-referrer") response.headers.setdefault( diff --git a/studio/backend/run.py b/studio/backend/run.py index 5ddf404370..f401bc2ec1 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -653,7 +653,7 @@ def run_server( from threading import Thread, Event import uvicorn - from main import app, setup_frontend + from main import app, setup_frontend, _IS_COLAB from utils.paths import ensure_studio_directories # Create all standard directories on startup @@ -737,14 +737,22 @@ def run_server( ready_event.set() # server_header=False suppresses uvicorn's "Server: uvicorn"; SecurityHeadersMiddleware sets its own. - config = uvicorn.Config( - app, + config_kwargs = dict( host = host, port = port, log_level = "info", access_log = False, server_header = False, ) + # Only in Colab: trust X-Forwarded-* from Colab's reverse proxy so the app + # sees the real https origin. forwarded_allow_ips="*" is fine inside Colab's + # single-user sandbox, but would be an unwanted security relaxation for a + # normal local/standalone Studio, so leave uvicorn's safe defaults + # (forwarded headers trusted from loopback only) in place there. + if _IS_COLAB: + config_kwargs["proxy_headers"] = True + config_kwargs["forwarded_allow_ips"] = "*" + config = uvicorn.Config(app, **config_kwargs) _server = _ReadyServer(config) _shutdown_event = Event() @@ -766,14 +774,21 @@ def run_server( app.state.trigger_shutdown = _trigger_shutdown - # Run server in a daemon thread + # Run server in a daemon thread. + # Use an explicit new_event_loop() + run_until_complete() instead of + # asyncio.run() to avoid nest_asyncio's global patches to asyncio.run + # interfering when called from a thread while Colab/IPython already has + # a running loop on the main thread. def _run(): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) try: - asyncio.run(_server.serve()) + loop.run_until_complete(_server.serve()) except BaseException as exc: startup_errors.append(exc) startup_failed.set() finally: + loop.close() if not ready_event.is_set(): startup_failed.set()