streaming ¶

@abstractmethod
def drop(self, blob_id: str) -> None:
    """Remove a blob. Missing ids are a no-op."""

@abstractmethod
def get(self, blob_id: str) -> bytes:
    """Load a previously stored blob. Raises ``KeyError`` if absent."""

@abstractmethod
def put(self, data: bytes, *, mime: str = "application/octet-stream") -> str:
    """Store ``data`` and return a blob id for later retrieval."""

async def encode(
    self,
    frames: list[np.ndarray] | AsyncIterator[np.ndarray],
) -> AsyncIterator[FragmentedMP4Chunk]:
    """Feed frames into ffmpeg and yield fMP4 chunks as they appear."""
    if self._proc is None:
        self._spawn()
    assert self._proc is not None and self._proc.stdin is not None
    proc = self._proc

    loop = asyncio.get_running_loop()

    async def _writer() -> None:
        try:
            if hasattr(frames, "__aiter__"):
                async for frame in frames:  # type: ignore[union-attr]
                    await loop.run_in_executor(None, _write_frame, proc.stdin, frame)
            else:
                for frame in frames:  # type: ignore[assignment]
                    await loop.run_in_executor(None, _write_frame, proc.stdin, frame)
        finally:
            with contextlib.suppress(BrokenPipeError):
                proc.stdin.close()

    writer_task = asyncio.create_task(_writer())
    try:
        reader = proc.stdout
        assert reader is not None
        # Read in reasonably-sized chunks; MSE tolerates any size
        # but we don't want to starve the event loop.
        chunk_size = 64 * 1024
        while True:
            data = await loop.run_in_executor(None, reader.read, chunk_size)
            if not data:
                break
            kind: Literal["init", "media"] = "init" if not self._init_emitted else "media"
            self._init_emitted = True
            yield FragmentedMP4Chunk(
                kind=kind,
                data=bytes(data),
                stream_id=self.stream_id,
                segment_idx=self.segment_idx,
            )
    finally:
        await writer_task

def register_provider(self, provider: LLMProvider, *, priority: int = -1) -> None:
    """Insert an additional provider. ``priority=0`` makes it primary;
    ``priority=-1`` (default) appends as a fallback."""
    if priority < 0:
        self._providers.append(provider)
    else:
        self._providers.insert(priority, provider)

def reload_system_prompts(self) -> None:
    """Re-read the system prompt files from ``system_prompt_dir``.

    The streaming server exposes this via a management endpoint so
    operators can iterate on prompt templates without restarting
    workers.
    """
    self._system_prompts = self._load_system_prompts()
    logger.info("prompt enhancer: reloaded system prompts from %s", self._system_prompt_dir or "defaults")

def transition(self, target: SessionState) -> None:
    """Move to ``target`` if the edge is allowed.

    Raises :class:`InvalidSessionTransition` on illegal moves. The
    self-loop on ``ACTIVE`` is legal so the server can re-assert
    ACTIVE on segment completion without special casing.
    """
    allowed = _VALID_TRANSITIONS.get(self.state, frozenset())
    if target not in allowed and target is not self.state:
        raise InvalidSessionTransition(f"{self.state.value} -> {target.value} is not a valid "
                                       f"session transition")
    self.state = target
    self.last_activity = time.monotonic()

def reap_timed_out(self, now: float | None = None) -> list[str]:
    """Return the ids of sessions that have exceeded the idle timeout.

    The caller is responsible for actually closing them — this
    method only *identifies* dead sessions so the server can emit
    ``session_timeout`` frames before dropping the WebSocket.

    TODO: unused until a background driver calls it. Per-connection
    idle enforcement currently happens via asyncio.wait_for on
    receive_json; this helper catches sessions stuck before any
    receive (e.g. future QUEUED state) and is expected to be wired
    into the GPU-pool reaper.
    """
    now = now if now is not None else time.monotonic()
    dead: list[str] = []
    for sid, session in self._sessions.items():
        if session.state in {
                SessionState.COMPLETE,
                SessionState.ERROR,
                SessionState.TIMEOUT,
                SessionState.REJECTED,
        }:
            continue
        if now - session.last_activity > self._session_timeout_seconds:
            dead.append(sid)
    return dead

@abstractmethod
def drop(self, session_id: str) -> None:
    """Forget a session. Missing ids are a no-op."""

@abstractmethod
def hydrate(
    self,
    state: ContinuationState,
    *,
    session_id: str | None = None,
) -> str:
    """Install ``state`` as the starting point for a session.

    When ``session_id`` is ``None`` the store allocates a fresh id
    (UUID4); when provided the store uses it verbatim, overwriting
    any prior state at that id.
    """

@abstractmethod
def snapshot(self, session_id: str) -> ContinuationState | None:
    """Return the current state for ``session_id`` (or ``None``)."""

@abstractmethod
def store(self, session_id: str, state: ContinuationState) -> None:
    """Persist ``state`` for ``session_id``, replacing any prior value."""

async def start(self) -> None:
    """Spawn worker processes and wait for each to report ready."""
    num_workers = self._pool_config.num_workers or 1
    for gpu_id in range(num_workers):
        handle = self._worker_factory(
            gpu_id=gpu_id,
            generator_config=self._generator_config,
            warmup_config=self._warmup_config,
        )
        self._workers.append(handle)
        self._worker_by_id[handle.worker_id] = handle

    # Wait for each worker's ready event in a thread to avoid
    # blocking the event loop.
    loop = asyncio.get_running_loop()
    await asyncio.gather(*[
        loop.run_in_executor(None, handle.ready.wait, self._warmup_config.timeout_seconds)
        for handle in self._workers
    ])

    # Start background result readers — one task per worker
    # drains its result queue and resolves futures in _pending.
    for handle in self._workers:
        task = asyncio.create_task(self._drain_results(handle))
        self._result_reader_tasks.append(task)

    # Only admit workers that successfully booted. Anything that
    # failed boot (timeout, crash, error sentinel) stays out of the
    # available queue so we never assign a session to it.
    for idx, handle in enumerate(self._workers):
        if handle.boot_ok.is_set():
            await self._available.put(idx)
        else:
            logger.error(
                "pool: worker %s failed to boot; skipping",
                handle.worker_id,
            )