sglang 0.5.1.post2__py3-none-any.whl → 0.5.2rc0__py3-none-any.whl

sglang/srt/managers/multi_tokenizer_mixin.py

@@ -0,0 +1,591 @@
+# 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.
+# ==============================================================================
+"""MultiTokenizerMixin is a class that provides nesscary methods for MultiTokenizerManager and DetokenizerManager."""
+import asyncio
+import dataclasses
+import json
+import logging
+import multiprocessing as multiprocessing
+import os
+import sys
+import threading
+from multiprocessing import shared_memory
+from typing import Dict
+
+import zmq
+import zmq.asyncio
+
+from sglang.srt.disaggregation.utils import DisaggregationMode, TransferBackend
+from sglang.srt.managers.io_struct import (
+    BatchEmbeddingOut,
+    BatchMultimodalOut,
+    BatchStrOut,
+    BatchTokenIDOut,
+    MultiTokenizerRegisterReq,
+    MultiTokenizerWarpper,
+)
+from sglang.srt.managers.tokenizer_manager import TokenizerManager, _Communicator
+from sglang.srt.server_args import PortArgs, ServerArgs
+from sglang.srt.utils import (
+    get_worker_ids_from_req_rids,
+    get_zmq_socket,
+    kill_process_tree,
+)
+from sglang.utils import get_exception_traceback
+
+logger = logging.getLogger(__name__)
+
+
+class MultiTokenizerMixin:
+    """Mixin class for MultiTokenizerManager and DetokenizerManager"""
+
+    def create_sockets_mapping(self):
+        if not hasattr(self, "tokenizer_mapping"):
+            self.tokenizer_mapping = {}
+        # Create ZMQ context if needed
+        if not hasattr(self, "_zmq_context"):
+            self._zmq_context = zmq.Context()
+
+    def init_tokenizer_mapping(
+        self, recv_obj: MultiTokenizerRegisterReq, worker_id: str
+    ):
+        """init tokenizer mapping from register request"""
+        ipc_name = recv_obj.ipc_name
+        worker_id_int = int(worker_id)
+
+        if worker_id_int not in self.tokenizer_mapping:
+            socket = get_zmq_socket(self._zmq_context, zmq.PUSH, ipc_name, False)
+            self.tokenizer_mapping[worker_id_int] = socket
+            self.tokenizer_mapping[worker_id_int].send_pyobj(recv_obj)
+            return True
+        else:
+            return False
+
+    def register_tokenizer_ipc(self, recv_obj, worker_id):
+        if worker_id not in self.tokenizer_mapping:
+            # register the worker if not already done
+            if isinstance(recv_obj, MultiTokenizerRegisterReq):
+                return self.init_tokenizer_mapping(recv_obj, worker_id)
+            else:
+                logger.error(
+                    f"Worker {worker_id} not registered and not found in tokenizer mapping . "
+                    "Please ensure the worker is registered correctly."
+                )
+        return False
+
+    def _handle_output_by_index(self, output, i):
+        """NOTE: A maintainable method is better here."""
+        if isinstance(output, BatchTokenIDOut):
+            new_output = BatchTokenIDOut(
+                rids=[output.rids[i]],
+                finished_reasons=(
+                    [output.finished_reasons[i]]
+                    if len(output.finished_reasons) > i
+                    else None
+                ),
+                decoded_texts=(
+                    [output.decoded_texts[i]] if len(output.decoded_texts) > i else None
+                ),
+                decode_ids=(
+                    [output.decode_ids[i]] if len(output.decode_ids) > i else None
+                ),
+                read_offsets=(
+                    [output.read_offsets[i]] if len(output.read_offsets) > i else None
+                ),
+                output_ids=(
+                    [output.output_ids[i]]
+                    if output.output_ids and len(output.output_ids) > i
+                    else None
+                ),
+                skip_special_tokens=(
+                    [output.skip_special_tokens[i]]
+                    if len(output.skip_special_tokens) > i
+                    else None
+                ),
+                spaces_between_special_tokens=(
+                    [output.spaces_between_special_tokens[i]]
+                    if len(output.spaces_between_special_tokens) > i
+                    else None
+                ),
+                no_stop_trim=(
+                    [output.no_stop_trim[i]] if len(output.no_stop_trim) > i else None
+                ),
+                prompt_tokens=(
+                    [output.prompt_tokens[i]] if len(output.prompt_tokens) > i else None
+                ),
+                completion_tokens=(
+                    [output.completion_tokens[i]]
+                    if len(output.completion_tokens) > i
+                    else None
+                ),
+                cached_tokens=(
+                    [output.cached_tokens[i]] if len(output.cached_tokens) > i else None
+                ),
+                spec_verify_ct=(
+                    [output.spec_verify_ct[i]]
+                    if len(output.spec_verify_ct) > i
+                    else None
+                ),
+                input_token_logprobs_val=(
+                    [output.input_token_logprobs_val[i]]
+                    if output.input_token_logprobs_val
+                    else None
+                ),
+                input_token_logprobs_idx=(
+                    [output.input_token_logprobs_idx[i]]
+                    if output.input_token_logprobs_idx
+                    else None
+                ),
+                output_token_logprobs_val=(
+                    [output.output_token_logprobs_val[i]]
+                    if output.output_token_logprobs_val
+                    else None
+                ),
+                output_token_logprobs_idx=(
+                    [output.output_token_logprobs_idx[i]]
+                    if output.output_token_logprobs_idx
+                    else None
+                ),
+                input_top_logprobs_val=(
+                    [output.input_top_logprobs_val[i]]
+                    if output.input_top_logprobs_val
+                    else None
+                ),
+                input_top_logprobs_idx=(
+                    [output.input_top_logprobs_idx[i]]
+                    if output.input_top_logprobs_idx
+                    else None
+                ),
+                output_top_logprobs_val=(
+                    [output.output_top_logprobs_val[i]]
+                    if output.output_top_logprobs_val
+                    else None
+                ),
+                output_top_logprobs_idx=(
+                    [output.output_top_logprobs_idx[i]]
+                    if output.output_top_logprobs_idx
+                    else None
+                ),
+                input_token_ids_logprobs_val=(
+                    [output.input_token_ids_logprobs_val[i]]
+                    if output.input_token_ids_logprobs_val
+                    else None
+                ),
+                input_token_ids_logprobs_idx=(
+                    [output.input_token_ids_logprobs_idx[i]]
+                    if output.input_token_ids_logprobs_idx
+                    else None
+                ),
+                output_token_ids_logprobs_val=(
+                    [output.output_token_ids_logprobs_val[i]]
+                    if output.output_token_ids_logprobs_val
+                    else None
+                ),
+                output_token_ids_logprobs_idx=(
+                    [output.output_token_ids_logprobs_idx[i]]
+                    if output.output_token_ids_logprobs_idx
+                    else None
+                ),
+                output_hidden_states=(
+                    [output.output_hidden_states[i]]
+                    if output.output_hidden_states
+                    else None
+                ),
+            )
+        elif isinstance(output, BatchEmbeddingOut):
+            new_output = BatchEmbeddingOut(
+                rids=[output.rids[i]],
+                finished_reasons=(
+                    [output.finished_reasons[i]]
+                    if len(output.finished_reasons) > i
+                    else None
+                ),
+                embeddings=(
+                    [output.embeddings[i]] if len(output.embeddings) > i else None
+                ),
+                prompt_tokens=(
+                    [output.prompt_tokens[i]] if len(output.prompt_tokens) > i else None
+                ),
+                cached_tokens=(
+                    [output.cached_tokens[i]] if len(output.cached_tokens) > i else None
+                ),
+            )
+        elif isinstance(output, BatchStrOut):
+            new_output = BatchStrOut(
+                rids=[output.rids[i]],
+                finished_reasons=(
+                    [output.finished_reasons[i]]
+                    if len(output.finished_reasons) > i
+                    else None
+                ),
+                output_strs=(
+                    [output.output_strs[i]] if len(output.output_strs) > i else None
+                ),
+                output_ids=(
+                    [output.output_ids[i]]
+                    if output.output_ids and len(output.output_ids) > i
+                    else None
+                ),
+                prompt_tokens=(
+                    [output.prompt_tokens[i]] if len(output.prompt_tokens) > i else None
+                ),
+                completion_tokens=(
+                    [output.completion_tokens[i]]
+                    if len(output.completion_tokens) > i
+                    else None
+                ),
+                cached_tokens=(
+                    [output.cached_tokens[i]] if len(output.cached_tokens) > i else None
+                ),
+                spec_verify_ct=(
+                    [output.spec_verify_ct[i]]
+                    if len(output.spec_verify_ct) > i
+                    else None
+                ),
+                input_token_logprobs_val=(
+                    [output.input_token_logprobs_val[i]]
+                    if output.input_token_logprobs_val
+                    else None
+                ),
+                input_token_logprobs_idx=(
+                    [output.input_token_logprobs_idx[i]]
+                    if output.input_token_logprobs_idx
+                    else None
+                ),
+                output_token_logprobs_val=(
+                    [output.output_token_logprobs_val[i]]
+                    if output.output_token_logprobs_val
+                    else None
+                ),
+                output_token_logprobs_idx=(
+                    [output.output_token_logprobs_idx[i]]
+                    if output.output_token_logprobs_idx
+                    else None
+                ),
+                input_top_logprobs_val=(
+                    [output.input_top_logprobs_val[i]]
+                    if output.input_top_logprobs_val
+                    else None
+                ),
+                input_top_logprobs_idx=(
+                    [output.input_top_logprobs_idx[i]]
+                    if output.input_top_logprobs_idx
+                    else None
+                ),
+                output_top_logprobs_val=(
+                    [output.output_top_logprobs_val[i]]
+                    if output.output_top_logprobs_val
+                    else None
+                ),
+                output_top_logprobs_idx=(
+                    [output.output_top_logprobs_idx[i]]
+                    if output.output_top_logprobs_idx
+                    else None
+                ),
+                input_token_ids_logprobs_val=(
+                    [output.input_token_ids_logprobs_val[i]]
+                    if output.input_token_ids_logprobs_val
+                    else None
+                ),
+                input_token_ids_logprobs_idx=(
+                    [output.input_token_ids_logprobs_idx[i]]
+                    if output.input_token_ids_logprobs_idx
+                    else None
+                ),
+                output_token_ids_logprobs_val=(
+                    [output.output_token_ids_logprobs_val[i]]
+                    if output.output_token_ids_logprobs_val
+                    else None
+                ),
+                output_token_ids_logprobs_idx=(
+                    [output.output_token_ids_logprobs_idx[i]]
+                    if output.output_token_ids_logprobs_idx
+                    else None
+                ),
+                output_hidden_states=(
+                    [output.output_hidden_states[i]]
+                    if output.output_hidden_states
+                    else None
+                ),
+            )
+        elif isinstance(output, BatchMultimodalOut):
+            new_output = BatchMultimodalOut(
+                rids=[output.rids[i]],
+                finished_reasons=(
+                    [output.finished_reasons[i]]
+                    if len(output.finished_reasons) > i
+                    else None
+                ),
+                outputs=([output.outputs[i]] if len(output.outputs) > i else None),
+                prompt_tokens=(
+                    [output.prompt_tokens[i]] if len(output.prompt_tokens) > i else None
+                ),
+                completion_tokens=(
+                    [output.completion_tokens[i]]
+                    if len(output.completion_tokens) > i
+                    else None
+                ),
+                cached_tokens=(
+                    [output.cached_tokens[i]] if len(output.cached_tokens) > i else None
+                ),
+            )
+        else:
+            new_output = output
+        return new_output
+
+    def clear_tokenizer_mapping(self):
+        if hasattr(self, "tokenizer_mapping"):
+            for socket in self.tokenizer_mapping.values():
+                try:
+                    socket.close()
+                except Exception as e:
+                    logger.warning(f"Failed to close socket: {e}")
+            self.tokenizer_mapping.clear()
+
+
+class MultiTokenizerRouter(TokenizerManager, MultiTokenizerMixin):
+    """A router to receive requests from MultiTokenizerManager"""
+
+    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
+        )
+        # Start handle_loop simultaneously
+        self._handle_task = asyncio.run_coroutine_threadsafe(
+            print_exception_wrapper(self.handle_loop), self._loop
+        )
+        self.init_disaggregation()
+
+    def _run_loop(self):
+        self._loop.run_forever()
+
+    async def router_worker_obj(self):
+        while True:
+            recv_obj = await self.receive_from_worker.recv_pyobj()
+            await self.send_to_scheduler.send_pyobj(recv_obj)
+
+    async def handle_loop(self):
+        # special reqs will recv from scheduler, need to route to right worker
+        self.create_sockets_mapping()
+        while True:
+            recv_obj = await self.recv_from_detokenizer.recv_pyobj()
+            await self._distribute_result_to_workers(recv_obj)
+
+    async def _distribute_result_to_workers(self, recv_obj):
+        """Distribute result to corresponding workers based on rid"""
+        if isinstance(recv_obj, MultiTokenizerWarpper):
+            worker_ids = [recv_obj.worker_id]
+            recv_obj = recv_obj.obj
+        else:
+            worker_ids = get_worker_ids_from_req_rids(recv_obj.rids)
+
+        if len(worker_ids) == 0:
+            logger.error(f"Cannot find worker_id from rids {recv_obj.rids}")
+            return
+
+        # Distribute result to each worker
+        for i, worker_id in enumerate(worker_ids):
+            if isinstance(recv_obj, MultiTokenizerRegisterReq):
+                if self.register_tokenizer_ipc(recv_obj, worker_id):
+                    logger.info(
+                        f"MultiTokenizerRouter Created ZMQ socket for worker {worker_id}"
+                    )
+                continue
+            else:
+                if worker_id not in self.tokenizer_mapping:
+                    logger.error(
+                        f"Tokenizer Worker ID {worker_id} not registered. Check if the server Process {worker_id} is alive"
+                    )
+                    continue
+                new_recv_obj = self._handle_output_by_index(recv_obj, i)
+                self.tokenizer_mapping[worker_id].send_pyobj(new_recv_obj)
+
+
+class MultiTokenizerManager(TokenizerManager, MultiTokenizerMixin):
+    """Multi Process Tokenizer Manager that tokenizes the text."""
+
+    def __init__(
+        self,
+        server_args: ServerArgs,
+        port_args: PortArgs,
+    ):
+        # prevent init prefill bootstrapserver again
+        disaggregation_mode = server_args.disaggregation_mode
+        server_args.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.disaggregation_mode = disaggregation_mode
+        self.disaggregation_mode = DisaggregationMode(
+            self.server_args.disaggregation_mode
+        )
+        self.disaggregation_transfer_backend = TransferBackend(
+            self.server_args.disaggregation_transfer_backend
+        )
+        # Communicator
+        self.register_multi_tokenizer_communicator = _Communicator(
+            self.send_to_scheduler, 2
+        )
+        self._result_dispatcher._mapping.append(
+            (
+                MultiTokenizerRegisterReq,
+                self.register_multi_tokenizer_communicator.handle_recv,
+            )
+        )
+
+    async def register_to_main_tokenizer_manager(self):
+        """Register this worker to the main TokenizerManager"""
+        # create a handle loop to receive messages from the main TokenizerManager
+        self.auto_create_handle_loop()
+        req = MultiTokenizerRegisterReq(rids=[f"{self.worker_id}_register"])
+        req.ipc_name = self.tokenizer_ipc_name
+        _Communicator.enable_multi_tokenizer = True
+        await self.register_multi_tokenizer_communicator(req)
+
+
+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 serialize_port_args(port_args: PortArgs) -> dict:
+    """Serialize PortArgs into a shareable dictionary"""
+    return {
+        "tokenizer_ipc_name": port_args.tokenizer_ipc_name,
+        "scheduler_input_ipc_name": port_args.scheduler_input_ipc_name,
+        "detokenizer_ipc_name": port_args.detokenizer_ipc_name,
+        "nccl_port": port_args.nccl_port,
+        "rpc_ipc_name": port_args.rpc_ipc_name,
+        "metrics_ipc_name": port_args.metrics_ipc_name,
+        "tokenizer_worker_ipc_name": port_args.tokenizer_worker_ipc_name,
+    }
+
+
+def deserialize_data(port_args: dict, server_args: dict):
+    """Deserialize data from shared dictionaries"""
+    return PortArgs(**port_args), ServerArgs(**server_args)
+
+
+def serialize_server_args(server_args: ServerArgs) -> dict:
+    """Serialize ServerArgs into a shareable dictionary"""
+    return dataclasses.asdict(server_args)
+
+
+def serialize_scheduler_info(scheduler_info: Dict) -> dict:
+    """Serialize scheduler_info into a shareable dictionary"""
+    return scheduler_info
+
+
+def deserialize_scheduler_info(data: dict) -> Dict:
+    """Deserialize scheduler_info from a shared dictionary"""
+    return data
+
+
+def write_to_shared_memory(data: dict, name: str) -> shared_memory.SharedMemory:
+    """Write data to shared memory"""
+    serialized = json.dumps(data).encode("utf-8")
+    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) -> dict:
+    """Read data from shared memory"""
+    try:
+        shm = shared_memory.SharedMemory(name=name)
+        data = json.loads(bytes(shm.buf).decode("utf-8"))
+        shm.close()
+        return data
+    except FileNotFoundError:
+        raise FileNotFoundError(f"Shared memory {name} not found")
+
+
+def get_main_process_id() -> int:
+    """Get the main process ID"""
+    return multiprocessing.current_process()._parent_pid
+
+
+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}")
+
+    # Write port_args to shared memory
+    port_args_shm = write_to_shared_memory(
+        serialize_port_args(port_args), f"port_args_{current_pid}"
+    )
+    # Write server_args to shared memory
+    server_args_shm = write_to_shared_memory(
+        serialize_server_args(server_args), f"server_args_{current_pid}"
+    )
+    # Write scheduler_info to shared memory
+    scheduler_info_shm = write_to_shared_memory(
+        serialize_scheduler_info(scheduler_info), f"scheduler_info_{current_pid}"
+    )
+
+    port_args_shm.close()
+    server_args_shm.close()
+    scheduler_info_shm.close()
+
+    return port_args_shm, server_args_shm, scheduler_info_shm