from __future__ import annotations

# Copyright 2023-2024 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

"""
Mixin classes and utils for multi-http-worker mode
This file uses multiple processes to handle requests and tokenization, reducing the overhead of python and http server.
"""

import asyncio
import logging
import multiprocessing as multiprocessing
import os
import pickle
import signal
import sys
import threading
import zlib
from multiprocessing import shared_memory
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type

import psutil
import setproctitle
import zmq
import zmq.asyncio

from sglang.srt.disaggregation.utils import DisaggregationMode, TransferBackend
from sglang.srt.managers.disagg_service import start_disagg_service
from sglang.srt.managers.io_struct import (
    BaseBatchReq,
    BaseReq,
    BatchEmbeddingOutput,
    BatchStrOutput,
    BatchTokenIDOutput,
    ContinueGenerationReqInput,
    FreezeGCReq,
    PauseContinueBroadcastReq,
    PauseGenerationReqInput,
    TokenizerWorkerRegistrationReq,
    async_sock_recv,
    async_sock_send,
    sock_recv,
    sock_send,
    unwrap_from_pickle,
    wrap_as_pickle,
)
from sglang.srt.managers.load_snapshot import (
    create_load_snapshot_reader,
    zmq_reader_owner,
)
from sglang.srt.managers.tokenizer_manager import TokenizerManager
from sglang.srt.server_args import PortArgs, ServerArgs
from sglang.srt.utils import (
    configure_logger,
    kill_itself_when_parent_died,
    kill_process_tree,
)
from sglang.srt.utils.network import get_zmq_socket
from sglang.utils import get_exception_traceback

if TYPE_CHECKING:
    from sglang.srt.managers.detokenizer_manager import DetokenizerManager

logger = logging.getLogger(__name__)


class SocketMapping:
    def __init__(self):
        self._zmq_context = zmq.Context()
        self._mapping: Dict[str, zmq.Socket] = {}

    def clear_all_sockets(self):
        for socket in self._mapping.values():
            socket.close()
        self._mapping.clear()

    def _register_ipc_mapping(self, ipc_name: str, is_tokenizer: bool):
        type_str = "tokenizer" if is_tokenizer else "detokenizer"
        if ipc_name in self._mapping:
            logger.warning(f"{type_str} already registered {ipc_name=}, skipping...")
            return
        logger.info(f"Registering {type_str} {ipc_name=} in SocketMapping...")
        socket = get_zmq_socket(self._zmq_context, zmq.PUSH, ipc_name, False)
        self._mapping[ipc_name] = socket

    def send_output(self, ipc_name: str, output: Any, is_tokenizer: bool = False):
        if ipc_name is None:
            # Some unhandled cases
            logger.warning(f"IPC name is None, output type={type(output)}, skipping...")
            return

        if ipc_name not in self._mapping:
            self._register_ipc_mapping(ipc_name, is_tokenizer=is_tokenizer)
        sock_send(self._mapping[ipc_name], output)


def _extract_field_by_index(
    output: Any, field_name: str, index: int, check_length: bool = True
) -> Any:
    """Extract a field value from output by index, handling None and length checks.

    Args:
        output: The output object containing the field
        field_name: The name of the field to extract
        index: The index to access in the field list
        check_length: If True, check both field existence and length. If False, only check field existence.

    Returns:
        A list containing the field value at index, or None if not available.
    """
    field = getattr(output, field_name, None)
    if field is None:
        return None

    should_wrap_result = field_name in ("customized_info", "time_stats")
    if should_wrap_result:
        field = unwrap_from_pickle(field)
        if field is None:
            return None

    if isinstance(field, dict):
        new_field = {}
        for k, v in field.items():
            if len(v) > index:
                new_field[k] = [v[index]] if should_wrap_result else v[index]
            else:
                new_field[k] = [None] if should_wrap_result else None
        if should_wrap_result:
            return wrap_as_pickle(new_field) if new_field else None
        return new_field

    if check_length:
        if len(field) <= index:
            return None

    new_field = [field[index]]
    return wrap_as_pickle(new_field) if should_wrap_result else new_field


