lightseekorg--tokenspeed/python/tokenspeed/runtime/pd/epd/encode_executor.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.

"""EPD encode-worker execution: run the vision tower, scatter its output back
onto each item, and hand the contiguous embeddings to the Mooncake sender.
"""

from __future__ import annotations

import logging

import torch

from tokenspeed.runtime.multimodal.embedder import _item_token_count
from tokenspeed.runtime.multimodal.inputs import Modality, MultimodalDataItem
from tokenspeed.runtime.pd.base.status import TransferPoll
from tokenspeed.runtime.pd.epd.embedding_transfer import (
    MooncakeEmbeddingSender,
)
from tokenspeed.runtime.utils.env import envs

logger = logging.getLogger(__name__)


def assign_encoded_embeddings(
    items: list[MultimodalDataItem],
    output: torch.Tensor,
    model,
) -> None:
    """Scatter a packed vision-tower output onto each item, in place.

    ``output`` is the tower's ``[sum_tokens, width]`` result for ``items`` in
    order (``width = hidden`` for plain models, or ``hidden * (1 + n_deepstack)``
    for deepstack models like Qwen3.5). Each item's row span is its post-merge
    token count (``_item_token_count``); the rows are split accordingly and,
    for deepstack models, column-split via ``model.separate_deepstack_embeds``
    into the main ``[N, hidden]`` and deepstack ``[N, hidden * n_deepstack]``
    halves. Results are made contiguous because a TP-gathered tower output may
    not be, and the transfer ships raw row-major bytes.

    Sets ``item.encoded`` (and ``item.encoded_deepstack`` when the model emits
    deepstack, else ``None``), which is exactly the ``skip-ViT`` form the
    prefill-side VisionEmbedder consumes.
    """
    output = output.reshape(-1, output.shape[-1])
    per_item_tokens = [_item_token_count(item) for item in items]
    total = sum(per_item_tokens)
    if output.shape[0] != total:
        raise ValueError(
            f"vision-tower output has {output.shape[0]} rows but items sum to "
            f"{total} post-merge tokens; check the token-count / grid contract"
        )

    has_deepstack = getattr(model, "num_deepstack_embeddings", 0) > 0
    per_item_embeds = torch.split(output, per_item_tokens, dim=0)
    for item, emb in zip(items, per_item_embeds):
        if has_deepstack:
            main, deep = model.separate_deepstack_embeds(emb)
            item.encoded = main.contiguous()
            item.encoded_deepstack = deep.contiguous()
        else:
            item.encoded = emb.contiguous()
            item.encoded_deepstack = None


