Docker Studio: keep unsloth-studio-update --ref from leaving Studio dead

[oobabooga]

unsloth-studio-update --ref main, the command the helper's own --help shows, kills Studio in the container it was asked to update. Studio never comes back, not even after docker restart, and running the helper again says it cannot help. This makes the command work, and makes every failure put the previous install back instead of leaving a dead container.

Why it breaks

--ref runs pip install git+https://github.com/unslothai/unsloth@main, which replaces the editable install at /opt/unsloth-studio/src with a build in site-packages. Only release wheels carry the built frontend, so unsloth studio exits 1 on start:

[ERROR] Unsloth frontend build not found.
Tried:
  - /opt/unsloth-studio/unsloth_studio/lib/python3.12/site-packages/studio/frontend/dist

The helper only checked import studio.backend.main, which still succeeds, so it restarted Studio into that tree. supervisord retries three times and parks the program in FATAL.

Two things then make it worse:

• Running the helper again says supervisor not managing 'studio' here; restart Studio yourself. supervisorctl status exits 3 for a program that exists but is not running, and the script read that as "not managed".
• The venv on disk is already replaced by then, so any update that cannot start leaves the container one restart away from a dead Studio.

The fix

• --ref now installs the same shape the image was built with instead of a wheel: the ref is fetched into a staging tree next to src, its frontend is built there with the bundled Node (npm ci, falling back to npm install, then npm run build), the oxc-validator runtime is installed, and only then is the tree swapped in and installed editable. Nothing is touched until the build succeeds.
• The swap follows src to where it really is (readlink -f), so an image that presents it as a link into its own copy of Studio gets the update in that copy rather than a directory written over the link.
• Before restarting, every path now checks what Studio actually needs: import studio.backend.main and a studio/frontend/dist/index.html next to the imported package. The checks run from /, so a studio directory in the caller's working directory cannot answer for the installed one.
• A failed check restores the previous install. The helper records how unsloth and unsloth_zoo are installed before it starts (from direct_url.json: an editable path, a git commit, or a released version), puts the old source tree back, reinstalls exactly those, says so, and exits non-zero without restarting.
• supervisorctl status exit 3 now counts as managed, so a Studio left FATAL by an earlier update gets restarted instead of a "not managing" message.
• After a restart the helper waits for /api/health and says whether Studio came back. UNSLOTH_STUDIO_UPDATE_HEALTH_WAIT bounds the wait (default 180 seconds, 0 skips it).

Before and after

Both sides are unsloth/unsloth:latest; B carries this branch's helper. Same container recipe, no GPU.

Check	A (published)	B (this PR)
--ref main	exit 7 after 49 s, Studio FATAL, /api/health unreachable	exit 0 after 80 s, Studio RUNNING, health 200
Studio after --ref main, Playwright	sign-in never loads, 0/1	5/5: sign in, pick a GGUF model, two chat turns
Code actually running after --ref main	site-packages copy without a frontend	/opt/unsloth-studio/src, inference.py md5 5ff793a5 to 0ad553ae (main)
studio/frontend/dist and oxc-validator/node_modules after the update	dist absent in the installed tree	both present
Plain update with Studio already FATAL	"supervisor not managing 'studio' here", stays FATAL	restarts it, health 200
Update that cannot start (--packages unsloth==2025.1.1)	exit 1, venv left replaced; the next restart lands in FATAL	exit 1, previous install restored, still RUNNING before and after a restart
--ref no-such-ref-xyz	exit 1	exit 1, Studio untouched and RUNNING

The --ref main run costs about 80 seconds, most of it the frontend build.

Verification

• tests/python/test_docker_update_helpers.py: 25 passed. The Studio helper's cases run the real script against a fake studio package, with pip, supervisorctl, git and the bundled npm as recording stubs, so the import check, the dist check and the rollback are exercised rather than asserted from source text. The linked-src case is covered too: the update must land in the tree the link points at, with no staging directory left in the Studio home.
• Six of those fail on main: the dist check, the rollback, the FATAL restart, the two --ref build cases and the rollback message. Run them from outside the repository: from the repo root, main's import studio.backend.main picks up the working tree's studio/ and fails for an unrelated reason.
• The llama.cpp updater cases in the same file are unchanged and still pass.

Not in this PR

This is the in-container helper only. unsloth studio update and the installer are unchanged.

--ref needs the bundled Node and network access for the frontend build. Without $UNSLOTH_STUDIO_HOME/node/bin/npm it refuses, rather than installing a tree Studio cannot serve.

A release-wheel update still installs into site-packages, where the wheel carries frontend/dist but no oxc-validator/node_modules, so code validation in Data Recipes degrades there as it does for any wheel install. Pre-existing, and separate from this failure.

The update still lives in the container's writable layer, so it survives docker restart and not docker rm. Pulling a new image is still the lasting path.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: e11d8c097e

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Propagate a failed Studio restart

When supervisorctl restart studio returns nonzero, this line discards the failure; the subsequent unreachable health check is only a warning, and with the documented UNSLOTH_STUDIO_UPDATE_HEALTH_WAIT=0 it is skipped entirely, so the updater prints done and exits successfully. Because the previous source tree has already been deleted, a ref that imports successfully but fails during actual startup can still leave Studio in FATAL with no rollback while callers believe the update succeeded; retain the rollback tree and return nonzero or restore it when the restart cannot be confirmed.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Stop when the frontend build command fails

When npm run build creates dist/index.html but subsequently exits nonzero—for example, a requested ref adds a failing post-build step—the update still succeeds. The surrounding subshell is the left side of || return 1, so Bash disables errexit for commands inside it; the following if can return zero, and the later file-existence check accepts the partial build. Check the build command's status explicitly before continuing so a failed ref build is never swapped into the live source tree.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Restore dependencies after a --with-deps failure

When an update is run with --with-deps, pip may upgrade any dependency before either pip itself or studio_tree_ok reports failure, but the rollback manifest records only unsloth and unsloth_zoo and this reinstall explicitly uses --no-deps. Reinstalling those two packages therefore cannot restore the previous environment; an upgraded incompatible dependency can keep the old Studio unable to import or start even though the failure path claims the previous install is back. Snapshot and restore the affected dependency versions as well when dependency updates are enabled.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Restore the swapped tree when the updater is interrupted

If the updater receives a trappable signal such as TERM or HUP after the source swap but while pip is still installing the Zoo ref, the EXIT trap runs only cleanup: it deletes the rollback manifest without calling restore, leaving the unvalidated new tree active and the previous tree stranded at .src-prev.*. A container shutdown or command timeout in this window can therefore leave lazy imports or the next Studio restart using a mismatched source/Zoo combination; make the exit trap restore while SWAPPED=1 and disarm that rollback only after the update is committed.

Useful? React with 👍 / 👎.

[danielhanchen]

Confirmed this hits docker/unsloth_studio_update.sh, where the post-update check only imports studio.backend.main and the supervisorctl exit 3 path reports Studio as unmanaged, so --ref can leave it FATAL. Will get this reviewed.