From 23eb7fc0a7d01749101a4b7fe0c1468dc3345ea7 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 25 Mar 2026 09:00:08 -0700 Subject: [PATCH] Fix Colab Studio launch and setup.ps1 box alignment (#4601) * Fix Colab Studio launch and setup.ps1 box alignment - colab.py: when the Studio venv is missing on Colab, pip-install backend dependencies (structlog, fastapi, etc.) from studio.txt into the current Python instead of failing with ModuleNotFoundError - setup.sh: on Colab without a venv, install backend deps into system Python and skip venv-dependent sections (Python stack update, llama.cpp build) that would otherwise fail - setup.ps1: use PadRight(47) for the done-line so "Setup Complete!" and "Update Complete!" both align with the box border * [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> --- studio/backend/colab.py | 31 +++++++++++++++++++++++++++++++ studio/setup.ps1 | 3 ++- studio/setup.sh | 38 ++++++++++++++++++++++++++++++++------ 3 files changed, 65 insertions(+), 7 deletions(-) diff --git a/studio/backend/colab.py b/studio/backend/colab.py index ecf9fc2907..ad9a7e98b2 100644 --- a/studio/backend/colab.py +++ b/studio/backend/colab.py @@ -18,6 +18,30 @@ if _backend_dir not in sys.path: import _platform_compat # noqa: F401 +def _is_colab() -> bool: + """Detect Google Colab by checking for COLAB_ prefixed env vars.""" + import os + + return any(k.startswith("COLAB_") for k in os.environ) + + +def _pip_install_backend_deps() -> None: + """Install Studio backend dependencies directly into the current Python. + + Used on Colab when the Studio venv does not exist (install.sh was not + run). Reads the requirements from studio.txt next to this file. + """ + import subprocess + + req_file = Path(__file__).parent / "requirements" / "studio.txt" + if not req_file.exists(): + return + print("Installing Studio backend dependencies ...") + subprocess.check_call( + [sys.executable, "-m", "pip", "install", "-q", "-r", str(req_file)], + ) + + def _bootstrap_studio_venv() -> None: """Expose the Studio venv's site-packages to the current interpreter. @@ -25,9 +49,16 @@ def _bootstrap_studio_venv() -> None: installing the full stack into system Python, we prepend the venv's site-packages so that packages like structlog, fastapi, etc. are importable from notebook cells and take priority over system copies. + + If the venv does not exist and we are running on Colab, fall back to + pip-installing the backend dependencies into the current environment + so that imports like structlog and fastapi succeed. """ venv_lib = Path.home() / ".unsloth" / "studio" / "unsloth_studio" / "lib" if not venv_lib.exists(): + if _is_colab(): + _pip_install_backend_deps() + return import warnings warnings.warn( diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 0ac54d3866..42bae42819 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -1746,8 +1746,9 @@ if (-not $NeedLlamaSourceBuild) { # ============================================ Write-Host "" $doneLine = if ($env:SKIP_STUDIO_BASE -eq "1") { "Setup Complete!" } else { "Update Complete!" } +$doneContent = " $doneLine" Write-Host "+===============================================+" -ForegroundColor Green -Write-Host "| $doneLine |" -ForegroundColor Green +Write-Host ("|" + $doneContent.PadRight(47) + "|") -ForegroundColor Green Write-Host "| |" -ForegroundColor Green Write-Host "| Launch with: |" -ForegroundColor Green Write-Host "| unsloth studio -H 0.0.0.0 -p 8888 |" -ForegroundColor Green diff --git a/studio/setup.sh b/studio/setup.sh index a7991b83be..dd758e49e8 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -292,15 +292,23 @@ VENV_T5_DIR="$STUDIO_HOME/.venv_t5" [ -d "$REPO_ROOT/.venv_t5" ] && rm -rf "$REPO_ROOT/.venv_t5" # Note: do NOT delete $STUDIO_HOME/.venv here — install.sh handles migration +_COLAB_NO_VENV=false if [ ! -x "$VENV_DIR/bin/python" ]; then - echo "❌ ERROR: Virtual environment not found at $VENV_DIR" - echo " Run install.sh first to create the environment:" - echo " curl -fsSL https://unsloth.ai/install.sh | sh" - exit 1 + if [ "$IS_COLAB" = true ]; then + # On Colab there is no Studio venv -- install backend deps into system Python + echo " Colab detected, installing Studio backend dependencies..." + pip install -q -r "$SCRIPT_DIR/backend/requirements/studio.txt" 2>/dev/null || true + _COLAB_NO_VENV=true + else + echo "❌ ERROR: Virtual environment not found at $VENV_DIR" + echo " Run install.sh first to create the environment:" + echo " curl -fsSL https://unsloth.ai/install.sh | sh" + exit 1 + fi +else + source "$VENV_DIR/bin/activate" fi -source "$VENV_DIR/bin/activate" - install_python_stack() { python "$SCRIPT_DIR/install_python_stack.py" } @@ -324,6 +332,24 @@ fast_install() { cd "$SCRIPT_DIR" +# On Colab without a venv, skip all venv-dependent sections (Python deps +# update, llama.cpp build) -- the backend deps were already installed above. +if [ "$_COLAB_NO_VENV" = true ]; then + echo "✅ Studio backend dependencies installed into system Python" + + echo "" + echo "╔══════════════════════════════════════╗" + echo "║ Setup Complete! ║" + echo "╠══════════════════════════════════════╣" + echo "║ Unsloth Studio is ready to start ║" + echo "║ in your Colab notebook! ║" + echo "║ ║" + echo "║ from colab import start ║" + echo "║ start() ║" + echo "╚══════════════════════════════════════╝" + exit 0 +fi + # ── Check if Python deps need updating ── # Compare installed package version against PyPI latest. # Skip all Python dependency work if versions match (fast update path).