Merge pull request #217 from JuliaParallel/jps/better-spawn

Add Dagger.spawn, allow on any node

Author: jpsamaroo

Diff for: Project.toml

@@ -1,6 +1,6 @@
 name = "Dagger"
 uuid = "d58978e5-989f-55fb-8d15-ea34adc7bf54"
-version = "0.11.2"
+version = "0.11.3"
 
 [deps]
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"

Diff for: docs/src/index.md

@@ -82,14 +82,14 @@ compute(top_node)
 
 ## Eager Execution
 
-Similar to `@par`, Dagger has an `@spawn` macro which works similarly to
-`@async` and `Threads.@spawn`: when called, it wraps the function call
-specified by the user in an `EagerThunk` object, and immediately places it onto
-a running scheduler, to be executed once its dependencies are fulfilled. This
-contrasts with `@par` in that `@par` does not begin executing its thunks until
-`collect` or `compute` are called on a given thunk or one of its downstream
-dependencies. Additionally, one fetches the result of an `@spawn` call with
-`fetch`. As a concrete example:
+Similar to `@par`, Dagger has an `@spawn` macro (and matching `Dagger.spawn`)
+which works similarly to `@async` and `Threads.@spawn`: when called, it wraps
+the function call specified by the user in an `EagerThunk` object, and
+immediately places it onto a running scheduler, to be executed once its
+dependencies are fulfilled. This contrasts with `@par` in that `@par` does not
+begin executing its thunks until `collect` or `compute` are called on a given
+thunk or one of its downstream dependencies. Additionally, one fetches the
+result of any `@spawn` call with `fetch`. As a concrete example:
 
 ```julia
 x = rand(400,400)
@@ -109,6 +109,17 @@ wait(x)
 @info "Done!"
 ```
 
+One can also safely call `@spawn` from another worker (not id 1), and it will
+be sent to worker 1 to schedule:
+
+```
+x = fetch(Distributed.@spawnat 2 Dagger.@spawn 1+2) # actually scheduled on worker 1
+x::EagerThunk
+@assert fetch(x) == 3
+```
+
+This is useful for nested execution, where an `@spawn`'d thunk calls `@spawn`.
+
 If a thunk errors while running under the eager scheduler, it will be marked as
 having failed, all dependent (downstream) thunks will be marked as failed, and
 any future thunks that use a failed thunk as input will fail. Failure can be

Diff for: src/processor.jl

@@ -312,16 +312,23 @@ rmprocs!(ctx::Context, xs::AbstractVector{<:OSProc}) = lock(ctx) do
 end
 
 "Gets the current processor executing the current thunk."
-thunk_processor() = task_local_storage(:processor)::Processor
+thunk_processor() = task_local_storage(:_dagger_processor)::Processor
+
+"Determines if we're currently in a thunk context."
+in_thunk() = haskey(task_local_storage(), :_dagger_sch_uid)
 
 "Gets all Dagger TLS variables as a NamedTuple."
 get_tls() = (
+    sch_uid=task_local_storage(:_dagger_sch_uid),
+    sch_handle=task_local_storage(:_dagger_sch_handle),
     processor=thunk_processor(),
-    sch_handle=task_local_storage(:sch_handle)
+    utilization=task_local_storage(:_dagger_utilization),
 )
 
 "Sets all Dagger TLS variables from a NamedTuple."
 function set_tls!(tls)
-    task_local_storage(:processor, tls.processor)
-    task_local_storage(:sch_handle, tls.sch_handle)
+    task_local_storage(:_dagger_sch_uid, tls.sch_uid)
+    task_local_storage(:_dagger_sch_handle, tls.sch_handle)
+    task_local_storage(:_dagger_processor, tls.processor)
+    task_local_storage(:_dagger_utilization, tls.utilization)
 end

Diff for: src/sch/Sch.jl

@@ -290,7 +290,9 @@ function compute_dag(ctx, d::Thunk; options=SchedulerOptions())
                 end
             end
         end
-        state.worker_pressure[pid][typeof(proc)] = metadata.pressure
+        if metadata !== nothing
+            state.worker_pressure[pid][typeof(proc)] = metadata.pressure
+        end
         node = state.thunk_dict[thunk_id]
         state.cache[node] = res
         if node.options !== nothing && node.options.checkpoint !== nothing
@@ -661,8 +663,12 @@ end
     res = nothing
     result_meta = try
         # Set TLS variables
