Progress bars when sampling multiple chains

Project.toml

@@ -3,7 +3,7 @@ uuid = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
 keywords = ["markov chain monte carlo", "probabilistic programming"]
 license = "MIT"
 desc = "A lightweight interface for common MCMC methods."
-version = "5.6.3"
+version = "5.7.0"
 
 [deps]
 BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"
@@ -18,6 +18,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 TerminalLoggers = "5d786b92-1e48-4d6f-9151-6b4477ca9bed"
 Transducers = "28d57a85-8fef-5791-bfe6-a80928e7c999"
+UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 
 [compat]
 BangBang = "0.3.19, 0.4"
@@ -29,6 +30,7 @@ ProgressLogging = "0.1"
 StatsBase = "0.32, 0.33, 0.34"
 TerminalLoggers = "0.1"
 Transducers = "0.4.30"
+UUIDs = "<0.0.1, 1"
 julia = "1.6"
 
 [extras]

docs/src/api.md

@@ -68,7 +68,7 @@ AbstractMCMC.MCMCSerial
 ## Common keyword arguments
 
 Common keyword arguments for regular and parallel sampling are:
-- `progress` (default: `AbstractMCMC.PROGRESS[]` which is `true` initially):  toggles progress logging
+- `progress` (default: `AbstractMCMC.PROGRESS[]` which is `true` initially): toggles progress logging. See the section on [Progress logging](#progress-logging) below for more details.
 - `chain_type` (default: `Any`): determines the type of the returned chain
 - `callback` (default: `nothing`): if `callback !== nothing`, then
   `callback(rng, model, sampler, sample, iteration)` is called after every sampling step,
@@ -90,12 +90,45 @@ However, multiple packages such as [EllipticalSliceSampling.jl](https://github.c
 To ensure that sampling multiple chains "just works" when sampling of a single chain is implemented, [we decided to support `initial_params` in the default implementations of the ensemble methods](https://github.com/TuringLang/AbstractMCMC.jl/pull/94):
 - `initial_params` (default: `nothing`): if `initial_params isa AbstractArray`, then the `i`th element of `initial_params` is used as initial parameters of the `i`th chain. If one wants to use the same initial parameters `x` for every chain, one can specify e.g. `initial_params = FillArrays.Fill(x, N)`.
 
-Progress logging can be enabled and disabled globally with `AbstractMCMC.setprogress!(progress)`.
+## Progress logging
+
+Progress logging is controlled in one of two ways:
+
+- by passing the `progress` keyword argument to the `sample(...)` function, or
+- by globally changing the defaults with `AbstractMCMC.setprogress!` and `AbstractMCMC.setmaxchainsprogress!`.
+
+### `progress` keyword argument
+
+For single-chain sampling (i.e., `sample([rng,] model, sampler, N)`), as well as multiple-chain sampling with `MCMCSerial`, the `progress` keyword argument should be a `Bool`.
+
+For multiple-chain sampling using `MCMCThreads`, there are several, more detailed, options:
+
+- `:overall`: create one progress bar for the overall sampling process, which tracks the percentage of samples that have been sampled across all chains
+- `:perchain`: in addition to `:overall`, also create one progress bar for each individual chain
+- `:none`: do not create any progress bar
+- `true` (the default): same as `:overall`, i.e. one progress bar for the overall sampling process
+- `false`: same as `:none`, i.e. no progress bar
+
+Multiple-chain sampling using `MCMCDistributed` behaves the same as `MCMCThreads`, except that `:perchain` is not (yet?) implemented.
+
+!!! warning "Do not override the `progress` keyword argument"
+    If you are implementing your own methods for `sample(...)`, you should make sure to not override the `progress` keyword argument if you want progress logging in multi-chain sampling to work correctly, as the multi-chain `sample()` call makes sure to specifically pass custom values of `progress` to the single-chain calls.
+
+### Global settings
+
+If you are sampling multiple times and would like to change the default behaviour, you can use this function to control progress logging globally:
 
 ```@docs
 AbstractMCMC.setprogress!
 ```
 
+`setprogress!` is more general, and applies to all types of sampling (both single- and multiple-chain).
+It only takes a boolean argument, which switches progress logging on or off.
+For example, `setprogress!(false)` will disable all progress logging.
+
+Note that `setprogress!` cannot be used to set the type of progress bar for multiple-chain sampling.
+If you want to use `:perchain`, it has to be set on each individual call to `sample`.
+
 ## Chains
 
 The `chain_type` keyword argument allows to set the type of the returned chain. A common

src/AbstractMCMC.jl

@@ -13,6 +13,7 @@ using FillArrays: FillArrays
 using Distributed: Distributed
 using Logging: Logging
 using Random: Random
+using UUIDs: UUIDs
 
 # Reexport sample
 using StatsBase: sample

src/logging.jl

@@ -1,19 +1,138 @@
-# 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(progress, exprs...)
+"""
+    AbstractProgressKwarg
+
+Abstract type representing the values that the `progress` keyword argument can
+internally take for single-chain sampling.
+"""
+abstract type AbstractProgressKwarg end
+
+"""
+    CreateNewProgressBar
+
+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())
+    end
+end
+function init_progress!(p::CreateNewProgressBar)
+    ProgressLogging.@logprogress p.name nothing _id = p.uuid
+end
+function update_progress!(p::CreateNewProgressBar, progress_frac)
+    ProgressLogging.@logprogress p.name progress_frac _id = p.uuid
+end
+function finish_progress!(p::CreateNewProgressBar)
+    ProgressLogging.@logprogress p.name "done" _id = p.uuid
+end
+get_n_updates(::CreateNewProgressBar) = 200
+
+"""
+    NoLogging
+
+Do not log progress at all.
+"""
+struct NoLogging <: AbstractProgressKwarg end
+init_progress!(::NoLogging) = nothing
+update_progress!(::NoLogging, ::Any) = nothing
+finish_progress!(::NoLogging) = nothing
+get_n_updates(::NoLogging) = 200
+
+"""
+    ExistingProgressBar
+Use an existing progress bar to log progress. This is used for tracking
+progress in a progress bar that has been previously generated elsewhere,
+specifically, during multi-threaded sampling where per-chain progress
+bars are requested. In this case we can use `@logprogress name progress_frac
+_id = uuid` to log progress.
+"""
+struct ExistingProgressBar{S<:AbstractString} <: AbstractProgressKwarg
+    name::S
+    uuid::UUIDs.UUID
+end
+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)
+    ProgressLogging.@logprogress p.name "done" _id = p.uuid
+end
+get_n_updates(::ExistingProgressBar) = 200
+
+"""
+    ChannelProgress
+
+Use a `Channel` to log progress. This is used for 'reporting' progress back to
+the main thread or worker when using multi-threaded or distributed sampling.
+
+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
+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
+get_n_updates(p::ChannelProgress) = p.n_updates
+
+"""
+    ChannelPlusExistingProgress
+
+Send updates to two places: a `Channel` as well as an existing progress bar.
+"""
+struct ChannelPlusExistingProgress{C<:ChannelProgress,E<:ExistingProgressBar} <:
+       AbstractProgressKwarg
+    channel_progress::C
+    existing_progress::E
+end
+function init_progress!(p::ChannelPlusExistingProgress)
+    init_progress!(p.channel_progress)
+    init_progress!(p.existing_progress)
+    return nothing
+end
+function update_progress!(p::ChannelPlusExistingProgress, progress_frac)
+    update_progress!(p.channel_progress, progress_frac)
+    update_progress!(p.existing_progress, progress_frac)
+    return nothing
+end
+function finish_progress!(p::ChannelPlusExistingProgress)
+    finish_progress!(p.channel_progress)
+    finish_progress!(p.existing_progress)
+    return nothing
+end
+get_n_updates(p::ChannelPlusExistingProgress) = get_n_updates(p.channel_progress)
+
+# Add a custom progress logger if the current logger does not seem to be able to handle
+# progress logs.
+macro maybewithricherlogger(expr)
     return esc(
         quote
-            if $progress
-                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
-                $(exprs[end])
+                $(expr)
             end
         end,
     )
@@ -39,8 +158,7 @@ end
 function progresslogger()
     # detect if code is running under IJulia since TerminalLogger does not work with IJulia
     # https://github.com/JuliaLang/IJulia.jl#detecting-that-code-is-running-under-ijulia
-    if (Sys.iswindows() && VERSION < v"1.5.3") ||
-        (isdefined(Main, :IJulia) && Main.IJulia.inited)
+    if (isdefined(Main, :IJulia) && Main.IJulia.inited)
         return ConsoleProgressMonitor.ProgressLogger()
     else
         return TerminalLoggers.TerminalLogger()