espnet3.parallel.parallel.wrap_func_with_worker_env

source

espnet3.parallel.parallel.wrap_func_with_worker_env(func: Callable) → Callable

Wrap a user-defined function for a WorkerPlugin.

This wrapper inspects the function signature and injects values from the worker environment when names match missing keyword parameters.

• Parameters:func (Callable) – The original user-defined function to be executed on the worker. It may have positional parameters, keyword parameters, and/or a **kwargs catch-all.
• Returns: Wrapped callable that pulls matching keyword values from : worker.plugins["env"] before invoking func.
• Return type: Callable
• Raises:ValueError – Raised when both the worker environment and explicit keyword arguments define the same parameter.

Example

>>> def process(idx, dataset, model):
...     return model(dataset[idx])
>>> wrapped_fn = wrap_func_with_worker_env(process)
>>> # On a Dask worker with env={"dataset": ds, "model": md},
>>> # wrapped_fn(0) injects dataset and model from the worker plugin.