sgl-project--sglang/python/sglang/multimodal_gen/runtime/server_args/disagg.py

# SPDX-License-Identifier: Apache-2.0
"""Disaggregated diffusion server argument helpers."""

from __future__ import annotations

from typing import ClassVar, Literal

from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
from sglang.multimodal_gen.runtime.utils.common import (
    format_tcp_endpoint,
    parse_tcp_host_port,
)
from sglang.multimodal_gen.utils import FlexibleArgumentParser


class DisaggServerArgsMixin:
    DISAGG_RESULT_PORT_OFFSETS: ClassVar[dict[RoleType, int]] = {
        RoleType.ENCODER: 1,
        RoleType.DENOISER: 2,
        RoleType.DECODER: 3,
    }

    def get_role_parallelism(self, role_type: RoleType) -> dict[str, int | None]:
        _none = {
            "tp_size": None,
            "sp_degree": None,
            "ulysses_degree": None,
            "ring_degree": None,
        }
        if role_type == RoleType.ENCODER:
            return {**_none, "tp_size": self.encoder_tp}
        if role_type == RoleType.DENOISER:
            return {
                "tp_size": self.denoiser_tp,
                "sp_degree": self.denoiser_sp,
                "ulysses_degree": self.denoiser_ulysses,
                "ring_degree": self.denoiser_ring,
            }
        if role_type == RoleType.DECODER:
            return {**_none, "sp_degree": self.decoder_sp}
        return _none

    def derive_pool_result_endpoint(self) -> str:
        host, base_port = parse_tcp_host_port(
            self.disagg_server_addr, "disagg_server_addr"
        )
        role = (
            self.disagg_role
            if isinstance(self.disagg_role, RoleType)
            else RoleType.from_string(self.disagg_role)
        )
        try:
            offset = self.DISAGG_RESULT_PORT_OFFSETS[role]
        except KeyError as exc:
            raise ValueError(
                "pool result endpoints are only defined for encoder, denoiser, "
                f"and decoder roles, got {role.value!r}"
            ) from exc
        return format_tcp_endpoint(host, base_port + offset, "pool_result_endpoint")

    def derive_pool_work_endpoint(self) -> str:
        return format_tcp_endpoint("0.0.0.0", self.scheduler_port, "pool_work_endpoint")

    def derive_pool_control_endpoint(self) -> str:
        return format_tcp_endpoint(
            "0.0.0.0", self.scheduler_port + 1, "pool_control_endpoint"
        )

    def derive_pool_control_advertised_endpoint(self) -> str:
        host = self.host or self.disagg_p2p_hostname or "127.0.0.1"
        if host == "0.0.0.0":
            host = self.disagg_p2p_hostname or "127.0.0.1"
        return format_tcp_endpoint(
            host, self.scheduler_port + 1, "pool_control_advertised_endpoint"
        )

    def resolved_role_device(self) -> Literal["cpu", "cuda"]:
        if self.disagg_role_device == "auto":
            return "cpu" if self.num_gpus <= 0 else "cuda"
        return self.disagg_role_device

    @classmethod
    def add_disagg_cli_args(cls, parser: FlexibleArgumentParser) -> None:
        role_default = (
            cls.disagg_role.value
            if isinstance(cls.disagg_role, RoleType)
            else cls.disagg_role
        )
        parser.add_argument(
            "--disagg-role",
            type=str,
            default=role_default,
            choices=RoleType.choices(),
            help="Role for disaggregated pipeline.",
        )
        parser.add_argument(
            "--disagg-timeout",
            type=int,
            default=cls.disagg_timeout,
            help="Timeout in seconds for pending disagg requests. "
            f"Default: {cls.disagg_timeout}.",
        )
        parser.add_argument(
            "--disagg-downstream-wait-timeout",
            type=int,
            default=cls.disagg_downstream_wait_timeout,
            help="Timeout in seconds while waiting for a downstream role slot. "
            f"Default: {cls.disagg_downstream_wait_timeout}.",
        )
        parser.add_argument(
            "--disagg-dispatch-policy",
            type=str,
            default=cls.disagg_dispatch_policy,
            choices=["round_robin", "max_free_slots"],
            help="Dispatch policy for pool mode disagg routing.",
        )
        parser.add_argument(
            "--disagg-instance-id",
            type=int,
            default=cls.disagg_instance_id,
            help="Stable per-role instance ID used by DiffusionServer registration.",
        )
        parser.add_argument(
            "--disagg-max-slots-per-instance",
            type=int,
            default=cls.disagg_max_slots_per_instance,
            help="Maximum concurrent transfer/computation slots tracked per instance.",
        )
        parser.add_argument(
            "--disagg-transfer-redundancy",
            type=float,
            default=cls.disagg_transfer_redundancy,
            help="Redundancy factor used when sizing transfer buffers from warmup.",
        )
        parser.add_argument(
            "--disagg-role-device",
            type=str,
            default=cls.disagg_role_device,
            choices=["auto", "cpu", "cuda"],
            help=(
                "Per-role device override. 'cpu' is intended for same-machine "
                "encoder roles."
            ),
        )
        parser.add_argument(
            "--disagg-transfer-backend",
            type=str,
            default=cls.disagg_transfer_backend,
            choices=["auto", "mock", "mooncake"],
            help="Transfer backend for multimodal diffusion disaggregation.",
        )
        parser.add_argument(
            "--disagg-transfer-pool-size",
            type=int,
            default=cls.disagg_transfer_pool_size,
            help="Size of the P2P transfer buffer pool in bytes.",
        )
        parser.add_argument(
            "--disagg-transfer-pin-memory",
            type=str,
            default=cls.disagg_transfer_pin_memory,
            choices=["auto", "off", "required"],
            help="CUDA host-register same-host shared-memory transfer buffers.",
        )
        parser.add_argument(
            "--disagg-p2p-hostname",
            type=str,
            default=cls.disagg_p2p_hostname,
            help="Hostname for P2P transfer engine.",
        )
        parser.add_argument(
            "--disagg-ib-device",
            type=str,
            default=cls.disagg_ib_device,
            help="InfiniBand device for P2P RDMA transfers.",
        )
        parser.add_argument(
            "--disagg-server-addr",
            type=str,
            default=cls.disagg_server_addr,
            help="DiffusionServer head node address for per-role launch mode.",
        )
        parser.add_argument(
            "--encoder-urls",
            type=str,
            default=cls.encoder_urls,
            help="Encoder instance work endpoints for DiffusionServer head mode.",
        )
        parser.add_argument(
            "--denoiser-urls",
            type=str,
            default=cls.denoiser_urls,
            help="Denoiser instance work endpoints for DiffusionServer head mode.",
        )
        parser.add_argument(
            "--decoder-urls",
            type=str,
            default=cls.decoder_urls,
            help="Decoder instance work endpoints for DiffusionServer head mode.",
        )
        parser.add_argument(
            "--encoder-tp",
            type=int,
            default=cls.encoder_tp,
            help="Tensor parallelism for encoder role.",
        )
        parser.add_argument(
            "--denoiser-tp",
            type=int,
            default=cls.denoiser_tp,
            help="Tensor parallelism for denoiser role.",
        )
        parser.add_argument(
            "--denoiser-sp",
            type=int,
            default=cls.denoiser_sp,
            help="Sequence parallelism for denoiser role.",
        )
        parser.add_argument(
            "--denoiser-ulysses",
            type=int,
            default=cls.denoiser_ulysses,
            help="Ulysses SP degree for denoiser role.",
        )
        parser.add_argument(
            "--denoiser-ring",
            type=int,
            default=cls.denoiser_ring,
            help="Ring SP degree for denoiser role.",
        )
        parser.add_argument(
            "--decoder-sp",
            type=int,
            default=cls.decoder_sp,
            help="Sequence parallelism for decoder role.",
        )
        parser.add_argument(
            "--decoder-tp",
            type=int,
            default=cls.decoder_tp,
            help="Deprecated alias for --decoder-sp.",
        )