Distributed DataWrangling

Project.toml

@@ -16,6 +16,7 @@ Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
 ImageMorphology = "787d08f9-d448-5407-9aad-5290dd7ab264"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
+MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
 Oceananigans = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 OrthogonalSphericalShellGrids = "c2be9673-fb75-4747-82dc-aa2bb9f4aed0"
@@ -40,12 +41,12 @@ JLD2 = "0.4, 0.5"
 KernelAbstractions = "0.9"
 NCDatasets = "0.12, 0.13, 0.14"
 Oceananigans = "0.93.1"
-OrthogonalSphericalShellGrids = "0.1.2"
+OrthogonalSphericalShellGrids = "0.1.2, 0.2"
 Scratch = "1"
 SeawaterPolynomials = "0.3.4"
 StaticArrays = "1"
 Statistics = "1.9"
-SurfaceFluxes = "0.11"
+SurfaceFluxes = "0.12"
 Thermodynamics = "0.12"
 julia = "1.9"
 

src/Bathymetry.jl

@@ -3,7 +3,7 @@ module Bathymetry
 export regrid_bathymetry, retrieve_bathymetry
 
 using ImageMorphology
-using ..DataWrangling: download_progress
+using ..DataWrangling: download_progress, @root
 
 using Oceananigans
 using Oceananigans.Architectures: architecture, on_architecture
@@ -87,26 +87,19 @@ function regrid_bathymetry(target_grid;
                            major_basins = Inf) # Allow an `Inf` number of ``lakes''
 
     filepath = joinpath(dir, filename)
-
-    if isfile(filepath)
-        @info "Regridding bathymetry from existing file $filepath."
-    else
-        @info "Downloading bathymetry..."
-        if !ispath(dir)
-            @info "Making bathymetry directory $dir..."
-            mkdir(dir)
-        end
-
-        fileurl = joinpath(url, filename)
-
-        try 
-            Downloads.download(fileurl, filepath; progress=download_progress, verbose=true)
-        catch 
-            cmd = `wget --no-check-certificate -O $filepath $fileurl`
-            run(cmd)
+    fileurl  = joinpath(url, filename)
+
+    @root begin # perform all this only on rank 0, aka the "root" rank
+        if !isfile(filepath)
+            try 
+                Downloads.download(fileurl, filepath; progress=download_progress, verbose=true)
+            catch 
+                cmd = `wget --no-check-certificate -O $filepath $fileurl`
+                @root run(cmd)
+            end
         end
     end
-
+    
     dataset = Dataset(filepath)
 
     FT = eltype(target_grid)

src/ClimaOcean.jl

@@ -18,7 +18,8 @@ export
     JRA55_field_time_series,
     ECCO_field, ECCOMetadata,
     ocean_simulation,
-    initialize!
+    initialize!,
+    @root
 
 using Oceananigans
 using Oceananigans.Operators: ℑxyᶠᶜᵃ, ℑxyᶜᶠᵃ

src/DataWrangling/DataWrangling.jl

@@ -1,8 +1,11 @@
 module DataWrangling
 
+export @root, @distribute
+
 using Oceananigans
 using Downloads
 using Printf
+using Downloads
 
 using Oceananigans.Architectures: architecture, on_architecture
 using Oceananigans.Grids: node
@@ -71,6 +74,7 @@ function save_field_time_series!(fts; path, name, overwrite_existing=false)
     return nothing
 end
 
+include("distributed_datawrangling.jl")
 include("inpaint_mask.jl")
 include("JRA55.jl")
 include("ECCO/ECCO.jl")

src/DataWrangling/ECCO/ECCO.jl

