scopes: Add Dagger.scope helper

Author: jpsamaroo

src/scopes.jl

@@ -205,3 +205,137 @@ constrain(x::ProcessScope, y::ExactScope) =
     x == y.parent ? y : InvalidScope(x, y)
 constrain(x::NodeScope, y::ExactScope) =
     x == y.parent.parent ? y : InvalidScope(x, y)
+
+### Scopes helper
+
+"""
+    scope(scs...) -> AbstractScope
+    scope(;scs...) -> AbstractScope
+
+Constructs an `AbstractScope` from a set of scope specifiers. Each element in
+`scs` is a separate specifier; if `scs` is empty, an empty `UnionScope()` is
+produced; if `scs` has one element, then exactly one specifier is constructed;
+if `scs` has more than one element, a `UnionScope` of the scopes specified by
+`scs` is constructed. A variety of specifiers can be passed to construct a
+scope:
+- `:any` - Constructs an `AnyScope()`
+- `:default` - Constructs a `DefaultScope()`
+- `(scs...,)` - Constructs a `UnionScope` of scopes, each specified by `scs`
+- `thread=tid` or `threads=[tids...]` - Constructs an `ExactScope` or `UnionScope` containing all `Dagger.ThreadProc`s with thread ID `tid`/`tids` across all workers.
+- `worker=wid` or `workers=[wids...]` - Constructs a `ProcessScope` or `UnionScope` containing all `Dagger.ThreadProc`s with worker ID `wid`/`wids` across all threads.
+- `thread=tid`/`threads=tids` and `worker=wid`/`workers=wids` - Constructs an `ExactScope`, `ProcessScope`, or `UnionScope` containing all `Dagger.ThreadProc`s with worker ID `wid`/`wids` and threads `tid`/`tids`.
+
+Aside from the worker and thread specifiers, it's possible to add custom
+specifiers for scoping to other kinds of processors (like GPUs) or providing
+different ways to specify a scope. Specifier selection is determined by a
+precedence ordering: by default, all specifiers have precedence `0`, which can
+be changed by defining `scope_key_precedence(::Val{spec}) = precedence` (where
+`spec` is the specifier as a `Symbol)`. The specifier with the highest
+precedence in a set of specifiers is used to determine the scope by calling
+`to_scope(::Val{spec}, sc::NamedTuple)` (where `sc` is the full set of
+specifiers), which should be overriden for each custom specifier, and which
+returns an `AbstractScope`. For example:
+
+```julia
+# Setup a GPU specifier
+Dagger.scope_key_precedence(::Val{:gpu}) = 1
+Dagger.to_scope(::Val{:gpu}, sc::NamedTuple) = ExactScope(MyGPUDevice(sc.worker, sc.gpu))
+
+# Generate an `ExactScope` for `MyGPUDevice` on worker 2, device 3
+Dagger.scope(gpu=3, worker=2)
+```
+"""
+scope(scs...) = simplified_union_scope(map(to_scope, scs))
+scope(; kwargs...) = to_scope((;kwargs...))
+
+function simplified_union_scope(scopes)
+    if length(scopes) == 1
+        return only(scopes)
+    else
+        return UnionScope(scopes)
+    end
+end
+to_scope(scope::AbstractScope) = scope
+function to_scope(sc::Symbol)
+    if sc == :any
+        return AnyScope()
+    elseif sc == :default
+        return DefaultScope()
+    else
+        throw(ArgumentError("Cannot construct scope from: $(repr(sc))"))
+    end
+end
+function to_scope(sc::NamedTuple)
+    if isempty(sc)
+        return UnionScope()
+    end
+
+    # FIXME: node and nodes
+    known_keys = (:worker, :workers, :thread, :threads)
+    unknown_keys = filter(key->!in(key, known_keys), keys(sc))
+    if length(unknown_keys) > 0
+        # Hand off construction if unknown members encountered
+        precs = map(scope_key_precedence, map(Val, unknown_keys))
+        max_prec = maximum(precs)
+        if length(findall(prec->prec==max_prec, precs)) > 1
+            throw(ArgumentError("Incompatible scope specifiers detected: $unknown_keys"))
+        end
+        max_prec_key = unknown_keys[argmax(precs)]
+        return to_scope(Val(max_prec_key), sc)
+    end
+
+    workers = if haskey(sc, :worker)
+        Int[sc.worker]
+    elseif haskey(sc, :workers)
+        Int[sc.workers...]
+    else
+        nothing
+    end
+    threads = if haskey(sc, :thread)
+        Int[sc.thread]
+    elseif haskey(sc, :threads)
+        Int[sc.threads...]
+    else
+        nothing
+    end
+
+    # Simple cases
+    if workers !== nothing && threads !== nothing
+        subscopes = AbstractScope[]
+        for w in workers, t in threads
+            push!(subscopes, ExactScope(ThreadProc(w, t)))
+        end
+        return simplified_union_scope(subscopes)
+    elseif workers !== nothing && threads === nothing
+        subscopes = AbstractScope[ProcessScope(w) for w in workers]
+        return simplified_union_scope(subscopes)
+    end
+
+    # More complex cases that require querying the cluster
+    # FIXME: Use per-field scope taint
+    if workers === nothing
+        workers = procs()
+    end
+    subscopes = AbstractScope[]
+    for w in workers
+        if threads === nothing
+            threads = map(c->c.tid,
+                          filter(c->c isa ThreadProc,
+                                 collect(children(OSProc(w)))))
+        end
+        for t in threads
+            push!(subscopes, ExactScope(ThreadProc(w, t)))
+        end
+    end
+    return simplified_union_scope(subscopes)
+end
+to_scope(scs::Tuple) =
+    simplified_union_scope(map(to_scope, scs))
+to_scope(sc) =
+    throw(ArgumentError("Cannot construct scope from: $sc"))
+
+to_scope(::Val{key}, sc::NamedTuple) where key =
+    throw(ArgumentError("Scope construction not implemented for key: $key"))
+
+# Base case for all Dagger-owned keys
+scope_key_precedence(::Val) = 0

