Add keyword argument support

APIs like `delayed` and `spawn` assumed that passed kwargs were to be
treated as options to the scheduler, which is both somewhat confusing
for users, and precludes passing kwargs to user functions.

This commit changes those APIs, as well as `@spawn`, to instead pass
kwargs directly to the user's function. Options are now passed in an
`Options` struct to `delayed` and `spawn` as the second argument (the
first being the function), while `@spawn` still keeps them before the
call (which is generally more convenient).

Internally, `Thunk`'s `inputs` field is now a
`Vector{Pair{Union{Symbol,Nothing},Any}}`, where the second element of
each pair is the argument, while the first element is a position; if
`nothing`, it's a positional argument, and if a `Symbol`, then it's a
kwarg.

Author: jpsamaroo

docs/src/checkpointing.md

@@ -54,17 +54,17 @@ Let's see how we'd modify the above example to use checkpointing:
 
 ```julia
 using Serialization
+
 X = compute(randn(Blocks(128,128), 1024, 1024))
-Y = [delayed(sum; options=Dagger.Sch.ThunkOptions(;
-checkpoint=(thunk,result)->begin
+Y = [delayed(sum; checkpoint=(thunk,result)->begin
     open("checkpoint-$idx.bin", "w") do io
         serialize(io, collect(result))
     end
 end, restore=(thunk)->begin
     open("checkpoint-$idx.bin", "r") do io
         Dagger.tochunk(deserialize(io))
     end
-end))(chunk) for (idx,chunk) in enumerate(X.chunks)]
+end)(chunk) for (idx,chunk) in enumerate(X.chunks)]
 inner(x...) = sqrt(sum(x))
 Z = delayed(inner)(Y...)
 z = collect(Z)

docs/src/index.md

@@ -2,32 +2,34 @@
 
 ## Usage
 
-The main function for using Dagger is `spawn`:
+The main entrypoint to Dagger is `@spawn`:
 
-`Dagger.spawn(f, args...; options...)`
+`Dagger.@spawn [option=value]... f(args...; kwargs...)`
 
-or `@spawn` for the more convenient macro form:
+or `spawn` if it's more convenient:
 
-`Dagger.@spawn [option=value]... f(args...)`
+`Dagger.spawn(f, Dagger.Options(options), args...; kwargs...)`
 
 When called, it creates an `EagerThunk` (also known as a "thunk" or "task")
-object representing a call to function `f` with the arguments `args`. If it is
-called with other thunks as inputs, such as in `Dagger.@spawn f(Dagger.@spawn
-g())`, then the function `f` gets passed the results of those input thunks. If
-those thunks aren't yet finished executing, then the execution of `f` waits on
-all of its input thunks to complete before executing.
+object representing a call to function `f` with the arguments `args` and
+keyword arguments `kwargs`. If it is called with other thunks as args/kwargs,
+such as in `Dagger.@spawn f(Dagger.@spawn g())`, then the function `f` gets
+passed the results of those input thunks, once they're available. If those
+thunks aren't yet finished executing, then the execution of `f` waits on all of
+its input thunks to complete before executing.
 
 The key point is that, for each argument to a thunk, if the argument is an
 `EagerThunk`, it'll be executed before this node and its result will be passed
 into the function `f`. If the argument is *not* an `EagerThunk` (instead, some
 other type of Julia object), it'll be passed as-is to the function `f`.
 
-Thunks don't accept regular keyword arguments for the function `f`. Instead,
-the `options` kwargs are passed to the scheduler to control its behavior:
+The `Options` struct in the second argument position is optional; if provided,
+it is passed to the scheduler to control its behavior. `Options` contains a
+`NamedTuple` of option key-value pairs, which can be any of:
 - Any field in `Dagger.Sch.ThunkOptions` (see [Scheduler and Thunk options](@ref))
 - `meta::Bool` -- Pass the input `Chunk` objects themselves to `f` and not the value contained in them
 
-There are also some extra kwargs that can be passed, although they're considered advanced options to be used only by developers or library authors:
+There are also some extra optionss that can be passed, although they're considered advanced options to be used only by developers or library authors:
 - `get_result::Bool` -- return the actual result to the scheduler instead of `Chunk` objects. Used when `f` explicitly constructs a Chunk or when return value is small (e.g. in case of reduce)
 - `persist::Bool` -- the result of this Thunk should not be released after it becomes unused in the DAG
 - `cache::Bool` -- cache the result of this Thunk such that if the thunk is evaluated again, one can just reuse the cached value. If it’s been removed from cache, recompute the value.
@@ -133,18 +135,18 @@ via `@par` or `delayed`. The above computation can be executed with the lazy
 API by substituting `@spawn` with `@par` and `fetch` with `collect`:
 
 ```julia
