Merge pull request #739 from CliMA/js/gpu_tests

tests on GPU - parent usage

Author: juliasloan25

.buildkite/pipeline.yml

@@ -133,6 +133,14 @@ steps:
         agents:
           slurm_mem: 16GB
 
+  - group: "GPU: unit tests and global bucket"
+    steps:
+      - label: "GPU runtests"
+        command: "julia --color=yes --project=test/ test/runtests.jl"
+        agents:
+          slurm_ntasks: 1
+          slurm_gres: "gpu:1"
+
   - group: "Integration Tests"
     steps:
 

test/checkpointer_tests.jl

@@ -20,9 +20,9 @@ Checkpointer.get_model_prog_state(sim::DummySimulation) = sim.state
 end
 
 @testset "checkpoint_model_state, restart_model_state!" begin
-    boundary_space = TestHelper.create_space(FT)
-    t = 1
     comms_ctx = ClimaComms.context(ClimaComms.CPUSingleThreaded())
+    boundary_space = TestHelper.create_space(FT; comms_ctx = comms_ctx)
+    t = 1
     # old sim run
     sim = DummySimulation(CC.Fields.FieldVector(T = ones(boundary_space)))
     Checkpointer.checkpoint_model_state(sim, comms_ctx, t, output_dir = "test_checkpoint")

test/component_model_tests/bucket_tests.jl

@@ -28,18 +28,17 @@ for FT in (Float32, Float64)
             ),
             t = FT(0),
         )
-        integrator_copy = deepcopy(integrator)
         sim = BucketSimulation(nothing, nothing, nothing, integrator, nothing)
 
         # make fields non-constant to check the impact of the dss step
-        for i in eachindex(parent(sim.integrator.u.state_field_2d))
-            parent(sim.integrator.u.state_field_2d)[i] = sin(i)
-        end
-        for i in eachindex(parent(sim.integrator.u.state_field_3d))
-            parent(sim.integrator.u.state_field_3d)[i] = sin(i)
-        end
+        coords_lat = CC.Fields.coordinate_field(sim.integrator.u.state_field_2d).lat
+        @. sim.integrator.u.state_field_2d = sin(coords_lat)
+
+        coords_lat = CC.Fields.coordinate_field(sim.integrator.u.state_field_3d).lat
+        @. sim.integrator.u.state_field_3d = sin(coords_lat)
 
         # apply DSS
+        integrator_copy = deepcopy(integrator)
         dss_state!(sim)
 
         # test that uniform field and cache are unchanged, non-constant is changed

test/component_model_tests/climaatmos_tests.jl

@@ -18,14 +18,13 @@ for FT in (Float32, Float64)
             ),
             p = (; cache_field = FT.(CC.Fields.zeros(boundary_space))),
         )
-        integrator_copy = deepcopy(integrator)
         sim = ClimaAtmosSimulation(nothing, nothing, nothing, integrator)
 
         # make field non-constant to check the impact of the dss step
-        for i in eachindex(parent(sim.integrator.u.state_field2))
-            parent(sim.integrator.u.state_field2)[i] = FT(sin(i))
-        end
+        coords_lat = CC.Fields.coordinate_field(sim.integrator.u.state_field2).lat
+        @. sim.integrator.u.state_field2 = sin(coords_lat)
 
+        integrator_copy = deepcopy(integrator)
         # apply DSS
         dss_state!(sim)
 

test/component_model_tests/prescr_seaice_tests.jl

@@ -75,15 +75,14 @@ for FT in (Float32, Float64)
             ),
             p = (; cache_field = FT.(CC.Fields.zeros(boundary_space)), dss_buffer = dss_buffer),
         )
-        integrator_copy = deepcopy(integrator)
         sim = PrescribedIceSimulation(nothing, nothing, nothing, integrator)
 
         # make field non-constant to check the impact of the dss step
-        for i in eachindex(parent(sim.integrator.u.state_field2))
-            parent(sim.integrator.u.state_field2)[i] = FT(sin(i))
-        end
+        coords_lat = CC.Fields.coordinate_field(sim.integrator.u.state_field2).lat
+        @. sim.integrator.u.state_field2 = sin(coords_lat)
 
         # apply DSS
