Function parallel_map

espnet3/parallel/parallel.py:382–450

Apply a function to an iterable of inputs in parallel using Dask. This helper takes care of: - Creating (or reusing) a Dask client according to the global or provided configuration. - Optionally registering a per-worker environment via `setup_fn`, making its returned

(
    func: Callable[[Any], Any],
    data: Iterable[Any],
    client: Optional[Client] = None,
    setup_fn: Optional[Callable[[], dict]] = None,
    **kwargs: Any,
)

@typechecked
def parallel_map(
    func: Callable[[Any], Any],
    data: Iterable[Any],
    client: Optional[Client] = None,
    setup_fn: Optional[Callable[[], dict]] = None,
    **kwargs: Any,
) -> list:
    """Apply a function to an iterable of inputs in parallel using Dask.

    This helper takes care of:
      - Creating (or reusing) a Dask client according to the global or
        provided configuration.
      - Optionally registering a per-worker environment via `setup_fn`,
        making its returned dictionary available to `func` automatically.
      - Detecting and preventing conflicts between explicit `kwargs` and
        environment-provided arguments on the **client side** before
        submitting tasks.
      - Wrapping `func` with `wrap_func_with_worker_env` so that missing
        keyword arguments can be injected from the worker environment.

    Args:
        func (Callable[[Any], Any]):
            The function to execute on each element of `data`. May take
            positional and/or keyword parameters.
        data (Iterable[Any]):
            Iterable of input elements to process.
        client (Optional[Client], default=None):
            An existing Dask client to use. If `None`, a temporary client
            will be created using `get_client` and shut down afterwards.
        setup_fn (Optional[Callable[[], dict]], default=None):
            A function run once per worker that returns a dictionary of
            environment variables. These variables are automatically
            injected into `func` if they match parameter names and are not
            explicitly provided.
        **kwargs:
            Additional keyword arguments to pass directly to `func` for all
            elements.

    Returns:
        list:
            The results of applying `func` to each element of `data`, in
            order. The list has the same length as `data`.

    Raises:
        ValueError:
            If any keyword argument name in `kwargs` is also present in the
            worker environment keys (conflict detected before submission).

    Example:
        >>> def setup_fn():
        ...     return {"bias": 10}
        >>> def add_bias(x, bias):
        ...     return x + bias
        >>> # Automatic injection of 'bias' from worker environment:
        >>> results = parallel_map(add_bias, [1, 2, 3], setup_fn=setup_fn)
        >>> results
        [11, 12, 13]
    """