Add tests for GPU distributed DatasetRestoring

.buildkite/pipeline.yml

@@ -40,7 +40,7 @@ steps:
     env:
       TEST_GROUP: "JRA55"
     commands:
-      - "julia --project -e 'using Pkg; Pkg.test()'"
+      - "julia --project --color=yes -e 'using Pkg; Pkg.test()'"
     agents:
       slurm_mem: 32G
       slurm_cpus_per_task: 8
@@ -51,7 +51,7 @@ steps:
     env:
       TEST_GROUP: "ecco2_monthly"
     commands:
-      - "julia --project -e 'using Pkg; Pkg.test()'"
+      - "julia --project --color=yes -e 'using Pkg; Pkg.test()'"
     agents:
       slurm_mem: 32G
       slurm_cpus_per_task: 8
@@ -62,7 +62,7 @@ steps:
     env:
       TEST_GROUP: "ecco2_daily"
     commands:
-      - "julia --project -e 'using Pkg; Pkg.test()'"
+      - "julia --project --color=yes -e 'using Pkg; Pkg.test()'"
     agents:
       slurm_mem: 32G
       slurm_cpus_per_task: 8
@@ -73,7 +73,7 @@ steps:
     env:
       TEST_GROUP: "ecco4_en4"
     commands:
-      - "julia --project -e 'using Pkg; Pkg.test()'"
+      - "julia --project --color=yes -e 'using Pkg; Pkg.test()'"
     agents:
       slurm_mem: 32G
       slurm_cpus_per_task: 8
@@ -84,7 +84,7 @@ steps:
     env:
       TEST_GROUP: "fluxes"
     commands:
-      - "julia --project -e 'using Pkg; Pkg.test()'"
+      - "julia --project --color=yes -e 'using Pkg; Pkg.test()'"
     agents:
       slurm_mem: 32G
       slurm_cpus_per_task: 8
@@ -95,7 +95,7 @@ steps:
     env:
       TEST_GROUP: "bathymetry"
     commands:
-      - "julia --project -e 'using Pkg; Pkg.test()'"
+      - "julia --project --color=yes -e 'using Pkg; Pkg.test()'"
     agents:
       slurm_mem: 32G
       slurm_cpus_per_task: 8
@@ -107,7 +107,7 @@ steps:
       TEST_GROUP: "JRA55"
       GPU_TEST: "true"
     commands:
-      - "julia --project -e 'using Pkg; Pkg.test()'"
+      - "julia --project --color=yes -e 'using Pkg; Pkg.test()'"
     agents:
       slurm_mem: 32G
       slurm_gpus: 1
@@ -121,7 +121,7 @@ steps:
       TEST_GROUP: "ecco2_monthly"
       GPU_TEST: "true"
     commands:
-      - "julia --project -e 'using Pkg; Pkg.test()'"
+      - "julia --project --color=yes -e 'using Pkg; Pkg.test()'"
     agents:
       slurm_mem: 8G
       slurm_gpus: 1
@@ -135,7 +135,7 @@ steps:
       TEST_GROUP: "ecco2_daily"
       GPU_TEST: "true"
     commands:
-      - "julia --project -e 'using Pkg; Pkg.test()'"
+      - "julia --project --color=yes -e 'using Pkg; Pkg.test()'"
     agents:
       slurm_mem: 8G
       slurm_gpus: 1
@@ -149,7 +149,7 @@ steps:
       TEST_GROUP: "ecco4_en4"
       GPU_TEST: "true"
     commands:
-      - "julia --project -e 'using Pkg; Pkg.test()'"
+      - "julia --project --color=yes -e 'using Pkg; Pkg.test()'"
     agents:
       slurm_mem: 8G
       slurm_gpus: 1
@@ -163,7 +163,7 @@ steps:
       TEST_GROUP: "fluxes"
       GPU_TEST: "true"
     commands:
-      - "julia --project -e 'using Pkg; Pkg.test()'"
+      - "julia --project --color=yes -e 'using Pkg; Pkg.test()'"
     agents:
       slurm_mem: 8G
       slurm_gpus: 1
@@ -177,7 +177,7 @@ steps:
       TEST_GROUP: "bathymetry"
       GPU_TEST: "true"
     commands:
-      - "julia --project -e 'using Pkg; Pkg.test()'"
+      - "julia --project --color=yes -e 'using Pkg; Pkg.test()'"
     agents:
       slurm_mem: 8G
       slurm_gpus: 1
@@ -191,24 +191,38 @@ steps:
       TEST_GROUP: "ocean_sea_ice_model"
       GPU_TEST: "true"
     commands:
-      - "julia --project -e 'using Pkg; Pkg.test()'"
+      - "julia --project --color=yes -e 'using Pkg; Pkg.test()'"
     agents:
       slurm_mem: 10G
       slurm_gpus: 1
       slurm_cpus_per_task: 8
       slurm_ntasks: 1
       slurm_gpus_per_task: 1
 
