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