Merge pull request #464 from JuliaParallel/jps/chained-dtors

Eagerly free Thunk cached results when unneeded

Author: jpsamaroo

Diff for: .buildkite/pipeline.yml

@@ -34,6 +34,16 @@ steps:
           julia_args: "--threads=1"
       - JuliaCI/julia-coverage#v1:
           codecov: true
+  - label: Julia 1.10
+    timeout_in_minutes: 60
+    <<: *test
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "1.10"
+      - JuliaCI/julia-test#v1:
+          julia_args: "--threads=1"
+      - JuliaCI/julia-coverage#v1:
+          codecov: true
   - label: Julia nightly
     timeout_in_minutes: 60
     <<: *test

Diff for: Manifest.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.8.5"
 manifest_format = "2.0"
-project_hash = "5333a6c200b6e6add81c46547527f66ddc0dc16c"
+project_hash = "8da7911e4788068aaea8c0ef8589d674bce0fb39"
 
 [[deps.Artifacts]]
 uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
@@ -12,9 +12,9 @@ uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
 
 [[deps.ChainRulesCore]]
 deps = ["Compat", "LinearAlgebra", "SparseArrays"]
-git-tree-sha1 = "b66b8f8e3db5d7835fb8cbe2589ffd1cd456e491"
+git-tree-sha1 = "2118cb2765f8197b08e5958cdd17c165427425ee"
 uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-version = "1.17.0"
+version = "1.19.0"
 
 [[deps.ChangesOfVariables]]
 deps = ["InverseFunctions", "LinearAlgebra", "Test"]
@@ -24,9 +24,9 @@ version = "0.1.8"
 
 [[deps.Compat]]
 deps = ["Dates", "LinearAlgebra", "UUIDs"]
-git-tree-sha1 = "8a62af3e248a8c4bad6b32cbbe663ae02275e32c"
+git-tree-sha1 = "886826d76ea9e72b35fcd000e535588f7b60f21d"
 uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
-version = "4.10.0"
+version = "4.10.1"
 
 [[deps.CompilerSupportLibraries_jll]]
 deps = ["Artifacts", "Libdl"]
@@ -100,19 +100,19 @@ uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"
 
 [[deps.MacroTools]]
 deps = ["Markdown", "Random"]
-git-tree-sha1 = "9ee1618cbf5240e6d4e0371d6f24065083f60c48"
+git-tree-sha1 = "b211c553c199c111d998ecdaf7623d1b89b69f93"
 uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
-version = "0.5.11"
+version = "0.5.12"
 
 [[deps.Markdown]]
 deps = ["Base64"]
 uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"
 
 [[deps.MemPool]]
-deps = ["DataStructures", "Distributed", "Mmap", "Random", "Serialization", "Sockets"]
-git-tree-sha1 = "b9c1a032c3c1310a857c061ce487c632eaa1faa4"
+deps = ["DataStructures", "Distributed", "Mmap", "Random", "ScopedValues", "Serialization", "Sockets"]
+git-tree-sha1 = "60dd4ac427d39e0b3f15b193845324523ee71c03"
 uuid = "f9f48841-c794-520a-933b-121f7ba6ed94"
-version = "0.4.4"
+version = "0.4.6"
 
 [[deps.Missings]]
 deps = ["DataAPI"]
@@ -133,9 +133,9 @@ uuid = "4536629a-c528-5b80-bd46-f80d51c5b363"
 version = "0.3.20+0"
 
 [[deps.OrderedCollections]]
-git-tree-sha1 = "2e73fe17cac3c62ad1aebe70d44c963c3cfdc3e3"
+git-tree-sha1 = "dfdf5519f235516220579f949664f1bf44e741c5"
 uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
-version = "1.6.2"
+version = "1.6.3"
 
 [[deps.PrecompileTools]]
 deps = ["Preferences"]
@@ -173,9 +173,9 @@ version = "0.7.0"
 
 [[deps.ScopedValues]]
 deps = ["HashArrayMappedTries", "Logging"]
