NEDAS.utils.parallel module

class NEDAS.utils.parallel.Comm

Bases: object

Communicator class supporting both serial and MPI programs.

When the python program is started with MPI environment, for example:

$ mpirun -n 10 python -m mpi4py program.py

A communicator can be obtained from the mpi4py package:

>>> from mpi4py import MPI
>>> comm = MPI.COMM_WORLD

However, when the program is run in

Variables:

parallel_io (bool) – If netCDF4.Dataset is built with parallel I/O support.

mpi_ready: bool = False

parallel_io: bool

init_file_lock(filename)

Register a filename that THIS rank will personally acquire/release the lock for. Actual chain construction is deferred to build_file_locks(), which must be called collectively afterward (once every rank has registered its own files) and before any acquire_file_lock()/release_file_lock() call.

Unlike the old RMA design, callers should register only the files each rank itself intends to write – NOT the global union of every file across all ranks (build_file_locks() does its own internal allgather to reconstruct the full per-file writer ordering).

Parameters:

filename (str) – Path to the file.

build_file_locks()

Build the handoff chain for every filename registered via init_file_lock() since the last build_file_locks()/ cleanup_file_locks(). Must be called collectively by every rank (each rank may have registered a different, possibly empty, set of files) before any acquire_file_lock()/release_file_lock() call.

See the class-level comment above self._locks for the design.

check_parallel_io() → bool

Check if netCDF4 is built with parallel I/O support.

Returns:

True if netCDF4 module support parallel I/O mode.

Return type:

bool

finish_file_locks()

Send the handoff to each registered file’s successor, once this rank is done with all its own writes for the current generation.

Must be called by every rank BEFORE any barrier that other ranks’ pending acquire_file_lock() calls might be blocking on. Sending the handoff from cleanup_file_locks() instead (i.e. after such a barrier) deadlocks whenever a file has more than one writer (e.g. nproc_mem<nproc, where several rec-groups sharing a member all write into that member’s one file): the successor, still blocked inside acquire_file_lock() waiting for this rank’s handoff, can never reach the barrier itself, so the predecessor never gets past it either to send anything.

cleanup_file_locks()

Clear all lock bookkeeping for the current generation. Purely local (no MPI calls) – call finish_file_locks() first to send any outstanding handoffs.

acquire_file_lock(filename)

release_file_lock(filename)

finalize()

Clean up MPI resources cleanly to avoid hangs on exit.

NEDAS.utils.parallel.abort_all_ranks(comm: Comm | None = None, code: int = 1) → None

Terminate the whole MPI job (all ranks) if running under MPI; a plain sys.exit() already works fine for a genuinely serial run, so this only takes the MPI path when comm says it’s actually ready for it.

An uncaught exception on one rank would otherwise just kill that rank’s own process; every other rank keeps running until it reaches the next collective call shared with the dead rank (e.g. a gather/bcast downstream) and blocks there forever, since the dead rank can never arrive to participate – turning a single-rank failure into a silent, multi-node hang that only a SLURM walltime timeout or a human noticing and cancelling the job would end. MPI.COMM_WORLD.Abort() reaches every rank regardless of which sub-communicator (comm_mem/comm_rec, see core/context.py) the code was using at the time.

comm should be the process’s existing Comm instance (e.g. scheme.c.comm), not a fresh one – Comm.mpi_ready reflects whether this process actually detected and successfully imported mpi4py at startup (core/context.py’s set_comm()). Probing for MPI from scratch here (a bare from mpi4py import MPI) is unsafe: a process that isn’t truly part of an MPI communicator (e.g. NEDAS’s own single-process driver, which merely inherits SLURM/PMI environment variables from the surrounding sbatch allocation without itself being launched under srun/mpirun) can still have those env vars set, so a fresh import triggers a real (and here, failing) MPI_Init attempt instead of a clean exit. Reusing comm._MPI (the module reference saved once mpi_ready detection already succeeded) avoids re-triggering that initialization.

Call this from a bare except: block in a script’s top-level main(), after printing/logging the traceback yourself (Abort() does not unwind normally, so the usual automatic traceback print does not happen).

class NEDAS.utils.parallel.DummyComm

Bases: object

Dummy communicator for python without mpi

Get_size()

Get_rank()

Barrier()

Abort(code: int)

Split(color=0, key=0)

bcast(obj, root=0)

send(obj, dest, tag)

recv(source, tag)

allgather(obj)

gather(obj, root=0)

allreduce(obj)

reduce(obj, root=0)

NEDAS.utils.parallel.by_rank(comm: Comm, rank: int) → Callable[[Callable[[P], T]], Callable[[P], T | None]]

Decorator for func() to be run only by rank 0 in comm

NEDAS.utils.parallel.bcast_by_root(comm: Comm) → Callable[[Callable[[P], T]], Callable[[P], T]]

Decorator for func() to be run only by rank 0 in comm, and result of func() is then broadcasted to all other ranks.

NEDAS.utils.parallel.distribute_tasks(comm: Comm, tasks: ndarray | Sequence, load: ndarray | Sequence | None = None) → dict[int, list]

Divide a list of task indices and assign a subset to each rank in comm

Parameters:

• comm (Comm) – MPI communicator

• tasks (ArrayLike) – List of task indices (to be distributed over the processors)

• load (np.ndarray, optional) – Amount of workload for each task element The default is None, we will let tasks have equal workload

Returns:

Dictionary {rank:list}, list is the subset of tasks for the processor rank

calling this function to work on

Return type:

dict

class NEDAS.utils.parallel.OfflineScheduler(c, nworker: int, walltime: int | None = None, check_dt: float = 0.1, debug: bool = False)

Bases: object

An offline scheduler class for queuing and running multiple jobs on available workers (group of processors). The jobs are submitted by one processor with the scheduler, while the job.run code is calling subprocess to be run on the worker

submit_job(name: str, job: Callable, *args, **kwargs) → None

Submit a job to the scheduler, hold info in jobs dict.

Parameters:

• name (str) – unique name to identify this job

• job (Callable) – callable with is_running and kill methods

• *args – passed into job()

• **kwargs – passed into job()

monitor_job_queue() → None

Monitor the available_workers and pending_jobs, assign a job to a worker if possible Monitor the running_jobs for jobs that are finished, kill jobs that exceed walltime, and move the finished jobs to completed_jobs

start_queue()

Start the job queue, and wait for jobs to complete

shutdown()