Merge pull request #223 from JuliaParallel/jps/eager-mem-leak-fix

Fix memory leaks in eager API

Author: jpsamaroo

Diff for: Project.toml

@@ -1,6 +1,6 @@
 name = "Dagger"
 uuid = "d58978e5-989f-55fb-8d15-ea34adc7bf54"
-version = "0.11.5"
+version = "0.11.6"
 
 [deps]
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"

Diff for: src/Dagger.jl

@@ -2,6 +2,8 @@ module Dagger
 
 using Distributed, SharedArrays
 
+using MemPool
+
 import Base: collect, adjoint, reduce
 import Distributed: procs
 

Diff for: src/chunks.jl

@@ -1,5 +1,3 @@
-
-using MemPool
 using Serialization
 
 export domain, UnitDomain, project, alignfirst, ArrayDomain

Diff for: src/sch/Sch.jl

@@ -4,7 +4,7 @@ using Distributed
 import MemPool: DRef
 
 import ..Dagger
-import ..Dagger: Context, Processor, Thunk, ThunkFuture, ThunkFailedException, Chunk, OSProc, order, free!, dependents, noffspring, istask, inputs, affinity, tochunk, @dbg, @logmsg, timespan_start, timespan_end, unrelease, procs, move, capacity, chunktype, default_enabled, get_processors, execute!, rmprocs!, addprocs!, thunk_processor
+import ..Dagger: Context, Processor, Thunk, ThunkFuture, ThunkFailedException, Chunk, OSProc, order, free!, dependents, noffspring, istask, inputs, unwrap_weak, affinity, tochunk, @dbg, @logmsg, timespan_start, timespan_end, unrelease, procs, move, capacity, chunktype, default_enabled, get_processors, execute!, rmprocs!, addprocs!, thunk_processor
 
 const OneToMany = Dict{Thunk, Set{Thunk}}
 
@@ -44,8 +44,6 @@ The internal state-holding struct of the scheduler.
 
 Fields:
 - `uid::UInt64` - Unique identifier for this scheduler instance
-- `dependents::Dict{Union{Thunk,Chunk},Set{Thunk}}` - The result of calling `dependents` on the DAG
-- `finished::Set{Thunk}` - The set of completed `Thunk`s
 - `waiting::OneToMany` - Map from downstream `Thunk` to upstream `Thunk`s that still need to execute
 - `waiting_data::Dict{Union{Thunk,Chunk},Set{Thunk}}` - Map from input `Chunk`/upstream `Thunk` to all unfinished downstream `Thunk`s, to retain caches
 - `ready::Vector{Thunk}` - The list of `Thunk`s that are ready to execute
@@ -60,20 +58,18 @@ Fields:
 - `worker_chans::Dict{Int, Tuple{RemoteChannel,RemoteChannel}}` - Communication channels between the scheduler and each worker
 - `procs_cache_list::Base.RefValue{Union{ProcessorCacheEntry,Nothing}}` - Cached linked list of processors ready to be used
 - `function_cost_cache::Dict{Type{<:Tuple},UInt}` - Cache of estimated CPU time required to compute the given signature
