diff --git a/experiments/AMIP/coupler_driver.jl b/experiments/AMIP/coupler_driver.jl index 69ebdd49f9..8a8c4a665f 100644 --- a/experiments/AMIP/coupler_driver.jl +++ b/experiments/AMIP/coupler_driver.jl @@ -29,22 +29,26 @@ redirect_stderr(IOContext(stderr, :stacktrace_types_limited => Ref(false))) using CUDA """ - show_memory_usage(comms_ctx, objects) + show_memory_usage(comms_ctx, objects, io::Union{Base.TTY, IOBuffer} = stdout) + Display the current memory footprint of the simulation, using an appropriate method based on the device being used. -In the GPU case, show and return the memory usage of the GPU. +In the GPU case, show and return the memory usage of the GPU. Note that the +return value in this case is a vector of unicode hex characters that must be +converted to text to read (e.g. via a call`write(file, alloc_info)`). 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 `comms_ctx`: the communication context being used to run the model `objects`: Dict mapping objects whose memory footprint is displayed in the CPU case to their names +`io`: The input/output stream to use. Currently, `stdout` and `IOBuffer` types are supported. """ -function show_memory_usage(comms_ctx, objects) +function show_memory_usage(comms_ctx, objects, io::Union{Base.TTY, IOBuffer} = stdout) if comms_ctx.device isa ClimaComms.CUDADevice - mem_status = CUDA.memory_status() - @info "Memory usage: $mem_status" - return mem_status + # If `io` is `stdout`, print the memory status, otherwise store in buffer to return + CUDA.memory_status(io) + return io != stdout ? take!(io) : nothing elseif comms_ctx.device isa ClimaComms.AbstractCPUDevice if ClimaComms.iamroot(comms_ctx) cumul_size = 0 @@ -830,12 +834,14 @@ sypd = CA.simulated_years_per_day(es) ## 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) + sypd_filename = joinpath(COUPLER_ARTIFACTS_DIR, "sypd.txt") + write(sypd_filename, "$sypd") - allocs = show_memory_usage(comms_ctx, Dict(cs => "cs")) + # For GPU allocation information, we have to save to a buffer before writing + alloc_buf = IOBuffer() + allocs = show_memory_usage(comms_ctx, Dict(cs => "cs"), alloc_buf) allocs_filename = joinpath(COUPLER_ARTIFACTS_DIR, "allocations.txt") - save_as_txt(allocs, allocs_filename) + write(allocs_filename, allocs) end #= diff --git a/experiments/AMIP/user_io/user_logging.jl b/experiments/AMIP/user_io/user_logging.jl index b67fcd5ecb..0d43338adb 100644 --- a/experiments/AMIP/user_io/user_logging.jl +++ b/experiments/AMIP/user_io/user_logging.jl @@ -12,18 +12,6 @@ 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, _)