Merge pull request #271 from JuliaParallel/jps/function-chunks

Wrap thunk functions in Chunk

Author: jpsamaroo

Diff for: Project.toml

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

Diff for: docs/src/processors.md

@@ -1,14 +1,56 @@
 # Processors
 
 Dagger contains a flexible mechanism to represent CPUs, GPUs, and other
-devices that the scheduler can place user work on. The indiviual devices that
+devices that the scheduler can place user work on. The individual devices that
 are capable of computing a user operation are called "processors", and are
 subtypes of `Dagger.Processor`. Processors are automatically detected by
 Dagger at scheduler initialization, and placed in a hierarchy reflecting the
 physical (network-, link-, or memory-based) boundaries between processors in
 the hierarchy. The scheduler uses the information in this hierarchy to
 efficiently schedule and partition user operations.
 
+Dagger's `Chunk` objects can have a processor associated with them that defines
+where the contained data "resides". Each processor has a set of functions that
+define the mechanisms and rules by which the data can be transferred between
+similar or different kinds of processors, and will be called by Dagger's
+scheduler automatically when fetching function arguments (or the function
+itself) for computation on a given processor.
+
+Setting the processor on a function argument is done by wrapping it in a
+`Chunk` with `Dagger.tochunk`:
+
+```julia
+a = 1
+b = 2
+# Let's say `b` "resides" on the second thread of the first worker:
+b_chunk = Dagger.tochunk(b, Dagger.ThreadProc(1, 2))::Dagger.Chunk
+c = Dagger.@spawn a + b_chunk
+fetch(c) == 3
+```
+
+It's also simple to set the processor of the function being passed; it will be
+automatically wrapped in a `Chunk` if necessary:
+
+```julia
+# `+` is treated as existing on the second thread of the first worker:
+Dagger.@spawn processor=Dagger.ThreadProc(1, 2) a + b
+```
+
+You can also tell Dagger about the processor type for the returned value of a
+task by making it a `Chunk`:
+
+```julia
+Dagger.spawn(a) do a
+    c = a + 1
+    return Dagger.tochunk(c, Dagger.ThreadProc(1, 2))
+end
+```
+
+Note that unless you know that your function, arguments, or return value are
+associated with a specific processor, you don't need to assign one to them.
+Dagger will treat them as being simple values with no processor association,
+and will serialize them to wherever they're used.
+
 ## Hardware capabilities, topology, and data locality
 
 The processor hierarchy is modeled as a multi-root tree, where each root is an

Diff for: docs/src/scopes.md

@@ -13,7 +13,10 @@ considered valid.
 
 ## Scope Basics
 
