Sch: Use MetricsTracker

Author: jpsamaroo

Project.toml

@@ -11,6 +11,7 @@ Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 MemPool = "f9f48841-c794-520a-933b-121f7ba6ed94"
+MetricsTracker = "9a9c6fec-044d-4a27-aa18-2b01ca4026eb"
 OnlineStats = "a15396b6-48d5-5d58-9928-6d29437db91e"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"

src/Dagger.jl

@@ -43,6 +43,10 @@ else
     import Distributed: Future, RemoteChannel, myid, workers, nworkers, procs, remotecall, remotecall_wait, remotecall_fetch, check_same_host
 end
 
+import MetricsTracker as MT
+const reuse_metrics = @load_preference("reuse-metrics", false)
+const metrics_path = @load_preference("metrics-path", "metrics.json")
+
 include("lib/util.jl")
 include("utils/dagdebug.jl")
 
@@ -67,6 +71,9 @@ include("submission.jl")
 include("chunks.jl")
 include("memory-spaces.jl")
 
+# Metrics
+include("utils/metrics.jl")
+
 # Task scheduling
 include("compute.jl")
 include("utils/clock.jl")
@@ -126,6 +133,30 @@ function set_distributed_package!(value)
     @info "Dagger.jl preference has been set, restart your Julia session for this change to take effect!"
 end
 
+"""
+    set_reuse_metrics!(value::Bool)
+
+Set a [preference](https://github.com/JuliaPackaging/Preferences.jl) for
+enabling or disabling the reuse of collected metrics across Julia sessions.
+You will need to restart Julia after setting a new preference.
+"""
+function set_reuse_metrics!(value::Bool)
+    @set_preferences!("reuse-metrics" => value)
+    @info "Dagger.jl preference has been set, restart your Julia session for this change to take effect!"
+end
+
+"""
+    set_metrics_path!(value::String)
+
+Set a [preference](https://github.com/JuliaPackaging/Preferences.jl) for
+the path to save and load metrics. You will need to restart Julia after setting
+a new preference.
+"""
+function set_metrics_path!(value::String)
+    @set_preferences!("metrics-path" => value)
+    @info "Dagger.jl preference has been set, restart your Julia session for this change to take effect!"
+end
+
 # Precompilation
 import PrecompileTools: @compile_workload
 include("precompile.jl")
@@ -189,6 +220,26 @@ function __init__()
     catch err
         @warn "Error parsing JULIA_DAGGER_DEBUG" exception=err
     end
+
+    if reuse_metrics
+        if isfile(metrics_path)
+            # Load metrics
+            @dagdebug nothing :metrics "Loading metrics"
+            try
+                MT.load_metrics!(metrics_path)
+            catch err
+                @warn "Error loading metrics" exception=(err, catch_backtrace())
+            end
+        else
+            @dagdebug nothing :metrics "Metrics file not found"
+        end
+
+        atexit() do
+            # Save metrics on exit
+            @dagdebug nothing :metrics "Saving metrics"
+            MT.save_metrics(metrics_path)
+        end
+    end
 end
 
 end # module

src/sch/Sch.jl

@@ -15,9 +15,12 @@ import Base: @invokelatest
 
 import ..Dagger
 import ..Dagger: Context, Processor, Thunk, WeakThunk, ThunkFuture, DTaskFailedException, Chunk, WeakChunk, OSProc, AnyScope, DefaultScope, LockedObject
-import ..Dagger: order, dependents, noffspring, istask, inputs, unwrap_weak_checked, affinity, tochunk, timespan_start, timespan_finish, procs, move, chunktype, processor, get_processors, get_parent, execute!, rmprocs!, task_processor, constrain, cputhreadtime
+import ..Dagger: order, dependents, noffspring, istask, inputs, unwrap_weak_checked, affinity, tochunk, timespan_start, timespan_finish, procs, move, chunktype, processor, get_processors, get_parent, execute!, rmprocs!, task_processor, constrain
 import ..Dagger: @dagdebug, @safe_lock_spin1
 import DataStructures: PriorityQueue, enqueue!, dequeue_pair!, peek
+import ScopedValues: @with
+
+import MetricsTracker as MT
 
 import ..Dagger
 
@@ -1648,6 +1651,12 @@ function do_task(to_proc, task_desc)
         end
     end
 
