KeyErrors running models in parallel with Distributed

[pauljmyers]

I've recently been trying to run models generated from Catalyst reaction networks in parallel using Distributed and keep running into a KeyError when using pmap directly or using the EnsembleDistributed interface. First noticed this trying to use Michaelis-Menten rate laws but have run into the same problem for reaction networks with standard mass-action, non-MM rate laws. This code

using Distributed
if nprocs() == 1
    addprocs(Threads.nthreads(), exeflags="--project=.")
end
@everywhere using DifferentialEquations, Catalyst

## Example model using built-in Michaelis-Menten rate function:
model1 = @reaction_network begin
    mm(X,v,K), ∅ → X
end v K
tspan = (0.0,1.0)
u0 = [1.0]
odeprob1 = ODEProblem(model1, u0, tspan, ones(2))
ps1 = [ones(2), ones(2), ones(2), ones(2), ones(2)].*collect(1:5)
batchsol1 = pmap(p->solve(odeprob1, Tsit5(), p=p), ps1)

## Same thing happens with EnsembleProblem:
@everywhere prob_func(prob, i, repeat) = remake(prob; p=ps1[i])
ensemble_prob = EnsembleProblem(odeprob1, prob_func=prob_func)
batchsol2 = solve(ensemble_prob, Tsit5(), EnsembleDistributed(); trajectories=length(ps1))  # distributed errors out
batchsol3 = solve(ensemble_prob, Tsit5(), EnsembleThreads(); trajectories=length(ps1)) # but multithreading works fine

produces the following error and stacktrace for pmap and EnsembleDistributed (but not EnsembleThreads):

ERROR: On worker 4:
KeyError: key (0xf66753a6, 0x8168c51d, 0xfa5c35ea, 0x941a4aa8, 0x39949471) not found
Stacktrace:
[1] getindex
@ .\dict.jl:482
[2] #8
@ C:\Users\pmyer.julia\packages\RuntimeGeneratedFunctions\3SZ1T\src\RuntimeGeneratedFunctions.jl:177
[3] lock
@ .\lock.jl:187
[4] _lookup_body
@ C:\Users\pmyer.julia\packages\RuntimeGeneratedFunctions\3SZ1T\src\RuntimeGeneratedFunctions.jl:175
[5] generated_callfunc_body
@ C:\Users\pmyer.julia\packages\RuntimeGeneratedFunctions\3SZ1T\src\RuntimeGeneratedFunctions.jl:120
[6] #s1#3
@ C:\Users\pmyer.julia\packages\RuntimeGeneratedFunctions\3SZ1T\src\RuntimeGeneratedFunctions.jl:208 [inlined]
[7] #s1#3
@ .\none:0
[8] GeneratedFunctionStub
@ .\boot.jl:571
[9] RuntimeGeneratedFunction
@ C:\Users\pmyer.julia\packages\RuntimeGeneratedFunctions\3SZ1T\src\RuntimeGeneratedFunctions.jl:112 [inlined]
[10] f
@ C:\Users\pmyer.julia\packages\ModelingToolkit\yGWxM\src\systems\diffeqs\abstractodesystem.jl:228 [inlined]
[11] ODEFunction
@ C:\Users\pmyer.julia\packages\SciMLBase\cU5k7\src\scimlfunctions.jl:334 [inlined]
[12] initialize!
@ C:\Users\pmyer.julia\packages\OrdinaryDiffEq\UcJTX\src\perform_step\low_order_rk_perform_step.jl:623
[13] #__init#467
@ C:\Users\pmyer.julia\packages\OrdinaryDiffEq\UcJTX\src\solve.jl:456
[14] #__solve#466
@ C:\Users\pmyer.julia\packages\OrdinaryDiffEq\UcJTX\src\solve.jl:4 [inlined]
[15] #solve_call#58
@ C:\Users\pmyer.julia\packages\DiffEqBase\niZxn\src\solve.jl:61 [inlined]
[16] solve_call
@ C:\Users\pmyer.julia\packages\DiffEqBase\niZxn\src\solve.jl:48 [inlined]
[17] #solve_up#60
@ C:\Users\pmyer.julia\packages\DiffEqBase\niZxn\src\solve.jl:85 [inlined]
[18] solve_up
@ C:\Users\pmyer.julia\packages\DiffEqBase\niZxn\src\solve.jl:78 [inlined]
[19] #solve#59
@ C:\Users\pmyer.julia\packages\DiffEqBase\niZxn\src\solve.jl:73
[20] #27
@ g:\my_script_to_test_pmap.jl:21
[21] #106
@ C:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.6\Distributed\src\process_messages.jl:278
[22] run_work_thunk
@ C:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.6\Distributed\src\process_messages.jl:63
[23] macro expansion
@ C:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.6\Distributed\src\process_messages.jl:278 [inlined]
[24] #105
@ .\task.jl:411

Stacktrace:

1. (::Base.var"Reorganisations/renamings to the docs #837#839")(x::Task) at .\asyncmap.jl
2. foreach(f::Base.var"Reorganisations/renamings to the docs #837#839", itr::Vector{Any}) at .\abstractarray.jl
3. maptwice(wrapped_f::Function, chnl::Channel{Any}, worker_tasks::Vector{Any}, c::Vector{Vector{Float64}}) at .\asyncmap.jl
4. wrap_n_exec_twice at .\asyncmap.jl
5. #async_usemap#822 at .\asyncmap.jl
6. async_usemap at .\asyncmap.jl
7. #asyncmap#821 at .\asyncmap.jl
8. asyncmap at .\asyncmap.jl
9. pmap(f::Function, p::WorkerPool, c::Vector{Vector{Float64}}; distributed::Bool, batch_size::Int64, on_error::Nothing, retry_delays::Vector{Any}, retry_check::Nothing) at C:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.6\Distributed\src\pmap.jl
10. pmap(f::Function, p::WorkerPool, c::Vector{Vector{Float64}}) at C:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.6\Distributed\src\pmap.jl
11. pmap(f::Function, c::Vector{Vector{Float64}}; kwargs::Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}) at C:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.6\Distributed\src\pmap.jl
12. pmap(f::Function, c::Vector{Vector{Float64}}) at C:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.6\Distributed\src\pmap.jl

This was run on Julia v1.6.2.

[ChrisRackauckas]

SciML/RuntimeGeneratedFunctions.jl#37 fixes distributed RGFs. Thanks!