"""Backends: what reads the backbone. Every backend gives, for a prompt's token ids, the hidden state at the last position after the final norm (float32, one vector). The package's readout table and calibration go on top in model.py, so the answers are the package's whichever backend reads the prompt. torch transformers; every input, pictures included (always installed) NVIDIA GPU, CPU, Apple (mps) vllm vLLM's pooling runner, LAST pooling with no normalization NVIDIA GPU openvino OpenVINO, bf16, one IR per static length (python -m decisionnode_edge.build) x86 CPU mlx mlx-lm's Qwen3.5 text model, reading the model directory as it is Apple silicon llamacpp llama.cpp's llama-server (--embeddings --pooling last) on a GGUF Q8_0 file any machine llama.cpp builds for The fast backends read text. A prompt one of them cannot take (pictures; a prompt longer than an OpenVINO build's largest length or a llama.cpp context) is read by torch, loaded on first use. Choosing: `backend` (EdgeModel.load) or DECISIONNODE_EDGE_BACKEND; "auto" (the default) takes the first that can run here, in the order of `ORDER` for the device class. A backend that needs built files (openvino, llamacpp) runs only with the directory the build tool wrote (`artifacts`, or DECISIONNODE_EDGE_ARTIFACTS). Options are flat, or under a backend's name ({"openvino": {"threads": 16}}) so one set can serve whichever backend auto takes. """ from __future__ import annotations import dataclasses import hashlib import importlib.util import json import logging import os import platform import shutil import socket import subprocess import sys import tempfile import time import urllib.error import urllib.request import weakref from typing import Any, Sequence NAMES = ("torch", "vllm", "openvino", "mlx", "llamacpp") # per device class, fastest first (the runtime README's table) ORDER = { "cuda": ("vllm", "llamacpp", "torch"), "apple": ("mlx", "llamacpp", "torch"), "cpu": ("openvino", "llamacpp", "torch"), } ENV_BACKEND = "DECISIONNODE_EDGE_BACKEND" ENV_ARTIFACTS = "DECISIONNODE_EDGE_ARTIFACTS" ENV_OPTIONS = "DECISIONNODE_EDGE_OPTIONS" BUILD_FILE = "decisionnode_edge_build.json" BUILD_FORMAT = "decisionnode-edge-build-1" # what a backend that reads the model directory needs from it: never readout.safetensors (vLLM loads every .safetensors file) MODEL_FILES = ("config.json", "model.safetensors", "tokenizer.json", "tokenizer_config.json", "preprocessor_config.json", "video_preprocessor_config.json") log = logging.getLogger("decisionnode_edge") # --------------------------------------------------------------------------- # Choosing # --------------------------------------------------------------------------- def apple_silicon() -> bool: return sys.platform == "darwin" and platform.machine() == "arm64" def linux() -> bool: return sys.platform.startswith("linux") def device_class(device: str | None) -> str: """"cuda", "apple" or "cpu": the device given, else CUDA when present, else Apple silicon, else the CPU.""" if device is not None: device = str(device) return "cuda" if device.startswith("cuda") else "apple" if device.startswith("mps") else "cpu" import torch if torch.cuda.is_available(): return "cuda" return "apple" if apple_silicon() else "cpu" def env_options() -> dict: raw = os.environ.get(ENV_OPTIONS) if not raw: return {} value = json.loads(raw) if not isinstance(value, dict): raise ValueError(f"{ENV_OPTIONS} must be a JSON object") return value def read_build(artifacts: str | None) -> dict | None: """The build tool's record in `artifacts` (decisionnode_edge_build.json), or None.""" if not artifacts or not os.path.isfile(os.path.join(artifacts, BUILD_FILE)): return None with open(os.path.join(artifacts, BUILD_FILE), encoding="utf-8") as f: build = json.load(f) if build.get("format") != BUILD_FORMAT: raise ValueError(f"{artifacts}: not a build of this runtime ({build.get('format')!r})") return build def own_options(name: str, options: dict) -> dict: """The options for one backend: the flat ones, then the ones under its name.""" return {**{k: v for k, v in options.items() if k not in NAMES}, **(options.get(name) or {})} def _installed(*modules: str) -> bool: return all(importlib.util.find_spec(m) is not None for m in modules) def unavailable(name: str, device: str | None, artifacts: str | None, options: dict) -> str | None: """Why `name` cannot run here, or None when it can.""" if name == "torch": return None if name == "vllm": if not linux() or not _installed("vllm"): return "vLLM is not installed (pip install 'decisionnode-edge[vllm]', Linux with an NVIDIA GPU)" return None if device_class(device) == "cuda" else "vLLM needs an NVIDIA GPU" if name == "mlx": if not apple_silicon(): return "MLX needs Apple silicon" return None if _installed("mlx", "mlx_lm") else "MLX is not installed (pip install 'decisionnode-edge[mlx]')" build = read_build(artifacts) if name == "openvino": if not _installed("openvino"): return "OpenVINO is not installed (pip install 'decisionnode-edge[openvino]')" if not build or build.get("backend") != "openvino": return "OpenVINO needs a build: python -m decisionnode_edge.build openvino --out , then artifacts=" return None if name == "llamacpp": options = own_options(name, options) if options.get("url"): return None if not build or build.get("backend") != "llamacpp": return "llama.cpp needs a GGUF build (python -m decisionnode_edge.build gguf) as artifacts, or options={'url': }" if not (options.get("server") or shutil.which("llama-server")): return "llama-server is not on PATH: give options={'server': } or {'url': }" return None raise ValueError(f"unknown backend {name!r}: one of {', '.join(NAMES)} or auto") def choose(requested: str | None, device: str | None, artifacts: str | None, options: dict) -> str: """The backend to load: the one asked for (argument, else DECISIONNODE_EDGE_BACKEND), or for "auto" the first of the device class's order that can run here.""" name = (requested or os.environ.get(ENV_BACKEND) or "auto").strip().lower() if name != "auto": reason = unavailable(name, device, artifacts, options) if reason: raise RuntimeError(f"backend {name!r}: {reason}") return name for candidate in ORDER[device_class(device)]: reason = unavailable(candidate, device, artifacts, options) if reason is None: return candidate log.info("auto: not %s: %s", candidate, reason) return "torch" def load(name: str, path: str, config: dict, artifacts: str | None, options: dict, device: str | None): """The backend object for a name other than torch (model.py loads torch itself).""" max_tokens = int(config["limits"]["max_input_tokens"]) options = own_options(name, options) if name == "vllm": return VLLM(path, max_tokens, options, device) if name == "mlx": return MLX(path, options) build = read_build(artifacts) if build is not None and build.get("model_fingerprint") != fingerprint(path): raise ValueError(f"{artifacts} was built from another model than {path}") if name == "openvino": return OpenVINO(artifacts, build, options) if name == "llamacpp": return LlamaCpp(artifacts, build, max_tokens, options) raise ValueError(f"unknown backend {name!r}") def fingerprint(path: str) -> str: """A quick identity of a model directory's weights, kept in every build: the file's size, its safetensors header and 1 MiB every 256 MiB plus the last MiB (hashing all of it would add seconds to every load).""" weights = os.path.join(path, "model.safetensors") size = os.path.getsize(weights) h = hashlib.sha256(str(size).encode()) with open(weights, "rb") as f: n = int.from_bytes(f.read(8), "little") h.update(f.read(n)) for offset in list(range(8 + n, size, 256 << 20)) + [max(0, size - (1 << 20))]: f.seek(offset) h.update(f.read(1 << 20)) return h.hexdigest() def link_model_files(path: str, target: str) -> None: """Links in `target` to the model directory's backbone files (MODEL_FILES), never its readout.""" for name in MODEL_FILES: src = os.path.join(os.path.abspath(path), name) if os.path.exists(src): os.symlink(src, os.path.join(target, name)) # --------------------------------------------------------------------------- # Backends: hidden(ids) gives the float32 hidden state of the last position; accepts(n) says whether a prompt of n # tokens is theirs (torch reads the others) # --------------------------------------------------------------------------- class VLLM: """vLLM's pooling runner (convert "embed", LAST pooling, use_activation False: the raw hidden state after the final norm) on a temporary directory of links to the model files. Text only: the multimodal limits are 0. Its engine runs in this process (VLLM_ENABLE_V1_MULTIPROCESSING 0 unless set): as fast at batch 1, and a child process would need the caller's script to be importable (vLLM starts it with spawn once CUDA is up).""" name = "vllm" KV_CACHE_BYTES = 2 << 30 # several prompts of the full 16,384 tokens HEADROOM_BYTES = 4 << 30 # activations at 16,384 tokens, CUDA graphs, the CUDA context def __init__(self, path: str, max_tokens: int, options: dict, device: str | None): if device is not None and str(device) not in ("cuda", "cuda:0"): raise ValueError("vLLM runs on the first visible GPU: choose the card with CUDA_VISIBLE_DEVICES") os.environ.setdefault("VLLM_ENABLE_V1_MULTIPROCESSING", "0") from vllm import LLM from vllm.config import PoolerConfig from vllm.engine.arg_utils import EngineArgs from vllm.platforms import current_platform self.dir = tempfile.mkdtemp(prefix="decisionnode-edge-vllm-") self._cleanup = weakref.finalize(self, shutil.rmtree, self.dir, True) link_model_files(path, self.dir) # the KV cache is sized, not a share of the card: the share only has to cover it, the weights and the headroom total = current_platform.get_device_total_memory(0) need = os.path.getsize(os.path.join(path, "model.safetensors")) + self.KV_CACHE_BYTES + self.HEADROOM_BYTES args: dict[str, Any] = { "model": self.dir, "runner": "pooling", "convert": "embed", "pooler_config": PoolerConfig(pooling_type="LAST", use_activation=False), "dtype": "bfloat16", "max_model_len": max_tokens, "kv_cache_memory_bytes": self.KV_CACHE_BYTES, "gpu_memory_utilization": round(min(0.92, need / total), 3), "enable_prefix_caching": False, "limit_mm_per_prompt": {"image": 0, "video": 0}, "seed": 0, "disable_log_stats": True, } fields = {f.name for f in dataclasses.fields(EngineArgs)} args.update({k: v for k, v in options.items() if k in fields}) # any vLLM engine argument, as given if set(options) - fields: log.warning("options vLLM does not know, not used: %s", sorted(set(options) - fields)) self.llm = LLM(**args) def accepts(self, n: int) -> bool: return True def hidden(self, ids: Sequence[int]): import torch out = self.llm.embed([{"prompt_token_ids": list(ids)}], use_tqdm=False)[0] return torch.tensor(out.outputs.embedding, dtype=torch.float32) class OpenVINO: """OpenVINO on the CPU, one compiled model per static length of the build. A prompt is padded on the right to the smallest length that holds it and the row of its last token is read: the backbone is causal (attention, the short convolution and the delta rule), so the padding never reaches the earlier rows. The weights are float32 in the IR (the bf16 weights, exactly) and run with OpenVINO's bf16 inference precision.""" name = "openvino" def __init__(self, artifacts: str, build: dict, options: dict): import openvino as ov config = {"PERFORMANCE_HINT": "LATENCY", "INFERENCE_PRECISION_HINT": options.get("precision", "bf16")} if options.get("threads"): config["INFERENCE_NUM_THREADS"] = int(options["threads"]) config.update(options.get("config") or {}) core = ov.Core() self.pad = int(build["pad_token_id"]) self.runners = {} for length, file in sorted(((int(L), f) for L, f in build["files"].items())): compiled = core.compile_model(os.path.join(artifacts, file), "CPU", config) self.runners[length] = (compiled.create_infer_request(), compiled.output(0)) self.lengths = sorted(self.runners) def accepts(self, n: int) -> bool: return n <= self.lengths[-1] def hidden(self, ids: Sequence[int]): import numpy as np import torch n = len(ids) length = next(L for L in self.lengths if L >= n) request, output = self.runners[length] x = np.full((1, length), self.pad, dtype=np.int64) x[0, :n] = ids out = request.infer({0: x})[output] return torch.from_numpy(np.asarray(out[0, n - 1], dtype=np.float32)) class MLX: """mlx-lm's Qwen3.5 model on the model directory (it reads model*.safetensors, leaves the picture encoder out and never touches readout.safetensors); the text model's output is the hidden state after the final norm.""" name = "mlx" cache_limit_gb = 2.0 def __init__(self, path: str, options: dict | None = None): from pathlib import Path import mlx.core as mx from mlx_lm.utils import load_model self.mx = mx # MLX keeps freed buffers for reuse, by default up to its memory limit, so a server reading long prompts grew its # cache to most of the machine's memory. A bound keeps the reuse that matters for speed. mx.set_cache_limit(int(float((options or {}).get("cache_limit_gb", self.cache_limit_gb)) * 2**30)) model, _ = load_model(Path(path)) self.text = model.language_model.model def accepts(self, n: int) -> bool: return True def hidden(self, ids: Sequence[int]): import numpy as np import torch h = self.text(self.mx.array([list(ids)]))[0, -1].astype(self.mx.float32) self.mx.eval(h) return torch.from_numpy(np.array(h)) class LlamaCpp: """llama.cpp's llama-server answering /embedding with the last position's hidden state after the final norm, unnormalized (embd_normalize -1). The prompt goes as token ids, so llama.cpp's tokenizer is never used. Either a server you run (options["url"]) or one started here on the build's GGUF, on 127.0.0.1, stopped with the model.""" name = "llamacpp" def __init__(self, artifacts: str | None, build: dict | None, max_tokens: int, options: dict): self.context = int(options.get("context", max_tokens)) self.process = None url = options.get("url") if not url: gguf = options.get("gguf") or os.path.join(artifacts, build["files"]["gguf"]) server = options.get("server") or shutil.which("llama-server") with socket.socket() as s: s.bind(("127.0.0.1", 0)) port = s.getsockname()[1] # the embedding endpoint needs the whole prompt in one micro-batch: batch and micro-batch are the context command = [server, "-m", gguf, "--embeddings", "--pooling", "last", "-c", str(self.context), "-b", str(self.context), "-ub", str(self.context), "-ngl", str(options.get("gpu_layers", 999)), "--parallel", "1", "--host", "127.0.0.1", "--port", str(port)] # fmt: skip if options.get("threads"): command += ["-t", str(int(options["threads"]))] self.output = tempfile.TemporaryFile() self.process = subprocess.Popen(command, stdout=self.output, stderr=subprocess.STDOUT, stdin=subprocess.DEVNULL) self._stop = weakref.finalize(self, _stop, self.process) url = f"http://127.0.0.1:{port}" self._wait(url, float(options.get("start_timeout", 600))) self.url = url.rstrip("/") + "/embedding" def _wait(self, url: str, timeout: float) -> None: deadline = time.monotonic() + timeout while time.monotonic() < deadline: if self.process.poll() is not None: self.output.seek(0) tail = self.output.read().decode("utf-8", "replace")[-2000:] raise RuntimeError(f"llama-server stopped (exit {self.process.returncode}):\n{tail}") try: with urllib.request.urlopen(url + "/health", timeout=5) as r: if r.status == 200: return except (urllib.error.URLError, OSError): pass time.sleep(0.5) raise RuntimeError(f"llama-server did not answer within {timeout:.0f} s") def accepts(self, n: int) -> bool: return n <= self.context def hidden(self, ids: Sequence[int]): import torch body = json.dumps({"content": list(ids), "embd_normalize": -1, "cache_prompt": False}).encode() request = urllib.request.Request(self.url, data=body, headers={"Content-Type": "application/json"}) with urllib.request.urlopen(request, timeout=600) as r: out = json.loads(r.read()) vector = (out[0] if isinstance(out, list) else out)["embedding"] if vector and isinstance(vector[0], list): # one row per position: the last vector = vector[-1] return torch.tensor(vector, dtype=torch.float32) def _stop(process: subprocess.Popen) -> None: if process.poll() is None: process.terminate() try: process.wait(10) except subprocess.TimeoutExpired: process.kill()