Rewrite asyncmap with one task per work item

Author: richiejp

base/asyncmap.jl

@@ -2,6 +2,8 @@
 
 using Base.Iterators: Enumerate
 
+struct AbortMapException <: Exception end
+
 """
     asyncmap(f, c...; ntasks=0, batch_size=nothing)
 
@@ -104,30 +106,120 @@ julia> asyncmap(batch_func, 1:5; ntasks=2, batch_size=2)
     worker invocation, etc. See [`Threads`](@ref) and `Distributed` for alternatives.
 
 """
-function asyncmap(f, c...; ntasks=0, batch_size=nothing)
-    return async_usemap(f, c...; ntasks=ntasks, batch_size=batch_size)
-end
+asyncmap(f, c...; kwargs...) = do_asyncmap(f, c...; kwargs...)
 
-function async_usemap(f, c...; ntasks=0, batch_size=nothing)
-    ntasks = verify_ntasks(c[1], ntasks)
+function do_asyncmap(f, c...; ntasks=0, batch_size=nothing)
+    ntasks = let n = verify_ntasks(c[1], ntasks)
+        n isa Number ? () -> n : n
+    end
     batch_size = verify_batch_size(batch_size)
 
-    if batch_size !== nothing
-        exec_func = batch -> begin
-            # extract the Refs from the input tuple
-            batch_refs = map(x->x[1], batch)
+    (tasks, echnl, res) = if batch_size === nothing
+        do_asyncmap(f, c, ntasks)
+    else
+        do_asyncmap(f, c, ntasks, batch_size)
+    end
 
-            # and the args tuple....
-            batched_args = map(x->x[2], batch)
+    foreach(_wait, tasks)
+    close(echnl)
 
-            results = f(batched_args)
-            foreach(x -> (batch_refs[x[1]].x = x[2]), enumerate(results))
-        end
+    isready(echnl) &&
+        throw(CompositeException(TaskFailedException.(echnl)))
+
+    if res isa Ref
+        # scalar case: map(f, x...) = f(x...)
+        res.x
     else
-        exec_func = (r,args) -> (r.x = f(args...))
+        map(ref -> ref.x, res)
+    end
+end
+
+mutable struct AsyncmapState
+    tasks::Vector{Task}
+    ntasks::Function
+    concur::Int
+    echnl::Channel{Task}
+end
+
+AsyncmapState(ntasks) =
+    AsyncmapState([], ntasks, 0, Channel{Task}(typemax(Int)))
+
+function do_asyncmap_task!(f::Function, s::AsyncmapState, xs...)
+    t = @task try
+        f(xs...)
+    catch
+        put!(s.echnl, current_task())
+        rethrow()
+    finally
+        s.concur -= 1
+    end
+
+    while s.concur >= s.ntasks() && !isready(s.echnl)
+        yield()
+    end
+
+    isready(s.echnl) && throw(AbortMapException());
+
+    s.concur += 1
+    schedule(t)
+    push!(s.tasks, t)
+end
+
+function do_asyncmap(f, c::Tuple, ntasks::Function)
+    s = AsyncmapState(ntasks)
+
+    res = try
+        map(c...) do x...
+            r = Ref{Any}(nothing)
+
+            do_asyncmap_task!(s) do
+                r.x = f(x...)
+            end
+
+            r
+        end
+    catch ex
+        ex isa AbortMapException || rethrow()
     end
-    chnl, err_chnl, worker_tasks = setup_chnl_and_tasks(exec_func, ntasks, batch_size)
-    return wrap_n_exec_twice(chnl, err_chnl, worker_tasks, ntasks, exec_func, c...)
+
+    s.tasks, s.echnl, res
+end
+
+function do_asyncmap(f, c::Tuple, ntasks::Function, batch_size)
+    s = AsyncmapState(ntasks)
+    local batch_in
+    local batch_out
+
+    function new_batch()
+        batch_in = sizehint!(Vector{Any}(), batch_size)
+        batch_out = sizehint!(Vector{Ref{Any}}(), batch_size)
+    end
+
+    exec_batch(args, res) = for (i, x) in enumerate(f(args))
+        res[i].x = x
+    end
+
+    new_batch()
+    res = try
+        map(c...) do x
+            r = Ref{Any}(undef)
+
+            push!(batch_in, x)
+            push!(batch_out, r)
+
+            if length(batch_in) >= batch_size
+                do_asyncmap_task!(exec_batch, s, batch_in, batch_out)
+                new_batch()
+            end
+
+            r
+        end
+    catch ex
+        ex isa AbortMapException || rethrow()
+    end
+    do_asyncmap_task!(exec_batch, s, batch_in, batch_out)
+
+    s.tasks, s.echnl, res
 end
 
 batch_size_err_str(batch_size) = string("batch_size must be specified as a positive integer. batch_size=", batch_size)
