Docker notebook safety hardening and vLLM startup timeout fix
pip shim (docker/unsloth_pip_shim.py): - Drop protected packages named via a VCS/URL #egg=NAME fragment so git+... #egg=torch no longer reinstalls into the baked venv. - Filter constraint files (-c/--constraint) through the same protected package filter as requirement files, so a pinned torch/transformers in a constraint cannot downgrade the baked stack during resolution. - Recursively filter nested -r/-c includes and absolutise their paths so the filtered /tmp copy still resolves them and no protected spec deep in the include tree slips past the keep list. - Remove an unused subprocess import. Notebook environment: - Scope the transformers-request marker per kernel (UNSLOTH_NB_TF_MARKER keyed on the kernel connection-file id) so concurrent notebooks no longer read each other's pin. - Install the IPython startup hook under IPYTHONDIR (set via ENV) so it loads for any uid, including docker run --user, not just root. - unsloth_nb_content_sig.py: only treat a %%capture / %%bash cell as install boilerplate when it carries an install command, so substantive captured/bash cells are hashed and upstream changes are not skipped. - unsloth_run.py: clean up the temp dir used to materialise a downloaded notebook. - unsloth_sync_notebooks.sh: honor UNSLOTH_KEEP_DELETED_NOTEBOOKS across GitHub refreshes so a deleted notebook is not restored when upstream advances. install_python_stack.py: the --local unsloth-zoo overlay now honors UNSLOTH_ZOO_REF (default main), matching the install.sh overlay. synthetic.py: preserve the timeout=None unbounded vLLM startup wait instead of coercing it to 1200s.
This commit is contained in:
Daniel Han
parent
08a1bf6680
commit
034fbc9785
8 changed files with 199 additions and 24 deletions
|
|
@ -683,11 +683,20 @@ RUN set -eux \
|
|||
&& ln -sf /opt/unsloth-nb/unsloth_run.py /usr/local/bin/unsloth-run \
|
||||
&& ln -sf /opt/unsloth-nb/unsloth_sync_notebooks.sh /usr/local/bin/unsloth-sync-notebooks \
|
||||
&& ln -sf /opt/unsloth-nb/unsloth_nb_content_sig.py /usr/local/bin/unsloth-nb-content-sig \
|
||||
&& mkdir -p /root/.ipython/profile_default/startup \
|
||||
&& cp /opt/unsloth-nb/unsloth_ipython_startup.py /root/.ipython/profile_default/startup/00-unsloth-nb.py \
|
||||
&& mkdir -p /opt/unsloth-nb/ipython/profile_default/startup \
|
||||
&& cp /opt/unsloth-nb/unsloth_ipython_startup.py /opt/unsloth-nb/ipython/profile_default/startup/00-unsloth-nb.py \
|
||||
&& chmod -R a+rX /opt/unsloth-nb/ipython \
|
||||
&& /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}
|
||||
# Load the notebook startup hook (sidecar activation + %pip/%uv magic re-point)
|
||||
# for EVERY kernel, whatever uid runs it. IPYTHONDIR (inherited by any user via
|
||||
# ENV) points IPython at this shared profile, so the hook still loads when the
|
||||
# container is started with `--user <uid>` and $HOME is not /root -- unlike a
|
||||
# /root/.ipython startup dir, which only a root kernel reads. Kernel-writable
|
||||
# state (history.sqlite) still lands under each user's own path, so a read-only
|
||||
# profile dir is fine.
|
||||
ENV IPYTHONDIR=/opt/unsloth-nb/ipython
|
||||
|
||||
# Pre-clone unslothai/notebooks so JupyterLab opens with the notebooks already
|
||||
# present (no git clone or wget needed). Baked here as a READ-ONLY template
|
||||
|
|
|
|||
|
|
@ -13,6 +13,27 @@ try:
|
|||
# `!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"
|
||||
|
||||
# Scope the transformers-request marker to THIS kernel so two notebooks
|
||||
# running concurrently in the same container (each its own kernel process)
|
||||
# do not read each other's pin. The pip/uv shim runs as a child of this
|
||||
# kernel and inherits UNSLOTH_NB_TF_MARKER, so writer (shim) and reader
|
||||
# (unsloth_nb_compat pre_run_cell hook, same process tree) agree on the
|
||||
# path. Falls back to the shared default when unset (e.g. `unsloth-run`,
|
||||
# which drives a single notebook per process).
|
||||
if not os.environ.get("UNSLOTH_NB_TF_MARKER"):
|
||||
# A kernel id that is stable for the kernel's lifetime and unique per
|
||||
# kernel: the ipykernel connection file name, else the kernel PID.
|
||||
_kid = ""
|
||||
try:
|
||||
from ipykernel import get_connection_file # type: ignore
|
||||
|
||||
_kid = os.path.splitext(os.path.basename(get_connection_file()))[0]
|
||||
except Exception:
|
||||
_kid = ""
|
||||
_kid = _kid or ("pid-%d" % os.getpid())
|
||||
os.environ["UNSLOTH_NB_TF_MARKER"] = "/tmp/unsloth_nb/requested_transformers." + _kid
|
||||
|
||||
import unsloth_nb_compat
|
||||
|
||||
unsloth_nb_compat.register_ipython()
|
||||
|
|
|
|||
|
|
@ -48,15 +48,32 @@ def _text(cell):
|
|||
return src.replace("\r\n", "\n").replace("\r", "\n")
|
||||
|
||||
|
||||
# Package-manager command fragments that mark a cell as the generated install
|
||||
# cell rather than substantive tutorial code.
|
||||
_INSTALL_MARKERS = (
|
||||
"pip install",
|
||||
"pip3-autoremove",
|
||||
"uv pip install",
|
||||
"conda install",
|
||||
"apt-get install",
|
||||
"apt install",
|
||||
)
|
||||
|
||||
|
||||
def _is_install_code(cell):
|
||||
if cell.get("cell_type") != "code":
|
||||
return False
|
||||
t = _text(cell)
|
||||
low = t.lower()
|
||||
if "pip install" in low or "pip3-autoremove" in low:
|
||||
if any(m in low for m in _INSTALL_MARKERS):
|
||||
return True
|
||||
first = t.lstrip().split("\n", 1)[0].strip().lower()
|
||||
return first.startswith("%%capture") or first.startswith("%%bash")
|
||||
# A %%capture / %%bash cell is boilerplate ONLY when it also carries an
|
||||
# install command. A bare %%capture (e.g. wrapping training to silence
|
||||
# output) or a %%bash cell doing real tutorial setup is substantive: hashing
|
||||
# it keeps the boot refresh from silently skipping an upstream fix to that
|
||||
# cell (a false SAME). The install markers above already catch the generated
|
||||
# install cell, which begins with %%capture.
|
||||
return False
|
||||
|
||||
|
||||
def _is_boilerplate_md(cell):
|
||||
|
|
|
|||
|
|
@ -21,7 +21,7 @@ are not intercepted -- the driven `unsloth-run` handles those by parsing the
|
|||
notebook directly.
|
||||
"""
|
||||
|
||||
import os, re, sys, subprocess, tempfile
|
||||
import os, re, sys, tempfile
|
||||
|
||||
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")
|
||||
|
|
@ -75,6 +75,11 @@ _VALUE_FLAGS = {
|
|||
# requirements file pulls real requirements. An index-url / find-links /
|
||||
# constraint / target value is an option, not something to install.
|
||||
_REQ_FILE_FLAGS = {"-r", "--requirement"}
|
||||
# Constraint files are not install targets, but pip applies their pins during
|
||||
# resolution, so a `-c constraints.txt` that pins torch/transformers/etc. can
|
||||
# still downgrade or reinstall a baked package when another target pulls it in.
|
||||
# Filter protected packages out of them the same way as requirement files.
|
||||
_CONSTRAINT_FILE_FLAGS = {"-c", "--constraint"}
|
||||
|
||||
|
||||
def _canon(token):
|
||||
|
|
@ -95,6 +100,15 @@ def _canon(token):
|
|||
if _dref:
|
||||
return _dref.group(1).lower().replace("_", "-") or None
|
||||
if re.match(r"^[a-z]+\+", token) or "://" in token or token.startswith((".", "/")):
|
||||
# A VCS / URL install can still name a protected package via the legacy
|
||||
# `#egg=NAME` (or `&egg=NAME`) fragment, e.g.
|
||||
# `git+https://github.com/unslothai/unsloth.git#egg=unsloth`. Pull that
|
||||
# name out so _KEEP can drop it; otherwise the shim would exec the URL
|
||||
# and reinstall a baked package into the venv. A non-protected egg name
|
||||
# is returned too, but the caller keeps it as a normal target either way.
|
||||
_egg = re.search(r"[#&]egg=([A-Za-z0-9][A-Za-z0-9._-]*)", token)
|
||||
if _egg:
|
||||
return _egg.group(1).lower().replace("_", "-") or None
|
||||
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()
|
||||
|
|
@ -107,7 +121,66 @@ def _version_pin(token):
|
|||
return m.group(1) if m else None
|
||||
|
||||
|
||||
def _filter_requirements_file(path):
|
||||
def _parse_include(stripped):
|
||||
"""If `stripped` is an `-r`/`--requirement`/`-c`/`--constraint` include,
|
||||
return (flag, target_path, inline_comment_or_None); else (None, None, None)."""
|
||||
body, sep, comment = stripped.partition(" #")
|
||||
body = body.rstrip()
|
||||
comment = ("#" + comment) if sep else None
|
||||
for flag in ("-r", "--requirement", "-c", "--constraint"):
|
||||
target = None
|
||||
if body == flag or body.startswith(flag + " "):
|
||||
target = body[len(flag):].strip()
|
||||
elif body.startswith(flag + "="):
|
||||
target = body[len(flag) + 1:].strip()
|
||||
elif not flag.startswith("--") and body.startswith(flag) and len(body) > len(flag):
|
||||
target = body[len(flag):].strip() # attached short form, e.g. `-rextras.txt`
|
||||
else:
|
||||
continue
|
||||
return flag, (target or None), comment
|
||||
return None, None, None
|
||||
|
||||
|
||||
def _rewrite_include(line, stripped, src_dir, depth):
|
||||
"""Rewrite a nested `-r`/`-c` include so pip still resolves it and its
|
||||
protected specs are filtered too.
|
||||
|
||||
pip resolves a nested include against the directory of the file it is
|
||||
READING; our filtered copy lives under /tmp, so a relative include would
|
||||
look in /tmp and fail. Recursively filter the included file (dropping
|
||||
protected packages there too, closing the multi-level bypass) and point the
|
||||
parent at that filtered copy. URLs and unreadable/absolute-unfiltered files
|
||||
fall back to an absolutised path so they still resolve. Returns
|
||||
(new_line, changed, recorded, dropped)."""
|
||||
flag, target, comment = _parse_include(stripped)
|
||||
if not target:
|
||||
return line, False, None, []
|
||||
newline_char = "\n" if line.endswith("\n") else ""
|
||||
|
||||
def _emit(new_target):
|
||||
rebuilt = flag + " " + new_target
|
||||
if comment:
|
||||
rebuilt += " " + comment
|
||||
return rebuilt + newline_char
|
||||
|
||||
# A URL include cannot be filtered locally; leave it verbatim.
|
||||
if "://" in target:
|
||||
return line, False, None, []
|
||||
abs_target = target if os.path.isabs(target) else os.path.join(src_dir, target)
|
||||
# Recursively filter the included file. Guard against cyclic / deep includes.
|
||||
if depth < 8:
|
||||
f_path, f_rec, f_drp = _filter_requirements_file(abs_target, _depth = depth + 1)
|
||||
if f_path != abs_target:
|
||||
# The include was rewritten (protected specs dropped and/or its own
|
||||
# nested includes absolutised); point at the filtered copy.
|
||||
return _emit(f_path), True, f_rec, f_drp
|
||||
# Nothing to filter inside; just make sure the path still resolves from /tmp.
|
||||
if not os.path.isabs(target):
|
||||
return _emit(abs_target), True, None, []
|
||||
return line, False, None, []
|
||||
|
||||
|
||||
def _filter_requirements_file(path, _depth = 0):
|
||||
"""Strip baked/protected packages out of a `-r` requirements file.
|
||||
|
||||
Returns (path_to_use, recorded_transformers_version, dropped_specs). The same
|
||||
|
|
@ -115,19 +188,35 @@ def _filter_requirements_file(path):
|
|||
line, so a notebook `pip install -r reqs.txt` cannot overwrite the cu128 torch
|
||||
/ vLLM / transformers stack with versions pinned inside the file. When nothing
|
||||
is protected, or the file cannot be read/written, the original path is returned
|
||||
unchanged. Comments, blank lines, option lines and nested `-r`/`-c` includes are
|
||||
kept verbatim (nested includes are passed through, i.e. filtered one level).
|
||||
unchanged. Comments, blank lines and option lines are kept verbatim; a nested
|
||||
`-r`/`-c` include is recursively filtered too (protected specs dropped at every
|
||||
level).
|
||||
"""
|
||||
try:
|
||||
with open(path, encoding = "utf-8") as f:
|
||||
lines = f.readlines()
|
||||
except OSError:
|
||||
return path, None, [] # remote URL / unreadable -> let the real tool handle it
|
||||
src_dir = os.path.dirname(os.path.abspath(path))
|
||||
out, dropped, recorded, changed = [], [], None, False
|
||||
for line in lines:
|
||||
stripped = line.strip()
|
||||
if not stripped or stripped.startswith(("#", "-")):
|
||||
out.append(line) # comment / blank / option / nested include -> keep
|
||||
if not stripped or stripped.startswith("#"):
|
||||
out.append(line) # comment / blank -> keep
|
||||
continue
|
||||
if stripped.startswith("-"):
|
||||
# Option or nested include. Recursively filter a nested `-r`/`-c`
|
||||
# include (so protected specs deep in the include tree cannot slip
|
||||
# past _KEEP) and repoint it so it still resolves from /tmp.
|
||||
new_line, rewrote, inc_rec, inc_drp = _rewrite_include(
|
||||
line, stripped, src_dir, _depth
|
||||
)
|
||||
out.append(new_line)
|
||||
if rewrote:
|
||||
changed = True
|
||||
if inc_rec and not recorded:
|
||||
recorded = inc_rec
|
||||
dropped.extend(inc_drp)
|
||||
continue
|
||||
spec = stripped.split(" #", 1)[0].strip() # drop any inline comment
|
||||
name = _canon(spec)
|
||||
|
|
@ -200,6 +289,14 @@ def main():
|
|||
if _req_rec and not recorded:
|
||||
recorded = _req_rec
|
||||
dropped.extend(_req_drp)
|
||||
elif prev_flag in _CONSTRAINT_FILE_FLAGS:
|
||||
# Strip protected pins from the constraint file so it cannot
|
||||
# downgrade the baked stack, but a constraint is not an install
|
||||
# target and its transformers pin is not an install request, so
|
||||
# do not set has_target / recorded here.
|
||||
_c_path, _c_rec, _c_drp = _filter_requirements_file(tok)
|
||||
keep_args.append(_c_path)
|
||||
dropped.extend(_c_drp)
|
||||
else:
|
||||
keep_args.append(tok)
|
||||
skip_next = False
|
||||
|
|
@ -220,6 +317,10 @@ def main():
|
|||
if _req_rec and not recorded:
|
||||
recorded = _req_rec
|
||||
dropped.extend(_req_drp)
|
||||
elif _flag in _CONSTRAINT_FILE_FLAGS:
|
||||
_c_path, _c_rec, _c_drp = _filter_requirements_file(_val)
|
||||
keep_args.append(_flag + "=" + _c_path)
|
||||
dropped.extend(_c_drp)
|
||||
else:
|
||||
keep_args.append(tok) # option with inline value, not a target
|
||||
continue
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@ Usage:
|
|||
A raw github URL (raw.githubusercontent.com/.../nb/Foo.ipynb) is fetched first.
|
||||
"""
|
||||
|
||||
import argparse, json, os, re, subprocess, sys, tempfile, urllib.request
|
||||
import argparse, json, os, re, shutil, subprocess, sys, tempfile, urllib.request
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
try:
|
||||
|
|
@ -69,10 +69,13 @@ def main():
|
|||
sidecar = compat.sidecar_for(want) if (compat and want) else None
|
||||
|
||||
# Materialise the notebook locally for nbconvert.
|
||||
tmp_dir = None
|
||||
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])
|
||||
)
|
||||
if args.out:
|
||||
src_path = args.out
|
||||
else:
|
||||
tmp_dir = tempfile.mkdtemp()
|
||||
src_path = os.path.join(tmp_dir, os.path.basename(args.notebook.split("?")[0]))
|
||||
with open(src_path, "w") as f:
|
||||
json.dump(nb, f)
|
||||
else:
|
||||
|
|
@ -109,7 +112,13 @@ def main():
|
|||
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))
|
||||
try:
|
||||
rc = subprocess.call(cmd, env = env)
|
||||
finally:
|
||||
# Clean up the temp dir we materialised a downloaded notebook into.
|
||||
if tmp_dir is not None:
|
||||
shutil.rmtree(tmp_dir, ignore_errors = True)
|
||||
sys.exit(rc)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
|
|
@ -183,6 +183,15 @@ while IFS= read -r -d '' f; do
|
|||
unchanged=$((unchanged + 1))
|
||||
continue
|
||||
fi
|
||||
elif [ -n "${LAST[$rel]:-}" ] && [ "${UNSLOTH_KEEP_DELETED_NOTEBOOKS:-0}" = "1" ]; then
|
||||
# We previously wrote this notebook and the user has since DELETED it.
|
||||
# With the opt-out set, honor the deletion instead of restoring it from
|
||||
# the fresh clone when upstream advances (otherwise the deletion only
|
||||
# held until the next remote refresh). Keep the record so it stays known
|
||||
# as managed-but-deleted.
|
||||
printf '%s %s\n' "${LAST[$rel]}" "$rel" >> "$TMPSTATE"
|
||||
kept=$((kept + 1))
|
||||
continue
|
||||
fi
|
||||
mkdir -p "$(dirname "$dst")" 2>/dev/null || true
|
||||
if cp -a "$f" "$dst" 2>/dev/null; then
|
||||
|
|
|
|||
|
|
@ -2045,6 +2045,12 @@ def install_python_stack() -> int:
|
|||
package_name = os.environ.get("STUDIO_PACKAGE_NAME", "unsloth")
|
||||
# --local overlays a local repo checkout after updating deps.
|
||||
local_repo = os.environ.get("STUDIO_LOCAL_REPO", "")
|
||||
# unsloth-zoo git ref for the --local overlay. Honor UNSLOTH_ZOO_REF (the
|
||||
# Docker publish workflow / unsloth-studio-update resolve one ref and forward
|
||||
# it) so the Studio venv can track the operator-requested zoo instead of
|
||||
# always main. Unset -> main, byte-identical to the previous bare git URL.
|
||||
zoo_ref = os.environ.get("UNSLOTH_ZOO_REF", "").strip() or "main"
|
||||
zoo_git_spec = "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo@" + zoo_ref
|
||||
base_total = 11 if IS_WINDOWS else 12 # +1 for the anyio repair check (step 8b)
|
||||
if IS_MACOS:
|
||||
base_total -= 1 # triton step is skipped on macOS
|
||||
|
|
@ -2154,13 +2160,13 @@ def install_python_stack() -> int:
|
|||
local_repo,
|
||||
constrain = False,
|
||||
)
|
||||
_step(_LABEL, "overlaying unsloth-zoo from git main")
|
||||
_step(_LABEL, f"overlaying unsloth-zoo from git {zoo_ref}")
|
||||
pip_install(
|
||||
"Overlaying unsloth-zoo from git main",
|
||||
f"Overlaying unsloth-zoo from git {zoo_ref}",
|
||||
"--no-cache-dir",
|
||||
"--no-deps",
|
||||
"--force-reinstall",
|
||||
"unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo",
|
||||
zoo_git_spec,
|
||||
constrain = False,
|
||||
)
|
||||
elif local_repo:
|
||||
|
|
@ -2185,13 +2191,13 @@ def install_python_stack() -> int:
|
|||
local_repo,
|
||||
constrain = False,
|
||||
)
|
||||
_step(_LABEL, "overlaying unsloth-zoo from git main")
|
||||
_step(_LABEL, f"overlaying unsloth-zoo from git {zoo_ref}")
|
||||
pip_install(
|
||||
"Overlaying unsloth-zoo from git main",
|
||||
f"Overlaying unsloth-zoo from git {zoo_ref}",
|
||||
"--no-cache-dir",
|
||||
"--no-deps",
|
||||
"--force-reinstall",
|
||||
"unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo",
|
||||
zoo_git_spec,
|
||||
constrain = False,
|
||||
)
|
||||
elif package_name != "unsloth":
|
||||
|
|
|
|||
|
|
@ -295,8 +295,11 @@ class SyntheticDataKit:
|
|||
# we don't print stderr to console but self.stderr_capture.tail(200) will print the last 200 lines
|
||||
|
||||
ready = False
|
||||
deadline = time.monotonic() + (timeout or 1200)
|
||||
while time.monotonic() < deadline:
|
||||
# timeout = None (or 0) preserves the previous Event.wait(None) escape
|
||||
# hatch: wait indefinitely for the readiness message (useful for large
|
||||
# models or slow first-time downloads). Any positive value is a deadline.
|
||||
deadline = (time.monotonic() + timeout) if timeout else None
|
||||
while deadline is None or time.monotonic() < deadline:
|
||||
if self.stdout_capture.wait_for_ready(timeout = 1) or self.stderr_capture.wait_for_ready(
|
||||
timeout = 0
|
||||
):
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue