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"
@@ -69,15 +69,6 @@ steps:
         command: "julia --color=yes --project=.buildkite experiments/standalone/Soil/richards_comparison.jl"
         artifact_paths: "experiments/standalone/Soil/cpu/output_active/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/output_active/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/output_active/*png"
@@ -144,6 +135,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/output_active/comparison*png"
+        agents:
+          slurm_ntasks: 1
+          slurm_gpus: 1
+        env:
+          CLIMACOMMS_DEVICE: "CUDA"
+
   - group: "ClimaLandSimulations"
     steps:
       - label: "Ozark figures Makie"
@@ -188,6 +188,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"

NEWS.md

@@ -3,7 +3,9 @@ ClimaLand.jl Release Notes
 
 main
 --------
-
+- Run unit tests on GPU, and update code for GPU compatibility
+  (including workaround for ClimaCore type inference failure)
+  PR[#739](https://github.com/CliMA/ClimaLand.jl/pull/739)
 
 v0.15.6
 -------

experiments/integrated/performance/integrated_timestep_test.jl

@@ -104,8 +104,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)
@@ -176,14 +177,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) = tuple(ps)
 
 # wrapper methods:
 P_ref(ps::ALP) = ps.MSLP

src/standalone/Soil/energy_hydrology.jl

@@ -623,6 +623,11 @@ function ClimaLand.source!(
     _ρ_l = FT(LP.ρ_cloud_liq(earth_param_set))
     _ρ_i = FT(LP.ρ_cloud_ice(earth_param_set))
     Δz = model.domain.fields.Δz # center face distance
+
+    # Wrap hydrology and earth parameters in one struct to avoid type inference failure
+    hydrology_earth_params =
+        ClimaLand.Soil.HydrologyEarthParameters.(hydrology_cm, earth_param_set)
+
     @. dY.soil.ϑ_l +=
         -phase_change_source(
             p.soil.θ_l,
@@ -640,8 +645,7 @@ function ClimaLand.source!(
             ),
             ν,
             θ_r,
-            hydrology_cm,
-            earth_param_set,
+            hydrology_earth_params,
         )
     @. dY.soil.θ_i +=
         (_ρ_l / _ρ_i) * phase_change_source(
@@ -660,12 +664,10 @@ function ClimaLand.source!(
             ),
             ν,
             θ_r,
-            hydrology_cm,
-            earth_param_set,
+            hydrology_earth_params,
         )
 end
 
-
 """
     SoilSublimation{FT} <: AbstractSoilSource{FT}
 

src/standalone/Soil/soil_heat_parameterizations.jl

@@ -31,12 +31,11 @@ end
         τ::FT,
         ν::FT,
         θ_r::FT,
-        hydrology_cm::C,
-        earth_param_set::EP,
-    ) where {FT, EP, C}
+        hydrology_earth_params::HEP
+    ) where {FT, HEP}
 Returns the source term (1/s) used for converting liquid water
 and ice into each other during phase changes. Note that
-there are unitless prefactors multiplying this term in the 
+there are unitless prefactors multiplying this term in the
 equations.
 
 Note that these equations match what is in Dall'Amico (for θstar,
@@ -50,9 +49,12 @@ function phase_change_source(
     τ::FT,
     ν::FT,
     θ_r::FT,
-    hydrology_cm::C,
-    earth_param_set::EP,
-) where {FT, EP, C}
+    hydrology_earth_params::HEP,
+) where {FT, HEP}
+    # Extract parameter sets from their container
+    hydrology_cm = hydrology_earth_params.hydrology_cm
+    earth_param_set = hydrology_earth_params.earth_param_set
+
     _ρ_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))
@@ -71,6 +73,28 @@ function phase_change_source(
     return (θ_l - θstar) / τ
 end
 
+"""
+    struct HydrologyEarthParameters{
+        HCM <: AbstractSoilHydrologyClosure,
+        EP <: ClimaLand.Parameters.AbstractLandParameters,
+    }
+
+A wrapper type around the hydrology closure model and land parameter
+structs. This is needed because of a type inference failure coming from
+ClimaCore when multiple structs and fields are broadcasted over.
+This struct circumvents that issue by wrapping the structs in a single
+type, that can be unpacked within the broadcasted function.
+
+See github.com/CliMA/ClimaCore.jl/issues/2065 for more information
+"""
+struct HydrologyEarthParameters{
+    HCM <: AbstractSoilHydrologyClosure,
+    EP <: ClimaLand.Parameters.AbstractLandParameters,
+}
+    hydrology_cm::HCM
+    earth_param_set::EP
+end
+Base.broadcastable(x::HydrologyEarthParameters) = tuple(x)
 
 """
     volumetric_heat_capacity(
@@ -205,7 +229,7 @@ end
             ν::FT
     ) where {FT}
 
-Compute the expression for relative saturation. 
+Compute the expression for relative saturation.
 This is referred to as θ_sat in Balland and Arp's paper.
 """
 function relative_saturation(θ_l::FT, θ_i::FT, ν::FT) where {FT}

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/standalone/Soil/soil_parameterizations.jl

@@ -307,6 +307,10 @@ for FT in (Float32, Float64)
         vg_n = FT(1.4)
         hcm = vanGenuchten{FT}(; α = vg_α, n = vg_n)
 
+        # Wrap hydrology and earth parameters in one struct to avoid type inference failure
+        hydrology_earth_params =
+            ClimaLand.Soil.HydrologyEarthParameters(hcm, param_set)
+
         K_sat = FT(1e-5)
         ν_ss_om = FT(0.1)
         ν_ss_gravel = FT(0.1)
@@ -336,11 +340,26 @@ for FT in (Float32, Float64)
         ρc_s = volumetric_heat_capacity.(θ_l, θ_i, parameters.ρc_ds, param_set)
         τ = thermal_time.(ρc_s, Δz, parameters.κ_dry)
         @test (
-            phase_change_source.(θ_l, θ_i, T, τ, ν, θ_r, hcm, param_set) ≈
-            (θ_l .- θ_star) ./ τ
+            phase_change_source.(
+                θ_l,
+                θ_i,
+                T,
+                τ,
+                ν,
+                θ_r,
+                hydrology_earth_params,
+            ) ≈ (θ_l .- θ_star) ./ τ
         )
         @test sum(
-            phase_change_source.(θ_l, θ_i, T, τ, ν, θ_r, hcm, param_set) .> 0.0,
+            phase_change_source.(
+                θ_l,
+                θ_i,
+                T,
+                τ,
+                ν,
+                θ_r,
+                hydrology_earth_params,
+            ) .> 0.0,
         ) == 3
         # try θ_l = 0.1
 
@@ -354,8 +373,15 @@ for FT in (Float32, Float64)
         ρc_s = volumetric_heat_capacity.(θ_l, θ_i, parameters.ρc_ds, param_set)
         τ = thermal_time.(ρc_s, Δz, parameters.κ_dry)
         @test (
-            phase_change_source.(θ_l, θ_i, T, τ, ν, θ_r, hcm, param_set) ≈
-            zeros(FT, 3)
+            phase_change_source.(
+                θ_l,
+                θ_i,
+                T,
+                τ,
+                ν,
+                θ_r,
+                hydrology_earth_params,
+            ) ≈ zeros(FT, 3)
         )
         @test (θ_star ≈ θ_l)
 
@@ -370,11 +396,26 @@ for FT in (Float32, Float64)
         ρc_s = volumetric_heat_capacity.(θ_l, θ_i, parameters.ρc_ds, param_set)
         τ = thermal_time.(ρc_s, Δz, parameters.κ_dry)
         @test (
-            phase_change_source.(θ_l, θ_i, T, τ, ν, θ_r, hcm, param_set) ≈
-            (θ_l .- θ_star) ./ τ
+            phase_change_source.(
+                θ_l,
+                θ_i,
+                T,
+                τ,
+                ν,
+                θ_r,
+                hydrology_earth_params,
+            ) ≈ (θ_l .- θ_star) ./ τ
         )
         @test sum(
-            phase_change_source.(θ_l, θ_i, T, τ, ν, θ_r, hcm, param_set) .< 0.0,
+            phase_change_source.(
+                θ_l,
+                θ_i,
+                T,
+                τ,
+                ν,
+                θ_r,
+                hydrology_earth_params,
+            ) .< 0.0,
         ) == 2
 
 
@@ -388,11 +429,26 @@ for FT in (Float32, Float64)
         ρc_s = volumetric_heat_capacity.(θ_l, θ_i, parameters.ρc_ds, param_set)
         τ = thermal_time.(ρc_s, Δz, parameters.κ_dry)
         @test (
-            phase_change_source.(θ_l, θ_i, T, τ, ν, θ_r, hcm, param_set) ≈
-            (θ_l .- θ_star) ./ τ
+            phase_change_source.(
+                θ_l,
+                θ_i,
+                T,
+                τ,
+                ν,
+                θ_r,
+                hydrology_earth_params,
+            ) ≈ (θ_l .- θ_star) ./ τ
         )
         @test sum(
-            phase_change_source.(θ_l, θ_i, T, τ, ν, θ_r, hcm, param_set) .> 0.0,
+            phase_change_source.(
+                θ_l,
+                θ_i,
+                T,
+                τ,
+                ν,
+                θ_r,
+                hydrology_earth_params,
+            ) .> 0.0,
         ) == 2
     end
 end