ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge branch 'main' into studio-composer

Browse source
This commit is contained in:
Michael Han 2026-05-30 05:11:22 -07:00 • committed by GitHub
commit a879ca1d82
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 1179 additions and 364 deletions
Download patch file Download diff file Expand all files Collapse all files

50
.github/workflows/lint-ci.yml vendored
View file

@ -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/<base> is reliably created (a bare
# `git fetch origin <ref>` 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

854
scripts/verify_import_hoist.py Normal file
View file

@ -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] <file>... # 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", "<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}.<lambda>", 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}.<comp>", 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())

142
studio/backend/core/_torchao_stub.py Normal file
View file

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

98
studio/backend/core/export/worker.py
View file

@ -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:

10
studio/backend/core/inference/llama_cpp.py
View file

@ -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/<pid>/exe on Linux when psutil is
not installed.
"""
import os
import signal
import sys
try:
# -- Build the ownership allowlist --------------------------------
# Two kinds of matches:

47
studio/backend/core/training/trainer.py
View file

@ -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()

132
studio/backend/core/training/worker.py
View file

@ -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)),

2
studio/backend/main.py
View file

@ -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")

60
studio/backend/utils/hardware/hardware.py
View file

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

47
studio/install_python_stack.py
View file

@ -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

14
studio/setup.ps1
View file

@ -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

12
studio/setup.sh
View file

@ -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"

75
tests/studio/install/test_rocm_support.py
View file

@ -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