[llm.serving] Reconfigure router to better perform under high concurrency (#50876) (#50884)

Cherry pick: #50876 Signed-off-by: Gene Su <e870252314@gmail.com>

Author: GeneDer

python/ray/llm/_internal/serve/builders/application_builders.py

@@ -64,4 +64,6 @@ def build_openai_app(llm_serving_args: LLMServingArgs) -> Application:
 
  llm_deployments = _get_llm_deployments(llm_configs)
 
- return LLMRouter.as_deployment().bind(llm_deployments=llm_deployments)
+ return LLMRouter.as_deployment(llm_configs=llm_configs).bind(
+ llm_deployments=llm_deployments
+ )

python/ray/llm/_internal/serve/configs/constants.py

@@ -66,3 +66,21 @@
 ENABLE_VERBOSE_TELEMETRY = bool(int(os.getenv("RAYLLM_ENABLE_VERBOSE_TELEMETRY", "0")))
 
 RAYLLM_VLLM_ENGINE_CLS_ENV = "RAYLLM_VLLM_ENGINE_CLS"
+
+# The ratio of number of router replicas to number of model replicas. Default to 2
+# meaning that there are 2 router replicas for every model replica.
+ROUTER_TO_MODEL_REPLICA_RATIO = float(
+ os.getenv("RAYLLM_ROUTER_TO_MODEL_REPLICA_RATIO", "2")
+)
+
+RAYLLM_ROUTER_MIN_REPLICAS = int(os.environ.get("RAYLLM_ROUTER_MIN_REPLICAS", 0))
+RAYLLM_ROUTER_INITIAL_REPLICAS = int(
+ os.environ.get("RAYLLM_ROUTER_INITIAL_REPLICAS", 2)
+)
+RAYLLM_ROUTER_MAX_REPLICAS = int(os.environ.get("RAYLLM_ROUTER_MAX_REPLICAS", 16))
+RAYLLM_ROUTER_TARGET_ONGOING_REQUESTS = int(
+ os.environ.get(
+ "RAYLLM_ROUTER_TARGET_ONGOING_REQUESTS",
+ DEFAULT_TARGET_ONGOING_REQUESTS, # 16
+ )
+)

python/ray/llm/_internal/serve/deployments/routers/router.py

@@ -22,7 +22,14 @@
 from ray.serve.handle import DeploymentHandle
 from starlette.responses import JSONResponse, Response, StreamingResponse
 
-from ray.llm._internal.serve.configs.constants import RAYLLM_ROUTER_HTTP_TIMEOUT
+from ray.llm._internal.serve.configs.constants import (
+ RAYLLM_ROUTER_HTTP_TIMEOUT,
+ ROUTER_TO_MODEL_REPLICA_RATIO,
+ RAYLLM_ROUTER_MIN_REPLICAS,
+ RAYLLM_ROUTER_INITIAL_REPLICAS,
+ RAYLLM_ROUTER_MAX_REPLICAS,
+ RAYLLM_ROUTER_TARGET_ONGOING_REQUESTS,
+)
 from ray.llm._internal.serve.observability.logging import get_logger
 from ray.llm._internal.serve.observability.metrics.fast_api_metrics import (
  add_http_metrics_middleware,
@@ -52,6 +59,7 @@
  LLMConfig,
  ModelData,
  Model,
+ AutoscalingConfig,
 )
 from ray.llm._internal.serve.deployments.routers.middleware import (
  SetRequestIdMiddleware,
@@ -397,30 +405,53 @@ async def chat(self, body: ChatCompletionRequest) -> Response:
  return JSONResponse(content=result.model_dump())
 
  @classmethod
- def as_deployment(cls) -> serve.Deployment:
+ def as_deployment(
+ cls, llm_configs: Optional[List[LLMConfig]] = None
+ ) -> serve.Deployment:
  """Converts this class to a Ray Serve deployment with ingress.
 
  Returns:
  A Ray Serve deployment.
  """
+ min_replicas = RAYLLM_ROUTER_MIN_REPLICAS
+ initial_replicas = RAYLLM_ROUTER_INITIAL_REPLICAS
+ max_replicas = RAYLLM_ROUTER_MAX_REPLICAS
+
+ # Note (genesu): Based on our internal benchmark, we are currently bottleneck
+ # by the router replicas during high concurrency situation. We are setting the
+ # router replicas to be ~2x the total model replicas and making it scale faster.
+ if llm_configs:
+ model_min_replicas = 0
+ model_initial_replicas = 0
+ model_max_replicas = 0
+ for llm_config in llm_configs:
+ if "autoscaling_config" in llm_config.deployment_config:
+ autoscaling_config = llm_config.deployment_config[
+ "autoscaling_config"
+ ]
+ if isinstance(autoscaling_config, dict):
+ autoscaling_config = AutoscalingConfig(
+ **llm_config.deployment_config["autoscaling_config"]
+ )
+ else:
+ # When autoscaling config is not provided, we use the default.
+ autoscaling_config = AutoscalingConfig()
+ model_min_replicas += autoscaling_config.min_replicas
+ model_initial_replicas += autoscaling_config.initial_replicas
+ model_max_replicas += autoscaling_config.max_replicas
+ min_replicas = int(model_min_replicas * ROUTER_TO_MODEL_REPLICA_RATIO)
+ initial_replicas = int(
+ model_initial_replicas * ROUTER_TO_MODEL_REPLICA_RATIO
+ )
+ max_replicas = int(model_max_replicas * ROUTER_TO_MODEL_REPLICA_RATIO)
 
  ingress_cls = serve.ingress(fastapi_router_app)(cls)
  deployment_decorator = serve.deployment(
- # TODO (Kourosh): make this configurable
  autoscaling_config={
- "min_replicas": int(os.environ.get("RAYLLM_ROUTER_MIN_REPLICAS", 0)),
- "initial_replicas": int(
- os.environ.get("RAYLLM_ROUTER_INITIAL_REPLICAS", 2)
- ),
- "max_replicas": int(os.environ.get("RAYLLM_ROUTER_MAX_REPLICAS", 16)),
- "target_ongoing_requests": int(
- os.environ.get(
- "RAYLLM_ROUTER_TARGET_ONGOING_REQUESTS",
- os.environ.get(
- "RAYLLM_ROUTER_TARGET_NUM_ONGOING_REQUESTS_PER_REPLICA", 200
- ),
- )
- ),
+ "min_replicas": min_replicas,
+ "initial_replicas": initial_replicas,
+ "max_replicas": max_replicas,
+ "target_ongoing_requests": RAYLLM_ROUTER_TARGET_ONGOING_REQUESTS,
  },
  ray_actor_options=json.loads(
  os.environ.get("RAYLLM_ROUTER_RAY_ACTOR_OPTIONS", "{}")

python/ray/llm/tests/serve/builders/test_application_builders.py

@@ -3,11 +3,17 @@
 
 from ray.llm._internal.serve.configs.server_models import (
  LLMServingArgs,
+ LLMConfig,
+ AutoscalingConfig,
+ ModelLoadingConfig,
 )
 from ray.llm._internal.serve.builders.application_builders import (
  build_openai_app,
  build_vllm_deployment,
 )
+from ray.llm._internal.serve.configs.constants import (
+ RAYLLM_ROUTER_TARGET_ONGOING_REQUESTS,
+)
 import subprocess
 import yaml
 import os
@@ -94,6 +100,51 @@ def deployments_healthy():
  p.send_signal(signal.SIGINT) # Equivalent to ctrl-C
  p.wait()
 
+ def test_router_built_with_autoscaling_configs(self):
+ """Test that the router is built with the correct autoscaling configs that
+ will scale.
+ """
+ llm_config_no_autoscaling_configured = LLMConfig(
+ model_loading_config=ModelLoadingConfig(model_id="model_id_1"),
+ accelerator_type="L4",
+ )
+ llm_config_autoscaling_default = LLMConfig(
+ model_loading_config=ModelLoadingConfig(model_id="model_id_2"),
+ accelerator_type="L4",
+ deployment_config={"autoscaling_config": AutoscalingConfig()},
+ )
+ llm_config_autoscaling_non_default = LLMConfig(
+ model_loading_config=ModelLoadingConfig(model_id="model_id_3"),
+ accelerator_type="L4",
+ deployment_config={
+ "autoscaling_config": AutoscalingConfig(
+ min_replicas=2,
+ initial_replicas=3,
+ max_replicas=4,
+ )
+ },
+ )
+
+ app = build_openai_app(
+ LLMServingArgs(
+ llm_configs=[
+ llm_config_no_autoscaling_configured,
+ llm_config_autoscaling_default,
+ llm_config_autoscaling_non_default,
+ ]
+ )
+ )
+ router_autoscaling_config = (
+ app._bound_deployment._deployment_config.autoscaling_config
+ )
+ assert router_autoscaling_config.min_replicas == 8 # (1 + 1 + 2) * 2
+ assert router_autoscaling_config.initial_replicas == 10 # (1 + 1 + 3) * 2
+ assert router_autoscaling_config.max_replicas == 408 # (100 + 100 + 4) * 2
+ assert (
+ router_autoscaling_config.target_ongoing_requests
+ == RAYLLM_ROUTER_TARGET_ONGOING_REQUESTS
+ )
+
 
 class TestBuildVllmDeployment:
  def test_build_vllm_deployment(