Eagerly free Thunk cached results when unneeded

Uses the new MemPool destructor functionality to eagerly clean up Thunk
cached results once no more EagerThunk handles exist for the Thunk. Will
either free the cache immediately, or attempt to do so upon finishing
each Thunk.

Author: jpsamaroo

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"]

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"

src/precompile.jl

@@ -3,20 +3,35 @@
     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)
 
-    # Shutdown scheduler and clean up
-    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), map(last, tracked))
             empty!(tracked)
@@ -26,6 +41,8 @@
             @warn "Waiting on $name"
             if t.state == :runnable
                 Base.throwto(t, InterruptException())
+            else
+                wait(t)
             end
         end
     end

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)
@@ -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))

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)