def _handle_output_by_index(output, i):
    """NOTE: A maintainable method is better here."""
    if isinstance(output, BatchTokenIDOutput):
        new_output = BatchTokenIDOutput(
            rids=[output.rids[i]],
            spec_verify_ct=_extract_field_by_index(output, "spec_verify_ct", i),
            spec_num_correct_drafts=_extract_field_by_index(
                output, "spec_num_correct_drafts", i
            ),
            spec_correct_drafts_histogram=_extract_field_by_index(
                output, "spec_correct_drafts_histogram", i
            ),
            spec_num_block_accept_tokens=_extract_field_by_index(
                output, "spec_num_block_accept_tokens", i
            ),
            spec_num_cap_tokens=_extract_field_by_index(
                output, "spec_num_cap_tokens", i
            ),
            spec_cap_lens_histogram=_extract_field_by_index(
                output, "spec_cap_lens_histogram", i
            ),
            time_stats=_extract_field_by_index(output, "time_stats", i),
            finished_reasons=_extract_field_by_index(output, "finished_reasons", i),
            decoded_texts=_extract_field_by_index(output, "decoded_texts", i),
            decode_ids=_extract_field_by_index(output, "decode_ids", i),
            read_offsets=_extract_field_by_index(output, "read_offsets", i),
            output_ids=_extract_field_by_index(output, "output_ids", i),
            skip_special_tokens=_extract_field_by_index(
                output, "skip_special_tokens", i
            ),
            spaces_between_special_tokens=_extract_field_by_index(
                output, "spaces_between_special_tokens", i
            ),
            no_stop_trim=_extract_field_by_index(output, "no_stop_trim", i),
            prompt_tokens=_extract_field_by_index(output, "prompt_tokens", i),
            completion_tokens=_extract_field_by_index(output, "completion_tokens", i),
            reasoning_tokens=_extract_field_by_index(output, "reasoning_tokens", i),
            cached_tokens=_extract_field_by_index(output, "cached_tokens", i),
            cached_tokens_details=_extract_field_by_index(
                output, "cached_tokens_details", i
            ),
            image_tokens=_extract_field_by_index(output, "image_tokens", i),
            audio_tokens=_extract_field_by_index(output, "audio_tokens", i),
            video_tokens=_extract_field_by_index(output, "video_tokens", i),
            input_token_logprobs_val=_extract_field_by_index(
                output, "input_token_logprobs_val", i, check_length=False
            ),
            input_token_logprobs_idx=_extract_field_by_index(
                output, "input_token_logprobs_idx", i, check_length=False
            ),
            output_token_logprobs_val=_extract_field_by_index(
                output, "output_token_logprobs_val", i, check_length=False
            ),
            output_token_logprobs_idx=_extract_field_by_index(
                output, "output_token_logprobs_idx", i, check_length=False
            ),
            input_top_logprobs_val=_extract_field_by_index(
                output, "input_top_logprobs_val", i, check_length=False
            ),
            input_top_logprobs_idx=_extract_field_by_index(
                output, "input_top_logprobs_idx", i, check_length=False
            ),
            output_top_logprobs_val=_extract_field_by_index(
                output, "output_top_logprobs_val", i, check_length=False
            ),
            output_top_logprobs_idx=_extract_field_by_index(
                output, "output_top_logprobs_idx", i, check_length=False
            ),
            input_token_ids_logprobs_val=_extract_field_by_index(
                output, "input_token_ids_logprobs_val", i, check_length=False
            ),
            input_token_ids_logprobs_idx=_extract_field_by_index(
                output, "input_token_ids_logprobs_idx", i, check_length=False
            ),
            output_token_ids_logprobs_val=_extract_field_by_index(
                output, "output_token_ids_logprobs_val", i, check_length=False
            ),
            output_token_ids_logprobs_idx=_extract_field_by_index(
                output, "output_token_ids_logprobs_idx", i, check_length=False
            ),
            output_token_entropy_val=_extract_field_by_index(
                output, "output_token_entropy_val", i, check_length=False
            ),
            output_hidden_states=_extract_field_by_index(
                output, "output_hidden_states", i, check_length=False
            ),
            routed_experts=_extract_field_by_index(
                output, "routed_experts", i, check_length=False
            ),
            indexer_topk=_extract_field_by_index(
                output, "indexer_topk", i, check_length=False
            ),
            retraction_counts=_extract_field_by_index(output, "retraction_counts", i),
            placeholder_tokens_idx=None,
            placeholder_tokens_val=None,
            token_steps=_extract_field_by_index(
                output, "token_steps", i, check_length=False
            ),
            customized_info=_extract_field_by_index(
                output, "customized_info", i, check_length=False
            ),
            dp_ranks=_extract_field_by_index(output, "dp_ranks", i, check_length=False),
        )
    elif isinstance(output, BatchEmbeddingOutput):
        new_output = BatchEmbeddingOutput(
            rids=[output.rids[i]],
            finished_reasons=_extract_field_by_index(output, "finished_reasons", i),
            embeddings=_extract_field_by_index(output, "embeddings", i),
            prompt_tokens=_extract_field_by_index(output, "prompt_tokens", i),
            cached_tokens=_extract_field_by_index(output, "cached_tokens", i),
            placeholder_tokens_idx=None,
            placeholder_tokens_val=None,
        )
    elif isinstance(output, BatchStrOutput):
        new_output = BatchStrOutput(
            rids=[output.rids[i]],
            spec_verify_ct=_extract_field_by_index(output, "spec_verify_ct", i),
            spec_num_correct_drafts=_extract_field_by_index(
                output, "spec_num_correct_drafts", i
            ),
            spec_correct_drafts_histogram=_extract_field_by_index(
                output, "spec_correct_drafts_histogram", i
            ),
            spec_num_block_accept_tokens=_extract_field_by_index(
                output, "spec_num_block_accept_tokens", i
            ),
            spec_num_cap_tokens=_extract_field_by_index(
                output, "spec_num_cap_tokens", i
            ),
            spec_cap_lens_histogram=_extract_field_by_index(
                output, "spec_cap_lens_histogram", i
            ),
            time_stats=_extract_field_by_index(output, "time_stats", i),
            finished_reasons=_extract_field_by_index(output, "finished_reasons", i),
            output_strs=_extract_field_by_index(output, "output_strs", i),
            output_ids=_extract_field_by_index(output, "output_ids", i),
            prompt_tokens=_extract_field_by_index(output, "prompt_tokens", i),
            completion_tokens=_extract_field_by_index(output, "completion_tokens", i),
            reasoning_tokens=_extract_field_by_index(output, "reasoning_tokens", i),
            cached_tokens=_extract_field_by_index(output, "cached_tokens", i),
            cached_tokens_details=_extract_field_by_index(
                output, "cached_tokens_details", i
            ),
            image_tokens=_extract_field_by_index(output, "image_tokens", i),
            audio_tokens=_extract_field_by_index(output, "audio_tokens", i),
            video_tokens=_extract_field_by_index(output, "video_tokens", i),
            input_token_logprobs_val=_extract_field_by_index(
                output, "input_token_logprobs_val", i, check_length=False
            ),
            input_token_logprobs_idx=_extract_field_by_index(
                output, "input_token_logprobs_idx", i, check_length=False
            ),
            output_token_logprobs_val=_extract_field_by_index(
                output, "output_token_logprobs_val", i, check_length=False
            ),
            output_token_logprobs_idx=_extract_field_by_index(
                output, "output_token_logprobs_idx", i, check_length=False
            ),
            input_top_logprobs_val=_extract_field_by_index(
                output, "input_top_logprobs_val", i, check_length=False
            ),
            input_top_logprobs_idx=_extract_field_by_index(
                output, "input_top_logprobs_idx", i, check_length=False
            ),
            output_top_logprobs_val=_extract_field_by_index(
                output, "output_top_logprobs_val", i, check_length=False
            ),
            output_top_logprobs_idx=_extract_field_by_index(
                output, "output_top_logprobs_idx", i, check_length=False
            ),
            input_token_ids_logprobs_val=_extract_field_by_index(
                output, "input_token_ids_logprobs_val", i, check_length=False
            ),
            input_token_ids_logprobs_idx=_extract_field_by_index(
                output, "input_token_ids_logprobs_idx", i, check_length=False
            ),
            output_token_ids_logprobs_val=_extract_field_by_index(
                output, "output_token_ids_logprobs_val", i, check_length=False
            ),
            output_token_ids_logprobs_idx=_extract_field_by_index(
                output, "output_token_ids_logprobs_idx", i, check_length=False
            ),
            output_token_entropy_val=_extract_field_by_index(
                output, "output_token_entropy_val", i, check_length=False
            ),
            output_hidden_states=_extract_field_by_index(
                output, "output_hidden_states", i, check_length=False
            ),
            routed_experts=_extract_field_by_index(
                output, "routed_experts", i, check_length=False
            ),
            indexer_topk=_extract_field_by_index(
                output, "indexer_topk", i, check_length=False
            ),
            customized_info=_extract_field_by_index(
                output, "customized_info", i, check_length=False
            ),
            dp_ranks=_extract_field_by_index(output, "dp_ranks", i, check_length=False),
            placeholder_tokens_idx=None,
            placeholder_tokens_val=None,
            retraction_counts=_extract_field_by_index(output, "retraction_counts", i),
            token_steps=_extract_field_by_index(
                output, "token_steps", i, check_length=False
            ),
        )
    else:
        new_output = output
    return new_output


