Sch: Minimize calls to reschedule_inputs

JuliaParallel · May 16, 2022 · c17b86d · c17b86d
1 parent 2f47217
commit c17b86d
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 1 deletion.
diff --git a/src/sch/Sch.jl b/src/sch/Sch.jl
@@ -47,6 +47,7 @@ Fields:
 - `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
 - `cache::WeakKeyDict{Thunk, Any}` - Maps from a finished `Thunk` to it's cached result, often a DRef
+- `valid::WeakKeyDict{Thunk, Nothing}` - Tracks all `Thunk`s that are in a valid scheduling state
 - `running::Set{Thunk}` - The set of currently-running `Thunk`s
 - `running_on::Dict{Thunk,OSProc}` - Map from `Thunk` to the OS process executing it
 - `thunk_dict::Dict{Int, WeakThunk}` - Maps from thunk IDs to a `Thunk`
@@ -70,6 +71,7 @@ struct ComputeState
     waiting_data::Dict{Union{Thunk,Chunk},Set{Thunk}}
     ready::Vector{Thunk}
     cache::WeakKeyDict{Thunk, Any}
+    valid::WeakKeyDict{Thunk, Nothing}
     running::Set{Thunk}
     running_on::Dict{Thunk,OSProc}
     thunk_dict::Dict{Int, WeakThunk}
@@ -93,7 +95,8 @@ function start_state(deps::Dict, node_order, chan)
                          OneToMany(),
                          deps,
                          Vector{Thunk}(undef, 0),
-                         Dict{Thunk, Any}(),
+                         WeakKeyDict{Thunk, Any}(),
+                         WeakKeyDict{Thunk, Nothing}(),
                          Set{Thunk}(),
                          Dict{Thunk,OSProc}(),
                          Dict{Int, WeakThunk}(),
@@ -119,6 +122,7 @@ function start_state(deps::Dict, node_order, chan)
             else
                 state.waiting[k] = waiting
             end
+            state.valid[k] = nothing
         end
     end
     state

diff --git a/src/sch/dynamic.jl b/src/sch/dynamic.jl
@@ -208,6 +208,7 @@ function _add_thunk!(ctx, state, task, tid, (f, args, kwargs, future, ref))
             # Preserve the `EagerThunkFinalizer` through `thunk`
             thunk.eager_ref = ref
         end
+        state.valid[thunk] = nothing
         put!(state.chan, RescheduleSignal())
         timespan_finish(ctx, :add_thunk, tid, 0)
         return thunk_id

diff --git a/src/sch/util.jl b/src/sch/util.jl
@@ -101,6 +101,9 @@ function reschedule_inputs!(state, thunk, seen=Set{Thunk}())
     while !isempty(to_visit)
         thunk = pop!(to_visit)
         push!(seen, thunk)
+        if haskey(state.valid, thunk)
+            continue
+        end
         if haskey(state.cache, thunk) || (thunk in state.ready) || (thunk in state.running)
             continue
         end