fixup! Add Task Monitor

Author: jpsamaroo

Project.toml

@@ -21,6 +21,7 @@ SharedArrays = "1a1011a3-84de-559e-8e89-a11a2f7dc383"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
+TaskLocalValues = "ed4db957-447d-4319-bfb6-7fa9ae7ecf34"
 TimespanLogging = "a526e669-04d3-4846-9525-c66122c55f63"
 UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 
@@ -46,6 +47,7 @@ Requires = "1"
 ScopedValues = "1.1"
 Statistics = "1"
 StatsBase = "0.28, 0.29, 0.30, 0.31, 0.32, 0.33, 0.34"
+TaskLocalValues = "0.1"
 TimespanLogging = "0.1"
 julia = "1.8"
 

src/queue.jl

@@ -14,13 +14,10 @@ enqueue!(::EagerTaskQueue, spec::Pair{EagerTaskSpec,EagerThunk}) =
 enqueue!(::EagerTaskQueue, specs::Vector{Pair{EagerTaskSpec,EagerThunk}}) =
     eager_launch!(specs)
 
-function enqueue!(spec::Pair{EagerTaskSpec,EagerThunk}) 
-    println("Hi There")
-    enqueue!(MONITOR_QUEUE,spec)
-end
-
+enqueue!(spec::Pair{EagerTaskSpec,EagerThunk}) =
+    enqueue!(get_options(:task_queue, MONITOR_QUEUE[]), spec)
 enqueue!(specs::Vector{Pair{EagerTaskSpec,EagerThunk}}) =
-    enqueue!(get_options(:task_queue, EagerTaskQueue()), specs)
+    enqueue!(get_options(:task_queue, MONITOR_QUEUE[]), specs)
 
 struct LazyTaskQueue <: AbstractTaskQueue
     tasks::Vector{Pair{EagerTaskSpec,EagerThunk}}

src/task_monitor.jl

@@ -1,36 +1,47 @@
 using ProgressMeter
+using TaskLocalValues
 
 struct MonitorTaskQueue <: AbstractTaskQueue
-    runningTasks::Vector{Pair{EagerTaskSpec,EagerThunk}}
-    finishedTasks::Vector{Pair{EagerTaskSpec,EagerThunk}}
-    MonitorTaskQueue() = new(Pair{EagerTaskSpec,EagerThunk}[], Pair{EagerTaskSpec,EagerThunk}[])
+    running_tasks::Vector{WeakRef}
+    MonitorTaskQueue() = new(WeakRef[])
 end
 function enqueue!(queue::MonitorTaskQueue, spec::Pair{EagerTaskSpec,EagerThunk})
-    push!(queue.runningTasks, spec)
+    push!(queue.running_tasks, WeakRef(spec[2]))
     upper = get_options(:task_queue, EagerTaskQueue())
-    enqueue!(upper,spec)    
+    enqueue!(upper, spec)
 end
 
 function enqueue!(queue::MonitorTaskQueue, specs::Vector{Pair{EagerTaskSpec,EagerThunk}})
-    append!(queue.runningTasks, specs)
+    for (_, task) in specs
+        push!(queue.running_tasks, WeakRef(task))
+    end
+    upper = get_options(:task_queue, EagerTaskQueue())
+    enqueue!(upper, specs)
 end
 
-const MONITOR_QUEUE = MonitorTaskQueue()
+const MONITOR_QUEUE = TaskLocalValue{MonitorTaskQueue}(MonitorTaskQueue)
 
+"Monitors and displays the progress of any still-executing tasks."
 function monitor()
-        errormonitor(Threads.@spawn begin
-            runningTasks = MONITOR_QUEUE.runningTasks
-            finishedTasks = MONITOR_QUEUE.finishedTasks
-            meter = Progress(length(runningTasks))
-            while !isempty(runningTasks)
-                for (i,task) in reverse(collect(enumerate(runningTasks)))
-                    if isready(task[2])
-                        next!(meter)
-                        push!(finishedTasks,task)
-                        deleteat!(runningTasks,i)
-                    end
-                end
-                sleep(0.1)
+    queue = MONITOR_QUEUE[]
+    running_tasks = queue.running_tasks
+    isempty(running_tasks) && return
+
+    ntasks = length(running_tasks)
+    meter = Progress(ntasks;
+                     desc="Waiting for $ntasks tasks...",
+                     dt=0.01, showspeed=true)
+    while !isempty(running_tasks)
+        for (i, task_weak) in reverse(collect(enumerate(running_tasks)))
+            task = task_weak.value
+            if task === nothing || isready(task)
+                next!(meter)
+                deleteat!(running_tasks, i)
             end
-        end)
-end
+        end
+        sleep(0.01)
+    end
+    finish!(meter)
+
+    return
+end

src/thunk.jl

@@ -309,7 +309,7 @@ function spawn(f, args...; kwargs...)
     args_kwargs = args_kwargs_to_pairs(args, kwargs)
 
     # Get task queue, and don't let it propagate
-    task_queue = get_options(:task_queue, MONITOR_QUEUE)
+    task_queue = get_options(:task_queue, MONITOR_QUEUE[])
     options = NamedTuple(filter(opt->opt[1] != :task_queue, Base.pairs(options)))
     propagates = filter(prop->prop != :task_queue, propagates)
     options = merge(options, (;propagates))