class DisaggEncodeExecutor:
    """Drives one encode worker: run the vision tower on a batch of items, then
    ship each item's embedding to its prefill peer over Mooncake.

    Python orchestration: this is invoked by the encode loop, not the
    C++ scheduler. ``execute`` groups items by modality, runs the tower once per
    modality via the model's ``get_image_feature`` / ``get_video_feature``,
    scatters the output onto ``item.encoded`` (see
    :func:`assign_encoded_embeddings`), and queues a transfer per item through
    the per-request :class:`MooncakeEmbeddingSender`. Each request must first be
    ``register``-ed with its prefill peer's bootstrap (host, port, room).
    """

    def __init__(
        self,
        manager,
        multimodal_model,
        device,
        *,
        ring_slots: int = 64,
        ring_bytes: int = 256 * 1024 * 1024,
    ):
        self.manager = manager
        self.model = multimodal_model
        self.device = device
        self.senders = {}
        # RDMA requires every transferred buffer to be a registered memory region,
        # and mooncake rejects OVERLAPPING registrations -- registering each
        # per-request ``item.encoded`` fails because the torch caching allocator
        # packs freed-but-still-registered tensors so a grown region straddles
        # others. Collapse every send through a fixed ring of pre-registered bounce
        # buffers: each slot is registered once at a fixed size (never grows, never
        # overlaps), and ``item.encoded`` is copied into a slot before its async
        # send. ``ring_slots`` / ``ring_bytes`` are injectable for tests and
        # env-tunable; total reservation is slots * slot_bytes PER ring (main, plus
        # deepstack if present), so depth and per-slot bytes must be sized to the
        # model and peak concurrency.
        self._ring_slots = int(
            envs.TOKENSPEED_EPD_ENCODE_RING_SLOTS.get_set_value_or(ring_slots)
        )
        slot_mb = envs.TOKENSPEED_EPD_ENCODE_RING_SLOT_MB.get()
        # Env override is in whole MiB; unset -> keep the exact ``ring_bytes`` arg.
        self._ring_bytes = slot_mb * 1024 * 1024 if slot_mb else ring_bytes
        self._main_ring = None  # lazily allocated on first send (device live by then)
        self._deep_ring = None
        self._ring_idx = 0
        # Per-slot lease: the room whose send last staged into the slot. A slot is
        # reusable only once that room's transfer is TERMINAL and not parked (a
        # parked chunk holds the slot's pointer until bootstrap_time_out and is
        # re-sent on late receiver registration; see _lease_slot), so a full ring
        # DEFERS the send rather than overwriting an in-flight slot.
        self._slot_rooms: list = [None] * self._ring_slots
        # Sends whose ViT output is ready but could not lease a free ring slot
        # (all slots still hold in-flight transfers). Retried non-blocking by
        # drain_deferred() each loop tick (a busy-wait here would GIL-starve the
        # daemon transfer-workers that free the slots and deadlock the loop).
        self._deferred_sends: list = []

    def register(self, request_id, bootstrap_host, bootstrap_port, bootstrap_room):
        self.senders[request_id] = MooncakeEmbeddingSender(
            self.manager, f"{bootstrap_host}:{bootstrap_port}", bootstrap_room
        )

    def _feature_fn(self, modality):
        # IMAGE dispatches through the model's ``image_encoder`` seam, NOT
        # ``get_image_feature`` directly: that seam is what the encoder CUDA-graph
        # wrapper overrides (see _maybe_install_encoder_cudagraph). When the graph
        # is disabled the model leaves ``image_encoder = get_image_feature`` (eager);
        # VIDEO has no captured graph -> always eager.
        if modality == Modality.IMAGE:
            return self.model.image_encoder
        if modality == Modality.VIDEO:
            return self.model.get_video_feature
        raise ValueError(f"unsupported modality for encode: {modality}")

    def execute(self, request_items: list[tuple[str, MultimodalDataItem]]) -> None:
        by_modality = {}
        for _, item in request_items:
            by_modality.setdefault(item.modality, []).append(item)
        with torch.inference_mode():
            for modality, items in by_modality.items():
                output = self._feature_fn(modality)(items)
                assign_encoded_embeddings(items, output, self.model)
        # Stage every embedding into its ring slot, then issue the async
        # Mooncake sends. See _stage_and_send for the copy/RDMA
        # overwrite-safety invariant (one CUDA event gates each transfer).
        self._stage_and_send(request_items)

    def _ensure_rings(self) -> None:
        """Lazily allocate + register the bounce-buffer ring (see ``__init__``)."""
        if self._main_ring is not None:
            return
        self._main_ring = [
            torch.empty(self._ring_bytes, dtype=torch.uint8, device=self.device)
            for _ in range(self._ring_slots)
        ]
        for buf in self._main_ring:
            self.manager.engine.register(buf.data_ptr(), self._ring_bytes)

    def _copy_into(self, ring, slot: int, src) -> tuple[int, int]:
        """Copy ``src``'s bytes into pre-registered ring ``slot``; return its
        (device pointer, byte length). Fails loud if an embedding exceeds a slot."""
        nbytes = src.numel() * src.element_size()
        if nbytes > self._ring_bytes:
            raise RuntimeError(
                f"EPD encode embedding {nbytes} B exceeds ring slot "
                f"{self._ring_bytes} B; raise TOKENSPEED_EPD_ENCODE_RING_SLOT_MB "
                "or the ring_bytes constructor argument"
            )
        buf = ring[slot]
        buf[:nbytes].view(src.dtype).copy_(src.reshape(-1))
        return buf.data_ptr(), nbytes

    def _lease_slot(self) -> "int | None":
        """Return a reusable ring-slot index, or ``None`` if every slot still
        holds an in-flight transfer. NON-BLOCKING: the caller DEFERS rather than
        spinning (a busy-wait would GIL-starve the daemon transfer-workers that
        mark rooms terminal and deadlock the single-threaded loop).

        A slot is reusable once the room it last staged is TERMINAL (Success,
        Failed, or None=already reaped) AND no parked chunk still holds its
        pointer -- the overwrite-safety invariant."""
        mgr = self.manager
        n = self._ring_slots
        for _ in range(n):
            slot = self._ring_idx % n
            self._ring_idx += 1
            room = self._slot_rooms[slot]
            if room is None:
                return slot
            status = mgr.room_status(room)
            if status is None or status in (
                TransferPoll.Success,
                TransferPoll.Failed,
            ):
                if not mgr.is_parked(room):
                    return slot
        return None

    def _stage_and_send(self, items: list[tuple[str, MultimodalDataItem]]) -> None:
        """Lease a ring slot per item and ship it; items that cannot lease a free
        slot (ring full) are DEFERRED for a later non-blocking retry rather than
        blocking the loop. Stages every leased item then issues ONE stream sync
        before the sends, so the one-sided RDMA reads never race the device-to-
        device copies (the ViT->send corruption hazard)."""
        self._ensure_rings()
        staged = []
        for rid, item in items:
            if rid not in self.senders:
                # Sender reaped (its room concluded/failed) -- drop this stale
                # deferred send instead of crashing on senders[rid].
                continue
            slot = self._lease_slot()
            if slot is None:
                self._deferred_sends.append((rid, item))
                continue
            try:
                send_args = self._stage_item(
                    item, self.senders[rid].bootstrap_room, slot
                )
            except Exception as e:
                # A staging error (most plausibly _copy_into rejecting an embedding
                # larger than a ring slot) must fail only THIS item's room, never
                # raise out of the single-threaded encode loop into the engine's
                # SIGUSR1 handler (which kills the whole worker and every other
                # in-flight image). Covers the unguarded send_item() /
                # drain_deferred() callers too; the leased slot returns to the ring
                # once the room is Failed (see _lease_slot).
                self._fail_staged_room(rid, e)
                continue
            staged.append((rid, send_args))
        if not staged:
            return
        # Record ONE CUDA event after all the ring copies above (they ran on the
        # current stream inside _stage_item) and hand it to each transfer rather
        # than host-syncing on this single encode-loop thread. The daemon
        # transfer-worker waits the event before its one-sided RDMA read
        # (embedding_transfer._transfer_worker), so the read never races the copy;
        # _lease_slot keeps the slot until its room is terminal (Success only after
        # the RDMA completes).
        copy_event = None
        if torch.cuda.is_available():
            copy_event = torch.cuda.Event()
            copy_event.record()
        for rid, send_args in staged:
            self.senders[rid].send(copy_event=copy_event, **send_args)

    def drain_deferred(self) -> None:
        """Retry deferred sends (ViT done, waiting for a free ring slot). Non-
        blocking: items that still cannot lease a slot stay deferred. Driven once
        per encode-loop tick; the loop yields the GIL between ticks so the daemon
        transfer-workers can free slots for the next drain."""
        if not self._deferred_sends:
            return
        pending = self._deferred_sends
        self._deferred_sends = []
        self._stage_and_send(pending)

    def has_deferred(self) -> bool:
        return bool(self._deferred_sends)

    def _conclude_room_failed(self, room: int, exc: Exception) -> None:
        """Push Failed to ``room``'s prefill receivers so they abort via the
        rank-synced admission path, instead of the error escaping the encode loop
        and SIGUSR1-ing the worker. The single seam every failure path goes
        through; delegates to the manager's public ``fail_room`` rather than
        reaching into its ``transfer_infos`` / status FSM."""
        self.manager.fail_room(room, str(exc))

    def _fail_staged_room(self, rid: str, exc: Exception) -> None:
        """Per-item staging failure (the unguarded ``send_item`` /
        ``drain_deferred`` callers): conclude ``rid``'s room Failed."""
        sender = self.senders.get(rid)
        if sender is None:
            return
        self._conclude_room_failed(sender.bootstrap_room, exc)
        logger.error(
            "encode staging failed for room %s: %s", sender.bootstrap_room, exc
        )

    def fail_rooms(self, request_ids, exc: Exception) -> int:
        """Conclude every room owned by ``request_ids`` Failed; return the count.
        The owning seam for a batch-level failure that fired before any send was
        issued (ViT / assign_encoded_embeddings): the worker hands its batch's
        request_ids and stays out of the sender/manager internals. Rooms are
        de-duped (a multi-image request shares one room)."""
        rooms = set()
        for rid in request_ids:
            sender = self.senders.get(rid)
            if sender is not None:
                rooms.add(sender.bootstrap_room)
        for room in rooms:
            self._conclude_room_failed(room, exc)
        return len(rooms)

    def reap_concluded_senders(self, pending_request_ids) -> None:
        """Drop per-request senders whose room reached a terminal transfer status
        (the ``senders`` dict otherwise grows forever). Senders whose request_id
        is still awaiting the tower (``pending_request_ids``) are kept -- their
        send has not been queued. Only the sender is dropped; the manager's
        terminal ``request_status`` tombstone stays (the transfer worker's
        straggler-drop and the ring-slot lease both key on it)."""
        for rid in list(self.senders):
            if rid in pending_request_ids:
                continue
            room = self.senders[rid].bootstrap_room
            if self.manager.room_status(room) in (
                TransferPoll.Success,
                TransferPoll.Failed,
            ):
                self.senders.pop(rid, None)

    def _stage_item(self, item: MultimodalDataItem, room, slot: int) -> dict:
        """Copy one item's embedding (and deepstack half, if any) into the leased
        ring ``slot`` and return the scalar ``send`` kwargs. The copy runs on the
        current stream; the CALLER must synchronize before handing these pointers
        to the transfer engine, so the one-sided RDMA read never races the device-
        to-device copy (same hazard class as the ViT->send race)."""
        enc = item.encoded
        self._slot_rooms[slot] = room
        send_ptr, nbytes = self._copy_into(self._main_ring, slot, enc)
        ds_ptr = ds_width = ds_nbytes = 0
        deep = item.encoded_deepstack
        if deep is not None and deep.numel() > 0:
            if self._deep_ring is None:
                self._deep_ring = [
                    torch.empty(self._ring_bytes, dtype=torch.uint8, device=self.device)
                    for _ in range(self._ring_slots)
                ]
                for buf in self._deep_ring:
                    self.manager.engine.register(buf.data_ptr(), self._ring_bytes)
            ds_width = deep.shape[1]
            ds_ptr, ds_nbytes = self._copy_into(self._deep_ring, slot, deep)
        return dict(
            src_embedding_ptr=send_ptr,
            n_tokens=enc.shape[0],
            hidden=enc.shape[1],
            dtype=str(enc.dtype),
            nbytes=nbytes,
            src_deepstack_ptr=ds_ptr,
            deepstack_width=ds_width,
            deepstack_nbytes=ds_nbytes,
        )

    def send_item(self, request_id, item: MultimodalDataItem) -> None:
        """Ship an already-encoded item (``item.encoded`` set) to its prefill peer.
        Used by the encode loop for cache hits, which skip the tower but still
        transfer. Routes through the same lease-or-defer path as ``execute`` so a
        full ring defers (non-blocking) instead of stalling the loop."""
        self._stage_and_send([(request_id, item)])