runtime/scheduler.py

import threading
from dataclasses import dataclass, field
from enum import IntEnum


class TokenPriority(IntEnum):
    HIGH = 0
    NORMAL = 1
    LOW = 2
    PREFETCH = 3


@dataclass
class TokenBatch:
    chunk_index: int
    layer_mask: int
    priority: TokenPriority = TokenPriority.NORMAL
    frames: list | None = None


class TokenScheduler:
    """Orchestrates which token chunks get decoded and when.

    Decouples the decoder from the raw chunk index by introducing
    priority-aware scheduling — critical frames, seek targets, and
    coarse layers can overtake bulk decode work.
    """
    def __init__(self, max_inflight: int = 4):
        self._max_inflight = max_inflight
        self._queue: list[TokenBatch] = []
        self._lock = threading.Lock()

    def submit(self, batch: TokenBatch):
        with self._lock:
            self._queue.append(batch)
            self._queue.sort(key=lambda b: b.priority)

    def acquire(self) -> TokenBatch | None:
        with self._lock:
            if not self._queue:
                return None
            return self._queue.pop(0)

    @property
    def pending(self) -> int:
        with self._lock:
            return len(self._queue)

    @property
    def max_inflight(self) -> int:
        return self._max_inflight