-git-tree-sha1 = "e3b5e4ccb1702db2ae9ac2a660d4b6b2a8595742"
+git-tree-sha1 = "c27d546a4749c81f70d1fabd604da6aa5054e3d2"
 uuid = "7e506255-f358-4e82-b7e4-beb19740aa63"
-version = "1.1.0"
+version = "1.2.0"
 
 [[deps.Serialization]]
 uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
@@ -189,9 +189,9 @@ uuid = "6462fe0b-24de-5631-8697-dd941f90decc"
 
 [[deps.SortingAlgorithms]]
 deps = ["DataStructures"]
-git-tree-sha1 = "c60ec5c62180f27efea3ba2908480f8055e17cee"
+git-tree-sha1 = "66e0a8e672a0bdfca2c3f5937efb8538b9ddc085"
 uuid = "a2af1166-a08f-5f64-846c-94a0d3cef48c"
-version = "1.1.1"
+version = "1.2.1"
 
 [[deps.SparseArrays]]
 deps = ["LinearAlgebra", "Random"]

Diff for: Project.toml

@@ -24,7 +24,7 @@ UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 [compat]
 DataStructures = "0.18"
 MacroTools = "0.5"
-MemPool = "0.4.5"
+MemPool = "0.4.6"
 PrecompileTools = "1.2"
 Requires = "1"
 ScopedValues = "1.1"

Diff for: src/precompile.jl

@@ -3,25 +3,47 @@
     add_processor_callback!("__cpu_thread_1__") do
         ThreadProc(1, 1)
     end
-    t1 = @spawn 1+1
+    # FIXME: t1 = @spawn 1+1
+    t1 = spawn(+, 1, 1)
+    fetch(t1)
     t2 = spawn(+, 1, t1)
     fetch(t2)
-    spawn() do
-        Sch.halt!(sch_handle())
+
+    # Clean up refs
+    t1 = nothing; t2 = nothing
+    state = Sch.EAGER_STATE[]
+    for i in 1:5
+        length(state.thunk_dict) == 1 && break
+        GC.gc()
+        yield()
     end
+    @assert length(state.thunk_dict) == 1
+
+    # Halt scheduler
+    notify(state.halt)
+    put!(state.chan, (1, nothing, nothing, (Sch.SchedulerHaltedException(), nothing)))
+    state = nothing
+
+    # Wait for halt
     while Sch.EAGER_INIT[]
         sleep(0.1)
     end
+
+    # Final clean-up
     Sch.EAGER_CONTEXT[] = nothing
-    GC.gc()
-    yield()
+    GC.gc(); yield()
     lock(Sch.ERRORMONITOR_TRACKED) do tracked
-        if all(t->istaskdone(t) || istaskfailed(t), tracked)
+        if all(t->istaskdone(t) || istaskfailed(t), map(last, tracked))
             empty!(tracked)
             return
         end
-        for t in tracked
-            Base.throwto(t, InterruptException())
+        for (name, t) in tracked
+            @warn "Waiting on $name"
+            if t.state == :runnable
+                Base.throwto(t, InterruptException())
+            else
+                wait(t)
+            end
         end
     end
     MemPool.exit_hook()

Diff for: src/sch/Sch.jl

@@ -72,6 +72,7 @@ Fields:
 - `lock::ReentrantLock` - Lock around operations which modify the state
 - `futures::Dict{Thunk, Vector{ThunkFuture}}` - Futures registered for waiting on the result of a thunk.
 - `errored::WeakKeyDict{Thunk,Bool}` - Indicates if a thunk's result is an error.
