Add pd disaggregated inference

[Bihan]

Testing Steps

1. Create (CPU node) in K8s cluster

2. Create gateway in the CPU node using below config

type: gateway
name: bihan-gateway

backend: kubernetes
region: any

domain: bihan-gateway.dstack.ai
router: sglang

3. Create GPU-node with 3 instances (1 Prefill, 1 Decode and 1 for testing scaling) in the same K8s cluster where gateway node exists.
   Note: See design doc for details on why the gateway and workers are required to be on the same network.

4. Apply below prefill-decode service configuration

type: service
name: prefill-decode
image: lmsysorg/sglang:latest

env:
  - HF_TOKEN
  - MODEL_ID=meta-llama/Llama-3.2-1B-Instruct

replicas:
  - count: 1..2
    scaling:
      metric: rps
      target: 3
    commands:
      - |
          python -m sglang.launch_server \
            --model-path $MODEL_ID \
            --disaggregation-mode prefill \
            --disaggregation-transfer-backend mooncake \
            --host 0.0.0.0 \
            --port 8000 \
            --disaggregation-bootstrap-port 8998 \
            --log-level debug \
            > worker-server.log 2>&1
    resources:
      gpu: 1

  - count: 1
    commands:
      - |
          python -m sglang.launch_server \
            --model-path $MODEL_ID \
            --disaggregation-mode decode \
            --disaggregation-transfer-backend mooncake \
            --host 0.0.0.0 \
            --port 8000 \
            --log-level debug \
            > worker-server.log 2>&1
    resources:
      gpu: 1

port: 8000
model: meta-llama/Llama-3.2-1B-Instruct

probes:
  - type: http
    url: /health_generate
    interval: 15s

router:
  type: sglang
  policy: round_robin
  pd_disaggregation: true

5. When rps>=3 prefill replica scales to 2.

Note: For testing you need to assign wheel to https://bihan-test-bucket.s3.eu-west-1.amazonaws.com/dstack_gateway-0.0.1-py3-none-any.whl

[jvstme]

This adds pd_disaggregation to both service configurations and gateway configurations. I'd advocate for adding it to service configurations only.

• Whether or not the service is configured to use PD disaggregation is clearly a service property, because it depends on the replica groups configuration. I don't think many users would want to configure that property at the gateway level, making assumptions about what services are going to run on that gateway in the future.

• Having two places for the same property complicates the interface — you'd have to explain in the docs how these places are related to each other, when and how one setting overrides the other, etc.

• Having the property at the gateway level can potentially complicate further development — that way, you can only tell whether a service is using PD disaggregation if the service is already registered, and you need to fetch the GatewayModel object from the database to do so. For example, this would complicate adding the default probe for services with PD disaggregation.

[Bihan]

Done

[jvstme]

(nit) If it now accepts service, it no longer has to accept domain. The domain is available as service.domain_safe

[jvstme]

(nit) Attribute access via getattr is not indexed by the IDE. I think it's almost never a good idea to use it, unless there's no other option.

You could do service.router.pd_disaggregation like in other places.

If you want to make it extra clear that the attribute exists, you could also do isinstance(service.router, SGLangRouterConfig), although it seems to be redundant until we have other router types.

[jvstme]

This change isn't backward compatible — if I have a previously created gateway with router: type: sglang, dstack fails to start new services:

[09:10:10] INFO     dstack._internal.server.services.events:205 Emitting event: Run submitted. Status: SUBMITTED. Event targets:                        
                    run(99e05f)pretty-insect-1. Actor: user(aeafc8)admin                                                                                
ERROR:    Exception in ASGI application
  + Exception Group Traceback (most recent call last):
[...]
  File "pydantic/main.py", line 364, in pydantic.main.BaseModel.__init__
pydantic.error_wrappers.ValidationError: 1 validation error for GatewayConfigurationResponse
router
  value is not a valid enumeration member; permitted: 'sglang' (type=type_error.enum; enum_values=[<RouterType.SGLANG: 'sglang'>])

And to delete the gateway:

