docker: notebook deps, image size cuts, per-notebook transformers
Notebook dependency coverage (base Dockerfile): - Bake omegaconf, einx, librosa, decord, ftfy so the TTS/STT and vision notebooks stop dying on a silent No module named X. Installed in the notebook-deps layer (after the torch/vLLM resolve) with an assertion that the resolve did not move torch 2.10.0 / numpy>=2.3 / numba>=0.65. Image size (no functional change): - Base: prune npp to the two libs torchcodec actually dlopens (libnppicc + libnppc), drop link-time-only .a archives and the nvshmem device bitcode. Headers (torch/include etc) are kept so causal-conv1d / mamba-ssm still build at notebook time with --no-build-isolation. - Studio: pin the Studio venv to Python 3.12 (matches base) so its nvidia-*-cu12 wheels are byte-identical to the base venv's, then symlink the heavy arch-independent CUDA libs (cudnn/cublas/nccl/...) into the base venv copy. cuda_nvrtc and cuda_runtime are excluded (the arm64 nvrtc swap mutates nvrtc in place). Also remove the build-only frontend node_modules (runtime serves the committed dist). Studio image drops ~4.8GB. Per-notebook transformers version, run notebooks unchanged: - Bake coherent transformers sidecars (4.57.6 default + 5.3.0/5.5.0/5.10.2), each transformers==X with its matched huggingface_hub/tokenizers/ safetensors installed --no-deps into its own dir. Companion versions are resolved at build time so they satisfy each transformers' requirements. - unsloth_nb_compat.py: pick the sidecar from the notebook's pin or the model name and activate it (prepend to sys.path) before any ML import, without touching the base cu128 torch/vLLM/unsloth stack. - pip/uv shim on PATH: a notebook install cell becomes safe and idempotent inside a kernel (keeps the baked stack, records the requested transformers for its sidecar); passthrough to the real tool everywhere else. - IPython startup hook for manual JupyterLab, and unsloth-run for the headless driven path.
This commit is contained in:
Daniel Han
parent
dec240ade7
commit
aba16af123
7 changed files with 528 additions and 5 deletions
|
|
@ -5,3 +5,7 @@
|
|||
!fetch_llama_prebuilt.py
|
||||
!supervisord.conf
|
||||
!studio_launch.sh
|
||||
!unsloth_nb_compat.py
|
||||
!unsloth_pip_shim.py
|
||||
!unsloth_ipython_startup.py
|
||||
!unsloth_run.py
|
||||
|
|
|
|||
|
|
@ -282,10 +282,20 @@ RUN set -eux \
|
|||
# langid DeepSeek-R1 GRPO reward's language-id check
|
||||
# easydict some vision trust_remote_code modeling files
|
||||
# protobuf slow->fast tokenizer conversion for sentencepiece models
|
||||
# omegaconf TTS families + both NeMo-Gym RL notebooks' config objects
|
||||
# einx TTS codec tensor-rearrange (Llasa / Oute / Spark TTS)
|
||||
# librosa Whisper audio feature extraction (pairs with soundfile + torchcodec)
|
||||
# decord ERNIE-VL vision notebook video decode
|
||||
# ftfy Oute TTS text normalisation
|
||||
# librosa pulls numba/soxr/audioread; numba is already pinned >=0.65 (numpy 2.4
|
||||
# compatible) by the vLLM pass, so the resolve must NOT move torch/numpy/numba --
|
||||
# the assertion below fails the build loudly if it did.
|
||||
RUN ${VENV}/bin/uv pip install \
|
||||
--python ${VENV}/bin/python \
|
||||
jupyterlab notebook ipywidgets matplotlib \
|
||||
soundfile evaluate jiwer tensorboard langid easydict protobuf
|
||||
soundfile evaluate jiwer tensorboard langid easydict protobuf \
|
||||
omegaconf einx librosa decord ftfy \
|
||||
&& ${VENV}/bin/python -c "import torch, numpy, numba; from packaging.version import Version; assert torch.__version__.startswith('2.10.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__)"
|
||||
|
||||
# Audio decode out of the box: the TTS/STT notebooks feed datasets' Audio
|
||||
# features, which decode through torchcodec. Three traps, all defended:
|
||||
|
|
@ -307,6 +317,40 @@ RUN set -eux \
|
|||
&& ${VENV}/bin/uv pip install --python ${VENV}/bin/python nvidia-npp-cu12; } \
|
||||
|| echo ">> torchcodec bake skipped (no matching wheel for ${TARGETARCH:-amd64})"
|
||||
|
||||
# Coherent transformers SIDECARS for per-notebook version activation (see
|
||||
# docker/unsloth_nb_compat.py). unslothai/notebooks pin many transformers
|
||||
# versions in their install cells; the base venv ships one (newest 5.x). Each
|
||||
# sidecar is `transformers==X` + its matched huggingface_hub/tokenizers/
|
||||
# safetensors, installed --no-deps into its own --target dir under
|
||||
# ${VENV}/tf-sidecars (rides along in the COPY to runtime). Activating one
|
||||
# (prepend to sys.path before any ML import) swaps transformers for a model
|
||||
# WITHOUT touching the cu128 torch/vLLM/unsloth base stack -- verified: base
|
||||
# unsloth loads + generates under both a 4.57.6 and a 5.5.0 sidecar on B200.
|
||||
# Versions mirror Unsloth Studio's tiers (4.57.6 default + 5.3.0/5.5.0/5.10.2).
|
||||
# The companion versions are RESOLVED at build time (not hardcoded) so they
|
||||
# always satisfy each transformers' hard requirements. ~300MB total after the
|
||||
# __pycache__/tests strip in the cleanup RUN below. Fail-soft per arch/wheel.
|
||||
RUN set -eux \
|
||||
&& 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}"}; \
|
||||
echo ">> sidecar transformers==${TFV} (hf_hub=${HFV} tokenizers=${TKV} safetensors=${SFV})"; \
|
||||
done \
|
||||
&& { du -sh ${VENV}/tf-sidecars || true; }
|
||||
|
||||
# 5) Emit an informational pin record so downstream consumers can see exactly
|
||||
# what was resolved. This is NOT a byte-reproducible lockfile -- `pip freeze`
|
||||
# captures version strings but not wheel hashes, and several deps (unsloth,
|
||||
|
|
@ -325,13 +369,33 @@ RUN ${VENV}/bin/pip freeze --exclude-editable > ${VENV}/requirements.lock.txt \
|
|||
# same `from numpy._core.tests._natype import pd_NA` ImportError on the
|
||||
# deployed image. Exclude numpy's tests directories explicitly so the
|
||||
# upgrade fix stays in effect; keep stripping the rest.
|
||||
RUN find ${VENV} -depth -type d -name __pycache__ -exec rm -rf {} + \
|
||||
# Also (size reductions, all verified runtime-safe):
|
||||
# * npp: torchcodec dlopens only libnppicc + libnppc; the other ~10 npp libs
|
||||
# (libnppif 163M, libnppist, libnppig, ...) are dead weight (~388MB). Nothing
|
||||
# else in the venv links them.
|
||||
# * static .a archives (~143MB): xgrammar/triton-cupti/nvperf/nvshmem ship .a
|
||||
# alongside the .so they actually load at runtime; .a are link-time only and
|
||||
# nothing in the image links venv archives (nvcc JIT links /usr/local/cuda).
|
||||
# * nvshmem device-side bitcode (~30MB): host .so kept; .bc is device-relink only.
|
||||
# We deliberately do NOT strip headers (torch/include etc.): the leave-to-pip
|
||||
# kernels causal-conv1d / mamba-ssm build against torch headers at notebook time
|
||||
# with --no-build-isolation, so torch/include must survive.
|
||||
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
|
||||
&& 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.
|
||||
#
|
||||
|
|
@ -564,6 +628,34 @@ ENV UNSLOTH_LLAMA_CPP_PATH=/opt/unsloth/llama.cpp
|
|||
WORKDIR /workspace
|
||||
RUN mkdir -p ${HF_HOME} ${TRITON_CACHE_DIR}
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Per-notebook transformers version activation -- run unslothai/notebooks
|
||||
# UNCHANGED. See docker/unsloth_nb_compat.py for the full rationale. Pieces:
|
||||
# * unsloth_nb_compat.py -> site-packages (importable everywhere): tier
|
||||
# detection + sidecar resolution + activation + the IPython hook.
|
||||
# * pip/uv shim on a PATH dir AHEAD of the venv bin: a notebook's
|
||||
# `!pip install ...` / `!uv pip install ...` cell becomes SAFE + idempotent
|
||||
# (keeps the baked torch/vLLM stack; records the requested transformers so
|
||||
# its sidecar is activated for the model cells).
|
||||
# * IPython startup hook: activates the right sidecar before the first model
|
||||
# cell in manual JupyterLab.
|
||||
# * unsloth-run: headless `unsloth-run <notebook|url>` that auto-picks the
|
||||
# sidecar and executes every cell -- the robust driven path.
|
||||
# ---------------------------------------------------------------------------
|
||||
COPY unsloth_nb_compat.py unsloth_pip_shim.py unsloth_ipython_startup.py unsloth_run.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" \
|
||||
&& chmod +x /opt/unsloth-nb/unsloth_pip_shim.py /opt/unsloth-nb/unsloth_run.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 \
|
||||
&& mkdir -p /root/.ipython/profile_default/startup \
|
||||
&& cp /opt/unsloth-nb/unsloth_ipython_startup.py /root/.ipython/profile_default/startup/00-unsloth-nb.py \
|
||||
&& /opt/unsloth-venv/bin/python -c "import sys, glob; sys.path.insert(0, '$SP'); import unsloth_nb_compat; print('nb-compat OK; baked sidecars:', sorted(glob.glob('/opt/unsloth-venv/tf-sidecars/t_*')))"
|
||||
# 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}
|
||||
|
||||
# JupyterLab lives in the venv (see builder stage). Persistent notebooks
|
||||
# should be bind-mounted onto /workspace.
|
||||
EXPOSE 8888
|
||||
|
|
|
|||
|
|
@ -82,6 +82,12 @@ RUN apt-get update \
|
|||
# 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).
|
||||
#
|
||||
# UNSLOTH_PYTHON=3.12 pins the Studio venv to the SAME Python minor as the base
|
||||
# venv (install.sh defaults Linux to 3.13). Matching minors makes the two venvs'
|
||||
# nvidia-*-cu12 CUDA wheels byte-identical, which lets the dedup RUN further down
|
||||
# replace the Studio venv's ~3.7GB of CUDA .so with symlinks into the base venv's
|
||||
# copies (cudnn/cublas/nccl/... are plain C libs, Python-minor independent).
|
||||
#
|
||||
# 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).
|
||||
|
|
@ -100,20 +106,23 @@ RUN set -eux \
|
|||
&& git checkout -q FETCH_HEAD \
|
||||
&& UNSLOTH_STUDIO_HOME="${UNSLOTH_STUDIO_HOME}" \
|
||||
UNSLOTH_TORCH_INDEX_FAMILY="${TORCH_FAMILY}" \
|
||||
UNSLOTH_PYTHON=3.12 \
|
||||
bash install.sh --local \
|
||||
# Fail loud if the Studio venv torch missed the pinned CUDA family (an
|
||||
# install.sh that ignores UNSLOTH_TORCH_INDEX_FAMILY falls back to
|
||||
# nvidia-smi probing, which cannot work at build time and lands on cu126
|
||||
# wheels with no sm_100/sm_120 kernels). metadata check only: importing
|
||||
# torch needs native libs, which QEMU arm64 builds cannot load.
|
||||
&& "${UNSLOTH_STUDIO_HOME}/unsloth_studio/bin/python" -c "from importlib.metadata import version; v = version('torch'); assert v.endswith('+${TORCH_FAMILY}'), 'Studio venv torch ' + v + ' does not match ${TORCH_FAMILY}'; print('Studio venv torch', v)" \
|
||||
&& "${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.endswith('+${TORCH_FAMILY}'), 'Studio venv torch ' + v + ' does not match ${TORCH_FAMILY}'; print('Studio venv python %d.%d' % sys.version_info[:2], 'torch', v)" \
|
||||
# setup.sh may relink the root llama-quantize into build/bin; prove the
|
||||
# relinked quantizer still resolves its libraries, or GGUF export breaks
|
||||
# at runtime with "No working quantizer found". Content check, not rc:
|
||||
# llama-quantize exits nonzero on --help, while a loader failure prints
|
||||
# "error while loading shared libraries" and no usage text.
|
||||
&& { "${UNSLOTH_STUDIO_HOME}/llama.cpp/llama-quantize" --help 2>&1 || true; } | grep -q "usage" \
|
||||
&& rm -rf "${UNSLOTH_STUDIO_HOME}/src/.git" /root/.cache \
|
||||
&& rm -rf "${UNSLOTH_STUDIO_HOME}/src/.git" \
|
||||
"${UNSLOTH_STUDIO_HOME}/src/studio/frontend/node_modules" \
|
||||
/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 \
|
||||
|
|
@ -121,6 +130,28 @@ RUN set -eux \
|
|||
ln -s /usr/local/cuda-13.0/lib64/libnvrtc.so.13 "${NVRTC_DIR}/libnvrtc.so.12"; \
|
||||
fi; \
|
||||
done; \
|
||||
fi \
|
||||
&& 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
|
||||
|
|
|
|||
18
docker/unsloth_ipython_startup.py
Normal file
18
docker/unsloth_ipython_startup.py
Normal file
|
|
@ -0,0 +1,18 @@
|
|||
"""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 running inside a notebook kernel, so a cell's
|
||||
# `!pip install ...` / `!uv pip install ...` (which inherits this env) gets
|
||||
# the safe-install behaviour. Unset everywhere else => shim is a passthrough.
|
||||
os.environ["UNSLOTH_NB_SHIM"] = "1"
|
||||
import unsloth_nb_compat
|
||||
unsloth_nb_compat.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)
|
||||
135
docker/unsloth_nb_compat.py
Normal file
135
docker/unsloth_nb_compat.py
Normal file
|
|
@ -0,0 +1,135 @@
|
|||
"""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 <dir> --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. Verified:
|
||||
base unsloth loads + generates under both a 4.57.6 and a 5.5.0 sidecar on B200.
|
||||
|
||||
Two activation paths:
|
||||
* driven/headless: `unsloth-run <notebook>` 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")
|
||||
|
||||
# Model-name -> minimum transformers tier, ported from Studio's
|
||||
# transformers_version.py (substring match on the lowered model id). Used as a
|
||||
# fallback when a notebook does not pin transformers but names a new-family model.
|
||||
_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 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.
|
||||
|
||||
Uses ceiling semantics: the smallest baked version >= the request, because a
|
||||
model added in version X needs *at least* X. If the request is newer than
|
||||
every baked sidecar, return None -> use the base venv (the newest 5.x)."""
|
||||
baked = _baked()
|
||||
if not baked or not version:
|
||||
return None
|
||||
if version in baked:
|
||||
return baked[version]
|
||||
try:
|
||||
from packaging.version import Version
|
||||
want = Version(version)
|
||||
except Exception:
|
||||
return None
|
||||
ge = sorted((Version(v), d) for v, d in baked.items() if Version(v) >= want)
|
||||
return ge[0][1] if ge else 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:
|
||||
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
|
||||
136
docker/unsloth_pip_shim.py
Normal file
136
docker/unsloth_pip_shim.py
Normal file
|
|
@ -0,0 +1,136 @@
|
|||
#!/opt/unsloth-venv/bin/python
|
||||
"""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 / triton / xformers / vllm / bitsandbytes /
|
||||
flashinfer / nvidia-* -> SKIPPED (the baked, ABI-matched versions are kept;
|
||||
a notebook reinstall here only ever breaks the GPU stack).
|
||||
* 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, subprocess
|
||||
|
||||
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.
|
||||
_KEEP = {
|
||||
"torch", "torchvision", "torchaudio", "triton", "triton-rocm", "pytorch-triton",
|
||||
"xformers", "vllm", "bitsandbytes", "flashinfer", "flashinfer-python",
|
||||
"unsloth", "unsloth-zoo", "unsloth_zoo",
|
||||
}
|
||||
_KEEP_PREFIX = ("nvidia-", "nvidia_")
|
||||
# pip/uv flags that consume the following token as a value (so we don't mistake
|
||||
# that value for a requirement).
|
||||
_VALUE_FLAGS = {
|
||||
"-r", "--requirement", "-c", "--constraint", "-i", "--index-url",
|
||||
"--extra-index-url", "-f", "--find-links", "--target", "-t", "--python", "-p",
|
||||
"--prefix", "--index-strategy", "--upgrade-package", "-P", "--no-binary",
|
||||
"--only-binary", "--platform", "--python-version", "--abi", "--implementation",
|
||||
}
|
||||
|
||||
|
||||
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
|
||||
if re.match(r"^[a-z]+\+", token) or "://" in token or token.startswith((".", "/")):
|
||||
return None # vcs / url / local path -> let it pass through
|
||||
# strip extras and any version/marker tail
|
||||
name = re.split(r"[<>=!~\[\s;@]", token, 1)[0].strip()
|
||||
return name.lower().replace("_", "-") or 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
|
||||
|
||||
|
||||
def main():
|
||||
tool = "uv" if os.path.basename(sys.argv[0]).startswith("uv") else "pip"
|
||||
argv = sys.argv[1:]
|
||||
|
||||
# Only intercept inside a notebook kernel (UNSLOTH_NB_SHIM is set by the baked
|
||||
# IPython startup and by `unsloth-run`). EVERYWHERE else -- install.sh during
|
||||
# the image build, internal tooling, an interactive shell -- behave exactly
|
||||
# like the real tool, so we never disturb the build or system package mgmt.
|
||||
if os.environ.get("UNSLOTH_NB_SHIM") != "1":
|
||||
os.execv(REAL[tool], [REAL[tool]] + argv)
|
||||
return
|
||||
|
||||
# Locate the `install` verb (uv: `uv pip install ...`; pip: `pip install ...`).
|
||||
try:
|
||||
if tool == "uv":
|
||||
# skip a leading `pip` subcommand
|
||||
i = argv.index("install")
|
||||
else:
|
||||
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
|
||||
skip_next = False
|
||||
for tok in tail:
|
||||
if skip_next:
|
||||
keep_args.append(tok)
|
||||
skip_next = False
|
||||
continue
|
||||
if tok in _VALUE_FLAGS:
|
||||
keep_args.append(tok)
|
||||
skip_next = True
|
||||
continue
|
||||
name = _canon(tok)
|
||||
if name is None:
|
||||
keep_args.append(tok) # flag / url / path
|
||||
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)
|
||||
|
||||
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 actually install? (a requirement, not just flags)
|
||||
real_reqs = [t for t in keep_args if _canon(t)]
|
||||
if not real_reqs:
|
||||
print("[unsloth-nb] nothing to install after keeping the baked stack; ok.")
|
||||
return
|
||||
cmd = [REAL[tool]] + head + keep_args
|
||||
sys.stdout.flush()
|
||||
os.execv(REAL[tool], cmd)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
107
docker/unsloth_run.py
Normal file
107
docker/unsloth_run.py
Normal file
|
|
@ -0,0 +1,107 @@
|
|||
#!/opt/unsloth-venv/bin/python
|
||||
"""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 <notebook.ipynb | URL> [--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, 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 locally for nbconvert.
|
||||
if args.notebook.startswith(("http://", "https://")) or args.out:
|
||||
src_path = args.out or os.path.join(
|
||||
tempfile.mkdtemp(), os.path.basename(args.notebook.split("?")[0]))
|
||||
with open(src_path, "w") as f:
|
||||
json.dump(nb, f)
|
||||
else:
|
||||
src_path = args.notebook
|
||||
out_path = args.out or src_path
|
||||
|
||||
env = dict(os.environ)
|
||||
env["UNSLOTH_NB_SHIM"] = "1" # enable safe-install for the notebook's cells
|
||||
# The pip/uv shim writes the marker; pre-seed it too so the kernel agrees.
|
||||
if want:
|
||||
marker = env.get("UNSLOTH_NB_TF_MARKER", "/tmp/unsloth_nb/requested_transformers")
|
||||
os.makedirs(os.path.dirname(marker), 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")
|
||||
|
||||
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(out_path),
|
||||
"--output-dir", os.path.dirname(os.path.abspath(out_path)) or ".",
|
||||
]
|
||||
print("[unsloth-run] executing:", os.path.basename(src_path))
|
||||
sys.exit(subprocess.call(cmd, env=env))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Loading…
Add table
Add a link
Reference in a new issue