fix buildkite run so that unit tests run on GPU

.buildkite/pipeline.yml

@@ -1,6 +1,6 @@
 agents:
   queue: new-central
-  slurm_mem: 8G
+  slurm_mem: 12G
   modules: climacommon/2024_10_09
 
 env:
@@ -26,7 +26,7 @@ steps:
       - "julia --project=.buildkite -e 'using Pkg; Pkg.status()'"
 
       - echo "--- Instantiate test"
-      - "julia --project=test -e 'using Pkg; Pkg.develop(;path=\".\"); Pkg.instantiate(;verbose=true)'"
+      - "julia --project=test -e 'using Pkg; Pkg.develop(;path=\".\"); Pkg.add(\"MPI\"); Pkg.add(\"CUDA\"); Pkg.instantiate(;verbose=true)'"
       - "julia --project=test -e 'using Pkg; Pkg.status()'"
 
       - echo "--- Instantiate lib/ClimaLandSimulations"
@@ -67,15 +67,6 @@ steps:
         command: "julia --color=yes --project=.buildkite experiments/standalone/Soil/richards_comparison.jl"
         artifact_paths: "experiments/standalone/Soil/cpu/comparison*png"
 
-      - label: "Richards comparison to Bonan: GPU"
-        command: "julia --color=yes --project=.buildkite experiments/standalone/Soil/richards_comparison.jl"
-        artifact_paths: "experiments/standalone/Soil/gpu/comparison*png"
-        agents:
-          slurm_ntasks: 1
-          slurm_gpus: 1
-        env:
-          CLIMACOMMS_DEVICE: "CUDA"
-
       - label: "vaira_test"
         command: "julia --color=yes --project=.buildkite experiments/integrated/fluxnet/run_fluxnet.jl US-Var"
         artifact_paths: "experiments/integrated/fluxnet/US-Var/out/*png"
@@ -142,6 +133,15 @@ steps:
         env:
           CLIMACOMMS_DEVICE: "CUDA"
 
+      - label: "Richards comparison to Bonan: GPU"
+        command: "julia --color=yes --project=.buildkite experiments/standalone/Soil/richards_comparison.jl"
+        artifact_paths: "experiments/standalone/Soil/gpu/comparison*png"
+        agents:
+          slurm_ntasks: 1
+          slurm_gpus: 1
+        env:
+          CLIMACOMMS_DEVICE: "CUDA"
+
   - group: "ClimaLandSimulations"
     steps:
       - label: "Ozark figures Makie"
@@ -181,6 +181,8 @@ steps:
         agents:
           slurm_ntasks: 1
           slurm_gpus: 1
+        env:
+          CLIMACOMMS_DEVICE: "CUDA"
 
       - label: "Global Bucket on GPU (functional albedo)"
         key: "global_bucket_function_gpu"

experiments/integrated/performance/integrated_timestep_test.jl

@@ -103,8 +103,9 @@ function set_initial_conditions(land, t0)
         )
 
     for i in 1:2
+        S_l = S_l_ini[i]
         Y.canopy.hydraulics.ϑ_l.:($i) .=
-            augmented_liquid_fraction.(canopy_params.ν, S_l_ini[i])
+            augmented_liquid_fraction.(canopy_params.ν, S_l)
     end
 
     Y.canopy.energy.T = FT(297.5)
@@ -175,14 +176,14 @@ function zenith_angle(
     t,
     start_date;
     cd_field = sfc_cds,
-    insol_params::Insolation.Parameters.InsolationParameters{FT} = earth_param_set.insol_params,
-) where {FT}
+    insol_params::Insolation.Parameters.InsolationParameters{_FT} = earth_param_set.insol_params,
+) where {_FT}
     # This should be time in UTC
     current_datetime = start_date + Dates.Second(round(t))
 
     # Orbital Data uses Float64, so we need to convert to our sim FT
     d, δ, η_UTC =
