Files
esh-pfi-infrastructure/services/semif-serve/src/semif_serve/engine.py
T
vh 47cad33dd1 fix(semif): 0.1.4 — object states ending in ) ; } no longer 422 (INV-7)
SemIf's shared scorer trims one token at the state boundary. When an object
state's last value ends in ')', ';' or '}', the JSON that follows re-merges two
tokens back, so score_shared refused the request with 422. The engine now wraps
semif_phase1.shared._state_prefix to keep only the tokens the full prompts
share. Each row scores the same token sequence; only the prefill/suffix split
moves.

Startup proves the fix is in effect, not just installed (heid bug hunt SKAL,
folded). It checks that the hook is callable and is what score_shared resolves,
that an ordinary state keeps upstream's whole prefix, and that a merge-prone
state scores through the shared path.

Real tokenizer: 154 states, 23 refused before and 0 after, with no ordinary or
authored144 prefix changed. Acceptance: 144/144 parity. Shared vs direct 71/72;
the miss is a bf16 tie that flipped across a plain restart (see README).
2026-09-27 10:23:14 -07:00

187 lines
10 KiB
Python

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