save sypd, allocs to txt

CliMA · Mar 21, 2024 · dd039f2 · dd039f2
1 parent 2c6573a
commit dd039f2
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 14 deletions.
diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
@@ -177,7 +177,7 @@ steps:
 
       - label: "Slabplanet: albedo from static map"
         command: "julia --color=yes --project=experiments/AMIP/ experiments/AMIP/coupler_driver.jl --config_file $CONFIG_PATH/slabplanet_albedo_static_map.yml"
-        artifact_paths: "experiments/AMIP/output/slabplanet/slabplanet_albedo_static_map_artifacts/total_energy*.png"
+        artifact_paths: "experiments/AMIP/output/slabplanet/slabplanet_albedo_static_map_artifacts/total_energy*"
         agents:
           slurm_mem: 20GB
 
@@ -260,36 +260,36 @@ steps:
 
       - label: "Moist earth with slab surface - default: monin gray no_sponge idealinsol freq_dt_cpl"
         command: "julia --color=yes --project=experiments/AMIP/ experiments/AMIP/coupler_driver.jl --config_file $CONFIG_PATH/default_mono.yml"
-        artifact_paths: "experiments/AMIP/output/slabplanet/default_mono_artifacts/total_energy*.png"
+        artifact_paths: "experiments/AMIP/output/slabplanet/default_mono_artifacts/total_energy*"
         agents:
           slurm_mem: 20GB
 
       - label: "Moist earth with slab surface - notmono: monin gray no_sponge idealinsol freq_dt_cpl notmono"
         command: "julia --color=yes --project=experiments/AMIP/ experiments/AMIP/coupler_driver.jl --config_file $CONFIG_PATH/default_notmono.yml"
-        artifact_paths: "experiments/AMIP/output/slabplanet/default_notmono_artifacts/total_energy*.png"
+        artifact_paths: "experiments/AMIP/output/slabplanet/default_notmono_artifacts/total_energy*"
         agents:
           slurm_mem: 20GB
 
       # - label: "Moist earth with slab surface - test: monin allsky sponge idealinsol infreq_dt_cpl"
         # command: "julia --color=yes --project=experiments/AMIP/ experiments/AMIP/coupler_driver.jl --FLOAT_TYPE Float64 --coupled true --surface_setup PrescribedSurface --moist equil --vert_diff true --rad allskywithclear --rayleigh_sponge true --alpha_rayleigh_uh 0 --alpha_rayleigh_w 10 --energy_check true --mode_name slabplanet --t_end 10days --dt_save_to_sol 3600secs --dt_cpl 21600 --dt 200secs --dt_rad 6hours --mono_surface true --h_elem 4 --precip_model 0M --run_name target_params_in_slab_test1 --job_id target_params_in_slab_test1" # Unconverged SF (reproduced locally); works with 200s dt_cpl
-        # artifact_paths: "experiments/AMIP/output/slabplanet/target_params_in_slab_test1_artifacts/total_energy*.png"
+        # artifact_paths: "experiments/AMIP/output/slabplanet/target_params_in_slab_test1_artifacts/total_energy*"
 
       - label: "Moist earth with slab surface - test: bulk allsky sponge realinsol infreq_dt_cpl"
         command: "julia --color=yes --project=experiments/AMIP/ experiments/AMIP/coupler_driver.jl --config_file $CONFIG_PATH/target_params_in_slab_test2.yml"
-        artifact_paths: "experiments/AMIP/output/slabplanet/target_params_in_slab_test2_artifacts/total_energy*.png"
+        artifact_paths: "experiments/AMIP/output/slabplanet/target_params_in_slab_test2_artifacts/total_energy*.g"
         agents:
           slurm_mem: 20GB
 
       - label: "Moist earth with slab surface - test: monin gray sponge realinsol infreq_dt_cpl"
         command: "julia --color=yes --project=experiments/AMIP/ experiments/AMIP/coupler_driver.jl --config_file $CONFIG_PATH/target_params_in_slab_test3.yml"
