"""Prefill-only decision engine with dynamic batching. A decision is one forward pass (no generation): the prompt holds the state and every question, and the head scores every option. Requests arriving close together are run as one GPU batch: the worker takes the first waiting pass, then keeps adding passes until the batch is full (token budget or batch size) or `max_wait_ms` has passed since it started filling. Jev limits are enforced before anything is queued: over 65,536 tokens per request, or over 32,768 for the state plus its longest question, is rejected. Nothing is truncated. A request too large for one pass of the trained length is answered in several passes with the full state. Used by the live API (server.py) and the Decision Index engine (decider_engine.py). """ import concurrent.futures import importlib.util import logging import os import queue import threading import time import torch import model as M import prompting import schema log = logging.getLogger("decider.engine") class RequestTooLarge(ValueError): """The request exceeds a Jev token limit; the message says which.""" class Overloaded(RuntimeError): """The queue is full; the caller should retry later.""" def fast_kernels(): """Which optimised kernels transformers will pick up for Qwen3.5's linear-attention layers.""" return {"flash_linear_attention": importlib.util.find_spec("fla") is not None, "causal_conv1d": importlib.util.find_spec("causal_conv1d") is not None} class Pending: """One request in flight: its passes, and the probabilities gathered so far.""" def __init__(self, questions, items, input_tokens): self.questions = questions self.items = items self.input_tokens = input_tokens self.probabilities = [None] * len(items) self.remaining = 0 self.future = concurrent.futures.Future() self.lock = threading.Lock() class Pass: def __init__(self, pending, state, indices, tokens): self.pending = pending self.x = {"state": state, "items": [pending.items[i] for i in indices]} self.indices = indices self.tokens = tokens class DecisionEngine: def __init__(self, adapter, device="cuda", name="decider-2b", max_batch_tokens=65536, max_batch_size=64, max_wait_ms=4.0, max_queue=2048, max_pass_tokens=schema.MAX_STATE_PLUS_QUESTION_TOKENS, require_fast_kernels=False): kernels = fast_kernels() if device.startswith("cuda") and not all(kernels.values()): message = f"fast linear-attention kernels missing: {kernels}; install flash-linear-attention and causal-conv1d" if require_fast_kernels: raise RuntimeError(message) log.warning(message) self.kernels = kernels self.model, self.config = M.load(adapter, device) # The abstention threshold tuned for this model ships in its config; the environment overrides it. if "abstain_threshold" in self.config and "DECIDER_ABSTAIN_THRESHOLD" not in os.environ: schema.ABSTAIN_THRESHOLD = float(self.config["abstain_threshold"]) self.device = device self.name = name self.trained_length = int(self.config.get("max_length", 8192)) self.max_pass_tokens = int(max_pass_tokens) self.max_batch_tokens = max(int(max_batch_tokens), self.max_pass_tokens) self.max_batch_size = int(max_batch_size) self.max_wait = max_wait_ms / 1000.0 self.queue = queue.Queue(maxsize=int(max_queue)) self.intro_tokens = len(self.model.tokenizer(prompting.INTRO, add_special_tokens=False)["input_ids"]) self.stats = {"batches": 0, "passes": 0, "last_batch_size": 0, "last_batch_ms": 0.0} self.observers = [] # callables(batch_passes, batch_tokens, seconds) for metrics self.ready = False self._carry = None self._worker = threading.Thread(target=self._run, name="decision-engine", daemon=True) self._worker.start() # ---- request side ------------------------------------------------------------------------- def prepare(self, body): """Validate and plan a request. Raises schema.SchemaError or RequestTooLarge.""" schema.validate_request(body) questions = body["questions"] state = schema.request_state_text(body["state"], questions) items = schema.items(questions) tokenizer = self.model.tokenizer _, state_tokens, sizes = prompting.plan(tokenizer, state, items, 1 << 30) total = self.intro_tokens + state_tokens + sum(sizes) if total > schema.MAX_REQUEST_TOKENS: raise RequestTooLarge(f"request is {total} tokens; the limit is {schema.MAX_REQUEST_TOKENS}") longest = self.intro_tokens + state_tokens + max(sizes) if longest > schema.MAX_STATE_PLUS_QUESTION_TOKENS: raise RequestTooLarge(f"state plus the longest question is {longest} tokens; " f"the limit is {schema.MAX_STATE_PLUS_QUESTION_TOKENS}") # Passes no longer than the trained length when the state allows it. limit = self.trained_length if longest <= self.trained_length else self.max_pass_tokens passes, _, _ = prompting.plan(tokenizer, state, items, limit) pending = Pending(questions, items, total) pending.remaining = len(passes) base = self.intro_tokens + state_tokens return pending, [Pass(pending, state, indices, base + sum(sizes[i] for i in indices)) for indices in passes] def submit(self, body): """Queue a request; returns a Future resolving to the SystemOne response.""" pending, passes = self.prepare(body) if self.queue.qsize() + len(passes) > self.queue.maxsize: raise Overloaded("the decision queue is full") for p in passes: self.queue.put_nowait(p) return pending.future def decide(self, body, timeout=None): """Blocking form of submit().""" return self.submit(body).result(timeout) def warmup(self): """Run a few requests of growing size so kernels are compiled before traffic arrives.""" for words in (8, 200, 1500): self.decide({"state": "warm up " * words, "questions": { "a": {"type": "noul", "instructions": "This is a warm-up statement."}, "b": {"type": "choice", "instructions": "Pick one.", "criteria": {"x": "first", "y": "second"}}, "c": {"type": "score", "instructions": "Rate it.", "criteria": ["low", "mid", "high"]}}}) self.ready = True # ---- worker side -------------------------------------------------------------------------- def _collect(self): """One batch: the oldest waiting pass, then whatever arrives within max_wait that fits.""" if self._carry is not None: batch, self._carry = [self._carry], None else: batch = [self.queue.get()] longest = batch[0].tokens deadline = time.monotonic() + self.max_wait while len(batch) < self.max_batch_size: timeout = deadline - time.monotonic() try: nxt = self.queue.get(timeout=timeout) if timeout > 0 else self.queue.get_nowait() except queue.Empty: break # Padded cost: every row in the batch is as long as the longest one. if (len(batch) + 1) * max(longest, nxt.tokens) > self.max_batch_tokens: self._carry = nxt # starts the next batch, keeping arrival order break batch.append(nxt) longest = max(longest, nxt.tokens) return batch def _run(self): while True: # Skip passes whose request already timed out or failed. batch = [p for p in self._collect() if not p.pending.future.done()] if not batch: continue started = time.monotonic() try: with torch.inference_mode(): tensors = self.model.encode([p.x for p in batch]) tensors = {k: v.to(self.device) for k, v in tensors.items()} logits = (self.model(**tensors).float() / self.model.temperature).cpu() except Exception as error: log.exception("batch of %d passes failed", len(batch)) for p in batch: if not p.pending.future.done(): p.pending.future.set_exception(error) continue elapsed = time.monotonic() - started row = 0 for p in batch: for index in p.indices: count = len(p.pending.items[index]["labels"]) p.pending.probabilities[index] = logits[row, :count].softmax(-1).tolist() row += 1 self._finish(p) self.stats.update(batches=self.stats["batches"] + 1, passes=self.stats["passes"] + len(batch), last_batch_size=len(batch), last_batch_ms=1000 * elapsed) for observe in self.observers: observe(len(batch), max(p.tokens for p in batch) * len(batch), elapsed) def _finish(self, p): pending = p.pending with pending.lock: pending.remaining -= 1 done = pending.remaining == 0 if done and not pending.future.done(): try: pending.future.set_result(schema.response(self.name, pending.questions, pending.probabilities, pending.input_tokens)) except Exception as error: pending.future.set_exception(error)