+        integrator_copy = deepcopy(integrator)
         dss_state!(sim)
 
         # test that uniform field and cache are unchanged, non-constant is changed

test/component_model_tests/slab_ocean_tests.jl

@@ -21,15 +21,14 @@ for FT in (Float32, Float64)
             ),
             p = (; cache_field = FT.(CC.Fields.zeros(boundary_space)), dss_buffer = dss_buffer),
         )
-        integrator_copy = deepcopy(integrator)
         sim = SlabOceanSimulation(nothing, nothing, nothing, integrator)
 
         # make field non-constant to check the impact of the dss step
-        for i in eachindex(parent(sim.integrator.u.state_field2))
-            parent(sim.integrator.u.state_field2)[i] = FT(sin(i))
-        end
+        coords_lat = CC.Fields.coordinate_field(sim.integrator.u.state_field2).lat
+        @. sim.integrator.u.state_field2 = sin(coords_lat)
 
         # apply DSS
+        integrator_copy = deepcopy(integrator)
         dss_state!(sim)
 
         # test that uniform field and cache are unchanged, non-constant is changed

test/field_exchanger_tests.jl

@@ -105,11 +105,11 @@ for FT in (Float32, Float64)
             coupler_fields =
                 NamedTuple{coupler_names}(ntuple(i -> CC.Fields.zeros(boundary_space), length(coupler_names)))
             FieldExchanger.import_atmos_fields!(coupler_fields, model_sims, boundary_space, t)
-            @test parent(coupler_fields.F_turb_energy)[1] == results[i]
-            @test parent(coupler_fields.F_turb_moisture)[1] == results[i]
-            @test parent(coupler_fields.F_radiative)[1] == results[1]
-            @test parent(coupler_fields.P_liq)[1] == results[1]
-            @test parent(coupler_fields.P_snow)[1] == results[1]
+            @test Array(parent(coupler_fields.F_turb_energy))[1] == results[i]
+            @test Array(parent(coupler_fields.F_turb_moisture))[1] == results[i]
+            @test Array(parent(coupler_fields.F_radiative))[1] == results[1]
+            @test Array(parent(coupler_fields.P_liq))[1] == results[1]
+            @test Array(parent(coupler_fields.P_snow))[1] == results[1]
         end
     end
 
@@ -138,12 +138,12 @@ for FT in (Float32, Float64)
             coupler_fields =
                 NamedTuple{coupler_names}(ntuple(i -> CC.Fields.zeros(boundary_space), length(coupler_names)))
             FieldExchanger.import_combined_surface_fields!(coupler_fields, sims, boundary_space, t)
-            @test parent(coupler_fields.T_S)[1] == results[1]
-            @test parent(coupler_fields.surface_direct_albedo)[1] == results[1]
-            @test parent(coupler_fields.surface_diffuse_albedo)[1] == results[1]
-            @test parent(coupler_fields.z0m_S)[1] == results[i]
-            @test parent(coupler_fields.z0b_S)[1] == results[i]
-            @test parent(coupler_fields.beta)[1] == results[i]
+            @test Array(parent(coupler_fields.T_S))[1] == results[1]
+            @test Array(parent(coupler_fields.surface_direct_albedo))[1] == results[1]
+            @test Array(parent(coupler_fields.surface_diffuse_albedo))[1] == results[1]
+            @test Array(parent(coupler_fields.z0m_S))[1] == results[i]
+            @test Array(parent(coupler_fields.z0b_S))[1] == results[i]
+            @test Array(parent(coupler_fields.beta))[1] == results[i]
         end
     end
 
@@ -203,31 +203,31 @@ for FT in (Float32, Float64)
             FieldExchanger.update_model_sims!(model_sims, coupler_fields, t)
 
             # test atmos
-            @test parent(model_sims.atmos_sim.cache.albedo_direct)[1] == results[2]
-            @test parent(model_sims.atmos_sim.cache.albedo_diffuse)[1] == results[3]
+            @test Array(parent(model_sims.atmos_sim.cache.albedo_direct))[1] == results[2]
+            @test Array(parent(model_sims.atmos_sim.cache.albedo_diffuse))[1] == results[3]
             if t isa FluxCalculator.CombinedStateFluxes
