Make Dagger.finish_stream() propagate downstream

.buildkite/pipeline.yml

@@ -6,6 +6,7 @@
     os: linux
     arch: x86_64
   command: "julia --project -e 'using Pkg; Pkg.develop(;path=\"lib/TimespanLogging\")'"
+
 .bench: &bench
   if: build.message =~ /\[run benchmarks\]/
   agents:
@@ -14,6 +15,7 @@
     os: linux
     arch: x86_64
     num_cpus: 16
+
 steps:
   - label: Julia 1.9
     timeout_in_minutes: 90
@@ -25,6 +27,7 @@ steps:
           julia_args: "--threads=1"
       - JuliaCI/julia-coverage#v1:
           codecov: true
+
   - label: Julia 1.10
     timeout_in_minutes: 90
     <<: *test
@@ -35,6 +38,7 @@ steps:
           julia_args: "--threads=1"
       - JuliaCI/julia-coverage#v1:
           codecov: true
+
   - label: Julia nightly
     timeout_in_minutes: 90
     <<: *test
@@ -77,6 +81,7 @@ steps:
       - JuliaCI/julia-coverage#v1:
           codecov: true
     command: "julia -e 'using Pkg; Pkg.develop(;path=pwd()); Pkg.develop(;path=\"lib/TimespanLogging\"); Pkg.develop(;path=\"lib/DaggerWebDash\"); include(\"lib/DaggerWebDash/test/runtests.jl\")'"
+
   - label: Benchmarks
     timeout_in_minutes: 120
     <<: *bench
@@ -93,6 +98,7 @@ steps:
       BENCHMARK_SCALE: "5:5:50"
     artifacts:
       - benchmarks/result*
+
   - label: DTables.jl stability test
     timeout_in_minutes: 20
     plugins:

Project.toml

@@ -52,7 +52,7 @@ GraphViz = "0.2"
 Graphs = "1"
 JSON3 = "1"
 MacroTools = "0.5"
-MemPool = "0.4.6"
+MemPool = "0.4.10"
 OnlineStats = "1"
 Plots = "1"
 PrecompileTools = "1.2"

docs/make.jl

