Sch: Fix WeakChunk usage

Construct WeakChunk within scheduler
Add anti-serialization method to WeakChunk
Add Chunk type to Chunk hashing method
Track DRef owner and ID in WeakChunk struct
Preserve old arguments in scheduler submission
Check for matching Chunk in state.waiting_data

Author: jpsamaroo

src/chunks.jl

@@ -65,7 +65,7 @@ affinity(c::Chunk) = affinity(c.handle)
 is_task_or_chunk(c::Chunk) = true
 
 Base.:(==)(c1::Chunk, c2::Chunk) = c1.handle == c2.handle
-Base.hash(c::Chunk, x::UInt64) = hash(c.handle, x)
+Base.hash(c::Chunk, x::UInt64) = hash(c.handle, hash(Chunk, x))
 
 collect_remote(chunk::Chunk) =
     move(chunk.processor, OSProc(), poolget(chunk.handle))
@@ -281,16 +281,22 @@ function savechunk(data, dir, f)
 end
 
 struct WeakChunk
+    wid::Int
+    id::Int
     x::WeakRef
+    function WeakChunk(c::Chunk)
+        return new(c.handle.owner, c.handle.id, WeakRef(c))
+    end
 end
-WeakChunk(c::Chunk) = WeakChunk(WeakRef(c))
 unwrap_weak(c::WeakChunk) = c.x.value
 function unwrap_weak_checked(c::WeakChunk)
-    c = unwrap_weak(c)
-    @assert c !== nothing
-    return c
+    cw = unwrap_weak(c)
+    @assert cw !== nothing "WeakChunk expired: ($(c.wid), $(c.id))"
+    return cw
 end
 is_task_or_chunk(c::WeakChunk) = true
+Serialization.serialize(io::AbstractSerializer, wc::WeakChunk) =
+    error("Cannot serialize a WeakChunk")
 
 Base.@deprecate_binding AbstractPart Union{Chunk, Thunk}
 Base.@deprecate_binding Part Chunk

src/sch/util.jl

@@ -107,6 +107,18 @@ function reschedule_syncdeps!(state, thunk, seen=Set{Thunk}())
         if haskey(state.cache, thunk) || (thunk in state.ready) || (thunk in state.running)
             continue
         end
+        for (_,input) in thunk.inputs
+            if input isa WeakChunk
+                input = unwrap_weak_checked(input)
+            end
+            if input isa 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
+                if !haskey(state.waiting_data, input)
+                    push!(get!(()->Set{Thunk}(), state.waiting_data, input), thunk)
+                end
+            end
+        end
         w = get!(()->Set{Thunk}(), state.waiting, thunk)
         for input in thunk.syncdeps
             input = unwrap_weak_checked(input)

src/submission.jl

@@ -27,6 +27,7 @@ function eager_submit_internal!(ctx, state, task, tid, payload; uid_to_tid=Dict{
     timespan_start(ctx, :add_thunk, tid, 0)
 
     # Lookup EagerThunk/ThunkID -> Thunk
+    old_args = copy(args)
     args::Vector{Any}
     syncdeps = if haskey(options, :syncdeps)
         collect(options.syncdeps)
@@ -47,6 +48,13 @@ function eager_submit_internal!(ctx, state, task, tid, payload; uid_to_tid=Dict{
             elseif arg isa Sch.ThunkID
                 arg_tid = arg.id
                 state.thunk_dict[arg_tid]
+            elseif arg isa 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
+                if haskey(state.waiting_data, arg)
+                    arg = only(filter(o->o isa Chunk && o.handle == arg.handle, keys(state.waiting_data)))::Chunk
+                end
+                WeakChunk(arg)
             else
                 arg
             end
@@ -76,7 +84,7 @@ function eager_submit_internal!(ctx, state, task, tid, payload; uid_to_tid=Dict{
         options = merge(options, (;syncdeps))
     end
 
-    GC.@preserve args begin
+    GC.@preserve old_args args begin
         # Create the `Thunk`
         thunk = Thunk(f, args...; options...)
 
@@ -146,8 +154,6 @@ function eager_process_elem_submission_to_local(id_map, x)
     @assert !isa(x, Thunk) "Cannot use `Thunk`s in `@spawn`/`spawn`"
     if x isa Dagger.EagerThunk && haskey(id_map, x.uid)
         return Sch.ThunkID(id_map[x.uid], x.thunk_ref)
-    elseif x isa Dagger.Chunk
-        return WeakChunk(x)
     else
         return x
     end