Distributed `AbstractMCMC.sample` fails

[sunxd3]

using Distributed

@everywhere begin
    using JuliaBUGS, LogDensityProblems, LogDensityProblemsAD, AbstractMCMC, AdvancedHMC, MCMCChains, ReverseDiff # also other packages one may need
    using Random
    Random.seed!(1234)
end

data = (
    r = [10, 23, 23, 26, 17, 5, 53, 55, 32, 46, 10, 8, 10, 8, 23, 0, 3, 22, 15, 32, 3],
    n = [39, 62, 81, 51, 39, 6, 74, 72, 51, 79, 13, 16, 30, 28, 45, 4, 12, 41, 30, 51, 7],
    x1 = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
    x2 = [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
    N = 21,
)

model_def = @bugs begin
    for i in 1:N
        r[i] ~ dbin(p[i], n[i])
        b[i] ~ dnorm(0.0, tau)
        p[i] = logistic(alpha0 + alpha1 * x1[i] + alpha2 * x2[i] + alpha12 * x1[i] * x2[i] + b[i])
    end
    alpha0 ~ dnorm(0.0, 1.0E-6)
    alpha1 ~ dnorm(0.0, 1.0E-6)
    alpha2 ~ dnorm(0.0, 1.0E-6)
    alpha12 ~ dnorm(0.0, 1.0E-6)
    tau ~ dgamma(0.001, 0.001)
    sigma = 1 / sqrt(tau)
end

model = compile(model_def, data)
ad_model = ADgradient(:ReverseDiff, model)

n_samples, n_adapts = 2000, 1000
D = LogDensityProblems.dimension(model)

n_chains = nprocs() > 1 ? nprocs() - 1 : 1 # use all the processes except the parent process
samples_and_stats = AbstractMCMC.sample(
    ad_model,
    AdvancedHMC.NUTS(0.65),
    AbstractMCMC.MCMCDistributed(),
    n_samples,
    n_chains;
    chain_type = Chains,
    n_adapts = n_adapts,
    init_params = [rand(D) for _ = 1:n_chains], # each chain has its own initial parameters
    discard_initial = n_adapts,
    progress = false, # Base.TTY creating problems in distributed setting
)

produce error

ERROR: TaskFailedException

    nested task error: On worker 2:
    UndefVarError: `###__compute_log_density__#230` not defined in `JuliaBUGS.Model`
    Suggestion: add an appropriate import or assignment. This global was declared but not assigned.
    Stacktrace:
      [1] deserialize_datatype
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Serialization/src/Serialization.jl:1471
      [2] handle_deserialize
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Serialization/src/Serialization.jl:897
      [3] deserialize
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Serialization/src/Serialization.jl:844
      [4] deserialize_datatype
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Serialization/src/Serialization.jl:1495
      [5] handle_deserialize
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Serialization/src/Serialization.jl:897
      [6] deserialize
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Serialization/src/Serialization.jl:844
      [7] deserialize_datatype
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Serialization/src/Serialization.jl:1495
      [8] handle_deserialize
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Serialization/src/Serialization.jl:897
      [9] deserialize
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Serialization/src/Serialization.jl:844
     [10] deserialize_datatype
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Serialization/src/Serialization.jl:1495
     [11] handle_deserialize
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Serialization/src/Serialization.jl:897
     [12] deserialize
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Serialization/src/Serialization.jl:844
     [13] deserialize_datatype
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Serialization/src/Serialization.jl:1495
     [14] handle_deserialize
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Serialization/src/Serialization.jl:897
     [15] deserialize
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Serialization/src/Serialization.jl:844
     [16] deserialize_datatype
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Serialization/src/Serialization.jl:1495
     [17] handle_deserialize
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Serialization/src/Serialization.jl:897
     [18] deserialize
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Serialization/src/Serialization.jl:844
     [19] handle_deserialize
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Serialization/src/Serialization.jl:904
     [20] deserialize
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Serialization/src/Serialization.jl:844
     [21] #5
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Serialization/src/Serialization.jl:1003
     [22] ntupleany
        @ ./ntuple.jl:43
     [23] deserialize_tuple
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Serialization/src/Serialization.jl:1003
     [24] handle_deserialize
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Serialization/src/Serialization.jl:887
     [25] deserialize
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Serialization/src/Serialization.jl:844
     [26] #5
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Serialization/src/Serialization.jl:1003
     [27] ntupleany
        @ ./ntuple.jl:43
     [28] deserialize_tuple
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Serialization/src/Serialization.jl:1003
     [29] handle_deserialize
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Serialization/src/Serialization.jl:887
     [30] deserialize
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Serialization/src/Serialization.jl:844 [inlined]
     [31] deserialize_msg
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Distributed/src/messages.jl:87
     [32] #invokelatest#2
        @ ./essentials.jl:1055 [inlined]
     [33] invokelatest
        @ ./essentials.jl:1052 [inlined]
     [34] message_handler_loop
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Distributed/src/process_messages.jl:176
     [35] process_tcp_streams
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Distributed/src/process_messages.jl:133
     [36] #103
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Distributed/src/process_messages.jl:121
    Stacktrace:
      [1] (::Base.var"#1170#1172")(x::Task)
        @ Base ./asyncmap.jl:171
      [2] foreach(f::Base.var"#1170#1172", itr::Vector{Any})
        @ Base ./abstractarray.jl:3187
      [3] maptwice(wrapped_f::Function, chnl::Channel{Any}, worker_tasks::Vector{Any}, c::Base.Iterators.Zip{Tuple{…}})
        @ Base ./asyncmap.jl:171
      [4] wrap_n_exec_twice
        @ ./asyncmap.jl:147 [inlined]
      [5] #async_usemap#1155
        @ ./asyncmap.jl:97 [inlined]
      [6] async_usemap
        @ ./asyncmap.jl:78 [inlined]
      [7] #asyncmap#1154
        @ ./asyncmap.jl:75 [inlined]
      [8] asyncmap
        @ ./asyncmap.jl:74 [inlined]
      [9] pmap(f::Function, p::CachingPool, c::Base.Iterators.Zip{…}; distributed::Bool, batch_size::Int64, on_error::Nothing, retry_delays::Vector{…}, retry_check::Nothing)
        @ Distributed ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Distributed/src/pmap.jl:126
     [10] pmap
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Distributed/src/pmap.jl:99 [inlined]
     [11] pmap
        @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Distributed/src/pmap.jl:155 [inlined]
     [12] (::AbstractMCMC.var"#67#79"{…})()
        @ AbstractMCMC ~/.julia/packages/AbstractMCMC/F50rb/src/sample.jl:583
Stacktrace:
  [1] sync_end(c::Channel{Any})
    @ Base ./task.jl:466
  [2] macro expansion
    @ ./task.jl:499 [inlined]
  [3] macro expansion
    @ ~/.julia/packages/AbstractMCMC/F50rb/src/sample.jl:539 [inlined]
  [4] macro expansion
    @ ~/.julia/packages/AbstractMCMC/F50rb/src/logging.jl:16 [inlined]
  [5] mcmcsample(rng::TaskLocalRNG, model::AbstractMCMC.LogDensityModel{…}, sampler::NUTS{…}, ::MCMCDistributed, N::Int64, nchains::Int64; progress::Bool, progressname::String, initial_params::Nothing, initial_state::Nothing, kwargs::@Kwargs{…})
    @ AbstractMCMC ~/.julia/packages/AbstractMCMC/F50rb/src/sample.jl:533
  [6] sample(rng::TaskLocalRNG, model::AbstractMCMC.LogDensityModel{…}, sampler::NUTS{…}, parallel::MCMCDistributed, N::Int64, nchains::Int64; n_adapts::Int64, progress::Bool, verbose::Bool, callback::Nothing, kwargs::@Kwargs{…})
    @ AdvancedHMC ~/.julia/packages/AdvancedHMC/xuhUJ/src/abstractmcmc.jl:115
  [7] sample
    @ ~/.julia/packages/AdvancedHMC/xuhUJ/src/abstractmcmc.jl:89 [inlined]
  [8] #sample#90
    @ ~/.julia/packages/AbstractMCMC/F50rb/src/logdensityproblems.jl:73 [inlined]
  [9] sample
    @ ~/.julia/packages/AbstractMCMC/F50rb/src/logdensityproblems.jl:64 [inlined]
 [10] #sample#22
    @ ~/.julia/packages/AbstractMCMC/F50rb/src/sample.jl:72 [inlined]
 [11] top-level scope
    @ ~/TuringLang/JuliaBUGS.jl.worktrees/sunxd/multi-threaded/dev.jl:38
Some type information was truncated. Use `show(err)` to see complete types.