-  - label: "Run distributed tests"
-    key: "test_distributed"
+  - label: "CPU distributed tests"
+    key: "cpu_distributed_tests"
     env:
       TEST_GROUP: "distributed"
     commands:
-      - "srun julia --project -e 'using Pkg; Pkg.test()'"
+      - "srun julia --project --color=yes -e 'using Pkg; Pkg.test()'"
     agents:
       slurm_mem: 10G
       slurm_cpus_per_task: 1
       slurm_ntasks: 4
 
+  - label: "GPU distributed tests"
+    key: "gpu_distributed_test"
+    env:
+      TEST_GROUP: "distributed"
+      GPU_TEST: "true"
+    commands:
+      - "srun julia --project --color=yes -e 'using Pkg; Pkg.test()'"
+    agents:
+      slurm_mem: 50G
+      slurm_gpus: 4
+      slurm_cpus_per_task: 8
+      slurm_ntasks: 4
+      slurm_gpus_per_task: 1
+
   - wait: ~
     continue_on_failure: true

Project.toml

@@ -1,8 +1,8 @@
 name = "ClimaOcean"
 uuid = "0376089a-ecfe-4b0e-a64f-9c555d74d754"
 license = "MIT"
-authors = ["Climate Modeling Alliance and contributors"]
 version = "0.8.7"
+authors = ["Climate Modeling Alliance and contributors"]
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -38,6 +38,9 @@ Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 SpeedyWeather = "9e226e20-d153-4fed-8a5b-493def4f21a9"
 XESMF = "2e0b0046-e7a1-486f-88de-807ee8ffabe5"
 
+[sources]
+ClimaSeaIce = {rev = "ss/update-to-new-oceananigans", url = "https://github.com/CliMA/ClimaSeaIce.jl.git"}
+
 [extensions]
 ClimaOceanCopernicusMarineExt = "CopernicusMarine"
 ClimaOceanReactantExt = "Reactant"
@@ -60,7 +63,7 @@ KernelAbstractions = "0.9"
 MPI = "0.20"
 MeshArrays = "0.3"
 NCDatasets = "0.12, 0.13, 0.14"
-Oceananigans = "0.101"
+Oceananigans = "0.102.1"
 OffsetArrays = "1.14"
 PrecompileTools = "1"
 Reactant = "0.2.45"

docs/Project.toml

@@ -16,6 +16,6 @@ XESMF = "2e0b0046-e7a1-486f-88de-807ee8ffabe5"
 [compat]
 Documenter = "1"
 DocumenterCitations = "1.3"
-Oceananigans = "0.101"
+Oceananigans = "0.102.1"
 XESMF = "0.1.6"
 Literate = "2.2"

examples/ecco_mixed_layer_depth.jl

@@ -6,7 +6,7 @@ using Printf
 using Dates
 
 using SeawaterPolynomials: TEOS10EquationOfState
-using Oceananigans.BuoyancyFormulations: buoyancy
+using Oceananigans.Models: buoyancy_force
 
 arch = CPU()
 Nx = 360
@@ -44,7 +44,7 @@ Nt = length(ht)
 for n = 1:Nt-1
     local tracers
     tracers = (T=Tt[n], S=St[n])
-    h.operand.buoyancy_perturbation = buoyancy(sb, grid, tracers)
+    h.operand.buoyancy_perturbation = buoyancy_force(sb, grid, tracers)
     @show n
     @time compute!(h)
     parent(ht[n]) .= parent(h)

src/DataWrangling/metadata.jl

@@ -1,7 +1,7 @@
 using CFTime
 using Dates
 using Base: @propagate_inbounds
-import Oceananigans.Utils: prettysummary
+import Oceananigans.Grids: prettysummary
 
 struct BoundingBox{X, Y, Z}
     longitude :: X

src/DataWrangling/restoring.jl

@@ -2,7 +2,7 @@ using Oceananigans: location
 using Oceananigans.Grids: node
 using Oceananigans.Fields: interpolate, instantiated_location
 using Oceananigans.OutputReaders: Cyclical
-using Oceananigans.Utils: Time
+using Oceananigans.Units: Time
 using Oceananigans.Architectures: AbstractArchitecture
 
 using JLD2

src/Diagnostics/Diagnostics.jl

@@ -4,7 +4,7 @@ export MixedLayerDepthField, MixedLayerDepthOperand
 
 using Oceananigans
 using Oceananigans.Architectures: architecture
-using Oceananigans.BuoyancyFormulations: buoyancy
+using Oceananigans.Models: buoyancy_operation
 using Oceananigans.Grids: new_data, inactive_cell, znode
 using Oceananigans.BoundaryConditions: FieldBoundaryConditions, fill_halo_regions!
 using Oceananigans.Fields: FieldStatus

src/Diagnostics/mixed_layer_depth.jl

@@ -8,7 +8,7 @@ end
 Base.summary(mldo::MixedLayerDepthOperand) = "MixedLayerDepthOperand"
 
 function MixedLayerDepthOperand(bm, grid, tracers; difference_criterion=1e-4)
