Merge pull request #409 from JuliaParallel/jps/mutable-auto-remote

at-mutable: Add remote execution like shard

Author: jpsamaroo

docs/src/data-management.md

@@ -41,28 +41,24 @@ argument to a task, the task will be forced to execute on the same worker that
 `@mutable` was called on. For example:
 
 ```julia
-x = remotecall_fetch(2) do
-    Dagger.@mutable Threads.Atomic{Int}(0)
-end
+Dagger.@mutable worker=2 Threads.Atomic{Int}(0)
 x::Dagger.Chunk # The result is always a `Chunk`
 
 # x is now considered mutable, and may only be accessed on worker 2:
-fetch(Dagger.@spawn Threads.atomic_add!(x, 3)) # Always executed on worker 2
-fetch(Dagger.@spawn single=1 Threads.atomic_add!(x, 3)) # SchedulingException
+wait(Dagger.@spawn Threads.atomic_add!(x, 1)) # Always executed on worker 2
+wait(Dagger.@spawn scope=Dagger.scope(worker=1) Threads.atomic_add!(x, 1)) # SchedulingException
 ```
 
-`@mutable`, when called as above, gain a scope of `ProcessorScope(myid())`,
-which means that any processor on that worker is allowed to execute tasks that
-use the object (subject to the usual scheduling rules).
+`@mutable`, when called as above, is constructed on worker 2, and the data
+gains a scope of `ProcessScope(myid())`, which means that any processor on that
+worker is allowed to execute tasks that use the object (subject to the usual
+scheduling rules).
 
-`@mutable` also has two other forms, allowing the processor and scope to be
-manually supplied:
+`@mutable` also allows the scope to be manually supplied, if more specific
+restrictions are desirable:
 
 ```julia
-proc1 = Dagger.ThreadProc(myid(), 3)
-proc2 = Dagger.ThreadProc(myid(), 4)
-scope = Dagger.UnionScope(ExactScope.([proc1, proc2]))
-x = @mutable OSProc() scope rand(100)
+x = @mutable scope=Dagger.scope(worker=1, threads=[3,4]) rand(100)
 # x is now scoped to threads 3 and 4 on worker `myid()`
 ```
 

docs/src/index.md

@@ -104,6 +104,12 @@ A = Dagger.spawn() do
 end
 ```
 
+or by specifying the `worker` argument to `@mutable`:
+
+```julia
+A = Dagger.@mutable worker=2 rand(1000, 1000)
+```
+
 ### Parallel reduction
 
 Reductions are often parallelized by reducing a set of partitions on each

src/chunks.jl

@@ -115,23 +115,42 @@ affinity(r::FileRef) = OSProc(1)=>r.size
 
 ### Mutation
 
-"Wraps `x` in a `Chunk` on `proc`, scoped to `scope`, which allows `x` to be mutated by tasks that use it."
-macro mutable(proc, scope, x)
-    :(Dagger.tochunk($(esc(x)), $(esc(proc)), $(esc(scope))))
+function _mutable_inner(@nospecialize(f), proc, scope)
+    result = f()
+    return Ref(Dagger.tochunk(result, proc, scope))
 end
-"Creates a mutable `Chunk` on `proc`, scoped to exactly `proc`."
-macro mutable(proc, x)
-    quote
-        let proc = $(esc(proc))
-            let scope = proc isa OSProc ? Dagger.ProcessScope(proc.pid) : Dagger.ExactScope(proc)
-                Dagger.@mutable proc scope $(esc(x))
-            end
+
+"""
+    mutable(f::Base.Callable; worker, processor, scope) -> Chunk
+    @mutable [worker=1] [processor=OSProc()] [scope=ProcessorScope()] f()
+
+Calls `f()` on the specified worker or processor, returning a `Chunk`
+referencing the result with the specified scope `scope`.
+"""
+function mutable(@nospecialize(f); worker=nothing, processor=nothing, scope=nothing)
+    if processor === nothing
+        if worker === nothing
+            processor = OSProc()
+        else
+            processor = OSProc(worker)
         end
+    else
+        @assert worker === nothing "mutable: Can't mix worker and processor"
+    end
+    if scope === nothing
+        scope = processor isa OSProc ? ProcessScope(processor) : ExactScope(processor)
     end
+    return fetch(Dagger.@spawn scope=scope _mutable_inner(f, processor, scope))[]
 end
-"Creates a mutable `Chunk` on the current worker."
-macro mutable(x)
-    :(Dagger.@mutable OSProc() Dagger.ProcessScope() $(esc(x)))
+
+macro mutable(exs...)
+    opts = esc.(exs[1:end-1])
+    ex = exs[end]
+    quote
+        let f = @noinline ()->$(esc(ex))
+            $mutable(f; $(opts...))
+        end
+    end
 end
 
 """
@@ -181,14 +200,11 @@ function shard(@nospecialize(f); procs=nothing, workers=nothing, per_thread=fals
     shard_dict = Dict{Processor,Chunk}()
     for proc in procs
         scope = proc isa OSProc ? ProcessScope(proc) : ExactScope(proc)
-        thunk = Dagger.@spawn scope=scope _shard_inner(f, proc, scope)
+        thunk = Dagger.@spawn scope=scope _mutable_inner(f, proc, scope)
         shard_dict[proc] = fetch(thunk)[]
     end
     return Shard(shard_dict)
 end
-function _shard_inner(@nospecialize(f), proc, scope)
-    Ref(Dagger.@mutable proc scope f())
-end
 
 "Creates a `Shard`. See [`Dagger.shard`](@ref) for details."
 macro shard(exs...)

test/mutation.jl

@@ -37,16 +37,19 @@
     end
 end
 
+@everywhere mutable_update!(x) = x[] = myid()
+
 @testset "Mutation" begin
 
 @testset "@mutable" begin
-    w = first(workers())
-    @assert w != 1 "Not enough workers to test mutability"
-    x = remotecall_fetch(w) do
-        Dagger.@mutable Ref{Int}()
+    w1 = first(workers())
+    @assert w1 != 1 "Not enough workers to test mutability"
+    for (w, wo) in [(1, w1), (w1, 1)]
+        x = Dagger.@mutable worker=w Ref{Int}()
+        @test fetch(Dagger.@spawn mutable_update!(x)) == w
+        wo_scope = Dagger.ProcessScope(wo)
+        @test_throws_unwrap Dagger.ThunkFailedException fetch(Dagger.@spawn scope=wo_scope mutable_update!(x))
     end
-    @test fetch(Dagger.@spawn (x->x[] = myid())(x)) == w
-    @test_throws_unwrap Dagger.ThunkFailedException fetch(Dagger.@spawn single=1 (x->x[] = myid())(x))
 end # @testset "@mutable"
 
 @testset "Shard" begin