cef7dcf160* Studio: log transformers version-switching decisions and stop swallowing MLX activation failures Two logging gaps in dynamic transformers version switching (issue #6103): 1. get_transformers_tier returned a tier with no trace of why. Add an info log at each decision point naming the model and the trigger (which substring matched, or which config check fired), so a model landing on the wrong tier is diagnosable. 2. The MLX fast-path in run_training_process activated the transformers version inside a bare 'except Exception: pass', silently swallowing failures while the non-MLX path reports them. A missing or broken version venv (e.g. Gemma-4 needing 5.5.0) left no trace and only a confusing downstream crash. Extract a small _activate_transformers_version_or_warn helper that logs a warning on failure while keeping the non-fatal fall-through, and call it from the MLX path. Adds tier-selection logging tests and helper warn/silent tests. * Studio: clarify path-prepend log, warn on venv version mismatch, log per-package install progress Completes the remaining logging items of #6103 in studio/backend/utils/transformers_version.py: - activate_transformers_for_subprocess: the early "Activated transformers X.X.X" line was misleading because at that point only the venv directory has been prepended to sys.path, not imported. It now says it prepended the venv to sys.path and notes the loaded version is confirmed later by "Subprocess loaded transformers ...". - _venv_dir_is_valid: a detected version mismatch is logged at warning instead of info, since it immediately triggers a full venv wipe and reinstall that should be visible in the logs. - _ensure_venv_dir: log each package as it starts installing with an N/M progress counter, so a slow runtime install is not mistaken for a hang (pip/uv output is piped and only surfaced on error). Adds tests covering all three behaviours; pre-existing unused imports are left untouched. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: make tier log-capture tests independent of import order The new issue #6103 caplog assertions in test_transformers_version.py relied on the module-level sys.modules.setdefault("loggers", stub) winning the import race. In a full backend pytest run another module (for example test_log_filter_no_truncation, collected earlier) imports the real loggers first, so the setdefault is a no-op and transformers_version.logger becomes a structlog/stdout logger that caplog cannot capture -- the tier, activation, venv-mismatch and install-progress log assertions then fail even though the line was emitted. Bind a real stdlib logger to transformers_version.logger for the duration of each test via an autouse fixture, so the module logs through logging and caplog captures them regardless of collection order. * Studio: log local checkpoint tier decisions and warn on MLX inference activation - get_transformers_tier: the local config.json fast path returned a tier without logging it, so local checkpoints stayed opaque while HF ids were traceable. Log each decision there too, with a caplog regression test. - inference worker: the MLX path swallowed _activate_transformers_version failures with a bare except, the same gap issue #6103 fixed for training. Warn instead, keeping the non-fatal fall-through. --------- Co-authored-by: Daniel Han <michaelhan2050@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
949 lines
34 KiB
Python
949 lines
34 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 structlog
|
|
from loggers import get_logger
|
|
import os
|
|
import queue as _queue
|
|
import sys
|
|
import threading
|
|
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
|
|
|
|
|
|
def _activate_transformers_version(model_name: str) -> None:
|
|
"""Activate the correct transformers version BEFORE any ML imports."""
|
|
# Ensure backend is on path for utils imports.
|
|
backend_path = str(Path(__file__).resolve().parent.parent.parent)
|
|
if backend_path not in sys.path:
|
|
sys.path.insert(0, backend_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."""
|
|
try:
|
|
resp_queue.put(response)
|
|
except (OSError, ValueError) as exc:
|
|
logger.error("Failed to send response: %s", exc)
|
|
|
|
|
|
def _build_model_config(config: dict):
|
|
"""Build a ModelConfig from the config dict."""
|
|
from utils.models import ModelConfig
|
|
|
|
model_name = config["model_name"]
|
|
hf_token = config.get("hf_token")
|
|
hf_token = hf_token if hf_token and hf_token.strip() else None
|
|
gguf_variant = config.get("gguf_variant")
|
|
|
|
mc = ModelConfig.from_identifier(
|
|
model_id = model_name,
|
|
hf_token = hf_token,
|
|
gguf_variant = gguf_variant,
|
|
)
|
|
if not mc:
|
|
raise ValueError(f"Invalid model identifier: {model_name}")
|
|
return mc
|
|
|
|
|
|
def _get_hf_download_state(model_names: list[str] | None = None) -> tuple[int, bool] | None:
|
|
"""Return (total_bytes, has_incomplete) for the HF Hub cache, or None on error.
|
|
|
|
With *model_names*, only those models' ``blobs/`` dirs are checked (faster);
|
|
accepts multiple names so LoRA loads can watch adapter + base repos at once.
|
|
*has_incomplete* is True when any ``*.incomplete`` files exist (download
|
|
active). None means state could not be determined, so callers skip stall logic.
|
|
"""
|
|
try:
|
|
from huggingface_hub.constants import HF_HUB_CACHE
|
|
|
|
cache = Path(HF_HUB_CACHE)
|
|
if not cache.exists():
|
|
return (0, False)
|
|
|
|
total = 0
|
|
has_incomplete = False
|
|
blobs_dirs: list[Path] = []
|
|
|
|
if model_names:
|
|
from utils.paths import resolve_cached_repo_id_case
|
|
for name in model_names:
|
|
if not name:
|
|
continue
|
|
# Skip local filesystem paths -- HF IDs (org/model) never start
|
|
# with / . ~ or contain backslashes.
|
|
if name.startswith(("/", ".", "~")) or "\\" in name:
|
|
continue
|
|
name = resolve_cached_repo_id_case(name)
|
|
# HF cache dir format: models--org--name (slashes -> --).
|
|
cache_dir_name = "models--" + name.replace("/", "--")
|
|
blobs_dir = cache / cache_dir_name / "blobs"
|
|
if blobs_dir.exists():
|
|
blobs_dirs.append(blobs_dir)
|
|
else:
|
|
blobs_dirs = list(cache.glob("models--*/blobs"))
|
|
|
|
for bdir in blobs_dirs:
|
|
for f in bdir.iterdir():
|
|
try:
|
|
if f.is_file():
|
|
total += f.stat().st_size
|
|
if f.name.endswith(".incomplete"):
|
|
has_incomplete = True
|
|
except OSError:
|
|
pass
|
|
|
|
return (total, has_incomplete)
|
|
except Exception as e:
|
|
logger.debug("Failed to determine HF download state: %s", e)
|
|
return None
|
|
|
|
|
|
def _start_heartbeat(
|
|
resp_queue: Any,
|
|
interval: float = 30.0,
|
|
stall_timeout: float = 180.0,
|
|
xet_disabled: bool = False,
|
|
model_names: list[str] | None = None,
|
|
) -> threading.Event:
|
|
"""Start a daemon thread that sends periodic status heartbeats.
|
|
|
|
A stall is reported only when ``*.incomplete`` files are present (download
|
|
active) AND cache size hasn't changed for *stall_timeout* seconds. When the
|
|
download finishes the timer resets, so post-download init (quantization, GPU
|
|
weight load) isn't misclassified as a stall. Returns a stop event.
|
|
"""
|
|
stop = threading.Event()
|
|
transport = "https" if xet_disabled else "xet"
|
|
|
|
def _beat():
|
|
state = _get_hf_download_state(model_names)
|
|
last_size = state[0] if state is not None else 0
|
|
last_change = time.monotonic()
|
|
|
|
while not stop.wait(interval):
|
|
state = _get_hf_download_state(model_names)
|
|
now = time.monotonic()
|
|
|
|
# Skip stall logic if we cannot measure the cache.
|
|
if state is None:
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "status",
|
|
"message": f"Loading model ({transport} transport)...",
|
|
"ts": time.time(),
|
|
},
|
|
)
|
|
continue
|
|
|
|
current_size, has_incomplete = state
|
|
|
|
if current_size != last_size:
|
|
last_size = current_size
|
|
last_change = now
|
|
|
|
# Only fire stall while .incomplete files confirm an active download;
|
|
# reset the timer otherwise so model init isn't counted as a stall.
|
|
if not has_incomplete:
|
|
last_change = now
|
|
elif now - last_change >= stall_timeout:
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "stall",
|
|
"message": (
|
|
f"Download appears stalled ({transport} transport) "
|
|
f"-- no progress for {int(now - last_change)}s"
|
|
),
|
|
"ts": time.time(),
|
|
},
|
|
)
|
|
# fire once -- the orchestrator will kill us
|
|
return
|
|
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "status",
|
|
"message": f"Loading model ({transport} transport)...",
|
|
"ts": time.time(),
|
|
},
|
|
)
|
|
|
|
t = threading.Thread(target = _beat, daemon = True)
|
|
t.start()
|
|
return stop
|
|
|
|
|
|
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 = config.get("hf_token")
|
|
hf_token = hf_token if hf_token and hf_token.strip() else None
|
|
|
|
# Auto-detect quantization for LoRA adapters.
|
|
load_in_4bit = config.get("load_in_4bit", True)
|
|
if mc.is_lora and mc.path:
|
|
import json
|
|
from pathlib import Path
|
|
|
|
adapter_cfg_path = Path(mc.path) / "adapter_config.json"
|
|
if adapter_cfg_path.exists():
|
|
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")
|
|
load_in_4bit = False
|
|
elif training_method == "qlora" and not load_in_4bit:
|
|
logger.info("adapter_config.json says qlora — setting load_in_4bit=True")
|
|
load_in_4bit = True
|
|
elif not training_method:
|
|
if (
|
|
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"
|
|
)
|
|
load_in_4bit = False
|
|
except Exception as e:
|
|
logger.warning("Could not read adapter_config.json: %s", e)
|
|
|
|
# Auto-enable trust_remote_code only for NemotronH/Nano (config parsing
|
|
# bugs require it). Must NOT match Llama-Nemotron (standard Llama arch).
|
|
_NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano")
|
|
trust_remote_code = config.get("trust_remote_code", False)
|
|
if not trust_remote_code:
|
|
model_name = config["model_name"]
|
|
_mn_lower = model_name.lower()
|
|
if any(sub in _mn_lower for sub in _NEMOTRON_TRUST_SUBSTRINGS) and (
|
|
_mn_lower.startswith("unsloth/") or _mn_lower.startswith("nvidia/")
|
|
):
|
|
trust_remote_code = True
|
|
logger.info(
|
|
"Auto-enabled trust_remote_code for Nemotron model: %s",
|
|
model_name,
|
|
)
|
|
|
|
# Heartbeat every 30s so the orchestrator knows we're alive during slow loads.
|
|
xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1"
|
|
|
|
# Watch model + base repos (base download is the LoRA bottleneck).
|
|
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_heartbeat(
|
|
resp_queue,
|
|
interval = 30.0,
|
|
xet_disabled = xet_disabled,
|
|
model_names = watch_repos,
|
|
)
|
|
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:
|
|
# Build model_info for the parent to mirror.
|
|
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),
|
|
}
|
|
try:
|
|
_bm = getattr(backend, "models", {}) or {}
|
|
_entry = (
|
|
_bm.get(mc.identifier)
|
|
or _bm.get(getattr(backend, "active_model_name", None))
|
|
or {}
|
|
)
|
|
_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:
|
|
_bm = getattr(backend, "models", {}) or {}
|
|
_entry = (
|
|
_bm.get(mc.identifier)
|
|
or _bm.get(getattr(backend, "active_model_name", None))
|
|
or {}
|
|
)
|
|
_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,
|
|
"ts": time.time(),
|
|
},
|
|
)
|
|
else:
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "loaded",
|
|
"success": False,
|
|
"error": "Failed to load model",
|
|
"ts": time.time(),
|
|
},
|
|
)
|
|
|
|
except Exception as exc:
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "loaded",
|
|
"success": False,
|
|
"error": str(exc),
|
|
"stack": traceback.format_exc(limit = 20),
|
|
"ts": time.time(),
|
|
},
|
|
)
|
|
|
|
|
|
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,
|
|
"ts": time.time(),
|
|
},
|
|
)
|
|
|
|
_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),
|
|
"ts": time.time(),
|
|
},
|
|
)
|
|
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),
|
|
"ts": time.time(),
|
|
},
|
|
)
|
|
|
|
|
|
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,
|
|
"ts": time.time(),
|
|
},
|
|
)
|
|
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),
|
|
"ts": time.time(),
|
|
},
|
|
)
|
|
|
|
|
|
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,
|
|
"ts": time.time(),
|
|
},
|
|
)
|
|
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "gen_done",
|
|
"request_id": request_id,
|
|
"ts": time.time(),
|
|
},
|
|
)
|
|
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),
|
|
"ts": time.time(),
|
|
},
|
|
)
|
|
|
|
|
|
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,
|
|
"ts": time.time(),
|
|
},
|
|
)
|
|
except Exception as exc:
|
|
logger.error("Unload error: %s", exc)
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "unloaded",
|
|
"model_name": model_name,
|
|
"error": str(exc),
|
|
"ts": time.time(),
|
|
},
|
|
)
|
|
|
|
|
|
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 ──
|
|
backend_path = str(Path(__file__).resolve().parent.parent.parent)
|
|
if backend_path not in sys.path:
|
|
sys.path.insert(0, backend_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...", "ts": time.time()},
|
|
)
|
|
_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),
|
|
"ts": time.time(),
|
|
},
|
|
)
|
|
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", "ts": time.time()})
|
|
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),
|
|
"ts": time.time(),
|
|
},
|
|
)
|
|
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),
|
|
"ts": time.time(),
|
|
},
|
|
)
|
|
return
|
|
|
|
# ── 1. Activate transformers version BEFORE any ML imports ──
|
|
try:
|
|
_activate_transformers_version(model_name)
|
|
except Exception as exc:
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "error",
|
|
"error": f"Failed to activate transformers version: {exc}",
|
|
"stack": traceback.format_exc(limit = 20),
|
|
"ts": time.time(),
|
|
},
|
|
)
|
|
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"'
|
|
)
|
|
|
|
# ── 2. Import ML libraries (fresh in this clean process) ──
|
|
try:
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "status",
|
|
"message": "Importing Unsloth...",
|
|
"ts": time.time(),
|
|
},
|
|
)
|
|
|
|
backend_path = str(Path(__file__).resolve().parent.parent.parent)
|
|
if backend_path not in sys.path:
|
|
sys.path.insert(0, backend_path)
|
|
|
|
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),
|
|
"ts": time.time(),
|
|
},
|
|
)
|
|
return
|
|
|
|
# ── 3. Create inference backend and load initial model ──
|
|
try:
|
|
backend = InferenceBackend()
|
|
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "status",
|
|
"message": "Loading model...",
|
|
"ts": time.time(),
|
|
},
|
|
)
|
|
|
|
_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),
|
|
"ts": time.time(),
|
|
},
|
|
)
|
|
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":
|
|
# Unload the current model before loading the new one.
|
|
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",
|
|
"ts": time.time(),
|
|
},
|
|
)
|
|
|
|
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),
|
|
"ts": time.time(),
|
|
},
|
|
)
|
|
|
|
elif cmd_type == "shutdown":
|
|
logger.info("Shutdown command received, exiting")
|
|
for model_name in list(backend.models.keys()):
|
|
try:
|
|
backend.unload_model(model_name)
|
|
except Exception:
|
|
pass
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "shutdown_ack",
|
|
"ts": time.time(),
|
|
},
|
|
)
|
|
return
|
|
|
|
else:
|
|
logger.warning("Unknown command type: %s", cmd_type)
|
|
_send_response(
|
|
resp_queue,
|
|
{
|
|
"type": "error",
|
|
"error": f"Unknown command type: {cmd_type}",
|
|
"ts": time.time(),
|
|
},
|
|
)
|
|
|
|
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),
|
|
"ts": time.time(),
|
|
},
|
|
)
|