Rework sharding tests

.github/workflows/ci.yml

@@ -40,22 +40,23 @@ jobs:
         arch:
           - x64
     steps:
-      - run: |
-          touch LocalPreferences.toml
-
-          echo "[Reactant]" >> LocalPreferences.toml
-          echo "xla_runtime = \"IFRT\"" >> LocalPreferences.toml
-
-          cat LocalPreferences.toml     
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v2
         with:
           version: ${{ matrix.version }}
           arch: ${{ matrix.arch }}
       - uses: julia-actions/cache@v2
+      - run: |
+          touch LocalPreferences.toml
+
+          echo "[Reactant]" >> LocalPreferences.toml
+          echo "xla_runtime = \"IFRT\"" >> LocalPreferences.toml
+
+          cat LocalPreferences.toml    
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
         env:
+          XLA_FLAGS: "--xla_force_host_platform_device_count=4"
           JULIA_DEBUG: "Reactant, Reactant_jll"
           REACTANT_TEST: true
           TEST_GROUP: "sharding"
@@ -110,31 +111,6 @@ jobs:
         env:
           TEST_GROUP: "turbulence_closures"
 
-  reactant:
-    name: Reactant - Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
-    runs-on: ${{ matrix.os }}
-    timeout-minutes: 120
-    strategy:
-      fail-fast: false
-      matrix:
-        version:
-          - '1.10'
-        os:
-          - ubuntu-latest
-        arch:
-          - x64
-    steps:
-      - uses: actions/checkout@v4
-      - uses: julia-actions/setup-julia@v2
-        with:
-          version: ${{ matrix.version }}
-          arch: ${{ matrix.arch }}
-      - uses: julia-actions/cache@v2
-      - uses: julia-actions/julia-buildpkg@v1
-      - uses: julia-actions/julia-runtest@v1
-        env:
-          TEST_GROUP: "reactant"
-
   metal:
     name: Metal - Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
     runs-on: ${{ matrix.os }}

ext/OceananigansReactantExt/Fields.jl

@@ -9,6 +9,7 @@ using Oceananigans.Fields: Field, interior
 using KernelAbstractions: @index, @kernel
 
 import Oceananigans.Fields: set_to_field!, set_to_function!, set!
+import Oceananigans.DistributedComputations: reconstruct_global_field
 
 import ..OceananigansReactantExt: deconcretize
 import ..Grids: ReactantGrid
@@ -17,6 +18,8 @@ import ..Grids: ShardedGrid
 const ReactantField{LX, LY, LZ, O} = Field{LX, LY, LZ, O, <:ReactantGrid}
 const ShardedDistributedField{LX, LY, LZ, O} = Field{LX, LY, LZ, O, <:ShardedGrid}
 
