"""A bounded, single-worker inference queue with cancellable visitor leases.""" import logging import threading import time from collections import OrderedDict, deque from dataclasses import dataclass, field class AdmissionError(Exception): def __init__(self, message, status=429, code="queue_full"): super().__init__(message) self.status = status self.code = code class RequestCancelled(Exception): pass @dataclass(eq=False) class Ticket: request_id: str payload: dict created: float lease_until: float | None is_disconnected: object = None events: deque = field(default_factory=deque) status: str = "waiting" cancelled: bool = False done: bool = False class InferenceQueue: def __init__(self, runner, inference_lock=None, capacity=2, lease_seconds=12, max_wait=180, buffer_size=8, clock=time.monotonic, autostart=True): self.runner = runner self.inference_lock = inference_lock or threading.Lock() self.capacity = capacity self.lease_seconds = lease_seconds self.max_wait = max_wait self.buffer_size = buffer_size self.clock = clock self._condition = threading.Condition() self._tickets = {} self._waiting = deque() self._active = None self._legacy = False self._closed = False self._cancelled_ids = OrderedDict() self._thread = None self._autostart = autostart def _start_locked(self): if self._autostart and self._thread is None: self._thread = threading.Thread(target=self._work, name="marvin-inference", daemon=True) self._thread.start() def _remember_cancelled_locked(self, request_id): self._cancelled_ids[request_id] = self.clock() + 30 self._cancelled_ids.move_to_end(request_id) while len(self._cancelled_ids) > 1024: self._cancelled_ids.popitem(last=False) def _positions_locked(self): for position, ticket in enumerate(self._waiting, 1): # Positions can change faster than a slow client reads. Only the # current position matters, and these updates must stay bounded. ticket.events = deque(event for event in ticket.events if event.get("type") != "queued") ticket.events.append({"type": "queued", "position": position}) self._condition.notify_all() def _cancel_locked(self, ticket, message="Request cancelled.", code="cancelled"): if ticket.done or ticket.cancelled: return ticket.cancelled = True ticket.events.clear() ticket.events.append({"type": "error", "error": message, "code": code}) self._remember_cancelled_locked(ticket.request_id) if ticket in self._waiting: self._waiting.remove(ticket) ticket.done = True self._positions_locked() self._condition.notify_all() def _expire_locked(self): now = self.clock() for request_id, expiry in list(self._cancelled_ids.items()): if expiry <= now: del self._cancelled_ids[request_id] for ticket in list(self._tickets.values()): expired = ticket.lease_until is not None and ticket.lease_until <= now if ticket.done: if expired or now - ticket.created > self.max_wait + 300: self._tickets.pop(ticket.request_id, None) continue if ticket.is_disconnected is not None and ticket.is_disconnected(): self._cancel_locked(ticket) elif expired: self._cancel_locked(ticket, "Your turn ended because you left Marvin. Come back and try again.", "lease_expired") elif ticket.status == "waiting" and now - ticket.created >= self.max_wait: self._cancel_locked(ticket, "Marvin's taking a while. Please try again in a moment.", "queue_expired") def submit(self, request_id, payload, leased=True, allow_wait=True, is_disconnected=None): with self._condition: self._expire_locked() if self._closed: raise AdmissionError("Marvin is restarting. Give him a moment.", 503, "unavailable") if request_id in self._cancelled_ids: raise AdmissionError("This request was cancelled. Send a new message when you're ready.", 409, "cancelled") if request_id in self._tickets: raise AdmissionError("That request is already with Marvin.", 409, "duplicate_request") occupied = self._active is not None or self._legacy if occupied and (not allow_wait or len(self._waiting) >= self.capacity): raise AdmissionError("Marvin's got a little crowd. Try again shortly.") now = self.clock() ticket = Ticket(request_id, payload, now, now + self.lease_seconds if leased else None, is_disconnected) self._tickets[request_id] = ticket if occupied: self._waiting.append(ticket) self._positions_locked() else: ticket.status = "starting" self._active = ticket self._start_locked() self._condition.notify_all() return ticket def keepalive(self, request_id): with self._condition: self._expire_locked() ticket = self._tickets.get(request_id) if ticket is None or ticket.cancelled or ticket.done: return False if ticket.lease_until is not None: ticket.lease_until = self.clock() + self.lease_seconds return True def cancel(self, request_id): with self._condition: # A pagehide beacon can overtake the original POST on the network. # Remember it even before admission, so that POST cannot start work. self._remember_cancelled_locked(request_id) ticket = self._tickets.get(request_id) if ticket is not None: self._cancel_locked(ticket) self._condition.notify_all() def checkpoint(self, ticket): with self._condition: self._expire_locked() if ticket.cancelled or ticket.done or self._closed: raise RequestCancelled() def emit(self, ticket, event): with self._condition: stalled_at = self.clock() while len(ticket.events) >= self.buffer_size: self.checkpoint(ticket) if self.clock() - stalled_at >= self.lease_seconds: self._cancel_locked(ticket, "The connection to Marvin stopped. Please try again.", "connection_stalled") raise RequestCancelled() self._condition.wait(0.25) self.checkpoint(ticket) ticket.events.append(event) self._condition.notify_all() def events(self, ticket): try: while True: with self._condition: self._expire_locked() if not ticket.events and not ticket.done and not ticket.cancelled: self._condition.wait(2) self._expire_locked() if ticket.events: event = ticket.events.popleft() self._condition.notify_all() elif ticket.done or ticket.cancelled: break else: event = None # An SSE heartbeat notices broken sockets. yield event finally: self.cancel(ticket.request_id) with self._condition: if ticket.done: self._tickets.pop(ticket.request_id, None) def _promote_locked(self): self._expire_locked() if self._active is None and not self._legacy and self._waiting: self._active = self._waiting.popleft() self._active.events.clear() self._active.status = "starting" self._positions_locked() def run_once(self): """Run one reserved job; also permits deterministic tests without threads.""" with self._condition: self._expire_locked() ticket = self._active if ticket is None or ticket.status != "starting": return False ticket.status = "running" locked = False try: self.checkpoint(ticket) locked = self.inference_lock.acquire(blocking=False) if not locked: raise RuntimeError("Inference admission lost its exclusive lock") self.runner(ticket, self) except RequestCancelled: pass except Exception: logging.exception("Diffusion generation failed") try: self.emit(ticket, {"type": "error", "error": "Generation failed. Please try again."}) except RequestCancelled: pass finally: if locked: self.inference_lock.release() with self._condition: ticket.done = True ticket.payload = None self._active = None self._promote_locked() self._condition.notify_all() return True def _work(self): while True: with self._condition: self._expire_locked() if self._closed: return if self._active is None: self._condition.wait(0.5) continue self.run_once() def acquire_legacy(self): with self._condition: self._expire_locked() if self._closed or self._active is not None or self._waiting or self._legacy: return False if not self.inference_lock.acquire(blocking=False): return False self._legacy = True return True def release_legacy(self): with self._condition: self.inference_lock.release() self._legacy = False self._promote_locked() self._condition.notify_all() def snapshot(self): with self._condition: self._expire_locked() return {"active": self._active is not None or self._legacy, "waiting": len(self._waiting), "capacity": self.capacity, "lease_seconds": self.lease_seconds} def close(self): with self._condition: self._closed = True for ticket in list(self._tickets.values()): self._cancel_locked(ticket) self._condition.notify_all() if self._thread: self._thread.join(timeout=5)