# 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.
# ==============================================================================
"""
The entry point of inference server. (SRT = SGLang Runtime)

This file implements python APIs for the inference engine.
"""

from __future__ import annotations

import asyncio
import atexit
import dataclasses
import gc
import logging
import multiprocessing as mp
import os
import random
import signal
import subprocess
import sys
import tempfile
import threading
import time
from typing import (
    Any,
    AsyncIterator,
    Callable,
    Dict,
    Iterator,
    List,
    Optional,
    Tuple,
    Union,
    cast,
)

import torch
import uvloop
import zmq

from sglang.srt.elastic_ep.expert_backup_manager import run_expert_backup_manager
from sglang.srt.entrypoints.engine_info_bootstrap_server import (
    EngineInfoBootstrapServer,
)
from sglang.srt.entrypoints.engine_score_mixin import EngineScoreMixin
from sglang.srt.entrypoints.EngineBase import EngineBase
from sglang.srt.environ import envs
from sglang.srt.managers.data_parallel_controller import (
    SCHEDULER_PIDS_ARG,
    run_data_parallel_controller_process,
)
from sglang.srt.managers.detokenizer_manager import run_detokenizer_process
from sglang.srt.managers.io_struct import (
    CloseSessionReqInput,
    DestroyWeightsUpdateGroupReqInput,
    EmbeddingReqInput,
    GenerateReqInput,
    GetWeightsByNameReqInput,
    InitWeightsUpdateGroupReqInput,
    LoadLoRAAdapterFromTensorsReqInput,
    LoadLoRAAdapterReqInput,
    MultimodalDataInputFormat,
    OpenSessionReqInput,
    ProfileReq,
    ProfileReqType,
    ReleaseMemoryOccupationReqInput,
    ResumeMemoryOccupationReqInput,
    ReturnHiddenStatesMode,
    RpcReqInput,
    RpcReqOutput,
    UnloadLoRAAdapterReqInput,
    UpdateWeightFromDiskReqInput,
    UpdateWeightsFromDistributedReqInput,
    UpdateWeightsFromIPCReqInput,
    UpdateWeightsFromTensorReqInput,
    sock_recv,
    sock_send,
)
from sglang.srt.managers.multi_tokenizer_mixin import (
    MultiTokenizerRouter,
    run_multi_detokenizer_router_process,
)
from sglang.srt.managers.scheduler import run_scheduler_process
from sglang.srt.managers.tokenizer_manager import TokenizerManager
from sglang.srt.observability.startup_time import build_engine_startup_time
from sglang.srt.observability.trace import process_tracing_init, trace_set_thread_info
from sglang.srt.parser.template_detection import resolve_auto_parsers
from sglang.srt.parser.template_manager import TemplateManager
from sglang.srt.plugins import load_plugins
from sglang.srt.server_args import PortArgs, ServerArgs
from sglang.srt.utils import (
    MultiprocessingSerializer,
    SerializedTensorPayload,
    assert_pkg_version,
    configure_logger,
    get_bool_env_var,
    is_cuda,
    is_mnnvl_fabric_device,
    kill_process_tree,
    launch_dummy_health_check_server,
    maybe_reindex_device_id,
    normalize_serialized_named_tensor_payloads,
    numa_utils,
    set_prometheus_multiproc_dir,
    set_ulimit,
)
from sglang.srt.utils.msgspec_utils import msgspec_to_builtins
from sglang.srt.utils.network import (
    NetworkAddress,
    get_free_port,
    get_zmq_socket,
    is_port_available,
)
from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
from sglang.srt.utils.watchdog import SubprocessWatchdog
from sglang.version import __version__

logger = logging.getLogger(__name__)
asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())

_is_cuda = is_cuda()


@dataclasses.dataclass
class SchedulerInitResult:
    """Result from launching schedulers."""

    scheduler_infos: List[Dict[str, Any]]
    all_child_pids: List[int] = dataclasses.field(default_factory=list)
    wait_for_ready: Callable[[], None] = lambda: None
    block_until_scheduler_exits: Callable[[], None] = lambda: None
    engine_info_bootstrap_server: Optional[Any] = None


def init_tokenizer_manager(
    server_args: ServerArgs,
    port_args: PortArgs,
    TokenizerManagerClass: Optional[TokenizerManager] = None,
) -> Tuple[TokenizerManager, TemplateManager]:
    # Launch tokenizer process
    TokenizerManagerClass = TokenizerManagerClass or TokenizerManager
    tokenizer_manager = TokenizerManagerClass(server_args, port_args)

    # Initialize templates
    template_manager = TemplateManager()
    template_manager.initialize_templates(
        tokenizer_manager=tokenizer_manager,
        model_path=server_args.model_path,
        chat_template=server_args.chat_template,
        completion_template=server_args.completion_template,
    )

    # Resolve any remaining auto parsers using template manager's detection results
    for attr, suggested, label in (
        (
            "reasoning_parser",
            template_manager.suggested_reasoning_parser,
            "reasoning parser",
        ),
        (
            "tool_call_parser",
            template_manager.suggested_tool_call_parser,
            "tool-call parser",
        ),
    ):
        if tokenizer_manager.config_value(attr) != "auto":
            continue
        if suggested is not None:
            tokenizer_manager.record_config_updates(
                "template-detection", **{attr: suggested}
            )
            logger.info(
                f"Auto-detected --{attr.replace('_', '-')} as '{suggested}' from chat template"
            )
        else:
            logger.warning(
                f"--{attr.replace('_', '-')}=auto specified but could not detect "
                f"{label} from chat template. Disabling {label}."
            )
            tokenizer_manager.record_config_updates(
                "template-detection", **{attr: None}
            )

    return tokenizer_manager, template_manager


