espnet3.parallel.parallel.parallel_map

source

espnet3.parallel.parallel.parallel_map(func: Callable[[Any], Any], data: Iterable[Any], client: Client | None = None, setup_fn: Callable[[], dict] | None = 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.

• Parameters:
  • 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: The results of applying func to each element of data, in order. The list has the same length as data.
• Return type: list
• 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]