class MultiHttpWorkerDetokenizerMixin:
    """Mixin class for DetokenizerManager"""

    def maybe_clear_socket_mapping(self: DetokenizerManager):
        if hasattr(self, "socket_mapping"):
            self.socket_mapping.clear_all_sockets()

    def multi_http_worker_event_loop(self: DetokenizerManager):
        """The event loop that handles requests, for multi multi-http-worker mode"""
        self.socket_mapping = SocketMapping()
        while True:
            recv_obj = sock_recv(self.recv_from_scheduler)
            output = self._request_dispatcher(recv_obj)
            if output is None:
                continue

            # Fan out the output back to the originating tokenizer worker(s).
            # In multi-detokenizer mode the upstream MultiDetokenizerRouter may
            # forward either batched or single requests, so handle both shapes.
            if isinstance(recv_obj, BaseBatchReq):
                for i, ipc_name in enumerate(recv_obj.http_worker_ipcs):
                    new_output = _handle_output_by_index(output, i)
                    self.socket_mapping.send_output(
                        ipc_name, new_output, is_tokenizer=True
                    )
            elif isinstance(recv_obj, BaseReq):
                self.socket_mapping.send_output(
                    recv_obj.http_worker_ipc, output, is_tokenizer=True
                )
            else:
                raise ValueError(
                    f"multi_http_worker_event_loop got unexpected req type {type(recv_obj)}"
                )


