From 6fe3e7655abb9cfc1edb5a5b5d06198dfa01019e Mon Sep 17 00:00:00 2001
From: Roland Tannous <rolandtannous@gravityq.ai>
Date: Thu, 26 Mar 2026 13:45:16 +0000
Subject: [PATCH] Fix setup.sh --local on Colab not installing unsloth_zoo

Two bugs prevented unsloth_zoo from being installed when running
setup.sh --local on Google Colab:

1. setup.sh had no argument parsing for --local. The flag was only
   handled by install.sh, which sets STUDIO_LOCAL_INSTALL and
   STUDIO_LOCAL_REPO as env vars before calling setup.sh. When the
   Colab notebook calls setup.sh --local directly, the flag was
   silently ignored.

2. The Colab no-venv path set _SKIP_PYTHON_DEPS=true, which skipped
   install_python_stack entirely. That function is the only place
   that installs unsloth + unsloth_zoo and applies the local editable
   overlay (install_python_stack.py lines 402-423).

Fix: add --local parsing to setup.sh, and re-enable install_python_stack
on Colab when --local is passed (the PyPI version check that uses
$VENV_DIR/bin/python is still correctly skipped).
---
 studio/setup.sh | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/studio/setup.sh b/studio/setup.sh
index 4e46a12fe0..b3f66ca4a1 100755
--- a/studio/setup.sh
+++ b/studio/setup.sh
@@ -7,6 +7,16 @@ set -euo pipefail
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 REPO_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
 
+# ── Parse arguments ──
+for arg in "$@"; do
+    case "$arg" in
+        --local)
+            export STUDIO_LOCAL_INSTALL=1
+            export STUDIO_LOCAL_REPO="$REPO_ROOT"
+            ;;
+    esac
+done
+
 # ── Helper: run command quietly, show output only on failure ──
 _run_quiet() {
     local on_fail=$1
@@ -373,6 +383,13 @@ print(version('$_PKG_NAME'))
     fi
 fi
 
+# On Colab with --local, still call install_python_stack to install
+# unsloth + unsloth_zoo and overlay the local repo — only the version check
+# above needed skipping (it requires $VENV_DIR/bin/python which doesn't exist).
+if [ "$_COLAB_NO_VENV" = true ] && [ "${STUDIO_LOCAL_INSTALL:-0}" = "1" ]; then
+    _SKIP_PYTHON_DEPS=false
+fi
+
 if [ "$_SKIP_PYTHON_DEPS" = false ]; then
     install_python_stack