"""The real engine: SemIf's torch scorers over one resident model. Needs the `model` extra. Contract: semif-serve.contract.md, INV-3 (fail-closed startup), INV-4 (VRAM cap + OOM), INV-5 (offline weights). load() is checked on the card at acceptance; the OOM path is unit-tested against a fake torch (tests/test_engine.py). """ from __future__ import annotations import gc import logging import traceback from typing import Any, Callable from .config import Settings from .errors import OutOfMemory, ScoringFailed log = logging.getLogger("semif_serve.engine") RELEASE_SLACK_BYTES = 512 * 2**20 WARMUP_ROW = { "id": "semif-serve-warmup", "state": "The deployment completed at 14:02 UTC. Health checks passed in all three zones.", "question": "Is there evidence that the deployment succeeded?", "options": [ {"id": "yes", "description": "The deployment succeeded."}, {"id": "no", "description": "The deployment did not succeed."}, ], } # INV-7 startup proof: a state SemIf's own prefix refuses (measured with the real tokenizer, # 2026-09-27), sent through score_shared itself, so the fix is shown to be IN EFFECT, not just # installed (bug hunt SKAL, R2/H1/H2). BOUNDARY_ROW = {**WARMUP_ROW, "id": "semif-serve-boundary-check", "state": {"person_said": "ok :)"}} PREFIX_PROBE_ROW = { "id": "semif-serve-prefix-probe", "question": "prefix boundary placeholder", "options": [{"id": "yes", "description": "Yes"}, {"id": "no", "description": "No"}], } def boundary_safe_prefix(state_prefix: Callable, messages: Callable) -> Callable: """INV-7: wrap SemIf's `shared._state_prefix` so the prefill never covers a token that the full prompts do not share. Upstream encodes the prompt up to the end of the state and drops ONE token, because the JSON punctuation that follows the state can merge with it. One is not always enough: an object state whose last value ends in ")", ";" or "}" re-tokenises TWO tokens back once `, "criterion"` follows, and score_shared then refused the request with 422 (found 2026-09-27). This keeps only the leading tokens that upstream's prefix shares with a real full prompt for the same state. The suffix starts that much earlier and scores the same token sequence; the cost is a few tokens of lost sharing. Only the punctuation run at the boundary can merge, and it is the same in every row, so a probe row stands for all of them. score_shared still checks every real row and fails closed if that ever stops holding.""" def prefix(tokenizer, state): ids = state_prefix(tokenizer, state) full = tokenizer.encode(tokenizer.apply_chat_template( messages({**PREFIX_PROBE_ROW, "state": state}), tokenize=False, add_generation_prompt=True, enable_thinking=False), add_special_tokens=False) shared = 0 while shared < min(len(ids), len(full)) and ids[shared] == full[shared]: shared += 1 return ids[:shared] prefix.semif_serve_wraps = state_prefix return prefix def _first_line(exc: BaseException) -> str: lines = str(exc).splitlines() return lines[0] if lines else "" class TorchEngine: def __init__(self, torch: Any, model: Any, tokenizer: Any, metadata: dict, settings: Settings, direct_fn: Callable, shared_fn: Callable, release_above_bytes: int | None = None): self._torch, self._model, self._tokenizer = torch, model, tokenizer self._metadata, self._settings = metadata, settings self._direct, self._shared = direct_fn, shared_fn self._release_above = release_above_bytes @classmethod def load(cls, settings: Settings) -> "TorchEngine": import torch import semif_phase1.shared as upstream_shared from semif_phase1.core import direct_messages, load_causal_model from semif_phase1.direct import score from semif_phase1.shared import score_shared # INV-7: score_shared looks `_state_prefix` up as a module global, so the wrapper goes there. # A SemIf bump that renames it, or moves score_shared so it resolves its globals elsewhere, # must stop startup, not silently bring the 422 back. original = getattr(upstream_shared, "_state_prefix", None) if not callable(original): raise RuntimeError("semif_phase1.shared._state_prefix is gone or not callable at this SemIf " "commit: re-check the boundary-safe prefix (INV-7) before serving") original = getattr(original, "semif_serve_wraps", original) # one wrapper, however many loads upstream_shared._state_prefix = boundary_safe_prefix(original, direct_messages) if getattr(score_shared, "__globals__", {}).get("_state_prefix") is not upstream_shared._state_prefix: raise RuntimeError("INV-7: score_shared does not resolve semif_phase1.shared._state_prefix, " "so the boundary-safe prefix would be inert") if settings.device == "cuda": if not torch.cuda.is_available(): raise RuntimeError("SEMIF_DEVICE=cuda but torch sees no CUDA device") major, minor = torch.cuda.get_device_capability(0) arch = f"sm_{major}{minor}" if arch not in torch.cuda.get_arch_list(): # INV-3: no silent PTX/CPU fallback raise RuntimeError(f"torch {torch.__version__} has no kernels for {arch}: {torch.cuda.get_arch_list()}") if settings.vram_cap_gib is not None: # INV-4: cap BEFORE the weights land total = torch.cuda.get_device_properties(0).total_memory fraction = settings.vram_cap_gib * 2**30 / total if not 0 < fraction <= 1: raise ValueError(f"SEMIF_VRAM_CAP_GIB={settings.vram_cap_gib} does not fit a {total / 2**30:.1f} GiB card") torch.cuda.set_per_process_memory_fraction(fraction, 0) elif settings.device != "cpu": raise ValueError(f"SEMIF_DEVICE must be cuda or cpu, not {settings.device!r}") model, tokenizer, metadata = load_causal_model(settings.model, settings.revision, settings.device, "bfloat16") placed = next(model.parameters()).device.type if placed != settings.device: # INV-3 raise RuntimeError(f"model landed on {placed}, expected {settings.device}") engine = cls(torch, model, tokenizer, metadata, settings, direct_fn=score, shared_fn=score_shared) engine.direct(WARMUP_ROW) # INV-3: one decision must score engine._prove_prefix_hook(upstream_shared._state_prefix, original) if settings.device == "cuda": # INV-4: the resting footprint engine._release_above = torch.cuda.memory_reserved(0) + RELEASE_SLACK_BYTES return engine def _prove_prefix_hook(self, hook: Callable, original: Callable) -> None: """INV-7, at startup: the wrapper keeps upstream's whole prefix on an ordinary state (a wrapper rendering the wrong prompt would silently give up all sharing), and a state upstream alone refuses scores through score_shared itself (a hook score_shared never calls would not).""" state = WARMUP_ROW["state"] if hook(self._tokenizer, state) != original(self._tokenizer, state): raise RuntimeError("INV-7: the boundary-safe prefix does not keep upstream's prefix on an ordinary state") try: self.shared([BOUNDARY_ROW]) except ValueError as exc: raise RuntimeError(f"INV-7: shared scoring still refuses a merge-prone state: {exc}") from None def health(self) -> dict: info = dict(self._metadata) if self._settings.device == "cuda": info["device_name"] = self._torch.cuda.get_device_name(0) info["allocated_gib"] = round(self._torch.cuda.memory_allocated(0) / 2**30, 2) info["reserved_gib"] = round(self._torch.cuda.memory_reserved(0) / 2**30, 2) return info def _release_burst(self) -> None: """INV-4: hand a burst back to the driver so the card's shared headroom (scriberr, the vLLM seats) returns after a big request, instead of sitting in torch's cache.""" if self._release_above is not None and self._torch.cuda.memory_reserved(0) > self._release_above: self._torch.cuda.empty_cache() def _guard(self, fn, *args): try: result = fn(*args) except ValueError: raise # validation: SemIf raises it before any GPU work except self._torch.cuda.OutOfMemoryError as exc: failure, message = OutOfMemory, _first_line(exc) or "CUDA out of memory" except Exception as exc: # noqa: BLE001 — every other failure is released and reported below message = _first_line(exc) if "out of memory" in message.lower(): # cuBLAS/cuDNN allocation failures failure = OutOfMemory else: failure, message = ScoringFailed, f"{type(exc).__name__}: {message}" # Formatted text, not exc_info: a log record that keeps the traceback object alive # (pytest's capture handler does; so would any buffering handler) pins the tensors. log.error("scorer failed:\n%s", traceback.format_exc()) else: self._release_burst() return result # INV-4, outside the except block on purpose: the exception's traceback holds the failed # scorer's frames, and with them its tensors (the replicated prefix cache). Raising inside # the block, or `from exc`, would chain to it and keep GiBs allocated after the response. gc.collect() self._torch.cuda.empty_cache() raise failure(message) def direct(self, row: dict) -> dict: return self._guard(self._direct, self._model, self._tokenizer, row, self._metadata, self._settings.max_tokens) def shared(self, rows: list[dict]) -> tuple[list[dict], dict]: return self._guard(self._shared, self._model, self._tokenizer, rows, self._metadata, self._settings.max_tokens)