@@ -22,6 +22,7 @@ makedocs(;
         "Task Spawning" => "task-spawning.md",
         "Data Management" => "data-management.md",
         "Distributed Arrays" => "darray.md",
+        "Streaming Tasks" => "streaming.md",
         "Scopes" => "scopes.md",
         "Processors" => "processors.md",
         "Task Queues" => "task-queues.md",

docs/src/streaming.md

@@ -0,0 +1,105 @@
+# Streaming Tasks
+
+Dagger tasks have a limited lifetime - they are created, execute, finish, and
+are eventually destroyed when they're no longer needed. Thus, if one wants
+to run the same kind of computations over and over, one might re-create a
+similar set of tasks for each unit of data that needs processing.
+
+This might be fine for computations which take a long time to run (thus
+dwarfing the cost of task creation, which is quite small), or when working with
+a limited set of data, but this approach is not great for doing lots of small
+computations on a large (or endless) amount of data. For example, processing
+image frames from a webcam, reacting to messages from a message bus, reading
+samples from a software radio, etc. All of these tasks are better suited to a
+"streaming" model of data processing, where data is simply piped into a
+continuously-running task (or DAG of tasks) forever, or until the data runs
+out.
+
+Thankfully, if you have a problem which is best modeled as a streaming system
+of tasks, Dagger has you covered! Building on its support for
+[Task Queues](@ref), Dagger provides a means to convert an entire DAG of
+tasks into a streaming DAG, where data flows into and out of each task
+asynchronously, using the `spawn_streaming` function:
+
+```julia
+Dagger.spawn_streaming() do # enters a streaming region
+  vals = Dagger.@spawn rand()
+  print_vals = Dagger.@spawn println(vals)
+end # exits the streaming region, and starts the DAG running
+```
+
+In the above example, `vals` is a Dagger task which has been transformed to run
+in a streaming manner - instead of just calling `rand()` once and returning its
+result, it will re-run `rand()` endlessly, continuously producing new random
+values. In typical Dagger style, `print_vals` is a Dagger task which depends on
+`vals`, but in streaming form - it will continuously `println` the random
+values produced from `vals`. Both tasks will run forever, and will run
+efficiently, only doing the work necessary to generate, transfer, and consume
+values.
+
+As the comments point out, `spawn_streaming` creates a streaming region, during
+which `vals` and `print_vals` are created and configured. Both tasks are halted
+until `spawn_streaming` returns, allowing large DAGs to be built all at once,
+without any task losing a single value. If desired, streaming regions can be
+connected, although some values might be lost while tasks are being connected:
+
+```julia
+vals = Dagger.spawn_streaming() do
+    Dagger.@spawn rand()
+end
+
+# Some values might be generated by `vals` but thrown away
+# before `print_vals` is fully setup and connected to it
+
+print_vals = Dagger.spawn_streaming() do
+    Dagger.@spawn println(vals)
+end
+```
+
+More complicated streaming DAGs can be easily constructed, without doing
+anything different. For example, we can generate multiple streams of random
+numbers, write them all to their own files, and print the combined results:
+
+```julia
+Dagger.spawn_streaming() do
+    all_vals = [Dagger.spawn(rand) for i in 1:4]
+    all_vals_written = map(1:4) do i
+        Dagger.spawn(all_vals[i]) do val
+            open("results_$i.txt"; write=true, create=true, append=true) do io
+                println(io, repr(val))
+            end
+            return val
+        end
+    end
+    Dagger.spawn(all_vals_written...) do all_vals_written...
+        vals_sum = sum(all_vals_written)
+        println(vals_sum)
+    end
+end
+```
+
+If you want to stop the streaming DAG and tear it all down, you can call
+`Dagger.kill!(all_vals[1])` (or `Dagger.kill!(all_vals_written[2])`, etc., the
+kill propagates throughout the DAG).
+
+Alternatively, tasks can stop themselves from the inside with
+`finish_streaming`, optionally returning a value that can be `fetch`'d. Let's
+do this when our randomly-drawn number falls within some arbitrary range:
+
+```julia
+vals = Dagger.spawn_streaming() do
+    Dagger.spawn() do
+        x = rand()
+        if x < 0.001
+            # That's good enough, let's be done
+            return Dagger.finish_streaming("Finished!")
+        end
+        return x
+    end
+end
+fetch(vals)
+```
+
+In this example, the call to `fetch` will hang (while random numbers continue
+to be drawn), until a drawn number is less than 0.001; at that point, `fetch`
+will return with "Finished!", and the task `vals` will have terminated.

src/Dagger.jl

@@ -23,6 +23,9 @@ if !isdefined(Base, :ScopedValues)
 else
     import Base.ScopedValues: ScopedValue, with
 end
+import TaskLocalValues: TaskLocalValue
+
+import TaskLocalValues: TaskLocalValue
 
 if !isdefined(Base, :get_extension)
     import Requires: @require
@@ -46,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")
@@ -67,6 +70,11 @@ include("sch/Sch.jl"); using .Sch
 # Data dependency task queue
 include("datadeps.jl")
 
+# Streaming
+include("stream.jl")
+include("stream-buffers.jl")
+include("stream-transfer.jl")
+
 # Array computations
 include("array/darray.jl")
 include("array/alloc.jl")
@@ -145,6 +153,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/compute.jl

@@ -36,12 +36,6 @@ end
 Base.@deprecate gather(ctx, x) collect(ctx, x)
 Base.@deprecate gather(x) collect(x)
 
-cleanup() = cleanup(Context(global_context()))
-function cleanup(ctx::Context)
-    Sch.cleanup(ctx)
-    nothing
-end
-
 function get_type(s::String)
     local T
     for t in split(s, ".")

src/dtask.jl

@@ -39,6 +39,16 @@ end
 Options(;options...) = Options((;options...))
 Options(options...) = Options((;options...))
 
+"""
+    DTaskMetadata
+
+Represents some useful metadata pertaining to a `DTask`:
+- `return_type::Type` - The inferred return type of the task
+"""
+mutable struct DTaskMetadata
+    return_type::Type
+end
+
 """
     DTask
 
@@ -50,9 +60,11 @@ more details.
 mutable struct DTask
     uid::UInt
     future::ThunkFuture
+    metadata::DTaskMetadata
     finalizer_ref::DRef
     thunk_ref::DRef
-    DTask(uid, future, finalizer_ref) = new(uid, future, finalizer_ref)
+
+    DTask(uid, future, metadata, finalizer_ref) = new(uid, future, metadata, finalizer_ref)
 end
 
 const EagerThunk = DTask

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