Commit graph

4 commits

Author SHA1 Message Date
Indrajit swain
5205a38f70
Studio: keep SSM kernel installs alive and reap the worker on load timeout (#9419)
---------

Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
2026-08-25 00:22:28 -03:00
Eyera
0b147fdd2a
fix(studio): harden training setup, lifecycle, and audio loading (#8103)
* fix(studio): harden training setup and model loading

Scan every model load root before approving remote code, and pin third-party codec sources to verified revisions.

Align dataset option validation across the UI and backend, preserve manual drafts, and include edits in training start identity.

Require job-scoped stop requests and retain bounded early-cancel tombstones without unsafe eviction.

* fix(studio): harden training lifecycle and audio loading

Preserve early start cancellations with bounded tombstones and explicit capacity handling for concurrent requests.

Pin and verify third-party audio sources and codec artifacts with safe archive extraction and offline cache migration.

Keep automatic evaluation data separate from every split included in a combined training instruction.

Keep training summaries and GPU progress state aligned with the active configuration through public feature exports.

Make route and lifecycle tests deterministic by isolating Hub reachability and executor behavior.

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

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

* fix(studio): correct completion masking and offline training

Apply Alpaca response markers safely when completion-only training is enabled.

Skip optional runtime dependency installs while Studio is offline.

Limit causal-conv1d hooks to recognized model families without dropping supported architectures.

Disable known-broken TileLang dispatch when offline repair is unavailable.

Use cached GGUF and model-size metadata without offline Hub retries.

* fix(studio): resolve causal conv kernels from model configs

Detect causal-conv1d requirements from resolved model architectures before loading model code. Keep name matching as a fallback while excluding unrelated renamed checkpoints.

* fix(studio): preserve registered training cancellation

Allow registered start cancellations to reclaim the oldest expiring tombstone when unknown cancellation capacity is full. Preserve the hard capacity limit and 429 response for unregistered request IDs.

* Fix reset job scoping, DAC fast path, scan target and subset splits for PR #8103

- /api/train/reset: an unscoped reset could force-terminate a run mid-cancel. The
  guard now refuses to touch a live run it cannot prove it owns. The field stays
  optional so pre-rework clients, which POST /reset with no body, keep working.
- ensure_dac_speech_weights: install the download into the pinned destination, so
  later loads hit the fast path instead of re-downloading and re-hashing 295 MB
  under the install lock.
- _requires_security_review_for_model: apply the same load_scan_target alias
  normalization the sibling remote-code check gained, else the Spark-TTS alias
  404s and fails open to "no review needed".
- _purge_package_bytecode: best effort. It runs without the install lock over a
  cache shared by the inference and training workers; 7 of 8 concurrent imports
  died on it. Also moved inside the try so a failure cannot strand the cache dir
  on sys.path, and snapshot sys.modules before the origin audit.
- commitSubset: clear the backing splits too, else the render-phase draft sync
  reads the previous subset's split back into the boxes it just reset.
- Fix two tests that fail on the branch: the DAC assertion pinned the old return
  value, and the causal-conv1d assertion matched call formatting.

* Make the SSM runtime tests Windows aware

ensure_ssm_runtime deliberately skips causal-conv1d on win32 (no prebuilt wheel),
so the two install-order assertions only hold off Windows. Caught on a real
windows-latest runner.

* Update two frontend source contracts the branch moved

- captureTrainingStartInputs now delegates to createTrainingStartInputIdentity,
  so the normalize/flags assertions belong against training-start-inputs.ts.
- resetTraining takes a RequiredTrainingJobScope and always sends the body, which
  is stronger than the hasScope branch the contract pinned.

Both fail on the branch today; caught by tests/studio, which the studio backend
job does not cover.

* Fix unscoped reset compat, pyc purge fail-open and DAC fallback for PR #8103

Corrects four things in my earlier commits on this branch.

Unscoped /api/train/reset returned "superseded" (HTTP 200) for a live run. The
pre-rework cancel dialog chains stopTrainingRun then a bodyless reset, so an older
client read that 200 as success and cleared its UI while training kept running.
Return "active" (409) instead: same answer a live run already gives, and one those
clients already handle. It still never force-terminates, so a bodyless reset landing
between current_job_id being set and _cancel_requested being cleared cannot kill the
run that just started.

The bytecode purge was made best-effort, but it is the only thing stopping a stale
or planted .pyc shadowing a verified .py: the manifest skips __pycache__ and the
origin audit reads __file__, which still names the .py. Tolerate only
FileNotFoundError, the real concurrent-purge race, and let PermissionError fail the
load again.

The DAC fast path copies 295 MB inside the hub cache with only Timeout caught, so a
full disk turned a hash-verified download into a hard failure. Fall back to the
verified hub path on OSError.

commitSubset cleared both splits, but setDatasetSubset already does that; the extra
setDatasetSplit(null) only cost a runDatasetCheck against an assumed "train" split.

* Restore the eval split reset and the cancelled-run dismiss for PR #8103

Two corrections to 4e5389c21.

commitSubset: I removed setDatasetEvalSplit(null) because setDatasetSubset already
nulls datasetEvalSplit. It does, but it never resets evalSteps, and
setDatasetEvalSplit is not a plain setter: it zeroes evalSteps and runs
streamingCompatiblePatch. Without it, changing the subset left evaluation armed with
no split, which routes/training.py rejects with 422 once streaming is on, and which
silently auto-detects an eval split otherwise. Restored. setDatasetSplit(null) stays
out, since its only unique effect was a runDatasetCheck against an assumed "train".

Unscoped reset: returning "active" for every live run was too broad. The pre-rework
cancel dialog only dismisses after stopTrainingRun succeeded, so _cancel_requested is
already set and clearing the UI is right; 409 there just wedged the overlay behind a
"Training still active" toast. Now 409 only when no stop was requested, which is the
stale-tab case the change was for. Still no force_terminate on an unscoped reset.

Also covers the __pycache__ branch of the purge, which is the route a planted .pyc
actually takes; the existing test only reached the top-level .pyc loop.

* Keep live start cancellations and fall back on a full disk for PR #8103

Two fixes plus the regression tests the earlier lifecycle fixes shipped without.

Cancelling the active start at tombstone capacity reclaimed a slot by deleting the
soonest-expiring entry. Expired ones are already pruned a few lines above, so that
entry was always live, and dropping it let a delayed /start spawn the job it had
cancelled. Reserve capacity instead: only the owner of the active start reaches that
branch and there is at most one, so the table lands at cap + 1 rather than forgetting
a cancellation. Unregistered ids still hit the hard cap.

ensure_dac_speech_weights migrated a pre-existing legacy file with an unguarded copy,
so a hub cache that cannot absorb a second 295 MB copy failed weights that had already
passed size and sha256. Same fallback the download branch below it already uses.

Tests: renamed and local SSM checkpoints resolving from config rather than name, the
owner of an active start staying cancellable at capacity, live cancellations surviving
that cancel, and the full-disk legacy DAC fallback.

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

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

* Opt the pinned source checkouts into Git long paths for PR #8103

Cross-OS CI caught 14 failures on windows-latest, all of them the pinned checkout
dying with "error: unable to write file ...: Filename too long". Git for Windows still
enforces MAX_PATH unless told otherwise, and the cache nests a 40-char revision, a
staging dir and .git/objects under the studio home. A venv-inferred home already
measures about 253 of the 260 characters, so a slightly longer user or install path
fails on a normal Windows machine, not just under the deeper pytest tmp dir.

Passed per invocation with -c so no user or system Git config is touched, and it is a
no-op off Windows.

* Bound pending cancels, reach legacy DAC weights, and delete read-only checkouts for PR #8103

Three fixes, two of them on my own previous commit.

Moving the owner cancel from evict-oldest to a one-slot overshoot also changed what the
hardcoded reclaim_capacity=True on the pending non-owner branch did: it used to evict,
so the table stayed at the cap, and it started overshooting instead. Start plus cancel
could then be repeated to grow it without bound (1224 entries against a cap of 1024).
That branch now takes the plain reservation, so the extra slot belongs to the owner of
the active run alone.

The DAC legacy fallback sat behind destination.parent.mkdir() and the install lock, both
of which need a writable cache, so a read-only or full hub cache raised before weights
we can already verify were ever looked at. Fall back at both points, and only to an
artifact that passes the same size and sha256 check.

Cross-OS CI then caught replacing a pinned checkout failing on windows-latest with
WinError 5: Git marks .git/objects read-only and Windows will not delete a read-only
file, so any repair or revision change died there. Clear the attribute and retry, only
when the path is genuinely not writable, so an open handle still surfaces.

* Delete the cached checkout the Windows-safe way in the migration test for PR #8103

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen <moonshotaisubstack@gmail.com>
2026-08-08 01:45:40 -07:00
Michael Han
a00fe86c13
Studio: read model text as utf-8 so umlauts survive on Windows (#7467)
* Studio: read model text as utf-8 so umlauts survive on Windows

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

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

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

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

Encoding arguments only, no logic changes.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

* Fail closed when a KFD topology node does not decode

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

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

* Tighten the comments added in this branch

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

* Stop reading a UTF-8 record a second time

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

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

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

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

* Pin the scanned source fixture's line endings

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

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

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

* Trim the newer comments to their point

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

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

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

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

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

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

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

* Gate the remaining tilelang install tests to Linux

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

* Gate the Windows-incompatible worker and ROCm tests

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

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

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

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

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

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

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

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

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

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

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

* Tighten encoding comments

* Follow subprocess calls through callable aliases in the encoding guard

---------

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

---------

Co-authored-by: Unsloth <michaelhan@Michaels-MacBook-Pro.local>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen <unslothshared@gmail.com>
2026-07-28 21:27:27 -07:00
Daniel Han
2a05426adb
Auto-install SSM kernels (causal-conv1d, mamba-ssm) for inference loads (#6535)
* Auto-install SSM kernels (causal-conv1d, mamba-ssm) for inference loads

Mamba/SSM hybrids (Nemotron-H/Nano, Falcon-H1, Granite-4.0-H, ...) lazily import
mamba_ssm / causal_conv1d during from_pretrained, so loading them for chat failed
with 'mamba-ssm is required by the Mamba model but cannot be imported'. The training
worker already wheel-first installs these before a fine-tune; the inference worker
did not. Add utils/ssm_runtime.ensure_ssm_runtime and call it from the inference load
path so the same models load for inference. Training worker is untouched; a drift
test keeps the shared detection and pinned versions in lockstep.

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

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

* ssm_runtime: invalidate import caches, skip MLX, cover LoRA base

- Invalidate importlib finder caches in _is_importable and after a successful
  wheel install, so a kernel installed earlier in this same process is actually
  importable when the modeling code lazy-imports it during from_pretrained.
- Skip the SSM kernel install entirely on the MLX (Apple Silicon) load path:
  these are CUDA/ROCm Torch kernels with no MLX use and no macOS prebuilt wheel,
  so the source build would fail before the MLX backend loads the model.
- For LoRA loads, also run detection over the resolved base model, since an
  adapter id like 'me/my-lora' won't match the SSM heuristics but its SSM base
  (Nemotron-H, ...) is what needs the kernels.

Adds tests for cache invalidation and the MLX-skip / LoRA-base worker wiring.

* Tighten SSM autoinstall comments

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

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

* ssm_runtime: verify wheel imports, HIP-aware source build, build heartbeat

Address review feedback:
- Verify a prebuilt wheel actually imports before trusting it; a CUDA/ABI-mismatched
  wheel now falls back to a source build instead of returning success and failing later
  with the cryptic lazy-import error.
- HIP-aware source build: require hipcc on ROCm, inject clang --gcc-install-dir, and use
  the 1800s timeout, mirroring the training worker (ROCm has no prebuilt wheel).
- Emit a status heartbeat every 60s during the source build so a long (ROCm) build does
  not trip the orchestrator's 300s inactivity timeout.

Tests cover the wheel-not-importable fallback and the missing-hipcc ROCm bail.

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

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

* Make causal-conv1d best-effort and harden the SSM source build

- causal-conv1d is a fast path: models that merely want it (Qwen3-Next, LFM2)
  fall back to torch, so a failed install must not reject an otherwise loadable
  chat model on Windows/CPU/macOS or an ABI without a wheel. Only a true SSM
  model's mamba-ssm requirement stays fatal, matching the training worker which
  treats causal-conv1d as best-effort.
- The source build is reached only when not importable, including a wheel that
  installed but failed to import; add --reinstall/--force-reinstall so it
  replaces the broken install instead of no-opping as already satisfied.
- Add --no-cache to the ROCm uv source build to avoid reusing stale artifacts
  from a partial HIP build, mirroring the training worker.

* Address review: install SSM kernels before transformers, harden import + Windows

Codex:
- Install the SSM kernels before importing transformers. run_inference_process
  imported core.inference.inference (which imports unsloth/transformers) before the
  load, and a sidecar transformers can evaluate its optional-backend gates against
  the import state; installing causal_conv1d/mamba_ssm afterwards left those gates
  unsatisfied and a Nemotron/Falcon/Granite load still failed with "mamba-ssm is
  required". The initial model's kernels are now installed in run_inference_process
  before the ML import, via a shared _ensure_ssm_kernels helper; _handle_load keeps
  calling it (idempotent) for a LoRA's base and for later in-process loads.
- _is_importable now treats any import failure as "not importable", not only
  ImportError. An ABI-incompatible native kernel (undefined symbol after a torch/CUDA
  upgrade) raises OSError/RuntimeError; letting those escape reported
  ssm_runtime_install_failed instead of falling back to reinstall/source build.
- Skip causal-conv1d on Windows (no prebuilt wheel), mirroring the training worker.
  A causal-conv1d-only model (Qwen3-Next/LFM2) no longer drops a chat load into a
  multi-minute untimed source build; it uses the torch fallback. mamba-ssm is still
  attempted for true SSM hybrids.

Tests: test_ssm_runtime.py +5 (broken-kernel exceptions read as not-importable;
causal-conv1d skipped on win32 while mamba-ssm still installs). 36 passed.

* Trim comments to be more succinct

* Run security gates before installing SSM kernels

The SSM kernel auto-install is name-based (model_is_ssm is a substring match, no
config fetch), so a model id merely containing an SSM substring triggered a
native-package install (possibly a slow source build) before the malware and
remote-code consent gates ran. Extract those gates into _run_security_gates and
call it before the kernel install in both the pre-import path of
run_inference_process and in _handle_load, so a blocked or nonexistent model is
refused before any build. The gates are metadata-only and do not import
transformers, so they are safe to run before the pre-import install.

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

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

* Resolve remote LoRA bases before importing transformers

_resolve_base_model only reads a local adapter_config.json, so a remote LoRA
adapter whose own id has no SSM substring but whose base is a Nemotron/Falcon/
Granite model had its base discovered only by ModelConfig in _handle_load, after
transformers was imported and its optional-backend availability snapshotted, so
the SSM kernel install there was too late. Add _remote_lora_base, a metadata-only
adapter_config.json fetch (no huggingface_hub / transformers import), and use it
in the pre-import path so the base is gated and its kernels pre-installed.

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

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

* Gate only loaded roots, tier on the resolved base, read offline LoRA cache

Three follow-ups to the pre-import resolution:

- The security gate reused the SSM target list, which for a local full fine-tune
  includes the config.json-recorded base. That base is never loaded, so scanning
  it could falsely block a safe local checkpoint. Gate only the model plus a
  genuine LoRA base (matching _handle_load's mc.is_lora), separate from the
  broader SSM-install list.

- Tier activation ran on the raw adapter id, so a remote LoRA whose base needs a
  sidecar transformers version imported the default and failed. Resolve the base
  once up front and activate on it.

- _remote_lora_base bailed on offline before checking the hub cache, missing a
  cached adapter's base. Read the cached adapter_config.json when offline or when
  the fetch fails.

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

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

* Keep the pre-import gate transformers-free; harden remote LoRA resolution

The pre-import security gate called security_load_subdirs, which imports
model_config and thus transformers, snapshotting optional-backend availability
before the SSM kernels are installed and defeating the ordering. Add
compute_subdirs to _run_security_gates and pass False in the preflight so it scans
from the root only (transformers-free); _handle_load still runs the authoritative
gate with full subdir scoping after the import.

_remote_lora_base now skips existing local relative paths (is_local_path) so a
checkpoint like outputs/run1 is never treated as a Hub repo, and distinguishes a
definitive 404 (not a LoRA -> None) from transient/offline failures (read the
cache), so a repo that is now a full model no longer resolves a stale cached base.

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

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

* Probe a real model id for SSM kernels; respect HF_ENDPOINT

model_is_ssm is a substring match, so an arbitrary name could false-match and
force a mamba-ssm install that fails the load for a non-SSM model:
- a LoRA adapter id like user/falcon-h1-lora (the SSM-relevant code is the base's);
- a local checkpoint under an SSM-named parent dir, e.g. /runs/falcon-h1/llama-ckpt.

Add ssm_probe_identifier, which resolves the base (or a bare local checkpoint's
basename) and feed that to ensure_ssm_runtime from both the pre-import path and
_handle_load, so detection runs against a real model id, never an adapter id or
parent folders.

_remote_lora_base now honors HF_ENDPOINT so enterprise/mirror deployments resolve
the adapter base instead of always hitting huggingface.co.

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

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

* Tighten comments in the pre-import SSM gate/install path

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <michaelhan2050@gmail.com>
2026-06-22 04:48:29 -07:00