From 556f396b3cd63da525e2cdac1a5f1c606f4494ce Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 27 May 2026 01:35:13 -0700
Subject: [PATCH] ci: install unsloth_zoo from git main in notebooks-ci +
 studio-backend-ci (#5802)

* ci: install unsloth_zoo from git main in notebooks-ci + studio-backend-ci

These were the only two workflows that still pulled unsloth_zoo from
PyPI; every other CI (Core, MLX, version-compat, install.sh-driven
Studio smokes) installs zoo from git main. Drift between PyPI and
main hides fixes-on-zoo-main and lets PR-time validation pass on a
stale zoo, then break for users on next release.

Both edits match the retry-with-backoff shape mlx-ci.yml already uses.

* ci: drop --no-deps from studio-backend-ci unsloth_zoo install

The prior PyPI line was `pip install 'unsloth_zoo>=2026.5.1'` (no
--no-deps), which pulled in triton and the rest of zoo's runtime deps.
I dropped that transitive resolve in the first commit, which broke
collection of 5 tests in Repo tests (CPU) with
ModuleNotFoundError: No module named 'triton'.

Match the prior dep-resolve shape, keeping the source-from-git change.
notebooks-ci keeps --no-deps because its original line also had it.
---
 .github/workflows/notebooks-ci.yml      | 10 +++++++++-
 .github/workflows/studio-backend-ci.yml | 16 +++++++++++++---
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/notebooks-ci.yml b/.github/workflows/notebooks-ci.yml
index 673b2f3cc5..2edcae8ab2 100644
--- a/.github/workflows/notebooks-ci.yml
+++ b/.github/workflows/notebooks-ci.yml
@@ -285,7 +285,15 @@ jobs:
           # The PR-time CI must validate the code in this PR; PyPI unsloth
           # may lag the in-repo CPU-torch fallback in unsloth/kernels/utils.py
           # (lines 162-170) that handles missing torch._C._cuda_getCurrentRawStream.
-          pip install --no-deps unsloth_zoo
+          # unsloth_zoo from git main mirrors every other CI (Core / MLX /
+          # install.sh) so PR-time validation sees the same zoo HEAD.
+          for attempt in 1 2 3; do
+            if pip install --no-deps "unsloth_zoo @ git+https://github.com/unslothai/unsloth-zoo"; then
+              break
+            fi
+            [ "$attempt" -eq 3 ] && { echo "::error::unsloth_zoo install failed after 3 attempts"; exit 1; }
+            sleep $((5 * attempt))
+          done
           pip install --no-deps -e ./unsloth
 
       - name: Convert notebooks for AST scan
diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml
index 63eb70f7f1..ee5bbe8633 100644
--- a/.github/workflows/studio-backend-ci.yml
+++ b/.github/workflows/studio-backend-ci.yml
@@ -144,9 +144,19 @@ jobs:
           # versions ship a CPU build that imports cleanly on Linux.
           pip install 'bitsandbytes>=0.45'
           # unsloth.device_type imports unsloth_zoo.utils.Version at module
-          # scope, so the conftest preload needs unsloth_zoo even though
-          # it is an optional dep of unsloth.
-          pip install 'unsloth_zoo>=2026.5.1'
+          # scope, so the conftest preload needs unsloth_zoo. Pull from
+          # git main so this job sees the same zoo HEAD as Core / MLX /
+          # install.sh do (otherwise a fix on zoo main hides until release).
+          # No --no-deps: matches prior `pip install 'unsloth_zoo>=2026.5.1'`
+          # behaviour so triton etc. still come in for the Repo tests CPU
+          # collection imports.
+          for attempt in 1 2 3; do
+            if pip install "unsloth_zoo @ git+https://github.com/unslothai/unsloth-zoo"; then
+              break
+            fi
+            [ "$attempt" -eq 3 ] && { echo "::error::unsloth_zoo install failed after 3 attempts"; exit 1; }
+            sleep $((5 * attempt))
+          done
           pip install -e . --no-deps
 
       - name: Repo tests (CPU, auto-discovered)