test/scopes.jl

@@ -124,7 +124,7 @@
         @test fetch(Dagger.@spawn exact_scope_test(us_es1_multi_ch)) == es1.processor
 
         # No inner scopes
-        @test_throws ArgumentError UnionScope()
+        @test UnionScope() isa UnionScope
 
         # Same inner scope
         @test fetch(Dagger.@spawn exact_scope_test(us_es1_ch, us_es1_ch)) == es1.processor
@@ -161,5 +161,82 @@
     end
     # TODO: Test scope propagation
 
+    @testset "scope helper" begin
+        @test Dagger.scope(:any) isa AnyScope
+        @test Dagger.scope(:default) == DefaultScope()
+        @test_throws ArgumentError Dagger.scope(:blah)
+        @test Dagger.scope(()) == UnionScope()
+
+        @test Dagger.scope(worker=wid1) ==
+              Dagger.scope(workers=[wid1]) ==
+              ProcessScope(wid1)
+        @test Dagger.scope(workers=[wid1,wid2]) == UnionScope([ProcessScope(wid1),
+                                                               ProcessScope(wid2)])
+        @test Dagger.scope(workers=[]) == UnionScope()
+
+        @test Dagger.scope(thread=1) ==
+              Dagger.scope(threads=[1]) ==
+              UnionScope([ExactScope(Dagger.ThreadProc(w,1)) for w in procs()])
+        @test Dagger.scope(threads=[1,2]) == UnionScope([ExactScope(Dagger.ThreadProc(w,t)) for t in [1,2] for w in procs()])
+        @test Dagger.scope(threads=[]) == UnionScope()
+
+        @test Dagger.scope(worker=wid1,thread=1) ==
+              Dagger.scope(thread=1,worker=wid1) ==
+              Dagger.scope(workers=[wid1],thread=1) ==
+              Dagger.scope(worker=wid1,threads=[1]) ==
+              Dagger.scope(workers=[wid1],threads=[1]) ==
+              ExactScope(Dagger.ThreadProc(wid1,1))
+
+        @test_throws ArgumentError Dagger.scope(blah=1)
+        @test_throws ArgumentError Dagger.scope(thread=1, blah=1)
+
+        @test Dagger.scope(worker=1,thread=1) ==
+              Dagger.scope((worker=1,thread=1)) ==
+              Dagger.scope(((worker=1,thread=1),))
+        @test Dagger.scope((worker=1,thread=1),(worker=wid1,thread=2)) ==
+              Dagger.scope(((worker=1,thread=1),(worker=wid1,thread=2),)) ==
+              Dagger.scope(((worker=1,thread=1),), ((worker=wid1,thread=2),)) ==
+              UnionScope([ExactScope(Dagger.ThreadProc(1, 1)),
+                          ExactScope(Dagger.ThreadProc(wid1, 2))])
+        @test_throws ArgumentError Dagger.scope((;blah=1))
+        @test_throws ArgumentError Dagger.scope((thread=1, blah=1))
+
+        @testset "custom handler" begin
+            @eval begin
+                Dagger.scope_key_precedence(::Val{:gpu}) = 1
+                Dagger.scope_key_precedence(::Val{:rocm}) = 2
+                Dagger.scope_key_precedence(::Val{:cuda}) = 2
+
+                # Some fake scopes to use as sentinels
+                Dagger.to_scope(::Val{:gpu}, sc::NamedTuple) = ExactScope(Dagger.ThreadProc(1, sc.device))
+                Dagger.to_scope(::Val{:rocm}, sc::NamedTuple) = ExactScope(Dagger.ThreadProc($wid1, sc.gpu))
+                Dagger.to_scope(::Val{:cuda}, sc::NamedTuple) = ExactScope(Dagger.ThreadProc($wid2, sc.gpu))
+            end
+
+            @test Dagger.scope(gpu=1,device=2) ==
+                  Dagger.scope(device=2,gpu=1) ==
+                  Dagger.scope(gpu=1,device=2,blah=3) ==
+                  Dagger.scope((gpu=1,device=2,blah=3)) ==
+                  ExactScope(Dagger.ThreadProc(1, 2))
+            @test Dagger.scope((gpu=1,device=2),(worker=1,thread=1)) ==
+                  Dagger.scope((worker=1,thread=1),(device=2,gpu=1)) ==
+                  UnionScope([ExactScope(Dagger.ThreadProc(1, 2)),
+                              ExactScope(Dagger.ThreadProc(1, 1))])
+            @test Dagger.scope((gpu=1,device=2),(device=3,gpu=1)) ==
+                  UnionScope([ExactScope(Dagger.ThreadProc(1, 2)),
+                              ExactScope(Dagger.ThreadProc(1, 3))])
+
+            @test Dagger.scope(rocm=1,gpu=2) ==
+                  Dagger.scope(gpu=2,rocm=1) ==
+                  ExactScope(Dagger.ThreadProc(wid1, 2))
+            @test Dagger.scope(cuda=1,gpu=2) ==
+                  Dagger.scope(gpu=2,cuda=1) ==
+                  ExactScope(Dagger.ThreadProc(wid2, 2))
+            @test_throws ArgumentError Dagger.scope(rocm=1,cuda=1,gpu=2)
+            @test_throws ArgumentError Dagger.scope(gpu=2,rocm=1,cuda=1)
+            @test_throws ArgumentError Dagger.scope((rocm=1,cuda=1,gpu=2))
+        end
+    end
+
     rmprocs([wid1, wid2])
 end