Parallelisation of loo using Mirai #308

Activity

1. n-kall commented on Oct 21, 2025

   

   n-kall

   on Oct 21, 2025

   Contributor

   More actions

   Forking is currently needed to parallelize moment matching because it relies on the model methods of Stan model objects, but other functions should be able to use other parallelization schemes
2. avehtari commented on Mar 1, 2026

   

   avehtari

   on Mar 1, 2026

   Member

   More actions

   As forking is not possible in Windows, does this mean that moment matching is not available on Windows?
3. n-kall commented on Mar 1, 2026

   

   n-kall

   on Mar 1, 2026 via email

   Contributor

   More actions

   I believe it fails maybe silently if cores > 1 on windows 1 Mar 2026, 14:26 by ***@***.***:

   …

   avehtari> left a comment > (stan-dev/loo#308) <#308 (comment)> As forking is not possible in Windows, does this mean that moment matching is not available on Windows? — Reply to this email directly, > view it on GitHub <#308 (comment)>> , or > unsubscribe <https://github.com/notifications/unsubscribe-auth/AIAFQS2M6XIUUFJ4NFKJFU34OQT53AVCNFSM6AAAAACJZYPKW2VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZTSNZZHA3DKMRWGA>> . You are receiving this because you commented.> Message ID: > <stan-dev/loo/issues/308/3979865260> @> github> .> com>
4. avehtari commented on Mar 1, 2026

   

   avehtari

   on Mar 1, 2026

   Member

   More actions

   Model methods can be used with mirai by re-compiling models once per daemon. There is the additional time cost of re-compilation, but that happens also parallely and is made only once per daemon. For example, with 4 daemons, each daemon re-compiles, but then mirai can send them many tasks. In moment matching, the re-compilation per daemon would happen once and as long as the daemons are alive they have valid model methods. Here is a proof-of-concept code (generated with help from Claude) using the Jacobian case study Stan code

   library(cmdstanr)
   data_bin <- list(N = 10, y = 9)
   code_binom <- "binom.stan"
   model_bin <- cmdstan_model(stan_file = code_binom,
                              compile_model_methods = TRUE, force_recompile = TRUE)
   fit_bin <- model_bin$sample(data = data_bin, refresh = 0)
   
   df_bin <- data.frame(theta = plogis(seq(-4, 6, length.out = 100))) |>
     dplyr::mutate(pdfbeta = dbeta(theta, 9 + 1, 1 + 1))
   
   # --- Parallel mirai_map ---
   library(mirai)
   mirai::daemons(n = 4)
   
   # Get the absolute path to the Stan file so daemons can find it
   stan_file_path <- normalizePath(code_binom)
   
   # Initialize each daemon: load cmdstanr, compile model methods,
   # fit the model, and store the fit in the daemon's global environment.
   # The Stan file and data are small so passing them via ... is fine.
   # everywhere() returns a mirai_map object (list of mirai, one per daemon).
   # Collect with [] to block until all daemons finish initialization.
   # Any daemon errors will surface here.
   everywhere(
     {
       library(cmdstanr)
       options(mc.cores = 1)
       model_daemon <- cmdstan_model(
         stan_file = stan_file_path,
         compile_model_methods = TRUE,
         force_recompile = TRUE
       )
       fit_daemon <<- model_daemon$sample(data = data_bin,
                                          fixed_param=TRUE,
                                          iter=1,
                                          iter_sampling=1)
       fit_daemon$init_model_methods()
     },
     stan_file_path = stan_file_path,
     data_bin = data_bin
   )[]
   
   # The mapped function uses 'fit_daemon' from the daemon's global environment
   # — no external pointers need to be serialized
   fit_pdf_daemon <- function(th) {
     exp(fit_daemon$log_prob(fit_daemon$unconstrain_variables(list(theta = th)),
                             jacobian = FALSE))
   }
   
   result_mirai <- mirai_map(df_bin$theta, fit_pdf_daemon)[.flat]
   print(result_mirai)
   
   mirai::daemons(0)
   
5. VisruthSK commented on Mar 2, 2026

   

   VisruthSK

   on Mar 2, 2026

   MemberAuthor

   More actions

   Seems like we don't need to recompile on every daemon too--seems like its tenable to compile once on main thread and share a cache to avoid model/method recompilations. Its a bit hack-y right now, but I think we could write a less brittle implementation with more careful thought.

   https://gist.github.com/VisruthSK/57eb18774d5ddd10b8b10643c3111d13
6. VisruthSK

   self-assigned this

   on May 26, 2026
7. florence-bockting

   changed the title [-]Look into using Mirai[/-] [+]Parallelisation of loo using Mirai[/+] on Jun 8, 2026
8. florence-bockting commented on Jun 8, 2026

   

   florence-bockting

   on Jun 8, 2026

   Contributor

   More actions

   linking here to Issue #169 as it seems to be related to a parallelisation issue on Windows. The Issue is already older so I don't know how relevant it still is, but it should be checked.
9. VisruthSK

   added this to the v3.0.0 milestone on Jun 9, 2026