class MultiTokenizerRouter:
    """A router between tokenizer managers and the scheduler/detokenizer manager.

    Forward: tokenizer managers → router → scheduler.
    Backward: detokenizer manager → router → tokenizer managers.
    Also broadcasts pause/continue to all tokenizer managers for consistent is_pause state.
    """

    def __init__(
        self,
        server_args: ServerArgs,
        port_args: PortArgs,
    ):
        self.server_args = server_args
        context = zmq.asyncio.Context(3)
        self.recv_from_detokenizer = get_zmq_socket(
            context, zmq.PULL, port_args.tokenizer_ipc_name, True
        )
        self.send_to_scheduler = get_zmq_socket(
            context, zmq.PUSH, port_args.scheduler_input_ipc_name, True
        )
        self.receive_from_worker = get_zmq_socket(
            context, zmq.PULL, port_args.tokenizer_worker_ipc_name, True
        )
        self._loop = asyncio.new_event_loop()
        self._thread = threading.Thread(target=self._run_loop, daemon=True)
        self._thread.start()
        self._task = asyncio.run_coroutine_threadsafe(
            self.router_worker_obj(), self._loop
        )
        self._handle_task = asyncio.run_coroutine_threadsafe(
            print_exception_wrapper(self.handle_loop), self._loop
        )

        # In multi-tokenizer mode the N TokenizerWorker processes cannot each
        # bind the zmq PULL socket used for load snapshots, so the single
        # MultiTokenizerRouter process owns it (zmq -> SHM) and the workers
        # read SHM only. Drain it event-driven via the socket's fd instead of
        # polling on a timer.
        self.load_snapshot_reader = None
        if zmq_reader_owner(server_args, "MultiTokenizerRouter"):
            self.load_snapshot_reader = create_load_snapshot_reader(
                server_args, port_args, caller="MultiTokenizerRouter"
            )
            self._loop.call_soon_threadsafe(self._register_load_snapshot_reader)

        self.disaggregation_bootstrap_server = start_disagg_service(self.server_args)

        # Worker IPC names for pause/continue broadcasting
        self.all_worker_ipcs: set[str] = set()
        # Shared socket mapping (both coroutines run on self._loop, so safe)
        self.socket_mapping = SocketMapping()

    def _run_loop(self):
        self._loop.run_forever()

    def _register_load_snapshot_reader(self):
        """Drain zmq load snapshots into SHM whenever the PULL socket is readable.

        zmq exposes an edge-triggered fd; ``poll()`` drains it until empty, which
        also re-arms the fd, so TokenizerWorkers reading SHM stay up to date
        without any timer.
        """
        assert self.load_snapshot_reader is not None
        self._loop.add_reader(
            self.load_snapshot_reader.fileno(), self.load_snapshot_reader.poll
        )
        # Drain anything already queued before the fd was registered.
        self.load_snapshot_reader.poll()

    async def router_worker_obj(self):
        """Forward path: workers → scheduler, with pause/continue broadcast."""
        while True:
            recv_obj = await async_sock_recv(self.receive_from_worker)

            if isinstance(recv_obj, TokenizerWorkerRegistrationReq):
                if recv_obj.worker_ipc_name not in self.all_worker_ipcs:
                    self.all_worker_ipcs.add(recv_obj.worker_ipc_name)
                    logger.info(
                        f"Router registered worker IPC: {recv_obj.worker_ipc_name} "
                        f"(total: {len(self.all_worker_ipcs)})"
                    )
                continue

            if isinstance(
                recv_obj, (PauseGenerationReqInput, ContinueGenerationReqInput)
            ):
                # Broadcast to ALL workers so every worker's is_pause is set
                is_pause = isinstance(recv_obj, PauseGenerationReqInput)
                broadcast = PauseContinueBroadcastReq(is_pause=is_pause)
                for ipc_name in self.all_worker_ipcs:
                    self.socket_mapping.send_output(ipc_name, broadcast)
                # Forward to scheduler rank 0 (it broadcasts to all TP/PP/DP
                # ranks internally). Skip for abort mode which drains via polling.
                if not (
                    isinstance(recv_obj, PauseGenerationReqInput)
                    and recv_obj.mode == "abort"
                ):
                    await async_sock_send(self.send_to_scheduler, recv_obj)
                continue

            await async_sock_send(self.send_to_scheduler, recv_obj)

    async def handle_loop(self):
        """Backward path: detokenizer → route results to correct worker."""
        while True:
            recv_obj = await async_sock_recv(self.recv_from_detokenizer)
            await self._distribute_result_to_workers(recv_obj)

    async def _distribute_result_to_workers(self, recv_obj):
        if isinstance(recv_obj, BaseReq):
            ipc_names = [recv_obj.http_worker_ipc]
        elif isinstance(recv_obj, BaseBatchReq):
            ipc_names = recv_obj.http_worker_ipcs
        else:
            raise ValueError(f"Unknown recv_obj type: {type(recv_obj)}")

        for i, ipc_name in enumerate(ipc_names):
            new_recv_obj = _handle_output_by_index(recv_obj, i)
            self.socket_mapping.send_output(ipc_name, new_recv_obj)


