Function get_compute_strategy

python/ray/data/_internal/util.py:589–696

Get `ComputeStrategy` based on the function or class, and concurrency information. Args: fn: The function or generator to apply to a record batch, or a class type that can be instantiated to create such a callable. fn_constructor_args: Positional arguments to pas

(
    fn: "UserDefinedFunction",
    fn_constructor_args: Optional[Iterable[Any]] = None,
    compute: Optional[Union[str, "ComputeStrategy"]] = None,
    concurrency: Optional[Union[int, Tuple[int, int], Tuple[int, int, int]]] = None,
)

def get_compute_strategy(
    fn: "UserDefinedFunction",
    fn_constructor_args: Optional[Iterable[Any]] = None,
    compute: Optional[Union[str, "ComputeStrategy"]] = None,
    concurrency: Optional[Union[int, Tuple[int, int], Tuple[int, int, int]]] = None,
) -> "ComputeStrategy":
    """Get `ComputeStrategy` based on the function or class, and concurrency
    information.

    Args:
        fn: The function or generator to apply to a record batch, or a class type
            that can be instantiated to create such a callable.
        fn_constructor_args: Positional arguments to pass to ``fn``'s constructor.
        compute: Either "tasks" (default) to use Ray Tasks or an
                :class:`~ray.data.ActorPoolStrategy` to use an autoscaling actor pool.
        concurrency: The number of Ray workers to use concurrently.

    Returns:
       The `ComputeStrategy` for execution.
    """
    # Lazily import these objects to avoid circular imports.
    from ray.data._internal.compute import ActorPoolStrategy, TaskPoolStrategy
    from ray.data.block import CallableClass

    if isinstance(fn, CallableClass):
        is_callable_class = True
    else:
        # TODO(chengsu): disallow object that is not a function. For example,
        # An object instance of class often indicates a bug in user code.
        is_callable_class = False
        if fn_constructor_args is not None:
            raise ValueError(
                "``fn_constructor_args`` can only be specified if providing a "
                f"callable class instance for ``fn``, but got: {fn}."
            )

    if compute is not None:
        if is_callable_class and (
            compute == "tasks" or isinstance(compute, TaskPoolStrategy)
        ):
            raise ValueError(
                f"You specified the callable class {fn} as your UDF with the compute "
                f"{compute}, but Ray Data can't schedule callable classes with the task "
                f"pool strategy. To fix this error, pass an ActorPoolStrategy to compute or "
                f"None to use the default compute strategy."
            )
        elif not is_callable_class and (
            compute == "actors" or isinstance(compute, ActorPoolStrategy)
        ):
            raise ValueError(
                f"You specified the function {fn} as your UDF with the compute "
                f"{compute}, but Ray Data can't schedule regular functions with the actor "
                f"pool strategy. To fix this error, pass a TaskPoolStrategy to compute or "
                f"None to use the default compute strategy."
            )
        return compute
    elif concurrency is not None:
        # Legacy code path to support `concurrency` argument.