@@ -216,48 +216,24 @@ function inpainted_ECCO_field(metadata::ECCOMetadata;
     inpaint_mask!(f, mask; maxiter)
 
     fill_halo_regions!(f)
-
+    
     return f
 end
 
 inpainted_ECCO_field(variable_name::Symbol; kw...) = inpainted_ECCO_field(ECCOMetadata(variable_name); kw...)
 
-function set!(field::DistributedField, ECCO_metadata::ECCOMetadata; kw...)
+function set!(field::Field, ecco_metadata::ECCOMetadata; kw...)
+
     # Fields initialized from ECCO
     grid = field.grid
-    arch = architecture(grid)
-    child_arch = child_architecture(arch)
-
-    f_ECCO = if arch.local_rank == 0 # Make sure we read/write the file using only one core
-        mask = ECCO_mask(ECCO_metadata, child_arch)
-        inpainted_ECCO_field(ECCO_metadata; mask, architecture = child_arch, kw...)
-    else
-        empty_ECCO_field(ECCO_metadata; architecture = child_arch)
-    end
-
-    barrier!(arch)
-
-    # Distribute ECCO field to all workers
-    parent(f_ECCO) .= all_reduce(+, parent(f_ECCO), arch)
+    arch = child_architecture(grid)
+    mask = ecco_mask(ecco_metadata, arch)
 
-    f_grid = Field(location(ECCO_metadata), grid)   
-    interpolate!(f_grid, f_ECCO)
-    set!(field, f_grid)
-
-    return field
-end
-
-function set!(field::Field, ECCO_metadata::ECCOMetadata; kw...)
+    f = inpainted_ecco_field(ecco_metadata; mask,
+                              architecture = arch,
+                              kw...)
 
-    # Fields initialized from ECCO
-    grid = field.grid
-    arch = architecture(grid)
-    mask = ECCO_mask(ECCO_metadata, arch)
-
-    f = inpainted_ECCO_field(ECCO_metadata; mask, architecture=arch, kw...)
-    f_grid = Field(location(ECCO_metadata), grid)   
-    interpolate!(f_grid, f)
-    set!(field, f_grid)
+    interpolate!(field, f)
 
     return field
 end

src/DataWrangling/ECCO/ECCO_metadata.jl

@@ -1,5 +1,8 @@
 using CFTime
 using Dates
+using ClimaOcean.DataWrangling: distributed_run
+
+import Dates: year, month, day
 
 using Base: @propagate_inbounds
 
@@ -96,6 +99,9 @@ all_ECCO_dates(::ECCO4Monthly) = DateTimeProlepticGregorian(1992, 1, 1) : Month(
 all_ECCO_dates(::ECCO2Monthly) = DateTimeProlepticGregorian(1992, 1, 1) : Month(1) : DateTimeProlepticGregorian(2023, 12, 1)
 all_ECCO_dates(::ECCO2Daily)   = DateTimeProlepticGregorian(1992, 1, 4) : Day(1)   : DateTimeProlepticGregorian(2023, 12, 31)
 
+# File names of metadata containing multiple dates
+metadata_filename(metadata) = [metadata_filename(metadatum) for metadatum in metadata]
+
 # File name generation specific to each Dataset version
 function metadata_filename(metadata::ECCOMetadata{<:AbstractCFDateTime, <:ECCO4Monthly})
     shortname = short_name(metadata)
@@ -190,28 +196,27 @@ ECCO_USERNAME=myuser ECCO_PASSWORD=mypasswrd julia
 function download_dataset!(metadata::ECCOMetadata; url = urls(metadata))
     username = get(ENV, "ECCO_USERNAME", nothing)
     password = get(ENV, "ECCO_PASSWORD", nothing)
-    dir = metadata.dir
 
-    for metadatum in metadata
-        filename = metadata_filename(metadatum)
+    @distribute for metadatum in metadata # Distribute the download among ranks if MPI is initialized
+
+        fileurl  = metadata_url(url, metadatum) 
         filepath = metadata_path(metadatum)
 
         if !isfile(filepath)
             instructions_msg = "\n See ClimaOcean.jl/src/ECCO/README.md for instructions."
             if isnothing(username)
                 msg = "Could not find the ECCO_PASSWORD environment variable. \
                        See ClimaOcean.jl/src/ECCO/README.md for instructions on obtaining \
-                       and setting your ECCO_USERNAME and ECCO_PASSWORD."
+                       and setting your ECCO_USERNAME and ECCO_PASSWORD." * instructions_msg
                 throw(ArgumentError(msg))
             elseif isnothing(password)
                 msg = "Could not find the ECCO_PASSWORD environment variable. \
                        See ClimaOcean.jl/src/ECCO/README.md for instructions on obtaining \
-                       and setting your ECCO_USERNAME and ECCO_PASSWORD."
+                       and setting your ECCO_USERNAME and ECCO_PASSWORD." * instructions_msg
                 throw(ArgumentError(msg))
             end
-
-            fileurl = metadata_url(url, metadatum) 
-            cmd = `wget --http-user=$(username) --http-passwd=$(password) --directory-prefix=$dir $fileurl`
+
+            cmd = `wget --http-user=$(username) --http-passwd=$(password) $(fileurl)`
             run(cmd)
         end
     end

src/DataWrangling/JRA55.jl

@@ -2,7 +2,7 @@ module JRA55
 
 using Oceananigans
 using Oceananigans.Units
- 
+
 using Oceananigans.Architectures: arch_array
 using Oceananigans.DistributedComputations
 using Oceananigans.DistributedComputations: child_architecture
@@ -11,6 +11,8 @@ using Oceananigans.Grids: λnodes, φnodes, on_architecture
 using Oceananigans.Fields: interpolate!
 using Oceananigans.OutputReaders: Cyclical, TotallyInMemory, AbstractInMemoryBackend, FlavorOfFTS, time_indices
 
+using ClimaOcean.DataWrangling: download_progress
+
 using ClimaOcean.OceanSeaIceModels:
     PrescribedAtmosphere,
     TwoBandDownwellingRadiation
@@ -370,8 +372,10 @@ function JRA55_field_time_series(variable_name;
     fts_name = field_time_series_short_names[variable_name]
 
     # Note, we don't re-use existing jld2 files.
-    isfile(filepath) || download(url, filepath)
-    isfile(jld2_filepath) && rm(jld2_filepath)
+    @root begin
+        isfile(filepath) || download(url, filepath)
+        isfile(jld2_filepath) && rm(jld2_filepath)
+    end
 
     # Determine default time indices
     if totally_in_memory

src/DataWrangling/distributed_datawrangling.jl

@@ -0,0 +1,60 @@
+using MPI
+
+#####
+##### Handle commands, typically downloading files
+##### which should be executed by only one rank or distributed among ranks
+#####
+
+function global_barrier()
+    if MPI.Initialized()
+        MPI.Barrier(MPI.COMM_WORLD)
+    end
+end
+
+"""
+Perform `exp` only on rank 0, otherwise know as "root" rank.
+Other ranks will wait for the root rank to finish before continuing
+"""
+macro root(exp)
+    if MPI.Initialized() && MPI.Comm_rank(MPI.COMM_WORLD) != 0
+        command = quote
+          global_barrier()
+        end
+    else
+        command = quote
+           $exp
+           global_barrier()
+        end
+    end
+
+    return esc(command)
+end
+
+""" Distributed a `for` loop among ranks """
+macro distribute(exp)
+    if exp.head != :for
+        error("The `@distribute` macro expects a `for` loop")
+    end
+
+    mpi_initialized = MPI.Initialized()
+    if !mpi_initialized
+        return esc(exp)
+    end
+
+    rank     = MPI.Comm_rank(MPI.COMM_WORLD)
+    nprocs   = MPI.Comm_size(MPI.COMM_WORLD)
+    iterable = exp.args[1].args[2]
+    variable = exp.args[1].args[1]
+    forbody  = exp.args[2]
+
+    new_loop = quote
+        for (counter, $variable) in enumerate($iterable)
+            if $variable % $nprocs == $rank
+                $forbody
+            end
+        end
+        global_barrier()
+    end
+
+    return esc(new_loop)
+end