-p = @par add1(4)
-q = @par add2(p)
-r = @par add1(3)
-s = @par combine(p, q, r)
+p = Dagger.@par add1(4)
+q = Dagger.@par add2(p)
+r = Dagger.@par add1(3)
+s = Dagger.@par combine(p, q, r)
 
 @assert collect(s) == 16
 ```
 
 or similarly, in block form:
 
 ```julia
-s = @par begin
+s = Dagger.@par begin
     p = add1(4)
     q = add2(p)
     r = add1(3)
@@ -159,7 +161,7 @@ operation, you can call `compute` on the thunk. This will return a `Chunk`
 object which references the result (see [Chunks](@ref) for more details):
 
 ```julia
-x = @par 1+2
+x = Dagger.@par 1+2
 cx = compute(x)
 cx::Chunk
 @assert collect(cx) == 3
@@ -207,7 +209,7 @@ Scheduler options can be constructed and passed to `collect()` or `compute()`
 as the keyword argument `options` for lazy API usage:
 
 ```julia
-t = @par 1+2
+t = Dagger.@par 1+2
 opts = Dagger.Sch.SchedulerOptions(;single=1) # Execute on worker 1
 
 compute(t; options=opts)
@@ -221,10 +223,9 @@ Thunk options can be passed to `@spawn/spawn`, `@par`, and `delayed` similarly:
 # Execute on worker 1
 
 Dagger.@spawn single=1 1+2
-Dagger.spawn(+, 1, 2; single=1)
+Dagger.spawn(+, Dagger.Options(;single=1), 1, 2)
 
