chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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)
|
||||
Reference in New Issue
Block a user