Add a Reactant Extension (#438)

* Add a Reactant Extension

* Clean up extension plus a few fixes

* Update .buildkite/pipeline.yml

* Update .buildkite/pipeline.yml

* Worry about IFRT later

* MAke sure GPU test has a GPU

* Fix silly test issue

* Import things needed

Author: glwagner

.buildkite/pipeline.yml

@@ -158,5 +158,32 @@ steps:
       slurm_cpus_per_task: 1
       slurm_ntasks: 4
 
+  - label: "CPU Reactant tests"
+    key: "cpu_reactant_tests"
+    env:
+      TEST_GROUP: "reactant"
+      GPU_TEST: "false"
+    commands:
+      - "srun julia --project -e 'using Pkg; Pkg.test()'"
+    agents:
+      slurm_mem: 10G
+      slurm_cpus_per_task: 1
+      slurm_ntasks: 4
+
+  - label: "GPU Reactant tests"
+    key: "gpu_reactant_tests"
+    env:
+      TEST_GROUP: "reactant"
+      GPU_TEST: "true"
+    commands:
+      - "srun julia --project -e 'using Pkg; Pkg.test()'"
+    agents:
+      slurm_mem: 8G
+      slurm_gpus: 1
+      slurm_cpus_per_task: 8
+      slurm_ntasks: 1
+      slurm_gpus_per_task: 1
+
   - wait: ~
     continue_on_failure: true
+

Project.toml

@@ -2,7 +2,7 @@ name = "ClimaOcean"
 uuid = "0376089a-ecfe-4b0e-a64f-9c555d74d754"
 license = "MIT"
 authors = ["Climate Modeling Alliance and contributors"]
-version = "0.5.6"
+version = "0.5.7"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -29,6 +29,12 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 SurfaceFluxes = "49b00bb7-8bd4-4f2b-b78c-51cd0450215f"
 Thermodynamics = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c"
 
+[weakdeps]
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
+
+[extensions]
+ClimaOceanReactantExt = "Reactant"
+
 [compat]
 Adapt = "4"
 CFTime = "0.1"
@@ -45,6 +51,7 @@ NCDatasets = "0.12, 0.13, 0.14"
 Oceananigans = "0.96 - 0.99"
 OffsetArrays = "1.14"
 PrecompileTools = "1"
+Reactant = "0.2.45"
 Scratch = "1"
 SeawaterPolynomials = "0.3.5"
 StaticArrays = "1"
@@ -60,4 +67,4 @@ MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Coverage", "Test", "MPIPreferences", "CUDA_Runtime_jll"]
+test = ["Coverage", "Test", "MPIPreferences", "CUDA_Runtime_jll", "Reactant"]

ext/ClimaOceanReactantExt.jl

@@ -0,0 +1,24 @@
+module ClimaOceanReactantExt
+
+using Reactant
+using Oceananigans.Architectures: ReactantState
+using Oceananigans.DistributedComputations: Distributed
+
+using ClimaOcean: OceanSeaIceModel
+
+import Oceananigans
+import Oceananigans.Models: initialization_update_state!
+
+const OceananigansReactantExt = Base.get_extension(
+     Oceananigans, :OceananigansReactantExt
+)
+
+const ReactantOSIM{I, A, O, F, C} = Union{
+    OceanSeaIceModel{I, A, O, F, C, <:ReactantState},
+    OceanSeaIceModel{I, A, O, F, C, <:Distributed{ReactantState}},
+}
+
+initialization_update_state!(model::ReactantOSIM) = nothing
+
+end # module ClimaOceanReactantExt
+

src/OceanSeaIceModels/InterfaceComputations/atmosphere_ocean_fluxes.jl

@@ -108,7 +108,7 @@ end
     downwelling_radiation = (; Qs, Qℓ)
 
     # Estimate initial interface state
-    FT = eltype(grid)
+    FT = typeof(Tᵢ)
     u★ = convert(FT, 1e-4)
 
     # Estimate interface specific humidity using interior temperature

src/OceanSeaIceModels/InterfaceComputations/interface_states.jl

@@ -47,7 +47,8 @@ SpecificHumidityFormulation(phase) = SpecificHumidityFormulation(phase, nothing)
     q★ = Thermodynamics.q_vap_saturation_from_density(ℂₐ, convert(CT, Tₛ), convert(CT, ρₛ), p★)
 
     # Compute saturation specific humidity according to Raoult's law
-    return q★ * x_H₂O
+    FT = eltype(Tₛ)
+    return convert(FT, q★ * x_H₂O)
 end
 
 struct SalinityConstituent{FT}
@@ -309,7 +310,7 @@ struct InterfaceState{FT}
     melting :: Bool
 end
 
-InterfaceState(u★, θ★, q★, u, v, T, S, q) =
+@inline InterfaceState(u★, θ★, q★, u, v, T, S, q) =
     InterfaceState(u★, θ★, q★, u, v, T, S, q, false)
 
 Base.eltype(::InterfaceState{FT}) where FT = FT

test/runtests.jl

@@ -64,3 +64,7 @@ if test_group == :distributed || test_group == :all
     include("test_distributed_utils.jl")
 end
 
+if test_group == :reactant || test_group == :all
+    include("test_reactant.jl")
+end
+

test/test_reactant.jl

@@ -0,0 +1,40 @@
+using Test
+using Reactant
+using Oceananigans.Models: initialization_update_state!
+using Oceananigans.Architectures: ReactantState
+using ClimaOcean
+
+gpu_test = get(ENV, "GPU_TEST", "false") == "true"
+
+if gpu_test
+    Reactant.set_default_backend("gpu")
+else
+    Reactant.set_default_backend("cpu")
+end
+
+@testset "Reactant extension tests" begin
+    arch = ReactantState()
+    grid = LatitudeLongitudeGrid(arch;
+                                 size = (256, 128, 10),
+                                 longitude = (0, 360),
+                                 latitude = (-80, 80),
+                                 halo = (7, 7, 7),
+                                 z = (-6000, 0))
+
+    ocean = ocean_simulation(grid)
+    backend = JRA55NetCDFBackend(4)
+    atmosphere = JRA55PrescribedAtmosphere(arch; backend)
+    radiation = Radiation(arch)
+    coupled_model = OceanSeaIceModel(ocean; atmosphere, radiation)
+
+    # Test that Reactant does _not_ initialize in the constructor for OceanSeaIceModel
+    exchange_state = coupled_model.interfaces.exchanger.exchange_atmosphere_state
+    atmos_exchanger = coupled_model.interfaces.exchanger.atmosphere_exchanger
+    @test all(atmos_exchanger.i .== 0)
+    @test all(atmos_exchanger.j .== 0)
+
+    # This tests that update_state! is not called
+    ue = exchange_state.u
+    @test all(ue .== 0) # not initialized with Reactant
+end
+