unsloth/docker/test_locally.sh
Daniel Han 23a5b43180 test_locally.sh: pre-flight check for docker daemon connectivity
If the user is not in the 'docker' group, every docker command after the
pre-flight returns "permission denied while trying to connect to the Docker
daemon socket at /var/run/docker.sock". This used to surface as a confusing
buildx failure mid-Block-2, but the actual problem is a host permissions
issue that's settable up front.

Detect by running 'docker info' and checking its exit code (not just grep
on its output -- a permission failure prints to stderr and returns non-zero,
so the old grep-based check was a silent skip).

Also clarify the nvidia-runtime WARN: on Docker 28+ with CDI mode this is
a false positive most of the time. The real GPU-attach test is the smoke
run in Block 3a, where the container entrypoint catches missing GPUs with
an actionable message.
2026-05-24 07:50:30 +00:00

278 lines
11 KiB
Bash
Executable file

#!/usr/bin/env bash
# End-to-end Docker validation for the unsloth-blackwell image.
#
# Runs three blocks:
# 1. Host pre-flight (docker, nvidia-smi, nvidia runtime registered)
# 2. Build the image (no GPU required at build time)
# 3a. Smoke test: 5-step LoRA on Llama-3.2-1B (~1-2 min)
# 3b. Real workload: gpt-oss-20B fine-tuning notebook with max_steps=10
# (~10 min, needs ~30GB free for the model cache)
#
# Usage:
# bash docker/test_locally.sh # all blocks
# bash docker/test_locally.sh --skip-notebook # blocks 1-3a only (fast)
# bash docker/test_locally.sh --skip-build # assume $TAG already built
# TAG=my-image:latest bash docker/test_locally.sh
# HF_TOKEN=hf_xxx bash docker/test_locally.sh # for gated models (optional)
#
# All output is teed to $LOG_DIR (default /tmp/unsloth-docker-test/).
# Paste the listed log snippets back if anything fails.
set -uo pipefail
TAG="${TAG:-unsloth-blackwell:test}"
LOG_DIR="${LOG_DIR:-/tmp/unsloth-docker-test}"
SKIP_BUILD=0
SKIP_NOTEBOOK=0
while [[ $# -gt 0 ]]; do
case "$1" in
--skip-build) SKIP_BUILD=1; shift ;;
--skip-notebook) SKIP_NOTEBOOK=1; shift ;;
--tag) TAG="$2"; shift 2 ;;
--log-dir) LOG_DIR="$2"; shift 2 ;;
--help|-h) sed -n '2,20p' "$0"; exit 0 ;;
*) echo "Unknown flag: $1" >&2; exit 2 ;;
esac
done
mkdir -p "$LOG_DIR"
GREEN='\033[1;32m'; RED='\033[1;31m'; YELLOW='\033[1;33m'; BLUE='\033[1;34m'; NC='\033[0m'
banner() { printf "\n${BLUE}==== %s ====${NC}\n" "$*"; }
ok() { printf "${GREEN}OK${NC} %s\n" "$*"; }
warn() { printf "${YELLOW}WARN${NC} %s\n" "$*"; }
err() { printf "${RED}ERROR${NC} %s\n" "$*" >&2; }
fail() { err "$*"; exit 1; }
# ============================================================================
# Block 1: pre-flight
# ============================================================================
banner "Block 1: host pre-flight"
command -v docker >/dev/null 2>&1 || fail "docker not found on PATH"
echo " docker: $(docker --version)"
# Verify we can actually talk to the docker daemon as the current user.
# This catches the "user not in docker group" case up front, instead of
# letting docker buildx blow up with a "permission denied on /var/run/docker.sock"
# error that looks like a build failure but is really a host permissions issue.
DOCKER_INFO_OUT=$(docker info 2>&1)
DOCKER_INFO_RC=$?
if [[ $DOCKER_INFO_RC -ne 0 ]]; then
err "Cannot talk to the docker daemon as user '$USER'."
cat >&2 <<MSG
docker info exited $DOCKER_INFO_RC. The most common cause is that your user
is not in the 'docker' group. Fix:
sudo usermod -aG docker \$USER
newgrp docker # activate the new group in this shell
docker info | head -3 # verify
Then re-run this script in the same shell (or any new login session).
Alternative: run the script with sudo, but be aware it will use root's
home directory for HF cache (~/root/.cache/huggingface) which is probably
not what you want.
Raw docker info output:
$DOCKER_INFO_OUT
MSG
fail "docker daemon unreachable"
fi
echo " daemon: reachable as '$USER'"
if command -v nvidia-smi >/dev/null 2>&1; then
echo " host gpu: $(nvidia-smi --query-gpu=name --format=csv,noheader | head -1)"
echo " host driver: $(nvidia-smi --query-gpu=driver_version --format=csv,noheader | head -1)"
else
warn "nvidia-smi not on the host -- you may not be able to run --gpus all"
fi
# This grep only makes sense once we know `docker info` succeeded above.
if echo "$DOCKER_INFO_OUT" | grep -qiE 'Runtimes:.*nvidia'; then
echo " nvidia runtime: registered with docker"
else
warn "docker info does not list 'nvidia' as a runtime"
warn "(on Docker 28+ with CDI this is often a false positive; the real"
warn " test is whether --gpus all works in Block 3a below)"
warn "if --gpus all fails, install nvidia-container-toolkit:"
warn " https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html"
warn " then: sudo systemctl restart docker"
fi
ok "pre-flight done"
# ============================================================================
# Block 2: build
# ============================================================================
if [[ $SKIP_BUILD -eq 1 ]]; then
warn "skipping build (--skip-build); expecting $TAG to exist"
else
banner "Block 2: build $TAG"
# Find the build context: current dir, docker/ subdir, or clone the PR branch
if [[ -f "Dockerfile" && -f "smoke_test.py" ]]; then
BUILD_CTX="$PWD"
elif [[ -f "docker/Dockerfile" ]]; then
BUILD_CTX="$PWD/docker"
else
BUILD_CTX="/tmp/unsloth-pr/docker"
if [[ ! -d /tmp/unsloth-pr/.git ]]; then
echo " cloning docker-blackwell-build branch..."
git clone --depth 1 -b docker-blackwell-build \
https://github.com/unslothai/unsloth.git /tmp/unsloth-pr 2>&1 | tail -3
else
git -C /tmp/unsloth-pr pull --ff-only 2>&1 | tail -2
fi
fi
echo " build context: $BUILD_CTX"
BUILD_LOG="$LOG_DIR/build.log"
echo " log: $BUILD_LOG"
# The Dockerfile uses BuildKit-only features ('# syntax=docker/dockerfile:1.7'
# and 'RUN ... <<\'PY\'' heredocs). Docker 28 removed the legacy builder
# entirely -- DOCKER_BUILDKIT=1 now delegates to buildx, so without the
# buildx component installed there is no fallback that works. Fail fast
# with install instructions before attempting the build.
if ! docker buildx version >/dev/null 2>&1; then
cat >&2 <<'MSG'
ERROR: docker buildx is not installed.
The Dockerfile requires BuildKit (syntax=docker/dockerfile:1.7 + RUN heredocs).
Docker 28 removed the legacy builder, so buildx is required for any build.
Install buildx, then re-run this script:
Ubuntu / Debian (apt):
sudo apt-get update && sudo apt-get install -y docker-buildx
Ubuntu / Debian (Docker's official repo, recommended):
# Follow https://docs.docker.com/engine/install/ubuntu/ -- the docker-ce
# package bundles docker-buildx-plugin and is what most production guides
# assume. The Ubuntu-shipped docker.io package omits buildx.
RHEL / Fedora (dnf):
sudo dnf install -y docker-buildx-plugin
Manual install (any distro):
https://github.com/docker/buildx/releases (download into ~/.docker/cli-plugins/)
Verify with: docker buildx version
MSG
fail "docker buildx required -- install per the message above"
fi
echo " builder: docker buildx ($(docker buildx version | head -1))"
docker buildx build --progress=plain --load -t "$TAG" "$BUILD_CTX" 2>&1 | tee "$BUILD_LOG"
rc=${PIPESTATUS[0]}
if [[ $rc -ne 0 ]]; then
fail "docker build exited $rc -- see $BUILD_LOG"
fi
# Sanity check the build's own self-test ran and passed
if grep -q "FAIL: missing wheels\|sm_100 (B200) missing\|sm_120 (RTX 5090) missing" "$BUILD_LOG"; then
fail "build-time sanity check failed -- see $BUILD_LOG"
fi
grep -E "OK: torch 2.10.0|OK: all required wheels|OK: xformers \+ bitsandbytes" "$BUILD_LOG" || \
warn "could not find 'OK:' lines in build log -- did the verification step run?"
ok "built $TAG"
fi
# ============================================================================
# Block 3a: smoke test
# ============================================================================
banner "Block 3a: smoke test (5-step LoRA on Llama-3.2-1B)"
SMOKE_LOG="$LOG_DIR/smoke.log"
echo " log: $SMOKE_LOG"
docker run --rm --gpus all "$TAG" python /workspace/smoke_test.py 2>&1 | tee "$SMOKE_LOG"
rc=${PIPESTATUS[0]}
if [[ $rc -ne 0 ]]; then
fail "smoke test exited $rc -- see $SMOKE_LOG"
fi
if ! grep -q "all checks passed" "$SMOKE_LOG"; then
fail "smoke test did not print 'all checks passed' -- see $SMOKE_LOG"
fi
ok "smoke test passed"
# ============================================================================
# Block 3b: gpt-oss-20B fine-tuning notebook
# ============================================================================
if [[ $SKIP_NOTEBOOK -eq 1 ]]; then
warn "skipping gpt-oss-20B notebook (--skip-notebook)"
else
banner "Block 3b: gpt-oss-20B fine-tuning notebook (10 LoRA steps)"
GPT_LOG="$LOG_DIR/gpt_oss.log"
HOST_RUN_DIR="$LOG_DIR/host"
mkdir -p "$HOST_RUN_DIR"
echo " log: $GPT_LOG"
echo " host dir: $HOST_RUN_DIR"
cat > "$HOST_RUN_DIR/run_notebook.sh" <<'INNER'
#!/bin/bash
set -e
cd /workspace/host
echo "=== install triton_kernels (MXFP4 support for unsloth/gpt-oss-20b) ==="
pip install -q 'git+https://github.com/triton-lang/triton.git@0add68262ab0a2e33b84524346cb27cbb2787356#subdirectory=python/triton_kernels' 2>&1 | tail -5
echo
echo "=== fetch + convert notebook ==="
pip install -q nbconvert
curl -fsSL 'https://raw.githubusercontent.com/unslothai/notebooks/main/nb/gpt-oss-(20B)-Fine-tuning.ipynb' -o nb.ipynb
jupyter nbconvert --to script nb.ipynb --output nb 2>/dev/null
echo " nb.py: $(wc -l < nb.py) lines"
echo
echo "=== patch nb.py: max_steps 30 -> 10, drop pre-train demo generations ==="
python - <<'PY'
import re
src = open('nb.py').read()
src = src.replace('max_steps = 30', 'max_steps = 10')
src = re.sub(
r'messages = \[\s*\{[\"\']role[\"\']: [\"\']user[\"\'], [\"\']content[\"\']: [\"\']Solve x\^5.*?\n_ = model\.generate.*?streamer = TextStreamer\(tokenizer\)\)\n',
'# (pre-train inference skipped)\n',
src, flags=re.DOTALL, count=3,
)
open('nb.py', 'w').write(src)
print(' patched. max_steps now:', re.search(r'max_steps = (\d+)', src).group(1))
PY
echo
echo "=== run gpt-oss-20B fine-tuning ==="
python -u nb.py
INNER
chmod +x "$HOST_RUN_DIR/run_notebook.sh"
docker run --rm \
--gpus all \
--ipc=host \
--ulimit memlock=-1 \
--ulimit stack=67108864 \
-v "$HOST_RUN_DIR:/workspace/host" \
-v "$HOME/.cache/huggingface:/workspace/.cache/huggingface" \
-e HF_TOKEN="${HF_TOKEN:-}" \
-e HF_HUB_ENABLE_HF_TRANSFER=1 \
"$TAG" \
bash /workspace/host/run_notebook.sh 2>&1 | tee "$GPT_LOG"
rc=${PIPESTATUS[0]}
if [[ $rc -ne 0 ]]; then
fail "gpt-oss-20B notebook exited $rc -- see $GPT_LOG"
fi
ok "gpt-oss-20B notebook completed"
fi
# ============================================================================
# Summary
# ============================================================================
banner "summary"
echo " image: $TAG"
echo " log dir: $LOG_DIR"
echo
echo " to paste back for PR validation:"
[[ $SKIP_BUILD -eq 0 ]] && echo " tail -40 $LOG_DIR/build.log"
echo " cat $LOG_DIR/smoke.log"
[[ $SKIP_NOTEBOOK -eq 0 ]] && echo " tail -100 $LOG_DIR/gpt_oss.log"
echo
ok "all blocks completed"