chore: import upstream snapshot with attribution

python/ray/llm/_internal/batch/processor/utils.py

@@ -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,
+        ),
+    )