From 40425c3615725a3e4fda0f5b4dfcb2d9a17246a6 Mon Sep 17 00:00:00 2001 From: Julia Sloan Date: Fri, 13 Dec 2024 10:19:02 -0800 Subject: [PATCH] initialize comms contexts --- experiments/benchmarks/bucket.jl | 1 + experiments/benchmarks/land.jl | 1 + experiments/benchmarks/richards.jl | 1 + experiments/benchmarks/snowy_land.jl | 1 + .../integrated/fluxnet/met_drivers_FLUXNET.jl | 1 + experiments/integrated/global/global_soil_canopy.jl | 4 ++-- .../performance/integrated_timestep_test.jl | 4 ++-- experiments/long_runs/bucket.jl | 1 + experiments/long_runs/land.jl | 1 + experiments/long_runs/land_region.jl | 1 + experiments/long_runs/leaderboard/data_sources.jl | 12 +++++------- experiments/long_runs/snowy_land.jl | 1 + experiments/long_runs/soil.jl | 1 + experiments/standalone/Bucket/bucket_era5.jl | 4 ++-- .../standalone/Bucket/global_bucket_function.jl | 4 ++-- .../standalone/Bucket/global_bucket_temporalmap.jl | 4 ++-- experiments/standalone/Snow/snowmip_simulation.jl | 1 + experiments/standalone/Soil/richards_comparison.jl | 4 ++-- experiments/standalone/Soil/richards_runoff.jl | 4 ++-- .../Fluxnet/fluxnet_utilities/inputs_dataframe.jl | 6 ++++-- lib/ClimaLandSimulations/src/Fluxnet/run_fluxnet.jl | 5 ++++- src/Artifacts.jl | 6 +++--- src/shared_utilities/Domains.jl | 11 ++++++++++- src/shared_utilities/checkpoints.jl | 3 +++ test/simulations/spatial_parameters.jl | 1 + test/standalone/Bucket/albedo_types.jl | 1 - 26 files changed, 55 insertions(+), 29 deletions(-) diff --git a/experiments/benchmarks/bucket.jl b/experiments/benchmarks/bucket.jl index 0b27e50846..7c3dd6f597 100644 --- 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 index ab27eef8f2..73a2674c45 100644 --- 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 index af5f4ca500..4346a7cbc9 100644 --- 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 index bcc7fc084d..c7cac4dc1b 100644 --- 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/fluxnet/met_drivers_FLUXNET.jl b/experiments/integrated/fluxnet/met_drivers_FLUXNET.jl index cd9291ddca..c7c3252c54 100644 --- a/experiments/integrated/fluxnet/met_drivers_FLUXNET.jl +++ b/experiments/integrated/fluxnet/met_drivers_FLUXNET.jl @@ -11,6 +11,7 @@ using ClimaLand import ClimaComms context = ClimaComms.context() +ClimaComms.init(context) # Methods for reading in the LAI data from MODIS data include( diff --git a/experiments/integrated/global/global_soil_canopy.jl b/experiments/integrated/global/global_soil_canopy.jl index 99851a2cbf..79999cd43e 100644 --- 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 index b73a34af37..5d9c443392 100644 --- 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 index f665d62cf2..a43485a607 100644 --- 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 index 505c97dbf0..139bd5cfca 100644 --- 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 index e38b5e9f05..129cb24755 100644 --- 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 index abf9276e01..fb6b2f292e 100644 --- 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 index 850560e08c..7e066c890e 100644 --- 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 index e778799508..fe61f8fe4e 100644 --- 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 index 17f52b078d..b4190dbd8f 100644 --- 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 index 67bbf8f1f1..244303b963 100644 --- 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 index e010a007fc..423c8d915f 100644 --- 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 index e9a1f11fb9..5ee24ea556 100644 --- 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 index 2eea990257..10028d7c54 100644 --- 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 index cd5db1f121..f6d556a103 100644 --- 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/lib/ClimaLandSimulations/src/Fluxnet/fluxnet_utilities/inputs_dataframe.jl b/lib/ClimaLandSimulations/src/Fluxnet/fluxnet_utilities/inputs_dataframe.jl index 930bfbbc67..2ce03af4d8 100644 --- a/lib/ClimaLandSimulations/src/Fluxnet/fluxnet_utilities/inputs_dataframe.jl +++ b/lib/ClimaLandSimulations/src/Fluxnet/fluxnet_utilities/inputs_dataframe.jl @@ -6,12 +6,12 @@ export make_inputs_df """ make_inputs_df(site_ID) -Returns DataFrames of inputs, one unitless, one in SI units, one in commonly used units. +Returns DataFrames of inputs, one unitless, one in SI units, one in commonly used units. """ function make_inputs_df( site_ID; FT = Float64, - context = ClimaComms.SingletonCommsContext(), + context = ClimaComms.context(), setup = make_setup(site_ID), domain = make_domain(setup, FT), config = make_config(site_ID), @@ -19,6 +19,8 @@ function make_inputs_df( drivers = make_drivers(site_ID, setup, config, params, context), timestepper = make_timestepper(setup), ) + # Initialize the communication context + ClimaComms.init(context) variables_name = ( "DateTime", diff --git a/lib/ClimaLandSimulations/src/Fluxnet/run_fluxnet.jl b/lib/ClimaLandSimulations/src/Fluxnet/run_fluxnet.jl index ca6b7242d8..dcbeb16eaa 100644 --- a/lib/ClimaLandSimulations/src/Fluxnet/run_fluxnet.jl +++ b/lib/ClimaLandSimulations/src/Fluxnet/run_fluxnet.jl @@ -8,7 +8,7 @@ Run ClimaLand at a fluxnet site. function run_fluxnet( site_ID; FT = Float64, - context = ClimaComms.SingletonCommsContext(), + context = ClimaComms.context(), setup = make_setup(site_ID), domain = make_domain(setup, FT), config = make_config(site_ID), @@ -16,6 +16,9 @@ function run_fluxnet( drivers = make_drivers(site_ID, setup, config, params, context), timestepper = make_timestepper(setup), ) + # Initialize the communication context + ClimaComms.init(context) + #earth_param_set = create_lsm_parameters(FT) # Now we set up the model. For the soil model, we pick diff --git a/src/Artifacts.jl b/src/Artifacts.jl index 6827d30516..dd7aa4e7ff 100644 --- a/src/Artifacts.jl +++ b/src/Artifacts.jl @@ -337,7 +337,7 @@ 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. @@ -345,8 +345,8 @@ 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 index 3173f9feb4..72b7da5179 100644 --- a/src/shared_utilities/Domains.jl +++ b/src/shared_utilities/Domains.jl @@ -254,6 +254,9 @@ function Plane(; comms_ctx = ClimaComms.context(), radius_earth = 6.378e6, ) where {FT} + # Initialize the communication context + ClimaComms.init(comms_ctx) + if isnothing(longlat) @assert xlim[1] < xlim[2] @assert ylim[1] < ylim[2] @@ -530,7 +533,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. @@ -550,6 +553,9 @@ function SphericalShell(; dz_tuple::Union{Tuple{FT, FT}, Nothing} = nothing, comms_ctx = ClimaComms.context(), ) where {FT} + # Initialize the communication context + ClimaComms.init(comms_ctx) + @assert 0 < radius @assert 0 < depth vertdomain = ClimaCore.Domains.IntervalDomain( @@ -646,6 +652,9 @@ function SphericalSurface(; comms_ctx = ClimaComms.context(), ) where {FT} @assert 0 < radius + # Initialize the communication context + ClimaComms.init(comms_ctx) + horzdomain = ClimaCore.Domains.SphereDomain(radius) horzmesh = Meshes.EquiangularCubedSphere(horzdomain, nelements) horztopology = Topologies.Topology2D(comms_ctx, horzmesh) diff --git a/src/shared_utilities/checkpoints.jl b/src/shared_utilities/checkpoints.jl index 17532b1a92..d64379ecaa 100644 --- a/src/shared_utilities/checkpoints.jl +++ b/src/shared_utilities/checkpoints.jl @@ -156,6 +156,9 @@ function read_checkpoint( context = isnothing(model) ? ClimaComms.context() : ClimaComms.context(model), ) + # Initialize the communication context + ClimaComms.init(context) + hdfreader = InputOutput.HDF5Reader(file_path, context) Y = InputOutput.read_field(hdfreader, "Y") attributes = InputOutput.read_attributes(hdfreader, "/") diff --git a/test/simulations/spatial_parameters.jl b/test/simulations/spatial_parameters.jl index a5610f8cec..fe3294c570 100644 --- 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 index a270d95909..26ba0f5178 100644 --- 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