Skip to content

Commit

Permalink
move show_mem to driver
Browse files Browse the repository at this point in the history
  • Loading branch information
juliasloan25 committed Mar 18, 2024
1 parent 2325d9c commit bad4ef7
Show file tree
Hide file tree
Showing 4 changed files with 30 additions and 34 deletions.
2 changes: 0 additions & 2 deletions Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@ authors = ["CliMA Contributors <[email protected]>"]
version = "0.1.0"

[deps]
CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
ClimaAtmos = "b2c96348-7fb7-4fe0-8da9-78d88439e717"
ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884"
Expand All @@ -25,7 +24,6 @@ TempestRemap_jll = "8573a8c5-1df0-515e-a024-abad257ee284"
Thermodynamics = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c"

[compat]
CUDA = "5.2"
ClimaAtmos = "0.22"
ClimaComms = "0.5.6"
ClimaCore = "0.13"
Expand Down
1 change: 0 additions & 1 deletion docs/src/utilities.md
Original file line number Diff line number Diff line change
Expand Up @@ -10,5 +10,4 @@ modules in the coupler.
ClimaCoupler.Utilities.swap_space!
ClimaCoupler.Utilities.get_comms_context
ClimaCoupler.Utilities.get_device
ClimaCoupler.Utilities.show_memory_usage
```
31 changes: 30 additions & 1 deletion experiments/AMIP/coupler_driver.jl
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,35 @@ so we force abbreviated stacktraces even in non-interactive runs.
=#
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).
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
"""
function show_memory_usage(comms_ctx, objects)
if comms_ctx.device isa ClimaComms.CUDADevice
@info "Memory usage: $(CUDA.memory_status())"
elseif comms_ctx.device isa ClimaComms.AbstractCPUDevice
if ClimaComms.iamroot(comms_ctx)
for (obj, name) in objects
@info "Memory footprint of `$(name)` in bytes: $(Base.summarysize(obj))"
end
end
else
@warn "Invalid device type $device; cannot show memory usage."
end
end

redirect_stderr(IOContext(stderr, :stacktrace_types_limited => Ref(false)))

#=
Expand Down Expand Up @@ -65,7 +94,7 @@ using ClimaCoupler.Regridder
using ClimaCoupler.Regridder: update_surface_fractions!, combine_surfaces!, binary_mask
using ClimaCoupler.TimeManager:
current_date, Monthly, EveryTimestep, HourlyCallback, MonthlyCallback, update_firstdayofmonth!, trigger_callback!
import ClimaCoupler.Utilities: get_comms_context, show_memory_usage
import ClimaCoupler.Utilities: get_comms_context

pkg_dir = pkgdir(ClimaCoupler)

Expand Down
30 changes: 0 additions & 30 deletions src/Utilities.jl
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@ module Utilities

import ClimaComms
using ClimaCore: Fields, Spaces
using CUDA

export swap_space!

Expand Down Expand Up @@ -70,33 +69,4 @@ function get_comms_context(parsed_args)
return comms_ctx
end

"""
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).
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
"""
function show_memory_usage(comms_ctx, objects)
if comms_ctx.device isa ClimaComms.CUDADevice
@info "Memory usage: $(CUDA.memory_status())"
elseif comms_ctx.device isa ClimaComms.AbstractCPUDevice
if ClimaComms.iamroot(comms_ctx)
for (obj, name) in objects
@info "Memory footprint of `$(name)` in bytes: $(Base.summarysize(obj))"
end
end
else
@warn "Invalid device type $device; cannot show memory usage."
end
end

end # module

0 comments on commit bad4ef7

Please sign in to comment.