runtime/src/decisionnode_edge/backends.py

"""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 <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, 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()