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
|
||||
Reference in New Issue
Block a user