class Engine(EngineScoreMixin, EngineBase):
    """
    The entry point to the inference engine.

    - The engine consists of three components:
        1. TokenizerManager: Tokenizes the requests and sends them to the scheduler.
        2. Scheduler (subprocess): Receives requests from the Tokenizer Manager, schedules batches, forwards them, and sends the output tokens to the Detokenizer Manager.
        3. DetokenizerManager (subprocess): Detokenizes the output tokens and sends the result back to the Tokenizer Manager.

    Note:
    1. The HTTP server, Engine, and TokenizerManager all run in the main process.
    2. Inter-process communication is done through IPC (each process uses a different port) via the ZMQ library.
    """

    # Some fields to allow people to override the server args
    # and launch processes for their private forks.
    server_args_class: ServerArgs = ServerArgs
    init_tokenizer_manager_func: Callable = staticmethod(init_tokenizer_manager)
    run_scheduler_process_func: Callable = staticmethod(run_scheduler_process)
    run_detokenizer_process_func: Callable = staticmethod(run_detokenizer_process)

    # Backend-specific launch handle: the Ray engine schedules its actors onto a
    # placement group. Not config — a live cluster object.
    _placement_group = None

    def __init__(self, **kwargs):
        """
        The arguments of this function is the same as `sglang/srt/server_args.py::ServerArgs`.
        Please refer to `ServerArgs` for the documentation.
        """

        # Ensure plugins are loaded before ServerArgs construction,
        # so hooks on ServerArgs.__post_init__ fire correctly.
        load_plugins()

        # Parse server_args
        if "server_args" in kwargs:
            # Directly load server_args
            server_args = kwargs["server_args"]
        else:
            # Construct server_args from kwargs
            if "log_level" not in kwargs:
                # Do not print logs by default
                kwargs["log_level"] = "error"
            server_args = self.server_args_class(**kwargs)
        self.server_args = server_args
        logger.info(f"{server_args=}")

        # Rust Server is not supported with the offline Engine API
        if envs.SGLANG_RUST_SERVER.get():
            raise ValueError(
                "SGLANG_RUST_SERVER is not supported with the offline Engine "
                "API; it only replaces the HTTP server path (`sglang serve`). "
                "Unset SGLANG_RUST_SERVER to use sgl.Engine."
            )

        # Pre-initialize tokenizer_manager so the atexit handler in
        # shutdown() won't hit AttributeError.
        self.tokenizer_manager = None

        # Shutdown the subprocesses automatically when the program exits
        atexit.register(self.shutdown)

        # Launch subprocesses
        (
            tokenizer_manager,
            template_manager,
            port_args,
            scheduler_init_result,
            subprocess_watchdog,
            weight_cache_daemon_procs,
        ) = self._launch_subprocesses(
            server_args=server_args,
            init_tokenizer_manager_func=self.init_tokenizer_manager_func,
            run_scheduler_process_func=self.run_scheduler_process_func,
            run_detokenizer_process_func=self.run_detokenizer_process_func,
            placement_group=self._placement_group,
        )
        self.tokenizer_manager = tokenizer_manager
        self.template_manager = template_manager
        self._scheduler_init_result = scheduler_init_result
        # Engine-spawned weight cache daemons owned by *this* instance (empty
        # unless --weight-cache-mode daemon). Kept per-instance so two Engines
        # in one process each reap only their own daemons in shutdown().
        self._weight_cache_daemon_procs = weight_cache_daemon_procs
        if tokenizer_manager is not None:
            tokenizer_manager._subprocess_watchdog = subprocess_watchdog
        self.port_args = port_args

        # Initialize ZMQ sockets
        context = zmq.Context(2)
        if self.server_args.node_rank == 0:
            self.send_to_rpc = get_zmq_socket(
                context, zmq.DEALER, self.port_args.rpc_ipc_name, True
            )
        else:
            self.send_to_rpc = None

        # Enable tracing
        if server_args.enable_trace:
            process_tracing_init(
                server_args.otlp_traces_endpoint,
                "sglang",
                trace_modules=server_args.trace_modules,
            )
            thread_label = "Tokenizer"
            if server_args.disaggregation_mode == "prefill":
                thread_label = "Prefill Tokenizer"
            elif server_args.disaggregation_mode == "decode":
                thread_label = "Decode Tokenizer"
            trace_set_thread_info(thread_label)

        try:
            self.loop = asyncio.get_running_loop()
        except RuntimeError:
            self.loop = asyncio.new_event_loop()
            asyncio.set_event_loop(self.loop)

    def get_all_child_pids(self) -> List[int]:
        """Returns a list of all child process PIDs."""
        return self._scheduler_init_result.all_child_pids

    def _resolve_routed_dp_rank(
        self,
        routed_dp_rank: Optional[int],
        data_parallel_rank: Optional[int],
    ) -> Optional[int]:
        if data_parallel_rank is not None:
            import warnings

            warnings.warn(
                "'data_parallel_rank' is deprecated, use 'routed_dp_rank' instead.",
                DeprecationWarning,
                stacklevel=3,
            )
            if routed_dp_rank is None:
                routed_dp_rank = data_parallel_rank

        if routed_dp_rank is not None:
            dp_size = self.server_args.dp_size
            if dp_size <= 1 and routed_dp_rank == 0:
                logger.debug(
                    f"routed_dp_rank={routed_dp_rank} is ignored because dp_size={dp_size}"
                )
                return None
            if routed_dp_rank < 0 or routed_dp_rank >= dp_size:
                raise ValueError(
                    f"routed_dp_rank={routed_dp_rank} out of range [0, {dp_size})"
                )

        logger.debug(f"routed_dp_rank: {routed_dp_rank}")
        return routed_dp_rank

    def generate(
        self,
        # The input prompt. It can be a single prompt or a batch of prompts.
        prompt: Optional[Union[List[str], str]] = None,
        sampling_params: Optional[Union[List[Dict], Dict]] = None,
        # The token ids for text; one can either specify text or input_ids.
        input_ids: Optional[Union[List[List[int]], List[int]]] = None,
        # The image input. It can be an image instance, file name, URL, or base64 encoded string.
        # Can be formatted as:
        # - Single image for a single request
        # - List of images (one per request in a batch)
        # - List of lists of images (multiple images per request)
        # - List of preprocessed outputs from a Huggingface processor, each as a dict containing `format`: 'processor_output' and other data
        # - List of precomputed image embeddings, each as a dict containing field `format`: 'precomputed_embedding' and `feature`: the precomputed embedding
        # See also python/sglang/srt/utils.py:load_image for more details.
        image_data: Optional[MultimodalDataInputFormat] = None,
        audio_data: Optional[MultimodalDataInputFormat] = None,
        video_data: Optional[MultimodalDataInputFormat] = None,
        # See GenerateReqInput.mm_hashes / async_generate for the contract.
        mm_hashes: Optional[Union[List[str], List[List[str]]]] = None,
        return_logprob: Optional[Union[List[bool], bool]] = False,
        logprob_start_len: Optional[Union[List[int], int]] = None,
        top_logprobs_num: Optional[Union[List[int], int]] = None,
        token_ids_logprob: Optional[Union[List[List[int]], List[int]]] = None,
        lora_path: Optional[List[Optional[str]]] = None,
        custom_logit_processor: Optional[Union[List[str], str]] = None,
        require_reasoning: bool = False,
        return_hidden_states: Union[
            ReturnHiddenStatesMode, List[ReturnHiddenStatesMode]
        ] = False,
        return_routed_experts: bool = False,
        routed_experts_start_len: int = 0,
        stream: bool = False,
        bootstrap_host: Optional[Union[List[str], str]] = None,
        bootstrap_port: Optional[Union[List[int], int]] = None,
        bootstrap_room: Optional[Union[List[int], int]] = None,
        routed_dp_rank: Optional[int] = None,
        disagg_prefill_dp_rank: Optional[int] = None,
        # Deprecated: use routed_dp_rank instead
        data_parallel_rank: Optional[int] = None,
        external_trace_header: Optional[Dict] = None,
        rid: Optional[Union[List[str], str]] = None,
        session_params: Optional[Dict] = None,
        priority: Optional[int] = None,
        session_id: Optional[str] = None,
    ) -> Union[Dict, Iterator[Dict]]:
        """
        The arguments of this function is the same as `sglang/srt/managers/io_struct.py::GenerateReqInput`.
        Please refer to `GenerateReqInput` for the documentation.
        """
        routed_dp_rank = self._resolve_routed_dp_rank(
            routed_dp_rank, data_parallel_rank
        )

        obj = GenerateReqInput(
            text=prompt,
            input_ids=input_ids,
            sampling_params=sampling_params,
            image_data=image_data,
            audio_data=audio_data,
            video_data=video_data,
            mm_hashes=mm_hashes,
            return_logprob=return_logprob,
            logprob_start_len=logprob_start_len,
            top_logprobs_num=top_logprobs_num,
            token_ids_logprob=token_ids_logprob,
            lora_path=lora_path,
            custom_logit_processor=custom_logit_processor,
            require_reasoning=require_reasoning,
            return_hidden_states=return_hidden_states,
            return_routed_experts=return_routed_experts,
            routed_experts_start_len=routed_experts_start_len,
            stream=stream,
            bootstrap_host=bootstrap_host,
            bootstrap_port=bootstrap_port,
            bootstrap_room=bootstrap_room,
            routed_dp_rank=routed_dp_rank,
            disagg_prefill_dp_rank=disagg_prefill_dp_rank,
            external_trace_header=external_trace_header,
            rid=rid,
            session_id=session_id,
            session_params=session_params,
            priority=priority,
        )
        generator = self.tokenizer_manager.generate_request(obj, None)

        if stream:

            def generator_wrapper():
                while True:
                    try:
                        chunk = self.loop.run_until_complete(generator.__anext__())
                        yield chunk
                    except StopAsyncIteration:
                        break

            return generator_wrapper()
        else:
            ret = self.loop.run_until_complete(generator.__anext__())
            return ret

    async def async_generate(
        self,
        # The input prompt. It can be a single prompt or a batch of prompts.
        prompt: Optional[Union[List[str], str]] = None,
        sampling_params: Optional[Union[List[Dict], Dict]] = None,
        # The token ids for text; one can either specify text or input_ids.
        input_ids: Optional[Union[List[List[int]], List[int]]] = None,
        # The image input. It can be an image instance, file name, URL, or base64 encoded string.
        # Can be formatted as:
        # - Single image for a single request
        # - List of images (one per request in a batch)
        # - List of lists of images (multiple images per request)
        # - List of preprocessed outputs from a Huggingface processor, each as a dict containing `format`: 'processor_output' and other data
        # - List of precomputed image embeddings, each as a dict containing field `format`: 'precomputed_embedding' and `feature`: the precomputed embedding
        # See also python/sglang/srt/utils.py:load_image for more details.
        image_data: Optional[MultimodalDataInputFormat] = None,
        audio_data: Optional[MultimodalDataInputFormat] = None,
        video_data: Optional[MultimodalDataInputFormat] = None,
        # Optional per-image hashes the caller has already computed (hex strings,
        # one per image in `image_data`). When supplied, each MultimodalDataItem's
        # `hash` is initialised from this list and `set_pad_value` skips the
        # internal `hash_feature()` recompute. Intended for external KV routers
        # that compute their own per-image hash for routing decisions and need
        # sglang's prefix-cache key to align. See GenerateReqInput.mm_hashes.
        mm_hashes: Optional[Union[List[str], List[List[str]]]] = None,
        return_logprob: Optional[Union[List[bool], bool]] = False,
        logprob_start_len: Optional[Union[List[int], int]] = None,
        top_logprobs_num: Optional[Union[List[int], int]] = None,
        token_ids_logprob: Optional[Union[List[List[int]], List[int]]] = None,
        lora_path: Optional[List[Optional[str]]] = None,
        custom_logit_processor: Optional[Union[List[str], str]] = None,
        require_reasoning: bool = False,
        return_hidden_states: Union[
            ReturnHiddenStatesMode, List[ReturnHiddenStatesMode]
        ] = False,
        return_routed_experts: bool = False,
        routed_experts_start_len: int = 0,
        stream: bool = False,
        bootstrap_host: Optional[Union[List[str], str]] = None,
        bootstrap_port: Optional[Union[List[int], int]] = None,
        bootstrap_room: Optional[Union[List[int], int]] = None,
        routed_dp_rank: Optional[int] = None,
        disagg_prefill_dp_rank: Optional[int] = None,
        # Deprecated: use routed_dp_rank instead
        data_parallel_rank: Optional[int] = None,
        external_trace_header: Optional[Dict] = None,
        rid: Optional[Union[List[str], str]] = None,
        session_params: Optional[Dict] = None,
        priority: Optional[int] = None,
        session_id: Optional[str] = None,
    ) -> Union[Dict, AsyncIterator[Dict]]:
        """
        The arguments of this function is the same as `sglang/srt/managers/io_struct.py::GenerateReqInput`.
        Please refer to `GenerateReqInput` for the documentation.
        """
        routed_dp_rank = self._resolve_routed_dp_rank(
            routed_dp_rank, data_parallel_rank
        )

        obj = GenerateReqInput(
            text=prompt,
            input_ids=input_ids,
            sampling_params=sampling_params,
            image_data=image_data,
            audio_data=audio_data,
            video_data=video_data,
            mm_hashes=mm_hashes,
            return_logprob=return_logprob,
            logprob_start_len=logprob_start_len,
            top_logprobs_num=top_logprobs_num,
            token_ids_logprob=token_ids_logprob,
            lora_path=lora_path,
            require_reasoning=require_reasoning,
            return_hidden_states=return_hidden_states,
            return_routed_experts=return_routed_experts,
            routed_experts_start_len=routed_experts_start_len,
            stream=stream,
            custom_logit_processor=custom_logit_processor,
            bootstrap_host=bootstrap_host,
            bootstrap_port=bootstrap_port,
            bootstrap_room=bootstrap_room,
            routed_dp_rank=routed_dp_rank,
            disagg_prefill_dp_rank=disagg_prefill_dp_rank,
            external_trace_header=external_trace_header,
            rid=rid,
            session_id=session_id,
            session_params=session_params,
            priority=priority,
        )
        generator = self.tokenizer_manager.generate_request(obj, None)

        if stream is True:
            return generator
        else:
            return await generator.__anext__()

    def encode(
        self,
        prompt: Union[str, List[str], List[Dict], List[List[Dict]]],
        image_data: Optional[MultimodalDataInputFormat] = None,
        audio_data: Optional[MultimodalDataInputFormat] = None,
        video_data: Optional[MultimodalDataInputFormat] = None,
        dimensions: Optional[int] = None,
        lora_path: Optional[Union[List[Optional[str]], Optional[str]]] = None,
        embed_override_token_id: Optional[int] = None,
        embed_overrides: Optional[List[List[torch.Tensor]]] = None,
        external_trace_header: Optional[Dict] = None,
        rid: Optional[Union[List[str], str]] = None,
    ) -> Dict:
        """
        The arguments of this function is the same as `sglang/srt/managers/io_struct.py::EmbeddingReqInput`.
        Please refer to `EmbeddingReqInput` for the documentation.
        """
        obj = EmbeddingReqInput(
            text=prompt,
            image_data=image_data,
            audio_data=audio_data,
            video_data=video_data,
            dimensions=dimensions,
            lora_path=lora_path,
            embed_override_token_id=embed_override_token_id,
            embed_overrides=embed_overrides,
            external_trace_header=external_trace_header,
            rid=rid,
        )
        generator = self.tokenizer_manager.generate_request(obj, None)
        ret = self.loop.run_until_complete(generator.__anext__())
        return ret

    async def async_encode(
        self,
        prompt: Union[str, List[str], List[Dict], List[List[Dict]]],
        image_data: Optional[MultimodalDataInputFormat] = None,
        audio_data: Optional[MultimodalDataInputFormat] = None,
        video_data: Optional[MultimodalDataInputFormat] = None,
        dimensions: Optional[int] = None,
        lora_path: Optional[Union[List[Optional[str]], Optional[str]]] = None,
        embed_override_token_id: Optional[int] = None,
        embed_overrides: Optional[List[List[torch.Tensor]]] = None,
        external_trace_header: Optional[Dict] = None,
        rid: Optional[Union[List[str], str]] = None,
    ) -> Dict:
        """
        Asynchronous version of encode method.

        The arguments of this function is the same as `sglang/srt/managers/io_struct.py::EmbeddingReqInput`.
        Please refer to `EmbeddingReqInput` for the documentation.
        """
        obj = EmbeddingReqInput(
            text=prompt,
            image_data=image_data,
            audio_data=audio_data,
            video_data=video_data,
            dimensions=dimensions,
            lora_path=lora_path,
            embed_override_token_id=embed_override_token_id,
            embed_overrides=embed_overrides,
            external_trace_header=external_trace_header,
            rid=rid,
        )
        generator = self.tokenizer_manager.generate_request(obj, None)
        return await generator.__anext__()

    def rerank(
        self,
        prompt: Union[List[List[str]]],
    ) -> Dict:
        """
        The arguments of this function is the same as `sglang/srt/managers/io_struct.py::EmbeddingReqInput`.
        Please refer to `EmbeddingReqInput` for the documentation.
        """
        obj = EmbeddingReqInput(text=prompt, is_cross_encoder_request=True)
        generator = self.tokenizer_manager.generate_request(obj, None)
        ret = self.loop.run_until_complete(generator.__anext__())
        return ret

    @classmethod
    def _launch_weight_cache_daemons(cls, server_args: ServerArgs):
        """Launch weight cache daemon processes for this node's PP×TP ranks.

        All daemon processes join the same NCCL distributed group so that
        TP-sharded model loading works correctly. Each daemon holds its
        rank's weight shard in GPU memory and serves IPC handles.

        Lifecycle: these daemons are *co-terminal* with the engine. They are
        children of this process (kill_itself_when_parent_died installs
        PR_SET_PDEATHSIG) and are gracefully reaped in ``shutdown()``. They do
        NOT persist across engine restarts, so ``--weight-cache-mode daemon``
        on its own does not deliver a faster restart -- the first start is in
        fact slower (disk-load into the daemon plus the IPC handshake). The
        fast-recovery story is the standalone launcher
        (``python -m sglang.srt.weight_cache.daemon``) plus
        ``--weight-cache-mode client``, where the daemon outlives the engine.
        """
        if server_args.dp_size > 1:
            raise ValueError(
                "Weight cache daemon mode does not support dp_size > 1. "
                "Please set --dp-size 1 when using --weight-cache-mode daemon."
            )

        # Multi-node needs an explicit rendezvous address; otherwise each node
        # picks its own local 127.0.0.1 port (below) and the per-node daemons
        # can never form the joint process group.
        if server_args.nnodes > 1 and not server_args.dist_init_addr:
            raise ValueError(
                "Multi-node weight cache daemons (nnodes > 1) require "
                "--dist-init-addr so all nodes rendezvous at the same endpoint."
            )

        tp_size = server_args.tp_size

        pp_rank_range, tp_rank_range, pp_size_per_node, tp_size_per_node = (
            _calculate_rank_ranges(
                server_args.nnodes,
                server_args.pp_size,
                tp_size,
                server_args.node_rank,
            )
        )

        # Build the distributed init method (multi-node uses the user-provided
        # dist_init_addr so all nodes reach the same endpoint).
        if server_args.dist_init_addr:
            host, port = server_args.dist_init_addr.rsplit(":", 1)
            dist_init_method = f"tcp://{host}:{port}"
        else:
            # Fresh free port for the daemons' own rendezvous, not the engine's
            # nccl_port: a pinned --nccl-port would otherwise collide with the
            # engine's own NCCL TCPStore.
            dist_init_method = NetworkAddress("127.0.0.1", get_free_port()).to_tcp()

        num_daemons = len(pp_rank_range) * len(tp_rank_range)
        daemon_procs = []
        logger.info(
            f"Launching {num_daemons} weight cache daemon(s) on node "
            f"{server_args.node_rank} for model={server_args.model_path}, "
            f"pp_ranks={pp_rank_range.start}..{pp_rank_range.stop - 1}, "
            f"tp_ranks={tp_rank_range.start}..{tp_rank_range.stop - 1}, "
            f"dist_init_method={dist_init_method}"
        )

        # Validate and clean up stale .ready/.sock files from prior runs.
        # If a daemon is still alive at this rank, raise instead of clobbering.
        from sglang.srt.weight_cache.protocol import (
            cleanup_stale_daemon_files,
            compute_global_rank,
            compute_local_gpu_id,
            get_ready_path,
        )

        for pp_rank in pp_rank_range:
            for tp_rank in tp_rank_range:
                global_rank = compute_global_rank(tp_size, pp_rank, tp_rank)
                cleanup_stale_daemon_files(global_rank)

        for pp_rank in pp_rank_range:
            for tp_rank in tp_rank_range:
                gpu_id = compute_local_gpu_id(
                    pp_rank,
                    tp_rank,
                    pp_size_per_node,
                    tp_size_per_node,
                    base_gpu_id=server_args.base_gpu_id,
                    gpu_id_step=server_args.gpu_id_step,
                )
                cmd = [
                    sys.executable,
                    "-m",
                    "sglang.srt.weight_cache.daemon",
                    "--model-path",
                    server_args.model_path,
                    "--gpu-id",
                    str(gpu_id),
                    "--tp-size",
                    str(tp_size),
                    "--tp-rank",
                    str(tp_rank),
                    "--pp-size",
                    str(server_args.pp_size),
                    "--pp-rank",
                    str(pp_rank),
                    "--dp-size",
                    "1",
                    "--ep-size",
                    str(server_args.ep_size),
                    "--load-format",
                    server_args.load_format,
                    "--dtype",
                    server_args.dtype,
                    "--dist-init-method",
                    dist_init_method,
                ]
                if server_args.quantization:
                    cmd += ["--quantization", server_args.quantization]
                if (
                    server_args.model_loader_extra_config
                    and server_args.model_loader_extra_config != "{}"
                ):
                    cmd += [
                        "--model-loader-extra-config",
                        server_args.model_loader_extra_config,
                    ]
                if server_args.trust_remote_code:
                    cmd += ["--trust-remote-code"]
                if server_args.revision:
                    cmd += ["--revision", server_args.revision]

                proc = subprocess.Popen(cmd)
                daemon_procs.append(proc)

        # Wait for all daemons to be ready (ready file exists). On any failure
        # (readiness timeout or a daemon exiting early) terminate the siblings
        # we already spawned before propagating, so a partial launch does not
        # leak GPU-resident daemons.
        timeout = server_args.weight_cache_timeout
        check_interval = 2
        start_time = time.time()
        try:
            for pp_rank in pp_rank_range:
                for tp_rank in tp_rank_range:
                    global_rank = compute_global_rank(tp_size, pp_rank, tp_rank)
                    ready_path = get_ready_path(global_rank)
                    while not os.path.exists(ready_path):
                        time.sleep(check_interval)
                        if time.time() - start_time > timeout:
                            raise TimeoutError(
                                f"Weight cache daemon for pp_rank={pp_rank} "
                                f"tp_rank={tp_rank} did not become ready "
                                f"within {timeout}s"
                            )
                        # Check if daemon process is still alive
                        for p in daemon_procs:
                            if p.poll() is not None:
                                raise RuntimeError(
                                    f"Weight cache daemon (pid={p.pid}) exited prematurely "
                                    f"with code {p.returncode}"
                                )
                    logger.info(
                        f"Weight cache daemon for pp_rank={pp_rank} "
                        f"tp_rank={tp_rank} is ready"
                    )
        except BaseException:
            cls._terminate_weight_cache_daemons(daemon_procs)
            raise

        logger.info(
            f"All {num_daemons} weight cache daemons on node "
            f"{server_args.node_rank} are ready"
        )
        return daemon_procs

    @staticmethod
    def _terminate_weight_cache_daemons(procs, timeout: float = 10.0):
        """Gracefully stop engine-spawned weight cache daemons.

        Send SIGTERM first so each daemon's signal handler can unlink its
        ``.sock``/``.ready`` files, then SIGKILL any straggler. This matters
        because ``shutdown()`` otherwise reaps children via
        ``kill_process_tree`` (SIGKILL), which would skip that cleanup and
        leave stale files that make the next client-mode boot fail with a
        confusing "socket exists but connection refused" instead of a clean
        "no daemon" path.
        """
        if not procs:
            return
        for p in procs:
            if p.poll() is None:
                p.terminate()  # SIGTERM -> daemon cleanup handler runs
        for p in procs:
            try:
                p.wait(timeout=timeout)
            except subprocess.TimeoutExpired:
                logger.warning(
                    f"Weight cache daemon (pid={p.pid}) did not exit within "
                    f"{timeout}s of SIGTERM; sending SIGKILL."
                )
                p.kill()

    @classmethod
    def _launch_scheduler_processes(
        cls,
        server_args: ServerArgs,
        port_args: PortArgs,
        run_scheduler_process_func: Callable,
        *,
        placement_group=None,
    ) -> Tuple[SchedulerInitResult, Optional[List]]:
        """Launch scheduler processes using multiprocessing.
        Override in subclasses for different backends (e.g. Ray).

        Returns:
            Tuple of (SchedulerInitResult, scheduler_procs).
            scheduler_procs is None for RayEngine (uses Ray actors instead).
        """
        scheduler_procs = []
        use_dp_controller = (
            server_args.dp_size > 1 or server_args.ep_join_mode == "scale"
        )

        if not use_dp_controller:
            # Launch tensor parallel scheduler processes
            memory_saver_adapter = TorchMemorySaverAdapter.create(
                enable=server_args.enable_memory_saver
            )
            scheduler_pipe_readers = []

            pp_rank_range, tp_rank_range, pp_size_per_node, tp_size_per_node = (
                _calculate_rank_ranges(
                    server_args.nnodes,
                    server_args.pp_size,
                    server_args.tp_size,
                    server_args.node_rank,
                )
            )

            for pp_rank in pp_rank_range:
                for tp_rank in tp_rank_range:
                    reader, writer = mp.Pipe(duplex=False)
                    gpu_id = (
                        server_args.base_gpu_id
                        + ((pp_rank % pp_size_per_node) * tp_size_per_node)
                        + (tp_rank % tp_size_per_node) * server_args.gpu_id_step
                    )
                    attn_cp_rank, moe_dp_rank, moe_ep_rank = _compute_parallelism_ranks(
                        server_args, tp_rank
                    )

                    with maybe_reindex_device_id(gpu_id) as gpu_id:
                        proc = mp.Process(
                            target=run_scheduler_process_func,
                            args=(
                                server_args,
                                port_args,
                                gpu_id,
                                tp_rank,
                                attn_cp_rank,
                                moe_dp_rank,
                                moe_ep_rank,
                                pp_rank,
                                None,
                                writer,
                            ),
                        )
                        with (
                            memory_saver_adapter.configure_subprocess(),
                            numa_utils.configure_subprocess(server_args, gpu_id),
                        ):
                            proc.start()

                    scheduler_procs.append(proc)
                    scheduler_pipe_readers.append(reader)
        else:
            # Launch the data parallel controller
            reader, writer = mp.Pipe(duplex=False)
            scheduler_pipe_readers = [reader]
            proc = mp.Process(
                target=run_data_parallel_controller_process,
                kwargs=dict(
                    server_args=server_args,
                    port_args=port_args,
                    pipe_writer=writer,
                    run_scheduler_process_func=run_scheduler_process_func,
                ),
            )
            proc.start()
            scheduler_procs.append(proc)

        all_child_pids = [proc.pid for proc in scheduler_procs]
        scheduler_infos = []

        def wait_for_ready():
            infos = _wait_for_scheduler_ready(scheduler_pipe_readers, scheduler_procs)
            scheduler_infos.extend(infos)
            if use_dp_controller:
                for info in infos:
                    if SCHEDULER_PIDS_ARG in info:
                        all_child_pids.extend(info[SCHEDULER_PIDS_ARG])

        def block_until_scheduler_exits():
            for proc in scheduler_procs:
                proc.join()
                logger.error(
                    f"Scheduler or DataParallelController {proc.pid} "
                    f"terminated with {proc.exitcode}"
                )

        return (
            SchedulerInitResult(
                scheduler_infos=scheduler_infos,
                all_child_pids=all_child_pids,
                wait_for_ready=wait_for_ready,
                block_until_scheduler_exits=block_until_scheduler_exits,
            ),
            scheduler_procs,
        )

    @classmethod
    def _launch_detokenizer_subprocesses(
        cls,
        server_args: ServerArgs,
        port_args: PortArgs,
        run_detokenizer_process_func: Callable,
    ) -> Tuple[List[mp.Process], List[str]]:
        """Launch detokenizer worker(s).

        - When ``detokenizer_worker_num == 1``: a single detokenizer process listens on
          ``port_args.detokenizer_ipc_name`` (the original behavior).
        - When ``detokenizer_worker_num > 1``: each detokenizer worker gets its own
          private IPC socket, and a ``MultiDetokenizerRouter`` process owns the
          original ``port_args.detokenizer_ipc_name`` and fans out to them.

        Returns (processes, names) for SubprocessWatchdog.
        """
        processes: List[mp.Process] = []
        names: List[str] = []

        if server_args.detokenizer_worker_num <= 1:
            proc = mp.Process(
                target=run_detokenizer_process_func,
                args=(server_args, port_args),
            )
            proc.start()
            processes.append(proc)
            names.append("detokenizer")
            return processes, names

        router_ipc_name = port_args.detokenizer_ipc_name
        worker_ipc_names: List[str] = []
        try:
            for i in range(server_args.detokenizer_worker_num):
                worker_ipc = f"ipc://{tempfile.NamedTemporaryFile(delete=False).name}"
                port_args.detokenizer_ipc_name = worker_ipc
                proc = mp.Process(
                    target=run_detokenizer_process_func,
                    args=(server_args, port_args),
                )
                proc.start()
                processes.append(proc)
                names.append(f"detokenizer_{i}")
                worker_ipc_names.append(worker_ipc)
        finally:
            port_args.detokenizer_ipc_name = router_ipc_name

        router_proc = mp.Process(
            target=run_multi_detokenizer_router_process,
            args=(worker_ipc_names, server_args, port_args),
        )
        router_proc.start()
        processes.append(router_proc)
        names.append("detokenizer_router")

        return processes, names

    @staticmethod
    def _set_startup_time(
        tokenizer_manager: Union[TokenizerManager, MultiTokenizerRouter],
        scheduler_init_result: SchedulerInitResult,
        startup_tic: float,
    ) -> None:
        startup_time = build_engine_startup_time(
            (
                info.get("startup_time")
                for info in scheduler_init_result.scheduler_infos
            ),
            tokenizer_e2e=time.perf_counter() - startup_tic,
        )
        tokenizer_manager.set_startup_time(startup_time)
        cuda_graph_timings = ", ".join(
            f"{phase}={duration:.2f}"
            for phase, duration in startup_time["cuda_graph"].items()
        )
        logger.info(
            "Engine startup timings (s): load_weight=%.2f, "
            "kv_cache_allocation=%.2f, scheduler_e2e=%.2f, "
            "cuda_graph={%s}, tokenizer_e2e=%.2f",
            startup_time["load_weight"],
            startup_time["kv_cache_allocation"],
            startup_time["scheduler_e2e"],
            cuda_graph_timings,
            startup_time["tokenizer_e2e"],
        )

    @classmethod
    def _launch_subprocesses(
        cls,
        server_args: ServerArgs,
        init_tokenizer_manager_func: Callable,
        run_scheduler_process_func: Callable,
        run_detokenizer_process_func: Callable,
        port_args: Optional[PortArgs] = None,
        placement_group=None,
    ) -> Tuple[
        TokenizerManager,
        TemplateManager,
        PortArgs,
        SchedulerInitResult,
        Optional[SubprocessWatchdog],
    ]:
        """Launch the TokenizerManager in the main process, the Scheduler in a subprocess, and the DetokenizerManager in another subprocess.

        Returns:
            Tuple of (tokenizer_manager, template_manager, port_args, scheduler_init_result, subprocess_watchdog, weight_cache_daemon_procs).
        """
        startup_tic = time.perf_counter()

        # Configure global environment
        configure_logger(server_args)
        _set_envs_and_config(server_args)

        # Defensive: ensure plugins loaded (may already be loaded by
        # Engine.__init__ or CLI entry).
        load_plugins()

        server_args.check_server_args()

        # Allocate ports for inter-process communications
        if port_args is None:
            port_args = PortArgs.init_new(server_args)
        logger.info(f"{server_args=}")

        # Start the engine info bootstrap server if per-rank info is needed.
        engine_info_bootstrap_server = None
        if (
            server_args.remote_instance_weight_loader_start_seed_via_transfer_engine
            and server_args.node_rank == 0
        ):
            bootstrap_port = server_args.engine_info_bootstrap_port
            if not is_port_available(bootstrap_port):
                raise RuntimeError(
                    f"engine_info_bootstrap_port {bootstrap_port} is already in use. "
                    f"When running multiple instances on the same node, each instance must use a "
                    f"different --engine-info-bootstrap-port."
                )
            engine_info_bootstrap_server = EngineInfoBootstrapServer(
                host=server_args.host, port=bootstrap_port
            )

        if (
            server_args.reasoning_parser == "auto"
            or server_args.tool_call_parser == "auto"
        ):
            resolve_auto_parsers(server_args)

        # Launch daemons (daemon mode only). Handles are threaded back to the
        # owning Engine instance (not a class attr) so two Engines in one process
        # don't clobber each other's daemon list.
        weight_cache_daemon_procs: List = []
        if server_args.weight_cache_mode == "daemon":
            weight_cache_daemon_procs = cls._launch_weight_cache_daemons(server_args)

        # Launch scheduler processes
        # Passed only when there is one: this hook is an override point, and a
        # subclass written against the three-argument signature must keep working.
        launch_kwargs = (
            {"placement_group": placement_group} if placement_group is not None else {}
        )
        scheduler_init_result, scheduler_procs = cls._launch_scheduler_processes(
            server_args, port_args, run_scheduler_process_func, **launch_kwargs
        )
        scheduler_init_result.engine_info_bootstrap_server = (
            engine_info_bootstrap_server
        )

        if (
            server_args.enable_elastic_expert_backup
            and server_args.elastic_ep_backend is not None
        ):
            run_expert_backup_manager(server_args, port_args)

        if server_args.node_rank >= 1:
            # Non-zero-rank nodes do not run tokenizer processes.
            scheduler_init_result.wait_for_ready()

            if os.getenv("SGLANG_BLOCK_NONZERO_RANK_CHILDREN") == "0":
                # When using `Engine` as a Python API, we don't want to block here.
                return (
                    None,
                    None,
                    port_args,
                    scheduler_init_result,
                    None,
                    weight_cache_daemon_procs,
                )

            launch_dummy_health_check_server(
                server_args.host, server_args.port, server_args.enable_metrics
            )

            scheduler_init_result.block_until_scheduler_exits()
            return (
                None,
                None,
                port_args,
                scheduler_init_result,
                None,
                weight_cache_daemon_procs,
            )

        # The embedded Rust server (started inside the rank-0 scheduler) owns
        # the API server, tokenization, and detokenization. In that mode we do
        # not start the Python detokenizer subprocess(es) or tokenizer manager.
        # Do not use RayEngine with the Rust server, as it is not supported.
        if envs.SGLANG_RUST_SERVER.get():
            scheduler_init_result.wait_for_ready()
            # Set up subprocess liveness watchdog to detect crashes
            processes = list(scheduler_procs or [])
            names = [f"scheduler_{i}" for i in range(len(processes))]
            subprocess_watchdog = SubprocessWatchdog(
                processes=processes, process_names=names
            )
            subprocess_watchdog.start()

            return (
                None,
                None,
                port_args,
                scheduler_init_result,
                subprocess_watchdog,
                None,
            )

        # Launch detokenizer process(es) — optionally fronted by a router when
        # detokenizer_worker_num > 1.
        detoken_procs, detoken_names = cls._launch_detokenizer_subprocesses(
            server_args=server_args,
            port_args=port_args,
            run_detokenizer_process_func=run_detokenizer_process_func,
        )
        for p in detoken_procs:
            scheduler_init_result.all_child_pids.append(p.pid)

        # Init tokenizer manager first, as the bootstrap server is initialized here
        if server_args.tokenizer_worker_num == 1:
            tokenizer_manager, template_manager = init_tokenizer_manager_func(
                server_args, port_args
            )
        else:
            # Launch multi-tokenizer router
            tokenizer_manager = MultiTokenizerRouter(server_args, port_args)
            template_manager = None

        startup_complete = False
        try:
            # Wait for the model to finish loading
            scheduler_init_result.wait_for_ready()

            cls._set_startup_time(tokenizer_manager, scheduler_init_result, startup_tic)

            # Get back some info from scheduler to tokenizer_manager
            tokenizer_manager.max_req_input_len = scheduler_init_result.scheduler_infos[
                0
            ]["max_req_input_len"]

            # Set up subprocess liveness watchdog to detect crashes
            # Note: RayEngine returns scheduler_procs=None as it uses Ray actors instead of mp.Process
            processes = list(scheduler_procs or [])
            names = [f"scheduler_{i}" for i in range(len(processes))]
            processes.extend(detoken_procs)
            names.extend(detoken_names)
            subprocess_watchdog = SubprocessWatchdog(
                processes=processes, process_names=names
            )
            subprocess_watchdog.start()
            startup_complete = True
        finally:
            if not startup_complete and isinstance(tokenizer_manager, TokenizerManager):
                tokenizer_manager.cuda_vmm_feature_transport.shutdown()

        return (
            tokenizer_manager,
            template_manager,
            port_args,
            scheduler_init_result,
            subprocess_watchdog,
            weight_cache_daemon_procs,
        )

    def shutdown(self):
        """Shutdown the engine; block until the scheduler subprocess releases
        its GPU context so the caller can immediately reallocate on the same
        device."""
        try:
            if (
                self.tokenizer_manager is not None
                and self.tokenizer_manager._subprocess_watchdog is not None
            ):
                self.tokenizer_manager._subprocess_watchdog.stop()

            send_to_rpc = getattr(self, "send_to_rpc", None)
            if send_to_rpc is not None:
                send_to_rpc.close(linger=0)
                self.send_to_rpc = None

            # Gracefully stop weight cache daemons *before* the blanket
            # kill_process_tree below, so their SIGTERM handlers can unlink the
            # .sock/.ready files instead of being SIGKILLed and leaving stale state.
            daemon_procs = getattr(self, "_weight_cache_daemon_procs", None)
            if daemon_procs:
                self._terminate_weight_cache_daemons(daemon_procs)
                self._weight_cache_daemon_procs = []

            kill_process_tree(os.getpid(), include_parent=False, wait_timeout=60)
        finally:
            if isinstance(self.tokenizer_manager, TokenizerManager):
                self.tokenizer_manager.cuda_vmm_feature_transport.shutdown()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.shutdown()
        return False

    def flush_cache(self):
        return self.loop.run_until_complete(self.tokenizer_manager.flush_cache())

    def open_session(
        self,
        capacity_of_str_len: int,
        session_id: Optional[str] = None,
        streaming: bool = False,
        timeout: Optional[float] = None,
    ) -> str:
        """Open a session for multi-turn conversation with shared context.

        Args:
            capacity_of_str_len: Maximum string length capacity for the session.
            session_id: Optional session ID. If not provided, a UUID will be generated.
            streaming: Use low-overhead path for realtime streaming (append-only mode).
            timeout: If set, the session is automatically closed after being inactive
                for this many seconds. Inactivity is measured from session open or the
                most recent request submission.

        Returns:
            The session ID (either the provided one or a newly generated UUID).
        """
        obj = OpenSessionReqInput(
            capacity_of_str_len=capacity_of_str_len,
            session_id=session_id,
            streaming=streaming,
            timeout=timeout,
        )
        return self.loop.run_until_complete(
            self.tokenizer_manager.open_session(obj, None)
        )

    def close_session(self, session_id: str) -> None:
        """Close a session and release its resources.

        Args:
            session_id: The session ID to close.
        """
        obj = CloseSessionReqInput(session_id=session_id)
        self.loop.run_until_complete(self.tokenizer_manager.close_session(obj, None))

    def start_profile(self, **kwargs):
        req = ProfileReq(req_type=ProfileReqType.START_PROFILE, **kwargs)
        self.loop.run_until_complete(self.tokenizer_manager.start_profile(req))

    def stop_profile(self):
        self.loop.run_until_complete(self.tokenizer_manager.stop_profile())

    def start_expert_distribution_record(self):
        self.loop.run_until_complete(
            self.tokenizer_manager.start_expert_distribution_record()
        )

    def stop_expert_distribution_record(self):
        self.loop.run_until_complete(
            self.tokenizer_manager.stop_expert_distribution_record()
        )

    def dump_expert_distribution_record(self):
        self.loop.run_until_complete(
            self.tokenizer_manager.dump_expert_distribution_record()
        )

    def get_server_info(self):
        internal_states = self.loop.run_until_complete(
            self.tokenizer_manager.get_internal_state()
        )
        return msgspec_to_builtins(
            {
                **self.tokenizer_manager.resolved_config_dict(
                    dataclasses.asdict(self.tokenizer_manager.server_args)
                ),
                **self._scheduler_init_result.scheduler_infos[0],
                "startup_time": self.tokenizer_manager.startup_time,
                "internal_states": internal_states,
                "version": __version__,
            }
        )

    def init_weights_update_group(
        self,
        master_address: str,
        master_port: int,
        rank_offset: int,
        world_size: int,
        group_name: str,
        backend: str = "nccl",
    ):
        """Initialize parameter update group."""
        obj = InitWeightsUpdateGroupReqInput(
            master_address=master_address,
            master_port=master_port,
            rank_offset=rank_offset,
            world_size=world_size,
            group_name=group_name,
            backend=backend,
        )
        return self.loop.run_until_complete(
            self.tokenizer_manager.init_weights_update_group(obj, None)
        )

    def destroy_weights_update_group(
        self,
        group_name: str,
    ):
        """Destroy parameter update group."""
        obj = DestroyWeightsUpdateGroupReqInput(
            group_name=group_name,
        )
        return self.loop.run_until_complete(
            self.tokenizer_manager.destroy_weights_update_group(obj, None)
        )

    def update_weights_from_distributed(
        self,
        names: list[str],
        dtypes: list[str],
        shapes: list[list[int]],
        group_name: str = "weight_update_group",
        flush_cache: bool = True,
        load_format: Optional[str] = None,
    ):
        """Update weights from distributed source."""
        obj = UpdateWeightsFromDistributedReqInput(
            names=names,
            dtypes=dtypes,
            shapes=shapes,
            group_name=group_name,
            flush_cache=flush_cache,
            load_format=load_format,
        )
        return self.loop.run_until_complete(
            self.tokenizer_manager.update_weights_from_distributed(obj, None)
        )

    def update_weights_from_tensor(
        self,
        named_tensors: Union[
            List[Tuple[str, torch.Tensor]],
            List[SerializedTensorPayload],
        ],
        load_format: Optional[str] = None,
        flush_cache: bool = True,
    ):
        """Update weights from distributed source. If there are going to be more updates, set `flush_cache` to be false
        to avoid duplicated cache cleaning operation."""
        serialized_named_tensors = self._serialize_tensors_per_rank(
            named_tensors, load_format
        )
        obj = UpdateWeightsFromTensorReqInput(
            serialized_named_tensors=serialized_named_tensors,
            load_format=load_format,
            flush_cache=flush_cache,
        )
        return self.loop.run_until_complete(
            self.tokenizer_manager.update_weights_from_tensor(obj, None)
        )

    def update_weights_from_disk(
        self,
        model_path: str,
        load_format: Optional[str] = None,
    ):
        """Update the weights from disk inplace without re-launching the engine.

        This method allows updating the model weights from disk without restarting
        the engine. It can be used to load a different model or update weights with
        new training.
        """
        obj = UpdateWeightFromDiskReqInput(
            model_path=model_path,
            load_format=load_format,
        )

        return self.loop.run_until_complete(
            self.tokenizer_manager.update_weights_from_disk(obj, None)
        )

    def update_weights_from_ipc(
        self,
        zmq_handles: Dict[str, str],
        flush_cache: bool = True,
    ):
        """Update weights from IPC for checkpoint-engine integration."""
        obj = UpdateWeightsFromIPCReqInput(
            zmq_handles=zmq_handles,
            flush_cache=flush_cache,
        )
        return self.loop.run_until_complete(
            self.tokenizer_manager.update_weights_from_ipc(obj, None)
        )

    def get_weights_by_name(self, name: str, truncate_size: int = 100):
        """Get weights by parameter name."""
        obj = GetWeightsByNameReqInput(name=name, truncate_size=truncate_size)
        return self.loop.run_until_complete(
            self.tokenizer_manager.get_weights_by_name(obj, None)
        )

    def _serialize_tensors_per_rank(
        self,
        tensors,
        load_format: Optional[str],
    ) -> List[bytes]:
        """One serialized payload per TP rank: each rank deserializes only its
        own copy, so producer-side CUDA-IPC refcounts drop cleanly after every
        load. flattened_bucket callers pass pre-serialized per-rank payloads."""
        if load_format == "flattened_bucket":
            return normalize_serialized_named_tensor_payloads(
                cast(List[SerializedTensorPayload], tensors)
            )
        else:
            return [
                MultiprocessingSerializer.serialize(tensors)
                for _ in range(self.server_args.tp_size)
            ]

    def load_lora_adapter_from_tensors(
        self,
        lora_name: str,
        tensors: Union[Dict[str, torch.Tensor], List[SerializedTensorPayload]],
        config_dict: Dict,
        load_format: Optional[str] = None,
    ):
        serialized_named_tensors = self._serialize_tensors_per_rank(
            tensors, load_format
        )
        lora_req = LoadLoRAAdapterFromTensorsReqInput(
            lora_name=lora_name,
            config_dict=config_dict,
            serialized_named_tensors=serialized_named_tensors,
            load_format=load_format,
        )
        return self.loop.run_until_complete(
            self.tokenizer_manager.load_lora_adapter_from_tensors(lora_req, None)
        )

    def load_lora_adapter(self, lora_name: str, lora_path: str, pinned: bool = False):
        """Load a new LoRA adapter without re-launching the engine."""

        obj = LoadLoRAAdapterReqInput(
            lora_name=lora_name,
            lora_path=lora_path,
            pinned=pinned,
        )

        return self.loop.run_until_complete(
            self.tokenizer_manager.load_lora_adapter(obj, None)
        )

    def unload_lora_adapter(self, lora_name: str):
        """Unload a LoRA adapter without re-launching the engine."""

        obj = UnloadLoRAAdapterReqInput(lora_name=lora_name)

        return self.loop.run_until_complete(
            self.tokenizer_manager.unload_lora_adapter(obj, None)
        )

    async def async_load_lora_adapter(
        self, lora_name: str, lora_path: str, pinned: bool = False
    ):
        """
        Asynchronous version of load_lora_adapter.

        See load_lora_adapter() for detailed documentation.
        """

        obj = LoadLoRAAdapterReqInput(
            lora_name=lora_name,
            lora_path=lora_path,
            pinned=pinned,
        )

        return await self.tokenizer_manager.load_lora_adapter(obj, None)

    async def async_unload_lora_adapter(self, lora_name: str):
        """
        Asynchronous version of unload_lora_adapter.

        See unload_lora_adapter() for detailed documentation.
        """

        obj = UnloadLoRAAdapterReqInput(lora_name=lora_name)

        return await self.tokenizer_manager.unload_lora_adapter(obj, None)

    def release_memory_occupation(self, tags: Optional[List[str]] = None):
        obj = ReleaseMemoryOccupationReqInput(tags=tags)
        return self.loop.run_until_complete(
            self.tokenizer_manager.release_memory_occupation(obj, None)
        )

    def resume_memory_occupation(self, tags: Optional[List[str]] = None):
        obj = ResumeMemoryOccupationReqInput(tags=tags)
        return self.loop.run_until_complete(
            self.tokenizer_manager.resume_memory_occupation(obj, None)
        )

    def freeze_gc(self):
        """
        To maintain a high performance server with low latency, we want to reduce the
        stalls caused by the garbage collector scanning through a large number of objects.

        It is usually helpful to start the server and warm it up with real requests to
        initialize many of the long-lived objects that do not need to be garbage collected.

        After sufficient warmup, we can call this function to freeze the garbage collector
        so that all objects created before this point are considered out of scope for garbage
        collection.
        """

        self.loop.run_until_complete(self.tokenizer_manager.freeze_gc())

    """
    Execute an RPC call on all scheduler processes.
    """

    def collective_rpc(self, method: str, **kwargs):
        obj = RpcReqInput(method=method, parameters=kwargs)
        sock_send(self.send_to_rpc, obj)
        recv_req = sock_recv(self.send_to_rpc, flags=zmq.BLOCKY)
        assert isinstance(recv_req, RpcReqOutput)
        assert recv_req.success, recv_req.message

    def save_remote_model(self, **kwargs):
        self.collective_rpc("save_remote_model", **kwargs)

    def save_sharded_model(self, **kwargs):
        self.collective_rpc("save_sharded_model", **kwargs)

    # score() and async_score() are provided by EngineScoreMixin


