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

Harden model fetching (#6391)

Browse source 
* Harden model fetching: consent gate for trust_remote_code

Add a load-path consent gate that scans a model's auto_map repository code
before it executes and blocks CRITICAL/HIGH findings unless the user pins
approval of that exact code version. Capability detection stays code-free,
reading raw config.json instead of AutoConfig.

- Scan config.json and tokenizer_config.json auto_map, nested local helpers,
  and external owner/name--module repos; fail closed on partial downloads.
- Gate inference, training, and export workers, including the MLX path and a
  LoRA's base model, and report requires_trust_remote_code from the raw config
  so chat and auto-load surface the dialog.
- Verify trusted-org auto-enable against the Hub with the request token and key
  the verdict cache by token; reject local-path and spoofed names.
- Add a consent dialog showing the flagged file, line, and surrounding code.
- Thread hf_token through the scan and load paths for gated repos.

* Address review: token handling, tokenizer/LoRA scan coverage, rollback

- Send the HF token for remote-code scans in the POST body, not the URL, so it
  never lands in a log or browser history.
- Collect tokenizer_config.json auto_map files directly instead of relying only
  on the repo file listing.
- Resolve a LoRA's base model for the validate flag and the scan endpoint so the
  dialog scans the code the workers actually gate.
- Pass the request token to the training YAML trusted-org auto-enable.
- Resend a previously approved fingerprint when rolling back to a custom-code
  model after a failed switch.

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

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

* Consent UX: drop legacy chat toggle, fix decline copy, purge declined downloads

The per-model consent dialog is now the single approval path for custom
(auto_map) code in chat, so three leftovers from before it existed are removed:

- Remove the "Enable custom code" switch from Chat Settings and stop persisting
  trust_remote_code, so a previously saved blanket-on cannot linger and load a
  model without going through per-version review. The flag stays as an internal
  YAML/preset default (e.g. first-party auto-enable); the load path still gates
  every custom-code load on a fingerprint only the dialog produces.
- Reword the decline message and the auto-load toast to describe approving the
  model's code from the dialog, not a missing settings toggle.
- On decline, purge the repo the scan downloaded so untrusted code is not left
  on disk. A new /api/models/discard-remote-code endpoint deletes only a
  metadata-only cache entry the scan created; it refuses local paths, loaded
  models, and any repo with weight files cached, so a model the user already had
  or pre-downloaded is always left untouched. The frontend only calls it when
  the scan reported created_by_scan.

Adds discard-endpoint tests (delete metadata-only, refuse on weights/gguf,
refuse local, no-op when not cached) and a created_by_scan payload assertion.

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

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

* Export: remove the user-facing trust remote code toggle

The Export page kept a "Trust remote code" switch (default on) next to the HF
token field. Like chat, custom (auto_map) code should be approved per model
through the load-time review dialog, not a persistent blanket switch, so the
toggle is removed. The export load path already routes through the same consent
dialog: an HF source now starts with trust_remote_code off and only enables it
when the user approves the scanned code in the dialog (a local checkpoint the
user exported stays trusted by default). With the dialog unreachable and no
approval, an HF source loads with trust_remote_code off, which fails closed
rather than running unreviewed code.

* Block loads of repos with unsafe files using Hugging Face's security scan

The trust_remote_code consent gate covers one load-time RCE vector (a repo's
auto_map Python). It does not cover the other: a malicious pickle inside a weight
file (pytorch_model.bin, *.pkl, *.dat) deserializes during from_pretrained even
with trust_remote_code False, so a repo with a normal config plus a poisoned
pickle slips past the existing gate.

Add a metadata-only malware gate that uses Hugging Face's own scan (picklescan +
ClamAV), read via model_info(securityStatus=True).security_repo_status. It never
downloads, opens, or unpickles the flagged files; it only reads the Hub's verdict
and surfaces the flagged file names. New evaluate_file_security runs
unconditionally (independent of trust_remote_code) in every load path (inference,
training SFT/MLX, export), blocking the load when a file is flagged
unsafe/suspicious/malicious. The /remote-code-scan preflight and the validate
endpoint also report the result so the consent dialog opens as a hard block (no
override) listing the flagged files, even for a repo with no custom code.

Policy: hard block with no user override; fail open when the scan is unavailable
(offline/unscanned) so legitimate loads are not broken; no first-party exemption
(a poisoned pickle in a compromised trusted repo still blocks); local paths and
GGUF are skipped (no Hub scan, non-pickle format). Blocking does not gate on
scansDone, since that is often false for clean repos and a file already flagged
unsafe is unsafe regardless.

Adds test_file_security.py covering the block/allow/fail-open/skip matrix.

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

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

* Address review: scan list-form tokenizer auto_map, gate unsafe files on all load paths

Fixes from a 10-reviewer pass on the model-fetching hardening:

- The remote-code scanner skipped tokenizer auto_map encoded as a [slow, fast]
  list (transformers' standard tokenizer shape, e.g.
  {"AutoTokenizer": ["owner/repo--tokenization_x.Slow", null]}). External
  tokenizer code in that form was never fetched, scanned, or fingerprinted, so an
  AutoTokenizer(trust_remote_code=True) load could run it. _auto_map_refs now
  flattens string, list, and nested values. Adds a regression test.

- Compare-mode chat loads and background auto-load only gated on
  requires_trust_remote_code, so a repo flagged unsafe by the Hub scan but with no
  custom code skipped the hard-block dialog. Both now also gate on
  requires_security_review, matching the main chat path.

- The /remote-code-scan and /validate routes collapsed a LoRA adapter to its base
  before the malware scan, so unsafe files in the adapter repo itself were missed
  in the pre-load review (the workers already scan both). Both routes now run the
  file-security scan over the adapter and the base.

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

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

* Require approval for all HIGH remote code, fail closed when unscannable

Tighten the load-time security gates based on review:

Consent gate
- HIGH-severity auto_map code now requires explicit, per-version approval for
  every repo, including first-party unsloth/nvidia. The org is no longer a
  blanket bypass: a compromised first-party repo with HIGH code still warrants
  review. CRITICAL stays a hard block; clean code still loads after the consent
  prompt.
- Fail closed when auto_map code is present but cannot be fully fetched or
  listed to scan (gated, offline, transient, or a repo-listing failure that
  could hide an imported helper). We cannot fingerprint code we cannot see, so
  this is a non-approvable block, retryable once the repo is reachable.
- Scan auto_map from every config that can carry one (model, tokenizer, image
  and feature processor, processor, video processor), not just config.json and
  tokenizer_config.json, so a custom-processor model is not missed. The file
  list is the single source of truth in remote_code_scan and is pinned to the
  transformers filename constants by a guard test.
- Distinguish a genuine 404 (config truly absent) from a transient error: only
  the latter forces a scan, so a repo with no config is correctly a no-op.

Malware gate
- Scan a remote repo even when its name ends in .gguf; only local paths skip the
  Hub scan, so a repo cannot dodge the scan by naming itself "*.gguf".
- Correct the docstring: a file already flagged unsafe blocks regardless of
  scansDone; the only fail-open path is an unavailable scan.

Coverage
- Resolve a remote LoRA adapter's base model (not just local directories) so the
  base, where the code and weights actually execute, is scanned in validate,
  the scan route, and the training and export workers.
- Gate the embedding training path (FastSentenceTransformer) with the malware
  and consent checks, matching the other load paths.

Tests updated and added for each change.

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

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

* Scope malware gate to the load-path vector; stop false-blocking first-party models