@@ -161,63 +253,6 @@ function verify_ntasks(iterable, ntasks)
     return ntasks
 end
 
-function wrap_n_exec_twice(chnl, err_chnl, worker_tasks, ntasks, exec_func, c...)
-    # The driver task, creates a Ref object and writes it and the args tuple to
-    # the communication channel for processing by a free worker task.
-    push_arg_to_channel = (x...) -> (r=Ref{Any}(nothing); put!(chnl,(r,x));r)
-
-    if isa(ntasks, Function)
-        map_f = (x...) -> begin
-            # check number of tasks every time, and start one if required.
-            # number_tasks > optimal_number is fine, the other way around is inefficient.
-            if length(worker_tasks) < ntasks()
-                start_worker_task!(worker_tasks, exec_func, chnl, err_chnl)
-            end
-            push_arg_to_channel(x...)
-        end
-    else
-        map_f = push_arg_to_channel
-    end
-    maptwice(map_f, chnl, err_chnl, worker_tasks, c...)
-end
-
-function maptwice(wrapped_f, chnl, err_chnl, worker_tasks, c...)
-    # first run, returns a collection of Refs
-    asyncrun = try
-        map(wrapped_f, c...)
-    catch ex
-        # the work channel was closed early and we have the error
-        ex isa InvalidStateException && isready(err_chnl) || rethrow()
-    finally
-        close(chnl)
-    end
-
-    # wait for all worker tasks to finish. We try not to throw
-    # task errors here because they would be in task creation
-    # order.
-    try
-        @sync foreach(t -> @sync_add(t), worker_tasks)
-    catch
-        isready(err_chnl) || rethrow()
-    finally
-        close(err_chnl)
-    end
-
-    # throw task errors in chronological order
-    if isready(err_chnl)
-        exs = [TaskFailedException(task) for task in err_chnl]
-        throw(CompositeException(exs))
-    end
-
-    if isa(asyncrun, Ref)
-        # scalar case
-        return asyncrun.x
-    else
-        # second run, extract values from the Refs and return
-        return map(ref->ref.x, asyncrun)
-    end
-end
-
 function setup_chnl_and_tasks(exec_func, ntasks, batch_size=nothing)
     if isa(ntasks, Function)
         nt = ntasks()
@@ -279,7 +314,7 @@ end
 
 # map on a single BitArray returns a BitArray if the mapping function is boolean.
 function asyncmap(f, b::BitArray; kwargs...)
-    b2 = async_usemap(f, b; kwargs...)
+    b2 = do_asyncmap(f, b; kwargs...)
     if eltype(b2) == Bool
         return BitArray(b2)
     end

test/asyncmap.jl

@@ -6,10 +6,10 @@ using Random
 @test allunique(asyncmap(x->(sleep(1.0);objectid(current_task())), 1:10))
 
 # num tasks
-@test length(unique(asyncmap(x->(yield();objectid(current_task())), 1:20; ntasks=5))) == 5
+@test length(unique(asyncmap(x->(yield();objectid(current_task())), 1:20; ntasks=5))) == 20
 
 # default num tasks
-@test length(unique(asyncmap(x->(yield();objectid(current_task())), 1:200))) == 100
+@test length(unique(asyncmap(x->(yield();objectid(current_task())), 1:200))) == 200
 
 # ntasks as a function
 let nt=0
@@ -18,7 +18,7 @@ let nt=0
                                            # nt_func() will be called initially once and then for every
                                            # iteration
 end
-@test length(unique(asyncmap(x->(yield();objectid(current_task())), 1:200; ntasks=nt_func))) == 7
+@test length(unique(asyncmap(x->(yield();objectid(current_task())), 1:200; ntasks=nt_func))) == 200
 
 # batch mode tests
 let ctr=0