unsloth/studio/backend/utils/security/file_security.py
Michael Han a00fe86c13
Studio: read model text as utf-8 so umlauts survive on Windows (#7467)
* Studio: read model text as utf-8 so umlauts survive on Windows

Chat rejects or mangles non-ASCII on Windows: "ä ö ü" in a prompt, a chat
template, or a model path comes back as mojibake, or the load dies with
UnicodeDecodeError.

open() and Path.read_text() fall back to locale.getencoding() when no encoding
is passed. On Windows that is the ANSI codepage (cp1252, cp932, cp1251, ... by
system locale), never UTF-8. Hugging Face writes these files as raw UTF-8, so
every read of one decodes with the wrong codec:

- tokenizer_config.json, which holds the chat template. Templates routinely
  carry -> arrows, smart quotes and CJK, so this is the common path into chat
- config.json and adapter_config.json
- modules.json, Ollama manifests, and the .py sources the remote-code scanner
  reads before a model is allowed to load

The llama-server and embedding-server stdout readers have the same problem via
subprocess(text = True); they now decode utf-8 with errors = "replace" so a
stray byte cannot kill a log reader.

Encoding arguments only, no logic changes.

tests/test_chat_text_encoding.py covers a config.json and a chat template
holding umlauts, arrows and CJK, plus the remote-code scanner reading a source
file with umlauts. Those pass anywhere the locale is already UTF-8, so a fourth
test re-runs the readers under -X warn_default_encoding and fails on any
platform if an encoding argument goes missing again.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: name utf-8 explicitly on the remaining text I/O, with an AST guard (#7465)

* Studio: name utf-8 explicitly on the remaining text I/O

Follow-up to the model-text reads in #7467, covering the rest of the backend:
system probes (nvidia-smi, amd-smi, powershell, git, node), package installers,
/proc and /sys readers, and internal marker files (pid, install id, bootstrap
password, Colab credentials).

Same reason as #7467. open(), Path.read_text()/write_text() and
subprocess(text = True) fall back to locale.getencoding(), which on Windows is
the ANSI codepage rather than UTF-8. These paths are mostly ASCII today, so this
is hardening, not a live bug. Encoding arguments only, no logic changes.

Adds tests/test_text_io_encoding.py: an AST guard walking every backend source
and asserting text I/O names its encoding, so the class of bug cannot creep back
in one call at a time. 275 files.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Catch aliased subprocess and positional Path.open, migrate legacy JSONL

The guard only matched a receiver literally named subprocess, so worker.py's
`import subprocess as _sp` hid three text = True installs that decode pip
output with the ANSI codepage. It also skipped any .open() with more than one
positional argument, though Path.open takes buffering/encoding/errors/newline
positionally.

Resuming a scrape written by an older release is the other half: those JSONL
lines are in the locale codepage, so the UTF-8 preload raised, the dedup keys
were silently forgotten and duplicates were appended to a now mixed-encoding
file. Decode with the locale codepage as fallback and rewrite as UTF-8 before
the append handle opens, since Windows cannot replace a file it holds open.

* Stream the JSONL preload and keep a torn line from relabelling the shard

Reading the whole shard to migrate it was wrong twice over. These files reach
gigabytes on a large scrape, so the preload now streams line by line and the
rewrite streams through a temp file.

Worse, one interrupted append used to condemn the file: the whole-file UTF-8
decode failed, every byte was retried as cp1252, and the rewrite persisted
mojibake over records that were fine. A line now counts as legacy only if the
locale codepage both decodes it and yields valid JSON, which a torn UTF-8 line
does not. Damaged lines are skipped and copied through byte for byte.

When the rewrite cannot be written at all, the append handle opens with the
legacy encoding rather than mixing UTF-8 into the file.

install_wheel takes run = subprocess.run as a parameter, so the guard cannot
see it. Both wheel installs there now name their encoding.

* Decide the shard's encoding from the file, not one line at a time

Some byte strings parse both ways. cp1251 `Р°` is D0 B0, which is also valid
UTF-8 for `а`, so a UTF-8-first parse quietly showed the wrong text instead of
migrating it.

A line now yields both readings, and the file decides. Any line that parses
under the codepage but not as UTF-8 is unambiguous evidence, and ambiguous lines
then follow that verdict, which is enough for any real shard: ordinary Cyrillic
or Japanese prose is invalid UTF-8 several times per line. Keys for ambiguous
lines are re-derived from the legacy reading during the rewrite.

A shard is undecidable only if every line is ambiguous, and nothing can tell
those apart.

latin-1 is also tried after the locale codepage, so a scrape carried from
Windows to a UTF-8 machine still has a reading rather than none. Requiring valid
JSON, not just a decode, keeps that from claiming torn lines.

* Weigh the whole shard, and never lose a record on the fallback path

One structurally valid JSON line carrying a stray 0x96 parses as cp1252, so a
single-line verdict let it relabel a healthy shard and mojibake every good
record in it. Each line with non-ASCII bytes now votes: parsing only under the
codepage is evidence for legacy, parsing as UTF-8 is evidence against, since
codepage text rarely forms valid multibyte UTF-8. Ties leave the file alone.

When the migration cannot be written the append handle uses the legacy codepage,
and errors = "replace" quietly turned characters it cannot hold into question
marks while write() still reported success. That path now escapes to \uXXXX
instead, which is ASCII, so every codepage holds it and json.loads returns the
exact characters. Nothing needs replacing, so errors = "strict" is safe.

stream_installer runs sys.executable, so its output is now decoded as UTF-8 by
utf8_child_env rather than read as the ANSI codepage.

* Only rewrite a shard we can attribute, and append ASCII when we cannot

latin-1 was doing too much work. It reads any byte, so it gave a moved shard a
reading, but it is the right text only for cp1252: cp1251 Привет came back as
Ïðèâåò and the rewrite made that permanent. The codepage is now trusted only
when it is the locale's, and an untrusted reading is never written back.

That leaves three cases where the file holds bytes UTF-8 cannot read and we are
not converting it: no codepage to attribute it to, ambiguous lines outvoting the
unambiguous ones, and a preload that could not read the file at all. All three
used to append UTF-8 into it. They now append pure ASCII, which every
ASCII-compatible codepage stores identically, so the file keeps decoding exactly
as it did and no record is lost.

Keys from the two readings are also kept apart. A damaged line in a healthy
shard was marked seen through its codepage reading, so the retry that would have
replaced the unreadable record was refused as a duplicate.

* Let the flash-attn install stub take the kwargs the installer now passes

_run_kwargs gained encoding and errors, so the one stub in this file that
spelled its signature out rejected the call. The other four here already take
**kwargs; this one now matches.

* Do not let a stuck temp file mask the migration failure

unlink() on the failure path could raise in its own right, on a stale
.utf8.tmp directory or a temp another process holds. That escaped the
constructor instead of returning False, so the caller never reached the ASCII
append fallback that keeps the shard single-encoding.

The pip fallback in install_wheel also spawns a Python child, so it gets
utf8_child_env like the probe above it already had. The uv and nvidia-smi
children are native binaries, where PYTHONIOENCODING would do nothing.

* Stop converting legacy shards; the encoding that wrote them is unknowable

trusted only ever meant that the bytes parse under this machine's codepage,
which for a single-byte codepage is nearly always true. A cp1251 shard opened on
a cp1252 Windows box decodes cleanly and would have been rewritten with Привет
as Ïðèâåò. That is the fourth way this rewrite could corrupt a shard, and the
common cause is that a file's encoding cannot be recovered from its bytes.

So the rewrite is gone. The shard is left exactly as found, and appends are pure
ASCII whenever it holds bytes UTF-8 cannot read, which is what actually
delivered the no-mixed-encoding guarantee the rewrite was added for. Dedup keys
still come from whichever reading parses, since ids are ASCII either way.

This also removes the temp file, so there is no longer any file mode or ACL to
carry across.

* Scan the sandbox shim; it is shipped code, not a build artifact

sandbox_site is on the sandboxed child's PYTHONPATH for every Python run
(tools.py:332, 2660), so excluding it let two unannotated text calls through in
code we ship. Both read and write the remap sidecar, which holds file paths.

The exclusion list is meant for build output only, so the directory comes off
it and the two calls name their encoding.

* Force the worker's pip children to UTF-8, and read DBCS keys with a DBCS codec

The three installer calls run sys.executable -m pip with an inherited
environment, so the parent decoded UTF-8 while the child emitted the ANSI
codepage. They now go through utf8_child_env like the other Python children.

Two tests asserted no env kwarg was passed as a stand-in for no HIP flag being
injected. They now assert the flag itself, which is the guarantee they were
written for and does not depend on how the env is delivered.

Separately, latin-1 cannot stand in for a double-byte codepage while recovering
dedup keys: cp932 表 is 95 5C, and the trail byte reads as a JSON backslash, so
the record failed to parse and its id was forgotten, appending a duplicate on
resume. cp932, cp936, cp949 and cp950 are tried too. The reading is still only
ever used for keys, which are ASCII and identical whichever codec parses.

* Require more than one legacy line before trusting its dedup keys

A shard whose valid records are all ASCII casts no UTF-8 votes, so a single
damaged line won the vote by itself, its key was remembered, and the retry that
would have replaced the unreadable record was refused.

One such line is genuinely undecidable: a legacy record with one accented
character and an ASCII record with one stray byte are the same shape. Reading it
as damage costs a duplicate; reading it as legacy loses the record for good.
Only one of those is recoverable, so it is now read as damage.

A real legacy shard has a legacy line for every record carrying an umlaut, so
its dedup is unaffected.

* Append ASCII whenever the shard already holds non-ASCII bytes

The gate asked whether any line was undecodable as UTF-8, which misses a shard
where every legacy line happens to be valid UTF-8 too. A cp1251 shard of Р°
records is bytes D0 B0 throughout, so appending 世界 as UTF-8 left a file where
cp1251 reads the old records correctly and the new one as mojibake, and UTF-8
does the reverse. No single decoding recovered the whole scrape.

The gate is now simply whether the shard holds any non-ASCII byte at all, which
covers both cases and is easier to reason about: if what is already there reads
differently under different encodings, do not add more bytes that do.

Appending ASCII costs only \uXXXX escapes, which json.loads turns back into the
exact characters, and it leaves the new record correct under either reading.

* Skip the two Linux-gated flash-attn tests off Linux

_should_try_runtime_flash_attn_install ends in sys.platform.startswith(
"linux"), and the threshold test one line above already asserts exactly that,
so the two tests that drive _ensure_flash_attn_for_long_context past the gate
cannot pass anywhere else: the call returns before it reports a status. They
were written on Linux and only surface once the suite actually runs on Windows
or macOS, where both fail on an empty status list. This PR is about making the
backend behave on Windows, so its own suite should be runnable there.

* Fail closed when a KFD topology node does not decode

This PR pins that read to utf-8, which turns an undecodable byte into
UnicodeDecodeError. That is a ValueError, not an OSError, so it slips past the
handler one line below and escapes a helper whose docstring promises to fail
closed on any unreadable node. The caller would then lose the whole HIP-order
map on a machine that has AMD GPUs, and the reason the helper fails closed is
that dropping a node shifts every later ordinal and lets a similar-capacity GPU
pass the total-size guard while showing another card's usage.

Widening the handler is the same one-line change main already made in #7487, so
the two agree and the eventual merge is clean.

* Tighten the comments added in this branch

* Treat an undecodable marker and undecodable metadata as malformed, not fatal

Two more places where pinning the decode changed the failure mode. A
UnicodeDecodeError is a ValueError, so neither `except OSError` nor
`except (JSONDecodeError, OSError)` catches it, and both sites had a documented
fallback that stopped being reached.

An undecodable .transport marker used to read as an unknown value, and the
caller then safely purged and restarted the partial download. It now aborts
prepare_cache_for_transport instead, so the transfer fails rather than retrying.

Undecodable .meta.json used to fall back to the file's own name, the same way
invalid JSON does. It now aborts URI construction for the entire unstructured
seed, so one corrupt byte in original_filename takes out the whole dataset.

Both handlers are widened, matching the KFD fix earlier on this branch.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Widen two more decode guards, and pin the kernel installer's pipe

Same shape as the ones already fixed here: the read was pinned to UTF-8 while
the handler around it still only catches OSError, and UnicodeDecodeError is a
ValueError.

hf_cache_snapshot_dir answers whether a model is already on disk, and the
offline embedding checks turn a raise into a 500. A torn refs/main used to
decode into a nonsense commit and miss the snapshot dir; it now skips that cache
root and keeps looking. _remove_pid_file runs first in _graceful_shutdown, so a
corrupt studio.pid raising there abandoned the inference, export, training and
tunnel children the rest of that function exists to kill.

ssm_runtime's source-build path builds its subprocess kwargs in a dict and
splats them through _run_with_heartbeat, so neither the encoding guard nor the
earlier sweep saw the text = True in it: pip's output was still decoded with the
Windows ANSI codepage, where a non-ASCII path or a compiler diagnostic mojibakes
or raises over an install that was going fine. It now pins the same
utf-8/replace pair install_wheel uses, and the HIP branch extends that env
rather than replacing it. The guard learned the dict-literal shape and reddens
on the old code (ssm_runtime.py:253).

* Tighten the comments around the UTF-8 text I/O pins

Collapse the multi-line rationales added with the encoding pins down to a
line or two each, drop what the code already says, and use one wording for
the repeated child-env note.

* Do not let an unreadable bootstrap password stop startup, and narrow the kwargs guard

ensure_default_admin calls _load_bootstrap_password for every existing admin and
the lifespan calls that with no handler, so pinning the decode turned a damaged
or pre-pin .bootstrap_password file into a backend that will not start. We write
that file ourselves in UTF-8, so a byte that will not decode belongs to a file
whose plaintext is worthless anyway; it now reads as no bootstrap password, the
same answer as an absent file. A readable one still loads.

The new kwargs check also judged every dict literal in the tree, so an unrelated
payload carrying "text": True would have been reported as subprocess
configuration with a misleading message, and a dict that fills in its encoding on
a later line would have been reported too. It now only judges a dict that
actually reaches a call, either splatted through a name or written at the call
site, and treats a later kw["encoding"] assignment as satisfying it. The
ssm_runtime shape it was written for is still caught, and a test pins both
directions.

* Stop reading a UTF-8 record a second time

_read_line always parsed the line under the codepage as well, even when it had
already read as UTF-8. Both callers take the UTF-8 reading when there is one and
never look at the other, so on a healthy shard the second parse is pure waste,
and this file reads all of one on every resume of a scrape it expects to reach
gigabytes. Measured on 200,000 records, 76 MB: 1.96s before, 0.81s after, so the
double reading was costing 2.8x.

The early return is limited to a record, since the key lookup deliberately falls
through to the codepage reading when UTF-8 yields something that is not one. A
line UTF-8 cannot read still tries the codepage, latin-1 and the double-byte
encodings as before, which is what the second reading is for.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Pin the scanned source fixture's line endings

test_remote_code_scan_reads_non_ascii_sources compared a file's contents against
the string it wrote, but wrote it in text mode, so Windows translated the line
ends on the way out and the read back differed by a carriage return. That is the
writer's doing, not the encoding the test is about, and it was the one failure on
the Windows runner that belonged to this branch. The fixture now writes with
newline = "" so the bytes on disk are the string on every platform.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Trim the newer comments to their point

Shorten the widened-guard and state store notes added since the last pass,
and collapse the line-ending note on the scanned source fixture.

* Read the scraper checkpoint as UTF-8 only, never as a codepage

A checkpoint holds nothing but base64 cursors and booleans, so one written by
an older locale-encoded release is byte-identical to a UTF-8 one and already
reads back. The codepage fallback can therefore only ever contribute non-ASCII:
if a single-byte reading of the file were all ASCII, the UTF-8 read would have
succeeded first.

So the only file it changes the answer for is a damaged one, and there it turns
a safe reset into a resume on a mojibaked cursor. GitHub answers that with
INVALID_CURSOR_ARGUMENTS at HTTP 200, gh_client returns the partial document,
and the scraper reads zero nodes and an empty pageInfo, which marks the stream
done. Every later resume then skips it entirely.

Reading UTF-8 only restores the earlier behaviour of dropping a checkpoint that
will not decode, which re-scrapes from the first page while the writers dedup
the replay. The shard scan below keeps its codepage reading; those records do
carry non-ASCII.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Gate the remaining tilelang install tests to Linux

_tilelang_platform_supported() returns False off Linux, so _ensure_tilelang_backend
returns before the install and the subprocess mock these six assert on is never
called. They fail on macOS runners for that reason alone. The rest of the file
already carries this marker; these were missed.

* Gate the Windows-incompatible worker and ROCm tests

Two different gates, because the production code has two. The causal-conv1d and
flash-linear-attention installers bail out on sys.platform == 'win32' alone and
run everywhere else including macOS, so those cases get not_on_windows; marking
them linux_only would skip tests that legitimately pass off Linux. The DRM and
KFD readers return early unless platform.system() is Linux, and their fixtures
build a fake sysfs tree needing PCI addresses like 0000:00:02.0 as directory
names, which Windows cannot represent, so those get linux_only.

The two visible-utilization cases failed for a different reason: on Windows
get_visible_gpu_utilization takes the AMD adapter branch ahead of the torch
fallback under test, and probing it imports torch, which the runner lacks.
Stubbing that branch empty leaves every other platform unchanged.

* Treat unparseable JSON nesting as a parse failure, and guard os.fdopen

json.loads answers nesting it cannot descend with RecursionError, a
RuntimeError, so _parse let it escape where the catch-all it replaced
discarded the record. Both callers run _parse outside any further handler,
so one damaged checkpoint or shard line aborted the scraper at startup.

The encoding guard also missed os.fdopen, which is open() on a descriptor
and takes the same locale default in text mode. It flags exactly the two
text-mode calls that were left unencoded; the swap lock file's reader was
already pinned to UTF-8 while its writer still used the codepage.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Write the non-ASCII source fixture without a 3.10-only argument

Path.write_text() only grew newline in 3.10, and pyproject declares
requires-python >=3.9, so this raised TypeError there. open() takes the same
argument on every supported version and pins the bytes on disk the same way.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Tighten encoding comments

* Follow subprocess calls through callable aliases in the encoding guard

---------

Co-authored-by: Unsloth <michaelhan@Michaels-MacBook-Pro.local>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen <unslothshared@gmail.com>

---------

Co-authored-by: Unsloth <michaelhan@Michaels-MacBook-Pro.local>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen <unslothshared@gmail.com>
2026-07-28 21:27:27 -07:00

611 lines
28 KiB
Python

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Malware / unsafe-file gate for model loads.
The ``trust_remote_code`` consent gate covers the ``auto_map`` Python vector; this
covers the other one -- a malicious pickle inside a weight file, which executes
during ``from_pretrained`` deserialization even with ``trust_remote_code=False``.
It reads Hugging Face's OWN scan (picklescan + ClamAV) via
``model_info(securityStatus=True).security_repo_status``. METADATA-ONLY: it never
downloads, opens, or unpickles the flagged files.
Policy:
* Hard block, non-approvable.
* Block whenever ``filesWithIssues`` lists a non-``safe`` level, regardless of
``scansDone`` (often false even for clean repos). Unknown/future levels fail
CLOSED (block) so Hub schema drift cannot silently allow a bad verdict; only a
small allowlist of clean / not-yet-scanned levels is non-blocking. The sole
fail-open path is an unavailable status (missing field / offline / error).
* Scope to the load-path RCE vector: a root-level (or load-subdir-level),
code-executing file. Inert formats (safetensors / gguf / config / text) and
subdirectory pickles that no root weight-index references are NOT loaded, so
they do not block; an index-referenced shard does, wherever it lives. This
blocks real malware (eicar's root ``*.pkl``/``*.dat``) without false-blocking
repos like ``nvidia/Nemotron-H-8B-Base-8K`` (flagged NeMo pickles under
``nemo/`` that no index lists).
* No first-party exemption (scoping is by load path/format, not org).
* Local paths are skipped (no Hub scan); a remote ``*.gguf``-named repo is still
scanned so a repo cannot dodge the gate by suffixing its name.
"""
import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional
from loggers import get_logger
logger = get_logger(__name__)
# Pickle-format weight files (plain or sharded) that execute code on load; safetensors/gguf
# are inert. Grouped by weight family so an inert safetensors only suppresses the pickle it
# actually replaces: the loader won't use an adapter's safetensors for pytorch_model.bin.
_PICKLE_WEIGHT_RE = re.compile(
r"^(model|pytorch_model|adapter_model|consolidated)(-\d+-of-\d+)?"
r"\.(bin|pt|pth|ckpt|pkl|pickle)$",
re.IGNORECASE,
)
# Non-blocking levels: clean or not-yet-finished. Anything else (unsafe/suspicious/
# malicious or a future label) blocks, so Hub schema drift fails CLOSED.
_NONBLOCKING_LEVELS = frozenset(
{"", "safe", "pending", "scanning", "queued", "unscanned", "error", "unknown", "none"}
)
# Suffixes that cannot execute code on load (tensor-only safetensors, non-pickle gguf,
# text/markup/images), so a flag on one is never an RCE vector.
_INERT_SUFFIXES = frozenset(
{
".safetensors",
".gguf",
".json",
".txt",
".md",
".rst",
".yaml",
".yml",
".png",
".jpg",
".jpeg",
".gif",
".webp",
".svg",
".bmp",
".gitattributes",
".gitignore",
}
)
# Source files are not deserialized by a weight load; executable repo code runs only
# via auto_map, which is the consent gate's domain. So a flag on a .py is not this
# gate's vector (else a flagged helper/train script would false-block).
_SOURCE_SUFFIXES = frozenset({".py", ".pyc", ".pyx", ".pyi"})
# Torch-family weight indexes: from_pretrained feeds each shard they name to load_state_dict, which
# torch.load()s (pickle) any shard whose name does not end in .safetensors, whatever its stem. A
# pytorch index is superseded when a base safetensors is present (the loader prefers it); a
# safetensors index IS the chosen archive, so a non-safetensors target it names still loads. tf/flax
# indexes load via non-pickle loaders, so they are not a torch.load vector here.
_TORCH_INDEX_FILES = ("pytorch_model.bin.index.json", "model.safetensors.index.json")
# Root weight-index files. from_pretrained reads these to find sharded weights, so a
# flagged subdir pickle is a load vector iff a root index references it.
_TRANSFORMERS_INDEX_FILES = (
"pytorch_model.bin.index.json",
"model.safetensors.index.json",
"tf_model.h5.index.json",
"flax_model.msgpack.index.json",
)
def _normalize_repo_path(path: str) -> str:
"""Strip ``./`` prefixes and normalize separators for repo-relative comparison."""
p = (path or "").strip().replace("\\", "/")
while p.startswith("./"):
p = p[2:]
return p
def _file_suffix(path: str) -> str:
"""Lowercase ``.ext`` of the basename, or ``""`` if none."""
base = _normalize_repo_path(path).rsplit("/", 1)[-1]
return "." + base.rsplit(".", 1)[1].lower() if "." in base else ""
def _hf_cache_snapshot_ref(local_path: str) -> Optional[tuple]:
"""``(repo_id, revision)`` for an HF-cache snapshot path, else None. An inactive Studio
cache loads by its snapshot path but keeps the ``models--org--repo/snapshots/<rev>``
layout, so the gate recovers its provenance and scans that exact commit instead of
exempting it (an older cached commit can hold a pickle since dropped from the branch)."""
try:
path = Path(local_path).resolve(strict = False)
except (OSError, ValueError):
return None
for parent in path.parents:
if parent.name != "snapshots":
continue
encoded = parent.parent.name
if not encoded.startswith("models--"):
return None
repo_id = encoded.removeprefix("models--").replace("--", "/")
if not repo_id:
return None
return repo_id, path.relative_to(parent).parts[0] # <rev> dir under snapshots/
return None
def _load_relative_path(norm: str, load_subdirs) -> str:
"""``norm`` relative to a ``from_pretrained`` load root. Some loads read from a
snapshot SUBDIRECTORY (Spark-TTS / BiCodec load ``<snapshot>/LLM``), where a file
directly under the subdir is root-level, not nested. Strips the matching load-subdir
prefix, or returns ``norm`` unchanged when it is not under one.
"""
for subdir in load_subdirs or ():
prefix = _normalize_repo_path(subdir).strip("/")
if prefix and norm.startswith(prefix + "/"):
return norm[len(prefix) + 1 :]
return norm
def _index_prefixes(load_subdirs) -> tuple:
"""Prefixes to look for weight-index files under: repo root plus each load subdir."""
prefixes = [""]
for subdir in load_subdirs or ():
p = _normalize_repo_path(subdir).strip("/")
if p:
prefixes.append(p + "/")
return tuple(prefixes)
def _indexed_shard_paths(
model_name: str,
hf_token: Optional[str],
load_subdirs = (),
revision: Optional[str] = None,
):
"""Repo-relative weight paths a load could fetch via weight-index files. Returns a
set (empty when the repo ships no index files -- a definitive "nothing sharded"), or
None when the lookup was inconclusive (transient error) so the caller treats a
flagged subdir pickle conservatively. Reads only small JSON indexes, never weights.
Indexes are looked up at the root and each ``load_subdirs`` root, with ``weight_map``
entries re-prefixed to repo-relative paths. ``revision`` scopes to a cached commit.
"""
import json
try:
from huggingface_hub import hf_hub_download
from huggingface_hub.utils import EntryNotFoundError
from utils.hf_cache_settings import active_hf_hub_cache
except Exception:
return None
paths: set = set()
inconclusive = False
for prefix in _index_prefixes(load_subdirs):
for filename in _TRANSFORMERS_INDEX_FILES:
try:
index_path = hf_hub_download(
model_name,
prefix + filename,
revision = revision,
token = hf_token or None,
cache_dir = active_hf_hub_cache(),
)
except EntryNotFoundError:
continue # definitively absent, not an error
except Exception:
inconclusive = True # transient: an index that might exist could not be read
continue
try:
weight_map = (json.loads(open(index_path, encoding = "utf-8-sig").read()) or {}).get(
"weight_map"
) or {}
for shard in weight_map.values():
shard_norm = _normalize_repo_path(str(shard))
# weight_map paths are relative to the index file's directory.
if prefix and not shard_norm.startswith(prefix):
shard_norm = prefix + shard_norm
paths.add(shard_norm)
except Exception:
inconclusive = True
# Any transient failure -> inconclusive (the shard could be listed only by the index
# we could not read), so fail closed (None) and let the caller block. Ships no index
# files -> EntryNotFoundError for each, empty set, a definitive "nothing sharded".
if inconclusive:
return None
return paths
# Two-timeout metadata fetch, mirroring hub.workers.hf_download._retry_metadata_fetch.
_REQUEST_TIMEOUT = 10.0
_RETRY_TIMEOUT = 20.0
@dataclass
class FileSecurityDecision:
"""Outcome of the Hub security scan for one model repo."""
model_name: str
blocked: bool
unsafe_files: list = field(default_factory = list) # [{"path", "level"}]
reason: str = ""
def response_payload(self) -> dict:
"""Machine-readable detail merged into the preflight payload the dialog reads."""
return {
"unsafe_files": self.unsafe_files,
"security_blocked": self.blocked,
"reason": self.reason,
}
def security_load_subdirs(model_name: str, hf_token: Optional[str] = None) -> tuple:
"""Snapshot subdirectories a load calls ``from_pretrained`` on, for scoping the scan.
Most models load from the root (``()``); Spark-TTS / BiCodec load ``<snapshot>/LLM``,
so ``LLM/`` is a load root for them. Metadata-only (tokenizer special tokens), cached.
"""
try:
from utils.models.model_config import detect_audio_type, load_model_defaults
if detect_audio_type(model_name, hf_token = hf_token) == "bicodec":
return ("LLM",)
# Tokenizer detection can fail (network/gated/unresolved alias); the YAML default
# also pins the audio type, so fall back to it (else a flagged LLM/ pickle is
# treated as an ignored subdir artifact).
if (load_model_defaults(model_name) or {}).get("audio_type") == "bicodec":
return ("LLM",)
except Exception:
pass
return ()
def _load_scan_target(model_name: str, load_subdirs: tuple) -> tuple:
"""Map a load alias to the ``(repo_id, load_subdirs)`` the load actually fetches. The
Spark-TTS / BiCodec alias ``<parent>/LLM`` is downloaded by the trainer as
``unsloth/<parent>`` and loaded from ``LLM/``, so scan that repo with ``LLM`` as a
load root (the literal alias 404s and fails open). Everything else is unchanged.
"""
try:
from utils.paths import is_local_path
if is_local_path(model_name):
return model_name, load_subdirs
except Exception:
return model_name, load_subdirs
name = (model_name or "").strip().strip("/")
# Rewrite ONLY a registry-known bicodec alias, never any repo ending in "/LLM"
# (e.g. "evil/LLM" would scan unsloth/evil and fail open on the real repo).
if name.endswith("/LLM") and name.count("/") == 1:
try:
from utils.models.model_config import load_model_defaults
if (load_model_defaults(name) or {}).get("audio_type") == "bicodec":
parent = name[: -len("/LLM")]
return f"unsloth/{parent}", tuple(dict.fromkeys((*load_subdirs, "LLM")))
except Exception:
pass
return model_name, load_subdirs
def _fetch_security_status(
model_name: str,
hf_token: Optional[str],
revision: Optional[str] = None,
):
"""``security_repo_status`` (a dict) or None if unavailable. Hub metadata only;
retries once on a transient error, then returns None so the caller fails open.
``revision`` scopes the scan to a specific cached commit (else the default branch).
"""
from huggingface_hub import model_info as hf_model_info
token_arg = hf_token if hf_token else False
last_exc = None
for attempt, timeout in enumerate((_REQUEST_TIMEOUT, _RETRY_TIMEOUT)):
try:
info = hf_model_info(
model_name,
revision = revision,
token = token_arg,
securityStatus = True,
timeout = timeout,
)
return getattr(info, "security_repo_status", None)
except Exception as exc: # network/offline/gated/404/unsupported-client
last_exc = exc
if attempt == 0:
continue
logger.debug(
"HF security scan unavailable for '%s' (%s); failing open.",
model_name,
type(last_exc).__name__ if last_exc else "unknown",
)
return None
def _st_load_roots(snapshot: Path) -> list:
"""Directories a SentenceTransformer load deserializes weights from: the snapshot root plus
each module path in modules.json. Local, no network. Mirrors the online gate (which ignores
unreferenced nested pickles ST never loads) so the offline gate doesn't over-block."""
roots = [snapshot]
try:
import json
modules = json.loads((snapshot / "modules.json").read_text(encoding = "utf-8-sig"))
except (OSError, ValueError):
return roots # no / invalid modules.json -> snapshot root is the only load root
for module in modules or ():
path = str((module or {}).get("path", "")).strip().strip("/")
# Relative module path only; ignore a crafted "../" escape.
if path and ".." not in path.split("/"):
candidate = snapshot / path
if candidate not in roots:
roots.append(candidate)
return roots
def _indexed_pickle_shards(index_path: Path, root: Path, snapshot: Path) -> list:
"""Shards a torch weight index points a ``from_pretrained`` load at that load_state_dict would
torch.load (pickle): every ``weight_map`` target NOT ending in ``.safetensors``, whatever its
stem (an arbitrary name like ``shards/payload`` still deserializes). Resolved relative to the
index dir (``root``) like the loader, so a shard in a nested dir is followed (iterdir misses it).
Lexical only, never ``Path.resolve()`` (HF snapshot files symlink into ``blobs/``, so resolving
escapes the snapshot and false-blocks every shard). Raises OSError -> caller fails CLOSED on an
unreadable/invalid index or a target escaping the snapshot."""
import json
import os
try:
# JSON is UTF-8 by spec; pin it so a non-ASCII index is not misdecoded (and needlessly
# blocked) under Windows' cp1252 default.
parsed = json.loads(index_path.read_text(encoding = "utf-8-sig"))
except (OSError, ValueError) as exc:
raise OSError(f"unreadable weight index: {index_path}") from exc
weight_map = parsed.get("weight_map") if isinstance(parsed, dict) else None
if not isinstance(weight_map, dict):
return [] # no dict weight_map -> the loader resolves no shards from this index
snapshot_norm = os.path.normpath(str(snapshot))
shards = []
for shard in weight_map.values():
raw = str(shard)
if not raw:
continue
# Join the RAW weight_map value like from_pretrained's os.path.join: on POSIX a backslash is a
# literal filename char (not a separator), so normalizing it would probe a different path than
# the loader opens. normpath + containment stay platform-aware (os.sep) to block "..".
joined = os.path.normpath(os.path.join(str(root), raw))
if joined != snapshot_norm and not joined.startswith(snapshot_norm + os.sep):
raise OSError(f"weight index escapes the snapshot: {index_path}")
shard_path = Path(joined)
# Case-SENSITIVE, mirroring load_state_dict's own endswith(".safetensors"): a shard named
# payload.SAFETENSORS is not treated as safetensors by the loader and falls to torch.load.
if not shard_path.name.endswith(".safetensors") and shard_path.is_file():
shards.append(shard_path)
return shards
def _loader_resolves(root: Path, name: str) -> bool:
"""True iff from_pretrained would open ``name`` under ``root``. ``is_file()`` honors the platform
(case-sensitive on Linux, case-insensitive on Windows/macOS), so it mirrors the loader's own
lookup: an oddly-cased decoy counts as an alternative only where the loader would truly open it.
A name-fold instead would let an uppercase MODEL.SAFETENSORS suppress the scan on Linux while the
loader, asking for the canonical lowercase name, silently falls through to a pickle index."""
return (root / name).is_file()
def _cached_pickle_weight_files(snapshot: Path) -> list:
"""Pickle weight files a SentenceTransformer/Transformers load deserializes from snapshot's ST
load roots, EXCLUDING those whose weight family also ships an inert safetensors in the same dir
(the loader prefers it): a base pickle is suppressed only by a base model.safetensors, an adapter
pickle only by adapter_model.safetensors -- an unrelated safetensors is no substitute. Covers
both direct-child pickles AND pickle shards referenced by a local weight index (which the loader
follows into nested dirs, matching the online gate). Raises OSError -- caller fails CLOSED -- if
the snapshot root or a weight index is unreadable, or an index reference escapes the snapshot."""
blocked = []
seen = set()
def _add(path: Path):
key = str(path)
if key not in seen:
seen.add(key)
blocked.append(path)
for root in _st_load_roots(snapshot):
try:
entries = [p for p in root.iterdir() if p.is_file()]
except OSError:
if root == snapshot:
raise # top-level unreadable -> fail closed
continue # unreadable module subdir: nothing loadable to attest here
# Safetensors alternatives the loader would actually resolve (never a bare name-fold, which
# fails OPEN: see _loader_resolves). A base pickle is replaced only by a base safetensors, an
# adapter pickle only by an adapter one. A single model.safetensors also outranks BOTH indexes.
has_direct_base_safetensors = _loader_resolves(root, "model.safetensors")
has_base_safetensors = has_direct_base_safetensors or _loader_resolves(
root, "model.safetensors.index.json"
)
has_adapter_safetensors = _loader_resolves(root, "adapter_model.safetensors")
for path in entries:
if not _PICKLE_WEIGHT_RE.match(path.name):
continue
is_adapter = path.name.lower().startswith("adapter_model")
has_alternative = has_adapter_safetensors if is_adapter else has_base_safetensors
if not has_alternative:
_add(path)
# A torch weight index makes from_pretrained load nested shards iterdir never sees; the loader
# torch.loads any not ending in .safetensors. Probe the canonical index name with the loader's
# own lookup (_loader_resolves), so an oddly-cased artifact it would never open does not block.
# A direct model.safetensors wins over BOTH indexes; failing that a base safetensors still
# outranks the pytorch index, while a safetensors index is itself the chosen archive.
for index_name in _TORCH_INDEX_FILES:
if not _loader_resolves(root, index_name):
continue
if has_direct_base_safetensors:
continue
if index_name == "pytorch_model.bin.index.json" and has_base_safetensors:
continue
for shard_path in _indexed_pickle_shards(root / index_name, root, snapshot):
_add(shard_path)
return blocked
def _evaluate_local_only(model_name: str) -> FileSecurityDecision:
"""Offline security gate. The Hub scan is unreachable, so inspect the local cache and fail
CLOSED on an unscanned pickle weight with no inert safetensors alternative, rather than
failing open or hanging. Safetensors/gguf-only cache loads; nothing cached -> allowed."""
from utils.utils import hf_cache_snapshot_dir
try:
snapshot = hf_cache_snapshot_dir(model_name)
except Exception:
logger.warning("Offline gate: could not resolve the cache for '%s'; blocking.", model_name)
return FileSecurityDecision(
model_name, True, reason = "offline; could not inspect the local cache"
)
if snapshot is None:
return FileSecurityDecision(model_name, False, reason = "offline; nothing cached to load")
try:
pickles = _cached_pickle_weight_files(snapshot)
except OSError:
logger.warning("Offline gate: could not read the cache for '%s'; blocking.", model_name)
return FileSecurityDecision(
model_name, True, reason = "offline; could not read the local cache"
)
if not pickles:
return FileSecurityDecision(
model_name, False, reason = "offline; cached weights are inert (safetensors/gguf)"
)
# Snapshot-relative posix paths (match the online gate; disambiguate same-named pickles).
rel_paths = sorted(p.relative_to(snapshot).as_posix() for p in pickles)
names = ", ".join(rel_paths)
logger.warning(
"Blocking offline load of '%s': cached pickle weight(s) cannot be malware-scanned "
"offline and have no safetensors alternative (%s).",
model_name,
names,
)
return FileSecurityDecision(
model_name,
True,
unsafe_files = [{"path": rel, "level": "unscanned"} for rel in rel_paths],
reason = f"offline; unscanned pickle weights with no safetensors alternative: {names}",
)
def evaluate_file_security(
model_name: str,
hf_token: Optional[str] = None,
*,
load_subdirs = (),
local_only_load: bool = False,
) -> FileSecurityDecision:
"""Block a load when HF's security scan flags unsafe serialized files.
Call UNCONDITIONALLY before any load (independent of trust_remote_code): a malicious
pickle deserializes during ``from_pretrained`` regardless. Metadata-only; fails open
when the scan is unavailable.
``load_subdirs`` names subdirs the load calls ``from_pretrained`` on (e.g. ``("LLM",)``
for Spark-TTS / BiCodec, loading ``<snapshot>/LLM``): a flagged file directly under one
is root-level there and blocks, and an index inside it is honored when scoping shards.
``local_only_load`` marks an offline load: with the Hub scan unreachable, inspect the local
cache and fail CLOSED on an unscanned pickle weight with no safetensors alternative.
"""
# Scan the repo the load actually fetches, not the literal alias (which 404s and
# fails open): the Spark-TTS "<parent>/LLM" alias is really unsloth/<parent> from LLM/.
model_name, load_subdirs = _load_scan_target(model_name, tuple(load_subdirs))
# Local paths have no Hub scan, EXCEPT an HF-cache snapshot whose canonical path
# encodes a repo id + commit: scan that exact commit so an inactive-cache load can't
# dodge the gate. A remote ref is scanned even if named "*.gguf" (name can't dodge it).
snapshot_revision = None
try:
from utils.paths import is_local_path
if is_local_path(model_name):
cache_ref = _hf_cache_snapshot_ref(model_name)
if cache_ref is None:
return FileSecurityDecision(model_name, False, reason = "local path; no Hub scan")
model_name, snapshot_revision = cache_ref
except Exception:
# Cannot classify the path -> do not block on that account.
return FileSecurityDecision(model_name, False, reason = "path check failed; not blocked")
# Offline: inspect the local cache and fail closed rather than hang on model_info or fail open.
if local_only_load:
return _evaluate_local_only(model_name)
status = _fetch_security_status(model_name, hf_token, revision = snapshot_revision)
if not isinstance(status, dict):
return FileSecurityDecision(
model_name, False, reason = "scan unavailable; allowed (fail-open)"
)
# Block a non-``safe`` flagged file scoped to the load-path RCE vector (root-level,
# code-executing). Not gated on ``scansDone`` (often false even when clean; a flagged
# file is flagged regardless). Unknown levels fail closed; in-progress/clean do not.
# Subdir pickles and inert formats (safetensors/gguf) are not loaded by
# from_pretrained and do not block. Unavailable status (above) is the only fail-open.
unsafe = []
skipped = [] # flagged, but not a load-path RCE vector (subdir artifact / inert)
maybe_shard = [] # flagged subdir pickle: a load vector ONLY if a root index lists it
for entry in status.get("filesWithIssues") or []:
if not isinstance(entry, dict):
continue
level = str(entry.get("level", "")).lower()
if level in _NONBLOCKING_LEVELS:
continue
path = entry.get("path", "")
norm = _normalize_repo_path(path)
suffix = _file_suffix(norm)
# Path relative to the load root: a file under a load subdir (e.g. LLM/) is
# root-level there, not nested.
load_rel = _load_relative_path(norm, load_subdirs)
if not norm or suffix in _INERT_SUFFIXES or suffix in _SOURCE_SUFFIXES:
# Inert formats cannot execute on load; source code is the consent gate's
# domain (auto_map), not a deserialization vector.
skipped.append({"path": path, "level": level})
elif "/" not in load_rel:
unsafe.append({"path": path, "level": level}) # root pickle -> load vector
else:
# Subdir pickle: deserialized only if a weight index references it.
maybe_shard.append({"path": path, "level": level, "norm": norm})
if maybe_shard:
indexed = _indexed_shard_paths(
model_name, hf_token, load_subdirs, revision = snapshot_revision
)
for m in maybe_shard:
# Block if a root index lists this shard, or if the lookup was inconclusive
# (transient error -> stay conservative). A definitive "no index / not listed"
# stays non-blocking (e.g. NeMo nemo/*.distcp).
if indexed is None or m["norm"] in indexed:
unsafe.append({"path": m["path"], "level": m["level"]})
else:
skipped.append({"path": m["path"], "level": m["level"]})
if not unsafe:
if skipped:
# Flagged files exist, but none the load deserializes (subdir pickle or inert
# format) -> allow, but log them so they stay visible.
logger.info(
"'%s': Hugging Face flagged files, but none are a load-path RCE "
"vector (subdir/inert); allowing the load. Flagged: %s",
model_name,
", ".join(f"{s['path']}({s['level']})" for s in skipped),
)
return FileSecurityDecision(model_name, False, reason = "no unsafe files in the load path")
names = ", ".join(u["path"] for u in unsafe if u["path"]) or "unknown files"
logger.warning(
"Blocking load of '%s': Hugging Face security scan flagged unsafe files (%s).",
model_name,
names,
)
return FileSecurityDecision(
model_name,
True,
unsafe_files = unsafe,
reason = f"Hugging Face security scan flagged unsafe files: {names}",
)