-        FT.(
+        _FT.(
             Insolation.helper_instantaneous_zenith_angle(
                 current_datetime,
                 start_date,

ext/CreateParametersExt.jl

@@ -66,7 +66,6 @@ function LP.LandParameters(toml_dict::CP.AbstractTOMLDict)
         insol_params,
     )
 end
-Base.broadcastable(ps::LP.LandParameters) = tuple(ps)
 
 """
     AutotrophicRespirationParameters(FT; kwargs...)

src/shared_utilities/Parameters.jl

@@ -28,6 +28,7 @@ Base.@kwdef struct LandParameters{FT, TP, SFP, IP} <: ALP
 end
 
 Base.eltype(::LandParameters{FT}) where {FT} = FT
+Base.broadcastable(ps::LandParameters) = Ref(ps)
 
 # wrapper methods:
 P_ref(ps::ALP) = ps.MSLP

src/standalone/Soil/retention_models.jl

@@ -1,10 +1,10 @@
 export AbstractSoilHydrologyClosure, vanGenuchten, BrooksCorey
 
 """
-    AbstractSoilHydrologyClosure{FT <: AbstractFloat} 
+    AbstractSoilHydrologyClosure{FT <: AbstractFloat}
 
 The abstract type of soil hydrology closure, of which
-vanGenuchten{FT} and BrooksCorey{FT} are the two supported 
+vanGenuchten{FT} and BrooksCorey{FT} are the two supported
 concrete types.
 
 To add a new parameterization, methods are required for:
@@ -20,7 +20,7 @@ Base.broadcastable(x::AbstractSoilHydrologyClosure) = tuple(x)
 """
     vanGenuchten{FT} <: AbstractSoilHydrologyClosure{FT}
 
-The van Genuchten soil hydrology closure, chosen when the 
+The van Genuchten soil hydrology closure, chosen when the
 hydraulic conductivity and matric potential are modeled
 using the van Genuchten parameterization (van Genuchten 1980;
 see also Table 8.2 of G. Bonan 2019).
@@ -41,13 +41,13 @@ struct vanGenuchten{FT} <: AbstractSoilHydrologyClosure{FT}
         S_c = (1 + ((n - 1) / n)^(1 - 2 * n))^(-m)
         return new{FT}(α, n, m, S_c)
     end
-
 end
+Base.broadcastable(vg::vanGenuchten) = tuple(vg)
 
 """
    BrooksCorey{FT} <: AbstractSoilHydrologyClosure{FT}
 
-The Brooks and Corey soil hydrology closure, chosen when the 
+The Brooks and Corey soil hydrology closure, chosen when the
 hydraulic conductivity and matric potential are modeled
 using the Brooks and Corey parameterization (Brooks and Corey,
 1964, 1966; see also Table 8.2 of G. Bonan 2019).
@@ -66,3 +66,4 @@ struct BrooksCorey{FT} <: AbstractSoilHydrologyClosure{FT}
         return new{FT}(c, ψb, S_c)
     end
 end
+Base.broadcastable(bc::BrooksCorey) = tuple(bc)

src/standalone/Vegetation/radiation.jl

@@ -23,7 +23,7 @@ struct ConstantGFunction{F <: Union{AbstractFloat, ClimaCore.Fields.Field}} <:
 end
 
 # Make the ConstantGFunction broadcastable
-Base.broadcastable(G::ConstantGFunction) = Ref(G)
+Base.broadcastable(G::ConstantGFunction) = tuple(G)
 
 """
     CLMGFunction
@@ -38,7 +38,7 @@ struct CLMGFunction{F <: Union{AbstractFloat, ClimaCore.Fields.Field}} <:
 end
 
 # Make the CLMGFunction broadcastable
-Base.broadcastable(G::CLMGFunction) = Ref(G)
+Base.broadcastable(G::CLMGFunction) = tuple(G)
 
 """
     BeerLambertParameters{FT <: AbstractFloat}

test/ref-tuple_repro.jl

@@ -0,0 +1,64 @@
+# Initially from soiltest.jl phase change source term test
+# To reproduce, run this script on GPU
+import ClimaComms
+ClimaComms.@import_required_backends
+using Statistics
+using ClimaCore
+import ClimaParams as CP
+using ClimaLand
+using ClimaLand.Domains: Column
+using ClimaLand.Soil
+import ClimaLand.Parameters as LP
+
+FT = Float64
+earth_param_set = LP.LandParameters(FT)
+ν = FT(0.495)
+θ_r = FT(0.1)
+hydrology_cm = vanGenuchten{FT}(; α = FT(2.6), n = FT(2.0))
+
+soil_domain = Column(; zlim = (FT(-1), FT(0)), nelements = 200)
+space_3d = soil_domain.space.subsurface # CenterFiniteDifferenceSpace
+
+θ_l = ClimaCore.Fields.ones(space_3d)
+θ_i = ClimaCore.Fields.ones(space_3d)
+T = ClimaCore.Fields.ones(space_3d)
+κ = ClimaCore.Fields.ones(space_3d)
+tau = ClimaCore.Fields.ones(space_3d)
+
+# This fails with dynamic function invocation when `LandParameters`
+# and `vanGenuchten` both use `tuple` for broadcasting, and it
+# passes when `Ref` is used for either `LandParameters` or `vanGenuchten` broadcasting
+@. -phase_change_source(θ_l, θ_i, T, tau, ν, θ_r, hydrology_cm, earth_param_set)
+
+
+# # Broadcasting f2 fails with the same error as the original matric_potential call
+# function f2(hydrology_cm, earth_param_set, ν, θ_r, θ_l)
+#     _ρ_i = FT(LP.ρ_cloud_ice(earth_param_set))
+#     _ρ_l = FT(LP.ρ_cloud_liq(earth_param_set))
+#     _LH_f0 = FT(LP.LH_f0(earth_param_set))
+#     _T_freeze = FT(LP.T_freeze(earth_param_set))
+#     _grav = FT(LP.grav(earth_param_set))
+#     # According to Dall'Amico (text above equation 1), ψw0 corresponds
+#     # to the matric potential corresponding to the total water content (liquid and ice).
+#     θtot = min(_ρ_i / _ρ_l * θ_i + θ_l, ν)
+
+#     matric_potential(hydrology_cm, effective_saturation(ν, θtot, θ_r))
+# end
+
+# @. f2(hydrology_cm, earth_param_set, ν, θ_r, θ_l)
+
+
+# # Defining θtot either of these ways, then calling f doesn't fail
+# function f(hydrology_cm, ν, θtot, θ_r)
+#     @. matric_potential(hydrology_cm, effective_saturation(ν, θtot, θ_r))
+# end
+
+# _ρ_i = FT(LP.ρ_cloud_ice(earth_param_set))
+# _ρ_l = FT(LP.ρ_cloud_liq(earth_param_set))
+# _LH_f0 = FT(LP.LH_f0(earth_param_set))
+# _T_freeze = FT(LP.T_freeze(earth_param_set))
+# _grav = FT(LP.grav(earth_param_set))
+# θtot = @. min(_ρ_i / _ρ_l * θ_i + θ_l, ν) # f returns [-0.0, -0.0, ..., -0.0]
+
+# # θtot = ClimaCore.Fields.ones(space_3d) # f returns [NaN, NaN, ..., NaN]
+# f(hydrology_cm, ν, θtot, θ_r)

test/standalone/Vegetation/canopy_model.jl

@@ -521,7 +521,7 @@ end
         set_canopy_prescribed_field!(Default{FT}(), p, t0)
         # Test that they are unchanged
         @test all(
-            parent(p.canopy.hydraulics.area_index.leaf) .==
+            Array(parent(p.canopy.hydraulics.area_index.leaf)) .==
             ClimaLand.Canopy.PlantHydraulics.clip(
                 FT(LAI * sin(200 * 2π / 365)),
                 FT(0.05),
@@ -1066,8 +1066,7 @@ end
         imp_tendency!(dY, Y, p, t0)
         jac = ImplicitEquationJacobian(Y)
         jacobian!(jac, Y, p, FT(1), t0)
-        jac_value =
-            parent(jac.matrix[@name(canopy.energy.T), @name(canopy.energy.T)])
+        jac_value = jac.matrix[@name(canopy.energy.T), @name(canopy.energy.T)]
         ΔT = FT(0.01)
 
         Y_2 = deepcopy(Y)
@@ -1099,17 +1098,19 @@ end
                 p.canopy.radiative_transfer.LW_n
             ) ./ ΔT
         estimated_LW = p.canopy.energy.∂LW_n∂Tc
-        @test parent(abs.(finitediff_LW .- estimated_LW) ./ finitediff_LW)[1] <
-              0.01
+        @test Array(
+            parent(abs.(finitediff_LW .- estimated_LW) ./ finitediff_LW),
+        )[1] < 0.01
 
         finitediff_SHF =
             (
                 p_2.canopy.energy.turbulent_fluxes.shf .-
                 p.canopy.energy.turbulent_fluxes.shf
             ) ./ ΔT
         estimated_SHF = p.canopy.energy.turbulent_fluxes.∂SHF∂Tc
-        @test parent(abs.(finitediff_SHF .- estimated_SHF) ./ finitediff_SHF)[1] <
-              0.15
+        @test Array(
+            parent(abs.(finitediff_SHF .- estimated_SHF) ./ finitediff_SHF),
+        )[1] < 0.15
 
         finitediff_LHF =
             (
@@ -1118,27 +1119,32 @@ end
             ) ./ ΔT
         estimated_LHF =
             p.canopy.energy.turbulent_fluxes.∂LHF∂qc .* p.canopy.energy.∂qc∂Tc
-        @test parent(abs.(finitediff_LHF .- estimated_LHF) ./ finitediff_LHF)[1] <
-              0.3
+        @test Array(
+            parent(abs.(finitediff_LHF .- estimated_LHF) ./ finitediff_LHF),
+        )[1] < 0.3
 
         # Error in `q` derivative is large
         finitediff_q = (q_sfc2 .- q_sfc) ./ ΔT
         estimated_q = p.canopy.energy.∂qc∂Tc
-        @test parent(abs.(finitediff_q .- estimated_q) ./ finitediff_q)[1] <
-              0.25
+        @test Array(
+            parent(abs.(finitediff_q .- estimated_q) ./ finitediff_q),
+        )[1] < 0.25
 
         # Im not sure why this is not smaller! There must be an error in ∂LHF∂qc also.
         estimated_LHF_with_correct_q =
             p.canopy.energy.turbulent_fluxes.∂LHF∂qc .* finitediff_q
-        @test parent(
-            abs.(finitediff_LHF .- estimated_LHF_with_correct_q) ./
-            finitediff_LHF,
+        @test Array(
+            parent(
+                abs.(finitediff_LHF .- estimated_LHF_with_correct_q) ./
+                finitediff_LHF,
+            ),
         )[1] < 0.5
 
         # Recall jac = ∂Ṫ∂T - 1 [dtγ = 1]
-        ∂Ṫ∂T = jac_value .+ 1
+        ∂Ṫ∂T = Array(parent(jac_value)) .+ 1
         @test (abs.(
-            parent((dY_2.canopy.energy.T .- dY.canopy.energy.T) ./ ΔT) - ∂Ṫ∂T
+            Array(parent((dY_2.canopy.energy.T .- dY.canopy.energy.T) ./ ΔT)) -
+            ∂Ṫ∂T
         ) / ∂Ṫ∂T)[1] < 0.25 # Error propagates here from ∂LHF∂T
     end
 end
@@ -1402,24 +1408,42 @@ end
                 t0 = FT(0.0)
                 set_initial_cache!(p, Y, t0)
 
-                @test all(parent(p.canopy.hydraulics.fa.:1) .== FT(0))
+                @test all(Array(parent(p.canopy.hydraulics.fa.:1)) .== FT(0))
+                @test all(
+                    Array(parent(p.canopy.energy.turbulent_fluxes.lhf)) .==
+                    FT(0),
+                )
+                @test all(
+                    Array(parent(p.canopy.energy.turbulent_fluxes.shf)) .==
+                    FT(0),
+                )
+                @test all(
+                    Array(parent(p.canopy.energy.fa_energy_roots)) .== FT(0),
+                )
+                @test all(Array(parent(p.canopy.hydraulics.fa_roots)) .== FT(0))
+                @test all(
+                    Array(
+                        parent(p.canopy.energy.turbulent_fluxes.transpiration),
+                    ) .== FT(0),
+                )
+                @test all(
+                    Array(parent(p.canopy.radiative_transfer.LW_n)) .== FT(0),
+                )
                 @test all(
-                    parent(p.canopy.energy.turbulent_fluxes.lhf) .== FT(0),
+                    Array(parent(p.canopy.radiative_transfer.SW_n)) .== FT(0),
                 )
                 @test all(
-                    parent(p.canopy.energy.turbulent_fluxes.shf) .== FT(0),
+                    Array(parent(p.canopy.radiative_transfer.par.abs)) .==
+                    FT(0),
+                )
+                @test all(
+                    Array(parent(p.canopy.radiative_transfer.nir.abs)) .==
+                    FT(0),
                 )
-                @test all(parent(p.canopy.energy.fa_energy_roots) .== FT(0))
-                @test all(parent(p.canopy.hydraulics.fa_roots) .== FT(0))
                 @test all(
-                    parent(p.canopy.energy.turbulent_fluxes.transpiration) .==
+                    Array(parent(p.canopy.autotrophic_respiration.Ra)) .==
                     FT(0),
                 )
-                @test all(parent(p.canopy.radiative_transfer.LW_n) .== FT(0))
-                @test all(parent(p.canopy.radiative_transfer.SW_n) .== FT(0))
-                @test all(parent(p.canopy.radiative_transfer.par.abs) .== FT(0))
-                @test all(parent(p.canopy.radiative_transfer.nir.abs) .== FT(0))
-                @test all(parent(p.canopy.autotrophic_respiration.Ra) .== FT(0))
             end
         end
     end