class MultiDetokenizerRouter:
    """Route scheduler outputs to one of N DetokenizerManager workers.

    Each request is pinned to a worker by hashing its ``http_worker_ipc`` with
    ``zlib.crc32`` (deterministic across runs), so all outputs of the same rid
    always land on the same detokenizer and ``decode_status`` stays consistent.
    """

    def __init__(self, ipc_name_list: List[str], port_args: PortArgs):
        self.ipc_name_list = ipc_name_list
        self.num_workers = len(ipc_name_list)
        self.socket_mapping = SocketMapping()
        context = zmq.Context(2)
        self.recv_from_scheduler = get_zmq_socket(
            context, zmq.PULL, port_args.detokenizer_ipc_name, True
        )

    def _pick(self, key: str) -> str:
        return self.ipc_name_list[zlib.crc32(key.encode()) % self.num_workers]

    def _send(self, ipc_name: str, obj: Any) -> None:
        self.socket_mapping.send_output(ipc_name, obj, is_tokenizer=False)

    def event_loop(self):
        while True:
            recv_obj = sock_recv(self.recv_from_scheduler)

            # FreezeGCReq must freeze every detokenizer process.
            if isinstance(recv_obj, FreezeGCReq):
                for ipc in self.ipc_name_list:
                    self._send(ipc, recv_obj)
                continue

            # Single request: route by its own http_worker_ipc.
            if isinstance(recv_obj, BaseReq):
                assert (
                    recv_obj.http_worker_ipc is not None
                ), f"Single req {recv_obj.rid=} missing http_worker_ipc"
                self._send(self._pick(recv_obj.http_worker_ipc), recv_obj)
                continue

            # Batch request.
            if isinstance(recv_obj, BaseBatchReq):
                # Idle/no-op batch (rids=[]): broadcast to all detokenizers
                if not recv_obj.rids:
                    for ipc in self.ipc_name_list:
                        self._send(ipc, recv_obj)
                    continue

                ipcs = recv_obj.http_worker_ipcs
                assert (
                    ipcs is not None
                    and len(ipcs) == len(recv_obj.rids)
                    and all(x is not None for x in ipcs)
                ), f"Batch req {recv_obj.rids=} has invalid http_worker_ipcs"

                # Split per-item and route each by its own ipc.
                for i, ipc_key in enumerate(ipcs):
                    one = _handle_output_by_index(recv_obj, i)
                    if one is recv_obj:
                        raise TypeError(f"Cannot split {type(recv_obj)}")
                    one.http_worker_ipcs = [ipc_key]
                    self._send(self._pick(ipc_key), one)
                continue

            raise ValueError(
                f"MultiDetokenizerRouter got unsupported type {type(recv_obj)}"
            )


