espnet3.parallel.parallel.build_local_gpu_cluster

source

espnet3.parallel.parallel.build_local_gpu_cluster(n_workers: int, options: dict) → Client

Create a Dask LocalCUDACluster using available GPUs.

This requires dask_cuda package.

• Parameters:
  • n_workers (int) – Number of Dask workers (must not exceed number of GPUs).
  • options (dict) – Additional options for the LocalCUDACluster.
• Returns: Dask client connected to the GPU cluster.
• Return type: Client
• Raises:
  • RuntimeError – If dask_cuda is not installed.
  • ValueError – If n_workers exceeds the number of available GPUs.

Example

>>> client = build_local_gpu_cluster(n_workers=2, options={})
>>> client.close()