chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,159 @@
|
||||
import pprint
|
||||
from typing import Any, Optional, Union
|
||||
|
||||
from pydantic import Field, field_validator
|
||||
|
||||
from ray import serve
|
||||
from ray.llm._internal.common.base_pydantic import BaseModelExtended
|
||||
from ray.llm._internal.common.dict_utils import deep_merge_dicts
|
||||
from ray.llm._internal.serve.constants import RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING
|
||||
from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
|
||||
from ray.llm._internal.serve.core.configs.openai_api_models import to_model_metadata
|
||||
from ray.llm._internal.serve.core.ingress.builder import (
|
||||
IngressClsConfig,
|
||||
_build_direct_streaming_llm_deployment,
|
||||
_build_openai_ingress_request_router,
|
||||
_validate_direct_streaming_ingress_config,
|
||||
)
|
||||
from ray.llm._internal.serve.core.ingress.ingress import (
|
||||
make_fastapi_ingress,
|
||||
)
|
||||
from ray.llm._internal.serve.core.server.builder import build_llm_deployment
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
from ray.llm._internal.serve.serving_patterns.data_parallel.dp_server import (
|
||||
DPServer,
|
||||
)
|
||||
from ray.serve.deployment import Application
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def build_dp_deployment(
|
||||
llm_config: LLMConfig,
|
||||
*,
|
||||
name_prefix: Optional[str] = None,
|
||||
bind_kwargs: Optional[dict] = None,
|
||||
override_serve_options: Optional[dict] = None,
|
||||
deployment_cls: Optional[type] = None,
|
||||
) -> Application:
|
||||
"""Build a data parallel attention LLM deployment.
|
||||
|
||||
Args:
|
||||
llm_config: The LLM configuration.
|
||||
name_prefix: The prefix to add to the deployment name.
|
||||
bind_kwargs: Optional extra kwargs to pass to the deployment constructor.
|
||||
Used by PD disaggregation to inject prefill_server handles.
|
||||
override_serve_options: The optional serve options to override the
|
||||
default options.
|
||||
deployment_cls: Optional deployment class to use. Defaults to DPServer.
|
||||
|
||||
Returns:
|
||||
The Ray Serve Application for the data parallel attention LLM deployment.
|
||||
"""
|
||||
return build_llm_deployment(
|
||||
llm_config,
|
||||
name_prefix=name_prefix,
|
||||
bind_kwargs=bind_kwargs,
|
||||
override_serve_options=override_serve_options,
|
||||
deployment_cls=deployment_cls or DPServer,
|
||||
)
|
||||
|
||||
|
||||
class DPOpenAiServingArgs(BaseModelExtended):
|
||||
"""Schema for DP OpenAI serving args."""
|
||||
|
||||
llm_config: Union[str, dict, LLMConfig] = Field(
|
||||
description="The LLM configuration",
|
||||
)
|
||||
ingress_cls_config: Union[dict, IngressClsConfig] = Field(
|
||||
default_factory=IngressClsConfig,
|
||||
description="The configuration for the ingress class.",
|
||||
)
|
||||
ingress_deployment_config: Optional[dict] = Field(
|
||||
default_factory=dict,
|
||||
description="The Ray @server.deployment options for the ingress server.",
|
||||
)
|
||||
|
||||
@field_validator("llm_config")
|
||||
@classmethod
|
||||
def _validate_llm_config(cls, value: Any) -> LLMConfig:
|
||||
if isinstance(value, str):
|
||||
return LLMConfig.from_file(value)
|
||||
elif isinstance(value, dict):
|
||||
return LLMConfig.model_validate(value)
|
||||
elif isinstance(value, LLMConfig):
|
||||
return value
|
||||
else:
|
||||
raise TypeError(f"Invalid LLMConfig type: {type(value)}")
|
||||
|
||||
@field_validator("ingress_cls_config")
|
||||
@classmethod
|
||||
def _validate_ingress_cls_config(cls, value: Any) -> IngressClsConfig:
|
||||
if isinstance(value, dict):
|
||||
return IngressClsConfig.model_validate(value)
|
||||
return value
|
||||
|
||||
|
||||
def build_dp_openai_app(builder_config: dict) -> Application:
|
||||
"""Build an OpenAI compatible app with the DP attention deployment
|
||||
setup from the given builder configuration.
|
||||
|
||||
Args:
|
||||
builder_config: The configuration for the builder. It has to conform
|
||||
to the DPOpenAiServingArgs pydantic model.
|
||||
|
||||
Returns:
|
||||
The configured Ray Serve Application.
|
||||
"""
|
||||
|
||||
builder_config = DPOpenAiServingArgs.model_validate(builder_config)
|
||||
llm_config = builder_config.llm_config
|
||||
|
||||
if RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING:
|
||||
_validate_direct_streaming_ingress_config(
|
||||
builder_config.ingress_deployment_config,
|
||||
builder_config.ingress_cls_config,
|
||||
)
|
||||
direct_deployment = _build_direct_streaming_llm_deployment(
|
||||
llm_config,
|
||||
deployment_cls=DPServer,
|
||||
)
|
||||
logger.info(
|
||||
"Direct streaming enabled for DP: "
|
||||
"DPServer=ingress, LLMRouter=ingress_request_router"
|
||||
)
|
||||
return direct_deployment._with_ingress_request_router(
|
||||
_build_openai_ingress_request_router(
|
||||
server=direct_deployment, llm_config=llm_config
|
||||
)
|
||||
)
|
||||
|
||||
dp_deployment = build_dp_deployment(llm_config)
|
||||
|
||||
ingress_cls_config = builder_config.ingress_cls_config
|
||||
ingress_options = ingress_cls_config.ingress_cls.get_deployment_options(
|
||||
[llm_config]
|
||||
)
|
||||
|
||||
if builder_config.ingress_deployment_config:
|
||||
ingress_options = deep_merge_dicts(
|
||||
ingress_options, builder_config.ingress_deployment_config
|
||||
)
|
||||
|
||||
ingress_cls = make_fastapi_ingress(ingress_cls_config.ingress_cls)
|
||||
|
||||
logger.info("============== Ingress Options ==============")
|
||||
logger.info(pprint.pformat(ingress_options))
|
||||
|
||||
model_id = llm_config.model_id
|
||||
lora_config = llm_config.lora_config
|
||||
return serve.deployment(ingress_cls, **ingress_options).bind(
|
||||
llm_deployments={model_id: dp_deployment},
|
||||
model_cards={model_id: to_model_metadata(model_id, llm_config)},
|
||||
lora_paths=(
|
||||
{model_id: lora_config.dynamic_lora_loading_path}
|
||||
if lora_config is not None
|
||||
else {}
|
||||
),
|
||||
**ingress_cls_config.ingress_extra_kwargs,
|
||||
)
|
||||
@@ -0,0 +1,297 @@
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import List, Optional, Tuple, Type
|
||||
|
||||
import ray
|
||||
from ray import serve
|
||||
from ray.experimental.internal_kv import (
|
||||
_internal_kv_del,
|
||||
_internal_kv_get,
|
||||
_internal_kv_put,
|
||||
)
|
||||
from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
|
||||
from ray.llm._internal.serve.core.engine.protocol import LLMEngine
|
||||
from ray.llm._internal.serve.core.server.llm_server import LLMServer
|
||||
from ray.llm._internal.serve.utils.lora_serve_utils import LoraModelLoader
|
||||
from ray.llm._internal.serve.utils.pg_utils import get_bundle_indices_sorted_by_node
|
||||
from ray.serve.config import (
|
||||
AutoscalingConfig,
|
||||
GangPlacementStrategy,
|
||||
GangRuntimeFailurePolicy,
|
||||
GangSchedulingConfig,
|
||||
)
|
||||
from ray.util.collective.collective import get_address_and_port
|
||||
from ray.util.placement_group import get_placement_group
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
TIMEOUT_SECONDS = 120
|
||||
POLL_INTERVAL_SECONDS = 0.5
|
||||
|
||||
|
||||
class GangMasterInfoRegistry:
|
||||
"""Registry for gang DP master info using GCS KV store."""
|
||||
|
||||
_KEY_PREFIX = "LLMServeRegistry:serve_global:gang_dp_master/"
|
||||
|
||||
@classmethod
|
||||
def _make_key(cls, gang_id: str) -> bytes:
|
||||
return (cls._KEY_PREFIX + gang_id).encode("utf-8")
|
||||
|
||||
@classmethod
|
||||
def register(cls, gang_id: str, address: str, port: int, node_id: str) -> None:
|
||||
"""Store the DP master info in GCS KV store."""
|
||||
key = cls._make_key(gang_id)
|
||||
value = json.dumps(
|
||||
{"address": address, "port": port, "node_id": node_id}
|
||||
).encode("utf-8")
|
||||
_internal_kv_put(key, value, overwrite=True)
|
||||
|
||||
@classmethod
|
||||
def unregister(cls, gang_id: str) -> None:
|
||||
"""Remove the DP master info from GCS KV store."""
|
||||
key = cls._make_key(gang_id)
|
||||
try:
|
||||
_internal_kv_del(key)
|
||||
except Exception:
|
||||
logger.warning(
|
||||
f"Failed to unregister gang master info for gang {gang_id}.",
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def get(
|
||||
cls,
|
||||
gang_id: str,
|
||||
timeout: float = TIMEOUT_SECONDS,
|
||||
poll_interval: float = POLL_INTERVAL_SECONDS,
|
||||
) -> Tuple[str, int, str]:
|
||||
"""Retrieve the DP master info for gang_id, polling until available.
|
||||
|
||||
Args:
|
||||
gang_id: The ID of the gang.
|
||||
timeout: The timeout in seconds.
|
||||
poll_interval: The poll interval in seconds.
|
||||
|
||||
Returns:
|
||||
A tuple of (address, port, node_id).
|
||||
|
||||
Raises:
|
||||
TimeoutError: If the info is not available within timeout_seconds seconds.
|
||||
"""
|
||||
key = cls._make_key(gang_id)
|
||||
deadline = time.monotonic() + timeout
|
||||
while True:
|
||||
data = _internal_kv_get(key)
|
||||
if data is not None:
|
||||
info = json.loads(data)
|
||||
return info["address"], info["port"], info["node_id"]
|
||||
if time.monotonic() >= deadline:
|
||||
raise TimeoutError(
|
||||
f"Timed out waiting for DP master info for gang {gang_id} "
|
||||
f"after {timeout}s."
|
||||
)
|
||||
await asyncio.sleep(poll_interval)
|
||||
|
||||
|
||||
class DPServer(LLMServer):
|
||||
"""
|
||||
Gang-scheduled Data Parallel LLM Server.
|
||||
|
||||
Uses Ray Serve's gang scheduling so that if any replica in a DP group deployment
|
||||
fails, the entire group is restarted atomically.
|
||||
"""
|
||||
|
||||
async def __init__(
|
||||
self,
|
||||
llm_config: LLMConfig,
|
||||
*,
|
||||
engine_cls: Optional[Type[LLMEngine]] = None,
|
||||
model_downloader: Optional[Type[LoraModelLoader]] = None,
|
||||
):
|
||||
ctx = serve.get_replica_context()
|
||||
gang_context = ctx.gang_context
|
||||
|
||||
if gang_context is None:
|
||||
raise RuntimeError(
|
||||
"DPServer requires gang scheduling to be enabled. "
|
||||
"Set gang_scheduling_config in the deployment options."
|
||||
)
|
||||
|
||||
self.dp_rank = gang_context.rank
|
||||
self.gang_id = gang_context.gang_id
|
||||
self.dp_size = gang_context.world_size
|
||||
|
||||
logger.info(
|
||||
f"DPServer replica initialized: dp_rank={self.dp_rank}, "
|
||||
f"dp_size={self.dp_size}, gang_id={self.gang_id}"
|
||||
)
|
||||
|
||||
if self.dp_rank == 0:
|
||||
self.dp_address, self.dp_rpc_port = get_address_and_port()
|
||||
# Record rank 0's node so every replica places vLLM's global rank 0
|
||||
# worker (which hosts the distributed rendezvous store) on the same
|
||||
# node whose address we advertise below. See the bundle-ordering
|
||||
# comment in __init__ for why this matters.
|
||||
self.dp_node_id = ray.get_runtime_context().get_node_id()
|
||||
GangMasterInfoRegistry.register(
|
||||
self.gang_id, self.dp_address, self.dp_rpc_port, self.dp_node_id
|
||||
)
|
||||
logger.info(
|
||||
f"DP rank {self.dp_rank} has set DP master info: "
|
||||
f"data_parallel_address={self.dp_address}, "
|
||||
f"data_parallel_rpc_port={self.dp_rpc_port}, "
|
||||
f"data_parallel_node_id={self.dp_node_id}"
|
||||
)
|
||||
else:
|
||||
timestamp = time.time()
|
||||
(
|
||||
self.dp_address,
|
||||
self.dp_rpc_port,
|
||||
self.dp_node_id,
|
||||
) = await GangMasterInfoRegistry.get(self.gang_id)
|
||||
logger.info(
|
||||
f"DP rank {self.dp_rank} got DP master info: "
|
||||
f"data_parallel_address={self.dp_address}, "
|
||||
f"data_parallel_rpc_port={self.dp_rpc_port}, "
|
||||
f"data_parallel_node_id={self.dp_node_id}, "
|
||||
f"waited {time.time() - timestamp:.3f} seconds"
|
||||
)
|
||||
|
||||
# Update the engine_kwargs to assign the DP information
|
||||
llm_config.update_engine_kwargs(
|
||||
data_parallel_rank=self.dp_rank,
|
||||
data_parallel_address=self.dp_address,
|
||||
data_parallel_rpc_port=self.dp_rpc_port,
|
||||
)
|
||||
|
||||
engine_config = llm_config.get_engine_config()
|
||||
|
||||
# Direct vLLM to use this replica's bundles within the gang placement group.
|
||||
# Gang placement group concatenates per-replica bundles for all ranks, so
|
||||
# rank i owns bundles [i*B, i*B+1, ..., i*B+B-1] where B is the number of
|
||||
# bundles per DP replica.
|
||||
#
|
||||
# However, adjacent bundle indices in a placement group don't necessarily
|
||||
# map to adjacent physical ranks. We use get_bundle_indices_sorted_by_node
|
||||
# to reorder bundle indices so that same-node bundles are adjacent and
|
||||
# rank 0's node bundles come first. This prevents us from scattering
|
||||
# adjacent TP ranks in the same DP rank across nodes.
|
||||
#
|
||||
# Ordering rank 0's node first is also required for correctness: vLLM
|
||||
# forms a single distributed group across all DP workers whose rendezvous
|
||||
# store is hosted by global rank 0 and reached at the advertised
|
||||
# data_parallel_address (set above from rank 0's node). vLLM pins global
|
||||
# rank 0 to sorted_indices[0], so that bundle must live on the same node
|
||||
# whose address we advertised. Sorting by the cluster head node instead
|
||||
# (the previous default) breaks this whenever the head node owns no
|
||||
# bundles in the gang (e.g. a CPU-only head in a GPU cluster): rank 0
|
||||
# then lands on an arbitrary node, the store binds there, and every
|
||||
# worker hangs connecting to the wrong (advertised) address until the
|
||||
# distributed-init timeout fires and the gang is restarted.
|
||||
#
|
||||
# Example: dp_size=2, tp_size=2, 2 GPUs per node for simplicity
|
||||
# Gang placement group = [{GPU: 1}, {GPU: 1}, {GPU: 1}, {GPU: 1}]
|
||||
# Physical rank location: ^^N0R0^^ ^^N1R1^^ ^^N0R1^^ ^^N1R0^^
|
||||
# DP placement: ^^DP0^^^ ^^DP1^^^ ^^DP0^^^ ^^DP1^^^
|
||||
#
|
||||
# placement_bundles below is the gang placement group, and therefore
|
||||
# get_current_placement_group from the actor yields the gang placement group,
|
||||
# not the per-replica placement group.
|
||||
bundles_per_replica = len(engine_config.placement_bundles)
|
||||
pg = get_placement_group(gang_context.pg_name)
|
||||
sorted_indices = get_bundle_indices_sorted_by_node(
|
||||
pg, driver_node_id=self.dp_node_id
|
||||
)
|
||||
os.environ["VLLM_RAY_BUNDLE_INDICES"] = self._compute_bundle_indices(
|
||||
self.dp_rank, bundles_per_replica, sorted_indices
|
||||
)
|
||||
|
||||
await super().__init__(
|
||||
llm_config,
|
||||
engine_cls=engine_cls,
|
||||
model_downloader=model_downloader,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _compute_bundle_indices(
|
||||
dp_rank: int, bundles_per_replica: int, sorted_indices: List[int]
|
||||
) -> str:
|
||||
"""Return the VLLM_RAY_BUNDLE_INDICES value for a given DP rank.
|
||||
|
||||
Slices into sorted_indices (bundle indices reordered so that
|
||||
same-node bundles are adjacent) to pick the bundles that belong to
|
||||
this DP rank.
|
||||
|
||||
Args:
|
||||
dp_rank: This replica's data-parallel rank.
|
||||
bundles_per_replica: Number of placement-group bundles each DP
|
||||
replica owns.
|
||||
sorted_indices: Bundle indices sorted by node.
|
||||
|
||||
Returns:
|
||||
Comma-separated bundle indices, e.g. "0,2".
|
||||
"""
|
||||
start = dp_rank * bundles_per_replica
|
||||
return ",".join(
|
||||
str(sorted_indices[start + i]) for i in range(bundles_per_replica)
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_deployment_options(cls, llm_config: "LLMConfig"):
|
||||
deployment_options = super().get_deployment_options(llm_config)
|
||||
|
||||
dp_size = llm_config.engine_kwargs.get("data_parallel_size", 1)
|
||||
if not (isinstance(dp_size, int) and dp_size > 0):
|
||||
raise ValueError(
|
||||
f"Invalid data_parallel_size: {dp_size}, expecting positive integer."
|
||||
)
|
||||
if dp_size != 1:
|
||||
num_replicas = deployment_options.get("num_replicas")
|
||||
has_autoscaling = num_replicas == "auto" or (
|
||||
num_replicas is None and "autoscaling_config" in deployment_options
|
||||
)
|
||||
if has_autoscaling:
|
||||
autoscaling_config = AutoscalingConfig.default().dict()
|
||||
user_config = deployment_options.get("autoscaling_config")
|
||||
if user_config is not None:
|
||||
autoscaling_config.update(user_config)
|
||||
|
||||
logger.warning(
|
||||
"In DP deployment, a replica refers to a DP group. "
|
||||
"Multiplying autoscaling_config's min_replicas, max_replicas, and "
|
||||
f"initial_replicas by data_parallel_size ({dp_size})."
|
||||
)
|
||||
for key in ["min_replicas", "max_replicas", "initial_replicas"]:
|
||||
if autoscaling_config.get(key) is not None:
|
||||
autoscaling_config[key] *= dp_size
|
||||
|
||||
deployment_options["autoscaling_config"] = autoscaling_config
|
||||
elif num_replicas is not None:
|
||||
logger.warning(
|
||||
"In DP deployment, num_replicas refers to the number of DP groups. "
|
||||
f"Multiplying num_replicas ({num_replicas}) by data_parallel_size ({dp_size}) "
|
||||
f"to get the total number of serve replicas ({num_replicas * dp_size})."
|
||||
)
|
||||
deployment_options["num_replicas"] = num_replicas * dp_size
|
||||
else:
|
||||
deployment_options["num_replicas"] = dp_size
|
||||
|
||||
deployment_options["gang_scheduling_config"] = GangSchedulingConfig(
|
||||
gang_size=dp_size,
|
||||
gang_placement_strategy=GangPlacementStrategy.PACK,
|
||||
runtime_failure_policy=GangRuntimeFailurePolicy.RESTART_GANG,
|
||||
)
|
||||
# Remove per-replica placement_group_strategy. Ray Serve raises an error
|
||||
# if both placement_group_strategy and gang_scheduling_config are provided.
|
||||
if "placement_group_strategy" in deployment_options:
|
||||
logger.warning(
|
||||
"placement_group_strategy configured in the deployment config is ignored. "
|
||||
"DP deployment uses PACK strategy for scheduling DP groups."
|
||||
)
|
||||
deployment_options.pop("placement_group_strategy", None)
|
||||
|
||||
return deployment_options
|
||||
@@ -0,0 +1,305 @@
|
||||
"""Using Ray Serve to deploy LLM models with P/D disaggregation.
|
||||
|
||||
3-tier graph: ingress -> PDDecodeServer (decode config + engine) -> PDPrefillServer.
|
||||
"""
|
||||
|
||||
import warnings
|
||||
from typing import Any, Optional, Union
|
||||
|
||||
from pydantic import Field, field_validator, model_validator
|
||||
|
||||
from ray import serve
|
||||
from ray.llm._internal.common.base_pydantic import BaseModelExtended
|
||||
from ray.llm._internal.common.dict_utils import (
|
||||
maybe_apply_llm_deployment_config_defaults,
|
||||
)
|
||||
from ray.llm._internal.serve.constants import RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING
|
||||
from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
|
||||
from ray.llm._internal.serve.core.configs.openai_api_models import to_model_metadata
|
||||
from ray.llm._internal.serve.core.ingress.builder import (
|
||||
IngressClsConfig,
|
||||
_build_direct_streaming_llm_deployment,
|
||||
_build_openai_ingress_request_router,
|
||||
_validate_direct_streaming_ingress_config,
|
||||
load_class,
|
||||
)
|
||||
from ray.llm._internal.serve.core.ingress.ingress import (
|
||||
make_fastapi_ingress,
|
||||
)
|
||||
from ray.llm._internal.serve.core.server.builder import build_llm_deployment
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
from ray.llm._internal.serve.serving_patterns.data_parallel.builder import (
|
||||
build_dp_deployment,
|
||||
)
|
||||
from ray.llm._internal.serve.serving_patterns.prefill_decode.pd_server import (
|
||||
DPPDDecodeServer,
|
||||
DPPDPrefillServer,
|
||||
PDDecodeServer,
|
||||
PDPrefillServer,
|
||||
PDProxyServer, # TODO(Kourosh): Deprecate, remove in Ray 2.58.
|
||||
)
|
||||
from ray.serve.deployment import Application
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Deprecated: ProxyClsConfig
|
||||
# TODO(Kourosh): Deprecate, remove in Ray 2.58.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class ProxyClsConfig(BaseModelExtended):
|
||||
"""Deprecated. Unused proxy configuration kept for backwards compatibility."""
|
||||
|
||||
proxy_cls: Union[str, type] = Field(
|
||||
default=PDProxyServer,
|
||||
description="Deprecated.",
|
||||
)
|
||||
|
||||
proxy_extra_kwargs: Optional[dict] = Field(
|
||||
default_factory=dict,
|
||||
description="Deprecated.",
|
||||
)
|
||||
|
||||
@field_validator("proxy_cls")
|
||||
@classmethod
|
||||
def validate_class(cls, value):
|
||||
if isinstance(value, str):
|
||||
return load_class(value)
|
||||
return value
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# PDServingArgs
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class PDServingArgs(BaseModelExtended):
|
||||
"""Schema for P/D serving args.
|
||||
|
||||
Defines the prefill and decode LLMConfigs plus ingress options.
|
||||
The deprecated ``proxy_cls_config`` and ``proxy_deployment_config``
|
||||
fields are accepted for backwards compatibility but ignored.
|
||||
"""
|
||||
|
||||
prefill_config: Union[str, dict, LLMConfig]
|
||||
decode_config: Union[str, dict, LLMConfig]
|
||||
|
||||
# TODO(Kourosh): Deprecated, remove in Ray 2.58.
|
||||
# Deprecated proxy fields — accepted for backwards compat, ignored at build time.
|
||||
proxy_cls_config: Optional[Union[dict, ProxyClsConfig]] = Field(
|
||||
default=None,
|
||||
description="Deprecated. Accepted but ignored.",
|
||||
)
|
||||
proxy_deployment_config: Optional[dict] = Field(
|
||||
default=None,
|
||||
description="Deprecated. Accepted but ignored.",
|
||||
)
|
||||
|
||||
ingress_cls_config: Union[dict, IngressClsConfig] = Field(
|
||||
default_factory=IngressClsConfig,
|
||||
description="The configuration for the ingress class.",
|
||||
)
|
||||
ingress_deployment_config: Optional[dict] = Field(
|
||||
default_factory=dict,
|
||||
description="The Ray @serve.deployment options for the ingress.",
|
||||
)
|
||||
|
||||
@field_validator("prefill_config", "decode_config")
|
||||
@classmethod
|
||||
def _validate_llm_config(cls, value: Any) -> LLMConfig:
|
||||
if isinstance(value, str):
|
||||
return LLMConfig.from_file(value)
|
||||
elif isinstance(value, dict):
|
||||
return LLMConfig.model_validate(value)
|
||||
elif isinstance(value, LLMConfig):
|
||||
return value
|
||||
else:
|
||||
raise TypeError(f"Invalid LLMConfig type: {type(value)}")
|
||||
|
||||
@field_validator("proxy_cls_config")
|
||||
@classmethod
|
||||
def _validate_proxy_cls_config(
|
||||
cls, value: Optional[Union[dict, ProxyClsConfig]]
|
||||
) -> Optional[ProxyClsConfig]:
|
||||
if value is not None:
|
||||
warnings.warn(
|
||||
"proxy_cls_config is deprecated and ignored. "
|
||||
"The proxy has been replaced by PDDecodeServer which "
|
||||
"orchestrates prefill and decode directly. "
|
||||
"See PDDecodeServer and PDPrefillServer.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
if isinstance(value, dict):
|
||||
return ProxyClsConfig.model_validate(value)
|
||||
return value
|
||||
|
||||
@field_validator("proxy_deployment_config")
|
||||
@classmethod
|
||||
def _validate_proxy_deployment_config(cls, value: Optional[dict]) -> Optional[dict]:
|
||||
if value is not None:
|
||||
warnings.warn(
|
||||
"proxy_deployment_config is deprecated and ignored. "
|
||||
"The proxy has been replaced by PDDecodeServer which "
|
||||
"orchestrates prefill and decode directly. "
|
||||
"See PDDecodeServer and PDPrefillServer.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
return value
|
||||
|
||||
@field_validator("ingress_cls_config")
|
||||
@classmethod
|
||||
def _validate_ingress_cls_config(
|
||||
cls, value: Union[dict, IngressClsConfig]
|
||||
) -> IngressClsConfig:
|
||||
if isinstance(value, dict):
|
||||
return IngressClsConfig.model_validate(value)
|
||||
return value
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_model_ids(self):
|
||||
"""Validate that prefill and decode configs use the same model ID."""
|
||||
if self.prefill_config.model_id != self.decode_config.model_id:
|
||||
raise ValueError("P/D model id mismatch")
|
||||
return self
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_kv_transfer_config(self):
|
||||
"""Validate that kv_transfer_config is set for both prefill and decode configs."""
|
||||
for config in [self.prefill_config, self.decode_config]:
|
||||
if config.engine_kwargs.get("kv_transfer_config") is None:
|
||||
raise ValueError(
|
||||
"kv_transfer_config is required for P/D disaggregation"
|
||||
)
|
||||
return self
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _default_decode_nixl_port_base(self):
|
||||
"""Shift decode's NIXL base off prefill's default (20000) so colocated replicas don't collide."""
|
||||
self.decode_config.experimental_configs.setdefault(
|
||||
"NIXL_SIDE_CHANNEL_PORT_BASE", 22000
|
||||
)
|
||||
return self
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _default_decode_moriio_port_base(self):
|
||||
"""Shift decode's MoRIIO handshake/notify bases off prefill's defaults.
|
||||
|
||||
Mirrors ``_default_decode_nixl_port_base``: a colocated P+D pair on one
|
||||
node would otherwise share MoRIIO's default handshake/notify ports. Only
|
||||
applies when the decode config uses the MoRIIO connector. The +1000
|
||||
stride is well above any realistic tp_size*pp_size offset added on top.
|
||||
"""
|
||||
kv_transfer_config = (
|
||||
self.decode_config.engine_kwargs.get("kv_transfer_config") or {}
|
||||
)
|
||||
if kv_transfer_config.get("kv_connector") != "MoRIIOConnector":
|
||||
return self
|
||||
|
||||
from ray.llm._internal.serve.engines.vllm.kv_transfer.moriio import (
|
||||
DEFAULT_HANDSHAKE_PORT_BASE,
|
||||
DEFAULT_NOTIFY_PORT_BASE,
|
||||
HANDSHAKE_PORT_BASE_KEY,
|
||||
NOTIFY_PORT_BASE_KEY,
|
||||
)
|
||||
|
||||
self.decode_config.experimental_configs.setdefault(
|
||||
HANDSHAKE_PORT_BASE_KEY, DEFAULT_HANDSHAKE_PORT_BASE + 1000
|
||||
)
|
||||
self.decode_config.experimental_configs.setdefault(
|
||||
NOTIFY_PORT_BASE_KEY, DEFAULT_NOTIFY_PORT_BASE + 1000
|
||||
)
|
||||
return self
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Builder
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def build_pd_openai_app(pd_serving_args: dict) -> Application:
|
||||
"""Build a deployable application utilizing prefill/decode disaggregation.
|
||||
|
||||
3-tier graph: ingress -> PDDecodeServer -> PDPrefillServer.
|
||||
"""
|
||||
pd_config = PDServingArgs.model_validate(pd_serving_args)
|
||||
|
||||
if RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING:
|
||||
_validate_direct_streaming_ingress_config(
|
||||
pd_config.ingress_deployment_config,
|
||||
pd_config.ingress_cls_config,
|
||||
)
|
||||
|
||||
prefill_dp_size = pd_config.prefill_config.engine_kwargs.get(
|
||||
"data_parallel_size", 1
|
||||
)
|
||||
decode_dp_size = pd_config.decode_config.engine_kwargs.get("data_parallel_size", 1)
|
||||
prefill_builder = (
|
||||
build_dp_deployment if prefill_dp_size > 1 else build_llm_deployment
|
||||
)
|
||||
|
||||
# When DP > 1, use combined DP+PD server classes that inherit from both
|
||||
# the PD server and DPServer (for gang scheduling, DP master info, etc.).
|
||||
prefill_cls = DPPDPrefillServer if prefill_dp_size > 1 else PDPrefillServer
|
||||
decode_cls = DPPDDecodeServer if decode_dp_size > 1 else PDDecodeServer
|
||||
|
||||
prefill_deployment = prefill_builder(
|
||||
pd_config.prefill_config,
|
||||
name_prefix="Prefill:",
|
||||
deployment_cls=prefill_cls,
|
||||
)
|
||||
|
||||
if RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING:
|
||||
# Direct streaming makes decode the ASGI ingress, so it must be built
|
||||
# with the ASGI wrapper while still receiving the prefill backend.
|
||||
decode_deployment = _build_direct_streaming_llm_deployment(
|
||||
pd_config.decode_config,
|
||||
name_prefix="Decode:",
|
||||
bind_kwargs={"prefill_server": prefill_deployment},
|
||||
deployment_cls=decode_cls,
|
||||
)
|
||||
logger.info(
|
||||
"Direct streaming enabled for PD: "
|
||||
f"{decode_cls.__name__}=ingress, LLMRouter=ingress_request_router"
|
||||
)
|
||||
return decode_deployment._with_ingress_request_router(
|
||||
_build_openai_ingress_request_router(
|
||||
server=decode_deployment, llm_config=pd_config.decode_config
|
||||
)
|
||||
)
|
||||
|
||||
decode_builder = build_dp_deployment if decode_dp_size > 1 else build_llm_deployment
|
||||
decode_deployment = decode_builder(
|
||||
pd_config.decode_config,
|
||||
name_prefix="Decode:",
|
||||
bind_kwargs={"prefill_server": prefill_deployment},
|
||||
deployment_cls=decode_cls,
|
||||
)
|
||||
|
||||
# -- Ingress: binds to decode only (the "model" the client sees) --
|
||||
ingress_cls_config = pd_config.ingress_cls_config
|
||||
default_ingress_options = ingress_cls_config.ingress_cls.get_deployment_options(
|
||||
[pd_config.decode_config]
|
||||
)
|
||||
|
||||
ingress_options = maybe_apply_llm_deployment_config_defaults(
|
||||
default_ingress_options, pd_config.ingress_deployment_config
|
||||
)
|
||||
|
||||
ingress_cls = make_fastapi_ingress(ingress_cls_config.ingress_cls)
|
||||
# Prefill and decode share the same model_id (validated in PDServingArgs).
|
||||
# Ingress binds to decode only (the "model" the client sees).
|
||||
model_id = pd_config.decode_config.model_id
|
||||
lora_config = pd_config.decode_config.lora_config
|
||||
return serve.deployment(ingress_cls, **ingress_options).bind(
|
||||
llm_deployments={model_id: decode_deployment},
|
||||
model_cards={model_id: to_model_metadata(model_id, pd_config.decode_config)},
|
||||
lora_paths=(
|
||||
{model_id: lora_config.dynamic_lora_loading_path}
|
||||
if lora_config is not None
|
||||
else {}
|
||||
),
|
||||
**ingress_cls_config.ingress_extra_kwargs,
|
||||
)
|
||||
@@ -0,0 +1,870 @@
|
||||
"""Prefill-Decode disaggregated LLM serving: decode-as-orchestrator architecture.
|
||||
|
||||
3-tier graph (ingress -> PDDecodeServer -> PDPrefillServer) where the
|
||||
decode deployment owns a real engine and orchestrates remote prefill.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import logging
|
||||
import uuid
|
||||
import warnings
|
||||
from typing import Any, AsyncGenerator, Dict, List, Optional, Union
|
||||
|
||||
from fastapi.routing import APIRoute
|
||||
from starlette.requests import Request
|
||||
from starlette.responses import JSONResponse, Response, StreamingResponse
|
||||
|
||||
from ray.llm._internal.common.patches.vllm.tokenize_once import (
|
||||
install as _install_tokenize_once,
|
||||
reuse_prompt_token_ids as _reuse_prompt_token_ids,
|
||||
)
|
||||
from ray.llm._internal.serve.constants import DEFAULT_MAX_ONGOING_REQUESTS
|
||||
from ray.llm._internal.serve.core.configs.openai_api_models import (
|
||||
ChatCompletionRequest,
|
||||
ChatCompletionResponse,
|
||||
CompletionRequest,
|
||||
CompletionResponse,
|
||||
EmbeddingRequest,
|
||||
EmbeddingResponse,
|
||||
ErrorResponse,
|
||||
)
|
||||
from ray.llm._internal.serve.core.ingress.utils import (
|
||||
NON_STREAMING_RESPONSE_TYPES,
|
||||
_openai_json_wrapper,
|
||||
_peek_at_generator,
|
||||
_sanitize_chat_completion_request,
|
||||
)
|
||||
from ray.llm._internal.serve.core.protocol import LLMServerProtocol, RawRequestInfo
|
||||
from ray.llm._internal.serve.core.server.llm_server import LLMServer
|
||||
from ray.llm._internal.serve.engines.vllm.kv_transfer.base import BaseConnectorBackend
|
||||
from ray.llm._internal.serve.serving_patterns.data_parallel.dp_server import DPServer
|
||||
from ray.llm._internal.serve.utils.broadcast import broadcast
|
||||
from ray.llm._internal.serve.utils.server_utils import (
|
||||
get_serve_request_id,
|
||||
)
|
||||
from ray.serve._private.http_util import session_id_from_headers
|
||||
from ray.serve.exceptions import DeploymentUnavailableError
|
||||
from ray.serve.handle import DeploymentHandle
|
||||
from ray.serve.llm import LLMConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
RequestType = Union[ChatCompletionRequest, CompletionRequest]
|
||||
|
||||
# TODO(Kourosh): Deprecate in Ray 2.56, remove in Ray 2.58.
|
||||
DEFAULT_PD_PROXY_SERVER_OPTIONS = {
|
||||
"max_ongoing_requests": DEFAULT_MAX_ONGOING_REQUESTS,
|
||||
}
|
||||
|
||||
_PREWARM_PROMPT = " x"
|
||||
_PREWARM_MAX_TOKENS = 1
|
||||
_PREWARM_RETRY_INTERVAL_S = 5.0
|
||||
_PREWARM_MAX_RETRIES = 60
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Direct-streaming route helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
#
|
||||
# Direct streaming exposes the engine-native ASGI app directly on the LLM
|
||||
# server replica (see ``LLMServer.__serve_build_asgi_app__``), eliminating the
|
||||
# separate ``OpenAiIngress`` deployment. For P/D, the engine-native
|
||||
# chat/completions routes would send traffic straight to the local decode
|
||||
# engine and bypass remote prefill, so ``PDOrchestratorMixin`` re-points just
|
||||
# those two routes at its own ``chat`` / ``completions`` (which orchestrate
|
||||
# prefill then decode). Every other route stays engine-native, identical to
|
||||
# non-P/D direct streaming.
|
||||
|
||||
|
||||
def _strip_routes(app, path: str) -> None:
|
||||
"""Remove the engine-native APIRoute(s) registered at ``path``."""
|
||||
app.routes[:] = [
|
||||
r for r in app.routes if not (isinstance(r, APIRoute) and r.path == path)
|
||||
]
|
||||
|
||||
|
||||
async def _pd_http_response(gen) -> Response:
|
||||
"""Shape a P/D orchestration generator into an OpenAI HTTP response.
|
||||
|
||||
Returns a JSON response when the first chunk is an error or a complete
|
||||
(non-streaming) response, otherwise an SSE stream. Uses the same response
|
||||
helpers as ``OpenAiIngress`` so the wire format matches the standard path.
|
||||
"""
|
||||
first, gen = await _peek_at_generator(gen)
|
||||
if isinstance(first, list):
|
||||
first = first[0]
|
||||
if isinstance(first, ErrorResponse):
|
||||
return JSONResponse(
|
||||
content=first.model_dump(), status_code=first.error.code or 400
|
||||
)
|
||||
if isinstance(first, NON_STREAMING_RESPONSE_TYPES):
|
||||
return JSONResponse(content=first.model_dump())
|
||||
return StreamingResponse(_openai_json_wrapper(gen), media_type="text/event-stream")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Mixin: PD Orchestration Logic
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class PDOrchestratorMixin:
|
||||
"""Mixin that adds prefill-decode orchestration to an LLMServer subclass.
|
||||
|
||||
For chat/completions requests it:
|
||||
1. Sends a modified prefill request (max_tokens=1, kv_transfer_params).
|
||||
2. Receives kv_transfer_params back from the first prefill chunk.
|
||||
3. Runs decode locally on its own engine with those params.
|
||||
"""
|
||||
|
||||
# Set by __init__ of the concrete class that mixes this in.
|
||||
_prefill_handle: DeploymentHandle
|
||||
|
||||
# Decode reuses prefill's prompt token ids. Set from experimental_configs in
|
||||
# PDDecodeServer.__init__. Default off.
|
||||
_pd_tokenize_once: bool = False
|
||||
|
||||
# ---- Connector backend resolution ----
|
||||
|
||||
def _get_connector_backend(self) -> BaseConnectorBackend:
|
||||
"""Return the connector backend that was set up during engine init.
|
||||
|
||||
``LLMConfig.setup_engine_backend()`` creates the backend and calls its
|
||||
``setup()`` during engine initialization, storing it on the config. By the
|
||||
time a request reaches the orchestrator it must already be there — a
|
||||
missing backend means engine init was skipped, which is a bug (and a
|
||||
freshly-created, un-``setup()`` backend would mis-shape traffic, e.g. a
|
||||
MultiConnector whose sub-connectors are populated only in ``setup()``).
|
||||
Cached on first access since the request path calls this.
|
||||
"""
|
||||
cached = getattr(self, "_connector_backend_cache", None)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
backend = getattr(self._llm_config, "kv_connector_backend", None)
|
||||
assert backend is not None, (
|
||||
"No KV-connector backend on the LLMConfig. setup_engine_backend() must "
|
||||
"run during engine init before the P/D orchestrator handles requests."
|
||||
)
|
||||
self._connector_backend_cache = backend
|
||||
return backend
|
||||
|
||||
# ---- Request Preparation ----
|
||||
#
|
||||
# Thin instance delegates to the resolved backend's protocol so existing
|
||||
# callers/tests that reference these names keep working. The orchestrator
|
||||
# itself goes through ``backend.prepare_*`` directly.
|
||||
|
||||
def _prepare_prefill_request(self, request: RequestType) -> RequestType:
|
||||
return self._get_connector_backend().prepare_prefill_request(
|
||||
request=request, peer=None
|
||||
)
|
||||
|
||||
def _prepare_decode_request(
|
||||
self,
|
||||
request: RequestType,
|
||||
prefill_chunk: Union[ChatCompletionResponse, CompletionResponse],
|
||||
) -> RequestType:
|
||||
return self._get_connector_backend().prepare_decode_request(
|
||||
request=request, peer=None, prefill_response=prefill_chunk
|
||||
)
|
||||
|
||||
def _decode_reuse_ids(self, prefill_chunk) -> Optional[list]:
|
||||
"""Prompt token ids for decode to reuse, or None when disabled or absent.
|
||||
|
||||
Chat carries them top-level. Completions carry them on the first choice as
|
||||
``CompletionResponseChoice.prompt_token_ids``."""
|
||||
if not self._pd_tokenize_once:
|
||||
return None
|
||||
ids = getattr(prefill_chunk, "prompt_token_ids", None)
|
||||
if ids is None:
|
||||
choices = getattr(prefill_chunk, "choices", None)
|
||||
if choices:
|
||||
ids = getattr(choices[0], "prompt_token_ids", None)
|
||||
return ids
|
||||
|
||||
def _request_prefill_token_ids(self, prefill_request) -> None:
|
||||
"""Ask prefill to echo its prompt token ids so decode can reuse them.
|
||||
|
||||
No-op when disabled or the request lacks the field. Used on sequential handoff
|
||||
only. Concurrent decode starts before prefill returns, so it has nothing to
|
||||
reuse."""
|
||||
if self._pd_tokenize_once and hasattr(prefill_request, "return_token_ids"):
|
||||
prefill_request.return_token_ids = True
|
||||
|
||||
# ---- Orchestrated Request Flow ----
|
||||
|
||||
async def _pd_handle_request(
|
||||
self,
|
||||
request: RequestType,
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[
|
||||
Union[str, ChatCompletionResponse, CompletionResponse, ErrorResponse], None
|
||||
]:
|
||||
"""Orchestrate prefill (remote) then decode (local engine).
|
||||
|
||||
Request shaping, peer addressing, and handoff discipline are delegated to
|
||||
the resolved KV-connector backend. With the default backend flags
|
||||
(``requires_peer_binding=False``, ``concurrent_handoff=False``) the
|
||||
control flow and calls are identical to the historical NIXL/default flow.
|
||||
|
||||
A connector that encodes coordination data in the request id (MoRIIO's
|
||||
dual-address id) just stamps ``request.request_id`` in ``prepare_*``; it
|
||||
then reaches both engines unchanged -- the LLMServer pipeline preserves
|
||||
an explicitly-set request_id (it no longer clobbers it with the Serve
|
||||
id) and the engine copies it into the ``X-Request-Id`` header it reads.
|
||||
"""
|
||||
|
||||
# Determine method name for the handle call
|
||||
if isinstance(request, ChatCompletionRequest):
|
||||
method = "chat"
|
||||
elif isinstance(request, CompletionRequest):
|
||||
method = "completions"
|
||||
else:
|
||||
raise ValueError(f"Unsupported request type: {type(request)}")
|
||||
|
||||
backend = self._get_connector_backend()
|
||||
|
||||
prefill_handle = self._prefill_handle
|
||||
if raw_request_info is not None:
|
||||
session_id = session_id_from_headers(raw_request_info.headers)
|
||||
if session_id:
|
||||
prefill_handle = prefill_handle.options(session_id=session_id)
|
||||
prefill_handle_method = getattr(prefill_handle, method)
|
||||
|
||||
if backend.requires_peer_binding:
|
||||
# Connector needs to bind to the selected prefill replica *before*
|
||||
# dispatch (e.g. request-id-addressed transfers). Reserve a replica
|
||||
# via choose_replica, expose its metadata to the backend, then
|
||||
# dispatch onto that exact selection.
|
||||
async with prefill_handle_method.choose_replica(request) as selection:
|
||||
# The selected replica's published metadata (empty dict if none).
|
||||
peer = getattr(selection, "replica_metadata", {})
|
||||
prefill_request = backend.prepare_prefill_request(
|
||||
request=request, peer=peer
|
||||
)
|
||||
|
||||
if backend.concurrent_handoff:
|
||||
# Concurrent handoff: start remote prefill and run local decode
|
||||
# together, draining prefill before leaving the choose_replica
|
||||
# context (so on_request_completed fires once).
|
||||
decode_request = backend.prepare_decode_request(
|
||||
request=request, peer=peer, prefill_response=None
|
||||
)
|
||||
prefill_resp = prefill_handle_method.dispatch(
|
||||
selection, prefill_request, raw_request_info
|
||||
)
|
||||
# dispatch()'s completion accounting fires when its result
|
||||
# completes, so the response must be drained to exhaustion
|
||||
# inside the choose_replica context — never cancelled
|
||||
# (prefill is clamped to a single token, so draining is
|
||||
# bounded).
|
||||
async for chunk in self._concurrent_decode(
|
||||
method,
|
||||
decode_request,
|
||||
prefill_resp,
|
||||
raw_request_info,
|
||||
cancel_on_failure=False,
|
||||
):
|
||||
yield chunk
|
||||
return
|
||||
|
||||
# Sequential handoff with peer binding: run prefill to its first
|
||||
# chunk, then drive local decode with the returned params.
|
||||
self._request_prefill_token_ids(prefill_request)
|
||||
prefill_gen = prefill_handle_method.dispatch(
|
||||
selection, prefill_request, raw_request_info
|
||||
)
|
||||
prefill_chunk = await prefill_gen.__anext__()
|
||||
# Drain the dispatched stream to exhaustion inside the
|
||||
# choose_replica context: dispatch()'s completion accounting
|
||||
# (queue-length cache decrement) fires when the result
|
||||
# completes. Prefill is clamped to a single token, so this is
|
||||
# at most one trivial extra iteration.
|
||||
async for _ in prefill_gen:
|
||||
pass
|
||||
if isinstance(prefill_chunk, ErrorResponse):
|
||||
logger.error(f"Prefill returned error: {prefill_chunk}")
|
||||
yield prefill_chunk
|
||||
return
|
||||
decode_request = backend.prepare_decode_request(
|
||||
request=request, peer=peer, prefill_response=prefill_chunk
|
||||
)
|
||||
with _reuse_prompt_token_ids(self._decode_reuse_ids(prefill_chunk)):
|
||||
local_gen = await getattr(super(), method)(
|
||||
decode_request, raw_request_info
|
||||
)
|
||||
async for chunk in local_gen:
|
||||
yield chunk
|
||||
return
|
||||
|
||||
# Default path: no pre-dispatch peer binding; dispatch prefill via the
|
||||
# standard handle path.
|
||||
prefill_request = backend.prepare_prefill_request(request=request, peer=None)
|
||||
|
||||
if backend.concurrent_handoff:
|
||||
# Concurrent handoff: dispatch via remote() and run local decode
|
||||
# together.
|
||||
decode_request = backend.prepare_decode_request(
|
||||
request=request, peer=None, prefill_response=None
|
||||
)
|
||||
prefill_resp = prefill_handle_method.remote(
|
||||
prefill_request, raw_request_info
|
||||
)
|
||||
async for chunk in self._concurrent_decode(
|
||||
method, decode_request, prefill_resp, raw_request_info
|
||||
):
|
||||
yield chunk
|
||||
return
|
||||
|
||||
# 1. Remote prefill
|
||||
self._request_prefill_token_ids(prefill_request)
|
||||
prefill_gen = prefill_handle_method.remote(prefill_request, raw_request_info)
|
||||
prefill_chunk = await prefill_gen.__anext__()
|
||||
|
||||
if isinstance(prefill_chunk, ErrorResponse):
|
||||
logger.error(f"Prefill returned error: {prefill_chunk}")
|
||||
yield prefill_chunk
|
||||
return
|
||||
|
||||
# 2. Local decode via super().chat / super().completions so the
|
||||
# standard LLMServer request pipeline (request_id, LoRA multiplex,
|
||||
# batch_output_stream) runs on the decode side.
|
||||
decode_request = backend.prepare_decode_request(
|
||||
request=request, peer=None, prefill_response=prefill_chunk
|
||||
)
|
||||
# Reuse prefill's ids for this decode so the render skips re-tokenizing.
|
||||
with _reuse_prompt_token_ids(self._decode_reuse_ids(prefill_chunk)):
|
||||
local_gen = await getattr(super(), method)(decode_request, raw_request_info)
|
||||
async for chunk in local_gen:
|
||||
yield chunk
|
||||
|
||||
async def _concurrent_decode(
|
||||
self,
|
||||
method: str,
|
||||
decode_request: RequestType,
|
||||
prefill_resp: AsyncGenerator,
|
||||
raw_request_info: Optional[RawRequestInfo],
|
||||
*,
|
||||
cancel_on_failure: bool = True,
|
||||
):
|
||||
"""Run local decode while a remote prefill drains concurrently.
|
||||
|
||||
While prefill is in flight, each decode chunk is raced against the
|
||||
prefill task so a prefill failure surfaces to the client as an
|
||||
``ErrorResponse`` (instead of a hung — decode may be waiting on KV that
|
||||
will never arrive — or seemingly-successful decode stream). The
|
||||
background prefill task is always awaited so it never leaks on the
|
||||
prefill/decode engines.
|
||||
|
||||
Args:
|
||||
method: The handle method name ("chat" or "completions").
|
||||
decode_request: The request to run on the local decode engine.
|
||||
prefill_resp: The in-flight remote prefill response stream.
|
||||
raw_request_info: Raw HTTP request info forwarded to the engine.
|
||||
cancel_on_failure: Whether to cancel the in-flight prefill when
|
||||
local decode does not complete. Must be False for
|
||||
``dispatch()``-based prefill (the choose_replica path): its
|
||||
completion accounting fires when the response completes, so the
|
||||
stream must be drained to exhaustion, never abandoned. Prefill
|
||||
is clamped to a single token, so draining is bounded either way.
|
||||
"""
|
||||
prefill_task = asyncio.ensure_future(_drain_prefill(prefill_resp))
|
||||
completed = False
|
||||
local_gen = None
|
||||
next_fut = None
|
||||
try:
|
||||
local_gen = await getattr(super(), method)(decode_request, raw_request_info)
|
||||
gen = local_gen.__aiter__()
|
||||
while True:
|
||||
# Surface a failed prefill as soon as it is observed.
|
||||
if prefill_task.done() and isinstance(
|
||||
prefill_task.result(), ErrorResponse
|
||||
):
|
||||
err = prefill_task.result()
|
||||
logger.error("Remote prefill returned error: %s", err)
|
||||
yield err
|
||||
return
|
||||
if next_fut is None:
|
||||
next_fut = asyncio.ensure_future(gen.__anext__())
|
||||
# Race the next decode chunk against the in-flight prefill;
|
||||
# once prefill has completed (successfully), just stream.
|
||||
awaitables = {next_fut}
|
||||
if not prefill_task.done():
|
||||
awaitables.add(prefill_task)
|
||||
done, _ = await asyncio.wait(
|
||||
awaitables, return_when=asyncio.FIRST_COMPLETED
|
||||
)
|
||||
if next_fut in done:
|
||||
try:
|
||||
chunk = next_fut.result()
|
||||
except StopAsyncIteration:
|
||||
break
|
||||
next_fut = None
|
||||
yield chunk
|
||||
# else: prefill finished first; loop back to inspect it.
|
||||
completed = True
|
||||
finally:
|
||||
if next_fut is not None and not next_fut.done():
|
||||
next_fut.cancel()
|
||||
with contextlib.suppress(BaseException):
|
||||
await next_fut
|
||||
if not completed:
|
||||
# Abort the local decode request if we bailed early.
|
||||
if local_gen is not None:
|
||||
with contextlib.suppress(BaseException):
|
||||
await local_gen.aclose()
|
||||
if cancel_on_failure:
|
||||
prefill_task.cancel()
|
||||
try:
|
||||
err = await prefill_task
|
||||
if isinstance(err, ErrorResponse):
|
||||
logger.error("Remote prefill returned error: %s", err)
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
except Exception as exc: # pragma: no cover - defensive
|
||||
logger.error("Remote prefill failed: %s", exc)
|
||||
|
||||
# ---- Direct-streaming ASGI app ----
|
||||
|
||||
async def __serve_build_asgi_app__(self):
|
||||
"""Serve direct-streaming HTTP through P/D orchestration.
|
||||
|
||||
Start from the engine-native app (same as non-P/D direct streaming)
|
||||
and re-point only ``/v1/chat/completions`` and ``/v1/completions`` at
|
||||
this server's ``chat`` / ``completions``, which run remote prefill
|
||||
then local decode. All other routes stay engine-native.
|
||||
"""
|
||||
app = await super().__serve_build_asgi_app__()
|
||||
_strip_routes(app, "/v1/chat/completions")
|
||||
_strip_routes(app, "/v1/completions")
|
||||
|
||||
@app.post("/v1/chat/completions")
|
||||
async def _pd_chat(body: ChatCompletionRequest, request: Request):
|
||||
body = _sanitize_chat_completion_request(body)
|
||||
raw_info = RawRequestInfo.from_starlette_request(request)
|
||||
return await _pd_http_response(await self.chat(body, raw_info))
|
||||
|
||||
@app.post("/v1/completions")
|
||||
async def _pd_completions(body: CompletionRequest, request: Request):
|
||||
raw_info = RawRequestInfo.from_starlette_request(request)
|
||||
return await _pd_http_response(await self.completions(body, raw_info))
|
||||
|
||||
return app
|
||||
|
||||
# ---- Pre-warm ----
|
||||
#
|
||||
# KV transfer connectors (e.g. NIXL) require a handshake between each
|
||||
# prefill and decode replica before real traffic can flow. Pre-warming
|
||||
# sends a tiny dummy request through the full prefill->decode path for
|
||||
# every prefill replica so that the connector establishes its connections
|
||||
# eagerly at startup rather than on the first user request.
|
||||
# Enable via: experimental_configs={"_prewarm_prefill_decode": True}
|
||||
|
||||
def _make_dummy_request(self, model_id: str) -> CompletionRequest:
|
||||
"""Build the smallest valid completion request for pre-warm."""
|
||||
return CompletionRequest(
|
||||
model=model_id,
|
||||
prompt=_PREWARM_PROMPT,
|
||||
max_tokens=_PREWARM_MAX_TOKENS,
|
||||
stream=False,
|
||||
request_id=f"prewarm-{uuid.uuid4()}",
|
||||
)
|
||||
|
||||
async def _maybe_prewarm(self) -> None:
|
||||
"""Run one prefill->decode round-trip per P replica to complete
|
||||
the connector handshake on both sides before traffic arrives."""
|
||||
prewarm_enabled = getattr(
|
||||
self, "_llm_config", None
|
||||
) and self._llm_config.experimental_configs.get("_prewarm_prefill_decode")
|
||||
if not prewarm_enabled:
|
||||
return
|
||||
|
||||
logger.info("[PDDecodeServer] Starting pre-warm across all P replicas.")
|
||||
|
||||
backend = self._get_connector_backend()
|
||||
if backend.requires_peer_binding:
|
||||
# Peer-binding connectors (e.g. MoRIIO) shape a prefill request
|
||||
# against a specific selected replica's metadata; a peerless
|
||||
# broadcast prewarm cannot bind one. The connector handshake
|
||||
# completes on the first real request instead.
|
||||
logger.info(
|
||||
"[PDDecodeServer] Skipping pre-warm: connector %s requires peer "
|
||||
"binding (handshake completes on the first real request).",
|
||||
type(backend).__name__,
|
||||
)
|
||||
return
|
||||
|
||||
model_id = self._llm_config.model_id
|
||||
dummy = self._make_dummy_request(model_id)
|
||||
prefill_req = backend.prepare_prefill_request(request=dummy, peer=None)
|
||||
|
||||
# Broadcast to every live P replica; retry until they are up.
|
||||
kv_params_list: List[Any] = []
|
||||
attempt = 0
|
||||
while attempt < _PREWARM_MAX_RETRIES:
|
||||
attempt += 1
|
||||
try:
|
||||
kv_params_list = await asyncio.get_event_loop().run_in_executor(
|
||||
None,
|
||||
lambda: broadcast(
|
||||
self._prefill_handle,
|
||||
method_name="prewarm_prefill",
|
||||
args=[prefill_req],
|
||||
),
|
||||
)
|
||||
break
|
||||
except DeploymentUnavailableError:
|
||||
logger.info(
|
||||
"[PDDecodeServer] PrefillServer not available yet "
|
||||
"(attempt %d/%d); retrying in %.0fs...",
|
||||
attempt,
|
||||
_PREWARM_MAX_RETRIES,
|
||||
_PREWARM_RETRY_INTERVAL_S,
|
||||
)
|
||||
await asyncio.sleep(_PREWARM_RETRY_INTERVAL_S)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"[PDDecodeServer] broadcast() attempt %d/%d failed; "
|
||||
"retrying in %.0fs...",
|
||||
attempt,
|
||||
_PREWARM_MAX_RETRIES,
|
||||
_PREWARM_RETRY_INTERVAL_S,
|
||||
exc_info=exc,
|
||||
)
|
||||
await asyncio.sleep(_PREWARM_RETRY_INTERVAL_S)
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"[PDDecodeServer] Pre-warm failed after {_PREWARM_MAX_RETRIES} "
|
||||
f"attempts ({_PREWARM_MAX_RETRIES * _PREWARM_RETRY_INTERVAL_S:.0f}s). "
|
||||
f"PrefillServer may be permanently unavailable."
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"[PDDecodeServer] broadcast() reached %d P replica(s); "
|
||||
"driving local decode to complete the handshake.",
|
||||
len(kv_params_list),
|
||||
)
|
||||
|
||||
# Build one decode request per P replica result.
|
||||
decode_reqs: List[CompletionRequest] = []
|
||||
for idx, kv_params in enumerate(kv_params_list):
|
||||
if not kv_params:
|
||||
logger.warning(
|
||||
"[PDDecodeServer] P replica %d returned empty kv_params; skipping.",
|
||||
idx,
|
||||
)
|
||||
continue
|
||||
req = dummy.model_copy(deep=True)
|
||||
req.kv_transfer_params = kv_params
|
||||
decode_reqs.append(req)
|
||||
|
||||
# Run all decode requests on the local engine concurrently to trigger
|
||||
# the connector handshake on D side for each P replica.
|
||||
async def _decode_one(req: CompletionRequest, idx: int) -> None:
|
||||
async for _ in self.engine.completions(req, None):
|
||||
pass
|
||||
logger.info(
|
||||
"[PDDecodeServer] Pre-warm handshake done for P replica %d.", idx
|
||||
)
|
||||
|
||||
await asyncio.gather(*[_decode_one(r, i) for i, r in enumerate(decode_reqs)])
|
||||
|
||||
logger.info("[PDDecodeServer] Pre-warm complete — all P replicas registered.")
|
||||
|
||||
|
||||
async def _drain_prefill(prefill_resp) -> Optional[ErrorResponse]:
|
||||
"""Consume a concurrent-handoff prefill response to completion.
|
||||
|
||||
In concurrent (e.g. WRITE-mode) handoff the remote prefill produces no useful
|
||||
tokens — it only needs to run so the connector pushes/registers the KV. We
|
||||
drain it so the response is fully awaited before the ``choose_replica``
|
||||
context (if any) exits. Returns an ``ErrorResponse`` if one is observed.
|
||||
Handles both streaming (``DeploymentResponseGenerator``) and non-streaming
|
||||
(single ``DeploymentResponse``) results.
|
||||
"""
|
||||
try:
|
||||
async for chunk in prefill_resp:
|
||||
if isinstance(chunk, ErrorResponse):
|
||||
return chunk
|
||||
except TypeError:
|
||||
result = await prefill_resp
|
||||
if isinstance(result, ErrorResponse):
|
||||
return result
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# PDPrefillServer
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class PDPrefillServer(LLMServer):
|
||||
"""Prefill-side LLM server for P/D disaggregation.
|
||||
|
||||
This is a standard LLMServer with an additional ``prewarm_prefill``
|
||||
method used during the pre-warm handshake.
|
||||
"""
|
||||
|
||||
async def record_replica_metadata(self) -> Dict[str, Any]:
|
||||
"""Publish this prefill replica's connector coordination metadata.
|
||||
|
||||
Read by the decode orchestrator via the replica-metadata hook
|
||||
(``ReplicaSelection.replica_metadata``) so peer-binding connectors (e.g.
|
||||
MoRIIO) can address the selected prefill replica. Returns ``{}`` for
|
||||
connectors that publish nothing (the ``BaseConnectorBackend`` default).
|
||||
|
||||
Returns the metadata of the backend that engine init
|
||||
(``setup_engine_backend``) created, ``setup()``-ed, and stored on this
|
||||
server's ``_llm_config``. The replica-metadata hook is captured after
|
||||
engine init, so for connector deployments the backend is present by
|
||||
then; with no backend stored there is nothing to publish.
|
||||
"""
|
||||
backend = getattr(self._llm_config, "kv_connector_backend", None)
|
||||
if backend is None:
|
||||
return {}
|
||||
return backend.replica_metadata()
|
||||
|
||||
async def prewarm_prefill(
|
||||
self, prefill_request: CompletionRequest
|
||||
) -> Optional[dict]:
|
||||
"""Run one prefill pass and return kv_transfer_params as a dict.
|
||||
|
||||
Returns None on error.
|
||||
"""
|
||||
async for chunk in self.engine.completions(prefill_request, None):
|
||||
if hasattr(chunk, "kv_transfer_params") and chunk.kv_transfer_params:
|
||||
return chunk.kv_transfer_params
|
||||
if isinstance(chunk, ErrorResponse):
|
||||
logger.warning("[PDPrefillServer] prewarm_prefill got error: %s", chunk)
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# PDDecodeServer
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class PDDecodeServer(PDOrchestratorMixin, LLMServer):
|
||||
"""Decode-side LLM server that orchestrates remote prefill.
|
||||
|
||||
This deployment owns a real engine (decode config) and holds a handle
|
||||
to the prefill deployment. For chat / completions it runs remote
|
||||
prefill first, then local decode.
|
||||
"""
|
||||
|
||||
async def __init__(
|
||||
self,
|
||||
llm_config: LLMConfig,
|
||||
*,
|
||||
prefill_server: DeploymentHandle,
|
||||
engine_cls=None,
|
||||
model_downloader=None,
|
||||
):
|
||||
self._prefill_handle = prefill_server.options(stream=True)
|
||||
await super().__init__(
|
||||
llm_config,
|
||||
engine_cls=engine_cls,
|
||||
model_downloader=model_downloader,
|
||||
)
|
||||
# Active only if enabled and the renderer wrap installs. The `and`
|
||||
# short-circuits so install() is not called when disabled.
|
||||
self._pd_tokenize_once = (
|
||||
bool(self._llm_config.experimental_configs.get("pd_tokenize_once"))
|
||||
and _install_tokenize_once()
|
||||
)
|
||||
await self._maybe_prewarm()
|
||||
|
||||
async def chat(
|
||||
self,
|
||||
request: ChatCompletionRequest,
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[Union[str, ChatCompletionResponse, ErrorResponse], None]:
|
||||
return self._pd_handle_request(request, raw_request_info)
|
||||
|
||||
async def completions(
|
||||
self,
|
||||
request: CompletionRequest,
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[Union[str, CompletionResponse, ErrorResponse], None]:
|
||||
return self._pd_handle_request(request, raw_request_info)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# DP + PD combined servers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class DPPDPrefillServer(PDPrefillServer, DPServer):
|
||||
"""PDPrefillServer with data-parallel gang scheduling.
|
||||
|
||||
MRO: DPPDPrefillServer -> PDPrefillServer -> DPServer -> LLMServer
|
||||
- get_deployment_options comes from DPServer (adds gang scheduling).
|
||||
- __init__ falls through to DPServer (DP master info, bundle indices)
|
||||
then LLMServer (engine setup).
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class DPPDDecodeServer(PDDecodeServer, DPServer):
|
||||
"""PDDecodeServer with data-parallel gang scheduling.
|
||||
|
||||
MRO: DPPDDecodeServer -> PDDecodeServer -> PDOrchestratorMixin
|
||||
-> DPServer -> LLMServer
|
||||
- get_deployment_options comes from DPServer (adds gang scheduling).
|
||||
- __init__ from PDDecodeServer sets _prefill_handle, then super().__init__
|
||||
flows through DPServer (DP setup) then LLMServer (engine setup).
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Deprecated: PDProxyServer
|
||||
# TODO(Kourosh): Deprecate, remove in Ray 2.58.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class PDProxyServer(LLMServerProtocol):
|
||||
"""Proxy between P/D LLM servers.
|
||||
|
||||
.. deprecated::
|
||||
``PDProxyServer`` is deprecated. Use ``PDDecodeServer`` instead.
|
||||
This class will be removed in a future release.
|
||||
"""
|
||||
|
||||
async def __init__(
|
||||
self,
|
||||
prefill_server: DeploymentHandle,
|
||||
decode_server: DeploymentHandle,
|
||||
):
|
||||
warnings.warn(
|
||||
"PDProxyServer is deprecated and will be removed in Ray 2.58. "
|
||||
"Use PDDecodeServer (decode orchestrator) and PDPrefillServer instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
self._llm_config = await prefill_server.llm_config.remote()
|
||||
self.prefill_server = prefill_server.options(stream=True)
|
||||
self.decode_server = decode_server.options(stream=True)
|
||||
|
||||
async def start(self) -> None:
|
||||
pass
|
||||
|
||||
async def check_health(self) -> None:
|
||||
pass
|
||||
|
||||
async def reset_prefix_cache(self) -> None:
|
||||
raise NotImplementedError(
|
||||
"reset_prefix_cache is not supported for P/D disaggregation"
|
||||
)
|
||||
|
||||
async def start_profile(self) -> None:
|
||||
raise NotImplementedError(
|
||||
"start_profile is not supported for P/D disaggregation"
|
||||
)
|
||||
|
||||
async def stop_profile(self) -> None:
|
||||
raise NotImplementedError(
|
||||
"stop_profile is not supported for P/D disaggregation"
|
||||
)
|
||||
|
||||
async def llm_config(self) -> Optional[LLMConfig]:
|
||||
return self._llm_config
|
||||
|
||||
def _prepare_prefill_request(self, request: RequestType) -> RequestType:
|
||||
assert (
|
||||
getattr(request, "kv_transfer_params", None) is None
|
||||
), "kv_transfer_params should be empty before proxy"
|
||||
prefill_request = request.model_copy(deep=True)
|
||||
prefill_request.kv_transfer_params = {
|
||||
"do_remote_decode": True,
|
||||
"do_remote_prefill": False,
|
||||
"remote_engine_id": None,
|
||||
"remote_block_ids": None,
|
||||
"remote_host": None,
|
||||
"remote_port": None,
|
||||
}
|
||||
prefill_request.max_tokens = 1
|
||||
prefill_request.stream = False
|
||||
return prefill_request
|
||||
|
||||
def _prepare_decode_request(
|
||||
self,
|
||||
request: RequestType,
|
||||
prefill_chunk: Union[ChatCompletionResponse, CompletionResponse],
|
||||
) -> RequestType:
|
||||
decode_request = request.model_copy(deep=True)
|
||||
decode_request.kv_transfer_params = prefill_chunk.kv_transfer_params
|
||||
return decode_request
|
||||
|
||||
def _maybe_add_request_id_to_request(
|
||||
self,
|
||||
request: Union[ChatCompletionRequest, CompletionRequest],
|
||||
) -> None:
|
||||
request_id = get_serve_request_id()
|
||||
if request_id:
|
||||
request.request_id = request_id
|
||||
|
||||
async def _handle_request(
|
||||
self,
|
||||
request: RequestType,
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[
|
||||
Union[str, ChatCompletionResponse, CompletionResponse, ErrorResponse], None
|
||||
]:
|
||||
self._maybe_add_request_id_to_request(request)
|
||||
|
||||
if isinstance(request, ChatCompletionRequest):
|
||||
method = "chat"
|
||||
elif isinstance(request, CompletionRequest):
|
||||
method = "completions"
|
||||
else:
|
||||
raise ValueError(f"Unsupported request type: {type(request)}")
|
||||
|
||||
prefill_request = self._prepare_prefill_request(request)
|
||||
prefill_gen = getattr(self.prefill_server, method).remote(
|
||||
prefill_request, raw_request_info
|
||||
)
|
||||
prefill_chunk = await prefill_gen.__anext__()
|
||||
|
||||
if isinstance(prefill_chunk, ErrorResponse):
|
||||
logger.error(f"Prefill returned error: {prefill_chunk}")
|
||||
yield prefill_chunk
|
||||
return
|
||||
|
||||
decode_request = self._prepare_decode_request(request, prefill_chunk)
|
||||
decode_gen = getattr(self.decode_server, method).remote(
|
||||
decode_request, raw_request_info
|
||||
)
|
||||
async for chunk in decode_gen:
|
||||
yield chunk
|
||||
|
||||
async def chat(
|
||||
self,
|
||||
request: ChatCompletionRequest,
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[Union[str, ChatCompletionResponse, ErrorResponse], None]:
|
||||
return self._handle_request(request, raw_request_info)
|
||||
|
||||
async def completions(
|
||||
self,
|
||||
request: CompletionRequest,
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[Union[str, CompletionResponse, ErrorResponse], None]:
|
||||
return self._handle_request(request, raw_request_info)
|
||||
|
||||
async def embeddings(
|
||||
self,
|
||||
request: EmbeddingRequest,
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[EmbeddingResponse, None]:
|
||||
raise NotImplementedError("Embedding is not supported for P/D disaggregation")
|
||||
|
||||
@classmethod
|
||||
def get_deployment_options(
|
||||
cls, prefill_config: "LLMConfig", decode_config: "LLMConfig"
|
||||
) -> Dict[str, Any]:
|
||||
return DEFAULT_PD_PROXY_SERVER_OPTIONS
|
||||
Reference in New Issue
Block a user