chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,13 @@
|
||||
from ray.llm._internal.batch.processor import (
|
||||
HttpRequestProcessorConfig,
|
||||
Processor,
|
||||
ProcessorBuilder,
|
||||
ProcessorConfig,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"Processor",
|
||||
"ProcessorConfig",
|
||||
"ProcessorBuilder",
|
||||
"HttpRequestProcessorConfig",
|
||||
]
|
||||
@@ -0,0 +1,635 @@
|
||||
#!/usr/bin/env python
|
||||
"""
|
||||
Benchmark Ray Data LLM offline batch inference throughput.
|
||||
|
||||
Sample usage:
|
||||
python ray.llm._internal.batch.benchmark.benchmark_processor --mode vllm_engine --batch-size 64 --concurrency 1 --num-prompts 10000 --model facebook/opt-1.3b
|
||||
--tensor-parallel-size 2 --pipeline-parallel-size 2 --distributed-executor-backend ray
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from time import perf_counter, sleep
|
||||
|
||||
import ray
|
||||
from .dataset import ShareGPTDataset
|
||||
from ray import data, serve
|
||||
from ray.data.llm import (
|
||||
ChatTemplateStageConfig,
|
||||
DetokenizeStageConfig,
|
||||
ServeDeploymentProcessorConfig,
|
||||
TokenizerStageConfig,
|
||||
build_processor,
|
||||
vLLMEngineProcessorConfig,
|
||||
)
|
||||
from ray.serve.llm import (
|
||||
LLMConfig,
|
||||
ModelLoadingConfig,
|
||||
build_llm_deployment,
|
||||
)
|
||||
from ray.serve.llm.openai_api_models import CompletionRequest
|
||||
|
||||
|
||||
class Mode(Enum):
|
||||
"""Processor to benchmark."""
|
||||
|
||||
VLLM_ENGINE = "vllm_engine"
|
||||
SHARED_VLLM_ENGINE = "shared_vllm_engine"
|
||||
SERVE_DEPLOYMENT = "serve_deployment"
|
||||
SHARED_SERVE_DEPLOYMENT = "shared_serve_deployment"
|
||||
CLASSIFY = "classify"
|
||||
|
||||
|
||||
# Default sampling parameters -- ensure a fair comparison by omitting sampling-induced variance
|
||||
VLLM_SAMPLING_PARAMS = {
|
||||
"top_p": 1.0,
|
||||
"temperature": 1.0,
|
||||
"max_tokens": 100,
|
||||
"ignore_eos": True,
|
||||
}
|
||||
|
||||
# Default vLLM engine kwargs
|
||||
VLLM_ENGINE_KWARGS = {
|
||||
"max_num_batched_tokens": 4096,
|
||||
}
|
||||
|
||||
# Default tokenization kwargs for classification -- truncate to max_model_len.
|
||||
CLASSIFY_TOKENIZATION_KWARGS_DEFAULT = {"truncation": True, "max_length": 512}
|
||||
|
||||
|
||||
def build_vllm_engine_kwargs(**kwargs) -> dict:
|
||||
"""Build vLLM engine kwargs from command line arguments."""
|
||||
engine_kwargs = VLLM_ENGINE_KWARGS.copy()
|
||||
engine_kwargs.update({k: v for k, v in kwargs.items() if v is not None})
|
||||
return engine_kwargs
|
||||
|
||||
|
||||
def _build_vllm_engine_config(
|
||||
model: str,
|
||||
batch_size: int,
|
||||
concurrency: int,
|
||||
pipeline_parallel_size: int = None,
|
||||
tensor_parallel_size: int = None,
|
||||
distributed_executor_backend: str = None,
|
||||
task_type: str = None,
|
||||
max_model_len: int = None,
|
||||
) -> vLLMEngineProcessorConfig:
|
||||
"""Helper to create vLLMEngineProcessorConfig."""
|
||||
engine_kwargs = build_vllm_engine_kwargs(
|
||||
pipeline_parallel_size=pipeline_parallel_size,
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
distributed_executor_backend=distributed_executor_backend,
|
||||
)
|
||||
if max_model_len is not None:
|
||||
engine_kwargs["max_model_len"] = max_model_len
|
||||
|
||||
config = vLLMEngineProcessorConfig(
|
||||
model_source=model,
|
||||
batch_size=batch_size,
|
||||
concurrency=concurrency,
|
||||
chat_template_stage=False,
|
||||
tokenize_stage=False,
|
||||
detokenize_stage=False,
|
||||
engine_kwargs=engine_kwargs,
|
||||
)
|
||||
|
||||
if task_type is not None:
|
||||
config.task_type = task_type
|
||||
|
||||
return config
|
||||
|
||||
|
||||
def _build_serve_deployment_config(
|
||||
batch_size: int,
|
||||
concurrency: int,
|
||||
deployment_name: str = None,
|
||||
app_name: str = None,
|
||||
) -> ServeDeploymentProcessorConfig:
|
||||
"""Helper to create ServeDeploymentProcessorConfig."""
|
||||
return ServeDeploymentProcessorConfig(
|
||||
deployment_name=deployment_name,
|
||||
app_name=app_name,
|
||||
dtype_mapping={
|
||||
"CompletionRequest": CompletionRequest,
|
||||
},
|
||||
batch_size=batch_size,
|
||||
concurrency=concurrency,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class BenchmarkResult:
|
||||
mode: Mode
|
||||
batch_size: int
|
||||
concurrency: int
|
||||
samples: int
|
||||
elapsed_s: float
|
||||
|
||||
@property
|
||||
def throughput(self) -> float:
|
||||
return self.samples / self.elapsed_s if self.elapsed_s else 0.0
|
||||
|
||||
def show(self) -> None:
|
||||
print("\n" + "=" * 60)
|
||||
print(f"BENCHMARK - {self.mode}")
|
||||
print("=" * 60)
|
||||
print(f"Samples : {self.samples}")
|
||||
print(f"Batch size : {self.batch_size}")
|
||||
print(f"Concurrency : {self.concurrency}")
|
||||
print(f"Time (s) : {self.elapsed_s:.2f}")
|
||||
print(f"Throughput : {self.throughput:.2f} req/s")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
def build_single_vllm_engine_processor(
|
||||
batch_size: int,
|
||||
concurrency: int,
|
||||
model: str,
|
||||
sampling_params: dict = VLLM_SAMPLING_PARAMS,
|
||||
pipeline_parallel_size: int = None,
|
||||
tensor_parallel_size: int = None,
|
||||
distributed_executor_backend: str = None,
|
||||
):
|
||||
"""Build vLLM engine processor for single-turn benchmark."""
|
||||
config = _build_vllm_engine_config(
|
||||
model,
|
||||
batch_size,
|
||||
concurrency,
|
||||
pipeline_parallel_size,
|
||||
tensor_parallel_size,
|
||||
distributed_executor_backend,
|
||||
)
|
||||
return build_processor(
|
||||
config,
|
||||
preprocess=lambda row: dict(
|
||||
prompt=row["prompt"],
|
||||
sampling_params=sampling_params,
|
||||
),
|
||||
postprocess=lambda row: row,
|
||||
)
|
||||
|
||||
|
||||
def build_shared_vllm_engine_processor(
|
||||
batch_size: int,
|
||||
concurrency: int,
|
||||
model: str,
|
||||
sampling_params: dict = VLLM_SAMPLING_PARAMS,
|
||||
pipeline_parallel_size: int = None,
|
||||
tensor_parallel_size: int = None,
|
||||
distributed_executor_backend: str = None,
|
||||
):
|
||||
"""Build vLLM engine processor for multi-turn benchmark."""
|
||||
config = _build_vllm_engine_config(
|
||||
model,
|
||||
batch_size,
|
||||
concurrency,
|
||||
pipeline_parallel_size,
|
||||
tensor_parallel_size,
|
||||
distributed_executor_backend,
|
||||
)
|
||||
|
||||
processor1 = build_processor(
|
||||
config,
|
||||
preprocess=lambda row: dict(
|
||||
prompt=row["prompt"],
|
||||
sampling_params=sampling_params,
|
||||
),
|
||||
postprocess=lambda row: {
|
||||
"prompt": row["generated_text"]
|
||||
if str(row.get("generated_text", "")).strip()
|
||||
else row["prompt"]
|
||||
},
|
||||
)
|
||||
|
||||
processor2 = build_processor(
|
||||
config,
|
||||
preprocess=lambda row: dict(
|
||||
prompt=row["prompt"],
|
||||
sampling_params=sampling_params,
|
||||
),
|
||||
postprocess=lambda row: row,
|
||||
)
|
||||
|
||||
def multi_turn_processor(dataset):
|
||||
return processor2(processor1(dataset))
|
||||
|
||||
return multi_turn_processor
|
||||
|
||||
|
||||
def build_classify_processor(
|
||||
batch_size: int,
|
||||
concurrency: int,
|
||||
model: str,
|
||||
tokenization_kwargs: dict = CLASSIFY_TOKENIZATION_KWARGS_DEFAULT,
|
||||
max_model_len: int = 512,
|
||||
distributed_executor_backend: str = None,
|
||||
):
|
||||
"""Build vLLM engine processor for classification benchmark."""
|
||||
|
||||
engine_kwargs = VLLM_ENGINE_KWARGS.copy()
|
||||
if distributed_executor_backend is not None:
|
||||
engine_kwargs["distributed_executor_backend"] = distributed_executor_backend
|
||||
|
||||
# Truncate prompts to max_model_len to avoid errors on long inputs.
|
||||
tokenization_kwargs = {**tokenization_kwargs, "max_length": max_model_len}
|
||||
|
||||
config = vLLMEngineProcessorConfig(
|
||||
model_source=model,
|
||||
task_type="classify",
|
||||
batch_size=batch_size,
|
||||
concurrency=concurrency,
|
||||
chat_template_stage=ChatTemplateStageConfig(enabled=False),
|
||||
tokenize_stage=TokenizerStageConfig(enabled=True),
|
||||
detokenize_stage=DetokenizeStageConfig(enabled=False),
|
||||
engine_kwargs=engine_kwargs,
|
||||
)
|
||||
return build_processor(
|
||||
config,
|
||||
preprocess=lambda row: dict(
|
||||
prompt=row["prompt"],
|
||||
tokenization_kwargs=tokenization_kwargs,
|
||||
),
|
||||
postprocess=lambda row: {
|
||||
"probs": float(row["embeddings"][0])
|
||||
if row.get("embeddings") is not None and len(row["embeddings"]) > 0
|
||||
else None,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def setup_serve_deployment(model: str, concurrency: int) -> tuple[str, str]:
|
||||
"""Set up Ray Serve deployment for hosting the LLM model."""
|
||||
deployment_name = "benchmark_deployment"
|
||||
app_name = "benchmark_app"
|
||||
|
||||
llm_config = LLMConfig(
|
||||
model_loading_config=ModelLoadingConfig(
|
||||
model_id=model,
|
||||
model_source=model,
|
||||
),
|
||||
deployment_config=dict(
|
||||
name=deployment_name,
|
||||
# To fairly compare with vLLM engine processor, fix the number of replicas to the concurrency level
|
||||
autoscaling_config=dict(
|
||||
min_replicas=concurrency,
|
||||
max_replicas=concurrency,
|
||||
),
|
||||
),
|
||||
engine_kwargs=dict(
|
||||
enable_prefix_caching=True,
|
||||
enable_chunked_prefill=True,
|
||||
max_num_batched_tokens=4096,
|
||||
),
|
||||
)
|
||||
|
||||
override_serve_options = dict(name=deployment_name)
|
||||
llm_app = build_llm_deployment(
|
||||
llm_config, override_serve_options=override_serve_options
|
||||
)
|
||||
serve.run(llm_app, name=app_name)
|
||||
|
||||
print("Waiting for Serve deployment to be ready...")
|
||||
max_wait_time = 120 # seconds
|
||||
wait_time = 0
|
||||
while not _is_app_ready(app_name) and wait_time < max_wait_time:
|
||||
sleep(5)
|
||||
wait_time += 5
|
||||
|
||||
if wait_time >= max_wait_time:
|
||||
raise TimeoutError("Deployment failed to become ready within timeout")
|
||||
|
||||
print("Deployment is ready!")
|
||||
return deployment_name, app_name
|
||||
|
||||
|
||||
def _is_app_ready(app_name: str) -> bool:
|
||||
try:
|
||||
serve_status = serve.status()
|
||||
|
||||
if app_name in serve_status.applications:
|
||||
app_status = serve_status.applications[app_name]
|
||||
if app_status.status == "RUNNING":
|
||||
print(f"Application '{app_name}' is RUNNING.")
|
||||
return True
|
||||
else:
|
||||
print(f"Application '{app_name}' status: {app_status.status}")
|
||||
return False
|
||||
else:
|
||||
print(f"Application '{app_name}' not found in Serve status.")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f"Error checking app status: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def build_single_serve_deployment_processor(
|
||||
batch_size: int,
|
||||
concurrency: int,
|
||||
model: str,
|
||||
sampling_params: dict = VLLM_SAMPLING_PARAMS,
|
||||
deployment_name: str = None,
|
||||
app_name: str = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""Build Serve deployment processor for single-turn benchmark."""
|
||||
config = _build_serve_deployment_config(
|
||||
batch_size,
|
||||
concurrency,
|
||||
deployment_name,
|
||||
app_name,
|
||||
)
|
||||
return build_processor(
|
||||
config,
|
||||
preprocess=lambda row: dict(
|
||||
method="completions",
|
||||
dtype="CompletionRequest",
|
||||
request_kwargs=dict(
|
||||
model=model,
|
||||
prompt=row["prompt"],
|
||||
**sampling_params,
|
||||
),
|
||||
),
|
||||
postprocess=lambda row: row,
|
||||
)
|
||||
|
||||
|
||||
def build_shared_serve_deployment_processor(
|
||||
batch_size: int,
|
||||
concurrency: int,
|
||||
model: str,
|
||||
sampling_params: dict = VLLM_SAMPLING_PARAMS,
|
||||
deployment_name: str = None,
|
||||
app_name: str = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""Build Serve deployment processor for multi-turn benchmark."""
|
||||
config = _build_serve_deployment_config(
|
||||
batch_size,
|
||||
concurrency,
|
||||
deployment_name,
|
||||
app_name,
|
||||
)
|
||||
|
||||
processor1 = build_processor(
|
||||
config,
|
||||
preprocess=lambda row: dict(
|
||||
method="completions",
|
||||
dtype="CompletionRequest",
|
||||
request_kwargs=dict(
|
||||
model=model,
|
||||
prompt=row["prompt"],
|
||||
stream=False,
|
||||
),
|
||||
),
|
||||
postprocess=lambda row: {
|
||||
# Fall back to original prompt if generated text is empty
|
||||
"prompt": (
|
||||
row["choices"][0]["text"]
|
||||
if row.get("choices") and str(row["choices"][0].get("text", "")).strip()
|
||||
else row["prompt"]
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
processor2 = build_processor(
|
||||
config,
|
||||
preprocess=lambda row: dict(
|
||||
method="completions",
|
||||
dtype="CompletionRequest",
|
||||
request_kwargs=dict(
|
||||
model=model,
|
||||
prompt=row["prompt"],
|
||||
stream=False,
|
||||
),
|
||||
),
|
||||
postprocess=lambda row: row,
|
||||
)
|
||||
|
||||
def multi_turn_processor(dataset):
|
||||
return processor2(processor1(dataset))
|
||||
|
||||
return multi_turn_processor
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Benchmark execution
|
||||
# -----------------------------------------------------------------------------
|
||||
def run_processor(
|
||||
mode: Mode,
|
||||
dataset: data.Dataset,
|
||||
builder,
|
||||
**kwargs,
|
||||
) -> BenchmarkResult:
|
||||
processor = builder(**kwargs)
|
||||
|
||||
total_samples = dataset.count()
|
||||
|
||||
start = perf_counter()
|
||||
processor(dataset).materialize()
|
||||
elapsed = perf_counter() - start
|
||||
|
||||
return BenchmarkResult(
|
||||
mode=mode,
|
||||
batch_size=kwargs.get("batch_size"),
|
||||
concurrency=kwargs.get("concurrency"),
|
||||
samples=total_samples,
|
||||
elapsed_s=elapsed,
|
||||
)
|
||||
|
||||
|
||||
def benchmark(
|
||||
mode: Mode,
|
||||
dataset: data.Dataset,
|
||||
*,
|
||||
batch_size: int,
|
||||
concurrency: int,
|
||||
model: str,
|
||||
sampling_params: dict = VLLM_SAMPLING_PARAMS,
|
||||
pipeline_parallel_size: int = None,
|
||||
tensor_parallel_size: int = None,
|
||||
distributed_executor_backend: str = None,
|
||||
) -> BenchmarkResult:
|
||||
mode_to_builder = {
|
||||
Mode.VLLM_ENGINE: build_single_vllm_engine_processor,
|
||||
Mode.SHARED_VLLM_ENGINE: build_shared_vllm_engine_processor,
|
||||
Mode.SERVE_DEPLOYMENT: build_single_serve_deployment_processor,
|
||||
Mode.SHARED_SERVE_DEPLOYMENT: build_shared_serve_deployment_processor,
|
||||
Mode.CLASSIFY: build_classify_processor,
|
||||
}
|
||||
|
||||
if mode not in mode_to_builder:
|
||||
raise ValueError(f"Unknown benchmark mode: {mode}")
|
||||
|
||||
builder = mode_to_builder[mode]
|
||||
|
||||
if mode in [Mode.SERVE_DEPLOYMENT, Mode.SHARED_SERVE_DEPLOYMENT]:
|
||||
deployment_name, app_name = setup_serve_deployment(model, concurrency)
|
||||
try:
|
||||
return run_processor(
|
||||
mode,
|
||||
dataset,
|
||||
builder,
|
||||
batch_size=batch_size,
|
||||
concurrency=concurrency,
|
||||
model=model,
|
||||
sampling_params=sampling_params,
|
||||
deployment_name=deployment_name,
|
||||
app_name=app_name,
|
||||
)
|
||||
finally:
|
||||
serve.delete(app_name)
|
||||
elif mode == Mode.CLASSIFY:
|
||||
return run_processor(
|
||||
mode,
|
||||
dataset,
|
||||
builder,
|
||||
batch_size=batch_size,
|
||||
concurrency=concurrency,
|
||||
model=model,
|
||||
distributed_executor_backend=distributed_executor_backend,
|
||||
)
|
||||
else:
|
||||
return run_processor(
|
||||
mode,
|
||||
dataset,
|
||||
builder,
|
||||
batch_size=batch_size,
|
||||
concurrency=concurrency,
|
||||
model=model,
|
||||
sampling_params=sampling_params,
|
||||
pipeline_parallel_size=pipeline_parallel_size,
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
distributed_executor_backend=distributed_executor_backend,
|
||||
)
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# CLI
|
||||
# -----------------------------------------------------------------------------
|
||||
def parse_args(argv: list[str]) -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="vLLM throughput benchmark")
|
||||
parser.add_argument(
|
||||
"--mode",
|
||||
choices=[mode.value for mode in Mode],
|
||||
default=Mode.VLLM_ENGINE.value,
|
||||
help="Ray Data LLM processor to run benchmarks for",
|
||||
)
|
||||
# Dataset configuration
|
||||
parser.add_argument(
|
||||
"--dataset-path",
|
||||
type=str,
|
||||
default="/home/ubuntu/datasets/Code-feedback-sharegpt-renamed",
|
||||
help="Path to dataset on disk",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-prompts", type=int, default=1000, help="Number of prompts to process"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hf-dataset-id",
|
||||
type=str,
|
||||
default="Crystalcareai/Code-feedback-sharegpt-renamed",
|
||||
help="Hugging Face dataset ID to download",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hf-split",
|
||||
type=str,
|
||||
default="train",
|
||||
help="Hugging Face dataset split to load",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--seed",
|
||||
type=int,
|
||||
default=0,
|
||||
help="Random seed for dataset sampling",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--truncate-prompt",
|
||||
type=int,
|
||||
default=512,
|
||||
help="Maximum prompt length",
|
||||
)
|
||||
# Engine configuration
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
required=True,
|
||||
help="LLM model to use",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pipeline-parallel-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Pipeline parallel size for vLLM engine",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tensor-parallel-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Tensor parallel size for vLLM engine",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--distributed-executor-backend",
|
||||
type=str,
|
||||
default=None,
|
||||
choices=["ray", "mp", "uni"],
|
||||
help="Distributed executor backend for vLLM engine",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-tokens",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Maximum number of tokens to generate per request (default: 100)",
|
||||
)
|
||||
# Ray Data worker configuration
|
||||
parser.add_argument(
|
||||
"--batch-size",
|
||||
type=int,
|
||||
required=True,
|
||||
help="Ray Data batch size for processing",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--concurrency", type=int, required=True, help="Ray Data concurrency level"
|
||||
)
|
||||
return parser.parse_args(argv)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args(sys.argv[1:])
|
||||
|
||||
ray.init()
|
||||
try:
|
||||
dataset = ShareGPTDataset(
|
||||
dataset_path=args.dataset_path,
|
||||
seed=args.seed,
|
||||
hf_dataset_id=args.hf_dataset_id,
|
||||
hf_split=args.hf_split,
|
||||
truncate_prompt=args.truncate_prompt,
|
||||
)
|
||||
prompts = dataset.sample(args.num_prompts)
|
||||
|
||||
dataset = data.from_items(prompts)
|
||||
|
||||
sampling_params = VLLM_SAMPLING_PARAMS.copy()
|
||||
if args.max_tokens is not None:
|
||||
sampling_params["max_tokens"] = args.max_tokens
|
||||
|
||||
result = benchmark(
|
||||
Mode(args.mode),
|
||||
dataset,
|
||||
batch_size=args.batch_size,
|
||||
concurrency=args.concurrency,
|
||||
model=args.model,
|
||||
sampling_params=sampling_params,
|
||||
pipeline_parallel_size=args.pipeline_parallel_size,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
distributed_executor_backend=args.distributed_executor_backend,
|
||||
)
|
||||
result.show()
|
||||
finally:
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,175 @@
|
||||
"""
|
||||
This module defines a dataset framework for sampling benchmark requests.
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from datasets import load_dataset, load_from_disk
|
||||
|
||||
|
||||
class BenchmarkDataset(ABC):
|
||||
DEFAULT_RANDOM_SEED = 0
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dataset_path: Optional[str] = None,
|
||||
random_seed: int = DEFAULT_RANDOM_SEED,
|
||||
) -> None:
|
||||
"""
|
||||
Abstract base class for benchmark datasets.
|
||||
|
||||
All benchmark datasets should inherit from this class and implement
|
||||
the required abstract methods.
|
||||
|
||||
Args:
|
||||
dataset_path: The path to the dataset on disk.
|
||||
random_seed: The seed for the random number generator.
|
||||
"""
|
||||
self._dataset_path = dataset_path
|
||||
self._random_seed = random_seed
|
||||
|
||||
@abstractmethod
|
||||
def load_data(self) -> None:
|
||||
"""
|
||||
Load data from the dataset source into memory.
|
||||
|
||||
Raises:
|
||||
NotImplementedError: If the method is not implemented in subclasses.
|
||||
"""
|
||||
raise NotImplementedError("load_data must be implemented in subclasses.")
|
||||
|
||||
@abstractmethod
|
||||
def sample(self, num_requests: int) -> List[Dict]:
|
||||
"""
|
||||
Sample prompts from the loaded dataset.
|
||||
|
||||
Args:
|
||||
num_requests: The number of prompts to sample from the dataset.
|
||||
|
||||
Returns:
|
||||
A list of sampled request dictionaries.
|
||||
|
||||
Raises:
|
||||
NotImplementedError: If the method is not implemented in subclasses.
|
||||
"""
|
||||
raise NotImplementedError("sample must be implemented in subclasses.")
|
||||
|
||||
|
||||
class ShareGPTDataset(BenchmarkDataset):
|
||||
"""Implements the ShareGPT dataset. The first human message of each conversation is used to build a prompt."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dataset_path: str,
|
||||
seed: int,
|
||||
hf_dataset_id: str = "Crystalcareai/Code-feedback-sharegpt-renamed",
|
||||
hf_split: str = "train",
|
||||
truncate_prompt: Optional[int] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Initializes the ShareGPTDataset.
|
||||
|
||||
Args:
|
||||
dataset_path: The path to the dataset on disk.
|
||||
seed: The seed for the random number generator.
|
||||
hf_dataset_id: The Hugging Face dataset ID to download if the dataset is not found on disk.
|
||||
hf_split: The Hugging Face split to load from the dataset.
|
||||
truncate_prompt: Maximum prompt length so that the prompt fits in the model's context window.
|
||||
"""
|
||||
super().__init__(dataset_path, seed)
|
||||
self._seed = seed
|
||||
|
||||
self._hf_dataset_id = hf_dataset_id
|
||||
self._hf_split = hf_split
|
||||
self._truncate_prompt = truncate_prompt
|
||||
|
||||
self._data: list[Dict] | None = None
|
||||
|
||||
def load_data(self) -> None:
|
||||
"""Load data from the dataset path into memory."""
|
||||
if self._data is None:
|
||||
self._data = self._load_dataset_data()
|
||||
|
||||
def sample(self, num_requests: int) -> List[Dict]:
|
||||
"""Sample prompts from the loaded dataset."""
|
||||
if self._data is None:
|
||||
self.load_data()
|
||||
|
||||
# Extract all valid prompts from the dataset
|
||||
all_prompts = []
|
||||
for item in self._data:
|
||||
prompt_data = self._extract_prompt(item)
|
||||
if prompt_data is not None:
|
||||
all_prompts.append(prompt_data)
|
||||
|
||||
if not all_prompts:
|
||||
raise ValueError("ShareGPT dataset yielded no usable prompts")
|
||||
|
||||
# Replicate samples if num_requests exceeds available samples
|
||||
if num_requests <= len(all_prompts):
|
||||
return all_prompts[:num_requests]
|
||||
|
||||
full_copies = num_requests // len(all_prompts)
|
||||
remainder = num_requests % len(all_prompts)
|
||||
prompts = all_prompts * full_copies + all_prompts[:remainder]
|
||||
return prompts
|
||||
|
||||
def _load_dataset(self):
|
||||
"""Load dataset from disk or Hugging Face."""
|
||||
path = Path(self._dataset_path)
|
||||
print(f"Attempting to load dataset from {path}")
|
||||
print(f"Dataset exists on disk: {path.exists()}")
|
||||
|
||||
try:
|
||||
if path.exists():
|
||||
dataset = load_from_disk(str(path))
|
||||
else:
|
||||
print(
|
||||
f"Dataset not found on disk, downloading from Hugging Face: {self._hf_dataset_id}"
|
||||
)
|
||||
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
dataset = load_dataset(self._hf_dataset_id, split=self._hf_split)
|
||||
dataset.save_to_disk(str(path))
|
||||
return dataset
|
||||
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Error loading ShareGPT dataset: {e}")
|
||||
|
||||
def _load_dataset_data(self) -> List[Dict]:
|
||||
"""Load and process dataset data into a list of dictionaries."""
|
||||
ds = self._load_dataset().shuffle(seed=self._seed)
|
||||
data = []
|
||||
|
||||
for i, row in enumerate(ds):
|
||||
data.append(row)
|
||||
|
||||
print(f"Loaded {len(data)} samples from dataset")
|
||||
return data
|
||||
|
||||
def _extract_prompt(self, item: Dict) -> Dict | None:
|
||||
"""
|
||||
Extracts the first human message of a conversation or None.
|
||||
|
||||
The ShareGPT schema uses {"role": "human", "value": ...} for user
|
||||
turns.
|
||||
"""
|
||||
messages = item.get("messages") or item.get("conversations") or []
|
||||
prompt = next(
|
||||
(
|
||||
str(msg.get("value", "")).strip()
|
||||
for msg in messages
|
||||
if msg.get("role") in {"human", "user"}
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
# Only return a valid prompt if it's not empty
|
||||
if prompt and prompt.strip():
|
||||
if self._truncate_prompt:
|
||||
prompt = prompt[: self._truncate_prompt]
|
||||
return {"prompt": prompt}
|
||||
|
||||
return None
|
||||
@@ -0,0 +1,39 @@
|
||||
from typing import Literal
|
||||
|
||||
|
||||
class TaskType:
|
||||
@classmethod
|
||||
def values(cls):
|
||||
"""Return a set of all valid task type values."""
|
||||
return {
|
||||
value
|
||||
for key, value in vars(cls).items()
|
||||
if not key.startswith("_") and isinstance(value, str)
|
||||
}
|
||||
|
||||
|
||||
class vLLMTaskType(TaskType):
|
||||
"""The type of task to run on the vLLM engine."""
|
||||
|
||||
# Generate text.
|
||||
GENERATE = "generate"
|
||||
|
||||
# Generate embeddings.
|
||||
EMBED = "embed"
|
||||
|
||||
# Classification (e.g., sequence classification models).
|
||||
CLASSIFY = "classify"
|
||||
|
||||
# Scoring (e.g., cross-encoder models).
|
||||
SCORE = "score"
|
||||
|
||||
|
||||
class SGLangTaskType(TaskType):
|
||||
"""The type of task to run on the SGLang engine."""
|
||||
|
||||
# Generate text.
|
||||
GENERATE = "generate"
|
||||
|
||||
|
||||
TypeVLLMTaskType = Literal[tuple(vLLMTaskType.values())]
|
||||
TypeSGLangTaskType = Literal[tuple(SGLangTaskType.values())]
|
||||
@@ -0,0 +1,23 @@
|
||||
from ray.llm._internal.batch.observability.logging.setup import (
|
||||
setup_logging,
|
||||
)
|
||||
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
|
||||
|
||||
_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,40 @@
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from ray._common.filters import CoreContextFilter
|
||||
|
||||
|
||||
def _setup_logger(logger_name: str) -> None:
|
||||
"""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 Data logger is already setup and has handlers.
|
||||
|
||||
Args:
|
||||
logger_name: logger name used to get the logger.
|
||||
"""
|
||||
logger = logging.getLogger(logger_name)
|
||||
llm_logger = logging.getLogger("ray.data")
|
||||
|
||||
# Skip setup if the logger already has handlers setup or if the parent (Data
|
||||
# logger) has handlers.
|
||||
if logger.handlers or llm_logger.handlers:
|
||||
return
|
||||
|
||||
# Set up stream handler, which logs to console as plaintext.
|
||||
stream_handler = logging.StreamHandler()
|
||||
stream_handler.addFilter(CoreContextFilter())
|
||||
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 Data logger.
|
||||
|
||||
Loggers by default are logging to stdout, and are expected to be scraped by an
|
||||
external process.
|
||||
"""
|
||||
logger_name = f"ray.data.{name}"
|
||||
_setup_logger(logger_name)
|
||||
return logging.getLogger(logger_name)
|
||||
@@ -0,0 +1,26 @@
|
||||
import logging
|
||||
|
||||
from ray._common.filters import CoreContextFilter
|
||||
from ray._common.formatters import JSONFormatter
|
||||
|
||||
|
||||
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.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,162 @@
|
||||
from enum import Enum
|
||||
from typing import Callable, Dict, Tuple, Union
|
||||
|
||||
import ray
|
||||
from ray._common.constants import HEAD_NODE_RESOURCE_NAME
|
||||
from ray._common.usage.usage_lib import record_extra_usage_tag
|
||||
from ray.llm._internal.batch.observability.logging import get_logger
|
||||
from ray.llm._internal.common.base_pydantic import BaseModelExtended
|
||||
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
|
||||
|
||||
LLM_BATCH_TELEMETRY_NAMESPACE = "llm_batch_telemetry"
|
||||
LLM_BATCH_TELEMETRY_ACTOR_NAME = "llm_batch_telemetry"
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class BatchModelTelemetry(BaseModelExtended):
|
||||
# Dedup identity only; never recorded as a tag value. A hash of model_source
|
||||
# so distinct models that share an architecture stay separate in the dedup
|
||||
# key while the cleartext model name never reaches the head-node actor.
|
||||
model_id_hash: str = ""
|
||||
processor_config_name: str = ""
|
||||
model_architecture: str = ""
|
||||
batch_size: int = 0
|
||||
accelerator_type: str = ""
|
||||
concurrency: Union[int, Tuple[int, int]] = 0
|
||||
task_type: str = ""
|
||||
|
||||
# For the parallel size, 0 means not supported.
|
||||
pipeline_parallel_size: int = 0
|
||||
tensor_parallel_size: int = 0
|
||||
data_parallel_size: int = 0
|
||||
|
||||
|
||||
class BatchTelemetryTags(str, Enum):
|
||||
"""Telemetry tags for RayLLM Batch."""
|
||||
|
||||
LLM_BATCH_PROCESSOR_CONFIG_NAME = "LLM_BATCH_PROCESSOR_CONFIG_NAME"
|
||||
LLM_BATCH_MODEL_ARCHITECTURE = "LLM_BATCH_MODEL_ARCHITECTURE"
|
||||
LLM_BATCH_SIZE = "LLM_BATCH_SIZE"
|
||||
LLM_BATCH_ACCELERATOR_TYPE = "LLM_BATCH_ACCELERATOR_TYPE"
|
||||
LLM_BATCH_CONCURRENCY = "LLM_BATCH_CONCURRENCY"
|
||||
LLM_BATCH_TASK_TYPE = "LLM_BATCH_TASK_TYPE"
|
||||
LLM_BATCH_PIPELINE_PARALLEL_SIZE = "LLM_BATCH_PIPELINE_PARALLEL_SIZE"
|
||||
LLM_BATCH_TENSOR_PARALLEL_SIZE = "LLM_BATCH_TENSOR_PARALLEL_SIZE"
|
||||
LLM_BATCH_DATA_PARALLEL_SIZE = "LLM_BATCH_DATA_PARALLEL_SIZE"
|
||||
|
||||
|
||||
@ray.remote(
|
||||
name=LLM_BATCH_TELEMETRY_ACTOR_NAME,
|
||||
namespace=LLM_BATCH_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 full telemetry identity (incl. model_id_hash) so repeated
|
||||
# identical processor builds overwrite while distinct models/configs
|
||||
# remain separate.
|
||||
self._tracking_telemetries: Dict[str, BatchModelTelemetry] = {}
|
||||
self._record_tag_func = record_extra_usage_tag
|
||||
|
||||
def _update_record_tag_func(self, record_tag_func: Callable) -> None:
|
||||
self._record_tag_func = record_tag_func
|
||||
|
||||
def _reset(self) -> None:
|
||||
"""Only used in tests to clear accumulated telemetries."""
|
||||
self._tracking_telemetries = {}
|
||||
|
||||
def generate_report(self) -> Dict[str, str]:
|
||||
return {
|
||||
BatchTelemetryTags.LLM_BATCH_PROCESSOR_CONFIG_NAME: ",".join(
|
||||
[t.processor_config_name for t in self._tracking_telemetries.values()]
|
||||
),
|
||||
BatchTelemetryTags.LLM_BATCH_MODEL_ARCHITECTURE: ",".join(
|
||||
[t.model_architecture for t in self._tracking_telemetries.values()]
|
||||
),
|
||||
BatchTelemetryTags.LLM_BATCH_SIZE: ",".join(
|
||||
[str(t.batch_size) for t in self._tracking_telemetries.values()]
|
||||
),
|
||||
BatchTelemetryTags.LLM_BATCH_ACCELERATOR_TYPE: ",".join(
|
||||
[t.accelerator_type for t in self._tracking_telemetries.values()]
|
||||
),
|
||||
BatchTelemetryTags.LLM_BATCH_CONCURRENCY: ",".join(
|
||||
[str(t.concurrency) for t in self._tracking_telemetries.values()]
|
||||
),
|
||||
BatchTelemetryTags.LLM_BATCH_TASK_TYPE: ",".join(
|
||||
[t.task_type for t in self._tracking_telemetries.values()]
|
||||
),
|
||||
BatchTelemetryTags.LLM_BATCH_PIPELINE_PARALLEL_SIZE: ",".join(
|
||||
[
|
||||
str(t.pipeline_parallel_size)
|
||||
for t in self._tracking_telemetries.values()
|
||||
]
|
||||
),
|
||||
BatchTelemetryTags.LLM_BATCH_TENSOR_PARALLEL_SIZE: ",".join(
|
||||
[
|
||||
str(t.tensor_parallel_size)
|
||||
for t in self._tracking_telemetries.values()
|
||||
]
|
||||
),
|
||||
BatchTelemetryTags.LLM_BATCH_DATA_PARALLEL_SIZE: ",".join(
|
||||
[str(t.data_parallel_size) for t in self._tracking_telemetries.values()]
|
||||
),
|
||||
}
|
||||
|
||||
def record(self, telemetry: BatchModelTelemetry) -> None:
|
||||
"""Upsert by identity and record telemetries."""
|
||||
from ray._common.usage.usage_lib import TagKey
|
||||
|
||||
self._tracking_telemetries[telemetry.model_dump_json()] = telemetry
|
||||
for key, value in self.generate_report().items():
|
||||
try:
|
||||
self._record_tag_func(TagKey.Value(key), value)
|
||||
except ValueError:
|
||||
# Tag not in the installed usage proto; skip rather than fail.
|
||||
continue
|
||||
|
||||
|
||||
class TelemetryAgent:
|
||||
"""Wrapper around the telemetry agent that calls the remote method until
|
||||
push_telemetry_report is called."""
|
||||
|
||||
def __init__(self):
|
||||
# Creation must never break processor construction, so swallow failures
|
||||
# (e.g. Ray not initialized, transient GCS issues) and degrade to a no-op.
|
||||
try:
|
||||
# get_if_exists makes creation atomic across concurrent drivers.
|
||||
self.remote_telemetry_agent = _TelemetryAgent.options(
|
||||
name=LLM_BATCH_TELEMETRY_ACTOR_NAME,
|
||||
namespace=LLM_BATCH_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()
|
||||
except Exception:
|
||||
self.remote_telemetry_agent = None
|
||||
logger.exception("Failed to initialize LLM batch telemetry agent")
|
||||
|
||||
def _update_record_tag_func(self, record_tag_func: Callable):
|
||||
if self.remote_telemetry_agent is not None:
|
||||
self.remote_telemetry_agent._update_record_tag_func.remote(record_tag_func)
|
||||
|
||||
def push_telemetry_report(self, telemetry: BatchModelTelemetry):
|
||||
# Telemetry must never break processor construction.
|
||||
if self.remote_telemetry_agent is None:
|
||||
return
|
||||
try:
|
||||
ray.get(self.remote_telemetry_agent.record.remote(telemetry))
|
||||
except Exception:
|
||||
logger.exception("Failed to push LLM batch telemetry")
|
||||
|
||||
|
||||
def get_or_create_telemetry_agent() -> TelemetryAgent:
|
||||
"""Helper to get or create the telemetry agent."""
|
||||
return TelemetryAgent()
|
||||
@@ -0,0 +1,94 @@
|
||||
"""Lazy re-exports for batch processors.
|
||||
|
||||
Each ``*_proc.py`` module pulls in heavy ML dependencies (transformers, vllm,
|
||||
sglang, ...). Eagerly importing all of them here causes a single
|
||||
``from ray.llm._internal.batch.processor import HttpRequestProcessorConfig``
|
||||
to load the entire ML stack and to fail when optional dependencies (e.g.
|
||||
sglang) are not installed.
|
||||
|
||||
We use PEP 562 ``__getattr__`` to load each engine-specific config only when
|
||||
it is first referenced. The ``base`` module is imported eagerly because it is
|
||||
cheap and exports types used everywhere else.
|
||||
|
||||
Note on registration side effects: each ``*_proc.py`` calls
|
||||
``ProcessorBuilder.register(...)`` at import time. With lazy loading the
|
||||
registration happens the first time the corresponding config is accessed via
|
||||
this package -- which is exactly when a user constructs the config and asks
|
||||
``ProcessorBuilder.build`` to build a processor for it, so the registry is
|
||||
populated in time for every realistic usage.
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from ray.llm._internal.batch.processor.base import (
|
||||
Processor,
|
||||
ProcessorBuilder,
|
||||
ProcessorConfig,
|
||||
)
|
||||
|
||||
# Mapping of public attribute name -> (submodule, attribute name in submodule).
|
||||
# Each entry is an engine-specific processor config whose defining submodule
|
||||
# transitively imports heavy optional dependencies.
|
||||
_LAZY_ATTRS = {
|
||||
"HttpRequestProcessorConfig": (
|
||||
"http_request_proc",
|
||||
"HttpRequestProcessorConfig",
|
||||
),
|
||||
"ServeDeploymentProcessorConfig": (
|
||||
"serve_deployment_proc",
|
||||
"ServeDeploymentProcessorConfig",
|
||||
),
|
||||
"SGLangEngineProcessorConfig": (
|
||||
"sglang_engine_proc",
|
||||
"SGLangEngineProcessorConfig",
|
||||
),
|
||||
"vLLMEngineProcessorConfig": (
|
||||
"vllm_engine_proc",
|
||||
"vLLMEngineProcessorConfig",
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def __getattr__(name):
|
||||
"""Lazily import engine-specific processor configs (PEP 562)."""
|
||||
try:
|
||||
submodule, attr = _LAZY_ATTRS[name]
|
||||
except KeyError:
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None
|
||||
|
||||
import importlib
|
||||
|
||||
module = importlib.import_module(f"{__name__}.{submodule}")
|
||||
value = getattr(module, attr)
|
||||
globals()[name] = value
|
||||
return value
|
||||
|
||||
|
||||
def __dir__():
|
||||
return sorted(set(globals()).union(_LAZY_ATTRS))
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.llm._internal.batch.processor.http_request_proc import ( # noqa: F401
|
||||
HttpRequestProcessorConfig,
|
||||
)
|
||||
from ray.llm._internal.batch.processor.serve_deployment_proc import ( # noqa: F401
|
||||
ServeDeploymentProcessorConfig,
|
||||
)
|
||||
from ray.llm._internal.batch.processor.sglang_engine_proc import ( # noqa: F401
|
||||
SGLangEngineProcessorConfig,
|
||||
)
|
||||
from ray.llm._internal.batch.processor.vllm_engine_proc import ( # noqa: F401
|
||||
vLLMEngineProcessorConfig,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ProcessorConfig",
|
||||
"ProcessorBuilder",
|
||||
"HttpRequestProcessorConfig",
|
||||
"vLLMEngineProcessorConfig",
|
||||
"SGLangEngineProcessorConfig",
|
||||
"ServeDeploymentProcessorConfig",
|
||||
"Processor",
|
||||
]
|
||||
@@ -0,0 +1,533 @@
|
||||
import logging
|
||||
from collections import OrderedDict
|
||||
from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
|
||||
|
||||
from pydantic import Field, field_validator, model_validator
|
||||
|
||||
from ray.data import Dataset
|
||||
from ray.data.block import UserDefinedFunction
|
||||
from ray.llm._internal.batch.stages import (
|
||||
StatefulStage,
|
||||
wrap_postprocess,
|
||||
wrap_preprocess,
|
||||
)
|
||||
from ray.llm._internal.common.base_pydantic import BaseModelExtended
|
||||
from ray.util.annotations import DeveloperAPI, PublicAPI
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ProcessorConfig(BaseModelExtended):
|
||||
"""The processor configuration."""
|
||||
|
||||
batch_size: int = Field(
|
||||
default=32,
|
||||
description="Large batch sizes are likely to saturate the compute resources "
|
||||
"and could achieve higher throughput. On the other hand, small batch sizes "
|
||||
"are more fault-tolerant and could reduce bubbles in the data pipeline. "
|
||||
"You can tune the batch size to balance the throughput and fault-tolerance "
|
||||
"based on your use case. Defaults to 32.",
|
||||
)
|
||||
resources_per_bundle: Optional[Dict[str, float]] = Field(
|
||||
default=None,
|
||||
description="[DEPRECATED] This parameter is deprecated and will be removed in a future version. ",
|
||||
deprecated=True,
|
||||
)
|
||||
accelerator_type: Optional[str] = Field(
|
||||
default=None,
|
||||
description="The accelerator type used by the LLM stage in a processor. "
|
||||
"Default to None, meaning that only the CPU will be used.",
|
||||
)
|
||||
concurrency: Union[int, Tuple[int, int]] = Field(
|
||||
default=1,
|
||||
description="The number of workers for data parallelism. Default to 1. "
|
||||
"If ``concurrency`` is a ``tuple`` ``(m, n)``, Ray creates an autoscaling "
|
||||
"actor pool that scales between ``m`` and ``n`` workers (``1 <= m <= n``). "
|
||||
"If ``concurrency`` is an ``int`` ``n``, Ray uses either a fixed pool of ``n`` "
|
||||
"workers or an autoscaling pool from ``1`` to ``n`` workers, depending on "
|
||||
"the processor and stage.",
|
||||
)
|
||||
|
||||
experimental: Dict[str, Any] = Field(
|
||||
default_factory=dict,
|
||||
description="[Experimental] Experimental configurations. "
|
||||
"Supported keys:\n"
|
||||
"`max_tasks_in_flight_per_actor`: [DEPRECATED] Prefer the top-level "
|
||||
"`max_tasks_in_flight_per_actor` field on `OfflineProcessorConfig`. "
|
||||
"Setting it here is still respected (and overridden by the top-level "
|
||||
"field if both are set), but logs a deprecation warning.",
|
||||
)
|
||||
|
||||
@field_validator("concurrency")
|
||||
def validate_concurrency(
|
||||
cls, concurrency: Union[int, Tuple[int, int]]
|
||||
) -> Union[int, Tuple[int, int]]:
|
||||
"""Validate that `concurrency` is either:
|
||||
- a positive int, or
|
||||
- a 2-tuple `(min, max)` of positive ints with `min <= max`.
|
||||
"""
|
||||
|
||||
def require(condition: bool, message: str) -> None:
|
||||
if not condition:
|
||||
raise ValueError(message)
|
||||
|
||||
if isinstance(concurrency, int):
|
||||
require(
|
||||
concurrency > 0,
|
||||
f"A positive integer for `concurrency` is expected! Got: `{concurrency}`.",
|
||||
)
|
||||
elif isinstance(concurrency, tuple):
|
||||
require(
|
||||
all(c > 0 for c in concurrency),
|
||||
f"`concurrency` tuple items must be positive integers! Got: `{concurrency}`.",
|
||||
)
|
||||
|
||||
min_concurrency, max_concurrency = concurrency
|
||||
require(
|
||||
min_concurrency <= max_concurrency,
|
||||
f"min > max in the concurrency tuple `{concurrency}`!",
|
||||
)
|
||||
return concurrency
|
||||
|
||||
def get_concurrency(self, autoscaling_enabled: bool = True) -> Dict[str, int]:
|
||||
"""Return a normalized dict of worker pool parameters from `self.concurrency`.
|
||||
|
||||
Behavior:
|
||||
- If `concurrency` is an int `n`:
|
||||
- `autoscaling_enabled` is True -> return `{"min_size": 1, "max_size": n}` (autoscaling).
|
||||
- `autoscaling_enabled` is False -> return `{"size": n}` (fixed-size pool).
|
||||
- If `concurrency` is a 2-tuple `(m, n)`, return `{"min_size": m, "max_size": n}`
|
||||
(the `autoscaling_enabled` flag is ignored).
|
||||
|
||||
Args:
|
||||
autoscaling_enabled: When False, treat an integer `concurrency` as fixed size;
|
||||
otherwise treat it as an autoscaling range from 1 to n. Defaults to True.
|
||||
|
||||
Returns:
|
||||
Dict[str, int]: A dictionary with either:
|
||||
- `{"size": n}` for fixed-size pools
|
||||
- `{"min_size": m, "max_size": n}` for autoscaling pools
|
||||
|
||||
Examples:
|
||||
>>> self.concurrency = (2, 4)
|
||||
>>> self.get_concurrency()
|
||||
{'min_size': 2, 'max_size': 4}
|
||||
>>> self.concurrency = 4
|
||||
>>> self.get_concurrency()
|
||||
{'min_size': 1, 'max_size': 4}
|
||||
>>> self.get_concurrency(autoscaling_enabled=False)
|
||||
{'size': 4}
|
||||
"""
|
||||
if isinstance(self.concurrency, int):
|
||||
if autoscaling_enabled:
|
||||
return {"min_size": 1, "max_size": self.concurrency}
|
||||
else:
|
||||
return {"size": self.concurrency}
|
||||
return {
|
||||
"min_size": self.concurrency[0],
|
||||
"max_size": self.concurrency[1],
|
||||
}
|
||||
|
||||
class Config:
|
||||
validate_assignment = True
|
||||
arbitrary_types_allowed = True
|
||||
|
||||
|
||||
class OfflineProcessorConfig(ProcessorConfig):
|
||||
"""The processor configuration for offline processing."""
|
||||
|
||||
model_source: str = Field(
|
||||
description="The model source to use for the offline processing.",
|
||||
)
|
||||
runtime_env: Optional[Dict[str, Any]] = Field(
|
||||
default=None,
|
||||
description="The runtime environment to use for the offline processing.",
|
||||
)
|
||||
max_pending_requests: Optional[int] = Field(
|
||||
default=None,
|
||||
description="The maximum number of pending requests. If not specified, "
|
||||
"will use the default value from the backend engine.",
|
||||
)
|
||||
max_concurrent_batches: int = Field(
|
||||
default=8,
|
||||
description="The maximum number of concurrent batches in the engine. "
|
||||
"This is to overlap the batch processing to avoid the tail latency of "
|
||||
"each batch. The default value may not be optimal when the batch size "
|
||||
"or the batch processing latency is too small, but it should be good "
|
||||
"enough for batch size >= 32. Sets the engine actor's Ray Core "
|
||||
"`max_concurrency`.",
|
||||
)
|
||||
max_tasks_in_flight_per_actor: Optional[int] = Field(
|
||||
default=None,
|
||||
description="Max tasks Ray Data submits concurrently to each engine "
|
||||
"actor. Passed through to `ray.data.ActorPoolStrategy`. If unset, Ray "
|
||||
"Data uses `ray.data.DataContext.max_tasks_in_flight_per_actor` if set "
|
||||
"globally. Otherwise, it defaults to `2 * max_concurrent_batches`; the "
|
||||
"factor can be overridden via the "
|
||||
"`RAY_DATA_ACTOR_DEFAULT_MAX_TASKS_IN_FLIGHT_TO_MAX_CONCURRENCY_FACTOR` "
|
||||
"env var. "
|
||||
"Setting this lower than `max_concurrent_batches` can underutilize the "
|
||||
"engine actor because Ray Data submits fewer tasks than the actor can "
|
||||
"process concurrently.",
|
||||
)
|
||||
should_continue_on_error: bool = Field(
|
||||
default=False,
|
||||
description="If True, continue processing when inference fails for a row "
|
||||
"instead of raising an exception. Failed rows will have a non-empty "
|
||||
"'__inference_error__' column containing the error message, and other "
|
||||
"output columns will be empty strings. Error rows bypass postprocess. "
|
||||
"If False (default), any inference error will raise an exception.",
|
||||
)
|
||||
|
||||
# Processor stage configurations (legacy booleans, will be deprecated).
|
||||
# TODO (jeffreywang): Remove apply_chat_template, chat_template, tokenize,
|
||||
# detokenize in Ray 2.57.0 in favor of the *_stage fields below.
|
||||
apply_chat_template: bool = Field(
|
||||
default=True,
|
||||
description="[DEPRECATED] Prefer `chat_template_stage`. Whether to apply chat template.",
|
||||
)
|
||||
chat_template: Optional[str] = Field(
|
||||
default=None,
|
||||
description="[DEPRECATED] Prefer `chat_template_stage.chat_template`. The chat template to use.",
|
||||
)
|
||||
tokenize: bool = Field(
|
||||
default=True,
|
||||
description="[DEPRECATED] Prefer `tokenize_stage`. Whether to tokenize input before engine.",
|
||||
)
|
||||
detokenize: bool = Field(
|
||||
default=True,
|
||||
description="[DEPRECATED] Prefer `detokenize_stage`. Whether to detokenize the output.",
|
||||
)
|
||||
# New nested stage configuration (bool | dict | typed config).
|
||||
chat_template_stage: Any = Field(
|
||||
default=True,
|
||||
description="Chat templating stage config (bool | dict | ChatTemplateStageConfig).",
|
||||
)
|
||||
tokenize_stage: Any = Field(
|
||||
default=True,
|
||||
description="Tokenizer stage config (bool | dict | TokenizerStageConfig).",
|
||||
)
|
||||
detokenize_stage: Any = Field(
|
||||
default=True,
|
||||
description="Detokenizer stage config (bool | dict | DetokenizeStageConfig).",
|
||||
)
|
||||
prepare_multimodal_stage: Any = Field(
|
||||
default=False,
|
||||
description="Prepare multimodal stage config (bool | dict | PrepareMultimodalStageConfig).",
|
||||
)
|
||||
|
||||
@model_validator(mode="before")
|
||||
def _coerce_legacy_to_stage_config(cls, values: Dict[str, Any]) -> Dict[str, Any]:
|
||||
# Only set stage fields if not explicitly provided.
|
||||
# Emit deprecation warnings when legacy boolean flags are used.
|
||||
|
||||
# Chat template stage: special case (handles both apply_chat_template and chat_template fields)
|
||||
if "chat_template_stage" not in values:
|
||||
if "apply_chat_template" in values or "chat_template" in values:
|
||||
logger.warning(
|
||||
"The `apply_chat_template` and `chat_template` fields are deprecated. "
|
||||
"Use `chat_template_stage` instead. For example: "
|
||||
"`chat_template_stage=ChatTemplateStageConfig(enabled=True, chat_template='...')` "
|
||||
"or `chat_template_stage={'enabled': True, 'chat_template': '...'}`. "
|
||||
"This will raise an error in a future version."
|
||||
)
|
||||
enabled_value = values.get("apply_chat_template")
|
||||
enabled = enabled_value if enabled_value is not None else True
|
||||
stage: Dict[str, Any] = {"enabled": enabled}
|
||||
if values.get("chat_template") is not None:
|
||||
stage["chat_template"] = values["chat_template"]
|
||||
values["chat_template_stage"] = stage
|
||||
|
||||
# Other stages: simple boolean-to-stage mapping
|
||||
stage_mappings = [
|
||||
("tokenize_stage", "tokenize", True, "TokenizerStageConfig"),
|
||||
("detokenize_stage", "detokenize", True, "DetokenizeStageConfig"),
|
||||
]
|
||||
for (
|
||||
stage_field,
|
||||
legacy_field,
|
||||
default_enabled,
|
||||
config_class_name,
|
||||
) in stage_mappings:
|
||||
if stage_field not in values and legacy_field in values:
|
||||
logger.warning(
|
||||
f"The `{legacy_field}` field is deprecated. "
|
||||
f"Use `{stage_field}` instead. For example: "
|
||||
f"`{stage_field}={config_class_name}(enabled=True)` "
|
||||
f"or `{stage_field}={{'enabled': True}}`. "
|
||||
"This will raise an error in a future version."
|
||||
)
|
||||
legacy_value = values.get(legacy_field)
|
||||
enabled = default_enabled if legacy_value is None else legacy_value
|
||||
values[stage_field] = {"enabled": enabled}
|
||||
|
||||
return values
|
||||
|
||||
@model_validator(mode="before")
|
||||
def _migrate_experimental_max_tasks_in_flight_per_actor(
|
||||
cls, values: Dict[str, Any]
|
||||
) -> Dict[str, Any]:
|
||||
"""Migrate deprecated `experimental[max_tasks_in_flight_per_actor]` to
|
||||
the top-level field; top-level wins if both are set."""
|
||||
experimental = values.get("experimental") or {}
|
||||
if "max_tasks_in_flight_per_actor" in experimental:
|
||||
logger.warning(
|
||||
"Setting `max_tasks_in_flight_per_actor` via `experimental` is "
|
||||
"deprecated; use the top-level `max_tasks_in_flight_per_actor` "
|
||||
"field on `OfflineProcessorConfig` instead. The value in "
|
||||
"`experimental` is still respected for now (and overridden by "
|
||||
"the top-level field if both are set), but will be removed in "
|
||||
"a future version."
|
||||
)
|
||||
if values.get("max_tasks_in_flight_per_actor") is None:
|
||||
values["max_tasks_in_flight_per_actor"] = experimental[
|
||||
"max_tasks_in_flight_per_actor"
|
||||
]
|
||||
return values
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _warn_if_max_tasks_in_flight_underutilizes_actor(self):
|
||||
if (
|
||||
self.max_tasks_in_flight_per_actor is not None
|
||||
and self.max_tasks_in_flight_per_actor < self.max_concurrent_batches
|
||||
):
|
||||
logger.warning(
|
||||
"Setting `max_tasks_in_flight_per_actor` (%s) lower than "
|
||||
"`max_concurrent_batches` (%s) can underutilize each engine "
|
||||
"actor because Ray Data will submit fewer tasks than the actor "
|
||||
"can process concurrently.",
|
||||
self.max_tasks_in_flight_per_actor,
|
||||
self.max_concurrent_batches,
|
||||
)
|
||||
return self
|
||||
|
||||
|
||||
@PublicAPI(stability="beta")
|
||||
class Processor:
|
||||
"""A processor is composed of a preprocess stage, followed by one or more
|
||||
processing stages, and finally a postprocess stage. We use processor as a
|
||||
paradigm for processing data using LLMs.
|
||||
|
||||
Args:
|
||||
config: The processor config.
|
||||
stages: List of processing stages.
|
||||
preprocess: An optional lambda function that takes a row (dict) as input
|
||||
and returns a preprocessed row (dict). The output row must contain the
|
||||
required fields for the following processing stages.
|
||||
postprocess: An optional lambda function that takes a row (dict) as input
|
||||
and returns a postprocessed row (dict).
|
||||
preprocess_map_kwargs: Optional kwargs to pass to Dataset.map() for the
|
||||
preprocess stage (e.g., num_cpus, memory, concurrency).
|
||||
postprocess_map_kwargs: Optional kwargs to pass to Dataset.map() for the
|
||||
postprocess stage (e.g., num_cpus, memory, concurrency).
|
||||
"""
|
||||
|
||||
# The internal used data column name ("__data"). Your input
|
||||
# dataset should not contain this column. If you want to use this column
|
||||
# in your input dataset, you have to derive and customize Processor.
|
||||
DATA_COLUMN: str = "__data"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: ProcessorConfig,
|
||||
stages: List[StatefulStage],
|
||||
preprocess: Optional[UserDefinedFunction] = None,
|
||||
postprocess: Optional[UserDefinedFunction] = None,
|
||||
preprocess_map_kwargs: Optional[Dict[str, Any]] = None,
|
||||
postprocess_map_kwargs: Optional[Dict[str, Any]] = None,
|
||||
):
|
||||
self.config = config
|
||||
self.preprocess = None
|
||||
self.postprocess = None
|
||||
self.preprocess_map_kwargs = preprocess_map_kwargs or {}
|
||||
self.postprocess_map_kwargs = postprocess_map_kwargs or {}
|
||||
self.stages: OrderedDict[str, StatefulStage] = OrderedDict()
|
||||
|
||||
# NOTE (Kourosh): If pre/postprocess is not provided, use the identity function.
|
||||
# Wrapping is required even if they are identity functions, b/c data_column
|
||||
# gets inserted/removed via wrap_preprocess/wrap_postprocess.
|
||||
preprocess = preprocess or (lambda row: row)
|
||||
postprocess = postprocess or (lambda row: row)
|
||||
|
||||
self.preprocess = wrap_preprocess(
|
||||
preprocess,
|
||||
self.DATA_COLUMN,
|
||||
)
|
||||
|
||||
# When should_continue_on_error is enabled, include __inference_error__ column
|
||||
# in all output rows for consistent schema (Empty string for success, message for error).
|
||||
include_error_column = getattr(config, "should_continue_on_error", False)
|
||||
self.postprocess = wrap_postprocess(
|
||||
postprocess,
|
||||
self.DATA_COLUMN,
|
||||
include_error_column=include_error_column,
|
||||
)
|
||||
|
||||
for stage in stages:
|
||||
self._append_stage(stage)
|
||||
|
||||
def __call__(self, dataset: Dataset) -> Dataset:
|
||||
"""Execute the processor:
|
||||
preprocess -> stages -> postprocess.
|
||||
Note that the dataset won't be materialized during the execution.
|
||||
|
||||
Args:
|
||||
dataset: The input dataset.
|
||||
|
||||
Returns:
|
||||
The output dataset.
|
||||
"""
|
||||
if self.preprocess is not None:
|
||||
dataset = dataset.map(self.preprocess, **self.preprocess_map_kwargs)
|
||||
|
||||
# Apply stages.
|
||||
for stage in self.stages.values():
|
||||
kwargs = stage.get_dataset_map_batches_kwargs(
|
||||
batch_size=self.config.batch_size,
|
||||
data_column=self.DATA_COLUMN,
|
||||
)
|
||||
dataset = dataset.map_batches(stage.fn, **kwargs)
|
||||
|
||||
if self.postprocess is not None:
|
||||
dataset = dataset.map(self.postprocess, **self.postprocess_map_kwargs)
|
||||
return dataset
|
||||
|
||||
def _append_stage(self, stage: StatefulStage) -> None:
|
||||
"""Append a stage before postprocess. The stage class name will be used as
|
||||
the stage name. If there are multiple stages with the same type, a suffix
|
||||
will be added to the stage name to avoid conflicts.
|
||||
|
||||
Args:
|
||||
stage: The stage to append.
|
||||
"""
|
||||
stage_name = type(stage).__name__
|
||||
|
||||
# When a processor has multiple stages with the same type,
|
||||
# append a index suffix to the stage name to avoid conflicts.
|
||||
if stage_name in self.stages:
|
||||
num_same_type_stage = len([s for s in self.stages.values() if s is stage])
|
||||
stage_name = f"{stage_name}_{num_same_type_stage + 1}"
|
||||
self.stages[stage_name] = stage
|
||||
|
||||
def list_stage_names(self) -> List[str]:
|
||||
"""List the stage names of this processor in order. Preprocess and postprocess
|
||||
are not included.
|
||||
|
||||
Returns:
|
||||
A list of stage names.
|
||||
"""
|
||||
return list(self.stages.keys())
|
||||
|
||||
def get_stage_by_name(self, name: str) -> StatefulStage:
|
||||
"""Get a particular stage by its name. If the stage is not found,
|
||||
a ValueError will be raised.
|
||||
|
||||
Args:
|
||||
name: The stage name.
|
||||
|
||||
Returns:
|
||||
The pipeline stage.
|
||||
"""
|
||||
if name in self.stages:
|
||||
return self.stages[name]
|
||||
raise ValueError(f"Stage {name} not found")
|
||||
|
||||
def log_input_column_names(self):
|
||||
"""Log.info the input stage and column names of this processor.
|
||||
If the input dataset does not contain these columns, you have to
|
||||
provide a preprocess function to bridge the gap.
|
||||
"""
|
||||
name, stage = list(self.stages.items())[0]
|
||||
expected_input_keys = stage.get_required_input_keys()
|
||||
optional_input_keys = stage.get_optional_input_keys()
|
||||
|
||||
message = f"The first stage of the processor is {name}."
|
||||
if expected_input_keys:
|
||||
message += "\nRequired input columns:\n"
|
||||
message += "\n".join(f"\t{k}: {v}" for k, v in expected_input_keys.items())
|
||||
if optional_input_keys:
|
||||
message += "\nOptional input columns:\n"
|
||||
message += "\n".join(f"\t{k}: {v}" for k, v in optional_input_keys.items())
|
||||
|
||||
logger.info(message)
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class ProcessorBuilder:
|
||||
"""Build a processor based on the configuration."""
|
||||
|
||||
_registry: Dict[str, Callable] = {}
|
||||
|
||||
@classmethod
|
||||
def register(cls, config_type: Type[ProcessorConfig], builder: Callable) -> None:
|
||||
"""A decorator to associate a particular pipeline config
|
||||
with its build function.
|
||||
"""
|
||||
type_name = config_type.__name__
|
||||
if type_name in cls._registry:
|
||||
raise ValueError(f"Processor config type {type_name} already registered.")
|
||||
cls._registry[type_name] = builder
|
||||
|
||||
@classmethod
|
||||
def clear_registry(cls) -> None:
|
||||
"""Clear the processor builder registry."""
|
||||
cls._registry.clear()
|
||||
|
||||
@classmethod
|
||||
def validate_builder_kwargs(cls, builder_kwargs: Optional[Dict[str, Any]]) -> None:
|
||||
"""Validate builder kwargs for conflicts with reserved keys.
|
||||
|
||||
Args:
|
||||
builder_kwargs: Optional additional kwargs to pass to the processor builder
|
||||
function.
|
||||
|
||||
Raises:
|
||||
ValueError: If builder_kwargs contains reserved keys that conflict with
|
||||
explicit arguments.
|
||||
"""
|
||||
if builder_kwargs is not None:
|
||||
# Check for conflicts with explicitly passed arguments
|
||||
reserved_keys = {
|
||||
"preprocess",
|
||||
"postprocess",
|
||||
"preprocess_map_kwargs",
|
||||
"postprocess_map_kwargs",
|
||||
}
|
||||
conflicting_keys = reserved_keys & builder_kwargs.keys()
|
||||
if conflicting_keys:
|
||||
raise ValueError(
|
||||
f"builder_kwargs cannot contain {conflicting_keys} as these are "
|
||||
"passed as explicit arguments to build_processor. "
|
||||
"Please pass these directly instead of in builder_kwargs."
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def build(
|
||||
cls,
|
||||
config: ProcessorConfig,
|
||||
override_stage_config_fn: Optional[Callable] = None,
|
||||
**kwargs,
|
||||
) -> Processor:
|
||||
"""Build a processor.
|
||||
|
||||
Args:
|
||||
config: The processor config.
|
||||
override_stage_config_fn: Custom stages configurations.
|
||||
**kwargs: Additional keyword arguments to pass through to the
|
||||
registered builder function. The builder function must accept
|
||||
these kwargs in its signature, otherwise a TypeError will be raised.
|
||||
|
||||
Returns:
|
||||
The built processor.
|
||||
"""
|
||||
type_name = type(config).__name__
|
||||
if type_name not in cls._registry:
|
||||
raise ValueError(
|
||||
f"Processor config type {type_name} not registered. "
|
||||
f"Available types: {cls._registry.keys()}"
|
||||
)
|
||||
processor = cls._registry[type_name](config, **kwargs)
|
||||
if override_stage_config_fn is not None:
|
||||
for name, stage in processor.stages.items():
|
||||
override_stage_config_fn(name, stage)
|
||||
return processor
|
||||
@@ -0,0 +1,146 @@
|
||||
"""The HTTP request processor."""
|
||||
|
||||
import hashlib
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
from ray.data.block import UserDefinedFunction
|
||||
from ray.llm._internal.batch.observability.usage_telemetry.usage import (
|
||||
BatchModelTelemetry,
|
||||
get_or_create_telemetry_agent,
|
||||
)
|
||||
from ray.llm._internal.batch.processor.base import (
|
||||
Processor,
|
||||
ProcessorBuilder,
|
||||
ProcessorConfig,
|
||||
)
|
||||
from ray.llm._internal.batch.processor.utils import build_cpu_stage_map_kwargs
|
||||
from ray.llm._internal.batch.stages import HttpRequestStage
|
||||
from ray.llm._internal.batch.stages.configs import (
|
||||
HttpRequestStageConfig,
|
||||
resolve_stage_config,
|
||||
)
|
||||
|
||||
|
||||
class HttpRequestProcessorConfig(ProcessorConfig):
|
||||
"""The configuration for the HTTP request processor."""
|
||||
|
||||
batch_size: int = Field(
|
||||
default=64,
|
||||
description="The batch size.",
|
||||
)
|
||||
url: str = Field(
|
||||
description="The URL to query.",
|
||||
)
|
||||
headers: Optional[Dict[str, Any]] = Field(
|
||||
default=None,
|
||||
description="The query header. Note that we will add "
|
||||
"'Content-Type: application/json' to be the header for sure "
|
||||
"because we only deal with requests body in JSON.",
|
||||
)
|
||||
qps: Optional[int] = Field(
|
||||
default=None,
|
||||
description="The maximum number of requests per second to avoid rate limit. "
|
||||
"If None, the request will be sent sequentially.",
|
||||
)
|
||||
max_retries: int = Field(
|
||||
default=0,
|
||||
description="The maximum number of retries per request in the event of failures.",
|
||||
)
|
||||
base_retry_wait_time_in_s: float = Field(
|
||||
default=1,
|
||||
description="The base wait time for a retry during exponential backoff.",
|
||||
)
|
||||
# Since `session_factory` is a callable, we use type Any to avoid pydantic serialization issues
|
||||
session_factory: Optional[Any] = Field(
|
||||
default=None,
|
||||
description="Optional session factory to be used for initializing a client session. Type: Callable[[], ClientSession]",
|
||||
# exclude from JSON serialization since `session_factory` is a callable
|
||||
exclude=True,
|
||||
)
|
||||
http_request_stage: Any = Field(
|
||||
default=True,
|
||||
description="Http request stage config (bool | dict | HttpRequestStageConfig).",
|
||||
)
|
||||
|
||||
|
||||
def build_http_request_processor(
|
||||
config: HttpRequestProcessorConfig,
|
||||
preprocess: Optional[UserDefinedFunction] = None,
|
||||
postprocess: Optional[UserDefinedFunction] = None,
|
||||
preprocess_map_kwargs: Optional[Dict[str, Any]] = None,
|
||||
postprocess_map_kwargs: Optional[Dict[str, Any]] = None,
|
||||
) -> Processor:
|
||||
"""Construct a Processor and configure stages.
|
||||
|
||||
Args:
|
||||
config: The configuration for the processor.
|
||||
preprocess: An optional lambda function that takes a row (dict) as input
|
||||
and returns a preprocessed row (dict). The output row must contain the
|
||||
required fields for the following processing stages.
|
||||
postprocess: An optional lambda function that takes a row (dict) as input
|
||||
and returns a postprocessed row (dict).
|
||||
preprocess_map_kwargs: Optional kwargs to pass to Dataset.map() for the
|
||||
preprocess stage (e.g., num_cpus, memory, concurrency).
|
||||
postprocess_map_kwargs: Optional kwargs to pass to Dataset.map() for the
|
||||
postprocess stage (e.g., num_cpus, memory, concurrency).
|
||||
|
||||
Returns:
|
||||
The constructed processor.
|
||||
"""
|
||||
|
||||
# Prepare processor defaults for merging into stage configs
|
||||
processor_defaults = {
|
||||
"batch_size": config.batch_size,
|
||||
"concurrency": config.concurrency,
|
||||
}
|
||||
|
||||
# Resolve and build HttpRequestStage if enabled
|
||||
http_request_stage_cfg = resolve_stage_config(
|
||||
config.http_request_stage,
|
||||
HttpRequestStageConfig,
|
||||
processor_defaults,
|
||||
)
|
||||
|
||||
if not http_request_stage_cfg.enabled:
|
||||
raise ValueError(
|
||||
"The HTTP request stage is required and cannot be disabled in HttpRequestProcessorConfig."
|
||||
)
|
||||
|
||||
stages = [
|
||||
HttpRequestStage(
|
||||
fn_constructor_kwargs=dict(
|
||||
url=config.url,
|
||||
additional_header=config.headers,
|
||||
qps=config.qps,
|
||||
max_retries=config.max_retries,
|
||||
base_retry_wait_time_in_s=config.base_retry_wait_time_in_s,
|
||||
session_factory=config.session_factory,
|
||||
),
|
||||
map_batches_kwargs=build_cpu_stage_map_kwargs(http_request_stage_cfg),
|
||||
)
|
||||
]
|
||||
telemetry_agent = get_or_create_telemetry_agent()
|
||||
telemetry_agent.push_telemetry_report(
|
||||
BatchModelTelemetry(
|
||||
# Hash the target URL so distinct endpoints stay separate in the
|
||||
# dedup key without the cleartext URL reaching the head-node actor.
|
||||
model_id_hash=hashlib.sha256(config.url.encode("utf-8")).hexdigest(),
|
||||
processor_config_name=type(config).__name__,
|
||||
batch_size=config.batch_size,
|
||||
concurrency=config.concurrency,
|
||||
)
|
||||
)
|
||||
processor = Processor(
|
||||
config,
|
||||
stages,
|
||||
preprocess=preprocess,
|
||||
postprocess=postprocess,
|
||||
preprocess_map_kwargs=preprocess_map_kwargs,
|
||||
postprocess_map_kwargs=postprocess_map_kwargs,
|
||||
)
|
||||
return processor
|
||||
|
||||
|
||||
ProcessorBuilder.register(HttpRequestProcessorConfig, build_http_request_processor)
|
||||
@@ -0,0 +1,107 @@
|
||||
"""The processor that runs serve deployment."""
|
||||
|
||||
from typing import Any, Dict, Optional, Type
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
from ray.data import ActorPoolStrategy
|
||||
from ray.data.block import UserDefinedFunction
|
||||
from ray.llm._internal.batch.processor.base import (
|
||||
Processor,
|
||||
ProcessorBuilder,
|
||||
ProcessorConfig,
|
||||
)
|
||||
from ray.llm._internal.batch.stages import (
|
||||
ServeDeploymentStage,
|
||||
)
|
||||
|
||||
|
||||
class ServeDeploymentProcessorConfig(ProcessorConfig):
|
||||
"""The configuration for the serve deployment processor."""
|
||||
|
||||
# Configurations used to build the serve deployment
|
||||
deployment_name: str = Field(
|
||||
description="The name of the serve deployment to use.",
|
||||
)
|
||||
app_name: str = Field(
|
||||
description="The name of the serve application to use.",
|
||||
default="default",
|
||||
)
|
||||
dtype_mapping: Dict[str, Type[Any]] = Field(
|
||||
description="A dictionary mapping data type names to their corresponding request classes for the serve deployment.",
|
||||
default=None,
|
||||
)
|
||||
should_continue_on_error: bool = Field(
|
||||
default=False,
|
||||
description="If True, continue processing when inference fails for a row "
|
||||
"instead of raising an exception. Failed rows will have a non-null "
|
||||
"'__inference_error__' column containing the error message. Error rows "
|
||||
"bypass postprocess. If False (default), any inference error raises.",
|
||||
)
|
||||
request_timeout_s: Optional[float] = Field(
|
||||
default=None,
|
||||
gt=0,
|
||||
description="Optional per-request timeout in seconds. When set, a request "
|
||||
"that does not return within this many seconds raises TimeoutError instead "
|
||||
"of blocking indefinitely (e.g. when replicas are saturated). TimeoutError "
|
||||
"is recoverable, so combine with should_continue_on_error=True to drop the "
|
||||
"slow row as an error instead of failing the job. If None (default), "
|
||||
"requests wait indefinitely.",
|
||||
)
|
||||
|
||||
|
||||
def build_serve_deployment_processor(
|
||||
config: ServeDeploymentProcessorConfig,
|
||||
preprocess: Optional[UserDefinedFunction] = None,
|
||||
postprocess: Optional[UserDefinedFunction] = None,
|
||||
preprocess_map_kwargs: Optional[Dict[str, Any]] = None,
|
||||
postprocess_map_kwargs: Optional[Dict[str, Any]] = None,
|
||||
) -> Processor:
|
||||
"""Construct a processor that runs a serve deployment.
|
||||
|
||||
Args:
|
||||
config: The configuration for the processor.
|
||||
preprocess: An optional lambda function that takes a row (dict) as input
|
||||
and returns a preprocessed row (dict). The output row must contain the
|
||||
required fields for the following processing stages.
|
||||
postprocess: An optional lambda function that takes a row (dict) as input
|
||||
and returns a postprocessed row (dict).
|
||||
preprocess_map_kwargs: Optional kwargs to pass to Dataset.map() for the
|
||||
preprocess stage (e.g., num_cpus, memory, concurrency).
|
||||
postprocess_map_kwargs: Optional kwargs to pass to Dataset.map() for the
|
||||
postprocess stage (e.g., num_cpus, memory, concurrency).
|
||||
|
||||
Returns:
|
||||
The constructed processor.
|
||||
"""
|
||||
stages = [
|
||||
ServeDeploymentStage(
|
||||
fn_constructor_kwargs=dict(
|
||||
deployment_name=config.deployment_name,
|
||||
app_name=config.app_name,
|
||||
dtype_mapping=config.dtype_mapping,
|
||||
should_continue_on_error=config.should_continue_on_error,
|
||||
request_timeout_s=config.request_timeout_s,
|
||||
),
|
||||
map_batches_kwargs=dict(
|
||||
compute=ActorPoolStrategy(
|
||||
**config.get_concurrency(autoscaling_enabled=False),
|
||||
)
|
||||
),
|
||||
)
|
||||
]
|
||||
# TODO (Kourosh): Add telemetry for ServeDeploymentStage
|
||||
processor = Processor(
|
||||
config,
|
||||
stages,
|
||||
preprocess=preprocess,
|
||||
postprocess=postprocess,
|
||||
preprocess_map_kwargs=preprocess_map_kwargs,
|
||||
postprocess_map_kwargs=postprocess_map_kwargs,
|
||||
)
|
||||
return processor
|
||||
|
||||
|
||||
ProcessorBuilder.register(
|
||||
ServeDeploymentProcessorConfig, build_serve_deployment_processor
|
||||
)
|
||||
@@ -0,0 +1,279 @@
|
||||
"""The SGLang engine processor."""
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import transformers
|
||||
from pydantic import Field, root_validator
|
||||
|
||||
import ray
|
||||
from ray.data.block import UserDefinedFunction
|
||||
from ray.llm._internal.batch.constants import SGLangTaskType, TypeSGLangTaskType
|
||||
from ray.llm._internal.batch.observability.usage_telemetry.usage import (
|
||||
BatchModelTelemetry,
|
||||
TelemetryAgent,
|
||||
get_or_create_telemetry_agent,
|
||||
)
|
||||
from ray.llm._internal.batch.processor.base import (
|
||||
OfflineProcessorConfig,
|
||||
Processor,
|
||||
ProcessorBuilder,
|
||||
)
|
||||
from ray.llm._internal.batch.processor.utils import (
|
||||
build_cpu_stage_map_kwargs,
|
||||
get_value_or_fallback,
|
||||
)
|
||||
from ray.llm._internal.batch.stages import (
|
||||
ChatTemplateStage,
|
||||
DetokenizeStage,
|
||||
SGLangEngineStage,
|
||||
TokenizeStage,
|
||||
)
|
||||
from ray.llm._internal.batch.stages.configs import (
|
||||
ChatTemplateStageConfig,
|
||||
DetokenizeStageConfig,
|
||||
TokenizerStageConfig,
|
||||
resolve_stage_config,
|
||||
)
|
||||
from ray.llm._internal.common.observability.telemetry_utils import DEFAULT_GPU_TYPE
|
||||
from ray.llm._internal.common.utils.download_utils import (
|
||||
NodeModelDownloadable,
|
||||
download_model_files,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
DEFAULT_MODEL_ARCHITECTURE = "UNKNOWN_MODEL_ARCHITECTURE"
|
||||
|
||||
|
||||
class SGLangEngineProcessorConfig(OfflineProcessorConfig):
|
||||
"""The configuration for the SGLang engine processor."""
|
||||
|
||||
# SGLang stage configurations.
|
||||
engine_kwargs: Dict[str, Any] = Field(
|
||||
default_factory=dict,
|
||||
description="The kwargs to pass to the SGLang engine. See "
|
||||
"https://docs.sglang.ai/backend/server_arguments.html "
|
||||
"for more details.",
|
||||
)
|
||||
task_type: TypeSGLangTaskType = Field(
|
||||
default=SGLangTaskType.GENERATE,
|
||||
description="The task type to use. If not specified, will use "
|
||||
"'generate' by default.",
|
||||
)
|
||||
|
||||
@root_validator(pre=True)
|
||||
def validate_task_type(cls, values):
|
||||
task_type = values.get("task_type", SGLangTaskType.GENERATE)
|
||||
if task_type not in SGLangTaskType.values():
|
||||
raise ValueError(f"Invalid task type: {task_type}")
|
||||
|
||||
engine_kwargs = values.get("engine_kwargs", {})
|
||||
engine_kwargs_task = engine_kwargs.get("task", "")
|
||||
if engine_kwargs_task != task_type:
|
||||
logger.warning(
|
||||
"The task set in engine kwargs (%s) is different from the "
|
||||
"stage (%s). Overriding the task in engine kwargs to %s.",
|
||||
engine_kwargs_task,
|
||||
task_type,
|
||||
task_type,
|
||||
)
|
||||
engine_kwargs["task"] = task_type
|
||||
values["engine_kwargs"] = engine_kwargs
|
||||
return values
|
||||
|
||||
|
||||
def build_sglang_engine_processor(
|
||||
config: SGLangEngineProcessorConfig,
|
||||
chat_template_kwargs: Optional[Dict[str, Any]] = None,
|
||||
preprocess: Optional[UserDefinedFunction] = None,
|
||||
postprocess: Optional[UserDefinedFunction] = None,
|
||||
preprocess_map_kwargs: Optional[Dict[str, Any]] = None,
|
||||
postprocess_map_kwargs: Optional[Dict[str, Any]] = None,
|
||||
telemetry_agent: Optional[TelemetryAgent] = None,
|
||||
) -> Processor:
|
||||
"""Construct a Processor and configure stages.
|
||||
|
||||
Args:
|
||||
config: The configuration for the processor.
|
||||
chat_template_kwargs: The optional kwargs to pass to apply_chat_template.
|
||||
preprocess: An optional lambda function that takes a row (dict) as input
|
||||
and returns a preprocessed row (dict). The output row must contain the
|
||||
required fields for the following processing stages.
|
||||
postprocess: An optional lambda function that takes a row (dict) as input
|
||||
and returns a postprocessed row (dict).
|
||||
preprocess_map_kwargs: Optional kwargs to pass to Dataset.map() for the
|
||||
preprocess stage (e.g., num_cpus, memory, concurrency).
|
||||
postprocess_map_kwargs: Optional kwargs to pass to Dataset.map() for the
|
||||
postprocess stage (e.g., num_cpus, memory, concurrency).
|
||||
telemetry_agent: An optional telemetry agent for collecting usage telemetry.
|
||||
|
||||
Returns:
|
||||
The constructed processor.
|
||||
"""
|
||||
ray.init(runtime_env=config.runtime_env, ignore_reinit_error=True)
|
||||
|
||||
stages = []
|
||||
|
||||
# Prepare processor defaults for merging into stage configs
|
||||
trust_remote_code = config.engine_kwargs.get("trust_remote_code", False)
|
||||
processor_defaults = {
|
||||
"batch_size": config.batch_size,
|
||||
"concurrency": config.concurrency,
|
||||
"runtime_env": config.runtime_env,
|
||||
"model_source": config.model_source,
|
||||
}
|
||||
|
||||
# Resolve and build ChatTemplateStage if enabled
|
||||
chat_template_stage_cfg = resolve_stage_config(
|
||||
config.chat_template_stage,
|
||||
ChatTemplateStageConfig,
|
||||
processor_defaults,
|
||||
)
|
||||
if chat_template_stage_cfg.enabled:
|
||||
stages.append(
|
||||
ChatTemplateStage(
|
||||
fn_constructor_kwargs=dict(
|
||||
model=chat_template_stage_cfg.model_source,
|
||||
chat_template=get_value_or_fallback(
|
||||
chat_template_stage_cfg.chat_template, config.chat_template
|
||||
),
|
||||
chat_template_kwargs=get_value_or_fallback(
|
||||
chat_template_stage_cfg.chat_template_kwargs,
|
||||
chat_template_kwargs,
|
||||
),
|
||||
trust_remote_code=trust_remote_code,
|
||||
),
|
||||
map_batches_kwargs=build_cpu_stage_map_kwargs(chat_template_stage_cfg),
|
||||
)
|
||||
)
|
||||
|
||||
# Resolve and build TokenizeStage if enabled
|
||||
tokenize_stage_cfg = resolve_stage_config(
|
||||
getattr(config, "tokenize_stage", config.tokenize),
|
||||
TokenizerStageConfig,
|
||||
processor_defaults,
|
||||
)
|
||||
if tokenize_stage_cfg.enabled:
|
||||
stages.append(
|
||||
TokenizeStage(
|
||||
fn_constructor_kwargs=dict(
|
||||
model=tokenize_stage_cfg.model_source,
|
||||
trust_remote_code=trust_remote_code,
|
||||
),
|
||||
map_batches_kwargs=build_cpu_stage_map_kwargs(tokenize_stage_cfg),
|
||||
)
|
||||
)
|
||||
|
||||
# Core stage -- the SGLang engine.
|
||||
stages.append(
|
||||
SGLangEngineStage(
|
||||
fn_constructor_kwargs=dict(
|
||||
model=config.model_source,
|
||||
engine_kwargs=config.engine_kwargs,
|
||||
task_type=config.task_type,
|
||||
max_pending_requests=config.max_pending_requests,
|
||||
),
|
||||
map_batches_kwargs=dict(
|
||||
zero_copy_batch=True,
|
||||
# The number of running replicas. This is a deprecated field, but
|
||||
# we need to set `max_tasks_in_flight_per_actor` through `compute`,
|
||||
# which initiates enough many overlapping UDF calls per actor, to
|
||||
# saturate `max_concurrency`.
|
||||
compute=ray.data.ActorPoolStrategy(
|
||||
**config.get_concurrency(autoscaling_enabled=True),
|
||||
max_tasks_in_flight_per_actor=config.max_tasks_in_flight_per_actor,
|
||||
),
|
||||
# The number of running batches "per actor" in Ray Core level.
|
||||
# This is used to make sure we overlap batches to avoid the tail
|
||||
# latency of each batch.
|
||||
max_concurrency=config.max_concurrent_batches,
|
||||
accelerator_type=config.accelerator_type,
|
||||
runtime_env=config.runtime_env,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
# Resolve and build DetokenizeStage if enabled
|
||||
detokenize_stage_cfg = resolve_stage_config(
|
||||
getattr(config, "detokenize_stage", config.detokenize),
|
||||
DetokenizeStageConfig,
|
||||
processor_defaults,
|
||||
)
|
||||
if detokenize_stage_cfg.enabled:
|
||||
stages.append(
|
||||
DetokenizeStage(
|
||||
fn_constructor_kwargs=dict(
|
||||
model=detokenize_stage_cfg.model_source,
|
||||
trust_remote_code=trust_remote_code,
|
||||
),
|
||||
map_batches_kwargs=build_cpu_stage_map_kwargs(detokenize_stage_cfg),
|
||||
)
|
||||
)
|
||||
|
||||
# Download model files for telemetry before engine init.
|
||||
# Use EXCLUDE_SAFETENSORS for trust_remote_code models so custom .py config
|
||||
# files are available locally.
|
||||
try:
|
||||
download_mode = (
|
||||
NodeModelDownloadable.EXCLUDE_SAFETENSORS
|
||||
if trust_remote_code
|
||||
else NodeModelDownloadable.TOKENIZER_ONLY
|
||||
)
|
||||
model_path_or_id = download_model_files(
|
||||
model_id=config.model_source,
|
||||
mirror_config=None,
|
||||
download_model=download_mode,
|
||||
download_extra_files=False,
|
||||
)
|
||||
|
||||
hf_config = transformers.AutoConfig.from_pretrained(
|
||||
model_path_or_id,
|
||||
trust_remote_code=trust_remote_code,
|
||||
)
|
||||
except Exception as e:
|
||||
# Failed to retrieve HuggingFace config for telemetry purposes.
|
||||
# This is non-fatal: we fall back to DEFAULT_MODEL_ARCHITECTURE for telemetry.
|
||||
# The actual model loading happens later in SGLang, which may support models
|
||||
# that aren't available via HuggingFace's AutoConfig.
|
||||
logger.warning(
|
||||
"Failed to retrieve HuggingFace config for %s: %s",
|
||||
config.model_source,
|
||||
e,
|
||||
)
|
||||
hf_config = None
|
||||
|
||||
architectures = getattr(hf_config, "architectures", [])
|
||||
architecture = architectures[0] if architectures else DEFAULT_MODEL_ARCHITECTURE
|
||||
|
||||
telemetry_agent = get_or_create_telemetry_agent()
|
||||
telemetry_agent.push_telemetry_report(
|
||||
BatchModelTelemetry(
|
||||
model_id_hash=hashlib.sha256(
|
||||
config.model_source.encode("utf-8")
|
||||
).hexdigest(),
|
||||
processor_config_name=type(config).__name__,
|
||||
model_architecture=architecture,
|
||||
batch_size=config.batch_size,
|
||||
accelerator_type=config.accelerator_type or DEFAULT_GPU_TYPE,
|
||||
concurrency=config.concurrency,
|
||||
task_type=config.task_type,
|
||||
tensor_parallel_size=config.engine_kwargs.get("tp_size", 1),
|
||||
data_parallel_size=config.engine_kwargs.get("dp_size", 1),
|
||||
)
|
||||
)
|
||||
|
||||
processor = Processor(
|
||||
config,
|
||||
stages,
|
||||
preprocess=preprocess,
|
||||
postprocess=postprocess,
|
||||
preprocess_map_kwargs=preprocess_map_kwargs,
|
||||
postprocess_map_kwargs=postprocess_map_kwargs,
|
||||
)
|
||||
return processor
|
||||
|
||||
|
||||
ProcessorBuilder.register(SGLangEngineProcessorConfig, build_sglang_engine_processor)
|
||||
@@ -0,0 +1,58 @@
|
||||
"""Shared utility functions for processor builders."""
|
||||
|
||||
from typing import Any, Dict, Optional, Tuple, Union
|
||||
|
||||
from ray.data import ActorPoolStrategy
|
||||
from ray.llm._internal.batch.stages.configs import _StageConfigBase
|
||||
|
||||
|
||||
def get_value_or_fallback(value: Any, fallback: Any) -> Any:
|
||||
"""Return value if not None, otherwise return fallback."""
|
||||
return value if value is not None else fallback
|
||||
|
||||
|
||||
def extract_resource_kwargs(
|
||||
runtime_env: Optional[Dict[str, Any]],
|
||||
num_cpus: Optional[float],
|
||||
memory: Optional[float],
|
||||
) -> Dict[str, Any]:
|
||||
"""Extract non-None resource kwargs for map_batches."""
|
||||
kwargs = {}
|
||||
if runtime_env is not None:
|
||||
kwargs["runtime_env"] = runtime_env
|
||||
if num_cpus is not None:
|
||||
kwargs["num_cpus"] = num_cpus
|
||||
if memory is not None:
|
||||
kwargs["memory"] = memory
|
||||
return kwargs
|
||||
|
||||
|
||||
def normalize_cpu_stage_concurrency(
|
||||
concurrency: Optional[Union[int, Tuple[int, int]]]
|
||||
) -> Dict[str, int]:
|
||||
"""Normalize concurrency for CPU stages (int -> (1, int) for autoscaling)."""
|
||||
if concurrency is None:
|
||||
return {"size": 1} # Default to minimal autoscaling pool
|
||||
if isinstance(concurrency, int):
|
||||
return {"min_size": 1, "max_size": concurrency}
|
||||
return {
|
||||
"min_size": concurrency[0],
|
||||
"max_size": concurrency[1],
|
||||
}
|
||||
|
||||
|
||||
def build_cpu_stage_map_kwargs(
|
||||
stage_cfg: _StageConfigBase,
|
||||
) -> Dict[str, Any]:
|
||||
"""Build map_batches_kwargs for CPU stages."""
|
||||
concurrency = normalize_cpu_stage_concurrency(stage_cfg.concurrency)
|
||||
return dict(
|
||||
zero_copy_batch=True,
|
||||
compute=ActorPoolStrategy(**concurrency),
|
||||
batch_size=stage_cfg.batch_size,
|
||||
**extract_resource_kwargs(
|
||||
stage_cfg.runtime_env,
|
||||
stage_cfg.num_cpus,
|
||||
stage_cfg.memory,
|
||||
),
|
||||
)
|
||||
@@ -0,0 +1,359 @@
|
||||
"""The vLLM engine processor."""
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import transformers
|
||||
from pydantic import Field, field_validator, model_validator
|
||||
|
||||
import ray
|
||||
from ray.data.block import UserDefinedFunction
|
||||
from ray.llm._internal.batch.constants import TypeVLLMTaskType, vLLMTaskType
|
||||
from ray.llm._internal.batch.observability.usage_telemetry.usage import (
|
||||
BatchModelTelemetry,
|
||||
TelemetryAgent,
|
||||
get_or_create_telemetry_agent,
|
||||
)
|
||||
from ray.llm._internal.batch.processor.base import (
|
||||
OfflineProcessorConfig,
|
||||
Processor,
|
||||
ProcessorBuilder,
|
||||
)
|
||||
from ray.llm._internal.batch.processor.utils import (
|
||||
build_cpu_stage_map_kwargs,
|
||||
get_value_or_fallback,
|
||||
)
|
||||
from ray.llm._internal.batch.stages import (
|
||||
ChatTemplateStage,
|
||||
DetokenizeStage,
|
||||
PrepareMultimodalStage,
|
||||
TokenizeStage,
|
||||
vLLMEngineStage,
|
||||
)
|
||||
from ray.llm._internal.batch.stages.configs import (
|
||||
ChatTemplateStageConfig,
|
||||
DetokenizeStageConfig,
|
||||
PrepareMultimodalStageConfig,
|
||||
TokenizerStageConfig,
|
||||
resolve_stage_config,
|
||||
)
|
||||
from ray.llm._internal.common.observability.telemetry_utils import DEFAULT_GPU_TYPE
|
||||
from ray.llm._internal.common.placement import PlacementGroupConfig
|
||||
from ray.llm._internal.common.utils.download_utils import (
|
||||
STREAMING_LOAD_FORMATS,
|
||||
NodeModelDownloadable,
|
||||
download_model_files,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
DEFAULT_MODEL_ARCHITECTURE = "UNKNOWN_MODEL_ARCHITECTURE"
|
||||
|
||||
|
||||
class vLLMEngineProcessorConfig(OfflineProcessorConfig):
|
||||
"""The configuration for the vLLM engine processor."""
|
||||
|
||||
# vLLM stage configurations.
|
||||
engine_kwargs: Dict[str, Any] = Field(
|
||||
default_factory=dict,
|
||||
description="The kwargs to pass to the vLLM engine. See "
|
||||
"https://docs.vllm.ai/en/latest/serving/engine_args.html "
|
||||
"for more details.",
|
||||
)
|
||||
task_type: TypeVLLMTaskType = Field(
|
||||
default=vLLMTaskType.GENERATE,
|
||||
description="The task type to use. If not specified, will use "
|
||||
"'generate' by default.",
|
||||
)
|
||||
log_engine_metrics: bool = Field(
|
||||
default=True,
|
||||
description="Enable vLLM engine metrics export via Ray's Prometheus endpoint. "
|
||||
"When enabled, metrics like prefix cache hit rate, TTFT, TPOT, KV cache "
|
||||
"utilization, and scheduler state are available at Ray's metrics endpoint. "
|
||||
"Requires Ray to be initialized with _metrics_export_port "
|
||||
"(e.g., ray.init(_metrics_export_port=8080)).",
|
||||
)
|
||||
# LoRA configurations.
|
||||
dynamic_lora_loading_path: Optional[str] = Field(
|
||||
default=None,
|
||||
description="The path to the dynamic LoRA adapter. It is expected "
|
||||
"to hold subfolders each for a different lora checkpoint. If not "
|
||||
"specified and LoRA is enabled, then the 'model' in LoRA "
|
||||
"requests will be interpreted as model ID used by HF transformers.",
|
||||
)
|
||||
# Custom placement group config for TP/PP.
|
||||
placement_group_config: Optional[Dict[str, Any]] = Field(
|
||||
default=None,
|
||||
description="Ray placement group configuration for scheduling vLLM engine workers. "
|
||||
"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'}.",
|
||||
)
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def validate_task_type(cls, values):
|
||||
task_type = values.get("task_type", vLLMTaskType.GENERATE)
|
||||
if task_type not in vLLMTaskType.values():
|
||||
raise ValueError(f"Invalid task type: {task_type}")
|
||||
|
||||
engine_kwargs = values.get("engine_kwargs", {})
|
||||
engine_kwargs_task_type = engine_kwargs.get("task_type", "")
|
||||
if engine_kwargs_task_type != task_type:
|
||||
if engine_kwargs_task_type:
|
||||
logger.warning(
|
||||
"The task_type set in engine kwargs (%s) is different from the "
|
||||
"config (%s). Overriding the task_type in engine kwargs to %s.",
|
||||
engine_kwargs_task_type,
|
||||
task_type,
|
||||
task_type,
|
||||
)
|
||||
engine_kwargs["task_type"] = task_type
|
||||
values["engine_kwargs"] = engine_kwargs
|
||||
return values
|
||||
|
||||
@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()
|
||||
|
||||
|
||||
def build_vllm_engine_processor(
|
||||
config: vLLMEngineProcessorConfig,
|
||||
chat_template_kwargs: Optional[Dict[str, Any]] = None,
|
||||
preprocess: Optional[UserDefinedFunction] = None,
|
||||
postprocess: Optional[UserDefinedFunction] = None,
|
||||
preprocess_map_kwargs: Optional[Dict[str, Any]] = None,
|
||||
postprocess_map_kwargs: Optional[Dict[str, Any]] = None,
|
||||
telemetry_agent: Optional[TelemetryAgent] = None,
|
||||
) -> Processor:
|
||||
"""Construct a Processor and configure stages.
|
||||
|
||||
Args:
|
||||
config: The configuration for the processor.
|
||||
chat_template_kwargs: The optional kwargs to pass to apply_chat_template.
|
||||
preprocess: An optional lambda function that takes a row (dict) as input
|
||||
and returns a preprocessed row (dict). The output row must contain the
|
||||
required fields for the following processing stages.
|
||||
postprocess: An optional lambda function that takes a row (dict) as input
|
||||
and returns a postprocessed row (dict).
|
||||
preprocess_map_kwargs: Optional kwargs to pass to Dataset.map() for the
|
||||
preprocess stage (e.g., num_cpus, memory, concurrency).
|
||||
postprocess_map_kwargs: Optional kwargs to pass to Dataset.map() for the
|
||||
postprocess stage (e.g., num_cpus, memory, concurrency).
|
||||
telemetry_agent: An optional telemetry agent for collecting usage telemetry.
|
||||
|
||||
Returns:
|
||||
The constructed processor.
|
||||
"""
|
||||
ray.init(runtime_env=config.runtime_env, ignore_reinit_error=True)
|
||||
|
||||
stages = []
|
||||
|
||||
# Prepare processor defaults for merging into stage configs
|
||||
trust_remote_code = config.engine_kwargs.get("trust_remote_code", False)
|
||||
processor_defaults = {
|
||||
"batch_size": config.batch_size,
|
||||
"concurrency": config.concurrency,
|
||||
"runtime_env": config.runtime_env,
|
||||
"model_source": config.model_source,
|
||||
}
|
||||
|
||||
# Resolve and build PrepareMultimodalStage if enabled.
|
||||
prepare_multimodal_stage_cfg = resolve_stage_config(
|
||||
config.prepare_multimodal_stage,
|
||||
PrepareMultimodalStageConfig,
|
||||
processor_defaults,
|
||||
)
|
||||
|
||||
if prepare_multimodal_stage_cfg.enabled:
|
||||
base_model_config_kwargs = (
|
||||
prepare_multimodal_stage_cfg.model_config_kwargs or {}
|
||||
)
|
||||
# Respect the model source from the processor
|
||||
model_config_kwargs = {
|
||||
**base_model_config_kwargs,
|
||||
"model": processor_defaults.get("model_source"),
|
||||
}
|
||||
stages.append(
|
||||
PrepareMultimodalStage(
|
||||
fn_constructor_kwargs=dict(
|
||||
model_config_kwargs=model_config_kwargs,
|
||||
chat_template_content_format=prepare_multimodal_stage_cfg.chat_template_content_format,
|
||||
apply_sys_msg_formatting=prepare_multimodal_stage_cfg.apply_sys_msg_formatting,
|
||||
),
|
||||
map_batches_kwargs=build_cpu_stage_map_kwargs(
|
||||
prepare_multimodal_stage_cfg
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
# Resolve and build ChatTemplateStage if enabled
|
||||
chat_template_stage_cfg = resolve_stage_config(
|
||||
getattr(config, "chat_template_stage", config.apply_chat_template),
|
||||
ChatTemplateStageConfig,
|
||||
processor_defaults,
|
||||
)
|
||||
if chat_template_stage_cfg.enabled:
|
||||
stages.append(
|
||||
ChatTemplateStage(
|
||||
fn_constructor_kwargs=dict(
|
||||
model=chat_template_stage_cfg.model_source,
|
||||
chat_template=get_value_or_fallback(
|
||||
chat_template_stage_cfg.chat_template, config.chat_template
|
||||
),
|
||||
chat_template_kwargs=get_value_or_fallback(
|
||||
chat_template_stage_cfg.chat_template_kwargs,
|
||||
chat_template_kwargs,
|
||||
),
|
||||
trust_remote_code=trust_remote_code,
|
||||
),
|
||||
map_batches_kwargs=build_cpu_stage_map_kwargs(chat_template_stage_cfg),
|
||||
)
|
||||
)
|
||||
|
||||
# Resolve and build TokenizeStage if enabled
|
||||
tokenize_stage_cfg = resolve_stage_config(
|
||||
getattr(config, "tokenize_stage", config.tokenize),
|
||||
TokenizerStageConfig,
|
||||
processor_defaults,
|
||||
)
|
||||
if tokenize_stage_cfg.enabled:
|
||||
stages.append(
|
||||
TokenizeStage(
|
||||
fn_constructor_kwargs=dict(
|
||||
model=tokenize_stage_cfg.model_source,
|
||||
trust_remote_code=trust_remote_code,
|
||||
),
|
||||
map_batches_kwargs=build_cpu_stage_map_kwargs(tokenize_stage_cfg),
|
||||
)
|
||||
)
|
||||
|
||||
# Core stage -- the vLLM engine.
|
||||
|
||||
stages.append(
|
||||
vLLMEngineStage(
|
||||
fn_constructor_kwargs=dict(
|
||||
batch_size=config.batch_size,
|
||||
max_concurrent_batches=config.max_concurrent_batches,
|
||||
model=config.model_source,
|
||||
engine_kwargs=config.engine_kwargs,
|
||||
task_type=config.task_type,
|
||||
max_pending_requests=config.max_pending_requests,
|
||||
dynamic_lora_loading_path=config.dynamic_lora_loading_path,
|
||||
placement_group_config=config.placement_group_config,
|
||||
should_continue_on_error=config.should_continue_on_error,
|
||||
log_engine_metrics=config.log_engine_metrics,
|
||||
),
|
||||
map_batches_kwargs=dict(
|
||||
zero_copy_batch=True,
|
||||
# The number of running replicas. This is a deprecated field, but
|
||||
# we need to set `max_tasks_in_flight_per_actor` through `compute`,
|
||||
# which initiates enough many overlapping UDF calls per actor, to
|
||||
# saturate `max_concurrency`.
|
||||
compute=ray.data.ActorPoolStrategy(
|
||||
**config.get_concurrency(autoscaling_enabled=True),
|
||||
max_tasks_in_flight_per_actor=config.max_tasks_in_flight_per_actor,
|
||||
),
|
||||
# The number of running batches "per actor" in Ray Core level.
|
||||
# This is used to make sure we overlap batches to avoid the tail
|
||||
# latency of each batch.
|
||||
max_concurrency=config.max_concurrent_batches,
|
||||
accelerator_type=config.accelerator_type,
|
||||
runtime_env=config.runtime_env,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
# Resolve and build DetokenizeStage if enabled
|
||||
detokenize_stage_cfg = resolve_stage_config(
|
||||
getattr(config, "detokenize_stage", config.detokenize),
|
||||
DetokenizeStageConfig,
|
||||
processor_defaults,
|
||||
)
|
||||
if detokenize_stage_cfg.enabled:
|
||||
stages.append(
|
||||
DetokenizeStage(
|
||||
fn_constructor_kwargs=dict(
|
||||
model=detokenize_stage_cfg.model_source,
|
||||
trust_remote_code=trust_remote_code,
|
||||
),
|
||||
map_batches_kwargs=build_cpu_stage_map_kwargs(detokenize_stage_cfg),
|
||||
)
|
||||
)
|
||||
|
||||
# We download the config files here so that we can report the underlying architecture to the telemetry system.
|
||||
# This should be a lightweight operation.
|
||||
# Use EXCLUDE_SAFETENSORS for streaming formats or trust_remote_code models,
|
||||
# since custom model architectures require Python config files to be downloaded.
|
||||
if config.engine_kwargs.get(
|
||||
"load_format", None
|
||||
) in STREAMING_LOAD_FORMATS or config.engine_kwargs.get("trust_remote_code", False):
|
||||
download_model_mode = NodeModelDownloadable.EXCLUDE_SAFETENSORS
|
||||
else:
|
||||
download_model_mode = NodeModelDownloadable.TOKENIZER_ONLY
|
||||
model_path = download_model_files(
|
||||
model_id=config.model_source,
|
||||
mirror_config=None,
|
||||
download_model=download_model_mode,
|
||||
download_extra_files=False,
|
||||
)
|
||||
|
||||
try:
|
||||
hf_config = transformers.AutoConfig.from_pretrained(
|
||||
model_path,
|
||||
trust_remote_code=config.engine_kwargs.get("trust_remote_code", False),
|
||||
)
|
||||
except Exception:
|
||||
# Failed to retrieve HuggingFace config for telemetry purposes.
|
||||
# This is non-fatal: we fall back to DEFAULT_MODEL_ARCHITECTURE for telemetry.
|
||||
# The actual model loading happens later in vLLM, which may support models
|
||||
# that aren't available via HuggingFace's AutoConfig.
|
||||
logger.warning(
|
||||
f"Failed to retrieve HuggingFace config for {config.model_source}"
|
||||
)
|
||||
hf_config = None
|
||||
|
||||
architectures = getattr(hf_config, "architectures", [])
|
||||
architecture = architectures[0] if architectures else DEFAULT_MODEL_ARCHITECTURE
|
||||
|
||||
telemetry_agent = get_or_create_telemetry_agent()
|
||||
telemetry_agent.push_telemetry_report(
|
||||
BatchModelTelemetry(
|
||||
model_id_hash=hashlib.sha256(
|
||||
config.model_source.encode("utf-8")
|
||||
).hexdigest(),
|
||||
processor_config_name=type(config).__name__,
|
||||
model_architecture=architecture,
|
||||
batch_size=config.batch_size,
|
||||
accelerator_type=config.accelerator_type or DEFAULT_GPU_TYPE,
|
||||
concurrency=config.concurrency,
|
||||
task_type=config.task_type,
|
||||
pipeline_parallel_size=config.engine_kwargs.get(
|
||||
"pipeline_parallel_size", 1
|
||||
),
|
||||
tensor_parallel_size=config.engine_kwargs.get("tensor_parallel_size", 1),
|
||||
data_parallel_size=config.engine_kwargs.get("data_parallel_size", 1),
|
||||
)
|
||||
)
|
||||
|
||||
processor = Processor(
|
||||
config,
|
||||
stages,
|
||||
preprocess=preprocess,
|
||||
postprocess=postprocess,
|
||||
preprocess_map_kwargs=preprocess_map_kwargs,
|
||||
postprocess_map_kwargs=postprocess_map_kwargs,
|
||||
)
|
||||
return processor
|
||||
|
||||
|
||||
ProcessorBuilder.register(vLLMEngineProcessorConfig, build_vllm_engine_processor)
|
||||
@@ -0,0 +1,108 @@
|
||||
"""Lazy re-exports for batch stages.
|
||||
|
||||
The stage modules transitively import heavy ML dependencies (transformers,
|
||||
vllm, sglang, mistral_common, etc.). Eagerly importing them all here causes
|
||||
even lightweight processors (e.g. ``HttpRequestProcessor``) to pay the cost of
|
||||
loading the entire ML stack, which is both slow and pulls in optional
|
||||
dependencies that may not be installed.
|
||||
|
||||
We use PEP 562 ``__getattr__`` to defer those imports until first access so
|
||||
that ``from ray.llm._internal.batch.stages import X`` only loads the submodule
|
||||
that defines ``X``.
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from ray.llm._internal.batch.stages.base import (
|
||||
StatefulStage,
|
||||
wrap_postprocess,
|
||||
wrap_preprocess,
|
||||
)
|
||||
from ray.llm._internal.batch.stages.configs import (
|
||||
ChatTemplateStageConfig,
|
||||
DetokenizeStageConfig,
|
||||
PrepareMultimodalStageConfig,
|
||||
TokenizerStageConfig,
|
||||
)
|
||||
|
||||
# Mapping of public attribute name -> (submodule, attribute name in submodule).
|
||||
# Each entry here is a stage class that is expensive to import (because the
|
||||
# defining submodule pulls in transformers / vllm / sglang / mistral_common /
|
||||
# pillow / etc.). They are loaded on first access via ``__getattr__`` below.
|
||||
_LAZY_ATTRS = {
|
||||
"ChatTemplateStage": ("chat_template_stage", "ChatTemplateStage"),
|
||||
"HttpRequestStage": ("http_request_stage", "HttpRequestStage"),
|
||||
"PrepareMultimodalStage": ("prepare_multimodal_stage", "PrepareMultimodalStage"),
|
||||
"ServeDeploymentStage": ("serve_deployment_stage", "ServeDeploymentStage"),
|
||||
"SGLangEngineStage": ("sglang_engine_stage", "SGLangEngineStage"),
|
||||
"TokenizeStage": ("tokenize_stage", "TokenizeStage"),
|
||||
"DetokenizeStage": ("tokenize_stage", "DetokenizeStage"),
|
||||
"vLLMEngineStage": ("vllm_engine_stage", "vLLMEngineStage"),
|
||||
}
|
||||
|
||||
|
||||
def __getattr__(name):
|
||||
"""Lazily import heavy stage classes on first access (PEP 562)."""
|
||||
try:
|
||||
submodule, attr = _LAZY_ATTRS[name]
|
||||
except KeyError:
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None
|
||||
|
||||
import importlib
|
||||
|
||||
module = importlib.import_module(f"{__name__}.{submodule}")
|
||||
value = getattr(module, attr)
|
||||
# Cache on the package so subsequent lookups skip __getattr__.
|
||||
globals()[name] = value
|
||||
return value
|
||||
|
||||
|
||||
def __dir__():
|
||||
return sorted(set(globals()).union(_LAZY_ATTRS))
|
||||
|
||||
|
||||
# Help static type checkers (and IDEs) see the lazily-exported names. These
|
||||
# imports are never executed at runtime, so they do not reintroduce the heavy
|
||||
# dependencies.
|
||||
if TYPE_CHECKING:
|
||||
from ray.llm._internal.batch.stages.chat_template_stage import ( # noqa: F401
|
||||
ChatTemplateStage,
|
||||
)
|
||||
from ray.llm._internal.batch.stages.http_request_stage import ( # noqa: F401
|
||||
HttpRequestStage,
|
||||
)
|
||||
from ray.llm._internal.batch.stages.prepare_multimodal_stage import ( # noqa: F401
|
||||
PrepareMultimodalStage,
|
||||
)
|
||||
from ray.llm._internal.batch.stages.serve_deployment_stage import ( # noqa: F401
|
||||
ServeDeploymentStage,
|
||||
)
|
||||
from ray.llm._internal.batch.stages.sglang_engine_stage import ( # noqa: F401
|
||||
SGLangEngineStage,
|
||||
)
|
||||
from ray.llm._internal.batch.stages.tokenize_stage import ( # noqa: F401
|
||||
DetokenizeStage,
|
||||
TokenizeStage,
|
||||
)
|
||||
from ray.llm._internal.batch.stages.vllm_engine_stage import ( # noqa: F401
|
||||
vLLMEngineStage,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"StatefulStage",
|
||||
"HttpRequestStage",
|
||||
"PrepareMultimodalStage",
|
||||
"ChatTemplateStage",
|
||||
"TokenizeStage",
|
||||
"DetokenizeStage",
|
||||
"vLLMEngineStage",
|
||||
"SGLangEngineStage",
|
||||
"ServeDeploymentStage",
|
||||
"wrap_preprocess",
|
||||
"wrap_postprocess",
|
||||
"ChatTemplateStageConfig",
|
||||
"DetokenizeStageConfig",
|
||||
"PrepareMultimodalStageConfig",
|
||||
"TokenizerStageConfig",
|
||||
]
|
||||
@@ -0,0 +1,327 @@
|
||||
"""The base class for all stages."""
|
||||
import logging
|
||||
from typing import Any, AsyncIterator, Callable, Dict, List, Optional, Type
|
||||
|
||||
import pyarrow
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from ray.data.block import UserDefinedFunction
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def wrap_preprocess(
|
||||
fn: UserDefinedFunction,
|
||||
processor_data_column: str,
|
||||
) -> Callable:
|
||||
"""Wrap the preprocess function, so that the output schema of the
|
||||
preprocess is normalized to {processor_data_column: fn(row), other input columns}.
|
||||
|
||||
Args:
|
||||
fn: The function to be applied.
|
||||
processor_data_column: The internal data column name of the processor.
|
||||
|
||||
Returns:
|
||||
The wrapped function.
|
||||
"""
|
||||
|
||||
def _preprocess(row: dict[str, Any]) -> dict[str, Any]:
|
||||
# First put everything into processor_data_column.
|
||||
outputs = {processor_data_column: row}
|
||||
|
||||
# Then apply the preprocess function and add its outputs.
|
||||
preprocess_output = fn(row)
|
||||
outputs[processor_data_column].update(preprocess_output)
|
||||
return outputs
|
||||
|
||||
return _preprocess
|
||||
|
||||
|
||||
def wrap_postprocess(
|
||||
fn: UserDefinedFunction,
|
||||
processor_data_column: str,
|
||||
include_error_column: bool = False,
|
||||
) -> Callable:
|
||||
"""Wrap the postprocess function to remove the processor_data_column.
|
||||
Note that we fully rely on users to determine which columns to carry over.
|
||||
|
||||
Error rows (with __inference_error__ set) bypass the user's postprocess
|
||||
function and return directly with the error information preserved.
|
||||
|
||||
Args:
|
||||
fn: The function to be applied.
|
||||
processor_data_column: The internal data column name of the processor.
|
||||
include_error_column: If True, always include __inference_error__ in output
|
||||
(Empty string for success rows, error message for failures). This ensures
|
||||
consistent schema across all output rows.
|
||||
|
||||
Returns:
|
||||
The wrapped function.
|
||||
"""
|
||||
|
||||
def _postprocess(row: dict[str, Any]) -> dict[str, Any]:
|
||||
if processor_data_column not in row:
|
||||
raise ValueError(
|
||||
f"[Internal] {processor_data_column} not found in row {row}"
|
||||
)
|
||||
|
||||
data = row[processor_data_column]
|
||||
|
||||
# Error rows bypass user postprocess to avoid crashes when
|
||||
# expected output fields are missing. Return entire data dict
|
||||
# to preserve debugging info (e.g., prompt).
|
||||
if data.get("__inference_error__", "") != "":
|
||||
return data
|
||||
|
||||
result = fn(data)
|
||||
if include_error_column:
|
||||
result["__inference_error__"] = ""
|
||||
return result
|
||||
|
||||
return _postprocess
|
||||
|
||||
|
||||
class StatefulStageUDF:
|
||||
"""A stage UDF wrapper that processes the input and output columns
|
||||
before and after the UDF.
|
||||
|
||||
Args:
|
||||
data_column: The internal data column name of the processor. The
|
||||
__call__ method takes the data column as the input of the UDF
|
||||
method, and encapsulates the output of the UDF method into the data
|
||||
column for the next stage.
|
||||
expected_input_keys: The expected input keys of the stage.
|
||||
"""
|
||||
|
||||
# The internal column name for the index of the row in the batch.
|
||||
# This is used to align the output of the UDF with the input batch.
|
||||
IDX_IN_BATCH_COLUMN: str = "__idx_in_batch"
|
||||
|
||||
def __init__(
|
||||
self, data_column: str, expected_input_keys: Optional[List[str]] = None
|
||||
):
|
||||
self.data_column = data_column
|
||||
self.expected_input_keys = set(expected_input_keys or [])
|
||||
|
||||
async def __call__(self, batch: Dict[str, Any]) -> AsyncIterator[Dict[str, Any]]:
|
||||
"""A stage UDF wrapper that processes the input and output columns
|
||||
before and after the UDF.
|
||||
|
||||
The expected schema of "batch" is:
|
||||
{data_column: {
|
||||
dataset columns,
|
||||
other intermediate columns
|
||||
},
|
||||
...other metadata columns...,
|
||||
}.
|
||||
|
||||
The input of the UDF will then [dataset columns and other intermediate columns].
|
||||
In addition, while the output of the UDF depends on the UDF implementation,
|
||||
the output schema is expected to be
|
||||
{data_column: {
|
||||
dataset columns,
|
||||
other intermediate columns,
|
||||
UDF output columns (will override above columns if they have the same name)
|
||||
},
|
||||
...other metadata columns...,
|
||||
}.
|
||||
And this will become the input of the next stage.
|
||||
|
||||
Examples:
|
||||
Input dataset columns: {A, B, C}
|
||||
Preprocess: (lambda row: {"D": row["A"] + 1})
|
||||
Input:
|
||||
UDF input: {A, B, C}
|
||||
UDF output: {D}
|
||||
Output: {__data: {A, B, C, D}}
|
||||
Stage 1:
|
||||
Input: {__data: {A, B, C, D}}
|
||||
UDF input: {A, B, C, D}
|
||||
UDF output: {E}
|
||||
Output: {__data: {A, B, C, D, E}}
|
||||
Stage 2:
|
||||
Input: {__data: {A, B, C, D, E}}
|
||||
UDF input: {A, B, C, D, E}
|
||||
UDF output: {F, E} # E is in-place updated.
|
||||
Output: {__data: {A, B, C, D, E, F}}
|
||||
Postprocess: (lambda row: {"G": row["F"], "A": row["A"], "E": row["E"]})
|
||||
Input: {__data: {A, B, C, D, E, F}}
|
||||
UDF input: {A, B, C, D, E, F}
|
||||
UDF output: {G, A, E}
|
||||
Output: {G, A, E} # User chooses to keep G, A, E.
|
||||
|
||||
Args:
|
||||
batch: The input batch.
|
||||
|
||||
Yields:
|
||||
(str, Any): An async iterator of the outputs.
|
||||
|
||||
TODO(MARK): The yield type should be `Dict[str, Any]`, pydoclint is bugged in 0.8.4 (https://github.com/jsh9/pydoclint/issues/288)
|
||||
"""
|
||||
# Handle the case where the batch is empty.
|
||||
# FIXME: This should not happen.
|
||||
if isinstance(batch, pyarrow.lib.Table) and batch.num_rows == 0:
|
||||
yield {}
|
||||
return
|
||||
|
||||
if self.data_column not in batch:
|
||||
raise ValueError(
|
||||
f"[Internal] {self.data_column} not found in batch {batch}"
|
||||
)
|
||||
|
||||
inputs = batch.pop(self.data_column)
|
||||
if hasattr(inputs, "tolist"):
|
||||
inputs = inputs.tolist()
|
||||
|
||||
# Separate error rows from normal rows BEFORE validation. Error rows
|
||||
# (those with __inference_error__ set) bypass the UDF to avoid crashes
|
||||
# when expected fields are missing (e.g., generated_tokens for DetokenizeUDF).
|
||||
normal_rows = []
|
||||
error_row_indices = set()
|
||||
for idx, row in enumerate(inputs):
|
||||
if row.get("__inference_error__", "") != "":
|
||||
error_row_indices.add(idx)
|
||||
else:
|
||||
normal_rows.append(row)
|
||||
|
||||
# Validate only normal rows - error rows may be missing required keys
|
||||
self.validate_inputs(normal_rows)
|
||||
|
||||
# Assign the index of the row in the batch to the idx_in_batch_column.
|
||||
# This is because the UDF output may be out-of-order (if asyncio.as_completed
|
||||
# is used internally for example), and we need to carry over unused input
|
||||
# columns to the next stage. Thus, we use the row index in batch to match
|
||||
# the output of the UDF with the input.
|
||||
for idx, row in enumerate(inputs):
|
||||
row[self.IDX_IN_BATCH_COLUMN] = idx
|
||||
|
||||
# Collect all outputs first, then return them in the original order
|
||||
# This is a requirement set by https://github.com/ray-project/ray/pull/54190/
|
||||
not_output_rows = set(range(len(inputs))) - error_row_indices
|
||||
if normal_rows:
|
||||
async for output in self.udf(normal_rows):
|
||||
if self.IDX_IN_BATCH_COLUMN not in output:
|
||||
raise ValueError(
|
||||
"The output of the UDF must contain the column "
|
||||
f"{self.IDX_IN_BATCH_COLUMN}."
|
||||
)
|
||||
idx_in_batch = output.pop(self.IDX_IN_BATCH_COLUMN)
|
||||
if idx_in_batch not in not_output_rows:
|
||||
raise ValueError(
|
||||
f"The row {idx_in_batch} is output twice. "
|
||||
"This is likely due to the UDF is not one-to-one."
|
||||
)
|
||||
not_output_rows.remove(idx_in_batch)
|
||||
|
||||
# Add stage outputs to the data column of the row.
|
||||
inputs[idx_in_batch].pop(self.IDX_IN_BATCH_COLUMN)
|
||||
inputs[idx_in_batch].update(output)
|
||||
|
||||
if not_output_rows:
|
||||
raise ValueError(f"The rows {not_output_rows} are not output.")
|
||||
|
||||
# Clean up idx column from error rows (normal rows already cleaned above)
|
||||
for idx in error_row_indices:
|
||||
inputs[idx].pop(self.IDX_IN_BATCH_COLUMN, None)
|
||||
|
||||
# Return all updated inputs in the original order
|
||||
yield {self.data_column: inputs}
|
||||
|
||||
def validate_inputs(self, inputs: List[Dict[str, Any]]):
|
||||
"""Validate the inputs to make sure the required keys are present.
|
||||
|
||||
Args:
|
||||
inputs: The inputs.
|
||||
|
||||
Raises:
|
||||
ValueError: If the required keys are not found.
|
||||
"""
|
||||
for inp in inputs:
|
||||
input_keys = set(inp.keys())
|
||||
|
||||
if self.IDX_IN_BATCH_COLUMN in input_keys:
|
||||
raise ValueError(
|
||||
f"The input column {self.IDX_IN_BATCH_COLUMN} is reserved "
|
||||
"for internal use."
|
||||
)
|
||||
|
||||
if not self.expected_input_keys:
|
||||
continue
|
||||
|
||||
missing_required = self.expected_input_keys - input_keys
|
||||
if missing_required:
|
||||
raise ValueError(
|
||||
f"Required input keys {missing_required} not found at the input of "
|
||||
f"{self.__class__.__name__}. Input keys: {input_keys}"
|
||||
)
|
||||
|
||||
async def udf(self, rows: List[Dict[str, Any]]) -> AsyncIterator[Dict[str, Any]]:
|
||||
raise NotImplementedError("StageUDF must implement the udf method")
|
||||
|
||||
|
||||
class StatefulStage(BaseModel):
|
||||
"""
|
||||
A basic building block to compose a Processor.
|
||||
"""
|
||||
|
||||
fn: Type[StatefulStageUDF] = Field(
|
||||
description="The well-optimized stateful UDF for this stage."
|
||||
)
|
||||
fn_constructor_kwargs: Dict[str, Any] = Field(
|
||||
default_factory=dict,
|
||||
description="The keyword arguments of the UDF constructor.",
|
||||
)
|
||||
map_batches_kwargs: Dict[str, Any] = Field(
|
||||
default_factory=lambda: {"concurrency": 1},
|
||||
description="The arguments of .map_batches(). Default {'concurrency': 1}.",
|
||||
)
|
||||
|
||||
def get_required_input_keys(self) -> Dict[str, str]:
|
||||
"""The required input keys of the stage and their descriptions."""
|
||||
return {}
|
||||
|
||||
def get_optional_input_keys(self) -> Dict[str, str]:
|
||||
"""The optional input keys of the stage and their descriptions."""
|
||||
return {}
|
||||
|
||||
def get_dataset_map_batches_kwargs(
|
||||
self,
|
||||
batch_size: int,
|
||||
data_column: str,
|
||||
) -> Dict[str, Any]:
|
||||
"""We separate fn and fn_constructor_kwargs in Stage for better UX,
|
||||
so we combine them with other map_batches_kwargs together in this method.
|
||||
|
||||
Args:
|
||||
batch_size: The batch size set by the processor config.
|
||||
data_column: The data column name set by the processor.
|
||||
|
||||
Returns:
|
||||
The dataset map_batches kwargs.
|
||||
"""
|
||||
kwargs = self.map_batches_kwargs.copy()
|
||||
batch_size_in_kwargs = kwargs.get("batch_size", batch_size)
|
||||
if batch_size_in_kwargs != batch_size:
|
||||
logger.warning(
|
||||
"batch_size is set to %d in map_batches_kwargs, but it will be "
|
||||
"overridden by the batch size configured by the processor %d.",
|
||||
batch_size_in_kwargs,
|
||||
batch_size,
|
||||
)
|
||||
kwargs["batch_size"] = batch_size
|
||||
|
||||
kwargs.update({"fn_constructor_kwargs": self.fn_constructor_kwargs.copy()})
|
||||
if "data_column" in kwargs["fn_constructor_kwargs"]:
|
||||
raise ValueError(
|
||||
"'data_column' cannot be used as in fn_constructor_kwargs."
|
||||
)
|
||||
|
||||
kwargs["fn_constructor_kwargs"]["data_column"] = data_column
|
||||
kwargs["fn_constructor_kwargs"]["expected_input_keys"] = list(
|
||||
self.get_required_input_keys().keys()
|
||||
)
|
||||
return kwargs
|
||||
|
||||
class Config:
|
||||
arbitrary_types_allowed = True
|
||||
validate_assignment = True
|
||||
@@ -0,0 +1,141 @@
|
||||
"""Apply chat template stage"""
|
||||
|
||||
from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, List, Optional, Type, Union
|
||||
|
||||
from ray.llm._internal.batch.stages.base import (
|
||||
StatefulStage,
|
||||
StatefulStageUDF,
|
||||
)
|
||||
from ray.llm._internal.common.utils.download_utils import (
|
||||
NodeModelDownloadable,
|
||||
download_model_files,
|
||||
)
|
||||
|
||||
|
||||
class ChatTemplateUDF(StatefulStageUDF):
|
||||
def __init__(
|
||||
self,
|
||||
data_column: str,
|
||||
expected_input_keys: List[str],
|
||||
model: str,
|
||||
chat_template: Optional[str] = None,
|
||||
chat_template_kwargs: Optional[Dict[str, Any]] = None,
|
||||
trust_remote_code: bool = False,
|
||||
):
|
||||
"""
|
||||
Initialize the ChatTemplateUDF.
|
||||
|
||||
Args:
|
||||
data_column: The data column name.
|
||||
expected_input_keys: The expected input keys of the stage.
|
||||
model: The model to use for the chat template.
|
||||
chat_template: The chat template in Jinja template format. This is
|
||||
usually not needed if the model checkpoint already contains the
|
||||
chat template.
|
||||
chat_template_kwargs: The optional kwargs to pass apply_chat_template.
|
||||
trust_remote_code: Whether to trust remote code when loading the model.
|
||||
"""
|
||||
from transformers import AutoProcessor
|
||||
|
||||
super().__init__(data_column, expected_input_keys)
|
||||
|
||||
# NOTE: We always use processor instead of tokenizer in this stage,
|
||||
# because tokenizers of VLM models may not have chat template attribute.
|
||||
# However, this may not be a reliable solution, because processors and
|
||||
# tokenizers are not standardized across different models.
|
||||
# Use EXCLUDE_SAFETENSORS for trust_remote_code models to ensure
|
||||
# Python config files are downloaded.
|
||||
download_mode = (
|
||||
NodeModelDownloadable.EXCLUDE_SAFETENSORS
|
||||
if trust_remote_code
|
||||
else NodeModelDownloadable.TOKENIZER_ONLY
|
||||
)
|
||||
model_path = download_model_files(
|
||||
model_id=model,
|
||||
mirror_config=None,
|
||||
download_model=download_mode,
|
||||
download_extra_files=False,
|
||||
)
|
||||
if TYPE_CHECKING:
|
||||
from transformers.processing_utils import ProcessorMixin
|
||||
from transformers.tokenization_utils_base import PreTrainedTokenizerBase
|
||||
|
||||
self.processor: Union[
|
||||
"PreTrainedTokenizerBase", "ProcessorMixin"
|
||||
] = AutoProcessor.from_pretrained(
|
||||
model_path, trust_remote_code=trust_remote_code
|
||||
)
|
||||
self.chat_template = chat_template
|
||||
self.chat_template_kwargs = chat_template_kwargs
|
||||
|
||||
async def udf(self, batch: List[Dict[str, Any]]) -> AsyncIterator[Dict[str, Any]]:
|
||||
"""
|
||||
Apply chat template to the given batch.
|
||||
|
||||
Args:
|
||||
batch: A list of rows to send.
|
||||
|
||||
Yields:
|
||||
Dict[str, Any]: A generator of rows with the chat template applied.
|
||||
"""
|
||||
prompts = []
|
||||
for row in batch:
|
||||
# PyArrow cannot handle the messages with images, so Ray Data
|
||||
# will fallback to use pickle for serialization. In this case,
|
||||
# the "messages" column is already a list of dicts and does not
|
||||
# have .tolist() method.
|
||||
if hasattr(row["messages"], "tolist"):
|
||||
conversation = row["messages"].tolist()
|
||||
else:
|
||||
conversation = row["messages"]
|
||||
add_generation_prompt = self._should_add_generation_prompt(conversation)
|
||||
# If we don't add a generation prompt, we should continue the final message.
|
||||
continue_final_message = not add_generation_prompt
|
||||
prompts.append(
|
||||
self.processor.apply_chat_template(
|
||||
conversation,
|
||||
tokenize=False,
|
||||
chat_template=self.chat_template,
|
||||
add_generation_prompt=add_generation_prompt,
|
||||
continue_final_message=continue_final_message,
|
||||
**(self.chat_template_kwargs or {}),
|
||||
)
|
||||
)
|
||||
assert len(batch) == len(prompts)
|
||||
|
||||
for row, prompt in zip(batch, prompts):
|
||||
yield {
|
||||
self.IDX_IN_BATCH_COLUMN: row[self.IDX_IN_BATCH_COLUMN],
|
||||
"prompt": prompt,
|
||||
}
|
||||
|
||||
def _should_add_generation_prompt(self, conversation: List[Dict[str, Any]]) -> bool:
|
||||
"""Determines if the generation prompt should be added for the given conversation.
|
||||
|
||||
Adds the generation prompt only if the last message is from the user.
|
||||
This is useful in cases where the user provides an assistant prefill
|
||||
message.
|
||||
|
||||
Args:
|
||||
conversation: The conversation to check.
|
||||
|
||||
Returns:
|
||||
True if the generation prompt should be added, False otherwise.
|
||||
"""
|
||||
return conversation[-1]["role"] == "user"
|
||||
|
||||
|
||||
class ChatTemplateStage(StatefulStage):
|
||||
"""
|
||||
A stage that applies chat template.
|
||||
"""
|
||||
|
||||
fn: Type[StatefulStageUDF] = ChatTemplateUDF
|
||||
|
||||
def get_required_input_keys(self) -> Dict[str, str]:
|
||||
"""The required input keys of the stage and their descriptions."""
|
||||
return {
|
||||
"messages": "A list of messages in OpenAI chat format. "
|
||||
"See https://platform.openai.com/docs/api-reference/chat/create "
|
||||
"for details."
|
||||
}
|
||||
@@ -0,0 +1,36 @@
|
||||
"""
|
||||
Shared utilities for stages.
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List, Union
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def truncate_str(s: str, max_length: int, suffix: str = "...[truncated]") -> str:
|
||||
"""Truncate a string to max_length, appending suffix if truncated."""
|
||||
if len(s) <= max_length:
|
||||
return s
|
||||
return s[:max_length] + suffix
|
||||
|
||||
|
||||
def maybe_convert_ndarray_to_list(
|
||||
params: Union[np.ndarray, List[Any], Dict[str, Any]]
|
||||
) -> Union[List[Any], Dict[str, Any]]:
|
||||
"""Convert all ndarray to list in the params. This is because Ray Data
|
||||
by default converts all lists to ndarrays when passing data around, but
|
||||
vLLM expects lists.
|
||||
|
||||
Args:
|
||||
params: The parameters to convert.
|
||||
|
||||
Returns:
|
||||
The converted parameters.
|
||||
"""
|
||||
if isinstance(params, dict):
|
||||
return {k: maybe_convert_ndarray_to_list(v) for k, v in params.items()}
|
||||
elif isinstance(params, list):
|
||||
return [maybe_convert_ndarray_to_list(v) for v in params]
|
||||
elif isinstance(params, np.ndarray):
|
||||
return params.tolist()
|
||||
return params
|
||||
@@ -0,0 +1,117 @@
|
||||
from typing import Any, Dict, Literal, Optional, Tuple, Type, TypeVar, Union
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
from ray.llm._internal.common.base_pydantic import BaseModelExtended
|
||||
|
||||
T = TypeVar("T", bound="_StageConfigBase")
|
||||
|
||||
|
||||
class _StageConfigBase(BaseModelExtended):
|
||||
enabled: bool = Field(default=True, description="Whether this stage is enabled.")
|
||||
# Optional overrides; processor-level defaults still apply
|
||||
batch_size: Optional[int] = Field(default=None, description="Rows per batch.")
|
||||
concurrency: Optional[Union[int, Tuple[int, int]]] = Field(
|
||||
default=None, description="Actor pool size or range for this stage."
|
||||
)
|
||||
runtime_env: Optional[Dict[str, Any]] = Field(
|
||||
default=None, description="Optional runtime env for this stage."
|
||||
)
|
||||
num_cpus: Optional[float] = Field(
|
||||
default=None,
|
||||
description="Number of CPUs to reserve for each map worker in this stage.",
|
||||
)
|
||||
memory: Optional[float] = Field(
|
||||
default=None,
|
||||
description="Heap memory in bytes to reserve for each map worker in this stage.",
|
||||
)
|
||||
|
||||
|
||||
class ChatTemplateStageConfig(_StageConfigBase):
|
||||
model_source: Optional[str] = Field(
|
||||
default=None, description="Model source/identifier for this stage."
|
||||
)
|
||||
chat_template: Optional[str] = Field(default=None)
|
||||
chat_template_kwargs: Optional[Dict[str, Any]] = Field(default=None)
|
||||
|
||||
|
||||
class TokenizerStageConfig(_StageConfigBase):
|
||||
model_source: Optional[str] = Field(
|
||||
default=None, description="Model source/identifier for this stage."
|
||||
)
|
||||
|
||||
|
||||
class DetokenizeStageConfig(_StageConfigBase):
|
||||
model_source: Optional[str] = Field(
|
||||
default=None, description="Model source/identifier for this stage."
|
||||
)
|
||||
|
||||
|
||||
class PrepareMultimodalStageConfig(_StageConfigBase):
|
||||
model_config_kwargs: Optional[Dict[str, Any]] = Field(
|
||||
default=None,
|
||||
description="Optional kwargs to pass to the model config. See available model config "
|
||||
"kwargs at https://docs.vllm.ai/en/latest/api/vllm/config/#vllm.config.ModelConfig",
|
||||
)
|
||||
chat_template_content_format: Optional[Literal["string", "openai"]] = Field(
|
||||
default="string",
|
||||
description="The content format to use for the chat template. "
|
||||
"This is used to format the chat template content according to a specific model.",
|
||||
)
|
||||
apply_sys_msg_formatting: Optional[bool] = Field(
|
||||
default=False,
|
||||
description="Whether to apply formatting system messages.",
|
||||
)
|
||||
|
||||
|
||||
class HttpRequestStageConfig(_StageConfigBase):
|
||||
pass
|
||||
|
||||
|
||||
def resolve_stage_config(
|
||||
stage_cfg_value: Union[bool, Dict[str, Any], _StageConfigBase],
|
||||
stage_config_cls: Type[T],
|
||||
processor_defaults: Optional[Dict[str, Any]] = None,
|
||||
) -> T:
|
||||
"""Resolve a stage config value (bool | dict | StageConfig) into a typed StageConfig.
|
||||
|
||||
Args:
|
||||
stage_cfg_value: The stage config value (bool, dict, or typed StageConfig).
|
||||
stage_config_cls: The StageConfig class to instantiate.
|
||||
processor_defaults: Optional dict of processor-level defaults to merge in.
|
||||
Expected keys: 'batch_size', 'concurrency', 'runtime_env', 'model_source'.
|
||||
|
||||
Returns:
|
||||
Resolved StageConfig instance with defaults merged.
|
||||
"""
|
||||
processor_defaults = processor_defaults or {}
|
||||
|
||||
# If already a typed config, create a copy to avoid mutating the input
|
||||
if isinstance(stage_cfg_value, stage_config_cls):
|
||||
resolved = stage_config_cls.model_validate(stage_cfg_value.model_dump())
|
||||
# If bool, create minimal config with enabled flag
|
||||
elif isinstance(stage_cfg_value, bool):
|
||||
resolved = stage_config_cls(enabled=stage_cfg_value)
|
||||
# If dict, parse it into the config class
|
||||
elif isinstance(stage_cfg_value, dict):
|
||||
resolved = stage_config_cls(**stage_cfg_value)
|
||||
else:
|
||||
raise TypeError(
|
||||
f"Unsupported type for stage config: {type(stage_cfg_value).__name__}. "
|
||||
f"Expected bool, dict, or {stage_config_cls.__name__} instance. "
|
||||
f"Got: {stage_cfg_value}"
|
||||
)
|
||||
|
||||
# Merge processor defaults for fields not explicitly set
|
||||
default_fields = ["batch_size", "concurrency", "runtime_env", "model_source"]
|
||||
for field_name in default_fields:
|
||||
# Skip if field doesn't exist on this config class (e.g., model_source only on some stages)
|
||||
if not hasattr(resolved, field_name):
|
||||
continue
|
||||
if (
|
||||
getattr(resolved, field_name, None) is None
|
||||
and field_name in processor_defaults
|
||||
):
|
||||
setattr(resolved, field_name, processor_defaults[field_name])
|
||||
|
||||
return resolved
|
||||
@@ -0,0 +1,177 @@
|
||||
"""HTTP Request Stage"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
import traceback
|
||||
from typing import Any, AsyncIterator, Callable, Dict, List, Optional, Type
|
||||
|
||||
import aiohttp
|
||||
import aiohttp.web_exceptions
|
||||
import numpy as np
|
||||
from aiohttp.client_exceptions import ClientPayloadError
|
||||
|
||||
from ray.llm._internal.batch.stages.base import StatefulStage, StatefulStageUDF
|
||||
|
||||
|
||||
class NumpyEncoder(json.JSONEncoder):
|
||||
def default(self, obj):
|
||||
if isinstance(obj, np.ndarray):
|
||||
return obj.tolist()
|
||||
elif isinstance(obj, (np.integer, np.floating)):
|
||||
return obj.item()
|
||||
elif isinstance(obj, np.bool_):
|
||||
return bool(obj)
|
||||
else:
|
||||
return super().default(obj)
|
||||
|
||||
|
||||
class HttpRequestUDF(StatefulStageUDF):
|
||||
RETRYABLE_STATUS_CODES = [429, 408, 504, 502, 503]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
data_column: str,
|
||||
expected_input_keys: List[str],
|
||||
url: str,
|
||||
additional_header: Optional[Dict[str, Any]] = None,
|
||||
qps: Optional[int] = None,
|
||||
max_retries: int = 0,
|
||||
base_retry_wait_time_in_s: float = 1.0,
|
||||
session_factory: Optional[Callable[[], aiohttp.ClientSession]] = None,
|
||||
):
|
||||
"""
|
||||
Initialize the HttpRequestUDF.
|
||||
|
||||
Args:
|
||||
data_column: The data column name.
|
||||
expected_input_keys: The expected input keys of the stage.
|
||||
url: The URL to send the HTTP request to.
|
||||
additional_header: The additional headers to send with the HTTP request.
|
||||
qps: The maximum number of requests per second.
|
||||
max_retries: The maximum number of retries per request in the event of failures. We retry with exponential backoff upto this specific maximum retries.
|
||||
base_retry_wait_time_in_s: The base retry wait time during exponential backoff.
|
||||
session_factory: Optional session factory to be used for initializing a client session.
|
||||
"""
|
||||
super().__init__(data_column, expected_input_keys)
|
||||
self.url = url
|
||||
self.additional_header = additional_header or {}
|
||||
self.qps = qps
|
||||
self.max_retries = max_retries
|
||||
self.base_retry_wait_time_in_s = base_retry_wait_time_in_s
|
||||
self.session_factory = session_factory or aiohttp.ClientSession
|
||||
|
||||
async def udf(self, batch: List[Dict[str, Any]]) -> AsyncIterator[Dict[str, Any]]:
|
||||
"""
|
||||
Send HTTP requests to the given URL.
|
||||
|
||||
Args:
|
||||
batch: A list of rows to send.
|
||||
|
||||
Yields:
|
||||
Dict[str, Any]: A generator of rows of the response of the HTTP request.
|
||||
"""
|
||||
# preprocess to get request body for the given batch
|
||||
request_bodies = [None] * len(batch)
|
||||
for row in batch:
|
||||
# Normalize the row to a JSON body.
|
||||
request_bodies[row[self.IDX_IN_BATCH_COLUMN]] = json.dumps(
|
||||
row["payload"], cls=NumpyEncoder
|
||||
)
|
||||
|
||||
async with self.session_factory() as session:
|
||||
start_time = time.time()
|
||||
request_count = 0
|
||||
pending_requests = []
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
**self.additional_header,
|
||||
}
|
||||
|
||||
# First send all requests based on QPS
|
||||
for row in batch:
|
||||
# Rate limit based on qps if specified
|
||||
if self.qps is not None:
|
||||
request_count += 1
|
||||
expected_time = request_count / self.qps
|
||||
elapsed = time.time() - start_time
|
||||
if elapsed < expected_time:
|
||||
await asyncio.sleep(expected_time - elapsed)
|
||||
|
||||
# self.IDX_IN_BATCH_COLUMN is the index of row in the batch
|
||||
json_body = request_bodies[row[self.IDX_IN_BATCH_COLUMN]]
|
||||
# Create request but don't await it yet
|
||||
request = session.post(
|
||||
self.url,
|
||||
headers=headers,
|
||||
data=json_body,
|
||||
)
|
||||
pending_requests.append((row[self.IDX_IN_BATCH_COLUMN], request))
|
||||
|
||||
# Now receive all responses
|
||||
for idx_in_batch_column, request in pending_requests:
|
||||
resp_json = None
|
||||
last_exception = None
|
||||
last_exception_traceback = None
|
||||
for retry_count in range(self.max_retries + 1):
|
||||
if retry_count > 0:
|
||||
json_body = request_bodies[idx_in_batch_column]
|
||||
request = session.post(
|
||||
self.url,
|
||||
headers=headers,
|
||||
data=json_body,
|
||||
)
|
||||
try:
|
||||
async with await request as response:
|
||||
status_code = response.status
|
||||
# check status and see if it's retry worthy
|
||||
if status_code in self.RETRYABLE_STATUS_CODES:
|
||||
last_exception = aiohttp.web_exceptions.HTTPException(
|
||||
reason=response.reason
|
||||
)
|
||||
last_exception.status_code = status_code
|
||||
wait_time = self.base_retry_wait_time_in_s * (
|
||||
2**retry_count
|
||||
)
|
||||
await asyncio.sleep(wait_time)
|
||||
continue
|
||||
resp_json = await response.json()
|
||||
if self.IDX_IN_BATCH_COLUMN in resp_json:
|
||||
raise ValueError(
|
||||
"The response of the HTTP request must not contain "
|
||||
f"the column {self.IDX_IN_BATCH_COLUMN}."
|
||||
)
|
||||
break
|
||||
except (
|
||||
asyncio.TimeoutError,
|
||||
aiohttp.ClientConnectionError,
|
||||
ClientPayloadError,
|
||||
) as e:
|
||||
last_exception_traceback = traceback.format_exc()
|
||||
last_exception = type(e).__name__
|
||||
wait_time = self.base_retry_wait_time_in_s * (2**retry_count)
|
||||
await asyncio.sleep(wait_time)
|
||||
continue
|
||||
if not resp_json:
|
||||
raise RuntimeError(
|
||||
f"Reached maximum retries of {self.max_retries} for input row {batch[idx_in_batch_column]}. Previous Exception: {last_exception}. Full Traceback: \n{last_exception_traceback}"
|
||||
)
|
||||
yield {
|
||||
self.IDX_IN_BATCH_COLUMN: idx_in_batch_column,
|
||||
"http_response": resp_json,
|
||||
}
|
||||
|
||||
|
||||
class HttpRequestStage(StatefulStage):
|
||||
"""
|
||||
A stage that sends HTTP requests.
|
||||
"""
|
||||
|
||||
fn: Type[StatefulStageUDF] = HttpRequestUDF
|
||||
|
||||
def get_required_input_keys(self) -> Dict[str, str]:
|
||||
"""The required input keys of the stage and their descriptions."""
|
||||
return {
|
||||
"payload": "The payload to send to the HTTP request. "
|
||||
"It should be in JSON format."
|
||||
}
|
||||
@@ -0,0 +1,157 @@
|
||||
"""Prepare Multimodal Stage"""
|
||||
|
||||
import asyncio
|
||||
from typing import Any, AsyncIterator, Dict, List
|
||||
|
||||
from ray.llm._internal.batch.stages.base import StatefulStage, StatefulStageUDF
|
||||
|
||||
|
||||
class PrepareMultimodalUDF(StatefulStageUDF):
|
||||
def __init__(
|
||||
self,
|
||||
data_column: str,
|
||||
expected_input_keys: List[str],
|
||||
model_config_kwargs: Dict[str, Any],
|
||||
chat_template_content_format: str,
|
||||
apply_sys_msg_formatting: bool = False,
|
||||
):
|
||||
"""
|
||||
Initialize the PrepareMultimodalUDF.
|
||||
|
||||
Args:
|
||||
data_column: The data column name.
|
||||
expected_input_keys: The expected input keys of the stage.
|
||||
model_config_kwargs: The kwargs to pass to the model config.
|
||||
chat_template_content_format: The format to render message content.
|
||||
apply_sys_msg_formatting: Whether to skip formatting system messages.
|
||||
"""
|
||||
super().__init__(data_column, expected_input_keys)
|
||||
|
||||
try:
|
||||
from vllm.config import ModelConfig
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"vLLM is not installed or failed to import. Please run "
|
||||
"`pip install ray[llm]` to install required dependencies."
|
||||
) from e
|
||||
|
||||
self.model_config = ModelConfig(**model_config_kwargs)
|
||||
self.chat_template_content_format = chat_template_content_format
|
||||
self.apply_sys_msg_formatting = apply_sys_msg_formatting
|
||||
|
||||
def _extract_system_messages(
|
||||
self, messages: List[Dict[str, Any]]
|
||||
) -> tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
|
||||
"""Extract system messages from the message list.
|
||||
|
||||
System messages are kept as strings (not converted to list format) to avoid
|
||||
issues with chat templates that expect string system messages, e.g. Pixtral.
|
||||
|
||||
Args:
|
||||
messages: The full message list.
|
||||
|
||||
Returns:
|
||||
A tuple of (system_messages, non_system_messages).
|
||||
"""
|
||||
system_messages = []
|
||||
non_system_messages = []
|
||||
|
||||
for msg in messages:
|
||||
if msg.get("role") == "system":
|
||||
system_content = msg.get("content")
|
||||
if isinstance(system_content, list):
|
||||
text_parts = []
|
||||
for part in system_content:
|
||||
if isinstance(part, dict) and part.get("type") == "text":
|
||||
text_value = part.get("text") or part.get("content")
|
||||
if text_value:
|
||||
text_parts.append(str(text_value))
|
||||
elif isinstance(part, str) and part:
|
||||
text_parts.append(part)
|
||||
system_content = "\n".join(text_parts) if text_parts else ""
|
||||
|
||||
system_messages.append({**msg, "content": system_content})
|
||||
else:
|
||||
non_system_messages.append(msg)
|
||||
|
||||
return system_messages, non_system_messages
|
||||
|
||||
async def udf(self, batch: List[Dict[str, Any]]) -> AsyncIterator[Dict[str, Any]]:
|
||||
"""
|
||||
Process multimodal data from input messages.
|
||||
|
||||
Args:
|
||||
batch: A list of rows to process.
|
||||
|
||||
Yields:
|
||||
Dict[str, Any]: A dictionary containing the multimodal data
|
||||
along with processing metadata.
|
||||
"""
|
||||
try:
|
||||
from vllm.entrypoints.chat_utils import parse_chat_messages_async
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"vLLM is not installed or failed to import. Please run "
|
||||
"`pip install ray[llm]` to install required dependencies."
|
||||
) from e
|
||||
|
||||
async def _process_row(row: Dict[str, Any]):
|
||||
# Extract system messages to keep them as strings (not converted to list format)
|
||||
# This avoids issues with chat templates that expect string system messages.
|
||||
system_messages = []
|
||||
messages_to_parse = row["messages"]
|
||||
|
||||
if self.apply_sys_msg_formatting:
|
||||
system_messages, messages_to_parse = self._extract_system_messages(
|
||||
row["messages"]
|
||||
)
|
||||
|
||||
# Users can provide stable IDs for each multimodal item from messages to
|
||||
# enable engine to cache and reuse work across requests.
|
||||
conversation, mm_data, mm_uuids = await parse_chat_messages_async(
|
||||
messages_to_parse,
|
||||
self.model_config,
|
||||
content_format=self.chat_template_content_format,
|
||||
)
|
||||
|
||||
if system_messages:
|
||||
conversation = system_messages + conversation
|
||||
|
||||
return row, conversation, mm_uuids, mm_data
|
||||
|
||||
tasks = [asyncio.create_task(_process_row(row)) for row in batch]
|
||||
|
||||
for task in asyncio.as_completed(tasks):
|
||||
row, conversation, uuid, multimodal_data = await task
|
||||
|
||||
output = {
|
||||
k: v
|
||||
for k, v in row.items()
|
||||
if k not in ("messages", self.IDX_IN_BATCH_COLUMN)
|
||||
}
|
||||
output.update(
|
||||
{
|
||||
self.IDX_IN_BATCH_COLUMN: row[self.IDX_IN_BATCH_COLUMN],
|
||||
"multimodal_data": multimodal_data,
|
||||
# Use the parsed conversation which has placeholders embedded instead of the original messages
|
||||
"messages": conversation,
|
||||
"multimodal_uuids": uuid,
|
||||
}
|
||||
)
|
||||
yield output
|
||||
|
||||
|
||||
class PrepareMultimodalStage(StatefulStage):
|
||||
"""
|
||||
A stage that prepares multimodal data from the input messages for a specific model.
|
||||
"""
|
||||
|
||||
fn: StatefulStageUDF = PrepareMultimodalUDF
|
||||
|
||||
def get_required_input_keys(self) -> Dict[str, str]:
|
||||
"""The required input keys of the stage and their descriptions."""
|
||||
return {
|
||||
"messages": "A list of messages in OpenAI chat format. "
|
||||
"See https://platform.openai.com/docs/api-reference/chat/create "
|
||||
"for details."
|
||||
}
|
||||
@@ -0,0 +1,246 @@
|
||||
"""The stage that runs serve deployment."""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
import uuid
|
||||
from typing import Any, AsyncIterator, Dict, List, Optional, Tuple, Type
|
||||
|
||||
from pydantic import BaseModel, ValidationError
|
||||
|
||||
from ray import serve
|
||||
from ray.exceptions import RayTaskError
|
||||
from ray.llm._internal.batch.stages.base import (
|
||||
StatefulStage,
|
||||
StatefulStageUDF,
|
||||
)
|
||||
from ray.llm._internal.batch.stages.common import truncate_str
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_MAX_PROMPT_LENGTH_IN_ERROR = 200
|
||||
|
||||
# Request-level errors safe to catch. Unknown errors are treated as fatal.
|
||||
# TimeoutError is included so a single stuck request (see request_timeout_s)
|
||||
# is dropped as an error row rather than blocking the batch forever.
|
||||
_SERVE_RECOVERABLE_ERRORS = (
|
||||
ValueError,
|
||||
TypeError,
|
||||
ValidationError,
|
||||
TimeoutError,
|
||||
)
|
||||
|
||||
|
||||
class ServeDeploymentStageUDF(StatefulStageUDF):
|
||||
def __init__(
|
||||
self,
|
||||
data_column: str,
|
||||
expected_input_keys: List[str],
|
||||
*,
|
||||
deployment_name: str,
|
||||
app_name: str,
|
||||
dtype_mapping: Dict[str, Type[Any]],
|
||||
should_continue_on_error: bool = False,
|
||||
request_timeout_s: Optional[float] = None,
|
||||
):
|
||||
"""
|
||||
Initialize the ServeDeploymentStageUDF.
|
||||
|
||||
Args:
|
||||
data_column: The data column name.
|
||||
expected_input_keys: The expected input keys of the stage.
|
||||
deployment_name: The name of the deployment.
|
||||
app_name: The name of the deployment app.
|
||||
dtype_mapping: The mapping of the request class name to the request class.
|
||||
should_continue_on_error: If True, continue processing when inference
|
||||
fails for a row instead of raising. Failed rows will have
|
||||
'__inference_error__' set to the error message.
|
||||
request_timeout_s: Optional per-request timeout in seconds. When set,
|
||||
a request that does not return within this many seconds raises
|
||||
TimeoutError instead of awaiting indefinitely. TimeoutError is
|
||||
recoverable, so with should_continue_on_error=True the row is
|
||||
emitted as an error row rather than crashing the batch.
|
||||
"""
|
||||
super().__init__(data_column, expected_input_keys)
|
||||
self._dtype_mapping = dtype_mapping
|
||||
self.should_continue_on_error = should_continue_on_error
|
||||
self.request_timeout_s = request_timeout_s
|
||||
|
||||
# Using stream=True as LLM serve deployments return async generators.
|
||||
# TODO (Kourosh): Generalize this to support non-streaming deployments.
|
||||
self._dh = serve.get_deployment_handle(deployment_name, app_name).options(
|
||||
stream=True
|
||||
)
|
||||
self.request_id = 0
|
||||
|
||||
def _prepare_request(
|
||||
self, row: Dict[str, Any]
|
||||
) -> Tuple[Dict[str, Any], Optional[Type[Any]], str]:
|
||||
"""
|
||||
Decorate the request with metadata related to the batch.
|
||||
|
||||
Args:
|
||||
row: The row.
|
||||
|
||||
Returns:
|
||||
A tuple of (decorated_request, dtype, method_name). dtype is the class of
|
||||
the request object and can be None if the serve deployment accepts a raw
|
||||
dict. method_name is the name of the method to invoke on the deployment.
|
||||
"""
|
||||
method = row.get("method")
|
||||
dtype_name = row.get("dtype")
|
||||
|
||||
dtype = None
|
||||
if dtype_name is not None:
|
||||
if not self._dtype_mapping or dtype_name not in self._dtype_mapping:
|
||||
raise ValueError(
|
||||
f"{dtype_name} must be provided in "
|
||||
"ServeDeploymentProcessorConfig's dtype_mapping."
|
||||
)
|
||||
dtype = self._dtype_mapping[dtype_name]
|
||||
|
||||
request_kwargs = row.pop("request_kwargs")
|
||||
request = {
|
||||
"request_id": str(self.request_id),
|
||||
"idx_in_batch": row[self.IDX_IN_BATCH_COLUMN],
|
||||
**request_kwargs,
|
||||
}
|
||||
self.request_id += 1
|
||||
|
||||
return request, dtype, method
|
||||
|
||||
async def generate_async(
|
||||
self, row: Dict[str, Any]
|
||||
) -> Tuple[Dict[str, Any], Dict[str, Any], float]:
|
||||
"""
|
||||
Run the serve deployment.
|
||||
|
||||
Args:
|
||||
row: The row to run the serve deployment on.
|
||||
|
||||
Returns:
|
||||
The response from the serve deployment.
|
||||
"""
|
||||
request, dtype, method = self._prepare_request(row)
|
||||
request_obj = dtype(**request) if dtype else request
|
||||
|
||||
if getattr(self._dh, method) is None:
|
||||
raise ValueError(f"Method {method} not found in the serve deployment.")
|
||||
|
||||
t = time.perf_counter()
|
||||
# Directly using anext() requires python3.10 and above
|
||||
response_gen = getattr(self._dh, method).remote(request_obj)
|
||||
if self.request_timeout_s is not None:
|
||||
try:
|
||||
output_data = await asyncio.wait_for(
|
||||
response_gen.__anext__(), timeout=self.request_timeout_s
|
||||
)
|
||||
except asyncio.TimeoutError as e:
|
||||
raise TimeoutError(
|
||||
f"Request timed out after {self.request_timeout_s}s waiting "
|
||||
f"for deployment '{method}' to return a response."
|
||||
) from e
|
||||
else:
|
||||
output_data = await response_gen.__anext__()
|
||||
time_taken = time.perf_counter() - t
|
||||
|
||||
# Convert the output data to a dict if it is a Pydantic model.
|
||||
if isinstance(output_data, BaseModel):
|
||||
output_data = output_data.model_dump()
|
||||
|
||||
return request, output_data, time_taken
|
||||
|
||||
def _is_recoverable_error(self, exc: Exception) -> bool:
|
||||
"""Check if exception is recoverable. Unknown errors are treated as fatal."""
|
||||
if isinstance(exc, _SERVE_RECOVERABLE_ERRORS):
|
||||
return True
|
||||
# RayTaskError wraps remote exceptions - check the cause
|
||||
if isinstance(exc, RayTaskError) and hasattr(exc, "cause"):
|
||||
if isinstance(exc.cause, _SERVE_RECOVERABLE_ERRORS):
|
||||
return True
|
||||
return False
|
||||
|
||||
async def _generate_with_error_handling(
|
||||
self,
|
||||
row: Dict[str, Any],
|
||||
batch_uuid: uuid.UUID,
|
||||
) -> Dict[str, Any]:
|
||||
"""Generate output for a row, yielding error row on recoverable failure."""
|
||||
idx_in_batch = row[self.IDX_IN_BATCH_COLUMN]
|
||||
# Save before generate_async pops it
|
||||
original_request_kwargs = row.get("request_kwargs", {})
|
||||
try:
|
||||
request, output, time_taken = await self.generate_async(row)
|
||||
|
||||
return {
|
||||
**output,
|
||||
"request_id": request["request_id"],
|
||||
self.IDX_IN_BATCH_COLUMN: request["idx_in_batch"],
|
||||
"batch_uuid": batch_uuid.hex,
|
||||
"time_taken": time_taken,
|
||||
"__inference_error__": "",
|
||||
}
|
||||
except Exception as e:
|
||||
# Only recover from known recoverable errors; unknown errors propagate
|
||||
if not self._is_recoverable_error(e) or not self.should_continue_on_error:
|
||||
raise
|
||||
|
||||
error_msg = f"{type(e).__name__}: {str(e)}"
|
||||
logger.warning(
|
||||
"[Serve Deployment] Inference failed for row %d in batch %s: %s",
|
||||
idx_in_batch,
|
||||
batch_uuid.hex,
|
||||
error_msg,
|
||||
)
|
||||
|
||||
# Include request_kwargs snippet for debuggability
|
||||
request_str = truncate_str(
|
||||
str(original_request_kwargs), _MAX_PROMPT_LENGTH_IN_ERROR
|
||||
)
|
||||
|
||||
return {
|
||||
self.IDX_IN_BATCH_COLUMN: idx_in_batch,
|
||||
"batch_uuid": batch_uuid.hex,
|
||||
"__inference_error__": error_msg,
|
||||
"request_kwargs": request_str,
|
||||
}
|
||||
|
||||
async def udf(self, batch: List[Dict[str, Any]]) -> AsyncIterator[Dict[str, Any]]:
|
||||
"""
|
||||
Run the serve deployment.
|
||||
|
||||
Args:
|
||||
batch: A list of rows to run the serve deployment on.
|
||||
|
||||
Yields:
|
||||
Dict[str, Any]: A dictionary containing the response from the serve
|
||||
deployment along with processing metadata.
|
||||
"""
|
||||
batch_uuid = uuid.uuid4()
|
||||
t = time.perf_counter()
|
||||
|
||||
tasks = [
|
||||
asyncio.create_task(self._generate_with_error_handling(row, batch_uuid))
|
||||
for row in batch
|
||||
]
|
||||
|
||||
for resp in asyncio.as_completed(tasks):
|
||||
yield await resp
|
||||
|
||||
batch_time_taken = time.perf_counter() - t
|
||||
logger.info(
|
||||
"[LLM Batch - Serve Deployment] Elapsed time for batch %s with size %d: %s",
|
||||
batch_uuid.hex,
|
||||
len(batch),
|
||||
batch_time_taken,
|
||||
)
|
||||
|
||||
|
||||
class ServeDeploymentStage(StatefulStage):
|
||||
fn: Type[StatefulStageUDF] = ServeDeploymentStageUDF
|
||||
|
||||
def get_required_input_keys(self) -> Dict[str, str]:
|
||||
return {
|
||||
"method": "Name of the method to invoke on the serve deployment.",
|
||||
"request_kwargs": "The request_kwargs to construct the request.",
|
||||
}
|
||||
@@ -0,0 +1,387 @@
|
||||
"""The stage that runs SGLang engine."""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
import uuid
|
||||
from contextlib import nullcontext
|
||||
from typing import Any, AsyncIterator, Dict, List, Optional, Tuple, Type
|
||||
|
||||
from pydantic import BaseModel, root_validator
|
||||
|
||||
from ray.llm._internal.batch.constants import SGLangTaskType, TypeSGLangTaskType
|
||||
from ray.llm._internal.batch.stages.base import (
|
||||
StatefulStage,
|
||||
StatefulStageUDF,
|
||||
)
|
||||
from ray.llm._internal.batch.stages.common import maybe_convert_ndarray_to_list
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class SGLangEngineRequest(BaseModel):
|
||||
"""A request to the SGLang engine."""
|
||||
|
||||
# The request ID for the LLM engine (unique per replica).
|
||||
request_id: int
|
||||
# The index of the request in the batch.
|
||||
idx_in_batch: int
|
||||
# The input prompt.
|
||||
prompt: Optional[str] = None
|
||||
# Alternative to text. Specify the input as token IDs instead of text.
|
||||
prompt_token_ids: Optional[List[int]] = None
|
||||
# The sampling parameters (more details can be seen in https://docs.sglang.ai/backend/sampling_params.html).
|
||||
params: Optional[Dict[str, Any]] = None
|
||||
|
||||
class Config:
|
||||
validate_assignment = True
|
||||
arbitrary_types_allowed = True
|
||||
|
||||
|
||||
class SGLangOutputData(BaseModel):
|
||||
"""The output of the SGLang engine."""
|
||||
|
||||
prompt: Optional[str] = None
|
||||
prompt_token_ids: Optional[List[int]] = None
|
||||
num_input_tokens: int
|
||||
|
||||
# Generate fields.
|
||||
generated_tokens: Optional[List[int]] = None
|
||||
generated_text: Optional[str] = None
|
||||
num_generated_tokens: int
|
||||
|
||||
# Metrics fields.
|
||||
metrics: Optional[Dict[str, Any]] = None
|
||||
|
||||
@classmethod
|
||||
def from_sglang_engine_output(cls, output: Dict[str, Any]) -> "SGLangOutputData":
|
||||
"""Create a SGLangOutputData from a SGLang engine output."""
|
||||
|
||||
# Set by `_generate_async`.
|
||||
assert "prompt" in output
|
||||
assert "prompt_token_ids" in output
|
||||
|
||||
# Returned in the native output of the SGLang engine.
|
||||
assert "meta_info" in output
|
||||
assert "prompt_tokens" in output["meta_info"]
|
||||
assert "completion_tokens" in output["meta_info"]
|
||||
|
||||
data = cls(
|
||||
prompt=output["prompt"],
|
||||
prompt_token_ids=output["prompt_token_ids"],
|
||||
num_input_tokens=output["meta_info"]["prompt_tokens"],
|
||||
generated_tokens=output["output_ids"] if "output_ids" in output else None,
|
||||
generated_text=output["text"] if "text" in output else None,
|
||||
num_generated_tokens=output["meta_info"]["completion_tokens"],
|
||||
)
|
||||
|
||||
return data
|
||||
|
||||
class Config:
|
||||
validate_assignment = True
|
||||
arbitrary_types_allowed = True
|
||||
|
||||
|
||||
class SGLangEngineWrapper:
|
||||
"""Wrapper around the SGLang engine to handle async requests.
|
||||
|
||||
Args:
|
||||
idx_in_batch_column: The column name for the index of the row in the batch.
|
||||
max_pending_requests: The maximum number of pending requests in the queue.
|
||||
**kwargs: The keyword arguments for the engine.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
idx_in_batch_column: str,
|
||||
max_pending_requests: int = -1,
|
||||
**kwargs,
|
||||
):
|
||||
self.request_id = 0
|
||||
self.idx_in_batch_column = idx_in_batch_column
|
||||
self.task_type = kwargs.pop("task", SGLangTaskType.GENERATE)
|
||||
self.model = kwargs.pop("model", None)
|
||||
assert self.model is not None
|
||||
# We need to rename the `model` to `model_path` for SGLang.
|
||||
kwargs["model_path"] = self.model
|
||||
|
||||
# Set the skip_tokenizer_init to True by default for SGLang engine
|
||||
# because we will not use the tokenizer/detokenizer in SGLang engine
|
||||
# by default.
|
||||
self.skip_tokenizer_init = kwargs.pop("skip_tokenizer_init", True)
|
||||
kwargs["skip_tokenizer_init"] = self.skip_tokenizer_init
|
||||
|
||||
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
|
||||
|
||||
# Initialize the SGLang engine
|
||||
self.engine = sglang.Engine(**kwargs)
|
||||
|
||||
# The performance gets really bad if there are too many requests in the pending queue.
|
||||
# We work around it with semaphore to limit the number of concurrent requests in the engine.
|
||||
self.max_pending_requests = max_pending_requests
|
||||
if self.max_pending_requests > 0:
|
||||
self.semaphore = asyncio.Semaphore(self.max_pending_requests)
|
||||
else:
|
||||
# Use contextlib.nullcontext which works for both sync and async contexts.
|
||||
self.semaphore = nullcontext()
|
||||
|
||||
async def _prepare_llm_request(self, row: Dict[str, Any]) -> SGLangEngineRequest:
|
||||
"""Prepare the inputs for LLM inference.
|
||||
|
||||
Args:
|
||||
row: The row.
|
||||
|
||||
Returns:
|
||||
A single SGLangEngineRequest.
|
||||
"""
|
||||
prompt = row.pop("prompt")
|
||||
|
||||
if "tokenized_prompt" in row:
|
||||
tokenized_prompt = row.pop("tokenized_prompt").tolist()
|
||||
else:
|
||||
tokenized_prompt = None
|
||||
|
||||
# Prepare sampling parameters.
|
||||
if self.task_type == SGLangTaskType.GENERATE:
|
||||
params = maybe_convert_ndarray_to_list(row.pop("sampling_params"))
|
||||
else:
|
||||
raise ValueError(f"Unsupported task type: {self.task_type}")
|
||||
|
||||
if tokenized_prompt is not None and not self.skip_tokenizer_init:
|
||||
raise ValueError(
|
||||
"To use a token-in-token-out mode of SGLang Engine, please set engine_kwargs['skip_tokenizer_init'] to True."
|
||||
)
|
||||
|
||||
request = SGLangEngineRequest(
|
||||
request_id=self.request_id,
|
||||
idx_in_batch=row[self.idx_in_batch_column],
|
||||
prompt=prompt,
|
||||
prompt_token_ids=tokenized_prompt,
|
||||
params=params,
|
||||
)
|
||||
self.request_id += 1
|
||||
return request
|
||||
|
||||
async def generate_async(
|
||||
self, row: Dict[str, Any]
|
||||
) -> Tuple[SGLangEngineRequest, Dict[str, Any], float]:
|
||||
"""Process a single request.
|
||||
|
||||
Args:
|
||||
row: The input row.
|
||||
|
||||
Returns:
|
||||
A tuple of index in batch, request output and bypassed custom fields, and time taken.
|
||||
"""
|
||||
request = await self._prepare_llm_request(row)
|
||||
t = time.perf_counter()
|
||||
|
||||
async with self.semaphore:
|
||||
output = await self._generate_async(request)
|
||||
|
||||
time_taken = time.perf_counter() - t
|
||||
|
||||
output_data = SGLangOutputData.from_sglang_engine_output(output)
|
||||
return request, output_data.model_dump(), time_taken
|
||||
|
||||
async def _generate_async(self, request: SGLangEngineRequest) -> Any:
|
||||
"""Process a single request.
|
||||
|
||||
Args:
|
||||
request: The request.
|
||||
|
||||
Returns:
|
||||
The output of the request.
|
||||
"""
|
||||
|
||||
# Send the request to the LLM engine.
|
||||
stream = await self.engine.async_generate(
|
||||
prompt=request.prompt,
|
||||
input_ids=request.prompt_token_ids,
|
||||
sampling_params=request.params,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
# Consume the stream until the request is finished.
|
||||
async for output in stream:
|
||||
if output["meta_info"]["finish_reason"] is not None:
|
||||
output["prompt"] = request.prompt
|
||||
output["prompt_token_ids"] = request.prompt_token_ids
|
||||
return output
|
||||
|
||||
raise RuntimeError(
|
||||
"[SGLang] The request is not finished. This should not happen. Please report this issue to the Ray team."
|
||||
)
|
||||
|
||||
def shutdown(self):
|
||||
"""Shutdown the SGLang engine."""
|
||||
if hasattr(self.engine, "shutdown"):
|
||||
logger.info("Shutting down SGLang engine")
|
||||
self.engine.shutdown()
|
||||
|
||||
|
||||
class SGLangEngineStageUDF(StatefulStageUDF):
|
||||
def __init__(
|
||||
self,
|
||||
data_column: str,
|
||||
expected_input_keys: List[str],
|
||||
model: str,
|
||||
engine_kwargs: Dict[str, Any],
|
||||
task_type: TypeSGLangTaskType = SGLangTaskType.GENERATE,
|
||||
max_pending_requests: Optional[int] = None,
|
||||
):
|
||||
"""
|
||||
Initialize the SGLangEngineStageUDF.
|
||||
|
||||
Args:
|
||||
data_column: The data column name.
|
||||
expected_input_keys: The expected input keys of the stage.
|
||||
model: The path to the model to use for the SGLang engine.
|
||||
engine_kwargs: The kwargs to pass to the SGLang engine.
|
||||
task_type: The task to use for the SGLang engine (e.g., "generate", "embed", "reward").
|
||||
max_pending_requests: The maximum number of pending requests. If None,
|
||||
it will be set to a default value based on engine settings.
|
||||
"""
|
||||
super().__init__(data_column, expected_input_keys)
|
||||
self.model = model
|
||||
|
||||
# Setup SGLang engine kwargs.
|
||||
self.task_type = task_type
|
||||
self.engine_kwargs = self.normalize_engine_kwargs(engine_kwargs)
|
||||
|
||||
# Set up the max pending requests.
|
||||
# Disable the semaphore if max_pending_requests is not set.
|
||||
self.max_pending_requests = max_pending_requests or -1
|
||||
if self.max_pending_requests > 0:
|
||||
logger.info("Max pending requests is set to %d", self.max_pending_requests)
|
||||
|
||||
# Create an LLM engine.
|
||||
self.llm = SGLangEngineWrapper(
|
||||
model=self.model,
|
||||
idx_in_batch_column=self.IDX_IN_BATCH_COLUMN,
|
||||
max_pending_requests=self.max_pending_requests,
|
||||
**self.engine_kwargs,
|
||||
)
|
||||
|
||||
def normalize_engine_kwargs(
|
||||
self,
|
||||
engine_kwargs: Dict[str, Any],
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Normalize the engine kwargs.
|
||||
|
||||
Args:
|
||||
engine_kwargs: The kwargs to normalize.
|
||||
|
||||
Returns:
|
||||
The normalized kwargs.
|
||||
"""
|
||||
# Copy to avoid mutating fn_constructor_kwargs. Ray Data generates UDF
|
||||
# instance keys before __init__, so in-place changes cause KeyError.
|
||||
engine_kwargs = engine_kwargs.copy()
|
||||
|
||||
# Remove model from engine kwargs if set.
|
||||
model = engine_kwargs.pop("model", None)
|
||||
if model is not None and model != self.model:
|
||||
logger.warning(
|
||||
"The model set in engine kwargs (%s) is different from the "
|
||||
"stage (%s). Please remove 'model' from engine kwargs.",
|
||||
model,
|
||||
self.model,
|
||||
)
|
||||
return engine_kwargs
|
||||
|
||||
async def udf(self, batch: List[Dict[str, Any]]) -> AsyncIterator[Dict[str, Any]]:
|
||||
"""Run the SGLang engine.
|
||||
|
||||
Args:
|
||||
batch: A list of rows to run the SGLang engine on.
|
||||
|
||||
Yields:
|
||||
Dict[str, Any]: The response of the SGLang engine.
|
||||
"""
|
||||
batch_uuid = uuid.uuid4()
|
||||
batch_start_time = time.perf_counter()
|
||||
|
||||
tasks = [asyncio.create_task(self.llm.generate_async(row)) for row in batch]
|
||||
|
||||
for resp in asyncio.as_completed(tasks):
|
||||
request, output, time_taken_llm = await resp
|
||||
|
||||
yield {
|
||||
**output,
|
||||
"request_id": request.request_id,
|
||||
self.IDX_IN_BATCH_COLUMN: request.idx_in_batch,
|
||||
"batch_uuid": batch_uuid.hex,
|
||||
"time_taken_llm": time_taken_llm,
|
||||
"params": str(request.params),
|
||||
}
|
||||
|
||||
batch_time_taken = time.perf_counter() - batch_start_time
|
||||
logger.info(
|
||||
"[SGLang] Elapsed time for batch %s with size %d: %s",
|
||||
batch_uuid.hex,
|
||||
len(batch),
|
||||
batch_time_taken,
|
||||
)
|
||||
|
||||
def __del__(self):
|
||||
if hasattr(self, "llm"):
|
||||
self.llm.shutdown()
|
||||
|
||||
|
||||
class SGLangEngineStage(StatefulStage):
|
||||
"""
|
||||
A stage that runs SGLang engine.
|
||||
"""
|
||||
|
||||
fn: Type[StatefulStageUDF] = SGLangEngineStageUDF
|
||||
|
||||
@root_validator(pre=True)
|
||||
def post_init(cls, values: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Post-initialize the stage. Specifically,
|
||||
this function determines the num_gpus and Ray remote args
|
||||
for the .map_batches() call in this stage.
|
||||
|
||||
Args:
|
||||
values: The raw stage values.
|
||||
Returns:
|
||||
The updated values.
|
||||
"""
|
||||
map_batches_kwargs = values["map_batches_kwargs"]
|
||||
accelerator_type = map_batches_kwargs.get("accelerator_type", "")
|
||||
fn_constructor_kwargs = values["fn_constructor_kwargs"]
|
||||
engine_kwargs = fn_constructor_kwargs.get("engine_kwargs", {})
|
||||
|
||||
ray_remote_args = {}
|
||||
if accelerator_type:
|
||||
ray_remote_args["accelerator_type"] = accelerator_type
|
||||
|
||||
# Set up num_gpus required
|
||||
tp_size = engine_kwargs.get("tp_size", 1)
|
||||
dp_size = engine_kwargs.get("dp_size", 1)
|
||||
num_gpus = tp_size * dp_size
|
||||
|
||||
ray_remote_args["num_gpus"] = num_gpus
|
||||
map_batches_kwargs.update(ray_remote_args)
|
||||
return values
|
||||
|
||||
def get_required_input_keys(self) -> Dict[str, str]:
|
||||
"""The required input keys of the stage and their descriptions."""
|
||||
ret = {"prompt": "The text prompt (str)."}
|
||||
task_type = self.fn_constructor_kwargs.get("task_type", SGLangTaskType.GENERATE)
|
||||
if task_type == SGLangTaskType.GENERATE:
|
||||
ret[
|
||||
"sampling_params"
|
||||
] = "The sampling parameters. See https://docs.sglang.ai/backend/sampling_params.html for details."
|
||||
return ret
|
||||
|
||||
def get_optional_input_keys(self) -> Dict[str, str]:
|
||||
"""The optional input keys of the stage and their descriptions."""
|
||||
return {}
|
||||
@@ -0,0 +1,160 @@
|
||||
"""Tokenize and detokenize stage"""
|
||||
|
||||
from typing import Any, AsyncIterator, Dict, List, Type
|
||||
|
||||
from ray.llm._internal.batch.stages.base import (
|
||||
StatefulStage,
|
||||
StatefulStageUDF,
|
||||
)
|
||||
from ray.llm._internal.batch.utils import get_cached_tokenizer
|
||||
from ray.llm._internal.common.utils.download_utils import (
|
||||
NodeModelDownloadable,
|
||||
download_model_files,
|
||||
)
|
||||
|
||||
|
||||
class TokenizeUDF(StatefulStageUDF):
|
||||
def __init__(
|
||||
self,
|
||||
data_column: str,
|
||||
expected_input_keys: List[str],
|
||||
model: str,
|
||||
trust_remote_code: bool = False,
|
||||
):
|
||||
"""
|
||||
Initialize the TokenizeUDF.
|
||||
|
||||
Args:
|
||||
data_column: The data column name.
|
||||
expected_input_keys: The expected input keys of the stage.
|
||||
model: The model to use for tokenization.
|
||||
trust_remote_code: Whether to trust remote code when loading the model.
|
||||
"""
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
super().__init__(data_column, expected_input_keys)
|
||||
# Use EXCLUDE_SAFETENSORS for trust_remote_code models to ensure
|
||||
# Python config files are downloaded.
|
||||
download_mode = (
|
||||
NodeModelDownloadable.EXCLUDE_SAFETENSORS
|
||||
if trust_remote_code
|
||||
else NodeModelDownloadable.TOKENIZER_ONLY
|
||||
)
|
||||
model_path = download_model_files(
|
||||
model_id=model,
|
||||
mirror_config=None,
|
||||
download_model=download_mode,
|
||||
download_extra_files=False,
|
||||
)
|
||||
self.tokenizer = get_cached_tokenizer(
|
||||
AutoTokenizer.from_pretrained(
|
||||
model_path,
|
||||
trust_remote_code=trust_remote_code,
|
||||
)
|
||||
)
|
||||
|
||||
async def udf(self, batch: List[Dict[str, Any]]) -> AsyncIterator[Dict[str, Any]]:
|
||||
"""
|
||||
Tokenize the given batch.
|
||||
|
||||
Args:
|
||||
batch: A list of rows to send.
|
||||
|
||||
Yields:
|
||||
Dict[str, Any]: A generator of rows with the tokenized prompt.
|
||||
"""
|
||||
for row, prompt_token_ids in zip(
|
||||
batch,
|
||||
self.tokenizer([row["prompt"] for row in batch])["input_ids"],
|
||||
):
|
||||
yield {
|
||||
self.IDX_IN_BATCH_COLUMN: row[self.IDX_IN_BATCH_COLUMN],
|
||||
"tokenized_prompt": prompt_token_ids,
|
||||
}
|
||||
|
||||
|
||||
class TokenizeStage(StatefulStage):
|
||||
"""
|
||||
A stage that tokenizes the input.
|
||||
"""
|
||||
|
||||
fn: Type[StatefulStageUDF] = TokenizeUDF
|
||||
|
||||
def get_required_input_keys(self) -> Dict[str, str]:
|
||||
"""The required input keys of the stage and their descriptions."""
|
||||
return {"prompt": "The text prompt (str) to tokenize."}
|
||||
|
||||
|
||||
class DetokenizeUDF(StatefulStageUDF):
|
||||
def __init__(
|
||||
self,
|
||||
data_column: str,
|
||||
expected_input_keys: List[str],
|
||||
model: str,
|
||||
trust_remote_code: bool = False,
|
||||
):
|
||||
"""
|
||||
Initialize the DetokenizeUDF.
|
||||
|
||||
Args:
|
||||
data_column: The data column name.
|
||||
expected_input_keys: The expected input keys of the stage.
|
||||
model: The model to use for detokenization.
|
||||
trust_remote_code: Whether to trust remote code when loading the model.
|
||||
"""
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
super().__init__(data_column, expected_input_keys)
|
||||
# Use EXCLUDE_SAFETENSORS for trust_remote_code models to ensure
|
||||
# Python config files are downloaded.
|
||||
download_mode = (
|
||||
NodeModelDownloadable.EXCLUDE_SAFETENSORS
|
||||
if trust_remote_code
|
||||
else NodeModelDownloadable.TOKENIZER_ONLY
|
||||
)
|
||||
model_path = download_model_files(
|
||||
model_id=model,
|
||||
mirror_config=None,
|
||||
download_model=download_mode,
|
||||
download_extra_files=False,
|
||||
)
|
||||
self.tokenizer = get_cached_tokenizer(
|
||||
AutoTokenizer.from_pretrained(
|
||||
model_path,
|
||||
trust_remote_code=trust_remote_code,
|
||||
)
|
||||
)
|
||||
|
||||
async def udf(self, batch: List[Dict[str, Any]]) -> AsyncIterator[Dict[str, Any]]:
|
||||
"""
|
||||
Detokenize the given batch.
|
||||
|
||||
Args:
|
||||
batch: A list of rows to send.
|
||||
|
||||
Yields:
|
||||
Dict[str, Any]: A generator of rows with the detokenized generated text.
|
||||
"""
|
||||
for row, generated_text in zip(
|
||||
batch,
|
||||
self.tokenizer.batch_decode(
|
||||
[row["generated_tokens"] for row in batch],
|
||||
skip_special_tokens=True,
|
||||
),
|
||||
):
|
||||
yield {
|
||||
self.IDX_IN_BATCH_COLUMN: row[self.IDX_IN_BATCH_COLUMN],
|
||||
"generated_text": generated_text,
|
||||
}
|
||||
|
||||
|
||||
class DetokenizeStage(StatefulStage):
|
||||
"""
|
||||
A stage that detokenizes the input.
|
||||
"""
|
||||
|
||||
fn: Type[StatefulStageUDF] = DetokenizeUDF
|
||||
|
||||
def get_required_input_keys(self) -> Dict[str, str]:
|
||||
"""The required input keys of the stage and their descriptions."""
|
||||
return {"generated_tokens": "A list of generated tokens (int) to detokenize."}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,68 @@
|
||||
"""Utility functions for batch processing."""
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Any, Union
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
|
||||
|
||||
AnyTokenizer = Union["PreTrainedTokenizer", "PreTrainedTokenizerFast", Any]
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def get_cached_tokenizer(tokenizer: AnyTokenizer) -> AnyTokenizer:
|
||||
"""Get tokenizer with cached properties.
|
||||
This will patch the tokenizer object in place.
|
||||
By default, transformers will recompute multiple tokenizer properties
|
||||
each time they are called, leading to a significant slowdown. This
|
||||
function caches these properties for faster access.
|
||||
Args:
|
||||
tokenizer: The tokenizer object.
|
||||
Returns:
|
||||
The patched tokenizer object.
|
||||
"""
|
||||
chat_template = getattr(tokenizer, "chat_template", None)
|
||||
# For VLM, the text tokenizer is wrapped by a processor.
|
||||
if hasattr(tokenizer, "tokenizer"):
|
||||
tokenizer = tokenizer.tokenizer
|
||||
# Some VLM's tokenizer has chat_template attribute (e.g. Qwen/Qwen2-VL-7B-Instruct),
|
||||
# however some other VLM's tokenizer does not have chat_template attribute (e.g.
|
||||
# mistral-community/pixtral-12b). Therefore, we cache the processor's chat_template.
|
||||
if chat_template is None:
|
||||
chat_template = getattr(tokenizer, "chat_template", None)
|
||||
|
||||
tokenizer_all_special_ids = set(tokenizer.all_special_ids)
|
||||
tokenizer_all_special_tokens = set(tokenizer.all_special_tokens)
|
||||
# all_special_tokens_extended is removed in transformers v5, used in latest
|
||||
# SGLang version. We require this SGLang version bc it's ABI compatible with
|
||||
# PyTorch 2.9, which is installed by vLLM.
|
||||
# TODO(seiji) remove the attribute completely once vLLM moves to transformers v5.
|
||||
tokenizer_all_special_tokens_extended = getattr(
|
||||
tokenizer, "all_special_tokens_extended", None
|
||||
)
|
||||
tokenizer_len = len(tokenizer)
|
||||
|
||||
class CachedTokenizer(tokenizer.__class__): # type: ignore
|
||||
@property
|
||||
def all_special_ids(self):
|
||||
return tokenizer_all_special_ids
|
||||
|
||||
@property
|
||||
def all_special_tokens(self):
|
||||
return tokenizer_all_special_tokens
|
||||
|
||||
@property
|
||||
def all_special_tokens_extended(self):
|
||||
return tokenizer_all_special_tokens_extended
|
||||
|
||||
@property
|
||||
def chat_template(self):
|
||||
return chat_template
|
||||
|
||||
def __len__(self):
|
||||
return tokenizer_len
|
||||
|
||||
CachedTokenizer.__name__ = f"Cached{tokenizer.__class__.__name__}"
|
||||
|
||||
tokenizer.__class__ = CachedTokenizer
|
||||
return tokenizer
|
||||
@@ -0,0 +1,31 @@
|
||||
from typing import Type, TypeVar
|
||||
|
||||
import yaml
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
ModelT = TypeVar("ModelT", bound=BaseModel)
|
||||
|
||||
|
||||
class BaseModelExtended(BaseModel):
|
||||
# NOTE(edoakes): Pydantic protects the namespace `model_` by default and prints
|
||||
# warnings if you define fields with that prefix. However, we added such fields
|
||||
# before this behavior existed. To avoid spamming user-facing logs, we mark the
|
||||
# namespace as not protected. This means we need to be careful about overriding
|
||||
# internal attributes starting with `model_`.
|
||||
# See: https://github.com/anyscale/ray-llm/issues/1425
|
||||
model_config = ConfigDict(
|
||||
protected_namespaces=tuple(),
|
||||
extra="forbid",
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def parse_yaml(cls: Type[ModelT], file, **kwargs) -> ModelT:
|
||||
kwargs.setdefault("Loader", yaml.SafeLoader)
|
||||
dict_args = yaml.load(file, **kwargs)
|
||||
return cls.model_validate(dict_args)
|
||||
|
||||
@classmethod
|
||||
def from_file(cls: Type[ModelT], path: str, **kwargs) -> ModelT:
|
||||
"""Load a model from a YAML file path."""
|
||||
with open(path, "r") as f:
|
||||
return cls.parse_yaml(f, **kwargs)
|
||||
@@ -0,0 +1,144 @@
|
||||
import asyncio
|
||||
import inspect
|
||||
import logging
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple, Type, Union
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.llm._internal.common.utils.download_utils import NodeModelDownloadable
|
||||
from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class CallbackCtx:
|
||||
"""
|
||||
Context object passed to all callback hooks.
|
||||
Callbacks can read and modify fields as needed.
|
||||
"""
|
||||
|
||||
worker_node_download_model: Optional["NodeModelDownloadable"] = None
|
||||
"""Model download configuration for worker nodes. Used to specify how
|
||||
models should be downloaded and cached on worker nodes in distributed
|
||||
deployments."""
|
||||
placement_group: Optional[Any] = None
|
||||
"""Ray placement group for resource allocation and scheduling. Controls
|
||||
where and how resources are allocated across the cluster."""
|
||||
runtime_env: Optional[Dict[str, Any]] = None
|
||||
"""Runtime environment configuration for the Ray workers. Includes
|
||||
dependencies, environment variables, and other runtime settings."""
|
||||
custom_data: Dict[str, Any] = field(default_factory=dict)
|
||||
"""Flexible dictionary for callback-specific state and data. Allows
|
||||
callbacks to store and share custom information during initialization."""
|
||||
run_init_node: bool = True
|
||||
"""Whether to run model downloads during initialization. Set to False
|
||||
to skip downloading models."""
|
||||
|
||||
|
||||
class CallbackBase:
|
||||
"""Base class for custom initialization implementations.
|
||||
|
||||
This class defines the interface for custom initialization logic
|
||||
for LLMEngine to be called in node_initialization.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
llm_config: "LLMConfig",
|
||||
raise_error_on_callback: bool = True,
|
||||
ctx_kwargs: Optional[Dict[str, Any]] = None,
|
||||
**kwargs,
|
||||
):
|
||||
self.raise_error_on_callback = raise_error_on_callback
|
||||
self.kwargs = kwargs
|
||||
self.llm_config = llm_config
|
||||
|
||||
# Create and store CallbackCtx internally using ctx_kwargs
|
||||
ctx_kwargs = ctx_kwargs or {}
|
||||
self.ctx = CallbackCtx(**ctx_kwargs)
|
||||
|
||||
async def on_before_node_init(self) -> None:
|
||||
"""Called before node initialization begins."""
|
||||
pass
|
||||
|
||||
async def on_after_node_init(self) -> None:
|
||||
"""Called after node initialization completes."""
|
||||
pass
|
||||
|
||||
def on_before_download_model_files_distributed(self) -> None:
|
||||
"""Called before model files are downloaded on each node."""
|
||||
pass
|
||||
|
||||
def _get_method(self, method_name: str) -> Tuple[Callable, bool]:
|
||||
"""Get a callback method."""
|
||||
if not hasattr(self, method_name):
|
||||
raise AttributeError(
|
||||
f"Callback {type(self).__name__} does not have method '{method_name}'"
|
||||
)
|
||||
return getattr(self, method_name), inspect.iscoroutinefunction(
|
||||
getattr(self, method_name)
|
||||
)
|
||||
|
||||
def _handle_callback_error(self, method_name: str, e: Exception) -> None:
|
||||
if self.raise_error_on_callback:
|
||||
raise Exception(
|
||||
f"Error running callback method '{method_name}' on {type(self).__name__}: {str(e)}"
|
||||
) from e
|
||||
else:
|
||||
logger.error(
|
||||
f"Error running callback method '{method_name}' on {type(self).__name__}: {str(e)}"
|
||||
)
|
||||
|
||||
async def run_callback(self, method_name: str) -> None:
|
||||
"""Run a callback method either synchronously or asynchronously.
|
||||
|
||||
Args:
|
||||
method_name: The name of the method to call on the callback
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
method, is_async = self._get_method(method_name)
|
||||
|
||||
try:
|
||||
if is_async:
|
||||
await method()
|
||||
else:
|
||||
method()
|
||||
except Exception as e:
|
||||
self._handle_callback_error(method_name, e)
|
||||
|
||||
def run_callback_sync(self, method_name: str) -> None:
|
||||
"""Run a callback method synchronously
|
||||
|
||||
Args:
|
||||
method_name: The name of the method to call on the callback
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
method, is_async = self._get_method(method_name)
|
||||
|
||||
try:
|
||||
if is_async:
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
loop.run_until_complete(method())
|
||||
except RuntimeError:
|
||||
asyncio.run(method())
|
||||
else:
|
||||
method()
|
||||
except Exception as e:
|
||||
self._handle_callback_error(method_name, e)
|
||||
|
||||
|
||||
@dataclass
|
||||
class CallbackConfig:
|
||||
"""Configuration for the callback to be used in LLMConfig"""
|
||||
|
||||
callback_class: Union[str, Type[CallbackBase]] = CallbackBase
|
||||
"""Class to use for the callback. Can be custom user defined class"""
|
||||
callback_kwargs: Dict[str, Any] = field(default_factory=dict)
|
||||
"""Keyword arguments to pass to the Callback class at construction."""
|
||||
raise_error_on_callback: bool = True
|
||||
"""Whether to raise an error if a callback method fails."""
|
||||
@@ -0,0 +1,86 @@
|
||||
import logging
|
||||
import time
|
||||
from typing import Any, List, Tuple
|
||||
|
||||
from pydantic import BaseModel, field_validator
|
||||
|
||||
from .base import CallbackBase
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class CloudDownloaderConfig(BaseModel):
|
||||
"""Model for validating CloudDownloader configuration."""
|
||||
|
||||
paths: List[Tuple[str, str]]
|
||||
|
||||
@field_validator("paths")
|
||||
@classmethod
|
||||
def validate_paths(cls, v: List[Tuple[str, str]]) -> List[Tuple[str, str]]:
|
||||
# Supported cloud storage URI schemes
|
||||
valid_schemes = ("s3://", "gs://", "abfss://", "azure://")
|
||||
|
||||
for i, (cloud_uri, _) in enumerate(v):
|
||||
if not any(cloud_uri.startswith(scheme) for scheme in valid_schemes):
|
||||
raise ValueError(
|
||||
f"paths[{i}][0] (cloud_uri) must start with one of {valid_schemes}, "
|
||||
f"got '{cloud_uri}'"
|
||||
)
|
||||
return v
|
||||
|
||||
|
||||
class CloudDownloader(CallbackBase):
|
||||
"""Callback that downloads files from cloud storage before model files are downloaded.
|
||||
|
||||
This callback expects self.kwargs to contain a 'paths' field which should be
|
||||
a list of tuples, where each tuple contains (cloud_uri, local_path) strings.
|
||||
|
||||
Supported cloud storage URIs: s3://, gs://, abfss://, azure://
|
||||
|
||||
Example:
|
||||
```
|
||||
from ray.llm._internal.common.callbacks.cloud_downloader import CloudDownloader
|
||||
from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
|
||||
config = LLMConfig(
|
||||
...
|
||||
callback_config={
|
||||
"callback_class": CloudDownloader,
|
||||
"callback_kwargs": {
|
||||
"paths": [
|
||||
("s3://bucket/path/to/file.txt", "/local/path/to/file.txt"),
|
||||
("gs://bucket/path/to/file.txt", "/local/path/to/file.txt"),
|
||||
]
|
||||
}
|
||||
}
|
||||
...
|
||||
)
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, **kwargs: Any) -> None:
|
||||
"""Initialize the CloudDownloader callback.
|
||||
|
||||
Args:
|
||||
**kwargs: Keyword arguments passed to the callback as a dictionary.
|
||||
Must contain a 'paths' field with a list of (cloud_uri, local_path) tuples.
|
||||
"""
|
||||
super().__init__(**kwargs)
|
||||
|
||||
# Validate configuration using Pydantic
|
||||
if "paths" not in self.kwargs:
|
||||
raise ValueError("CloudDownloader requires 'paths' field in kwargs")
|
||||
|
||||
CloudDownloaderConfig.model_validate(self.kwargs)
|
||||
|
||||
def on_before_download_model_files_distributed(self) -> None:
|
||||
"""Download files from cloud storage to local paths before model files are downloaded."""
|
||||
from ray.llm._internal.common.utils.cloud_utils import CloudFileSystem
|
||||
|
||||
paths = self.kwargs["paths"]
|
||||
start_time = time.monotonic()
|
||||
for cloud_uri, local_path in paths:
|
||||
CloudFileSystem.download_files(path=local_path, bucket_uri=cloud_uri)
|
||||
end_time = time.monotonic()
|
||||
logger.info(
|
||||
f"CloudDownloader: Files downloaded in {end_time - start_time} seconds"
|
||||
)
|
||||
@@ -0,0 +1,13 @@
|
||||
"""
|
||||
Generic constants for common utilities.
|
||||
|
||||
These constants are used by generic utilities and should not contain
|
||||
serve-specific or batch-specific values.
|
||||
"""
|
||||
|
||||
# Cloud object caching timeouts (in seconds)
|
||||
CLOUD_OBJECT_EXISTS_EXPIRE_S = 300 # 5 minutes
|
||||
CLOUD_OBJECT_MISSING_EXPIRE_S = 30 # 30 seconds
|
||||
|
||||
# LoRA adapter configuration file name
|
||||
LORA_ADAPTER_CONFIG_NAME = "adapter_config.json"
|
||||
@@ -0,0 +1,60 @@
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
|
||||
def maybe_apply_llm_deployment_config_defaults(
|
||||
defaults: Dict[str, Any],
|
||||
user_deployment_config: Optional[Dict[str, Any]],
|
||||
) -> Dict[str, Any]:
|
||||
"""Apply defaults and merge with user-provided deployment config.
|
||||
|
||||
If the user has explicitly set 'num_replicas' in their deployment config,
|
||||
we remove 'autoscaling_config' from the defaults since Ray Serve
|
||||
does not allow both to be set simultaneously. Then merges the defaults
|
||||
with the user config.
|
||||
|
||||
Args:
|
||||
defaults: The default deployment options dictionary.
|
||||
user_deployment_config: The user-provided deployment configuration.
|
||||
|
||||
Returns:
|
||||
The merged deployment options with conflicts resolved.
|
||||
"""
|
||||
if user_deployment_config and "num_replicas" in user_deployment_config:
|
||||
defaults = defaults.copy()
|
||||
defaults.pop("autoscaling_config", None)
|
||||
return deep_merge_dicts(defaults, user_deployment_config or {})
|
||||
|
||||
|
||||
def deep_merge_dicts(base: Dict[str, Any], override: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""
|
||||
Merge two dictionaries hierarchically, creating a new dictionary without modifying inputs.
|
||||
|
||||
For each key:
|
||||
- If the key exists in both dicts and both values are dicts, recursively merge them
|
||||
- Otherwise, the value from override takes precedence
|
||||
|
||||
Args:
|
||||
base: The base dictionary
|
||||
override: The dictionary with values that should override the base
|
||||
|
||||
Returns:
|
||||
A new merged dictionary
|
||||
|
||||
Example:
|
||||
>>> base = {"a": 1, "b": {"c": 2, "d": 3}}
|
||||
>>> override = {"b": {"c": 10}, "e": 5}
|
||||
>>> result = deep_merge_dicts(base, override)
|
||||
>>> result
|
||||
{'a': 1, 'b': {'c': 10, 'd': 3}, 'e': 5}
|
||||
"""
|
||||
result = base.copy()
|
||||
|
||||
for key, value in override.items():
|
||||
if key in result and isinstance(result[key], dict) and isinstance(value, dict):
|
||||
# Recursively merge nested dictionaries
|
||||
result[key] = deep_merge_dicts(result[key], value)
|
||||
else:
|
||||
# Override the value (or add new key)
|
||||
result[key] = value
|
||||
|
||||
return result
|
||||
@@ -0,0 +1,14 @@
|
||||
"""Fatal engine error definitions shared by serve and batch layers."""
|
||||
|
||||
from typing import Tuple, Type
|
||||
|
||||
# vLLM fatal errors that should always be re-raised, never swallowed.
|
||||
# EngineDeadError indicates the vLLM engine process has crashed and is
|
||||
# unrecoverable — all subsequent requests would fail anyway.
|
||||
VLLM_FATAL_ERRORS: Tuple[Type[Exception], ...] = ()
|
||||
try:
|
||||
from vllm.v1.engine.exceptions import EngineDeadError
|
||||
|
||||
VLLM_FATAL_ERRORS = (EngineDeadError,)
|
||||
except ImportError:
|
||||
pass
|
||||
@@ -0,0 +1,49 @@
|
||||
"""
|
||||
Generic model definitions for common utilities.
|
||||
|
||||
These models represent generic concepts that can be used by both
|
||||
serve and batch components.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
from functools import partial
|
||||
from typing import Awaitable, Callable, TypeVar
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
# DiskMultiplexConfig removed - it's serve-specific and belongs in serve/configs/server_models.py
|
||||
|
||||
|
||||
class GlobalIdManager:
|
||||
"""Thread-safe global ID manager for assigning unique IDs."""
|
||||
|
||||
def __init__(self):
|
||||
self._counter = 0
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def next(self) -> int:
|
||||
"""Get the next unique ID."""
|
||||
with self._lock:
|
||||
self._counter += 1
|
||||
return self._counter
|
||||
|
||||
|
||||
# Global instance
|
||||
global_id_manager = GlobalIdManager()
|
||||
|
||||
|
||||
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
|
||||
@@ -0,0 +1,38 @@
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from ray._common.filters import CoreContextFilter
|
||||
|
||||
|
||||
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 logger is already setup and has handlers.
|
||||
|
||||
Args:
|
||||
logger_name: logger name used to get the logger.
|
||||
"""
|
||||
logger = logging.getLogger(logger_name)
|
||||
llm_logger = logging.getLogger("ray.llm")
|
||||
|
||||
# Skip setup if the logger already has handlers setup or if the parent (Data
|
||||
# logger) has handlers.
|
||||
if not (logger.handlers or llm_logger.handlers):
|
||||
# Set up stream handler, which logs to console as plaintext.
|
||||
stream_handler = logging.StreamHandler()
|
||||
stream_handler.addFilter(CoreContextFilter())
|
||||
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 Data logger.
|
||||
|
||||
Loggers by default are logging to stdout, and are expected to be scraped by an
|
||||
external process.
|
||||
"""
|
||||
logger_name = f"ray.llm.{name}"
|
||||
_setup_logger(logger_name)
|
||||
return logging.getLogger(logger_name)
|
||||
@@ -0,0 +1,26 @@
|
||||
import logging
|
||||
|
||||
from ray._private.ray_logging.filters import CoreContextFilter
|
||||
from ray._private.ray_logging.formatters import JSONFormatter
|
||||
|
||||
|
||||
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.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,48 @@
|
||||
"""Utilities for logging."""
|
||||
|
||||
import logging
|
||||
|
||||
import ray
|
||||
|
||||
|
||||
def disable_vllm_custom_ops_logger_on_cpu_nodes():
|
||||
"""This disables a log line in the "vllm._custom_ops" logger on CPU nodes.
|
||||
|
||||
vllm._custom_ops is automatically imported when vllm is imported. It checks
|
||||
for CUDA binaries that don't exist in CPU-only nodes. This makes rayllm
|
||||
raise a scary-looking (but harmless) warning when imported on CPU nodes,
|
||||
such as when running the generate_config.py script or running the
|
||||
build-app task.
|
||||
"""
|
||||
|
||||
class SkipVLLMWarningFilter(logging.Filter):
|
||||
def filter(self, record: logging.LogRecord):
|
||||
"""Only allow CRITICAL logs from the datasets/config.py file."""
|
||||
|
||||
log_fragment = "Failed to import from vllm._C with"
|
||||
|
||||
return log_fragment not in record.getMessage()
|
||||
|
||||
if not ray.is_initialized() or len(ray.get_gpu_ids()) == 0:
|
||||
logging.getLogger("vllm._custom_ops").addFilter(SkipVLLMWarningFilter())
|
||||
|
||||
|
||||
def disable_datasets_logger():
|
||||
"""This disables "datasets" logs from its "config.py" file.
|
||||
|
||||
Upon import, rayllm imports vllm which calls datasets. The datasets package
|
||||
emits a log from its config.py file.
|
||||
|
||||
The file that emits this log uses the root "datasets" logger, so we use a
|
||||
filter to prevent logs from only the config.py file.
|
||||
"""
|
||||
|
||||
class SkipDatasetsConfigLogFilter(logging.Filter):
|
||||
def filter(self, record: logging.LogRecord):
|
||||
"""Only allow CRITICAL logs from the datasets/config.py file."""
|
||||
return (
|
||||
record.levelno >= logging.CRITICAL
|
||||
or "datasets/config.py" not in record.pathname
|
||||
)
|
||||
|
||||
logging.getLogger("datasets").addFilter(SkipDatasetsConfigLogFilter())
|
||||
@@ -0,0 +1,45 @@
|
||||
"""Utilities for telemetry."""
|
||||
|
||||
from threading import Lock
|
||||
from typing import Callable
|
||||
|
||||
DEFAULT_GPU_TYPE = "UNSPECIFIED"
|
||||
|
||||
|
||||
class Once:
|
||||
"""Execute a function exactly once and block all callers until the function returns
|
||||
|
||||
Same as golang's `sync.Once <https://pkg.go.dev/sync#Once>`_
|
||||
|
||||
Took this directly from OpenTelemetry's Python SDK:
|
||||
Ref: https://github.com/open-telemetry/opentelemetry-python/blob
|
||||
/c6fab7d4c339dc5bf9eb9ef2723caad09d69bfca/opentelemetry-api/src/opentelemetry
|
||||
/util/_once.py
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._lock = Lock()
|
||||
self._done = False
|
||||
|
||||
def do_once(self, func: Callable[[], None]) -> bool:
|
||||
"""Execute ``func`` if it hasn't been executed or return.
|
||||
|
||||
Will block until ``func`` has been called by one thread.
|
||||
|
||||
Args:
|
||||
func: The function to execute exactly once.
|
||||
|
||||
Returns:
|
||||
Whether or not ``func`` was executed in this call
|
||||
"""
|
||||
|
||||
# fast path, try to avoid locking
|
||||
if self._done:
|
||||
return False
|
||||
|
||||
with self._lock:
|
||||
if not self._done:
|
||||
func()
|
||||
self._done = True
|
||||
return True
|
||||
return False
|
||||
@@ -0,0 +1,7 @@
|
||||
"""Temporary runtime patches of third-party libraries (currently vLLM).
|
||||
|
||||
Each module here is a stopgap for behavior Ray Serve LLM needs before it exists
|
||||
upstream. The goal is for this package to trend toward empty: every patch states
|
||||
its removal condition in its module docstring and tracks the upstream work that
|
||||
will retire it. Patches are deleted as soon as that upstream lands.
|
||||
"""
|
||||
@@ -0,0 +1,2 @@
|
||||
"""Runtime patches of vLLM internals. See the parent package policy: each patch
|
||||
is temporary and tracked to an upstream issue/PR for removal."""
|
||||
@@ -0,0 +1,95 @@
|
||||
"""Decode-stage reuse of prefill's prompt token ids (P/D tokenize-once).
|
||||
|
||||
A P/D chat prompt is otherwise tokenized once per stage. ``install()`` wraps
|
||||
``BaseRenderer.tokenize_prompts_async`` to inject the ids prefill already produced.
|
||||
``reuse_prompt_token_ids`` publishes those ids per async task. vLLM skips the encode
|
||||
when ``prompt_token_ids`` is present, so the rest of the pipeline runs unchanged.
|
||||
``install()`` is idempotent and fails safe.
|
||||
|
||||
Decode tokenization must run within the reuse block, on the same async task or one
|
||||
spawned from it, so the contextvar reaches it.
|
||||
|
||||
Temporary patch. The intended end state is native pre-tokenized input on the
|
||||
chat-completions path (a ``prompt_token_ids`` field on ``ChatCompletionRequest``),
|
||||
at which point decode passes the ids through a request field and this wrap is
|
||||
deleted. See vllm-project/vllm#22817 (token-in/token-out) for the upstream
|
||||
direction.
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import contextvars
|
||||
import functools
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Per-async-task reused prompt token ids. contextvars don't leak across tasks, so
|
||||
# concurrent requests can't cross-contaminate.
|
||||
_reused_token_ids: contextvars.ContextVar = contextvars.ContextVar(
|
||||
"pd_reused_prompt_token_ids", default=None
|
||||
)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def reuse_prompt_token_ids(token_ids):
|
||||
"""Publish ``token_ids`` so the chat render inside this block skips tokenize.
|
||||
|
||||
No-op when ``token_ids`` is falsy, so callers need no separate enabled check.
|
||||
"""
|
||||
if not token_ids:
|
||||
yield
|
||||
return
|
||||
_reused_token_ids.set(list(token_ids))
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
# Use set(None), not reset(token). The finally may run in a different
|
||||
# Context than the enter during off-task generator finalization, where
|
||||
# reset() would raise.
|
||||
_reused_token_ids.set(None)
|
||||
|
||||
|
||||
def install() -> bool:
|
||||
"""Wrap ``BaseRenderer.tokenize_prompts_async`` to honor the contextvar.
|
||||
|
||||
Idempotent and fails safe. Returns False and leaves tokenization untouched when
|
||||
vLLM's renderer is missing or differs, so it never crashes startup.
|
||||
"""
|
||||
try:
|
||||
from vllm.renderers.base import BaseRenderer
|
||||
|
||||
orig = getattr(BaseRenderer, "tokenize_prompts_async", None)
|
||||
if orig is None:
|
||||
logger.debug("pd-tokenize-once: BaseRenderer.tokenize_prompts_async absent")
|
||||
return False
|
||||
if getattr(orig, "_pd_tokonce_wrapped", False):
|
||||
return True
|
||||
|
||||
@functools.wraps(orig)
|
||||
async def tokenize_prompts_async(self, prompts, params, *args, **kwargs):
|
||||
ids = _reused_token_ids.get()
|
||||
# Inject the reused ids into the lone rendered prompt so vLLM skips the
|
||||
# encode but still preserves multi_modal_data and runs detok/validation.
|
||||
# Anything else falls through to a normal tokenize.
|
||||
if (
|
||||
ids
|
||||
and len(prompts) == 1
|
||||
and isinstance(prompts[0], dict)
|
||||
and "prompt_token_ids" not in prompts[0]
|
||||
and "prompt_embeds" not in prompts[0]
|
||||
and "encoder_prompt" not in prompts[0]
|
||||
):
|
||||
prompts = [{**prompts[0], "prompt_token_ids": ids}]
|
||||
return await orig(self, prompts, params, *args, **kwargs)
|
||||
|
||||
tokenize_prompts_async._pd_tokonce_wrapped = True
|
||||
BaseRenderer.tokenize_prompts_async = tokenize_prompts_async
|
||||
except Exception as e: # pragma: no cover - defensive
|
||||
logger.debug("pd-tokenize-once: install failed (%s)", e)
|
||||
return False
|
||||
|
||||
logger.info(
|
||||
"pd-tokenize-once: wrapped BaseRenderer.tokenize_prompts_async "
|
||||
"(decode stage will reuse prefill's prompt token ids)"
|
||||
)
|
||||
return True
|
||||
@@ -0,0 +1,79 @@
|
||||
from typing import List, Literal, Optional
|
||||
|
||||
from pydantic import ConfigDict, Field, model_validator
|
||||
|
||||
from ray.llm._internal.common.base_pydantic import BaseModelExtended
|
||||
|
||||
|
||||
class BundleConfig(BaseModelExtended):
|
||||
"""Configuration for a single placement group bundle.
|
||||
|
||||
Resource counts are floats to align with Ray's internal resource
|
||||
representation, which supports fractional values (e.g. GPU=0.5).
|
||||
Both CPU and GPU default to 0.0 — the schema does not inject hidden
|
||||
resource requests. Extra fields are allowed for custom Ray resources (e.g. TPU,
|
||||
accelerator_type:L4).
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
CPU: float = Field(
|
||||
default=0.0,
|
||||
ge=0.0,
|
||||
description="The number of CPUs per bundle.",
|
||||
)
|
||||
GPU: float = Field(
|
||||
default=0.0,
|
||||
ge=0.0,
|
||||
description="The number of GPUs per bundle.",
|
||||
)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_extra_resources(self):
|
||||
"""Ensure custom resources (TPU, accelerator_type, etc.) are non-negative."""
|
||||
extra_resources = self.model_extra
|
||||
if not extra_resources:
|
||||
return self
|
||||
for key, value in extra_resources.items():
|
||||
if not isinstance(value, (int, float)):
|
||||
raise ValueError(
|
||||
f"Resource '{key}' must be a number, got {type(value).__name__}"
|
||||
)
|
||||
if value < 0:
|
||||
raise ValueError(f"Resource '{key}' must be non-negative, got {value}")
|
||||
return self
|
||||
|
||||
|
||||
class PlacementGroupConfig(BaseModelExtended):
|
||||
"""Configuration for placement group."""
|
||||
|
||||
bundle_per_worker: Optional[BundleConfig] = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Resource bundle specification for each worker. "
|
||||
"Auto-replicated based on tensor_parallel_size * pipeline_parallel_size. "
|
||||
"Cannot be used together with 'bundles'."
|
||||
),
|
||||
)
|
||||
bundles: Optional[List[BundleConfig]] = Field(
|
||||
default=None, description="List of resource bundles"
|
||||
)
|
||||
strategy: Literal["PACK", "SPREAD", "STRICT_PACK", "STRICT_SPREAD"] = Field(
|
||||
default="PACK", description="Placement group strategy"
|
||||
)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_bundle_options(self):
|
||||
if self.bundle_per_worker is not None and self.bundles is not None:
|
||||
raise ValueError(
|
||||
"Cannot specify both 'bundle_per_worker' and 'bundles' in "
|
||||
"placement_group_config. Use 'bundle_per_worker' for simple "
|
||||
"per-worker resource specification (auto-replicated by tp*pp), "
|
||||
"or 'bundles' for full control."
|
||||
)
|
||||
if self.bundle_per_worker is None and self.bundles is None:
|
||||
raise ValueError(
|
||||
"placement_group_config must specify either 'bundle_per_worker' "
|
||||
"or 'bundles'."
|
||||
)
|
||||
return self
|
||||
@@ -0,0 +1,29 @@
|
||||
"""Cloud filesystem module for provider-specific implementations.
|
||||
|
||||
This module provides a unified interface for cloud storage operations across
|
||||
different providers (S3, GCS, Azure) while allowing provider-specific optimizations.
|
||||
"""
|
||||
|
||||
from ray.llm._internal.common.utils.cloud_filesystem.azure_filesystem import (
|
||||
AzureFileSystem,
|
||||
)
|
||||
from ray.llm._internal.common.utils.cloud_filesystem.base import (
|
||||
BaseCloudFileSystem,
|
||||
)
|
||||
from ray.llm._internal.common.utils.cloud_filesystem.gcs_filesystem import (
|
||||
GCSFileSystem,
|
||||
)
|
||||
from ray.llm._internal.common.utils.cloud_filesystem.pyarrow_filesystem import (
|
||||
PyArrowFileSystem,
|
||||
)
|
||||
from ray.llm._internal.common.utils.cloud_filesystem.s3_filesystem import (
|
||||
S3FileSystem,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"BaseCloudFileSystem",
|
||||
"PyArrowFileSystem",
|
||||
"GCSFileSystem",
|
||||
"AzureFileSystem",
|
||||
"S3FileSystem",
|
||||
]
|
||||
@@ -0,0 +1,82 @@
|
||||
"""Azure-specific filesystem implementation.
|
||||
|
||||
This module provides an Azure-specific implementation that delegates to PyArrowFileSystem.
|
||||
This maintains backward compatibility while allowing for future optimizations using
|
||||
native Azure tools (azcopy, azure-storage-blob SDK).
|
||||
"""
|
||||
|
||||
from typing import List, Optional, Union
|
||||
|
||||
from ray.llm._internal.common.utils.cloud_filesystem.base import BaseCloudFileSystem
|
||||
from ray.llm._internal.common.utils.cloud_filesystem.pyarrow_filesystem import (
|
||||
PyArrowFileSystem,
|
||||
)
|
||||
|
||||
|
||||
class AzureFileSystem(BaseCloudFileSystem):
|
||||
"""Azure-specific implementation of cloud filesystem operations.
|
||||
|
||||
**Note**: This implementation currently delegates to PyArrowFileSystem to maintain
|
||||
stability. Optimized implementation using azure-storage-blob SDK and azcopy
|
||||
will be added in a future PR.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def get_file(
|
||||
object_uri: str, decode_as_utf_8: bool = True
|
||||
) -> Optional[Union[str, bytes]]:
|
||||
"""Download a file from cloud storage into memory.
|
||||
|
||||
Args:
|
||||
object_uri: URI of the file (abfss:// or azure://)
|
||||
decode_as_utf_8: If True, decode the file as UTF-8
|
||||
|
||||
Returns:
|
||||
File contents as string or bytes, or None if file doesn't exist
|
||||
"""
|
||||
return PyArrowFileSystem.get_file(object_uri, decode_as_utf_8)
|
||||
|
||||
@staticmethod
|
||||
def list_subfolders(folder_uri: str) -> List[str]:
|
||||
"""List the immediate subfolders in a cloud directory.
|
||||
|
||||
Args:
|
||||
folder_uri: URI of the directory (abfss:// or azure://)
|
||||
|
||||
Returns:
|
||||
List of subfolder names (without trailing slashes)
|
||||
"""
|
||||
return PyArrowFileSystem.list_subfolders(folder_uri)
|
||||
|
||||
@staticmethod
|
||||
def download_files(
|
||||
path: str,
|
||||
bucket_uri: str,
|
||||
substrings_to_include: Optional[List[str]] = None,
|
||||
suffixes_to_exclude: Optional[List[str]] = None,
|
||||
) -> None:
|
||||
"""Download files from cloud storage to a local directory.
|
||||
|
||||
Args:
|
||||
path: Local directory where files will be downloaded
|
||||
bucket_uri: URI of cloud directory
|
||||
substrings_to_include: Only include files containing these substrings
|
||||
suffixes_to_exclude: Exclude certain files from download (e.g .safetensors)
|
||||
"""
|
||||
PyArrowFileSystem.download_files(
|
||||
path, bucket_uri, substrings_to_include, suffixes_to_exclude
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def upload_files(
|
||||
local_path: str,
|
||||
bucket_uri: str,
|
||||
) -> None:
|
||||
"""Upload files to cloud storage.
|
||||
|
||||
Args:
|
||||
local_path: The local path of the files to upload.
|
||||
bucket_uri: The bucket uri to upload the files to, must start with
|
||||
`abfss://` or `azure://`.
|
||||
"""
|
||||
PyArrowFileSystem.upload_files(local_path, bucket_uri)
|
||||
@@ -0,0 +1,81 @@
|
||||
"""Abstract base class for cloud filesystem implementations.
|
||||
|
||||
This module defines the interface that all cloud storage provider implementations
|
||||
must follow, ensuring consistency across different providers while allowing
|
||||
provider-specific optimizations.
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import List, Optional, Union
|
||||
|
||||
|
||||
class BaseCloudFileSystem(ABC):
|
||||
"""Abstract base class for cloud filesystem implementations.
|
||||
|
||||
This class defines the interface that all cloud storage provider implementations
|
||||
must implement. Provider-specific classes (S3FileSystem, GCSFileSystem, etc.)
|
||||
will inherit from this base class and provide optimized implementations for
|
||||
their respective cloud storage platforms.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
@abstractmethod
|
||||
def get_file(
|
||||
object_uri: str, decode_as_utf_8: bool = True
|
||||
) -> Optional[Union[str, bytes]]:
|
||||
"""Download a file from cloud storage into memory.
|
||||
|
||||
Args:
|
||||
object_uri: URI of the file (s3://, gs://, abfss://, or azure://)
|
||||
decode_as_utf_8: If True, decode the file as UTF-8
|
||||
|
||||
Returns:
|
||||
File contents as string or bytes, or None if file doesn't exist
|
||||
"""
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
@abstractmethod
|
||||
def list_subfolders(folder_uri: str) -> List[str]:
|
||||
"""List the immediate subfolders in a cloud directory.
|
||||
|
||||
Args:
|
||||
folder_uri: URI of the directory (s3://, gs://, abfss://, or azure://)
|
||||
|
||||
Returns:
|
||||
List of subfolder names (without trailing slashes)
|
||||
"""
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
@abstractmethod
|
||||
def download_files(
|
||||
path: str,
|
||||
bucket_uri: str,
|
||||
substrings_to_include: Optional[List[str]] = None,
|
||||
suffixes_to_exclude: Optional[List[str]] = None,
|
||||
) -> None:
|
||||
"""Download files from cloud storage to a local directory.
|
||||
|
||||
Args:
|
||||
path: Local directory where files will be downloaded
|
||||
bucket_uri: URI of cloud directory
|
||||
substrings_to_include: Only include files containing these substrings
|
||||
suffixes_to_exclude: Exclude certain files from download (e.g .safetensors)
|
||||
"""
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
@abstractmethod
|
||||
def upload_files(
|
||||
local_path: str,
|
||||
bucket_uri: str,
|
||||
) -> None:
|
||||
"""Upload files to cloud storage.
|
||||
|
||||
Args:
|
||||
local_path: The local path of the files to upload.
|
||||
bucket_uri: The bucket uri to upload the files to, must start with
|
||||
`s3://`, `gs://`, `abfss://`, or `azure://`.
|
||||
"""
|
||||
pass
|
||||
@@ -0,0 +1,81 @@
|
||||
"""GCS-specific filesystem implementation.
|
||||
|
||||
This module provides a GCS-specific implementation.
|
||||
This maintains backward compatibility while allowing for future optimizations using
|
||||
native GCS tools (gsutil, google-cloud-storage SDK).
|
||||
"""
|
||||
|
||||
from typing import List, Optional, Union
|
||||
|
||||
from ray.llm._internal.common.utils.cloud_filesystem.base import BaseCloudFileSystem
|
||||
from ray.llm._internal.common.utils.cloud_filesystem.pyarrow_filesystem import (
|
||||
PyArrowFileSystem,
|
||||
)
|
||||
|
||||
|
||||
class GCSFileSystem(BaseCloudFileSystem):
|
||||
"""GCS-specific implementation of cloud filesystem operations.
|
||||
|
||||
**Note**: This implementation currently delegates to PyArrowFileSystem to maintain
|
||||
stability. Optimized implementation using google-cloud-storage SDK and gsutil
|
||||
will be added in a future PR.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def get_file(
|
||||
object_uri: str, decode_as_utf_8: bool = True
|
||||
) -> Optional[Union[str, bytes]]:
|
||||
"""Download a file from cloud storage into memory.
|
||||
|
||||
Args:
|
||||
object_uri: URI of the file (gs://)
|
||||
decode_as_utf_8: If True, decode the file as UTF-8
|
||||
|
||||
Returns:
|
||||
File contents as string or bytes, or None if file doesn't exist
|
||||
"""
|
||||
return PyArrowFileSystem.get_file(object_uri, decode_as_utf_8)
|
||||
|
||||
@staticmethod
|
||||
def list_subfolders(folder_uri: str) -> List[str]:
|
||||
"""List the immediate subfolders in a cloud directory.
|
||||
|
||||
Args:
|
||||
folder_uri: URI of the directory (gs://)
|
||||
|
||||
Returns:
|
||||
List of subfolder names (without trailing slashes)
|
||||
"""
|
||||
return PyArrowFileSystem.list_subfolders(folder_uri)
|
||||
|
||||
@staticmethod
|
||||
def download_files(
|
||||
path: str,
|
||||
bucket_uri: str,
|
||||
substrings_to_include: Optional[List[str]] = None,
|
||||
suffixes_to_exclude: Optional[List[str]] = None,
|
||||
) -> None:
|
||||
"""Download files from cloud storage to a local directory.
|
||||
|
||||
Args:
|
||||
path: Local directory where files will be downloaded
|
||||
bucket_uri: URI of cloud directory
|
||||
substrings_to_include: Only include files containing these substrings
|
||||
suffixes_to_exclude: Exclude certain files from download (e.g .safetensors)
|
||||
"""
|
||||
PyArrowFileSystem.download_files(
|
||||
path, bucket_uri, substrings_to_include, suffixes_to_exclude
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def upload_files(
|
||||
local_path: str,
|
||||
bucket_uri: str,
|
||||
) -> None:
|
||||
"""Upload files to cloud storage.
|
||||
|
||||
Args:
|
||||
local_path: The local path of the files to upload.
|
||||
bucket_uri: The bucket uri to upload the files to, must start with `gs://`.
|
||||
"""
|
||||
PyArrowFileSystem.upload_files(local_path, bucket_uri)
|
||||
@@ -0,0 +1,389 @@
|
||||
"""PyArrow-based filesystem implementation for cloud storage.
|
||||
|
||||
This module provides a PyArrow-based implementation of the cloud filesystem
|
||||
interface, supporting S3, GCS, and Azure storage providers.
|
||||
"""
|
||||
|
||||
import os
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import List, Optional, Tuple, Union
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import pyarrow.fs as pa_fs
|
||||
|
||||
from ray.llm._internal.common.observability.logging import get_logger
|
||||
from ray.llm._internal.common.utils.cloud_filesystem.base import BaseCloudFileSystem
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class PyArrowFileSystem(BaseCloudFileSystem):
|
||||
"""PyArrow-based implementation of cloud filesystem operations.
|
||||
|
||||
This class provides a unified interface for cloud storage operations using
|
||||
PyArrow's filesystem abstraction. It supports S3, GCS, and Azure storage
|
||||
providers.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def get_fs_and_path(object_uri: str) -> Tuple[pa_fs.FileSystem, str]:
|
||||
"""Get the appropriate filesystem and path from a URI.
|
||||
|
||||
Args:
|
||||
object_uri: URI of the file (s3://, gs://, abfss://, or azure://)
|
||||
If URI contains 'anonymous@', anonymous access is used.
|
||||
Example: s3://anonymous@bucket/path
|
||||
|
||||
Returns:
|
||||
Tuple of (filesystem, path)
|
||||
"""
|
||||
|
||||
if object_uri.startswith("pyarrow-"):
|
||||
object_uri = object_uri[8:]
|
||||
|
||||
anonymous = False
|
||||
# Check for anonymous access pattern (only for S3/GCS)
|
||||
# e.g. s3://anonymous@bucket/path
|
||||
if "@" in object_uri and not (
|
||||
object_uri.startswith("abfss://") or object_uri.startswith("azure://")
|
||||
):
|
||||
parts = object_uri.split("@", 1)
|
||||
# Check if the first part ends with "anonymous"
|
||||
if parts[0].endswith("anonymous"):
|
||||
anonymous = True
|
||||
# Remove the anonymous@ part, keeping the scheme
|
||||
scheme = parts[0].split("://")[0]
|
||||
object_uri = f"{scheme}://{parts[1]}"
|
||||
|
||||
if object_uri.startswith("s3://"):
|
||||
endpoint = os.getenv("AWS_ENDPOINT_URL_S3", None)
|
||||
virtual_hosted_style = os.getenv("AWS_S3_ADDRESSING_STYLE", None)
|
||||
fs = pa_fs.S3FileSystem(
|
||||
anonymous=anonymous,
|
||||
endpoint_override=endpoint,
|
||||
force_virtual_addressing=(virtual_hosted_style == "virtual"),
|
||||
)
|
||||
path = object_uri[5:] # Remove "s3://"
|
||||
elif object_uri.startswith("gs://"):
|
||||
fs = pa_fs.GcsFileSystem(anonymous=anonymous)
|
||||
path = object_uri[5:] # Remove "gs://"
|
||||
elif object_uri.startswith("abfss://"):
|
||||
fs, path = PyArrowFileSystem._create_abfss_filesystem(object_uri)
|
||||
elif object_uri.startswith("azure://"):
|
||||
fs, path = PyArrowFileSystem._create_azure_filesystem(object_uri)
|
||||
else:
|
||||
raise ValueError(f"Unsupported URI scheme: {object_uri}")
|
||||
|
||||
return fs, path
|
||||
|
||||
@staticmethod
|
||||
def _create_azure_filesystem(object_uri: str) -> Tuple[pa_fs.FileSystem, str]:
|
||||
"""Create an Azure filesystem for Azure Blob Storage or ABFSS.
|
||||
|
||||
Args:
|
||||
object_uri: Azure URI (azure://container@account.blob.core.windows.net/path or
|
||||
abfss://container@account.dfs.core.windows.net/path)
|
||||
|
||||
Returns:
|
||||
Tuple of (PyArrow FileSystem, path without scheme prefix)
|
||||
|
||||
Raises:
|
||||
ImportError: If required dependencies are not installed.
|
||||
ValueError: If the Azure URI format is invalid.
|
||||
"""
|
||||
try:
|
||||
import adlfs
|
||||
from azure.identity import DefaultAzureCredential
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"You must `pip install adlfs azure-identity` "
|
||||
"to use Azure/ABFSS URIs. "
|
||||
"Note that these must be preinstalled on all nodes in the Ray cluster."
|
||||
)
|
||||
|
||||
# Parse and validate the Azure URI
|
||||
parsed = urlparse(object_uri)
|
||||
scheme = parsed.scheme.lower()
|
||||
|
||||
# Validate URI format: scheme://container@account.domain/path
|
||||
if not parsed.netloc or "@" not in parsed.netloc:
|
||||
raise ValueError(
|
||||
f"Invalid {scheme.upper()} URI format - missing container@account: {object_uri}"
|
||||
)
|
||||
|
||||
container_part, hostname_part = parsed.netloc.split("@", 1)
|
||||
|
||||
# Validate container name (must be non-empty)
|
||||
if not container_part:
|
||||
raise ValueError(
|
||||
f"Invalid {scheme.upper()} URI format - empty container name: {object_uri}"
|
||||
)
|
||||
|
||||
# Validate hostname format based on scheme
|
||||
valid_hostname = False
|
||||
if scheme == "abfss":
|
||||
valid_hostname = hostname_part.endswith(".dfs.core.windows.net")
|
||||
expected_domains = ".dfs.core.windows.net"
|
||||
elif scheme == "azure":
|
||||
valid_hostname = hostname_part.endswith(
|
||||
".blob.core.windows.net"
|
||||
) or hostname_part.endswith(".dfs.core.windows.net")
|
||||
expected_domains = ".blob.core.windows.net or .dfs.core.windows.net"
|
||||
|
||||
if not hostname_part or not valid_hostname:
|
||||
raise ValueError(
|
||||
f"Invalid {scheme.upper()} URI format - invalid hostname (must end with {expected_domains}): {object_uri}"
|
||||
)
|
||||
|
||||
# Extract and validate account name
|
||||
azure_storage_account_name = hostname_part.split(".")[0]
|
||||
if not azure_storage_account_name:
|
||||
raise ValueError(
|
||||
f"Invalid {scheme.upper()} URI format - empty account name: {object_uri}"
|
||||
)
|
||||
|
||||
# Create the adlfs filesystem
|
||||
adlfs_fs = adlfs.AzureBlobFileSystem(
|
||||
account_name=azure_storage_account_name,
|
||||
credential=DefaultAzureCredential(),
|
||||
)
|
||||
|
||||
# Wrap with PyArrow's PyFileSystem for compatibility
|
||||
fs = pa_fs.PyFileSystem(pa_fs.FSSpecHandler(adlfs_fs))
|
||||
|
||||
# Return the path without the scheme prefix
|
||||
path = f"{container_part}{parsed.path}"
|
||||
|
||||
return fs, path
|
||||
|
||||
@staticmethod
|
||||
def _create_abfss_filesystem(object_uri: str) -> Tuple[pa_fs.FileSystem, str]:
|
||||
"""Create an ABFSS filesystem for Azure Data Lake Storage Gen2.
|
||||
|
||||
This is a wrapper around _create_azure_filesystem for backward compatibility.
|
||||
|
||||
Args:
|
||||
object_uri: ABFSS URI (abfss://container@account.dfs.core.windows.net/path)
|
||||
|
||||
Returns:
|
||||
Tuple of (PyArrow FileSystem, path without abfss:// prefix)
|
||||
"""
|
||||
return PyArrowFileSystem._create_azure_filesystem(object_uri)
|
||||
|
||||
@staticmethod
|
||||
def _filter_files(
|
||||
fs: pa_fs.FileSystem,
|
||||
source_path: str,
|
||||
destination_path: str,
|
||||
substrings_to_include: Optional[List[str]] = None,
|
||||
suffixes_to_exclude: Optional[List[str]] = None,
|
||||
) -> List[Tuple[str, str]]:
|
||||
"""Filter files from cloud storage based on inclusion and exclusion criteria.
|
||||
|
||||
Args:
|
||||
fs: PyArrow filesystem instance
|
||||
source_path: Source path in cloud storage
|
||||
destination_path: Local destination path
|
||||
substrings_to_include: Only include files containing these substrings
|
||||
suffixes_to_exclude: Exclude files ending with these suffixes
|
||||
|
||||
Returns:
|
||||
List of tuples containing (source_file_path, destination_file_path)
|
||||
"""
|
||||
file_selector = pa_fs.FileSelector(source_path, recursive=True)
|
||||
file_infos = fs.get_file_info(file_selector)
|
||||
|
||||
path_pairs = []
|
||||
for file_info in file_infos:
|
||||
if file_info.type != pa_fs.FileType.File:
|
||||
continue
|
||||
|
||||
rel_path = file_info.path[len(source_path) :].lstrip("/")
|
||||
|
||||
# Apply filters
|
||||
if substrings_to_include:
|
||||
if not any(
|
||||
substring in rel_path for substring in substrings_to_include
|
||||
):
|
||||
continue
|
||||
|
||||
if suffixes_to_exclude:
|
||||
if any(rel_path.endswith(suffix) for suffix in suffixes_to_exclude):
|
||||
continue
|
||||
|
||||
path_pairs.append(
|
||||
(file_info.path, os.path.join(destination_path, rel_path))
|
||||
)
|
||||
|
||||
return path_pairs
|
||||
|
||||
@staticmethod
|
||||
def get_file(
|
||||
object_uri: str, decode_as_utf_8: bool = True
|
||||
) -> Optional[Union[str, bytes]]:
|
||||
"""Download a file from cloud storage into memory.
|
||||
|
||||
Args:
|
||||
object_uri: URI of the file (s3://, gs://, abfss://, or azure://)
|
||||
decode_as_utf_8: If True, decode the file as UTF-8
|
||||
|
||||
Returns:
|
||||
File contents as string or bytes, or None if file doesn't exist
|
||||
"""
|
||||
try:
|
||||
fs, path = PyArrowFileSystem.get_fs_and_path(object_uri)
|
||||
|
||||
# Check if file exists
|
||||
if not fs.get_file_info(path).type == pa_fs.FileType.File:
|
||||
logger.info(f"URI {object_uri} does not exist.")
|
||||
return None
|
||||
|
||||
# Read file
|
||||
with fs.open_input_file(path) as f:
|
||||
body = f.read()
|
||||
|
||||
if decode_as_utf_8:
|
||||
body = body.decode("utf-8")
|
||||
return body
|
||||
except Exception as e:
|
||||
logger.warning(f"Error reading {object_uri}: {e}")
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def list_subfolders(folder_uri: str) -> List[str]:
|
||||
"""List the immediate subfolders in a cloud directory.
|
||||
|
||||
Args:
|
||||
folder_uri: URI of the directory (s3://, gs://, abfss://, or azure://)
|
||||
|
||||
Returns:
|
||||
List of subfolder names (without trailing slashes)
|
||||
"""
|
||||
# Ensure that the folder_uri has a trailing slash.
|
||||
folder_uri = f"{folder_uri.rstrip('/')}/"
|
||||
|
||||
try:
|
||||
fs, path = PyArrowFileSystem.get_fs_and_path(folder_uri)
|
||||
|
||||
# List directory contents
|
||||
file_infos = fs.get_file_info(pa_fs.FileSelector(path, recursive=False))
|
||||
|
||||
# Filter for directories and extract subfolder names
|
||||
subfolders = []
|
||||
for file_info in file_infos:
|
||||
if file_info.type == pa_fs.FileType.Directory:
|
||||
# Extract just the subfolder name without the full path
|
||||
subfolder = os.path.basename(file_info.path.rstrip("/"))
|
||||
subfolders.append(subfolder)
|
||||
|
||||
return subfolders
|
||||
except Exception as e:
|
||||
logger.error(f"Error listing subfolders in {folder_uri}: {e}")
|
||||
return []
|
||||
|
||||
@staticmethod
|
||||
def download_files(
|
||||
path: str,
|
||||
bucket_uri: str,
|
||||
substrings_to_include: Optional[List[str]] = None,
|
||||
suffixes_to_exclude: Optional[List[str]] = None,
|
||||
max_concurrency: int = 10,
|
||||
chunk_size: int = 64 * 1024 * 1024,
|
||||
) -> None:
|
||||
"""Download files from cloud storage to a local directory.
|
||||
|
||||
Args:
|
||||
path: Local directory where files will be downloaded
|
||||
bucket_uri: URI of cloud directory
|
||||
substrings_to_include: Only include files containing these substrings
|
||||
suffixes_to_exclude: Exclude certain files from download
|
||||
max_concurrency: Maximum number of concurrent files to download (default: 10)
|
||||
chunk_size: Size of transfer chunks (default: 64MB)
|
||||
"""
|
||||
try:
|
||||
fs, source_path = PyArrowFileSystem.get_fs_and_path(bucket_uri)
|
||||
|
||||
# Ensure destination exists
|
||||
os.makedirs(path, exist_ok=True)
|
||||
|
||||
# If no filters, use direct copy_files
|
||||
if not substrings_to_include and not suffixes_to_exclude:
|
||||
pa_fs.copy_files(
|
||||
source=source_path,
|
||||
destination=path,
|
||||
source_filesystem=fs,
|
||||
destination_filesystem=pa_fs.LocalFileSystem(),
|
||||
use_threads=True,
|
||||
chunk_size=chunk_size,
|
||||
)
|
||||
return
|
||||
|
||||
# List and filter files
|
||||
files_to_download = PyArrowFileSystem._filter_files(
|
||||
fs, source_path, path, substrings_to_include, suffixes_to_exclude
|
||||
)
|
||||
|
||||
if not files_to_download:
|
||||
logger.info("Filters do not match any of the files, skipping download")
|
||||
return
|
||||
|
||||
def download_single_file(file_paths):
|
||||
source_file_path, dest_file_path = file_paths
|
||||
# Create destination directory if needed
|
||||
dest_dir = os.path.dirname(dest_file_path)
|
||||
if dest_dir:
|
||||
os.makedirs(dest_dir, exist_ok=True)
|
||||
|
||||
# Use PyArrow's copy_files for individual files,
|
||||
pa_fs.copy_files(
|
||||
source=source_file_path,
|
||||
destination=dest_file_path,
|
||||
source_filesystem=fs,
|
||||
destination_filesystem=pa_fs.LocalFileSystem(),
|
||||
use_threads=True,
|
||||
chunk_size=chunk_size,
|
||||
)
|
||||
return dest_file_path
|
||||
|
||||
max_workers = min(max_concurrency, len(files_to_download))
|
||||
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
futures = [
|
||||
executor.submit(download_single_file, file_paths)
|
||||
for file_paths in files_to_download
|
||||
]
|
||||
|
||||
for future in futures:
|
||||
try:
|
||||
future.result()
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to download file: {e}")
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"Error downloading files from {bucket_uri}: {e}")
|
||||
raise
|
||||
|
||||
@staticmethod
|
||||
def upload_files(
|
||||
local_path: str,
|
||||
bucket_uri: str,
|
||||
) -> None:
|
||||
"""Upload files to cloud storage.
|
||||
|
||||
Args:
|
||||
local_path: The local path of the files to upload.
|
||||
bucket_uri: The bucket uri to upload the files to, must start with
|
||||
`s3://`, `gs://`, `abfss://`, or `azure://`.
|
||||
"""
|
||||
try:
|
||||
fs, dest_path = PyArrowFileSystem.get_fs_and_path(bucket_uri)
|
||||
|
||||
pa_fs.copy_files(
|
||||
source=local_path,
|
||||
destination=dest_path,
|
||||
source_filesystem=pa_fs.LocalFileSystem(),
|
||||
destination_filesystem=fs,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.exception(f"Error uploading files to {bucket_uri}: {e}")
|
||||
raise
|
||||
@@ -0,0 +1,397 @@
|
||||
"""S3-specific filesystem implementation using boto3.
|
||||
|
||||
This module provides an S3-specific implementation that uses boto3 (AWS SDK for Python)
|
||||
for reliable and efficient S3 operations.
|
||||
"""
|
||||
|
||||
import os
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from pathlib import Path
|
||||
from typing import Any, List, Optional, Union
|
||||
|
||||
try:
|
||||
import boto3
|
||||
from botocore import UNSIGNED
|
||||
from botocore.config import Config
|
||||
except ImportError:
|
||||
boto3 = None # type: ignore[assignment]
|
||||
UNSIGNED = None # type: ignore[assignment]
|
||||
Config = None # type: ignore[assignment]
|
||||
|
||||
from ray.llm._internal.common.observability.logging import get_logger
|
||||
from ray.llm._internal.common.utils.cloud_filesystem.base import BaseCloudFileSystem
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def _check_boto3() -> None:
|
||||
"""Raise a clear error if boto3/botocore are not installed."""
|
||||
if boto3 is None:
|
||||
raise ImportError(
|
||||
"boto3 and botocore are required for S3 operations but are not installed. "
|
||||
"Install them with: pip install boto3"
|
||||
)
|
||||
|
||||
|
||||
class S3FileSystem(BaseCloudFileSystem):
|
||||
"""S3-specific implementation of cloud filesystem operations using boto3.
|
||||
|
||||
This implementation uses boto3 (AWS SDK for Python) for reliable and efficient
|
||||
operations with S3 storage.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def _parse_s3_uri(uri: str) -> tuple[str, str, bool]:
|
||||
"""Parse S3 URI into bucket and key.
|
||||
|
||||
Args:
|
||||
uri: S3 URI (e.g., s3://bucket/path/to/object or s3://anonymous@bucket/path/to/object)
|
||||
|
||||
Returns:
|
||||
Tuple of (bucket_name, key, is_anonymous)
|
||||
|
||||
Raises:
|
||||
ValueError: If URI is not a valid S3 URI
|
||||
"""
|
||||
# Check if anonymous@ prefix exists
|
||||
is_anonymous = False
|
||||
if uri.startswith("s3://anonymous@"):
|
||||
is_anonymous = True
|
||||
uri = uri.replace("s3://anonymous@", "s3://", 1)
|
||||
|
||||
if not uri.startswith("s3://"):
|
||||
raise ValueError(f"Invalid S3 URI: {uri}")
|
||||
|
||||
# Remove s3:// prefix and split into bucket and key
|
||||
path = uri[5:] # Remove "s3://"
|
||||
parts = path.split("/", 1)
|
||||
bucket = parts[0]
|
||||
key = parts[1] if len(parts) > 1 else ""
|
||||
|
||||
return bucket, key, is_anonymous
|
||||
|
||||
@staticmethod
|
||||
def _get_s3_client(max_pool_connections: int = 50, anonymous: bool = False):
|
||||
"""Create a new S3 client instance with connection pooling.
|
||||
|
||||
Args:
|
||||
max_pool_connections: Maximum number of connections in the pool.
|
||||
Should be >= max_workers for optimal performance.
|
||||
anonymous: Whether to use anonymous access to S3.
|
||||
|
||||
Returns:
|
||||
boto3 S3 client with connection pooling configured
|
||||
"""
|
||||
_check_boto3()
|
||||
# Configure connection pooling for better concurrent performance
|
||||
config = Config(
|
||||
max_pool_connections=max_pool_connections,
|
||||
# Retry configuration for transient failures
|
||||
retries={
|
||||
"max_attempts": 3,
|
||||
"mode": "adaptive", # Adapts retry behavior based on error type
|
||||
},
|
||||
# TCP keepalive helps with long-running connections
|
||||
tcp_keepalive=True,
|
||||
signature_version=UNSIGNED if anonymous else None,
|
||||
)
|
||||
|
||||
return boto3.client("s3", config=config)
|
||||
|
||||
@staticmethod
|
||||
def get_file(
|
||||
object_uri: str, decode_as_utf_8: bool = True
|
||||
) -> Optional[Union[str, bytes]]:
|
||||
"""Download a file from cloud storage into memory.
|
||||
|
||||
Args:
|
||||
object_uri: URI of the file (s3://)
|
||||
decode_as_utf_8: If True, decode the file as UTF-8
|
||||
|
||||
Returns:
|
||||
File contents as string or bytes, or None if file doesn't exist
|
||||
"""
|
||||
_check_boto3()
|
||||
try:
|
||||
bucket, key, is_anonymous = S3FileSystem._parse_s3_uri(object_uri)
|
||||
s3_client = S3FileSystem._get_s3_client(anonymous=is_anonymous)
|
||||
|
||||
# Download file directly into memory
|
||||
response = s3_client.get_object(Bucket=bucket, Key=key)
|
||||
body = response["Body"].read()
|
||||
|
||||
if decode_as_utf_8:
|
||||
return body.decode("utf-8")
|
||||
return body
|
||||
except Exception as e:
|
||||
logger.error(f"Error reading {object_uri}: {e}")
|
||||
|
||||
@staticmethod
|
||||
def list_subfolders(folder_uri: str) -> List[str]:
|
||||
"""List the immediate subfolders in a cloud directory.
|
||||
|
||||
Args:
|
||||
folder_uri: URI of the directory (s3://)
|
||||
|
||||
Returns:
|
||||
List of subfolder names (without trailing slashes)
|
||||
"""
|
||||
_check_boto3()
|
||||
try:
|
||||
bucket, prefix, is_anonymous = S3FileSystem._parse_s3_uri(folder_uri)
|
||||
|
||||
# Ensure that the prefix has a trailing slash
|
||||
if prefix and not prefix.endswith("/"):
|
||||
prefix = f"{prefix}/"
|
||||
|
||||
s3_client = S3FileSystem._get_s3_client(anonymous=is_anonymous)
|
||||
|
||||
# List objects with delimiter to get only immediate subfolders
|
||||
response = s3_client.list_objects_v2(
|
||||
Bucket=bucket, Prefix=prefix, Delimiter="/"
|
||||
)
|
||||
|
||||
subfolders = []
|
||||
# CommonPrefixes contains the subdirectories
|
||||
for common_prefix in response.get("CommonPrefixes", []):
|
||||
folder_path = common_prefix["Prefix"]
|
||||
# Extract the folder name from the full prefix
|
||||
# Remove the parent prefix and trailing slash
|
||||
folder_name = folder_path[len(prefix) :].rstrip("/")
|
||||
if folder_name:
|
||||
subfolders.append(folder_name)
|
||||
|
||||
return subfolders
|
||||
except Exception as e:
|
||||
logger.error(f"Error listing subfolders in {folder_uri}: {e}")
|
||||
return []
|
||||
|
||||
@staticmethod
|
||||
def _calculate_optimal_workers(
|
||||
num_files: int, total_size: int, default_max: int = 100, default_min: int = 10
|
||||
) -> int:
|
||||
"""Calculate optimal number of workers based on file characteristics.
|
||||
|
||||
Args:
|
||||
num_files: Number of files to download
|
||||
total_size: Total size of all files in bytes
|
||||
default_max: Maximum workers to cap at
|
||||
default_min: Minimum workers to use
|
||||
|
||||
Returns:
|
||||
Optimal number of workers between default_min and default_max
|
||||
"""
|
||||
if num_files == 0:
|
||||
return default_min
|
||||
|
||||
avg_file_size = total_size / num_files if total_size > 0 else 0
|
||||
|
||||
# Strategy: More workers for smaller files, fewer for larger files
|
||||
if avg_file_size < 1024 * 1024: # < 1MB (small files)
|
||||
# Use more workers for many small files
|
||||
workers = min(num_files, default_max)
|
||||
elif avg_file_size < 10 * 1024 * 1024: # 1-10MB (medium files)
|
||||
# Use moderate workers
|
||||
workers = min(num_files // 2, default_max // 2)
|
||||
else: # > 10MB (large files)
|
||||
# Use fewer workers since each download is bandwidth-intensive
|
||||
workers = min(20, num_files)
|
||||
|
||||
# Ensure workers is between min and max
|
||||
return max(default_min, min(workers, default_max))
|
||||
|
||||
@staticmethod
|
||||
def _download_single_file(
|
||||
s3_client: Any, bucket: str, key: str, local_file_path: str
|
||||
) -> tuple[str, bool]:
|
||||
"""Download a single file from S3.
|
||||
|
||||
Args:
|
||||
s3_client: Shared boto3 S3 client
|
||||
bucket: S3 bucket name
|
||||
key: S3 object key
|
||||
local_file_path: Local path where file will be saved
|
||||
|
||||
Returns:
|
||||
Tuple of (key, success)
|
||||
"""
|
||||
try:
|
||||
# Create parent directories if needed
|
||||
os.makedirs(os.path.dirname(local_file_path), exist_ok=True)
|
||||
|
||||
s3_client.download_file(bucket, key, local_file_path)
|
||||
logger.debug(f"Downloaded {key} to {local_file_path}")
|
||||
return key, True
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to download {key}: {e}")
|
||||
return key, False
|
||||
|
||||
@staticmethod
|
||||
def download_files(
|
||||
path: str,
|
||||
bucket_uri: str,
|
||||
substrings_to_include: Optional[List[str]] = None,
|
||||
suffixes_to_exclude: Optional[List[str]] = None,
|
||||
max_workers: Optional[int] = None,
|
||||
) -> None:
|
||||
"""Download files from cloud storage to a local directory concurrently.
|
||||
|
||||
Args:
|
||||
path: Local directory where files will be downloaded
|
||||
bucket_uri: URI of cloud directory
|
||||
substrings_to_include: Only include files containing these substrings
|
||||
suffixes_to_exclude: Exclude certain files from download (e.g .safetensors)
|
||||
max_workers: Maximum number of concurrent downloads. If None, automatically
|
||||
calculated based on file count and sizes (min: 10, max: 100)
|
||||
"""
|
||||
try:
|
||||
bucket, prefix, is_anonymous = S3FileSystem._parse_s3_uri(bucket_uri)
|
||||
|
||||
# Ensure the destination directory exists
|
||||
os.makedirs(path, exist_ok=True)
|
||||
|
||||
# Ensure prefix has trailing slash for directory listing
|
||||
if prefix and not prefix.endswith("/"):
|
||||
prefix = f"{prefix}/"
|
||||
|
||||
# Create initial client for listing (will recreate with proper pool size later)
|
||||
s3_client = S3FileSystem._get_s3_client(anonymous=is_anonymous)
|
||||
|
||||
# List all objects in the bucket with the given prefix
|
||||
paginator = s3_client.get_paginator("list_objects_v2")
|
||||
pages = paginator.paginate(Bucket=bucket, Prefix=prefix)
|
||||
|
||||
# Collect all files to download and track total size
|
||||
files_to_download = []
|
||||
total_size = 0
|
||||
|
||||
for page in pages:
|
||||
for obj in page.get("Contents", []):
|
||||
key = obj["Key"]
|
||||
size = obj.get("Size", 0)
|
||||
|
||||
# Skip if it's a directory marker
|
||||
if key.endswith("/"):
|
||||
continue
|
||||
|
||||
# Get the relative path (remove the prefix)
|
||||
relative_path = key[len(prefix) :]
|
||||
|
||||
# Apply include filters
|
||||
if substrings_to_include:
|
||||
if not any(
|
||||
substr in relative_path for substr in substrings_to_include
|
||||
):
|
||||
continue
|
||||
|
||||
# Apply exclude filters
|
||||
if suffixes_to_exclude:
|
||||
if any(
|
||||
relative_path.endswith(suffix.lstrip("*"))
|
||||
for suffix in suffixes_to_exclude
|
||||
):
|
||||
continue
|
||||
|
||||
# Construct local file path
|
||||
local_file_path = os.path.join(path, relative_path)
|
||||
files_to_download.append((bucket, key, local_file_path))
|
||||
total_size += size
|
||||
|
||||
# Download files concurrently
|
||||
if not files_to_download:
|
||||
logger.info(f"No files matching filters to download from {bucket_uri}")
|
||||
return
|
||||
|
||||
# Dynamically calculate workers if not provided
|
||||
if max_workers is None:
|
||||
max_workers = S3FileSystem._calculate_optimal_workers(
|
||||
num_files=len(files_to_download),
|
||||
total_size=total_size,
|
||||
default_max=100,
|
||||
default_min=10,
|
||||
)
|
||||
|
||||
# Create shared client with proper connection pool size for downloads
|
||||
s3_client = S3FileSystem._get_s3_client(
|
||||
max_pool_connections=max_workers + 10, anonymous=is_anonymous
|
||||
)
|
||||
|
||||
failed_downloads = []
|
||||
|
||||
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
# Submit all download tasks with shared client
|
||||
future_to_key = {
|
||||
executor.submit(
|
||||
S3FileSystem._download_single_file,
|
||||
s3_client, # Pass shared client to each worker
|
||||
bucket,
|
||||
key,
|
||||
local_path,
|
||||
): key
|
||||
for bucket, key, local_path in files_to_download
|
||||
}
|
||||
|
||||
# Process completed downloads
|
||||
for future in as_completed(future_to_key):
|
||||
key, success = future.result()
|
||||
if not success:
|
||||
failed_downloads.append(key)
|
||||
|
||||
# Report any failures
|
||||
if failed_downloads:
|
||||
logger.error(
|
||||
f"Failed to download {len(failed_downloads)} files: {failed_downloads[:5]}..."
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"Error downloading files from {bucket_uri}: {e}")
|
||||
raise
|
||||
|
||||
@staticmethod
|
||||
def upload_files(
|
||||
local_path: str,
|
||||
bucket_uri: str,
|
||||
) -> None:
|
||||
"""Upload files to cloud storage.
|
||||
|
||||
Args:
|
||||
local_path: The local path of the files to upload.
|
||||
bucket_uri: The bucket uri to upload the files to, must start with `s3://`.
|
||||
"""
|
||||
try:
|
||||
bucket, prefix, is_anonymous = S3FileSystem._parse_s3_uri(bucket_uri)
|
||||
|
||||
# Ensure prefix has trailing slash for directory upload
|
||||
if prefix and not prefix.endswith("/"):
|
||||
prefix = f"{prefix}/"
|
||||
|
||||
s3_client = S3FileSystem._get_s3_client(anonymous=is_anonymous)
|
||||
|
||||
local_path_obj = Path(local_path)
|
||||
|
||||
# Walk through the local directory and upload each file
|
||||
if local_path_obj.is_file():
|
||||
# Upload a single file
|
||||
file_name = local_path_obj.name
|
||||
s3_key = f"{prefix}{file_name}" if prefix else file_name
|
||||
s3_client.upload_file(str(local_path_obj), bucket, s3_key)
|
||||
logger.debug(f"Uploaded {local_path_obj} to s3://{bucket}/{s3_key}")
|
||||
elif local_path_obj.is_dir():
|
||||
# Upload directory recursively
|
||||
for file_path in local_path_obj.rglob("*"):
|
||||
if file_path.is_file():
|
||||
# Calculate relative path from local_path
|
||||
relative_path = file_path.relative_to(local_path_obj)
|
||||
# Construct S3 key
|
||||
s3_key = f"{prefix}{relative_path.as_posix()}"
|
||||
# Upload file
|
||||
s3_client.upload_file(str(file_path), bucket, s3_key)
|
||||
logger.debug(f"Uploaded {file_path} to s3://{bucket}/{s3_key}")
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Path {local_path} does not exist or is not a file/directory"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"Error uploading files to {bucket_uri}: {e}")
|
||||
raise
|
||||
@@ -0,0 +1,546 @@
|
||||
import asyncio
|
||||
import inspect
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import (
|
||||
Any,
|
||||
Awaitable,
|
||||
Callable,
|
||||
Dict,
|
||||
List,
|
||||
NamedTuple,
|
||||
Optional,
|
||||
TypeVar,
|
||||
Union,
|
||||
)
|
||||
|
||||
from pydantic import Field, field_validator
|
||||
|
||||
from ray.llm._internal.common.base_pydantic import BaseModelExtended
|
||||
from ray.llm._internal.common.observability.logging import get_logger
|
||||
from ray.llm._internal.common.utils.cloud_filesystem import (
|
||||
AzureFileSystem,
|
||||
GCSFileSystem,
|
||||
PyArrowFileSystem,
|
||||
S3FileSystem,
|
||||
)
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def is_remote_path(path: str) -> bool:
|
||||
"""Check if the path is a remote path.
|
||||
|
||||
Args:
|
||||
path: The path to check.
|
||||
|
||||
Returns:
|
||||
True if the path is a remote path, False otherwise.
|
||||
"""
|
||||
return (
|
||||
path.startswith("s3://")
|
||||
or path.startswith("gs://")
|
||||
or path.startswith("abfss://")
|
||||
or path.startswith("azure://")
|
||||
or path.startswith("pyarrow-")
|
||||
)
|
||||
|
||||
|
||||
class ExtraFiles(BaseModelExtended):
|
||||
bucket_uri: str
|
||||
destination_path: str
|
||||
|
||||
|
||||
class CloudMirrorConfig(BaseModelExtended):
|
||||
"""Unified mirror config for cloud storage (S3, GCS, or Azure).
|
||||
|
||||
Args:
|
||||
bucket_uri: URI of the bucket (s3://, gs://, abfss://, or azure://)
|
||||
extra_files: Additional files to download
|
||||
"""
|
||||
|
||||
bucket_uri: Optional[str] = None
|
||||
extra_files: List[ExtraFiles] = Field(default_factory=list)
|
||||
|
||||
@field_validator("bucket_uri")
|
||||
@classmethod
|
||||
def check_uri_format(cls, value):
|
||||
if value is None:
|
||||
return value
|
||||
|
||||
if not is_remote_path(value):
|
||||
raise ValueError(
|
||||
f'Got invalid value "{value}" for bucket_uri. '
|
||||
'Expected a URI that starts with "s3://", "gs://", "abfss://", or "azure://".'
|
||||
)
|
||||
return value
|
||||
|
||||
@property
|
||||
def storage_type(self) -> str:
|
||||
"""Returns the storage type ('s3', 'gcs', 'abfss', or 'azure') based on the URI prefix."""
|
||||
if self.bucket_uri is None:
|
||||
return None
|
||||
elif self.bucket_uri.startswith("s3://"):
|
||||
return "s3"
|
||||
elif self.bucket_uri.startswith("gs://"):
|
||||
return "gcs"
|
||||
elif self.bucket_uri.startswith("abfss://"):
|
||||
return "abfss"
|
||||
elif self.bucket_uri.startswith("azure://"):
|
||||
return "azure"
|
||||
return None
|
||||
|
||||
|
||||
class LoraMirrorConfig(BaseModelExtended):
|
||||
lora_model_id: str
|
||||
bucket_uri: str
|
||||
max_total_tokens: Optional[int]
|
||||
sync_args: Optional[List[str]] = None
|
||||
|
||||
@field_validator("bucket_uri")
|
||||
@classmethod
|
||||
def check_uri_format(cls, value):
|
||||
if value is None:
|
||||
return value
|
||||
|
||||
if not is_remote_path(value):
|
||||
raise ValueError(
|
||||
f'Got invalid value "{value}" for bucket_uri. '
|
||||
'Expected a URI that starts with "s3://", "gs://", "abfss://", or "azure://".'
|
||||
)
|
||||
return value
|
||||
|
||||
@property
|
||||
def _bucket_name_and_path(self) -> str:
|
||||
for prefix in ["s3://", "gs://", "abfss://", "azure://"]:
|
||||
if self.bucket_uri.startswith(prefix):
|
||||
return self.bucket_uri[len(prefix) :]
|
||||
return self.bucket_uri
|
||||
|
||||
@property
|
||||
def bucket_name(self) -> str:
|
||||
bucket_part = self._bucket_name_and_path.split("/")[0]
|
||||
|
||||
# For ABFSS and Azure URIs, extract container name from container@account format
|
||||
if self.bucket_uri.startswith(("abfss://", "azure://")) and "@" in bucket_part:
|
||||
return bucket_part.split("@")[0]
|
||||
|
||||
return bucket_part
|
||||
|
||||
@property
|
||||
def bucket_path(self) -> str:
|
||||
return "/".join(self._bucket_name_and_path.split("/")[1:])
|
||||
|
||||
|
||||
class CloudFileSystem:
|
||||
"""A unified interface for cloud file system operations.
|
||||
|
||||
This class provides a simple interface for common operations on cloud storage
|
||||
systems (S3, GCS, Azure) by delegating to provider-specific implementations
|
||||
for optimal performance.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def _get_provider_fs(bucket_uri: str):
|
||||
"""Get the appropriate provider-specific filesystem class based on URI.
|
||||
|
||||
Args:
|
||||
bucket_uri: URI of the cloud storage (s3://, gs://, abfss://, or azure://)
|
||||
|
||||
Returns:
|
||||
The appropriate filesystem class (S3FileSystem, GCSFileSystem, or AzureFileSystem)
|
||||
|
||||
Raises:
|
||||
ValueError: If the URI scheme is not supported
|
||||
"""
|
||||
if bucket_uri.startswith("pyarrow-"):
|
||||
return PyArrowFileSystem
|
||||
elif bucket_uri.startswith("s3://"):
|
||||
return S3FileSystem
|
||||
elif bucket_uri.startswith("gs://"):
|
||||
return GCSFileSystem
|
||||
elif bucket_uri.startswith(("abfss://", "azure://")):
|
||||
return AzureFileSystem
|
||||
else:
|
||||
raise ValueError(f"Unsupported URI scheme: {bucket_uri}")
|
||||
|
||||
@staticmethod
|
||||
def get_file(
|
||||
object_uri: str, decode_as_utf_8: bool = True
|
||||
) -> Optional[Union[str, bytes]]:
|
||||
"""Download a file from cloud storage into memory.
|
||||
|
||||
Args:
|
||||
object_uri: URI of the file (s3://, gs://, abfss://, or azure://)
|
||||
decode_as_utf_8: If True, decode the file as UTF-8
|
||||
|
||||
Returns:
|
||||
File contents as string or bytes, or None if file doesn't exist
|
||||
"""
|
||||
fs_class = CloudFileSystem._get_provider_fs(object_uri)
|
||||
return fs_class.get_file(object_uri, decode_as_utf_8)
|
||||
|
||||
@staticmethod
|
||||
def list_subfolders(folder_uri: str) -> List[str]:
|
||||
"""List the immediate subfolders in a cloud directory.
|
||||
|
||||
Args:
|
||||
folder_uri: URI of the directory (s3://, gs://, abfss://, or azure://)
|
||||
|
||||
Returns:
|
||||
List of subfolder names (without trailing slashes)
|
||||
"""
|
||||
fs_class = CloudFileSystem._get_provider_fs(folder_uri)
|
||||
return fs_class.list_subfolders(folder_uri)
|
||||
|
||||
@staticmethod
|
||||
def download_files(
|
||||
path: str,
|
||||
bucket_uri: str,
|
||||
substrings_to_include: Optional[List[str]] = None,
|
||||
suffixes_to_exclude: Optional[List[str]] = None,
|
||||
) -> None:
|
||||
"""Download files from cloud storage to a local directory.
|
||||
|
||||
Args:
|
||||
path: Local directory where files will be downloaded
|
||||
bucket_uri: URI of cloud directory
|
||||
substrings_to_include: Only include files containing these substrings
|
||||
suffixes_to_exclude: Exclude certain files from download (e.g .safetensors)
|
||||
"""
|
||||
fs_class = CloudFileSystem._get_provider_fs(bucket_uri)
|
||||
fs_class.download_files(
|
||||
path, bucket_uri, substrings_to_include, suffixes_to_exclude
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def download_model(
|
||||
destination_path: str,
|
||||
bucket_uri: str,
|
||||
tokenizer_only: bool,
|
||||
exclude_safetensors: bool = False,
|
||||
) -> None:
|
||||
"""Download a model from cloud storage.
|
||||
|
||||
This downloads a model in the format expected by the HuggingFace transformers
|
||||
library.
|
||||
|
||||
Args:
|
||||
destination_path: Path where the model will be stored
|
||||
bucket_uri: URI of the cloud directory containing the model
|
||||
tokenizer_only: If True, only download tokenizer-related files
|
||||
exclude_safetensors: If True, skip download of safetensor files
|
||||
"""
|
||||
try:
|
||||
# Get the provider-specific filesystem
|
||||
fs_class = CloudFileSystem._get_provider_fs(bucket_uri)
|
||||
|
||||
# Construct hash file URI
|
||||
hash_uri = bucket_uri.rstrip("/") + "/hash"
|
||||
|
||||
# Try to download and read hash file
|
||||
hash_content = fs_class.get_file(hash_uri, decode_as_utf_8=True)
|
||||
|
||||
if hash_content is not None:
|
||||
f_hash = hash_content.strip()
|
||||
logger.info(
|
||||
f"Detected hash file in bucket {bucket_uri}. "
|
||||
f"Using {f_hash} as the hash."
|
||||
)
|
||||
else:
|
||||
f_hash = "0000000000000000000000000000000000000000"
|
||||
logger.info(
|
||||
f"Hash file does not exist in bucket {bucket_uri}. "
|
||||
f"Using default hash {f_hash} - expected behavior - a hash file is optional. "
|
||||
)
|
||||
|
||||
# Write hash to refs/main
|
||||
main_dir = os.path.join(destination_path, "refs")
|
||||
os.makedirs(main_dir, exist_ok=True)
|
||||
with open(os.path.join(main_dir, "main"), "w") as f:
|
||||
f.write(f_hash)
|
||||
|
||||
# Create destination directory
|
||||
destination_dir = os.path.join(destination_path, "snapshots", f_hash)
|
||||
os.makedirs(destination_dir, exist_ok=True)
|
||||
|
||||
logger.info(f'Downloading model files to directory "{destination_dir}".')
|
||||
|
||||
# Download files
|
||||
tokenizer_file_substrings = (
|
||||
["tokenizer", "config.json", "chat_template"] if tokenizer_only else []
|
||||
)
|
||||
|
||||
safetensors_to_exclude = [".safetensors"] if exclude_safetensors else None
|
||||
|
||||
CloudFileSystem.download_files(
|
||||
path=destination_dir,
|
||||
bucket_uri=bucket_uri,
|
||||
substrings_to_include=tokenizer_file_substrings,
|
||||
suffixes_to_exclude=safetensors_to_exclude,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"Error downloading model from {bucket_uri}: {e}")
|
||||
raise
|
||||
|
||||
@staticmethod
|
||||
def upload_files(
|
||||
local_path: str,
|
||||
bucket_uri: str,
|
||||
) -> None:
|
||||
"""Upload files to cloud storage.
|
||||
|
||||
Args:
|
||||
local_path: The local path of the files to upload.
|
||||
bucket_uri: The bucket uri to upload the files to, must start with
|
||||
`s3://`, `gs://`, `abfss://`, or `azure://`.
|
||||
"""
|
||||
fs_class = CloudFileSystem._get_provider_fs(bucket_uri)
|
||||
fs_class.upload_files(local_path, bucket_uri)
|
||||
|
||||
@staticmethod
|
||||
def upload_model(
|
||||
local_path: str,
|
||||
bucket_uri: str,
|
||||
) -> None:
|
||||
"""Upload a model to cloud storage.
|
||||
|
||||
Args:
|
||||
local_path: The local path of the model.
|
||||
bucket_uri: The bucket uri to upload the model to, must start with `s3://` or `gs://`.
|
||||
"""
|
||||
try:
|
||||
# If refs/main exists, upload as hash, and treat snapshots/<hash> as the model.
|
||||
# Otherwise, this is a custom model, we do not assume folder hierarchy.
|
||||
refs_main = Path(local_path, "refs", "main")
|
||||
if refs_main.exists():
|
||||
model_path = os.path.join(
|
||||
local_path, "snapshots", refs_main.read_text().strip()
|
||||
)
|
||||
CloudFileSystem.upload_files(
|
||||
local_path=model_path, bucket_uri=bucket_uri
|
||||
)
|
||||
CloudFileSystem.upload_files(
|
||||
local_path=str(refs_main),
|
||||
bucket_uri=os.path.join(bucket_uri, "hash"),
|
||||
)
|
||||
else:
|
||||
CloudFileSystem.upload_files(
|
||||
local_path=local_path, bucket_uri=bucket_uri
|
||||
)
|
||||
logger.info(f"Uploaded model files to {bucket_uri}.")
|
||||
except Exception as e:
|
||||
logger.exception(f"Error uploading model to {bucket_uri}: {e}")
|
||||
raise
|
||||
|
||||
|
||||
class _CacheEntry(NamedTuple):
|
||||
value: Any
|
||||
expire_time: Optional[float]
|
||||
|
||||
|
||||
class CloudObjectCache:
|
||||
"""A cache that works with both sync and async fetch functions.
|
||||
|
||||
The purpose of this data structure is to cache the result of a function call
|
||||
usually used to fetch a value from a cloud object store.
|
||||
|
||||
The idea is this:
|
||||
- Cloud operations are expensive
|
||||
- In LoRA specifically, we would fetch remote storage to download the model weights
|
||||
at each request.
|
||||
- If the same model is requested many times, we don't want to inflate the time to first token.
|
||||
- We control the cache via not only the least recently used eviction policy, but also
|
||||
by expiring cache entries after a certain time.
|
||||
- If the object is missing, we cache the missing status for a small duration while if
|
||||
the object exists, we cache the object for a longer duration.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_size: int,
|
||||
fetch_fn: Union[Callable[[str], Any], Callable[[str], Awaitable[Any]]],
|
||||
missing_expire_seconds: Optional[int] = None,
|
||||
exists_expire_seconds: Optional[int] = None,
|
||||
missing_object_value: Any = object(),
|
||||
):
|
||||
"""Initialize the cache.
|
||||
|
||||
Args:
|
||||
max_size: Maximum number of items to store in cache
|
||||
fetch_fn: Function to fetch values (can be sync or async)
|
||||
missing_expire_seconds: How long to cache missing objects (None for no expiration)
|
||||
exists_expire_seconds: How long to cache existing objects (None for no expiration)
|
||||
missing_object_value: Sentinel value used to represent a missing object in the cache.
|
||||
"""
|
||||
self._cache: Dict[str, _CacheEntry] = {}
|
||||
self._max_size = max_size
|
||||
self._fetch_fn = fetch_fn
|
||||
self._missing_expire_seconds = missing_expire_seconds
|
||||
self._exists_expire_seconds = exists_expire_seconds
|
||||
self._is_async = inspect.iscoroutinefunction(fetch_fn) or (
|
||||
callable(fetch_fn) and inspect.iscoroutinefunction(fetch_fn.__call__)
|
||||
)
|
||||
self._missing_object_value = missing_object_value
|
||||
# Lock for thread-safe cache access
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def aget(self, key: str) -> Any:
|
||||
"""Async get value from cache or fetch it if needed."""
|
||||
|
||||
if not self._is_async:
|
||||
raise ValueError("Cannot use async get() with sync fetch function")
|
||||
|
||||
async with self._lock:
|
||||
value, should_fetch = self._check_cache(key)
|
||||
if not should_fetch:
|
||||
return value
|
||||
|
||||
# Fetch new value
|
||||
value = await self._fetch_fn(key)
|
||||
self._update_cache(key, value)
|
||||
return value
|
||||
|
||||
def get(self, key: str) -> Any:
|
||||
"""Sync get value from cache or fetch it if needed."""
|
||||
if self._is_async:
|
||||
raise ValueError("Cannot use sync get() with async fetch function")
|
||||
|
||||
# For sync access, we use a simple check-then-act pattern
|
||||
# This is safe because sync functions are not used in async context
|
||||
value, should_fetch = self._check_cache(key)
|
||||
if not should_fetch:
|
||||
return value
|
||||
|
||||
# Fetch new value
|
||||
value = self._fetch_fn(key)
|
||||
self._update_cache(key, value)
|
||||
return value
|
||||
|
||||
def _check_cache(self, key: str) -> tuple[Any, bool]:
|
||||
"""Check if key exists in cache and is valid.
|
||||
|
||||
Args:
|
||||
key: The cache key to check.
|
||||
|
||||
Returns:
|
||||
Tuple of (value, should_fetch)
|
||||
where should_fetch is True if we need to fetch a new value
|
||||
"""
|
||||
now = time.monotonic()
|
||||
|
||||
if key in self._cache:
|
||||
value, expire_time = self._cache[key]
|
||||
if expire_time is None or now < expire_time:
|
||||
return value, False
|
||||
|
||||
return None, True
|
||||
|
||||
def _update_cache(self, key: str, value: Any) -> None:
|
||||
"""Update cache with new value."""
|
||||
now = time.monotonic()
|
||||
|
||||
# Calculate expiration
|
||||
expire_time = None
|
||||
if (
|
||||
self._missing_expire_seconds is not None
|
||||
or self._exists_expire_seconds is not None
|
||||
):
|
||||
if value is self._missing_object_value:
|
||||
expire_time = (
|
||||
now + self._missing_expire_seconds
|
||||
if self._missing_expire_seconds
|
||||
else None
|
||||
)
|
||||
else:
|
||||
expire_time = (
|
||||
now + self._exists_expire_seconds
|
||||
if self._exists_expire_seconds
|
||||
else None
|
||||
)
|
||||
|
||||
# Enforce size limit by removing oldest entry if needed
|
||||
# This is an O(n) operation but it's fine since the cache size is usually small.
|
||||
if len(self._cache) >= self._max_size:
|
||||
oldest_key = min(
|
||||
self._cache, key=lambda k: self._cache[k].expire_time or float("inf")
|
||||
)
|
||||
del self._cache[oldest_key]
|
||||
|
||||
self._cache[key] = _CacheEntry(value, expire_time)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self._cache)
|
||||
|
||||
|
||||
class CloudModelAccessor:
|
||||
"""Unified accessor for models stored in cloud storage (S3 or GCS).
|
||||
|
||||
Args:
|
||||
model_id: The model id to download or upload.
|
||||
mirror_config: The mirror config for the model.
|
||||
"""
|
||||
|
||||
def __init__(self, model_id: str, mirror_config: CloudMirrorConfig):
|
||||
self.model_id = model_id
|
||||
self.mirror_config = mirror_config
|
||||
|
||||
def _get_lock_path(self, suffix: str = "") -> Path:
|
||||
return Path(
|
||||
"~", f"{self.model_id.replace('/', '--')}{suffix}.lock"
|
||||
).expanduser()
|
||||
|
||||
def _get_model_path(self) -> Path:
|
||||
if Path(self.model_id).exists():
|
||||
return Path(self.model_id)
|
||||
# Delayed import to avoid circular dependencies
|
||||
from huggingface_hub.constants import HF_HUB_CACHE
|
||||
|
||||
return Path(
|
||||
HF_HUB_CACHE, f"models--{self.model_id.replace('/', '--')}"
|
||||
).expanduser()
|
||||
|
||||
|
||||
def remote_object_cache(
|
||||
max_size: int,
|
||||
missing_expire_seconds: Optional[int] = None,
|
||||
exists_expire_seconds: Optional[int] = None,
|
||||
missing_object_value: Any = None,
|
||||
) -> Callable[[Callable[..., T]], Callable[..., T]]:
|
||||
"""A decorator that provides async caching using CloudObjectCache.
|
||||
|
||||
This is a direct replacement for the remote_object_cache/cachetools combination,
|
||||
using CloudObjectCache internally to maintain cache state.
|
||||
|
||||
Args:
|
||||
max_size: Maximum number of items to store in cache
|
||||
missing_expire_seconds: How long to cache missing objects
|
||||
exists_expire_seconds: How long to cache existing objects
|
||||
missing_object_value: Value to use for missing objects
|
||||
|
||||
Returns:
|
||||
A decorator that wraps an async function with cache lookup.
|
||||
"""
|
||||
|
||||
def decorator(func: Callable[..., T]) -> Callable[..., T]:
|
||||
# Create a single cache instance for this function
|
||||
cache = CloudObjectCache(
|
||||
max_size=max_size,
|
||||
fetch_fn=func,
|
||||
missing_expire_seconds=missing_expire_seconds,
|
||||
exists_expire_seconds=exists_expire_seconds,
|
||||
missing_object_value=missing_object_value,
|
||||
)
|
||||
|
||||
async def wrapper(*args, **kwargs):
|
||||
# Extract the key from either first positional arg or object_uri kwarg
|
||||
key = args[0] if args else kwargs.get("object_uri")
|
||||
return await cache.aget(key)
|
||||
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
@@ -0,0 +1,324 @@
|
||||
import enum
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import List, Optional
|
||||
|
||||
from filelock import FileLock
|
||||
|
||||
from ray.llm._internal.common.callbacks.base import CallbackBase
|
||||
from ray.llm._internal.common.observability.logging import get_logger
|
||||
from ray.llm._internal.common.utils.cloud_utils import (
|
||||
CloudFileSystem,
|
||||
CloudMirrorConfig,
|
||||
CloudModelAccessor,
|
||||
is_remote_path,
|
||||
)
|
||||
from ray.llm._internal.common.utils.import_utils import try_import
|
||||
|
||||
torch = try_import("torch")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
STREAMING_LOAD_FORMATS = ["runai_streamer", "runai_streamer_sharded", "tensorizer"]
|
||||
|
||||
|
||||
class NodeModelDownloadable(enum.Enum):
|
||||
"""Defines which files to download from cloud storage."""
|
||||
|
||||
MODEL_AND_TOKENIZER = enum.auto()
|
||||
TOKENIZER_ONLY = enum.auto()
|
||||
EXCLUDE_SAFETENSORS = enum.auto()
|
||||
NONE = enum.auto()
|
||||
|
||||
def __bool__(self):
|
||||
return self != NodeModelDownloadable.NONE
|
||||
|
||||
def union(self, other: "NodeModelDownloadable") -> "NodeModelDownloadable":
|
||||
"""Return a NodeModelDownloadable that is a union of this and the other."""
|
||||
if (
|
||||
self == NodeModelDownloadable.MODEL_AND_TOKENIZER
|
||||
or other == NodeModelDownloadable.MODEL_AND_TOKENIZER
|
||||
):
|
||||
return NodeModelDownloadable.MODEL_AND_TOKENIZER
|
||||
if (
|
||||
self == NodeModelDownloadable.EXCLUDE_SAFETENSORS
|
||||
or other == NodeModelDownloadable.EXCLUDE_SAFETENSORS
|
||||
):
|
||||
return NodeModelDownloadable.EXCLUDE_SAFETENSORS
|
||||
if (
|
||||
self == NodeModelDownloadable.TOKENIZER_ONLY
|
||||
or other == NodeModelDownloadable.TOKENIZER_ONLY
|
||||
):
|
||||
return NodeModelDownloadable.TOKENIZER_ONLY
|
||||
|
||||
return NodeModelDownloadable.NONE
|
||||
|
||||
|
||||
def get_model_entrypoint(model_id: str) -> str:
|
||||
"""Get the path to entrypoint of the model on disk if it exists, otherwise return the model id as is.
|
||||
|
||||
Entrypoint is typically <HF_HUB_CACHE>/models--<model_id>/
|
||||
|
||||
Args:
|
||||
model_id: Hugging Face model ID.
|
||||
|
||||
Returns:
|
||||
The path to the entrypoint of the model on disk if it exists, otherwise the model id as is.
|
||||
"""
|
||||
from huggingface_hub.constants import HF_HUB_CACHE
|
||||
|
||||
model_dir = Path(
|
||||
HF_HUB_CACHE, f"models--{model_id.replace('/', '--')}"
|
||||
).expanduser()
|
||||
if not model_dir.exists():
|
||||
return model_id
|
||||
return str(model_dir.absolute())
|
||||
|
||||
|
||||
def get_model_location_on_disk(model_id: str) -> str:
|
||||
"""Get the location of the model on disk if exists, otherwise return the model id as is.
|
||||
|
||||
Args:
|
||||
model_id: Hugging Face model ID.
|
||||
|
||||
Returns:
|
||||
The path to the model on disk if it exists, otherwise the model id as is.
|
||||
"""
|
||||
model_dir = Path(get_model_entrypoint(model_id))
|
||||
model_id_or_path = model_id
|
||||
|
||||
model_dir_refs_main = Path(model_dir, "refs", "main")
|
||||
|
||||
if model_dir.exists():
|
||||
if model_dir_refs_main.exists():
|
||||
# If refs/main exists, use the snapshot hash to find the model
|
||||
# and check if *config.json (could be config.json for general models
|
||||
# or adapter_config.json for LoRA adapters) exists to make sure it
|
||||
# follows HF model repo structure.
|
||||
with open(model_dir_refs_main, "r") as f:
|
||||
snapshot_hash = f.read().strip()
|
||||
|
||||
snapshot_hash_path = Path(model_dir, "snapshots", snapshot_hash)
|
||||
if snapshot_hash_path.exists() and list(
|
||||
Path(snapshot_hash_path).glob("*config.json")
|
||||
):
|
||||
model_id_or_path = str(snapshot_hash_path.absolute())
|
||||
else:
|
||||
# If it doesn't have refs/main, it is a custom model repo
|
||||
# and we can just return the model_dir.
|
||||
model_id_or_path = str(model_dir.absolute())
|
||||
|
||||
return model_id_or_path
|
||||
|
||||
|
||||
class CloudModelDownloader(CloudModelAccessor):
|
||||
"""Unified downloader for models stored in cloud storage (S3 or GCS).
|
||||
|
||||
Args:
|
||||
model_id: The model id to download.
|
||||
mirror_config: The mirror config for the model.
|
||||
"""
|
||||
|
||||
def get_model(
|
||||
self,
|
||||
tokenizer_only: bool,
|
||||
exclude_safetensors: bool = False,
|
||||
) -> str:
|
||||
"""Gets a model from cloud storage and stores it locally.
|
||||
|
||||
Args:
|
||||
tokenizer_only: whether to download only the tokenizer files.
|
||||
exclude_safetensors: whether to download safetensors files to disk.
|
||||
|
||||
Returns:
|
||||
File path of model if downloaded, else the model id.
|
||||
"""
|
||||
bucket_uri = self.mirror_config.bucket_uri
|
||||
|
||||
if bucket_uri is None:
|
||||
return self.model_id
|
||||
|
||||
# Use different lock paths for different download types to avoid race conditions
|
||||
# where a tokenizer-only download completes and subsequent full model downloads
|
||||
# incorrectly assume the model weights are already cached.
|
||||
if tokenizer_only:
|
||||
lock_suffix = "-tokenizer"
|
||||
elif exclude_safetensors:
|
||||
lock_suffix = "-exclude-safetensors"
|
||||
else:
|
||||
lock_suffix = "-full"
|
||||
lock_path = self._get_lock_path(suffix=lock_suffix)
|
||||
path = self._get_model_path()
|
||||
storage_type = self.mirror_config.storage_type
|
||||
|
||||
try:
|
||||
# Timeout 0 means there will be only one attempt to acquire
|
||||
# the file lock. If it cannot be acquired, a TimeoutError
|
||||
# will be thrown.
|
||||
# This ensures that subsequent processes don't duplicate work.
|
||||
with FileLock(lock_path, timeout=0):
|
||||
try:
|
||||
if exclude_safetensors:
|
||||
logger.info("Skipping download of safetensors files.")
|
||||
CloudFileSystem.download_model(
|
||||
destination_path=path,
|
||||
bucket_uri=bucket_uri,
|
||||
tokenizer_only=tokenizer_only,
|
||||
exclude_safetensors=exclude_safetensors,
|
||||
)
|
||||
logger.info(
|
||||
"Finished downloading %s for %s from %s storage",
|
||||
"tokenizer" if tokenizer_only else "model and tokenizer",
|
||||
self.model_id,
|
||||
storage_type.upper() if storage_type else "cloud",
|
||||
)
|
||||
except RuntimeError:
|
||||
logger.exception(
|
||||
"Failed to download files for model %s from %s storage",
|
||||
self.model_id,
|
||||
storage_type.upper() if storage_type else "cloud",
|
||||
)
|
||||
except TimeoutError:
|
||||
# If the directory is already locked, then wait but do not do anything.
|
||||
with FileLock(lock_path, timeout=-1):
|
||||
pass
|
||||
return get_model_location_on_disk(self.model_id)
|
||||
|
||||
def get_extra_files(self) -> List[str]:
|
||||
"""Gets user-specified extra files from cloud storage and stores them in
|
||||
provided paths.
|
||||
|
||||
Returns: list of file paths of extra files if downloaded.
|
||||
"""
|
||||
paths = []
|
||||
extra_files = self.mirror_config.extra_files or []
|
||||
if not extra_files:
|
||||
return paths
|
||||
|
||||
lock_path = self._get_lock_path(suffix="-extra_files")
|
||||
storage_type = self.mirror_config.storage_type
|
||||
|
||||
logger.info(
|
||||
f"Downloading extra files for {self.model_id} from {storage_type} storage"
|
||||
)
|
||||
try:
|
||||
# Timeout 0 means there will be only one attempt to acquire
|
||||
# the file lock. If it cannot be acquired, a TimeoutError
|
||||
# will be thrown.
|
||||
# This ensures that subsequent processes don't duplicate work.
|
||||
with FileLock(lock_path, timeout=0):
|
||||
for extra_file in extra_files:
|
||||
path = Path(
|
||||
os.path.expandvars(extra_file.destination_path)
|
||||
).expanduser()
|
||||
paths.append(path)
|
||||
CloudFileSystem.download_files(
|
||||
path=path,
|
||||
bucket_uri=extra_file.bucket_uri,
|
||||
)
|
||||
except TimeoutError:
|
||||
# If the directory is already locked, then wait but do not do anything.
|
||||
with FileLock(lock_path, timeout=-1):
|
||||
pass
|
||||
return paths
|
||||
|
||||
|
||||
def _log_download_info(
|
||||
*, source: str, download_model: NodeModelDownloadable, download_extra_files: bool
|
||||
):
|
||||
if download_model == NodeModelDownloadable.NONE:
|
||||
if download_extra_files:
|
||||
logger.info("Downloading extra files from %s", source)
|
||||
else:
|
||||
logger.info("Not downloading anything from %s", source)
|
||||
elif download_model == NodeModelDownloadable.TOKENIZER_ONLY:
|
||||
if download_extra_files:
|
||||
logger.info("Downloading tokenizer and extra files from %s", source)
|
||||
else:
|
||||
logger.info("Downloading tokenizer from %s", source)
|
||||
elif download_model == NodeModelDownloadable.MODEL_AND_TOKENIZER:
|
||||
if download_extra_files:
|
||||
logger.info("Downloading model, tokenizer, and extra files from %s", source)
|
||||
else:
|
||||
logger.info("Downloading model and tokenizer from %s", source)
|
||||
|
||||
|
||||
def download_model_files(
|
||||
model_id: Optional[str] = None,
|
||||
mirror_config: Optional[CloudMirrorConfig] = None,
|
||||
download_model: NodeModelDownloadable = NodeModelDownloadable.MODEL_AND_TOKENIZER,
|
||||
download_extra_files: bool = True,
|
||||
callback: Optional[CallbackBase] = None,
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Download the model files from the cloud storage. We support two ways to specify
|
||||
the remote model path in the cloud storage:
|
||||
Approach 1:
|
||||
- model_id: The vanilla model id such as "meta-llama/Llama-3.1-8B-Instruct".
|
||||
- mirror_config: Config for downloading model from cloud storage.
|
||||
|
||||
Approach 2:
|
||||
- model_id: The remote path (s3:// or gs://) in the cloud storage.
|
||||
- mirror_config: None.
|
||||
In this approach, we will create a CloudMirrorConfig from the model_id and use that
|
||||
to download the model.
|
||||
|
||||
Args:
|
||||
model_id: The model id.
|
||||
mirror_config: Config for downloading model from cloud storage.
|
||||
download_model: What parts of the model to download.
|
||||
download_extra_files: Whether to download extra files specified in the mirror config.
|
||||
callback: Callback to run before downloading model files.
|
||||
|
||||
Returns:
|
||||
The local path to the downloaded model, or the original model ID
|
||||
if no cloud storage mirror is configured or if the model is not downloaded.
|
||||
"""
|
||||
|
||||
# Create the torch cache kernels directory if it doesn't exist.
|
||||
# This is a workaround for a torch issue, where the kernels directory
|
||||
# cannot be created by torch if the parent directory doesn't exist.
|
||||
torch_cache_home = torch.hub._get_torch_home()
|
||||
os.makedirs(os.path.join(torch_cache_home, "kernels"), exist_ok=True)
|
||||
model_path_or_id = model_id
|
||||
|
||||
if callback is not None:
|
||||
callback.run_callback_sync("on_before_download_model_files_distributed")
|
||||
|
||||
if model_id is None:
|
||||
return None
|
||||
|
||||
if mirror_config is None:
|
||||
if is_remote_path(model_id):
|
||||
logger.info(
|
||||
"Creating a CloudMirrorConfig from remote model path %s", model_id
|
||||
)
|
||||
mirror_config = CloudMirrorConfig(bucket_uri=model_id)
|
||||
else:
|
||||
logger.info("No cloud storage mirror configured")
|
||||
return model_id
|
||||
|
||||
storage_type = mirror_config.storage_type
|
||||
source = (
|
||||
f"{storage_type.upper()} mirror" if storage_type else "Cloud storage mirror"
|
||||
)
|
||||
|
||||
_log_download_info(
|
||||
source=source,
|
||||
download_model=download_model,
|
||||
download_extra_files=download_extra_files,
|
||||
)
|
||||
|
||||
downloader = CloudModelDownloader(model_id, mirror_config)
|
||||
|
||||
if download_model != NodeModelDownloadable.NONE:
|
||||
model_path_or_id = downloader.get_model(
|
||||
tokenizer_only=download_model == NodeModelDownloadable.TOKENIZER_ONLY,
|
||||
exclude_safetensors=download_model
|
||||
== NodeModelDownloadable.EXCLUDE_SAFETENSORS,
|
||||
)
|
||||
|
||||
if download_extra_files:
|
||||
downloader.get_extra_files()
|
||||
|
||||
return model_path_or_id
|
||||
@@ -0,0 +1,90 @@
|
||||
"""Utility functions for importing modules in the LLM module."""
|
||||
import importlib
|
||||
import logging
|
||||
from types import ModuleType
|
||||
from typing import Any, NoReturn, Optional, Type
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def try_import(name: str, error: bool = False) -> Optional[ModuleType]:
|
||||
"""Try importing the module and returns the module (or None).
|
||||
|
||||
Args:
|
||||
name: The name of the module to import.
|
||||
error: Whether to raise an error if the module cannot be imported.
|
||||
|
||||
Returns:
|
||||
The module, or None if it cannot be imported.
|
||||
|
||||
Raises:
|
||||
ImportError: If error=True and the module is not installed.
|
||||
"""
|
||||
try:
|
||||
return importlib.import_module(name)
|
||||
except ImportError:
|
||||
if error:
|
||||
raise ImportError(f"Could not import {name}")
|
||||
else:
|
||||
logger.warning("Could not import %s", name)
|
||||
return None
|
||||
|
||||
|
||||
def raise_llm_engine_import_error(
|
||||
vllm_error: ImportError,
|
||||
sglang_error: ImportError,
|
||||
) -> NoReturn:
|
||||
"""Raise a descriptive ImportError when both vLLM and SGLang fail to import.
|
||||
|
||||
Distinguishes between a package not being installed (ModuleNotFoundError
|
||||
whose .name matches the top-level package) and a broken installation
|
||||
(any other ImportError, e.g. a missing .so or a missing transitive dep).
|
||||
|
||||
Args:
|
||||
vllm_error: The ImportError raised when importing vLLM.
|
||||
sglang_error: The ImportError raised when importing SGLang.
|
||||
"""
|
||||
vllm_not_installed = (
|
||||
isinstance(vllm_error, ModuleNotFoundError) and vllm_error.name == "vllm"
|
||||
)
|
||||
sglang_not_installed = (
|
||||
isinstance(sglang_error, ModuleNotFoundError) and sglang_error.name == "sglang"
|
||||
)
|
||||
|
||||
if vllm_not_installed and sglang_not_installed:
|
||||
raise ImportError(
|
||||
"Neither vLLM nor SGLang is installed. At least one is required "
|
||||
"for Ray Serve LLM protocol models. Install with: "
|
||||
"`pip install ray[llm]` or `pip install sglang[all]`"
|
||||
)
|
||||
|
||||
messages = []
|
||||
if not vllm_not_installed:
|
||||
messages.append(
|
||||
"vLLM is installed but failed to import. This may indicate a "
|
||||
"CUDA version mismatch or a missing vLLM dependency. "
|
||||
f"Original error: {vllm_error}"
|
||||
)
|
||||
if not sglang_not_installed:
|
||||
messages.append(
|
||||
"SGLang is installed but failed to import. This may indicate a "
|
||||
"missing SGLang dependency. "
|
||||
f"Original error: {sglang_error}"
|
||||
)
|
||||
# Chain to the error that is actually relevant: vLLM's if it is broken,
|
||||
# otherwise sglang's (i.e. vLLM was simply not installed).
|
||||
cause = vllm_error if not vllm_not_installed else sglang_error
|
||||
raise ImportError("\n".join(messages)) from cause
|
||||
|
||||
|
||||
def load_class(path: str) -> Type[Any]:
|
||||
"""Load class from string path."""
|
||||
if ":" in path:
|
||||
module_path, class_name = path.rsplit(":", 1)
|
||||
else:
|
||||
module_path, class_name = path.rsplit(".", 1)
|
||||
|
||||
module = try_import(module_path, error=True)
|
||||
callback_class = getattr(module, class_name)
|
||||
|
||||
return callback_class
|
||||
@@ -0,0 +1,233 @@
|
||||
"""
|
||||
Generic LoRA utilities and abstractions.
|
||||
|
||||
This module provides canonical LoRA utility functions for both serve and batch components.
|
||||
It serves as the single source of truth for LoRA operations and builds on the generic
|
||||
download primitives from download_utils.py.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
from functools import wraps
|
||||
from typing import Any, Callable, List, Optional, TypeVar, Union
|
||||
|
||||
from ray.llm._internal.common.constants import (
|
||||
CLOUD_OBJECT_EXISTS_EXPIRE_S,
|
||||
CLOUD_OBJECT_MISSING_EXPIRE_S,
|
||||
LORA_ADAPTER_CONFIG_NAME,
|
||||
)
|
||||
|
||||
# Import the global ID manager from common models
|
||||
from ray.llm._internal.common.models import make_async
|
||||
from ray.llm._internal.common.observability.logging import get_logger
|
||||
from ray.llm._internal.common.utils.cloud_utils import (
|
||||
CloudFileSystem,
|
||||
is_remote_path,
|
||||
remote_object_cache,
|
||||
)
|
||||
from ray.llm._internal.common.utils.download_utils import (
|
||||
CloudMirrorConfig,
|
||||
CloudModelDownloader,
|
||||
)
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
# Sentinel object for missing cloud objects
|
||||
CLOUD_OBJECT_MISSING = object()
|
||||
|
||||
DEFAULT_LORA_MAX_TOTAL_TOKENS = 4096
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
def get_base_model_id(model_id: str) -> str:
|
||||
"""Get base model id for a given model id."""
|
||||
return model_id.split(":")[0]
|
||||
|
||||
|
||||
def get_lora_id(lora_model_id: str) -> str:
|
||||
"""Get lora id for a given lora model id."""
|
||||
return ":".join(lora_model_id.split(":")[1:])
|
||||
|
||||
|
||||
def clean_model_id(model_id: str) -> str:
|
||||
"""Clean model ID for filesystem usage by replacing slashes with dashes."""
|
||||
return model_id.replace("/", "--")
|
||||
|
||||
|
||||
def clear_directory(dir: str) -> None:
|
||||
"""Clear a directory recursively, ignoring missing directories."""
|
||||
try:
|
||||
subprocess.run(f"rm -r {dir}", shell=True, check=False)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
|
||||
|
||||
def retry_with_exponential_backoff(
|
||||
max_tries: int,
|
||||
exception_to_check: type[Exception],
|
||||
base_delay: float = 1,
|
||||
max_delay: float = 32,
|
||||
exponential_base: float = 2,
|
||||
) -> Callable[[Callable[..., T]], Callable[..., T]]:
|
||||
"""Retry decorator with exponential backoff."""
|
||||
|
||||
def decorator(func: Callable[..., T]) -> Callable[..., T]:
|
||||
@wraps(func)
|
||||
def wrapper(*args: Any, **kwargs: Any) -> T:
|
||||
delay = base_delay
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_tries):
|
||||
try:
|
||||
return func(*args, **kwargs)
|
||||
except exception_to_check as e:
|
||||
last_exception = e
|
||||
if attempt == max_tries - 1: # Last attempt
|
||||
raise last_exception
|
||||
|
||||
# Log the failure and retry
|
||||
logger.warning(
|
||||
f"Attempt {attempt + 1}/{max_tries} failed: {str(e)}. "
|
||||
f"Retrying in {delay} seconds..."
|
||||
)
|
||||
time.sleep(delay)
|
||||
# Calculate next delay with exponential backoff
|
||||
delay = min(delay * exponential_base, max_delay)
|
||||
|
||||
# This should never be reached due to the raise in the loop
|
||||
raise last_exception if last_exception else RuntimeError(
|
||||
"Unexpected error in retry logic"
|
||||
)
|
||||
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def sync_files_with_lock(
|
||||
bucket_uri: str,
|
||||
local_path: str,
|
||||
timeout: Optional[float] = None,
|
||||
substrings_to_include: Optional[List[str]] = None,
|
||||
) -> None:
|
||||
"""Sync files from bucket_uri to local_path with file locking."""
|
||||
from filelock import FileLock
|
||||
|
||||
logger.info("Downloading %s to %s", bucket_uri, local_path)
|
||||
|
||||
with FileLock(local_path + ".lock", timeout=timeout or -1):
|
||||
try:
|
||||
CloudFileSystem.download_files(
|
||||
path=local_path,
|
||||
bucket_uri=bucket_uri,
|
||||
substrings_to_include=substrings_to_include,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Failed to sync files from %s to %s: %s",
|
||||
bucket_uri,
|
||||
local_path,
|
||||
str(e),
|
||||
)
|
||||
raise
|
||||
|
||||
|
||||
@make_async
|
||||
def _get_object_from_cloud(object_uri: str) -> Union[str, object]:
|
||||
"""Gets an object from the cloud."""
|
||||
if object_uri.endswith("/"):
|
||||
raise ValueError(f'object_uri {object_uri} must not end with a "/".')
|
||||
|
||||
body_str = CloudFileSystem.get_file(object_uri)
|
||||
|
||||
if body_str is None:
|
||||
logger.info(f"{object_uri} does not exist.")
|
||||
return CLOUD_OBJECT_MISSING
|
||||
else:
|
||||
return body_str
|
||||
|
||||
|
||||
@remote_object_cache(
|
||||
max_size=4096,
|
||||
missing_expire_seconds=CLOUD_OBJECT_MISSING_EXPIRE_S,
|
||||
exists_expire_seconds=CLOUD_OBJECT_EXISTS_EXPIRE_S,
|
||||
missing_object_value=CLOUD_OBJECT_MISSING,
|
||||
)
|
||||
async def get_object_from_cloud(object_uri: str) -> Union[str, object]:
|
||||
"""Gets an object from the cloud with caching."""
|
||||
return await _get_object_from_cloud(object_uri)
|
||||
|
||||
|
||||
async def get_lora_finetuned_context_length(bucket_uri: str) -> Optional[int]:
|
||||
"""Gets the sequence length used to tune the LoRA adapter."""
|
||||
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:
|
||||
logger.debug(f"LoRA adapter config file not found at {object_uri}")
|
||||
return None
|
||||
|
||||
try:
|
||||
adapter_config_str = object_str_or_missing_message
|
||||
adapter_config = json.loads(adapter_config_str)
|
||||
return adapter_config.get("max_length")
|
||||
except (json.JSONDecodeError, AttributeError) as e:
|
||||
logger.warning(f"Failed to parse LoRA adapter config at {object_uri}: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def get_lora_model_ids(
|
||||
dynamic_lora_loading_path: str,
|
||||
base_model_id: str,
|
||||
) -> List[str]:
|
||||
"""Get the model IDs of all the LoRA models.
|
||||
|
||||
The dynamic_lora_loading_path is expected to hold subfolders each for
|
||||
a different lora checkpoint. Each subfolder name will correspond to
|
||||
the unique identifier for the lora checkpoint. The lora model is
|
||||
accessible via <base_model_id>:<lora_id>. Therefore, we prepend
|
||||
the base_model_id to each subfolder name.
|
||||
|
||||
Args:
|
||||
dynamic_lora_loading_path: the cloud folder that contains all the LoRA
|
||||
weights.
|
||||
base_model_id: model ID of the base model.
|
||||
|
||||
Returns:
|
||||
List of LoRA fine-tuned model IDs. Does not include the base model
|
||||
itself.
|
||||
"""
|
||||
lora_subfolders = CloudFileSystem.list_subfolders(dynamic_lora_loading_path)
|
||||
|
||||
lora_model_ids = []
|
||||
for subfolder in lora_subfolders:
|
||||
lora_model_ids.append(f"{base_model_id}:{subfolder}")
|
||||
|
||||
return lora_model_ids
|
||||
|
||||
|
||||
def download_lora_adapter(
|
||||
lora_name: str,
|
||||
remote_path: Optional[str] = None,
|
||||
) -> str:
|
||||
"""Download a LoRA adapter from remote storage.
|
||||
|
||||
This maintains backward compatibility with existing code.
|
||||
"""
|
||||
|
||||
assert not is_remote_path(
|
||||
lora_name
|
||||
), "lora_name cannot be a remote path (s3:// or gs://)"
|
||||
|
||||
if remote_path is None:
|
||||
return lora_name
|
||||
|
||||
lora_path = os.path.join(remote_path, lora_name)
|
||||
mirror_config = CloudMirrorConfig(bucket_uri=lora_path)
|
||||
downloader = CloudModelDownloader(lora_name, mirror_config)
|
||||
return downloader.get_model(tokenizer_only=False)
|
||||
@@ -0,0 +1,126 @@
|
||||
from pathlib import Path
|
||||
|
||||
import typer
|
||||
from filelock import FileLock
|
||||
from typing_extensions import Annotated
|
||||
|
||||
from ray.llm._internal.common.observability.logging import get_logger
|
||||
from ray.llm._internal.common.utils.cloud_utils import (
|
||||
CloudFileSystem,
|
||||
CloudMirrorConfig,
|
||||
CloudModelAccessor,
|
||||
is_remote_path,
|
||||
)
|
||||
from ray.llm._internal.common.utils.download_utils import (
|
||||
get_model_entrypoint,
|
||||
)
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class CloudModelUploader(CloudModelAccessor):
|
||||
"""Unified uploader to upload models to cloud storage (S3 or GCS).
|
||||
|
||||
Args:
|
||||
model_id: The model id to upload.
|
||||
mirror_config: The mirror config for the model.
|
||||
"""
|
||||
|
||||
def upload_model(self) -> str:
|
||||
"""Upload the model to cloud storage (s3 or gcs).
|
||||
|
||||
Returns:
|
||||
The remote path of the uploaded model.
|
||||
"""
|
||||
bucket_uri = self.mirror_config.bucket_uri
|
||||
|
||||
lock_path = self._get_lock_path()
|
||||
path = self._get_model_path()
|
||||
storage_type = self.mirror_config.storage_type
|
||||
|
||||
try:
|
||||
# Timeout 0 means there will be only one attempt to acquire
|
||||
# the file lock. If it cannot be acquired, a TimeoutError
|
||||
# will be thrown.
|
||||
# This ensures that subsequent processes don't duplicate work.
|
||||
with FileLock(lock_path, timeout=0):
|
||||
try:
|
||||
CloudFileSystem.upload_model(
|
||||
local_path=path,
|
||||
bucket_uri=bucket_uri,
|
||||
)
|
||||
logger.info(
|
||||
"Finished uploading %s to %s storage",
|
||||
self.model_id,
|
||||
storage_type.upper() if storage_type else "cloud",
|
||||
)
|
||||
except RuntimeError:
|
||||
logger.exception(
|
||||
"Failed to upload model %s to %s storage",
|
||||
self.model_id,
|
||||
storage_type.upper() if storage_type else "cloud",
|
||||
)
|
||||
except TimeoutError:
|
||||
# If the directory is already locked, then wait but do not do anything.
|
||||
with FileLock(lock_path, timeout=-1):
|
||||
pass
|
||||
return bucket_uri
|
||||
|
||||
|
||||
def upload_model_files(model_id: str, bucket_uri: str) -> str:
|
||||
"""Upload the model files to cloud storage (s3 or gcs).
|
||||
|
||||
If `model_id` is a local path, the files will be uploaded to the cloud storage.
|
||||
If `model_id` is a huggingface model id, the model will be downloaded from huggingface
|
||||
and then uploaded to the cloud storage.
|
||||
|
||||
Args:
|
||||
model_id: The huggingface model id, or local model path to upload.
|
||||
bucket_uri: The bucket uri to upload the model to, must start with `s3://` or `gs://`.
|
||||
|
||||
Returns:
|
||||
The remote path of the uploaded model.
|
||||
"""
|
||||
assert not is_remote_path(
|
||||
model_id
|
||||
), f"model_id must NOT be a remote path: {model_id}"
|
||||
assert is_remote_path(bucket_uri), f"bucket_uri must be a remote path: {bucket_uri}"
|
||||
|
||||
if not Path(model_id).exists():
|
||||
maybe_downloaded_model_path = get_model_entrypoint(model_id)
|
||||
if not Path(maybe_downloaded_model_path).exists():
|
||||
logger.info(
|
||||
"Assuming %s is huggingface model id, and downloading it.", model_id
|
||||
)
|
||||
import huggingface_hub
|
||||
|
||||
huggingface_hub.snapshot_download(repo_id=model_id)
|
||||
# Try to get the model path again after downloading.
|
||||
maybe_downloaded_model_path = get_model_entrypoint(model_id)
|
||||
assert Path(
|
||||
maybe_downloaded_model_path
|
||||
).exists(), f"Failed to download the model {model_id} to {maybe_downloaded_model_path}"
|
||||
return upload_model_files(maybe_downloaded_model_path, bucket_uri)
|
||||
else:
|
||||
return upload_model_files(maybe_downloaded_model_path, bucket_uri)
|
||||
|
||||
uploader = CloudModelUploader(model_id, CloudMirrorConfig(bucket_uri=bucket_uri))
|
||||
return uploader.upload_model()
|
||||
|
||||
|
||||
def upload_model_cli(
|
||||
model_source: Annotated[
|
||||
str,
|
||||
typer.Option(
|
||||
help="HuggingFace model ID to download, or local model path to upload",
|
||||
),
|
||||
],
|
||||
bucket_uri: Annotated[
|
||||
str,
|
||||
typer.Option(
|
||||
help="The bucket uri to upload the model to, must start with `s3://` or `gs://`",
|
||||
),
|
||||
],
|
||||
):
|
||||
"""Upload the model files to cloud storage (s3 or gcs)."""
|
||||
upload_model_files(model_source, bucket_uri)
|
||||
@@ -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
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user