+    # Compute signature
+    @warn "Fix kwargs" maxlog=1
+    sig = DataType[Tf, map(fetched_args) do x
+        chunktype(x)
+    end...]
+
     #= FIXME: If MaxUtilization, stop processors and wait
     if (est_time_util isa MaxUtilization) && (real_time_util > 0)
         # FIXME: Stop processors
@@ -1660,8 +1669,11 @@ function do_task(to_proc, task_desc)
     timespan_start(ctx, :compute, (;thunk_id, processor=to_proc), (;f))
     res = nothing
 
-    # Start counting time and GC allocations
-    threadtime_start = cputhreadtime()
+    # Setup metrics for time monitoring
+    mspec = MT.MetricsSpec(MT.TimeMetric(), Dagger.SignatureMetric(), Dagger.ProcessorMetric())
+    local_cache = MT.MetricsCache()
+
+    # Start counting GC allocations
     # FIXME
     #gcnum_start = Base.gc_num()
 
@@ -1677,9 +1689,13 @@ function do_task(to_proc, task_desc)
             cancel_token=Dagger.DTASK_CANCEL_TOKEN[],
         ))
 
+        # Execute
         res = Dagger.with_options(propagated) do
-            # Execute
-            execute!(to_proc, f, fetched_args...; fetched_kwargs...)
+            @with Dagger.TASK_SIGNATURE=>sig Dagger.TASK_PROCESSOR=>to_proc begin
+                MT.@with_metrics mspec Dagger :execute! thunk_id MT.SyncInto(local_cache) begin
+                    execute!(to_proc, f, fetched_args...; fetched_kwargs...)
+                end
+            end
         end
 
         # Check if result is safe to store
@@ -1705,10 +1721,16 @@ function do_task(to_proc, task_desc)
         RemoteException(myid(), CapturedException(ex, bt))
     end
 
-    threadtime = cputhreadtime() - threadtime_start
+    lock(MT.GLOBAL_METRICS_CACHE) do global_cache
+        MT.sync_results_into!(global_cache, local_cache)
+    end
+
     # FIXME: This is not a realistic measure of max. required memory
     #gc_allocd = min(max(UInt64(Base.gc_num().allocd) - UInt64(gcnum_start.allocd), UInt64(0)), UInt64(1024^4))
     timespan_finish(ctx, :compute, (;thunk_id, processor=to_proc), (;f, result=result_meta))
+
+    threadtime = MT.cache_lookup(local_cache, Dagger, :execute!, thunk_id, MT.TimeMetric())
+
     lock(TASK_SYNC) do
         real_time_util[] -= est_time_util
         pop!(TASKS_RUNNING, thunk_id)
@@ -1723,7 +1745,7 @@ function do_task(to_proc, task_desc)
         storage_pressure=real_alloc_util,
         storage_capacity=storage_cap,
         loadavg=((Sys.loadavg()...,) ./ Sys.CPU_THREADS),
-        threadtime=threadtime,
+        threadtime,
         # FIXME: Add runtime allocation tracking
         gc_allocd=(isa(result_meta, Chunk) ? result_meta.handle.size : 0),
         transfer_rate=(transfer_size[] > 0 && transfer_time[] > 0) ? round(UInt64, transfer_size[] / (transfer_time[] / 10^9)) : nothing,

src/utils/metrics.jl

@@ -0,0 +1,15 @@
+const TASK_SIGNATURE = ScopedValue{Union{Vector{DataType}, Nothing}}(nothing)
+struct SignatureMetric <: MT.AbstractMetric end
+MT.metric_applies(::SignatureMetric, ::Val{:execute!}) = true
+MT.metric_type(::SignatureMetric) = Union{Vector{DataType}, Nothing}
+MT.start_metric(::SignatureMetric) = nothing
+MT.stop_metric(::SignatureMetric, _) = TASK_SIGNATURE[]
+
+const TASK_PROCESSOR = ScopedValue{Union{Processor, Nothing}}(nothing)
+struct ProcessorMetric <: MT.AbstractMetric end
+MT.metric_applies(::ProcessorMetric, ::Val{:execute!}) = true
+MT.metric_type(::ProcessorMetric) = Union{Processor, Nothing}
+MT.start_metric(::ProcessorMetric) = nothing
+MT.stop_metric(::ProcessorMetric, _) = TASK_PROCESSOR[]
+
+# FIXME: struct TransferTimeMetric <: MT.AbstractMetric end