progress bars for EnsembleProblem

ODE->Ensemble

optionally aggregate progress bars

handle integer loglevel (sundials)

avoid lock contention

only log significant progress

improve progress performance

add version constraint

Update Project.toml

Update Project.toml

ignore derivatives of logging

more AD fixing attempts

only pass progress_id if needed

use ignore_deriviative and fix rrule for with_logger

delete rules moved to ChainRules

remove using

import as opposed to using

more import fixes

more missing Logging qualifiers

Author: Pepijn de Vos

src/ensemble/basic_ensemble_solve.jl

@@ -29,6 +29,64 @@ function merge_stats(us)
     reduce(merge, st)
 end
 
+mutable struct AggregateLogger{T<:Logging.AbstractLogger} <: Logging.AbstractLogger
+    progress::Dict{Symbol, Float64}
+    done_counter::Int
+    total::Float64
+    print_time::Float64
+    lock::ReentrantLock
+    logger::T
+end
+AggregateLogger(logger::Logging.AbstractLogger) = AggregateLogger(Dict{Symbol, Float64}(),0 , 0.0, 0.0, ReentrantLock(), logger)
+
+function Logging.handle_message(l::AggregateLogger, level, message, _module, group, id, file, line; kwargs...)
+    if convert(Logging.LogLevel, level) == Logging.LogLevel(-1) && haskey(kwargs, :progress)
+        pr = kwargs[:progress]
+        if trylock(l.lock) || (pr == "done" && lock(l.lock)===nothing)
+            try
+                if pr == "done"
+                    pr = 1.0
+                    l.done_counter += 1
+                end
+                len = length(l.progress)
+                if haskey(l.progress, id)
+                    l.total += (pr-l.progress[id])/len
+                else
+                    l.total = l.total*(len/(len+1)) + pr/(len+1)
+                    len += 1
+                end
+                l.progress[id] = pr
+                # validation check (slow)
+                # tot = sum(values(l.progress))/length(l.progress)
+                # @show tot l.total l.total ≈ tot
+                curr_time = time()
+                if l.done_counter >= len
+                    tot="done"
+                    empty!(l.progress)
+                    l.done_counter = 0
+                    l.print_time = 0.0
+                elseif curr_time-l.print_time > 0.1
+                    tot = l.total
+                    l.print_time = curr_time
+                else
+                    return
+                end
+                id=:total
+                message="Total"
+                kwargs=merge(values(kwargs), (progress=tot,))
+            finally
+                unlock(l.lock)
+            end
+        else
+            return
+        end
+    end
+    Logging.handle_message(l.logger, level, message, _module, group, id, file, line; kwargs...)
+end
+Logging.shouldlog(l::AggregateLogger, args...) = Logging.shouldlog(l.logger, args...)
+Logging.min_enabled_level(l::AggregateLogger) = Logging.min_enabled_level(l.logger)
+Logging.catch_exceptions(l::AggregateLogger) = Logging.catch_exceptions(l.logger)
+
 function __solve(prob::AbstractEnsembleProblem,
     alg::Union{AbstractDEAlgorithm, Nothing};
     kwargs...)
@@ -59,51 +117,71 @@ end
 function __solve(prob::AbstractEnsembleProblem,
     alg::A,
     ensemblealg::BasicEnsembleAlgorithm;
-    trajectories, batch_size = trajectories,
+    trajectories, batch_size = trajectories, progress_aggregate=true,
     pmap_batch_size = batch_size ÷ 100 > 0 ? batch_size ÷ 100 : 1, kwargs...) where {A}
