chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,14 @@
|
||||
from ray.llm._internal.common.observability.logging_utils import (
|
||||
disable_vllm_custom_ops_logger_on_cpu_nodes,
|
||||
)
|
||||
from ray.llm._internal.serve.observability import setup_observability
|
||||
|
||||
# Set up observability
|
||||
disable_vllm_custom_ops_logger_on_cpu_nodes()
|
||||
setup_observability()
|
||||
|
||||
|
||||
def _worker_process_setup_hook():
|
||||
"""Noop setup hook used for ENABLE_WORKER_PROCESS_SETUP_HOOK
|
||||
(see python/ray/llm/_internal/serve/configs/constants.py)."""
|
||||
pass
|
||||
@@ -0,0 +1,196 @@
|
||||
# Multi-Turn LLM Benchmark
|
||||
|
||||
A benchmark tool for OpenAI-compatible LLM inference servers that supports
|
||||
multi-turn conversations with configurable prefix cache hit rates, input/output
|
||||
sequence lengths, and cross-session prefix sharing.
|
||||
|
||||
## Entry Point
|
||||
|
||||
```
|
||||
python -m ray.llm._internal.serve.benchmark.cli [OPTIONS]
|
||||
```
|
||||
|
||||
## Modes
|
||||
|
||||
| Command | Mode | Description |
|
||||
|---------|------|-------------|
|
||||
| `... -s` | Smoke | Single request health check |
|
||||
| `... --concurrency 8 ...` | Direct (concurrency) | Closed-loop concurrency benchmark |
|
||||
| `... --request-rate 10 ...` | Direct (rate) | Constant-QPS benchmark |
|
||||
| `... -i` | Interactive server | Long-running server with UNIX socket control |
|
||||
| `... -i --client` | Interactive client | Connect to server; REPL or `--cmd` one-shot |
|
||||
|
||||
## Quick Examples
|
||||
|
||||
### Smoke test
|
||||
|
||||
```bash
|
||||
python -m ray.llm._internal.serve.benchmark.cli -s \
|
||||
-u http://localhost:8000 -m my-model
|
||||
```
|
||||
|
||||
### Concurrency benchmark
|
||||
|
||||
```bash
|
||||
python -m ray.llm._internal.serve.benchmark.cli \
|
||||
-u http://localhost:8000 -m meta-llama/Llama-3-8B-Instruct \
|
||||
--concurrency 8 --num-sessions 200 \
|
||||
--isl 2000 --osl 200 --hit-rate 0.85 --num-turns 5 \
|
||||
--think-time 1.0 --save-result results.json
|
||||
```
|
||||
|
||||
### Rate benchmark
|
||||
|
||||
```bash
|
||||
python -m ray.llm._internal.serve.benchmark.cli \
|
||||
-u http://localhost:8000 -m meta-llama/Llama-3-8B-Instruct \
|
||||
--request-rate 10 --duration 120 \
|
||||
--isl 2000 --osl 200 --hit-rate 0.85 --num-turns 5 \
|
||||
--warm-up 10 --save-result results.json
|
||||
```
|
||||
|
||||
### Interactive server
|
||||
|
||||
```bash
|
||||
python -m ray.llm._internal.serve.benchmark.cli -i \
|
||||
-u http://localhost:8000 -m meta-llama/Llama-3-8B-Instruct \
|
||||
--isl 2000 --osl 200 --hit-rate 0.85 --num-turns 5
|
||||
```
|
||||
|
||||
### Interactive client (REPL)
|
||||
|
||||
```bash
|
||||
python -m ray.llm._internal.serve.benchmark.cli -i --client
|
||||
```
|
||||
|
||||
### Interactive client (one-shot)
|
||||
|
||||
```bash
|
||||
python -m ray.llm._internal.serve.benchmark.cli -i --client --cmd "rate 10"
|
||||
python -m ray.llm._internal.serve.benchmark.cli -i --client --cmd "status"
|
||||
```
|
||||
|
||||
## Workload Parameters
|
||||
|
||||
All workload parameters use **simple mode**: you specify user-facing values
|
||||
and the tool derives internal parameters (per-turn user tokens `u` and
|
||||
system prompt tokens `s`) automatically.
|
||||
|
||||
| Parameter | Flag | Description |
|
||||
|-----------|------|-------------|
|
||||
| ISL | `--isl` | Average input sequence length (tokens) across all turns |
|
||||
| OSL | `--osl` | Output tokens per turn |
|
||||
| Hit rate | `--hit-rate` | Target prefix cache hit rate [0, 1] |
|
||||
| Shared system prompt ratio | `--shared-system-prompt-ratio` | Fraction of system prompt shared across sessions (default: 0.0) |
|
||||
| Num turns | `--num-turns` | Number of turns per conversation session |
|
||||
| Think time | `--think-time` | Simulated user think-time between turns in seconds (default: 0) |
|
||||
| First chunk threshold | `--first-chunk-threshold` | Number of SSE content chunks before recording first-chunk latency (default: 16) |
|
||||
|
||||
The solver derives `user_tokens` (new user tokens per turn) and `sys_tokens`
|
||||
(total system prompt tokens) from these inputs. The `print_summary()` output
|
||||
shows the resolved per-turn token breakdown including cached vs. new tokens at
|
||||
each turn.
|
||||
|
||||
## Tokenizer
|
||||
|
||||
By default, `--tokenizer` is `None`, which causes the tool to use the
|
||||
`--model` value as the HuggingFace tokenizer name. This works when `--model`
|
||||
is a valid HuggingFace model ID (e.g., `meta-llama/Llama-3-8B-Instruct`).
|
||||
|
||||
Provide `--tokenizer` explicitly when:
|
||||
- The `--model` value is an alias or deployment name that is not a valid
|
||||
HuggingFace repo (e.g., `--model my-deployment --tokenizer meta-llama/Llama-3-8B-Instruct`).
|
||||
- You want to use a local tokenizer path.
|
||||
|
||||
## Warm-Up Strategies
|
||||
|
||||
### Concurrency mode
|
||||
|
||||
Warm-up is **automatic** using entropy-based detection. The tool monitors the
|
||||
distribution of active turns across concurrent sessions. Once the Shannon
|
||||
entropy of the turn distribution reaches 50% of its theoretical maximum, the
|
||||
pool is considered at steady state and measurement begins. All requests
|
||||
dispatched before that point are discarded.
|
||||
|
||||
### Rate mode
|
||||
|
||||
Warm-up is **time-based** via the `--warm-up` flag (in seconds). All requests
|
||||
whose dispatch time falls within the warm-up window are excluded from reported
|
||||
metrics. Set this to allow the server's KV cache to fill and stabilize.
|
||||
|
||||
### Interactive mode
|
||||
|
||||
Warm-up is **manual**. The operator starts traffic with `rate <qps>`, waits
|
||||
for the system to stabilize, then explicitly starts a measurement window with
|
||||
`start` or `measure <n>`.
|
||||
|
||||
## Interactive Commands
|
||||
|
||||
| Command | Description |
|
||||
|---------|-------------|
|
||||
| `help` | Show available commands |
|
||||
| `rate <qps>` | Set target request rate (0 to pause) |
|
||||
| `start` | Start open-ended measurement window |
|
||||
| `measure <n>` | Start measurement capturing next `n` completed requests |
|
||||
| `stop` | Stop measurement and print summary |
|
||||
| `status` | Show current state: QPS, inflight, completed, measured |
|
||||
| `workload [k=v ...]` | Show or update workload parameters (e.g., `workload isl=3000 osl=300`) |
|
||||
| `save [path]` | Save last measurement window to JSON |
|
||||
| `save-dir <path>` | Set default directory for saved results |
|
||||
| `quit` | Stop the benchmark server |
|
||||
|
||||
## JSON Result Schema
|
||||
|
||||
Results saved with `--save-result` (direct mode) contain these top-level keys:
|
||||
|
||||
| Key | Description |
|
||||
|-----|-------------|
|
||||
| `config` | Run configuration (concurrency/rate, model, etc.) |
|
||||
| `spec` | Resolved workload spec with per-turn token breakdown |
|
||||
| `first_chunk_threshold` | Number of chunks before recording first-chunk latency |
|
||||
| `benchmark` | Run metadata: total requests, duration, warm-up info |
|
||||
| `stats` | Aggregate latency statistics (avg, P50, P90, P99 for TTFT, FC, TPOT, latency) |
|
||||
| `per_turn` | Per-turn breakdown of count, avg ISL, and latency percentiles |
|
||||
| `raw_metrics` | Array of per-request metrics (session_id, turn, all latency fields, token counts) |
|
||||
|
||||
Interactive mode saves with `save` produce a similar structure with a `window`
|
||||
summary instead of `benchmark`/`stats`/`per_turn`.
|
||||
|
||||
## Typical Workflow
|
||||
|
||||
1. **Smoke test** to verify connectivity:
|
||||
```bash
|
||||
python -m ray.llm._internal.serve.benchmark.cli -s -u http://localhost:8000 -m my-model
|
||||
```
|
||||
|
||||
2. **Direct benchmark** for a fixed workload:
|
||||
```bash
|
||||
python -m ray.llm._internal.serve.benchmark.cli \
|
||||
--concurrency 8 --num-sessions 200 \
|
||||
--isl 2000 --osl 200 --hit-rate 0.85 --num-turns 5 \
|
||||
-u http://localhost:8000 -m meta-llama/Llama-3-8B-Instruct \
|
||||
--save-result concurrency_8.json
|
||||
```
|
||||
|
||||
3. **Interactive mode** for exploratory testing:
|
||||
```bash
|
||||
# Terminal 1: start server
|
||||
python -m ray.llm._internal.serve.benchmark.cli -i \
|
||||
--isl 2000 --osl 200 --hit-rate 0.85 --num-turns 5 \
|
||||
-u http://localhost:8000 -m meta-llama/Llama-3-8B-Instruct
|
||||
|
||||
# Terminal 2: control
|
||||
python -m ray.llm._internal.serve.benchmark.cli -i --client
|
||||
benchctl> rate 5
|
||||
benchctl> measure 500
|
||||
benchctl> status
|
||||
benchctl> save results_qps5.json
|
||||
benchctl> rate 10
|
||||
benchctl> measure 500
|
||||
benchctl> save results_qps10.json
|
||||
benchctl> quit
|
||||
```
|
||||
|
||||
4. **Sweep** over multiple configurations: write an external script that loops
|
||||
over the CLI with different parameters. The tool does not include built-in
|
||||
sweep orchestration.
|
||||
@@ -0,0 +1,7 @@
|
||||
"""CLI entry point for the multi-turn OpenAI-compatible HTTP benchmark.
|
||||
|
||||
Example: python -m ray.llm._internal.serve.benchmark --help
|
||||
"""
|
||||
from ray.llm._internal.serve.benchmark.cli import main
|
||||
|
||||
main()
|
||||
@@ -0,0 +1,233 @@
|
||||
"""CLI entry point for the multi-turn OpenAI-compatible HTTP benchmark."""
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="python -m ray.llm._internal.serve.benchmark",
|
||||
description="Multi-turn OpenAI-compatible HTTP benchmark",
|
||||
)
|
||||
|
||||
## Mode flags ##
|
||||
mode = parser.add_argument_group("mode")
|
||||
mode.add_argument(
|
||||
"-s",
|
||||
"--smoke",
|
||||
action="store_true",
|
||||
help="Smoke test (single request, exit)",
|
||||
)
|
||||
mode.add_argument(
|
||||
"-i",
|
||||
"--interactive",
|
||||
action="store_true",
|
||||
help="Interactive mode (server by default)",
|
||||
)
|
||||
mode.add_argument(
|
||||
"--client",
|
||||
action="store_true",
|
||||
help="Interactive client mode (used with -i)",
|
||||
)
|
||||
|
||||
## Server / API ##
|
||||
server = parser.add_argument_group("server/API")
|
||||
server.add_argument(
|
||||
"-u",
|
||||
"--base-url",
|
||||
default="http://127.0.0.1:8000",
|
||||
help="Base URL of the OpenAI-compatible API (default: %(default)s)",
|
||||
)
|
||||
server.add_argument(
|
||||
"-m",
|
||||
"--model",
|
||||
default=None,
|
||||
help="Model name to send in requests (required except for -i --client)",
|
||||
)
|
||||
server.add_argument(
|
||||
"--tokenizer",
|
||||
default=None,
|
||||
help="HuggingFace tokenizer name/path (default: same as --model)",
|
||||
)
|
||||
server.add_argument(
|
||||
"--api-key",
|
||||
default=None,
|
||||
help="API key for Authorization header (default: None)",
|
||||
)
|
||||
|
||||
## Workload ##
|
||||
workload = parser.add_argument_group("workload")
|
||||
workload.add_argument(
|
||||
"--isl",
|
||||
type=int,
|
||||
default=1000,
|
||||
help="Average input sequence length (default: %(default)s)",
|
||||
)
|
||||
workload.add_argument(
|
||||
"--hit-rate",
|
||||
type=float,
|
||||
default=0.5,
|
||||
help="Prefix cache hit rate [0, 1] (default: %(default)s)",
|
||||
)
|
||||
workload.add_argument(
|
||||
"--num-turns",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of turns per session (default: %(default)s)",
|
||||
)
|
||||
workload.add_argument(
|
||||
"--osl",
|
||||
type=int,
|
||||
default=100,
|
||||
help="Output tokens per turn (default: %(default)s)",
|
||||
)
|
||||
workload.add_argument(
|
||||
"--shared-system-prompt-ratio",
|
||||
dest="shared_system_prompt_ratio",
|
||||
type=float,
|
||||
default=1.0,
|
||||
help="Fraction of the system prompt shared across all sessions "
|
||||
"(1.0 = identical, 0.0 = all unique) (default: %(default)s)",
|
||||
)
|
||||
workload.add_argument(
|
||||
"--think-time",
|
||||
type=float,
|
||||
default=0.0,
|
||||
help="Simulated user think-time between turns in seconds (default: %(default)s)",
|
||||
)
|
||||
workload.add_argument(
|
||||
"-fc",
|
||||
"--first-chunk-threshold",
|
||||
type=int,
|
||||
default=16,
|
||||
help="Number of content chunks before recording first-chunk latency (default: %(default)s)",
|
||||
)
|
||||
|
||||
## Traffic ##
|
||||
traffic = parser.add_argument_group("traffic")
|
||||
traffic.add_argument(
|
||||
"--concurrency",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Number of concurrent sessions",
|
||||
)
|
||||
traffic.add_argument(
|
||||
"--request-rate",
|
||||
type=float,
|
||||
default=None,
|
||||
help="Request rate (requests per second)",
|
||||
)
|
||||
traffic.add_argument(
|
||||
"--duration",
|
||||
type=float,
|
||||
default=None,
|
||||
help="Duration in seconds",
|
||||
)
|
||||
traffic.add_argument(
|
||||
"--num-sessions",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Total number of sessions to run",
|
||||
)
|
||||
traffic.add_argument(
|
||||
"--warm-up",
|
||||
type=float,
|
||||
default=0,
|
||||
help="Warm-up period in seconds (default: %(default)s)",
|
||||
)
|
||||
traffic.add_argument(
|
||||
"--warmup-jitter-max",
|
||||
type=float,
|
||||
default=10.0,
|
||||
help="Max random delay (seconds) between turns during entropy warm-up "
|
||||
"in concurrency mode. Jitter desynchronizes sessions so the benchmark "
|
||||
"reaches steady-state faster (default: %(default)s)",
|
||||
)
|
||||
traffic.add_argument(
|
||||
"--ramp-interval",
|
||||
type=float,
|
||||
default=-1,
|
||||
help="Seconds between launching successive sessions at benchmark start. "
|
||||
"Use this to avoid a thundering-herd of simultaneous first requests. "
|
||||
"-1 = auto-derive from request rate or concurrency (default: %(default)s)",
|
||||
)
|
||||
|
||||
## Interactive-only ##
|
||||
interactive = parser.add_argument_group("interactive-only")
|
||||
interactive.add_argument(
|
||||
"--status-interval",
|
||||
type=int,
|
||||
default=5,
|
||||
help="Status reporting interval in seconds (default: %(default)s)",
|
||||
)
|
||||
interactive.add_argument(
|
||||
"--cmd",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Command to send in interactive client mode",
|
||||
)
|
||||
interactive.add_argument(
|
||||
"--log-failures",
|
||||
action="store_true",
|
||||
help="Log individual request failures",
|
||||
)
|
||||
interactive.add_argument(
|
||||
"--seed",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Random seed for reproducibility",
|
||||
)
|
||||
interactive.add_argument(
|
||||
"--save-result",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Filename to save results",
|
||||
)
|
||||
interactive.add_argument(
|
||||
"--save-dir",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Directory to save results",
|
||||
)
|
||||
interactive.add_argument(
|
||||
"--num-workers",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of process-pool workers for conversation generation (default: %(default)s)",
|
||||
)
|
||||
|
||||
return parser
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = build_parser()
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.interactive and args.client:
|
||||
from ray.llm._internal.serve.benchmark.interactive import run_interactive_client
|
||||
|
||||
sys.exit(run_interactive_client(args))
|
||||
|
||||
# All other modes require --model
|
||||
if not args.model:
|
||||
parser.error("--model is required (except for -i --client mode)")
|
||||
|
||||
if args.smoke:
|
||||
from ray.llm._internal.serve.benchmark.runners import run_smoke
|
||||
|
||||
sys.exit(run_smoke(args))
|
||||
elif args.interactive:
|
||||
from ray.llm._internal.serve.benchmark.interactive import run_interactive_server
|
||||
|
||||
sys.exit(run_interactive_server(args))
|
||||
elif args.concurrency or args.request_rate:
|
||||
from ray.llm._internal.serve.benchmark.runners import run_direct
|
||||
|
||||
sys.exit(run_direct(args))
|
||||
else:
|
||||
parser.print_help()
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,122 @@
|
||||
"""HTTP client for OpenAI-compatible chat completion endpoints."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
import aiohttp
|
||||
|
||||
from ray.llm._internal.serve.benchmark.models import TurnResult
|
||||
|
||||
|
||||
async def send_chat_completion(
|
||||
session: aiohttp.ClientSession,
|
||||
base_url: str,
|
||||
model: str,
|
||||
messages: list[dict[str, str]],
|
||||
session_id: str = "",
|
||||
max_tokens: int = 256,
|
||||
first_chunk_threshold: int = 16,
|
||||
timeout_sec: int = 300,
|
||||
api_key: Optional[str] = None,
|
||||
) -> TurnResult:
|
||||
"""Send a streaming chat completion request and collect metrics."""
|
||||
url = f"{base_url}/v1/chat/completions"
|
||||
payload = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"max_tokens": max_tokens,
|
||||
"stream": True,
|
||||
"stream_options": {"include_usage": True},
|
||||
"temperature": 0.0,
|
||||
}
|
||||
|
||||
headers: dict[str, str] = {
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
if api_key:
|
||||
headers["Authorization"] = f"Bearer {api_key}"
|
||||
if session_id:
|
||||
headers["X-Session-Id"] = session_id
|
||||
|
||||
timeout = aiohttp.ClientTimeout(total=timeout_sec)
|
||||
|
||||
start_ns = time.perf_counter_ns()
|
||||
ttft_ns: Optional[int] = None
|
||||
fc_ns: Optional[int] = None
|
||||
content_chunk_count = 0
|
||||
chunk_times: list[int] = []
|
||||
generated_text = ""
|
||||
input_tokens = 0
|
||||
output_tokens = 0
|
||||
prev_ts = start_ns
|
||||
|
||||
async with session.post(
|
||||
url, json=payload, headers=headers, timeout=timeout
|
||||
) as resp:
|
||||
if resp.status != 200:
|
||||
body = await resp.text()
|
||||
raise RuntimeError(f"HTTP {resp.status}: {body[:500]}")
|
||||
|
||||
async for raw_line in resp.content:
|
||||
line = raw_line.strip()
|
||||
if not line:
|
||||
continue
|
||||
text = line.decode("utf-8", errors="replace")
|
||||
if not text.startswith("data: "):
|
||||
continue
|
||||
data_str = text[6:]
|
||||
if data_str == "[DONE]":
|
||||
continue
|
||||
|
||||
try:
|
||||
data = json.loads(data_str)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
usage = data.get("usage")
|
||||
if usage:
|
||||
input_tokens = usage.get("prompt_tokens", input_tokens)
|
||||
output_tokens = usage.get("completion_tokens", output_tokens)
|
||||
|
||||
choices = data.get("choices", [])
|
||||
if not choices:
|
||||
continue
|
||||
|
||||
delta = choices[0].get("delta", {})
|
||||
content = delta.get("content") or delta.get("reasoning")
|
||||
if content:
|
||||
now_ns = time.perf_counter_ns()
|
||||
content_chunk_count += 1
|
||||
if ttft_ns is None:
|
||||
ttft_ns = now_ns - start_ns
|
||||
else:
|
||||
chunk_times.append(now_ns - prev_ts)
|
||||
if fc_ns is None and content_chunk_count >= first_chunk_threshold:
|
||||
fc_ns = now_ns - start_ns
|
||||
prev_ts = now_ns
|
||||
generated_text += content
|
||||
|
||||
end_ns = time.perf_counter_ns()
|
||||
latency_ns = end_ns - start_ns
|
||||
|
||||
if ttft_ns is None:
|
||||
ttft_ns = latency_ns
|
||||
if fc_ns is None:
|
||||
fc_ns = latency_ns
|
||||
|
||||
itl_ms_list = [t / 1e6 for t in chunk_times]
|
||||
itl_ms = sum(itl_ms_list) / len(itl_ms_list) if itl_ms_list else 0.0
|
||||
|
||||
return TurnResult(
|
||||
ttft_ms=ttft_ns / 1e6,
|
||||
fc_ms=fc_ns / 1e6,
|
||||
itl_ms=itl_ms,
|
||||
e2e_latency_ms=latency_ns / 1e6,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
generated_text=generated_text,
|
||||
itl_ms_list=itl_ms_list,
|
||||
)
|
||||
@@ -0,0 +1,824 @@
|
||||
"""Interactive server and client for the multi-turn benchmark.
|
||||
|
||||
The interactive server runs a long-lived benchmark loop whose QPS, workload
|
||||
parameters, and measurement windows are controlled at runtime via a UNIX
|
||||
domain socket. The interactive client connects to that socket (either as an
|
||||
interactive REPL or for one-shot commands).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import aiohttp
|
||||
import numpy as np
|
||||
|
||||
from ray.llm._internal.serve.benchmark.metrics import (
|
||||
serialize_raw_metrics,
|
||||
summarize_metrics,
|
||||
)
|
||||
from ray.llm._internal.serve.benchmark.models import TurnMetric, WorkloadSpec
|
||||
from ray.llm._internal.serve.benchmark.text_gen import (
|
||||
Conversation,
|
||||
TextGenerator,
|
||||
conversation_factory,
|
||||
)
|
||||
from ray.llm._internal.serve.benchmark.turn import execute_single_turn
|
||||
|
||||
try:
|
||||
from prompt_toolkit import PromptSession
|
||||
from prompt_toolkit.history import FileHistory
|
||||
except ImportError:
|
||||
PromptSession = None # type: ignore[assignment,misc]
|
||||
FileHistory = None # type: ignore[assignment,misc]
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Control socket path
|
||||
# ---------------------------------------------------------------------------
|
||||
_DEFAULT_CONTROL_SOCKET = "/tmp/interactive_rate_bench.sock"
|
||||
|
||||
|
||||
def _control_socket_path() -> str:
|
||||
return os.environ.get("RAY_BENCH_CONTROL_SOCKET", _DEFAULT_CONTROL_SOCKET)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Process-pool worker helpers (module-level so they are picklable)
|
||||
# ---------------------------------------------------------------------------
|
||||
_worker_tokenizer = None
|
||||
_worker_text_gen: Optional[TextGenerator] = None
|
||||
|
||||
|
||||
def _pool_initializer(tokenizer_name: str, base_seed: int) -> None:
|
||||
"""Called once per worker process to load the tokenizer and seed RNG."""
|
||||
global _worker_tokenizer, _worker_text_gen
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
_worker_tokenizer = AutoTokenizer.from_pretrained(
|
||||
tokenizer_name, trust_remote_code=True
|
||||
)
|
||||
_worker_text_gen = TextGenerator(_worker_tokenizer)
|
||||
proc_seed = (base_seed + os.getpid()) % (2**32)
|
||||
random.seed(proc_seed)
|
||||
np.random.seed(proc_seed)
|
||||
|
||||
|
||||
def _create_conv_in_worker(
|
||||
session_idx: int,
|
||||
spec: WorkloadSpec,
|
||||
shared_system_text: str,
|
||||
) -> Conversation:
|
||||
"""Create a Conversation inside a worker process."""
|
||||
return conversation_factory(session_idx, spec, shared_system_text, _worker_text_gen)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Interactive-mode runtime state & helpers
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@dataclass
|
||||
class RuntimeState:
|
||||
current_qps: float = 0.0
|
||||
total_completed: int = 0
|
||||
total_failed: int = 0
|
||||
inflight: int = 0
|
||||
measurement_active: bool = False
|
||||
measurement_start_ns: Optional[int] = None
|
||||
measurement_metrics: list[TurnMetric] = field(default_factory=list)
|
||||
measurement_target_requests: Optional[int] = None
|
||||
last_window_metrics: list[TurnMetric] = field(default_factory=list)
|
||||
last_window_elapsed_s: float = 0.0
|
||||
last_notice: Optional[str] = None
|
||||
save_dir: Optional[str] = None
|
||||
|
||||
|
||||
def _save_window_result(
|
||||
path: str,
|
||||
args: argparse.Namespace,
|
||||
spec: WorkloadSpec,
|
||||
metrics: list[TurnMetric],
|
||||
elapsed_s: float,
|
||||
runtime_qps: float = 0.0,
|
||||
) -> None:
|
||||
payload = {
|
||||
"mode": "interactive_rate",
|
||||
"saved_at_epoch_s": time.time(),
|
||||
"config": {
|
||||
"base_url": args.base_url,
|
||||
"model": args.model,
|
||||
"tokenizer": getattr(args, "tokenizer", None) or args.model,
|
||||
"first_chunk_threshold": args.first_chunk_threshold,
|
||||
"num_turns": args.num_turns,
|
||||
"osl": args.osl,
|
||||
"shared_system_prompt_ratio": args.shared_system_prompt_ratio,
|
||||
"isl": args.isl,
|
||||
"hit_rate": args.hit_rate,
|
||||
"runtime_qps": runtime_qps,
|
||||
},
|
||||
"spec": spec.summary(),
|
||||
"window": summarize_metrics(metrics, elapsed_s),
|
||||
"raw_metrics": serialize_raw_metrics(metrics),
|
||||
}
|
||||
p = Path(path)
|
||||
p.parent.mkdir(parents=True, exist_ok=True)
|
||||
with p.open("w") as f:
|
||||
json.dump(payload, f, indent=2)
|
||||
print(f"Saved measurement window to {path}")
|
||||
|
||||
|
||||
def _build_spec(
|
||||
args: argparse.Namespace, overrides: Optional[dict] = None
|
||||
) -> WorkloadSpec:
|
||||
"""Build and resolve a WorkloadSpec from args, optionally merging overrides."""
|
||||
kw = dict(
|
||||
num_sessions=1,
|
||||
duration_s=1.0,
|
||||
num_turns=args.num_turns,
|
||||
osl=args.osl,
|
||||
think_time=0.0,
|
||||
concurrency=None,
|
||||
request_rate=1.0,
|
||||
ramp_interval=0.0,
|
||||
shared_system_prompt_ratio=args.shared_system_prompt_ratio,
|
||||
isl=args.isl,
|
||||
hit_rate=args.hit_rate,
|
||||
)
|
||||
if overrides:
|
||||
kw.update(overrides)
|
||||
spec = WorkloadSpec(**kw)
|
||||
spec.resolve()
|
||||
return spec
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Command handler (extracted for testability)
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class CommandHandler:
|
||||
"""Handles interactive benchmark commands.
|
||||
|
||||
Extracted from the ``run_interactive`` closure so that command parsing,
|
||||
state mutation, and response formatting can be unit-tested without
|
||||
starting a real server or HTTP session.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
runtime: RuntimeState,
|
||||
workload: dict,
|
||||
args: argparse.Namespace,
|
||||
text_gen: Optional[TextGenerator] = None,
|
||||
rate_changed: Optional[asyncio.Event] = None,
|
||||
workload_changed: Optional[asyncio.Event] = None,
|
||||
stop_event: Optional[asyncio.Event] = None,
|
||||
):
|
||||
self.runtime = runtime
|
||||
self.workload = workload
|
||||
self.args = args
|
||||
self.text_gen = text_gen
|
||||
self.rate_changed = rate_changed or asyncio.Event()
|
||||
self.workload_changed = workload_changed or asyncio.Event()
|
||||
self.stop_event = stop_event or asyncio.Event()
|
||||
|
||||
def resolve_save_path(self, raw: Optional[str]) -> str:
|
||||
if raw:
|
||||
expanded = str(Path(raw).expanduser())
|
||||
if "/" in expanded or expanded.startswith("."):
|
||||
return expanded
|
||||
return str(Path(self.runtime.save_dir) / expanded)
|
||||
|
||||
if self.args.save_result:
|
||||
return str(Path(self.args.save_result).expanduser())
|
||||
|
||||
ts = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
qps_label = f"{self.runtime.current_qps:.2f}".replace(".", "p")
|
||||
return str(
|
||||
Path(self.runtime.save_dir)
|
||||
/ f"interactive_measure_qps{qps_label}_{ts}.json"
|
||||
)
|
||||
|
||||
async def handle(self, cmd: str) -> str: # noqa: C901
|
||||
"""Process a single command string and return the response."""
|
||||
cmd = cmd.strip()
|
||||
if not cmd:
|
||||
return "empty command"
|
||||
|
||||
parts = cmd.split()
|
||||
op = parts[0].lower()
|
||||
|
||||
if op == "help":
|
||||
return (
|
||||
"Commands: help, rate <qps>, start, measure <n>, stop, "
|
||||
"status, save [path|name], save-dir <path>, quit\n"
|
||||
"Workload: workload [isl=N] [osl=N] [hit-rate=F] "
|
||||
"[sharing=F] [num-turns=N]\n"
|
||||
" e.g. workload isl=2000 osl=200 hit-rate=0.5\n"
|
||||
" All params optional; unspecified ones keep their current values.\n"
|
||||
" workload (no args) prints current workload spec."
|
||||
)
|
||||
if op == "rate":
|
||||
if len(parts) != 2:
|
||||
return "Usage: rate <qps>"
|
||||
try:
|
||||
new_qps = float(parts[1])
|
||||
if new_qps < 0:
|
||||
raise ValueError()
|
||||
except ValueError:
|
||||
return "QPS must be a non-negative number."
|
||||
self.runtime.current_qps = new_qps
|
||||
self.rate_changed.set()
|
||||
return f"Set target qps={new_qps:.3f}"
|
||||
if op == "start":
|
||||
self.runtime.measurement_active = True
|
||||
self.runtime.measurement_start_ns = time.perf_counter_ns()
|
||||
self.runtime.measurement_metrics = []
|
||||
self.runtime.measurement_target_requests = None
|
||||
self.runtime.last_notice = None
|
||||
return "Measurement started."
|
||||
if op == "measure":
|
||||
if len(parts) != 2:
|
||||
return "Usage: measure <num_requests>"
|
||||
try:
|
||||
tgt = int(parts[1])
|
||||
if tgt <= 0:
|
||||
raise ValueError()
|
||||
except ValueError:
|
||||
return "measure requires a positive integer."
|
||||
self.runtime.measurement_active = True
|
||||
self.runtime.measurement_start_ns = time.perf_counter_ns()
|
||||
self.runtime.measurement_metrics = []
|
||||
self.runtime.measurement_target_requests = tgt
|
||||
self.runtime.last_notice = None
|
||||
return f"Measurement started: capturing next {tgt} completed requests."
|
||||
if op == "stop":
|
||||
if not self.runtime.measurement_active:
|
||||
return "Measurement is not active."
|
||||
self.runtime.measurement_active = False
|
||||
end_ns = time.perf_counter_ns()
|
||||
start_ns = self.runtime.measurement_start_ns or end_ns
|
||||
self.runtime.last_window_elapsed_s = (end_ns - start_ns) / 1e9
|
||||
self.runtime.last_window_metrics = list(self.runtime.measurement_metrics)
|
||||
self.runtime.measurement_target_requests = None
|
||||
summary = summarize_metrics(
|
||||
list(self.runtime.last_window_metrics),
|
||||
self.runtime.last_window_elapsed_s,
|
||||
)
|
||||
return f"Measurement stopped.\n{json.dumps(summary, indent=2)}"
|
||||
if op == "status":
|
||||
cur = self.workload["spec"]
|
||||
status = (
|
||||
f"qps={self.runtime.current_qps:.2f} "
|
||||
f"inflight={self.runtime.inflight} "
|
||||
f"completed={self.runtime.total_completed} "
|
||||
f"failed={self.runtime.total_failed} "
|
||||
f"measured={len(self.runtime.measurement_metrics)} "
|
||||
f"active={self.runtime.measurement_active} "
|
||||
f"target={self.runtime.measurement_target_requests} "
|
||||
f"save_dir={self.runtime.save_dir}\n"
|
||||
f"workload: isl={cur.isl} osl={cur.osl} hit-rate={cur.hit_rate} "
|
||||
f"sharing={cur.shared_system_prompt_ratio} num-turns={cur.num_turns}"
|
||||
)
|
||||
if self.runtime.last_notice:
|
||||
status += f"\n{self.runtime.last_notice}"
|
||||
self.runtime.last_notice = None
|
||||
return status
|
||||
if op == "save-dir":
|
||||
if len(parts) != 2:
|
||||
return "Usage: save-dir <path>"
|
||||
new_dir = str(Path(parts[1]).expanduser())
|
||||
self.runtime.save_dir = new_dir
|
||||
return f"Set save_dir={self.runtime.save_dir}"
|
||||
if op == "save":
|
||||
if len(parts) > 2:
|
||||
return "Usage: save [path.json|name.json]"
|
||||
|
||||
if (
|
||||
self.runtime.measurement_active
|
||||
and self.runtime.measurement_start_ns is not None
|
||||
):
|
||||
el = (time.perf_counter_ns() - self.runtime.measurement_start_ns) / 1e9
|
||||
mlist = list(self.runtime.measurement_metrics)
|
||||
else:
|
||||
el = self.runtime.last_window_elapsed_s
|
||||
mlist = list(self.runtime.last_window_metrics)
|
||||
if not mlist:
|
||||
return "No measured window data to save."
|
||||
save_path = self.resolve_save_path(parts[1] if len(parts) == 2 else None)
|
||||
_save_window_result(
|
||||
save_path,
|
||||
self.args,
|
||||
self.workload["spec"],
|
||||
mlist,
|
||||
el,
|
||||
runtime_qps=self.runtime.current_qps,
|
||||
)
|
||||
return f"Saved measurement window to {save_path}"
|
||||
if op == "workload":
|
||||
cur = self.workload["spec"]
|
||||
if len(parts) == 1:
|
||||
return (
|
||||
f"isl={cur.isl} osl={cur.osl} hit-rate={cur.hit_rate} "
|
||||
f"sharing={cur.shared_system_prompt_ratio} num-turns={cur.num_turns}"
|
||||
)
|
||||
_param_aliases = {
|
||||
"isl": "isl",
|
||||
"osl": "osl",
|
||||
"hit-rate": "hit_rate",
|
||||
"hitrate": "hit_rate",
|
||||
"hit_rate": "hit_rate",
|
||||
"sharing": "shared_system_prompt_ratio",
|
||||
"shared-system-prompt-ratio": "shared_system_prompt_ratio",
|
||||
"shared_system_prompt_ratio": "shared_system_prompt_ratio",
|
||||
"num-turns": "num_turns",
|
||||
"num_turns": "num_turns",
|
||||
}
|
||||
overrides: dict = {}
|
||||
errors: list[str] = []
|
||||
for token in parts[1:]:
|
||||
if "=" not in token:
|
||||
errors.append(f"bad token {token!r} (expected key=value)")
|
||||
continue
|
||||
k, _, v = token.partition("=")
|
||||
mapped = _param_aliases.get(k.lower())
|
||||
if mapped is None:
|
||||
errors.append(f"unknown param {k!r}")
|
||||
continue
|
||||
try:
|
||||
overrides[mapped] = (
|
||||
int(v) if mapped in ("isl", "osl", "num_turns") else float(v)
|
||||
)
|
||||
except ValueError:
|
||||
errors.append(f"invalid value for {k}: {v!r}")
|
||||
if errors:
|
||||
return "Error: " + "; ".join(errors)
|
||||
merged = dict(
|
||||
isl=cur.isl,
|
||||
osl=cur.osl,
|
||||
hit_rate=cur.hit_rate,
|
||||
shared_system_prompt_ratio=cur.shared_system_prompt_ratio,
|
||||
num_turns=cur.num_turns,
|
||||
)
|
||||
merged.update(overrides)
|
||||
try:
|
||||
new_spec = _build_spec(self.args, merged)
|
||||
except Exception as e:
|
||||
return f"Invalid workload spec: {e}"
|
||||
if self.text_gen is not None:
|
||||
new_sst = self.text_gen.generate(new_spec.shared_s)
|
||||
else:
|
||||
new_sst = ""
|
||||
self.workload["spec"] = new_spec
|
||||
self.workload["shared_system_text"] = new_sst
|
||||
self.workload_changed.set()
|
||||
new_spec.print_summary()
|
||||
return (
|
||||
f"Workload updated: isl={new_spec.isl} osl={new_spec.osl} "
|
||||
f"hit-rate={new_spec.hit_rate} "
|
||||
f"sharing={new_spec.shared_system_prompt_ratio} "
|
||||
f"num-turns={new_spec.num_turns}"
|
||||
)
|
||||
if op in ("quit", "exit"):
|
||||
self.stop_event.set()
|
||||
return "Stopping benchmark..."
|
||||
return f"Unknown command: {op}"
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Interactive server
|
||||
# ============================================================================
|
||||
|
||||
|
||||
async def run_interactive(args: argparse.Namespace) -> None:
|
||||
spec = _build_spec(args)
|
||||
spec.print_summary()
|
||||
print("Interactive mode: starts idle. Use 'rate <qps>' to begin sending traffic.")
|
||||
|
||||
from concurrent.futures import ProcessPoolExecutor
|
||||
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
tokenizer_name: str = args.tokenizer if args.tokenizer else args.model
|
||||
|
||||
if args.seed is None:
|
||||
args.seed = random.randint(0, 2**31 - 1)
|
||||
random.seed(args.seed)
|
||||
np.random.seed(args.seed % (2**32))
|
||||
print(f"Seed: {args.seed}")
|
||||
|
||||
print(f"Loading tokenizer: {tokenizer_name}")
|
||||
tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, trust_remote_code=True)
|
||||
text_gen = TextGenerator(tokenizer)
|
||||
|
||||
shared_system_text = text_gen.generate(spec.shared_s)
|
||||
bench_start_ns = time.perf_counter_ns()
|
||||
|
||||
workload: dict = {"spec": spec, "shared_system_text": shared_system_text}
|
||||
workload_changed = asyncio.Event()
|
||||
|
||||
default_save_dir = args.save_dir
|
||||
if default_save_dir is None and args.save_result:
|
||||
default_save_dir = str(Path(args.save_result).parent)
|
||||
if default_save_dir is None:
|
||||
default_save_dir = os.getcwd()
|
||||
|
||||
runtime = RuntimeState(
|
||||
current_qps=0.0,
|
||||
save_dir=str(Path(default_save_dir).expanduser()),
|
||||
)
|
||||
stop_event = asyncio.Event()
|
||||
rate_changed = asyncio.Event()
|
||||
ready_queue: asyncio.Queue[tuple[Conversation, int]] = asyncio.Queue()
|
||||
next_session_idx = 0
|
||||
running_tasks: set[asyncio.Task] = set()
|
||||
|
||||
num_workers = args.num_workers
|
||||
print(f"Starting process pool with {num_workers} workers")
|
||||
cpu_pool = ProcessPoolExecutor(
|
||||
max_workers=num_workers,
|
||||
initializer=_pool_initializer,
|
||||
initargs=(tokenizer_name, args.seed),
|
||||
)
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
def _next_session_idx() -> int:
|
||||
nonlocal next_session_idx
|
||||
idx = next_session_idx
|
||||
next_session_idx += 1
|
||||
return idx
|
||||
|
||||
async def next_conv_async() -> Conversation:
|
||||
idx = _next_session_idx()
|
||||
s = workload["spec"]
|
||||
sst = workload["shared_system_text"]
|
||||
return await loop.run_in_executor(
|
||||
cpu_pool,
|
||||
_create_conv_in_worker,
|
||||
idx,
|
||||
s,
|
||||
sst,
|
||||
)
|
||||
|
||||
async def prefill_queue() -> None:
|
||||
while not stop_event.is_set():
|
||||
if workload_changed.is_set():
|
||||
workload_changed.clear()
|
||||
drained = 0
|
||||
while not ready_queue.empty():
|
||||
try:
|
||||
ready_queue.get_nowait()
|
||||
drained += 1
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
if drained:
|
||||
print(
|
||||
f"[workload] drained {drained} stale conversations from queue.",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
qps = runtime.current_qps
|
||||
if qps <= 0:
|
||||
await asyncio.sleep(0.2)
|
||||
continue
|
||||
s = workload["spec"]
|
||||
sst = workload["shared_system_text"]
|
||||
target = max(8, int(qps * 2))
|
||||
current = ready_queue.qsize()
|
||||
if current < target:
|
||||
batch_size = min(target - current, num_workers * 2)
|
||||
idxs = [_next_session_idx() for _ in range(batch_size)]
|
||||
futs = [
|
||||
loop.run_in_executor(
|
||||
cpu_pool,
|
||||
_create_conv_in_worker,
|
||||
idx,
|
||||
s,
|
||||
sst,
|
||||
)
|
||||
for idx in idxs
|
||||
]
|
||||
for fut in asyncio.as_completed(futs):
|
||||
try:
|
||||
conv = await fut
|
||||
await ready_queue.put((conv, 0))
|
||||
except Exception as e:
|
||||
logger.warning("Failed to create conversation in worker: %s", e)
|
||||
await asyncio.sleep(0.02)
|
||||
|
||||
async def execute_turn(
|
||||
conv: Conversation, turn_idx: int, http_session: aiohttp.ClientSession
|
||||
) -> None:
|
||||
cur_spec = workload["spec"]
|
||||
runtime.inflight += 1
|
||||
try:
|
||||
outcome = await execute_single_turn(
|
||||
http_session=http_session,
|
||||
conv=conv,
|
||||
turn_idx=turn_idx,
|
||||
base_url=args.base_url,
|
||||
model=args.model,
|
||||
max_tokens=cur_spec.osl,
|
||||
bench_start_ns=bench_start_ns,
|
||||
first_chunk_threshold=args.first_chunk_threshold,
|
||||
api_key=getattr(args, "api_key", None),
|
||||
)
|
||||
metric = outcome.metric
|
||||
auto_complete_summary: Optional[str] = None
|
||||
runtime.total_completed += 1
|
||||
if runtime.measurement_active:
|
||||
target = runtime.measurement_target_requests
|
||||
if target is None:
|
||||
runtime.measurement_metrics.append(metric)
|
||||
elif len(runtime.measurement_metrics) < target:
|
||||
runtime.measurement_metrics.append(metric)
|
||||
|
||||
if target is not None and len(runtime.measurement_metrics) >= target:
|
||||
runtime.measurement_active = False
|
||||
end_ns = time.perf_counter_ns()
|
||||
start_ns = runtime.measurement_start_ns or end_ns
|
||||
runtime.last_window_elapsed_s = (end_ns - start_ns) / 1e9
|
||||
runtime.last_window_metrics = list(
|
||||
runtime.measurement_metrics[:target]
|
||||
)
|
||||
runtime.measurement_target_requests = None
|
||||
summary = summarize_metrics(
|
||||
runtime.last_window_metrics,
|
||||
runtime.last_window_elapsed_s,
|
||||
)
|
||||
auto_complete_summary = json.dumps(summary, indent=2)
|
||||
runtime.last_notice = (
|
||||
f"measurement auto-complete ({target} req):\n"
|
||||
f"{auto_complete_summary}"
|
||||
)
|
||||
|
||||
if auto_complete_summary is not None:
|
||||
print("Measurement auto-complete:")
|
||||
print(auto_complete_summary)
|
||||
|
||||
next_turn = turn_idx + 1
|
||||
if not stop_event.is_set():
|
||||
if next_turn < cur_spec.num_turns:
|
||||
await ready_queue.put((conv, next_turn))
|
||||
else:
|
||||
conv = await next_conv_async()
|
||||
await ready_queue.put((conv, 0))
|
||||
except Exception as e:
|
||||
if args.log_failures:
|
||||
print(
|
||||
f"[request-failed] session={conv.session_id} turn={turn_idx}: {e}"
|
||||
)
|
||||
runtime.total_failed += 1
|
||||
if not stop_event.is_set():
|
||||
conv = await next_conv_async()
|
||||
await ready_queue.put((conv, 0))
|
||||
finally:
|
||||
runtime.inflight -= 1
|
||||
|
||||
async def pacer(http_session: aiohttp.ClientSession) -> None:
|
||||
next_dispatch = time.perf_counter()
|
||||
while not stop_event.is_set():
|
||||
qps = runtime.current_qps
|
||||
|
||||
if qps <= 0:
|
||||
await asyncio.sleep(0.1)
|
||||
next_dispatch = time.perf_counter() + 0.05
|
||||
continue
|
||||
|
||||
if rate_changed.is_set():
|
||||
rate_changed.clear()
|
||||
next_dispatch = time.perf_counter() + (1.0 / qps)
|
||||
|
||||
now = time.perf_counter()
|
||||
if next_dispatch < now - 1.0:
|
||||
next_dispatch = now
|
||||
|
||||
wait = next_dispatch - now
|
||||
if wait > 0:
|
||||
await asyncio.sleep(wait)
|
||||
if stop_event.is_set():
|
||||
break
|
||||
|
||||
try:
|
||||
conv, turn_idx = ready_queue.get_nowait()
|
||||
except asyncio.QueueEmpty:
|
||||
next_dispatch += 1.0 / qps
|
||||
continue
|
||||
|
||||
t = asyncio.create_task(execute_turn(conv, turn_idx, http_session))
|
||||
running_tasks.add(t)
|
||||
t.add_done_callback(running_tasks.discard)
|
||||
next_dispatch += 1.0 / qps
|
||||
|
||||
async def reporter() -> None:
|
||||
if args.status_interval <= 0:
|
||||
return
|
||||
while not stop_event.is_set():
|
||||
await asyncio.sleep(args.status_interval)
|
||||
print(
|
||||
"status: "
|
||||
f"qps={runtime.current_qps:.2f} "
|
||||
f"inflight={runtime.inflight} "
|
||||
f"completed={runtime.total_completed} "
|
||||
f"failed={runtime.total_failed} "
|
||||
f"measured={len(runtime.measurement_metrics)} "
|
||||
f"active={runtime.measurement_active}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
cmd_handler = CommandHandler(
|
||||
runtime=runtime,
|
||||
workload=workload,
|
||||
args=args,
|
||||
text_gen=text_gen,
|
||||
rate_changed=rate_changed,
|
||||
workload_changed=workload_changed,
|
||||
stop_event=stop_event,
|
||||
)
|
||||
handle_command = cmd_handler.handle
|
||||
|
||||
async def stdin_command_loop() -> None:
|
||||
print(
|
||||
"Interactive commands: help | rate <qps> | start | measure <n> | "
|
||||
"stop | status | workload [k=v ...] | save [path|name] | "
|
||||
"save-dir <path> | quit"
|
||||
)
|
||||
while not stop_event.is_set():
|
||||
raw = await asyncio.to_thread(input, "bench> ")
|
||||
resp = await handle_command(raw)
|
||||
if resp:
|
||||
print(resp)
|
||||
|
||||
async def socket_command_handler(
|
||||
reader: asyncio.StreamReader, writer: asyncio.StreamWriter
|
||||
) -> None:
|
||||
try:
|
||||
data = await reader.read(4096)
|
||||
cmd = data.decode("utf-8", errors="replace").strip()
|
||||
resp = await handle_command(cmd)
|
||||
writer.write((resp + "\n").encode("utf-8"))
|
||||
await writer.drain()
|
||||
finally:
|
||||
writer.close()
|
||||
await writer.wait_closed()
|
||||
|
||||
# Seed the ready queue
|
||||
seed_count = 2
|
||||
seed_idxs = [_next_session_idx() for _ in range(seed_count)]
|
||||
seed_futs = [
|
||||
loop.run_in_executor(
|
||||
cpu_pool,
|
||||
_create_conv_in_worker,
|
||||
idx,
|
||||
spec,
|
||||
shared_system_text,
|
||||
)
|
||||
for idx in seed_idxs
|
||||
]
|
||||
for conv in await asyncio.gather(*seed_futs):
|
||||
await ready_queue.put((conv, 0))
|
||||
|
||||
control_socket = _control_socket_path()
|
||||
if os.path.exists(control_socket):
|
||||
os.unlink(control_socket)
|
||||
socket_server = await asyncio.start_unix_server(
|
||||
socket_command_handler,
|
||||
path=control_socket,
|
||||
)
|
||||
print(f"Control socket listening at: {control_socket}")
|
||||
|
||||
stdin_control = getattr(args, "stdin_control", False)
|
||||
if not stdin_control:
|
||||
print("Use a second terminal with --client to send commands.")
|
||||
|
||||
connector = aiohttp.TCPConnector(limit=0)
|
||||
async with aiohttp.ClientSession(connector=connector) as http_session:
|
||||
background_tasks = [
|
||||
asyncio.create_task(pacer(http_session)),
|
||||
asyncio.create_task(reporter()),
|
||||
asyncio.create_task(socket_server.serve_forever()),
|
||||
asyncio.create_task(prefill_queue()),
|
||||
]
|
||||
stdin_task = (
|
||||
asyncio.create_task(stdin_command_loop()) if stdin_control else None
|
||||
)
|
||||
|
||||
if stdin_task is not None:
|
||||
await stdin_task
|
||||
else:
|
||||
await stop_event.wait()
|
||||
|
||||
stop_event.set()
|
||||
socket_server.close()
|
||||
await socket_server.wait_closed()
|
||||
await asyncio.gather(*background_tasks, return_exceptions=True)
|
||||
|
||||
if running_tasks:
|
||||
print(f"Waiting for {len(running_tasks)} in-flight request task(s)...")
|
||||
await asyncio.gather(*list(running_tasks), return_exceptions=True)
|
||||
|
||||
cpu_pool.shutdown(wait=False)
|
||||
if os.path.exists(control_socket):
|
||||
os.unlink(control_socket)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Interactive client
|
||||
# ============================================================================
|
||||
|
||||
|
||||
async def _send_command_once(control_socket: str, cmd: str) -> str:
|
||||
reader, writer = await asyncio.open_unix_connection(control_socket)
|
||||
writer.write(cmd.encode("utf-8"))
|
||||
await writer.drain()
|
||||
if writer.can_write_eof():
|
||||
writer.write_eof()
|
||||
data = await reader.read()
|
||||
writer.close()
|
||||
await writer.wait_closed()
|
||||
return data.decode("utf-8", errors="replace").strip()
|
||||
|
||||
|
||||
async def run_client(args: argparse.Namespace) -> None:
|
||||
control_socket = _control_socket_path()
|
||||
print(f"Connected to control socket: {control_socket}")
|
||||
print(
|
||||
"Type commands: help, rate <qps>, start, measure <n>, stop, status, "
|
||||
"save [path|name], save-dir <path>, quit"
|
||||
)
|
||||
session = None
|
||||
if PromptSession is not None and FileHistory is not None:
|
||||
history_path = str(Path("~/.interactive_rate_bench_history").expanduser())
|
||||
session = PromptSession(history=FileHistory(history_path))
|
||||
else:
|
||||
print(
|
||||
"prompt_toolkit not installed; using basic input(). "
|
||||
"Install with: pip install prompt_toolkit"
|
||||
)
|
||||
while True:
|
||||
if session is not None:
|
||||
raw = await asyncio.to_thread(session.prompt, "benchctl> ")
|
||||
else:
|
||||
raw = await asyncio.to_thread(input, "benchctl> ")
|
||||
cmd = raw.strip()
|
||||
if not cmd:
|
||||
continue
|
||||
try:
|
||||
resp = await _send_command_once(control_socket, cmd)
|
||||
except (FileNotFoundError, ConnectionRefusedError) as e:
|
||||
print(f"Failed to connect to server socket: {e}")
|
||||
return
|
||||
print(resp)
|
||||
if cmd.lower() in ("quit", "exit"):
|
||||
return
|
||||
|
||||
|
||||
async def run_client_oneshot(args: argparse.Namespace) -> None:
|
||||
if not args.cmd:
|
||||
raise ValueError("Client mode with --cmd requires a command string.")
|
||||
control_socket = _control_socket_path()
|
||||
try:
|
||||
resp = await _send_command_once(control_socket, args.cmd)
|
||||
except (FileNotFoundError, ConnectionRefusedError) as e:
|
||||
raise RuntimeError(f"Failed to connect to server socket: {e}") from e
|
||||
print(resp)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Entry points for cli.py
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def run_interactive_server(args: argparse.Namespace) -> int:
|
||||
"""Entry point for interactive server mode."""
|
||||
try:
|
||||
asyncio.run(run_interactive(args))
|
||||
return 0
|
||||
except Exception as e:
|
||||
logger.error("Interactive server failed: %s", e)
|
||||
return 1
|
||||
|
||||
|
||||
def run_interactive_client(args: argparse.Namespace) -> int:
|
||||
"""Entry point for interactive client mode."""
|
||||
try:
|
||||
if args.cmd:
|
||||
asyncio.run(run_client_oneshot(args))
|
||||
else:
|
||||
asyncio.run(run_client(args))
|
||||
return 0
|
||||
except Exception as e:
|
||||
logger.error("Interactive client failed: %s", e)
|
||||
return 1
|
||||
@@ -0,0 +1,81 @@
|
||||
"""Metrics computation and serialization for the benchmark."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from statistics import mean
|
||||
|
||||
import numpy as np
|
||||
|
||||
from ray.llm._internal.serve.benchmark.models import TurnMetric
|
||||
|
||||
|
||||
def percentile(values: list[float], p: float) -> float:
|
||||
"""Compute the p-th percentile (0-100)."""
|
||||
if not values:
|
||||
return 0.0
|
||||
return float(np.percentile(values, p))
|
||||
|
||||
|
||||
def summarize_metrics(metrics: list[TurnMetric], elapsed_s: float) -> dict:
|
||||
"""Compute aggregate statistics from a list of TurnMetrics.
|
||||
|
||||
ITL (inter-token latency) statistics are computed from raw per-token values
|
||||
flattened across all requests, capturing the full distribution including variance.
|
||||
"""
|
||||
if not metrics:
|
||||
return {"requests": 0, "elapsed_s": round(elapsed_s, 2)}
|
||||
|
||||
ttft = [m.ttft_ms for m in metrics]
|
||||
fc = [m.fc_ms for m in metrics]
|
||||
# Flatten per-token ITL values across all requests for accurate distribution stats
|
||||
itl_all = [v for m in metrics for v in m.itl_ms_list]
|
||||
latency = [m.e2e_latency_ms for m in metrics]
|
||||
out_tok = [m.output_tokens for m in metrics]
|
||||
in_tok = [m.input_tokens for m in metrics]
|
||||
total_output_tokens = sum(out_tok)
|
||||
|
||||
return {
|
||||
"requests": len(metrics),
|
||||
"elapsed_s": round(elapsed_s, 2),
|
||||
"request_rate": round(len(metrics) / elapsed_s, 2) if elapsed_s > 0 else 0.0,
|
||||
"throughput_tok_s": round(total_output_tokens / elapsed_s, 1)
|
||||
if elapsed_s > 0
|
||||
else 0.0,
|
||||
"avg_input_tokens": round(mean(in_tok), 1),
|
||||
"avg_output_tokens": round(mean(out_tok), 1),
|
||||
"avg_ttft_ms": round(mean(ttft), 2),
|
||||
"p50_ttft_ms": round(percentile(ttft, 50), 2),
|
||||
"p90_ttft_ms": round(percentile(ttft, 90), 2),
|
||||
"p99_ttft_ms": round(percentile(ttft, 99), 2),
|
||||
"avg_fc_ms": round(mean(fc), 2),
|
||||
"p50_fc_ms": round(percentile(fc, 50), 2),
|
||||
"p90_fc_ms": round(percentile(fc, 90), 2),
|
||||
"p99_fc_ms": round(percentile(fc, 99), 2),
|
||||
"avg_itl_ms": round(float(np.mean(itl_all)), 2) if itl_all else 0.0,
|
||||
"std_itl_ms": round(float(np.std(itl_all)), 2) if itl_all else 0.0,
|
||||
"p50_itl_ms": round(percentile(itl_all, 50), 2) if itl_all else 0.0,
|
||||
"p90_itl_ms": round(percentile(itl_all, 90), 2) if itl_all else 0.0,
|
||||
"p99_itl_ms": round(percentile(itl_all, 99), 2) if itl_all else 0.0,
|
||||
"avg_e2e_latency_ms": round(mean(latency), 2),
|
||||
"p50_e2e_latency_ms": round(percentile(latency, 50), 2),
|
||||
"p90_e2e_latency_ms": round(percentile(latency, 90), 2),
|
||||
"p99_e2e_latency_ms": round(percentile(latency, 99), 2),
|
||||
}
|
||||
|
||||
|
||||
def serialize_raw_metrics(metrics: list[TurnMetric]) -> list[dict]:
|
||||
"""Serialize TurnMetrics to dicts suitable for JSON output."""
|
||||
return [
|
||||
{
|
||||
"session_id": m.session_id,
|
||||
"turn": m.turn,
|
||||
"ttft_ms": round(m.ttft_ms, 2),
|
||||
"fc_ms": round(m.fc_ms, 2),
|
||||
"itl_ms": round(m.itl_ms, 2),
|
||||
"e2e_latency_ms": round(m.e2e_latency_ms, 2),
|
||||
"input_tokens": m.input_tokens,
|
||||
"output_tokens": m.output_tokens,
|
||||
"start_time_ms": round(m.start_time_ms, 2),
|
||||
}
|
||||
for m in metrics
|
||||
]
|
||||
@@ -0,0 +1,397 @@
|
||||
"""Data models for the multi-turn benchmark."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass, field
|
||||
from typing import List, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class TurnResult:
|
||||
"""Result of a single turn's HTTP request."""
|
||||
|
||||
ttft_ms: float # time to first token
|
||||
fc_ms: float # first-chunk latency (time to N-th content chunk)
|
||||
itl_ms: float # mean inter-token latency across output tokens
|
||||
e2e_latency_ms: float # total request latency
|
||||
input_tokens: int # reported by server (usage.prompt_tokens)
|
||||
output_tokens: int # reported by server (usage.completion_tokens)
|
||||
generated_text: str # generated text
|
||||
itl_ms_list: List[float] = field(default_factory=list) # per-token ITL values
|
||||
|
||||
|
||||
@dataclass
|
||||
class TurnMetric:
|
||||
"""Metrics for a single turn."""
|
||||
|
||||
session_id: str
|
||||
turn: int # 0-indexed
|
||||
ttft_ms: float
|
||||
fc_ms: float # first-chunk latency
|
||||
itl_ms: float # mean inter-token latency
|
||||
e2e_latency_ms: float
|
||||
input_tokens: int
|
||||
output_tokens: int
|
||||
start_time_ms: float # relative to benchmark start
|
||||
itl_ms_list: List[float] = field(default_factory=list) # per-token ITL values
|
||||
|
||||
|
||||
@dataclass
|
||||
class WorkloadSpec:
|
||||
"""Workload specification for multi-turn session benchmarks.
|
||||
|
||||
Supports simple mode: specify isl + hit_rate, derive user_tokens and sys_tokens.
|
||||
All parameters are scalar (fixed) values -- no distributions.
|
||||
"""
|
||||
|
||||
# Core parameters
|
||||
num_sessions: Optional[int] = None # total unique sessions (None = duration-based)
|
||||
num_turns: int = 1 # turns per session
|
||||
osl: int = 1 # output sequence length per turn
|
||||
think_time: float = 0.0 # seconds between turns within a session
|
||||
|
||||
# Traffic (use either concurrency or request_rate, not both)
|
||||
concurrency: Optional[int] = None # max concurrent in-flight requests
|
||||
request_rate: Optional[float] = None # requests per second (constant rate mode)
|
||||
ramp_interval: float = -1.0 # seconds between session launches (-1 = auto)
|
||||
|
||||
# Duration-based mode (used with request_rate)
|
||||
duration_s: float = 0.0 # seconds to run benchmark (0 = use num_sessions)
|
||||
|
||||
# Fraction of system prompt shared across all sessions
|
||||
# 1.0 = identical system prompt, 0.0 = all unique
|
||||
shared_system_prompt_ratio: float = 1.0
|
||||
|
||||
# Simple mode inputs (derive user_tokens, sys_tokens)
|
||||
isl: Optional[int] = None
|
||||
hit_rate: Optional[float] = None
|
||||
|
||||
# Resolved values (computed by resolve())
|
||||
_user_tokens: int = field(default=0, init=False, repr=False)
|
||||
_sys_tokens: int = field(default=0, init=False, repr=False)
|
||||
|
||||
def resolve(self) -> "WorkloadSpec":
|
||||
"""Resolve the spec: derive user_tokens and sys_tokens from inputs. Call after init."""
|
||||
if self.isl is None or self.hit_rate is None:
|
||||
raise ValueError("Simple mode requires both --isl and --hit-rate.")
|
||||
self._validate()
|
||||
self._derive_from_simple()
|
||||
return self
|
||||
|
||||
def _derive_from_simple(self) -> None:
|
||||
"""Derive user_tokens and sys_tokens from (ISL, hit_rate, num_turns, OSL, shared_system_prompt_ratio).
|
||||
|
||||
Two equations, two unknowns (u = user_tokens, s = sys_tokens):
|
||||
|
||||
(1) ISL = s + (n+1)/2 · u + (n-1)/2 · a [average input length]
|
||||
(2) (1-h)·ISL = (1-f)·s/n + u [average new-token fraction]
|
||||
|
||||
where n = num_turns, a = osl, f = shared_system_prompt_ratio, h = hit_rate.
|
||||
|
||||
Substituting s from (1) into (2) and solving for u:
|
||||
|
||||
u = [ (1-h)·ISL - (1-f)/n · (ISL - (n-1)·a/2) ]
|
||||
/ [ 1 - (1-f)·(n+1)/(2n) ]
|
||||
|
||||
Then s = ISL - (n+1)/2 · u - (n-1)/2 · a.
|
||||
|
||||
Special case: when n=1 and f=0, equations (1) and (2) collapse to
|
||||
s + u = ISL with h = s/(s+u), giving s = h·ISL and u = (1-h)·ISL.
|
||||
"""
|
||||
isl = self.isl
|
||||
h = self.hit_rate
|
||||
n = self.num_turns
|
||||
a = self.osl
|
||||
f = self.shared_system_prompt_ratio
|
||||
|
||||
denom = 1 - (1 - f) * (n + 1) / (2 * n)
|
||||
if abs(denom) < 1e-9:
|
||||
# n=1, f=0, h=0 (validated earlier): s=0, u=ISL.
|
||||
sys_tokens = 0.0
|
||||
user_tokens = float(isl)
|
||||
else:
|
||||
numer = (1 - h) * isl - (1 - f) / n * (isl - (n - 1) * a / 2)
|
||||
user_tokens = numer / denom
|
||||
sys_tokens = isl - (n + 1) / 2 * user_tokens - (n - 1) / 2 * a
|
||||
|
||||
if user_tokens < 0.5 or sys_tokens < -0.5:
|
||||
suggestions = self._feasibility_suggestions()
|
||||
which = "user_tokens" if user_tokens < 0.5 else "sys_tokens"
|
||||
val = user_tokens if user_tokens < 0.5 else sys_tokens
|
||||
raise ValueError(
|
||||
f"Derived {which} = {val:.1f} is infeasible with "
|
||||
f"(ISL={isl}, hit_rate={h}, num_turns={n}, "
|
||||
f"OSL={a}, shared_system_prompt_ratio={f}).\n"
|
||||
f"To fix, try one of:\n{suggestions}"
|
||||
)
|
||||
|
||||
self._user_tokens = max(1, int(round(user_tokens)))
|
||||
self._sys_tokens = max(0, int(round(sys_tokens)))
|
||||
|
||||
def _feasibility_suggestions(self) -> str:
|
||||
"""Compute feasible boundary values for each parameter and return suggestions.
|
||||
|
||||
For each workload parameter, search for a boundary value that makes
|
||||
the solver yield user_tokens >= 0.5 and sys_tokens >= -0.5 (the
|
||||
minimum values that round to physically meaningful token counts:
|
||||
at least 1 user token and non-negative system tokens).
|
||||
"""
|
||||
isl = self.isl
|
||||
hit_rate = self.hit_rate
|
||||
num_turns = self.num_turns
|
||||
osl = self.osl
|
||||
sharing = self.shared_system_prompt_ratio
|
||||
lines = []
|
||||
|
||||
def _try_solve(isl_, hit_rate_, num_turns_, osl_, sharing_):
|
||||
"""Solve for (user_tokens, sys_tokens) or return None if degenerate."""
|
||||
denom = 1 - (1 - sharing_) * (num_turns_ + 1) / (2 * num_turns_)
|
||||
if abs(denom) < 1e-9:
|
||||
if hit_rate_ > 1e-9:
|
||||
return None
|
||||
return (float(isl_), 0.0)
|
||||
numer = (1 - hit_rate_) * isl_ - (1 - sharing_) / num_turns_ * (
|
||||
isl_ - (num_turns_ - 1) * osl_ / 2
|
||||
)
|
||||
user_tokens = numer / denom
|
||||
sys_tokens = (
|
||||
isl_ - (num_turns_ + 1) / 2 * user_tokens - (num_turns_ - 1) / 2 * osl_
|
||||
)
|
||||
return (user_tokens, sys_tokens)
|
||||
|
||||
def _feasible(isl_, hit_rate_, num_turns_, osl_, sharing_):
|
||||
result = _try_solve(isl_, hit_rate_, num_turns_, osl_, sharing_)
|
||||
# user_tokens >= 0.5 rounds to at least 1 token per turn;
|
||||
# sys_tokens >= -0.5 rounds to at least 0 system prompt tokens.
|
||||
return result is not None and result[0] >= 0.5 and result[1] >= -0.5
|
||||
|
||||
# Min ISL (binary search)
|
||||
lo, hi = isl, isl * 20
|
||||
if _feasible(hi, hit_rate, num_turns, osl, sharing):
|
||||
while hi - lo > 1:
|
||||
mid = (lo + hi) // 2
|
||||
if _feasible(mid, hit_rate, num_turns, osl, sharing):
|
||||
hi = mid
|
||||
else:
|
||||
lo = mid
|
||||
lines.append(f" - ISL >= {hi} (with current params)")
|
||||
|
||||
# Max OSL
|
||||
lo, hi = 1, osl
|
||||
if _feasible(isl, hit_rate, num_turns, lo, sharing):
|
||||
while hi - lo > 1:
|
||||
mid = (lo + hi) // 2
|
||||
if _feasible(isl, hit_rate, num_turns, mid, sharing):
|
||||
lo = mid
|
||||
else:
|
||||
hi = mid
|
||||
lines.append(f" - OSL <= {lo} (with current ISL={isl})")
|
||||
|
||||
# Min hit_rate / max hit_rate (search in 0.01 steps)
|
||||
for h_try in range(0, 100):
|
||||
h_val = h_try / 100.0
|
||||
if _feasible(isl, h_val, num_turns, osl, sharing):
|
||||
if h_val != hit_rate:
|
||||
if h_val > hit_rate:
|
||||
lines.append(
|
||||
f" - hit_rate >= {h_val:.2f} (with current ISL/OSL)"
|
||||
)
|
||||
else:
|
||||
lines.append(
|
||||
f" - hit_rate <= {h_val:.2f} (with current ISL/OSL)"
|
||||
)
|
||||
break
|
||||
|
||||
# Max num_turns
|
||||
for n_try in range(num_turns, 0, -1):
|
||||
if _feasible(isl, hit_rate, n_try, osl, sharing):
|
||||
if n_try != num_turns:
|
||||
lines.append(f" - num_turns <= {n_try} (with current ISL/OSL)")
|
||||
break
|
||||
|
||||
# Min shared_system_prompt_ratio
|
||||
if sharing < 1.0:
|
||||
for f_try in range(int(sharing * 100), 101):
|
||||
f_val = f_try / 100.0
|
||||
if _feasible(isl, hit_rate, num_turns, osl, f_val):
|
||||
if f_val != sharing:
|
||||
lines.append(f" - shared_system_prompt_ratio >= {f_val:.2f}")
|
||||
break
|
||||
|
||||
return "\n".join(lines) if lines else " (no single-parameter fix found)"
|
||||
|
||||
def _validate(self) -> None:
|
||||
"""Validate resolved parameters."""
|
||||
if self.num_turns < 1:
|
||||
raise ValueError("num_turns must be >= 1.")
|
||||
if self.osl < 1:
|
||||
raise ValueError("osl must be >= 1.")
|
||||
if self.num_sessions is not None and self.num_sessions < 1:
|
||||
raise ValueError("num_sessions must be >= 1.")
|
||||
if self.num_sessions is None and self.duration_s <= 0:
|
||||
raise ValueError(
|
||||
"Must specify either --num-sessions or --duration (> 0) for rate-based mode."
|
||||
)
|
||||
if not (0 <= self.shared_system_prompt_ratio <= 1):
|
||||
raise ValueError("shared_system_prompt_ratio must be in [0, 1].")
|
||||
if self.think_time < 0:
|
||||
raise ValueError("think_time must be >= 0.")
|
||||
if (
|
||||
self.num_turns == 1
|
||||
and self.shared_system_prompt_ratio == 0
|
||||
and self.hit_rate is not None
|
||||
and self.hit_rate > 1e-9
|
||||
):
|
||||
raise ValueError(
|
||||
f"Cannot achieve hit_rate={self.hit_rate} with num_turns=1 and "
|
||||
f"shared_system_prompt_ratio=0. There is no caching source "
|
||||
f"(no multi-turn history, no shared prefix). "
|
||||
f"Set shared_system_prompt_ratio > 0 to enable cross-session "
|
||||
f"prefix caching, or use num_turns > 1 for multi-turn caching."
|
||||
)
|
||||
|
||||
if self.concurrency is None and self.request_rate is None:
|
||||
raise ValueError("Must specify either --concurrency or --request-rate.")
|
||||
if self.concurrency is not None and self.request_rate is not None:
|
||||
raise ValueError("Cannot specify both --concurrency and --request-rate.")
|
||||
if self.concurrency is not None and self.concurrency < 1:
|
||||
raise ValueError("concurrency must be >= 1.")
|
||||
if self.request_rate is not None and self.request_rate <= 0:
|
||||
raise ValueError("request_rate must be > 0.")
|
||||
|
||||
if self.ramp_interval < 0:
|
||||
if self.concurrency is not None:
|
||||
if self.think_time > 0:
|
||||
self.ramp_interval = self.think_time / self.concurrency
|
||||
else:
|
||||
self.ramp_interval = 0.0
|
||||
else:
|
||||
self.ramp_interval = 0.0
|
||||
|
||||
if (
|
||||
self.concurrency is not None
|
||||
and self.think_time > 0
|
||||
and self.num_sessions is not None
|
||||
and self.num_sessions < self.concurrency * 2
|
||||
):
|
||||
logger.warning(
|
||||
"num_sessions=%d may be too low to sustain concurrency=%d "
|
||||
"with think_time=%.1f. Consider increasing num_sessions.",
|
||||
self.num_sessions,
|
||||
self.concurrency,
|
||||
self.think_time,
|
||||
)
|
||||
|
||||
@property
|
||||
def user_tokens(self) -> int:
|
||||
return self._user_tokens
|
||||
|
||||
@property
|
||||
def sys_tokens(self) -> int:
|
||||
return self._sys_tokens
|
||||
|
||||
@property
|
||||
def shared_s(self) -> int:
|
||||
return int(round(self._sys_tokens * self.shared_system_prompt_ratio))
|
||||
|
||||
@property
|
||||
def unique_s(self) -> int:
|
||||
return self._sys_tokens - self.shared_s
|
||||
|
||||
def turn_input_tokens(self, k: int) -> int:
|
||||
"""Total input tokens at turn k (1-indexed)."""
|
||||
return self._sys_tokens + k * self._user_tokens + (k - 1) * self.osl
|
||||
|
||||
@property
|
||||
def effective_isl(self) -> float:
|
||||
n = self.num_turns
|
||||
return (
|
||||
self._sys_tokens + self._user_tokens * (n + 1) / 2 + self.osl * (n - 1) / 2
|
||||
)
|
||||
|
||||
@property
|
||||
def effective_h(self) -> float:
|
||||
f = self.shared_system_prompt_ratio
|
||||
n = self.num_turns
|
||||
avg_new = (1 - f) * self._sys_tokens / n + self._user_tokens
|
||||
isl = self.effective_isl
|
||||
return 1.0 - avg_new / isl if isl > 0 else 0.0
|
||||
|
||||
def summary(self) -> dict:
|
||||
per_turn = []
|
||||
for k in range(1, self.num_turns + 1):
|
||||
total = self.turn_input_tokens(k)
|
||||
if k == 1:
|
||||
cached = int(round(self._sys_tokens * self.shared_system_prompt_ratio))
|
||||
else:
|
||||
cached = (
|
||||
self._sys_tokens + (k - 1) * self._user_tokens + (k - 1) * self.osl
|
||||
)
|
||||
new = total - cached
|
||||
h_k = cached / total if total > 0 else 0.0
|
||||
per_turn.append(
|
||||
{
|
||||
"turn": k,
|
||||
"total": total,
|
||||
"cached": cached,
|
||||
"new": new,
|
||||
"hit_rate": round(h_k, 4),
|
||||
}
|
||||
)
|
||||
|
||||
return {
|
||||
"num_sessions": self.num_sessions,
|
||||
"duration_s": self.duration_s,
|
||||
"num_turns": self.num_turns,
|
||||
"osl": self.osl,
|
||||
"think_time": self.think_time,
|
||||
"concurrency": self.concurrency,
|
||||
"request_rate": self.request_rate,
|
||||
"shared_system_prompt_ratio": self.shared_system_prompt_ratio,
|
||||
"user_tokens_per_turn": self._user_tokens,
|
||||
"system_prompt_tokens": self._sys_tokens,
|
||||
"shared_system_prompt": self.shared_s,
|
||||
"unique_system_prompt": self.unique_s,
|
||||
"effective_isl": round(self.effective_isl, 1),
|
||||
"effective_hit_rate": round(self.effective_h, 4),
|
||||
"per_turn": per_turn,
|
||||
}
|
||||
|
||||
def print_summary(self) -> None:
|
||||
s = self.summary()
|
||||
print("=" * 70)
|
||||
print("Workload Spec (resolved)")
|
||||
print("=" * 70)
|
||||
if s["num_sessions"] is not None:
|
||||
print(f" Sessions (N_s): {s['num_sessions']}")
|
||||
else:
|
||||
print(" Sessions (N_s): unlimited (duration-based)")
|
||||
if s["duration_s"] > 0:
|
||||
print(f" Duration: {s['duration_s']}s")
|
||||
print(f" Turns per session (N_t): {s['num_turns']}")
|
||||
print(f" User tokens/turn (u): {s['user_tokens_per_turn']}")
|
||||
print(
|
||||
f" System prompt (s): {s['system_prompt_tokens']} "
|
||||
f"(shared={s['shared_system_prompt']}, unique={s['unique_system_prompt']})"
|
||||
)
|
||||
print(f" Output tokens (o): {s['osl']}")
|
||||
print(f" Think time: {s['think_time']}s")
|
||||
if self.concurrency is not None:
|
||||
print(f" Concurrency (C): {self.concurrency}")
|
||||
print(f" Ramp interval: {self.ramp_interval:.3f}s")
|
||||
if self.request_rate is not None:
|
||||
print(f" Request rate (QPS): {self.request_rate}")
|
||||
print(f" Shared sys prompt ratio: {s['shared_system_prompt_ratio']}")
|
||||
print(f" Effective avg ISL: {s['effective_isl']}")
|
||||
print(f" Effective avg hit rate: {s['effective_hit_rate']:.1%}")
|
||||
print("-" * 70)
|
||||
print(f" {'Turn':<6} {'Total':<8} {'Cached':<8} {'New':<8} {'Hit Rate':<10}")
|
||||
for t in s["per_turn"]:
|
||||
print(
|
||||
f" {t['turn']:<6} {t['total']:<8} {t['cached']:<8} "
|
||||
f"{t['new']:<8} {t['hit_rate']:.1%}"
|
||||
)
|
||||
print("=" * 70)
|
||||
@@ -0,0 +1,154 @@
|
||||
"""Reporting and result persistence for the benchmark."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from statistics import mean
|
||||
from typing import Optional
|
||||
|
||||
from ray.llm._internal.serve.benchmark.metrics import (
|
||||
percentile,
|
||||
serialize_raw_metrics,
|
||||
summarize_metrics,
|
||||
)
|
||||
from ray.llm._internal.serve.benchmark.models import TurnMetric, WorkloadSpec
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def report_results(
|
||||
metrics: list[TurnMetric],
|
||||
spec: WorkloadSpec,
|
||||
bench_elapsed_s: float,
|
||||
first_chunk_threshold: int = 16,
|
||||
save_path: Optional[str] = None,
|
||||
warmup_s: float = 0.0,
|
||||
discarded_warmup_requests: int = 0,
|
||||
) -> None:
|
||||
"""Print and optionally save benchmark results."""
|
||||
if not metrics:
|
||||
print("No metrics collected.")
|
||||
return
|
||||
|
||||
all_ttft = [m.ttft_ms for m in metrics]
|
||||
all_fc = [m.fc_ms for m in metrics]
|
||||
all_itl = [v for m in metrics for v in m.itl_ms_list]
|
||||
all_latency = [m.e2e_latency_ms for m in metrics]
|
||||
all_input = [m.input_tokens for m in metrics]
|
||||
all_output = [m.output_tokens for m in metrics]
|
||||
|
||||
total_output_tokens = sum(all_output)
|
||||
throughput = total_output_tokens / bench_elapsed_s if bench_elapsed_s > 0 else 0
|
||||
|
||||
print()
|
||||
print("=" * 70)
|
||||
print("BENCHMARK RESULTS")
|
||||
print("=" * 70)
|
||||
print(f" Total requests: {len(metrics)}")
|
||||
print(f" Unique sessions: {len({m.session_id for m in metrics})}")
|
||||
print(f" Duration: {bench_elapsed_s:.1f}s")
|
||||
if warmup_s > 0:
|
||||
print(f" Warm-up excluded: {warmup_s:.1f}s")
|
||||
if discarded_warmup_requests > 0:
|
||||
print(f" Warm-up requests: {discarded_warmup_requests} (discarded)")
|
||||
print(f" Throughput: {throughput:.1f} output tok/s")
|
||||
print(f" Request rate: {len(metrics) / bench_elapsed_s:.1f} req/s")
|
||||
print(
|
||||
f" Avg input tokens: {mean(all_input):.0f} "
|
||||
f"(target ISL: {spec.effective_isl:.0f})"
|
||||
)
|
||||
print(f" Avg output tokens: {mean(all_output):.0f} (target OSL: {spec.osl})")
|
||||
print()
|
||||
|
||||
fc_label = f"FC({first_chunk_threshold})"
|
||||
print(" Latency Statistics:")
|
||||
for name, values in [
|
||||
("TTFT", all_ttft),
|
||||
(fc_label, all_fc),
|
||||
("ITL", all_itl),
|
||||
("Latency", all_latency),
|
||||
]:
|
||||
if not values:
|
||||
continue
|
||||
print(
|
||||
f" {name:>8}: avg={mean(values):>8.1f}ms "
|
||||
f"P50={percentile(values, 50):>8.1f}ms "
|
||||
f"P90={percentile(values, 90):>8.1f}ms "
|
||||
f"P99={percentile(values, 99):>8.1f}ms"
|
||||
)
|
||||
print()
|
||||
|
||||
print(" Per-Turn Breakdown:")
|
||||
print(
|
||||
f" {'Turn':<6} {'Count':<7} {'Avg ISL':<9} {'Avg TTFT':<10} "
|
||||
f"{'Avg FC':<10} {'Avg ITL':<10} {'Avg Lat':<10}"
|
||||
)
|
||||
for t in range(spec.num_turns):
|
||||
turn_metrics = [m for m in metrics if m.turn == t]
|
||||
if not turn_metrics:
|
||||
continue
|
||||
t_ttft = mean([m.ttft_ms for m in turn_metrics])
|
||||
t_fc = mean([m.fc_ms for m in turn_metrics])
|
||||
t_itl_all = [v for m in turn_metrics for v in m.itl_ms_list]
|
||||
t_itl = mean(t_itl_all) if t_itl_all else 0.0
|
||||
t_lat = mean([m.e2e_latency_ms for m in turn_metrics])
|
||||
t_isl = mean([m.input_tokens for m in turn_metrics])
|
||||
print(
|
||||
f" {t + 1:<6} {len(turn_metrics):<7} {t_isl:<9.0f} "
|
||||
f"{t_ttft:<10.1f} {t_fc:<10.1f} {t_itl:<10.1f} {t_lat:<10.1f}"
|
||||
)
|
||||
print("=" * 70)
|
||||
|
||||
if save_path:
|
||||
stats = summarize_metrics(metrics, bench_elapsed_s)
|
||||
result = {
|
||||
"config": {
|
||||
"concurrency": spec.concurrency,
|
||||
"request_rate": spec.request_rate,
|
||||
},
|
||||
"spec": spec.summary(),
|
||||
"first_chunk_threshold": first_chunk_threshold,
|
||||
"benchmark": {
|
||||
"total_requests": len(metrics),
|
||||
"duration_s": round(bench_elapsed_s, 2),
|
||||
"warmup_s": round(warmup_s, 2),
|
||||
"discarded_warmup_requests": discarded_warmup_requests,
|
||||
},
|
||||
"stats": {
|
||||
("measured_request_rate" if k == "request_rate" else k): v
|
||||
for k, v in stats.items()
|
||||
if k not in ("requests", "elapsed_s")
|
||||
},
|
||||
"per_turn": [],
|
||||
"raw_metrics": serialize_raw_metrics(metrics),
|
||||
}
|
||||
|
||||
for t in range(spec.num_turns):
|
||||
turn_metrics = [m for m in metrics if m.turn == t]
|
||||
if not turn_metrics:
|
||||
continue
|
||||
t_ttft = [m.ttft_ms for m in turn_metrics]
|
||||
t_fc = [m.fc_ms for m in turn_metrics]
|
||||
t_itl = [v for m in turn_metrics for v in m.itl_ms_list]
|
||||
t_isl = [m.input_tokens for m in turn_metrics]
|
||||
result["per_turn"].append(
|
||||
{
|
||||
"turn": t + 1,
|
||||
"count": len(turn_metrics),
|
||||
"avg_isl": round(mean(t_isl), 1),
|
||||
"avg_ttft_ms": round(mean(t_ttft), 2),
|
||||
"avg_fc_ms": round(mean(t_fc), 2),
|
||||
"avg_itl_ms": round(mean(t_itl), 2) if t_itl else 0,
|
||||
"p50_fc_ms": round(percentile(t_fc, 50), 2),
|
||||
"p99_ttft_ms": round(percentile(t_ttft, 99), 2),
|
||||
"p99_fc_ms": round(percentile(t_fc, 99), 2),
|
||||
"p99_itl_ms": (round(percentile(t_itl, 99), 2) if t_itl else 0),
|
||||
}
|
||||
)
|
||||
|
||||
Path(save_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(save_path, "w") as f:
|
||||
json.dump(result, f, indent=2)
|
||||
logger.info("Results saved to %s", save_path)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,175 @@
|
||||
"""Text generation and conversation management for the benchmark."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import numpy as np
|
||||
|
||||
from ray.llm._internal.serve.benchmark.models import WorkloadSpec
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from transformers import PreTrainedTokenizerBase
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Conversation:
|
||||
"""A single multi-turn conversation with a unique session ID."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
session_id: str,
|
||||
system_prompt: str,
|
||||
user_messages: list[str],
|
||||
num_turns: int,
|
||||
):
|
||||
self.session_id = session_id
|
||||
self.system_prompt = system_prompt
|
||||
self.user_messages = user_messages
|
||||
self.num_turns = num_turns
|
||||
self._assistant_responses: list[str] = []
|
||||
|
||||
def get_turn_messages(self, turn_idx: int) -> list[dict[str, str]]:
|
||||
"""Build the messages list for turn `turn_idx` (0-indexed)."""
|
||||
messages: list[dict[str, str]] = []
|
||||
if self.system_prompt:
|
||||
messages.append({"role": "system", "content": self.system_prompt})
|
||||
|
||||
for i in range(turn_idx + 1):
|
||||
messages.append({"role": "user", "content": self.user_messages[i]})
|
||||
if i < turn_idx:
|
||||
if i < len(self._assistant_responses):
|
||||
messages.append(
|
||||
{"role": "assistant", "content": self._assistant_responses[i]}
|
||||
)
|
||||
else:
|
||||
messages.append({"role": "assistant", "content": "(placeholder)"})
|
||||
return messages
|
||||
|
||||
def inject_assistant_response(self, turn_idx: int, content: str) -> None:
|
||||
"""Record the server's response for turn `turn_idx`."""
|
||||
if turn_idx == len(self._assistant_responses):
|
||||
self._assistant_responses.append(content)
|
||||
elif turn_idx < len(self._assistant_responses):
|
||||
self._assistant_responses[turn_idx] = content
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Cannot inject response for turn {turn_idx}: "
|
||||
f"only {len(self._assistant_responses)} responses recorded."
|
||||
)
|
||||
|
||||
|
||||
class TextGenerator:
|
||||
"""Generates random text with exact token counts using a tokenizer."""
|
||||
|
||||
def __init__(self, tokenizer: "PreTrainedTokenizerBase"):
|
||||
self._tokenizer = tokenizer
|
||||
self._vocab_size = tokenizer.vocab_size
|
||||
logger.info(
|
||||
"TextGenerator using tokenizer (vocab_size=%d) for exact token counts.",
|
||||
self._vocab_size,
|
||||
)
|
||||
|
||||
def generate(self, num_tokens: int) -> str:
|
||||
if num_tokens <= 0:
|
||||
return ""
|
||||
return self._generate_exact(num_tokens)
|
||||
|
||||
def generate_token_ids(self, num_tokens: int) -> list[int]:
|
||||
if num_tokens <= 0:
|
||||
return []
|
||||
return np.random.randint(0, self._vocab_size, size=num_tokens).tolist()
|
||||
|
||||
def _generate_exact(self, target_tokens: int) -> str:
|
||||
tokenizer = self._tokenizer
|
||||
token_ids = np.random.randint(
|
||||
0, self._vocab_size, size=target_tokens + 20
|
||||
).tolist()
|
||||
|
||||
text = tokenizer.decode(token_ids, skip_special_tokens=True)
|
||||
actual_ids = tokenizer.encode(text, add_special_tokens=False)
|
||||
actual_len = len(actual_ids)
|
||||
|
||||
if actual_len == target_tokens:
|
||||
return text
|
||||
|
||||
if actual_len > target_tokens:
|
||||
trimmed_ids = actual_ids[:target_tokens]
|
||||
text = tokenizer.decode(trimmed_ids, skip_special_tokens=True)
|
||||
final_len = len(tokenizer.encode(text, add_special_tokens=False))
|
||||
if final_len != target_tokens:
|
||||
text = self._binary_search_trim(actual_ids, target_tokens)
|
||||
return text
|
||||
|
||||
deficit = target_tokens - actual_len
|
||||
extra_ids = np.random.randint(0, self._vocab_size, size=deficit + 20).tolist()
|
||||
extra_text = tokenizer.decode(extra_ids, skip_special_tokens=True)
|
||||
combined = text + " " + extra_text
|
||||
combined_ids = tokenizer.encode(combined, add_special_tokens=False)
|
||||
|
||||
if len(combined_ids) >= target_tokens:
|
||||
trimmed = combined_ids[:target_tokens]
|
||||
text = tokenizer.decode(trimmed, skip_special_tokens=True)
|
||||
final_len = len(tokenizer.encode(text, add_special_tokens=False))
|
||||
if final_len != target_tokens:
|
||||
text = self._binary_search_trim(combined_ids, target_tokens)
|
||||
return text
|
||||
|
||||
while len(tokenizer.encode(combined, add_special_tokens=False)) < target_tokens:
|
||||
combined += " hello"
|
||||
combined_ids = tokenizer.encode(combined, add_special_tokens=False)
|
||||
return self._binary_search_trim(combined_ids, target_tokens)
|
||||
|
||||
def _binary_search_trim(self, token_ids: list[int], target: int) -> str:
|
||||
tokenizer = self._tokenizer
|
||||
lo, hi = target, len(token_ids)
|
||||
best_text = tokenizer.decode(token_ids[:target], skip_special_tokens=True)
|
||||
|
||||
while lo <= hi:
|
||||
mid = (lo + hi) // 2
|
||||
text = tokenizer.decode(token_ids[:mid], skip_special_tokens=True)
|
||||
actual = len(tokenizer.encode(text, add_special_tokens=False))
|
||||
if actual == target:
|
||||
return text
|
||||
elif actual < target:
|
||||
lo = mid + 1
|
||||
else:
|
||||
hi = mid - 1
|
||||
best_text = text
|
||||
|
||||
for n in range(target, len(token_ids) + 1):
|
||||
text = tokenizer.decode(token_ids[:n], skip_special_tokens=True)
|
||||
if len(tokenizer.encode(text, add_special_tokens=False)) == target:
|
||||
return text
|
||||
return best_text
|
||||
|
||||
|
||||
def conversation_factory(
|
||||
session_idx: int,
|
||||
spec: WorkloadSpec,
|
||||
shared_system_text: str,
|
||||
text_gen: Optional[TextGenerator],
|
||||
) -> Conversation:
|
||||
"""Create a single conversation on-demand (lazy generation)."""
|
||||
session_id = f"session-{session_idx:06d}"
|
||||
|
||||
if spec.unique_s > 0 and text_gen is not None:
|
||||
unique_text = text_gen.generate(spec.unique_s)
|
||||
system_prompt = shared_system_text + " " + unique_text
|
||||
else:
|
||||
system_prompt = shared_system_text
|
||||
|
||||
user_messages = (
|
||||
[text_gen.generate(spec.user_tokens) for _ in range(spec.num_turns)]
|
||||
if text_gen is not None
|
||||
else ["" for _ in range(spec.num_turns)]
|
||||
)
|
||||
|
||||
return Conversation(
|
||||
session_id=session_id,
|
||||
system_prompt=system_prompt,
|
||||
user_messages=user_messages,
|
||||
num_turns=spec.num_turns,
|
||||
)
|
||||
@@ -0,0 +1,75 @@
|
||||
"""Single-turn execution primitive for the benchmark.
|
||||
|
||||
This module provides the pure core of turn execution: send an HTTP request,
|
||||
build a TurnMetric, and inject the response. It has NO side effects — callers
|
||||
are responsible for inflight tracking, metric recording, and queue management.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
import aiohttp
|
||||
|
||||
from ray.llm._internal.serve.benchmark.http_client import send_chat_completion
|
||||
from ray.llm._internal.serve.benchmark.models import TurnMetric, TurnResult
|
||||
from ray.llm._internal.serve.benchmark.text_gen import Conversation
|
||||
|
||||
|
||||
@dataclass
|
||||
class TurnOutcome:
|
||||
"""Result of executing a single benchmark turn."""
|
||||
|
||||
metric: TurnMetric
|
||||
result: TurnResult
|
||||
|
||||
|
||||
async def execute_single_turn(
|
||||
http_session: aiohttp.ClientSession,
|
||||
conv: Conversation,
|
||||
turn_idx: int,
|
||||
base_url: str,
|
||||
model: str,
|
||||
max_tokens: int,
|
||||
bench_start_ns: int,
|
||||
first_chunk_threshold: int = 16,
|
||||
api_key: Optional[str] = None,
|
||||
) -> TurnOutcome:
|
||||
"""Execute a single benchmark turn: HTTP call, build metric, inject response.
|
||||
|
||||
This is the pure core shared by all three benchmark engines (concurrency,
|
||||
rate-based, interactive). The caller handles inflight tracking, warmup
|
||||
filtering, measurement windows, and queue re-enqueue.
|
||||
"""
|
||||
messages = conv.get_turn_messages(turn_idx)
|
||||
req_start_ns = time.perf_counter_ns()
|
||||
|
||||
result = await send_chat_completion(
|
||||
session=http_session,
|
||||
base_url=base_url,
|
||||
model=model,
|
||||
messages=messages,
|
||||
session_id=conv.session_id,
|
||||
max_tokens=max_tokens,
|
||||
first_chunk_threshold=first_chunk_threshold,
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
metric = TurnMetric(
|
||||
session_id=conv.session_id,
|
||||
turn=turn_idx,
|
||||
ttft_ms=result.ttft_ms,
|
||||
fc_ms=result.fc_ms,
|
||||
itl_ms=result.itl_ms,
|
||||
e2e_latency_ms=result.e2e_latency_ms,
|
||||
input_tokens=result.input_tokens,
|
||||
output_tokens=result.output_tokens,
|
||||
start_time_ms=(req_start_ns - bench_start_ns) / 1e6,
|
||||
itl_ms_list=result.itl_ms_list,
|
||||
)
|
||||
|
||||
conv.inject_assistant_response(turn_idx, result.generated_text)
|
||||
|
||||
return TurnOutcome(metric=metric, result=result)
|
||||
@@ -0,0 +1,111 @@
|
||||
import os
|
||||
|
||||
ALLOW_NEW_PLACEMENT_GROUPS_IN_DEPLOYMENT = int(
|
||||
os.getenv("RAYLLM_ALLOW_NEW_PLACEMENT_GROUPS_IN_DEPLOYMENT", "1")
|
||||
)
|
||||
|
||||
|
||||
# Timeout before download in multiplex deployment fails. <=0 means no timeout.
|
||||
DEFAULT_MULTIPLEX_DOWNLOAD_TIMEOUT_S = float(
|
||||
os.getenv("DEFAULT_MULTIPLEX_DOWNLOAD_TIMEOUT_S", "30")
|
||||
)
|
||||
if DEFAULT_MULTIPLEX_DOWNLOAD_TIMEOUT_S <= 0:
|
||||
DEFAULT_MULTIPLEX_DOWNLOAD_TIMEOUT_S = None
|
||||
|
||||
|
||||
# Number of retries for downloading a model in multiplex deployment.
|
||||
DEFAULT_MULTIPLEX_DOWNLOAD_TRIES = int(
|
||||
os.getenv("DEFAULT_MULTIPLEX_DOWNLOAD_RETRIES", "3")
|
||||
)
|
||||
|
||||
|
||||
# If true, a default runtime_env will be injected to import rayllm on worker startup.
|
||||
# This is a startup time optimization to avoid the latency penalty of sequentially
|
||||
# importing rayllm in multiple layers of worker processes.
|
||||
ENABLE_WORKER_PROCESS_SETUP_HOOK = (
|
||||
os.environ.get("RAYLLM_ENABLE_WORKER_PROCESS_SETUP_HOOK", "1") == "1"
|
||||
)
|
||||
|
||||
|
||||
CLOUD_OBJECT_MISSING_EXPIRE_S = 30
|
||||
CLOUD_OBJECT_EXISTS_EXPIRE_S = 60 * 60
|
||||
|
||||
# Sentinel object used to indicate that a LoRA adapter config file is missing.
|
||||
LORA_ADAPTER_CONFIG_NAME = "adapter_config.json"
|
||||
|
||||
DEFAULT_HEALTH_CHECK_PERIOD_S = int(
|
||||
os.getenv("RAY_SERVE_LLM_DEFAULT_HEALTH_CHECK_PERIOD_S", "10")
|
||||
)
|
||||
DEFAULT_HEALTH_CHECK_TIMEOUT_S = int(
|
||||
os.getenv("RAY_SERVE_LLM_DEFAULT_HEALTH_CHECK_TIMEOUT_S", "10")
|
||||
)
|
||||
DEFAULT_MAX_ONGOING_REQUESTS = int(
|
||||
os.getenv("RAY_SERVE_LLM_DEFAULT_MAX_ONGOING_REQUESTS", str(int(1e9)))
|
||||
)
|
||||
DEFAULT_MAX_REPLICAS = int(os.getenv("RAY_SERVE_LLM_DEFAULT_MAX_REPLICAS", "10"))
|
||||
DEFAULT_MAX_TARGET_ONGOING_REQUESTS = int(
|
||||
os.getenv("RAY_SERVE_LLM_DEFAULT_MAX_TARGET_ONGOING_REQUESTS", str(int(1e9)))
|
||||
)
|
||||
|
||||
|
||||
ENGINE_START_TIMEOUT_S = int(os.getenv("RAYLLM_ENGINE_START_TIMEOUT_S", str(60 * 60)))
|
||||
|
||||
MIN_NUM_TOPLOGPROBS_ALLOWED = 0
|
||||
MAX_NUM_TOPLOGPROBS_ALLOWED = 5
|
||||
MODEL_RESPONSE_BATCH_TIMEOUT_MS = float(
|
||||
os.getenv("RAYLLM_MODEL_RESPONSE_BATCH_TIMEOUT_MS", "50")
|
||||
)
|
||||
RAYLLM_ENABLE_REQUEST_PROMPT_LOGS = (
|
||||
os.environ.get("RAYLLM_ENABLE_REQUEST_PROMPT_LOGS", "1") == "1"
|
||||
)
|
||||
RAYLLM_GUIDED_DECODING_BACKEND = os.environ.get(
|
||||
"RAYLLM_GUIDED_DECODING_BACKEND", "xgrammar"
|
||||
)
|
||||
|
||||
RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING = (
|
||||
os.environ.get("RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING", "0") == "1"
|
||||
)
|
||||
|
||||
MAX_NUM_STOPPING_SEQUENCES = int(os.getenv("RAYLLM_MAX_NUM_STOPPING_SEQUENCES", "8"))
|
||||
ENV_VARS_TO_PROPAGATE = {
|
||||
"HUGGING_FACE_HUB_TOKEN",
|
||||
"HF_TOKEN",
|
||||
}
|
||||
# timeout in 10 minutes. Streaming can take longer than 3 min
|
||||
DEFAULT_LLM_ROUTER_HTTP_TIMEOUT = float(
|
||||
os.environ.get("RAY_SERVE_LLM_ROUTER_HTTP_TIMEOUT", 600)
|
||||
)
|
||||
|
||||
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.
|
||||
DEFAULT_ROUTER_TO_MODEL_REPLICA_RATIO = float(
|
||||
os.getenv("RAY_SERVE_LLM_ROUTER_TO_MODEL_REPLICA_RATIO", "2")
|
||||
)
|
||||
|
||||
DEFAULT_LLM_ROUTER_MIN_REPLICAS = int(
|
||||
os.environ.get("RAY_SERVE_LLM_ROUTER_MIN_REPLICAS", 2)
|
||||
)
|
||||
DEFAULT_LLM_ROUTER_INITIAL_REPLICAS = int(
|
||||
os.environ.get("RAY_SERVE_LLM_ROUTER_INITIAL_REPLICAS", 2)
|
||||
)
|
||||
DEFAULT_LLM_ROUTER_MAX_REPLICAS = int(
|
||||
os.environ.get("RAY_SERVE_LLM_ROUTER_MAX_REPLICAS", 1000)
|
||||
)
|
||||
DEFAULT_LLM_ROUTER_TARGET_ONGOING_REQUESTS = int(
|
||||
os.environ.get(
|
||||
"RAY_SERVE_LLM_ROUTER_TARGET_ONGOING_REQUESTS",
|
||||
DEFAULT_MAX_TARGET_ONGOING_REQUESTS,
|
||||
)
|
||||
)
|
||||
# Minimum interval (seconds) between full tracebacks for fatal engine errors
|
||||
DEFAULT_FATAL_ERROR_COOLDOWN_S = float(
|
||||
os.getenv("RAY_SERVE_LLM_ERROR_LOG_COOLDOWN_S", "10")
|
||||
)
|
||||
|
||||
|
||||
# HOME DIR
|
||||
RAYLLM_HOME_DIR = os.environ.get("RAYLLM_HOME_DIR", os.path.expanduser("~/.ray/llm"))
|
||||
@@ -0,0 +1,305 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, List, Literal, Optional, Union
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from typing_extensions import Annotated
|
||||
|
||||
import ray.util.accelerators.accelerators as accelerators
|
||||
from ray._private.accelerators.tpu import get_chips_per_host
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
from ray.util.placement_group import PlacementGroup, placement_group
|
||||
from ray.util.tpu import (
|
||||
get_tpu_version_from_type,
|
||||
slice_placement_group,
|
||||
)
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
AcceleratorType = Enum("AcceleratorType", vars(accelerators))
|
||||
|
||||
# Set of TPU string values from Ray's known accelerators.
|
||||
TPU_ACCELERATOR_VALUES = {
|
||||
member.value
|
||||
for name, member in AcceleratorType.__members__.items()
|
||||
if name.startswith("GOOGLE_TPU")
|
||||
}
|
||||
|
||||
|
||||
def format_ray_accelerator_resource(accelerator_type_str: str) -> str:
|
||||
"""Formats the accelerator type into a Ray custom resource string."""
|
||||
return f"accelerator_type:{accelerator_type_str}"
|
||||
|
||||
|
||||
def infer_hardware_kind_from_bundles(
|
||||
placement_group_config: Optional[Dict[str, Any]]
|
||||
) -> Optional[str]:
|
||||
"""Inspects placement group bundles and returns the inferred hardware kind."""
|
||||
if not placement_group_config:
|
||||
return None
|
||||
|
||||
bundle_per_worker = placement_group_config.get("bundle_per_worker") or {}
|
||||
bundles = placement_group_config.get("bundles") or []
|
||||
all_bundles = [bundle_per_worker] + bundles
|
||||
|
||||
if any(b.get("TPU", 0) > 0 for b in all_bundles):
|
||||
return "tpu"
|
||||
if any(b.get("GPU", 0) > 0 for b in all_bundles):
|
||||
return "gpu"
|
||||
|
||||
# If a config was provided but lacks GPUs or TPUs, it is a CPU deployment
|
||||
return "cpu"
|
||||
|
||||
|
||||
class AcceleratorConfig(BaseModel):
|
||||
kind: str
|
||||
|
||||
|
||||
class CPUConfig(AcceleratorConfig):
|
||||
kind: Literal["cpu"] = "cpu"
|
||||
|
||||
|
||||
class GPUConfig(AcceleratorConfig):
|
||||
kind: Literal["gpu"] = "gpu"
|
||||
|
||||
|
||||
class TPUConfig(AcceleratorConfig):
|
||||
kind: Literal["tpu"] = "tpu"
|
||||
topology: Optional[str] = None
|
||||
|
||||
|
||||
AnyAcceleratorConfig = Annotated[
|
||||
Union[CPUConfig, GPUConfig, TPUConfig],
|
||||
Field(discriminator="kind"),
|
||||
]
|
||||
|
||||
|
||||
class AcceleratorBackend(ABC):
|
||||
@abstractmethod
|
||||
def default_bundles(
|
||||
self,
|
||||
*,
|
||||
num_devices: int,
|
||||
accelerator_type_str: Optional[str] = None,
|
||||
) -> List[Dict[str, float]]:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def create_placement_group(
|
||||
self,
|
||||
*,
|
||||
bundles: List[Dict[str, float]],
|
||||
strategy: str,
|
||||
name: str,
|
||||
accelerator_type_str: Optional[str] = None,
|
||||
) -> PlacementGroup:
|
||||
pass
|
||||
|
||||
@property
|
||||
def requires_deferred_placement_group(self) -> bool:
|
||||
"""
|
||||
If True, Ray Serve will not provision a placement group for the deployment.
|
||||
Instead, creation is deferred to the replica at runtime.
|
||||
Defaults to False.
|
||||
"""
|
||||
return False
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def requires_remote_initialization(self) -> bool:
|
||||
"""Boolean indicating whether this backend needs a remote Ray task to query hardware during init."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_remote_options(self, accelerator_type_str: str = None) -> Dict[str, Any]:
|
||||
"""Returns the hardware-specific kwargs for ray.remote().options()."""
|
||||
pass
|
||||
|
||||
def shutdown(self) -> None:
|
||||
"""Release any resources owned by this backend. Idempotent."""
|
||||
return
|
||||
|
||||
|
||||
class CPUAccelerator(AcceleratorBackend):
|
||||
# stateless — no __init__
|
||||
def default_bundles(
|
||||
self, *, num_devices: int, accelerator_type_str: Optional[str] = None
|
||||
):
|
||||
return [{"CPU": 1} for _ in range(num_devices)]
|
||||
|
||||
def create_placement_group(
|
||||
self,
|
||||
*,
|
||||
bundles: List[Dict[str, float]],
|
||||
strategy: str,
|
||||
name: str,
|
||||
accelerator_type_str: Optional[str] = None,
|
||||
):
|
||||
return placement_group(bundles=bundles, strategy=strategy, name=name)
|
||||
|
||||
@property
|
||||
def requires_remote_initialization(self) -> bool:
|
||||
return False
|
||||
|
||||
def get_remote_options(self, accelerator_type_str: str = None):
|
||||
return {}
|
||||
|
||||
|
||||
class GPUAccelerator(AcceleratorBackend):
|
||||
# stateless — no __init__
|
||||
def default_bundles(
|
||||
self, *, num_devices: int, accelerator_type_str: Optional[str] = None
|
||||
):
|
||||
bundle = {"GPU": 1}
|
||||
if accelerator_type_str:
|
||||
bundle[format_ray_accelerator_resource(accelerator_type_str)] = 0.001
|
||||
return [bundle.copy() for _ in range(num_devices)]
|
||||
|
||||
def create_placement_group(
|
||||
self,
|
||||
*,
|
||||
bundles: List[Dict[str, float]],
|
||||
strategy: str,
|
||||
name: str,
|
||||
accelerator_type_str: Optional[str] = None,
|
||||
):
|
||||
return placement_group(bundles=bundles, strategy=strategy, name=name)
|
||||
|
||||
@property
|
||||
def requires_remote_initialization(self) -> bool:
|
||||
return True
|
||||
|
||||
def get_remote_options(self, accelerator_type_str: str = None):
|
||||
options = {"num_gpus": 0.001}
|
||||
if accelerator_type_str:
|
||||
options["accelerator_type"] = accelerator_type_str
|
||||
return options
|
||||
|
||||
|
||||
class TPUAccelerator(AcceleratorBackend):
|
||||
def __init__(self, config: TPUConfig):
|
||||
self._config = config
|
||||
self._slice_pg_wrapper = None
|
||||
|
||||
def default_bundles(
|
||||
self, *, num_devices: int, accelerator_type_str: Optional[str] = None
|
||||
):
|
||||
if not self._config.topology:
|
||||
# Fallback to per-chip bundles if no topology is specified
|
||||
bundle = {"TPU": 1}
|
||||
if accelerator_type_str:
|
||||
bundle[format_ray_accelerator_resource(accelerator_type_str)] = 0.001
|
||||
return [bundle.copy() for _ in range(num_devices)]
|
||||
|
||||
# Topology is specified, compute per-host bundles
|
||||
if not accelerator_type_str:
|
||||
raise ValueError(
|
||||
"`accelerator_type` must be specified when `topology` is present "
|
||||
"in order to compute TPU resource requirements."
|
||||
)
|
||||
version = get_tpu_version_from_type(accelerator_type_str)
|
||||
chips_per_host = get_chips_per_host(self._config.topology, version)
|
||||
|
||||
if num_devices > chips_per_host and num_devices % chips_per_host != 0:
|
||||
raise ValueError(
|
||||
f"num_devices ({num_devices}) must be a multiple of "
|
||||
f"chips_per_host ({chips_per_host}) for TPU topologies."
|
||||
)
|
||||
|
||||
num_hosts = max(1, num_devices // chips_per_host)
|
||||
|
||||
tpu_resources = min(num_devices, chips_per_host)
|
||||
bundle = {"TPU": tpu_resources}
|
||||
bundle[format_ray_accelerator_resource(accelerator_type_str)] = 0.001
|
||||
|
||||
return [bundle.copy() for _ in range(num_hosts)]
|
||||
|
||||
def create_placement_group(
|
||||
self,
|
||||
*,
|
||||
bundles: List[Dict[str, float]],
|
||||
strategy: str,
|
||||
name: str,
|
||||
accelerator_type_str: Optional[str] = None,
|
||||
) -> PlacementGroup:
|
||||
|
||||
if not self._config.topology:
|
||||
return placement_group(bundles=bundles, strategy=strategy, name=name)
|
||||
|
||||
if not accelerator_type_str:
|
||||
raise ValueError(
|
||||
"accelerator_type must be provided for TPU slice provisioning."
|
||||
)
|
||||
|
||||
version = get_tpu_version_from_type(accelerator_type_str)
|
||||
|
||||
if bundles:
|
||||
# Filter for bundles that actually specify TPU resources
|
||||
tpu_bundles = [b for b in bundles if b.get("TPU", 0) > 0]
|
||||
|
||||
if not tpu_bundles:
|
||||
worker_bundle = {"TPU": 1}
|
||||
else:
|
||||
worker_bundle = tpu_bundles[0]
|
||||
|
||||
# Ensure all TPU bundles are homogeneous
|
||||
if any(b != worker_bundle for b in tpu_bundles):
|
||||
raise ValueError(
|
||||
"Heterogeneous TPU bundles are not supported when `topology` is set. "
|
||||
"A multi-host TPU slice requires homogeneous resource bundles across all workers. "
|
||||
"Please use `bundle_per_worker` in `placement_group_config` to define uniform worker resources."
|
||||
)
|
||||
else:
|
||||
# Default to 1 TPU per bundle.
|
||||
worker_bundle = {"TPU": 1}
|
||||
|
||||
if self._slice_pg_wrapper is not None:
|
||||
logger.debug(
|
||||
"Existing TPU slice PG found. Shutting it down before creating a new one."
|
||||
)
|
||||
self.shutdown()
|
||||
|
||||
self._slice_pg_wrapper = slice_placement_group(
|
||||
topology=self._config.topology,
|
||||
accelerator_version=version,
|
||||
resources_per_bundle=worker_bundle,
|
||||
strategy=strategy,
|
||||
name=name,
|
||||
)
|
||||
return self._slice_pg_wrapper.placement_group
|
||||
|
||||
@property
|
||||
def requires_deferred_placement_group(self) -> bool:
|
||||
"""
|
||||
If a TPU topology is specified, we defer PG creation so the replica can
|
||||
provision a `SlicePlacementGroup` at runtime. This ensures multi-host
|
||||
TPU slices are gang-scheduled atomically according to their physical
|
||||
topology rather than fragmented across the cluster.
|
||||
"""
|
||||
return bool(self._config.topology)
|
||||
|
||||
@property
|
||||
def requires_remote_initialization(self) -> bool:
|
||||
return True
|
||||
|
||||
def get_remote_options(self, accelerator_type_str: str = None):
|
||||
# The PlacementGroupSchedulingStrategy natively handles routing the task to
|
||||
# the correct hardware. We omit TPU resource requests to avoid consuming
|
||||
# chips that the model engine workers must use.
|
||||
options: Dict[str, Any] = {"resources": {}}
|
||||
if accelerator_type_str:
|
||||
# Pin the task to the TPU accelerator to avoid scheduling on a CPU bundle.
|
||||
options["label_selector"] = {
|
||||
"ray.io/accelerator-type": accelerator_type_str
|
||||
}
|
||||
return options
|
||||
|
||||
def shutdown(self):
|
||||
if self._slice_pg_wrapper is not None:
|
||||
try:
|
||||
logger.info("Shutting down TPU slice PG for server replica.")
|
||||
self._slice_pg_wrapper.shutdown()
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to shut down TPU slice PG: {e}")
|
||||
finally:
|
||||
self._slice_pg_wrapper = None
|
||||
@@ -0,0 +1,658 @@
|
||||
from enum import Enum
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Dict,
|
||||
Optional,
|
||||
TypeVar,
|
||||
Union,
|
||||
)
|
||||
|
||||
from pydantic import (
|
||||
BaseModel,
|
||||
Field,
|
||||
PositiveInt,
|
||||
PrivateAttr,
|
||||
field_validator,
|
||||
model_validator,
|
||||
)
|
||||
|
||||
from ray.llm._internal.common.base_pydantic import BaseModelExtended
|
||||
from ray.llm._internal.common.callbacks.base import (
|
||||
CallbackBase,
|
||||
CallbackConfig,
|
||||
)
|
||||
from ray.llm._internal.common.utils.cloud_utils import (
|
||||
CloudMirrorConfig,
|
||||
is_remote_path,
|
||||
)
|
||||
from ray.llm._internal.common.utils.download_utils import (
|
||||
STREAMING_LOAD_FORMATS,
|
||||
NodeModelDownloadable,
|
||||
)
|
||||
from ray.llm._internal.common.utils.import_utils import load_class, try_import
|
||||
from ray.llm._internal.serve.constants import (
|
||||
DEFAULT_MULTIPLEX_DOWNLOAD_TIMEOUT_S,
|
||||
DEFAULT_MULTIPLEX_DOWNLOAD_TRIES,
|
||||
MODEL_RESPONSE_BATCH_TIMEOUT_MS,
|
||||
)
|
||||
from ray.llm._internal.serve.core.configs.accelerators import (
|
||||
TPU_ACCELERATOR_VALUES,
|
||||
AcceleratorType,
|
||||
AnyAcceleratorConfig,
|
||||
CPUConfig,
|
||||
GPUConfig,
|
||||
TPUConfig,
|
||||
infer_hardware_kind_from_bundles,
|
||||
)
|
||||
from ray.llm._internal.serve.engines.vllm.kv_transfer.factory import (
|
||||
KVConnectorBackendFactory,
|
||||
)
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
from ray.serve._private.config import DeploymentConfig, handle_num_replicas_auto
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.llm._internal.serve.engines.vllm.kv_transfer.base import (
|
||||
BaseConnectorBackend,
|
||||
)
|
||||
|
||||
transformers = try_import("transformers")
|
||||
|
||||
ModelT = TypeVar("ModelT", bound=BaseModel)
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class ServeMultiplexConfig(BaseModelExtended):
|
||||
max_num_models_per_replica: PositiveInt = Field(
|
||||
..., description="The maximum number of models to be loaded on each replica."
|
||||
)
|
||||
download_timeout_s: Optional[float] = Field(
|
||||
DEFAULT_MULTIPLEX_DOWNLOAD_TIMEOUT_S,
|
||||
description="How much time the download subprocess has to download a single LoRA before a timeout. None means no timeout.",
|
||||
)
|
||||
max_download_tries: int = Field(
|
||||
DEFAULT_MULTIPLEX_DOWNLOAD_TRIES,
|
||||
description="The maximum number of download retries.",
|
||||
)
|
||||
|
||||
|
||||
class InputModality(str, Enum):
|
||||
text = "text"
|
||||
image = "image"
|
||||
|
||||
|
||||
class LLMEngine(str, Enum):
|
||||
"""Enum that represents an LLMEngine."""
|
||||
|
||||
vLLM = "vLLM"
|
||||
|
||||
|
||||
class LoraConfig(BaseModelExtended):
|
||||
dynamic_lora_loading_path: Optional[str] = Field(
|
||||
default=None,
|
||||
description="Cloud storage path where LoRA adapter weights are stored.",
|
||||
)
|
||||
max_num_adapters_per_replica: PositiveInt = Field(
|
||||
default=16,
|
||||
description="The maximum number of adapters to load on each replica.",
|
||||
)
|
||||
download_timeout_s: Optional[float] = Field(
|
||||
DEFAULT_MULTIPLEX_DOWNLOAD_TIMEOUT_S,
|
||||
description=(
|
||||
"How much time the download subprocess has to download a single "
|
||||
"LoRA before a timeout. None means no timeout."
|
||||
),
|
||||
)
|
||||
max_download_tries: int = Field(
|
||||
DEFAULT_MULTIPLEX_DOWNLOAD_TRIES,
|
||||
description="The maximum number of download retries.",
|
||||
)
|
||||
|
||||
@field_validator("dynamic_lora_loading_path")
|
||||
def validate_dynamic_lora_loading_path(cls, value: Optional[str]):
|
||||
if value is None:
|
||||
return value
|
||||
|
||||
assert is_remote_path(value), (
|
||||
"Only AWS S3, Google Cloud Storage, and Azure Storage are supported. The "
|
||||
'dynamic_lora_loading_path must start with "s3://", "gs://", "abfss://", or "azure://". '
|
||||
f'Got "{value}" instead.'
|
||||
)
|
||||
return value.rstrip("/")
|
||||
|
||||
|
||||
class ModelLoadingConfig(BaseModelExtended):
|
||||
|
||||
model_id: str = Field(
|
||||
description="The ID that should be used by end users to access this model.",
|
||||
)
|
||||
model_source: Optional[Union[str, CloudMirrorConfig]] = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Where to obtain the model weights from. "
|
||||
"Should be a HuggingFace model ID, S3 mirror config, GCS mirror config, "
|
||||
"or a local path. When omitted, defaults to the model_id as a "
|
||||
"HuggingFace model ID."
|
||||
),
|
||||
)
|
||||
tokenizer_source: Optional[str] = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Where to obtain the tokenizer from. If None, tokenizer is "
|
||||
"obtained from the model source. Only HuggingFace IDs are "
|
||||
"supported for now."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
EngineConfigType = Union[None, "VLLMEngineConfig"] # noqa: F821
|
||||
|
||||
|
||||
class LLMConfig(BaseModelExtended):
|
||||
|
||||
runtime_env: Optional[Dict[str, Any]] = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"The runtime_env to use for the model deployment replica "
|
||||
"and the engine workers."
|
||||
),
|
||||
)
|
||||
|
||||
model_loading_config: Union[Dict[str, Any], ModelLoadingConfig] = Field(
|
||||
description="The settings for how to download and expose the model. Validated against ModelLoadingConfig."
|
||||
)
|
||||
|
||||
llm_engine: str = Field(
|
||||
default=LLMEngine.vLLM.value,
|
||||
description=f"The LLMEngine that should be used to run the model. Only the following values are supported: {str([t.value for t in LLMEngine])}",
|
||||
)
|
||||
|
||||
engine_kwargs: Dict[str, Any] = Field(
|
||||
default={},
|
||||
description=(
|
||||
"Additional keyword arguments for the engine. In case of vLLM, "
|
||||
"this will include all the configuration knobs they provide out "
|
||||
"of the box"
|
||||
),
|
||||
)
|
||||
|
||||
accelerator_type: Optional[str] = Field(
|
||||
default=None,
|
||||
description=f"The type of accelerator runs the model on. Only the following values are supported: {str([t.value for t in AcceleratorType])}",
|
||||
)
|
||||
|
||||
accelerator_config: Optional[AnyAcceleratorConfig] = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Hardware-specific configuration parameters for the chosen accelerator. "
|
||||
"The expected schema is dynamically typed based on the 'kind' discriminator."
|
||||
),
|
||||
)
|
||||
|
||||
placement_group_config: Optional[Dict[str, Any]] = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Ray placement group configuration for scheduling vLLM engine workers. "
|
||||
"Defines resource bundles and placement strategy for multi-node deployments. "
|
||||
"Can specify either 'bundle_per_worker' (auto-replicated by tp*pp) or 'bundles' "
|
||||
"(full list of resource dicts). Optionally include 'strategy' key "
|
||||
"('PACK', 'STRICT_PACK', 'SPREAD', or 'STRICT_SPREAD'). "
|
||||
"Example with bundle_per_worker: {'bundle_per_worker': {'CPU': 1, 'GPU': 1}, 'strategy': 'SPREAD'}. "
|
||||
"Example with bundles: {'bundles': [{'CPU': 1, 'GPU': 1}] * 4, 'strategy': 'SPREAD'}."
|
||||
),
|
||||
)
|
||||
|
||||
lora_config: Optional[Union[Dict[str, Any], LoraConfig]] = Field(
|
||||
default=None,
|
||||
description="Settings for LoRA adapter. Validated against LoraConfig.",
|
||||
)
|
||||
|
||||
deployment_config: Dict[str, Any] = Field(
|
||||
default_factory=dict,
|
||||
description="""
|
||||
The Ray @server.deployment options.
|
||||
Supported fields are:
|
||||
`name`, `num_replicas`, `ray_actor_options`, `max_ongoing_requests`,
|
||||
`autoscaling_config`, `max_queued_requests`, `user_config`,
|
||||
`health_check_period_s`, `health_check_timeout_s`,
|
||||
`graceful_shutdown_wait_loop_s`, `graceful_shutdown_timeout_s`,
|
||||
`logging_config`, `request_router_config`.
|
||||
For more details, see the `Ray Serve Documentation <https://docs.ray.io/en/latest/serve/configure-serve-deployment.html>`_.
|
||||
""",
|
||||
)
|
||||
|
||||
server_cls: Optional[Union[str, Any]] = Field(
|
||||
default=None,
|
||||
description="The serve class to use.(e.g., LLMServer, SGLangServer or other Server backends).",
|
||||
)
|
||||
|
||||
@field_validator("server_cls")
|
||||
@classmethod
|
||||
def validate_server_cls(cls, value):
|
||||
if isinstance(value, str):
|
||||
return load_class(value)
|
||||
return value
|
||||
|
||||
experimental_configs: Dict[str, Any] = Field(
|
||||
default_factory=dict,
|
||||
description="Experimental configurations for Ray Serve LLM. This is a "
|
||||
"dictionary of key-value pairs. Current supported keys are:\n"
|
||||
"- `stream_batching_interval_ms`: Ray Serve LLM batches streaming "
|
||||
"requests together. This config decides how long to wait for the "
|
||||
"batch before processing the requests. Defaults to "
|
||||
f"{MODEL_RESPONSE_BATCH_TIMEOUT_MS}.\n"
|
||||
"- `num_ingress_replicas`: The number of replicas for the router. Ray "
|
||||
"Serve will take the max amount all the replicas. Default would be 2 "
|
||||
"router replicas per model replica.\n",
|
||||
)
|
||||
|
||||
log_engine_metrics: Optional[bool] = Field(
|
||||
default=True,
|
||||
description="Enable additional engine metrics via Ray Prometheus port.",
|
||||
)
|
||||
|
||||
callback_config: CallbackConfig = Field(
|
||||
default_factory=CallbackConfig,
|
||||
description="Callback configuration to use for model initialization. Can be a string path to a class or a Callback subclass.",
|
||||
)
|
||||
|
||||
_supports_vision: bool = PrivateAttr(False)
|
||||
_model_architecture: str = PrivateAttr("UNSPECIFIED")
|
||||
_engine_config: EngineConfigType = PrivateAttr(None)
|
||||
_callback_instance: Optional[CallbackBase] = PrivateAttr(None)
|
||||
_kv_connector_backend: Optional["BaseConnectorBackend"] = PrivateAttr(None)
|
||||
|
||||
def _load_hf_config(self, model_id_or_path: str, trust_remote_code: bool = False):
|
||||
"""Load the HuggingFace config for a model.
|
||||
|
||||
Uses AutoConfig which loads the model-specific config class (e.g.
|
||||
DeepseekV3Config) instead of the generic PretrainedConfig. The generic
|
||||
base class can fail for models whose config.json contains fields (like
|
||||
``rope_scaling``) that require model-specific post-init logic.
|
||||
"""
|
||||
try:
|
||||
return transformers.AutoConfig.from_pretrained(
|
||||
model_id_or_path, trust_remote_code=trust_remote_code
|
||||
)
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
f"Failed to load Hugging Face config for "
|
||||
f"model_id='{model_id_or_path}'. Ensure `model_id` is a valid "
|
||||
f"Hugging Face repo or a local path that contains a valid "
|
||||
f"`config.json` file. Original error: {repr(e)}"
|
||||
) from e
|
||||
|
||||
def _infer_supports_vision(
|
||||
self, model_id_or_path: str, trust_remote_code: bool = False
|
||||
) -> None:
|
||||
"""Called in llm node initializer together with other transformers calls. It
|
||||
loads the model config from huggingface and sets the supports_vision
|
||||
attribute based on whether the config has `vision_config`. All LVM models has
|
||||
`vision_config` setup.
|
||||
"""
|
||||
hf_config = self._load_hf_config(
|
||||
model_id_or_path, trust_remote_code=trust_remote_code
|
||||
)
|
||||
self._supports_vision = hasattr(hf_config, "vision_config")
|
||||
|
||||
def _set_model_architecture(
|
||||
self,
|
||||
model_id_or_path: Optional[str] = None,
|
||||
model_architecture: Optional[str] = None,
|
||||
trust_remote_code: bool = False,
|
||||
) -> None:
|
||||
"""Called in llm node initializer together with other transformers calls. It
|
||||
loads the model config from huggingface and sets the model_architecture
|
||||
attribute based on whether the config has `architectures`.
|
||||
"""
|
||||
if model_id_or_path:
|
||||
hf_config = self._load_hf_config(
|
||||
model_id_or_path, trust_remote_code=trust_remote_code
|
||||
)
|
||||
if (
|
||||
hf_config
|
||||
and hasattr(hf_config, "architectures")
|
||||
and hf_config.architectures
|
||||
):
|
||||
self._model_architecture = hf_config.architectures[0]
|
||||
|
||||
if model_architecture:
|
||||
self._model_architecture = model_architecture
|
||||
|
||||
def apply_checkpoint_info(
|
||||
self, model_id_or_path: str, trust_remote_code: bool = False
|
||||
) -> None:
|
||||
"""Apply the checkpoint info to the model config."""
|
||||
self._infer_supports_vision(
|
||||
model_id_or_path, trust_remote_code=trust_remote_code
|
||||
)
|
||||
self._set_model_architecture(
|
||||
model_id_or_path, trust_remote_code=trust_remote_code
|
||||
)
|
||||
|
||||
def get_or_create_callback(self) -> Optional[CallbackBase]:
|
||||
"""Get or create the callback instance for this process.
|
||||
|
||||
This ensures one callback instance per process (singleton pattern).
|
||||
The instance is cached so the same object is used across all hooks.
|
||||
|
||||
Returns:
|
||||
Instance of class that implements Callback
|
||||
""" # Return cached instance if exists
|
||||
if self._callback_instance is not None:
|
||||
return self._callback_instance
|
||||
|
||||
engine_config = self.get_engine_config()
|
||||
assert engine_config is not None
|
||||
pg = engine_config.get_or_create_pg()
|
||||
runtime_env = engine_config.get_runtime_env_with_local_env_vars()
|
||||
if self.engine_kwargs.get("load_format", None) in STREAMING_LOAD_FORMATS:
|
||||
worker_node_download_model = NodeModelDownloadable.NONE
|
||||
else:
|
||||
worker_node_download_model = NodeModelDownloadable.MODEL_AND_TOKENIZER
|
||||
|
||||
# Create new instance
|
||||
if isinstance(self.callback_config.callback_class, str):
|
||||
callback_class = load_class(self.callback_config.callback_class)
|
||||
else:
|
||||
callback_class = self.callback_config.callback_class
|
||||
|
||||
self._callback_instance = callback_class(
|
||||
raise_error_on_callback=self.callback_config.raise_error_on_callback,
|
||||
llm_config=self,
|
||||
ctx_kwargs={
|
||||
"worker_node_download_model": worker_node_download_model,
|
||||
"placement_group": pg,
|
||||
"runtime_env": runtime_env,
|
||||
},
|
||||
**self.callback_config.callback_kwargs,
|
||||
)
|
||||
return self._callback_instance
|
||||
|
||||
@property
|
||||
def supports_vision(self) -> bool:
|
||||
return self._supports_vision
|
||||
|
||||
@property
|
||||
def model_architecture(self) -> str:
|
||||
return self._model_architecture
|
||||
|
||||
@property
|
||||
def input_modality(self) -> str:
|
||||
"""Returns the input modality of the model. There could be more types in the
|
||||
future. Right now assumes if the model doesn't support version, it'll be text.
|
||||
"""
|
||||
if self.supports_vision:
|
||||
return InputModality.image.value
|
||||
|
||||
return InputModality.text.value
|
||||
|
||||
@property
|
||||
def model_id(self) -> str:
|
||||
return self.model_loading_config.model_id
|
||||
|
||||
@property
|
||||
def max_request_context_length(self) -> Optional[int]:
|
||||
return self.engine_kwargs.get("max_model_len")
|
||||
|
||||
@field_validator("accelerator_type")
|
||||
def validate_accelerator_type(cls, value: Optional[str]):
|
||||
if value is None:
|
||||
return value
|
||||
|
||||
# Ensure A10 is converted to A10G.
|
||||
if value == "A10":
|
||||
value = "A10G"
|
||||
|
||||
if value not in [t.value for t in AcceleratorType]:
|
||||
raise ValueError(f"Unsupported accelerator type: {value}")
|
||||
|
||||
return value
|
||||
|
||||
@field_validator("llm_engine")
|
||||
def validate_llm_engine(cls, value: str) -> str:
|
||||
"""Validates the llm_engine string value."""
|
||||
try:
|
||||
# Validate the engine
|
||||
LLMEngine(value)
|
||||
except ValueError as e:
|
||||
raise ValueError(f"Unsupported engine: {value}") from e
|
||||
return value
|
||||
|
||||
@field_validator("deployment_config")
|
||||
def validate_deployment_config(cls, value: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Validates the deployment config dictionary."""
|
||||
try:
|
||||
# Resolve "auto" for num_replicas before validating against DeploymentConfig
|
||||
if value.get("num_replicas") == "auto":
|
||||
resolved = {**value, "num_replicas": None}
|
||||
_, autoscaling_config = handle_num_replicas_auto(
|
||||
resolved.get("max_ongoing_requests"),
|
||||
resolved.get("autoscaling_config"),
|
||||
)
|
||||
resolved["autoscaling_config"] = autoscaling_config
|
||||
DeploymentConfig(**resolved)
|
||||
else:
|
||||
DeploymentConfig(**value)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Invalid deployment config: {value}") from e
|
||||
|
||||
return value
|
||||
|
||||
@field_validator("model_loading_config")
|
||||
def validate_model_loading_config(
|
||||
cls, value: Union[Dict[str, Any], ModelLoadingConfig]
|
||||
) -> ModelLoadingConfig:
|
||||
"""Validates the model loading config dictionary."""
|
||||
if isinstance(value, ModelLoadingConfig):
|
||||
return value
|
||||
|
||||
try:
|
||||
model_loading_config = ModelLoadingConfig(**value)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Invalid model_loading_config: {value}") from e
|
||||
|
||||
return model_loading_config
|
||||
|
||||
@field_validator("lora_config")
|
||||
def validate_lora_config(
|
||||
cls, value: Optional[Union[Dict[str, Any], LoraConfig]]
|
||||
) -> Optional[LoraConfig]:
|
||||
"""Validates the lora config dictionary."""
|
||||
if value is None or isinstance(value, LoraConfig):
|
||||
return value
|
||||
|
||||
try:
|
||||
lora_config = LoraConfig(**value)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Invalid lora_config: {value}") from e
|
||||
|
||||
return lora_config
|
||||
|
||||
@field_validator("experimental_configs")
|
||||
def validate_experimental_configs(cls, value: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Validates the experimental configs dictionary."""
|
||||
# TODO(Kourosh): Remove this deprecation check after users have
|
||||
# migrated.
|
||||
if "num_router_replicas" in value:
|
||||
raise ValueError(
|
||||
"The 'num_router_replicas' key in experimental_configs has "
|
||||
"been renamed to 'num_ingress_replicas'. Please update "
|
||||
"your configuration to use 'num_ingress_replicas' instead."
|
||||
)
|
||||
return value
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _check_log_stats_with_metrics(self):
|
||||
"""Validate that disable_log_stats isn't enabled when log_engine_metrics is enabled."""
|
||||
if self.log_engine_metrics and self.engine_kwargs.get("disable_log_stats"):
|
||||
raise ValueError(
|
||||
"disable_log_stats cannot be set to True when log_engine_metrics is enabled. "
|
||||
"Engine metrics require log stats to be enabled."
|
||||
)
|
||||
|
||||
return self
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _resolve_and_validate_accelerator(self):
|
||||
"""Resolves the accelerator configuration and validates it."""
|
||||
self._resolve_accelerator_config()
|
||||
self._check_accelerator_type_matches_hardware()
|
||||
return self
|
||||
|
||||
def _resolve_accelerator_config(self) -> None:
|
||||
"""Infers and populates accelerator_config if omitted by the user."""
|
||||
if self.accelerator_config is not None:
|
||||
return
|
||||
|
||||
# Infer hardware from placement_group_config bundles
|
||||
inferred_kind = infer_hardware_kind_from_bundles(self.placement_group_config)
|
||||
|
||||
if inferred_kind == "tpu":
|
||||
self.accelerator_config = TPUConfig(kind="tpu")
|
||||
return
|
||||
if inferred_kind == "gpu":
|
||||
self.accelerator_config = GPUConfig(kind="gpu")
|
||||
return
|
||||
if inferred_kind == "cpu":
|
||||
self.accelerator_config = CPUConfig(kind="cpu")
|
||||
return
|
||||
|
||||
# Infer hardware from accelerator_type string
|
||||
if self.accelerator_type:
|
||||
accel_str = getattr(
|
||||
self.accelerator_type, "value", str(self.accelerator_type)
|
||||
)
|
||||
if accel_str in TPU_ACCELERATOR_VALUES:
|
||||
self.accelerator_config = TPUConfig(kind="tpu")
|
||||
return
|
||||
|
||||
self.accelerator_config = GPUConfig(kind="gpu")
|
||||
return
|
||||
|
||||
# Default to GPUConfig if not otherwise specified
|
||||
self.accelerator_config = GPUConfig(kind="gpu")
|
||||
|
||||
def _check_accelerator_type_matches_hardware(self) -> None:
|
||||
"""Validate that accelerator_type aligns with the hardware configuration."""
|
||||
if isinstance(self.accelerator_config, TPUConfig):
|
||||
# For TPU slices, both accelerator_type and topology must be provided.
|
||||
if self.accelerator_config.topology and not self.accelerator_type:
|
||||
raise ValueError(
|
||||
"accelerator_type must be provided when specifying a TPU topology "
|
||||
"for TPU slice provisioning."
|
||||
)
|
||||
|
||||
if not self.accelerator_type:
|
||||
return
|
||||
|
||||
if isinstance(self.accelerator_config, CPUConfig):
|
||||
raise ValueError(
|
||||
f"accelerator_type='{self.accelerator_type}' cannot be used with "
|
||||
"CPU-only configurations. Either remove accelerator_type, or provide an accelerator_config."
|
||||
)
|
||||
|
||||
# Determine what hardware kind the string implies to check for kind mismatch
|
||||
accel_str = getattr(self.accelerator_type, "value", str(self.accelerator_type))
|
||||
expected_kind = "tpu" if accel_str in TPU_ACCELERATOR_VALUES else "gpu"
|
||||
|
||||
if self.accelerator_config.kind != expected_kind:
|
||||
raise ValueError(
|
||||
f"Hardware mismatch: accelerator_type='{self.accelerator_type}' requires a "
|
||||
f"{expected_kind.upper()} backend, but the configuration resolved to a "
|
||||
f"{self.accelerator_config.kind.upper()} backend. Please ensure your "
|
||||
f"bundles and accelerator_type align."
|
||||
)
|
||||
|
||||
def multiplex_config(self) -> ServeMultiplexConfig:
|
||||
multiplex_config = None
|
||||
if self.lora_config:
|
||||
multiplex_config = ServeMultiplexConfig(
|
||||
max_num_models_per_replica=self.lora_config.max_num_adapters_per_replica,
|
||||
download_timeout_s=self.lora_config.download_timeout_s,
|
||||
max_download_tries=self.lora_config.max_download_tries,
|
||||
)
|
||||
return multiplex_config
|
||||
|
||||
def get_engine_config(self) -> EngineConfigType:
|
||||
"""Returns the engine config for the given LLM config.
|
||||
|
||||
LLMConfig not only has engine config but also deployment config, etc.
|
||||
"""
|
||||
# Note (genesu): This is important that we cache the engine config as the
|
||||
# `hf_model_id` attribute on the engine config will be set based on whether
|
||||
# the model is downloaded from a remote storage and will be set to the
|
||||
# local path of the model. This is important for vLLM not going to Hugging
|
||||
# Face to download the model again after it's already downloaded during node
|
||||
# initialization step.
|
||||
if self._engine_config:
|
||||
return self._engine_config
|
||||
|
||||
if self.llm_engine == LLMEngine.vLLM:
|
||||
from ray.llm._internal.serve.engines.vllm.vllm_models import (
|
||||
VLLMEngineConfig,
|
||||
)
|
||||
|
||||
self._engine_config = VLLMEngineConfig.from_llm_config(self)
|
||||
else:
|
||||
# Note (genesu): This should never happen because we validate the engine
|
||||
# in the config.
|
||||
raise ValueError(f"Unsupported engine: {self.llm_engine}")
|
||||
|
||||
return self._engine_config
|
||||
|
||||
def update_engine_kwargs(self, **kwargs: Any) -> None:
|
||||
"""Update the engine_kwargs and the engine_config engine_kwargs.
|
||||
|
||||
This is typically called during engine starts, when certain engine_kwargs
|
||||
(e.g., data_parallel_rank) become available.
|
||||
"""
|
||||
self.engine_kwargs.update(kwargs)
|
||||
# engine_config may be created before engine starts, this makes sure
|
||||
# the engine_config is updated with the latest engine_kwargs.
|
||||
if self._engine_config:
|
||||
self._engine_config.engine_kwargs.update(kwargs)
|
||||
|
||||
def setup_engine_backend(self):
|
||||
self._setup_kv_connector_backend()
|
||||
|
||||
def _setup_kv_connector_backend(self):
|
||||
"""Private method to setup kv connector depending on the local deployment state"""
|
||||
# 1. validate that the backend is one of the backends supported (Nixl or LMCache)
|
||||
kv_transfer_config = self.engine_kwargs.get("kv_transfer_config")
|
||||
if not kv_transfer_config:
|
||||
return
|
||||
|
||||
kv_connector = kv_transfer_config.get("kv_connector")
|
||||
if not kv_connector:
|
||||
raise ValueError("Connector type is not specified.")
|
||||
|
||||
# 2. Setup the backend using factory
|
||||
kv_connector_backend = KVConnectorBackendFactory.create_backend(
|
||||
kv_connector, self
|
||||
)
|
||||
kv_connector_backend.setup()
|
||||
# 3. Stash the instance so the P/D orchestrator can reach the connector's
|
||||
# coordination protocol (request shaping, peer binding, handoff
|
||||
# discipline) without re-creating it. May be None on configs that never
|
||||
# call setup_engine_backend(); the orchestrator falls back to the factory.
|
||||
self._kv_connector_backend = kv_connector_backend
|
||||
|
||||
@property
|
||||
def kv_connector_backend(self) -> Optional["BaseConnectorBackend"]:
|
||||
"""The KV-connector backend instance created by ``setup_engine_backend``.
|
||||
|
||||
Returns None if no KV transfer connector is configured, or if the
|
||||
backend has not been set up yet on this config copy.
|
||||
"""
|
||||
return self._kv_connector_backend
|
||||
|
||||
|
||||
class DiskMultiplexConfig(BaseModelExtended):
|
||||
model_id: str
|
||||
max_total_tokens: Optional[int]
|
||||
local_path: str
|
||||
|
||||
# this is a per process id assigned to the model
|
||||
lora_assigned_int_id: int
|
||||
@@ -0,0 +1,326 @@
|
||||
"""This module contains wrapper classes for OpenAI-compatible protocol models.
|
||||
|
||||
Supports both vLLM and SGLang as the underlying engine. vLLM is tried first;
|
||||
on ImportError, SGLang models are imported as a fallback. If neither is
|
||||
installed, an ImportError is raised at import time.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from typing import TYPE_CHECKING, Any, AsyncGenerator, Dict, List, Optional, Union
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from ray.llm._internal.common.utils.import_utils import raise_llm_engine_import_error
|
||||
|
||||
try:
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import (
|
||||
ChatCompletionRequest as _ChatCompletionRequest,
|
||||
ChatCompletionResponse as _ChatCompletionResponse,
|
||||
ChatCompletionStreamResponse as _ChatCompletionStreamResponse,
|
||||
)
|
||||
from vllm.entrypoints.openai.completion.protocol import (
|
||||
CompletionRequest as _CompletionRequest,
|
||||
CompletionResponse as _CompletionResponse,
|
||||
CompletionStreamResponse as _CompletionStreamResponse,
|
||||
)
|
||||
from vllm.entrypoints.openai.engine.protocol import (
|
||||
ErrorInfo as _ErrorInfo,
|
||||
ErrorResponse as _ErrorResponse,
|
||||
)
|
||||
from vllm.entrypoints.pooling.embed.protocol import (
|
||||
EmbeddingChatRequest as _EmbeddingChatRequest,
|
||||
EmbeddingCompletionRequest as _EmbeddingCompletionRequest,
|
||||
EmbeddingResponse as _EmbeddingResponse,
|
||||
)
|
||||
from vllm.entrypoints.pooling.scoring.protocol import (
|
||||
ScoreResponse as _ScoreResponse,
|
||||
ScoreTextRequest as _ScoreTextRequest,
|
||||
)
|
||||
from vllm.entrypoints.serve.tokenize.protocol import (
|
||||
DetokenizeRequest as _DetokenizeRequest,
|
||||
DetokenizeResponse as _DetokenizeResponse,
|
||||
TokenizeChatRequest as _TokenizeChatRequest,
|
||||
TokenizeCompletionRequest as _TokenizeCompletionRequest,
|
||||
TokenizeResponse as _TokenizeResponse,
|
||||
)
|
||||
from vllm.entrypoints.speech_to_text.transcription.protocol import (
|
||||
TranscriptionRequest as _TranscriptionRequest,
|
||||
TranscriptionResponse as _TranscriptionResponse,
|
||||
TranscriptionStreamResponse as _TranscriptionStreamResponse,
|
||||
)
|
||||
|
||||
except ImportError as _vllm_import_error:
|
||||
try:
|
||||
from sglang.srt.entrypoints.openai.protocol import (
|
||||
ChatCompletionRequest as _ChatCompletionRequest,
|
||||
ChatCompletionResponse as _ChatCompletionResponse,
|
||||
ChatCompletionStreamResponse as _ChatCompletionStreamResponse,
|
||||
CompletionRequest as _CompletionRequest,
|
||||
CompletionResponse as _CompletionResponse,
|
||||
CompletionStreamResponse as _CompletionStreamResponse,
|
||||
DetokenizeRequest as _DetokenizeRequest,
|
||||
DetokenizeResponse as _DetokenizeResponse,
|
||||
EmbeddingRequest as _EmbeddingCompletionRequest,
|
||||
EmbeddingResponse as _EmbeddingResponse,
|
||||
ScoringRequest as _ScoreTextRequest,
|
||||
ScoringResponse as _ScoreResponse,
|
||||
TokenizeRequest as _TokenizeCompletionRequest,
|
||||
TokenizeResponse as _TokenizeResponse,
|
||||
)
|
||||
except ImportError as _sglang_import_error:
|
||||
raise_llm_engine_import_error(_vllm_import_error, _sglang_import_error)
|
||||
|
||||
def _unsupported_model(name: str, feature: str = ""):
|
||||
"""Create a BaseModel stub that raises NotImplementedError on instantiation."""
|
||||
msg = f"{name} is not supported with the current backend." + (
|
||||
f" {feature}" if feature else ""
|
||||
)
|
||||
|
||||
class _Stub(BaseModel):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
raise NotImplementedError(msg)
|
||||
|
||||
_Stub.__name__ = _Stub.__qualname__ = name
|
||||
return _Stub
|
||||
|
||||
# SGLang does not provide transcription protocol models.
|
||||
_vllm_hint = "Install vLLM to use transcription endpoints."
|
||||
_TranscriptionRequest = _unsupported_model("TranscriptionRequest", _vllm_hint)
|
||||
_TranscriptionResponse = _unsupported_model("TranscriptionResponse", _vllm_hint)
|
||||
_TranscriptionStreamResponse = _unsupported_model(
|
||||
"TranscriptionStreamResponse", _vllm_hint
|
||||
)
|
||||
|
||||
# SGLang has no equivalent to vLLM's nested ErrorResponse.error -> ErrorInfo
|
||||
# pattern, so we define our own.
|
||||
|
||||
class _ErrorInfo(BaseModel):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
message: str
|
||||
type: str
|
||||
param: Optional[str] = None
|
||||
code: int
|
||||
|
||||
class _ErrorResponse(BaseModel):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
error: _ErrorInfo
|
||||
|
||||
_EmbeddingChatRequest = _EmbeddingCompletionRequest
|
||||
_TokenizeChatRequest = _TokenizeCompletionRequest
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
|
||||
|
||||
|
||||
class ChatCompletionRequest(_ChatCompletionRequest):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
|
||||
class ChatCompletionResponse(_ChatCompletionResponse):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
|
||||
class ChatCompletionStreamResponse(_ChatCompletionStreamResponse):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
|
||||
class ErrorInfo(_ErrorInfo):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
|
||||
class ErrorResponse(_ErrorResponse):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
|
||||
# TODO (Kourosh): Upstream
|
||||
class CompletionRequest(_CompletionRequest):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
|
||||
class CompletionResponse(_CompletionResponse):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
|
||||
class CompletionStreamResponse(_CompletionStreamResponse):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
|
||||
# TODO (Kourosh): Upstream
|
||||
class EmbeddingCompletionRequest(_EmbeddingCompletionRequest):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
|
||||
class EmbeddingChatRequest(_EmbeddingChatRequest):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
|
||||
class EmbeddingResponse(_EmbeddingResponse):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
|
||||
class TranscriptionRequest(_TranscriptionRequest):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
request_id: str = Field(
|
||||
default_factory=lambda: str(uuid.uuid4()),
|
||||
description=(
|
||||
"The request_id related to this request. If the caller does "
|
||||
"not set it, a random_uuid will be generated. This id is used "
|
||||
"through out the inference process and return in response."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class TranscriptionResponse(_TranscriptionResponse):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
|
||||
class TranscriptionStreamResponse(_TranscriptionStreamResponse):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
|
||||
class ScoreRequest(_ScoreTextRequest):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
|
||||
class ScoreResponse(_ScoreResponse):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
|
||||
class TokenizeCompletionRequest(_TokenizeCompletionRequest):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
|
||||
class TokenizeChatRequest(_TokenizeChatRequest):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
|
||||
class TokenizeResponse(_TokenizeResponse):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
|
||||
class DetokenizeRequest(_DetokenizeRequest):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
|
||||
class DetokenizeResponse(_DetokenizeResponse):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
|
||||
EmbeddingRequest = Union[EmbeddingCompletionRequest, EmbeddingChatRequest]
|
||||
|
||||
TokenizeRequest = Union[TokenizeCompletionRequest, TokenizeChatRequest]
|
||||
|
||||
LLMEmbeddingsResponse = Union[
|
||||
AsyncGenerator[Union[EmbeddingResponse, ErrorResponse], None],
|
||||
]
|
||||
|
||||
LLMScoreResponse = Union[
|
||||
AsyncGenerator[Union[ScoreResponse, ErrorResponse], None],
|
||||
]
|
||||
|
||||
LLMTokenizeResponse = Union[
|
||||
AsyncGenerator[Union[TokenizeResponse, ErrorResponse], None],
|
||||
]
|
||||
|
||||
LLMDetokenizeResponse = Union[
|
||||
AsyncGenerator[Union[DetokenizeResponse, ErrorResponse], None],
|
||||
]
|
||||
|
||||
LLMChatResponse = Union[
|
||||
AsyncGenerator[
|
||||
Union[str, ChatCompletionStreamResponse, ChatCompletionResponse, ErrorResponse],
|
||||
None,
|
||||
],
|
||||
]
|
||||
|
||||
LLMCompletionsResponse = Union[
|
||||
AsyncGenerator[
|
||||
Union[str, CompletionStreamResponse, CompletionResponse, ErrorResponse], None
|
||||
],
|
||||
]
|
||||
|
||||
LLMTranscriptionResponse = Union[
|
||||
AsyncGenerator[
|
||||
Union[str, TranscriptionStreamResponse, TranscriptionResponse, ErrorResponse],
|
||||
None,
|
||||
],
|
||||
]
|
||||
|
||||
|
||||
# TODO: remove this class
|
||||
class OpenAIHTTPException(Exception):
|
||||
def __init__(
|
||||
self,
|
||||
status_code: int,
|
||||
message: str,
|
||||
type: str = "Unknown",
|
||||
internal_message: Optional[str] = None,
|
||||
) -> None:
|
||||
self.status_code = status_code
|
||||
self.message = message
|
||||
self.type = type
|
||||
self.internal_message = internal_message
|
||||
|
||||
|
||||
# TODO: upstream metadata for ModelData
|
||||
# Compared to vLLM this has a metadata field.
|
||||
class ModelCard(BaseModel):
|
||||
model_config = ConfigDict(
|
||||
protected_namespaces=tuple(), arbitrary_types_allowed=True
|
||||
)
|
||||
|
||||
id: str
|
||||
object: str
|
||||
owned_by: str
|
||||
permission: List[str]
|
||||
metadata: Dict[str, Any]
|
||||
|
||||
@property
|
||||
def model_type(self) -> str:
|
||||
return self.metadata["engine_config"]["model_type"]
|
||||
|
||||
|
||||
class ModelList(BaseModel):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
data: List[ModelCard]
|
||||
object: str = "list"
|
||||
|
||||
|
||||
def to_model_metadata(
|
||||
model_id: str,
|
||||
model_config: "LLMConfig",
|
||||
overrides: Optional[Dict[str, Any]] = None,
|
||||
) -> ModelCard:
|
||||
"""Creates an OpenAI-compatible ModelData object.
|
||||
|
||||
Args:
|
||||
model_id: The ID of the model. Should contain the suffix if the model
|
||||
is LoRA fine-tuned. For example:
|
||||
meta-llama/Llama-2-7b-chat-hf:my_suffix:aBc1234
|
||||
model_config: The model's YAML config.
|
||||
overrides: should only be set for LoRA fine-tuned models. The
|
||||
overrides of the fine-tuned model metadata.
|
||||
|
||||
Returns:
|
||||
A ModelCard object.
|
||||
"""
|
||||
metadata = {
|
||||
"model_id": model_config.model_id,
|
||||
"input_modality": model_config.input_modality,
|
||||
"max_request_context_length": model_config.max_request_context_length,
|
||||
}
|
||||
|
||||
if overrides:
|
||||
metadata.update(overrides)
|
||||
|
||||
return ModelCard(
|
||||
id=model_id,
|
||||
object="model",
|
||||
owned_by="organization-owner",
|
||||
permission=[],
|
||||
metadata=metadata,
|
||||
)
|
||||
@@ -0,0 +1,333 @@
|
||||
import abc
|
||||
from typing import TYPE_CHECKING, Any, AsyncGenerator, Optional, Union
|
||||
|
||||
from ray.llm._internal.serve.core.configs.llm_config import (
|
||||
DiskMultiplexConfig,
|
||||
LLMConfig,
|
||||
)
|
||||
from ray.llm._internal.serve.core.protocol import RawRequestInfo
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.llm._internal.serve.core.configs.openai_api_models import (
|
||||
ChatCompletionRequest,
|
||||
ChatCompletionResponse,
|
||||
CompletionRequest,
|
||||
CompletionResponse,
|
||||
DetokenizeRequest,
|
||||
DetokenizeResponse,
|
||||
EmbeddingRequest,
|
||||
EmbeddingResponse,
|
||||
ErrorResponse,
|
||||
TokenizeRequest,
|
||||
TokenizeResponse,
|
||||
TranscriptionRequest,
|
||||
TranscriptionResponse,
|
||||
)
|
||||
|
||||
|
||||
class LLMEngine(abc.ABC):
|
||||
"""Base protocol class for all LLM engines."""
|
||||
|
||||
@abc.abstractmethod
|
||||
def __init__(self, llm_config: LLMConfig):
|
||||
"""Initialize the engine with the llm config"""
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
async def start(self):
|
||||
"""Start the engine"""
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
def routing_stats(self) -> dict:
|
||||
"""Replica routing stats surfaced to Serve's request router via
|
||||
``record_routing_stats`` (e.g. the KV-events endpoint for KV-aware
|
||||
routing)."""
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
async def resolve_lora(self, lora_model: DiskMultiplexConfig):
|
||||
"""Mounts the LoRA model on the engine, given the local disk path."""
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
async def chat(
|
||||
self,
|
||||
request: "ChatCompletionRequest",
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[Union[str, "ChatCompletionResponse", "ErrorResponse"], None]:
|
||||
"""Run a ChatCompletion with the engine.
|
||||
|
||||
To implement this method, you need to take a openAI compatible chat request, internally cast it to the target engine request type, and then call the engine's chat method.
|
||||
|
||||
This method is an async generator, so it yields chunks of response and when it is done, it returns None. We have the following convention:
|
||||
|
||||
- In case of streaming, yield a string representing data: <json_str>\n\n for each chunk. This should be already openAI compatible, so the higher level can just yield it to the client.
|
||||
- In case of non-streaming, yield a single object of type ChatCompletionResponse.
|
||||
- In case of error, yield a single object of type ErrorResponse.
|
||||
|
||||
Args:
|
||||
request: The chat completion request.
|
||||
raw_request_info: Optional RawRequestInfo containing data from the original
|
||||
HTTP request.
|
||||
|
||||
Yields:
|
||||
Union[str, ChatCompletionResponse, ErrorResponse]: A string representing a chunk of the response, a ChatCompletionResponse object, or an ErrorResponse object.
|
||||
|
||||
Returns:
|
||||
None when the generator is done.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
async def completions(
|
||||
self,
|
||||
request: "CompletionRequest",
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[Union[str, "CompletionResponse", "ErrorResponse"], None]:
|
||||
"""Run a Completion with the engine.
|
||||
|
||||
Similar to chat, this method is an async generator, so it yields chunks
|
||||
of response and when it is done, it returns None. We have the following
|
||||
convention:
|
||||
|
||||
* In case of streaming, yield a string representing data:
|
||||
<json_str>\n\n for each chunk. This should be already openAI compatible
|
||||
with completion response format, so the higher level can just yield it
|
||||
directly to the client.
|
||||
* In case of non-streaming, yield a single object of type
|
||||
CompletionResponse.
|
||||
* In case of error, yield a single object of type ErrorResponse.
|
||||
|
||||
Args:
|
||||
request: The completion request.
|
||||
raw_request_info: Optional RawRequestInfo containing data from the original
|
||||
HTTP request.
|
||||
|
||||
Yields:
|
||||
Union[str, CompletionResponse, ErrorResponse]: A string
|
||||
representing a chunk of the response, a CompletionResponse object,
|
||||
or an ErrorResponse object.
|
||||
|
||||
Returns:
|
||||
None when the generator is done.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
async def embeddings(
|
||||
self,
|
||||
request: "EmbeddingRequest",
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[Union["EmbeddingResponse", "ErrorResponse"], None]:
|
||||
"""Run an Embedding with the engine.
|
||||
|
||||
This method is different from chat and completion in that it does not
|
||||
have streaming, but still it is an async generator that yields response
|
||||
objects and when it is done, it returns None. We have the following
|
||||
convention:
|
||||
|
||||
* yield a single object of type EmbeddingResponse.
|
||||
* For errors, yield a single object of type ErrorResponse.
|
||||
|
||||
Args:
|
||||
request: The embedding request.
|
||||
raw_request_info: Optional RawRequestInfo containing data from the original
|
||||
HTTP request.
|
||||
|
||||
Returns:
|
||||
An async generator that yields EmbeddingResponse objects or ErrorResponse objects, and returns None when the generator is done.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
async def transcriptions(
|
||||
self,
|
||||
request: "TranscriptionRequest",
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[Union[str, "TranscriptionResponse", "ErrorResponse"], None]:
|
||||
"""Run a Transcription with the engine.
|
||||
|
||||
Similar to chat and completion, this method is an async generator,
|
||||
so it yields chunks of response and when it is done, it returns None.
|
||||
We have the following convention:
|
||||
|
||||
* In case of streaming, yield a string representing data:
|
||||
<json_str>\n\n for each chunk. This should be already openAI compatible,
|
||||
so the higher level can just yield it to the client.
|
||||
* In case of non-streaming, yield a single object of type TranscriptionResponse.
|
||||
* In case of error, yield a single object of type ErrorResponse.
|
||||
|
||||
Args:
|
||||
request: The transcription request.
|
||||
raw_request_info: Optional RawRequestInfo containing data from the original
|
||||
HTTP request.
|
||||
|
||||
Yields:
|
||||
Union[str, TranscriptionResponse, ErrorResponse]: A string
|
||||
representing a chunk of the response, a TranscriptionResponse object,
|
||||
or an ErrorResponse object.
|
||||
|
||||
Returns:
|
||||
None when the generator is done.
|
||||
"""
|
||||
pass
|
||||
|
||||
async def tokenize(
|
||||
self,
|
||||
request: "TokenizeRequest",
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[Union["TokenizeResponse", "ErrorResponse"], None]:
|
||||
"""Tokenize the input text.
|
||||
|
||||
This method tokenizes the input prompt or chat messages and returns
|
||||
the token IDs and optionally token strings.
|
||||
|
||||
Args:
|
||||
request: The tokenize request containing the text to tokenize.
|
||||
raw_request_info: Optional RawRequestInfo containing data from the original
|
||||
HTTP request.
|
||||
|
||||
Yields:
|
||||
Union[TokenizeResponse, ErrorResponse]: A TokenizeResponse object
|
||||
containing the tokens, or an ErrorResponse object.
|
||||
|
||||
Returns:
|
||||
None when the generator is done.
|
||||
"""
|
||||
yield # type: ignore
|
||||
|
||||
async def detokenize(
|
||||
self,
|
||||
request: "DetokenizeRequest",
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[Union["DetokenizeResponse", "ErrorResponse"], None]:
|
||||
"""Detokenize the input token IDs.
|
||||
|
||||
This method converts token IDs back into text.
|
||||
|
||||
Args:
|
||||
request: The detokenize request containing the token IDs.
|
||||
raw_request_info: Optional RawRequestInfo containing data from the original
|
||||
HTTP request.
|
||||
|
||||
Yields:
|
||||
Union[DetokenizeResponse, ErrorResponse]: A DetokenizeResponse object
|
||||
containing the text, or an ErrorResponse object.
|
||||
|
||||
Returns:
|
||||
None when the generator is done.
|
||||
"""
|
||||
yield # type: ignore
|
||||
|
||||
async def check_health(self) -> None:
|
||||
"""Check the health of the engine.
|
||||
|
||||
Does not return anything. Raise error when the engine is dead and needs
|
||||
to be restarted.
|
||||
"""
|
||||
return
|
||||
|
||||
async def build_asgi_app(self) -> Any:
|
||||
"""Build an ASGI app that serves directly from this engine's frontend.
|
||||
|
||||
Used by direct streaming, which serves traffic from the LLMServer
|
||||
replica's own ASGI ingress instead of the OpenAiIngress deployment.
|
||||
Engines that do not support direct serving should keep the default,
|
||||
which raises NotImplementedError.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
f"{type(self).__name__} does not support direct ASGI serving."
|
||||
)
|
||||
|
||||
##############################################################
|
||||
# Optional methods
|
||||
# These methods will be implemented in the future to allow
|
||||
# more granular life-cycle management of the engine.
|
||||
# e.g. in usecases like RL training, we need to put the engine
|
||||
# to sleep during training and wake up during rollouts.
|
||||
##############################################################
|
||||
|
||||
@abc.abstractmethod
|
||||
async def reset_prefix_cache(self) -> None:
|
||||
"""Reset the prefix cache of the underlying engine"""
|
||||
|
||||
async def sleep(self, **kwargs: Any) -> None:
|
||||
"""Put the engine to sleep.
|
||||
|
||||
The caller should guarantee that no requests are being processed
|
||||
during the sleep period, before `wakeup` is called.
|
||||
|
||||
Args:
|
||||
**kwargs: Engine-specific sleep options. See the concrete engine
|
||||
implementation for available options.
|
||||
"""
|
||||
pass
|
||||
|
||||
async def wakeup(self, **kwargs: Any) -> None:
|
||||
"""Wake up the engine from sleep mode.
|
||||
|
||||
Args:
|
||||
**kwargs: Engine-specific wakeup options. See the concrete engine
|
||||
implementation for available options.
|
||||
"""
|
||||
pass
|
||||
|
||||
async def is_sleeping(self) -> bool:
|
||||
"""Check whether the engine is currently sleeping.
|
||||
|
||||
Returns:
|
||||
True if the engine is sleeping, False otherwise.
|
||||
"""
|
||||
return False
|
||||
|
||||
async def collective_rpc(
|
||||
self,
|
||||
method: str,
|
||||
timeout: Optional[float] = None,
|
||||
args: tuple = (),
|
||||
kwargs: Optional[dict] = None,
|
||||
) -> list:
|
||||
"""Execute a collective RPC call on all workers.
|
||||
|
||||
This is used for RLHF workflows where a trainer needs to execute
|
||||
methods on all TP/PP workers (e.g., for weight synchronization).
|
||||
|
||||
Args:
|
||||
method: Name of the worker method to execute.
|
||||
timeout: Maximum time in seconds to wait for execution.
|
||||
args: Positional arguments to pass to the worker method.
|
||||
kwargs: Keyword arguments to pass to the worker method.
|
||||
|
||||
Returns:
|
||||
A list containing the results from each worker.
|
||||
"""
|
||||
raise NotImplementedError("collective_rpc is not implemented for this engine")
|
||||
|
||||
async def pause(self, **kwargs: Any) -> None:
|
||||
"""Pause the engine.
|
||||
|
||||
Args:
|
||||
**kwargs: Engine-specific pause options. Passed through to the engine.
|
||||
"""
|
||||
pass
|
||||
|
||||
async def resume(self, **kwargs: Any) -> None:
|
||||
"""Resume the engine.
|
||||
|
||||
Args:
|
||||
**kwargs: Engine-specific resume options. Passed through to the engine.
|
||||
"""
|
||||
pass
|
||||
|
||||
async def is_paused(self) -> bool:
|
||||
"""Check whether the engine is currently paused.
|
||||
|
||||
Returns:
|
||||
True if the engine is paused, False otherwise.
|
||||
"""
|
||||
return False
|
||||
|
||||
def shutdown(self) -> None:
|
||||
"""Shuts down the engine"""
|
||||
pass
|
||||
@@ -0,0 +1,281 @@
|
||||
import os
|
||||
import pprint
|
||||
from typing import Any, Dict, List, Optional, Type, 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.common.utils.import_utils import load_class
|
||||
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.ingress import (
|
||||
OpenAiIngress,
|
||||
make_fastapi_ingress,
|
||||
)
|
||||
from ray.llm._internal.serve.core.server.builder import (
|
||||
build_llm_deployment,
|
||||
)
|
||||
from ray.llm._internal.serve.core.server.llm_server import LLMServer
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
from ray.llm._internal.serve.routing_policies.kv_aware.kv_aware_router import (
|
||||
is_kv_aware,
|
||||
)
|
||||
from ray.serve.config import RequestRouterConfig
|
||||
from ray.serve.deployment import Application
|
||||
from ray.serve.experimental.round_robin_router import RoundRobinRouter
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def _get_direct_streaming_serve_options(
|
||||
llm_config: LLMConfig,
|
||||
override_serve_options: Optional[dict] = None,
|
||||
) -> dict:
|
||||
override_serve_options = dict(override_serve_options or {})
|
||||
if (
|
||||
"request_router_config" not in llm_config.deployment_config
|
||||
and "request_router_config" not in override_serve_options
|
||||
):
|
||||
override_serve_options["request_router_config"] = RequestRouterConfig(
|
||||
request_router_class=RoundRobinRouter,
|
||||
)
|
||||
return override_serve_options
|
||||
|
||||
|
||||
def _build_direct_streaming_llm_deployment(
|
||||
llm_config: LLMConfig,
|
||||
*,
|
||||
name_prefix: Optional[str] = None,
|
||||
bind_kwargs: Optional[dict] = None,
|
||||
override_serve_options: Optional[dict] = None,
|
||||
deployment_cls: Optional[Type[LLMServer]] = None,
|
||||
) -> Application:
|
||||
"""Build an LLM deployment with late-bound ASGI ingress enabled.
|
||||
|
||||
Used by the OpenAI, DP, and PD builders to wrap their respective server
|
||||
class (``LLMServer``, ``DPServer``, ``PDDecodeServer``/``DPPDDecodeServer``)
|
||||
as the ingress. The real ASGI app (vLLM FastAPI) is constructed inside
|
||||
``LLMServer.__serve_build_asgi_app__`` after the engine starts; subclasses
|
||||
inherit this hook.
|
||||
|
||||
Replica selection is driven by the deployment's ``request_router_config``.
|
||||
Default to ``RoundRobinRouter`` when the user hasn't set one, and otherwise
|
||||
leave their configured value untouched.
|
||||
"""
|
||||
server_cls = deployment_cls or llm_config.server_cls or LLMServer
|
||||
return build_llm_deployment(
|
||||
llm_config,
|
||||
name_prefix=name_prefix,
|
||||
bind_kwargs=bind_kwargs,
|
||||
deployment_cls=serve.ingress()(server_cls),
|
||||
override_serve_options=_get_direct_streaming_serve_options(
|
||||
llm_config, override_serve_options
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _build_openai_ingress_request_router(
|
||||
*, server: Application, llm_config: LLMConfig
|
||||
) -> Application:
|
||||
"""Build the ingress request router peer for OpenAI compatible LLM apps.
|
||||
|
||||
The returned Application is attached to the ingress application with
|
||||
``Application._with_ingress_request_router``.
|
||||
|
||||
``num_replicas`` is pinned to 1 because HAProxy's ingress request router
|
||||
backend currently expects a single endpoint. TODO(eicherseiji): expose
|
||||
these as a user-overridable IngressRequestRouterConfig once HAProxy
|
||||
supports multiple router replicas.
|
||||
|
||||
Pre-routing tokenization is wired on only when ``llm_config`` configures a
|
||||
KVAwareRouter, the sole policy that scores replicas on prompt token IDs.
|
||||
"""
|
||||
from ray.llm._internal.serve.core.ingress.router import LLMRouter
|
||||
|
||||
deployment = serve.deployment(
|
||||
LLMRouter,
|
||||
num_replicas=1,
|
||||
max_ongoing_requests=1000,
|
||||
)
|
||||
return deployment.bind(
|
||||
server=server,
|
||||
pre_routing_tokenization=is_kv_aware(llm_config),
|
||||
)
|
||||
|
||||
|
||||
class IngressClsConfig(BaseModelExtended):
|
||||
ingress_cls: Union[str, Type[OpenAiIngress]] = Field(
|
||||
default=OpenAiIngress,
|
||||
description="The class name of the ingress to use. It can be in form of `module_name.class_name` or `module_name:class_name` or the class itself. The class constructor should take the following arguments: `(llm_deployments: Dict[str, DeploymentHandle], model_cards: Dict[str, ModelCard], lora_paths: Optional[Dict[str, str]] = None, **extra_kwargs)` where the dicts are keyed by base model ID.",
|
||||
)
|
||||
|
||||
ingress_extra_kwargs: Optional[dict] = Field(
|
||||
default_factory=dict,
|
||||
description="""The kwargs to bind to the ingress deployment. This will be passed to the ingress class constructor.""",
|
||||
)
|
||||
|
||||
@field_validator("ingress_cls")
|
||||
@classmethod
|
||||
def validate_class(
|
||||
cls, value: Union[str, Type[OpenAiIngress]]
|
||||
) -> Type[OpenAiIngress]:
|
||||
if isinstance(value, str):
|
||||
return load_class(value)
|
||||
return value
|
||||
|
||||
|
||||
class LLMServingArgs(BaseModelExtended):
|
||||
llm_configs: List[Union[str, dict, LLMConfig]] = Field(
|
||||
description="A list of LLMConfigs, or dicts representing LLMConfigs, or paths to yaml files defining LLMConfigs.",
|
||||
)
|
||||
ingress_cls_config: Union[dict, IngressClsConfig] = Field(
|
||||
default_factory=IngressClsConfig,
|
||||
description="The configuration for the ingress class. It can be a dict representing the ingress class configuration, or an IngressClsConfig object.",
|
||||
)
|
||||
ingress_deployment_config: Dict[str, Any] = Field(
|
||||
default_factory=dict,
|
||||
description="""
|
||||
The Ray @server.deployment options for the ingress server.
|
||||
""",
|
||||
)
|
||||
|
||||
@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
|
||||
|
||||
@field_validator("llm_configs")
|
||||
@classmethod
|
||||
def _validate_llm_configs(
|
||||
cls, value: List[Union[str, dict, LLMConfig]]
|
||||
) -> List[LLMConfig]:
|
||||
llm_configs = []
|
||||
for config in value:
|
||||
if isinstance(config, str):
|
||||
if not os.path.exists(config):
|
||||
raise ValueError(
|
||||
f"Could not load model config from {config}, as the file does not exist."
|
||||
)
|
||||
llm_configs.append(LLMConfig.from_file(config))
|
||||
elif isinstance(config, dict):
|
||||
llm_configs.append(LLMConfig.model_validate(config))
|
||||
elif isinstance(config, LLMConfig):
|
||||
llm_configs.append(config)
|
||||
else:
|
||||
raise TypeError(f"Invalid LLMConfig type: {type(config)}")
|
||||
return llm_configs
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_model_ids(self):
|
||||
"""Validate that model IDs are unique and at least one model is configured."""
|
||||
if len({m.model_id for m in self.llm_configs}) != len(self.llm_configs):
|
||||
raise ValueError("Duplicate models found. Make sure model ids are unique.")
|
||||
|
||||
if len(self.llm_configs) == 0:
|
||||
raise ValueError(
|
||||
"List of models is empty. Maybe some parameters cannot be parsed into the LLMConfig config."
|
||||
)
|
||||
return self
|
||||
|
||||
|
||||
def _validate_direct_streaming_ingress_config(
|
||||
ingress_deployment_config: Optional[dict],
|
||||
ingress_cls_config: IngressClsConfig,
|
||||
) -> None:
|
||||
if ingress_deployment_config:
|
||||
raise ValueError(
|
||||
"RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING does not support "
|
||||
"ingress_deployment_config because the LLM server class is used "
|
||||
"directly as the ingress deployment. Configure the server through "
|
||||
"each LLMConfig.deployment_config instead."
|
||||
)
|
||||
|
||||
if (
|
||||
ingress_cls_config.ingress_cls != OpenAiIngress
|
||||
or ingress_cls_config.ingress_extra_kwargs
|
||||
):
|
||||
raise ValueError(
|
||||
"RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING does not support "
|
||||
"ingress_cls_config because the LLM server class is used directly "
|
||||
"as the ingress deployment."
|
||||
)
|
||||
|
||||
|
||||
def build_openai_app(builder_config: dict) -> Application:
|
||||
"""Build an OpenAI compatible app with the llm deployment setup from
|
||||
the given builder configuration.
|
||||
|
||||
Args:
|
||||
builder_config: The configuration for the builder. It has to conform
|
||||
to the LLMServingArgs pydantic model.
|
||||
|
||||
Returns:
|
||||
The configured Ray Serve Application router.
|
||||
"""
|
||||
|
||||
builder_config = LLMServingArgs.model_validate(builder_config)
|
||||
llm_configs = builder_config.llm_configs
|
||||
|
||||
# Direct streaming attaches LLMRouter as the ingress request router and
|
||||
# uses the LLMServer deployment itself as the ingress app, so it returns
|
||||
# before the regular OpenAiIngress wiring.
|
||||
if RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING:
|
||||
if len(llm_configs) > 1:
|
||||
raise ValueError(
|
||||
"RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING currently supports exactly "
|
||||
"one LLM config. Multi-model direct streaming requires composing "
|
||||
"multiple LLMServer deployments into the main application graph, "
|
||||
"which is not supported yet."
|
||||
)
|
||||
_validate_direct_streaming_ingress_config(
|
||||
builder_config.ingress_deployment_config,
|
||||
builder_config.ingress_cls_config,
|
||||
)
|
||||
direct_deployment = _build_direct_streaming_llm_deployment(llm_configs[0])
|
||||
logger.info(
|
||||
"Direct streaming enabled: "
|
||||
"LLMServer=ingress, LLMRouter=ingress_request_router"
|
||||
)
|
||||
return direct_deployment._with_ingress_request_router(
|
||||
_build_openai_ingress_request_router(
|
||||
server=direct_deployment, llm_config=llm_configs[0]
|
||||
)
|
||||
)
|
||||
|
||||
llm_deployments = {c.model_id: build_llm_deployment(c) for c in llm_configs}
|
||||
model_cards = {c.model_id: to_model_metadata(c.model_id, c) for c in llm_configs}
|
||||
lora_paths = {
|
||||
c.model_id: c.lora_config.dynamic_lora_loading_path
|
||||
for c in llm_configs
|
||||
if c.lora_config is not None
|
||||
}
|
||||
|
||||
ingress_cls_config = builder_config.ingress_cls_config
|
||||
default_ingress_options = ingress_cls_config.ingress_cls.get_deployment_options(
|
||||
llm_configs
|
||||
)
|
||||
|
||||
ingress_options = maybe_apply_llm_deployment_config_defaults(
|
||||
default_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))
|
||||
|
||||
return serve.deployment(ingress_cls, **ingress_options).bind(
|
||||
llm_deployments=llm_deployments,
|
||||
model_cards=model_cards,
|
||||
lora_paths=lora_paths,
|
||||
**ingress_cls_config.ingress_extra_kwargs,
|
||||
)
|
||||
@@ -0,0 +1,138 @@
|
||||
"""Development/RL-focused ingress with control plane endpoints.
|
||||
|
||||
This module provides DevIngress, an extension of OpenAiIngress that adds
|
||||
control plane endpoints for managing engine lifecycle. These endpoints
|
||||
are useful for RL training workflows where engines need to be put to sleep
|
||||
during training and woken up for inference.
|
||||
|
||||
Endpoints:
|
||||
POST /sleep: Put engine to sleep (frees GPU memory)
|
||||
POST /wakeup: Wake up engine from sleep
|
||||
GET /is_sleeping: Check if engine is sleeping
|
||||
POST /pause: Pause generation (keeps weights in GPU)
|
||||
POST /resume: Resume generation after pause
|
||||
GET /is_paused: Check if engine is paused
|
||||
POST /reset_prefix_cache: Reset the KV prefix cache
|
||||
POST /collective_rpc: Execute collective RPC on all workers
|
||||
"""
|
||||
|
||||
import pprint
|
||||
from typing import Dict
|
||||
|
||||
from ray import serve
|
||||
from ray.llm._internal.common.dict_utils import (
|
||||
maybe_apply_llm_deployment_config_defaults,
|
||||
)
|
||||
from ray.llm._internal.serve.core.configs.openai_api_models import to_model_metadata
|
||||
from ray.llm._internal.serve.core.ingress.builder import LLMServingArgs
|
||||
from ray.llm._internal.serve.core.ingress.ingress import (
|
||||
DEFAULT_ENDPOINTS,
|
||||
OpenAiIngress,
|
||||
make_fastapi_ingress,
|
||||
)
|
||||
from ray.llm._internal.serve.core.ingress.mixins import (
|
||||
CacheManagerIngressMixin,
|
||||
CollectiveRpcIngressMixin,
|
||||
PausableIngressMixin,
|
||||
SleepableIngressMixin,
|
||||
)
|
||||
from ray.llm._internal.serve.core.server.builder import build_llm_deployment
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
from ray.serve.deployment import Application
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
# Endpoint map for DevIngress - includes all default endpoints plus control plane
|
||||
DEV_ENDPOINTS = {
|
||||
**CacheManagerIngressMixin.ENDPOINTS,
|
||||
**CollectiveRpcIngressMixin.ENDPOINTS,
|
||||
**PausableIngressMixin.ENDPOINTS,
|
||||
**SleepableIngressMixin.ENDPOINTS,
|
||||
**DEFAULT_ENDPOINTS,
|
||||
}
|
||||
|
||||
|
||||
class DevIngress(
|
||||
OpenAiIngress,
|
||||
SleepableIngressMixin,
|
||||
PausableIngressMixin,
|
||||
CacheManagerIngressMixin,
|
||||
CollectiveRpcIngressMixin,
|
||||
):
|
||||
"""OpenAI-compatible ingress with additional control plane endpoints.
|
||||
|
||||
This ingress extends the standard OpenAI endpoints with control plane
|
||||
operations for managing engine lifecycle. These are useful for:
|
||||
- RL training: Put engines to sleep during training, wake up for rollouts
|
||||
- Memory management: Free GPU memory between inference workloads
|
||||
- Benchmarking: Reset prefix cache between benchmark rounds
|
||||
- RLHF: Execute collective RPC on all workers for weight updates
|
||||
|
||||
Control plane endpoints provided by mixins:
|
||||
- SleepableIngressMixin: /sleep, /wakeup, /is_sleeping
|
||||
- PausableIngressMixin: /pause, /resume, /is_paused
|
||||
- CacheManagerIngressMixin: /reset_prefix_cache
|
||||
- CollectiveRpcIngressMixin: /collective_rpc
|
||||
|
||||
WARNING: These endpoints are intended for development and trusted
|
||||
environments. Consider access control in production deployments.
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
def build_dev_openai_app(builder_config: Dict) -> Application:
|
||||
"""Build an OpenAI compatible app with dev/control plane endpoints.
|
||||
|
||||
This is similar to build_openai_app but uses DevIngress with
|
||||
additional control plane endpoints:
|
||||
- /sleep, /wakeup, /is_sleeping (sleep mode - offloads weights to CPU)
|
||||
- /pause, /resume, /is_paused (pause mode - keeps weights in GPU)
|
||||
- /reset_prefix_cache (cache management)
|
||||
- /collective_rpc (RLHF - execute RPC on all workers)
|
||||
|
||||
Args:
|
||||
builder_config: Configuration conforming to LLMServingArgs.
|
||||
See LLMServingArgs for details on the expected structure.
|
||||
|
||||
Returns:
|
||||
The configured Ray Serve Application.
|
||||
|
||||
Example:
|
||||
config = {
|
||||
"llm_configs": [llm_config],
|
||||
"ingress_deployment_config": {}
|
||||
}
|
||||
app = build_dev_openai_app(config)
|
||||
serve.run(app)
|
||||
"""
|
||||
config = LLMServingArgs.model_validate(builder_config)
|
||||
llm_configs = config.llm_configs
|
||||
|
||||
llm_deployments = {c.model_id: build_llm_deployment(c) for c in llm_configs}
|
||||
model_cards = {c.model_id: to_model_metadata(c.model_id, c) for c in llm_configs}
|
||||
lora_paths = {
|
||||
c.model_id: c.lora_config.dynamic_lora_loading_path
|
||||
for c in llm_configs
|
||||
if c.lora_config is not None
|
||||
}
|
||||
|
||||
ingress_cls_config = config.ingress_cls_config
|
||||
default_ingress_options = DevIngress.get_deployment_options(llm_configs)
|
||||
|
||||
ingress_options = maybe_apply_llm_deployment_config_defaults(
|
||||
default_ingress_options, config.ingress_deployment_config
|
||||
)
|
||||
|
||||
ingress_cls = make_fastapi_ingress(DevIngress, endpoint_map=DEV_ENDPOINTS)
|
||||
|
||||
logger.info("============== Ingress Options ==============")
|
||||
logger.info(pprint.pformat(ingress_options))
|
||||
|
||||
return serve.deployment(ingress_cls, **ingress_options).bind(
|
||||
llm_deployments=llm_deployments,
|
||||
model_cards=model_cards,
|
||||
lora_paths=lora_paths,
|
||||
**ingress_cls_config.ingress_extra_kwargs,
|
||||
)
|
||||
@@ -0,0 +1,694 @@
|
||||
import asyncio
|
||||
import copy
|
||||
import sys
|
||||
from contextlib import asynccontextmanager
|
||||
from enum import Enum
|
||||
from typing import (
|
||||
Annotated,
|
||||
Any,
|
||||
AsyncGenerator,
|
||||
Awaitable,
|
||||
Callable,
|
||||
Dict,
|
||||
List,
|
||||
Optional,
|
||||
Type,
|
||||
Union,
|
||||
)
|
||||
|
||||
from fastapi import FastAPI, Form, HTTPException, Request, status
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from starlette.responses import JSONResponse, Response, StreamingResponse
|
||||
|
||||
from ray import serve
|
||||
from ray.llm._internal.common.utils.lora_utils import (
|
||||
get_base_model_id,
|
||||
get_lora_model_ids,
|
||||
)
|
||||
from ray.llm._internal.serve.constants import (
|
||||
DEFAULT_LLM_ROUTER_HTTP_TIMEOUT,
|
||||
DEFAULT_MAX_ONGOING_REQUESTS,
|
||||
DEFAULT_MAX_TARGET_ONGOING_REQUESTS,
|
||||
)
|
||||
from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
|
||||
from ray.llm._internal.serve.core.configs.openai_api_models import (
|
||||
ChatCompletionRequest,
|
||||
CompletionRequest,
|
||||
DetokenizeRequest,
|
||||
DetokenizeResponse,
|
||||
EmbeddingRequest,
|
||||
EmbeddingResponse,
|
||||
ErrorResponse,
|
||||
LLMChatResponse,
|
||||
LLMCompletionsResponse,
|
||||
LLMEmbeddingsResponse,
|
||||
LLMScoreResponse,
|
||||
LLMTranscriptionResponse,
|
||||
ModelCard,
|
||||
ModelList,
|
||||
OpenAIHTTPException,
|
||||
ScoreRequest,
|
||||
ScoreResponse,
|
||||
TokenizeCompletionRequest,
|
||||
TokenizeResponse,
|
||||
TranscriptionRequest,
|
||||
)
|
||||
from ray.llm._internal.serve.core.ingress.middleware import (
|
||||
SetRequestIdMiddleware,
|
||||
add_exception_handling_middleware,
|
||||
)
|
||||
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 DeploymentProtocol, RawRequestInfo
|
||||
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,
|
||||
metrics_lifespan,
|
||||
)
|
||||
from ray.llm._internal.serve.utils.lora_serve_utils import (
|
||||
get_lora_model_metadata,
|
||||
)
|
||||
from ray.llm._internal.serve.utils.server_utils import replace_prefix
|
||||
from ray.serve._private.http_util import session_id_from_headers
|
||||
from ray.serve.handle import DeploymentHandle
|
||||
|
||||
# Import asyncio timeout depends on python version
|
||||
if sys.version_info >= (3, 11):
|
||||
from asyncio import timeout
|
||||
else:
|
||||
from async_timeout import timeout
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
DEFAULT_INGRESS_OPTIONS = {
|
||||
"max_ongoing_requests": DEFAULT_MAX_ONGOING_REQUESTS,
|
||||
"autoscaling_config": {
|
||||
"target_ongoing_requests": DEFAULT_MAX_TARGET_ONGOING_REQUESTS,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _get_min_replicas_from_llm_config(config: LLMConfig) -> Optional[int]:
|
||||
autoscaling_config = config.deployment_config.get("autoscaling_config")
|
||||
if autoscaling_config is None:
|
||||
return None
|
||||
if isinstance(autoscaling_config, dict):
|
||||
return autoscaling_config.get("min_replicas")
|
||||
return getattr(autoscaling_config, "min_replicas", None)
|
||||
|
||||
|
||||
def _all_models_scale_to_zero(llm_configs: Optional[List[LLMConfig]]) -> bool:
|
||||
"""Check if all models are configured with min_replicas == 0."""
|
||||
if not llm_configs:
|
||||
return False
|
||||
return all(_get_min_replicas_from_llm_config(config) == 0 for config in llm_configs)
|
||||
|
||||
|
||||
# These methods correspond to functions defined in the LLMEngine class in python/ray/llm/_internal/serve/deployments/llm/llm_engine.py
|
||||
class CallMethod(Enum):
|
||||
CHAT = "chat"
|
||||
COMPLETIONS = "completions"
|
||||
TRANSCRIPTIONS = "transcriptions"
|
||||
|
||||
|
||||
DEFAULT_ENDPOINTS = {
|
||||
"models": lambda app: app.get("/v1/models", response_model=ModelList),
|
||||
"model_data": lambda app: app.get(
|
||||
"/v1/models/{model:path}", response_model=ModelCard
|
||||
),
|
||||
"completions": lambda app: app.post("/v1/completions"),
|
||||
"chat": lambda app: app.post("/v1/chat/completions"),
|
||||
"embeddings": lambda app: app.post("/v1/embeddings"),
|
||||
"transcriptions": lambda app: app.post(
|
||||
"/v1/audio/transcriptions",
|
||||
),
|
||||
"score": lambda app: app.post("/v1/score"),
|
||||
"tokenize": lambda app: app.post("/tokenize"),
|
||||
"detokenize": lambda app: app.post("/detokenize"),
|
||||
}
|
||||
|
||||
|
||||
def init() -> FastAPI:
|
||||
_fastapi_router_app = FastAPI(lifespan=metrics_lifespan)
|
||||
|
||||
# NOTE: PLEASE READ CAREFULLY BEFORE MODIFYING
|
||||
#
|
||||
# FastAPI middleware is executed in LIFO (last-in, first-out) order,
|
||||
# hence maintaining current ordering is crucial as some of the middleware
|
||||
# might have data dependency on the other: for ex, telemetry middleware
|
||||
# depends on middleware generating request-id
|
||||
#
|
||||
# Add exception handling middleware
|
||||
# NOTE: This middleware should be added first such that it's intercepting
|
||||
# exceptions from the handlers, avoiding them propagating to other
|
||||
# middleware (for ex, telemetry)
|
||||
add_exception_handling_middleware(_fastapi_router_app)
|
||||
# Configure CORS middleware
|
||||
_fastapi_router_app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
# Add HTTP metrics middleware
|
||||
add_http_metrics_middleware(_fastapi_router_app)
|
||||
|
||||
# Inject unique per-request ID
|
||||
#
|
||||
# NOTE: This middleware should be executed among the last (since
|
||||
# middleware is executed in LIFO).
|
||||
_fastapi_router_app.add_middleware(SetRequestIdMiddleware)
|
||||
|
||||
return _fastapi_router_app
|
||||
|
||||
|
||||
def make_fastapi_ingress(
|
||||
cls: Type,
|
||||
*,
|
||||
endpoint_map: Optional[Dict[str, Callable[[FastAPI], Callable]]] = None,
|
||||
app: Optional[FastAPI] = None,
|
||||
):
|
||||
"""
|
||||
Create a Ray Serve ingress deployment from a class and endpoint mapping.
|
||||
|
||||
Args:
|
||||
cls: The class to convert into an ingress deployment
|
||||
endpoint_map: Dictionary mapping method names to FastAPI route
|
||||
decorators. Each value is a lambda that takes a FastAPI app and
|
||||
returns a route decorator.
|
||||
app: Optional FastAPI app to use for the ingress deployment. If not
|
||||
provided, a new FastAPI app will be created.
|
||||
|
||||
Returns:
|
||||
A class decorated with @serve.ingress
|
||||
|
||||
Example:
|
||||
endpoint_map = {
|
||||
"increment": lambda app: app.post("/increment"),
|
||||
"get_counter": lambda app: app.get("/counter"),
|
||||
}
|
||||
|
||||
# With additional FastAPI parameters:
|
||||
endpoint_map = {
|
||||
"increment": lambda app: app.post("/increment", status_code=201, tags=["counter"]),
|
||||
"get_counter": lambda app: app.get("/counter", response_model=CounterResponse),
|
||||
}
|
||||
"""
|
||||
|
||||
if app is None:
|
||||
app = init()
|
||||
|
||||
if endpoint_map is None:
|
||||
endpoint_map = DEFAULT_ENDPOINTS
|
||||
|
||||
# Create a new class that inherits from the original to avoid modifying it
|
||||
# in-place. We populate the new class's __dict__ with decorated methods.
|
||||
class_dict = {}
|
||||
|
||||
# Apply route decorators to the class methods and store them in class_dict
|
||||
for method_name, route_factory in endpoint_map.items():
|
||||
# Get the route decorator from the lambda
|
||||
route_decorator = route_factory(app)
|
||||
# Get the original method from the class
|
||||
original_method = getattr(cls, method_name)
|
||||
# Apply the decorator to the original method
|
||||
decorated_method = route_decorator(original_method)
|
||||
# Store in the class dict so it will be properly bound to new_cls
|
||||
class_dict[method_name] = decorated_method
|
||||
|
||||
# Create new class with the decorated methods in its __dict__.
|
||||
# We keep the same __name__ and __qualname__ as the original class
|
||||
# so that the new class properly represents the input class.
|
||||
new_cls = type(cls.__name__, (cls,), class_dict)
|
||||
new_cls.__qualname__ = cls.__qualname__
|
||||
|
||||
# Apply the serve.ingress decorator to the new class
|
||||
return serve.ingress(app)(new_cls)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def router_request_timeout(timeout_duration: float):
|
||||
try:
|
||||
async with timeout(timeout_duration):
|
||||
yield
|
||||
except asyncio.TimeoutError as e:
|
||||
raise OpenAIHTTPException(
|
||||
status_code=status.HTTP_408_REQUEST_TIMEOUT,
|
||||
message="Request server side timeout",
|
||||
internal_message=str(e),
|
||||
)
|
||||
|
||||
|
||||
class OpenAiIngress(DeploymentProtocol):
|
||||
def __init__(
|
||||
self,
|
||||
llm_deployments: Dict[str, DeploymentHandle],
|
||||
model_cards: Dict[str, ModelCard],
|
||||
*,
|
||||
lora_paths: Optional[Dict[str, str]] = None,
|
||||
_get_lora_model_metadata_func: Optional[
|
||||
Callable[[str, str], Awaitable[Dict[str, Any]]]
|
||||
] = None,
|
||||
):
|
||||
if set(llm_deployments) != set(model_cards):
|
||||
raise ValueError(
|
||||
"llm_deployments and model_cards must have the same model IDs. "
|
||||
f"Got llm_deployments={sorted(llm_deployments)}, "
|
||||
f"model_cards={sorted(model_cards)}."
|
||||
)
|
||||
|
||||
self._default_serve_handles: Dict[str, DeploymentHandle] = dict(llm_deployments)
|
||||
self._model_cards: Dict[str, ModelCard] = dict(model_cards)
|
||||
self._lora_paths: Dict[str, str] = dict(lora_paths or {})
|
||||
|
||||
# Configuring a ServeHandle with .options() creates a new ServeHandle
|
||||
# object, which contains a new metrics pusher and long-polling call.
|
||||
# Creating too many ServeHandles can impact event-loop and Serve Controller
|
||||
# performance, so we save configured ServeHandles here and reuse them.
|
||||
self._configured_serve_handles: Dict[str, DeploymentHandle] = {}
|
||||
self._get_lora_model_metadata_func = (
|
||||
_get_lora_model_metadata_func or self._default_get_lora_model_metadata_func
|
||||
)
|
||||
|
||||
async def _default_get_lora_model_metadata_func(
|
||||
self, model_id: str, base_path: str
|
||||
) -> Dict[str, Any]:
|
||||
return await get_lora_model_metadata(model_id, base_path)
|
||||
|
||||
async def check_health(self):
|
||||
pass
|
||||
|
||||
def _get_configured_serve_handle(self, model_id: str):
|
||||
"""Gets a ServeHandle to a model deployment.
|
||||
|
||||
Configures the handle's options, and stores it in a cache.
|
||||
|
||||
If the model_id includes LoRA suffix, we set the model ID as
|
||||
the multiplexed_model_id, so the request uses Serve's multiplexed
|
||||
routing logic.
|
||||
|
||||
If the model_id is a base model- even if the model has LoRA
|
||||
adapters- we don't set multiplexed_model_id. Setting
|
||||
multiplexed_model_id would cause base model requests to be
|
||||
sent to a single model replica, instead of being load
|
||||
balanced across all replicas. This is undesirable for base
|
||||
model requests (unlike LoRA requests) because all the replicas
|
||||
have a copy of the base model.
|
||||
"""
|
||||
|
||||
if model_id not in self._configured_serve_handles:
|
||||
base_model_id = get_base_model_id(model_id)
|
||||
if base_model_id in self._default_serve_handles:
|
||||
if model_id == base_model_id:
|
||||
default_handle = self._default_serve_handles[model_id]
|
||||
configured_handle = default_handle.options(stream=True)
|
||||
self._configured_serve_handles[model_id] = configured_handle
|
||||
else:
|
||||
default_handle = self._default_serve_handles[base_model_id]
|
||||
configured_handle = default_handle.options(
|
||||
stream=True,
|
||||
multiplexed_model_id=model_id,
|
||||
)
|
||||
self._configured_serve_handles[model_id] = configured_handle
|
||||
else:
|
||||
raise HTTPException(
|
||||
status.HTTP_404_NOT_FOUND,
|
||||
f'Could not find model with id "{model_id}".',
|
||||
)
|
||||
|
||||
return self._configured_serve_handles[model_id]
|
||||
|
||||
async def _get_model_id(self, model: Optional[str]) -> str:
|
||||
# Default to the only configured model if no model specified
|
||||
if model is None:
|
||||
if len(self._model_cards) == 1:
|
||||
model = next(iter(self._model_cards.keys()))
|
||||
else:
|
||||
raise HTTPException(
|
||||
status.HTTP_400_BAD_REQUEST,
|
||||
"Model parameter is required when multiple models are configured. "
|
||||
f"Available models: {list(self._model_cards.keys())}",
|
||||
)
|
||||
|
||||
base_model_id = get_base_model_id(model)
|
||||
if base_model_id not in self._model_cards:
|
||||
raise HTTPException(
|
||||
status.HTTP_404_NOT_FOUND,
|
||||
f'Got request for model "{model}". '
|
||||
f'Could not find base model with ID "{base_model_id}".',
|
||||
)
|
||||
|
||||
# Return original model ID so multiplexed routing works correctly.
|
||||
return model
|
||||
|
||||
async def _get_response(
|
||||
self,
|
||||
*,
|
||||
body: Union[
|
||||
CompletionRequest,
|
||||
ChatCompletionRequest,
|
||||
EmbeddingRequest,
|
||||
TranscriptionRequest,
|
||||
ScoreRequest,
|
||||
],
|
||||
call_method: str,
|
||||
raw_request: Optional[Request] = None,
|
||||
) -> AsyncGenerator[
|
||||
Union[
|
||||
LLMChatResponse,
|
||||
LLMCompletionsResponse,
|
||||
LLMEmbeddingsResponse,
|
||||
LLMTranscriptionResponse,
|
||||
LLMScoreResponse,
|
||||
],
|
||||
None,
|
||||
]:
|
||||
"""Calls the model deployment and returns the stream."""
|
||||
model_id = await self._get_model_id(body.model)
|
||||
model_handle = self._get_configured_serve_handle(model_id)
|
||||
|
||||
# Propagate the session id from the client request to the downstream
|
||||
# LLMServer handle. The Serve HTTP proxy attaches session_id to the
|
||||
# *ingress* deployment handle (proxy.py:_setup_request_context), but
|
||||
# that does NOT carry over to a second handle hop (here -> LLMServer).
|
||||
# Re-read the configured session header from the raw request and apply
|
||||
# it via .options(session_id=...) so session-aware request routers
|
||||
# (e.g. ConsistentHashRouter) on the LLMServer deployment see it.
|
||||
# Uses the same case-insensitive, separator-tolerant matcher as
|
||||
# proxy.py so a `-`/`_` rewrite by an intermediate proxy doesn't
|
||||
# silently drop session affinity on this second hop.
|
||||
if raw_request is not None:
|
||||
session_id = session_id_from_headers(raw_request.headers)
|
||||
if session_id:
|
||||
model_handle = model_handle.options(session_id=session_id)
|
||||
|
||||
# TODO(seiji): Remove when we update to Pydantic v2.11+ with the fix
|
||||
# for tool calling ValidatorIterator serialization issue.
|
||||
if isinstance(body, ChatCompletionRequest):
|
||||
body = _sanitize_chat_completion_request(body)
|
||||
|
||||
# Convert Starlette request to serializable RawRequestInfo
|
||||
raw_request_info: Optional[RawRequestInfo] = None
|
||||
if raw_request is not None:
|
||||
raw_request_info = RawRequestInfo.from_starlette_request(raw_request)
|
||||
|
||||
async for response in getattr(model_handle, call_method).remote(
|
||||
body, raw_request_info
|
||||
):
|
||||
yield response
|
||||
|
||||
async def model(self, model_id: str) -> Optional[ModelCard]:
|
||||
if model_id in self._model_cards:
|
||||
return self._model_cards[model_id]
|
||||
|
||||
base_model_id = get_base_model_id(model_id)
|
||||
base_path = self._lora_paths.get(base_model_id)
|
||||
if base_path is not None:
|
||||
try:
|
||||
overrides = await self._get_lora_model_metadata_func(
|
||||
model_id, base_path
|
||||
)
|
||||
base_card = self._model_cards[base_model_id]
|
||||
return ModelCard(
|
||||
id=model_id,
|
||||
object="model",
|
||||
owned_by=base_card.owned_by,
|
||||
permission=list(base_card.permission),
|
||||
metadata={**base_card.metadata, **overrides},
|
||||
)
|
||||
except HTTPException:
|
||||
logger.exception(
|
||||
"Unable to retrieve LoRA adapter config file for "
|
||||
f'"{model_id}". Omitting it from list of available models. '
|
||||
"Check that adapter config file exists in cloud bucket."
|
||||
)
|
||||
|
||||
async def models(self) -> ModelList:
|
||||
"""OpenAI API-compliant endpoint to get all rayllm models."""
|
||||
all_models = dict()
|
||||
for base_model_id in self._model_cards:
|
||||
# Add the base model.
|
||||
all_models[base_model_id] = await self.model(base_model_id)
|
||||
|
||||
base_path = self._lora_paths.get(base_model_id)
|
||||
if base_path is not None:
|
||||
# Add all the fine-tuned models.
|
||||
lora_model_ids = get_lora_model_ids(
|
||||
dynamic_lora_loading_path=base_path,
|
||||
base_model_id=base_model_id,
|
||||
)
|
||||
for lora_id in lora_model_ids:
|
||||
model_data = await self.model(lora_id)
|
||||
if model_data is not None:
|
||||
all_models[lora_id] = model_data
|
||||
|
||||
return ModelList(data=list(all_models.values()))
|
||||
|
||||
async def model_data(self, model: str) -> ModelCard:
|
||||
"""OpenAI API-compliant endpoint to get one rayllm model.
|
||||
|
||||
Args:
|
||||
model: The model ID (e.g. "amazon/LightGPT").
|
||||
|
||||
Returns:
|
||||
The ``ModelCard`` for ``model``.
|
||||
"""
|
||||
model = replace_prefix(model)
|
||||
model_data = await self.model(model)
|
||||
if model_data is None:
|
||||
raise OpenAIHTTPException(
|
||||
message=f"Unable to find {model}. Please ensure that the model exists and you have permission.",
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
type="InvalidModel",
|
||||
)
|
||||
return model_data
|
||||
|
||||
async def _process_llm_request(
|
||||
self,
|
||||
body: Union[CompletionRequest, ChatCompletionRequest, TranscriptionRequest],
|
||||
call_method: str,
|
||||
raw_request: Optional[Request] = None,
|
||||
) -> Response:
|
||||
|
||||
async with router_request_timeout(DEFAULT_LLM_ROUTER_HTTP_TIMEOUT):
|
||||
|
||||
gen = self._get_response(
|
||||
body=body, call_method=call_method, raw_request=raw_request
|
||||
)
|
||||
|
||||
# In streaming with batching enabled, this first response can be a list of chunks.
|
||||
initial_response, gen = await _peek_at_generator(gen)
|
||||
|
||||
if isinstance(initial_response, list):
|
||||
first_chunk = initial_response[0]
|
||||
else:
|
||||
first_chunk = initial_response
|
||||
|
||||
if isinstance(first_chunk, ErrorResponse):
|
||||
raise OpenAIHTTPException(
|
||||
message=first_chunk.error.message,
|
||||
status_code=first_chunk.error.code,
|
||||
type=first_chunk.error.type,
|
||||
)
|
||||
|
||||
if isinstance(first_chunk, NON_STREAMING_RESPONSE_TYPES):
|
||||
# Not streaming, first chunk should be a single response
|
||||
return JSONResponse(content=first_chunk.model_dump())
|
||||
|
||||
# In case of streaming we need to iterate over the chunks and yield them
|
||||
openai_stream_generator = _openai_json_wrapper(gen)
|
||||
|
||||
return StreamingResponse(
|
||||
openai_stream_generator, media_type="text/event-stream"
|
||||
)
|
||||
|
||||
async def completions(self, body: CompletionRequest, request: Request) -> Response:
|
||||
"""Given a prompt, the model will return one or more predicted completions,
|
||||
and can also return the probabilities of alternative tokens at each position.
|
||||
|
||||
Args:
|
||||
body: The completion request.
|
||||
request: The raw FastAPI request object.
|
||||
|
||||
Returns:
|
||||
A response object with completions.
|
||||
"""
|
||||
return await self._process_llm_request(
|
||||
body, call_method=CallMethod.COMPLETIONS.value, raw_request=request
|
||||
)
|
||||
|
||||
async def chat(self, body: ChatCompletionRequest, request: Request) -> Response:
|
||||
"""Given a prompt, the model will return one or more predicted completions,
|
||||
and can also return the probabilities of alternative tokens at each position.
|
||||
|
||||
Args:
|
||||
body: The chat completion request.
|
||||
request: The raw FastAPI request object.
|
||||
|
||||
Returns:
|
||||
A response object with completions.
|
||||
"""
|
||||
return await self._process_llm_request(
|
||||
body, call_method=CallMethod.CHAT.value, raw_request=request
|
||||
)
|
||||
|
||||
async def embeddings(self, body: EmbeddingRequest, request: Request) -> Response:
|
||||
"""Create embeddings for the provided input.
|
||||
|
||||
Args:
|
||||
body: The embedding request.
|
||||
request: The raw FastAPI request object.
|
||||
|
||||
Returns:
|
||||
A response object with embeddings.
|
||||
"""
|
||||
async with router_request_timeout(DEFAULT_LLM_ROUTER_HTTP_TIMEOUT):
|
||||
results = self._get_response(
|
||||
body=body, call_method="embeddings", raw_request=request
|
||||
)
|
||||
result = await results.__anext__()
|
||||
if isinstance(result, ErrorResponse):
|
||||
raise OpenAIHTTPException(
|
||||
message=result.error.message,
|
||||
status_code=result.error.code,
|
||||
type=result.error.type,
|
||||
)
|
||||
|
||||
if isinstance(result, EmbeddingResponse):
|
||||
return JSONResponse(content=result.model_dump())
|
||||
|
||||
# Annotated[..., Form()] is wrapper that is used to handle multiple form data, which is how audio is sent in transcription requests.
|
||||
# vLLM implementation for handling transcription requests: https://github.com/vllm-project/vllm/blob/0825197bee8dea547f2ab25f48afd8aea0cd2578/vllm/entrypoints/openai/api_server.py#L839.
|
||||
async def transcriptions(
|
||||
self, body: Annotated[TranscriptionRequest, Form()], request: Request
|
||||
) -> Response:
|
||||
"""Create transcription for the provided audio input.
|
||||
|
||||
Args:
|
||||
body: The TranscriptionRequest object.
|
||||
request: The raw FastAPI request object.
|
||||
|
||||
Returns:
|
||||
A response object with transcriptions.
|
||||
"""
|
||||
|
||||
return await self._process_llm_request(
|
||||
body, call_method=CallMethod.TRANSCRIPTIONS.value, raw_request=request
|
||||
)
|
||||
|
||||
async def score(self, body: ScoreRequest, request: Request) -> Response:
|
||||
"""Create scores for the provided text pairs.
|
||||
|
||||
Note: This is a vLLM specific endpoint.
|
||||
|
||||
Args:
|
||||
body: The score request containing input text pairs to score.
|
||||
request: The raw FastAPI request object.
|
||||
|
||||
Returns:
|
||||
A response object with scores.
|
||||
"""
|
||||
|
||||
async with router_request_timeout(DEFAULT_LLM_ROUTER_HTTP_TIMEOUT):
|
||||
results = self._get_response(
|
||||
body=body, call_method="score", raw_request=request
|
||||
)
|
||||
result = await results.__anext__()
|
||||
if isinstance(result, ErrorResponse):
|
||||
raise OpenAIHTTPException(
|
||||
message=result.error.message,
|
||||
status_code=result.error.code,
|
||||
type=result.error.type,
|
||||
)
|
||||
|
||||
if isinstance(result, ScoreResponse):
|
||||
return JSONResponse(content=result.model_dump())
|
||||
|
||||
async def tokenize(
|
||||
self, body: TokenizeCompletionRequest, request: Request
|
||||
) -> Response:
|
||||
"""Tokenize text into token IDs.
|
||||
|
||||
This endpoint tokenizes the provided text prompt and returns the token IDs,
|
||||
counts, and optionally token strings.
|
||||
|
||||
Note: This is a vLLM specific endpoint.
|
||||
|
||||
Args:
|
||||
body: The tokenize request containing the text to tokenize.
|
||||
request: The raw FastAPI request object.
|
||||
|
||||
Returns:
|
||||
A response object with token IDs and metadata.
|
||||
"""
|
||||
async with router_request_timeout(DEFAULT_LLM_ROUTER_HTTP_TIMEOUT):
|
||||
results = self._get_response(
|
||||
body=body, call_method="tokenize", raw_request=request
|
||||
)
|
||||
result = await results.__anext__()
|
||||
if isinstance(result, ErrorResponse):
|
||||
raise OpenAIHTTPException(
|
||||
message=result.error.message,
|
||||
status_code=result.error.code,
|
||||
type=result.error.type,
|
||||
)
|
||||
|
||||
if isinstance(result, TokenizeResponse):
|
||||
return JSONResponse(content=result.model_dump())
|
||||
|
||||
async def detokenize(self, body: DetokenizeRequest, request: Request) -> Response:
|
||||
"""Convert token IDs back to text.
|
||||
|
||||
This endpoint detokenizes the provided token IDs and returns the
|
||||
corresponding text.
|
||||
|
||||
Note: This is a vLLM specific endpoint.
|
||||
|
||||
Args:
|
||||
body: The detokenize request containing the token IDs.
|
||||
request: The raw FastAPI request object.
|
||||
|
||||
Returns:
|
||||
A response object with the detokenized text.
|
||||
"""
|
||||
async with router_request_timeout(DEFAULT_LLM_ROUTER_HTTP_TIMEOUT):
|
||||
results = self._get_response(
|
||||
body=body, call_method="detokenize", raw_request=request
|
||||
)
|
||||
result = await results.__anext__()
|
||||
if isinstance(result, ErrorResponse):
|
||||
raise OpenAIHTTPException(
|
||||
message=result.error.message,
|
||||
status_code=result.error.code,
|
||||
type=result.error.type,
|
||||
)
|
||||
|
||||
if isinstance(result, DetokenizeResponse):
|
||||
return JSONResponse(content=result.model_dump())
|
||||
|
||||
@classmethod
|
||||
def get_deployment_options(
|
||||
cls, llm_configs: Optional[List[LLMConfig]] = None
|
||||
) -> Dict[str, Any]:
|
||||
"""Get the deployment options for the ingress deployment.
|
||||
|
||||
If all models are configured with min_replicas=0 (scale-to-zero),
|
||||
the ingress will also be configured with min_replicas=0 so that
|
||||
the worker node/GPU instance can be fully released when idle.
|
||||
|
||||
Args:
|
||||
llm_configs: The LLM configs to infer the number of ingress replicas from.
|
||||
|
||||
Returns:
|
||||
A dictionary containing the deployment options for the ingress deployment.
|
||||
"""
|
||||
options = copy.deepcopy(DEFAULT_INGRESS_OPTIONS)
|
||||
if _all_models_scale_to_zero(llm_configs):
|
||||
options.setdefault("autoscaling_config", {})["min_replicas"] = 0
|
||||
return options
|
||||
@@ -0,0 +1,181 @@
|
||||
import uuid
|
||||
from asyncio import CancelledError
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import FastAPI, Request, status
|
||||
from fastapi.exceptions import RequestValidationError
|
||||
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
|
||||
from starlette.responses import JSONResponse, Response
|
||||
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
from ray.llm._internal.serve.utils.server_utils import (
|
||||
get_response_for_error,
|
||||
)
|
||||
|
||||
logger = get_logger(__file__)
|
||||
|
||||
|
||||
def get_request_id(request: Request) -> str:
|
||||
"""Fetches request-id from Starlette's request object.
|
||||
|
||||
NOTE: This method relies on "request_id" value to be injected into the
|
||||
Starlette's ``request.state`` via ``inject_request_id`` middleware.
|
||||
|
||||
Args:
|
||||
request: Starlette request object.
|
||||
|
||||
Returns:
|
||||
Id allowing to identify the particular request, or ``None`` if not set.
|
||||
"""
|
||||
return getattr(request.state, "request_id", None)
|
||||
|
||||
|
||||
async def _handle_validation_error(
|
||||
request: Request, exc: RequestValidationError
|
||||
) -> JSONResponse:
|
||||
"""Handle pydantic validation errors in an OpenAI-like format."""
|
||||
error_details = exc.errors()[0] if exc.errors() else {"msg": "Invalid request"}
|
||||
|
||||
error_msg = error_details.get("msg", "Unknown validation error")
|
||||
error_loc = error_details.get("loc", ("body"))
|
||||
error_input = error_details.get("input", None)
|
||||
msg = f"Invalid request format: {error_msg} at {error_loc}"
|
||||
|
||||
error_response = {
|
||||
"error": {
|
||||
"message": msg,
|
||||
"type": error_details.get("type", "invalid_request_error"),
|
||||
"param": error_input,
|
||||
"code": "invalid_parameter",
|
||||
}
|
||||
}
|
||||
|
||||
return JSONResponse(status_code=status.HTTP_400_BAD_REQUEST, content=error_response)
|
||||
|
||||
|
||||
def _uncaught_exception_handler(request: Request, e: Exception):
|
||||
"""This method serves as an uncaught exception handler being
|
||||
the last resort to return properly formatted response.
|
||||
|
||||
NOTE: Exceptions from application handlers should NOT be reaching this point,
|
||||
this handler is here to intercept "fly-away" exceptions and should not
|
||||
be handled for handling of converting application exceptions into
|
||||
appropriate responses
|
||||
"""
|
||||
|
||||
if isinstance(e, CancelledError):
|
||||
return JSONResponse(content={}, status_code=204)
|
||||
|
||||
request_id = get_request_id(request)
|
||||
|
||||
logger.error(f"Uncaught exception while handling request {request_id}", exc_info=e)
|
||||
|
||||
error_response = get_response_for_error(e, request_id)
|
||||
|
||||
return JSONResponse(
|
||||
content=error_response.model_dump(), status_code=error_response.error.code
|
||||
)
|
||||
|
||||
|
||||
def add_exception_handling_middleware(router: FastAPI):
|
||||
# NOTE: PLEASE READ CAREFULLY BEFORE CHANGING
|
||||
#
|
||||
# Starlette has different behavior depending on the Exception class being handled
|
||||
# that we unfortunately have to take into account here:
|
||||
#
|
||||
# - Handler for `Exception` will be added as uncaught exception handler (of last resort)
|
||||
# that is going to be executed absolute last, making sure that in case of any fly-away
|
||||
# (uncaught) exception
|
||||
# - Handlers for any other classes of exceptions will be executed as last middleware layer,
|
||||
# therefore being to intercept any exceptions originating from the handler before it
|
||||
# propagates to the middleware above it
|
||||
#
|
||||
# As such we're aiming for 2 goals here:
|
||||
# - Intercepting exceptions from the handlers, converting them into proper user-facing
|
||||
# response (avoiding exception propagation up the middleware stack)
|
||||
# - Adding uncaught exception handler (of last resort) to intercept any exceptions that
|
||||
# might be originating from the middleware itself
|
||||
|
||||
async def _handle_application_exceptions(
|
||||
request: Request, call_next: RequestResponseEndpoint
|
||||
) -> Response:
|
||||
"""This method intercepts application level exceptions not handled by the
|
||||
application code converting them into appropriately formatted (JSON) response
|
||||
"""
|
||||
|
||||
try:
|
||||
return await call_next(request)
|
||||
except CancelledError as ce:
|
||||
# NOTE: We re-raise CancelledError as is to let other middleware handle it.
|
||||
# Since no response is expected in this case, it's deferred to uncaught
|
||||
# exception handler to ultimately handle it
|
||||
raise ce
|
||||
except RequestValidationError as e:
|
||||
return await _handle_validation_error(request, e)
|
||||
except Exception as e:
|
||||
request_id = get_request_id(request)
|
||||
error_response = get_response_for_error(e, request_id)
|
||||
|
||||
return JSONResponse(
|
||||
content=error_response.model_dump(),
|
||||
status_code=error_response.error.code,
|
||||
)
|
||||
|
||||
# This adds last-resort uncaught exception handler into Starlette
|
||||
router.add_exception_handler(Exception, _uncaught_exception_handler)
|
||||
# Add validation error handler
|
||||
router.add_exception_handler(RequestValidationError, _handle_validation_error)
|
||||
# This adds application exception handler, allowing to convert application
|
||||
# exceptions into properly formatted responses
|
||||
router.add_middleware(
|
||||
BaseHTTPMiddleware,
|
||||
dispatch=_handle_application_exceptions,
|
||||
)
|
||||
|
||||
|
||||
class SetRequestIdMiddleware:
|
||||
"""Injects request ID into the request's state.
|
||||
|
||||
The ID is either:
|
||||
1. the value of the request's "x-request-id" header, set by Ray
|
||||
Serve's Proxy, or
|
||||
2. if "x-request-id" header is unavailable, this middleware creates
|
||||
a UUIDv4 request ID.
|
||||
"""
|
||||
|
||||
def __init__(self, app):
|
||||
self.app = app
|
||||
|
||||
async def __call__(self, scope, receive, send):
|
||||
if scope["type"] == "http":
|
||||
headers = list(scope.get("headers", []))
|
||||
request_id = None
|
||||
for name, value in headers:
|
||||
if name.lower() == b"x-request-id" and value:
|
||||
request_id = value.decode()
|
||||
break
|
||||
|
||||
if request_id is None:
|
||||
request_id = str(uuid.uuid4())
|
||||
headers.append((b"x-request-id", request_id.encode()))
|
||||
|
||||
scope["headers"] = headers
|
||||
request = Request(scope)
|
||||
request.state.request_id = request_id
|
||||
|
||||
return await self.app(scope, receive, send)
|
||||
|
||||
|
||||
def get_user_id(request: Request) -> Optional[str]:
|
||||
"""Fetches user id inside Starlette's request object.
|
||||
|
||||
NOTE: This method relies on "user_id" value to be injected into the
|
||||
Starlette's ``request.state`` via authentication middleware.
|
||||
|
||||
Args:
|
||||
request: Starlette request object.
|
||||
|
||||
Returns:
|
||||
Id identifying the particular user, or ``None`` if not set.
|
||||
"""
|
||||
return getattr(request.state, "user_id", None)
|
||||
@@ -0,0 +1,44 @@
|
||||
"""Ingress capability mixins.
|
||||
|
||||
Provides HTTP endpoint mixins for control plane operations.
|
||||
"""
|
||||
|
||||
from ray.llm._internal.serve.core.ingress.mixins.cache_manager import (
|
||||
CacheManagerIngressMixin,
|
||||
ResetPrefixCacheRequest,
|
||||
)
|
||||
from ray.llm._internal.serve.core.ingress.mixins.collective_rpc import (
|
||||
CollectiveRpcIngressMixin,
|
||||
CollectiveRpcRequest,
|
||||
CollectiveRpcResponse,
|
||||
ReplicaResult,
|
||||
)
|
||||
from ray.llm._internal.serve.core.ingress.mixins.pausable import (
|
||||
IsPausedResponse,
|
||||
PausableIngressMixin,
|
||||
PauseRequest,
|
||||
ResumeRequest,
|
||||
)
|
||||
from ray.llm._internal.serve.core.ingress.mixins.sleepable import (
|
||||
IsSleepingResponse,
|
||||
SleepableIngressMixin,
|
||||
SleepRequest,
|
||||
WakeupRequest,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"CacheManagerIngressMixin",
|
||||
"CollectiveRpcIngressMixin",
|
||||
"PausableIngressMixin",
|
||||
"SleepableIngressMixin",
|
||||
"CollectiveRpcRequest",
|
||||
"CollectiveRpcResponse",
|
||||
"ReplicaResult",
|
||||
"ResetPrefixCacheRequest",
|
||||
"PauseRequest",
|
||||
"ResumeRequest",
|
||||
"IsPausedResponse",
|
||||
"SleepRequest",
|
||||
"WakeupRequest",
|
||||
"IsSleepingResponse",
|
||||
]
|
||||
@@ -0,0 +1,26 @@
|
||||
import asyncio
|
||||
from typing import Any, List
|
||||
|
||||
from ray.llm._internal.serve.utils.broadcast import broadcast
|
||||
|
||||
|
||||
class ReplicaBroadcastable:
|
||||
async def _broadcast_to_replicas(
|
||||
self, model: str, method: str, kwargs: dict | None = None
|
||||
) -> List[Any]:
|
||||
"""Broadcast a command to all replicas and return their results.
|
||||
|
||||
Args:
|
||||
model: The model ID to broadcast to.
|
||||
method: The method name to call on each replica.
|
||||
kwargs: Optional kwargs to pass to the method.
|
||||
|
||||
Returns:
|
||||
List of results from each replica.
|
||||
"""
|
||||
model_id = await self._get_model_id(model)
|
||||
handle = self._get_configured_serve_handle(model_id)
|
||||
# Run blocking broadcast() in a thread to avoid blocking the event loop.
|
||||
# broadcast() uses ray.get() internally which is synchronous.
|
||||
results = await asyncio.to_thread(broadcast, handle, method, kwargs=kwargs)
|
||||
return results
|
||||
@@ -0,0 +1,53 @@
|
||||
"""Cache manager ingress mixin.
|
||||
|
||||
Provides HTTP endpoints for cache management control plane operations.
|
||||
"""
|
||||
|
||||
from pydantic import BaseModel
|
||||
from starlette.responses import Response
|
||||
|
||||
from ray.llm._internal.serve.core.ingress.mixins.broadcastable import (
|
||||
ReplicaBroadcastable,
|
||||
)
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
# --- Pydantic Models ---
|
||||
|
||||
|
||||
class ResetPrefixCacheRequest(BaseModel):
|
||||
"""Request to reset the prefix cache."""
|
||||
|
||||
model: str
|
||||
|
||||
|
||||
# --- Mixin ---
|
||||
|
||||
|
||||
class CacheManagerIngressMixin(ReplicaBroadcastable):
|
||||
"""Ingress mixin for /reset_prefix_cache endpoint.
|
||||
|
||||
Adds control plane endpoint for managing the KV prefix cache.
|
||||
"""
|
||||
|
||||
ENDPOINTS = {
|
||||
"reset_prefix_cache": lambda app: app.post("/reset_prefix_cache"),
|
||||
}
|
||||
|
||||
async def reset_prefix_cache(self, body: ResetPrefixCacheRequest) -> Response:
|
||||
"""Reset the KV prefix cache on all replicas for the specified model.
|
||||
|
||||
Clears cached key-value pairs from previous requests. Useful for
|
||||
benchmarking or when cache invalidation is needed.
|
||||
|
||||
Args:
|
||||
body: Request containing the model ID.
|
||||
|
||||
Returns:
|
||||
200 OK on success.
|
||||
"""
|
||||
logger.info("Resetting prefix cache for model: %s", body.model)
|
||||
await self._broadcast_to_replicas(body.model, "reset_prefix_cache")
|
||||
return Response(status_code=200)
|
||||
@@ -0,0 +1,100 @@
|
||||
"""Collective RPC ingress mixin.
|
||||
|
||||
Provides HTTP endpoint for collective RPC operations across all replicas
|
||||
and their workers, enabling RLHF workflows where a trainer forms a single
|
||||
NCCL process group with all TP/PP workers across all replicas.
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from ray.llm._internal.serve.core.ingress.mixins.broadcastable import (
|
||||
ReplicaBroadcastable,
|
||||
)
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
# --- Pydantic Models ---
|
||||
|
||||
|
||||
class CollectiveRpcRequest(BaseModel):
|
||||
"""Request to execute a collective RPC on all replicas."""
|
||||
|
||||
model: str
|
||||
method: str
|
||||
args: List[Any] = Field(default_factory=list)
|
||||
kwargs: Dict[str, Any] = Field(default_factory=dict)
|
||||
timeout: Optional[float] = None
|
||||
|
||||
|
||||
class ReplicaResult(BaseModel):
|
||||
"""Result from a single replica containing all worker results."""
|
||||
|
||||
replica: int
|
||||
worker_results: List[Any]
|
||||
|
||||
|
||||
class CollectiveRpcResponse(BaseModel):
|
||||
"""Response containing results from all replicas."""
|
||||
|
||||
results: List[ReplicaResult]
|
||||
|
||||
|
||||
# --- Mixin ---
|
||||
|
||||
|
||||
class CollectiveRpcIngressMixin(ReplicaBroadcastable):
|
||||
"""Ingress mixin for /collective_rpc endpoint.
|
||||
|
||||
Adds control plane endpoint for executing collective RPC calls across
|
||||
all replicas and their workers. This is used for RLHF workflows where
|
||||
a trainer needs to communicate with all TP/PP workers across all replicas.
|
||||
"""
|
||||
|
||||
ENDPOINTS = {
|
||||
"collective_rpc": lambda app: app.post("/collective_rpc"),
|
||||
}
|
||||
|
||||
async def collective_rpc(self, body: CollectiveRpcRequest) -> CollectiveRpcResponse:
|
||||
"""Execute a collective RPC on all replicas for the specified model.
|
||||
|
||||
This broadcasts the RPC call to all replicas, and each replica
|
||||
executes the call on all its workers (TP/PP ranks).
|
||||
|
||||
Args:
|
||||
body: Request containing the model ID, method name, args, kwargs,
|
||||
and optional timeout.
|
||||
|
||||
Returns:
|
||||
CollectiveRpcResponse with results from all replicas.
|
||||
"""
|
||||
logger.info(
|
||||
"Executing collective_rpc '%s' for model %s with args=%s, kwargs=%s",
|
||||
body.method,
|
||||
body.model,
|
||||
body.args,
|
||||
body.kwargs,
|
||||
)
|
||||
|
||||
# Broadcast to all replicas - each replica returns a list of worker results
|
||||
replica_results = await self._broadcast_to_replicas(
|
||||
body.model,
|
||||
"collective_rpc",
|
||||
kwargs={
|
||||
"method": body.method,
|
||||
"args": tuple(body.args),
|
||||
"kwargs": body.kwargs,
|
||||
"timeout": body.timeout,
|
||||
},
|
||||
)
|
||||
|
||||
# Format results with replica index for debugging
|
||||
results = [
|
||||
ReplicaResult(replica=i, worker_results=worker_results or [])
|
||||
for i, worker_results in enumerate(replica_results or [])
|
||||
]
|
||||
|
||||
return CollectiveRpcResponse(results=results)
|
||||
@@ -0,0 +1,115 @@
|
||||
"""Pausable ingress mixin.
|
||||
|
||||
Provides HTTP endpoints for pause/resume control plane operations.
|
||||
"""
|
||||
|
||||
from typing import Any, Dict
|
||||
|
||||
from fastapi import Query
|
||||
from pydantic import BaseModel, Field
|
||||
from starlette.responses import Response
|
||||
|
||||
from ray.llm._internal.serve.core.ingress.mixins.broadcastable import (
|
||||
ReplicaBroadcastable,
|
||||
)
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
# --- Pydantic Models ---
|
||||
|
||||
|
||||
class PauseRequest(BaseModel):
|
||||
"""Request to pause generation on an engine."""
|
||||
|
||||
model: str
|
||||
options: Dict[str, Any] = Field(
|
||||
default_factory=dict,
|
||||
description="Engine-specific pause options (e.g., mode, clear_cache)",
|
||||
)
|
||||
|
||||
|
||||
class ResumeRequest(BaseModel):
|
||||
"""Request to resume generation on an engine."""
|
||||
|
||||
model: str
|
||||
options: Dict[str, Any] = Field(
|
||||
default_factory=dict,
|
||||
description="Engine-specific resume options",
|
||||
)
|
||||
|
||||
|
||||
class IsPausedResponse(BaseModel):
|
||||
"""Response indicating whether the engine is paused."""
|
||||
|
||||
is_paused: bool
|
||||
|
||||
|
||||
# --- Mixin ---
|
||||
|
||||
|
||||
class PausableIngressMixin(ReplicaBroadcastable):
|
||||
"""Ingress mixin for /pause, /resume, /is_paused endpoints.
|
||||
|
||||
Adds control plane endpoints for managing engine pause state.
|
||||
Pause mode halts generation/encoding while keeping weights in GPU memory.
|
||||
Unlike sleep mode, pause does not offload weights to CPU.
|
||||
"""
|
||||
|
||||
ENDPOINTS = {
|
||||
"pause": lambda app: app.post("/pause"),
|
||||
"resume": lambda app: app.post("/resume"),
|
||||
"is_paused": lambda app: app.get("/is_paused"),
|
||||
}
|
||||
|
||||
async def pause(self, body: PauseRequest) -> Response:
|
||||
"""Pause generation on all replicas for the specified model.
|
||||
|
||||
This halts generation/encoding requests while keeping model weights
|
||||
in GPU memory. New requests are blocked until resume is called.
|
||||
Unlike sleep mode, pause does not offload weights to CPU.
|
||||
|
||||
Args:
|
||||
body: Request containing the model ID and engine-specific options.
|
||||
Options may include:
|
||||
- mode (str): "abort" (default), "wait", or "keep".
|
||||
- clear_cache (bool): Clear KV cache after draining. Default True.
|
||||
|
||||
Returns:
|
||||
200 OK on success.
|
||||
"""
|
||||
logger.info("Pausing model %s with options: %s", body.model, body.options)
|
||||
await self._broadcast_to_replicas(body.model, "pause", kwargs=body.options)
|
||||
return Response(status_code=200)
|
||||
|
||||
async def resume(self, body: ResumeRequest) -> Response:
|
||||
"""Resume generation on all replicas for the specified model.
|
||||
|
||||
Args:
|
||||
body: Request containing the model ID and engine-specific options.
|
||||
|
||||
Returns:
|
||||
200 OK on success.
|
||||
"""
|
||||
logger.info("Resuming model %s with options: %s", body.model, body.options)
|
||||
await self._broadcast_to_replicas(body.model, "resume", kwargs=body.options)
|
||||
return Response(status_code=200)
|
||||
|
||||
async def is_paused(
|
||||
self, model: str = Query(..., description="The model ID to check")
|
||||
) -> IsPausedResponse:
|
||||
"""Check if the engine is paused for the specified model.
|
||||
|
||||
This checks the pause status across all replicas. Returns True if
|
||||
ANY replica is paused (uses logical OR across replicas).
|
||||
|
||||
Args:
|
||||
model: The model ID to check.
|
||||
|
||||
Returns:
|
||||
IsPausedResponse with is_paused boolean.
|
||||
"""
|
||||
results = await self._broadcast_to_replicas(model, "is_paused")
|
||||
is_paused_result = any(results) if results else False
|
||||
return IsPausedResponse(is_paused=is_paused_result)
|
||||
@@ -0,0 +1,112 @@
|
||||
"""Sleepable ingress mixin.
|
||||
|
||||
Provides HTTP endpoints for sleep/wakeup control plane operations.
|
||||
"""
|
||||
|
||||
from typing import Any, Dict
|
||||
|
||||
from fastapi import Query
|
||||
from pydantic import BaseModel, Field
|
||||
from starlette.responses import Response
|
||||
|
||||
from ray.llm._internal.serve.core.ingress.mixins.broadcastable import (
|
||||
ReplicaBroadcastable,
|
||||
)
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
# --- Pydantic Models ---
|
||||
|
||||
|
||||
class SleepRequest(BaseModel):
|
||||
"""Request to put an engine to sleep."""
|
||||
|
||||
model: str
|
||||
options: Dict[str, Any] = Field(
|
||||
default_factory=dict,
|
||||
description="Engine-specific sleep options (e.g., level for vLLM)",
|
||||
)
|
||||
|
||||
|
||||
class WakeupRequest(BaseModel):
|
||||
"""Request to wake up an engine from sleep."""
|
||||
|
||||
model: str
|
||||
options: Dict[str, Any] = Field(
|
||||
default_factory=dict,
|
||||
description="Engine-specific wakeup options (e.g., tags for vLLM)",
|
||||
)
|
||||
|
||||
|
||||
class IsSleepingResponse(BaseModel):
|
||||
"""Response indicating whether the engine is sleeping."""
|
||||
|
||||
is_sleeping: bool
|
||||
|
||||
|
||||
# --- Mixin ---
|
||||
|
||||
|
||||
class SleepableIngressMixin(ReplicaBroadcastable):
|
||||
"""Ingress mixin for /sleep, /wakeup, /is_sleeping endpoints.
|
||||
|
||||
Adds control plane endpoints for managing engine sleep state.
|
||||
Sleep mode offloads model weights to CPU and discards KV cache.
|
||||
"""
|
||||
|
||||
ENDPOINTS = {
|
||||
"sleep": lambda app: app.post("/sleep"),
|
||||
"wakeup": lambda app: app.post("/wakeup"),
|
||||
"is_sleeping": lambda app: app.get("/is_sleeping"),
|
||||
}
|
||||
|
||||
async def sleep(self, body: SleepRequest) -> Response:
|
||||
"""Put the engine to sleep on all replicas for the specified model.
|
||||
|
||||
This offloads model weights to CPU and discards KV cache, freeing
|
||||
GPU memory. The engine cannot process requests while sleeping.
|
||||
|
||||
Args:
|
||||
body: Request containing the model ID and engine-specific options.
|
||||
|
||||
Returns:
|
||||
200 OK on success.
|
||||
"""
|
||||
logger.info(
|
||||
"Putting model %s to sleep with options: %s", body.model, body.options
|
||||
)
|
||||
await self._broadcast_to_replicas(body.model, "sleep", kwargs=body.options)
|
||||
return Response(status_code=200)
|
||||
|
||||
async def wakeup(self, body: WakeupRequest) -> Response:
|
||||
"""Wake up the engine from sleep on all replicas for the specified model.
|
||||
|
||||
Args:
|
||||
body: Request containing the model ID and engine-specific options.
|
||||
|
||||
Returns:
|
||||
200 OK on success.
|
||||
"""
|
||||
logger.info("Waking up model %s with options: %s", body.model, body.options)
|
||||
await self._broadcast_to_replicas(body.model, "wakeup", kwargs=body.options)
|
||||
return Response(status_code=200)
|
||||
|
||||
async def is_sleeping(
|
||||
self, model: str = Query(..., description="The model ID to check")
|
||||
) -> IsSleepingResponse:
|
||||
"""Check if the engine is sleeping for the specified model.
|
||||
|
||||
This checks the sleep status across all replicas. Returns True if
|
||||
ANY replica is sleeping (uses logical OR across replicas).
|
||||
|
||||
Args:
|
||||
model: The model ID to check.
|
||||
|
||||
Returns:
|
||||
IsSleepingResponse with is_sleeping boolean.
|
||||
"""
|
||||
results = await self._broadcast_to_replicas(model, "is_sleeping")
|
||||
is_sleeping_result = any(results) if results else False
|
||||
return IsSleepingResponse(is_sleeping=is_sleeping_result)
|
||||
@@ -0,0 +1,210 @@
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
from fastapi import FastAPI, HTTPException, Request
|
||||
|
||||
from ray import serve
|
||||
from ray.llm._internal.serve.core.ingress.tokenizer import (
|
||||
REQUEST_TOKEN_IDS_KWARG,
|
||||
TokenizeError,
|
||||
Tokenizer,
|
||||
)
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
from ray.serve._private.http_util import _matches_session_id_header
|
||||
from ray.serve.exceptions import DeploymentUnavailableError
|
||||
from ray.serve.handle import DeploymentHandle
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
_BODY_TRUNCATED_HEADER = "x-body-truncated"
|
||||
|
||||
# A request body routes on one of these fields. Body-aware routers read it off
|
||||
# the namespace; a body without any of them degrades to load-balancing. Extend
|
||||
# as routers learn to route additional request types.
|
||||
_ROUTING_KEY_FIELDS = ("messages", "prompt")
|
||||
|
||||
router_app = FastAPI()
|
||||
|
||||
|
||||
def _parse_routing_payload(body: bytes) -> Optional[SimpleNamespace]:
|
||||
"""Wrap a request body as a namespace a body-aware router routes on.
|
||||
|
||||
Routers read a routing field (``messages`` or ``prompt``) off the first
|
||||
positional routing arg, the parsed request the normal ingress forwards.
|
||||
Direct streaming has only the raw body, so this wraps the parsed body in a
|
||||
namespace exposing every field by attribute, which a router reads the same
|
||||
way regardless of request type. Returns ``None`` for an empty, non-object,
|
||||
unparseable, or keyless body, so the caller falls back to load-balancing.
|
||||
"""
|
||||
if not body:
|
||||
return None
|
||||
try:
|
||||
data = json.loads(body)
|
||||
except (ValueError, TypeError):
|
||||
return None
|
||||
if not isinstance(data, dict):
|
||||
return None
|
||||
if not any(data.get(field) for field in _ROUTING_KEY_FIELDS):
|
||||
return None
|
||||
return SimpleNamespace(**data)
|
||||
|
||||
|
||||
@serve.ingress(router_app)
|
||||
class LLMRouter:
|
||||
"""Ingress request router for direct streaming.
|
||||
|
||||
When direct streaming is enabled, HAProxy calls /internal/route on this
|
||||
deployment to get a data plane replica, then forwards traffic directly
|
||||
to the matching LLMServer replica's backend HTTP port.
|
||||
|
||||
Replica selection is delegated to the underlying deployment's configured
|
||||
request router, and this class translates the resulting pick into a backend
|
||||
HTTP endpoint.
|
||||
|
||||
/internal/route HTTP contract
|
||||
-----------------------------
|
||||
Request:
|
||||
POST /internal/route
|
||||
Content-Type: application/json
|
||||
Body: the target ChatCompletions or Completions request payload.
|
||||
Wrapped in a namespace by ``_parse_routing_payload`` and passed to
|
||||
``choose_replica`` positionally, exposing the request fields the way
|
||||
the parsed request does. Body-aware policies then score replicas the
|
||||
same way on both paths.
|
||||
|
||||
Truncated bodies:
|
||||
HAProxy may forward only a prefix of the body for routing and sets the
|
||||
``x-body-truncated`` header. A truncated prefix is usually not valid
|
||||
JSON, so no routing key is derived and the request falls back to the
|
||||
default load-balanced pick.
|
||||
|
||||
Session affinity:
|
||||
If the client request carried the session-id header configured by
|
||||
``RAY_SERVE_SESSION_ID_HEADER_KEY`` (default ``x-session-id``),
|
||||
HAProxy's Lua action forwards it to ``/internal/route`` on the same
|
||||
name. This handler reads it and applies
|
||||
``handle.options(session_id=...)`` before calling
|
||||
``choose_replica`` so session-aware policies (e.g.
|
||||
``ConsistentHashRouter``) pin all turns of a session to one replica.
|
||||
|
||||
Responses:
|
||||
200 ``{"host": str, "port": int, "replica_id": str}``: pick
|
||||
succeeded.
|
||||
4xx/5xx FastAPI ``{"detail": str}``: informational only; HAProxy
|
||||
treats any non-200 as a routing failure. When using KV aware routing,
|
||||
a pre-routing ``/tokenize`` rejection is surfaced here.
|
||||
|
||||
Health:
|
||||
``GET /health`` is exposed as a human-operator convenience.
|
||||
Serve uses ``check_health()`` for replica readiness, not HTTP.
|
||||
"""
|
||||
|
||||
# Warn once per replica when no routing key is derived. Class-level default
|
||||
# keeps the guard safe before __init__ runs.
|
||||
_warned_no_routing_key: bool = False
|
||||
|
||||
async def __init__(
|
||||
self, server: DeploymentHandle, pre_routing_tokenization: bool = False
|
||||
):
|
||||
self._handle: DeploymentHandle = server
|
||||
self._handle._init()
|
||||
# Pre-routing tokenization is only useful to a KV-aware request router,
|
||||
# which scores replicas based on the prompt token IDs.
|
||||
self._tokenizer = Tokenizer(self._handle) if pre_routing_tokenization else None
|
||||
|
||||
@router_app.post("/internal/route")
|
||||
async def route(self, request: Request):
|
||||
body = await request.body()
|
||||
body_truncated = _BODY_TRUNCATED_HEADER in request.headers
|
||||
routing_payload = _parse_routing_payload(body)
|
||||
if routing_payload is None and not self._warned_no_routing_key:
|
||||
self._warned_no_routing_key = True
|
||||
logger.warning(
|
||||
"Could not derive a routing key from the request body. "
|
||||
"body_truncated=%s. Falling back to load-balanced replica "
|
||||
"selection. A configured body-aware router such as "
|
||||
"PrefixCacheAffinityRouter cannot take effect for these "
|
||||
"requests. For truncated bodies, raise HAProxy's routing body "
|
||||
"limit.",
|
||||
body_truncated,
|
||||
)
|
||||
# Tokenize only a parseable, routable body; a truncated or unparseable
|
||||
# body has no routing payload, so fall back to token-less routing.
|
||||
request_token_ids = None
|
||||
if self._tokenizer is not None and routing_payload is not None:
|
||||
try:
|
||||
request_token_ids = await self._tokenizer.tokenize(
|
||||
vars(routing_payload)
|
||||
)
|
||||
except TokenizeError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.message)
|
||||
# HAProxy forwards the configured session header on the same name,
|
||||
# but use the same case-insensitive, separator-tolerant matcher as
|
||||
# proxy.py / ingress.py so a `-`/`_` rewrite anywhere in the path
|
||||
# doesn't silently drop session affinity.
|
||||
session_id = next(
|
||||
(v for k, v in request.headers.items() if _matches_session_id_header(k)),
|
||||
None,
|
||||
)
|
||||
handle = (
|
||||
self._handle.options(session_id=session_id) if session_id else self._handle
|
||||
)
|
||||
try:
|
||||
host, port, replica_id = await self._pick_replica(
|
||||
handle=handle,
|
||||
routing_payload=routing_payload,
|
||||
request_token_ids=request_token_ids,
|
||||
)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
except (RuntimeError, DeploymentUnavailableError) as e:
|
||||
raise HTTPException(status_code=503, detail=str(e))
|
||||
return {"host": host, "port": port, "replica_id": replica_id}
|
||||
|
||||
@router_app.get("/health")
|
||||
async def health(self):
|
||||
return {"status": "ok"}
|
||||
|
||||
async def _pick_replica(
|
||||
self,
|
||||
handle: DeploymentHandle,
|
||||
routing_payload: Optional[SimpleNamespace] = None,
|
||||
request_token_ids: Optional[List[int]] = None,
|
||||
) -> Tuple[str, int, str]:
|
||||
"""Pick a backend HTTP replica via the deployment's request router.
|
||||
|
||||
``handle`` is the LLMServer deployment handle, optionally configured
|
||||
with ``.options(session_id=...)`` by the caller so session-aware
|
||||
routers see the session id on ``RequestMetadata``.
|
||||
|
||||
``routing_payload``, when present, is passed to ``choose_replica``
|
||||
positionally. It lands in ``pending_request.args`` where the normal
|
||||
ingress puts the parsed request, so a body-aware policy scores replicas
|
||||
as on the normal path. When ``None``, nothing is forwarded. The router
|
||||
sees empty ``args`` and falls back to its default load-balanced pick.
|
||||
|
||||
``request_token_ids``, when present, is forwarded as a keyword arg so a
|
||||
KV-aware request router can score replicas on prompt-prefix overlap.
|
||||
|
||||
``_reserve=False`` short-circuits the replica-side ``reserve_slot``
|
||||
actor RPC and the rejection-retry loop: the real request goes out via
|
||||
HAProxy, so Serve's capacity semaphore isn't load-bearing here, and
|
||||
the extra RPC + retry introduced burstiness compared to the prior
|
||||
local round-robin implementation.
|
||||
"""
|
||||
route_args = (routing_payload,) if routing_payload is not None else ()
|
||||
choose_replica_kwargs = {"_reserve": False}
|
||||
if request_token_ids is not None:
|
||||
choose_replica_kwargs[REQUEST_TOKEN_IDS_KWARG] = request_token_ids
|
||||
async with handle.choose_replica(
|
||||
*route_args, **choose_replica_kwargs
|
||||
) as selection:
|
||||
replica = selection._replica
|
||||
endpoint = replica.backend_http_endpoint
|
||||
if endpoint is None:
|
||||
raise RuntimeError(
|
||||
f"replica {selection.replica_id} has no backend HTTP endpoint"
|
||||
)
|
||||
host, port = endpoint
|
||||
return host, port, replica.replica_id.to_full_id_str()
|
||||
@@ -0,0 +1,124 @@
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
from ray.llm._internal.serve.core.configs.openai_api_models import (
|
||||
ErrorResponse,
|
||||
TokenizeChatRequest,
|
||||
TokenizeCompletionRequest,
|
||||
)
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
from ray.serve.handle import DeploymentHandle
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
# choose_replica kwarg carrying the prompt token IDs to KV-aware routers.
|
||||
REQUEST_TOKEN_IDS_KWARG = "request_token_ids"
|
||||
|
||||
|
||||
class TokenizeError(Exception):
|
||||
"""The ``/tokenize`` endpoint rejected the request.
|
||||
|
||||
Carries vLLM's HTTP ``status_code``, ``message`` and error ``type``.
|
||||
"""
|
||||
|
||||
def __init__(self, message: str, *, status_code: int, type: str):
|
||||
super().__init__(message)
|
||||
self.message = message
|
||||
self.status_code = status_code
|
||||
self.type = type
|
||||
|
||||
|
||||
class Tokenizer:
|
||||
"""Tokenizes incoming requests via the replica's ``/tokenize`` endpoint.
|
||||
|
||||
Args:
|
||||
handle: A handle to the LLMServer deployment.
|
||||
"""
|
||||
|
||||
def __init__(self, handle: DeploymentHandle):
|
||||
self._handle = handle
|
||||
|
||||
async def tokenize(self, payload: Dict[str, Any]) -> Optional[List[int]]:
|
||||
"""Tokenize a request ``payload`` into prompt token IDs.
|
||||
|
||||
Args:
|
||||
payload: The request body, already parsed into a dict by ``LLMRouter``.
|
||||
|
||||
Returns:
|
||||
The prompt token IDs, or ``None`` for bodies that are not routed on.
|
||||
|
||||
Raises:
|
||||
TokenizeError: The ``/tokenize`` endpoint rejected the request.
|
||||
"""
|
||||
tok_req = self._build_tokenize_request(payload)
|
||||
if tok_req is None:
|
||||
return None
|
||||
|
||||
# /tokenize yields a single response; drain the stream fully so the
|
||||
# handle response is cleaned up.
|
||||
resp = None
|
||||
async for chunk in self._handle.options(stream=True).tokenize.remote(
|
||||
tok_req, None
|
||||
):
|
||||
resp = chunk
|
||||
if resp is None:
|
||||
raise TokenizeError(
|
||||
"/tokenize returned no response",
|
||||
status_code=500,
|
||||
type="internal_error",
|
||||
)
|
||||
if isinstance(resp, ErrorResponse):
|
||||
raise TokenizeError(
|
||||
resp.error.message,
|
||||
status_code=resp.error.code,
|
||||
type=resp.error.type,
|
||||
)
|
||||
return list(resp.tokens)
|
||||
|
||||
def _build_tokenize_request(
|
||||
self, payload: Dict[str, Any]
|
||||
) -> Optional[Union[TokenizeChatRequest, TokenizeCompletionRequest]]:
|
||||
"""Build the Tokenize* request for ``payload``.
|
||||
|
||||
KV-aware routing sends each request to one replica, scored on a single
|
||||
prompt's token sequence, so we return ``None`` (the caller falls back to
|
||||
token-less routing) for bodies that don't have exactly one prompt:
|
||||
- A non-string ``prompt``: an OpenAI *batch* completion where ``prompt``
|
||||
is a list, e.g. ``{"prompt": ["q1", "q2"]}`` (or pre-tokenized id
|
||||
lists). N prompts give N token sequences, so there's no single key to
|
||||
route the one request on.
|
||||
|
||||
TODO (jeffreywang): Support multi-prompt tokenization.
|
||||
"""
|
||||
try:
|
||||
if "messages" in payload:
|
||||
# Forward every request field the engine renders the prompt from
|
||||
# so the routing token IDs match the prefill tokens.
|
||||
return TokenizeChatRequest.model_validate(
|
||||
{
|
||||
k: v
|
||||
for k, v in payload.items()
|
||||
if k in TokenizeChatRequest.model_fields
|
||||
}
|
||||
)
|
||||
if "prompt" in payload:
|
||||
if not isinstance(payload["prompt"], str):
|
||||
# TODO (jeffreywang): Multi-prompt (list) tokenization is unsupported;
|
||||
# fall back to token-less routing.
|
||||
return None
|
||||
return TokenizeCompletionRequest.model_validate(
|
||||
{
|
||||
k: v
|
||||
for k, v in payload.items()
|
||||
if k in TokenizeCompletionRequest.model_fields
|
||||
}
|
||||
)
|
||||
# Should be unreachable: LLMRouter only routes bodies with messages
|
||||
# or a prompt (see _parse_routing_payload).
|
||||
logger.warning(
|
||||
"Tokenizer got a payload with neither messages nor prompt; "
|
||||
"falling back to token-less routing."
|
||||
)
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.debug("Unsupported tokenize request, falling back: %s", e)
|
||||
return None
|
||||
@@ -0,0 +1,153 @@
|
||||
"""Shared helpers for OpenAI ingress, reused by the P/D direct-streaming path."""
|
||||
|
||||
import json
|
||||
from typing import AsyncGenerator, List, Tuple, TypeVar, Union
|
||||
|
||||
from ray.llm._internal.serve.core.configs.openai_api_models import (
|
||||
ChatCompletionRequest,
|
||||
ChatCompletionResponse,
|
||||
ChatCompletionStreamResponse,
|
||||
CompletionResponse,
|
||||
CompletionStreamResponse,
|
||||
TranscriptionResponse,
|
||||
TranscriptionStreamResponse,
|
||||
)
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
NON_STREAMING_RESPONSE_TYPES = (
|
||||
ChatCompletionResponse,
|
||||
CompletionResponse,
|
||||
TranscriptionResponse,
|
||||
)
|
||||
|
||||
StreamResponseType = Union[
|
||||
ChatCompletionStreamResponse, CompletionStreamResponse, TranscriptionStreamResponse
|
||||
]
|
||||
BatchedStreamResponseType = List[StreamResponseType]
|
||||
|
||||
|
||||
def _sanitize_chat_completion_request(
|
||||
request: ChatCompletionRequest,
|
||||
) -> ChatCompletionRequest:
|
||||
"""Sanitize ChatCompletionRequest to fix Pydantic ValidatorIterator serialization issue.
|
||||
|
||||
This addresses a known Pydantic bug where fields typed as ``Iterable[...]``
|
||||
on OpenAI message TypedDicts (notably ``content`` on every message variant
|
||||
and ``tool_calls`` on assistant messages) become ValidatorIterator objects
|
||||
that cannot be pickled for Ray remote calls.
|
||||
|
||||
Workaround logic adapted from vLLM (credits: @gcalmettes):
|
||||
- vLLM PR: https://github.com/vllm-project/vllm/pull/9951
|
||||
- Pydantic Issue: https://github.com/pydantic/pydantic/issues/9467
|
||||
- Related Issue: https://github.com/pydantic/pydantic/issues/9541
|
||||
- Official Workaround: https://github.com/pydantic/pydantic/issues/9467#issuecomment-2442097291
|
||||
|
||||
Note: still reproducible on Pydantic 2.12 for the ``Iterable[...]`` arm of
|
||||
a ``Union``, so this sanitizer is required regardless of Pydantic version.
|
||||
"""
|
||||
for i, message in enumerate(request.messages):
|
||||
# SGLang messages are Pydantic BaseModels (no .get()); convert to dicts
|
||||
# so the same logic works for both vLLM (TypedDict) and SGLang.
|
||||
if not isinstance(message, dict):
|
||||
request.messages[i] = message = message.model_dump()
|
||||
|
||||
# `content` is typed `Union[str, Iterable[ContentPart], None]` on every
|
||||
# OpenAI message variant. When the iterable arm matches, Pydantic stores
|
||||
# a non-picklable ValidatorIterator. Materialize it for any role.
|
||||
content_val = message.get("content")
|
||||
if content_val is not None and not isinstance(content_val, str):
|
||||
try:
|
||||
message["content"] = list(content_val)
|
||||
except (TypeError, ValueError) as e:
|
||||
raise ValueError(
|
||||
"Validating message `content` raised an error. Please "
|
||||
"ensure `content` is a string, None, or an iterable of "
|
||||
"content parts."
|
||||
) from e
|
||||
|
||||
if message.get("role") == "assistant":
|
||||
tool_calls_val = message.get("tool_calls")
|
||||
if tool_calls_val is not None:
|
||||
try:
|
||||
message["tool_calls"] = list(tool_calls_val)
|
||||
except (TypeError, ValueError) as e:
|
||||
raise ValueError(
|
||||
"Validating messages' `tool_calls` raised an error. "
|
||||
"Please ensure `tool_calls` are iterable of tool calls."
|
||||
) from e
|
||||
|
||||
return request
|
||||
|
||||
|
||||
def _apply_openai_json_format(
|
||||
response: Union[StreamResponseType, BatchedStreamResponseType],
|
||||
) -> str:
|
||||
"""Converts the stream response to OpenAI format.
|
||||
|
||||
Each model response is converted to the string:
|
||||
data: <response-json1>\n\n
|
||||
|
||||
The converted strings are concatenated and returned:
|
||||
data: <response-json1>\n\ndata: <response-json2>\n\n...
|
||||
"""
|
||||
if isinstance(response, list):
|
||||
first_response = next(iter(response))
|
||||
if isinstance(first_response, str):
|
||||
return "".join(response)
|
||||
if isinstance(first_response, dict):
|
||||
return "".join(f"data: {json.dumps(r)}\n\n" for r in response)
|
||||
if hasattr(first_response, "model_dump_json"):
|
||||
return "".join(f"data: {r.model_dump_json()}\n\n" for r in response)
|
||||
raise ValueError(
|
||||
f"Unexpected response type: {type(first_response)}, {first_response=}"
|
||||
)
|
||||
if hasattr(response, "model_dump_json"):
|
||||
return f"data: {response.model_dump_json()}\n\n"
|
||||
if isinstance(response, str):
|
||||
return response
|
||||
raise ValueError(f"Unexpected response type: {type(response)}, {response=}")
|
||||
|
||||
|
||||
async def _peek_at_generator(
|
||||
gen: AsyncGenerator[T, None],
|
||||
) -> Tuple[T, AsyncGenerator[T, None]]:
|
||||
# Peek at the first element
|
||||
first_item = await gen.__anext__()
|
||||
|
||||
# Create a new generator that yields the peeked item first
|
||||
async def new_generator() -> AsyncGenerator[T, None]:
|
||||
yield first_item
|
||||
async for item in gen:
|
||||
yield item
|
||||
|
||||
return first_item, new_generator()
|
||||
|
||||
|
||||
async def _openai_json_wrapper(
|
||||
generator: AsyncGenerator[
|
||||
Union[StreamResponseType, BatchedStreamResponseType], None
|
||||
],
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""Wrapper that converts stream responses into OpenAI JSON strings.
|
||||
|
||||
Args:
|
||||
generator: an async generator that yields either individual stream responses
|
||||
(StreamResponseType) or batches of stream responses (BatchedStreamResponseType).
|
||||
Each response is converted into OpenAI JSON format and streamed to the client.
|
||||
For batched responses, the items are concatenated together as a single string.
|
||||
|
||||
Yields:
|
||||
String chunks in OpenAI SSE format: "data: {json}\n\n", with a final
|
||||
"data: [DONE]\n\n" to indicate completion. If the upstream generator
|
||||
already yields a "data: [DONE]" sentinel, it is not duplicated.
|
||||
"""
|
||||
done_sent = False
|
||||
async for response in generator:
|
||||
packet = _apply_openai_json_format(response)
|
||||
if packet.strip().endswith("data: [DONE]"):
|
||||
done_sent = True
|
||||
yield packet
|
||||
|
||||
if not done_sent:
|
||||
yield "data: [DONE]\n\n"
|
||||
@@ -0,0 +1,208 @@
|
||||
from dataclasses import dataclass, field
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
AsyncGenerator,
|
||||
Dict,
|
||||
List,
|
||||
Optional,
|
||||
Protocol,
|
||||
Union,
|
||||
)
|
||||
|
||||
from starlette.requests import Request
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
|
||||
from ray.llm._internal.serve.core.configs.openai_api_models import (
|
||||
ChatCompletionRequest,
|
||||
ChatCompletionResponse,
|
||||
CompletionRequest,
|
||||
CompletionResponse,
|
||||
DetokenizeRequest,
|
||||
DetokenizeResponse,
|
||||
ErrorResponse,
|
||||
TokenizeRequest,
|
||||
TokenizeResponse,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RawRequestInfo:
|
||||
"""A serializable representation of important fields from a Starlette Request.
|
||||
|
||||
This dataclass captures key request data that needs to be passed through
|
||||
RPC boundaries (e.g., from ingress to LLMServer). The Starlette Request
|
||||
object itself is not serializable, so we extract the needed fields here.
|
||||
|
||||
Usage:
|
||||
raw_request = RawRequestInfo.from_starlette_request(starlette_request)
|
||||
# Pass raw_request through RPC...
|
||||
starlette_request = raw_request.to_starlette_request()
|
||||
"""
|
||||
|
||||
headers: Dict[str, str] = field(default_factory=dict)
|
||||
|
||||
@classmethod
|
||||
def from_starlette_request(cls, request: Request) -> "RawRequestInfo":
|
||||
"""Create a RawRequestInfo from a Starlette Request object."""
|
||||
return cls(headers=dict(request.headers))
|
||||
|
||||
def to_starlette_request(self) -> Request:
|
||||
"""Create a minimal Starlette Request from this RawRequestInfo."""
|
||||
scope = {
|
||||
"type": "http",
|
||||
"method": "POST",
|
||||
"path": "/",
|
||||
"headers": [
|
||||
(k.lower().encode(), (v or "").encode())
|
||||
for k, v in self.headers.items()
|
||||
],
|
||||
"query_string": b"",
|
||||
}
|
||||
return Request(scope)
|
||||
|
||||
@classmethod
|
||||
def to_starlette_request_optional(
|
||||
cls, raw_request_info: Optional["RawRequestInfo"] = None
|
||||
) -> Optional[Request]:
|
||||
"""Convert RawRequestInfo to Starlette Request, or return None if input is None."""
|
||||
if raw_request_info is not None:
|
||||
return raw_request_info.to_starlette_request()
|
||||
return None
|
||||
|
||||
|
||||
class DeploymentProtocol(Protocol):
|
||||
@classmethod
|
||||
def get_deployment_options(cls, **kwargs) -> Dict[str, Any]:
|
||||
"""Get the default deployment options for the this deployment."""
|
||||
|
||||
|
||||
class LLMServerProtocol(DeploymentProtocol):
|
||||
"""
|
||||
This is the common interface between all the llm deployment. All llm deployments
|
||||
need to implement a sync constructor, an async start method, and check_health method.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""
|
||||
Constructor takes basic setup that doesn't require async operations.
|
||||
"""
|
||||
|
||||
async def start(self) -> None:
|
||||
"""
|
||||
Start the underlying engine. This handles async initialization.
|
||||
"""
|
||||
|
||||
async def chat(
|
||||
self,
|
||||
request: "ChatCompletionRequest",
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[Union[str, "ChatCompletionResponse", "ErrorResponse"], None]:
|
||||
"""
|
||||
Inferencing to the engine for chat, and return the response.
|
||||
"""
|
||||
|
||||
async def completions(
|
||||
self,
|
||||
request: "CompletionRequest",
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[
|
||||
Union[List[Union[str, "ErrorResponse"]], "CompletionResponse"], None
|
||||
]:
|
||||
"""
|
||||
Inferencing to the engine for completion api, and return the response.
|
||||
"""
|
||||
|
||||
async def check_health(self) -> None:
|
||||
"""
|
||||
Check the health of the replica. Does not return anything.
|
||||
Raise error when the engine is dead and needs to be restarted.
|
||||
"""
|
||||
|
||||
async def reset_prefix_cache(self) -> None:
|
||||
"""Reset the prefix cache of the underlying engine"""
|
||||
|
||||
async def start_profile(self) -> None:
|
||||
"""Start profiling"""
|
||||
|
||||
async def stop_profile(self) -> None:
|
||||
"""Stop profiling"""
|
||||
|
||||
async def sleep(self, **kwargs: Any) -> None:
|
||||
"""Put the engine to sleep.
|
||||
|
||||
Args:
|
||||
**kwargs: Engine-specific sleep options. Passed through to the engine.
|
||||
"""
|
||||
|
||||
async def wakeup(self, **kwargs: Any) -> None:
|
||||
"""Wake up the engine from sleep mode.
|
||||
|
||||
Args:
|
||||
**kwargs: Engine-specific wakeup options. Passed through to the engine.
|
||||
"""
|
||||
|
||||
async def is_sleeping(self) -> bool:
|
||||
"""Check whether the engine is currently sleeping.
|
||||
|
||||
Returns:
|
||||
True if the engine is sleeping, False otherwise.
|
||||
"""
|
||||
|
||||
async def pause(self, **kwargs: Any) -> None:
|
||||
"""Pause the engine.
|
||||
|
||||
Args:
|
||||
**kwargs: Engine-specific pause options. Passed through to the engine.
|
||||
"""
|
||||
|
||||
async def resume(self, **kwargs: Any) -> None:
|
||||
"""Resume the engine.
|
||||
|
||||
Args:
|
||||
**kwargs: Engine-specific resume options. Passed through to the engine.
|
||||
"""
|
||||
|
||||
async def is_paused(self) -> bool:
|
||||
"""Check whether the engine is currently paused.
|
||||
|
||||
Returns:
|
||||
True if the engine is paused, False otherwise.
|
||||
"""
|
||||
|
||||
async def tokenize(
|
||||
self,
|
||||
request: "TokenizeRequest",
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[Union["TokenizeResponse", "ErrorResponse"], None]:
|
||||
"""Tokenize the input text.
|
||||
|
||||
Args:
|
||||
request: The tokenize request containing the text to tokenize.
|
||||
raw_request_info: Optional RawRequestInfo containing data from the original
|
||||
HTTP request.
|
||||
|
||||
Returns:
|
||||
An AsyncGenerator yielding TokenizeResponse or ErrorResponse objects.
|
||||
"""
|
||||
|
||||
async def detokenize(
|
||||
self,
|
||||
request: "DetokenizeRequest",
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[Union["DetokenizeResponse", "ErrorResponse"], None]:
|
||||
"""Detokenize the input token IDs.
|
||||
|
||||
Args:
|
||||
request: The detokenize request containing the token IDs.
|
||||
raw_request_info: Optional RawRequestInfo containing data from the original
|
||||
HTTP request.
|
||||
|
||||
Returns:
|
||||
An AsyncGenerator yielding DetokenizeResponse or ErrorResponse objects.
|
||||
"""
|
||||
|
||||
# TODO (Kourosh): This does not belong here.
|
||||
async def llm_config(self) -> Optional["LLMConfig"]:
|
||||
"""Get the LLM config"""
|
||||
@@ -0,0 +1,89 @@
|
||||
import pprint
|
||||
from typing import Optional, Type
|
||||
|
||||
from ray import serve
|
||||
from ray.llm._internal.common.dict_utils import (
|
||||
maybe_apply_llm_deployment_config_defaults,
|
||||
)
|
||||
from ray.llm._internal.serve.constants import (
|
||||
DEFAULT_HEALTH_CHECK_PERIOD_S,
|
||||
DEFAULT_HEALTH_CHECK_TIMEOUT_S,
|
||||
DEFAULT_MAX_ONGOING_REQUESTS,
|
||||
DEFAULT_MAX_TARGET_ONGOING_REQUESTS,
|
||||
)
|
||||
from ray.llm._internal.serve.core.configs.llm_config import (
|
||||
LLMConfig,
|
||||
)
|
||||
from ray.llm._internal.serve.core.server.llm_server import LLMServer
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
from ray.llm._internal.serve.routing_policies.kv_aware.utils import (
|
||||
_maybe_setup_kv_aware_routing,
|
||||
)
|
||||
from ray.serve.deployment import Application
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
DEFAULT_DEPLOYMENT_OPTIONS = {
|
||||
"max_ongoing_requests": DEFAULT_MAX_ONGOING_REQUESTS,
|
||||
"health_check_period_s": DEFAULT_HEALTH_CHECK_PERIOD_S,
|
||||
"health_check_timeout_s": DEFAULT_HEALTH_CHECK_TIMEOUT_S,
|
||||
"autoscaling_config": {
|
||||
"target_ongoing_requests": DEFAULT_MAX_TARGET_ONGOING_REQUESTS,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _get_deployment_name(llm_config: LLMConfig) -> str:
|
||||
return llm_config.model_id.replace("/", "--").replace(".", "_")
|
||||
|
||||
|
||||
def build_llm_deployment(
|
||||
llm_config: LLMConfig,
|
||||
*,
|
||||
name_prefix: Optional[str] = None,
|
||||
bind_kwargs: Optional[dict] = None,
|
||||
override_serve_options: Optional[dict] = None,
|
||||
deployment_cls: Optional[Type[LLMServer]] = None,
|
||||
) -> Application:
|
||||
"""Build an LLMServer deployment.
|
||||
|
||||
Args:
|
||||
llm_config: The LLMConfig to build the deployment.
|
||||
name_prefix: The prefix to add to the deployment name.
|
||||
bind_kwargs: The optional extra kwargs to pass to the deployment.
|
||||
Used for customizing the deployment.
|
||||
override_serve_options: The optional serve options to override the
|
||||
default options.
|
||||
deployment_cls: The deployment class to use. Defaults to LLMServer.
|
||||
|
||||
Returns:
|
||||
The Ray Serve Application for the LLMServer deployment.
|
||||
"""
|
||||
deployment_cls = deployment_cls or llm_config.server_cls or LLMServer
|
||||
name_prefix = name_prefix or f"{deployment_cls.__name__}:"
|
||||
bind_kwargs = bind_kwargs or {}
|
||||
|
||||
deployment_options = deployment_cls.get_deployment_options(llm_config)
|
||||
|
||||
# Set the name of the deployment config to map to the model ID.
|
||||
deployment_name = deployment_options.get("name", _get_deployment_name(llm_config))
|
||||
|
||||
if name_prefix:
|
||||
deployment_options["name"] = name_prefix + deployment_name
|
||||
|
||||
if override_serve_options:
|
||||
deployment_options.update(override_serve_options)
|
||||
|
||||
deployment_options = maybe_apply_llm_deployment_config_defaults(
|
||||
DEFAULT_DEPLOYMENT_OPTIONS, deployment_options
|
||||
)
|
||||
|
||||
_maybe_setup_kv_aware_routing(deployment_options, llm_config)
|
||||
|
||||
logger.info("============== Deployment Options ==============")
|
||||
logger.info(pprint.pformat(deployment_options))
|
||||
|
||||
return serve.deployment(deployment_cls, **deployment_options).bind(
|
||||
llm_config=llm_config, **bind_kwargs
|
||||
)
|
||||
@@ -0,0 +1,791 @@
|
||||
import asyncio
|
||||
import copy
|
||||
import os
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
AsyncGenerator,
|
||||
Dict,
|
||||
List,
|
||||
Optional,
|
||||
Type,
|
||||
TypeVar,
|
||||
Union,
|
||||
)
|
||||
|
||||
import ray
|
||||
from ray import serve
|
||||
from ray._common.usage.usage_lib import TagKey, record_extra_usage_tag
|
||||
from ray._common.utils import import_attr
|
||||
from ray.llm._internal.serve.constants import (
|
||||
ENABLE_WORKER_PROCESS_SETUP_HOOK,
|
||||
ENGINE_START_TIMEOUT_S,
|
||||
MODEL_RESPONSE_BATCH_TIMEOUT_MS,
|
||||
RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING,
|
||||
RAYLLM_VLLM_ENGINE_CLS_ENV,
|
||||
)
|
||||
from ray.llm._internal.serve.core.configs.llm_config import (
|
||||
DiskMultiplexConfig,
|
||||
LLMConfig,
|
||||
)
|
||||
from ray.llm._internal.serve.core.engine.protocol import LLMEngine
|
||||
from ray.llm._internal.serve.core.protocol import LLMServerProtocol, RawRequestInfo
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
from ray.llm._internal.serve.observability.usage_telemetry.usage import (
|
||||
push_telemetry_report_for_all_models,
|
||||
)
|
||||
from ray.llm._internal.serve.utils.batcher import Batcher
|
||||
from ray.llm._internal.serve.utils.lora_serve_utils import (
|
||||
LoraModelLoader,
|
||||
)
|
||||
from ray.llm._internal.serve.utils.server_utils import (
|
||||
get_serve_request_id,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.llm._internal.serve.core.configs.openai_api_models import (
|
||||
ChatCompletionRequest,
|
||||
ChatCompletionResponse,
|
||||
CompletionRequest,
|
||||
CompletionResponse,
|
||||
DetokenizeRequest,
|
||||
DetokenizeResponse,
|
||||
EmbeddingRequest,
|
||||
EmbeddingResponse,
|
||||
ErrorResponse,
|
||||
ScoreRequest,
|
||||
ScoreResponse,
|
||||
TokenizeRequest,
|
||||
TokenizeResponse,
|
||||
TranscriptionRequest,
|
||||
TranscriptionResponse,
|
||||
)
|
||||
|
||||
logger = get_logger(__name__)
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
def _merge_replica_actor_and_child_actor_bundles(
|
||||
child_actor_bundles: List[Dict[str, float]],
|
||||
replica_actor_bundle: Dict[str, float],
|
||||
) -> List[Dict[str, float]]:
|
||||
"""Sum up the bundles from replica actor bundles with the first bundle from child actor bundles.
|
||||
|
||||
This is because the replica actor will use the first bundle in the list, and we want to collocate the replica actor with the child actor.
|
||||
So we need to group them together.
|
||||
|
||||
So for example:
|
||||
child_actor_bundles = [{"GPU": 1, "CPU": 1}, {"GPU": 1, "CPU": 1}]
|
||||
replica_actor_bundle = {"GPU": 0, "CPU": 1, "memory": 100}
|
||||
return [{"GPU": 1, "CPU": 2, "memory": 100}, {"GPU": 1, "CPU": 1}]
|
||||
"""
|
||||
|
||||
if not child_actor_bundles:
|
||||
return [copy.copy(replica_actor_bundle)]
|
||||
|
||||
if not replica_actor_bundle:
|
||||
return [copy.copy(bundle) for bundle in child_actor_bundles]
|
||||
|
||||
original_first_bundle = child_actor_bundles[0]
|
||||
bundle_key_set = set(original_first_bundle.keys()) | set(
|
||||
replica_actor_bundle.keys()
|
||||
)
|
||||
|
||||
merged_first_bundle = {
|
||||
key: original_first_bundle.get(key, 0) + replica_actor_bundle.get(key, 0)
|
||||
for key in bundle_key_set
|
||||
}
|
||||
|
||||
return [merged_first_bundle] + [
|
||||
copy.copy(bundle) for bundle in child_actor_bundles[1:]
|
||||
]
|
||||
|
||||
|
||||
class LLMServer(LLMServerProtocol):
|
||||
"""This is a shim layer to decouple the LLM engine from the ingress
|
||||
deployment.
|
||||
|
||||
It has a very similar API as the engine. Almost all of the abstractions are
|
||||
implemented by the engine. This class just a little bit more logic on top:
|
||||
|
||||
1. Logic for serve multiplexing (e.g. LoRA loading).
|
||||
2. Request id handing from serve context.
|
||||
3. Batching in case of streaming (only for chat and completions).
|
||||
4. Telemetry reporting.
|
||||
|
||||
Usage Patterns:
|
||||
|
||||
1. Basic pattern (for testing):
|
||||
server = LLMServer.sync_init(llm_config) # Sync constructor, unstarted
|
||||
await server.start() # Must explicitly start
|
||||
|
||||
2. Async context (default, used by Ray Serve):
|
||||
server = await LLMServer(llm_config) # Async constructor, fully started
|
||||
|
||||
3. Ray Serve deployment:
|
||||
# Ray Serve calls the async constructor directly
|
||||
deployment = serve.deployment(LLMServer).bind(llm_config)
|
||||
"""
|
||||
|
||||
_default_engine_cls = None
|
||||
|
||||
async def __init__(
|
||||
self,
|
||||
llm_config: LLMConfig,
|
||||
*,
|
||||
engine_cls: Optional[Type[LLMEngine]] = None,
|
||||
model_downloader: Optional[Type[LoraModelLoader]] = None,
|
||||
):
|
||||
"""Asynchronous constructor that returns a fully started instance.
|
||||
|
||||
This is the default constructor used by Ray Serve deployments.
|
||||
|
||||
Args:
|
||||
llm_config: LLMConfig for the model.
|
||||
engine_cls: Dependency injection for the vllm engine class.
|
||||
Defaults to `VLLMEngine`.
|
||||
model_downloader: Dependency injection for the model downloader.
|
||||
Defaults to `LoraModelLoader`.
|
||||
"""
|
||||
super().__init__()
|
||||
self._init_shared(llm_config, engine_cls, model_downloader)
|
||||
await self.start()
|
||||
|
||||
def _init_shared(
|
||||
self,
|
||||
llm_config: LLMConfig,
|
||||
engine_cls: Optional[Type[LLMEngine]] = None,
|
||||
model_downloader: Optional[Type[LoraModelLoader]] = None,
|
||||
):
|
||||
"""Shared initialization logic between constructors."""
|
||||
self._llm_config = llm_config
|
||||
self._engine_cls = engine_cls or self._get_default_engine_class()
|
||||
self.engine: Optional[LLMEngine] = None
|
||||
self._init_multiplex_loader(model_downloader)
|
||||
|
||||
@classmethod
|
||||
def sync_init(
|
||||
cls,
|
||||
llm_config: LLMConfig,
|
||||
*,
|
||||
engine_cls: Optional[Type[LLMEngine]] = None,
|
||||
model_downloader: Optional[Type[LoraModelLoader]] = None,
|
||||
) -> "LLMServer":
|
||||
"""Synchronous constructor that returns an unstarted instance.
|
||||
|
||||
This is used for testing the new pattern where initialization
|
||||
and starting are explicitly separated.
|
||||
|
||||
Args:
|
||||
llm_config: LLMConfig for the model.
|
||||
engine_cls: Dependency injection for the vllm engine class.
|
||||
Defaults to `VLLMEngine`.
|
||||
model_downloader: Dependency injection for the model downloader.
|
||||
Defaults to `LoraModelLoader`.
|
||||
|
||||
Returns:
|
||||
An unstarted LLMServer instance. Caller must call await start().
|
||||
"""
|
||||
instance = cls.__new__(cls)
|
||||
LLMServerProtocol.__init__(instance)
|
||||
instance._init_shared(llm_config, engine_cls, model_downloader)
|
||||
return instance
|
||||
|
||||
async def start(self):
|
||||
"""Start the underlying engine. This handles async initialization."""
|
||||
if self._engine_cls is not None:
|
||||
self.engine = self._engine_cls(self._llm_config)
|
||||
await asyncio.wait_for(self._start_engine(), timeout=ENGINE_START_TIMEOUT_S)
|
||||
|
||||
async def __serve_build_asgi_app__(self):
|
||||
from fastapi import HTTPException
|
||||
|
||||
from ray.llm._internal.serve.core.configs.openai_api_models import (
|
||||
ModelCard,
|
||||
to_model_metadata,
|
||||
)
|
||||
|
||||
app = await self.engine.build_asgi_app()
|
||||
|
||||
# vLLM's native ASGI app only exposes `GET /v1/models` (list); add
|
||||
# `GET /v1/models/{id}` so direct-streaming clients can call
|
||||
# `openai_client.models.retrieve(...)` like the OpenAiIngress path.
|
||||
model_id = self._llm_config.model_id
|
||||
model_card = to_model_metadata(model_id, self._llm_config)
|
||||
|
||||
@app.get("/v1/models/{model:path}", response_model=ModelCard)
|
||||
async def _get_model(model: str):
|
||||
if model != model_id:
|
||||
raise HTTPException(status_code=404, detail=f"Unknown model: {model}")
|
||||
return model_card
|
||||
|
||||
return app
|
||||
|
||||
def _init_multiplex_loader(
|
||||
self, model_downloader_cls: Optional[Type[LoraModelLoader]] = None
|
||||
):
|
||||
"""Initialize the multiplex loader."""
|
||||
|
||||
model_downloader_cls = model_downloader_cls or LoraModelLoader
|
||||
mx_config = self._llm_config.multiplex_config()
|
||||
|
||||
if mx_config is not None:
|
||||
model_downloader = model_downloader_cls(
|
||||
download_timeout_s=mx_config.download_timeout_s,
|
||||
max_tries=mx_config.max_download_tries,
|
||||
)
|
||||
|
||||
async def _load_model(lora_model_id: str) -> DiskMultiplexConfig:
|
||||
return await model_downloader.load_model_from_config(
|
||||
lora_model_id=lora_model_id,
|
||||
llm_config=self._llm_config,
|
||||
)
|
||||
|
||||
self._load_model = serve.multiplexed(
|
||||
max_num_models_per_replica=mx_config.max_num_models_per_replica
|
||||
)(_load_model)
|
||||
else:
|
||||
|
||||
async def _load_model(lora_model_id: str) -> DiskMultiplexConfig:
|
||||
raise ValueError("LoRA config is not set in the LLMConfig")
|
||||
|
||||
self._load_model = _load_model
|
||||
|
||||
def _get_default_engine_class(self) -> Type[LLMEngine]:
|
||||
"""Helper to load the engine class from the environment variable.
|
||||
This is used for testing or escape-hatch for patching purposes.
|
||||
If env variable is not set, it will fallback to the default engine class
|
||||
(VLLMEngine, imported lazily to avoid a hard module-level dependency).
|
||||
"""
|
||||
engine_cls_path = os.environ.get(RAYLLM_VLLM_ENGINE_CLS_ENV)
|
||||
if engine_cls_path:
|
||||
return import_attr(engine_cls_path)
|
||||
if self._default_engine_cls is not None:
|
||||
return self._default_engine_cls
|
||||
from ray.llm._internal.serve.engines.vllm.vllm_engine import VLLMEngine
|
||||
|
||||
return VLLMEngine
|
||||
|
||||
async def _start_engine(self):
|
||||
if self.engine is None:
|
||||
raise ValueError("Engine is not set")
|
||||
|
||||
await self.engine.start()
|
||||
|
||||
# Push telemetry reports for the model in the current deployment.
|
||||
push_telemetry_report_for_all_models(all_models=[self._llm_config])
|
||||
if RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING:
|
||||
# Cluster-wide adoption signal: written from each replica on engine
|
||||
# start, but last-write-wins so it reports one value per cluster.
|
||||
record_extra_usage_tag(TagKey.LLM_SERVE_DIRECT_STREAMING_ENABLED, "1")
|
||||
|
||||
def _get_batch_interval_ms(self, stream: bool = True) -> int:
|
||||
"""Calculate the batching interval for responses."""
|
||||
stream_batching_interval_ms = self._llm_config.experimental_configs.get(
|
||||
"stream_batching_interval_ms"
|
||||
)
|
||||
if stream_batching_interval_ms is None:
|
||||
stream_batching_interval_ms = MODEL_RESPONSE_BATCH_TIMEOUT_MS
|
||||
return stream_batching_interval_ms if stream else None
|
||||
|
||||
async def _maybe_add_request_id_to_request(
|
||||
self,
|
||||
request: Union[
|
||||
"ChatCompletionRequest",
|
||||
"CompletionRequest",
|
||||
"EmbeddingRequest",
|
||||
"TranscriptionRequest",
|
||||
],
|
||||
):
|
||||
"""Stamp the Serve request id, unless the caller set request_id explicitly.
|
||||
|
||||
request_id defaults to a random uuid (never None), so use model_fields_set
|
||||
to avoid clobbering an id a caller deliberately set (e.g. a P/D connector's
|
||||
coordination id). Some request types (tokenize/detokenize) have no
|
||||
request_id field at all -- skip those.
|
||||
"""
|
||||
if not hasattr(request, "request_id"):
|
||||
return
|
||||
if "request_id" in request.model_fields_set:
|
||||
return
|
||||
request_id = get_serve_request_id()
|
||||
if request_id:
|
||||
request.request_id = request_id
|
||||
|
||||
async def _maybe_resolve_lora_from_multiplex(self) -> None:
|
||||
"""Handle the lora model for the request."""
|
||||
multiplexed_model_id = serve.get_multiplexed_model_id()
|
||||
if multiplexed_model_id:
|
||||
if self._llm_config.lora_config is None:
|
||||
raise ValueError("Must setup lora config for multiplexed requests.")
|
||||
disk_lora_model = await self._load_model(multiplexed_model_id)
|
||||
await self.engine.resolve_lora(disk_lora_model)
|
||||
|
||||
def _batch_output_stream(
|
||||
self, generator: AsyncGenerator[T, None]
|
||||
) -> AsyncGenerator[List[T], None]:
|
||||
return Batcher(
|
||||
generator,
|
||||
interval_ms=self._get_batch_interval_ms(),
|
||||
).stream()
|
||||
|
||||
async def _run_request(
|
||||
self,
|
||||
request: Union[
|
||||
"ChatCompletionRequest",
|
||||
"CompletionRequest",
|
||||
"EmbeddingRequest",
|
||||
"TranscriptionRequest",
|
||||
"ScoreRequest",
|
||||
],
|
||||
*,
|
||||
engine_method: str,
|
||||
batch_output_stream: bool = False,
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[Any, None]:
|
||||
"""Run the engine method on the request + perform batching when stream=True.
|
||||
|
||||
Args:
|
||||
request: The request to run.
|
||||
engine_method: The method to call on the engine.
|
||||
batch_output_stream: Whether to batch the output stream.
|
||||
raw_request_info: Optional RawRequestInfo containing data from the original
|
||||
HTTP request.
|
||||
|
||||
Returns:
|
||||
An AsyncGenerator of the response. If stream is True and batching is enabled, then the generator will yield a list of streaming responses (strings of the format data: {response_json}\n\n). Otherwise, it will yield the non-streaming response from engine directly.
|
||||
"""
|
||||
|
||||
await self._maybe_add_request_id_to_request(request)
|
||||
await self._maybe_resolve_lora_from_multiplex()
|
||||
|
||||
is_stream = hasattr(request, "stream") and request.stream
|
||||
engine_stream = getattr(self.engine, engine_method)(request, raw_request_info)
|
||||
|
||||
if is_stream and batch_output_stream:
|
||||
stream = self._batch_output_stream(engine_stream)
|
||||
else:
|
||||
stream = engine_stream
|
||||
|
||||
return stream
|
||||
|
||||
async def chat(
|
||||
self,
|
||||
request: "ChatCompletionRequest",
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[
|
||||
Union[List[Union[str, "ErrorResponse"]], "ChatCompletionResponse"], None
|
||||
]:
|
||||
"""Runs a chat request to the LLM engine and returns the response.
|
||||
|
||||
Args:
|
||||
request: A ChatCompletionRequest object.
|
||||
raw_request_info: Optional RawRequestInfo containing data from the original
|
||||
HTTP request.
|
||||
|
||||
Returns:
|
||||
An AsyncGenerator of the response. If stream is True and batching
|
||||
is enabled, then the generator will yield a list of chat streaming
|
||||
responses (strings of the format data: {response_json}\\n\\n).
|
||||
Otherwise, it will yield the ChatCompletionResponse object directly.
|
||||
"""
|
||||
return await self._run_request(
|
||||
request,
|
||||
engine_method="chat",
|
||||
batch_output_stream=True,
|
||||
raw_request_info=raw_request_info,
|
||||
)
|
||||
|
||||
async def completions(
|
||||
self,
|
||||
request: "CompletionRequest",
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[
|
||||
Union[List[Union[str, "ErrorResponse"]], "CompletionResponse"], None
|
||||
]:
|
||||
"""Runs a completion request to the LLM engine and returns the response.
|
||||
|
||||
Args:
|
||||
request: A CompletionRequest object.
|
||||
raw_request_info: Optional RawRequestInfo containing data from the original
|
||||
HTTP request.
|
||||
|
||||
Returns:
|
||||
An AsyncGenerator of the response. If stream is True and batching
|
||||
is enabled, then the generator will yield a list of completion
|
||||
streaming responses (strings of the format data: {response_json}\\n\\n).
|
||||
Otherwise, it will yield the CompletionResponse object directly.
|
||||
"""
|
||||
return await self._run_request(
|
||||
request,
|
||||
engine_method="completions",
|
||||
batch_output_stream=True,
|
||||
raw_request_info=raw_request_info,
|
||||
)
|
||||
|
||||
async def embeddings(
|
||||
self,
|
||||
request: "EmbeddingRequest",
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[Union[List["ErrorResponse"], "EmbeddingResponse"], None]:
|
||||
"""Runs an embeddings request to the engine and returns the response.
|
||||
|
||||
Returns an AsyncGenerator over the EmbeddingResponse object. This is so that the caller can have a consistent interface across all the methods of chat, completions, embeddings and transcriptions.
|
||||
|
||||
Args:
|
||||
request: An EmbeddingRequest object.
|
||||
raw_request_info: Optional RawRequestInfo containing data from the original
|
||||
HTTP request.
|
||||
|
||||
Returns:
|
||||
An AsyncGenerator over the EmbeddingResponse object.
|
||||
"""
|
||||
# NOTE: Embeddings does not need batching.
|
||||
return await self._run_request(
|
||||
request,
|
||||
engine_method="embeddings",
|
||||
batch_output_stream=False,
|
||||
raw_request_info=raw_request_info,
|
||||
)
|
||||
|
||||
async def transcriptions(
|
||||
self,
|
||||
request: "TranscriptionRequest",
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[
|
||||
Union[List[Union[str, "ErrorResponse"]], "TranscriptionResponse"], None
|
||||
]:
|
||||
"""Runs an transcriptions request to the engine and returns the response.
|
||||
|
||||
Returns an AsyncGenerator over the TranscriptionResponse object. This is so that the caller can have a consistent interface across all the methods of chat, completions, embeddings and transcriptions.
|
||||
|
||||
Args:
|
||||
request: A TranscriptionRequest object.
|
||||
raw_request_info: Optional RawRequestInfo containing data from the original
|
||||
HTTP request.
|
||||
|
||||
Returns:
|
||||
An AsyncGenerator over the TranscriptionResponse object.
|
||||
"""
|
||||
return await self._run_request(
|
||||
request,
|
||||
engine_method="transcriptions",
|
||||
batch_output_stream=True,
|
||||
raw_request_info=raw_request_info,
|
||||
)
|
||||
|
||||
async def score(
|
||||
self,
|
||||
request: "ScoreRequest",
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[Union["ScoreResponse", "ErrorResponse"], None]:
|
||||
"""Runs a score request to the engine and returns the response.
|
||||
|
||||
Returns an AsyncGenerator over the ScoreResponse object. This is so that the caller can have a consistent interface across all the methods of chat, completions, embeddings, and score.
|
||||
|
||||
Args:
|
||||
request: A ScoreRequest object.
|
||||
raw_request_info: Optional RawRequestInfo containing data from the original
|
||||
HTTP request.
|
||||
|
||||
Returns:
|
||||
An AsyncGenerator over the ScoreResponse object.
|
||||
"""
|
||||
# NOTE: Score does not need batching, similar to embeddings.
|
||||
return await self._run_request(
|
||||
request,
|
||||
engine_method="score",
|
||||
batch_output_stream=False,
|
||||
raw_request_info=raw_request_info,
|
||||
)
|
||||
|
||||
async def tokenize(
|
||||
self,
|
||||
request: "TokenizeRequest",
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[Union["TokenizeResponse", "ErrorResponse"], None]:
|
||||
"""Tokenize the input text.
|
||||
|
||||
Args:
|
||||
request: A TokenizeRequest object (TokenizeCompletionRequest or TokenizeChatRequest).
|
||||
raw_request_info: Optional RawRequestInfo containing data from the original
|
||||
HTTP request.
|
||||
|
||||
Returns:
|
||||
An AsyncGenerator over the TokenizeResponse object.
|
||||
"""
|
||||
# NOTE: Tokenize does not need batching.
|
||||
return await self._run_request(
|
||||
request,
|
||||
engine_method="tokenize",
|
||||
batch_output_stream=False,
|
||||
raw_request_info=raw_request_info,
|
||||
)
|
||||
|
||||
async def detokenize(
|
||||
self,
|
||||
request: "DetokenizeRequest",
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[Union["DetokenizeResponse", "ErrorResponse"], None]:
|
||||
"""Detokenize the input token IDs.
|
||||
|
||||
Args:
|
||||
request: A DetokenizeRequest object.
|
||||
raw_request_info: Optional RawRequestInfo containing data from the original
|
||||
HTTP request.
|
||||
|
||||
Returns:
|
||||
An AsyncGenerator over the DetokenizeResponse object.
|
||||
"""
|
||||
# NOTE: Detokenize does not need batching.
|
||||
return await self._run_request(
|
||||
request,
|
||||
engine_method="detokenize",
|
||||
batch_output_stream=False,
|
||||
raw_request_info=raw_request_info,
|
||||
)
|
||||
|
||||
async def check_health(self) -> None:
|
||||
"""
|
||||
Check the health of the replica. Does not return anything. Raise error when
|
||||
the engine is dead and needs to be restarted.
|
||||
"""
|
||||
if self.engine is None:
|
||||
return
|
||||
try:
|
||||
return await self.engine.check_health()
|
||||
except Exception as e:
|
||||
logger.error("Engine health check failed in LLMServer.check_health: %s", e)
|
||||
raise e
|
||||
|
||||
async def record_routing_stats(self) -> Dict[str, Any]:
|
||||
"""Serve request-router hook, polled by the controller.
|
||||
|
||||
Surfaces this replica's routing stats (the engine's KV-events endpoint
|
||||
for KV-aware routing); the deployment's ``KVRouterActor`` reads them off
|
||||
the ``LongPoll`` replica snapshot to register the worker.
|
||||
"""
|
||||
if self.engine is None:
|
||||
return {}
|
||||
return self.engine.routing_stats()
|
||||
|
||||
async def sleep(self, **kwargs: Any) -> None:
|
||||
"""Put the engine to sleep.
|
||||
|
||||
Args:
|
||||
**kwargs: Engine-specific sleep options. Passed through to the engine.
|
||||
"""
|
||||
if self.engine is None:
|
||||
return
|
||||
try:
|
||||
await self.engine.sleep(**kwargs)
|
||||
except Exception as e:
|
||||
logger.error("Engine sleep failed in LLMServer.sleep: %s", e)
|
||||
raise e
|
||||
|
||||
async def wakeup(self, **kwargs: Any) -> None:
|
||||
"""Wake up the engine from sleep mode.
|
||||
|
||||
Args:
|
||||
**kwargs: Engine-specific wakeup options. Passed through to the engine.
|
||||
"""
|
||||
if self.engine is None:
|
||||
return
|
||||
try:
|
||||
await self.engine.wakeup(**kwargs)
|
||||
except Exception as e:
|
||||
logger.error("Engine wakeup failed in LLMServer.wakeup: %s", e)
|
||||
raise e
|
||||
|
||||
async def is_sleeping(self) -> bool:
|
||||
"""Check whether the engine is currently sleeping.
|
||||
|
||||
Returns:
|
||||
True if the engine is sleeping, False otherwise.
|
||||
"""
|
||||
if self.engine is None:
|
||||
return False
|
||||
try:
|
||||
return await self.engine.is_sleeping()
|
||||
except Exception as e:
|
||||
logger.error("Engine is_sleeping failed in LLMServer.is_sleeping: %s", e)
|
||||
raise e
|
||||
|
||||
async def reset_prefix_cache(self) -> None:
|
||||
"""Reset the KV prefix cache on the engine.
|
||||
|
||||
Clears cached key-value pairs from previous requests.
|
||||
"""
|
||||
if self.engine is None:
|
||||
return
|
||||
try:
|
||||
await self.engine.reset_prefix_cache()
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Engine reset_prefix_cache failed in LLMServer.reset_prefix_cache: %s",
|
||||
e,
|
||||
)
|
||||
raise e
|
||||
|
||||
async def pause(self, **kwargs: Any) -> None:
|
||||
"""Pause generation on the engine.
|
||||
|
||||
This halts generation requests while keeping model weights
|
||||
in GPU memory. New requests are blocked until resume is called.
|
||||
|
||||
Args:
|
||||
**kwargs: Engine-specific pause options. Passed through to the engine.
|
||||
"""
|
||||
if self.engine is None:
|
||||
return
|
||||
try:
|
||||
await self.engine.pause(**kwargs)
|
||||
except Exception as e:
|
||||
logger.error("Engine pause failed in LLMServer.pause: %s", e)
|
||||
raise e
|
||||
|
||||
async def resume(self, **kwargs: Any) -> None:
|
||||
"""Resume generation on the engine after pause.
|
||||
|
||||
Args:
|
||||
**kwargs: Engine-specific resume options. Passed through to the engine.
|
||||
"""
|
||||
if self.engine is None:
|
||||
return
|
||||
try:
|
||||
await self.engine.resume(**kwargs)
|
||||
except Exception as e:
|
||||
logger.error("Engine resume failed in LLMServer.resume: %s", e)
|
||||
raise e
|
||||
|
||||
async def is_paused(self) -> bool:
|
||||
"""Check whether the engine is currently paused.
|
||||
|
||||
Returns:
|
||||
True if the engine is paused, False otherwise.
|
||||
"""
|
||||
if self.engine is None:
|
||||
return False
|
||||
try:
|
||||
return await self.engine.is_paused()
|
||||
except Exception as e:
|
||||
logger.error("Engine is_paused failed in LLMServer.is_paused: %s", e)
|
||||
raise e
|
||||
|
||||
async def start_profile(self) -> None:
|
||||
"""Start profiling"""
|
||||
if self.engine is None:
|
||||
return
|
||||
try:
|
||||
await self.engine.start_profile()
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Engine start profile failed in LLMServer.start_profile: %s", e
|
||||
)
|
||||
raise e
|
||||
|
||||
async def stop_profile(self) -> None:
|
||||
"""Stop profiling"""
|
||||
if self.engine is None:
|
||||
return
|
||||
try:
|
||||
await self.engine.stop_profile()
|
||||
except Exception as e:
|
||||
logger.error("Engine stop profile failed in LLMServer.stop_profile: %s", e)
|
||||
raise e
|
||||
|
||||
async def collective_rpc(
|
||||
self,
|
||||
method: str,
|
||||
timeout: Optional[float] = None,
|
||||
args: tuple = (),
|
||||
kwargs: Optional[dict] = None,
|
||||
) -> list:
|
||||
"""Execute a collective RPC call on all workers.
|
||||
|
||||
This is used for RLHF workflows where a trainer needs to execute
|
||||
methods on all TP/PP workers (e.g., for weight synchronization).
|
||||
|
||||
Args:
|
||||
method: Name of the worker method to execute.
|
||||
timeout: Maximum time in seconds to wait for execution.
|
||||
args: Positional arguments to pass to the worker method.
|
||||
kwargs: Keyword arguments to pass to the worker method.
|
||||
|
||||
Returns:
|
||||
A list containing the results from each worker.
|
||||
"""
|
||||
if self.engine is None:
|
||||
return []
|
||||
try:
|
||||
return await self.engine.collective_rpc(
|
||||
method=method,
|
||||
timeout=timeout,
|
||||
args=args,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Engine collective_rpc failed in LLMServer.collective_rpc: %s", e
|
||||
)
|
||||
raise e
|
||||
|
||||
async def llm_config(self) -> Optional[LLMConfig]:
|
||||
return self._llm_config
|
||||
|
||||
@classmethod
|
||||
def get_deployment_options(cls, llm_config: "LLMConfig"):
|
||||
engine_config = llm_config.get_engine_config()
|
||||
deployment_options = copy.deepcopy(llm_config.deployment_config)
|
||||
|
||||
if (
|
||||
"placement_group_bundles" in llm_config.deployment_config
|
||||
or "placement_group_strategy" in llm_config.deployment_config
|
||||
):
|
||||
raise ValueError(
|
||||
"placement_group_bundles and placement_group_strategy must not be specified in deployment_config. You can override the default values by setting the `placement_group_config` in the LLMConfig."
|
||||
)
|
||||
|
||||
# Handle the ray_actor_options that could be passed in to
|
||||
# deployment_options
|
||||
ray_actor_options = deployment_options.get("ray_actor_options", {})
|
||||
|
||||
if not engine_config.accelerator.requires_deferred_placement_group:
|
||||
replica_actor_resources = {
|
||||
"CPU": ray_actor_options.get("num_cpus", 1),
|
||||
"GPU": ray_actor_options.get("num_gpus", 0),
|
||||
**ray_actor_options.get("resources", {}),
|
||||
}
|
||||
if "memory" in ray_actor_options:
|
||||
replica_actor_resources["memory"] = ray_actor_options["memory"]
|
||||
|
||||
# TODO: Move this _merge_replica_actor_and_child_actor_bundles to a
|
||||
# more generic place.
|
||||
pg_bundles = _merge_replica_actor_and_child_actor_bundles(
|
||||
engine_config.placement_bundles, replica_actor_resources
|
||||
)
|
||||
|
||||
deployment_options.update(
|
||||
{
|
||||
"placement_group_bundles": pg_bundles,
|
||||
"placement_group_strategy": engine_config.placement_strategy,
|
||||
}
|
||||
)
|
||||
|
||||
# Handle env vars from runtime_env
|
||||
default_runtime_env = ray.get_runtime_context().runtime_env
|
||||
if ENABLE_WORKER_PROCESS_SETUP_HOOK:
|
||||
default_runtime_env[
|
||||
"worker_process_setup_hook"
|
||||
] = "ray.llm._internal.serve._worker_process_setup_hook"
|
||||
|
||||
ray_actor_options = deployment_options.get("ray_actor_options", {})
|
||||
ray_actor_options["runtime_env"] = {
|
||||
**default_runtime_env,
|
||||
# Existing runtime_env should take precedence over the default.
|
||||
**ray_actor_options.get("runtime_env", {}),
|
||||
**(llm_config.runtime_env if llm_config.runtime_env else {}),
|
||||
}
|
||||
deployment_options["ray_actor_options"] = ray_actor_options
|
||||
|
||||
return deployment_options
|
||||
@@ -0,0 +1,3 @@
|
||||
from ray.llm._internal.serve.engines.sglang.sglang_engine import SGLangServer
|
||||
|
||||
__all__ = ["SGLangServer"]
|
||||
@@ -0,0 +1,773 @@
|
||||
"""SGLang engine integration for Ray Serve LLM.
|
||||
|
||||
Provides ``SGLangServer``, a custom server class that wraps SGLang's
|
||||
in-process engine and exposes chat, completions, embeddings, tokenize,
|
||||
and detokenize endpoints through the standard Ray Serve LLM protocol.
|
||||
|
||||
Community SGLang support is in early development. Track progress and
|
||||
provide feedback at https://github.com/ray-project/ray/issues/61114.
|
||||
"""
|
||||
|
||||
import copy
|
||||
import json
|
||||
import signal
|
||||
import time
|
||||
import uuid
|
||||
from typing import (
|
||||
Any,
|
||||
AsyncGenerator,
|
||||
List,
|
||||
Literal,
|
||||
Optional,
|
||||
Union,
|
||||
)
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from ray.llm._internal.serve.constants import ENABLE_WORKER_PROCESS_SETUP_HOOK
|
||||
from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
|
||||
from ray.llm._internal.serve.core.configs.openai_api_models import (
|
||||
ChatCompletionRequest,
|
||||
ChatCompletionResponse,
|
||||
CompletionRequest,
|
||||
CompletionResponse,
|
||||
DetokenizeRequest,
|
||||
DetokenizeResponse,
|
||||
EmbeddingCompletionRequest,
|
||||
EmbeddingRequest,
|
||||
EmbeddingResponse,
|
||||
TokenizeCompletionRequest,
|
||||
TokenizeRequest,
|
||||
TokenizeResponse,
|
||||
)
|
||||
from ray.llm._internal.serve.core.protocol import RawRequestInfo
|
||||
from ray.llm._internal.serve.core.server.llm_server import (
|
||||
_merge_replica_actor_and_child_actor_bundles,
|
||||
)
|
||||
|
||||
|
||||
class SGLangPauseConfig(BaseModel):
|
||||
"""SGLang-specific configuration for pause operation."""
|
||||
|
||||
mode: Literal["abort", "in_place", "retract"] = "abort"
|
||||
"""Pause mode:
|
||||
- "abort" (default): Terminate all in-flight requests immediately.
|
||||
- "in_place": Freeze requests in queue, preserve kv cache.
|
||||
- "retract": Freeze requests in queue, free corresponding KV cache.
|
||||
"""
|
||||
|
||||
|
||||
class SGLangSleepConfig(BaseModel):
|
||||
"""SGLang-specific configuration for sleep operation"""
|
||||
|
||||
tags: Optional[List[Literal["kv_cache", "weights", "cuda_graph"]]] = None
|
||||
|
||||
"""Sleep tags:
|
||||
- "kv_cache": Discard KV cache
|
||||
- "weights": Offload to CPU RAM
|
||||
- "cuda_graph": Discard CUDA graph
|
||||
- None: Discard/Offload everything
|
||||
"""
|
||||
|
||||
|
||||
class SGLangWakeupConfig(BaseModel):
|
||||
"""SGLang-specific configuration for wakeup operation"""
|
||||
|
||||
tags: Optional[List[Literal["kv_cache", "weights", "cuda_graph"]]] = None
|
||||
"""Optional tags to selectively wake up components:
|
||||
- "kv_cache": Restore KV cache only
|
||||
- "weights": Restore weights only
|
||||
- "cuda_graph": Restore CUDA graph only
|
||||
- None: Restore everything
|
||||
"""
|
||||
|
||||
|
||||
_SLEEP_TAGS: frozenset[str] = frozenset({"kv_cache", "weights", "cuda_graph"})
|
||||
|
||||
|
||||
class SGLangServer:
|
||||
def __init__(self, llm_config: LLMConfig):
|
||||
|
||||
self._llm_config = llm_config
|
||||
self.engine_kwargs = llm_config.engine_kwargs
|
||||
self._is_paused = False
|
||||
self._sleeping_tags: set[str] = set()
|
||||
|
||||
try:
|
||||
import sglang
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"SGLang is not installed or failed to import. Please run "
|
||||
"`pip install sglang[all]` to install required dependencies."
|
||||
) from e
|
||||
|
||||
# TODO(issue-61108): remove this once sglang#18752 is merged and included
|
||||
# in the minimum supported SGLang version for this example.
|
||||
original_signal_func = signal.signal
|
||||
|
||||
def noop_signal_handler(sig, action):
|
||||
# Returns default handler to satisfy signal.signal() return signature
|
||||
return signal.SIG_DFL
|
||||
|
||||
try:
|
||||
# Override signal.signal with our no-op function
|
||||
signal.signal = noop_signal_handler
|
||||
self.engine = sglang.Engine(**self.engine_kwargs)
|
||||
finally:
|
||||
signal.signal = original_signal_func
|
||||
|
||||
@staticmethod
|
||||
def _build_sampling_params(request: Any) -> dict[str, Any]:
|
||||
sampling_params: dict[str, Any] = {}
|
||||
model_fields_set = getattr(request, "model_fields_set", None)
|
||||
has_model_fields_set = model_fields_set is not None
|
||||
fields_set = set(model_fields_set) if has_model_fields_set else set()
|
||||
|
||||
def was_explicitly_set(field_name: str) -> bool:
|
||||
# Use model_fields_set when available to avoid injecting defaults for
|
||||
# fields omitted by the caller.
|
||||
if has_model_fields_set:
|
||||
return field_name in fields_set
|
||||
return getattr(request, field_name, None) is not None
|
||||
|
||||
temperature = getattr(request, "temperature", None)
|
||||
top_p = getattr(request, "top_p", None)
|
||||
max_tokens = getattr(request, "max_tokens", None)
|
||||
stop = getattr(request, "stop", None)
|
||||
|
||||
if was_explicitly_set("temperature") and temperature is not None:
|
||||
sampling_params["temperature"] = temperature
|
||||
if was_explicitly_set("top_p") and top_p is not None:
|
||||
sampling_params["top_p"] = top_p
|
||||
if was_explicitly_set("max_tokens") and max_tokens is not None:
|
||||
sampling_params["max_new_tokens"] = max_tokens
|
||||
if was_explicitly_set("stop") and stop is not None:
|
||||
sampling_params["stop"] = stop
|
||||
|
||||
return sampling_params
|
||||
|
||||
@staticmethod
|
||||
def _parse_finish_reason(finish_reason_info: Any) -> str:
|
||||
"""Parse finish_reason from SGLang metadata."""
|
||||
if isinstance(finish_reason_info, dict):
|
||||
return finish_reason_info.get("type", "length")
|
||||
return str(finish_reason_info)
|
||||
|
||||
@staticmethod
|
||||
def _build_chat_messages(messages: List[Any]) -> List[dict[str, Any]]:
|
||||
converted_messages: List[dict[str, Any]] = []
|
||||
for message in messages:
|
||||
if isinstance(message, dict):
|
||||
message_dict = dict(message)
|
||||
elif hasattr(message, "model_dump") and callable(message.model_dump):
|
||||
message_dict = dict(message.model_dump())
|
||||
else:
|
||||
message_dict = {
|
||||
"role": getattr(message, "role", "user"),
|
||||
"content": getattr(message, "content", ""),
|
||||
}
|
||||
|
||||
message_dict["role"] = str(message_dict.get("role", "user"))
|
||||
converted_messages.append(message_dict)
|
||||
return converted_messages
|
||||
|
||||
@staticmethod
|
||||
def _build_chat_template_kwargs(request: Any) -> dict[str, Any]:
|
||||
"""
|
||||
Build optional chat-template kwargs using request fields when present.
|
||||
This mirrors SGLang's chat-serving pipeline semantics without directly
|
||||
coupling to its internal server classes.
|
||||
|
||||
Works with both ChatCompletionRequest and TokenizeChatRequest since
|
||||
both expose tools and chat_template_kwargs fields.
|
||||
"""
|
||||
kwargs: dict[str, Any] = {}
|
||||
|
||||
tools = getattr(request, "tools", None)
|
||||
if tools is not None:
|
||||
kwargs["tools"] = tools
|
||||
|
||||
reasoning_effort = getattr(request, "reasoning_effort", None)
|
||||
if reasoning_effort is not None:
|
||||
kwargs["reasoning_effort"] = reasoning_effort
|
||||
|
||||
chat_template_kwargs = getattr(request, "chat_template_kwargs", None)
|
||||
if isinstance(chat_template_kwargs, dict):
|
||||
kwargs.update(chat_template_kwargs)
|
||||
|
||||
return kwargs
|
||||
|
||||
def _render_chat_prompt(
|
||||
self,
|
||||
messages: List[dict[str, Any]],
|
||||
add_generation_prompt: bool = True,
|
||||
template_kwargs: Optional[dict[str, Any]] = None,
|
||||
) -> str:
|
||||
tokenizer = self.engine.tokenizer_manager.tokenizer
|
||||
# SGLang supports --skip-tokenizer-init, where tokenizer is intentionally
|
||||
# None and text prompt rendering is not available.
|
||||
if tokenizer is None:
|
||||
return self._render_fallback_prompt(
|
||||
messages, add_generation_prompt=add_generation_prompt
|
||||
)
|
||||
|
||||
return tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=add_generation_prompt,
|
||||
**(template_kwargs or {}),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _render_fallback_prompt(
|
||||
messages: List[dict[str, Any]],
|
||||
add_generation_prompt: bool = True,
|
||||
) -> str:
|
||||
# Fallback prompt format for tokenizers without chat-template support.
|
||||
prompt_lines: List[str] = []
|
||||
for message in messages:
|
||||
role = str(message.get("role", "user"))
|
||||
content = message.get("content", "")
|
||||
if content is None:
|
||||
content = ""
|
||||
prompt_lines.append(f"{role}: {content}")
|
||||
if add_generation_prompt:
|
||||
prompt_lines.append("assistant:")
|
||||
return "\n".join(prompt_lines)
|
||||
|
||||
async def start(self) -> None:
|
||||
# Engine is initialized in __init__; keep start idempotent for protocol
|
||||
# compatibility.
|
||||
return
|
||||
|
||||
async def check_health(self) -> None:
|
||||
# SGLang's in-process Engine API does not expose a health-check method.
|
||||
# Its health endpoints exist only in HTTP/gRPC server entrypoints, which
|
||||
# this integration does not run. Keep the protocol hook as a no-op.
|
||||
return
|
||||
|
||||
def _build_generate_kwargs(
|
||||
self, request: Any, prompt: Any, stream: bool
|
||||
) -> dict[str, Any]:
|
||||
"""Build kwargs dict for engine.async_generate."""
|
||||
generate_kwargs: dict[str, Any] = {
|
||||
"prompt": prompt,
|
||||
"stream": stream,
|
||||
}
|
||||
sampling_params = self._build_sampling_params(request)
|
||||
if sampling_params:
|
||||
generate_kwargs["sampling_params"] = sampling_params
|
||||
return generate_kwargs
|
||||
|
||||
async def _generate_raw(
|
||||
self,
|
||||
request: Any,
|
||||
prompt: Any,
|
||||
) -> dict[str, Any]:
|
||||
"""Run generation and return raw engine output payload."""
|
||||
generate_kwargs = self._build_generate_kwargs(request, prompt, stream=False)
|
||||
return await self.engine.async_generate(**generate_kwargs)
|
||||
|
||||
@staticmethod
|
||||
def _extract_generation_metadata(raw: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Extract normalized generation metadata from one raw engine payload."""
|
||||
text: str = raw.get("text", "")
|
||||
meta: dict[str, Any] = raw.get("meta_info", {}) or {}
|
||||
finish_reason_info = meta.get("finish_reason", {}) or {}
|
||||
finish_reason = SGLangServer._parse_finish_reason(finish_reason_info)
|
||||
|
||||
prompt_tokens = int(meta.get("prompt_tokens", 0))
|
||||
completion_tokens = int(meta.get("completion_tokens", 0))
|
||||
total_tokens = prompt_tokens + completion_tokens
|
||||
|
||||
return {
|
||||
"text": text.strip(),
|
||||
"id": meta.get("id", f"sglang-gen-{uuid.uuid4().hex}"),
|
||||
"created": int(time.time()),
|
||||
"finish_reason": finish_reason,
|
||||
"prompt_tokens": prompt_tokens,
|
||||
"completion_tokens": completion_tokens,
|
||||
"total_tokens": total_tokens,
|
||||
}
|
||||
|
||||
async def _generate_and_extract_metadata(
|
||||
self,
|
||||
request: Any,
|
||||
prompt: Union[str, List[str]],
|
||||
) -> Union[dict[str, Any], List[dict[str, Any]]]:
|
||||
"""
|
||||
Handles parameter extraction, calls the SGLang engine, and processes the
|
||||
raw response to extract common metadata and generated text.
|
||||
|
||||
Accepts either a single prompt string or a list of prompts. When a list
|
||||
is provided, all prompts are sent to SGLang in one batched call, letting
|
||||
SGLang's scheduler handle concurrency natively via async_generate.
|
||||
"""
|
||||
raw = await self._generate_raw(request, prompt)
|
||||
|
||||
# Batch case — SGLang returns a list of results, one per prompt
|
||||
if isinstance(prompt, list):
|
||||
if not raw:
|
||||
raise RuntimeError(
|
||||
"SGLang engine returned an empty response list during generation."
|
||||
)
|
||||
return [self._extract_generation_metadata(r) for r in raw]
|
||||
|
||||
# Single prompt case
|
||||
if isinstance(raw, list):
|
||||
if not raw:
|
||||
raise RuntimeError(
|
||||
"SGLang engine returned an empty response list during generation."
|
||||
)
|
||||
raw = raw[0]
|
||||
return self._extract_generation_metadata(raw)
|
||||
|
||||
async def _stream_generate(
|
||||
self,
|
||||
request: Any,
|
||||
prompt: Any,
|
||||
) -> AsyncGenerator[tuple[str, Optional[str]], None]:
|
||||
"""Stream from SGLang engine, yielding (delta_text, finish_reason) tuples.
|
||||
|
||||
SGLang returns cumulative text in each chunk, so this method
|
||||
tracks the previous text and yields only the incremental delta.
|
||||
"""
|
||||
generate_kwargs = self._build_generate_kwargs(request, prompt, stream=True)
|
||||
stream = await self.engine.async_generate(**generate_kwargs)
|
||||
|
||||
previous_text = ""
|
||||
async for chunk in stream:
|
||||
text = chunk.get("text", "")
|
||||
meta = chunk.get("meta_info", {}) or {}
|
||||
|
||||
delta_text = text[len(previous_text) :]
|
||||
previous_text = text
|
||||
|
||||
finish_reason_info = meta.get("finish_reason", None)
|
||||
finish_reason = (
|
||||
self._parse_finish_reason(finish_reason_info)
|
||||
if finish_reason_info is not None
|
||||
else None
|
||||
)
|
||||
yield delta_text, finish_reason
|
||||
|
||||
@staticmethod
|
||||
def _build_sse_chunk(
|
||||
gen_id: str,
|
||||
object_type: str,
|
||||
created: int,
|
||||
model: str,
|
||||
choice: dict[str, Any],
|
||||
) -> str:
|
||||
"""Build an SSE-formatted chunk string from a single choice payload."""
|
||||
chunk_data = {
|
||||
"id": gen_id,
|
||||
"object": object_type,
|
||||
"created": created,
|
||||
"model": model,
|
||||
"choices": [choice],
|
||||
}
|
||||
return f"data: {json.dumps(chunk_data)}\n\n"
|
||||
|
||||
async def chat(
|
||||
self,
|
||||
request: ChatCompletionRequest,
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[Union[str, ChatCompletionResponse], None]:
|
||||
chat_messages = self._build_chat_messages(request.messages)
|
||||
template_kwargs = self._build_chat_template_kwargs(request)
|
||||
prompt = self._render_chat_prompt(
|
||||
chat_messages, template_kwargs=template_kwargs
|
||||
)
|
||||
|
||||
if request.stream:
|
||||
gen_id = f"sglang-gen-{uuid.uuid4().hex}"
|
||||
created = int(time.time())
|
||||
first_chunk = True
|
||||
async for delta_text, finish_reason in self._stream_generate(
|
||||
request, prompt
|
||||
):
|
||||
delta: dict[str, Any] = {"content": delta_text}
|
||||
if first_chunk:
|
||||
delta["role"] = "assistant"
|
||||
first_chunk = False
|
||||
yield self._build_sse_chunk(
|
||||
gen_id,
|
||||
"chat.completion.chunk",
|
||||
created,
|
||||
request.model,
|
||||
{"index": 0, "delta": delta, "finish_reason": finish_reason},
|
||||
)
|
||||
return
|
||||
|
||||
metadata = await self._generate_and_extract_metadata(request, prompt)
|
||||
|
||||
usage_data = {
|
||||
"prompt_tokens": metadata["prompt_tokens"],
|
||||
"completion_tokens": metadata["completion_tokens"],
|
||||
"total_tokens": metadata["total_tokens"],
|
||||
}
|
||||
|
||||
choice_data = {
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": metadata["text"]},
|
||||
"finish_reason": metadata["finish_reason"],
|
||||
}
|
||||
|
||||
resp = ChatCompletionResponse(
|
||||
id=metadata["id"],
|
||||
object="chat.completion",
|
||||
created=metadata["created"],
|
||||
model=request.model,
|
||||
choices=[choice_data],
|
||||
usage=usage_data,
|
||||
)
|
||||
|
||||
yield resp
|
||||
|
||||
async def completions(
|
||||
self,
|
||||
request: CompletionRequest,
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[Union[str, CompletionResponse], None]:
|
||||
prompt_input = request.prompt
|
||||
|
||||
# Normalize prompt input.
|
||||
if isinstance(prompt_input, list):
|
||||
if not prompt_input:
|
||||
raise ValueError(
|
||||
"The 'prompt' list cannot be empty for completion requests."
|
||||
)
|
||||
prompts_to_process = prompt_input
|
||||
else:
|
||||
prompts_to_process = [prompt_input]
|
||||
|
||||
if request.stream:
|
||||
gen_id = f"sglang-gen-{uuid.uuid4().hex}"
|
||||
created = int(time.time())
|
||||
for i, prompt_string in enumerate(prompts_to_process):
|
||||
async for delta_text, finish_reason in self._stream_generate(
|
||||
request, prompt_string
|
||||
):
|
||||
yield self._build_sse_chunk(
|
||||
gen_id,
|
||||
"text_completion",
|
||||
created,
|
||||
request.model,
|
||||
{
|
||||
"index": i,
|
||||
"text": delta_text,
|
||||
"logprobs": None,
|
||||
"finish_reason": finish_reason,
|
||||
},
|
||||
)
|
||||
return
|
||||
|
||||
results = await self._generate_and_extract_metadata(request, prompts_to_process)
|
||||
|
||||
all_choices = []
|
||||
total_prompt_tokens = 0
|
||||
total_completion_tokens = 0
|
||||
|
||||
for index, metadata in enumerate(results):
|
||||
total_prompt_tokens += metadata["prompt_tokens"]
|
||||
total_completion_tokens += metadata["completion_tokens"]
|
||||
choice_data = {
|
||||
"index": index,
|
||||
"text": metadata["text"],
|
||||
"logprobs": None,
|
||||
"finish_reason": metadata["finish_reason"],
|
||||
}
|
||||
all_choices.append(choice_data)
|
||||
|
||||
usage_data = {
|
||||
"prompt_tokens": total_prompt_tokens,
|
||||
"completion_tokens": total_completion_tokens,
|
||||
"total_tokens": total_prompt_tokens + total_completion_tokens,
|
||||
}
|
||||
|
||||
last_metadata = results[-1]
|
||||
|
||||
resp = CompletionResponse(
|
||||
id=last_metadata["id"],
|
||||
object="text_completion",
|
||||
created=last_metadata.get("created", int(time.time())),
|
||||
model=getattr(request, "model", "default_model"),
|
||||
choices=all_choices,
|
||||
usage=usage_data,
|
||||
)
|
||||
|
||||
yield resp
|
||||
|
||||
async def embeddings(
|
||||
self,
|
||||
request: EmbeddingRequest,
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[EmbeddingResponse, None]:
|
||||
# Input handling follows SGLang's OpenAIServingEmbedding pattern:
|
||||
# https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/entrypoints/openai/serving_embedding.py
|
||||
if isinstance(request, EmbeddingCompletionRequest):
|
||||
prompt = request.input
|
||||
else:
|
||||
# Chat embedding request - join messages without the trailing
|
||||
# "assistant:" generation cue that _render_fallback_prompt adds.
|
||||
chat_messages = self._build_chat_messages(request.messages)
|
||||
prompt = "\n".join(
|
||||
f"{m.get('role', 'user')}: {m.get('content') or ''}"
|
||||
for m in chat_messages
|
||||
)
|
||||
|
||||
# async_encode handles both single strings and lists of strings
|
||||
results = await self.engine.async_encode(prompt)
|
||||
if not isinstance(results, list):
|
||||
results = [results]
|
||||
|
||||
if not results:
|
||||
raise RuntimeError(
|
||||
"SGLang engine returned an empty response for embedding request."
|
||||
)
|
||||
|
||||
# Build response following SGLang's _build_embedding_response pattern
|
||||
data = []
|
||||
total_prompt_tokens = 0
|
||||
|
||||
for idx, ret_item in enumerate(results):
|
||||
data.append(
|
||||
{
|
||||
"index": idx,
|
||||
"object": "embedding",
|
||||
"embedding": ret_item.get("embedding", []),
|
||||
}
|
||||
)
|
||||
meta = ret_item.get("meta_info", {}) or {}
|
||||
total_prompt_tokens += int(meta.get("prompt_tokens", 0))
|
||||
|
||||
resp = EmbeddingResponse(
|
||||
object="list",
|
||||
model=request.model or "",
|
||||
data=data,
|
||||
usage={
|
||||
"prompt_tokens": total_prompt_tokens,
|
||||
"total_tokens": total_prompt_tokens,
|
||||
"completion_tokens": 0,
|
||||
},
|
||||
)
|
||||
|
||||
yield resp
|
||||
|
||||
async def tokenize(
|
||||
self,
|
||||
request: TokenizeRequest,
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[TokenizeResponse, None]:
|
||||
tokenizer = self.engine.tokenizer_manager.tokenizer
|
||||
if tokenizer is None:
|
||||
raise RuntimeError(
|
||||
"Tokenizer is not available. The tokenize endpoint is not "
|
||||
"supported when SGLang is initialized with --skip-tokenizer-init."
|
||||
)
|
||||
|
||||
if isinstance(request, TokenizeCompletionRequest):
|
||||
prompt = request.prompt
|
||||
else:
|
||||
# Chat tokenize request - render messages to prompt string
|
||||
chat_messages = self._build_chat_messages(request.messages)
|
||||
add_generation_prompt = getattr(request, "add_generation_prompt", True)
|
||||
template_kwargs = self._build_chat_template_kwargs(request)
|
||||
prompt = self._render_chat_prompt(
|
||||
chat_messages,
|
||||
add_generation_prompt=add_generation_prompt,
|
||||
template_kwargs=template_kwargs,
|
||||
)
|
||||
|
||||
add_special_tokens = getattr(request, "add_special_tokens", True)
|
||||
tokens = tokenizer.encode(prompt, add_special_tokens=add_special_tokens)
|
||||
|
||||
max_model_len = (
|
||||
getattr(self.engine.tokenizer_manager, "context_len", None)
|
||||
or getattr(self.engine.server_args, "context_length", None)
|
||||
or 0
|
||||
)
|
||||
|
||||
yield TokenizeResponse(
|
||||
tokens=tokens,
|
||||
count=len(tokens),
|
||||
max_model_len=max_model_len,
|
||||
)
|
||||
|
||||
async def detokenize(
|
||||
self,
|
||||
request: DetokenizeRequest,
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[DetokenizeResponse, None]:
|
||||
tokenizer = self.engine.tokenizer_manager.tokenizer
|
||||
if tokenizer is None:
|
||||
raise RuntimeError(
|
||||
"Tokenizer is not available. The detokenize endpoint is not "
|
||||
"supported when SGLang is initialized with --skip-tokenizer-init."
|
||||
)
|
||||
prompt = tokenizer.decode(request.tokens)
|
||||
|
||||
yield DetokenizeResponse(text=prompt)
|
||||
|
||||
async def llm_config(self) -> Optional[LLMConfig]:
|
||||
return self._llm_config
|
||||
|
||||
@classmethod
|
||||
def get_deployment_options(cls, llm_config: "LLMConfig"):
|
||||
deployment_options = copy.deepcopy(llm_config.deployment_config)
|
||||
pg_config = llm_config.placement_group_config or {}
|
||||
|
||||
ray_actor_options = deployment_options.get("ray_actor_options", {})
|
||||
|
||||
tp_size = llm_config.engine_kwargs.get("tp_size", 1)
|
||||
pp_size = llm_config.engine_kwargs.get("pp_size", 1)
|
||||
num_devices = tp_size * pp_size
|
||||
|
||||
if tp_size < 1 or pp_size < 1:
|
||||
raise ValueError(
|
||||
f"Invalid configuration: tp_size={tp_size} and pp_size={pp_size}. "
|
||||
f"Both must be >= 1."
|
||||
)
|
||||
|
||||
if "placement_group_bundles" not in pg_config:
|
||||
child_bundles = [{"GPU": 1} for _ in range(num_devices)]
|
||||
|
||||
replica_bundle = {
|
||||
"CPU": ray_actor_options.get("num_cpus", 1),
|
||||
}
|
||||
|
||||
if ray_actor_options.get("num_gpus"):
|
||||
replica_bundle["GPU"] = ray_actor_options["num_gpus"]
|
||||
|
||||
replica_bundle.update(ray_actor_options.get("resources", {}))
|
||||
|
||||
if "memory" in ray_actor_options:
|
||||
replica_bundle["memory"] = ray_actor_options["memory"]
|
||||
|
||||
pg_bundles = _merge_replica_actor_and_child_actor_bundles(
|
||||
child_actor_bundles=child_bundles,
|
||||
replica_actor_bundle=replica_bundle,
|
||||
)
|
||||
pg_strategy = "PACK"
|
||||
else:
|
||||
pg_bundles = pg_config.get("placement_group_bundles")
|
||||
pg_strategy = pg_config.get("placement_group_strategy", "PACK")
|
||||
|
||||
deployment_options.update(
|
||||
{
|
||||
"placement_group_bundles": pg_bundles,
|
||||
"placement_group_strategy": pg_strategy,
|
||||
}
|
||||
)
|
||||
|
||||
runtime_env = ray_actor_options.setdefault("runtime_env", {})
|
||||
|
||||
if ENABLE_WORKER_PROCESS_SETUP_HOOK:
|
||||
runtime_env.setdefault(
|
||||
"worker_process_setup_hook",
|
||||
"ray.llm._internal.serve._worker_process_setup_hook",
|
||||
)
|
||||
|
||||
if llm_config.runtime_env:
|
||||
runtime_env.update(llm_config.runtime_env)
|
||||
|
||||
deployment_options["ray_actor_options"] = ray_actor_options
|
||||
|
||||
return deployment_options
|
||||
|
||||
async def pause(self, **kwargs: Any) -> None:
|
||||
"""Pause generation on the SGlang server
|
||||
|
||||
This halts generation/encoding requests while keeping model weights in GPU memory. New requests are blocked until resume is called.
|
||||
|
||||
Args:
|
||||
**kwargs: Options parsed into SGLangPauseConfig.
|
||||
- mode (str): "abort" (default), "in_place", or "retract"
|
||||
"""
|
||||
|
||||
assert self.engine is not None, "server is not initialized"
|
||||
config = SGLangPauseConfig(**kwargs)
|
||||
from sglang.srt.managers.io_struct import PauseGenerationReqInput
|
||||
|
||||
await self.engine.tokenizer_manager.pause_generation(
|
||||
PauseGenerationReqInput(mode=config.mode)
|
||||
)
|
||||
self._is_paused = True
|
||||
|
||||
async def resume(self, **kwargs: Any) -> None:
|
||||
"""Resume generation on the SGLang server after pause.
|
||||
|
||||
Args:
|
||||
**kwargs: Reserved for future options.
|
||||
"""
|
||||
assert self.engine is not None, "server is not initialized"
|
||||
from sglang.srt.managers.io_struct import ContinueGenerationReqInput
|
||||
|
||||
await self.engine.tokenizer_manager.continue_generation(
|
||||
ContinueGenerationReqInput()
|
||||
)
|
||||
self._is_paused = False
|
||||
|
||||
async def is_paused(self) -> bool:
|
||||
"""Check whether the SGLang server is currently paused.
|
||||
|
||||
Returns:
|
||||
True if the server is paused, False otherwise.
|
||||
"""
|
||||
return self._is_paused
|
||||
|
||||
async def sleep(self, **kwargs: Any) -> None:
|
||||
"""Put SGLang server to sleep.
|
||||
|
||||
Args:
|
||||
**kwargs: Options parsed into SGLangSleepConfig
|
||||
- tags (List[str], optional): Components to put to sleep.
|
||||
"""
|
||||
|
||||
assert self.engine is not None, "server is not initialized"
|
||||
config = SGLangSleepConfig(**kwargs)
|
||||
|
||||
# release_memory_occupation() calls loop.run_until_complete() internally, which fails
|
||||
# inside an async context. Await the underlying coroutine directly.
|
||||
from sglang.srt.entrypoints.engine import ReleaseMemoryOccupationReqInput
|
||||
|
||||
obj = ReleaseMemoryOccupationReqInput(tags=config.tags)
|
||||
await self.engine.tokenizer_manager.release_memory_occupation(obj, None)
|
||||
self._sleeping_tags |= set(config.tags) if config.tags else set(_SLEEP_TAGS)
|
||||
|
||||
async def wakeup(self, **kwargs: Any) -> None:
|
||||
"""Wake up the SGLang server from sleep mode.
|
||||
|
||||
Args:
|
||||
**kwargs: Options parsed into SGLangWakeupConfig
|
||||
- tags (List[str], optional): Components to wake up.
|
||||
"""
|
||||
|
||||
assert self.engine is not None, "server is not initialized"
|
||||
config = SGLangWakeupConfig(**kwargs)
|
||||
# resume_memory_occupation() release_memory_occupation() calls loop.run_until_complete() internally, which fails
|
||||
# inside an async context. Await the underlying coroutine directly.
|
||||
from sglang.srt.entrypoints.engine import ResumeMemoryOccupationReqInput
|
||||
|
||||
obj = ResumeMemoryOccupationReqInput(tags=config.tags)
|
||||
await self.engine.tokenizer_manager.resume_memory_occupation(obj, None)
|
||||
|
||||
if config.tags is None:
|
||||
self._sleeping_tags.clear()
|
||||
else:
|
||||
self._sleeping_tags -= set(config.tags)
|
||||
|
||||
async def is_sleeping(self) -> bool:
|
||||
"""Check whether the SGLang server is currently sleeping.
|
||||
|
||||
Returns:
|
||||
True if any component is currently offloaded/discarded, False otherwise.
|
||||
"""
|
||||
return bool(self._sleeping_tags)
|
||||
|
||||
async def reset_prefix_cache(self, timeout: Optional[float] = None) -> None:
|
||||
assert self.engine is not None, "server is not initialized"
|
||||
# flush_cache() calls loop.run_until_complete() internally, which fails
|
||||
# inside an async context. Await the underlying coroutine directly.
|
||||
await self.engine.tokenizer_manager.flush_cache()
|
||||
@@ -0,0 +1,5 @@
|
||||
"""KV Transfer connector backends for Ray Serve LLM.
|
||||
|
||||
This package provides connector backends for KV cache transfer in vLLM.
|
||||
All backends are lazily loaded through the factory to avoid circular imports.
|
||||
"""
|
||||
@@ -0,0 +1,264 @@
|
||||
import abc
|
||||
import random
|
||||
import string
|
||||
from typing import TYPE_CHECKING, Any, Dict, Optional, Union
|
||||
|
||||
from ray import serve
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
|
||||
from ray.llm._internal.serve.core.configs.openai_api_models import (
|
||||
ChatCompletionRequest,
|
||||
CompletionRequest,
|
||||
)
|
||||
|
||||
# The two OpenAI request models the P/D orchestrator shapes. Defined under
|
||||
# TYPE_CHECKING (and used as a string annotation) to avoid an import cycle
|
||||
# between this module and the config/openai-models modules.
|
||||
RequestType = Union[ChatCompletionRequest, CompletionRequest]
|
||||
|
||||
|
||||
def base_prefill_kv_transfer_params() -> Dict[str, Any]:
|
||||
"""The ``kv_transfer_params`` common to a prefill (producer) request.
|
||||
|
||||
Tells the prefill engine to produce KV for a remote decode. Connectors layer
|
||||
their own keys (e.g. a transfer id, DP/TP routing) on top of these.
|
||||
"""
|
||||
return {
|
||||
"do_remote_decode": True,
|
||||
"do_remote_prefill": False,
|
||||
"remote_engine_id": None,
|
||||
"remote_block_ids": None,
|
||||
}
|
||||
|
||||
|
||||
def clamp_request_to_single_token(request: "RequestType") -> None:
|
||||
"""Clamp a prefill request to a single, non-streaming token (in place)."""
|
||||
request.max_tokens = 1
|
||||
if hasattr(request, "max_completion_tokens"):
|
||||
request.max_completion_tokens = 1
|
||||
request.stream = False
|
||||
if hasattr(request, "stream_options"):
|
||||
request.stream_options = None
|
||||
|
||||
|
||||
class BaseConnectorBackend(abc.ABC):
|
||||
# ---- P/D coordination protocol ----
|
||||
#
|
||||
# These class attributes and methods let the P/D orchestrator
|
||||
# (``PDOrchestratorMixin``) delegate request shaping, peer addressing, and
|
||||
# handoff discipline to the connector. They are connector-agnostic: a
|
||||
# connector picks a quadrant of (``requires_peer_binding``,
|
||||
# ``concurrent_handoff``) and implements ``prepare_prefill_request`` /
|
||||
# ``prepare_decode_request`` accordingly.
|
||||
#
|
||||
# ``requires_peer_binding``:
|
||||
# * False -> the orchestrator dispatches prefill via the standard handle
|
||||
# path; the peer (if any) is resolved post-hoc from the prefill response.
|
||||
# * True -> the orchestrator selects the prefill replica first
|
||||
# (``choose_replica``) and passes its ``replica_metadata`` to the backend
|
||||
# as ``peer`` (pre-dispatch addressing).
|
||||
#
|
||||
# ``concurrent_handoff``:
|
||||
# * False -> prefill runs to its first chunk before local decode starts
|
||||
# (sequential handoff).
|
||||
# * True -> prefill dispatch and local decode run concurrently.
|
||||
#
|
||||
# The two flags are independent; the known combos:
|
||||
# * (False, False) — e.g. NixlConnector / LMCacheConnectorV1: decode learns
|
||||
# everything it needs (remote engine id / block ids) from the prefill
|
||||
# response, so it must wait for that response (sequential).
|
||||
# * (True, True) — e.g. MoRIIO WRITE mode: the peer address is bound
|
||||
# into the request id before dispatch and prefill *pushes* KV to decode,
|
||||
# so decode needs nothing from prefill's response and can start
|
||||
# immediately (concurrent). Concurrent handoff is only possible because
|
||||
# peer binding happens up front.
|
||||
# * (True, False) — e.g. MoRIIO READ mode: the peer is bound up front,
|
||||
# but decode *pulls* KV using block ids returned in prefill's response,
|
||||
# so it still waits for prefill to finish (sequential).
|
||||
requires_peer_binding: bool = False
|
||||
concurrent_handoff: bool = False
|
||||
|
||||
def __init__(self, llm_config: "LLMConfig"):
|
||||
"""Base class for connector backends.
|
||||
|
||||
Args:
|
||||
llm_config: The llm configuration for this engine
|
||||
"""
|
||||
self.llm_config = llm_config
|
||||
|
||||
@property
|
||||
def kv_transfer_config(self) -> Dict[str, Any]:
|
||||
engine_kwargs = self.llm_config.engine_kwargs
|
||||
kv_transfer_config = engine_kwargs.get("kv_transfer_config")
|
||||
assert (
|
||||
kv_transfer_config is not None
|
||||
), "In Connector backend, kv_transfer_config is not set"
|
||||
return kv_transfer_config
|
||||
|
||||
def _get_unique_suffix(self, len: int = 6) -> str:
|
||||
"""Generates unique alphanumeric suffix.
|
||||
|
||||
Args:
|
||||
len: Length of the suffix to generate.
|
||||
Returns:
|
||||
A unique alphanumeric suffix string of specified length.
|
||||
"""
|
||||
return "".join(random.choices(string.ascii_letters + string.digits, k=len))
|
||||
|
||||
def _compute_port_offset(self) -> int:
|
||||
"""Compute a deterministic port offset for this replica.
|
||||
|
||||
Uses data_parallel_rank if DP case, otherwise falls back to
|
||||
the replica rank assigned by Ray Serve (TP/PP case).
|
||||
|
||||
For TP/PP cases, multiply by num_devices (tp × pp) to reserve
|
||||
sufficient port space, since each worker needs a unique port.
|
||||
Each TP worker adds its tp_rank (0, 1, ..., tp_size-1) to the
|
||||
base port at bind time, and PP stages also need separate ports.
|
||||
|
||||
Returns:
|
||||
Non-negative integer offset to add to a base port.
|
||||
"""
|
||||
# Prefer explicit DP rank when available
|
||||
dp_rank = self.llm_config.engine_kwargs.get("data_parallel_rank")
|
||||
if isinstance(dp_rank, int) and dp_rank >= 0:
|
||||
# vLLM already accounts for TP spacing in DP offset calculation
|
||||
# (data_parallel_rank × tp_size), don't multiply here
|
||||
return dp_rank
|
||||
|
||||
# NOTE (jeffreywang): A missing replica context must fail loudly, not
|
||||
# silently return a 0 offset that collides colocated replicas on the
|
||||
# same NIXL side-channel port. get_replica_context() raises RayServeException
|
||||
# outside a replica.
|
||||
rc = serve.get_replica_context()
|
||||
engine_config = self.llm_config.get_engine_config()
|
||||
num_devices = engine_config.num_devices
|
||||
return rc.rank.rank * num_devices
|
||||
|
||||
@abc.abstractmethod
|
||||
def prepare_prefill_request(
|
||||
self, *, request: "RequestType", peer: Optional[Dict[str, Any]]
|
||||
) -> "RequestType":
|
||||
"""Shape the request sent to the remote prefill engine.
|
||||
|
||||
Args:
|
||||
request: The incoming chat/completion request.
|
||||
peer: The selected prefill replica's ``replica_metadata`` dict when
|
||||
the connector opted into pre-dispatch peer binding
|
||||
(``requires_peer_binding=True``), else None.
|
||||
|
||||
Returns:
|
||||
A new request object to dispatch to the prefill engine.
|
||||
"""
|
||||
...
|
||||
|
||||
@abc.abstractmethod
|
||||
def prepare_decode_request(
|
||||
self,
|
||||
*,
|
||||
request: "RequestType",
|
||||
peer: Optional[Dict[str, Any]],
|
||||
prefill_response: Optional[Any],
|
||||
) -> "RequestType":
|
||||
"""Shape the request run on the local decode engine.
|
||||
|
||||
Args:
|
||||
request: The incoming chat/completion request.
|
||||
peer: The selected prefill replica's ``replica_metadata`` dict when
|
||||
the connector opted into pre-dispatch peer binding, else None.
|
||||
prefill_response: The captured prefill response chunk whose
|
||||
``kv_transfer_params`` may be forwarded, or None when no chunk is
|
||||
captured before decode starts (concurrent-handoff mode).
|
||||
|
||||
Returns:
|
||||
A new request object to run on the local decode engine.
|
||||
"""
|
||||
...
|
||||
|
||||
def setup(self) -> None:
|
||||
"""Setup the connector backend.
|
||||
|
||||
This method is called to setup the connector backend.
|
||||
"""
|
||||
pass
|
||||
|
||||
def replica_metadata(self) -> Dict[str, Any]:
|
||||
"""Static per-replica coordination data published to the orchestrator.
|
||||
|
||||
Surfaced via the replica-metadata hook on ``ReplicaSelection`` so that a
|
||||
connector opting into ``requires_peer_binding`` can address the selected
|
||||
prefill peer. The default backend publishes nothing; connectors that need
|
||||
to advertise an address (e.g. MoRIIO's zmq endpoint) override this.
|
||||
|
||||
Returns:
|
||||
A JSON-serializable dict of per-replica metadata (empty by default).
|
||||
"""
|
||||
return {}
|
||||
|
||||
|
||||
class DefaultPDProtocolMixin:
|
||||
"""The default P/D protocol policy: no peer binding, sequential handoff.
|
||||
|
||||
Implements ``prepare_prefill_request`` / ``prepare_decode_request`` for
|
||||
connectors that follow the standard policy: the prefill engine is told to
|
||||
produce KV for a remote decode (clamped to a single non-streaming token),
|
||||
and the decode engine forwards the ``kv_transfer_params`` that the prefill
|
||||
engine returned on its first response chunk.
|
||||
|
||||
Mix this in *before* ``BaseConnectorBackend`` in a backend's bases so its
|
||||
concrete methods satisfy the abstract methods.
|
||||
"""
|
||||
|
||||
def prepare_prefill_request(
|
||||
self, *, request: "RequestType", peer: Optional[Dict[str, Any]]
|
||||
) -> "RequestType":
|
||||
"""Shape the prefill request under the default P/D protocol policy.
|
||||
|
||||
Deep-copies the request, stamps the standard ``kv_transfer_params`` that
|
||||
tell the prefill engine to produce KV for a remote decode, and clamps it
|
||||
to a single, non-streaming token. ``peer`` is ignored.
|
||||
"""
|
||||
assert (
|
||||
getattr(request, "kv_transfer_params", None) is None
|
||||
), "kv_transfer_params should be empty before orchestrator"
|
||||
prefill_request = request.model_copy(deep=True)
|
||||
prefill_request.kv_transfer_params = {
|
||||
**base_prefill_kv_transfer_params(),
|
||||
"remote_host": None,
|
||||
"remote_port": None,
|
||||
}
|
||||
clamp_request_to_single_token(prefill_request)
|
||||
return prefill_request
|
||||
|
||||
def prepare_decode_request(
|
||||
self,
|
||||
*,
|
||||
request: "RequestType",
|
||||
peer: Optional[Dict[str, Any]],
|
||||
prefill_response: Optional[Any],
|
||||
) -> "RequestType":
|
||||
"""Shape the decode request under the default P/D protocol policy.
|
||||
|
||||
Deep-copies the request and, only when a prefill response chunk was
|
||||
captured, forwards its ``kv_transfer_params`` so the decode engine
|
||||
pulls/receives the KV produced by prefill. In concurrent-handoff mode
|
||||
``prefill_response`` is None and the request is left unmodified. ``peer``
|
||||
is ignored.
|
||||
"""
|
||||
decode_request = request.model_copy(deep=True)
|
||||
if prefill_response is not None:
|
||||
decode_request.kv_transfer_params = prefill_response.kv_transfer_params
|
||||
return decode_request
|
||||
|
||||
|
||||
class DefaultConnectorBackend(DefaultPDProtocolMixin, BaseConnectorBackend):
|
||||
"""Concrete connector backend using the default P/D protocol policy.
|
||||
|
||||
Used as the factory fallback for connectors that are not registered with a
|
||||
dedicated backend class: they get a no-op ``setup()`` and the default
|
||||
request-shaping policy. ``BaseConnectorBackend`` is abstract, so the factory
|
||||
must return a concrete class like this one.
|
||||
"""
|
||||
|
||||
pass
|
||||
@@ -0,0 +1,140 @@
|
||||
"""Factory for lazy-loading KV connector backends.
|
||||
|
||||
This module provides a factory pattern for registering and instantiating
|
||||
KV connector backends without eagerly importing all implementations.
|
||||
This avoids circular import issues and improves startup performance.
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING, Type, Union
|
||||
|
||||
from ray.llm._internal.serve.engines.vllm.kv_transfer.base import (
|
||||
BaseConnectorBackend,
|
||||
DefaultConnectorBackend,
|
||||
)
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
from ray.llm._internal.serve.utils.registry import get_registry
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
|
||||
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
# Get the registry instance for KV connector backends
|
||||
_kv_backend_registry = get_registry("kv_connector_backend")
|
||||
|
||||
|
||||
class KVConnectorBackendFactory:
|
||||
"""Factory for creating KV connector backend instances with lazy loading."""
|
||||
|
||||
@classmethod
|
||||
def register_backend(
|
||||
cls,
|
||||
name: str,
|
||||
backend_class_or_path: Union[Type["BaseConnectorBackend"], str],
|
||||
) -> None:
|
||||
"""Register a connector backend.
|
||||
|
||||
This enables the backend to be accessed on every Ray process in the cluster.
|
||||
|
||||
Args:
|
||||
name: The name of the connector (e.g., "LMCacheConnectorV1")
|
||||
backend_class_or_path: Either:
|
||||
- The backend class object directly (preferred), or
|
||||
- A string in the format "module_path:class_name" for lazy loading
|
||||
|
||||
Examples:
|
||||
# Register with class directly (recommended):
|
||||
KVConnectorBackendFactory.register_backend("MyConnector", MyConnectorClass)
|
||||
|
||||
# Register with module path string (for lazy loading):
|
||||
KVConnectorBackendFactory.register_backend("MyConnector", "my.module:MyClass")
|
||||
"""
|
||||
_kv_backend_registry.register(name, backend_class_or_path)
|
||||
|
||||
@classmethod
|
||||
def get_backend_class(cls, name: str) -> Type["BaseConnectorBackend"]:
|
||||
"""Get the connector backend class by name.
|
||||
|
||||
For registered connectors, returns the registered backend class.
|
||||
For unregistered connectors, returns DefaultConnectorBackend (a concrete
|
||||
backend with a no-op setup() and the default P/D protocol policy),
|
||||
allowing connectors that don't require Ray Serve orchestration to work
|
||||
without registration. (BaseConnectorBackend itself is abstract and
|
||||
cannot be instantiated.)
|
||||
|
||||
Args:
|
||||
name: The name of the connector backend
|
||||
|
||||
Returns:
|
||||
The connector backend class
|
||||
|
||||
Raises:
|
||||
ImportError: If a registered backend fails to load
|
||||
"""
|
||||
try:
|
||||
return _kv_backend_registry.get(name)
|
||||
except ValueError:
|
||||
logger.warning(
|
||||
f"Unsupported connector backend: {name}. "
|
||||
f"Using default: {DefaultConnectorBackend.__name__}."
|
||||
)
|
||||
return DefaultConnectorBackend
|
||||
except Exception as e:
|
||||
raise ImportError(
|
||||
f"Failed to load connector backend '{name}': {type(e).__name__}: {e}"
|
||||
) from e
|
||||
|
||||
@classmethod
|
||||
def create_backend(
|
||||
cls, name: str, llm_config: "LLMConfig"
|
||||
) -> "BaseConnectorBackend":
|
||||
"""Create a connector backend instance.
|
||||
|
||||
Args:
|
||||
name: The name of the connector backend
|
||||
llm_config: The LLM configuration
|
||||
|
||||
Returns:
|
||||
An instance of the connector backend
|
||||
"""
|
||||
return cls.get_backend_class(name)(llm_config)
|
||||
|
||||
@classmethod
|
||||
def is_registered(cls, name: str) -> bool:
|
||||
"""Check if a connector backend is registered."""
|
||||
return _kv_backend_registry.contains(name)
|
||||
|
||||
@classmethod
|
||||
def unregister_backend(cls, name: str) -> None:
|
||||
"""Unregister a connector backend.
|
||||
|
||||
Removes the backend from the registry across all Ray processes.
|
||||
|
||||
Args:
|
||||
name: The name of the connector backend to unregister
|
||||
"""
|
||||
_kv_backend_registry.unregister(name)
|
||||
|
||||
|
||||
BUILTIN_BACKENDS = {
|
||||
"LMCacheConnectorV1": "ray.llm._internal.serve.engines.vllm.kv_transfer.lmcache:LMCacheConnectorV1Backend",
|
||||
"NixlConnector": "ray.llm._internal.serve.engines.vllm.kv_transfer.nixl:NixlConnectorBackend",
|
||||
"MultiConnector": "ray.llm._internal.serve.engines.vllm.kv_transfer.multi_connector:MultiConnectorBackend",
|
||||
"MoRIIOConnector": "ray.llm._internal.serve.engines.vllm.kv_transfer.moriio:MoRIIOConnectorBackend",
|
||||
}
|
||||
|
||||
|
||||
def _initialize_registry() -> None:
|
||||
"""Initialize the registry with built-in backends.
|
||||
|
||||
This function is called when the module is imported to ensure
|
||||
built-in backends are registered.
|
||||
"""
|
||||
for name, backend_path in BUILTIN_BACKENDS.items():
|
||||
if not KVConnectorBackendFactory.is_registered(name):
|
||||
KVConnectorBackendFactory.register_backend(name, backend_path)
|
||||
|
||||
|
||||
# Initialize registry when module is imported
|
||||
_initialize_registry()
|
||||
@@ -0,0 +1,62 @@
|
||||
from ray.llm._internal.serve.engines.vllm.kv_transfer.base import (
|
||||
BaseConnectorBackend,
|
||||
DefaultPDProtocolMixin,
|
||||
)
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def _check_lmcache_installed():
|
||||
try:
|
||||
import lmcache # noqa: F401
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"LMCache is not installed. Please install it with `pip install lmcache`."
|
||||
)
|
||||
|
||||
|
||||
class LMCacheConnectorV1Backend(DefaultPDProtocolMixin, BaseConnectorBackend):
|
||||
|
||||
KV_CONNECTOR_EXTRA_CONFIG_FIELD_NAME = "kv_connector_extra_config"
|
||||
LMCACHE_RPC_PORT_FIELD_NAME = "lmcache_rpc_port"
|
||||
DEFAULT_LMCACHE_RPC_PORT_NAME = "lmcache_rpc_port"
|
||||
|
||||
def setup(self) -> None:
|
||||
"""Initialize the LMCache connector backend.
|
||||
|
||||
Creates a unique LMCache RPC port name across replicas by appending
|
||||
a random suffix to the base port name.
|
||||
|
||||
Raises:
|
||||
ImportError: If LMCache is not installed.
|
||||
"""
|
||||
_check_lmcache_installed()
|
||||
|
||||
if (
|
||||
LMCacheConnectorV1Backend.KV_CONNECTOR_EXTRA_CONFIG_FIELD_NAME
|
||||
not in self.kv_transfer_config
|
||||
):
|
||||
return
|
||||
|
||||
kv_connector_extra_config = self.kv_transfer_config[
|
||||
LMCacheConnectorV1Backend.KV_CONNECTOR_EXTRA_CONFIG_FIELD_NAME
|
||||
]
|
||||
base_value = kv_connector_extra_config.get(
|
||||
LMCacheConnectorV1Backend.LMCACHE_RPC_PORT_FIELD_NAME,
|
||||
LMCacheConnectorV1Backend.DEFAULT_LMCACHE_RPC_PORT_NAME,
|
||||
)
|
||||
|
||||
# Append random suffix for uniqueness
|
||||
lmcache_rpc_port_value = str(base_value) + self._get_unique_suffix()
|
||||
if (
|
||||
LMCacheConnectorV1Backend.LMCACHE_RPC_PORT_FIELD_NAME
|
||||
in kv_connector_extra_config
|
||||
):
|
||||
logger.info(
|
||||
f"Setting unique lmcache_rpc_port={lmcache_rpc_port_value} for current replica."
|
||||
)
|
||||
|
||||
kv_connector_extra_config[
|
||||
LMCacheConnectorV1Backend.LMCACHE_RPC_PORT_FIELD_NAME
|
||||
] = lmcache_rpc_port_value
|
||||
@@ -0,0 +1,344 @@
|
||||
"""MoRIIO connector backend for Ray Serve LLM (analogue of nixl.py).
|
||||
|
||||
Configures a vLLM engine's ``kv_transfer_config.kv_connector_extra_config`` for
|
||||
the MoRIIO connector and computes per-replica handshake/notify ports so colocated
|
||||
replicas don't collide. Also builds the engine's advertised zmq address so the
|
||||
P/D orchestrator can discover it via the replica-metadata hook
|
||||
(``ReplicaSelection.replica_metadata``), and implements the PD connector protocol
|
||||
(``requires_peer_binding`` / ``concurrent_handoff`` / ``prepare_prefill_request`` /
|
||||
``prepare_decode_request``) so the decode orchestrator can address the selected
|
||||
prefill peer by request id.
|
||||
|
||||
Unlike NIXL/LMCache, MoRIIO does NOT use ``DefaultPDProtocolMixin``: it has custom
|
||||
request shaping (a dual-address request_id + transfer_id) and therefore IMPLEMENTS
|
||||
the abstract ``prepare_*`` methods directly on ``BaseConnectorBackend``.
|
||||
|
||||
Two transfer disciplines, selected by ``read_mode``:
|
||||
* WRITE (default): prefill PUSHES KV to decode -> concurrent handoff.
|
||||
* READ: decode PULLS KV from prefill -> sequential handoff; the decode request
|
||||
forwards the ``remote_block_ids`` / ``remote_engine_id`` the prefill engine
|
||||
returned.
|
||||
|
||||
The dual-address request_id and the transfer_id are derived DETERMINISTICALLY
|
||||
from the incoming request id (a hash), so ``prepare_prefill_request`` and
|
||||
``prepare_decode_request`` produce identical ids across their two separate calls
|
||||
without per-request backend state (the backend instance is shared across
|
||||
requests).
|
||||
|
||||
Registered with Ray's public connector registry via the factory.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
import re
|
||||
from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple
|
||||
|
||||
import ray
|
||||
from ray.llm._internal.serve.engines.vllm.kv_transfer.base import (
|
||||
BaseConnectorBackend,
|
||||
base_prefill_kv_transfer_params,
|
||||
clamp_request_to_single_token,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.llm._internal.serve.engines.vllm.kv_transfer.base import RequestType
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Defaults mirror vLLM's MoRIIOConstants (DEFAULT_HANDSHAKE_PORT / NOTIFY_PORT).
|
||||
# Prefill uses these bases; decode is shifted (see builder.py) so a colocated
|
||||
# P+D pair on one node doesn't collide.
|
||||
DEFAULT_HANDSHAKE_PORT_BASE = 6301
|
||||
DEFAULT_NOTIFY_PORT_BASE = 61005
|
||||
|
||||
# experimental_configs keys understood by this backend.
|
||||
HANDSHAKE_PORT_BASE_KEY = "MORI_HANDSHAKE_PORT_BASE"
|
||||
NOTIFY_PORT_BASE_KEY = "MORI_NOTIFY_PORT_BASE"
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dual-address request_id / zmq address encoding.
|
||||
#
|
||||
# These MUST stay byte-compatible with the regexes vLLM's MoRIIO connector uses
|
||||
# to recover peer addresses from the request_id:
|
||||
#
|
||||
# vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py
|
||||
# _PREFILL_ZMQ_RE = re.compile(r"___prefill_addr_(.+?)___decode_addr_")
|
||||
# _DECODE_ZMQ_RE = re.compile(r"___decode_addr_(.+)_[0-9a-f]{32}(?:-.*)?$")
|
||||
# # zmq address: "host:IP,handshake:PORT,notify:PORT"
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_PREFILL_PREFIX = "___prefill_addr_"
|
||||
_DECODE_PREFIX = "___decode_addr_"
|
||||
_TRANSFER_PREFIX = "tx"
|
||||
|
||||
# Copies of vLLM's regexes for local validation / round-trip tests.
|
||||
_PREFILL_ZMQ_RE = re.compile(r"___prefill_addr_(.+?)___decode_addr_")
|
||||
_DECODE_ZMQ_RE = re.compile(r"___decode_addr_(.+)_[0-9a-f]{32}(?:-.*)?$")
|
||||
|
||||
|
||||
def build_zmq_address(host: str, handshake_port: int, notify_port: int) -> str:
|
||||
"""Build the MORI zmq address string ``host:IP,handshake:PORT,notify:PORT``."""
|
||||
return f"host:{host},handshake:{handshake_port},notify:{notify_port}"
|
||||
|
||||
|
||||
def parse_zmq_address(zmq_address: str) -> Tuple[str, int, int]:
|
||||
"""Inverse of :func:`build_zmq_address` -> ``(host, handshake_port, notify_port)``."""
|
||||
parts = {}
|
||||
for segment in zmq_address.split(","):
|
||||
key, _, val = segment.partition(":")
|
||||
parts[key.strip()] = val.strip()
|
||||
return parts["host"], int(parts["handshake"]), int(parts["notify"])
|
||||
|
||||
|
||||
def parse_peer_zmq(request_id: str, is_producer: bool) -> str:
|
||||
"""Recover the peer's zmq address from a request id (for tests/debugging).
|
||||
|
||||
Producer (prefill) wants the *decode* address; consumer wants the *prefill*.
|
||||
"""
|
||||
rex = _DECODE_ZMQ_RE if is_producer else _PREFILL_ZMQ_RE
|
||||
m = rex.search(request_id)
|
||||
if not m:
|
||||
raise ValueError(f"No peer zmq address in request_id: {request_id!r}")
|
||||
return m.group(1)
|
||||
|
||||
|
||||
def _read_mode_enabled(extra_config: Dict[str, Any]) -> bool:
|
||||
"""Mirror vLLM's ``get_moriio_mode`` parse of ``read_mode``.
|
||||
|
||||
true / 1 -> READ; anything else -> WRITE (default).
|
||||
"""
|
||||
return str(extra_config.get("read_mode", "false")).lower().strip() in (
|
||||
"true",
|
||||
"1",
|
||||
)
|
||||
|
||||
|
||||
class MoRIIOConnectorBackend(BaseConnectorBackend):
|
||||
"""Set up MoRIIO ports/extra_config and implement the PD connector protocol."""
|
||||
|
||||
# The advertised zmq address ("host:IP,handshake:PORT,notify:PORT"),
|
||||
# computed by setup(); consumers reach it via this backend instance.
|
||||
_zmq_address: Optional[str] = None
|
||||
|
||||
# MORI addresses peers by the dual-address request id, so the orchestrator
|
||||
# must bind to the selected prefill replica BEFORE dispatch.
|
||||
requires_peer_binding: bool = True
|
||||
|
||||
def _extra_config(self) -> dict:
|
||||
cfg = self.kv_transfer_config.setdefault("kv_connector_extra_config", {})
|
||||
return cfg
|
||||
|
||||
@property
|
||||
def _read_mode(self) -> bool:
|
||||
"""True iff this engine's MoRIIO connector is configured for READ mode."""
|
||||
extra = self._extra_config()
|
||||
return _read_mode_enabled(extra)
|
||||
|
||||
@property
|
||||
def concurrent_handoff(self) -> bool:
|
||||
"""WRITE -> concurrent (prefill pushes); READ -> sequential (decode pulls)."""
|
||||
return not self._read_mode
|
||||
|
||||
def setup(self) -> None:
|
||||
offset = self._compute_port_offset()
|
||||
|
||||
handshake_base = int(
|
||||
self.llm_config.experimental_configs.get(
|
||||
HANDSHAKE_PORT_BASE_KEY, DEFAULT_HANDSHAKE_PORT_BASE
|
||||
)
|
||||
)
|
||||
notify_base = int(
|
||||
self.llm_config.experimental_configs.get(
|
||||
NOTIFY_PORT_BASE_KEY, DEFAULT_NOTIFY_PORT_BASE
|
||||
)
|
||||
)
|
||||
|
||||
# NOTE: vLLM internally adds get_port_offset(dp_rank, tp_rank) on top of
|
||||
# these bases. For TP/DP>1, reserve a stride >= tp_size*pp_size when
|
||||
# shifting decode's base in the builder so the two offset schemes never
|
||||
# overlap.
|
||||
handshake_port = handshake_base + offset
|
||||
notify_port = notify_base + offset
|
||||
|
||||
extra = self._extra_config()
|
||||
# Required keys for vLLM's config parser (KeyError otherwise) -- proxyless.
|
||||
extra.setdefault("proxy_ip", "") # empty => ping/registration thread disabled
|
||||
extra.setdefault("proxy_ping_port", "0")
|
||||
# TODO: real Serve replica HTTP port. Harmless placeholder while
|
||||
# proxy_ip="" (only used to build request_address for the disabled ping).
|
||||
extra.setdefault("http_port", str(8000 + offset))
|
||||
# WRITE mode (prefill pushes). READ would be "true".
|
||||
extra.setdefault("read_mode", "false")
|
||||
extra["handshake_port"] = str(handshake_port)
|
||||
extra["notify_port"] = str(notify_port)
|
||||
|
||||
# Advertise the Ray internal cluster IP as the zmq host.
|
||||
host = ray.util.get_node_ip_address()
|
||||
zmq_address = build_zmq_address(host, handshake_port, notify_port)
|
||||
# Stash so replica_metadata() can publish it; the decode
|
||||
# orchestrator reads the selected prefill replica's copy off the peer.
|
||||
self._zmq_address = zmq_address
|
||||
# Cross-node correctness: vLLM's MoRIIO worker otherwise binds/advertises
|
||||
# get_ip(), which on a Ray cluster is the node's unroutable public IP, and
|
||||
# VLLM_HOST_IP cannot be propagated to workers (it is excluded from vLLM's
|
||||
# driver->worker env copy). Pass the routable node IP through the connector
|
||||
# config, which does reach the workers; vLLM's MoRIIO connector honors
|
||||
# "host_ip" over get_ip(). Requires vllm-project/vllm#45488.
|
||||
extra["host_ip"] = host
|
||||
|
||||
# ---- parallelism (data/tensor) ----
|
||||
|
||||
def _dp_rank(self) -> int:
|
||||
rank = self.llm_config.engine_kwargs.get("data_parallel_rank")
|
||||
return rank if isinstance(rank, int) and rank >= 0 else 0
|
||||
|
||||
def _dp_size(self) -> int:
|
||||
return int(self.llm_config.engine_kwargs.get("data_parallel_size") or 1)
|
||||
|
||||
def _tp_size(self) -> int:
|
||||
return int(self.llm_config.engine_kwargs.get("tensor_parallel_size") or 1)
|
||||
|
||||
# ---- replica metadata (published via the replica-metadata hook) ----
|
||||
|
||||
def replica_metadata(self) -> dict:
|
||||
"""Static per-replica coordination data published to the orchestrator.
|
||||
|
||||
The prefill replica publishes its MORI zmq address and its parallelism
|
||||
(DP rank/size, TP size); the decode orchestrator reads them off the
|
||||
selected prefill replica's ``ReplicaSelection.replica_metadata`` and uses
|
||||
them to address the right remote (dp_rank, tp) workers.
|
||||
"""
|
||||
return {
|
||||
"mori_zmq_address": self._zmq_address,
|
||||
"dp_rank": self._dp_rank(),
|
||||
"dp_size": self._dp_size(),
|
||||
"tp_size": self._tp_size(),
|
||||
}
|
||||
|
||||
def _remote_routing(self, remote: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""``kv_transfer_params`` keys telling vLLM which remote workers to reach.
|
||||
|
||||
``remote`` is the metadata of the *other* side of the transfer: the
|
||||
decode (this orchestrator) for a prefill request, the selected prefill
|
||||
peer for a decode request. vLLM addresses a remote worker at
|
||||
``advertised_base + get_port_offset(remote_dp_rank, tp_index)`` and
|
||||
handshakes all ``remote_dp_size`` ranks, so both must match the target
|
||||
replica. ``tp_size`` is the remote's TP (symmetric across P/D).
|
||||
"""
|
||||
return {
|
||||
"remote_dp_rank": int(remote.get("dp_rank", 0)),
|
||||
"remote_dp_size": int(remote.get("dp_size", 1)),
|
||||
"tp_size": int(remote.get("tp_size", self._tp_size())),
|
||||
}
|
||||
|
||||
def _own_routing(self) -> Dict[str, Any]:
|
||||
"""``_remote_routing`` input describing this (decode orchestrator) replica
|
||||
-- the remote side for a prefill request."""
|
||||
return {
|
||||
"dp_rank": self._dp_rank(),
|
||||
"dp_size": self._dp_size(),
|
||||
"tp_size": self._tp_size(),
|
||||
}
|
||||
|
||||
# ---- request shaping (PD connector protocol) ----
|
||||
|
||||
def _dual_ids(
|
||||
self, request: Any, peer: Optional[Dict[str, Any]]
|
||||
) -> Tuple[str, str]:
|
||||
"""Compute the (dual-address request_id, transfer_id) for this request.
|
||||
|
||||
``prepare_prefill_request`` and ``prepare_decode_request`` are two
|
||||
independent, stateless calls for the same request, so both ids are
|
||||
derived deterministically (hash of a stable per-request seed) — no
|
||||
per-request backend state.
|
||||
"""
|
||||
prefill_zmq = (peer or {}).get("mori_zmq_address")
|
||||
decode_zmq = self._zmq_address
|
||||
if not prefill_zmq:
|
||||
raise ValueError(
|
||||
"MoRIIO peer is missing 'mori_zmq_address': the selected prefill "
|
||||
"replica did not publish its address (is MoRIIOConnector "
|
||||
"configured on the prefill deployment?)."
|
||||
)
|
||||
if not decode_zmq:
|
||||
raise ValueError(
|
||||
"MoRIIO decode zmq address is not set: setup() must run on this "
|
||||
"engine before requests are shaped."
|
||||
)
|
||||
# The incoming request_id (always populated -- OpenAI models default it
|
||||
# to a uuid) is the seed. Both prepare_* calls run on the same request
|
||||
# object, so they agree; uniqueness per request is inherited from it.
|
||||
seed = str(request.request_id)
|
||||
# 32 hex chars (the trailing uid _PREFILL_ZMQ_RE / _DECODE_ZMQ_RE
|
||||
# anchor on); a hash of the seed, so both prepare_* calls agree.
|
||||
uid = hashlib.sha256(seed.encode()).hexdigest()[:32]
|
||||
# Wire format consumed by vLLM's MoRIIO connector.
|
||||
request_id = f"{_PREFILL_PREFIX}{prefill_zmq}{_DECODE_PREFIX}{decode_zmq}_{uid}"
|
||||
transfer_id = f"{_TRANSFER_PREFIX}-{uid}"
|
||||
return request_id, transfer_id
|
||||
|
||||
def prepare_prefill_request(
|
||||
self, *, request: "RequestType", peer: Optional[Dict[str, Any]]
|
||||
) -> "RequestType":
|
||||
request_id, transfer_id = self._dual_ids(request, peer)
|
||||
prefill_request = request.model_copy(deep=True)
|
||||
# The dual-address id (peer zmq encoded in it) must reach the engine:
|
||||
# setting request_id explicitly makes the LLMServer pipeline preserve it
|
||||
# (not clobber it with the Serve id) and the engine copies it into the
|
||||
# X-Request-Id header that vLLM's MoRIIO connector parses.
|
||||
prefill_request.request_id = request_id
|
||||
prefill_request.kv_transfer_params = {
|
||||
**base_prefill_kv_transfer_params(),
|
||||
"transfer_id": transfer_id,
|
||||
# The prefill engine's remote is the decode (this orchestrator).
|
||||
**self._remote_routing(self._own_routing()),
|
||||
}
|
||||
clamp_request_to_single_token(prefill_request)
|
||||
return prefill_request
|
||||
|
||||
def prepare_decode_request(
|
||||
self,
|
||||
*,
|
||||
request: "RequestType",
|
||||
peer: Optional[Dict[str, Any]],
|
||||
prefill_response: Optional[Any],
|
||||
) -> "RequestType":
|
||||
request_id, transfer_id = self._dual_ids(request, peer)
|
||||
decode_request = request.model_copy(deep=True)
|
||||
decode_request.request_id = request_id
|
||||
# The decode engine's remote is the selected prefill peer.
|
||||
remote_routing = self._remote_routing(peer or {})
|
||||
|
||||
if not self._read_mode:
|
||||
# WRITE: prefill pushes KV; decode just needs do_remote_prefill + the
|
||||
# shared transfer_id (no block ids -- they are pushed, not pulled).
|
||||
decode_request.kv_transfer_params = {
|
||||
"do_remote_prefill": True,
|
||||
"do_remote_decode": False,
|
||||
"remote_engine_id": None,
|
||||
"remote_block_ids": None,
|
||||
"transfer_id": transfer_id,
|
||||
**remote_routing,
|
||||
}
|
||||
return decode_request
|
||||
|
||||
# READ: decode PULLS KV; forward the remote_block_ids / remote_engine_id
|
||||
# the prefill engine returned on its response. If absent (e.g. prompt <
|
||||
# block_size / full prefix hit), fall back to a local recompute.
|
||||
prefill_kv_params = getattr(prefill_response, "kv_transfer_params", None)
|
||||
params = dict(prefill_kv_params) if prefill_kv_params else {}
|
||||
if params.get("remote_block_ids") and params.get("remote_engine_id"):
|
||||
params.setdefault("transfer_id", transfer_id)
|
||||
params["do_remote_prefill"] = True
|
||||
params["do_remote_decode"] = False
|
||||
# Address the prefill peer's (dp_rank, dp_size, tp) workers.
|
||||
params.update(remote_routing)
|
||||
decode_request.kv_transfer_params = params
|
||||
else:
|
||||
logger.warning(
|
||||
"[MORI][READ] prefill returned no remote_block_ids/remote_engine_id "
|
||||
"(kv_transfer_params=%s); decode will recompute locally.",
|
||||
prefill_kv_params,
|
||||
)
|
||||
decode_request.kv_transfer_params = None
|
||||
return decode_request
|
||||
@@ -0,0 +1,94 @@
|
||||
import copy
|
||||
from typing import TYPE_CHECKING, List
|
||||
|
||||
from ray.llm._internal.serve.engines.vllm.kv_transfer.base import (
|
||||
BaseConnectorBackend,
|
||||
)
|
||||
from ray.llm._internal.serve.engines.vllm.kv_transfer.factory import (
|
||||
KVConnectorBackendFactory,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
|
||||
|
||||
|
||||
class MultiConnectorBackend(BaseConnectorBackend):
|
||||
"""Wraps multiple sub-connectors.
|
||||
|
||||
The P/D protocol (``prepare_prefill_request`` / ``prepare_decode_request`` and
|
||||
the ``requires_peer_binding`` / ``concurrent_handoff`` policy) is delegated to
|
||||
the *first* (top-most) sub-connector listed in ``connectors`` — that
|
||||
connector's policy governs request shaping and handoff for the group. Each
|
||||
sub-connector's ``setup()`` still runs.
|
||||
"""
|
||||
|
||||
def __init__(self, llm_config: "LLMConfig"):
|
||||
super().__init__(llm_config)
|
||||
self._connector_backends: List[BaseConnectorBackend] = []
|
||||
|
||||
def setup(self) -> None:
|
||||
"""Setup all connectors listed in the kv_transfer_config."""
|
||||
kv_transfer_config = self.kv_transfer_config
|
||||
connectors = kv_transfer_config.get("kv_connector_extra_config", {}).get(
|
||||
"connectors", []
|
||||
)
|
||||
if not connectors:
|
||||
# Fail fast at setup rather than with a cryptic error when the
|
||||
# orchestrator later delegates to a (missing) top-most sub-connector.
|
||||
raise ValueError(
|
||||
"MultiConnector requires at least one sub-connector in "
|
||||
"kv_connector_extra_config.connectors."
|
||||
)
|
||||
|
||||
for connector in connectors:
|
||||
connector_backend_str = connector.get("kv_connector")
|
||||
if connector_backend_str is None:
|
||||
raise ValueError("kv_connector is not set in the connector")
|
||||
|
||||
if connector_backend_str == "MultiConnector":
|
||||
raise ValueError(
|
||||
"Nesting MultiConnector within MultiConnector is not supported."
|
||||
)
|
||||
|
||||
# Merge parent config with connector-specific config
|
||||
sub_llm_config = copy.deepcopy(self.llm_config)
|
||||
sub_llm_config.engine_kwargs["kv_transfer_config"] = {
|
||||
**{
|
||||
k: v
|
||||
for k, v in kv_transfer_config.items()
|
||||
if k != "kv_connector_extra_config"
|
||||
},
|
||||
**connector,
|
||||
}
|
||||
|
||||
# Use factory to get backend class lazily
|
||||
connector_backend = KVConnectorBackendFactory.create_backend(
|
||||
connector_backend_str, sub_llm_config
|
||||
)
|
||||
connector_backend.setup()
|
||||
self._connector_backends.append(connector_backend)
|
||||
|
||||
@property
|
||||
def _primary(self) -> BaseConnectorBackend:
|
||||
"""The top-most sub-connector, whose protocol governs the group."""
|
||||
if not self._connector_backends:
|
||||
raise ValueError(
|
||||
"MultiConnectorBackend has no sub-connectors; was setup() called?"
|
||||
)
|
||||
return self._connector_backends[0]
|
||||
|
||||
@property
|
||||
def requires_peer_binding(self) -> bool:
|
||||
return bool(self._connector_backends) and self._primary.requires_peer_binding
|
||||
|
||||
@property
|
||||
def concurrent_handoff(self) -> bool:
|
||||
return bool(self._connector_backends) and self._primary.concurrent_handoff
|
||||
|
||||
def prepare_prefill_request(self, *, request, peer):
|
||||
return self._primary.prepare_prefill_request(request=request, peer=peer)
|
||||
|
||||
def prepare_decode_request(self, *, request, peer, prefill_response):
|
||||
return self._primary.prepare_decode_request(
|
||||
request=request, peer=peer, prefill_response=prefill_response
|
||||
)
|
||||
@@ -0,0 +1,67 @@
|
||||
import os
|
||||
|
||||
import ray
|
||||
from ray.llm._internal.serve.engines.vllm.kv_transfer.base import (
|
||||
BaseConnectorBackend,
|
||||
DefaultPDProtocolMixin,
|
||||
)
|
||||
|
||||
|
||||
class NixlConnectorBackend(DefaultPDProtocolMixin, BaseConnectorBackend):
|
||||
def _set_side_channel_port(self):
|
||||
from vllm import envs as vllm_envs
|
||||
|
||||
if vllm_envs.is_set("VLLM_NIXL_SIDE_CHANNEL_PORT"):
|
||||
return
|
||||
|
||||
base_port = int(
|
||||
self.llm_config.experimental_configs.get(
|
||||
"NIXL_SIDE_CHANNEL_PORT_BASE", 20000
|
||||
)
|
||||
)
|
||||
port = base_port + self._compute_port_offset()
|
||||
os.environ["VLLM_NIXL_SIDE_CHANNEL_PORT"] = str(port)
|
||||
|
||||
def _set_side_channel_host(self):
|
||||
from vllm import envs as vllm_envs
|
||||
|
||||
if not vllm_envs.is_set("VLLM_NIXL_SIDE_CHANNEL_HOST"):
|
||||
# Use Ray's node IP (internal/cluster IP) instead of vLLM's
|
||||
# get_ip() which can return external/public IPs on hostNetwork
|
||||
# pods, causing cross-node NIXL handshakes to fail.
|
||||
os.environ["VLLM_NIXL_SIDE_CHANNEL_HOST"] = ray.util.get_node_ip_address()
|
||||
|
||||
def setup(self) -> None:
|
||||
"""Initialize the NIXL connector backend.
|
||||
|
||||
This method sets up the NIXL (NVIDIA Inference Xfer Library) connector by:
|
||||
1. Verifying that the required vLLM environment variables are supported
|
||||
2. Configuring the side channel port and host if not already set
|
||||
3. Creating a unique engine ID across replicas
|
||||
|
||||
The side channel is used for KV cache transfer between vLLM instances.
|
||||
|
||||
Raises:
|
||||
ValueError: If the current vLLM version doesn't support the required
|
||||
NIXL environment variables.
|
||||
"""
|
||||
from vllm import envs as vllm_envs
|
||||
|
||||
if (
|
||||
"VLLM_NIXL_SIDE_CHANNEL_PORT" not in vllm_envs.environment_variables
|
||||
or "VLLM_NIXL_SIDE_CHANNEL_HOST" not in vllm_envs.environment_variables
|
||||
):
|
||||
raise ValueError(
|
||||
"This vLLM version does not support VLLM_NIXL_SIDE_CHANNEL_PORT"
|
||||
"or VLLM_NIXL_SIDE_CHANNEL_HOST environment variable. It's likely"
|
||||
"that you are using an older version of vLLM."
|
||||
)
|
||||
|
||||
self._set_side_channel_port()
|
||||
self._set_side_channel_host()
|
||||
|
||||
# We need to overwrite the engine_id to make it unique across replicas.
|
||||
engine_id = self.kv_transfer_config.get("engine_id", self._get_unique_suffix())
|
||||
host = vllm_envs.VLLM_NIXL_SIDE_CHANNEL_HOST
|
||||
port = vllm_envs.VLLM_NIXL_SIDE_CHANNEL_PORT
|
||||
self.kv_transfer_config["engine_id"] = "-".join([engine_id, host, str(port)])
|
||||
@@ -0,0 +1,960 @@
|
||||
import argparse
|
||||
import dataclasses
|
||||
import inspect
|
||||
import json
|
||||
import typing
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
AsyncGenerator,
|
||||
Dict,
|
||||
List,
|
||||
Literal,
|
||||
Optional,
|
||||
Tuple,
|
||||
Union,
|
||||
)
|
||||
|
||||
from pydantic import BaseModel, field_validator
|
||||
from starlette.datastructures import State
|
||||
from starlette.requests import Request
|
||||
from vllm.engine.arg_utils import AsyncEngineArgs
|
||||
from vllm.entrypoints.openai.cli_args import FrontendArgs
|
||||
from vllm.entrypoints.openai.engine.protocol import ErrorResponse as VLLMErrorResponse
|
||||
|
||||
import ray
|
||||
from ray.llm._internal.common.callbacks.base import CallbackCtx
|
||||
from ray.llm._internal.common.utils.import_utils import try_import
|
||||
from ray.llm._internal.serve.core.configs.llm_config import (
|
||||
DiskMultiplexConfig,
|
||||
LLMConfig,
|
||||
)
|
||||
from ray.llm._internal.serve.core.configs.openai_api_models import (
|
||||
ChatCompletionRequest,
|
||||
ChatCompletionResponse,
|
||||
CompletionRequest,
|
||||
CompletionResponse,
|
||||
DetokenizeRequest,
|
||||
DetokenizeResponse,
|
||||
EmbeddingRequest,
|
||||
EmbeddingResponse,
|
||||
ErrorInfo,
|
||||
ErrorResponse,
|
||||
ScoreRequest,
|
||||
ScoreResponse,
|
||||
TokenizeRequest,
|
||||
TokenizeResponse,
|
||||
TranscriptionRequest,
|
||||
TranscriptionResponse,
|
||||
)
|
||||
from ray.llm._internal.serve.core.engine.protocol import LLMEngine
|
||||
from ray.llm._internal.serve.core.protocol import RawRequestInfo
|
||||
from ray.llm._internal.serve.engines.vllm.vllm_models import (
|
||||
VLLMEngineConfig,
|
||||
)
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
from ray.llm._internal.serve.routing_policies.kv_aware.kv_aware_router import (
|
||||
is_kv_aware,
|
||||
)
|
||||
from ray.llm._internal.serve.routing_policies.kv_aware.vllm.kv_events import (
|
||||
assign_replica_kv_events_endpoint,
|
||||
get_kv_event_routing_stats,
|
||||
)
|
||||
from ray.llm._internal.serve.routing_policies.kv_aware.vllm.token_tracking import (
|
||||
enable_token_tracking,
|
||||
)
|
||||
from ray.llm._internal.serve.utils.node_initialization_utils import (
|
||||
initialize_node,
|
||||
)
|
||||
from ray.util.placement_group import PlacementGroup
|
||||
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.engine.protocol import EngineClient
|
||||
from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
|
||||
from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion
|
||||
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
|
||||
from vllm.entrypoints.pooling.embed.serving import ServingEmbedding
|
||||
from vllm.entrypoints.pooling.scoring.serving import ServingScores
|
||||
from vllm.entrypoints.serve.tokenize.serving import OpenAIServingTokenization
|
||||
from vllm.entrypoints.speech_to_text.transcription.serving import (
|
||||
OpenAIServingTranscription,
|
||||
)
|
||||
|
||||
vllm = try_import("vllm")
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def _canonicalize_request_id_header(
|
||||
request: Any, raw_request_info: Optional[RawRequestInfo]
|
||||
) -> Optional[RawRequestInfo]:
|
||||
"""Ensure raw_request_info carries X-Request-Id == request.request_id so vLLM's
|
||||
OpenAI layer (which prefers the header) sees the authoritative request id.
|
||||
|
||||
Returns a RawRequestInfo (new or mutated) with a single, correctly-cased header.
|
||||
If the request has no request_id, this is a no-op and returns raw_request_info
|
||||
unchanged (so embeddings/score/transcription requests are unaffected).
|
||||
"""
|
||||
rid = getattr(request, "request_id", None)
|
||||
if not rid:
|
||||
return raw_request_info
|
||||
headers = dict(raw_request_info.headers) if raw_request_info is not None else {}
|
||||
# Drop any existing variant of the header (case- and separator-insensitive,
|
||||
# e.g. "X-Request-Id" or "x_request_id") before setting the canonical one.
|
||||
headers = {
|
||||
k: v
|
||||
for k, v in headers.items()
|
||||
if k.replace("_", "-").lower() != "x-request-id"
|
||||
}
|
||||
headers["x-request-id"] = str(rid)
|
||||
if raw_request_info is None:
|
||||
return RawRequestInfo(headers=headers)
|
||||
# Preserve any non-header fields RawRequestInfo carries (now or in the future).
|
||||
return dataclasses.replace(raw_request_info, headers=headers)
|
||||
|
||||
|
||||
def _convert_config_dicts(merged: dict) -> dict:
|
||||
"""Convert dict values to their proper vLLM config classes based on type hints.
|
||||
|
||||
vLLM's AsyncEngineArgs has fields like structured_outputs_config,
|
||||
compilation_config, etc. that expect dataclass instances. When users pass
|
||||
dicts for these fields, we need to convert them to the proper config classes
|
||||
so that default values are populated correctly.
|
||||
|
||||
Without this conversion, dicts get converted to argparse.Namespace objects
|
||||
which lack the default field values, causing AttributeError when vLLM code
|
||||
tries to access those fields.
|
||||
"""
|
||||
fields_by_name = {f.name: f for f in dataclasses.fields(AsyncEngineArgs)}
|
||||
|
||||
for key, value in list(merged.items()):
|
||||
if not isinstance(value, dict) or key not in fields_by_name:
|
||||
continue
|
||||
|
||||
hint = fields_by_name[key].type
|
||||
if isinstance(hint, str):
|
||||
continue
|
||||
|
||||
# Handle Optional[X] (Union[X, None]) -> X
|
||||
origin = typing.get_origin(hint)
|
||||
if origin is Union:
|
||||
args = typing.get_args(hint)
|
||||
hint = next((a for a in args if a is not type(None)), hint)
|
||||
|
||||
# Convert dict to dataclass if the field expects a dataclass type
|
||||
if isinstance(hint, type) and dataclasses.is_dataclass(hint):
|
||||
try:
|
||||
merged[key] = hint(**value)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Failed to convert {key} dict to {hint.__name__}: {e}. "
|
||||
"Using dict as-is."
|
||||
)
|
||||
|
||||
return merged
|
||||
|
||||
|
||||
def _dict_to_namespace(obj: Any) -> Any:
|
||||
"""Recursively converts dictionaries to argparse.Namespace."""
|
||||
if isinstance(obj, dict):
|
||||
return argparse.Namespace(**{k: _dict_to_namespace(v) for k, v in obj.items()})
|
||||
elif isinstance(obj, list):
|
||||
return [_dict_to_namespace(item) for item in obj]
|
||||
else:
|
||||
return obj
|
||||
|
||||
|
||||
def _get_vllm_engine_config(
|
||||
llm_config: LLMConfig,
|
||||
) -> Tuple["AsyncEngineArgs", "VllmConfig"]:
|
||||
engine_config = llm_config.get_engine_config()
|
||||
|
||||
# Resolve to local cache path if model was downloaded from S3/GCS mirror
|
||||
# Only do this if mirror_config was specified (intentional S3/GCS download)
|
||||
if engine_config.mirror_config:
|
||||
from ray.llm._internal.common.utils.download_utils import (
|
||||
get_model_location_on_disk,
|
||||
)
|
||||
|
||||
local_path = get_model_location_on_disk(engine_config.actual_hf_model_id)
|
||||
if local_path and local_path != engine_config.actual_hf_model_id:
|
||||
engine_config.hf_model_id = local_path
|
||||
logger.info(f"Resolved model from mirror to local path: {local_path}")
|
||||
|
||||
from vllm.usage.usage_lib import UsageContext
|
||||
|
||||
try:
|
||||
async_engine_args = vllm.engine.arg_utils.AsyncEngineArgs(
|
||||
**engine_config.get_initialization_kwargs()
|
||||
)
|
||||
vllm_engine_config = async_engine_args.create_engine_config(
|
||||
usage_context=UsageContext.OPENAI_API_SERVER
|
||||
)
|
||||
except Exception as e:
|
||||
# vLLM's ModelConfig is a pydantic dataclass; its ValidationError holds an
|
||||
# unpicklable ArgsKwargs and cannot cross the Ray task boundary. Re-raise as
|
||||
# a plain error so the real message propagates instead of a pickling failure.
|
||||
raise RuntimeError(f"Failed to create vLLM engine config: {e}") from None
|
||||
return async_engine_args, vllm_engine_config
|
||||
|
||||
|
||||
def _clear_current_platform_cache():
|
||||
"""Clear the cache of the current platform.
|
||||
|
||||
vllm current has an lru cache for getting device compatibility
|
||||
that will not have the correct returned value if
|
||||
CUDA_VISIBLE_DEVICES is not set properly. In RayLLM eventually
|
||||
when we want to create the engine the env will be set properly,
|
||||
but till then, upon the import of vllm somewhere
|
||||
(which is a mystery) the lru cache will have the wrong value.
|
||||
This function will clear the cache so that the next time the
|
||||
cache is accessed, it will be re-evaluated.
|
||||
|
||||
Related issues:
|
||||
https://github.com/vllm-project/vllm/issues/8402
|
||||
https://github.com/vllm-project/vllm/issues/7890
|
||||
"""
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
# This check is just to future proof this implementation
|
||||
# in case vllm removes their lru_cache decorator
|
||||
if hasattr(current_platform.get_device_capability, "cache_clear"):
|
||||
logger.info("Clearing the current platform cache ...")
|
||||
current_platform.get_device_capability.cache_clear()
|
||||
|
||||
|
||||
class VLLMSleepConfig(BaseModel):
|
||||
"""vLLM-specific configuration for sleep operation."""
|
||||
|
||||
level: int = 1
|
||||
"""Sleep level:
|
||||
- Level 1: Offload weights to CPU RAM, discard KV cache
|
||||
- Level 2: Discard both model weights and KV cache (deeper sleep)
|
||||
"""
|
||||
|
||||
@field_validator("level")
|
||||
@classmethod
|
||||
def validate_level(cls, v: Any) -> int:
|
||||
if v not in (1, 2):
|
||||
raise ValueError("level must be 1 or 2")
|
||||
return v
|
||||
|
||||
|
||||
class VLLMWakeupConfig(BaseModel):
|
||||
"""vLLM-specific configuration for wakeup operation."""
|
||||
|
||||
tags: Optional[List[str]] = None
|
||||
"""Optional tags to selectively wake up components:
|
||||
- "weights": Restore model weights only
|
||||
- "kv_cache": Restore KV cache only
|
||||
- None: Restore everything
|
||||
"""
|
||||
|
||||
@field_validator("tags")
|
||||
@classmethod
|
||||
def validate_tags(cls, v: Any) -> Optional[List[str]]:
|
||||
if v is not None:
|
||||
valid_tags = {"weights", "kv_cache"}
|
||||
for tag in v:
|
||||
if tag not in valid_tags:
|
||||
raise ValueError(
|
||||
f"Invalid tag '{tag}'. Must be one of: {valid_tags}"
|
||||
)
|
||||
return v
|
||||
|
||||
|
||||
class VLLMPauseConfig(BaseModel):
|
||||
"""vLLM-specific configuration for pause operation."""
|
||||
|
||||
mode: Literal["abort", "wait", "keep"] = "abort"
|
||||
"""Pause mode:
|
||||
- "abort" (default): Abort all in-flight requests immediately.
|
||||
- "wait": Wait for in-flight requests to complete before pausing.
|
||||
- "keep": Freeze requests in queue; they resume on resume_generation().
|
||||
"""
|
||||
|
||||
clear_cache: bool = True
|
||||
"""Whether to clear KV and prefix caches after draining.
|
||||
Set to False to preserve cache for faster resume.
|
||||
"""
|
||||
|
||||
|
||||
class VLLMEngine(LLMEngine):
|
||||
def __init__(
|
||||
self,
|
||||
llm_config: LLMConfig,
|
||||
):
|
||||
"""Create a vLLM Engine class
|
||||
|
||||
Args:
|
||||
llm_config: The llm configuration for this engine
|
||||
"""
|
||||
super().__init__(llm_config)
|
||||
|
||||
self.llm_config = llm_config
|
||||
|
||||
if vllm is None:
|
||||
raise ImportError(
|
||||
"vLLM is not installed. Please install it with `pip install ray[llm]`."
|
||||
)
|
||||
assign_replica_kv_events_endpoint(self.llm_config)
|
||||
self.llm_config.setup_engine_backend()
|
||||
|
||||
self._running = False
|
||||
# Routing stats advertised to Serve's request router; populated in
|
||||
# start() once the engine's KV-events endpoint is bound.
|
||||
self._routing_stats: Dict[str, Any] = {}
|
||||
|
||||
# vLLM Integration points. Will be set through .start()
|
||||
self._engine_client = None
|
||||
self._oai_models: Optional["OpenAIServingModels"] = None
|
||||
self._oai_serving_chat: Optional["OpenAIServingChat"] = None
|
||||
self._oai_serving_completion: Optional["OpenAIServingCompletion"] = None
|
||||
self._oai_serving_embedding: Optional["ServingEmbedding"] = None
|
||||
self._oai_serving_transcription: Optional["OpenAIServingTranscription"] = None
|
||||
self._oai_serving_scores: Optional["ServingScores"] = None
|
||||
self._oai_serving_tokenization: Optional["OpenAIServingTokenization"] = None
|
||||
|
||||
async def build_asgi_app(self):
|
||||
from vllm.entrypoints.openai.api_server import build_app, init_app_state
|
||||
|
||||
supported_tasks = ("generate",)
|
||||
if hasattr(self._engine_client, "get_supported_tasks"):
|
||||
supported_tasks = await self._engine_client.get_supported_tasks()
|
||||
|
||||
# Pass model_config so vLLM mounts the pooling routers (/pooling, /classify,
|
||||
# /embed, /score) on the native ASGI app to enable direct streaming for pooling
|
||||
# classify, embed, and score.
|
||||
app = build_app(
|
||||
self._vllm_args,
|
||||
supported_tasks=supported_tasks,
|
||||
model_config=self._engine_client.model_config,
|
||||
)
|
||||
await init_app_state(
|
||||
self._engine_client,
|
||||
app.state,
|
||||
self._vllm_args,
|
||||
supported_tasks=supported_tasks,
|
||||
)
|
||||
return app
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start the vLLM engine.
|
||||
|
||||
If the engine is already running, do nothing.
|
||||
"""
|
||||
|
||||
if self._running:
|
||||
# The engine is already running!
|
||||
logger.info("Skipping engine restart because the engine is already running")
|
||||
return
|
||||
|
||||
from vllm.entrypoints.openai.api_server import init_app_state
|
||||
|
||||
callback = self.llm_config.get_or_create_callback()
|
||||
await callback.run_callback("on_before_node_init")
|
||||
if callback.ctx.run_init_node:
|
||||
await initialize_node(self.llm_config)
|
||||
await callback.run_callback("on_after_node_init")
|
||||
|
||||
(
|
||||
vllm_engine_args,
|
||||
vllm_frontend_args,
|
||||
vllm_engine_config,
|
||||
) = self._prepare_engine_config(callback.ctx)
|
||||
|
||||
# Apply checkpoint info to the llm_config.
|
||||
# This is needed for capturing model capabilities
|
||||
# (e.g. supports vision, etc.) on the llm_config.
|
||||
config = self.llm_config.get_engine_config()
|
||||
self.llm_config.apply_checkpoint_info(
|
||||
vllm_engine_config.model_config.model,
|
||||
trust_remote_code=config.trust_remote_code,
|
||||
)
|
||||
|
||||
self._engine_client = self._start_async_llm_engine(
|
||||
vllm_engine_args,
|
||||
vllm_engine_config,
|
||||
callback.ctx.placement_group,
|
||||
)
|
||||
|
||||
state = State()
|
||||
# TODO (Kourosh): There might be some variables that needs protection?
|
||||
merged = vllm_frontend_args.__dict__ | vllm_engine_args.__dict__
|
||||
|
||||
# Convert dict values to proper vLLM config classes (e.g., StructuredOutputsConfig)
|
||||
# so that default field values are populated correctly.
|
||||
merged = _convert_config_dicts(merged)
|
||||
|
||||
args = _dict_to_namespace(merged)
|
||||
self._vllm_args = args
|
||||
|
||||
# Query supported tasks from the engine so init_app_state initializes the correct serving objects.
|
||||
# Without this, vLLM falls back to 'generate' only.
|
||||
init_kwargs: dict[str, Any] = dict(
|
||||
state=state,
|
||||
args=args,
|
||||
)
|
||||
if "supported_tasks" in inspect.signature(init_app_state).parameters:
|
||||
if hasattr(self._engine_client, "get_supported_tasks"):
|
||||
supported_tasks = await self._engine_client.get_supported_tasks()
|
||||
init_kwargs["supported_tasks"] = supported_tasks
|
||||
|
||||
if "vllm_config" in inspect.signature(init_app_state).parameters:
|
||||
init_kwargs["vllm_config"] = vllm_engine_config
|
||||
|
||||
await init_app_state(self._engine_client, **init_kwargs)
|
||||
|
||||
self._oai_models = getattr(state, "openai_serving_models", None)
|
||||
self._oai_serving_chat = getattr(state, "openai_serving_chat", None)
|
||||
self._oai_serving_completion = getattr(state, "openai_serving_completion", None)
|
||||
self._oai_serving_embedding = getattr(state, "serving_embedding", None)
|
||||
self._oai_serving_transcription = getattr(
|
||||
state, "openai_serving_transcription", None
|
||||
)
|
||||
self._oai_serving_scores = getattr(state, "serving_scores", None)
|
||||
self._oai_serving_tokenization = getattr(
|
||||
state, "openai_serving_tokenization", None
|
||||
)
|
||||
|
||||
self._validate_openai_serving_models()
|
||||
self._validate_engine_client()
|
||||
|
||||
self._routing_stats = get_kv_event_routing_stats(
|
||||
self.llm_config,
|
||||
vllm_engine_config.cache_config.block_size,
|
||||
vllm_engine_config.scheduler_config.max_num_batched_tokens,
|
||||
)
|
||||
|
||||
self._running = True
|
||||
|
||||
logger.info("Started vLLM engine.")
|
||||
|
||||
def routing_stats(self) -> Dict[str, Any]:
|
||||
"""Returns KV event and replay endpoints for KV-aware routing."""
|
||||
return self._routing_stats
|
||||
|
||||
def _validate_openai_serving_models(self):
|
||||
assert self._oai_models is not None, "oai_models is not initialized"
|
||||
assert hasattr(
|
||||
self._oai_models, "lora_requests"
|
||||
), "oai_models must have a lora_requests attribute"
|
||||
assert hasattr(
|
||||
self._oai_models, "load_lora_adapter"
|
||||
), "oai_models must have a load_lora_adapter attribute"
|
||||
|
||||
@staticmethod
|
||||
def _make_error(message: str) -> ErrorResponse:
|
||||
return ErrorResponse(
|
||||
error=ErrorInfo(message=message, type="invalid_request_error", code=400)
|
||||
)
|
||||
|
||||
def _validate_openai_serving_chat(self) -> Optional[ErrorResponse]:
|
||||
if self._oai_serving_chat is None:
|
||||
return self._make_error(
|
||||
"This model does not support the 'generate' task. "
|
||||
"The chat completion endpoint is not available for this model."
|
||||
)
|
||||
|
||||
def _validate_openai_serving_completion(self) -> Optional[ErrorResponse]:
|
||||
if self._oai_serving_completion is None:
|
||||
return self._make_error(
|
||||
"This model does not support the 'generate' task. "
|
||||
"The completion endpoint is not available for this model."
|
||||
)
|
||||
|
||||
def _validate_openai_serving_embedding(self) -> Optional[ErrorResponse]:
|
||||
if self._oai_serving_embedding is None:
|
||||
return self._make_error(
|
||||
"This model does not support the 'embed' task. "
|
||||
"The embedding endpoint is not available for this model."
|
||||
)
|
||||
|
||||
def _validate_openai_serving_transcription(self) -> Optional[ErrorResponse]:
|
||||
if self._oai_serving_transcription is None:
|
||||
return self._make_error(
|
||||
"This model does not support the 'transcription' task. "
|
||||
"The transcription endpoint is not available for this model."
|
||||
)
|
||||
|
||||
def _validate_openai_serving_scores(self) -> Optional[ErrorResponse]:
|
||||
if self._oai_serving_scores is None:
|
||||
return self._make_error(
|
||||
"This model does not support the 'score' task. "
|
||||
"The score endpoint is not available for this model."
|
||||
)
|
||||
|
||||
def _validate_openai_serving_tokenization(self) -> Optional[ErrorResponse]:
|
||||
if self._oai_serving_tokenization is None:
|
||||
return self._make_error(
|
||||
"This model does not support the 'tokenization' task. "
|
||||
"The tokenization endpoint is not available for this model."
|
||||
)
|
||||
|
||||
def _validate_engine_client(self):
|
||||
assert hasattr(
|
||||
self._engine_client, "check_health"
|
||||
), "engine_client must have a check_health attribute"
|
||||
|
||||
def _prepare_engine_config(
|
||||
self, callback_ctx: CallbackCtx
|
||||
) -> Tuple["AsyncEngineArgs", "FrontendArgs", "VllmConfig"]:
|
||||
"""Prepare the engine config to start the engine.
|
||||
|
||||
Args:
|
||||
callback_ctx: The callback context.
|
||||
|
||||
Returns:
|
||||
A tuple of:
|
||||
engine_args: The vLLM's internal engine arguments that is flattened.
|
||||
frontend_args: The vLLM's internal frontend arguments that is flattened.
|
||||
engine_config: The vLLM's internal engine config that is nested.
|
||||
"""
|
||||
|
||||
engine_config: VLLMEngineConfig = self.llm_config.get_engine_config()
|
||||
|
||||
# If the backend is anything other than CPU, we need to create the
|
||||
# engine config on a task with hardware access.
|
||||
if engine_config.accelerator.requires_remote_initialization:
|
||||
accelerator = engine_config.accelerator
|
||||
accelerator_type = self.llm_config.accelerator_type
|
||||
|
||||
# Initialize options required for the remote task and hardware backend
|
||||
remote_options = {
|
||||
"num_cpus": 0,
|
||||
"runtime_env": callback_ctx.runtime_env,
|
||||
"scheduling_strategy": PlacementGroupSchedulingStrategy(
|
||||
placement_group=callback_ctx.placement_group,
|
||||
),
|
||||
**accelerator.get_remote_options(accelerator_type),
|
||||
}
|
||||
|
||||
ref = (
|
||||
ray.remote(_get_vllm_engine_config)
|
||||
.options(**remote_options)
|
||||
.remote(self.llm_config)
|
||||
)
|
||||
vllm_engine_args, vllm_engine_config = ray.get(ref)
|
||||
else:
|
||||
vllm_engine_args, vllm_engine_config = _get_vllm_engine_config(
|
||||
self.llm_config
|
||||
)
|
||||
|
||||
vllm_frontend_args = FrontendArgs(**engine_config.frontend_kwargs)
|
||||
return vllm_engine_args, vllm_frontend_args, vllm_engine_config
|
||||
|
||||
def _start_async_llm_engine(
|
||||
self,
|
||||
vllm_engine_args: "AsyncEngineArgs",
|
||||
vllm_engine_config: "VllmConfig",
|
||||
placement_group: PlacementGroup,
|
||||
) -> "EngineClient":
|
||||
"""Creates an async LLM engine from the engine arguments."""
|
||||
|
||||
from vllm.v1.engine.async_llm import AsyncLLM
|
||||
from vllm.v1.executor.abstract import Executor
|
||||
|
||||
vllm_engine_config.parallel_config.placement_group = placement_group
|
||||
|
||||
_clear_current_platform_cache()
|
||||
|
||||
custom_stat_loggers = None
|
||||
if self.llm_config.log_engine_metrics:
|
||||
from vllm.v1.metrics.ray_wrappers import RayPrometheusStatLogger
|
||||
|
||||
# V1 AsyncLLM does not yet support add_logger: https://github.com/vllm-project/vllm/issues/17702
|
||||
# Use `disable_log_stats: False` and `log_engine_metrics: False` as
|
||||
# a workaround to enable PrometheusStatLogger instead.
|
||||
custom_stat_loggers = [RayPrometheusStatLogger]
|
||||
|
||||
executor_class = Executor.get_class(vllm_engine_config)
|
||||
logger.info(f"Using executor class: {executor_class}")
|
||||
# Report per-request token progress to the deployment's KV router actor,
|
||||
# but only on KV-aware deployments: elsewhere the actor never exists and
|
||||
# resolving it per request would block the engine's event loop.
|
||||
engine_cls = AsyncLLM
|
||||
if is_kv_aware(self.llm_config):
|
||||
engine_cls = enable_token_tracking(AsyncLLM)
|
||||
engine_client = engine_cls(
|
||||
vllm_config=vllm_engine_config,
|
||||
executor_class=executor_class,
|
||||
log_requests=vllm_engine_args.enable_log_requests,
|
||||
log_stats=not vllm_engine_args.disable_log_stats,
|
||||
stat_loggers=custom_stat_loggers,
|
||||
)
|
||||
|
||||
return engine_client
|
||||
|
||||
async def resolve_lora(self, disk_lora_model: DiskMultiplexConfig):
|
||||
from vllm.entrypoints.serve.lora.protocol import LoadLoRAAdapterRequest
|
||||
|
||||
self._validate_openai_serving_models()
|
||||
|
||||
if disk_lora_model.model_id in self._oai_models.lora_requests:
|
||||
# Lora is already loaded, return
|
||||
return
|
||||
|
||||
lora_request = await self._oai_models.load_lora_adapter( # type: ignore[attr-defined]
|
||||
request=LoadLoRAAdapterRequest(
|
||||
lora_name=disk_lora_model.model_id,
|
||||
lora_path=disk_lora_model.local_path,
|
||||
)
|
||||
)
|
||||
|
||||
if isinstance(lora_request, VLLMErrorResponse):
|
||||
raise ValueError(f"Failed to load lora model: {lora_request.error.message}")
|
||||
|
||||
@staticmethod
|
||||
def _make_error_response(
|
||||
serving: Any,
|
||||
exc: Exception,
|
||||
) -> ErrorResponse:
|
||||
"""Convert an exception to an ErrorResponse and map exception types to
|
||||
the appropriate HTTP status codes (e.g. VLLMValidationError -> 400).
|
||||
"""
|
||||
try:
|
||||
vllm_error = serving.create_error_response(exc)
|
||||
return ErrorResponse(error=ErrorInfo(**vllm_error.error.model_dump()))
|
||||
except Exception:
|
||||
raise exc # re-raise the original so it surfaces as a 500
|
||||
|
||||
async def chat(
|
||||
self,
|
||||
request: ChatCompletionRequest,
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[Union[str, ChatCompletionResponse, ErrorResponse], None]:
|
||||
if error := self._validate_openai_serving_chat():
|
||||
yield error
|
||||
return
|
||||
|
||||
raw_request_info = _canonicalize_request_id_header(request, raw_request_info)
|
||||
raw_request: Optional[Request] = RawRequestInfo.to_starlette_request_optional(
|
||||
raw_request_info
|
||||
)
|
||||
try:
|
||||
chat_response = await self._oai_serving_chat.create_chat_completion( # type: ignore[attr-defined]
|
||||
request,
|
||||
raw_request=raw_request,
|
||||
)
|
||||
except ValueError as e:
|
||||
yield self._make_error_response(self._oai_serving_chat, e)
|
||||
return
|
||||
|
||||
if isinstance(chat_response, AsyncGenerator):
|
||||
async for response in chat_response:
|
||||
if not isinstance(response, str):
|
||||
raise ValueError(
|
||||
f"Expected create_chat_completion to return a stream of strings, got an item with type {type(response)}"
|
||||
)
|
||||
yield response
|
||||
else:
|
||||
if isinstance(chat_response, VLLMErrorResponse):
|
||||
yield ErrorResponse(error=ErrorInfo(**chat_response.error.model_dump()))
|
||||
else:
|
||||
yield ChatCompletionResponse(**chat_response.model_dump())
|
||||
|
||||
async def completions(
|
||||
self,
|
||||
request: CompletionRequest,
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[Union[str, CompletionResponse, ErrorResponse], None]:
|
||||
if error := self._validate_openai_serving_completion():
|
||||
yield error
|
||||
return
|
||||
|
||||
raw_request_info = _canonicalize_request_id_header(request, raw_request_info)
|
||||
raw_request: Optional[Request] = RawRequestInfo.to_starlette_request_optional(
|
||||
raw_request_info
|
||||
)
|
||||
try:
|
||||
completion_response = await self._oai_serving_completion.create_completion( # type: ignore[attr-defined]
|
||||
request,
|
||||
raw_request=raw_request,
|
||||
)
|
||||
except ValueError as e:
|
||||
yield self._make_error_response(self._oai_serving_completion, e)
|
||||
return
|
||||
|
||||
if isinstance(completion_response, AsyncGenerator):
|
||||
async for response in completion_response:
|
||||
if not isinstance(response, str):
|
||||
raise ValueError(
|
||||
f"Expected create_completion to return a stream of strings, got an item with type {type(response)}"
|
||||
)
|
||||
yield response
|
||||
else:
|
||||
if isinstance(completion_response, VLLMErrorResponse):
|
||||
yield ErrorResponse(
|
||||
error=ErrorInfo(**completion_response.error.model_dump())
|
||||
)
|
||||
else:
|
||||
yield CompletionResponse(**completion_response.model_dump())
|
||||
|
||||
async def embeddings(
|
||||
self,
|
||||
request: EmbeddingRequest,
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[Union[EmbeddingResponse, ErrorResponse], None]:
|
||||
if error := self._validate_openai_serving_embedding():
|
||||
yield error
|
||||
return
|
||||
|
||||
raw_request_info = _canonicalize_request_id_header(request, raw_request_info)
|
||||
raw_request: Optional[Request] = RawRequestInfo.to_starlette_request_optional(
|
||||
raw_request_info
|
||||
)
|
||||
try:
|
||||
embedding_response = await self._oai_serving_embedding(
|
||||
request,
|
||||
raw_request=raw_request,
|
||||
)
|
||||
except ValueError as e:
|
||||
yield self._make_error_response(self._oai_serving_embedding, e)
|
||||
return
|
||||
|
||||
# vLLM 0.18+ returns a starlette Response object
|
||||
content = json.loads(embedding_response.body)
|
||||
yield EmbeddingResponse(**content)
|
||||
|
||||
async def transcriptions(
|
||||
self,
|
||||
request: TranscriptionRequest,
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[Union[str, TranscriptionResponse, ErrorResponse], None]:
|
||||
if error := self._validate_openai_serving_transcription():
|
||||
yield error
|
||||
return
|
||||
|
||||
# Extract audio data from the request file
|
||||
audio_data = await request.file.read()
|
||||
|
||||
raw_request_info = _canonicalize_request_id_header(request, raw_request_info)
|
||||
raw_request: Optional[Request] = RawRequestInfo.to_starlette_request_optional(
|
||||
raw_request_info
|
||||
)
|
||||
try:
|
||||
transcription_response = await self._oai_serving_transcription.create_transcription( # type: ignore[attr-defined]
|
||||
audio_data,
|
||||
request,
|
||||
raw_request=raw_request,
|
||||
)
|
||||
except ValueError as e:
|
||||
yield self._make_error_response(self._oai_serving_transcription, e)
|
||||
return
|
||||
|
||||
if isinstance(transcription_response, AsyncGenerator):
|
||||
async for response in transcription_response:
|
||||
if not isinstance(response, str):
|
||||
raise ValueError(
|
||||
f"Expected create_transcription to return a stream of strings, got an item with type {type(response)}"
|
||||
)
|
||||
yield response
|
||||
else:
|
||||
if isinstance(transcription_response, VLLMErrorResponse):
|
||||
yield ErrorResponse(
|
||||
error=ErrorInfo(**transcription_response.error.model_dump())
|
||||
)
|
||||
else:
|
||||
yield TranscriptionResponse(**transcription_response.model_dump())
|
||||
|
||||
async def score(
|
||||
self,
|
||||
request: ScoreRequest,
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[Union[ScoreResponse, ErrorResponse], None]:
|
||||
if error := self._validate_openai_serving_scores():
|
||||
yield error
|
||||
return
|
||||
|
||||
raw_request_info = _canonicalize_request_id_header(request, raw_request_info)
|
||||
raw_request: Optional[Request] = RawRequestInfo.to_starlette_request_optional(
|
||||
raw_request_info
|
||||
)
|
||||
try:
|
||||
assert self._oai_serving_scores is not None
|
||||
score_response = await self._oai_serving_scores(
|
||||
request,
|
||||
raw_request=raw_request,
|
||||
)
|
||||
except ValueError as e:
|
||||
yield self._make_error_response(self._oai_serving_scores, e)
|
||||
return
|
||||
|
||||
content = json.loads(score_response.body)
|
||||
yield ScoreResponse(**content)
|
||||
|
||||
async def tokenize(
|
||||
self,
|
||||
request: TokenizeRequest,
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[Union[TokenizeResponse, ErrorResponse], None]:
|
||||
if error := self._validate_openai_serving_tokenization():
|
||||
yield error
|
||||
return
|
||||
|
||||
raw_request_info = _canonicalize_request_id_header(request, raw_request_info)
|
||||
raw_request: Optional[Request] = RawRequestInfo.to_starlette_request_optional(
|
||||
raw_request_info
|
||||
)
|
||||
try:
|
||||
tokenize_response = await self._oai_serving_tokenization.create_tokenize(
|
||||
request,
|
||||
raw_request=raw_request,
|
||||
)
|
||||
except ValueError as e:
|
||||
yield self._make_error_response(self._oai_serving_tokenization, e)
|
||||
return
|
||||
|
||||
if isinstance(tokenize_response, VLLMErrorResponse):
|
||||
yield ErrorResponse(error=ErrorInfo(**tokenize_response.error.model_dump()))
|
||||
else:
|
||||
yield TokenizeResponse(**tokenize_response.model_dump())
|
||||
|
||||
async def detokenize(
|
||||
self,
|
||||
request: DetokenizeRequest,
|
||||
raw_request_info: Optional[RawRequestInfo] = None,
|
||||
) -> AsyncGenerator[Union[DetokenizeResponse, ErrorResponse], None]:
|
||||
if error := self._validate_openai_serving_tokenization():
|
||||
yield error
|
||||
return
|
||||
|
||||
raw_request_info = _canonicalize_request_id_header(request, raw_request_info)
|
||||
raw_request: Optional[Request] = RawRequestInfo.to_starlette_request_optional(
|
||||
raw_request_info
|
||||
)
|
||||
try:
|
||||
detokenize_response = (
|
||||
await self._oai_serving_tokenization.create_detokenize(
|
||||
request,
|
||||
raw_request=raw_request,
|
||||
)
|
||||
)
|
||||
except ValueError as e:
|
||||
yield self._make_error_response(self._oai_serving_tokenization, e)
|
||||
return
|
||||
|
||||
if isinstance(detokenize_response, VLLMErrorResponse):
|
||||
yield ErrorResponse(
|
||||
error=ErrorInfo(**detokenize_response.error.model_dump())
|
||||
)
|
||||
else:
|
||||
yield DetokenizeResponse(**detokenize_response.model_dump())
|
||||
|
||||
async def check_health(self) -> None:
|
||||
assert self._engine_client is not None, "engine_client is not initialized"
|
||||
|
||||
try:
|
||||
await self._engine_client.check_health()
|
||||
except BaseException as e:
|
||||
logger.error("Healthcheck failed. The replica will be restarted")
|
||||
raise e from None
|
||||
|
||||
async def reset_prefix_cache(self) -> None:
|
||||
assert self._engine_client is not None, "engine_client is not initialized"
|
||||
await self._engine_client.reset_prefix_cache()
|
||||
|
||||
async def sleep(self, **kwargs: Any) -> None:
|
||||
"""Put the vLLM engine to sleep.
|
||||
|
||||
Args:
|
||||
**kwargs: Options parsed into VLLMSleepConfig.
|
||||
- level (int): Sleep level (1 or 2). Default 1.
|
||||
"""
|
||||
assert self._engine_client is not None, "engine_client is not initialized"
|
||||
config = VLLMSleepConfig(**kwargs)
|
||||
await self._engine_client.sleep(level=config.level)
|
||||
|
||||
async def wakeup(self, **kwargs: Any) -> None:
|
||||
"""Wake up the vLLM engine from sleep mode.
|
||||
|
||||
Args:
|
||||
**kwargs: Options parsed into VLLMWakeupConfig.
|
||||
- tags (List[str], optional): Components to wake up.
|
||||
"""
|
||||
assert self._engine_client is not None, "engine_client is not initialized"
|
||||
config = VLLMWakeupConfig(**kwargs)
|
||||
await self._engine_client.wake_up(tags=config.tags)
|
||||
|
||||
async def is_sleeping(self) -> bool:
|
||||
"""Check whether the vLLM engine is currently sleeping.
|
||||
|
||||
Returns:
|
||||
True if the engine is sleeping, False otherwise.
|
||||
"""
|
||||
assert self._engine_client is not None, "engine_client is not initialized"
|
||||
return await self._engine_client.is_sleeping()
|
||||
|
||||
async def pause(self, **kwargs: Any) -> None:
|
||||
"""Pause generation on the vLLM engine.
|
||||
|
||||
This halts generation/encoding requests while keeping model weights
|
||||
in GPU memory. New requests are blocked until resume is called.
|
||||
|
||||
Args:
|
||||
**kwargs: Options parsed into VLLMPauseConfig.
|
||||
- mode (str): "abort" (default), "wait", or "keep".
|
||||
- clear_cache (bool): Clear KV cache after draining. Default True.
|
||||
"""
|
||||
assert self._engine_client is not None, "engine_client is not initialized"
|
||||
config = VLLMPauseConfig(**kwargs)
|
||||
await self._engine_client.pause_generation(
|
||||
mode=config.mode,
|
||||
clear_cache=config.clear_cache,
|
||||
)
|
||||
|
||||
async def resume(self, **kwargs: Any) -> None:
|
||||
"""Resume generation on the vLLM engine after pause.
|
||||
|
||||
Args:
|
||||
**kwargs: Reserved for future options.
|
||||
"""
|
||||
assert self._engine_client is not None, "engine_client is not initialized"
|
||||
await self._engine_client.resume_generation()
|
||||
|
||||
async def is_paused(self) -> bool:
|
||||
"""Check whether the vLLM engine is currently paused.
|
||||
|
||||
Returns:
|
||||
True if the engine is paused, False otherwise.
|
||||
"""
|
||||
assert self._engine_client is not None, "engine_client is not initialized"
|
||||
return await self._engine_client.is_paused()
|
||||
|
||||
async def start_profile(self) -> None:
|
||||
assert self._engine_client is not None, "engine_client is not initialized"
|
||||
await self._engine_client.start_profile()
|
||||
|
||||
async def stop_profile(self) -> None:
|
||||
assert self._engine_client is not None, "engine_client is not initialized"
|
||||
await self._engine_client.stop_profile()
|
||||
|
||||
async def collective_rpc(
|
||||
self,
|
||||
method: str,
|
||||
timeout: Optional[float] = None,
|
||||
args: tuple = (),
|
||||
kwargs: Optional[dict] = None,
|
||||
) -> list:
|
||||
"""Execute a collective RPC call on all vLLM workers.
|
||||
|
||||
This is used for RLHF workflows where a trainer needs to execute
|
||||
methods on all TP/PP workers (e.g., for weight synchronization).
|
||||
|
||||
Args:
|
||||
method: Name of the worker method to execute.
|
||||
timeout: Maximum time in seconds to wait for execution.
|
||||
args: Positional arguments to pass to the worker method.
|
||||
kwargs: Keyword arguments to pass to the worker method.
|
||||
|
||||
Returns:
|
||||
A list containing the results from each worker.
|
||||
"""
|
||||
assert self._engine_client is not None, "engine_client is not initialized"
|
||||
return await self._engine_client.collective_rpc(
|
||||
method=method,
|
||||
timeout=timeout,
|
||||
args=args,
|
||||
kwargs=kwargs or {},
|
||||
)
|
||||
@@ -0,0 +1,361 @@
|
||||
import dataclasses
|
||||
import os
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from pydantic import ConfigDict, Field, PrivateAttr, field_validator, model_validator
|
||||
from vllm.engine.arg_utils import AsyncEngineArgs
|
||||
from vllm.entrypoints.openai.cli_args import FrontendArgs
|
||||
|
||||
from ray.llm._internal.common.base_pydantic import BaseModelExtended
|
||||
from ray.llm._internal.common.placement import PlacementGroupConfig
|
||||
from ray.llm._internal.common.utils.cloud_utils import CloudMirrorConfig, is_remote_path
|
||||
from ray.llm._internal.common.utils.import_utils import try_import
|
||||
from ray.llm._internal.serve.constants import (
|
||||
ALLOW_NEW_PLACEMENT_GROUPS_IN_DEPLOYMENT,
|
||||
ENV_VARS_TO_PROPAGATE,
|
||||
)
|
||||
from ray.llm._internal.serve.core.configs.accelerators import (
|
||||
AcceleratorBackend,
|
||||
AnyAcceleratorConfig,
|
||||
CPUAccelerator,
|
||||
CPUConfig,
|
||||
GPUAccelerator,
|
||||
TPUAccelerator,
|
||||
TPUConfig,
|
||||
format_ray_accelerator_resource,
|
||||
)
|
||||
from ray.llm._internal.serve.core.configs.llm_config import (
|
||||
AcceleratorType,
|
||||
LLMConfig,
|
||||
)
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
from ray.util.placement_group import (
|
||||
PlacementGroup,
|
||||
get_current_placement_group,
|
||||
placement_group_table,
|
||||
)
|
||||
|
||||
# The key for the kv_transfer_params in the internal metadata.
|
||||
KV_TRANSFER_PARAMS_KEY = "kv_transfer_params"
|
||||
vllm = try_import("vllm")
|
||||
logger = get_logger(__name__)
|
||||
|
||||
# Executor backend constants
|
||||
EXECUTOR_BACKEND_RAY = "ray"
|
||||
EXECUTOR_BACKEND_MP = "mp"
|
||||
|
||||
|
||||
class VLLMEngineConfig(BaseModelExtended):
|
||||
model_config = ConfigDict(
|
||||
use_enum_values=True,
|
||||
)
|
||||
|
||||
model_id: str = Field(
|
||||
description="The identifier for the model. This is the id that will be used to query the model.",
|
||||
)
|
||||
hf_model_id: Optional[str] = Field(
|
||||
None, description="The Hugging Face model identifier."
|
||||
)
|
||||
mirror_config: Optional[CloudMirrorConfig] = Field(
|
||||
None,
|
||||
description="Configuration for cloud storage mirror. This is for where the weights are downloaded from.",
|
||||
)
|
||||
accelerator_type: Optional[AcceleratorType] = Field(
|
||||
None,
|
||||
description="The type of accelerator to use. This is used to determine the placement group strategy.",
|
||||
)
|
||||
accelerator_config: Optional[AnyAcceleratorConfig] = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Hardware-specific configuration parameters for the chosen accelerator. "
|
||||
"The expected schema is dynamically typed based on the 'kind' discriminator."
|
||||
),
|
||||
)
|
||||
|
||||
placement_group_config: Optional[Dict[str, Any]] = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Ray placement group configuration for scheduling vLLM engine workers. "
|
||||
"Defines resource bundles and placement strategy for multi-node deployments. "
|
||||
"Defaults to PACK strategy with automatic bundle generation based on TP/PP sizes."
|
||||
),
|
||||
)
|
||||
|
||||
@field_validator("placement_group_config")
|
||||
@classmethod
|
||||
def validate_placement_group_config(cls, value):
|
||||
if value is None:
|
||||
return None
|
||||
# Validate through PlacementGroupConfig, then dump back to dict
|
||||
validated = PlacementGroupConfig(**value)
|
||||
return validated.model_dump(exclude_unset=True)
|
||||
|
||||
runtime_env: Optional[Dict[str, Any]] = None
|
||||
engine_kwargs: Dict[str, Any] = {}
|
||||
frontend_kwargs: Dict[str, Any] = {}
|
||||
|
||||
_accelerator: AcceleratorBackend = PrivateAttr(default=None)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _build_accelerator(self):
|
||||
"""Instantiates the accelerator backend based on the resolved config."""
|
||||
cfg = self.accelerator_config
|
||||
|
||||
if self.accelerator_type and isinstance(cfg, CPUConfig):
|
||||
raise ValueError(
|
||||
f"accelerator_type='{self.accelerator_type}' cannot be used with "
|
||||
"CPU-only configurations. Either remove accelerator_type, or provide an accelerator_config."
|
||||
)
|
||||
|
||||
# LLMConfig has already resolved and validated accelerator_config
|
||||
if isinstance(cfg, TPUConfig):
|
||||
self._accelerator = TPUAccelerator(cfg)
|
||||
elif isinstance(cfg, CPUConfig):
|
||||
self._accelerator = CPUAccelerator()
|
||||
else:
|
||||
# Default to GPU if it's GPUConfig or isn't set
|
||||
self._accelerator = GPUAccelerator()
|
||||
return self
|
||||
|
||||
@property
|
||||
def accelerator(self) -> AcceleratorBackend:
|
||||
return self._accelerator
|
||||
|
||||
@property
|
||||
def actual_hf_model_id(self) -> str:
|
||||
return self.hf_model_id or self.model_id
|
||||
|
||||
@property
|
||||
def trust_remote_code(self) -> bool:
|
||||
return self.engine_kwargs.get("trust_remote_code", False)
|
||||
|
||||
def get_initialization_kwargs(self) -> dict:
|
||||
"""
|
||||
Get kwargs that will be actually passed to the LLMInitializer
|
||||
constructor.
|
||||
"""
|
||||
engine_kwargs = self.engine_kwargs.copy()
|
||||
|
||||
if "model" in engine_kwargs or "served_model_name" in engine_kwargs:
|
||||
raise ValueError(
|
||||
"model or served_model_name is not allowed in engine_kwargs when using Ray Serve LLM. Please use `model_loading_config` in LLMConfig instead."
|
||||
)
|
||||
|
||||
engine_kwargs["model"] = self.actual_hf_model_id
|
||||
engine_kwargs["served_model_name"] = [self.model_id]
|
||||
|
||||
# Handle distributed_executor_backend based on backend type
|
||||
if isinstance(self.accelerator, CPUAccelerator):
|
||||
executor_backend = EXECUTOR_BACKEND_MP
|
||||
else:
|
||||
executor_backend = EXECUTOR_BACKEND_RAY
|
||||
|
||||
if (
|
||||
"distributed_executor_backend" in engine_kwargs
|
||||
and engine_kwargs["distributed_executor_backend"] != executor_backend
|
||||
and executor_backend == EXECUTOR_BACKEND_RAY
|
||||
):
|
||||
raise ValueError(
|
||||
"distributed_executor_backend != 'ray' is not allowed in engine_kwargs when using Ray Serve LLM Configs."
|
||||
)
|
||||
|
||||
engine_kwargs["distributed_executor_backend"] = executor_backend
|
||||
|
||||
if "enable_log_requests" not in engine_kwargs:
|
||||
engine_kwargs["enable_log_requests"] = False
|
||||
|
||||
return engine_kwargs
|
||||
|
||||
def get_runtime_env_with_local_env_vars(self) -> dict:
|
||||
runtime_env = self.runtime_env or {}
|
||||
runtime_env.setdefault("env_vars", {})
|
||||
env_vars = runtime_env["env_vars"]
|
||||
|
||||
# Propagate env vars to the runtime env
|
||||
for env_var in ENV_VARS_TO_PROPAGATE:
|
||||
if env_var in os.environ:
|
||||
env_vars[env_var] = os.getenv(env_var)
|
||||
|
||||
if "VLLM_RAY_PER_WORKER_GPUS" not in env_vars:
|
||||
fractional_gpu = self._detect_fractional_gpu_from_pg(
|
||||
self.placement_group_config
|
||||
)
|
||||
if fractional_gpu is not None:
|
||||
env_vars["VLLM_RAY_PER_WORKER_GPUS"] = str(fractional_gpu)
|
||||
return runtime_env
|
||||
|
||||
@classmethod
|
||||
def from_llm_config(cls, llm_config: LLMConfig) -> "VLLMEngineConfig":
|
||||
"""Converts the LLMConfig to a VLLMEngineConfig."""
|
||||
# Set up the model downloading configuration.
|
||||
hf_model_id, mirror_config = None, None
|
||||
if llm_config.model_loading_config.model_source is None:
|
||||
hf_model_id = llm_config.model_id
|
||||
elif isinstance(llm_config.model_loading_config.model_source, str):
|
||||
model_source = llm_config.model_loading_config.model_source
|
||||
if is_remote_path(model_source):
|
||||
# Remote URIs (s3://, gs://, …) are download addresses,
|
||||
# not HuggingFace IDs. Using the URI verbatim as
|
||||
# hf_model_id propagates the scheme and slashes into the
|
||||
# cache directory name (``models--s3:----bucket--…``).
|
||||
# Use the user-supplied model_id as the identifier and
|
||||
# treat the URI as a bucket mirror instead.
|
||||
hf_model_id = llm_config.model_id
|
||||
mirror_config = CloudMirrorConfig(bucket_uri=model_source)
|
||||
else:
|
||||
hf_model_id = model_source
|
||||
else:
|
||||
# If it's a CloudMirrorConfig (or subtype)
|
||||
mirror_config = llm_config.model_loading_config.model_source
|
||||
|
||||
all_engine_kwargs = llm_config.engine_kwargs.copy()
|
||||
engine_kwargs = {}
|
||||
frontend_kwargs = {}
|
||||
|
||||
# Get field names from dataclasses
|
||||
frontend_field_names = {
|
||||
field.name for field in dataclasses.fields(FrontendArgs)
|
||||
}
|
||||
async_engine_field_names = {
|
||||
field.name for field in dataclasses.fields(AsyncEngineArgs)
|
||||
}
|
||||
|
||||
for key, value in all_engine_kwargs.items():
|
||||
if key in frontend_field_names:
|
||||
frontend_kwargs[key] = value
|
||||
elif key in async_engine_field_names:
|
||||
engine_kwargs[key] = value
|
||||
else:
|
||||
raise ValueError(f"Unknown engine argument: {key}")
|
||||
|
||||
# placement_group_config is already validated and stored as dict in LLMConfig
|
||||
placement_group_config = llm_config.placement_group_config
|
||||
return VLLMEngineConfig(
|
||||
model_id=llm_config.model_id,
|
||||
hf_model_id=hf_model_id,
|
||||
mirror_config=mirror_config,
|
||||
accelerator_type=llm_config.accelerator_type,
|
||||
accelerator_config=llm_config.accelerator_config,
|
||||
engine_kwargs=engine_kwargs,
|
||||
frontend_kwargs=frontend_kwargs,
|
||||
runtime_env=llm_config.runtime_env,
|
||||
placement_group_config=placement_group_config,
|
||||
)
|
||||
|
||||
@property
|
||||
def tensor_parallel_degree(self) -> int:
|
||||
return self.engine_kwargs.get("tensor_parallel_size", 1)
|
||||
|
||||
@property
|
||||
def pipeline_parallel_degree(self) -> int:
|
||||
return self.engine_kwargs.get("pipeline_parallel_size", 1)
|
||||
|
||||
@property
|
||||
def num_devices(self) -> int:
|
||||
return self.tensor_parallel_degree * self.pipeline_parallel_degree
|
||||
|
||||
@property
|
||||
def placement_strategy(self) -> str:
|
||||
# Use custom strategy if placement_group_config is provided
|
||||
if self.placement_group_config:
|
||||
return self.placement_group_config.get("strategy", "PACK")
|
||||
# Default to PACK (cross-node best-effort placement)
|
||||
# DP deployments overridden to STRICT_PACK in Serve config
|
||||
return "PACK"
|
||||
|
||||
@property
|
||||
def placement_bundles(self) -> List[Dict[str, float]]:
|
||||
if self.placement_group_config:
|
||||
bundle_per_worker = self.placement_group_config.get("bundle_per_worker")
|
||||
|
||||
if bundle_per_worker is not None:
|
||||
# Expand bundle_per_worker to num_devices bundles
|
||||
bundles = []
|
||||
for _ in range(self.num_devices):
|
||||
bundle = bundle_per_worker.copy()
|
||||
if self.accelerator_type:
|
||||
res_key = format_ray_accelerator_resource(self.accelerator_type)
|
||||
bundle.setdefault(res_key, 0.001)
|
||||
bundles.append(bundle)
|
||||
return bundles
|
||||
|
||||
# Otherwise use explicit bundles list
|
||||
bundles = []
|
||||
explicit_bundles = self.placement_group_config.get("bundles") or []
|
||||
for bundle_dict in explicit_bundles:
|
||||
bundle = bundle_dict.copy()
|
||||
if self.accelerator_type:
|
||||
# Use setdefault to add accelerator hint WITHOUT overriding explicit user values
|
||||
res_key = format_ray_accelerator_resource(self.accelerator_type)
|
||||
bundle.setdefault(res_key, 0.001)
|
||||
bundles.append(bundle)
|
||||
return bundles
|
||||
|
||||
# Default bundles based on the accelerator backend.
|
||||
return self.accelerator.default_bundles(
|
||||
num_devices=self.num_devices, accelerator_type_str=self.accelerator_type
|
||||
)
|
||||
|
||||
def get_or_create_pg(self) -> PlacementGroup:
|
||||
"""Gets or creates a placement group.
|
||||
|
||||
If we are already in a placement group, return the existing placement group.
|
||||
Else, delegate PG creation to the accelerator backend.
|
||||
"""
|
||||
dp_rank = self.engine_kwargs.get("data_parallel_rank", None)
|
||||
pg = get_current_placement_group()
|
||||
if pg:
|
||||
logger.debug(
|
||||
"Using existing placement group %s, details: %s",
|
||||
pg.id,
|
||||
placement_group_table(pg),
|
||||
)
|
||||
else:
|
||||
if not ALLOW_NEW_PLACEMENT_GROUPS_IN_DEPLOYMENT:
|
||||
raise RuntimeError(
|
||||
"Creating new placement groups is not allowed. "
|
||||
"Change RAYLLM_ALLOW_NEW_PLACEMENT_GROUPS_IN_DEPLOYMENT "
|
||||
"if this is not intended."
|
||||
)
|
||||
name = "" if dp_rank is None else f"dp_{dp_rank}"
|
||||
|
||||
pg = self.accelerator.create_placement_group(
|
||||
bundles=self.placement_bundles,
|
||||
strategy=self.placement_strategy,
|
||||
name=name,
|
||||
accelerator_type_str=self.accelerator_type,
|
||||
)
|
||||
|
||||
logger.info(f"Using new placement group {pg}. {placement_group_table(pg)}")
|
||||
|
||||
return pg
|
||||
|
||||
@staticmethod
|
||||
def _detect_fractional_gpu_from_pg(
|
||||
placement_group_config: Optional[Dict[str, Any]]
|
||||
) -> Optional[float]:
|
||||
if not placement_group_config:
|
||||
return None
|
||||
|
||||
# Check bundle_per_worker first
|
||||
bundle_per_worker = placement_group_config.get("bundle_per_worker")
|
||||
if bundle_per_worker:
|
||||
gpu_value = bundle_per_worker.get("GPU", 0)
|
||||
if 0 < gpu_value < 1:
|
||||
return gpu_value
|
||||
return None
|
||||
|
||||
# Fall back to bundles list
|
||||
bundles = placement_group_config.get("bundles") or []
|
||||
|
||||
for bundle in bundles:
|
||||
if "GPU" not in bundle:
|
||||
continue
|
||||
|
||||
gpu_value = bundle["GPU"]
|
||||
if gpu_value <= 0 or gpu_value >= 1:
|
||||
return None
|
||||
|
||||
return gpu_value
|
||||
|
||||
return None
|
||||
@@ -0,0 +1,23 @@
|
||||
from ray.llm._internal.common.observability.logging_utils import (
|
||||
disable_datasets_logger,
|
||||
disable_vllm_custom_ops_logger_on_cpu_nodes,
|
||||
)
|
||||
from ray.llm._internal.common.observability.telemetry_utils import Once
|
||||
from ray.llm._internal.serve.observability.logging.setup import (
|
||||
setup_logging,
|
||||
)
|
||||
|
||||
_setup_observability_once = Once()
|
||||
|
||||
|
||||
def _setup_observability():
|
||||
setup_logging()
|
||||
disable_datasets_logger()
|
||||
disable_vllm_custom_ops_logger_on_cpu_nodes()
|
||||
|
||||
|
||||
def setup_observability():
|
||||
_setup_observability_once.do_once(_setup_observability)
|
||||
|
||||
|
||||
__all__ = ["setup_observability"]
|
||||
@@ -0,0 +1,39 @@
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from ray._common.filters import CoreContextFilter
|
||||
from ray.serve._private.logging_utils import ServeContextFilter
|
||||
|
||||
|
||||
def _setup_logger(logger_name: str):
|
||||
"""Setup logger given the logger name.
|
||||
|
||||
This function is idempotent and won't set up the same logger multiple times. It will
|
||||
Also skip the setup if Serve logger is already setup and has handlers.
|
||||
"""
|
||||
logger = logging.getLogger(logger_name)
|
||||
serve_logger = logging.getLogger("ray.serve")
|
||||
|
||||
# Skip setup if the logger already has handlers setup or if the parent (Serve
|
||||
# logger) has handlers.
|
||||
if logger.handlers or serve_logger.handlers:
|
||||
return
|
||||
|
||||
# Set up stream handler, which logs to console as plaintext.
|
||||
stream_handler = logging.StreamHandler()
|
||||
stream_handler.addFilter(CoreContextFilter())
|
||||
stream_handler.addFilter(ServeContextFilter())
|
||||
logger.addHandler(stream_handler)
|
||||
logger.setLevel(logging.INFO)
|
||||
logger.propagate = False
|
||||
|
||||
|
||||
def get_logger(name: Optional[str] = None):
|
||||
"""Get a structured logger inherited from the Ray Serve logger.
|
||||
|
||||
Loggers by default are logging to stdout, and are expected to be scraped by an
|
||||
external process.
|
||||
"""
|
||||
logger_name = f"ray.serve.{name}"
|
||||
_setup_logger(logger_name)
|
||||
return logging.getLogger(logger_name)
|
||||
@@ -0,0 +1,28 @@
|
||||
import logging
|
||||
|
||||
from ray._common.filters import CoreContextFilter
|
||||
from ray._common.formatters import JSONFormatter
|
||||
from ray.serve._private.logging_utils import ServeContextFilter
|
||||
|
||||
|
||||
def _configure_stdlib_logging():
|
||||
"""Configures stdlib root logger to make sure stdlib loggers (created as
|
||||
`logging.getLogger(...)`) are using Ray's `JSONFormatter` with Core and Serve
|
||||
context filters.
|
||||
"""
|
||||
|
||||
handler = logging.StreamHandler()
|
||||
handler.addFilter(CoreContextFilter())
|
||||
handler.addFilter(ServeContextFilter())
|
||||
handler.setFormatter(JSONFormatter())
|
||||
|
||||
root_logger = logging.getLogger()
|
||||
# NOTE: It's crucial we reset all the handlers of the root logger,
|
||||
# to make sure that logs aren't emitted twice
|
||||
root_logger.handlers = []
|
||||
root_logger.addHandler(handler)
|
||||
root_logger.setLevel(logging.INFO)
|
||||
|
||||
|
||||
def setup_logging():
|
||||
_configure_stdlib_logging()
|
||||
@@ -0,0 +1,77 @@
|
||||
import asyncio
|
||||
import time
|
||||
from typing import Dict
|
||||
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
from ray.util import metrics
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
_METRICS_LOOP_INTERVAL = 5 # 5 seconds
|
||||
EVENT_LOOP_LATENCY_HISTOGRAM_BOUNDARIES = [
|
||||
0.05,
|
||||
0.1,
|
||||
0.15,
|
||||
0.20,
|
||||
0.25,
|
||||
0.5,
|
||||
0.75,
|
||||
1.0,
|
||||
1.5,
|
||||
2.0,
|
||||
3.0,
|
||||
5.0,
|
||||
10.0,
|
||||
15.0,
|
||||
20.0,
|
||||
30.0,
|
||||
45.0,
|
||||
60.0,
|
||||
90.0,
|
||||
120.0,
|
||||
150.0,
|
||||
180.0,
|
||||
300.0,
|
||||
600.0,
|
||||
]
|
||||
|
||||
|
||||
def setup_event_loop_monitoring(
|
||||
loop: asyncio.AbstractEventLoop,
|
||||
scheduling_latency: metrics.Histogram,
|
||||
iterations: metrics.Counter,
|
||||
tasks: metrics.Gauge,
|
||||
tags: Dict[str, str],
|
||||
) -> asyncio.Task:
|
||||
return asyncio.create_task(
|
||||
_run_monitoring_loop(
|
||||
loop,
|
||||
schedule_latency=scheduling_latency,
|
||||
iterations=iterations,
|
||||
task_gauge=tasks,
|
||||
tags=tags,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def _run_monitoring_loop(
|
||||
loop: asyncio.AbstractEventLoop,
|
||||
schedule_latency: metrics.Histogram,
|
||||
iterations: metrics.Counter,
|
||||
task_gauge: metrics.Gauge,
|
||||
tags: Dict[str, str],
|
||||
) -> None:
|
||||
while loop.is_running():
|
||||
iterations.inc(1, tags)
|
||||
num_tasks = len(asyncio.all_tasks())
|
||||
task_gauge.set(num_tasks, tags)
|
||||
yield_time = time.monotonic()
|
||||
await asyncio.sleep(_METRICS_LOOP_INTERVAL)
|
||||
elapsed_time = time.monotonic() - yield_time
|
||||
|
||||
# Historically, Ray's implementation of histograms are extremely finicky
|
||||
# with non-positive values (https://github.com/ray-project/ray/issues/26698).
|
||||
# Technically it shouldn't be possible for this to be negative, add the
|
||||
# max just to be safe.
|
||||
latency = max(0.0, elapsed_time - _METRICS_LOOP_INTERVAL)
|
||||
schedule_latency.observe(latency, tags)
|
||||
@@ -0,0 +1,133 @@
|
||||
import asyncio
|
||||
import os
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from fastapi import FastAPI
|
||||
|
||||
from ray.llm._internal.serve.constants import ENABLE_VERBOSE_TELEMETRY
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
from ray.llm._internal.serve.observability.metrics.event_loop_monitoring import (
|
||||
EVENT_LOOP_LATENCY_HISTOGRAM_BOUNDARIES,
|
||||
setup_event_loop_monitoring,
|
||||
)
|
||||
from ray.llm._internal.serve.observability.metrics.fastapi_utils import (
|
||||
FASTAPI_API_SERVER_TAG_KEY,
|
||||
FASTAPI_BASE_HTTP_METRIC_TAG_KEYS,
|
||||
get_app_name,
|
||||
)
|
||||
from ray.llm._internal.serve.observability.metrics.middleware import (
|
||||
MeasureHTTPRequestMetricsMiddleware,
|
||||
)
|
||||
from ray.util import metrics
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
ray_llm_build_info_gauge = metrics.Gauge(
|
||||
"ray_llm_build_info",
|
||||
description="Metadata about the ray-llm build.",
|
||||
tag_keys=("git_commit",),
|
||||
)
|
||||
|
||||
|
||||
_HTTP_HANDLER_LATENCY_S_HISTOGRAM_BUCKETS = [
|
||||
0.01,
|
||||
0.05,
|
||||
0.1,
|
||||
0.25,
|
||||
0.5,
|
||||
0.75,
|
||||
1,
|
||||
1.5,
|
||||
2,
|
||||
5,
|
||||
10,
|
||||
30,
|
||||
60,
|
||||
120,
|
||||
300,
|
||||
]
|
||||
|
||||
|
||||
async def add_fastapi_event_loop_monitoring(app: FastAPI):
|
||||
tags = {FASTAPI_API_SERVER_TAG_KEY: get_app_name(app)}
|
||||
tag_keys = tuple(tags.keys())
|
||||
|
||||
# Store the task handle to prevent it from being garbage collected
|
||||
app.state.fastapi_event_loop_schedule_latency = metrics.Histogram(
|
||||
"fastapi_event_loop_schedule_latency",
|
||||
description="Latency of getting yielded control on the FastAPI event loop in seconds",
|
||||
boundaries=EVENT_LOOP_LATENCY_HISTOGRAM_BOUNDARIES,
|
||||
tag_keys=tag_keys,
|
||||
)
|
||||
app.state.fastapi_event_loop_monitoring_iterations = metrics.Counter(
|
||||
"fastapi_event_loop_monitoring_iterations",
|
||||
description="Number of times the FastAPI event loop has iterated to get anyscale_fastapi_event_loop_schedule_latency.",
|
||||
tag_keys=tag_keys,
|
||||
)
|
||||
app.state.fastapi_event_loop_monitoring_tasks = metrics.Gauge(
|
||||
"fastapi_event_loop_monitoring_tasks",
|
||||
description="Number of outstanding tasks on the FastAPI event loop.",
|
||||
tag_keys=tag_keys,
|
||||
)
|
||||
|
||||
app.state.fastapi_event_loop_schedule_latency_metrics_task = (
|
||||
setup_event_loop_monitoring(
|
||||
asyncio.get_running_loop(),
|
||||
app.state.fastapi_event_loop_schedule_latency,
|
||||
app.state.fastapi_event_loop_monitoring_iterations,
|
||||
app.state.fastapi_event_loop_monitoring_tasks,
|
||||
tags,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def add_http_metrics_middleware(app: FastAPI):
|
||||
if not ENABLE_VERBOSE_TELEMETRY:
|
||||
logger.debug(
|
||||
"ENABLE_VERBOSE_TELEMETRY is false, not setting up FastAPI telemetry"
|
||||
)
|
||||
return
|
||||
|
||||
logger.info("ENABLE_VERBOSE_TELEMETRY is true, setting up FastAPI telemetry")
|
||||
base_tag_keys = FASTAPI_BASE_HTTP_METRIC_TAG_KEYS
|
||||
|
||||
logger.debug("Setting up FastAPI telemetry")
|
||||
|
||||
app.state.http_requests_metrics = metrics.Counter(
|
||||
"http_requests",
|
||||
description=(
|
||||
"Total number of HTTP requests by status code, handler and method."
|
||||
),
|
||||
tag_keys=base_tag_keys,
|
||||
)
|
||||
# NOTE: Custom decorators are not applied to histogram-based metrics
|
||||
# to make sure we can keep cardinality of those in check
|
||||
app.state.http_requests_latency_metrics = metrics.Histogram(
|
||||
"http_request_duration_seconds",
|
||||
description="Duration in seconds of HTTP requests.",
|
||||
boundaries=_HTTP_HANDLER_LATENCY_S_HISTOGRAM_BUCKETS,
|
||||
tag_keys=base_tag_keys,
|
||||
)
|
||||
|
||||
app.add_middleware(MeasureHTTPRequestMetricsMiddleware)
|
||||
|
||||
logger.debug("Setting up FastAPI telemetry completed")
|
||||
|
||||
|
||||
async def set_ray_llm_build_info():
|
||||
git_commit = os.environ.get("GIT_COMMIT")
|
||||
if git_commit:
|
||||
tags = {"git_commit": git_commit}
|
||||
ray_llm_build_info_gauge.set(1, tags)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def metrics_lifespan(app: FastAPI):
|
||||
"""Lifespan for a FastAPI app that sets up metrics observability."""
|
||||
|
||||
if ENABLE_VERBOSE_TELEMETRY:
|
||||
await add_fastapi_event_loop_monitoring(app)
|
||||
|
||||
await set_ray_llm_build_info()
|
||||
|
||||
yield
|
||||
@@ -0,0 +1,26 @@
|
||||
"""This file contains constants and utility functions for FastAPI."""
|
||||
|
||||
from fastapi import FastAPI
|
||||
|
||||
# These tag keys are used in metrics for the FastAPI app.
|
||||
FASTAPI_HTTP_RESPONSE_CODE_TAG_KEY = "code"
|
||||
FASTAPI_HTTP_HANDLER_TAG_KEY = "handler"
|
||||
FASTAPI_HTTP_METHOD_TAG_KEY = "method"
|
||||
FASTAPI_HTTP_PATH_TAG_KEY = "path"
|
||||
FASTAPI_HTTP_USER_ID_TAG_KEY = "user_id"
|
||||
FASTAPI_API_SERVER_TAG_KEY = "api_server"
|
||||
|
||||
FASTAPI_BASE_HTTP_METRIC_TAG_KEYS = (
|
||||
FASTAPI_HTTP_RESPONSE_CODE_TAG_KEY,
|
||||
FASTAPI_HTTP_HANDLER_TAG_KEY,
|
||||
FASTAPI_HTTP_METHOD_TAG_KEY,
|
||||
FASTAPI_HTTP_PATH_TAG_KEY,
|
||||
FASTAPI_HTTP_USER_ID_TAG_KEY,
|
||||
FASTAPI_API_SERVER_TAG_KEY,
|
||||
)
|
||||
|
||||
|
||||
def get_app_name(app: FastAPI) -> str:
|
||||
"""Gets the FastAPI app name."""
|
||||
|
||||
return getattr(app.state, "name", "unknown")
|
||||
@@ -0,0 +1,150 @@
|
||||
import time
|
||||
from asyncio import CancelledError
|
||||
from typing import Dict, Optional
|
||||
|
||||
from fastapi import FastAPI
|
||||
from starlette.requests import Request
|
||||
from starlette.types import Message
|
||||
|
||||
from ray.llm._internal.serve.core.ingress.middleware import (
|
||||
get_request_id,
|
||||
get_user_id,
|
||||
)
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
from ray.llm._internal.serve.observability.metrics.fastapi_utils import (
|
||||
FASTAPI_API_SERVER_TAG_KEY,
|
||||
FASTAPI_HTTP_HANDLER_TAG_KEY,
|
||||
FASTAPI_HTTP_METHOD_TAG_KEY,
|
||||
FASTAPI_HTTP_PATH_TAG_KEY,
|
||||
FASTAPI_HTTP_RESPONSE_CODE_TAG_KEY,
|
||||
FASTAPI_HTTP_USER_ID_TAG_KEY,
|
||||
get_app_name,
|
||||
)
|
||||
from ray.serve._private.thirdparty.get_asgi_route_name import _get_route_name
|
||||
|
||||
logger = get_logger("ray.serve")
|
||||
|
||||
|
||||
class MeasureHTTPRequestMetricsMiddleware:
|
||||
"""Measures and stores HTTP request metrics."""
|
||||
|
||||
def __init__(self, app):
|
||||
self.app = app
|
||||
|
||||
async def __call__(self, scope, receive, send):
|
||||
if scope["type"] not in ("http", "websocket"):
|
||||
await self.app(scope, receive, send)
|
||||
|
||||
# If the status_code isn't set by send_wrapper,
|
||||
# we should consider that an error.
|
||||
status_code = 500
|
||||
send_wrapper_failed_exc_info = "Status code was never set by send_wrapper."
|
||||
exception_info = send_wrapper_failed_exc_info
|
||||
|
||||
async def send_wrapper(message: Message) -> None:
|
||||
"""Wraps the send message.
|
||||
|
||||
Enables this middleware to access the response headers.
|
||||
"""
|
||||
|
||||
nonlocal status_code, exception_info
|
||||
if message["type"] == "http.response.start":
|
||||
status_code = message.get("status", 500)
|
||||
|
||||
# Clear the send_wrapper_failed_exc_info.
|
||||
if exception_info == send_wrapper_failed_exc_info:
|
||||
exception_info = None
|
||||
|
||||
await send(message)
|
||||
|
||||
request = Request(scope)
|
||||
req_id = get_request_id(request)
|
||||
now = time.monotonic()
|
||||
try:
|
||||
logger.info(f"Starting handling of the request {req_id}")
|
||||
await self.app(scope, receive, send_wrapper)
|
||||
|
||||
except CancelledError as ce:
|
||||
status_code = -1
|
||||
exception_info = ce
|
||||
raise
|
||||
|
||||
except BaseException as e:
|
||||
status_code = 500
|
||||
exception_info = e
|
||||
raise
|
||||
|
||||
finally:
|
||||
duration_s = time.monotonic() - now
|
||||
|
||||
tags = _get_tags(request, status_code, request.app)
|
||||
# NOTE: Custom decorators are not applied to histogram-based metrics
|
||||
# to make sure we can keep cardinality of those in check
|
||||
truncated_tags = {
|
||||
**tags,
|
||||
FASTAPI_HTTP_USER_ID_TAG_KEY: "truncated",
|
||||
}
|
||||
|
||||
request.app.state.http_requests_metrics.inc(1, tags)
|
||||
request.app.state.http_requests_latency_metrics.observe(
|
||||
duration_s, truncated_tags
|
||||
)
|
||||
|
||||
extra_context = {
|
||||
"status_code": status_code,
|
||||
"duration_ms": duration_s * 1000,
|
||||
}
|
||||
|
||||
if status_code >= 400:
|
||||
log = logger.error if status_code >= 500 else logger.warning
|
||||
log(
|
||||
f"Handling of the request {req_id} failed",
|
||||
exc_info=exception_info,
|
||||
extra={"ray_serve_extra_fields": extra_context},
|
||||
)
|
||||
elif status_code == -1:
|
||||
logger.info(
|
||||
f"Handling of the request {req_id} have been cancelled",
|
||||
extra={"ray_serve_extra_fields": extra_context},
|
||||
)
|
||||
else:
|
||||
logger.info(
|
||||
f"Handling of the request {req_id} successfully completed",
|
||||
extra={"ray_serve_extra_fields": extra_context},
|
||||
)
|
||||
|
||||
|
||||
def _get_route_details(scope: dict) -> Optional[str]:
|
||||
"""
|
||||
Function to retrieve Starlette route from scope.
|
||||
TODO: there is currently no way to retrieve http.route from
|
||||
a starlette application from scope.
|
||||
See: https://github.com/encode/starlette/pull/804
|
||||
Args:
|
||||
scope: A Starlette scope
|
||||
Returns:
|
||||
A string containing the route or None
|
||||
"""
|
||||
# Delegate to Serve's shared route-name resolver, which walks the route tree
|
||||
# and handles FastAPI >= 0.137 `_IncludedRouter` nodes (added by
|
||||
# `include_router`) that have no `.path` attribute of their own. Accessing
|
||||
# `.path` on such a node previously raised AttributeError here (#64245).
|
||||
return _get_route_name(scope, scope["app"].routes)
|
||||
|
||||
|
||||
def _get_tags(request: Request, status_code: int, app: FastAPI) -> Dict[str, str]:
|
||||
"""Generates tags for the request's metrics."""
|
||||
|
||||
route = str(_get_route_details(request.scope)) or "unknown"
|
||||
path = str(request.url.path) or "unknown"
|
||||
method = str(request.method) or "unknown"
|
||||
user_id = str(get_user_id(request) or "unknown")
|
||||
|
||||
return {
|
||||
FASTAPI_API_SERVER_TAG_KEY: get_app_name(app),
|
||||
FASTAPI_HTTP_RESPONSE_CODE_TAG_KEY: str(status_code),
|
||||
FASTAPI_HTTP_PATH_TAG_KEY: path,
|
||||
FASTAPI_HTTP_HANDLER_TAG_KEY: route,
|
||||
FASTAPI_HTTP_METHOD_TAG_KEY: method,
|
||||
FASTAPI_HTTP_USER_ID_TAG_KEY: user_id,
|
||||
}
|
||||
@@ -0,0 +1,144 @@
|
||||
import time
|
||||
from enum import Enum
|
||||
from typing import (
|
||||
AsyncGenerator,
|
||||
Callable,
|
||||
List,
|
||||
Set,
|
||||
TypeVar,
|
||||
)
|
||||
|
||||
from ray.util import metrics
|
||||
|
||||
# Histogram buckets for short-range latencies measurements:
|
||||
# Min=1ms, Max=30s
|
||||
#
|
||||
# NOTE: Number of buckets have to be bounded (and not exceed 30)
|
||||
# to avoid overloading metrics sub-system
|
||||
SHORT_RANGE_LATENCY_HISTOGRAM_BUCKETS_MS: List[float] = [
|
||||
1,
|
||||
5,
|
||||
10,
|
||||
25,
|
||||
50,
|
||||
100,
|
||||
150,
|
||||
250,
|
||||
500,
|
||||
1000,
|
||||
1500,
|
||||
2500,
|
||||
5000,
|
||||
7500,
|
||||
10000,
|
||||
20000,
|
||||
30000,
|
||||
]
|
||||
|
||||
# Histogram buckets for long-range latencies measurements:
|
||||
# Min=10ms, Max=300s
|
||||
LONG_RANGE_LATENCY_HISTOGRAM_BUCKETS_MS = [
|
||||
x * 10 for x in SHORT_RANGE_LATENCY_HISTOGRAM_BUCKETS_MS
|
||||
]
|
||||
|
||||
|
||||
class ClockUnit(int, Enum):
|
||||
ms = 1000
|
||||
s = 1
|
||||
|
||||
|
||||
class MsClock:
|
||||
"""A clock that tracks intervals in milliseconds"""
|
||||
|
||||
def __init__(self, unit: ClockUnit = ClockUnit.ms):
|
||||
self.reset()
|
||||
self.unit = unit.value
|
||||
self.start_time = time.perf_counter()
|
||||
|
||||
def reset(self):
|
||||
self.start_time = time.perf_counter()
|
||||
|
||||
def interval(self):
|
||||
return (time.perf_counter() - self.start_time) * self.unit
|
||||
|
||||
def reset_interval(self):
|
||||
interval = self.interval()
|
||||
self.reset()
|
||||
return interval
|
||||
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
class InstrumentTokenAsyncGenerator:
|
||||
"""This class instruments an asynchronous generator.
|
||||
|
||||
It gathers 3 metrics:
|
||||
1. Time to first time
|
||||
2. Time between tokens
|
||||
3. Total completion time
|
||||
|
||||
Usage:
|
||||
|
||||
@InstrumentTokenAsyncGenerator("my_special_fn")
|
||||
async def to_instrument():
|
||||
yield ...
|
||||
"""
|
||||
|
||||
all_instrument_names: Set[str] = set()
|
||||
|
||||
def __init__(
|
||||
self, generator_name: str, latency_histogram_buckets: List[float] = None
|
||||
):
|
||||
self.generator_name = f"rayllm_{generator_name}"
|
||||
|
||||
target_latency_histogram_buckets = (
|
||||
latency_histogram_buckets or SHORT_RANGE_LATENCY_HISTOGRAM_BUCKETS_MS
|
||||
)
|
||||
|
||||
assert (
|
||||
self.generator_name not in self.all_instrument_names
|
||||
), "This generator name was already used elsewhere. Please specify another one."
|
||||
self.all_instrument_names.add(self.generator_name)
|
||||
|
||||
self.token_latency_histogram = metrics.Histogram(
|
||||
f"{self.generator_name}_per_token_latency_ms",
|
||||
f"Generator metrics for {self.generator_name}",
|
||||
boundaries=target_latency_histogram_buckets,
|
||||
)
|
||||
|
||||
self.first_token_latency_histogram = metrics.Histogram(
|
||||
f"{self.generator_name}_first_token_latency_ms",
|
||||
f"Generator metrics for {self.generator_name}",
|
||||
boundaries=target_latency_histogram_buckets,
|
||||
)
|
||||
self.total_latency_histogram = metrics.Histogram(
|
||||
f"{self.generator_name}_total_latency_ms",
|
||||
f"Generator metrics for {self.generator_name}",
|
||||
boundaries=target_latency_histogram_buckets,
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self, async_generator_fn: Callable[..., AsyncGenerator[T, None]]
|
||||
) -> Callable[..., AsyncGenerator[T, None]]:
|
||||
async def new_gen(*args, **kwargs):
|
||||
interval_clock = MsClock()
|
||||
total_clock = MsClock()
|
||||
is_first_token = True
|
||||
try:
|
||||
async for x in async_generator_fn(*args, **kwargs):
|
||||
if is_first_token:
|
||||
self.first_token_latency_histogram.observe(
|
||||
total_clock.interval()
|
||||
)
|
||||
interval_clock.reset()
|
||||
is_first_token = False
|
||||
else:
|
||||
self.token_latency_histogram.observe(
|
||||
interval_clock.reset_interval()
|
||||
)
|
||||
yield x
|
||||
finally:
|
||||
self.total_latency_histogram.observe(total_clock.interval())
|
||||
|
||||
return new_gen
|
||||
@@ -0,0 +1,353 @@
|
||||
import hashlib
|
||||
import random
|
||||
import time
|
||||
from enum import Enum
|
||||
from typing import TYPE_CHECKING, Callable, Dict, Optional, Sequence
|
||||
|
||||
import ray
|
||||
from ray import serve
|
||||
from ray._common.constants import HEAD_NODE_RESOURCE_NAME
|
||||
from ray._common.usage.usage_lib import (
|
||||
get_hardware_usages_to_report,
|
||||
record_extra_usage_tag,
|
||||
)
|
||||
from ray.llm._internal.common.base_pydantic import BaseModelExtended
|
||||
from ray.llm._internal.common.observability.telemetry_utils import DEFAULT_GPU_TYPE
|
||||
from ray.llm._internal.common.utils.lora_utils import get_lora_model_ids
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
|
||||
|
||||
LLM_SERVE_TELEMETRY_NAMESPACE = "llm_serve_telemetry"
|
||||
LLM_SERVE_TELEMETRY_ACTOR_NAME = "llm_serve_telemetry"
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class TelemetryTags(str, Enum):
|
||||
"""Telemetry tags for LLM SERVE."""
|
||||
|
||||
LLM_SERVE_SERVE_MULTIPLE_MODELS = "LLM_SERVE_SERVE_MULTIPLE_MODELS"
|
||||
LLM_SERVE_SERVE_MULTIPLE_APPS = "LLM_SERVE_SERVE_MULTIPLE_APPS"
|
||||
LLM_SERVE_LORA_BASE_MODELS = "LLM_SERVE_LORA_BASE_MODELS"
|
||||
LLM_SERVE_INITIAL_NUM_LORA_ADAPTERS = "LLM_SERVE_INITIAL_NUM_LORA_ADAPTERS"
|
||||
LLM_SERVE_AUTOSCALING_ENABLED_MODELS = "LLM_SERVE_AUTOSCALING_ENABLED_MODELS"
|
||||
LLM_SERVE_AUTOSCALING_MIN_REPLICAS = "LLM_SERVE_AUTOSCALING_MIN_REPLICAS"
|
||||
LLM_SERVE_AUTOSCALING_MAX_REPLICAS = "LLM_SERVE_AUTOSCALING_MAX_REPLICAS"
|
||||
LLM_SERVE_TENSOR_PARALLEL_DEGREE = "LLM_SERVE_TENSOR_PARALLEL_DEGREE"
|
||||
LLM_SERVE_NUM_REPLICAS = "LLM_SERVE_NUM_REPLICAS"
|
||||
LLM_SERVE_MODELS = "LLM_SERVE_MODELS"
|
||||
LLM_SERVE_GPU_TYPE = "LLM_SERVE_GPU_TYPE"
|
||||
LLM_SERVE_NUM_GPUS = "LLM_SERVE_NUM_GPUS"
|
||||
|
||||
|
||||
class TelemetryModel(BaseModelExtended):
|
||||
"""Telemetry model for LLM Serve.
|
||||
|
||||
``model_id_hash`` is the dedup identity used by the telemetry agent and is
|
||||
never recorded as a tag value. It is a hash of the model id so the raw model
|
||||
name never reaches the head-node actor.
|
||||
"""
|
||||
|
||||
model_id_hash: str
|
||||
model_architecture: str
|
||||
num_replicas: int
|
||||
use_lora: bool
|
||||
initial_num_lora_adapters: int
|
||||
use_autoscaling: bool
|
||||
min_replicas: int
|
||||
max_replicas: int
|
||||
tensor_parallel_degree: int
|
||||
gpu_type: str
|
||||
num_gpus: int
|
||||
|
||||
|
||||
@ray.remote(
|
||||
name=LLM_SERVE_TELEMETRY_ACTOR_NAME,
|
||||
namespace=LLM_SERVE_TELEMETRY_NAMESPACE,
|
||||
num_cpus=0,
|
||||
lifetime="detached",
|
||||
)
|
||||
class TelemetryAgent:
|
||||
"""Named Actor to keep the state of all deployed models and record telemetry."""
|
||||
|
||||
def __init__(self):
|
||||
# Keyed by model_id_hash so repeated reports from replicas/restarts of
|
||||
# the same model overwrite rather than accumulate.
|
||||
self.models: Dict[str, TelemetryModel] = {}
|
||||
self.record_tag_func = record_extra_usage_tag
|
||||
|
||||
def _update_record_tag_func(self, record_tag_func: Callable) -> None:
|
||||
"""This method is only used in tests to record the telemetry tags to a different
|
||||
object than Ray's default `record_extra_usage_tag` function."""
|
||||
self.record_tag_func = record_tag_func
|
||||
|
||||
def _reset_models(self):
|
||||
"""This method is only used in tests to clean up the models list."""
|
||||
self.models = {}
|
||||
|
||||
def _multiple_models(self) -> str:
|
||||
unique_models = {model.model_architecture for model in self.models.values()}
|
||||
return "1" if len(unique_models) > 1 else "0"
|
||||
|
||||
@staticmethod
|
||||
def _multiple_apps() -> str:
|
||||
try:
|
||||
try:
|
||||
serve_status = serve.status()
|
||||
except ray.exceptions.ActorDiedError:
|
||||
# In a workspace with multiple Serve sessions, the long-running
|
||||
# telemetry agent may still be connected to a previous, now-dead
|
||||
# session. Shut down so we can reconnect to the live one.
|
||||
serve.shutdown()
|
||||
serve_status = serve.status()
|
||||
return "1" if len(serve_status.applications) > 1 else "0"
|
||||
except Exception:
|
||||
# Telemetry must never fail; fall back to "not multiple".
|
||||
logger.debug("Failed to query serve.status() for telemetry", exc_info=True)
|
||||
return "0"
|
||||
|
||||
def _lora_base_nodes(self) -> str:
|
||||
return ",".join(
|
||||
[
|
||||
model.model_architecture
|
||||
for model in self.models.values()
|
||||
if model.use_lora
|
||||
]
|
||||
)
|
||||
|
||||
def _lora_initial_num_adaptors(self) -> str:
|
||||
return ",".join(
|
||||
[
|
||||
str(model.initial_num_lora_adapters)
|
||||
for model in self.models.values()
|
||||
if model.use_lora
|
||||
]
|
||||
)
|
||||
|
||||
def _autoscaling_enabled_models(self) -> str:
|
||||
return ",".join(
|
||||
[
|
||||
model.model_architecture
|
||||
for model in self.models.values()
|
||||
if model.use_autoscaling
|
||||
]
|
||||
)
|
||||
|
||||
def _autoscaling_min_replicas(self) -> str:
|
||||
return ",".join(
|
||||
[
|
||||
str(model.min_replicas)
|
||||
for model in self.models.values()
|
||||
if model.use_autoscaling
|
||||
]
|
||||
)
|
||||
|
||||
def _autoscaling_max_replicas(self) -> str:
|
||||
return ",".join(
|
||||
[
|
||||
str(model.max_replicas)
|
||||
for model in self.models.values()
|
||||
if model.use_autoscaling
|
||||
]
|
||||
)
|
||||
|
||||
def _model_architectures(self) -> str:
|
||||
return ",".join([model.model_architecture for model in self.models.values()])
|
||||
|
||||
def _tensor_parallel_degree(self) -> str:
|
||||
return ",".join(
|
||||
[str(model.tensor_parallel_degree) for model in self.models.values()]
|
||||
)
|
||||
|
||||
def _num_replicas(self) -> str:
|
||||
return ",".join([str(model.num_replicas) for model in self.models.values()])
|
||||
|
||||
def _gpu_type(self) -> str:
|
||||
return ",".join([model.gpu_type for model in self.models.values()])
|
||||
|
||||
def _num_gpus(self) -> str:
|
||||
return ",".join([str(model.num_gpus) for model in self.models.values()])
|
||||
|
||||
def generate_report(self) -> Dict[str, str]:
|
||||
return {
|
||||
TelemetryTags.LLM_SERVE_SERVE_MULTIPLE_MODELS: self._multiple_models(),
|
||||
TelemetryTags.LLM_SERVE_SERVE_MULTIPLE_APPS: self._multiple_apps(),
|
||||
TelemetryTags.LLM_SERVE_LORA_BASE_MODELS: self._lora_base_nodes(),
|
||||
TelemetryTags.LLM_SERVE_INITIAL_NUM_LORA_ADAPTERS: self._lora_initial_num_adaptors(),
|
||||
TelemetryTags.LLM_SERVE_AUTOSCALING_ENABLED_MODELS: self._autoscaling_enabled_models(),
|
||||
TelemetryTags.LLM_SERVE_AUTOSCALING_MIN_REPLICAS: self._autoscaling_min_replicas(),
|
||||
TelemetryTags.LLM_SERVE_AUTOSCALING_MAX_REPLICAS: self._autoscaling_max_replicas(),
|
||||
TelemetryTags.LLM_SERVE_MODELS: self._model_architectures(),
|
||||
TelemetryTags.LLM_SERVE_TENSOR_PARALLEL_DEGREE: self._tensor_parallel_degree(),
|
||||
TelemetryTags.LLM_SERVE_NUM_REPLICAS: self._num_replicas(),
|
||||
TelemetryTags.LLM_SERVE_GPU_TYPE: self._gpu_type(),
|
||||
TelemetryTags.LLM_SERVE_NUM_GPUS: self._num_gpus(),
|
||||
}
|
||||
|
||||
def record(self, model: Optional[TelemetryModel] = None) -> None:
|
||||
"""Record telemetry model."""
|
||||
from ray._common.usage.usage_lib import TagKey
|
||||
|
||||
if model:
|
||||
self.models[model.model_id_hash] = model
|
||||
|
||||
for key, value in self.generate_report().items():
|
||||
try:
|
||||
self.record_tag_func(TagKey.Value(key), value)
|
||||
except ValueError:
|
||||
# If the key doesn't exist in the TagKey enum, skip it.
|
||||
continue
|
||||
|
||||
|
||||
def _get_or_create_telemetry_agent() -> TelemetryAgent:
|
||||
"""Get or create the detached telemetry agent.
|
||||
|
||||
``get_if_exists`` makes creation atomic, so concurrent replicas converge on a
|
||||
single actor without racing on the name.
|
||||
"""
|
||||
return TelemetryAgent.options(
|
||||
name=LLM_SERVE_TELEMETRY_ACTOR_NAME,
|
||||
namespace=LLM_SERVE_TELEMETRY_NAMESPACE,
|
||||
get_if_exists=True,
|
||||
# Ensure the actor is created on the head node.
|
||||
resources={HEAD_NODE_RESOURCE_NAME: 0.001},
|
||||
# Ensure the actor is not scheduled with the existing placement group.
|
||||
scheduling_strategy=PlacementGroupSchedulingStrategy(placement_group=None),
|
||||
).remote()
|
||||
|
||||
|
||||
def _retry_get_telemetry_agent(
|
||||
max_retries: int = 5, base_delay: float = 0.1
|
||||
) -> TelemetryAgent:
|
||||
"""Get-or-create the telemetry agent, retrying transient failures."""
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
return _get_or_create_telemetry_agent()
|
||||
except Exception as e:
|
||||
logger.info(
|
||||
"Attempt %s/%s to get telemetry agent failed", attempt + 1, max_retries
|
||||
)
|
||||
if attempt == max_retries - 1:
|
||||
raise e
|
||||
# Exponential backoff with jitter; ~3.5s total over 5 attempts.
|
||||
time.sleep(base_delay * (2**attempt) + random.uniform(0, 0.5))
|
||||
|
||||
|
||||
def _push_telemetry_report(model: Optional[TelemetryModel] = None) -> None:
|
||||
"""Push telemetry report for a model."""
|
||||
telemetry_agent = _retry_get_telemetry_agent()
|
||||
assert telemetry_agent is not None
|
||||
ray.get(telemetry_agent.record.remote(model))
|
||||
|
||||
|
||||
class HardwareUsage:
|
||||
"""Hardware usage class to report telemetry."""
|
||||
|
||||
def __init__(self, get_hardware_fn: Callable = get_hardware_usages_to_report):
|
||||
self._get_hardware_fn = get_hardware_fn
|
||||
|
||||
def infer_gpu_from_hardware(self) -> str:
|
||||
"""Infer the GPU type from the hardware when the accelerator type on llm config is
|
||||
not specified.
|
||||
"""
|
||||
from ray.llm._internal.serve.core.configs.accelerators import AcceleratorType
|
||||
|
||||
all_accelerator_types = [t.value for t in AcceleratorType]
|
||||
gcs_client = ray.experimental.internal_kv.internal_kv_get_gcs_client()
|
||||
hardwares = self._get_hardware_fn(gcs_client)
|
||||
for hardware in hardwares:
|
||||
if hardware in all_accelerator_types:
|
||||
return hardware
|
||||
|
||||
return DEFAULT_GPU_TYPE
|
||||
|
||||
|
||||
def push_telemetry_report_for_all_models(
|
||||
all_models: Optional[Sequence["LLMConfig"]] = None,
|
||||
get_lora_model_func: Callable = get_lora_model_ids,
|
||||
get_hardware_fn: Callable = get_hardware_usages_to_report,
|
||||
):
|
||||
"""Push a telemetry report for each model. Never raises."""
|
||||
if not all_models:
|
||||
return
|
||||
|
||||
for model in all_models:
|
||||
# Telemetry must never break the caller (e.g. engine start).
|
||||
try:
|
||||
_push_model_telemetry(model, get_lora_model_func, get_hardware_fn)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"Failed to push telemetry for model %s",
|
||||
getattr(model, "model_id", "<unknown>"),
|
||||
)
|
||||
|
||||
|
||||
def _push_model_telemetry(
|
||||
model: "LLMConfig",
|
||||
get_lora_model_func: Callable,
|
||||
get_hardware_fn: Callable,
|
||||
) -> None:
|
||||
use_lora = (
|
||||
model.lora_config is not None
|
||||
and model.lora_config.dynamic_lora_loading_path is not None
|
||||
)
|
||||
initial_num_lora_adapters = 0
|
||||
if use_lora:
|
||||
lora_model_ids = get_lora_model_func(
|
||||
dynamic_lora_loading_path=model.lora_config.dynamic_lora_loading_path,
|
||||
base_model_id=model.model_id,
|
||||
)
|
||||
initial_num_lora_adapters = len(lora_model_ids)
|
||||
|
||||
deployment_config = model.deployment_config
|
||||
autoscaling_config = deployment_config.get("autoscaling_config")
|
||||
if deployment_config.get("num_replicas") == "auto":
|
||||
# "auto" resolves to an autoscaling config; mirror LLMConfig validation
|
||||
# since the stored deployment_config keeps the literal "auto".
|
||||
from ray.serve._private.config import handle_num_replicas_auto
|
||||
|
||||
_, autoscaling_config = handle_num_replicas_auto(
|
||||
deployment_config.get("max_ongoing_requests"), autoscaling_config
|
||||
)
|
||||
|
||||
use_autoscaling = autoscaling_config is not None
|
||||
if use_autoscaling:
|
||||
from ray.serve.config import AutoscalingConfig
|
||||
|
||||
if isinstance(autoscaling_config, dict):
|
||||
autoscaling_config = AutoscalingConfig(**autoscaling_config)
|
||||
num_replicas = (
|
||||
autoscaling_config.initial_replicas or autoscaling_config.min_replicas
|
||||
)
|
||||
min_replicas = autoscaling_config.min_replicas
|
||||
max_replicas = autoscaling_config.max_replicas
|
||||
else:
|
||||
# Fixed replica count; honor the configured value (including 0),
|
||||
# defaulting to 1 only when unset.
|
||||
num_replicas = deployment_config.get("num_replicas")
|
||||
if num_replicas is None:
|
||||
num_replicas = 1
|
||||
min_replicas = max_replicas = num_replicas
|
||||
|
||||
engine_config = model.get_engine_config()
|
||||
hardware_usage = HardwareUsage(get_hardware_fn)
|
||||
|
||||
telemetry_model = TelemetryModel(
|
||||
# Hash so the cleartext model name (possibly proprietary) never reaches
|
||||
# the head-node actor; deterministic across replicas/restarts so dedup holds.
|
||||
model_id_hash=hashlib.sha256(model.model_id.encode("utf-8")).hexdigest(),
|
||||
model_architecture=model.model_architecture,
|
||||
num_replicas=num_replicas,
|
||||
use_lora=use_lora,
|
||||
initial_num_lora_adapters=initial_num_lora_adapters,
|
||||
use_autoscaling=use_autoscaling,
|
||||
min_replicas=min_replicas,
|
||||
max_replicas=max_replicas,
|
||||
tensor_parallel_degree=engine_config.tensor_parallel_degree,
|
||||
gpu_type=model.accelerator_type or hardware_usage.infer_gpu_from_hardware(),
|
||||
num_gpus=engine_config.num_devices,
|
||||
)
|
||||
_push_telemetry_report(telemetry_model)
|
||||
@@ -0,0 +1,18 @@
|
||||
# experimental_configs key overriding the per-node base port.
|
||||
KV_EVENTS_PORT_BASE_KEY = "KV_EVENTS_PORT_BASE"
|
||||
DEFAULT_KV_EVENTS_PORT_BASE = 5557
|
||||
|
||||
# experimental_configs key overriding the selection service's KV-indexer thread count.
|
||||
KV_INDEXER_THREADS_KEY = "KV_INDEXER_THREADS"
|
||||
DEFAULT_KV_INDEXER_THREADS = 4
|
||||
|
||||
# The engine's KV-event replay (ROUTER) socket sits this many ports above its PUB
|
||||
# port, a separate range so it never collides with the PUB ports of colocated
|
||||
# replicas (PORT_BASE + replica rank). Dynamo's selection service dials it to recover
|
||||
# events missed before its SUB connected.
|
||||
DEFAULT_KV_EVENTS_REPLAY_PORT_OFFSET = 1000
|
||||
|
||||
# TTL for a request's lifecycle tracking on the KV router actor. A live
|
||||
# replica whose completion event was lost (e.g. a batch dropped on a
|
||||
# transient actor outage) would otherwise leave its entry tracked forever.
|
||||
REQUEST_TRACKING_TTL_S = 3600
|
||||
@@ -0,0 +1,483 @@
|
||||
import asyncio
|
||||
import hashlib
|
||||
import logging
|
||||
import math
|
||||
import time
|
||||
from collections import OrderedDict
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, List, Optional, Set, TypedDict
|
||||
|
||||
import ray
|
||||
from ray import serve
|
||||
from ray.llm._internal.serve.routing_policies.kv_aware.constants import (
|
||||
DEFAULT_KV_INDEXER_THREADS,
|
||||
REQUEST_TRACKING_TTL_S,
|
||||
)
|
||||
from ray.serve._private.common import DeploymentTargetInfo
|
||||
from ray.serve._private.constants import (
|
||||
SERVE_CONTROLLER_NAME,
|
||||
SERVE_LOGGER_NAME,
|
||||
SERVE_NAMESPACE,
|
||||
)
|
||||
from ray.serve._private.long_poll import LongPollClient, LongPollNamespace
|
||||
|
||||
logger = logging.getLogger(SERVE_LOGGER_NAME)
|
||||
|
||||
KV_ROUTER_ACTOR_NAME = "serve_llm_kv_router"
|
||||
|
||||
# Dynamo's selection service keys all worker, indexer, and load state by
|
||||
# (model_name, tenant_id). KVRouterActor is a deployment-scoped actor that
|
||||
# instantiates a selection service and serves exactly one model, so a single
|
||||
# fixed key scopes all of its workers together.
|
||||
_MODEL_NAME = "default"
|
||||
_TENANT_ID = "default"
|
||||
|
||||
# Hooks a replica may invoke through ``KVRouterActor.on_lifecycle_events``.
|
||||
LIFECYCLE_HOOKS = frozenset(
|
||||
{
|
||||
"on_request_added",
|
||||
"on_prefill_complete",
|
||||
"on_decode_progress",
|
||||
"on_request_completed",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def get_worker_id(replica_unique_id: str) -> int:
|
||||
"""Deterministically derive a Dynamo worker id from a replica's unique id."""
|
||||
return int.from_bytes(
|
||||
hashlib.blake2b(replica_unique_id.encode(), digest_size=8).digest(), "big"
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RequestLifecycle:
|
||||
"""In-flight request load state while the request is served by a replica."""
|
||||
|
||||
worker_id: int
|
||||
prompt_tokens: int = 0
|
||||
# Client-provided output-length estimate (``sampling_params.max_tokens``);
|
||||
# weights each decode block's load by how much generation remains.
|
||||
expected_output_tokens: Optional[int] = None
|
||||
prefill_completed: bool = False
|
||||
output_tokens: int = 0
|
||||
# Running count of KV blocks (prompt + output) the request occupies; the
|
||||
# cursor for booking each newly crossed decode block.
|
||||
total_blocks: int = 0
|
||||
# Monotonic admission time, for the TTL eviction sweep.
|
||||
created_at: float = field(default_factory=time.monotonic)
|
||||
|
||||
|
||||
class WorkerSelection(TypedDict):
|
||||
"""The worker chosen by ``KVRouterActor.select_worker`` for a request."""
|
||||
|
||||
# The chosen worker.
|
||||
worker_id: int
|
||||
# Data-parallel rank within the worker.
|
||||
dp_rank: int
|
||||
# Matched prompt tokens available on the selected worker.
|
||||
overlap_tokens: int
|
||||
# Prompt tokens that still need prefill on the selected worker.
|
||||
effective_prefill_tokens: int
|
||||
|
||||
|
||||
class KVRouterActor:
|
||||
"""Deployment-scoped Ray actor backing KV-aware routing.
|
||||
|
||||
Attached to the LLMServer deployment via Serve's ``DeploymentActorConfig``,
|
||||
independent of any replica's lifetime.
|
||||
|
||||
1. Created once per deployment, attached to the LLMServer deployment via
|
||||
Serve's ``DeploymentActorConfig`` (independent of any replica's lifetime).
|
||||
2. Owns an in-process Dynamo ``SelectionService``.
|
||||
3. Tracks live replicas via a ``LongPollClient`` on ``DEPLOYMENT_TARGETS``,
|
||||
mapping each running replica to a Dynamo worker id.
|
||||
4. The ``SelectionService`` maintains a global KV index radix tree, fed by
|
||||
every replica's KV events; each node records which workers hold that KV block.
|
||||
5. Scoring (``select_worker``) ranks candidate workers by KV-cache overlap
|
||||
and prefill/decode load.
|
||||
6. Books each request's lifecycle into the service's active-load tracker, so
|
||||
in-flight load feeds back into scoring for subsequent requests.
|
||||
"""
|
||||
|
||||
def __init__(self, indexer_threads: int = DEFAULT_KV_INDEXER_THREADS):
|
||||
# KV-cache block size, learned once from the first replica's reported
|
||||
# engine config and passed to the selection service, which uses it to
|
||||
# track the worker's active load and index its KV blocks for overlap.
|
||||
self._block_size: Optional[int] = None
|
||||
self._indexer_threads = indexer_threads
|
||||
# _replica_id_by_worker maps a Dynamo worker id to the running replica's full
|
||||
# id string, kept in sync with the deployment's live replicas over LongPoll.
|
||||
# NOTE (jeffreywang): _replica_id_by_worker is later used by select_worker
|
||||
# to get candidate workers to route among.
|
||||
self._replica_id_by_worker: Dict[int, str] = {}
|
||||
# Per-request state that the lifecycle hooks need, keyed by request id, serves
|
||||
# the following purposes:
|
||||
# 1. Block cursor: Turn cumulative decode tokens into add_output_block deltas.
|
||||
# 2. expected_output_tokens for decode-block decay weighting.
|
||||
# 3. In-flight request set: Free reservation exactly once.
|
||||
# Ordered oldest-first so the TTL sweep pops stale entries off the front.
|
||||
self._requests: "OrderedDict[str, RequestLifecycle]" = OrderedDict()
|
||||
# Reverse index of in-flight request ids per worker, kept in lockstep with
|
||||
# _requests, so remove_worker is O(k) in the worker's requests, not O(N).
|
||||
self._request_ids_by_worker: Dict[int, Set[str]] = {}
|
||||
# Carries the effective prefill tokens select() computed at routing time to
|
||||
# on_request_added, which books them via the explicit create_reservation.
|
||||
# TODO(jeffreywang): this map is only needed because create_reservation
|
||||
# requires the effective prefill tokens to be passed in explicitly. Once the
|
||||
# selection service caches each select() result and create_reservation can
|
||||
# look it up by request id, Ray no longer needs to forward it.
|
||||
self._effective_prefill_tokens_by_request: Dict[str, int] = {}
|
||||
self._pending_tasks: Set[asyncio.Task] = set()
|
||||
self._long_poll_client: Optional[LongPollClient] = None
|
||||
self._create_selection_service()
|
||||
self._start_replica_tracking()
|
||||
|
||||
async def ready(self) -> None:
|
||||
"""Readiness probe for KVAwareRouter to confirm KVRouterActor is initialized
|
||||
before it starts routing requests to it.
|
||||
"""
|
||||
|
||||
def get_block_size(self) -> int:
|
||||
"""Return the KV-cache block size used for decode-block accounting."""
|
||||
return self._block_size
|
||||
|
||||
def _create_selection_service(self) -> None:
|
||||
"""Create the in-process Dynamo selection service for this deployment."""
|
||||
# Imported here, not at module scope: Ray pickles this actor class by
|
||||
# value, and Dynamo's pyo3 classes cannot be pickled as its globals.
|
||||
try:
|
||||
from dynamo.llm import SelectionService
|
||||
except ImportError:
|
||||
self._svc = None
|
||||
logger.warning(
|
||||
"ai-dynamo is not installed; KV-aware routing requires ai-dynamo."
|
||||
)
|
||||
return
|
||||
|
||||
self._svc = SelectionService(indexer_threads=self._indexer_threads)
|
||||
logger.info(
|
||||
"Dynamo SelectionService created (indexer threads %d).",
|
||||
self._indexer_threads,
|
||||
)
|
||||
|
||||
def _start_replica_tracking(self) -> None:
|
||||
"""Subscribe to this deployment's running replicas via LongPollClient."""
|
||||
deployment_id = serve.get_deployment_actor_context().deployment_id
|
||||
controller = ray.get_actor(SERVE_CONTROLLER_NAME, namespace=SERVE_NAMESPACE)
|
||||
self._long_poll_client = LongPollClient(
|
||||
controller,
|
||||
{
|
||||
(
|
||||
LongPollNamespace.DEPLOYMENT_TARGETS,
|
||||
deployment_id,
|
||||
): self._on_deployment_targets,
|
||||
},
|
||||
# Relies on KVRouterActor being an async actor (it defines async
|
||||
# methods), so Ray runs __init__ inside the actor's event loop.
|
||||
call_in_event_loop=asyncio.get_running_loop(),
|
||||
client_id=f"{type(self).__name__}:{deployment_id}",
|
||||
)
|
||||
|
||||
def _schedule(self, coro) -> None:
|
||||
"""Run a coroutine on the actor's event loop, holding a reference until
|
||||
it completes.
|
||||
"""
|
||||
task = asyncio.ensure_future(coro)
|
||||
self._pending_tasks.add(task)
|
||||
task.add_done_callback(self._pending_tasks.discard)
|
||||
|
||||
def _register_block_size(self, block_size: int, replica_id: str) -> None:
|
||||
"""Pin the deployment's KV-cache block size from the first replica's
|
||||
reported engine config.
|
||||
"""
|
||||
if self._block_size is None:
|
||||
self._block_size = block_size
|
||||
logger.info("KV router block size set to %d.", block_size)
|
||||
elif block_size != self._block_size:
|
||||
# Replicas of a deployment are expected to resolve the same block
|
||||
# size, so a mismatch is unexpected. We still register the worker so
|
||||
# the selection service spawns its KV-event listener, but the indexer
|
||||
# only ingests blocks whose size matches the pinned block size, so a
|
||||
# genuinely mismatched replica's KV events would be dropped (its KV
|
||||
# cache never indexed).
|
||||
logger.error(
|
||||
"Replica %s reports KV block size %d but the KV router is "
|
||||
"pinned at %d; registering it at the pinned size (replicas of a "
|
||||
"deployment are expected to agree).",
|
||||
replica_id,
|
||||
block_size,
|
||||
self._block_size,
|
||||
)
|
||||
|
||||
def _on_deployment_targets(self, target_info: DeploymentTargetInfo) -> None:
|
||||
"""LongPoll listener: reconcile tracked workers against the running-replica
|
||||
snapshot.
|
||||
|
||||
Each replica advertises its KV-events endpoint via ``record_routing_stats``
|
||||
(carried in ``RunningReplicaInfo.routing_stats``); newly advertised replicas
|
||||
are registered with the selection service and departed ones evicted.
|
||||
"""
|
||||
members: Dict[int, tuple] = {}
|
||||
for replica in target_info.running_replicas:
|
||||
worker_id = get_worker_id(replica.replica_id.unique_id)
|
||||
kv_event_metadata = replica.routing_stats.get("kv_event_metadata")
|
||||
if kv_event_metadata is not None:
|
||||
members[worker_id] = (
|
||||
replica.replica_id.to_full_id_str(),
|
||||
kv_event_metadata,
|
||||
)
|
||||
|
||||
registered = set(self._replica_id_by_worker)
|
||||
added = members.keys() - registered
|
||||
removed = registered - members.keys()
|
||||
|
||||
for worker_id in removed:
|
||||
self.remove_worker(worker_id)
|
||||
self._replica_id_by_worker.pop(worker_id, None)
|
||||
for worker_id in added:
|
||||
replica_id, kv_event_metadata = members[worker_id]
|
||||
self._register_block_size(kv_event_metadata["block_size"], replica_id)
|
||||
self._replica_id_by_worker[worker_id] = replica_id
|
||||
self._schedule(
|
||||
self._upsert_worker(worker_id, replica_id, kv_event_metadata)
|
||||
)
|
||||
|
||||
if added or removed:
|
||||
logger.info(
|
||||
"KV router replica membership updated: +%d -%d, tracking %d worker(s).",
|
||||
len(added),
|
||||
len(removed),
|
||||
len(self._replica_id_by_worker),
|
||||
)
|
||||
|
||||
def remove_worker(self, worker_id: int) -> None:
|
||||
"""Evict a departed replica's worker and its KV blocks from the
|
||||
selection service.
|
||||
"""
|
||||
# Drop the departed replica's in-flight requests; their completions can
|
||||
# never arrive, so they would otherwise leak. delete_worker below frees
|
||||
# their load in the service, so no per-request free_reservation is needed.
|
||||
for request_id in self._request_ids_by_worker.pop(worker_id, set()):
|
||||
self._requests.pop(request_id, None)
|
||||
if self._svc is None:
|
||||
return
|
||||
self._schedule(self._svc.delete_worker(worker_id))
|
||||
|
||||
async def _upsert_worker(
|
||||
self, worker_id: int, replica_id: str, kv_event_metadata: Dict[str, Any]
|
||||
) -> None:
|
||||
"""Register a replica's KV-event endpoint with the selection service.
|
||||
|
||||
The selection service spawns a connect-out ZMQ listener to the
|
||||
replica's ``endpoint`` and indexes its live KV events.
|
||||
"""
|
||||
if self._svc is None:
|
||||
return
|
||||
dp_rank = kv_event_metadata["dp_rank"]
|
||||
await self._svc.upsert_worker(
|
||||
{
|
||||
"worker_id": worker_id,
|
||||
"model_name": _MODEL_NAME,
|
||||
"tenant_id": _TENANT_ID,
|
||||
# NOTE: SelectionService requires endpoint to be non-empty although it's left
|
||||
# unused under an external runtime like Ray Serve LLM.
|
||||
# TODO (jeffreywang): Allow empty endpoints upstream.
|
||||
"endpoint": f"ray://{replica_id}",
|
||||
"block_size": self._block_size,
|
||||
# NOTE: max_num_batched_tokens is a proxy of load capacity for load-based
|
||||
# scoring in the selection service.
|
||||
"max_num_batched_tokens": kv_event_metadata["max_num_batched_tokens"],
|
||||
"data_parallel_start_rank": dp_rank,
|
||||
# TODO (jeffreywang): Support KV-aware routing for data parallel deployments.
|
||||
"data_parallel_size": 1,
|
||||
"kv_events_endpoints": {dp_rank: kv_event_metadata["endpoint"]},
|
||||
# The listener dials this on a sequence gap (slow-joiner) to replay
|
||||
# the events it missed before its SUB connected; without it those
|
||||
# events are dropped and never indexed.
|
||||
"replay_endpoint": kv_event_metadata.get("replay_endpoint"),
|
||||
}
|
||||
)
|
||||
logger.info(
|
||||
"Registered KV event worker %d for replica %s at %s.",
|
||||
worker_id,
|
||||
replica_id,
|
||||
kv_event_metadata["endpoint"],
|
||||
)
|
||||
|
||||
async def select_worker(
|
||||
self,
|
||||
request_id: str,
|
||||
token_ids: List[int],
|
||||
allowed_worker_ids: List[int],
|
||||
) -> WorkerSelection:
|
||||
"""Score the allowed workers for a request based on KV-cache overlap and
|
||||
load and pick the best one.
|
||||
|
||||
Args:
|
||||
request_id: Unique identifier for the request being routed.
|
||||
token_ids: Prompt token ids used to compute KV-cache overlap.
|
||||
allowed_worker_ids: Candidate worker ids the router may select from.
|
||||
|
||||
Returns:
|
||||
The selected worker (see ``WorkerSelection``).
|
||||
"""
|
||||
if token_ids is None or len(token_ids) == 0:
|
||||
raise ValueError("KV aware routing requires non-empty token_ids.")
|
||||
|
||||
if self._svc is None:
|
||||
# ai-dynamo is not installed, so this deployment cannot score requests.
|
||||
# Fail fast and surface RuntimeError to the client as a 503 via LLMRouter.
|
||||
raise RuntimeError(
|
||||
"KV-aware routing is unavailable because ai-dynamo is not "
|
||||
"installed in the deployment's environment."
|
||||
)
|
||||
selection = await self._svc.select(
|
||||
{
|
||||
"model_name": _MODEL_NAME,
|
||||
"tenant_id": _TENANT_ID,
|
||||
"selection_id": request_id,
|
||||
"token_ids": token_ids,
|
||||
"allowed_worker_ids": allowed_worker_ids,
|
||||
}
|
||||
)
|
||||
self._effective_prefill_tokens_by_request[request_id] = selection[
|
||||
"effective_prefill_tokens"
|
||||
]
|
||||
return {
|
||||
"worker_id": selection["worker_id"],
|
||||
"dp_rank": selection["dp_rank"],
|
||||
"overlap_tokens": selection["overlap"]["longest_matched"],
|
||||
"effective_prefill_tokens": selection["effective_prefill_tokens"],
|
||||
}
|
||||
|
||||
async def on_lifecycle_events(self, events: List[tuple]) -> None:
|
||||
"""Apply a replica's ``(hook_name, args)`` lifecycle events in order.
|
||||
|
||||
The hooks are order-sensitive (e.g. a completion arriving before its
|
||||
admission would resurrect an evicted request) so a replica sends its
|
||||
events in submission order, batched into one call.
|
||||
"""
|
||||
if self._svc is None or self._block_size is None:
|
||||
return
|
||||
for hook_name, args in events:
|
||||
if hook_name not in LIFECYCLE_HOOKS:
|
||||
logger.warning("Ignoring unknown lifecycle hook %s", hook_name)
|
||||
continue
|
||||
try:
|
||||
await getattr(self, hook_name)(*args)
|
||||
except Exception:
|
||||
# One hook raising must not abort the batch and drop other events.
|
||||
logger.exception(
|
||||
"KV lifecycle hook %s failed; skipping it and continuing.",
|
||||
hook_name,
|
||||
)
|
||||
|
||||
async def on_request_added(
|
||||
self,
|
||||
request_id: str,
|
||||
worker_id: int,
|
||||
token_ids: List[int],
|
||||
expected_output_tokens: Optional[int] = None,
|
||||
) -> None:
|
||||
"""Admit a routed request into ``worker_id``'s active load, booking it
|
||||
into the selection service which computes the worker's KV overlap from
|
||||
``token_ids``, so the recorded prefill excludes the cached prefix."""
|
||||
await self._evict_stale_requests()
|
||||
prompt_tokens = len(token_ids)
|
||||
self._requests[request_id] = RequestLifecycle(
|
||||
worker_id=worker_id,
|
||||
prompt_tokens=prompt_tokens,
|
||||
expected_output_tokens=expected_output_tokens,
|
||||
total_blocks=math.ceil(prompt_tokens / self._block_size),
|
||||
)
|
||||
self._request_ids_by_worker.setdefault(worker_id, set()).add(request_id)
|
||||
effective_prefill_tokens = self._effective_prefill_tokens_by_request.pop(
|
||||
request_id, None
|
||||
)
|
||||
|
||||
await self._svc.create_reservation(
|
||||
{
|
||||
"model_name": _MODEL_NAME,
|
||||
"tenant_id": _TENANT_ID,
|
||||
"reservation_id": request_id,
|
||||
"worker_id": worker_id,
|
||||
"token_ids": token_ids,
|
||||
"expected_output_tokens": expected_output_tokens,
|
||||
"effective_prefill_tokens": effective_prefill_tokens,
|
||||
}
|
||||
)
|
||||
if request_id not in self._requests:
|
||||
await self._svc.free_reservation(request_id)
|
||||
|
||||
async def on_prefill_complete(self, request_id: str) -> None:
|
||||
"""Record a request's prefill -> decode transition, dropping its prefill
|
||||
load in the selection service."""
|
||||
state = self._requests.get(request_id)
|
||||
if state is None:
|
||||
return
|
||||
state.prefill_completed = True
|
||||
await self._svc.prefill_complete(request_id)
|
||||
|
||||
async def on_decode_progress(
|
||||
self, request_id: str, cumulative_output_tokens: int
|
||||
) -> None:
|
||||
"""Advance ``request_id`` to an exact cumulative output-token count,
|
||||
booking one decode block in the selection service per crossed boundary.
|
||||
"""
|
||||
state = self._requests.get(request_id)
|
||||
if state is None:
|
||||
return
|
||||
state.output_tokens = cumulative_output_tokens
|
||||
new_total_blocks = math.ceil(
|
||||
(state.prompt_tokens + cumulative_output_tokens) / self._block_size
|
||||
)
|
||||
decay_fraction = self._get_decay_fraction(state)
|
||||
while new_total_blocks > state.total_blocks:
|
||||
state.total_blocks += 1
|
||||
self._svc.add_output_block(request_id, decay_fraction=decay_fraction)
|
||||
|
||||
async def on_request_completed(self, request_id: str) -> None:
|
||||
"""Free ``request_id`` from the selection service's active load and the
|
||||
local view."""
|
||||
self._effective_prefill_tokens_by_request.pop(request_id, None)
|
||||
state = self._requests.pop(request_id, None)
|
||||
if state is not None:
|
||||
self._untrack_worker_request(request_id, state.worker_id)
|
||||
await self._svc.free_reservation(request_id)
|
||||
|
||||
def _untrack_worker_request(self, request_id: str, worker_id: int) -> None:
|
||||
"""Drop a request from the per-worker reverse index, keeping it in
|
||||
lockstep with ``_requests``."""
|
||||
request_ids = self._request_ids_by_worker.get(worker_id)
|
||||
if request_ids is not None:
|
||||
request_ids.discard(request_id)
|
||||
if not request_ids:
|
||||
del self._request_ids_by_worker[worker_id]
|
||||
|
||||
async def _evict_stale_requests(self) -> None:
|
||||
"""Backstop for a lost completion on a live replica: evict requests tracked
|
||||
past ``REQUEST_TRACKING_TTL_S``, freeing their reservations.
|
||||
"""
|
||||
cutoff = time.monotonic() - REQUEST_TRACKING_TTL_S
|
||||
while self._requests:
|
||||
request_id, state = next(iter(self._requests.items()))
|
||||
if state.created_at > cutoff:
|
||||
break
|
||||
self._requests.popitem(last=False)
|
||||
self._untrack_worker_request(request_id, state.worker_id)
|
||||
self._effective_prefill_tokens_by_request.pop(request_id, None)
|
||||
logger.warning(
|
||||
"Evicting stale KV request %s (tracked > %ds without completion); "
|
||||
"freeing its reservation.",
|
||||
request_id,
|
||||
REQUEST_TRACKING_TTL_S,
|
||||
)
|
||||
await self._svc.free_reservation(request_id)
|
||||
|
||||
def _get_decay_fraction(self, state: RequestLifecycle) -> Optional[float]:
|
||||
"""Fraction of output still expected, or ``None`` without an estimate;
|
||||
weights each decode block by how much generation remains."""
|
||||
if not state.expected_output_tokens:
|
||||
return None
|
||||
return max(0.0, 1.0 - state.output_tokens / state.expected_output_tokens)
|
||||
@@ -0,0 +1,122 @@
|
||||
import logging
|
||||
import random
|
||||
from typing import List, Optional
|
||||
|
||||
import ray
|
||||
from ray.actor import ActorHandle
|
||||
from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
|
||||
from ray.llm._internal.serve.core.ingress.tokenizer import REQUEST_TOKEN_IDS_KWARG
|
||||
from ray.llm._internal.serve.routing_policies.kv_aware.kv_aware_actor import (
|
||||
KV_ROUTER_ACTOR_NAME,
|
||||
get_worker_id,
|
||||
)
|
||||
from ray.serve._private.constants import (
|
||||
SERVE_DEPLOYMENT_ACTOR_PREFIX,
|
||||
SERVE_LOGGER_NAME,
|
||||
SERVE_NAMESPACE,
|
||||
)
|
||||
from ray.serve._private.request_router.common import PendingRequest
|
||||
from ray.serve._private.request_router.replica_wrapper import RunningReplica
|
||||
from ray.serve._private.request_router.request_router import RequestRouter
|
||||
from ray.serve.config import RequestRouterConfig
|
||||
|
||||
logger = logging.getLogger(SERVE_LOGGER_NAME)
|
||||
|
||||
|
||||
class KVAwareRouter(RequestRouter):
|
||||
"""Routes each request to the candidate that best balances expected KV-cache
|
||||
overlap against the worker's current prefill/decode load.
|
||||
|
||||
Scoring is delegated to the deployment-scoped ``KVRouterActor`` (which owns the
|
||||
Dynamo selection service and the global KV index); this per-handle router stays
|
||||
thin and simply maps candidate replicas to/from Dynamo worker ids.
|
||||
"""
|
||||
|
||||
def initialize_state(self):
|
||||
"""Resolve the deployment's ``KVRouterActor``.
|
||||
|
||||
The actor is attached to this deployment via ``DeploymentActorConfig``
|
||||
whenever the request router is a ``KVAwareRouter``, so it exists by the time
|
||||
requests route. We resolve its Serve-generated name and block on a cheap
|
||||
call to confirm it finished initializing, so the first routed request finds
|
||||
a ready scorer.
|
||||
"""
|
||||
self._kv_router_actor = self._discover_kv_router_actor()
|
||||
# Synchronization barrier: Ray defers actor methods until __init__ completes,
|
||||
# so awaiting any method blocks until KVRouterActor is constructed.
|
||||
ray.get(self._kv_router_actor.ready.remote())
|
||||
|
||||
def _discover_kv_router_actor(self) -> ActorHandle:
|
||||
"""Handle to this deployment's ``KVRouterActor`` by its Serve-scoped name."""
|
||||
prefix = (
|
||||
f"{SERVE_DEPLOYMENT_ACTOR_PREFIX}"
|
||||
f"{self._deployment_id.app_name}::{self._deployment_id.name}::"
|
||||
)
|
||||
suffix = f"::{KV_ROUTER_ACTOR_NAME}"
|
||||
for entry in ray.util.list_named_actors(all_namespaces=True):
|
||||
name = entry.get("name") or ""
|
||||
if (
|
||||
entry.get("namespace") == SERVE_NAMESPACE
|
||||
and name.startswith(prefix)
|
||||
and name.endswith(suffix)
|
||||
):
|
||||
return ray.get_actor(name, namespace=SERVE_NAMESPACE)
|
||||
raise RuntimeError(
|
||||
f"KVRouterActor for deployment {self._deployment_id} not found; it must "
|
||||
"be attached via DeploymentActorConfig when using KVAwareRouter."
|
||||
)
|
||||
|
||||
async def choose_replicas(
|
||||
self,
|
||||
candidate_replicas: List[RunningReplica],
|
||||
pending_request: Optional[PendingRequest] = None,
|
||||
) -> List[List[RunningReplica]]:
|
||||
"""Choose the candidate replica(s) to route ``pending_request`` to.
|
||||
|
||||
Maps the candidate replicas to their Dynamo worker ids, asks the
|
||||
``KVRouterActor`` to rank them via ``select_worker``, and routes to
|
||||
the chosen worker's replica. With direct streaming enabled, HAProxy
|
||||
then forwards the original request to that replica.
|
||||
|
||||
Requests with no prompt token ids have nothing to score on, so they route
|
||||
to a random candidate. This covers the pre-routing ``/tokenize`` RPC (routed
|
||||
before token ids exist) and token-less fallbacks (batch prompts,
|
||||
truncated/unparseable bodies).
|
||||
TODO (jeffreywang): Move pre-routing tokenization to KVRouterActor while
|
||||
ensuring tokenization correctness.
|
||||
|
||||
Args:
|
||||
candidate_replicas: The replicas eligible to serve the request.
|
||||
pending_request: The request being routed.
|
||||
|
||||
Returns:
|
||||
Ranked groups of replicas.
|
||||
"""
|
||||
token_ids = (
|
||||
pending_request.kwargs.get(REQUEST_TOKEN_IDS_KWARG)
|
||||
if pending_request is not None
|
||||
else None
|
||||
)
|
||||
if not token_ids:
|
||||
return [[random.choice(candidate_replicas)]] if candidate_replicas else []
|
||||
|
||||
worker_id_to_replica = {
|
||||
get_worker_id(replica.replica_id.unique_id): replica
|
||||
for replica in candidate_replicas
|
||||
}
|
||||
selection = await self._kv_router_actor.select_worker.remote(
|
||||
pending_request.metadata.request_id,
|
||||
token_ids,
|
||||
list(worker_id_to_replica),
|
||||
)
|
||||
return [[worker_id_to_replica[selection["worker_id"]]]]
|
||||
|
||||
|
||||
def is_kv_aware(llm_config: LLMConfig) -> bool:
|
||||
"""Whether ``llm_config`` selects a ``KVAwareRouter`` for replica selection."""
|
||||
request_router_config = llm_config.deployment_config.get("request_router_config")
|
||||
if isinstance(request_router_config, dict):
|
||||
request_router_config = RequestRouterConfig(**request_router_config)
|
||||
return isinstance(request_router_config, RequestRouterConfig) and issubclass(
|
||||
request_router_config.get_request_router_class(), KVAwareRouter
|
||||
)
|
||||
@@ -0,0 +1,66 @@
|
||||
"""Helpers for wiring KV-aware routing into an LLM deployment."""
|
||||
|
||||
import logging
|
||||
|
||||
import ray
|
||||
from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
|
||||
from ray.llm._internal.serve.routing_policies.kv_aware.constants import (
|
||||
DEFAULT_KV_INDEXER_THREADS,
|
||||
KV_INDEXER_THREADS_KEY,
|
||||
)
|
||||
from ray.llm._internal.serve.routing_policies.kv_aware.kv_aware_actor import (
|
||||
KV_ROUTER_ACTOR_NAME,
|
||||
KVRouterActor,
|
||||
)
|
||||
from ray.llm._internal.serve.routing_policies.kv_aware.kv_aware_router import (
|
||||
is_kv_aware,
|
||||
)
|
||||
from ray.llm._internal.serve.routing_policies.kv_aware.vllm.kv_events import (
|
||||
configure_kv_events_for_kv_routing,
|
||||
)
|
||||
from ray.serve._private.constants import SERVE_LOGGER_NAME
|
||||
from ray.serve.config import DeploymentActorConfig
|
||||
|
||||
logger = logging.getLogger(SERVE_LOGGER_NAME)
|
||||
|
||||
|
||||
def _maybe_setup_kv_aware_routing(
|
||||
deployment_options: dict, llm_config: LLMConfig
|
||||
) -> None:
|
||||
"""Set up KV-aware routing when the deployment's request router is a
|
||||
KVAwareRouter.
|
||||
|
||||
Attaches the KVRouterActor, which owns the deployment's global KV radix
|
||||
tree, and enables the engine KV events that feed it.
|
||||
"""
|
||||
if not is_kv_aware(llm_config):
|
||||
if llm_config.engine_kwargs.get("kv_events_config") is not None:
|
||||
logger.warning(
|
||||
"engine_kwargs['kv_events_config'] is set but the deployment's "
|
||||
"request router is not a KVAwareRouter, so the engine's KV events "
|
||||
"will not be consumed. To use them, configure KVAwareRouter via "
|
||||
"deployment_config.request_router_config."
|
||||
)
|
||||
return
|
||||
|
||||
# Keep the engine's token-tracking gate which reads llm_config consistent
|
||||
# with the router resolved here from the merged deployment options.
|
||||
llm_config.deployment_config["request_router_config"] = deployment_options[
|
||||
"request_router_config"
|
||||
]
|
||||
|
||||
deployment_options["deployment_actors"] = [
|
||||
*deployment_options.get("deployment_actors", []),
|
||||
DeploymentActorConfig(
|
||||
name=KV_ROUTER_ACTOR_NAME,
|
||||
actor_class=ray.remote(KVRouterActor),
|
||||
actor_options={"num_cpus": 0},
|
||||
init_kwargs={
|
||||
"indexer_threads": llm_config.experimental_configs.get(
|
||||
KV_INDEXER_THREADS_KEY, DEFAULT_KV_INDEXER_THREADS
|
||||
),
|
||||
},
|
||||
),
|
||||
]
|
||||
|
||||
configure_kv_events_for_kv_routing(llm_config)
|
||||
@@ -0,0 +1 @@
|
||||
"""vLLM-specific KV-aware routing helpers."""
|
||||
@@ -0,0 +1,163 @@
|
||||
import logging
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import ray
|
||||
from ray import serve
|
||||
from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
|
||||
from ray.llm._internal.serve.routing_policies.kv_aware.constants import (
|
||||
DEFAULT_KV_EVENTS_PORT_BASE,
|
||||
DEFAULT_KV_EVENTS_REPLAY_PORT_OFFSET,
|
||||
KV_EVENTS_PORT_BASE_KEY,
|
||||
)
|
||||
from ray.llm._internal.serve.routing_policies.kv_aware.kv_aware_router import (
|
||||
is_kv_aware,
|
||||
)
|
||||
from ray.serve._private.constants import SERVE_LOGGER_NAME
|
||||
|
||||
logger = logging.getLogger(SERVE_LOGGER_NAME)
|
||||
|
||||
|
||||
def configure_kv_events_for_kv_routing(llm_config: LLMConfig) -> None:
|
||||
"""
|
||||
Enable engine KV-cache events for a KV-aware-routed deployment.
|
||||
"""
|
||||
engine_kwargs = llm_config.engine_kwargs
|
||||
if engine_kwargs.get("enable_prefix_caching") is False:
|
||||
logger.warning(
|
||||
"KV-aware routing is configured but enable_prefix_caching is False; "
|
||||
"the engine will not emit KV-cache events."
|
||||
)
|
||||
|
||||
llm_config.update_engine_kwargs(
|
||||
kv_events_config={
|
||||
"enable_kv_cache_events": True,
|
||||
"publisher": "zmq",
|
||||
"endpoint": _default_kv_events_endpoint(llm_config),
|
||||
# Bind a replay (ROUTER) socket so the selection service can recover
|
||||
# the events it missed before its SUB connected (slow-joiner gap).
|
||||
"replay_endpoint": _default_kv_events_replay_endpoint(llm_config),
|
||||
}
|
||||
)
|
||||
|
||||
_pin_block_hash_seed(llm_config)
|
||||
|
||||
|
||||
def _pin_block_hash_seed(llm_config: LLMConfig) -> None:
|
||||
"""Make engine block hashes content-deterministic across replicas.
|
||||
|
||||
The KV router's global indexer chains and dedups blocks by the engines'
|
||||
block hashes, so identical content must hash identically on every
|
||||
replica. vLLM salts its block-hash chain root per process unless
|
||||
``PYTHONHASHSEED`` is set, so pin it deployment-wide.
|
||||
"""
|
||||
runtime_env = dict(llm_config.runtime_env or {})
|
||||
env_vars = dict(runtime_env.get("env_vars") or {})
|
||||
env_vars.setdefault("PYTHONHASHSEED", "0")
|
||||
runtime_env["env_vars"] = env_vars
|
||||
llm_config.runtime_env = runtime_env
|
||||
|
||||
|
||||
def assign_replica_kv_events_endpoint(llm_config: LLMConfig) -> None:
|
||||
"""Pin the engine's KV-events endpoint to a per-replica port.
|
||||
|
||||
Replicas of a deployment share one ``engine_kwargs``, so colocated
|
||||
replicas would otherwise bind the same ZMQ port. Offsets the configured
|
||||
base port by the replica's rank.
|
||||
"""
|
||||
if not is_kv_aware(llm_config):
|
||||
return
|
||||
kv_events_config = llm_config.engine_kwargs.get("kv_events_config")
|
||||
if kv_events_config is None:
|
||||
return
|
||||
updated = dict(kv_events_config)
|
||||
endpoint = updated.pop("endpoint")
|
||||
replay_endpoint = updated.pop("replay_endpoint")
|
||||
# With data parallelism the engine offsets ports by dp_rank under ZmqEventPublisher;
|
||||
# otherwise offset by the replica's node-local rank so colocated replicas don't collide.
|
||||
if llm_config.engine_kwargs.get("data_parallel_rank") is not None:
|
||||
offset = 0
|
||||
else:
|
||||
offset = _get_replica_rank()
|
||||
updated["endpoint"] = _get_offset_endpoint_port(endpoint, offset)
|
||||
updated["replay_endpoint"] = _get_offset_endpoint_port(replay_endpoint, offset)
|
||||
llm_config.update_engine_kwargs(kv_events_config=updated)
|
||||
|
||||
|
||||
def resolve_kv_event_source_endpoint(llm_config: LLMConfig) -> Optional[str]:
|
||||
"""This replica's node-routable KV-events endpoint, for the selection
|
||||
service to connect out to.
|
||||
|
||||
The engine's KV-events endpoint at the replica's node IP; ``None`` when
|
||||
KV-cache events are not enabled.
|
||||
"""
|
||||
if not is_kv_aware(llm_config):
|
||||
return None
|
||||
kv_events_config = llm_config.engine_kwargs.get("kv_events_config")
|
||||
if kv_events_config is None:
|
||||
return None
|
||||
return _get_node_routable_endpoint(llm_config, kv_events_config["endpoint"])
|
||||
|
||||
|
||||
def get_kv_event_routing_stats(
|
||||
llm_config: LLMConfig, block_size: int, max_num_batched_tokens: int
|
||||
) -> Dict[str, Any]:
|
||||
"""Returns this replica's routing-stats payload advertising its KV-events
|
||||
endpoint and the engine's resolved KV-cache block size."""
|
||||
if not is_kv_aware(llm_config):
|
||||
return {}
|
||||
kv_events_config = llm_config.engine_kwargs.get("kv_events_config")
|
||||
if kv_events_config is None:
|
||||
return {}
|
||||
kv_event_metadata = {
|
||||
"endpoint": _get_node_routable_endpoint(
|
||||
llm_config, kv_events_config["endpoint"]
|
||||
),
|
||||
"block_size": block_size,
|
||||
"max_num_batched_tokens": max_num_batched_tokens,
|
||||
"dp_rank": llm_config.engine_kwargs.get("data_parallel_rank") or 0,
|
||||
"replay_endpoint": _get_node_routable_endpoint(
|
||||
llm_config, kv_events_config["replay_endpoint"]
|
||||
),
|
||||
}
|
||||
return {"kv_event_metadata": kv_event_metadata}
|
||||
|
||||
|
||||
def _get_node_routable_endpoint(llm_config: LLMConfig, endpoint: str) -> str:
|
||||
"""Rewrite a wildcard-bound engine endpoint to this replica's node IP.
|
||||
|
||||
Offsets by ``data_parallel_rank`` to match the engine's own per-rank offset,
|
||||
so the selection service dials the right socket.
|
||||
"""
|
||||
dp_rank = llm_config.engine_kwargs.get("data_parallel_rank")
|
||||
if dp_rank is not None:
|
||||
endpoint = _get_offset_endpoint_port(endpoint, dp_rank)
|
||||
port = endpoint.rsplit(":", 1)[1]
|
||||
return f"tcp://{ray.util.get_node_ip_address()}:{port}"
|
||||
|
||||
|
||||
def _get_replica_rank() -> int:
|
||||
return serve.get_replica_context().rank.local_rank
|
||||
|
||||
|
||||
def _default_kv_events_endpoint(llm_config: LLMConfig) -> str:
|
||||
port_base = int(
|
||||
llm_config.experimental_configs.get(
|
||||
KV_EVENTS_PORT_BASE_KEY, DEFAULT_KV_EVENTS_PORT_BASE
|
||||
)
|
||||
)
|
||||
return f"tcp://*:{port_base}"
|
||||
|
||||
|
||||
def _default_kv_events_replay_endpoint(llm_config: LLMConfig) -> str:
|
||||
port_base = int(
|
||||
llm_config.experimental_configs.get(
|
||||
KV_EVENTS_PORT_BASE_KEY, DEFAULT_KV_EVENTS_PORT_BASE
|
||||
)
|
||||
)
|
||||
return f"tcp://*:{port_base + DEFAULT_KV_EVENTS_REPLAY_PORT_OFFSET}"
|
||||
|
||||
|
||||
def _get_offset_endpoint_port(endpoint: str, offset: int) -> str:
|
||||
"""Offset a TCP endpoint's port with vLLM's ZmqEventPublisher convention."""
|
||||
base, port = endpoint.rsplit(":", 1)
|
||||
return f"{base}:{int(port) + offset}"
|
||||
@@ -0,0 +1,192 @@
|
||||
import asyncio
|
||||
from typing import Any, List, Optional, Type
|
||||
|
||||
from vllm.outputs import RequestOutput
|
||||
from vllm.sampling_params import RequestOutputKind
|
||||
from vllm.v1.engine.async_llm import AsyncLLM
|
||||
|
||||
from ray import serve
|
||||
from ray.actor import ActorHandle
|
||||
from ray.exceptions import RayActorError, RayTaskError
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
from ray.llm._internal.serve.routing_policies.kv_aware.kv_aware_actor import (
|
||||
KV_ROUTER_ACTOR_NAME,
|
||||
get_worker_id,
|
||||
)
|
||||
from ray.llm._internal.serve.utils.server_utils import get_serve_request_id
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def _get_prompt_token_ids(prompt: Any) -> List[int]:
|
||||
"""The prompt's pre-tokenized token ids."""
|
||||
try:
|
||||
return list(prompt["prompt_token_ids"])
|
||||
except (KeyError, TypeError) as e:
|
||||
raise ValueError(
|
||||
"KV-aware token tracking requires a pre-tokenized prompt "
|
||||
f"(dict with 'prompt_token_ids'); got {type(prompt).__name__}"
|
||||
) from e
|
||||
|
||||
|
||||
class LifecycleEventForwarder:
|
||||
"""Ordered, non-blocking bridge from the engine to the KV router actor
|
||||
which maintains per-replica token load statistics.
|
||||
|
||||
``report`` only enqueues locally, so generation never blocks on the actor.
|
||||
A single delivery task per replica drains the queue to the actor, awaiting
|
||||
one ``on_lifecycle_events`` call at a time so events arrive in the order they
|
||||
were reported. Events that pile up during a call are sent together in the
|
||||
next one.
|
||||
"""
|
||||
|
||||
def __init__(self, actor: ActorHandle, worker_id: int):
|
||||
self.actor = actor
|
||||
self.worker_id = worker_id
|
||||
self._events: asyncio.Queue = asyncio.Queue()
|
||||
self._delivery_task: Optional[asyncio.Task] = None
|
||||
|
||||
def report(self, method_name: str, *args) -> None:
|
||||
if self._delivery_task is None or self._delivery_task.done():
|
||||
self._delivery_task = asyncio.get_running_loop().create_task(
|
||||
self._deliver()
|
||||
)
|
||||
self._events.put_nowait((method_name, args))
|
||||
|
||||
async def _deliver(self) -> None:
|
||||
while True:
|
||||
# Wait for the next event, then drain whatever queued up behind it
|
||||
# into the same batch.
|
||||
batch = [await self._events.get()]
|
||||
while not self._events.empty():
|
||||
batch.append(self._events.get_nowait())
|
||||
try:
|
||||
await self.actor.on_lifecycle_events.remote(batch)
|
||||
except (RayActorError, RayTaskError) as e:
|
||||
logger.warning("Dropping KV lifecycle events: %s", e)
|
||||
finally:
|
||||
for _ in batch:
|
||||
self._events.task_done()
|
||||
|
||||
async def flush(self) -> None:
|
||||
"""Wait until every reported event has been delivered."""
|
||||
await self._events.join()
|
||||
|
||||
def close(self) -> None:
|
||||
"""Cancel the delivery task on engine shutdown."""
|
||||
if self._delivery_task is not None:
|
||||
self._delivery_task.cancel()
|
||||
self._delivery_task = None
|
||||
|
||||
|
||||
class RequestTokenTracker:
|
||||
"""Drives the request lifecycle hooks for one ``generate()`` stream."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
forwarder: LifecycleEventForwarder,
|
||||
request_id: str,
|
||||
prompt_token_ids: List[int],
|
||||
expected_output_tokens: Optional[int],
|
||||
):
|
||||
self._forwarder = forwarder
|
||||
self._request_id = request_id
|
||||
self._cumulative = 0
|
||||
self._prefill_marked = False
|
||||
self._finished = False
|
||||
forwarder.report(
|
||||
"on_request_added",
|
||||
request_id,
|
||||
forwarder.worker_id,
|
||||
prompt_token_ids,
|
||||
expected_output_tokens,
|
||||
)
|
||||
|
||||
def on_output(self, output: RequestOutput) -> None:
|
||||
"""Observe one engine ``RequestOutput`` (forwarded to the caller as-is).
|
||||
|
||||
vLLM streams either DELTA chunks or a single FINAL_ONLY chunk; both
|
||||
carry only new tokens, so output progress simply accumulates.
|
||||
"""
|
||||
step_tokens = sum(len(o.token_ids or []) for o in output.outputs)
|
||||
if step_tokens == 0:
|
||||
# No new tokens this step (e.g. a finish-only chunk).
|
||||
return
|
||||
self._cumulative += step_tokens
|
||||
if not self._prefill_marked:
|
||||
# The first output token signals prefill completion.
|
||||
self._prefill_marked = True
|
||||
self._forwarder.report("on_prefill_complete", self._request_id)
|
||||
self._forwarder.report("on_decode_progress", self._request_id, self._cumulative)
|
||||
|
||||
def finish(self) -> None:
|
||||
"""Report completion exactly once."""
|
||||
if not self._finished:
|
||||
self._finished = True
|
||||
self._forwarder.report("on_request_completed", self._request_id)
|
||||
|
||||
|
||||
def enable_token_tracking(engine_cls: Type[AsyncLLM]) -> Type[AsyncLLM]:
|
||||
"""Decorator adding KV-router request lifecycle tracking."""
|
||||
|
||||
class TokenTrackingEngine(engine_cls):
|
||||
_lifecycle_forwarder: Optional[LifecycleEventForwarder] = None
|
||||
_resolve_warned: bool = False
|
||||
|
||||
def _resolve_lifecycle_forwarder(self) -> Optional[LifecycleEventForwarder]:
|
||||
if self._lifecycle_forwarder is None:
|
||||
try:
|
||||
actor = serve.get_deployment_actor(KV_ROUTER_ACTOR_NAME)
|
||||
worker_id = get_worker_id(
|
||||
serve.get_replica_context().replica_id.unique_id
|
||||
)
|
||||
self._lifecycle_forwarder = LifecycleEventForwarder(
|
||||
actor, worker_id
|
||||
)
|
||||
except Exception as e:
|
||||
# Warn once: resolution is retried per request until it succeeds.
|
||||
if not self._resolve_warned:
|
||||
self._resolve_warned = True
|
||||
logger.warning("KV token tracking disabled: %s", e)
|
||||
return self._lifecycle_forwarder
|
||||
|
||||
def shutdown(self, *args, **kwargs):
|
||||
if self._lifecycle_forwarder is not None:
|
||||
self._lifecycle_forwarder.close()
|
||||
return super().shutdown(*args, **kwargs)
|
||||
|
||||
async def generate(self, prompt, sampling_params, request_id, *args, **kwargs):
|
||||
stream = super().generate(
|
||||
prompt, sampling_params, request_id, *args, **kwargs
|
||||
)
|
||||
forwarder = self._resolve_lifecycle_forwarder()
|
||||
# CUMULATIVE repeats output-so-far per chunk; our accounting sums
|
||||
# deltas, so skip it rather than over-count. vLLM's OpenAI layer only
|
||||
# uses DELTA/FINAL_ONLY (*Request.to_sampling_params):
|
||||
# https://github.com/vllm-project/vllm/tree/main/vllm/entrypoints/openai
|
||||
if forwarder is None or (
|
||||
sampling_params.output_kind == RequestOutputKind.CUMULATIVE
|
||||
):
|
||||
async for output in stream:
|
||||
yield output
|
||||
return
|
||||
|
||||
lifecycle_request_id = get_serve_request_id() or request_id
|
||||
tracker = RequestTokenTracker(
|
||||
forwarder,
|
||||
lifecycle_request_id,
|
||||
_get_prompt_token_ids(prompt),
|
||||
# The request's own output cap is its expected length; weights
|
||||
# the selection service's decode-block decay.
|
||||
# TODO(jeffreywang): Use an agent-provided expected-OSL hint for
|
||||
# more accurate decode-load estimation.
|
||||
sampling_params.max_tokens,
|
||||
)
|
||||
try:
|
||||
async for output in stream:
|
||||
tracker.on_output(output)
|
||||
yield output
|
||||
finally:
|
||||
tracker.finish()
|
||||
|
||||
return TokenTrackingEngine
|
||||
@@ -0,0 +1,420 @@
|
||||
# These imports are used for metrics tracking, will remove for PR
|
||||
import logging
|
||||
import time
|
||||
from typing import (
|
||||
Any,
|
||||
List,
|
||||
Optional,
|
||||
)
|
||||
|
||||
import ray
|
||||
from ray.actor import ActorHandle
|
||||
from ray.llm._internal.serve.routing_policies.prefix_aware.prefix_tree import (
|
||||
PrefixTreeActor,
|
||||
)
|
||||
from ray.serve._private.common import ReplicaID
|
||||
from ray.serve._private.constants import (
|
||||
SERVE_CONTROLLER_NAME,
|
||||
SERVE_LOGGER_NAME,
|
||||
SERVE_NAMESPACE,
|
||||
)
|
||||
from ray.serve._private.replica_result import ReplicaResult
|
||||
from ray.serve._private.request_router import (
|
||||
PowerOfTwoChoicesRequestRouter,
|
||||
)
|
||||
from ray.serve._private.request_router.common import (
|
||||
PendingRequest,
|
||||
)
|
||||
from ray.serve._private.request_router.replica_wrapper import (
|
||||
RunningReplica,
|
||||
)
|
||||
from ray.serve._private.request_router.request_router import (
|
||||
LocalityMixin,
|
||||
MultiplexMixin,
|
||||
RequestRouter,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(SERVE_LOGGER_NAME)
|
||||
|
||||
|
||||
class PrefixCacheAffinityRouter(LocalityMixin, MultiplexMixin, RequestRouter):
|
||||
"""Extends the PowerOfTwoChoicesRequestRouter with prefix-matching capabilities.
|
||||
|
||||
This request router optimizes replica selection by considering input text prefixes:
|
||||
|
||||
1. Mixes between three strategies to balance prefix cache hit rate and load balancing:
|
||||
- When load is balanced (queue length difference < threshold), it selects replicas
|
||||
with the highest prefix match rate for the input text
|
||||
- When load is balanced but match rate is below 10%, it falls back to the smallest tenants
|
||||
- When load is imbalanced, it uses the default Power of Two selection
|
||||
|
||||
2. Maintains a prefix tree to track which replicas have processed similar inputs:
|
||||
- Inserts prompt text into the prefix tree after routing
|
||||
- Uses this history to inform future routing decisions
|
||||
|
||||
This approach improves performance by routing related requests to the same replicas,
|
||||
increasing cache locality and reducing overhead for language model inference.
|
||||
"""
|
||||
|
||||
def initialize_state(
|
||||
self,
|
||||
imbalanced_threshold: Optional[float] = float("inf"),
|
||||
match_rate_threshold: Optional[float] = 0.1,
|
||||
do_eviction: Optional[bool] = False,
|
||||
eviction_threshold_chars: Optional[int] = 400_000,
|
||||
eviction_target_chars: Optional[int] = 360_000,
|
||||
eviction_interval_secs: Optional[int] = 10,
|
||||
tree_actor: Optional[ActorHandle] = None,
|
||||
):
|
||||
"""Initialize the prefix-aware routing state and configuration.
|
||||
|
||||
Args:
|
||||
imbalanced_threshold: Threshold for queue length difference to consider
|
||||
load balanced. When the difference between replica queue lengths is
|
||||
less than this value, prefix-aware routing is used.
|
||||
match_rate_threshold: Minimum prefix match rate (0.0-1.0) required to
|
||||
use prefix-aware routing. If match rate is below this threshold,
|
||||
falls back to smallest tenant selection.
|
||||
do_eviction: Whether to enable automatic eviction of old prefix tree
|
||||
entries to manage memory usage.
|
||||
eviction_threshold_chars: Maximum number of characters in the prefix
|
||||
tree before eviction is triggered.
|
||||
eviction_target_chars: Target number of characters to reduce the
|
||||
prefix tree to during eviction.
|
||||
eviction_interval_secs: Interval in seconds between eviction checks
|
||||
when eviction is enabled.
|
||||
tree_actor: The actor to use for the prefix tree in a test environment.
|
||||
If None, a detached actor will be created/retrieved.
|
||||
"""
|
||||
# === Prefix-aware routing logic hyperparameters ===
|
||||
self._imbalanced_threshold = imbalanced_threshold
|
||||
self._match_rate_threshold = match_rate_threshold
|
||||
|
||||
# === Eviction policy ===
|
||||
self._do_eviction = do_eviction
|
||||
self._eviction_loop_running = False
|
||||
self._eviction_threshold_chars = eviction_threshold_chars
|
||||
# Default eviction_target_chars to eviction_threshold_chars if not specified
|
||||
self._eviction_target_chars = (
|
||||
eviction_target_chars
|
||||
if eviction_target_chars is not None
|
||||
else eviction_threshold_chars
|
||||
)
|
||||
self._eviction_interval_secs = eviction_interval_secs
|
||||
|
||||
if tree_actor is None:
|
||||
# Create deployment-specific detached actor to avoid replica ID conflicts
|
||||
# in multi-deployment scenarios (e.g., PD disaggregation with DP)
|
||||
deployment_name = self._deployment_id.name if self._deployment_id else None
|
||||
app_name = self._deployment_id.app_name if self._deployment_id else None
|
||||
actor_name = "LlmPrefixTreeActor"
|
||||
actor_namespace_components = [SERVE_NAMESPACE]
|
||||
|
||||
if app_name:
|
||||
actor_namespace_components.append(app_name)
|
||||
if deployment_name:
|
||||
actor_namespace_components.append(deployment_name)
|
||||
actor_namespace = "::".join(actor_namespace_components)
|
||||
|
||||
self._tree_actor = PrefixTreeActor.options(
|
||||
name=actor_name,
|
||||
namespace=actor_namespace,
|
||||
get_if_exists=True,
|
||||
lifetime="detached",
|
||||
).remote()
|
||||
|
||||
# Register the actor with the controller for cleanup on serve.shutdown()
|
||||
controller = ray.get_actor(SERVE_CONTROLLER_NAME, namespace=SERVE_NAMESPACE)
|
||||
ray.get(
|
||||
controller._register_shutdown_cleanup_actor.remote(self._tree_actor)
|
||||
)
|
||||
else:
|
||||
self._tree_actor = tree_actor
|
||||
|
||||
def _extract_text_from_request(self, pending_request: PendingRequest) -> str:
|
||||
"""Extracts the text content from a pending request for prefix matching.
|
||||
|
||||
Searches through request arguments for either 'messages' or 'prompt' attributes,
|
||||
then normalizes the content to a single string representation that can be used
|
||||
for prefix tree operations.
|
||||
|
||||
Args:
|
||||
pending_request: The request to extract text from
|
||||
|
||||
Returns:
|
||||
A string containing the prompt text or concatenated message contents
|
||||
|
||||
Raises:
|
||||
ValueError: If no prompt or messages attribute is found in the request
|
||||
"""
|
||||
prompt = None
|
||||
for arg in pending_request.args:
|
||||
valid_input_types = ["messages", "prompt"]
|
||||
for valid_input_type in valid_input_types:
|
||||
if hasattr(arg, valid_input_type):
|
||||
prompt = (
|
||||
arg.prompt if valid_input_type == "prompt" else arg.messages
|
||||
)
|
||||
break
|
||||
if prompt is not None:
|
||||
break
|
||||
if prompt is None:
|
||||
raise ValueError(
|
||||
"No request with message or prompt attribute found in pending_request.args"
|
||||
)
|
||||
|
||||
return self._normalize_prompt_to_string(prompt)
|
||||
|
||||
def _coerce_to_text(self, value: Any) -> str:
|
||||
if value is None:
|
||||
return ""
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
if isinstance(value, list):
|
||||
return "".join(self._coerce_to_text(item) for item in value)
|
||||
if isinstance(value, dict):
|
||||
text_value = value.get("text")
|
||||
if isinstance(text_value, str):
|
||||
return text_value
|
||||
if "content" in value:
|
||||
return self._coerce_to_text(value["content"])
|
||||
|
||||
return ""
|
||||
|
||||
def _normalize_prompt_to_string(self, prompt: Any) -> str:
|
||||
"""Normalize prompt/messages a single string of characters.
|
||||
This is not exhaustive (e.g. thinking parts, multimodal are not supported).
|
||||
TODO(seiji): find a more maintainable way to normalize the prompt/messages.
|
||||
|
||||
Supported:
|
||||
- string → return as-is
|
||||
- list of strings → concat
|
||||
- list of message dicts with 'content' as string → concat
|
||||
- list of message dicts with 'content' as list of dicts → concat the 'text' fields from those parts
|
||||
"""
|
||||
if isinstance(prompt, str):
|
||||
return prompt
|
||||
|
||||
if isinstance(prompt, list):
|
||||
return "".join(
|
||||
self._coerce_to_text(
|
||||
message.get("content") if isinstance(message, dict) else message
|
||||
)
|
||||
for message in prompt
|
||||
)
|
||||
|
||||
return ""
|
||||
|
||||
async def _prefix_match_best_replicas(
|
||||
self,
|
||||
pending_request: Optional[PendingRequest],
|
||||
candidate_replicas: List[RunningReplica],
|
||||
) -> List[RunningReplica]:
|
||||
"""
|
||||
Returns a set of candidate replicas, of which the one with the smallest replica queue will be chosen.
|
||||
0. Default: same as pow 2 request router, return 2 replicas at random.
|
||||
1. If load is balanced, choose replica(s) with highest prefix match rate. If highest hit rate is below 10% or no match found, use replicas with smallest KV cache usage.
|
||||
2. If load is imbalanced, use default.
|
||||
"""
|
||||
chosen_replica_id_strings = []
|
||||
if (
|
||||
pending_request is not None
|
||||
and pending_request.args is not None
|
||||
and len(pending_request.args) > 0
|
||||
):
|
||||
input_text = self._extract_text_from_request(pending_request)
|
||||
if input_text is not None:
|
||||
# Start Sphinx tag: __begin_load_balance_component__
|
||||
# Check for imbalanced load.
|
||||
highest_queue_len = 0
|
||||
lowest_queue_len = float("inf")
|
||||
not_in_cache: List[ReplicaID] = []
|
||||
if self._use_replica_queue_len_cache:
|
||||
# Populate available queue lens from the cache.
|
||||
for r in candidate_replicas:
|
||||
queue_len = self._replica_queue_len_cache.get(r.replica_id)
|
||||
if queue_len is None or queue_len >= r.max_ongoing_requests:
|
||||
not_in_cache.append(r)
|
||||
else:
|
||||
highest_queue_len = max(highest_queue_len, queue_len)
|
||||
lowest_queue_len = min(lowest_queue_len, queue_len)
|
||||
else:
|
||||
not_in_cache = candidate_replicas
|
||||
if len(not_in_cache) > 0:
|
||||
for r, queue_len in await self._probe_queue_lens(
|
||||
not_in_cache,
|
||||
0,
|
||||
):
|
||||
if queue_len is None:
|
||||
continue
|
||||
highest_queue_len = max(highest_queue_len, queue_len)
|
||||
lowest_queue_len = min(lowest_queue_len, queue_len)
|
||||
|
||||
is_imbalanced = (
|
||||
highest_queue_len - lowest_queue_len > self._imbalanced_threshold
|
||||
)
|
||||
# End Sphinx tag: __end_load_balance_component__
|
||||
# Start Sphinx tag: __begin_prefix_match_component__
|
||||
if not is_imbalanced:
|
||||
# Convert candidate replica IDs to strings for prefix matching.
|
||||
candidate_replica_ids_strings = [
|
||||
r.replica_id.to_full_id_str() for r in candidate_replicas
|
||||
]
|
||||
(matched_text, matched_tenant_id_strings,) = ray.get(
|
||||
self._tree_actor.prefix_match.remote(
|
||||
input_text, candidate_replica_ids_strings
|
||||
)
|
||||
)
|
||||
match_rate = len(matched_text) / len(input_text)
|
||||
if match_rate < self._match_rate_threshold:
|
||||
smallest_tenants_id_strings = ray.get(
|
||||
self._tree_actor.get_smallest_tenants.remote()
|
||||
)
|
||||
if (
|
||||
smallest_tenants_id_strings is not None
|
||||
and len(smallest_tenants_id_strings) > 0
|
||||
):
|
||||
chosen_replica_id_strings = smallest_tenants_id_strings
|
||||
else:
|
||||
if (
|
||||
matched_tenant_id_strings is not None
|
||||
and len(matched_tenant_id_strings) > 0
|
||||
):
|
||||
chosen_replica_id_strings = matched_tenant_id_strings
|
||||
# End Sphinx tag: __end_prefix_match_component__
|
||||
return [
|
||||
[
|
||||
self._replicas[ReplicaID.from_full_id_str(chosen_id_string)]
|
||||
for chosen_id_string in chosen_replica_id_strings
|
||||
]
|
||||
]
|
||||
|
||||
# Start Sphinx tag: __begin_on_replica_actor_died__
|
||||
def on_replica_actor_died(self, replica_id: ReplicaID):
|
||||
"""Drop replica from replica set so it's not considered for future requests."""
|
||||
super().on_replica_actor_died(replica_id)
|
||||
ray.get(self._tree_actor.remove_tenants.remote([replica_id.to_full_id_str()]))
|
||||
|
||||
# End Sphinx tag: __end_on_replica_actor_died__
|
||||
|
||||
def update_replicas(self, replicas: List[RunningReplica]):
|
||||
"""Update the set of available replicas to be considered for routing.
|
||||
|
||||
When the set of replicas changes, we may spawn additional routing tasks
|
||||
if there are pending requests.
|
||||
"""
|
||||
# 1) Record the old replica IDs
|
||||
old_ids = set(self._replica_id_set)
|
||||
|
||||
# 2) Run the default update_replicas logic
|
||||
super().update_replicas(replicas)
|
||||
|
||||
# 3) Figure out which replicas were added / removed
|
||||
new_ids = set(self._replica_id_set)
|
||||
added = new_ids - old_ids
|
||||
removed = old_ids - new_ids
|
||||
|
||||
# 4) Update the prefix tree with the changes
|
||||
if added:
|
||||
added_strings = [rid.to_full_id_str() for rid in added]
|
||||
ray.get(self._tree_actor.add_tenants.remote(added_strings, time.time()))
|
||||
|
||||
if removed:
|
||||
removed_strings = [rid.to_full_id_str() for rid in removed]
|
||||
ray.get(self._tree_actor.remove_tenants.remote(removed_strings))
|
||||
|
||||
# === Start tasks (if enabled and not already running) ===
|
||||
if self._do_eviction and not self._eviction_loop_running:
|
||||
ray.get(
|
||||
self._tree_actor.start_eviction_loop.remote(
|
||||
self._eviction_threshold_chars,
|
||||
self._eviction_target_chars,
|
||||
self._eviction_interval_secs,
|
||||
)
|
||||
)
|
||||
self._eviction_loop_running = True
|
||||
|
||||
async def choose_replicas(
|
||||
self,
|
||||
candidate_replicas: List[RunningReplica],
|
||||
pending_request: Optional[PendingRequest] = None,
|
||||
) -> List[RunningReplica]:
|
||||
"""One iteration of the power of two choices procedure that chooses
|
||||
(at most) two random available replicas.
|
||||
|
||||
For multiplexing, this will first attempt to choose replicas that have the
|
||||
requested model ID for a configured timeout. If no replicas with the matching
|
||||
model ID are available after that timeout, it will fall back to the regular
|
||||
procedure.
|
||||
"""
|
||||
# Start Sphinx tag: __begin_pow2_router_base__
|
||||
# Get fallback replicas from PowerOfTwoChoicesRequestRouter
|
||||
fallback_replicas = await PowerOfTwoChoicesRequestRouter.choose_replicas(
|
||||
self,
|
||||
candidate_replicas=candidate_replicas,
|
||||
pending_request=pending_request,
|
||||
)
|
||||
if pending_request is None or not fallback_replicas:
|
||||
return fallback_replicas
|
||||
# End Sphinx tag: __end_pow2_router_base__
|
||||
|
||||
if (
|
||||
pending_request is not None
|
||||
and pending_request.metadata.multiplexed_model_id
|
||||
):
|
||||
# Get candidates for multiplexed model ID.
|
||||
candidate_replica_ids = self.apply_multiplex_routing(
|
||||
pending_request=pending_request,
|
||||
)
|
||||
else:
|
||||
# Get candidates for locality preference.
|
||||
candidate_replica_ids = self.apply_locality_routing(
|
||||
pending_request=pending_request,
|
||||
)
|
||||
if not candidate_replica_ids:
|
||||
return fallback_replicas
|
||||
|
||||
# Convert candidate replica IDs to RunningReplica objects.
|
||||
replica_id_to_replica_map = {
|
||||
replica.replica_id: replica for replica in candidate_replicas
|
||||
}
|
||||
candidate_replicas = [
|
||||
replica_id_to_replica_map[candidate_replica_id]
|
||||
for candidate_replica_id in candidate_replica_ids
|
||||
]
|
||||
chosen_replicas = await self._prefix_match_best_replicas(
|
||||
pending_request, candidate_replicas
|
||||
)
|
||||
if chosen_replicas[0]:
|
||||
return chosen_replicas
|
||||
|
||||
return fallback_replicas
|
||||
|
||||
# Start Sphinx tag: __begin_on_request_routed__
|
||||
def on_request_routed(
|
||||
self,
|
||||
pending_request: PendingRequest,
|
||||
replica_id: ReplicaID,
|
||||
result: ReplicaResult,
|
||||
):
|
||||
"""Called when a request is routed to a replica.
|
||||
|
||||
This is used as a callback to update the state of the request router
|
||||
after a response is generated.
|
||||
"""
|
||||
# Right now this only inserts the prompt into the prefix tree, not the response (streaming response makes things complicated)
|
||||
if (
|
||||
pending_request is not None
|
||||
and pending_request.args is not None
|
||||
and len(pending_request.args) > 0
|
||||
):
|
||||
input_text = self._extract_text_from_request(pending_request)
|
||||
if input_text is not None:
|
||||
# Insert into prefix tree
|
||||
ray.get(
|
||||
self._tree_actor.insert.remote(
|
||||
input_text, replica_id.to_full_id_str(), time.time()
|
||||
)
|
||||
)
|
||||
|
||||
# End Sphinx tag: __end_on_request_routed__
|
||||
@@ -0,0 +1,620 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
from threading import RLock
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
import ray
|
||||
from ray.serve._private.constants import (
|
||||
SERVE_LOGGER_NAME,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(SERVE_LOGGER_NAME)
|
||||
|
||||
|
||||
class Node:
|
||||
"""
|
||||
Node in a prefix tree that represents a segment of text and can belong to multiple tenants.
|
||||
Each node also tracks the last access time for each tenant.
|
||||
Simple example of root node connected to two children Nodes:
|
||||
root = Node(text="", parent=None, edge_label_to_child={"f": fooNode, "b": barNode}, tenant_to_last_access_time={"tenant_1": 2})
|
||||
fooNode = Node(text="foo", parent=root, edge_label_to_child={}, tenant_to_last_access_time={"tenant_1": 1})
|
||||
barNode = Node(text="bar", parent=root, edge_label_to_child={}, tenant_to_last_access_time={"tenant_1": 2})
|
||||
|
||||
In the above example, "foo" was inserted at time 1, and "bar" was inserted at time 2.
|
||||
It follows that root was last accessed at time 2.
|
||||
"""
|
||||
|
||||
def __init__(self, text: str = "", parent: Optional[Node] = None) -> None:
|
||||
"""
|
||||
Initialize a node in the prefix tree.
|
||||
|
||||
Args:
|
||||
text: The text segment this node represents
|
||||
parent: The parent node of this node
|
||||
"""
|
||||
self.text: str = text
|
||||
self.parent: Optional[Node] = parent
|
||||
|
||||
# Maps first character to child node
|
||||
self.edge_label_to_child: Dict[str, Node] = {}
|
||||
# For each tenant that has inserted text matching this node, track its last access timestamp (in seconds)
|
||||
self.tenant_to_last_access_time: Dict[str, float] = {}
|
||||
# Doubly linked list pointers for LRU tracking per tenant
|
||||
# Points to the less recently used node (toward tail)
|
||||
self.tenant_to_older_node: Dict[str, Optional[Node]] = {}
|
||||
# Points to the more recently used node (toward head)
|
||||
self.tenant_to_newer_node: Dict[str, Optional[Node]] = {}
|
||||
|
||||
|
||||
class PrefixTree:
|
||||
"""
|
||||
Thread-safe multi-tenant prefix tree (approximate radix tree).
|
||||
|
||||
Features:
|
||||
1. Stores data for multiple tenants in the same tree structure
|
||||
2. Thread-safe with node-level locking for concurrent access
|
||||
3. LRU eviction based on tenant access time
|
||||
4. Efficient prefix matching across multiple tenants
|
||||
|
||||
|
||||
Example tree structure:
|
||||
Representing the strings inserted in order:
|
||||
- "helloworld" at time 1 by tenant_1
|
||||
- "hellothere" at time 2 by tenant_2
|
||||
- "hellothomas" at time 3 by tenant_2
|
||||
|
||||
root: [] {tenant_1: 1, tenant_2: 3}
|
||||
(h) → [hello] {tenant_1: 1, tenant_2: 3}
|
||||
(w) → [world] {tenant_1: 1}
|
||||
(t) → [th] {tenant_2: 3}
|
||||
(e) → [ere] {tenant_2: 2}
|
||||
(o) → [omas] {tenant_2: 3}
|
||||
|
||||
Legend for each node:
|
||||
- [text] = Node.text
|
||||
- {tenant, timestamp} = Node.tenant_to_last_access_time
|
||||
- (x) = edge label (first character used as key for parent's children)
|
||||
|
||||
PrefixTree instance variables:
|
||||
self.tenant_to_char_count = {"tenant_1": 10, "tenant_2": 14}
|
||||
self.tenant_to_lru_tail = {"tenant_1": Node("world"), "tenant_2": Node("ere")}
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
"""Initialize an empty prefix tree."""
|
||||
self.lock: RLock = RLock()
|
||||
|
||||
# Root is always the head of the LRU list for each tenant.
|
||||
self.root: Node = Node()
|
||||
|
||||
# Tracks total character count per tenant. Can be used by the replica request router to determine which tenant to evict, and by how much.
|
||||
# Also uses the keys to track the active tenants in the tree.
|
||||
self.tenant_to_char_count: Dict[str, int] = {}
|
||||
|
||||
# LRU tracking - root is always the head, tail is the least recently used.
|
||||
self.tenant_to_lru_tail: Dict[str, Optional[Node]] = {}
|
||||
self._eviction_thread: Optional[threading.Thread] = None
|
||||
self._eviction_stop_event: threading.Event = threading.Event()
|
||||
|
||||
@staticmethod
|
||||
def _shared_prefix_count(a: str, b: str) -> int:
|
||||
"""
|
||||
Count the number of shared characters at the beginning of two strings.
|
||||
|
||||
Args:
|
||||
a: First string
|
||||
b: Second string
|
||||
|
||||
Returns:
|
||||
Number of matching characters at the beginning
|
||||
"""
|
||||
return len(os.path.commonprefix([a, b]))
|
||||
|
||||
def _get_lru_chain(self, tenant: str) -> List[Node]:
|
||||
"""
|
||||
Get the LRU chain for a given tenant by traversing from the root to the oldest node.
|
||||
Note: This method is intended to be used only in tests.
|
||||
"""
|
||||
with self.lock:
|
||||
if tenant not in self.tenant_to_char_count:
|
||||
return []
|
||||
nodes = []
|
||||
current_node = self.root
|
||||
while current_node:
|
||||
nodes.append(current_node)
|
||||
current_node = current_node.tenant_to_older_node.get(tenant)
|
||||
return nodes
|
||||
|
||||
def _insert_node_into_linked_list(
|
||||
self,
|
||||
node: Node,
|
||||
newer_neighbor: Optional[Node],
|
||||
older_neighbor: Optional[Node],
|
||||
tenant: str,
|
||||
) -> None:
|
||||
"""
|
||||
Insert a node into the LRU list between two neighbors. Updates the neighbors' pointers and the tail pointer, if that changes.
|
||||
"""
|
||||
with self.lock:
|
||||
if tenant not in self.tenant_to_char_count:
|
||||
logger.debug(f"Tenant '{tenant}' does not exist. No action taken.")
|
||||
return
|
||||
|
||||
# Skip if node is the root
|
||||
if node == self.root:
|
||||
return
|
||||
|
||||
node.tenant_to_newer_node[tenant] = newer_neighbor
|
||||
node.tenant_to_older_node[tenant] = older_neighbor
|
||||
|
||||
if newer_neighbor:
|
||||
newer_neighbor.tenant_to_older_node[tenant] = node
|
||||
|
||||
if older_neighbor:
|
||||
older_neighbor.tenant_to_newer_node[tenant] = node
|
||||
|
||||
if self.tenant_to_lru_tail[tenant] == newer_neighbor:
|
||||
self.tenant_to_lru_tail[tenant] = node
|
||||
|
||||
def _remove_node_from_linked_list(self, node: Node, tenant: str) -> None:
|
||||
"""
|
||||
Remove a node from the LRU list. Updates the neighbors' pointers and the tail pointer, if that changes.
|
||||
"""
|
||||
with self.lock:
|
||||
if tenant not in self.tenant_to_char_count:
|
||||
logger.debug(f"Tenant '{tenant}' does not exist. No action taken.")
|
||||
return
|
||||
|
||||
# Skip if node is the root
|
||||
if node == self.root:
|
||||
return
|
||||
|
||||
# Connect older and newer neighbors
|
||||
older_neighbor = node.tenant_to_older_node.get(tenant)
|
||||
newer_neighbor = node.tenant_to_newer_node.get(tenant)
|
||||
|
||||
if older_neighbor:
|
||||
older_neighbor.tenant_to_newer_node[tenant] = newer_neighbor
|
||||
|
||||
if newer_neighbor:
|
||||
newer_neighbor.tenant_to_older_node[tenant] = older_neighbor
|
||||
|
||||
# Update tail pointer if necessary
|
||||
if self.tenant_to_lru_tail[tenant] == node:
|
||||
self.tenant_to_lru_tail[tenant] = newer_neighbor
|
||||
|
||||
# Remove node from list
|
||||
node.tenant_to_newer_node.pop(tenant, None)
|
||||
node.tenant_to_older_node.pop(tenant, None)
|
||||
|
||||
def _remove_tenant_single_node(self, tenant: str, node: Node) -> int:
|
||||
"""
|
||||
Remove a tenant from a single node.
|
||||
|
||||
Args:
|
||||
tenant: Tenant to remove
|
||||
node: Node to remove tenant from
|
||||
|
||||
Returns:
|
||||
Number of characters removed (0 if preconditions not met)
|
||||
"""
|
||||
with self.lock:
|
||||
if tenant not in self.tenant_to_char_count:
|
||||
logger.debug(f"Tenant '{tenant}' does not exist. No action taken.")
|
||||
return 0
|
||||
if tenant not in node.tenant_to_last_access_time:
|
||||
logger.debug(
|
||||
f"Tenant '{tenant}' does not have node '{node.text}'. No action taken."
|
||||
)
|
||||
return 0
|
||||
|
||||
removed_chars_len: int = len(node.text)
|
||||
self.tenant_to_char_count[tenant] -= removed_chars_len
|
||||
node.tenant_to_last_access_time.pop(tenant, None)
|
||||
|
||||
self._remove_node_from_linked_list(node, tenant)
|
||||
|
||||
# Clean up empty nodes
|
||||
if not node.tenant_to_last_access_time and node.parent:
|
||||
if (
|
||||
node.text and node.text[0] in node.parent.edge_label_to_child
|
||||
): # Defensive check
|
||||
node.parent.edge_label_to_child.pop(node.text[0], None)
|
||||
|
||||
return removed_chars_len
|
||||
|
||||
def add_tenants(self, tenants: List[str], time_s: float) -> None:
|
||||
"""
|
||||
Add multiple new tenants to the tree. Also inserts an empty string for each tenant into the tree.
|
||||
|
||||
For each tenant that already exists, a warning is logged and that tenant is skipped.
|
||||
|
||||
Args:
|
||||
tenants: List of tenants to add
|
||||
time_s: Current timestamp in seconds
|
||||
"""
|
||||
with self.lock:
|
||||
for tenant in tenants:
|
||||
if tenant in self.tenant_to_char_count:
|
||||
logger.debug(f"Tenant '{tenant}' already exists. Skipping.")
|
||||
continue
|
||||
|
||||
self.tenant_to_char_count[tenant] = 0
|
||||
self.tenant_to_lru_tail[tenant] = self.root
|
||||
|
||||
# Initialize the root node as the head of the LRU list for this tenant
|
||||
self.root.tenant_to_newer_node[tenant] = None
|
||||
self.root.tenant_to_older_node[tenant] = None
|
||||
self.insert("", tenant, time_s)
|
||||
|
||||
def insert(self, text: str, tenant: str, time_s: float) -> None:
|
||||
"""
|
||||
Insert text into tree for a specific tenant, but only if the tenant already exists.
|
||||
|
||||
If the tenant doesn't exist in the tree, this will log a warning and return without
|
||||
inserting anything. Use add_tenants() first to add a new tenant.
|
||||
|
||||
Args:
|
||||
text: Text to insert
|
||||
tenant: Tenant
|
||||
time_s: Current timestamp in seconds
|
||||
|
||||
Loop structure:
|
||||
1. We update the current node at the start of each iteration of the while loop.
|
||||
This includes updating tenant_to_char_count and tenant_to_last_access_time, and moving the node to the front of the LRU list.
|
||||
2. Each iteration then either:
|
||||
a. Breaks (if we've processed the entire string).
|
||||
b. Processes the next segment of text by:
|
||||
1. If no child exists for the first character, create a new leaf node that matches the current text.
|
||||
2. Then, match the current text with the child's text:
|
||||
a. If they share a prefix (partial match), split the node and traverse into the new parent.
|
||||
b. If they fully match, traverse into the child node.
|
||||
"""
|
||||
with self.lock:
|
||||
if tenant not in self.tenant_to_char_count:
|
||||
logger.debug(
|
||||
f"Tenant '{tenant}' does not exist. Use add_tenants() first."
|
||||
)
|
||||
return
|
||||
|
||||
curr_node: Node = self.root
|
||||
i: int = 0
|
||||
while i <= len(text):
|
||||
# Invariant at beginning of each iteration: assume curr_node has not been visited by tenant yet.
|
||||
# Update tenant info for current node.
|
||||
if tenant not in curr_node.tenant_to_last_access_time:
|
||||
self.tenant_to_char_count[tenant] += len(curr_node.text)
|
||||
|
||||
curr_node.tenant_to_last_access_time[tenant] = time_s
|
||||
if curr_node != self.root:
|
||||
self._remove_node_from_linked_list(curr_node, tenant)
|
||||
self._insert_node_into_linked_list(
|
||||
curr_node,
|
||||
self.root,
|
||||
self.root.tenant_to_older_node.get(tenant),
|
||||
tenant,
|
||||
)
|
||||
if i == len(text):
|
||||
break
|
||||
|
||||
first_char: str = text[i]
|
||||
curr_text: str = text[i:]
|
||||
|
||||
if first_char not in curr_node.edge_label_to_child:
|
||||
# No match, create new node. Don't update new node as "visited" by tenant yet; it will be done at the beginning of the next iteration.
|
||||
# e.g. curr_node.edge_label_to_child = {}, curr_text = "hello" -> curr_node.edge_label_to_child = {"h": Node("hello")}
|
||||
new_node: Node = Node(text=curr_text, parent=curr_node)
|
||||
curr_node.edge_label_to_child[first_char] = new_node
|
||||
# Add the node to the back of the LRU list; it will be moved to the front in the next iteration.
|
||||
self._insert_node_into_linked_list(
|
||||
new_node, self.tenant_to_lru_tail[tenant], None, tenant
|
||||
)
|
||||
|
||||
# Match found, check if we need to split
|
||||
matched_node: Node = curr_node.edge_label_to_child[first_char]
|
||||
shared_count: int = self._shared_prefix_count(
|
||||
matched_node.text, curr_text
|
||||
)
|
||||
if shared_count < len(matched_node.text):
|
||||
# Partial match, split node at matched point
|
||||
# Example:
|
||||
## Before update:
|
||||
### curr_node.edge_label_to_child = {"h": Node("helloworld")}, curr_text = "hellothere" -> shared_count = 5
|
||||
### matched_node = Node("helloworld")
|
||||
|
||||
## After update:
|
||||
### curr_node.edge_label_to_child = {"h": Node("hello", edge_label_to_child = {"w": Node("world")})}
|
||||
### parent_node = Node("hello"), matched_node = Node("world")
|
||||
### Copy matched_node.tenant_to_last_access_time to parent_node.tenant_to_last_access_time
|
||||
### Insert parent_node into the back of the LRU list; it will be moved to the front in the next iteration. (for the current tenant)
|
||||
### Insert parent_node between matched_node and matched_node's newer neighbor (for all other tenants)
|
||||
### (new) curr_text = "there", (new) curr_node = parent_node
|
||||
### Continue adding "there" to tree in next iteration
|
||||
|
||||
matched_text: str = matched_node.text[:shared_count]
|
||||
remaining_text: str = matched_node.text[shared_count:]
|
||||
|
||||
# Create new intermediate node
|
||||
# Note that we don't update new_parent.tenant_to_last_access_time yet; it will be done at the beginning of the next iteration.
|
||||
new_parent: Node = Node(text=matched_text, parent=curr_node)
|
||||
new_parent.tenant_to_last_access_time = (
|
||||
matched_node.tenant_to_last_access_time.copy()
|
||||
)
|
||||
# Insert new_parent into the back of the LRU list; it will be moved to the front in the next iteration. (for the current tenant)
|
||||
self._insert_node_into_linked_list(
|
||||
new_parent, self.tenant_to_lru_tail[tenant], None, tenant
|
||||
)
|
||||
# Insert new_parent between matched_node and matched_node's newer neighbor (for all other tenants)
|
||||
for existing_tenant in new_parent.tenant_to_last_access_time:
|
||||
if existing_tenant != tenant:
|
||||
self._insert_node_into_linked_list(
|
||||
new_parent,
|
||||
matched_node.tenant_to_newer_node.get(existing_tenant),
|
||||
matched_node,
|
||||
existing_tenant,
|
||||
)
|
||||
|
||||
# Update existing matched node
|
||||
matched_node.text = remaining_text
|
||||
matched_node.parent = new_parent
|
||||
|
||||
# Connect nodes
|
||||
new_parent.edge_label_to_child[remaining_text[0]] = matched_node
|
||||
curr_node.edge_label_to_child[first_char] = new_parent
|
||||
|
||||
# Continue traversal
|
||||
curr_node = new_parent
|
||||
i += shared_count
|
||||
else:
|
||||
# Full match, continue down the tree
|
||||
curr_node = matched_node
|
||||
i += shared_count
|
||||
|
||||
def prefix_match(
|
||||
self, text: str, available_tenants: Optional[List[str]] = None
|
||||
) -> Tuple[str, Optional[List[str]]]:
|
||||
"""
|
||||
Match text against tree and return matched text and matching tenants.
|
||||
|
||||
Args:
|
||||
text: Text to match
|
||||
available_tenants: List of tenants to match against (or None for all)
|
||||
|
||||
Returns:
|
||||
Tuple of (matched_text, matched_tenants):
|
||||
If the list of available tenants doesn't match any tenants in the tree: returns ("", None)
|
||||
When no prefix match is found (does not traverse further than the root node): returns ("", list of available tenants)
|
||||
When a prefix match is found: returns (matched_prefix, list of tenants that own the matched node)
|
||||
"""
|
||||
with self.lock:
|
||||
if available_tenants:
|
||||
# Filter available_tenants to only include those in the tree
|
||||
available_tenants = [
|
||||
tenant
|
||||
for tenant in available_tenants
|
||||
if tenant in self.tenant_to_char_count
|
||||
]
|
||||
if not available_tenants:
|
||||
return "", None
|
||||
else:
|
||||
available_tenants = list(self.tenant_to_char_count.keys())
|
||||
|
||||
curr_node: Node = self.root
|
||||
i: int = 0
|
||||
text_len: int = len(text)
|
||||
|
||||
while i < text_len:
|
||||
first_char: str = text[i]
|
||||
curr_text: str = text[i:]
|
||||
|
||||
if first_char in curr_node.edge_label_to_child:
|
||||
matched_node: Node = curr_node.edge_label_to_child[first_char]
|
||||
|
||||
# Check if any available tenants match this node
|
||||
if not any(
|
||||
tenant in matched_node.tenant_to_last_access_time
|
||||
for tenant in available_tenants
|
||||
):
|
||||
break
|
||||
|
||||
shared_count: int = self._shared_prefix_count(
|
||||
matched_node.text, curr_text
|
||||
)
|
||||
i += shared_count
|
||||
curr_node = matched_node
|
||||
|
||||
if shared_count < len(matched_node.text):
|
||||
# Partial match, stop here
|
||||
break
|
||||
else:
|
||||
# No match found, stop here
|
||||
break
|
||||
|
||||
# Find tenants in current node that match available tenants
|
||||
matched_tenants = [
|
||||
tenant
|
||||
for tenant in available_tenants
|
||||
if tenant in curr_node.tenant_to_last_access_time
|
||||
] or None
|
||||
|
||||
matched_text: str = text[:i]
|
||||
|
||||
return matched_text, matched_tenants
|
||||
|
||||
def remove_tenants(self, tenants: List[str]) -> Dict[str, int]:
|
||||
"""
|
||||
Remove multiple tenants and all their nodes from the tree.
|
||||
Time complexity: O(n) where n is the total number of nodes owned by all tenants.
|
||||
|
||||
Args:
|
||||
tenants: List of tenants to remove
|
||||
|
||||
Returns:
|
||||
Dictionary mapping each tenant to the number of characters removed
|
||||
(0 if tenant doesn't exist)
|
||||
"""
|
||||
chars_removed: Dict[str, int] = {}
|
||||
|
||||
with self.lock:
|
||||
for tenant in tenants:
|
||||
if tenant not in self.tenant_to_char_count:
|
||||
logger.debug(f"Tenant '{tenant}' does not exist. Skipping.")
|
||||
chars_removed[tenant] = 0
|
||||
continue
|
||||
|
||||
tenant_chars_removed: int = 0
|
||||
|
||||
# Start from the tail and remove all nodes
|
||||
current_tail = self.tenant_to_lru_tail.get(tenant)
|
||||
while current_tail:
|
||||
newer_neighbor = current_tail.tenant_to_newer_node.get(tenant)
|
||||
tenant_chars_removed += self._remove_tenant_single_node(
|
||||
tenant, current_tail
|
||||
)
|
||||
current_tail = newer_neighbor
|
||||
|
||||
# Clean up tenant references
|
||||
self.tenant_to_char_count.pop(tenant, None)
|
||||
self.tenant_to_lru_tail.pop(tenant, None)
|
||||
|
||||
chars_removed[tenant] = tenant_chars_removed
|
||||
|
||||
return chars_removed
|
||||
|
||||
def evict_tenant_by_lru(self, tenant: str, min_remove_size: int) -> int:
|
||||
"""
|
||||
Evict least recently used nodes for a tenant until minimum size is freed.
|
||||
Time complexity: O(m) where m is the number of nodes removed.
|
||||
|
||||
Args:
|
||||
tenant: The tenant to evict nodes from
|
||||
min_remove_size: Minimum number of characters to remove
|
||||
|
||||
Returns:
|
||||
Actual number of characters removed (0 if tenant doesn't exist)
|
||||
|
||||
Note:
|
||||
- All nodes with the same oldest access time are removed together to maintain tree integrity, even if only removing a subset of them satisfies the min_remove_size.
|
||||
- For more predictable eviction, use unique timestamps for each insertion.
|
||||
- The root node is never evicted as it serves as the permanent head of the LRU list.
|
||||
"""
|
||||
with self.lock:
|
||||
if tenant not in self.tenant_to_char_count:
|
||||
logger.debug(
|
||||
f"Cannot evict tenant '{tenant}': tenant does not exist. No action taken."
|
||||
)
|
||||
return 0
|
||||
|
||||
if self.tenant_to_char_count[tenant] < min_remove_size:
|
||||
logger.debug(
|
||||
f"Cannot evict {min_remove_size} characters from tenant '{tenant}', which has only "
|
||||
f"{self.tenant_to_char_count[tenant]} characters. Will remove all available characters."
|
||||
)
|
||||
min_remove_size = self.tenant_to_char_count[tenant]
|
||||
|
||||
total_chars_removed: int = 0
|
||||
|
||||
# Start removing from the tail (least recently used)
|
||||
current_tail = self.tenant_to_lru_tail.get(tenant)
|
||||
|
||||
# Continue until we've freed enough space or run out of nodes
|
||||
while total_chars_removed < min_remove_size and current_tail:
|
||||
# Stop if we've reached the root - the root is never evicted
|
||||
if current_tail == self.root:
|
||||
break
|
||||
|
||||
# Get the current timestamp to remove all nodes with this timestamp
|
||||
current_timestamp = current_tail.tenant_to_last_access_time[tenant]
|
||||
|
||||
# Collect all nodes with the same timestamp (guaranteed to be contiguous in our LRU list)
|
||||
while (
|
||||
current_tail != self.root # Never include the root
|
||||
and current_tail.tenant_to_last_access_time[tenant]
|
||||
== current_timestamp
|
||||
):
|
||||
newer_neighbor = current_tail.tenant_to_newer_node.get(tenant)
|
||||
total_chars_removed += self._remove_tenant_single_node(
|
||||
tenant, current_tail
|
||||
)
|
||||
current_tail = newer_neighbor
|
||||
|
||||
return total_chars_removed
|
||||
|
||||
def get_smallest_tenants(self) -> Optional[List[str]]:
|
||||
"""
|
||||
Get the tenants with the smallest total character count.
|
||||
|
||||
Returns:
|
||||
Tenants with smallest character count, or None if no tenants
|
||||
"""
|
||||
with self.lock:
|
||||
if not self.tenant_to_char_count:
|
||||
return None
|
||||
|
||||
min_count = min(self.tenant_to_char_count.values())
|
||||
return [
|
||||
tenant
|
||||
for tenant, count in self.tenant_to_char_count.items()
|
||||
if count == min_count
|
||||
]
|
||||
|
||||
def start_eviction_loop(
|
||||
self, eviction_threshold: int, eviction_target: int, interval_secs: float
|
||||
) -> bool:
|
||||
"""Start a single eviction loop within the actor itself.
|
||||
|
||||
Args:
|
||||
eviction_threshold: Minimum number of characters a tenant must have to be evicted
|
||||
eviction_target: The maximum number of characters a tenant should have after eviction
|
||||
interval_secs: Number of seconds between eviction checks
|
||||
|
||||
Returns:
|
||||
True if the loop was started, False if it was already running
|
||||
"""
|
||||
self._eviction_stop_event.clear()
|
||||
with self.lock:
|
||||
if self._eviction_thread is None:
|
||||
self._eviction_thread = threading.Thread(
|
||||
target=self._run_eviction_loop,
|
||||
args=(eviction_threshold, eviction_target, interval_secs),
|
||||
daemon=True,
|
||||
)
|
||||
self._eviction_thread.start()
|
||||
return True
|
||||
else:
|
||||
logger.debug("Eviction loop already running")
|
||||
return False
|
||||
|
||||
def _run_eviction_loop(self, eviction_threshold, eviction_target, interval_secs):
|
||||
while not self._eviction_stop_event.is_set():
|
||||
if self._eviction_stop_event.wait(interval_secs):
|
||||
# Stop event was set, exit loop asap
|
||||
break
|
||||
|
||||
with self.lock:
|
||||
for tenant, char_count in self.tenant_to_char_count.items():
|
||||
if char_count > eviction_threshold:
|
||||
excess = char_count - eviction_target
|
||||
self.evict_tenant_by_lru(tenant, excess)
|
||||
|
||||
def stop_eviction_loop(self):
|
||||
self._eviction_stop_event.set()
|
||||
if self._eviction_thread:
|
||||
self._eviction_thread.join()
|
||||
self._eviction_thread = None
|
||||
|
||||
|
||||
@ray.remote
|
||||
class PrefixTreeActor(PrefixTree):
|
||||
def getattr(self, attribute: str) -> Any:
|
||||
"""
|
||||
Get an attribute of the PrefixTree.
|
||||
Note: This method is intended to be used only in tests.
|
||||
"""
|
||||
return getattr(self, attribute)
|
||||
|
||||
def setattr(self, attribute: str, value: Any) -> None:
|
||||
setattr(self, attribute, value)
|
||||
@@ -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
|
||||
@@ -0,0 +1,111 @@
|
||||
import asyncio
|
||||
from typing import AsyncGenerator, Generic, Iterable, List, Optional, TypeVar
|
||||
|
||||
from ray.llm._internal.serve.constants import (
|
||||
MODEL_RESPONSE_BATCH_TIMEOUT_MS,
|
||||
)
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
class Batcher(Generic[T]):
|
||||
"""This class batches multiple responses from a generator into a list of
|
||||
single responses, at some time interval.
|
||||
|
||||
Args:
|
||||
generator: the async generator that this class pulls responses
|
||||
from.
|
||||
interval_ms: the interval at which this class yields the current batch.
|
||||
If None, this class will batch all responses from the generator
|
||||
together and yield the entire batch once.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
generator: AsyncGenerator[T, None],
|
||||
interval_ms: Optional[float] = MODEL_RESPONSE_BATCH_TIMEOUT_MS,
|
||||
):
|
||||
self.generator = generator
|
||||
self.queue: asyncio.Queue = asyncio.Queue()
|
||||
|
||||
if interval_ms is None:
|
||||
self.interval_s = None
|
||||
else:
|
||||
self.interval_s = interval_ms / 1000
|
||||
|
||||
if interval_ms == 0:
|
||||
return
|
||||
|
||||
self.done_event: asyncio.Event = asyncio.Event()
|
||||
|
||||
# We are okay with this task getting cancelled (to propagate cancellations)
|
||||
self.read_task = asyncio.create_task(self.read())
|
||||
|
||||
def _merge_results(self, results: List[T]) -> Iterable[T]:
|
||||
return results
|
||||
|
||||
async def stream(self) -> AsyncGenerator[Iterable[T], None]:
|
||||
"""Drain from the queue every interval_ms and yield the merged results"""
|
||||
|
||||
if self.interval_s == 0:
|
||||
async for item in self.generator:
|
||||
yield [item]
|
||||
|
||||
return
|
||||
|
||||
try:
|
||||
while True:
|
||||
# Wait for the interval or until we finish, whichever is faster.
|
||||
# We use an event to avoid asyncio.wait_for cancelling the real task on timeout.
|
||||
try:
|
||||
if self.interval_s is None:
|
||||
await self.done_event.wait()
|
||||
else:
|
||||
await asyncio.wait_for(
|
||||
self.done_event.wait(), timeout=self.interval_s
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
pass
|
||||
|
||||
# Get all elements from the queue
|
||||
results, is_done = self.check_done_and_drain()
|
||||
|
||||
# If there are results, merge and yield them
|
||||
if results:
|
||||
output = self._merge_results(results)
|
||||
yield output
|
||||
|
||||
# If the read task is done, exit the stream task
|
||||
if is_done:
|
||||
# Raise exception, if any
|
||||
self.read_task.result()
|
||||
break
|
||||
finally:
|
||||
# If the stream task is done, make sure to exit the read task
|
||||
if not self.read_task.done():
|
||||
self.read_task.cancel()
|
||||
|
||||
def check_done_and_drain(self):
|
||||
results = self.drain_queue()
|
||||
return results, self.read_task.done()
|
||||
|
||||
async def read(self):
|
||||
"""Read from the generator and put into the queue in a tight loop"""
|
||||
try:
|
||||
async for x in self.generator:
|
||||
self.queue.put_nowait(x)
|
||||
finally:
|
||||
self.done_event.set()
|
||||
|
||||
def drain_queue(self):
|
||||
"""Drain all results currently in the queue"""
|
||||
results = []
|
||||
try:
|
||||
while True:
|
||||
results.append(self.queue.get_nowait())
|
||||
except asyncio.QueueEmpty:
|
||||
pass
|
||||
return results
|
||||
@@ -0,0 +1,153 @@
|
||||
import pickle
|
||||
import time
|
||||
import uuid
|
||||
from typing import Any, Callable, Dict, List, Optional, Union
|
||||
|
||||
import ray
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
from ray.serve._private.common import RequestMetadata
|
||||
from ray.serve.handle import DeploymentHandle
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
# Timeout in seconds for waiting for deployment replicas to be populated
|
||||
BROADCAST_REPLICA_POPULATION_TIMEOUT_S = 30
|
||||
|
||||
|
||||
def broadcast(
|
||||
handle: DeploymentHandle,
|
||||
method_name: str,
|
||||
args: Union[Any, Callable[[Any], Any]] = None,
|
||||
kwargs: Union[Dict[str, Any], Callable[[Any], Dict[str, Any]]] = None,
|
||||
combine: Optional[Callable[[List[Any]], Any]] = None,
|
||||
) -> Any:
|
||||
"""
|
||||
Broadcasts a method call to all replicas of the given handle.
|
||||
|
||||
This is useful for broadcasting a control plane message such as kv-cache
|
||||
reset or weight update to all replicas of the given handle.
|
||||
|
||||
NOTE: This API is experimental and may later be promoted to a public API in
|
||||
Ray Serve directly. For now, it is only available in Ray LLM and is
|
||||
intended to enable control plane operations during RL training which is
|
||||
required when orchestrating trianing and inference loops.
|
||||
|
||||
Args:
|
||||
handle: The DeploymentHandle to broadcast to.
|
||||
method_name: The name of the method to call on the deployment.
|
||||
args: The arguments to pass to the method. Can be a list/tuple of args,
|
||||
or a callable that takes the replica object and returns args.
|
||||
kwargs: The keyword arguments to pass to the method. Can be a dict,
|
||||
or a callable that takes the replica object and returns kwargs.
|
||||
combine: An optional callable that takes the list of results from all
|
||||
replicas and returns an aggregated result. If not provided, returns
|
||||
the list of results. The default combine function is to return the
|
||||
list of results.
|
||||
|
||||
Returns:
|
||||
The result of the method call to all replicas. If combine is provided,
|
||||
returns the aggregated result. Otherwise, returns the list of results.
|
||||
"""
|
||||
if args is None:
|
||||
args = ()
|
||||
if kwargs is None:
|
||||
kwargs = {}
|
||||
|
||||
if not handle.is_initialized:
|
||||
# If the handle is not initialized, we initialize it here.
|
||||
# We enforce running the router in a separate loop to ensure it can
|
||||
# update its replica set asynchronously while we might be blocking or
|
||||
# waiting.
|
||||
handle._init(_run_router_in_separate_loop=True)
|
||||
|
||||
router = handle._router
|
||||
if router is None:
|
||||
raise RuntimeError("DeploymentHandle router is None.")
|
||||
|
||||
# Wait for both the replica set AND the request router to be populated.
|
||||
# `running_replicas_populated()` flips when DEPLOYMENT_TARGETS long-poll
|
||||
# arrives; `request_router` becomes non-None only after DEPLOYMENT_CONFIG
|
||||
# long-poll arrives and sets `_request_router_class`. These are independent
|
||||
# long-polls, so polling only the former races with the latter.
|
||||
#
|
||||
# In normal request flow this is hidden because `assign_request` awaits
|
||||
# `_request_router_initialized` before routing — but `broadcast()` bypasses
|
||||
# `assign_request` and pokes `_replica_id_set` directly, so it has to
|
||||
# synchronize itself.
|
||||
def _get_request_router():
|
||||
if hasattr(router, "_asyncio_router"):
|
||||
return router._asyncio_router.request_router
|
||||
if hasattr(router, "request_router"):
|
||||
return router.request_router
|
||||
return None
|
||||
|
||||
start_time = time.time()
|
||||
while not handle.running_replicas_populated() or _get_request_router() is None:
|
||||
if time.time() - start_time > BROADCAST_REPLICA_POPULATION_TIMEOUT_S:
|
||||
raise TimeoutError(
|
||||
"Timed out waiting for deployment router/replicas to initialize."
|
||||
)
|
||||
time.sleep(0.1)
|
||||
|
||||
request_router = _get_request_router()
|
||||
|
||||
replica_set = request_router._replica_id_set
|
||||
|
||||
# Execute calls
|
||||
futures = []
|
||||
|
||||
# We copy the set to avoid modification during iteration if that happens
|
||||
replicas = list(replica_set)
|
||||
|
||||
for replica in replicas:
|
||||
actor_name = replica.to_full_id_str()
|
||||
try:
|
||||
actor_handle = ray.get_actor(actor_name, namespace="serve")
|
||||
except ValueError:
|
||||
# Actor might be dead or not found
|
||||
continue
|
||||
|
||||
# Prepare args
|
||||
call_args = args
|
||||
call_kwargs = kwargs
|
||||
|
||||
if callable(args):
|
||||
call_args = args(replica)
|
||||
if callable(kwargs):
|
||||
call_kwargs = kwargs(replica)
|
||||
|
||||
if not isinstance(call_args, (list, tuple)):
|
||||
raise ValueError(f"args must be a list or tuple, got {type(call_args)}")
|
||||
|
||||
if not isinstance(call_kwargs, dict):
|
||||
# Fallback if callable returned something else or initial was not dict
|
||||
# But initial default is dict.
|
||||
if call_kwargs is None:
|
||||
call_kwargs = {}
|
||||
else:
|
||||
raise ValueError(f"kwargs must be a dict, got {type(call_kwargs)}")
|
||||
|
||||
# Prepare Metadata
|
||||
request_id = f"broadcast-{uuid.uuid4()}"
|
||||
dummy_rm = RequestMetadata(
|
||||
request_id=request_id,
|
||||
internal_request_id=request_id,
|
||||
call_method=method_name,
|
||||
)
|
||||
|
||||
pickled_rm = pickle.dumps(dummy_rm)
|
||||
|
||||
# Fire remote call
|
||||
# We collect futures to wait for them
|
||||
futures.append(
|
||||
actor_handle.handle_request.remote(pickled_rm, *call_args, **call_kwargs)
|
||||
)
|
||||
|
||||
# Wait for all calls to complete
|
||||
results = []
|
||||
if futures:
|
||||
results = ray.get(futures)
|
||||
|
||||
if combine:
|
||||
return combine(results)
|
||||
return results
|
||||
@@ -0,0 +1,237 @@
|
||||
"""
|
||||
Serve-specific LoRA utilities that use generic abstractions from lora_utils.py.
|
||||
|
||||
This module provides serve-specific functionality while using the generic
|
||||
LoRA abstractions from common/lora_utils.py. This ensures clean separation
|
||||
between generic and serve-specific concerns.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
from ray.llm._internal.common.constants import LORA_ADAPTER_CONFIG_NAME
|
||||
from ray.llm._internal.common.models import global_id_manager, make_async
|
||||
from ray.llm._internal.common.utils.cloud_utils import (
|
||||
LoraMirrorConfig,
|
||||
)
|
||||
from ray.llm._internal.common.utils.lora_utils import (
|
||||
CLOUD_OBJECT_MISSING,
|
||||
clean_model_id,
|
||||
clear_directory,
|
||||
get_base_model_id,
|
||||
get_lora_id,
|
||||
get_object_from_cloud,
|
||||
retry_with_exponential_backoff,
|
||||
sync_files_with_lock,
|
||||
)
|
||||
from ray.llm._internal.serve.core.configs.llm_config import (
|
||||
DiskMultiplexConfig,
|
||||
LLMConfig,
|
||||
)
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
async def get_lora_finetuned_context_length(bucket_uri: str) -> Optional[int]:
|
||||
"""Gets the sequence length used to tune the LoRA adapter.
|
||||
|
||||
Return: Returns the max sequence length for the adapter, if it exists.
|
||||
|
||||
Raises: HTTPException if the LoRA adapter config file isn't available
|
||||
in the cloud storage repository.
|
||||
"""
|
||||
if bucket_uri.endswith("/"):
|
||||
bucket_uri = bucket_uri.rstrip("/")
|
||||
object_uri = f"{bucket_uri}/{LORA_ADAPTER_CONFIG_NAME}"
|
||||
|
||||
object_str_or_missing_message = await get_object_from_cloud(object_uri)
|
||||
|
||||
if object_str_or_missing_message is CLOUD_OBJECT_MISSING:
|
||||
raise HTTPException(
|
||||
404,
|
||||
f"Unable to find LoRA adapter config file "
|
||||
f'"{LORA_ADAPTER_CONFIG_NAME}" in folder {bucket_uri}. '
|
||||
"Check that the file exists and that you have read permissions.",
|
||||
)
|
||||
else:
|
||||
adapter_config_str = object_str_or_missing_message
|
||||
adapter_config = json.loads(adapter_config_str)
|
||||
return adapter_config.get("max_length")
|
||||
|
||||
|
||||
async def download_multiplex_config_info(
|
||||
model_id: str, base_path: str
|
||||
) -> tuple[str, Optional[int]]:
|
||||
"""Downloads info needed to create a multiplex config.
|
||||
|
||||
Downloads objects using cloud storage provider APIs.
|
||||
|
||||
Returns: 2-tuple containing
|
||||
1. A bucket_uri for the bucket containing LoRA weights and config.
|
||||
2. The maximum LoRA sequence length.
|
||||
|
||||
Raises: HTTPException if the LoRA adapter config file isn't available
|
||||
in the cloud storage repository.
|
||||
"""
|
||||
bucket_uri = f"{base_path}/{model_id}"
|
||||
ft_context_length = await get_lora_finetuned_context_length(bucket_uri)
|
||||
return bucket_uri, ft_context_length
|
||||
|
||||
|
||||
async def get_lora_model_metadata(model_id: str, base_path: str) -> Dict[str, Any]:
|
||||
"""Get the lora model metadata for a given model id and base path.
|
||||
|
||||
Args:
|
||||
model_id: A lora model id of the form ``base_model_id:suffix:id``.
|
||||
base_path: The cloud storage path under which LoRA adapters are stored
|
||||
(typically ``llm_config.lora_config.dynamic_lora_loading_path``).
|
||||
|
||||
Returns:
|
||||
A dict with keys ``model_id``, ``base_model_id``,
|
||||
``max_request_context_length``, and ``bucket_uri``.
|
||||
"""
|
||||
base_model_id = get_base_model_id(model_id)
|
||||
lora_id = get_lora_id(model_id)
|
||||
|
||||
# Examples of the variables:
|
||||
# model_id: "meta-llama/Meta-Llama-3.1-8B-Instruct:my_suffix:aBc1234"
|
||||
# base_path: "s3://ray-llama-weights"
|
||||
# bucket_uri: "s3://ray-llama-weights/my_suffix:aBc1234"
|
||||
(
|
||||
bucket_uri,
|
||||
ft_context_length,
|
||||
) = await download_multiplex_config_info(lora_id, base_path)
|
||||
|
||||
return {
|
||||
"model_id": model_id,
|
||||
"base_model_id": base_model_id,
|
||||
"max_request_context_length": ft_context_length,
|
||||
# Note (genesu): `bucket_uri` affects where the lora weights are downloaded
|
||||
# from remote location.
|
||||
"bucket_uri": bucket_uri,
|
||||
}
|
||||
|
||||
|
||||
async def get_lora_mirror_config(
|
||||
model_id: str,
|
||||
llm_config: LLMConfig,
|
||||
) -> LoraMirrorConfig:
|
||||
"""Get LoRA mirror configuration for serve-specific LLM config."""
|
||||
metadata = await get_lora_model_metadata(
|
||||
model_id, llm_config.lora_config.dynamic_lora_loading_path
|
||||
)
|
||||
|
||||
return LoraMirrorConfig(
|
||||
lora_model_id=model_id,
|
||||
bucket_uri=metadata["bucket_uri"],
|
||||
max_total_tokens=metadata["max_request_context_length"],
|
||||
sync_args=None,
|
||||
)
|
||||
|
||||
|
||||
class LoraModelLoader:
|
||||
"""Download LoRA weights from remote storage and manage disk cache.
|
||||
|
||||
This class is serve-specific as it depends on DiskMultiplexConfig and
|
||||
other serve-specific concepts.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
lora_root: Optional[str] = None,
|
||||
download_timeout_s: Optional[float] = None,
|
||||
max_tries: int = 1,
|
||||
):
|
||||
self.lora_root = lora_root or "/tmp/ray/llm/lora/cache"
|
||||
self.disk_cache: Dict[str, DiskMultiplexConfig] = {}
|
||||
self.active_syncing_tasks: Dict[str, asyncio.Task[DiskMultiplexConfig]] = {}
|
||||
if download_timeout_s is not None and download_timeout_s <= 0:
|
||||
raise ValueError(
|
||||
f"download_timeout_s must be None or >0, got {download_timeout_s}"
|
||||
)
|
||||
self.download_timeout_s = download_timeout_s
|
||||
if max_tries < 1:
|
||||
raise ValueError(f"max_tries must be >=1, got {max_tries}")
|
||||
self.max_tries = max_tries
|
||||
|
||||
async def load_model_from_config(
|
||||
self, lora_model_id: str, llm_config
|
||||
) -> DiskMultiplexConfig:
|
||||
"""Load a LoRA model by first fetching its mirror config from S3."""
|
||||
lora_mirror_config = await get_lora_mirror_config(lora_model_id, llm_config)
|
||||
return await self.load_model(lora_model_id, lora_mirror_config)
|
||||
|
||||
async def load_model(
|
||||
self, lora_model_id: str, lora_mirror_config: LoraMirrorConfig
|
||||
) -> DiskMultiplexConfig:
|
||||
"""Load a LoRA model."""
|
||||
if lora_model_id in self.disk_cache:
|
||||
return self.disk_cache[lora_model_id]
|
||||
|
||||
if lora_model_id not in self.active_syncing_tasks:
|
||||
task = asyncio.create_task(self._load_model_async(lora_mirror_config))
|
||||
task.add_done_callback(
|
||||
lambda result: self.active_syncing_tasks.pop(lora_model_id, None)
|
||||
)
|
||||
self.active_syncing_tasks[lora_model_id] = task
|
||||
else:
|
||||
task = self.active_syncing_tasks[lora_model_id]
|
||||
|
||||
disk_config = await asyncio.shield(task)
|
||||
self.disk_cache[lora_model_id] = disk_config
|
||||
return disk_config
|
||||
|
||||
async def _load_model_async(
|
||||
self, lora_mirror_config: LoraMirrorConfig
|
||||
) -> DiskMultiplexConfig:
|
||||
return await self._load_model(lora_mirror_config)
|
||||
|
||||
@make_async
|
||||
def _load_model(self, lora_mirror_config: LoraMirrorConfig) -> DiskMultiplexConfig:
|
||||
return self._load_model_sync(lora_mirror_config)
|
||||
|
||||
@make_async
|
||||
def clear_cache(self):
|
||||
"""Clear the disk cache."""
|
||||
clear_directory(self.lora_root)
|
||||
|
||||
def _model_dir_path(self, model_id: str) -> str:
|
||||
"""Construct the path for the lora weight."""
|
||||
lora_id = get_lora_id(clean_model_id(model_id))
|
||||
path = os.path.join(self.lora_root, lora_id)
|
||||
os.makedirs(path, exist_ok=True)
|
||||
return path
|
||||
|
||||
def _download_lora(self, lora_mirror_config: LoraMirrorConfig) -> str:
|
||||
"""Download LoRA weights using generic download primitives."""
|
||||
model_local_path = self._model_dir_path(lora_mirror_config.lora_model_id)
|
||||
sync_files_with_lock(
|
||||
lora_mirror_config.bucket_uri,
|
||||
model_local_path,
|
||||
timeout=self.download_timeout_s,
|
||||
)
|
||||
return model_local_path
|
||||
|
||||
def _load_model_sync(
|
||||
self, lora_mirror_config: LoraMirrorConfig
|
||||
) -> DiskMultiplexConfig:
|
||||
"""Load a model from the given mirror configuration."""
|
||||
download_with_retries = retry_with_exponential_backoff(
|
||||
max_tries=self.max_tries,
|
||||
exception_to_check=Exception,
|
||||
)(lambda config: self._download_lora(config))
|
||||
|
||||
local_path = download_with_retries(lora_mirror_config)
|
||||
return DiskMultiplexConfig.model_validate(
|
||||
{
|
||||
"model_id": lora_mirror_config.lora_model_id,
|
||||
"max_total_tokens": lora_mirror_config.max_total_tokens,
|
||||
"local_path": local_path,
|
||||
"lora_assigned_int_id": global_id_manager.next(),
|
||||
}
|
||||
)
|
||||
@@ -0,0 +1,75 @@
|
||||
import asyncio
|
||||
from typing import Optional
|
||||
|
||||
import ray
|
||||
from ray.llm._internal.common.utils.download_utils import (
|
||||
download_model_files,
|
||||
)
|
||||
from ray.llm._internal.common.utils.import_utils import try_import
|
||||
from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
|
||||
torch = try_import("torch")
|
||||
transformers = try_import("transformers")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def initialize_remote_node(llm_config: LLMConfig) -> Optional[str]:
|
||||
|
||||
callback = llm_config.get_or_create_callback()
|
||||
engine_config = llm_config.get_engine_config()
|
||||
|
||||
local_path = download_model_files(
|
||||
model_id=engine_config.actual_hf_model_id,
|
||||
mirror_config=engine_config.mirror_config,
|
||||
download_model=callback.ctx.worker_node_download_model,
|
||||
download_extra_files=True,
|
||||
callback=callback,
|
||||
)
|
||||
|
||||
# Validate that the binary exists
|
||||
if local_path and local_path != engine_config.actual_hf_model_id:
|
||||
engine_config.hf_model_id = local_path
|
||||
|
||||
return local_path
|
||||
|
||||
|
||||
async def initialize_node(llm_config: LLMConfig):
|
||||
"""Implements node initialization for LLM engines.
|
||||
|
||||
Downloads model, tokenizer, and extra files as necessary.
|
||||
"""
|
||||
# Get callback instance (if configured) with context information
|
||||
callback = llm_config.get_or_create_callback()
|
||||
ctx = callback.ctx
|
||||
pg_table = ray.util.placement_group_table(ctx.placement_group)
|
||||
|
||||
node_set = set(pg_table["bundles_to_node_id"].values())
|
||||
download_tasks = []
|
||||
for node_id in node_set:
|
||||
node_affinity_strategy = (
|
||||
ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy(
|
||||
node_id=node_id,
|
||||
soft=False,
|
||||
)
|
||||
)
|
||||
download_tasks.append(
|
||||
ray.remote(initialize_remote_node).options(
|
||||
num_cpus=0,
|
||||
scheduling_strategy=node_affinity_strategy,
|
||||
runtime_env=ctx.runtime_env,
|
||||
)
|
||||
)
|
||||
|
||||
logger.info("Running tasks to download model files on worker nodes")
|
||||
paths = await asyncio.gather(
|
||||
*[download_task.remote(llm_config) for download_task in download_tasks]
|
||||
)
|
||||
|
||||
# assume that all paths are the same
|
||||
assert paths, "No paths returned from download_model_files"
|
||||
assert (
|
||||
len(set(paths)) == 1
|
||||
), "Paths returned from download_model_files are not the same"
|
||||
llm_config.get_engine_config().hf_model_id = paths[0]
|
||||
@@ -0,0 +1,68 @@
|
||||
"""Placement group utilities for Ray LLM Serve."""
|
||||
|
||||
from collections import defaultdict
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from ray.serve._private.utils import get_head_node_id
|
||||
from ray.util.placement_group import PlacementGroup, placement_group_table
|
||||
|
||||
|
||||
def _sort_bundle_indices_by_node(
|
||||
bundles_to_node_id: Dict[int, str],
|
||||
driver_node_id: str,
|
||||
) -> List[int]:
|
||||
"""Sort bundle indices so that same-node bundles are adjacent, driver node first.
|
||||
|
||||
Args:
|
||||
bundles_to_node_id: Mapping from bundle index to node ID.
|
||||
driver_node_id: The node ID of the driver node.
|
||||
|
||||
Returns:
|
||||
List of bundle indices sorted with driver node bundles first,
|
||||
then remaining nodes in deterministic order.
|
||||
"""
|
||||
node_to_bundles: Dict[str, List[int]] = defaultdict(list)
|
||||
# bundle_idx is already in ascending order: created sequentially during
|
||||
# placement group creation and preserved by the GCS protobuf.
|
||||
for bundle_idx, node_id in bundles_to_node_id.items():
|
||||
node_to_bundles[node_id].append(bundle_idx)
|
||||
|
||||
result: List[int] = []
|
||||
if driver_node_id in node_to_bundles:
|
||||
result.extend(node_to_bundles.pop(driver_node_id))
|
||||
for node_id in sorted(node_to_bundles.keys()):
|
||||
result.extend(node_to_bundles[node_id])
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def get_bundle_indices_sorted_by_node(
|
||||
pg: PlacementGroup,
|
||||
driver_node_id: Optional[str] = None,
|
||||
) -> List[int]:
|
||||
"""Return bundle indices sorted such that same-node bundles are adjacent, driver node first.
|
||||
|
||||
When a placement group is provisioned, adjacent bundle indices don't
|
||||
necessarily map to the same physical node. This utility reorders bundle
|
||||
indices so that bundles on the same node are grouped together.
|
||||
|
||||
The driver node's bundles come first so that global rank 0 (which hosts the
|
||||
distributed rendezvous store) is co-located with the driver.
|
||||
|
||||
Args:
|
||||
pg: A ready placement group.
|
||||
driver_node_id: Node ID whose bundles should be ordered first. Callers
|
||||
should pass the node that advertises the distributed master address
|
||||
so that global rank 0 is co-located with it. Defaults to the cluster
|
||||
head node, which may not own any of the placement group's bundles
|
||||
(e.g. a CPU-only head node in a GPU cluster), in which case rank 0
|
||||
falls back to the lexicographically-first node.
|
||||
|
||||
Returns:
|
||||
List of bundle indices sorted such that same-node bundles are adjacent, driver node first.
|
||||
"""
|
||||
table = placement_group_table(pg)
|
||||
bundles_to_node_id = table["bundles_to_node_id"]
|
||||
if driver_node_id is None:
|
||||
driver_node_id = get_head_node_id()
|
||||
return _sort_bundle_indices_by_node(bundles_to_node_id, driver_node_id)
|
||||
@@ -0,0 +1,328 @@
|
||||
"""Generic registry for LLM serving components using Ray's internal KV store.
|
||||
|
||||
This module provides a reusable registry mechanism that enables components to be
|
||||
registered in the driver process and accessed across all Ray processes in the cluster,
|
||||
including Ray Serve child processes.
|
||||
|
||||
Similar to RLlib/Tune's registry but with a fixed global prefix for cross-job access.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
from typing import Any, Callable
|
||||
|
||||
import ray._private.worker as worker
|
||||
import ray.cloudpickle as pickle
|
||||
from ray.experimental.internal_kv import (
|
||||
_internal_kv_del,
|
||||
_internal_kv_exists,
|
||||
_internal_kv_get,
|
||||
_internal_kv_initialized,
|
||||
_internal_kv_put,
|
||||
)
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
# Fixed prefix for cross-job accessibility (Serve deployments run in different jobs)
|
||||
_SERVE_REGISTRY_PREFIX = "serve_global"
|
||||
|
||||
|
||||
def _make_key(category: str, name: str) -> bytes:
|
||||
"""Generate a binary key for the KV store.
|
||||
|
||||
Args:
|
||||
category: The component category (e.g., "kv_connector_backend")
|
||||
name: The component name
|
||||
|
||||
Returns:
|
||||
The key to use for storing the value
|
||||
"""
|
||||
return (
|
||||
b"LLMServeRegistry:"
|
||||
+ _SERVE_REGISTRY_PREFIX.encode("ascii")
|
||||
+ b":"
|
||||
+ category.encode("ascii")
|
||||
+ b"/"
|
||||
+ name.encode("ascii")
|
||||
)
|
||||
|
||||
|
||||
def _create_loader(value: Any) -> Callable[[], Any]:
|
||||
"""Create a loader callable for a value.
|
||||
|
||||
Handles both direct objects/classes and string paths for lazy loading.
|
||||
|
||||
Args:
|
||||
value: Either:
|
||||
- A class, object, or callable (returns lambda: value)
|
||||
- A string in format "module_path:class_name" (creates import loader)
|
||||
|
||||
Returns:
|
||||
A callable that returns the value when called
|
||||
|
||||
Raises:
|
||||
ValueError: If value is a string but doesn't have the correct format
|
||||
"""
|
||||
if isinstance(value, str):
|
||||
if ":" not in value:
|
||||
raise ValueError(
|
||||
f"Invalid format for string value: '{value}'. "
|
||||
f"Expected format: 'module_path:class_name' or a class/object."
|
||||
)
|
||||
module_path, class_name = value.rsplit(":", 1)
|
||||
# Create a loader callable that imports on demand
|
||||
def loader():
|
||||
module = importlib.import_module(module_path)
|
||||
return getattr(module, class_name)
|
||||
|
||||
return loader
|
||||
else:
|
||||
# For direct objects/classes, create a simple loader
|
||||
return lambda: value
|
||||
|
||||
|
||||
class ComponentRegistry:
|
||||
"""Generic registry for LLM serving components using Ray's internal KV store.
|
||||
|
||||
This registry enables components to be registered in the driver process and
|
||||
accessed across all Ray processes in the cluster, including Ray Serve child processes.
|
||||
|
||||
Similar to RLlib/Tune's registry but with a fixed global prefix for cross-job access.
|
||||
|
||||
**Usage Pattern:**
|
||||
This registry is designed for a "register once, read many" pattern:
|
||||
- Components are typically registered in the driver process before deployment
|
||||
- Ray Serve replicas read from the KV store during initialization
|
||||
- Once a component is resolved and cached in a process, subsequent `get()` calls return the cached value without checking the KV store for updates
|
||||
|
||||
Example:
|
||||
# Create a registry for a component category
|
||||
registry = ComponentRegistry("my_component")
|
||||
|
||||
# Register a component
|
||||
registry.register("my_component", MyComponentClass)
|
||||
|
||||
# Get a registered component
|
||||
component = registry.get("my_component")
|
||||
|
||||
# Check if registered
|
||||
if registry.contains("my_component"):
|
||||
...
|
||||
"""
|
||||
|
||||
def __init__(self, category: str):
|
||||
"""Initialize a registry for a specific component category.
|
||||
|
||||
Args:
|
||||
category: The category name (e.g., "kv_connector_backend")
|
||||
"""
|
||||
self.category = category
|
||||
self._loader_cache: dict[str, Callable[[], Any]] = {}
|
||||
self._resolved_cache: dict[str, Any] = {}
|
||||
self._pending: dict[str, bytes] = {}
|
||||
|
||||
def register(self, name: str, value: Any) -> None:
|
||||
"""Register a component.
|
||||
|
||||
Args:
|
||||
name: The name to register under
|
||||
value: The component to register. Can be:
|
||||
- A class, object, or callable (serialized directly)
|
||||
- A string in format "module_path:class_name" (lazy-loaded via import)
|
||||
|
||||
Raises:
|
||||
ValueError: If the component is already registered. Use unregister() first if you need to change the registration.
|
||||
|
||||
Examples:
|
||||
# Register a class directly
|
||||
registry.register("MyClass", MyClass)
|
||||
|
||||
# Register via module path (lazy loading)
|
||||
registry.register("MyClass", "my.module:MyClass")
|
||||
"""
|
||||
# Prevent double registration to avoid cache inconsistencies
|
||||
if self.contains(name):
|
||||
raise ValueError(
|
||||
f"{self.category} '{name}' is already registered. "
|
||||
f"Use unregister() first if you need to change the registration."
|
||||
)
|
||||
|
||||
# Create a loader callable (handles both direct values and string paths)
|
||||
loader = _create_loader(value)
|
||||
|
||||
# Serialize the loader callable
|
||||
serialized = pickle.dumps(loader)
|
||||
|
||||
# Store loader in cache
|
||||
self._loader_cache[name] = loader
|
||||
|
||||
# Store in KV store if Ray is initialized, otherwise queue for later
|
||||
if _internal_kv_initialized():
|
||||
try:
|
||||
key = _make_key(self.category, name)
|
||||
_internal_kv_put(key, serialized, overwrite=True)
|
||||
logger.debug(f"Registered {self.category} '{name}' in KV store")
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Failed to register {self.category} '{name}' in KV store: {e}",
|
||||
exc_info=True,
|
||||
)
|
||||
self._pending[name] = serialized
|
||||
else:
|
||||
self._pending[name] = serialized
|
||||
|
||||
def get(self, name: str) -> Any:
|
||||
"""Get a registered component.
|
||||
|
||||
Args:
|
||||
name: The name of the component
|
||||
|
||||
Returns:
|
||||
The registered component. If registered with a string path,
|
||||
returns the imported class/object. If registered directly,
|
||||
returns the original value.
|
||||
|
||||
Raises:
|
||||
ValueError: If the component is not registered
|
||||
"""
|
||||
# Check resolved cache first.
|
||||
if name in self._resolved_cache:
|
||||
return self._resolved_cache[name]
|
||||
|
||||
loader = self._loader_cache.get(name)
|
||||
# If not in local loader cache, try fetching from KV store.
|
||||
if loader is None and _internal_kv_initialized():
|
||||
try:
|
||||
key = _make_key(self.category, name)
|
||||
serialized = _internal_kv_get(key)
|
||||
if serialized is not None:
|
||||
loader = pickle.loads(serialized)
|
||||
# Cache the loader for future gets.
|
||||
self._loader_cache[name] = loader
|
||||
logger.debug(f"Loaded {self.category} '{name}' from KV store")
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Failed to load {self.category} '{name}' from KV store: {e}",
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
if loader is not None:
|
||||
value = loader()
|
||||
self._resolved_cache[name] = value
|
||||
return value
|
||||
|
||||
# Not found
|
||||
raise ValueError(
|
||||
f"{self.category} '{name}' not found. "
|
||||
f"Registered: {list(self._loader_cache.keys())}"
|
||||
)
|
||||
|
||||
def contains(self, name: str) -> bool:
|
||||
"""Check if a component is registered.
|
||||
|
||||
Args:
|
||||
name: The name to check
|
||||
|
||||
Returns:
|
||||
True if registered, False otherwise
|
||||
"""
|
||||
if name in self._loader_cache:
|
||||
return True
|
||||
|
||||
if _internal_kv_initialized():
|
||||
try:
|
||||
key = _make_key(self.category, name)
|
||||
return _internal_kv_exists(key)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Failed to check if {self.category} '{name}' exists in KV store: {e}",
|
||||
exc_info=True,
|
||||
)
|
||||
return False
|
||||
|
||||
return False
|
||||
|
||||
def unregister(self, name: str) -> None:
|
||||
"""Unregister a component.
|
||||
|
||||
Removes the component from local cache, pending registrations, and KV store.
|
||||
|
||||
Args:
|
||||
name: The name of the component to unregister
|
||||
"""
|
||||
# Remove from local caches
|
||||
if name in self._loader_cache:
|
||||
del self._loader_cache[name]
|
||||
if name in self._resolved_cache:
|
||||
del self._resolved_cache[name]
|
||||
|
||||
# Remove from pending if present
|
||||
if name in self._pending:
|
||||
del self._pending[name]
|
||||
|
||||
# Remove from KV store if Ray is initialized
|
||||
if _internal_kv_initialized():
|
||||
try:
|
||||
key = _make_key(self.category, name)
|
||||
_internal_kv_del(key)
|
||||
logger.debug(f"Unregistered {self.category} '{name}' from KV store")
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Failed to unregister {self.category} '{name}' from KV store: {e}",
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
def flush_pending(self) -> None:
|
||||
"""Flush pending registrations to KV store.
|
||||
|
||||
This is called automatically when Ray initializes via _post_init_hooks.
|
||||
"""
|
||||
if not _internal_kv_initialized() or not self._pending:
|
||||
return
|
||||
|
||||
for name, serialized in self._pending.items():
|
||||
try:
|
||||
key = _make_key(self.category, name)
|
||||
_internal_kv_put(key, serialized, overwrite=True)
|
||||
logger.debug(
|
||||
f"Flushed pending registration for {self.category} '{name}'"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Failed to flush {self.category} '{name}': {e}", exc_info=True
|
||||
)
|
||||
|
||||
self._pending.clear()
|
||||
|
||||
|
||||
# Global registry instances for different component categories
|
||||
_registries: dict[str, ComponentRegistry] = {}
|
||||
|
||||
|
||||
def get_registry(category: str) -> ComponentRegistry:
|
||||
"""Get or create a registry for a component category.
|
||||
|
||||
Args:
|
||||
category: The component category name
|
||||
|
||||
Returns:
|
||||
The ComponentRegistry instance for this category
|
||||
"""
|
||||
if category not in _registries:
|
||||
_registries[category] = ComponentRegistry(category)
|
||||
return _registries[category]
|
||||
|
||||
|
||||
def _flush_all_registries():
|
||||
"""Flush all pending registrations to KV store.
|
||||
|
||||
This is registered as a Ray post-init hook to ensure registrations
|
||||
made before Ray initialization are available across processes.
|
||||
"""
|
||||
for registry in _registries.values():
|
||||
registry.flush_pending()
|
||||
|
||||
|
||||
if _flush_all_registries not in worker._post_init_hooks:
|
||||
worker._post_init_hooks.append(_flush_all_registries)
|
||||
@@ -0,0 +1,218 @@
|
||||
import asyncio
|
||||
import threading
|
||||
import time
|
||||
import traceback
|
||||
from functools import partial
|
||||
from typing import Awaitable, Callable, TypeVar
|
||||
|
||||
from fastapi import HTTPException, status
|
||||
from httpx import HTTPStatusError as HTTPXHTTPStatusError
|
||||
from pydantic import ValidationError as PydanticValidationError
|
||||
|
||||
from ray import serve
|
||||
from ray.llm._internal.common.errors import VLLM_FATAL_ERRORS
|
||||
from ray.llm._internal.serve.constants import DEFAULT_FATAL_ERROR_COOLDOWN_S
|
||||
from ray.llm._internal.serve.core.configs.openai_api_models import (
|
||||
ErrorInfo,
|
||||
ErrorResponse,
|
||||
OpenAIHTTPException,
|
||||
)
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
def _is_fatal_engine_error(e: Exception) -> bool:
|
||||
"""Detect fatal engine errors via isinstance check."""
|
||||
if not VLLM_FATAL_ERRORS:
|
||||
return False
|
||||
return isinstance(e, VLLM_FATAL_ERRORS)
|
||||
|
||||
|
||||
class _FatalEngineErrorLogHandler:
|
||||
"""Rate limits logging for fatal engine errors.
|
||||
|
||||
- First fatal error: logged with full traceback.
|
||||
- Subsequent occurences within ``cooldown_s``: suppressed.
|
||||
- Next fatal error after ``cooldown_s``: emits a summary with suppressed errors.
|
||||
- Fatal error after ``2 * cooldown_s`` of quiet: logs full traceback again.
|
||||
- Non-fatal errors: always logged, unaffected by rate limiting.
|
||||
"""
|
||||
|
||||
def __init__(self, cooldown_s: float = DEFAULT_FATAL_ERROR_COOLDOWN_S):
|
||||
self._cooldown_s = cooldown_s
|
||||
self._first_logged = False
|
||||
self._suppressed_count = 0
|
||||
self._last_summary_time = 0.0
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def log(
|
||||
self,
|
||||
e: Exception,
|
||||
request_id: str,
|
||||
status_code: int,
|
||||
) -> None:
|
||||
"""Log the error, rate limiting fatal engine errors."""
|
||||
is_fatal = _is_fatal_engine_error(e)
|
||||
|
||||
if not is_fatal:
|
||||
log_fn = logger.error if status_code >= 500 else logger.warning
|
||||
log_fn(
|
||||
f"Encountered failure while handling request {request_id}",
|
||||
exc_info=e,
|
||||
extra={"ray_serve_extra_fields": {"status_code": status_code}},
|
||||
)
|
||||
return
|
||||
|
||||
with self._lock:
|
||||
now = time.monotonic()
|
||||
|
||||
# If enough quiet time has passed, treat this as a new failure
|
||||
# event. The suppressed count is intentionally dropped since the
|
||||
# original fatal error's full traceback was already emitted.
|
||||
if (
|
||||
self._first_logged
|
||||
and (now - self._last_summary_time) >= 2 * self._cooldown_s
|
||||
):
|
||||
self._first_logged = False
|
||||
self._suppressed_count = 0
|
||||
|
||||
if not self._first_logged:
|
||||
self._first_logged = True
|
||||
self._last_summary_time = now
|
||||
logger.error(
|
||||
"Encountered failure while handling request %s",
|
||||
request_id,
|
||||
exc_info=e,
|
||||
extra={"ray_serve_extra_fields": {"status_code": status_code}},
|
||||
)
|
||||
return
|
||||
|
||||
self._suppressed_count += 1
|
||||
elapsed = now - self._last_summary_time
|
||||
if elapsed >= self._cooldown_s:
|
||||
logger.error(
|
||||
"Suppressed %d fatal engine error(s) in the last %.0fs. "
|
||||
"Engine is dead, awaiting replica restart.",
|
||||
self._suppressed_count,
|
||||
elapsed,
|
||||
)
|
||||
self._suppressed_count = 0
|
||||
self._last_summary_time = now
|
||||
|
||||
|
||||
_fatal_error_log_handler = _FatalEngineErrorLogHandler()
|
||||
|
||||
|
||||
def make_async(_func: Callable[..., T]) -> Callable[..., Awaitable[T]]:
|
||||
"""Take a blocking function, and run it on in an executor thread.
|
||||
|
||||
This function prevents the blocking function from blocking the asyncio event loop.
|
||||
The code in this function needs to be thread safe.
|
||||
"""
|
||||
|
||||
def _async_wrapper(*args, **kwargs) -> asyncio.Future:
|
||||
loop = asyncio.get_event_loop()
|
||||
func = partial(_func, *args, **kwargs)
|
||||
return loop.run_in_executor(executor=None, func=func)
|
||||
|
||||
return _async_wrapper
|
||||
|
||||
|
||||
def extract_message_from_exception(e: Exception) -> str:
|
||||
# If the exception is a Ray exception, we need to dig through the text to get just
|
||||
# the exception message without the stack trace
|
||||
# This also works for normal exceptions (we will just return everything from
|
||||
# format_exception_only in that case)
|
||||
message_lines = traceback.format_exception_only(type(e), e)[-1].strip().split("\n")
|
||||
message = ""
|
||||
# The stack trace lines will be prefixed with spaces, so we need to start from the bottom
|
||||
# and stop at the last line before a line with a space
|
||||
found_last_line_before_stack_trace = False
|
||||
for line in reversed(message_lines):
|
||||
if not line.startswith(" "):
|
||||
found_last_line_before_stack_trace = True
|
||||
if found_last_line_before_stack_trace and line.startswith(" "):
|
||||
break
|
||||
message = line + "\n" + message
|
||||
message = message.strip()
|
||||
return message
|
||||
|
||||
|
||||
def _extract_message(e):
|
||||
if isinstance(e, OpenAIHTTPException) and e.internal_message is not None:
|
||||
internal_message = e.internal_message
|
||||
else:
|
||||
internal_message = extract_message_from_exception(e)
|
||||
|
||||
if isinstance(e, HTTPException):
|
||||
message = e.detail
|
||||
elif isinstance(e, OpenAIHTTPException):
|
||||
message = e.message
|
||||
else:
|
||||
message = internal_message
|
||||
|
||||
return internal_message, message
|
||||
|
||||
|
||||
def get_response_for_error(
|
||||
e: Exception,
|
||||
request_id: str,
|
||||
) -> ErrorResponse:
|
||||
if isinstance(e, HTTPException):
|
||||
status_code = e.status_code
|
||||
elif isinstance(e, OpenAIHTTPException):
|
||||
status_code = e.status_code
|
||||
elif isinstance(e, PydanticValidationError):
|
||||
status_code = 400
|
||||
elif isinstance(e, HTTPXHTTPStatusError):
|
||||
status_code = e.response.status_code
|
||||
else:
|
||||
# Try to get the status code attribute from exception,
|
||||
# if not present, fallback to generic 500
|
||||
status_code = int(
|
||||
getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR)
|
||||
)
|
||||
|
||||
_fatal_error_log_handler.log(e, request_id, status_code)
|
||||
|
||||
if status_code == status.HTTP_500_INTERNAL_SERVER_ERROR:
|
||||
internal_message = message = "Internal Server Error"
|
||||
exc_type = "InternalServerError"
|
||||
else:
|
||||
internal_message, message = _extract_message(e)
|
||||
exc_type = e.__class__.__name__
|
||||
|
||||
# TODO make this more robust
|
||||
if "(Request ID: " not in message:
|
||||
message += f" (Request ID: {request_id})"
|
||||
|
||||
if "(Request ID: " not in internal_message:
|
||||
internal_message += f" (Request ID: {request_id})"
|
||||
|
||||
error_info = ErrorInfo(
|
||||
message=f"Message: {message}, Internal exception: {internal_message}, original exception: {str(e)}",
|
||||
code=status_code,
|
||||
type=exc_type,
|
||||
)
|
||||
error_response = ErrorResponse(error=error_info)
|
||||
return error_response
|
||||
|
||||
|
||||
def get_serve_request_id() -> str:
|
||||
"""Get request id from serve request context."""
|
||||
context = serve.context._serve_request_context.get()
|
||||
if context is not None:
|
||||
return context.request_id
|
||||
return ""
|
||||
|
||||
|
||||
def get_model_request_id(model: str):
|
||||
return model + "-" + get_serve_request_id()
|
||||
|
||||
|
||||
def replace_prefix(model: str) -> str:
|
||||
"""Replace -- with / in model name to handle slashes within the URL path segment"""
|
||||
return model.replace("--", "/")
|
||||
Reference in New Issue
Block a user