Follow-up hardening from a second review pass + a broad live model matrix
(unsloth/* , nvidia/* , third-party, and the eicar malware repo).

Malware / unsafe-file gate
- Scope the block to the actual RCE vector: a root-level file in a code-executing
  format. from_pretrained deserializes weight files at the repo ROOT, so a flag is
  only a load-path pickle vector there. Two exclusions, because neither is loaded:
  inert formats (safetensors is tensor-only, gguf is non-pickle, configs/text/
  images) and files in subdirectories. This keeps eicar blocked (its *.pkl/*.dat/
  eicar_test_file sit at the repo root) while no longer false-blocking legitimate
  first-party repos: nvidia/Nemotron-H-8B-Base-8K ships root safetensors plus NeMo
  pickle checkpoints under nemo/ that the loader never touches, and the Hub flags
  both; the gate previously hard-blocked it.
- Unknown / future non-"safe" levels now fail closed (block) instead of being
  silently allowed, so Hub schema drift cannot introduce a bypass; in-progress
  ("pending"/"scanning"/"error") levels stay non-blocking to avoid false blocks.

Consent gate
- Ignore a STALE own-repo auto_map target that is absent from the repo listing (an
  older config pointing at a file the repo no longer ships) instead of failing the
  whole repo closed as unscannable. The present .py are still fully scanned, which
  is the stronger coverage, and a file that is not there cannot execute. This
  unblocks first-party models like unsloth/PaddleOCR-VL (its tokenizer_config.json
  names processing_ppocrvl.py while the repo ships processing_paddleocr_vl.py). A
  referenced .py that IS present but cannot be fetched, and a repo-listing failure,
  still fail closed.

Remote LoRA base resolution
- Distinguish a genuine 404 (not a LoRA / repo absent -> None) from a transient
  error: the transient case is retried once, then logged as a WARNING (a missed
  base is scanned by neither gate) rather than silently skipped.

Discard endpoint
- Treat .onnx and .ckpt as weights so a repo whose only heavy artifact is one of
  those is never eligible for the declined-download purge.

Tests added for each: load-path scoping (safetensors/subdir/Nemotron-H shapes,
unknown-level fail-closed, pending non-block), stale own-repo auto_map ref, remote
LoRA transient retry, and the empty-config-list (all-404 -> []) semantics.

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

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

* Make LoRA-base transient-warning test robust to logging backend

Assert on the logger object directly instead of capsys, so the test does not
depend on whether the real structlog logger or the module-stub logger is active
(which varies with test collection order).

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

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

* Allow a repo with auto_map but no executable code (e.g. GGUF) instead of blocking

A config can declare an auto_map yet the repo ship NO executable .py -- most
commonly a GGUF repo whose config.json carries an auto_map copied from the original
model (e.g. unsloth/Llama-3_1-Nemotron-Ultra-253B-v1-GGUF references
modeling_decilm.py, which the GGUF-only repo does not contain). A GGUF model loads
through llama.cpp, which never executes auto_map, and transformers cannot run a file
that is not present, so there is nothing to scan and trust_remote_code is a no-op.

The fail-closed change treated this empty result the same as "code is present but we
could not fetch it" and hard-blocked the load. Distinguish the two: repo_remote_code_files
now RAISES RemoteCodeUnscannable when code is present but cannot be fully fetched or
listed (offline / gated / transient / a present .py that 404s / a listing failure),
and returns an empty dict only when the listing succeeded and the repo genuinely ships
no executable .py. The consent gate blocks on the exception (fail closed) and allows the
empty case as a no-op. Real unscannable code still hard-blocks; eicar and CRITICAL/HIGH
custom code are unaffected.

Verified against all 37 unsloth/*Nemotron* models (two GGUF repos were false-blocked,
now load) and the existing matrix (eicar still blocks; DeepSeek-OCR / NVLM-D-72B still
prompt approvable consent). Tests updated to expect the raise for unscannable cases and
added for the no-executable-code no-op.

* Ignore vestigial auto_map in GGUF repos (llama.cpp never runs it)

A GGUF repo's config.json is often copied verbatim from the original
transformers model, auto_map and all, but a GGUF load goes through
llama.cpp which never executes auto_map, so the config is inert. Treat
a direct .gguf reference, and a repo that ships .gguf weights with no
.safetensors, as having no remote code so the consent flow is never
triggered. A mixed repo with both .gguf and .safetensors is still gated,
since the safetensors variant would load through transformers where
auto_map does run. The check sits behind the existing auto_map-present
gate so normal models pay no extra repo listing.

* Add scanner-result copy to the remote-code consent dialog

Make the consent dialog state the scan outcome in plain language for
every model. When the static scan finds nothing, reassure the user with
'Our automatic scanner did not flag any worrying files, but please
double check.' (shown only for the clean, approvable case). When the
scan flags custom code or unsafe files, label the list with 'Our
automatic scanner flagged issues including:'. The Hugging Face
attribution for unsafe files stays in the dialog description.

* Close GGUF-suffix consent bypass for repo ids ending in .gguf

The .gguf short-circuit in _config_has_auto_map skipped the scan for any
model name ending in .gguf, including a bare two-segment repo id like
'evil/model.gguf'. Such a repo can still ship safetensors plus auto_map
Python that transformers would execute, so skipping the scan was an
asymmetric bypass (file_security already scans those repos). Restrict the
short-circuit to genuine direct GGUF file references via
_is_direct_gguf_file_ref: a local .gguf path, or a remote repo_id plus
filename (three or more segments). A two-segment repo id named *.gguf now
falls through to the config scan and _is_gguf_repo file inspection, so it
only skips consent when it actually ships .gguf weights and no safetensors.

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

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

* Align consent dialog body with the title and fix narrow-width overflow

The scan results (the 'Our automatic scanner...' label, finding/unsafe
cards, and the clean-scan reassurance) sat at the dialog's left padding
while the title and description were indented past the status icon, so
the body did not line up under the description. Move the title,
description and results into one column to the right of the icon so they
share a left edge, and let that column fill its width so the description
no longer wraps early.

Also stop a wide code snippet from pushing the dialog off-screen on
narrow viewports: AlertDialogHeader is a grid with place-items-center,
which sized the content row to its content; give the row w-full so it
fills the track, and add min-w-0 down the results chain so the snippet
scrolls inside its card instead of widening the dialog. Verified aligned
and contained from mobile portrait through ultrawide.

* Treat a repo as GGUF-only only when it ships no transformers weights

_is_gguf_repo excluded only .safetensors, so a repo with a .gguf and a
pytorch_model.bin (or .pt/.pth/.h5/.msgpack/.onnx/.ckpt) and no
safetensors was treated as GGUF-only and skipped the consent scan, even
though transformers can load that weight set and execute the repo's
auto_map code. Require the absence of ANY transformers-loadable weight
before treating the repo as a llama.cpp-only GGUF load. A genuine
GGUF-only repo (only .gguf) is still inert; a mixed repo with any pickle
or safetensors weight is gated. Adds a regression test across all the
non-safetensors weight formats.

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

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

* Block flagged subdir weight shards referenced by a root index

The malware gate treated every subdirectory file as non-loadable, but
from_pretrained deserializes a subdir shard a root index references
(pytorch_model.bin.index.json -> shards/...-00001-of-00002.bin). Read the
root weight indexes and block a flagged subdir pickle the weight_map
points at; a flagged subdir pickle no index lists (NeMo nemo/*.distcp)
stays non-blocking, and an inconclusive index lookup fails closed.

* Pass hf_token to the export checkpoint load

ExportBackend.load_checkpoint scanned with hf_token in the worker but
loaded the weights unauthenticated, so a gated/private checkpoint passed
preflight then 401'd at from_pretrained. Add hf_token to load_checkpoint
and forward token to every from_pretrained branch; the worker passes the
command's hf_token.

* Scope created_by_scan to every HF cache the discard searches

created_by_scan used get_cache_path (active HF_HUB_CACHE only) while
/discard-remote-code deletes across active, legacy, and default caches. A
repo the user already had in a legacy/default cache was marked
scan-created and deleted on decline. Check all three caches for the repo
dir before declaring the scan created it.

* Scan the full .py closure of external auto_map repos

An auto_map cross-repo ref (owner/name--module.Class) only had its entry
file downloaded, but transformers also fetches that file's relative
imports from the same repo, so a dangerous helper.py was left outside the
scanned fingerprint. List each external repo's .py and scan the whole set
(plus the referenced entry files); fail closed if the repo cannot be
listed or fetched.

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

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

* Fail closed when a weight index cannot be fully read

_indexed_shard_paths treated a partial result as definitive: if one weight
index read cleanly but another failed transiently, it returned the shard
paths it did see. A flagged subdirectory pickle listed only by the index we
could not read would then be classed as "not a load input" and skipped,
re-opening the very fail-open this guard was added to close.

Return None whenever any index read is inconclusive, even if another read
cleanly, so the caller blocks the already-flagged subdir pickle. A repo that
ships no index files raises EntryNotFoundError for each (never inconclusive)
and still returns an empty set.

* Match cached repos case-insensitively in the created_by_scan guard

_repo_in_any_hf_cache resolved casing only against the active cache and then
probed every cache with an exact directory name. A case-variant already
present in a legacy or default cache (models--Unsloth--Foo for a scan of
unsloth/foo) was missed, so the repo was marked created_by_scan and deleted
on decline -- but discard_remote_code_download deletes case-insensitively,
so that delete would hit the user's pre-existing cache entry. Detect
case-insensitively too, mirroring the deletion path.

* Skip remote-code and security review for selected GGUF variants

validate_model ran the trust_remote_code and Hugging Face security-scan
preflight against the repo even when the selected artifact is a .gguf. A
GGUF loads through llama.cpp, which never executes the repo's auto_map
Python and never deserializes root pickle weights, so repo-level Transformers
artifacts (a config.json with auto_map, or an unsafe pytorch_model.bin next
to the .gguf in a mixed repo) are inert for that load. Gating the GGUF on
them is a false positive. Run both preflights only for non-GGUF loads.

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

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

* Scope the malware gate to actual load roots and serialized files

Two fixes to evaluate_file_security so it neither misses a load-path pickle nor
false-blocks an inert file:

- Honor subdirectory load roots. Spark-TTS / BiCodec call from_pretrained on the
  snapshot's LLM subdirectory, so a flagged pickle directly under it is a
  root-level load artifact there. A new load_subdirs parameter (set from the
  model's audio type via security_load_subdirs) reclassifies those files relative
  to the load root and looks for weight indexes under it, so a flagged shard in
  that subdir is no longer skipped as "not root-level".
- Exempt source files. A root .py is never deserialized by from_pretrained;
  executable repo code runs only through auto_map, which the remote-code consent
  gate scans. Flagging a Python helper here would false-block a repo that merely
  ships a build or train script.

* Scan a LoRA adapter and base as one consent unit, and gate MEDIUM code

A LoRA load runs both the adapter's and the base's repo code. The consent gate
scanned them separately and pinned one fingerprint per repo, so an adapter that
shipped its own auto_map code was either never shown in the dialog (which only
saw the base) or impossible to approve with the base's fingerprint.

evaluate_remote_code_consent_for_targets now scans all of a load's repos as a
single combined unit and pins ONE fingerprint over the union of their code, so
approving the load approves every repo's code together. evaluate_remote_code_consent
becomes a thin single-target wrapper, and an unscannable target fails the whole
load closed.

Also gate MEDIUM findings: like HIGH they now block pending pinned approval, so a
direct API caller cannot run flagged code by setting trust_remote_code=True
without consenting. Only a clean scan loads without a fingerprint.

* Preflight a LoRA load's adapter and base as one combined consent scan

scan_model_remote_code rewrote a LoRA adapter to its base and scanned only the
base for remote code, so the dialog never surfaced an adapter's own auto_map
code. Scan the adapter and base together through
preflight_remote_code_consent_for_targets, which pins one combined fingerprint
the worker gate accepts. The malware preflight is also scoped to each target's
load subdirectories.

* Apply combined consent and subdir-aware malware scan in load workers

Each load worker (inference, export, training) evaluated remote-code consent
once per target with a single shared fingerprint, so a LoRA adapter that ships
its own auto_map code could not be approved by the base's fingerprint. They now
scan the adapter and base together via evaluate_remote_code_consent_for_targets,
which pins one combined fingerprint over the union of their code. The malware
scan in each worker is also scoped to the model's load subdirectories so a
flagged pickle under a from_pretrained load subdir is not missed.

* Report a consistent trust_remote_code requirement after a model loads

validate_model reports requires_trust_remote_code from the YAML default OR the
raw auto_map, but the load, already-loaded, and status responses reported only
the YAML default. A custom-code model approved and loaded via auto_map was then
reported as not requiring trust_remote_code, so the frontend stored false and a
later retry or rollback sent trust_remote_code=false and failed.

A shared resolver reports the same requirement for a loaded model (a value
stored at load time, else the trust_remote_code the load used, else the YAML
default, else the raw auto_map check), and the load response persists it so the
status and already-loaded paths stay consistent. The selected-GGUF security
review is also scoped to the model's load subdirectories.

* Run the consent gate on training resume and for YAML-only trust_remote_code

Three frontend gaps left a model loading without the trust_remote_code it needs:

- The shared consent helper returned early when the scan found no auto_map and no
  unsafe files, dropping a requirement that comes from a model's Studio YAML
  default (e.g. GLM-4.7-Flash). It now grants the caller's requirement with an
  empty pin instead of sending trust_remote_code=false.
- Resume-from-history called startTraining directly with no consent gate, so a
  resumed run whose model needs custom code (or an old run with no approved
  fingerprint) hit the worker block with no dialog. It now runs the same gate as
  a fresh start.
- HF export passed requiresTrustRemoteCode=false for every HF source, so a
  YAML-only model could not flip the flag before export. It now signals the
  requirement for HF sources.

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

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

* Cover both LoRA repos in validate, report GGUF as inert, purge all declined repos

Three follow-on gaps from the combined adapter+base consent work:

- validate_model resolved requires_trust_remote_code from the base alone, so a
  LoRA adapter that ships its OWN auto_map code (with a plain base) was reported
  as not needing trust_remote_code and the consent dialog never opened. It now
  checks the [adapter, base] target set, matching the scan route and the workers
  (which already gate both) and the security review already running over both.

- The already-loaded, loaded, and status responses for a selected GGUF reported
  requires_trust_remote_code from the model's YAML default. A GGUF loads through
  llama.cpp, which never executes the repo's auto_map Python, so the requirement
  is inert for that load. They now report False, matching validate_model (which
  already skips both gates for GGUF) so a status refresh cannot flip the flag
  back on.

- The remote-code scan downloads both the adapter's and the base's config, but
  created_by_scan tracked only the primary, so a base the scan was first to pull
  into the cache was left on disk when the user declined. The scan now reports
  scan_created_repos (every repo it newly cached) and the decline cleanup purges
  each; created_by_scan stays for older clients. The frontend falls back to the
  primary flag when the list is absent.

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

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

* Scan the repo the load fetches, purge external code on decline, harden consent pins

Six follow-on hardening fixes from a fresh review pass over the gate:

- The malware gate scanned the literal "Spark-TTS-0.5B/LLM" alias, but the trainer
  downloads it as unsloth/Spark-TTS-0.5B and loads LLM/, so the alias 404'd and
  failed open, missing a flagged LLM/ pickle. evaluate_file_security now resolves
  the alias to the repo the loader fetches and scans LLM/ as a load root.

- security_load_subdirs relied only on tokenizer detection, which fails on an
  unresolved alias or offline; it now also honors the Studio YAML audio_type
  default, so a BiCodec LLM/ load root is not missed.

- The remote-code scan downloads external auto_map repos (owner/name--module.Class),
  but the decline cleanup tracked only the model/adapter/base, leaving the external
  untrusted code cached. The scan now enumerates external auto_map repos and reports
  the ones it created in scan_created_repos, so a decline purges them too.

- External auto_map refs failed the whole load closed on a stale or mis-derived
  dotted ref (sub.mod.py vs the real sub/mod.py) even though the actual file was
  present and scanned. They now drop such refs when the repo listing is real, exactly
  like the own-repo path; an empty/incomplete listing still fetches and fails closed.

- The combined consent fingerprint keyed code by the raw target string, so the scan
  endpoint's canonicalized casing and a worker's raw user input produced different
  pins for identical code, rejecting a valid approval. Hub repo ids are now folded to
  lowercase in the key (local paths stay case-sensitive), so the pin tracks the code.

- Export threaded hf_token into the weight load but not into detect_audio_type /
  is_vision_model, so a gated multimodal base 404'd in detection and fell through to
  the text loader. Both probes now use the same token.

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

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

* Thread the token through check-vision and guard the gate's parallel sites

The /check-vision endpoint classified a model without the hf_token, so a gated or
private vision model 404'd in the probe and was reported as a plain text model --
the same dropped-token shape as the export probes, at a sibling site. It now passes
the token like the neighboring /check-embedding endpoint.

Add deterministic consistency guards (tests/test_security_gate_consistency.py) that
enumerate the gate's parallel sites mechanically instead of relying on a review to
spot a missed sibling: every is_vision_model / is_embedding_model / detect_audio_type
caller under routes/ and core/ must thread the token, every GGUF response must report
trust_remote_code via the resolver or False (never the raw YAML default), and every
load worker that runs the malware or consent gate must resolve the LoRA base. A new
site that drops the token or mis-reports the requirement now fails CI directly.

* Narrow the LLM alias rewrite and make audio detection token-aware

Three fixes from the confirmatory review, one a regression from the previous round:

- _load_scan_target rewrote EVERY remote repo ending in "/LLM" to unsloth/<parent>,
  so a real third-party repo named "<owner>/LLM" was scanned as unsloth/<owner>
  while the loader still fetched the real repo -- a fail-open hole introduced when
  the Spark-TTS alias handling was added. It now rewrites only a registry-known
  bicodec alias; every other "/LLM" repo is scanned as itself.

- detect_audio_type cached results under the bare model name, so an unauthenticated
  probe of a gated/private repo cached None and poisoned a later authenticated call
  with the token. The cache is now keyed by (normalized_name, token_fingerprint),
  matching the vision cache.

- The training fallback /check-vision call dropped the hf_token, misclassifying a
  gated/private VLM when the config endpoint failed. It now passes the token, like
  the getModelConfig call it falls back from; checkEmbeddingModel takes the token too.

Extend the consistency guards: every capability cache must be keyed by a tuple
including the token, so a cache re-declared as Dict[str, ...] fails CI.

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

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

* Document the broad .py scan as deliberate and enforce it with a test

The remote-code scanner scans every .py in a repo once an auto_map exists, not
just the auto_map entry's static import closure. This is intentional: the entry
module can reach a sibling via an absolute import, importlib, or exec, none of
which a static relative-import closure follows, so closure-only scanning would be
a real bypass of a load-time RCE gate. The broad scan never under-scans; the cost
is that an unrelated benign script can over-block, which is the safe failure
direction (HIGH stays approvable; only CRITICAL hard-blocks).

Spell this out at both the local and remote scan sites so the choice reads as
deliberate, and add a test asserting an unrelated, never-imported .py is still
scanned -- so a future narrowing to the static closure fails CI.

* Purge a declined remote LoRA adapter the scan downloaded

scan_model_remote_code probed the created-by-scan state AFTER resolving the base,
but get_base_model_from_lora_identifier downloads a remote adapter's own
adapter_config.json, so the adapter looked already-cached and was dropped from
scan_created_repos. On decline the adapter -- including the auto_map .py the
preflight fetched -- was left on disk, defeating the "untrusted code is not left
on disk" guarantee for the adapter itself.

Snapshot the primary's cache state BEFORE base resolution and use it when marking
the adapter scan-created; on any probe error treat it as pre-existing so a decline
never deletes it. The base and external repos are unaffected (their configs are not
downloaded before their own probe). Add a test that models the mid-scan download
side effect, which the prior static-stub tests did not.

* Clear remote-code approval when the training model changes

Switching the training model from an approved custom-code model to a clean one
kept the previous model's trust_remote_code=true and approved fingerprint in the
store: setSelectedModel reset visionImageSize on a true switch but not the
remote-code approval. The clean model then trained with trust_remote_code=true,
which bypasses the compiler and disables fused cross-entropy.

Reset trustRemoteCode and approvedRemoteCodeFingerprint on a true model switch.
The new model's own YAML default is re-applied by loadAndApplyModelDefaults, and a
custom-code model still re-opens the consent dialog before training starts, so the
only change is that a clean model no longer inherits a stale approval.

* Trim verbose comments across the model-fetching hardening changes

Condense the explanatory comments and docstrings introduced across the
trust_remote_code consent gate, the malware/unsafe-file gate, the remote-code
scanner, the load workers, the model routes, and the security frontend into
fewer, tighter lines while preserving every security rationale (fail-open vs
fail-closed direction, the deliberate broad-scan anti-bypass note, the
empty-vs-unscannable distinction, stale-ref handling, and the alias-rewrite
spoof guard).

Comments and docstrings only. No code, logic, identifiers, or test behaviour
changed; verified comment-only via the AST/TypeScript checker (40/40), with the
backend test suite and frontend tsc green.

* Do not cache transient audio-detection failures

detect_audio_type cached _detect_audio_from_tokenizer's result
unconditionally, so a transient read failure (network error or 5xx,
returned as None) poisoned the cache and the later successful probe never
ran. Mirror the vision cache: _detect_audio_from_tokenizer now returns
(audio_type, definitive) and the caller caches only definitive results.

A read that succeeds with no audio tokens, or clean 404s for every
tokenizer path, stays a cacheable None; only a genuine transient failure
(connection error, timeout, 5xx, malformed body) skips the cache so the
next call retries.

---------

Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-06-18 05:39:52 -07:00 • committed by GitHub
commit 0533efe3f8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
53 changed files with 6243 additions and 257 deletions
Download patch file Download diff file Expand all files Collapse all files

19
studio/backend/core/export/export.py
View file

@ -145,13 +145,19 @@ class ExportBackend:
max_seq_length: int = 2048,
load_in_4bit: bool = True,
trust_remote_code: bool = False,
hf_token: Optional[str] = None,
) -> Tuple[bool, str]:
"""
Load a checkpoint for export.
``hf_token`` authenticates the actual weight load for gated/private
checkpoints, matching the token the worker used for the security preflight
(otherwise a gated repo passes scanning then 401s at from_pretrained).
Returns:
Tuple of (success: bool, message: str)
"""
token = hf_token if hf_token and hf_token.strip() else None
try:
logger.info(f"Loading checkpoint: {checkpoint_path}")
@ -169,8 +175,10 @@ class ExportBackend:
model_id = base_model or checkpoint_path
self._audio_type = detect_audio_type(model_id)
self.is_vision = not self._audio_type and is_vision_model(model_id)
# Token the type-detection probes too, else a gated multimodal base
# 404s here and falls through to the text loader.
self._audio_type = detect_audio_type(model_id, hf_token = token)
self.is_vision = not self._audio_type and is_vision_model(model_id, hf_token = token)
if self._audio_type == "csm":
from unsloth import FastModel
@ -184,6 +192,7 @@ class ExportBackend:
auto_model = CsmForConditionalGeneration,
load_in_4bit = False,
trust_remote_code = trust_remote_code,
token = token,
)
elif self._audio_type == "whisper":
@ -197,6 +206,7 @@ class ExportBackend:
load_in_4bit = False,
auto_model = WhisperForConditionalGeneration,
trust_remote_code = trust_remote_code,
token = token,
)
elif self._audio_type == "snac":
@ -207,6 +217,7 @@ class ExportBackend:
dtype = None,
load_in_4bit = load_in_4bit,
trust_remote_code = trust_remote_code,
token = token,
)
elif self._audio_type == "bicodec":
@ -218,6 +229,7 @@ class ExportBackend:
dtype = None if _IS_MLX else torch.float32,
load_in_4bit = False,
trust_remote_code = trust_remote_code,
token = token,
)
elif self._audio_type == "dac":
@ -228,6 +240,7 @@ class ExportBackend:
max_seq_length = max_seq_length,
load_in_4bit = False,
trust_remote_code = trust_remote_code,
token = token,
)
elif self.is_vision:
@ -238,6 +251,7 @@ class ExportBackend:
dtype = None,
load_in_4bit = load_in_4bit,
trust_remote_code = trust_remote_code,
token = token,
)
tokenizer = processor # vision: processor acts as tokenizer
@ -249,6 +263,7 @@ class ExportBackend:
dtype = None,
load_in_4bit = load_in_4bit,
trust_remote_code = trust_remote_code,
token = token,
)
if _IS_MLX:

2
studio/backend/core/export/orchestrator.py
View file

@ -301,6 +301,7 @@ class ExportOrchestrator:
max_seq_length: int = 2048,
load_in_4bit: bool = True,
trust_remote_code: bool = False,
approved_remote_code_fingerprint: Optional[str] = None,
hf_token: Optional[str] = None,
) -> Tuple[bool, str]:
"""Load a checkpoint for export.
@ -312,6 +313,7 @@ class ExportOrchestrator:
"max_seq_length": max_seq_length,
"load_in_4bit": load_in_4bit,
"trust_remote_code": trust_remote_code,
"approved_remote_code_fingerprint": approved_remote_code_fingerprint,
"hf_token": hf_token,
}

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

@ -188,10 +188,16 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None:
# Auto-enable trust_remote_code for NemotronH/Nano models.
if not trust_remote_code:
from utils.security.trusted_org import is_trusted_org_repo
_NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano")
_cp_lower = checkpoint_path.lower()
if any(sub in _cp_lower for sub in _NEMOTRON_TRUST_SUBSTRINGS) and (
_cp_lower.startswith("unsloth/") or _cp_lower.startswith("nvidia/")
if (
any(sub in _cp_lower for sub in _NEMOTRON_TRUST_SUBSTRINGS)
and (_cp_lower.startswith("unsloth/") or _cp_lower.startswith("nvidia/"))
# Genuine first-party Hub repo only (not a local/spoof name starting
# with "unsloth/"); authenticated so private repos resolve.
and is_trusted_org_repo(checkpoint_path, hf_token = cmd.get("hf_token"))
):
trust_remote_code = True
logger.info(
@ -199,6 +205,81 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None:
checkpoint_path,
)
# Malware gate: a poisoned pickle deserializes on load even with
# trust_remote_code False, so check HF's security scan (metadata-only) every
# load. Local checkpoints have no Hub scan and are skipped in the helper; a
# LoRA merges its base weights, so gate that repo too.
from utils.security import evaluate_file_security, security_load_subdirs
malware_targets = [checkpoint_path]
try:
from utils.models.model_config import get_base_model_from_lora_identifier
# Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too.
_base = get_base_model_from_lora_identifier(checkpoint_path, cmd.get("hf_token"))
if _base:
malware_targets.append(_base)
except Exception as exc:
logger.debug("Could not resolve LoRA base for malware scan: %s", exc)
_hf_token = cmd.get("hf_token")
for target in dict.fromkeys(malware_targets):
_fs = evaluate_file_security(
target, hf_token = _hf_token, load_subdirs = security_load_subdirs(target, _hf_token)
)
if _fs.blocked:
_send_response(
resp_queue,
{
"type": "loaded",
"success": False,
"message": _fs.reason,
"error_kind": "malware_blocked",
"security": _fs.response_payload(),
"ts": time.time(),
},
)
return
# Consent gate: scan auto_map code before it runs; block CRITICAL/HIGH unless
# pinned-approved. A LoRA merges its base model, whose code runs, so gate it too.
if trust_remote_code:
from utils.security import evaluate_remote_code_consent_for_targets
consent_targets = [checkpoint_path]
try:
from utils.models.model_config import get_base_model_from_lora_identifier
# Resolve a local or remote adapter's base so its base repo is gated too.
base_model = get_base_model_from_lora_identifier(checkpoint_path, cmd.get("hf_token"))
if base_model:
consent_targets.append(base_model)
except Exception as exc:
logger.debug("Could not resolve LoRA base for consent scan: %s", exc)
# Scan adapter + base as one combined unit, pinned by a single fingerprint.
_rc = evaluate_remote_code_consent_for_targets(
consent_targets,
hf_token = cmd.get("hf_token"),
trust_remote_code = True,
approved_fingerprint = cmd.get("approved_remote_code_fingerprint"),
)
if _rc.blocked:
_send_response(
resp_queue,
{
"type": "loaded",
"success": False,
"message": (
f"Checkpoint '{_rc.model_name}' ships custom code flagged as "
f"{_rc.max_severity} by the security scan. Review and "
f"approve it to proceed."
),
"error_kind": "remote_code_blocked",
"remote_code": _rc.response_payload(),
"ts": time.time(),
},
)
return
try:
_send_response(
resp_queue,
@ -214,6 +295,7 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None:
max_seq_length = max_seq_length,
load_in_4bit = load_in_4bit,
trust_remote_code = trust_remote_code,
hf_token = cmd.get("hf_token"),
)
_send_response(

5
studio/backend/core/inference/orchestrator.py
View file

@ -636,6 +636,7 @@ class InferenceOrchestrator:
load_in_4bit: bool = True,
hf_token: Optional[str] = None,
trust_remote_code: bool = False,
approved_remote_code_fingerprint: Optional[str] = None,
gpu_ids: Optional[list[int]] = None,
) -> bool:
"""Load a model for inference.
@ -659,6 +660,7 @@ class InferenceOrchestrator:
"hf_token": hf_token or "",
"gguf_variant": getattr(config, "gguf_variant", None),
"trust_remote_code": trust_remote_code,
"approved_remote_code_fingerprint": approved_remote_code_fingerprint,
"gpu_ids": gpu_ids,
}
resolved_gpu_ids, gpu_selection = prepare_gpu_selection(
@ -738,7 +740,8 @@ class InferenceOrchestrator:
logger.info("Model '%s' loaded successfully in subprocess", model_name)
return True
else:
error = resp.get("error", "Failed to load model")
# Worker reports failures (consent gate included) under "message".
error = resp.get("message") or resp.get("error") or "Failed to load model"
self.loading_models.discard(model_name)
self.active_model_name = None
self.models.clear()

69
studio/backend/core/inference/worker.py
View file

@ -260,13 +260,19 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
# Auto-enable trust_remote_code only for NemotronH/Nano (config parsing
# bugs require it). Must NOT match Llama-Nemotron (standard Llama arch).
from utils.security.trusted_org import is_trusted_org_repo
_NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano")
trust_remote_code = config.get("trust_remote_code", False)
if not trust_remote_code:
model_name = config["model_name"]
_mn_lower = model_name.lower()
if any(sub in _mn_lower for sub in _NEMOTRON_TRUST_SUBSTRINGS) and (
_mn_lower.startswith("unsloth/") or _mn_lower.startswith("nvidia/")
if (
any(sub in _mn_lower for sub in _NEMOTRON_TRUST_SUBSTRINGS)
and (_mn_lower.startswith("unsloth/") or _mn_lower.startswith("nvidia/"))
# Genuine first-party Hub repo only (not a local/spoof name starting
# with "unsloth/"); authenticated so private repos resolve.
and is_trusted_org_repo(model_name, hf_token = hf_token)
):
trust_remote_code = True
logger.info(
@ -274,6 +280,65 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
model_name,
)
# Malware gate: a poisoned pickle deserializes during from_pretrained even
# with trust_remote_code False, so check HF's security scan (metadata-only)
# every load. For a LoRA, gate the base whose weights deserialize.
from utils.security import evaluate_file_security, security_load_subdirs
malware_targets = [config["model_name"]]
if mc.is_lora and getattr(mc, "base_model", None):
malware_targets.append(str(mc.base_model))
for target in dict.fromkeys(malware_targets):
_fs = evaluate_file_security(
target, hf_token = hf_token, load_subdirs = security_load_subdirs(target, hf_token)
)
if _fs.blocked:
_send_response(
resp_queue,
{
"type": "loaded",
"success": False,
"message": _fs.reason,
"error_kind": "malware_blocked",
"security": _fs.response_payload(),
"ts": time.time(),
},
)
return
# Consent gate: scan auto_map code before it runs; block CRITICAL/HIGH
# unless pinned-approved. For a LoRA, gate the base whose code runs.
if trust_remote_code:
from utils.security import evaluate_remote_code_consent_for_targets
consent_targets = [config["model_name"]]
if mc.is_lora and getattr(mc, "base_model", None):
consent_targets.append(str(mc.base_model))
# Scan adapter + base as one unit, pinned by a single fingerprint.
_rc = evaluate_remote_code_consent_for_targets(
consent_targets,
hf_token = hf_token,
trust_remote_code = True,
approved_fingerprint = config.get("approved_remote_code_fingerprint"),
)
if _rc.blocked:
_send_response(
resp_queue,
{
"type": "loaded",
"success": False,
"message": (
f"Model '{_rc.model_name}' ships custom code flagged as "
f"{_rc.max_severity} by the security scan. Review "
f"and approve it to proceed."
),
"error_kind": "remote_code_blocked",
"remote_code": _rc.response_payload(),
"ts": time.time(),
},
)
return
# Heartbeat every 30s so the orchestrator knows we're alive during slow loads.
xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1"

1
studio/backend/core/training/training.py
View file

@ -315,6 +315,7 @@ class TrainingBackend:
"tensorboard_dir": kwargs.get("tensorboard_dir", "runs"),
"resume_from_checkpoint": kwargs.get("resume_from_checkpoint"),
"trust_remote_code": kwargs.get("trust_remote_code", False),
"approved_remote_code_fingerprint": kwargs.get("approved_remote_code_fingerprint"),
"gpu_ids": kwargs.get("gpu_ids"),
"s3_config": kwargs.get("s3_config"),
# Flipped to True only by the HTTP-fallback respawn after a stall.

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

@ -1453,6 +1453,75 @@ def _run_mlx_training(event_queue, stop_queue, config):
model_random_state = random_seed if _model_seed is None else int(_model_seed)
_lora_seed = config.get("lora_random_state")
lora_random_state = random_seed if _lora_seed is None else int(_lora_seed)
# Malware gate (MLX): a poisoned pickle deserializes on load even with
# trust_remote_code False, so check HF's security scan (metadata-only) first.
# For a LoRA, gate the base whose weights deserialize.
from utils.security import evaluate_file_security
malware_targets = [model_name]
try:
from utils.models.model_config import get_base_model_from_lora_identifier
# Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too.
_base = get_base_model_from_lora_identifier(model_name, config.get("hf_token") or None)
if _base:
malware_targets.append(_base)
except Exception as exc:
logger.debug("Could not resolve LoRA base for malware scan: %s", exc)
from utils.security import security_load_subdirs
for target in dict.fromkeys(malware_targets):
_fs = evaluate_file_security(
target, hf_token = hf_token, load_subdirs = security_load_subdirs(target, hf_token)
)
if _fs.blocked:
_send(
"error",
error = _fs.reason,
error_kind = "malware_blocked",
security = _fs.response_payload(),
)
return
# Consent gate (MLX): the CUDA path gates in run_training_process, but MLX returns
# before that, so scan auto_map code here before FastMLXModel runs it. Block
# CRITICAL/HIGH unless pinned-approved; for a LoRA, gate the base whose code runs.
if config.get("trust_remote_code", False):
from utils.security import evaluate_remote_code_consent_for_targets
consent_targets = [model_name]
try:
from utils.models.model_config import get_base_model_from_lora_identifier
# Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too.
base_model = get_base_model_from_lora_identifier(
model_name, config.get("hf_token") or None
)
if base_model:
consent_targets.append(base_model)
except Exception as exc:
logger.debug("Could not resolve LoRA base for consent scan: %s", exc)
# Scan adapter + base as one combined unit, pinned by a single fingerprint.
_rc = evaluate_remote_code_consent_for_targets(
consent_targets,
hf_token = hf_token,
trust_remote_code = True,
approved_fingerprint = config.get("approved_remote_code_fingerprint"),
)
if _rc.blocked:
_send(
"error",
error = (
f"Model '{_rc.model_name}' ships custom code flagged as "
f"{_rc.max_severity} by the security scan. Review it and "
f"re-run with approval to proceed.\n\n{_rc.findings_summary}"
),
error_kind = "remote_code_blocked",
remote_code = _rc.response_payload(),
)
return
model, tokenizer = FastMLXModel.from_pretrained(
model_name,
load_in_4bit = config.get("load_in_4bit", True),
@ -2100,11 +2169,16 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
# NemotronH needs trust_remote_code=True to work around config-parsing bugs.
# Other 5.x models are native and don't need it (it bypasses the compiler,
# disabling fused CE). Must NOT match Llama-Nemotron (standard Llama arch).
from utils.security.trusted_org import is_trusted_org_repo
_NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano")
_lowered = model_name.lower()
if (
any(sub in _lowered for sub in _NEMOTRON_TRUST_SUBSTRINGS)
and (_lowered.startswith("unsloth/") or _lowered.startswith("nvidia/"))
# Confirm a genuine first-party Hub repo (not a local/spoofed name starting
# with "unsloth/"); authenticated so private first-party repos resolve.
and is_trusted_org_repo(model_name, hf_token = config.get("hf_token") or None)
and not config.get("trust_remote_code", False)
):
config["trust_remote_code"] = True
@ -2113,6 +2187,81 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
model_name,
)
# 1a. Malware gate: a poisoned pickle deserializes on load even with
# trust_remote_code False, so check HF's security scan (metadata-only) first.
# For a LoRA, gate the base whose weights deserialize.
from utils.security import evaluate_file_security
malware_targets = [model_name]
try:
from utils.models.model_config import get_base_model_from_lora_identifier
# Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too.
_base = get_base_model_from_lora_identifier(model_name, config.get("hf_token") or None)
if _base:
malware_targets.append(_base)
except Exception as exc:
logger.debug("Could not resolve LoRA base for malware scan: %s", exc)
from utils.security import security_load_subdirs
_ls_hf = config.get("hf_token") or None
for target in dict.fromkeys(malware_targets):
_fs = evaluate_file_security(
target, hf_token = _ls_hf, load_subdirs = security_load_subdirs(target, _ls_hf)
)
if _fs.blocked:
event_queue.put(
{
"type": "error",
"error": _fs.reason,
"error_kind": "malware_blocked",
"security": _fs.response_payload(),
"ts": time.time(),
}
)
return
# 1a'. Consent gate: scan auto_map Python before it runs; refuse CRITICAL/HIGH
# unless pinned-approved.
if config.get("trust_remote_code", False):
from utils.security import evaluate_remote_code_consent_for_targets
# A LoRA adapter's base is where custom code runs, so gate it too.
consent_targets = [model_name]
try:
from utils.models.model_config import get_base_model_from_lora_identifier
# Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too.
base_model = get_base_model_from_lora_identifier(
model_name, config.get("hf_token") or None
)
if base_model:
consent_targets.append(base_model)
except Exception as exc:
logger.debug("Could not resolve LoRA base for consent scan: %s", exc)
# Scan adapter + base as one combined unit, pinned by a single fingerprint.
_rc = evaluate_remote_code_consent_for_targets(
consent_targets,
hf_token = config.get("hf_token") or None,
trust_remote_code = True,
approved_fingerprint = config.get("approved_remote_code_fingerprint"),
)
if _rc.blocked:
event_queue.put(
{
"type": "error",
"error": (
f"Model '{_rc.model_name}' ships custom code flagged as "
f"{_rc.max_severity} by the security scan. Review it and "
f"re-run with approval to proceed.\n\n{_rc.findings_summary}"
),
"error_kind": "remote_code_blocked",
"remote_code": _rc.response_payload(),
"ts": time.time(),
}
)
return
# ── 1b. Install fast-path kernel libraries for the chosen model.
# 1) causal-conv1d ALWAYS runs eagerly via the substring path: some SSM
# modeling files lazy_load it without calling is_causal_conv1d_available.
@ -3064,6 +3213,73 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
training_type = config.get("training_type", "LoRA/QLoRA")
use_lora = training_type == "LoRA/QLoRA"
# Malware gate (embedding): a poisoned pickle deserializes on load even with
# trust_remote_code False, so check HF's security scan (metadata-only) first.
# For a LoRA, gate the base whose weights deserialize.
from utils.security import evaluate_file_security
malware_targets = [model_name]
try:
from utils.models.model_config import get_base_model_from_lora_identifier
_base = get_base_model_from_lora_identifier(model_name, hf_token)
if _base:
malware_targets.append(_base)
except Exception as exc:
logger.debug("Could not resolve LoRA base for malware scan: %s", exc)
from utils.security import security_load_subdirs
for target in dict.fromkeys(malware_targets):
_fs = evaluate_file_security(
target, hf_token = hf_token, load_subdirs = security_load_subdirs(target, hf_token)
)
if _fs.blocked:
event_queue.put(
{
"type": "error",
"error": _fs.reason,
"error_kind": "malware_blocked",
"security": _fs.response_payload(),
"ts": time.time(),
}
)
return
# Consent gate (embedding): scan any auto_map code before it runs; block
# CRITICAL/HIGH unless pinned-approved. A no-op without auto_map.
if config.get("trust_remote_code", False):
from utils.security import evaluate_remote_code_consent_for_targets
consent_targets = [model_name]
try:
from utils.models.model_config import get_base_model_from_lora_identifier
_cbase = get_base_model_from_lora_identifier(model_name, hf_token)
if _cbase:
consent_targets.append(_cbase)
except Exception as exc:
logger.debug("Could not resolve LoRA base for consent scan: %s", exc)
# Scan adapter + base as one combined unit, pinned by a single fingerprint.
_rc = evaluate_remote_code_consent_for_targets(
consent_targets,
hf_token = hf_token,
trust_remote_code = True,
approved_fingerprint = config.get("approved_remote_code_fingerprint"),
)
if _rc.blocked:
event_queue.put(
{
"type": "error",
"error": (
f"Model '{_rc.model_name}' ships custom code flagged as "
f"{_rc.max_severity} by the security scan. Review it and "
f"re-run with approval to proceed.\n\n{_rc.findings_summary}"
),
"error_kind": "remote_code_blocked",
"remote_code": _rc.response_payload(),
"ts": time.time(),
}
)
return
model = FastSentenceTransformer.from_pretrained(
model_name = model_name,
max_seq_length = max_seq_length,

8
studio/backend/models/export.py
View file

@ -51,6 +51,14 @@ class LoadCheckpointRequest(BaseModel):
False,
description = "Allow loading models with custom code. Only enable for checkpoints/base models you trust.",
)
approved_remote_code_fingerprint: Optional[str] = Field(
None,
description = "sha256 fingerprint from the remote-code scan, pinning user approval of this exact custom-code version.",
)
hf_token: Optional[str] = Field(
None,
description = "Hugging Face token used to scan/load gated checkpoints and their base models.",
)
class ExportStatusResponse(BaseModel):

9
studio/backend/models/inference.py
View file

@ -42,6 +42,10 @@ class LoadRequest(BaseModel):
False,
description = "Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust.",
)
approved_remote_code_fingerprint: Optional[str] = Field(
None,
description = "sha256 fingerprint from the remote-code scan, pinning user approval of this exact custom-code version.",
)
chat_template_override: Optional[str] = Field(
None,
description = "Custom Jinja2 chat template to use instead of the model's default",
@ -149,6 +153,11 @@ class ValidateModelResponse(BaseModel):
False,
description = "Whether the model defaults require trust_remote_code to be enabled for loading.",
)
requires_security_review: bool = Field(
False,
description = "Whether Hugging Face's security scan flagged unsafe files (e.g. a "
"malicious pickle), so the load is hard-blocked pending review.",
)
context_length: Optional[int] = Field(
None,
description = "Native training context length, read from the GGUF header when the file "

4
studio/backend/models/training.py
View file

@ -108,6 +108,10 @@ class TrainingStartRequest(BaseModel):
False,
description = "Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust.",
)
approved_remote_code_fingerprint: Optional[str] = Field(
None,
description = "sha256 fingerprint from the remote-code scan, pinning user approval of this exact custom-code version.",
)
# Dataset parameters
hf_dataset: Optional[str] = Field(None, description = "HuggingFace dataset identifier")

2
studio/backend/routes/export.py
View file

@ -93,6 +93,8 @@ async def load_checkpoint(
max_seq_length = request.max_seq_length,
load_in_4bit = request.load_in_4bit,
trust_remote_code = request.trust_remote_code,
approved_remote_code_fingerprint = request.approved_remote_code_fingerprint,
hf_token = request.hf_token,
)
if not success:

143
studio/backend/routes/inference.py
View file

@ -2037,9 +2037,8 @@ async def load_model(
audio_type = _gguf_audio,
has_audio_input = getattr(llama_backend, "_has_audio_input", False),
inference = inference_config,
requires_trust_remote_code = bool(
inference_config.get("trust_remote_code", False)
),
# GGUF loads via llama.cpp: auto_map never executes, so inert (matches validate_model).
requires_trust_remote_code = False,
context_length = llama_backend.context_length,
max_context_length = llama_backend.max_context_length,
native_context_length = llama_backend.native_context_length,
@ -2086,8 +2085,8 @@ async def load_model(
audio_type = _model_info.get("audio_type"),
has_audio_input = _model_info.get("has_audio_input", False),
inference = inference_config,
requires_trust_remote_code = bool(
inference_config.get("trust_remote_code", False)
requires_trust_remote_code = _resolve_loaded_trust_remote_code(
backend.active_model_name, _model_info, inference_config
),
supports_reasoning = _sf_supports_reasoning,
reasoning_style = _sf_reasoning_style,
@ -2325,7 +2324,8 @@ async def load_model(
audio_type = _gguf_audio,
has_audio_input = llama_backend._has_audio_input,
inference = inference_config,
requires_trust_remote_code = bool(inference_config.get("trust_remote_code", False)),
# GGUF loads via llama.cpp: auto_map never executes, so inert (matches validate_model).
requires_trust_remote_code = False,
context_length = llama_backend.context_length,
max_context_length = llama_backend.max_context_length,
native_context_length = llama_backend.native_context_length,
@ -2418,6 +2418,7 @@ async def load_model(
load_in_4bit = load_in_4bit,
hf_token = request.hf_token,
trust_remote_code = request.trust_remote_code,
approved_remote_code_fingerprint = request.approved_remote_code_fingerprint,
gpu_ids = effective_gpu_ids,
)
@ -2458,6 +2459,23 @@ async def load_model(
# Classify reasoning/tool flags via the GGUF sniffer.
_sf_flags = _detect_safetensors_features(backend, _chat_template)
# Report validate_model's requirement (raw auto_map OR YAML) plus the value the
# load used, and persist it, so a later retry/rollback doesn't send
# trust_remote_code=false for a custom-code model (and status reports it too).
_requires_rc = _resolve_loaded_trust_remote_code(
config.identifier,
None,
inference_config,
request.hf_token,
trust_remote_code_used = bool(getattr(request, "trust_remote_code", False)),
)
try:
backend.models.setdefault(config.identifier, {})["requires_trust_remote_code"] = (
_requires_rc
)
except Exception:
pass
return LoadResponse(
status = "loaded",
model = model_log_label if native_grant_backed else config.identifier,
@ -2469,7 +2487,7 @@ async def load_model(
audio_type = config.audio_type,
has_audio_input = config.has_audio_input,
inference = inference_config,
requires_trust_remote_code = bool(inference_config.get("trust_remote_code", False)),
requires_trust_remote_code = _requires_rc,
supports_reasoning = _sf_flags["supports_reasoning"],
reasoning_style = _sf_flags["reasoning_style"],
reasoning_always_on = _sf_flags["reasoning_always_on"],
@ -2522,6 +2540,76 @@ async def load_model(
raise HTTPException(status_code = 500, detail = f"Failed to load model: {msg}")
def _requires_trust_remote_code_for_model(
model_identifier: str, hf_token: Optional[str] = None
) -> bool:
"""Whether loading this model would execute custom repo code, so the consent
dialog must run first. True if the Studio YAML default enables
``trust_remote_code`` OR the raw config declares an ``auto_map`` (Hub/local,
config.json or tokenizer_config.json). Reads raw JSON only; never imports
model code."""
from utils.inference import load_inference_config
try:
if bool(load_inference_config(model_identifier).get("trust_remote_code", False)):
return True
except Exception:
pass
try:
from utils.security.consent import _config_has_auto_map
return _config_has_auto_map(model_identifier, hf_token) is True
except Exception:
return False
def _resolve_loaded_trust_remote_code(
model_id,
model_info,
inference_config,
hf_token = None,
trust_remote_code_used = False,
) -> bool:
"""TRC requirement to report for an ALREADY-LOADED model, consistent with
``validate_model``.
``validate_model`` reports ``requires_trust_remote_code`` from
``_requires_trust_remote_code_for_model`` (YAML default OR raw ``auto_map``), but
the load / already-loaded / status responses historically reported only the YAML
default. That dropped raw-``auto_map`` models: after approving and loading one, the
response said ``false``, so the frontend stored ``false`` and a later retry/rollback
sent ``trust_remote_code=false`` and failed.
Resolution order: a value stored on the model at load time (so a status refresh does
not re-derive it) -> the trust_remote_code the load actually used -> the YAML default
-> the raw ``auto_map`` check (reads the loaded model's cached config; no network)."""
stored = (model_info or {}).get("requires_trust_remote_code")
if stored is not None:
return bool(stored)
if trust_remote_code_used or bool((inference_config or {}).get("trust_remote_code", False)):
return True
try:
return bool(_requires_trust_remote_code_for_model(model_id, hf_token))
except Exception:
return False
def _requires_security_review_for_model(
model_identifier: str, hf_token: Optional[str] = None
) -> bool:
"""Whether Hugging Face's security scan flagged unsafe files for this repo, so
the consent dialog must open as a hard block before loading. Metadata-only;
never downloads the flagged files. Fails open (False) on any error."""
try:
from utils.security import evaluate_file_security, security_load_subdirs
return evaluate_file_security(
model_identifier,
hf_token,
load_subdirs = security_load_subdirs(model_identifier, hf_token),
).blocked
except Exception:
return False
@router.post("/validate", response_model = ValidateModelResponse)
async def validate_model(
request: ValidateModelRequest, current_subject: str = Depends(get_current_subject)
@ -2550,7 +2638,34 @@ async def validate_model(
detail = f"Invalid model identifier: {model_log_label}",
)
# Both checks cover the [adapter, base] set (matching the scan route and workers):
# either repo can ship auto_map code or a poisoned pickle.
security_targets = [config.identifier]
try:
from utils.models.model_config import get_base_model_from_lora_identifier
# Resolve a LOCAL or REMOTE adapter's base so its code/weights are reviewed too.
_base = get_base_model_from_lora_identifier(model_identifier, request.hf_token)
if _base:
security_targets.append(_base)
except Exception:
pass
security_targets = list(dict.fromkeys(security_targets))
is_gguf = getattr(config, "is_gguf", False)
# A selected GGUF loads via llama.cpp: auto_map Python and root pickle weights in a
# mixed repo are inert for this load, so gating on them is a false positive. Only
# run the remote-code/security preflight for non-GGUF loads.
requires_trust_remote_code = False
requires_security_review = False
if not is_gguf:
requires_trust_remote_code = any(
_requires_trust_remote_code_for_model(_t, request.hf_token)
for _t in security_targets
)
requires_security_review = any(
_requires_security_review_for_model(_t, request.hf_token) for _t in security_targets
)
# Native context length, read from the local GGUF header when present.
# Lets the staged ("Load on selection" off) flow populate the context
# slider before the GPU load; None until the file is downloaded.
@ -2587,9 +2702,8 @@ async def validate_model(
is_gguf = is_gguf,
is_lora = getattr(config, "is_lora", False),
is_vision = getattr(config, "is_vision", False),
requires_trust_remote_code = bool(
load_inference_config(config.identifier).get("trust_remote_code", False)
),
requires_trust_remote_code = requires_trust_remote_code,
requires_security_review = requires_security_review,
context_length = context_length,
)
@ -2898,9 +3012,8 @@ async def get_status(current_subject: str = Depends(get_current_subject)):
loading = [],
loaded = [_display_model_id] if _display_model_id else [],
inference = _inference_cfg,
requires_trust_remote_code = bool(
(_inference_cfg or {}).get("trust_remote_code", False)
),
# GGUF status: auto_map never executes, so inert (matches validate_model).
requires_trust_remote_code = False,
supports_reasoning = llama_backend.supports_reasoning,
reasoning_style = llama_backend.reasoning_style,
reasoning_always_on = llama_backend.reasoning_always_on,
@ -2958,8 +3071,8 @@ async def get_status(current_subject: str = Depends(get_current_subject)):
loading = list(getattr(backend, "loading_models", set())),
loaded = list(backend.models.keys()),
inference = inference_config,
requires_trust_remote_code = bool(
(inference_config or {}).get("trust_remote_code", False)
requires_trust_remote_code = _resolve_loaded_trust_remote_code(
backend.active_model_name, model_info, inference_config
),
supports_reasoning = _sf_flags["supports_reasoning"],
reasoning_style = _sf_flags["reasoning_style"],

263
studio/backend/routes/models.py
View file

@ -1532,16 +1532,22 @@ async def get_model_config(
except Exception:
pass
# Fallback: try AutoConfig directly.
# Fallback: read raw config.json (declarative fields only) -- a selection-time
# metadata probe that must never execute a repo's auto_map Python.
if max_position_embeddings is None:
try:
from transformers import AutoConfig as _AutoConfig
from utils.transformers_version import _load_config_json
from types import SimpleNamespace
_trust = model_name.lower().startswith("unsloth/")
_ac = _AutoConfig.from_pretrained(
model_name, trust_remote_code = _trust, token = hf_token
)
max_position_embeddings = _get_max_position_embeddings(_ac)
_cfg = _load_config_json(model_name, hf_token = hf_token)
if _cfg is not None:
def _to_ns(d):
if isinstance(d, dict):
return SimpleNamespace(**{k: _to_ns(v) for k, v in d.items()})
return d
max_position_embeddings = _get_max_position_embeddings(_to_ns(_cfg))
except Exception:
pass
@ -1574,6 +1580,195 @@ async def get_model_config(
)
@router.post("/remote-code-scan")
async def scan_model_remote_code(
model_name: str = Body(..., embed = True),
hf_token: Optional[str] = Body(None, embed = True),
current_subject: str = Depends(get_current_subject),
):
"""Scan a model's ``auto_map`` custom code so the UI can show findings before
the user enables ``trust_remote_code``. Code-free: reads ``config.json`` and
statically scans the repo ``.py`` (never loads the model). Returns
``has_remote_code`` plus the severity-tagged findings + a pinning fingerprint.
POST (not GET) so the ``hf_token`` for gated repos travels in the body and
never lands in a URL, browser history, or access log.
"""
try:
from utils.security import preflight_remote_code_consent_for_targets
if not is_local_path(model_name):
model_name = resolve_cached_repo_id_case(model_name)
# Scan the adapter AND the base together (a LoRA runs both repos' code; a pickle
# can live in either), pinned by one combined fingerprint. Snapshot the primary's
# cache state BEFORE resolving the base: for a remote adapter that resolve
# downloads adapter_config.json, which would otherwise hide the adapter from
# cleanup on decline. On error treat as pre-existing so a decline never deletes it.
try:
_primary_preexisting = is_local_path(model_name) or _repo_in_any_hf_cache(model_name)
except Exception:
_primary_preexisting = True
security_targets = [model_name]
try:
from utils.models.model_config import get_base_model_from_lora_identifier
# Resolve a LOCAL or REMOTE adapter's base so its code/weights are scanned too.
_base = get_base_model_from_lora_identifier(model_name, hf_token)
if _base:
security_targets.append(_base)
except Exception:
pass
security_targets = list(dict.fromkeys(security_targets))
# Record every repo OUR scan is first to pull into the cache (adapter, base, and
# external auto_map repos like owner/name--module.Class), so a decline purges
# exactly what was downloaded. Computed BEFORE the preflight downloads, against
# every cache the discard searches, so a repo the user already had is not deleted.
from utils.security.remote_code_scan import external_auto_map_repos
scan_created_repos: list = []
_seen_created: set = set()
def _mark_scan_created(repo: str, *, preexisting: Optional[bool] = None) -> None:
if not repo or repo in _seen_created:
return
_seen_created.add(repo)
try:
already = (
preexisting
if preexisting is not None
else (is_local_path(repo) or _repo_in_any_hf_cache(repo))
)
if not already:
scan_created_repos.append(repo)
except Exception:
pass
for _target in security_targets:
# Use the pre-base-resolution snapshot for the primary (see above).
_mark_scan_created(
_target, preexisting = _primary_preexisting if _target == model_name else None
)
for _ext in external_auto_map_repos(_target, hf_token):
_mark_scan_created(_ext)
decision = preflight_remote_code_consent_for_targets(security_targets, hf_token = hf_token)
payload = decision.response_payload()
payload["requires_trust_remote_code"] = decision.has_remote_code
# created_by_scan = primary flag (older clients); scan_created_repos drives cleanup.
payload["created_by_scan"] = model_name in scan_created_repos
payload["scan_created_repos"] = scan_created_repos
# Malware gate (metadata-only): surface HF-flagged unsafe files so the dialog can
# hard-block. Orthogonal to remote code -- a poisoned pickle needs no auto_map.
from utils.security import evaluate_file_security, security_load_subdirs
unsafe_files: list = []
security_blocked = False
for _target in security_targets:
_sec = evaluate_file_security(
_target, hf_token = hf_token, load_subdirs = security_load_subdirs(_target, hf_token)
)
security_blocked = security_blocked or _sec.blocked
unsafe_files.extend(_sec.unsafe_files)
payload["unsafe_files"] = unsafe_files
payload["security_blocked"] = security_blocked
if security_blocked:
# Non-approvable hard block: approvable False hides "Enable and continue", and
# requires_trust_remote_code forces the dialog open even with no custom code.
payload["approvable"] = False
payload["requires_trust_remote_code"] = True
payload["error_kind"] = "malware_blocked"
return payload
except Exception as e:
raise log_and_http_error(
e,
500,
"Failed to scan model remote code",
event = "models.remote_code_scan_failed",
log = logger,
)
@router.post("/discard-remote-code")
async def discard_remote_code_download(
model_name: str = Body(..., embed = True), current_subject: str = Depends(get_current_subject)
):
"""Purge a repo the consent scan downloaded after the user DECLINED its custom
code, so untrusted code is not left on disk.
Safety: only ever deletes a metadata-only cache entry the scan created. It
refuses a local path (never touches user files), a currently-loaded model, and
any repo that has weight files cached (``*.safetensors`` / ``*.bin`` /
``*.gguf``) -- i.e. a model the user actually downloaded. The frontend only
calls this when the scan reported ``created_by_scan``.
"""
if is_local_path(model_name):
return {"deleted": False, "reason": "local"}
if not _is_valid_repo_id(model_name):
return {"deleted": False, "reason": "invalid"}
# Never delete a model that is loaded for inference.
try:
from routes.inference import get_llama_cpp_backend
llama_backend = get_llama_cpp_backend()
if llama_backend.is_loaded and llama_backend.model_identifier:
loaded = llama_backend.model_identifier.lower()
if loaded == model_name.lower() or loaded.startswith(model_name.lower()):
return {"deleted": False, "reason": "loaded"}
except Exception:
pass
try:
inference_backend = get_inference_backend()
if inference_backend.active_model_name:
active = inference_backend.active_model_name.lower()
if active == model_name.lower() or active.startswith(model_name.lower()):
return {"deleted": False, "reason": "loaded"}
except Exception:
pass
_WEIGHTS = (
".safetensors",
".bin",
".pt",
".pth",
".h5",
".msgpack",
".gguf",
".onnx",
".ckpt",
)
try:
target_repo = None
hf_cache = None
for cache in _all_hf_cache_scans():
for repo_info in cache.repos:
if repo_info.repo_type != "model":
continue
if repo_info.repo_id.lower() == model_name.lower():
target_repo, hf_cache = repo_info, cache
break
if target_repo is not None:
break
if target_repo is None:
return {"deleted": False, "reason": "not_cached"}
# Hard guard: a repo with weights is a real model the user has -- leave it.
for rev in target_repo.revisions:
for f in rev.files:
if f.file_name.lower().endswith(_WEIGHTS):
return {"deleted": False, "reason": "has_weights"}
revision_hashes = [rev.commit_hash for rev in target_repo.revisions]
if not revision_hashes:
return {"deleted": False, "reason": "not_cached"}
hf_cache.delete_revisions(*revision_hashes).execute()
logger.info("Discarded declined remote-code download: %s", model_name)
return {"deleted": True}
except Exception as e:
logger.warning("Could not discard remote-code download for %s: %s", model_name, e)
return {"deleted": False, "reason": "error"}
@router.get("/loras")
async def scan_loras(
outputs_dir: str = Query(
@ -1990,7 +2185,11 @@ async def get_lora_base_model(lora_path: str, current_subject: str = Depends(get
@router.get("/check-vision/{model_name:path}", response_model = VisionCheckResponse)
async def check_vision_model(model_name: str, current_subject: str = Depends(get_current_subject)):
async def check_vision_model(
model_name: str,
hf_token: Optional[str] = Query(None),
current_subject: str = Depends(get_current_subject),
):
"""
Check if a model is a vision model.
@ -1998,7 +2197,8 @@ async def check_vision_model(model_name: str, current_subject: str = Depends(get
"""
try:
logger.info(f"Checking if vision model: {model_name}")
is_vision = is_vision_model(model_name)
# Authenticate so a gated/private VLM classifies correctly (else 404 -> non-vision).
is_vision = is_vision_model(model_name, hf_token = hf_token)
logger.info(f"Vision check result for {model_name}: is_vision={is_vision}")
return VisionCheckResponse(
@ -2348,6 +2548,51 @@ def _get_repo_size_cached(repo_id: str) -> int:
return 0
def _repo_in_any_hf_cache(model_name: str) -> bool:
"""Whether ``model_name`` already exists in ANY HF cache the discard searches
(active, legacy, default).
``created_by_scan`` must be True only when the scan itself first pulled the repo;
checking just the active cache (``get_cache_path``) would mark a repo the user
already had in a legacy/default cache as scan-created, so declining the consent
would delete a model they did not download via the scan. Mirrors the cache set in
``_all_hf_cache_scans`` but only probes for the one repo dir (cheap, no full scan).
"""
from utils.paths import (
hf_default_cache_dir,
legacy_hf_cache_dir,
resolve_cached_repo_id_case,
)
dirname = f"models--{resolve_cached_repo_id_case(model_name).replace('/', '--')}"
dirname_lower = dirname.lower()
candidates = []
try:
from huggingface_hub.constants import HF_HUB_CACHE
candidates.append(Path(HF_HUB_CACHE))
except Exception:
pass
for fn in (legacy_hf_cache_dir, hf_default_cache_dir):
try:
candidates.append(fn())
except Exception:
continue
# resolve_cached_repo_id_case only normalizes the ACTIVE cache, but discard deletes
# case-insensitively across all caches, so detect case-insensitively too -- else a
# pre-existing case-variant repo is misreported as scan-created and deleted on decline.
for cache in candidates:
try:
if (cache / dirname).exists():
return True
if cache.is_dir():
for entry in cache.iterdir():
if entry.name.lower() == dirname_lower and entry.is_dir():
return True
except Exception:
continue
return False
def _all_hf_cache_scans():
"""scan_cache_dir for the active, legacy, and default HF caches.

18
studio/backend/routes/training.py
View file

@ -248,18 +248,30 @@ async def start_training(
"output_dir": resume_output_dir,
"resume_from_checkpoint": request.resume_from_checkpoint,
"trust_remote_code": request.trust_remote_code,
"approved_remote_code_fingerprint": request.approved_remote_code_fingerprint,
"gpu_ids": request.gpu_ids,
"s3_config": request.s3_config.model_dump() if request.s3_config else None,
}
# Training page has no trust_remote_code toggle; as a safety net consult
# YAML model defaults directly so models that need it always get it.
# Training page has no trust_remote_code toggle, so honor the YAML default
# -- but only for genuine first-party (unsloth/nvidia) Hub repos, never a
# local path or a name merely starting with "unsloth/".
if not training_kwargs["trust_remote_code"]:
from utils.security.trusted_org import is_trusted_org_repo
model_defaults = load_model_defaults(request.model_name)
yaml_trust = model_defaults.get("training", {}).get("trust_remote_code", False)
if yaml_trust:
if yaml_trust and is_trusted_org_repo(
request.model_name, hf_token = request.hf_token or None
):
logger.info(f"YAML config sets trust_remote_code=True for {request.model_name}")
training_kwargs["trust_remote_code"] = True
elif yaml_trust:
logger.warning(
"YAML sets trust_remote_code=True for %s but it is not a trusted "
"first-party repo; leaving disabled (user can opt in explicitly).",
request.model_name,
)
# Free GPU memory: shut down any running inference/export subprocesses
# before training (they'd compete for VRAM otherwise).

361
studio/backend/tests/test_capability_detection.py Normal file
View file

@ -0,0 +1,361 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Component A tests: capability detection must never execute model repo code.
Covers: load_model_config defaults trust_remote_code False; the _VISION_CHECK_SCRIPT
subprocess literal keeps remote code off; registry-backed vision/audio detection from
raw config.json (repo-code VLMs detected without execution; ForConditionalGeneration
false positives fixed); and the model-details / GPU probes never enable remote code.
"""
import json
from types import SimpleNamespace
from unittest.mock import patch, MagicMock
import pytest
from utils.models.model_config import (
load_model_config,
is_vision_model,
_is_vlm,
_raw_config_has_vision_config,
_vision_detection_cache,
_VISION_CHECK_SCRIPT,
_VLM_MODEL_TYPES,
_AUDIO_ONLY_MODEL_TYPES,
_VLM_CLASS_NAMES,
)
@pytest.fixture(autouse = True)
def _clear_vision_cache():
_vision_detection_cache.clear()
yield
_vision_detection_cache.clear()
def _write_model_dir(
tmp_path,
cfg,
with_evil_module = False,
):
"""Write a local model dir, optionally with an auto_map module that writes a sentinel
on import so accidental code execution during detection shows up on disk."""
(tmp_path / "config.json").write_text(json.dumps(cfg))
if with_evil_module:
sentinel = tmp_path / "PWNED_SENTINEL"
(tmp_path / "modeling_evil.py").write_text(
"import os\n"
f"open({str(sentinel)!r}, 'w').write('pwned')\n"
"class EvilConfig: pass\n"
"class EvilModel: pass\n"
)
return str(tmp_path)
# load_model_config default
class TestLoadModelConfigDefault:
@patch("transformers.AutoConfig.from_pretrained")
def test_default_off_with_token(self, fp):
load_model_config("org/m", token = "hf_x")
assert fp.call_args.kwargs["trust_remote_code"] is False
@patch("utils.models.model_config.without_hf_auth")
@patch("transformers.AutoConfig.from_pretrained")
def test_default_off_public(self, fp, no_auth):
from contextlib import nullcontext
no_auth.return_value = nullcontext()
load_model_config("org/m", use_auth = False)
assert fp.call_args.kwargs["trust_remote_code"] is False
@patch("transformers.AutoConfig.from_pretrained")
def test_default_off_cached_auth(self, fp):
load_model_config("org/m", use_auth = True)
assert fp.call_args.kwargs["trust_remote_code"] is False
@patch("transformers.AutoConfig.from_pretrained")
def test_explicit_true_forwarded(self, fp):
load_model_config("org/m", token = "t", trust_remote_code = True)
assert fp.call_args.kwargs["trust_remote_code"] is True
# subprocess script literal
def test_vision_check_script_disables_remote_code():
assert '"trust_remote_code": False' in _VISION_CHECK_SCRIPT
assert '"trust_remote_code": True' not in _VISION_CHECK_SCRIPT
# _is_vlm matrix (pure function, registry-backed)
def _cfg(**kw):
return SimpleNamespace(**kw)
class TestIsVlm:
def test_deepseek_ocr_vision_via_vision_config(self):
# auto_map repo-code model; vision-ness is declarative.
c = _cfg(
model_type = "deepseek_vl_v2",
architectures = ["DeepseekOCRForCausalLM"],
vision_config = {},
projector_config = {},
)
assert _is_vlm(c) is True
def test_kimi_vision_via_vision_config(self):
c = _cfg(
model_type = "kimi_k25",
architectures = ["KimiK25ForConditionalGeneration"],
vision_config = {},
)
assert _is_vlm(c) is True
def test_glm_flash_text_is_not_vision(self):
c = _cfg(model_type = "glm4_moe_lite", architectures = ["Glm4MoeLiteForCausalLM"])
assert _is_vlm(c) is False
def test_gemma4_vision_via_vision_config(self):
c = _cfg(
model_type = "gemma4_unified",
architectures = ["Gemma4UnifiedForConditionalGeneration"],
vision_config = {},
image_token_id = 1,
)
assert _is_vlm(c) is True
def test_t5_not_misclassified_as_vision(self):
# Regression: ForConditionalGeneration must NOT be a vision signal.
c = _cfg(model_type = "t5", architectures = ["T5ForConditionalGeneration"])
assert _is_vlm(c) is False
def test_bart_not_misclassified_as_vision(self):
c = _cfg(model_type = "bart", architectures = ["BartForConditionalGeneration"])
assert _is_vlm(c) is False
def test_whisper_audio_not_vision(self):
c = _cfg(model_type = "whisper", architectures = ["WhisperForConditionalGeneration"])
assert _is_vlm(c) is False
def test_csm_audio_not_vision(self):
c = _cfg(model_type = "csm", architectures = ["CsmForConditionalGeneration"])
assert _is_vlm(c) is False
def test_native_vlm_via_registry_model_type(self):
# llava is in the transformers vision registry.
assert "llava" in _VLM_MODEL_TYPES
c = _cfg(model_type = "llava", architectures = ["LlavaForConditionalGeneration"])
assert _is_vlm(c) is True
def test_native_vlm_via_registry_class_name(self):
# Class-name match works even if model_type were unknown.
cls = next(iter(_VLM_CLASS_NAMES))
c = _cfg(model_type = "something_unlisted", architectures = [cls])
assert _is_vlm(c) is True
def test_omni_audio_plus_vision_is_vision(self):
# An audio-registry model_type with an explicit vision sub-config is still vision.
audio_mt = next(iter(_AUDIO_ONLY_MODEL_TYPES - _VLM_MODEL_TYPES))
c = _cfg(model_type = audio_mt, architectures = ["X"], vision_config = {})
assert _is_vlm(c) is True
# _raw_config_has_vision_config (code-free reader, mocked HF download)
def _mock_raw_config(tmp_path, payload):
p = tmp_path / "config.json"
p.write_text(json.dumps(payload))
return p
class TestRawConfigVisionReader:
@pytest.mark.parametrize(
"payload,expected",
[
(
{
"model_type": "deepseek_vl_v2",
"architectures": ["DeepseekOCRForCausalLM"],
"auto_map": {"AutoConfig": "modeling_deepseekocr.DeepseekOCRConfig"},
"vision_config": {},
"projector_config": {},
},
True,
),
(
{
"model_type": "kimi_k25",
"architectures": ["KimiK25ForConditionalGeneration"],
"auto_map": {"AutoConfig": "configuration_kimi_k25.KimiK25Config"},
"vision_config": {},
},
True,
),
({"model_type": "glm4_moe_lite", "architectures": ["Glm4MoeLiteForCausalLM"]}, False),
(
{
"model_type": "gemma4_unified",
"architectures": ["Gemma4UnifiedForConditionalGeneration"],
"vision_config": {},
},
True,
),
({"model_type": "t5", "architectures": ["T5ForConditionalGeneration"]}, False),
(
{"model_type": "whisper", "architectures": ["WhisperForConditionalGeneration"]},
False,
),
({"model_type": "csm", "architectures": ["CsmForConditionalGeneration"]}, False),
],
)
def test_reader(self, tmp_path, payload, expected):
cfg_path = _mock_raw_config(tmp_path, payload)
with (
patch("utils.models.model_config.is_local_path", return_value = False),
patch("huggingface_hub.hf_hub_download", return_value = str(cfg_path)),
):
assert _raw_config_has_vision_config("org/model") is expected
def test_reader_never_executes_remote_code(self, tmp_path):
# Even with auto_map present, the reader only parses JSON: no AutoConfig touched.
cfg_path = _mock_raw_config(
tmp_path,
{
"model_type": "deepseek_vl_v2",
"architectures": ["DeepseekOCRForCausalLM"],
"auto_map": {"AutoConfig": "modeling_deepseekocr.DeepseekOCRConfig"},
"vision_config": {},
},
)
with (
patch("utils.models.model_config.is_local_path", return_value = False),
patch("huggingface_hub.hf_hub_download", return_value = str(cfg_path)),
patch(
"transformers.AutoConfig.from_pretrained",
side_effect = AssertionError("AutoConfig must not be called"),
),
):
assert _raw_config_has_vision_config("org/deepseek-ocr") is True
# Probes: model-details + GPU estimate never execute remote code
def test_gpu_estimate_probe_is_code_free():
from utils.hardware import hardware
cfg = {
"model_type": "glm4_moe_lite",
"hidden_size": 4096,
"num_hidden_layers": 40,
"max_position_embeddings": 8192,
}
with (
patch("utils.transformers_version._load_config_json", return_value = cfg),
patch(
"transformers.AutoConfig.from_pretrained",
side_effect = AssertionError("AutoConfig must not be called"),
),
):
out = hardware._load_config_for_gpu_estimate("unsloth/GLM-4.7-Flash")
assert out.max_position_embeddings == 8192
assert out.hidden_size == 4096
def test_models_route_source_has_no_remote_code_probe():
# The metadata probe must never build a trust_remote_code=True loader; referencing
# the static consent scanner or the requires_trust_remote_code flag is fine.
import inspect
import routes.models as models_route
src = inspect.getsource(models_route)
assert "trust_remote_code = True" not in src
assert "trust_remote_code=True" not in src
# Adversarial end-to-end: is_vision_model + the two metadata probes never run auto_map.
def test_no_code_execution_on_detection(tmp_path):
# A malicious local auto_map -> modeling_evil must not execute through any probe.
cfg = {
"model_type": "deepseek_vl_v2",
"architectures": ["DeepseekOCRForCausalLM"],
"auto_map": {
"AutoConfig": "modeling_evil.EvilConfig",
"AutoModel": "modeling_evil.EvilModel",
},
"vision_config": {"image_size": 1024},
"max_position_embeddings": 4096,
}
path = _write_model_dir(tmp_path, cfg, with_evil_module = True)
sentinel = tmp_path / "PWNED_SENTINEL"
from utils.hardware.hardware import _load_config_for_gpu_estimate
from utils.transformers_version import _load_config_json
result = is_vision_model(path)
ns = _load_config_for_gpu_estimate(path)
raw = _load_config_json(path)
assert not sentinel.exists(), "SECURITY FAILURE: auto_map code executed during detection"
assert result is True # detected as vision via raw vision_config, no exec
assert ns is not None and getattr(ns, "max_position_embeddings", None) == 4096
assert raw is not None and raw.get("model_type") == "deepseek_vl_v2"
@pytest.mark.parametrize(
"cfg, expected",
[
# repo-code VLMs (auto_map) detected via declarative vision_config
(
{
"model_type": "deepseek_vl_v2",
"architectures": ["DeepseekOCRForCausalLM"],
"auto_map": {"AutoConfig": "x.Y"},
"vision_config": {},
},
True,
),
(
{
"model_type": "kimi_k25",
"architectures": ["KimiK25ForConditionalGeneration"],
"auto_map": {"AutoConfig": "x.Y"},
"vision_config": {},
},
True,
),
# newer-native vision via vision_config
(
{
"model_type": "gemma4_unified",
"architectures": ["Gemma4UnifiedForConditionalGeneration"],
"vision_config": {},
"image_token_id": 7,
},
True,
),
# text / seq2seq / audio that share the ForConditionalGeneration suffix
({"model_type": "glm4_moe_lite", "architectures": ["Glm4MoeLiteForCausalLM"]}, False),
({"model_type": "t5", "architectures": ["T5ForConditionalGeneration"]}, False),
({"model_type": "bart", "architectures": ["BartForConditionalGeneration"]}, False),
({"model_type": "whisper", "architectures": ["WhisperForConditionalGeneration"]}, False),
({"model_type": "csm", "architectures": ["CsmForConditionalGeneration"]}, False),
# registry-native VLMs via model_type
({"model_type": "qwen2_vl", "architectures": ["Qwen2VLForConditionalGeneration"]}, True),
({"model_type": "llava", "architectures": ["LlavaForConditionalGeneration"]}, True),
],
)
def test_is_vision_model_end_to_end(tmp_path, cfg, expected):
path = _write_model_dir(tmp_path, cfg)
assert is_vision_model(path) is expected, f"{cfg['model_type']} expected vision={expected}"
def test_registry_derivation():
# Registry-derived sets are large and include the curated repo-code VLMs.
assert len(_VLM_MODEL_TYPES) >= 50, f"_VLM_MODEL_TYPES too small: {len(_VLM_MODEL_TYPES)}"
assert (
len(_AUDIO_ONLY_MODEL_TYPES) >= 20
), f"_AUDIO_ONLY too small: {len(_AUDIO_ONLY_MODEL_TYPES)}"
for repo_vlm in ("deepseek_vl_v2", "kimi_k25", "phi3_v", "cogvlm2", "minicpmv"):
assert repo_vlm in _VLM_MODEL_TYPES, f"curated repo-code VLM {repo_vlm} missing"
for native in ("llava", "qwen2_vl"):
assert native in _VLM_MODEL_TYPES, f"registry-native VLM {native} missing"
for audio in ("whisper", "csm"):
assert audio in _AUDIO_ONLY_MODEL_TYPES, f"audio type {audio} missing"

1414
studio/backend/tests/test_consent_gate.py Normal file
View file

File diff suppressed because it is too large Load diff

534
studio/backend/tests/test_file_security.py Normal file
View file

@ -0,0 +1,534 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tests for the malware / unsafe-file gate (utils.security.file_security).
The gate reads HF's security scan (model_info securityStatus) metadata-only and never
downloads flagged files; only the Hub call is stubbed. Policy: block a non-"safe" level
(unknown levels fail closed), fail open when the scan is unavailable, skip local paths
only, no first-party exemption. The block is scoped to the load-path RCE vector (a
root-level code-executing file), so flagged safetensors and subdir pickles do not block.
"""
from types import SimpleNamespace
from unittest.mock import patch
import pytest
from utils.security import evaluate_file_security
def _patch_status(status):
"""Patch huggingface_hub.model_info to return one fixed security_repo_status."""
def _mi(*_args, **_kwargs):
return SimpleNamespace(security_repo_status = status)
return patch("huggingface_hub.model_info", side_effect = _mi)
def _patch_raises(exc = RuntimeError("offline")):
return patch("huggingface_hub.model_info", side_effect = exc)
def _patch_no_index():
"""Make the weight-index lookup find no index files (definitive: nothing sharded)."""
from huggingface_hub.utils import EntryNotFoundError
def _dl(
repo_id = None,
filename = None,
token = None,
**kw,
):
raise EntryNotFoundError(filename or "")
return patch("huggingface_hub.hf_hub_download", side_effect = _dl)
def _patch_index(weight_map, index_filename = "pytorch_model.bin.index.json"):
"""Serve a root weight index mapping tensor names -> shard paths; others 404."""
import json
import tempfile
from pathlib import Path
from huggingface_hub.utils import EntryNotFoundError
def _dl(
repo_id = None,
filename = None,
token = None,
**kw,
):
if filename == index_filename:
p = Path(tempfile.mkdtemp()) / filename
p.write_text(json.dumps({"weight_map": weight_map}))
return str(p)
raise EntryNotFoundError(filename or "")
return patch("huggingface_hub.hf_hub_download", side_effect = _dl)
def _patch_index_unreadable():
"""Make every index fetch fail transiently (inconclusive lookup -> fail closed)."""
def _dl(
repo_id = None,
filename = None,
token = None,
**kw,
):
raise RuntimeError("transient network error")
return patch("huggingface_hub.hf_hub_download", side_effect = _dl)
def _patch_index_mixed(weight_map, readable_index, failing_index):
"""Serve one index cleanly while another fails transiently: the flagged shard is
listed only by the index we could not read, so a naive "read any index?" check breaks."""
import json
import tempfile
from pathlib import Path
from huggingface_hub.utils import EntryNotFoundError
def _dl(
repo_id = None,
filename = None,
token = None,
**kw,
):
if filename == readable_index:
p = Path(tempfile.mkdtemp()) / filename
p.write_text(json.dumps({"weight_map": weight_map}))
return str(p)
if filename == failing_index:
raise RuntimeError("transient network error")
raise EntryNotFoundError(filename or "")
return patch("huggingface_hub.hf_hub_download", side_effect = _dl)
@pytest.mark.parametrize("level", ["unsafe", "suspicious", "malicious"])
def test_blocks_each_blocking_level(level):
status = {"scansDone": True, "filesWithIssues": [{"path": "pytorch_model.bin", "level": level}]}
with _patch_status(status):
d = evaluate_file_security("evil/repo")
assert d.blocked is True
assert d.unsafe_files == [{"path": "pytorch_model.bin", "level": level}]
assert d.response_payload()["security_blocked"] is True
def test_ignores_safe_only():
status = {
"scansDone": True,
"filesWithIssues": [{"path": "model.safetensors", "level": "safe"}],
}
with _patch_status(status):
d = evaluate_file_security("good/repo")
assert d.blocked is False
assert d.unsafe_files == []
def test_blocks_unsafe_even_when_scans_not_done():
# scansDone is often False for clean repos; an already-flagged file must still block.
status = {"scansDone": False, "filesWithIssues": [{"path": "x.pkl", "level": "unsafe"}]}
with _patch_status(status):
d = evaluate_file_security("evil/repo")
assert d.blocked is True
def test_fail_open_when_scan_unavailable():
# model_info returns no security_repo_status -> unknown -> allow.
with _patch_status(None):
d = evaluate_file_security("unknown/repo")
assert d.blocked is False
def test_fail_open_on_exception_offline():
with _patch_raises():
d = evaluate_file_security("offline/repo")
assert d.blocked is False
def test_fail_open_scans_done_no_issues():
with _patch_status({"scansDone": True, "filesWithIssues": []}):
d = evaluate_file_security("clean/repo")
assert d.blocked is False
def test_skips_local_path():
# A local path has no Hub scan; must not even call model_info.
with patch("huggingface_hub.model_info", side_effect = AssertionError("should not be called")):
d = evaluate_file_security("/tmp/some/local/model")
assert d.blocked is False
assert "local" in d.reason
def test_remote_gguf_named_repo_is_still_scanned():
# Only LOCAL paths skip the Hub scan, so a remote .gguf repo is still scanned and a
# poisoned pickle smuggled into it is blocked.
status = {
"scansDone": True,
"filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}],
}
with _patch_status(status):
d = evaluate_file_security("evil/model.gguf")
assert d.blocked is True
assert d.unsafe_files
def test_skips_local_gguf_file():
# A local .gguf path is caught by is_local_path -- no Hub call.
with patch("huggingface_hub.model_info", side_effect = AssertionError("should not be called")):
d = evaluate_file_security("/tmp/models/model.gguf")
assert d.blocked is False
assert "local" in d.reason
def test_no_first_party_exemption():
# A poisoned pickle in a first-party repo still blocks (compromised-repo defense).
status = {
"scansDone": True,
"filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}],
}
with _patch_status(status):
d = evaluate_file_security("unsloth/some-model")
assert d.blocked is True
def test_malformed_entries_are_ignored():
status = {
"scansDone": True,
"filesWithIssues": ["not-a-dict", {"path": "ok.pkl", "level": "unsafe"}],
}
with _patch_status(status):
d = evaluate_file_security("evil/repo")
assert d.blocked is True
assert d.unsafe_files == [{"path": "ok.pkl", "level": "unsafe"}]
def test_response_payload_shape():
status = {"scansDone": True, "filesWithIssues": [{"path": "a.pkl", "level": "malicious"}]}
with _patch_status(status):
payload = evaluate_file_security("evil/repo").response_payload()
assert set(payload) == {"unsafe_files", "security_blocked", "reason"}
assert payload["security_blocked"] is True
assert payload["unsafe_files"] == [{"path": "a.pkl", "level": "malicious"}]
# ── Load-path RCE scoping: block only files a load would actually deserialize ──
def test_flagged_safetensors_does_not_block():
# safetensors is tensor-only and cannot execute code, so a flag on one (often
# picklescan tripping on a sibling pickle) is not an RCE vector and must not block.
status = {
"scansDone": False,
"filesWithIssues": [{"path": "model-00001-of-00004.safetensors", "level": "unsafe"}],
}
with _patch_status(status):
d = evaluate_file_security("nvidia/some-model")
assert d.blocked is False
assert d.unsafe_files == []
def test_flagged_subdirectory_pickle_does_not_block():
# from_pretrained reads only root weights; a flagged subdir pickle no root index
# references (e.g. a NeMo checkpoint) is never loaded, so it must not block.
status = {
"scansDone": False,
"filesWithIssues": [
{"path": "nemo/weights/common.pt", "level": "unsafe"},
{"path": "nemo/weights/__0_0.distcp", "level": "unsafe"},
],
}
with _patch_status(status), _patch_no_index():
d = evaluate_file_security("nvidia/some-model")
assert d.blocked is False
assert d.unsafe_files == []
def test_nemotron_h_shaped_status_loads():
# Real Nemotron-H-8B-Base-8K shape: flagged root safetensors + unreferenced nemo/
# pickles. None is a load-path vector, so it must load.
status = {
"scansDone": False,
"filesWithIssues": [
{"path": "nemo/weights/.metadata", "level": "unsafe"},
{"path": "nemo/weights/__0_0.distcp", "level": "unsafe"},
{"path": "nemo/weights/common.pt", "level": "unsafe"},
{"path": "model-00001-of-00004.safetensors", "level": "unsafe"},
{"path": "model-00002-of-00004.safetensors", "level": "unsafe"},
],
}
with _patch_status(status), _patch_no_index():
d = evaluate_file_security("nvidia/Nemotron-H-8B-Base-8K")
assert d.blocked is False
assert d.unsafe_files == []
def test_indexed_subdir_shard_blocks():
# A flagged subdir shard that a root index references IS deserialized, so it blocks.
status = {
"scansDone": False,
"filesWithIssues": [
{"path": "shards/pytorch_model-00001-of-00002.bin", "level": "unsafe"},
],
}
weight_map = {
"layer.0.weight": "shards/pytorch_model-00001-of-00002.bin",
"layer.1.weight": "shards/pytorch_model-00002-of-00002.bin",
}
with _patch_status(status), _patch_index(weight_map):
d = evaluate_file_security("evil/sharded")
assert d.blocked is True
assert d.unsafe_files == [
{"path": "shards/pytorch_model-00001-of-00002.bin", "level": "unsafe"}
]
def test_unindexed_subdir_pickle_does_not_block_when_index_present():
# An index exists but does not list the flagged subdir pickle -> not loaded -> no block.
status = {
"scansDone": False,
"filesWithIssues": [{"path": "extras/notes.bin", "level": "unsafe"}],
}
weight_map = {"layer.0.weight": "pytorch_model-00001-of-00001.bin"}
with _patch_status(status), _patch_index(weight_map):
d = evaluate_file_security("org/has-index")
assert d.blocked is False
assert d.unsafe_files == []
def test_inconclusive_index_lookup_blocks_subdir_pickle():
# An unreadable index can't rule out that the flagged subdir pickle is a shard -> block.
status = {
"scansDone": False,
"filesWithIssues": [{"path": "weights/model_part.bin", "level": "unsafe"}],
}
with _patch_status(status), _patch_index_unreadable():
d = evaluate_file_security("org/transient")
assert d.blocked is True
assert d.unsafe_files == [{"path": "weights/model_part.bin", "level": "unsafe"}]
def test_partial_index_read_with_transient_failure_blocks_subdir_pickle():
# The bin index (which would list the flagged shard) fails transiently; a partial
# path set is not definitive, so fail closed.
status = {
"scansDone": False,
"filesWithIssues": [
{"path": "shards/pytorch_model-00001-of-00002.bin", "level": "unsafe"},
],
}
# The readable index lists only benign shards; the flagged .bin is in the unread index.
safetensors_map = {"layer.0.weight": "model-00001-of-00001.safetensors"}
with (
_patch_status(status),
_patch_index_mixed(
safetensors_map,
readable_index = "model.safetensors.index.json",
failing_index = "pytorch_model.bin.index.json",
),
):
d = evaluate_file_security("evil/mixed-index")
assert d.blocked is True
assert d.unsafe_files == [
{"path": "shards/pytorch_model-00001-of-00002.bin", "level": "unsafe"}
]
def test_root_pickle_alongside_safetensors_still_blocks():
# A real root pickle blocks even alongside a flagged safetensors; it is a load-path vector.
status = {
"scansDone": False,
"filesWithIssues": [
{"path": "model.safetensors", "level": "unsafe"},
{"path": "pytorch_model.bin", "level": "unsafe"},
],
}
with _patch_status(status):
d = evaluate_file_security("evil/repo")
assert d.blocked is True
assert d.unsafe_files == [{"path": "pytorch_model.bin", "level": "unsafe"}]
def test_eicar_shaped_root_files_block():
# The canonical eicar repo ships its dangerous files at the ROOT, so it stays blocked.
status = {
"scansDone": True,
"filesWithIssues": [
{"path": "model_broken_X.pkl", "level": "unsafe"},
{"path": "danger.dat", "level": "unsafe"},
{"path": "eicar_test_file", "level": "unsafe"},
],
}
with _patch_status(status):
d = evaluate_file_security("mcpotato/42-eicar-street")
assert d.blocked is True
assert len(d.unsafe_files) == 3
def test_unknown_future_level_fails_closed():
# Hub schema drift: an unrecognized non-"safe" level (e.g. "infected") on a root pickle must block.
status = {"scansDone": True, "filesWithIssues": [{"path": "weights.bin", "level": "infected"}]}
with _patch_status(status):
d = evaluate_file_security("evil/repo")
assert d.blocked is True
assert d.unsafe_files == [{"path": "weights.bin", "level": "infected"}]
def test_pending_or_scanning_level_does_not_block():
# A not-yet-finished per-file scan state must not false-block.
for lvl in ("pending", "scanning", "queued", "unscanned", "error"):
status = {
"scansDone": False,
"filesWithIssues": [{"path": "pytorch_model.bin", "level": lvl}],
}
with _patch_status(status):
d = evaluate_file_security("some/repo")
assert d.blocked is False, lvl
# -- Subdir load roots: Spark-TTS / BiCodec load from_pretrained(<snapshot>/LLM) --
def test_flagged_pickle_under_load_subdir_blocks():
# A flagged pickle directly under a declared load subdir is a root-level artifact there.
status = {
"scansDone": False,
"filesWithIssues": [{"path": "LLM/pytorch_model.bin", "level": "unsafe"}],
}
with _patch_status(status), _patch_no_index():
d = evaluate_file_security("org/spark-tts", load_subdirs = ("LLM",))
assert d.blocked is True
assert d.unsafe_files == [{"path": "LLM/pytorch_model.bin", "level": "unsafe"}]
def test_flagged_pickle_under_subdir_without_load_root_does_not_block():
# Same file, but NOT declared a load root and not indexed -> not deserialized.
status = {
"scansDone": False,
"filesWithIssues": [{"path": "LLM/pytorch_model.bin", "level": "unsafe"}],
}
with _patch_status(status), _patch_no_index():
d = evaluate_file_security("org/not-a-load-root")
assert d.blocked is False
assert d.unsafe_files == []
def test_indexed_shard_under_load_subdir_blocks():
# An index inside the load subdir referencing a flagged shard makes it a vector.
status = {
"scansDone": False,
"filesWithIssues": [
{"path": "LLM/shards/pytorch_model-00001-of-00002.bin", "level": "unsafe"}
],
}
weight_map = {
"layer.0.weight": "shards/pytorch_model-00001-of-00002.bin",
"layer.1.weight": "shards/pytorch_model-00002-of-00002.bin",
}
with (
_patch_status(status),
_patch_index(weight_map, index_filename = "LLM/pytorch_model.bin.index.json"),
):
d = evaluate_file_security("org/spark-tts", load_subdirs = ("LLM",))
assert d.blocked is True
# -- Source files are the consent gate's domain, not a deserialization vector --
def test_flagged_root_python_helper_does_not_block():
# A root .py is never deserialized; repo code runs only via auto_map under the consent
# gate, so flagging it here would false-block.
status = {
"scansDone": True,
"filesWithIssues": [
{"path": "build_pickles.py", "level": "unsafe"},
{"path": "train.py", "level": "suspicious"},
],
}
with _patch_status(status):
d = evaluate_file_security("org/has-helper-scripts")
assert d.blocked is False
assert d.unsafe_files == []
def test_root_pickle_still_blocks_with_flagged_python_sibling():
# The .py exemption must not mask a genuine root pickle in the same repo.
status = {
"scansDone": True,
"filesWithIssues": [
{"path": "convert.py", "level": "unsafe"},
{"path": "pytorch_model.bin", "level": "unsafe"},
],
}
with _patch_status(status):
d = evaluate_file_security("evil/mixed")
assert d.blocked is True
# -- Alias resolution: scan the repo the loader actually fetches from --
def _patch_status_capture(status):
"""Like _patch_status, but records the repo id model_info was queried with."""
seen = {}
def _mi(repo, *_a, **_k):
seen["repo"] = repo
return SimpleNamespace(security_repo_status = status)
return patch("huggingface_hub.model_info", side_effect = _mi), seen
def test_spark_tts_llm_alias_scans_real_repo():
# "Spark-TTS-0.5B/LLM" loads as unsloth/Spark-TTS-0.5B with LLM as load root; scanning
# the literal alias 404s and fails open, missing a flagged LLM/ pickle.
status = {"filesWithIssues": [{"path": "LLM/pytorch_model.bin", "level": "unsafe"}]}
cap, seen = _patch_status_capture(status)
with cap, patch("utils.paths.is_local_path", return_value = False), _patch_no_index():
d = evaluate_file_security("Spark-TTS-0.5B/LLM", load_subdirs = ())
assert seen["repo"] == "unsloth/Spark-TTS-0.5B" # scanned the real repo, not the alias
assert d.model_name == "unsloth/Spark-TTS-0.5B"
assert d.blocked is True
assert d.unsafe_files == [{"path": "LLM/pytorch_model.bin", "level": "unsafe"}]
def test_non_llm_alias_is_not_rewritten():
# A normal repo id with one slash must be scanned as-is (no spurious rewrite).
status = {"filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}]}
cap, seen = _patch_status_capture(status)
with cap, patch("utils.paths.is_local_path", return_value = False):
d = evaluate_file_security("org/model")
assert seen["repo"] == "org/model"
assert d.model_name == "org/model"
def test_generic_slash_llm_repo_is_scanned_as_itself():
# A third-party repo merely ending in "/LLM" is not a bicodec alias, so it must be
# scanned as itself; rewriting to unsloth/<parent> would scan the wrong repo.
status = {"filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}]}
cap, seen = _patch_status_capture(status)
with cap, patch("utils.paths.is_local_path", return_value = False):
d = evaluate_file_security("evil/LLM")
assert seen["repo"] == "evil/LLM" # scanned the real repo, not unsloth/evil
assert d.model_name == "evil/LLM"
assert d.blocked is True
def test_security_load_subdirs_yaml_fallback(monkeypatch):
# Tokenizer detection failed, but a YAML default of audio_type=bicodec still yields LLM.
import utils.models.model_config as mc
from utils.security import security_load_subdirs
monkeypatch.setattr(mc, "detect_audio_type", lambda *_a, **_k: None)
monkeypatch.setattr(mc, "load_model_defaults", lambda *_a, **_k: {"audio_type": "bicodec"})
assert security_load_subdirs("unsloth/Spark-TTS-0.5B") == ("LLM",)
# A non-bicodec default contributes no subdir.
monkeypatch.setattr(mc, "load_model_defaults", lambda *_a, **_k: {"audio_type": None})
assert security_load_subdirs("unsloth/Llama-3.2-1B") == ()

26
studio/backend/tests/test_models_get_model_config_case_resolution.py
View file

@ -77,3 +77,29 @@ def test_get_model_config_resolves_cached_case_before_model_checks(monkeypatch):
assert calls["is_embedding_model"] == "Org/Model"
assert calls["detect_audio_type"] == "Org/Model"
assert calls["from_identifier"] == "Org/Model"
def test_repo_in_any_hf_cache_matches_case_variant_in_legacy_cache(tmp_path, monkeypatch):
# A case-variant in a legacy/default cache must read as present (case resolution only
# covers the active cache; discard deletes case-insensitively, so detection must too,
# else a decline deletes a pre-existing user repo).
import utils.paths as paths_pkg
import huggingface_hub.constants as hf_constants
active = tmp_path / "active"
legacy = tmp_path / "legacy"
default = tmp_path / "default"
for d in (active, legacy, default):
d.mkdir()
# Differently-cased entry in the legacy cache only.
(legacy / "models--Unsloth--Foo").mkdir()
# No active-cache variant; case resolution is a no-op here.
monkeypatch.setattr(paths_pkg, "resolve_cached_repo_id_case", lambda name: name)
monkeypatch.setattr(paths_pkg, "legacy_hf_cache_dir", lambda: legacy)
monkeypatch.setattr(paths_pkg, "hf_default_cache_dir", lambda: default)
monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(active))
assert models_route._repo_in_any_hf_cache("unsloth/foo") is True
# Absent from every cache -> reported absent.
assert models_route._repo_in_any_hf_cache("unsloth/not-cached") is False

101
studio/backend/tests/test_security_gate_consistency.py Normal file
View file

@ -0,0 +1,101 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Deterministic consistency guards for the model-load security gate.
The gate spans many parallel sites (validate/load/status, the inference/training/export
workers, the preflight route); past regressions were a fix at one site with a sibling
left behind. These guards enumerate the sites mechanically (AST + source) so a new site
that drops the token or mis-reports the requirement fails here, not in a later review.
"""
import ast
from pathlib import Path
_BACKEND = Path(__file__).resolve().parent.parent
# Probes read Hub config to classify a model; a token-less call 404s on a gated repo.
# Scan callers under routes/ and core/ (probe definitions live in utils/).
_PROBE_FUNCS = {"is_vision_model", "is_embedding_model", "detect_audio_type"}
_PROBE_CALLER_ROOTS = ("routes", "core")
def _iter_caller_files():
for root in _PROBE_CALLER_ROOTS:
yield from (_BACKEND / root).rglob("*.py")
def _passes_token(call: ast.Call) -> bool:
"""True if the call passes an hf_token (keyword, or the 2nd positional slot)."""
if any(kw.arg in ("hf_token", "token") for kw in call.keywords if kw.arg is not None):
return True
return len(call.args) >= 2
def _call_name(call: ast.Call):
fn = call.func
return fn.id if isinstance(fn, ast.Name) else getattr(fn, "attr", None)
def test_capability_probes_thread_the_hf_token():
"""Every capability-probe caller passes the token; a token-less probe misclassifies
a gated model (the /check-vision regression)."""
offenders = []
for path in _iter_caller_files():
try:
tree = ast.parse(path.read_text())
except SyntaxError:
continue
for node in ast.walk(tree):
if isinstance(node, ast.Call) and _call_name(node) in _PROBE_FUNCS:
if not _passes_token(node):
rel = path.relative_to(_BACKEND)
offenders.append(f"{rel}:{node.lineno} {_call_name(node)}() drops the hf_token")
assert not offenders, (
"A capability probe must pass the hf_token so gated/private models classify "
"correctly:\n " + "\n ".join(offenders)
)
def test_gguf_trust_remote_code_reported_inert_not_from_yaml():
"""GGUF never executes auto_map, so requires_trust_remote_code is reported via the
resolver or False, never the raw YAML bool() (the round-6 regression)."""
src = (_BACKEND / "routes" / "inference.py").read_text()
assert "requires_trust_remote_code = bool(" not in src, (
"Report requires_trust_remote_code via _resolve_loaded_trust_remote_code "
"(non-GGUF) or set it False (GGUF); never bool(inference_config.get(...))."
)
def test_capability_detection_caches_are_token_aware():
"""Every capability cache is keyed by (model, token_fingerprint) so an unauthenticated
miss cannot poison a later authenticated lookup (the audio-cache regression)."""
src = (_BACKEND / "utils" / "models" / "model_config.py").read_text()
offenders = []
for line in src.splitlines():
stripped = line.strip()
if "_detection_cache:" in stripped and stripped.endswith("= {}"):
if "Dict[Tuple" not in stripped and "Dict[tuple" not in stripped:
offenders.append(stripped)
assert not offenders, (
"A capability cache must be keyed by (model, token_fingerprint), not the bare "
"model name:\n " + "\n ".join(offenders)
)
def test_malware_and_consent_gates_cover_the_lora_base():
"""Every worker that runs a load gate also resolves the LoRA base, so a poisoned or
custom-code base is never skipped."""
gated_workers = [
"core/inference/worker.py",
"core/export/worker.py",
"core/training/worker.py",
]
offenders = []
for rel in gated_workers:
src = (_BACKEND / rel).read_text()
runs_gate = "evaluate_file_security(" in src or "evaluate_remote_code_consent" in src
resolves_base = "get_base_model_from_lora_identifier(" in src or "base_model" in src
if runs_gate and not resolves_base:
offenders.append(f"{rel} runs a load gate but never resolves the LoRA base")
assert not offenders, "\n".join(offenders)

78
studio/backend/tests/test_trained_model_scan.py
View file

@ -24,6 +24,7 @@ from utils.models.model_config import (
ModelConfig,
get_base_model_from_checkpoint,
get_base_model_from_lora,
get_base_model_from_lora_identifier,
scan_trained_models,
)
@ -76,6 +77,83 @@ def test_get_base_model_from_lora_rejects_full_finetune_dirs(tmp_path: Path):
assert get_base_model_from_lora(str(tmp_path)) is None
def test_lora_identifier_resolves_local_dir_like_the_local_helper(tmp_path: Path):
# Local path: behaves like the directory reader, no Hub call.
(tmp_path / "adapter_config.json").write_text(
json.dumps({"base_model_name_or_path": "HuggingFaceTB/SmolLM-135M"})
)
(tmp_path / "adapter_model.safetensors").write_bytes(b"")
with patch("huggingface_hub.hf_hub_download", side_effect = AssertionError("no Hub call")):
assert get_base_model_from_lora_identifier(str(tmp_path)) == "HuggingFaceTB/SmolLM-135M"
def test_lora_identifier_resolves_remote_adapter_base(tmp_path: Path):
# Remote adapter: the identifier helper fetches adapter_config.json from the Hub so
# the gate can scan the base, where the local helper returns None.
cfg = tmp_path / "adapter_config.json"
cfg.write_text(json.dumps({"base_model_name_or_path": "unsloth/Llama-3.2-1B-Instruct"}))
def _dl(
repo,
fn,
token = None,
):
assert repo == "someone/my-remote-lora"
assert fn == "adapter_config.json"
return str(cfg)
assert get_base_model_from_lora("someone/my-remote-lora") is None # local-only: misses it
with patch("huggingface_hub.hf_hub_download", side_effect = _dl):
base = get_base_model_from_lora_identifier("someone/my-remote-lora")
assert base == "unsloth/Llama-3.2-1B-Instruct"
def test_lora_identifier_returns_none_for_non_adapter_remote_repo():
# Non-LoRA remote repo: a 404 on adapter_config.json returns None without retrying.
from huggingface_hub.utils import EntryNotFoundError
mock = patch("huggingface_hub.hf_hub_download", side_effect = EntryNotFoundError("404"))
with mock as m:
assert get_base_model_from_lora_identifier("unsloth/Llama-3.2-1B-Instruct") is None
assert m.call_count == 1 # 404 is definitive -> no retry
def test_lora_identifier_retries_transient_then_resolves(tmp_path: Path):
# A transient error is retried (not treated as "not a LoRA"); the retry resolves the base.
cfg = tmp_path / "adapter_config.json"
cfg.write_text(json.dumps({"base_model_name_or_path": "unsloth/Llama-3.2-1B-Instruct"}))
calls = {"n": 0}
def _dl(
repo,
fn,
token = None,
):
calls["n"] += 1
if calls["n"] == 1:
raise RuntimeError("transient network blip")
return str(cfg)
with patch("huggingface_hub.hf_hub_download", side_effect = _dl):
base = get_base_model_from_lora_identifier("someone/remote-lora")
assert base == "unsloth/Llama-3.2-1B-Instruct"
assert calls["n"] == 2 # retried once
def test_lora_identifier_persistent_transient_returns_none():
# Two transient errors -> None, logged at WARNING (a missed base is gated by neither).
# Assert on the logger directly: robust to the logging backend (structlog vs stub).
from utils.models import model_config as _mc
with (
patch("huggingface_hub.hf_hub_download", side_effect = RuntimeError("down")),
patch.object(_mc.logger, "warning") as mock_warn,
):
assert get_base_model_from_lora_identifier("someone/remote-lora") is None
assert any(
"Could not resolve remote LoRA base" in str(c.args[0]) for c in mock_warn.call_args_list
)
@patch("utils.models.model_config.is_audio_input_type", return_value = False)
@patch("utils.models.model_config.detect_audio_type", return_value = None)
@patch("utils.models.model_config.is_vision_model", return_value = False)

127
studio/backend/tests/test_validate_model_error.py
View file

@ -84,3 +84,130 @@ def test_empty_runtime_error_falls_back_to_generic(monkeypatch):
http = _provoke(monkeypatch, RuntimeError(""))
assert http.status_code == 400
assert http.detail == "Invalid model"
def _drive_validate(monkeypatch, *, is_gguf: bool):
"""Run validate_model with both security helpers forced True; return the response."""
from types import SimpleNamespace
import utils.models.model_config as mc
monkeypatch.setattr(
inf,
"_resolve_model_identifier_for_request",
lambda request, operation: ("org/mixed-repo", "org/mixed-repo", False),
)
config = SimpleNamespace(
identifier = "org/mixed-repo",
display_name = "org/mixed-repo",
is_gguf = is_gguf,
is_lora = False,
is_vision = False,
gguf_file = None,
)
monkeypatch.setattr(inf.ModelConfig, "from_identifier", staticmethod(lambda **_kw: config))
# No LoRA base to resolve; keep it offline.
monkeypatch.setattr(mc, "get_base_model_from_lora_identifier", lambda *_a, **_k: None)
# Both gates WOULD flag this repo (mixed repo with auto_map + an unsafe pickle).
monkeypatch.setattr(inf, "_requires_trust_remote_code_for_model", lambda *_a, **_k: True)
monkeypatch.setattr(inf, "_requires_security_review_for_model", lambda *_a, **_k: True)
req = ValidateModelRequest(model_path = "org/mixed-repo")
return asyncio.run(inf.validate_model(req, current_subject = "tester"))
def test_selected_gguf_variant_skips_trc_and_security_review(monkeypatch):
# GGUF loads via llama.cpp: auto_map and root pickles are inert, so neither gate fires.
resp = _drive_validate(monkeypatch, is_gguf = True)
assert resp.is_gguf is True
assert resp.requires_trust_remote_code is False
assert resp.requires_security_review is False
def test_non_gguf_load_still_runs_trc_and_security_review(monkeypatch):
# Control: a Transformers (non-GGUF) load must still honor both gates.
resp = _drive_validate(monkeypatch, is_gguf = False)
assert resp.is_gguf is False
assert resp.requires_trust_remote_code is True
assert resp.requires_security_review is True
def test_resolve_loaded_trc_prefers_stored_value():
# A value stored at load time wins, so a status refresh does not re-derive it.
assert (
inf._resolve_loaded_trust_remote_code("org/m", {"requires_trust_remote_code": True}, {})
is True
)
assert (
inf._resolve_loaded_trust_remote_code(
"org/m", {"requires_trust_remote_code": False}, {"trust_remote_code": True}
)
is False
)
def test_resolve_loaded_trc_uses_runtime_and_yaml():
# No stored value: the trust_remote_code the load used, then the YAML default.
assert (
inf._resolve_loaded_trust_remote_code("org/m", {}, {}, trust_remote_code_used = True) is True
)
assert inf._resolve_loaded_trust_remote_code("org/m", {}, {"trust_remote_code": True}) is True
def test_resolve_loaded_trc_falls_back_to_raw_auto_map(monkeypatch):
# No stored value or runtime/YAML signal: fall back to the raw auto_map check.
monkeypatch.setattr(inf, "_requires_trust_remote_code_for_model", lambda *_a, **_k: True)
assert inf._resolve_loaded_trust_remote_code("org/custom", {}, {}) is True
monkeypatch.setattr(inf, "_requires_trust_remote_code_for_model", lambda *_a, **_k: False)
assert inf._resolve_loaded_trust_remote_code("org/plain", {}, {}) is False
def _drive_validate_lora(monkeypatch, *, adapter_needs_trc, base_needs_trc):
"""Run validate_model for a LoRA adapter whose base resolves, with per-target
trust_remote_code answers; return the response."""
from types import SimpleNamespace
import utils.models.model_config as mc
adapter, base = "org/lora-adapter", "org/base-model"
monkeypatch.setattr(
inf,
"_resolve_model_identifier_for_request",
lambda request, operation: (adapter, adapter, False),
)
config = SimpleNamespace(
identifier = adapter,
display_name = adapter,
is_gguf = False,
is_lora = True,
is_vision = False,
gguf_file = None,
)
monkeypatch.setattr(inf.ModelConfig, "from_identifier", staticmethod(lambda **_kw: config))
monkeypatch.setattr(mc, "get_base_model_from_lora_identifier", lambda *_a, **_k: base)
trc = {adapter: adapter_needs_trc, base: base_needs_trc}
monkeypatch.setattr(
inf,
"_requires_trust_remote_code_for_model",
lambda target, *_a, **_k: trc.get(target, False),
)
monkeypatch.setattr(inf, "_requires_security_review_for_model", lambda *_a, **_k: False)
req = ValidateModelRequest(model_path = adapter)
return asyncio.run(inf.validate_model(req, current_subject = "tester"))
def test_validate_lora_flags_trc_from_adapter_only(monkeypatch):
# Adapter ships auto_map, base does not: the requirement follows either repo.
resp = _drive_validate_lora(monkeypatch, adapter_needs_trc = True, base_needs_trc = False)
assert resp.requires_trust_remote_code is True
def test_validate_lora_flags_trc_from_base_only(monkeypatch):
# The classic case: the base ships custom code, the adapter does not.
resp = _drive_validate_lora(monkeypatch, adapter_needs_trc = False, base_needs_trc = True)
assert resp.requires_trust_remote_code is True
def test_validate_lora_clean_when_neither_needs_trc(monkeypatch):
resp = _drive_validate_lora(monkeypatch, adapter_needs_trc = False, base_needs_trc = False)
assert resp.requires_trust_remote_code is False

109
studio/backend/tests/test_vision_cache.py
View file

@ -96,11 +96,13 @@ class TestVisionCacheSubprocessPath:
The cache should spawn the subprocess at most once per model per
process."""
@patch("utils.models.model_config._raw_config_has_vision_config", return_value = None)
@patch("utils.models.model_config._is_vision_model_subprocess", return_value = True)
@patch("utils.transformers_version.needs_transformers_5", return_value = True)
def test_subprocess_called_once_with_cache(self, mock_needs_t5, mock_subprocess):
"""Subprocess fires only on the first call; second is cached."""
# First call: uncached → subprocess
def test_subprocess_called_once_with_cache(self, mock_needs_t5, mock_subprocess, mock_raw):
"""When the raw-config reader is inconclusive (None), the transformers
5.x subprocess fires only on the first call; the second is cached."""
# First call: raw None -> subprocess
assert is_vision_model("unsloth/Qwen3.5-2B") is True
# Second call: cache hit, no subprocess
assert is_vision_model("unsloth/Qwen3.5-2B") is True
@ -111,14 +113,16 @@ class TestVisionCacheSubprocessPath:
@patch("utils.models.model_config._raw_config_has_vision_config", return_value = True)
@patch("utils.models.model_config._is_vision_model_subprocess", return_value = None)
@patch("utils.transformers_version.needs_transformers_5", return_value = True)
def test_subprocess_none_falls_back_to_raw_vision_config(
def test_raw_config_primary_skips_subprocess(
self, mock_needs_t5, mock_subprocess, mock_raw_config
):
# The raw config.json read is the primary path; a definitive answer there never
# reaches the transformers-5.x subprocess or needs_transformers_5 routing.
assert is_vision_model("unsloth/gemma-4-E4B-it") is True
assert is_vision_model("unsloth/gemma-4-E4B-it") is True
mock_subprocess.assert_called_once()
mock_raw_config.assert_called_once_with("unsloth/gemma-4-E4B-it", hf_token = None)
mock_subprocess.assert_not_called()
# ---------------------------------------------------------------------------
@ -259,9 +263,10 @@ class TestVisionCacheDirectPath:
"""Models that do NOT need transformers 5.x detect via
load_model_config directly. The cache must work the same way."""
@patch("utils.models.model_config._raw_config_has_vision_config", return_value = None)
@patch("utils.transformers_version.needs_transformers_5", return_value = False)
@patch("utils.models.model_config.load_model_config")
def test_direct_vlm_detection_cached(self, mock_load_config, mock_needs_t5):
def test_direct_vlm_detection_cached(self, mock_load_config, mock_needs_t5, mock_raw):
"""A standard VLM detected via architecture suffix should be cached."""
cfg = MagicMock(spec = []) # strict: only explicitly set attrs exist
cfg.model_type = "gemma3"
@ -273,9 +278,10 @@ class TestVisionCacheDirectPath:
# load_model_config should only be called once
mock_load_config.assert_called_once()
@patch("utils.models.model_config._raw_config_has_vision_config", return_value = None)
@patch("utils.transformers_version.needs_transformers_5", return_value = False)
@patch("utils.models.model_config.load_model_config")
def test_direct_non_vlm_detection_cached(self, mock_load_config, mock_needs_t5):
def test_direct_non_vlm_detection_cached(self, mock_load_config, mock_needs_t5, mock_raw):
"""A standard text model (no VLM indicators) should cache False."""
cfg = MagicMock(spec = []) # spec=[] means no attributes at all
cfg.model_type = "llama"
@ -287,9 +293,12 @@ class TestVisionCacheDirectPath:
assert is_vision_model("meta-llama/Llama-3-8B") is False
mock_load_config.assert_called_once()
@patch("utils.models.model_config._raw_config_has_vision_config", return_value = None)
@patch("utils.transformers_version.needs_transformers_5", return_value = False)
@patch("utils.models.model_config.load_model_config")
def test_vision_config_attr_detected_and_cached(self, mock_load_config, mock_needs_t5):
def test_vision_config_attr_detected_and_cached(
self, mock_load_config, mock_needs_t5, mock_raw
):
"""Models with vision_config (LLaVA, Qwen2-VL, etc.) should be cached as True."""
cfg = MagicMock(spec = []) # strict: only explicitly set attrs exist
cfg.model_type = "qwen2_vl"
@ -301,9 +310,10 @@ class TestVisionCacheDirectPath:
assert is_vision_model("Qwen/Qwen2-VL-7B") is True
mock_load_config.assert_called_once()
@patch("utils.models.model_config._raw_config_has_vision_config", return_value = None)
@patch("utils.transformers_version.needs_transformers_5", return_value = False)
@patch("utils.models.model_config.load_model_config")
def test_gemma4_model_type_detected_and_cached(self, mock_load_config, mock_needs_t5):
def test_gemma4_model_type_detected_and_cached(self, mock_load_config, mock_needs_t5, mock_raw):
cfg = MagicMock(spec = [])
cfg.model_type = "gemma4"
cfg.architectures = ["Gemma4ForConditionalGeneration"]
@ -313,9 +323,12 @@ class TestVisionCacheDirectPath:
assert is_vision_model("google/gemma-4-E4B-it") is True
mock_load_config.assert_called_once()
@patch("utils.models.model_config._raw_config_has_vision_config", return_value = None)
@patch("utils.transformers_version.needs_transformers_5", return_value = False)
@patch("utils.models.model_config.load_model_config")
def test_gemma4_audio_subconfig_not_detected_as_vision(self, mock_load_config, mock_needs_t5):
def test_gemma4_audio_subconfig_not_detected_as_vision(
self, mock_load_config, mock_needs_t5, mock_raw
):
cfg = MagicMock(spec = [])
cfg.model_type = "gemma4_audio"
cfg.architectures = ["Gemma4AudioModel"]
@ -325,9 +338,12 @@ class TestVisionCacheDirectPath:
assert is_vision_model("local/gemma4-audio-encoder") is False
mock_load_config.assert_called_once()
@patch("utils.models.model_config._raw_config_has_vision_config", return_value = None)
@patch("utils.transformers_version.needs_transformers_5", return_value = False)
@patch("utils.models.model_config.load_model_config")
def test_gemma4_text_subconfig_not_detected_as_vision(self, mock_load_config, mock_needs_t5):
def test_gemma4_text_subconfig_not_detected_as_vision(
self, mock_load_config, mock_needs_t5, mock_raw
):
cfg = MagicMock(spec = [])
cfg.model_type = "gemma4_text"
cfg.architectures = ["Gemma4ForCausalLM"]
@ -337,9 +353,10 @@ class TestVisionCacheDirectPath:
assert is_vision_model("local/gemma-4-text") is False
mock_load_config.assert_called_once()
@patch("utils.models.model_config._raw_config_has_vision_config", return_value = None)
@patch("utils.transformers_version.needs_transformers_5", return_value = False)
@patch("utils.models.model_config.load_model_config")
def test_audio_model_excluded_and_cached(self, mock_load_config, mock_needs_t5):
def test_audio_model_excluded_and_cached(self, mock_load_config, mock_needs_t5, mock_raw):
"""Audio-only models (csm, whisper) with ForConditionalGeneration
should be excluded from VLM detection and cached as False."""
cfg = MagicMock(spec = []) # strict: only explicitly set attrs exist
@ -489,7 +506,8 @@ class TestVlmAudioExclusion:
fallback, and the inlined subprocess helper too."""
def test_audio_only_set_canonical(self):
assert _AUDIO_ONLY_MODEL_TYPES == {"csm", "whisper"}
# Derived from the transformers audio registry, so a superset of {csm, whisper}.
assert {"csm", "whisper"} <= _AUDIO_ONLY_MODEL_TYPES
def test_is_vlm_excludes_whisper(self):
cfg = MagicMock(spec = [])
@ -528,3 +546,68 @@ class TestVlmAudioExclusion:
},
)
assert is_vision_model(str(tmp_path)) is False
class TestAudioDetectionCacheTokenAware:
"""The audio cache mirrors the vision cache: keyed by (model, token_fingerprint)
so an unauthenticated miss cannot poison a later authenticated lookup."""
def test_audio_cache_is_token_aware(self, monkeypatch):
import utils.models.model_config as mc
mc._audio_detection_cache.clear()
calls = []
def _fake(name, hf_token = None):
calls.append(hf_token)
# Gated repo: only an authenticated probe can read the tokenizer.
return ("bicodec", True) if hf_token else (None, True)
monkeypatch.setattr(mc, "_detect_audio_from_tokenizer", _fake)
monkeypatch.setattr(mc, "is_local_path", lambda *_a, **_k: False)
monkeypatch.setattr(mc, "resolve_cached_repo_id_case", lambda n, *_a, **_k: n)
# Unauthenticated miss caches None under (name, None)...
assert mc.detect_audio_type("private/spark") is None
# ...but the authenticated call uses a different key and is NOT poisoned.
assert mc.detect_audio_type("private/spark", hf_token = "hf_x") == "bicodec"
assert calls == [None, "hf_x"]
# Same (model, token) is served from cache (no third probe).
assert mc.detect_audio_type("private/spark", hf_token = "hf_x") == "bicodec"
assert calls == [None, "hf_x"]
mc._audio_detection_cache.clear()
def test_transient_none_is_not_cached_but_definitive_none_is(self, monkeypatch):
"""A transient probe failure (definitive=False) must retry; a clean
'not audio' read (definitive=True) caches so we don't re-probe."""
import utils.models.model_config as mc
mc._audio_detection_cache.clear()
monkeypatch.setattr(mc, "is_local_path", lambda *_a, **_k: False)
monkeypatch.setattr(mc, "resolve_cached_repo_id_case", lambda n, *_a, **_k: n)
transient_calls = []
def _transient(name, hf_token = None):
transient_calls.append(hf_token)
return (None, False) # network/5xx -- not cacheable
monkeypatch.setattr(mc, "_detect_audio_from_tokenizer", _transient)
assert mc.detect_audio_type("flaky/model") is None
assert mc.detect_audio_type("flaky/model") is None
# Re-probed both times: the transient None was never cached.
assert transient_calls == [None, None]
definitive_calls = []
def _definitive(name, hf_token = None):
definitive_calls.append(hf_token)
return (None, True) # read the config, no audio tokens
monkeypatch.setattr(mc, "_detect_audio_from_tokenizer", _definitive)
assert mc.detect_audio_type("plain/text-model") is None
assert mc.detect_audio_type("plain/text-model") is None
# Probed once: the definitive None was cached.
assert definitive_calls == [None]
mc._audio_detection_cache.clear()

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

@ -1136,14 +1136,22 @@ def _get_hf_safetensors_total_params(
def _load_config_for_gpu_estimate(model_name: str, hf_token: Optional[str] = None):
# Estimation needs only declarative config.json fields, and this probe runs
# on model selection, so read raw config.json (never run auto_map Python) and
# expose it as an attribute namespace for downstream getattr access.
try:
from transformers import AutoConfig
trust_remote_code = model_name.lower().startswith("unsloth/")
return AutoConfig.from_pretrained(
model_name,
token = hf_token,
trust_remote_code = trust_remote_code,
)
from utils.transformers_version import _load_config_json
cfg = _load_config_json(model_name, hf_token = hf_token)
if cfg is None:
return None
def _to_ns(d):
if isinstance(d, dict):
return types.SimpleNamespace(**{k: _to_ns(v) for k, v in d.items()})
return d
return _to_ns(cfg)
except Exception as e:
logger.warning("Could not load config for '%s': %s", model_name, e)
return None

2
studio/backend/utils/models/__init__.py
View file

@ -16,6 +16,7 @@ from .model_config import (
get_base_model_from_checkpoint,
load_model_defaults,
get_base_model_from_lora,
get_base_model_from_lora_identifier,
load_model_config,
list_gguf_variants,
extract_model_size_b,
@ -40,6 +41,7 @@ __all__ = [
"get_base_model_from_checkpoint",
"load_model_defaults",
"get_base_model_from_lora",
"get_base_model_from_lora_identifier",
"load_model_config",
"list_gguf_variants",
"extract_model_size_b",

324
studio/backend/utils/models/model_config.py
View file

@ -470,9 +470,15 @@ def load_model_config(
model_name: str,
use_auth: bool = False,
token: Optional[str] = None,
trust_remote_code: bool = True,
trust_remote_code: bool = False,
):
"""Load model config with optional authentication control."""
"""Load model config with optional authentication control.
``trust_remote_code`` defaults to ``False``: capability detection and
metadata lookups must never execute a model repo's ``auto_map`` Python.
Deliberate remote-code loads pass the flag explicitly through
``FastLanguageModel.from_pretrained`` with the user's own consent.
"""
from transformers import AutoConfig
if token:
@ -496,22 +502,72 @@ def load_model_config(
)
# VLM architecture suffixes and known VLM model_type values.
_VLM_ARCH_SUFFIXES = ("ForConditionalGeneration", "ForVisionText2Text")
_VLM_MODEL_TYPES = {
"phi3_v",
"llava",
"llava_next",
"llava_onevision",
"internvl_chat",
"cogvlm2",
"minicpmv",
"gemma4",
}
# Detection sets come from the installed transformers registry, unioned with a
# small curated set of auto_map VLMs (DeepSeek-OCR, Kimi, phi3_v) whose arch is
# repo-defined and absent from the registry. ForConditionalGeneration is NOT a
# vision signal (overloaded across text/audio/vision); ForVisionText2Text is.
_VLM_ARCH_SUFFIXES = ("ForVisionText2Text",)
# Audio-only models that share the ForConditionalGeneration suffix
# (e.g. CsmForConditionalGeneration, WhisperForConditionalGeneration).
_AUDIO_ONLY_MODEL_TYPES = {"csm", "whisper"}
_CURATED_REMOTE_VLM_TYPES = frozenset(
{
"phi3_v",
"llava",
"llava_next",
"llava_onevision",
"internvl_chat",
"cogvlm2",
"minicpmv",
"gemma4",
"deepseek_vl_v2",
"kimi_k25",
}
)
# Fallbacks used only if the transformers registry import fails.
_FALLBACK_AUDIO_MODEL_TYPES = frozenset({"csm", "whisper"})
def _build_detection_sets():
"""Return (vlm_model_types, vlm_class_names, audio_model_types) from the
installed transformers registry, unioned with the curated repo-code VLM
set. Reads only static name dicts -- no model is loaded, no code runs.
Falls back to curated/hardcoded values if transformers is unavailable.
"""
try:
from transformers.models.auto import modeling_auto as _ma
def _names(attr):
d = getattr(_ma, attr, None)
return dict(d) if d else {}
itt = _names("MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES")
v2s = _names("MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES")
vlm_types = set(itt) | set(v2s) | set(_CURATED_REMOTE_VLM_TYPES)
vlm_classes = set(itt.values()) | set(v2s.values())
audio_types: set = set()
for attr in (
"MODEL_FOR_CTC_MAPPING_NAMES",
"MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES",
"MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES",
"MODEL_FOR_TEXT_TO_WAVEFORM_MAPPING_NAMES",
"MODEL_FOR_TEXT_TO_SPECTROGRAM_MAPPING_NAMES",
"MODEL_FOR_AUDIO_XVECTOR_MAPPING_NAMES",
):
audio_types |= set(_names(attr))
audio_types |= set(_FALLBACK_AUDIO_MODEL_TYPES)
return frozenset(vlm_types), frozenset(vlm_classes), frozenset(audio_types)
except Exception as exc: # pragma: no cover - defensive
logger.warning("Could not build detection sets from transformers: %s", exc)
return (
frozenset(_CURATED_REMOTE_VLM_TYPES),
frozenset(),
frozenset(_FALLBACK_AUDIO_MODEL_TYPES),
)
_VLM_MODEL_TYPES, _VLM_CLASS_NAMES, _AUDIO_ONLY_MODEL_TYPES = _build_detection_sets()
# Pre-computed .venv_t5 paths and backend dir for subprocess version switching.
# Vision check uses the Gemma 4 5.5 sidecar for existing Gemma 4 architectures.
@ -524,13 +580,19 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent.parent)
def _is_vlm(config) -> bool:
architectures = getattr(config, "architectures", None) or []
model_type = getattr(config, "model_type", None)
if model_type in _AUDIO_ONLY_MODEL_TYPES:
return False
return (
any(x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures)
or hasattr(config, "vision_config")
explicit_vision = (
hasattr(config, "vision_config")
or hasattr(config, "img_processor")
or hasattr(config, "image_token_index")
or hasattr(config, "projector_config")
)
# Audio-only models are vision only if they carry an explicit vision sub-config.
if model_type in _AUDIO_ONLY_MODEL_TYPES and not explicit_vision:
return False
return (
explicit_vision
or any(x in _VLM_CLASS_NAMES for x in architectures)
or any(isinstance(x, str) and x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures)
or model_type in _VLM_MODEL_TYPES
)
@ -553,13 +615,19 @@ def _raw_config_has_vision_config(
config = json.loads(config_path.read_text())
architectures = config.get("architectures") or []
model_type = config.get("model_type")
if model_type in _AUDIO_ONLY_MODEL_TYPES:
return False
return (
any(isinstance(x, str) and x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures)
or "vision_config" in config
explicit_vision = (
"vision_config" in config
or "img_processor" in config
or "image_token_index" in config
or "projector_config" in config
)
# Audio-only models are vision only if they carry an explicit vision sub-config.
if model_type in _AUDIO_ONLY_MODEL_TYPES and not explicit_vision:
return False
return (
explicit_vision
or any(isinstance(x, str) and x in _VLM_CLASS_NAMES for x in architectures)
or any(isinstance(x, str) and x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures)
or model_type in _VLM_MODEL_TYPES
)
except Exception as exc:
@ -570,19 +638,25 @@ def _raw_config_has_vision_config(
# why: inline _is_vlm and constants are prepended so the subprocess stays
# self-contained and does not import the parent backend module graph.
_VISION_CHECK_INLINE_HELPERS = (
"_VLM_ARCH_SUFFIXES = " + repr(_VLM_ARCH_SUFFIXES) + "\n"
"_VLM_MODEL_TYPES = " + repr(_VLM_MODEL_TYPES) + "\n"
"_AUDIO_ONLY_MODEL_TYPES = " + repr(_AUDIO_ONLY_MODEL_TYPES) + "\n"
"_VLM_ARCH_SUFFIXES = " + repr(tuple(_VLM_ARCH_SUFFIXES)) + "\n"
"_VLM_MODEL_TYPES = " + repr(set(_VLM_MODEL_TYPES)) + "\n"
"_VLM_CLASS_NAMES = " + repr(set(_VLM_CLASS_NAMES)) + "\n"
"_AUDIO_ONLY_MODEL_TYPES = " + repr(set(_AUDIO_ONLY_MODEL_TYPES)) + "\n"
"def _is_vlm(config):\n"
" architectures = getattr(config, 'architectures', None) or []\n"
" model_type = getattr(config, 'model_type', None)\n"
" if model_type in _AUDIO_ONLY_MODEL_TYPES:\n"
" return False\n"
" return (\n"
" any(x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures)\n"
" or hasattr(config, 'vision_config')\n"
" explicit_vision = (\n"
" hasattr(config, 'vision_config')\n"
" or hasattr(config, 'img_processor')\n"
" or hasattr(config, 'image_token_index')\n"
" or hasattr(config, 'projector_config')\n"
" )\n"
" if model_type in _AUDIO_ONLY_MODEL_TYPES and not explicit_vision:\n"
" return False\n"
" return (\n"
" explicit_vision\n"
" or any(x in _VLM_CLASS_NAMES for x in architectures)\n"
" or any(isinstance(x, str) and x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures)\n"
" or model_type in _VLM_MODEL_TYPES\n"
" )\n"
)
@ -610,7 +684,8 @@ if backend_dir not in sys.path:
try:
from transformers import AutoConfig
kwargs = {"trust_remote_code": True}
# Capability detection never executes model repo code.
kwargs = {"trust_remote_code": False}
if token:
kwargs["token"] = token
config = AutoConfig.from_pretrained(model_name, **kwargs)
@ -780,8 +855,15 @@ def _is_vision_model_uncached(model_name: str, hf_token: Optional[str] = None) -
Returns True/False for definitive results, or None on transient errors
(network, timeout, subprocess failure) so the caller knows not to cache.
"""
# Models needing transformers 5.x must be checked in a subprocess: the main
# process (transformers 4.57.x) doesn't recognize their architectures.
# Try the raw-config reader FIRST (code-free, version-independent): it classifies
# repo-code VLMs like DeepSeek-OCR via declarative vision_config with no remote-code
# execution or transformers-5.x subprocess.
raw = _raw_config_has_vision_config(model_name, hf_token = hf_token)
if raw is not None:
return raw
# Raw read failed transiently: fall back to AutoConfig with remote code DISABLED
# (in a transformers-5.x subprocess when the main process can't parse the arch).
from utils.transformers_version import needs_transformers_5
if needs_transformers_5(model_name):
@ -789,21 +871,13 @@ def _is_vision_model_uncached(model_name: str, hf_token: Optional[str] = None) -
"Model '%s' needs transformers 5.x -- checking vision via subprocess",
model_name,
)
result = _is_vision_model_subprocess(model_name, hf_token = hf_token)
if result is not None:
return result
return _raw_config_has_vision_config(model_name, hf_token = hf_token)
return _is_vision_model_subprocess(model_name, hf_token = hf_token)
try:
config = load_model_config(model_name, use_auth = True, token = hf_token)
# Exclude audio-only models sharing the ForConditionalGeneration suffix
# (e.g. CsmForConditionalGeneration, WhisperForConditionalGeneration)
model_type = getattr(config, "model_type", None)
if model_type in _AUDIO_ONLY_MODEL_TYPES:
return False
if _is_vlm(config):
model_type = getattr(config, "model_type", None)
archs = getattr(config, "architectures", None) or []
logger.info(
"Model %s detected as VLM (model_type=%s, architectures=%s)",
@ -840,8 +914,9 @@ def _is_vision_model_uncached(model_name: str, hf_token: Optional[str] = None) -
VALID_AUDIO_TYPES = ("snac", "csm", "bicodec", "dac", "whisper", "audio_vlm")
# Cache detection per session to avoid repeated API calls
_audio_detection_cache: Dict[str, Optional[str]] = {}
# Keyed by (normalized_name, token_fingerprint) like the vision cache, so an
# unauthenticated miss (None) cannot poison a later authenticated lookup.
_audio_detection_cache: Dict[Tuple[str, Optional[str]], Optional[str]] = {}
# Tokenizer token patterns → audio_type (all 6 types from tokenizer_config.json)
_AUDIO_TOKEN_PATTERNS = {
@ -867,22 +942,38 @@ def detect_audio_type(model_name: str, hf_token: Optional[str] = None) -> Option
Returns an audio_type string ('snac', 'csm', 'bicodec', 'dac', 'whisper',
'audio_vlm') or None.
"""
if model_name in _audio_detection_cache:
return _audio_detection_cache[model_name]
# Normalize casing + include the token fingerprint (mirrors is_vision_model).
try:
if is_local_path(model_name):
resolved_name = normalize_path(model_name)
else:
resolved_name = resolve_cached_repo_id_case(model_name)
except Exception:
resolved_name = model_name
cache_key = (resolved_name, _token_fingerprint(hf_token))
if cache_key in _audio_detection_cache:
return _audio_detection_cache[cache_key]
result = _detect_audio_from_tokenizer(model_name, hf_token)
_audio_detection_cache[model_name] = result
result, definitive = _detect_audio_from_tokenizer(model_name, hf_token)
# Cache only definitive results; a transient read failure stays None and retries.
if definitive:
_audio_detection_cache[cache_key] = result
if result:
logger.info(f"Model {model_name} detected as audio model: audio_type={result}")
return result
def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None) -> Optional[str]:
def _detect_audio_from_tokenizer(
model_name: str, hf_token: Optional[str] = None
) -> Tuple[Optional[str], bool]:
"""Detect audio type from tokenizer special tokens.
Checks local HF cache first, then fetches tokenizer_config.json from HF;
examines added_tokens_decoder for distinctive patterns.
Returns (audio_type_or_None, definitive). definitive is False only on a
transient read failure (network/timeout/5xx) so the caller skips caching and
retries; a successful read with no audio tokens is a definitive None.
"""
def _check_token_patterns(tok_config: dict) -> Optional[str]:
@ -895,6 +986,8 @@ def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None
return audio_type
return None
read_any = False # parsed at least one tokenizer_config -> a None is definitive
# 1) Local HF cache first (works for gated/offline models)
try:
repo_dir = get_cache_path(model_name)
@ -909,9 +1002,10 @@ def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None
tok_file = snapshot / tok_path
if tok_file.exists():
tok_config = json.loads(tok_file.read_text())
read_any = True
result = _check_token_patterns(tok_config)
if result:
return result
return result, True
except Exception as e:
logger.debug(f"Could not check local cache for {model_name}: {e}")
@ -919,28 +1013,40 @@ def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None
try:
import requests
import os
except Exception:
return None, read_any
paths_to_try = ["tokenizer_config.json", "LLM/tokenizer_config.json"]
token = hf_token or os.environ.get("HF_TOKEN")
headers = {}
if token:
headers["Authorization"] = f"Bearer {token}"
paths_to_try = ["tokenizer_config.json", "LLM/tokenizer_config.json"]
token = hf_token or os.environ.get("HF_TOKEN")
headers = {"Authorization": f"Bearer {token}"} if token else {}
for tok_path in paths_to_try:
url = f"https://huggingface.co/{model_name}/resolve/main/{tok_path}"
transient = False # a fetch failed for a non-404 reason (network/5xx)
for tok_path in paths_to_try:
url = f"https://huggingface.co/{model_name}/resolve/main/{tok_path}"
try:
resp = requests.get(url, headers = headers, timeout = 15)
if not resp.ok:
continue
except Exception as e:
logger.debug(f"Could not fetch {tok_path} for {model_name}: {e}")
transient = True
continue
if resp.status_code == 404:
continue # genuinely absent on this path
if not resp.ok:
transient = True # 5xx/403/etc -- can't tell, don't cache
continue
try:
tok_config = resp.json()
result = _check_token_patterns(tok_config)
if result:
return result
except Exception as e:
logger.debug(f"Bad tokenizer_config for {model_name}/{tok_path}: {e}")
transient = True
continue
read_any = True
result = _check_token_patterns(tok_config)
if result:
return result, True
return None
except Exception as e:
logger.debug(f"Could not detect audio type from tokenizer for {model_name}: {e}")
return None
# No audio tokens: definitive unless every attempt failed transiently.
return None, (read_any or not transient)
def is_audio_input_type(audio_type: Optional[str]) -> bool:
@ -2135,7 +2241,7 @@ def get_base_model_from_lora(lora_path: str) -> Optional[str]:
return base_model
# Fallback: try training_args.bin (requires torch)
# TODO: torch.load default weights_only=True (torch >= 2.6) rejects pickled TrainingArguments; also an RCE sink for third-party LoRAs via this route, re-enable behind a trust check if needed.
# TODO: torch.load default weights_only=True (torch >= 2.6) rejects pickled TrainingArguments; also an remote code execution sink for third-party LoRAs via this route, re-enable behind a trust check if needed.
# training_args_path = lora_path_obj / "training_args.bin"
# if training_args_path.exists():
# try:
@ -2169,6 +2275,74 @@ def get_base_model_from_lora(lora_path: str) -> Optional[str]:
return None
def get_base_model_from_lora_identifier(
identifier: str, hf_token: Optional[str] = None
) -> Optional[str]:
"""Resolve a LoRA adapter's base model for a LOCAL dir OR a REMOTE HF repo.
``get_base_model_from_lora`` only reads a local adapter directory (it requires
``is_dir()``). The SECURITY gates must also follow a *remote* adapter's base,
because the base model's code / weights are what execute on load: an attacker's
adapter repo can point ``base_model_name_or_path`` at a base carrying a poisoned
pickle or HIGH auto_map code. For a remote repo id we fetch ONLY the small
``adapter_config.json`` (metadata; never a weight file) and read the base. Use
this in the gate paths so a remote LoRA base is scanned, not just the adapter.
Returns the base model id, or ``None`` when the identifier is not a LoRA adapter
or the base cannot be determined (the caller still scans the identifier itself).
A genuine 404 (no ``adapter_config.json`` / repo absent) is distinguished from a
transient error: the latter is retried once, then logged as a WARNING (a missed
base would be scanned by neither gate), so a network blip does not silently and
invisibly skip the base.
"""
# Local path: reuse the existing directory reader (identical behavior).
try:
if is_local_path(identifier):
return get_base_model_from_lora(identifier)
except Exception:
return get_base_model_from_lora(identifier)
# Remote repo id: read base_model_name_or_path from adapter_config.json only.
from huggingface_hub import hf_hub_download
from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError
last_exc = None
for _attempt in range(2): # one retry: a transient blip must not skip the base
try:
cfg_path = hf_hub_download(
identifier, "adapter_config.json", token = hf_token if hf_token else None
)
except (EntryNotFoundError, RepositoryNotFoundError):
# No adapter_config.json -> not a resolvable LoRA; caller scans the identifier.
return None
except Exception as exc: # transient / auth / network -> retry once
last_exc = exc
continue
try:
with open(cfg_path, "r") as f:
base_model = json.load(f).get("base_model_name_or_path")
except Exception as exc:
logger.warning("Could not parse adapter_config.json for '%s': %s", identifier, exc)
return None
if base_model:
logger.info(
"Detected base model from remote adapter_config.json (%s): %s",
identifier,
base_model,
)
return base_model # may be None if the key is absent (still a valid answer)
# Both attempts failed transiently: log loudly -- a missed base is gated by neither gate.
logger.warning(
"Could not resolve remote LoRA base for '%s' after retry (%s); its base, if "
"any, will not be added to the security scan targets.",
identifier,
type(last_exc).__name__ if last_exc else "unknown",
)
return None
# Status indicators that appear in UI dropdowns
UI_STATUS_INDICATORS = [" (Ready)", " (Loading...)", " (Active)", "↓ "]

107
studio/backend/utils/security/__init__.py Normal file
View file

@ -0,0 +1,107 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Security helpers for the ``trust_remote_code`` boundary.
Two orthogonal questions: ``trusted_org.is_trusted_org_repo`` (may we AUTO-enable
remote code for this name?) and ``remote_code_scan`` (WHAT would run if the user
opts in?). The load paths try ``trust_remote_code=False`` first and, on the
transformers "requires trust_remote_code" error, scan the repo's ``auto_map``,
surface findings + a pinning fingerprint, and require explicit consent before
retrying with it enabled. Detection (is-vision / version / size) reads raw
``config.json`` and never enters this flow.
"""
from utils.security.consent import ( # noqa: F401
RemoteCodeDecision,
evaluate_remote_code_consent,
evaluate_remote_code_consent_for_targets,
)
from utils.security.file_security import ( # noqa: F401
FileSecurityDecision,
evaluate_file_security,
security_load_subdirs,
)
from utils.security.remote_code_scan import ( # noqa: F401
CRITICAL,
HIGH,
MEDIUM,
Finding,
RemoteCodeUnscannable,
ScanResult,
remote_code_fingerprint,
repo_remote_code_files,
scan_remote_code_files,
)
from utils.security.trusted_org import is_trusted_org_repo # noqa: F401
__all__ = [
"is_trusted_org_repo",
"scan_remote_code_files",
"repo_remote_code_files",
"RemoteCodeUnscannable",
"remote_code_fingerprint",
"should_block_remote_code",
"evaluate_remote_code_consent",
"evaluate_remote_code_consent_for_targets",
"preflight_remote_code_consent",
"preflight_remote_code_consent_for_targets",
"evaluate_file_security",
"security_load_subdirs",
"FileSecurityDecision",
"RemoteCodeDecision",
"ScanResult",
"Finding",
"CRITICAL",
"HIGH",
"MEDIUM",
]
def preflight_remote_code_consent(
model_name: str,
hf_token = None,
*,
trust_remote_code: bool = True,
approved_fingerprint = None,
trusted_org = None,
) -> "RemoteCodeDecision":
"""Scan a model's ``auto_map`` for the consent dialog. Thin wrapper over
``evaluate_remote_code_consent`` defaulting ``trust_remote_code=True`` so the scan
runs whenever the repo declares custom code; the start routes pass the user's real
value + approved fingerprint to enforce consent before any state mutation.
"""
return evaluate_remote_code_consent(
model_name,
hf_token,
trust_remote_code = trust_remote_code,
approved_fingerprint = approved_fingerprint,
trusted_org = trusted_org,
)
def preflight_remote_code_consent_for_targets(
targets,
hf_token = None,
*,
trust_remote_code: bool = True,
approved_fingerprint = None,
) -> "RemoteCodeDecision":
"""Preflight consent over multiple repos (a LoRA adapter plus its base) scanned as
one combined unit with a single pinning fingerprint. Wrapper defaulting
``trust_remote_code=True``; the load passes the user's real value + fingerprint.
"""
return evaluate_remote_code_consent_for_targets(
targets,
hf_token,
trust_remote_code = trust_remote_code,
approved_fingerprint = approved_fingerprint,
)
def should_block_remote_code(result: "ScanResult") -> bool:
"""Recommend blocking by default on CRITICAL/HIGH findings. Advisory only: the
caller still surfaces findings and takes explicit consent.
"""
sev = result.max_severity
return sev in (CRITICAL, HIGH)

358
studio/backend/utils/security/consent.py Normal file
View file

@ -0,0 +1,358 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Consent gate for loads that would execute model repo code.
The LOAD-path counterpart to the capability probes (which read raw config and
never need remote code). A deliberate load calls ``evaluate_remote_code_consent``
right before passing ``trust_remote_code=True``, and decides by the severity of a
static scan of the repo's ``auto_map`` ``.py``:
* No ``auto_map`` in any config (model/tokenizer/processor) -> nothing runs; allow.
* CRITICAL (reverse shell, IMDS, credential theft, droppers) -> hard block, never
approvable, even first-party (defends a compromised trusted repo).
* HIGH/MEDIUM (subprocess/exec/eval/network/b64decode, or a large embedded blob) ->
block but user-approvable: the dialog pins approval to the scanned ``fingerprint``.
Applies to EVERY repo; first-party is not a blanket bypass.
* ``auto_map`` present but unscannable (gated/offline/listing failure) -> fail
closed: hard block, since we cannot verify or fingerprint unseen code.
Hardening + consent, not a sandbox: static patterns are evadable, so subprocess /
venv isolation remains the containment layer.
"""
from dataclasses import dataclass, field
from typing import Optional
from loggers import get_logger
from utils.security.remote_code_scan import (
CRITICAL,
HIGH,
MEDIUM,
REMOTE_CODE_CONFIG_FILES,
RemoteCodeUnscannable,
remote_code_fingerprint,
repo_remote_code_files,
scan_remote_code_files,
)
logger = get_logger(__name__)
@dataclass
class RemoteCodeDecision:
"""Outcome of the consent gate for one (model, trust_remote_code) load."""
model_name: str
has_remote_code: bool
blocked: bool
fingerprint: Optional[str]
max_severity: Optional[str]
findings_summary: str
reason: str
findings: list = field(default_factory = list) # structured [{severity,file,check,evidence}]
approvable: bool = True # False only for CRITICAL (user cannot override)
def response_payload(self) -> dict:
"""Machine-readable detail for the frontend. ``error_kind`` splits a
user-approvable prompt (``remote_code_consent_required``) from a CRITICAL hard
block (``remote_code_blocked``).
"""
return {
"error_kind": (
"remote_code_consent_required" if self.approvable else "remote_code_blocked"
),
"model_name": self.model_name,
"has_remote_code": self.has_remote_code,
"approvable": self.approvable,
"fingerprint": self.fingerprint,
"max_severity": self.max_severity,
"findings": self.findings,
"findings_summary": self.findings_summary,
"reason": self.reason,
}
# trust_remote_code runs auto_map from ANY of these configs (model/tokenizer/
# processor), so all of them gate consent (scanning only config.json/tokenizer would
# miss a custom-processor VLM). The list lives in remote_code_scan so the gate and
# scanner stay in lockstep.
_REMOTE_CODE_CONFIG_FILES = REMOTE_CODE_CONFIG_FILES
def _config_has_auto_map(model_name: str, hf_token: Optional[str] = None) -> Optional[bool]:
"""Whether any config (model/tokenizer/processor) declares an ``auto_map`` the load
would execute. Reads raw JSON with ``hf_token``; returns None when a config is
unreadable (transient/auth) so the caller treats it as "unknown" and scans, False
when the repo genuinely ships none. GGUF is False (llama.cpp never runs auto_map);
this is the single chokepoint for that rule, shared by validate / scan / worker.
"""
# A direct .gguf FILE loads via llama.cpp (auto_map inert). A bare repo id ending in
# .gguf can still ship safetensors + auto_map, so it falls through to the scan.
if _is_direct_gguf_file_ref(model_name):
return False
configs = _load_remote_code_configs(model_name, hf_token)
if configs is None:
return None
if not any(bool((cfg or {}).get("auto_map")) for cfg in configs):
return False
# auto_map present but a GGUF repo -> inert. Checked only when auto_map exists, so
# normal models skip the extra listing.
if _is_gguf_repo(model_name, hf_token):
logger.debug("Ignoring auto_map for GGUF repo '%s' (llama.cpp never runs it).", model_name)
return False
return True
def _is_direct_gguf_file_ref(model_name: str) -> bool:
"""Whether ``model_name`` names a specific ``.gguf`` FILE (llama.cpp), not a repo:
a local ``.gguf`` path or a remote ``org/repo/.../file.gguf`` (>= 2 slashes). A bare
``org/name.gguf`` is a repo id that can still ship safetensors + auto_map, so it
falls through to the scan.
"""
name = model_name or ""
if not name.lower().endswith(".gguf"):
return False
try:
from utils.paths import is_local_path
if is_local_path(name):
return True
except Exception:
pass
# Remote: a file reference is repo_id ("org/name") + filename => >= 2 slashes.
return name.count("/") >= 2
# Weight formats transformers can load (and thus run auto_map for). A repo shipping any
# of these is not GGUF-only -- the user could load it through transformers -- so consent
# still applies even if it also ships a .gguf.
_TRANSFORMERS_WEIGHT_SUFFIXES = (
".safetensors",
".bin",
".pt",
".pth",
".h5",
".msgpack",
".onnx",
".ckpt",
)
def _is_gguf_repo(model_name: str, hf_token: Optional[str] = None) -> bool:
"""Whether a remote repo loads only through llama.cpp (GGUF weights and NO
transformers-loadable weights), making its config inert. A repo that also ships
transformers weights is NOT GGUF (auto_map could run, so still gate). A listing
failure is treated as "not known-GGUF" (fall through to scan).
"""
try:
from utils.paths import is_local_path
if is_local_path(model_name):
return False
from huggingface_hub import list_repo_files
files = [f.lower() for f in list_repo_files(model_name, token = hf_token)]
has_gguf = any(f.endswith(".gguf") for f in files)
has_transformers_weights = any(f.endswith(_TRANSFORMERS_WEIGHT_SUFFIXES) for f in files)
return has_gguf and not has_transformers_weights
except Exception:
return False
def _load_remote_code_configs(model_name: str, hf_token: Optional[str] = None) -> Optional[list]:
"""Read every config that can declare ``auto_map`` (model/tokenizer/processor) as
raw dicts. Returns the configs present (``[]`` when all 404, a definitive "no
auto_map"), or None when one is unreadable (transient/auth) so the caller scans.
The 404-vs-error split matters: real absence is "allow"; unreadable is "unknown".
"""
import json
from pathlib import Path
try:
from utils.paths import is_local_path, normalize_path
if is_local_path(model_name):
root = Path(normalize_path(model_name)).expanduser()
configs = []
for name in _REMOTE_CODE_CONFIG_FILES:
p = root / name
if p.is_file():
configs.append(json.loads(p.read_text()))
return configs
from huggingface_hub import hf_hub_download
from huggingface_hub.utils import EntryNotFoundError
configs = []
for name in _REMOTE_CODE_CONFIG_FILES:
try:
p = hf_hub_download(repo_id = model_name, filename = name, token = hf_token)
except EntryNotFoundError:
continue # genuine 404 -> truly absent
except Exception:
# Transient/auth failure is not "absent" -> fail closed to "unknown" so
# the caller scans (a tokenizer/processor-only auto_map must not slip by).
return None
configs.append(json.loads(Path(p).read_text()))
# Every config was read or a genuine 404 -> an empty list is a definitive
# "no auto_map", not "unknown".
return configs
except Exception as exc:
logger.debug("auto_map check could not read config for %s: %s", model_name, exc)
return None
def evaluate_remote_code_consent(
model_name: str,
hf_token: Optional[str] = None,
*,
trust_remote_code: bool,
approved_fingerprint: Optional[str] = None,
trusted_org: Optional[bool] = None,
) -> RemoteCodeDecision:
"""Single-repo consent; thin wrapper over the for_targets form. ``trusted_org`` is
accepted for backward compatibility but no longer changes the decision.
"""
return evaluate_remote_code_consent_for_targets(
[model_name],
hf_token,
trust_remote_code = trust_remote_code,
approved_fingerprint = approved_fingerprint,
)
def _fingerprint_target_key(target: str) -> str:
"""Namespace key for a target in the combined fingerprint. The pin is over CODE
BYTES, not the repo-id spelling: the scan canonicalizes a cached repo's casing while
workers pass raw input, so lowercase Hub ids (keep local paths as-is) or ``Org/Model``
vs ``org/model`` would fingerprint differently and reject a valid approval.
"""
try:
from utils.paths import is_local_path
if is_local_path(target):
return target
except Exception:
return target
return target.lower()
def evaluate_remote_code_consent_for_targets(
targets,
hf_token: Optional[str] = None,
*,
trust_remote_code: bool,
approved_fingerprint: Optional[str] = None,
) -> RemoteCodeDecision:
"""Decide whether a ``trust_remote_code=True`` load may proceed, over every repo whose
code the load would execute. A LoRA load runs adapter AND base code, so all targets
are scanned as ONE unit and pinned by ONE fingerprint over the union of their ``.py``
-- one approval covers every repo, and a base-only fingerprint can't leave an
adapter's own ``auto_map`` unreviewed. On ``blocked``, the caller surfaces
``response_payload()`` and retries with ``approved_fingerprint`` if the user accepts.
"""
targets = [t for t in dict.fromkeys(targets) if t]
primary = targets[0] if targets else ""
if not trust_remote_code:
return RemoteCodeDecision(
primary, False, False, None, None, "", "trust_remote_code disabled"
)
# Gather executable .py from every target that ships auto_map. A definitively
# auto_map-free target contributes nothing; an unreadable config is scanned anyway.
# If ANY target's code is present but unscannable, fail the whole load closed.
combined: dict = {}
has_remote_code = False
for target in targets:
if _config_has_auto_map(target, hf_token) is False:
continue
has_remote_code = True
try:
files = repo_remote_code_files(target, hf_token = hf_token)
except RemoteCodeUnscannable:
logger.warning(
"Blocking trust_remote_code load of '%s': remote code present (auto_map) "
"but could not be downloaded and scanned.",
target,
)
return RemoteCodeDecision(
target,
True,
True,
None,
None,
"Remote code is present (auto_map) but could not be downloaded and "
"scanned. Retry when the repo is reachable and the correct Hugging Face "
"token is set.",
"blocked: remote code could not be scanned",
approvable = False,
)
# Namespace filenames by (casing-normalized) target so two repos' same-named
# files stay distinct and the pin tracks code, not the repo-id spelling.
target_key = _fingerprint_target_key(target)
for filename, body in files.items():
combined[f"{target_key}\0{filename}"] = body
if not has_remote_code:
return RemoteCodeDecision(
primary, False, False, None, None, "", "no auto_map; trust_remote_code is a no-op"
)
if not combined:
# auto_map declared but no executable .py (e.g. a GGUF repo's vestigial
# auto_map) -> nothing to run -> allow.
return RemoteCodeDecision(
primary,
False,
False,
None,
None,
"",
"auto_map declared but no executable code present; trust_remote_code is a no-op",
)
result = scan_remote_code_files(combined)
fingerprint = remote_code_fingerprint(combined)
sev = result.max_severity
# CRITICAL is never approvable; a fingerprint pins approval for lower severities only.
approvable = sev != CRITICAL
approved = (
approvable and approved_fingerprint is not None and approved_fingerprint == fingerprint
)
if sev == CRITICAL:
blocked, reason = True, "blocked: scan found CRITICAL patterns"
elif approved:
blocked, reason = False, "approved by fingerprint"
elif sev == HIGH:
# HIGH is user-approvable but must pin the fingerprint via the dialog, for every
# repo including first-party (a compromised trusted repo still needs review).
blocked, reason = True, "blocked: scan found HIGH patterns; approval required"
elif sev == MEDIUM:
# MEDIUM (e.g. a big embedded base64 blob) also pins approval like HIGH, so a
# direct API caller can't run flagged code by just setting trust_remote_code=True.
blocked, reason = True, "blocked: scan found MEDIUM patterns; approval required"
else:
blocked, reason = False, "allowed: no high-risk patterns"
if blocked:
logger.warning(
"Blocking trust_remote_code load of '%s': scan severity %s (fingerprint %s)",
primary,
sev,
fingerprint[:12],
)
return RemoteCodeDecision(
primary,
True,
blocked,
fingerprint,
sev,
result.summary(),
reason,
findings = result.findings_payload(),
approvable = approvable,
)

368
studio/backend/utils/security/file_security.py Normal file
View file

@ -0,0 +1,368 @@
# 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.
"""
from dataclasses import dataclass, field
from typing import Optional
from loggers import get_logger
logger = get_logger(__name__)
# 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"})
# 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 _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 = (),
):
"""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.
"""
import json
try:
from huggingface_hub import hf_hub_download
from huggingface_hub.utils import EntryNotFoundError
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, token = hf_token or None)
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).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]):
"""``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.
"""
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,
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 evaluate_file_security(
model_name: str,
hf_token: Optional[str] = None,
*,
load_subdirs = (),
) -> 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.
"""
# 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 (including a local .gguf) have no Hub scan. A remote ref is scanned
# even if named "*.gguf", so a repo cannot dodge the scan via its name.
try:
from utils.paths import is_local_path
if is_local_path(model_name):
return FileSecurityDecision(model_name, False, reason = "local path; no Hub scan")
except Exception:
# Cannot classify the path -> do not block on that account.
return FileSecurityDecision(model_name, False, reason = "path check failed; not blocked")
status = _fetch_security_status(model_name, hf_token)
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)
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}",
)

679
studio/backend/utils/security/remote_code_scan.py Normal file
View file

@ -0,0 +1,679 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Static scan of a model's ``auto_map`` remote code, for the consent gate.
When a user opts into ``trust_remote_code``, the repo's ``auto_map`` Python
(``modeling_*.py`` etc.) is scanned BEFORE execution and suspicious patterns are
surfaced to inform consent. A warning aid, not a hard boundary: a determined
attacker can obfuscate past regexes, so the job is to raise the bar and inform the
hash-pinned consent. Containment (subprocess/venv) is separate; execution still
requires opt-in.
Single source of truth: ``scripts/scan_packages.py`` (the scanner CI runs via
``security-audit.yml``). We import its ``check_py_file`` so the gate inherits every
CI improvement with no drift; its heuristics are deliberately low-false-positive
(combinations flag, not bare ``subprocess``/``eval``). When ``scripts/`` is absent
(stripped install) we fall back to ``_FALLBACK_PATTERNS`` below; a test asserts the
canonical scanner loads in-repo so the fallback never silently takes over.
"""
from __future__ import annotations
import hashlib
import importlib.util
import pathlib
import re
import sys
from dataclasses import dataclass, field
from typing import Optional
from loggers import get_logger
logger = get_logger(__name__)
CRITICAL = "CRITICAL"
HIGH = "HIGH"
MEDIUM = "MEDIUM"
_SEVERITY_ORDER = {CRITICAL: 0, HIGH: 1, MEDIUM: 2}
# Configs that can carry an ``auto_map`` pointing at executable repo ``.py``.
# ``trust_remote_code`` runs code from ANY of these, so scanner and gate must read the
# same set (scanning only config.json/tokenizer would miss a custom-processor VLM).
REMOTE_CODE_CONFIG_FILES = (
"config.json",
"tokenizer_config.json",
"preprocessor_config.json",
"processor_config.json",
"video_preprocessor_config.json",
)
class RemoteCodeUnscannable(Exception):
"""The repo's executable code could not be fully fetched/listed to scan.
Raised (not returned empty) so the gate distinguishes "code PRESENT but unreadable"
(offline/gated/transient/404/listing failure) -> fail CLOSED, from "repo has NO .py"
(empty result) -> trust_remote_code is a no-op -> allow. Conflating them would block
a code-free repo or fail open on code we could not see.
"""
# Fallback patterns (used only if scripts/scan_packages.py is absent): (regex, check,
# severity). A flat subset of the canonical scanner so a stripped install still scans;
# the canonical scanner (imported below) supersedes it whenever the repo is present.
_FALLBACK_PATTERNS: tuple[tuple[re.Pattern, str, str], ...] = (
(
re.compile(
r"\bexec\s*\(\s*(?:urllib|requests|httpx|urlopen)"
r"|\bexec\s*\([^)]*\.(?:text|content|read)\s*\("
r"|\beval\s*\([^)]*\.(?:text|content|read)\s*\("
r"|\b__import__\s*\([^)]*\+",
re.DOTALL,
),
"loads-and-executes-remote-code",
CRITICAL,
),
(
re.compile(
r"\bsocket\b.*\bconnect\b.*\bsubprocess\b"
r"|\bsocket\b.*\bconnect\b.*\b(?:sh|bash|cmd)\b"
r"|\bpty\s*\.\s*spawn\b|\bos\s*\.\s*dup2\s*\(",
re.DOTALL,
),
"reverse/bind-shell",
CRITICAL,
),
(
re.compile(
r"169\.254\.169\.254|metadata\.google\.internal|/latest/meta-data"
r"|/metadata/identity|169\.254\.170\.2"
),
"cloud-metadata/IMDS-access",
CRITICAL,
),
(
re.compile(
r"(?:open|Path|read_text|read_bytes)\s*\([^)]*?"
r"(?:\.ssh[/\\]|\.aws[/\\]|\.kube[/\\]|\.gnupg[/\\]|id_rsa|id_ed25519"
r"|credentials\.json|\.git-credentials|\.npmrc|\.pypirc|/etc/shadow)"
r"|(?:open|Path)\(\s*['\"]\.env['\"]\s*[,)]",
re.DOTALL,
),
"credential-file-access",
CRITICAL,
),
(
re.compile(r"/tmp/\S+.*(?:subprocess|os\.system|os\.popen|Popen|chmod.*\+x)", re.DOTALL),
"tmp-staged-dropper",
CRITICAL,
),
(
re.compile(r"\bopenssl\s+(enc|rand|rsautl|pkeyutl|genrsa|dgst|s_client)\b"),
"openssl-cli-exfil",
HIGH,
),
(
re.compile(
r"\bsubprocess\s*\.\s*(Popen|call|run|check_call|check_output)\b"
r"|\bos\s*\.\s*(system|popen|exec[lv]p?e?)\b"
),
"subprocess/os-exec",
HIGH,
),
# Bare exec()/eval() only; the (?<![\w.]) excludes torch's ``module.eval()``.
(re.compile(r"(?<![\w.])(?:exec|eval)\s*\("), "exec/eval", HIGH),
(
re.compile(
r"\burllib\.request\b|\burlopen\s*\("
r"|\brequests\s*\.\s*(get|post|put|patch|delete|head|Session)\b"
r"|\bhttpx\s*\.\s*(get|post|put|patch|delete|Client|AsyncClient)\b"
r"|\bsocket\s*\.\s*(socket|create_connection)\b|\bhttp\.client\b"
),
"network-access",
HIGH,
),
(
re.compile(
r"\bmarshal\s*\.\s*(loads|load)\b"
r"|\bcompile\s*\([^)]*['\"]exec['\"]\s*\)"
r"|\b__import__\s*\(|\bgetattr\s*\(\s*__builtins__"
r"|\b(?:b64decode|decodebytes)\s*\(.*(?:b64decode|decodebytes)\s*\(",
re.DOTALL,
),
"obfuscation",
HIGH,
),
(
re.compile(
r"-----BEGIN\s+(?:RSA\s+)?(?:PUBLIC|PRIVATE|ENCRYPTED|EC|DSA|OPENSSH)\s+KEY-----"
r"|\bMII[A-Za-z0-9+/]{20,}",
re.DOTALL,
),
"embedded-key-material",
HIGH,
),
(
re.compile(
r"\bos\.environ\s*\.\s*copy\s*\(|\bdict\s*\(\s*os\.environ\s*\)"
r"|\bjson\.dumps\s*\(\s*(?:dict\s*\(\s*)?os\.environ",
re.IGNORECASE,
),
"environment-harvest",
HIGH,
),
(
re.compile(
r"\bbase64\s*\.\s*(b64decode|decodebytes|b32decode|b16decode)\b"
r"|\bcodecs\s*\.\s*decode\b"
),
"base64/encoding-decode",
MEDIUM,
),
(re.compile(r"[A-Za-z0-9+/=]{200,}"), "large-base64-blob", MEDIUM),
)
@dataclass
class Finding:
severity: str
filename: str
check: str
evidence: str = ""
# 1-based match line + surrounding code window (see _attach_location). For the UI;
# None/[] when unlocatable.
line: Optional[int] = None
snippet: list = field(default_factory = list)
@dataclass
class ScanResult:
findings: list[Finding] = field(default_factory = list)
fingerprint: str = ""
@property
def max_severity(self) -> Optional[str]:
if not self.findings:
return None
return min((f.severity for f in self.findings), key = lambda s: _SEVERITY_ORDER[s])
@property
def clean(self) -> bool:
return not self.findings
def summary(self) -> str:
if self.clean:
return "no suspicious patterns found"
by = {}
for f in self.findings:
by.setdefault(f.severity, set()).add(f.check)
parts = []
for sev in (CRITICAL, HIGH, MEDIUM):
if sev in by:
parts.append(f"{sev}: {', '.join(sorted(by[sev]))}")
return "; ".join(parts)
def findings_payload(self) -> list[dict]:
"""Structured findings for the UI: one record per match, with line + snippet."""
return [
{
"severity": f.severity,
"file": f.filename,
"check": f.check,
"evidence": f.evidence,
"line": f.line,
"snippet": f.snippet,
}
for f in self.findings
]
# Canonical scanner: import scripts/scan_packages.py by file path (scripts/ is not an
# importable package from the backend root). It imports only stdlib at module level and
# guards its CLI under __main__, so importing it is side-effect-free.
_CANON_SENTINEL = object()
_canon_cache = _CANON_SENTINEL
def _load_canonical_scanner():
"""Return the ``scripts/scan_packages.py`` module, or None if unavailable."""
global _canon_cache
if _canon_cache is not _CANON_SENTINEL:
return _canon_cache
module = None
# Walk up from this file to a repo root that contains scripts/scan_packages.py.
here = pathlib.Path(__file__).resolve()
for parent in here.parents:
candidate = parent / "scripts" / "scan_packages.py"
if candidate.is_file():
try:
spec = importlib.util.spec_from_file_location("unsloth_scan_packages", candidate)
mod = importlib.util.module_from_spec(spec)
sys.modules.setdefault("unsloth_scan_packages", mod)
spec.loader.exec_module(mod) # type: ignore[union-attr]
if hasattr(mod, "check_py_file"):
module = mod
except Exception as exc: # pragma: no cover - defensive
logger.warning("Could not load canonical scan_packages.py: %s", exc)
break
if module is None:
logger.warning(
"scripts/scan_packages.py not found; remote-code scan using the "
"vendored fallback patterns."
)
_canon_cache = module
return module
# Model-context-strict patterns. The canonical scanner only flags bare
# ``subprocess``/``eval`` in combinations (common in package build scripts), but a
# model's modeling_*.py never legitimately shells out, so the gate flags them alone
# (e.g. a bare ``subprocess.Popen`` in a config ``__init__``).
_MODEL_STRICT_PATTERNS: tuple[tuple[re.Pattern, str, str], ...] = (
(
re.compile(
r"\bsubprocess\s*\.\s*(Popen|call|run|check_call|check_output)\b"
r"|\bos\s*\.\s*(system|popen|exec[lv]p?e?)\b"
),
"subprocess/os-exec (model code)",
HIGH,
),
# Bare exec()/eval(); excludes attribute calls like torch ``module.eval()``.
(re.compile(r"(?<![\w.])(?:exec|eval)\s*\("), "exec/eval (model code)", HIGH),
)
# Lines of context shown on each side of a flagged line in the dialog snippet.
_SNIPPET_CONTEXT = 3
# Cap a rendered line so a minified/one-line file can't bloat the payload.
_SNIPPET_MAX_LINE = 240
# Canonical evidence is formatted as "L<n>: <source line>" by _extract_evidence.
_EVIDENCE_LINE_RE = re.compile(r"^L(\d+):")
def _snippet_rows(
content: str,
line: int,
col: Optional[int] = None,
match_len: int = 0,
) -> list[dict]:
"""A `±_SNIPPET_CONTEXT`-line window around `line` (1-based). Rows are
{number, text, is_match}; the matched row adds match_start/match_end for an inline
highlight when a precise column span is known."""
lines = content.splitlines()
if not lines or line < 1:
return []
line = min(line, len(lines))
lo = max(1, line - _SNIPPET_CONTEXT)
hi = min(len(lines), line + _SNIPPET_CONTEXT)
rows: list[dict] = []
for n in range(lo, hi + 1):
text = lines[n - 1]
clipped = len(text) > _SNIPPET_MAX_LINE
if clipped:
text = text[:_SNIPPET_MAX_LINE] + " ..."
row = {"number": n, "text": text, "is_match": n == line}
if n == line and col is not None and match_len > 0 and not clipped:
row["match_start"] = col
row["match_end"] = min(col + match_len, len(text))
rows.append(row)
return rows
def _attach_location(
content: str,
finding: Finding,
match: "Optional[re.Match]" = None,
) -> None:
"""Populate finding.line + finding.snippet. A regex match gives a precise
line+column; canonical findings are located via the `L<n>:` prefix in their evidence."""
if match is not None:
before = content[: match.start()]
line = before.count("\n") + 1
col = match.start() - (before.rfind("\n") + 1)
finding.line = line
finding.snippet = _snippet_rows(content, line, col, match.end() - match.start())
return
first = finding.evidence.splitlines()[0] if finding.evidence else ""
tag = _EVIDENCE_LINE_RE.match(first)
if tag:
line = int(tag.group(1))
finding.line = line
finding.snippet = _snippet_rows(content, line)
def _scan_content(content: str, filename: str) -> list[Finding]:
findings: list[Finding] = []
canon = _load_canonical_scanner()
if canon is not None:
# Canonical Finding is (severity, package, filename, check, evidence); adapt to
# the gate's (severity, filename, check, evidence).
for f in canon.check_py_file(content, filename, ""):
finding = Finding(f.severity, f.filename, f.check, (f.evidence or "")[:120])
_attach_location(content, finding)
findings.append(finding)
else:
for pat, check, sev in _FALLBACK_PATTERNS:
m = pat.search(content)
if m:
finding = Finding(sev, filename, check, m.group(0)[:120])
_attach_location(content, finding, m)
findings.append(finding)
# Augment with the model-context-strict patterns the package scanner omits.
have = {f.check for f in findings}
for pat, check, sev in _MODEL_STRICT_PATTERNS:
if check in have:
continue
m = pat.search(content)
if m:
finding = Finding(sev, filename, check, m.group(0)[:120])
_attach_location(content, finding, m)
findings.append(finding)
return findings
def scan_remote_code_files(files: dict[str, str]) -> ScanResult:
"""Scan a mapping of {filename: content} and return aggregated findings."""
result = ScanResult(fingerprint = remote_code_fingerprint(files))
for name, content in files.items():
if not name.endswith(".py"):
continue
result.findings.extend(_scan_content(content or "", name))
return result
def remote_code_fingerprint(files: dict[str, str]) -> str:
"""Stable sha256 over the (sorted) file contents, for pinning consent."""
h = hashlib.sha256()
for name in sorted(files):
h.update(name.encode("utf-8"))
h.update(b"\0")
h.update((files[name] or "").encode("utf-8"))
h.update(b"\0")
return h.hexdigest()
def repo_remote_code_files(model_name: str, hf_token: Optional[str] = None) -> dict[str, str]:
"""Download a repo's executable ``.py`` (auto_map targets + modeling/config).
Returns {filename: content}. An EMPTY dict means the repo ships no executable ``.py``
(trust_remote_code is a no-op). Raises ``RemoteCodeUnscannable`` when code is present
but cannot be fully fetched/listed (offline/gated/404/listing failure), so the caller
fails closed rather than fingerprint a partial view of code transformers would run in
full. The empty-vs-raise split lets the gate allow a code-free repo while still
blocking unscannable code.
"""
import json
from pathlib import Path
files: dict[str, str] = {}
try:
from utils.paths import is_local_path, normalize_path
if is_local_path(model_name):
root = Path(normalize_path(model_name)).expanduser()
# Walk ALL .py, not just the auto_map entry's static import closure. This is
# DELIBERATE (see the remote-branch note): the entry can reach a sibling via
# an absolute import, importlib, or exec, which a relative-import closure
# misses, so closure-only scanning is a real bypass. Broad scan never
# under-scans; the cost is a benign script can over-block, the safe direction
# for an RCE gate (HIGH stays approvable; only CRITICAL hard-blocks).
for p in root.rglob("*.py"):
if p.is_file():
files[str(p.relative_to(root))] = p.read_text(errors = "replace")
# A local config can still point auto_map at an EXTERNAL Hub repo
# (owner/name--module.Class) that executes on load, so fetch it. Every config
# that can declare auto_map is checked, so a custom processor's external code
# is not missed.
ext_refs = set()
for name in REMOTE_CODE_CONFIG_FILES:
p = root / name
if p.is_file():
try:
ext_refs |= _auto_map_refs(json.loads(p.read_text()))
except Exception:
pass
if not _add_external_refs(files, ext_refs, hf_token, model_name):
raise RemoteCodeUnscannable(f"{model_name}: external auto_map code unreachable")
return files
from huggingface_hub import hf_hub_download, list_repo_files
from huggingface_hub.utils import EntryNotFoundError
# Collect auto_map refs from EVERY config that can declare one. A 404
# (EntryNotFoundError) means the config is absent -> skip; any other failure is
# transient/auth and could hide an auto_map, so fail closed (unscannable).
refs = set()
for cfg_name in REMOTE_CODE_CONFIG_FILES:
try:
cfg_path = hf_hub_download(model_name, cfg_name, token = hf_token)
except EntryNotFoundError:
continue
except Exception as exc:
raise RemoteCodeUnscannable(
f"{model_name}: config {cfg_name} could not be fetched ({exc})"
) from exc
try:
refs |= _auto_map_refs(json.loads(Path(cfg_path).read_text()))
except Exception:
pass
own_refs = {fn for repo, fn in refs if repo is None}
# The full file list catches helper .py the auto_map code imports but does not
# name. If we cannot list the repo, an imported module could be missed and the
# fingerprint cover less than transformers runs, so fail closed (unscannable).
try:
repo_files = list_repo_files(model_name, token = hf_token)
except Exception as exc:
raise RemoteCodeUnscannable(f"{model_name}: could not list repo files ({exc})") from exc
repo_file_set = set(repo_files)
# Scan every present .py PLUS own-repo auto_map targets that ACTUALLY EXIST in
# this revision. Scanning EVERY .py (not just the closure) is DELIBERATE: the
# entry can reach a sibling via absolute import / importlib / exec, which a
# relative-import closure misses, so closure-only scanning is a real bypass.
# Broad scan never under-scans; the cost is a benign script can over-block, the
# safe direction for an RCE gate (HIGH approvable; only CRITICAL hard-blocks). An
# auto_map target absent from the listing is a STALE ref (an older config naming a
# since-removed file, e.g. unsloth/PaddleOCR-VL names processing_ppocrvl.py but
# ships processing_paddleocr_vl.py). transformers cannot execute an absent file,
# so drop the stale ref rather than fail closed; present .py are still fully
# scanned. This also absorbs a mis-derived dotted name (sub.mod.py vs sub/mod.py):
# the bad name drops as stale while the real present file is scanned.
present_py = {f for f in repo_files if f.endswith(".py")}
stale_refs = own_refs - repo_file_set
for fn in sorted(stale_refs):
logger.info(
"repo_remote_code_files(%s): ignoring stale own-repo auto_map target "
"%s (absent from the repo listing; it cannot execute)",
model_name,
fn,
)
wanted = present_py | (own_refs & repo_file_set)
for fn in sorted(wanted):
try:
fp = hf_hub_download(model_name, fn, token = hf_token)
except Exception as exc:
# A .py CONFIRMED PRESENT could not be fetched. A partial set would
# fingerprint "clean" while transformers later runs this file, so fail
# closed. (Stale/absent refs were dropped above, so this only fires on a
# present-file fetch failure.)
raise RemoteCodeUnscannable(
f"{model_name}: present file {fn} could not be fetched ({exc})"
) from exc
files[fn] = Path(fp).read_text(errors = "replace")
# Code referenced from another repo executes too: scan it or fail closed.
if not _add_external_refs(files, refs, hf_token, model_name):
raise RemoteCodeUnscannable(f"{model_name}: external auto_map code unreachable")
except RemoteCodeUnscannable:
logger.warning("repo_remote_code_files(%s): unscannable; failing closed", model_name)
raise
except Exception as exc:
# An unexpected error mid-scan means we could not complete it -> unscannable.
raise RemoteCodeUnscannable(f"{model_name}: scan failed ({exc})") from exc
# An empty dict here means the listing succeeded and the repo ships no executable .py
# (nor fetchable external refs) -> trust_remote_code is a no-op for the caller.
return files
def _iter_auto_map_strings(value):
"""Yield every string class-ref inside one ``auto_map`` value.
A value is a bare string (``"modeling_x.Cls"``) or a list/tuple (transformers encodes
a tokenizer as ``"AutoTokenizer": [slow, fast]``, possibly nested or with nulls).
Flatten all forms so external tokenizer code in the list shape is scanned.
"""
if isinstance(value, str):
yield value
elif isinstance(value, (list, tuple, set)):
for item in value:
yield from _iter_auto_map_strings(item)
elif isinstance(value, dict):
for item in value.values():
yield from _iter_auto_map_strings(item)
def _auto_map_refs(cfg: dict) -> set:
"""``(repo, filename)`` pairs referenced by config auto_map.
``repo`` is ``None`` for own-repo code. An external ``owner/name--module.Class`` ref
(transformers' cross-repo form) yields ``("owner/name", "module.py")`` so cross-repo
code is scanned + fingerprinted too.
"""
out = set()
am = cfg.get("auto_map") or {}
if isinstance(am, dict):
for value in am.values():
# A value may be a string OR a [slow, fast] list (tokenizers); cover both.
for ref in _iter_auto_map_strings(value):
# ref like "modeling_deepseekocr.Cls" or "owner/name--modeling.Cls"
if "." not in ref:
continue
module = ref.rsplit(".", 1)[0] # drop trailing .ClassName
if "--" in module:
repo, mod = module.split("--", 1)
out.add((repo or None, mod + ".py"))
else:
out.add((None, module + ".py"))
return out
def _auto_map_py(cfg: dict) -> set[str]:
"""Own-repo ``.py`` filenames referenced by auto_map (external refs excluded)."""
return {fn for repo, fn in _auto_map_refs(cfg) if repo is None}
def external_auto_map_repos(model_name: str, hf_token: Optional[str] = None) -> set:
"""External Hub repos referenced by any of this model's auto_map configs.
The ``owner/name`` repos ``_add_external_refs`` downloads. The scan route uses this so
declining consent purges them too, not leaving untrusted external code cached.
Best-effort, config/metadata-only: returns whatever can be read, never raises.
"""
repos: set = set()
try:
import json
from pathlib import Path
from utils.paths import is_local_path, normalize_path
if is_local_path(model_name):
root = Path(normalize_path(model_name)).expanduser()
for cfg_name in REMOTE_CODE_CONFIG_FILES:
p = root / cfg_name
if not p.is_file():
continue
try:
refs = _auto_map_refs(json.loads(p.read_text()))
except Exception:
continue
repos.update(repo for repo, _fn in refs if repo)
return repos
from huggingface_hub import hf_hub_download
from huggingface_hub.utils import EntryNotFoundError
for cfg_name in REMOTE_CODE_CONFIG_FILES:
try:
cfg_path = hf_hub_download(model_name, cfg_name, token = hf_token)
except EntryNotFoundError:
continue
except Exception:
continue
try:
refs = _auto_map_refs(json.loads(Path(cfg_path).read_text()))
except Exception:
continue
repos.update(repo for repo, _fn in refs if repo)
except Exception:
return repos
return repos
def _add_external_refs(files: dict, refs, hf_token, model_name: str) -> bool:
"""Download external-repo auto_map code into ``files`` (keyed ``repo--file``).
transformers fetches the entry file AND its relative imports from the same external
repo, so scanning only the entry would miss code in a ``helper.py`` it imports. Mirror
the own-repo path: enumerate each external repo's ``.py`` and scan the whole set (plus
the referenced entry files). Returns False if any external repo cannot be listed or a
file cannot be fetched, so the caller fails closed.
"""
from pathlib import Path
from huggingface_hub import hf_hub_download, list_repo_files
# Group the explicit entry refs by external repo.
entries: dict = {}
for repo, fn in refs:
if repo is None:
continue
entries.setdefault(repo, set()).add(fn)
for repo, entry_files in entries.items():
try:
repo_files = list_repo_files(repo, token = hf_token)
except Exception as exc:
logger.warning(
"repo_remote_code_files(%s): external repo %s unlistable (%s); failing closed",
model_name,
repo,
exc,
)
return False
# The loader's executable closure = every present .py plus any referenced entry
# file. With a REAL (non-empty) listing, present_py covers the code, so an entry
# ref absent from it is stale/mis-derived and is dropped rather than failing
# closed (like the own-repo path). With an EMPTY listing we cannot prove the ref
# stale, so keep fetching it and fail closed if unreachable; never under-scan. A
# PRESENT file that cannot be fetched still fails closed below.
repo_file_set = set(repo_files)
present_py = {f for f in repo_files if f.endswith(".py")}
if repo_file_set:
for fn in sorted(set(entry_files) - repo_file_set):
logger.info(
"repo_remote_code_files(%s): ignoring stale external auto_map target "
"%s:%s (absent from the repo listing; it cannot execute)",
model_name,
repo,
fn,
)
wanted = present_py | (set(entry_files) & repo_file_set)
else:
wanted = present_py | set(entry_files)
for fn in sorted(wanted):
try:
fp = hf_hub_download(repo, fn, token = hf_token)
except Exception as exc:
logger.warning(
"repo_remote_code_files(%s): external %s:%s unscannable (%s)",
model_name,
repo,
fn,
exc,
)
return False
files[f"{repo}--{fn}"] = Path(fp).read_text(errors = "replace")
return True

120
studio/backend/utils/security/trusted_org.py Normal file
View file

@ -0,0 +1,120 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Trusted-org checks for the ``trust_remote_code`` auto-enable paths.
A bare ``name.startswith("unsloth/")`` is spoofable by a local path like
``./unsloth/evil``. ``is_trusted_org_repo`` rejects local paths, requires an
``org/repo`` under a trusted org, and (online) confirms via the Hub. Fails CLOSED
on any uncertainty and never raises; a False just means "do not auto-enable".
"""
from __future__ import annotations
import hashlib
import os
from typing import Optional
from loggers import get_logger
from utils.paths import is_local_path
logger = get_logger(__name__)
# Orgs we auto-enable remote code for.
TRUSTED_ORGS: frozenset[str] = frozenset({"unsloth", "nvidia"})
# Keyed on (name, verify_remote, token) so an unauthenticated failure can't poison
# a later authenticated lookup; token is hashed, never stored raw.
_verdict_cache: dict[tuple[str, bool, str], bool] = {}
def _token_key(hf_token: Optional[str]) -> str:
"""Non-reversible cache discriminator; empty when no token, never the raw token."""
if not hf_token:
return ""
return hashlib.sha256(hf_token.encode("utf-8")).hexdigest()[:12]
def _env_offline() -> bool:
return os.environ.get("HF_HUB_OFFLINE", "").lower() in ("1", "true", "yes") or os.environ.get(
"TRANSFORMERS_OFFLINE", ""
).lower() in ("1", "true", "yes")
def is_trusted_org_repo(
name: str,
hf_token: Optional[str] = None,
*,
verify_remote: bool = True,
) -> bool:
"""True only if *name* is a genuine HF repo under a trusted org. Fails closed
(local paths, malformed names, untrusted namespaces, Hub errors); never raises.
Offline trusts the namespace shape, since the Hub is unreachable by design.
"""
if not name or not isinstance(name, str):
return False
cache_key = (name, verify_remote, _token_key(hf_token))
if cache_key in _verdict_cache:
return _verdict_cache[cache_key]
verdict = _evaluate(name, hf_token, verify_remote)
_verdict_cache[cache_key] = verdict
return verdict
def _namespace(name: str) -> Optional[str]:
"""Lowercased org of an ``org/repo`` id, else None."""
parts = name.split("/")
if len(parts) != 2 or not parts[0] or not parts[1]:
return None
return parts[0].lower()
def _evaluate(name: str, hf_token: Optional[str], verify_remote: bool) -> bool:
# Local paths are never a trusted remote repo (the spoof this guards against).
try:
if is_local_path(name):
logger.debug("is_trusted_org_repo(%s): local path -> not trusted", name)
return False
except Exception:
return False
ns = _namespace(name)
if ns is None or ns not in TRUSTED_ORGS:
return False
# Offline: trust the shape (Hub intentionally unreachable).
if not verify_remote or _env_offline():
return True
# Online: confirm the id resolves to a trusted-org repo.
try:
from huggingface_hub import HfApi
info = HfApi().model_info(name, token = hf_token)
resolved_id = getattr(info, "id", None) or name
resolved_ns = _namespace(resolved_id)
author = getattr(info, "author", None)
if resolved_ns in TRUSTED_ORGS:
return True
if author and str(author).lower() in TRUSTED_ORGS:
return True
logger.warning(
"is_trusted_org_repo(%s): resolved id %r not under a trusted org",
name,
resolved_id,
)
return False
except Exception as exc: # network/404/auth -> fail closed
logger.warning(
"is_trusted_org_repo(%s): Hub verification failed (%s) -> not trusted",
name,
exc,
)
return False
def clear_cache() -> None:
"""Test helper: drop the memoized verdicts."""
_verdict_cache.clear()

36
studio/backend/utils/transformers_version.py
View file

@ -111,8 +111,8 @@ _TRANSFORMERS_5_TOKENIZER_CLASSES: set[str] = {
# Cache for dynamic tokenizer_config.json lookups (avoids repeated fetches).
_tokenizer_class_cache: dict[str, bool] = {}
# Cache for dynamic config.json lookups (architecture/model_type checks).
_config_json_cache: dict[str, dict | None] = {}
# config.json cache keyed on (model_name, token-hash) so authed/unauthed reads stay separate.
_config_json_cache: dict[tuple[str, str | None], dict | None] = {}
_config_needs_510_cache: dict[str, bool] = {}
_config_needs_550_cache: dict[str, bool] = {}
@ -328,39 +328,51 @@ def _check_tokenizer_config_needs_v5(model_name: str) -> bool:
return False
def _load_config_json(model_name: str) -> dict | None:
"""Return parsed ``config.json`` for *model_name*, checking local files first."""
if model_name in _config_json_cache:
return _config_json_cache[model_name]
def _load_config_json(model_name: str, hf_token: str | None = None) -> dict | None:
"""Return parsed ``config.json`` for *model_name*, checking local files first.
``hf_token`` authenticates the raw fetch so gated/private repos resolve. The
cache is keyed on the token so an unauthenticated miss never poisons a later
authenticated read.
"""
import hashlib
tok = hashlib.sha256(hf_token.encode()).hexdigest()[:16] if hf_token else None
cache_key = (model_name, tok)
if cache_key in _config_json_cache:
return _config_json_cache[cache_key]
local_cfg = Path(model_name) / "config.json"
if local_cfg.is_file():
try:
with open(local_cfg) as f:
cfg = json.load(f)
_config_json_cache[model_name] = cfg
_config_json_cache[cache_key] = cfg
return cfg
except Exception as exc:
logger.debug("Could not read %s: %s", local_cfg, exc)
_config_json_cache[model_name] = None
_config_json_cache[cache_key] = None
return None
if _env_offline():
_config_json_cache[model_name] = None
_config_json_cache[cache_key] = None
return None
import urllib.request
url = f"https://huggingface.co/{model_name}/raw/main/config.json"
headers = {"User-Agent": "unsloth-studio"}
if hf_token:
headers["Authorization"] = f"Bearer {hf_token}"
try:
req = urllib.request.Request(url, headers = {"User-Agent": "unsloth-studio"})
req = urllib.request.Request(url, headers = headers)
with urllib.request.urlopen(req, timeout = 10) as resp:
cfg = json.loads(resp.read().decode())
_config_json_cache[model_name] = cfg
_config_json_cache[cache_key] = cfg
return cfg
except Exception as exc:
logger.debug("Could not fetch config.json for '%s': %s", model_name, exc)
_config_json_cache[model_name] = None
_config_json_cache[cache_key] = None
return None

2
studio/frontend/src/app/routes/__root.tsx
View file

@ -6,6 +6,7 @@ import { Navbar } from "@/components/navbar";
import { fetchDeviceType, usePlatformStore } from "@/config/env";
import { SidebarInset, SidebarProvider } from "@/components/ui/sidebar";
import { SettingsDialog, useSettingsDialogStore } from "@/features/settings";
import { RemoteCodeConsentDialog } from "@/features/security";
import { clearNewChatDraft, useChatRuntimeStore } from "@/features/chat";
import { useTrainingUnloadGuard } from "@/features/training";
import { useSidebarPin } from "@/hooks/use-sidebar-pin";
@ -142,6 +143,7 @@ function RootLayout() {
return (
<AppProvider>
<SettingsDialog />
<RemoteCodeConsentDialog />
{hideNavbar ? (
<main className="flex-1">
<Suspense fallback={<RouteFallback />}>

11
studio/frontend/src/features/chat/api/chat-adapter.ts
View file

@ -1215,7 +1215,12 @@ async function autoLoadSmallestModel(): Promise<{
load_in_4bit: true,
trust_remote_code: trustRemoteCode,
});
if (validation.requires_trust_remote_code && !trustRemoteCode) {
// Background auto-load never runs a repo's custom code or loads Hub-flagged unsafe
// files on its own; both are deferred to the explicit consent dialog instead.
if (
validation.requires_trust_remote_code ||
validation.requires_security_review
) {
blockedByTrustRemoteCode = true;
return false;
}
@ -1567,11 +1572,11 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
if (!loaded) {
toast.error(
blockedByTrustRemoteCode
? "Enable custom code to auto-load this model"
? "This model needs custom code approval"
: "No model loaded",
{
description: blockedByTrustRemoteCode
? 'Turn on "Enable custom code" in Chat Settings, or pick another model in the top bar.'
? "Select it from the top bar to review and approve its custom code, or pick another model."
: "Pick a model in the top bar, then retry.",
},
);

41
studio/frontend/src/features/chat/chat-settings-sheet.tsx
View file

@ -530,9 +530,6 @@ export function ChatSettingsPanel({
const loadedSpecDraftNMax = useChatRuntimeStore(
(s) => s.loadedSpecDraftNMax,
);
const modelRequiresTrustRemoteCode = useChatRuntimeStore(
(s) => s.modelRequiresTrustRemoteCode,
);
const currentCheckpoint = params.checkpoint;
const ggufContextLength = useChatRuntimeStore((s) => s.ggufContextLength);
const ggufMaxContextLength = useChatRuntimeStore(
@ -635,10 +632,6 @@ export function ChatSettingsPanel({
[activePreset, hasUnsavedPresetChanges, presetNameInput, presets],
);
const systemPromptEditorDirty = systemPromptDraft !== params.systemPrompt;
const trustRemoteCodeMissing =
Boolean(currentCheckpoint) &&
modelRequiresTrustRemoteCode &&
!(params.trustRemoteCode ?? false);
const showPromptCacheTtlControl = Boolean(
activeExternalProvider &&
supportsProviderPromptCacheTtl(activeExternalProvider.providerType),
@ -1094,38 +1087,8 @@ export function ChatSettingsPanel({
</div>
</>
)}
{!isGguf && params.checkpoint && (
<>
<div className="flex items-center justify-between gap-3">
<div className="flex min-w-0 items-center gap-1.5">
<span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
Enable custom code
</span>
<InfoHint>
Run custom Python from the model repo (e.g. Nemotron).
Only enable for trusted sources.
</InfoHint>
</div>
<Switch
className="panel-switch shrink-0"
checked={params.trustRemoteCode ?? false}
onCheckedChange={set("trustRemoteCode")}
/>
</div>
{trustRemoteCodeMissing && (
<Alert className="rounded-[14px] border-amber-200/70 bg-amber-50/70 px-3 py-2 text-amber-950 dark:border-amber-900/70 dark:bg-amber-950/35 dark:text-amber-100">
<AlertTitle className="text-[12px] font-medium">
Keep custom code enabled for this model
</AlertTitle>
<AlertDescription className="text-[11.5px] leading-[1.45] text-amber-800 dark:text-amber-200">
This model requires custom code to load. You can edit the
toggle, but loading will stay blocked until it is turned
back on.
</AlertDescription>
</Alert>
)}
</>
)}
{/* No persistent "enable custom code" toggle: it is consented per model
via the load-time review dialog. */}
{/* Apply/Reset belongs to the model-reload settings above (context
length, KV cache, speculative decoding). Render it here, before
the Chat Template row, so it never reads as attached to Chat

42
studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts
View file

@ -3,6 +3,7 @@
import { createElement, useCallback, useRef, useState } from "react";
import { toast } from "@/lib/toast";
import { confirmRemoteCodeIfNeeded } from "@/features/security";
import { consumeNativePathToken } from "@/features/native-intents/api";
import {
notifyNative,
@ -68,6 +69,16 @@ export type SelectedModelInput = {
keepSpeculative?: boolean;
};
// Approved fingerprints by checkpoint, so a rollback after a failed switch can resend
// the pinned approval the worker requires instead of being blocked.
const approvedRemoteCodeFingerprints = new Map<string, string>();
function rememberApprovedRemoteCode(
checkpoint: string,
fingerprint: string | null,
): void {
if (fingerprint) approvedRemoteCodeFingerprints.set(checkpoint, fingerprint);
}
const MODEL_LOAD_TOAST_CLASSNAMES = {
toast: "chat-model-load-toast items-center gap-2.5",
content: "gap-0.5 flex-1 min-w-0",
@ -221,7 +232,7 @@ function toLoraSummary(lora: {
}
function getTrustRemoteCodeRequiredMessage(modelName: string): string {
return `${modelName} needs custom code enabled to load. Turn on "Enable custom code" in Chat Settings, then try again.`;
return `${modelName} was not loaded because its custom code was not approved. Load it again to review the code and approve it.`;
}
export function useChatModelRuntime() {
@ -481,7 +492,8 @@ export function useChatModelRuntime() {
const currentCheckpoint =
useChatRuntimeStore.getState().params.checkpoint;
const stateBeforeUnload = useChatRuntimeStore.getState();
const trustRemoteCode = stateBeforeUnload.params.trustRemoteCode ?? false;
let trustRemoteCode = stateBeforeUnload.params.trustRemoteCode ?? false;
let approvedRemoteCodeFingerprint: string | null = null;
const maxSeqLength = stateBeforeUnload.params.maxSeqLength;
const previousIsGguf =
previousModel?.isGguf === true
@ -510,8 +522,25 @@ export function useChatModelRuntime() {
is_lora: isLora,
gguf_variant: ggufVariant ?? null,
});
if (validation.requires_trust_remote_code && !trustRemoteCode) {
throw new Error(getTrustRemoteCodeRequiredMessage(displayName));
// Open the consent dialog when the model needs custom-code consent or has a
// flagged unsafe file. Fires even when trustRemoteCode is preset on, since the
// worker requires a matching fingerprint that only the dialog produces.
if (
validation.requires_trust_remote_code
|| validation.requires_security_review
) {
const approved = await confirmRemoteCodeIfNeeded({
modelName: modelId,
hfToken,
requiresTrustRemoteCode: true,
onApprove: (fp) => {
trustRemoteCode = true;
approvedRemoteCodeFingerprint = fp;
},
});
if (!approved) {
throw new Error(getTrustRemoteCodeRequiredMessage(displayName));
}
}
if (abortCtrl.signal.aborted) throw new Error("Cancelled");
const loadNativePathLease = nativePathToken
@ -575,6 +604,7 @@ export function useChatModelRuntime() {
is_lora: isLora,
gguf_variant: ggufVariant ?? null,
trust_remote_code: trustRemoteCode,
approved_remote_code_fingerprint: approvedRemoteCodeFingerprint,
chat_template_override: effectiveChatTemplateOverride,
cache_type_kv: kvCacheDtype,
speculative_type: speculativeType,
@ -646,6 +676,7 @@ export function useChatModelRuntime() {
: reloadingSameModel && supportsReasoning
? stateBeforeUnload.reasoningEnabled
: reasoningDefault;
rememberApprovedRemoteCode(modelId, approvedRemoteCodeFingerprint);
useChatRuntimeStore.setState({
ggufContextLength: nativeCtx,
ggufMaxContextLength,
@ -746,6 +777,9 @@ export function useChatModelRuntime() {
gguf_variant: previousVariant,
trust_remote_code:
previousModelRequiresTrustRemoteCode || trustRemoteCode,
// Resend the previous model's pinned approval so restoring it is not re-blocked.
approved_remote_code_fingerprint:
approvedRemoteCodeFingerprints.get(previousCheckpoint) ?? null,
// Restore the previous model in the split mode it was running,
// not the default layer split.
tensor_parallel: stateBeforeUnload.loadedTensorParallel ?? false,

47
studio/frontend/src/features/chat/shared-composer.tsx
View file

@ -63,6 +63,7 @@ import { BypassPermissionsMenuItem } from "./bypass-permissions-menu-item";
import { KnowledgeBaseComposerButton } from "@/features/rag/components/knowledge-base-composer-button";
import { NewProjectDialog } from "./components/new-project-dialog";
import { useChatProjects } from "./hooks/use-chat-projects";
import { confirmRemoteCodeIfNeeded } from "@/features/security";
import { loadModel, validateModel } from "./api/chat-api";
import {
parseExternalModelId,
@ -932,24 +933,43 @@ export function SharedComposer({
sel: CompareModelSelection,
): Promise<string> {
const currentStore = useChatRuntimeStore.getState();
let loadTrustRemoteCode = trustRemoteCode;
let approvedRemoteCodeFingerprint: string | null = null;
const isAlreadyActive =
currentStore.params.checkpoint === sel.id &&
(currentStore.activeGgufVariant ?? null) ===
(sel.ggufVariant ?? null);
if (!isAlreadyActive) {
const validation = await validateModel({
model_path: sel.id,
hf_token: currentStore.hfToken || null,
max_seq_length: maxSeqLength,
load_in_4bit: true,
is_lora: sel.isLora,
gguf_variant: sel.ggufVariant ?? null,
trust_remote_code: trustRemoteCode,
chat_template_override: effectiveChatTemplateOverride,
// Already loaded (gate passed at first load): skip a redundant reload that would
// re-trigger the gate without the approval fingerprint and fail for HIGH custom code.
if (isAlreadyActive) {
return "ready";
}
const validation = await validateModel({
model_path: sel.id,
hf_token: currentStore.hfToken || null,
max_seq_length: maxSeqLength,
load_in_4bit: true,
is_lora: sel.isLora,
gguf_variant: sel.ggufVariant ?? null,
trust_remote_code: loadTrustRemoteCode,
chat_template_override: effectiveChatTemplateOverride,
});
if (
validation.requires_trust_remote_code ||
validation.requires_security_review
) {
const approved = await confirmRemoteCodeIfNeeded({
modelName: sel.id,
hfToken: currentStore.hfToken || null,
requiresTrustRemoteCode: true,
onApprove: (fp) => {
loadTrustRemoteCode = true;
approvedRemoteCodeFingerprint = fp;
},
});
if (validation.requires_trust_remote_code && !trustRemoteCode) {
if (!approved) {
throw new Error(
`${modelDisplayName(sel.id)} needs custom code enabled to load. Turn on "Enable custom code" in Chat Settings, then try again.`,
`${modelDisplayName(sel.id)} needs custom code approval to load.`,
);
}
}
@ -960,7 +980,8 @@ export function SharedComposer({
load_in_4bit: true,
is_lora: sel.isLora,
gguf_variant: sel.ggufVariant ?? null,
trust_remote_code: trustRemoteCode,
trust_remote_code: loadTrustRemoteCode,
approved_remote_code_fingerprint: approvedRemoteCodeFingerprint,
chat_template_override: effectiveChatTemplateOverride,
speculative_type: specSettings.speculativeType,
spec_draft_n_max: specSettings.specDraftNMax,

4
studio/frontend/src/features/chat/types/api.ts
View file

@ -41,6 +41,8 @@ export interface LoadModelRequest {
gguf_variant?: string | null;
/** Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust. */
trust_remote_code?: boolean;
/** sha256 fingerprint pinning user approval of this exact custom-code version. */
approved_remote_code_fingerprint?: string | null;
chat_template_override?: string | null;
cache_type_kv?: string | null;
/**
@ -72,6 +74,8 @@ export interface ValidateModelResponse {
is_lora?: boolean;
is_vision?: boolean;
requires_trust_remote_code?: boolean;
// HF flagged unsafe files, so the load is hard-blocked pending dialog review.
requires_security_review?: boolean;
/** Native context length from the local GGUF header; null until downloaded. */
context_length?: number | null;
}

