gpc1_server/runtime.py

from __future__ import annotations

import json
import hashlib
import os
from pathlib import Path
from typing import Any, Protocol, Sequence

from .assets import load_codebook, load_runtime_manifest, sha256_file
from .compiler import NumericField, compile_numeric, compile_vision_numeric, decode_numeric
from .native_onepass import compile_joint, score_joint


class Backend(Protocol):
    model_id: str
    def numeric(self, context: str, fields: Sequence[NumericField], image: Any | None = None) -> dict[str, Any]: ...
    def atomic(self, context: str, field_key: str, choices: Sequence[Any]) -> dict[str, Any]: ...
    def finite_joint(self, context: str, schema: dict[str, Any], candidates: Sequence[dict[str, Any]]) -> dict[str, Any]: ...


class TransformersBackend:
    """Pinned Transformers/PEFT backend; no custom CUDA kernels or generation."""

    @staticmethod
    def _context_limits(manifest: dict[str, Any], model_config: dict[str, Any],
                        environ: dict[str, str]) -> tuple[int, int]:
        capacity = int(model_config["text_config"]["max_position_embeddings"])
        default = int(manifest["admission"]["default_max_input_tokens"])
        limit = int(environ.get("GPC1_MAX_INPUT_TOKENS", str(default)))
        if not 0 < limit <= capacity:
            raise ValueError(f"GPC1_MAX_INPUT_TOKENS must be between 1 and {capacity}")
        padded = int(environ.get("GPC1_MAX_PADDED_TOKENS", str(max(131072, 2 * limit))))
        if padded <= 0:
            raise ValueError("GPC1_MAX_PADDED_TOKENS must be positive")
        return limit, padded

    @staticmethod
    def _verify_model_file_closure(model_path: Path, expected_files: set[str]) -> None:
        top_files={p.name for p in model_path.iterdir() if p.is_file()}
        suspicious={name for name in top_files if (
            name.endswith((".safetensors",".json",".jinja",".txt"))
            or "config" in name or "tokenizer" in name or "processor" in name or "index" in name)}
        if suspicious != expected_files:
            raise RuntimeError(f"loader-visible model file closure mismatch: {sorted(suspicious ^ expected_files)}")

    @staticmethod
    def _verify_adapter_file_closure(adapter_path: Path) -> None:
        adapter_files={p.name for p in adapter_path.iterdir() if p.is_file()}
        if adapter_files != {"adapter_config.json","adapter_model.safetensors"}:
            raise RuntimeError("adapter directory must contain exactly the two bound release files")

    @staticmethod
    def _adapter_path(manifest: dict[str, Any], package_root: Path,
                      environ: dict[str, str]) -> Path:
        if manifest.get("weight_format") != "adapter" or not isinstance(manifest.get("adapter"), dict) \
                or manifest.get("merged_full_weights") is not None:
            raise RuntimeError("release runtime requires adapter weight metadata only")
        return Path(environ.get("GPC1_ADAPTER_PATH", package_root / "adapter"))

    def __init__(self) -> None:
        import torch
        from transformers import AutoProcessor
        from transformers import Qwen3_5MoeForConditionalGeneration

        self.torch = torch
        torch.backends.cuda.matmul.allow_tf32 = False
        torch.backends.cudnn.allow_tf32 = False
        torch.set_float32_matmul_precision("highest")
        self.manifest = load_runtime_manifest()
        package_root = Path(__file__).resolve().parents[1]
        model_path = Path(os.environ["GPC1_MODEL_PATH"])
        self.max_input_tokens, self.max_padded_tokens = self._context_limits(
            self.manifest, json.loads((model_path / "config.json").read_text()), os.environ)
        adapter_path = self._adapter_path(self.manifest, package_root, os.environ)
        self._verify_artifacts(model_path, adapter_path)
        self.processor = AutoProcessor.from_pretrained(model_path, local_files_only=True,
                                                       trust_remote_code=False)
        base = Qwen3_5MoeForConditionalGeneration.from_pretrained(
            model_path, local_files_only=True, trust_remote_code=False,
            dtype=torch.bfloat16, attn_implementation="sdpa",
            low_cpu_mem_usage=True, device_map={"": 0})
        from peft import PeftModel
        self.model = PeftModel.from_pretrained(base, adapter_path, is_trainable=False,
                                               autocast_adapter_dtype=True)
        self.model.eval()
        for parameter in self.model.parameters():
            parameter.requires_grad_(False)
        if base.__class__.__name__ != "Qwen3_5MoeForConditionalGeneration" \
                or getattr(base.config, "model_type", None) != "qwen3_5_moe":
            raise RuntimeError("loaded base is not the pinned Qwen3.5 MoE architecture")
        adapter_parameters = {name: p for name, p in self.model.named_parameters() if "lora_" in name}
        base_dtypes = {p.dtype for name, p in self.model.named_parameters()
                       if p.is_floating_point() and "lora_" not in name}
        if base_dtypes != {torch.bfloat16} or not adapter_parameters \
                or {p.dtype for p in adapter_parameters.values()} != {torch.float32}:
            raise RuntimeError("loaded base/adapter dtype inventory differs from BF16/FP32 contract")
        from peft.utils import get_peft_model_state_dict
        digest = hashlib.sha256()
        state = get_peft_model_state_dict(self.model)
        for name in sorted(state):
            tensor = state[name].detach().cpu().contiguous()
            metadata = json.dumps({"dtype":str(tensor.dtype),"name":name,"shape":list(tensor.shape)},
                                  sort_keys=True,separators=(",", ":")).encode()
            raw = tensor.view(torch.uint8).numpy().tobytes()
            digest.update(len(metadata).to_bytes(8,"big")); digest.update(metadata)
            digest.update(len(raw).to_bytes(8,"big")); digest.update(raw)
        if digest.hexdigest() != self.manifest["adapter"]["adapter_tensor_digest"]:
            raise RuntimeError("loaded adapter tensor digest mismatch")
        self.model_id = self.manifest["served_model_id"]
        self.codebook = load_codebook()
        self.max_joint_rows = int(os.environ.get("GPC1_MAX_JOINT_ROWS", "32"))
        self.max_logit_vectors = int(os.environ.get("GPC1_MAX_LOGIT_VECTORS", "512"))
        self.max_logits_bytes = int(os.environ.get("GPC1_MAX_LOGITS_BYTES", str(512 * 1024 * 1024)))
        accepted_modes = set(self.manifest["admission"]["accepted_modes"])
        self.admit_vision = "numeric101_image" in accepted_modes \
            and os.environ.get("GPC1_ADMIT_VISION_NUMERIC", "accepted") == "accepted"
        self.admit_joint = "finite_joint" in accepted_modes \
            and os.environ.get("GPC1_ADMIT_FINITE_JOINT", "accepted") == "accepted"
        cache_mb = int(os.environ.get("GPC1_IMAGE_CACHE_MB", "64"))
        if not 0 <= cache_mb <= 64:
            raise ValueError("GPC1_IMAGE_CACHE_MB must be an integer between 0 and 64")
        self._vision_cache = None
        if cache_mb:
            from .vision_cache import VisionFeatureCache
            self._vision_cache = VisionFeatureCache(
                self.model.base_model.model.model.visual, torch, cache_mb * 1024 * 1024)

    def _verify_artifacts(self, model_path: Path, adapter_path: Path) -> None:
        expected = self.manifest
        package_root = Path(__file__).resolve().parents[1]
        inventory_path = package_root / "assets" / "base_files.sha256.json"
        inventory = json.loads(inventory_path.read_text())
        if inventory["model_id"] != expected["base"]["model_id"] or inventory["revision"] != expected["base"]["revision"]:
            raise RuntimeError("model-file inventory identity mismatch")
        self._verify_model_file_closure(model_path, set(inventory["files"]))
        for relative, digest in inventory["files"].items():
            path = model_path / relative
            if not path.is_file() or sha256_file(path) != digest:
                raise RuntimeError(f"pinned model file identity mismatch: {relative}")
        self._verify_adapter_file_closure(adapter_path)
        if sha256_file(adapter_path / "adapter_model.safetensors") != expected["adapter"]["adapter_model_sha256"]:
            raise RuntimeError("bundled adapter identity mismatch")
        if sha256_file(adapter_path / "adapter_config.json") != expected["adapter"]["adapter_config_sha256"]:
            raise RuntimeError("bundled adapter config identity mismatch")

    def _forward(self, inputs: dict[str, Any], positions: Sequence[int]) -> Any:
        torch = self.torch
        device = next(self.model.parameters()).device
        moved = {k: v.to(device) if hasattr(v, "to") else v for k, v in inputs.items()}
        with torch.inference_mode():
            # Project logits only at requested positions.
            output = self.model(**moved, logits_to_keep=[int(p) for p in positions],
                                use_cache=False, return_dict=True)
        if output.logits.ndim != 3 or tuple(output.logits.shape[:2]) != (1, len(positions)):
            raise RuntimeError("native sparse logits do not align to query positions")
        return output.logits[0].float()

    def numeric(self, context: str, fields: Sequence[NumericField], image: Any | None = None) -> dict[str, Any]:
        vocab = int(self.model.config.text_config.vocab_size)
        if len(fields) * vocab * 4 > self.max_logits_bytes:
            raise ValueError("numeric sparse logits exceed byte admission")
        if image is None:
            compiled = compile_numeric(self.processor.tokenizer, context, fields, self.codebook,
                                       self.max_input_tokens)
            ids = self.torch.tensor([compiled.input_ids], dtype=self.torch.long)
            mask = self.torch.ones_like(ids)
            inputs = {"input_ids": ids, "attention_mask": mask,
                      "position_ids": mask.cumsum(dim=-1) - 1}
            positions = compiled.positions
        else:
            if not self.admit_vision:
                raise ValueError("vision numeric requires an accepted same-hardware parity gate")
            vision = compile_vision_numeric(self.processor, image, context, fields, self.codebook,
                                            self.max_input_tokens, self.model.config)
            compiled = vision.numeric; inputs = vision.processor_inputs; positions = vision.positions
        logits = self._forward(inputs, positions)
        token_ids = self.torch.tensor(compiled.candidate_token_ids, device=logits.device)
        results = {}
        for index, field in enumerate(fields):
            raw = logits[index].index_select(0, token_ids).float()
            probs = self.torch.softmax(raw, dim=-1).cpu().tolist()
            results[field.key] = decode_numeric(probs, field)
        return {"fields": results, "accounting": {"backbone_calls": 1,
                "physical_rows": 1, "field_slots": len(fields), "truncated": False},
                "semantics": "per-field marginals; not a learned full joint distribution"}

    def atomic(self, context: str, field_key: str, choices: Sequence[Any]) -> dict[str, Any]:
        if not 2 <= len(choices) <= 255 or len({json.dumps(v,sort_keys=True) for v in choices}) != len(choices):
            raise ValueError("atomic choices must contain 2..255 unique JSON values")
        schema={"field_order":[field_key],"fields":{field_key:{"type":"enum","choices":list(choices)}}}
        support="\n".join(f'{self.codebook[i]["code"]} = '+json.dumps(v,ensure_ascii=False,sort_keys=True,separators=(",",":")) for i,v in enumerate(choices))
        user=("Classify the active field from the complete input.\n\nContext:\n"+context+"\n\n"
              "Registered schema gpc1-public-atomic@1:\n"+json.dumps(schema,ensure_ascii=False,sort_keys=True,separators=(",",":"))
              +"\n\nAllowed complete outputs in registered semantic order:\n"+support+"\n\nReturn exactly one allowed output.")
        prompt=self.processor.tokenizer.apply_chat_template([
            {"role":"system","content":"You are a fixed-schema lookup model. Return exactly one listed opaque code and no other text."},
            {"role":"user","content":user}],tokenize=False,add_generation_prompt=True,enable_thinking=False)
        ids=tuple(self.processor.tokenizer.encode(prompt,add_special_tokens=False))
        if len(ids)>self.max_input_tokens: raise ValueError("atomic request exceeds token admission; no truncation performed")
        for row in self.codebook[:len(choices)]:
            combined=tuple(self.processor.tokenizer.encode(prompt+str(row["code"]),add_special_tokens=False))
            if combined != ids + (int(row["token_id"]),):
                raise ValueError("atomic response-boundary tokenization changed")
        tensor=self.torch.tensor([ids]); logits=self._forward({"input_ids":tensor,"attention_mask":self.torch.ones_like(tensor)},[len(ids)-1])[0]
        indices=self.torch.tensor([int(r["token_id"]) for r in self.codebook[:len(choices)]],device=logits.device)
        probs=self.torch.softmax(logits.index_select(0,indices).float(),dim=-1).cpu().tolist(); selected=max(range(len(choices)),key=probs.__getitem__)
        return {"value":choices[selected],"probabilities":probs,"selected_index":selected,"calibrated":False,
                "accounting":{"backbone_calls":1,"physical_rows":1,"truncated":False},
                "semantics":"conditional distribution over the complete caller-declared choice set"}

    def finite_joint(self, context: str, schema: dict[str, Any], candidates: Sequence[dict[str, Any]]) -> dict[str, Any]:
        if not self.admit_joint:
            raise ValueError("finite_joint requires an accepted same-hardware parity gate")
        tokenizer = self.processor.tokenizer
        canonical = [json.dumps(v, ensure_ascii=False, sort_keys=True,
                                separators=(",", ":"), allow_nan=False) for v in candidates]
        support = "\n".join(f"{value} = {value}" for value in canonical)
        user = ("Select one complete allowed JSON record from the complete input.\n\nContext:\n" + context
                + "\n\nRegistered schema gpc1-public-finite-joint@1:\n"
                + json.dumps(schema, ensure_ascii=False, sort_keys=True, separators=(",", ":"), allow_nan=False)
                + "\n\nAllowed complete outputs in registered semantic order:\n" + support
                + "\n\nReturn exactly one allowed output.")
        prompt = tokenizer.apply_chat_template([
            {"role":"system","content":"You are a fixed-schema lookup model. Return exactly one listed complete JSON record and no other text."},
            {"role":"user","content":user}], tokenize=False, add_generation_prompt=True,
            enable_thinking=False)
        prefix = tuple(tokenizer.encode(prompt, add_special_tokens=False))
        terminal_text = tokenizer.eos_token or "<|im_end|>"
        terminal = tuple(tokenizer.encode(terminal_text, add_special_tokens=False))
        suffixes = []
        for value in canonical:
            complete = tuple(tokenizer.encode(prompt + value + terminal_text, add_special_tokens=False))
            if complete[:len(prefix)] != prefix: raise ValueError("candidate tokenization changed prompt prefix")
            suffixes.append(complete[len(prefix):])
        compiled = compile_joint(prefix, tuple(str(i) for i in range(len(candidates))), suffixes,
            terminal, pad_token_id=int(tokenizer.pad_token_id), max_rows=self.max_joint_rows,
            max_padded_tokens=self.max_padded_tokens, max_logit_vectors=self.max_logit_vectors,
            max_input_tokens=self.max_input_tokens)
        vocab = int(self.model.config.text_config.vocab_size)
        if compiled.logit_vectors * vocab * 4 > self.max_logits_bytes:
            raise ValueError("joint sparse logits exceed byte admission")
        device = next(self.model.parameters()).device
        result = score_joint(self.model, compiled, device=device)
        result["value"] = candidates[result["selected_index"]]
        result["semantics"] = "conditional distribution over the complete caller-declared joint support"
        return result