$ dstack gateway delete -y sglang-gateway
Unexpected error: status code 500 when requesting http://localhost:3000/api/project/main/gateways/get. Check the server logs for backend issues, and the
CLI logs at (~/.dstack/logs/cli/latest.log) local CLI output

Apart from backward compatibility, having router: type: sglang is useful in case we ever support specifying any gateway-level router properties, such as the SGLang package version.

router:
  type: sglang
  version: 0.3.2

I can suggest to introduce two different pydantic models for service-level (type, policy, pd_disaggregation) and gateway-level (type, policy) router configuration. Before we can drop the gateway-level router.policy in a future version, we mention that it is deprecated in its description

[jvstme]

Does that mean the SGLang router will not be enabled if the service config does not specify router: type: sglang? I'm not sure if this is expected or not.

• On one hand, I like these semantics more:
  • To use the SGLang router, you explicitly request it in the service configuration. This prevents misconfiguration cases when a non-SGLang service gets assigned to a gateway with SGLang enabled.
  • This allows to use the same gateway for both SGLang and non-SGLang services, without the need to create two separate gateways.
• On the other hand, this change of semantics is not backward-compatible — if someone is already using SGLang-enabled gateways, they expect SGLang to be enabled automatically.

Maybe the safest option for now is to keep the old semantics (automatically enable SGLang if specified in the gateway configuration). Later, we can discuss the change of semantics with the team (cc @peterschmidt85) and decide if it's worth it

[Bihan]

@jvstme
cc @peterschmidt85

Maybe we keep GatewayConfiguration's router to Optional[AnyRouterConfig] with properties type and policy without pd_disaggregation as before. This will eliminate backward compatibility issue.

Create new GatewayConfig class as

class GatewayConfig(CoreModel):
    name: Annotated[
        Optional[str],
        Field(description="The name of the gateway. Omit for the default gateway."),
    ] = None
    pd_disaggregation: Annotated[
        bool,
        Field(
            description=(
                "Enable PD disaggregation mode. Requires a gateway with SGLang router. "
                "Uses internal IP instead of SSH tunnels for replica connections."
            ),
        ),
    ] = False

Update ServiceConfigurationParams's gateway field to

gateway: Annotated[
    Optional[Union[bool, str, GatewayConfig]],
    Field(
        description=(
            "Gateway configuration. Use `false` to run without a gateway. "
            "Use `true` to run with the default gateway. "
            "Use a string for the gateway name. "
            "Use an object for name + options (e.g. pd_disaggregation). "
            "Omit to use the default gateway if one exists, otherwise no gateway."
        ),
    ),
    ] = None

YAML configs with gateway..

# Boolean shortcuts (unchanged)
gateway: true   # default gateway
gateway: false  # no gateway

# String shortcut
gateway: "my-gateway"

# Object form
gateway:
  name: "my-gateway"
  pd_disaggregation: true

# Default gateway + pd_disaggregation
gateway:
  pd_disaggregation: true

Validation in _register_service_in_server would be

if isinstance(run_spec.configuration.gateway, GatewayConfig):
    raise ServerClientError(
        "Service is configured with gateway options but no gateway is available. "
        "Please configure a gateway or set gateway to false to run without a gateway."
    )

Validation in _register_service_in_gateway

gateway_configuration = get_gateway_configuration(gateway)
service_gateway = run_spec.configuration.gateway
if (
    isinstance(service_gateway, GatewayConfig)
    and service_gateway.pd_disaggregation
):
    if gateway_configuration.router != RouterType.SGLANG:
        raise ServerClientError(
            f"Service requires a SGLang gateway (pd_disaggregation) but gateway '{gateway.name}' "
            "does not have the SGLang router configured."
        )

Finally router construction in _register_service_in_gateway becomes

router = None
if gateway_configuration.router == RouterType.SGLANG:
    service_gateway = run_spec.configuration.gateway
    pd_disaggregation = (
        isinstance(service_gateway, GatewayConfig) and service_gateway.pd_disaggregation
    )
    policy = gateway_configuration.router_policy  # or gateway_configuration.router_config.policy
    router = ServiceRouterConfig(
        type=RouterType.SGLANG.value,
        policy=policy,
        pd_disaggregation=pd_disaggregation,
    )