-- `halt::Base.RefValue{Bool}` - Flag indicating, when set, that the scheduler should halt immediately
-- `lock::ReentrantLock()` - Lock around operations which modify the state
+- `halt::Base.Event` - Event indicating that the scheduler is halting
+- `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::Set{Thunk}` - Thunks that threw an error
-- `chan::RemoteChannel{Channel{Any}}` - Channel for receiving completed thunks
+- `errored::WeakKeyDict{Thunk,Bool}` - Indicates if a thunk's result is due to an error.
+- `chan::RemoteChannel{Channel{Any}}` - Channel for receiving completed thunks.
 """
 struct ComputeState
     uid::UInt64
-    dependents::Dict{Union{Thunk,Chunk},Set{Thunk}}
-    finished::Set{Thunk}
     waiting::OneToMany
     waiting_data::Dict{Union{Thunk,Chunk},Set{Thunk}}
     ready::Vector{Thunk}
-    cache::Dict{Thunk, Any}
+    cache::WeakKeyDict{Thunk, Any}
     running::Set{Thunk}
     running_on::Dict{Thunk,OSProc}
     thunk_dict::Dict{Int, Any}
@@ -84,19 +80,17 @@ struct ComputeState
     worker_chans::Dict{Int, Tuple{RemoteChannel,RemoteChannel}}
     procs_cache_list::Base.RefValue{Union{ProcessorCacheEntry,Nothing}}
     function_cost_cache::Dict{Type{<:Tuple},UInt}
-    halt::Base.RefValue{Bool}
+    halt::Base.Event
     lock::ReentrantLock
     futures::Dict{Thunk, Vector{ThunkFuture}}
-    errored::Set{Thunk}
+    errored::WeakKeyDict{Thunk,Bool}
     chan::RemoteChannel{Channel{Any}}
 end
 
 function start_state(deps::Dict, node_order, chan)
     state = ComputeState(rand(UInt64),
-                         deps,
-                         Set{Thunk}(),
                          OneToMany(),
-                         Dict{Union{Thunk,Chunk},Set{Thunk}}(),
+                         deps,
                          Vector{Thunk}(undef, 0),
                          Dict{Thunk, Any}(),
                          Set{Thunk}(),
@@ -109,18 +103,13 @@ function start_state(deps::Dict, node_order, chan)
                          Dict{Int, Tuple{RemoteChannel,RemoteChannel}}(),
                          Ref{Union{ProcessorCacheEntry,Nothing}}(nothing),
                          Dict{Type{<:Tuple},UInt}(),
-                         Ref{Bool}(false),
+                         Base.Event(),
                          ReentrantLock(),
                          Dict{Thunk, Vector{ThunkFuture}}(),
-                         Set{Thunk}(),
+                         WeakKeyDict{Thunk,Bool}(),
                          chan)
 
-    nodes = sort(collect(keys(deps)), by=node_order)
-    # N.B. Using merge! here instead would modify deps
-    for (key,val) in deps
-        state.waiting_data[key] = copy(val)
-    end
-    for k in nodes
+    for k in sort(collect(keys(deps)), by=node_order)
         if istask(k)
             waiting = Set{Thunk}(Iterators.filter(istask, inputs(k)))
             if isempty(waiting)
@@ -392,7 +381,11 @@ function compute_dag(ctx, d::Thunk; options=SchedulerOptions())
         end
 
         isempty(state.running) && continue
-        pid, proc, thunk_id, (res, metadata) = take!(chan) # get result of completed thunk
+        chan_value = take!(chan) # get result of completed thunk
+        if chan_value isa RescheduleSignal
+            continue
+        end
+        pid, proc, thunk_id, (res, metadata) = chan_value
         gproc = OSProc(pid)
         lock(newtasks_lock) # This waits for any assign_new_procs! above to complete and then shuts down the task
         safepoint(state)
@@ -424,6 +417,7 @@ function compute_dag(ctx, d::Thunk; options=SchedulerOptions())
             state.function_cost_cache[sig] = (metadata.threadtime + get(state.function_cost_cache, sig, 0)) ÷ 2
         end
         state.cache[node] = res
+        state.errored[node] = thunk_failed
         if node.options !== nothing && node.options.checkpoint !== nothing
             try
                 node.options.checkpoint(node, res)
@@ -441,12 +435,14 @@ function compute_dag(ctx, d::Thunk; options=SchedulerOptions())
 
         safepoint(state)
     end
-    state.halt[] = true
+    @assert !isready(state.chan)
+    close(state.chan)
+    notify(state.halt)
     @sync for p in procs_to_use(ctx)
         @async cleanup_proc(state, p)
     end
     value = state.cache[d] # TODO: move(OSProc(), state.cache[d])
-    if d in state.errored
+    if state.errored[d]
         throw(value)
     end
     if options.checkpoint !== nothing
@@ -556,6 +552,7 @@ function schedule!(ctx, state, procs=procs_to_use(ctx))
         while !isempty(state.ready)
             # Select a new task and get its options
             task = pop!(state.ready)
+            @assert !haskey(state.cache, task)
             opts = merge(ctx.options, task.options)
             sig = signature(task, state)
 
@@ -693,62 +690,69 @@ function pop_with_affinity!(ctx, tasks, proc)
     return nothing
 end
 
-function finish_task!(ctx, state, node, thunk_failed; free=true)
+function finish_task!(ctx, state, node, thunk_failed)
     pop!(state.running, node)
     delete!(state.running_on, node)
-    if !thunk_failed
-        push!(state.finished, node)
-    else
+    if thunk_failed
         set_failed!(state, node)
     end
-    if istask(node) && node.cache
+    if node.cache
         node.cache_ref = state.cache[node]
     end
-    if !thunk_failed
-        for dep in sort!(collect(state.dependents[node]), by=state.node_order)
+    for dep in sort!(collect(get(()->Set{Thunk}(), state.waiting_data, node)), by=state.node_order)
+        dep_isready = false
+        if haskey(state.waiting, dep)
             set = state.waiting[dep]
             node in set && pop!(set, node)
-            if isempty(set)
-                pop!(state.waiting, dep)
-                push!(state.ready, dep)
+            dep_isready = isempty(set)
+            if dep_isready
+                delete!(state.waiting, dep)
             end
-            # todo: free data
+        else
+            dep_isready = true
         end
-        if haskey(state.futures, node)
-            # Notify any listening thunks
-            for future in state.futures[node]
-                if istask(node) && haskey(state.cache, node)
-                    put!(future, state.cache[node])
-                else
-                    put!(future, nothing)
-                end
+        if dep_isready
+            if !thunk_failed
+                push!(state.ready, dep)
             end
-            delete!(state.futures, node)
         end
     end
+    if haskey(state.futures, node)
+        # Notify any listening thunks
+        for future in state.futures[node]
+            put!(future, state.cache[node]; error=thunk_failed)
+        end
+        delete!(state.futures, node)
+    end
 
     # Chunk clean-up
     to_evict = Set{Chunk}()
-    for inp in filter(t->istask(t) || (t isa Chunk), inputs(node))
+    for inp in filter(t->istask(t) || (t isa Chunk), unwrap_weak.(node.inputs))
         if inp in keys(state.waiting_data)
-            s = state.waiting_data[inp]
-            if node in s
-                pop!(s, node)
+            w = state.waiting_data[inp]
+            if node in w
+                pop!(w, node)
             end
-            if free && isempty(s)
+            if isempty(w)
+                delete!(state.waiting_data, inp)
                 if istask(inp) && haskey(state.cache, inp)
                     _node = state.cache[inp]
                     if _node isa Chunk
                         push!(to_evict, _node)
                     end
-                    free!(_node, force=false, cache=(istask(inp) && inp.cache))
-                    pop!(state.cache, inp)
+                    GC.@preserve inp begin
+                        pop!(state.cache, inp)
+                        pop!(state.errored, inp)
+                    end
                 elseif inp isa Chunk
                     push!(to_evict, inp)
                 end
             end
         end
     end
+    if haskey(state.waiting_data, node) && isempty(state.waiting_data[node])
+        delete!(state.waiting_data, node)
+    end
     if !isempty(to_evict)
         @sync for w in map(p->p.pid, procs_to_use(ctx))
             @async remote_do(evict_chunks!, w, to_evict)
@@ -779,8 +783,8 @@ function fire_tasks!(ctx, thunks::Vector{<:Tuple}, (gproc, proc), state)
             if data !== nothing
                 # cache hit
                 state.cache[thunk] = data
-                thunk_failed = thunk in state.errored
-                finish_task!(ctx, state, thunk, thunk_failed; free=false)
+                thunk_failed = state.errored[thunk]
+                finish_task!(ctx, state, thunk, thunk_failed)
                 continue
             else
                 # cache miss
@@ -791,19 +795,20 @@ function fire_tasks!(ctx, thunks::Vector{<:Tuple}, (gproc, proc), state)
             try
                 result = thunk.options.restore(thunk)
                 state.cache[thunk] = result
-                finish_task!(ctx, state, thunk, false; free=false)
+                state.errored[thunk] = false
+                finish_task!(ctx, state, thunk, false)
                 continue
             catch err
                 @error "Thunk restore failed" exception=(err,catch_backtrace())
             end
         end
 
         ids = map(enumerate(thunk.inputs)) do (idx,x)
-            istask(x) ? x.id : -idx
+            istask(x) ? unwrap_weak(x).id : -idx
         end
 
         data = map(thunk.inputs) do x
-            istask(x) ? state.cache[x] : x
+            istask(x) ? state.cache[unwrap_weak(x)] : x
         end
         toptions = thunk.options !== nothing ? thunk.options : ThunkOptions()
         options = merge(ctx.options, toptions)
@@ -829,7 +834,8 @@ function do_tasks(to_proc, chan, tasks)
     for task in tasks
         @async begin
             try
-                put!(chan, (myid(), to_proc, task[2], do_task(to_proc, task...)))
+                result = do_task(to_proc, task...)
+                put!(chan, (myid(), to_proc, task[2], result))
             catch ex
                 bt = catch_backtrace()
                 put!(chan, (myid(), to_proc, task[2], (CapturedException(ex, bt), nothing)))