Parallel simulations using multiple nodes/MPI

[fspereira1]

Dear All,

I am new to oceananigans, so my apologies for any basic questions. I want to run oceananigans on an HPC machine using several nodes (MPI). I wonder if anyone could help me with this. What is the best way of doing this?

One of your reference cases is quite close to my problem:

[https://github.com/CliMA/Oceananigans.jl/blob/main/examples/ocean_wind_mixing_and_convection.jl]

Best,
Filipe

[glwagner]

@fspereira1, do you mind if I convert this to a discussion? We like to use discussions for usage and syntax questions (in part to make it easier to find for future users, and in part because then we won't need to "close" the discussion).

[fspereira1]

Sure.

[fspereira1]

and I can provide the .jl file I am trying to run in parallel

[glwagner]

The short answer is that distributed support is experimental, which is why there is little documentation or information about it.

Here's a short summary:

• The HydrostaticFreeSurfaceModel can be run on multiple GPUs via MultiRegion, but only in a single process / on one node. With 4 V100s, for example, the maximum simulation size will be something like 1024^2 x 512.

• The NonhydrostaticModel support distributed CPU simulations, but not distributed GPU simulations. I believe there are some usability issues there (see Usability concerns and possible improvements for MultiArch #2349). Some prototypes can be found in https://github.com/CliMA/Oceananigans.jl/tree/main/validation/distributed_simulations.

• Note that on one V100 with 32 GB, you can run NonhydrostaticModel with one tracer up to 512 x 512 x 384 (with more tracers, the size will have to be reduced). I think @tomchor has done some exploration running the absolutely largest simulations that can be afforded on one GPU. Slightly larger models can be run if you have access to an A100 with 80 GB (which may also be ~3x faster).

• There is also a PR (Distributed mixed FFT / vertical tridiagonal solver #2538, getting a bit ancient now), which extends the distributed FFT algorithm to cover all cases for NonhydrostaticModel. Last I checked the code on this PR seemed to work, but the tests were not passing for some reason.

• You'll likely get best results on GPUs, since we haven't focused much on CPU performance. However, if you don't have access to GPUs, then you of course have to consider CPU solutions.

Hope that helps.

[fspereira1]

I am trying to run

pressure_solver = FourierTridiagonalPoissonSolver(grid, FFTW.ESTIMATE)
model = NonhydrostaticModel(; grid, pressure_solver, buoyancy,
advection = UpwindBiasedFifthOrder(),
timestepper = :RungeKutta3,
tracers = (:T, :S),
coriolis = FPlane(f=1e-4),
closure = AnisotropicMinimumDissipation(),
boundary_conditions = (u=u_bcs, T=T_bcs, S=S_bcs))

just on CPUs

[glwagner]

Check out this script for inspiration:

https://github.com/CliMA/Oceananigans.jl/blob/main/validation/distributed_simulations/mpi_nonhydrostatic_two_dimensional_turbulence.jl

I believe on main you can only distribute grids that are regular in all directions. I think you'll need #2538 to distribute a stretched vertical grid.

[tomchor]

• The NonhydrostaticModel support distributed CPU simulations, but not distributed GPU simulations. I believe there are some usability issues there (see Usability concerns and possible improvements for MultiArch #2349). Some prototypes can be found in https://github.com/CliMA/Oceananigans.jl/tree/main/validation/distributed_simulations.

It's a shame we haven't been able to merge that PR. IIRC there's only one test that hangs and never completes, forcing us to kill the test manually. I tried debugging it at some point, but since it doesn't throw any errors I couldn't quite figure it out. Maybe we should try to revive that PR since it's such a nice addition.

[glwagner]

Contributions very welcome on #2538 !

[tomchor]

@fspereira1 if you're trying to run large models on a GPU let me know and I can give you some tips. There's some things you can do when defining diagnostics that can save quite a bit of space in my experience (i.e. reusing scratch space).

[iuryt]

Hi all!
I know that it is being a while since this discussion started and maybe things have changed.
I noticed that in the example Nranks is set to the y direction,

Oceananigans.jl/validation/distributed_simulations/distributed_nonhydrostatic_turbulence.jl

Line 29 in 2946027

arch = Distributed(CPU(); topology, ranks=(1, Nranks, 1), communicator=comm)

but each rank is running a chunk in the x direction

Oceananigans.jl/validation/distributed_simulations/distributed_nonhydrostatic_turbulence.jl

Line 30 in 2946027

grid = RectilinearGrid(arch; topology, size=(Nx ÷ Nranks, Ny), halo=(3, 3), x=(0, 2π), y=(0, 2π))

Is this correct?

[iuryt]

hm.. I will check that, but I was having the following error while trying to run
Attempting to use an MPI routine before initializing MPICH

[iuryt]

Okay, so, MPI needs to be initialized before accessing its communication variables.

Thus, despite calling Distributed initializes MPI, we need to access Nranks = MPI.Comm_size(comm) before to define Partition.

Oceananigans.jl/validation/distributed_simulations/distributed_nonhydrostatic_turbulence.jl

Lines 20 to 22 in 2946027

	comm = MPI.COMM_WORLD
	rank = MPI.Comm_rank(comm)
	Nranks = MPI.Comm_size(comm)

Should we initialize MPI when importing DistributedComputations?

[iuryt]

I will transfer this discussion to a PR.

[iuryt]

We can continue the discussion on #3429

[glwagner]

Taking a step back, we should not have to write using MPI to do distributed simulations.

If we do require it for some specialized purpose (or to demonstrate something, as in a validation example), then perhaps we can right a comment about this in the script.

[Gaspartcho]

Hello, I have basically the same problem as @fspereira1, but thanks to all the resources on this page, I managed to run a few simulations in parallels. (Using up to 4 nodes).

My problem comes after: If I understood well, how it works is that Oceanigans splits the main grid (here RectilinearGrid) into multiple "chunks" on which each node will do the simulation, and then each node saves its results on its own file. However, I want to combine all these results in one big image with the CairoMakie library (as seen in the examples), and so I need to combine all the results.

Is there an easy way to do that?

[Gaspartcho]

Hello @simone-silvestri, thanks for your answer. Unfortunately, the workaround does not seem to work for me.

My code is:

time_series = (
    w  = FieldTimeSeries(filename, "w"),
    T  = FieldTimeSeries(filename, "T"),
    S  = FieldTimeSeries(filename, "S"),
    νₑ = FieldTimeSeries(filename, "νₑ")
)

# Start Time
times = time_series.w.times
intro = searchsortedfirst(times, 10minutes)

n = Observable(intro)

wₙ  = @lift begin
      data = interior(time_series.w[$n], :, :, 1)
      global_data = construct_global_array(arch, data, size(data))
      global_data
end

With

arch = Distributed(
    CPU();
    partition = Partition(Nranks, 1, 1),
    communicator = comm
)

The rest of the code is greatly inspired from this example.

When I try to run it using 4 cores, I get the following error (at the line defining global_data):

ERROR: LoadError: BoundsError: attempt to access Tuple{Vector{Int64}, Vector{Int64}} at index [3]

What am I doing wrong here?

[Gaspartcho]

UPDATE: I managed to make it work: it turns out that the construct_global_array function has a bug where it assumes that the data is one dimension superior to what it actually is.

I changed my code to this and now everything seems to work as expected.

wₙ  = @lift begin
      data_w   = interior(time_series.w[$n], :, 1, :)
      global_w = construct_global_array(arch, data_w, (size(data_w)..., 1))
      global_w
end

[tomchor]

we are in the process of developing a utility

Don't we simply want to extend FieldTimeSeries to work when a field is saved across multiple files? I don't think we want a new utility, just to extend the utilities we already have...

It would be nice to come up with a utility that merges everything at some point, at least for NetCDFs. My impression is that most Oceananigans users still use NetCDF since its become the default in geosciences over the years. (Ofc, depending on how the chunks are written, some NetCDFs readers may alreayd be capable of doing that automatically when reading; I haven't tested.)

[glwagner]

sure, we also need FieldTimeSeries for NetCDF

[simone-silvestri]

thanks @Gaspartcho I ll take a look at the code!