a00fe86c13* Studio: read model text as utf-8 so umlauts survive on Windows Chat rejects or mangles non-ASCII on Windows: "ä ö ü" in a prompt, a chat template, or a model path comes back as mojibake, or the load dies with UnicodeDecodeError. open() and Path.read_text() fall back to locale.getencoding() when no encoding is passed. On Windows that is the ANSI codepage (cp1252, cp932, cp1251, ... by system locale), never UTF-8. Hugging Face writes these files as raw UTF-8, so every read of one decodes with the wrong codec: - tokenizer_config.json, which holds the chat template. Templates routinely carry -> arrows, smart quotes and CJK, so this is the common path into chat - config.json and adapter_config.json - modules.json, Ollama manifests, and the .py sources the remote-code scanner reads before a model is allowed to load The llama-server and embedding-server stdout readers have the same problem via subprocess(text = True); they now decode utf-8 with errors = "replace" so a stray byte cannot kill a log reader. Encoding arguments only, no logic changes. tests/test_chat_text_encoding.py covers a config.json and a chat template holding umlauts, arrows and CJK, plus the remote-code scanner reading a source file with umlauts. Those pass anywhere the locale is already UTF-8, so a fourth test re-runs the readers under -X warn_default_encoding and fails on any platform if an encoding argument goes missing again. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: name utf-8 explicitly on the remaining text I/O, with an AST guard (#7465) * Studio: name utf-8 explicitly on the remaining text I/O Follow-up to the model-text reads in #7467, covering the rest of the backend: system probes (nvidia-smi, amd-smi, powershell, git, node), package installers, /proc and /sys readers, and internal marker files (pid, install id, bootstrap password, Colab credentials). Same reason as #7467. open(), Path.read_text()/write_text() and subprocess(text = True) fall back to locale.getencoding(), which on Windows is the ANSI codepage rather than UTF-8. These paths are mostly ASCII today, so this is hardening, not a live bug. Encoding arguments only, no logic changes. Adds tests/test_text_io_encoding.py: an AST guard walking every backend source and asserting text I/O names its encoding, so the class of bug cannot creep back in one call at a time. 275 files. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Catch aliased subprocess and positional Path.open, migrate legacy JSONL The guard only matched a receiver literally named subprocess, so worker.py's `import subprocess as _sp` hid three text = True installs that decode pip output with the ANSI codepage. It also skipped any .open() with more than one positional argument, though Path.open takes buffering/encoding/errors/newline positionally. Resuming a scrape written by an older release is the other half: those JSONL lines are in the locale codepage, so the UTF-8 preload raised, the dedup keys were silently forgotten and duplicates were appended to a now mixed-encoding file. Decode with the locale codepage as fallback and rewrite as UTF-8 before the append handle opens, since Windows cannot replace a file it holds open. * Stream the JSONL preload and keep a torn line from relabelling the shard Reading the whole shard to migrate it was wrong twice over. These files reach gigabytes on a large scrape, so the preload now streams line by line and the rewrite streams through a temp file. Worse, one interrupted append used to condemn the file: the whole-file UTF-8 decode failed, every byte was retried as cp1252, and the rewrite persisted mojibake over records that were fine. A line now counts as legacy only if the locale codepage both decodes it and yields valid JSON, which a torn UTF-8 line does not. Damaged lines are skipped and copied through byte for byte. When the rewrite cannot be written at all, the append handle opens with the legacy encoding rather than mixing UTF-8 into the file. install_wheel takes run = subprocess.run as a parameter, so the guard cannot see it. Both wheel installs there now name their encoding. * Decide the shard's encoding from the file, not one line at a time Some byte strings parse both ways. cp1251 `Р°` is D0 B0, which is also valid UTF-8 for `а`, so a UTF-8-first parse quietly showed the wrong text instead of migrating it. A line now yields both readings, and the file decides. Any line that parses under the codepage but not as UTF-8 is unambiguous evidence, and ambiguous lines then follow that verdict, which is enough for any real shard: ordinary Cyrillic or Japanese prose is invalid UTF-8 several times per line. Keys for ambiguous lines are re-derived from the legacy reading during the rewrite. A shard is undecidable only if every line is ambiguous, and nothing can tell those apart. latin-1 is also tried after the locale codepage, so a scrape carried from Windows to a UTF-8 machine still has a reading rather than none. Requiring valid JSON, not just a decode, keeps that from claiming torn lines. * Weigh the whole shard, and never lose a record on the fallback path One structurally valid JSON line carrying a stray 0x96 parses as cp1252, so a single-line verdict let it relabel a healthy shard and mojibake every good record in it. Each line with non-ASCII bytes now votes: parsing only under the codepage is evidence for legacy, parsing as UTF-8 is evidence against, since codepage text rarely forms valid multibyte UTF-8. Ties leave the file alone. When the migration cannot be written the append handle uses the legacy codepage, and errors = "replace" quietly turned characters it cannot hold into question marks while write() still reported success. That path now escapes to \uXXXX instead, which is ASCII, so every codepage holds it and json.loads returns the exact characters. Nothing needs replacing, so errors = "strict" is safe. stream_installer runs sys.executable, so its output is now decoded as UTF-8 by utf8_child_env rather than read as the ANSI codepage. * Only rewrite a shard we can attribute, and append ASCII when we cannot latin-1 was doing too much work. It reads any byte, so it gave a moved shard a reading, but it is the right text only for cp1252: cp1251 Привет came back as Ïðèâåò and the rewrite made that permanent. The codepage is now trusted only when it is the locale's, and an untrusted reading is never written back. That leaves three cases where the file holds bytes UTF-8 cannot read and we are not converting it: no codepage to attribute it to, ambiguous lines outvoting the unambiguous ones, and a preload that could not read the file at all. All three used to append UTF-8 into it. They now append pure ASCII, which every ASCII-compatible codepage stores identically, so the file keeps decoding exactly as it did and no record is lost. Keys from the two readings are also kept apart. A damaged line in a healthy shard was marked seen through its codepage reading, so the retry that would have replaced the unreadable record was refused as a duplicate. * Let the flash-attn install stub take the kwargs the installer now passes _run_kwargs gained encoding and errors, so the one stub in this file that spelled its signature out rejected the call. The other four here already take **kwargs; this one now matches. * Do not let a stuck temp file mask the migration failure unlink() on the failure path could raise in its own right, on a stale .utf8.tmp directory or a temp another process holds. That escaped the constructor instead of returning False, so the caller never reached the ASCII append fallback that keeps the shard single-encoding. The pip fallback in install_wheel also spawns a Python child, so it gets utf8_child_env like the probe above it already had. The uv and nvidia-smi children are native binaries, where PYTHONIOENCODING would do nothing. * Stop converting legacy shards; the encoding that wrote them is unknowable trusted only ever meant that the bytes parse under this machine's codepage, which for a single-byte codepage is nearly always true. A cp1251 shard opened on a cp1252 Windows box decodes cleanly and would have been rewritten with Привет as Ïðèâåò. That is the fourth way this rewrite could corrupt a shard, and the common cause is that a file's encoding cannot be recovered from its bytes. So the rewrite is gone. The shard is left exactly as found, and appends are pure ASCII whenever it holds bytes UTF-8 cannot read, which is what actually delivered the no-mixed-encoding guarantee the rewrite was added for. Dedup keys still come from whichever reading parses, since ids are ASCII either way. This also removes the temp file, so there is no longer any file mode or ACL to carry across. * Scan the sandbox shim; it is shipped code, not a build artifact sandbox_site is on the sandboxed child's PYTHONPATH for every Python run (tools.py:332, 2660), so excluding it let two unannotated text calls through in code we ship. Both read and write the remap sidecar, which holds file paths. The exclusion list is meant for build output only, so the directory comes off it and the two calls name their encoding. * Force the worker's pip children to UTF-8, and read DBCS keys with a DBCS codec The three installer calls run sys.executable -m pip with an inherited environment, so the parent decoded UTF-8 while the child emitted the ANSI codepage. They now go through utf8_child_env like the other Python children. Two tests asserted no env kwarg was passed as a stand-in for no HIP flag being injected. They now assert the flag itself, which is the guarantee they were written for and does not depend on how the env is delivered. Separately, latin-1 cannot stand in for a double-byte codepage while recovering dedup keys: cp932 表 is 95 5C, and the trail byte reads as a JSON backslash, so the record failed to parse and its id was forgotten, appending a duplicate on resume. cp932, cp936, cp949 and cp950 are tried too. The reading is still only ever used for keys, which are ASCII and identical whichever codec parses. * Require more than one legacy line before trusting its dedup keys A shard whose valid records are all ASCII casts no UTF-8 votes, so a single damaged line won the vote by itself, its key was remembered, and the retry that would have replaced the unreadable record was refused. One such line is genuinely undecidable: a legacy record with one accented character and an ASCII record with one stray byte are the same shape. Reading it as damage costs a duplicate; reading it as legacy loses the record for good. Only one of those is recoverable, so it is now read as damage. A real legacy shard has a legacy line for every record carrying an umlaut, so its dedup is unaffected. * Append ASCII whenever the shard already holds non-ASCII bytes The gate asked whether any line was undecodable as UTF-8, which misses a shard where every legacy line happens to be valid UTF-8 too. A cp1251 shard of Р° records is bytes D0 B0 throughout, so appending 世界 as UTF-8 left a file where cp1251 reads the old records correctly and the new one as mojibake, and UTF-8 does the reverse. No single decoding recovered the whole scrape. The gate is now simply whether the shard holds any non-ASCII byte at all, which covers both cases and is easier to reason about: if what is already there reads differently under different encodings, do not add more bytes that do. Appending ASCII costs only \uXXXX escapes, which json.loads turns back into the exact characters, and it leaves the new record correct under either reading. * Skip the two Linux-gated flash-attn tests off Linux _should_try_runtime_flash_attn_install ends in sys.platform.startswith( "linux"), and the threshold test one line above already asserts exactly that, so the two tests that drive _ensure_flash_attn_for_long_context past the gate cannot pass anywhere else: the call returns before it reports a status. They were written on Linux and only surface once the suite actually runs on Windows or macOS, where both fail on an empty status list. This PR is about making the backend behave on Windows, so its own suite should be runnable there. * Fail closed when a KFD topology node does not decode This PR pins that read to utf-8, which turns an undecodable byte into UnicodeDecodeError. That is a ValueError, not an OSError, so it slips past the handler one line below and escapes a helper whose docstring promises to fail closed on any unreadable node. The caller would then lose the whole HIP-order map on a machine that has AMD GPUs, and the reason the helper fails closed is that dropping a node shifts every later ordinal and lets a similar-capacity GPU pass the total-size guard while showing another card's usage. Widening the handler is the same one-line change main already made in #7487, so the two agree and the eventual merge is clean. * Tighten the comments added in this branch * Treat an undecodable marker and undecodable metadata as malformed, not fatal Two more places where pinning the decode changed the failure mode. A UnicodeDecodeError is a ValueError, so neither `except OSError` nor `except (JSONDecodeError, OSError)` catches it, and both sites had a documented fallback that stopped being reached. An undecodable .transport marker used to read as an unknown value, and the caller then safely purged and restarted the partial download. It now aborts prepare_cache_for_transport instead, so the transfer fails rather than retrying. Undecodable .meta.json used to fall back to the file's own name, the same way invalid JSON does. It now aborts URI construction for the entire unstructured seed, so one corrupt byte in original_filename takes out the whole dataset. Both handlers are widened, matching the KFD fix earlier on this branch. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Widen two more decode guards, and pin the kernel installer's pipe Same shape as the ones already fixed here: the read was pinned to UTF-8 while the handler around it still only catches OSError, and UnicodeDecodeError is a ValueError. hf_cache_snapshot_dir answers whether a model is already on disk, and the offline embedding checks turn a raise into a 500. A torn refs/main used to decode into a nonsense commit and miss the snapshot dir; it now skips that cache root and keeps looking. _remove_pid_file runs first in _graceful_shutdown, so a corrupt studio.pid raising there abandoned the inference, export, training and tunnel children the rest of that function exists to kill. ssm_runtime's source-build path builds its subprocess kwargs in a dict and splats them through _run_with_heartbeat, so neither the encoding guard nor the earlier sweep saw the text = True in it: pip's output was still decoded with the Windows ANSI codepage, where a non-ASCII path or a compiler diagnostic mojibakes or raises over an install that was going fine. It now pins the same utf-8/replace pair install_wheel uses, and the HIP branch extends that env rather than replacing it. The guard learned the dict-literal shape and reddens on the old code (ssm_runtime.py:253). * Tighten the comments around the UTF-8 text I/O pins Collapse the multi-line rationales added with the encoding pins down to a line or two each, drop what the code already says, and use one wording for the repeated child-env note. * Do not let an unreadable bootstrap password stop startup, and narrow the kwargs guard ensure_default_admin calls _load_bootstrap_password for every existing admin and the lifespan calls that with no handler, so pinning the decode turned a damaged or pre-pin .bootstrap_password file into a backend that will not start. We write that file ourselves in UTF-8, so a byte that will not decode belongs to a file whose plaintext is worthless anyway; it now reads as no bootstrap password, the same answer as an absent file. A readable one still loads. The new kwargs check also judged every dict literal in the tree, so an unrelated payload carrying "text": True would have been reported as subprocess configuration with a misleading message, and a dict that fills in its encoding on a later line would have been reported too. It now only judges a dict that actually reaches a call, either splatted through a name or written at the call site, and treats a later kw["encoding"] assignment as satisfying it. The ssm_runtime shape it was written for is still caught, and a test pins both directions. * Stop reading a UTF-8 record a second time _read_line always parsed the line under the codepage as well, even when it had already read as UTF-8. Both callers take the UTF-8 reading when there is one and never look at the other, so on a healthy shard the second parse is pure waste, and this file reads all of one on every resume of a scrape it expects to reach gigabytes. Measured on 200,000 records, 76 MB: 1.96s before, 0.81s after, so the double reading was costing 2.8x. The early return is limited to a record, since the key lookup deliberately falls through to the codepage reading when UTF-8 yields something that is not one. A line UTF-8 cannot read still tries the codepage, latin-1 and the double-byte encodings as before, which is what the second reading is for. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Pin the scanned source fixture's line endings test_remote_code_scan_reads_non_ascii_sources compared a file's contents against the string it wrote, but wrote it in text mode, so Windows translated the line ends on the way out and the read back differed by a carriage return. That is the writer's doing, not the encoding the test is about, and it was the one failure on the Windows runner that belonged to this branch. The fixture now writes with newline = "" so the bytes on disk are the string on every platform. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Trim the newer comments to their point Shorten the widened-guard and state store notes added since the last pass, and collapse the line-ending note on the scanned source fixture. * Read the scraper checkpoint as UTF-8 only, never as a codepage A checkpoint holds nothing but base64 cursors and booleans, so one written by an older locale-encoded release is byte-identical to a UTF-8 one and already reads back. The codepage fallback can therefore only ever contribute non-ASCII: if a single-byte reading of the file were all ASCII, the UTF-8 read would have succeeded first. So the only file it changes the answer for is a damaged one, and there it turns a safe reset into a resume on a mojibaked cursor. GitHub answers that with INVALID_CURSOR_ARGUMENTS at HTTP 200, gh_client returns the partial document, and the scraper reads zero nodes and an empty pageInfo, which marks the stream done. Every later resume then skips it entirely. Reading UTF-8 only restores the earlier behaviour of dropping a checkpoint that will not decode, which re-scrapes from the first page while the writers dedup the replay. The shard scan below keeps its codepage reading; those records do carry non-ASCII. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Gate the remaining tilelang install tests to Linux _tilelang_platform_supported() returns False off Linux, so _ensure_tilelang_backend returns before the install and the subprocess mock these six assert on is never called. They fail on macOS runners for that reason alone. The rest of the file already carries this marker; these were missed. * Gate the Windows-incompatible worker and ROCm tests Two different gates, because the production code has two. The causal-conv1d and flash-linear-attention installers bail out on sys.platform == 'win32' alone and run everywhere else including macOS, so those cases get not_on_windows; marking them linux_only would skip tests that legitimately pass off Linux. The DRM and KFD readers return early unless platform.system() is Linux, and their fixtures build a fake sysfs tree needing PCI addresses like 0000:00:02.0 as directory names, which Windows cannot represent, so those get linux_only. The two visible-utilization cases failed for a different reason: on Windows get_visible_gpu_utilization takes the AMD adapter branch ahead of the torch fallback under test, and probing it imports torch, which the runner lacks. Stubbing that branch empty leaves every other platform unchanged. * Treat unparseable JSON nesting as a parse failure, and guard os.fdopen json.loads answers nesting it cannot descend with RecursionError, a RuntimeError, so _parse let it escape where the catch-all it replaced discarded the record. Both callers run _parse outside any further handler, so one damaged checkpoint or shard line aborted the scraper at startup. The encoding guard also missed os.fdopen, which is open() on a descriptor and takes the same locale default in text mode. It flags exactly the two text-mode calls that were left unencoded; the swap lock file's reader was already pinned to UTF-8 while its writer still used the codepage. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Write the non-ASCII source fixture without a 3.10-only argument Path.write_text() only grew newline in 3.10, and pyproject declares requires-python >=3.9, so this raised TypeError there. open() takes the same argument on every supported version and pins the bytes on disk the same way. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten encoding comments * Follow subprocess calls through callable aliases in the encoding guard --------- Co-authored-by: Unsloth <michaelhan@Michaels-MacBook-Pro.local> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen <unslothshared@gmail.com> --------- Co-authored-by: Unsloth <michaelhan@Michaels-MacBook-Pro.local> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen <unslothshared@gmail.com>
1202 lines
46 KiB
Python
1202 lines
46 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""
|
|
Inference subprocess entry point.
|
|
|
|
Each session runs in a persistent spawn subprocess, giving a clean interpreter
|
|
with no stale module state (solves transformers version-switching). It stays
|
|
alive while a model is loaded, taking commands (generate, load, unload) via
|
|
mp.Queue, and exits on shutdown or unload. Pattern follows core/training/worker.py.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import json
|
|
from loggers import get_logger
|
|
import os
|
|
import queue as _queue
|
|
import sys
|
|
import time
|
|
import traceback
|
|
from io import BytesIO
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
logger = get_logger(__name__)
|
|
from utils.hardware import apply_gpu_ids, is_apple_silicon
|
|
|
|
_SHARE_OBJECT_MAX_BYTES = 1 << 20
|
|
_SHARE_OBJECT_ERROR_SIZE = -1
|
|
|
|
# studio/backend root, prepended to sys.path so the spawned subprocess can
|
|
# import the utils/core packages.
|
|
_BACKEND_PATH = str(Path(__file__).resolve().parent.parent.parent)
|
|
|
|
|
|
def _ensure_backend_on_path() -> None:
|
|
if _BACKEND_PATH not in sys.path:
|
|
sys.path.insert(0, _BACKEND_PATH)
|
|
|
|
|
|
def _activate_transformers_version(model_name: str, hf_token: str | None = None) -> None:
|
|
"""Activate the correct transformers version BEFORE any ML imports."""
|
|
_ensure_backend_on_path()
|
|
|
|
from utils.transformers_version import activate_transformers_for_subprocess
|
|
|
|
activate_transformers_for_subprocess(model_name, hf_token)
|
|
|
|
|
|
def _decode_image(image_base64: str):
|
|
"""Decode base64 string to PIL.Image."""
|
|
from PIL import Image
|
|
|
|
image_data = base64.b64decode(image_base64)
|
|
return Image.open(BytesIO(image_data))
|
|
|
|
|
|
def _resize_image(img, max_size: int = 800):
|
|
"""Resize image while maintaining aspect ratio."""
|
|
if img is None:
|
|
return None
|
|
if img.size[0] > max_size or img.size[1] > max_size:
|
|
from PIL import Image
|
|
|
|
ratio = min(max_size / img.size[0], max_size / img.size[1])
|
|
new_size = (int(img.size[0] * ratio), int(img.size[1] * ratio))
|
|
return img.resize(new_size, Image.Resampling.LANCZOS)
|
|
return img
|
|
|
|
|
|
def _send_response(resp_queue: Any, response: dict) -> None:
|
|
"""Send a response to the parent process; stamps ``ts`` if absent."""
|
|
response.setdefault("ts", time.time())
|
|
try:
|
|
resp_queue.put(response)
|
|
except (OSError, ValueError) as exc:
|
|
logger.error("Failed to send response: %s", exc)
|
|
|
|
|
|
def _encode_share_object(obj: Any) -> bytes:
|
|
data = json.dumps(obj, separators = (",", ":"), ensure_ascii = False).encode("utf-8")
|
|
if len(data) > _SHARE_OBJECT_MAX_BYTES:
|
|
raise ValueError("Distributed object share payload is too large")
|
|
return data
|
|
|
|
|
|
def _decode_share_object(data: Any) -> Any:
|
|
return json.loads(bytes(data.tolist()).decode("utf-8"))
|
|
|
|
|
|
def _clean_token(value: str | None) -> str | None:
|
|
"""Normalize an HF token: blank or whitespace-only becomes None."""
|
|
return value if value and value.strip() else None
|
|
|
|
|
|
def _build_model_config(config: dict):
|
|
"""Build a ModelConfig from the config dict."""
|
|
from utils.models import ModelConfig
|
|
|
|
model_name = config["model_name"]
|
|
mc = ModelConfig.from_identifier(
|
|
model_id = model_name,
|
|
hf_token = _clean_token(config.get("hf_token")),
|
|
gguf_variant = config.get("gguf_variant"),
|
|
)
|
|
if not mc:
|
|
raise ValueError(f"Invalid model identifier: {model_name}")
|
|
return mc
|
|
|
|
|
|
_NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano")
|
|
|
|
|
|
def _needs_nemotron_trust(model_name: str, hf_token: str | None = None) -> bool:
|
|
"""Whether *model_name* is a NemotronH/Nano model that needs trust_remote_code.
|
|
|
|
NemotronH/Nano have config-parsing bugs that require it. Must NOT match
|
|
Llama-Nemotron (standard Llama arch), so also require the unsloth/ or nvidia/
|
|
namespace, and a genuine first-party Hub repo (not a local path or a spoof
|
|
name starting with "unsloth/"). The repo check is authenticated so private
|
|
first-party repos still resolve, and runs only after the cheap checks pass.
|
|
"""
|
|
mn = model_name.lower()
|
|
if not (
|
|
any(sub in mn for sub in _NEMOTRON_TRUST_SUBSTRINGS)
|
|
and (mn.startswith("unsloth/") or mn.startswith("nvidia/"))
|
|
):
|
|
return False
|
|
|
|
from utils.security.trusted_org import is_trusted_org_repo
|
|
|
|
return is_trusted_org_repo(model_name, hf_token = hf_token)
|
|
|
|
|
|
def _resolve_lora_4bit(mc, load_in_4bit: bool) -> bool:
|
|
"""Reconcile load_in_4bit with a LoRA adapter's recorded training method.
|
|
|
|
lora -> base is full precision (4bit off); qlora -> base is quantized (4bit
|
|
on); unknown method -> force off only when the base is not a -bnb-4bit repo.
|
|
A missing or unreadable adapter_config.json leaves the value unchanged.
|
|
"""
|
|
if not (mc.is_lora and mc.path):
|
|
return load_in_4bit
|
|
|
|
adapter_cfg_path = Path(mc.path) / "adapter_config.json"
|
|
if not adapter_cfg_path.exists():
|
|
return load_in_4bit
|
|
|
|
import json
|
|
|
|
try:
|
|
with open(adapter_cfg_path, encoding = "utf-8-sig") as f:
|
|
adapter_cfg = json.load(f)
|
|
training_method = adapter_cfg.get("unsloth_training_method")
|
|
if training_method == "lora" and load_in_4bit:
|
|
logger.info("adapter_config.json says lora — setting load_in_4bit=False")
|
|
return False
|
|
if training_method == "qlora" and not load_in_4bit:
|
|
logger.info("adapter_config.json says qlora — setting load_in_4bit=True")
|
|
return True
|
|
if (
|
|
not training_method
|
|
and mc.base_model
|
|
and "-bnb-4bit" not in mc.base_model.lower()
|
|
and load_in_4bit
|
|
):
|
|
logger.info(
|
|
"No training method, base model has no -bnb-4bit — setting load_in_4bit=False"
|
|
)
|
|
return False
|
|
except Exception as e:
|
|
logger.warning("Could not read adapter_config.json: %s", e)
|
|
return load_in_4bit
|
|
|
|
|
|
def _ensure_ssm_kernels(targets: list, resp_queue: Any) -> bool:
|
|
"""Install the SSM kernels the given model(s) lazy-import in from_pretrained; no-op for
|
|
non-SSM models, idempotent. Returns True on success; on a fatal mamba-ssm failure sends a
|
|
'loaded' failure response and returns False. Call BEFORE importing transformers, which
|
|
snapshots its optional-backend gates at import (a later install may not be picked up).
|
|
"""
|
|
try:
|
|
from utils.ssm_runtime import ensure_ssm_runtime
|
|
except Exception as exc:
|
|
logger.debug("ssm_runtime unavailable (%s); skipping SSM kernel pre-install", exc)
|
|
return True
|
|
|
|
_ssm_status = lambda m: _send_response(resp_queue, {"type": "status", "message": m})
|
|
try:
|
|
for ssm_target in dict.fromkeys(t for t in targets if t):
|
|
ensure_ssm_runtime(ssm_target, status_cb = _ssm_status)
|
|
return True
|
|
except Exception as exc:
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "loaded",
|
|
"success": False,
|
|
"message": (
|
|
f"This model needs SSM kernel libraries (causal-conv1d / "
|
|
f"mamba-ssm) that could not be installed: {exc}"
|
|
),
|
|
"error_kind": "ssm_runtime_install_failed",
|
|
},
|
|
)
|
|
return False
|
|
|
|
|
|
def _run_security_gates(
|
|
targets: list,
|
|
*,
|
|
trust_remote_code: bool,
|
|
hf_token: str | None,
|
|
approved_fingerprint: str | None,
|
|
resp_queue: Any,
|
|
compute_subdirs: bool = True,
|
|
subject: str | None = None,
|
|
) -> bool:
|
|
"""Malware + (when trust_remote_code) remote-code consent gates over *targets*
|
|
(model + base). Sends the matching 'loaded' failure and returns False if blocked; True
|
|
when every target is clear.
|
|
|
|
``compute_subdirs=False`` keeps the gate transformers-free (``security_load_subdirs``
|
|
imports ``model_config`` -> ``transformers``, which would snapshot optional-backend
|
|
availability before the SSM kernels are installed): used for the pre-import preflight,
|
|
where ``_handle_load`` re-runs the authoritative gate with full subdir scoping.
|
|
"""
|
|
targets = list(dict.fromkeys(t for t in targets if t))
|
|
|
|
# A poisoned pickle deserializes during from_pretrained even with trust_remote_code
|
|
# False, so check HF's security scan every load (for a LoRA, the base deserializes).
|
|
from utils.security import evaluate_file_security
|
|
|
|
if compute_subdirs:
|
|
from utils.security import security_load_subdirs
|
|
|
|
for target in targets:
|
|
_subdirs = security_load_subdirs(target, hf_token) if compute_subdirs else ()
|
|
_fs = evaluate_file_security(target, hf_token = hf_token, load_subdirs = _subdirs)
|
|
if _fs.blocked:
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "loaded",
|
|
"success": False,
|
|
"message": _fs.reason,
|
|
"error_kind": "malware_blocked",
|
|
"security": _fs.response_payload(),
|
|
},
|
|
)
|
|
return False
|
|
|
|
# Scan auto_map code before it runs; block CRITICAL/HIGH unless pinned-approved. Adapter
|
|
# and base are scanned as one unit, pinned by a single fingerprint.
|
|
if trust_remote_code:
|
|
from utils.security import evaluate_remote_code_consent_for_targets
|
|
_rc = evaluate_remote_code_consent_for_targets(
|
|
targets,
|
|
hf_token = hf_token,
|
|
trust_remote_code = True,
|
|
approved_fingerprint = approved_fingerprint,
|
|
subject = subject,
|
|
)
|
|
if _rc.blocked:
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "loaded",
|
|
"success": False,
|
|
"message": (
|
|
f"Model '{_rc.model_name}' ships custom code flagged as "
|
|
f"{_rc.max_severity} by the security scan. Review "
|
|
f"and approve it to proceed."
|
|
),
|
|
"error_kind": "remote_code_blocked",
|
|
"remote_code": _rc.response_payload(),
|
|
},
|
|
)
|
|
return False
|
|
|
|
return True
|
|
|
|
|
|
def _handle_load(backend, config: dict, resp_queue: Any) -> None:
|
|
"""Handle a load command: load a model into the backend."""
|
|
try:
|
|
mc = _build_model_config(config)
|
|
|
|
hf_token = _clean_token(config.get("hf_token"))
|
|
load_in_4bit = _resolve_lora_4bit(mc, config.get("load_in_4bit", True))
|
|
|
|
# Latest-transformers sidecar models load 16-bit: bnb 4-bit feeds quantized
|
|
# expert weights into unvalidated paths (e.g. grouped-MoE torch._grouped_mm).
|
|
if load_in_4bit:
|
|
from utils.transformers_version import latest_tier_active_for
|
|
if latest_tier_active_for(config["model_name"], hf_token):
|
|
load_in_4bit = False
|
|
logger.info(
|
|
"Latest-transformers sidecar active for %s - forcing a 16-bit "
|
|
"load (4-bit is disabled for brand-new architectures)",
|
|
config["model_name"],
|
|
)
|
|
|
|
trust_remote_code = config.get("trust_remote_code", False)
|
|
if not trust_remote_code and _needs_nemotron_trust(config["model_name"], hf_token = hf_token):
|
|
trust_remote_code = True
|
|
logger.info(
|
|
"Auto-enabled trust_remote_code for Nemotron model: %s", config["model_name"]
|
|
)
|
|
|
|
# Authoritative gates over the model + the LoRA base resolved via mc. Must run before
|
|
# the SSM install so a blocked model never triggers a native kernel build.
|
|
targets = [config["model_name"]]
|
|
if mc.is_lora and getattr(mc, "base_model", None):
|
|
targets.append(str(mc.base_model))
|
|
if not _run_security_gates(
|
|
targets,
|
|
trust_remote_code = trust_remote_code,
|
|
hf_token = hf_token,
|
|
approved_fingerprint = config.get("approved_remote_code_fingerprint"),
|
|
resp_queue = resp_queue,
|
|
subject = config.get("subject"),
|
|
):
|
|
return
|
|
|
|
# Install SSM/Mamba kernels: a no-op for the initial load (pre-installed before import)
|
|
# but still needed for a LoRA's base (resolved only now via mc) and in-process loads.
|
|
# Skip on MLX (no macOS wheel). Probe the base, not the adapter id / local path.
|
|
if getattr(backend, "device", None) != "mlx":
|
|
from utils.ssm_runtime import ssm_probe_identifier
|
|
|
|
_ssm_base = (
|
|
str(mc.base_model) if (mc.is_lora and getattr(mc, "base_model", None)) else None
|
|
)
|
|
ssm_targets = [ssm_probe_identifier(config["model_name"], _ssm_base)]
|
|
if not _ensure_ssm_kernels(ssm_targets, resp_queue):
|
|
return
|
|
|
|
# Heartbeat keeps the orchestrator's inactivity deadline alive during slow
|
|
# loads; a no-progress Xet download is reported as a stall so the parent
|
|
# can respawn over HTTP. Watch model + base repos (base is the LoRA
|
|
# download bottleneck).
|
|
from utils.hf_xet_fallback import start_watchdog
|
|
|
|
watch_repos = [mc.identifier]
|
|
base = getattr(mc, "base_model", None)
|
|
if base and str(base) != mc.identifier:
|
|
watch_repos.append(str(base))
|
|
|
|
heartbeat_stop = start_watchdog(
|
|
repo_ids = watch_repos,
|
|
on_stall = lambda msg: _send_response(resp_queue, {"type": "stall", "message": msg}),
|
|
on_heartbeat = lambda msg: _send_response(resp_queue, {"type": "status", "message": msg}),
|
|
xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1",
|
|
)
|
|
try:
|
|
load_kwargs = {
|
|
"config": mc,
|
|
"max_seq_length": config.get("max_seq_length", 2048),
|
|
"load_in_4bit": load_in_4bit,
|
|
"hf_token": hf_token,
|
|
"trust_remote_code": trust_remote_code,
|
|
"gpu_ids": config.get("resolved_gpu_ids"),
|
|
}
|
|
if getattr(backend, "device", None) == "mlx":
|
|
load_kwargs["parallel_mode"] = config.get("mlx_parallel_mode")
|
|
load_kwargs["distributed_group"] = config.get("_mlx_distributed_group")
|
|
success = backend.load_model(**load_kwargs)
|
|
finally:
|
|
heartbeat_stop.set()
|
|
|
|
if success:
|
|
model_info = {
|
|
"identifier": mc.identifier,
|
|
"display_name": mc.display_name,
|
|
"is_vision": mc.is_vision,
|
|
"is_lora": mc.is_lora,
|
|
"is_gguf": False,
|
|
# MLX backend sets device="mlx"; lets the UI tag MLX models.
|
|
"is_mlx": getattr(backend, "device", None) == "mlx",
|
|
"is_audio": getattr(mc, "is_audio", False),
|
|
"audio_type": getattr(mc, "audio_type", None),
|
|
"has_audio_input": getattr(mc, "has_audio_input", False),
|
|
}
|
|
_bm = getattr(backend, "models", {}) or {}
|
|
_entry = (
|
|
_bm.get(mc.identifier) or _bm.get(getattr(backend, "active_model_name", None)) or {}
|
|
)
|
|
try:
|
|
_context_length = _entry.get("context_length")
|
|
if _context_length is not None:
|
|
model_info["context_length"] = int(_context_length)
|
|
except Exception as _ctx_exc:
|
|
logger.warning("context_length forward failed: %s", _ctx_exc)
|
|
# Forward chat_template_info so the parent can classify capabilities.
|
|
try:
|
|
_tpl_info = _entry.get("chat_template_info")
|
|
if isinstance(_tpl_info, dict):
|
|
model_info["chat_template_info"] = {
|
|
"has_template": bool(_tpl_info.get("has_template", False)),
|
|
"template": _tpl_info.get("template"),
|
|
"format_type": _tpl_info.get("format_type", "generic"),
|
|
"template_name": _tpl_info.get("template_name"),
|
|
"special_tokens": _tpl_info.get("special_tokens", {}) or {},
|
|
}
|
|
except Exception as _tpl_exc:
|
|
logger.warning("chat_template_info forward failed: %s", _tpl_exc)
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "loaded",
|
|
"success": True,
|
|
"model_info": model_info,
|
|
},
|
|
)
|
|
else:
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "loaded",
|
|
"success": False,
|
|
"error": "Failed to load model",
|
|
},
|
|
)
|
|
|
|
except Exception as exc:
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "loaded",
|
|
"success": False,
|
|
"error": str(exc),
|
|
"stack": traceback.format_exc(limit = 20),
|
|
},
|
|
)
|
|
|
|
|
|
def _drain_skip_generate(cmd: dict, resp_queue: Any, drain_event) -> bool:
|
|
"""Skip a generate queued behind a cancelled one during an unload.
|
|
|
|
The parent sets ``drain_event`` for the whole unload. Because the parent's
|
|
per-token ``cancel_event`` is cleared at the start of every generate, a cancel
|
|
set while this generate was still queued would otherwise be lost when it is
|
|
dequeued. If the drain is in effect, emit an immediate (empty) ``gen_done`` so
|
|
the parent's stream/mailbox drains fast and the switch stays fast, and report
|
|
the generate was skipped so the caller does not clear the cancel or run it.
|
|
"""
|
|
if drain_event is None or not drain_event.is_set():
|
|
return False
|
|
request_id = cmd.get("request_id", "")
|
|
logger.info("Skipping generate for request %s: unload draining", request_id)
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "gen_done",
|
|
"request_id": request_id,
|
|
"cancelled": True,
|
|
"stats": None,
|
|
},
|
|
)
|
|
return True
|
|
|
|
|
|
def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None:
|
|
"""Handle a generate command: stream tokens back via resp_queue.
|
|
|
|
cancel_event is an mp.Event the parent can set anytime (user stop, or new
|
|
model load mid-generate); generation stops within 1-2 tokens.
|
|
"""
|
|
request_id = cmd.get("request_id", "")
|
|
|
|
try:
|
|
image = None
|
|
image_b64 = cmd.get("image_base64")
|
|
if image_b64:
|
|
image = _decode_image(image_b64)
|
|
image = _resize_image(image)
|
|
|
|
gen_kwargs = {
|
|
"messages": cmd["messages"],
|
|
"system_prompt": cmd.get("system_prompt", ""),
|
|
"image": image,
|
|
"temperature": cmd.get("temperature", 0.7),
|
|
"top_p": cmd.get("top_p", 0.9),
|
|
"top_k": cmd.get("top_k", 40),
|
|
"min_p": cmd.get("min_p", 0.0),
|
|
"max_new_tokens": cmd.get("max_new_tokens", 256),
|
|
"repetition_penalty": cmd.get("repetition_penalty", 1.0),
|
|
"presence_penalty": cmd.get("presence_penalty", 0.0),
|
|
"cancel_event": cancel_event,
|
|
}
|
|
|
|
# Forward only present optional keys so the backend signature can evolve.
|
|
for opt_key in (
|
|
"tools",
|
|
"enable_thinking",
|
|
"reasoning_effort",
|
|
"preserve_thinking",
|
|
):
|
|
if opt_key in cmd:
|
|
gen_kwargs[opt_key] = cmd[opt_key]
|
|
|
|
use_adapter = cmd.get("use_adapter")
|
|
if use_adapter is not None:
|
|
generator = backend.generate_with_adapter_control(
|
|
use_adapter = use_adapter,
|
|
**gen_kwargs,
|
|
)
|
|
else:
|
|
generator = backend.generate_chat_response(**gen_kwargs)
|
|
|
|
logger.info("Starting text generation for request_id=%s", request_id)
|
|
|
|
try:
|
|
for cumulative_text in generator:
|
|
# cancel_event is an mp.Event — checked instantly, no queue polling.
|
|
if cancel_event.is_set():
|
|
logger.info("Generation cancelled for request %s", request_id)
|
|
break
|
|
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "token",
|
|
"request_id": request_id,
|
|
"text": cumulative_text,
|
|
},
|
|
)
|
|
finally:
|
|
close = getattr(generator, "close", None)
|
|
if callable(close):
|
|
close()
|
|
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "gen_done",
|
|
"request_id": request_id,
|
|
# usage/timings from the MLX backend (None elsewhere).
|
|
"stats": getattr(backend, "last_generation_stats", None),
|
|
},
|
|
)
|
|
logger.info("Finished text generation for request_id=%s", request_id)
|
|
|
|
except Exception as exc:
|
|
logger.error("Generation error: %s", exc, exc_info = True)
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "gen_error",
|
|
"request_id": request_id,
|
|
"error": str(exc),
|
|
"stack": traceback.format_exc(limit = 20),
|
|
},
|
|
)
|
|
|
|
|
|
def _handle_share_object(backend, cmd: dict, resp_queue: Any) -> None:
|
|
"""Share a small Python object across MLX distributed ranks."""
|
|
request_id = cmd.get("request_id", "")
|
|
group = getattr(backend, "_distributed_group", None)
|
|
rank = int(getattr(backend, "_distributed_rank", 0) or 0)
|
|
world_size = int(getattr(backend, "_distributed_world_size", 1) or 1)
|
|
obj = cmd.get("object")
|
|
|
|
try:
|
|
if group is None or world_size <= 1:
|
|
shared = obj
|
|
else:
|
|
import mlx.core as mx
|
|
if rank == 0:
|
|
if obj is None:
|
|
mx.eval(mx.distributed.all_sum(mx.array(0), group = group))
|
|
shared = None
|
|
else:
|
|
try:
|
|
data = mx.array(_encode_share_object(obj), dtype = mx.uint8)
|
|
except Exception:
|
|
mx.eval(
|
|
mx.distributed.all_sum(
|
|
mx.array(_SHARE_OBJECT_ERROR_SIZE),
|
|
group = group,
|
|
)
|
|
)
|
|
raise
|
|
mx.eval(mx.distributed.all_sum(mx.array(data.size), group = group))
|
|
mx.eval(mx.distributed.all_sum(data, group = group))
|
|
shared = obj
|
|
else:
|
|
size = int(mx.distributed.all_sum(mx.array(0), group = group).item())
|
|
if size == _SHARE_OBJECT_ERROR_SIZE:
|
|
raise RuntimeError("Failed to share distributed object")
|
|
if size == 0:
|
|
shared = None
|
|
else:
|
|
data = mx.zeros(size, dtype = mx.uint8)
|
|
data = mx.distributed.all_sum(data, group = group)
|
|
shared = _decode_share_object(data)
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "shared",
|
|
"request_id": request_id,
|
|
"object": shared,
|
|
},
|
|
)
|
|
except Exception as exc:
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "share_error",
|
|
"request_id": request_id,
|
|
"error": str(exc),
|
|
"stack": traceback.format_exc(limit = 20),
|
|
},
|
|
)
|
|
|
|
|
|
def _handle_generate_audio(backend, cmd: dict, resp_queue: Any) -> None:
|
|
"""Handle TTS audio generation — returns WAV bytes + sample_rate."""
|
|
request_id = cmd.get("request_id", "")
|
|
try:
|
|
logger.info("Starting audio generation for request_id=%s", request_id)
|
|
wav_bytes, sample_rate = backend.generate_audio_response(
|
|
text = cmd["text"],
|
|
temperature = cmd.get("temperature", 0.6),
|
|
top_p = cmd.get("top_p", 0.95),
|
|
top_k = cmd.get("top_k", 50),
|
|
min_p = cmd.get("min_p", 0.0),
|
|
max_new_tokens = cmd.get("max_new_tokens", 2048),
|
|
repetition_penalty = cmd.get("repetition_penalty", 1.0),
|
|
use_adapter = cmd.get("use_adapter"),
|
|
)
|
|
|
|
# Send WAV bytes as base64 (bytes can't go through mp.Queue directly).
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "audio_done",
|
|
"request_id": request_id,
|
|
"wav_base64": base64.b64encode(wav_bytes).decode("ascii"),
|
|
"sample_rate": sample_rate,
|
|
},
|
|
)
|
|
logger.info("Finished audio generation for request_id=%s", request_id)
|
|
|
|
except Exception as exc:
|
|
logger.error("Audio generation error: %s", exc, exc_info = True)
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "audio_error",
|
|
"request_id": request_id,
|
|
"error": str(exc),
|
|
"stack": traceback.format_exc(limit = 20),
|
|
},
|
|
)
|
|
|
|
|
|
def _handle_generate_audio_input(backend, cmd: dict, resp_queue: Any, cancel_event) -> None:
|
|
"""Handle audio input generation (ASR/Whisper) — streams text tokens back."""
|
|
request_id = cmd.get("request_id", "")
|
|
|
|
try:
|
|
import numpy as np
|
|
|
|
# numpy arrays can't go through mp.Queue, so decode from list.
|
|
audio_array = np.array(cmd["audio_data"], dtype = np.float32)
|
|
|
|
audio_type = cmd.get("audio_type")
|
|
|
|
if audio_type == "whisper":
|
|
generator = backend.generate_whisper_response(
|
|
audio_array = audio_array,
|
|
cancel_event = cancel_event,
|
|
)
|
|
else:
|
|
generator = backend.generate_audio_input_response(
|
|
messages = cmd.get("messages", []),
|
|
system_prompt = cmd.get("system_prompt", ""),
|
|
audio_array = audio_array,
|
|
temperature = cmd.get("temperature", 0.7),
|
|
top_p = cmd.get("top_p", 0.9),
|
|
top_k = cmd.get("top_k", 40),
|
|
min_p = cmd.get("min_p", 0.0),
|
|
max_new_tokens = cmd.get("max_new_tokens", 512),
|
|
repetition_penalty = cmd.get("repetition_penalty", 1.0),
|
|
cancel_event = cancel_event,
|
|
)
|
|
|
|
logger.info("Starting audio input generation for request_id=%s", request_id)
|
|
|
|
for text_chunk in generator:
|
|
if cancel_event.is_set():
|
|
logger.info("Audio input generation cancelled for request %s", request_id)
|
|
break
|
|
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "token",
|
|
"request_id": request_id,
|
|
"text": text_chunk,
|
|
},
|
|
)
|
|
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "gen_done",
|
|
"request_id": request_id,
|
|
},
|
|
)
|
|
logger.info("Finished audio input generation for request_id=%s", request_id)
|
|
|
|
except Exception as exc:
|
|
logger.error("Audio input generation error: %s", exc, exc_info = True)
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "gen_error",
|
|
"request_id": request_id,
|
|
"error": str(exc),
|
|
"stack": traceback.format_exc(limit = 20),
|
|
},
|
|
)
|
|
|
|
|
|
def _handle_unload(backend, cmd: dict, resp_queue: Any) -> None:
|
|
"""Handle an unload command."""
|
|
model_name = cmd.get("model_name", "")
|
|
try:
|
|
if model_name and model_name in backend.models:
|
|
backend.unload_model(model_name)
|
|
elif backend.active_model_name:
|
|
backend.unload_model(backend.active_model_name)
|
|
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "unloaded",
|
|
"model_name": model_name,
|
|
},
|
|
)
|
|
except Exception as exc:
|
|
logger.error("Unload error: %s", exc)
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "unloaded",
|
|
"model_name": model_name,
|
|
"error": str(exc),
|
|
},
|
|
)
|
|
|
|
|
|
def run_inference_process(
|
|
*,
|
|
cmd_queue: Any,
|
|
resp_queue: Any,
|
|
cancel_event,
|
|
config: dict,
|
|
drain_event = None,
|
|
) -> None:
|
|
"""Subprocess entrypoint. Persistent — runs the command loop until shutdown.
|
|
|
|
Args:
|
|
cmd_queue: mp.Queue for receiving commands from parent.
|
|
resp_queue: mp.Queue for sending responses to parent.
|
|
cancel_event: mp.Event the parent sets to cancel generation.
|
|
config: Initial configuration dict with model info.
|
|
drain_event: mp.Event the parent sets for the duration of an unload. Unlike
|
|
cancel_event (cleared at the start of every generate), it is never cleared
|
|
here, so a generate still queued behind a cancelled one is skipped rather
|
|
than run — the cancel survives the queue handoff.
|
|
"""
|
|
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
|
os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports
|
|
|
|
if config.get("disable_xet"):
|
|
os.environ["HF_HUB_DISABLE_XET"] = "1"
|
|
logger.info("Xet transport disabled (HF_HUB_DISABLE_XET=1)")
|
|
|
|
import warnings
|
|
from loggers.config import LogConfig
|
|
|
|
if os.getenv("ENVIRONMENT_TYPE", "production") == "production":
|
|
warnings.filterwarnings("ignore")
|
|
|
|
LogConfig.setup_logging(
|
|
service_name = "unsloth-studio-inference-worker",
|
|
env = os.getenv("ENVIRONMENT_TYPE", "production"),
|
|
)
|
|
|
|
apply_gpu_ids(config.get("resolved_gpu_ids"), backend = config.get("device_backend"))
|
|
|
|
model_name = config["model_name"]
|
|
|
|
# ── 0. MLX fast-path — skip torch/transformers ──
|
|
_ensure_backend_on_path()
|
|
|
|
if is_apple_silicon():
|
|
# Non-fatal: fall through with the installed version, but log the cause
|
|
# instead of swallowing it (issue #6103).
|
|
try:
|
|
_activate_transformers_version(model_name, config.get("hf_token") or None)
|
|
except Exception as exc:
|
|
logger.warning(
|
|
"Failed to activate transformers version for '%s' (MLX inference); "
|
|
"inference may fail if this model requires a specific version. Error: %s",
|
|
model_name,
|
|
exc,
|
|
)
|
|
|
|
from utils.hardware import hardware as _hw
|
|
|
|
_hw.detect_hardware()
|
|
if _hw.DEVICE == _hw.DeviceType.MLX:
|
|
try:
|
|
from core.inference.mlx_inference import MLXInferenceBackend, _init_mlx_distributed
|
|
|
|
backend = MLXInferenceBackend()
|
|
if config.get("mlx_distributed"):
|
|
group, rank, size = _init_mlx_distributed()
|
|
config["_mlx_distributed_group"] = group
|
|
if size <= 1:
|
|
# A singleton group (MLX built without distributed support,
|
|
# or an invalid launch env/hostfile) would leave nonzero ranks
|
|
# looping forever on share_distributed_object. Fail the load
|
|
# instead of silently continuing without sharding.
|
|
raise RuntimeError(
|
|
"MLX distributed launch requested but initialized a singleton "
|
|
"group (size 1). Ensure the installed MLX has distributed "
|
|
"support and the launch environment/hostfile is valid, or run "
|
|
"without distributed."
|
|
)
|
|
logger.info(
|
|
"MLX distributed initialized in worker: rank=%s size=%s mode=%s",
|
|
rank,
|
|
size,
|
|
config.get("mlx_parallel_mode"),
|
|
)
|
|
_send_response(
|
|
resp_queue,
|
|
{"type": "status", "message": "Loading model..."},
|
|
)
|
|
_handle_load(backend, config, resp_queue)
|
|
except Exception as exc:
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "error",
|
|
"error": f"MLX inference init failed: {exc}",
|
|
"stack": traceback.format_exc(limit = 20),
|
|
},
|
|
)
|
|
return
|
|
|
|
# Enter the same command loop as the GPU path.
|
|
logger.info("MLX inference subprocess ready, entering command loop")
|
|
while True:
|
|
try:
|
|
cmd = cmd_queue.get(timeout = 1.0)
|
|
except _queue.Empty:
|
|
continue
|
|
except (EOFError, OSError):
|
|
return
|
|
if cmd is None:
|
|
continue
|
|
cmd_type = cmd.get("type", "")
|
|
try:
|
|
if cmd_type == "generate":
|
|
if _drain_skip_generate(cmd, resp_queue, drain_event):
|
|
continue
|
|
cancel_event.clear()
|
|
# Re-check the drain after clearing: the parent sets drain_event
|
|
# then cancel_event for an unload, so if that pair landed between
|
|
# the check above and this clear, the clear just erased the unload's
|
|
# cancel. Skip here so the outgoing model is not run to completion,
|
|
# which would stall the switch until the dispatcher idle-timeout.
|
|
if _drain_skip_generate(cmd, resp_queue, drain_event):
|
|
continue
|
|
_handle_generate(backend, cmd, resp_queue, cancel_event)
|
|
elif cmd_type == "share_object":
|
|
_handle_share_object(backend, cmd, resp_queue)
|
|
elif cmd_type == "load":
|
|
if backend.active_model_name:
|
|
backend.unload_model(backend.active_model_name)
|
|
_handle_load(backend, cmd, resp_queue)
|
|
elif cmd_type == "unload":
|
|
_handle_unload(backend, cmd, resp_queue)
|
|
elif cmd_type == "cancel":
|
|
cancel_event.set()
|
|
elif cmd_type == "reset":
|
|
cancel_event.set()
|
|
backend.reset_generation_state()
|
|
_send_response(resp_queue, {"type": "reset_ack"})
|
|
elif cmd_type == "status":
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "status_response",
|
|
"active_model": backend.active_model_name,
|
|
"models": {
|
|
k: {kk: vv for kk, vv in v.items() if kk != "model"}
|
|
for k, v in backend.models.items()
|
|
},
|
|
"loading": list(backend.loading_models),
|
|
},
|
|
)
|
|
elif cmd_type == "shutdown":
|
|
return
|
|
except Exception as exc:
|
|
logger.error("MLX command error (%s): %s", cmd_type, exc)
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "gen_error" if cmd_type == "generate" else "error",
|
|
"request_id": cmd.get("request_id"),
|
|
"error": str(exc),
|
|
"stack": traceback.format_exc(limit = 20),
|
|
},
|
|
)
|
|
return
|
|
|
|
# ── Windows: check Triton availability ──
|
|
# Placed ahead of the torchao stub below (which imports torch on win32 to detect ROCm),
|
|
# matching the training and export workers' gate-then-stub ordering.
|
|
if sys.platform == "win32":
|
|
try:
|
|
import triton # noqa: F401
|
|
logger.info("Triton available — torch.compile enabled")
|
|
except ImportError:
|
|
os.environ["TORCHDYNAMO_DISABLE"] = "1"
|
|
logger.warning(
|
|
"Triton not found on Windows — torch.compile disabled. "
|
|
'Install for better performance: pip install "triton-windows<3.7"'
|
|
)
|
|
|
|
# ── Stub torchao on Windows ROCm before ANY transformers import ──
|
|
# Must precede every path that pulls transformers, not just the ML imports in section 2:
|
|
# a local LoRA adapter with no recorded base reaches transformers here via
|
|
# _resolve_base_model -> utils.models. See core/_torchao_stub.py; no-op off Windows ROCm.
|
|
from core._torchao_stub import install_torchao_windows_rocm_stub
|
|
|
|
install_torchao_windows_rocm_stub()
|
|
|
|
# ── Resolve the effective base once, before activation/gates/install ──
|
|
# No ML import on the common path; a local adapter with no recorded base pulls
|
|
# transformers via utils.models, which is why the stub above precedes this.
|
|
# A remote LoRA's base is in its Hub adapter_config.json (else surfaced only by ModelConfig
|
|
# after import). _lora_base is set only for a genuine adapter, never a full fine-tune's base.
|
|
import json as _json
|
|
|
|
_ensure_backend_on_path()
|
|
from utils.transformers_version import _remote_lora_base, _resolve_base_model
|
|
|
|
_hf_token = _clean_token(config.get("hf_token"))
|
|
_lora_base = None
|
|
_local_adapter_cfg = Path(model_name) / "adapter_config.json"
|
|
if _local_adapter_cfg.is_file():
|
|
try:
|
|
_lora_base = (
|
|
_json.loads(_local_adapter_cfg.read_text(encoding = "utf-8-sig")).get(
|
|
"base_model_name_or_path"
|
|
)
|
|
or None
|
|
)
|
|
except Exception:
|
|
_lora_base = None
|
|
if not _lora_base:
|
|
_lora_base = _remote_lora_base(model_name, hf_token = _hf_token)
|
|
# Base for tier activation + the SSM-kernel heuristic: the LoRA base if any, else a full
|
|
# fine-tune's recorded base from config.json (its name reveals the SSM/sidecar arch).
|
|
_base = _lora_base or _resolve_base_model(model_name)
|
|
|
|
# ── 1. Activate transformers version (on the resolved base) BEFORE any ML imports ──
|
|
try:
|
|
_activate_transformers_version(_base, _hf_token)
|
|
except Exception as exc:
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "error",
|
|
"error": f"Failed to activate transformers version: {exc}",
|
|
"stack": traceback.format_exc(limit = 20),
|
|
},
|
|
)
|
|
return
|
|
|
|
# ── 1b. Security gates, then SSM/Mamba kernels, BEFORE importing transformers ──
|
|
# transformers snapshots its optional-backend gates at import, so a hybrid model's kernels
|
|
# must be installed before the import below ("mamba-ssm is required" otherwise). The gates
|
|
# are metadata-only, so run them first and refuse a blocked model before any native build.
|
|
# Gate only the model + a genuine LoRA base (matching _handle_load), never a full fine-tune's
|
|
# unloaded base; _handle_load re-runs the authoritative gates with the mc base.
|
|
_gate_targets = [model_name]
|
|
if _lora_base:
|
|
_gate_targets.append(_lora_base)
|
|
_trust_remote_code = config.get("trust_remote_code", False) or _needs_nemotron_trust(
|
|
model_name, hf_token = _hf_token
|
|
)
|
|
if not _run_security_gates(
|
|
_gate_targets,
|
|
trust_remote_code = _trust_remote_code,
|
|
hf_token = _hf_token,
|
|
approved_fingerprint = config.get("approved_remote_code_fingerprint"),
|
|
resp_queue = resp_queue,
|
|
compute_subdirs = False, # stay transformers-free until the SSM kernels are installed
|
|
subject = config.get("subject"),
|
|
):
|
|
return
|
|
# Probe the resolved base for SSM kernels, not the adapter id / local checkpoint path
|
|
# (arbitrary names must not match the SSM substrings).
|
|
from utils.ssm_runtime import ssm_probe_identifier
|
|
|
|
_ssm_targets = [ssm_probe_identifier(model_name, _base)]
|
|
if not _ensure_ssm_kernels(_ssm_targets, resp_queue):
|
|
return
|
|
|
|
# ── 2. Import ML libraries (fresh in this clean process) ──
|
|
try:
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "status",
|
|
"message": "Importing Unsloth...",
|
|
},
|
|
)
|
|
|
|
_ensure_backend_on_path()
|
|
|
|
# Recover from any namespace-package shadow before importing Unsloth.
|
|
from core.import_guards import ensure_real_packages
|
|
|
|
ensure_real_packages("unsloth_zoo", "unsloth")
|
|
|
|
from core.inference.inference import InferenceBackend
|
|
|
|
import transformers
|
|
|
|
logger.info("Subprocess loaded transformers %s", transformers.__version__)
|
|
|
|
except Exception as exc:
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "error",
|
|
"error": f"Failed to import ML libraries: {exc}",
|
|
"stack": traceback.format_exc(limit = 20),
|
|
},
|
|
)
|
|
return
|
|
|
|
# ── 3. Create inference backend and load initial model ──
|
|
try:
|
|
backend = InferenceBackend()
|
|
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "status",
|
|
"message": "Loading model...",
|
|
},
|
|
)
|
|
|
|
_handle_load(backend, config, resp_queue)
|
|
|
|
except Exception as exc:
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "error",
|
|
"error": f"Failed to initialize inference backend: {exc}",
|
|
"stack": traceback.format_exc(limit = 20),
|
|
},
|
|
)
|
|
return
|
|
|
|
# ── 4. Command loop — process commands until shutdown ──
|
|
# cancel_event is an mp.Event the parent can set anytime to cancel
|
|
# generation instantly (no queue polling needed).
|
|
logger.info("Inference subprocess ready, entering command loop")
|
|
|
|
while True:
|
|
try:
|
|
cmd = cmd_queue.get(timeout = 1.0)
|
|
except _queue.Empty:
|
|
continue
|
|
except (EOFError, OSError):
|
|
logger.info("Command queue closed, shutting down")
|
|
return
|
|
|
|
if cmd is None:
|
|
continue
|
|
|
|
cmd_type = cmd.get("type", "")
|
|
logger.info("Received command: %s", cmd_type)
|
|
|
|
try:
|
|
if cmd_type == "generate":
|
|
if _drain_skip_generate(cmd, resp_queue, drain_event):
|
|
continue
|
|
cancel_event.clear()
|
|
# Re-check the drain after clearing: the parent sets drain_event then
|
|
# cancel_event for an unload, so if that pair landed between the check
|
|
# above and this clear, the clear just erased the unload's cancel. Skip
|
|
# here so the outgoing model is not run to completion, which would stall
|
|
# the switch until the dispatcher idle-timeout tears the subprocess down.
|
|
if _drain_skip_generate(cmd, resp_queue, drain_event):
|
|
continue
|
|
_handle_generate(backend, cmd, resp_queue, cancel_event)
|
|
|
|
elif cmd_type == "share_object":
|
|
_handle_share_object(backend, cmd, resp_queue)
|
|
|
|
elif cmd_type == "load":
|
|
if backend.active_model_name:
|
|
backend.unload_model(backend.active_model_name)
|
|
_handle_load(backend, cmd, resp_queue)
|
|
|
|
elif cmd_type == "generate_audio":
|
|
cancel_event.clear()
|
|
_handle_generate_audio(backend, cmd, resp_queue)
|
|
|
|
elif cmd_type == "generate_audio_input":
|
|
cancel_event.clear()
|
|
_handle_generate_audio_input(backend, cmd, resp_queue, cancel_event)
|
|
|
|
elif cmd_type == "unload":
|
|
_handle_unload(backend, cmd, resp_queue)
|
|
|
|
elif cmd_type == "cancel":
|
|
# Redundant with mp.Event but handle gracefully.
|
|
cancel_event.set()
|
|
logger.info("Cancel command received")
|
|
|
|
elif cmd_type == "reset":
|
|
cancel_event.set()
|
|
backend.reset_generation_state()
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "reset_ack",
|
|
},
|
|
)
|
|
|
|
elif cmd_type == "status":
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "status_response",
|
|
"active_model": backend.active_model_name,
|
|
"models": {
|
|
name: {
|
|
"is_vision": info.get("is_vision", False),
|
|
"is_lora": info.get("is_lora", False),
|
|
"context_length": info.get("context_length"),
|
|
}
|
|
for name, info in backend.models.items()
|
|
},
|
|
"loading": list(backend.loading_models),
|
|
},
|
|
)
|
|
|
|
elif cmd_type == "shutdown":
|
|
logger.info("Shutdown command received, exiting")
|
|
for name in list(backend.models.keys()):
|
|
try:
|
|
backend.unload_model(name)
|
|
except Exception:
|
|
pass
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "shutdown_ack",
|
|
},
|
|
)
|
|
return
|
|
|
|
else:
|
|
logger.warning("Unknown command type: %s", cmd_type)
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "error",
|
|
"error": f"Unknown command type: {cmd_type}",
|
|
},
|
|
)
|
|
|
|
except Exception as exc:
|
|
logger.error("Error handling command '%s': %s", cmd_type, exc, exc_info = True)
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "error",
|
|
"error": f"Command '{cmd_type}' failed: {exc}",
|
|
"stack": traceback.format_exc(limit = 20),
|
|
},
|
|
)
|