Cancellation and test fixes

Author: jpsamaroo

src/cancellation.jl

@@ -100,19 +100,30 @@ function _cancel!(state, tid, force, graceful, halt_sch)
         @dagdebug tid :cancel "Cancelling ready task"
         ex = DTaskFailedException(task, task, InterruptException())
         Sch.store_result!(state, task, ex; error=true)
-        Sch.set_failed!(state, task)
+        Sch.finish_failed!(state, task, task)
+    end
+    if tid === nothing
+        empty!(state.ready)
+    else
+        idx = findfirst(t->t.id == tid, state.ready)
+        idx !== nothing && deleteat!(state.ready, idx)
     end
-    empty!(state.ready)
 
     # Cancel waiting tasks
     for task in keys(state.waiting)
         tid !== nothing && task.id != tid && continue
         @dagdebug tid :cancel "Cancelling waiting task"
         ex = DTaskFailedException(task, task, InterruptException())
         Sch.store_result!(state, task, ex; error=true)
-        Sch.set_failed!(state, task)
+        Sch.finish_failed!(state, task, task)
+    end
+    if tid === nothing
+        empty!(state.waiting)
+    else
+        if haskey(state.waiting, tid)
+            delete!(state.waiting, tid)
+        end
     end
-    empty!(state.waiting)
 
     # Cancel running tasks at the processor level
     wids = unique(map(root_worker_id, values(state.running_on)))
@@ -155,6 +166,7 @@ function _cancel!(state, tid, force, graceful, halt_sch)
             return
         end
     end
+    put!(state.chan, Sch.RescheduleSignal())
 
     if halt_sch
         unlock(state.lock)

src/sch/util.jl

@@ -587,15 +587,12 @@ end
         # TODO: Measure and model processor move overhead
         tx_cost = impute_sum(affinity(chunk)[2] for chunk in chunks_filt)
 
-        # Estimate total cost to move data and get task running after currently-scheduled tasks
-        est_business = get(state.worker_time_pressure[get_parent(proc).pid], proc, 0)
-
         # Add fixed cost for cross-worker task transfer (esimated at 1ms)
         # TODO: Actually estimate/benchmark this
         task_xfer_cost = gproc.pid != myid() ? 1_000_000 : 0 # 1ms
 
         # Compute final cost
-        costs[proc] = est_time_util + est_business + (tx_cost/tx_rate) + task_xfer_cost
+        costs[proc] = est_time_util + (tx_cost/tx_rate) + task_xfer_cost
     end
     chunks_cleanup()
 

test/scheduler.jl

@@ -400,13 +400,18 @@ end
         end
 
         state = Dagger.Sch.EAGER_STATE[]
-        tproc1 = Dagger.ThreadProc(1, 1)
-        tproc2 = Dagger.ThreadProc(first(workers()), 1)
-        procs = [tproc1, tproc2]
+        tproc1_1 = Dagger.ThreadProc(1, 1)
+        tproc2_1 = Dagger.ThreadProc(first(workers()), 1)
+        procs = [tproc1_1, tproc2_1]
 
-        pres1 = state.worker_time_pressure[1][tproc1]
-        pres2 = state.worker_time_pressure[first(workers())][tproc2]
+        # Ensure that this worker has been used at least once
+        fetch(Dagger.@spawn scope=Dagger.ExactScope(tproc2_1) 1+1)
+
+        #pres1_1 = state.worker_time_pressure[1][tproc1_1]
+        #pres2_1 = state.worker_time_pressure[first(workers())][tproc2_1]
         tx_rate = state.transfer_rate[]
+        tx_xfer_cost = 1e6
+        sig_unknown_cost = 1e9
 
         for (args, tx_size) in [
             ([1, 2], 0),
@@ -433,18 +438,18 @@ end
             Dagger.Sch.collect_task_inputs!(state, t)
             sorted_procs, costs = Dagger.Sch.estimate_task_costs(state, procs, t)
 
-            @test tproc1 in sorted_procs
-            @test tproc2 in sorted_procs
+            @test tproc1_1 in sorted_procs
+            @test tproc2_1 in sorted_procs
             if length(cargs) > 0
-                @test sorted_procs[1] == tproc1
-                @test sorted_procs[2] == tproc2
+                @test sorted_procs[1] == tproc1_1
+                @test sorted_procs[2] == tproc2_1
             end
 
-            @test haskey(costs, tproc1)
-            @test haskey(costs, tproc2)
-            @test costs[tproc1] ≈ pres1 # All chunks are local
+            @test haskey(costs, tproc1_1)
+            @test haskey(costs, tproc2_1)
+            @test costs[tproc1_1] ≈ #=pres1_1 +=# sig_unknown_cost # All chunks are local, and this signature is unknown
             if nprocs() > 1
-                @test costs[tproc2] ≈ (tx_size/tx_rate) + pres2 # All chunks are remote
+                @test costs[tproc2_1] ≈ (tx_size/tx_rate) + tx_xfer_cost + #=pres2_1 +=# sig_unknown_cost # All chunks are remote, and this signature is unknown
             end
         end
     end
@@ -564,12 +569,32 @@ end
 end
 
 @testset "Cancellation" begin
+    # Ready task cancellation
+    start_time = time_ns()
     t = Dagger.@spawn scope=Dagger.scope(worker=1, thread=1) sleep(100)
+    Dagger.cancel!(t)
+    @test timedwait(()->istaskdone(t), 10) == :ok
+    if istaskdone(t)
+        @test_throws_unwrap (Dagger.DTaskFailedException, InterruptException) fetch(t)
+        @test (time_ns() - start_time) * 1e-9 < 100
+    end
+
+    # Running task cancellation
     start_time = time_ns()
+    t = Dagger.@spawn scope=Dagger.scope(worker=1, thread=1) sleep(100)
+    sleep(0.1) # Give the scheduler a chance to schedule the task
     Dagger.cancel!(t)
-    @test_throws_unwrap (Dagger.DTaskFailedException, InterruptException) fetch(t)
+    @test timedwait(()->istaskdone(t), 10) == :ok
+    if istaskdone(t)
+        @test_throws_unwrap (Dagger.DTaskFailedException, InterruptException) fetch(t)
+        @test (time_ns() - start_time) * 1e-9 < 100
+    end
+
+    # Normal task execution
+    start_time = time_ns()
     t = Dagger.@spawn scope=Dagger.scope(worker=1, thread=1) yield()
-    fetch(t)
-    finish_time = time_ns()
-    @test (finish_time - start_time) * 1e-9 < 100
+    @test timedwait(()->istaskdone(t), 10) == :ok
+    if istaskdone(t)
+        @test (time_ns() - start_time) * 1e-9 < 100
+    end
 end