+reconstruct_global_field(field::ShardedDistributedField) = field
+
 deconcretize(field::Field{LX, LY, LZ}) where {LX, LY, LZ} =
     Field{LX, LY, LZ}(field.grid,
                       deconcretize(field.data),

ext/OceananigansReactantExt/Grids/serial_grids.jl

@@ -33,7 +33,7 @@ function constant_with_arch(cpu_grid::AbstractUnderlyingGrid, arch)
 end
 
 function constant_with_arch(cpu_ibg::CPUImmersedBoundaryGrid, arch)
-    underlying_grid = constant_with_reactant_state(cpu_ibg.underlying_grid, arch)
+    underlying_grid = constant_with_arch(cpu_ibg.underlying_grid, arch)
     TX, TY, TZ = Oceananigans.Grids.topology(cpu_ibg)
     return ImmersedBoundaryGrid{TX, TY, TZ}(underlying_grid,
                                             cpu_ibg.immersed_boundary,

src/DistributedComputations/distributed_grids.jl

@@ -346,7 +346,7 @@ insert_connected_topology(::Type{Bounded}, R, r) = ifelse(R == 1, Bounded,
 insert_connected_topology(::Type{Periodic}, R, r) = ifelse(R == 1, Periodic, FullyConnected)
 
 """ 
-    reconstruct_global_topology(T, R, r, comm)
+    reconstruct_global_topology(T, R, r, r1, r2, comm)
 
 reconstructs the global topology associated with the local topologies `T`, the amount of ranks 
 in `T` direction (`R`) and the local rank index `r`. If all ranks hold a `FullyConnected` topology,
@@ -362,7 +362,7 @@ function reconstruct_global_topology(T, R, r, r1, r2, comm)
         topologies[r] = 1
     end
 
-    topologies = all_reduce(topologies, +, comm)
+    all_reduce!(topologies, +, comm)
 
     if sum(topologies) == R
         return Periodic

src/DistributedComputations/partition_assemble.jl

@@ -5,6 +5,9 @@ import Oceananigans.Architectures: on_architecture
 all_reduce(op, val, arch::Distributed) = MPI.Allreduce(val, op, arch.communicator)
 all_reduce(op, val, arch) = val
 
+all_reduce!(op, val, arch::Distributed) = MPI.Allreduce!(val, op, arch.communicator)
+all_reduce!(op, val, arch) = nothing
+
 # MPI Barrier
 barrier!(arch) = nothing
 barrier!(arch::Distributed) = MPI.Barrier(arch.communicator)

test/distributed_tests_utils.jl

@@ -104,17 +104,21 @@ function run_distributed_latitude_longitude_grid(arch, filename)
 
     distributed_grid = ImmersedBoundaryGrid(distributed_grid, GridFittedBottom(bottom_height))
     model = run_distributed_simulation(distributed_grid)
-
+    
     η = reconstruct_global_field(model.free_surface.η)
     u = reconstruct_global_field(model.velocities.u)
     v = reconstruct_global_field(model.velocities.v)
     c = reconstruct_global_field(model.tracers.c)
 
+    # Check also that the bottom height is reconstructed correctly!
+    b = reconstruct_global_field(model.grid.immersed_boundary.bottom_height)
+
     if arch.local_rank == 0
         jldsave(filename; u = Array(interior(u, :, :, 10)),
                           v = Array(interior(v, :, :, 10)), 
                           c = Array(interior(c, :, :, 10)),
-                          η = Array(interior(η, :, :, 1))) 
+                          η = Array(interior(η, :, :, 1)),
+                          b = Array(parent(b))[:, :, 1]) 
     end
 
     return nothing
@@ -124,19 +128,19 @@ end
 function run_distributed_simulation(grid)
 
     model = HydrostaticFreeSurfaceModel(; grid = grid,
-                                          free_surface = SplitExplicitFreeSurface(grid; substeps = 20),
+                                          free_surface = ExplicitFreeSurface(), # SplitExplicitFreeSurface(grid; substeps = 20),
                                           tracers = :c,
                                           buoyancy = nothing, 
-                                          tracer_advection = WENO(), 
-                                          momentum_advection = WENOVectorInvariant(order=3),
-                                          coriolis = HydrostaticSphericalCoriolis())
+                                          tracer_advection = nothing, #WENO(), 
+                                          momentum_advection = nothing, #WENOVectorInvariant(order=3),
+                                          coriolis = nothing) # HydrostaticSphericalCoriolis())
 
     # Setup the model with a gaussian sea surface height
     # near the physical north poles and one near the equator
     ηᵢ(λ, φ, z) = exp(- (φ - 90)^2 / 10^2) + exp(- φ^2 / 10^2)
     set!(model, c=ηᵢ, η=ηᵢ)
 
-    Δt = 5minutes
+    Δt = 10 # 5minutes
     arch = architecture(grid)
     if arch isa ReactantState || arch isa Distributed{<:ReactantState}
         @info "Compiling first_time_step..."

test/run_sharding_tests.jl

@@ -0,0 +1,30 @@
+# We need to initiate MPI for sharding because we are using a multi-host implementation: 
+# i.e. we are launching the tests with `mpiexec` and on Github actions the default MPI 
+# implementation is MPICH which requires calling MPI.Init(). In the case of OpenMPI,
+# MPI.Init() is not necessary.
+
+using MPI 
+MPI.Init()
+include("distributed_tests_utils.jl")
+
+if Base.ARGS[1] == "tripolar"
+    run_function = run_distributed_tripolar_grid
+    suffix = "trg"
+else
+    run_function = run_distributed_latitude_longitude_grid
+    suffix = "llg"
+end
+
+Reactant.Distributed.initialize(; single_gpu_per_process=false)
+
+arch = Distributed(ReactantState(), partition = Partition(4, 1))
+filename = "distributed_xslab_$(suffix).jld2"
+run_function(arch, filename)
+
+arch = Distributed(ReactantState(), partition = Partition(1, 4))
+filename = "distributed_yslab_$(suffix).jld2"
+run_function(arch, filename)
+
+arch = Distributed(ReactantState(), partition = Partition(2, 2))
+filename = "distributed_pencil_$(suffix).jld2"
+run_function(arch, filename)

test/test_reactant.jl

@@ -108,7 +108,7 @@ ridge(λ, φ) = 0.1 * exp((λ - 2)^2 / 2)
                                         latitude = [0, 1, 2, 3, 4],
                                         z = (0, 1))
 
-        constant_llg = OceananigansReactantExt.constant_with_reactant_state(cpu_llg)
+        constant_llg = OceananigansReactantExt.constant_with_arch(cpu_llg)
 
         for name in propertynames(constant_llg)
             p = getproperty(constant_llg, name)
@@ -124,7 +124,7 @@ ridge(λ, φ) = 0.1 * exp((λ - 2)^2 / 2)
 
         @info "  Testing constantified ImmersedBoundaryGrid construction [$FT]..."
         cpu_ibg = ImmersedBoundaryGrid(cpu_llg, GridFittedBottom(ridge))
-        constant_ibg = OceananigansReactantExt.constant_with_reactant_state(cpu_ibg)
+        constant_ibg = OceananigansReactantExt.constant_with_arch(cpu_ibg)
         @test architecture(constant_ibg) isa ReactantState
         @test architecture(constant_ibg.immersed_boundary.bottom_height) isa CPU
 
@@ -153,7 +153,7 @@ ridge(λ, φ) = 0.1 * exp((λ - 2)^2 / 2)
                                                 z = (0, 1))
 
         @info "    Replacing architecture with ReactantState [$FT]..."
-        constant_rllg = OceananigansReactantExt.constant_with_reactant_state(cpu_rllg)
+        constant_rllg = OceananigansReactantExt.constant_with_arch(cpu_rllg)
 
         for name in propertynames(constant_rllg)
             p = getproperty(constant_rllg, name)
@@ -164,7 +164,7 @@ ridge(λ, φ) = 0.1 * exp((λ - 2)^2 / 2)
 
         @info "  Testing constantified immersed OrthogonalSphericalShellGrid construction [$FT]..."
         cpu_ribg = ImmersedBoundaryGrid(cpu_rllg, GridFittedBottom(ridge))
-        constant_ribg = OceananigansReactantExt.constant_with_reactant_state(cpu_ribg)
+        constant_ribg = OceananigansReactantExt.constant_with_arch(cpu_ribg)
         @test architecture(constant_ribg) isa ReactantState
         @test architecture(constant_ribg.immersed_boundary.bottom_height) isa CPU
     end

test/test_sharded_lat_lon.jl

@@ -1,41 +1,9 @@
-using JLD2
-using Oceananigans
-using Oceananigans.DistributedComputations: reconstruct_global_field, reconstruct_global_grid
-using Oceananigans.Units
-using Reactant
-using Random
-using Test
-
+include("dependencies_for_runtests.jl")
 include("distributed_tests_utils.jl")
 
-run_xslab_distributed_grid = """
-    using MPI 
-    MPI.Init()
-    include("distributed_tests_utils.jl")
-    Reactant.Distributed.initialize(; single_gpu_per_process=false)
-    arch = Distributed(ReactantState(), partition = Partition(4, 1))
-    run_distributed_latitude_longitude_grid(arch, "distributed_xslab_llg.jld2")
-"""
-
-run_yslab_distributed_grid = """
-    using MPI 
-    MPI.Init()
-    include("distributed_tests_utils.jl")
-    Reactant.Distributed.initialize(; single_gpu_per_process=false)
-    arch = Distributed(ReactantState(), partition = Partition(1, 4))
-    run_distributed_latitude_longitude_grid(arch, "distributed_yslab_llg.jld2")
-"""
+Nhosts = 1
 
-run_pencil_distributed_grid = """
-    using MPI 
-    MPI.Init()
-    include("distributed_tests_utils.jl")
-    Reactant.Distributed.initialize(; single_gpu_per_process=false)
-    arch = Distributed(ReactantState(), partition = Partition(2, 2))
-    run_distributed_latitude_longitude_grid(arch, "distributed_pencil_llg.jld2")
-"""
-
-@testset "Test distributed LatitudeLongitudeGrid simulations..." begin
+@testset "Test sharded LatitudeLongitudeGrid simulations..." begin
     # Run the serial computation    
     Random.seed!(1234)
     bottom_height = - rand(40, 40, 1) .* 500 .- 500
@@ -48,67 +16,54 @@ run_pencil_distributed_grid = """
     us, vs, ws = model.velocities
     cs = model.tracers.c
     ηs = model.free_surface.η
+    bs = model.grid.immersed_boundary.bottom_height
 
     us = interior(us, :, :, 10)
     vs = interior(vs, :, :, 10)
     cs = interior(cs, :, :, 10)
     ηs = interior(ηs, :, :, 1)
+    bs = parent(bs)[:, :, 1]
 
-    # Run the distributed grid simulation with a pencil configuration
-    write("distributed_xslab_llg_tests.jl", run_xslab_distributed_grid)
-    run(`$(mpiexec()) -n 4 $(Base.julia_cmd()) --project -O0 distributed_xslab_llg_tests.jl`)
-    rm("distributed_xslab_llg_tests.jl")
-
-    # Retrieve Parallel quantities
-    up = jldopen("distributed_xslab_llg.jld2")["u"]
-    vp = jldopen("distributed_xslab_llg.jld2")["v"]
-    ηp = jldopen("distributed_xslab_llg.jld2")["η"]
-    cp = jldopen("distributed_xslab_llg.jld2")["c"]
-
-    # rm("distributed_xslab_llg.jld2")
-
-    @test all(us .≈ up)
-    @test all(vs .≈ vp)
-    @test all(cs .≈ cp)
-    @test all(ηs .≈ ηp)
-
-    # Run the distributed grid simulation with a slab configuration
-    write("distributed_yslab_llg_tests.jl", run_yslab_distributed_grid)
-    run(`$(mpiexec()) -n 4 $(Base.julia_cmd()) --project -O0 distributed_yslab_llg_tests.jl`)
-    rm("distributed_yslab_llg_tests.jl")
-
-    # Retrieve Parallel quantities
-    up = jldopen("distributed_yslab_llg.jld2")["u"]
-    vp = jldopen("distributed_yslab_llg.jld2")["v"]
-    cp = jldopen("distributed_yslab_llg.jld2")["c"]
-    ηp = jldopen("distributed_yslab_llg.jld2")["η"]
-
-    # rm("distributed_yslab_llg.jld2")
-
-    # Test slab partitioning
-    @test all(us .≈ up)
-    @test all(vs .≈ vp)
-    @test all(cs .≈ cp)
-    @test all(ηs .≈ ηp)
-
-    # We try now with more ranks in the x-direction. This is not a trivial
-    # test as we are now splitting, not only where the singularities are, but
-    # also in the middle of the north fold. This is a more challenging test
-    write("distributed_pencil_llg_tests.jl", run_pencil_distributed_grid)
-    run(`$(mpiexec()) -n 4 julia --project -O0 distributed_pencil_llg_tests.jl`)
-    rm("distributed_pencil_llg_tests.jl")
+    # Run the distributed grid simulations in all the configurations
+    run(`$(mpiexec()) -n $(Nhosts) $(Base.julia_cmd()) --project -O0 run_sharding_tests.jl "latlon"`)
 
     # Retrieve Parallel quantities
-    up = jldopen("distributed_pencil_llg.jld2")["u"]
-    vp = jldopen("distributed_pencil_llg.jld2")["v"]
-    ηp = jldopen("distributed_pencil_llg.jld2")["η"]
-    cp = jldopen("distributed_pencil_llg.jld2")["c"]
-
-    # rm("distributed_pencil_llg.jld2")
-
-    @test all(us .≈ up)
-    @test all(vs .≈ vp)
-    @test all(cs .≈ cp)
-    @test all(ηs .≈ ηp)
-end
-
+    bp1 = jldopen("distributed_xslab_llg.jld2")["b"]
+    up1 = jldopen("distributed_xslab_llg.jld2")["u"]
+    vp1 = jldopen("distributed_xslab_llg.jld2")["v"]
+    cp1 = jldopen("distributed_xslab_llg.jld2")["c"]
+    ηp1 = jldopen("distributed_xslab_llg.jld2")["η"]
+
+    bp2 = jldopen("distributed_yslab_llg.jld2")["b"]
+    up2 = jldopen("distributed_yslab_llg.jld2")["u"]
+    vp2 = jldopen("distributed_yslab_llg.jld2")["v"]
+    cp2 = jldopen("distributed_yslab_llg.jld2")["c"]
+    ηp2 = jldopen("distributed_yslab_llg.jld2")["η"]
+
+    bp3 = jldopen("distributed_pencil_llg.jld2")["b"]
+    up3 = jldopen("distributed_pencil_llg.jld2")["u"]
+    vp3 = jldopen("distributed_pencil_llg.jld2")["v"]
+    cp3 = jldopen("distributed_pencil_llg.jld2")["c"]
+    ηp3 = jldopen("distributed_pencil_llg.jld2")["η"]
+
+    @info "Testing xslab partitioning..."
+    @test all(bs .≈ bp1)
+    @test all(us .≈ up1)
+    @test all(vs .≈ vp1)
+    @test all(cs .≈ cp1)
+    @test all(ηs .≈ ηp1)
+
+    @info "Testing yslab partitioning..."
+    @test all(bs .≈ bp2)
+    @test all(us .≈ up2)
+    @test all(vs .≈ vp2)
+    @test all(cs .≈ cp2)
+    @test all(ηs .≈ ηp2)
+
+    @info "Testing pencil partitioning..."
+    @test all(bs .≈ bp3)
+    @test all(us .≈ up3)
+    @test all(vs .≈ vp3)
+    @test all(cs .≈ cp3)
+    @test all(ηs .≈ ηp3)
+end