sgl-project--sglang/python/sglang/multimodal_gen/runtime/disaggregation/scheduler_mixin.py

# SPDX-License-Identifier: Apache-2.0
"""Mixin that adds disaggregated diffusion scheduling to the Scheduler.

Extracted from scheduler.py to keep the core scheduler lean.
All transfer, compute, and event-loop logic for disaggregated roles
(encoder / denoiser / decoder) lives here.
"""

from __future__ import annotations

import contextlib
import dataclasses
import inspect
import json
import logging
import pickle
import queue
import threading
import time
from typing import TYPE_CHECKING, Any

import torch
import zmq

from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
from sglang.multimodal_gen.runtime.disaggregation.transport.buffer import (
    TransferTensorBuffer,
)
from sglang.multimodal_gen.runtime.disaggregation.transport.codec import (
    send_tensors,
)
from sglang.multimodal_gen.runtime.disaggregation.transport.engine import (
    create_transfer_engine,
)
from sglang.multimodal_gen.runtime.disaggregation.transport.manager import (
    DiffusionTransferManager,
)
from sglang.multimodal_gen.runtime.disaggregation.transport.protocol import (
    TRANSFER_MAGIC,
    TransferAllocatedMsg,
    TransferDoneMsg,
    TransferMsgType,
    TransferPushedMsg,
    TransferRegisterMsg,
    decode_transfer_msg,
    encode_transfer_msg,
    is_transfer_message,
)
from sglang.multimodal_gen.runtime.pipelines_core import Req
from sglang.multimodal_gen.runtime.pipelines_core.diffusion_scheduler_utils import (
    clone_scheduler_runtime,
)
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.utils.common import get_zmq_socket
from sglang.multimodal_gen.runtime.utils.distributed import broadcast_pyobj
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.runtime.utils.trace_wrapper import DiffStage, trace_slice
from sglang.srt.observability.trace import TraceReqContext

if TYPE_CHECKING:
    from sglang.multimodal_gen.runtime.managers.scheduler import Scheduler

logger = init_logger(__name__)

# ---------------------------------------------------------------------------
# Field extraction: split Req into tensors (transfer buffer) and scalars (JSON)
# ---------------------------------------------------------------------------

# Fields that should never be transferred (non-serializable, internal, or receiver rebuilds)
_EXCLUDE_FIELDS = frozenset(
    {
        "sampling_params",
        "generator",
        "modules",
        "metrics",
        "extra_step_kwargs",
        "extra",
        "condition_image",
        "vae_image",
        "pixel_values",
        "preprocessed_image",
        "image_embeds",
        "original_condition_image_size",
        "vae_image_sizes",
        "output",
        "audio",
        "audio_sample_rate",
        "trajectory_timesteps",
        "trajectory_latents",
        "trajectory_audio_latents",
        "timestep",
        "step_index",
        # Request scheduler is a local runtime object cloned from the pipeline
        # scheduler template. It may hold live mutable state and is not JSON-safe.
        "scheduler",
        "prompt_template",
        "max_sequence_length",
        # trace_ctx holds live OTel SDK objects that aren't JSON-serializable.
        # We propagate tracing across the JSON hop via a separate, JSON-safe
        # ``_trace_state`` scalar field built from ``TraceReqContext.__getstate__``
        # (same W3C carrier SRT relies on for pickle transport) and rebuild it
        # on the receiver in ``_build_disagg_req``.
        "trace_ctx",
    }
)

# Sampling-params fields that should never be transferred across roles:
#   - data_type / supported_resolutions: enums / non-JSON classvars reconstructed on the receiver
#   - teacache_params: model-specific object, not JSON-safe
#   - output_* / save_output / return_*: output-side concerns owned by the decoder role
#
# Everything else on SamplingParams is forwarded automatically via a field-walk
# below; this keeps new request-level features (e.g. Qwen-Image's
# true_cfg_scale, guidance_rescale, cfg_normalization, ...) from silently
# getting dropped just because nobody remembered to add them to a whitelist.
_SAMPLING_PARAMS_EXCLUDE_FIELDS = frozenset(
    {
        "data_type",
        "supported_resolutions",
        "teacache_params",
    }
)

_BASE_SP_DEFAULTS: dict[str, Any] = {}
for _f in dataclasses.fields(SamplingParams):
    if _f.default is not dataclasses.MISSING:
        _BASE_SP_DEFAULTS[_f.name] = _f.default


def _is_tensor_like(value) -> bool:
    if isinstance(value, torch.Tensor):
        return True
    if isinstance(value, list) and value and isinstance(value[0], torch.Tensor):
        return True
    return False


def _to_json_serializable(value):
    if isinstance(value, torch.Tensor):
        return value.tolist()
    if isinstance(value, (list, tuple)):
        converted = []
        for item in value:
            if isinstance(item, torch.Tensor):
                converted.append(item.tolist())
            else:
                converted.append(item)
        return converted
    return value


def _is_default(value, field_info) -> bool:
    if field_info.default is not dataclasses.MISSING:
        return value == field_info.default
    if field_info.default_factory is not dataclasses.MISSING:
        if isinstance(value, (list, dict)) and len(value) == 0:
            return True
    return False


def _extract_extra_fields(extra: dict, scalar_fields: dict) -> None:
    """Extract JSON-serializable entries from Req.extra into scalar_fields."""
    for key, value in extra.items():
        if key.startswith("_"):
            continue
        try:
            json.dumps(value)
            scalar_fields[f"_extra_{key}"] = value
        except (TypeError, ValueError, OverflowError):
            pass


def _init_request_scheduler_from_template(
    scheduler_template: Any, req: Req, device: torch.device
) -> None:
    scheduler = clone_scheduler_runtime(scheduler_template)
    extra_kwargs = {}
    mu = req.extra.get("mu") if hasattr(req, "extra") else None
    if mu is not None:
        extra_kwargs["mu"] = mu

    set_timesteps_params = inspect.signature(scheduler.set_timesteps).parameters
    timesteps = getattr(req, "timesteps", None)
    sigmas = getattr(req, "sigmas", None)
    num_steps = getattr(req, "num_inference_steps", None)

    if sigmas is not None and "sigmas" in set_timesteps_params:
        if isinstance(sigmas, torch.Tensor):
            sigmas = sigmas.detach().cpu()
        scheduler.set_timesteps(sigmas=sigmas, device=device, **extra_kwargs)
    elif timesteps is not None and "timesteps" in set_timesteps_params:
        if isinstance(timesteps, torch.Tensor):
            timesteps = timesteps.detach().cpu()
        scheduler.set_timesteps(timesteps=timesteps, device=device, **extra_kwargs)
    elif num_steps is not None:
        scheduler.set_timesteps(num_steps, device=device, **extra_kwargs)
    else:
        return

    req.scheduler = scheduler
    req.timesteps = scheduler.timesteps


def _init_disagg_request_scheduler(self: Scheduler, req: Req) -> None:
    scheduler_template = self.worker.pipeline.get_module("scheduler")
    if scheduler_template is None:
        return
    device = torch.device(f"{current_platform.device_type}:{self.worker.local_rank}")
    _init_request_scheduler_from_template(scheduler_template, req, device)


