lightseekorg--tokenspeed/python/tokenspeed/runtime/engine/event_loop.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.

import faulthandler
import signal
import time
from collections import OrderedDict
from dataclasses import dataclass

import psutil
import setproctitle
import torch
import torch.distributed as dist
import zmq
from tokenspeed_scheduler import PD, Cache, ExecutionEvent, ForwardEvent, Scheduler

from tokenspeed.runtime.cache.executor.flat_memory_executor import (
    FlatMemoryExecutor,
)
from tokenspeed.runtime.cache.executor.memory_executor import (
    MemoryExecutor,
    MemoryExecutorConfig,
)
from tokenspeed.runtime.cache.transfer.types import CacheKind
from tokenspeed.runtime.configs.model_config import ModelConfig
from tokenspeed.runtime.configs.paged_cache_spec import (
    scheduler_ext_flat_kvcache,
    validate_flat_scheduler_config,
)
from tokenspeed.runtime.distributed.process_group_manager import (
    process_group_manager as pg_manager,
)
from tokenspeed.runtime.engine.generation_output_processor import OutputProcesser
from tokenspeed.runtime.engine.memory_occupation import MemoryOccupationController
from tokenspeed.runtime.engine.pause import PauseController
from tokenspeed.runtime.engine.request_handler import RequestHandler
from tokenspeed.runtime.engine.scheduler_utils import (
    advance_forward,
    cache_event_from_payload,
    cache_event_key,
    cache_event_to_payload,
    cache_sync_debug_enabled,
    make_config,
    pool_to_paged_cache_groups,
    pool_to_prefix_cache_adjunct_spec,
    pop_common_cache_event_payloads,
    should_use_overlap_schedule,
)
from tokenspeed.runtime.execution.distributed_initializer import (
    DistributedConfig,
    DistributedInitializer,
)
from tokenspeed.runtime.execution.factory import (
    ModelExecutorConfig,
    create_model_executor,
    create_model_runner,
)
from tokenspeed.runtime.execution.forward_batch_info import ForwardMode
from tokenspeed.runtime.execution.types import ModelExecutionResult
from tokenspeed.runtime.grammar.capturable_grammar import GrammarStepInputs
from tokenspeed.runtime.layers.attention.registry import create_attn_components
from tokenspeed.runtime.metrics.collector import EngineMetrics
from tokenspeed.runtime.pd.decode_executor import DisaggDecodeExecutor
from tokenspeed.runtime.pd.factory import (
    create_kv_transfer,
    get_kv_args,
)
from tokenspeed.runtime.pd.kv_events import (
    EventPublisherFactory,
    KVEventBatch,
    NullEventPublisher,
    drain_scheduler_kv_events,
    scheduler_kv_events_to_wire_events,
)
from tokenspeed.runtime.pd.mooncake.entities import KVManagerArgs
from tokenspeed.runtime.pd.prefill_executor import DisaggPrefillExecutor
from tokenspeed.runtime.sampling.sampling_params import SamplingParams
from tokenspeed.runtime.utils import (
    configure_logger,
    get_colorful_logger,
    get_zmq_socket,
)
from tokenspeed.runtime.utils.exceptions import get_exception_traceback
from tokenspeed.runtime.utils.nvtx import nvtx_range
from tokenspeed.runtime.utils.process import register_usr_signal
from tokenspeed.runtime.utils.server_args import PortArgs, ServerArgs
from tokenspeed.runtime.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter

logger = get_colorful_logger(__name__)


def calc_l3_query_hashes(scheduler, tokens: list[int]) -> list[str]:
    return scheduler.calc_rolling_hash(tokens, apply_match=True)


# Sleep between iterations while frozen (PAUSED_ALL) so the keep-mode pause does
# not busy-spin a CPU core waiting for /resume.
_PAUSED_IDLE_SLEEP_S = 0.001


def _forward_op_executes_model_forward(forward_op, *, is_disagg_decode: bool) -> bool:
    """Return whether ``forward_op`` will enter the model forward path.

    On decode-side PD, EXTEND ops only start remote KV receive; the model
    forward runs after the remote prefill completes and the scheduler advances
    the request into decode. Treating those EXTEND ops as model work makes
    idle DP ranks enter dummy collectives that the active rank will not match.
    """
    if forward_op is None:
        return False
    if sum(forward_op.input_lengths) <= 0:
        return False
    if is_disagg_decode and forward_op.num_extends() > 0:
        return False
    return True


class _NullSender:
    """No-op ZMQ sender for non-rank-0 workers."""

    @staticmethod
    def send_pyobj(x):
        return None


@dataclass(frozen=True)
class DpForwardMetadata:
    global_num_tokens: list[int]
    global_batch_size: list[int]
    global_forward_mode: list[int]
    all_decode_or_idle: bool
    all_extend: bool
    need_idle_forward: bool