-        artifact_paths: "experiments/AMIP/output/slabplanet/target_params_in_slab_test3_artifacts/total_energy*.png"
+        artifact_paths: "experiments/AMIP/output/slabplanet/target_params_in_slab_test3_artifacts/total_energy*"
         agents:
           slurm_mem: 20GB
 
       # breaking:
       # - label: "Moist earth with slab surface - monin allsky no_sponge idealinsol infreq_dt_cpl"
       #   command: "julia --color=yes --project=experiments/AMIP/ experiments/AMIP/coupler_driver.jl --coupled true --surface_setup PrescribedSurface --moist equil --vert_diff true --rad allskywithclear --rayleigh_sponge false --energy_check true --mode_name slabplanet --t_end 10days --dt_save_to_sol 3600secs --dt_cpl 21600 --dt 200secs --dt_rad 6hours --idealized_insolation true --mono_surface true --h_elem 4 --precip_model 0M --run_name target_params_in_slab1"
-      #   artifact_paths: "experiments/AMIP/output/slabplanet/target_params_in_slab1_artifacts/total_energy*.png"
+      #   artifact_paths: "experiments/AMIP/output/slabplanet/target_params_in_slab1_artifacts/total_energy*"
 
       - label: "AMIP target: albedo from function"
         key: "target_amip_albedo_function"

diff --git a/experiments/AMIP/coupler_driver.jl b/experiments/AMIP/coupler_driver.jl
@@ -32,8 +32,8 @@ using CUDA
     show_memory_usage(comms_ctx, objects)
 Display the current memory footprint of the simulation, using an appropriate
 method based on the device being used.
-In the GPU case, show the memory usage of the GPU.
-In the CPU case, show the memory footprint of the provided object(s).
+In the GPU case, show and return the memory usage of the GPU.
+In the CPU case, show and return the memory footprint of the provided object(s) in GB.
 Note that these two cases provide different information, and should not be
 directly compared.
 # Arguments
@@ -42,15 +42,21 @@ directly compared.
 """
 function show_memory_usage(comms_ctx, objects)
     if comms_ctx.device isa ClimaComms.CUDADevice
-        @info "Memory usage: $(CUDA.memory_status())"
+        mem_status = CUDA.memory_status()
+        @info "Memory usage: $mem_status"
+        return mem_status
     elseif comms_ctx.device isa ClimaComms.AbstractCPUDevice
         if ClimaComms.iamroot(comms_ctx)
+            cumul_size = 0
             for (obj, name) in objects
-                @info "Memory footprint of `$(name)` in bytes: $(Base.summarysize(obj))"
+                cumul_size += Base.summarysize(obj)
+                @info "Memory footprint of `$(name)` in GB: $(Base.summarysize(obj) / 1e9)"
             end
+            return cumul_size / 1e9
         end
     else
         @warn "Invalid device type $device; cannot show memory usage."
+        return nothing
     end
 end
 
@@ -816,11 +822,21 @@ end #hide
 
 ## run the coupled simulation
 walltime = solve_coupler!(cs);
-show_memory_usage(comms_ctx, Dict(cs => "cs"))
 
-## Use ClimaAtmos calculation to show the simulated years per day of the simulation
+## Use ClimaAtmos calculation to show the simulated years per day of the simulation (SYPD)
 es = CA.EfficiencyStats(tspan, walltime)
-@info "SYPD: $(CA.simulated_years_per_day(es))"
+sypd = CA.simulated_years_per_day(es)
+@info "SYPD: $sypd"
+
+## Save the SYPD and allocation information
+if ClimaComms.iamroot(comms_ctx)
+    sypd_filename = joinpath(COUPLER_ARTIFACTS_DIR, "SYPD.txt")
+    save_as_txt(sypd, sypd_filename)
+
+    allocs = show_memory_usage(comms_ctx, Dict(cs => "cs"))
+    allocs_filename = joinpath(COUPLER_ARTIFACTS_DIR, "allocations.txt")
+    save_as_txt(allocs, allocs_filename)
+end
 
 #=
 ## Postprocessing

diff --git a/experiments/AMIP/user_io/user_logging.jl b/experiments/AMIP/user_io/user_logging.jl
@@ -1,4 +1,5 @@
 using ClimaCoupler.Checkpointer: checkpoint_model_state
+import DelimitedFiles as DLM
 
 """
     Base.show(io::IO, dict::Dict)
@@ -11,6 +12,18 @@ function Base.show(io::IO, dict::Dict)
     end
 end
 
+"""
+    save_as_txt(var, filename::String)
+
+Save the input variable to a text file at `filename` using DelimitedFiles.jl.
+`filename` should end in `.txt`.
+"""
+function save_as_txt(var, filename::String)
+    open(filename, "w") do io
+        DLM.writedlm(io, var, ',')
+    end
+end
+
 # user callbacks
 """
     checkpoint_sims(cs::CoupledSimulation, _)