Support / roadmap for parallel MODFLOW 6 via MODFLOW API (initialize_mpi)?

[ifariasg]

Question
Is running a simulation through modflowapi (using callbacks + run_simulation + BMI/XMI loop) supported when using a parallel (MPI-enabled) build of MODFLOW 6? If not, is this on the roadmap, and is there a recommended pattern (e.g. optional initialize_mpi in the runner)?

I did some digging and found the following:

1. MODFLOW 6 exposes an XMI entry point initialize_mpi (C name initialize_mpi) when built with __WITH_MPI__. It sets the Fortran MPI communicator and simulation_mode = 'PARALLEL', then calls the normal BMI initialize (bmi_initialize). See MODFLOW 6 mf6xmi.F90 – xmi_initialize_mpi / bind(C, name="initialize_mpi").

2. xmipy XmiWrapper wraps that symbol as initialize_mpi(self, value: int) (passing the Fortran communicator handle, e.g. from mpi4py). See Deltares xmipy xmiwrapper.py – initialize_mpi.

3. run_simulation currently always calls mf6.initialize() and never initialize_mpi. modflowapi extensions/runner.py (initialization right after ModflowApi(...) is constructed).

All in all, it all comes down to whether run_simulation is expected to support MPI, and if a small extension such as run_simulation(..., mpi_comm_fortran=None, number_mpi_ranks = int) that calls initialize_mpi would be enough.

I'm not sure if this is more of an API question or an extendend MF6 question but maybe @mjr-deltares could also shed some light?

Thanks in advance!

[mjr-deltares]

Hi @ifariasg , thanks for posting this. You are right that this is currently not supported in the modflowapi (Python). Your suggestion to extend the run_simulation procedure makes sense, except that the number of ranks is probably not needed because the MPI world should be created before calling the routine, by running your script with mpirun -np <nranks> python my_parallel_script.py.
Now the code inside MODFLOW has been organized to operate in lock step across multiple MPI processes to avoid dead locks and race conditions. I guess the callbacks will have to be considered with similar care. Other than that, it might just work. Maybe you can give it a try yourself, locally modifying the run_simulation routine to included the MPI communicator, and then running your script with mpirun. Make sure to use the shared library from the extended build though. Also, to create a communicator in Python and pass it on into the MODFLOW core, check out this sample:

from mpi4py import MPI
...
mf6_api = modflowapi.ModflowApi(lib_path="path/to/libmf6.so")
comm_f90 = int(MPI.COMM_WORLD.py2f())
mf6_api.initialize_mpi(comm_f90)
...

Let me know if that works for you

[ifariasg]

Ok, so good news @mjr-deltares, I got a working version of the runner for the API using parallel execution. It was indeed sort of straightforward. The most difficult part was getting MPI to work on windows but it ended up working also using the precompiled executable from the nightly builds.

The one thing I'm not sure is working as intended or if its by design, with initialize_mpi alone, the BMI/XMI driver does not synchronize timestep length across ranks. After each prepare_time_step, DELT can differ between partitions when ATS or other local constraints apply. Ranks then advance TOTIM / PERTIM inconsistently and can fall out of sync on stress period and timestep index. The run may continue without crashing, but exchange coupling is then inconsistent and results are not trustworthy.

The workaround was to take the minimum DELT over ranks after each prepare_time_step and use the API to set DELT, TOTIM, and PERTIM on ranks that used a larger step. Then, after finalize_time_step, use an MPI_Barrier and broadcast get_current_time() so the Python driver uses one time on all ranks.

I had to rely heavily on AI assisted coding to get this to work so I'm not at all sure how "clean" this solution is, but I can share it if you want.

[mjr-deltares]

