Add pd disaggregated inference

docs/docs/reference/dstack.yml/gateway.md

@@ -14,7 +14,7 @@ The `gateway` configuration type allows creating and updating [gateways](../../c
 
 === "SGLang Model Gateway"
 
-    #SCHEMA# dstack._internal.core.models.routers.SGLangRouterConfig
+    #SCHEMA# dstack._internal.core.models.routers.SGLangGatewayRouterConfig
         overrides:
           show_root_heading: false
           type:

gateway/pyproject.toml

@@ -15,7 +15,7 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
-sglang = ["sglang-router==0.2.1"]
+sglang = ["sglang-router==0.3.2"]
 
 [tool.setuptools.package-data]
 "dstack.gateway" = [

src/dstack/_internal/core/backends/base/compute.py

@@ -39,7 +39,7 @@
     SSHKey,
 )
 from dstack._internal.core.models.placement import PlacementGroup, PlacementGroupProvisioningData
-from dstack._internal.core.models.routers import AnyRouterConfig
+from dstack._internal.core.models.routers import AnyGatewayRouterConfig
 from dstack._internal.core.models.runs import Job, JobProvisioningData, Requirements, Run
 from dstack._internal.core.models.volumes import (
     Volume,
@@ -924,7 +924,9 @@ def get_run_shim_script(
     ]
 
 
-def get_gateway_user_data(authorized_key: str, router: Optional[AnyRouterConfig] = None) -> str:
+def get_gateway_user_data(
+    authorized_key: str, router: Optional[AnyGatewayRouterConfig] = None
+) -> str:
     return get_cloud_config(
         package_update=True,
         packages=[
@@ -1036,7 +1038,7 @@ def get_latest_runner_build() -> Optional[str]:
     return None
 
 
-def get_dstack_gateway_wheel(build: str, router: Optional[AnyRouterConfig] = None) -> str:
+def get_dstack_gateway_wheel(build: str, router: Optional[AnyGatewayRouterConfig] = None) -> str:
     channel = "release" if settings.DSTACK_RELEASE else "stgn"
     base_url = f"https://dstack-gateway-downloads.s3.amazonaws.com/{channel}"
     if build == "latest":
@@ -1049,7 +1051,7 @@ def get_dstack_gateway_wheel(build: str, router: Optional[AnyRouterConfig] = Non
     return f"dstack-gateway @ {wheel}"
 
 
-def get_dstack_gateway_commands(router: Optional[AnyRouterConfig] = None) -> List[str]:
+def get_dstack_gateway_commands(router: Optional[AnyGatewayRouterConfig] = None) -> List[str]:
     build = get_dstack_runner_version() or "latest"
     gateway_package = get_dstack_gateway_wheel(build, router)
     return [

src/dstack/_internal/core/backends/kubernetes/compute.py

@@ -66,7 +66,7 @@
 )
 from dstack._internal.core.models.placement import PlacementGroup
 from dstack._internal.core.models.resources import CPUSpec, GPUSpec
-from dstack._internal.core.models.routers import AnyRouterConfig
+from dstack._internal.core.models.routers import AnyGatewayRouterConfig
 from dstack._internal.core.models.runs import Job, JobProvisioningData, Requirements, Run
 from dstack._internal.core.models.volumes import Volume
 from dstack._internal.utils.common import get_or_error
@@ -864,7 +864,7 @@ def _wait_for_load_balancer_address(
 
 
 def _get_gateway_commands(
-    authorized_keys: List[str], router: Optional[AnyRouterConfig] = None
+    authorized_keys: List[str], router: Optional[AnyGatewayRouterConfig] = None
 ) -> List[str]:
     authorized_keys_content = "\n".join(authorized_keys).strip()
     gateway_commands = " && ".join(get_dstack_gateway_commands(router=router))

src/dstack/_internal/core/compatibility/gateways.py

@@ -31,9 +31,7 @@ def _get_gateway_configuration_excludes(
 ) -> IncludeExcludeDictType:
     configuration_excludes: IncludeExcludeDictType = {}
 
-    # Add excludes like this:
-    #
-    # if configuration.tags is None:
-    #     configuration_excludes["tags"] = True
+    if configuration.router is None:
+        configuration_excludes["router"] = True
 
     return configuration_excludes

src/dstack/_internal/core/compatibility/runs.py

@@ -2,6 +2,7 @@
 
 from dstack._internal.core.models.common import IncludeExcludeDictType, IncludeExcludeSetType
 from dstack._internal.core.models.configurations import ServiceConfiguration
+from dstack._internal.core.models.routers import SGLangServiceRouterConfig
 from dstack._internal.core.models.runs import (
     DEFAULT_PROBE_UNTIL_READY,
     DEFAULT_REPLICA_GROUP_NAME,
@@ -72,6 +73,12 @@ def get_run_spec_excludes(run_spec: RunSpec) -> IncludeExcludeDictType:
             # Servers prior to 0.20.8 do not support probes=None
             configuration_excludes["probes"] = True
 
+        router = run_spec.configuration.router
+        if router is None:
+            configuration_excludes["router"] = True
+        elif isinstance(router, SGLangServiceRouterConfig) and router.pd_disaggregation is False:
+            configuration_excludes["router"] = {"pd_disaggregation": True}
+
     if configuration_excludes:
         spec_excludes["configuration"] = configuration_excludes
     if profile_excludes:

src/dstack/_internal/core/models/configurations.py

@@ -28,6 +28,7 @@
     parse_off_duration,
 )
 from dstack._internal.core.models.resources import Range, ResourcesSpec
+from dstack._internal.core.models.routers import AnyServiceRouterConfig
 from dstack._internal.core.models.services import AnyModel, OpenAIChatModel
 from dstack._internal.core.models.unix import UnixUser
 from dstack._internal.core.models.volumes import MountPoint, VolumeConfiguration, parse_mount_point
@@ -887,6 +888,14 @@ class ServiceConfigurationParams(CoreModel):
             )
         ),
     ] = None
+    router: Annotated[
+        Optional[AnyServiceRouterConfig],
+        Field(
+            description=(
+                "Router configuration for the service. Requires a gateway with matching router enabled. "
+            ),
+        ),
+    ] = None
 
     @validator("port")
     def convert_port(cls, v) -> PortMapping:

src/dstack/_internal/core/models/gateways.py

@@ -8,7 +8,7 @@
 
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.common import CoreModel
-from dstack._internal.core.models.routers import AnyRouterConfig
+from dstack._internal.core.models.routers import AnyGatewayRouterConfig
 from dstack._internal.utils.tags import tags_validator
 
 
@@ -63,8 +63,13 @@ class GatewayConfiguration(CoreModel):
         ),
     ] = None
     router: Annotated[
-        Optional[AnyRouterConfig],
-        Field(description="The router configuration"),
+        Optional[AnyGatewayRouterConfig],
+        Field(
+            description=(
+                "The router configuration for this gateway. "
+                "E.g. `{ type: sglang, policy: round_robin }`."
+            ),
+        ),
     ] = None
     domain: Annotated[
         Optional[str], Field(description="The gateway domain, e.g. `example.com`")
@@ -134,7 +139,7 @@ class GatewayComputeConfiguration(CoreModel):
     ssh_key_pub: str
     certificate: Optional[AnyGatewayCertificate] = None
     tags: Optional[Dict[str, str]] = None
-    router: Optional[AnyRouterConfig] = None
+    router: Optional[AnyGatewayRouterConfig] = None
 
 
 class GatewayProvisioningData(CoreModel):

src/dstack/_internal/core/models/routers.py

@@ -11,14 +11,37 @@ class RouterType(str, Enum):
     SGLANG = "sglang"
 
 
-class SGLangRouterConfig(CoreModel):
+class SGLangGatewayRouterConfig(CoreModel):
+    """Gateway-level router configuration. type and policy only. pd_disaggregation is service-level."""
+
+    type: Annotated[
+        Literal["sglang"],
+        Field(description="The router type enabled on this gateway."),
+    ] = "sglang"
+    policy: Annotated[
+        Literal["random", "round_robin", "cache_aware", "power_of_two"],
+        Field(
+            description=(
+                "The routing policy. Deprecated: prefer setting policy in the service's router config. "
+                "Options: `random`, `round_robin`, `cache_aware`, `power_of_two`"
+            ),
+        ),
+    ] = "cache_aware"
+
+
+class SGLangServiceRouterConfig(CoreModel):
     type: Annotated[Literal["sglang"], Field(description="The router type")] = "sglang"
     policy: Annotated[
         Literal["random", "round_robin", "cache_aware", "power_of_two"],
         Field(
             description="The routing policy. Options: `random`, `round_robin`, `cache_aware`, `power_of_two`"
         ),
     ] = "cache_aware"
+    pd_disaggregation: Annotated[
+        bool,
+        Field(description="Enable PD disaggregation mode for the SGLang router"),
+    ] = False
 
 
-AnyRouterConfig = SGLangRouterConfig
+AnyServiceRouterConfig = SGLangServiceRouterConfig
+AnyGatewayRouterConfig = SGLangGatewayRouterConfig

src/dstack/_internal/proxy/gateway/routers/registry.py

@@ -80,6 +80,7 @@ async def register_replica(
         ssh_proxy=body.ssh_proxy,
         ssh_head_proxy=body.ssh_head_proxy,
         ssh_head_proxy_private_key=body.ssh_head_proxy_private_key,
+        internal_ip=body.internal_ip,
         repo=repo,
         nginx=nginx,
         service_conn_pool=service_conn_pool,

src/dstack/_internal/proxy/gateway/schemas/registry.py

@@ -3,7 +3,7 @@
 from pydantic import BaseModel, Field
 
 from dstack._internal.core.models.instances import SSHConnectionParams
-from dstack._internal.core.models.routers import AnyRouterConfig
+from dstack._internal.core.models.routers import AnyServiceRouterConfig
 from dstack._internal.proxy.lib.models import RateLimit
 
 
@@ -45,7 +45,7 @@ class RegisterServiceRequest(BaseModel):
     options: Options
     ssh_private_key: str
     rate_limits: tuple[RateLimit, ...] = ()
-    router: Optional[AnyRouterConfig] = None
+    router: Optional[AnyServiceRouterConfig] = None
 
 
 class RegisterReplicaRequest(BaseModel):
@@ -56,6 +56,7 @@ class RegisterReplicaRequest(BaseModel):
     ssh_proxy: Optional[SSHConnectionParams]
     ssh_head_proxy: Optional[SSHConnectionParams]
     ssh_head_proxy_private_key: Optional[str]
+    internal_ip: Optional[str] = None
 
 
 class RegisterEntrypointRequest(BaseModel):

src/dstack/_internal/proxy/gateway/services/model_routers/__init__.py

@@ -1,11 +1,11 @@
-from dstack._internal.core.models.routers import AnyRouterConfig, RouterType
+from dstack._internal.core.models.routers import AnyServiceRouterConfig, RouterType
 from dstack._internal.proxy.gateway.services.model_routers.sglang import SglangRouter
 from dstack._internal.proxy.lib.errors import ProxyError
 
 from .base import Router, RouterContext
 
 
-def get_router(router: AnyRouterConfig, context: RouterContext) -> Router:
+def get_router(router: AnyServiceRouterConfig, context: RouterContext) -> Router:
     if router.type == RouterType.SGLANG:
         return SglangRouter(config=router, context=context)
     raise ProxyError(f"Router type '{router.type}' is not available")

src/dstack/_internal/proxy/gateway/services/model_routers/base.py

@@ -4,7 +4,7 @@
 
 from pydantic import BaseModel
 
-from dstack._internal.core.models.routers import AnyRouterConfig
+from dstack._internal.core.models.routers import AnyServiceRouterConfig
 
 
 class RouterContext(BaseModel):
@@ -29,7 +29,7 @@ class Router(ABC):
     def __init__(
         self,
         context: RouterContext,
-        config: Optional[AnyRouterConfig] = None,
+        config: Optional[AnyServiceRouterConfig] = None,
     ):
         """Initialize router with context.