comms_context() added to utilities

.buildkite/pipeline.yml

@@ -108,6 +108,23 @@ steps:
           slurm_ntasks: 2
           slurm_mem: 16GB
 
+      - label: "MPI Utilities unit tests"
+        key: "utilities_mpi_tests"
+        command: "srun julia --color=yes --project=test/ test/utilities_tests.jl --run_name utilities_mpi --job_id utilities_mpi"
+        timeout_in_minutes: 20
+        env:
+          CLIMACORE_DISTRIBUTED: "MPI"
+        agents:
+          slurm_ntasks: 2
+          slurm_mem: 16GB
+
+      - label: "GPU Utilities unit tests"
+        key: "utilities_gpu_tests"
+        command: "srun julia --color=yes --project=test/ test/utilities_tests.jl --run_name utilities_gpu --job_id utilities_gpu"
+        agents:
+          slurm_mem: 5GB
+          slurm_gpus: 1
+
       - label: "Perf flame graph diff tests"
         command: "julia --color=yes --project=perf/ perf/flame_test.jl --run_name flame_test --job_id flame_perf_target"
         timeout_in_minutes: 5

experiments/AMIP/cli_options.jl

@@ -62,6 +62,10 @@ function argparse_settings()
         help = "Boolean flag indicating whether to use a dynamic slab ocean model or constant surface temperatures"
         arg_type = Bool
         default = true
+        "--device"
+        help = "Device type to use [`auto` (default) `CPUSingleThreaded`, `CPUMultiThreaded`, `CUDADevice`]"
+        arg_type = String
+        default = "auto"
         # ClimaAtmos specific
         "--surface_setup"
         help = "Triggers ClimaAtmos into the coupled mode [`PrescribedSurface` (default)]" # retained here for standalone Atmos benchmarks

experiments/AMIP/coupler_driver.jl

@@ -5,11 +5,6 @@
 # (See also Base.type_limited_string_from_context())
 redirect_stderr(IOContext(stderr, :stacktrace_types_limited => Ref(false)))
 
-# Set up context for CPU single thread/CPU with MPI/GPU
-using ClimaComms
-comms_ctx = ClimaComms.context()
-const pid, nprocs = ClimaComms.init(comms_ctx)
-
 # # AMIP Driver
 
 #=
@@ -80,7 +75,7 @@ import ClimaCoupler.Regridder:
     update_surface_fractions!, combine_surfaces!, combine_surfaces_from_sol!, dummmy_remap!, binary_mask
 import ClimaCoupler.ConservationChecker:
     EnergyConservationCheck, WaterConservationCheck, check_conservation!, plot_global_conservation
-import ClimaCoupler.Utilities: swap_space!
+import ClimaCoupler.Utilities: swap_space!, get_comms_context
 import ClimaCoupler.BCReader:
     bcfile_info_init, float_type_bcf, update_midmonth_data!, next_date_in_file, interpolate_midmonth_to_daily
 import ClimaCoupler.TimeManager:
@@ -159,6 +154,11 @@ config_dict_atmos = get_atmos_config(config_dict)
 # (if there are common keys, the last dictorionary in the `merge` arguments takes precedence)
 config_dict = merge(config_dict_atmos, config_dict)
 
+# Set up context for CPU single thread/CPU with MPI/GPU
+using ClimaComms
+comms_ctx = get_comms_context(parsed_args)
+const pid, nprocs = ClimaComms.init(comms_ctx)
+
 ## read in some parsed command line arguments
 mode_name = config_dict["mode_name"]
 run_name = config_dict["run_name"]

src/Utilities.jl

@@ -6,6 +6,7 @@ modules in the coupler.
 """
 module Utilities
 
+import ClimaComms
 using ClimaCore: Fields, Spaces
 
 export swap_space!
@@ -23,4 +24,49 @@ function swap_space!(field_out, field_in::Fields.Field)
     field_out = Fields.Field(Fields.field_values(field_in), axes(field_out))
 end
 
+"""
+    get_device(parsed_args) 
+
+Returns the device on which the model is being run 
+
+# Arguments 
+- `parsed_args`: dictionary containing a "device" flag which decides which device to run on 
+"""
+function get_device(parsed_args)
+    if parsed_args["device"] == "auto"
+        return ClimaComms.device()
+    elseif parsed_args["device"] == "CUDADevice"
+        return ClimaComms.CUDADevice()
+    elseif parsed_args["device"] == "CPUMultiThreaded" || Threads.nthreads() > 1
+        return ClimaComms.CPUMultiThreaded()
+    else
+        return ClimaComms.CPUSingleThreaded()
+    end
+end
+
+
+"""
+    get_comms_context(parsed_args) 
+
+Sets up the appropriate ClimaComms context for the device the model is to be run on 
+
+# Arguments 
+`parsed_args`: dictionary containing a "device" flag whcih decides which device context is needed
+"""
+function get_comms_context(parsed_args)
+    device = get_device(parsed_args)
+    comms_ctx = ClimaComms.context(device)
+    ClimaComms.init(comms_ctx)
+
+    @info "Running on $(nameof(typeof(device)))."
+
+    if comms_ctx isa ClimaComms.SingletonCommsContext
+        @info "Setting up single-process ClimaCoupler run"
+    else
+        @info "Setting up distributed ClimaCoupler run" nprocs = ClimaComms.nprocs(comms_ctx)
+    end
+
+    return comms_ctx
+end
+
 end # module

test/utilities_tests.jl

@@ -6,6 +6,8 @@ using Test
 using ClimaCoupler: Utilities, TestHelper
 using ClimaCore: Fields
 
+import ClimaComms
+
 for FT in (Float32, Float64)
     @testset "test swap_space!" begin
         space1 = TestHelper.create_space(FT, R = FT(6371e3))
@@ -19,4 +21,36 @@ for FT in (Float32, Float64)
         @test parent(field1) == parent(field2)
         @test axes(field2) == space2
     end
+
+    @testset "test comms_ctx" begin
+        parsed_args = Dict("device" => "auto")
+
+        this_device = ClimaComms.device()
+        this_ctx = ClimaComms.context(this_device)
+
+        @test typeof(Utilities.get_comms_context(parsed_args)) == typeof(this_ctx)
+
+        if typeof(this_device) == typeof(ClimaComms.CUDADevice())
+            parsed_args["device"] = "CUDADevice"
+        elseif typeof(this_device) == typeof(ClimaComms.CPUMultiThreaded())
+            parsed_args["device"] = "CPUMultiThreaded"
+        else
+            parsed_args["device"] = "CPUSingleThreaded"
+        end
+
+        ClimaComms.init(this_ctx)
+        @test typeof(Utilities.get_comms_context(parsed_args)) == typeof(this_ctx)
+
+        # Additional test calls for code coverage since Github Actions only exercises the SingleThreaded calculates
+        # More meaningful testing performed on buildkite
+        # Cannot test CUDADevice on CPU 
+        # Test other devices: 
+        parsed_args["device"] = "CPUMultiThreaded"
+        @test typeof(Utilities.get_comms_context(parsed_args)) ==
+              typeof(ClimaComms.context(ClimaComms.CPUMultiThreaded()))
+
+        parsed_args["device"] = ""
+        @test typeof(Utilities.get_comms_context(parsed_args)) ==
+              typeof(ClimaComms.context(ClimaComms.CPUSingleThreaded()))
+    end
 end