Function get_client

espnet3/parallel/parallel.py:301–343

Context manager to yield a Dask client from the global singleton cluster. Args: config (DictConfig, optional): Cluster config. setup_fn (Callable[[], dict], optional): A setup function that runs on each worker and returns a dictionary of environment variables. Y

(
    config: DictConfig = None, setup_fn: Optional[Callable[[], dict]] = None
)

@contextmanager
def get_client(
    config: DictConfig = None, setup_fn: Optional[Callable[[], dict]] = None
) -> Generator[Client, None, None]:
    """Context manager to yield a Dask client from the global singleton cluster.

    Args:
        config (DictConfig, optional): Cluster config.
        setup_fn (Callable[[], dict], optional): A setup function that runs
            on each worker and returns a dictionary of environment variables.

    Yields:
        Client: A Dask client instance tied to the global cluster.

    Example:
        >>> with get_client() as client:
        ...     results = client.map(lambda x: x**2, range(10))
    """
    client = build_client(config)
    if setup_fn is not None:
        plugin = DictReturnWorkerPlugin(setup_fn)
        reg = getattr(client, "register_worker_plugin", None)
        if reg is None:
            raise RuntimeError(
                "This Dask version lacks register_worker_plugin; please upgrade."
            )
        reg(plugin, name="env")
    try:
        yield client
    finally:
        # Avoid shutdown for LocalCluster and LocalCUDACluster
        cluster = getattr(client, "cluster", None)

        # always close the client first
        client.close()

        if cluster is not None:
            close = getattr(cluster, "close", None)
            if callable(close):
                close()
            else:
                shutdown = getattr(cluster, "shutdown", None)
                if callable(shutdown):
                    shutdown()


@contextmanager