From ed08673bb414c31f80e0e8296f72ee0523b7911b Mon Sep 17 00:00:00 2001
From: nefrathenrici <nat.henrici@gmail.com>
Date: Fri, 18 Aug 2023 11:32:19 -0700
Subject: [PATCH] Add status callback

---
 config/default_configs/default_config.yml |  3 ++
 src/callbacks/callbacks.jl                | 41 +++++++++++++++++++++++
 src/solver/type_getters.jl                | 12 +++++++
 3 files changed, 56 insertions(+)

diff --git a/config/default_configs/default_config.yml b/config/default_configs/default_config.yml
index 720b5a0d32a..bc9ede12020 100644
--- a/config/default_configs/default_config.yml
+++ b/config/default_configs/default_config.yml
@@ -153,6 +153,9 @@ dt_save_restart:
 dt_save_to_sol:
   help: "Time between saving solution. Examples: [`10days`, `1hours`, `Inf` (do not save)]"
   value: "1days"
+dt_show_progress:
+  help: "Time between displaying progress update"
+  value: "600secs"
 moist:
   help: "Moisture model [`dry` (default), `equil`, `non_equil`]"
   value: "dry"
diff --git a/src/callbacks/callbacks.jl b/src/callbacks/callbacks.jl
index 57cd7445cef..7e9d3cfdbc2 100644
--- a/src/callbacks/callbacks.jl
+++ b/src/callbacks/callbacks.jl
@@ -16,6 +16,47 @@ import ClimaCore.Fields: ColumnField
 
 include("callback_helpers.jl")
 
+"""
+    display_status_callback!(::Type{tType})
+
+Given typeof(dt), returns a callback to display:
+ - percentage of work completed, 
+ - total wallclock time elapsed,
+ - estimated wallclock time remaining.
+Adapted from ClimaTimeSteppers.jl #89.
+"""
+function display_status_callback!(::Type{tType}) where {tType}
+    start_time = Ref{Float64}()
+    prev_time = Ref{Float64}()
+    current_time = Ref{Float64}()
+    prev_t = Ref{tType}()
+    # milliseconds
+    speed = Ref{Float64}()
+    eta = Ref{Float64}()
+    is_first_step = Ref{Bool}(true)
+
+    return function (integrator)
+        # speed = wallclock time / simulation time
+        # Print ETA = speed * remaining simulation time
+        t = integrator.t
+        t_end = integrator.p.simulation.t_end
+        current_time[] = round(time_ns()) / 1e9
+        speed[] = (current_time[] - prev_time[]) / (t - prev_t[])
+        eta[] = speed[] * (t_end - t)
+        if is_first_step[]
+            println("Time Remaining: ...")
+            is_first_step[] = false
+            start_time[] = current_time[]
+        else
+            println(Dates.now())
+            println("$(round(t / t_end * 100, digits=2))% complete in $(round(current_time[] - start_time[], digits=2)) seconds")
+            println("Time Remaining: $(round(eta[], digits=2)) seconds")
+        end
+        prev_t[] = t
+        prev_time[] = current_time[]
+    end
+end
+
 function dss_callback!(integrator)
     Y = integrator.u
     ghost_buffer = integrator.p.ghost_buffer
diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
index 459dfc57589..067a5406923 100644
--- a/src/solver/type_getters.jl
+++ b/src/solver/type_getters.jl
@@ -491,6 +491,18 @@ function get_callbacks(parsed_args, simulation, atmos, params)
             (callbacks..., call_every_dt(save_restart_func, dt_save_restart))
     end
 
+    dt_show_progress = time_to_seconds(parsed_args["dt_show_progress"])
+    if !(dt_show_progress == Inf)
+        callbacks = (
+            callbacks...,
+            call_every_dt(
+                display_status_callback!(typeof(dt)),
+                dt_show_progress;
+                skip_first = true,
+            ),
+        )
+    end
+
     if is_distributed(simulation.comms_ctx)
         callbacks = (
             callbacks...,