diff --git a/docker/Dockerfile b/docker/Dockerfile index 33901f4b55..9a3fc7ee90 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -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 ` 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 diff --git a/docker/unsloth_ipython_startup.py b/docker/unsloth_ipython_startup.py index 3e07c9f890..accb8be781 100644 --- a/docker/unsloth_ipython_startup.py +++ b/docker/unsloth_ipython_startup.py @@ -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() diff --git a/docker/unsloth_nb_content_sig.py b/docker/unsloth_nb_content_sig.py index 6d1444342d..d3dfcb738d 100644 --- a/docker/unsloth_nb_content_sig.py +++ b/docker/unsloth_nb_content_sig.py @@ -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): diff --git a/docker/unsloth_pip_shim.py b/docker/unsloth_pip_shim.py index 080d071ba5..927c8bc791 100644 --- a/docker/unsloth_pip_shim.py +++ b/docker/unsloth_pip_shim.py @@ -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 diff --git a/docker/unsloth_run.py b/docker/unsloth_run.py index 7daa81e7d9..5bd0be7793 100644 --- a/docker/unsloth_run.py +++ b/docker/unsloth_run.py @@ -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__": diff --git a/docker/unsloth_sync_notebooks.sh b/docker/unsloth_sync_notebooks.sh index 2727e4427b..12c28f063d 100644 --- a/docker/unsloth_sync_notebooks.sh +++ b/docker/unsloth_sync_notebooks.sh @@ -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 diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 2d98988e50..f19f08fff0 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -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": diff --git a/unsloth/dataprep/synthetic.py b/unsloth/dataprep/synthetic.py index 39814de4d6..e009e3caab 100644 --- a/unsloth/dataprep/synthetic.py +++ b/unsloth/dataprep/synthetic.py @@ -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 ):