Function _submit_tasks

espnet3/parallel/parallel.py:347–378

Submit job for parallel_map and parallel_for. - Creates/reuses a client (and handles its lifecycle if internal). - Registers the setup plugin (client- and worker-side). - Performs client-side conflict checks. - Wraps the function to inject worker env. - Submits tasks and yields

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

@contextmanager
def _submit_tasks(
    func: Callable[[Any], Any],
    iterable: Iterable[Any],
    client: Optional[Client] = None,
    setup_fn: Optional[Callable[[], dict]] = None,
    **kwargs: Any,
):
    """Submit job for parallel_map and parallel_for.

    - Creates/reuses a client (and handles its lifecycle if internal).
    - Registers the setup plugin (client- and worker-side).
    - Performs client-side conflict checks.
    - Wraps the function to inject worker env.
    - Submits tasks and yields (client, futures).
    """
    _ensure_dask()
    internal = client is None
    if internal:
        ctx = get_client(setup_fn=setup_fn)
        client_cm = ctx
        client = ctx.__enter__()
    elif setup_fn is not None:
        plugin = DictReturnWorkerPlugin(setup_fn)
        getattr(client, "register_worker_plugin")(plugin, name="env")

    try:
        wrapped_func = wrap_func_with_worker_env(func)
        futures = client.map(wrapped_func, iterable, **kwargs)
        yield client, futures
    finally:
        if internal:
            client_cm.__exit__(None, None, None)


@typechecked