Add pd disaggregated inference (#3558)

* Initial PD disaggregation implementation

Test2

Internal IP Test

Add worker with internal_ip

Check status and register

Add Status Ready Log

Add Prefill-Decode

Add PD to dstack

Test register worker without poll

Add router config in service config

Update remove worker

Clean Up router code

Clean Up

Further Cleanup

* Add pd disaggregation service

* Move router configuration to service

* Resolve major comments

* Resolve Lint Error

* Minor Update

* Resolve Minor Comments

* Update wheel url

* Resolve backward incompatibility

* Update RouterConfigs

* Resolve Lint Error

* Update gateway wheel

* Minor Update

---------

Co-authored-by: Bihan  Rana

Author: Bihan

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,