Sch: Check iscompatible_* in can_use_proc

Author: jpsamaroo

src/sch/Sch.jl

@@ -774,7 +774,7 @@ function schedule!(ctx, state, procs=procs_to_use(ctx))
 
         for proc in local_procs
             gproc = get_parent(proc)
-            can_use, scope = can_use_proc(task, gproc, proc, opts, scope)
+            can_use, scope = can_use_proc(state, task, gproc, proc, opts, scope)
             if can_use
                 has_cap, est_time_util, est_alloc_util, est_occupancy =
                     has_capacity(state, proc, gproc.pid, opts.time_util, opts.alloc_util, opts.occupancy, sig)
@@ -806,7 +806,7 @@ function schedule!(ctx, state, procs=procs_to_use(ctx))
         cap, extra_util = nothing, nothing
         procs_found = false
         # N.B. if we only have one processor, we need to select it now
-        can_use, scope = can_use_proc(task, entry.gproc, entry.proc, opts, scope)
+        can_use, scope = can_use_proc(state, task, entry.gproc, entry.proc, opts, scope)
         if can_use
             has_cap, est_time_util, est_alloc_util, est_occupancy =
                 has_capacity(state, entry.proc, entry.gproc.pid, opts.time_util, opts.alloc_util, opts.occupancy, sig)
@@ -832,7 +832,7 @@ function schedule!(ctx, state, procs=procs_to_use(ctx))
                 @goto pop_task
             end
 
-            can_use, scope = can_use_proc(task, entry.gproc, entry.proc, opts, scope)
+            can_use, scope = can_use_proc(state, task, entry.gproc, entry.proc, opts, scope)
             if can_use
                 has_cap, est_time_util, est_alloc_util, est_occupancy =
                     has_capacity(state, entry.proc, entry.gproc.pid, opts.time_util, opts.alloc_util, opts.occupancy, sig)

src/sch/util.jl

@@ -322,7 +322,7 @@ function signature(f, args)
     return sig
 end
 
-function can_use_proc(task, gproc, proc, opts, scope)
+function can_use_proc(state, task, gproc, proc, opts, scope)
     # Check against proclist
     if opts.proclist !== nothing
         @warn "The `proclist` option is deprecated, please use scopes instead\nSee https://juliaparallel.org/Dagger.jl/stable/scopes/ for details" maxlog=1
@@ -369,6 +369,24 @@ function can_use_proc(task, gproc, proc, opts, scope)
         return false, scope
     end
 
+    # Check against f/args
+    Tf = chunktype(task.f)
+    if !Dagger.iscompatible_func(proc, opts, Tf)
+        @dagdebug task :scope "Rejected $proc: Not compatible with function type ($Tf)"
+        return false, scope
+    end
+    for (_, arg) in task.inputs
+        arg = unwrap_weak_checked(arg)
+        if arg isa Thunk
+            arg = state.cache[arg]
+        end
+        Targ = chunktype(arg)
+        if !Dagger.iscompatible_arg(proc, opts, Targ)
+            @dagdebug task :scope "Rejected $proc: Not compatible with argument type ($Targ)"
+            return false, scope
+        end
+    end
+
     @label accept
 
     @dagdebug task :scope "Accepted $proc"