-        task_local_storage(:processor, to_proc)
-        task_local_storage(:sch_handle, sch_handle)
+        Dagger.set_tls!((
+            sch_uid=uid,
+            sch_handle=sch_handle,
+            processor=to_proc,
+            utilization=extra_util,
+        ))
 
         # Execute
         res = execute!(to_proc, f, fetched...)

Diff for: src/sch/dynamic.jl

@@ -14,7 +14,7 @@ struct SchedulerHandle
 end
 
 "Gets the scheduler handle for the currently-executing thunk."
-sch_handle() = task_local_storage(:sch_handle)::SchedulerHandle
+sch_handle() = task_local_storage(:_dagger_sch_handle)::SchedulerHandle
 
 "Thrown when the scheduler halts before finishing processing the DAG."
 struct SchedulerHaltedException <: Exception end
@@ -33,7 +33,7 @@ end
 
 "Processes dynamic messages from worker-executing thunks."
 function dynamic_listener!(ctx, state)
-    task = current_task()
+    task = current_task() # The scheduler's main task
     for tid in keys(state.worker_chans)
         inp_chan, out_chan = state.worker_chans[tid]
         @async begin
@@ -71,10 +71,14 @@ end
 
 ## Worker-side methods for dynamic communication
 
+const DYNAMIC_EXEC_LOCK = Threads.ReentrantLock()
+
 "Executes an arbitrary function within the scheduler, returning the result."
 function exec!(f, h::SchedulerHandle, args...)
-    put!(h.out_chan, (h.thunk_id.id, f, args))
-    failed, res = take!(h.inp_chan)
+    failed, res = lock(DYNAMIC_EXEC_LOCK) do
+        put!(h.out_chan, (h.thunk_id.id, f, args))
+        take!(h.inp_chan)
+    end
     failed && throw(res)
     res
 end
@@ -138,13 +142,13 @@ end
 
 "Adds a new Thunk to the DAG."
 add_thunk!(f, h::SchedulerHandle, args...; kwargs...) =
-    ThunkID(exec!(_add_thunk!, h, f, args, kwargs)::Int)
+    ThunkID(exec!(_add_thunk!, h, f, args, kwargs))
 function _add_thunk!(ctx, state, task, tid, (f, args, kwargs))
     _args = map(arg->arg isa ThunkID ? state.thunk_dict[arg.id] : arg, args)
     thunk = Thunk(f, _args...; kwargs...)
     state.thunk_dict[thunk.id] = thunk
     state.dependents[thunk] = Set{Thunk}()
     @assert reschedule_inputs!(state, thunk) || (thunk in state.errored)
     schedule!(ctx, state)
-    return thunk.id
+    return thunk.id::Int
 end

Diff for: src/sch/eager.jl

@@ -29,18 +29,44 @@ function init_eager()
     end
 end
 
-"Sets the scheduler's cached pressure indicator for the specified worker."
-set_pressure!(h::SchedulerHandle, pid::Int, proctype::Type, pressure::Int) =
-    exec!(_set_pressure!, h, pid, proctype, pressure)
-function _set_pressure!(ctx, state, task, tid, (pid, proctype, pressure))
-    state.worker_pressure[pid][proctype] = pressure
-    ACTIVE_TASKS[state.uid][proctype] = pressure # HACK-ish
-    @show state.worker_pressure[pid]
+"Adjusts the scheduler's cached pressure indicator for the specified worker by
+the specified amount."
+function adjust_pressure!(h::SchedulerHandle, proctype::Type, pressure)
+    uid = Dagger.get_tls().sch_uid
+    lock(ACTIVE_TASKS_LOCK) do
+        ACTIVE_TASKS[uid][proctype][] += pressure
+    end
+    exec!(_adjust_pressure!, h, myid(), proctype, pressure)
+end
+function _adjust_pressure!(ctx, state, task, tid, (pid, proctype, pressure))
+    state.worker_pressure[pid][proctype] += pressure
+    nothing
+end
+
+"Allows a thunk to safely wait on another thunk, by temporarily reducing its
+effective pressure to 0."
+function thunk_yield(f)
+    if Dagger.in_thunk()
+        h = sch_handle()
+        tls = Dagger.get_tls()
+        proctype = typeof(tls.processor)
+        util = tls.utilization
+        adjust_pressure!(h, proctype, -util)
+        try
+            f()
+        finally
+            adjust_pressure!(h, proctype, util)
+        end
+    else
+        f()
+    end
 end
 
 function eager_thunk()
     h = sch_handle()
