diff --git a/docker/.dockerignore b/docker/.dockerignore index ae1f499a29..1bd005c9c1 100644 --- a/docker/.dockerignore +++ b/docker/.dockerignore @@ -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 diff --git a/docker/Dockerfile b/docker/Dockerfile index 2d577d0559..fd3a8d1cdc 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -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 ` 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 diff --git a/docker/Dockerfile.studio b/docker/Dockerfile.studio index 8fb488c9e1..89609ef6fc 100644 --- a/docker/Dockerfile.studio +++ b/docker/Dockerfile.studio @@ -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 diff --git a/docker/unsloth_ipython_startup.py b/docker/unsloth_ipython_startup.py new file mode 100644 index 0000000000..fcc9f8c472 --- /dev/null +++ b/docker/unsloth_ipython_startup.py @@ -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) diff --git a/docker/unsloth_nb_compat.py b/docker/unsloth_nb_compat.py new file mode 100644 index 0000000000..70a905b757 --- /dev/null +++ b/docker/unsloth_nb_compat.py @@ -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 --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 ` 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 diff --git a/docker/unsloth_pip_shim.py b/docker/unsloth_pip_shim.py new file mode 100644 index 0000000000..4005141167 --- /dev/null +++ b/docker/unsloth_pip_shim.py @@ -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() diff --git a/docker/unsloth_run.py b/docker/unsloth_run.py new file mode 100644 index 0000000000..6cf2309eef --- /dev/null +++ b/docker/unsloth_run.py @@ -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 [--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()