def _set_envs_and_config(server_args: ServerArgs):
    # Set global environments
    # MNNVL fabric (GB200/GB300) multi-node: cross-node NVLink needs NCCL's
    # cuMem-based buffers and MNNVL transport. Default them on (user-set
    # values win; the symm-mem override below only fires when unset).
    if server_args.nnodes > 1 and is_mnnvl_fabric_device():
        os.environ.setdefault("NCCL_CUMEM_ENABLE", "1")
        os.environ.setdefault("NCCL_MNNVL_ENABLE", "1")
    if "NCCL_CUMEM_ENABLE" not in os.environ or server_args.enable_symm_mem:
        os.environ["NCCL_CUMEM_ENABLE"] = str(int(server_args.enable_symm_mem))
    if (
        "NCCL_NVLS_ENABLE" not in os.environ
        or server_args.enable_nccl_nvls
        or server_args.enable_symm_mem
    ):
        os.environ["NCCL_NVLS_ENABLE"] = str(
            int(server_args.enable_nccl_nvls or server_args.enable_symm_mem)
        )
    if "NCCL_GRAPH_MIXING_SUPPORT" not in os.environ or server_args.enable_symm_mem:
        # Note(wh): NCCL_GRAPH_MIXING_SUPPORT=0 can help improve performance for symmetric kernels.
        # details in https://github.com/NVIDIA/nccl-tests/issues/333#issuecomment-3103636985
        if server_args.dcp_size > 1:
            os.environ["NCCL_GRAPH_MIXING_SUPPORT"] = "0"
    os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] = "8"
    os.environ["CUDA_MODULE_LOADING"] = "AUTO"

    if os.environ.get("TRTLLM_ENABLE_PDL", "1") != "0":
        # flashinfer uses this environment variable for various kernels from MoE to quant kernels
        os.environ["TRTLLM_ENABLE_PDL"] = "1"

    if os.environ.get("CUTE_DSL_LOG_LEVEL") is None:
        # Default to warning level, to avoid too many logs
        os.environ["CUTE_DSL_LOG_LEVEL"] = "30"

    if os.environ.get("CUTE_DSL_LOG_TO_CONSOLE") is None:
        # Need to set log to console, otherwise the log level won't take effect
        os.environ["CUTE_DSL_LOG_TO_CONSOLE"] = "1"

    # Can also be passed as argument
    os.environ["SGLANG_RUN_ID"] = (
        f"sglang-run-{time.time()}-{random.randint(0, 100000000)}"
    )

    # Set prometheus env vars
    if server_args.enable_metrics:
        set_prometheus_multiproc_dir()

    # Set ulimit
    set_ulimit()

    # Check flashinfer version
    if not get_bool_env_var("SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK"):
        if server_args.attention_backend == "flashinfer":
            assert_pkg_version(
                "flashinfer_python",
                "0.6.15.post1",
                "Please uninstall the old version and "
                "reinstall the latest version by following the instructions "
                "at https://docs.flashinfer.ai/installation.html.",
            )
        if _is_cuda:
            assert_pkg_version(
                "sglang-kernel",
                "0.4.6.post1",
                "Please reinstall the latest version with `pip install sglang-kernel --force-reinstall`",
            )

    # Signal handlers can only be registered from the main thread.
    if threading.current_thread() is threading.main_thread():
        if server_args.custom_sigquit_handler is None:
            # Register the signal handler.
            # The child processes will send SIGQUIT to this process when any error happens
            # This process then clean up the whole process tree
            # Note: This sigquit handler is used in the launch phase, and may be replaced by
            # the running_phase_sigquit_handler in the tokenizer manager after the grpc server is launched.
            def launch_phase_sigquit_handler(signum, frame):
                logger.error(
                    "Received sigquit from a child process. It usually means the child failed."
                )
                kill_process_tree(os.getpid())

            signal.signal(signal.SIGQUIT, launch_phase_sigquit_handler)
        else:
            # Allow users to register a custom SIGQUIT handler for things like crash dump
            logger.error(
                f"Using custom SIGQUIT handler: {server_args.custom_sigquit_handler}"
            )
            signal.signal(signal.SIGQUIT, server_args.custom_sigquit_handler)
    else:
        logger.warning(
            "Signal handler is not added because the engine is not in the "
            "main thread. This disables the SIGQUIT handler for cleaning up "
            "the process tree when a child process fails."
        )

    # Set mp start method
    mp.set_start_method("spawn", force=True)

    # Set gc threshold
    if gc_threshold := server_args.gc_threshold:
        gc.set_threshold(*gc_threshold)

    _log_legacy_kernel_cache_dirs()


