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 2026-06-15 03:13:16 +00:00
commit aba16af123
7 changed files with 528 additions and 5 deletions

View 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)