scripts/lint_risky_loaders.py
Rama main · réplica de lectura del repositorio original.
#!/usr/bin/env python3
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""Fail CI on a new call site that loads code or objects in a way a model repo can steer.
`lint_exec_literals.py` covers `exec`, `eval` and `compile`. These are the other
shapes behind holes fixed here since July 2026, each as one rule:
dynamic-import import_module / __import__ / import_from_string /
get_class_from_dynamic_module / spec_from_file_location /
runpy / pydoc.locate handed a name that is not written out
(sentence-transformers modules.json, mlx `model_type`).
trust-remote-code `trust_remote_code` switched on in code, rather than by the
caller: a default, an assignment, a dict entry or a literal
keyword (the Studio defaults that ran Hub code unasked).
unpinned-code-fetch Python fetched from a branch head (`.../main/x.py`), a
`git clone` with no checkout of a fixed revision, or a Hub
download with no `revision` in a function that then puts
the result on `sys.path` or loads a file from it (Spark-TTS).
unsafe-deserialize pickle / dill / joblib / marshal / shelve, `torch.load` with
`weights_only` not left at the safe default, `np.load` with
`allow_pickle=True`, `yaml.load` without a safe loader.
Every rule is one AST pass and a few set lookups, so the whole tree takes seconds.
Existing call sites are recorded in a baseline beside this script with a reason
each, keyed on the call's own text rather than its line, with a count, exactly as
in `lint_exec_literals.py`: only a new site fails.
python scripts/lint_risky_loaders.py # check, exit 1 on a new site
python scripts/lint_risky_loaders.py --update # rewrite the baseline
python scripts/lint_risky_loaders.py --self-test # prove the rules still fire
"""
from __future__ import annotations
import argparse
import ast
import hashlib
import json
import re
import sys
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parents[1]
BASELINE_PATH = Path(__file__).resolve().parent / "risky_loaders_baseline.json"
EXCLUDED_PARTS = frozenset(
{
"tests",
"node_modules",
"build",
"dist",
".venv",
"venv",
"site-packages",
".git",
".tox",
".mypy_cache",
".pytest_cache",
"__pycache__",
".ipynb_checkpoints",
".eggs",
"unsloth_compiled_cache",
}
)
DYNAMIC_IMPORTS = {
"importlib.import_module": "import_module",
"importlib.__import__": "__import__",
"importlib.util.spec_from_file_location": "spec_from_file_location",
"importlib.machinery.SourceFileLoader": "SourceFileLoader",
"runpy.run_path": "runpy.run_path",
"runpy.run_module": "runpy.run_module",
"pydoc.locate": "pydoc.locate",
"pkgutil.resolve_name": "pkgutil.resolve_name",
}
# Matched on the last name alone: distinctive enough, and reached through several
# import paths (`transformers.dynamic_module_utils`, `sentence_transformers.util`).
DYNAMIC_IMPORT_TAILS = {
"import_from_string": "import_from_string",
"get_class_from_dynamic_module": "get_class_from_dynamic_module",
"get_class_in_module": "get_class_in_module",
}
PICKLE_LIKE = {
f"{module}.{function}"
for module in ("pickle", "_pickle", "cPickle", "dill", "cloudpickle")
for function in ("load", "loads", "Unpickler")
} | {
"joblib.load",
"marshal.load",
"marshal.loads",
"shelve.open",
"pandas.read_pickle",
"yaml.unsafe_load",
"yaml.unsafe_load_all",
}
SAFE_YAML_LOADERS = {"SafeLoader", "CSafeLoader", "BaseLoader", "CBaseLoader"}
HUB_DOWNLOADS = {"snapshot_download", "hf_hub_download"}
CODE_LOADERS = {
"sys.path.insert",
"sys.path.append",
"importlib.util.spec_from_file_location",
"importlib.machinery.SourceFileLoader",
"runpy.run_path",
}
BRANCH_HEAD_PY = re.compile(
r"(?:raw\.githubusercontent\.com/[^/\s]+/[^/\s]+/(?:refs/heads/)?(?:main|master)/"
r"|github\.com/[^/\s]+/[^/\s]+/raw/(?:refs/heads/)?(?:main|master)/)\S*\.py\b"
)
PIN_WORDS = {"checkout", "--revision", "reset"}
def _imports(tree: ast.AST) -> dict:
"""Local name -> dotted origin, so aliases (`import importlib as il`) resolve."""
table = {}
for node in ast.walk(tree):
if isinstance(node, ast.ImportFrom) and node.module and not node.level:
for alias in node.names:
table[alias.asname or alias.name] = f"{node.module}.{alias.name}"
elif isinstance(node, ast.Import):
for alias in node.names:
if alias.asname:
table[alias.asname] = alias.name
else:
head = alias.name.split(".")[0]
table[head] = head
return table
def _dotted(node: ast.AST) -> str:
parts = []
while isinstance(node, ast.Attribute):
parts.append(node.attr)
node = node.value
if not isinstance(node, ast.Name):
return ""
parts.append(node.id)
return ".".join(reversed(parts))
def _qualified(func: ast.AST, table: dict) -> str:
name = _dotted(func)
if not name:
return ""
head, _, rest = name.partition(".")
if head == "__import__":
return "importlib.__import__"
if head in table:
return table[head] + ("." + rest if rest else "")
return name
def _is_written_out(node: ast.AST) -> bool:
if isinstance(node, ast.Constant):
return isinstance(node.value, (str, bytes))
if isinstance(node, ast.JoinedStr):
return not any(isinstance(part, ast.FormattedValue) for part in node.values)
if isinstance(node, ast.BinOp) and isinstance(node.op, ast.Add):
return _is_written_out(node.left) and _is_written_out(node.right)
return False
def _is_true(node: ast.AST) -> bool:
return isinstance(node, ast.Constant) and node.value is True
def _keyword(call: ast.Call, name: str):
for keyword in call.keywords:
if keyword.arg == name:
return keyword.value
return None
def _first_argument(call: ast.Call, keyword: str):
if call.args:
return call.args[0]
return _keyword(call, keyword)
# Sinks whose second argument is the file that gets executed, and its keyword.
_PATH_KEYWORDS = {
"spec_from_file_location": "location",
"SourceFileLoader": "path",
"get_class_in_module": "module_path",
}
_TARGET_KEYWORDS = {
"runpy.run_path": "path_name",
"runpy.run_module": "mod_name",
"pydoc.locate": "path",
}
def _dynamic_import(call: ast.Call, qualified: str):
sink = DYNAMIC_IMPORTS.get(qualified) or DYNAMIC_IMPORT_TAILS.get(qualified.split(".")[-1])
if sink is None:
return None
if sink in _PATH_KEYWORDS:
keyword = _PATH_KEYWORDS[sink]
target = call.args[1] if len(call.args) > 1 else _keyword(call, keyword)
else:
target = _first_argument(call, _TARGET_KEYWORDS.get(sink, "name"))
if target is None or _is_written_out(target):
return None
return sink
def _unsafe_deserialize(call: ast.Call, qualified: str):
if qualified in PICKLE_LIKE:
return qualified
if qualified == "torch.load":
# Omitted is unsafe too: torch < 2.6, still supported, defaults to weights_only=False.
if not _is_true(_keyword(call, "weights_only")):
return "torch.load"
elif qualified == "numpy.load":
if _is_true(_keyword(call, "allow_pickle")):
return "numpy.load"
elif qualified in ("yaml.load", "yaml.load_all"):
loader = _keyword(call, "Loader")
if loader is None and len(call.args) > 1:
loader = call.args[1]
if loader is None or _dotted(loader).split(".")[-1] not in SAFE_YAML_LOADERS:
return qualified
return None
def _true_field_default(value: ast.AST) -> bool:
"""`Field(True, ...)` / `field(default = True)`: the schema default Studio's request models use."""
if not (
isinstance(value, ast.Call) and _dotted(value.func).split(".")[-1] in ("Field", "field")
):
return False
return (bool(value.args) and _is_true(value.args[0])) or _is_true(_keyword(value, "default"))
def _trust_remote_code(node: ast.AST):
"""`trust_remote_code` turned on by the code itself."""
if isinstance(node, ast.Call):
if _is_true(_keyword(node, "trust_remote_code")):
return "keyword"
func = node.func
if (
isinstance(func, ast.Attribute)
and func.attr == "setdefault"
and len(node.args) == 2
and isinstance(node.args[0], ast.Constant)
and node.args[0].value == "trust_remote_code"
and _is_true(node.args[1])
):
return "setdefault"
elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)):
arguments = node.args
positional = arguments.posonlyargs + arguments.args
pairs = list(
zip(positional[len(positional) - len(arguments.defaults) :], arguments.defaults)
)
pairs += [
(a, d) for a, d in zip(arguments.kwonlyargs, arguments.kw_defaults) if d is not None
]
if any(a.arg == "trust_remote_code" and _is_true(d) for a, d in pairs):
return "default"
elif (
isinstance(node, (ast.Assign, ast.AnnAssign))
and node.value is not None
and (_is_true(node.value) or _true_field_default(node.value))
):
targets = node.targets if isinstance(node, ast.Assign) else [node.target]
for target in targets:
if isinstance(target, (ast.Name, ast.Attribute)) and (
getattr(target, "id", None) == "trust_remote_code"
or getattr(target, "attr", None) == "trust_remote_code"
):
return "assignment"
if (
isinstance(target, ast.Subscript)
and isinstance(target.slice, ast.Constant)
and target.slice.value == "trust_remote_code"
):
return "item"
elif isinstance(node, ast.Dict):
for key, value in zip(node.keys, node.values):
if (
isinstance(key, ast.Constant)
and key.value == "trust_remote_code"
and _is_true(value)
):
return "dict"
return None
def _git_words(command) -> list:
"""The written-out words of a git command, from an argv list or a shell string; [] otherwise."""
if isinstance(command, (ast.List, ast.Tuple)):
first = command.elts[0] if command.elts else None
if not (isinstance(first, ast.Constant) and first.value == "git"):
return []
words = [
e.value
for e in command.elts
if isinstance(e, ast.Constant) and isinstance(e.value, str)
]
else:
if isinstance(command, ast.JoinedStr) and command.values:
command = command.values[0]
if (
not (isinstance(command, ast.Constant) and isinstance(command.value, str))
or "git " not in command.value
):
return []
# The first git command in a chain, so `cd llama.cpp && git reset --hard X` counts.
segments = (part.split() for part in re.split(r"&&|;|\|\|", command.value))
words = next((w for w in segments if w[:1] == ["git"]), [])
return words if words[:1] == ["git"] else []
def _unpinned_fetches(calls: list, git_commands: list) -> list:
"""Per function: a clone with no pin, or a revision-less Hub download feeding a code loader."""
found = []
loads_code = any(qualified in CODE_LOADERS for _, qualified in calls)
for call, qualified in calls:
if (
qualified.split(".")[-1] in HUB_DOWNLOADS
and loads_code
and _keyword(call, "revision") is None
):
found.append((call, "hub-download-loaded-as-code"))
# A pin is a git command that moves to a revision, not any string mentioning one.
if not any(set(words) & PIN_WORDS for _, words in git_commands):
found += [
(node, "git-clone-unpinned") for node, words in git_commands if words[1:2] == ["clone"]
]
return found
def _key(relative: str, rule: str, sink: str, node: ast.AST) -> dict:
return {
"file": relative,
"rule": rule,
"sink": sink,
"digest": hashlib.sha256(ast.unparse(node).encode()).hexdigest()[:16],
"line": getattr(node, "lineno", 0),
}
_FUNCTIONS = (ast.FunctionDef, ast.AsyncFunctionDef)
_TRUST_SHAPES = (ast.Call, ast.Lambda, ast.Assign, ast.AnnAssign, ast.Dict) + _FUNCTIONS
def _notebook_source(path: Path, relative: str) -> str:
"""A notebook's code cells as one module; `%magic` / `!command` lines are blanked to keep line numbers."""
try:
document = json.loads(path.read_text(encoding = "utf-8"))
except (OSError, UnicodeDecodeError, json.JSONDecodeError) as error:
raise SystemExit(f"{relative}: could not be read ({error.__class__.__name__})")
lines = []
for cell in document.get("cells") or []:
if cell.get("cell_type") != "code":
continue
source = cell.get("source") or []
for line in (source if isinstance(source, str) else "".join(source)).splitlines():
lines.append("" if line.lstrip().startswith(("%", "!")) else line)
lines.append("")
return "\n".join(lines)
def scan_file(path: Path, relative: str) -> list:
source = _notebook_source(path, relative) if path.suffix == ".ipynb" else path.read_bytes()
try:
tree = ast.parse(source, filename = str(path))
except (SyntaxError, ValueError, MemoryError, RecursionError) as error:
if path.suffix == ".ipynb":
return []
# Unparsed is unchecked; reporting it clean would be a bypass.
raise SystemExit(f"{relative}: could not be parsed ({error.__class__.__name__})")
found = []
imports = []
calls_by_owner: dict = {}
git_by_owner: dict = {}
covered = set()
stack = [(tree, tree)]
while stack:
node, owner = stack.pop()
for child in ast.iter_child_nodes(node):
if isinstance(child, _TRUST_SHAPES):
shape = _trust_remote_code(child)
if shape:
subject = child.args if isinstance(child, _FUNCTIONS) else child
entry = _key(relative, "trust-remote-code", shape, subject)
entry["line"] = child.lineno
found.append(entry)
if isinstance(child, ast.Call):
calls_by_owner.setdefault(owner, []).append(child)
elif isinstance(child, _FUNCTIONS):
stack.append((child, child))
continue
elif isinstance(child, (ast.Import, ast.ImportFrom)):
imports.append(child)
elif isinstance(child, ast.Constant) and isinstance(child.value, str):
if "/ma" in child.value and BRANCH_HEAD_PY.search(child.value):
found.append(_key(relative, "unpinned-code-fetch", "branch-head-url", child))
if (
isinstance(child, (ast.List, ast.Tuple, ast.JoinedStr, ast.Constant))
and id(child) not in covered
):
words = _git_words(child)
if words:
git_by_owner.setdefault(owner, []).append((child, words))
if isinstance(child, ast.JoinedStr):
covered.add(id(child.values[0]))
stack.append((child, owner))
table = _imports(ast.Module(body = imports, type_ignores = []))
for owner in list(calls_by_owner) + [o for o in git_by_owner if o not in calls_by_owner]:
calls = [(call, _qualified(call.func, table)) for call in calls_by_owner.get(owner, ())]
for call, qualified in calls:
sink = _dynamic_import(call, qualified)
if sink:
found.append(_key(relative, "dynamic-import", sink, call))
sink = _unsafe_deserialize(call, qualified)
if sink:
found.append(_key(relative, "unsafe-deserialize", sink, call))
for call, sink in _unpinned_fetches(calls, git_by_owner.get(owner, [])):
found.append(_key(relative, "unpinned-code-fetch", sink, call))
return found
def collect(targets: list) -> list:
found = []
seen = set()
for target in targets:
root = REPO_ROOT / target
if root.is_file() and root.suffix in (".py", ".ipynb"):
paths = [root]
elif root.is_dir():
paths = sorted(q for suffix in ("*.py", "*.ipynb") for q in root.rglob(suffix))
else:
raise SystemExit(
f"{target}: scan target does not exist, so nothing under it was checked"
)
for path in paths:
# This file's self-test sources spell out every bad shape on purpose.
if not path.is_file() or path.resolve() == Path(__file__).resolve():
continue
relative = _relative(path)
if relative in seen or EXCLUDED_PARTS & set(Path(relative).parts):
continue
seen.add(relative)
found.extend(scan_file(path, relative))
return found
def _relative(path: Path) -> str:
try:
return path.resolve().relative_to(REPO_ROOT).as_posix()
except ValueError:
return path.as_posix()
def _identity(entry: dict) -> tuple:
return (entry["file"], entry["rule"], entry["sink"], entry["digest"])
def _counted(entries: list) -> dict:
counts: dict = {}
for entry in entries:
counts[_identity(entry)] = counts.get(_identity(entry), 0) + 1
return counts
def main() -> int:
parser = argparse.ArgumentParser(
description = __doc__, formatter_class = argparse.RawDescriptionHelpFormatter
)
parser.add_argument("--update", action = "store_true", help = "rewrite the baseline")
parser.add_argument("--self-test", action = "store_true", help = "check the rules still fire")
parser.add_argument("--paths", nargs = "*", help = "scan these instead of the defaults")
arguments = parser.parse_args()
if arguments.self_test:
return self_test()
if arguments.update and arguments.paths:
parser.error("--update rewrites the whole baseline; run it without --paths")
document = json.loads(BASELINE_PATH.read_text(encoding = "utf-8"))
found = collect(arguments.paths or document["targets"])
baseline_rel = BASELINE_PATH.relative_to(REPO_ROOT).as_posix()
if arguments.update:
# Carry reasons over so a reviewed entry never becomes unreviewed.
reasons = {_identity(e): e.get("reason", "") for e in document["entries"]}
document["entries"] = sorted(
(
{
"file": f,
"rule": r,
"sink": s,
"digest": d,
"count": n,
"reason": reasons.get((f, r, s, d), "REVIEW ME"),
}
for (f, r, s, d), n in _counted(found).items()
),
key = _identity,
)
BASELINE_PATH.write_text(json.dumps(document, indent = 2) + "\n", encoding = "utf-8")
print(f"baseline: {len(document['entries'])} entries, {len(found)} call sites")
return 0
entries = document["entries"]
if arguments.paths:
scopes = [_relative(REPO_ROOT / path) for path in arguments.paths]
if "." not in scopes:
entries = [
e
for e in entries
if any(e["file"] == p or e["file"].startswith(p.rstrip("/") + "/") for p in scopes)
]
allowed = {_identity(e): e["count"] for e in entries}
observed = _counted(found)
lines = {_identity(e): e["line"] for e in found}
new = sorted(k for k, n in observed.items() if n > allowed.get(k, 0))
if new:
print(f"{len(new)} risky loader call site(s) not in the baseline:\n")
for key in new:
f, r, s, _ = key
print(f" {f}:{lines[key]} [{r}] {s}")
print(
"\nEither load from a written-out name, a pinned revision or a safe format, or - if "
"the input really is trusted - record it:\n"
" 1. python scripts/lint_risky_loaders.py --update\n"
f' 2. replace the new entr(y/ies)\' "REVIEW ME" reason in {baseline_rel} with why it is safe.'
)
return 1
unreviewed = sorted(
(e["file"], e["rule"], e["sink"])
for e in entries
if e.get("reason", "REVIEW ME") in ("", "REVIEW ME")
)
if unreviewed:
print(f"{len(unreviewed)} baseline entr(y/ies) carry no justification:\n")
for f, r, s in unreviewed:
print(f" {f} [{r}] {s}")
print(f"\nSay why the input is trusted in the entry's `reason` field in {baseline_rel}.")
return 1
# Fewer calls than allowed counts too, so a removed duplicate cannot make room for a new one.
stale = sorted(k for k in allowed if observed.get(k, 0) < allowed[k])
if stale:
print(f"{len(stale)} baseline entr(y/ies) match fewer calls than recorded:\n")
for f, r, s, d in stale:
print(f" {f} [{r}] {s} {d}")
print(
"\nA risky call was removed or rewritten; run `python scripts/lint_risky_loaders.py --update`."
)
return 1
print(f"ok: {len(found)} risky loader call site(s), all recorded")
return 0
_BAD = {
"dynamic-import": """
import importlib
from importlib import import_module as im
from sentence_transformers.util import import_from_string
from transformers.dynamic_module_utils import get_class_from_dynamic_module
def f(config, name, path):
importlib.import_module(f"mlx_lm.models.{config.model_type}")
im(name)
__import__(name)
import_from_string(config["type"])
get_class_from_dynamic_module(config.auto_map["AutoModel"], name)
importlib.util.spec_from_file_location("x", path)
importlib.machinery.SourceFileLoader("x", path)
""",
"trust-remote-code": """
def load(name, trust_remote_code = True):
AutoConfig.from_pretrained(name, trust_remote_code = True)
trust_remote_code = True
kwargs["trust_remote_code"] = True
kwargs.setdefault("trust_remote_code", True)
return {"trust_remote_code": True}
class Request(BaseModel):
trust_remote_code: bool = Field(True, description = "x")
""",
"unpinned-code-fetch": """
import sys, subprocess
from huggingface_hub import snapshot_download
URL = "https://github.com/ggml-org/llama.cpp/raw/refs/heads/master/convert_hf_to_gguf.py"
def codec():
path = snapshot_download("org/codec")
sys.path.insert(0, path)
def clone():
subprocess.run(["git", "clone", "https://github.com/org/repo"])
def shell_clone(folder):
run(f"git clone https://github.com/org/repo {folder}")
def held_clone():
commands = ["git clone --recursive https://github.com/org/repo", "pip install x"]
try_execute(commands)
""",
"unsafe-deserialize": """
import pickle, torch, yaml, joblib
import numpy as np
def f(path, flag):
pickle.loads(open(path, "rb").read())
torch.load(path, weights_only = False)
torch.load(path, weights_only = flag)
torch.load(path)
np.load(path, allow_pickle = True)
yaml.load(open(path))
joblib.load(path)
""",
}
_BAD_COUNTS = {
"dynamic-import": 7,
"trust-remote-code": 7,
"unpinned-code-fetch": 5,
"unsafe-deserialize": 7,
}
_GOOD = """
import importlib, subprocess, sys, torch, yaml
import numpy as np
from huggingface_hub import snapshot_download
URL = "https://github.com/ggml-org/llama.cpp/raw/b1234/convert_hf_to_gguf.py"
def f(name, path, trust_remote_code = False):
importlib.import_module("torch.nn")
importlib.util.find_spec(name)
__import__("os")
AutoConfig.from_pretrained(name, trust_remote_code = trust_remote_code)
torch.load(path, weights_only = True)
np.load(path)
yaml.load(open(path), Loader = yaml.SafeLoader)
yaml.safe_load(open(path))
snapshot_download(name)
def chained_pin(version):
try_execute(["git clone https://github.com/org/repo", f"cd repo && git reset --hard {version}"])
def pinned():
subprocess.run(["git", "clone", "https://github.com/org/repo"])
subprocess.run(["git", "checkout", "0123abc"])
def codec():
path = snapshot_download("org/codec", revision = "0123abc")
sys.path.insert(0, path)
"""
def self_test() -> int:
"""Each rule fires on its bad shapes and stays quiet on the good ones, on a runner with nothing but Python."""
import tempfile
failures = []
with tempfile.TemporaryDirectory() as directory:
cases = [(rule, source, _BAD_COUNTS[rule], rule) for rule, source in _BAD.items()]
cases.append(("good", _GOOD, 0, None))
for label, source, expected, rule in cases:
path = Path(directory) / f"{label.replace('-', '_')}.py"
path.write_text(source, encoding = "utf-8")
hits = scan_file(path, path.name)
count = len([h for h in hits if rule is None or h["rule"] == rule])
if count != expected or (rule is not None and len(hits) != expected):
failures.append(
f"{label}: expected {expected} finding(s), got {[(h['rule'], h['sink']) for h in hits]}"
)
if failures:
print("self-test FAILED:\n " + "\n ".join(failures))
return 1
print("self-test: ok")
return 0
if __name__ == "__main__":
sys.exit(main())