-Let's take the example of a webcam handle generated by VideoIO.jl. This handle is a C pointer, and thus has process scope. We can open the handle on a given process, and set the scope of the resulting data to a `ProcessScope()`, which defaults to the current Julia process:
+Let's take the example of a webcam handle generated by VideoIO.jl. This handle
+is a C pointer, and thus has process scope. We can open the handle on a given
+process, and set the scope of the resulting data to a `ProcessScope()`, which
+defaults to the current Julia process:
 
 ```julia
 using VideoIO
@@ -41,6 +44,22 @@ the `cam_handle` task was executed on. Of course, the resulting camera frame is
 *not* scoped to anywhere specific (denoted as `AnyScope()`), and thus
 computations on it may execute anywhere.
 
+You may also encounter situations where you want to use a callable struct (such
+as a closure, or a Flux.jl layer) only within a certain scope; you can specify
+the scope of the function pretty easily:
+
+```julia
+using Flux
+m = Chain(...)
+# If `m` is only safe to transfer to and execute on this process,
+# we can set a `ProcessScope` on it:
+result = Dagger.@spawn scope=ProcessScope() m(rand(8,8))
+```
+
+Setting a scope on the function treats it as a regular piece of data (like the
+arguments to the function), so it participates in the scoping rules described
+in the following sections all the same.
+
 Now, let's try out some other kinds of scopes, starting with `NodeScope`. This
 scope encompasses the server that one or more Julia processes may be running
 on. Say we want to use memory mapping (mmap) to more efficiently send arrays

Diff for: src/Dagger.jl

@@ -79,12 +79,10 @@ function __init__()
             include("ui/video.jl")
         end
     end
-    @static if VERSION >= v"1.3.0-DEV.573"
-        for tid in 1:Threads.nthreads()
-            push!(PROCESSOR_CALLBACKS, ()->ThreadProc(myid(), tid))
+    for tid in 1:Threads.nthreads()
+        add_processor_callback!("__cpu_thread_$(tid)__") do
+            ThreadProc(myid(), tid)
         end
-    else
-        push!(PROCESSOR_CALLBACKS, ()->ThreadProc(myid(), 1))
     end
 end
 

Diff for: src/processor.jl

@@ -12,7 +12,7 @@ make it easy to transfer data to/from other types of `Processor` at runtime.
 """
 abstract type Processor end
 
-const PROCESSOR_CALLBACKS = []
+const PROCESSOR_CALLBACKS = Dict{Symbol,Any}()
 
 """
     execute!(proc::Processor, f, args...) -> Any
@@ -104,7 +104,8 @@ const OSPROC_CACHE = Dict{Int,Vector{Processor}}()
 children(proc::OSProc) = OSPROC_CACHE[proc.pid]
 function get_proc_hierarchy()
     children = Processor[]
-    for cb in PROCESSOR_CALLBACKS
+    for name in keys(PROCESSOR_CALLBACKS)
+        cb = PROCESSOR_CALLBACKS[name]
         try
             child = Base.invokelatest(cb)
             if (child isa Tuple) || (child isa Vector)
@@ -113,13 +114,21 @@ function get_proc_hierarchy()
                 push!(children, child)
             end
         catch err
-            @error "Error in processor callback" exception=(err,catch_backtrace())
+            @error "Error in processor callback: $name" exception=(err,catch_backtrace())
         end
     end
     children
 end
-function add_callback!(func)
-    push!(Dagger.PROCESSOR_CALLBACKS, func)
+add_processor_callback!(func, name::String) =
+    add_processor_callback!(func, Symbol(name))
+function add_processor_callback!(func, name::Symbol)
+    Dagger.PROCESSOR_CALLBACKS[name] = func
+    empty!(OSPROC_CACHE)
+end
+delete_processor_callback!(name::String) =
+    delete_processor_callback!(Symbol(name))
+function delete_processor_callback!(name::Symbol)
+    delete!(Dagger.PROCESSOR_CALLBACKS, name)
     empty!(OSPROC_CACHE)
 end
 Base.:(==)(proc1::OSProc, proc2::OSProc) = proc1.pid == proc2.pid

Diff for: src/sch/Sch.jl

@@ -589,7 +589,11 @@ function schedule!(ctx, state, procs=procs_to_use(ctx))
         sig = signature(task, state)
 
         # Calculate scope
-        scope = AnyScope()
+        scope = if task.f isa Chunk
+            task.f.scope
+        else
+            AnyScope()
+        end
         for input in unwrap_weak_checked.(task.inputs)
             chunk = if istask(input)
                 state.cache[input]
@@ -763,6 +767,10 @@ function finish_task!(ctx, state, node, thunk_failed)
     fill_registered_futures!(state, node, thunk_failed)
 
     to_evict = cleanup_inputs!(state, node)
+    if node.f isa Chunk
+        # FIXME: Check the graph for matching chunks
+        push!(to_evict, node.f)
+    end
     if haskey(state.waiting_data, node) && isempty(state.waiting_data[node])
         delete!(state.waiting_data, node)
     end
@@ -826,13 +834,15 @@ function fire_tasks!(ctx, thunks::Vector{<:Tuple}, (gproc, proc), state)
             end
         end
 
-        ids = map(enumerate(thunk.inputs)) do (idx,x)
+        ids = convert(Vector{Int}, map(enumerate(thunk.inputs)) do (idx,x)
             istask(x) ? unwrap_weak_checked(x).id : -idx
-        end
+        end)
+        pushfirst!(ids, 0)
 
-        data = map(thunk.inputs) do x
+        data = convert(Vector{Any}, map(Any[thunk.inputs...]) do x
             istask(x) ? state.cache[unwrap_weak_checked(x)] : x
-        end
+        end)
+        pushfirst!(data, thunk.f)
         toptions = thunk.options !== nothing ? thunk.options : ThunkOptions()
         options = merge(ctx.options, toptions)
         @assert (options.single == 0) || (gproc.pid == options.single)
@@ -841,7 +851,7 @@ function fire_tasks!(ctx, thunks::Vector{<:Tuple}, (gproc, proc), state)
         state.worker_pressure[gproc.pid][typeof(proc)] += util
 
         # TODO: De-dup common fields (log_sink, uid, etc.)
-        push!(to_send, (util, thunk.id, thunk.f, data, thunk.get_result,
+        push!(to_send, (util, thunk.id, fn_type(thunk.f), data, thunk.get_result,
                         thunk.persist, thunk.cache, thunk.meta, options, ids,
                         (log_sink=ctx.log_sink, profile=ctx.profile),
                         sch_handle, state.uid))
@@ -894,11 +904,12 @@ function do_tasks(to_proc, chan, tasks)
     end
 end
 "Executes a single task on `to_proc`."
-function do_task(to_proc, extra_util, thunk_id, f, data, send_result, persist, cache, meta, options, ids, ctx_vars, sch_handle, uid)
+function do_task(to_proc, extra_util, thunk_id, Tf, data, send_result, persist, cache, meta, options, ids, ctx_vars, sch_handle, uid)
     ctx = Context(Processor[]; log_sink=ctx_vars.log_sink, profile=ctx_vars.profile)
 
     from_proc = OSProc()
     Tdata = map(x->x isa Chunk ? chunktype(x) : x, data)
+    f = isdefined(Tf, :instance) ? Tf.instance : nothing
 
     # Fetch inputs
     fetched = if meta
@@ -929,6 +940,7 @@ function do_task(to_proc, extra_util, thunk_id, f, data, send_result, persist, c
             end
         end)
     end
+    f = popfirst!(fetched)
 
     # Check if we'll go over capacity from running this thunk
     real_util = lock(TASK_SYNC) do

Diff for: src/sch/eager.jl

@@ -110,3 +110,6 @@ function eager_cleanup(state, uid)
         delete!(state.thunk_dict, tid)
     end
 end
+
+_find_thunk(e::Dagger.EagerThunk) =
+    unwrap_weak_checked(EAGER_STATE[].thunk_dict[EAGER_ID_MAP[e.uid]])

Diff for: src/sch/util.jl

@@ -203,9 +203,11 @@ function fetch_report(task)
     end
 end
 
+fn_type(x::Chunk) = x.chunktype
+fn_type(x) = typeof(x)
 function signature(task::Thunk, state)
     inputs = map(x->istask(x) ? state.cache[x] : x, unwrap_weak_checked.(task.inputs))
-    Tuple{typeof(task.f), map(x->x isa Chunk ? x.chunktype : typeof(x), inputs)...}
+    Tuple{fn_type(task.f), map(x->x isa Chunk ? x.chunktype : typeof(x), inputs)...}
 end
 
 function report_catch_error(err, desc=nothing)

Diff for: src/scopes.jl

@@ -32,7 +32,7 @@ function ProcessScope(wid)
         ProcessScope(NodeScope(system_uuid(wid)), wid)
     end
 end
-ProcessScope() = ProcessScope(myid)
+ProcessScope() = ProcessScope(myid())
 
 "Scoped to a specific processor."
 struct ExactScope <: AbstractScope

Diff for: src/thunk.jl

@@ -11,17 +11,10 @@ created through a call to `delayed` or its macro equivalent `@par`.
 
 ## Constructors
 ```julia
