Remove distributed macros from ClimaOcean (#403)

* go ahead

* new oceananigans version

* remove distributed utils

* remove distributed utils

* fix tests

Author: simone-silvestri

Project.toml

@@ -41,7 +41,7 @@ JLD2 = "0.4, 0.5"
 KernelAbstractions = "0.9"
 MPI = "0.20"
 NCDatasets = "0.12, 0.13, 0.14"
-Oceananigans = "0.95.25 - 0.99"
+Oceananigans = "0.95.26 - 0.99"
 OffsetArrays = "1.14"
 Scratch = "1"
 SeawaterPolynomials = "0.3.4"

src/Bathymetry.jl

@@ -4,7 +4,7 @@ export regrid_bathymetry, retrieve_bathymetry, download_bathymetry
 
 using ImageMorphology
 using ..DataWrangling: download_progress
-using ..DistributedUtils: @root
+using Oceananigans.DistributedComputations
 
 using Oceananigans
 using Oceananigans.Architectures: architecture, on_architecture
@@ -282,7 +282,7 @@ function regrid_bathymetry(target_grid::DistributedGrid; kw...)
     end
 
     # Synchronize
-    ClimaOcean.global_barrier(arch.communicator)
+    Oceananigans.DistributedComputations.global_barrier(arch.communicator)
 
     # Share the result (can we share SubArrays?)
     bottom_height = all_reduce(+, bottom_height, arch)

src/ClimaOcean.jl

@@ -32,11 +32,7 @@ export
     LinearlyTaperedPolarMask,
     ocean_simulation,
     sea_ice_simulation,
-    initialize!,
-    @root, 
-    @onrank,
-    @distribute,
-    @handshake
+    initialize!
 
 using Oceananigans
 using Oceananigans.Operators: ℑxyᶠᶜᵃ, ℑxyᶜᶠᵃ
@@ -73,7 +69,6 @@ end
     return NamedTuple{names}(vals)
 end
 
-include("DistributedUtils.jl")
 include("OceanSimulations/OceanSimulations.jl")
 include("SeaIceSimulations.jl")
 include("OceanSeaIceModels/OceanSeaIceModels.jl")
@@ -83,7 +78,6 @@ include("DataWrangling/DataWrangling.jl")
 include("Bathymetry.jl")
 include("Diagnostics/Diagnostics.jl")
 
-using .DistributedUtils
 using .VerticalGrids
 using .Bathymetry
 using .DataWrangling

src/DataWrangling/DataWrangling.jl

@@ -15,7 +15,7 @@ using Oceananigans: pretty_filesize, location
 using Oceananigans.Utils: launch!
 using KernelAbstractions: @kernel, @index
 
-using ClimaOcean.DistributedUtils
+using Oceananigans.DistributedComputations
 
 #####
 ##### Downloading utilities

src/DataWrangling/ECCO/ECCO.jl

@@ -5,7 +5,6 @@ export ECCO2Monthly, ECCO4Monthly, ECCO2Daily
 export ECCOFieldTimeSeries, ECCORestoring, LinearlyTaperedPolarMask
 
 using ClimaOcean
-using ClimaOcean.DistributedUtils
 using ClimaOcean.DataWrangling
 using ClimaOcean.DataWrangling: inpaint_mask!, NearestNeighborInpainting, download_progress
 using ClimaOcean.InitialConditions: three_dimensional_regrid!, interpolate!
@@ -14,6 +13,7 @@ using Oceananigans
 using Oceananigans: location
 using Oceananigans.Architectures: architecture, child_architecture
 using Oceananigans.BoundaryConditions
+using Oceananigans.DistributedComputations
 using Oceananigans.DistributedComputations: DistributedField, all_reduce, barrier!
 using Oceananigans.Utils
 

src/DataWrangling/ECCO/ECCO_metadata.jl

@@ -2,6 +2,7 @@ using CFTime
 using Dates
 using ClimaOcean.DataWrangling
 using ClimaOcean.DataWrangling: netrc_downloader, metadata_path, AnyDateTime
+using Oceananigans.DistributedComputations
 
 import Dates: year, month, day
 using Downloads

src/DataWrangling/JRA55/JRA55.jl

@@ -14,7 +14,6 @@ using Oceananigans.Fields: interpolate!
 using Oceananigans.OutputReaders: Cyclical, TotallyInMemory, AbstractInMemoryBackend, FlavorOfFTS, time_indices
 
 using ClimaOcean
-using ClimaOcean.DistributedUtils
 
 using ClimaOcean.OceanSeaIceModels:
     PrescribedAtmosphere,

src/DataWrangling/JRA55/JRA55_metadata.jl

@@ -2,6 +2,8 @@ using CFTime
 using Dates
 using Downloads
 
+using Oceananigans.DistributedComputations
+
 using ClimaOcean.DataWrangling
 using ClimaOcean.DataWrangling: Metadata, metadata_path, download_progress, AnyDateTime
 

src/DistributedUtils.jl

@@ -5,68 +5,11 @@ MPI.Init()
 
 using NCDatasets
 using ClimaOcean.ECCO: download_dataset, metadata_path
+using Oceananigans.DistributedComputations
 using Oceananigans.DistributedComputations: reconstruct_global_grid
 using CFTime
 using Dates
 
-@testset begin
-    rank = MPI.Comm_rank(MPI.COMM_WORLD)
-
-    @onrank 0 begin
-        @test rank == 0
-    end
-
-    @root begin
-        @test rank == 0
-    end
-
-    @onrank 1 begin
-        @test rank == 1
-    end
-
-    @onrank 2 begin
-        @test rank == 2
-    end
-
-    @onrank 3 begin
-        @test rank == 3
-    end
-
-    a = Int[]
-    
-    @distribute for i in 1:10
-        push!(a, i)
-    end
-
-    @root begin
-        @test a == [1, 5, 9]
-    end
-
-    @onrank 1 begin
-        @test a == [2, 6, 10]
-    end
-
-    @onrank 2 begin
-        @test a == [3, 7]
-    end
-
-    @onrank 3 begin
-        @test a == [4, 8]
-    end
-  
-
-    split_comm = MPI.Comm_split(MPI.COMM_WORLD, rank % 2, rank)
-
-    a = Int[]
-    
-    @distribute split_comm for i in 1:10
-        push!(a, i)
-    end
-
-    @onrank split_comm 0 @test a == [1, 3, 5, 7, 9]
-    @onrank split_comm 1 @test a == [2, 4, 6, 8, 10]
-end
-
 @testset "Distributed ECCO download" begin
     dates = DateTimeProlepticGregorian(1992, 1, 1) : Month(1) : DateTimeProlepticGregorian(1994, 4, 1)
     metadata = Metadata(:u_velocity; dataset=ECCO4Monthly(), dates)
@@ -115,9 +58,9 @@ end
                                         z = (0, 1))
 
     global_height = regrid_bathymetry(global_grid; 
-                                    dir = "./",
-                                    filename = "trivial_bathymetry.nc",
-                                    interpolation_passes=10)
+                                      dir = "./",
+                                      filename = "trivial_bathymetry.nc",
+                                      interpolation_passes=10)
 
     arch_x  = Distributed(CPU(), partition=Partition(4, 1))
     arch_y  = Distributed(CPU(), partition=Partition(1, 4))