Sch: Track unique Chunk per DRef

Author: jpsamaroo

src/sch/Sch.jl

@@ -58,6 +58,7 @@ Fields:
 - `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`
 - `node_order::Any` - Function that returns the order of a thunk
+- `equiv_chunks::WeakKeyDict{DRef,Chunk}` - Cache mapping from `DRef` to a `Chunk` which contains it
 - `worker_time_pressure::Dict{Int,Dict{Processor,UInt64}}` - Maps from worker ID to processor pressure
 - `worker_storage_pressure::Dict{Int,Dict{Union{StorageResource,Nothing},UInt64}}` - Maps from worker ID to storage resource pressure
 - `worker_storage_capacity::Dict{Int,Dict{Union{StorageResource,Nothing},UInt64}}` - Maps from worker ID to storage resource capacity
@@ -84,6 +85,7 @@ struct ComputeState
     running_on::Dict{Thunk,OSProc}
     thunk_dict::Dict{Int, WeakThunk}
     node_order::Any
+    equiv_chunks::WeakKeyDict{DRef,Chunk}
     worker_time_pressure::Dict{Int,Dict{Processor,UInt64}}
     worker_storage_pressure::Dict{Int,Dict{Union{StorageResource,Nothing},UInt64}}
     worker_storage_capacity::Dict{Int,Dict{Union{StorageResource,Nothing},UInt64}}
@@ -113,6 +115,7 @@ function start_state(deps::Dict, node_order, chan)
                          Dict{Thunk,OSProc}(),
                          Dict{Int, WeakThunk}(),
                          node_order,
+                         WeakKeyDict{DRef,Chunk}(),
                          Dict{Int,Dict{Processor,UInt64}}(),
                          Dict{Int,Dict{Union{StorageResource,Nothing},UInt64}}(),
                          Dict{Int,Dict{Union{StorageResource,Nothing},UInt64}}(),
@@ -428,6 +431,11 @@ function scheduler_run(ctx, state::ComputeState, d::Thunk, options::SchedulerOpt
                     state.transfer_rate[] = (state.transfer_rate[] + metadata.transfer_rate) ÷ 2
                 end
             end
+            if res isa Chunk
+                if !haskey(state.equiv_chunks, res)
+                    state.equiv_chunks[res.handle::DRef] = res
+                end
+            end
             store_result!(state, node, res; error=thunk_failed)
             if node.options !== nothing && node.options.checkpoint !== nothing
                 try

src/submission.jl

@@ -101,18 +101,17 @@ function eager_submit_internal!(ctx, state, task, tid, payload::AnyPayload; uid_
                 elseif valuetype(arg) <: Chunk
                     # N.B. Different Chunks with the same DRef handle will hash to the same slot,
                     # so we just pick an equivalent Chunk as our upstream
+                    # FIXME: It's super slow to iterate over state.waiting_data, can we cache this?
                     chunk = value(arg)::Chunk
-                    if haskey(state.waiting_data, chunk)
-                        newchunk = nothing
-                        for other in keys(state.waiting_data)
-                            if other isa Chunk && other.handle == chunk.handle
-                                newchunk = other
-                                break
-                            end
+                    function find_equivalent_chunk(state, chunk::C) where {C<:Chunk}
+                        if haskey(state.equiv_chunks, chunk.handle)
+                            return state.equiv_chunks[chunk.handle]::C
+                        else
+                            state.equiv_chunks[chunk.handle] = chunk
+                            return chunk
                         end
-                        @assert newchunk !== nothing
-                        chunk = newchunk::Chunk
                     end
+                    chunk = find_equivalent_chunk(state, chunk)
                     #=FIXME:UNIQUE=#
                     @inbounds fargs[idx] = Argument(arg.pos, WeakChunk(chunk))
                 end
@@ -193,6 +192,7 @@ function eager_submit_internal!(ctx, state, task, tid, payload::AnyPayload; uid_
 
             return thunk_id
         end
+        empty!(equiv_chunks)
     end
 end
 struct UnrefThunk
@@ -238,16 +238,12 @@ function eager_submit!(payload::AnyPayload)
         return remotecall_fetch(1, payload) do payload
             Sch.init_eager()
             state = Dagger.Sch.EAGER_STATE[]
-            @lock state.lock begin
-                eager_submit_internal!(payload)
-            end
+            @lock state.lock eager_submit_internal!(payload)
         end
     else
         Sch.init_eager()
         state = Dagger.Sch.EAGER_STATE[]
-        return lock(state.lock) do
-            eager_submit_internal!(payload)
-        end
+        return @lock state.lock eager_submit_internal!(payload)
     end
 end