-delayed(f; options=nothing)(args...)
+delayed(f; kwargs...)(args...)
 @par [option=value]... f(args...)
 ```
 
-## Arguments
-- `f`: The function to be called upon execution of the `Thunk`.
-- `args`: The arguments to be passed to the `Thunk`.
-- `options`: A `Sch.ThunkOptions` struct providing the options for the `Thunk`,
-or `nothing`.
-- `option=value`: The same options available in `Sch.ThunkOptions`.
-
 ## Examples
 ```julia
 julia> t = delayed(sin)(π)  # creates a Thunk to be computed later
@@ -30,6 +23,30 @@ Thunk(sin, (π,))
 julia> collect(t)  # computes the result and returns it to the current process
 1.2246467991473532e-16
 ```
+
+## Arguments
+- `f`: The function to be called upon execution of the `Thunk`.
+- `args`: The arguments to be passed to the `Thunk`.
+- `kwargs`: The properties describing unique behavior of this `Thunk`. Details
+for each property are described in the next section.
+- `option=value`: The same as passing `kwargs` to `delayed`.
+
+## Public Properties
+- `meta::Bool=false`: If `true`, instead of fetching cached arguments from
+`Chunk`s and passing the raw arguments to `f`, instead pass the `Chunk`. Useful
+for doing manual fetching or manipulation of `Chunk` references. Non-`Chunk`
+arguments are still passed as-is.
+- `processor::Processor=OSProc()` - The processor associated with `f`. Useful if
+`f` is a callable struct that exists on a given processor and should be
+transferred appropriately.
+- `scope::Dagger.AbstractScope=AnyScope()` - The scope associated with `f`.
+Useful if `f` is a function or callable struct that may only be transferred to,
+and executed within, the specified scope.
+
+## Options
+- `options`: A `Sch.ThunkOptions` struct providing the options for the `Thunk`.
+If omitted, options can also be specified by passing key-value pairs as
+`kwargs`.
 """
 mutable struct Thunk
     f::Any # usually a Function, but could be any callable
@@ -53,9 +70,16 @@ mutable struct Thunk
                    cache_ref=nothing,
                    affinity=nothing,
                    eager_ref=nothing,
+                   processor=nothing,
+                   scope=nothing,
                    options=nothing,
                    kwargs...
                   )
+        if !isnothing(processor) || !isnothing(scope)
+            f = tochunk(f,
+                        something(processor, OSProc()),
+                        something(scope, AnyScope()))
+        end
         if options !== nothing
             @assert isempty(kwargs)
             new(f, xs, id, get_result, meta, persist, cache, cache_ref,
@@ -103,11 +127,10 @@ function affinity(t::Thunk)
 end
 
 """
-    delayed(f; options)(args...)
+    delayed(f; kwargs...)(args...)
 
 Creates a [`Thunk`](@ref) object which can be executed later, which will call
-`f` with `args`. Options are typically either `nothing` or of type
-`Sch.ThunkOptions`.
+`f` with `args`. `kwargs` controls various properties of the resulting `Thunk`.
 """
 function delayed(f; kwargs...)
     (args...) -> Thunk(f, args...; kwargs...)
@@ -218,8 +241,16 @@ end
 
 Spawns a task with `f` as the function and `args` as the arguments, returning
 an `EagerThunk`. Uses a scheduler running in the background to execute code.
+
+Note that `kwargs` are passed to the `Thunk` constructor, and are documented in
+its docstring.
 """
-function spawn(f, args...; kwargs...)
+function spawn(f, args...; processor=nothing, scope=nothing, kwargs...)
+    if !isnothing(processor) || !isnothing(scope)
+        f = tochunk(f,
+                    something(processor, OSProc()),
+                    something(scope, AnyScope()))
+    end
     uid, future, finalizer_ref, thunk_ref = if myid() == 1
         _spawn(f, args...; kwargs...)
     else
@@ -316,7 +347,17 @@ Base.isequal(x::Thunk, y::Thunk) = x.id==y.id
 
 function Base.show(io::IO, z::Thunk)
     lvl = get(io, :lazy_level, 2)
-    print(io, "Thunk[$(z.id)]($(z.f), ")
+    f = if z.f isa Chunk
+        Tf = z.f.chunktype
+        if isdefined(Tf, :instance)
+            Tf.instance
+        else
+            "instance of $Tf"
+        end
+    else
+        z.f
+    end
+    print(io, "Thunk[$(z.id)]($f, ")
     if lvl > 0
         show(IOContext(io, :lazy_level => lvl-1), z.inputs)
     else

Diff for: test/processors.jl

@@ -54,9 +54,9 @@ end
     @testset "Add callback in same world" begin
         function addcb()
             cb = @eval ()->FakeProc(myid())
-            @everywhere Dagger.add_callback!($cb)
+            @everywhere Dagger.add_processor_callback!($cb, :fakeproc)
             @test any(x->x isa FakeProc, Dagger.children(OSProc()))
-            @everywhere pop!(Dagger.PROCESSOR_CALLBACKS); empty!(Dagger.OSPROC_CACHE)
+            @everywhere Dagger.delete_processor_callback!(:fakeproc)
         end
         addcb()
     end