2a05426adb* Auto-install SSM kernels (causal-conv1d, mamba-ssm) for inference loads Mamba/SSM hybrids (Nemotron-H/Nano, Falcon-H1, Granite-4.0-H, ...) lazily import mamba_ssm / causal_conv1d during from_pretrained, so loading them for chat failed with 'mamba-ssm is required by the Mamba model but cannot be imported'. The training worker already wheel-first installs these before a fine-tune; the inference worker did not. Add utils/ssm_runtime.ensure_ssm_runtime and call it from the inference load path so the same models load for inference. Training worker is untouched; a drift test keeps the shared detection and pinned versions in lockstep. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * ssm_runtime: invalidate import caches, skip MLX, cover LoRA base - Invalidate importlib finder caches in _is_importable and after a successful wheel install, so a kernel installed earlier in this same process is actually importable when the modeling code lazy-imports it during from_pretrained. - Skip the SSM kernel install entirely on the MLX (Apple Silicon) load path: these are CUDA/ROCm Torch kernels with no MLX use and no macOS prebuilt wheel, so the source build would fail before the MLX backend loads the model. - For LoRA loads, also run detection over the resolved base model, since an adapter id like 'me/my-lora' won't match the SSM heuristics but its SSM base (Nemotron-H, ...) is what needs the kernels. Adds tests for cache invalidation and the MLX-skip / LoRA-base worker wiring. * Tighten SSM autoinstall comments * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * ssm_runtime: verify wheel imports, HIP-aware source build, build heartbeat Address review feedback: - Verify a prebuilt wheel actually imports before trusting it; a CUDA/ABI-mismatched wheel now falls back to a source build instead of returning success and failing later with the cryptic lazy-import error. - HIP-aware source build: require hipcc on ROCm, inject clang --gcc-install-dir, and use the 1800s timeout, mirroring the training worker (ROCm has no prebuilt wheel). - Emit a status heartbeat every 60s during the source build so a long (ROCm) build does not trip the orchestrator's 300s inactivity timeout. Tests cover the wheel-not-importable fallback and the missing-hipcc ROCm bail. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Make causal-conv1d best-effort and harden the SSM source build - causal-conv1d is a fast path: models that merely want it (Qwen3-Next, LFM2) fall back to torch, so a failed install must not reject an otherwise loadable chat model on Windows/CPU/macOS or an ABI without a wheel. Only a true SSM model's mamba-ssm requirement stays fatal, matching the training worker which treats causal-conv1d as best-effort. - The source build is reached only when not importable, including a wheel that installed but failed to import; add --reinstall/--force-reinstall so it replaces the broken install instead of no-opping as already satisfied. - Add --no-cache to the ROCm uv source build to avoid reusing stale artifacts from a partial HIP build, mirroring the training worker. * Address review: install SSM kernels before transformers, harden import + Windows Codex: - Install the SSM kernels before importing transformers. run_inference_process imported core.inference.inference (which imports unsloth/transformers) before the load, and a sidecar transformers can evaluate its optional-backend gates against the import state; installing causal_conv1d/mamba_ssm afterwards left those gates unsatisfied and a Nemotron/Falcon/Granite load still failed with "mamba-ssm is required". The initial model's kernels are now installed in run_inference_process before the ML import, via a shared _ensure_ssm_kernels helper; _handle_load keeps calling it (idempotent) for a LoRA's base and for later in-process loads. - _is_importable now treats any import failure as "not importable", not only ImportError. An ABI-incompatible native kernel (undefined symbol after a torch/CUDA upgrade) raises OSError/RuntimeError; letting those escape reported ssm_runtime_install_failed instead of falling back to reinstall/source build. - Skip causal-conv1d on Windows (no prebuilt wheel), mirroring the training worker. A causal-conv1d-only model (Qwen3-Next/LFM2) no longer drops a chat load into a multi-minute untimed source build; it uses the torch fallback. mamba-ssm is still attempted for true SSM hybrids. Tests: test_ssm_runtime.py +5 (broken-kernel exceptions read as not-importable; causal-conv1d skipped on win32 while mamba-ssm still installs). 36 passed. * Trim comments to be more succinct * Run security gates before installing SSM kernels The SSM kernel auto-install is name-based (model_is_ssm is a substring match, no config fetch), so a model id merely containing an SSM substring triggered a native-package install (possibly a slow source build) before the malware and remote-code consent gates ran. Extract those gates into _run_security_gates and call it before the kernel install in both the pre-import path of run_inference_process and in _handle_load, so a blocked or nonexistent model is refused before any build. The gates are metadata-only and do not import transformers, so they are safe to run before the pre-import install. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Resolve remote LoRA bases before importing transformers _resolve_base_model only reads a local adapter_config.json, so a remote LoRA adapter whose own id has no SSM substring but whose base is a Nemotron/Falcon/ Granite model had its base discovered only by ModelConfig in _handle_load, after transformers was imported and its optional-backend availability snapshotted, so the SSM kernel install there was too late. Add _remote_lora_base, a metadata-only adapter_config.json fetch (no huggingface_hub / transformers import), and use it in the pre-import path so the base is gated and its kernels pre-installed. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Gate only loaded roots, tier on the resolved base, read offline LoRA cache Three follow-ups to the pre-import resolution: - The security gate reused the SSM target list, which for a local full fine-tune includes the config.json-recorded base. That base is never loaded, so scanning it could falsely block a safe local checkpoint. Gate only the model plus a genuine LoRA base (matching _handle_load's mc.is_lora), separate from the broader SSM-install list. - Tier activation ran on the raw adapter id, so a remote LoRA whose base needs a sidecar transformers version imported the default and failed. Resolve the base once up front and activate on it. - _remote_lora_base bailed on offline before checking the hub cache, missing a cached adapter's base. Read the cached adapter_config.json when offline or when the fetch fails. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Keep the pre-import gate transformers-free; harden remote LoRA resolution The pre-import security gate called security_load_subdirs, which imports model_config and thus transformers, snapshotting optional-backend availability before the SSM kernels are installed and defeating the ordering. Add compute_subdirs to _run_security_gates and pass False in the preflight so it scans from the root only (transformers-free); _handle_load still runs the authoritative gate with full subdir scoping after the import. _remote_lora_base now skips existing local relative paths (is_local_path) so a checkpoint like outputs/run1 is never treated as a Hub repo, and distinguishes a definitive 404 (not a LoRA -> None) from transient/offline failures (read the cache), so a repo that is now a full model no longer resolves a stale cached base. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Probe a real model id for SSM kernels; respect HF_ENDPOINT model_is_ssm is a substring match, so an arbitrary name could false-match and force a mamba-ssm install that fails the load for a non-SSM model: - a LoRA adapter id like user/falcon-h1-lora (the SSM-relevant code is the base's); - a local checkpoint under an SSM-named parent dir, e.g. /runs/falcon-h1/llama-ckpt. Add ssm_probe_identifier, which resolves the base (or a bare local checkpoint's basename) and feed that to ensure_ssm_runtime from both the pre-import path and _handle_load, so detection runs against a real model id, never an adapter id or parent folders. _remote_lora_base now honors HF_ENDPOINT so enterprise/mirror deployments resolve the adapter base instead of always hitting huggingface.co. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten comments in the pre-import SSM gate/install path --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han <michaelhan2050@gmail.com>
1003 lines
37 KiB
Python
1003 lines
37 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
|
|
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
|
|
|
|
# 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) -> 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)
|
|
|
|
|
|
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 _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) 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,
|
|
) -> 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,
|
|
)
|
|
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))
|
|
|
|
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,
|
|
):
|
|
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:
|
|
success = backend.load_model(
|
|
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"),
|
|
)
|
|
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 _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),
|
|
"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)
|
|
|
|
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,
|
|
},
|
|
)
|
|
|
|
_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_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) -> 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.
|
|
"""
|
|
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"))
|
|
|
|
model_name = config["model_name"]
|
|
|
|
# ── 0. MLX fast-path — skip torch/transformers ──
|
|
_ensure_backend_on_path()
|
|
|
|
from utils.hardware import hardware as _hw
|
|
|
|
_hw.detect_hardware()
|
|
if _hw.DEVICE == _hw.DeviceType.MLX:
|
|
# Non-fatal: fall through with the installed version, but log the cause
|
|
# instead of swallowing it (issue #6103).
|
|
try:
|
|
_activate_transformers_version(model_name)
|
|
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,
|
|
)
|
|
try:
|
|
from core.inference.mlx_inference import MLXInferenceBackend
|
|
|
|
backend = MLXInferenceBackend()
|
|
_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":
|
|
cancel_event.clear()
|
|
_handle_generate(backend, cmd, resp_queue, cancel_event)
|
|
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
|
|
|
|
# ── Resolve the effective base once, before activation/gates/install (no ML import) ──
|
|
# 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()).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)
|
|
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. Windows: check Triton availability (must precede import torch) ──
|
|
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"'
|
|
)
|
|
|
|
# ── 1c. 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
|
|
):
|
|
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":
|
|
cancel_event.clear()
|
|
_handle_generate(backend, cmd, resp_queue, cancel_event)
|
|
|
|
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),
|
|
},
|
|
)
|