def run_multi_detokenizer_router_process(
    ipc_name_list: List[str],
    server_args: ServerArgs,
    port_args: PortArgs,
):
    kill_itself_when_parent_died()
    setproctitle.setproctitle("sglang::detokenizer_router")
    configure_logger(server_args)
    parent_process = psutil.Process().parent()

    router = None
    try:
        router = MultiDetokenizerRouter(ipc_name_list, port_args)
        router.event_loop()
    except Exception:
        traceback = get_exception_traceback()
        logger.error(f"MultiDetokenizerRouter hit an exception: {traceback}")
        if router is not None:
            router.socket_mapping.clear_all_sockets()
        parent_process.send_signal(signal.SIGQUIT)


class TokenizerWorker(TokenizerManager):
    """Tokenizer Worker in multi-http-worker mode"""

    def __init__(
        self,
        server_args: ServerArgs,
        port_args: PortArgs,
    ):
        setproctitle.setproctitle(f"sglang::tokenizer_worker:{os.getpid()}")
        import torch

        torch.set_num_threads(1)
        # prevent init prefill bootstrapserver again
        disaggregation_mode = server_args.disaggregation_mode
        server_args.override(
            "tokenizer_worker.suppress_bootstrap", disaggregation_mode="null"
        )
        super().__init__(server_args, port_args)

        self.worker_id = os.getpid()
        self.tokenizer_ipc_name = port_args.tokenizer_ipc_name

        # For PD disaggregtion
        self.server_args.override(
            "tokenizer_worker.restore_disaggregation_mode",
            disaggregation_mode=disaggregation_mode,
        )
        self.disaggregation_mode = DisaggregationMode(
            self.server_args.disaggregation_mode
        )
        self.disaggregation_transfer_backend = TransferBackend(
            self.server_args.disaggregation_transfer_backend
        )

        # Register this worker with the router for pause/continue broadcasting
        reg = TokenizerWorkerRegistrationReq(worker_ipc_name=self.tokenizer_ipc_name)
        self._dispatch_to_scheduler(reg)

        # Future for awaiting pause/continue broadcast confirmation
        self._pause_continue_future: Optional[asyncio.Future] = None

        # Register PauseContinueBroadcastReq in the result dispatcher so
        # handle_loop routes it to _handle_pause_continue_broadcast
        from sglang.utils import TypeBasedDispatcher

        self._result_dispatcher += TypeBasedDispatcher(
            [(PauseContinueBroadcastReq, self._handle_pause_continue_broadcast)]
        )

    async def pause_generation(self, obj: PauseGenerationReqInput):
        loop = asyncio.get_event_loop()
        self._pause_continue_future = loop.create_future()
        # Send to router which will broadcast to all workers
        # (router also handles forwarding to scheduler for non-abort modes)
        self._dispatch_to_scheduler(obj)
        await self._pause_continue_future

        if obj.mode == "abort":
            # Abort polling: only the originator checks its own lock state
            while True:
                self.abort_request(abort_all=True)
                is_locked = await self.model_update_lock.is_locked()
                if not is_locked:
                    break
                await asyncio.sleep(1.0)

    async def continue_generation(self, obj: ContinueGenerationReqInput):
        loop = asyncio.get_event_loop()
        self._pause_continue_future = loop.create_future()
        self._dispatch_to_scheduler(obj)
        await self._pause_continue_future

    def _handle_pause_continue_broadcast(self, obj: PauseContinueBroadcastReq):
        """Called from handle_loop when a broadcast arrives from the router."""
        loop = asyncio.get_event_loop()
        loop.create_task(self._apply_pause_continue_broadcast(obj))

    async def _apply_pause_continue_broadcast(self, obj: PauseContinueBroadcastReq):
        """Apply pause/continue state under the condition lock."""
        async with self.is_pause_cond:
            if obj.is_pause:
                self.is_pause = True
            else:
                self.is_pause = False
                self.is_pause_cond.notify_all()

        # Resolve the pending future if this worker initiated the pause/continue
        if self._pause_continue_future and not self._pause_continue_future.done():
            self._pause_continue_future.set_result(True)
            self._pause_continue_future = None