def extract_transfer_fields(req) -> tuple[dict, dict]:
    """Extract all transferable fields from a Req, split into tensors and scalars."""
    tensor_fields = {}
    scalar_fields = {}
    _debug_transfer = logger.isEnabledFor(logging.DEBUG)

    for f in dataclasses.fields(req):
        if f.name in _EXCLUDE_FIELDS:
            continue

        value = getattr(req, f.name, None)
        if value is None:
            continue
        if _is_default(value, f):
            continue

        if _is_tensor_like(value):
            tensor_fields[f.name] = value
        else:
            try:
                scalar_fields[f.name] = _to_json_serializable(value)
            except (TypeError, ValueError):
                pass

    extra = getattr(req, "extra", None)
    if extra:
        _extract_extra_fields(extra, scalar_fields)

    sp = getattr(req, "sampling_params", None)
    if sp is not None:
        # Forward every non-default, JSON-safe SamplingParams field, not a
        # narrow whitelist. Previously only a handful of fields were carried
        # across roles, which silently dropped per-request config like
        # Qwen-Image's true_cfg_scale (and any future feature added to
        # SamplingParams). Using a field-walk keeps the disagg boundary
        # feature-complete without needing to edit this list.
        for f in dataclasses.fields(sp):
            name = f.name
            if name in _SAMPLING_PARAMS_EXCLUDE_FIELDS:
                continue
            if name in scalar_fields:
                # Req-level field already took precedence (or upstream Req
                # explicitly set it).
                continue
            value = getattr(sp, name, None)
            if value is None:
                continue
            base_default = _BASE_SP_DEFAULTS.get(name, dataclasses.MISSING)
            if base_default is not dataclasses.MISSING and value == base_default:
                continue
            try:
                scalar_fields[name] = _to_json_serializable(value)
            except (TypeError, ValueError):
                pass

    if getattr(req, "generator", None) is not None:
        seed = getattr(req, "seed", None)
        if seed is not None:
            scalar_fields["seed"] = _to_json_serializable(seed)

    if _debug_transfer:
        import torch as _torch

        for _n, _t in tensor_fields.items():
            if isinstance(_t, _torch.Tensor):
                _sz = _t.nelement() * _t.element_size()
                logger.debug(
                    "transfer_field %s shape=%s dtype=%s size=%d",
                    _n,
                    list(_t.shape),
                    _t.dtype,
                    _sz,
                )
            elif isinstance(_t, list):
                for _i, _ti in enumerate(_t):
                    if isinstance(_ti, _torch.Tensor):
                        _sz = _ti.nelement() * _ti.element_size()
                        logger.debug(
                            "transfer_field %s[%d] shape=%s dtype=%s size=%d",
                            _n,
                            _i,
                            list(_ti.shape),
                            _ti.dtype,
                            _sz,
                        )

    # Propagate OTel trace context over the JSON hop. TraceReqContext.__getstate__
    # reduces the live context to a JSON-safe dict (W3C traceparent/tracestate in
    # root_span_context). Receiver rebuilds via __setstate__ in _build_disagg_req.
    trace_ctx = getattr(req, "trace_ctx", None)
    if trace_ctx is not None and getattr(trace_ctx, "tracing_enable", False):
        try:
            trace_state = trace_ctx.__getstate__()
            if trace_state and trace_state.get("tracing_enable"):
                scalar_fields["_trace_state"] = trace_state
        except Exception as e:
            logger.debug("Failed to export trace state: %s", e)

    return tensor_fields, scalar_fields


# ---------------------------------------------------------------------------
# Helpers for broadcasting Req contents across SP/CFG/TP ranks
# ---------------------------------------------------------------------------

# Sentinel marker key used to distinguish "list of tensors" from a regular
# nested dict when round-tripping through GroupCoordinator.broadcast_tensor_dict
# (which only natively understands tensor / nested-dict values).
_LIST_MARKER_KEY = "__is_list__"


def _pack_tensor_fields_for_broadcast(tensor_fields: dict) -> dict:
    """Pack ``tensor_fields`` into a structure ``broadcast_tensor_dict`` accepts.

    ``broadcast_tensor_dict`` understands dict-of-tensor values (recursively),
    but not list-of-tensor values. Several Req fields (``prompt_embeds``,
    ``image_embeds``, ...) are lists of tensors, so we encode each list as a
    nested dict whose tensors are keyed by their stringified index, with a
    sentinel ``__is_list__`` flag to disambiguate from real nested dicts.
    """
    packed: dict = {}
    for key, value in tensor_fields.items():
        if isinstance(value, torch.Tensor):
            packed[key] = value
        elif isinstance(value, list):
            sub: dict = {_LIST_MARKER_KEY: True}
            for i, item in enumerate(value):
                if isinstance(item, torch.Tensor):
                    sub[str(i)] = item
            packed[key] = sub
        # Anything else (e.g. None, scalars) is intentionally dropped — the
        # scalar_fields broadcast covers non-tensor metadata.
    return packed


def _unpack_tensor_fields_from_broadcast(packed: dict) -> dict:
    """Inverse of :func:`_pack_tensor_fields_for_broadcast`."""
    out: dict = {}
    for key, value in packed.items():
        if isinstance(value, dict) and value.get(_LIST_MARKER_KEY) is True:
            indexed = [(int(k), v) for k, v in value.items() if k != _LIST_MARKER_KEY]
            indexed.sort(key=lambda kv: kv[0])
            out[key] = [v for _, v in indexed]
        else:
            out[key] = value
    return out