def _log_legacy_kernel_cache_dirs():
    """Note the pre-SGLANG_CACHE_DIR cache dirs without touching them: other
    frameworks on the box may still be using them."""
    # TODO(shuwang21): drop once SGLANG_CACHE_DIR has been the default for a
    # few releases.
    legacy_dirs = [
        d
        for d in (
            os.path.expanduser("~/.triton"),
            os.path.expanduser("~/.cache/flashinfer"),
            os.path.expanduser("~/.cache/deep_gemm"),
        )
        if os.path.isdir(d)
    ]
    if not legacy_dirs:
        return
    logger.info(
        "Compiled-kernel caches now live under SGLANG_CACHE_DIR (%s). These "
        "older directories are no longer used by sglang, but may still be "
        "used by other frameworks on this machine, so they were left alone: "
        "%s. Remove them yourself if nothing else needs them.",
        envs.SGLANG_CACHE_DIR.get(),
        ", ".join(legacy_dirs),
    )


def _scheduler_died_error(rank: int, proc) -> RuntimeError:
    """Build a descriptive error for a scheduler process that died during init."""
    proc.join(timeout=10)
    return RuntimeError(
        f"Rank {rank} scheduler died during initialization "
        f"(exit code: {proc.exitcode}). "
        f"If exit code is -9 (SIGKILL), a common cause is the OS OOM killer. "
        f"Run `dmesg -T | grep -i oom` to check."
    )