-                @test parent(model_sims.atmos_sim.cache.roughness_momentum)[1] == results[2]
+                @test Array(parent(model_sims.atmos_sim.cache.roughness_momentum))[1] == results[2]
             else
-                @test parent(model_sims.atmos_sim.cache.roughness_momentum)[1] == results[1]
+                @test Array(parent(model_sims.atmos_sim.cache.roughness_momentum))[1] == results[1]
             end
 
             # unspecified variables
-            @test parent(model_sims.atmos_sim.cache.surface_temperature)[1] == results[1]
-            @test parent(model_sims.atmos_sim.cache.beta)[1] == results[1]
-            @test parent(model_sims.atmos_sim.cache.roughness_buoyancy)[1] == results[1]
+            @test Array(parent(model_sims.atmos_sim.cache.surface_temperature))[1] == results[1]
+            @test Array(parent(model_sims.atmos_sim.cache.beta))[1] == results[1]
+            @test Array(parent(model_sims.atmos_sim.cache.roughness_buoyancy))[1] == results[1]
 
             # test surface
-            @test parent(model_sims.land_sim.cache.turbulent_energy_flux)[1] == results[2] # assuming units / m2
-            @test parent(model_sims.land_sim.cache.turbulent_moisture_flux)[1] == results[2]
+            @test Array(parent(model_sims.land_sim.cache.turbulent_energy_flux))[1] == results[2] # assuming units / m2
+            @test Array(parent(model_sims.land_sim.cache.turbulent_moisture_flux))[1] == results[2]
 
             # unspecified variables
-            @test parent(model_sims.land_sim.cache.radiative_energy_flux_sfc)[1] == results[1]
-            @test parent(model_sims.land_sim.cache.liquid_precipitation)[1] == results[1]
-            @test parent(model_sims.land_sim.cache.snow_precipitation)[1] == results[1]
+            @test Array(parent(model_sims.land_sim.cache.radiative_energy_flux_sfc))[1] == results[1]
+            @test Array(parent(model_sims.land_sim.cache.liquid_precipitation))[1] == results[1]
+            @test Array(parent(model_sims.land_sim.cache.snow_precipitation))[1] == results[1]
 
             # test stub - albedo should be updated by update_sim!
-            @test parent(model_sims.stub_sim.cache.albedo_direct)[1] == results[2]
-            @test parent(model_sims.stub_sim.cache.albedo_diffuse)[1] == results[2]
+            @test Array(parent(model_sims.stub_sim.cache.albedo_direct))[1] == results[2]
+            @test Array(parent(model_sims.stub_sim.cache.albedo_diffuse))[1] == results[2]
 
         end
     end

test/flux_calculator_tests.jl

@@ -1,6 +1,7 @@
 import Test: @test, @testset, @test_throws
 import StaticArrays
 import ClimaCore as CC
+import ClimaParams
 import Thermodynamics as TD
 import Thermodynamics.Parameters.ThermodynamicsParameters
 import SurfaceFluxes.Parameters.SurfaceFluxesParameters
@@ -149,7 +150,7 @@ for FT in (Float32, Float64)
         for (i, t) in enumerate(flux_types)
             sim.cache.flux .= FT(0)
             FluxCalculator.combined_turbulent_fluxes!(model_sims, coupler_fields, t)
-            @test parent(sim.cache.flux)[1] ≈ results[i]
+            @test Array(parent(sim.cache.flux))[1] ≈ results[i]
         end
         sim2 = DummySimulation2((; cache = (; flux = zeros(boundary_space))))
         model_sims = (; atmos_sim = sim2)
@@ -270,7 +271,7 @@ for FT in (Float32, Float64)
                 @test parent(fields.F_turb_energy[colidx]) == -parent(atmos_sim.integrator.p.energy_bc[colidx])
 
             end
-            @test parent(fields.F_turb_moisture)[1] ≈ FT(0)
+            @test Array(parent(fields.F_turb_moisture))[1] ≈ FT(0)
         end
     end
 

test/interfacer_tests.jl

@@ -57,7 +57,7 @@ for FT in (Float32, Float64)
         for sim in (DummySimulation(space), DummySimulation2(space), DummySimulation3(space))
             # field
             colidx = CC.Fields.ColumnIndex{2}((1, 1), 73)
