Initial support for robustly migrating streaming tasks

src/Dagger.jl

@@ -25,6 +25,8 @@ else
 end
 import TaskLocalValues: TaskLocalValue
 
+import TaskLocalValues: TaskLocalValue
+
 if !isdefined(Base, :get_extension)
     import Requires: @require
 end
@@ -47,16 +49,16 @@ include("processor.jl")
 include("threadproc.jl")
 include("context.jl")
 include("utils/processors.jl")
+include("dtask.jl")
+include("cancellation.jl")
 include("task-tls.jl")
 include("scopes.jl")
 include("utils/scopes.jl")
-include("dtask.jl")
 include("queue.jl")
 include("thunk.jl")
 include("submission.jl")
 include("chunks.jl")
 include("memory-spaces.jl")
-include("cancellation.jl")
 
 # Task scheduling
 include("compute.jl")
@@ -69,9 +71,9 @@ include("sch/Sch.jl"); using .Sch
 include("datadeps.jl")
 
 # Streaming
-include("stream-buffers.jl")
-include("stream-fetchers.jl")
 include("stream.jl")
+include("stream-buffers.jl")
+include("stream-transfer.jl")
 
 # Array computations
 include("array/darray.jl")
@@ -152,6 +154,20 @@ function __init__()
             ThreadProc(myid(), tid)
         end
     end
+
+    # Set up @dagdebug categories, if specified
+    try
+        if haskey(ENV, "JULIA_DAGGER_DEBUG")
+            empty!(DAGDEBUG_CATEGORIES)
+            for category in split(ENV["JULIA_DAGGER_DEBUG"], ",")
+                if category != ""
+                    push!(DAGDEBUG_CATEGORIES, Symbol(category))
+                end
+            end
+        end
+    catch err
+        @warn "Error parsing JULIA_DAGGER_DEBUG" exception=err
+    end
 end
 
 end # module

src/array/indexing.jl

@@ -1,5 +1,3 @@
-import TaskLocalValues: TaskLocalValue
-
 ### getindex
 
 struct GetIndex{T,N} <: ArrayOp{T,N}

src/cancellation.jl

@@ -1,3 +1,51 @@
+# DTask-level cancellation
+
+mutable struct CancelToken
+    @atomic cancelled::Bool
+    @atomic graceful::Bool
+    event::Base.Event
+end
+CancelToken() = CancelToken(false, false, Base.Event())
+function cancel!(token::CancelToken; graceful::Bool=true)
+    if !graceful
+        @atomic token.graceful = false
+    end
+    @atomic token.cancelled = true
+    notify(token.event)
+    return
+end
+function is_cancelled(token::CancelToken; must_force::Bool=false)
+    if token.cancelled[]
+        if must_force && token.graceful[]
+            # If we're only responding to forced cancellation, ignore graceful cancellations
+            return false
+        end
+        return true
+    end
+    return false
+end
+Base.wait(token::CancelToken) = wait(token.event)
+# TODO: Enable this for safety
+#Serialization.serialize(io::AbstractSerializer, ::CancelToken) =
+#    throw(ConcurrencyViolationError("Cannot serialize a CancelToken"))
+
+const DTASK_CANCEL_TOKEN = TaskLocalValue{Union{CancelToken,Nothing}}(()->nothing)
+
+function clone_cancel_token_remote(orig_token::CancelToken, wid::Integer)
+    remote_token = remotecall_fetch(wid) do
+        return poolset(CancelToken())
+    end
+    errormonitor_tracked("remote cancel_token communicator", Threads.@spawn begin
+        wait(orig_token)
+        @dagdebug nothing :cancel "Cancelling remote token on worker $wid"
+        MemPool.access_ref(remote_token) do remote_token
+            cancel!(remote_token)
+        end
+    end)
+end
+
+# Global-level cancellation
+
 """
     cancel!(task::DTask; force::Bool=false, halt_sch::Bool=false)
 
@@ -48,7 +96,7 @@ function _cancel!(state, tid, force, halt_sch)
     for task in state.ready
         tid !== nothing && task.id != tid && continue
         @dagdebug tid :cancel "Cancelling ready task"
-        state.cache[task] = InterruptException()
+        state.cache[task] = DTaskFailedException(task, task, InterruptException())
         state.errored[task] = true
         Sch.set_failed!(state, task)
     end
@@ -58,7 +106,7 @@ function _cancel!(state, tid, force, halt_sch)
     for task in keys(state.waiting)
         tid !== nothing && task.id != tid && continue
         @dagdebug tid :cancel "Cancelling waiting task"
-        state.cache[task] = InterruptException()
+        state.cache[task] = DTaskFailedException(task, task, InterruptException())
         state.errored[task] = true
         Sch.set_failed!(state, task)
     end
@@ -80,11 +128,11 @@ function _cancel!(state, tid, force, halt_sch)
                             Tf === typeof(Sch.eager_thunk) && continue
                             istaskdone(task) && continue
                             any_cancelled = true
-                            @dagdebug tid :cancel "Cancelling running task ($Tf)"
                             if force
                                 @dagdebug tid :cancel "Interrupting running task ($Tf)"
                                 Threads.@spawn Base.throwto(task, InterruptException())
                             else
+                                @dagdebug tid :cancel "Cancelling running task ($Tf)"
                                 # Tell the processor to just drop this task
                                 task_occupancy = task_spec[4]
                                 time_util = task_spec[2]
@@ -93,6 +141,7 @@ function _cancel!(state, tid, force, halt_sch)
                                 push!(istate.cancelled, tid)
                                 to_proc = istate.proc
                                 put!(istate.return_queue, (myid(), to_proc, tid, (InterruptException(), nothing)))
+                                cancel!(istate.cancel_tokens[tid]; graceful=false)
                             end
                         end
                     end

src/options.jl

@@ -20,6 +20,12 @@ function with_options(f, options::NamedTuple)
 end
 with_options(f; options...) = with_options(f, NamedTuple(options))
 
+function _without_options(f)
+    with(options_context => NamedTuple()) do
+        f()
+    end
+end
+
 """
     get_options(key::Symbol, default) -> Any
     get_options(key::Symbol) -> Any