-    num_batches = trajectories ÷ batch_size
-    num_batches < 1 &&
-        error("trajectories ÷ batch_size cannot be less than 1, got $num_batches")
-    num_batches * batch_size != trajectories && (num_batches += 1)
-
-    if num_batches == 1 && prob.reduction === DEFAULT_REDUCTION
-        elapsed_time = @elapsed u = solve_batch(prob, alg, ensemblealg, 1:trajectories,
-            pmap_batch_size; kwargs...)
-        _u = tighten_container_eltype(u)
-        stats = merge_stats(_u)
-        return EnsembleSolution(_u, elapsed_time, true, stats)
-    end
+    logger = progress_aggregate ? AggregateLogger(Logging.current_logger()) : Logging.current_logger()
+    
+    Logging.with_logger(logger) do
+        num_batches = trajectories ÷ batch_size
+        num_batches < 1 &&
+            error("trajectories ÷ batch_size cannot be less than 1, got $num_batches")
+        num_batches * batch_size != trajectories && (num_batches += 1)
 
-    converged::Bool = false
-    elapsed_time = @elapsed begin
-        i = 1
-        II = (batch_size * (i - 1) + 1):(batch_size * i)
+        if get(kwargs, :progress, false)
+            name = get(kwargs, :progress_name, "Ensemble")
+            for i in 1:trajectories
+                msg = "$name #$i"
+                Logging.@logmsg(Logging.LogLevel(-1), msg, _id=Symbol("SciMLBase_$i"), progress=0)
+            end
+        end
+        
 
-        batch_data = solve_batch(prob, alg, ensemblealg, II, pmap_batch_size; kwargs...)
+        if num_batches == 1 && prob.reduction === DEFAULT_REDUCTION
+            elapsed_time = @elapsed u = solve_batch(prob, alg, ensemblealg, 1:trajectories,
+                pmap_batch_size; kwargs...)
+            _u = tighten_container_eltype(u)
+            return EnsembleSolution(_u, elapsed_time, true)
+        end
+
+        converged::Bool = false
+        elapsed_time = @elapsed begin
+            i = 1
+            II = (batch_size * (i - 1) + 1):(batch_size * i)
 
-        u = prob.u_init === nothing ? similar(batch_data, 0) : prob.u_init
-        u, converged = prob.reduction(u, batch_data, II)
-        for i in 2:num_batches
-            converged && break
-            if i == num_batches
-                II = (batch_size * (i - 1) + 1):trajectories
-            else
-                II = (batch_size * (i - 1) + 1):(batch_size * i)
-            end
             batch_data = solve_batch(prob, alg, ensemblealg, II, pmap_batch_size; kwargs...)
+
+            u = prob.u_init === nothing ? similar(batch_data, 0) : prob.u_init
             u, converged = prob.reduction(u, batch_data, II)
+            for i in 2:num_batches
+                converged && break
+                if i == num_batches
+                    II = (batch_size * (i - 1) + 1):trajectories
+                else
+                    II = (batch_size * (i - 1) + 1):(batch_size * i)
+                end
+                batch_data = solve_batch(prob, alg, ensemblealg, II, pmap_batch_size; kwargs...)
+                u, converged = prob.reduction(u, batch_data, II)
+            end
         end
+        _u = tighten_container_eltype(u)
+        stats = merge_stats(_u)
+        return EnsembleSolution(_u, elapsed_time, converged, stats)
     end
-    _u = tighten_container_eltype(u)
-    stats = merge_stats(_u)
-    return EnsembleSolution(_u, elapsed_time, converged, stats)
 end
 
 function batch_func(i, prob, alg; kwargs...)
     iter = 1
     _prob = prob.safetycopy ? deepcopy(prob.prob) : prob.prob
     new_prob = prob.prob_func(_prob, i, iter)
     rerun = true
+
+    progress = get(kwargs, :progress, false)
+    if progress
+        name = get(kwargs, :progress_name, "Ensemble")
+        progress_name = "$name #$i"
+        progress_id = Symbol("SciMLBase_$i")
+        kwargs = (kwargs..., progress_name=progress_name, progress_id=progress_id)
+    end
     x = prob.output_func(solve(new_prob, alg; kwargs...), i)
     if !(x isa Tuple)
         rerun_warn()