| """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) |
| |
| 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 = [] |
| self.ready = False |
| self._carry = None |
| self._worker = threading.Thread(target=self._run, name="decision-engine", daemon=True) |
| self._worker.start() |
|
|
| |
|
|
| 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}") |
| |
| 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 |
|
|
| |
|
|
| 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 |
| |
| if (len(batch) + 1) * max(longest, nxt.tokens) > self.max_batch_tokens: |
| self._carry = nxt |
| break |
| batch.append(nxt) |
| longest = max(longest, nxt.tokens) |
| return batch |
|
|
| def _run(self): |
| while True: |
| |
| 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) |
|
|