src/sch/Sch.jl

@@ -1170,6 +1170,7 @@ struct ProcessorInternalState
     proc_occupancy::Base.RefValue{UInt32}
     time_pressure::Base.RefValue{UInt64}
     cancelled::Set{Int}
+    cancel_tokens::Dict{Int,Dagger.CancelToken}
     done::Base.RefValue{Bool}
 end
 struct ProcessorState
@@ -1189,7 +1190,7 @@ function proc_states(f::Base.Callable, uid::UInt64)
     end
 end
 proc_states(f::Base.Callable) =
-    proc_states(f, task_local_storage(:_dagger_sch_uid)::UInt64)
+    proc_states(f, Dagger.get_tls().sch_uid)
 
 task_tid_for_processor(::Processor) = nothing
 task_tid_for_processor(proc::Dagger.ThreadProc) = proc.tid
@@ -1318,7 +1319,14 @@ function start_processor_runner!(istate::ProcessorInternalState, uid::UInt64, re
 
             # Execute the task and return its result
             t = @task begin
+                # Set up cancellation
+                cancel_token = Dagger.CancelToken()
+                Dagger.DTASK_CANCEL_TOKEN[] = cancel_token
+                lock(istate.queue) do _
+                    istate.cancel_tokens[thunk_id] = cancel_token
+                end
                 was_cancelled = false
+
                 result = try
                     do_task(to_proc, task)
                 catch err
@@ -1335,6 +1343,7 @@ function start_processor_runner!(istate::ProcessorInternalState, uid::UInt64, re
                             # Task was cancelled, so occupancy and pressure are
                             # already reduced
                             pop!(istate.cancelled, thunk_id)
+                            delete!(istate.cancel_tokens, thunk_id)
                             was_cancelled = true
                         end
                     end
@@ -1352,6 +1361,9 @@ function start_processor_runner!(istate::ProcessorInternalState, uid::UInt64, re
                     else
                         rethrow(err)
                     end
+                finally
+                    # Ensure that any spawned tasks get cleaned up
+                    Dagger.cancel!(cancel_token)
                 end
             end
             lock(istate.queue) do _
@@ -1401,6 +1413,7 @@ function do_tasks(to_proc, return_queue, tasks)
                                             Dict{Int,Vector{Any}}(),
                                             Ref(UInt32(0)), Ref(UInt64(0)),
                                             Set{Int}(),
+                                            Dict{Int,Dagger.CancelToken}(),
                                             Ref(false))
             runner = start_processor_runner!(istate, uid, return_queue)
             @static if VERSION < v"1.9"
@@ -1640,6 +1653,7 @@ function do_task(to_proc, task_desc)
             sch_handle,
             processor=to_proc,
             task_spec=task_desc,
+            cancel_token=Dagger.DTASK_CANCEL_TOKEN[],
         ))
 
         res = Dagger.with_options(propagated) do

src/sch/dynamic.jl

@@ -17,7 +17,7 @@ struct SchedulerHandle
 end
 
 "Gets the scheduler handle for the currently-executing thunk."
-sch_handle() = task_local_storage(:_dagger_sch_handle)::SchedulerHandle
+sch_handle() = Dagger.get_tls().sch_handle::SchedulerHandle
 
 "Thrown when the scheduler halts before finishing processing the DAG."
 struct SchedulerHaltedException <: Exception end

src/sch/eager.jl

@@ -141,3 +141,6 @@ function _find_thunk(e::Dagger.DTask)
         unwrap_weak_checked(EAGER_STATE[].thunk_dict[tid])
     end
 end
+Dagger.task_id(t::Dagger.DTask) = lock(EAGER_ID_MAP) do id_map
+    id_map[t.uid]
+end

src/sch/util.jl

@@ -29,6 +29,8 @@ unwrap_nested_exception(err::CapturedException) =
     unwrap_nested_exception(err.ex)
 unwrap_nested_exception(err::RemoteException) =
     unwrap_nested_exception(err.captured)
+unwrap_nested_exception(err::DTaskFailedException) =
+    unwrap_nested_exception(err.ex)
 unwrap_nested_exception(err) = err
 
 "Gets a `NamedTuple` of options propagated by `thunk`."