fix multithreaded thread safety and clean up

Author: ChrisRackauckas

src/ensemble/basic_ensemble_solve.jl

@@ -174,56 +174,27 @@ function solve_batch(prob,alg,ensemblealg::EnsembleThreads,II,pmap_batch_size;kw
     probs = prob.prob
   end
 
-  function multithreaded_batch(batch_idx)
-    i = II[batch_idx]
-    iter = 1
-    _prob = if prob.safetycopy
-      probs isa Vector ? deepcopy(probs[Threads.threadid()]) : probs
-    else
-      probs isa Vector ? probs[Threads.threadid()] : probs
-    end
-    new_prob = prob.prob_func(_prob,i,iter)
-    x = prob.output_func(solve(new_prob,alg;kwargs...),i)
-    if !(typeof(x) <: Tuple)
-        @warn("output_func should return (out,rerun). See docs for updated details")
-        _x = (x,false)
-    else
-      _x = x
-    end
-    rerun = _x[2]
-
-    while rerun
-        iter += 1
-        _prob = if prob.safetycopy
-          probs isa Vector ? deepcopy(probs[Threads.threadid()]) : probs
-        else
-          probs isa Vector ? probs[Threads.threadid()] : probs
-        end
-        new_prob = prob.prob_func(_prob,i,iter)
-        x = prob.output_func(solve(new_prob,alg;alias_jumps=true,kwargs...),i)
-        if !(typeof(x) <: Tuple)
-            @warn("output_func should return (out,rerun). See docs for updated details")
-            _x = (x,false)
-        else
-          _x = x
-        end
-        rerun = _x[2]
-    end
-    _x[1]
-  end
-
   #batch_data = Vector{Core.Compiler.return_type(multithreaded_batch,Tuple{typeof(first(II))})}(undef,length(II))
-  batch_data = Vector{Any}(undef,length(II))
 
+  local batch_data
   let
     if length(II) == 1 || Threads.nthreads() == 1
+      batch_data = Vector{Any}(undef,length(II))
       for batch_idx in axes(batch_data, 1)
-          batch_data[batch_idx] = multithreaded_batch(batch_idx)
+          batch_data[batch_idx] = multithreaded_batch(batch_idx,probs,alg,II)
       end
     else
-      Threads.@threads for batch_idx in axes(batch_data, 1)
-          batch_data[batch_idx] = multithreaded_batch(batch_idx)
+      batch_data = Vector{Any}(undef,Threads.nthreads())
+      batch_size = length(II)÷Threads.nthreads()
+      Threads.@threads for i in 1:Threads.nthreads()
+          if i == Threads.nthreads()
+            I_local = II[(batch_size*(i-1)+1):end]
+          else
+            I_local = II[(batch_size*(i-1)+1):(batch_size*i)]
+          end
+          batch_data[i] = solve_batch(prob,alg,EnsembleSerial(),I_local,pmap_batch_size;kwargs...)
       end
+      batch_data = reduce(vcat,batch_data)
     end
   end
   tighten_container_eltype(batch_data)
@@ -240,71 +211,8 @@ function solve_batch(prob,alg,::EnsembleSplitThreads,II,pmap_batch_size;kwargs..
       else
         I_local = II[(batch_size*(i-1)+1):(batch_size*i)]
       end
-      thread_monte(prob,I_local,alg,i;kwargs...)
+      solve_batch(prob,alg,EnsembleThreads(),I_local,pmap_batch_size;kwargs...)
     end
   end
   reduce(vcat,batch_data)
 end
-
-function thread_monte(prob,II,alg,procid;kwargs...)
-
-  if typeof(prob.prob) <: AbstractJumpProblem && length(II) != 1
-    probs = [deepcopy(prob.prob) for i in 1:Threads.nthreads()]
-  else
-    probs = prob.prob
-  end
-
-  function multithreaded_batch(j)
-    i = II[j]
-    iter = 1
-    _prob = if prob.safetycopy
-      probs isa Vector ? deepcopy(probs[Threads.threadid()]) : probs
-    else
-      probs isa Vector ? probs[Threads.threadid()] : probs
-    end
-    new_prob = prob.prob_func(_prob,i,iter)
-    rerun = true
-    x = prob.output_func(solve(new_prob,alg;alias_jumps=true,kwargs...),i)
-    if !(typeof(x) <: Tuple)
-        @warn("output_func should return (out,rerun). See docs for updated details")
-        _x = (x,false)
-    else
-      _x = x
-    end
-    rerun = _x[2]
-    while rerun
-        iter += 1
-        _prob = if prob.safetycopy
-          probs isa Vector ? deepcopy(probs[Threads.threadid()]) : probs
-        else
-          probs isa Vector ? probs[Threads.threadid()] : probs
-        end
-        new_prob = prob.prob_func(_prob,i,iter)
-        x = prob.output_func(solve(new_prob,alg;alias_jumps=true,kwargs...),i)
-        if !(typeof(x) <: Tuple)
-            @warn("output_func should return (out,rerun). See docs for updated details")
-            _x = (x,false)
-        else
-          _x = x
-        end
-        rerun = _x[2]
-    end
-    _x[1]
-  end
-
-  #batch_data = Vector{Core.Compiler.return_type(multithreaded_batch,Tuple{typeof(first(II))})}(undef,length(II))
-  batch_data = Vector{Any}(undef,length(II))
-
-  let
-    if length(II) == 1 || Threads.nthreads() == 1
-      for batch_idx in axes(batch_data, 1)
-          batch_data[batch_idx] = multithreaded_batch(batch_idx)
-      end
-    else
-      Threads.@threads for batch_idx in axes(batch_data, 1)
-          batch_data[batch_idx] = multithreaded_batch(batch_idx)
-      end
-    end
-  end
-  tighten_container_eltype(batch_data)
-end

test/downstream/ensemble_thread_safety.jl

@@ -0,0 +1,26 @@
+using OrdinaryDiffEq
+function f(du,u,p,t)
+    du[1] = 1.01*u[1]
+end
+u0 = [0.0, 0.0]
+tspan = (0.0,1.0)
+prob = ODEProblem(f,u0,tspan)
+n = 100
+
+initial_conditions = range(0, stop=1, length=n)
+function prob_func(prob,i,repeat)
+    prob.u0[1] = initial_conditions[i]
+    prob
+end
+ensemble_prob = EnsembleProblem(prob,prob_func=prob_func)
+sim_1 = solve(ensemble_prob,Tsit5()
+            ,EnsembleThreads(),
+            trajectories=100)
+sim_2 = solve(ensemble_prob,Tsit5()
+            ,EnsembleDistributed(),
+            trajectories=100)
+ss_sol_1 = hcat(collect(EnsembleAnalysis.get_timepoint(sim_1,0))...);
+ss_sol_2 = hcat(collect(EnsembleAnalysis.get_timepoint(sim_2,0))...);
+
+ss_sol_1[1,:] == initial_conditions
+ss_sol_2[1,:] == initial_conditions

test/runtests.jl

@@ -38,6 +38,7 @@ if !is_APPVEYOR && GROUP == "Downstream"
     @time @safetestset "Null Parameters" begin include("downstream/null_params_test.jl") end
     @time @safetestset "Ensemble Simulations" begin include("downstream/ensemble.jl") end
     @time @safetestset "Ensemble Analysis" begin include("downstream/ensemble_analysis.jl") end
+    @time @safetestset "Ensemble Thread Safety" begin include("downstream/ensemble_thread_safety.jl") end
     @time @safetestset "Inference Tests" begin include("downstream/inference.jl") end
     @time @safetestset "Default linsolve with structure" begin include("downstream/default_linsolve_structure.jl") end
     @time @safetestset "Callback Merging Tests" begin include("downstream/callback_merging.jl") end