-            @test parent(Interfacer.get_field(sim, Val(:var), colidx))[1] == FT(1)
+            @test Array(parent(Interfacer.get_field(sim, Val(:var), colidx)))[1] == FT(1)
             # float
             @test Interfacer.get_field(sim, Val(:var_float), colidx) == FT(2)
         end
@@ -108,10 +108,10 @@ for FT in (Float32, Float64)
         Interfacer.update_field!(stub, Val(:surface_direct_albedo), ones(boundary_space) .* 3)
         Interfacer.update_field!(stub, Val(:surface_diffuse_albedo), ones(boundary_space) .* 4)
 
-        @test parent(Interfacer.get_field(stub, Val(:area_fraction)))[1] == FT(1)
-        @test parent(Interfacer.get_field(stub, Val(:surface_temperature)))[1] == FT(2)
-        @test parent(Interfacer.get_field(stub, Val(:surface_direct_albedo)))[1] == FT(3)
-        @test parent(Interfacer.get_field(stub, Val(:surface_diffuse_albedo)))[1] == FT(4)
+        @test Array(parent(Interfacer.get_field(stub, Val(:area_fraction))))[1] == FT(1)
+        @test Array(parent(Interfacer.get_field(stub, Val(:surface_temperature))))[1] == FT(2)
+        @test Array(parent(Interfacer.get_field(stub, Val(:surface_direct_albedo))))[1] == FT(3)
+        @test Array(parent(Interfacer.get_field(stub, Val(:surface_diffuse_albedo))))[1] == FT(4)
     end
 end
 

test/runtests.jl

@@ -1,18 +1,21 @@
 import SafeTestsets: @safetestset
+import ClimaComms
+
+gpu_broken = ClimaComms.device() isa ClimaComms.CUDADevice
 
 @safetestset "Aqua tests" begin
     include("aqua.jl")
 end
 @safetestset "Interfacer tests" begin
     include("interfacer_tests.jl")
 end
-@safetestset "Regridder tests" begin
+gpu_broken || @safetestset "Regridder tests" begin
     include("regridder_tests.jl")
 end
 @safetestset "ConservationChecker tests" begin
     include("conservation_checker_tests.jl")
 end
-@safetestset "BCReader tests" begin
+gpu_broken || @safetestset "BCReader tests" begin
     include("bcreader_tests.jl")
 end
 @safetestset "Utilities tests" begin
@@ -24,19 +27,19 @@ end
 @safetestset "FieldExchanger tests" begin
     include("field_exchanger_tests.jl")
 end
-@safetestset "FluxCalculator tests" begin
+gpu_broken || @safetestset "FluxCalculator tests" begin
     include("flux_calculator_tests.jl")
 end
-@safetestset "Diagnostics tests" begin
+gpu_broken || @safetestset "Diagnostics tests" begin
     include("diagnostics_tests.jl")
 end
-@safetestset "PostProcessor tests" begin
+gpu_broken || @safetestset "PostProcessor tests" begin
     include("postprocessor_tests.jl")
 end
 @safetestset "Checkpointer tests" begin
     include("checkpointer_tests.jl")
 end
-@safetestset "experiment test: CoupledSims tests" begin
+gpu_broken || @safetestset "experiment test: CoupledSims tests" begin
     include("experiment_tests/coupled_sims.jl")
 end
 @safetestset "experiment test: Leaderboard" begin
@@ -51,12 +54,12 @@ end
 @safetestset "component model test: prescr. sea ice" begin
     include("component_model_tests/prescr_seaice_tests.jl")
 end
-@safetestset "component model test: eisenman sea ice" begin
+gpu_broken || @safetestset "component model test: eisenman sea ice" begin
     include("component_model_tests/eisenman_seaice_tests.jl")
 end
 @safetestset "component model test: slab ocean" begin
     include("component_model_tests/slab_ocean_tests.jl")
 end
-@safetestset "debug diagnostics: amip plots" begin
+gpu_broken || @safetestset "debug diagnostics: amip plots" begin
     include("debug/debug_amip_plots.jl")
 end