-    buoyancy_perturbation = buoyancy(bm, grid, tracers)
+    buoyancy_perturbation = buoyancy_operation(bm, grid, tracers)
     difference_criterion = convert(eltype(grid), difference_criterion)
     return MixedLayerDepthOperand(buoyancy_perturbation, difference_criterion)
 end

src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl

@@ -20,7 +20,8 @@ using Oceananigans: HydrostaticFreeSurfaceModel, architecture
 using Oceananigans.Grids: inactive_node, node, topology
 using Oceananigans.BoundaryConditions: fill_halo_regions!
 using Oceananigans.Fields: ConstantField, interpolate, FractionalIndices
-using Oceananigans.Utils: launch!, Time, KernelParameters
+using Oceananigans.Utils: launch!, KernelParameters
+using Oceananigans.Units: Time
 using Oceananigans.Operators: ℑxᶜᵃᵃ, ℑyᵃᶜᵃ, ℑxᶠᵃᵃ, ℑyᵃᶠᵃ
 
 using KernelAbstractions: @kernel, @index

src/OceanSeaIceModels/InterfaceComputations/similarity_theory_turbulent_fluxes.jl

@@ -1,4 +1,4 @@
-using Oceananigans.Utils: prettysummary
+using Oceananigans.Grids: prettysummary
 using Oceananigans.Grids: AbstractGrid
 
 using Adapt

src/OceanSeaIceModels/InterfaceComputations/tabulated_albedo.jl

@@ -1,6 +1,6 @@
 using Oceananigans.Fields: interpolator
 using Oceananigans.Grids: on_architecture
-using Oceananigans.Utils: Time
+using Oceananigans.Units: Time
 using Base
 
 using ClimaOcean.OceanSeaIceModels:

src/OceanSeaIceModels/OceanSeaIceModels.jl

@@ -14,7 +14,8 @@ using SeawaterPolynomials
 
 using Oceananigans
 using Oceananigans.Operators
-using Oceananigans.Utils: launch!, Time, KernelParameters
+using Oceananigans.Utils: launch!, KernelParameters
+using Oceananigans.Units: Time
 using Oceananigans.Architectures: architecture
 using Oceananigans.BoundaryConditions: fill_halo_regions!, BoundaryCondition
 using Oceananigans.Grids: architecture

src/OceanSeaIceModels/PrescribedAtmospheres.jl

@@ -2,10 +2,10 @@ module PrescribedAtmospheres
 
 using Oceananigans
 using Oceananigans.Fields: Center
-using Oceananigans.Grids: grid_name
+using Oceananigans.Grids: grid_name, prettysummary
 using Oceananigans.OutputReaders: FieldTimeSeries, update_field_time_series!, extract_field_time_series
 using Oceananigans.TimeSteppers: Clock, tick!
-using Oceananigans.Utils: prettysummary, Time
+using Oceananigans.Units: Time
 
 using Adapt
 using Thermodynamics.Parameters: AbstractThermodynamicsParameters

src/OceanSimulations/ocean_simulation.jl

@@ -98,6 +98,7 @@ end
 """
     ocean_simulation(grid;
                      Δt = estimate_maximum_Δt(grid),
+                     clock = Clock(grid),
                      closure = default_ocean_closure(),
                      tracers = (:T, :S),
                      free_surface = default_free_surface(grid),
@@ -124,7 +125,6 @@ consistent defaults for advection, closures, the equation of state, surface flux
 barotropic pressure–gradient forcing, boundary conditions, and optional biogeochemistry.
 It then wraps the model into an Oceananigans's `Simulation` with the specified timestepping options.
 
-
 ## Behaviour and automatic configuration
 
 ### Coriolis
@@ -159,6 +159,7 @@ defaults on a per-field basis.
 ## Keyword Arguments
 
 - `Δt`: Timestep used by the `Simulation`. Defaults to the maximum stable timestep estimated from the `grid`.
+- `clock`: Clock object. Defaults to `Clock(grid)`.
 - `closure`: A turbulence or mixing closure. Defaults to `default_ocean_closure()`.
 - `tracers`: Tuple of tracer names. Defaults to `(:T, :S)`.
 - `free_surface`: Free–surface solver. Defaults to `default_free_surface(grid)`.
@@ -180,6 +181,7 @@ defaults on a per-field basis.
 """
 function ocean_simulation(grid;
                           Δt = estimate_maximum_Δt(grid),
+                          clock = Clock(grid),
                           closure = default_ocean_closure(),
                           tracers = (:T, :S),
                           free_surface = default_free_surface(grid),
@@ -307,6 +309,7 @@ function ocean_simulation(grid;
     end
 
     ocean_model = HydrostaticFreeSurfaceModel(; grid,
+                                              clock,
                                               buoyancy,
                                               closure,
                                               biogeochemistry,

test/runtests.jl

@@ -108,6 +108,7 @@ if test_group == :ocean_sea_ice_model || test_group == :all
 end
 
 if test_group == :distributed || test_group == :all
+    distributed_archs = distributed_test_architectures()
     include("test_distributed_utils.jl")
 end