| """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") |
| |
| 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" |
| |
| MODEL_FILES = ("config.json", "model.safetensors", "tokenizer.json", "tokenizer_config.json", "preprocessor_config.json", "video_preprocessor_config.json") |
|
|
| log = logging.getLogger("decisionnode_edge") |
|
|
|
|
| |
| |
| |
|
|
|
|
| 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 <model> --out <dir>, then artifacts=<dir>" |
| 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': <a running llama-server>}" |
| if not (options.get("server") or shutil.which("llama-server")): |
| return "llama-server is not on PATH: give options={'server': <path to llama-server>} or {'url': <a running llama-server>}" |
| 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) |
| 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)) |
|
|
|
|
| |
| |
| |
| |
|
|
|
|
| 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 |
| HEADROOM_BYTES = 4 << 30 |
|
|
| 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) |
| |
| 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}) |
| 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" |
|
|
| def __init__(self, path: str): |
| from pathlib import Path |
|
|
| import mlx.core as mx |
| from mlx_lm.utils import load_model |
|
|
| self.mx = mx |
| 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] |
| |
| 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)] |
| 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): |
| 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() |
|
|