class EventLoop:
    def __init__(
        self,
        server_args: ServerArgs,
        port_args: PortArgs,
        gpu_id: int,
        attn_tp_rank: int,
        dp_rank: int,
        global_rank: int,
    ) -> None:
        # Do not pass server_args further down the stack after this point.

        self.server_args = server_args
        self.port_args = port_args
        self.gpu_id = gpu_id
        self.global_rank = global_rank

        self.model_config = self._load_model_config(server_args.model)
        if server_args.speculative_draft_model_path is not None:
            draft_model_config = self._load_model_config(
                server_args.speculative_draft_model_path,
                is_draft_worker=True,
            )
        else:
            draft_model_config = None

        min_per_gpu_mem = self._init_distributed()

        target, draft = create_model_runner(
            server_args, self.model_config, draft_model_config, gpu_id, global_rank
        )
        self.use_overlap_schedule = should_use_overlap_schedule(
            disable_overlap_schedule=server_args.disable_overlap_schedule,
            disaggregation_mode=server_args.disaggregation_mode,
        )
        self.overlap_schedule_depth = int(self.use_overlap_schedule)
        decode_input_tokens = (
            server_args.speculative_num_draft_tokens
            if server_args.speculative_algorithm is not None
            else 1
        )

        (
            attn_backend,
            token_to_kv_pool,
            draft_attn_backend,
            draft_token_to_kv_pool,
            self.max_total_num_tokens,
            mamba_pool_total_chunks,
            mamba_pool,
        ) = create_attn_components(
            server_args,
            self.model_config,
            gpu_id,
            global_rank,
            min_per_gpu_mem,
            server_args.enable_memory_saver,
            draft_model_config,
            decode_input_tokens=decode_input_tokens,
            overlap_schedule_depth=self.overlap_schedule_depth,
        )

        num_total_pages = self.max_total_num_tokens // server_args.block_size
        hf_config = getattr(self.model_config, "hf_config", None)
        text_config = getattr(hf_config, "text_config", None) if hf_config else None
        has_mamba = getattr(self.model_config, "mambaish_config", None) is not None or (
            text_config is not None and hasattr(text_config, "mamba2_cache_params")
        )

        mapping = server_args.mapping
        # The C++ scheduler's req_pool_idx range is rank-local and 1-based:
        # real rows are 1..max_batch_size, row 0 is reserved, and CUDA graph
        # padding needs one non-real sink row after the scheduler-owned range.
        per_rank_max_batch = server_args.max_num_seqs // max(mapping.attn.dp_size, 1)
        req_pool_padding_index = per_rank_max_batch + 1

        model_executor_config = ModelExecutorConfig.from_server_args(
            server_args=server_args,
            model_config=self.model_config,
            max_req_pool_size=req_pool_padding_index,
            gpu_id=gpu_id,
            global_rank=global_rank,
            num_total_pages=num_total_pages,
            overlap_schedule_depth=self.overlap_schedule_depth,
        )
        self.model_executor = create_model_executor(
            server_args=server_args,
            config=model_executor_config,
            model_runner=target,
            draft_model_runner=draft,
            attn_backend=attn_backend,
            token_to_kv_pool=token_to_kv_pool,
            draft_attn_backend=draft_attn_backend,
            draft_token_to_kv_pool=draft_token_to_kv_pool,
            mamba_pool=mamba_pool,
        )

        # Reserve one token slot because request validation uses a strict
        # ``< max_req_len`` check against the model context length.
        self.max_req_input_len = self.model_config.context_len - 1
        self.attn_tp_size = server_args.attn_tp_size or mapping.attn.tp_size
        self.world_size = server_args.world_size or mapping.world_size
        self.attn_tp_rank = attn_tp_rank
        self.attn_tp_cpu_group = pg_manager.get_process_group(
            "gloo", server_args.mapping.attn.tp_group
        )
        self._pending_cache_event_payloads: OrderedDict[tuple[str, int], dict] = (
            OrderedDict()
        )
        # All ranks submit identical cache plans (the C++ scheduler is mirrored),
        # so a local in-flight counter mirrors across ranks: if it's 0 here, no
        # rank has anything pending. Lets us skip the TP collective in
        # _commit_cache_results entirely when nothing is in flight.
        self._num_inflight_cache_ops = 0
        self.dp_rank = dp_rank
        self.dp_size = mapping.attn.dp_size
        self.has_dp = mapping.has_attn_dp
        if self.has_dp:
            self.world_cpu_group = pg_manager.get_process_group(
                "gloo", mapping.world_group
            )
            self._dp_local_info = torch.zeros(1, 3, dtype=torch.int32)
            self._dp_global_info = torch.zeros(mapping.world_size, 3, dtype=torch.int32)
        if not server_args.enable_kvstore:
            logger.warning(
                "KVStore L2 cache will not be used during normal execution, but it will still be used when retraction happens."
            )

        mamba_l2_host_slots = 0
        if has_mamba and server_args.enable_mamba_l2:
            if server_args.mamba_l2_host_slots > 0:
                mamba_l2_host_slots = server_args.mamba_l2_host_slots
            elif server_args.mamba_l2_host_gb > 0 and mamba_pool is not None:
                slot_bytes = int(
                    mamba_pool.conv_state.shape[0]
                    * (
                        mamba_pool.conv_state[0, 0].nbytes
                        + mamba_pool.ssm_state[0, 0].nbytes
                    )
                )
                mamba_l2_host_slots = int(
                    server_args.mamba_l2_host_gb * (1024**3) // max(slot_bytes, 1)
                )
            else:
                mamba_l2_host_slots = max(
                    int(mamba_pool_total_chunks * server_args.mamba_l2_ratio), 1
                )

        mem_cfg = MemoryExecutorConfig(
            layer_num=self.model_config.num_hidden_layers,
            page_size=server_args.block_size,
            host_ratio=server_args.kvstore_ratio,
            host_size_gb=server_args.kvstore_size,
            host_parallel_count=max(
                int(getattr(server_args.mapping, "nprocs_per_node", 1) or 1), 1
            ),
            io_backend=server_args.kvstore_io_backend,
            host_layout=server_args.kvstore_mem_layout,
            storage_backend=server_args.kvstore_storage_backend,
            storage_backend_extra_config=server_args.kvstore_storage_backend_extra_config,
            model_name=server_args.model,
            enable_mamba_l2=server_args.enable_mamba_l2,
            mamba_l2_host_slots=mamba_l2_host_slots,
            mamba_l2_layout=server_args.mamba_l2_layout,
            mamba_l2_io_backend=server_args.mamba_l2_io_backend,
        )
        if scheduler_ext_flat_kvcache() and server_args.enable_kvstore:
            if server_args.kvstore_storage_backend is not None:
                raise NotImplementedError(
                    "flat scheduler build (TOKENSPEED_FLAT_KVCACHE) has no L3 "
                    "storage tier yet; unset --kvstore-storage-backend."
                )
            self.memory_executor = FlatMemoryExecutor(
                device_pool=token_to_kv_pool,
                host_ratio=server_args.kvstore_ratio,
                host_size_gb=server_args.kvstore_size,
            )
            num_host_pages = self.memory_executor.num_host_pages
        elif not token_to_kv_pool.supports_hierarchical_kv_cache:
            if server_args.enable_kvstore:
                raise NotImplementedError(
                    "This KV cache pool does not support hierarchical cache "
                    "(kvstore); pass --disable-kvstore."
                )
            self.memory_executor = None
            num_host_pages = 0
        else:
            self.memory_executor = MemoryExecutor(
                device_pool=token_to_kv_pool,
                config=mem_cfg,
                is_dp_attention_enabled=self.has_dp,
                tp_group=self.attn_tp_cpu_group,
                draft_device_pool=draft_token_to_kv_pool,
                mamba_pool=mamba_pool,
            )
            num_host_pages = self.memory_executor.host_pool.page_num

        # Flat host tier acks loadbacks (LoadBackDoneEvent), so they join the
        # inflight accounting in _submit_cache_ops; radix loadbacks never ack.
        self._loadback_acks_expected = getattr(
            self.memory_executor, "emits_loadback_acks", False
        )

        self._kv_events_enabled = (
            EventPublisherFactory.is_enabled(server_args.kv_events_config)
            and attn_tp_rank == 0
        )

        if has_mamba and server_args.max_mamba_cache_size is None:
            logger.info(
                f"Mamba radix cache enabled without explicit max_mamba_cache_size. "
                f"Auto-derived mamba_pool_total_chunks={mamba_pool_total_chunks} "
                f"(ratio={server_args.mamba_full_memory_ratio})."
            )

        # Adjunct enabled only when pool opts in AND prefix-caching switch is on.
        paged_cache_groups = pool_to_paged_cache_groups(token_to_kv_pool)
        validate_flat_scheduler_config(
            flat_kvcache_ext=scheduler_ext_flat_kvcache(),
            paged_cache_groups=paged_cache_groups,
            attn_backend=attn_backend,
            kv_pool=token_to_kv_pool,
            speculative_algorithm=server_args.speculative_algorithm,
        )
        self._paged_cache_groups = paged_cache_groups
        prefix_cache_adjunct = None
        required_groups = token_to_kv_pool.prefix_cache_required_group_ids
        if required_groups is not None and server_args.enable_prefix_caching:
            prefix_cache_adjunct = pool_to_prefix_cache_adjunct_spec(required_groups)
        scheduler_cfg = make_config(
            num_device_pages=self.max_total_num_tokens // server_args.block_size,
            max_scheduled_tokens=server_args.chunked_prefill_size,
            max_batch_size=per_rank_max_batch,
            page_size=server_args.block_size,
            num_host_pages=num_host_pages,
            disable_l2_cache=not server_args.enable_kvstore,
            enable_l3_storage=server_args.kvstore_storage_backend is not None,
            prefetch_threshold=4,  # Keep this hard-coded until it becomes configurable.
            role=server_args.disaggregation_mode,
            enable_kv_cache_events=self._kv_events_enabled,
            decode_input_tokens=decode_input_tokens,
            overlap_schedule_depth=self.overlap_schedule_depth,
            disable_prefix_cache=not server_args.enable_prefix_caching,
            enable_mamba=has_mamba,
            mamba_cache_chunk_size=server_args.mamba_cache_chunk_size,
            mamba_pool_total_chunks=mamba_pool_total_chunks,
            enable_mamba_l2=server_args.enable_mamba_l2,
            mamba_l2_host_slots=mamba_l2_host_slots,
            paged_cache_groups=paged_cache_groups,
            enable_mixed_prefill_decode=server_args.enable_mixed_batch,
            prefix_cache_adjunct=prefix_cache_adjunct,
        )
        logger.info(
            "Scheduler config: block_size=%s num_device_pages=%s "
            "max_scheduled_tokens=%s decode_input_tokens=%s "
            "overlap_schedule_depth=%s disable_l2_cache=%s "
            "max_batch_size=%s (global max_num_seqs=%s, dp_size=%s) "
            "mamba_pool_total_chunks=%s enable_mamba=%s "
            "disable_prefix_cache=%s paged_cache_groups=%s",
            scheduler_cfg.block_size,
            scheduler_cfg.num_device_pages,
            scheduler_cfg.max_scheduled_tokens,
            scheduler_cfg.decode_input_tokens,
            scheduler_cfg.overlap_schedule_depth,
            scheduler_cfg.disable_l2_cache,
            scheduler_cfg.max_batch_size,
            server_args.max_num_seqs,
            self.dp_size,
            mamba_pool_total_chunks,
            has_mamba,
            scheduler_cfg.disable_prefix_cache,
            [group.group_id for group in paged_cache_groups],
        )
        self.scheduler = Scheduler(scheduler_cfg)
        token_to_kv_pool.bind_paged_cache_scheduler(self.scheduler)
        if attn_tp_rank == 0:
            self.kv_event_publisher = EventPublisherFactory.create(
                server_args.kv_events_config,
                attn_dp_rank=dp_rank,
            )
        else:
            self.kv_event_publisher = NullEventPublisher(attn_dp_rank=dp_rank)

        self._init_interprocess_comm()

        # Pause/resume control state. Shared with the request handler, which
        # drives the control-request side; the event loop reads the gate.
        self._pause = PauseController(self.send_to_tokenizer)

        # GPU-memory data plane (release/resume_memory_occupation). Reuses the
        # pause controller's drain machinery; frees memory via the memory-saver
        # adapter once the scheduler drains. See memory_occupation.py.
        # Releasing KV is only safe if any prefix cache it backs can be cleared:
        # either prefix caching is off, or the scheduler exposes a reset. Decide
        # once here (static config) and let the controller reject unsafe releases.
        kv_cache_release_allowed = (
            not self.server_args.enable_prefix_caching
            or callable(getattr(self.scheduler, "reset_prefix_cache", None))
        )
        self._memory = MemoryOccupationController(
            send_func=self.send_to_tokenizer,
            pause_controller=self._pause,
            adapter=TorchMemorySaverAdapter.create(
                enable=self.server_args.enable_memory_saver
            ),
            enabled=self.server_args.enable_memory_saver,
            reset_caches_fn=self._reset_caches_for_release,
            kv_repair_fn=self._kv_repair_after_wake,
            kv_cache_release_allowed=kv_cache_release_allowed,
        )

        self.metrics = EngineMetrics(
            labels={
                "model_name": server_args.served_model_name,
                "app_key": server_args.app_key or "",
                "dp_rank": str(dp_rank),
            },
            enabled=(
                server_args.enable_metrics
                and attn_tp_rank == 0
                and "prometheus" in (server_args.metrics_reporters or [])
            ),
        )

        self.request_handler = RequestHandler(
            server_args=self.server_args,
            hf_eos_token_id=self.model_config.hf_eos_token_id,
            max_req_len=self.model_config.context_len - 1,
            vocab_size=self.model_config.vocab_size,
            recv_func=self.recv_from_tokenizer,
            send_func=self.send_to_tokenizer,
            get_load_fn=self._get_load,
            architectures=self.model_config.hf_config.architectures,
            pause_controller=self._pause,
            memory_controller=self._memory,
        )

        self.output_processor = OutputProcesser(
            send_to_tokenizer=self.send_to_tokenizer,
            attn_tp_rank=attn_tp_rank,
            spec_algorithm=self.server_args.speculative_algorithm,
            spec_num_tokens=(
                self.server_args.speculative_num_draft_tokens
                if self.server_args.speculative_algorithm is not None
                else None
            ),
            stream_interval=self.server_args.stream_interval,
            enable_log_request_stats=self.server_args.enable_log_request_stats,
            metrics=self.metrics,
        )
        self.prefetch_threshold = scheduler_cfg.prefetch_threshold

        if server_args.disaggregation_mode != "null":
            kv_args = get_kv_args(
                global_rank,
                global_rank,
                server_args.disaggregation_ib_device,
                token_to_kv_pool,
                draft_token_to_kv_pool,
                mamba_pool,
            )
            pd_manager_args = KVManagerArgs(
                bootstrap_port=server_args.disaggregation_bootstrap_port,
                dist_init_addr=server_args.dist_init_addr,
                world_size=server_args.world_size or mapping.world_size,
                dp_size=server_args.data_parallel_size or mapping.attn.dp_size,
                attn_tp_rank=attn_tp_rank,
                attn_dp_rank=dp_rank,
                is_mla_backend=False,
                draft_is_mla_backend=False,
                enable_metrics=False,
                enable_mla_l1_5_cache=server_args.enable_mla_l1_5_cache,
                served_model_name=server_args.served_model_name,
                app_key=server_args.app_key,
                metrics_reporters=server_args.metrics_reporters,
                enable_dp_attention=self.has_dp,
            )
            self.kv_transfer = create_kv_transfer(
                mode=server_args.disaggregation_mode,
                backend=server_args.disaggregation_transfer_backend,
                args=pd_manager_args,
                kv_args=kv_args,
                gloo_group=self.attn_tp_cpu_group,
                page_size=token_to_kv_pool.page_size,
            )
            self._setup_pd_layerwise_transfer(
                server_args.disaggregation_layerwise_interval
            )
            # EPD: a multimodal prefill node is also the encode->prefill embedding
            # SINK (independent of kv_transfer, its P->D KV source) -- it receives
            # each image's embedding from encode workers over Mooncake so the
            # prefill skips the vision tower. The admission controller owns the
            # receive jobs, the rank-synced admission drain, and the optional NCCL
            # row-shard reassembly; None for decode/encode/text-only nodes.
            from tokenspeed.runtime.pd.epd.prefill_receiver import (
                make_epd_prefill_admission,
            )

            self.epd_admission = make_epd_prefill_admission(
                server_args,
                global_rank,
                model_config=self.model_config,
                model_executor=self.model_executor,
                mapping=mapping,
                attn_tp_rank=self.attn_tp_rank,
                attn_tp_size=self.attn_tp_size,
                attn_tp_cpu_group=self.attn_tp_cpu_group,
                pg_manager=pg_manager,
            )
            # Staged EPD request payloads (request_id -> (spec, state, bootstrap)),
            # held here while the controller (rid-keyed, like kv_transfer) runs the
            # async receive; popped in _drain_ready_epd_embeddings on admit/abort.
            self._epd_staged: dict = {}
        else:
            self.kv_transfer = None
            self.epd_admission = None
            self._epd_staged: dict = {}

    def _setup_pd_layerwise_transfer(self, interval: int) -> None:
        if not isinstance(self.kv_transfer, DisaggPrefillExecutor):
            return
        if interval <= 0:
            return

        from tokenspeed.runtime.pd.utils import StepCounter

        step_counter = StepCounter(self.model_executor.device, self.gpu_id)
        self.model_executor.attn_backend.register_step_counter(step_counter)
        if self.model_executor.draft_attn_backend is not None:
            self.model_executor.draft_attn_backend.register_step_counter(step_counter)
        self.kv_transfer.register_layerwise_step_counter(step_counter, interval)

    def _is_epd_request(self, state) -> bool:
        """True iff this request's images are encode-routed (smg injected per-image
        encode handshakes) -- it must wait for its embeddings (staged via the EPD
        admission controller, polled in _drain_ready_epd_embeddings) before being
        scheduled. Caller guards on self.epd_admission (only a multimodal prefill
        node has one); everything else admits immediately.
        """
        mm = getattr(state, "multimodal_inputs", None)
        return mm is not None and any(
            getattr(it, "encode_handshake", None) for it in mm.mm_items
        )

    def _assert_epd_embeddings_received(self, multimodal_context) -> None:
        """EPD invariant: every handshaked item is filled with its embedding by the
        async EPD admission drain (EpdPrefillAdmission.drain) BEFORE admission, so by
        it is already encoded. This is a defensive check, not a receive: a handshaked
        item that reached the forward un-received leaked past async admission (the
        only EPD admission path) -- fail loud instead of running the tower or
        publishing shard-only rows. No-op for non-EPD / text-only requests.
        """
        if (
            self.epd_admission is None
            or multimodal_context is None
            or not multimodal_context.has_extend_inputs()
        ):
            return
        for mm in multimodal_context.mm_inputs:
            if mm is None:
                continue
            missing = [
                i
                for i, item in enumerate(mm.mm_items)
                if getattr(item, "encode_handshake", None) is not None
                and item.encoded is None
            ]
            if missing:
                raise RuntimeError(
                    f"EPD: handshaked items {missing} reached the prefill forward "
                    "un-received; they must be admitted via the EPD admission drain"
                )

    def _drain_ready_epd_embeddings(self) -> None:
        """Admit EPD requests whose async embedding receives completed this cycle.

        The EpdPrefillAdmission controller DECIDES (poll + rank-lockstep MIN
        all-reduce + reassemble) and returns (admitted, failed); here we ACT on
        those decisions with the EventLoop's collaborators -- register/abort the
        P->D sender, submit admitted requests, finish failed ones. No-op (and no
        collective) on non-EPD nodes.
        """
        if self.epd_admission is None:
            return
        # Pause gate: withhold EPD admission while paused, mirroring the non-EPD
        # admit_blocked gate -- else the drain below would submit and RUN reassembled
        # specs during the pause. Staged receives wait in _pending until resume.
        # Rank-safe: admit_blocked is rank-identical, so all ranks skip together.
        if self._pause.admit_blocked:
            return
        admitted_ids, failed_ids = self.epd_admission.drain()
        for rid in failed_ids:
            spec, state, bootstrap = self._epd_staged.pop(rid)
            # Signal the dual-dispatched decode that this request failed so its KV
            # receiver fails (FailedEvent -> _process_kv_transfer_events abort)
            # instead of waiting forever for KV the prefill will never send. The
            # prefill never registered a P->D sender (deferred to admission), so the
            # decode has no other reliable way to learn (heartbeat only trips on a
            # dead prefill /health). Best-effort: only reaches decodes that already
            # pre-allocated.
            if (
                isinstance(self.kv_transfer, DisaggPrefillExecutor)
                and bootstrap is not None
            ):
                try:
                    self.kv_transfer.abort(rid, bootstrap)
                except Exception as exc:  # never let it wedge the loop
                    logger.warning(
                        "EPD abort->decode signal failed for rid=%s: %s",
                        rid,
                        exc,
                    )
            state.set_finish_with_abort("EPD embedding receive failed or timed out")
            self.output_processor.publish_finished_at_admission(rid, state)
        admitted_specs = []
        for rid in admitted_ids:
            spec, state, bootstrap = self._epd_staged.pop(rid)
            # Aborted mid-receive (no abort path, so drain still returns it admitted):
            # don't register the P->D sender or submit -- that runs a wasted forward
            # and leaks the sender. Stream its finish instead.
            if state.finished:
                self.output_processor.publish_finished_at_admission(rid, state)
                continue
            # Register the P->D sender now (deferred from admission) -- the request
            # is about to enter the scheduler.
            if self.kv_transfer is not None:
                self.kv_transfer.register(rid, bootstrap)
            admitted_specs.append(spec)
        if admitted_specs:
            self.scheduler.submit_requests(admitted_specs)
        elif self.epd_admission.has_pending():
            # Nothing advanced this cycle but requests are still receiving; yield the
            # GIL so the Python daemon transfer/recv threads run (rank-consistent:
            # admitted/leftover are rank-identical here).
            time.sleep(0.0005)

    def _commit_cache_results(self) -> None:
        if self.memory_executor is None:
            return
        cache_results = self.memory_executor.poll_results()
        self._num_inflight_cache_ops -= len(cache_results)
        for event in cache_results:
            payload = cache_event_to_payload(event)
            self._pending_cache_event_payloads[cache_event_key(payload)] = payload

        # The gather below is a collective, but cache-op completion is async and
        # not lock-step across ranks, so local state (_num_inflight_cache_ops /
        # _pending_cache_event_payloads) diverges transiently. A rank-local skip
        # would let some ranks gather while others return, deadlocking the group.
        # Agree on the skip via a cheap single-int all_reduce.
        # NOTE: For non-DFLASH algorithms, cache ops are deterministic across
        # ranks, so the local short-circuit is safe and avoids collective overhead.
        local_has_work = bool(
            self._num_inflight_cache_ops != 0 or self._pending_cache_event_payloads
        )
        if self.server_args.speculative_algorithm == "DFLASH":
            if not self._cache_group_has_work(local_has_work):
                return
        else:
            if not local_has_work:
                return

        ready_payloads = self._pop_ready_cache_event_payloads()
        if not ready_payloads:
            return
        logger.debug(
            "[cache_poll] got %s synchronized results, advancing scheduler",
            len(ready_payloads),
        )
        ec = ExecutionEvent()
        for payload in ready_payloads:
            e = cache_event_from_payload(payload)
            logger.debug(
                "[cache_poll] event: op_id=%s success=%s type=%s request_id=%s",
                e.op_id,
                e.success,
                type(e).__name__,
                getattr(e, "request_id", "N/A"),
            )
            ec.add_event(e)
        self.scheduler.advance(ec)
        logger.debug("[cache_poll] scheduler.advance() done")
        self._publish_scheduler_kv_events()

    def _publish_scheduler_kv_events(self) -> None:
        raw_events = drain_scheduler_kv_events(
            self.scheduler,
            enabled=self._kv_events_enabled,
        )
        if not raw_events:
            return

        events = scheduler_kv_events_to_wire_events(raw_events)
        if not events:
            return

        self.kv_event_publisher.publish(
            KVEventBatch(ts=time.time(), events=events, attn_dp_rank=self.dp_rank)
        )

    def _cache_group_has_work(self, local_has_work: bool) -> bool:
        """Whether ANY attn-tp rank has cache work this step (unanimous via a
        single-int MAX all_reduce, far cheaper than the payload gather it
        guards). Deciding from rank-local state alone deadlocks the group; see
        _commit_cache_results.

        Args:
            local_has_work: This rank's view of whether any cache op is in
                flight or any polled payload awaits commit.

        Returns:
            ``True`` if any rank has work (all must gather); ``False`` only when
            every rank is idle.
        """
        if self.attn_tp_size == 1:
            return local_has_work
        flag = torch.tensor([1 if local_has_work else 0], dtype=torch.int32)
        dist.all_reduce(flag, op=dist.ReduceOp.MAX, group=self.attn_tp_cpu_group)
        return bool(flag.item())

    def _pop_ready_cache_event_payloads(self) -> list[dict]:
        local_payloads = list(self._pending_cache_event_payloads.values())
        if self.attn_tp_size == 1:
            ready_payloads = local_payloads
        else:
            gathered_payloads = [None] * self.attn_tp_size
            dist.all_gather_object(
                gathered_payloads,
                local_payloads,
                group=self.attn_tp_cpu_group,
            )
            ready_payloads = pop_common_cache_event_payloads(gathered_payloads)
            if self.attn_tp_rank == 0 and cache_sync_debug_enabled():
                pending_ops = [
                    [(payload["kind"], payload["op_id"]) for payload in rank_payloads]
                    for rank_payloads in gathered_payloads
                ]
                if len({tuple(rank_ops) for rank_ops in pending_ops}) > 1:
                    logger.info(
                        "[cache_sync] rank=%s pending_ops=%s ready_ops=%s",
                        self.global_rank,
                        pending_ops,
                        [
                            (payload["kind"], payload["op_id"])
                            for payload in ready_payloads
                        ],
                    )

        for payload in ready_payloads:
            self._pending_cache_event_payloads.pop(cache_event_key(payload), None)
        return ready_payloads

    def _dispatch_forward(
        self,
        forward_op,
        sampling_params_list,
        execution_plan,
        dp_metadata=None,
        stats=None,
        grammar_inputs=None,
    ):
        """Execute one forward step; return (results, on_first_token).

        results is None when the step produces no model output (Path 2/3).
        Both event_loop and event_loop_overlap call this method; they differ
        only in *when* they call post_process on the returned results.

        Path 1 — no PD:              run forward, return (results, None)
        Path 2 — decode, extend:     trigger RDMA receive, return (None, None)
        Path 3 — prefill, decode:    send KV to decode side, return (None, None)
        Path 4 — prefill, extend:    run prefill forward, return (results, on_first_token)
        """
        if stats is None:
            stats = {}
        dp_global_num_tokens = (
            dp_metadata.global_num_tokens if dp_metadata is not None else None
        )
        dp_global_bs = (
            dp_metadata.global_batch_size if dp_metadata is not None else None
        )
        dp_all_decode_or_idle = (
            dp_metadata.all_decode_or_idle if dp_metadata is not None else False
        )
        dp_all_extend = dp_metadata.all_extend if dp_metadata is not None else False
        multimodal_context = self._get_multimodal_context_for_forward(forward_op)

        self.model_executor.update_block_table(forward_op)

        if self.kv_transfer is None:
            # Path 1: normal (no disaggregation)
            self.model_executor.reset_valid_cache_length(forward_op)
            return (
                self.model_executor.execute_forward_op_with_log(
                    forward_op,
                    sampling_params_list,
                    dp_global_num_tokens=dp_global_num_tokens,
                    dp_global_bs=dp_global_bs,
                    dp_all_decode_or_idle=dp_all_decode_or_idle,
                    dp_all_extend=dp_all_extend,
                    grammar_inputs=grammar_inputs,
                    multimodal_context=multimodal_context,
                    **stats,
                ),
                None,
            )

        elif isinstance(self.kv_transfer, DisaggDecodeExecutor):
            # Decode node
            if forward_op.num_extends() > 0:
                # Path 2: new requests waiting for remote KV — trigger RDMA receive
                self.kv_transfer.reset_valid_cache_length(
                    forward_op,
                    self.model_executor.runtime_states,
                    self.model_executor.execution_stream,
                    self.model_executor.device,
                )
                self.kv_transfer.execute(forward_op)
                self.model_executor.reset_remote_prefill_mamba_inputs(forward_op)
                return None, None
            else:
                # Path 3b: decode batch — normal forward
                self.model_executor.reset_valid_cache_length(forward_op)
                return (
                    self.model_executor.execute_forward_op_with_log(
                        forward_op,
                        sampling_params_list,
                        dp_global_num_tokens=dp_global_num_tokens,
                        dp_global_bs=dp_global_bs,
                        dp_all_decode_or_idle=dp_all_decode_or_idle,
                        dp_all_extend=dp_all_extend,
                        multimodal_context=multimodal_context,
                        **stats,
                    ),
                    None,
                )

        else:
            # Prefill node (only reached from event_loop, never event_loop_overlap)
            if not isinstance(self.kv_transfer, DisaggPrefillExecutor):
                raise TypeError("kv_transfer must be a DisaggPrefillExecutor.")
            if forward_op.num_extends() == 0:
                # Path 3: all prefill done — send KV to decode side
                self.kv_transfer.execute(forward_op)
                return None, None
            else:
                # Path 4: extend batch — run prefill forward
                self.model_executor.reset_valid_cache_length(forward_op)
                self.kv_transfer.prepare_prefill(forward_op)
                # EPD invariant: handshaked items are filled by the async
                # EPD admission drain before admission; assert none reached
                # the forward un-received (no-op for non-EPD / text-only requests).
                self._assert_epd_embeddings_received(multimodal_context)
                return (
                    self.model_executor.execute_forward_op_with_log(
                        forward_op,
                        sampling_params_list,
                        dp_global_num_tokens=dp_global_num_tokens,
                        dp_global_bs=dp_global_bs,
                        dp_all_decode_or_idle=dp_all_decode_or_idle,
                        dp_all_extend=dp_all_extend,
                        grammar_inputs=grammar_inputs,
                        multimodal_context=multimodal_context,
                        capture_next_input_ids=True,
                        **stats,
                    ),
                    self.kv_transfer.store_prefill_token,
                )

    def _get_multimodal_context_for_forward(self, forward_op):
        if not self.model_config.is_multimodal_active:
            return None

        num_extends = forward_op.num_extends()
        mm_inputs = []
        has_mm = False
        for index, rid in enumerate(forward_op.request_ids):
            state = self.output_processor.rid_to_state.get(rid)
            if state is not None and index < num_extends:
                state.maybe_extend_multimodal_mrope_positions()
            item = getattr(state, "multimodal_inputs", None) if state else None
            mm_inputs.append(item)
            has_mm = has_mm or item is not None
        if not has_mm:
            return None

        from tokenspeed.runtime.multimodal.inputs import MultimodalForwardContext

        return MultimodalForwardContext(
            mm_inputs=mm_inputs,
            extend_prefix_lens=list(forward_op.extend_prefix_lens),
            extend_seq_lens=list(forward_op.input_lengths[:num_extends]),
        )

    def _build_mamba_layerwise_cow(
        self, execution_plan, forward_op
    ) -> dict[int, list[int]]:
        if forward_op is None:
            return {}
        loaded_mamba_slots: set[int] = set()
        for cache_op in execution_plan.cache:
            if not isinstance(cache_op, Cache.LoadBackOp):
                continue
            dst_by_kind = getattr(cache_op, "dst_pages_by_kind", None)
            if dst_by_kind is None:
                dst_groups = getattr(cache_op, "dst_pages", [])
            else:
                dst_groups = dst_by_kind.get(CacheKind.MAMBA.value, [])
            for dst_pages in dst_groups:
                loaded_mamba_slots.update(int(page) for page in dst_pages)
        if not loaded_mamba_slots:
            return {}

        cow_src_indices = getattr(forward_op, "mamba_cow_src_indices", None)
        working_indices = getattr(forward_op, "mamba_pool_indices", None)
        if cow_src_indices is None or working_indices is None:
            return {}

        cow_by_src: dict[int, list[int]] = {}
        for cow_src, working in zip(list(cow_src_indices), list(working_indices)):
            cow_src = int(cow_src)
            working = int(working)
            if cow_src < 0 or working < 0 or cow_src not in loaded_mamba_slots:
                continue
            cow_dsts = cow_by_src.setdefault(cow_src, [])
            if working not in cow_dsts:
                cow_dsts.append(working)
        return cow_by_src

    def _submit_cache_ops(self, execution_plan) -> None:
        if self.memory_executor is None:
            return
        forward_op = self._get_forward_op(execution_plan)
        mamba_layerwise_cow = self._build_mamba_layerwise_cow(
            execution_plan, forward_op
        )
        if mamba_layerwise_cow:
            self.model_executor.set_layerwise_mamba_cow_done(mamba_layerwise_cow)
            self.memory_executor.set_mamba_layerwise_cow(mamba_layerwise_cow)
        self.memory_executor.submit_plan(execution_plan)
        for op in execution_plan.cache:
            if isinstance(op, Cache.WriteBackOp):
                self._num_inflight_cache_ops += len(op.op_ids)
            elif isinstance(op, Cache.LoadBackOp):
                # Radix loadbacks are fire-and-forget (no ack, nothing in
                # flight); the flat host tier acks one LoadBackDone per op_id.
                if self._loadback_acks_expected:
                    self._num_inflight_cache_ops += len(op.op_ids)
            elif isinstance(op, (Cache.PrefetchOp, Cache.BackUpOp)):
                self._num_inflight_cache_ops += 1
            else:
                raise ValueError(f"unsupported cache op kind: {type(op).__name__}")
        self._setup_layerwise_loadback(execution_plan)

    def _setup_layerwise_loadback(self, execution_plan) -> None:
        host_exec = getattr(self.memory_executor, "host_exec", None)
        available_pools = (
            getattr(host_exec, "pools", {}) if host_exec is not None else {}
        )
        consumer_indices_by_kind: dict[CacheKind, list[int]] = {
            kind: [] for kind in available_pools
        }
        for cache_op in execution_plan.cache:
            if isinstance(cache_op, Cache.LoadBackOp):
                for op_id in cache_op.op_ids:
                    for kind in consumer_indices_by_kind:
                        producer_idx = self.memory_executor.get_producer_index(
                            kind, op_id
                        )
                        if (
                            producer_idx is not None
                            and producer_idx not in consumer_indices_by_kind[kind]
                        ):
                            consumer_indices_by_kind[kind].append(producer_idx)
        for kind, consumer_indices in consumer_indices_by_kind.items():
            self.memory_executor.set_consumer(
                kind, consumer_indices if consumer_indices else -1
            )

    def _flush_mamba_retract_states(self, forward_op) -> None:
        """Copy draft->working mamba states when retract occurred (no forward scheduled)."""
        if forward_op is not None:
            return
        if self.model_executor.drafter is None:
            return
        if self.model_executor.runtime_states.mamba_pool is None:
            return
        self.model_executor.flush_mamba_draft_to_working_on_retract()

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _load_model_config(
        self, model_path: str, is_draft_worker: bool = False
    ) -> ModelConfig:
        server_args = self.server_args
        quantization = server_args.quantization
        if is_draft_worker:
            quantization = server_args.speculative_draft_model_quantization
        return ModelConfig(
            model_path,
            trust_remote_code=server_args.trust_remote_code,
            revision=server_args.revision,
            context_length=server_args.max_model_len,
            model_override_args=server_args.hf_overrides,
            dtype=server_args.dtype,
            quantization=quantization,
            server_args=server_args,
            is_draft_worker=is_draft_worker,
        )

    def _init_distributed(self) -> float:
        max_num_input_tokens = (
            self.server_args.chunked_prefill_size
            if self.server_args.chunked_prefill_size > 0
            else self.server_args.max_prefill_tokens + self.server_args.max_model_len
        )
        distributed_config = DistributedConfig.from_server_args(
            server_args=self.server_args,
            port_args=self.port_args,
            gpu_id=self.gpu_id,
            global_rank=self.global_rank,
            hidden_size=self.model_config.hidden_size,
            max_num_tokens=max_num_input_tokens,
        )
        return DistributedInitializer.initialize(distributed_config)

    def _init_interprocess_comm(self):
        context = zmq.Context(2)
        if self.attn_tp_rank == 0:
            self.recv_from_tokenizer = get_zmq_socket(
                context, zmq.PULL, self.port_args.scheduler_input_ipc_name, False
            )
            self.send_to_tokenizer = get_zmq_socket(
                context, zmq.PUSH, self.port_args.tokenizer_ipc_name, False
            )
        else:
            self.recv_from_tokenizer = None
            self.send_to_tokenizer = _NullSender()

    # ------------------------------------------------------------------
    # Shared step helpers
    # ------------------------------------------------------------------

    def _reap_or_keep_buffered_spec(self, spec) -> bool:
        """Resolve a buffered spec on resume; return True if it should be admitted.

        A buffered spec was already registered in ``rid_to_state`` before it was
        withheld, so if it was aborted while paused it never reached the
        scheduler and the forward path can never reap it. Handle that here:

        - state missing  -> already published and reaped; drop silently.
        - state finished -> aborted in place. Stream a terminating finish for
          pause-initiated aborts (the passive client is still waiting) and drop
          the registered state so the rid does not leak; client-initiated aborts
          already tore down their own state, so just reap.
        - otherwise      -> still live; admit it.
        """
        state = self.output_processor.rid_to_state.get(spec.request_id)
        if state is None:
            return False
        if state.finished:
            self.output_processor.reap_finished_orphan(spec.request_id, state)
            return False
        return True

    def _process_new_requests(self):
        recv_reqs = self.request_handler.recv_reqs()
        # Snapshot the pause state before dispatch: process_requests may flip it
        # mid-batch. If it was not blocked before but is after, a pause control
        # message was processed in this very batch — which is what makes the
        # FIFO edge below detectable (see TODO(pause-fifo)).
        pause_blocked_before = self._pause.admit_blocked
        new_req_specs, new_req_states, bootstrap_infos, abort_rids = (
            self.request_handler.process_requests(recv_reqs)
        )
        # Sweep TTL-expired abort markers every iteration. Without this
        # the map only gets cleaned inside ``mark_abort``, so a burst of
        # stale-cancel traffic followed by silence leaves the last batch
        # of entries sitting past their TTL (and potentially re-aborting
        # reused rids). Amortized O(1): expired entries are always at
        # the front of the insertion-ordered dict.
        self.output_processor.sweep_pending_aborts()
        # Abort both registered and grammar-queued requests. Without the
        # grammar_manager.mark_abort call, a request aborted mid-compile
        # would finish compiling and get admitted before being noticed.
        grammar_manager = self.request_handler.grammar_manager
        for rid in abort_rids:
            self.output_processor.mark_abort(rid)
            grammar_manager.mark_abort(rid)

        # A pause(mode="abort") cancels every in-flight request through the same
        # marker path as a client abort; they finish on their next scheduled
        # step, then the drain check resolves the pause reply.
        if self._pause.consume_abort_all():
            for rid in list(self.output_processor.rid_to_state.keys()):
                # notify_client=True: pause aborts a passive client's request,
                # so it must receive a terminating finish (unlike a client abort).
                self.output_processor.mark_abort(rid, notify_client=True)
                grammar_manager.mark_abort(rid)

        # abort/wait also cancel requests still compiling in the grammar queue:
        # they are not yet in rid_to_state or the scheduler, so the sweep above
        # and the drain check both miss them. A finished state makes the next
        # get_ready_grammar_requests pass publish them instead of admitting, so
        # they never run under post-resume weights or strand the drain.
        if self._pause.consume_cancel_grammar():
            for _, state, _ in grammar_manager.grammar_queue:
                state.set_finish_with_abort("Aborted by pause", notify_client=True)

        # On resume, flush specs buffered while paused even when no new request
        # arrives this iteration. This must run before the ``if not ready:
        # return`` guard below, which would otherwise strand buffered specs
        # until the next inbound request. Specs aborted while paused are reaped
        # in place (terminating finish + state cleanup) rather than admitted, so
        # they don't burn a scheduler slot or leak their rid — see
        # ``_reap_or_keep_buffered_spec``.
        if not self._pause.admit_blocked and self._pause.buffered_specs:
            specs = [
                spec
                for spec in self._pause.take_buffered_specs()
                if self._reap_or_keep_buffered_spec(spec)
            ]
            if specs:
                self.scheduler.submit_requests(specs)

        # Partition new requests by grammar readiness. Compile-bound requests
        # are queued in GrammarManager and admitted in a later iteration when
        # their futures resolve (see _drain_ready_grammar_requests below).
        ready = []
        for spec, state, bootstrap in zip(
            new_req_specs, new_req_states, bootstrap_infos
        ):
            # Requests pre-marked finished (e.g. invalid session ID aborted
            # in RequestHandler) skip grammar compilation entirely — we'd
            # just be wasting a compile slot on a response we're about to
            # abort anyway, and the terminal response would be delayed by
            # the compile/timeout window.
            if state.finished:
                ready.append((spec, state, bootstrap))
                continue
            if grammar_manager.process_req_with_grammar(state):
                ready.append((spec, state, bootstrap))
            else:
                grammar_manager.add_to_queue(spec, state, bootstrap)

        # Drain any previously-queued requests whose grammar just finished
        # compiling. With attn_tp > 1 this also drives the per-iter all_gather
        # that keeps grammar admission in sync across ranks.
        ready.extend(grammar_manager.get_ready_grammar_requests())

        if not ready:
            return

        admitted_specs = []
        for spec, state, bootstrap in ready:
            # Grammar-aborted (invalid grammar, timed-out compile, or missing
            # backend) requests must not enter the scheduler — they have no
            # valid grammar to mask logits with, and we don't want to spend a
            # prefill slot on a request that's already finished. Publish the
            # finish_reason directly so the client still gets a response.
            if state.finished:
                self.output_processor.publish_finished_at_admission(
                    spec.request_id, state
                )
                continue

            if isinstance(self.kv_transfer, DisaggDecodeExecutor):
                state.computed_length = state.input_length
            self.output_processor.register(spec.request_id, state)
            is_epd = self.epd_admission is not None and self._is_epd_request(state)
            # EPD: DEFER the P->D sender registration to admission (in
            # EpdPrefillAdmission.drain, just before submit_requests). Registering
            # it now -- while the request is staged and NOT yet in the C++ scheduler
            # -- would let DisaggPrefillExecutor.generate_events poll the sender and
            # emit a BootstrappedEvent that the scheduler's requests_.at(rid) THROWS
            # on (no such request yet). Non-EPD requests register now (submitted this
            # same call).
            if self.kv_transfer is not None and not is_epd:
                self.kv_transfer.register(spec.request_id, bootstrap)

            if self.memory_executor is not None:
                hashes = calc_l3_query_hashes(self.scheduler, spec.tokens)
                if hashes and len(hashes) > self.prefetch_threshold:
                    hit_pages = self.memory_executor.query_l3_pages(hashes)
                    logger.debug(
                        "[cache_op] L3 query: rid=%s hash_pages=%s hit_pages=%s threshold=%s",
                        spec.request_id,
                        len(hashes),
                        hit_pages,
                        self.prefetch_threshold,
                    )
                    spec.rolling_hashes = hashes
                    spec.storage_hit_pages = hit_pages
            # EPD prefill: hold a request whose images are encode-routed OUT of the
            # scheduler until its per-image embeddings have been received (started
            # here, polled in EpdPrefillAdmission.drain, which registers the P->D
            # sender + submits once ready). It is output_processor-registered above;
            # the sender registration + submission are both deferred. Non-EPD
            # requests admit immediately as before. Rank-identical because `ready` is
            # rank-synced (recv_reqs broadcast + grammar gather).
            if is_epd:
                self.epd_admission.stage(
                    spec.request_id, state.multimodal_inputs.mm_items
                )
                self._epd_staged[spec.request_id] = (spec, state, bootstrap)
            else:
                admitted_specs.append(spec)

        # Pause gate: while paused, withhold new requests from the scheduler
        # (running requests keep stepping); buffered specs are flushed on resume
        # above, ahead of any newly-admitted ones, preserving FIFO order.
        #
        # TODO(pause-fifo): recv_reqs() drains the socket non-blocking, so a
        # generate request that arrived *before* a pause control message can be
        # coalesced into the same batch and reach here after the pause flipped
        # admit_blocked. Such a pre-pause request is buffered as post-pause work
        # instead of running (wait) / being aborted (abort). Correct handling
        # needs the batch processed as an ordered stream that respects the
        # control request's FIFO position. Tracked as a follow-up; until then we
        # warn when the coalescing condition is observed so it is not silent.
        if self._pause.admit_blocked:
            if admitted_specs and not pause_blocked_before:
                logger.warning(
                    "Pause engaged in the same recv batch as %d generate "
                    "request(s) (rids=%s); their FIFO order relative to the "
                    "pause is not preserved, so a pre-pause request may be "
                    "buffered as post-pause work and run only after resume. "
                    "See TODO(pause-fifo).",
                    len(admitted_specs),
                    [spec.request_id for spec in admitted_specs],
                )
            self._pause.buffer_specs(admitted_specs)
            return

        if admitted_specs:
            self.scheduler.submit_requests(admitted_specs)

    @nvtx_range("loop:commit", color="rapids")
    def _commit_forward_results(
        self,
        forward_op,
        results: ModelExecutionResult,
        on_first_token=None,
    ):
        self.request_handler.forward_ct += 1
        forward_mode = ForwardMode.from_num_extends(
            forward_op.num_extends(),
            len(forward_op.request_ids),
        )
        self.request_handler._profile_batch_predicate(forward_mode)

        # post_process_forward_op calls sync() — after this, CPU tensors are ready
        is_prefill_instance = isinstance(self.kv_transfer, DisaggPrefillExecutor)
        request_changes = self.output_processor.post_process_forward_op(
            forward_op,
            results,
            is_prefill_instance=is_prefill_instance,
            on_first_token=on_first_token,
        )

        # Accumulate decode stats from synced results (no GPU sync)
        if forward_op.num_extends() <= 0:
            bs = len(forward_op.request_ids)
            self.model_executor.accumulate_decode_stats(results, bs)

        return request_changes

    def _get_forward_op(self, execution_plan):
        """Return the next forward op from the given plan, or None if there is nothing to run."""
        forward_ops = execution_plan.forward
        if len(forward_ops) == 0 or len(forward_ops[0].request_ids) == 0:
            return None
        return forward_ops[0]

    def _handle_flat_oom_terminals(self, execution_plan) -> None:
        """Surface flat-KV OOM terminals to their clients as abort finishes.

        The C++ flat scheduler terminalizes a request that can never fit the
        flat pool (AbortEvent inside the scheduler; the reaper reclaims its
        resources) and reports its id on ``plan.flat_oom_request_ids``
        (always empty on radix builds). The scheduler side is already fully
        torn down — do NOT send a ForwardEvent.Abort back — but the client is
        still waiting on the response stream, so finish the request with an
        abort here (mirrors the PD FailedEvent handling above, minus the
        scheduler abort).
        """
        oom_rids = getattr(execution_plan, "flat_oom_request_ids", None)
        if not oom_rids:
            return
        for rid in oom_rids:
            state = self.output_processor.rid_to_state.get(rid)
            if state is None:
                # rid already gone (e.g. a client abort raced ahead).
                logger.debug(
                    "flat OOM terminal for rid=%s: state missing; skipping", rid
                )
                continue
            if state.finished:
                # Already carries a finish (an abort raced ahead of the
                # terminal). C++ reports this rid exactly once and no future
                # forward op will reap it, so resolve it here (same orphan
                # rule as _reap_or_keep_buffered_spec).
                self.output_processor.reap_finished_orphan(rid, state)
                continue
            state.set_finish_with_abort(
                "flat KV cache cannot fit this request: prompt exceeds pool "
                "capacity (OOM)"
            )
            self.output_processor.publish_finished_at_admission(rid, state)

    def _process_kv_transfer_events(self, kv_transfer_events: list) -> list:
        processed = []
        for event in kv_transfer_events:
            processed.append(event)
            if isinstance(event, PD.SucceededEvent) and isinstance(
                self.kv_transfer, DisaggPrefillExecutor
            ):
                req_id = event.request_id
                processed.extend(self.output_processor.finish_prefill_request(req_id))
            elif isinstance(event, PD.RemotePrefillDoneEvent):
                req_id = event.request_id
                bootstrap_token = event.bootstrap_token
                self.output_processor.on_remote_prefill_done(req_id, bootstrap_token)
                if isinstance(self.kv_transfer, DisaggDecodeExecutor):
                    candidate_info = self.kv_transfer.pop_remote_spec_candidate_ids(
                        req_id
                    )
                    if candidate_info is not None:
                        req_pool_idx, candidate_ids = candidate_info
                        self.model_executor.write_remote_spec_candidate_ids(
                            req_pool_idx, candidate_ids
                        )
            elif isinstance(event, PD.FailedEvent):
                # A PD/EPD transfer failed: the decode KV receiver timed out (e.g. the
                # prefill aborted on embedding timeout so the KV never arrives), or a
                # transfer errored. The C++ scheduler's FailedEvent handler is a no-op,
                # so without this the request is never finished and the CLIENT HANGS
                # FOREVER (the decode is its response stream). Finish it with an abort
                # (streams the error to the client) and abort it in the scheduler so its
                # slot/KV is freed. AbortEvent is valid from the decode-waiting state
                # (forward_events.cpp: AbortEvent(Prefilling&&) -> Finished).
                req_id = event.request_id
                state = self.output_processor.rid_to_state.get(req_id)
                if state is not None:
                    state.set_finish_with_abort(
                        "PD/EPD remote transfer failed or timed out"
                    )
                    self.output_processor.publish_finished_at_admission(req_id, state)
                    abort = ForwardEvent.Abort()
                    abort.request_id = req_id
                    processed.append(abort)

        return processed

    def _get_load(self):
        """Return load metrics for the DP load balancer."""
        from tokenspeed.runtime.engine.io_struct import GetLoadReqOutput

        available = self.scheduler.available_kv_pages()
        num_total_pages = self.max_total_num_tokens // self.server_args.block_size
        num_used_pages = num_total_pages - available
        num_waiting = self.scheduler.waiting_size()
        # num_reqs: running + waiting (used by SHORTEST_QUEUE balancing)
        num_running = len(self.output_processor.rid_to_state)
        return GetLoadReqOutput(
            dp_rank=self.dp_rank,
            num_reqs=num_running + num_waiting,
            num_waiting_reqs=num_waiting,
            num_pages=num_used_pages,
        )

    def _dp_sync_and_check(self, forward_op) -> DpForwardMetadata:
        """Synchronize DP ranks with CPU-only metadata.

        All ranks call this before GPU forward work. The gathered metadata is
        used for eager token-aware collectives and for choosing a common padded
        CUDA graph shape during decode.
        """
        import torch.distributed as dist

        executes_model_forward = _forward_op_executes_model_forward(
            forward_op,
            is_disagg_decode=isinstance(self.kv_transfer, DisaggDecodeExecutor),
        )
        num_tokens = sum(forward_op.input_lengths) if executes_model_forward else 0
        batch_size = len(forward_op.request_ids) if executes_model_forward else 0
        if not executes_model_forward:
            forward_mode = ForwardMode.IDLE
        else:
            forward_mode = ForwardMode.from_num_extends(
                forward_op.num_extends(),
                batch_size,
            )

        self._dp_local_info[0, 0] = num_tokens
        self._dp_local_info[0, 1] = batch_size
        self._dp_local_info[0, 2] = int(forward_mode)
        dist.all_gather_into_tensor(
            self._dp_global_info,
            self._dp_local_info,
            group=self.world_cpu_group,
        )
        global_num_tokens = self._dp_global_info[:, 0].tolist()
        global_batch_size = self._dp_global_info[:, 1].tolist()
        global_forward_mode = self._dp_global_info[:, 2].tolist()
        any_rank_has_work = max(global_num_tokens) > 0
        need_idle_forward = num_tokens == 0 and any_rank_has_work
        all_decode_or_idle = all(
            mode
            in (
                int(ForwardMode.DECODE),
                int(ForwardMode.IDLE),
            )
            for mode in global_forward_mode
        )
        # Replicated prefill-graph gate (see PrefillGraph._select_bucket).
        all_extend = all(
            mode == int(ForwardMode.EXTEND) for mode in global_forward_mode
        )
        return DpForwardMetadata(
            global_num_tokens=global_num_tokens,
            global_batch_size=global_batch_size,
            global_forward_mode=global_forward_mode,
            all_decode_or_idle=all_decode_or_idle,
            all_extend=all_extend,
            need_idle_forward=need_idle_forward,
        )

    def _get_scheduler_stats(self):
        """Query scheduler for page usage and queue depth."""
        available = self.scheduler.available_kv_pages()
        active = self.scheduler.active_kv_pages()
        num_total_pages = self.max_total_num_tokens // self.server_args.block_size
        return {
            "num_active_pages": active,
            "num_cached_pages": num_total_pages - available,
            "num_queue_reqs": self.scheduler.waiting_size(),
        }

    def _record_scheduler_iteration_metrics(
        self, stats: dict, num_iteration_tokens: int
    ) -> None:
        self.metrics.record_scheduler_iteration(
            running=len(self.output_processor.rid_to_state),
            waiting=stats["num_queue_reqs"],
            num_active_pages=stats["num_active_pages"],
            num_total_pages=self.max_total_num_tokens // self.server_args.block_size,
            num_iteration_tokens=num_iteration_tokens,
        )

    # ------------------------------------------------------------------
    # Pause / resume helpers
    # ------------------------------------------------------------------

    def _reset_caches_for_release(self) -> None:
        """Invalidate the prefix/radix cache before KV is discarded on release.

        KV pages are re-mapped + zeroed on wake, so any retained prefix entry
        would be stale. The unsafe case (prefix caching on with no reset) is
        rejected up front in ``MemoryOccupationController.handle_release`` via
        ``kv_cache_release_allowed``, so by the time we get here either a reset
        exists or prefix caching is off (nothing to invalidate).
        """
        reset = getattr(self.scheduler, "reset_prefix_cache", None)
        if callable(reset):
            reset()

    def _kv_pools(self) -> list:
        """All KV pools whose pages are tagged ``kv_cache`` — the target pool and
        the draft pool in speculative-decoding runs. Release/repair must walk the
        SAME set, so both derive it here rather than enumerating pools by hand."""
        pools = []
        for attr in ("token_to_kv_pool", "draft_token_to_kv_pool"):
            pool = getattr(self.model_executor, attr, None)
            if pool is not None:
                pools.append(pool)
        return pools

    def _kv_repair_after_wake(self) -> None:
        """Zero re-mapped KV buffers (garbage after re-map) for every KV pool,
        including the draft pool in spec-decode runs — its allocations are tagged
        ``kv_cache`` too, so a wake that skipped it would feed the draft model
        stale KV. FP8 KV scales ride with the weights region, so no scale reset
        is needed here."""
        for pool in self._kv_pools():
            if hasattr(pool, "clear_kv_buffers"):
                pool.clear_kv_buffers()

    def _paused_idle_step(self, prev_forward_op=None, prev_results=None) -> None:
        """Run one iteration under ``PAUSED_ALL`` (keep mode): no new forward
        work, but keep DP ranks in lockstep, service the drain check, and yield
        the CPU so the freeze does not busy-spin a core."""
        if prev_results is not None:
            request_changes = self._commit_forward_results(
                prev_forward_op, prev_results
            )
            advance_forward(self.scheduler, request_changes)
            self._publish_scheduler_kv_events()

        if self.has_dp:
            dp_metadata = self._dp_sync_and_check(None)
            # While memory is released the weights region is unmapped; an idle
            # forward runs the model and would read freed memory. All DP ranks
            # release together, so skipping the idle forward stays consistent
            # across ranks (the small DP sync above still runs to keep lockstep).
            if dp_metadata.need_idle_forward and not self._pause.released:
                self.model_executor.execute_idle_forward(
                    dp_metadata.global_num_tokens,
                    dp_metadata.global_batch_size,
                    dp_metadata.all_decode_or_idle,
                )

        self._pause.maybe_finish_drain(self.scheduler)
        time.sleep(_PAUSED_IDLE_SLEEP_S)

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

    def event_loop(self):
        """Non-overlapping scheduler loop."""
        while True:
            self._process_new_requests()
            # EPD prefill: admit requests whose async embedding receives completed
            # this cycle (rank-synced). Fixed position right after
            # _process_new_requests so the drain's TP collective ordering is
            # rank-identical every cycle.
            self._drain_ready_epd_embeddings()
            self._commit_cache_results()
            if self._pause.forward_blocked:
                self._paused_idle_step()
                continue
            execution_plan = self.scheduler.next_execution_plan()
            self._publish_scheduler_kv_events()
            self._handle_flat_oom_terminals(execution_plan)
            self._submit_cache_ops(execution_plan)

            forward_op = self._get_forward_op(execution_plan)
            self._flush_mamba_retract_states(forward_op)

            stats = self._get_scheduler_stats()
            num_iter_tokens = (
                sum(forward_op.input_lengths) if forward_op is not None else 0
            )

            # DP sync: all ranks must participate even when idle.
            dp_metadata = None
            if self.has_dp:
                dp_metadata = self._dp_sync_and_check(forward_op)
                if dp_metadata.need_idle_forward:
                    self.model_executor.execute_idle_forward(
                        dp_metadata.global_num_tokens,
                        dp_metadata.global_batch_size,
                        dp_metadata.all_decode_or_idle,
                    )
                    self._record_scheduler_iteration_metrics(stats, num_iter_tokens)
                    continue

            request_changes = []

            if forward_op is not None:
                sampling_params_list = self._gather_sampling_params(forward_op)
                grammar_inputs = self._gather_grammar_state(forward_op)
                self._mark_stats_scheduled(forward_op)
                results, on_first_token = self._dispatch_forward(
                    forward_op,
                    sampling_params_list,
                    execution_plan,
                    dp_metadata=dp_metadata,
                    stats=stats,
                    grammar_inputs=grammar_inputs,
                )
                if results is not None:
                    request_changes.extend(
                        self._commit_forward_results(
                            forward_op, results, on_first_token
                        )
                    )

            if self.kv_transfer is not None:
                kv_transfer_events = self.kv_transfer.generate_events()
                request_changes.extend(
                    self._process_kv_transfer_events(kv_transfer_events)
                )

            if request_changes:
                advance_forward(self.scheduler, request_changes)
                self._publish_scheduler_kv_events()

            # Resolve a deferred abort/wait pause reply once in-flight work drains.
            self._pause.maybe_finish_drain(self.scheduler)

            self._record_scheduler_iteration_metrics(stats, num_iter_tokens)

    def _mark_stats_scheduled(self, forward_op) -> None:
        # Stamp the pre-forward "scheduled" time on each request's stats tracker
        # so the queue/prefill split is anchored before the forward (idempotent:
        # only the first forward a request appears in sets it). --enable-log-request-stats.
        if not self.server_args.enable_log_request_stats or forward_op is None:
            return
        now = time.time()
        rid_to_state = self.output_processor.rid_to_state
        for rid in forward_op.request_ids:
            st = rid_to_state.get(rid)
            if st is not None:
                st.stats.mark_scheduled(now)

    def _gather_sampling_params(self, forward_op) -> list[SamplingParams]:
        """Look up per-request SamplingParams from the output processor. The
        sampling backend does its own flip detection + RNG state management
        internally, so we only need the scalar params here."""
        return [
            self.output_processor.rid_to_state[rid].sampling_params
            for rid in forward_op.request_ids
        ]

    def _gather_grammar_state(self, forward_op) -> GrammarStepInputs | None:
        """Build ``GrammarStepInputs`` for the current batch, or ``None``.

        Returns ``None`` when no request in this batch has a grammar — the
        model_executor short-circuits then. Otherwise carries the grammars
        list + per-EXTEND-slot ``advance_mask`` (False on intermediate
        chunked-prefill chunks, since the sampled token is discarded by
        post_process and must not advance the matcher).
        """
        rid_to_state = self.output_processor.rid_to_state
        grammars = [rid_to_state[rid].grammar for rid in forward_op.request_ids]
        if not any(grammars):
            return None

        advance_mask = None
        num_extends = forward_op.num_extends()
        if num_extends > 0:
            bs = len(forward_op.request_ids)
            extend_prefix_lens = forward_op.extend_prefix_lens
            extend_input_lengths = forward_op.input_lengths[:num_extends]
            advance_mask = [True] * bs
            for i in range(num_extends):
                rid = forward_op.request_ids[i]
                # This chunk completes prefill iff it processes the final
                # token of the prompt; intermediate chunks don't.
                advance_mask[i] = (
                    extend_prefix_lens[i] + extend_input_lengths[i]
                    >= rid_to_state[rid].input_length
                )

        return GrammarStepInputs(grammars=grammars, advance_mask=advance_mask)

    def event_loop_overlap(self):
        """
        Overlapping scheduler loop: post-process the previous step's results
        while the current step's forward pass is in flight.
        """
        # EPD invariant: the async embedding drain (EpdPrefillAdmission.drain)
        # that admits EPD requests runs ONLY in event_loop(), never here. A
        # prefill node that receives encode embeddings must therefore run the
        # non-overlap loop -- should_use_overlap_schedule enforces this by forcing
        # prefill -> non-overlap. Assert it rather than trusting that external
        # coupling: if a prefill ever reached this loop, every EPD request would
        # stage into the admission controller and hang forever with no drain.
        assert self.epd_admission is None, (
            "EPD prefill must run the non-overlap event_loop(); the embedding "
            "drain is not wired into event_loop_overlap()"
        )
        prev_results: ModelExecutionResult = None
        prev_forward_op = None

        while True:
            # Order this iter's default-stream writes (KVAllocator,
            # update_block_table, prefix_cache writes to req_to_page)
            # after the prev iter's forward on execution_stream that
            # reads the same tensor. Non-blocking on host.
            torch.cuda.default_stream().wait_stream(
                self.model_executor.execution_stream
            )
            self._process_new_requests()
            self._commit_cache_results()
            if self._pause.forward_blocked:
                # Freeze: commit any in-flight (overlapped) step — a forward
                # already on the GPU can't be un-launched — then idle.
                self._paused_idle_step(prev_forward_op, prev_results)
                prev_results = None
                prev_forward_op = None
                continue
            execution_plan = self.scheduler.next_execution_plan()
            self._publish_scheduler_kv_events()
            self._handle_flat_oom_terminals(execution_plan)

            self._submit_cache_ops(execution_plan)

            forward_op = self._get_forward_op(execution_plan)
            self._flush_mamba_retract_states(forward_op)

            stats = self._get_scheduler_stats()
            num_iter_tokens = (
                sum(forward_op.input_lengths) if forward_op is not None else 0
            )

            grammar_inputs = None
            if forward_op is not None:
                # Gather both sampling params and grammar state BEFORE the
                # prev_results commit below — that commit can finish requests
                # and pop them from output_processor.rid_to_state, which would
                # KeyError when we look up rids that are still in the current
                # forward_op.
                sampling_params_list = self._gather_sampling_params(forward_op)
                grammar_inputs = self._gather_grammar_state(forward_op)

            # DP sync: all ranks must participate even when idle.
            dp_metadata = None
            if self.has_dp:
                dp_metadata = self._dp_sync_and_check(forward_op)
                if dp_metadata.need_idle_forward:
                    if prev_results is not None:
                        request_changes = self._commit_forward_results(
                            prev_forward_op, prev_results
                        )
                        advance_forward(self.scheduler, request_changes)
                        self._publish_scheduler_kv_events()
                        prev_results = None
                        prev_forward_op = None
                    self.model_executor.execute_idle_forward(
                        dp_metadata.global_num_tokens,
                        dp_metadata.global_batch_size,
                        dp_metadata.all_decode_or_idle,
                    )
                    self._record_scheduler_iteration_metrics(stats, num_iter_tokens)
                    continue

            # ---- dispatch current forward first (async GPU launch) ----
            # Issue curr's forward before committing prev so the GPU runs curr
            # while the CPU syncs/post-processes prev. Committing prev first
            # would block the CPU on prev's copy_event and leave the GPU idle
            # until dispatch — visible as a gap between forwards in the trace.
            #
            # Eager grammar exception: setup_grammar_step reads each matcher's
            # current state to fill the bitmask. Under the overlap pattern the
            # matcher hasn't been advanced yet by prev's accept_token (commit
            # below), so the fill would use a one-step-stale state and let the
            # model sample a token the matcher then rejects. Capturable
            # grammar dodges this with an in-graph hostfunc that advances
            # before fill; eager has no equivalent, so we commit prev first
            # whenever this batch carries grammars. Costs the dispatch/commit
            # overlap for grammar batches but is correct.
            request_changes = []
            curr_has_grammar = grammar_inputs is not None
            eager_grammar_needs_advance = (
                curr_has_grammar
                and prev_results is not None
                and self.model_executor.eager_grammar_buffers is not None
            )
            if eager_grammar_needs_advance:
                request_changes.extend(
                    self._commit_forward_results(prev_forward_op, prev_results)
                )
                prev_results = None
                prev_forward_op = None

            curr_results = None
            if forward_op is not None:
                self._mark_stats_scheduled(forward_op)
                curr_results, _ = self._dispatch_forward(
                    forward_op,
                    sampling_params_list,
                    execution_plan,
                    dp_metadata=dp_metadata,
                    stats=stats,
                    grammar_inputs=grammar_inputs,
                )

            # ---- post-process previous step (overlapped with current forward) ----
            if prev_results is not None:
                request_changes.extend(
                    self._commit_forward_results(prev_forward_op, prev_results)
                )

            # ---- collect KV transfer events ----
            if self.kv_transfer is not None:
                kv_transfer_events = self.kv_transfer.generate_events()
                request_changes.extend(
                    self._process_kv_transfer_events(kv_transfer_events)
                )

            if request_changes:
                advance_forward(self.scheduler, request_changes)
                self._publish_scheduler_kv_events()

            # Resolve a deferred abort/wait pause reply once in-flight work drains.
            self._pause.maybe_finish_drain(self.scheduler)

            self._record_scheduler_iteration_metrics(stats, num_iter_tokens)

            prev_results = curr_results
            prev_forward_op = forward_op