def _wait_for_scheduler_ready(
    scheduler_pipe_readers: List,
    scheduler_procs: List,
) -> List[Dict]:
    """Wait for the model to finish loading and return scheduler infos.

    Uses poll() with timeout instead of blocking recv(), so that child process
    death (e.g. OOM SIGKILL) is detected promptly instead of hanging forever.
    """
    scheduler_infos = []
    for i in range(len(scheduler_pipe_readers)):
        while True:
            if scheduler_pipe_readers[i].poll(timeout=5.0):
                try:
                    data = scheduler_pipe_readers[i].recv()
                except EOFError:
                    raise _scheduler_died_error(i, scheduler_procs[i])
                if data["status"] != "ready":
                    raise RuntimeError(
                        "Initialization failed. Please see the error messages above."
                    )
                scheduler_infos.append(data)
                break

            # Poll timed out — check all processes for early death
            for j in range(len(scheduler_procs)):
                if not scheduler_procs[j].is_alive():
                    raise _scheduler_died_error(j, scheduler_procs[j])

    return scheduler_infos


def _calculate_rank_ranges(
    nnodes: int, pp_size: int, tp_size: int, node_rank: int
) -> Tuple[range, range, int, int]:
    """Calculate pp_rank_range and tp_rank_range for a given node.

    Args:
        nnodes: Total number of nodes.
        pp_size: Pipeline parallel size.
        tp_size: Tensor parallel size.
        node_rank: The rank of the node to compute ranges for.

    Returns:
        A tuple of (pp_rank_range, tp_rank_range, pp_size_per_node, tp_size_per_node):
        - pp_rank_range: range of pipeline-parallel ranks assigned to this node.
        - tp_rank_range: range of tensor-parallel ranks assigned to this node.
        - pp_size_per_node: number of PP ranks per node.
        - tp_size_per_node: number of TP ranks per node.
    """
    pp_size_per_node = max(pp_size // nnodes, 1)
    nnodes_per_pp_rank = max(nnodes // pp_size, 1)
    pp_rank_range = range(
        pp_size_per_node * (node_rank // nnodes_per_pp_rank),
        pp_size_per_node * (node_rank // nnodes_per_pp_rank + 1),
    )

    nnodes_per_tp_group = nnodes_per_pp_rank
    tp_size_per_node = tp_size // nnodes_per_tp_group
    tp_rank_range = range(
        tp_size_per_node * (node_rank % nnodes_per_tp_group),
        tp_size_per_node * (node_rank % nnodes_per_tp_group + 1),
    )

    return pp_rank_range, tp_rank_range, pp_size_per_node, tp_size_per_node


def _compute_parallelism_ranks(
    server_args: ServerArgs, tp_rank: int
) -> Tuple[int, int, int]:
    """Compute attention-CP, MoE-DP, and MoE-EP ranks for a TP rank."""
    attn_dp_size = server_args.dp_size if server_args.enable_dp_attention else 1

    # Parallelism hierarchy (outermost to innermost):
    # - Attention: Global(TP) -> DP -> ATTN_CP -> ATTN_TP (innermost)
    # - MoE: Global(TP) -> MOE_DP -> EP -> MOE_TP (innermost)
    attn_tp_size = server_args.tp_size // attn_dp_size // server_args.attn_cp_size
    attn_cp_rank = (tp_rank // attn_tp_size) % server_args.attn_cp_size
    moe_dp_rank = tp_rank // (server_args.tp_size // server_args.moe_dp_size)
    moe_ep_rank = (
        tp_rank
        % (server_args.tp_size // server_args.moe_dp_size)
        // (server_args.tp_size // server_args.moe_dp_size // server_args.ep_size)
    )
    return attn_cp_rank, moe_dp_rank, moe_ep_rank
