parallel

Dask helper functions.

SchedulerAddress = str | Cluster | Literal['sequential'] | None module-attribute

Dask scheduler address.

• "sequential": Run sequentially without Dask.
• None: Create a local Dask cluster.
• str: Address of an existing Dask scheduler.
• Cluster: An existing Dask cluster object.

MapWithProgressOptions

Bases: TypedDict

Options for progress bar and Dask scheduler in map_with_progress.

Attributes:

Name	Type	Description
dask_scheduler_name	str	Name for the Dask scheduler (default: "map_with_progress").
tqdm_desc	str	Description for the tqdm progress bar (default: "").
tqdm_unit	str	Unit for the tqdm progress bar (default: "it").

MyProgressBar

Bases: ProgressBar

Show progress of Dask computations.

Copy of distributed.diagnostics.progressbar.TextProgressBar that:

• prints to stderr instead of stdout
• Can have its interval (in seconds) set with TQDM_MININTERVAL environment variable

configure_dask_scheduler(scheduler_address, name, nproc=1)

Context manager that offers a Dask cluster.

If scheduler_address is None then creates a local Dask cluster else returns scheduler_address unchanged and the callee is responsible for cluster cleanup.

Parameters:

Name	Type	Description	Default
scheduler_address	str | Cluster | None	Address of the Dask scheduler to connect to, or None for local cluster.	required
name	str	Name for the Dask cluster.	required
nproc	int	Number of processes to use per worker for CPU support.	1

Yields:

Type	Description
Generator[str | Cluster]	The scheduler address as a string or a cluster.

dask_map_with_progress(client, func, iterable, *args, **kwargs)

Wrapper for map, progress, and gather of Dask that returns a correctly typed list.

Shared positional and keyword arguments are scattered once with broadcast before mapping so Dask does not embed the full payload in every task graph entry.

Environment variables

• Set interval (in seconds) of progress updates with TQDM_MININTERVAL
• Disabled by setting TQDM_DISABLE to any value

Parameters:

Name	Type	Description	Default
client	Client	Dask client.	required
func	Callable[Concatenate[T, P], R]	Function to map; first parameter comes from iterable and any additional parameters can be provided positionally via *args or as keyword arguments via **kwargs.	required
iterable	Collection[T]	Collection of arguments to map over.	required
*args	P.args	Additional positional arguments to scatter once and pass to each mapped call.	()
**kwargs	P.kwargs	Additional keyword arguments to scatter once and pass to each mapped call.	{}

Returns:

Type	Description
list[R]	List of results of type returned by func function.

map_with_progress(scheduler_address, func, iterable, map_with_progress_options=None, *args, **kwargs)

Map a function over an iterable with optional progress bar.

Wraps sequential execution with tqdm and parallel execution with Dask.

Parameters:

Name	Type	Description	Default
scheduler_address	SchedulerAddress	"sequential" for local execution, None for local cluster, or existing cluster.	required
func	Callable[Concatenate[T, P], R]	Function to map.	required
iterable	Collection[T]	Collection of items to map over.	required
map_with_progress_options	MapWithProgressOptions | None	Options for progress bar and Dask scheduler.	None
*args	P.args	Positional arguments passed to func.	()
**kwargs	P.kwargs	Keyword arguments passed to func.	{}

Returns:

Type	Description
list[R]	List of results from func.

nr_cpus()

Determine the number of CPU cores to use.

If the environment variables SLURM_CPUS_PER_TASK or OMP_NUM_THREADS are set, their value is used. Otherwise, the number of physical CPU cores is returned.

Returns:

Type	Description
int	The number of CPU cores to use.

Raises:

Type	Description
ValueError	If the number of physical CPU cores cannot be determined.