Hi @ifariasg , thanks for being our most-valued beta-tester on this! ;-) It certainly looks like you have stumbled upon a bit of unsupported functionality. ATS is generally not compatible with parallel simulations yet, and we have to add that to the release notes, I guess. Additionally, a recent issue has come up when using the API combined with ATS as well (MODFLOW-ORG/modflow6#2756), so it might be worth checking if your setup and work around is ready to deal with that. If you want you can post your python simulation loop here, so we can have a look.

I will discuss with the team and create issues for both.

[ifariasg]

Hi @mjr-deltares, indeed I experienced the same issue described in MODFLOW-ORG/modflow6#2756 and this solution is not capable of overcoming it. I had to rely on a sensible selection of the ATS + IMS settings. If the solver settings are not robust the simulation will just continue to attempt iterations endlessly.

So, the base logic would be the following:
in modflowapi extensions/runner.py
Add dependencies:

from mpi4py import MPI
import numpy as np.

1. initialize_mpi (Fortran communicator from MPI_COMM_WORLD)
   Use this instead of mf6.initialize() when running under MPI. Pass the Fortran handle from mpi4py:

mf6 = ModflowApi(dll, working_directory=sim_path)
mf6.initialize_mpi(int(MPI.COMM_WORLD.py2f()))
sim = ApiSimulation.load(mf6)

2. MPI timestep sync: global DELT + Barrier / bcast for the driver clock
   Before the main time loop (optional but useful if ranks disagree on the initial time):

current_time = float(mf6.get_current_time())
end_time = float(mf6.get_end_time())
end_tol = max(1e-6, abs(end_time) * 1e-9) ## needed to prevent extra SPs, but needs tweaking

if MPI.COMM_WORLD.Get_size() > 1:
  current_time = float(MPI.COMM_WORLD.bcast(current_time, root=0))

Immediately after mf6.prepare_time_step(dt) (inside the loop):

if MPI.COMM_WORLD.Get_size() > 1:
  local_dt = float(mf6.get_time_step())
  send = np.array(local_dt, dtype=np.float64)
  recv = np.empty(1, dtype=np.float64)
  MPI.COMM_WORLD.Allreduce(send, recv, op=MPI.MIN)
  global_dt = float(recv[0])

  if abs(global_dt - local_dt) >= 1e-15 * max(abs(global_dt), 1.0): ## Condition to evaluate if DELT is different across ranks 
      diff = local_dt - global_dt
      addr_delt = mf6.get_var_address("DELT", "TDIS")
      addr_totim = mf6.get_var_address("TOTIM", "TDIS")
      addr_pertim = mf6.get_var_address("PERTIM", "TDIS")

      mf6.set_value(addr_delt, np.array([global_dt], dtype=np.float64))

      totim_arr = mf6.get_value(addr_totim)
      totim_arr[0] -= diff
      mf6.set_value(addr_totim, totim_arr)

      pertim_arr = mf6.get_value(addr_pertim)
      pertim_arr[0] -= diff
      mf6.set_value(addr_pertim, pertim_arr)

Immediately after mf6.finalize_time_step() (still inside the loop):

if MPI.COMM_WORLD.Get_size() > 1:
  MPI.COMM_WORLD.Barrier()
  buf = float(mf6.get_current_time()) if MPI.COMM_WORLD.Get_rank() == 0 else None
  current_time = float(MPI.COMM_WORLD.bcast(buf, root=0))
else:
  current_time = float(mf6.get_current_time())

if abs(current_time - end_time) < end_tol:
  current_time = end_time

Annoyingly, using the API to run the simulation I'm using to test (partitioned or not), had a secondary issue of float accumulation drift over many timesteps in TOTIM which causes a potential extra stress period. That is the reason for the extra "railguards" with the times.
This probably should be enforced in the API at the ATS level instead of the MPI runner, but as of today, it is required. Sadly, since its a problem that accumulates over long runs, its difficult to create a minimal reproducible example for it, but since there is some discussion already about ATS in the API due to MODFLOW-ORG/modflow6#2756, I think it should also be addressed as part of it.

[mjr-deltares]

issue created: MODFLOW-ORG/modflow6#2762

[ifariasg]

Hi @mjr-deltares , I just wanted to update that with the fix you implemented for the ATS in parallel https://github.com/MODFLOW-ORG/modflow6/pull/2765 most of the nasty issues from trying to combine the API+Parallel are solved. Then, only the runner needs to be updated with a call to initialize_mpi instead of the normal initialize and it works! I attached the runner I'm currently using if that helps to get a more robust fix.

run_simulation_api_parallel.py

The only issue that I see left has to do with TOTIM floating point drift when using ATS. This affects the usual loop driver while current_time < end_time: which ends up generating and extra stress period. I do believe this is a more general modflow 6 problem, so I will open an issue there (EDIT: MODFLOW-ORG/modflow6#2779).
If that is not fixed though, the driving loop of the callback function should be either updated to use the ENDOFSIMULATION flag or KPER. For example:

addr_eos = mf6.get_var_address("ENDOFSIMULATION", "TDIS")
while not bool(mf6.get_value(addr_eos)[0]):
    ...
    mf6.finalize_time_step()

[mjr-deltares]

thanks @ifariasg for reporting back on this, and glad this (partly) helps your issue! I will leave this issue open because the original part about extending run_simulation for parallel still stands.