lightseekorg--tokenspeed/python/tokenspeed/runtime/pd/epd/encode_worker.py

# Copyright (c) 2026 LightSeek Foundation
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""Encode-worker control loop for EPD (Python orchestration).

This is the body the engine's encode event loop drives: it sits between request
arrival and the vision tower. On ``submit`` it registers the request's transfer
peer and, per item, either resolves the embedding from the cache (skip the
tower, still transfer) or queues it on the scheduler. Each ``step`` pulls one
deterministic batch off the scheduler, runs the tower + ships it via the
executor, and populates the cache.

The model load, mooncake manager construction, request transport and the
event-loop wiring are supplied by the engine integration; this class only
orchestrates them, so it is unit-testable with fakes.
"""

from __future__ import annotations

import dataclasses

from tokenspeed.runtime.cache.embedding_cache import (
    EmbeddingCache,
    TieredEmbeddingCache,
)
from tokenspeed.runtime.multimodal.embedder import _item_token_count
from tokenspeed.runtime.multimodal.inputs import MultimodalDataItem
from tokenspeed.runtime.multimodal.shm_transport import ShmTensorHandle
from tokenspeed.runtime.pd.epd.encode_scheduler import (
    EncodeScheduler,
    PendingEncodeItem,
)
from tokenspeed.runtime.utils import get_colorful_logger

logger = get_colorful_logger(__name__)


@dataclasses.dataclass(frozen=True)
class EncodeRequest:
    """One encode request: a transfer peer plus its vision items.

    ``bootstrap_host``/``port``/``room`` identify the prefill peer this
    request's embeddings are shipped to (assigned upstream, per request).
    """

    request_id: str
    bootstrap_host: str
    bootstrap_port: int
    bootstrap_room: int
    items: list[MultimodalDataItem]


def _nbytes(tensor) -> int:
    return tensor.numel() * tensor.element_size()


class EncodeWorker:
    """Orchestrates cache + scheduler + executor for the encode role.

    Injected with the executor (real ``DisaggEncodeExecutor`` or a fake), an
    ``EncodeScheduler`` and an embedding cache (single-tier ``EmbeddingCache`` or
    the two-tier ``TieredEmbeddingCache``; only ``get``/``put`` are used) so the
    control flow is testable without a GPU or transport.
    """

    def __init__(
        self,
        executor,
        scheduler: EncodeScheduler,
        cache: EmbeddingCache | TieredEmbeddingCache,
    ):
        self.executor = executor
        self.scheduler = scheduler
        self.cache = cache
        # (request_id, item_index) -> item awaiting the tower
        self._pending: dict = {}

    def submit(self, request: EncodeRequest) -> None:
        self.executor.register(
            request.request_id,
            request.bootstrap_host,
            request.bootstrap_port,
            request.bootstrap_room,
        )
        for idx, item in enumerate(request.items):
            cached = self.cache.get(item.hash)
            if isinstance(item.feature, ShmTensorHandle):
                # EPD pixel-SHM: the servicer published pixels to POSIX SHM and
                # the ZMQ hop carried only this handle (hash/pad_value were set on
                # the real tensor before publish). consume() unlinks, so segments
                # never outlive the item: materialize on a miss, and on a hit still
                # consume-and-drop to unlink the unused segment.
                handle, item.feature = item.feature, None
                handle.attach()
                if cached is None:
                    item.feature = handle.consume()
                else:
                    handle.consume()
            if cached is not None:
                # Cache hit: tower skipped, but the embedding still must reach
                # the prefill peer, so ship it directly. Entries are
                # (main, deepstack) pairs and BOTH halves must be restored, else
                # the prefill publishes a never-written deepstack buffer. Tolerate
                # a bare tensor for legacy/test-seeded entries.
                if isinstance(cached, tuple):
                    item.encoded, item.encoded_deepstack = cached
                else:
                    item.encoded = cached
                self.executor.send_item(request.request_id, item)
            else:
                self.scheduler.add(
                    PendingEncodeItem(
                        request_id=request.request_id,
                        item_index=idx,
                        cost=_item_token_count(item),
                    )
                )
                self._pending[(request.request_id, idx)] = item

    def step(self) -> int:
        """Run one scheduler batch through the tower + transfer. Returns the
        number of items encoded (0 when nothing is pending)."""
        self.executor.reap_concluded_senders({rid for (rid, _idx) in self._pending})
        # Retry sends that couldn't lease a ring slot last tick (non-blocking).
        self.executor.drain_deferred()
        # Backpressure: if sends are STILL deferred after the drain, the bounce
        # ring is saturated (all slots hold in-flight transfers). Pulling more ViT
        # now would only pile fresh embeddings into _deferred_sends -- each pins a
        # GPU embedding tensor with no slot to ship it, growing an unbounded
        # backlog into an OOM. Skip this tick; the loop yields the GIL (encode_loop
        # sees has_deferred) so the transfer daemons free slots, then we resume.
        if self.executor.has_deferred():
            return 0
        batch = self.scheduler.next_batch()
        if not batch:
            return 0
        request_items = [(p.request_id, self._pending[p.key]) for p in batch]
        try:
            self.executor.execute(request_items)
        except Exception as e:
            # A tower-step contract violation (ViT output not matching the items'
            # post-merge token count, or the forward itself) must fail only the
            # rooms in THIS batch, not propagate out into the engine's SIGUSR1
            # handler, which would kill the worker and drop every other request's
            # in-flight image. These raises fire before any send is issued, so
            # concluding the batch Failed never poisons an already-shipped room.
            # Per-item STAGING errors are handled finer-grained inside
            # _stage_and_send -> _fail_staged_room.
            n_failed = self.executor.fail_rooms((rid for rid, _ in request_items), e)
            for p in batch:
                self._pending.pop(p.key, None)
            logger.error(
                "encode batch failed (%d rooms concluded Failed): %s", n_failed, e
            )
            return 0
        for p in batch:
            item = self._pending.pop(p.key)
            if item.encoded is not None:
                # Cache the (main, deepstack) PAIR: caching only the main half
                # would make every later hit ship a deepstack-less transfer,
                # publishing uninitialized rows on the prefill.
                deep = item.encoded_deepstack
                nbytes = _nbytes(item.encoded) + (
                    _nbytes(deep) if deep is not None else 0
                )
                self.cache.put(item.hash, (item.encoded, deep), nbytes)
        return len(batch)

    def has_pending(self) -> bool:
        return self.scheduler.pending_size() > 0

    def has_deferred(self) -> bool:
        """True while sends are queued waiting for a free ring slot (executor)."""
        return self.executor.has_deferred()