diff --git a/.github/workflows/docker-publish.yml b/.github/workflows/docker-publish.yml new file mode 100644 index 0000000000..d91f61a14b --- /dev/null +++ b/.github/workflows/docker-publish.yml @@ -0,0 +1,583 @@ +# Builds and publishes the Blackwell-compatible Unsloth Docker image. +# +# Runs on free GPU-less GitHub Ubuntu runners: cu128 wheels are fat binaries +# (sm_70..sm_120 amd64, sm_80;90;100;120 aarch64), the Dockerfile pins explicit +# wheel URLs, the build-time check uses torch._C._cuda_getArchFlags() (no CUDA +# device needed), and UNSLOTH_COMPILE_DISABLE=1 blocks GPU-keyed JIT. +# +# Multi-arch: amd64 + arm64 build in parallel on native runners (ubuntu-latest + +# ubuntu-24.04-arm), then merge per-arch digests into one manifest. Native arm64 +# is ~3x faster and less flaky than QEMU; DGX Spark / Grace pull the arm64 child. +# +# Required secrets: DOCKERHUB_USERNAME, DOCKERHUB_TOKEN +# Optional variable HAS_GPU_RUNNER='true' gates the smoke-test job. + +name: Publish Blackwell Docker image + +on: + push: + branches: [main] + tags: ['v*'] + schedule: + - cron: '17 4 * * 1' # weekly Mon 04:17 UTC (off-the-hour on purpose) + workflow_dispatch: + inputs: + unsloth_ref: + # Blank means "the dispatched branch" (resolver falls back to sha, then + # main). The stable-tag gates require this EMPTY, so a non-blank default + # would make every UI-default dispatch publish SHA tags only. + description: 'unsloth git ref override (blank = dispatched branch + stable tags)' + required: false + default: '' + unsloth_zoo_ref: + description: 'unsloth-zoo git ref to bake in' + required: false + default: 'main' + llama_prebuilt_tag: + description: 'unslothai/llama.cpp prebuilt release tag to bake (blank = newest)' + required: false + default: '' + notebooks_ref: + description: 'unslothai/notebooks git ref to bake (resolved to one commit)' + required: false + default: 'main' + +env: + REGISTRY: docker.io + IMAGE_NAME: unsloth/unsloth + +# Serialise per-ref runs so two pushes don't both retag :latest from different +# commits. Don't cancel in-progress -- the build is expensive and a half-built +# image is worse than a briefly stale :latest. +concurrency: + group: docker-publish-${{ github.ref }} + cancel-in-progress: false + +# Least-privilege default for GITHUB_TOKEN. Pushes use Docker Hub registry creds, +# not GITHUB_TOKEN, so read is enough; jobs needing more declare packages: write. +permissions: + contents: read + +jobs: + # Resolve every upstream ref ONCE (llama tag + unsloth/zoo shas + notebooks + # commit) so both arch legs and Studio bake identical bits. A dispatch input + # pins a frozen value; else a branch/tag is frozen to a sha via ls-remote, and + # llama "latest" follows the /releases/latest redirect (mirrors build.sh). + prepare: + runs-on: ubuntu-latest + timeout-minutes: 5 + permissions: + contents: read + outputs: + llama_tag: ${{ steps.llama.outputs.tag }} + # Resolved once, shared by every consumer -- see the job header. + unsloth_ref: ${{ steps.unsloth_ref.outputs.ref }} + zoo_ref: ${{ steps.zoo_ref.outputs.ref }} + notebooks_commit: ${{ steps.notebooks.outputs.commit }} + steps: + - name: Resolve llama.cpp prebuilt tag + id: llama + env: + INPUT_TAG: ${{ github.event.inputs.llama_prebuilt_tag }} + run: | + TAG="$INPUT_TAG" + if [ -z "$TAG" ]; then + # Same rule as the three ref resolvers below. This step has no + # explicit `shell:`, so it runs under `bash -e` WITHOUT pipefail and + # a failing curl inside `curl | sed` is lost: the step exited 0 and + # published tag=latest. Every consumer resolves that MUTABLE tag + # again -- fetch_llama_prebuilt.py once per arch leg, Dockerfile. + # studio once more -- so a release cut mid-run can put different + # llama.cpp bundles under one manifest. Fail the job instead. + if ! REDIRECT="$(curl -fsSL -o /dev/null -w '%{url_effective}' \ + https://github.com/unslothai/llama.cpp/releases/latest)"; then + echo "::error::unslothai/llama.cpp unreachable; cannot resolve the newest prebuilt tag" + exit 1 + fi + TAG="$(printf '%s\n' "$REDIRECT" | sed -n 's#.*/releases/tag/##p')" + if [ -z "$TAG" ]; then + echo "::error::/releases/latest did not redirect to a release tag (landed on ${REDIRECT})" + exit 1 + fi + fi + echo "tag=${TAG}" >> "$GITHUB_OUTPUT" + echo "llama.cpp prebuilt tag: ${TAG}" + + # Requested-ref precedence: dispatch input, else pushed tag, else trigger + # sha, else main -- then frozen to one sha per the job header. + - name: Resolve unsloth ref + id: unsloth_ref + env: + INPUT_REF: ${{ github.event.inputs.unsloth_ref }} + TAG_REF: ${{ startsWith(github.ref, 'refs/tags/') && github.ref_name || '' }} + PUSH_SHA: ${{ github.sha }} + run: | + REF="$INPUT_REF" + [ -n "$REF" ] || REF="$TAG_REF" + [ -n "$REF" ] || REF="$PUSH_SHA" + REF="${REF:-main}" + if printf '%s' "$REF" | grep -Eq '^[0-9a-f]{40}$'; then + SHA="$REF" + else + # ls-remote exits 0 whether or not a ref matched, so a non-zero exit + # means we never reached the remote. The pipe into awk would hide it + # (no pipefail under the default `bash -e` shell) and the fallback + # below would then hand a MUTABLE name to the amd64, arm64 and Studio + # builds, which each resolve it again -- the exact split this job + # exists to prevent. Fail the run instead. + if ! LS_OUT="$(git ls-remote https://github.com/unslothai/unsloth "$REF")"; then + echo "::error::unslothai/unsloth unreachable; cannot freeze ref '${REF}' to a sha" + exit 1 + fi + SHA="$(printf '%s\n' "$LS_OUT" | awk 'NR==1{print $1}')" + [ -n "$SHA" ] || SHA="$REF" + fi + echo "ref=${SHA}" >> "$GITHUB_OUTPUT" + echo "unsloth ref: ${SHA}" + + # Mirror the unsloth tag into the zoo ONLY when that tag exists there: + # unsloth's v* tags are Studio releases the zoo never cuts, so blindly + # mirroring github.ref_name made every tag publish fail at zoo install. + - name: Resolve unsloth-zoo ref + id: zoo_ref + run: | + REF="${{ github.event.inputs.unsloth_zoo_ref }}" + if [ -z "$REF" ] && [ "${{ startsWith(github.ref, 'refs/tags/') }}" = "true" ]; then + if git ls-remote --exit-code --tags https://github.com/unslothai/unsloth-zoo \ + "refs/tags/${{ github.ref_name }}" >/dev/null 2>&1; then + REF="${{ github.ref_name }}" + fi + fi + REF="${REF:-main}" + # Freeze to one sha per the job header; a 40-char sha already is one. + if printf '%s' "$REF" | grep -Eq '^[0-9a-f]{40}$'; then + SHA="$REF" + else + # Same rule as the unsloth ref above: a non-zero ls-remote is a + # transport failure, not "no such ref", and forwarding the branch + # name would let the three builds each pick a different commit. + if ! LS_OUT="$(git ls-remote https://github.com/unslothai/unsloth-zoo "$REF")"; then + echo "::error::unslothai/unsloth-zoo unreachable; cannot freeze ref '${REF}' to a sha" + exit 1 + fi + SHA="$(printf '%s\n' "$LS_OUT" | awk 'NR==1{print $1}')" + [ -n "$SHA" ] || SHA="$REF" + fi + echo "ref=${SHA}" >> "$GITHUB_OUTPUT" + echo "unsloth-zoo ref: ${SHA}" + + # Freeze notebooks to ONE commit per the job header, so baked templates + + # .unsloth_template_commit are identical across legs and reruns. + - name: Resolve unsloth/notebooks commit + id: notebooks + env: + INPUT_REF: ${{ github.event.inputs.notebooks_ref }} + run: | + REF="${INPUT_REF:-main}" + if printf '%s' "$REF" | grep -Eq '^[0-9a-f]{40}$'; then + SHA="$REF" + else + # Same rule as the two refs above: only a reachable remote with no + # matching ref may fall through to the literal "$REF". + if ! LS_OUT="$(git ls-remote https://github.com/unslothai/notebooks "$REF")"; then + echo "::error::unslothai/notebooks unreachable; cannot freeze ref '${REF}' to a sha" + exit 1 + fi + SHA="$(printf '%s\n' "$LS_OUT" | awk 'NR==1{print $1}')" + [ -n "$SHA" ] || SHA="$REF" + fi + echo "commit=${SHA}" >> "$GITHUB_OUTPUT" + echo "notebooks commit: ${SHA}" + + # Per-arch build: two parallel jobs on native runners, each pushing a single-arch + # image by digest (no tag); the merge job stitches them into one manifest. Avoids + # the "last push wins" race of two jobs pushing the same tag. + build: + needs: prepare + strategy: + fail-fast: false + matrix: + include: + - platform: linux/amd64 + runner: ubuntu-latest + - platform: linux/arm64 + runner: ubuntu-24.04-arm + runs-on: ${{ matrix.runner }} + timeout-minutes: 90 + permissions: + contents: read + packages: write + + steps: + - uses: actions/checkout@v4 + + # Free up ~20GB so cu128 wheels + cudnn fit. Runner layouts differ (arm64 + # lacks /usr/share/dotnet), hence `|| true`. + - name: Reclaim disk + run: | + # None of these toolchains are used; paths differ across runners, hence `|| true`. + sudo rm -rf /usr/share/dotnet /usr/local/lib/android /opt/ghc \ + /opt/hostedtoolcache "$AGENT_TOOLSDIRECTORY" \ + /usr/local/.ghcup /usr/share/swift \ + /usr/local/share/powershell /usr/local/lib/node_modules \ + /usr/local/julia* /opt/microsoft /usr/share/miniconda \ + /opt/az /usr/local/share/boost /usr/local/share/chromium || true + sudo docker image prune -af >/dev/null 2>&1 || true + df -h / + + - uses: docker/setup-buildx-action@v3 + + - name: Log in to Docker Hub + uses: docker/login-action@v3 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_TOKEN }} + + # Labels/annotations for the FINAL manifest. No tags here -- each per-arch + # build pushes by digest only; tags are attached by the merge job. + - name: Resolve labels + id: meta + uses: docker/metadata-action@v5 + with: + images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }} + + - name: Build and push (per-arch by digest) + id: build + uses: docker/build-push-action@v6 + with: + context: ./docker + file: ./docker/Dockerfile + platforms: ${{ matrix.platform }} + labels: ${{ steps.meta.outputs.labels }} + # Per-arch build cache: the platform suffix keeps the two legs from colliding. + cache-from: type=gha,scope=build-${{ matrix.platform }} + cache-to: type=gha,scope=build-${{ matrix.platform }},mode=max + outputs: type=image,name=${{ env.REGISTRY }}/${{ env.IMAGE_NAME }},push-by-digest=true,name-canonical=true,push=true + # Keep prose OUT of build-args -- build-push-action forwards every + # non-empty line verbatim, so a #-line becomes a bogus --build-arg. All + # four values come from the prepare job (resolved once). + build-args: | + CUDA_VERSION=12.8.1 + UBUNTU_VERSION=24.04 + PYTHON_VERSION=3.12 + UNSLOTH_REF=${{ needs.prepare.outputs.unsloth_ref }} + UNSLOTH_ZOO_REF=${{ needs.prepare.outputs.zoo_ref }} + LLAMA_PREBUILT_TAG=${{ needs.prepare.outputs.llama_tag }} + UNSLOTH_NOTEBOOKS_REF=${{ needs.prepare.outputs.notebooks_commit }} + + # Stash the per-arch digest as an artifact for the merge job. `platform` + # has a slash, so substitute a dash for a unique filename. + - name: Export digest + run: | + mkdir -p /tmp/digests + digest='${{ steps.build.outputs.digest }}' + touch "/tmp/digests/${digest#sha256:}" + + - name: Upload digest + uses: actions/upload-artifact@v4 + with: + name: digests-core-${{ matrix.platform == 'linux/amd64' && 'amd64' || 'arm64' }} + path: /tmp/digests/* + if-no-files-found: error + retention-days: 1 + + # Merge the two per-arch digests into a multi-platform manifest under the real + # user-facing tag(s). Runs only after both build legs succeed. + merge: + runs-on: ubuntu-latest + needs: build + timeout-minutes: 15 + permissions: + contents: read + packages: write + outputs: + # Manifest digest of the just-published base image; build-studio FROMs this + # exact digest so Studio layers on THIS run's bits, not whatever `base` + # points at later. + digest: ${{ steps.manifest_digest.outputs.digest }} + steps: + - uses: actions/download-artifact@v4 + with: + path: /tmp/digests + pattern: digests-core-* + merge-multiple: true + + - uses: docker/setup-buildx-action@v3 + + - name: Log in to Docker Hub + uses: docker/login-action@v3 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_TOKEN }} + + - name: Resolve tags + id: meta + uses: docker/metadata-action@v5 + with: + images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }} + # The base image must NEVER claim :latest. metadata-action defaults to + # flavor latest=auto, which would tag :latest on a v* (semver) tag push + # and collide with the Studio image that legitimately owns :latest. + flavor: latest=false + tags: | + # The lean training image publishes under the core- prefix; the + # full Studio image (build-studio/merge-studio below) owns + # :latest, matching what the previous production image shipped. + # Only tag :core when the workflow ran on the default branch + # AND the operator did NOT override ANY baked input on dispatch + # (unsloth_ref, unsloth_zoo_ref, notebooks_ref, llama_prebuilt_tag; + # push/schedule leave inputs null == '', and the 'main' defaults + # are accepted explicitly). Without these conditions a maintainer + # testing a feature ref could overwrite :core with non-main bits. + type=raw,value=core,enable=${{ github.ref == format('refs/heads/{0}', github.event.repository.default_branch) && github.event.inputs.unsloth_ref == '' && (github.event.inputs.unsloth_zoo_ref == '' || github.event.inputs.unsloth_zoo_ref == 'main') && (github.event.inputs.notebooks_ref == '' || github.event.inputs.notebooks_ref == 'main') && github.event.inputs.llama_prebuilt_tag == '' }} + type=ref,event=tag,prefix=core- + type=schedule,pattern=core-nightly + type=sha,prefix=core-sha-,format=short + + - name: Create multi-arch manifest + working-directory: /tmp/digests + run: | + docker buildx imagetools create \ + $(jq -cr '.tags | map("-t " + .) | join(" ")' <<<"$DOCKER_METADATA_OUTPUT_JSON") \ + $(printf '${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}@sha256:%s ' *) + + - name: Inspect the result + run: | + for tag in $(jq -r '.tags[]' <<<"$DOCKER_METADATA_OUTPUT_JSON"); do + echo "=== $tag ===" + docker buildx imagetools inspect "$tag" + done + + - name: Export manifest digest + id: manifest_digest + run: | + TAG="$(jq -r '.tags[0]' <<<"$DOCKER_METADATA_OUTPUT_JSON")" + DIGEST="$(docker buildx imagetools inspect "$TAG" --format '{{json .Manifest.Digest}}' | tr -d '"')" + test -n "$DIGEST" + echo "digest=${DIGEST}" >> "$GITHUB_OUTPUT" + echo "base manifest: ${TAG} @ ${DIGEST}" + + # Full image: base + Unsloth Studio + JupyterLab + sshd (Dockerfile.studio). + # This is :latest. Same by-digest build + merge pattern as the base, FROMing the + # base manifest digest from the merge job. The arm64 leg builds Studio's vite + # frontend natively (the long pole), hence the larger timeout. + build-studio: + # `merge` for the freshly-published base manifest digest; `prepare` for the + # one resolved zoo ref (job outputs only flow through direct `needs`). + needs: [prepare, merge] + strategy: + fail-fast: false + matrix: + include: + - platform: linux/amd64 + runner: ubuntu-latest + - platform: linux/arm64 + runner: ubuntu-24.04-arm + runs-on: ${{ matrix.runner }} + timeout-minutes: 150 + permissions: + contents: read + packages: write + steps: + - uses: actions/checkout@v4 + + - name: Reclaim disk + run: | + # None of these toolchains are used; paths differ across runners, hence `|| true`. + sudo rm -rf /usr/share/dotnet /usr/local/lib/android /opt/ghc \ + /opt/hostedtoolcache "$AGENT_TOOLSDIRECTORY" \ + /usr/local/.ghcup /usr/share/swift \ + /usr/local/share/powershell /usr/local/lib/node_modules \ + /usr/local/julia* /opt/microsoft /usr/share/miniconda \ + /opt/az /usr/local/share/boost /usr/local/share/chromium || true + sudo docker image prune -af >/dev/null 2>&1 || true + df -h / + + - uses: docker/setup-buildx-action@v3 + + - name: Log in to Docker Hub + uses: docker/login-action@v3 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_TOKEN }} + + - name: Resolve labels + id: meta + uses: docker/metadata-action@v5 + with: + images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }} + + - name: Build and push (per-arch by digest) + id: build + uses: docker/build-push-action@v6 + with: + context: ./docker + file: ./docker/Dockerfile.studio + platforms: ${{ matrix.platform }} + labels: ${{ steps.meta.outputs.labels }} + # mode=min (final layers only): mode=max on this ~24GB image would blow + # the 10GB GHA cache quota and evict the base build's cache for no gain. + cache-from: type=gha,scope=studio-${{ matrix.platform }} + cache-to: type=gha,scope=studio-${{ matrix.platform }},mode=min + outputs: type=image,name=${{ env.REGISTRY }}/${{ env.IMAGE_NAME }},push-by-digest=true,name-canonical=true,push=true + # All three pins are the SAME values the base build baked (prepare job), + # so Studio, its zoo overlay and its llama.cpp match the base even if + # upstream moved mid-run. (build-args must be KEY=VALUE only.) + build-args: | + BASE_IMAGE=${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}@${{ needs.merge.outputs.digest }} + UNSLOTH_STUDIO_REF=${{ needs.prepare.outputs.unsloth_ref }} + UNSLOTH_STUDIO_ZOO_REF=${{ needs.prepare.outputs.zoo_ref }} + LLAMA_PREBUILT_TAG=${{ needs.prepare.outputs.llama_tag }} + + - name: Export digest + run: | + mkdir -p /tmp/digests + digest='${{ steps.build.outputs.digest }}' + touch "/tmp/digests/${digest#sha256:}" + + - name: Upload digest + uses: actions/upload-artifact@v4 + with: + name: digests-studio-${{ matrix.platform == 'linux/amd64' && 'amd64' || 'arm64' }} + path: /tmp/digests/* + if-no-files-found: error + retention-days: 1 + + merge-studio: + runs-on: ubuntu-latest + needs: build-studio + timeout-minutes: 15 + permissions: + contents: read + packages: write + steps: + - uses: actions/download-artifact@v4 + with: + path: /tmp/digests + pattern: digests-studio-* + merge-multiple: true + + - uses: docker/setup-buildx-action@v3 + + - name: Log in to Docker Hub + uses: docker/login-action@v3 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_TOKEN }} + + - name: Resolve tags + id: meta + uses: docker/metadata-action@v5 + with: + images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }} + # latest=false disables metadata-action's implicit latest=auto, which + # would otherwise emit :latest on a v* tag push and bypass the + # default-branch-only gate below. :latest is published only by the + # explicit type=raw rule (default-branch pushes), matching the base job. + flavor: latest=false + tags: | + # The full Studio image owns the unprefixed namespace, headed by + # :latest plus a stable :studio alias (default branch only). Tag + # pushes publish the version tag. Same gating rationale as the core job. + type=raw,value=latest,enable=${{ github.ref == format('refs/heads/{0}', github.event.repository.default_branch) && github.event.inputs.unsloth_ref == '' && (github.event.inputs.unsloth_zoo_ref == '' || github.event.inputs.unsloth_zoo_ref == 'main') && (github.event.inputs.notebooks_ref == '' || github.event.inputs.notebooks_ref == 'main') && github.event.inputs.llama_prebuilt_tag == '' }} + type=raw,value=studio,enable=${{ github.ref == format('refs/heads/{0}', github.event.repository.default_branch) && github.event.inputs.unsloth_ref == '' && (github.event.inputs.unsloth_zoo_ref == '' || github.event.inputs.unsloth_zoo_ref == 'main') && (github.event.inputs.notebooks_ref == '' || github.event.inputs.notebooks_ref == 'main') && github.event.inputs.llama_prebuilt_tag == '' }} + type=ref,event=tag + type=schedule,pattern=nightly + type=sha,prefix=sha-,format=short + + - name: Create multi-arch manifest + working-directory: /tmp/digests + run: | + docker buildx imagetools create \ + $(jq -cr '.tags | map("-t " + .) | join(" ")' <<<"$DOCKER_METADATA_OUTPUT_JSON") \ + $(printf '${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}@sha256:%s ' *) + + - name: Inspect the result + run: | + for tag in $(jq -r '.tags[]' <<<"$DOCKER_METADATA_OUTPUT_JSON"); do + echo "=== $tag ===" + docker buildx imagetools inspect "$tag" + done + + # Optional: pull the freshly published image onto a self-hosted GPU runner and + # run smoke_test.py. Skipped when no GPU runner is registered. + smoke-test: + needs: [merge, merge-studio] + if: ${{ vars.HAS_GPU_RUNNER == 'true' }} + runs-on: [self-hosted, gpu] + timeout-minutes: 30 + steps: + - uses: actions/checkout@v4 + + # Re-compute the tag list from the same metadata-action config the merge job + # used, so a run pulls the image it just published. IMPORTANT: keep the + # `enable=` expressions byte-identical to the merge jobs' gates above, else + # smoke could pull a previously-published :latest instead of the merged image. + - name: Resolve published base tag + id: meta_base + uses: docker/metadata-action@v5 + with: + images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }} + # Keep the base image off :latest here too (this recomputes the same + # tag list the merge step pushed, so the smoke test pulls the right ref). + flavor: latest=false + tags: | + type=raw,value=core,enable=${{ github.ref == format('refs/heads/{0}', github.event.repository.default_branch) && github.event.inputs.unsloth_ref == '' && (github.event.inputs.unsloth_zoo_ref == '' || github.event.inputs.unsloth_zoo_ref == 'main') && (github.event.inputs.notebooks_ref == '' || github.event.inputs.notebooks_ref == 'main') && github.event.inputs.llama_prebuilt_tag == '' }} + type=ref,event=tag,prefix=core- + type=schedule,pattern=core-nightly + type=sha,prefix=core-sha-,format=short + + - name: Pull and smoke-test the base image + run: | + # Use the first tag from the metadata output -- that is the image we + # just published. Falls back to :core only when the metadata is + # empty (defensive; should not happen on default-branch runs). + TAG="$(jq -r '.tags[0] // ""' <<<"$STEPS_META_BASE_JSON")" + if [ -z "$TAG" ]; then + TAG="${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:core" + fi + echo "smoke-testing $TAG" + docker pull "$TAG" + docker run --rm --gpus all "$TAG" python /workspace/smoke_test.py + env: + STEPS_META_BASE_JSON: ${{ steps.meta_base.outputs.json }} + + - name: Resolve published studio tag + id: meta_studio + uses: docker/metadata-action@v5 + with: + images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }} + # Mirror the studio tag rules (incl. latest=false) so the smoke test + # pulls the tag just published, not an implicit latest=auto :latest. + flavor: latest=false + tags: | + type=raw,value=latest,enable=${{ github.ref == format('refs/heads/{0}', github.event.repository.default_branch) && github.event.inputs.unsloth_ref == '' && (github.event.inputs.unsloth_zoo_ref == '' || github.event.inputs.unsloth_zoo_ref == 'main') && (github.event.inputs.notebooks_ref == '' || github.event.inputs.notebooks_ref == 'main') && github.event.inputs.llama_prebuilt_tag == '' }} + type=raw,value=studio,enable=${{ github.ref == format('refs/heads/{0}', github.event.repository.default_branch) && github.event.inputs.unsloth_ref == '' && (github.event.inputs.unsloth_zoo_ref == '' || github.event.inputs.unsloth_zoo_ref == 'main') && (github.event.inputs.notebooks_ref == '' || github.event.inputs.notebooks_ref == 'main') && github.event.inputs.llama_prebuilt_tag == '' }} + type=ref,event=tag + type=schedule,pattern=nightly + type=sha,prefix=sha-,format=short + + - name: Boot the full image and probe Studio + Jupyter + run: | + TAG="$(jq -r '.tags[0] // ""' <<<"$STEPS_META_STUDIO_JSON")" + if [ -z "$TAG" ]; then + TAG="${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:latest" + fi + echo "booting $TAG" + docker pull "$TAG" + CID="$(docker run -d --gpus all -p 18000:8000 -p 18888:8888 "$TAG")" + trap 'docker logs --tail 100 "$CID"; docker rm -f "$CID"' EXIT + ok_studio=0; ok_jupyter=0 + for i in $(seq 1 60); do + if curl -fsS http://localhost:18000/api/health >/dev/null 2>&1; then ok_studio=1; fi + # Probe /login, not /api: the launcher sets a password hash so /api + # returns 403; /login is unauthenticated and 200s once up. + if curl -fsS http://localhost:18888/login >/dev/null 2>&1; then ok_jupyter=1; fi + [ "$ok_studio" = 1 ] && [ "$ok_jupyter" = 1 ] && break + sleep 5 + done + [ "$ok_studio" = 1 ] || { echo "Studio /api/health never went healthy"; exit 1; } + [ "$ok_jupyter" = 1 ] || { echo "Jupyter /login never responded"; exit 1; } + echo "Studio + Jupyter healthy" + env: + STEPS_META_STUDIO_JSON: ${{ steps.meta_studio.outputs.json }} diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml index ec437e0c32..b99112a881 100644 --- a/.github/workflows/studio-backend-ci.yml +++ b/.github/workflows/studio-backend-ci.yml @@ -30,6 +30,9 @@ on: - 'unsloth/**' - 'unsloth_cli/**' - 'tests/**' + # The validate_studio_features.py step below guards docker/jupyter and the + # docker notebook helpers, so a docker-only change must trigger this CI. + - 'docker/**' # The root installers: tests/sh/*.sh and tests/studio/install/* assert # against these two files, so a change here must run the suite that # covers it. Without them an install-only edit (the shape most AMD/ROCm @@ -253,3 +256,7 @@ jobs: [ "$found" -gt 0 ] || { echo "::error::no shell tests discovered under tests/sh"; exit 1; } echo "ran $found shell installer test files" + - name: Docker JupyterLab/notebook feature validation + # Named validate_studio_features.py (not test_*.py) so pytest skips it; + # run explicitly so notebook/Colab/branding regressions fail CI. + run: python tests/validate_studio_features.py diff --git a/.gitignore b/.gitignore index fafd17aa95..0a5b53975b 100644 --- a/.gitignore +++ b/.gitignore @@ -237,6 +237,8 @@ package-lock.json !studio/backend/core/data_recipe/oxc-validator/package-lock.json !studio/package-lock.json llama.cpp/ +async_task_outputs/ +individual_reviews/ # Stray "~" dir some tools create from a literal ~ TMPDIR; never part of the repo. ~/ /temp/ diff --git a/docker/.dockerignore b/docker/.dockerignore new file mode 100644 index 0000000000..18ff6a52eb --- /dev/null +++ b/docker/.dockerignore @@ -0,0 +1,37 @@ +** +!Dockerfile +!entrypoint.sh +!smoke_test.py +!fetch_llama_prebuilt.py +!supervisord.conf +!studio_launch.sh +!unsloth_studio_update.sh +!unsloth_llama_update.sh +!unsloth_jupyter_tunnel.sh +!unsloth_nb_compat.py +!unsloth_pip_shim.py +!unsloth_nb_pip_magic.py +!unsloth_ipython_startup.py +!unsloth_run.py +!unsloth_sync_notebooks.sh +!unsloth_nb_content_sig.py +!unsloth_nb_view.py +!unsloth_nb_strip_colab.py +!unsloth_colab_compat.py +!jupyter +!jupyter/unsloth_branding.py +!jupyter/jupyter_server_config.d +!jupyter/jupyter_server_config.d/** +!jupyter/overrides.json +!jupyter/favicon.ico +!jupyter/logo.png +!jupyter/login.html +!jupyter/install_sloth_stickers.py +!jupyter/unsloth_labext +!jupyter/unsloth_labext/package.json +!jupyter/unsloth_labext/tsconfig.json +!jupyter/unsloth_labext/.yarnrc.yml +!jupyter/unsloth_labext/src +!jupyter/unsloth_labext/src/** +!jupyter/unsloth_labext/style +!jupyter/unsloth_labext/style/** diff --git a/docker/Dockerfile b/docker/Dockerfile new file mode 100644 index 0000000000..06f654ca7a --- /dev/null +++ b/docker/Dockerfile @@ -0,0 +1,623 @@ +# syntax=docker/dockerfile:1.7 +# ----------------------------------------------------------------------------- +# Unsloth + unsloth-zoo for every current NVIDIA arch (Turing -> Blackwell), +# on linux/amd64 and linux/arm64. +# +# Why it works: +# * cu128 wheels ship native SASS (no PTX), verified via `cuobjdump --list-elf`: +# amd64: sm_70 sm_75 sm_80 sm_86 sm_90 sm_100 sm_120 +# arm64: sm_80 sm_90 sm_90a sm_100 sm_100a sm_120 sm_120a +# * SASS is forward-compatible within a major: sm_86->sm_89 (Ada), +# sm_100->sm_103 (B300/GB300), sm_120->sm_121 (DGX Spark/GB10), so every +# non-Jetson GPU on https://developer.nvidia.com/cuda/gpus runs precompiled +# SASS (torch, llama.cpp, source-built ops). +# * Triton kernels JIT per-device at first run; the bundled cu12.8 ptxas/NVRTC +# cannot emit compute_103/compute_121, so the cu13 override below handles +# amd64 sm_103 and arm64 sm_121 (SASS still runs there via forward-compat, so +# only JIT-heavy paths need it). +# * Rare source builds compile against +# TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;8.9;9.0;10.0;12.0+PTX"; the host GPU is +# irrelevant, nvcc emits whatever the arch list says. +# +# Cross-arch build (arm64 / sm_121): built via QEMU binfmt on an x86_64 host +# (`docker run --privileged --rm tonistiigi/binfmt --install all` once, then +# `docker buildx build --platform linux/arm64 ...`). QEMU is build-time only; +# the image runs natively on aarch64. xformers has no cu128 aarch64 wheel, so +# arm64 falls back to Unsloth's SDPA (~5-10% slower, functionally complete). +# +# Build host needs Docker buildkit + buildx, and QEMU binfmt for arm64-on-x86_64; +# nvidia-container-toolkit only for test-time `--gpus all`. No GPU at build time. +# ----------------------------------------------------------------------------- + +ARG CUDA_VERSION=12.8.1 +ARG UBUNTU_VERSION=24.04 +ARG PYTHON_VERSION=3.12 + +# Stage 1: builder -- toolkit + dev headers, builds any source extensions. +FROM nvidia/cuda:${CUDA_VERSION}-cudnn-devel-ubuntu${UBUNTU_VERSION} AS builder + +# TARGETARCH (buildx: amd64/arm64) selects the unsloth extras matching the +# wheels available for the platform (xformers aarch64 gap -- see header). +ARG TARGETARCH +ARG PYTHON_VERSION +ENV DEBIAN_FRONTEND=noninteractive \ + PIP_NO_CACHE_DIR=1 \ + PIP_DISABLE_PIP_VERSION_CHECK=1 \ + PYTHONDONTWRITEBYTECODE=1 \ + PYTHONUNBUFFERED=1 \ + # Cross-compile for every current NVIDIA arch (developer.nvidia.com/cuda/gpus): + # sm_75 Turing (T4, RTX 20xx) | sm_80 A100/A30 | sm_86 A40/RTX 30xx + # sm_89 Ada (L4/L40/RTX 40xx) | sm_90 Hopper (H100/H200/GH200) + # sm_100 Blackwell DC (B100/B200/GB200) | sm_120 Blackwell (RTX 50xx, RTX PRO 6000) + # sm_103 (B300/GB300) and sm_121 (GB10) omitted: CUDA 12.8 nvcc can't compile + # them; sm_100/sm_120 SASS covers them via forward-compat. +PTX lets future + # revisions JIT. Same list on both arches. + TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;8.9;9.0;10.0;12.0+PTX" \ + MAX_JOBS=4 \ + CUDA_HOME=/usr/local/cuda \ + # Build-host-independence guards: the build must NEVER introspect a GPU so all + # hosts yield byte-identical images. + # 1) no JIT-compiled sm_NNN blob into unsloth_compiled_cache/ at import. + UNSLOTH_COMPILE_DISABLE=1 \ + UNSLOTH_COMPILE_OVERWRITE=0 \ + # 2) don't probe torch.cuda.is_available() at setup (would silently skip wheels). + UNSLOTH_DISABLE_GPU_PROBE=1 \ + # 3) empty CUDA_VISIBLE_DEVICES so stray torch.cuda calls see no devices + # (re-enabled at runtime via `docker run --gpus all`). + CUDA_VISIBLE_DEVICES="" + +RUN apt-get update && apt-get install -y --no-install-recommends \ + software-properties-common ca-certificates curl git build-essential \ + ninja-build cmake pkg-config \ + && add-apt-repository -y ppa:deadsnakes/ppa \ + && apt-get update && apt-get install -y --no-install-recommends \ + python${PYTHON_VERSION} python${PYTHON_VERSION}-venv python${PYTHON_VERSION}-dev \ + && ln -sf /usr/bin/python${PYTHON_VERSION} /usr/local/bin/python \ + && ln -sf /usr/bin/python${PYTHON_VERSION} /usr/local/bin/python3 \ + && rm -rf /var/lib/apt/lists/* + +# Isolated prefix; never touch the system Python (PEP 668 externally-managed). +# The venv bootstraps pip via ensurepip and gets uv a few lines below. +ENV VENV=/opt/unsloth-venv +RUN python -m venv ${VENV} && ${VENV}/bin/pip install -U pip wheel setuptools + +# Unified install: torch + triton + bitsandbytes + unsloth + unsloth_zoo in a +# SINGLE uv pass. Mandatory -- splitting it lets bnb's transitive `cuda-toolkit` +# silently upgrade torch to 2.12.0+cu130, breaking the pinned cu128 xformers wheel. +# +# Flags: +# --index-strategy unsafe-best-match: the PyTorch index serves an old +# requests==2.28.1 conflicting with datasets>=2.32.2; both indexes are equally +# trusted, so override uv's first-wins. +# --extra-index-url .../cu128: torch +cu128 wheels + the xformers/cu128 URLs. +# +# Plain `huggingface` extra + explicit xformers pin (amd64): the cu128 extras on +# main stop at torch2100, conflicting with the torch 2.11.0 held below. Pinning +# xformers==0.0.35 (untied to torch) keeps this self-contained; arm64 stays +# xformers-less (no cu128 aarch64 wheel). +# +# No flash-attn: FA3 is hard-refused on Blackwell (Dao-AILab/flash-attention#1810); +# FA2 has no cu128+torch2.11+cp312 wheel and Unsloth falls back to xformers/SDPA. +# Ampere/Ada/Hopper users can `pip install flash-attn` at deploy time. +ARG UNSLOTH_REF=main +ARG UNSLOTH_ZOO_REF=main +RUN set -eux \ + && case "${TARGETARCH:-amd64}" in \ + amd64) UNSLOTH_EXTRA="huggingface"; XFORMERS_PIN="xformers==0.0.35" ;; \ + arm64) UNSLOTH_EXTRA="huggingface"; XFORMERS_PIN="" ;; \ + *) echo "ERROR: unsupported TARGETARCH=${TARGETARCH}" >&2; exit 1 ;; \ + esac \ + && echo ">> TARGETARCH=${TARGETARCH:-amd64}, unsloth extra=[${UNSLOTH_EXTRA}], xformers=[${XFORMERS_PIN}]" \ + && ${VENV}/bin/pip install uv \ + && ${VENV}/bin/uv pip install \ + --python ${VENV}/bin/python \ + --index-strategy unsafe-best-match \ + --extra-index-url https://download.pytorch.org/whl/cu128 \ + "torch==2.11.0" "torchvision==0.26.0" "torchaudio==2.11.0" \ + ${XFORMERS_PIN} \ + "triton>=3.6.0" \ + "bitsandbytes>=0.49.2,!=0.46.0,!=0.48.0" \ + "unsloth_zoo @ git+https://github.com/unslothai/unsloth-zoo@${UNSLOTH_ZOO_REF}" \ + "unsloth[${UNSLOTH_EXTRA}] @ git+https://github.com/unslothai/unsloth@${UNSLOTH_REF}" \ + `# structlog is a studio backend dep, not an unsloth[huggingface] dep,` \ + `# but unsloth_cli's train / export / chat / list-checkpoints all import` \ + `# studio.backend.core.*, so without it every one of them dies on` \ + `# ModuleNotFoundError. The last builder stage imports it as a guard.` \ + "timm>=1.0.11" "addict" "structlog" + +# vLLM: required by Unsloth's GRPO path (fast_inference=True). A SECOND uv pass so +# torch 2.11.0 settles first; with torch held, uv picks the newest compatible vLLM +# (0.20+ pins torch 2.11.0). PyPI ships x86_64 + aarch64 wheels since 0.17. amd64 +# failures abort, arm64 is fail-soft (aarch64 kernels validated on Spark, not CI). +# https://docs.vllm.ai/en/latest/getting_started/installation/gpu/ +# https://wheels.vllm.ai/nightly +ARG INSTALL_VLLM=auto +RUN set -eux \ + && WANT_VLLM=0 \ + && case "${INSTALL_VLLM}" in \ + auto|1|true|yes) WANT_VLLM=1 ;; \ + 0|false|no) WANT_VLLM=0 ;; \ + *) echo "ERROR: invalid INSTALL_VLLM=${INSTALL_VLLM}" >&2; exit 1 ;; \ + esac \ + && if [ "${WANT_VLLM}" = "1" ]; then \ + echo ">> installing vLLM (TARGETARCH=${TARGETARCH:-amd64})"; \ + # Explicit && chain, not `set -e` -- POSIX shells disable errexit inside a + # condition context (verified on dash), masking install failures. + # 1: uv resolves vLLM's deps with torch==2.11.0 held (fails loudly if none). + # 2: vLLM pulls numpy down to 2.2.6 with a broken numpy.testing that breaks + # `import unsloth`; upgrade numpy back to a self-consistent release. + # 3: vLLM pins numba 0.61.2 (refuses numpy>=2.3); lift numba to one + # supporting numpy 2.4 (0.65 imports cleanly, vllm still imports). + { ${VENV}/bin/uv pip install \ + --python ${VENV}/bin/python \ + --pre \ + --index-strategy unsafe-best-match \ + --extra-index-url https://wheels.vllm.ai/nightly \ + --extra-index-url https://download.pytorch.org/whl/cu128 \ + "torch==2.11.0" \ + vllm \ + && ${VENV}/bin/uv pip install \ + --python ${VENV}/bin/python \ + --upgrade "numpy>=2.4" \ + && ${VENV}/bin/uv pip install \ + --python ${VENV}/bin/python \ + --upgrade "numba>=0.62" \ + && ${VENV}/bin/python -c "import vllm; print('vllm', vllm.__version__)" \ + && ${VENV}/bin/python -c "import numpy.testing, numpy; print('numpy', numpy.__version__, 'testing ok')" \ + && ${VENV}/bin/python -c "import numba; print('numba', numba.__version__, 'imports ok')" \ + # flashinfer-jit-cache: precompiled cubins so flashinfer ops skip the JIT + # path (standalone `vllm serve` dies there for fmha_gen on sm_100a). ~1.5 GB. + # The version MUST equal the flashinfer-python vLLM resolved: flashinfer + # raises at import when the two disagree, which takes the vLLM EngineCore + # down with it and breaks Unsloth's GRPO fast_inference path. So read the + # resolved version instead of pinning a literal that drifts. + && FI_VER="$(${VENV}/bin/python -c 'from importlib.metadata import version; print(version("flashinfer-python"))')" \ + && echo ">> flashinfer-python ${FI_VER}, matching flashinfer-jit-cache" \ + && { ${VENV}/bin/uv pip install \ + --python ${VENV}/bin/python \ + --index-url https://flashinfer.ai/whl/cu128 \ + "flashinfer-jit-cache==${FI_VER}" \ + || echo ">> flashinfer-jit-cache ${FI_VER} unavailable for ${TARGETARCH:-amd64}; vllm serve may require nvcc for uncached ops"; } \ + # Whatever happened above, flashinfer has to import: a version mismatch + # here is silent until the first vLLM engine start. + && ${VENV}/bin/python -c \ + "import flashinfer; print('OK: flashinfer', flashinfer.__version__, 'imports')" \ + && echo ">> vLLM installed (numpy + numba re-upgraded post-vllm)"; \ + } || { \ + if [ "${TARGETARCH:-amd64}" != "amd64" ]; then \ + echo ">> vLLM skipped on ${TARGETARCH}: install or import check failed (fail-soft on non-amd64)"; \ + # A partial install must not poison the base stack: drop vllm and + # restore the numpy/numba floor it may have moved. arm64 staging CI + # re-verifies `import unsloth` after this. + ${VENV}/bin/uv pip uninstall --python ${VENV}/bin/python vllm || true; \ + ${VENV}/bin/uv pip install --python ${VENV}/bin/python \ + --upgrade "numpy>=2.4" "numba>=0.62"; \ + ${VENV}/bin/python -c "import numpy.testing, numba; print('numpy/numba restored')"; \ + else \ + echo "ERROR: vLLM install failed on amd64" >&2; exit 1; \ + fi; \ + }; \ + else \ + echo ">> vLLM skipped (INSTALL_VLLM=${INSTALL_VLLM}, TARGETARCH=${TARGETARCH:-amd64})"; \ + fi + +# JupyterLab so the image runs unslothai/notebooks out of the box: +# docker run --gpus all -p 8888:8888 unsloth/unsloth \ +# jupyter lab --ip 0.0.0.0 --port 8888 --allow-root --no-browser +# Separate pass AFTER the torch pin: pure-Python, never names torch, so uv can't +# disturb the cu128 pin set. Declared by notebook install cells, so bake them: +# matplotlib plotting; some trust_remote_code files import it (DeepSeek-OCR) +# soundfile TTS audio read/write (bundles libsndfile) +# evaluate+jiwer Whisper WER metric +# tensorboard default TrainingArguments report_to backend +# langid DeepSeek-R1 GRPO reward language-id check +# easydict some vision trust_remote_code modeling files +# protobuf slow->fast tokenizer conversion for sentencepiece +# omegaconf TTS + NeMo-Gym RL notebook configs +# einx TTS codec tensor-rearrange (Llasa/Oute/Spark) +# librosa Whisper audio features (pulls numba, already pinned >=0.65) +# ftfy Oute TTS text normalisation +# decord is separate below (no aarch64 wheel). Pinned (==) for reproducible +# rebuilds. The resolve must NOT move torch/numpy/numba (asserted below). +RUN ${VENV}/bin/uv pip install \ + --python ${VENV}/bin/python \ + "jupyterlab==4.6.0" "notebook==7.6.0" "ipywidgets==8.1.8" "matplotlib==3.11.0" \ + "soundfile==0.14.0" "evaluate==0.4.6" "jiwer==4.0.0" "tensorboard==2.20.0" \ + "langid==1.1.6" "easydict==1.13" "protobuf==6.33.6" \ + "omegaconf==2.3.1" "einx==0.4.3" "librosa==0.11.0" "ftfy==6.3.1" \ + && ${VENV}/bin/python -c "import torch, numpy, numba; from packaging.version import Version; assert torch.__version__.startswith('2.11.0'), torch.__version__; assert Version(numpy.__version__) >= Version('2.3'), numpy.__version__; assert Version(numba.__version__) >= Version('0.65'), numba.__version__; print('notebook-deps pins OK:', torch.__version__, numpy.__version__, numba.__version__)" + +# decord (ERNIE-VL video decode) has wheels only for x86_64. Installed alone: +# HARD on amd64 (a missing wheel is a real regression), fail-soft elsewhere. +RUN if [ "${TARGETARCH:-amd64}" = "amd64" ]; then \ + ${VENV}/bin/uv pip install --python ${VENV}/bin/python "decord==0.6.0"; \ + else \ + ${VENV}/bin/uv pip install --python ${VENV}/bin/python "decord==0.6.0" \ + || echo ">> decord skipped (no matching wheel for ${TARGETARCH:-}); ERNIE-VL video decode unavailable"; \ + fi + +# Audio decode out of the box (torchcodec). Three traps: (1) torchcodec 0.11 must +# pair with torch 2.11; (2) the wheel must come from cu128, not the PyPI cu13 +# default; (3) its libs dlopen venv torch/NVIDIA libs registered via ld.so.conf.d +# in the runtime stage. Fail-soft on arches without a matching wheel. +RUN set -eux \ + && { ${VENV}/bin/uv pip install \ + --python ${VENV}/bin/python \ + --index-url https://download.pytorch.org/whl/cu128 \ + "torchcodec==0.11.0" \ + && ${VENV}/bin/uv pip install --python ${VENV}/bin/python nvidia-npp-cu12; } \ + || echo ">> torchcodec bake skipped (no matching wheel for ${TARGETARCH:-amd64})" + +# transformers SIDECARS for per-notebook version activation (see +# unsloth_nb_compat.py). Each sidecar is transformers==X + matched +# huggingface_hub/tokenizers/safetensors, --no-deps into its own --target under +# ${VENV}/tf-sidecars. Prepending one to sys.path swaps transformers without +# touching the cu128 base. Candidate versions mirror Studio's tiers (4.57.6 + +# 5.3.0/5.5.0/5.10.2). Fail-soft per arch/wheel. +# +# Every candidate is then VERIFIED against the baked vLLM and dropped if it does +# not survive, because vLLM is version-locked to transformers and a sidecar it +# cannot import does not give the notebook an older transformers -- it gives it +# an ImportError at `import unsloth`, before the first model cell. Measured on +# this image (vLLM 0.26.0): 4.57.6 raises "Support for Transformers v4 ... was +# removed in vLLM v0.24.0" and 5.3.0 raises "cannot import name +# 'ALLOWED_LAYER_TYPES'", between them breaking 254 of the 433 shipped notebooks, +# whose transformers pins select exactly those two. 5.5.0 and 5.10.2 pass. +# +# vllm.transformers_utils.config is the gate because it is the vLLM module that +# reads the transformers API, it reproduces BOTH failures, and it imports without +# a GPU (the build host has none, so `import unsloth` cannot be used here). +# Deriving the kept set instead of hardcoding it means a later vLLM bump that +# widens or narrows the supported range re-tunes the image by itself. The lowest +# survivor is recorded as the selection FLOOR read by unsloth_nb_compat. +RUN set -eux \ + && if ${VENV}/bin/python -c "import vllm" >/dev/null 2>&1; then HAVE_VLLM=1; else HAVE_VLLM=0; fi \ + && echo ">> sidecar verification: baked vLLM importable=${HAVE_VLLM}" \ + && KEPT="" \ + && for TFV in 4.57.6 5.3.0 5.5.0 5.10.2; do \ + SCRATCH="$(mktemp -d)"; \ + if ! ${VENV}/bin/uv pip install --python ${VENV}/bin/python \ + --target "$SCRATCH" "transformers==${TFV}" >/dev/null 2>&1; then \ + echo ">> sidecar resolve failed for ${TFV}; skipping"; rm -rf "$SCRATCH"; continue; \ + fi; \ + pin() { ls -d "$SCRATCH/$1"-*.dist-info 2>/dev/null \ + | sed -E "s@.*/$1-([0-9][0-9A-Za-z.]*)\.dist-info@\1@" | head -1; }; \ + HFV="$(pin huggingface_hub)"; TKV="$(pin tokenizers)"; SFV="$(pin safetensors)"; \ + rm -rf "$SCRATCH"; \ + DEST="${VENV}/tf-sidecars/t_$(echo "${TFV}" | tr . _)"; \ + ${VENV}/bin/uv pip install --python ${VENV}/bin/python --target "$DEST" --no-deps \ + "transformers==${TFV}" \ + ${HFV:+"huggingface_hub==${HFV}"} \ + ${TKV:+"tokenizers==${TKV}"} \ + ${SFV:+"safetensors==${SFV}"}; \ + if [ "$HAVE_VLLM" = "1" ] && ! PYTHONPATH="$DEST" ${VENV}/bin/python \ + -c "import vllm.transformers_utils.config" >/dev/null 2>&1; then \ + echo ">> sidecar transformers==${TFV} DROPPED -- the baked vLLM cannot import under it:"; \ + PYTHONPATH="$DEST" ${VENV}/bin/python \ + -c "import vllm.transformers_utils.config" 2>&1 | tail -2 || true; \ + rm -rf "$DEST"; \ + continue; \ + fi; \ + KEPT="${KEPT} ${TFV}"; \ + echo ">> sidecar transformers==${TFV} kept (hf_hub=${HFV} tokenizers=${TKV} safetensors=${SFV})"; \ + done \ + && if [ -z "$KEPT" ]; then \ + echo ">> FATAL: no transformers sidecar survived vLLM verification"; exit 1; \ + fi \ + && if [ "$HAVE_VLLM" = "1" ]; then \ + printf '%s\n' $KEPT | sort -V | head -1 > ${VENV}/tf-sidecars/.vllm_min_transformers; \ + fi \ + && echo ">> sidecars kept:${KEPT} floor=$(cat ${VENV}/tf-sidecars/.vllm_min_transformers 2>/dev/null || echo '(none)')" \ + && { du -sh ${VENV}/tf-sidecars || true; } + +# Informational pin record (NOT byte-reproducible: pip freeze omits wheel hashes +# and unsloth/vllm --pre float from VCS/nightly). +RUN ${VENV}/bin/pip freeze --exclude-editable > ${VENV}/requirements.lock.txt \ + && head -50 ${VENV}/requirements.lock.txt + +# Strip pip cache & __pycache__ to shrink the runtime layer. The `-name tests` +# strip excludes numpy's tests dirs (numpy 2.4 needs numpy/_core/tests/ or +# `import numpy` breaks). Other verified-safe cuts: +# * npp: torchcodec dlopens only libnppicc + libnppc; drop the rest (~388MB). +# * static .a archives (~143MB): link-time only. +# * nvshmem device .bc (~30MB): device-relink only; host .so kept. +# Do NOT strip headers (torch/include): causal-conv1d / mamba-ssm build against +# them at notebook time with --no-build-isolation. +RUN set -eux \ + && find ${VENV} -depth -type d -name __pycache__ -exec rm -rf {} + \ + && find ${VENV} -depth -type d -name tests \ + ! -path "*numpy/_core/tests*" \ + ! -path "*numpy/tests*" \ + ! -path "*numpy/ma/tests*" \ + -exec rm -rf {} + \ + && rm -rf /root/.cache/pip /root/.cache/uv \ + && SP=${VENV}/lib/python${PYTHON_VERSION}/site-packages \ + && if [ -d "$SP/nvidia/npp/lib" ]; then \ + find "$SP/nvidia/npp/lib" -maxdepth 1 -name 'libnpp*.so.*' \ + ! -name 'libnppicc.so.*' ! -name 'libnppc.so.*' -delete; \ + fi \ + && find ${VENV} -name '*.a' -delete \ + && rm -f "$SP"/nvidia/nvshmem/lib/libnvshmem_device.bc \ + && echo "venv size after prune:" && du -sh ${VENV} + +# Build-time verification. +# (1) arch-list check uses the RAW C++ accessor: torch.cuda.get_arch_list() +# returns [] with no GPU visible (CUDA_VISIBLE_DEVICES is empty here). +# (2) required packages verified via metadata only -- we do NOT import unsloth/ +# unsloth_zoo (their __init__ needs a real CUDA device). Import correctness is +# exercised at deploy time by smoke_test.py with --gpus all. +RUN TARGETARCH="${TARGETARCH:-amd64}" ${VENV}/bin/python - <<'PY' +import os, platform +target = os.environ.get("TARGETARCH", "amd64") +mach = platform.machine() +print(f"build target: TARGETARCH={target} platform.machine()={mach}") + +import torch +arches = torch._C._cuda_getArchFlags().split() +print("torch", torch.__version__, "cuda", torch.version.cuda) +print("arches:", arches) +assert torch.__version__.startswith("2.11.0"), f"torch silently moved: {torch.__version__}" +assert "+cu128" in torch.__version__, f"cu build silently changed: {torch.__version__}" +assert "sm_100" in arches, f"sm_100 (B200/GB200) missing: {arches}" +# cu128 wheels ship sm_120 native SASS on both amd64 and aarch64. On arm64 DGX +# Spark (sm_121) runs it via forward-compat; sm_121 is never in a cu128 wheel. +assert "sm_120" in arches, f"sm_120 missing: {arches}" +print(f"OK: torch 2.11.0+cu128 with sm_100 + sm_120 native SASS intact ({target})") + +from importlib.metadata import version, PackageNotFoundError +# xformers is amd64-only (aarch64 wheel gap -- see header). +REQUIRED = ["torch", "triton", "bitsandbytes", "unsloth", + "unsloth_zoo", "transformers", "trl", "peft", "accelerate"] +if target == "amd64": + REQUIRED.insert(2, "xformers") +missing = [] +for pkg in REQUIRED: + try: + v = version(pkg.replace("_", "-")) + print(f" {pkg:14s} {v}") + except PackageNotFoundError: + missing.append(pkg) +if missing: + raise SystemExit(f"FAIL: missing wheels: {missing}") +print("OK: all required wheels present") + +# Lightweight imports: these init without touching CUDA, unlike unsloth. +import importlib +LIGHT_IMPORTS = ["bitsandbytes", "triton"] +if target == "amd64": + LIGHT_IMPORTS.insert(0, "xformers") +for pkg in LIGHT_IMPORTS: + importlib.import_module(pkg) +print(f"OK: {' + '.join(LIGHT_IMPORTS)} import cleanly on no-GPU host") + +# Guard for the studio.backend.core.* closure the unsloth CLI needs (structlog, +# plus starlette via the logging handlers). Runs last in the builder, after vLLM, +# because that is what pulls starlette in. +from studio.backend.core.export import ExportBackend # noqa: F401 +print("OK: the unsloth CLI can reach the studio export backend") +PY + +# Stage 2: runtime -- slim, no nvcc, no cuDNN/cuBLAS layers. +# The "-base-" variant drops ~2.7 GB of system CUDA libs we never load: torch +# wheels bake their own cuDNN/cuBLAS into torch/lib/ and resolve via RPATH. The +# base still provides nvidia-smi + libcuda stubs + libnvidia-ml. +FROM nvidia/cuda:${CUDA_VERSION}-base-ubuntu${UBUNTU_VERSION} AS runtime + +# The base manifest is multi-arch; buildx picks the right one for +# TARGETPLATFORM at this FROM line, no conditional needed. +ARG TARGETARCH +ARG PYTHON_VERSION +ARG CUDA_VERSION +ENV DEBIAN_FRONTEND=noninteractive \ + PIP_NO_CACHE_DIR=1 \ + PYTHONDONTWRITEBYTECODE=1 \ + PYTHONUNBUFFERED=1 \ + PATH=/opt/unsloth-venv/bin:${PATH} \ + HF_HOME=/workspace/.cache/huggingface \ + TRITON_CACHE_DIR=/workspace/.cache/triton \ + # Keep the arch list at runtime so an in-container source build gets the same + # SASS coverage as the builder (10.3 omitted; cu12.8 can't emit it). + TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;8.9;9.0;10.0;12.0+PTX" + +# System packages needed by the notebooks: +# zstd Ollama installer (`curl ollama.com/install.sh | sh`) extracts a zstd tarball +# ffmpeg torchcodec dlopens system FFmpeg libs (not bundled in the wheel) +# wget notebooks fetch assets with `!wget URL` +# ninja-build flashinfer cpp_ext JIT shells out to ninja +# cuda-nvcc + cudart-dev flash-linear-attention TileLang JIT-compiles CUDA +# kernels via nvcc, absent from the -base image +RUN CUDA_PKG="$(echo "${CUDA_VERSION}" | awk -F. '{print $1"-"$2}')" \ + && apt-get update && apt-get install -y --no-install-recommends \ + software-properties-common ca-certificates curl wget git libgomp1 \ + gcc g++ zstd ffmpeg ninja-build \ + "cuda-nvcc-${CUDA_PKG}" "cuda-cudart-dev-${CUDA_PKG}" \ + && add-apt-repository -y ppa:deadsnakes/ppa \ + && apt-get update && apt-get install -y --no-install-recommends \ + python${PYTHON_VERSION} python${PYTHON_VERSION}-venv python${PYTHON_VERSION}-dev \ + && ln -sf /usr/bin/python${PYTHON_VERSION} /usr/local/bin/python \ + && ln -sf /usr/bin/python${PYTHON_VERSION} /usr/local/bin/python3 \ + && test -x /usr/local/cuda/bin/nvcc \ + && rm -rf /var/lib/apt/lists/* +# gcc + g++ + python3.12-dev in runtime: Triton's nvidia backend compiles a C +# extension (CudaUtils) on first GPU access; without a compiler + headers the +# first forward pass dies with "Failed to find C compiler". ~250MB. + +COPY --from=builder /opt/unsloth-venv /opt/unsloth-venv + +# Blackwell JIT fix for sm_103 (amd64) and sm_121 (arm64) -- the cu12.8 JIT gap. +# Two JIT paths need the cu13 override: +# (1) torch's bundled libnvrtc.so.12 errors on sm_103/sm_121. Fix: stage a cu13 +# NVRTC alias beside the cu12.8 default. +# (2) Triton's bundled ptxas (12.8) rejects sm_103, downgrades sm_121 to sm_80 +# (triton-lang/triton#8335). Fix: cu13 ptxas via TRITON_PTXAS_PATH. +# Both cu13 tools are CPU-side compilers, but their cubin needs a >=580 driver to +# LOAD, so neither is a global default (would break 570-579 drivers). +# select_cuda_jit_tools in entrypoint.sh activates them per device, only for +# sm_103/sm_121 (>=580 drivers). Both arches carry the ~400 MB. +RUN set -eux; \ + # The base already configures the CUDA apt repo with its own Signed-By + # keyring; a second cuda-keyring would make apt-get update refuse the repo. + # The base repo serves 13.x too, so install cu13 packages directly. + apt-get update; \ + apt-get install -y --no-install-recommends \ + cuda-nvrtc-13-0 \ + cuda-nvcc-13-0; \ + # cu13's postinst flips /usr/local/cuda to cuda-13.0; pin it back (cpp + # builds resolve /usr/local/cuda/bin/nvcc, and cu13 cubins need driver + # >= 580 while this image supports 570+). The cu13 tools stay reachable by + # absolute path; --set also stops later apt ops flipping it again. + update-alternatives --set cuda /usr/local/cuda-12.8; \ + rm -rf /var/lib/apt/lists/*; \ + # (1) NVRTC staging: keep the wheel's cu12.8 lib as .cu128.orig, point + # libnvrtc.so.12 at it, stage .cu13 -> the cu13 lib; + # select_cuda_jit_tools retargets the symlink only on sm_103/sm_121. + NVRTC_DIR=/opt/unsloth-venv/lib/python${PYTHON_VERSION}/site-packages/nvidia/cuda_nvrtc/lib; \ + if [ -f "${NVRTC_DIR}/libnvrtc.so.12" ] && [ ! -L "${NVRTC_DIR}/libnvrtc.so.12" ]; then \ + mv "${NVRTC_DIR}/libnvrtc.so.12" "${NVRTC_DIR}/libnvrtc.so.12.cu128.orig"; \ + ln -s libnvrtc.so.12.cu128.orig "${NVRTC_DIR}/libnvrtc.so.12"; \ + ln -s /usr/local/cuda-13.0/lib64/libnvrtc.so.13 "${NVRTC_DIR}/libnvrtc.so.12.cu13"; \ + fi +# (2) ptxas: the cu13 nvcc package above provides it; TRITON_PTXAS_PATH is set +# per device at boot (select_cuda_jit_tools) for the same driver-floor reason. + +# Register the venv's torch + NVIDIA lib dirs with the loader so torchcodec can +# dlopen them. ld.so.conf.d, NOT LD_LIBRARY_PATH: the cache is consulted after +# DT_RUNPATH, so llama.cpp keeps resolving its own $ORIGIN libs first. +# cublas/lib and cu13/lib are here for llama.cpp's libggml-cuda.so, which links +# against libcublas but does not ship it (see the guard after the fetch below). +RUN set -eux \ + && SP=/opt/unsloth-venv/lib/python${PYTHON_VERSION}/site-packages \ + && printf "%s\n" "$SP/torch/lib" "$SP/nvidia/cuda_nvrtc/lib" \ + "$SP/nvidia/cuda_runtime/lib" "$SP/nvidia/npp/lib" \ + "$SP/nvidia/cublas/lib" "$SP/nvidia/cu13/lib" \ + > /etc/ld.so.conf.d/zz-unsloth-venv.conf \ + && ldconfig \ + && { /opt/unsloth-venv/bin/python -c \ + "import torchcodec; print('torchcodec', torchcodec.__version__)" \ + || echo ">> torchcodec unavailable on this arch (audio decode falls back)"; } + +# Prebuilt llama.cpp so GGUF export works out of the box; without it the first +# export hits install_llama_cpp()'s prompt + slow source build. +# +# NOT studio/install_llama_prebuilt.py: it selects a bundle for the CURRENT host, +# but the build must never introspect the host, so release + asset are pinned by +# build target instead (see fetch_llama_prebuilt.py). +# +# /opt (not /root) so it survives `docker run --user`. Default "latest" resolves +# the newest release; build.sh pins a concrete tag so the cache busts only on new +# releases. --build-arg LLAMA_PREBUILT_TAG= for a frozen build. +ARG LLAMA_PREBUILT_TAG=latest +COPY fetch_llama_prebuilt.py /tmp/fetch_llama_prebuilt.py +RUN /opt/unsloth-venv/bin/python /tmp/fetch_llama_prebuilt.py \ + "${LLAMA_PREBUILT_TAG}" "${TARGETARCH:-amd64}" /opt/unsloth/llama.cpp \ + && rm -f /tmp/fetch_llama_prebuilt.py \ + && cat /opt/unsloth/llama.cpp/UNSLOTH_PREBUILT_INFO.json + +# libggml-cuda.so is loaded with dlopen (ggml_backend_dl), links against +# libcublas, and does not ship it; the CUDA runtime base only carries libcudart. +# A missing libcublas therefore makes the backend fail to load SILENTLY and +# llama.cpp runs on the CPU: measured 1.6 tok/s instead of 222 tok/s for +# gemma-4-E2B UD-Q4_K_XL on a B200, with `--list-devices` printing nothing. +# torch's wheels already ship libcublas for their own CUDA major (registered +# with the loader above); install the bundle's major when it differs. Then fail +# the build on any dependency that is still unresolved, so a silent CPU fallback +# can never ship again. libcuda.so.1 is exempt: that is the driver stub, injected +# by nvidia-container-toolkit at `docker run --gpus`, never present in the image. +# ldd needs no GPU, so this keeps the build host-independent. +RUN set -eux \ + && CUDA_SO=/opt/unsloth/llama.cpp/libggml-cuda.so \ + && if [ -f "$CUDA_SO" ]; then \ + want="$(ldd "$CUDA_SO" | sed -n 's/^[[:space:]]*\(libcublas\.so\.[0-9]*\)[[:space:]]*=> not found$/\1/p' | head -n1)"; \ + if [ -n "$want" ]; then \ + major="${want##*.}"; \ + echo ">> $want missing, installing nvidia-cublas-cu${major}"; \ + /opt/unsloth-venv/bin/uv pip install --python /opt/unsloth-venv/bin/python \ + "nvidia-cublas-cu${major}"; \ + ldconfig; \ + fi; \ + missing="$(ldd "$CUDA_SO" | grep 'not found' | grep -v 'libcuda\.so\.1 ' || true)"; \ + if [ -n "$missing" ]; then \ + echo "ERROR: llama.cpp CUDA backend has unresolved libraries:"; \ + echo "$missing"; \ + echo "GGUF inference would silently fall back to the CPU."; \ + exit 1; \ + fi; \ + echo "OK: llama.cpp CUDA backend dependencies all resolve"; \ + else \ + echo ">> no libggml-cuda.so in this bundle (CPU-only build)"; \ + fi +ENV UNSLOTH_LLAMA_CPP_PATH=/opt/unsloth/llama.cpp + +WORKDIR /workspace +# World-writable so `docker run --user ` (documented non-root use) can +# create notebooks and populate the default caches without a bind mount. +RUN mkdir -p ${HF_HOME} ${TRITON_CACHE_DIR} \ + && chmod -R a+rwX /workspace + +# Per-notebook transformers version activation -- run unslothai/notebooks +# UNCHANGED (see unsloth_nb_compat.py). Pieces: +# * unsloth_nb_compat.py: tier detection + sidecar resolution + IPython hook. +# * pip/uv shim on a PATH dir AHEAD of the venv bin: makes `!pip install` cells +# safe + idempotent (keeps the baked stack, records requested transformers). +# * unsloth_nb_pip_magic.py: re-points `%pip`/`%uv` and `!python -m pip` at the +# same shim so in-process installs can't bypass PATH. +# * IPython startup hook: activates the right sidecar before the first model cell. +# * unsloth-run: headless `unsloth-run `, the robust driven path. +COPY unsloth_nb_compat.py unsloth_pip_shim.py unsloth_nb_pip_magic.py unsloth_ipython_startup.py unsloth_run.py unsloth_sync_notebooks.sh unsloth_nb_content_sig.py unsloth_nb_view.py unsloth_nb_strip_colab.py unsloth_colab_compat.py /opt/unsloth-nb/ +RUN set -eux \ + && SP=/opt/unsloth-venv/lib/python${PYTHON_VERSION}/site-packages \ + && cp /opt/unsloth-nb/unsloth_nb_compat.py "$SP/unsloth_nb_compat.py" \ + && cp /opt/unsloth-nb/unsloth_nb_pip_magic.py "$SP/unsloth_nb_pip_magic.py" \ + && cp /opt/unsloth-nb/unsloth_colab_compat.py "$SP/unsloth_colab_compat.py" \ + && chmod +x /opt/unsloth-nb/unsloth_pip_shim.py /opt/unsloth-nb/unsloth_run.py /opt/unsloth-nb/unsloth_sync_notebooks.sh /opt/unsloth-nb/unsloth_nb_content_sig.py /opt/unsloth-nb/unsloth_nb_view.py /opt/unsloth-nb/unsloth_nb_strip_colab.py \ + && mkdir -p /opt/unsloth-nb/bin \ + && for t in pip pip3 uv; do ln -sf /opt/unsloth-nb/unsloth_pip_shim.py /opt/unsloth-nb/bin/$t; done \ + && ln -sf /opt/unsloth-nb/unsloth_run.py /usr/local/bin/unsloth-run \ + && ln -sf /opt/unsloth-nb/unsloth_sync_notebooks.sh /usr/local/bin/unsloth-sync-notebooks \ + && ln -sf /opt/unsloth-nb/unsloth_nb_content_sig.py /usr/local/bin/unsloth-nb-content-sig \ + && ln -sf /opt/unsloth-nb/unsloth_nb_view.py /usr/local/bin/unsloth-nb-view \ + && ln -sf /opt/unsloth-nb/unsloth_nb_strip_colab.py /usr/local/bin/unsloth-nb-strip-colab \ + && mkdir -p /opt/unsloth-nb/ipython/profile_default/startup \ + && cp /opt/unsloth-nb/unsloth_ipython_startup.py /opt/unsloth-nb/ipython/profile_default/startup/00-unsloth-nb.py \ + && chmod -R a+rX /opt/unsloth-nb/ipython \ + && /opt/unsloth-venv/bin/python -c "import sys, glob; sys.path.insert(0, '$SP'); import unsloth_nb_compat, unsloth_colab_compat; print('nb-compat OK; baked sidecars:', sorted(glob.glob('/opt/unsloth-venv/tf-sidecars/t_*')))" \ + && /opt/unsloth-venv/bin/python /opt/unsloth-nb/unsloth_pip_shim.py --unsloth-selfcheck-value-flags +# Shim dir AHEAD of the venv bin so `!pip`/`!uv` resolve to the shim, not the real tool. +ENV PATH=/opt/unsloth-nb/bin:${PATH} +# Load the notebook startup hook for EVERY kernel, any uid: IPYTHONDIR points +# IPython at this shared profile, so it loads under `--user ` too (unlike +# /root/.ipython). Writable state (history.sqlite) still lands per-user. +ENV IPYTHONDIR=/opt/unsloth-nb/ipython + +# Pre-clone unslothai/notebooks so JupyterLab opens with them present. Baked as a +# READ-ONLY template (~206MB, .git stripped); on boot the entrypoint copies it to +# /workspace/unsloth-notebooks and best-effort refreshes from GitHub, never +# overwriting a user-touched notebook (see unsloth_sync_notebooks.sh). +# +# UNSLOTH_NOTEBOOKS_REF pins ONE commit/branch/tag so a multi-arch publish bakes +# identical templates into both legs; default "main" tracks the tip. +ARG UNSLOTH_NOTEBOOKS_REF=main +RUN set -eux \ + && git init -q /opt/unsloth-notebooks \ + && git -C /opt/unsloth-notebooks remote add origin https://github.com/unslothai/notebooks \ + && git -C /opt/unsloth-notebooks fetch -q --depth 1 origin "${UNSLOTH_NOTEBOOKS_REF}" \ + && git -C /opt/unsloth-notebooks checkout -q FETCH_HEAD \ + && git -C /opt/unsloth-notebooks rev-parse HEAD > /opt/unsloth-notebooks/.unsloth_template_commit \ + && rm -rf /opt/unsloth-notebooks/.git \ + && du -sh /opt/unsloth-notebooks + +# Mount a volume on /workspace to persist the notebooks and caches. +EXPOSE 8888 + +COPY smoke_test.py /workspace/smoke_test.py +COPY entrypoint.sh /usr/local/bin/unsloth-entrypoint +RUN chmod +x /usr/local/bin/unsloth-entrypoint + +# Fast GPU pre-flight checks before user code, each with an actionable error (see +# entrypoint.sh). Bypass for offline tooling: docker run -e UNSLOTH_SKIP_GPU_CHECK=1 +ENTRYPOINT ["/usr/local/bin/unsloth-entrypoint"] + +# Override examples: +# docker run --gpus all unsloth/unsloth:latest python /workspace/smoke_test.py +# docker run --gpus all -it unsloth/unsloth:latest bash +CMD ["python"] diff --git a/docker/Dockerfile.studio b/docker/Dockerfile.studio new file mode 100644 index 0000000000..d8e1766c6d --- /dev/null +++ b/docker/Dockerfile.studio @@ -0,0 +1,210 @@ +# Full Unsloth image: base training stack + Studio + JupyterLab + sshd. +# Published as unsloth/unsloth:studio (and default :latest); layers Studio on the +# lean core image and runs Studio:8000, JupyterLab:8888, sshd:22. +# +# Build (local): +# docker buildx build --build-arg BASE_IMAGE=unsloth-blackwell:test \ +# -f docker/Dockerfile.studio -t unsloth-blackwell:studio docker/ +# Run: +# docker run --rm --gpus all -p 8000:8000 -p 8888:8888 \ +# -v $HOME/.cache/huggingface:/workspace/.cache/huggingface unsloth-blackwell:studio +# +# Studio on :8000 (first-boot admin password in the logs, persisted under +# /opt/unsloth-studio/auth/); JupyterLab on :8888 (JUPYTER_PASSWORD env, else a +# random one is printed). Without GPU passthrough add -e UNSLOTH_ALLOW_CPU=1: +# training is unavailable but Studio chat / Data Recipes / GGUF / Jupyter work. +# CI pins BASE_IMAGE to the published base digest so both images ship the same stack. + +ARG BASE_IMAGE=unsloth-blackwell:test + +# Builds the "Unsloth Dark" (Monokai) theme + Colab-style cell-nav keymap. Node +# lives only in this throwaway stage; the final image copies just the prebuilt +# labextension (runtime stays Node-free). Uses the base's bundled jlpm+jupyterlab. +FROM ${BASE_IMAGE} AS labext-builder +ENV DEBIAN_FRONTEND=noninteractive +# JupyterLab 4.6 needs Node >=20; Ubuntu 24.04 ships 18, so pull Node 20 LTS from +# NodeSource. This stage is thrown away, so the apt sources never reach runtime. +RUN apt-get update \ + && apt-get install -y --no-install-recommends ca-certificates curl gnupg git \ + && curl -fsSL https://deb.nodesource.com/setup_20.x | bash - \ + && apt-get install -y --no-install-recommends nodejs \ + && rm -rf /var/lib/apt/lists/* +COPY jupyter/unsloth_labext /opt/labext-src +RUN cd /opt/labext-src \ + && /opt/unsloth-venv/bin/jlpm install \ + && /opt/unsloth-venv/bin/jlpm build:prod + +FROM ${BASE_IMAGE} + +# Studio source ref to clone. Defaults to main; CI pins it (same UNSLOTH_REF as +# the base) so the published image is reproducible. +ARG UNSLOTH_STUDIO_REF=main +# unsloth-zoo ref overlaid into the Studio venv by install.sh --local. The publish +# workflow passes ONE zoo ref to both builds, so Studio runs the same zoo as base. +ARG UNSLOTH_STUDIO_ZOO_REF=main +# The SAME llama.cpp tag the base baked. setup.sh honours UNSLOTH_LLAMA_TAG; +# without the pin the Studio build could re-resolve "latest" and diverge. +ARG LLAMA_PREBUILT_TAG=latest +ARG TARGETARCH + +# Services run as root here (non-root parity is a follow-up). sshd is key-only, +# disabled unless PUBLIC_KEY/SSH_KEY is set (see studio_launch.sh). The +# JUPYTER_PORT / UNSLOTH_ENABLE_SSHD defaults let supervisord's %(ENV_*)s resolve. +USER root +ENV UNSLOTH_STUDIO_HOME=/opt/unsloth-studio \ + JUPYTER_PORT=8888 \ + UNSLOTH_ENABLE_SSHD=false \ + DEBIAN_FRONTEND=noninteractive + +# install.sh needs curl + git; supervisor + openssh-server run the service +# trio. The base image already has python + uv + pip. +RUN apt-get update \ + && apt-get install -y --no-install-recommends \ + curl git ca-certificates supervisor openssh-server \ + && rm -rf /var/lib/apt/lists/* + +# Clone + install Studio into a dedicated venv under $UNSLOTH_STUDIO_HOME. +# --local is editable, so the source MUST persist -- keep it at $STUDIO_HOME/src, +# strip .git (~120MB). +# +# The llama.cpp symlink BEFORE install.sh points Studio's prebuilt dir at the +# base's baked bundle so the installer skips a second ~400MB download; the +# .unsloth-studio-owned marker satisfies setup.sh's ownership assertion. +# +# UNSLOTH_TORCH_INDEX_FAMILY pins the Studio venv's torch index (no nvidia-smi at +# build time would land on cpu/cu126). cu128 on both arches, mirroring the base. +# Blackwell JIT (sm_103/sm_121) comes from the same cu13 NVRTC swap, repeated below. +# +# UNSLOTH_PYTHON=3.12 pins the Studio venv to the base's Python minor so the +# nvidia-*-cu12 wheels are byte-identical and the dedup below can symlink them. +# +# fetch+checkout FETCH_HEAD, not `clone --branch`: CI passes a commit SHA. +RUN set -eux \ + && case "${TARGETARCH:-amd64}" in \ + amd64|arm64) TORCH_FAMILY="cu128" ;; \ + *) echo "ERROR: unsupported TARGETARCH=${TARGETARCH}" >&2; exit 1 ;; \ + esac \ + && mkdir -p "${UNSLOTH_STUDIO_HOME}" \ + && ln -s /opt/unsloth/llama.cpp "${UNSLOTH_STUDIO_HOME}/llama.cpp" \ + && touch /opt/unsloth/llama.cpp/.unsloth-studio-owned \ + && git init -q "${UNSLOTH_STUDIO_HOME}/src" \ + && cd "${UNSLOTH_STUDIO_HOME}/src" \ + && git remote add origin https://github.com/unslothai/unsloth \ + && git fetch -q --depth 1 origin "${UNSLOTH_STUDIO_REF}" \ + && git checkout -q FETCH_HEAD \ + && UNSLOTH_STUDIO_HOME="${UNSLOTH_STUDIO_HOME}" \ + UNSLOTH_TORCH_INDEX_FAMILY="${TORCH_FAMILY}" \ + UNSLOTH_ZOO_REF="${UNSLOTH_STUDIO_ZOO_REF}" \ + UNSLOTH_LLAMA_TAG="${LLAMA_PREBUILT_TAG}" \ + UNSLOTH_PYTHON=3.12 \ + bash install.sh --local \ + # Fail loud unless the Studio venv torch EXACTLY matches the base (version AND + # CUDA family) before the dedup symlinks their CUDA libs. Compare to the base's + # own torch (no hardcoded version); metadata only (QEMU arm64 can't import torch). + && BASE_TORCH="$(/opt/unsloth-venv/bin/python -c "from importlib.metadata import version; print(version('torch'))")" \ + && "${UNSLOTH_STUDIO_HOME}/unsloth_studio/bin/python" -c "import sys; from importlib.metadata import version; assert sys.version_info[:2] == (3, 12), 'Studio venv python %d.%d is not 3.12 (UNSLOTH_PYTHON pin ignored) -- CUDA dedup below depends on it' % sys.version_info[:2]; v = version('torch'); assert v == '${BASE_TORCH}', 'Studio venv torch ' + v + ' does not match base venv torch ${BASE_TORCH} (CUDA dedup would link mismatched libs)'; print('Studio venv python %d.%d torch' % sys.version_info[:2], v, '== base', '${BASE_TORCH}')" \ + # setup.sh may relink llama-quantize into build/bin; prove it still resolves its + # libraries. Content check, not rc: --help exits nonzero but prints usage. + && { "${UNSLOTH_STUDIO_HOME}/llama.cpp/llama-quantize" --help 2>&1 || true; } | grep -q "usage" \ + && rm -rf "${UNSLOTH_STUDIO_HOME}/src/.git" \ + "${UNSLOTH_STUDIO_HOME}/src/studio/frontend/node_modules" \ + /root/.cache \ + # Stage the Studio venv's NVRTC like the base (.cu128.orig default + .cu13 + # alias, retargeted per device by select_cuda_jit_tools). Both arches. + && for NVRTC_DIR in "${UNSLOTH_STUDIO_HOME}"/unsloth_studio/lib/python*/site-packages/nvidia/cuda_nvrtc/lib; do \ + if [ -f "${NVRTC_DIR}/libnvrtc.so.12" ] && [ ! -L "${NVRTC_DIR}/libnvrtc.so.12" ]; then \ + mv "${NVRTC_DIR}/libnvrtc.so.12" "${NVRTC_DIR}/libnvrtc.so.12.cu128.orig"; \ + ln -s libnvrtc.so.12.cu128.orig "${NVRTC_DIR}/libnvrtc.so.12"; \ + ln -s /usr/local/cuda-13.0/lib64/libnvrtc.so.13 "${NVRTC_DIR}/libnvrtc.so.12.cu13"; \ + fi; \ + done \ + && BASE_NV=/opt/unsloth-venv/lib/python3.12/site-packages/nvidia \ + && STU_NV="${UNSLOTH_STUDIO_HOME}/unsloth_studio/lib/python3.12/site-packages/nvidia" \ + && if [ ! -d "${STU_NV}" ] || [ ! -d "${BASE_NV}" ]; then \ + echo ">> nvidia dir missing (STU=${STU_NV} BASE=${BASE_NV}); skipping CUDA dedup"; \ + else \ + find "${UNSLOTH_STUDIO_HOME}/unsloth_studio" -name '*.a' -delete; \ + rm -f "${STU_NV}/nvshmem/lib/libnvshmem_device.bc"; \ + for c in cudnn cublas cusparselt nccl cusolver cusparse cufft curand nvjitlink cuda_cupti nvshmem npp; do \ + b="${BASE_NV}/${c}/lib"; s="${STU_NV}/${c}/lib"; \ + { [ -d "$b" ] && [ -d "$s" ]; } || { echo ">> skip ${c} (dir missing)"; continue; }; \ + if [ "${c}" = "npp" ]; then \ + rm -rf "$s" && ln -s "$b" "$s" && readlink -e "$s" >/dev/null; \ + echo ">> deduped npp -> base (pruned)"; \ + elif [ "$(cd "$s" && ls | sort | tr '\n' ' ')" = "$(cd "$b" && ls | sort | tr '\n' ' ')" ]; then \ + rm -rf "$s" && ln -s "$b" "$s" && readlink -e "$s" >/dev/null; \ + echo ">> deduped ${c} -> base"; \ + else \ + echo ">> skip ${c} (file set differs base vs studio)"; \ + fi; \ + done; \ + echo "studio venv size after dedup:"; du -sh "${UNSLOTH_STUDIO_HOME}/unsloth_studio"; \ + fi + +COPY supervisord.conf /etc/supervisor/supervisord.conf +COPY studio_launch.sh /usr/local/bin/unsloth-studio-launch +# In-place updaters (no image pull): +# unsloth-studio-update refresh Studio packages (backend + frontend) and restart +# unsloth-llama-update swap the baked llama.cpp prebuilt to the latest release +COPY unsloth_studio_update.sh /usr/local/bin/unsloth-studio-update +COPY unsloth_llama_update.sh /usr/local/bin/unsloth-llama-update +# unsloth-llama-update reuses the build-time fetcher (redirect-based, not rate- +# limited; deterministic portable bundle) rather than the host-probing installer. +COPY fetch_llama_prebuilt.py /usr/local/lib/unsloth/fetch_llama_prebuilt.py +# Optional public Cloudflare tunnel for JupyterLab (UNSLOTH_JUPYTER_CLOUDFLARE=1, +# or `unsloth-jupyter-tunnel --force`); supervisord runs it as jupyter-cloudflare. +COPY unsloth_jupyter_tunnel.sh /usr/local/bin/unsloth-jupyter-tunnel +# JupyterLab defaults baked for every container (theme, non-advancing run button, +# labeled "Restart & Run All", windowing off, cell-nav keymap, news prompt off). +# overrides.json is the settings override; theme + keymap + logo ship as the +# prebuilt labextension from labext-builder above. +COPY jupyter/overrides.json /opt/unsloth-venv/share/jupyter/lab/settings/overrides.json +COPY --from=labext-builder /opt/labext-src/unsloth-jupyterlab/labextension /opt/unsloth-venv/share/jupyter/labextensions/unsloth-jupyterlab +# Unsloth branding (applied to jupyter_server's site-packages): replace favicon + +# logo, brand login.html, disable+lock the stock top-left logo. Only the +# sloth-sticker install is fail-soft (`|| echo`); the copies above stay fatal. +COPY jupyter/favicon.ico /tmp/unsloth-branding/favicon.ico +COPY jupyter/logo.png /tmp/unsloth-branding/logo.png +COPY jupyter/login.html /tmp/unsloth-branding/login.html +COPY jupyter/install_sloth_stickers.py /tmp/unsloth-branding/install_sloth_stickers.py +RUN JS="$(/opt/unsloth-venv/bin/python -c 'import os, jupyter_server; print(os.path.dirname(jupyter_server.__file__))')" \ + && for n in favicon.ico favicon-notebook.ico favicon-file.ico favicon-terminal.ico; do \ + cp /tmp/unsloth-branding/favicon.ico "${JS}/static/favicons/${n}"; \ + done \ + && cp /tmp/unsloth-branding/logo.png "${JS}/static/logo/logo.png" \ + && cp /tmp/unsloth-branding/login.html "${JS}/templates/login.html" \ + && { /opt/unsloth-venv/bin/python /tmp/unsloth-branding/install_sloth_stickers.py \ + --src "${UNSLOTH_STUDIO_HOME}/src/studio/frontend/public/Sloth emojis" \ + --dest "${JS}/static/sloth" \ + || echo ">> sloth stickers not installed (login falls back to the Unsloth logo)"; } \ + && rm -rf /tmp/unsloth-branding \ + && /opt/unsloth-venv/bin/jupyter labextension disable @jupyterlab/application-extension:logo \ + && /opt/unsloth-venv/bin/jupyter labextension lock @jupyterlab/application-extension:logo \ + && /opt/unsloth-venv/bin/jupyter labextension disable @jupyterlab/apputils-extension:splash \ + && /opt/unsloth-venv/bin/jupyter labextension lock @jupyterlab/apputils-extension:splash \ + && /opt/unsloth-venv/bin/jupyter labextension lock unsloth-jupyterlab +# Branding integrity guard: the attribution checker (a jupyter_server extension), +# the AGPLv3 license text, and its enabling config, into the base venv. --verify +# FAILS the build if any attribution / license asset is missing or altered. +COPY jupyter/unsloth_branding.py /tmp/unsloth-branding-guard/unsloth_branding.py +COPY jupyter/jupyter_server_config.d/unsloth_branding_guard.json /tmp/unsloth-branding-guard/unsloth_branding_guard.json +RUN SP="$(/opt/unsloth-venv/bin/python -c 'import sysconfig; print(sysconfig.get_path("purelib"))')" \ + && cp /tmp/unsloth-branding-guard/unsloth_branding.py "${SP}/unsloth_branding.py" \ + && mkdir -p /opt/unsloth-venv/etc/jupyter/jupyter_server_config.d \ + && cp /tmp/unsloth-branding-guard/unsloth_branding_guard.json \ + /opt/unsloth-venv/etc/jupyter/jupyter_server_config.d/unsloth_branding_guard.json \ + && cp "${UNSLOTH_STUDIO_HOME}/src/studio/LICENSE.AGPL-3.0" \ + /opt/unsloth-venv/share/jupyter/UNSLOTH_LICENSE.AGPL-3.0 \ + && rm -rf /tmp/unsloth-branding-guard \ + && /opt/unsloth-venv/bin/python -m unsloth_branding --verify +RUN chmod +x /usr/local/bin/unsloth-studio-launch \ + /usr/local/bin/unsloth-studio-update \ + /usr/local/bin/unsloth-llama-update \ + /usr/local/bin/unsloth-jupyter-tunnel + +# Studio, JupyterLab, sshd. All bind 0.0.0.0 in the container; publish with -p. +EXPOSE 8000 8888 22 + +# The base ENTRYPOINT (unsloth-entrypoint) still runs its GPU pre-flight +# first, then hands off to the service launcher. +CMD ["/usr/local/bin/unsloth-studio-launch"] diff --git a/docker/NOTICE b/docker/NOTICE new file mode 100644 index 0000000000..df23375f7a --- /dev/null +++ b/docker/NOTICE @@ -0,0 +1,41 @@ +Unsloth Docker Studio and JupyterLab image +========================================== + +This directory builds the Unsloth Docker Studio and JupyterLab image. The image +bundles Unsloth Studio, which is licensed under the GNU Affero General Public +License v3.0 (see /studio/LICENSE.AGPL-3.0). Unsloth Core is licensed under the +Apache License 2.0 (see /LICENSE). + + +Additional terms under AGPLv3 Section 7 +--------------------------------------- + +As permitted by Section 7(b) of the GNU Affero General Public License v3.0, and +in support of the "Appropriate Legal Notices" requirement for interactive user +interfaces, the following author attributions and legal notices are designated +as required Appropriate Legal Notices for this image. If you convey, modify, or +make the image (or any work based on it) available to users over a network, you +must keep these notices intact and displayed to those users: + + * The attribution "Built by the Unsloth team". + * The copyright line "Copyright 2026-Present the Unsloth team". + * The license notice "Licensed under Apache 2.0 and the GNU AGPLv3". + * The Unsloth logo and the "Unsloth Dark" theme shown in the JupyterLab top + bar and on the loading splash. + * The Help > About dialog, including the following links: + - Source: https://github.com/unslothai/unsloth + - Website: https://unsloth.ai + - License: https://github.com/unslothai/unsloth#license + - AGPLv3: https://www.gnu.org/licenses/agpl-3.0.html + - Apache: https://www.apache.org/licenses/LICENSE-2.0 + +These notices are displayed on the JupyterLab login page, the Help > About +dialog, the loading splash and the top bar. They are enforced at build time and +at runtime by docker/jupyter/unsloth_branding.py (see docker/jupyter/BRANDING.md +for details). Removing or altering them, whether by editing the build workflow, +the branding sources or the integrity guard, does not remove this license +condition. + +"Unsloth" and the Unsloth logo are trademarks of the Unsloth team. This NOTICE +governs copyright attribution under the AGPLv3 and does not grant any trademark +license. diff --git a/docker/build.sh b/docker/build.sh new file mode 100755 index 0000000000..ad295295b7 --- /dev/null +++ b/docker/build.sh @@ -0,0 +1,73 @@ +#!/usr/bin/env bash +# Build the unsloth-blackwell image on this B200 host (or any Linux host with Docker). +# The build host's GPU is NOT used -- nvcc cross-compiles for sm_100 + sm_120. +# +# Usage: +# ./build.sh # builds unsloth-blackwell:latest pinned to unsloth main +# TAG=2026.05.1 ./build.sh # custom tag +# UNSLOTH_REF=v2026.5.6 UNSLOTH_ZOO_REF=v2026.5.4 ./build.sh # pin git refs +set -euo pipefail + +cd "$(dirname "$0")" + +IMAGE_NAME="${IMAGE_NAME:-unsloth-blackwell}" +TAG="${TAG:-latest}" +CUDA_VERSION="${CUDA_VERSION:-12.8.1}" +UBUNTU_VERSION="${UBUNTU_VERSION:-24.04}" +PYTHON_VERSION="${PYTHON_VERSION:-3.12}" +UNSLOTH_REF="${UNSLOTH_REF:-main}" +UNSLOTH_ZOO_REF="${UNSLOTH_ZOO_REF:-main}" + +# llama.cpp prebuilt: default to the newest release, resolved here to a concrete +# tag so the build-arg changes only on a new release (correct layer caching). +# Pin for a frozen build: LLAMA_PREBUILT_TAG=b9596-mix-e6f2453 ./build.sh +resolve_latest_llama_tag() { + curl -fsSL -o /dev/null -w '%{url_effective}' \ + "https://github.com/unslothai/llama.cpp/releases/latest" 2>/dev/null \ + | sed -n 's#.*/releases/tag/##p' +} +if [ -z "${LLAMA_PREBUILT_TAG:-}" ]; then + LLAMA_PREBUILT_TAG="$(resolve_latest_llama_tag || true)" + if [ -n "$LLAMA_PREBUILT_TAG" ]; then + echo "Resolved latest llama.cpp release: ${LLAMA_PREBUILT_TAG}" + else + LLAMA_PREBUILT_TAG="latest" + echo "Could not resolve latest llama.cpp tag here; passing 'latest' (resolved inside the build)" + fi +fi + +echo "Building ${IMAGE_NAME}:${TAG}" +echo " CUDA ${CUDA_VERSION} Ubuntu ${UBUNTU_VERSION} Python ${PYTHON_VERSION}" +echo " unsloth @${UNSLOTH_REF}" +echo " unsloth-zoo @${UNSLOTH_ZOO_REF}" +echo " llama.cpp ${LLAMA_PREBUILT_TAG}" +# Read the arch list back out of the Dockerfile rather than repeating it: the +# hand-copied banner had already drifted, dropping 7.5 and so under-reporting +# Turing support to anyone reading this output. +# Bare filename: the script cd'd to its own directory above, so $0's dirname +# would be applied a second time and break every relative invocation. +ARCH_LIST="$(sed -n 's/^[[:space:]]*TORCH_CUDA_ARCH_LIST="\([^"]*\)".*/\1/p' \ + Dockerfile | head -n1)" +echo " arch list ${ARCH_LIST:-unknown}" +echo + +DOCKER_BUILDKIT=1 docker build \ + --progress=plain \ + --build-arg CUDA_VERSION="${CUDA_VERSION}" \ + --build-arg UBUNTU_VERSION="${UBUNTU_VERSION}" \ + --build-arg PYTHON_VERSION="${PYTHON_VERSION}" \ + --build-arg UNSLOTH_REF="${UNSLOTH_REF}" \ + --build-arg UNSLOTH_ZOO_REF="${UNSLOTH_ZOO_REF}" \ + --build-arg LLAMA_PREBUILT_TAG="${LLAMA_PREBUILT_TAG}" \ + -t "${IMAGE_NAME}:${TAG}" \ + . + +echo +echo "Built ${IMAGE_NAME}:${TAG}" +echo +echo "Smoke test on this host (B200, sm_100):" +echo " docker run --rm --gpus all ${IMAGE_NAME}:${TAG} python /workspace/smoke_test.py" +echo +echo "Smoke test on an RTX 5090 host (sm_120):" +echo " docker pull ${IMAGE_NAME}:${TAG} # or load .tar" +echo " docker run --rm --gpus all ${IMAGE_NAME}:${TAG} python /workspace/smoke_test.py" diff --git a/docker/entrypoint.sh b/docker/entrypoint.sh new file mode 100755 index 0000000000..ea430cb222 --- /dev/null +++ b/docker/entrypoint.sh @@ -0,0 +1,214 @@ +#!/usr/bin/env bash +# Container startup checks for Unsloth. Fails fast with actionable errors when the +# host GPU isn't reachable, catching the three modes behind ~95% of tickets: +# 1. nvidia-smi sees no GPU (missing --gpus all or nvidia-container-toolkit) +# 2. nvidia-smi works but torch.cuda.is_available() is False (driver too old) +# 3. GPU older than Ampere (sm < 80; Unsloth requires sm_80+) +# Bypass for offline tooling/docs/CI: docker run -e UNSLOTH_SKIP_GPU_CHECK=1 ... +set -euo pipefail + +# The image bakes CUDA 13 ptxas + NVRTC only for sm_103 (B300/GB300) and sm_121 +# (GB10/DGX Spark), which cu12.8 can't target. Both ship on >=580 drivers, which a +# cu13 cubin needs. Every other arch uses cu12.8 on the 570-579 floor, where a +# cu13 cubin can't load. Pick per DEVICE at boot: cu12.8 is the immutable default, +# only sm_103/sm_121 switch Triton to cu13 ptxas and retarget the NVRTC symlink. +# Best-effort: the default needs no write; only a non-root datacenter host can't switch. +select_cuda_jit_tools() { + local caps="" cc nvrtc_dir need_cu13=0 + if command -v nvidia-smi >/dev/null 2>&1; then + caps="$( { nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null || true; } )" + fi + # Scan EVERY visible GPU (a sm_103/sm_121 part can sit behind an H100). If ANY + # needs cu13, switch the whole process -- those hosts run >=580 drivers. + while IFS= read -r cc || [[ -n "${cc}" ]]; do + cc="$(printf '%s' "${cc}" | tr -d '[:space:]')" + case "${cc}" in + 10.3|12.1) need_cu13=1 ;; + esac + done <<< "${caps}" + # Non-datacenter / undetectable / CPU host: keep cu12.8 (needs no write). One + # exception: an earlier sm_103/sm_121 boot left libnvrtc.so.12 -> .cu13 that a + # 570-579 driver can't load -- reverse that (best-effort). + if [[ "${need_cu13}" -ne 1 ]]; then + for nvrtc_dir in \ + /opt/unsloth-venv/lib/python*/site-packages/nvidia/cuda_nvrtc/lib \ + "${UNSLOTH_STUDIO_HOME:-/opt/unsloth-studio}"/unsloth_studio/lib/python*/site-packages/nvidia/cuda_nvrtc/lib; do + [[ -e "${nvrtc_dir}/libnvrtc.so.12.cu128.orig" ]] || continue + [[ "$(readlink "${nvrtc_dir}/libnvrtc.so.12" 2>/dev/null)" == "libnvrtc.so.12.cu13" ]] || continue + ln -sf libnvrtc.so.12.cu128.orig "${nvrtc_dir}/libnvrtc.so.12" 2>/dev/null || true + done + return 0 + fi + # Blackwell datacenter present: point Triton at cu13 ptxas and retarget each + # venv's libnvrtc.so.12 -> the cu13 alias. -z guard lets an explicit + # TRITON_PTXAS_PATH win. Covers the base + Studio venvs. + if [[ -x /usr/local/cuda-13.0/bin/ptxas && -z "${TRITON_PTXAS_PATH:-}" ]]; then + export TRITON_PTXAS_PATH=/usr/local/cuda-13.0/bin/ptxas + fi + for nvrtc_dir in \ + /opt/unsloth-venv/lib/python*/site-packages/nvidia/cuda_nvrtc/lib \ + "${UNSLOTH_STUDIO_HOME:-/opt/unsloth-studio}"/unsloth_studio/lib/python*/site-packages/nvidia/cuda_nvrtc/lib; do + [[ -e "${nvrtc_dir}/libnvrtc.so.12.cu13" ]] || continue + ln -sf libnvrtc.so.12.cu13 "${nvrtc_dir}/libnvrtc.so.12" 2>/dev/null || true + done +} +# Best-effort: never let JIT-tool selection block container startup. +select_cuda_jit_tools || true + +# Make unslothai/notebooks available under /workspace before the user command. +# Best-effort, gated by UNSLOTH_SKIP_NOTEBOOK_SYNC, never blocks the container +# (see unsloth_sync_notebooks.sh). +sync_notebooks() { + if [[ -x /usr/local/bin/unsloth-sync-notebooks ]]; then + /usr/local/bin/unsloth-sync-notebooks || true + fi +} + +if [[ "${UNSLOTH_SKIP_GPU_CHECK:-0}" == "1" ]]; then + sync_notebooks + exec "$@" +fi + +err() { printf "\033[1;31mERROR:\033[0m %s\n" "$*" >&2; } +warn() { printf "\033[1;33mWARN:\033[0m %s\n" "$*" >&2; } + +# CPU mode for hosts that can't pass a GPU (Docker Desktop, CPU Linux, CI). Covers +# Jupyter, GGUF tooling, Studio chat; NOT training or loading a model. With +# UNSLOTH_ALLOW_CPU=1 a missing GPU warns instead of failing; a visible GPU still +# runs the checks below. +if [[ "${UNSLOTH_ALLOW_CPU:-0}" == "1" ]]; then + if ! command -v nvidia-smi >/dev/null 2>&1 || ! nvidia-smi -L 2>/dev/null | grep -q '^GPU'; then + warn "UNSLOTH_ALLOW_CPU=1 and no GPU visible -- continuing on CPU." + warn "CPU mode covers Jupyter, GGUF tooling and llama.cpp (GGUF) Studio chat." + warn "Training and loading Unsloth models (FastLanguageModel) still require an NVIDIA GPU." + sync_notebooks + exec "$@" + fi +fi + +# Check 1: nvidia-smi is injected by nvidia-container-toolkit on a GPU request, +# not baked in; a missing binary means "no GPU attached", same as an empty -L. +if ! command -v nvidia-smi >/dev/null 2>&1 || ! nvidia-smi -L 2>/dev/null | grep -q '^GPU'; then + err "No GPU visible inside the container." + cat >&2 <<'MSG' + +Likely causes (in order of frequency): + + 1. You started the container without --gpus all. + Re-launch with: + docker run --gpus all unsloth/unsloth:latest + Or use the bundled wrapper: + bash docker/run.sh + + 2. Host is missing nvidia-container-toolkit. + Install: https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html + Then: sudo systemctl restart docker + + 3. nvidia-container-toolkit is installed but the Docker daemon was not + restarted after install. Run: + sudo systemctl restart docker + + 4. You are using Podman / Kubernetes / a managed container service that + needs a different GPU flag than --gpus all. See the relevant docs: + podman: --device nvidia.com/gpu=all + k8s: nvidia.com/gpu resource request + GPU operator + + 5. This host has no NVIDIA GPU at all (Docker Desktop on macOS, Windows + without WSL2 GPU support, CPU-only Linux). Training and loading Unsloth + models need a GPU, but Jupyter, GGUF tooling and llama.cpp (GGUF) Studio + chat work on CPU: + docker run -e UNSLOTH_ALLOW_CPU=1 ... + +To bypass this check entirely (e.g. offline tooling), set UNSLOTH_SKIP_GPU_CHECK=1. +MSG + exit 1 +fi + +# Check 2: torch can use the GPU. Catches host-driver-too-old (nvidia-smi +# enumerates but CUDA contexts fail). +python - >&2 <<'PY' || exit 1 +import sys +import torch +if torch.cuda.is_available(): + sys.exit(0) +print("ERROR: torch.cuda.is_available() is False despite nvidia-smi working.") +print() +print("This image bakes in CUDA 12.8, so the host driver MUST be:") +print(" >= 570.26 (toolkit floor for cu128, applies to every GPU)") +print() +print("Two GPUs need an even newer driver because their launch driver was") +print("released after cu128's:") +print(" >= 580 B300 / GB300 (sm_103)") +print(" >= 580 GB10 / DGX Spark (sm_121)") +print() +print("Check the host (NOT the container) with: nvidia-smi") +print("Then upgrade the driver to match.") +sys.exit(1) +PY + +# Check 3: compute capability is supported. +python - >&2 <<'PY' || exit 1 +import sys +import torch +major, minor = torch.cuda.get_device_capability(0) +name = torch.cuda.get_device_name(0) +n = torch.cuda.device_count() +print(f"Unsloth container: {n} GPU(s). Primary: {name} sm_{major}{minor} bf16={torch.cuda.is_bf16_supported()}") + +# Image targets every current NVIDIA arch from Turing onward. +SUPPORTED = ( + ("sm_75", "Turing", "T4, RTX 20-series, Quadro RTX"), + ("sm_80", "Ampere DC", "A100, A30"), + ("sm_86", "Ampere", "A40, RTX A6000, RTX 30-series"), + ("sm_89", "Ada", "L4, L40, L40S, RTX 40-series"), + ("sm_90", "Hopper", "H100, H200, GH200"), + ("sm_100", "Blackwell DC", "B100, B200, GB200"), + ("sm_103", "Blackwell DC", "B300, GB300"), + ("sm_120", "Blackwell", "RTX 50-series, RTX PRO 6000 Blackwell"), + ("sm_121", "Blackwell", "GB10 (DGX Spark)"), +) +if major < 7 or (major == 7 and minor < 5): + print() + print(f"ERROR: Unsloth image requires Turing or newer (sm_75+). Got {name} sm_{major}{minor}.") + print() + print("Supported architectures in this image:") + for arch, fam, ex in SUPPORTED: + print(f" {arch:7s} {fam:13s} ({ex})") + sys.exit(1) +if major < 8: + print(f"NOTE: {name} is Turing (sm_{major}{minor}) -- bfloat16 is not supported.") + print(" Unsloth will fall back to fp16. Training works but is slightly slower.") + +# Secondary devices: all GPUs are exposed by default, so an unsupported later +# device only surfaces when a job pins to it. Device 0 is fatal above; +# secondaries warn now while excluding them is still cheap. +for d in range(1, n): + dmaj, dmin = torch.cuda.get_device_capability(d) + if dmaj < 7 or (dmaj == 7 and dmin < 5): + dname = torch.cuda.get_device_name(d) + print(f"WARNING: GPU {d} ({dname}, sm_{dmaj}{dmin}) is below this image's sm_75 floor.") + print(" Multi-GPU runs that include it, or jobs pinned to it, will fail;") + print(" exclude it with CUDA_VISIBLE_DEVICES or --gpus device=.") +PY + +# Upstream ships no CUDA 12 arm64 llama.cpp, so the arm64 image bakes cu13 while +# torch (cu128) runs on 570+. A cu13 cubin can't load on 570-579, so below 580 +# GGUF export / Studio chat fail even though training works -- warn up front. +if [ "$(uname -m)" = "aarch64" ]; then + _drv="$(nvidia-smi --query-gpu=driver_version --format=csv,noheader 2>/dev/null | head -1)" + _drv_major="${_drv%%.*}" + case "$_drv_major" in + *[!0-9]* | "") ;; # unreadable driver version -> no claim to make + *) + if [ "$_drv_major" -lt 580 ]; then + echo "WARNING: this arm64 image bakes a CUDA 13 llama.cpp (upstream ships no CUDA 12 arm64 build)." >&2 + echo " Host driver $_drv is < 580, which cannot load CUDA 13 binaries:" >&2 + echo " training (torch cu128) works, but GGUF export / Studio chat will fail" >&2 + echo " until the host driver is upgraded to >= 580." >&2 + fi + ;; + esac +fi + +sync_notebooks +exec "$@" diff --git a/docker/fetch_llama_prebuilt.py b/docker/fetch_llama_prebuilt.py new file mode 100644 index 0000000000..03a4af484c --- /dev/null +++ b/docker/fetch_llama_prebuilt.py @@ -0,0 +1,228 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +"""Bake a pinned llama.cpp prebuilt into the Docker image, deterministically. + +Why not studio/install_llama_prebuilt.py: that resolver selects a bundle for +the CURRENT host (nvidia-smi, /proc/driver/nvidia, installed CUDA runtime), +which is exactly what an image build must not do -- a B200 build host, a +GPU-less CI runner and a laptop must all produce byte-identical layers. This +script instead pins release + asset by build target only: + + amd64 -> app--linux-x64-cuda12-portable.tar.gz (sm_70..sm_120) + arm64 -> app--linux-arm64-cuda13-portable.tar.gz (sm_90..sm_121) + +The portable bundles carry their own CUDA runtime libs and dynamically load +the CUDA backend at runtime, so the same binaries also run CPU-only. + +Every download is sha256-verified against the release's own +llama-prebuilt-sha256.json. The converter (convert_hf_to_gguf.py) and its +gguf-py library are hydrated from the SAME release's source tarball so the +tensor mappings match the binaries -- the layout unsloth_zoo's +check_llama_cpp() expects: binaries, converter and gguf-py/ at the install +dir root. + +The tag may be the literal "latest" (or empty), in which case the newest +published release of RELEASE_REPO is resolved at build time by following the +/releases/latest redirect (no API token, no API rate limit). Pass a concrete +tag for a reproducible build. + +Usage (in the Dockerfile): + python fetch_llama_prebuilt.py +""" + +import hashlib +import json +import os +import shutil +import subprocess +import sys +import tarfile +import tempfile +import urllib.request + +RELEASE_REPO = "unslothai/llama.cpp" + + +def resolve_latest_tag(repo: str) -> str: + # Follow the /releases/latest redirect: no API token or rate limit. + url = f"https://github.com/{repo}/releases/latest" + request = urllib.request.Request(url, headers = {"User-Agent": "unsloth-docker-build"}) + with urllib.request.urlopen(request, timeout = 60) as response: + final_url = response.geturl() + marker = "/releases/tag/" + if marker not in final_url: + raise SystemExit( + f"FAIL: could not resolve latest release of {repo} (landed on {final_url})" + ) + return final_url.rsplit(marker, 1)[1].strip("/") + + +def fetch(url: str, dest: str) -> None: + request = urllib.request.Request(url, headers = {"User-Agent": "unsloth-docker-build"}) + with urllib.request.urlopen(request, timeout = 600) as response, open(dest, "wb") as f: + shutil.copyfileobj(response, f, length = 1 << 20) + + +def sha256_file(path: str) -> str: + digest = hashlib.sha256() + with open(path, "rb") as f: + for chunk in iter(lambda: f.read(1 << 20), b""): + digest.update(chunk) + return digest.hexdigest() + + +def fetch_verified(base_url: str, name: str, sums: dict, work: str) -> str: + path = os.path.join(work, name) + fetch(f"{base_url}/{name}", path) + expected = sums.get(name, {}).get("sha256") + if not expected: + raise SystemExit(f"FAIL: {name} not listed in llama-prebuilt-sha256.json") + actual = sha256_file(path) + if actual != expected: + raise SystemExit(f"FAIL: sha256 mismatch for {name}: expected {expected}, got {actual}") + print(f"verified {name} sha256={actual[:16]}...") + return path + + +def extracted_root(extract_dir: str) -> str: + children = os.listdir(extract_dir) + if len(children) == 1 and os.path.isdir(os.path.join(extract_dir, children[0])): + return os.path.join(extract_dir, children[0]) + return extract_dir + + +def main() -> None: + tag, target_arch, install_dir = sys.argv[1], sys.argv[2] or "amd64", sys.argv[3] + if tag in ("", "latest"): + tag = resolve_latest_tag(RELEASE_REPO) + print(f"resolved latest {RELEASE_REPO} release: {tag}") + base_url = f"https://github.com/{RELEASE_REPO}/releases/download/{tag}" + assets = { + "amd64": f"app-{tag}-linux-x64-cuda12-portable.tar.gz", + "arm64": f"app-{tag}-linux-arm64-cuda13-portable.tar.gz", + } + if target_arch not in assets: + raise SystemExit(f"FAIL: unsupported TARGETARCH={target_arch}") + bundle_name = assets[target_arch] + source_name = f"llama.cpp-source-{tag}.tar.gz" + + with tempfile.TemporaryDirectory() as work: + sha_path = os.path.join(work, "llama-prebuilt-sha256.json") + fetch(f"{base_url}/llama-prebuilt-sha256.json", sha_path) + sums = json.load(open(sha_path))["artifacts"] + + # Binaries: flat tarball, llama-quantize / llama-server / lib*.so at root. + bundle_path = fetch_verified(base_url, bundle_name, sums, work) + bundle_dir = os.path.join(work, "bundle") + os.makedirs(bundle_dir) + with tarfile.open(bundle_path) as tf: + tf.extractall(bundle_dir, filter = "tar") + os.makedirs(install_dir, exist_ok = True) + root = extracted_root(bundle_dir) + for entry in os.listdir(root): + target = os.path.join(install_dir, entry) + shutil.move(os.path.join(root, entry), target) + if os.path.isfile(target) and not entry.startswith("lib") and ".so" not in entry: + os.chmod(target, 0o755) + + # Converter + gguf-py from the same-tag source tarball so tensor mappings + # match the binaries (mirrors unsloth_zoo's _hydrate_converter_sources). + source_path = fetch_verified(base_url, source_name, sums, work) + source_dir = os.path.join(work, "source") + os.makedirs(source_dir) + with tarfile.open(source_path) as tf: + tf.extractall(source_dir, filter = "tar") + src_root = extracted_root(source_dir) + converter = os.path.join(src_root, "convert_hf_to_gguf.py") + gguf_py = os.path.join(src_root, "gguf-py") + if not (os.path.isfile(converter) and os.path.isdir(gguf_py)): + raise SystemExit(f"FAIL: source tarball for {tag} is missing converter files") + for script in os.listdir(src_root): + if script.startswith("convert_") and script.endswith(".py"): + shutil.copy2(os.path.join(src_root, script), os.path.join(install_dir, script)) + shutil.copytree(gguf_py, os.path.join(install_dir, "gguf-py"), dirs_exist_ok = True) + conversion = os.path.join(src_root, "conversion") + if os.path.isdir(conversion): + shutil.copytree(conversion, os.path.join(install_dir, "conversion"), dirs_exist_ok = True) + + # Make the baked marker readable by Studio's freshness check. The tarball keys + # off upstream_tag/source_repo, but the reader wants tag/release_tag/ + # published_repo (the install_llama_prebuilt.py schema). setdefault() leaves an + # already-populated tarball untouched; no timestamp, so layers stay identical. + marker_path = os.path.join(install_dir, "UNSLOTH_PREBUILT_INFO.json") + try: + with open(marker_path) as f: + marker = json.load(f) + except (OSError, ValueError): + marker = {} + marker.setdefault("tag", tag) + marker.setdefault("release_tag", tag) + marker.setdefault("published_repo", RELEASE_REPO) + with open(marker_path, "w") as f: + json.dump(marker, f, indent = 2) + f.write("\n") + print(f"marker augmented for freshness: tag={tag} published_repo={RELEASE_REPO}") + + # Mirror the install into build/bin/ via hardlinks (zero extra bytes) so + # Studio's setup.sh treats it as a complete local build and skips its + # source-build fallback (which would compile CPU-only llama.cpp over the baked + # CUDA bundle). Hardlinks keep $ORIGIN rpath and avoid a cycle when setup.sh + # relinks the root quantizer to build/bin/llama-quantize. + build_bin = os.path.join(install_dir, "build", "bin") + os.makedirs(build_bin, exist_ok = True) + for entry in os.listdir(install_dir): + source = os.path.join(install_dir, entry) + if os.path.isfile(source) and not os.path.islink(source): + try: + os.link(source, os.path.join(build_bin, entry)) + except OSError: + shutil.copy2(source, os.path.join(build_bin, entry)) + elif os.path.islink(source): + # Mirror same-dir soname symlinks (libllama.so.0 -> ...); without them + # a binary relinked into build/bin fails $ORIGIN (loader wants soname). + target = os.readlink(source) + dest = os.path.join(build_bin, entry) + if "/" not in target and not os.path.lexists(dest): + os.symlink(target, dest) + + # Sanity: the server must run on a GPU-less host (CUDA backend is a dlopen'd + # plugin). Check the quantizer from both roots: setup.sh relinks the root copy + # to build/bin, so build/bin must resolve standalone. + checks = ( + # llama-quantize has no --version: healthy run prints usage (rc 0), + # loader failure rc 127. + (os.path.join(install_dir, "llama-server"), "version"), + (os.path.join(install_dir, "llama-quantize"), "usage"), + (os.path.join(build_bin, "llama-quantize"), "usage"), + ) + for binary, expect in checks: + out = subprocess.run( + [binary, "--version"], + capture_output = True, + text = True, + timeout = 120, + ) + banner = (out.stdout + out.stderr).strip() + print( + os.path.relpath(binary, install_dir), + "->", + banner.splitlines()[0] if banner else "(no output)", + ) + if expect not in banner: + raise SystemExit( + f"FAIL: {binary} did not print '{expect}': rc={out.returncode}\n{banner[:400]}" + ) + for required in ( + "llama-quantize", + "convert_hf_to_gguf.py", + "gguf-py", + "UNSLOTH_PREBUILT_INFO.json", + ): + if not os.path.exists(os.path.join(install_dir, required)): + raise SystemExit(f"FAIL: {required} missing from {install_dir}") + print(f"OK: llama.cpp {tag} ({bundle_name}) installed at {install_dir}") + + +if __name__ == "__main__": + main() diff --git a/docker/jupyter/BRANDING.md b/docker/jupyter/BRANDING.md new file mode 100644 index 0000000000..6660575736 --- /dev/null +++ b/docker/jupyter/BRANDING.md @@ -0,0 +1,50 @@ +# Unsloth Docker Studio branding + +The Unsloth Docker Studio and JupyterLab image ships Unsloth attribution across +several files. Preserving it is a license condition, not just a build check. See +[../NOTICE](../NOTICE) and [/studio/LICENSE.AGPL-3.0](../../studio/LICENSE.AGPL-3.0). + +## What must stay + +- `Built by the Unsloth team` (login page and the labextension). +- `Copyright 2026-Present the Unsloth team`. +- `Licensed under Apache 2.0 and the GNU AGPLv3`. +- The Unsloth logo and the `Unsloth Dark` theme in the top bar and on the splash. +- The Help > About dialog with the Source, Website, License, AGPLv3 and Apache + links. + +The canonical strings live in `unsloth_branding.py` and its TypeScript mirror +`unsloth_labext/src/branding.ts`. The `PHRASE` literal must be byte-identical +between the two, because the guard greps the built labextension bundle for it. + +## Where it lives + +| File | Carries | +| --- | --- | +| `login.html` | JupyterLab login page and attribution line. | +| `unsloth_labext/src/branding.ts` | Canonical attribution strings (TS mirror). | +| `unsloth_labext/src/about.ts` | Help > About dialog and the license links. | +| `unsloth_labext/src/splash.ts` | Loading-splash caption. | +| `unsloth_labext/src/logo.ts` | Embedded Unsloth logo data URI. | +| `unsloth_branding.py` | Canonical strings and the integrity guard. | + +## How it is enforced + +`unsloth_branding.py` verifies the attribution is present and unaltered in three +places (see [../Dockerfile.studio](../Dockerfile.studio) and +[../studio_launch.sh](../studio_launch.sh)): + +1. **Build time:** `python -m unsloth_branding --verify` fails the image build if + any attribution asset is missing or altered. +2. **Whole image:** `studio_launch.sh` re-runs the same check before starting + supervisord; a failure refuses to start the container. +3. **JupyterLab:** the module is also a `jupyter_server` extension that re-checks + on load and refuses to serve JupyterLab if attribution was stripped after the + container started. + +The guard is a tripwire, not a lock. Anyone who forks the source controls the +build and can edit any of these files. It exists to make accidental removal fail +loudly and to make deliberate removal unambiguous. The attribution is protected +by the AGPLv3 as an Appropriate Legal Notice (see [../NOTICE](../NOTICE)), and +removing it before conveying or network-serving the image is a license +violation. diff --git a/docker/jupyter/favicon.ico b/docker/jupyter/favicon.ico new file mode 100644 index 0000000000..f922d8df9e Binary files /dev/null and b/docker/jupyter/favicon.ico differ diff --git a/docker/jupyter/install_sloth_stickers.py b/docker/jupyter/install_sloth_stickers.py new file mode 100644 index 0000000000..4ca9c54b4a --- /dev/null +++ b/docker/jupyter/install_sloth_stickers.py @@ -0,0 +1,79 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +"""Install the Unsloth Studio sloth stickers for the JupyterLab login screen. + +The branded login page (login.html) shows a different sloth sticker on each +visit, the same curated set Studio offers as profile avatars. The PNGs live in +the Studio frontend (`studio/frontend/public/Sloth emojis/`), which is present +in the studio image after install.sh runs. This copies the curated subset into +jupyter_server's static dir as `sloth/01.png .. sloth/20.png` so the template +can reference stable, space-free, auth-free URLs via `static_url(...)`. + +Usage: + install_sloth_stickers.py --src "" --dest "/sloth" + +Fail-soft: a missing source file is skipped (login.html's onerror falls back to +the Unsloth logo), and the script still exits 0 as long as at least one sticker +was installed. Stdlib only. +""" + +import argparse +import os +import shutil +import sys + +# Curated, in display order -> NN.png. Mirrors Studio's SLOTH_AVATARS: the square, +# low-whitespace stickers that frame cleanly. Synced by hand; missing names skipped. +CURATED = [ + "large sloth yay.png", + "large sloth heart.png", + "large sloth wave.png", + "large sloth thumbs.png", + "large sloth cheeky.png", + "large sloth glasses.png", + "large sloth fire.png", + "large sloth drink.png", + "large sloth sad.png", + "Large sloth Question mark.png", + "sloth shy large.png", + "sloth shock large.png", + "sloth sir large.png", + "sloth huglove large.png", + "sloth headphones.png", + "sloth pc square.png", + "sloth on phone.png", + "sloth magnify final.png", + "Sloth loca pc.png", + "UnSloth GPU Front square.png", +] + + +def main() -> int: + parser = argparse.ArgumentParser(description = __doc__) + parser.add_argument("--src", required = True, help = "Studio 'Sloth emojis' dir") + parser.add_argument("--dest", required = True, help = "output dir (static/sloth)") + args = parser.parse_args() + + os.makedirs(args.dest, exist_ok = True) + installed = 0 + for index, name in enumerate(CURATED, start = 1): + source = os.path.join(args.src, name) + target = os.path.join(args.dest, "%02d.png" % index) + if not os.path.isfile(source): + print(" skip (missing): %s" % name) + continue + try: + shutil.copyfile(source, target) + installed += 1 + except OSError as error: + print(" skip (%s): %s" % (error, name)) + + print("installed %d/%d sloth stickers into %s" % (installed, len(CURATED), args.dest)) + # Non-fatal, but an empty copy usually means a wrong --src, so signal it. + return 0 if installed else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/docker/jupyter/jupyter_server_config.d/unsloth_branding_guard.json b/docker/jupyter/jupyter_server_config.d/unsloth_branding_guard.json new file mode 100644 index 0000000000..592d6ad6de --- /dev/null +++ b/docker/jupyter/jupyter_server_config.d/unsloth_branding_guard.json @@ -0,0 +1,7 @@ +{ + "ServerApp": { + "jpserver_extensions": { + "unsloth_branding": true + } + } +} diff --git a/docker/jupyter/login.html b/docker/jupyter/login.html new file mode 100644 index 0000000000..ac029434f7 --- /dev/null +++ b/docker/jupyter/login.html @@ -0,0 +1,118 @@ +{# Unsloth-branded JupyterLab login page. Overwrites jupyter_server's default + login.html (same overwrite pattern as the favicon/logo). Extends the stock + page.html so favicon (already the Unsloth icon) and form plumbing stay intact; + we override the title, hide the stock header, and render a dark centered card + matching the "Unsloth Dark" (Monokai) theme. The card logo reads + static/logo/logo.png, which the image build replaces with the Unsloth logo. #} +{% extends "page.html" %} + +{% block title %}Unsloth{% endblock %} + +{% block stylesheet %} + +{% endblock %} + +{% block site %} +{# A different Unsloth Studio sloth sticker each visit (matches Studio's login). + The PNGs are copied into static/sloth/NN.png by the image build; if one is + missing the onerror handler falls back to the Unsloth logo so the page never + shows a broken image. #} +{% set sloths = [ + "01.png", "02.png", "03.png", "04.png", "05.png", "06.png", "07.png", + "08.png", "09.png", "10.png", "11.png", "12.png", "13.png", "14.png", + "15.png", "16.png", "17.png", "18.png", "19.png", "20.png" +] %} + +
+ Built by the Unsloth team. + Apache 2.0, AGPLv3 License Link
+ Copyright 2026-Present the Unsloth team.
+ github.com/unslothai/unsloth + · + unsloth.ai +
+{% endblock %} + +{% block script %}{% endblock %} diff --git a/docker/jupyter/logo.png b/docker/jupyter/logo.png new file mode 100644 index 0000000000..8fc411695d Binary files /dev/null and b/docker/jupyter/logo.png differ diff --git a/docker/jupyter/overrides.json b/docker/jupyter/overrides.json new file mode 100644 index 0000000000..4fe9590fdb --- /dev/null +++ b/docker/jupyter/overrides.json @@ -0,0 +1,40 @@ +{ + "@jupyterlab/apputils-extension:themes": { + "theme": "Unsloth Dark", + "theme-scrollbars": true, + "adaptive-theme": true, + "preferred-light-theme": "JupyterLab Light", + "preferred-dark-theme": "Unsloth Dark" + }, + "@jupyterlab/notebook-extension:tracker": { + "windowingMode": "none", + "scrollPastEnd": true, + "codeCellConfig": { + "autoClosingBrackets": true + } + }, + "@jupyterlab/cell-toolbar-extension:plugin": { + "toolbar": [ + { + "name": "run-cell-no-advance", + "command": "notebook:run-cell", + "icon": "ui-components:run", + "rank": 0 + } + ] + }, + "@jupyterlab/notebook-extension:panel": { + "toolbar": [ + { + "name": "restart-and-run", + "command": "notebook:restart-run-all", + "label": "Restart & Run All", + "rank": 33 + } + ] + }, + "@jupyterlab/apputils-extension:notification": { + "fetchNews": "false", + "checkForUpdates": false + } +} diff --git a/docker/jupyter/unsloth_branding.py b/docker/jupyter/unsloth_branding.py new file mode 100644 index 0000000000..ffbb6b3410 --- /dev/null +++ b/docker/jupyter/unsloth_branding.py @@ -0,0 +1,295 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 +"""Unsloth Docker Studio branding + AGPLv3 attribution integrity guard. + +This image is built by Unsloth and is licensed under the GNU AGPLv3. The +attribution (the Unsloth logo + theme, the Help > About dialog, the spinning +splash, the AGPLv3 notice and the source/website links) is shipped across +several independent files on purpose, so a reseller cannot white-label the image +with a shallow find-and-replace. This module is the canonical, plain-text source +of truth for those strings AND the checker that verifies they are still present. + +Everything here is plain readable text -- there are no base64/encoded/obfuscated +copies of the attribution (those would trip antivirus scanners and are pointless +for an open-source image). The single base64 blob in the build is the logo +*image* data URI in the labextension, which is an image, not hidden text. + +The guard runs in three places (see docker/Dockerfile.studio, docker/studio_launch.sh): + * build time -- `python -m unsloth_branding --verify` fails the image build + if any attribution asset is missing or altered. + * whole image -- studio_launch.sh runs the same check before launching + supervisord; a failure refuses to start the container. + * JupyterLab -- this module is also a jupyter_server extension; on load it + re-checks and refuses to serve JupyterLab if attribution was + stripped after the container started. +""" + +import json +import os +import sys + +# Canonical attribution strings. Plain text; keep in sync with the TS mirror +# unsloth_labext/src/branding.ts (the guard greps the built bundle for these). +PRODUCT = "Unsloth Docker Studio" +SHORT_LABEL = "Built by the Unsloth team" +# Loading-splash caption; distinct from SHORT_LABEL (see branding.ts). +SPLASH_LABEL = "Loading Unsloth Docker" +COPYRIGHT = "Copyright 2026-Present the Unsloth team" +AGPL_NOTICE = "Licensed under Apache 2.0 and the GNU AGPLv3" +WEBSITE_URL = "https://unsloth.ai" +DOCS_URL = "https://unsloth.ai/docs" +SOURCE_URL = "https://github.com/unslothai/unsloth" +LICENSE_URL = "https://github.com/unslothai/unsloth#license" +AGPL_URL = "https://www.gnu.org/licenses/agpl-3.0.html" +APACHE_URL = "https://www.apache.org/licenses/LICENSE-2.0" +# ONE plain literal, byte-identical to PHRASE in unsloth_labext/src/branding.ts; +# the guard greps the built bundle for it verbatim. +PHRASE = ( + "Unsloth Docker Studio and JupyterLab image. Built by the Unsloth team. " + "Licensed under Apache 2.0 and the GNU AGPLv3. " + "Source: https://github.com/unslothai/unsloth Website: https://unsloth.ai" +) + +THEME_NAME = "Unsloth Dark" +LABEXT_NAME = "unsloth-jupyterlab" +ABOUT_PLUGIN_ID = "unsloth-jupyterlab:about" +SPLASH_PLUGIN_ID = "unsloth-jupyterlab:splash" +# Prefix of the embedded logo image data URI in unsloth_labext/src/logo.ts. +# Removing the logo (a load-bearing ~19KB literal) breaks the top bar + splash. +LOGO_DATA_URI_PREFIX = "data:image/png;base64,iVBOR" + + +def resolve_paths( + venv_share = None, + jupyter_server_dir = None, + config_dirs = None, +): + """Resolve the installed locations of every checked branding asset. + + Defaults point at the live venv + the installed jupyter_server package. Tests + pass explicit roots so the checker can run against a staged temp tree. + """ + if venv_share is None: + venv_share = os.path.join(sys.prefix, "share", "jupyter") + if jupyter_server_dir is None: + import jupyter_server # local import: only needed for live resolution + jupyter_server_dir = os.path.dirname(jupyter_server.__file__) + labext_dir = os.path.join(venv_share, "labextensions", LABEXT_NAME) + + # Every page_config.json JupyterLab merges for disabledExtensions (app-settings + # + a labconfig/ file per config dir). Tests pass config_dirs=[] for hermeticity. + if config_dirs is None: + try: + from jupyter_core.paths import jupyter_config_path + config_dirs = jupyter_config_path() + except Exception: + config_dirs = [] + page_configs = [os.path.join(venv_share, "lab", "settings", "page_config.json")] + page_configs += [os.path.join(d, "labconfig", "page_config.json") for d in config_dirs] + + return { + "license": os.path.join(venv_share, "UNSLOTH_LICENSE.AGPL-3.0"), + "login": os.path.join(jupyter_server_dir, "templates", "login.html"), + "overrides": os.path.join(venv_share, "lab", "settings", "overrides.json"), + "labext_dir": labext_dir, + "labext_pkg": os.path.join(labext_dir, "package.json"), + "labext_static": os.path.join(labext_dir, "static"), + "favicon": os.path.join(jupyter_server_dir, "static", "favicons", "favicon.ico"), + "logo": os.path.join(jupyter_server_dir, "static", "logo", "logo.png"), + "page_configs": page_configs, + } + + +def _read(path): + try: + with open(path, encoding = "utf-8", errors = "replace") as f: + return f.read() + except OSError: + return None + + +def _nonempty_file(path): + try: + return os.path.getsize(path) > 0 + except OSError: + return False + + +def _bundle_text(static_dir): + """Concatenate every built .js chunk under the labextension static dir. + + The webpack production build splits the extension into several chunks but + keeps string literals verbatim (only identifiers are minified), so the + canonical attribution strings appear in one of these files. + """ + if not os.path.isdir(static_dir): + return "" + parts = [] + for name in sorted(os.listdir(static_dir)): + if name.endswith(".js"): + text = _read(os.path.join(static_dir, name)) + if text: + parts.append(text) + return "\n".join(parts) + + +def verify_branding(paths = None): + """Return a list of human-readable problems; empty list means all good.""" + if paths is None: + paths = resolve_paths() + problems = [] + + # 1. Full AGPLv3 license text shipped in the image. + license_text = _read(paths["license"]) + if license_text is None: + problems.append("missing AGPLv3 license file: " + paths["license"]) + elif "GNU AFFERO GENERAL PUBLIC LICENSE" not in license_text or "Version 3" not in license_text: + problems.append("AGPLv3 license file is not the GNU AGPL v3 text: " + paths["license"]) + + # 2. Branded login page carries the attribution + copyright + source link. + login = _read(paths["login"]) + if login is None: + problems.append("missing branded login page: " + paths["login"]) + else: + for marker in (SHORT_LABEL, COPYRIGHT, SOURCE_URL, "AGPLv3"): + if marker not in login: + problems.append("login page missing attribution marker: " + marker) + + # 3. The Unsloth Dark theme is the configured default. + overrides = _read(paths["overrides"]) + if not overrides or THEME_NAME not in overrides: + problems.append("overrides.json missing the '" + THEME_NAME + "' theme") + + # 4. The prebuilt labextension is installed and is ours. + pkg = _read(paths["labext_pkg"]) + if pkg is None: + problems.append("missing labextension: " + paths["labext_pkg"]) + else: + try: + if json.loads(pkg).get("name") != LABEXT_NAME: + problems.append("labextension package.json name is not " + LABEXT_NAME) + except ValueError: + problems.append("labextension package.json is not valid JSON") + + # 5. The built bundle still carries the visible attribution strings + plugins. + bundle = _bundle_text(paths["labext_static"]) + if not bundle: + problems.append("missing built labextension bundle: " + paths["labext_static"]) + else: + for marker in ( + PHRASE, + SHORT_LABEL, + COPYRIGHT, + AGPL_URL, + ABOUT_PLUGIN_ID, + SPLASH_PLUGIN_ID, + LOGO_DATA_URI_PREFIX, + ): + if marker not in bundle: + problems.append("labextension bundle missing: " + marker) + + # 6. Favicon + logo images present and non-empty. + if not _nonempty_file(paths["favicon"]): + problems.append("missing or empty favicon: " + paths["favicon"]) + if not _nonempty_file(paths["logo"]): + problems.append("missing or empty logo: " + paths["logo"]) + + # 7. No page_config.json disables the Unsloth extension or its plugins. + # disabledExtensions leaves the bundle on disk (check 5 passes) but strips + # it at load, so reject it. Only flag unsloth-jupyterlab ids. + for pc_path in paths.get("page_configs", []): + text = _read(pc_path) + if not text: + continue + try: + disabled = json.loads(text).get("disabledExtensions", {}) + except ValueError: + problems.append("page_config.json is not valid JSON: " + pc_path) + continue + # Modern JupyterLab uses a {id: bool} map; older configs used a list. + if isinstance(disabled, dict): + disabled_ids = [k for k, v in disabled.items() if v] + elif isinstance(disabled, (list, tuple)): + disabled_ids = list(disabled) + else: + disabled_ids = [] + for ident in disabled_ids: + if not isinstance(ident, str): + continue + if ident == LABEXT_NAME or ident.startswith(LABEXT_NAME + ":"): + problems.append( + "page_config.json disables Unsloth attribution '" + ident + "': " + pc_path + ) + + return problems + + +def banner(problems): + """A loud, plain-text failure banner naming what was stripped.""" + lines = [ + "", + "=" * 72, + "ERROR: Unsloth Docker Studio attribution / license integrity check failed.", + "", + "This image is built by Unsloth and ships under the GNU AGPLv3. It will not", + "start because required attribution or license assets are missing or altered:", + "", + ] + for p in problems: + lines.append(" - " + p) + lines += [ + "", + SHORT_LABEL + ". " + COPYRIGHT + ".", + "Website: " + WEBSITE_URL, + "Source: " + SOURCE_URL, + "License: GNU AGPLv3 (" + AGPL_URL + ")", + "=" * 72, + "", + ] + return "\n".join(lines) + + +# --- jupyter_server extension (Layer B: refuse to serve JupyterLab) ---------- +def _jupyter_server_extension_points(): + return [{"module": "unsloth_branding"}] + + +def _load_jupyter_server_extension(serverapp): + problems = verify_branding() + if not problems: + return + msg = banner(problems) + print(msg, file = sys.stderr, flush = True) + try: + serverapp.log.critical(msg) + except Exception: + pass + # Stop the server cleanly, then force exit if that's swallowed. Layer A + # (studio_launch.sh) refuses the container first; this backstops a direct run. + try: + serverapp.exit(1) + except Exception: + pass + raise SystemExit(1) + + +def main(argv = None): + import argparse + + parser = argparse.ArgumentParser(description = "Unsloth branding integrity check") + parser.add_argument("--verify", action = "store_true", help = "verify and exit nonzero on failure") + parser.add_argument("--venv-share", default = None) + parser.add_argument("--jupyter-server-dir", default = None) + args = parser.parse_args(argv) + + paths = resolve_paths(args.venv_share, args.jupyter_server_dir) + problems = verify_branding(paths) + if problems: + print(banner(problems), file = sys.stderr, flush = True) + return 1 + print("Unsloth branding integrity check passed (" + PRODUCT + ", AGPLv3).") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/docker/jupyter/unsloth_labext/.gitignore b/docker/jupyter/unsloth_labext/.gitignore new file mode 100644 index 0000000000..a51e4ca6ca --- /dev/null +++ b/docker/jupyter/unsloth_labext/.gitignore @@ -0,0 +1,7 @@ +node_modules/ +lib/ +*.tsbuildinfo +unsloth-jupyterlab/ +.yarn/ +.pnp.* +yarn.lock diff --git a/docker/jupyter/unsloth_labext/.yarnrc.yml b/docker/jupyter/unsloth_labext/.yarnrc.yml new file mode 100644 index 0000000000..3186f3f079 --- /dev/null +++ b/docker/jupyter/unsloth_labext/.yarnrc.yml @@ -0,0 +1 @@ +nodeLinker: node-modules diff --git a/docker/jupyter/unsloth_labext/package.json b/docker/jupyter/unsloth_labext/package.json new file mode 100644 index 0000000000..80284ed2bb --- /dev/null +++ b/docker/jupyter/unsloth_labext/package.json @@ -0,0 +1,54 @@ +{ + "name": "unsloth-jupyterlab", + "version": "0.1.0", + "description": "Unsloth Dark (Monokai) theme + Colab-style cell navigation for JupyterLab.", + "keywords": [ + "jupyter", + "jupyterlab", + "jupyterlab-extension", + "theme" + ], + "license": "AGPL-3.0-only", + "author": "Unsloth AI", + "private": true, + "main": "lib/index.js", + "types": "lib/index.d.ts", + "style": "style/index.css", + "files": [ + "lib/**/*.{d.ts,js,js.map}", + "style/**/*.{css,js,eot,gif,html,jpg,json,png,svg,woff2,ttf}", + "schema/*.json" + ], + "scripts": { + "build": "jlpm build:lib && jlpm build:labextension:dev", + "build:prod": "jlpm clean && jlpm build:lib:prod && jlpm build:labextension", + "build:lib": "tsc --sourceMap", + "build:lib:prod": "tsc", + "build:labextension": "jupyter labextension build .", + "build:labextension:dev": "jupyter labextension build --development True .", + "clean": "rimraf lib tsconfig.tsbuildinfo unsloth-jupyterlab/labextension" + }, + "dependencies": { + "@codemirror/state": "^6.0.0", + "@codemirror/view": "^6.0.0", + "@jupyterlab/application": "^4.5.0", + "@jupyterlab/apputils": "^4.5.0", + "@jupyterlab/cells": "^4.5.0", + "@jupyterlab/codemirror": "^4.5.0", + "@jupyterlab/mainmenu": "^4.5.0", + "@jupyterlab/notebook": "^4.5.0", + "@jupyterlab/theme-dark-extension": "^4.5.0", + "@lumino/disposable": "^2.0.0", + "@lumino/widgets": "^2.0.0" + }, + "devDependencies": { + "@jupyterlab/builder": "^4.5.0", + "rimraf": "^5.0.0", + "typescript": "~5.5.0" + }, + "jupyterlab": { + "extension": true, + "themePath": "style/index.css", + "outputDir": "unsloth-jupyterlab/labextension" + } +} diff --git a/docker/jupyter/unsloth_labext/src/about.ts b/docker/jupyter/unsloth_labext/src/about.ts new file mode 100644 index 0000000000..70bf028279 --- /dev/null +++ b/docker/jupyter/unsloth_labext/src/about.ts @@ -0,0 +1,95 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 +// +// "About Unsloth Docker Studio" command -> Help menu + command palette. Surfaces +// the AGPLv3 license, copyright and source/website links inside JupyterLab. + +import { + JupyterFrontEnd, + JupyterFrontEndPlugin +} from '@jupyterlab/application'; +import { Dialog, ICommandPalette, showDialog } from '@jupyterlab/apputils'; +import { IMainMenu } from '@jupyterlab/mainmenu'; +import { Widget } from '@lumino/widgets'; +import { UNSLOTH_LOGO_DATA_URI } from './logo'; +import { + AGPL_NOTICE, + AGPL_URL, + APACHE_URL, + COPYRIGHT, + DOCS_URL, + LICENSE_URL, + PHRASE, + PRODUCT, + SHORT_LABEL, + SOURCE_URL, + WEBSITE_URL +} from './branding'; + +const COMMAND_ID = 'unsloth:about'; + +/** + * Build the About dialog body from the trusted branding.ts constants only (no + * user input, so innerHTML has no injection surface). PHRASE is stamped as a data + * attribute so it's bundled verbatim for the integrity guard. + */ +function aboutBody(): Widget { + const body = new Widget(); + const el = body.node; + el.style.textAlign = 'center'; + el.style.padding = '4px 10px 10px'; + el.style.maxWidth = '430px'; + el.setAttribute('data-unsloth-attribution', PHRASE); + // Link rows in a left-aligned inline-block centered in the dialog, so the + // labels line up instead of each row centering independently. + el.innerHTML = ` + Unsloth +
${PRODUCT}
+
${SHORT_LABEL}
+
${AGPL_NOTICE}.
+
+ + +
Unsloth Reference: ${DOCS_URL}
+
Licenses
+
+
Unsloth Studio: AGPLv3
+
Unsloth Core: Apache 2.0
+
Unsloth license: ${LICENSE_URL}
+
+
+
${COPYRIGHT}
+ `; + return body; +} + +const aboutPlugin: JupyterFrontEndPlugin = { + id: 'unsloth-jupyterlab:about', + description: 'About Unsloth Docker Studio (AGPLv3 attribution).', + autoStart: true, + optional: [IMainMenu, ICommandPalette], + activate: ( + app: JupyterFrontEnd, + mainMenu: IMainMenu | null, + palette: ICommandPalette | null + ): void => { + app.commands.addCommand(COMMAND_ID, { + label: 'About ' + PRODUCT, + execute: () => + showDialog({ + title: 'About ' + PRODUCT, + body: aboutBody(), + buttons: [Dialog.okButton({ label: 'Close' })] + }) + }); + if (mainMenu) { + mainMenu.helpMenu.addGroup([{ command: COMMAND_ID }], 20); + } + if (palette) { + palette.addItem({ command: COMMAND_ID, category: 'Help' }); + } + } +}; + +export default aboutPlugin; diff --git a/docker/jupyter/unsloth_labext/src/branding.ts b/docker/jupyter/unsloth_labext/src/branding.ts new file mode 100644 index 0000000000..a17b1d108c --- /dev/null +++ b/docker/jupyter/unsloth_labext/src/branding.ts @@ -0,0 +1,24 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 +// +// Canonical attribution strings, mirrored from unsloth_branding.py. Imported by +// the About and splash plugins so they're bundled verbatim; the Python guard +// checks the built bundle still contains them. Plain text only, never encoded. + +export const PRODUCT = 'Unsloth Docker Studio'; +export const SHORT_LABEL = 'Built by the Unsloth team'; +// Loading-splash caption; distinct from SHORT_LABEL (says what's loading). +export const SPLASH_LABEL = 'Loading Unsloth Docker'; +export const COPYRIGHT = 'Copyright 2026-Present the Unsloth team'; +export const AGPL_NOTICE = 'Licensed under Apache 2.0 and the GNU AGPLv3'; +export const WEBSITE_URL = 'https://unsloth.ai'; +export const DOCS_URL = 'https://unsloth.ai/docs'; +export const SOURCE_URL = 'https://github.com/unslothai/unsloth'; +export const LICENSE_URL = 'https://github.com/unslothai/unsloth#license'; +export const AGPL_URL = 'https://www.gnu.org/licenses/agpl-3.0.html'; +export const APACHE_URL = 'https://www.apache.org/licenses/LICENSE-2.0'; + +// Must equal PHRASE in unsloth_branding.py (the guard greps the bundle for it). +// ONE plain literal, not a concatenation, so webpack keeps it contiguous. +export const PHRASE = + 'Unsloth Docker Studio and JupyterLab image. Built by the Unsloth team. Licensed under Apache 2.0 and the GNU AGPLv3. Source: https://github.com/unslothai/unsloth Website: https://unsloth.ai'; diff --git a/docker/jupyter/unsloth_labext/src/cellNav.ts b/docker/jupyter/unsloth_labext/src/cellNav.ts new file mode 100644 index 0000000000..a398b71fe4 --- /dev/null +++ b/docker/jupyter/unsloth_labext/src/cellNav.ts @@ -0,0 +1,138 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +import { + JupyterFrontEnd, + JupyterFrontEndPlugin +} from '@jupyterlab/application'; +import { CodeMirrorEditor } from '@jupyterlab/codemirror'; +import { INotebookTracker } from '@jupyterlab/notebook'; + +/** + * Colab-style cell navigation in BOTH command and edit mode. + * + * ArrowDown on a cell's last line (edit) or while selected (command) moves to the + * next cell and aligns its TOP to the viewport; ArrowUp mirrors it. JupyterLab + * centers tall cells, dropping the view mid-output. Settings can't fix this, so + * we listen in the CAPTURE phase, detect a cell boundary, and scroll-to-top. + */ +const cellNavPlugin: JupyterFrontEndPlugin = { + id: 'unsloth-jupyterlab:cell-nav', + description: + 'ArrowDown/ArrowUp move to the TOP of the next/previous cell (command + edit mode).', + autoStart: true, + requires: [INotebookTracker], + activate: (app: JupyterFrontEnd, tracker: INotebookTracker): void => { + const handler = (event: KeyboardEvent): void => { + if (event.key !== 'ArrowDown' && event.key !== 'ArrowUp') { + return; + } + if (event.ctrlKey || event.metaKey || event.altKey || event.shiftKey) { + return; + } + const panel = tracker.currentWidget; + if (!panel || !panel.isVisible) { + return; + } + if (!panel.node.contains(event.target as Node)) { + return; + } + // Never hijack arrows belonging to an interactive output (ipywidgets) or a + // form control; only the cell editor and command-mode cell nav. + const targetEl = event.target as HTMLElement | null; + if (targetEl) { + if (targetEl.closest('.jp-OutputArea')) { + return; + } + const tag = targetEl.tagName; + if (tag === 'INPUT' || tag === 'TEXTAREA' || tag === 'SELECT') { + return; + } + } + const notebook = panel.content; + const direction = event.key === 'ArrowDown' ? 1 : -1; + const editing = notebook.mode === 'edit'; + if (editing) { + const editor = notebook.activeCell?.editor; + if (!editor) { + return; + } + // While a completion popup is open the arrows belong to it; don't take + // over even at a cell boundary (common in one-line setup cells). + if ( + document.querySelector( + '.jp-Completer:not(.lm-mod-hidden), .cm-tooltip-autocomplete' + ) + ) { + return; + } + // Only take over at the cell boundary; else let CodeMirror move the + // cursor. `lineCount` counts LOGICAL lines, but JupyterLab wraps + // markdown and raw cell editors by default (StaticNotebook + // .defaultEditorConfig: markdown/raw lineWrap true), so the first and + // last logical line can own several visual rows -- the one-line markdown + // header every notebook opens with wraps to ~7. Ask CodeMirror whether + // it can still move one VISUAL line first, else those rows are + // unreachable: every arrow leaves the cell. + const view = editor instanceof CodeMirrorEditor ? editor.editor : null; + if (view) { + const range = view.state.selection.main; + const moved = view.moveVertically(range, direction === 1); + const from = view.coordsAtPos(range.head); + const to = + moved.head === range.head ? from : view.coordsAtPos(moved.head); + // moveVertically only returns the unchanged head at offset 0 / + // doc.length; elsewhere it clamps to the document edge, so a move that + // stays on the same visual row IS the editor edge and the cell + // boundary is the next stop. + if (from && to && Math.abs(to.top - from.top) > 1) { + return; + } + } else { + const line = editor.getCursorPosition().line; + if (direction === 1 && line !== editor.lineCount - 1) { + return; + } + if (direction === -1 && line !== 0) { + return; + } + } + } + const target = notebook.activeCellIndex + direction; + if (target < 0 || target >= notebook.widgets.length) { + return; + } + // We own this key: stop CodeMirror and Lumino from also handling it and + // re-triggering the centering scroll we replace. + event.preventDefault(); + event.stopPropagation(); + notebook.activeCellIndex = target; + const cell = notebook.activeCell; + const targetEditor = cell?.editor; + if (editing && cell && targetEditor) { + notebook.mode = 'edit'; + const lastLine = Math.max(0, targetEditor.lineCount - 1); + targetEditor.setCursorPosition({ + line: direction === 1 ? 0 : lastLine, + column: 0 + }); + } + if (cell) { + const node = cell.node; + // Defer so this runs AFTER JupyterLab's own ensureFocus/centering scroll + // and wins the last write. block:'start' puts the cell input at the top. + requestAnimationFrame(() => { + try { + node.scrollIntoView({ block: 'start' }); + } catch { + /* no-op */ + } + }); + } + }; + // Capture phase: decide before CodeMirror / Lumino consume the arrow keys. + document.addEventListener('keydown', handler, true); + } +}; + +export default cellNavPlugin; diff --git a/docker/jupyter/unsloth_labext/src/colabTitle.ts b/docker/jupyter/unsloth_labext/src/colabTitle.ts new file mode 100644 index 0000000000..4c589cc6a2 --- /dev/null +++ b/docker/jupyter/unsloth_labext/src/colabTitle.ts @@ -0,0 +1,153 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +import { + JupyterFrontEnd, + JupyterFrontEndPlugin +} from '@jupyterlab/application'; +import { INotebookTracker, NotebookPanel } from '@jupyterlab/notebook'; +import { Cell } from '@jupyterlab/cells'; + +/** + * Colab "#@title" form cells. A code cell whose first line is `#@title Some Title` + * renders in Colab as a titled, collapsed form. JupyterLab has no equivalent, so + * inject a clickable title bar and hide the input via a CSS class (not + * source_hidden, so metadata is never mutated). Clicking toggles the code. + */ + +const TITLE_RE = /^\s*#\s*@title\b[ \t]*(.*)$/; +const STYLE_ID = 'unsloth-colab-title-style'; + +function injectStyle(): void { + if (document.getElementById(STYLE_ID)) { + return; + } + const style = document.createElement('style'); + style.id = STYLE_ID; + style.textContent = ` +.unsloth-title-bar { + display: flex; + align-items: center; + gap: 8px; + cursor: pointer; + padding: 4px 8px; + /* Indent past the cell collapser + prompt gutter so the title aligns with the + cell's input/output content column instead of the far-left edge. */ + margin: 2px 0 2px var(--jp-cell-prompt-width, 64px); + user-select: none; + border-radius: 4px; + /* Heading-2-sized so a #@title form reads like a section heading (matches the + rendered-markdown h2 scale, --jp-content-font-size4); the caret inherits + this size so it grows too. */ + font-size: var(--jp-content-font-size4, 1.728em); + color: var(--jp-content-font-color1, inherit); +} +.unsloth-title-bar:hover { + background: var(--jp-layout-color2, rgba(128, 128, 128, 0.12)); +} +.unsloth-title-caret { + display: inline-block; + width: 1em; + line-height: 1; + opacity: 0.8; + transition: transform 0.12s ease; +} +.unsloth-title-bar.unsloth-collapsed .unsloth-title-caret { + transform: rotate(-90deg); +} +.unsloth-title-text { + font-weight: 700; + line-height: 1.25; +} +.jp-Cell.unsloth-code-collapsed > .jp-Cell-inputWrapper { + display: none; +} +`; + document.head.appendChild(style); +} + +function firstLineOf(cell: Cell): string { + try { + const raw = cell.model.toJSON().source as string | string[]; + const text = Array.isArray(raw) ? raw.join('') : String(raw || ''); + return text.split('\n', 1)[0] || ''; + } catch { + return ''; + } +} + +function applyTitle(cell: Cell): void { + let node: HTMLElement; + try { + node = cell.node; + } catch { + return; + } + if (cell.model?.type !== 'code') { + return; + } + const match = TITLE_RE.exec(firstLineOf(cell)); + let bar = node.querySelector(':scope > .unsloth-title-bar') as HTMLElement | null; + if (!match) { + if (bar) { + bar.remove(); + } + node.classList.remove('unsloth-titled', 'unsloth-code-collapsed'); + return; + } + // Drop trailing Colab form annotations, e.g. `{ display-mode: "form" }`. + const title = + (match[1] || '').replace(/\s*\{[^}]*\}\s*$/, '').trim() || 'Title'; + if (!bar) { + const barEl = document.createElement('div'); + barEl.className = 'unsloth-title-bar unsloth-collapsed'; + const caret = document.createElement('span'); + caret.className = 'unsloth-title-caret'; + caret.textContent = '▾'; + const text = document.createElement('span'); + text.className = 'unsloth-title-text'; + barEl.appendChild(caret); + barEl.appendChild(text); + barEl.addEventListener('click', () => { + const collapsed = node.classList.toggle('unsloth-code-collapsed'); + barEl.classList.toggle('unsloth-collapsed', collapsed); + }); + node.insertBefore(barEl, node.firstChild); + // Collapsed by default the first time we decorate this cell (Colab default). + node.classList.add('unsloth-code-collapsed'); + bar = barEl; + } + const label = bar.querySelector('.unsloth-title-text') as HTMLElement | null; + if (label) { + label.textContent = title; + } + node.classList.add('unsloth-titled'); +} + +const colabTitlePlugin: JupyterFrontEndPlugin = { + id: 'unsloth-jupyterlab:colab-title', + description: 'Render Colab #@title code cells as collapsed, titled forms.', + autoStart: true, + requires: [INotebookTracker], + activate: (app: JupyterFrontEnd, tracker: INotebookTracker): void => { + injectStyle(); + const decorate = (panel: NotebookPanel): void => { + const scan = (): void => { + panel.content.widgets.forEach(applyTitle); + }; + panel.revealed.then(scan).catch(() => undefined); + // Re-scan on cell add/remove/move or active-cell switch (covers editing a + // #@title line). applyTitle never re-collapses an existing bar, so manual + // expansions are preserved. + const model = panel.content.model; + if (model) { + model.cells.changed.connect(() => window.setTimeout(scan, 0)); + } + panel.content.activeCellChanged.connect(() => window.setTimeout(scan, 0)); + }; + tracker.widgetAdded.connect((_, panel) => decorate(panel)); + tracker.forEach(decorate); + } +}; + +export default colabTitlePlugin; diff --git a/docker/jupyter/unsloth_labext/src/index.ts b/docker/jupyter/unsloth_labext/src/index.ts new file mode 100644 index 0000000000..75a493c389 --- /dev/null +++ b/docker/jupyter/unsloth_labext/src/index.ts @@ -0,0 +1,77 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +import { + ILabShell, + JupyterFrontEnd, + JupyterFrontEndPlugin +} from '@jupyterlab/application'; +import { IThemeManager } from '@jupyterlab/apputils'; +import { Widget } from '@lumino/widgets'; +import { UNSLOTH_LOGO_DATA_URI } from './logo'; +import aboutPlugin from './about'; +import cellNavPlugin from './cellNav'; +import colabTitlePlugin from './colabTitle'; +import outputSelectPlugin from './outputSelect'; +import splashPlugin from './splash'; +import uiChromePlugin from './uiChrome'; + +/** + * The "Unsloth Dark" theme: JupyterLab Dark repainted with the Monokai palette + * (style/variables.css). A named theme so it appears in Settings > Theme and + * works with the adaptive light/dark switch in overrides.json. + */ +const themePlugin: JupyterFrontEndPlugin = { + id: 'unsloth-jupyterlab:theme', + description: 'Unsloth Dark (Monokai) theme.', + autoStart: true, + requires: [IThemeManager], + activate: (app: JupyterFrontEnd, manager: IThemeManager): void => { + const style = 'unsloth-jupyterlab/index.css'; + manager.register({ + name: 'Unsloth Dark', + isLight: false, + themeScrollbars: true, + load: () => manager.loadCSS(style), + unload: () => Promise.resolve(undefined) + }); + } +}; + +/** + * Replace the top-left Jupyter logo with the Unsloth logo. The stock logo plugin + * is disabled + locked at build, so this is the only logo widget. An with + * inline styles (not a LabIcon) so branding shows in any theme. + */ +const logoPlugin: JupyterFrontEndPlugin = { + id: 'unsloth-jupyterlab:logo', + description: 'Replace the top-left Jupyter logo with the Unsloth logo.', + autoStart: true, + requires: [ILabShell], + activate: (app: JupyterFrontEnd, shell: ILabShell): void => { + const logo = new Widget(); + const img = document.createElement('img'); + img.src = UNSLOTH_LOGO_DATA_URI; + img.alt = 'Unsloth'; + img.style.height = '24px'; + img.style.width = 'auto'; + img.style.margin = '1px 6px 1px 8px'; + img.style.display = 'block'; + logo.node.appendChild(img); + logo.node.style.display = 'flex'; + logo.node.style.alignItems = 'center'; + logo.id = 'jp-MainLogo'; + shell.add(logo, 'top', { rank: 0 }); + } +}; + +export default [ + themePlugin, + cellNavPlugin, + logoPlugin, + colabTitlePlugin, + outputSelectPlugin, + uiChromePlugin, + aboutPlugin, + splashPlugin +]; diff --git a/docker/jupyter/unsloth_labext/src/logo.ts b/docker/jupyter/unsloth_labext/src/logo.ts new file mode 100644 index 0000000000..3a6e83a3ae --- /dev/null +++ b/docker/jupyter/unsloth_labext/src/logo.ts @@ -0,0 +1,7 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +// Auto-generated: Unsloth circle logo as a base64 PNG data URI, embedded so the +// logo plugin has no runtime asset dependency and renders in any theme. +export const UNSLOTH_LOGO_DATA_URI = + "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAGQAAABhCAYAAAAgLwTnAAAKMWlDQ1BJQ0MgUHJvZmlsZQAAeJydlndUU9kWh8+9N71QkhCKlNBraFICSA29SJEuKjEJEErAkAAiNkRUcERRkaYIMijggKNDkbEiioUBUbHrBBlE1HFwFBuWSWStGd+8ee/Nm98f935rn73P3Wfvfda6AJD8gwXCTFgJgAyhWBTh58WIjYtnYAcBDPAAA2wA4HCzs0IW+EYCmQJ82IxsmRP4F726DiD5+yrTP4zBAP+flLlZIjEAUJiM5/L42VwZF8k4PVecJbdPyZi2NE3OMErOIlmCMlaTc/IsW3z2mWUPOfMyhDwZy3PO4mXw5Nwn4405Er6MkWAZF+cI+LkyviZjg3RJhkDGb+SxGXxONgAoktwu5nNTZGwtY5IoMoIt43kA4EjJX/DSL1jMzxPLD8XOzFouEiSniBkmXFOGjZMTi+HPz03ni8XMMA43jSPiMdiZGVkc4XIAZs/8WRR5bRmyIjvYODk4MG0tbb4o1H9d/JuS93aWXoR/7hlEH/jD9ld+mQ0AsKZltdn6h21pFQBd6wFQu/2HzWAvAIqyvnUOfXEeunxeUsTiLGcrq9zcXEsBn2spL+jv+p8Of0NffM9Svt3v5WF485M4knQxQ143bmZ6pkTEyM7icPkM5p+H+B8H/nUeFhH8JL6IL5RFRMumTCBMlrVbyBOIBZlChkD4n5r4D8P+pNm5lona+BHQllgCpSEaQH4eACgqESAJe2Qr0O99C8ZHA/nNi9GZmJ37z4L+fVe4TP7IFiR/jmNHRDK4ElHO7Jr8WgI0IABFQAPqQBvoAxPABLbAEbgAD+ADAkEoiARxYDHgghSQAUQgFxSAtaAYlIKtYCeoBnWgETSDNnAYdIFj4DQ4By6By2AE3AFSMA6egCnwCsxAEISFyBAVUod0IEPIHLKFWJAb5AMFQxFQHJQIJUNCSAIVQOugUqgcqobqoWboW+godBq6AA1Dt6BRaBL6FXoHIzAJpsFasBFsBbNgTzgIjoQXwcnwMjgfLoK3wJVwA3wQ7oRPw5fgEVgKP4GnEYAQETqiizARFsJGQpF4JAkRIauQEqQCaUDakB6kH7mKSJGnyFsUBkVFMVBMlAvKHxWF4qKWoVahNqOqUQdQnag+1FXUKGoK9RFNRmuizdHO6AB0LDoZnYsuRlegm9Ad6LPoEfQ4+hUGg6FjjDGOGH9MHCYVswKzGbMb0445hRnGjGGmsVisOtYc64oNxXKwYmwxtgp7EHsSewU7jn2DI+J0cLY4X1w8TogrxFXgWnAncFdwE7gZvBLeEO+MD8Xz8MvxZfhGfA9+CD+OnyEoE4wJroRIQiphLaGS0EY4S7hLeEEkEvWITsRwooC4hlhJPEQ8TxwlviVRSGYkNimBJCFtIe0nnSLdIr0gk8lGZA9yPFlM3kJuJp8h3ye/UaAqWCoEKPAUVivUKHQqXFF4pohXNFT0VFysmK9YoXhEcUjxqRJeyUiJrcRRWqVUo3RU6YbStDJV2UY5VDlDebNyi/IF5UcULMWI4kPhUYoo+yhnKGNUhKpPZVO51HXURupZ6jgNQzOmBdBSaaW0b2iDtCkVioqdSrRKnkqNynEVKR2hG9ED6On0Mvph+nX6O1UtVU9Vvuom1TbVK6qv1eaoeajx1UrU2tVG1N6pM9R91NPUt6l3qd/TQGmYaYRr5Grs0Tir8XQObY7LHO6ckjmH59zWhDXNNCM0V2ju0xzQnNbS1vLTytKq0jqj9VSbru2hnaq9Q/uE9qQOVcdNR6CzQ+ekzmOGCsOTkc6oZPQxpnQ1df11Jbr1uoO6M3rGelF6hXrtevf0Cfos/ST9Hfq9+lMGOgYhBgUGrQa3DfGGLMMUw12G/YavjYyNYow2GHUZPTJWMw4wzjduNb5rQjZxN1lm0mByzRRjyjJNM91tetkMNrM3SzGrMRsyh80dzAXmu82HLdAWThZCiwaLG0wS05OZw2xljlrSLYMtCy27LJ9ZGVjFW22z6rf6aG1vnW7daH3HhmITaFNo02Pzq62ZLde2xvbaXPJc37mr53bPfW5nbse322N3055qH2K/wb7X/oODo4PIoc1h0tHAMdGx1vEGi8YKY21mnXdCO3k5rXY65vTW2cFZ7HzY+RcXpkuaS4vLo3nG8/jzGueNueq5clzrXaVuDLdEt71uUnddd457g/sDD30PnkeTx4SnqWeq50HPZ17WXiKvDq/XbGf2SvYpb8Tbz7vEe9CH4hPlU+1z31fPN9m31XfKz95vhd8pf7R/kP82/xsBWgHcgOaAqUDHwJWBfUGkoAVB1UEPgs2CRcE9IXBIYMj2kLvzDecL53eFgtCA0O2h98KMw5aFfR+OCQ8Lrwl/GGETURDRv4C6YMmClgWvIr0iyyLvRJlESaJ6oxWjE6Kbo1/HeMeUx0hjrWJXxl6K04gTxHXHY+Oj45vipxf6LNy5cDzBPqE44foi40V5iy4s1licvvj4EsUlnCVHEtGJMYktie85oZwGzvTSgKW1S6e4bO4u7hOeB28Hb5Lvyi/nTyS5JpUnPUp2Td6ePJninlKR8lTAFlQLnqf6p9alvk4LTduf9ik9Jr09A5eRmHFUSBGmCfsytTPzMoezzLOKs6TLnJftXDYlChI1ZUPZi7K7xTTZz9SAxESyXjKa45ZTk/MmNzr3SJ5ynjBvYLnZ8k3LJ/J9879egVrBXdFboFuwtmB0pefK+lXQqqWrelfrry5aPb7Gb82BtYS1aWt/KLQuLC98uS5mXU+RVtGaorH1futbixWKRcU3NrhsqNuI2ijYOLhp7qaqTR9LeCUXS61LK0rfb+ZuvviVzVeVX33akrRlsMyhbM9WzFbh1uvb3LcdKFcuzy8f2x6yvXMHY0fJjpc7l+y8UGFXUbeLsEuyS1oZXNldZVC1tep9dUr1SI1XTXutZu2m2te7ebuv7PHY01anVVda926vYO/Ner/6zgajhop9mH05+x42Rjf2f836urlJo6m06cN+4X7pgYgDfc2Ozc0tmi1lrXCrpHXyYMLBy994f9Pdxmyrb6e3lx4ChySHHn+b+O31w0GHe4+wjrR9Z/hdbQe1o6QT6lzeOdWV0iXtjusePhp4tLfHpafje8vv9x/TPVZzXOV42QnCiaITn07mn5w+lXXq6enk02O9S3rvnIk9c60vvG/wbNDZ8+d8z53p9+w/ed71/LELzheOXmRd7LrkcKlzwH6g4wf7HzoGHQY7hxyHui87Xe4Znjd84or7ldNXva+euxZw7dLI/JHh61HXb95IuCG9ybv56Fb6ree3c27P3FlzF3235J7SvYr7mvcbfjT9sV3qID0+6j068GDBgztj3LEnP2X/9H686CH5YcWEzkTzI9tHxyZ9Jy8/Xvh4/EnWk5mnxT8r/1z7zOTZd794/DIwFTs1/lz0/NOvm1+ov9j/0u5l73TY9P1XGa9mXpe8UX9z4C3rbf+7mHcTM7nvse8rP5h+6PkY9PHup4xPn34D94Tz+6TMXDkAAC1vSURBVHic1Z15gBTlmf8/VdX3PRczwHCIqCCCCgjihQY1SoyoURE12cQrRvHKiq5Zo25cNUFdzU9jsiasyXoCHuCBIhgUATmUS24BOWaAYY6evs+q9/dHdVV3z/RAz4gm+8W2q956662q59vP8T7v+9ZIE5e/RE+gCYEqNASgCoGW29YAhEAIoVeUQACy0HckQJL0LcksBxkJOVcuyzIKkiQjnWCxW08QmjYATRwnSXI9EjUCqkA4hUDSrwYCIWmItNBoE4gWTWj7NVXdrvg8O9vWbd3avGTN6tiufYnmT75AS6V79MzfBSz/0KtLhjQFQkJGUkYrsnyOIsujZEk+TZGkvqgCWZYRCoBOtMF1jgj9nwANgYA6DYGQQFgUUuEI3uGDqT5vbCgVjn4Wb2ha0/zp6k+jW3cvCS5dF0nuOfCPe/4SkP5RGoIQSJKEgny2LEk/sirKlRZJ6aXIMrKUO1Xo50Je+Ln/iogQCDTzW+j3ITS03H1qCIQmEIqE7HVhqalAFVo2ebB1bmjlpplNsxbMCy36IiZUrUeyOJL4TgkROZWQJSoVWbndKsm3WGSlxqLIKEjISJ2uI3L/FyK/rYGuKYWaIYRJgEAnQRPCJCT/Te5eBcJmxVJbAVZLJrL+q5faF6x8qu31RV9mGpt7JJMjge+MEDSBJMlDLZL0nxZFucwqK1gkGRkJRTJ0pxiiiAhMQRdrhVakHSY5BQR1TY5A0zSEBHKVF7nST6q1fU34w5UPtf1pztuZr/f3SDbfBN86IRoCCWm4IsmP2mT5IkuOCEWSkKWutKK0nzAEX6wVhWbqMMIvql+inqaB04albw3ZWGJzeOZH97T/ae67Wmu4RzLqCb41QjShAfgVSX7KKis/s0gKFllCMcjIESFRrp+g6NdvEKSR9xXiMCQUEZLTIBVNJ9g0gTliXHaUAbVk9rV8Fpkx7+exlxZ8Kb6D6OxbIUQTGpKQfmZV5L/YJEW2SDKKLOdC25xjp5iIQj+hi6bAT2D8+jGFLjoIWiBy91Hw6899q4AQWpEfMeoXBgJm20YdVUPyu1EG1ZFatvFPoV/95dbstr3fquc/goQI/dcqqLBI8rtWWT7NIilYJAk5pxUdiYAu/IQhoBLCKvQTapGp0kyHrXUqE8UEdNCSQiKKfhS5OkITWOprkNyOcPSp2ZOiz839uIfyPizkI9WQQCAJ6Yd2WTnoUJTT7LIFmyxjkWQsUt5ldxS8ViAAXYAdTEyOeDVHgPHJmtsaqtBy+1oXZRpZoaGSb0czvzVU8ppTbBYBSUJSZNSGZtTGFp932lWLKmbc859ywHukRFeEI6IhGaEhwX/ZJOWuQvOkSKXD2M5+Im+eOvcn8r9eVYiCCKprZ50VgoymkhEqqhB6NCdJqJpWEDJ3jtoQ+fspiVykYRncF6059GnwF/91fmb9zmSPBNgFvjEhqhDISAsssnyuVVaK/ERHdDRPeTORF75JQoF5MYgotPmFBAAEM0n2JyOk1Cw22YLf6sCuKAigNRUnmYqBYqHW6ccmK2SEmv9RFPiqspBVUfpWI3mc7aFfPndm4t3PNvRIiCXQ49RJLuNRbZHkj62yMkyR9AjK8BNGHdAFj+hgrgp+mSIXLRnRTsdfv+F3MkIloWZJqBlTO1vScZJqluO9NVzXfyT1Dh9DvNUc5QrgUmxYJZndiXZWBfexLnyAWQ3rsSlWKuwuMkLNa0e3pKag7mtFrvQF/E/dulaurZgYmzHvw57KshA90hABaJqolyVphVWW+xj9ic5+omOfokM0U0CEGSUhkJDIaBoxNU1bOkEkm0ICAlYHFTYXFVYHEhBR04z29+GCXoM5rbIfAavjsPe+oHkHly5/BU2S8Fjs3SejEKqG5HVhGVhHZPqrk6N/mDOr543p6JmGCPpaZXmlRZJ7S2WHsR172MV+wthvScdpTcdxW2z0tnsYW9uXfk4/g1wVDPFU08fhpcLmRJYkRM4/lEIoFKKhoYFYLEbv3r3p168fAOfVHM0TJ3yfX6yZi8tiAyhpXsuCIiOicbLbG/Hed/VMFJno/3vzG5FiUbp4oFLICb2PJPGFIsm1Hc3TYfNOFAo/71zTQmV3vJ2spnGCrxc/7nciF/YazEBXBd6c0ErfUF6Ue/bs4Y033mDOnDns2bOHgwcPEo/HzeNjxoxh4cKFeL1eLqo7jgdcfmJqBqdiRfANSJFlRDJNdstevPdOmSniqVTsL+/N7WlzllKpi0NAkSU+kJFqjYKOROTLiqMm0SGkRIIsGl/HgmSExmW9h3Jt/QjGVNR3eUfJZJLdu3fT0NBAMBjE6XTS0NDACy+8wIoVKw554ytXrmTTpk2MHTsWm2IhnIyQSkaxBXrjlBUymtpzYmQJkcqQ3daA975r5mgtodMTc5Ys60lT0iUrXi2/siTNk+DCwjJRkojSYazhqBVJZn8yQnM6zg9rj+Ouo8dyvLdXp+t99dVXLFiwgHfeeYcdO3bQ1NREONz9vFIgEOCNN97ge9/7nlm2qGUnT+xYzrzGLwGo81TnIrlv0BFXNeQqH5LH2RT8l8dOTK/+qqm7TUiXrXyt3LrTgWnGTsf+BOQ1oDDvlE+LayBJZIXGpnAzg92V/O74czmzakDRRXbt2sVTTz1lmp4jgdNOO425c+dSXV3d6diGcBPXrZ3LqsaNOD1V+O1uspra84tlVZT+vdBaw6taL39ojNbWvR9QuYRMAubAIYgoCmPzqW8hBCq68z2YirE3EeKGAaN46LjxOBWreYEVK1Zw//33s3Dhwm49QHdQXV3NJZdcwl133cXxxx9fdOylhvVMXTOHUCJElb8uN37TswhMZFSswwaS+njt88Hrp/+8O+eWQ0g/YCPgLW2eOvoJzfQXhpmySgrbo604FStPD7+A82qONhtvaGjg9ttv56233urOfX9jTJw4kccff7yImIOpGJNXvsbHjRvw+nrhVCy6v+sJhMA6bCCR/5p9TfT3b7xS7mnlhFjPCPAW5p30jptWkA/K55w0QVHuSUFmc6SZOoeH98ddW0TGM888Q79+/b5zMgDmzZvHsGHDuOmmm8yxm152N4vOvJ7bh51HJNJMOJNC7kYU2hHZXU24r5/4rG34oM4OsgvIZh+h9L+bBWJSx9xPYeJPMxJ0hck+9DKLJLMmtJ8T/XV8csZ19Hf6AYhEIpx++uncfvvtPX7YI4U///nP+Hw+lixZYpb9fvhEnjnlSpKxNkLpRM9IkSREJA6yXOH592v/WO5ph7pSQAj+szgbmzdDqgA117vWv/VsalZoZDWdjHXhA4yvHsg7Y6dgzT3UmjVr6N27N8uW9Sgq/FYQjUY588wzuf/++82yqYNO5YVxPyaVCBLJJJFKJEoPC0VG3dOEbezQy9zXTZyUKz1kQ7I+XlDy85iAqqK0eMe0taZ1SHkLskLDKiusCzdxsr83b54y2UyqvPfee4wcOZJYLNb9h/sO8Mgjj3DllVea+z/tfzJ/OmUyiVgbKTVbcty/HKj7WnBfd+Fvlb7VNoqHgzqhKw0ZB+LmPBEUmSTTRKGbp2wBGYok81W0lX5OP6+Outxs8O233+aiiy7q0QN9l5g9ezYXX3yxuf/zo8byr8efSzhyENETPiQJEYwi960e4vrx+bcYpV1Vl83MKgXDl2h3qQYRBeZIM82TyA/2aPkBICEEbek4WaHxpxMvosrmBOCDDz5g0qRJXd3DPx3eeecdJk+ebO4/MXwi5/cdQUu4CUVWut+gIqPubsI56fQ7LYP7ujmElsjmMGV+8Ge0Jrgi77y1ghG2gtE3rXiULpsbs94abeWB48YzJtAXgJ07d3L55Zd3df1/WsyaNYt7773X3H/plCvp6+1FU7wduQf+REQTyL0CA1xTvndrrqhkI52iLE2IuzsOnZYaRs2SJ8swVZsizVzaeyi/GHgKANlslu9///tH1Gc4HA5qa2sPX/EIYPr06cyapSdva+xunjvxh6BmSGta97uMioza2IL9vNHXK/U1FrrQElnTBAWffiriR4YmmARQoA2mI89FVLl6B1NRau1unhh2vtn41KlT2b59ew/FUQyv18tVV11Fe3s711133RFpsxzccMMNHDigz/+9uPdQLqsfTnu0ucu0/6EgInEsfWuOdf3orGtzRZ20xFLYExUSPxUCS8dkYOEUHGMQyRhkUnPa8XUsyNPDLzT7Gq+//jr//d//3e2b9vv9jBs3jgEDBlBfX09dXR3Dhw9nzJgxZuj5/vvvl9VW//79mTRpEn/84x/JZrPdvhfQ+0w33ngj77zzDgDTh32fT1p20Z5J4rHYuhV3SZKM1hzEMWHUT2L/8/5ftUi8k5ZI5y79m74BCNimIY4pHM0zptEUT0DTzBkaAA2JMEM81Sw94zokSSKZTDJ06FB27dpV9s1ee+21TJ06tUjwpdDY2MiAAQNQ1XwC0Ov1kkqlSKeLJ7LNnj2byy+/nOeee45bb721Y1PdwowZM0zNvGvdOzy9ZRG1FX3JauphSZEKt4TAemw9wV/+YULi7WV/71hX1swwVjtbRRxT6KSzwghpi525Ef5mNRVV04hl09x19KmmIB999NGyyRg9ejSbNm3ixRdfZOzYsYftgH355ZdFZNx44428+eabncgYMWKEGUzccsst/OUvfynrfrrCY489ZqZYbht8GlXeGkLpQ084KRzAK5r5n1Wxjxn6/VLnyKa/QEw0/UaBnyj2FfntbG6SQUMyzIn+Wq7oMwyAYDDIs88+W9ZD3nrrraxatYqhQ4eWVR/gqKOOMrfvvfdenn/++SKCDNxyyy1F+9dffz2ff/55yRR8Odi+fbtpgge5q5jc9wSSiWDJiKszEZJORI4RrS2MbeSxF8p+d6dzZVXTUDXNogoxJZ+L0kwHni2IrIqO5QgJZ5Jc3Xe42eBjjz1GMBg87APed999ZRNXiOOOO441a9awadMmfvvb3wJ00g6ACy+8sFPZqFGjaG5u5uqrr+72dQGeeuopc/un/U5CsTpJdYi4isyT0bfvwJkWjmM9uu9wx4RRZ3S8hpE6OUkVot40U5rImSvRKZpSc8c0IJpN09vh5ZLeQwAIh8P89a9/PeyDTZ06lUcffdTcF0KQzWb1Sc5l4KSTTirSqsGDBxcdnzBhAv379+/y/Jdffpm33noLp9NZ1vUMbNu2jfnz5wNwSmU/zqo+inAyAhRqRZ4IqVhVckclpNxYsW34oE6/GllVNVTEmYW//LwZK9jWCqdj6r3yxng7Z1f2Z6CrwnzQ5ubDL3Y57rjjeOaZZxg/fjwnnHACAwcOpL6+nr59+zJkyBCmTZvG/v3lr80YOnQoo0aNMvdPO+20w55zySWXEI/Hi3JX5aBwqODS3kMhm0SRpJyWHIYIozCXTrEdP/AUyVo88UfOqllUoY3MO3CRd+amAxe5lHpOO0RuoEpTObtmkNnYzJkzD/tAsixzxx13cPvtt7N48WI2btzInj17aGpq4sCBA2zdupUnnniCPn368Nhjj5UtKCM1oygKN910U9nnzZw5kzfffLPs+i+//DJtbW0AnF7ZH5vNQ0bT8uI+FBHkNUnEk1j6Vo+w9O/Vp7C+LLlsqIjRhvCzRZpRvG+kRzQhaM8m6e+t4qK6YwF9LHzx4sWHfSBN08o2Tb/61a944IEHyqo7bdo07rnnHmbNmkV9fX1Z5xi49NJLWbFiRVkp9mg0ysqVKwEYWdGX8dVHEUzFSmhFaSJMZLIoFZ5ay9F9ji4sltOJ5FBVloYURlImORTvm6OC6IT0s3vp4/AB+i/NXMbWA1gsFu68805eeeUVLrjgArP84YcfZt26dYc93+Fw8Lvf/Y7LLrusR9cfM2YM//u//1tW3QULFpjbw7zVoKaQTLPVkYpiIiQpv+gVi4Jt2FHDCw4jJw62DdMUSdcCI6rqMG3fSDIaGWFVCEQmyRBvjdnQZ5991k0RFGPhwoU89dRTTJkyhffff5/bbrvNPPb6669/o7bLxbXXXsu555572Hpbtmwxt4/x5sJoURRXAcVaYRJReDSjYhtcP7qwVLbUVtZnUmlUrWDkz9AErcB/ULxCCVXlRH8dAIlEgqVLl/ZABDquuOIKxo8fX1RWSIgR2XxTNDU18dFHHzFz5kw+/vjjknO8HnnkkcO2s2LFCnNW5Eh/H2w2N1mhURxn5dGJCONoKoPscw0sPGrJqtmBwqrotr1gFom55s6c+FY8HxdZZmAub/XVV1/R0tJShkhKo0+fvF+bMWMG8+fPx2rNTxEq1fHrDsLhMDfccAOzZ88uKrdYLFx88cU8/PDD5uyTMWPGcOaZZ/Lpp5922V5rayvr16/n1FNPpZfdjd9iJy0EFulQREAhYQAimUap9NUrlT6f2qb/OmRVlgZk1YLVRZphnrT8qiLDZOXICGWSHOXrxRlVAwHYsOGbLY8wQsk333zTFNwrr+RnznTUnu6gpaWFE088sRMZoA8PvPnmmwwbNowpU6bQ1KRPNDznnHMO2+7evXsBCNicuC121IKcVknzhFTasQc8lUq1v8ooklVEvSY0NA2TCNUkguLlXcaDZBKcWtGPityIYOGMjZ5gz549DBw4kOuvv77k8VJDv21tbTz55JPMmDHjkJncW2+91cyrOZ1OrrjiiqI+i4HXXnuNuro65s+fz7XXXtvpeKnrA1RanfS2u4mrmS6JADo5eQlAFUg2W4XktJuEWDQhagrNk1a4lqMDEWbDapa+do9Ztnnz5sM+wOGwe/fukuXXXXdd0ZxcgE8//ZTx48ebUd2jjz7K0qVLqaurK6q3efNmc4Bp2LBhLFmyhEAgANBlBviCCy5gzJgxBAIB2tvbu7zfwkE3n9WBpi/rMzuIBjparY70SJIkI+E39mVVaM5C8yQ6mKeu0NtltkFra+shauYRCASoqKgoqy7oCcIZM2Z0Kp86dWpRiL1z586iKTwG3njjDXP72WefNckw2h4zZkzJ665cufKQZEBx/sxjtaEvCJcQXTl2io2WhP4uF2QJSZLsRj2LhlAMh218dwXziGLjKKf+cEKIw968gbPOOou5c+fS1NREa2srDQ0NbNy4kV27drF//34OHjyILMv069eP2267jdGjR3dqQwhRMnlZyo8ZP5SqqirGjh1rnv/+++8Ti8UIhUJl3XcpFAYaVTYXaDoheUp0dDZYBXsSuWyxZOZP9BFDg5AybkQTAhQLlTYXAPF4nGg0WtZDfPTRR7S0tFBbW0ttbS3HH388559//uFPLHwYSWL8+PG89FLxUrxx48Z1qmuYsNbWVpYtW8aECROYNm0aTz75ZLeuWQqKkp99Es9mQZKROkiwk3kynyHfeZT1A+Z4sCyEEFo35nlrQmCVFZyKTmooFCISiZR1biwWM216uXjxxRd5+OGHO5X94Ac/MPenTJlSlBo3MHnyZDMdcuutt3L55ZcfETJAzwwY2JcIISv5JGFJ85Tb0JcAGklIWd/XXzYB6G9N6tZgs4Y+M9GRW0rQ3t5uRjnHHHMMVVVVhzqdP//5z925HA8++CAPPPBAp/TJu+++SyKRYO/evbz88sslzx04cKCZaNy6dWuRT/mmcLv1wSUhBOFsCquka8yhiDAWxsqSpO/n3/liciADcboBAUUvBbBY8r+Mq6+++rBOe+3atTz99NNlX88Y9yjsR7S3tzN16lQ+/PBD6uvri5KCkUikKAL605/+1G2zWA6MH15TKkZzOo5TUYq8R9486UKXje2chsi5bUWAJETCOE8Gyn5blwRYJIVoJkFjQk87HHvssZx00klMmDCB8847j6+//vqw7Tz44INlj3f8+te/BvSUhqEJs2bN4g9/+AM33XRTUR+ktbUVn8/HeeedV9TG/Pnzefrpp/F4PBwpGNmFg+kYoUwSa84vdzZPRkdRMl+qIJskSchIWcDM4cgIsbtTsHwIWCQJsik+PJifb7V69WoWLlyIw+EoK80RDof5l3/5l7Kud+qppzJtmr6S7tprr8Xtdpv9h6amJpYvXw7oJumMM/QR0UL7buCOO+6gubmZ9957j/vuu48RI0aUdf1SCAQCZqplR6yN9kwSm6x0Mk+mVlBARq5cRkKxWiCZbtWiCVMpZBRlr9kbLAMC8LoqeGrHZyzIkWKYjO5M2V+wYAF33nlnWXWnT5/OK6+8Qn19PfF4nLq6OgYN0gfGzjzzTGRZZsiQIWzZsoUBAwbw2mulV4U5HA4mTpzIo48+yrp169i5cye//OUv6dWr7PU0AIwcORKfTx92WBVshGxGf9uRRBER+jslC7XCePuq/u4Vi82K1h4NZpvz/QZZC4Z3Sk576St3AY/Fhk1W+MFnL3LD2reJZvVOkuHoysXvf/977rnnnrLqTpkyhb1799LW1sbevXvZsWMHd999NzabDSEEVquVadOmsWvXrrIFfNRRR/Hkk0/S1NRkmsZyMGTIEHO7IRkCixVyWtDRPBlRlQzF5kqSUJx21NbQ7mwoGiH3Egc5vXTjPrkm0LmP3wUkIKOpVDu8SIqFGRs/ZFP7PkBXZbu9e+Q+/vjj3RrXLgwaHn/8cbMflEwmmT59ereubSASifD222+XXd/IhUWyKdaEDlBhc+raIUnIyJ3NU+6dYXkfojtvi8OOCMd35pqVAWTL0P7rRTojumu2MpqKz2IHTxWpnKmqra3tlE8qB7Nnz2bIkCHs2LGj2+cqioLb7UaWe7YWcNu2bRxzzDFljUoaOPXUUwFY1rqHDeGDVNpcZt+jpHkq0BaTHElGsVpJbN69OtesCiDLHtdO7UDbWrr5QAL9BQCoGQ6k8h3Dvn37dqsdA1u3bmXw4MH87ne/69H5PcF7773HiBEjzLR7ORg1apTp0L9OhEBophYoJfxEIUGKqTUSsqIgpdLENmzfmGtaA5DDj7yI5LSvlezWkjdQChKYr01CzbI92mYe6927d9ntlMK//du/0adPn2736LuDZDLJNddcw0UXXUQqlerWuYXR4YZwE7JswfAfpfyESYShKegmS7FbUYPRPcmd+4x+gj7up+49SHbHvpWSr3sOWZgtyHwZyvcpvikhAPv372fy5MlUVlbyH//xH+ZygG+KlpYW7r77bpxOZ9EAWHdgDJa1pRMsbN5JrcNj9r51oec0Q85rh/mvYN/q9ZD8ev+6xO79BwDT8coiniK9assS2eM8zHLEztCEAIuNhmQ+uVhqnq4kST2y88FgkIceeojevXtz8sknc++997Jw4UIymUzZbWzfvp0nnniCM844g5qamm7lsmRZprq62rzvESNGmP2XD5t3sCMepMJiLwp35cJtCp08ZupEAhwVXhJbvjbemGNK3gKQ/fLrDSKV3oIiDSnXueuDMQK7YmNvMkxbOk6lzcVZZ53Vqe5jjz3Gvffey+OPP152mNsRa9euZe3atUyfPh2fz0f//v2pra2lX79+1NTU4PP5yGazBINBGhsbaWxsZO/evTQ2NvboeqCPr59zzjnmhL2rrrrKPLaseRfZTAZZgKaJ/EihJPIagcgnE3Pj7ZKsa5JIZkRoxSbjLXRmusHo0VH50r//2nJc/W+05vLHCHJTVGmJh3j3tJ/wg7rjAH2q6LZt23JNS+zatcuca3vaaacdcspQIBBACPGNxiqOFB566CHq6+u54YYb8Hg8bN++ndraWvZFQ5y35AUimRRei033pwUvhzb7I0j59IlsbMvYagIkd+1fuuZH084WqmZHzyfqHgAAIUgv2/Ch7HWbb94sF7IkgZphTYEfKUzm2Wy2oknNh0qZjBs3jmAwSGtra1Hn6x+FSZMmce655+JyOfnDs8/qaxsF/OWrZWzat41UKsn+cJv+CbXSGGzWP+3N7G9roTnSTjQVJ51Jo6YyiGQWkU5j8ToJrdywSKhalvzSdMn8H4D1+IFUvHDPehFNDBfp7i3/aknFOMXfm5Vn3wzob2sYOXIkoPfe9+zZQ2VlJQD79u3j6KOPJpksXuwydOhQVq1aZfb2t2zZ0q11I0ca5593HvM//BDDDiSTSfbu3kskkeCxzYtJZzIMclVikxVsVgsWRUGWZTQZoiJDXNZoySbYF2olrKVJ2WQyNhmL00FFnzq23fTb84LL1i1Cd+ip3IWEmTvPbNpFZu32/7adfsKz6p6DpSYVdYleDi+rWnbxZuMGLut7AieffDLPP/88N910E3feeadJBuhZ0ssuu6woyunduzeLFy8uSr0MGTKE2267jWeeeaase+jfvz/pdPqIRWQ//unPCEbjLF22nKwmCEViqKpKUqhc6RxMv+oANpsNh82Gw27HYbfhdNqx2WzIdhuRRJz1mzaw/mCKllSMoJwl5rViGVrP5g9WvRdctm4l0AtoR9cSDTq8m8A56fRq/+M378zu2Ff2a5sFYJVkDsSDnFE1gE/PKj3zXLezul1dvXq1mX6ora3liy++MDuUK4MN9HF4qXf6aWxsZNCgQSUX5BhwuVy88MILXHnllTQ0NDBo0KBuRWEGJElm0LHHceKoMZwydhz1AwZwoOkgVosFl9OJy2HH6XTgcjhwO+zYbFadEIcdh92ORZHRO9cSBw828+X69TTu3YvIqLhsTrweH+1tbSxd9ulLHy2a//toOhlE9x1hdKeeBbRiNVBkqt74zXNKv16/0JraytISw7E7FCsN4QOc13soH5720071rv58FhfXDeWqen1u8d69e1m0aBEXXnghNTX6HOHVoX2Men86j468lPuO1eP9Sy+9lDlz5pS89ujRo1m8eHGRj7ryyitLTorrCgMGDeasCd9n6PAT6T9gIH5/gEwmRSTUjs1qwe1yYrfasNut2O12nHYbLpcTl8OB1WpFCEEqnSaVzpJIpdm9ezcbN3xJS0szLpcHb6CClmAbKz9b8tbyxYv+lozFPwdc6FrRCiSBDCUJAVzXnHu076Gfbs1+1aCUa7YEem9UkRT2Rw4yyNuLHw8YyQneGjZHW3i9cSPrGzfRv3oAK86+mTpHZwVcFz7A9z6dQVuwkR8OPp23T70G0JdAT5w4sVP9iy++mLlz53Yq76o+gEVReOTh34Cs8Prb8xh31jmMHHMqFRWVJGIRMukUFkXGbtNNkMOhE+Bw2HE7nXjcLuw2G5IE6XSGWCJBOBojkUwTiyf4eucOtm3bRjabpaqmlmCoXV2+dPH7a1Yuez0RjX0OWNG7Gu1AFN13pHJkqCUJAah85f7nrccPvFHd11K2LzGHdmWZtmSUdCqin6uBxe6lrzvA7kgTvRxefj/iIq7qewIA7ZkkLzas474N84mpWeqcfhLZFMvPupEh3hqEEAwaNKhoVe+FF17IvHnzzP1XG9ZTa/fwvZpBRKNRBg4c2GmumCzBazNnMvGSK1iycjXNbe0okkQqEUOWME2Pw27H6dA/LqcDn8eN2+XEYrGgqirxRJL2cIRQOEoskSCb1QhHIny1bQtNBw7g9gWIpxKs+XzFJ18sW/paPBJehe64LUAoR0QCSBd8spTyIQacPzrrWP+Tv9ic3bRb7o5z70SS0EcYLbKMACyywsF4iGQmzglVA6m2utgWbWFfuIkKdwU+qxNVaDREmrn/+Ak8PFRfGvDCCy+Ya8QfeughHnzwQfMacw9s4ZIPn+LekZfx2xP0lcZ33XVXp3H7Rx9/mjMmnM8HCz7C5/VQUxnAarFgzwnfmSPC5XDi8bjwe904HQ4kJFLpFJFojGA4QjgSIxaPk0xlyKgqLc3N7N2zm3QmSyweY+OXaxav+3zlnGh7+wpAyZFhEJFE14g0upkqIoOuNASg8sX7HrGOHvIrdXsjWA7/BpzC4X2pQ6FVks3cgCLJ5husU2oWv9WBz2rPvWlIn44ZyiSpsNhZffYv6JWbsrp48WICgUDR0Ousxo1MXv4yaFnOrz+B+eN+AugzD42JcX37D+TmO+7hqGOHsGf3TqoCAZwOO3a7DafDgcupf9xOB163G7/Pi92uD3rF4gmCoRDt4QjRaJxYIkE8kSSdyZJKpWltbSUYDHKw+QDbt25Zse7z5TOj7e2foWuDDYiQN03p3LfhL0wzZZBRLMcOsJ58jFQ5Y9oWLRw/VsSTXZquQxFhQJYklIJXcxmv5ujYiHGaIik0RFu4sv5EZp5SevDq3zYt4HebFlLpqsBrddCUCDHv1B9zTi99aPeNN97gf155nYsvn4Lb7SLU1oLb5cTpcOjRktOB2+nE7XLi83oI+NzYbTayqkYkEqUtFNaJiOWIiCdJJFNksiqqqhIKh9m54yu2bP7yiw1frHwl2t6+HN1RGxoRI+8jDI3I5EgoJMLM05YQXTG8d0++3HP7j2Zn1u8ApTgx2HFypNS50DxiRGJSwZULZ/GValWRZBpibYzw13H34DM4vbIfTsXK31t28dtti9nQ3kgfdxUORUEI2BNu5cRAHcvPuZFdG3ey8uu9ZAXEwu1YZBm3W4+MXE4nHrcTt8uF3+sh4PPidNhIZ7KEIlHagiHCkSjRWJxoPEE8kSCRTJJKpZFkBUmxsGPHVyz9eOHy9Z8vfzWdSKzLCdiOHsLGyJumQrPUUSOKiCiLEIDKl/79eduoY2/M5kxXd4g41MW6IsIkCgmbrHAgESKRThBwBbBIMi2JEE6rg94Or7nKS1E1EopGUyTInfFBDGqWCSpZ6mqqsNt05+x2OXG7nXicTvw+L5UBP06nnUwmS3s4QlswRCgaIxqLEYsliCeTJBJJ4vEEsqLg8gU4cGAff//w3VXLFy2cmU4kVpEnIkqeCCOM7RYRpWRUEtbjBzgqX/n1BhFLHi3aI2bnpyvz1JGKjocPR0ThfmFyLqXq72Z3WSxIgKppIASKCs1SEm9M5coDldTFFKwVbgI+Lw67HbdLD1e9bl0jKgMBPB4nqqoRCkdoCbYTCuc1IpHzE4lEAiQZb6CK1mALiz+av3LJwvfnRtvbl6IL3EneRyQpjppKmaaylh6XFUK5rjn3tMBjNy7NbNkLWVWf7lOGVnStEfmjpYgo1bjxh4sBEAJJgKQJmklQExb8eH81x8h+In4Fh8OOx+XC43bi9bjxedxUVQTw+7zIkkQ4GqOlLUgwFCYSjRONx4nHEySSKeIJPcfmCVTQHg6z4tNFaz6Z/+5rweaDy9AFa0fXhgidiei2RpSWShnwP/TTO9w/ueDp7OZdJfzJt0NEsZ/JjS8ITDJaSNK3XfDTA72os3mJ+S14nE48Lhdejwuvx02l30dVZQCHzUYsnqSlLUhbKOcn4gli8YRumhIJhCZw+yuIxCKs/nz5xr+/N+flln37PkUXqBPdR3TUCMNZfyMiSsnusKh++f4Z9lOPvy69eTeSReFImqeSx40xBQqmaAqQVY2DUpKBrXD9gVoCLjdJvxWfy4XH7cbrcRHweamprMTv9ZBVs7S0tdPSFiQcjRGJxojFddMUTyTQVA1PoJJEKsWKZZ98ufCdN15ua2pagS5oBzoJpfoRR4yIrmRzSChVPqpf/81HSm3F97I79+dI+S6IyMVqAmRV0Cwl6RcUOTI8ZHw2fG7dNPk8bioq/NRUBrDb7ISjUQ62tBIMRYhEY0RjceKJBPF4gqyq4fL5SSSTrF+zavuiD96d2bDjq4/QBW2Ypq561hlyf0eTYjK+EbrdDbcM6m2tefWBFZLTcbK6p8kkBY6UeSqY0G9s5xbly5qgTSSpCqncvL+OapeXtN8gw0OFz0NNVSUBv4+sqtLc2kZLW5BQpFAr4mTSWdz+AKlMhjWfL9++8L23Xt6/a9didA1wkdeIjkQUOmuVI6ARpSXWTVgG9bH3eu3B5bLHeVJmRyOS1VKWVpQqLSTC2M+vMMqTo6iCdlLYoml+sa+O/o4Aab8Nf04rKgN+elVX4XTYCUdjHGxpIxgKE47qEVQsFiedzuD2BcgiWL3ys53z58x+af+eXZ+ip8GduW/DWRumqdBZF6Y5jigRpeTVLViO6u3s9bf7PrL0qR6X3rrH9CmlGj4UEZ3Mk2QYKb3bK6E78ITIkkwm+Nm+akYpNSQqdc3wez1UV1ZQXRkAoLmtnda2IOFIjHAsRiQaJZPO4vL5SasqX6xYuv3jD959a/dXW/+OnnV10zUR37pGdETPM4eApbZCqv7ztHdtwwdNzGzVF9J3nAFftp8wNaTgDxjnwltN0ziYjXHpAT8XZ/sQrbLh9bjx+zz0qqzE7/OSTKV0rQiHdccdjpJKp3F5fCQyGT5fvmT7Jwvef2v3ti0foRPhQTdJYYpNU0dn/Z0Q0Uk+PW7Aaafmj7/8H+eEUT/Lbm9AJNMgy52IKNwr8hPoc5YoJML4CN1v7NdijGqxc2O8P2qFE4fXSYXfR01lBU6Hw+xXhMJRwpEoyWQKm8tNKpth/drVjQveefPVr7dsmoeuCYVEHC5qMj7fGb4xIQYC/zr55sCdV/xRPdCK1twOHRZBQt43GGXFCyANMqQiUxXWUrgjWe5o70+dtwLNZ6fS76O6IoCiKARDYYLter8inkhgd3lJZbOsXbNq37y3Zr28a8vmhejmqNA0GRph5JtKEfGdaERHHDFCAFwTRp5Y/fgt78oeV33m6/1IQuiRktmxM9IhnVMjBikyumYgQNVUWlNRrm/rzXhLb6IVVip9Pir8PjRNoz0UoT0cJh5PYnW6yAhYuXzJnnlvzJy5a9uWT4A2dGedJE9EYdLvn4YIA0eUEAClJiDX/PbnL7gvGPuTbEMzWlsESZE7+QlT+JQgBr2/cUCNMTrk5OfpQWT8drx+Dz6Ph2w2S3soTCyewO5ykxXwxepVre/MfvXVzWu+mAcEyZumCLpmdAxf/yE+4nA44oQY8E468+yqf//xS9a6yr6ZnfsglUGSZX1tNkYEVTzDTwYzTxVX0yjxNHdGB9HfVYHmt+NxOkmn00SjcZweL0KxsGL50oNvvfbSrI2rP18INKGbpjT6mIShEaVyTR171/8U+NYIAVAqvNaqO6+4r+Lq8++XbVZrZvcBSKSRFLlgVjidtEPKauzPRLk0XMOPpAFE/Qo2m5VsJovL7QWrjc+WfXrg7dmvvrP6s6XzgEbAiy50QyM6EmGEsP8UpqkrfKuEGHCOOLpvxTXfvyfww9Nvsvo8jsyeJrRYElmWCzQknzQMZZME4oJ/TR2L1+kia5UIVFQiWW2sWL6kac7rM99e+cmi94B96ERkKNaIDPnZHP8niDDwnRBiwDlkQK/qn/1gWuB7o6+396mpUFtDqMGI/mewJUkPAlRBUybKFZFafiDVk63z4fb6WL92ddtrL77w2rK/L3wbfT6TD13ghUR0jJqMnvU/PREGvlNCDFhrKlw1V58/ueL8sTe7jxswxuJykm1pJ9sWpj0dpzou8aBjNFW1dazbtiE468UX/vejd9+egx41VaELvDBqKhW+/lP6iMPhH0KIeXFZxjt66En+M06cUHn2qEnOY/qdGXFI/CTShyEbwweefum55z7+4L35QtVa0OfBKnSd4uiYazK+/0/hH0pIISSLgn3E0cMHjRs9aVjS7Z4/87W/haORvcAA9DGJQgJKOev/kxrREf8shBj3USjICiCALnjjF69Rel7T/3kiDPx/tcXfsY70TpIAAAAASUVORK5CYII="; diff --git a/docker/jupyter/unsloth_labext/src/outputSelect.ts b/docker/jupyter/unsloth_labext/src/outputSelect.ts new file mode 100644 index 0000000000..1c31ade442 --- /dev/null +++ b/docker/jupyter/unsloth_labext/src/outputSelect.ts @@ -0,0 +1,126 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +import { + JupyterFrontEnd, + JupyterFrontEndPlugin +} from '@jupyterlab/application'; + +/** + * Colab-style Ctrl/Cmd+A inside a cell output. + * + * Clicking an output leaves the notebook in command mode, so Ctrl/Cmd+A fires + * `notebook:select-all` (every cell). Colab selects only the clicked output's + * text; reproduce that and stop the event. Listens in the CAPTURE phase, acts + * only on exactly Ctrl/Cmd+A (no Alt) outside an editor/input, keyed off the + * target or last pointer-down (not the stale selection anchor). + */ + +// Output containers, widest first: a single output, then the whole output column +// (covers a click on padding between outputs). +const OUTPUT_SELECTORS = ['.jp-OutputArea-output', '.jp-Cell-outputWrapper']; + +function closestOutput(node: Node | null): HTMLElement | null { + const el = + node == null + ? null + : node.nodeType === Node.ELEMENT_NODE + ? (node as HTMLElement) + : node.parentElement; + if (!el) { + return null; + } + for (const sel of OUTPUT_SELECTORS) { + const hit = el.closest(sel) as HTMLElement | null; + if (hit) { + return hit; + } + } + return null; +} + +function inEditableContext(): boolean { + const ae = document.activeElement as HTMLElement | null; + if (!ae) { + return false; + } + if (ae.isContentEditable) { + return true; + } + const tag = ae.tagName; + if (tag === 'INPUT' || tag === 'TEXTAREA') { + return true; + } + // CodeMirror 6 editor (cell input in edit mode). + return !!ae.closest('.cm-editor'); +} + +const outputSelectPlugin: JupyterFrontEndPlugin = { + id: 'unsloth-jupyterlab:output-select-all', + description: + 'Ctrl/Cmd+A inside a cell output selects only that output, not every cell.', + autoStart: true, + activate: (_app: JupyterFrontEnd): void => { + // Remember the last pointer-down: a click on an image/widget output leaves no + // text selection, so the anchor alone can't tell which output is meant. + let lastPointerOutput: HTMLElement | null = null; + // ...but only trust it while that output is still in the document AND still + // inside the ACTIVE cell. Keyboard cell navigation (J/K, arrows) fires no + // pointer event, so an unvalidated value would make the chord on a later cell + // select the previously clicked output and swallow `notebook:select-all`; and + // a re-executed cell replaces the node, leaving a detached range that selects + // nothing at all while still suppressing the shortcut. + const rememberedOutput = (): HTMLElement | null => { + const output = lastPointerOutput; + if (!output || !output.isConnected) { + return null; + } + const cell = output.closest('.jp-Cell'); + return cell && cell.classList.contains('jp-mod-active') ? output : null; + }; + document.addEventListener( + 'pointerdown', + (event: PointerEvent): void => { + lastPointerOutput = closestOutput(event.target as Node | null); + }, + true + ); + + const handler = (event: KeyboardEvent): void => { + if (event.key !== 'a' && event.key !== 'A') { + return; + } + if (!(event.ctrlKey || event.metaKey) || event.altKey) { + return; + } + if (inEditableContext()) { + return; + } + // Own the chord only when in an output: the target, else the last click + // (not the stale selection anchor; see the header). + const output = + closestOutput(event.target as Node | null) ?? rememberedOutput(); + if (!output) { + return; + } + // We own this key: prevent Lumino's `notebook:select-all` from also running. + event.preventDefault(); + event.stopPropagation(); + try { + const range = document.createRange(); + range.selectNodeContents(output); + const sel = window.getSelection(); + if (sel) { + sel.removeAllRanges(); + sel.addRange(range); + } + } catch { + /* no-op */ + } + }; + // Capture phase: decide before Lumino's keybindings consume Ctrl/Cmd+A. + document.addEventListener('keydown', handler, true); + } +}; + +export default outputSelectPlugin; diff --git a/docker/jupyter/unsloth_labext/src/splash.ts b/docker/jupyter/unsloth_labext/src/splash.ts new file mode 100644 index 0000000000..6215bf0877 --- /dev/null +++ b/docker/jupyter/unsloth_labext/src/splash.ts @@ -0,0 +1,88 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 +// +// Replace the JupyterLab loading splash with a spinning Unsloth logo. Provides +// the core ISplashScreen token; the stock splash is disabled + locked at build, +// so this is the only provider. Animation honors prefers-reduced-motion. + +import { JupyterFrontEndPlugin } from '@jupyterlab/application'; +import { ISplashScreen } from '@jupyterlab/apputils'; +import { DisposableDelegate, IDisposable } from '@lumino/disposable'; +import { UNSLOTH_LOGO_DATA_URI } from './logo'; +import { SPLASH_LABEL } from './branding'; + +const STYLE_ID = 'unsloth-splash-style'; +const SPLASH_ID = 'unsloth-splash'; + +function ensureStyle(): void { + if (document.getElementById(STYLE_ID)) { + return; + } + const style = document.createElement('style'); + style.id = STYLE_ID; + style.textContent = ` +#${SPLASH_ID} { + position: fixed; + inset: 0; + z-index: 1000; + display: flex; + flex-direction: column; + align-items: center; + justify-content: center; + background: var(--jp-layout-color0, hsl(70, 8%, 12%)); +} +#${SPLASH_ID} img { + height: 72px; + width: 72px; + animation: unsloth-splash-spin 1.2s linear infinite; +} +#${SPLASH_ID} .unsloth-splash-label { + margin-top: 14px; + font-size: 13px; + opacity: 0.7; + font-family: sans-serif; + color: var(--jp-ui-font-color1, hsl(60, 30%, 92%)); +} +@keyframes unsloth-splash-spin { + from { transform: rotate(0deg); } + to { transform: rotate(360deg); } +} +@media (prefers-reduced-motion: reduce) { + #${SPLASH_ID} img { animation: none; } +} +`; + document.head.appendChild(style); +} + +const splashPlugin: JupyterFrontEndPlugin = { + id: 'unsloth-jupyterlab:splash', + description: 'Unsloth spinning-logo loading splash.', + autoStart: true, + provides: ISplashScreen, + activate: (): ISplashScreen => { + return { + show: (): IDisposable => { + ensureStyle(); + const overlay = document.createElement('div'); + overlay.id = SPLASH_ID; + + const img = document.createElement('img'); + img.src = UNSLOTH_LOGO_DATA_URI; + img.alt = 'Unsloth'; + overlay.appendChild(img); + + const label = document.createElement('div'); + label.className = 'unsloth-splash-label'; + label.textContent = SPLASH_LABEL; + overlay.appendChild(label); + + document.body.appendChild(overlay); + return new DisposableDelegate(() => { + overlay.remove(); + }); + } + }; + } +}; + +export default splashPlugin; diff --git a/docker/jupyter/unsloth_labext/src/uiChrome.ts b/docker/jupyter/unsloth_labext/src/uiChrome.ts new file mode 100644 index 0000000000..d7abaa083c --- /dev/null +++ b/docker/jupyter/unsloth_labext/src/uiChrome.ts @@ -0,0 +1,55 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +import { + ILabShell, + JupyterFrontEnd, + JupyterFrontEndPlugin +} from '@jupyterlab/application'; + +/** + * Colab-like chrome tweaks applied image-wide. + * + * Hide the right activity bar (Property Inspector / Debugger) by default. + * JupyterLab has no settings key for this, so hide the strip with CSS and + * collapse the right panel once on startup. Reopen from the View menu. + */ + +const STYLE_ID = 'unsloth-ui-chrome-style'; + +function injectStyle(): void { + if (document.getElementById(STYLE_ID)) { + return; + } + const style = document.createElement('style'); + style.id = STYLE_ID; + style.textContent = ` +/* Hide the right-hand activity bar strip (Property Inspector / Debugger tabs). */ +.jp-SideBar.jp-mod-right { + display: none !important; +} +`; + document.head.appendChild(style); +} + +const uiChromePlugin: JupyterFrontEndPlugin = { + id: 'unsloth-jupyterlab:ui-chrome', + description: 'Hide the right activity bar by default (Colab-like chrome).', + autoStart: true, + requires: [ILabShell], + activate: (app: JupyterFrontEnd, shell: ILabShell): void => { + injectStyle(); + // Collapse the right area once restored so an expanded panel doesn't linger. + app.restored + .then(() => { + try { + shell.collapseRight(); + } catch { + /* no-op */ + } + }) + .catch(() => undefined); + } +}; + +export default uiChromePlugin; diff --git a/docker/jupyter/unsloth_labext/style/index.css b/docker/jupyter/unsloth_labext/style/index.css new file mode 100644 index 0000000000..0046a04a36 --- /dev/null +++ b/docker/jupyter/unsloth_labext/style/index.css @@ -0,0 +1,6 @@ +/* "Unsloth Dark" theme entry point. + * Start from the built-in JupyterLab Dark theme (theme.css pulls in its full + * variable set + base rules), then override the palette with the Sublime/Colab + * Monokai colors in variables.css. */ +@import url('@jupyterlab/theme-dark-extension/style/theme.css'); +@import url('./variables.css'); diff --git a/docker/jupyter/unsloth_labext/style/variables.css b/docker/jupyter/unsloth_labext/style/variables.css new file mode 100644 index 0000000000..c95d90fcf4 --- /dev/null +++ b/docker/jupyter/unsloth_labext/style/variables.css @@ -0,0 +1,97 @@ +/* Unsloth Dark = Sublime/Colab "Monokai" palette, overriding JupyterLab Dark. + * Applied on :root because the theme manager only loads this file while the + * "Unsloth Dark" theme is active, so it never affects the light theme. + * + * Exact HSL from Sublime "Monokai": + * bg hsl(70,8%,15%) fg hsl(60,30%,96%) selection hsla(55,8%,31%,.7) + * comment hsl(50,11%,41%) string hsl(54,70%,68%) number hsl(261,100%,75%) + * keyword hsl(338,95%,56%) function hsl(80,76%,53%) builtin hsl(190,81%,67%) + * param hsl(32,98%,56%) error hsl(0,93%,59%) + */ +:root { + /* surfaces */ + --jp-layout-color0: hsl(70, 8%, 12%); + --jp-layout-color1: hsl(70, 8%, 15%); + --jp-layout-color2: hsl(70, 8%, 10%); + --jp-layout-color3: hsl(70, 8%, 8%); + --jp-layout-color4: hsl(70, 8%, 6%); + --jp-toolbar-background: hsl(70, 8%, 13%); + --jp-cell-editor-background: hsl(70, 8%, 15%); + --jp-cell-editor-active-background: hsl(70, 8%, 15%); + --jp-cell-editor-border-color: hsl(70, 8%, 22%); + --jp-rendermime-host-background: hsl(70, 8%, 15%); + --jp-rendermime-error-background: hsla(338, 50%, 56%, 0.15); + --jp-cell-prompt-not-active-font-color: hsl(60, 8%, 55%); + --jp-notebook-multiselected-color: hsla(80, 40%, 40%, 0.18); + + /* inverse surfaces */ + --jp-inverse-layout-color0: hsl(60, 30%, 98%); + --jp-inverse-layout-color1: hsl(60, 30%, 96%); + --jp-inverse-layout-color2: hsl(60, 10%, 72%); + --jp-inverse-layout-color3: hsl(60, 8%, 55%); + + /* text */ + --jp-ui-font-color0: hsl(60, 30%, 98%); + --jp-ui-font-color1: hsl(60, 18%, 90%); + --jp-ui-font-color2: hsl(60, 8%, 66%); + --jp-ui-font-color3: hsl(60, 6%, 46%); + --jp-content-font-color0: hsl(60, 30%, 98%); + --jp-content-font-color1: hsl(60, 30%, 96%); + --jp-content-font-color2: hsl(60, 12%, 72%); + --jp-content-font-color3: hsl(60, 8%, 52%); + + /* borders */ + --jp-border-color0: hsl(70, 8%, 26%); + --jp-border-color1: hsl(70, 8%, 22%); + --jp-border-color2: hsl(70, 8%, 18%); + --jp-border-color3: hsl(70, 8%, 14%); + + /* accent / links / brand */ + --jp-content-link-color: hsl(190, 81%, 67%); + --jp-brand-color0: hsl(190, 81%, 72%); + --jp-brand-color1: hsl(190, 70%, 58%); + --jp-brand-color2: hsl(190, 60%, 46%); + --jp-brand-color3: hsl(190, 55%, 36%); + --jp-accent-color1: hsl(80, 76%, 48%); + --jp-warn-color1: hsl(32, 98%, 56%); + --jp-error-color1: hsl(0, 93%, 59%); + --jp-success-color1: hsl(80, 76%, 45%); + + /* selection / cursor */ + --jp-editor-selected-background: hsla(55, 8%, 31%, 0.55); + --jp-editor-selected-focused-background: hsla(55, 8%, 31%, 0.75); + --jp-editor-cursor-color: hsl(60, 36%, 96%); + + /* CodeMirror 6 syntax tokens (Monokai) */ + --jp-mirror-editor-keyword-color: hsl(338, 95%, 56%); + --jp-mirror-editor-atom-color: hsl(261, 100%, 75%); + --jp-mirror-editor-number-color: hsl(261, 100%, 75%); + --jp-mirror-editor-def-color: hsl(80, 76%, 53%); + --jp-mirror-editor-variable-color: hsl(60, 30%, 96%); + --jp-mirror-editor-variable-2-color: hsl(32, 98%, 56%); + --jp-mirror-editor-variable-3-color: hsl(190, 81%, 67%); + --jp-mirror-editor-punctuation-color: hsl(60, 18%, 85%); + --jp-mirror-editor-property-color: hsl(80, 76%, 53%); + --jp-mirror-editor-operator-color: hsl(338, 95%, 56%); + --jp-mirror-editor-comment-color: hsl(50, 11%, 41%); + --jp-mirror-editor-string-color: hsl(54, 70%, 68%); + --jp-mirror-editor-string-2-color: hsl(54, 70%, 68%); + --jp-mirror-editor-meta-color: hsl(190, 81%, 67%); + --jp-mirror-editor-builtin-color: hsl(190, 81%, 67%); + --jp-mirror-editor-tag-color: hsl(338, 95%, 56%); + --jp-mirror-editor-attribute-color: hsl(80, 76%, 53%); + --jp-mirror-editor-header-color: hsl(338, 95%, 56%); + --jp-mirror-editor-quote-color: hsl(80, 76%, 53%); + --jp-mirror-editor-link-color: hsl(190, 81%, 67%); + --jp-mirror-editor-error-color: hsl(0, 93%, 59%); + --jp-mirror-editor-activeline-background: hsl(55, 11%, 22%); + --jp-mirror-editor-matchingbracket-color: hsl(54, 70%, 68%); +} + +/* Active line tint inside the code editor (Monokai line_highlight). */ +.cm-editor .cm-activeLine { + background-color: hsla(55, 11%, 30%, 0.35); +} +.cm-editor .cm-activeLineGutter { + background-color: hsla(55, 11%, 30%, 0.35); +} diff --git a/docker/jupyter/unsloth_labext/tsconfig.json b/docker/jupyter/unsloth_labext/tsconfig.json new file mode 100644 index 0000000000..a26bcc1a5a --- /dev/null +++ b/docker/jupyter/unsloth_labext/tsconfig.json @@ -0,0 +1,26 @@ +{ + "compilerOptions": { + "allowSyntheticDefaultImports": true, + "composite": true, + "declaration": true, + "esModuleInterop": true, + "incremental": true, + "jsx": "react", + "lib": ["DOM", "ES2018", "ES2020.Promise"], + "module": "esnext", + "moduleResolution": "node", + "noEmitOnError": true, + "noImplicitAny": true, + "noUnusedLocals": true, + "preserveWatchOutput": true, + "resolveJsonModule": true, + "outDir": "lib", + "rootDir": "src", + "skipLibCheck": true, + "strict": true, + "strictNullChecks": true, + "target": "ES2018", + "types": [] + }, + "include": ["src/*"] +} diff --git a/docker/run.sh b/docker/run.sh new file mode 100755 index 0000000000..c52cb22320 --- /dev/null +++ b/docker/run.sh @@ -0,0 +1,115 @@ +#!/usr/bin/env bash +# Convenience wrapper for `docker run unsloth/unsloth`. Sets the easily-forgotten +# flags behind the most confusing failures: +# --gpus all attach a GPU (entrypoint refuses to start without one) +# --ipc=host ample /dev/shm; the default 64MB crashes DataLoader workers +# --ulimit memlock=-1 unlimited pinned memory (else multi-GPU training stalls) +# --ulimit stack=64MB larger libtorch thread stack (some kernels OOM the 8MB default) +# Plus mounts the host HF + Triton caches so downloads and kernels persist. +# +# Usage: +# bash docker/run.sh # interactive python REPL +# bash docker/run.sh bash # shell in the container +# bash docker/run.sh python /workspace/smoke_test.py # run the smoke test +# bash docker/run.sh python /workspace/host/train.py # run your training script +# ($PWD is mounted at +# /workspace/host) +# +# The full image (unsloth/unsloth:latest) starts Studio (8000) + JupyterLab +# (8888) by default; publish the ports when you want them: +# UNSLOTH_PORTS="-p 8000:8000 -p 8888:8888" bash docker/run.sh +# JupyterLab on the lean core image (unsloth/unsloth:core): +# UNSLOTH_PORTS="-p 8888:8888" UNSLOTH_IMAGE=unsloth/unsloth:core \ +# bash docker/run.sh jupyter lab --ip 0.0.0.0 --port 8888 --allow-root +# CPU-only hosts (Docker Desktop on macOS, Windows without WSL2 GPU, plain +# CPU Linux): no --gpus and set UNSLOTH_ALLOW_CPU=1. Training is unavailable +# but Studio chat / Data Recipes, Jupyter and GGUF tooling work: +# UNSLOTH_GPUS=none UNSLOTH_ALLOW_CPU=1 \ +# UNSLOTH_PORTS="-p 8000:8000 -p 8888:8888" bash docker/run.sh +# +# Overridable env: +# UNSLOTH_IMAGE=unsloth/unsloth:latest image and tag to pull/run +# UNSLOTH_GPUS=all GPUs to expose ("all" | "0" | "0,1" +# | "none" to run without GPU) +# UNSLOTH_ALLOW_CPU= set to 1 to allow GPU-less runs +# UNSLOTH_PORTS= extra -p publish flags, e.g. +# "-p 8000:8000 -p 8888:8888" +# HF_HOME=$HOME/.cache/huggingface host HF cache dir to mount +# TRITON_CACHE_DIR=$HOME/.cache/unsloth-triton +# host Triton cache dir to mount +# UNSLOTH_WORKDIR=$PWD host dir mounted at /workspace/host +set -euo pipefail + +IMAGE="${UNSLOTH_IMAGE:-unsloth/unsloth:latest}" +GPUS="${UNSLOTH_GPUS:-all}" +# Translate index selectors to Docker's `device=` form: a bare integer is a COUNT +# not an INDEX, so `UNSLOTH_GPUS=0` would expose zero GPUs. `all`/quoted `device=` +# pass through; "none" omits --gpus (CPU mode). +GPU_FLAG=(--gpus "$GPUS") +case "$GPUS" in + none) GPU_FLAG=() ;; + all|"") ;; + \"device=*) ;; + device=*,*) GPU_FLAG=(--gpus "\"${GPUS}\"") ;; # native comma list: docker needs the quotes + device=*) ;; # single device, fine unquoted + *[!0-9]*) GPU_FLAG=(--gpus "\"device=${GPUS}\"") ;; # comma list / UUID + *) GPU_FLAG=(--gpus "\"device=${GPUS}\"") ;; # bare integer index +esac +HF_CACHE="${HF_HOME:-$HOME/.cache/huggingface}" +TRITON_CACHE="${TRITON_CACHE_DIR:-$HOME/.cache/unsloth-triton}" +WORK_DIR="${UNSLOTH_WORKDIR:-$PWD}" + +mkdir -p "$HF_CACHE" "$TRITON_CACHE" + +# Warn early if the host has no nvidia runtime registered. Let `docker run` fail +# loudly rather than abort -- some setups report runtimes differently. +if ! docker info 2>/dev/null | grep -qi 'Runtimes:.*nvidia'; then + printf "\033[1;33mWARN:\033[0m 'docker info' does not list 'nvidia' as a runtime.\n" >&2 + printf " If --gpus all fails below, install nvidia-container-toolkit:\n" >&2 + printf " https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html\n\n" >&2 +fi + +# Forward common secrets only if set (empty strings would shadow the image's). +# Use the dash-only `-e VAR` form: Docker reads the value from the parent shell, +# so the secret never lands in argv (visible via `ps auxe` / /proc//cmdline). +declare -a ENV_FORWARD=(-e HF_HUB_ENABLE_HF_TRANSFER=1) +[[ -n "${HF_TOKEN:-}" ]] && ENV_FORWARD+=(-e HF_TOKEN) +[[ -n "${WANDB_API_KEY:-}" ]] && ENV_FORWARD+=(-e WANDB_API_KEY) +[[ -n "${UNSLOTH_LICENSE:-}" ]] && ENV_FORWARD+=(-e UNSLOTH_LICENSE) +[[ -n "${UNSLOTH_ALLOW_CPU:-}" ]] && ENV_FORWARD+=(-e UNSLOTH_ALLOW_CPU) +# Studio/Jupyter service config read by studio_launch.sh. Dash-only -e VAR so +# JUPYTER_PASSWORD never lands in argv. Without these the launcher gets a random +# password and no sshd/tunnel. +[[ -n "${JUPYTER_PASSWORD:-}" ]] && ENV_FORWARD+=(-e JUPYTER_PASSWORD) +[[ -n "${PUBLIC_KEY:-}" ]] && ENV_FORWARD+=(-e PUBLIC_KEY) +[[ -n "${SSH_KEY:-}" ]] && ENV_FORWARD+=(-e SSH_KEY) +[[ -n "${UNSLOTH_JUPYTER_CLOUDFLARE:-}" ]] && ENV_FORWARD+=(-e UNSLOTH_JUPYTER_CLOUDFLARE) + +# Extra publish flags for the service ports (Studio 8000, Jupyter 8888). +declare -a PORT_FLAGS=() +if [[ -n "${UNSLOTH_PORTS:-}" ]]; then + # shellcheck disable=SC2206 # intentional word splitting of "-p X -p Y" + PORT_FLAGS=(${UNSLOTH_PORTS}) +fi + +# Only attach -t when our own stdin/stdout are a TTY; CI / piped invocations +# otherwise hit `the input device is not a TTY` and never reach the entrypoint. +TTY_FLAG=() +if [ -t 0 ] && [ -t 1 ]; then + TTY_FLAG=(-it) +fi + +# No `set -x` here: it would echo HF_TOKEN / WANDB_API_KEY / UNSLOTH_LICENSE to +# CI logs. The ${arr[@]+"${arr[@]}"} form keeps empty arrays nounset-safe on +# bash 3.2 (macOS), where a bare "${empty[@]}" trips set -u. +exec docker run --rm ${TTY_FLAG[@]+"${TTY_FLAG[@]}"} \ + ${GPU_FLAG[@]+"${GPU_FLAG[@]}"} \ + --ipc=host \ + --ulimit memlock=-1 \ + --ulimit stack=67108864 \ + -v "$HF_CACHE":/workspace/.cache/huggingface \ + -v "$TRITON_CACHE":/workspace/.cache/triton \ + -v "$WORK_DIR":/workspace/host \ + "${ENV_FORWARD[@]}" \ + ${PORT_FLAGS[@]+"${PORT_FLAGS[@]}"} \ + "$IMAGE" "$@" diff --git a/docker/smoke_test.py b/docker/smoke_test.py new file mode 100644 index 0000000000..b763b53527 --- /dev/null +++ b/docker/smoke_test.py @@ -0,0 +1,171 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +""" +Smoke test for the unsloth-blackwell image. + +What this checks (in order, fail-fast): + 1. torch sees the GPU and the arch list contains sm_100 + sm_120. + 2. The runtime device's compute capability is supported. + 3. xformers / bitsandbytes / triton import without ImportError. + 4. unsloth imports and exposes FastLanguageModel. + 5. A 5-step LoRA train on a tiny model actually runs forward + backward. + +Run inside the container: + docker run --rm --gpus all unsloth-blackwell:latest python /workspace/smoke_test.py + +Skip step 5 (faster, no model download): + docker run --rm --gpus all unsloth-blackwell:latest python /workspace/smoke_test.py --skip-train +""" + +from __future__ import annotations + +import argparse +import sys + + +def banner(title: str) -> None: + print(f"\n=== {title} ===", flush = True) + + +def check_torch() -> tuple[int, int]: + banner("torch + arch list") + import torch + + # Raw C++ accessor works even without CUDA (partial smoke test on no-GPU host). + arches = torch._C._cuda_getArchFlags().split() + print(f"torch {torch.__version__}") + print(f"cuda build {torch.version.cuda}") + print(f"arches {arches}") + assert "sm_100" in arches, f"sm_100 missing: {arches}" + assert "sm_120" in arches, f"sm_120 missing: {arches}" + + assert torch.cuda.is_available(), "CUDA not visible -- did you pass --gpus all?" + cap = torch.cuda.get_device_capability(0) + name = torch.cuda.get_device_name(0) + print(f"device 0 {name} sm_{cap[0]}{cap[1]}") + # cu128 wheels ship SASS down to sm_75 (Turing); match the entrypoint floor so + # a Turing-only runner doesn't false-fail (Turing falls back to fp16). + if cap[0] < 7 or (cap[0] == 7 and cap[1] < 5): + sys.exit(f"FAIL: pre-Turing GPU {name} is not supported by this image") + if cap[0] < 8: + print(f"NOTE: {name} is Turing (sm_{cap[0]}{cap[1]}) -- bf16 unavailable, fp16 fallback.") + return cap + + +def check_imports() -> None: + banner("dep imports") + import triton + + print(f"triton {triton.__version__}") + # Import order matters: unsloth before transformers/trl/peft (so its patches + # land) and before unsloth_zoo (which needs the UNSLOTH_IS_PRESENT marker). + import unsloth + + print(f"unsloth {unsloth.__version__}") + import unsloth_zoo + + print(f"unsloth_zoo {unsloth_zoo.__version__}") + # xformers has no aarch64 cu128 wheel; arm64 omits it. Best-effort so one + # script covers both arches. + try: + import xformers + print(f"xformers {xformers.__version__}") + except ImportError: + print("xformers (missing -- expected on arm64 [huggingface] extras)") + import bitsandbytes as bnb + + print(f"bnb {bnb.__version__}") + import transformers + + print(f"transformers {transformers.__version__}") + import trl + + print(f"trl {trl.__version__}") + import peft + + print(f"peft {peft.__version__}") + + +def check_unsloth_import() -> None: + banner("unsloth FastLanguageModel reachable") + # Already imported in check_imports(); this re-import is a no-op. + import unsloth + from unsloth import FastLanguageModel + + print(f"unsloth {unsloth.__version__}") + print(f"FastLanguageModel {FastLanguageModel}") + + +def check_tiny_train(cap: tuple[int, int]) -> None: + banner("tiny LoRA train (5 steps)") + import os + + # Unsloth must be imported first. + import unsloth # noqa: F401 + from unsloth import FastLanguageModel + import torch + + # Small, public, no-gate. + model_name = "unsloth/Llama-3.2-1B-Instruct-bnb-4bit" + print(f"loading {model_name}") + model, tokenizer = FastLanguageModel.from_pretrained( + model_name = model_name, + max_seq_length = 512, + dtype = None, + load_in_4bit = True, + ) + model = FastLanguageModel.get_peft_model( + model, + r = 8, + lora_alpha = 16, + target_modules = ["q_proj", "k_proj", "v_proj", "o_proj"], + lora_dropout = 0.0, + bias = "none", + use_gradient_checkpointing = "unsloth", + random_state = 0, + ) + + prompts = [ + "Q: What is the capital of France?\nA:", + "Q: 2 + 2 = ?\nA:", + "Q: Name a primary color.\nA:", + "Q: Hello, who are you?\nA:", + ] * 2 + enc = tokenizer(prompts, return_tensors = "pt", padding = True, truncation = True, max_length = 64) + enc = {k: v.cuda() for k, v in enc.items()} + labels = enc["input_ids"].clone() + + model.train() + optim = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr = 1e-4) + for step in range(5): + out = model(**enc, labels = labels) + out.loss.backward() + optim.step() + optim.zero_grad(set_to_none = True) + print(f"step {step} loss={out.loss.item():.4f}", flush = True) + + print("OK: 5 LoRA steps completed") + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument( + "--skip-train", + action = "store_true", + help = "Skip the tiny LoRA training step (no HF download).", + ) + args = ap.parse_args() + + cap = check_torch() + check_imports() + check_unsloth_import() + if not args.skip_train: + check_tiny_train(cap) + + banner("all checks passed") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/docker/studio_launch.sh b/docker/studio_launch.sh new file mode 100644 index 0000000000..a2f0d7dc3b --- /dev/null +++ b/docker/studio_launch.sh @@ -0,0 +1,118 @@ +#!/usr/bin/env bash +# Default CMD of the full Unsloth image (Dockerfile.studio). +# +# Bootstraps the three services managed by supervisord: +# studio port 8000 first-boot admin password printed in `docker logs` +# jupyter port 8888 password from JUPYTER_PASSWORD, or a random one +# printed in `docker logs` when unset +# sshd port 22 key-only; enabled when PUBLIC_KEY / SSH_KEY is set +# +# Environment: +# JUPYTER_PORT Jupyter port inside the container (default 8888) +# JUPYTER_PASSWORD Jupyter login password (unset: generated and printed) +# PUBLIC_KEY/SSH_KEY OpenSSH public key for root login; sshd stays disabled +# when neither is set (nothing to authenticate with -- +# password login is never enabled for root) +set -euo pipefail + +export JUPYTER_PORT="${JUPYTER_PORT:-8888}" +export UNSLOTH_STUDIO_HOME="${UNSLOTH_STUDIO_HOME:-/opt/unsloth-studio}" +# Default off so supervisord's %(ENV_UNSLOTH_JUPYTER_CLOUDFLARE)s autostart gate +# resolves; set to 1 (docker run -e) to expose JupyterLab on a trycloudflare URL. +export UNSLOTH_JUPYTER_CLOUDFLARE="${UNSLOTH_JUPYTER_CLOUDFLARE:-0}" + +# Make the runtime env visible to SSH login shells (which lack the `docker run -e` +# vars). Secrets are excluded on purpose -- they stay in process env, never on +# disk. shlex.quote() each value since this file is sourced by every login shell. +python - > /etc/profile.d/unsloth_env.sh <<'PY' || true +import os, re, shlex +keep = re.compile(r"^(HF_|CUDA_|NCCL_|JUPYTER_|UNSLOTH_|WANDB_|TRITON_)|^PATH$") +secret = re.compile(r"(_TOKEN|_API_KEY|_PASSWORD|_SECRET|_LICENSE)$") +for key, value in sorted(os.environ.items()): + if keep.search(key) and not secret.search(key): + print(f"export {key}={shlex.quote(value)}") +PY + +# Hash the Jupyter password with jupyter's helper; never store plaintext. No fixed +# default: when JUPYTER_PASSWORD is unset, generate a random one and print it once. +JUPYTER_CONFIG_DIR=/root/.jupyter +JUPYTER_NOTE="password from JUPYTER_PASSWORD env" +if [[ -f "${JUPYTER_CONFIG_DIR}/jupyter_lab_config.py" ]]; then + JUPYTER_NOTE="existing jupyter config reused" +else + if [[ -z "${JUPYTER_PASSWORD:-}" ]]; then + JUPYTER_PASSWORD="$(python -c 'import secrets; print(secrets.token_urlsafe(12))')" + JUPYTER_NOTE="generated password: ${JUPYTER_PASSWORD}" + fi + export JUPYTER_PASSWORD + mkdir -p "${JUPYTER_CONFIG_DIR}" + HASH=$(python - < "${JUPYTER_CONFIG_DIR}/jupyter_lab_config.py" <> "${JUPYTER_CONFIG_DIR}/jupyter_lab_config.py" </dev/null 2>&1; then + mkdir -p /root/.ssh && chmod 700 /root/.ssh + echo "${PUBLIC_SSH_KEY}" > /root/.ssh/authorized_keys + chmod 600 /root/.ssh/authorized_keys + ssh-keygen -A + mkdir -p /run/sshd + export UNSLOTH_ENABLE_SSHD=true +fi + +mkdir -p /workspace + +# This image ships under the GNU AGPLv3. Refuse to start if the Unsloth +# attribution (Help/About, splash, login, theme, AGPLv3 license + source links) +# is stripped or altered. The same checker runs as a jupyter_server extension and +# at build time. Bypass for local dev: UNSLOTH_SKIP_BRANDING_CHECK=1 (not resale). +if [[ "${UNSLOTH_SKIP_BRANDING_CHECK:-0}" != "1" ]]; then + if ! /opt/unsloth-venv/bin/python -m unsloth_branding --verify; then + echo "Refusing to start the container." >&2 + exit 1 + fi +fi + +echo "Unsloth Studio -> http://localhost:8000 (first-boot password below)" +echo "JupyterLab -> http://localhost:${JUPYTER_PORT} (${JUPYTER_NOTE})" +if [[ "${UNSLOTH_JUPYTER_CLOUDFLARE}" == "1" ]]; then + echo "JupyterLab tunnel-> enabled; public trycloudflare URL appears below once it is up" +else + echo "JupyterLab tunnel-> off (set UNSLOTH_JUPYTER_CLOUDFLARE=1 for a public link)" +fi +if [[ "${UNSLOTH_ENABLE_SSHD}" == "true" ]]; then + echo "sshd -> port 22 (key-only)" +fi + +exec supervisord -c /etc/supervisor/supervisord.conf diff --git a/docker/supervisord.conf b/docker/supervisord.conf new file mode 100644 index 0000000000..d2be57fe33 --- /dev/null +++ b/docker/supervisord.conf @@ -0,0 +1,76 @@ +# Service manager for the full Unsloth image (Dockerfile.studio). +# +# Mirrors the service set of the production docker.io/unsloth/unsloth image: +# studio Unsloth Studio web UI port 8000 +# jupyter JupyterLab for the notebooks port $JUPYTER_PORT (default 8888) +# sshd key-only SSH for cloud hosts port 22 +# +# All three log to stdout/stderr so `docker logs` shows everything, including +# Studio's first-boot password and Jupyter's startup line. + +[unix_http_server] +file=/run/supervisor.sock +chmod=0700 + +[supervisorctl] +serverurl=unix:///run/supervisor.sock + +[rpcinterface:supervisor] +supervisor.rpcinterface_factory = supervisor.rpcinterface:make_main_rpcinterface + +[supervisord] +nodaemon=true +pidfile=/run/supervisord.pid +logfile=/dev/null +logfile_maxbytes=0 +loglevel=info + +[program:studio] +command=%(ENV_UNSLOTH_STUDIO_HOME)s/bin/unsloth studio -H 0.0.0.0 -p 8000 +directory=/workspace +autostart=true +autorestart=true +startretries=3 +startsecs=5 +environment=HOME="/root",USER="root" +stdout_logfile=/dev/stdout +stdout_logfile_maxbytes=0 +stderr_logfile=/dev/stderr +stderr_logfile_maxbytes=0 + +[program:jupyter] +command=jupyter lab --no-browser --ip=0.0.0.0 --port=%(ENV_JUPYTER_PORT)s --allow-root --notebook-dir=/workspace +directory=/workspace +autostart=true +autorestart=true +; HOME pins config lookup to /root/.jupyter (where the launcher wrote the +; password config); without it an unset HOME falls back to token auth. +environment=HOME="/root",USER="root" +stdout_logfile=/dev/stdout +stdout_logfile_maxbytes=0 +stderr_logfile=/dev/stderr +stderr_logfile_maxbytes=0 + +; Optional public Cloudflare quick-tunnel for JupyterLab. Started only when +; UNSLOTH_JUPYTER_CLOUDFLARE=1 (studio_launch.sh exports a 0 default so this +; expands). The trycloudflare URL is printed to docker logs by cloudflared. +[program:jupyter-cloudflare] +command=/usr/local/bin/unsloth-jupyter-tunnel +directory=/workspace +autostart=%(ENV_UNSLOTH_JUPYTER_CLOUDFLARE)s +autorestart=true +startsecs=5 +environment=HOME="/root",USER="root" +stdout_logfile=/dev/stdout +stdout_logfile_maxbytes=0 +stderr_logfile=/dev/stderr +stderr_logfile_maxbytes=0 + +[program:sshd] +command=/usr/sbin/sshd -D -e +autostart=%(ENV_UNSLOTH_ENABLE_SSHD)s +autorestart=true +stdout_logfile=/dev/stdout +stdout_logfile_maxbytes=0 +stderr_logfile=/dev/stderr +stderr_logfile_maxbytes=0 diff --git a/docker/unsloth_colab_compat.py b/docker/unsloth_colab_compat.py new file mode 100644 index 0000000000..cb35a5088d --- /dev/null +++ b/docker/unsloth_colab_compat.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +"""Colab cell-magic compatibility for the Unsloth Docker notebooks. + +Colab cells often look like: + + #@title Colab Extra Install { display-mode: "form" } + %%capture + !pip install ... + +In IPython a cell magic (`%%capture`, `%%bash`, ...) is only recognised when it +is the VERY FIRST line of the cell. A leading Colab `#@title`/`#@param` form (or +any comment/blank line) pushes the `%%magic` to line 2, so IPython treats it as a +line magic and raises `UsageError: Line magic function `%%capture` not found.` +and the cell fails. + +Fix: register an `input_transformers_cleanup` (runs before magic detection) that +hoists a `%%` cell magic above any leading blank/comment (`#...`, incl. `#@...`) +lines, so the magic lands on line 0 and fires normally. The skipped comment lines +stay in the cell (still inert), just below the magic -- so `%%capture` now also +captures them. Idempotent and fully guarded: any problem returns the input +unchanged, so a cell never breaks because of this helper. + +The hoist is restricted to cell magics whose body is executed as code (Python or +shell), where a moved-down `#@title`/comment line stays an inert comment. Magics +that treat the body as literal content (`%%writefile`, `%%file`, `%%html`, +`%%javascript`, `%%latex`, `%%markdown`, `%%svg`, ...) are left untouched: moving +the Colab form comment into their body would write/render it and corrupt the +generated file or output. + +This mirrors unsloth_nb_compat.register_ipython(): it is wired from the baked +IPython startup file (docker/unsloth_ipython_startup.py). +""" + +from __future__ import annotations +import sys + + +# Cell magics whose body runs as code, so a hoisted comment stays inert. Only +# these; content/data magics (%%writefile, %%html, ...) untouched (see docstring). +_SAFE_CELL_MAGICS = frozenset( + { + "capture", # Colab install pattern: suppress pip output + "time", + "timeit", + "prun", + "debug", + "bash", + "sh", + "shell", + "python", + "python2", + "python3", + "pypy", + } +) + + +def colab_cell_magic_fix(lines): + """Hoist a safe `%%` cell magic above leading blank/comment lines. + + `lines` is the IPython cell as a list of strings (each ending in '\\n'). + Returns a (possibly reordered) list of the same lines. + """ + try: + skipped = [] + for i, line in enumerate(lines): + stripped = line.strip() + if stripped == "" or stripped.startswith("#"): + skipped.append(line) # blank or comment (incl. #@title) + continue + # First real line. Act only if it's a cell magic not already on top. + if stripped.startswith("%%") and i > 0: + name = stripped[2:].split(maxsplit = 1) + name = name[0] if name else "" + if name in _SAFE_CELL_MAGICS: + return [line] + skipped + lines[i + 1 :] + # Content/data magic: don't move the comment into its body. + return lines + return lines # already on top, or not a magic + return lines # all blank/comment -> nothing to do + except Exception: + return lines + + +def register_ipython(): + """Append the transformer to the running IPython (called from startup).""" + try: + ip = get_ipython() # noqa: F821 (provided by IPython) + except NameError: + return + if ip is None or getattr(ip, "_unsloth_colab_fix", False): + return + try: + ip.input_transformers_cleanup.append(colab_cell_magic_fix) + ip._unsloth_colab_fix = True + except Exception as e: # never break a kernel because of the helper + print(f"[unsloth-nb] colab-compat hook skipped: {e!r}", file = sys.stderr) diff --git a/docker/unsloth_ipython_startup.py b/docker/unsloth_ipython_startup.py new file mode 100644 index 0000000000..9939b2ece3 --- /dev/null +++ b/docker/unsloth_ipython_startup.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +"""Baked IPython startup hook (copied to the profile's startup/ dir). + +Runs once per kernel. Registers a pre_run_cell event that activates the right +transformers sidecar before the first model cell, using the version the +notebook's own install cell asked for (recorded by the pip/uv shim). Safe no-op +outside IPython, when no version was requested, or once transformers is imported. +""" + +try: + import os + + # Tell the pip/uv shim it's inside a notebook kernel, so a cell's + # `!pip install ...` gets safe-install behaviour. Unset elsewhere => passthrough. + os.environ["UNSLOTH_NB_SHIM"] = "1" + + # Scope the transformers-request marker to THIS kernel so concurrent notebooks + # don't read each other's pin. The shim (a child) inherits UNSLOTH_NB_TF_MARKER, + # so writer and reader agree. Unset => shared default (one notebook/process). + if not os.environ.get("UNSLOTH_NB_TF_MARKER"): + # Stable, unique kernel id: the ipykernel connection file name, else the PID. + _kid = "" + try: + from ipykernel import get_connection_file # type: ignore + _kid = os.path.splitext(os.path.basename(get_connection_file()))[0] + except Exception: + _kid = "" + _kid = _kid or ("pid-%d" % os.getpid()) + os.environ["UNSLOTH_NB_TF_MARKER"] = "/tmp/unsloth_nb/requested_transformers." + _kid + + import unsloth_nb_compat + + unsloth_nb_compat.register_ipython() + + # Re-point %pip / %uv and `!python -m pip` at the same shim so in-process + # installs can't bypass it and overwrite the baked torch/vLLM stack. + import unsloth_nb_pip_magic + + unsloth_nb_pip_magic.register_ipython() +except Exception as _e: # never break a kernel because of the helper + import sys + print(f"[unsloth-nb] startup hook skipped: {_e!r}", file = sys.stderr) + +# Colab cell-magic compatibility (hoist `%%capture` above a leading `#@title` +# form). Separate try/except so it can't disable the hook above, or vice versa. +try: + import unsloth_colab_compat + unsloth_colab_compat.register_ipython() +except Exception as _e: # never break a kernel because of the helper + import sys + print(f"[unsloth-nb] colab-compat hook skipped: {_e!r}", file = sys.stderr) diff --git a/docker/unsloth_jupyter_tunnel.sh b/docker/unsloth_jupyter_tunnel.sh new file mode 100755 index 0000000000..d30218412f --- /dev/null +++ b/docker/unsloth_jupyter_tunnel.sh @@ -0,0 +1,60 @@ +#!/usr/bin/env bash +# Optional public Cloudflare quick-tunnel for JupyterLab, mirroring the tunnel +# Studio creates for its own UI. Off by default. Two ways to use it: +# +# * at run time: docker run -e UNSLOTH_JUPYTER_CLOUDFLARE=1 ... unsloth/unsloth +# -> the https://.trycloudflare.com URL is printed in +# `docker logs` once JupyterLab is up. +# * on demand: docker exec unsloth-jupyter-tunnel --force +# +# The tunnel gives a public https URL that works from anywhere with no account +# or open inbound port. JupyterLab still requires its password, so the notebook +# is not open to the world; treat the URL as sensitive all the same. +set -u + +FORCE=0 +[ "${1:-}" = "--force" ] && FORCE=1 +if [ "$FORCE" != "1" ] && [ "${UNSLOTH_JUPYTER_CLOUDFLARE:-0}" != "1" ]; then + echo "[jupyter-tunnel] disabled (set UNSLOTH_JUPYTER_CLOUDFLARE=1, or run with --force)" + exit 0 +fi + +PORT="${JUPYTER_PORT:-8888}" + +echo "[jupyter-tunnel] waiting for JupyterLab on port ${PORT} ..." +for _ in $(seq 1 90); do + if curl -fsS -o /dev/null "http://localhost:${PORT}/login" 2>/dev/null; then + break + fi + sleep 2 +done + +# Reuse a cloudflared already on the host (Studio caches one for its own +# tunnel); otherwise fetch the static binary for this arch. No account needed. +CFD="" +for cand in \ + "${UNSLOTH_STUDIO_HOME:-/opt/unsloth-studio}/bin/cloudflared" \ + /usr/local/bin/cloudflared \ + cloudflared; do + if command -v "$cand" >/dev/null 2>&1; then CFD="$(command -v "$cand")"; break; fi + [ -x "$cand" ] && { CFD="$cand"; break; } +done +if [ -z "$CFD" ]; then + case "$(uname -m)" in + x86_64|amd64) A=amd64;; + aarch64|arm64) A=arm64;; + *) A=amd64;; + esac + CFD=/usr/local/bin/cloudflared + echo "[jupyter-tunnel] downloading cloudflared (${A}) ..." + if ! curl -fsSL -o "$CFD" \ + "https://github.com/cloudflare/cloudflared/releases/latest/download/cloudflared-linux-${A}"; then + echo "[jupyter-tunnel] could not download cloudflared" >&2 + exit 1 + fi + chmod +x "$CFD" +fi + +echo "[jupyter-tunnel] starting Cloudflare quick-tunnel to JupyterLab (port ${PORT})." +echo "[jupyter-tunnel] the https://.trycloudflare.com URL appears below; log in with your Jupyter password." +exec "$CFD" tunnel --no-autoupdate --url "http://localhost:${PORT}" diff --git a/docker/unsloth_llama_update.sh b/docker/unsloth_llama_update.sh new file mode 100755 index 0000000000..8ab0e1f068 --- /dev/null +++ b/docker/unsloth_llama_update.sh @@ -0,0 +1,222 @@ +#!/usr/bin/env bash +# Update the baked llama.cpp prebuilt in place, inside a running container, +# without pulling a new image. Downloads the newest portable llama.cpp bundle +# (the same target-pinned, sha256-verified bundle the image is built with) and +# atomically swaps it into $UNSLOTH_LLAMA_CPP_PATH, so the next GGUF export / +# model load uses it. +# +# docker exec unsloth-llama-update # latest release +# docker exec unsloth-llama-update --tag b9773-mix-1f1aaa4 +# docker exec unsloth-llama-update --check # report only, no download +# +# This reuses the build-time fetcher, which resolves the latest release via the +# GitHub /releases/latest redirect (no API token, not rate-limited) and installs +# the portable CUDA bundle that runs on CPU and every supported GPU. That makes +# it work the same in a CPU-only or a --gpus container, unlike the host-probing +# installer behind the in-app banner. +# +# Persistence: unmounted, the swap lands in the container's writable layer +# (survives docker restart). To keep it across a full recreate, mount the dir +# on a named volume (-v unsloth_llama:/opt/unsloth/llama.cpp); the updater +# detects the mount and swaps the bundle contents inside the volume. +set -euo pipefail + +INSTALL_DIR="${UNSLOTH_LLAMA_CPP_PATH:-/opt/unsloth/llama.cpp}" +STUDIO_HOME="${UNSLOTH_STUDIO_HOME:-/opt/unsloth-studio}" +FETCHER="${UNSLOTH_LLAMA_FETCHER:-/usr/local/lib/unsloth/fetch_llama_prebuilt.py}" +REPO="unslothai/llama.cpp" +TAG="latest" +CHECK_ONLY=0 + +usage() { sed -n '2,21p' "$0"; } + +while [ $# -gt 0 ]; do + case "$1" in + --tag) TAG="$2"; shift 2;; + --install-dir) INSTALL_DIR="$2"; shift 2;; + --check) CHECK_ONLY=1; shift;; + -h|--help) usage; exit 0;; + *) echo "unsloth-llama-update: unknown argument: $1" >&2; usage; exit 2;; + esac +done + +[ -f "$FETCHER" ] || { echo "unsloth-llama-update: fetcher not found at $FETCHER" >&2; exit 1; } + +# Any python works (the fetcher is stdlib-only); prefer the Studio venv, then base. +PY="" +for cand in \ + "$STUDIO_HOME/unsloth_studio/bin/python" \ + /opt/unsloth-venv/bin/python \ + python3 python; do + command -v "$cand" >/dev/null 2>&1 && { PY="$cand"; break; } + [ -x "$cand" ] && { PY="$cand"; break; } +done +[ -n "$PY" ] || { echo "unsloth-llama-update: no python found" >&2; exit 1; } + +# amd64 -> linux-x64-cuda12 portable; arm64 -> linux-arm64-cuda13 portable. +case "$(uname -m)" in + x86_64|amd64) ARCH="amd64";; + aarch64|arm64) ARCH="arm64";; + *) echo "unsloth-llama-update: unsupported arch $(uname -m)" >&2; exit 1;; +esac + +installed_tag() { + "$PY" - "$INSTALL_DIR" <<'PY' 2>/dev/null || echo "unknown" +import json, os, sys +p = os.path.join(sys.argv[1], "UNSLOTH_PREBUILT_INFO.json") +try: + d = json.load(open(p)); print(d.get("tag") or d.get("release_tag") or d.get("upstream_tag") or "unknown") +except Exception: + print("unknown") +PY +} + +resolve_latest() { + "$PY" - "$FETCHER" "$REPO" <<'PY' 2>/dev/null || echo "" +import importlib.util, sys +spec = importlib.util.spec_from_file_location("flp", sys.argv[1]) +m = importlib.util.module_from_spec(spec); spec.loader.exec_module(m) +print(m.resolve_latest_tag(sys.argv[2])) +PY +} + +CUR="$(installed_tag)" +echo "[llama-update] install dir: $INSTALL_DIR" +echo "[llama-update] installed: $CUR" + +if [ "$CHECK_ONLY" = "1" ]; then + LATEST="$(resolve_latest)" + echo "[llama-update] latest: ${LATEST:-unknown}" + # resolve_latest swallows every failure into "" (line 75), so an empty value + # means the lookup did not happen -- no network, proxy, GitHub down. Printing + # "up to date" there is the one answer --check must never give: it reports a + # state it could not observe. Say unknown and exit non-zero instead. + if [ -z "$LATEST" ]; then + echo "[llama-update] could not reach the release feed; update status UNKNOWN" >&2 + echo "[llama-update] (retry once the container has network access)" >&2 + exit 1 + fi + if [ "$LATEST" != "$CUR" ]; then + echo "[llama-update] an update is available (run without --check to apply)" + else + echo "[llama-update] up to date" + fi + exit 0 +fi + +# Fetch into a sibling temp dir (same filesystem as INSTALL_DIR, so the swap is +# an atomic rename), then swap. On any failure the existing install is untouched. +parent="$(dirname "$INSTALL_DIR")" + +# A named volume mounted AT the install dir can't be renamed (EBUSY), so the +# whole-dir swap below would fail; detect the mount and swap the CONTENTS inside +# the tree. UNSLOTH_LLAMA_UPDATE_IN_PLACE=1/0 overrides autodetection. +IN_PLACE="${UNSLOTH_LLAMA_UPDATE_IN_PLACE:-}" +if [ -z "$IN_PLACE" ]; then + IN_PLACE=0 + if command -v mountpoint >/dev/null 2>&1 && mountpoint -q "$INSTALL_DIR" 2>/dev/null; then + IN_PLACE=1 + elif [ "$(stat -c %d "$INSTALL_DIR" 2>/dev/null)" != "$(stat -c %d "$parent" 2>/dev/null)" ]; then + IN_PLACE=1 # filesystem boundary at the dir = a volume without mountpoint(1) + fi +fi +if [ "$IN_PLACE" = "1" ]; then + # Keep every move inside the mounted filesystem: work + backup live UNDER + # the install dir so each swap step is a same-fs rename within the volume. + work="$(mktemp -d "$INSTALL_DIR/.llamaupd.XXXXXX")" + backup="$INSTALL_DIR/.old.$$" +else + work="$(mktemp -d "$parent/.llamaupd.XXXXXX")" + backup="${INSTALL_DIR}.old.$$" +fi +swap_done=0 +drained=0 +# The exit handler must never delete $backup while it's the ONLY copy: restore the +# old tree first, remove it only after the new tree is active. Signal traps run +# the EXIT trap on HUP/INT/TERM too. +cleanup() { + if [ "$swap_done" -ne 1 ]; then + if [ "$IN_PLACE" = "1" ]; then + # Contents-swap restore. Every old entry lives in exactly one of + # $backup / $INSTALL_DIR, so a same-named entry in the install dir is a + # half-moved NEW one: drop it, then move the old one back. + if [ -d "$backup" ]; then + _restore_fail=0 + # The per-name loop below only sees entries the OLD tree had, so a + # file the new release introduced survives it and the "restored" + # dir ends up mixed-version -- ggml dlopens every libggml-*.so it + # finds next to the binaries. Once the drain finished, every + # remaining entry is a half-moved NEW one, so clear them all. + # Gated on "drained": before the drain completes an entry here can + # still be the ONLY copy of an old one, and deleting it loses data. + if [ "$drained" = "1" ]; then + find "$INSTALL_DIR" -mindepth 1 -maxdepth 1 \ + ! -path "$work" ! -path "$backup" \ + -exec rm -rf {} + 2>/dev/null || true + fi + for _e in "$backup"/* "$backup"/.[!.]* "$backup"/..?*; do + { [ -e "$_e" ] || [ -L "$_e" ]; } || continue + _b="$(basename "$_e")" + if [ -e "$INSTALL_DIR/$_b" ] || [ -L "$INSTALL_DIR/$_b" ]; then + rm -rf "${INSTALL_DIR:?}/$_b" 2>/dev/null || true + fi + mv "$_e" "$INSTALL_DIR/" 2>/dev/null || _restore_fail=1 + done + if [ "$_restore_fail" -eq 0 ]; then + rmdir "$backup" 2>/dev/null || true + else + echo "[llama-update] CRITICAL: restore failed; previous install preserved at $backup" >&2 + fi + fi + elif [ ! -e "$INSTALL_DIR" ] && [ -e "$backup" ]; then + if ! mv "$backup" "$INSTALL_DIR" 2>/dev/null; then + echo "[llama-update] CRITICAL: restore failed; previous install preserved at $backup" >&2 + fi + fi + fi + rm -rf "$work" 2>/dev/null || true + if [ "$swap_done" = "1" ]; then + rm -rf "$backup" 2>/dev/null || true + fi +} +trap cleanup EXIT +trap 'exit 129' HUP +trap 'exit 130' INT +trap 'exit 143' TERM +new="$work/llama.cpp" + +echo "[llama-update] fetching llama.cpp '$TAG' ($ARCH portable) ..." +"$PY" "$FETCHER" "$TAG" "$ARCH" "$new" + +# Preserve the Studio ownership marker so setup.sh keeps recognising the dir. +[ -e "$INSTALL_DIR/.unsloth-studio-owned" ] && touch "$new/.unsloth-studio-owned" + +echo "[llama-update] swapping into place ..." +if [ "$IN_PLACE" = "1" ]; then + # The install dir is a mount point: swap its CONTENTS (all same-fs renames + # inside the volume). The trap's contents-restore covers any mid-swap abort. + mkdir "$backup" + find "$INSTALL_DIR" -mindepth 1 -maxdepth 1 \ + ! -path "$work" ! -path "$backup" -exec mv -t "$backup" {} + + # Every old entry now lives in $backup, so from here the trap may clear the + # install dir before restoring. set -e means a failed drain never gets here. + drained=1 + if find "$new" -mindepth 1 -maxdepth 1 -exec mv -t "$INSTALL_DIR" {} +; then + swap_done=1 + else + echo "[llama-update] swap failed; restoring previous install" >&2 + exit 1 + fi +else + mv "$INSTALL_DIR" "$backup" + if mv "$new" "$INSTALL_DIR"; then + swap_done=1 + else + echo "[llama-update] swap failed; restoring previous install" >&2 + mv "$backup" "$INSTALL_DIR" + exit 1 + fi +fi + +echo "[llama-update] installed now: $(installed_tag)" +echo "[llama-update] done (reload your model / re-run export to use it)" diff --git a/docker/unsloth_nb_compat.py b/docker/unsloth_nb_compat.py new file mode 100644 index 0000000000..36cd1266c5 --- /dev/null +++ b/docker/unsloth_nb_compat.py @@ -0,0 +1,247 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +"""Per-notebook transformers version activation for the Unsloth Docker image. + +Problem: unslothai/notebooks pin many different transformers versions in their +install cells (transformers==4.56.2 on ~115, 5.5.0/5.3.0/5.10.x on newer model +families). The baked base venv ships ONE transformers (latest 5.x). Running an +old-model notebook against it, or letting the install cell pip-install a pinned +version on top, either breaks the model or clobbers the cu128 torch/vLLM stack. + +Solution (mirrors Unsloth Studio's studio/backend/utils/transformers_version.py): +keep the base venv intact and ship coherent transformers "sidecars" -- each is a +`pip install --target --no-deps transformers==X` plus the matched +huggingface_hub/tokenizers/safetensors. To use version X we just prepend its +sidecar dir to sys.path BEFORE transformers is imported; the rest of the stack +(torch, vllm, unsloth, peft, trl) comes from the base venv unchanged. + +That "rest of the stack" is the catch, and it is why selection has a FLOOR as +well as a ceiling (see sidecar_for): vLLM is version-locked to transformers, so a +sidecar older than what the baked vLLM accepts does not give the notebook an +older transformers, it gives it an ImportError at `import unsloth`. The image +therefore only ships sidecars whose vLLM import has been verified at build time, +and records the lowest of them as the floor. + +Two activation paths: + * driven/headless: `unsloth-run ` sets PYTHONPATH at kernel launch. + * manual JupyterLab: an IPython pre_run_cell hook (registered by the baked + startup file) activates the sidecar before the first model cell, using the + version the notebook's own install cell asked for (recorded by the pip shim). +""" + +from __future__ import annotations +import os, sys, glob, json + +SIDECAR_ROOT = os.environ.get("UNSLOTH_TF_SIDECAR_ROOT", "/opt/unsloth-venv/tf-sidecars") +# The pip/uv shim writes the transformers version a notebook asked for here. +MARKER = os.environ.get("UNSLOTH_NB_TF_MARKER", "/tmp/unsloth_nb/requested_transformers") + +# Lowest transformers the image's baked vLLM can import. A sidecar below this is +# not "an older transformers", it is a BROKEN image: `import unsloth` dies before +# the first model cell. Written by the Dockerfile's sidecar verification step +# (which imports vllm.transformers_utils.config under every candidate and drops +# the ones that raise), so it tracks whatever vLLM the image actually bakes +# instead of a literal that rots on the next bump. Measured on vLLM 0.26.0: +# +# transformers 4.57.6 FAIL "Support for Transformers v4 ... removed in vLLM v0.24.0" +# transformers 5.3.0 FAIL "cannot import name 'ALLOWED_LAYER_TYPES'" +# transformers 5.5.0 OK +# transformers 5.10.2 OK +# transformers 5.14.1 OK (the baked one, no sidecar) +FLOOR_FILE = os.path.join(SIDECAR_ROOT, ".vllm_min_transformers") + + +def _logging_enabled() -> bool: + """Sidecar activation is silent by default; users found the per-cell + `[unsloth-nb] activated transformers sidecar ...` line noisy. Set + UNSLOTH_ENABLE_LOGGING=1 to surface it (and other [unsloth-nb] diagnostics).""" + return os.environ.get("UNSLOTH_ENABLE_LOGGING", "").strip().lower() not in ( + "", + "0", + "false", + "no", + "off", + ) + + +# Model-name -> minimum transformers tier (substring match on the lowered id), +# ported from Studio. Fallback when a notebook names a new model but pins nothing. +_TIER_SUBSTRINGS = { + "5.10.2": ("gemma-4-12b", "gemma4-12b"), + "5.5.0": ("gemma-4", "gemma4", "qwen3.6"), + "5.3.0": ( + "ministral-3", + "glm-4.7-flash", + "qwen3-30b-a3b", + "qwen3.5", + "qwen3-next", + "qwen3_5", + "lfm2.5-vl", + ), +} + + +def _baked(): + """Return {version_str: dir} for every baked sidecar.""" + out = {} + for d in sorted(glob.glob(os.path.join(SIDECAR_ROOT, "t_*"))): + out[os.path.basename(d)[2:].replace("_", ".")] = d + return out + + +def min_version(): + """Lowest transformers this image's vLLM can import, or None if unrecorded. + + UNSLOTH_TF_SIDECAR_MIN overrides, so a hand-mounted sidecar root can declare + its own floor. Returns None when neither is set, which keeps the pre-floor + behaviour for any environment that never ran the build-time verification.""" + v = os.environ.get("UNSLOTH_TF_SIDECAR_MIN", "").strip() + if v: + return v + try: + with open(FLOOR_FILE) as f: + return f.read().strip() or None + except OSError: + return None + + +def _eligible(): + """Baked sidecars the floor allows, as a sorted [(Version, version_str, dir)]. + + Returns None when the versions cannot be parsed (no packaging available).""" + baked = _baked() + if not baked: + return [] + try: + from packaging.version import Version + except Exception: + return None + floor = min_version() + try: + low = Version(floor) if floor else None + except Exception: + low = None + rows = [] + for v, d in baked.items(): + try: + ver = Version(v) + except Exception: + continue + if low is not None and ver < low: + continue # vLLM cannot import it; activating it only breaks the run + rows.append((ver, v, d)) + rows.sort() + return rows + + +def tier_for_model(model_name: str): + """Best-effort minimum transformers version for a model id (or None).""" + if not model_name: + return None + low = model_name.lower() + # check newest tiers first so gemma-4-12b wins over gemma-4 + for ver in ("5.10.2", "5.5.0", "5.3.0"): + if any(s in low for s in _TIER_SUBSTRINGS[ver]): + return ver + return None + + +def sidecar_for(version: str): + """Map a requested/needed transformers version to a baked sidecar dir. + + FLOOR then CEILING, in that order: + + * floor -- a sidecar the baked vLLM cannot import is never eligible, no + matter what the notebook pinned. Selecting one used to break `import + unsloth` in 254 of the 433 shipped notebooks, because the two common pin + families (4.5x -> the 4.57.6 sidecar, 5.2/5.3 -> the 5.3.0 sidecar) both + landed on a sidecar vLLM 0.26.0 refuses. A request below the floor is + clamped UP to the lowest eligible sidecar: that is the closest version to + what the notebook asked for that this image can actually run. + * ceiling -- among the eligible sidecars pick the smallest >= the request, + because a model added in version X needs *at least* X. + + A request newer than every eligible sidecar returns None -> use the base venv + (the newest 5.x), which is always vLLM-compatible.""" + if not version: + return None + rows = _eligible() + if rows is None: # no packaging: only an exact, still-eligible match is safe + baked = _baked() + d = baked.get(version) + floor = min_version() + return d if (d and (not floor or version == floor)) else None + if not rows: + return None + for _ver, v, d in rows: + if v == version: + return d + try: + from packaging.version import Version + want = Version(version) + except Exception: + return None + for ver, _v, d in rows: + if ver >= want: + return d + return None + + +def requested_version(): + """transformers version a notebook asked for (recorded by the pip shim).""" + try: + with open(MARKER) as f: + v = f.read().strip() + return v or None + except OSError: + return None + + +def activate(version: str | None, *, quiet: bool = False): + """Prepend the matching sidecar to sys.path if transformers isn't imported yet. + + Returns the activated dir, or None if the base venv is used / activation is + no longer possible (transformers already imported).""" + if not version: + return None + d = sidecar_for(version) + if not d: + return None + if "transformers" in sys.modules: + if not quiet: + print( + f"[unsloth-nb] transformers already imported; cannot switch to " + f"{version} in-process (restart the kernel, or use `unsloth-run`).", + file = sys.stderr, + ) + return None + if d not in sys.path: + sys.path.insert(0, d) + os.environ["PYTHONPATH"] = d + os.pathsep + os.environ.get("PYTHONPATH", "") + if not quiet and _logging_enabled(): + print(f"[unsloth-nb] activated transformers sidecar for {version}: {d}") + return d + + +def resolve(model_name: str | None = None): + """Resolve the version to use: the notebook's pin first, else the model tier.""" + return requested_version() or tier_for_model(model_name or "") + + +# -- manual JupyterLab integration: activate before the first model cell -------- +def _pre_run_cell(_info = None): + v = requested_version() + if v and "transformers" not in sys.modules: + activate(v) + + +def register_ipython(): + """Register the pre_run_cell hook (called from the baked IPython startup).""" + try: + ip = get_ipython() # noqa: F821 (provided by IPython) + except NameError: + return + if ip is not None and not getattr(ip, "_unsloth_tf_hook", False): + ip.events.register("pre_run_cell", _pre_run_cell) + ip._unsloth_tf_hook = True diff --git a/docker/unsloth_nb_content_sig.py b/docker/unsloth_nb_content_sig.py new file mode 100644 index 0000000000..c9918540ee --- /dev/null +++ b/docker/unsloth_nb_content_sig.py @@ -0,0 +1,113 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +import hashlib +import json +import sys + +# Lowercased substrings that mark a markdown cell as top/bottom boilerplate. +_BOILERPLATE_MD = ( + "to run this, press", # Colab/AMD run announcement + 'press "*runtime*"', + "### news", # News heading + "introducing **unsloth studio**", # rotating announcement body + "you will learn how to do", # announcement tail + "this notebook is licensed", # announcement license line + "and we're done", # footer opener + "this notebook and all unsloth notebooks are licensed", # footer license + "join discord if you need help", # footer + "star us on", # footer + "some other resources", # footer resources block +) + + +def _text(cell): + src = cell.get("source", "") + if isinstance(src, list): + src = "".join(src) + return src.replace("\r\n", "\n").replace("\r", "\n") + + +# Command fragments that mark a cell as the generated install cell. +_INSTALL_MARKERS = ( + "pip install", + "pip3-autoremove", + "uv pip install", + "conda install", + "apt-get install", + "apt install", +) + + +def _is_install_code(cell): + if cell.get("cell_type") != "code": + return False + t = _text(cell) + low = t.lower() + if any(m in low for m in _INSTALL_MARKERS): + return True + # A %%capture / %%bash cell is boilerplate only if it also carries an install + # command (caught above); a bare one doing real setup is substantive, so hash + # it to avoid a false SAME on the boot refresh. + return False + + +def _is_boilerplate_md(cell): + if cell.get("cell_type") != "markdown": + return False + low = _text(cell).lower() + return any(m in low for m in _BOILERPLATE_MD) + + +def _is_boilerplate(cell): + return _is_install_code(cell) or _is_boilerplate_md(cell) + + +def middle_digest(path): + """sha256 over the (type, source) of every non-boilerplate cell, or None.""" + try: + with open(path, "r", encoding = "utf-8") as f: + nb = json.load(f) + except Exception: + return None + cells = nb.get("cells") + if not isinstance(cells, list): + return None + h = hashlib.sha256() + for cell in cells: + if not isinstance(cell, dict): + continue + if _is_boilerplate(cell): + continue + h.update(b"\x00") + h.update(str(cell.get("cell_type", "")).encode("utf-8")) + h.update(b"\x01") + h.update(_text(cell).encode("utf-8")) + return h.hexdigest() + + +def main(argv): + if len(argv) == 2: + d = middle_digest(argv[1]) + if d is None: + print("ERR") + return 0 + print(d) + return 0 + if len(argv) == 3: + a = middle_digest(argv[1]) + b = middle_digest(argv[2]) + if a is None or b is None: + print("ERR") + elif a == b: + print("SAME") + else: + print("DIFF") + return 0 + print("ERR") + return 0 + + +if __name__ == "__main__": + sys.exit(main(sys.argv)) diff --git a/docker/unsloth_nb_pip_magic.py b/docker/unsloth_nb_pip_magic.py new file mode 100644 index 0000000000..49c91852c9 --- /dev/null +++ b/docker/unsloth_nb_pip_magic.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +"""Route notebook `%pip` / `%uv` / `python -m pip` installs through the shim. + +The PATH shim (/opt/unsloth-nb/bin/{pip,pip3,uv} -> unsloth_pip_shim.py) only +intercepts `!pip` / `!uv` shell cells. IPython's `%pip` / `%uv` LINE MAGICS run +pip in-process, and `python -m pip` runs pip as a module -- both bypass PATH, so +a notebook could still reinstall torch / transformers / vLLM and clobber the +baked cu128 stack the shim is meant to protect. + +This closes that gap two ways, with no clobbering of the shell-escape path: + * `%pip` / `%pip3` / `%uv` are re-registered as line magics that delegate to + the shell (`get_ipython().system("pip ...")`); since /opt/unsloth-nb/bin is + first on PATH, that resolves to the shim. Overriding the real magic (rather + than rewriting cell text) means we only act when IPython actually dispatches + the magic -- a `%pip` inside a string is left untouched. + * a narrow input transformer rewrites an explicit `!python -m pip` / + `!python -m uv` shell line to `!pip` / `!uv`, so that form hits the shim too. + +UNSLOTH_NB_SHIM=1 is already exported by the startup hook and inherited by the +subprocess, so the shim applies. Safe no-op outside IPython. +""" + +import re + +# Only the explicit `! -m pip|uv ...` shell form. Transformers see the RAW +# cell text (IPython expands `{sys.executable}` later), so the braced form and +# quoted/bare interpreter paths must be matched here too, else module-pip bypasses +# the shim. +_PY_M_PIP = re.compile( + r"""^(\s*)!\s* + (?: + (?:python[0-9.]*|py) # literal python / py + | ["']?\{\s*sys\.executable\s*\}["']? # {sys.executable}, opt. quoted + | "(?:[^"]*[/\\])python[0-9.]*(?:\.exe)?" # quoted interpreter path + | '(?:[^']*[/\\])python[0-9.]*(?:\.exe)?' + | \S*[/\\]python[0-9.]*(?:\.exe)? # bare interpreter path + ) + \s+-m\s+(pip|uv)\b(.*)$""", + re.VERBOSE, +) + + +def _rewrite_python_dash_m(lines): + """`!python -m pip install X` -> `!pip install X` (so it hits the PATH shim).""" + try: + out = [] + for line in lines: + body = line.rstrip("\n") + tail = line[len(body) :] # preserve the trailing newline(s), if any + m = _PY_M_PIP.match(body) + if m: + out.append(m.group(1) + "!" + m.group(2) + m.group(3) + tail) + else: + out.append(line) + return out + except Exception: + return lines + + +def register_ipython(): + try: + ip = get_ipython() # noqa: F821 (provided by IPython) + except Exception: + ip = None + if ip is None or getattr(ip, "_unsloth_pip_magic", False): + return + + def _make(tool): + def _magic(line): + # /opt/unsloth-nb/bin is first on PATH, so `pip`/`uv` here is the shim. + return ip.system(tool + " " + line) + + return _magic + + ip.register_magic_function(_make("pip"), "line", "pip") + ip.register_magic_function(_make("pip"), "line", "pip3") + ip.register_magic_function(_make("uv"), "line", "uv") + + if _rewrite_python_dash_m not in ip.input_transformers_cleanup: + ip.input_transformers_cleanup.append(_rewrite_python_dash_m) + + ip._unsloth_pip_magic = True diff --git a/docker/unsloth_nb_strip_colab.py b/docker/unsloth_nb_strip_colab.py new file mode 100644 index 0000000000..8af17f2b79 --- /dev/null +++ b/docker/unsloth_nb_strip_colab.py @@ -0,0 +1,242 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +# Remove the Colab-only "how to run" sentence from Unsloth notebooks for Docker. +# +# Each generated notebook's first markdown cell opens with a Colab instruction +# ("To run this, press Runtime > Run all ...") that is wrong inside Docker. Strip +# only that leading sentence and keep the rest (badge row, install link, etc). +# Docker-only, applied at sync time; NOT pushed upstream. +# +# Two modes: +# unsloth_nb_strip_colab.py [b.ipynb ...] strip in place (idempotent) +# unsloth_nb_strip_colab.py --state --dest +# STATE-aware migration: strip + rehash each owned+unedited notebook (one +# whose hash still matches STATE); user-edited ones are left untouched. +# +# Safe with refresh: content_sig classifies the intro cell as boilerplate, so the +# body digest is unchanged. Exit code is always 0. +import argparse +import hashlib +import json +import os +import sys + +# Stable identifier for the offending line (all GPU/Cloud variants). +_INTRO_PREFIX = "to run this, press" + +# Baked notebooks ship tqdm widget outputs + a metadata.widgets block that +# JupyterLab can't rebuild, so they render as a stuck "Loading widget...". Drop +# them (the cell recreates a fresh widget). Outputs aren't in the refresh +# signature (content_sig hashes type+source), so this is safe. +_WIDGET_VIEW_MIME = "application/vnd.jupyter.widget-view+json" + + +def _is_intro_line(line): + """True for the Colab run announcement in either shipped spelling. + + Most notebooks open the line with the sentence itself, but two (NeMo-Gym-*) + ship it inside a single-line HTML comment: + + + + Only a comment that OPENS AND CLOSES on the same line is matched, so + dropping it can never leave a dangling `"): + return stripped[4:-3].strip().lower().startswith(_INTRO_PREFIX) + return False + + +def _strip_lines(lines): + """Drop the intro line (and an immediately-following blank). Return new list + or None if there was nothing to strip.""" + for i, line in enumerate(lines): + if _is_intro_line(line): + out = lines[:i] + lines[i + 1 :] + if i < len(out) and out[i].strip() == "": + out = out[:i] + out[i + 1 :] + return out + return None + + +def _strip_cell(cell): + """Strip the intro line out of ONE markdown cell. Return True if changed.""" + src = cell.get("source") + if isinstance(src, str): + lines = src.splitlines(keepends = True) + as_str = True + elif isinstance(src, list): + lines = list(src) + as_str = False + else: + return False + new_lines = _strip_lines(lines) + if new_lines is None: + return False + cell["source"] = "".join(new_lines) if as_str else new_lines + return True + + +def _strip_intro(nb): + """Strip the Colab intro sentence from the LEADING markdown block. + + Scanning cells[0] alone missed 23 of the 433 shipped notebooks: 21 put the + Colab badge `` in cells[0] and the sentence in cells[1] + (Advanced_Llama3_2_(3B)_GRPO_LoRA, Falcon_H1-Alpaca, gpt-oss-(20B)-GRPO, + ...), and 2 (NeMo-Gym-*) wrap it in an HTML comment cells[0]-only matching + never saw. The scan stops at the first non-markdown cell, so it only ever + touches the header block a notebook opens with (at most 5 cells across the + shipped set) and can never reach explanatory prose between code cells. + Return True if any cell changed.""" + cells = nb.get("cells") + if not isinstance(cells, list): + return False + changed = False + for cell in cells: + if not isinstance(cell, dict) or cell.get("cell_type") != "markdown": + break # the first code cell ends the header block + if _strip_cell(cell): + changed = True + return changed + + +def _clean_widgets(nb): + """Drop baked ipywidget outputs + the orphan widget-state metadata that + otherwise render as "Loading widget...". Return True if changed.""" + changed = False + cells = nb.get("cells") + if isinstance(cells, list): + for cell in cells: + if not isinstance(cell, dict): + continue + outs = cell.get("outputs") + if not isinstance(outs, list): + continue + kept = [ + o + for o in outs + if not (isinstance(o, dict) and _WIDGET_VIEW_MIME in (o.get("data") or {})) + ] + if len(kept) != len(outs): + cell["outputs"] = kept + changed = True + md = nb.get("metadata") + if isinstance(md, dict) and "widgets" in md: + del md["widgets"] + changed = True + return changed + + +def strip_notebook(path): + """Return True if the notebook was modified and written back.""" + try: + before = _sha256(path) + with open(path, "r", encoding = "utf-8") as f: + nb = json.load(f) + except Exception: + return False + + # Apply both transforms; write back if either changed. + changed = _strip_intro(nb) + changed = _clean_widgets(nb) or changed + if not changed: + return False + + tmp = path + ".tmp" + try: + with open(tmp, "w", encoding = "utf-8") as f: + json.dump(nb, f, indent = 1, ensure_ascii = False) + f.write("\n") + # The refresh child re-arms this cleanup AFTER the entrypoint has execed + # the container command, so JupyterLab is already serving the tree: a save + # landing between the read above and this replace would be silently + # overwritten, and migrate() would then record the cleaned hash and mark + # the notebook pristine forever. Re-read the live file once the staged + # copy is complete (the same rule the refresh publish in + # unsloth_sync_notebooks.sh follows) and let their edit win. + if _sha256(path) != before: + os.remove(tmp) + return False + os.replace(tmp, path) + except Exception: + try: + os.remove(tmp) + except OSError: + pass + return False + return True + + +def _sha256(path): + h = hashlib.sha256() + with open(path, "rb") as f: + for chunk in iter(lambda: f.read(65536), b""): + h.update(chunk) + return h.hexdigest() + + +def migrate(state_path, dest): + """Strip owned+unedited notebooks listed in STATE and update their hashes.""" + try: + with open(state_path, "r", encoding = "utf-8") as f: + lines = f.read().splitlines() + except OSError: + return 0 + + out = [] + changed = 0 + for line in lines: + parts = line.split(" ", 1) # " " + if len(parts) != 2: + out.append(line) + continue + rec, rel = parts + path = os.path.join(dest, rel) + if rel.endswith(".ipynb") and os.path.isfile(path): + try: + if _sha256(path) == rec: # we own it and it is unedited + if strip_notebook(path): + rec = _sha256(path) + changed += 1 + except OSError: + pass + out.append("%s %s" % (rec, rel)) + + if changed: + tmp = state_path + ".tmp" + try: + with open(tmp, "w", encoding = "utf-8") as f: + f.write("\n".join(out) + "\n") + os.replace(tmp, state_path) + except OSError: + pass + print(f"[unsloth-nb] cleaned {changed} notebook(s) (Colab intro + widget outputs)") + return 0 + + +def main(argv): + ap = argparse.ArgumentParser(description = "Strip the Colab-only intro sentence.") + ap.add_argument("--state", help = "sync state file (enables migration mode)") + ap.add_argument("--dest", help = "notebooks dir (with --state)") + ap.add_argument("paths", nargs = "*", help = "notebooks to strip in place") + args = ap.parse_args(argv) + + if args.state: + if not args.dest: + ap.error("--state requires --dest") + return migrate(args.state, args.dest) + + changed = sum(1 for p in args.paths if strip_notebook(p)) + if changed: + print(f"[unsloth-nb] cleaned {changed} notebook(s) (Colab intro + widget outputs)") + return 0 + + +if __name__ == "__main__": + sys.exit(main(sys.argv[1:])) diff --git a/docker/unsloth_nb_view.py b/docker/unsloth_nb_view.py new file mode 100644 index 0000000000..cc4d244def --- /dev/null +++ b/docker/unsloth_nb_view.py @@ -0,0 +1,242 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +# Build a categorized, Colab-like folder VIEW of the Unsloth notebooks. +# +# The canonical notebooks live flat under DEST/nb/.ipynb (kept by +# unsloth_sync_notebooks.sh). This builds a sibling dir of *relative symlinks* +# grouped into folders mirroring the README headers: +# /01 Main Notebooks/Llama3_2_(1B_and_3B)_Conversational.ipynb +# /99 Other Notebooks/ +# Symlinks so real files never move (the sync state machine skips them); the VIEW +# is a disposable sibling of DEST, rebuilt on every boot. +# +# Categorization rules: +# * Section = nearest preceding `###` header; a header repeated across domains +# merges into one folder (first order). +# * Folder names cleaned (dashes/slashes -> spaces) and numbered `NN ` by first +# appearance so JupyterLab's sort keeps README order; "Other" is last. +# * A notebook linked under several sections lands in its first. +# * AMD-*.ipynb hidden unless --amd; unlinked nb/*.ipynb go to "Other Notebooks". +# +# Usage: +# unsloth_nb_view.py [--amd] build the symlink view +# unsloth_nb_view.py --print [--amd] print "section\tfile" rows +# Exits nonzero on error (caller falls back to the raw tree). +import argparse +import os +import re +import sys +import urllib.parse + +# nb/.ipynb in any link form. Filenames use [\w.()-] plus %-escapes. +_NB_RE = re.compile(r"nb/([\w.()%\-]+?\.ipynb)") +_OTHER = "Other Notebooks" + + +def clean_section(title): + """README header text -> a filesystem-friendly folder label.""" + title = title.strip().strip("#").strip() + # Strip a leading emoji/symbol run so the folder label is clean text. + title = re.sub(r"^[^\w]+", "", title) + title = title.replace("-", " ").replace("/", " ") + title = re.sub(r"\s+", " ", title).strip() + return title + + +def parse_readme(readme_path): + """Return an ordered list of (section_label, filename) pairs. + + A notebook is intentionally cross-listed under several `###` headers in the + README (e.g. ModernBert under both "Embedding" and "BERT"), so that every + header becomes a populated folder. We therefore dedup per (section, file) -- + a file shows up once in EACH section that lists it -- rather than globally. + Repeated headers across the Fine-tuning / Kaggle / AMD domains share a label + and so merge into one folder downstream. + + filename is the urldecoded basename under nb/ (literal parens, matching disk). + """ + with open(readme_path, "r", encoding = "utf-8") as f: + text = f.read() + + rows = [] + seen_pairs = set() + section = None + # Reset on ANY markdown heading, not just `###`: `#`/`##` domain headers carry + # their own nb/*.ipynb tables, so matching only `###` mis-filed those links. + for line in text.splitlines(): + m = re.match(r"^#{1,6}\s+(.*)$", line) + if m: + section = clean_section(m.group(1)) + continue + if section is None: + continue + for raw in _NB_RE.findall(line): + fname = urllib.parse.unquote(raw) + key = (section, fname) + if key in seen_pairs: + continue + seen_pairs.add(key) + rows.append((section, fname)) + return rows + + +def _ordered_sections(rows): + """Section labels in first-appearance order, with Other Notebooks last.""" + order = [] + for section, _ in rows: + if section not in order: + order.append(section) + # Force the catch-all to the end even if the README defines it earlier. + order = [s for s in order if s != _OTHER] + [_OTHER] + return order + + +def build_view( + dest, + view, + amd = False, +): + nb_dir = os.path.join(dest, "nb") + readme = os.path.join(dest, "README.md") + if not os.path.isdir(nb_dir): + raise SystemExit(f"no nb/ dir under {dest}") + + # The VIEW may be a symlink to mounted storage; build inside its target. + if os.path.islink(view): + resolved = os.path.realpath(view) + if not os.path.isdir(resolved): + raise SystemExit(f"view symlink has no directory target: {view} -> {resolved}") + view = resolved + + rows = parse_readme(readme) if os.path.isfile(readme) else [] + + def allowed(fname): + return amd or not fname.startswith("AMD-") + + # section -> [filenames], preserving README order, AMD-filtered, on-disk only. + by_section = {} + placed = set() + for section, fname in rows: + if not allowed(fname): + continue + if not os.path.isfile(os.path.join(nb_dir, fname)): + continue + by_section.setdefault(section, []).append(fname) + placed.add(fname) + + # Everything on disk that the README never linked -> Other Notebooks. + for fname in sorted(os.listdir(nb_dir)): + if not fname.endswith(".ipynb"): + continue + if fname in placed or not allowed(fname): + continue + by_section.setdefault(_OTHER, []).append(fname) + + order = [s for s in _ordered_sections(rows) if s in by_section] + if _OTHER in by_section and _OTHER not in order: + order.append(_OTHER) + + # Rebuild VIEW: drop our own symlinks/empty folders, never the user's files + # (VIEW is also JupyterLab's landing dir). Ownership is keyed on DEST/nb -- + # the only place our links ever point -- so a shortcut the user made to their + # own file elsewhere in the checkout survives the rebuild. + nb_real = os.path.realpath(nb_dir) + _clear_view(view, nb_real) + os.makedirs(view, exist_ok = True) + + n_links = 0 + for i, section in enumerate(order, start = 1): + folder = os.path.join(view, f"{i:02d} {section}") + os.makedirs(folder, exist_ok = True) + for fname in by_section[section]: + link = os.path.join(folder, fname) + target = os.path.join(nb_dir, fname) + rel = os.path.relpath(target, folder) # ../../unsloth-notebooks/nb/ + try: + if os.path.islink(link) and _points_into(link, nb_real): + os.remove(link) # replace our own stale symlink + elif os.path.islink(link) or os.path.exists(link): + # a real user file occupies this name: keep it, skip linking. + print(f"[unsloth-nb] view: keep user file, skip link {fname}", file = sys.stderr) + continue + os.symlink(rel, link) + n_links += 1 + except OSError as e: + print(f"[unsloth-nb] view: skip {fname}: {e}", file = sys.stderr) + return len(order), n_links + + +def _points_into(link, nb_real): + """True when a symlink resolves into DEST/nb, the dir we link FROM. + + Every link this tool creates points at DEST/nb/, so this is the + ownership test for cleanup: a user's own symlink (to a dataset, project, + mounted dir, or their own notebook saved elsewhere in the checkout) resolves + outside DEST/nb and must survive a rebuild -- matching on all of DEST deleted + those. realpath resolves a broken link's path string too, so stale links to + since-removed notebooks are still recognised as ours. + """ + try: + target = os.path.realpath(link) + except OSError: + return False + return target == nb_real or target.startswith(nb_real + os.sep) + + +def _clear_view(path, nb_real): + # Tear down a previously built VIEW in place. It is also JupyterLab's landing + # dir, so user files/symlinks must survive: unlink only symlinks we own (see + # _points_into) and rmdir only emptied folders. The VIEW root is never unlinked. + if os.path.islink(path) or not os.path.isdir(path): + return + for root, dirs, files in os.walk(path, topdown = False): + for name in files: + p = os.path.join(root, name) + if os.path.islink(p) and _points_into(p, nb_real): + try: + os.remove(p) + except OSError: + pass + # a regular file / user symlink here is user-created -> keep it + for name in dirs: + p = os.path.join(root, name) + try: + if os.path.islink(p): + if _points_into(p, nb_real): + os.remove(p) # our symlinked dir: unlink, never recurse + else: + os.rmdir(p) # succeeds only if we emptied it + except OSError: + pass # holds user files -> keep + + +def main(argv): + ap = argparse.ArgumentParser(description = "Build the categorized notebook view.") + ap.add_argument("dest", help = "notebooks dir (contains README.md and nb/)") + ap.add_argument("view", nargs = "?", help = "output view dir (omit with --print)") + ap.add_argument("--amd", action = "store_true", help = "include AMD-* notebooks") + ap.add_argument( + "--print", + dest = "do_print", + action = "store_true", + help = "print sectionfile rows instead of building", + ) + args = ap.parse_args(argv) + + if args.do_print: + for section, fname in parse_readme(os.path.join(args.dest, "README.md")): + if args.amd or not fname.startswith("AMD-"): + print(f"{section}\t{fname}") + return 0 + + if not args.view: + ap.error("view dir is required unless --print is given") + n_sections, n_links = build_view(args.dest, args.view, amd = args.amd) + print(f"[unsloth-nb] view: {n_links} notebooks in {n_sections} folders -> {args.view}") + return 0 + + +if __name__ == "__main__": + sys.exit(main(sys.argv[1:])) diff --git a/docker/unsloth_pip_shim.py b/docker/unsloth_pip_shim.py new file mode 100644 index 0000000000..35bc2ffe29 --- /dev/null +++ b/docker/unsloth_pip_shim.py @@ -0,0 +1,822 @@ +#!/opt/unsloth-venv/bin/python +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +"""pip / uv shim for the Unsloth Docker notebook environment. + +Installed earlier on PATH than the real tools so a notebook's `!pip install ...` +or `!uv pip install ...` cell becomes SAFE + idempotent instead of clobbering the +carefully-resolved cu128 torch/vLLM/transformers stack: + + * `transformers==X` -> NOT installed into the base venv. The version X is + recorded so the sidecar mechanism (unsloth_nb_compat) activates it for the + model cells. The base stack stays intact. + * torch / torchvision / torchaudio / torchao / torchcodec / triton / xformers / + vllm / bitsandbytes / flashinfer / nvidia-* -> SKIPPED (the baked, + ABI-matched versions are kept; a notebook reinstall here only ever breaks + the GPU stack). + * trl / peft / datasets / accelerate / huggingface_hub / tokenizers / + safetensors -> SKIPPED for the same reason one level up: 382 of the shipped + notebooks end their install cell with `pip install --no-deps trl==0.22.2`, + which used to walk straight past this shim and downgrade the tested + trl 0.24.0 / peft 0.19.1 / datasets 4.3.0 on every single run. + * everything else (omegaconf, snac, causal-conv1d, ...) -> passed through to the + real tool unchanged, so notebooks that genuinely need extra packages still + get them. + +Real tools are at /opt/unsloth-venv/bin/{pip,uv}; this shim invokes them by +absolute path so there is no recursion. `python -m pip` / `%pip` bypass PATH and +are not intercepted -- the driven `unsloth-run` handles those by parsing the +notebook directly. +""" + +import os, re, sys, tempfile + +REAL = {"pip": "/opt/unsloth-venv/bin/pip", "uv": "/opt/unsloth-venv/bin/uv"} +MARKER = os.environ.get("UNSLOTH_NB_TF_MARKER", "/tmp/unsloth_nb/requested_transformers") + +# Packages whose baked version must never be changed by a notebook install cell. +# +# Membership criterion: replacing this package silently invalidates the stack the +# image was BUILT and TESTED against, or breaks unsloth outright. That is either +# (a) an ABI/CUDA-matched wheel the Dockerfile resolved deliberately (a PyPI +# reinstall swaps a +cu128 build for a generic or cu13 one), or (b) a library +# unsloth/unsloth_zoo monkey-patches by version at import time. Anything else -- +# including packages the notebook genuinely needs and the image does not bake +# (snac, causal-conv1d, omegaconf, mamba-ssm, ...) -- installs normally. +# +# Measured over the 433 shipped notebooks (probe_notebook_pins.py), the entries +# below the original torch/vLLM group cover: +# trl 382 notebooks pin an older release (0.22.2 x378, 0.15.2 x4) vs baked 0.24.0 +# torchao 2 pin 0.15.0, and 271 more reinstall it, replacing 0.17.0+cu128 +# torchcodec 26 pin 0.5 / 0.7.0, replacing the 0.11.0+cu128 wheel paired with torch 2.11 +# datasets 254 reinstall it; a trl 0.22.2 resolve pulled it back to 3.0.0 from 4.3.0 +# peft 225 reinstall it; observed dropping 0.19.1 -> 0.14.0 +# accelerate 225 reinstall it (Trainer/torch glue, patched by unsloth_zoo) +# hf hub 240 reinstall it; tokenizers 64. Both are version-locked to +# transformers, and the sidecars ship their own matched copies, so a +# base-venv swap desynchronises every sidecar at once. +_KEEP = { + "torch", + "torchvision", + "torchaudio", + "torchao", + "torchcodec", + "triton", + "triton-rocm", + "pytorch-triton", + "xformers", + "vllm", + "bitsandbytes", + "flashinfer", + "flashinfer-python", + "unsloth", + "unsloth-zoo", + "unsloth_zoo", + "trl", + "peft", + "datasets", + "accelerate", + "huggingface-hub", + "huggingface_hub", + "tokenizers", + "safetensors", +} +_KEEP_PREFIX = ("nvidia-", "nvidia_") +# pip/uv flags that consume the next token as a value (not a requirement). +_VALUE_FLAGS = { + "-r", + "--requirement", + "--requirements", + "-c", + "--constraint", + "--constraints", + "-i", + "--index-url", + "--extra-index-url", + "-f", + "--find-links", + "--target", + "-t", + "--python", + "-p", + "--prefix", + "--index-strategy", + "--upgrade-strategy", + "--upgrade-package", + "-P", + "--reinstall-package", + "--no-binary", + "--only-binary", + "--platform", + "--python-version", + "--abi", + "--implementation", + "-e", + "--editable", + # Every remaining value-taking flag of pip/uv install (from both --help). A + # missing one makes the scanner misread its VALUE. uv: + "--allow-insecure-host", + "--build-constraints", + "-b", + "--cache-dir", + "--color", + "--config-file", + "--config-setting", + "-C", + "--config-settings-package", + "--default-index", + "--directory", + "--exclude-newer", + "--exclude-newer-package", + "--excludes", + "--extra", + "--fork-strategy", + "--group", + "--index", + "--keyring-provider", + "--link-mode", + "--no-build-isolation-package", + "--no-sources-package", + "--overrides", + "--prerelease", + "--project", + "--python-platform", + "--refresh-package", + "--resolution", + "--torch-backend", + # newer uv (0.10+): + "--no-editable-package", + "--upgrade-group", + # pip: + "--build-constraint", + "--cert", + "--client-cert", + "--config-settings", + "--exists-action", + "--log", + "--progress-bar", + "--proxy", + "--report", + "--resume-retries", + "--retries", + "--root", + "--root-user-action", + "--src", + "--timeout", + "--trusted-host", + "--use-deprecated", + "--use-feature", + # newer pip (26+): + "--all-releases", + "--only-final", + "--requirements-from-script", + "--uploaded-prior-to", +} +# Value-flags whose VALUE is itself an install target (a requirements file pulls +# real requirements). uv spells the long forms plural; include both. +_REQ_FILE_FLAGS = {"-r", "--requirement", "--requirements"} +# Constraint files aren't install targets, but pip applies their pins, so a -c +# pinning torch/transformers can downgrade a baked package. Filter like -r files. +_CONSTRAINT_FILE_FLAGS = {"-c", "--constraint", "--constraints"} +# -e/--editable takes the next token as a real install target. A protected +# editable must drop BOTH flag and value, else a dangling -e swallows the next +# kept package and fails the cell. +_EDITABLE_FLAGS = {"-e", "--editable"} +# -P/--upgrade-package/--reinstall-package are uv's selective upgrade flags: +# filter the value through _KEEP, dropping the flag+value pair for a protected +# name. Unlike -e, none is itself an install target. +_UPGRADE_PKG_FLAGS = {"-P", "--upgrade-package", "--reinstall-package"} +# Short value-flags accepted ATTACHED (-rreqs.txt, -cX, -epath, -Pname). Split +# flag from value so it's filtered, else -r no-ops and -c/-e/-P bypass _KEEP. +_ATTACHED_SHORT_FLAGS = {"-r", "-c", "-e", "-P"} +# Resolver-wide reinstall/ignore-installed switches (pip --force-reinstall, +# --ignore-installed, -I; uv --reinstall) rebuild baked deps; drop them (the kept +# target still installs). uv's --exact removes everything outside the closure, so +# drop it too. +_REINSTALL_FLAGS = {"--force-reinstall", "--ignore-installed", "-I", "--reinstall", "--exact"} +# Value-flags dropped outright with their value. --upgrade-strategy eager would +# upgrade every dep of a kept target; dropping it falls back to only-if-needed. +_DROP_VALUE_FLAGS = {"--upgrade-strategy"} + + +# Source-distribution / archive suffixes pip accepts as an install target. +_ARCHIVE_EXTS = (".tar.gz", ".tgz", ".tar.bz2", ".tbz2", ".tar.xz", ".txz", ".tar", ".zip") + + +def _sdist_name(basename): + """Distribution name from a source-archive basename ({name}-{version}.ext), + or None if it is not a recognised archive. Splits at the first hyphen that + precedes a digit so legacy hyphenated names (flashinfer-python-1.0, + pytorch-triton-2.0) resolve correctly, not just PEP 625-normalised ones.""" + low = basename.lower() + stem = None + for ext in _ARCHIVE_EXTS: + if low.endswith(ext): + stem = basename[: -len(ext)] + break + if stem is None: + return None + m = re.match(r"^(.+?)-\d", stem) + name = (m.group(1) if m else stem).strip().lower().replace("_", "-") + return name or None + + +def _canon(token): + """Extract the lowercased distribution name from a requirement token, or None + if the token is not a plain pkg spec (url / path / vcs / option).""" + if token.startswith("-"): + return None + # PEP 508 direct reference: "name [extras] @ ". Pull the name out BEFORE + # the url/vcs guard below, else a protected package pinned via URL slips _KEEP. + _dref = re.match( + r"^([A-Za-z0-9][A-Za-z0-9._-]*)\s*(?:\[[^\]]*\])?\s*@(?:\s|git\+|hg\+|bzr\+|svn\+|[a-z]+://)", + token, + ) + if _dref: + return _dref.group(1).lower().replace("_", "-") or None + if re.match(r"^[a-z]+\+", token) or "://" in token or token.startswith((".", "/")): + # A VCS/URL install can name a protected package via the #egg=NAME + # fragment; pull it out so _KEEP can drop it. + _egg = re.search(r"[#&]egg=([A-Za-z0-9][A-Za-z0-9._-]*)", token) + if _egg: + return _egg.group(1).lower().replace("_", "-") or None + # A wheel URL/path names its distribution in the PEP 427 filename (leading + # dash-split of the basename), so a bare torch-*.whl would slip _KEEP. + _whl = re.search(r"([^/\\#?]+)\.whl(?:[#?]|$)", token) + if _whl: + dist = _whl.group(1).split("-", 1)[0].strip().lower().replace("_", "-") + if dist: + return dist + # A source archive ({name}-{version}.tar.gz) names its distribution too; + # match it against _KEEP instead of passing it through as opaque. + _arch = _sdist_name(token.split("#", 1)[0].split("?", 1)[0].rstrip("/").rsplit("/", 1)[-1]) + if _arch: + return _arch + # A VCS URL without #egg= still installs a named project; the repo basename + # equals the distribution for our protected packages. Infer from the last + # path segment so an egg-less git+ URL can't reinstall past _KEEP. + if re.match(r"^[a-z]+\+", token): + _rest = token.split("#", 1)[0].split("?", 1)[0] + # Drop the @ref before the basename (a ref may contain a slash). Split + # path from authority first so an SSH userinfo @ isn't the ref; like + # pip, the ref is everything after the LAST @. + if "://" in _rest: + _authority, _slash, _path = _rest.partition("://")[2].partition("/") + if "@" in _path: + _path = _path.rsplit("@", 1)[0] + _rest = _path if _slash else _authority + _seg = _rest.rstrip("/").rsplit("/", 1)[-1] + _seg = _seg.split("@", 1)[0] # schemeless fallback: drop a plain @ref + if _seg.endswith(".git"): + _seg = _seg[:-4] + _seg = _seg.strip().lower().replace("_", "-") + if _seg: + return _seg + # A local project DIRECTORY installs the project it contains; resolve its + # name from metadata so _KEEP applies. Metadata-less dirs pass through. + _local = _local_project_name(token) + if _local: + return _local + return None # plain url / metadata-less local path -> let it pass through + # A local project dir referenced without ./ or / is still a path target when + # it exists on disk; classify it before the spec parse mangles the separator. + if "/" in token or os.sep in token: + _local = _local_project_name(token) + if _local: + return _local + # A bare wheel filename from the CWD is a valid pip target; parse its PEP 427 + # distribution like the URL/path wheel case above, else it misses _KEEP. + if token.lower().endswith(".whl"): + dist = token.rsplit("/", 1)[-1][:-4].split("-", 1)[0].strip().lower().replace("_", "-") + if dist: + return dist + # A bare source-archive filename from the CWD is a valid target too; parse it. + _barch = _sdist_name(token.rsplit("/", 1)[-1]) + if _barch: + return _barch + # strip extras and any version/marker tail + name = re.split(r"[<>=!~\[\s;@]", token, 1)[0].strip() + return name.lower().replace("_", "-") or None + + +def _local_project_name(token): + """Distribution name of a local project directory install target, else None. + + Reads the name pip/uv would build: pyproject.toml [project].name, falling + back to setup.cfg [metadata] name, falling back to the directory basename + when a setup.py exists (a bare basename guess is used ONLY when the dir is + an installable project at all). A directory without any project metadata is + not a pip target and returns None so ordinary paths pass through untouched. + Names are exact after normalization: a user's own `my-torch-utils` dir never + matches the protected `torch`. + """ + path = token.split("#", 1)[0] + if not os.path.isdir(path): + return None + _pyproject = os.path.join(path, "pyproject.toml") + if os.path.isfile(_pyproject): + try: + import tomllib + with open(_pyproject, "rb") as f: + _name = (tomllib.load(f).get("project") or {}).get("name") + if _name: + return _name.strip().lower().replace("_", "-") or None + except Exception: + pass # unparseable metadata -> fall through to the other signals + _setup_cfg = os.path.join(path, "setup.cfg") + if os.path.isfile(_setup_cfg): + try: + import configparser + + _cp = configparser.ConfigParser() + _cp.read(_setup_cfg) + _name = _cp.get("metadata", "name", fallback = None) + if _name: + return _name.strip().lower().replace("_", "-") or None + except Exception: + pass + if os.path.isfile(os.path.join(path, "setup.py")) or os.path.isfile(_pyproject): + _base = os.path.basename(os.path.normpath(path)) + return _base.strip().lower().replace("_", "-") or None + return None + + +def _version_pin(token): + """Return the pinned version for a `pkg==X` token, else None.""" + m = re.search(r"==\s*([0-9][0-9A-Za-z.\-]*)", token) + return m.group(1) if m else None + + +# pip expands ${UPPERCASE_NAME} in requirements files, so `${PKG}==...` with +# PKG=torch would slip _KEEP. Expand for CLASSIFICATION only; kept lines verbatim. +_ENV_REF_RE = re.compile(r"\$\{([A-Z0-9_]+)\}") + + +def _expand_env_refs(text): + return _ENV_REF_RE.sub(lambda m: os.environ.get(m.group(1), m.group(0)), text) + + +def _classify_flag_target(spec): + """Classify the value that rides on -e/--editable or -P/--upgrade-package. + + Returns ("drop", version_or_None) when the value names a protected package + (so the flag+value pair must be dropped, closing the same bypass the bare + positional spec closes) or ("keep", None) when it is safe to forward. + transformers is reported as "drop" with any pinned version so its sidecar + marker is still recorded, mirroring the bare-spec handling in main().""" + name = _canon(spec) + if name == "transformers": + return "drop", _version_pin(spec) + if name is not None and (name in _KEEP or name.startswith(_KEEP_PREFIX)): + return "drop", None + return "keep", None + + +def _parse_flag_line(stripped, flags): + """If `stripped` is a ` ` requirements-file line for one of + `flags`, return (flag, target_or_None, inline_comment_or_None); else + (None, None, None). + + Shared by the `-r`/`--requirement`/`-c`/`--constraint` include parse and + the `-e`/`--editable` install-line parse. Handles the separated + (`-r ` / `--editable `), inline (`--editable=` / `-e=`) and + attached short (`-rextras.txt`, `-egit+...`) forms pip accepts from a + requirement file, so a protected include or editable there is handled + exactly like the command-line case.""" + body, sep, comment = stripped.partition(" #") + body = body.rstrip() + comment = ("#" + comment) if sep else None + for flag in flags: + if body == flag or body.startswith(flag + " "): + target = body[len(flag) :].strip() + elif body.startswith(flag + "="): + target = body[len(flag) + 1 :].strip() + elif not flag.startswith("--") and body.startswith(flag) and len(body) > len(flag): + target = body[len(flag) :].strip() # attached short form + else: + continue + return flag, (target or None), comment + return None, None, None + + +def _rewrite_include(line, stripped, src_dir, depth): + """Rewrite a nested `-r`/`-c` include so pip still resolves it and its + protected specs are filtered too. + + pip resolves a nested include against the directory of the file it is + READING; our filtered copy lives under /tmp, so a relative include would + look in /tmp and fail. Recursively filter the included file (dropping + protected packages there too, closing the multi-level bypass) and point the + parent at that filtered copy. URLs and unreadable/absolute-unfiltered files + fall back to an absolutised path so they still resolve. Returns + (new_line, changed, recorded, dropped).""" + flag, raw_target, comment = _parse_flag_line( + stripped, ("-r", "--requirement", "-c", "--constraint") + ) + if not raw_target: + return line, False, None, [] + # Resolve pip's ${VAR} references so the include we read/filter is the file + # pip would actually read (a literal `${DIR}/reqs.txt` never resolves here). + target = _expand_env_refs(raw_target) + newline_char = "\n" if line.endswith("\n") else "" + + def _emit(new_target): + rebuilt = flag + " " + new_target + if comment: + rebuilt += " " + comment + return rebuilt + newline_char + + # A remote (URL) nested include can't be filtered here, so drop it rather than + # let pip pull unfiltered pins off the network (mirrors main's top-level + # refusal). new_line=None tells the caller to remove the line. + if "://" in target: + return None, True, None, [flag + " " + raw_target] + abs_target = target if os.path.isabs(target) else os.path.join(src_dir, target) + # Recursively filter the included file. Guard against cyclic / deep includes. + if depth < 8: + f_path, f_rec, f_drp = _filter_requirements_file(abs_target, _depth = depth + 1) + # A nested -c include is a resolver CONSTRAINT, not an install request, so + # don't record its transformers pin (mirrors main's -c path). Only -r + # includes carry real requests, so keep their pin. + if flag in _CONSTRAINT_FILE_FLAGS: + f_rec = None + if f_path != abs_target: + # The include was rewritten; point at the filtered copy. + return _emit(f_path), True, f_rec, f_drp + # Nothing to filter inside; just make sure the path still resolves from /tmp. + if not os.path.isabs(target): + return _emit(abs_target), True, None, [] + return line, False, None, [] + + +def _filter_requirements_file(path, _depth = 0): + """Strip baked/protected packages out of a `-r` requirements file. + + Returns (path_to_use, recorded_transformers_version, dropped_specs). The same + _KEEP / transformers rules the inline args get are applied to each requirement + line, so a notebook `pip install -r reqs.txt` cannot overwrite the cu128 torch + / vLLM / transformers stack with versions pinned inside the file. When nothing + is protected, or the file cannot be read/written, the original path is returned + unchanged. Comments, blank lines and option lines are kept verbatim; a nested + `-r`/`-c` include is recursively filtered too (protected specs dropped at every + level). + """ + try: + with open(path, encoding = "utf-8") as f: + lines = f.readlines() + except OSError: + return path, None, [] # remote URL / unreadable -> let the real tool handle it + src_dir = os.path.dirname(os.path.abspath(path)) + out, dropped, recorded, changed = [], [], None, False + for line in lines: + stripped = line.strip() + if not stripped or stripped.startswith("#"): + out.append(line) # comment / blank -> keep + continue + if stripped.startswith("-"): + # An -e/--editable in the file is a real install target, so a + # protected editable would reinstall the baked stack. Classify through + # _KEEP like the command-line -e case; drop the whole line when + # protected (a transformers pin is still recorded). + e_flag, e_target, _e_comment = _parse_flag_line(stripped, ("-e", "--editable")) + if e_target is not None: + _action, _ver = _classify_flag_target(_expand_env_refs(e_target)) + if _action == "drop": + if _ver and not recorded: + recorded = _ver + dropped.append(e_flag + " " + e_target) + changed = True + continue + out.append(line) # kept editable -> forward the line verbatim + continue + # Option or nested include. Recursively filter a nested `-r`/`-c` + # include (protected specs deep in the tree) and repoint it for /tmp. + new_line, rewrote, inc_rec, inc_drp = _rewrite_include(line, stripped, src_dir, _depth) + if new_line is not None: + out.append(new_line) # None -> a remote include was dropped + if rewrote: + changed = True + if inc_rec and not recorded: + recorded = inc_rec + dropped.extend(inc_drp) + continue + spec = stripped.split(" #", 1)[0].strip() # drop any inline comment + classified = _expand_env_refs(spec) # classify what pip will SEE + name = _canon(classified) + if name is None: + out.append(line) # url / path / vcs / unparseable -> keep + continue + if name == "transformers": + v = _version_pin(classified) + if v and not recorded: + recorded = v + dropped.append(spec) + changed = True + continue + if name in _KEEP or name.startswith(_KEEP_PREFIX): + dropped.append(spec) + changed = True + continue + out.append(line) + if not changed: + return path, None, [] + try: + fd, tmp = tempfile.mkstemp(prefix = "unsloth-nb-req-", suffix = ".txt") + with os.fdopen(fd, "w", encoding = "utf-8") as f: + f.writelines(out) + except OSError as exc: + # Fail CLOSED: protected requirements were detected, so forwarding the + # original would hand pip the specs we must filter. Abort instead. + raise SystemExit( + f"[unsloth-nb] could not write a filtered copy of {path} ({exc}); " + "refusing to forward a requirements file that pins protected packages." + ) + return tmp, recorded, dropped + + +def _protected_constraints_file(): + """Write `name==version` pins for every INSTALLED protected package to a + temp constraints file and return its path (None when nothing is pinned or + the file cannot be written). + + Argument filtering alone does not constrain pip/uv's RESOLVER: a kept + package may declare e.g. `torch==99.0` as a dependency and the tool would + replace the baked torch to satisfy it. Pinning the protected set on every + forwarded install makes such an install fail loudly instead. This is + belt-and-braces on top of the argument filtering, so a failure here keeps + the install usable rather than aborting it. + """ + try: + from importlib.metadata import distributions + + pins = {} + for dist in distributions(): + raw = (dist.metadata["Name"] or "").strip() + name = raw.lower().replace("_", "-") + if not name or name in pins: + continue + if name == "transformers" or name in _KEEP or name.startswith(_KEEP_PREFIX): + pins[name] = f"{raw}=={dist.version}" + if not pins: + return None + fd, tmp = tempfile.mkstemp(prefix = "unsloth-nb-protected-", suffix = ".txt") + with os.fdopen(fd, "w", encoding = "utf-8") as f: + f.write("\n".join(pins[name] for name in sorted(pins)) + "\n") + return tmp + except Exception: + return None + + +def _selfcheck_value_flags(): + """Assert every value-taking flag the REAL pip/uv document is classified. + + A value flag missing from _VALUE_FLAGS makes the scanner misread its VALUE + (see --torch-backend in the header of the added block above). Run at image + build time against the BAKED tools -- the exact versions the shim fronts -- + so a pip/uv bump that adds a value flag fails the build, not a user's cell. + Exits 0 when clean, 1 with the missing flags listed. + """ + import subprocess + + known = _VALUE_FLAGS | _DROP_VALUE_FLAGS + missing = {} + for label, cmd in ( + ("pip", [REAL["pip"], "install", "--help"]), + ("uv", [REAL["uv"], "pip", "install", "--help"]), + ): + try: + out = subprocess.run(cmd, capture_output = True, text = True).stdout + except OSError: + continue # tool absent (e.g. a pip-only environment) + flags = set() + for m in re.finditer(r"^\s+(-\w)?,?\s*(--[\w-]+)[= ]<", out, re.M): + if m.group(1): + flags.add(m.group(1)) + flags.add(m.group(2)) + for m in re.finditer(r"^\s+(-\w) <", out, re.M): + flags.add(m.group(1)) + gap = flags - known + if gap: + missing[label] = sorted(gap) + if missing: + print(f"[unsloth-nb] value flags missing from _VALUE_FLAGS: {missing}", file = sys.stderr) + sys.exit(1) + print("[unsloth-nb] value-flag selfcheck OK") + sys.exit(0) + + +def main(): + tool = "uv" if os.path.basename(sys.argv[0]).startswith("uv") else "pip" + argv = sys.argv[1:] + + if argv[:1] == ["--unsloth-selfcheck-value-flags"]: + _selfcheck_value_flags() + + # Only intercept inside a notebook kernel (UNSLOTH_NB_SHIM); everywhere else + # behave exactly like the real tool. + if os.environ.get("UNSLOTH_NB_SHIM") != "1": + os.execv(REAL[tool], [REAL[tool]] + argv) + return + + # Locate the `install` verb (pip: `pip install ...`; uv: `uv pip install ...` + # -- index() already skips uv's leading `pip` subcommand). + try: + i = argv.index("install") + except ValueError: + os.execv(REAL[tool], [REAL[tool]] + argv) # not an install -> passthrough + return + + head, tail = argv[: i + 1], argv[i + 1 :] + keep_args, dropped, recorded = [], [], None + has_target = False + skip_next = False + prev_flag = None + for tok in tail: + if skip_next: + # -r/--requirement's value pulls real requirements (a target); an + # index-url / find-links / constraint value is an option, not a target. + if prev_flag in _REQ_FILE_FLAGS or prev_flag in _CONSTRAINT_FILE_FLAGS: + if "://" in tok: + # Remote requirement/constraint file: can't be filtered, so + # refuse it rather than fetch protected pins off the network. + # Pop the flag we appended so pip/uv has no dangling -r/-c. + if keep_args and keep_args[-1] == prev_flag: + keep_args.pop() + dropped.append(prev_flag + " " + tok) + elif prev_flag in _REQ_FILE_FLAGS: + # Filter protected packages out of the requirements file so + # `pip install -r reqs.txt` can't clobber the cu128 stack. + _req_path, _req_rec, _req_drp = _filter_requirements_file(tok) + keep_args.append(_req_path) + has_target = True + if _req_rec and not recorded: + recorded = _req_rec + dropped.extend(_req_drp) + else: + # Strip protected pins from the constraint file so it can't + # downgrade the baked stack; a constraint isn't an install + # target, so don't set has_target / recorded here. + _c_path, _c_rec, _c_drp = _filter_requirements_file(tok) + keep_args.append(_c_path) + dropped.extend(_c_drp) + elif prev_flag in _DROP_VALUE_FLAGS: + # --upgrade-strategy (eager): drop the pair so pip falls back to + # only-if-needed. + if keep_args and keep_args[-1] == prev_flag: + keep_args.pop() + dropped.append(prev_flag + " " + tok) + elif prev_flag in _EDITABLE_FLAGS or prev_flag in _UPGRADE_PKG_FLAGS: + # Flag held back: its value is an install target (-e) or upgrade + # selector (-P), filtered through _KEEP. A protected value drops + # the flag too. A kept editable sets has_target; -P does not. + _action, _ver = _classify_flag_target(tok) + if _action == "drop": + if _ver and not recorded: + recorded = _ver + dropped.append(prev_flag + " " + tok) + else: + keep_args.append(prev_flag) + keep_args.append(tok) + if prev_flag in _EDITABLE_FLAGS: + has_target = True + else: + keep_args.append(tok) + skip_next = False + prev_flag = None + continue + # --flag=value form (--requirement=reqs.txt / --index-url=URL as one token). + # Without this the -r file is never filtered and a file-only cell no-ops. + if tok.startswith("--") and "=" in tok: + _flag, _, _val = tok.partition("=") + if _flag in _VALUE_FLAGS: + if (_flag in _REQ_FILE_FLAGS or _flag in _CONSTRAINT_FILE_FLAGS) and "://" in _val: + # Remote requirement/constraint file in `--flag=URL` form: + # refuse it (dropping the token leaves nothing dangling). + dropped.append(tok) + elif _flag in _REQ_FILE_FLAGS: + _req_path, _req_rec, _req_drp = _filter_requirements_file(_val) + keep_args.append(_flag + "=" + _req_path) + has_target = True + if _req_rec and not recorded: + recorded = _req_rec + dropped.extend(_req_drp) + elif _flag in _DROP_VALUE_FLAGS: + dropped.append(tok) # --upgrade-strategy=eager -> drop the pair + elif _flag in _CONSTRAINT_FILE_FLAGS: + _c_path, _c_rec, _c_drp = _filter_requirements_file(_val) + keep_args.append(_flag + "=" + _c_path) + dropped.extend(_c_drp) + elif _flag in _EDITABLE_FLAGS or _flag in _UPGRADE_PKG_FLAGS: + # --editable= / --upgrade-package=: filter the + # inline value through _KEEP, dropping the token if protected. + _action, _ver = _classify_flag_target(_val) + if _action == "drop": + if _ver and not recorded: + recorded = _ver + dropped.append(tok) + else: + keep_args.append(tok) + if _flag in _EDITABLE_FLAGS: + has_target = True + else: + keep_args.append(tok) # option with inline value, not a target + continue + # Attached short value-flag form (-rreqs.txt, -cX, -epath, -Pname as ONE + # token). Split flag from value and reuse the separated-form handling, + # else -r no-ops and -c/-e/-P bypass _KEEP. + if len(tok) > 2 and tok[0] == "-" and tok[1] != "-" and tok[:2] in _ATTACHED_SHORT_FLAGS: + _sflag, _sval = tok[:2], tok[2:] + if (_sflag in _REQ_FILE_FLAGS or _sflag in _CONSTRAINT_FILE_FLAGS) and "://" in _sval: + # Remote requirement/constraint file in attached `-rURL`/`-cURL` + # form: refuse it (nothing appended yet, drop the whole token). + dropped.append(_sflag + " " + _sval) + elif _sflag in _REQ_FILE_FLAGS: + _req_path, _req_rec, _req_drp = _filter_requirements_file(_sval) + keep_args.append(_sflag) + keep_args.append(_req_path) + has_target = True + if _req_rec and not recorded: + recorded = _req_rec + dropped.extend(_req_drp) + elif _sflag in _CONSTRAINT_FILE_FLAGS: + _c_path, _c_rec, _c_drp = _filter_requirements_file(_sval) + keep_args.append(_sflag) + keep_args.append(_c_path) + dropped.extend(_c_drp) + else: # -e / -P: the attached value is an install target / selector + _action, _ver = _classify_flag_target(_sval) + if _action == "drop": + if _ver and not recorded: + recorded = _ver + dropped.append(_sflag + " " + _sval) + else: + keep_args.append(_sflag) + keep_args.append(_sval) + if _sflag in _EDITABLE_FLAGS: + has_target = True + continue + if tok in _REINSTALL_FLAGS: + # Resolver-wide reinstall / ignore-installed switch: drop it so pip/uv + # can't rebuild satisfied baked deps. The kept target still installs. + dropped.append(tok) + continue + if tok in _VALUE_FLAGS: + # -e/--editable and -P/--upgrade-package carry a potential install + # target, so hold the flag back and let skip_next emit or drop the + # pair together. Every other value-flag keeps its flag verbatim; only + # its value is an opaque option. + if tok not in _EDITABLE_FLAGS and tok not in _UPGRADE_PKG_FLAGS: + keep_args.append(tok) + skip_next = True + prev_flag = tok + continue + name = _canon(tok) + if name is None: + keep_args.append(tok) # bare flag, or a positional url / path / vcs + if not tok.startswith("-"): + has_target = True # standalone . / ./pkg / git+... / *.whl + continue + if name == "transformers": + v = _version_pin(tok) + if v: + recorded = v + dropped.append(tok) + continue + if name in _KEEP or name.startswith(_KEEP_PREFIX): + dropped.append(tok) + continue + keep_args.append(tok) + has_target = True # a kept package spec + + if recorded: + try: + os.makedirs(os.path.dirname(MARKER), exist_ok = True) + with open(MARKER, "w") as f: + f.write(recorded) + print( + f"[unsloth-nb] notebook requested transformers=={recorded}; will " + f"activate its sidecar for the model cells (base stack kept)." + ) + except OSError: + pass + if dropped: + print("[unsloth-nb] kept baked versions, skipped: " + " ".join(dropped)) + + # Anything left to install? A line with only baked packages + option flags + # leaves no target, so no-op instead of exec'ing a bare install that fails. + if not has_target: + print("[unsloth-nb] nothing to install after keeping the baked stack; ok.") + return + cmd = [REAL[tool]] + head + keep_args + # Constrain the resolver too: an allowed target could pull an incompatible + # torch/transformers in as a dependency and replace the baked wheel. + constraints = _protected_constraints_file() + if constraints: + cmd += ["--constraint", constraints] + sys.stdout.flush() + os.execv(REAL[tool], cmd) + + +if __name__ == "__main__": + main() diff --git a/docker/unsloth_run.py b/docker/unsloth_run.py new file mode 100644 index 0000000000..5a68644fd9 --- /dev/null +++ b/docker/unsloth_run.py @@ -0,0 +1,162 @@ +#!/opt/unsloth-venv/bin/python +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +"""unsloth-run: execute an unslothai/notebooks notebook unchanged, headless. + +The robust driven path for the Docker image: it reads the notebook, figures out +which transformers version it wants (its install-cell pin, else the model-name +tier), launches the kernel with that sidecar on PYTHONPATH so the whole kernel +process uses a coherent transformers, and executes every cell with nbconvert. +The notebook's own install cell still runs through the pip/uv shim, so it is safe +and idempotent (the baked torch/vLLM stack is never clobbered). + +Usage: + unsloth-run [--out OUT.ipynb] [--timeout SECONDS] + [--transformers X.Y.Z] # force a version, skip auto-detect + +A raw github URL (raw.githubusercontent.com/.../nb/Foo.ipynb) is fetched first. +""" + +import argparse, json, os, re, shutil, subprocess, sys, tempfile, urllib.request + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +try: + import unsloth_nb_compat as compat +except Exception: + compat = None + +_PIN_RE = re.compile(r"transformers\s*==\s*([0-9][0-9A-Za-z.\-]*)") +_MODEL_RE = re.compile(r"""from_pretrained\(\s*['"]([^'"]+)['"]""") +_MODEL_NAME_RE = re.compile(r"""model_name\s*=\s*['"]([^'"]+)['"]""") + + +def _load(path_or_url): + if path_or_url.startswith(("http://", "https://")): + with urllib.request.urlopen(path_or_url) as r: # nosec - user-provided nb + data = r.read().decode() + return json.loads(data) + with open(path_or_url) as f: + return json.load(f) + + +def _scan(nb): + """Return (pinned_transformers, first_model_name) from the notebook source.""" + pin = model = None + for cell in nb.get("cells", []): + if cell.get("cell_type") != "code": + continue + src = "".join(cell.get("source", [])) + if pin is None: + m = _PIN_RE.search(src) + if m: + pin = m.group(1) + if model is None: + m = _MODEL_RE.search(src) or _MODEL_NAME_RE.search(src) + if m: + model = m.group(1) + return pin, model + + +def main(): + ap = argparse.ArgumentParser(prog = "unsloth-run") + ap.add_argument("notebook") + ap.add_argument("--out") + ap.add_argument("--timeout", type = int, default = 3600) + ap.add_argument("--transformers", dest = "tf") + args = ap.parse_args() + + nb = _load(args.notebook) + pin, model = _scan(nb) + want = args.tf or pin or (compat.tier_for_model(model) if compat else None) + sidecar = compat.sidecar_for(want) if (compat and want) else None + + # Materialise the notebook for nbconvert. With --out, stage input + result as + # temp files next to the destination (same dir => atomic os.replace publish) + # and publish only on success, so a failed run can't destroy the old output. + tmp_dir = None + tmp_files = [] + publish_from = None + if args.out: + out_path = os.path.abspath(args.out) + out_dir = os.path.dirname(out_path) or "." + os.makedirs(out_dir, exist_ok = True) + fd, src_path = tempfile.mkstemp(prefix = ".unsloth-run-in-", suffix = ".ipynb", dir = out_dir) + with os.fdopen(fd, "w") as f: + json.dump(nb, f) + tmp_files.append(src_path) + fd, publish_from = tempfile.mkstemp( + prefix = ".unsloth-run-out-", suffix = ".ipynb", dir = out_dir + ) + os.close(fd) + tmp_files.append(publish_from) + elif args.notebook.startswith(("http://", "https://")): + tmp_dir = tempfile.mkdtemp() + src_path = os.path.join(tmp_dir, os.path.basename(args.notebook.split("?")[0])) + with open(src_path, "w") as f: + json.dump(nb, f) + out_path = src_path + else: + src_path = args.notebook + out_path = src_path + + env = dict(os.environ) + env["UNSLOTH_NB_SHIM"] = "1" # enable safe-install for the notebook's cells + # Per-run marker unless the caller pinned one: the shared default would leak + # this run's transformers pin into concurrent/later runs. An empty marker + # reads as "no pin", so pre-creating it is safe. + marker = env.get("UNSLOTH_NB_TF_MARKER") + if not marker: + fd, marker = tempfile.mkstemp(prefix = ".unsloth-run-tfmarker-") + os.close(fd) + env["UNSLOTH_NB_TF_MARKER"] = marker + tmp_files.append(marker) + # The pip/uv shim writes the marker; pre-seed it too so the kernel agrees. + if want: + os.makedirs(os.path.dirname(marker) or ".", exist_ok = True) + open(marker, "w").write(want) + if sidecar: + env["PYTHONPATH"] = sidecar + os.pathsep + env.get("PYTHONPATH", "") + print(f"[unsloth-run] transformers {want} -> sidecar {sidecar}") + elif want: + print(f"[unsloth-run] transformers {want}: no sidecar (using base venv's newest)") + else: + print("[unsloth-run] no transformers pin/model tier detected; using base venv") + + nbconvert_out = publish_from if publish_from is not None else out_path + cmd = [ + "/opt/unsloth-venv/bin/jupyter", + "nbconvert", + "--to", + "notebook", + "--execute", + f"--ExecutePreprocessor.timeout={args.timeout}", + "--ExecutePreprocessor.kernel_name=python3", + src_path, + "--output", + os.path.basename(nbconvert_out), + "--output-dir", + os.path.dirname(os.path.abspath(nbconvert_out)) or ".", + ] + print( + "[unsloth-run] executing:", + os.path.basename(args.notebook.split("?")[0]) if args.out else os.path.basename(src_path), + ) + try: + rc = subprocess.call(cmd, env = env) + if rc == 0 and publish_from is not None: + os.replace(publish_from, out_path) + finally: + # Clean up the temp dir and any staging files (already gone when published). + if tmp_dir is not None: + shutil.rmtree(tmp_dir, ignore_errors = True) + for p in tmp_files: + try: + os.remove(p) + except OSError: + pass + sys.exit(rc) + + +if __name__ == "__main__": + main() diff --git a/docker/unsloth_studio_update.sh b/docker/unsloth_studio_update.sh new file mode 100755 index 0000000000..8855e7b19e --- /dev/null +++ b/docker/unsloth_studio_update.sh @@ -0,0 +1,122 @@ +#!/usr/bin/env bash +# Update Unsloth Studio in place, inside a running container, without pulling a +# new image. Updates ONLY the Studio Python packages (the backend code and the +# pre-built frontend, which ships inside the unsloth wheel) and restarts the +# Studio service. The torch/CUDA stack is left untouched. +# +# docker exec unsloth-studio-update # latest PyPI release +# docker exec unsloth-studio-update --ref main # latest git main +# docker exec unsloth-studio-update --with-deps # also update deps +# docker exec unsloth-studio-update --no-restart # update, restart later +# +# Why not `unsloth studio update`: that command re-runs the full installer, +# which re-probes the host GPU to pick torch wheels. In a CPU-only container +# (run without --gpus) it finds no GPU and can downgrade torch to CPU/cu126, +# breaking CUDA. This helper only touches the Studio packages, so it is safe in +# both GPU and CPU containers. +# +# Persistence: the update is written to the container's writable layer, so it +# survives `docker restart`. To keep it across a full `docker rm` + `docker run` +# (and to keep your chats/users/models), run Studio with its home on a named +# volume: -v unsloth_studio_home:/opt/unsloth-studio +set -euo pipefail + +STUDIO_HOME="${UNSLOTH_STUDIO_HOME:-/opt/unsloth-studio}" +REF="" +ZOO_REF="" +NO_DEPS="--no-deps" +RESTART=1 +PACKAGES="unsloth unsloth_zoo" + +usage() { sed -n '2,21p' "$0"; } + +while [ $# -gt 0 ]; do + case "$1" in + --ref) REF="$2"; shift 2;; + --zoo-ref) ZOO_REF="$2"; shift 2;; + --with-deps) NO_DEPS=""; shift;; + --no-restart) RESTART=0; shift;; + --packages) PACKAGES="$2"; shift 2;; + -h|--help) usage; exit 0;; + *) echo "unsloth-studio-update: unknown argument: $1" >&2; usage; exit 2;; + esac +done + +# Resolve the Studio venv python. Prefer the venv directly; fall back to +# following the launcher symlink ($STUDIO_HOME/bin/unsloth -> venv/bin/unsloth). +PY="" +for cand in \ + "$STUDIO_HOME/unsloth_studio/bin/python" \ + "$STUDIO_HOME/unsloth_studio/bin/python3"; do + [ -x "$cand" ] && { PY="$cand"; break; } +done +if [ -z "$PY" ] && [ -L "$STUDIO_HOME/bin/unsloth" ]; then + venv_bin="$(dirname "$(readlink -f "$STUDIO_HOME/bin/unsloth")")" + [ -x "$venv_bin/python" ] && PY="$venv_bin/python" +fi +[ -n "$PY" ] || { echo "unsloth-studio-update: could not find the Studio venv under $STUDIO_HOME" >&2; exit 1; } + +version_of() { "$PY" -c "from importlib.metadata import version; print(version('unsloth'))" 2>/dev/null || echo "unknown"; } + +echo "[studio-update] Studio venv: $PY" +echo "[studio-update] before: unsloth $(version_of)" + +# Build the package specs. With --ref, install from git so you can track main +# (or any branch/tag/sha); otherwise take the latest PyPI release. +if [ -n "$REF" ]; then + SPECS="git+https://github.com/unslothai/unsloth.git@${REF}#egg=unsloth" + # unsloth-zoo does NOT track unsloth's tags (different cadence). Use --zoo-ref + # if given; else the unsloth ref only when the zoo repo has it, falling back to + # main. + _zoo_ref="$ZOO_REF" + if [ -z "$_zoo_ref" ]; then + if git ls-remote --exit-code https://github.com/unslothai/unsloth-zoo.git \ + "$REF" >/dev/null 2>&1; then + _zoo_ref="$REF" + else + _zoo_ref="main" + echo "[studio-update] unsloth-zoo has no ref '${REF}'; using zoo main" + fi + fi + SPECS="$SPECS git+https://github.com/unslothai/unsloth-zoo.git@${_zoo_ref}#egg=unsloth_zoo" + echo "[studio-update] installing from git: unsloth @${REF}, unsloth-zoo @${_zoo_ref}" +else + SPECS="$PACKAGES" + echo "[studio-update] installing latest release of: $PACKAGES" +fi + +# shellcheck disable=SC2086 +"$PY" -m pip install -U $NO_DEPS $SPECS + +echo "[studio-update] after: unsloth $(version_of)" + +# Sanity: the backend must still import after the swap (a missing --no-deps +# transitive dep shows up here). Restarting into code that cannot import kills a +# process that is serving fine and leaves supervisord's studio program in FATAL +# after startretries, which it never leaves on its own. Keep the running service +# and fail instead, so the operator can add the dep or roll back with Studio up. +if ! "$PY" -c "import studio.backend.main" >/dev/null 2>&1; then + echo "[studio-update] ERROR: 'import studio.backend.main' failed after update." >&2 + echo "[studio-update] A new dependency may be missing. Re-run with --with-deps:" >&2 + echo "[studio-update] unsloth-studio-update --with-deps" >&2 + echo "[studio-update] NOT restarting Studio: the running process keeps serving." >&2 + echo "[studio-update] Once fixed: supervisorctl restart studio" >&2 + exit 1 +fi + +if [ "$RESTART" = "1" ]; then + SUPCTL="$(command -v supervisorctl || true)" + [ -n "$SUPCTL" ] || SUPCTL="/opt/unsloth-venv/bin/supervisorctl" + if [ -x "$SUPCTL" ] && "$SUPCTL" status studio >/dev/null 2>&1; then + echo "[studio-update] restarting the studio service" + "$SUPCTL" restart studio + else + echo "[studio-update] supervisor not managing 'studio' here; restart Studio yourself" + echo "[studio-update] (e.g. 'docker restart ')" + fi +else + echo "[studio-update] --no-restart: restart Studio to load the update" + echo "[studio-update] docker exec supervisorctl restart studio" +fi + +echo "[studio-update] done" diff --git a/docker/unsloth_sync_notebooks.sh b/docker/unsloth_sync_notebooks.sh new file mode 100644 index 0000000000..2c23daf984 --- /dev/null +++ b/docker/unsloth_sync_notebooks.sh @@ -0,0 +1,350 @@ +#!/usr/bin/env bash +# Populate and refresh /workspace/unsloth-notebooks from unslothai/notebooks. +# +# On boot this copies the baked read-only template into /workspace/unsloth-notebooks +# (first run), then best-effort refreshes from GitHub when upstream advances. +# +# The user's edits ALWAYS win: each written file's hash is recorded; on refresh a +# file whose hash differs is left untouched. So a refresh only updates unchanged +# files and adds new ones. +# +# Opt-out / tuning (all optional): +# UNSLOTH_SKIP_NOTEBOOK_SYNC=1 do nothing (no populate, no refresh) +# UNSLOTH_SKIP_NOTEBOOK_REFRESH=1 populate from the baked template only; +# never touch the network +# UNSLOTH_KEEP_DELETED_NOTEBOOKS=1 do not restore notebooks the user deleted +# (default: deleted files are healed back) +# UNSLOTH_NOTEBOOKS_DIR= target dir (default /workspace/unsloth-notebooks) +# UNSLOTH_NOTEBOOKS_REPO= source repo (default unslothai/notebooks) +# UNSLOTH_NOTEBOOK_FETCH_TIMEOUT=N seconds for each network op (default 60) +# UNSLOTH_SKIP_NOTEBOOK_VIEW=1 do not build the categorized folder view +# UNSLOTH_NOTEBOOKS_VIEW_DIR= categorized view dir +# (default "/workspace/Unsloth Notebooks") +# UNSLOTH_NB_GPU=amd|cuda force AMD-* notebook visibility (default: +# autodetect; AMD-* shown only on AMD/HIP) +# UNSLOTH_KEEP_COLAB_INTRO=1 keep the Colab "Run all on Colab" sentence +# (default: strip it for the Docker image) +set -u + +TEMPLATE="${UNSLOTH_NOTEBOOKS_TEMPLATE:-/opt/unsloth-notebooks}" +DEST="${UNSLOTH_NOTEBOOKS_DIR:-/workspace/unsloth-notebooks}" +REMOTE="${UNSLOTH_NOTEBOOKS_REPO:-https://github.com/unslothai/notebooks}" +STATE="$DEST/.unsloth_sync_state" # "sha256 relpath" of what we last wrote +SYNCED="$DEST/.unsloth_sync_commit" # upstream commit we last synced to +LOCK="$DEST/.unsloth_sync.lock" # serialises this script against itself +TIMEOUT="${UNSLOTH_NOTEBOOK_FETCH_TIMEOUT:-60}" +LOCK_WAIT="${UNSLOTH_NOTEBOOK_LOCK_TIMEOUT:-600}" + +# Resolve a helper script ($1 override, $2 PATH command, $3 sibling filename), +# echoing the path or nothing. Used for SIG, VIEW and STRIP helpers. +PYBIN="$(command -v python3 2>/dev/null || command -v python 2>/dev/null || true)" +_self_dir="$(cd "$(dirname "$0")" 2>/dev/null && pwd)" +resolve_helper() { + if [ -n "$1" ]; then printf '%s' "$1"; return 0; fi + if command -v "$2" >/dev/null 2>&1; then command -v "$2"; return 0; fi + [ -n "$_self_dir" ] && [ -f "$_self_dir/$3" ] && printf '%s' "$_self_dir/$3" + return 0 +} +SIG_HELPER="$(resolve_helper "${UNSLOTH_NB_SIG_HELPER:-}" unsloth-nb-content-sig unsloth_nb_content_sig.py)" +VIEW_HELPER="$(resolve_helper "${UNSLOTH_NB_VIEW_HELPER:-}" unsloth-nb-view unsloth_nb_view.py)" +STRIP_HELPER="$(resolve_helper "${UNSLOTH_NB_STRIP_HELPER:-}" unsloth-nb-strip-colab unsloth_nb_strip_colab.py)" + +# True only when both are .ipynb and the SIG helper reports the non-boilerplate +# middle identical, so a refresh doesn't rewrite a notebook when only boilerplate +# moved. Any failure returns false. +middle_unchanged() { + case "$1" in *.ipynb) : ;; *) return 1 ;; esac + [ -n "$PYBIN" ] && [ -n "$SIG_HELPER" ] || return 1 + [ "${UNSLOTH_NOTEBOOK_BODY_AWARE:-1}" = "1" ] || return 1 + [ "$("$PYBIN" "$SIG_HELPER" "$1" "$2" 2>/dev/null)" = "SAME" ] || return 1 + return 0 +} + +[ "${UNSLOTH_SKIP_NOTEBOOK_SYNC:-0}" = "1" ] && exit 0 +[ -d "$TEMPLATE" ] || exit 0 +mkdir -p "$DEST" 2>/dev/null || exit 0 + +hash_of() { sha256sum "$1" 2>/dev/null | cut -d' ' -f1; } + +# --- mutual exclusion -------------------------------------------------------- +# Every phase below mutates $DEST and rewrites $STATE, and the GitHub refresh +# runs in a DETACHED child of this same script, so two copies are live at once by +# design. Without a lock the parent's strip/view pass interleaved with the child's +# `cp -a` + state rewrite: six identical boots reported "cleaned" 279/289/293/297/ +# 300/306/307/315/330 notebooks, and every notebook the child copied while the +# parent was hashing it ended up permanently marked user-edited (its recorded +# hash no longer matched), so it was skipped by every later strip. +# +# One exclusive lock covers a whole invocation. The child therefore cannot start +# until the parent has finished and exited, which also fixes the ORDER: strip and +# view rebuild always run over a quiesced tree. flock is best-effort -- when it is +# unavailable, or $DEST cannot hold the lock file, we fall back to running +# unlocked (the parent still finalizes before forking, see below). +_LOCK_HELD=0 +lock_acquire() { + [ "$_LOCK_HELD" = "1" ] && return 0 + command -v flock >/dev/null 2>&1 || return 0 + # Group-redirect, not `exec ... 2>/dev/null`: bash reports a failed exec + # redirection before the redirection it was given applies, so a read-only + # $DEST would print "Permission denied" into the container log. + { exec 9>>"$LOCK"; } 2>/dev/null || return 0 + flock -w "$LOCK_WAIT" 9 2>/dev/null || return 0 + _LOCK_HELD=1 + return 0 +} +lock_release() { + [ "$_LOCK_HELD" = "1" ] || return 0 + _LOCK_HELD=0 + flock -u 9 2>/dev/null || true + exec 9>&- 2>/dev/null || true + return 0 +} + +# --- categorized folder view + Docker-only Colab cleanups -------------------- +# AMD/HIP detection: AMD-*.ipynb are shown only on an AMD GPU. UNSLOTH_NB_GPU +# forces it (amd|cuda); otherwise probe nvidia-smi then the ROCm tools. +nb_gpu_is_amd() { + case "${UNSLOTH_NB_GPU:-}" in + amd|AMD|hip|HIP|rocm|ROCm|ROCM) return 0 ;; + cuda|CUDA|nvidia|NVIDIA|nv|NV) return 1 ;; + esac + if command -v nvidia-smi >/dev/null 2>&1 \ + && nvidia-smi -L 2>/dev/null | grep -q '^GPU'; then + return 1 + fi + if command -v rocm-smi >/dev/null 2>&1 || command -v rocminfo >/dev/null 2>&1; then + return 0 + fi + return 1 # default: treat as non-AMD (hide AMD-* notebooks) +} + +# Rebuild the sibling symlink VIEW from scratch. Symlinks live OUTSIDE $DEST, so +# the sync state machine (find -type f) never sees them. +build_categorized_view() { + [ "${UNSLOTH_SKIP_NOTEBOOK_VIEW:-0}" = "1" ] && return 0 + [ -n "$PYBIN" ] && [ -n "$VIEW_HELPER" ] || return 0 + [ -d "$DEST/nb" ] || return 0 + _view="${UNSLOTH_NOTEBOOKS_VIEW_DIR:-/workspace/Unsloth Notebooks}" + if nb_gpu_is_amd; then + "$PYBIN" "$VIEW_HELPER" "$DEST" "$_view" --amd 2>/dev/null || true + else + "$PYBIN" "$VIEW_HELPER" "$DEST" "$_view" 2>/dev/null || true + fi +} + +# Strip the Colab-only "Run all on Colab" sentence from notebooks WE own and the +# user has not edited (STATE-aware), updating their recorded hashes in place. +strip_colab_intros() { + [ "${UNSLOTH_KEEP_COLAB_INTRO:-0}" = "1" ] && return 0 + [ -n "$PYBIN" ] && [ -n "$STRIP_HELPER" ] || return 0 + [ -f "$STATE" ] || return 0 + "$PYBIN" "$STRIP_HELPER" --state "$STATE" --dest "$DEST" 2>/dev/null || true +} + +# Apply both on EVERY exit after the basic guards, so the view + cleanups also +# run on the common "nothing to refresh" / offline paths. Both are idempotent. +# Run-once: the parent calls this explicitly BEFORE it forks the refresh child +# (so the strip can never overlap the child's copy even where flock is missing), +# and the EXIT trap then has nothing left to do. +_FINALIZED=0 +finalize() { + [ "$_FINALIZED" = "1" ] && return 0 + _FINALIZED=1 + strip_colab_intros + build_categorized_view + return 0 +} +trap 'finalize; lock_release' EXIT + +# Everything past this point mutates $DEST / $STATE, so hold the lock for the +# whole run. A detached refresh child blocks here until its parent has exited. +lock_acquire + +# Record " " for every file currently under DEST (skip metadata). +record_state() { + : > "$STATE.tmp" 2>/dev/null || return 0 + ( cd "$DEST" && find . -type f -print0 ) | while IFS= read -r -d '' rel; do + rel="${rel#./}" + case "$rel" in + .unsloth_sync_state|.unsloth_sync_state.tmp|.unsloth_sync_commit) continue ;; + .unsloth_sync.lock) continue ;; + esac + printf '%s %s\n' "$(hash_of "$DEST/$rel")" "$rel" >> "$STATE.tmp" + done + mv "$STATE.tmp" "$STATE" 2>/dev/null || rm -f "$STATE.tmp" +} + +# 1) First-boot populate from the baked template (instant, works offline). +if [ ! -f "$STATE" ]; then + : > "$STATE.tmp" 2>/dev/null || true + ( cd "$TEMPLATE" && find . -type f -print0 ) | while IFS= read -r -d '' rel; do + rel="${rel#./}" + case "$rel" in .unsloth_template_commit) continue ;; esac + mkdir -p "$DEST/$(dirname "$rel")" 2>/dev/null || true + # A pre-existing file (bind-mounted or hand-created) is user data: keep it + # and do NOT record it, else the refresh below would treat it as pristine + # and overwrite it. Only files we lay down are recorded as managed. + if [ -e "$DEST/$rel" ]; then + if [ "$(hash_of "$DEST/$rel")" != "$(hash_of "$TEMPLATE/$rel")" ]; then + echo "[unsloth-nb] kept existing user file: $DEST/$rel" + continue + fi + # Same bytes already on disk (a bind-mounted checkout of the same + # notebooks). cp -a is --preserve=all, so copying would only stamp the + # baked root:root ownership, mode and build mtime onto the host user's + # own file and lock them out of editing it. Record it as managed -- the + # hash is identical, so the state is byte-for-byte what cp would write. + printf '%s %s\n' "$(hash_of "$DEST/$rel")" "$rel" >> "$STATE.tmp" + continue + fi + if cp -a "$TEMPLATE/$rel" "$DEST/$rel" 2>/dev/null; then + printf '%s %s\n' "$(hash_of "$DEST/$rel")" "$rel" >> "$STATE.tmp" + fi + done + mv "$STATE.tmp" "$STATE" 2>/dev/null || rm -f "$STATE.tmp" + cp -a "$TEMPLATE/.unsloth_template_commit" "$SYNCED" 2>/dev/null || true + echo "[unsloth-nb] notebooks ready at $DEST" +fi + +# 1b) Every-boot OFFLINE restore of deleted notebooks: a file we wrote that the +# user DELETED comes back from the baked template (no network). Existing files are +# never touched. Opt out with UNSLOTH_KEEP_DELETED_NOTEBOOKS=1. +if [ -f "$STATE" ] && [ "${UNSLOTH_KEEP_DELETED_NOTEBOOKS:-0}" != "1" ]; then + restored=0 + RS_TMP="$(mktemp)" + while IFS= read -r line; do + h="${line%% *}"; rel="${line#* }" + if [ -n "$rel" ] && [ "$rel" != "$line" ] \ + && [ ! -e "$DEST/$rel" ] && [ -f "$TEMPLATE/$rel" ]; then + mkdir -p "$DEST/$(dirname "$rel")" 2>/dev/null || true + if cp -a "$TEMPLATE/$rel" "$DEST/$rel" 2>/dev/null; then + printf '%s %s\n' "$(hash_of "$DEST/$rel")" "$rel" >> "$RS_TMP" + restored=$((restored + 1)) + continue + fi + fi + printf '%s\n' "$line" >> "$RS_TMP" + done < "$STATE" + mv "$RS_TMP" "$STATE" 2>/dev/null || rm -f "$RS_TMP" + [ "$restored" -gt 0 ] \ + && echo "[unsloth-nb] restored $restored deleted notebook(s) from the baked set" +fi + +# 2) Best-effort GitHub refresh -- only when upstream has advanced. Edits win. +# Detached: the local populate above already ran, and the refresh can spend up +# to 2x TIMEOUT on ls-remote + clone when offline, which must not delay +# container startup. The child re-enters past phase 1 (hash state makes it a +# no-op) and the flag keeps it from forking again. +[ "${UNSLOTH_SKIP_NOTEBOOK_REFRESH:-0}" = "1" ] && exit 0 +command -v git >/dev/null 2>&1 || exit 0 +command -v sha256sum >/dev/null 2>&1 || exit 0 +if [ "${UNSLOTH_NB_REFRESH_CHILD:-0}" != "1" ]; then + # Finalize BEFORE the fork, not from the EXIT trap after it: the trap used to + # fire while the child was already copying refreshed notebooks in, which is + # what made "cleaned N" differ on every boot. Doing it here also keeps the + # ordering deterministic on hosts without flock. Container startup is not + # delayed any further -- the trap ran exactly this work in the parent before. + finalize + lock_release + UNSLOTH_NB_REFRESH_CHILD=1 "$0" >/dev/null 2>&1 & + exit 0 +fi + +# --- refresh child ----------------------------------------------------------- +# The parent has already stripped + built the view for the tree as it stands, so +# suppress the EXIT-trap finalize; it is re-armed below only if this refresh +# actually rewrites notebooks, which keeps an up-to-date boot a true no-op. +_FINALIZED=1 + +last="$(cat "$SYNCED" 2>/dev/null || true)" +remote="$(timeout "$TIMEOUT" git ls-remote "$REMOTE" HEAD 2>/dev/null | cut -f1)" +[ -z "$remote" ] && exit 0 # offline / unreachable -> keep what we have +[ "$remote" = "$last" ] && exit 0 # nothing new since last sync -> done + +TMP="$(mktemp -d)" +if ! timeout "$TIMEOUT" git clone -q --depth 1 "$REMOTE" "$TMP" 2>/dev/null; then + rm -rf "$TMP"; exit 0 # network died mid-way -> keep what we have +fi + +declare -A LAST +if [ -f "$STATE" ]; then + while read -r h p; do + [ -n "${p:-}" ] && LAST["$p"]="$h" + done < "$STATE" +fi + +TMPSTATE="$(mktemp)" +updated=0; kept=0; unchanged=0 +while IFS= read -r -d '' f; do + rel="${f#"$TMP"/}" + case "$rel" in .git|.git/*) continue ;; esac + dst="$DEST/$rel" + if [ -e "$dst" ]; then + rec="${LAST[$rel]:-}" + if [ -z "$rec" ]; then + # In DEST but never recorded -> a pre-existing user/bind-mounted file. + # Keep it and don't adopt it into the state (stays protected). + kept=$((kept + 1)) + continue + fi + if [ -n "$rec" ] && [ "$(hash_of "$dst")" != "$rec" ]; then + # User changed this file since we wrote it -> keep theirs, keep marker. + printf '%s %s\n' "$rec" "$rel" >> "$TMPSTATE" + kept=$((kept + 1)) + continue + fi + if [ -n "$rec" ] && middle_unchanged "$dst" "$f"; then + # Untouched notebook whose only upstream change is the install header/ + # announcements/footer. Body identical, so keep it and its marker. + printf '%s %s\n' "$rec" "$rel" >> "$TMPSTATE" + unchanged=$((unchanged + 1)) + continue + fi + elif [ -n "${LAST[$rel]:-}" ] && [ "${UNSLOTH_KEEP_DELETED_NOTEBOOKS:-0}" = "1" ]; then + # We wrote this notebook and the user DELETED it; with the opt-out set, + # honor the deletion. Keep the record as managed-but-deleted. + printf '%s %s\n' "${LAST[$rel]}" "$rel" >> "$TMPSTATE" + kept=$((kept + 1)) + continue + fi + mkdir -p "$(dirname "$dst")" 2>/dev/null || true + # Publish through a same-dir temp + rename. This child is forked before the + # entrypoint execs the container command, so JupyterLab is already serving + # $DEST while this loop runs: cp -a writes in place (the inode is reused), so + # a reader can catch half-written JSON, and a save made between the recorded- + # hash check above and this write is destroyed and then recorded as pristine. + # rename(2) is atomic, and re-reading the hash once the temp is complete + # shrinks the check-to-write window to the rename itself. The staging name is + # dot-prefixed and per-PID so a killed refresh leaves nothing visible in the + # file browser; unsloth_nb_strip_colab.py already publishes these same files + # this way. + new="$(dirname "$dst")/.unsloth_nb_new.$$" + if cp -a "$f" "$new" 2>/dev/null; then + if [ -e "$dst" ] && [ "$(hash_of "$dst")" != "${LAST[$rel]:-}" ]; then + # Saved while we were copying -> their edit wins, keep the marker. + rm -f "$new" + printf '%s %s\n' "${LAST[$rel]:-}" "$rel" >> "$TMPSTATE" + kept=$((kept + 1)) + continue + fi + # A single-FILE bind mount cannot be renamed over (EBUSY); fall back to the + # previous in-place copy there so that setup keeps working as it does today. + if mv -f "$new" "$dst" 2>/dev/null || { rm -f "$new"; cp -a "$f" "$dst" 2>/dev/null; }; then + printf '%s %s\n' "$(hash_of "$dst")" "$rel" >> "$TMPSTATE" + updated=$((updated + 1)) + fi + fi +done < <(find "$TMP" -type f -print0) + +mv "$TMPSTATE" "$STATE" 2>/dev/null || rm -f "$TMPSTATE" +echo "$remote" > "$SYNCED" 2>/dev/null || true +rm -rf "$TMP" +echo "[unsloth-nb] notebooks refreshed from GitHub: $updated updated, $kept kept (your edits), $unchanged kept (only header/footer changed upstream)" +# Freshly copied notebooks arrive with the upstream Colab intro, and new files +# have to enter the view, so re-arm the finalize -- but only when something was +# actually copied. Still under the lock, so nothing else is touching the tree. +if [ "$updated" -gt 0 ]; then + _FINALIZED=0 + finalize +fi +exit 0 diff --git a/install.sh b/install.sh index fece7b173b..7cd7ff3976 100755 --- a/install.sh +++ b/install.sh @@ -2270,6 +2270,12 @@ TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.11.0" # ── Resolve repo root (for --local installs) ── _REPO_ROOT="$(cd "$(dirname "$0" 2>/dev/null || echo ".")" && pwd)" +# ── unsloth-zoo overlay ref (for --local installs) ── +# Honor UNSLOTH_ZOO_REF so the Studio venv tracks the requested zoo (the Docker +# publish workflow forwards one ref to both builds). Unset -> main. +_ZOO_REF="${UNSLOTH_ZOO_REF:-main}" +_ZOO_GIT_SPEC="unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo@${_ZOO_REF}" + # ── Helper: find no-torch-runtime.txt (local repo or site-packages) ── _find_no_torch_runtime() { # Check local repo first (for --local installs) @@ -3692,10 +3698,10 @@ if [ "$_MIGRATED" = true ]; then if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps - substep "overlaying unsloth-zoo from git main..." - run_install_cmd_retry "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \ + substep "overlaying unsloth-zoo from git ${_ZOO_REF}..." + run_install_cmd_retry "overlay unsloth-zoo (git ${_ZOO_REF})" uv pip install --python "$_VENV_PY" \ --no-deps --reinstall-package unsloth-zoo \ - "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" + "$_ZOO_GIT_SPEC" fi # AMD ROCm: install bitsandbytes even in migrated environments so # existing ROCm installs gain the AMD bitsandbytes build without a @@ -3930,10 +3936,10 @@ elif [ -n "$TORCH_INDEX_URL" ]; then if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps - substep "overlaying unsloth-zoo from git main..." - run_install_cmd_retry "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \ + substep "overlaying unsloth-zoo from git ${_ZOO_REF}..." + run_install_cmd_retry "overlay unsloth-zoo (git ${_ZOO_REF})" uv pip install --python "$_VENV_PY" \ --no-deps --reinstall-package unsloth-zoo \ - "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" + "$_ZOO_GIT_SPEC" fi elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \ @@ -3941,10 +3947,10 @@ elif [ -n "$TORCH_INDEX_URL" ]; then --upgrade-package unsloth "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps - substep "overlaying unsloth-zoo from git main..." - run_install_cmd_retry "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \ + substep "overlaying unsloth-zoo from git ${_ZOO_REF}..." + run_install_cmd_retry "overlay unsloth-zoo (git ${_ZOO_REF})" uv pip install --python "$_VENV_PY" \ --no-deps --reinstall-package unsloth-zoo \ - "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" + "$_ZOO_GIT_SPEC" else run_install_cmd_retry "install unsloth" uv pip install --python "$_VENV_PY" \ ${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \ @@ -3970,10 +3976,10 @@ else run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.6" "unsloth>=2026.7.5" --torch-backend=auto substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps - substep "overlaying unsloth-zoo from git main..." - run_install_cmd_retry "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \ + substep "overlaying unsloth-zoo from git ${_ZOO_REF}..." + run_install_cmd_retry "overlay unsloth-zoo (git ${_ZOO_REF})" uv pip install --python "$_VENV_PY" \ --no-deps --reinstall-package unsloth-zoo \ - "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" + "$_ZOO_GIT_SPEC" else run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" --torch-backend=auto -- "$PACKAGE_NAME" fi diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 9b787dbb15..f018367926 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -429,6 +429,17 @@ _os_error_messages = _core._os_error_messages is_busy_lock_error = _core.is_busy_lock_error +def is_cross_device_error(exc: BaseException) -> bool: + """True for an EXDEV "cross-device link" rename failure. + + os.replace / os.rename cannot move across filesystems -- e.g. inside a Docker + build where the staging tree and the install dir land on different overlayfs + layers (Errno 18). Unlike a busy/in-use error, a cross-device move is safely + completed by a copy + remove of the (idle) source. + """ + return isinstance(exc, OSError) and exc.errno == errno.EXDEV + + # Status logs default to stderr so resolver modes keep stdout machine-readable # (setup.sh json.load()s the whole stdout). main() flips this for the install # path, where PowerShell otherwise renders stderr as NativeCommandError noise. @@ -3969,13 +3980,51 @@ def activate_staged_dir(staging_dir: Path, dst: Path) -> None: try: os.replace(staging_dir, dst) except OSError as exc: - if not is_busy_lock_error(exc): + # Busy/in-use (Windows AV) or cross-device (Docker overlayfs): both safe to + # complete by copy + remove. Anything else (disk full, missing path) re-raises. + if not (is_busy_lock_error(exc) or is_cross_device_error(exc)): raise log(f"os.replace failed ({exc!r}); falling back to file-by-file copy of staging tree") shutil.copytree(staging_dir, dst, dirs_exist_ok = True) remove_tree(staging_dir) +def move_install_dir_aside(src: Path, dst: Path) -> None: + """Move an existing install dir to ``dst`` (a unique, non-existent sibling). + + os.replace is the fast path. On a cross-device link (EXDEV -- e.g. moving the + base-image llama.cpp aside during a Docker studio build, where the rollback + path is on a different overlay) fall back to copy + remove. A busy/in-use + failure is deliberately NOT copy-faked here: the source is a live install and + a partial copy + rmtree would be worse than failing, so it re-raises. + + The copy never writes into ``dst`` directly: callers treat ``dst.exists()`` + as proof of a complete tree (activation recovery restores a rollback dir + whenever it exists), so a copy that dies halfway (ENOSPC, I/O error) must + not leave a partial tree at ``dst``. Copy to a temp sibling and publish it + with one atomic rename; on failure remove the temp copy and leave ``src`` + untouched. + """ + try: + os.replace(src, dst) + except OSError as exc: + if not is_cross_device_error(exc): + raise + copy_tmp = dst.with_name(dst.name + ".copying") + counter = 0 + while copy_tmp.exists(): + counter += 1 + copy_tmp = dst.with_name(f"{dst.name}.copying-{counter}") + log(f"os.replace cross-device ({exc!r}); copy+publish {src} -> {dst}") + try: + shutil.copytree(src, copy_tmp) + os.replace(copy_tmp, dst) + except BaseException: + remove_tree(copy_tmp) + raise + remove_tree(src) + + def activate_install_tree(staging_dir: Path, install_dir: Path, host: HostInfo) -> None: rollback_dir: Path | None = None failed_dir: Path | None = None @@ -3983,7 +4032,7 @@ def activate_install_tree(staging_dir: Path, install_dir: Path, host: HostInfo) if install_dir.exists(): rollback_dir = unique_install_side_path(install_dir, "rollback") log(f"moving existing install to rollback path {rollback_dir}") - os.replace(install_dir, rollback_dir) + move_install_dir_aside(install_dir, rollback_dir) log(f"moved existing install to rollback path {rollback_dir.name}") log(f"activating staged install {staging_dir} -> {install_dir}") @@ -3998,7 +4047,7 @@ def activate_install_tree(staging_dir: Path, install_dir: Path, host: HostInfo) if install_dir.exists(): failed_dir = unique_install_side_path(install_dir, "failed") log(f"moving failed active install to {failed_dir}") - os.replace(install_dir, failed_dir) + move_install_dir_aside(install_dir, failed_dir) elif staging_dir.exists(): failed_dir = staging_dir staging_dir = None @@ -4006,7 +4055,7 @@ def activate_install_tree(staging_dir: Path, install_dir: Path, host: HostInfo) if rollback_dir and rollback_dir.exists(): log(f"restoring rollback path {rollback_dir} -> {install_dir}") - os.replace(rollback_dir, install_dir) + move_install_dir_aside(rollback_dir, install_dir) log(f"restored previous install from rollback path {rollback_dir.name}") if is_busy_lock_error(exc): raise BusyInstallConflict( diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 2883f30b20..0c062ba770 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -2847,6 +2847,10 @@ def install_python_stack() -> int: package_name = os.environ.get("STUDIO_PACKAGE_NAME", "unsloth") # --local overlays a local repo checkout after updating deps. local_repo = os.environ.get("STUDIO_LOCAL_REPO", "") + # unsloth-zoo git ref for the --local overlay. Honor UNSLOTH_ZOO_REF so the + # Studio venv tracks the requested zoo, not always main. Unset -> main. + zoo_ref = os.environ.get("UNSLOTH_ZOO_REF", "").strip() or "main" + zoo_git_spec = "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo@" + zoo_ref base_total = 11 if IS_WINDOWS else 12 # +1 for the anyio repair check (step 8b) if IS_MACOS: base_total -= 1 # triton step is skipped on macOS @@ -2963,13 +2967,13 @@ def install_python_stack() -> int: local_repo, constrain = False, ) - _step(_LABEL, "overlaying unsloth-zoo from git main") + _step(_LABEL, f"overlaying unsloth-zoo from git {zoo_ref}") pip_install( - "Overlaying unsloth-zoo from git main", + f"Overlaying unsloth-zoo from git {zoo_ref}", "--no-cache-dir", "--no-deps", "--force-reinstall", - "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo", + zoo_git_spec, constrain = False, ) elif local_repo: @@ -2994,13 +2998,13 @@ def install_python_stack() -> int: local_repo, constrain = False, ) - _step(_LABEL, "overlaying unsloth-zoo from git main") + _step(_LABEL, f"overlaying unsloth-zoo from git {zoo_ref}") pip_install( - "Overlaying unsloth-zoo from git main", + f"Overlaying unsloth-zoo from git {zoo_ref}", "--no-cache-dir", "--no-deps", "--force-reinstall", - "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo", + zoo_git_spec, constrain = False, ) elif package_name != "unsloth": diff --git a/tests/python/test_docker_labext_cell_nav.py b/tests/python/test_docker_labext_cell_nav.py new file mode 100644 index 0000000000..4ec3e3eebb --- /dev/null +++ b/tests/python/test_docker_labext_cell_nav.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +"""Colab-style arrow navigation must not swallow wrapped-line movement. + +`cellNav.ts` owns ArrowUp/ArrowDown in the capture phase and jumps to the +previous/next cell when the cursor sits on the first/last line of the editor. +That test used `editor.getCursorPosition().line` against `editor.lineCount`, +both of which are LOGICAL (JupyterLab's CodeMirrorEditor: `get lineCount() { +return this.doc.lines }`), while JupyterLab wraps markdown and raw cell editors +by default (`StaticNotebook.defaultEditorConfig` -> `markdown: { lineWrap: true +}`, `raw: { lineWrap: true }`; the image's `docker/jupyter/overrides.json` only +sets `autoClosingBrackets`). + +So for a one-line markdown header -- what every Unsloth notebook opens with -- +`lineCount === 1`, the cursor is on line 0 == lineCount - 1 from every visual +row, and BOTH arrows leave the cell: the wrapped rows in between cannot be +reached at all. Measured in Chromium with CodeMirror 6 + EditorView.lineWrapping +at the notebook's editor width: 1 logical line renders as 7 visual rows and the +logical test hijacks the arrows on 7 of 7 rows, in both directions. The same +measurement on an unwrapped code cell shows the visual test agreeing with the +logical one on every row, so the Colab-style jump is unchanged there. + +CodeMirror's own answer is `EditorView.moveVertically(range, forward)`, which +moves "to the next line (including wrapped lines)"; it returns the unchanged +head only at offset 0 / doc.length, so a move that stays on the same visual row +(same `coordsAtPos().top`) is the real editor edge. + +Static source guard: the labextension is only built inside Dockerfile.studio +(`jlpm install && jlpm build:prod`), so there is no TS test runner in-repo. +""" + +from __future__ import annotations + +import re +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +CELL_NAV = REPO_ROOT / "docker" / "jupyter" / "unsloth_labext" / "src" / "cellNav.ts" + + +@pytest.fixture(scope = "module") +def source() -> str: + assert CELL_NAV.is_file(), f"missing {CELL_NAV}" + return CELL_NAV.read_text(encoding = "utf-8") + + +def test_the_edit_mode_boundary_test_asks_codemirror_for_a_visual_line(source: str): + assert "moveVertically" in source, ( + "the edit-mode boundary check must ask CodeMirror whether it can still " + "move one VISUAL line (EditorView.moveVertically); a logical lineCount " + "test makes the wrapped rows of a markdown cell unreachable" + ) + + +def test_the_visual_check_compares_screen_rows(source: str): + assert "coordsAtPos" in source, ( + "moveVertically clamps to the document edge instead of returning the " + "same head, so the two positions have to be compared by visual row" + ) + + +def test_the_logical_line_test_is_only_a_fallback(source: str): + body = source[source.index("const editing = notebook.mode === 'edit'") :] + logical = re.search(r"editor\.lineCount - 1", body) + assert logical, "the non-CodeMirror fallback should still exist" + visual = re.search(r"moveVertically", body) + assert visual and visual.start() < logical.start(), ( + "the visual-line test has to run first; the logical one is only for an " + "editor that is not a CodeMirrorEditor" + ) diff --git a/tests/python/test_docker_llama_cuda_backend.py b/tests/python/test_docker_llama_cuda_backend.py new file mode 100644 index 0000000000..d2059bf1a7 --- /dev/null +++ b/tests/python/test_docker_llama_cuda_backend.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +"""Regression guard for the llama.cpp CUDA backend inside the Docker image. + +The portable llama.cpp bundle ships libggml-cuda.so and loads it with dlopen +(ggml_backend_dl), but the bundle does NOT carry the CUDA math libraries it +links against, and the CUDA runtime base image only carries libcudart. With no +libcublas on the loader path the backend fails to load SILENTLY: llama.cpp +prints nothing, `--list-devices` comes back empty and every GGUF request runs on +the CPU. Measured on a B200 with gemma-4-E2B UD-Q4_K_XL: 1.6 tok/s instead of +224 tok/s, a 140x regression that no functional test would have caught. + +The Dockerfile therefore has to do two things, and these tests pin both: + * put torch's bundled libcublas on the loader path (ld.so.conf.d, not + LD_LIBRARY_PATH, so llama.cpp's own $ORIGIN libs keep winning); + * fail the build when any non-driver dependency of libggml-cuda.so is still + unresolved, so a CPU-only image can never be published again. + +Static: parses the Dockerfile only. No docker, no GPU, no network. +""" + +from __future__ import annotations + +import re +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +DOCKERFILE = REPO_ROOT / "docker" / "Dockerfile" + + +@pytest.fixture(scope = "module") +def dockerfile() -> str: + assert DOCKERFILE.is_file(), f"missing {DOCKERFILE}" + return DOCKERFILE.read_text() + + +def test_cublas_dir_is_registered_with_the_loader(dockerfile: str): + conf = re.search( + r"ld\.so\.conf\.d/zz-unsloth-venv\.conf", + dockerfile, + ) + assert conf, "the venv loader-config layer disappeared" + block = dockerfile[: conf.end()] + assert "$SP/nvidia/cublas/lib" in block, ( + "libggml-cuda.so links against libcublas, which only exists in the venv's " + "wheel copy; without this entry the CUDA backend fails to dlopen and GGUF " + "silently runs on the CPU" + ) + + +def test_loader_config_is_not_ld_library_path(dockerfile: str): + # LD_LIBRARY_PATH is consulted BEFORE DT_RUNPATH, so it would let the venv's + # copies shadow llama.cpp's own $ORIGIN libs. ld.so.conf.d is consulted after. + assert "ld.so.conf.d/zz-unsloth-venv.conf" in dockerfile + assert not re.search( + r"ENV\s+LD_LIBRARY_PATH=.*site-packages/nvidia", + dockerfile, + ), "the venv nvidia libs must not go on LD_LIBRARY_PATH" + + +def test_build_fails_on_an_unresolved_cuda_backend(dockerfile: str): + assert "libggml-cuda.so" in dockerfile, "the CUDA backend guard disappeared" + guard = dockerfile[dockerfile.index("CUDA_SO=") :] + assert "ldd" in guard, "the guard must inspect the backend's dependencies" + assert "not found" in guard + assert "exit 1" in guard, "an unresolved backend must fail the build" + # The driver stub is injected by nvidia-container-toolkit at `docker run + # --gpus`, so it is never resolvable inside the build and must be exempt. + assert re.search( + r"grep -v .libcuda\\?\.so\\?\.1", guard + ), "libcuda.so.1 must be exempt from the guard or every build fails" + + +def test_guard_installs_the_matching_cublas_major(dockerfile: str): + # The amd64 bundle is CUDA 12 and torch already ships libcublas.so.12, but + # the arm64 bundle is CUDA 13. Deriving the major from ldd keeps the two + # legs correct without hardcoding either. + guard = dockerfile[dockerfile.index("CUDA_SO=") :] + assert ( + "nvidia-cublas-cu${major}" in guard + ), "the guard must install the cublas major the bundle actually asks for" + assert "libcublas" in guard + + +def test_guard_runs_after_the_prebuilt_is_fetched(dockerfile: str): + fetch = dockerfile.index("fetch_llama_prebuilt.py") + guard = dockerfile.index("CUDA_SO=") + assert fetch < guard, "the guard can only inspect a bundle that already exists" + + +def test_flashinfer_jit_cache_tracks_flashinfer(dockerfile: str): + # flashinfer raises at import when flashinfer-jit-cache and flashinfer-python + # disagree, and that exception kills the vLLM EngineCore, which is what + # Unsloth's GRPO fast_inference path runs on. A literal pin drifts the moment + # vLLM bumps its flashinfer requirement, so the version has to be derived. + assert ( + "flashinfer-jit-cache==${FI_VER}" in dockerfile + ), "flashinfer-jit-cache must be pinned to the resolved flashinfer-python version" + assert not re.search( + r"flashinfer-jit-cache==[0-9]", dockerfile + ), "a literal flashinfer-jit-cache version will drift away from flashinfer-python" + assert "import flashinfer" in dockerfile, ( + "the build must prove flashinfer imports, or a mismatch stays silent " + "until the first vLLM engine start" + ) + + +def test_cli_can_reach_the_studio_backend(dockerfile: str): + # unsloth_cli's train / export / chat / list-checkpoints import + # studio.backend.core.*, which needs structlog. It is a studio backend + # requirement rather than an unsloth[huggingface] one, so the base venv has + # to ask for it explicitly or the whole CLI dies on ModuleNotFoundError. + assert '"structlog"' in dockerfile, "the base venv must install structlog for unsloth_cli" + assert ( + "from studio.backend.core.export import ExportBackend" in dockerfile + ), "a build-time import guard must prove the CLI can reach the studio backend" diff --git a/tests/python/test_docker_nb_strip_colab_race.py b/tests/python/test_docker_nb_strip_colab_race.py new file mode 100644 index 0000000000..3af2590e12 --- /dev/null +++ b/tests/python/test_docker_nb_strip_colab_race.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +"""The Colab-intro cleanup must not overwrite a save it did not see. + +`unsloth_sync_notebooks.sh` forks the GitHub refresh into a DETACHED child before +the entrypoint execs the container command, so JupyterLab is already serving +$DEST while that child runs. When the refresh copied anything the child re-arms +`finalize()`, which runs `unsloth_nb_strip_colab.py --state ... --dest ...`, i.e. +`migrate()` -> `strip_notebook()` over every owned+unedited notebook. + +`strip_notebook` read the file, parsed it, serialised the cleaned copy and then +`os.replace`d it unconditionally. A user save that landed in that window was +destroyed, and `migrate` then recorded the cleaned file's hash, so the state +machine treats the notebook as pristine forever after -- the same +check-then-write hole that was closed in the refresh loop itself (the publish +there now re-reads the hash immediately before the rename). + +Behavioural: the save is injected inside the window, while the helper serialises +the cleaned copy (the widest part of it: json parse + dump of a notebook that is +often megabytes). No docker, no network. +""" + +from __future__ import annotations + +import copy +import importlib.util +import json +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +STRIP_PATH = REPO_ROOT / "docker" / "unsloth_nb_strip_colab.py" + +INTRO = 'To run this, press "*Runtime*" and press "*Run all*" on a **free** Tesla T4 Google Colab instance!\n' + + +@pytest.fixture(scope = "module") +def strip(): + assert STRIP_PATH.is_file(), f"missing {STRIP_PATH}" + spec = importlib.util.spec_from_file_location("unsloth_nb_strip_race", STRIP_PATH) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +def notebook(*sources): + return { + "cells": [ + {"cell_type": "markdown", "metadata": {}, "source": list(src)} for src in sources + ], + "metadata": {}, + "nbformat": 4, + "nbformat_minor": 5, + } + + +def write(path: Path, nb) -> None: + path.write_text(json.dumps(nb, indent = 1, ensure_ascii = False) + "\n", encoding = "utf-8") + + +@pytest.fixture +def racing(strip, tmp_path: Path): + """Fire a user save inside the window: after strip_notebook read the file, + while it is serialising the cleaned copy.""" + real_dump = strip.json.dump + state = {"save": None, "path": None, "fired": 0} + + def dump(obj, fp, *args, **kwargs): + out = real_dump(obj, fp, *args, **kwargs) + if state["save"] is not None and state["fired"] == 0: + state["fired"] = 1 + Path(state["path"]).write_text(state["save"], encoding = "utf-8") # Ctrl+S + return out + + strip.json.dump = dump + try: + yield state + finally: + strip.json.dump = real_dump + + +def test_a_save_during_the_cleanup_is_not_overwritten(strip, racing, tmp_path: Path): + path = tmp_path / "Llama.ipynb" + write(path, notebook([INTRO, "\n", "# Llama\n"])) + + edited = notebook([INTRO, "\n", "# Llama\n", "\n", "my own notes, saved from JupyterLab\n"]) + racing["save"] = json.dumps(edited, indent = 1, ensure_ascii = False) + "\n" + racing["path"] = str(path) + + strip.strip_notebook(str(path)) + + on_disk = json.loads(path.read_text(encoding = "utf-8")) + assert on_disk == edited, ( + "the user's save landed after strip_notebook read the file and was " + "overwritten by the cleaned copy of the OLD content; the sync contract " + "is that user edits always win" + ) + + +def test_the_recorded_hash_still_matches_the_file_after_a_racing_save( + strip, racing, tmp_path: Path +): + # migrate() rewrites STATE with the post-strip hash. If the write above is + # allowed to clobber a save, the state ALSO says "pristine", so every later + # refresh happily overwrites the notebook again. + dest = tmp_path / "unsloth-notebooks" + dest.mkdir() + path = dest / "Llama.ipynb" + write(path, notebook([INTRO, "\n", "# Llama\n"])) + before = strip._sha256(str(path)) + state = tmp_path / ".unsloth_sync_state" + state.write_text(f"{before} Llama.ipynb\n", encoding = "utf-8") + + edited = notebook([INTRO, "\n", "# Llama\n", "\n", "my own notes\n"]) + racing["save"] = json.dumps(edited, indent = 1, ensure_ascii = False) + "\n" + racing["path"] = str(path) + + strip.migrate(str(state), str(dest)) + + recorded = state.read_text(encoding = "utf-8").split(" ", 1)[0] + on_disk = strip._sha256(str(path)) + assert json.loads(path.read_text(encoding = "utf-8")) == edited + assert recorded != on_disk, ( + "a file the user saved during the cleanup must NOT end up recorded as " + "managed-and-pristine, or the next refresh overwrites it too" + ) + + +def test_the_normal_no_race_cleanup_still_strips_and_rewrites(strip, tmp_path: Path): + # Guard the fix from over-reaching: with nobody else writing, the cleanup + # must still strip the Colab sentence and publish the result. + path = tmp_path / "Llama.ipynb" + original = notebook([INTRO, "\n", "# Llama\n"]) + write(path, copy.deepcopy(original)) + + assert strip.strip_notebook(str(path)) is True + cleaned = json.loads(path.read_text(encoding = "utf-8")) + assert cleaned["cells"][0]["source"] == ["# Llama\n"] + assert strip.strip_notebook(str(path)) is False # idempotent diff --git a/tests/python/test_docker_nb_strip_colab_scope.py b/tests/python/test_docker_nb_strip_colab_scope.py new file mode 100644 index 0000000000..6eeb457e67 --- /dev/null +++ b/tests/python/test_docker_nb_strip_colab_scope.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +"""Regression guard for the Colab-intro strip in the Unsloth Docker image. + +Every generated Unsloth notebook opens with a Colab-only instruction ("To run +this, press Runtime > Run all ...") that is wrong inside Docker, so the image +strips it at sync time. The strip only ever inspected cells[0], and that missed +23 of the 433 shipped notebooks: + + * 21 put the Colab badge `` in + cells[0] and the sentence in cells[1] -- Advanced_Llama3_2_(3B)_GRPO_LoRA, + Falcon_H1-Alpaca, FunctionGemma_(270M)-LMStudio, gpt-oss-(20B)-GRPO, ... + * 2 (NeMo-Gym-Multi-Environment, NeMo-Gym-Sudoku) wrap the sentence in a + single-line HTML comment, so a "line starts with the sentence" match never + fired even though the sentence IS in cells[0]. + +Measured against the pristine baked template: a cells[0]-only strip left 23 of +433 notebooks carrying the line, a leading-markdown-block strip leaves 0, and +neither changes unsloth_nb_content_sig's middle digest for any of the 433 (which +matters, because a changed digest makes the boot refresh re-copy and re-strip the +notebook forever). + +The widening also has to stay narrow: the scan stops at the first non-markdown +cell so it can never reach explanatory prose between code cells, and it stays +idempotent so a second boot is a no-op. + +Static: imports the helper and feeds it in-memory notebooks. No docker, no GPU, +no network. +""" + +from __future__ import annotations + +import copy +import importlib.util +import json +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +STRIP_PATH = REPO_ROOT / "docker" / "unsloth_nb_strip_colab.py" + +INTRO = 'To run this, press "*Runtime*" and press "*Run all*" on a **free** Tesla T4 Google Colab instance!\n' +BADGE = 'badge\n' + + +@pytest.fixture(scope = "module") +def strip(): + assert STRIP_PATH.is_file(), f"missing {STRIP_PATH}" + spec = importlib.util.spec_from_file_location("unsloth_nb_strip_under_test", STRIP_PATH) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +def md(*lines): + return {"cell_type": "markdown", "metadata": {}, "source": list(lines)} + + +def code(src): + return { + "cell_type": "code", + "metadata": {}, + "execution_count": None, + "outputs": [], + "source": [src], + } + + +def nb(*cells): + return {"cells": list(cells), "metadata": {}, "nbformat": 4, "nbformat_minor": 5} + + +def text(cell): + src = cell.get("source", "") + return "".join(src) if isinstance(src, list) else src + + +def has_intro(notebook): + return any("to run this, press" in text(c).lower() for c in notebook["cells"]) + + +def test_intro_in_cell_zero_is_still_stripped(strip): + # The 386-notebook majority case must not regress. + doc = nb(md(INTRO, "\n", BADGE), code("print(1)")) + assert strip._strip_intro(doc) is True + assert not has_intro(doc) + assert BADGE in text(doc["cells"][0]), "the badge row must survive the strip" + + +def test_intro_in_cell_one_behind_the_badge_is_stripped(strip): + # 21 shipped notebooks; a cells[0]-only scan left every one of them. + doc = nb(md(BADGE), md(INTRO, "\n", "You will learn how to do data prep.\n"), code("print(1)")) + assert strip._strip_intro(doc) is True + assert not has_intro(doc) + assert "You will learn how to do data prep.\n" in text(doc["cells"][1]) + + +def test_intro_inside_a_single_line_html_comment_is_stripped(strip): + # NeMo-Gym-Multi-Environment / NeMo-Gym-Sudoku ship exactly this shape. + commented = "\n" + doc = nb(md(commented, '
\n'), code("print(1)")) + assert strip._strip_intro(doc) is True + assert not has_intro(doc) + assert '
\n' in text(doc["cells"][0]) + + +def test_multi_line_html_comment_is_left_alone(strip): + # A comment that does NOT close on the same line must not be half-removed, + # or the surviving `\n"), code("print(1)")) + assert strip._strip_intro(doc) is False + assert has_intro(doc) + + +def test_strip_stops_at_the_first_code_cell(strip): + # A markdown cell AFTER code is prose, not the header block: never touched. + later = md("Explanation.\n", INTRO) + doc = nb(md(BADGE), code("print(1)"), later) + assert strip._strip_intro(doc) is False + assert text(doc["cells"][2]) == "Explanation.\n" + INTRO + + +def test_strip_is_idempotent(strip): + doc = nb(md(BADGE), md(INTRO, "\n", "rest\n"), code("print(1)")) + assert strip._strip_intro(doc) is True + once = copy.deepcopy(doc) + assert strip._strip_intro(doc) is False, "a second boot must be a no-op" + assert doc == once + + +def test_a_notebook_without_the_intro_is_untouched(strip): + doc = nb(md(BADGE, "# Title\n"), code("print(1)")) + before = copy.deepcopy(doc) + assert strip._strip_intro(doc) is False + assert doc == before + + +def test_source_given_as_a_string_is_handled(strip): + doc = nb( + {"cell_type": "markdown", "metadata": {}, "source": BADGE}, + {"cell_type": "markdown", "metadata": {}, "source": INTRO + "\nrest\n"}, + code("print(1)"), + ) + assert strip._strip_intro(doc) is True + assert not has_intro(doc) + assert isinstance(doc["cells"][1]["source"], str) + + +def test_end_to_end_write_back_is_valid_json(strip, tmp_path): + p = tmp_path / "N.ipynb" + p.write_text(json.dumps(nb(md(BADGE), md(INTRO, "\n", "rest\n"), code("print(1)")))) + assert strip.strip_notebook(str(p)) is True + reloaded = json.loads(p.read_text()) + assert not has_intro(reloaded) + assert strip.strip_notebook(str(p)) is False diff --git a/tests/python/test_docker_nb_sync_race.py b/tests/python/test_docker_nb_sync_race.py new file mode 100644 index 0000000000..a7347edf83 --- /dev/null +++ b/tests/python/test_docker_nb_sync_race.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +"""Regression guard for the notebook-sync race in the Unsloth Docker image. + +unsloth_sync_notebooks.sh populates /workspace/unsloth-notebooks on boot and then +refreshes from GitHub in a DETACHED child, so container start is never blocked on +a network fetch. The parent forked that child and exited immediately, which fired +its `trap finalize EXIT` -- the Colab-intro strip plus the categorized-view +rebuild -- while the child was concurrently `cp -a`-ing refreshed notebooks into +the same tree and rewriting the same state file. Both processes also ran +build_categorized_view, which tears down and rebuilds the symlink farm. + +Six identical fresh-container boots reported "cleaned" 279 / 289 / 293 / 297 / +300 / 306 / 307 / 315 / 330 notebooks; two consecutive `docker run`s of the same +image printed 378 and 372. Worse than the noise, the lost writes were permanent: +a notebook the child copied while the parent was hashing it ended up with a +recorded hash that no longer matched the file, so the strip treated it as +user-edited and skipped it on every later boot. That is where 10 of the 23 +notebooks still carrying the Colab intro came from. Setting +UNSLOTH_SKIP_NOTEBOOK_REFRESH=1 -- i.e. never forking the child -- made the +result stable and correctly idempotent, which is what pinned the cause. + +The fix keeps the refresh detached and fixes the ORDERING instead: one exclusive +lock covers a whole invocation so the child cannot start work until the parent +has exited, the parent runs the finalize explicitly BEFORE it forks (so the order +holds even on a host without flock), the finalize is run-once, and the child +re-arms it only when the refresh actually copied something. + +Static: parses the shell script. No docker, no GPU, no network. +""" + +from __future__ import annotations + +import re +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +SYNC = REPO_ROOT / "docker" / "unsloth_sync_notebooks.sh" + + +@pytest.fixture(scope = "module") +def sync() -> str: + assert SYNC.is_file(), f"missing {SYNC}" + return SYNC.read_text() + + +def test_the_refresh_is_still_detached(sync: str): + # The whole point of the child is that a 60s ls-remote + clone must not delay + # container startup. A fix that simply made the refresh synchronous would + # pass every other test here and regress boot time. + assert re.search( + r'UNSLOTH_NB_REFRESH_CHILD=1 "\$0" >/dev/null 2>&1 &', sync + ), "the GitHub refresh must stay a detached child" + + +def test_an_exclusive_lock_serialises_the_two_processes(sync: str): + assert "lock_acquire()" in sync and "lock_release()" in sync + assert re.search( + r"flock -w \"\$LOCK_WAIT\" 9", sync + ), "the lock must be a real exclusive flock, and must not block forever" + + +def test_the_lock_is_taken_before_anything_mutates_the_tree(sync: str): + lock = sync.index("\nlock_acquire\n") + populate = sync.index("# 1) First-boot populate") + assert lock < populate, ( + "populate / restore / refresh all rewrite the state file; the lock has to " + "cover them, not just the strip" + ) + + +def test_a_missing_flock_degrades_instead_of_hanging(sync: str): + block = sync[sync.index("lock_acquire()") : sync.index("lock_release()")] + assert "command -v flock" in block and "return 0" in block, ( + "a host without flock, or a $DEST that cannot hold the lock file, must " + "fall back to running unlocked rather than failing the boot" + ) + + +def test_the_parent_finalizes_before_it_forks(sync: str): + fork = sync.index('UNSLOTH_NB_REFRESH_CHILD=1 "$0"') + block = sync[sync.index('if [ "${UNSLOTH_NB_REFRESH_CHILD:-0}" != "1" ]; then') : fork] + assert re.search(r"^\s*finalize\s*$", block, re.M), ( + "the strip and view rebuild must be done BEFORE the child exists; running " + "them from the EXIT trap after the fork is the race itself" + ) + + +def test_finalize_runs_at_most_once(sync: str): + block = sync[sync.index("finalize() {") : sync.index("trap 'finalize; lock_release' EXIT")] + assert ( + '[ "$_FINALIZED" = "1" ] && return 0' in block + ), "the explicit pre-fork call and the EXIT trap must not strip twice" + assert "_FINALIZED=1" in block + + +def test_the_exit_trap_still_covers_the_early_exits(sync: str): + # Offline / no-git / UNSLOTH_SKIP_NOTEBOOK_REFRESH all exit before the fork + # site, and still need the view built. + assert "trap 'finalize; lock_release' EXIT" in sync + + +def test_the_child_does_not_repeat_the_parents_finalize(sync: str): + tail = sync[sync.index("# --- refresh child ---") :] + assert re.search(r"^_FINALIZED=1\s*$", tail, re.M), ( + "the parent already stripped and built the view for the tree as it " + "stands; an unconditional second pass makes an up-to-date boot noisy" + ) + + +def test_the_child_re_arms_the_finalize_only_after_it_copies(sync: str): + tail = sync[sync.index("refreshed from GitHub") :] + assert re.search( + r'if \[ "\$updated" -gt 0 \]; then\s*\n\s*_FINALIZED=0\s*\n\s*finalize', tail + ), ( + "freshly copied notebooks arrive with the upstream Colab intro and have " + "to be stripped, but only when something was actually copied" + ) + + +def test_the_lock_file_is_not_recorded_as_a_notebook(sync: str): + block = sync[sync.index("record_state() {") :] + block = block[: block.index("\n}")] + assert ".unsloth_sync.lock) continue" in block, ( + "the lock file lives in $DEST next to the state file and must be excluded " + "from the managed-file state like the other metadata" + ) + + +def test_the_lock_lives_beside_the_state_it_protects(sync: str): + assert re.search(r'^LOCK="\$DEST/\.unsloth_sync\.lock"', sync, re.M), ( + "keeping the lock in $DEST also serialises two containers sharing the " + "notebooks volume, which /tmp would not" + ) + + +# --- concurrent-publish safety ------------------------------------------------ +# The detach above is deliberate, but entrypoint.sh runs `sync_notebooks` and then +# `exec "$@"`, so the child is still copying while JupyterLab serves the same tree. +# `cp -a` writes THROUGH the destination inode, so it both exposes half-written +# JSON to a reader and destroys a save made after the recorded-hash check. The +# publish therefore has to go via a same-dir temp plus an atomic rename. + + +def test_the_refresh_publishes_each_notebook_atomically(sync: str): + block = sync[sync.index("while IFS= read -r -d '' f; do") :] + block = block[: block.index("done < <(find")] + assert re.search( + r'cp -a "\$f" "\$new"', block + ), "the refresh must copy into a staging file, not onto the live notebook" + assert re.search( + r'mv -f "\$new" "\$dst"', block + ), "the staged copy must be published with an atomic rename" + + +def test_the_staging_file_is_hidden_and_beside_the_destination(sync: str): + assert re.search(r'new="\$\(dirname "\$dst"\)/\.unsloth_nb_new\.\$\$"', sync), ( + "the staging file must be dot-prefixed (invisible in the file browser), " + "per-PID (two containers on one volume) and in the destination directory " + "(a rename cannot cross filesystems)" + ) + + +def test_the_recorded_hash_is_rechecked_immediately_before_publishing(sync: str): + block = sync[sync.index("while IFS= read -r -d '' f; do") :] + block = block[: block.index("done < <(find")] + recheck = block.index('cp -a "$f" "$new"') + assert re.search( + r'if \[ -e "\$dst" \] && \[ "\$\(hash_of "\$dst"\)" != "\$\{LAST\[\$rel\]:-\}" \]', + block[recheck:], + ), ( + "the earlier check sits before middle_unchanged (a python subprocess), so " + "the hash has to be re-read once the staging copy is complete or a save " + "made in between is silently overwritten" + ) + + +def test_a_pristine_pre_existing_file_is_not_rewritten_on_first_boot(sync: str): + block = sync[sync.index('if [ ! -f "$STATE" ]; then') :] + block = block[: block.index('mv "$STATE.tmp" "$STATE"')] + assert "kept existing user file" in block + # A bind-mounted file whose bytes already match the template used to fall + # through to `cp -a`, i.e. --preserve=all stamping root:root, the baked mode + # and the build mtime onto the host user's own file. Record, don't copy. + same = block.index("kept existing user file") + tail = block[same:] + assert tail.index("$STATE.tmp") < tail.index('cp -a "$TEMPLATE/$rel"'), ( + "an existing file with the template's exact bytes must be recorded as " + "managed without being copied over" + ) diff --git a/tests/python/test_docker_nb_view_ownership.py b/tests/python/test_docker_nb_view_ownership.py new file mode 100644 index 0000000000..5e645f197b --- /dev/null +++ b/tests/python/test_docker_nb_view_ownership.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +"""The categorized notebook VIEW may only delete the links it created. + +`unsloth_nb_view.py` rebuilds "/workspace/Unsloth Notebooks" on every boot, and +that directory is also JupyterLab's landing dir, so `_clear_view()` promises to +remove only the tool's own symlinks. Every link the tool creates points at +DEST/nb/, but the ownership predicate accepted ANY target under DEST, so a +user's own symlink into the notebooks checkout -- e.g. a shortcut to their own +notebook saved beside it, which the sync script explicitly supports ("kept +existing user file" / "In DEST but never recorded") -- was classified as +tool-owned and deleted on the next boot. + +Behavioural: builds a real DEST/VIEW pair on disk and runs build_view twice. +No docker, no network. +""" + +from __future__ import annotations + +import importlib.util +import os +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +VIEW_PATH = REPO_ROOT / "docker" / "unsloth_nb_view.py" + +README = ( + "### Main Notebooks\n" + "[Llama](nb/Llama3_2_%281B_and_3B%29_Conversational.ipynb)\n" + "### Gemma\n" + "[Gemma](nb/Gemma3_%284B%29.ipynb)\n" +) + + +@pytest.fixture(scope = "module") +def view_mod(): + assert VIEW_PATH.is_file(), f"missing {VIEW_PATH}" + spec = importlib.util.spec_from_file_location("unsloth_nb_view_under_test", VIEW_PATH) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +@pytest.fixture +def tree(tmp_path: Path): + dest = tmp_path / "unsloth-notebooks" + view = tmp_path / "Unsloth Notebooks" + (dest / "nb").mkdir(parents = True) + view.mkdir() + for name in ("Llama3_2_(1B_and_3B)_Conversational.ipynb", "Gemma3_(4B).ipynb"): + (dest / "nb" / name).write_text("{}", encoding = "utf-8") + (dest / "README.md").write_text(README, encoding = "utf-8") + # The user's own notebook, saved inside the checkout (supported by the sync + # script), plus their own folder of shortcuts in the landing dir. + (dest / "my_work").mkdir() + (dest / "my_work" / "experiment.ipynb").write_text("{}", encoding = "utf-8") + return dest, view + + +def link(target: Path, at: Path) -> None: + at.parent.mkdir(parents = True, exist_ok = True) + os.symlink(os.path.relpath(target, at.parent), at) + + +def test_a_user_link_to_their_own_file_in_the_checkout_survives(view_mod, tree): + dest, view = tree + own = view / "00 My favourites" / "experiment.ipynb" + link(dest / "my_work" / "experiment.ipynb", own) + + view_mod.build_view(str(dest), str(view)) + + assert os.path.islink(own), ( + "a symlink the user created in the landing dir, pointing at their own " + "file inside the notebooks checkout, was deleted by _clear_view" + ) + assert os.path.realpath(own) == os.path.realpath(dest / "my_work" / "experiment.ipynb") + + +def test_a_user_link_outside_the_checkout_survives(view_mod, tree, tmp_path: Path): + dest, view = tree + outside = tmp_path / "datasets" + outside.mkdir() + own = view / "datasets" + link(outside, own) + + view_mod.build_view(str(dest), str(view)) + + assert os.path.islink(own) + + +def test_the_tools_own_stale_links_are_still_cleaned_up(view_mod, tree): + dest, view = tree + view_mod.build_view(str(dest), str(view)) + generated = view / "02 Gemma" / "Gemma3_(4B).ipynb" + assert os.path.islink(generated) + + # Upstream drops the notebook: its generated link (now stale, and pointing + # into DEST/nb) has to go, and the emptied folder with it. + (dest / "nb" / "Gemma3_(4B).ipynb").unlink() + (dest / "README.md").write_text( + "### Main Notebooks\n[Llama](nb/Llama3_2_%281B_and_3B%29_Conversational.ipynb)\n", + encoding = "utf-8", + ) + view_mod.build_view(str(dest), str(view)) + + assert not os.path.islink(generated) and not os.path.exists(generated) + assert not (view / "02 Gemma").exists() + + +def test_a_rebuild_is_stable_for_the_links_it_owns(view_mod, tree): + dest, view = tree + view_mod.build_view(str(dest), str(view)) + first = sorted(str(p.relative_to(view)) for p in view.rglob("*")) + view_mod.build_view(str(dest), str(view)) + assert sorted(str(p.relative_to(view)) for p in view.rglob("*")) == first diff --git a/tests/python/test_docker_pip_shim_training_stack.py b/tests/python/test_docker_pip_shim_training_stack.py new file mode 100644 index 0000000000..e04548a6b5 --- /dev/null +++ b/tests/python/test_docker_pip_shim_training_stack.py @@ -0,0 +1,235 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +"""Regression guard for what the Docker pip shim protects. + +The shim fronts pip/uv inside the notebook kernel so a `!pip install` cell cannot +replace the baked, ABI-matched stack. It protected torch/vLLM/unsloth and stopped +there, which left the training stack wide open. Measured over the 433 shipped +notebooks (probe_notebook_pins.py against the baked image): + + trl 382 notebooks pin an older release -- 378 of them end their + install cell with `!pip install --no-deps trl==0.22.2`, against + a baked and tested trl 0.24.0 + torchao 273 reinstall it, 2 pin 0.15.0, replacing 0.17.0+cu128 with a + generic PyPI build + torchcodec 92 reinstall it, 26 pin 0.5 / 0.7.0, replacing the 0.11.0+cu128 + wheel the Dockerfile deliberately paired with torch 2.11 + datasets 254 reinstall it; a trl 0.22.2 resolve was observed pulling it + back from 4.3.0 to 3.0.0 + peft 225 reinstall it; observed dropping 0.19.1 -> 0.14.0 + accelerate 225 reinstall it + hf_hub 240 reinstall it, tokenizers 64 -- both version-locked to + transformers, and the sidecars ship their own matched copies + +So EVERY notebook run silently mutated the stack the image was validated with, +and printed "Successfully installed trl-0.22.2 peft-0.14.0 datasets-3.0.0" while +the shim reported it was keeping the baked versions. + +The criterion for _KEEP is "replacing this invalidates the tested stack or breaks +unsloth", not "any package a notebook mentions": a package the notebook genuinely +needs and the image does not bake still has to install normally. + +Static: drives the shim's main() with os.execv captured. No docker, no GPU, no +network. +""" + +from __future__ import annotations + +import importlib.util +import os +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +SHIM_PATH = REPO_ROOT / "docker" / "unsloth_pip_shim.py" + +# The install cell 378 of the 433 shipped notebooks actually end on. +SHIPPED_TRL_CELL = ["--no-deps", "trl==0.22.2"] +# A package the image does NOT bake: must keep installing normally. +UNBAKED = "snac" + + +class _Exec(Exception): + def __init__(self, path, argv): + self.path = path + self.argv = list(argv) + + +@pytest.fixture() +def shim(tmp_path, monkeypatch): + monkeypatch.setenv("UNSLOTH_NB_TF_MARKER", str(tmp_path / "requested_transformers")) + monkeypatch.setenv("UNSLOTH_NB_SHIM", "1") + assert SHIM_PATH.is_file(), f"missing shim: {SHIM_PATH}" + spec = importlib.util.spec_from_file_location("unsloth_pip_shim_stack_test", SHIM_PATH) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + + def _fake_execv(path, argv): + raise _Exec(path, argv) + + monkeypatch.setattr(mod.os, "execv", _fake_execv) + return mod + + +def _run( + shim, + args, + tool = "pip", +): + """Return the args that reached the real tool after `install`, or None when + the shim no-op'd. The always-injected protected-constraints pair is dropped.""" + argv = ["uv", "pip", "install", *args] if tool == "uv" else ["pip", "install", *args] + with pytest.MonkeyPatch.context() as mp: + mp.setattr(shim.sys, "argv", argv) + try: + shim.main() + return None + except _Exec as exc: + i = exc.argv.index("install") + execd = exc.argv[i + 1 :] + if ( + len(execd) >= 2 + and execd[-2] == "--constraint" + and os.path.basename(execd[-1]).startswith("unsloth-nb-protected-") + ): + execd = execd[:-2] + return execd + + +# -------------------------------------------------------------------------- +# Membership +# -------------------------------------------------------------------------- +@pytest.mark.parametrize( + "pkg", + [ + "trl", + "peft", + "datasets", + "accelerate", + "torchao", + "torchcodec", + "huggingface-hub", + "tokenizers", + "safetensors", + ], +) +def test_training_stack_is_protected(shim, pkg): + assert ( + pkg in shim._KEEP + ), f"{pkg} is baked and tested; a notebook pin replacing it invalidates the image" + + +def test_the_original_gpu_stack_is_still_protected(shim): + for pkg in [ + "torch", + "torchvision", + "torchaudio", + "triton", + "xformers", + "vllm", + "bitsandbytes", + "unsloth", + "unsloth-zoo", + ]: + assert pkg in shim._KEEP + + +def test_unrelated_packages_are_not_swept_in(shim): + # The criterion is "invalidates the tested stack", not "a notebook mentions + # it". These are all installed by shipped notebooks and must stay installable. + for pkg in [ + "snac", + "causal-conv1d", + "mamba-ssm", + "omegaconf", + "timm", + "librosa", + "trackio", + "open-spiel", + "protobuf", + "sentencepiece", + ]: + assert pkg not in shim._KEEP, f"{pkg} must still install for the notebooks that need it" + + +# -------------------------------------------------------------------------- +# Behaviour +# -------------------------------------------------------------------------- +def test_the_shipped_trl_cell_installs_nothing(shim): + # `!pip install --no-deps trl==0.22.2` is the last line of 378 notebooks. + assert _run(shim, SHIPPED_TRL_CELL) is None + + +def test_a_mixed_cell_keeps_only_the_unbaked_package(shim): + execd = _run( + shim, + [ + "--no-deps", + "trl==0.22.2", + "peft==0.14.0", + "datasets==3.0.0", + "accelerate==1.0.0", + UNBAKED, + ], + ) + assert execd == ["--no-deps", UNBAKED], execd + + +def test_cuda_matched_wheels_are_not_replaced_by_pypi_builds(shim): + # torchao 0.17.0+cu128 and torchcodec 0.11.0+cu128 are resolved from the + # cu128 index; a PyPI pin swaps in a generic (or cu13) build. + assert _run(shim, ["torchao==0.15.0", "torchcodec==0.5"]) is None + + +def test_transformers_companions_cannot_desynchronise_the_sidecars(shim): + # Each sidecar ships its own matched huggingface_hub/tokenizers/safetensors; + # replacing the base-venv copies desynchronises every sidecar at once. + assert ( + _run(shim, ["huggingface_hub==0.30.0", "tokenizers==0.20.0", "safetensors==0.4.0"]) is None + ) + + +def test_an_unbaked_package_still_installs(shim): + assert _run(shim, [UNBAKED]) == [UNBAKED] + assert _run(shim, [UNBAKED], tool = "uv") == [UNBAKED] + + +def test_protection_survives_a_requirements_file(shim, tmp_path): + req = tmp_path / "requirements.txt" + req.write_text(f"trl==0.22.2\npeft==0.14.0\ndatasets==3.0.0\n{UNBAKED}\n") + execd = _run(shim, ["-r", str(req)]) + assert execd is not None and execd[0] == "-r" + filtered = Path(execd[1]).read_text() + assert UNBAKED in filtered + for dropped in ("trl", "peft", "datasets"): + assert dropped not in filtered, f"{dropped} slipped through the requirements file" + + +def test_protection_survives_a_direct_wheel_url(shim): + url = "https://files.pythonhosted.org/x/trl-0.22.2-py3-none-any.whl" + assert _run(shim, [url, UNBAKED]) == [UNBAKED] + + +def test_protection_survives_an_editable_vcs_install(shim): + assert _run(shim, ["-e", "git+https://github.com/huggingface/trl.git", UNBAKED]) == [UNBAKED] + + +def test_forwarded_installs_pin_the_protected_set_for_the_resolver(shim): + # Argument filtering alone does not stop a dependency of the kept target from + # dragging peft/datasets back down -- which is how peft 0.19.1 became 0.14.0 + # with no notebook ever naming peft. Every forwarded install carries pins. + with pytest.MonkeyPatch.context() as mp: + mp.setattr(shim.sys, "argv", ["pip", "install", UNBAKED]) + with pytest.raises(_Exec) as exc: + shim.main() + argv = exc.value.argv + assert "--constraint" in argv + pins = Path(argv[argv.index("--constraint") + 1]).read_text() + names = {line.split("==")[0].lower().replace("_", "-") for line in pins.splitlines() if line} + # only the installed subset is pinned, but nothing outside the protected set + assert names, "the constraints file must not be empty" + assert all( + n in shim._KEEP or n == "transformers" or n.startswith("nvidia-") for n in names + ), sorted(names) diff --git a/tests/python/test_docker_publish_ref_freeze.py b/tests/python/test_docker_publish_ref_freeze.py new file mode 100644 index 0000000000..78b47b5eae --- /dev/null +++ b/tests/python/test_docker_publish_ref_freeze.py @@ -0,0 +1,227 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +"""The docker publish workflow must never forward an unfrozen ref. + +`prepare` resolves unsloth, unsloth-zoo and notebooks to ONE commit each so the +amd64 leg, the arm64 leg and the Studio build all bake identical source; that is +the whole reason the job exists. Each resolver was + + SHA="$(git ls-remote "$REF" | awk 'NR==1{print $1}')" + [ -n "$SHA" ] || SHA="$REF" + +`git ls-remote` exits 0 whether or not a ref matched, so a non-zero exit means +the remote was never reached. That exit was lost twice over: it is the first +element of a pipeline, and a `run:` step with no explicit `shell:` runs under +`bash -e` WITHOUT pipefail, so the step exited 0 and published `ref=main`. Each +build then resolved `main` independently, and a branch advance between them +would ship one multi-arch tag containing different revisions. The stable-tag +gates key off the inputs, not off whether resolution worked, so `:latest` would +still be moved onto it. + +Static plus behavioural: the resolver `run:` blocks are executed under `bash -e` +with a `git` stub. No docker, no network. +""" + +from __future__ import annotations + +import os +import re +import shutil +import subprocess +from pathlib import Path + +import pytest +import yaml + +REPO_ROOT = Path(__file__).resolve().parents[2] +WORKFLOW = REPO_ROOT / ".github" / "workflows" / "docker-publish.yml" + +RESOLVER_STEPS = ("unsloth_ref", "zoo_ref", "notebooks") + +pytestmark = pytest.mark.skipif( + shutil.which("bash") is None, + reason = "needs bash", +) + + +@pytest.fixture(scope = "module") +def steps() -> dict: + assert WORKFLOW.is_file(), f"missing {WORKFLOW}" + doc = yaml.safe_load(WORKFLOW.read_text(encoding = "utf-8")) + found = {} + for step in doc["jobs"]["prepare"]["steps"]: + if step.get("id") in RESOLVER_STEPS: + found[step["id"]] = step["run"] + missing = set(RESOLVER_STEPS) - set(found) + assert not missing, f"resolver steps missing from the prepare job: {missing}" + return found + + +def test_the_workflow_never_pins_a_shell_so_bash_e_has_no_pipefail(steps: dict): + # If someone later adds `shell: bash` the runner switches to + # `bash --noprofile --norc -eo pipefail`, which would make the guards below + # redundant rather than wrong -- but until then they are the only protection. + doc = yaml.safe_load(WORKFLOW.read_text(encoding = "utf-8")) + assert "shell" not in doc.get("defaults", {}).get("run", {}), ( + "this test models the default `bash -e` shell; update it if a default " + "shell with pipefail is introduced" + ) + + +@pytest.mark.parametrize("step_id", RESOLVER_STEPS) +def test_an_unreachable_remote_fails_the_step(steps: dict, step_id: str, tmp_path: Path): + script = _expand(steps[step_id]) + res = _run_with_failing_ls_remote(script, tmp_path) + assert res.returncode != 0, ( + "a transport failure must fail the prepare job, not fall through to the " + f"mutable ref:\nstdout={res.stdout}\nstderr={res.stderr}" + ) + + +@pytest.mark.parametrize("step_id", RESOLVER_STEPS) +def test_an_unreachable_remote_never_emits_a_mutable_ref(steps: dict, step_id: str, tmp_path: Path): + script = _expand(steps[step_id]) + res = _run_with_failing_ls_remote(script, tmp_path) + emitted = ( + (tmp_path / "github_output").read_text(encoding = "utf-8") + if (tmp_path / "github_output").exists() + else "" + ) + for line in emitted.splitlines(): + key, _, value = line.partition("=") + assert re.fullmatch(r"[0-9a-f]{40}", value), ( + f"{step_id} published {key}={value!r}, which the three builds each " + "resolve again, so they can bake different revisions" + ) + assert res.returncode != 0 + + +# --- the llama.cpp prebuilt tag ---------------------------------------------- +# Same hole, same job, different resolver: the tag step is +# +# TAG="$(curl -fsSL -o /dev/null -w '%{url_effective}' .../releases/latest \ +# | sed -n 's#.*/releases/tag/##p')" +# echo "tag=${TAG:-latest}" >> "$GITHUB_OUTPUT" +# +# `bash -e` without pipefail takes the exit status of `sed`, so an unreachable +# github.com made the step emit `tag=latest`. That value is NOT a pin: both +# matrix legs pass it to docker/fetch_llama_prebuilt.py, whose main() re-resolves +# "latest" per build, and Dockerfile.studio re-resolves it a third time, so a +# release published mid-run can put two different llama.cpp bundles under one +# multi-arch manifest -- with `:latest` moved onto it, because the stable-tag +# gates key off the dispatch inputs, not off whether resolution worked. + + +@pytest.fixture(scope = "module") +def llama_step() -> str: + doc = yaml.safe_load(WORKFLOW.read_text(encoding = "utf-8")) + for step in doc["jobs"]["prepare"]["steps"]: + if step.get("id") == "llama": + return step["run"] + raise AssertionError("the llama tag resolver step is missing from the prepare job") + + +def test_an_unresolvable_llama_release_fails_the_step(llama_step: str, tmp_path: Path): + res = _run_llama_step(llama_step, tmp_path, curl_exit = 6) + assert res.returncode != 0, ( + "a failed /releases/latest lookup must fail the prepare job:\n" + f"stdout={res.stdout}\nstderr={res.stderr}" + ) + + +def test_an_unresolvable_llama_release_never_emits_a_mutable_tag(llama_step: str, tmp_path: Path): + res = _run_llama_step(llama_step, tmp_path, curl_exit = 6) + emitted = (tmp_path / "github_output").read_text(encoding = "utf-8") + assert "latest" not in emitted, ( + f"the step published {emitted.strip()!r}; every consumer resolves that " + "mutable tag again, so the two arch legs and Studio can bake different " + "llama.cpp versions under one manifest" + ) + assert res.returncode != 0 + + +def test_a_resolved_llama_release_is_forwarded_verbatim(llama_step: str, tmp_path: Path): + # The fix must not break the normal path. + res = _run_llama_step(llama_step, tmp_path, curl_exit = 0) + assert res.returncode == 0, f"stdout={res.stdout}\nstderr={res.stderr}" + assert (tmp_path / "github_output").read_text(encoding = "utf-8").strip() == ( + "tag=b10107-mix-1911198" + ) + + +def _run_llama_step(script: str, tmp_path: Path, *, curl_exit: int): + bin_dir = tmp_path / "bin" + bin_dir.mkdir(parents = True, exist_ok = True) + stub = bin_dir / "curl" + if curl_exit: + # How curl reports an unreachable github.com: nothing on stdout, non-zero. + stub.write_text( + "#!/usr/bin/env bash\n" + 'echo "curl: (6) Could not resolve host: github.com" >&2\n' + f"exit {curl_exit}\n", + encoding = "utf-8", + ) + else: + stub.write_text( + "#!/usr/bin/env bash\n" + "printf '%s' " + "'https://github.com/unslothai/llama.cpp/releases/tag/b10107-mix-1911198'\n", + encoding = "utf-8", + ) + stub.chmod(0o755) + out = tmp_path / "github_output" + out.write_text("", encoding = "utf-8") + env = dict(os.environ) + env["PATH"] = f"{bin_dir}{os.pathsep}" + env["PATH"] + env["GITHUB_OUTPUT"] = str(out) + env["INPUT_TAG"] = "" # the default (push / schedule) trigger + path = tmp_path / "llama_step.sh" + path.write_text(_expand(script), encoding = "utf-8") + return subprocess.run( + ["bash", "-e", str(path)], + capture_output = True, + text = True, + env = env, + timeout = 60, + ) + + +def _expand(run: str) -> str: + """Replace the `${{ ... }}` expressions with the empty string the default + (push to main, no dispatch inputs) trigger produces.""" + return re.sub(r"\$\{\{[^}]*\}\}", "", run) + + +def _run_with_failing_ls_remote(script: str, tmp_path: Path): + bin_dir = tmp_path / "bin" + bin_dir.mkdir(parents = True, exist_ok = True) + stub = bin_dir / "git" + stub.write_text( + "#!/usr/bin/env bash\n" + 'if [ "$1" = "ls-remote" ]; then\n' + ' echo "fatal: unable to access: Could not resolve host" >&2\n' + " exit 128\n" + "fi\n" + "exit 0\n", + encoding = "utf-8", + ) + stub.chmod(0o755) + out = tmp_path / "github_output" + out.write_text("", encoding = "utf-8") + env = dict(os.environ) + env["PATH"] = f"{bin_dir}{os.pathsep}" + env["PATH"] + env["GITHUB_OUTPUT"] = str(out) + # Whatever the expansions above blanked out; the resolvers default to "main". + for name in ("INPUT_REF", "TAG_REF", "PUSH_SHA"): + env[name] = "" + path = tmp_path / "step.sh" + path.write_text(script, encoding = "utf-8") + # Exactly how the runner invokes a `run:` step with no explicit `shell:`. + return subprocess.run( + ["bash", "-e", str(path)], + capture_output = True, + text = True, + env = env, + timeout = 60, + ) diff --git a/tests/python/test_docker_tf_sidecar_vllm_floor.py b/tests/python/test_docker_tf_sidecar_vllm_floor.py new file mode 100644 index 0000000000..23b00c9add --- /dev/null +++ b/tests/python/test_docker_tf_sidecar_vllm_floor.py @@ -0,0 +1,238 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +"""Regression guard for transformers-sidecar selection in the Unsloth Docker image. + +The image runs unslothai/notebooks unchanged by refusing a notebook's +`transformers==X` install and activating a baked "sidecar" (transformers X plus +its matched huggingface_hub/tokenizers/safetensors) on sys.path instead. The +selection was a pure CEILING -- smallest baked version >= the request -- which +ignored that vLLM is version-locked to transformers. Two of the four baked +sidecars could not be imported by the baked vLLM 0.26.0 at all, and they were +exactly the two the common pins selected: + + sidecar 4.57.6 ImportError: Support for Transformers v4 is deprecated and + was removed in vLLM v0.24.0 + <- pins 4.48 / 4.52.3 / 4.55.4 / 4.56.1 / 4.56.2 / 4.57.x + = 241 of the 433 shipped notebooks + sidecar 5.3.0 ImportError: cannot import name 'ALLOWED_LAYER_TYPES' from + transformers.configuration_utils + <- pins 5.2.0 / 5.3.0 = 13 more notebooks + +254 of 433 notebooks therefore died at `from unsloth import FastModel`, before +the first model cell. Pointing UNSLOTH_TF_SIDECAR_ROOT at an empty directory, +changing nothing else, turned two of them into clean 22/22 and 25/25 passes. + +The fix is a FLOOR in front of the ceiling. Which versions are above the floor is +not hardcoded: the Dockerfile imports vllm.transformers_utils.config under every +candidate sidecar (the vLLM module that reads the transformers API -- it +reproduces both failures and needs no GPU, which matters because the build host +has none), deletes the ones that raise, and records the lowest survivor. A +request below the floor is clamped UP to the lowest eligible sidecar, which is +the closest thing to the notebook's pin the image can actually run. + +Static: parses the Dockerfile and drives unsloth_nb_compat against a synthetic +sidecar root. No docker, no GPU, no network. +""" + +from __future__ import annotations + +import importlib.util +import re +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +DOCKERFILE = REPO_ROOT / "docker" / "Dockerfile" +COMPAT_PATH = REPO_ROOT / "docker" / "unsloth_nb_compat.py" + +# Every distinct transformers pin across the 433 shipped notebooks, and the +# sidecar each must resolve to once 4.57.6 and 5.3.0 are gone. +SHIPPED_PINS = [ + "4.48", + "4.52.3", + "4.55.4", + "4.56.1", + "4.56.2", + "4.57.0", + "4.57.1", + "4.57.3", + "5.2.0", + "5.3.0", + "5.5.0", + "5.10.1", + "5.11.0", +] + + +@pytest.fixture(scope = "module") +def dockerfile() -> str: + assert DOCKERFILE.is_file(), f"missing {DOCKERFILE}" + return DOCKERFILE.read_text() + + +@pytest.fixture(scope = "module") +def sidecar_block(dockerfile: str) -> str: + start = dockerfile.index("tf-sidecars/t_$(echo") + block = dockerfile[dockerfile.rindex("RUN set -eux", 0, start) :] + return block[: block.index("\n\n")] + + +def _load_compat(root, floor = None): + """Import a fresh unsloth_nb_compat bound to a synthetic sidecar root.""" + import os + + prev_root = os.environ.get("UNSLOTH_TF_SIDECAR_ROOT") + prev_min = os.environ.get("UNSLOTH_TF_SIDECAR_MIN") + os.environ["UNSLOTH_TF_SIDECAR_ROOT"] = str(root) + os.environ.pop("UNSLOTH_TF_SIDECAR_MIN", None) + try: + spec = importlib.util.spec_from_file_location("unsloth_nb_compat_under_test", COMPAT_PATH) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + finally: + if prev_root is None: + os.environ.pop("UNSLOTH_TF_SIDECAR_ROOT", None) + else: + os.environ["UNSLOTH_TF_SIDECAR_ROOT"] = prev_root + if prev_min is not None: + os.environ["UNSLOTH_TF_SIDECAR_MIN"] = prev_min + return mod + + +@pytest.fixture() +def fixed_root(tmp_path): + """The sidecar root the fixed Dockerfile produces: only verified sidecars, + plus the recorded floor.""" + for name in ("t_5_5_0", "t_5_10_2"): + (tmp_path / name).mkdir() + (tmp_path / ".vllm_min_transformers").write_text("5.5.0\n") + return tmp_path + + +@pytest.fixture() +def stale_root(tmp_path): + """A root that still carries the incompatible sidecars (a bind-mounted or + pre-fix directory). The recorded floor must keep them unselectable.""" + for name in ("t_4_57_6", "t_5_3_0", "t_5_5_0", "t_5_10_2"): + (tmp_path / name).mkdir() + (tmp_path / ".vllm_min_transformers").write_text("5.5.0\n") + return tmp_path + + +# -------------------------------------------------------------------------- +# The build must decide eligibility by measurement, not by a literal. +# -------------------------------------------------------------------------- +def test_build_verifies_every_sidecar_against_the_baked_vllm(sidecar_block: str): + assert "import vllm.transformers_utils.config" in sidecar_block, ( + "each baked sidecar must be proven importable by the baked vLLM; this is " + "the module that reads the transformers API and it reproduces both the " + "v4 refusal and the ALLOWED_LAYER_TYPES break" + ) + + +def test_build_verification_needs_no_gpu(sidecar_block: str): + # `import unsloth` raises NotImplementedError("cannot find any torch + # accelerator") on the build host, so it can never be the gate. + assert ( + "import unsloth" not in sidecar_block + ), "the sidecar gate must not import unsloth: the build host has no GPU" + + +def test_an_unverifiable_sidecar_is_deleted_not_shipped(sidecar_block: str): + assert re.search(r"DROPPED", sidecar_block), "a failed candidate must be reported" + assert re.search(r'rm -rf "\$DEST"', sidecar_block), ( + "a sidecar the baked vLLM cannot import must be removed, not shipped: it " + "can never be selected safely and it costs image size" + ) + + +def test_build_records_the_selection_floor(sidecar_block: str): + assert ( + ".vllm_min_transformers" in sidecar_block + ), "the lowest verified version must be recorded for unsloth_nb_compat" + assert "sort -V | head -1" in sidecar_block, "the floor is the LOWEST survivor" + + +def test_build_fails_when_no_sidecar_survives(sidecar_block: str): + assert "exit 1" in sidecar_block, ( + "an empty sidecar set means the whole per-notebook mechanism is dead; " + "that must fail the build rather than ship silently" + ) + + +def test_build_skips_the_gate_when_vllm_is_absent(sidecar_block: str): + # The vLLM install is fail-soft per arch; with no vLLM there is no constraint + # and every sidecar must survive rather than the build exploding. + assert "HAVE_VLLM" in sidecar_block + + +def test_compat_reads_the_floor_the_build_writes(): + assert ".vllm_min_transformers" in COMPAT_PATH.read_text(), ( + "unsloth_nb_compat must read the floor the Dockerfile records, not a " + "literal that rots on the next vLLM bump" + ) + + +# -------------------------------------------------------------------------- +# Selection: floor, then ceiling. +# -------------------------------------------------------------------------- +def test_floor_is_read_back(fixed_root): + assert _load_compat(fixed_root).min_version() == "5.5.0" + + +@pytest.mark.parametrize( + "pin, expected", + [ + # every pin below the floor clamps UP to the lowest eligible sidecar + ("4.48", "t_5_5_0"), + ("4.52.3", "t_5_5_0"), + ("4.55.4", "t_5_5_0"), + ("4.56.1", "t_5_5_0"), + ("4.56.2", "t_5_5_0"), + ("4.57.0", "t_5_5_0"), + ("4.57.1", "t_5_5_0"), + ("4.57.3", "t_5_5_0"), + ("5.2.0", "t_5_5_0"), + ("5.3.0", "t_5_5_0"), + # at and above the floor, the ceiling still decides + ("5.5.0", "t_5_5_0"), + ("5.10.1", "t_5_10_2"), + # newer than every sidecar -> the baked transformers + ("5.11.0", None), + ], +) +def test_every_shipped_pin_resolves_to_a_vllm_compatible_sidecar(fixed_root, pin, expected): + got = _load_compat(fixed_root).sidecar_for(pin) + assert (Path(got).name if got else None) == expected + + +def test_no_shipped_pin_can_reach_an_incompatible_sidecar(stale_root): + compat = _load_compat(stale_root) + for pin in SHIPPED_PINS: + got = compat.sidecar_for(pin) + name = Path(got).name if got else None + assert name not in ( + "t_4_57_6", + "t_5_3_0", + ), f"pin {pin} selected {name}, which the baked vLLM cannot import" + + +def test_model_tier_fallback_is_clamped_too(stale_root): + # tier_for_model maps qwen3-next and friends to 5.3.0; that tier must not + # reach the 5.3.0 sidecar either. + compat = _load_compat(stale_root) + tier = compat.tier_for_model("unsloth/Qwen3-Next-80B-A3B") + assert tier == "5.3.0" + assert Path(compat.sidecar_for(tier)).name == "t_5_5_0" + + +def test_an_unrecorded_floor_keeps_the_old_ceiling_behaviour(tmp_path): + # No .vllm_min_transformers (an environment that never ran the build-time + # verification): selection must not silently start dropping sidecars. + for name in ("t_4_57_6", "t_5_5_0"): + (tmp_path / name).mkdir() + compat = _load_compat(tmp_path) + assert compat.min_version() is None + assert Path(compat.sidecar_for("4.56.2")).name == "t_4_57_6" diff --git a/tests/python/test_docker_update_helpers.py b/tests/python/test_docker_update_helpers.py new file mode 100644 index 0000000000..2bd939f7ec --- /dev/null +++ b/tests/python/test_docker_update_helpers.py @@ -0,0 +1,275 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +"""Behavioural guards for the two in-container update helpers of the Docker image. + +Both are `docker exec` entry points that mutate a running container, so a wrong +answer costs an outage or a mixed-version install: + +* `unsloth-studio-update` swaps the Studio Python packages and then restarts the + service. It verifies the new backend imports first, but only warned -- so a + release that pulls in a dependency `--no-deps` did not install got the healthy + old process killed and replaced by one that cannot start. supervisord retries + `startretries` times, lands in FATAL and never leaves it on its own, so the + container serves nothing until someone exec's in. +* `unsloth-llama-update --check` reported "up to date" when it could not reach + the release feed at all, and its in-place rollback only removed entries whose + names the OLD tree also had, leaving new-release-only shared objects beside + the restored files. ggml dlopen()s every `libggml-*.so` next to the binaries, + so that mix is loaded on the next GGUF run. + +These drive the real scripts with stub `pip` / `supervisorctl` / `python` / +`mv` on PATH. No docker, no GPU, no network. +""" + +from __future__ import annotations + +import os +import shutil +import subprocess +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +STUDIO_UPDATE = REPO_ROOT / "docker" / "unsloth_studio_update.sh" +LLAMA_UPDATE = REPO_ROOT / "docker" / "unsloth_llama_update.sh" + +pytestmark = pytest.mark.skipif( + shutil.which("bash") is None, + reason = "needs bash", +) + + +def _stub(directory: Path, name: str, body: str) -> None: + directory.mkdir(parents = True, exist_ok = True) + path = directory / name + path.write_text("#!/usr/bin/env bash\n" + body, encoding = "utf-8") + path.chmod(0o755) + + +def _run( + script: Path, + args, + env, + cwd = None, +): + return subprocess.run( + ["bash", str(script), *args], + capture_output = True, + text = True, + env = env, + cwd = cwd, + timeout = 120, + ) + + +# --- unsloth-studio-update ---------------------------------------------------- + + +def _studio_env(tmp_path: Path, *, import_ok: bool) -> dict: + home = tmp_path / "studio" + venv_bin = home / "unsloth_studio" / "bin" + venv_bin.mkdir(parents = True) + _stub( + venv_bin, + "python", + 'if [ "$1" = "-c" ]; then\n' + + ( + " exit 0\n" + if import_ok + else ' case "$2" in *studio.backend.main*) exit 1;; esac\n exit 0\n' + ) + + "fi\n" + 'if [ "$1" = "-m" ] && [ "$2" = "pip" ]; then\n' + ' if [ "$3" = "show" ]; then echo "Version: 2026.7.5"; exit 0; fi\n' + ' echo "STUB-PIP $*" >> "$STUB_LOG"; exit 0\n' + "fi\n" + "exit 0\n", + ) + bin_dir = tmp_path / "bin" + _stub( + bin_dir, + "supervisorctl", + 'echo "STUB-SUPERVISORCTL $*" >> "$STUB_LOG"\n' + 'if [ "$1" = "status" ]; then exit 0; fi\nexit 0\n', + ) + env = dict(os.environ) + env["PATH"] = f"{bin_dir}{os.pathsep}" + env["PATH"] + env["UNSLOTH_STUDIO_HOME"] = str(home) + env["STUB_LOG"] = str(tmp_path / "calls.log") + return env + + +def test_studio_update_restarts_when_the_backend_imports(tmp_path: Path): + env = _studio_env(tmp_path, import_ok = True) + res = _run(STUDIO_UPDATE, [], env) + calls = Path(env["STUB_LOG"]).read_text() if Path(env["STUB_LOG"]).exists() else "" + assert res.returncode == 0, res.stderr + assert "STUB-SUPERVISORCTL restart studio" in calls, calls + + +def test_studio_update_does_not_restart_into_a_backend_that_cannot_import(tmp_path: Path): + env = _studio_env(tmp_path, import_ok = False) + res = _run(STUDIO_UPDATE, [], env) + calls = Path(env["STUB_LOG"]).read_text() if Path(env["STUB_LOG"]).exists() else "" + assert "STUB-SUPERVISORCTL restart studio" not in calls, ( + "restarting into code that cannot import kills a process that is serving " + "fine and parks supervisord's studio program in FATAL:\n" + calls + ) + assert res.returncode != 0, "a broken update must not report success" + assert "--with-deps" in res.stderr, "the remedy must still be printed" + + +# --- unsloth-llama-update ----------------------------------------------------- + + +def _llama_env(tmp_path: Path, *, latest: str | None) -> dict: + install = tmp_path / "llama.cpp" + install.mkdir(parents = True) + (install / "UNSLOTH_PREBUILT_INFO.json").write_text( + '{"tag": "b1111-old"}\n', + encoding = "utf-8", + ) + fetcher = tmp_path / "fetch_llama_prebuilt.py" + resolve = ( + " raise RuntimeError('unreachable')\n" if latest is None else f" return {latest!r}\n" + ) + fetcher.write_text( + "def resolve_latest_tag(repo):\n" + resolve, + encoding = "utf-8", + ) + env = dict(os.environ) + env["UNSLOTH_LLAMA_CPP_PATH"] = str(install) + env["UNSLOTH_LLAMA_FETCHER"] = str(fetcher) + return env + + +def _llama_check(tmp_path: Path, latest): + env = _llama_env(tmp_path, latest = latest) + return _run(LLAMA_UPDATE, ["--check"], env) + + +def test_llama_check_reports_an_available_update(tmp_path: Path): + res = _llama_check(tmp_path, "b2222-new") + assert res.returncode == 0, res.stderr + assert "an update is available" in res.stdout + + +def test_llama_check_reports_up_to_date(tmp_path: Path): + res = _llama_check(tmp_path, "b1111-old") + assert res.returncode == 0, res.stderr + assert "up to date" in res.stdout + + +def test_llama_check_does_not_claim_up_to_date_when_it_could_not_look(tmp_path: Path): + res = _llama_check(tmp_path, None) + assert "up to date" not in res.stdout, ( + "--check exists to report update status; saying 'up to date' for a lookup " + "that never happened is the one answer it must never give:\n" + res.stdout + ) + assert res.returncode != 0, "an unperformed check must not exit 0" + assert "UNKNOWN" in res.stdout + res.stderr + + +def _llama_inplace_env(tmp_path: Path, old: list[str], new: list[str]) -> dict: + """An in-place (volume-mounted) install whose activation fails part-way.""" + install = tmp_path / "llama.cpp" + install.mkdir(parents = True) + for name in old: + (install / name).write_text("OLD\n", encoding = "utf-8") + (install / "UNSLOTH_PREBUILT_INFO.json").write_text( + '{"tag": "b1111-old"}\n', + encoding = "utf-8", + ) + fetcher = tmp_path / "fetch_llama_prebuilt.py" + fetcher.write_text( + "import os, sys\n" + "def resolve_latest_tag(repo):\n" + " return 'b2222-new'\n" + "if __name__ == '__main__':\n" + " dest = sys.argv[3]\n" + " os.makedirs(dest, exist_ok = True)\n" + f" for name in {new!r}:\n" + " open(os.path.join(dest, name), 'w').write('NEW\\n')\n" + " open(os.path.join(dest, 'UNSLOTH_PREBUILT_INFO.json'), 'w')" + '.write(\'{"tag": "b2222-new"}\\n\')\n', + encoding = "utf-8", + ) + # Fail the ACTIVATION move (-t ) AFTER it has moved the files, so + # the install dir is populated with the new tree and `find` still reports the + # failure -- the mid-swap abort the rollback exists for. The drain + # (-t ) and the rollback's own per-file moves must keep working, so + # only that one invocation is broken. + bin_dir = tmp_path / "bin" + _stub( + bin_dir, + "mv", + 'if [ "$1" = "-t" ] && [ "$2" = "$FAIL_MV_TARGET" ]; then\n' + " shift 2\n" + ' for _s in "$@"; do /bin/mv "$_s" "$FAIL_MV_TARGET/"; done\n' + " exit 1\n" + "fi\n" + 'exec /bin/mv "$@"\n', + ) + env = dict(os.environ) + env["PATH"] = f"{bin_dir}{os.pathsep}" + env["PATH"] + env["UNSLOTH_LLAMA_CPP_PATH"] = str(install) + env["UNSLOTH_LLAMA_FETCHER"] = str(fetcher) + env["UNSLOTH_LLAMA_UPDATE_IN_PLACE"] = "1" + env["FAIL_MV_TARGET"] = str(install) + return env + + +def test_llama_rollback_leaves_no_new_release_files_behind(tmp_path: Path): + # "libggml-hexagon.so" exists only in the new release, so the rollback loop -- + # which iterates the BACKUP's entries -- cannot see it. ggml dlopen()s every + # libggml-*.so sitting next to the binaries, so a leftover is loaded against + # the restored older libggml-base.so. + old = ["libggml-base.so", "libggml-cpu-icelake.so", "llama-cli"] + new = [ + "libggml-base.so", + "libggml-cpu-icelake.so", + "llama-cli", + "libggml-hexagon.so", + "llama-mtmd-cli", + ] + env = _llama_inplace_env(tmp_path, old, new) + res = _run(LLAMA_UPDATE, [], env) + assert res.returncode != 0, "a failed swap must not report success" + install = tmp_path / "llama.cpp" + present = sorted(p.name for p in install.iterdir()) + leftovers = [n for n in ("libggml-hexagon.so", "llama-mtmd-cli") if n in present] + assert not leftovers, f"new-release-only files survived the rollback: {leftovers} in {present}" + for name in old: + assert ( + install / name + ).read_text() == "OLD\n", f"{name} was not restored from the backup: {present}" + + +def test_llama_rollback_keeps_every_old_file_when_the_drain_is_interrupted(tmp_path: Path): + # The mirror image: abort while the OLD tree is still being moved into the + # backup. The entries left in the install dir are then the only copy of those + # old files, so clearing the directory before restoring would destroy them. + old = ["libggml-base.so", "libggml-cpu-icelake.so", "llama-cli", "llama-quantize"] + env = _llama_inplace_env(tmp_path, old, old) + install = tmp_path / "llama.cpp" + # Fail the DRAIN (-t /.old.) after moving only the first + # source, so half the old tree is still sitting in the install dir when the + # rollback runs. Those entries are then the only copy there is. + _stub( + tmp_path / "bin", + "mv", + 'case "${1:-}:${2:-}" in\n' + " -t:*/.old.*)\n" + ' _t="$2"; shift 2\n' + ' [ $# -gt 0 ] && /bin/mv "$1" "$_t/"\n' + " exit 1;;\n" + "esac\n" + 'exec /bin/mv "$@"\n', + ) + res = _run(LLAMA_UPDATE, [], env) + assert res.returncode != 0 + survivors = sorted(p.name for p in install.rglob("*") if p.is_file()) + for name in old: + assert name in survivors, f"{name} was lost during an interrupted drain: {survivors}" diff --git a/tests/python/test_unsloth_nb_pip_magic.py b/tests/python/test_unsloth_nb_pip_magic.py new file mode 100644 index 0000000000..adbab11f2b --- /dev/null +++ b/tests/python/test_unsloth_nb_pip_magic.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +"""Regression tests for docker/unsloth_nb_pip_magic.py. + +The input transformer rewrites explicit `! -m pip|uv ...` shell lines +to `!pip|uv ...` so they resolve to the PATH shim. IPython input transformers +see the RAW cell text (brace expansion like `{sys.executable}` happens later, +in the system() execution path), so the braced and absolute-interpreter forms +notebooks use to target the running kernel must be rewritten too (item +3567875025); only matching literal `python`/`py` let module-pip bypass the +shim entirely. +""" + +import importlib.util +import pathlib + +_MOD_PATH = pathlib.Path(__file__).resolve().parents[2] / "docker" / "unsloth_nb_pip_magic.py" +_spec = importlib.util.spec_from_file_location("unsloth_nb_pip_magic", _MOD_PATH) +magic = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(magic) + + +def _rewrite(line): + return magic._rewrite_python_dash_m([line])[0] + + +def test_literal_python_rewritten(): + assert _rewrite("!python -m pip install peft\n") == "!pip install peft\n" + + +def test_literal_python_version_rewritten(): + assert _rewrite("!python3.12 -m pip install peft") == "!pip install peft" + + +def test_sys_executable_braces_rewritten(): + assert _rewrite("!{sys.executable} -m pip install peft\n") == "!pip install peft\n" + + +def test_sys_executable_braces_quoted_rewritten(): + assert _rewrite('!"{sys.executable}" -m pip install peft') == "!pip install peft" + + +def test_sys_executable_braces_spaced_rewritten(): + assert _rewrite("!{ sys.executable } -m pip install peft") == "!pip install peft" + + +def test_absolute_interpreter_path_rewritten(): + assert _rewrite("!/opt/unsloth-venv/bin/python -m pip install peft\n") == "!pip install peft\n" + + +def test_absolute_interpreter_versioned_path_rewritten(): + assert _rewrite("!/usr/bin/python3.11 -m uv pip install peft") == "!uv pip install peft" + + +def test_quoted_interpreter_path_rewritten(): + assert _rewrite('!"/opt/unsloth venv/bin/python" -m pip install peft') == "!pip install peft" + + +def test_indent_preserved(): + assert _rewrite(" !{sys.executable} -m pip install peft") == " !pip install peft" + + +def test_python_script_not_rewritten(): + line = "!python train.py --epochs 3" + assert _rewrite(line) == line + + +def test_module_other_than_pip_not_rewritten(): + line = "!python -m venv .venv" + assert _rewrite(line) == line + + +def test_non_shell_line_not_rewritten(): + line = "x = '{sys.executable} -m pip install peft'" + assert _rewrite(line) == line diff --git a/tests/python/test_unsloth_pip_shim.py b/tests/python/test_unsloth_pip_shim.py new file mode 100644 index 0000000000..dd993898f3 --- /dev/null +++ b/tests/python/test_unsloth_pip_shim.py @@ -0,0 +1,809 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +"""Regression tests for docker/unsloth_pip_shim.py. + +The shim sits ahead of the real pip/uv on PATH inside the Unsloth Docker +notebook environment so a notebook `!pip install ...` / `!uv pip install ...` +cell cannot clobber the baked, ABI-matched cu128 torch/vLLM/transformers stack. +These tests drive main() with UNSLOTH_NB_SHIM=1 and capture the command it would +os.execv, so we can assert what actually reaches the real tool. They cover: + + * -e/--editable paired with its target (a protected editable drops the flag + too, so pip is never left a dangling `-e`); + * -P/--upgrade-package values filtered through the protected set (uv cannot be + told to refresh a baked package); + * direct wheel URL / local wheel path basenames parsed for protected + distribution names before URL passthrough. + +No GPU or network is required. +""" + +from __future__ import annotations + +import importlib.util +import os +import sys +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +SHIM_PATH = REPO_ROOT / "docker" / "unsloth_pip_shim.py" + +TORCH_WHEEL_URL = ( + "https://download.pytorch.org/whl/cu128/torch-2.11.0%2Bcu128-cp312-cp312-linux_x86_64.whl" +) + + +class _Exec(Exception): + """Raised by the patched os.execv so main() stops at the exec point and the + intended command is captured instead of replacing the test process.""" + + def __init__(self, path, argv): + self.path = path + self.argv = list(argv) + + +@pytest.fixture() +def shim(tmp_path, monkeypatch): + """Load a fresh copy of the shim with the transformers marker pointed at a + temp file and os.execv patched to capture (not perform) the exec.""" + marker = tmp_path / "requested_transformers" + monkeypatch.setenv("UNSLOTH_NB_TF_MARKER", str(marker)) + monkeypatch.setenv("UNSLOTH_NB_SHIM", "1") + + assert SHIM_PATH.is_file(), f"missing shim: {SHIM_PATH}" + spec = importlib.util.spec_from_file_location("unsloth_pip_shim_under_test", SHIM_PATH) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + + def _fake_execv(path, argv): + raise _Exec(path, argv) + + monkeypatch.setattr(mod.os, "execv", _fake_execv) + mod._marker_path = marker # convenience for assertions + return mod + + +def _run(shim, tool, args): + """Invoke the shim as `tool install ` and return (execd_tail, marker). + + execd_tail is the argument list after the `install` verb that reached the + real tool, or None when the shim no-op'd (nothing left to install). marker is + the recorded transformers version, or None. + """ + if tool == "uv": + argv = ["uv", "pip", "install", *args] + else: + argv = ["pip", "install", *args] + with pytest.MonkeyPatch.context() as mp: + mp.setattr(shim.sys, "argv", argv) + try: + shim.main() + execd = None + except _Exec as exc: + # main() builds [REAL[tool]] + head + keep_args + the protected + # constraints pair; head ends with `install`, so everything after it + # is what we assert on. The trailing `--constraint <...>` pair is + # injected on EVERY install; strip it here so each test asserts on its + # own args (dedicated tests below cover the pair). + i = exc.argv.index("install") + execd = exc.argv[i + 1 :] + if ( + len(execd) >= 2 + and execd[-2] == "--constraint" + and os.path.basename(execd[-1]).startswith("unsloth-nb-protected-") + ): + execd = execd[:-2] + marker = shim._marker_path.read_text() if shim._marker_path.exists() else None + return execd, marker + + +# -------------------------------------------------------------------------- +# Item 3541142907 -- pair -e/--editable with its target. A protected editable +# drops the flag WITH its value (never `pip install -e snac`); an unprotected +# editable is forwarded verbatim. +# -------------------------------------------------------------------------- +UNSLOTH_VCS = "git+https://github.com/unslothai/unsloth.git#egg=unsloth" + +# Sentinel expectation: the whole command line is forwarded verbatim (execd == args). +KEPT = object() + + +@pytest.mark.parametrize( + "args, expected", + [ + pytest.param(["-e", UNSLOTH_VCS, "snac"], ["snac"], id = "sep-protected"), + # nothing left to install -> no-op, no dangling -e + pytest.param(["-e", UNSLOTH_VCS], None, id = "sep-only-protected-noop"), + pytest.param(["-e", "./localpkg"], KEPT, id = "sep-unprotected-kept"), + pytest.param(["--editable=" + UNSLOTH_VCS, "snac"], ["snac"], id = "inline-protected"), + pytest.param(["--editable=./localpkg"], KEPT, id = "inline-unprotected-kept"), + pytest.param(["-e" + UNSLOTH_VCS, "snac"], ["snac"], id = "attached-protected"), + ], +) +def test_editable_forms(shim, args, expected): + execd, _ = _run(shim, "pip", args) + assert execd == (args if expected is KEPT else expected), execd + + +# -------------------------------------------------------------------------- +# Item 3541142906 -- filter uv -P/--upgrade-package values. `uv pip install +# -P torch snac` must not let uv refresh baked torch; a pinned transformers +# upgrade selector still feeds the sidecar marker. +# -------------------------------------------------------------------------- +@pytest.mark.parametrize( + "args, expected, expected_marker", + [ + pytest.param(["-P", "torch", "snac"], ["snac"], None, id = "protected-dropped"), + pytest.param(["--upgrade-package=transformers", "snac"], ["snac"], None, id = "inline"), + pytest.param(["-P", "transformers==4.55.0", "snac"], ["snac"], "4.55.0", id = "tf-pin"), + pytest.param(["-P", "requests", "requests"], KEPT, None, id = "unprotected-kept"), + # -P is not itself a target + pytest.param(["-P", "torch"], None, None, id = "only-protected-noop"), + ], +) +def test_upgrade_package_forms(shim, args, expected, expected_marker): + execd, marker = _run(shim, "uv", args) + assert execd == (args if expected is KEPT else expected), execd + assert marker == expected_marker, marker + + +# -------------------------------------------------------------------------- +# Item 3541142908 -- parse protected wheel basenames before URL passthrough +# (a recognised protected wheel URL/path is dropped -> no-op). +# -------------------------------------------------------------------------- +NUMPY_WHEEL_URL = "https://example.com/wheels/numpy-2.1.0-cp312-cp312-linux_x86_64.whl" + + +@pytest.mark.parametrize( + "args, expected", + [ + pytest.param([TORCH_WHEEL_URL], None, id = "direct-url"), + pytest.param( + ["/tmp/torch-2.11.0+cu128-cp312-cp312-linux_x86_64.whl"], None, id = "local-path" + ), + # unsloth_zoo-*.whl normalises to unsloth-zoo, which is protected. + pytest.param( + ["https://example.com/unsloth_zoo-1.0-py3-none-any.whl"], None, id = "normalised" + ), + pytest.param([NUMPY_WHEEL_URL], KEPT, id = "unprotected-kept"), + ], +) +def test_wheel_url_and_path_forms(shim, args, expected): + execd, _ = _run(shim, "pip", args) + assert execd == (args if expected is KEPT else expected), execd + + +def test_protected_wheel_in_requirements_file_dropped(shim, tmp_path): + req = tmp_path / "reqs.txt" + req.write_text( + TORCH_WHEEL_URL + "\n" + "snac==1.2.0\n", + encoding = "utf-8", + ) + execd, _ = _run(shim, "pip", ["-r", str(req)]) + # The filtered requirements copy still installs snac; torch's wheel line is + # stripped. execd is `-r `. + assert execd is not None and execd[0] == "-r" + filtered = Path(execd[1]).read_text(encoding = "utf-8") + assert "snac==1.2.0" in filtered + assert "torch" not in filtered + + +# -------------------------------------------------------------------------- +# Guardrails: the ordinary happy paths still work unchanged. +# -------------------------------------------------------------------------- +def test_plain_package_passes_through(shim): + execd, _ = _run(shim, "pip", ["omegaconf==2.3.1"]) + assert execd == ["omegaconf==2.3.1"], execd + + +def test_bare_transformers_recorded_and_dropped(shim): + execd, marker = _run(shim, "pip", ["transformers==4.55.0"]) + assert execd is None + assert marker == "4.55.0" + + +def test_index_url_value_flag_kept_verbatim(shim): + execd, _ = _run(shim, "pip", ["--extra-index-url", "https://example.com/simple", "snac"]) + assert execd == ["--extra-index-url", "https://example.com/simple", "snac"], execd + + +# -------------------------------------------------------------------------- +# Item 3541404842 -- filter editable entries INSIDE a requirements file. +# -------------------------------------------------------------------------- +def test_editable_protected_in_requirements_file_dropped(shim, tmp_path): + req = tmp_path / "reqs.txt" + req.write_text( + "-e git+https://github.com/unslothai/unsloth.git#egg=unsloth\nsnac==1.2.0\n", + encoding = "utf-8", + ) + execd, _ = _run(shim, "pip", ["-r", str(req)]) + assert execd is not None and execd[0] == "-r", execd + filtered = Path(execd[1]).read_text(encoding = "utf-8") + assert "snac==1.2.0" in filtered + assert "unsloth" not in filtered # protected editable line stripped + + +def test_editable_attached_protected_in_requirements_file_dropped(shim, tmp_path): + req = tmp_path / "reqs.txt" + req.write_text( + "-egit+https://github.com/unslothai/unsloth.git#egg=unsloth\nsnac==1.2.0\n", + encoding = "utf-8", + ) + execd, _ = _run(shim, "pip", ["-r", str(req)]) + assert execd is not None and execd[0] == "-r", execd + filtered = Path(execd[1]).read_text(encoding = "utf-8") + assert "snac==1.2.0" in filtered + assert "unsloth" not in filtered + + +def test_editable_unprotected_in_requirements_file_kept(shim, tmp_path): + # An unprotected editable survives even when the file is otherwise rewritten + # (torch dropped); only protected editables are stripped. + req = tmp_path / "reqs.txt" + req.write_text( + "-e ./localpkg\ntorch==2.11.0\nsnac==1.2.0\n", + encoding = "utf-8", + ) + execd, _ = _run(shim, "pip", ["-r", str(req)]) + assert execd is not None and execd[0] == "-r", execd + filtered = Path(execd[1]).read_text(encoding = "utf-8") + assert "./localpkg" in filtered + assert "snac==1.2.0" in filtered + assert "torch" not in filtered + + +# -------------------------------------------------------------------------- +# Item 3541404849 -- a nested -c constraint pin is not recorded as a request. +# -------------------------------------------------------------------------- +def test_nested_constraint_transformers_pin_not_recorded(shim, tmp_path): + constraints = tmp_path / "constraints.txt" + constraints.write_text("transformers==4.55.0\n", encoding = "utf-8") + req = tmp_path / "reqs.txt" + req.write_text("-c constraints.txt\nsnac==1.2.0\n", encoding = "utf-8") + execd, marker = _run(shim, "pip", ["-r", str(req)]) + assert execd is not None and execd[0] == "-r", execd + # A constraint pin is not an install request -> no sidecar marker written. + assert marker is None, marker + + +def test_nested_requirement_transformers_pin_recorded(shim, tmp_path): + # Contrast: a nested -r requirement DOES carry install requests, so its + # transformers pin is still recorded for the sidecar. + nested = tmp_path / "nested.txt" + nested.write_text("transformers==4.55.0\n", encoding = "utf-8") + req = tmp_path / "reqs.txt" + req.write_text("-r nested.txt\nsnac==1.2.0\n", encoding = "utf-8") + execd, marker = _run(shim, "pip", ["-r", str(req)]) + assert execd is not None and execd[0] == "-r", execd + assert marker == "4.55.0", marker + + +# -------------------------------------------------------------------------- +# Item 3541404845 -- handle pip's attached short options (-rfile / -cfile / +# etc). The attached `-e` case lives in test_editable_forms above. +# -------------------------------------------------------------------------- +def test_attached_short_requirement_file_filtered(shim, tmp_path): + # `pip install -rreqs.txt` (attached) must filter the file AND count as a + # target -- before the fix it fell through as an opaque option and no-op'd. + req = tmp_path / "reqs.txt" + req.write_text("torch==2.11.0\nsnac==1.2.0\n", encoding = "utf-8") + execd, _ = _run(shim, "pip", ["-r" + str(req)]) + assert execd is not None and execd[0] == "-r", execd + filtered = Path(execd[1]).read_text(encoding = "utf-8") + assert "snac==1.2.0" in filtered + assert "torch" not in filtered + + +def test_attached_short_constraint_file_filtered(shim, tmp_path): + constraints = tmp_path / "constraints.txt" + constraints.write_text("torch==2.11.0\n", encoding = "utf-8") + execd, _ = _run(shim, "pip", ["-c" + str(constraints), "snac"]) + assert execd is not None and execd[0] == "-c", execd + assert "snac" in execd + filtered = Path(execd[1]).read_text(encoding = "utf-8") + assert "torch" not in filtered + + +def test_attached_short_upgrade_package_protected_dropped(shim): + execd, _ = _run(shim, "uv", ["-Ptorch", "snac"]) + assert execd == ["snac"], execd + assert "torch" not in execd and "-P" not in execd + + +# -------------------------------------------------------------------------- +# Item 3541773143 -- a bare wheel filename (no ./ or / prefix) is still a pip +# target from the CWD, so its protected distribution must be parsed too +# (`pip install torch-2.11.0-...whl` must not reinstall torch). +# -------------------------------------------------------------------------- +@pytest.mark.parametrize( + "args, expected", + [ + pytest.param(["torch-2.11.0+cu128-cp312-cp312-linux_x86_64.whl"], None, id = "bare-torch"), + pytest.param(["dist/torch-2.11.0-cp312-cp312-linux_x86_64.whl"], None, id = "subdir-torch"), + pytest.param(["numpy-2.1.0-cp312-cp312-linux_x86_64.whl"], KEPT, id = "unprotected-kept"), + ], +) +def test_bare_wheel_filename_forms(shim, args, expected): + execd, _ = _run(shim, "pip", args) + assert execd == (args if expected is KEPT else expected), execd + + +# -------------------------------------------------------------------------- +# Item 3541773157 -- a protected VCS URL WITHOUT an #egg= fragment (the egg-less +# form this repo recommends) must be dropped via its repo basename. +# -------------------------------------------------------------------------- +def test_vcs_url_without_egg_protected_dropped(shim): + # git+https://github.com/huggingface/transformers.git -> transformers. + execd, _ = _run(shim, "pip", ["git+https://github.com/huggingface/transformers.git", "snac"]) + assert execd == ["snac"], execd + + +def test_vcs_url_without_egg_with_ref_dropped(shim): + execd, _ = _run(shim, "pip", ["git+https://github.com/unslothai/unsloth-zoo.git@main", "snac"]) + assert execd == ["snac"], execd + + +def test_vcs_url_without_egg_unprotected_kept(shim): + url = "git+https://github.com/someone/coolpkg.git" + execd, _ = _run(shim, "pip", [url]) + assert execd == [url], execd + + +# -------------------------------------------------------------------------- +# Item 3541773153 -- refuse remote (URL) requirement / constraint files in shim +# mode; their protected pins cannot be inspected before the real tool installs. +# -------------------------------------------------------------------------- +R_URL = "https://example.com/reqs.txt" + + +@pytest.mark.parametrize( + "args, expected", + [ + # dropped, and no dangling -r left behind + pytest.param(["-r", R_URL], None, id = "sep-r-only-noop"), + pytest.param(["-r", R_URL, "snac"], ["snac"], id = "sep-r-target-kept"), + pytest.param(["--requirement=" + R_URL, "snac"], ["snac"], id = "inline-r"), + pytest.param(["-r" + R_URL, "snac"], ["snac"], id = "attached-r"), + pytest.param(["-c", "https://example.com/constraints.txt", "snac"], ["snac"], id = "sep-c"), + ], +) +def test_remote_requirement_and_constraint_urls_refused(shim, args, expected): + execd, _ = _run(shim, "pip", args) + assert execd == expected, execd + + +def test_nested_remote_include_dropped(shim, tmp_path): + # A local reqs file that pulls a remote include must have that include + # stripped, not passed through for the real pip to fetch unfiltered. + req = tmp_path / "reqs.txt" + req.write_text("-r https://example.com/evil.txt\nsnac==1.2.0\n", encoding = "utf-8") + execd, _ = _run(shim, "pip", ["-r", str(req)]) + assert execd is not None and execd[0] == "-r", execd + filtered = Path(execd[1]).read_text(encoding = "utf-8") + assert "snac==1.2.0" in filtered + assert "example.com" not in filtered and "://" not in filtered + + +# -------------------------------------------------------------------------- +# Item 3541773164 -- resolver-wide reinstall / ignore-installed flags are +# stripped so they cannot rebuild already-satisfied baked deps. +# -------------------------------------------------------------------------- +def test_force_reinstall_flag_stripped(shim): + execd, _ = _run(shim, "pip", ["--force-reinstall", "snac"]) + assert execd == ["snac"], execd + + +def test_ignore_installed_short_flag_stripped(shim): + execd, _ = _run(shim, "pip", ["-I", "snac"]) + assert execd == ["snac"], execd + + +def test_uv_reinstall_flag_stripped(shim): + execd, _ = _run(shim, "uv", ["--reinstall", "snac"]) + assert execd == ["snac"], execd + + +# -------------------------------------------------------------------------- +# Item 3541773168 -- uv's --reinstall-package selector is filtered through _KEEP +# exactly like -P/--upgrade-package (both forms, no dangling flag). +# -------------------------------------------------------------------------- +@pytest.mark.parametrize( + "args, expected, expected_marker", + [ + pytest.param(["--reinstall-package", "torch", "snac"], ["snac"], None, id = "sep-protected"), + pytest.param(["--reinstall-package=torch", "snac"], ["snac"], None, id = "inline-protected"), + pytest.param(["--reinstall-package", "requests", "requests"], KEPT, None, id = "unprotected"), + pytest.param( + ["--reinstall-package", "transformers==4.55.0", "snac"], ["snac"], "4.55.0", id = "tf-pin" + ), + ], +) +def test_reinstall_package_forms(shim, args, expected, expected_marker): + execd, marker = _run(shim, "uv", args) + assert execd == (args if expected is KEPT else expected), execd + assert marker == expected_marker, marker + + +# -------------------------------------------------------------------------- +# Item 3542096750 -- parse protected source archives (sdist / zip) too. +# -------------------------------------------------------------------------- +SDIST_URL = "https://files.pythonhosted.org/packages/aa/unsloth-2026.7.1.tar.gz" + + +@pytest.mark.parametrize( + "args, expected", + [ + pytest.param([SDIST_URL, "snac"], ["snac"], id = "url-protected"), + pytest.param(["torch-2.11.0.tar.gz"], None, id = "bare-protected"), + pytest.param(["./transformers-4.55.0.zip", "snac"], ["snac"], id = "zip-protected"), + # flashinfer-python is protected; the name must survive the hyphen split. + pytest.param(["flashinfer-python-0.5.0.tar.gz"], None, id = "hyphenated-name"), + pytest.param(["numpy-2.1.0.tar.gz"], KEPT, id = "unprotected-kept"), + ], +) +def test_source_archive_forms(shim, args, expected): + execd, _ = _run(shim, "pip", args) + assert execd == (args if expected is KEPT else expected), execd + + +# -------------------------------------------------------------------------- +# Item 3542096760 -- uv's PLURAL --requirements / --constraints go through the +# same filter as the pip-style singular names. +# -------------------------------------------------------------------------- +def test_uv_plural_requirements_filtered(shim, tmp_path): + req = tmp_path / "reqs.txt" + req.write_text("torch==2.11.0\nsnac==1.2.0\n", encoding = "utf-8") + execd, _ = _run(shim, "uv", ["--requirements", str(req)]) + assert execd is not None and execd[0] == "--requirements", execd + filtered = Path(execd[1]).read_text(encoding = "utf-8") + assert "snac==1.2.0" in filtered + assert "torch" not in filtered + + +def test_uv_plural_constraints_filtered(shim, tmp_path): + constraints = tmp_path / "constraints.txt" + constraints.write_text("torch==2.11.0\n", encoding = "utf-8") + execd, _ = _run(shim, "uv", ["--constraints", str(constraints), "snac"]) + assert execd is not None and execd[0] == "--constraints", execd + assert "snac" in execd + filtered = Path(execd[1]).read_text(encoding = "utf-8") + assert "torch" not in filtered + + +# -------------------------------------------------------------------------- +# Item 3542096764 -- neutralise --upgrade-strategy eager so a kept target cannot +# eagerly rebuild already-satisfied baked deps. +# -------------------------------------------------------------------------- +@pytest.mark.parametrize( + "args, expected", + [ + pytest.param(["-U", "--upgrade-strategy", "eager", "snac"], ["-U", "snac"], id = "eager"), + pytest.param(["--upgrade-strategy=eager", "snac"], ["snac"], id = "inline-eager"), + # only-if-needed is pip's default, so dropping it is a harmless no-op that + # keeps the kept target installing normally. + pytest.param( + ["--upgrade-strategy", "only-if-needed", "snac"], ["snac"], id = "only-if-needed" + ), + ], +) +def test_upgrade_strategy_forms(shim, args, expected): + execd, _ = _run(shim, "pip", args) + assert execd == expected, execd + + +# -------------------------------------------------------------------------- +# Resolver-level protection: every forwarded install carries a constraints file +# pinning the installed protected packages, so a kept target's dependency on an +# incompatible torch/transformers fails loudly instead of replacing the wheel. +# -------------------------------------------------------------------------- +def _raw_execd(shim, tool, args): + """Like _run but WITHOUT stripping the injected constraint pair.""" + argv = ["uv", "pip", "install", *args] if tool == "uv" else ["pip", "install", *args] + with pytest.MonkeyPatch.context() as mp: + mp.setattr(shim.sys, "argv", argv) + try: + shim.main() + return None + except _Exec as exc: + return exc.argv[exc.argv.index("install") + 1 :] + + +class _FakeDist: + """Minimal stand-in for an importlib.metadata Distribution.""" + + def __init__(self, name, version): + self.metadata = {"Name": name} + self.version = version + + +def _fake_distributions(monkeypatch, *pairs): + """Pin what _protected_constraints_file sees as INSTALLED. + + It reads the ambient environment via importlib.metadata.distributions, so + without this the outcome depends on whatever happens to be in the venv: + with no protected package installed it correctly returns None (see its + docstring) and no --constraint pair is appended. That made the assertion + below environment-dependent, and it surfaced as an IndexError on execd[-2] + rather than a readable failure. The shim imports the symbol inside the + function, so patch it at its source. + """ + monkeypatch.setattr( + "importlib.metadata.distributions", + lambda: [_FakeDist(n, v) for n, v in pairs], + ) + + +def test_forwarded_install_carries_protected_constraints(shim, monkeypatch): + _fake_distributions(monkeypatch, ("transformers", "5.14.1"), ("trl", "0.24.0")) + execd = _raw_execd(shim, "pip", ["snac"]) + assert execd is not None, "an unprotected target must still be forwarded" + assert len(execd) >= 2 and execd[-2] == "--constraint", execd + pins = Path(execd[-1]).read_text(encoding = "utf-8").strip().splitlines() + assert pins, "constraints file must pin the installed protected packages" + assert all("==" in pin for pin in pins), pins + names = {pin.split("==", 1)[0].lower().replace("_", "-") for pin in pins} + protected = {"transformers"} | shim._KEEP | {"nvidia-"} + assert all( + n in shim._KEEP or n == "transformers" or n.startswith("nvidia-") for n in names + ), names + + +def test_forwarded_install_without_protected_packages_has_no_constraints(shim, monkeypatch): + # The other half of the contract: with nothing protected installed there is + # nothing to pin, so the install must still be forwarded, just bare. This is + # the case a bare venv actually hits. + _fake_distributions(monkeypatch, ("snac", "1.2.1")) + execd = _raw_execd(shim, "pip", ["snac"]) + assert execd is not None, "the install must still be forwarded" + assert "--constraint" not in execd, execd + + +def test_noop_install_gets_no_constraints(shim): + # A cell whose only target is protected still no-ops (no exec at all). + execd = _raw_execd(shim, "pip", ["torch"]) + assert execd is None + + +# -------------------------------------------------------------------------- +# pip expands ${UPPERCASE} in requirements files AFTER the shim classifies the +# literal text; classification must expand the same way or `${PKG}==...` with +# PKG=torch walks straight past _KEEP. +# -------------------------------------------------------------------------- +def test_env_expanded_protected_requirement_dropped(shim, tmp_path, monkeypatch): + monkeypatch.setenv("PKG", "torch") + req = tmp_path / "reqs.txt" + req.write_text("${PKG}==2.11.0\nsnac==1.2.0\n", encoding = "utf-8") + execd, _ = _run(shim, "pip", ["-r", str(req)]) + assert execd is not None and execd[0] == "-r", execd + filtered = Path(execd[1]).read_text(encoding = "utf-8") + assert "snac==1.2.0" in filtered + assert "${PKG}" not in filtered and "torch" not in filtered + + +def test_env_expanded_transformers_pin_recorded(shim, tmp_path, monkeypatch): + monkeypatch.setenv("TF_PKG", "transformers") + req = tmp_path / "reqs.txt" + req.write_text("${TF_PKG}==4.56.2\nsnac==1.2.0\n", encoding = "utf-8") + _, marker = _run(shim, "pip", ["-r", str(req)]) + assert marker == "4.56.2" + + +def test_unset_env_reference_left_verbatim(shim, tmp_path, monkeypatch): + monkeypatch.delenv("NOT_SET_ANYWHERE", raising = False) + req = tmp_path / "reqs.txt" + req.write_text("${NOT_SET_ANYWHERE}==1.0\nsnac==1.2.0\n", encoding = "utf-8") + execd, _ = _run(shim, "pip", ["-r", str(req)]) + # Nothing protected detected -> the original file is forwarded unchanged + # (pip forwards unset references verbatim too). + assert execd == ["-r", str(req)], execd + + +# -------------------------------------------------------------------------- +# Filtered-copy write failures fail CLOSED: the original file pins protected +# packages, so forwarding it would hand pip exactly what must be filtered. +# -------------------------------------------------------------------------- +def test_filter_write_failure_refuses_original_file(shim, tmp_path, monkeypatch): + req = tmp_path / "reqs.txt" + req.write_text("torch==2.11.0\nsnac==1.2.0\n", encoding = "utf-8") + + def denied(*args, **kwargs): + raise OSError(30, "Read-only file system") + + monkeypatch.setattr(shim.tempfile, "mkstemp", denied) + with pytest.raises(SystemExit, match = "refusing to forward"): + shim._filter_requirements_file(str(req)) + + +def test_filter_write_failure_clean_file_passes_through(shim, tmp_path, monkeypatch): + # A file with nothing protected never needs the temp copy, so a broken + # TMPDIR must not block it. + req = tmp_path / "reqs.txt" + req.write_text("snac==1.2.0\n", encoding = "utf-8") + + def denied(*args, **kwargs): + raise OSError(30, "Read-only file system") + + monkeypatch.setattr(shim.tempfile, "mkstemp", denied) + path, recorded, dropped = shim._filter_requirements_file(str(req)) + assert path == str(req) and recorded is None and dropped == [] + + +# -------------------------------------------------------------------------- +# Item 3567875029 -- uv's --exact performs an exact SYNC (removes packages +# outside the kept target's closure), so it is stripped like the other +# resolver-wide destructive switches. +# -------------------------------------------------------------------------- +def test_uv_exact_flag_stripped(shim): + execd, _ = _run(shim, "uv", ["--exact", "snac"]) + assert execd == ["snac"], execd + + +# -------------------------------------------------------------------------- +# Item 3567875023 -- a local project directory naming a protected package +# (pip install ./transformers, pip install -e ./unsloth) is filtered like the +# wheel/sdist/VCS forms: a same-version dev build slips past the constraints +# file, so the name must come from the project metadata. +# -------------------------------------------------------------------------- +def _make_local_project(tmp_path, dirname, project_name): + proj = tmp_path / dirname + proj.mkdir() + (proj / "pyproject.toml").write_text(f'[project]\nname = "{project_name}"\nversion = "1.0"\n') + return str(proj) + + +def test_local_dir_protected_by_metadata_dropped(shim, tmp_path): + # Directory name is innocuous; pyproject names a protected package. + path = _make_local_project(tmp_path, "my-checkout", "transformers") + execd, _ = _run(shim, "pip", [path, "snac"]) + assert execd == ["snac"], execd + + +def test_local_dir_protected_editable_dropped(shim, tmp_path): + path = _make_local_project(tmp_path, "unsloth", "unsloth") + execd, _ = _run(shim, "pip", ["-e", path, "snac"]) + assert execd == ["snac"], execd + assert "-e" not in execd + + +def test_local_dir_basename_fallback_setup_py(shim, tmp_path): + # No parseable name in metadata: setup.py + protected basename still drops. + proj = tmp_path / "torch" + proj.mkdir() + (proj / "setup.py").write_text("from setuptools import setup\nsetup()\n") + execd, _ = _run(shim, "pip", [str(proj), "snac"]) + assert execd == ["snac"], execd + + +def test_local_dir_unprotected_kept(shim, tmp_path): + path = _make_local_project(tmp_path, "my-torch-utils", "my-torch-utils") + execd, _ = _run(shim, "pip", [path]) + assert execd == [path], execd + + +def test_local_dir_without_metadata_passes_through(shim, tmp_path): + plain = tmp_path / "datadir" + plain.mkdir() + execd, _ = _run(shim, "pip", [str(plain)]) + assert execd == [str(plain)], execd + + +# -------------------------------------------------------------------------- +# Item 3592835033 -- every uv/pip value-taking flag must be in _VALUE_FLAGS. +# `--torch-backend cu128 torch` used to drop torch but keep the separated flag +# pair, exec'ing uv with no target; `--extra torch snac` misread the extra NAME +# "torch" as a target, leaving a dangling `--extra` that swallowed snac. + + +@pytest.mark.parametrize( + "tool, flag, value", + [ + pytest.param("uv", "--torch-backend", "cu128", id = "uv-torch-backend"), + pytest.param("uv", "--resolution", "lowest", id = "uv-resolution"), + pytest.param("uv", "--default-index", "https://mirror/simple", id = "uv-default-index"), + pytest.param("uv", "--exclude-newer", "2026-01-01", id = "uv-exclude-newer"), + pytest.param("uv", "-b", "build-constraints.txt", id = "uv-build-constraints-short"), + pytest.param("pip", "--proxy", "http://proxy:3128", id = "pip-proxy"), + pytest.param("pip", "--retries", "3", id = "pip-retries"), + pytest.param("pip", "--trusted-host", "mirror.internal", id = "pip-trusted-host"), + ], +) +def test_value_flag_protected_only_noops(shim, tool, flag, value): + # The value must not be mistaken for an install target: with only a + # protected target the cell is a clean no-op, never a broken exec. + execd, _ = _run(shim, tool, [flag, value, "torch"]) + assert execd is None, execd + + +@pytest.mark.parametrize( + "tool, flag, value", + [ + pytest.param("uv", "--torch-backend", "cu128", id = "uv-torch-backend"), + pytest.param("uv", "--resolution", "lowest", id = "uv-resolution"), + pytest.param("pip", "--proxy", "http://proxy:3128", id = "pip-proxy"), + ], +) +def test_value_flag_pair_forwarded_with_kept_target(shim, tool, flag, value): + execd, _ = _run(shim, tool, [flag, value, "torch", "snac"]) + assert execd == [flag, value, "snac"], execd + + +def test_extra_value_is_not_a_protected_target(shim): + # `--extra torch` names an EXTRA, not the torch package: the pair stays and + # snac is not swallowed by a dangling --extra. + execd, _ = _run(shim, "uv", ["--extra", "torch", "snac"]) + assert execd == ["--extra", "torch", "snac"], execd + + +def _value_flags_from_help(cmd): + import re + import subprocess + + out = subprocess.run(cmd, capture_output = True, text = True).stdout + flags = set() + for m in re.finditer(r"^\s+(-\w)?,?\s*(--[\w-]+)[= ]<", out, re.M): + if m.group(1): + flags.add(m.group(1)) + flags.add(m.group(2)) + for m in re.finditer(r"^\s+(-\w) <", out, re.M): + flags.add(m.group(1)) + return flags + + +# The help-derived drift guards are OPT-IN: repo CI runs whatever pip/uv are +# current, so a hard assert would turn every upstream flag addition into a red +# PR. The authoritative check runs at image BUILD time against the baked tools +# (--unsloth-selfcheck-value-flags); set UNSLOTH_SHIM_FLAG_DRIFT_CHECK=1 locally. +_DRIFT_OPT_IN = os.environ.get("UNSLOTH_SHIM_FLAG_DRIFT_CHECK") == "1" + + +@pytest.mark.skipif(not _DRIFT_OPT_IN, reason = "opt-in: UNSLOTH_SHIM_FLAG_DRIFT_CHECK=1") +def test_pip_help_value_flags_all_classified(shim): + # Drift guard: every value-taking flag `pip install --help` documents must + # be classified as value-taking by the shim, or its VALUE is misread as an + # install target (see --torch-backend above). + known = shim._VALUE_FLAGS | shim._DROP_VALUE_FLAGS + missing = _value_flags_from_help([sys.executable, "-m", "pip", "install", "--help"]) - known + assert not missing, f"value flags missing from _VALUE_FLAGS: {sorted(missing)}" + + +@pytest.mark.skipif( + not _DRIFT_OPT_IN or not __import__("shutil").which("uv"), + reason = "opt-in: UNSLOTH_SHIM_FLAG_DRIFT_CHECK=1 (and uv installed)", +) +def test_uv_help_value_flags_all_classified(shim): + known = shim._VALUE_FLAGS | shim._DROP_VALUE_FLAGS + missing = _value_flags_from_help(["uv", "pip", "install", "--help"]) - known + assert not missing, f"value flags missing from _VALUE_FLAGS: {sorted(missing)}" + + +# -------------------------------------------------------------------------- +# Item 3592947879 -- a VCS @ref may contain a slash (@feature/foo); strip it +# before the last-segment split, else the ref's basename dodges _KEEP. + + +@pytest.mark.parametrize( + "url", + [ + pytest.param( + "git+https://github.com/unslothai/unsloth.git@feature/foo", id = "https-slash-ref" + ), + pytest.param( + "git+ssh://git@github.com/unslothai/unsloth.git@feature/foo", + id = "ssh-userinfo-and-slash-ref", + ), + pytest.param("git+https://github.com/unslothai/unsloth.git@v2026.7", id = "plain-tag-ref"), + pytest.param("git+https://github.com/unslothai/unsloth.git", id = "no-ref"), + ], +) +def test_vcs_slash_ref_still_protected(shim, url): + execd, _ = _run(shim, "pip", [url, "snac"]) + assert execd == ["snac"], execd + + +def test_vcs_slash_ref_unprotected_kept(shim): + url = "git+https://github.com/someorg/sometool.git@feature/foo" + execd, _ = _run(shim, "pip", [url]) + assert execd == [url], execd diff --git a/tests/sh/test_select_cuda_jit_tools.sh b/tests/sh/test_select_cuda_jit_tools.sh new file mode 100755 index 0000000000..acd7716103 --- /dev/null +++ b/tests/sh/test_select_cuda_jit_tools.sh @@ -0,0 +1,122 @@ +#!/bin/bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +# Unit tests for select_cuda_jit_tools() from docker/entrypoint.sh. +# +# cu12.8 is the immutable baked default; the cu13 tools switch on ONLY for sm_103 +# and sm_121 (>= 580 drivers). The function picks per device via nvidia-smi +# compute_cap: those two arches retarget libnvrtc.so.12 -> the .cu13 alias (and +# point Triton at cu13 ptxas); every other arch keeps cu12.8 and leaves ptxas unset. +set -e + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +ENTRYPOINT_SH="$SCRIPT_DIR/../../docker/entrypoint.sh" +PASS=0 +FAIL=0 + +# The fixtures below stage libnvrtc as symlinks and assert through readlink, +# because retargeting that symlink is exactly what the function under test does. +# git-bash copies instead of symlinking unless MSYS=winsymlinks:nativestrict and +# the user is elevated, so readlink comes back empty and all 14 assertions fail +# for reasons that have nothing to do with the code. That code only ever runs +# inside a Linux container, so skip rather than pretend: an unconditional run +# breaks tests/run_all.sh for every Windows contributor. +_probe=$(mktemp -d) +: > "$_probe/target" +if ! ln -s target "$_probe/link" 2>/dev/null || [ "$(readlink "$_probe/link")" != "target" ]; then + rm -rf "$_probe" + echo "=== test_select_cuda_jit_tools ===" + echo " SKIP: this filesystem does not honour symlinks (readlink cannot observe them)" + echo "PASS=0 FAIL=0 SKIPPED" + exit 0 +fi +rm -rf "$_probe" + +# Extract just the helper function (same sed range as the other function tests). +_FUNC_FILE=$(mktemp) +sed -n '/^select_cuda_jit_tools()/,/^}/p' "$ENTRYPOINT_SH" > "$_FUNC_FILE" + +assert_eq() { + _label="$1"; _expected="$2"; _actual="$3" + if [ "$_actual" = "$_expected" ]; then + echo " PASS: $_label" + PASS=$((PASS + 1)) + else + echo " FAIL: $_label (expected '$_expected', got '$_actual')" + FAIL=$((FAIL + 1)) + fi +} + +# $1 = compute_cap(s) the mock nvidia-smi reports, ONE PER LINE ("none" -> no +# nvidia-smi; multi-line models a mixed-GPU host). $2 (optional) = the initial +# libnvrtc.so.12 target (default cu12.8; "libnvrtc.so.12.cu13" models a stale +# link). Builds a fake Studio venv NVRTC dir. Prints " ". +run_select() { + _cap="$1" + _init="${2:-libnvrtc.so.12.cu128.orig}" + _tmp=$(mktemp -d) + mkdir -p "$_tmp/bin" + if [ "$_cap" != "none" ]; then + # nvidia-smi --query-gpu=compute_cap prints one cap per line; cat a file + # so an embedded newline in $_cap survives into the mock's output. + printf '%s\n' "$_cap" > "$_tmp/caps.txt" + printf '#!/bin/sh\ncat "%s"\n' "$_tmp/caps.txt" > "$_tmp/bin/nvidia-smi" + chmod +x "$_tmp/bin/nvidia-smi" + fi + _nvrtc="$_tmp/studio/unsloth_studio/lib/python3.12/site-packages/nvidia/cuda_nvrtc/lib" + mkdir -p "$_nvrtc" + : > "$_nvrtc/libnvrtc.so.12.cu128.orig" # real cu12.8 lib + : > "$_nvrtc/libnvrtc.so.13.stub" # stand-in cu13 lib + ln -sf libnvrtc.so.13.stub "$_nvrtc/libnvrtc.so.12.cu13" # staged cu13 alias + ln -sf "$_init" "$_nvrtc/libnvrtc.so.12" # cu12.8 default (or stale cu13) + bash -c ' + set -euo pipefail + export PATH="'"$_tmp"'/bin:/usr/bin:/bin" + export UNSLOTH_STUDIO_HOME="'"$_tmp"'/studio" + unset TRITON_PTXAS_PATH || true + . "'"$_FUNC_FILE"'" + select_cuda_jit_tools || true + printf "%s %s\n" "${TRITON_PTXAS_PATH:-UNSET}" "$(readlink "'"$_nvrtc"'/libnvrtc.so.12")" + ' + rm -rf "$_tmp" +} + +echo "=== test_select_cuda_jit_tools ===" + +# Non-DC arches: cu12.8 default is left untouched (no write) and ptxas unset +# (Triton keeps its bundled cu12.8 ptxas), so a 570-579 driver host -- root or +# --user -- is unaffected. +assert_eq "sm_80 Ampere -> cu128 default kept" "UNSET libnvrtc.so.12.cu128.orig" "$(run_select 8.0)" +assert_eq "sm_90 Hopper -> cu128 default kept" "UNSET libnvrtc.so.12.cu128.orig" "$(run_select 9.0)" +assert_eq "sm_100 B200 -> cu128 default kept" "UNSET libnvrtc.so.12.cu128.orig" "$(run_select 10.0)" +assert_eq "sm_120 RTX50 -> cu128 default kept" "UNSET libnvrtc.so.12.cu128.orig" "$(run_select 12.0)" +assert_eq "no nvidia-smi -> cu128 default kept" "UNSET libnvrtc.so.12.cu128.orig" "$(run_select none)" + +# Blackwell datacenter: retarget libnvrtc.so.12 -> the .cu13 alias. ptxas stays +# UNSET here only because the test host has no /usr/local/cuda-13.0/bin/ptxas; +# the assertion that matters is that the NVRTC switched to cu13 for these arches. +assert_eq "sm_103 B300 -> cu13 NVRTC selected" "UNSET libnvrtc.so.12.cu13" "$(run_select 10.3)" +assert_eq "sm_121 DGX Spark -> cu13 NVRTC selected" "UNSET libnvrtc.so.12.cu13" "$(run_select 12.1)" + +# Mixed-GPU hosts: a datacenter Blackwell (sm_103 / sm_121) sitting BEHIND an +# H100/B200 in the nvidia-smi ordering must still switch to cu13 -- every visible +# cap is scanned, not just the first. A host with no datacenter Blackwell at all +# keeps the cu12.8 default regardless of order. +assert_eq "H100 then B300 -> cu13 NVRTC selected" "UNSET libnvrtc.so.12.cu13" "$(run_select "$(printf '9.0\n10.3')")" +assert_eq "B200 then GB10 -> cu13 NVRTC selected" "UNSET libnvrtc.so.12.cu13" "$(run_select "$(printf '10.0\n12.1')")" +assert_eq "B300 then H100 -> cu13 NVRTC selected" "UNSET libnvrtc.so.12.cu13" "$(run_select "$(printf '10.3\n9.0')")" +assert_eq "H100 then A100 -> cu128 default kept" "UNSET libnvrtc.so.12.cu128.orig" "$(run_select "$(printf '9.0\n8.0')")" + +# Stateful transition: a cu13 selection left in the same container's writable +# layer by an earlier sm_103/sm_121 boot must be reversed when the container +# later starts on an ordinary GPU (or none) -- a 570-579 driver cannot load +# cu13-produced cubins -- and kept when the datacenter Blackwell is still there. +assert_eq "A100 after B300 -> cu128 restored" "UNSET libnvrtc.so.12.cu128.orig" "$(run_select 8.0 libnvrtc.so.12.cu13)" +assert_eq "no GPU after B300 -> cu128 restored" "UNSET libnvrtc.so.12.cu128.orig" "$(run_select none libnvrtc.so.12.cu13)" +assert_eq "B300 after B300 -> cu13 kept" "UNSET libnvrtc.so.12.cu13" "$(run_select 10.3 libnvrtc.so.12.cu13)" + +rm -f "$_FUNC_FILE" + +echo "" +echo "Results: $PASS passed, $FAIL failed" +[ "$FAIL" -eq 0 ] || exit 1 diff --git a/tests/studio/test_branding_guard.py b/tests/studio/test_branding_guard.py new file mode 100644 index 0000000000..c968cc4797 --- /dev/null +++ b/tests/studio/test_branding_guard.py @@ -0,0 +1,252 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 +"""Tests for the Unsloth Docker Studio branding / AGPLv3 integrity guard. + +verify_branding() is exercised against a staged temp tree that mirrors the +installed image layout, so no container or built labextension is required: + * positive: a faithful tree passes (no problems). + * negative: removing/altering each attribution marker is detected. + * no-encoding: the attribution sources carry no base64/decoder obfuscation + (plain readable strings only -- the only data URI is the logo *image*). +""" + +import json +import os +import sys + +import pytest + +HERE = os.path.dirname(os.path.abspath(__file__)) +REPO = os.path.abspath(os.path.join(HERE, "..", "..")) +sys.path.insert(0, os.path.join(REPO, "docker", "jupyter")) + +import unsloth_branding as ub # noqa: E402 + + +def _stage(tmp_path): + """Create a faithful copy of the installed branding layout; return paths.""" + venv_share = tmp_path / "venv-share" + js_dir = tmp_path / "jupyter_server" + + (venv_share).mkdir(parents = True) + (venv_share / "UNSLOTH_LICENSE.AGPL-3.0").write_text( + " GNU AFFERO GENERAL PUBLIC LICENSE\n" + " Version 3, 19 November 2007\n" + " Copyright (C) 2007 Free Software Foundation, Inc.\n", + encoding = "utf-8", + ) + + (venv_share / "lab" / "settings").mkdir(parents = True) + (venv_share / "lab" / "settings" / "overrides.json").write_text( + json.dumps({"@jupyterlab/apputils-extension:themes": {"theme": ub.THEME_NAME}}), + encoding = "utf-8", + ) + + labext = venv_share / "labextensions" / ub.LABEXT_NAME + (labext / "static").mkdir(parents = True) + (labext / "package.json").write_text(json.dumps({"name": ub.LABEXT_NAME}), encoding = "utf-8") + bundle = " ".join( + [ + ub.PHRASE, + ub.SHORT_LABEL, + ub.COPYRIGHT, + ub.AGPL_URL, + ub.ABOUT_PLUGIN_ID, + ub.SPLASH_PLUGIN_ID, + ub.LOGO_DATA_URI_PREFIX + "AAAAdummyimagebytes", + ] + ) + (labext / "static" / "remoteEntry.abc123.js").write_text(bundle, encoding = "utf-8") + + (js_dir / "templates").mkdir(parents = True) + (js_dir / "templates" / "login.html").write_text( + "Built by the Unsloth team. Apache 2.0, AGPLv3 License Link\n" + "Copyright 2026-Present the Unsloth team.\n" + "https://github.com/unslothai/unsloth#license\n" + "https://github.com/unslothai/unsloth\n", + encoding = "utf-8", + ) + + (js_dir / "static" / "favicons").mkdir(parents = True) + (js_dir / "static" / "favicons" / "favicon.ico").write_bytes(b"\x00\x00\x01\x00icon") + (js_dir / "static" / "logo").mkdir(parents = True) + (js_dir / "static" / "logo" / "logo.png").write_bytes(b"\x89PNG\r\n\x1a\nlogo") + + # config_dirs = [] keeps the tree hermetic (no host jupyter config scanned); + # page_config tests write to the app-settings page_config.json directly. + return ub.resolve_paths( + venv_share = str(venv_share), + jupyter_server_dir = str(js_dir), + config_dirs = [], + ) + + +def test_positive_clean_tree_passes(tmp_path): + paths = _stage(tmp_path) + assert ub.verify_branding(paths) == [] + + +# --- negative mutations: each strips one attribution marker -------------------- +def _remove_license(paths): + os.remove(paths["license"]) + + +def _blank_license(paths): + with open(paths["license"], "w", encoding = "utf-8") as f: + f.write("All rights reserved. Proprietary. Resold by someone else.\n") + + +def _remove_login(paths): + os.remove(paths["login"]) + + +def _strip_login_source(paths): + with open(paths["login"], encoding = "utf-8") as f: + text = f.read() + with open(paths["login"], "w", encoding = "utf-8") as f: + f.write(text.replace(ub.SOURCE_URL, "https://example.com/forks")) + + +def _strip_login_copyright(paths): + with open(paths["login"], encoding = "utf-8") as f: + text = f.read() + with open(paths["login"], "w", encoding = "utf-8") as f: + f.write(text.replace(ub.COPYRIGHT, "Copyright someone else")) + + +def _drop_theme(paths): + with open(paths["overrides"], "w", encoding = "utf-8") as f: + f.write("{}") + + +def _rebrand_labext(paths): + with open(paths["labext_pkg"], "w", encoding = "utf-8") as f: + f.write(json.dumps({"name": "totally-not-unsloth"})) + + +def _strip_bundle_phrase(paths): + import glob + for path in glob.glob(os.path.join(paths["labext_static"], "*.js")): + with open(path, encoding = "utf-8") as f: + text = f.read() + with open(path, "w", encoding = "utf-8") as f: + f.write(text.replace(ub.PHRASE, "").replace(ub.SHORT_LABEL, "")) + + +def _strip_bundle_logo(paths): + import glob + for path in glob.glob(os.path.join(paths["labext_static"], "*.js")): + with open(path, encoding = "utf-8") as f: + text = f.read() + with open(path, "w", encoding = "utf-8") as f: + f.write(text.replace(ub.LOGO_DATA_URI_PREFIX, "data:image/png;base64,XXXX")) + + +def _remove_logo_png(paths): + os.remove(paths["logo"]) + + +def _empty_favicon(paths): + open(paths["favicon"], "w").close() + + +def _disable_unsloth_ext(paths): + with open(paths["page_configs"][0], "w", encoding = "utf-8") as f: + json.dump({"disabledExtensions": {ub.LABEXT_NAME: True}}, f) + + +def _disable_unsloth_plugin(paths): + with open(paths["page_configs"][0], "w", encoding = "utf-8") as f: + json.dump({"disabledExtensions": {ub.ABOUT_PLUGIN_ID: True}}, f) + + +def _disable_unsloth_ext_list_form(paths): + # Older JupyterLab configs used a list of ids rather than an {id: bool} map. + with open(paths["page_configs"][0], "w", encoding = "utf-8") as f: + json.dump({"disabledExtensions": [ub.SPLASH_PLUGIN_ID]}, f) + + +@pytest.mark.parametrize( + "mutate", + [ + _remove_license, + _blank_license, + _remove_login, + _strip_login_source, + _strip_login_copyright, + _drop_theme, + _rebrand_labext, + _strip_bundle_phrase, + _strip_bundle_logo, + _remove_logo_png, + _empty_favicon, + _disable_unsloth_ext, + _disable_unsloth_plugin, + _disable_unsloth_ext_list_form, + ], +) +def test_negative_each_marker_is_enforced(tmp_path, mutate): + paths = _stage(tmp_path) + assert ub.verify_branding(paths) == [], "baseline should be clean before mutation" + mutate(paths) + problems = ub.verify_branding(paths) + assert problems, "stripping " + mutate.__name__ + " must be detected" + + +def test_disabling_stock_plugins_is_allowed(tmp_path): + """We disable the stock logo/splash ourselves -- the guard must not flag those.""" + paths = _stage(tmp_path) + with open(paths["page_configs"][0], "w", encoding = "utf-8") as f: + json.dump( + { + "disabledExtensions": { + "@jupyterlab/application-extension:logo": True, + "@jupyterlab/apputils-extension:splash": True, + } + }, + f, + ) + assert ub.verify_branding(paths) == [] + + +def test_attribution_sources_have_no_encoded_obfuscation(): + """Plain readable strings only -- no base64/decoder tricks (antivirus-safe).""" + src_dir = os.path.join(REPO, "docker", "jupyter") + files = [ + os.path.join(src_dir, "unsloth_branding.py"), + os.path.join(src_dir, "unsloth_labext", "src", "branding.ts"), + os.path.join(src_dir, "unsloth_labext", "src", "about.ts"), + os.path.join(src_dir, "unsloth_labext", "src", "splash.ts"), + ] + forbidden = [ + "b64decode", + "b64encode", + "atob(", + "btoa(", + "fromCharCode", + "unescape(", + "rot13", + "codecs.decode", + ] + for path in files: + with open(path, encoding = "utf-8") as f: + text = f.read() + for token in forbidden: + assert token not in text, path + " uses obfuscation token: " + token + + +def test_canonical_phrase_is_plain_text_in_definition_files(): + """The attribution lives as plain readable text in both definition files. + + branding.ts holds the full PHRASE as ONE contiguous literal (so webpack keeps + it whole in the bundle for the guard to grep). unsloth_branding.py keeps the + markers as plain constants (the runtime PHRASE value matches, even though the + source wraps it across adjacent literals).""" + src_dir = os.path.join(REPO, "docker", "jupyter") + ts = open( + os.path.join(src_dir, "unsloth_labext", "src", "branding.ts"), encoding = "utf-8" + ).read() + assert ub.PHRASE in ts, "branding.ts must hold the full PHRASE as one literal" + py = open(os.path.join(src_dir, "unsloth_branding.py"), encoding = "utf-8").read() + for marker in (ub.SHORT_LABEL, ub.COPYRIGHT, ub.SOURCE_URL, ub.AGPL_URL, ub.THEME_NAME): + assert marker in py, "unsloth_branding.py missing plain marker: " + marker diff --git a/tests/test_generate_kwarg_gate.py b/tests/test_generate_kwarg_gate.py index 00b3ddf6ee..43eb489cef 100644 --- a/tests/test_generate_kwarg_gate.py +++ b/tests/test_generate_kwarg_gate.py @@ -130,6 +130,52 @@ def test_generate_kwarg_gate(): assert got is expected, f"{name}: got {got}, expected {expected}" +# --- v5 logits-to-keep filtering ------------------------------------------ +# transformers >= 5 injects logits_to_keep=1 in generate() itself, but the +# injection is guarded by `"logits_to_keep" not in model_kwargs`, so it is a +# DEFAULT. An explicit caller value must survive: popping unconditionally turns +# logits_to_keep=0 (give me the full sequence) into 1 without telling anyone. +# The only values that must be stripped are the ones the strict validator would +# raise on, which is exactly what the gate above predicts. + + +def _filter_logits_kwargs(model, kwargs): + """The v5 branch of unsloth_base_fast_generate, as a testable function.""" + for key in ("logits_to_keep", "num_logits_to_keep"): + if key in kwargs and not accepts(model, key): + kwargs.pop(key, None) + return kwargs + + +def test_v5_preserves_a_supported_caller_value(): + model = PrepHasKwargs_ForwardHasKey() + # 0 means "all logits"; silently rewriting it to 1 changes the output shape. + assert _filter_logits_kwargs(model, {"logits_to_keep": 0}) == {"logits_to_keep": 0} + assert _filter_logits_kwargs(model, {"logits_to_keep": 5}) == {"logits_to_keep": 5} + + +def test_v5_strips_a_value_the_model_would_reject(): + # num_logits_to_keep was renamed away in v5, so the validator raises on it. + model = PrepHasKwargs_ForwardHasKey() + assert _filter_logits_kwargs(model, {"num_logits_to_keep": 1}) == {} + # A VLM whose top-level forward has no logits_to_keep at all. + assert _filter_logits_kwargs(NoPrepare(), {"logits_to_keep": 1}) == {} + + +def test_v5_leaves_other_kwargs_alone(): + model = PrepHasKwargs_ForwardHasKey() + out = _filter_logits_kwargs(model, {"logits_to_keep": 2, "max_new_tokens": 8}) + assert out == {"logits_to_keep": 2, "max_new_tokens": 8} + + +def test_source_has_no_unconditional_pop(): + src = open(VISION).read() + assert ( + 'kwargs.pop("logits_to_keep", None)\n kwargs.pop("num_logits_to_keep", None)' + not in src + ), "the v5 branch must not drop caller-supplied logits_to_keep unconditionally" + + if __name__ == "__main__": test_generate_kwarg_gate() for name, _, _, _ in CASES: diff --git a/tests/validate_studio_features.py b/tests/validate_studio_features.py new file mode 100644 index 0000000000..73930f9017 --- /dev/null +++ b/tests/validate_studio_features.py @@ -0,0 +1,278 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +"""Cross-platform validation of the Unsloth Docker JupyterLab/notebook features. + +Runs WITHOUT Docker or a GPU, so it can execute on the Linux/macOS/Windows CI +lanes. It exercises the actual notebook-helper logic (not just py_compile) and +checks the shipped JupyterLab config + labextension source, so a regression in +the notebook organisation, Colab compatibility, Colab-intro/widget stripping, +sidecar-log gating, the labextension plugins, the JupyterLab defaults, or the +login branding fails CI on every device. + +Usage: python tests/validate_studio_features.py +Exit 0 = all checks pass; non-zero = at least one failed. +""" + +from __future__ import annotations + +import importlib +import json +import os +import sys + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DOCKER = os.path.join(ROOT, "docker") +JUPYTER = os.path.join(DOCKER, "jupyter") +LABEXT = os.path.join(JUPYTER, "unsloth_labext") +sys.path.insert(0, DOCKER) + +_failures: list[str] = [] + + +def check( + name: str, + cond: bool, + detail: str = "", +) -> None: + status = "PASS" if cond else "FAIL" + print(f" [{status}] {name}" + (f" -- {detail}" if detail and not cond else "")) + if not cond: + _failures.append(name) + + +# 1. Colab cell-magic compatibility (#@title then %%capture) +def test_colab_compat() -> None: + print("colab cell-magic compat (unsloth_colab_compat):") + m = importlib.import_module("unsloth_colab_compat") + out = m.colab_cell_magic_fix(["#@title Setup\n", "%%capture\n", "!pip install x\n"]) + check("magic hoisted above #@title", out[0] == "%%capture\n" and "#@title Setup\n" in out) + # idempotent / already on top + same = ["%%capture\n", "print(1)\n"] + check("no-op when magic already first", m.colab_cell_magic_fix(same) == same) + # non-magic cell untouched + plain = ["x = 1\n", "y = 2\n"] + check("plain cell untouched", m.colab_cell_magic_fix(plain) == plain) + # content magic (%%writefile) NOT hoisted into the written file body + wf = ["#@title Config\n", "%%writefile config.json\n", "{}\n"] + check("content magic (%%writefile) left untouched", m.colab_cell_magic_fix(wf) == wf) + # safe magic with arg still hoisted + bash = ["#@title Run\n", "%%bash\n", "echo hi\n"] + check("safe magic (%%bash) hoisted", m.colab_cell_magic_fix(bash)[0] == "%%bash\n") + + +# 2. Notebook categorisation (clean_section) + README parsing +def test_nb_view() -> None: + print("notebook view (unsloth_nb_view):") + v = importlib.import_module("unsloth_nb_view") + check( + "clean_section dash/slash -> space", + v.clean_section("### GRPO-Reinforcement/Learning Notebooks") + == "GRPO Reinforcement Learning Notebooks", + v.clean_section("### GRPO-Reinforcement/Learning Notebooks"), + ) + check( + "clean_section strips hashes/space", + v.clean_section("## Main Notebooks ") == "Main Notebooks", + ) + + +# 3. Colab-intro + stale-widget stripping +def test_strip() -> None: + print("notebook strip (unsloth_nb_strip_colab):") + s = importlib.import_module("unsloth_nb_strip_colab") + nb = { + "metadata": {"widgets": {"application/vnd.jupyter.widget-state+json": {"x": 1}}}, + "cells": [ + { + "cell_type": "markdown", + "source": [ + 'To run this, press "Runtime" ... Tesla T4 Google Colab instance!\n', + "\n", + "You will learn how to ...\n", + ], + }, + { + "cell_type": "code", + "source": ["print(1)\n"], + "outputs": [ + {"output_type": "stream", "name": "stdout", "text": "ok\n"}, + { + "output_type": "display_data", + "data": { + "application/vnd.jupyter.widget-view+json": {"model_id": "abc"}, + "text/plain": "0%| | 0/10", + }, + }, + ], + }, + ], + } + changed1 = s._strip_intro(nb) + changed2 = s._clean_widgets(nb) + check( + "intro line stripped", + changed1 and not any("to run this, press" in (l.lower()) for l in nb["cells"][0]["source"]), + ) + check("intro body kept", any("You will learn" in l for l in nb["cells"][0]["source"])) + wv = sum( + 1 + for c in nb["cells"] + for o in (c.get("outputs", []) or []) + if "application/vnd.jupyter.widget-view+json" in (o.get("data", {}) or {}) + ) + check("widget-view outputs removed", changed2 and wv == 0) + check( + "non-widget outputs kept", + any( + o.get("output_type") == "stream" + for c in nb["cells"] + for o in (c.get("outputs", []) or []) + ), + ) + check("metadata.widgets removed", "widgets" not in nb["metadata"]) + # idempotent + check("strip idempotent", not s._strip_intro(nb) and not s._clean_widgets(nb)) + + +# 4. Sidecar-log gating +def test_sidecar_log_gate() -> None: + print("sidecar log gate (unsloth_nb_compat):") + c = importlib.import_module("unsloth_nb_compat") + old = os.environ.pop("UNSLOTH_ENABLE_LOGGING", None) + try: + check("logging off by default", c._logging_enabled() is False) + os.environ["UNSLOTH_ENABLE_LOGGING"] = "1" + check("logging on with env=1", c._logging_enabled() is True) + os.environ["UNSLOTH_ENABLE_LOGGING"] = "0" + check("logging off with env=0", c._logging_enabled() is False) + finally: + os.environ.pop("UNSLOTH_ENABLE_LOGGING", None) + if old is not None: + os.environ["UNSLOTH_ENABLE_LOGGING"] = old + + +# 5. JupyterLab defaults (overrides.json) +def test_overrides() -> None: + print("jupyterlab defaults (jupyter/overrides.json):") + path = os.path.join(JUPYTER, "overrides.json") + check("overrides.json exists", os.path.isfile(path)) + if not os.path.isfile(path): + return + with open(path, encoding = "utf-8") as f: + d = json.load(f) # raises -> CI fails if invalid JSON + themes = d.get("@jupyterlab/apputils-extension:themes", {}) + check( + "default theme = Unsloth Dark", + themes.get("theme") == "Unsloth Dark", + str(themes.get("theme")), + ) + check("adaptive theme on", themes.get("adaptive-theme") is True) + check("preferred dark = Unsloth Dark", themes.get("preferred-dark-theme") == "Unsloth Dark") + tracker = d.get("@jupyterlab/notebook-extension:tracker", {}) + check( + "windowingMode none", + tracker.get("windowingMode") == "none", + str(tracker.get("windowingMode")), + ) + notif = d.get("@jupyterlab/apputils-extension:notification", {}) + check( + "news prompt off", + str(notif.get("fetchNews")) == "false" and notif.get("checkForUpdates") is False, + ) + panel = d.get("@jupyterlab/notebook-extension:panel", {}) + labels = [t.get("label", "") for t in panel.get("toolbar", [])] + check( + "Restart & Run All label (single >>)", + any(l == "Restart & Run All" for l in labels) and not any(">>" in l for l in labels), + str(labels), + ) + + +# 6. Labextension source (plugins) + login branding assets +def test_labext_and_branding() -> None: + print("labextension + branding assets:") + pkg = os.path.join(LABEXT, "package.json") + check("labext package.json exists", os.path.isfile(pkg)) + if os.path.isfile(pkg): + with open(pkg, encoding = "utf-8") as f: + p = json.load(f) + check("labext name unsloth-jupyterlab", p.get("name") == "unsloth-jupyterlab") + check("labext themePath set", bool(p.get("jupyterlab", {}).get("themePath"))) + # Concatenate every .ts module under src/ so plugins defined in their own + # files (cellNav, colabTitle, outputSelect, uiChrome) are all covered. + src_dir = os.path.join(LABEXT, "src") + all_src = "" + if os.path.isdir(src_dir): + for fn in sorted(os.listdir(src_dir)): + if fn.endswith(".ts"): + with open(os.path.join(src_dir, fn), encoding = "utf-8") as f: + all_src += f.read() + "\n" + for plug in [ + "unsloth-jupyterlab:theme", + "unsloth-jupyterlab:cell-nav", + "unsloth-jupyterlab:logo", + "unsloth-jupyterlab:colab-title", + "unsloth-jupyterlab:output-select-all", + "unsloth-jupyterlab:ui-chrome", + ]: + check(f"plugin present: {plug}", plug in all_src) + # The two newest plugins are also exported from index.ts (wired in). + index = os.path.join(src_dir, "index.ts") + index_src = open(index, encoding = "utf-8").read() if os.path.isfile(index) else "" + check("outputSelect wired in index.ts", "outputSelectPlugin" in index_src) + check("uiChrome wired in index.ts", "uiChromePlugin" in index_src) + # uiChrome hides the right activity bar; CTRL+A output-select selects nodes. + check("right activity bar hidden", "jp-mod-right" in all_src and "display: none" in all_src) + check("ctrl+A output select", "selectNodeContents" in all_src) + # The remembered pointer-down is only replaced by another pointer-down, but + # J/K/arrow cell navigation fires none, so it has to be revalidated (still in + # the document, still in the ACTIVE cell) before it is used as the fallback -- + # otherwise Ctrl+A on a later cell selects the old output and swallows + # JupyterLab's notebook:select-all. + check( + "ctrl+A fallback revalidated", + "isConnected" in all_src and "jp-mod-active" in all_src, + ) + # branding assets + login = os.path.join(JUPYTER, "login.html") + login_src = open(login, encoding = "utf-8").read() if os.path.isfile(login) else "" + check("login.html branded", "unsloth-login-card" in login_src) + check( + "login.html uses sloth stickers", + 'static_url("sloth/' in login_src or "static_url('sloth/" in login_src, + ) + check("favicon.ico present", os.path.isfile(os.path.join(JUPYTER, "favicon.ico"))) + check("logo.png present", os.path.isfile(os.path.join(JUPYTER, "logo.png"))) + check( + "sloth sticker installer present", + os.path.isfile(os.path.join(JUPYTER, "install_sloth_stickers.py")), + ) + + +def main() -> int: + print("=== Unsloth Studio/notebook feature validation ===") + for t in ( + test_colab_compat, + test_nb_view, + test_strip, + test_sidecar_log_gate, + test_overrides, + test_labext_and_branding, + ): + try: + t() + except Exception as e: # a thrown exception is a failure, not a crash + _failures.append(f"{t.__name__}: {e!r}") + print(f" [FAIL] {t.__name__} raised {e!r}") + print() + if _failures: + print(f"FAILED ({len(_failures)}): " + ", ".join(_failures)) + return 1 + print("ALL CHECKS PASSED") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/unsloth/_gpu_init.py b/unsloth/_gpu_init.py index 984057e9f7..8c9353df20 100644 --- a/unsloth/_gpu_init.py +++ b/unsloth/_gpu_init.py @@ -114,6 +114,24 @@ del maybe_set_windows_rocm_bnb_version # Fixes https://github.com/unslothai/unsloth/issues/1266 os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python" +# `docker --gpus '"device=N"'` sets NVIDIA_VISIBLE_DEVICES but not +# CUDA_VISIBLE_DEVICES, so Inductor's compile-worker pool can't enumerate the +# cgroup-pinned GPU ("Could not find an active GPU backend"). Force a single +# in-process compile thread. Trigger only on pinned ids, not "all"/"none"/"void"/"" +# (the `--gpus all` default). Opt out with UNSLOTH_FORCE_SINGLE_COMPILE_WORKER=0. +_nvd = os.environ.get("NVIDIA_VISIBLE_DEVICES", "").strip().lower() +_cgroup_pinned = _nvd not in ("", "all", "none", "void") +if ( + os.environ.get("UNSLOTH_FORCE_SINGLE_COMPILE_WORKER", "auto") != "0" + and _cgroup_pinned + and "CUDA_VISIBLE_DEVICES" not in os.environ +): + # Honour an existing thread count; always plant the sentinel for the zoo patch. + if os.environ.get("TORCHINDUCTOR_COMPILE_THREADS") in (None, "", "1"): + os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = "1" + os.environ["UNSLOTH_FORCE_SINGLE_COMPILE_WORKER"] = "1" +del _nvd, _cgroup_pinned + # [TODO] Check why some GPUs don't work # "pinned_use_cuda_host_register:True,"\ # "pinned_num_register_threads:8" @@ -158,6 +176,25 @@ except ModuleNotFoundError: except: raise +# Re-assert single-compile-worker after unsloth_zoo's patch_torch_compile (which +# historically popped TORCHINDUCTOR_COMPILE_THREADS). Set the Inductor config +# directly and patch the zoo's determine_compile_threads so every options dict +# sees 1. No-op when the user opted out. +if os.environ.get("UNSLOTH_FORCE_SINGLE_COMPILE_WORKER", "0") == "1": + try: + torch._inductor.config.compile_threads = 1 + except Exception: + pass + os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = "1" + try: + setattr( + importlib.import_module("unsloth_zoo.temporary_patches.common"), + "determine_compile_threads", + lambda: 1, + ) + except Exception: + pass + from unsloth_zoo.device_type import ( is_hip, get_device_type, @@ -261,7 +298,12 @@ del patch_peft_weight_converter_compatibility del patch_accelerate_recursively_apply # Torch 2.4 has including_emulation -if DEVICE_TYPE == "cuda": +if DEVICE_TYPE == "cuda" and not torch.cuda.is_available(): + # UNSLOTH_ALLOW_CPU=1 keeps DEVICE_TYPE "cuda" on driverless hosts; probing + # would raise. bf16 on (CPU bf16 kernels exist, fp16 largely don't). + SUPPORTS_BFLOAT16 = True + torch.cuda.is_bf16_supported = lambda *args, **kwargs: True +elif DEVICE_TYPE == "cuda": major_version, minor_version = torch.cuda.get_device_capability() SUPPORTS_BFLOAT16 = major_version >= 8 diff --git a/unsloth/dataprep/synthetic.py b/unsloth/dataprep/synthetic.py index 9c529d5a03..351533c00a 100644 --- a/unsloth/dataprep/synthetic.py +++ b/unsloth/dataprep/synthetic.py @@ -267,7 +267,9 @@ class SyntheticDataKit: stderr = subprocess.PIPE, start_new_session = True, ) - ready_re = re.compile(r"Starting vLLM API server(?:\s+\d+)?\s+on\b") + # Accept both "Starting vLLM API server on" (<= 0.18) and "Starting vLLM + # server on" (0.19), with the optional server index some versions insert. + ready_re = re.compile(r"Starting vLLM(?:\s+API)?\s+server(?:\s+\d+)?\s+on\b") self.vllm_process = vllm_process self.stdout_capture = PipeCapture( vllm_process.stdout, @@ -282,12 +284,29 @@ class SyntheticDataKit: keep_lines = 2000, echo = False, name = "vLLM STDERR", - ready_regex = None, + # vLLM >= 0.19 logs startup lines to STDERR; watching stdout alone + # makes a healthy server look like a timeout and get killed. + ready_regex = ready_re, text = False, ) # we don't print stderr to console but self.stderr_capture.tail(200) will print the last 200 lines - ready = self.stdout_capture.wait_for_ready(timeout = timeout) + ready = False + # timeout None/0 waits indefinitely (large models / slow downloads); + # a positive value is a deadline. + deadline = (time.monotonic() + timeout) if timeout else None + while True: + # Cap the wait to the remaining budget so we don't overshoot the deadline. + _wait = 1 if deadline is None else min(1, deadline - time.monotonic()) + if _wait <= 0: + break + if self.stdout_capture.wait_for_ready( + timeout = _wait + ) or self.stderr_capture.wait_for_ready(timeout = 0): + ready = True + break + if self.vllm_process.poll() is not None: + break if not ready: if self.stdout_capture.has_closed() or self.vllm_process.poll() is not None: print("Stdout stream ended before readiness message detected.") diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index dfb8082c46..bb0f56731b 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -1991,7 +1991,11 @@ SUPPORTS_BFLOAT16 = False HAS_FLASH_ATTENTION = False HAS_FLASH_ATTENTION_SOFTCAPPING = False -if DEVICE_TYPE == "cuda": +if DEVICE_TYPE == "cuda" and not torch.cuda.is_available(): + # UNSLOTH_ALLOW_CPU=1 keeps DEVICE_TYPE "cuda" on driverless hosts; + # bf16 CPU kernels exist, fp16 ones largely do not. + SUPPORTS_BFLOAT16 = True +elif DEVICE_TYPE == "cuda": major_version, minor_version = torch.cuda.get_device_capability() torch.cuda.get_device_capability = functools.cache(torch.cuda.get_device_capability) @@ -2100,7 +2104,7 @@ try: # causing sm_90a kernels to be attempted on non-Hopper GPUs (CUDA error in # flash_fwd_launch_template.h:188). Fixed in 0.0.33 with `<= (9, 0)`. # See https://github.com/facebookresearch/xformers/issues/1329 - if DEVICE_TYPE == "cuda": + if DEVICE_TYPE == "cuda" and torch.cuda.is_available(): major_version, minor_version = torch.cuda.get_device_capability() if (f"{major_version}.{minor_version}" in ("10.0", "11.0", "12.0")) and ( Version(xformers_version) <= Version("0.0.32.post2") diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 0d55f272e3..5509a8a8e7 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -510,26 +510,41 @@ def unsloth_base_fast_generate(self, *args, **kwargs): ): kwargs.pop("mm_token_type_ids", None) - # VLMs do not allow logits_to_keep - global NUM_LOGITS_TO_KEEP - if arch not in NUM_LOGITS_TO_KEEP: - m = self - # Find which is used: num_logits_to_keep or logits_to_keep - while hasattr(m, "model"): - if hasattr(m, "forward"): - keys = inspect.signature(m.forward).parameters.keys() - if "num_logits_to_keep" in keys: - NUM_LOGITS_TO_KEEP[arch] = "num_logits_to_keep" - break - elif "logits_to_keep" in keys: - NUM_LOGITS_TO_KEEP[arch] = "logits_to_keep" - break - m = m.model + # VLMs do not allow logits_to_keep. transformers >= 5.0 sets it itself in + # generate(), so pre-injecting is redundant there, and the arch walk below + # can pick a key the top-level model rejects. Skip the injection on v5+. + if Version(transformers_version) < Version("5.0.0.dev0"): + global NUM_LOGITS_TO_KEEP if arch not in NUM_LOGITS_TO_KEEP: - NUM_LOGITS_TO_KEEP[arch] = None - key = NUM_LOGITS_TO_KEEP[arch] - if key is not None and key not in kwargs and _unsloth_generate_accepts_kwarg(self, key): - kwargs[key] = 1 + m = self + # Find which is used: num_logits_to_keep or logits_to_keep + while hasattr(m, "model"): + if hasattr(m, "forward"): + keys = inspect.signature(m.forward).parameters.keys() + if "num_logits_to_keep" in keys: + NUM_LOGITS_TO_KEEP[arch] = "num_logits_to_keep" + break + elif "logits_to_keep" in keys: + NUM_LOGITS_TO_KEEP[arch] = "logits_to_keep" + break + m = m.model + if arch not in NUM_LOGITS_TO_KEEP: + NUM_LOGITS_TO_KEEP[arch] = None + key = NUM_LOGITS_TO_KEEP[arch] + if key is not None and key not in kwargs and _unsloth_generate_accepts_kwarg(self, key): + kwargs[key] = 1 + else: + # v5's own injection (generation/utils.py) is guarded by + # `"logits_to_keep" not in model_kwargs`, so it is a default, not an + # override: an explicit caller value survives and must not be dropped. + # Popping unconditionally silently rewrites logits_to_keep=0 (full + # sequence) into 1. Only strip a key this model would reject, which is + # what the strict validator raises on: num_logits_to_keep everywhere + # (renamed away in v5), and logits_to_keep on the VLMs whose top-level + # forward does not take it. + for _logits_kwarg in ("logits_to_keep", "num_logits_to_keep"): + if _logits_kwarg in kwargs and not _unsloth_generate_accepts_kwarg(self, _logits_kwarg): + kwargs.pop(_logits_kwarg, None) model_eos_token_id = getattr(self.config, "eos_token_id", None) if model_eos_token_id is not None and hasattr(model_eos_token_id, "__iter__"):