-opts = Dagger.Sch.ThunkOptions(;single=1)
-delayed(+)(1, 2; options=opts)
+delayed(+; single=1)(1, 2)
 ```
 
 ### Core vs. Worker Schedulers

docs/src/propagation.md

@@ -1,6 +1,6 @@
 # Option Propagation
 
-Most options passed to Dagger are passed via `delayed` or `Dagger.@spawn`
+Most options passed to Dagger are passed via `@spawn/spawn` or `delayed`
 directly. This works well when an option only needs to be set for a single
 thunk, but is cumbersome when the same option needs to be set on multiple
 thunks, or set recursively on thunks spawned within other thunks. Thankfully,

src/compute.jl

@@ -97,7 +97,7 @@ function dependents(node::Thunk)
         if !haskey(deps, next)
             deps[next] = Set{Thunk}()
         end
-        for inp in inputs(next)
+        for (_, inp) in next.inputs
             if istask(inp) || (inp isa Chunk)
                 s = get!(()->Set{Thunk}(), deps, inp)
                 push!(s, next)
@@ -165,7 +165,7 @@ function order(node::Thunk, ndeps)
         haskey(output, next) && continue
         s += 1
         output[next] = s
-        parents = filter(istask, inputs(next))
+        parents = filter(istask, map(last, next.inputs))
         if !isempty(parents)
             # If parents is empty, sort! should be a no-op, but raises an ambiguity error
             # when InlineStrings.jl is loaded (at least, version 1.1.0), because InlineStrings

src/processor.jl

@@ -31,11 +31,11 @@ function delete_processor_callback!(name::Symbol)
 end
 
 """
-    execute!(proc::Processor, f, args...) -> Any
+    execute!(proc::Processor, f, args...; kwargs...) -> Any
 
-Executes the function `f` with arguments `args` on processor `proc`. This
-function can be overloaded by `Processor` subtypes to allow executing function
-calls differently than normal Julia.
+Executes the function `f` with arguments `args` and keyword arguments `kwargs`
+on processor `proc`. This function can be overloaded by `Processor` subtypes to
+allow executing function calls differently than normal Julia.
 """
 function execute! end
 
@@ -154,12 +154,12 @@ end
 iscompatible(proc::ThreadProc, opts, f, args...) = true
 iscompatible_func(proc::ThreadProc, opts, f) = true
 iscompatible_arg(proc::ThreadProc, opts, x) = true
-function execute!(proc::ThreadProc, @nospecialize(f), @nospecialize(args...))
+function execute!(proc::ThreadProc, @nospecialize(f), @nospecialize(args...); @nospecialize(kwargs...))
     tls = get_tls()
     task = Task() do
         set_tls!(tls)
         TimespanLogging.prof_task_put!(tls.sch_handle.thunk_id.id)
-        @invokelatest f(args...)
+        @invokelatest f(args...; kwargs...)
     end
     task.sticky = true
     ret = ccall(:jl_set_task_tid, Cint, (Any, Cint), task, proc.tid-1)

src/sch/Sch.jl

@@ -130,7 +130,7 @@ function start_state(deps::Dict, node_order, chan)
 
     for k in sort(collect(keys(deps)), by=node_order)
         if istask(k)
-            waiting = Set{Thunk}(Iterators.filter(istask, inputs(k)))
+            waiting = Set{Thunk}(Iterators.filter(istask, map(last, inputs(k))))
             if isempty(waiting)
                 push!(state.ready, k)
             else
@@ -659,7 +659,7 @@ function schedule!(ctx, state, procs=procs_to_use(ctx))
                 DefaultScope()
             end
         end
-        for input in task.inputs
+        for (_,input) in task.inputs
             input = unwrap_weak_checked(input)
             chunk = if istask(input)
                 state.cache[input]
@@ -688,7 +688,7 @@ function schedule!(ctx, state, procs=procs_to_use(ctx))
             @goto fallback
         end
 
-        inputs = collect_task_inputs(state, task)
+        inputs = map(last, collect_task_inputs(state, task))
         opts = populate_defaults(opts, chunktype(task.f), map(chunktype, inputs))
         local_procs, costs = estimate_task_costs(state, local_procs, task, inputs)
         scheduled = false
@@ -945,10 +945,13 @@ function fire_tasks!(ctx, thunks::Vector{<:Tuple}, (gproc, proc), state)
 
         ids = Int[0]
         data = Any[thunk.f]
-        for (idx, x) in enumerate(thunk.inputs)
+        positions = Union{Symbol,Nothing}[]
+        for (idx, pos_x) in enumerate(thunk.inputs)
+            pos, x = pos_x
             x = unwrap_weak_checked(x)
             push!(ids, istask(x) ? x.id : -idx)
             push!(data, istask(x) ? state.cache[x] : x)
+            push!(positions, pos)
         end
         toptions = thunk.options !== nothing ? thunk.options : ThunkOptions()
         options = merge(ctx.options, toptions)
@@ -961,7 +964,7 @@ function fire_tasks!(ctx, thunks::Vector{<:Tuple}, (gproc, proc), state)
         push!(to_send, Any[thunk.id, time_util, alloc_util, occupancy,
                            scope, chunktype(thunk.f), data,
                            thunk.get_result, thunk.persist, thunk.cache, thunk.meta, options,
-                           propagated, ids,
+                           propagated, ids, positions,
                            (log_sink=ctx.log_sink, profile=ctx.profile),
                            sch_handle, state.uid])
     end
@@ -1093,6 +1096,7 @@ function start_processor_runner!(istate::ProcessorInternalState, uid::UInt64, re
 
         while isopen(return_queue)
             # Wait for new tasks
+            @dagdebug nothing :processor "Waiting for tasks"
             timespan_start(ctx, :proc_run_wait, to_proc, nothing)
             wait(istate.reschedule)
             @static if VERSION >= v"1.9"
@@ -1101,22 +1105,27 @@ function start_processor_runner!(istate::ProcessorInternalState, uid::UInt64, re
             timespan_finish(ctx, :proc_run_wait, to_proc, nothing)
 
             # Fetch a new task to execute
+            @dagdebug nothing :processor "Trying to dequeue"
             timespan_start(ctx, :proc_run_fetch, to_proc, nothing)
             task_and_occupancy = lock(istate.queue) do queue
                 # Only steal if there are multiple queued tasks, to prevent
                 # ping-pong of tasks between empty queues
                 if length(queue) == 0
+                    @dagdebug nothing :processor "Nothing to dequeue"
                     return nothing
                 end
                 _, occupancy = peek(queue)
-                if proc_has_occupancy(proc_occupancy[], occupancy)
-                    return dequeue_pair!(queue)
+                if !proc_has_occupancy(proc_occupancy[], occupancy)
+                    @dagdebug nothing :processor "Insufficient occupancy" proc_occupancy=proc_occupancy[] task_occupancy=occupancy
+                    return nothing
                 end
-                return nothing
+                return dequeue_pair!(queue)
             end
             if task_and_occupancy === nothing
                 timespan_finish(ctx, :proc_run_fetch, to_proc, nothing)
 
+                @dagdebug nothing :processor "Failed to dequeue"
+
                 if !stealing_permitted(to_proc)
                     continue
                 end
@@ -1125,6 +1134,8 @@ function start_processor_runner!(istate::ProcessorInternalState, uid::UInt64, re
                     continue
                 end
 
+                @dagdebug nothing :processor "Trying to steal"
+
                 # Try to steal a task
                 timespan_start(ctx, :steal_local, to_proc, nothing)
 
@@ -1159,7 +1170,7 @@ function start_processor_runner!(istate::ProcessorInternalState, uid::UInt64, re
                     if task_and_occupancy !== nothing
                         from_proc = other_istate.proc
                         thunk_id = task[1]
-                        @dagdebug thunk_id :execute "Stolen from $from_proc by $to_proc"
+                        @dagdebug thunk_id :processor "Stolen from $from_proc by $to_proc"
                         timespan_finish(ctx, :steal_local, to_proc, (;from_proc, thunk_id))
                         # TODO: Keep stealing until we hit full occupancy?
                         @goto execute
@@ -1177,6 +1188,7 @@ function start_processor_runner!(istate::ProcessorInternalState, uid::UInt64, re
             thunk_id = task[1]
             time_util = task[2]
             timespan_finish(ctx, :proc_run_fetch, to_proc, (;thunk_id, proc_occupancy=proc_occupancy[], task_occupancy))
+            @dagdebug thunk_id :processor "Dequeued task"
 
             # Execute the task and return its result
             t = @task begin
@@ -1234,10 +1246,12 @@ Executes a batch of tasks on `to_proc`, returning their results through
 `return_queue`.
 """
 function do_tasks(to_proc, return_queue, tasks)
+    @dagdebug nothing :processor "Enqueuing task batch" batch_size=length(tasks)
+
     # FIXME: This is terrible
-    ctx_vars = first(tasks)[15]
+    ctx_vars = first(tasks)[16]
     ctx = Context(Processor[]; log_sink=ctx_vars.log_sink, profile=ctx_vars.profile)
-    uid = first(tasks)[17]
+    uid = first(tasks)[18]
     state = proc_states(uid) do states
         get!(states, to_proc) do
             queue = PriorityQueue{Vector{Any}, UInt32}()
@@ -1272,6 +1286,7 @@ function do_tasks(to_proc, return_queue, tasks)
             should_launch || continue
             enqueue!(queue, task, occupancy)
             timespan_finish(ctx, :enqueue, (;to_proc, thunk_id), nothing)
+            @dagdebug thunk_id :processor "Enqueued task"
         end
     end
     notify(istate.reschedule)
@@ -1287,6 +1302,7 @@ function do_tasks(to_proc, return_queue, tasks)
         end
         notify(other_istate.reschedule)
     end
+    @dagdebug nothing :processor "Kicked processors"
 end
 
 """
@@ -1298,7 +1314,7 @@ function do_task(to_proc, task_desc)
     thunk_id, est_time_util, est_alloc_util, est_occupancy,
         scope, Tf, data,
         send_result, persist, cache, meta,
-        options, propagated, ids, ctx_vars, sch_handle, uid = task_desc
+        options, propagated, ids, positions, ctx_vars, sch_handle, uid = task_desc
     ctx = Context(Processor[]; log_sink=ctx_vars.log_sink, profile=ctx_vars.profile)
 
     from_proc = OSProc()
@@ -1442,6 +1458,16 @@ function do_task(to_proc, task_desc)
     end
     f = popfirst!(fetched)
     @assert !(f isa Chunk) "Failed to unwrap thunk function"
+    fetched_args = Any[]
+    fetched_kwargs = Pair{Symbol,Any}[]
+    for (idx, x) in enumerate(fetched)
+        pos = positions[idx]
+        if pos === nothing
+            push!(fetched_args, x)
+        else
+            push!(fetched_kwargs, pos => x)
+        end
+    end
 
     #= FIXME: If MaxUtilization, stop processors and wait
     if (est_time_util isa MaxUtilization) && (real_time_util > 0)
@@ -1473,7 +1499,7 @@ function do_task(to_proc, task_desc)
 
         res = Dagger.with_options(propagated) do
             # Execute
-            execute!(to_proc, f, fetched...)
+            execute!(to_proc, f, fetched_args...; fetched_kwargs...)
         end
 
         # Check if result is safe to store

src/sch/dynamic.jl

@@ -144,7 +144,7 @@ function _register_future!(ctx, state, task, tid, (future, id, check)::Tuple{Thu
                 if t == target
                     return true
                 end
-                for input in t.inputs
+                for (_, input) in t.inputs
                     # N.B. Skips expired tasks
                     input = Dagger.unwrap_weak(input)
                     istask(input) || continue
@@ -195,13 +195,13 @@ function _get_dag_ids(ctx, state, task, tid, _)
 end
 
 "Adds a new Thunk to the DAG."
-add_thunk!(f, h::SchedulerHandle, args...; future=nothing, ref=nothing, kwargs...) =
-    exec!(_add_thunk!, h, f, args, kwargs, future, ref)
-function _add_thunk!(ctx, state, task, tid, (f, args, kwargs, future, ref))
+add_thunk!(f, h::SchedulerHandle, args...; future=nothing, ref=nothing, options...) =
+    exec!(_add_thunk!, h, f, args, options, future, ref)
+function _add_thunk!(ctx, state, task, tid, (f, args, options, future, ref))
     timespan_start(ctx, :add_thunk, tid, 0)
-    _args = map(arg->arg isa ThunkID ? state.thunk_dict[arg.id] : arg, args)
+    _args = map(pos_arg->pos_arg[1] => (pos_arg[2] isa ThunkID ? state.thunk_dict[pos_arg[2].id] : pos_arg[2]), args)
     GC.@preserve _args begin
-        thunk = Thunk(f, _args...; kwargs...)
+        thunk = Thunk(f, _args...; options...)
         # Create a `DRef` to `thunk` so that the caller can preserve it
         thunk_ref = poolset(thunk; size=64, device=MemPool.CPURAMDevice())
         thunk_id = ThunkID(thunk.id, thunk_ref)