diff --git a/.github/workflows/lint-ci.yml b/.github/workflows/lint-ci.yml index 00e6e357e2..8cb3535075 100644 --- a/.github/workflows/lint-ci.yml +++ b/.github/workflows/lint-ci.yml @@ -79,6 +79,56 @@ jobs: run: | ruff check unsloth unsloth_cli studio tests cli.py unsloth-cli.py + - name: Import-hoist verifier self-test + # scripts/verify_import_hoist.py is a scope-aware (LEGB) AST + # resolver that gates import-hoisting / alias-rename refactors + # against two bugs ruff and pyflakes both miss: + # 1. dangling alias -- `from a import b as _b` hoisted to + # `from a import b` but a leftover `_b` reference now + # resolves to nothing (or to some other module-level `_b`). + # 2. rename clash -- `_b -> b` silently re-points at a + # different object already named `b` in that scope. + # This step runs the tool's 8 negative-control cases so a + # regression in the verifier itself fails before we trust it on + # a diff. Hermetic, stdlib-only, sub-second. Hard gate. + run: | + python scripts/verify_import_hoist.py --self-test + + - name: Import-hoist / alias-rename safety (changed Python files) + # Runs the verifier in compare mode on every in-place-modified + # .py in the PR: parses each file BEFORE (base branch) and AFTER + # (this diff), resolves every name load, and fails on a BLOCKER + # (dangling alias / rename clash / re-pointed import). INFO + # findings (a helper relocated to another file) do not fail. + # + # --diff-filter=M (in-place edits only) is deliberate: that is + # exactly where a hoist refactor lives, and it skips brand-new + # files whose re-export imports would otherwise look "unused". + # + # actions/checkout uses fetch-depth: 1, so the base branch is not + # present locally. Fetch the single base commit with an explicit + # refspec so origin/ is reliably created (a bare + # `git fetch origin ` only updates FETCH_HEAD in some + # configs). Two-dot diff avoids needing a merge-base on a shallow + # clone. + if: github.event_name == 'pull_request' + run: | + git fetch --no-tags --depth=1 origin \ + "${{ github.base_ref }}:refs/remotes/origin/${{ github.base_ref }}" + mapfile -t CHANGED < <( + git diff --name-only --diff-filter=M \ + "origin/${{ github.base_ref }}" HEAD -- '*.py' \ + | grep -vE '(^|/)(unsloth_compiled_cache|node_modules|build|dist)/' || true + ) + if [ "${#CHANGED[@]}" -eq 0 ]; then + echo "no in-place-modified Python files to check" + exit 0 + fi + printf 'checking %d file(s):\n' "${#CHANGED[@]}" + printf ' %s\n' "${CHANGED[@]}" + python scripts/verify_import_hoist.py \ + --before "origin/${{ github.base_ref }}" --after HEAD "${CHANGED[@]}" + - name: No leftover debugger / pdb / breakpoint calls # Catches the "I'll just stick a breakpoint() here" mistake # before it ships. AST-based so commented-out debugger diff --git a/scripts/verify_import_hoist.py b/scripts/verify_import_hoist.py new file mode 100644 index 0000000000..606488cc7f --- /dev/null +++ b/scripts/verify_import_hoist.py @@ -0,0 +1,854 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. +"""Deterministic, scope-aware verifier for import-hoisting / alias-rename refactors. + +The risk when moving `from a import b as _b` (or `import b as _b`) to module top +and normalizing `_b` -> `b` is twofold: + + 1. DANGLING ALIAS - a `_b` reference is left un-normalized; it now resolves to + nothing (NameError) or, worse, to some *other* module-level `_b`. + 2. RENAME CLASH - `_b` was an alias on purpose because `b` already meant + something else in that scope; normalizing `_b` -> `b` silently re-points the + reference at the wrong object (no NameError, no pyflakes warning). + +This tool parses BEFORE (a git ref, default origin/main) and AFTER (default HEAD) +for each file, builds a real LEGB scope model (functions, classes, lambdas, +comprehensions, global/nonlocal, args, walrus, star-imports), and resolves every +Name load to its binding. It then compares, PER SCOPE: + + * UNRESOLVED-NEW : loads that resolve to nothing in AFTER but did in BEFORE + (or are newly present) -> catches dangling aliases. + * TARGET-MISSING : an import *target* (e.g. module `glob`, or + `importlib.metadata.version`) that a function resolved to + in BEFORE but no longer resolves to in AFTER -> catches a + function that lost access to a module it still uses. + Robust to alias renames because it compares the *target*, + not the local name. + * TARGET-CHANGED : a load whose resolved import target differs BEFORE vs + AFTER -> catches a rename that re-points to a different + module (the clash case). + * AMBIGUOUS-BIND : a name bound by BOTH an import and a non-import in the same + scope in AFTER (and not in BEFORE) -> the "alias was on + purpose / now collides" smell. + * MODULE-DUP-IMPORT: a module-level name imported and also defined/assigned at + module level (introduced by the change). + * NEW-UNUSED-IMPORT: a module-level import added in AFTER that nothing resolves + to (informational; re-exports are a known false positive). + +Usage: + verify_import_hoist.py [--before REF] [--after REF] ... # compare + verify_import_hoist.py --self-test # prove it catches bugs +Exit code 1 if any non-informational finding. +""" + +from __future__ import annotations + +import argparse +import ast +import builtins +import re as _re_mod +import subprocess +import sys +from dataclasses import dataclass, field + +_BUILTINS = set(dir(builtins)) | { + "__file__", + "__name__", + "__doc__", + "__package__", + "__spec__", + "__loader__", + "__builtins__", + "__class__", + "__annotations__", + "__dict__", + "__qualname__", + "__module__", + "__path__", + "__debug__", + "__import__", + "NotImplemented", + "Ellipsis", + "copyright", + "credits", + "license", + "help", + "exit", + "quit", + "__build_class__", + "__cached__", + "reveal_type", + "reveal_locals", +} + + +# ---------------------------------------------------------------- scope model + + +@dataclass +class Binding: + kind: str # 'import' | 'importfrom' | 'def' | 'class' | 'other' + target: str | None = None # canonical import target id, else None + + +@dataclass +class Scope: + kind: str # 'module' | 'function' | 'class' | 'lambda' | 'comp' + qualname: str + parent: "Scope | None" + bindings: dict[str, list[Binding]] = field(default_factory = dict) + globals: set[str] = field(default_factory = set) + nonlocals: set[str] = field(default_factory = set) + star_import: bool = False + + def add(self, name: str, b: Binding) -> None: + self.bindings.setdefault(name, []).append(b) + + +def _import_target(node: ast.AST, alias: ast.alias) -> tuple[str, str]: + """Return (bound_name, canonical_target_id) for one import alias.""" + if isinstance(node, ast.Import): + bound = alias.asname or alias.name.split(".")[0] + return bound, f"import:{alias.name}" + # ImportFrom + bound = alias.asname or alias.name + mod = ("." * (node.level or 0)) + (node.module or "") + return bound, f"from:{mod}:{alias.name}" + + +class _Builder(ast.NodeVisitor): + """Builds the scope tree + bindings, and records every (scope, Name-load).""" + + def __init__(self): + self.module = Scope("module", "", None) + self.uses: list[tuple[Scope, str, int]] = [] # (scope, name, lineno) hard loads + self.soft_uses: list[ + tuple[Scope, str, int] + ] = [] # annotations: count as "used" + # but never as "unresolved" + # (forward refs / string annos) + + def _visit_annotation(self, node, scope: Scope) -> None: + """Annotation context: with `from __future__ import annotations` these are + never evaluated (strings), and even otherwise they routinely contain forward + references. Record contained names as SOFT uses so an import used only in an + annotation still counts as used, but a forward-ref name is never 'unresolved'.""" + if node is None: + return + for n in ast.walk(node): + if isinstance(n, ast.Name) and isinstance(n.ctx, ast.Load): + self.soft_uses.append((scope, n.id, n.lineno)) + + # -- binding helpers -- + def _bind_targets(self, scope: Scope, target: ast.AST) -> None: + for n in ast.walk(target): + if isinstance(n, ast.Name) and isinstance(n.ctx, (ast.Store, ast.Del)): + self._bind_name(scope, n.id, Binding("other")) + elif isinstance(n, ast.Starred): + pass + + def _bind_name(self, scope: Scope, name: str, b: Binding) -> None: + if name in scope.globals: + self.module.add(name, b) + elif name in scope.nonlocals: + p = scope.parent + while p is not None and p.kind not in ("function", "lambda"): + p = p.parent + (p or self.module).add(name, b) + else: + scope.add(name, b) + + # -- generic dispatch within a scope -- + def _visit_body(self, stmts, scope: Scope) -> None: + for s in stmts: + self._visit_stmt(s, scope) + + def _visit_stmt(self, node: ast.AST, scope: Scope) -> None: + if isinstance(node, (ast.Import, ast.ImportFrom)): + star = isinstance(node, ast.ImportFrom) and any( + a.name == "*" for a in node.names + ) + if star: + scope.star_import = True + for alias in node.names: + if alias.name == "*": + continue + bound, target = _import_target(node, alias) + kind = "import" if isinstance(node, ast.Import) else "importfrom" + self._bind_name(scope, bound, Binding(kind, target)) + return + if isinstance(node, ast.Global): + scope.globals.update(node.names) + return + if isinstance(node, ast.Nonlocal): + scope.nonlocals.update(node.names) + return + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + self._bind_name(scope, node.name, Binding("def")) + # decorators / defaults evaluate in the ENCLOSING scope + for d in node.decorator_list: + self._visit_expr(d, scope) + self._visit_arg_defaults(node.args, scope) + child = Scope("function", f"{scope.qualname}.{node.name}", scope) + self._bind_type_params(node, child) + self._bind_args(node.args, child) + # arg + return annotations: soft uses (may be strings / forward refs) + for a in self._all_args(node.args): + self._visit_annotation(a.annotation, child) + self._visit_annotation(getattr(node, "returns", None), child) + self._visit_body(node.body, child) + return + if isinstance(node, ast.ClassDef): + self._bind_name(scope, node.name, Binding("class")) + for d in node.decorator_list: + self._visit_expr(d, scope) + for b in node.bases: + self._visit_expr(b, scope) + for kw in node.keywords: + self._visit_expr(kw.value, scope) + child = Scope("class", f"{scope.qualname}.{node.name}", scope) + self._bind_type_params(node, child) + self._visit_body(node.body, child) + return + if isinstance(node, ast.Match): + self._visit_expr(node.subject, scope) + for case in node.cases: + self._bind_pattern(case.pattern, scope) + if case.guard is not None: + self._visit_expr(case.guard, scope) + self._visit_body(case.body, scope) + return + if isinstance(node, getattr(ast, "TryStar", ())): # py3.11 except* + self._visit_body(node.body, scope) + for h in node.handlers: + if h.type is not None: + self._visit_expr(h.type, scope) + if h.name: + self._bind_name(scope, h.name, Binding("other")) + self._visit_body(h.body, scope) + self._visit_body(node.orelse, scope) + self._visit_body(node.finalbody, scope) + return + if isinstance(node, getattr(ast, "TypeAlias", ())): # py3.12 `type X = ...` + if isinstance(node.name, ast.Name): + self._bind_name(scope, node.name.id, Binding("other")) + self._visit_annotation(node.value, scope) + return + if isinstance(node, (ast.Assign, ast.AnnAssign, ast.AugAssign)): + targets = node.targets if isinstance(node, ast.Assign) else [node.target] + val = node.value + if val is not None: + self._visit_expr(val, scope) + if isinstance(node, ast.AnnAssign) and node.annotation is not None: + self._visit_annotation(node.annotation, scope) + for t in targets: + self._bind_targets(scope, t) + # AugAssign target is also a load + if isinstance(node, ast.AugAssign): + self._record_loads(t, scope) + return + if isinstance(node, (ast.For, ast.AsyncFor)): + self._visit_expr(node.iter, scope) + self._bind_targets(scope, node.target) + self._visit_body(node.body, scope) + self._visit_body(node.orelse, scope) + return + if isinstance(node, (ast.With, ast.AsyncWith)): + for item in node.items: + self._visit_expr(item.context_expr, scope) + if item.optional_vars is not None: + self._bind_targets(scope, item.optional_vars) + self._visit_body(node.body, scope) + return + if isinstance(node, ast.Try): + self._visit_body(node.body, scope) + for h in node.handlers: + if h.type is not None: + self._visit_expr(h.type, scope) + if h.name: + self._bind_name(scope, h.name, Binding("other")) + self._visit_body(h.body, scope) + self._visit_body(node.orelse, scope) + self._visit_body(node.finalbody, scope) + return + # generic statement: visit all child expressions/stmts in same scope + for child in ast.iter_child_nodes(node): + if isinstance(child, ast.stmt): + self._visit_stmt(child, scope) + else: + self._visit_expr(child, scope) + + # -- expressions -- + def _visit_arg_defaults(self, args: ast.arguments, scope: Scope) -> None: + for d in list(args.defaults) + [d for d in args.kw_defaults if d is not None]: + self._visit_expr(d, scope) + + def _all_args(self, args: ast.arguments) -> list[ast.arg]: + out = list(args.posonlyargs) + list(args.args) + list(args.kwonlyargs) + if args.vararg: + out.append(args.vararg) + if args.kwarg: + out.append(args.kwarg) + return out + + def _bind_args(self, args: ast.arguments, scope: Scope) -> None: + for a in self._all_args(args): + scope.add(a.arg, Binding("other")) + + def _bind_type_params(self, node, scope: Scope) -> None: + for tp in getattr(node, "type_params", []) or []: + name = getattr(tp, "name", None) + if isinstance(name, str): + scope.add(name, Binding("other")) + self._visit_annotation(getattr(tp, "bound", None), scope) + self._visit_annotation(getattr(tp, "default_value", None), scope) + + def _bind_pattern(self, pat, scope: Scope) -> None: + if pat is None: + return + if isinstance(pat, ast.MatchValue): + self._visit_expr(pat.value, scope) + elif isinstance(pat, ast.MatchSingleton): + pass + elif isinstance(pat, ast.MatchSequence): + for p in pat.patterns: + self._bind_pattern(p, scope) + elif isinstance(pat, ast.MatchStar): + if pat.name: + self._bind_name(scope, pat.name, Binding("other")) + elif isinstance(pat, ast.MatchMapping): + for k in pat.keys: + self._visit_expr(k, scope) + for p in pat.patterns: + self._bind_pattern(p, scope) + if pat.rest: + self._bind_name(scope, pat.rest, Binding("other")) + elif isinstance(pat, ast.MatchClass): + self._visit_expr(pat.cls, scope) + for p in pat.patterns: + self._bind_pattern(p, scope) + for p in pat.kwd_patterns: + self._bind_pattern(p, scope) + elif isinstance(pat, ast.MatchAs): + self._bind_pattern(pat.pattern, scope) + if pat.name: + self._bind_name(scope, pat.name, Binding("other")) + elif isinstance(pat, ast.MatchOr): + for p in pat.patterns: + self._bind_pattern(p, scope) + + def _record_loads(self, node: ast.AST, scope: Scope) -> None: + for n in ast.walk(node): + if isinstance(n, ast.Name) and isinstance(n.ctx, ast.Load): + self.uses.append((scope, n.id, n.lineno)) + + def _visit_expr(self, node: ast.AST, scope: Scope) -> None: + if isinstance(node, ast.Name): + if isinstance(node.ctx, ast.Load): + self.uses.append((scope, node.id, node.lineno)) + elif isinstance(node.ctx, (ast.Store, ast.Del)): + self._bind_name(scope, node.id, Binding("other")) + return + if isinstance(node, ast.Lambda): + self._visit_arg_defaults(node.args, scope) + child = Scope("lambda", f"{scope.qualname}.", scope) + self._bind_args(node.args, child) + self._visit_expr(node.body, child) + return + if isinstance( + node, (ast.ListComp, ast.SetComp, ast.GeneratorExp, ast.DictComp) + ): + child = Scope("comp", f"{scope.qualname}.", scope) + for i, gen in enumerate(node.generators): + # first iterable is evaluated in the enclosing scope + self._visit_expr(gen.iter, scope if i == 0 else child) + self._bind_targets(child, gen.target) + for cond in gen.ifs: + self._visit_expr(cond, child) + if isinstance(node, ast.DictComp): + self._visit_expr(node.key, child) + self._visit_expr(node.value, child) + else: + self._visit_expr(node.elt, child) + return + if isinstance(node, ast.NamedExpr): # walrus binds in enclosing scope + self._visit_expr(node.value, scope) + if isinstance(node.target, ast.Name): + self._bind_name(scope, node.target.id, Binding("other")) + return + for child in ast.iter_child_nodes(node): + if isinstance(child, ast.stmt): + self._visit_stmt(child, scope) + else: + self._visit_expr(child, scope) + + def run(self, tree: ast.Module) -> None: + self._visit_body(tree.body, self.module) + + +# ---------------------------------------------------------------- resolution + + +def _any_star(scope: Scope) -> bool: + c = scope + while c is not None: + if c.star_import: + return True + c = c.parent + return False + + +def _resolve(scope: Scope, name: str): + """LEGB resolution. Returns (status, bindings) where status in + {'local','import','other','builtin','star','unresolved'}.""" + # global / nonlocal redirection + start = scope + if name in scope.globals: + chain = [_module_of(scope)] + elif name in scope.nonlocals: + chain = _enclosing_functions(scope) + else: + chain = _legb_chain(scope) + for i, sc in enumerate(chain): + if sc is None: + continue + if name in sc.bindings: + binds = sc.bindings[name] + if any(b.kind in ("import", "importfrom") for b in binds): + return "import", binds + return "other", binds + if name in _BUILTINS: + return "builtin", [] + if _any_star(start): + return "star", [] + return "unresolved", [] + + +def _module_of(scope: Scope) -> Scope: + while scope.parent is not None: + scope = scope.parent + return scope + + +def _enclosing_functions(scope: Scope) -> list[Scope]: + out = [] + p = scope.parent + while p is not None: + if p.kind in ("function", "lambda"): + out.append(p) + p = p.parent + out.append(_module_of(scope)) + return out + + +def _legb_chain(scope: Scope) -> list[Scope]: + """Immediate scope, then enclosing scopes skipping class scopes, then module.""" + chain = [scope] + p = scope.parent + while p is not None: + if ( + p.kind != "class" or p.parent is None + ): # module-level class never happens; keep module + if p.kind != "class": + chain.append(p) + p = p.parent + return chain + + +# ---------------------------------------------------------------- analysis + + +def _analyze(src: str): + tree = ast.parse(src) + b = _Builder() + b.run(tree) + # Per-scope: unresolved load names, and import targets it resolves to. + unresolved: dict[str, set[str]] = {} + targets_by_scope: dict[str, set[str]] = {} + target_by_use: dict[tuple[str, str], set[str]] = {} + for scope, name, _ln in b.uses: + status, binds = _resolve(scope, name) + if status == "unresolved": + unresolved.setdefault(scope.qualname, set()).add(name) + elif status == "import": + tids = {bd.target for bd in binds if bd.target} + targets_by_scope.setdefault(scope.qualname, set()).update(tids) + target_by_use.setdefault((scope.qualname, name), set()).update(tids) + # soft uses (annotations): only contribute to "used", never to "unresolved" + for scope, name, _ln in b.soft_uses: + status, binds = _resolve(scope, name) + if status == "import": + tids = {bd.target for bd in binds if bd.target} + targets_by_scope.setdefault(scope.qualname, set()).update(tids) + # module-level binding info for clash checks + module = b.module + module_imports = { + n: bs + for n, bs in module.bindings.items() + if any(x.kind in ("import", "importfrom") for x in bs) + } + module_dup = { + n + for n, bs in module.bindings.items() + if any(x.kind in ("import", "importfrom") for x in bs) + and any(x.kind not in ("import", "importfrom") for x in bs) + } + # ambiguous: any scope where a name is bound by import AND non-import + ambiguous: dict[str, set[str]] = {} + + def walk_scopes(scope: Scope): + for n, bs in scope.bindings.items(): + if any(x.kind in ("import", "importfrom") for x in bs) and any( + x.kind not in ("import", "importfrom") for x in bs + ): + ambiguous.setdefault(scope.qualname, set()).add(n) + # scope tree isn't stored; rebuild via uses is hard. We approximate with module only. + + walk_scopes(module) + return { + "unresolved": unresolved, + "targets_by_scope": targets_by_scope, + "target_by_use": target_by_use, + "module_import_targets": { + n: {x.target for x in bs if x.target} for n, bs in module_imports.items() + }, + "module_dup": module_dup, + "ambiguous": ambiguous, + } + + +def _git_show(ref: str, path: str) -> str | None: + try: + return subprocess.run( + ["git", "show", f"{ref}:{path}"], capture_output = True, text = True, check = True + ).stdout + except subprocess.CalledProcessError: + return None + + +def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]]: + """Return list of (severity, message). severity in BLOCKER/WARN/INFO. + + Blocker signals (precise, no relocation false-positives): + UNRESOLVED-NEW - a load became undefined (dangling alias / removed import). + NEW-UNUSED-HOIST - a module-level import added by THIS change is resolved by + NO load. A correct hoist always wires its new import to a + reference; if the alias was left un-normalized OR renamed + to the wrong name, the hoisted import ends up unused. This + single signal catches BOTH user-described failure modes and + does NOT fire for code merely relocated to another file + (that removes the import, it doesn't add an unused one). + TARGET-CHANGED - the same (scope, name) load resolves to a different import + target before vs after (a same-name re-point). + """ + a = _analyze(before_src) + b = _analyze(after_src) + findings: list[tuple[str, str]] = [] + + def used_targets(analysis) -> set[str]: + out: set[str] = set() + for tids in analysis["targets_by_scope"].values(): + out |= tids + return out + + before_used = used_targets(a) + after_used = used_targets(b) + before_module_targets: set[str] = set() + for tids in a["module_import_targets"].values(): + before_module_targets |= tids + after_module_targets: set[str] = set() + for tids in b["module_import_targets"].values(): + after_module_targets |= tids + added_module_targets = after_module_targets - before_module_targets + + # 1. UNRESOLVED-NEW + for scope, names in b["unresolved"].items(): + new = names - a["unresolved"].get(scope, set()) + for n in sorted(new): + findings.append( + ( + "BLOCKER", + f"{path}: UNRESOLVED-NEW '{n}' in scope {scope} " + f"(undefined after change -> dangling alias / removed import)", + ) + ) + + # 2. HOISTED-IMPORT-UNUSED (the core botched-hoist / wrong-rename signal) + # A module-level import in AFTER that NO load resolves to, and which was + # either newly added by this change OR was actually used before. Excludes: + # - relocation (the import is REMOVED, so it's not in after at all) + # - stable pre-existing re-exports (unused before AND after, not newly added) + for n, tids in b["module_import_targets"].items(): + if tids & after_used: + continue # resolved by something -> fine + newly_added = bool(tids - before_module_targets) + was_used_before = bool(tids & before_used) + if newly_added or was_used_before: + why = ( + "added but unused" + if newly_added + else "was used before, now unused (references re-pointed)" + ) + findings.append( + ( + "BLOCKER", + f"{path}: HOISTED-IMPORT-UNUSED '{n}' ({sorted(tids)}) " + f"{why} -> un-normalized alias or wrong rename target?", + ) + ) + + # 3. TARGET-CHANGED (same scope+name resolves to a different import target) + for key, tafter in b["target_by_use"].items(): + tbefore = a["target_by_use"].get(key) + if tbefore and tbefore != tafter: + findings.append( + ( + "BLOCKER", + f"{path}: TARGET-CHANGED name '{key[1]}' in {key[0]} " + f"{sorted(tbefore)} -> {sorted(tafter)} (rename re-points module)", + ) + ) + + # 4. MODULE-DUP-IMPORT introduced + for n in sorted(b["module_dup"] - a["module_dup"]): + findings.append( + ( + "WARN", + f"{path}: MODULE-DUP-IMPORT '{n}' bound by import AND non-import " + f"at module level (possible clash)", + ) + ) + + # 5. AMBIGUOUS-BIND introduced (module scope) + for scope, names in b["ambiguous"].items(): + new = names - a["ambiguous"].get(scope, set()) + for n in sorted(new): + findings.append( + ("WARN", f"{path}: AMBIGUOUS-BIND '{n}' import+non-import in {scope}") + ) + + # 6. TARGET-MISSING (informational): a scope stopped resolving to an import + # target. Real bugs are already covered above; remaining cases are code + # relocated to another file (e.g. a moved helper). Shown for transparency. + for scope, tbefore in a["targets_by_scope"].items(): + tafter = b["targets_by_scope"].get(scope, set()) + for t in sorted(tbefore - tafter): + relocated = ( + "" + if t in added_module_targets + else " [target not re-added here -> likely relocated/deleted]" + ) + findings.append( + ("INFO", f"{path}: TARGET-MISSING {t} in scope {scope}{relocated}") + ) + return findings + + +# ---------------------------------------------------------------- self-test + +_SELF_TESTS = { + "dangling_alias": ( + # before: inline aliased import, used as _b + "import os\n" + "def f():\n" + " import glob as _b\n" + " return _b.glob('*')\n", + # after: hoisted to canonical, but reference NOT normalized -> _b dangles + "import os\n" "import glob\n" "def f():\n" " return _b.glob('*')\n", + "BLOCKER", + ), + "rename_clash": ( + # before: _b is a deliberate alias; `b` already means something else + "import re as _b\n" "b = 123\n" "def f():\n" " return _b.compile('x'), b\n", + # after: someone normalized _b -> b ; now f().b is the int, re is lost + "import re\n" "b = 123\n" "def f():\n" " return b.compile('x'), b\n", + "BLOCKER", # TARGET-MISSING from:.. or import:re in f + ), + "clean_rename": ( + "def f():\n" " import glob as _g\n" " return _g.glob('*')\n", + "import glob\n" "def f():\n" " return glob.glob('*')\n", + None, # expect NO blocker + ), + "clean_dedup_redundant": ( + "import sys\n" "def f():\n" " import sys\n" " return sys.argv\n", + "import sys\n" "def f():\n" " return sys.argv\n", + None, + ), + "from_import_dangling": ( + # from-import alias left un-normalized + "def f():\n" + " from importlib.metadata import version as _v\n" + " return _v('x')\n", + "from importlib.metadata import version\n" "def f():\n" " return _v('x')\n", + "BLOCKER", + ), + "local_var_clash": ( + # _b renamed to b, but b is a LOCAL variable in f -> import silently unused + "def f(b):\n" " import re as _b\n" " return _b.compile(b)\n", + "import re\n" + "def f(b):\n" + " return b.compile(b)\n", # 'b' is the param, not the module + "BLOCKER", + ), + "substring_safe": ( + # correct _copy->copy rename while a config_copy var exists: NO false positive + "def f(config):\n" + " import copy as _copy\n" + " config_copy = _copy.deepcopy(config)\n" + " return config_copy\n", + "import copy\n" + "def f(config):\n" + " config_copy = copy.deepcopy(config)\n" + " return config_copy\n", + None, + ), + "attr_access_not_a_use": ( + # x._b is attribute access, not a use of name _b; removing import _b is fine + "import os\n" + "def f(x):\n" + " import sys as _b\n" + " return x._b + _b.argv[0]\n", + "import os\n" "import sys\n" "def f(x):\n" " return x._b + sys.argv[0]\n", + None, + ), +} + + +def _self_test() -> int: + ok = True + for name, (before, after, expect) in _SELF_TESTS.items(): + findings = compare(before, after, f"<{name}>") + blockers = [m for sev, m in findings if sev == "BLOCKER"] + got = "BLOCKER" if blockers else None + passed = got == expect + ok = ok and passed + print(f"[{'PASS' if passed else 'FAIL'}] {name}: expect={expect} got={got}") + for sev, m in findings: + print(f" ({sev}) {m}") + print("\nSELF-TEST:", "ALL PASS" if ok else "FAILURES") + return 0 if ok else 1 + + +def _pyflakes_undefined(path: str) -> set[str] | None: + """Return the set of names pyflakes reports as 'undefined name' for `path`, + or None if pyflakes failed to run/parse the file.""" + try: + proc = subprocess.run( + [sys.executable, "-m", "pyflakes", path], capture_output = True, text = True + ) + except Exception: + return None + if "syntax error" in (proc.stdout + proc.stderr).lower(): + return None + names = set() + for line in proc.stdout.splitlines(): + m = _re_mod.search(r"undefined name '([^']+)'", line) + if m: + names.add(m.group(1)) + return names + + +def audit_files(paths: list[str]) -> int: + """Single-version robustness audit. For every file: confirm the analyzer does + not crash, then cross-check its 'unresolved' names against pyflakes. Any name + the resolver flags that pyflakes does NOT call undefined is a tool FALSE + POSITIVE (a resolver gap to fix).""" + n_files = n_err = n_fp = n_syntax = 0 + fp_detail: dict[str, set[str]] = {} + err_detail: dict[str, str] = {} + for path in paths: + n_files += 1 + try: + src = open(path, encoding = "utf-8").read() + except Exception as e: # unreadable + n_err += 1 + err_detail[path] = f"read: {e}" + continue + try: + res = _analyze(src) + except SyntaxError: + n_syntax += 1 + continue + except Exception as e: # analyzer crash -> robustness bug + n_err += 1 + err_detail[path] = f"{type(e).__name__}: {e}" + continue + tool_unresolved = set() + for names in res["unresolved"].values(): + tool_unresolved |= names + if not tool_unresolved: + continue + pf = _pyflakes_undefined(path) + if pf is None: + continue # pyflakes couldn't adjudicate; skip cross-check + false_pos = tool_unresolved - pf + if false_pos: + n_fp += 1 + fp_detail[path] = false_pos + print(f"audited files : {n_files}") + print(f"syntax-skipped : {n_syntax}") + print(f"analyzer errors : {n_err}") + for p, e in sorted(err_detail.items()): + print(f" ERROR {p}: {e}") + print(f"false-positive files: {n_fp} (resolver flagged a name pyflakes accepts)") + for p, names in sorted(fp_detail.items()): + print(f" FP {p}: {sorted(names)}") + ok = n_err == 0 and n_fp == 0 + print( + "\nAUDIT:", + "ROBUST (no crashes, no false positives vs pyflakes)" + if ok + else "NEEDS WORK (see above)", + ) + return 0 if ok else 1 + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--before", default = "origin/main") + ap.add_argument("--after", default = "HEAD") + ap.add_argument("--self-test", action = "store_true") + ap.add_argument( + "--audit", + action = "store_true", + help = "single-version robustness audit on filesystem paths", + ) + ap.add_argument("files", nargs = "*") + args = ap.parse_args() + + if args.self_test: + return _self_test() + if args.audit: + return audit_files(args.files) + + any_blocker = False + for path in args.files: + before = _git_show(args.before, path) + after = _git_show(args.after, path) + if after is None: + print(f"SKIP {path}: not found at {args.after}") + continue + if before is None: + before = "" # new file + findings = compare(before, after, path) + blockers = [f for f in findings if f[0] == "BLOCKER"] + warns = [f for f in findings if f[0] == "WARN"] + infos = [f for f in findings if f[0] == "INFO"] + status = ( + "CLEAN" + if not blockers and not warns + else ("BLOCKERS" if blockers else "WARNINGS") + ) + print(f"\n=== {path}: {status} ===") + for sev, m in blockers + warns + infos: + print(f" [{sev}] {m}") + any_blocker = any_blocker or bool(blockers) + print( + "\nOVERALL:", "FAIL (blockers found)" if any_blocker else "PASS (no blockers)" + ) + return 1 if any_blocker else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/studio/backend/core/_torchao_stub.py b/studio/backend/core/_torchao_stub.py new file mode 100644 index 0000000000..5650a60ee2 --- /dev/null +++ b/studio/backend/core/_torchao_stub.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Shared torchao Windows-ROCm import stub. + +torchao (pulled in by transformers.quantizers) imports +torch.distributed._functional_collectives at module level, which imports +distributed_c10d.py unconditionally — that file crashes on Windows ROCm because +torch._C._distributed_c10d (the RCCL backend) is absent. +torch/distributed/__init__.py itself is guarded by `if is_available()` so +`import torch.distributed` alone is safe; the crash only comes via torchao's +import chain. Stubbing torchao short-circuits it entirely. +_StubSubpackageFinder handles any depth of torchao.xxx.yyy imports. + +This logic used to be duplicated inline inside run_export_process() and +run_training_process(); it now lives here so both worker subprocesses call the +single `install_torchao_windows_rocm_stub()` entrypoint before importing +transformers / unsloth_zoo. +""" + +from __future__ import annotations + +import sys +import types +import importlib.abc +import importlib.machinery + +_STUB_SENTINEL = object() + + +# Metaclass for stub types so that isinstance(x, StubClass) returns False +# instead of raising TypeError ("arg 2 must be a type"). +# peft/tuners/lora/torchao.py does: +# from torchao.dtypes import AffineQuantizedTensor, LinearActivationQuantizedTensor +# isinstance(weight, (AffineQuantizedTensor, LinearActivationQuantizedTensor)) +# If those names resolve to stub modules rather than types, isinstance() raises. +class _StubTypeMeta(type): + def __instancecheck__(cls, instance): + return False + + def __subclasscheck__(cls, subclass): + return False + + def __getattr__(cls, attr): + if attr.startswith("__"): + raise AttributeError(attr) + child = _StubTypeMeta(attr, (), {}) + setattr(cls, attr, child) + return child + + def __call__(cls, *args, **kwargs): + return None + + +def _make_stub_type(name): + """Stub class: accepted by isinstance() (always False), supports attr access.""" + return _StubTypeMeta(name, (), {}) + + +def _make_mod_stub(mod_name): + m = types.ModuleType(mod_name) + m.__path__ = [] + m.__package__ = mod_name + m._unsloth_stub = _STUB_SENTINEL + m.__spec__ = importlib.machinery.ModuleSpec(mod_name, loader = None, is_package = True) + + def _ga(attr, _m = m, _n = mod_name): + if attr.startswith("__"): + raise AttributeError(attr) + # Return a stub CLASS (not a module) so that isinstance(x, attr) + # works and returns False instead of raising TypeError. + child = _make_stub_type(f"{_n}.{attr}") + setattr(_m, attr, child) + return child + + m.__getattr__ = _ga + return m + + +class _StubSubpackageLoader(importlib.abc.Loader): + def __init__(self, mod_name): + self._mod_name = mod_name + + def create_module(self, spec): + return _make_mod_stub(self._mod_name) + + def exec_module(self, module): + pass + + +class _StubSubpackageFinder(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path, target = None): + if "." not in fullname: + return None + parent = sys.modules.get(fullname.rsplit(".", 1)[0]) + if parent is None: + return None + if getattr(parent, "_unsloth_stub", None) is not _STUB_SENTINEL: + return None + return importlib.machinery.ModuleSpec( + fullname, _StubSubpackageLoader(fullname), is_package = True + ) + + +def install_torchao_windows_rocm_stub() -> None: + """Pre-stub torchao on Windows ROCm so transformers/peft imports don't crash. + + No-op on every other platform (Windows CUDA included — there torchao is real + and shadowing it would break torchao-based quantization paths). Must run + before any import of transformers / unsloth_zoo. Safe to call once per worker + process. + """ + # Gate on the active torch runtime, not env-var presence -- HIP_PATH / + # ROCM_PATH stay set after a user installs the HIP SDK and reverts to a + # CUDA torch wheel. AMD SDK / Radeon ROCm wheels may not set torch.version.hip + # but still encode "rocm" in torch.__version__, so accept either. + _is_win32_rocm = False + if sys.platform == "win32": + try: + import torch as _torch_probe + + _is_win32_rocm = bool( + getattr(getattr(_torch_probe, "version", None), "hip", None) + or "rocm" in getattr(_torch_probe, "__version__", "").lower() + ) + del _torch_probe + except Exception: + pass + if _is_win32_rocm: + # Register the finder only on Windows ROCm -- on other platforms there + # are no stub modules seeded, so appending is a pure accumulation. + sys.meta_path.append(_StubSubpackageFinder()) + # Seed torchao top-level + key submodules; the finder handles the rest. + for _tao_name in ( + "torchao", + "torchao.quantization", + "torchao.dtypes", + "torchao.float8", + "torchao.utils", + ): + if _tao_name not in sys.modules: + sys.modules[_tao_name] = _make_mod_stub(_tao_name) diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py index 77205bc298..defcff924b 100644 --- a/studio/backend/core/export/worker.py +++ b/studio/backend/core/export/worker.py @@ -440,101 +440,13 @@ def run_export_process( ) # ── 1c. Stub torchao on Windows ROCm ── - # torchao (pulled in by transformers.quantizers) imports - # torch.distributed._functional_collectives at module level, which imports - # distributed_c10d.py unconditionally — that file crashes on Windows ROCm - # because torch._C._distributed_c10d (the RCCL backend) is absent. - # Stubbing torchao short-circuits the crash entirely. + # Shared with the training worker; see core/_torchao_stub.py for the full + # rationale (torchao -> torch.distributed._functional_collectives crashes on + # Windows ROCm because the RCCL backend is absent). No-op off Windows ROCm. # Must run before any import of transformers / unsloth_zoo. - import types as _types - import importlib.machinery as _ilm - import importlib.abc as _ilabc + from core._torchao_stub import install_torchao_windows_rocm_stub - _STUB_SENTINEL = object() - - class _StubTypeMeta(type): - def __instancecheck__(cls, instance): - return False - - def __subclasscheck__(cls, subclass): - return False - - def __getattr__(cls, attr): - if attr.startswith("__"): - raise AttributeError(attr) - child = _StubTypeMeta(attr, (), {}) - setattr(cls, attr, child) - return child - - def __call__(cls, *args, **kwargs): - return None - - def _make_stub_type(name): - return _StubTypeMeta(name, (), {}) - - def _make_mod_stub(mod_name): - m = _types.ModuleType(mod_name) - m.__path__ = [] - m.__package__ = mod_name - m._unsloth_stub = _STUB_SENTINEL - m.__spec__ = _ilm.ModuleSpec(mod_name, loader = None, is_package = True) - - def _ga(attr, _m = m, _n = mod_name): - if attr.startswith("__"): - raise AttributeError(attr) - child = _make_stub_type(f"{_n}.{attr}") - setattr(_m, attr, child) - return child - - m.__getattr__ = _ga - return m - - class _StubSubpackageLoader(_ilabc.Loader): - def __init__(self, mod_name): - self._mod_name = mod_name - - def create_module(self, spec): - return _make_mod_stub(self._mod_name) - - def exec_module(self, module): - pass - - class _StubSubpackageFinder(_ilabc.MetaPathFinder): - def find_spec(self, fullname, path, target = None): - if "." not in fullname: - return None - parent = sys.modules.get(fullname.rsplit(".", 1)[0]) - if parent is None: - return None - if getattr(parent, "_unsloth_stub", None) is not _STUB_SENTINEL: - return None - return _ilm.ModuleSpec( - fullname, _StubSubpackageLoader(fullname), is_package = True - ) - - _is_win32_rocm = False - if sys.platform == "win32": - try: - import torch as _torch_probe - - _is_win32_rocm = bool( - getattr(getattr(_torch_probe, "version", None), "hip", None) - or "rocm" in getattr(_torch_probe, "__version__", "").lower() - ) - del _torch_probe - except Exception: - pass - if _is_win32_rocm: - sys.meta_path.append(_StubSubpackageFinder()) - for _tao_name in ( - "torchao", - "torchao.quantization", - "torchao.dtypes", - "torchao.float8", - "torchao.utils", - ): - if _tao_name not in sys.modules: - sys.modules[_tao_name] = _make_mod_stub(_tao_name) + install_torchao_windows_rocm_stub() # ── 2. Import ML libraries (fresh in this clean process) ── try: diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 83c4b4f4da..4b17f495ea 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -17,6 +17,7 @@ import struct import structlog from loggers import get_logger import shutil +import signal import socket import subprocess import sys @@ -965,9 +966,6 @@ class LlamaCppBackend: 7. llama-server on PATH (system install) 8. ./bin/llama-server (legacy: extracted binary) """ - import os - import sys - binary_name = "llama-server.exe" if sys.platform == "win32" else "llama-server" # 1. Env var — direct path to binary @@ -1282,8 +1280,6 @@ class LlamaCppBackend: Returns list of (gpu_index, free_mib) sorted by index. Empty list if no supported GPU is reachable. """ - import os - # ── NVIDIA via nvidia-smi ──────────────────────────────────── try: result = subprocess.run( @@ -3928,10 +3924,6 @@ class LlamaCppBackend: Falls back to pgrep + /proc//exe on Linux when psutil is not installed. """ - import os - import signal - import sys - try: # -- Build the ownership allowlist -------------------------------- # Two kinds of matches: diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index f1faaee4fb..0365b3ffd8 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -6,8 +6,10 @@ Unsloth Training Backend Integrates Unsloth training capabilities with the FastAPI backend """ +import gc import os import sys +import types # Prevent tokenizer parallelism deadlocks when datasets uses multiprocessing fork os.environ["TOKENIZERS_PARALLELISM"] = "false" @@ -420,8 +422,6 @@ class UnslothTrainer: in sys.modules. When the next training run calls dataset.map(num_proc=N), forked child processes inherit this stale state and deadlock. """ - import sys as _sys - # Remove cloned audio repo paths from sys.path base_dir = os.path.dirname(os.path.abspath(__file__)) audio_paths = [ @@ -436,15 +436,15 @@ class UnslothTrainer: removed_paths = [] for path in audio_paths: - if path in _sys.path: - _sys.path.remove(path) + if path in sys.path: + sys.path.remove(path) removed_paths.append(path) # Remove stale audio modules from sys.modules prefixes = ("snac", "whisper", "sparktts", "outetts") - removed_modules = [key for key in _sys.modules if key.startswith(prefixes)] + removed_modules = [key for key in sys.modules if key.startswith(prefixes)] for key in removed_modules: - del _sys.modules[key] + del sys.modules[key] if removed_paths or removed_modules: logger.info( @@ -541,10 +541,9 @@ class UnslothTrainer: # clear_unsloth_compiled_cache() deletes the disk cache, but the flag # prevents re-compilation — leaving missing cache files. Reloading # restores original class definitions so Unsloth can re-compile cleanly. - import sys as _sys import importlib - for _key, _mod in list(_sys.modules.items()): + for _key, _mod in list(sys.modules.items()): if "transformers.models." in _key and ".modeling_" in _key: if hasattr(_mod, "__UNSLOTH_PATCHED__"): try: @@ -660,14 +659,22 @@ class UnslothTrainer: f"Using device_map='{device_map}' ({get_visible_gpu_count()} GPU(s) visible)" ) - # On hardware without native bfloat16 support (e.g. RDNA2 / gfx103x), - # passing dtype=None lets unsloth auto-detect and incorrectly choose - # bf16, triggering an LLVM error at the first bf16 kernel dispatch. - # Explicitly pass float16 as the fallback so unsloth never reaches - # that path. Modern NVIDIA (Ampere+) and RDNA3+ return True here so - # they are unaffected — dtype stays None and unsloth picks bf16 as - # before. - _auto_dtype = None if is_bfloat16_supported() else torch.float16 + # AMD ROCm hardware without native bfloat16 (e.g. RDNA2 / gfx103x) + # crashes with an LLVM error at the first bf16 kernel dispatch if + # dtype=None lets unsloth auto-pick bf16. Force float16 there so that + # path is never reached. NVIDIA keeps dtype=None so unsloth's own + # bf16/fp16/float32 auto-detection (including FORCE_FLOAT32 models) is + # honored -- older NVIDIA without bf16 (T4/V100) must NOT be coerced to + # float16 here, which the previous unconditional branch did wrongly. + # Derive ROCm inline (not hardware.IS_ROCM) because that flag is unset + # until detect_hardware() runs, which isn't guaranteed in this subprocess. + _is_rocm = ( + bool(getattr(torch.version, "hip", None)) + or "rocm" in torch.__version__.lower() + ) + _auto_dtype = ( + torch.float16 if (_is_rocm and not is_bfloat16_supported()) else None + ) # Branch based on model type if self._audio_type == "csm": @@ -1200,7 +1207,6 @@ class UnslothTrainer: We patch at both instance AND class level for maximum reliability, and strip non-TransformersKwargs params that Unsloth/PEFT inject. """ - import types import torch import torch.nn as nn from transformers.models.csm.modeling_csm import ( @@ -1742,7 +1748,6 @@ class UnslothTrainer: logger.info("Freeing SNAC codec model from GPU...\n") snac_model.to("cpu") del snac_model - import gc gc.collect() torch.cuda.empty_cache() @@ -1766,13 +1771,10 @@ class UnslothTrainer: Mirrors Spark_TTS_(0_5B).ipynb: encode audio with BiCodec (semantic + global tokens), format as special-token text strings for SFTTrainer with dataset_text_field="text". """ - import sys import torch import numpy as np import torchaudio.transforms as T - import subprocess - device = "cuda" if torch.cuda.is_available() else "cpu" # The sparktts Python package lives in the SparkAudio/Spark-TTS GitHub repo, @@ -1972,7 +1974,6 @@ class UnslothTrainer: audio_tokenizer.model.cpu() audio_tokenizer.feature_extractor.cpu() del audio_tokenizer - import gc gc.collect() torch.cuda.empty_cache() @@ -2001,7 +2002,6 @@ class UnslothTrainer: OuteTTS AudioProcessor for speaker representations, PromptProcessor for training prompts. Outputs text strings for SFTTrainer with dataset_text_field="text". """ - import sys import io import tempfile import torch @@ -2185,7 +2185,6 @@ class UnslothTrainer: del whisper_model del audio_processor del prompt_processor - import gc gc.collect() torch.cuda.empty_cache() diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index a59def8751..a825321597 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -21,6 +21,9 @@ import shutil import sys import time import traceback +import gc +import re +import types import subprocess as _sp from pathlib import Path from typing import Any, Callable @@ -1259,7 +1262,6 @@ def _run_mlx_training(event_queue, stop_queue, config): Mirrors the event_queue protocol so the parent process pump works unchanged. """ import time - import gc import math import threading import queue as _queue @@ -2020,121 +2022,13 @@ def run_training_process( ) # ── 1d. Stub torchao on Windows ROCm ── - # torchao (pulled in by transformers.quantizers) imports - # torch.distributed._functional_collectives at module level, which imports - # distributed_c10d.py unconditionally — that file crashes on Windows ROCm - # because torch._C._distributed_c10d (the RCCL backend) is absent. - # torch/distributed/__init__.py itself is guarded by `if is_available()` - # so `import torch.distributed` alone is safe; the crash only comes via - # torchao's import chain. Stubbing torchao short-circuits it entirely. - # _StubSubpackageFinder handles any depth of torchao.xxx.yyy imports. - import types as _types - import importlib.machinery as _ilm - import importlib.abc as _ilabc + # Shared with the export worker; see core/_torchao_stub.py for the full + # rationale (torchao -> torch.distributed._functional_collectives crashes on + # Windows ROCm because the RCCL backend is absent). No-op off Windows ROCm. + # Must run before any import of transformers / unsloth_zoo. + from core._torchao_stub import install_torchao_windows_rocm_stub - _STUB_SENTINEL = object() - - # Metaclass for stub types so that isinstance(x, StubClass) returns False - # instead of raising TypeError ("arg 2 must be a type"). - # peft/tuners/lora/torchao.py does: - # from torchao.dtypes import AffineQuantizedTensor, LinearActivationQuantizedTensor - # isinstance(weight, (AffineQuantizedTensor, LinearActivationQuantizedTensor)) - # If those names resolve to stub modules rather than types, isinstance() raises. - class _StubTypeMeta(type): - def __instancecheck__(cls, instance): - return False - - def __subclasscheck__(cls, subclass): - return False - - def __getattr__(cls, attr): - if attr.startswith("__"): - raise AttributeError(attr) - child = _StubTypeMeta(attr, (), {}) - setattr(cls, attr, child) - return child - - def __call__(cls, *args, **kwargs): - return None - - def _make_stub_type(name): - """Stub class: accepted by isinstance() (always False), supports attr access.""" - return _StubTypeMeta(name, (), {}) - - def _make_mod_stub(mod_name): - m = _types.ModuleType(mod_name) - m.__path__ = [] - m.__package__ = mod_name - m._unsloth_stub = _STUB_SENTINEL - m.__spec__ = _ilm.ModuleSpec(mod_name, loader = None, is_package = True) - - def _ga(attr, _m = m, _n = mod_name): - if attr.startswith("__"): - raise AttributeError(attr) - # Return a stub CLASS (not a module) so that isinstance(x, attr) - # works and returns False instead of raising TypeError. - child = _make_stub_type(f"{_n}.{attr}") - setattr(_m, attr, child) - return child - - m.__getattr__ = _ga - return m - - class _StubSubpackageLoader(_ilabc.Loader): - def __init__(self, mod_name): - self._mod_name = mod_name - - def create_module(self, spec): - return _make_mod_stub(self._mod_name) - - def exec_module(self, module): - pass - - class _StubSubpackageFinder(_ilabc.MetaPathFinder): - def find_spec(self, fullname, path, target = None): - if "." not in fullname: - return None - parent = sys.modules.get(fullname.rsplit(".", 1)[0]) - if parent is None: - return None - if getattr(parent, "_unsloth_stub", None) is not _STUB_SENTINEL: - return None - return _ilm.ModuleSpec( - fullname, _StubSubpackageLoader(fullname), is_package = True - ) - - # Only stub torchao on Windows ROCm hosts -- on Windows CUDA (NVIDIA) torchao - # is real and shadowing it breaks torchao-based quantization paths. - # Gate on the active torch runtime, not env-var presence -- HIP_PATH / - # ROCM_PATH stay set after a user installs the HIP SDK and reverts to a - # CUDA torch wheel. AMD SDK / Radeon ROCm wheels may not set torch.version.hip - # but still encode "rocm" in torch.__version__, so accept either. - _is_win32_rocm = False - if sys.platform == "win32": - try: - import torch as _torch_probe - - _is_win32_rocm = bool( - getattr(getattr(_torch_probe, "version", None), "hip", None) - or "rocm" in getattr(_torch_probe, "__version__", "").lower() - ) - del _torch_probe - except Exception: - pass - if _is_win32_rocm: - # Register the finder only on Windows ROCm -- on other platforms there - # are no stub modules seeded, so appending is a pure accumulation. - sys.meta_path.append(_StubSubpackageFinder()) - # Seed torchao top-level + key submodules; the finder handles the rest. - for _tao_name in ( - "torchao", - "torchao.quantization", - "torchao.dtypes", - "torchao.float8", - "torchao.utils", - ): - if _tao_name not in sys.modules: - sys.modules[_tao_name] = _make_mod_stub(_tao_name) + install_torchao_windows_rocm_stub() # ── 1e. Ensure torch.distributed helper attrs are present ── # Single-GPU training never initialises the process group, so these helpers @@ -2155,7 +2049,7 @@ def run_training_process( if not hasattr(_td, _name): setattr(_td, _name, _stub) except Exception: - _td_mock = _types.ModuleType("torch.distributed") + _td_mock = types.ModuleType("torch.distributed") for _name, _stub in _td_stubs.items(): setattr(_td_mock, _name, _stub) sys.modules["torch.distributed"] = _td_mock @@ -2269,17 +2163,13 @@ def run_training_process( # (e.g. "2.11.0+rocm7.13.0" or "2.9.0+rocmsdk20251116"); fall back # to that string when version.hip is missing. def _hip_ver_at_least(major: int, minor: int) -> bool: - import re as _re_ver - _hip_str = getattr( getattr(_torch_for_rocm, "version", None), "hip", None ) if not _hip_str: # Try the standard "+rocmX.Y.Z" embedded version first # (e.g. "2.11.0+rocm7.13.0"). - _ver_match = _re_ver.search( - r"rocm(\d+)\.(\d+)", _build_version_for_rocm - ) + _ver_match = re.search(r"rocm(\d+)\.(\d+)", _build_version_for_rocm) if _ver_match: return ( int(_ver_match.group(1)), diff --git a/studio/backend/main.py b/studio/backend/main.py index be327bc5a8..6b8ac438c0 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -866,8 +866,6 @@ def _strip_crossorigin(html_bytes: bytes) -> bytes: @font-face downloads to fail silently. Stripping the attribute makes them regular same-origin fetches that work on any protocol. """ - import re as _re - html = html_bytes.decode("utf-8") html = _re.sub(r'\s+crossorigin(?:="[^"]*")?', "", html) return html.encode("utf-8") diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index ebac6a357c..180fde8f13 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -16,8 +16,16 @@ Usage: ... """ +import copy +import gc +import glob import os import platform +import re +import subprocess +import sys +import types +from importlib.metadata import PackageNotFoundError, version as pkg_version import structlog from loggers import get_logger from enum import Enum @@ -178,8 +186,6 @@ def clear_gpu_cache(): Clear GPU memory cache for the current device. Safe to call on any platform — no-ops gracefully. """ - import gc - gc.collect() device = get_device() @@ -361,8 +367,6 @@ def get_package_versions() -> Dict[str, Optional[str]]: Returns dict with keys: unsloth, torch, transformers, cuda. Missing packages yield None. """ - from importlib.metadata import version as pkg_version, PackageNotFoundError - packages = ("unsloth", "torch", "transformers") versions: Dict[str, Optional[str]] = {} @@ -481,9 +485,6 @@ def _read_apple_gpu_stats() -> Dict[str, Any]: Returns dict with utilization_pct, vram_used_bytes (system-wide GPU memory). Returns empty dict on failure. """ - import subprocess - import re - try: result = subprocess.run( ["ioreg", "-r", "-c", "AGXAccelerator"], @@ -510,12 +511,10 @@ def _read_apple_gpu_stats() -> Dict[str, Any]: def _rocm_linux_sysfs_gpu_busy_pct() -> Optional[float]: """Query AMD GPU compute utilization via Linux DRM sysfs gpu_busy_percent.""" - import glob as _glob - if platform.system() != "Linux": return None try: - files = _glob.glob("/sys/class/drm/card*/device/gpu_busy_percent") + files = glob.glob("/sys/class/drm/card*/device/gpu_busy_percent") if not files: return None values = [int(open(f).read().strip()) for f in files] @@ -526,12 +525,10 @@ def _rocm_linux_sysfs_gpu_busy_pct() -> Optional[float]: def _rocm_linux_sysfs_temp_c() -> Optional[float]: """Query AMD GPU edge temperature via Linux DRM hwmon sysfs (temp1_input, millidegrees C).""" - import glob as _glob - if platform.system() != "Linux": return None try: - files = _glob.glob("/sys/class/drm/card*/device/hwmon/hwmon*/temp1_input") + files = glob.glob("/sys/class/drm/card*/device/hwmon/hwmon*/temp1_input") if not files: return None temps = [int(open(f).read().strip()) / 1000.0 for f in files] @@ -542,8 +539,6 @@ def _rocm_linux_sysfs_temp_c() -> Optional[float]: def _rocm_linux_sysfs_power_w() -> Optional[float]: """Query AMD GPU average power draw via Linux DRM hwmon sysfs (microwatts).""" - import glob as _glob - if platform.system() != "Linux": return None try: @@ -551,7 +546,7 @@ def _rocm_linux_sysfs_power_w() -> Optional[float]: "/sys/class/drm/card*/device/hwmon/hwmon*/power1_average", "/sys/class/drm/card*/device/hwmon/hwmon*/power1_input", ): - files = _glob.glob(pattern) + files = glob.glob(pattern) if files: watts = sum(int(open(f).read().strip()) / 1_000_000.0 for f in files) return round(watts, 1) @@ -562,8 +557,6 @@ def _rocm_linux_sysfs_power_w() -> Optional[float]: def _rocm_windows_perf_counter_gpu_util_pct() -> Optional[float]: """Query AMD GPU compute utilization via Windows Performance Counters (3D engine nodes).""" - import subprocess as _sp - if platform.system() != "Windows": return None try: @@ -572,7 +565,7 @@ def _rocm_windows_perf_counter_gpu_util_pct() -> Optional[float]: " -ErrorAction SilentlyContinue).CounterSamples;" "if($s){[math]::Min(($s|Measure-Object CookedValue -Sum).Sum,100)}else{-1}" ) - r = _sp.run( + r = subprocess.run( ["powershell", "-NoProfile", "-NonInteractive", "-Command", ps], capture_output = True, text = True, @@ -593,13 +586,11 @@ def _rocm_linux_sysfs_vram_gb() -> tuple[Optional[float], Optional[float]]: updates in real-time across all processes. No tools required. Returns (used_gb, total_gb) or (None, None) on failure. """ - import glob as _glob - if platform.system() != "Linux": return None, None try: - used_files = _glob.glob("/sys/class/drm/card*/device/mem_info_vram_used") - total_files = _glob.glob("/sys/class/drm/card*/device/mem_info_vram_total") + used_files = glob.glob("/sys/class/drm/card*/device/mem_info_vram_used") + total_files = glob.glob("/sys/class/drm/card*/device/mem_info_vram_total") if not used_files or not total_files: return None, None used_bytes = sum(int(open(f).read().strip()) for f in used_files) @@ -618,8 +609,6 @@ def _rocm_windows_perf_counter_vram_gb() -> tuple[Optional[float], Optional[floa usage accurately. Works for any GPU vendor without amd-smi or nvidia-smi. Returns (used_gb, total_gb) or (None, None) on failure. """ - import subprocess as _sp - if platform.system() != "Windows": return None, None try: @@ -628,7 +617,7 @@ def _rocm_windows_perf_counter_vram_gb() -> tuple[Optional[float], Optional[floa " -ErrorAction SilentlyContinue).CounterSamples;" "if($s){($s|Measure-Object CookedValue -Sum).Sum}else{-1}" ) - r = _sp.run( + r = subprocess.run( ["powershell", "-NoProfile", "-NonInteractive", "-Command", ps], capture_output = True, text = True, @@ -1160,17 +1149,12 @@ def _load_config_for_gpu_estimate(model_name: str, hf_token: Optional[str] = Non def _determine_attention_impl_for_gpu_estimate(config) -> str: - import copy as _copy - # torch.distributed is incomplete on Windows ROCm — torch._C is a C # extension (not a package), so Python cannot import the submodule # torch._C._distributed_c10d that torch.distributed depends on. # Inject an empty stub into sys.modules BEFORE importing torch.distributed # so the import succeeds, then patch the missing process-group helpers. - import sys as _sys - import types as _types - - if _sys.platform == "win32" and IS_ROCM: + if sys.platform == "win32" and IS_ROCM: # Dummy class for any name torch.distributed tries to import from these stubs class _Dummy: pass @@ -1180,8 +1164,8 @@ def _determine_attention_impl_for_gpu_estimate(config) -> str: "torch._C._distributed_autograd", "torch._C._distributed_rpc", ): - if _c10d_name not in _sys.modules: - _stub = _types.ModuleType(_c10d_name) + if _c10d_name not in sys.modules: + _stub = types.ModuleType(_c10d_name) # torch.distributed imports these names from _distributed_c10d; # provide no-op dummies so the import doesn't raise AttributeError. for _sym in ( @@ -1200,7 +1184,7 @@ def _determine_attention_impl_for_gpu_estimate(config) -> str: "BuiltinCommHookType", ): setattr(_stub, _sym, _Dummy) - _sys.modules[_c10d_name] = _stub + sys.modules[_c10d_name] = _stub try: import torch.distributed as _td @@ -1225,7 +1209,7 @@ def _determine_attention_impl_for_gpu_estimate(config) -> str: # `sub_configs` and propagates to nested text_config / sub-configs, so a # shallow copy still mutates those shared inner objects on the cached # config returned by _load_config_for_gpu_estimate. Deepcopy isolates them. - config_copy = _copy.deepcopy(config) + config_copy = copy.deepcopy(config) model_class = None for auto_model in (AutoModelForCausalLM, AutoModel): @@ -2026,8 +2010,6 @@ def safe_num_proc(desired: Optional[int] = None) -> int: Returns: A safe integer ≥ 1. """ - import sys - # Windows and macOS use 'spawn' for multiprocessing -- the overhead of # re-importing torch/transformers/unsloth per worker is typically slower # than single-process. @@ -2078,8 +2060,6 @@ def dataset_map_num_proc(desired: Optional[int] = None) -> Optional[int]: ``datasets`` treats ``num_proc=1`` as multiprocessing (creates ``Pool(1)``). Only ``num_proc=None`` guarantees in-process execution. """ - import sys - if sys.platform in ("win32", "darwin"): return None return safe_num_proc(desired) diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 786b668e7b..5e785c934d 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -12,8 +12,10 @@ PATH to point at the venv. from __future__ import annotations +import glob import os import platform +import re import shutil import subprocess import sys @@ -209,8 +211,6 @@ def _detect_rocm_version() -> tuple[int, int] | None: # for the rocm-core package version. Matches the chain in # install.sh::get_torch_index_url so `unsloth studio update` behaves # the same as a fresh `curl | sh` install. - import re as _re_pkg - for cmd in ( ["dpkg-query", "-W", "-f=${Version}\n", "rocm-core"], ["rpm", "-q", "--qf", "%{VERSION}\n", "rocm-core"], @@ -232,8 +232,8 @@ def _detect_rocm_version() -> tuple[int, int] | None: continue raw = result.stdout.strip() # dpkg can prepend an epoch ("1:6.3.0-1"); strip it before parsing. - raw = _re_pkg.sub(r"^\d+:", "", raw) - m = _re_pkg.match(r"(\d+)[.-](\d+)", raw) + raw = re.sub(r"^\d+:", "", raw) + m = re.match(r"(\d+)[.-](\d+)", raw) if m: return int(m.group(1)), int(m.group(2)) @@ -274,8 +274,6 @@ def _detect_windows_gfx_arch() -> str | None: enumeration order) and HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES selects which one to install for. The first GPU is used when no env var is set. """ - import re - # 1. Explicit override (matches PowerShell installer's env-var path). _override = os.environ.get("UNSLOTH_ROCM_GFX_ARCH") if _override and _override.strip(): @@ -366,9 +364,7 @@ def _detect_bnb_rocm_dll_ver() -> str | None: bitsandbytes — uses importlib.util.find_spec so it is safe to call before BNB is imported. """ - import glob import importlib.util - import re spec = importlib.util.find_spec("bitsandbytes") if spec is None or not spec.submodule_search_locations: @@ -387,8 +383,6 @@ def _detect_bnb_rocm_dll_ver() -> str | None: def _has_rocm_gpu() -> bool: """Return True only if an actual AMD GPU is visible (not just ROCm tools installed).""" - import re - for cmd, check_fn in ( # rocminfo: look for a real gfx GPU id (3-4 chars, nonzero first digit). # gfx000 is the CPU agent; ROCm 6.1+ also emits generic ISA lines like @@ -469,7 +463,6 @@ def _detect_amd_gfx_codes() -> list[str]: amd-smi but no rocminfo. Returns an empty list when no probe yields a gfx target. """ - import re def _extract(text: str) -> list[str]: codes = re.findall(r"gfx([1-9][0-9a-z]{2,3})", text.lower()) @@ -510,28 +503,24 @@ def _install_bnb_windows_rocm() -> bool: The continuous-release wheel is intentionally mismatched: the filename encodes version 1.33.7.preview (parsed as 1.33.7rc0 by PEP 440) while the wheel metadata reports 0.50.0.dev0. uv rejects this filename/metadata - mismatch; set UV_SKIP_WHEEL_FILENAME_CHECK=1 to bypass that check, then - restore the previous value (or remove the var) when done. + mismatch -- and bypassing it with UV_SKIP_WHEEL_FILENAME_CHECK still leaves + uv mangling the bitsandbytes install. Per the AMD install guide + (https://unsloth.ai/docs/get-started/install/amd/amd-hackathon) the wheel + must be installed with plain pip, not uv, so we force pip here + (force_pip=True). plain pip performs no wheel filename/metadata check. """ _bnb_win_url = _BNB_ROCM_PRERELEASE_URLS.get("win_amd64") if _bnb_win_url is None: return False - _old = os.environ.get("UV_SKIP_WHEEL_FILENAME_CHECK") - os.environ["UV_SKIP_WHEEL_FILENAME_CHECK"] = "1" - try: - _ok = pip_install_try( - "bitsandbytes (AMD Windows, pre-release main)", - "--force-reinstall", - "--no-cache-dir", - "--no-deps", - _bnb_win_url, - constrain = False, - ) - finally: - if _old is None: - os.environ.pop("UV_SKIP_WHEEL_FILENAME_CHECK", None) - else: - os.environ["UV_SKIP_WHEEL_FILENAME_CHECK"] = _old + _ok = pip_install_try( + "bitsandbytes (AMD Windows, pre-release main)", + "--force-reinstall", + "--no-cache-dir", + "--no-deps", + _bnb_win_url, + constrain = False, + force_pip = True, + ) if not _ok: return False # After install: detect the actual ROCm DLL suffix shipped in the wheel and diff --git a/studio/setup.ps1 b/studio/setup.ps1 index e276b2151d..6c5f4171c4 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -818,11 +818,11 @@ if (-not $HasNvidiaSmi) { # Ordered most-specific first; first match wins. elseif ($ROCmGpuLabel) { $nameArchTable = @( - @{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4 + @{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4 @{ P = "9070|9060"; A = "gfx1200" } # RDNA 4 - @{ P = "8060S|890M|Strix Halo|HX 37[05]|HX 38[05]|AI 9 HX"; A = "gfx1151" } # RDNA 3.5 iGPU (Strix Halo / Radeon 8060S retail) - @{ P = "880M|Strix Point|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]"; A = "gfx1150" } # RDNA 3.5 iGPU (Strix Point) - @{ P = "RX 7900|RX 7800|RX 7700(?! S)"; A = "gfx1100" } # RDNA 3 desktop + @{ P = "8060S|890M|Strix Halo|HX 37[05]|HX 38[05]|AI 9 HX"; A = "gfx1151" } # RDNA 3.5 iGPU (Strix Halo / Radeon 8060S retail) + @{ P = "880M|Strix Point|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]"; A = "gfx1150" } # RDNA 3.5 iGPU (Strix Point) + @{ P = "RX 7900|RX 7800|RX 7700(?! S)"; A = "gfx1100" } # RDNA 3 desktop @{ P = "RX 7600"; A = "gfx1102" } # RDNA 3 @{ P = "780M|760M|740M|Phoenix"; A = "gfx1103" } # RDNA 3 iGPU (Phoenix) ) @@ -2037,10 +2037,10 @@ if ($HasROCm -and $CuTag -eq "cpu") { $amdIndexBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.amd.com/rocm/whl" } $archFamilyMap = @{ "gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all" # RDNA 4 - "gfx1151" = "gfx1151"; "gfx1150" = "gfx1150" # RDNA 3.5 (Strix Halo/Point) - "gfx1103" = "gfx110X-all"; "gfx1102" = "gfx110X-all" # RDNA 3 + "gfx1151" = "gfx1151"; "gfx1150" = "gfx1150" # RDNA 3.5 (Strix Halo/Point) + "gfx1103" = "gfx110X-all"; "gfx1102" = "gfx110X-all" # RDNA 3 "gfx1101" = "gfx110X-all"; "gfx1100" = "gfx110X-all" - "gfx90a" = "gfx90a"; "gfx908" = "gfx908" # MI200/MI100 + "gfx90a" = "gfx90a"; "gfx908" = "gfx908" # MI200/MI100 } # gfx120X and Strix have a null _grouped_mm kernel on torch <2.11.0. # Mirrors the $torchFloorMap in install.ps1 so both installers enforce diff --git a/studio/setup.sh b/studio/setup.sh index 3b263a08da..89489f38b0 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -707,13 +707,13 @@ elif [ "$_setup_amd_detected" = true ]; then # Name-based arch inference when tools don't report gfx (mirrors setup.ps1 nameArchTable) elif [ -z "$_setup_gfx" ] && [ -n "$_setup_mkt" ]; then case "$_setup_mkt" in - *"9070 XT"*|*9080*) _setup_gfx="gfx1201" ;; # RDNA 4 - *9070*|*9060*) _setup_gfx="gfx1200" ;; # RDNA 4 + *"9070 XT"*|*9080*) _setup_gfx="gfx1201" ;; # RDNA 4 + *9070*|*9060*) _setup_gfx="gfx1200" ;; # RDNA 4 *"8060S"*|*"890M"*|*"Strix Halo"*|*"HX 37"*|*"HX 38"*|*"AI 9 HX"*) _setup_gfx="gfx1151" ;; # RDNA 3.5 iGPU - *"880M"*|*"Strix Point"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*) _setup_gfx="gfx1150" ;; # RDNA 3.5 iGPU - *"RX 7900"*|*"RX 7800"*|*"RX 7700"*) _setup_gfx="gfx1100" ;; # RDNA 3 desktop - *"RX 7600"*) _setup_gfx="gfx1102" ;; # RDNA 3 - *"780M"*|*"760M"*|*"740M"*|*"Phoenix"*) _setup_gfx="gfx1103" ;; # RDNA 3 iGPU + *"880M"*|*"Strix Point"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*) _setup_gfx="gfx1150" ;; # RDNA 3.5 iGPU + *"RX 7900"*|*"RX 7800"*|*"RX 7700"*) _setup_gfx="gfx1100" ;; # RDNA 3 desktop + *"RX 7600"*) _setup_gfx="gfx1102" ;; # RDNA 3 + *"780M"*|*"760M"*|*"740M"*|*"Phoenix"*) _setup_gfx="gfx1103" ;; # RDNA 3 iGPU esac if [ -n "$_setup_gfx" ]; then substep "gfx arch inferred from GPU name: $_setup_gfx" diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index ecb1b9be15..ec99e9a8d9 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -1096,6 +1096,12 @@ class TestLiveRegression: # Load worker.py module _WORKER_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "training" / "worker.py" +_EXPORT_WORKER_PATH = ( + PACKAGE_ROOT / "studio" / "backend" / "core" / "export" / "worker.py" +) +# The torchao Windows-ROCm stub was de-duplicated out of the export/training +# workers into a shared module; both workers now call into it. +_TORCHAO_STUB_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "_torchao_stub.py" # The wheel-probe subprocess was hoisted out of worker.py into wheel_utils # during the wheel-resolver refactor; the probe script literal lives there. _WHEEL_UTILS_PATH = PACKAGE_ROOT / "studio" / "backend" / "utils" / "wheel_utils.py" @@ -1654,53 +1660,46 @@ class TestInstallBnbWindowsRocm: """Verify AMD Windows BNB wheel install helper.""" def test_calls_pip_install_try_with_win_amd64_url(self): - """Should call pip_install_try with the win_amd64 wheel URL.""" + """Should call pip_install_try with the win_amd64 wheel URL via plain pip.""" with patch.object(stack_mod, "pip_install_try", return_value = True) as mock_pip: stack_mod._install_bnb_windows_rocm() assert mock_pip.call_count == 1 call_args = str(mock_pip.call_args_list[0]) assert "bitsandbytes" in call_args assert "win_amd64" in call_args + # Must force plain pip (uv mangles the bitsandbytes wheel) -- see + # https://unsloth.ai/docs/get-started/install/amd/amd-hackathon + assert mock_pip.call_args.kwargs.get("force_pip") is True - def test_sets_uv_skip_env_var_during_install(self): - """UV_SKIP_WHEEL_FILENAME_CHECK must be '1' when pip_install_try runs.""" + def test_forces_plain_pip_not_uv(self): + """The bnb wheel must be installed with plain pip, never uv.""" + with patch.object(stack_mod, "pip_install_try", return_value = True) as mock_pip: + stack_mod._install_bnb_windows_rocm() + assert mock_pip.call_args.kwargs.get("force_pip") is True + + def test_does_not_touch_uv_skip_env_var(self): + """The UV_SKIP_WHEEL_FILENAME_CHECK hack is gone; the env must be untouched.""" observed = {} def _capture(*args, **kwargs): - observed["val"] = os.environ.get("UV_SKIP_WHEEL_FILENAME_CHECK") + observed["during"] = os.environ.get("UV_SKIP_WHEEL_FILENAME_CHECK") return True - with patch.object(stack_mod, "pip_install_try", side_effect = _capture): - stack_mod._install_bnb_windows_rocm() - assert observed.get("val") == "1" - - def test_restores_uv_skip_env_var_after_install(self): - """UV_SKIP_WHEEL_FILENAME_CHECK should be removed after install if it wasn't set before.""" with patch.dict(os.environ, {}, clear = False): os.environ.pop("UV_SKIP_WHEEL_FILENAME_CHECK", None) - with patch.object(stack_mod, "pip_install_try", return_value = True): + with patch.object(stack_mod, "pip_install_try", side_effect = _capture): stack_mod._install_bnb_windows_rocm() + assert observed.get("during") is None assert "UV_SKIP_WHEEL_FILENAME_CHECK" not in os.environ - def test_restores_previous_uv_skip_value(self): - """If UV_SKIP_WHEEL_FILENAME_CHECK was already set, restore it afterwards.""" - with patch.dict(os.environ, {"UV_SKIP_WHEEL_FILENAME_CHECK": "0"}): - with patch.object(stack_mod, "pip_install_try", return_value = True): - stack_mod._install_bnb_windows_rocm() - assert os.environ.get("UV_SKIP_WHEEL_FILENAME_CHECK") == "0" - - def test_restores_env_even_if_install_raises(self): - """UV_SKIP_WHEEL_FILENAME_CHECK must be cleaned up even on pip failure.""" + def test_returns_false_on_pip_failure(self): + """A failed pip_install_try must surface as a False return, not BNB_ROCM_VERSION.""" with patch.dict(os.environ, {}, clear = False): - os.environ.pop("UV_SKIP_WHEEL_FILENAME_CHECK", None) - with patch.object( - stack_mod, "pip_install_try", side_effect = RuntimeError("pip failed") - ): - try: - stack_mod._install_bnb_windows_rocm() - except RuntimeError: - pass - assert "UV_SKIP_WHEEL_FILENAME_CHECK" not in os.environ + os.environ.pop("BNB_ROCM_VERSION", None) + with patch.object(stack_mod, "pip_install_try", return_value = False): + result = stack_mod._install_bnb_windows_rocm() + assert result is False + assert "BNB_ROCM_VERSION" not in os.environ def test_no_op_when_win_amd64_url_missing(self): """Should be silent no-op if win_amd64 key absent from _BNB_ROCM_PRERELEASE_URLS.""" @@ -1906,24 +1905,34 @@ class TestWorkerWindowsRocmPatches: assert "offs_list" in source assert "offs.tolist()" in source + def test_worker_calls_shared_torchao_stub(self): + """worker.py must invoke the shared torchao stub entrypoint.""" + source = _WORKER_PATH.read_text(encoding = "utf-8") + assert "install_torchao_windows_rocm_stub()" in source + + def test_export_worker_calls_shared_torchao_stub(self): + """export/worker.py must invoke the same shared torchao stub entrypoint.""" + source = _EXPORT_WORKER_PATH.read_text(encoding = "utf-8") + assert "install_torchao_windows_rocm_stub()" in source + def test_torchao_stub_uses_stub_type_meta(self): """Torchao stub must use _StubTypeMeta so isinstance() returns False not TypeError.""" - source = _WORKER_PATH.read_text(encoding = "utf-8") + source = _TORCHAO_STUB_PATH.read_text(encoding = "utf-8") assert "_StubTypeMeta" in source def test_stub_type_meta_has_instancecheck(self): """_StubTypeMeta must define __instancecheck__ returning False.""" - source = _WORKER_PATH.read_text(encoding = "utf-8") + source = _TORCHAO_STUB_PATH.read_text(encoding = "utf-8") assert "__instancecheck__" in source def test_stub_subpackage_finder_registered(self): """_StubSubpackageFinder must be appended to sys.meta_path.""" - source = _WORKER_PATH.read_text(encoding = "utf-8") + source = _TORCHAO_STUB_PATH.read_text(encoding = "utf-8") assert "sys.meta_path.append(_StubSubpackageFinder())" in source def test_torchao_key_submodules_pre_stubbed(self): """Key torchao submodules (dtypes, quantization) must be pre-stubbed.""" - source = _WORKER_PATH.read_text(encoding = "utf-8") + source = _TORCHAO_STUB_PATH.read_text(encoding = "utf-8") assert "torchao.dtypes" in source assert "torchao.quantization" in source