From 4eabc74f3431e4af439a1b95dff07e9f0d4563a1 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Thu, 5 Mar 2026 17:47:57 +0000 Subject: [PATCH] feat: subprocess-based inference for transformers version switching Inference now runs in a persistent subprocess, solving the same transformers version-switching problem that was fixed for training. The subprocess stays alive between requests (model in GPU memory) and is only restarted when switching transformers versions. New files: - core/inference/worker.py: subprocess entry point with command loop - core/inference/orchestrator.py: parent-side proxy with same API Modified: - core/inference/__init__.py: exports orchestrator as default backend - routes/inference.py: removed in-process ensure_transformers_version() --- studio/backend/core/inference/__init__.py | 10 +- studio/backend/core/inference/orchestrator.py | 537 ++++++++++++++++++ studio/backend/core/inference/worker.py | 464 +++++++++++++++ studio/backend/routes/inference.py | 5 +- 4 files changed, 1012 insertions(+), 4 deletions(-) create mode 100644 studio/backend/core/inference/orchestrator.py create mode 100644 studio/backend/core/inference/worker.py diff --git a/studio/backend/core/inference/__init__.py b/studio/backend/core/inference/__init__.py index ff8b75d36a..6d44742d9c 100644 --- a/studio/backend/core/inference/__init__.py +++ b/studio/backend/core/inference/__init__.py @@ -1,11 +1,19 @@ """ Inference submodule - Inference backend for model loading and generation + +The default get_inference_backend() returns an InferenceOrchestrator that +delegates to a subprocess. The original InferenceBackend runs inside +the subprocess and can be imported directly from .inference when needed. """ -from .inference import InferenceBackend, get_inference_backend +from .orchestrator import InferenceOrchestrator, get_inference_backend from .llama_cpp import LlamaCppBackend +# Expose InferenceOrchestrator as InferenceBackend for backward compat +InferenceBackend = InferenceOrchestrator + __all__ = [ 'InferenceBackend', + 'InferenceOrchestrator', 'get_inference_backend', 'LlamaCppBackend', ] diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py new file mode 100644 index 0000000000..a3637cfb84 --- /dev/null +++ b/studio/backend/core/inference/orchestrator.py @@ -0,0 +1,537 @@ +""" +Inference orchestrator — subprocess-based. + +Provides the same API as InferenceBackend, but delegates all ML work +to a persistent subprocess. The subprocess is spawned on first model load +and stays alive for subsequent requests. + +When switching between models that need different transformers versions +(e.g. GLM-4.7-Flash needs 5.x, Qwen needs 4.57.x), the old subprocess +is killed and a new one is spawned with the correct version. + +Pattern follows core/training/training.py. +""" +import atexit +import base64 +import logging +import multiprocessing as mp +import queue +import threading +import time +import uuid +from io import BytesIO +from pathlib import Path +from typing import Any, Generator, Optional, Tuple, Union + +logger = logging.getLogger(__name__) + +_CTX = mp.get_context("spawn") + + +class InferenceOrchestrator: + """ + Inference backend orchestrator — subprocess-based. + + Exposes the same API surface as InferenceBackend so routes/inference.py + needs minimal changes. Internally, all heavy ML operations happen in + a persistent subprocess. + """ + + def __init__(self): + # Subprocess state + self._proc: Optional[mp.Process] = None + self._cmd_queue: Any = None + self._resp_queue: Any = None + self._lock = threading.Lock() + + # Local state mirrors (updated from subprocess responses) + self.active_model_name: Optional[str] = None + self.models: dict = {} + self.loading_models: set = set() + self.loaded_local_models: list = [] + self.default_models = [ + "unsloth/Qwen3-4B-Instruct-2507", + "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit", + "unsloth/Mistral-Nemo-Instruct-2407-bnb-4bit", + "unsloth/Phi-3.5-mini-instruct", + "unsloth/Gemma-3-4B-it", + "unsloth/Qwen2-VL-2B-Instruct-bnb-4bit", + ] + + # Version tracking for subprocess reuse + self._current_transformers_major: Optional[str] = None # "4" or "5" + + atexit.register(self._cleanup) + logger.info("InferenceOrchestrator initialized (subprocess mode)") + + # ------------------------------------------------------------------ + # Subprocess lifecycle + # ------------------------------------------------------------------ + + def _spawn_subprocess(self, config: dict) -> None: + """Spawn a new inference subprocess.""" + from .worker import run_inference_process + + self._cmd_queue = _CTX.Queue() + self._resp_queue = _CTX.Queue() + + self._proc = _CTX.Process( + target=run_inference_process, + kwargs={ + "cmd_queue": self._cmd_queue, + "resp_queue": self._resp_queue, + "config": config, + }, + daemon=True, + ) + self._proc.start() + logger.info("Inference subprocess started (pid=%s)", self._proc.pid) + + def _shutdown_subprocess(self, timeout: float = 10.0) -> None: + """Gracefully shut down the inference subprocess.""" + with self._lock: + if self._proc is None or not self._proc.is_alive(): + self._proc = None + return + + # Send shutdown command + try: + self._cmd_queue.put({"type": "shutdown"}) + except (OSError, ValueError): + pass + + # Wait for graceful shutdown + try: + self._proc.join(timeout=timeout) + except Exception: + pass + + # Force kill if still alive + if self._proc is not None and self._proc.is_alive(): + logger.warning("Inference subprocess did not exit gracefully, terminating") + try: + self._proc.terminate() + self._proc.join(timeout=5) + except Exception: + pass + if self._proc is not None and self._proc.is_alive(): + logger.warning("Subprocess still alive after terminate, killing") + try: + self._proc.kill() + self._proc.join(timeout=3) + except Exception: + pass + + self._proc = None + self._cmd_queue = None + self._resp_queue = None + logger.info("Inference subprocess shut down") + + def _cleanup(self): + """atexit handler.""" + self._shutdown_subprocess(timeout=5.0) + + def _ensure_subprocess_alive(self) -> bool: + """Check if subprocess is alive.""" + return self._proc is not None and self._proc.is_alive() + + # ------------------------------------------------------------------ + # Queue helpers + # ------------------------------------------------------------------ + + def _send_cmd(self, cmd: dict) -> None: + """Send a command to the subprocess.""" + if self._cmd_queue is None: + raise RuntimeError("No inference subprocess running") + try: + self._cmd_queue.put(cmd) + except (OSError, ValueError) as exc: + raise RuntimeError(f"Failed to send command to subprocess: {exc}") + + def _read_resp(self, timeout: float = 1.0) -> Optional[dict]: + """Read a response from the subprocess (non-blocking with timeout).""" + if self._resp_queue is None: + return None + try: + return self._resp_queue.get(timeout=timeout) + except queue.Empty: + return None + except (EOFError, OSError, ValueError): + return None + + def _wait_response( + self, expected_type: str, timeout: float = 120.0 + ) -> dict: + """Block until a response of the expected type arrives. + + Also handles 'status' and 'error' events during the wait. + Returns the matching response dict. + Raises RuntimeError on timeout or subprocess crash. + """ + deadline = time.monotonic() + timeout + + while time.monotonic() < deadline: + remaining = max(0.1, deadline - time.monotonic()) + resp = self._read_resp(timeout=min(remaining, 1.0)) + + if resp is None: + # Check subprocess health + if not self._ensure_subprocess_alive(): + raise RuntimeError("Inference subprocess crashed during wait") + continue + + rtype = resp.get("type", "") + + if rtype == expected_type: + return resp + + if rtype == "error": + error_msg = resp.get("error", "Unknown error") + raise RuntimeError(f"Subprocess error: {error_msg}") + + if rtype == "status": + logger.info("Subprocess status: %s", resp.get("message", "")) + continue + + # Other response types during wait — skip + logger.debug("Skipping response type '%s' while waiting for '%s'", rtype, expected_type) + + raise RuntimeError( + f"Timeout waiting for '{expected_type}' response after {timeout}s" + ) + + def _drain_queue(self) -> list: + """Drain all pending responses.""" + events = [] + if self._resp_queue is None: + return events + while True: + try: + events.append(self._resp_queue.get_nowait()) + except queue.Empty: + return events + except (EOFError, OSError, ValueError): + return events + + # ------------------------------------------------------------------ + # Public API — same interface as InferenceBackend + # ------------------------------------------------------------------ + + def load_model( + self, + config, # ModelConfig + max_seq_length: int = 2048, + dtype=None, + load_in_4bit: bool = True, + hf_token: Optional[str] = None, + ) -> bool: + """Load a model for inference. + + May spawn a new subprocess if no subprocess exists or if the model + needs a different transformers version than what's currently running. + """ + from utils.transformers_version import needs_transformers_5 + + model_name = config.identifier + self.loading_models.add(model_name) + + try: + needed_major = "5" if needs_transformers_5(model_name) else "4" + project_root = str(Path(__file__).resolve().parent.parent.parent.parent.parent) + + # Build config dict for subprocess + sub_config = { + "project_root": project_root, + "model_name": model_name, + "max_seq_length": max_seq_length, + "load_in_4bit": load_in_4bit, + "hf_token": hf_token or "", + "gguf_variant": getattr(config, "gguf_variant", None), + } + + if self._ensure_subprocess_alive(): + if needed_major == self._current_transformers_major: + # Reuse existing subprocess — send load command + logger.info( + "Reusing inference subprocess for '%s' (transformers %s.x)", + model_name, needed_major, + ) + self._send_cmd({"type": "load", **sub_config}) + resp = self._wait_response("loaded", timeout=180) + else: + # Version mismatch — kill and respawn + logger.info( + "Transformers version mismatch (have %s.x, need %s.x) — " + "restarting subprocess for '%s'", + self._current_transformers_major, needed_major, model_name, + ) + self._shutdown_subprocess() + self._spawn_subprocess(sub_config) + resp = self._wait_response("loaded", timeout=180) + else: + # No subprocess running — spawn new + logger.info( + "Spawning inference subprocess for '%s' (transformers %s.x)", + model_name, needed_major, + ) + if self._proc is not None: + # Dead subprocess — clean up + self._shutdown_subprocess(timeout=2) + self._spawn_subprocess(sub_config) + resp = self._wait_response("loaded", timeout=180) + + # Update local state from response + if resp.get("success"): + self._current_transformers_major = needed_major + model_info = resp.get("model_info", {}) + self.active_model_name = model_info.get("identifier", model_name) + self.models[self.active_model_name] = { + "is_vision": model_info.get("is_vision", False), + "is_lora": model_info.get("is_lora", False), + "display_name": model_info.get("display_name", model_name), + } + self.loading_models.discard(model_name) + logger.info("Model '%s' loaded successfully in subprocess", model_name) + return True + else: + error = resp.get("error", "Failed to load model") + self.loading_models.discard(model_name) + raise Exception(error) + + except Exception: + self.loading_models.discard(model_name) + raise + + def unload_model(self, model_name: str) -> bool: + """Unload a model from the subprocess.""" + if not self._ensure_subprocess_alive(): + # No subprocess — just clear local state + self.models.pop(model_name, None) + if self.active_model_name == model_name: + self.active_model_name = None + return True + + try: + self._send_cmd({ + "type": "unload", + "model_name": model_name, + }) + resp = self._wait_response("unloaded", timeout=30) + + # Update local state + self.models.pop(model_name, None) + if self.active_model_name == model_name: + self.active_model_name = None + + logger.info("Model '%s' unloaded from subprocess", model_name) + return True + + except Exception as exc: + logger.error("Error unloading model '%s': %s", model_name, exc) + # Clear local state anyway + self.models.pop(model_name, None) + if self.active_model_name == model_name: + self.active_model_name = None + return False + + def generate_chat_response( + self, + messages: list, + system_prompt: str = "", + image=None, + temperature: float = 0.7, + top_p: float = 0.9, + top_k: int = 40, + min_p: float = 0.0, + max_new_tokens: int = 256, + repetition_penalty: float = 1.1, + cancel_event=None, + ) -> Generator[str, None, None]: + """Generate response, streaming tokens from subprocess.""" + yield from self._generate_inner( + messages=messages, + system_prompt=system_prompt, + image=image, + temperature=temperature, + top_p=top_p, + top_k=top_k, + min_p=min_p, + max_new_tokens=max_new_tokens, + repetition_penalty=repetition_penalty, + cancel_event=cancel_event, + use_adapter=None, + ) + + def generate_with_adapter_control( + self, + use_adapter: Optional[Union[bool, str]] = None, + cancel_event=None, + **gen_kwargs, + ) -> Generator[str, None, None]: + """Generate with adapter control, streaming tokens from subprocess.""" + yield from self._generate_inner( + use_adapter=use_adapter, + cancel_event=cancel_event, + **gen_kwargs, + ) + + def _generate_inner( + self, + messages: list = None, + system_prompt: str = "", + image=None, + temperature: float = 0.7, + top_p: float = 0.9, + top_k: int = 40, + min_p: float = 0.0, + max_new_tokens: int = 256, + repetition_penalty: float = 1.1, + cancel_event=None, + use_adapter=None, + ) -> Generator[str, None, None]: + """Inner generation logic — sends command to subprocess, yields tokens.""" + if not self._ensure_subprocess_alive(): + yield "Error: Inference subprocess is not running" + return + + if not self.active_model_name: + yield "Error: No active model" + return + + request_id = str(uuid.uuid4()) + + # Convert PIL Image to base64 if needed + image_b64 = None + if image is not None: + image_b64 = self._pil_to_base64(image) + + cmd = { + "type": "generate", + "request_id": request_id, + "messages": messages or [], + "system_prompt": system_prompt, + "image_base64": image_b64, + "temperature": temperature, + "top_p": top_p, + "top_k": top_k, + "min_p": min_p, + "max_new_tokens": max_new_tokens, + "repetition_penalty": repetition_penalty, + } + + if use_adapter is not None: + cmd["use_adapter"] = use_adapter + + try: + self._send_cmd(cmd) + except RuntimeError as exc: + yield f"Error: {exc}" + return + + # Yield tokens from response queue + while True: + resp = self._read_resp(timeout=30.0) + + if resp is None: + # Check subprocess health + if not self._ensure_subprocess_alive(): + yield "Error: Inference subprocess crashed during generation" + return + continue + + # Skip responses for other requests + resp_rid = resp.get("request_id") + rtype = resp.get("type", "") + + # Status messages don't have request_id + if rtype == "status": + continue + + # Error without request_id = subprocess-level error + if rtype == "error" and not resp_rid: + yield f"Error: {resp.get('error', 'Unknown error')}" + return + + # Skip responses for other requests + if resp_rid and resp_rid != request_id: + continue + + if rtype == "token": + # Check cancel from route + if cancel_event is not None and cancel_event.is_set(): + try: + self._send_cmd({"type": "cancel", "request_id": request_id}) + except RuntimeError: + pass + return + yield resp.get("text", "") + + elif rtype == "gen_done": + return + + elif rtype == "gen_error": + yield f"Error: {resp.get('error', 'Unknown error')}" + return + + def reset_generation_state(self): + """Send cancel/reset to subprocess.""" + if not self._ensure_subprocess_alive(): + return + try: + self._send_cmd({"type": "reset"}) + except RuntimeError: + pass + + # ------------------------------------------------------------------ + # Local helpers (no subprocess needed) + # ------------------------------------------------------------------ + + def resize_image(self, img, max_size: int = 800): + """Resize image while maintaining aspect ratio. + No ML imports needed — runs locally in parent process. + """ + 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 + + @staticmethod + def _pil_to_base64(img) -> str: + """Convert a PIL Image to base64 string for IPC.""" + buf = BytesIO() + img.save(buf, format="PNG") + return base64.b64encode(buf.getvalue()).decode("ascii") + + def get_current_model(self) -> Optional[str]: + """Get currently active model name.""" + return self.active_model_name + + def is_model_loading(self) -> bool: + """Check if any model is currently loading.""" + return len(self.loading_models) > 0 + + def get_loading_model(self) -> Optional[str]: + """Get name of currently loading model.""" + return next(iter(self.loading_models)) if self.loading_models else None + + def check_vision_model_compatibility(self) -> bool: + """Check if current model supports vision.""" + if self.active_model_name and self.active_model_name in self.models: + return self.models[self.active_model_name].get("is_vision", False) + return False + + +# ========== GLOBAL INSTANCE ========== +_inference_backend = None + + +def get_inference_backend() -> InferenceOrchestrator: + """Get global inference backend instance (orchestrator).""" + global _inference_backend + if _inference_backend is None: + _inference_backend = InferenceOrchestrator() + return _inference_backend diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py new file mode 100644 index 0000000000..6fcc4546be --- /dev/null +++ b/studio/backend/core/inference/worker.py @@ -0,0 +1,464 @@ +""" +Inference subprocess entry point. + +Each inference session runs in a persistent subprocess (mp.get_context("spawn")). +This gives us a clean Python interpreter with no stale module state — +solving the transformers version-switching problem completely. + +The subprocess stays alive while a model is loaded, accepting commands +(generate, load, unload) via mp.Queue. It exits on shutdown or unload. + +Pattern follows core/training/worker.py. +""" +from __future__ import annotations + +import base64 +import logging +import os +import queue as _queue +import sys +import threading +import time +import traceback +from io import BytesIO +from typing import Any + +logger = logging.getLogger(__name__) + + +def _activate_transformers_version(model_name: str, project_root: str) -> None: + """Activate the correct transformers version BEFORE any ML imports. + + If the model needs transformers 5.x, prepend the pre-installed .venv_t5/ + directory to sys.path. Otherwise do nothing (default 4.57.x in .venv/). + """ + # Ensure backend is on path for utils imports + backend_path = os.path.join(project_root, "studio", "backend") + if backend_path not in sys.path: + sys.path.insert(0, backend_path) + + from utils.transformers_version import needs_transformers_5, _resolve_base_model + + resolved = _resolve_base_model(model_name) + if needs_transformers_5(resolved): + venv_t5 = os.path.join(project_root, ".venv_t5") + if os.path.isdir(venv_t5): + sys.path.insert(0, venv_t5) + logger.info("Activated transformers 5.x from %s", venv_t5) + else: + # Fallback: pip install at runtime (slower, ~10-15s) + logger.warning(".venv_t5 not found at %s — installing at runtime", venv_t5) + import subprocess as sp + os.makedirs(venv_t5, exist_ok=True) + sp.run( + [sys.executable, "-m", "pip", "install", "--target", venv_t5, + "--no-deps", "transformers==5.1.0"], + stdout=sp.PIPE, stderr=sp.STDOUT, + ) + sp.run( + [sys.executable, "-m", "pip", "install", "--target", venv_t5, + "--no-deps", "huggingface_hub>=1.3.0"], + stdout=sp.PIPE, stderr=sp.STDOUT, + ) + if os.path.isdir(venv_t5): + sys.path.insert(0, venv_t5) + else: + logger.info("Using default transformers (4.57.x) for %s", 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 _check_cancel(cmd_queue: Any, cancel_event: threading.Event) -> None: + """Non-blocking poll of cmd_queue for cancel commands.""" + try: + cmd = cmd_queue.get_nowait() + if cmd and cmd.get("type") == "cancel": + cancel_event.set() + logger.info("Cancel signal received for request %s", cmd.get("request_id", "*")) + except _queue.Empty: + pass + except (EOFError, OSError): + pass + + +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 _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) + + 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, + ) + + 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, + } + _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, + cmd_queue: Any, + cancel_event: threading.Event, +) -> None: + """Handle a generate command: stream tokens back via resp_queue.""" + request_id = cmd.get("request_id", "") + + try: + # Decode image if provided + image = None + image_b64 = cmd.get("image_base64") + if image_b64: + image = _decode_image(image_b64) + image = _resize_image(image) + + # Build generation kwargs + 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.1), + "cancel_event": cancel_event, + } + + # Choose generation path + 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) + + for cumulative_text in generator: + # Check for cancel commands (non-blocking) + _check_cancel(cmd_queue, cancel_event) + if cancel_event.is_set(): + 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, + "ts": time.time(), + }) + + 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_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, + config: dict, +) -> None: + """Subprocess entrypoint. Persistent — runs command loop until shutdown. + + Args: + cmd_queue: mp.Queue for receiving commands from parent. + resp_queue: mp.Queue for sending responses to parent. + config: Initial configuration dict with model info and project_root. + """ + os.environ["TOKENIZERS_PARALLELISM"] = "false" + + project_root = config["project_root"] + model_name = config["model_name"] + + # ── 1. Activate correct transformers version BEFORE any ML imports ── + try: + _activate_transformers_version(model_name, project_root) + 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 + + # ── 2. Import ML libraries (fresh in this clean process) ── + try: + _send_response(resp_queue, { + "type": "status", + "message": "Importing ML libraries...", + "ts": time.time(), + }) + + backend_path = os.path.join(project_root, "studio", "backend") + 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 = threading.Event() + 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, cmd_queue, cancel_event) + + elif cmd_type == "load": + # Load a new model (reusing this subprocess) + # First unload current model + 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() + logger.info("Cancel signal 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": + # Return current 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), + } + 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") + # Unload all models + 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(), + }) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 7e1cff7a7d..521ff79d9a 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -86,9 +86,8 @@ async def load_model( GGUF models are loaded via llama-server (llama.cpp) instead of Unsloth. """ try: - # Ensure correct transformers version for this model architecture - from utils.transformers_version import ensure_transformers_version - ensure_transformers_version(request.model_path) + # Version switching is handled automatically by the subprocess-based + # inference backend — no need for ensure_transformers_version() here. # Create config using clean factory method # is_lora is auto-detected from adapter_config.json on disk/HF