-    set_pressure!(h, 1, Dagger.ThreadProc, 0) # HACK: Don't apply pressure from this thunk
+    util = Dagger.get_tls().utilization
+    # Don't apply pressure from this thunk
+    adjust_pressure!(h, Dagger.ThreadProc, -util)
     while isopen(EAGER_THUNK_CHAN)
         try
             future, uid, f, args, opts = take!(EAGER_THUNK_CHAN)

Diff for: src/thunk.jl

@@ -90,9 +90,13 @@ end
 ThunkFuture(x::Integer) = ThunkFuture(Future(x))
 ThunkFuture() = ThunkFuture(Future())
 Base.isready(t::ThunkFuture) = isready(t.future)
-Base.wait(t::ThunkFuture) = wait(t.future)
+Base.wait(t::ThunkFuture) = Dagger.Sch.thunk_yield() do
+    wait(t.future)
+end
 function Base.fetch(t::ThunkFuture; proc=OSProc())
-    error, value = move(proc, fetch(t.future))
+    error, value = Dagger.Sch.thunk_yield() do
+        move(proc, fetch(t.future))
+    end
     if error
         throw(value)
     end
@@ -123,6 +127,18 @@ function Base.show(io::IO, t::EagerThunk)
     print(io, "EagerThunk ($(isready(t) ? "finished" : "running"))")
 end
 
+function spawn(f, args...; kwargs...)
+    if myid() == 1
+        Dagger.Sch.init_eager()
+        future = ThunkFuture()
+        uid = next_id()
+        put!(Dagger.Sch.EAGER_THUNK_CHAN, (future, uid, f, (args...,), (kwargs...,)))
+        EagerThunk(future, uid)
+    else
+        remotecall_fetch(spawn, 1, f, args...; kwargs...)
+    end
+end
+
 """
     @par [opts] f(args...) -> Thunk
 
@@ -172,11 +188,8 @@ function _par(ex::Expr; lazy=true, recur=true, opts=())
             return :(Dagger.delayed($(esc(f)); $(opts...))($(_par.(args; lazy=lazy, recur=false)...)))
         else
             return quote
-                Dagger.Sch.init_eager()
-                future = $ThunkFuture()
-                uid = $next_id()
-                put!(Dagger.Sch.EAGER_THUNK_CHAN, (future, uid, $(esc(f)), ($(_par.(args; lazy=lazy, recur=false)...),), ($(opts...),)))
-                EagerThunk(future, uid)
+                args = ($(_par.(args; lazy=lazy, recur=false)...),)
+                $spawn($(esc(f)), args...; $(opts...))
             end
         end
     else

Diff for: test/thunk.jl

@@ -1,6 +1,13 @@
-import Dagger: @par, @spawn
+import Dagger: @par, @spawn, spawn
 
-@everywhere checkwid() = myid()==1
+@everywhere begin
+    checkwid() = myid()==1
+    function dynamic_fib(n)
+        n <= 1 && return n
+        t = Dagger.spawn(dynamic_fib, n-1)
+        return (fetch(t)::Int) + dynamic_fib(n-2)
+    end
+end
 
 @testset "@par" begin
     @testset "per-call" begin
@@ -32,7 +39,8 @@ end
         a = @spawn x + x
         @test a isa Dagger.EagerThunk
         b = @spawn sum([x,1,2])
-        c = @spawn a * b
+        c = spawn(*, a, b)
+        @test c isa Dagger.EagerThunk
         @test fetch(a) == 4
         @test fetch(b) == 5
         @test fetch(c) == 20
@@ -131,4 +139,14 @@ end
             @test_throws_unwrap Dagger.ThunkFailedException fetch(d)
         end
     end
+    @testset "remote spawn" begin
+        a = fetch(Distributed.@spawnat 2 Dagger.spawn(+, 1, 2))
+        @test Dagger.Sch.EAGER_INIT[]
+        @test fetch(Distributed.@spawnat 2 !(Dagger.Sch.EAGER_INIT[]))
+        @test a isa Dagger.EagerThunk
+        @test fetch(a) == 3
+
+        # Mild stress-test
+        @test fetch(dynamic_fib(10)) == 55
+    end
 end