def run_event_loop(
    server_args: ServerArgs,
    port_args: PortArgs,
    pipe_writer,
):
    mapping = server_args.mapping
    gpu_id = mapping.rank % mapping.nprocs_per_node + server_args.base_gpu_id
    attn_tp_rank = mapping.attn.tp_rank
    dp_rank = mapping.attn.dp_rank
    global_rank = mapping.rank

    setproctitle.setproctitle(f"tokenspeed::scheduler_{dp_rank}")
    faulthandler.enable()
    parent_process = psutil.Process().parent()
    register_usr_signal()

    prefix = f" ATTN TP RANK {attn_tp_rank}"
    configure_logger(server_args, prefix=prefix)

    try:
        if server_args.disaggregation_mode == "encode":
            # The encode role is LM-free; run the lightweight vision-tower loop
            # instead of building the full EventLoop (KV/LM scheduler).
            from tokenspeed.runtime.pd.epd.encode_loop import (
                run_encode_loop,
            )

            run_encode_loop(server_args, port_args, pipe_writer, gpu_id, global_rank)
            return

        event_loop = EventLoop(
            server_args,
            port_args,
            gpu_id,
            attn_tp_rank,
            dp_rank,
            global_rank,
        )
        pipe_writer.send(
            {
                "status": "ready",
                "max_total_num_tokens": event_loop.max_total_num_tokens,
                "max_req_input_len": event_loop.max_req_input_len,
                "max_num_seqs": server_args.max_num_seqs,
                "chunked_prefill_size": server_args.chunked_prefill_size,
                "max_model_len": event_loop.model_config.context_len,
            }
        )

        if event_loop.has_dp:
            # All DP schedulers must finish initialization before any rank enters
            # the loop and starts the first DP metadata collective.
            dist.barrier(group=event_loop.world_cpu_group)

        if event_loop.use_overlap_schedule:
            event_loop.event_loop_overlap()
        else:
            event_loop.event_loop()

    except Exception:
        traceback = get_exception_traceback()
        logger.error("Scheduler hit an exception: %s", traceback)
        parent_process.send_signal(signal.SIGUSR1)