wehub-resource-sync
59 lines
1.8 KiB
Python
59 lines
1.8 KiB
Python
"""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,
|
|
),
|
|
)
|