class SchedulerDisaggMixin:
    """Disaggregated diffusion scheduling: transfer, compute, event loops."""

    # ------------------------------------------------------------------
    # Initialization
    # ------------------------------------------------------------------

    def _init_disagg_state(self: Scheduler, server_args, local_rank: int) -> None:
        """Initialize all disaggregation state, sockets, and transfer infrastructure."""
        from sglang.multimodal_gen.runtime.disaggregation.metrics import DisaggMetrics

        self._disagg_role = server_args.disagg_role
        self._disagg_timeout_s = float(getattr(server_args, "disagg_timeout", 600))
        self._disagg_metrics = None
        self._disagg_mode = getattr(server_args, "disagg_mode", False)
        self._pool_work_pull = None
        self._pool_result_push = None
        self._transfer_manager = None
        self._transfer_stream = None
        self._rdma_push_queue = None
        self._rdma_push_thread = None
        self._rdma_push_zmq = None
        self._compute_ready_queue = None
        self._recv_prefetch_thread = None

        if self._disagg_role != RoleType.MONOLITHIC:
            self._disagg_metrics = DisaggMetrics(role=self._disagg_role.value)
            device = torch.device(f"{current_platform.device_type}:{local_rank}")
            self._transfer_stream = torch.get_device_module().Stream(device=device)
            self._init_disagg_sockets()
            self._init_disagg_transfer_manager()

    def _init_disagg_sockets(self: Scheduler):
        """Initialize ZMQ sockets for disaggregated mode (DiffusionServer-mediated).

        Only rank 0 creates ZMQ sockets. Non-rank-0 processes participate
        via NCCL broadcast from rank 0 (see _disagg_recv_work).
        """
        if self.gpu_id != 0:
            logger.info(
                "Pool mode %s rank %d: no ZMQ sockets (non-rank-0)",
                self._disagg_role.value.upper(),
                self.gpu_id,
            )
            return

        sa = self.server_args

        # PULL: receive work from DiffusionServer
        self._pool_work_pull, _ = get_zmq_socket(
            self.context,
            zmq.PULL,
            sa.pool_work_endpoint,
            bind=True,
            max_bind_retries=5,
            same_port=True,
        )
        # PUSH: send results to DiffusionServer
        self._pool_result_push, _ = get_zmq_socket(
            self.context, zmq.PUSH, sa.pool_result_endpoint, bind=False
        )
        logger.info(
            "Disagg %s rank 0: work_pull=%s, result_push=%s",
            self._disagg_role.value.upper(),
            sa.pool_work_endpoint,
            sa.pool_result_endpoint,
        )

    def _init_disagg_transfer_manager(self: Scheduler):
        """Initialize TransferManager for transfer mode (rank 0 only).

        Creates a TransferTensorBuffer (pinned memory pool) and a
        BaseTransferEngine, then wraps them in a DiffusionTransferManager.
        Also sends a transfer_register message to DiffusionServer.
        """
        if self.gpu_id != 0:
            return

        sa = self.server_args

        # Pool size: configurable, default 256 MiB
        pool_size = getattr(sa, "disagg_transfer_pool_size", 256 * 1024 * 1024)

        # Create transfer engine.
        # NOTE: self.gpu_id is the role-internal rank (0..num_role_gpus-1),
        # not the physical GPU index. In disagg mode with --base-gpu-id > 0,
        # the physical device is self.worker.local_rank. Mooncake needs the
        # physical index to pin the right NIC and register GPUDirect buffers.
        hostname = getattr(sa, "disagg_p2p_hostname", "127.0.0.1")
        ib_device = getattr(sa, "disagg_ib_device", None)
        physical_gpu_id = self.worker.local_rank
        engine = create_transfer_engine(
            hostname=hostname,
            gpu_id=physical_gpu_id,
            ib_device=ib_device,
        )

        # Use GPU buffer when engine supports GPUDirect RDMA, CPU pinned otherwise
        device = (
            f"{current_platform.device_type}:{physical_gpu_id}"
            if engine.supports_gpu_direct
            else "cpu"
        )
        buffer = TransferTensorBuffer(
            pool_size=pool_size, device=device, role_name=self._disagg_role.value
        )

        # Create transfer manager
        self._transfer_manager = DiffusionTransferManager(engine=engine, buffer=buffer)

        # Pre-allocate receive slots for receivers (denoiser/decoder)
        self._preallocated_slots: dict[int, object] = {}
        preallocated_slot_info = []
        if self._disagg_role in (RoleType.DENOISER, RoleType.DECODER):
            capacity = getattr(sa, "disagg_prealloc_slots", 2)
            typical_size = 64 * 1024 * 1024  # 64 MiB per slot
            for i in range(capacity):
                slot = buffer.allocate(typical_size, f"prealloc_{i}")
                if slot is not None:
                    self._preallocated_slots[i] = slot
                    preallocated_slot_info.append(
                        {
                            "offset": slot.offset,
                            "size": slot.size,
                            "slot_id": i,
                            "addr": self._transfer_manager.pool_data_ptr + slot.offset,
                        }
                    )
            if preallocated_slot_info:
                logger.info(
                    "Transfer %s: pre-allocated %d receive slots",
                    self._disagg_role.value.upper(),
                    len(preallocated_slot_info),
                )

        # Register with DiffusionServer.
        # Include our own work_endpoint so DS can key the peer by URL index,
        # not by registration order (startup order is not guaranteed to match
        # --encoder/denoiser/decoder-urls ordering).
        register_msg = TransferRegisterMsg(
            role=self._disagg_role.value,
            session_id=self._transfer_manager.session_id,
            pool_ptr=self._transfer_manager.pool_data_ptr,
            pool_size=self._transfer_manager.pool_size,
            work_endpoint=sa.pool_work_endpoint,
            preallocated_slots=preallocated_slot_info,
        )
        self._pool_result_push.send_multipart(encode_transfer_msg(register_msg))
        logger.info(
            "Transfer %s: registered with DS (session=%s, pool=%d bytes, prealloc=%d)",
            self._disagg_role.value.upper(),
            self._transfer_manager.session_id,
            pool_size,
            len(preallocated_slot_info),
        )

        # RDMA push thread for sender roles (encoder/denoiser)
        if self._disagg_role in (RoleType.ENCODER, RoleType.DENOISER):
            self._rdma_push_queue = queue.Queue(maxsize=4)
            self._rdma_push_zmq, _ = get_zmq_socket(
                self.context,
                zmq.PUSH,
                sa.pool_result_endpoint,
                bind=False,
            )
            self._rdma_push_thread = threading.Thread(
                target=self._rdma_push_loop,
                daemon=True,
                name=f"rdma-push-{self._disagg_role.value}",
            )
            self._rdma_push_thread.start()
            logger.info(
                "Transfer %s: RDMA push thread started",
                self._disagg_role.value.upper(),
            )

        # Recv prefetch thread for receiver roles (denoiser/decoder)
        # Rank 0 only (bg thread does ZMQ recv + load; multi-rank gets
        # scalar fields via broadcast_pyobj from the main thread).
        if self._disagg_role in (RoleType.DENOISER, RoleType.DECODER):
            self._compute_ready_queue = queue.Queue(maxsize=4)
            self._recv_prefetch_thread = threading.Thread(
                target=self._recv_prefetch_loop,
                daemon=True,
                name=f"recv-prefetch-{self._disagg_role.value}",
            )
            self._recv_prefetch_thread.start()
            logger.info(
                "Transfer %s: recv prefetch thread started",
                self._disagg_role.value.upper(),
            )

    # ------------------------------------------------------------------
    # Background threads
    # ------------------------------------------------------------------

    def _rdma_push_loop(self: Scheduler):
        """Background thread: execute RDMA push + notify DS.

        Runs push_to_peer (blocking RDMA) on a dedicated thread so the
        main event loop can immediately start processing the next request.
        """
        role_name = self._disagg_role.value.upper()
        while True:
            item = self._rdma_push_queue.get()
            if item is None:
                break  # Shutdown signal
            request_id, dest_session_id, dest_addr, transfer_size = item
            try:
                success = self._transfer_manager.push_to_peer(
                    request_id=request_id,
                    dest_session_id=dest_session_id,
                    dest_addr=dest_addr,
                    transfer_size=transfer_size,
                )
                if success:
                    self._transfer_manager.free_staged(request_id)

                pushed_msg = TransferPushedMsg(request_id=request_id)
                self._rdma_push_zmq.send_multipart(encode_transfer_msg(pushed_msg))

                if not success:
                    logger.error(
                        "Transfer %s: RDMA push failed for %s", role_name, request_id
                    )
            except Exception:
                logger.exception(
                    "Transfer %s: RDMA push thread error for %s", role_name, request_id
                )

    def _recv_prefetch_loop(self: Scheduler):
        """Background thread: recv transfer messages and prefetch tensor loads.

        For transfer_ready: loads tensors + builds Req in this thread, then
        enqueues the ready-to-compute item. This allows loading of request N+1
        to overlap with compute of request N on the main thread.

        For transfer_alloc/push: passes them through to the main thread for handling.
        """
        role_name = self._disagg_role.value.upper()
        while self._running:
            try:
                raw_frames = self._pool_work_pull.recv_multipart()
                frames = [bytes(f) for f in raw_frames]

                msg = decode_transfer_msg(frames)
                msg_type = msg.get("msg_type", "")

                if msg_type == TransferMsgType.READY:
                    # Prefetch: load tensors + build Req in this thread
                    item = self._prefetch_transfer_ready(msg)
                    self._compute_ready_queue.put(("transfer_compute", item))
                elif msg_type == TransferMsgType.PUSH:
                    # Handle push directly in prefetch thread — it only
                    # enqueues to the RDMA bg thread (thread-safe queue).
                    # Critical for pipeline parallelism: if deferred to the
                    # main thread, this gets blocked behind the next request's
                    # GPU compute, preventing the previous request's output
                    # from reaching the decoder.
                    self._handle_transfer_push(msg)
                else:
                    # alloc and other messages: pass to main thread
                    # (alloc sends on _pool_result_push which isn't thread-safe)
                    self._compute_ready_queue.put(("transfer_control", frames))

            except zmq.ZMQError as e:
                if not self._running:
                    break
                logger.error("Transfer %s recv prefetch: ZMQ error: %s", role_name, e)
            except Exception:
                logger.exception("Transfer %s recv prefetch: error", role_name)

    def _prefetch_transfer_ready(self: Scheduler, msg: dict) -> tuple:
        """Load tensors from transfer buffer and build Req for a transfer_ready message.

        Called from the recv prefetch thread. Loads on _transfer_stream
        and builds the Req, so the main thread can start compute immediately.

        Returns (req, load_event, request_id, role_name, prealloc_slot_id).
        """
        request_id = msg["request_id"]
        manifest = msg.get("manifest", {})
        scalar_fields = msg.get("scalar_fields", {})
        role_name = self._disagg_role.value.upper()

        if self._disagg_metrics:
            self._disagg_metrics.record_request_start(request_id)

        # Pre-allocated slot handling
        prealloc_slot_id = scalar_fields.pop("_prealloc_slot_id", None)
        if (
            prealloc_slot_id is not None
            and prealloc_slot_id in self._preallocated_slots
        ):
            slot = self._preallocated_slots[prealloc_slot_id]
            self._transfer_manager.register_prealloc_as_receive(request_id, slot)

        # Load tensors on transfer_stream (non-blocking)
        local_device = f"{current_platform.device_type}:{self.worker.local_rank}"
        tensors, load_event = self._transfer_manager.load_tensors_async(
            request_id,
            manifest,
            device=local_device,
            stream=self._transfer_stream,
        )

        # NOTE: Do NOT free the receive slot here. The async load is still
        # in progress. The slot must remain valid until the main thread waits
        # on load_event. Freeing is done in _disagg_prefetch_event_loop.

        # Build Req (CPU work, overlapped with load)
        req = self._build_disagg_req(scalar_fields, tensors)

        # NOTE: Do NOT call scheduler_mod.set_timesteps() here!
        # This runs on the prefetch thread. set_timesteps mutates shared
        # scheduler state (self.sigmas), which would corrupt the currently
        # running denoising loop on the main thread. Deferred to main thread
        # in _disagg_prefetch_event_loop, right before compute.

        return (req, load_event, request_id, role_name, prealloc_slot_id, scalar_fields)

    # ------------------------------------------------------------------
    # Broadcast
    # ------------------------------------------------------------------

    def _broadcast_to_all_ranks(self: Scheduler, data):
        """Broadcast *data* from rank 0 to all other ranks.

        Rank 0 passes the real payload; non-rank-0 passes ``None``.
        Broadcasts through all applicable groups (SP, CFG, TP).
        """
        sa = self.server_args

        if sa.sp_degree != 1:
            data = broadcast_pyobj(
                data,
                self.worker.sp_group.rank,
                self.worker.sp_cpu_group,
                src=self.worker.sp_group.ranks[0],
            )

        if sa.enable_cfg_parallel:
            data = broadcast_pyobj(
                data,
                self.worker.cfg_group.rank,
                self.worker.cfg_cpu_group,
                src=self.worker.cfg_group.ranks[0],
            )

        if sa.tp_size > 1:
            data = broadcast_pyobj(
                data,
                self.worker.tp_group.rank,
                self.worker.tp_cpu_group,
                src=self.worker.tp_group.ranks[0],
            )

        return data

    def _is_multi_rank(self: Scheduler) -> bool:
        sa = self.server_args
        return sa.sp_degree != 1 or sa.tp_size > 1 or sa.enable_cfg_parallel

    def _broadcast_tensor_dict_to_all_ranks(
        self: Scheduler, tensor_dict: dict | None
    ) -> dict | None:
        """Broadcast a tensor dict from rank 0 to non-rank-0 via NCCL.

        Uses ``GroupCoordinator.broadcast_tensor_dict`` which ships tensor
        metadata over the CPU group and the tensor payload over the device
        (NCCL) group, so large GPU buffers never bounce through CPU.
        """
        sa = self.server_args

        if sa.sp_degree != 1:
            tensor_dict = self.worker.sp_group.broadcast_tensor_dict(tensor_dict, src=0)
        if sa.enable_cfg_parallel:
            tensor_dict = self.worker.cfg_group.broadcast_tensor_dict(
                tensor_dict, src=0
            )
        if sa.tp_size > 1:
            tensor_dict = self.worker.tp_group.broadcast_tensor_dict(tensor_dict, src=0)
        return tensor_dict

    def _broadcast_req_to_all_ranks(self: Scheduler, req: Req | None) -> Req | None:
        """Broadcast a fully-loaded Req (scalars + GPU tensors) from rank 0.

        Required for multi-rank denoiser/decoder in disagg mode: only rank 0
        owns the TransferManager and RDMA-loads tensors into GPU memory. All
        other ranks must see the same Req before entering ``execute_forward``,
        otherwise REPLICATED stages (e.g. denoising) blow up on empty tensor
        fields because ``ParallelExecutor`` never broadcasts the batch for
        that paradigm.

        Tensor fields travel over NCCL (stays on GPU); scalar fields travel
        as a small pickled object over the CPU group.
        """
        if not self._is_multi_rank():
            return req

        is_rank0 = self.gpu_id == 0

        if is_rank0:
            assert req is not None, "rank 0 must pass a loaded Req"
            tensor_fields, scalar_fields = extract_transfer_fields(req)
            packed_tensors = _pack_tensor_fields_for_broadcast(tensor_fields)
        else:
            scalar_fields = None
            packed_tensors = None

        # 1. Scalars via CPU pyobj broadcast.
        scalar_fields = self._broadcast_to_all_ranks(scalar_fields)

        # 2. Tensors via NCCL broadcast — keeps GPU buffers on device.
        packed_tensors = self._broadcast_tensor_dict_to_all_ranks(packed_tensors)

        if is_rank0:
            return req

        tensor_fields = _unpack_tensor_fields_from_broadcast(packed_tensors or {})
        # Move tensors onto this rank's physical device. The broadcast
        # allocates receive tensors on the receiver's default CUDA device
        # (set via torch.cuda.set_device(local_rank) during init), which is
        # already the right physical GPU — the .to() is effectively a no-op
        # but makes the invariant explicit for future readers.
        local_device = torch.device(
            f"{current_platform.device_type}:{self.worker.local_rank}"
        )
        for key, value in list(tensor_fields.items()):
            if isinstance(value, torch.Tensor):
                tensor_fields[key] = value.to(local_device, non_blocking=True)
            elif isinstance(value, list):
                tensor_fields[key] = [
                    (
                        t.to(local_device, non_blocking=True)
                        if isinstance(t, torch.Tensor)
                        else t
                    )
                    for t in value
                ]
        return self._build_disagg_req(scalar_fields or {}, tensor_fields)

    # ------------------------------------------------------------------
    # Event loops
    # ------------------------------------------------------------------

    def _disagg_recv_work(self: Scheduler) -> list[bytes] | None:
        """Receive work frames in pool mode, with multi-rank broadcast.

        Rank 0: recv from ZMQ PULL socket, broadcast to other ranks.
        Non-rank-0: receive via NCCL broadcast from rank 0.

        Returns list of bytes frames, or None on shutdown.
        """
        if self.gpu_id == 0:
            raw_frames = self._pool_work_pull.recv_multipart()
            frames = [bytes(f) for f in raw_frames]
        else:
            frames = None

        return self._broadcast_to_all_ranks(frames)

    def _disagg_prefetch_event_loop(self: Scheduler, role_name: str) -> None:
        """Event loop for transfer receiver roles with recv prefetch thread (rank 0).

        The recv thread reads from ZMQ and prefetches tensor loads.
        This loop reads from _compute_ready_queue:
          - "transfer_compute": load already done, wait_event + free slot
              → broadcast scalar_fields to non-rank-0 → compute
          - "transfer_control": alloc/push messages, handle on main thread
              → broadcast "skip" so non-rank-0 doesn't hang
          - queue timeout: broadcast "skip"
          - shutdown: broadcast None
        """
        is_multi_rank = (
            self.server_args.sp_degree != 1
            or self.server_args.tp_size > 1
            or self.server_args.enable_cfg_parallel
        )

        while self._running:
            try:
                try:
                    msg_type, data = self._compute_ready_queue.get(timeout=1.0)
                except queue.Empty:
                    if is_multi_rank:
                        self._broadcast_to_all_ranks(("skip",))
                    continue

                if msg_type == "transfer_compute":
                    # Load already done by recv thread
                    req, load_event, request_id, rn, prealloc_slot_id, scalar_fields = (
                        data
                    )
                    # Wait for load to complete on compute stream
                    if load_event is not None:
                        torch.get_device_module().current_stream().wait_event(
                            load_event
                        )
                    # Now safe to free the receive slot
                    if prealloc_slot_id is not None:
                        with self._transfer_manager._lock:
                            self._transfer_manager._pending_receives.pop(
                                request_id, None
                            )
                    else:
                        self._transfer_manager.free_receive_slot(request_id)
                    # Broadcast the full Req (scalar + tensor fields) to
                    # non-rank-0 ranks. Tensors ride NCCL on the SP/CFG/TP
                    # groups so downstream REPLICATED stages (e.g. denoising)
                    # see identical inputs on every rank — without this, the
                    # non-rank-0 ranks would enter execute_forward with empty
                    # prompt_embeds and fail verify_input.
                    if is_multi_rank:
                        self._broadcast_to_all_ranks(("compute",))
                        self._broadcast_req_to_all_ranks(req)
                    # Init scheduler timesteps on main thread (safe — no
                    # concurrent denoising loop can be running here).
                    if self._disagg_role == RoleType.DENOISER:
                        _init_disagg_request_scheduler(self, req)
                    # Run compute
                    if self._disagg_role == RoleType.DENOISER:
                        self._disagg_denoiser_compute(req, request_id, rn)
                    elif self._disagg_role == RoleType.DECODER:
                        self._disagg_decoder_compute(req, request_id, rn)

                elif msg_type == "transfer_control":
                    # alloc, push messages — handle on main thread (rank 0 only)
                    if is_multi_rank:
                        self._broadcast_to_all_ranks(("skip",))
                    self._handle_transfer_msg(data)

                self._consecutive_error_count = 0

            except Exception as e:
                self._consecutive_error_count += 1
                logger.error(
                    "Pool %s rank %d prefetch loop: error (attempt %d/%d): %s",
                    role_name,
                    self.gpu_id,
                    self._consecutive_error_count,
                    self._max_consecutive_errors,
                    e,
                    exc_info=True,
                )
                if self._consecutive_error_count >= self._max_consecutive_errors:
                    raise RuntimeError(
                        f"Pool {role_name} rank {self.gpu_id} terminated after "
                        f"{self._max_consecutive_errors} consecutive errors: {e}"
                    ) from e

        # Shutdown: notify non-rank-0 to exit
        if is_multi_rank:
            self._broadcast_to_all_ranks(None)
        self._cleanup_disagg()

    def _disagg_non_rank0_event_loop(self: Scheduler) -> None:
        """Event loop for non-rank-0 receivers in multi-rank prefetch mode.

        Blocks on broadcast from rank 0:
          - ("compute", scalar_fields): build minimal Req → execute_forward
          - ("skip",): continue (rank 0 handled a control msg or timed out)
          - None: shutdown, exit loop
        """
        role_name = self._disagg_role.value.upper()
        logger.info(
            "Pool %s rank %d: entering non-rank-0 prefetch loop",
            role_name,
            self.gpu_id,
        )

        while True:
            try:
                msg = self._broadcast_to_all_ranks(None)

                if msg is None:
                    # Shutdown signal
                    break

                if isinstance(msg, tuple) and len(msg) >= 1 and msg[0] == "compute":
                    # Participate in the companion tensor broadcast so this
                    # rank sees the full Req (scalars + GPU tensors). Without
                    # the tensor half, REPLICATED stages would see empty
                    # prompt_embeds on non-rank-0 and fail verify_input.
                    req = self._broadcast_req_to_all_ranks(None)
                    self._disagg_compute_non_rank0(req)
                # else: ("skip",) — continue

            except Exception as e:
                self._consecutive_error_count += 1
                logger.error(
                    "Pool %s rank %d non-rank-0 loop: error (attempt %d/%d): %s",
                    role_name,
                    self.gpu_id,
                    self._consecutive_error_count,
                    self._max_consecutive_errors,
                    e,
                    exc_info=True,
                )
                if self._consecutive_error_count >= self._max_consecutive_errors:
                    raise RuntimeError(
                        f"Pool {role_name} rank {self.gpu_id} terminated after "
                        f"{self._max_consecutive_errors} consecutive errors: {e}"
                    ) from e

        self._cleanup_disagg()

    def _disagg_event_loop(self: Scheduler) -> None:
        """Event loop for all roles in pool mode (DiffusionServer-mediated).

        Multi-rank support:
        - Rank 0 receives from ZMQ, broadcasts to other ranks via NCCL
        - All ranks process work (execute_forward with SP/TP sharding)
        - Only rank 0 sends results back to DiffusionServer

        Transfer:
        - Transfer control messages (transfer_alloc, transfer_push) are rank-0-only.
        - transfer_ready is broadcast to all ranks for compute.
        - Encoder receives pickled Req, runs compute, stages output for transfer.
        - Denoiser/decoder only receive transfer control messages.

        Receiver prefetch paths:
        - Rank 0: _disagg_prefetch_event_loop (reads from compute_ready_queue)
        - Non-rank-0 in multi-rank: _disagg_non_rank0_event_loop (broadcast)
        - Encoder (any rank): existing _disagg_recv_work while loop below
        """
        role_name = self._disagg_role.value.upper()
        is_rank0 = self.gpu_id == 0
        is_multi_rank = (
            self.server_args.sp_degree != 1
            or self.server_args.tp_size > 1
            or self.server_args.enable_cfg_parallel
        )
        use_prefetch = self._compute_ready_queue is not None
        logger.info(
            "Pool mode %s rank %d event loop started " "(multi_rank=%s, prefetch=%s)",
            role_name,
            self.gpu_id,
            is_multi_rank,
            use_prefetch,
        )

        # Rank 0 receiver with prefetch queue → prefetch event loop
        if use_prefetch:
            self._disagg_prefetch_event_loop(role_name)
            return

        # Non-rank-0 receiver in multi-rank → broadcast-based loop
        if (
            not is_rank0
            and is_multi_rank
            and self._disagg_role in (RoleType.DENOISER, RoleType.DECODER)
        ):
            self._disagg_non_rank0_event_loop()
            return

        while self._running:
            try:
                # All ranks receive work (rank 0 via ZMQ, others via broadcast)
                frames = self._disagg_recv_work()

                # Transfer dispatch: check on ALL ranks (frames are broadcast)
                if self._is_transfer_frames(frames):
                    if is_rank0:
                        # Rank 0: handle all transfer messages
                        self._handle_transfer_msg(frames)
                    else:
                        # Non-rank-0: only participate in transfer_ready compute
                        self._handle_transfer_non_rank0(frames)
                elif self._disagg_role == RoleType.ENCODER:
                    self._disagg_encoder_step(
                        send_tensors,
                        frames=frames,
                    )

                self._consecutive_error_count = 0

            except Exception as e:
                self._consecutive_error_count += 1
                logger.error(
                    "Pool %s rank %d: error (attempt %d/%d): %s",
                    role_name,
                    self.gpu_id,
                    self._consecutive_error_count,
                    self._max_consecutive_errors,
                    e,
                    exc_info=True,
                )
                if self._consecutive_error_count >= self._max_consecutive_errors:
                    raise RuntimeError(
                        f"Pool {role_name} rank {self.gpu_id} terminated after "
                        f"{self._max_consecutive_errors} consecutive errors: {e}"
                    ) from e

        self._cleanup_disagg()

    def _cleanup_disagg(self: Scheduler):
        """Clean up all pool mode resources (sockets, threads, transfer manager)."""
        # Shutdown RDMA push thread
        if self._rdma_push_queue is not None:
            self._rdma_push_queue.put(None)
        if self._rdma_push_thread is not None:
            self._rdma_push_thread.join(timeout=5)
        if self._rdma_push_zmq is not None:
            self._rdma_push_zmq.close()
        # Recv prefetch thread stops when self._running = False
        if self._recv_prefetch_thread is not None:
            self._recv_prefetch_thread.join(timeout=5)
        if self._transfer_manager is not None:
            self._transfer_manager.cleanup()
        if self._pool_work_pull is not None:
            self._pool_work_pull.close()
        if self._pool_result_push is not None:
            self._pool_result_push.close()

    # ------------------------------------------------------------------
    # Transfer message handling
    # ------------------------------------------------------------------

    @staticmethod
    def _is_transfer_frames(frames: list) -> bool:
        """Check if ZMQ multipart frames carry a transfer control message."""
        return is_transfer_message(frames)

    def _handle_transfer_msg(self: Scheduler, frames: list) -> None:
        """Dispatch a transfer control message to the appropriate handler (rank 0)."""
        msg = decode_transfer_msg(frames)
        msg_type = msg.get("msg_type", "")
        request_id = msg.get("request_id", "")

        logger.debug(
            "Transfer %s: received %s for %s",
            self._disagg_role.value.upper(),
            msg_type,
            request_id,
        )

        if msg_type == TransferMsgType.ALLOC:
            self._handle_transfer_alloc(msg)
        elif msg_type == TransferMsgType.PUSH:
            self._handle_transfer_push(msg)
        elif msg_type == TransferMsgType.READY:
            self._handle_transfer_ready(msg)
        else:
            logger.warning(
                "Transfer %s: unknown message type %s",
                self._disagg_role.value.upper(),
                msg_type,
            )

    def _handle_transfer_non_rank0(self: Scheduler, frames: list) -> None:
        """Handle transfer messages on non-rank-0 workers.

        Only transfer_ready requires non-rank-0 participation (for compute).
        transfer_alloc and transfer_push are rank-0-only operations — skip them.
        """
        msg = decode_transfer_msg(frames)
        msg_type = msg.get("msg_type", "")

        if msg_type == TransferMsgType.READY:
            # Non-rank-0 has no TransferManager, so rank 0 loads tensors from
            # the RDMA buffer and broadcasts the full Req (scalars + tensors)
            # over NCCL. Participate in the matching broadcast here.
            req = self._broadcast_req_to_all_ranks(None)
            self._disagg_compute_non_rank0(req)
        # else: transfer_alloc, transfer_push — skip (rank-0-only operations)

    def _handle_transfer_alloc(self: Scheduler, msg: dict) -> None:
        """Handle transfer_alloc: allocate a receive slot and reply with transfer_allocated."""
        request_id = msg["request_id"]
        data_size = msg.get("data_size", 0)

        pending = self._transfer_manager.allocate_receive_slot(request_id, data_size)
        if pending is None:
            logger.error(
                "Transfer %s: failed to allocate receive slot for %s (%d bytes)",
                self._disagg_role.value.upper(),
                request_id,
                data_size,
            )
            return

        allocated_msg = TransferAllocatedMsg(
            request_id=request_id,
            session_id=self._transfer_manager.session_id,
            pool_ptr=self._transfer_manager.pool_data_ptr,
            slot_offset=pending.slot.offset,
            slot_size=pending.slot.size,
        )
        self._pool_result_push.send_multipart(encode_transfer_msg(allocated_msg))

        logger.debug(
            "Transfer %s: allocated receive slot for %s (offset=%d, size=%d)",
            self._disagg_role.value.upper(),
            request_id,
            pending.slot.offset,
            pending.slot.size,
        )

    def _handle_transfer_push(self: Scheduler, msg: dict) -> None:
        """Handle transfer_push: RDMA push staged data to peer, reply with transfer_pushed.

        If RDMA push thread is active, enqueue non-blocking.
        Otherwise fall back to blocking push (e.g., during shutdown).
        """
        request_id = msg["request_id"]
        dest_session_id = msg.get("dest_session_id", "")
        dest_addr = msg.get("dest_addr", 0)
        transfer_size = msg.get("transfer_size", 0)

        if self._rdma_push_queue is not None:
            # Non-blocking: enqueue to RDMA push thread
            self._rdma_push_queue.put(
                (
                    request_id,
                    dest_session_id,
                    dest_addr,
                    transfer_size,
                )
            )
            return

        # Fallback: blocking push on main thread
        success = self._transfer_manager.push_to_peer(
            request_id=request_id,
            dest_session_id=dest_session_id,
            dest_addr=dest_addr,
            transfer_size=transfer_size,
        )

        if success:
            self._transfer_manager.free_staged(request_id)

        pushed_msg = TransferPushedMsg(request_id=request_id)
        self._pool_result_push.send_multipart(encode_transfer_msg(pushed_msg))

        if not success:
            logger.error(
                "Transfer %s: RDMA push failed for %s",
                self._disagg_role.value.upper(),
                request_id,
            )

    def _handle_transfer_ready(self: Scheduler, msg: dict) -> None:
        """Handle transfer_ready: load tensors from buffer, run compute, send result.

        Overlap tensor load with Req construction and scheduler init.
        After the RDMA data arrives:
        1. Start load on transfer_stream (non-blocking)
        2. Build Req from scalar fields + tensors (CPU, overlapped)
        3. Init scheduler timesteps if denoiser (CPU, overlapped)
        4. Wait for load before compute
        5. Run the role's compute
        """

        request_id = msg["request_id"]
        manifest = msg.get("manifest", {})
        scalar_fields = msg.get("scalar_fields", {})
        role_name = self._disagg_role.value.upper()

        if self._disagg_metrics:
            self._disagg_metrics.record_request_start(request_id)

        # If using a pre-allocated slot, register it as pending receive
        prealloc_slot_id = scalar_fields.pop("_prealloc_slot_id", None)
        if (
            prealloc_slot_id is not None
            and prealloc_slot_id in self._preallocated_slots
        ):
            slot = self._preallocated_slots[prealloc_slot_id]
            self._transfer_manager.register_prealloc_as_receive(request_id, slot)

        # 1. Start load on transfer_stream (non-blocking)
        local_device = f"{current_platform.device_type}:{self.worker.local_rank}"
        tensors, load_event = self._transfer_manager.load_tensors_async(
            request_id,
            manifest,
            device=local_device,
            stream=self._transfer_stream,
        )

        # 2. Build Req from scalar fields + tensors (CPU work, overlapped)
        req = self._build_disagg_req(scalar_fields, tensors)

        # 3. Init scheduler timesteps if denoiser (CPU work, overlapped)
        if self._disagg_role == RoleType.DENOISER:
            _init_disagg_request_scheduler(self, req)

        # 4. Wait for load before compute (GPU must see the data)
        if load_event is not None:
            torch.get_device_module().current_stream().wait_event(load_event)

        # 5. Free receive slot after load completes (data is on compute GPU)
        if prealloc_slot_id is not None:
            # Pre-allocated slot: just remove from pending receives, don't free buffer
            with self._transfer_manager._lock:
                self._transfer_manager._pending_receives.pop(request_id, None)
        else:
            self._transfer_manager.free_receive_slot(request_id)

        # 6. In multi-rank mode, broadcast the fully-loaded Req to the other
        # ranks so REPLICATED stages see identical inputs everywhere. See
        # the prefetch-loop variant for the matching receiver broadcast.
        if self._is_multi_rank():
            self._broadcast_req_to_all_ranks(req)

        # 7. Run compute
        if self._disagg_role == RoleType.DENOISER:
            self._disagg_denoiser_compute(req, request_id, role_name)
        elif self._disagg_role == RoleType.DECODER:
            self._disagg_decoder_compute(req, request_id, role_name)

    # ------------------------------------------------------------------
    # Compute
    # ------------------------------------------------------------------

    def _disagg_compute_non_rank0(self: Scheduler, req: Req) -> None:
        """Non-rank-0 compute: enter execute_forward with a Req received via
        NCCL broadcast from rank 0.

        The Req already contains tensor fields materialized on this rank's
        GPU (see ``_broadcast_req_to_all_ranks``), so REPLICATED stages such
        as denoising have non-empty prompt_embeds and verify_input passes.

        Used by both the non-prefetch path (:meth:`_handle_transfer_non_rank0`)
        and the prefetch non-rank-0 loop
        (:meth:`_disagg_non_rank0_event_loop`).
        """
        if self._disagg_role == RoleType.DENOISER:
            # Initialize scheduler timesteps (same as rank 0)
            _init_disagg_request_scheduler(self, req)

            with self._disagg_trace_dispatch(req):
                self.worker.execute_forward([req], return_req=True)

        elif self._disagg_role == RoleType.DECODER:
            req.save_output = False
            req.return_file_paths_only = False
            with self._disagg_trace_dispatch(req):
                self.worker.execute_forward([req])

    def _build_disagg_req(self: Scheduler, scalar_fields: dict, tensors: dict) -> Req:
        """Reconstruct a Req from transfer scalar fields and loaded GPU tensors.

        Initializes all dataclass field defaults first, then overlays
        scalar and tensor fields from the transfer message.
        """
        # Pop _trace_state before the generic setattr loop so it doesn't land
        # on the Req as a stray attribute.
        trace_state = scalar_fields.pop("_trace_state", None)

        req = object.__new__(Req)
        # Initialize all dataclass fields with their defaults
        for f in dataclasses.fields(Req):
            if f.default is not dataclasses.MISSING:
                object.__setattr__(req, f.name, f.default)
            elif f.default_factory is not dataclasses.MISSING:
                object.__setattr__(req, f.name, f.default_factory())
        # Ensure sampling_params is not None so __getattr__ delegation works
        object.__setattr__(req, "sampling_params", SamplingParams())
        # Restore _extra_* prefixed fields into req.extra dict
        extra_keys = [k for k in scalar_fields if k.startswith("_extra_")]
        for key in extra_keys:
            req.extra[key[len("_extra_") :]] = scalar_fields.pop(key)
        for key, value in scalar_fields.items():
            setattr(req, key, value)
        # Set tensor fields
        for key, value in tensors.items():
            setattr(req, key, value)
        # Recreate torch.Generator from seed (not serializable over transfer)
        seed = scalar_fields.get("seed")
        if seed is not None:
            if isinstance(seed, list):
                req.generator = [
                    torch.Generator(device="cpu").manual_seed(int(item))
                    for item in seed
                ]
            else:
                req.generator = torch.Generator(device="cpu").manual_seed(int(seed))
        # Rebuild trace_ctx from the propagated __getstate__ dict so this role's
        # spans nest under the sender's trace (same mechanism SRT uses via pickle).
        if trace_state and trace_state.get("tracing_enable"):
            try:
                ctx = object.__new__(TraceReqContext)
                ctx.__setstate__(trace_state)
                req.trace_ctx = ctx
            except Exception as e:
                logger.debug("Failed to rebuild trace_ctx from _trace_state: %s", e)
        req.validate()
        return req

    @contextlib.contextmanager
    def _disagg_trace_dispatch(self: Scheduler, req: Req):
        """Wrap a disagg role's worker.execute_forward in the tracing lifecycle.

        Mirrors the monolithic path in ``scheduler._handle_generation``: rebuild
        the thread context under the (potentially remote) root_span_context that
        was propagated in via ``_trace_state`` / pickle, then emit a
        ``scheduler_dispatch`` span for this role with ``thread_finish_flag``
        so the thread span closes when compute returns. If tracing is disabled
        (TraceNullContext), everything is a no-op.
        """
        ctx = getattr(req, "trace_ctx", None)
        if ctx is None:
            yield
            return
        # Disagg receive (__setstate__) and compute may run on different
        # threads (e.g. recv-prefetch vs scheduler main). Align the ctx's pid
        # with the current compute thread so __create_thread_context's
        # threads_info lookup resolves via the local registration.
        if getattr(ctx, "tracing_enable", False):
            ctx.pid = threading.get_native_id()
        ctx.rebuild_thread_context()
        with trace_slice(ctx, DiffStage.SCHEDULER_DISPATCH, thread_finish_flag=True):
            yield

    def _disagg_denoiser_compute(
        self: Scheduler, req: Req, request_id: str, role_name: str
    ) -> None:
        """Run denoiser compute in transfer mode, then stage output for decoder.

        Note: Scheduler timestep init is done in _handle_transfer_ready
        to overlap with tensor loading.
        """
        # Run denoising
        start_time = time.monotonic()
        with self._disagg_trace_dispatch(req):
            result = self.worker.execute_forward([req], return_req=True)
        duration_s = time.monotonic() - start_time

        if not isinstance(result, Req):
            error_msg = getattr(result, "error", "denoiser error")
            done_msg = TransferDoneMsg(request_id=request_id, error=str(error_msg))
            self._pool_result_push.send_multipart(encode_transfer_msg(done_msg))
            if self._disagg_metrics:
                self._disagg_metrics.record_request_failed(request_id)
            return

        # Stage denoiser output for decoder transfer (async staging)
        tensor_fields, scalar_fields = extract_transfer_fields(result)

        # 1. Stage tensors on transfer_stream (non-blocking)
        staged, stage_event = self._transfer_manager.stage_tensors_async(
            request_id=request_id,
            tensor_fields=tensor_fields,
            scalar_fields=scalar_fields,
            stream=self._transfer_stream,
        )

        if staged is None:
            done_msg = TransferDoneMsg(
                request_id=request_id,
                error="Failed to stage denoiser output for decoder",
            )
            self._pool_result_push.send_multipart(encode_transfer_msg(done_msg))
            if self._disagg_metrics:
                self._disagg_metrics.record_request_failed(request_id)
            return

        # 2. Build done_data dict while staging runs (CPU work, overlapped)
        done_data = {
            "msg_type": "transfer_done",
            "request_id": request_id,
            "staged_for_decoder": True,
            "session_id": self._transfer_manager.session_id,
            "pool_ptr": self._transfer_manager.pool_data_ptr,
            "slot_offset": staged.slot.offset if staged.slot else 0,
            "data_size": staged.slot.size if staged.slot else 0,
            "manifest": staged.manifest,
            "scalar_fields": staged.scalar_fields,
        }
        msg_bytes = json.dumps(done_data, separators=(",", ":")).encode("utf-8")

        # 3. Wait for staging to complete before sending
        if stage_event is not None:
            stage_event.synchronize()

        # 4. Send transfer_done with staged info
        self._pool_result_push.send_multipart([TRANSFER_MAGIC, msg_bytes])

        if self._disagg_metrics:
            self._disagg_metrics.record_request_complete(request_id)

        logger.debug(
            "Transfer DENOISER: processed %s in %.2f s, staged for decoder",
            request_id,
            duration_s,
        )

    def _disagg_decoder_compute(
        self: Scheduler, req: Req, request_id: str, role_name: str
    ) -> None:
        """Run decoder compute in transfer mode, send result to DS.

        Decoder result is sent as raw ZMQ multipart frames (same format as
        relay mode) so DiffusionServer handles it via _handle_decoder_result_frames
        without hex/JSON overhead.
        """

        # Check for upstream error
        disagg_error = getattr(req, "_disagg_error", None)
        if disagg_error:
            if self._pool_result_push is not None:
                send_tensors(
                    self._pool_result_push,
                    {},
                    {
                        "request_id": request_id,
                        "error": f"Upstream error: {disagg_error}",
                    },
                )
            return

        req.save_output = False
        req.return_file_paths_only = False

        start_time = time.monotonic()
        with self._disagg_trace_dispatch(req):
            output_batch = self.worker.execute_forward([req])
        duration_s = time.monotonic() - start_time

        # Send result as raw ZMQ frames (no TRANSFER_MAGIC prefix).
        # DiffusionServer will route it through _handle_decoder_result_frames,
        # the same path as relay mode.
        tensor_fields = {}
        scalar_fields = {"request_id": request_id}
        if output_batch.output is not None:
            tensor_fields["output"] = output_batch.output
        if output_batch.audio is not None:
            tensor_fields["audio"] = output_batch.audio
        if output_batch.audio_sample_rate is not None:
            scalar_fields["audio_sample_rate"] = output_batch.audio_sample_rate
        if output_batch.error is not None:
            scalar_fields["error"] = output_batch.error

        if self._pool_result_push is not None:
            send_tensors(self._pool_result_push, tensor_fields, scalar_fields)

        if self._disagg_metrics:
            if output_batch.error:
                self._disagg_metrics.record_request_failed(request_id)
            else:
                self._disagg_metrics.record_request_complete(request_id)

        logger.debug("Transfer DECODER: processed %s in %.2f s", request_id, duration_s)

    def _disagg_encoder_step(
        self: Scheduler,
        send_tensors_fn,
        frames=None,
    ):
        """Single encoder step in pool mode."""
        # Receive: [request_id_bytes, pickled_req_bytes]
        if frames is None:
            frames = self._pool_work_pull.recv_multipart()
        pickled_req = frames[-1]
        reqs = pickle.loads(pickled_req)
        if not isinstance(reqs, list):
            reqs = [reqs]

        req = reqs[0]
        request_id = getattr(req, "request_id", "unknown")

        if self._disagg_metrics:
            self._disagg_metrics.record_request_start(request_id)

        # Run encoder stages
        with self._disagg_trace_dispatch(req):
            req_result = self.worker.execute_forward(reqs, return_req=True)

        if not isinstance(req_result, Req):
            # Error — send error via scalar fields (rank 0 only)
            if self._pool_result_push is not None:
                error_msg = getattr(req_result, "error", "encoder error")
                send_tensors_fn(
                    self._pool_result_push,
                    {},
                    {"request_id": request_id, "_disagg_error": str(error_msg)},
                )
            if self._disagg_metrics:
                self._disagg_metrics.record_request_failed(request_id)
            return

        # Pack and send encoder output (rank 0 only sends)
        tensor_fields, scalar_fields = extract_transfer_fields(req_result)

        if self._pool_result_push is not None:
            if self._transfer_manager is not None:
                # Transfer mode: stage tensors to TransferBuffer, send transfer_staged
                self._disagg_encoder_transfer_stage(
                    request_id, tensor_fields, scalar_fields
                )
            else:
                # Fallback: send error (transfer manager not initialized)
                send_tensors_fn(
                    self._pool_result_push,
                    {},
                    {"request_id": request_id, "_disagg_error": "No transfer manager"},
                )

        if self._disagg_metrics:
            self._disagg_metrics.record_request_complete(request_id)

        logger.debug("Pool ENCODER: processed %s", request_id)

    def _disagg_encoder_transfer_stage(
        self: Scheduler, request_id: str, tensor_fields: dict, scalar_fields: dict
    ) -> None:
        """Stage encoder output and send transfer_staged to DS.

        Overlap staging with metadata JSON serialization.
        """
        # 1. Stage tensors on transfer_stream (non-blocking)
        staged, stage_event = self._transfer_manager.stage_tensors_async(
            request_id=request_id,
            tensor_fields=tensor_fields,
            scalar_fields=scalar_fields,
            stream=self._transfer_stream,
        )

        if staged is None:
            # Staging failed — send error via relay as fallback
            send_tensors(
                self._pool_result_push,
                {},
                {"request_id": request_id, "_disagg_error": "Transfer staging failed"},
            )
            if self._disagg_metrics:
                self._disagg_metrics.record_request_failed(request_id)
            return

        # 2. Build transfer metadata dict while staging runs (CPU work, overlapped)
        staged_data = {
            "msg_type": "transfer_staged",
            "request_id": request_id,
            "data_size": staged.slot.size if staged.slot else 0,
            "manifest": staged.manifest,
            "session_id": self._transfer_manager.session_id,
            "pool_ptr": self._transfer_manager.pool_data_ptr,
            "slot_offset": staged.slot.offset if staged.slot else 0,
            "scalar_fields": staged.scalar_fields,
        }
        msg_bytes = json.dumps(staged_data, separators=(",", ":")).encode("utf-8")

        # 3. Wait for staging to complete before sending (buffer must be ready)
        if stage_event is not None:
            stage_event.synchronize()

        # 4. Send transfer staged message
        self._pool_result_push.send_multipart([TRANSFER_MAGIC, msg_bytes])