Merge pull request #954 from CliMA/js/init_ctx

juliasloan25 · web-flow · commit b24f868d8053 · 2024-12-13T17:26:33.000-08:00
initialize comms contexts
diff --git a/NEWS.md b/NEWS.md
@@ -6,6 +6,9 @@ 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)
+- Initialize communications contexts, enabling experiments
+  to run with MPI.
+  PR[#954](https://github.com/CliMA/ClimaLand.jl/pull/954)
 
 v0.15.6
 -------
diff --git a/experiments/benchmarks/bucket.jl b/experiments/benchmarks/bucket.jl
@@ -37,6 +37,7 @@ import Profile, ProfileCanvas
 const FT = Float64;
 
 context = ClimaComms.context()
+ClimaComms.init(context)
 device = ClimaComms.device()
 device_suffix = device isa ClimaComms.CPUSingleThreaded ? "cpu" : "gpu"
 
diff --git a/experiments/benchmarks/land.jl b/experiments/benchmarks/land.jl
@@ -47,6 +47,7 @@ import Profile, ProfileCanvas
 const FT = Float64;
 time_interpolation_method = LinearInterpolation(PeriodicCalendar())
 context = ClimaComms.context()
+ClimaComms.init(context)
 device = ClimaComms.device()
 device_suffix = device isa ClimaComms.CPUSingleThreaded ? "cpu" : "gpu"
 
diff --git a/experiments/benchmarks/richards.jl b/experiments/benchmarks/richards.jl
@@ -51,6 +51,7 @@ const FT = Float64;
 
 regridder_type = :InterpolationsRegridder
 context = ClimaComms.context()
+ClimaComms.init(context)
 device = ClimaComms.device()
 device_suffix = device isa ClimaComms.CPUSingleThreaded ? "cpu" : "gpu"
 outdir = "richards_benchmark_$(device_suffix)"
diff --git a/experiments/benchmarks/snowy_land.jl b/experiments/benchmarks/snowy_land.jl
@@ -44,6 +44,7 @@ using Test
 const FT = Float64
 time_interpolation_method = LinearInterpolation(PeriodicCalendar())
 context = ClimaComms.context()
+ClimaComms.init(context)
 device = ClimaComms.device()
 device_suffix = device isa ClimaComms.CPUSingleThreaded ? "cpu" : "gpu"
 outdir = "snowy_land_benchmark_$(device_suffix)"
diff --git a/experiments/integrated/global/global_soil_canopy.jl b/experiments/integrated/global/global_soil_canopy.jl
@@ -28,11 +28,11 @@ import ClimaAnalysis.Visualize as viz
 import ClimaUtilities
 time_interpolation_method = LinearInterpolation(PeriodicCalendar())
 context = ClimaComms.context()
+ClimaComms.init(context)
 outdir = generate_output_path("experiments/integrated/global")
 
 device_suffix =
-    typeof(ClimaComms.context().device) <: ClimaComms.CPUSingleThreaded ?
-    "cpu" : "gpu"
+    typeof(context.device) <: ClimaComms.CPUSingleThreaded ? "cpu" : "gpu"
 
 FT = Float64
 earth_param_set = LP.LandParameters(FT)
diff --git a/experiments/integrated/performance/integrated_timestep_test.jl b/experiments/integrated/performance/integrated_timestep_test.jl
@@ -115,9 +115,9 @@ function set_initial_conditions(land, t0)
 end
 
 context = ClimaComms.context()
+ClimaComms.init(context)
 device_suffix =
-    typeof(ClimaComms.context().device) <: ClimaComms.CPUSingleThreaded ?
-    "cpu" : "gpu"
+    typeof(context.device) <: ClimaComms.CPUSingleThreaded ? "cpu" : "gpu"
 const FT = Float64
 earth_param_set = LP.LandParameters(FT)
 
diff --git a/experiments/long_runs/bucket.jl b/experiments/long_runs/bucket.jl
@@ -44,6 +44,7 @@ const FT = Float64;
 time_interpolation_method = LinearInterpolation(PeriodicCalendar())
 regridder_type = :InterpolationsRegridder
 context = ClimaComms.context()
+ClimaComms.init(context)
 device = ClimaComms.device()
 device_suffix = device isa ClimaComms.CPUSingleThreaded ? "cpu" : "gpu"
 root_path = "bucket_longrun_$(device_suffix)"
diff --git a/experiments/long_runs/land.jl b/experiments/long_runs/land.jl
@@ -48,6 +48,7 @@ using Poppler_jll: pdfunite
 const FT = Float64;
 time_interpolation_method = LinearInterpolation(PeriodicCalendar())
 context = ClimaComms.context()
+ClimaComms.init(context)
 device = ClimaComms.device()
 device_suffix = device isa ClimaComms.CPUSingleThreaded ? "cpu" : "gpu"
 root_path = "land_longrun_$(device_suffix)"
diff --git a/experiments/long_runs/land_region.jl b/experiments/long_runs/land_region.jl
@@ -50,6 +50,7 @@ using Poppler_jll: pdfunite
 const FT = Float64;
 time_interpolation_method = LinearInterpolation(PeriodicCalendar())
 context = ClimaComms.context()
+ClimaComms.init(context)
 device = ClimaComms.device()
 device_suffix = device isa ClimaComms.CPUSingleThreaded ? "cpu" : "gpu"
 root_path = "california_longrun_$(device_suffix)"
diff --git a/experiments/long_runs/leaderboard/data_sources.jl b/experiments/long_runs/leaderboard/data_sources.jl
@@ -92,8 +92,8 @@ function get_obs_var_dict()
     obs_var_dict["et"] =
         (start_date) -> begin
             obs_var = ClimaAnalysis.OutputVar(
-                ClimaLand.Artifacts.ilamb_dataset_path(;
-                    context = "evspsbl_MODIS_et_0.5x0.5.nc",
+                ClimaLand.Artifacts.ilamb_dataset_path(
+                    "evspsbl_MODIS_et_0.5x0.5.nc",
                 ),
                 "et",
                 new_start_date = start_date,
@@ -113,9 +113,7 @@ function get_obs_var_dict()
     obs_var_dict["gpp"] =
         (start_date) -> begin
             obs_var = ClimaAnalysis.OutputVar(
-                ClimaLand.Artifacts.ilamb_dataset_path(;
-                    context = "gpp_FLUXCOM_gpp.nc",
-                ),
+                ClimaLand.Artifacts.ilamb_dataset_path("gpp_FLUXCOM_gpp.nc"),
                 "gpp",
                 new_start_date = start_date,
                 shift_by = Dates.firstdayofmonth,
@@ -131,8 +129,8 @@ function get_obs_var_dict()
     obs_var_dict["lwu"] =
         (start_date) -> begin
             obs_var = ClimaAnalysis.OutputVar(
-                ClimaLand.Artifacts.ilamb_dataset_path(;
-                    context = "rlus_CERESed4.2_rlus.nc",
+                ClimaLand.Artifacts.ilamb_dataset_path(
+                    "rlus_CERESed4.2_rlus.nc",
                 ),
                 "rlus",
                 new_start_date = start_date,
diff --git a/experiments/long_runs/snowy_land.jl b/experiments/long_runs/snowy_land.jl
@@ -52,6 +52,7 @@ using Poppler_jll: pdfunite
 const FT = Float64;
 time_interpolation_method = LinearInterpolation(PeriodicCalendar())
 context = ClimaComms.context()
+ClimaComms.init(context)
 device = ClimaComms.device()
 device_suffix = device isa ClimaComms.CPUSingleThreaded ? "cpu" : "gpu"
 root_path = "snowy_land_longrun_$(device_suffix)"
diff --git a/experiments/long_runs/soil.jl b/experiments/long_runs/soil.jl
@@ -46,6 +46,7 @@ using Poppler_jll: pdfunite
 const FT = Float64;
 time_interpolation_method = LinearInterpolation(PeriodicCalendar())
 context = ClimaComms.context()
+ClimaComms.init(context)
 device = ClimaComms.device()
 device_suffix = device isa ClimaComms.CPUSingleThreaded ? "cpu" : "gpu"
 root_path = "soil_longrun_$(device_suffix)"
diff --git a/experiments/standalone/Bucket/bucket_era5.jl b/experiments/standalone/Bucket/bucket_era5.jl
@@ -78,11 +78,11 @@ time_interpolation_method = LinearInterpolation(PeriodicCalendar())
 regridder_type = :InterpolationsRegridder
 FT = Float64;
 context = ClimaComms.context()
+ClimaComms.init(context)
 earth_param_set = LP.LandParameters(FT);
 # Use separate output directory for CPU and GPU runs to avoid race condition
 device_suffix =
-    typeof(ClimaComms.context().device) <: ClimaComms.CPUSingleThreaded ?
-    "cpu" : "gpu"
+    typeof(context.device) <: ClimaComms.CPUSingleThreaded ? "cpu" : "gpu"
 outdir = "experiments/standalone/Bucket/artifacts_era5$(regional_str)_$(device_suffix)"
 output_dir = ClimaUtilities.OutputPathGenerator.generate_output_path(outdir)
 
diff --git a/experiments/standalone/Bucket/global_bucket_function.jl b/experiments/standalone/Bucket/global_bucket_function.jl
@@ -63,11 +63,11 @@ end
 anim_plots = false
 FT = Float64;
 context = ClimaComms.context()
+ClimaComms.init(context)
 earth_param_set = LP.LandParameters(FT);
 # Use separate output directory for CPU and GPU runs to avoid race condition
 device_suffix =
-    typeof(ClimaComms.context().device) <: ClimaComms.CPUSingleThreaded ?
-    "cpu" : "gpu"
+    typeof(context.device) <: ClimaComms.CPUSingleThreaded ? "cpu" : "gpu"
 outdir = "experiments/standalone/Bucket/artifacts_function_$(device_suffix)"
 
 # Construct simulation domain
diff --git a/experiments/standalone/Bucket/global_bucket_temporalmap.jl b/experiments/standalone/Bucket/global_bucket_temporalmap.jl
@@ -71,11 +71,11 @@ end
 
 FT = Float64;
 context = ClimaComms.context()
+ClimaComms.init(context)
 earth_param_set = LP.LandParameters(FT);
 # Use separate output directory for CPU and GPU runs to avoid race condition
 device_suffix =
-    typeof(ClimaComms.context().device) <: ClimaComms.CPUSingleThreaded ?
-    "cpu" : "gpu"
+    typeof(context.device) <: ClimaComms.CPUSingleThreaded ? "cpu" : "gpu"
 outdir = "experiments/standalone/Bucket/artifacts_temporalmap_$(device_suffix)"
 t0 = 0.0;
 # run for 50 days to test monthly file update
diff --git a/experiments/standalone/Snow/snowmip_simulation.jl b/experiments/standalone/Snow/snowmip_simulation.jl
@@ -33,6 +33,7 @@ climaland_dir = pkgdir(ClimaLand)
 FT = Float32
 param_set = LP.LandParameters(FT)
 context = ClimaComms.context()
+ClimaComms.init(context)
 
 # This reads in the data and sets up the drivers, as well as computes the IC from the data
 include(
diff --git a/experiments/standalone/Soil/richards_comparison.jl b/experiments/standalone/Soil/richards_comparison.jl
@@ -21,9 +21,9 @@ clay_datapath = joinpath(bonan_data_folder, "bonan_data_clay.txt")
 sand_datapath = joinpath(bonan_data_folder, "bonan_data_sand.txt")
 
 context = ClimaComms.context()
+ClimaComms.init(context)
 device_suffix =
-    typeof(ClimaComms.context().device) <: ClimaComms.CPUSingleThreaded ?
-    "cpu" : "gpu"
+    typeof(context.device) <: ClimaComms.CPUSingleThreaded ? "cpu" : "gpu"
 outdir = generate_output_path(
     joinpath("experiments", "standalone", "Soil", device_suffix),
 )
diff --git a/experiments/standalone/Soil/richards_runoff.jl b/experiments/standalone/Soil/richards_runoff.jl
@@ -23,9 +23,9 @@ import ClimaLand.Parameters as LP
 
 regridder_type = :InterpolationsRegridder
 context = ClimaComms.context()
+ClimaComms.init(context)
 device_suffix =
-    typeof(ClimaComms.context().device) <: ClimaComms.CPUSingleThreaded ?
-    "cpu" : "gpu"
+    typeof(context.device) <: ClimaComms.CPUSingleThreaded ? "cpu" : "gpu"
 outdir = generate_output_path(
     joinpath("experiments", "standalone", "Soil", "artifacts", device_suffix),
 )
diff --git a/src/Artifacts.jl b/src/Artifacts.jl
@@ -337,16 +337,16 @@ function bareground_albedo_dataset_path(; context = nothing)
 end
 
 """
-    ilamb_dataset_path(; context = nothing)
+    ilamb_dataset_path(filename; context = nothing)
 
 Triggers the download of the ILAMB dataset, if not already downloaded, using
 Julia Artifacts, and returns the path to this file.
 
 There are only three datasets available which are "rlus_CERESed4.2_rlus.nc",
 "gpp_FLUXCOM_gpp.nc", and "evspsbl_MODIS_et_0.5x0.5.nc".
 """
-function ilamb_dataset_path(; context = nothing)
-    return joinpath(@clima_artifact("ilamb_data"), context)
+function ilamb_dataset_path(filename; context = nothing)
+    return joinpath(@clima_artifact("ilamb_data", context), filename)
 end
 
 end
diff --git a/src/shared_utilities/Domains.jl b/src/shared_utilities/Domains.jl
@@ -530,7 +530,7 @@ end
         nelements::Tuple{Int, Int},
         npolynomial::Int,
         dz_tuple::Union{Tuple{FT, FT}, Nothing} = nothing,
-        comms_ctx = ClimaComms.SingletonCommsContext(),
+        comms_ctx = ClimaComms.context(),
     ) where {FT}
 Outer constructor for the `SphericalShell` domain, using keyword arguments.
 
diff --git a/test/simulations/spatial_parameters.jl b/test/simulations/spatial_parameters.jl
@@ -18,6 +18,7 @@ regridder_type = :InterpolationsRegridder
 extrapolation_bc =
     (Interpolations.Periodic(), Interpolations.Flat(), Interpolations.Flat())
 context = ClimaComms.context()
+ClimaComms.init(context)
 
 earth_param_set = LP.LandParameters(FT)
 radius = FT(6378.1e3)
diff --git a/test/standalone/Bucket/albedo_types.jl b/test/standalone/Bucket/albedo_types.jl
@@ -255,7 +255,6 @@ end
             set_initial_cache!(p, Y, FT(0.0))
 
             # Read in data manually
-            comms_ctx = ClimaComms.context(surface_space)
             α_bareground = SpaceVaryingInput(path, varname, surface_space)
             @test p.bucket.α_sfc == α_bareground
         else
