docker: second review pass fixes

- Dockerfile: lift numba past vllm's 0.61.2 pin after the numpy>=2.4
  re-upgrade; 0.61.2 refuses numpy 2.3+ at import time and the stack
  cannot move numpy down. Verified numba 0.65 + numpy 2.4.6 + vllm
  import cleanly together.
- docker-publish.yml: resolve UNSLOTH_ZOO_REF in a step that mirrors
  the pushed tag only when the tag exists in unsloth-zoo (the zoo
  currently cuts no tags, so blind mirroring broke every tag publish);
  falls back to main.
- Dockerfile.studio: Studio venv stays on cu128 for arm64 too, matching
  the base venv (cu130 wheels would lift the driver floor to 580+), and
  gets the same NVRTC cu13 swap for DGX Spark / GB10 sm_121 support.
- docker_confirm.sh: do not drop to CPU mode when docker info lacks a
  nvidia runtime entry; CDI installs and Docker Desktop WSL2 expose
  GPUs without one. The phase 3 --gpus probe is now the authority.
- docker_confirm.ps1: GPU selector built as an args array; comma device
  lists get version-aware CSV quoting (native arg passing changed in
  PowerShell 7.3).
- studio_launch.sh: no fixed Jupyter default password; generate a
  random one and print it when JUPYTER_PASSWORD is unset. Env snapshot
  for SSH sessions now written via shlex.quote instead of sed so
  values with quotes or command substitution cannot break or inject
  into /etc/profile.d.
- install.ps1: honour UNSLOTH_TORCH_INDEX_FAMILY like install.sh does.
This commit is contained in:
Daniel Han 2026-06-12 05:59:52 +00:00
commit 81b0d1ef10
7 changed files with 120 additions and 40 deletions

View file

@ -111,6 +111,24 @@ jobs:
with:
images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
# Mirror the unsloth tag into the zoo ONLY when that tag actually
# exists there. unsloth's v* tags are Studio releases the zoo never
# cuts (the zoo repo currently has no tags at all), so blindly
# mirroring github.ref_name made every tag publish fail inside the
# Dockerfile's 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
echo "ref=${REF:-main}" >> "$GITHUB_OUTPUT"
echo "unsloth-zoo ref: ${REF:-main}"
- name: Build and push (per-arch by digest)
id: build
uses: docker/build-push-action@v6
@ -134,12 +152,11 @@ jobs:
# scheduled runs: bake the triggering commit SHA. Falls back
# to `main` for any other event class.
UNSLOTH_REF=${{ github.event.inputs.unsloth_ref || (startsWith(github.ref, 'refs/tags/') && github.ref_name) || github.sha || 'main' }}
# UNSLOTH_ZOO_REF mirrors the tag case (unsloth-zoo cuts the same
# release tag, e.g. 2026.5.8, alongside unsloth) so release-tag
# images install a matched zoo. SHA-based branch pushes can't be
# mirrored -- the SHA doesn't exist in the zoo repo -- so they
# fall through to `main`. Workflow-dispatch can override.
UNSLOTH_ZOO_REF=${{ github.event.inputs.unsloth_zoo_ref || (startsWith(github.ref, 'refs/tags/') && github.ref_name) || 'main' }}
# UNSLOTH_ZOO_REF comes from the resolve step above: explicit
# workflow-dispatch input, else the pushed tag IF the zoo repo
# has it, else `main`. SHA-based branch pushes always fall to
# `main` -- the SHA doesn't exist in the zoo repo.
UNSLOTH_ZOO_REF=${{ steps.zoo_ref.outputs.ref }}
# Stash the per-arch digest as an artifact for the merge job to pick up.
# Filenames need to be unique across the matrix; `platform` contains a

View file

