From 0e92ede59ecd6177bde44970e73f222397c3fb49 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 26 Mar 2026 14:02:36 +0000
Subject: [PATCH 1/2] Fix Colab training failing with "Please install
 unsloth_zoo"

The Colab no-venv path in setup.sh (added in #4601, modified in #4618)
only installs studio.txt (fastapi, uvicorn, structlog, etc.) into
system Python. It skips install_python_stack entirely via
_SKIP_PYTHON_DEPS=true, which means base.txt (unsloth + unsloth-zoo)
and extras.txt (trl, transformers, etc.) are never installed.

When the training worker subprocess spawns and tries to import unsloth,
unsloth/__init__.py calls importlib.metadata.version("unsloth_zoo")
which raises PackageNotFoundError, surfaced as:
  "Failed to import ML libraries: Please install unsloth_zoo"

Fix: install base.txt and extras.txt in the Colab block, right after
studio.txt. extras.txt gets the same version-constraint stripping as
studio.txt to avoid conflicts with Colab's pre-installed packages.
base.txt has no version pins so it is installed as-is.

Only the Colab code path is changed. Non-Colab paths are untouched.
---
 studio/setup.sh | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/studio/setup.sh b/studio/setup.sh
index 4e46a12fe0..0dd93bf799 100755
--- a/studio/setup.sh
+++ b/studio/setup.sh
@@ -303,6 +303,17 @@ if [ ! -x "$VENV_DIR/bin/python" ]; then
         sed 's/[><=!~;].*//' "$SCRIPT_DIR/backend/requirements/studio.txt" \
             | grep -v '^#' | grep -v '^$' \
             | pip install -q -r /dev/stdin 2>/dev/null || true
+
+        # Install core ML packages (unsloth + unsloth-zoo) and training
+        # extras (trl, transformers, etc.) into system Python.  Without
+        # these the training/inference/export workers fail at import time
+        # with "Failed to import ML libraries: Please install unsloth_zoo".
+        echo "   Installing core ML packages for training..."
+        pip install -q -r "$SCRIPT_DIR/backend/requirements/base.txt" 2>/dev/null || true
+        sed 's/[><=!~;].*//' "$SCRIPT_DIR/backend/requirements/extras.txt" \
+            | grep -v '^#' | grep -v '^$' \
+            | pip install -q -r /dev/stdin 2>/dev/null || true
+
         _COLAB_NO_VENV=true
     else
         echo "❌ ERROR: Virtual environment not found at $VENV_DIR"

From be4c57b052cb2288674c35f9e3831fe0587aac3f Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 26 Mar 2026 14:18:49 +0000
Subject: [PATCH 2/2] Fix Colab training failing with "Please install
 unsloth_zoo"

The Colab no-venv path in setup.sh (added in #4601, modified in #4618)
only installs studio.txt (fastapi, uvicorn, structlog, etc.) into
system Python. It skips install_python_stack entirely via
_SKIP_PYTHON_DEPS=true, which means base.txt (unsloth + unsloth-zoo),
extras.txt, extras-no-deps.txt (trl, transformers, sentence_transformers),
and overrides.txt (torchao) are never installed.

When the training worker subprocess spawns and tries to import unsloth,
unsloth/__init__.py calls importlib.metadata.version("unsloth_zoo")
which raises PackageNotFoundError, surfaced as:
  "Failed to import ML libraries: Please install unsloth_zoo"

Even if unsloth_zoo were installed, the next failure would be trl
(imported at trainer.py:63) since it lives in extras-no-deps.txt.

Fix: install all four requirement groups in the Colab block, matching
what install_python_stack.py does for the normal venv path:
  - base.txt (unsloth + unsloth-zoo) -- fail-fast on error
  - extras.txt (version-stripped to avoid Colab conflicts)
  - extras-no-deps.txt (trl, transformers, sentence_transformers)
  - overrides.txt (torchao)

Only the Colab code path is changed. Non-Colab paths are untouched.
---
 studio/setup.sh | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/studio/setup.sh b/studio/setup.sh
index 0dd93bf799..b89e446a67 100755
--- a/studio/setup.sh
+++ b/studio/setup.sh
@@ -304,15 +304,24 @@ if [ ! -x "$VENV_DIR/bin/python" ]; then
             | grep -v '^#' | grep -v '^$' \
             | pip install -q -r /dev/stdin 2>/dev/null || true
 
-        # Install core ML packages (unsloth + unsloth-zoo) and training
-        # extras (trl, transformers, etc.) into system Python.  Without
-        # these the training/inference/export workers fail at import time
-        # with "Failed to import ML libraries: Please install unsloth_zoo".
+        # Install the same ML package groups that install_python_stack.py
+        # uses for the normal (venv) path.  Without these the training,
+        # inference, and export workers fail at import time with:
+        #   "Failed to import ML libraries: Please install unsloth_zoo"
+        # base.txt must succeed -- it has unsloth + unsloth-zoo.
         echo "   Installing core ML packages for training..."
-        pip install -q -r "$SCRIPT_DIR/backend/requirements/base.txt" 2>/dev/null || true
+        pip install -q -r "$SCRIPT_DIR/backend/requirements/base.txt" || {
+            echo "Failed to install base ML packages (unsloth + unsloth-zoo)"
+            exit 1
+        }
         sed 's/[><=!~;].*//' "$SCRIPT_DIR/backend/requirements/extras.txt" \
             | grep -v '^#' | grep -v '^$' \
             | pip install -q -r /dev/stdin 2>/dev/null || true
+        # trl, transformers, sentence_transformers live in extras-no-deps.txt
+        pip install -q --no-deps \
+            -r "$SCRIPT_DIR/backend/requirements/extras-no-deps.txt" 2>/dev/null || true
+        pip install -q --force-reinstall \
+            -r "$SCRIPT_DIR/backend/requirements/overrides.txt" 2>/dev/null || true
 
         _COLAB_NO_VENV=true
     else