def get_tokenizer_worker_class(server_args: ServerArgs) -> Type[TokenizerWorker]:
    worker_class = server_args.get_tokenizer_worker_class()
    if not isinstance(worker_class, type) or not issubclass(
        worker_class, TokenizerWorker
    ):
        raise TypeError(
            "ServerArgs.get_tokenizer_worker_class() must return a TokenizerWorker "
            f"subclass, got {worker_class!r}"
        )

    return worker_class


async def print_exception_wrapper(func):
    """
    Sometimes an asyncio function does not print exception.
    We do another wrapper to handle the exception.
    """
    try:
        await func()
    except Exception:
        traceback = get_exception_traceback()
        logger.error(f"MultiTokenizerRouter hit an exception: {traceback}")
        if hasattr(func, "__self__") and isinstance(
            func.__self__, MultiTokenizerRouter
        ):
            func.__self__.dump_requests_before_crash()
        kill_process_tree(os.getpid(), include_parent=True)
        sys.exit(1)


def get_main_process_id() -> int:
    """Get the main process ID."""
    return multiprocessing.current_process()._parent_pid


def write_to_shared_memory(obj, name: str) -> shared_memory.SharedMemory:
    """Write data to shared memory"""
    serialized = pickle.dumps(obj)
    size = len(serialized)
    try:
        # Try to open existing shared memory
        shm = shared_memory.SharedMemory(name=name)
        # If size is insufficient, close and recreate
        if shm.size < size:
            shm.close()
            shm.unlink()
            shm = shared_memory.SharedMemory(create=True, size=size, name=name)
    except FileNotFoundError:
        # If not present, create new shared memory
        shm = shared_memory.SharedMemory(create=True, size=size, name=name)

    shm.buf[:size] = serialized
    return shm


def read_from_shared_memory(name: str) -> Any:
    """Read data from shared memory"""
    try:
        shm = shared_memory.SharedMemory(name=name)
        data = pickle.loads(bytes(shm.buf))
        shm.close()
        return data
    except FileNotFoundError:
        raise FileNotFoundError(f"Shared memory {name} not found")


def write_data_for_multi_tokenizer(
    port_args: PortArgs, server_args: ServerArgs, scheduler_info: Dict
):
    """Write args information to share memory for multi-tokenizer"""
    # get main process ID
    main_pid = get_main_process_id()
    current_pid = os.getpid()
    logger.info(f"main process ID: {main_pid}, current process ID: {current_pid}")
    args = (port_args, server_args, scheduler_info)
    args_shm = write_to_shared_memory(args, f"multi_tokenizer_args_{current_pid}")
    args_shm.close()

    return args_shm