@ -215,9 +215,18 @@ RUN set -eux \
${VENV}/bin/uv pip install \
--python ${VENV}/bin/python \
--upgrade "numpy>=2.4"; \
echo ">> vLLM installed (numpy re-upgraded post-vllm):"; \
# vLLM pins numba==0.61.2, which hard-refuses numpy >= 2.3 at import
# time -- and the rest of the stack needs numpy >= 2.3, so the numpy
# ceiling cannot move down. Lift numba to a release that supports
# numpy 2.4 (verified: numba 0.65 imports cleanly and vllm still
# imports). Same intentional-override class as the numpy bump above.
${VENV}/bin/uv pip install \
--python ${VENV}/bin/python \
--upgrade "numba>=0.62"; \
echo ">> vLLM installed (numpy + numba re-upgraded post-vllm):"; \
${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')"; \
else \
echo ">> vLLM skipped (INSTALL_VLLM=${INSTALL_VLLM}, TARGETARCH=${TARGETARCH:-amd64})"; \
fi

View file

@ -18,10 +18,11 @@
#
# Open http://localhost:8000 for Studio (first-boot admin password is printed
# in the container logs and persisted under /opt/unsloth-studio/auth/) and
# http://localhost:8888 for JupyterLab (password: JUPYTER_PASSWORD env,
# default `unsloth`). On hosts without GPU passthrough (Docker Desktop on
# macOS, Windows without WSL2 GPU) add -e UNSLOTH_ALLOW_CPU=1: training is
# unavailable but Studio chat / Data Recipes / GGUF tooling / Jupyter work.
# http://localhost:8888 for JupyterLab (password: JUPYTER_PASSWORD env; when
# unset a random one is generated and printed in the container logs). On
# hosts without GPU passthrough (Docker Desktop on macOS, Windows without
# WSL2 GPU) add -e UNSLOTH_ALLOW_CPU=1: training is unavailable but Studio
# chat / Data Recipes / GGUF tooling / Jupyter work.
#
# CI pins BASE_IMAGE to the just-published multi-arch base digest so the two
# images always ship the same stack.
@ -74,18 +75,20 @@ RUN apt-get update \
# UNSLOTH_TORCH_INDEX_FAMILY pins the torch wheel index for the Studio
# venv: at build time there is no GPU and no nvidia-smi, so install.sh's
# probing would land on cpu or cu126 wheels depending on which host built
# the image. The image targets CUDA: cu128 on amd64 (Turing..Blackwell,
# same line as the base venv), cu130 on arm64 (DGX Spark / Grace, the
# aarch64 CUDA wheel line).
# the image. cu128 on BOTH arches, mirroring the base venv: cu130 wheels
# would silently lift the arm64 driver floor to 580+ while the base venv
# keeps the documented 570+ floor. DGX Spark / GB10 (sm_121) support comes
# from the same NVRTC cu13 swap the base image applies to its venv --
# repeated below for the Studio venv's own bundled libnvrtc (the base's
# arm64 layer already installed cuda-nvrtc-13-0, so the cu13 .so exists).
#
# fetch+checkout FETCH_HEAD instead of `clone --branch` because the CI
# pipeline passes a commit SHA as the ref (clone --branch only accepts
# branch/tag names).
RUN set -eux \
&& case "${TARGETARCH:-amd64}" in \
amd64) TORCH_FAMILY="cu128" ;; \
arm64) TORCH_FAMILY="cu130" ;; \
*) echo "ERROR: unsupported TARGETARCH=${TARGETARCH}" >&2; exit 1 ;; \
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" \
@ -98,7 +101,15 @@ RUN set -eux \
&& UNSLOTH_STUDIO_HOME="${UNSLOTH_STUDIO_HOME}" \
UNSLOTH_TORCH_INDEX_FAMILY="${TORCH_FAMILY}" \
bash install.sh --local \
&& rm -rf "${UNSLOTH_STUDIO_HOME}/src/.git" /root/.cache
&& rm -rf "${UNSLOTH_STUDIO_HOME}/src/.git" /root/.cache \
&& if [ "${TARGETARCH:-amd64}" = "arm64" ]; then \
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" ]; then \
mv "${NVRTC_DIR}/libnvrtc.so.12" "${NVRTC_DIR}/libnvrtc.so.12.cu128.orig"; \
ln -s /usr/local/cuda-13.0/lib64/libnvrtc.so.13 "${NVRTC_DIR}/libnvrtc.so.12"; \
fi; \
done; \
fi
COPY supervisord.conf /etc/supervisor/supervisord.conf
COPY studio_launch.sh /usr/local/bin/unsloth-studio-launch

View file