+- `thunks_to_delete::Set{Thunk}` - The list of `Thunk`s ready to be deleted upon completion.
 - `chan::RemoteChannel{Channel{Any}}` - Channel for receiving completed thunks.
 """
 struct ComputeState
@@ -98,6 +99,7 @@ struct ComputeState
     lock::ReentrantLock
     futures::Dict{Thunk, Vector{ThunkFuture}}
     errored::WeakKeyDict{Thunk,Bool}
+    thunks_to_delete::Set{Thunk}
     chan::RemoteChannel{Channel{Any}}
 end
 
@@ -127,6 +129,7 @@ function start_state(deps::Dict, node_order, chan)
                          ReentrantLock(),
                          Dict{Thunk, Vector{ThunkFuture}}(),
                          WeakKeyDict{Thunk,Bool}(),
+                         Set{Thunk}(),
                          chan)
 
     for k in sort(collect(keys(deps)), by=node_order)
@@ -366,7 +369,7 @@ function init_proc(state, p, log_sink)
                         end
                     end
                 end
-                errormonitor_tracked(t)
+                errormonitor_tracked("worker monitor $wid", t)
                 WORKER_MONITOR_TASKS[wid] = t
                 WORKER_MONITOR_CHANS[wid] = Dict{UInt64,RemoteChannel}()
             end
@@ -590,6 +593,8 @@ function scheduler_run(ctx, state::ComputeState, d::Thunk, options)
             timespan_start(ctx, :finish, thunk_id, (;thunk_id))
             finish_task!(ctx, state, node, thunk_failed)
             timespan_finish(ctx, :finish, thunk_id, (;thunk_id))
+
+            delete_unused_tasks!(state)
         end
 
         safepoint(state)
@@ -931,6 +936,39 @@ function finish_task!(ctx, state, node, thunk_failed)
     evict_all_chunks!(ctx, to_evict)
 end
 
+function delete_unused_tasks!(state)
+    to_delete = Thunk[]
+    for thunk in state.thunks_to_delete
+        if task_unused(state, thunk)
+            # Finished and nobody waiting on us, we can be deleted
+            push!(to_delete, thunk)
+        end
+    end
+    for thunk in to_delete
+        # Delete all cached data
+        task_delete!(state, thunk)
+
+        pop!(state.thunks_to_delete, thunk)
+    end
+end
+function delete_unused_task!(state, thunk)
+    if task_unused(state, thunk)
+        # Will not be accessed further, delete all cached data
+        task_delete!(state, thunk)
+        return true
+    else
+        return false
+    end
+end
+task_unused(state, thunk) =
+    haskey(state.cache, thunk) && !haskey(state.waiting_data, thunk)
+function task_delete!(state, thunk)
+    delete!(state.cache, thunk)
+    delete!(state.errored, thunk)
+    delete!(state.valid, thunk)
+    delete!(state.thunk_dict, thunk.id)
+end
+
 function evict_all_chunks!(ctx, to_evict)
     if !isempty(to_evict)
         @sync for w in map(p->p.pid, procs_to_use(ctx))
@@ -1290,7 +1328,7 @@ function start_processor_runner!(istate::ProcessorInternalState, uid::UInt64, re
                 else
                     t.sticky = false
                 end
-                tasks[thunk_id] = errormonitor_tracked(schedule(t))
+                tasks[thunk_id] = errormonitor_tracked("thunk $thunk_id", schedule(t))
                 proc_occupancy[] += task_occupancy
                 time_pressure[] += time_util
             end
@@ -1302,7 +1340,7 @@ function start_processor_runner!(istate::ProcessorInternalState, uid::UInt64, re
     else
         proc_run_task.sticky = false
     end
-    return errormonitor_tracked(schedule(proc_run_task))
+    return errormonitor_tracked("processor $to_proc", schedule(proc_run_task))
 end
 
 """

Diff for: src/sch/dynamic.jl

@@ -87,8 +87,8 @@ function dynamic_listener!(ctx, state, wid)
             end
         end
     end
