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.
This commit is contained in:
Daniel Han 2026-03-26 14:02:36 +00:00
commit 0e92ede59e

View file

@ -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"