Refactor logging, throttle per-chain updates

Author: penelopeysm

src/logging.jl

@@ -14,25 +14,19 @@ Create a new logger for progress logging.
 struct CreateNewProgressBar{S<:AbstractString} <: AbstractProgressKwarg
     name::S
     uuid::UUIDs.UUID
-
     function CreateNewProgressBar(name::AbstractString)
-        return new{typeof{name}}(name, UUIDs.uuid4())
+        return new{typeof(name)}(name, UUIDs.uuid4())
     end
 end
 function init_progress(p::CreateNewProgressBar)
-    if hasprogresslevel(Logging.current_logger())
-        ProgressLogging.@withprogress $(exprs...)
-    else
-        $with_progresslogger($Base.@__MODULE__, $Logging.current_logger()) do
-            $ProgressLogging.@withprogress $(exprs...)
-        end
-    end
     ProgressLogging.@logprogress p.name nothing _id = p.uuid
 end
-function update_progress(p::CreateNewProgressBar, progress_frac, ::Bool)
+function update_progress(p::CreateNewProgressBar, progress_frac)
     ProgressLogging.@logprogress p.name progress_frac _id = p.uuid
 end
-finish_progress(::CreateNewProgressBar) = ProgressLogging.@logprogress "done"
+function finish_progress(p::CreateNewProgressBar)
+    ProgressLogging.@logprogress p.name "done" _id = p.uuid
+end
 
 """
     NoLogging
@@ -41,7 +35,7 @@ Do not log progress at all.
 """
 struct NoLogging <: AbstractProgressKwarg end
 init_progress(::NoLogging) = nothing
-update_progress(::NoLogging, ::Any, ::Bool) = nothing
+update_progress(::NoLogging, ::Any) = nothing
 finish_progress(::NoLogging) = nothing
 
 """
@@ -57,8 +51,21 @@ struct ExistingProgressBar{S<:AbstractString} <: AbstractProgressKwarg
     name::S
     uuid::UUIDs.UUID
 end
-init_progress(::ExistingProgressBar) = nothing
-function update_progress(p::ExistingProgressBar, progress_frac, ::Bool)
+function init_progress(p::ExistingProgressBar)
+    # Hacky code to reset the start timer if called from a multi-chain sampling
+    # process. We need this because the progress bar is constructed in the
+    # multi-chain method, i.e. if we don't do this the progress bar shows the
+    # time elapsed since _all_ sampling began, not since the current chain
+    # started.
+    try
+        bartrees = Logging.current_logger().loggers[1].logger.bartrees
+        bar = TerminalLoggers.findbar(bartrees, p.uuid).data
+        bar.tfirst = time()
+    catch
+    end
+    ProgressLogging.@logprogress p.name nothing _id = p.uuid
+end
+function update_progress(p::ExistingProgressBar, progress_frac)
     ProgressLogging.@logprogress p.name progress_frac _id = p.uuid
 end
 function finish_progress(p::ExistingProgressBar)
@@ -71,39 +78,32 @@ end
 Use a `Channel` to log progress. This is used for 'reporting' progress back
 to the main thread or worker when using `progress=:overall` with MCMCThreads or
 MCMCDistributed.
+
+n_updates is the number of updates that each child thread is expected to report
+back to the main thread.
 """
 struct ChannelProgress{T<:Union{Channel{Bool},Distributed.RemoteChannel{Channel{Bool}}}} <:
        AbstractProgressKwarg
     channel::T
+    n_updates::Int
 end
 init_progress(::ChannelProgress) = nothing
-function update_progress(p::ChannelProgress, ::Any, update_channel::Bool)
-    return update_channel && put!(p.channel, true)
-end
+update_progress(p::ChannelProgress, ::Any) = put!(p.channel, true)
 # Note: We don't want to `put!(p.channel, false)`, because that would stop the
 # channel from being used for further updates e.g. from other chains.
 finish_progress(::ChannelProgress) = nothing
 
-# avoid creating a progress bar with @withprogress if progress logging is disabled
-# and add a custom progress logger if the current logger does not seem to be able to handle
-# progress logs
-macro ifwithprogresslogger(cond, exprs...)
+# Add a custom progress logger if the current logger does not seem to be able to handle
+# progress logs.
+macro withprogresslogger(expr)
     return esc(
         quote
-            if $cond
-                # Create a new logger
-                if $hasprogresslevel($Logging.current_logger())
-                    $ProgressLogging.@withprogress $(exprs...)
-                else
-                    $with_progresslogger($Base.@__MODULE__, $Logging.current_logger()) do
-                        $ProgressLogging.@withprogress $(exprs...)
-                    end
+            if !($hasprogresslevel($Logging.current_logger()))
+                $with_progresslogger($Base.@__MODULE__, $Logging.current_logger()) do
+                    $(expr)
                 end
             else
-                # Don't create a new logger, either because progress logging
-                # was disabled, or because it's otherwise being manually
-                # managed.
-                $(exprs[end])
+                $(expr)
             end
         end,
     )