@ -98,14 +98,28 @@ Hr
Bold "3) Container runtime check"
# Mirror docker_confirm.sh's GPU selector translation: bare indices and
# comma lists become device= selectors (Docker reads a bare integer for
# --gpus as a COUNT, not an index).
# --gpus as a COUNT, not an index). Built as an args array so every docker
# run call splats it identically.
#
# Comma lists are special: docker CSV-parses the --gpus value, so a list
# must arrive as a literal "device=0,1" INCLUDING the double quotes. How
# PowerShell passes embedded quotes to native commands changed in 7.3
# (PSNativeCommandArgumentPassing), so pick the escaping per version;
# single selectors need no quoting anywhere.
$GPU_SELECTOR = "all"
if ($GPUS -notin @("auto", "all", "none")) {
$GPU_SELECTOR = if ($GPUS -like "device=*") { $GPUS } else { "`"device=$GPUS`"" }
$sel = $GPUS -replace "^device=", ""
if ($sel -match ",") {
if ($PSVersionTable.PSVersion -ge [version]"7.3") { $GPU_SELECTOR = '"device=' + $sel + '"' }
else { $GPU_SELECTOR = '\"device=' + $sel + '\"' }
} else {
$GPU_SELECTOR = "device=$sel"
}
}
$GpuRunArgs = @("--gpus", $GPU_SELECTOR)
if ($GPU_MODE) {
$log = Join-Path $WORK "gpu_check.log"
docker run --rm --gpus $GPU_SELECTOR $BASE_IMAGE python -c "import torch; assert torch.cuda.is_available(); print('torch', torch.__version__, '-', torch.cuda.get_device_name(0))" *> $log
docker run --rm @GpuRunArgs $BASE_IMAGE python -c "import torch; assert torch.cuda.is_available(); print('torch', torch.__version__, '-', torch.cuda.get_device_name(0))" *> $log
if ($LASTEXITCODE -eq 0) {
Ok ("torch.cuda available in-container: " + (Get-Content $log -Tail 1))
} else {
@ -131,7 +145,7 @@ Bold "4) Training smoke"
if ($GPU_MODE -and -not $SKIP_TRAIN) {
$log = Join-Path $WORK "train_smoke.log"
$hfArgs = @(); if ($env:HF_TOKEN) { $hfArgs = @("-e", "HF_TOKEN") }
docker run --rm --gpus $GPU_SELECTOR --ipc=host @hfArgs $BASE_IMAGE python /workspace/smoke_test.py *> $log
docker run --rm @GpuRunArgs --ipc=host @hfArgs $BASE_IMAGE python /workspace/smoke_test.py *> $log
if ($LASTEXITCODE -eq 0) {
Ok "smoke_test.py: 5 LoRA steps completed"
Select-String -Path $log -Pattern "^step|loss" | Select-Object -Last 5 | ForEach-Object { Info $_.Line }
@ -160,7 +174,7 @@ Hr
# 6) Studio + JupyterLab ------------------------------------------------------
Bold "6) Studio + JupyterLab (full image)"
$runArgs = @("-d", "-p", "${PORT_STUDIO}:8000", "-p", "${PORT_JUPYTER}:8888")
if ($GPU_MODE) { $runArgs += @("--gpus", $GPU_SELECTOR) } else { $runArgs += @("-e", "UNSLOTH_ALLOW_CPU=1") }
if ($GPU_MODE) { $runArgs += $GpuRunArgs } else { $runArgs += @("-e", "UNSLOTH_ALLOW_CPU=1") }
$script:STUDIO_CID = (docker run @runArgs $IMAGE 2>(Join-Path $WORK "studio_run.err"))
if (-not $script:STUDIO_CID) {
Bad ("full image failed to start (see " + (Join-Path $WORK "studio_run.err") + ")")

View file

@ -83,22 +83,24 @@ fi
ok "docker daemon reachable ($(docker --version 2>/dev/null))"
GPU_MODE=0
NVRT_LISTED=0
if [ "$GPUS" = "none" ]; then
info "GPU mode : disabled by GPUS=none"
elif command -v nvidia-smi >/dev/null 2>&1 && nvidia-smi -L 2>/dev/null | grep -q '^GPU'; then
info "GPU(s) :"
nvidia-smi --query-gpu=index,name,compute_cap --format=csv,noheader 2>/dev/null | sed 's/^/ - /'
# `docker info | grep Runtimes:.*nvidia` misses CDI setups (docker 25+
# with nvidia-ctk cdi) and Docker Desktop's WSL2 backend, both of which
# expose GPUs without a host-visible runtime entry. Treat the listing as
# a hint only; phase 3 probes --gpus for real and demotes to CPU mode if
# the probe fails.
if docker info 2>/dev/null | grep -qi 'Runtimes:.*nvidia'; then
ok "NVIDIA GPU visible and docker has the nvidia runtime"
GPU_MODE=1
elif [ "$OS" = "Linux" ] && [ "$IS_WSL" = "1" ]; then
# Docker Desktop's WSL2 backend exposes GPUs without a host-visible
# nvidia runtime entry; --gpus all still works. Probe it for real below.
warn "nvidia runtime not listed by docker info (normal under Docker Desktop WSL2) - probing --gpus all directly"
GPU_MODE=1
ok "NVIDIA GPU visible and docker lists the nvidia runtime"
NVRT_LISTED=1
else
warn "NVIDIA GPU present but docker lacks the nvidia runtime - install nvidia-container-toolkit; falling back to CPU mode"
warn "nvidia runtime not listed by docker info (normal under CDI or Docker Desktop WSL2) - probing --gpus directly in phase 3"
fi
GPU_MODE=1
else
info "no NVIDIA GPU on the host (or nvidia-smi missing)"
fi
@ -137,7 +139,11 @@ if [ "$GPU_MODE" = "1" ]; then
>"$WORK/gpu_check.log" 2>&1; then
ok "torch.cuda available in-container: $(tail -1 "$WORK/gpu_check.log")"
else
bad "GPU passthrough failed (see $WORK/gpu_check.log) - falling back to CPU mode"
if [ "$NVRT_LISTED" = "1" ]; then
bad "GPU passthrough failed despite a listed nvidia runtime (see $WORK/gpu_check.log) - falling back to CPU mode"
else
warn "--gpus probe failed - docker has no nvidia runtime or CDI spec (install nvidia-container-toolkit); falling back to CPU mode"
fi
tail -5 "$WORK/gpu_check.log" | sed 's/^/ /'
GPU_MODE=0
fi

View file

@ -3,12 +3,13 @@
#
# 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 (default: unsloth)
# 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 (default unsloth)
# 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)
@ -21,19 +22,36 @@ export UNSLOTH_STUDIO_HOME="${UNSLOTH_STUDIO_HOME:-/opt/unsloth-studio}"
# without the `docker run -e` vars. Secrets are excluded on purpose: tokens,
# API keys and passwords stay in process env only, never on disk where an
# SSH session (or anything reading /etc/profile.d) could pick them up.
printenv | grep -E '^(HF_|CUDA_|NCCL_|JUPYTER_|UNSLOTH_|WANDB_|PATH=|TRITON_)' | \
grep -vE '^[^=]*(_TOKEN|_API_KEY|_PASSWORD|_SECRET|_LICENSE)=' | \
sed 's/^\([^=]*\)=\(.*\)$/export \1="\2"/' > /etc/profile.d/unsloth_env.sh || true
# shlex.quote() each value: env vars can contain quotes, $, backticks etc,
# and 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
# --- Jupyter -----------------------------------------------------------------
# Hash the password with jupyter's own helper; never store the plaintext.
# No fixed default password: when JUPYTER_PASSWORD is unset we generate a
# random one and print it once in the boot banner (docker logs).
JUPYTER_CONFIG_DIR=/root/.jupyter
if [[ ! -f "${JUPYTER_CONFIG_DIR}/jupyter_lab_config.py" ]]; then
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 - <<PY
from jupyter_server.auth import passwd
import os
print(passwd(os.environ.get("JUPYTER_PASSWORD", "unsloth")))
print(passwd(os.environ["JUPYTER_PASSWORD"]))
PY
)
cat > "${JUPYTER_CONFIG_DIR}/jupyter_lab_config.py" <<EOF
@ -60,7 +78,7 @@ fi
mkdir -p /workspace
echo "Unsloth Studio -> http://localhost:8000 (first-boot password below)"
echo "JupyterLab -> http://localhost:${JUPYTER_PORT} (password: JUPYTER_PASSWORD env, default 'unsloth')"
echo "JupyterLab -> http://localhost:${JUPYTER_PORT} (${JUPYTER_NOTE})"
if [[ "${UNSLOTH_ENABLE_SSHD}" == "true" ]]; then
echo "sshd -> port 22 (key-only)"
fi

View file

@ -1738,6 +1738,11 @@ shell.Run cmd, 0, False
# Mirrors Get-PytorchCudaTag in setup.ps1.
function Get-TorchIndexUrl {
$baseUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" }
# Explicit override (parity with install.sh):
# UNSLOTH_TORCH_INDEX_FAMILY=cu128|cu130|cu126|cpu|... pins the wheel
# index when probing is wrong or impossible (no GPU on the build host,
# containerised installs, CI).
if ($env:UNSLOTH_TORCH_INDEX_FAMILY) { return "$baseUrl/$($env:UNSLOTH_TORCH_INDEX_FAMILY)" }
if (-not $NvidiaSmiExe) { return "$baseUrl/cpu" }
try {
$output = Invoke-NvidiaSmiBounded $NvidiaSmiExe