-    errormonitor_tracked(listener_task)
-    errormonitor_tracked(@async begin
+    errormonitor_tracked("dynamic_listener! $wid", listener_task)
+    errormonitor_tracked("dynamic_listener! (halt+throw) $wid", @async begin
         wait(state.halt)
         # TODO: Not sure why we need the @async here, but otherwise we
         # don't stop all the listener tasks

Diff for: src/sch/eager.jl

@@ -20,7 +20,7 @@ function init_eager()
         return
     end
     ctx = eager_context()
-    errormonitor_tracked(Threads.@spawn try
+    errormonitor_tracked("eager compute()", Threads.@spawn try
         sopts = SchedulerOptions(;allow_errors=true)
         opts = Dagger.Options((;scope=Dagger.ExactScope(Dagger.ThreadProc(1, 1)),
                                 occupancy=Dict(Dagger.ThreadProc=>0)))
@@ -101,7 +101,7 @@ function thunk_yield(f)
 end
 
 eager_cleanup(t::Dagger.EagerThunkFinalizer) =
-    errormonitor_tracked(Threads.@spawn eager_cleanup(EAGER_STATE[], t.uid))
+    errormonitor_tracked("eager_cleanup $(t.uid)", Threads.@spawn eager_cleanup(EAGER_STATE[], t.uid))
 function eager_cleanup(state, uid)
     tid = nothing
     lock(EAGER_ID_MAP) do id_map

Diff for: src/sch/util.jl

@@ -1,15 +1,15 @@
 "Like `errormonitor`, but tracks how many outstanding tasks are running."
-function errormonitor_tracked(t::Task)
+function errormonitor_tracked(name::String, t::Task)
     errormonitor(t)
     @safe_lock_spin1 ERRORMONITOR_TRACKED tracked begin
-        push!(tracked, t)
+        push!(tracked, name => t)
     end
     errormonitor(Threads.@spawn begin
         try
             wait(t)
         finally
             lock(ERRORMONITOR_TRACKED) do tracked
-                idx = findfirst(o->o===t, tracked)
+                idx = findfirst(o->o[2]===t, tracked)
                 # N.B. This may be nothing if precompile emptied these
                 if idx !== nothing
                     deleteat!(tracked, idx)
@@ -18,7 +18,7 @@ function errormonitor_tracked(t::Task)
         end
     end)
 end
-const ERRORMONITOR_TRACKED = LockedObject(Task[])
+const ERRORMONITOR_TRACKED = LockedObject(Pair{String,Task}[])
 
 """
     unwrap_nested_exception(err::Exception) -> Bool

Diff for: src/submission.jl

@@ -89,7 +89,8 @@ function eager_submit_internal!(ctx, state, task, tid, payload; uid_to_tid=Dict{
         thunk = Thunk(f, args...; options...)
 
         # Create a `DRef` to `thunk` so that the caller can preserve it
-        thunk_ref = poolset(thunk; size=64, device=MemPool.CPURAMDevice())
+        thunk_ref = poolset(thunk; size=64, device=MemPool.CPURAMDevice(),
+                            destructor=UnrefThunkByUser(thunk))
         thunk_id = Sch.ThunkID(thunk.id, thunk_ref)
 
         # Attach `thunk` within the scheduler
@@ -122,6 +123,29 @@ function eager_submit_internal!(ctx, state, task, tid, payload; uid_to_tid=Dict{
         return thunk_id
     end
 end
+struct UnrefThunkByUser
+    thunk::Thunk
+end
+function (unref::UnrefThunkByUser)()
+    Sch.errormonitor_tracked("unref thunk $(unref.thunk.id)", Threads.@spawn begin
+        # This thunk is no longer referenced by the user, mark it as ready to be
+        # cleaned up as eagerly as possible (or do so now)
+        thunk = unref.thunk
+        state = Sch.EAGER_STATE[]
+        if state === nothing
+            return
+        end
+
+        @lock state.lock begin
+            if !Sch.delete_unused_task!(state, thunk)
+                # Register for deletion upon thunk completion
+                push!(state.thunks_to_delete, thunk)
+            end
+            # TODO: On success, walk down to children, as a fast-path
+        end
+    end)
+end
+
 
 # Local -> Remote
 function eager_submit!(ntasks, uid, future, finalizer_ref, f, args, options)