Enhance task affinity support in Dagger.jl

AkhilAkkapelli · AkhilAkkapelli · commit 6ca66cddd16c · 2025-06-18T22:50:38.000+05:30
- Updated documentation for the `@spawn` macro to clarify the usage of `scope`, `compute_scope`, and `result_scope`, including examples with the new syntax.
- Improved error messages in the scheduling logic to provide clearer feedback when scopes are incompatible.
- Refactored test cases for task affinity to ensure they align with the new scope handling and provide better coverage for edge cases.
- Removed deprecated comments and cleaned up the code for better readability.
diff --git a/docs/src/task-affinity.md b/docs/src/task-affinity.md
@@ -1,123 +1,127 @@
 # Task Affinity 
 
-Dagger.jl's `@spawn` macro allows precise control over task execution and result access using `scope`, `compute_scope`, and `result_scope`.
+Dagger.jl's `@spawn` macro allows precise control over task execution and result accessibility using `scope`, `compute_scope`, and `result_scope`, which specify various chunk scopes of the task.
+
+For more information on how these scopes work, see [Scopes](scopes.md#Scopes).
 
 ---
 
 ## Key Terms
 
 ### Scope  
-`scope` defines the general set of locations where a Dagger task can execute. If `compute_scope` and `result_scope` are not explicitly set, the task's `compute_scope` defaults to its `scope`, and its `result_scope` defaults to `AnyScope()`, meaning the result can be accessed by any processor. Execution occurs on any processor within the defined scope.
+`scope` defines the general set of locations where a Dagger task can execute. If `scope` is not explicitly set, the task runs within the `compute_scope`. If both `scope` and `compute_scope` both are unspecified, the task falls back to `DefaultScope()`, allowing it to run wherever execution is possible. Execution occurs on any worker within the defined scope.
 
 **Example:**
 ```julia
-g = Dagger.@spawn scope=ExactScope(Dagger.OSProc(3)) f(x,y)
+g = Dagger.@spawn scope=Dagger.scope(worker=3) f(x,y)
 ```
-Task `g` executes only on Processor 3. Its result can be accessed by any processor.
+Task `g` executes only on worker 3. Its result can be accessed by any worker.
 
 ---
 
 ### Compute Scope  
-`compute_scope` also specifies where a Dagger task can execute. The key difference is if both `compute_scope` and `scope` are provided, `compute_scope` takes precedence over `scope` for execution placement. If `result_scope` isn't specified, it defaults to `AnyScope()`, allowing the result to be accessed by any processor.
+Like `scope`,`compute_scope` also specifies where a Dagger task can execute. The key difference is if both `compute_scope` and `scope` are provided, `compute_scope` takes precedence over `scope` for execution placement. If neither is specified, the they default to `DefaultScope()`. 
 
 **Example:**
 ```julia
-g1 = Dagger.@spawn scope=ExactScope(Dagger.ThreadProc(2, 3)) compute_scope=Dagger.UnionScope(ExactScope(Dagger.ThreadProc(1, 2)), ExactScope(Dagger.ThreadProc(3, 1))) f(x,y)
-g2 = Dagger.@spawn compute_scope=Dagger.UnionScope(ExactScope(Dagger.ThreadProc(1, 2)), ExactScope(Dagger.ThreadProc(3, 1))) f(x,y)
+g1 = Dagger.@spawn scope=Dagger.scope(worker=2,thread=3) compute_scope=Dagger.scope((worker=1, thread=2), (worker=3, thread=1))  f(x,y)
+g2 = Dagger.@spawn compute_scope=Dagger.scope((worker=1, thread=2), (worker=3, thread=1)) f(x,y)
 ```
-Task `g1` and `g2` execute on either thread 2 of processor 1, or thread 1 of processor 3. Their result can be accessed by any processor.
+Tasks `g1` and `g2` execute on either thread 2 of worker 1, or thread 1 of worker 3. The `scope` argument to `g1` is ignored. Their result can be accessed by any worker.
 
 ---
 
 ### Result Scope  
 
-`result_scope` restricts where a task's result can be fetched or moved. This is crucial for managing data locality and minimizing transfers. If only `result_scope` is specified, the `compute_scope` defaults to `Dagger.DefaultScope()`, meaning computation may happen on any processor.
+The result_scope limits the workers from which a task's result can be accessed. This is crucial for managing data locality and minimizing transfers. If `result_scope` is not specified, it defaults to `AnyScope()`, meaning the result can be accessed by any worker.
 
 **Example:**
 ```julia
-g = Dagger.@spawn result_scope=ExactScope(Dagger.OSProc(3)) f(x,y)
+g = Dagger.@spawn result_scope=Dagger.scope(worker=3, threads=[1,3, 4]) f(x,y)
 ```
-The result of `g` is accessible only from worker process 3. The task's execution may happen anywhere.
+The result of `g` is accessible only from threads 1, 3 and 4 of worker process 3. The task's execution may happen anywhere on threads 1, 3 and 4 of worker 3.
 
 ---
 
-## Interaction of compute_scope and result_scope
+## Interaction of `compute_scope` and `result_scope`
 
-When `scope`, `compute_scope`, and `result_scope` are all used, the scheduler executes the task on the intersection of the effective compute scope (which will be `compute_scope` if provided, otherwise `scope`) and the `result_scope`. If intersection does not exist then Scheduler throws Exception error.
+When `scope`, `compute_scope`, and `result_scope` are all used, the scheduler executes the task on the intersection of the effective compute scope (which will be `compute_scope` if provided, otherwise `scope`) and the `result_scope`. If the intersection is empty then the scheduler throws a `Dagger.Sch.SchedulerException` error.
 
 **Example:**
 ```julia
-g = Dagger.@spawn scope=ExactScope(Dagger.ThreadProc(3, 2)) compute_scope=Dagger.ProcessScope(2) result_scope=Dagger.UnionScope(ExactScope(Dagger.ThreadProc(2, 2)), ExactScope(Dagger.ThreadProc(4, 2))) f(x,y)
+g = Dagger.@spawn scope=Dagger.scope(worker=3,thread=2) compute_scope=Dagger.scope(worker=2) result_scope=Dagger.scope((worker=2, thread=2), (worker=4, thread=2)) f(x,y)
 ```
-The task `g` computes on `Dagger.ThreadProc(2, 2)` (as it's the intersection of compute and result scopes), and its result access is also restricted to `Dagger.ThreadProc(2, 2)`.
+The task `g` computes on thread 2 of worker 2 (as it's the intersection of compute and result scopes), and its result access is also restricted to thread 2 of worker 2.
 
 ---
 
 ## Chunk Inputs to Tasks  
 
-This section explains how `scope`, `compute_scope`, and `result_scope` affect tasks when a `Chunk` is the primary input to `@spawn` (e.g., `Dagger.tochunk(...)`).
+This section explains how `scope`, `compute_scope`, and `result_scope` affect tasks when a `Chunk` is the primary input to `@spawn` (e.g. created via `Dagger.tochunk(...)` or by calling `fetch(task; raw=true)` on a task).
 
-Assume `g` is some function, e.g., `g(x, y) = x * 2 + y * 3` and . `chunk_proc` is the chunk's processor, and `chunk_scope` is its defined accessibility.
+Assume `g` is some function, e.g. `g(x, y) = x * 2 + y * 3`, `chunk_proc` is the chunk's processor, and `chunk_scope` is its defined accessibility.
 
 When `Dagger.tochunk(...)` is directly spawned:
 - The task executes on `chunk_proc`.
 - The result is accessible only within `chunk_scope`.
 - This behavior occurs irrespective of the `scope`, `compute_scope`, and `result_scope` values provided in the `@spawn` macro.
-- Dagger validates that there is an intersection between the effective `compute_scope` (derived from `@spawn`'s `compute_scope` or `scope`) and the `result_scope`. If no intersection exists, the Scheduler throws an exception.
+- Dagger validates that there is an intersection between the effective `compute_scope` (derived from `@spawn`'s `compute_scope` or `scope`) and the `result_scope`. If no intersection exists, the scheduler throws an exception.
+
+!!! info While `chunk_proc` is currently required when constructing a chunk, it is largely unused in actual scheduling logic. It exists primarily for backward compatibility and may be deprecated in the future.
 
 **Usage:**
 ```julia
-h1 = Dagger.@spawn scope=ExactScope(Dagger.OSProc(3)) Dagger.tochunk(g(10, 11), chunk_proc, chunk_scope)
-h2 = Dagger.@spawn compute_scope=Dagger.UnionScope(ExactScope(Dagger.ThreadProc(1, 2)), ExactScope(Dagger.ThreadProc(3, 1))) Dagger.tochunk(g(20, 21), chunk_proc, chunk_scope)
-h3 = Dagger.@spawn scope=ExactScope(Dagger.ThreadProc(2, 3)) compute_scope=Dagger.UnionScope(ExactScope(Dagger.ThreadProc(1, 2)), ExactScope(Dagger.ThreadProc(3, 1))) Dagger.tochunk(g(30, 31), chunk_proc, chunk_scope)
-h4 = Dagger.@spawn result_scope=ExactScope(Dagger.OSProc(3)) Dagger.tochunk(g(40, 41), chunk_proc, chunk_scope)
-h5 = Dagger.@spawn scope=ExactScope(Dagger.ThreadProc(3, 2)) compute_scope=Dagger.ProcessScope(2) result_scope=Dagger.UnionScope(ExactScope(Dagger.ThreadProc(2, 2)), ExactScope(Dagger.ThreadProc(4, 2))) Dagger.tochunk(g(50, 51), chunk_proc, chunk_scope)
+h1 = Dagger.@spawn scope=Dagger.scope(worker=3) Dagger.tochunk(g(10, 11), chunk_proc, chunk_scope)
+h2 = Dagger.@spawn compute_scope=Dagger.scope((worker=1, thread=2), (worker=3, thread=1)) Dagger.tochunk(g(20, 21), chunk_proc, chunk_scope)
+h3 = Dagger.@spawn scope=Dagger.scope(worker=2,thread=3) compute_scope=Dagger.scope((worker=1, thread=2), (worker=3, thread=1)) Dagger.tochunk(g(30, 31), chunk_proc, chunk_scope)
+h4 = Dagger.@spawn result_scope=Dagger.scope(worker=3) Dagger.tochunk(g(40, 41), chunk_proc, chunk_scope)
+h5 = Dagger.@spawn scope=Dagger.scope(worker=3,thread=2) compute_scope=Dagger.ProcessScope(2) result_scope=Dagger.scope(worker=2,threads=[2,3]) Dagger.tochunk(g(50, 51), chunk_proc, chunk_scope)
 ```
-In all these cases (`h1` through `h5`), the task gets executed on `chunk_proc`, and its result is accessible only within `chunk_scope`.
+In all these cases (`h1` through `h5`), the tasks get executed on processor `chunk_proc` of chunk, and its result is accessible only within `chunk_scope`.
 
 ---
 
-## Function with Chunked Arguments as Tasks  
+## Function with Chunk Arguments as Tasks  
 
-This section details behavior when `scope`, `compute_scope`, and `result_scope` are used with tasks where a function is the input, and its arguments include `Chunks`.
+This section details behavior when `scope`, `compute_scope`, and `result_scope` are used with tasks where a function is the input, and its arguments include `Chunk`s.
 
-Assume `g(x, y) = x * 2 + y * 3` is a function, and `arg = Dagger.tochunk(g(1, 2), arg_proc, arg_scope)` is a chunked argument, where `arg_proc` is the chunk's processor and `arg_scope` is its defined scope.
+Assume `g(x, y) = x * 2 + y * 3` is a function, and `arg = Dagger.tochunk(g(1, 2), arg_proc, arg_scope)` is a chunk argument, where `arg_proc` is the chunk's processor and `arg_scope` is its defined scope.
 
 ### Scope  
-If `arg_scope` and `scope` do not intersect, the Scheduler throws an exception. Otherwise, `compute_scope` defaults to `scope`, and `result_scope` defaults to `AnyScope()`. Execution occurs on the intersection of `scope` and `arg_scope`.
+If `arg_scope` and `scope` do not intersect, the scheduler throws an exception. Execution occurs on the intersection of `scope` and `arg_scope`.
 
 ```julia
-h = Dagger.@spawn scope=ExactScope(Dagger.OSProc(3)) g(arg, 11)
+h = Dagger.@spawn scope=Dagger.scope(worker=3) g(arg, 11)
 ```
-Task `h` executes on any processor within the intersection of `scope` and `arg_scope`. The result is stored and accessible from anywhere.
+Task `h` executes on any worker within the intersection of `scope` and `arg_scope`. The result is accessible from any worker.
 
 ---
 
-### Compute Scope  
-If `arg_scope` and `compute_scope` do not intersect, the Scheduler throws an exception. Otherwise, execution happens on the intersection of the effective compute scope (which will be `compute_scope` if provided, otherwise `scope`) and `arg_scope`. `result_scope` defaults to `AnyScope()`.
+### Compute scope and Chunk argument scopes interaction 
+If `arg_scope` and `compute_scope` do not intersect, the scheduler throws an exception. Otherwise, execution happens on the intersection of the effective compute scope (which will be `compute_scope` if provided, otherwise `scope`) and `arg_scope`. `result_scope` defaults to `AnyScope()`.
 
 ```julia
-h1 = Dagger.@spawn compute_scope=Dagger.UnionScope(ExactScope(Dagger.ThreadProc(1, 2)), ExactScope(Dagger.ThreadProc(3, 1))) g(arg, 11)
-h2 = Dagger.@spawn scope=ExactScope(Dagger.ThreadProc(2, 3)) compute_scope=Dagger.UnionScope(ExactScope(Dagger.ThreadProc(1, 2)), ExactScope(Dagger.ThreadProc(3, 1))) g(arg, 21)
+h1 = Dagger.@spawn compute_scope=Dagger.scope((worker=1, thread=2), (worker=3, thread=1)) g(arg, 11)
+h2 = Dagger.@spawn scope=Dagger.scope(worker=2,thread=3) compute_scope=Dagger.scope((worker=1, thread=2), (worker=3, thread=1)) g(arg, 21)
 ```
-Task `h1` and `h2` execute on any processor within the intersection of the `compute_scope` and `arg_scope`. `scope` is ignored if `compute_scope` is specified. The result is stored and accessible from anywhere.
+Tasks `h1` and `h2` execute on any worker within the intersection of the `compute_scope` and `arg_scope`. `scope` is ignored if `compute_scope` is specified. The result is stored and accessible from anywhere.
 
 ---
 
-### Result Scope  
-If only `result_scope` is specified, computation happens on any processor within `arg_scope`, and the result is only accessible from `result_scope`.
+### Result scope and Chunk argument scopes interaction
+If only `result_scope` is specified, computation happens on any worker within `arg_scope`, and the result is only accessible from `result_scope`.
 
 ```julia
-h = Dagger.@spawn result_scope=ExactScope(Dagger.OSProc(3)) g(arg, 11)
+h = Dagger.@spawn result_scope=Dagger.scope(worker=3) g(arg, 11)
 ```
-Task `h` executes on any processor within `arg_scope`. The result is accessible from `result_scope`.
+Task `h` executes on any worker within `arg_scope`. The result is accessible from `result_scope`.
 
 ---
 
-### Compute and Result Scope  
-When `scope`, `compute_scope`, and `result_scope` are all used, the scheduler executes the task on the intersection of `arg_scope`, the effective compute scope (which is `compute_scope` if provided, otherwise `scope`), and `result_scope`. If no intersection exists, the Scheduler throws an exception.
+### Compute, result, and chunk argument scopes interaction  
+When `scope`, `compute_scope`, and `result_scope` are all used, the scheduler executes the task on the intersection of `arg_scope`, the effective compute scope (which is `compute_scope` if provided, otherwise `scope`), and `result_scope`. If no intersection exists, the scheduler throws an exception.
 
 ```julia
-h = Dagger.@spawn scope=ExactScope(Dagger.ThreadProc(3, 2)) compute_scope=Dagger.ProcessScope(2) result_scope=Dagger.UnionScope(ExactScope(Dagger.ThreadProc(2, 2)), ExactScope(Dagger.ThreadProc(4, 2))) g(arg, 31)
+h = Dagger.@spawn scope=Dagger.scope(worker=3,thread=2) compute_scope=Dagger.ProcessScope(2) result_scope=Dagger.scope((worker=2, thread=2), (worker=4, thread=2)) g(arg, 31)
 ```
-Task `h` computes on `Dagger.ThreadProc(2, 2)` (as it's the intersection of `arg`, `compute`, and `result` scopes), and its result access is also restricted to `Dagger.ThreadProc(2, 2)`.
+Task `h` computes on thread 2 of worker 2 (as it's the intersection of `arg`, `compute`, and `result` scopes), and its result access is also restricted to thread 2 of worker 2.
diff --git a/src/sch/Sch.jl b/src/sch/Sch.jl
@@ -728,7 +728,7 @@ function schedule!(ctx, state, procs=procs_to_use(ctx))
         # Calculate scope
         scope = constrain(task.compute_scope, task.result_scope)
         if scope isa InvalidScope
-            ex = SchedulingException("Compute and Result Scopes are not compatible: $(scope.x), $(scope.y)")
+            ex = SchedulingException("compute_scope and result_scope are not compatible: $(scope.x), $(scope.y)")
             state.cache[task] = ex
             state.errored[task] = true
             set_failed!(state, task)
@@ -737,21 +737,14 @@ function schedule!(ctx, state, procs=procs_to_use(ctx))
         if task.f isa Chunk
             scope = constrain(scope, task.f.scope)
             if scope isa InvalidScope
-                ex = SchedulingException("Compute and Chunk Scopes are not compatible: $(scope.x), $(scope.y)")
+                ex = SchedulingException("Current scope and function Chunk Scope are not compatible: $(scope.x), $(scope.y)")
                 state.cache[task] = ex
                 state.errored[task] = true
                 set_failed!(state, task)
                 @goto pop_task
             end
         end
 
-        # if task.options.proclist !== nothing
-        #     # proclist overrides scope selection
-        #     AnyScope()
-        # else
-        #     DefaultScope()
-        # end
-
         for (_,input) in task.inputs
             input = unwrap_weak_checked(input)
             chunk = if istask(input)
@@ -764,7 +757,7 @@ function schedule!(ctx, state, procs=procs_to_use(ctx))
             chunk isa Chunk || continue
             scope = constrain(scope, chunk.scope)
             if scope isa InvalidScope
-                ex = SchedulingException("Final Compute and Argument Chunk Scopes are not compatible: $(scope.x), $(scope.y)")
+                ex = SchedulingException("Current scope and argument Chunk scope are not compatible: $(scope.x), $(scope.y)")
                 state.cache[task] = ex
                 state.errored[task] = true
                 set_failed!(state, task)
@@ -1721,8 +1714,6 @@ function do_task(to_proc, task_desc)
         RemoteException(myid(), CapturedException(ex, bt))
     end
 
-    # @dagdebug thunk_id :scope "Result scope is $result_scope"
-
     threadtime = cputhreadtime() - threadtime_start
     # FIXME: This is not a realistic measure of max. required memory
     #gc_allocd = min(max(UInt64(Base.gc_num().allocd) - UInt64(gcnum_start.allocd), UInt64(0)), UInt64(1024^4))
diff --git a/src/thunk.jl b/src/thunk.jl
@@ -93,11 +93,6 @@ mutable struct Thunk
                    propagates=(),
                    kwargs...
                   )
-        # if !isa(f, Chunk) && (!isnothing(processor) || !isnothing(scope))
-        #     f = tochunk(f,
-        #                 something(processor, OSProc()),
-        #                 something(scope, DefaultScope()))
-        # end
 
         xs = Base.mapany(identity, xs)
         syncdeps_set = Set{Any}(filterany(is_task_or_chunk, Base.mapany(last, xs)))
@@ -481,17 +476,6 @@ function spawn(f, args...; kwargs...)
         args = args[2:end]
     end
 
-    # Wrap f in a Chunk if necessary
-    # processor = haskey(options, :processor) ? options.processor : nothing
-    # compute_scope = haskey(options, :compute_scope) ? options.compute_scope : (haskey(options, :scope) ? options.scope : nothing)
-    # result_scope = haskey(options, :result_scope) ? options.result_scope : nothing
-    
-    # if !isnothing(processor) || !isnothing(scope)
-    #     f = tochunk(f,
-    #                 something(processor, get_options(:processor, OSProc())),
-    #                 something(scope, get_options(:scope, DefaultScope())))
-    # end
-
     # Process the args and kwargs into Pair form
     args_kwargs = args_kwargs_to_pairs(args, kwargs)
 
diff --git a/test/task-affinity.jl b/test/task-affinity.jl
