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()
This commit is contained in:
Roland Tannous
parent
1e04149ddf
commit
4eabc74f34
4 changed files with 1012 additions and 4 deletions
|
|
@ -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',
|
||||
]
|
||||
|
|
|
|||
537
studio/backend/core/inference/orchestrator.py
Normal file
537
studio/backend/core/inference/orchestrator.py
Normal file
|
|
@ -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
|
||||
464
studio/backend/core/inference/worker.py
Normal file
464
studio/backend/core/inference/worker.py
Normal file
|
|
@ -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(),
|
||||
})
|
||||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue