Studio: keep training from failing when a namespace-package shadows unsloth (#6269)
* Studio: guard the training import against a namespace-package shadow of unsloth A directory named unsloth without an __init__.py on sys.path (a stray checkout, a partial clone, or a polluted PYTHONPATH) makes the path finder return a namespace package, so the worker's 'from unsloth import FastLanguageModel' fails with 'cannot import name ... (unknown location)'. A normal site-packages install always wins this race, so only source/editable installs are exposed. Before the import, drop the offending sys.path entries, bind the real packages, then restore sys.path so other modules on those entries keep importing. It is a no-op when unsloth already resolves to a real package. * Studio: import unsloth before unsloth_zoo in the namespace-shadow guard _ensure_real_packages imported the requested names in argument order, so a unsloth_zoo namespace shadow made it import unsloth_zoo directly before unsloth. That skips unsloth.__init__ -> _gpu_init, which runs its ROCm and Windows bitsandbytes fixes before its own import unsloth_zoo, so the recovery path could import zoo with those guards skipped and fail on the ROCm/Windows cases those fixes handle. Import parent-first via reversed(names) so unsloth is imported first and pulls in the real unsloth_zoo after _gpu_init has run; the later cached import is a no-op. This also covers the case where only unsloth_zoo is shadowed: the bad sys.path entry is still dropped and unsloth owns the zoo import. Detection, sys.path pruning, shadow-cache clearing, and restoration are unchanged. Add tests/test_namespace_shadow_guard_pr6269.py: CPU-only subprocess scenarios (only zoo shadowed, both shadowed, only unsloth shadowed, healthy no-op, real package absent, multiple shadow entries) that assert unsloth imports before unsloth_zoo and that sys.path is restored. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: restore sys.path even if invalidate_caches fails in the shadow guard Move importlib.invalidate_caches() inside the try/finally so a failure there still restores sys.path, and tighten the import-order comment. Add a test that forces invalidate_caches to raise and asserts sys.path is restored. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han
GitHub
parent
785d446fc1
commit
f297593a7c
2 changed files with 338 additions and 0 deletions
|
|
@ -44,6 +44,62 @@ from utils.hardware import (
|
|||
# doesn't crash on RDNA2/RDNA3 with older ROCm wheels.
|
||||
if hasattr(torch._dynamo.config, "recompile_limit"):
|
||||
torch._dynamo.config.recompile_limit = 64
|
||||
|
||||
|
||||
def _ensure_real_packages(*names: str) -> None:
|
||||
"""Stop `import <name>` from binding to a namespace-package shadow.
|
||||
|
||||
A directory named like the package but missing __init__.py on sys.path (a
|
||||
stray checkout, a partial clone, or a polluted PYTHONPATH) makes the path
|
||||
finder return a namespace package, so `from unsloth import FastLanguageModel`
|
||||
dies with "cannot import name ... (unknown location)". A normal
|
||||
site-packages install always wins, so only source/editable installs are
|
||||
exposed. Drop the offending entries, import the real packages, then restore
|
||||
sys.path so other modules on those entries keep importing.
|
||||
"""
|
||||
import importlib
|
||||
import importlib.util
|
||||
|
||||
bad: set = set()
|
||||
shadowed: list = []
|
||||
for name in names:
|
||||
try:
|
||||
spec = importlib.util.find_spec(name)
|
||||
except (ImportError, ValueError, AttributeError):
|
||||
spec = None
|
||||
# a real package exposes its __init__ via spec.origin; a namespace
|
||||
# shadow has origin None/"namespace" and only search locations
|
||||
if spec is None or spec.origin not in (None, "namespace"):
|
||||
continue
|
||||
dirs = {os.path.realpath(d) for d in (spec.submodule_search_locations or [])}
|
||||
if not dirs:
|
||||
continue
|
||||
shadowed.append(name)
|
||||
for entry in sys.path:
|
||||
pkg = os.path.join(entry or os.getcwd(), name)
|
||||
if os.path.realpath(pkg) in dirs and not os.path.isfile(
|
||||
os.path.join(pkg, "__init__.py")
|
||||
):
|
||||
bad.add(entry)
|
||||
if not bad:
|
||||
return
|
||||
saved = list(sys.path)
|
||||
sys.path[:] = [e for e in sys.path if e not in bad]
|
||||
for name in shadowed:
|
||||
for cached in [m for m in list(sys.modules) if m == name or m.startswith(name + ".")]:
|
||||
del sys.modules[cached]
|
||||
try:
|
||||
importlib.invalidate_caches()
|
||||
# Import unsloth before unsloth_zoo (names are dependency-first):
|
||||
# unsloth.__init__ runs ROCm/Windows bnb fixes before it imports zoo,
|
||||
# so importing zoo first here would skip them. Repeat import is a no-op.
|
||||
for name in reversed(names):
|
||||
importlib.import_module(name)
|
||||
finally:
|
||||
sys.path[:] = saved
|
||||
|
||||
|
||||
_ensure_real_packages("unsloth_zoo", "unsloth")
|
||||
from unsloth import FastLanguageModel, FastVisionModel, is_bfloat16_supported
|
||||
from unsloth.chat_templates import get_chat_template
|
||||
|
||||
|
|
|
|||
282
studio/backend/tests/test_namespace_shadow_guard_pr6269.py
Normal file
282
studio/backend/tests/test_namespace_shadow_guard_pr6269.py
Normal file
|
|
@ -0,0 +1,282 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Verification tests for PR #6269 (training-worker namespace-shadow guard).
|
||||
|
||||
`_ensure_real_packages` (core/training/trainer.py) drops namespace-package
|
||||
shadow dirs (a `unsloth`/`unsloth_zoo` dir with no __init__.py on sys.path)
|
||||
before `from unsloth import ...`. Order matters: `unsloth.__init__` runs its
|
||||
ROCm/Windows bnb fixes before importing unsloth_zoo, so the guard must import
|
||||
unsloth first. Each test runs the real guard (ast-extracted from source, no
|
||||
GPU/torch) in a subprocess, with fake packages reachable only via a meta path
|
||||
finder to mimic an editable/PEP 660 install where the shadow wins.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import textwrap
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
TRAINER_PY = Path(__file__).resolve().parents[1] / "core" / "training" / "trainer.py"
|
||||
|
||||
|
||||
# ── fake package bodies ──────────────────────────────────────────────
|
||||
# Real `unsloth` sets a sentinel (mimics _gpu_init's pre-zoo fixes) then
|
||||
# imports unsloth_zoo, which records whether that sentinel was set first.
|
||||
|
||||
_REAL_UNSLOTH_INIT = textwrap.dedent(
|
||||
"""
|
||||
import os
|
||||
with open(os.environ["GUARD_ORDER_FILE"], "a") as _f:
|
||||
_f.write("unsloth\\n")
|
||||
# mimic _gpu_init: pre-zoo ROCm/Windows fixes run before importing zoo
|
||||
os.environ["UNSLOTH_GPU_INIT_RAN"] = "1"
|
||||
import unsloth_zoo # noqa: F401
|
||||
REAL = True
|
||||
"""
|
||||
)
|
||||
|
||||
_REAL_ZOO_INIT = textwrap.dedent(
|
||||
"""
|
||||
import os
|
||||
with open(os.environ["GUARD_ORDER_FILE"], "a") as _f:
|
||||
_f.write("unsloth_zoo\\n")
|
||||
with open(os.environ["GUARD_SENTINEL_FILE"], "w") as _f:
|
||||
_f.write(os.environ.get("UNSLOTH_GPU_INIT_RAN", "0"))
|
||||
REAL = True
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
# ── subprocess driver ────────────────────────────────────────────────
|
||||
# Reads a JSON config, rebuilds sys.path / sys.meta_path to model the
|
||||
# scenario, runs the real guard, and writes the observed result as JSON.
|
||||
|
||||
_DRIVER = textwrap.dedent(
|
||||
"""
|
||||
import ast, importlib.abc, importlib.util, json, os, sys
|
||||
|
||||
cfg = json.load(open(sys.argv[1]))
|
||||
|
||||
# Extract the real _ensure_real_packages from trainer.py source without
|
||||
# importing the heavy module or its `from unsloth import ...` line.
|
||||
src = open(cfg["trainer_py"]).read()
|
||||
tree = ast.parse(src)
|
||||
fn = next(n for n in tree.body
|
||||
if isinstance(n, ast.FunctionDef) and n.name == "_ensure_real_packages")
|
||||
mod = ast.Module(body=[fn], type_ignores=[])
|
||||
ast.fix_missing_locations(mod)
|
||||
ns = {"os": os, "sys": sys}
|
||||
exec(compile(mod, cfg["trainer_py"], "exec"), ns)
|
||||
_ensure_real_packages = ns["_ensure_real_packages"]
|
||||
|
||||
# Under -S site-packages is off, so a shadow root placed first wins the
|
||||
# path finder; the real packages come only from the meta finder below.
|
||||
for root in reversed(cfg["shadow_roots"]):
|
||||
sys.path.insert(0, root)
|
||||
|
||||
# Real packages are reachable only via a meta path finder appended AFTER
|
||||
# the standard PathFinder -- the editable / PEP 660 install shape.
|
||||
real_root = cfg.get("real_root")
|
||||
if real_root:
|
||||
class _RealFinder(importlib.abc.MetaPathFinder):
|
||||
def find_spec(self, fullname, path=None, target=None):
|
||||
if fullname.split(".")[0] not in ("unsloth", "unsloth_zoo"):
|
||||
return None
|
||||
parts = fullname.split(".")
|
||||
init = os.path.join(real_root, *parts, "__init__.py")
|
||||
if os.path.isfile(init):
|
||||
return importlib.util.spec_from_file_location(
|
||||
fullname, init,
|
||||
submodule_search_locations=[os.path.dirname(init)],
|
||||
)
|
||||
return None
|
||||
sys.meta_path.append(_RealFinder())
|
||||
|
||||
# optionally force importlib.invalidate_caches to raise, to prove the guard
|
||||
# still restores sys.path in that case
|
||||
if cfg.get("raise_on_invalidate"):
|
||||
def _boom():
|
||||
raise RuntimeError("invalidate_caches failed")
|
||||
importlib.invalidate_caches = _boom
|
||||
|
||||
path_before = list(sys.path)
|
||||
result = {"error": None}
|
||||
try:
|
||||
_ensure_real_packages(*cfg["names"])
|
||||
except Exception as e:
|
||||
result["error"] = type(e).__name__
|
||||
|
||||
result["path_restored"] = list(sys.path) == path_before
|
||||
of = os.environ["GUARD_ORDER_FILE"]
|
||||
result["order"] = open(of).read().split() if os.path.isfile(of) else []
|
||||
sf = os.environ["GUARD_SENTINEL_FILE"]
|
||||
result["sentinel_when_zoo_imported"] = (
|
||||
open(sf).read().strip() if os.path.isfile(sf) else None
|
||||
)
|
||||
|
||||
def _real(name):
|
||||
m = sys.modules.get(name)
|
||||
return bool(m and getattr(m, "REAL", False))
|
||||
result["unsloth_real"] = _real("unsloth")
|
||||
result["unsloth_zoo_real"] = _real("unsloth_zoo")
|
||||
|
||||
json.dump(result, open(cfg["out"], "w"))
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _make_namespace_shadow(root: Path, name: str) -> None:
|
||||
"""A directory named `name` with no __init__.py -> namespace portion."""
|
||||
(root / name).mkdir(parents = True, exist_ok = True)
|
||||
|
||||
|
||||
def _make_real_pkg(root: Path) -> None:
|
||||
(root / "unsloth").mkdir(parents = True, exist_ok = True)
|
||||
(root / "unsloth" / "__init__.py").write_text(_REAL_UNSLOTH_INIT)
|
||||
(root / "unsloth_zoo").mkdir(parents = True, exist_ok = True)
|
||||
(root / "unsloth_zoo" / "__init__.py").write_text(_REAL_ZOO_INIT)
|
||||
|
||||
|
||||
def _run(
|
||||
tmp_path: Path,
|
||||
*,
|
||||
shadow_roots,
|
||||
real: bool,
|
||||
names = ("unsloth_zoo", "unsloth"),
|
||||
trainer_py: Path = TRAINER_PY,
|
||||
raise_on_invalidate: bool = False,
|
||||
):
|
||||
order_file = tmp_path / "order.txt"
|
||||
sentinel_file = tmp_path / "sentinel.txt"
|
||||
out = tmp_path / "result.json"
|
||||
real_root = tmp_path / "real"
|
||||
if real:
|
||||
_make_real_pkg(real_root)
|
||||
|
||||
cfg = {
|
||||
"trainer_py": str(trainer_py),
|
||||
"shadow_roots": [str(r) for r in shadow_roots],
|
||||
"real_root": str(real_root) if real else None,
|
||||
"names": list(names),
|
||||
"out": str(out),
|
||||
"raise_on_invalidate": raise_on_invalidate,
|
||||
}
|
||||
cfg_path = tmp_path / "cfg.json"
|
||||
cfg_path.write_text(json.dumps(cfg))
|
||||
driver = tmp_path / "driver.py"
|
||||
driver.write_text(_DRIVER)
|
||||
|
||||
env = dict(os.environ)
|
||||
env["GUARD_ORDER_FILE"] = str(order_file)
|
||||
env["GUARD_SENTINEL_FILE"] = str(sentinel_file)
|
||||
env.pop("UNSLOTH_GPU_INIT_RAN", None)
|
||||
# Drop PYTHONPATH too so nothing re-introduces the real packages onto the
|
||||
# path; -S already keeps site-packages off.
|
||||
env.pop("PYTHONPATH", None)
|
||||
|
||||
proc = subprocess.run(
|
||||
# -S: skip site-packages so the installed unsloth_zoo can't shadow-beat
|
||||
# the namespace portion; the real package is served by the meta finder.
|
||||
[sys.executable, "-S", str(driver), str(cfg_path)],
|
||||
env = env,
|
||||
capture_output = True,
|
||||
text = True,
|
||||
timeout = 120,
|
||||
)
|
||||
assert (
|
||||
out.is_file()
|
||||
), f"driver did not produce a result\nstdout:\n{proc.stdout}\nstderr:\n{proc.stderr}"
|
||||
return json.loads(out.read_text())
|
||||
|
||||
|
||||
# ── scenarios ────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_only_zoo_shadowed_imports_unsloth_first(tmp_path):
|
||||
"""Core concern: a unsloth_zoo shadow must not make the guard import
|
||||
unsloth_zoo before unsloth -- that would skip _gpu_init's pre-zoo fixes."""
|
||||
shadow = tmp_path / "shadow"
|
||||
_make_namespace_shadow(shadow, "unsloth_zoo")
|
||||
res = _run(tmp_path, shadow_roots = [shadow], real = True)
|
||||
|
||||
assert res["error"] is None
|
||||
assert res["order"] == ["unsloth", "unsloth_zoo"], res
|
||||
assert res["sentinel_when_zoo_imported"] == "1", res
|
||||
assert res["unsloth_real"] and res["unsloth_zoo_real"], res
|
||||
assert res["path_restored"], res
|
||||
|
||||
|
||||
def test_both_shadowed_imports_unsloth_first(tmp_path):
|
||||
shadow = tmp_path / "shadow"
|
||||
_make_namespace_shadow(shadow, "unsloth")
|
||||
_make_namespace_shadow(shadow, "unsloth_zoo")
|
||||
res = _run(tmp_path, shadow_roots = [shadow], real = True)
|
||||
|
||||
assert res["error"] is None
|
||||
assert res["order"] == ["unsloth", "unsloth_zoo"], res
|
||||
assert res["sentinel_when_zoo_imported"] == "1", res
|
||||
assert res["unsloth_real"] and res["unsloth_zoo_real"], res
|
||||
assert res["path_restored"], res
|
||||
|
||||
|
||||
def test_only_unsloth_shadowed_recovers(tmp_path):
|
||||
shadow = tmp_path / "shadow"
|
||||
_make_namespace_shadow(shadow, "unsloth")
|
||||
res = _run(tmp_path, shadow_roots = [shadow], real = True)
|
||||
|
||||
assert res["error"] is None
|
||||
assert res["order"] == ["unsloth", "unsloth_zoo"], res
|
||||
assert res["unsloth_real"] and res["unsloth_zoo_real"], res
|
||||
assert res["path_restored"], res
|
||||
|
||||
|
||||
def test_healthy_install_is_noop(tmp_path):
|
||||
"""No shadow on sys.path: the guard imports nothing and leaves sys.path."""
|
||||
res = _run(tmp_path, shadow_roots = [], real = True)
|
||||
|
||||
assert res["error"] is None
|
||||
assert res["order"] == [], res # guard short-circuits before importing
|
||||
assert res["path_restored"], res
|
||||
|
||||
|
||||
def test_real_package_absent_surfaces_module_not_found(tmp_path):
|
||||
shadow = tmp_path / "shadow"
|
||||
_make_namespace_shadow(shadow, "unsloth_zoo")
|
||||
res = _run(tmp_path, shadow_roots = [shadow], real = False)
|
||||
|
||||
assert res["error"] == "ModuleNotFoundError", res
|
||||
assert res["path_restored"], res # restored even on failure
|
||||
|
||||
|
||||
def test_multiple_shadow_entries_all_removed(tmp_path):
|
||||
shadow_a = tmp_path / "shadow_a"
|
||||
shadow_b = tmp_path / "shadow_b"
|
||||
_make_namespace_shadow(shadow_a, "unsloth_zoo")
|
||||
_make_namespace_shadow(shadow_b, "unsloth_zoo")
|
||||
res = _run(tmp_path, shadow_roots = [shadow_a, shadow_b], real = True)
|
||||
|
||||
assert res["error"] is None
|
||||
# if only one offending entry were dropped, zoo would still resolve to a
|
||||
# namespace shadow and unsloth_zoo_real would be False
|
||||
assert res["unsloth_zoo_real"], res
|
||||
assert res["order"] == ["unsloth", "unsloth_zoo"], res
|
||||
assert res["path_restored"], res
|
||||
|
||||
|
||||
def test_sys_path_restored_if_invalidate_caches_raises(tmp_path):
|
||||
"""sys.path is restored even if importlib.invalidate_caches raises."""
|
||||
shadow = tmp_path / "shadow"
|
||||
_make_namespace_shadow(shadow, "unsloth_zoo")
|
||||
res = _run(tmp_path, shadow_roots = [shadow], real = True, raise_on_invalidate = True)
|
||||
|
||||
assert res["error"] == "RuntimeError", res
|
||||
assert res["path_restored"], res
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(pytest.main([__file__, "-v"]))
|
||||
Loading…
Add table
Add a link
Reference in a new issue