6
studio/frontend/src/features/chat/utils/chat-settings-storage.ts
View file

@ -140,11 +140,7 @@ function sanitizeInferenceParams(
if (typeof value.systemPrompt === "string") {
params.systemPrompt = value.systemPrompt;
}
if (typeof value.trustRemoteCode === "boolean") {
params.trustRemoteCode = value.trustRemoteCode;
}
// Mirror trustRemoteCode handling so the toggle survives reload
// and the /api/chat/settings round-trip.
// trustRemoteCode is no longer persisted: custom code is consented per model via the dialog.
if (typeof value.fastMode === "boolean") {
params.fastMode = value.fastMode;
}

4
studio/frontend/src/features/export/api/export-api.ts
View file

@ -54,6 +54,10 @@ export async function loadCheckpoint(params: {
load_in_4bit?: boolean;
/** Allow loading models with custom code. Only enable for checkpoints you trust. */
trust_remote_code?: boolean;
/** sha256 fingerprint pinning user approval of this exact custom-code version. */
approved_remote_code_fingerprint?: string | null;
/** HF token so the worker scans/loads gated checkpoints and base models with the same auth as preflight. */
hf_token?: string | null;
}): Promise<ExportOperationResponse> {
const response = await authFetch("/api/export/load-checkpoint", {
method: "POST",

71
studio/frontend/src/features/export/export-page.tsx
View file

@ -25,7 +25,6 @@ import {
SelectValue,
} from "@/components/ui/select";
import { Separator } from "@/components/ui/separator";
import { Switch } from "@/components/ui/switch";
import { Spinner } from "@/components/ui/spinner";
import {
Tooltip,
@ -37,6 +36,7 @@ import {
type LocalModelInfo,
useTrainingConfigStore,
} from "@/features/training";
import { confirmRemoteCodeIfNeeded } from "@/features/security";
import {
useDebouncedValue,
useHfModelSearch,
@ -129,8 +129,6 @@ export function ExportPage() {
"checkpoint",
);
const [modelSource, setModelSource] = useState<"hf" | "local">("hf");
const [hfExportTrustRemoteCode, setHfExportTrustRemoteCode] =
useState(true);
const [modelInput, setModelInput] = useState("");
const [selectedSourceModel, setSelectedSourceModel] = useState<string | null>(
null,
@ -490,13 +488,34 @@ export function ExportPage() {
// 1. Load model source
if (sourceMode === "checkpoint") {
if (!checkpointPath) return;
await loadCheckpoint({ checkpoint_path: checkpointPath });
await loadCheckpoint({
checkpoint_path: checkpointPath,
hf_token: hfToken || null,
});
} else {
// Consent gate for an HF source's custom (auto_map) code: the only way to enable
// trust_remote_code here. A local checkpoint the user exported is trusted by default.
let trustRemoteCode = modelSource !== "hf";
let approvedRemoteCodeFingerprint: string | null = null;
const remoteCodeOk = await confirmRemoteCodeIfNeeded({
modelName: source,
hfToken: hfToken || null,
// An HF source can need trust_remote_code via its YAML default with no auto_map
// to review; signal it so a YAML-only model does not export with it false.
requiresTrustRemoteCode: modelSource === "hf",
onApprove: (fingerprint) => {
trustRemoteCode = true;
approvedRemoteCodeFingerprint = fingerprint;
},
});
if (!remoteCodeOk) return;
await loadCheckpoint({
checkpoint_path: source,
load_in_4bit: false,
trust_remote_code:
modelSource === "hf" ? hfExportTrustRemoteCode : true,
trust_remote_code: trustRemoteCode,
approved_remote_code_fingerprint: approvedRemoteCodeFingerprint,
hf_token: hfToken || null,
});
}
@ -578,7 +597,6 @@ export function ExportPage() {
hfToken,
privateRepo,
modelSource,
hfExportTrustRemoteCode,
]);
// ---- Render ----
@ -868,43 +886,8 @@ export function ExportPage() {
</p>
)}
</div>
<div className="flex items-center gap-2">
<Switch
id="hf-export-trust-remote-code"
size="sm"
checked={hfExportTrustRemoteCode}
onCheckedChange={setHfExportTrustRemoteCode}
disabled={exporting}
/>
<label
htmlFor="hf-export-trust-remote-code"
className="cursor-pointer text-xs font-medium text-muted-foreground hover:text-foreground"
>
Trust remote code
</label>
<Tooltip>
<TooltipTrigger asChild={true}>
<button
type="button"
className="text-muted-foreground hover:text-foreground -m-1 inline-flex rounded p-1"
aria-label="About trust remote code"
>
<HugeiconsIcon
icon={InformationCircleIcon}
className="size-3.5"
/>
</button>
</TooltipTrigger>
<TooltipContent
side="top"
className="max-w-[260px] text-xs"
>
Loads custom Python from the repo if the model
needs it. Turn off if you do not trust the
source.
</TooltipContent>
</Tooltip>
</div>
{/* No persistent "trust remote code" toggle: custom code is
consented per model via the load-time review dialog. */}
<div className="flex flex-col gap-1.5">
<label className="text-xs font-medium text-muted-foreground">
Hugging Face Token (Optional)

111
studio/frontend/src/features/security/api/remote-code-api.ts Normal file
View file

@ -0,0 +1,111 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { authFetch } from "@/features/auth";
import { readFastApiError } from "@/lib/format-fastapi-error";
import type {
RemoteCodeFinding,
RemoteCodeScan,
RemoteCodeSeverity,
RemoteCodeSnippetRow,
UnsafeFile,
} from "../types";
interface SnippetRowResponse {
number: number;
text: string;
is_match: boolean;
match_start?: number;
match_end?: number;
}
interface RemoteCodeScanResponse {
requires_trust_remote_code?: boolean;
has_remote_code?: boolean;
approvable?: boolean;
max_severity?: string | null;
fingerprint?: string | null;
findings?: Array<{
severity: string;
file: string;
check: string;
evidence?: string;
line?: number | null;
snippet?: SnippetRowResponse[];
}>;
findings_summary?: string;
model_name?: string;
created_by_scan?: boolean;
scan_created_repos?: string[];
unsafe_files?: Array<{ path?: string; level?: string }>;
security_blocked?: boolean;
}
/** Scan a model's auto_map code for the consent dialog (backend reads config + repo
* .py, never loads the model). Token rides in the POST body, never the URL. */
export async function getRemoteCodeScan(
modelName: string,
hfToken?: string | null,
): Promise<RemoteCodeScan> {
const response = await authFetch(`/api/models/remote-code-scan`, {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({ model_name: modelName, hf_token: hfToken ?? null }),
});
if (!response.ok) {
throw new Error(await readFastApiError(response));
}
const data = (await response.json()) as RemoteCodeScanResponse;
const findings: RemoteCodeFinding[] = (data.findings ?? []).map((f) => ({
severity: f.severity as RemoteCodeSeverity,
file: f.file,
check: f.check,
evidence: f.evidence,
line: f.line ?? null,
snippet: (f.snippet ?? []).map(
(r): RemoteCodeSnippetRow => ({
number: r.number,
text: r.text,
isMatch: r.is_match,
matchStart: r.match_start,
matchEnd: r.match_end,
}),
),
}));
const unsafeFiles: UnsafeFile[] = (data.unsafe_files ?? []).map((u) => ({
path: u.path ?? "",
level: u.level ?? "unsafe",
}));
return {
requiresTrustRemoteCode: Boolean(
data.requires_trust_remote_code ?? data.has_remote_code,
),
approvable: data.approvable ?? true,
maxSeverity: (data.max_severity as RemoteCodeSeverity | null) ?? null,
fingerprint: data.fingerprint ?? null,
findings,
findingsSummary: data.findings_summary ?? "",
modelName: data.model_name ?? modelName,
createdByScan: Boolean(data.created_by_scan),
// Fall back to the primary flag for an older backend.
scanCreatedRepos:
data.scan_created_repos ??
(data.created_by_scan ? [data.model_name ?? modelName] : []),
unsafeFiles,
securityBlocked: Boolean(data.security_blocked),
};
}
/** Decline cleanup: purge what the scan downloaded. Fire-and-forget; the backend only
* removes a metadata-only cache entry it created (never weights, a loaded model, or a local path). */
export async function discardRemoteCodeDownload(modelName: string): Promise<void> {
try {
await authFetch(`/api/models/discard-remote-code`, {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({ model_name: modelName }),
});
} catch {
// Best-effort cleanup; ignore failures.
}
}

290
studio/frontend/src/features/security/components/remote-code-consent-dialog.tsx Normal file
View file

@ -0,0 +1,290 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import {
AlertDialog,
AlertDialogAction,
AlertDialogCancel,
AlertDialogContent,
AlertDialogDescription,
AlertDialogFooter,
AlertDialogHeader,
AlertDialogTitle,
} from "@/components/ui/alert-dialog";
import { Badge } from "@/components/ui/badge";
import { cn } from "@/lib/utils";
import {
SecurityWarningIcon,
ShieldBanIcon,
SourceCodeSquareIcon,
} from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { severityTone } from "../lib/severity-tone";
import { useRemoteCodeConsentDialogStore } from "../stores/remote-code-consent-dialog-store";
import type {
RemoteCodeFinding,
RemoteCodeSeverity,
RemoteCodeSnippetRow,
UnsafeFile,
} from "../types";
function matchTone(severity: RemoteCodeSeverity | string): {
row: string;
span: string;
} {
switch (severity) {
case "CRITICAL":
case "HIGH":
return {
row: "bg-red-500/10 border-l-2 border-red-500/60",
span: "bg-red-500/25 text-foreground",
};
case "MEDIUM":
return {
row: "bg-amber-500/10 border-l-2 border-amber-500/60",
span: "bg-amber-500/25 text-foreground",
};
default:
return {
row: "bg-muted border-l-2 border-border",
span: "bg-foreground/15 text-foreground",
};
}
}
function SnippetLine({
row,
severity,
}: {
row: RemoteCodeSnippetRow;
severity: RemoteCodeSeverity;
}) {
const tone = matchTone(severity);
const hasSpan =
row.isMatch &&
typeof row.matchStart === "number" &&
typeof row.matchEnd === "number" &&
row.matchEnd > row.matchStart;
let body;
if (hasSpan) {
const start = row.matchStart as number;
const end = row.matchEnd as number;
body = (
<>
{row.text.slice(0, start)}
<span className={cn("rounded-[3px] px-0.5 font-semibold", tone.span)}>
{row.text.slice(start, end)}
</span>
{row.text.slice(end)}
</>
);
} else {
body = row.text || " ";
}
return (
<div
className={cn(
"flex",
row.isMatch ? tone.row : "border-l-2 border-transparent",
)}
>
<span className="w-10 shrink-0 select-none pr-3 text-right tabular-nums text-muted-foreground/50">
{row.number}
</span>
<code
className={cn(
"whitespace-pre pr-3",
row.isMatch ? "text-foreground" : "text-muted-foreground",
)}
>
{body}
</code>
</div>
);
}
function UnsafeFileCard({ file }: { file: UnsafeFile }) {
return (
<div className="flex flex-wrap items-center gap-x-2 gap-y-1 rounded-lg border bg-muted/40 px-3 py-2">
<Badge
variant="outline"
className={cn(
"shrink-0 text-[10px] font-semibold uppercase tracking-wide",
severityTone("CRITICAL"),
)}
>
{file.level}
</Badge>
<span className="flex min-w-0 items-center gap-1 font-mono text-xs text-muted-foreground">
<HugeiconsIcon icon={SourceCodeSquareIcon} className="size-3.5 shrink-0" />
<span className="truncate">{file.path}</span>
</span>
</div>
);
}
function FindingCard({ finding }: { finding: RemoteCodeFinding }) {
const fileLabel = finding.line ? `${finding.file}:${finding.line}` : finding.file;
const snippet = finding.snippet ?? [];
return (
<div className="min-w-0 overflow-hidden rounded-lg border">
<div className="flex flex-wrap items-center gap-x-2 gap-y-1 border-b bg-muted/40 px-3 py-2">
<Badge
variant="outline"
className={cn(
"shrink-0 text-[10px] font-semibold tracking-wide",
severityTone(finding.severity),
)}
>
{finding.severity}
</Badge>
<span className="text-sm font-medium">{finding.check}</span>
<span className="ml-auto flex shrink-0 items-center gap-1 font-mono text-xs text-muted-foreground">
<HugeiconsIcon icon={SourceCodeSquareIcon} className="size-3.5" />
{fileLabel}
</span>
</div>
{snippet.length > 0 ? (
<div className="min-w-0 overflow-x-auto bg-background py-2 font-mono text-xs leading-relaxed">
{snippet.map((row) => (
<SnippetLine key={row.number} row={row} severity={finding.severity} />
))}
</div>
) : finding.evidence ? (
<div className="px-3 py-2 font-mono text-xs text-muted-foreground">
{finding.evidence}
</div>
) : null}
</div>
);
}
/** App-wide consent dialog for trust_remote_code loads: shows scan findings with the
* flagged code in context; CRITICAL is a hard block. Mounted once in the root layout. */
export function RemoteCodeConsentDialog() {
const open = useRemoteCodeConsentDialogStore((s) => s.open);
const scan = useRemoteCodeConsentDialogStore((s) => s.scan);
const resolve = useRemoteCodeConsentDialogStore((s) => s.resolve);
const displayName = scan?.modelName?.split("/").pop() || "This model";
const blocked = scan ? !scan.approvable : false;
const findings = scan?.findings ?? [];
const unsafeFiles = scan?.unsafeFiles ?? [];
// Malware (an HF-flagged unsafe serialized file) is a hard block with its own copy.
const malware = unsafeFiles.length > 0;
return (
<AlertDialog
open={open}
onOpenChange={(next) => {
if (!next) resolve(false);
}}
>
<AlertDialogContent className="max-w-2xl">
<AlertDialogHeader className="min-w-0">
{/* w-full: AlertDialogHeader is a grid with place-items-center, which
otherwise sizes this row to its content and lets a wide code snippet
push past the dialog. Filling the track keeps everything inside. */}
<div className="flex w-full min-w-0 items-start gap-3">
<div
className={cn(
"flex size-9 shrink-0 items-center justify-center rounded-full",
blocked
? "bg-red-500/10 text-red-600 dark:text-red-400"
: "bg-amber-500/10 text-amber-600 dark:text-amber-400",
)}
>
<HugeiconsIcon
icon={blocked ? ShieldBanIcon : SecurityWarningIcon}
className="size-5"
/>
</div>
{/* Title, description and scan results share one column to the right
of the icon, so the body aligns under the title instead of jumping
back to the dialog's left padding. */}
<div className="min-w-0 flex-1 space-y-3">
<div className="space-y-1">
<AlertDialogTitle>
{malware
? "Unsafe files detected"
: blocked
? "Custom code blocked"
: "Enable custom code for this model?"}
</AlertDialogTitle>
<AlertDialogDescription>
{malware ? (
<>
<span className="font-medium text-foreground">
{displayName}
</span>{" "}
contains files that Hugging Face's security scan flagged as
unsafe (for example, a malicious pickle that would run code
when the model loads). It cannot be loaded. The flagged
files were never downloaded.
</>
) : (
<>
<span className="font-medium text-foreground">
{displayName}
</span>{" "}
declares custom Python code in its repository.{" "}
{blocked
? "A security scan flagged CRITICAL issues, so it cannot be enabled."
: findings.length > 0
? "Review the security scan below. Continue only if you trust the model source."
: "Continue only if you trust the model source."}
</>
)}
</AlertDialogDescription>
</div>
{malware ? (
<div className="min-w-0 space-y-2">
<p className="text-xs font-medium text-muted-foreground">
Our automatic scanner flagged issues including:
</p>
<div className="max-h-[14rem] min-w-0 space-y-2 overflow-y-auto pr-1">
{unsafeFiles.map((f, i) => (
<UnsafeFileCard key={`${f.path}-${i}`} file={f} />
))}
</div>
</div>
) : null}
{findings.length > 0 ? (
<div className="min-w-0 space-y-2">
<p className="text-xs font-medium text-muted-foreground">
Our automatic scanner flagged issues including:
</p>
<div className="max-h-[22rem] min-w-0 space-y-3 overflow-y-auto pr-1">
{findings.map((f, i) => (
<FindingCard key={i} finding={f} />
))}
</div>
</div>
) : null}
{!malware && !blocked && findings.length === 0 ? (
<p className="text-xs text-muted-foreground">
Our automatic scanner did not flag any worrying files, but
please double check.
</p>
) : null}
</div>
</div>
</AlertDialogHeader>
<AlertDialogFooter>
<AlertDialogCancel>Cancel</AlertDialogCancel>
{blocked ? null : (
<AlertDialogAction onClick={() => resolve(true)}>
Enable and continue
</AlertDialogAction>
)}
</AlertDialogFooter>
</AlertDialogContent>
</AlertDialog>
);
}

73
studio/frontend/src/features/security/hooks/use-remote-code-consent.ts Normal file
View file

@ -0,0 +1,73 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import {
discardRemoteCodeDownload,
getRemoteCodeScan,
} from "../api/remote-code-api";
import { useRemoteCodeConsentDialogStore } from "../stores/remote-code-consent-dialog-store";
import type { RemoteCodeScan } from "../types";
interface ConfirmArgs {
modelName: string;
// Resolves the same findings + fingerprint for gated/private repos.
hfToken?: string | null;
// Coarse fallback when the scan endpoint is unreachable.
requiresTrustRemoteCode?: boolean;
// Called on approval with the pinning fingerprint.
onApprove: (fingerprint: string | null) => void;
}
/** Gate a load that may need trust_remote_code: scan, show the consent dialog, and on
* approval call onApprove with the pinning fingerprint. Returns false if declined. */
export async function confirmRemoteCodeIfNeeded({
modelName,
hfToken,
requiresTrustRemoteCode,
onApprove,
}: ConfirmArgs): Promise<boolean> {
let scan: RemoteCodeScan;
try {
scan = await getRemoteCodeScan(modelName, hfToken);
} catch {
scan = {
requiresTrustRemoteCode: Boolean(requiresTrustRemoteCode),
approvable: true,
maxSeverity: null,
fingerprint: null,
findings: [],
findingsSummary: "",
modelName,
createdByScan: false,
scanCreatedRepos: [],
unsafeFiles: [],
securityBlocked: false,
};
}
// Open the dialog for custom-code consent OR flagged unsafe files. Otherwise a model
// can still need trust_remote_code via its YAML default (no auto_map, e.g. GLM-4.7-Flash):
// propagate the caller's requirement with an empty pin instead of sending false.
if (!scan.requiresTrustRemoteCode && scan.unsafeFiles.length === 0) {
if (requiresTrustRemoteCode) onApprove(null);
return true;
}
const confirmed = await useRemoteCodeConsentDialogStore
.getState()
.requestConsent(scan);
if (!confirmed) {
// Declined: purge every repo our scan first downloaded (a LoRA scan pulls adapter +
// base) so untrusted code is not left on disk. Fall back to the primary flag for an older backend.
const toPurge =
scan.scanCreatedRepos.length > 0
? scan.scanCreatedRepos
: scan.createdByScan
? [scan.modelName]
: [];
for (const repo of toPurge) void discardRemoteCodeDownload(repo);
return false;
}
onApprove(scan.fingerprint);
return true;
}

8
studio/frontend/src/features/security/index.ts Normal file
View file

@ -0,0 +1,8 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
export { RemoteCodeConsentDialog } from "./components/remote-code-consent-dialog";
export { confirmRemoteCodeIfNeeded } from "./hooks/use-remote-code-consent";
export { getRemoteCodeScan } from "./api/remote-code-api";
export { useRemoteCodeConsentDialogStore } from "./stores/remote-code-consent-dialog-store";
export type { RemoteCodeFinding, RemoteCodeScan, RemoteCodeSeverity } from "./types";

17
studio/frontend/src/features/security/lib/severity-tone.ts Normal file
View file

@ -0,0 +1,17 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import type { RemoteCodeSeverity } from "../types";
/** Tailwind classes for a severity badge. */
export function severityTone(severity: RemoteCodeSeverity | string): string {
switch (severity) {
case "CRITICAL":
case "HIGH":
return "border-red-500/30 text-red-700 dark:text-red-300";
case "MEDIUM":
return "border-amber-500/30 text-amber-700 dark:text-amber-300";
default:
return "border-border text-muted-foreground";
}
}

37
studio/frontend/src/features/security/stores/remote-code-consent-dialog-store.ts Normal file
View file

@ -0,0 +1,37 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { create } from "zustand";
import type { RemoteCodeScan } from "../types";
type Resolver = (confirmed: boolean) => void;
// One in-flight consent at a time; a new request resolves any prior pending one as
// declined so its promise never leaks.
let pendingResolver: Resolver | null = null;
interface RemoteCodeConsentDialogStore {
open: boolean;
scan: RemoteCodeScan | null;
requestConsent: (scan: RemoteCodeScan) => Promise<boolean>;
resolve: (confirmed: boolean) => void;
}
export const useRemoteCodeConsentDialogStore = create<RemoteCodeConsentDialogStore>()(
(set) => ({
open: false,
scan: null,
requestConsent: (scan) =>
new Promise<boolean>((resolve) => {
pendingResolver?.(false);
pendingResolver = resolve;
set({ open: true, scan });
}),
resolve: (confirmed) => {
const resolver = pendingResolver;
pendingResolver = null;
set({ open: false, scan: null });
resolver?.(confirmed);
},
}),
);

46
studio/frontend/src/features/security/types.ts Normal file
View file

@ -0,0 +1,46 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
export type RemoteCodeSeverity = "CRITICAL" | "HIGH" | "MEDIUM" | "LOW";
/** One line of the code window shown for a finding. */
export interface RemoteCodeSnippetRow {
number: number; // 1-based source line
text: string;
isMatch: boolean; // the flagged line
matchStart?: number; // match column span, when known
matchEnd?: number;
}
export interface RemoteCodeFinding {
severity: RemoteCodeSeverity;
file: string;
check: string;
evidence?: string;
line?: number | null; // 1-based
snippet?: RemoteCodeSnippetRow[];
}
/** A repo file Hugging Face's security scan flagged (e.g. a malicious pickle). */
export interface UnsafeFile {
path: string;
level: string; // HF level: unsafe | suspicious | malicious
}
/** Backend remote-code scan result; drives the consent dialog and pins approval to the scanned code version. */
export interface RemoteCodeScan {
requiresTrustRemoteCode: boolean;
approvable: boolean; // false for CRITICAL (no override)
maxSeverity: RemoteCodeSeverity | null;
fingerprint: string | null;
findings: RemoteCodeFinding[];
findingsSummary: string;
modelName: string;
// True when our scan first downloaded this repo, so a decline may purge it.
createdByScan: boolean;
// Every repo the scan first cached (a LoRA scan pulls adapter + base); a decline
// purges each. Supersedes createdByScan, which tracks only the primary.
scanCreatedRepos: string[];
unsafeFiles: UnsafeFile[]; // files HF flagged unsafe; non-empty => hard block
securityBlocked: boolean; // blocked specifically by the malware gate
}

1
studio/frontend/src/features/training/api/mappers.ts
View file

@ -82,6 +82,7 @@ export function buildTrainingStartPayload(
? config.visionImageSize
: null,
trust_remote_code: config.trustRemoteCode ?? false,
approved_remote_code_fingerprint: config.approvedRemoteCodeFingerprint ?? null,
hf_dataset: hfDataset,
subset: hfDataset ? config.datasetSubset : null,
train_split: hfDataset ? config.datasetSplit : null,

16
studio/frontend/src/features/training/api/models-api.ts
View file

@ -92,10 +92,14 @@ interface LocalModelListResponse {
models: LocalModelInfo[];
}
/** Ask the backend whether a model is a vision model (GET /api/models/check-vision/{model_name}). */
export async function checkVisionModel(modelName: string): Promise<boolean> {
/** GET /api/models/check-vision; pass the token so a gated/private VLM is not misread as non-vision. */
export async function checkVisionModel(
modelName: string,
hfToken?: string | null,
): Promise<boolean> {
const encoded = encodeURIComponent(modelName);
const response = await authFetch(`/api/models/check-vision/${encoded}`);
const query = hfToken?.trim() ? `?hf_token=${encodeURIComponent(hfToken.trim())}` : "";
const response = await authFetch(`/api/models/check-vision/${encoded}${query}`);
if (!response.ok) {
// If the check fails (e.g. network error), default to non-vision
return false;
@ -104,12 +108,14 @@ export async function checkVisionModel(modelName: string): Promise<boolean> {
return data.is_vision;
}
/** Ask the backend whether a model is an embedding model (GET /api/models/check-embedding/{model_name}). */
/** GET /api/models/check-embedding; pass the token for gated/private repos. */
export async function checkEmbeddingModel(
modelName: string,
hfToken?: string | null,
): Promise<boolean> {
const encoded = encodeURIComponent(modelName);
const response = await authFetch(`/api/models/check-embedding/${encoded}`);
const query = hfToken?.trim() ? `?hf_token=${encodeURIComponent(hfToken.trim())}` : "";
const response = await authFetch(`/api/models/check-embedding/${encoded}${query}`);
if (!response.ok) {
// If the check fails (e.g. network error), default to non-embedding
return false;

42
studio/frontend/src/features/training/hooks/use-training-actions.ts
View file

@ -2,6 +2,7 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { primeNativeNotificationPermission } from "@/lib/native-notifications";
import { confirmRemoteCodeIfNeeded } from "@/features/security";
import { useCallback } from "react";
import { toast } from "@/lib/toast";
import { checkDatasetFormat } from "../api/datasets-api";
@ -127,6 +128,24 @@ export function useTrainingActions() {
return false;
}
// Consent gate for the selected model's custom (auto_map) code.
if (config.selectedModel) {
const remoteCodeOk = await confirmRemoteCodeIfNeeded({
modelName: config.selectedModel,
hfToken: config.hfToken.trim() || null,
requiresTrustRemoteCode: config.trustRemoteCode,
onApprove: (fingerprint) =>
useTrainingConfigStore.setState({
trustRemoteCode: true,
approvedRemoteCodeFingerprint: fingerprint,
}),
});
if (!remoteCodeOk) {
runtimeStore.setStarting(false);
return false;
}
}
// Re-read config after potential store updates from dataset check
const payload = buildTrainingStartPayload(useTrainingConfigStore.getState());
runtimeStore.setStartResources(payload.model_name, payload.hf_dataset, false);
@ -199,6 +218,29 @@ export function useTrainingActions() {
runtimeStore.setStartResources(payload.model_name, payload.hf_dataset, true);
// Resume goes straight to startTraining, so it runs the same consent gate as a
// fresh start; otherwise a resumed custom-code run hits the worker block with no dialog.
if (payload.model_name) {
let trustRemoteCode = Boolean(payload.trust_remote_code);
let approvedRemoteCodeFingerprint =
payload.approved_remote_code_fingerprint ?? null;
const remoteCodeOk = await confirmRemoteCodeIfNeeded({
modelName: payload.model_name,
hfToken: payload.hf_token ?? null,
requiresTrustRemoteCode: trustRemoteCode,
onApprove: (fingerprint) => {
trustRemoteCode = true;
approvedRemoteCodeFingerprint = fingerprint;
},
});
if (!remoteCodeOk) {
runtimeStore.setStarting(false);
return false;
}
payload.trust_remote_code = trustRemoteCode;
payload.approved_remote_code_fingerprint = approvedRemoteCodeFingerprint;
}
const response = await startTraining(payload);
if (response.status === "error") {
throw new Error(response.error || response.message);

17
studio/frontend/src/features/training/stores/training-config-store.ts
View file

@ -390,8 +390,8 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
visionImageSize: DEFAULT_HYPERPARAMS.visionImageSize,
});
// Fallback vision check if config endpoint fails.
void checkVisionModel(modelName)
// Fallback vision check; pass the token so a gated/private VLM classifies right.
void checkVisionModel(modelName, get().hfToken || undefined)
.then((isVision) => {
if (get().selectedModel !== modelName) return;
set({
@ -498,12 +498,23 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
const previousModel = get().selectedModel;
// Reset vision_image_size on a true switch only; same-model reloads
// go through the mapper, which preserves the user's choice.
const patch: { selectedModel: string | null; modelDefaultsError: null; visionImageSize?: number | null } = {
const patch: {
selectedModel: string | null;
modelDefaultsError: null;
visionImageSize?: number | null;
trustRemoteCode?: boolean;
approvedRemoteCodeFingerprint?: string | null;
} = {
selectedModel,
modelDefaultsError: null,
};
if (selectedModel !== previousModel) {
patch.visionImageSize = DEFAULT_HYPERPARAMS.visionImageSize;
// Clear the prior model's approval so a clean model is not trained with a
// stale trust_remote_code=true (disables fused CE). Its own YAML default is
// re-applied below, and a custom-code model re-opens the dialog before start.
patch.trustRemoteCode = false;
patch.approvedRemoteCodeFingerprint = null;
}
set(patch);

2
studio/frontend/src/features/training/types/api.ts
View file

@ -12,6 +12,8 @@ export interface TrainingStartRequest {
vision_image_size?: number | null;
/** Allow loading models with custom code. Only enable for repos you trust. */
trust_remote_code?: boolean;
/** sha256 fingerprint pinning user approval of this exact custom-code version. */
approved_remote_code_fingerprint?: string | null;
hf_dataset: string | null;
subset: string | null;
train_split: string | null;

1
studio/frontend/src/features/training/types/config.ts
View file

@ -77,6 +77,7 @@ export interface TrainingConfigState {
isDatasetImage: boolean | null;
isDatasetAudio: boolean;
trustRemoteCode: boolean;
approvedRemoteCodeFingerprint?: string | null;
finetuneVisionLayers: boolean;
finetuneLanguageLayers: boolean;
finetuneAttentionModules: boolean;