Skip to content

Parallel - #419

Draft
VisruthSK wants to merge 25 commits into
loo-v3.0.0from
parallel
Draft

VisruthSK wants to merge 25 commits into
loo-v3.0.0from
parallel

Conversation

@VisruthSK

@VisruthSK VisruthSK commented Nov 5, 2026 •

Copy link
Copy Markdown
Member

Supersedes #378. Assisted by Claude Code.

In progress, I need to rerun benchmarks, also need to run benchmarks on native Linux and ideally macOS too.

Also will need to open PRs and discuss rstanarm, rstan, etc. other downstream packages if we do fully deprecate cores.

I'm leaving this as the full implementation and once its hashed out, will stack a PR for a parallelization vignette. Also consider the current documentation to be temporary.

NB: The following is LLM generated text


Replace cores with mirai-based parallel

loo's parallel loops now run on persistent mirai daemons instead of parallel::mclapply() (macOS, Linux) and parLapply() (Windows). The cores argument is replaced by parallel, which defaults to TRUE.

# before
loo(log_lik, cores = 4)

# after
mirai::daemons(4)
loo(log_lik)

Old and new approach

Before After
Switch cores = n on each call, or options(mc.cores) / options(loo.cores) parallel = TRUE/FALSE on each call, mirai::daemons(n) once per session
Workers Forked on every call on Unix, a new PSOCK cluster on every call on Windows Daemons that start once and serve every later call
Data to workers Copy-on-write fork on Unix, serialized copies on Windows Shared memory (mori) for local daemons, copies for remote ones
Code A cores == 1 / mclapply / parLapply branch at every call site One internal helper, loo_map()
Machines The local machine The local machine, or remote daemons

Breaking changes

  • cores is removed from every function that had it. Use parallel = TRUE or FALSE. The deprecated psislw() keeps the argument, ignores it and runs sequentially.
  • With parallel = TRUE (the default) you need daemons first: mirai::daemons(n). Without them the call errors.
  • options(mc.cores) and options(loo.cores) are no longer read, and loo_cores() is gone.
  • mirai and mori are added to Imports and parallel is dropped.
  • In loo_moment_match(), a daemon that dies now raises an error. It used to be reported as a failed observation.

How it works

loo_map(X, FUN, ..., parallel, batch = TRUE) runs every parallel loop. With parallel = FALSE it is lapply(). Otherwise:

  • X is split into blocks that run on the daemons, and results come back in order. batch = TRUE sends each daemon one block of consecutive elements. batch = FALSE sends each element as its own task so an idle daemon takes the next one; loo_moment_match() uses it because its tasks are expensive and uneven.
  • Call sites map over row or column indices and pass large inputs as arguments, so no closure captures a big object.
  • Local daemons read arguments from shared memory with mori::share(). Remote daemons receive copies. The dispatcher URL decides which applies.
  • Matrices and 3-d arrays are split into slices before sharing, because mori gives each daemon a private copy of a whole array as soon as it is indexed but reads a shared list of vectors in place.
  • Warnings raised on a daemon are signalled again on the host, and the first error stops the map.
  • Calls made inside a worker use parallel = FALSE.

Results

Windows Linux (NRP TIDE)
System Windows 11, i7-12700H (14 cores, 20 threads), 16 GB Ubuntu 24.04 container, Xeon Gold 6430 (64 CPUs per node), one worker count per node, 6 GiB memory
R 4.6.2, mirai 2.7.3, mori 0.2.2 4.5.3
Baseline loo-v3.0.0 with cores = N (PSOCK cluster per call) same
Candidate this PR with mirai::daemons(N) same
Worker counts 2, 4, 8 2, 4, 8, 24, 32
Peak memory exact (job object) not measured (cgroup is read-only on the cluster)

foil ran 26 paired baseline and candidate benchmarks, interleaved, 20 repetitions (10 for the seven slowest, 80% CrI only). Windows peak memory is the total across the host and its workers, read with a locally patched foil1 that gives memory the same paired bootstrap and credible intervals as time.

On Windows, where the baseline pays for a PSOCK cluster on every call, the candidate is faster in 23 of 26 benchmarks at 2 workers and 24 of 26 at 4 and 8. On Linux, where the baseline forks, it is faster in 13 to 15 of 26 at every worker count.

The slower Linux ones are short jobs, and the gap grows with the worker count. foil times the whole script, and each candidate script starts its daemons inside it, while the baseline forks for free. Start-up took 0.3 to 0.4 s at 2 to 8 workers on a laptop under WSL and more at higher worker counts. With the daemons already running, waic() on the function interface takes 0.069 s against 0.146 s for forking at 8 workers. To get both numbers from one run, each candidate script times its own daemon start-up and it is subtracted run by run (WSL, 10 repetitions, 80% CrI, change against the old code):

Benchmark Workers Total Without start-up
diamonds_function 4 -0.5% (-2.8%, +1.9%) -5.8% (-8.0%, -3.4%)
loo_function 4 +1.0% (-1.4%, +3.1%) -10.6% (-12.6%, -8.6%)
loo_subsample 4 +50.7% (+49.0%, +52.6%) +4.3% (+3.1%, +5.7%)
relative_eff 4 +15.5% (+14.0%, +17.0%) +4.8% (+3.4%, +6.1%)
waic_function 4 +43.9% (+42.2%, +45.7%) -13.9% (-14.5%, -13.2%)
diamonds_function 8 -0.7% (-2.6%, +1.2%) -8.1% (-9.9%, -6.3%)
loo_function 8 -3.7% (-5.5%, -1.6%) -18.3% (-20.0%, -16.6%)
loo_subsample 8 +50.3% (+47.1%, +53.3%) +3.4% (+0.7%, +6.1%)
relative_eff 8 +7.6% (+4.1%, +11.5%) -6.4% (-9.4%, -3.2%)
waic_function 8 +60.4% (+58.4%, +62.3%) -13.5% (-15.0%, -12.1%)

The function-based paths do not ship large closures: the log-likelihood function serializes to 412 bytes. loo_subsample() and relative_eff() stay a few percent slower without start-up, which I have not looked into. Scaling plots, relative to sequential and in plain seconds, are dev-notes/foil_results/scaling_{linux,windows}_{relative,absolute}.png.

Windows

Each cell is the change in time / the change in peak memory against the baseline. Negative is faster or smaller.

Benchmark 2 workers 4 workers 8 workers
ap_matrix -24.32% / +11.18% -34.11% / +12.86% -39.39% / +13.78%
brms_toplevel -47.51% / +20.84% -54.39% / +32.94% -63.30% / +43.05%
diamonds_function -4.16% / +9.10% -3.11% / +11.65% -7.13% / +11.65%
diamonds_matrix -26.94% / +13.27% -40.81% / +16.17% -52.69% / +13.18%
loo_five_calls -26.14% / -1.86% -39.69% / -1.58% -51.93% / -0.78%
loo_function +0.13% / +5.25% -11.35% / +6.85% -15.89% / +7.51%
loo_function_heavy_draws +0.69% / +10.54% -2.28% / +13.61% +3.14% / +14.05%
loo_subsample -7.79% / +4.33% -10.71% / +16.49% -18.14% / +27.59%
metrics -19.41% / +29.74% -26.86% / +29.29% -26.63% / +41.99%
moment_matching -25.93% / +17.71% -40.90% / +20.84% -56.52% / +20.89%
psis_matrix_large -27.97% / -23.60% -46.86% / -33.33% -56.11% / -35.66%
psis_matrix_small -75.11% / +56.05% -80.23% / +71.21% -84.44% / +73.77%
radon_loo_ap +0.95% / +0.94% +3.87% / -0.64% +2.32% / -2.11%
radon_loo_function -0.54% / +0.55% +7.89% / -0.94% -0.87% / -0.92%
relative_eff -25.65% / -41.98% -49.60% / -46.17% -64.26% / -54.00%
sar_loo -90.07% / +42.86% -93.00% / +54.91% -93.90% / +60.90%
subsample_methods -50.29% / +33.54% -61.92% / +33.08% -69.13% / +38.07%
vig_large_data -32.39% / +8.73% -43.15% / +7.75% -53.58% / +5.15%
vig_milk_kline -85.36% / +41.28% -87.85% / +44.35% -89.22% / +38.38%
vig_moment_matching -46.26% / +43.14% -64.61% / +20.74% -77.35% / +5.36%
vig_roaches -51.11% / +26.04% -59.96% / +30.50% -65.70% / +37.35%
vig_wells_rstan -32.82% / -33.93% -51.30% / -44.37% -68.14% / -46.29%
waic_function -5.98% / -5.10% -9.15% / -6.37% -9.22% / +1.24%
wells_loo_subsample -57.10% / +9.98% -70.89% / +27.22% -79.02% / +35.50%
wells_relative_eff_loo -0.88% / +7.03% -14.46% / +8.79% -22.00% / +13.14%
wells4_weights -38.33% / -1.92% -55.79% / +7.03% -64.39% / +15.02%
Exact numbers
Benchmark Workers Baseline Candidate Time change, 80% and 90% CrI Baseline peak Candidate peak Memory change 80% CrI 90% CrI
ap_matrix 2 5.7s 4.3s -24.32% (-25.47%, -23.13%) (-25.85%, -22.75%) 880.0MB 978.5MB +11.18% (+10.79%, +11.55%) (+10.68%, +11.65%)
ap_matrix 4 6.3s 4.1s -34.11% (-35.68%, -32.31%) (-36.10%, -31.73%) 1.2GB 1.4GB +12.86% (+12.44%, +13.31%) (+12.33%, +13.44%)
ap_matrix 8 7.1s 4.3s -39.39% (-41.52%, -37.04%) (-42.15%, -36.30%) 2.0GB 2.3GB +13.78% (+12.88%, +14.75%) (+12.65%, +15.03%)
brms_toplevel 2 13.1s 6.9s -47.51% (-48.26%, -46.82%) (-48.49%, -46.59%) 0.9GB 1.1GB +20.84% (+20.54%, +21.19%) (+20.47%, +21.31%)
brms_toplevel 4 17.7s 8.0s -54.39% (-56.15%, -52.47%) (-56.61%, -51.95%) 1.2GB 1.6GB +32.94% (+32.30%, +33.56%) (+32.10%, +33.72%)
brms_toplevel 8 21.0s 7.7s -63.30% (-65.03%, -61.48%) (-65.49%, -60.96%) 1.7GB 2.4GB +43.05% (+42.56%, +43.52%) (+42.42%, +43.65%)
diamonds_function 2 5.0s 4.8s -4.16% (-6.12%, -2.06%) (-6.78%, -1.34%) 669.6MB 730.6MB +9.10% (+9.02%, +9.19%) (+9.00%, +9.21%)
diamonds_function 4 4.3s 4.2s -3.11% (-6.34%, -0.03%) (-7.35%, +0.94%) 1.0GB 1.2GB +11.65% (+11.42%, +11.90%) (+11.35%, +11.97%)
diamonds_function 8 4.6s 4.2s -7.13% (-11.26%, -2.73%) (-12.52%, -1.49%) 1.8GB 2.0GB +11.65% (+11.50%, +11.80%) (+11.46%, +11.84%)
diamonds_matrix 2 9.0s 6.6s -26.94% (-27.99%, -26.02%) (-28.29%, -25.79%) 1.2GB 1.4GB +13.27% (+12.56%, +14.00%) (+12.35%, +14.20%)
diamonds_matrix 4 9.6s 5.7s -40.81% (-42.22%, -39.31%) (-42.65%, -38.84%) 1.6GB 1.9GB +16.17% (+15.40%, +16.93%) (+15.18%, +17.13%)
diamonds_matrix 8 11.5s 5.4s -52.69% (-54.89%, -50.33%) (-55.43%, -49.64%) 2.5GB 2.8GB +13.18% (+12.59%, +13.83%) (+12.42%, +14.02%)
loo_five_calls 2 9.9s 7.3s -26.14% (-28.36%, -24.14%) (-29.02%, -23.63%) 542.4MB 532.3MB -1.86% (-1.97%, -1.76%) (-2.00%, -1.73%)
loo_five_calls 4 9.0s 5.5s -39.69% (-41.95%, -37.22%) (-42.54%, -36.46%) 954.7MB 939.6MB -1.58% (-1.78%, -1.38%) (-1.83%, -1.32%)
loo_five_calls 8 10.2s 4.9s -51.93% (-54.31%, -49.50%) (-54.95%, -48.75%) 1.8GB 1.7GB -0.78% (-0.93%, -0.65%) (-0.97%, -0.62%)
loo_function 2 6.9s 7.0s +0.13% (-2.88%, +3.35%) (-3.79%, +4.23%) 532.6MB 560.6MB +5.25% (+5.10%, +5.41%) (+5.06%, +5.46%)
loo_function 4 6.0s 5.3s -11.35% (-14.67%, -8.67%) (-15.72%, -8.05%) 933.8MB 997.7MB +6.85% (+6.60%, +7.09%) (+6.53%, +7.17%)
loo_function 8 5.6s 4.7s -15.89% (-18.96%, -12.57%) (-19.79%, -11.58%) 1.7GB 1.8GB +7.51% (+7.41%, +7.61%) (+7.37%, +7.63%)
loo_function_heavy_draws 2 7.7s 7.8s +0.69% (-2.66%, +4.20%) (-3.60%, +5.22%) 481.9MB 532.8MB +10.54% (+10.23%, +10.84%) (+10.14%, +10.91%)
loo_function_heavy_draws 4 5.9s 5.7s -2.28% (-5.89%, +1.14%) (-7.09%, +2.11%) 822.0MB 933.7MB +13.61% (+13.14%, +14.03%) (+13.01%, +14.15%)
loo_function_heavy_draws 8 4.9s 5.0s +3.14% (-1.18%, +8.75%) (-2.48%, +10.60%) 1.5GB 1.7GB +14.05% (+13.89%, +14.19%) (+13.83%, +14.22%)
loo_subsample 2 1.3s 1.2s -7.79% (-8.04%, -7.55%) (-8.11%, -7.49%) 523.2MB 545.9MB +4.33% (+4.30%, +4.36%) (+4.29%, +4.36%)
loo_subsample 4 1.6s 1.4s -10.71% (-11.65%, -9.43%) (-11.90%, -8.97%) 695.0MB 809.6MB +16.49% (+16.13%, +16.79%) (+16.02%, +16.86%)
loo_subsample 8 2.2s 1.8s -18.14% (-21.31%, -15.45%) (-22.25%, -14.80%) 1.2GB 1.5GB +27.59% (+27.41%, +27.77%) (+27.36%, +27.82%)
metrics 2 15.1s 12.2s -19.41% (-19.74%, -19.17%) n/a 589.7MB 765.1MB +29.74% (+29.32%, +29.95%) n/a
metrics 4 16.5s 12.1s -26.86% (-29.45%, -24.70%) n/a 0.9GB 1.2GB +29.29% (+28.87%, +29.96%) n/a
metrics 8 17.9s 13.1s -26.63% (-28.73%, -24.20%) n/a 1.4GB 2.0GB +41.99% (+41.59%, +42.30%) n/a
moment_matching 2 18.8s 14.0s -25.93% (-28.31%, -23.31%) (-28.98%, -22.43%) 601.5MB 708.0MB +17.71% (+17.25%, +18.14%) (+17.12%, +18.25%)
moment_matching 4 18.1s 10.7s -40.90% (-44.37%, -37.29%) (-45.39%, -36.23%) 0.9GB 1.1GB +20.84% (+20.39%, +21.28%) (+20.25%, +21.40%)
moment_matching 8 20.5s 8.9s -56.52% (-59.07%, -53.35%) (-59.67%, -52.37%) 1.6GB 1.9GB +20.89% (+20.61%, +21.19%) (+20.53%, +21.27%)
psis_matrix_large 2 16.0s 11.5s -27.97% (-28.34%, -27.48%) n/a 2.2GB 1.7GB -23.60% (-24.47%, -22.85%) n/a
psis_matrix_large 4 16.9s 9.0s -46.86% (-49.18%, -43.98%) n/a 2.9GB 1.9GB -33.33% (-33.63%, -33.02%) n/a
psis_matrix_large 8 18.6s 8.1s -56.11% (-59.01%, -52.94%) n/a 4.2GB 2.7GB -35.66% (-35.92%, -35.38%) n/a
psis_matrix_small 2 16.8s 4.2s -75.11% (-75.71%, -74.48%) (-75.86%, -74.30%) 428.5MB 668.7MB +56.05% (+55.70%, +56.43%) (+55.60%, +56.54%)
psis_matrix_small 4 20.8s 4.1s -80.23% (-80.93%, -79.44%) (-81.12%, -79.20%) 0.7GB 1.2GB +71.21% (+70.96%, +71.47%) (+70.89%, +71.54%)
psis_matrix_small 8 29.4s 4.6s -84.44% (-85.18%, -83.58%) (-85.37%, -83.32%) 1.2GB 2.0GB +73.77% (+73.50%, +74.06%) (+73.42%, +74.14%)
radon_loo_ap 2 6.7s 6.8s +0.95% (-0.96%, +2.86%) (-1.59%, +3.49%) 540.9MB 545.9MB +0.94% (+0.57%, +1.24%) (+0.46%, +1.31%)
radon_loo_ap 4 5.3s 5.5s +3.87% (-0.42%, +8.88%) (-1.61%, +10.43%) 952.3MB 946.2MB -0.64% (-0.92%, -0.35%) (-1.00%, -0.27%)
radon_loo_ap 8 4.6s 4.7s +2.32% (-2.46%, +7.25%) (-3.80%, +8.70%) 1.8GB 1.7GB -2.11% (-2.20%, -2.01%) (-2.23%, -1.98%)
radon_loo_function 2 6.2s 6.2s -0.54% (-2.14%, +1.29%) (-2.61%, +1.94%) 536.5MB 539.4MB +0.55% (+0.46%, +0.65%) (+0.43%, +0.68%)
radon_loo_function 4 4.6s 5.0s +7.89% (+5.22%, +11.47%) (+4.59%, +12.66%) 949.6MB 940.8MB -0.94% (-1.15%, -0.71%) (-1.20%, -0.64%)
radon_loo_function 8 4.4s 4.3s -0.87% (-5.90%, +4.08%) (-7.40%, +5.71%) 1.8GB 1.7GB -0.92% (-1.00%, -0.83%) (-1.02%, -0.80%)
relative_eff 2 6.2s 4.6s -25.65% (-26.82%, -24.91%) (-27.23%, -24.75%) 1.5GB 0.8GB -41.98% (-42.06%, -41.86%) (-42.08%, -41.83%)
relative_eff 4 8.0s 4.0s -49.60% (-51.38%, -47.80%) (-51.85%, -47.30%) 2.4GB 1.3GB -46.17% (-46.30%, -46.06%) (-46.34%, -46.02%)
relative_eff 8 11.2s 4.0s -64.26% (-65.73%, -62.70%) (-66.11%, -62.23%) 4.6GB 2.1GB -54.00% (-54.08%, -53.92%) (-54.11%, -53.90%)
sar_loo 2 25.6s 2.5s -90.07% (-90.44%, -89.56%) n/a 433.2MB 618.9MB +42.86% (+42.64%, +43.08%) n/a
sar_loo 4 35.4s 2.5s -93.00% (-93.38%, -92.48%) n/a 0.7GB 1.0GB +54.91% (+54.78%, +55.02%) n/a
sar_loo 8 52.0s 3.2s -93.90% (-94.15%, -93.61%) n/a 1.2GB 1.9GB +60.90% (+60.60%, +61.14%) n/a
subsample_methods 2 8.3s 4.1s -50.29% (-51.78%, -48.97%) (-52.23%, -48.58%) 423.4MB 565.4MB +33.54% (+33.50%, +33.57%) (+33.49%, +33.58%)
subsample_methods 4 8.9s 3.4s -61.92% (-63.40%, -60.31%) (-63.84%, -59.85%) 666.4MB 886.8MB +33.08% (+32.73%, +33.48%) (+32.64%, +33.60%)
subsample_methods 8 11.4s 3.5s -69.13% (-70.96%, -67.26%) (-71.44%, -66.74%) 1.1GB 1.6GB +38.07% (+37.81%, +38.36%) (+37.75%, +38.45%)
vig_large_data 2 28.2s 19.1s -32.39% (-34.30%, -29.78%) n/a 539.5MB 586.6MB +8.73% (+8.47%, +8.99%) n/a
vig_large_data 4 25.1s 14.2s -43.15% (-47.35%, -38.79%) n/a 1.0GB 1.0GB +7.75% (+7.27%, +8.19%) n/a
vig_large_data 8 26.7s 12.4s -53.58% (-57.87%, -49.59%) n/a 1.8GB 1.9GB +5.15% (+5.07%, +5.28%) n/a
vig_milk_kline 2 8.3s 1.2s -85.36% (-85.67%, -85.02%) (-85.76%, -84.92%) 416.3MB 588.2MB +41.28% (+41.24%, +41.33%) (+41.22%, +41.34%)
vig_milk_kline 4 11.6s 1.4s -87.85% (-88.35%, -87.32%) (-88.48%, -87.15%) 652.7MB 942.2MB +44.35% (+44.18%, +44.54%) (+44.14%, +44.59%)
vig_milk_kline 8 17.2s 1.9s -89.22% (-89.73%, -88.68%) (-89.87%, -88.51%) 1.1GB 1.6GB +38.38% (+38.26%, +38.50%) (+38.23%, +38.54%)
vig_moment_matching 2 19.6s 10.5s -46.26% (-46.90%, -45.69%) n/a 1.5GB 2.1GB +43.14% (+42.22%, +44.02%) n/a
vig_moment_matching 4 29.3s 10.3s -64.61% (-65.16%, -64.13%) n/a 2.3GB 2.8GB +20.74% (+20.24%, +21.20%) n/a
vig_moment_matching 8 47.1s 10.7s -77.35% (-77.54%, -77.18%) n/a 3.9GB 4.1GB +5.36% (+5.10%, +5.68%) n/a
vig_roaches 2 3.7s 1.8s -51.11% (-51.55%, -50.78%) (-51.67%, -50.71%) 549.0MB 692.0MB +26.04% (+25.84%, +26.22%) (+25.77%, +26.27%)
vig_roaches 4 4.7s 1.9s -59.96% (-60.62%, -59.39%) (-60.82%, -59.22%) 0.9GB 1.1GB +30.50% (+30.02%, +31.01%) (+29.90%, +31.15%)
vig_roaches 8 6.5s 2.2s -65.70% (-67.08%, -64.52%) (-67.49%, -64.18%) 1.4GB 1.9GB +37.35% (+37.02%, +37.71%) (+36.93%, +37.82%)
vig_wells_rstan 2 17.9s 12.0s -32.82% (-33.48%, -32.45%) n/a 2.1GB 1.4GB -33.93% (-34.14%, -33.73%) n/a
vig_wells_rstan 4 20.5s 10.0s -51.30% (-54.05%, -48.55%) n/a 3.2GB 1.8GB -44.37% (-44.75%, -44.02%) n/a
vig_wells_rstan 8 25.8s 8.2s -68.14% (-69.08%, -67.10%) n/a 5.1GB 2.7GB -46.29% (-46.44%, -46.14%) n/a
waic_function 2 1.6s 1.5s -5.98% (-6.23%, -5.69%) (-6.29%, -5.61%) 456.1MB 432.8MB -5.10% (-5.14%, -5.06%) (-5.14%, -5.04%)
waic_function 4 1.5s 1.4s -9.15% (-11.07%, -7.50%) (-11.77%, -6.90%) 794.1MB 743.5MB -6.37% (-6.50%, -6.15%) (-6.52%, -6.06%)
waic_function 8 1.6s 1.5s -9.22% (-10.50%, -7.82%) (-10.83%, -7.36%) 1.3GB 1.3GB +1.24% (+0.79%, +1.71%) (+0.67%, +1.83%)
wells_loo_subsample 2 12.8s 5.5s -57.10% (-58.41%, -55.68%) (-58.77%, -55.29%) 499.9MB 549.8MB +9.98% (+9.91%, +10.04%) (+9.89%, +10.06%)
wells_loo_subsample 4 14.7s 4.3s -70.89% (-72.44%, -69.25%) (-72.86%, -68.76%) 724.2MB 921.3MB +27.22% (+27.06%, +27.40%) (+27.02%, +27.46%)
wells_loo_subsample 8 19.7s 4.1s -79.02% (-79.95%, -77.85%) (-80.17%, -77.47%) 1.1GB 1.6GB +35.50% (+35.34%, +35.65%) (+35.29%, +35.69%)
wells_relative_eff_loo 2 8.5s 8.5s -0.88% (-3.31%, +1.94%) (-3.91%, +2.79%) 514.1MB 550.3MB +7.03% (+6.91%, +7.16%) (+6.87%, +7.19%)
wells_relative_eff_loo 4 7.0s 6.0s -14.46% (-17.39%, -11.79%) (-18.31%, -11.04%) 895.0MB 973.7MB +8.79% (+8.61%, +8.96%) (+8.57%, +9.01%)
wells_relative_eff_loo 8 6.6s 5.1s -22.00% (-25.32%, -19.03%) (-26.25%, -18.27%) 1.6GB 1.8GB +13.14% (+12.88%, +13.42%) (+12.81%, +13.50%)
wells4_weights 2 44.7s 27.6s -38.33% (-39.33%, -37.57%) n/a 1.3GB 1.2GB -1.92% (-2.07%, -1.77%) n/a
wells4_weights 4 48.3s 21.4s -55.79% (-56.87%, -54.36%) n/a 1.6GB 1.8GB +7.03% (+6.79%, +7.21%) n/a
wells4_weights 8 54.7s 19.5s -64.39% (-66.15%, -62.52%) n/a 2.4GB 2.8GB +15.02% (+14.09%, +15.76%) n/a

Linux

Change in time against the baseline, 80% CrI. Positive is slower. Raw reports are in dev-notes/foil_results/linux.

Benchmark 2 workers 4 workers 8 workers 24 workers 32 workers
ap_matrix +2.22% (+1.47%, +2.87%) -3.58% (-3.71%, -3.45%) -4.48% (-4.83%, -4.14%) -0.34% (-1.44%, +0.80%) +0.83% (+0.19%, +1.48%)
brms_toplevel -9.41% (-10.25%, -8.62%) +0.42% (+0.27%, +0.59%) +2.85% (+2.60%, +3.10%) +1.87% (+1.27%, +2.55%) -7.63% (-8.86%, -6.50%)
diamonds_function +5.91% (-0.04%, +12.72%) +8.28% (+5.60%, +10.62%) +11.29% (+9.95%, +12.59%) +23.47% (+17.44%, +29.69%) +15.07% (+13.46%, +16.76%)
diamonds_matrix +2.66% (+2.34%, +3.00%) -3.90% (-4.06%, -3.71%) -6.74% (-6.90%, -6.58%) -3.83% (-4.26%, -3.32%) -7.28% (-7.75%, -6.81%)
loo_five_calls -16.77% (-17.15%, -16.39%) -25.27% (-25.46%, -25.11%) -32.19% (-32.45%, -31.92%) -24.93% (-25.48%, -24.30%) -26.45% (-27.00%, -25.80%)
loo_function_heavy_draws +10.73% (+0.93%, +20.89%) -1.07% (-3.14%, +1.27%) +5.80% (+3.63%, +7.87%) +25.95% (+22.51%, +30.85%) +22.34% (+20.47%, +24.35%)
loo_function +6.68% (+6.11%, +7.28%) +12.22% (+11.99%, +12.48%) +18.37% (+17.96%, +18.81%) +51.12% (+49.29%, +52.84%) +54.84% (+53.00%, +56.62%)
loo_subsample +58.53% (+57.35%, +59.88%) +86.71% (+86.03%, +87.32%) +102.16% (+100.61%, +103.89%) +162.07% (+157.97%, +165.46%) +158.16% (+153.24%, +163.09%)
metrics -2.44% (-2.75%, -2.09%) -2.20% (-2.54%, -1.89%) -1.87% (-2.03%, -1.74%) +0.43% (-0.06%, +1.06%) -2.34% (-2.83%, -1.87%)
moment_matching +9.89% (+9.63%, +10.17%) +11.77% (+11.53%, +12.01%) +14.36% (+14.02%, +14.77%) +59.03% (+57.53%, +60.41%) +52.55% (+51.57%, +53.77%)
psis_matrix_large -11.16% (-11.52%, -10.83%) -24.78% (-25.21%, -24.20%) -28.99% (-31.31%, -24.92%) -33.96% (-34.36%, -33.40%) -31.22% (-31.80%, -30.62%)
psis_matrix_small -61.22% (-61.46%, -60.96%) -67.11% (-67.76%, -66.66%) -62.34% (-62.49%, -62.20%) -41.45% (-41.76%, -41.12%) -28.87% (-29.68%, -28.07%)
radon_loo_ap +6.70% (+6.33%, +7.05%) +13.62% (+13.43%, +13.85%) +29.45% (+28.84%, +30.14%) +73.97% (+71.87%, +76.23%) +84.88% (+82.98%, +86.86%)
radon_loo_function +8.91% (+8.49%, +9.30%) +12.82% (+12.54%, +13.06%) +27.24% (+26.72%, +27.74%) +88.91% (+85.80%, +91.55%) +106.03% (+104.06%, +107.92%)
relative_eff +25.82% (+25.27%, +26.32%) +32.04% (+31.55%, +33.02%) +38.72% (+38.36%, +39.09%) +57.67% (+56.57%, +58.93%) +58.90% (+57.58%, +60.08%)
sar_loo -84.95% (-85.04%, -84.86%) -86.99% (-87.06%, -86.91%) -86.20% (-86.29%, -86.11%) -79.16% (-79.40%, -78.91%) -76.05% (-76.28%, -75.74%)
subsample_methods -30.92% (-32.26%, -29.82%) -31.18% (-31.71%, -30.67%) -28.37% (-28.75%, -27.97%) -25.03% (-25.61%, -24.38%) -23.35% (-23.78%, -22.92%)
vig_large_data -12.98% (-19.86%, -6.71%) -13.07% (-15.24%, -11.09%) -12.85% (-13.90%, -11.64%) -8.51% (-10.28%, -6.86%) -5.28% (-6.63%, -4.23%)
vig_milk_kline -76.69% (-76.88%, -76.51%) -78.52% (-78.60%, -78.43%) -77.10% (-77.19%, -77.01%) -72.05% (-72.29%, -71.81%) -70.34% (-70.60%, -70.09%)
vig_moment_matching -13.53% (-14.06%, -12.82%) +3.60% (+3.43%, +3.79%) +18.28% (+17.91%, +18.56%) +32.80% (+32.44%, +33.24%) +26.82% (+26.06%, +27.51%)
vig_roaches -22.78% (-23.14%, -22.38%) -23.75% (-23.91%, -23.58%) -23.69% (-24.06%, -23.31%) -13.53% (-14.41%, -12.56%) -12.97% (-13.64%, -12.31%)
vig_wells_rstan -8.32% (-8.77%, -7.82%) -13.44% (-13.58%, -13.35%) -20.23% (-20.33%, -20.14%) -21.68% (-21.95%, -21.37%) -24.38% (-24.83%, -23.90%)
waic_function +48.12% (+47.16%, +49.15%) +82.35% (+81.50%, +83.30%) +135.86% (+133.15%, +138.86%) +244.35% (+238.82%, +250.04%) +266.76% (+263.13%, +271.03%)
wells4_weights -22.14% (-22.56%, -21.77%) -33.97% (-34.09%, -33.86%) -43.43% (-43.59%, -43.27%) -51.00% (-51.28%, -50.70%) -53.37% (-53.57%, -53.16%)
wells_loo_subsample -41.21% (-42.37%, -40.01%) -37.79% (-38.44%, -37.22%) -32.86% (-33.11%, -32.59%) -27.93% (-28.51%, -27.39%) -26.36% (-26.99%, -25.79%)
wells_relative_eff_loo +14.01% (+5.29%, +22.16%) +2.71% (+0.55%, +4.59%) +5.06% (+3.48%, +6.92%) +16.42% (+13.90%, +19.05%) +11.14% (+10.07%, +12.20%)

Typical usage

The loo calls (relative_eff, loo, waic where the source uses it, loo_compare) from two Bayesian Workflow scripts (roaches, arsenic) and the loo2-example, loo2-with-rstan and loo2-weights vignettes, with the fits done beforehand. Median seconds of 7 repetitions on WSL (20 threads), daemons already running. 1 worker is cores = 1 against parallel = FALSE. Code and data setup are in dev-notes/vignette-benchmarks.

Workload Workers Before This PR This PR / before
bayesian-workflow/arsenic_loo 1 10.60 6.91 0.65
bayesian-workflow/arsenic_loo 2 8.84 5.63 0.64
bayesian-workflow/arsenic_loo 4 8.21 4.79 0.58
bayesian-workflow/arsenic_loo 8 7.87 4.72 0.60
bayesian-workflow/arsenic_loo 16 7.90 3.98 0.50
loo2-with-rstan 1 8.21 8.71 1.06
loo2-with-rstan 2 6.78 7.78 1.15
loo2-with-rstan 4 5.89 5.55 0.94
loo2-with-rstan 8 5.99 3.78 0.63
loo2-with-rstan 16 6.41 3.62 0.56
bayesian-workflow/roaches 1 1.10 0.93 0.84
bayesian-workflow/roaches 2 1.47 0.77 0.52
bayesian-workflow/roaches 4 1.31 0.54 0.41
bayesian-workflow/roaches 8 1.54 0.55 0.36
bayesian-workflow/roaches 16 2.09 0.52 0.25
loo2-example 1 0.74 0.76 1.02
loo2-example 2 1.14 0.67 0.59
loo2-example 4 0.98 0.41 0.42
loo2-example 8 1.14 0.40 0.35
loo2-example 16 1.54 0.38 0.25
loo2-weights 1 0.14 0.12 0.85
loo2-weights 2 0.44 0.15 0.34
loo2-weights 4 0.70 0.17 0.25
loo2-weights 8 1.18 0.29 0.25
loo2-weights 16 1.47 0.50 0.34

On the small workloads the old code gets slower than serial as workers are added (loo2-weights goes from 0.14 s to 1.47 s at 16); the new code rises less (0.12 s to 0.50 s). On loo2-with-rstan (3020 observations) the new code is 1.06x and 1.15x the old time at 1 and 2 workers and faster from 4. I did not look into why arsenic_loo is faster at 1 worker. loo2-elpd, loo2-mixis and loo2-non-factorized are not covered.

One daemon against lapply

loo_map() on one mirai daemon against parallel = FALSE, real job sizes, Linux, time change with 80% CrI. Positive is slower; this is the fixed cost of the daemon round trip.

Benchmark Change
ap_matrix +30.96% (+30.71%, +31.21%)
brms_toplevel +5.89% (+5.76%, +6.04%)
diamonds_function +9.18% (+9.04%, +9.32%)
diamonds_matrix +27.21% (+27.02%, +27.45%)
loo_five_calls +7.62% (+7.32%, +7.97%)
loo_function_heavy_draws +6.93% (+6.57%, +7.29%)
loo_function +13.75% (+13.43%, +14.03%)
loo_subsample +36.50% (+35.50%, +37.33%)
metrics +7.85% (+7.22%, +8.48%)
moment_matching +19.14% (+18.96%, +19.28%)
psis_matrix_large +37.59% (+37.24%, +37.83%)
psis_matrix_small +29.77% (+29.65%, +29.89%)
radon_loo_ap +8.04% (+7.92%, +8.17%)
radon_loo_function +7.34% (+7.24%, +7.44%)
relative_eff +32.10% (+32.02%, +32.17%)
sar_loo +51.00% (+50.74%, +51.27%)
subsample_methods +10.49% (+10.40%, +10.58%)
vig_large_data +2.79% (+2.70%, +2.88%)
vig_milk_kline +53.78% (+53.45%, +54.12%)
vig_moment_matching -21.32% (-21.44%, -21.22%)
vig_roaches +38.82% (+38.58%, +39.08%)
vig_wells_rstan +27.80% (+27.62%, +28.00%)
waic_function +37.17% (+36.89%, +37.45%)
wells4_weights +29.67% (+29.53%, +29.89%)
wells_loo_subsample +8.68% (+8.59%, +8.77%)
wells_relative_eff_loo +4.73% (+4.62%, +4.84%)

mori against plain copies

The same six benchmarks at 4 workers on Windows with loo_map() as written, against loo_map() with mori::share() replaced by the identity (arguments copied to each daemon).

Benchmark Time with mori vs copies (80% CrI) Peak without mori Peak with mori Peak change (80% CrI)
diamonds_matrix -16.50% (-18.60%, -14.07%) 3.0GB 1.9GB -37.10% (-37.31%, -36.87%)
loo_function -9.77% (-13.32%, -6.69%) 997.4MB 999.3MB +0.19% (+0.10%, +0.28%)
psis_matrix_large -33.82% (-37.42%, -29.77%) 5.7GB 1.9GB -66.10% (-66.33%, -65.80%)
psis_matrix_small -4.16% (-7.12%, -1.76%) 1.2GB 1.2GB +0.37% (+0.19%, +0.57%)
relative_eff -12.73% (-15.42%, -9.86%) 2.5GB 1.3GB -49.56% (-49.80%, -49.32%)
waic_function +0.77% (-0.37%, +3.05%) 743.5MB 743.8MB +0.04% (+0.01%, +0.07%)

Footnotes

  1. This is a one-off/hot-patched version of foil which supports memory measurement, not the full and verified implementation. The results seem tenable but there's a chance this isn't correct, but I won't get to updating foil fully for a while -Visruth ↩

mizu_map() over shared memory replaces mirai_map() when the user has set
a default pool with mizu::mizu_set_default_pool(); mirai daemons stay the
default and the path for remote workers.
The mizu pool matched mirai on compute-bound loo workloads and only
won on dispatch-bound ones that loo does not produce, so the extra
dependency and code path are not worth keeping. Benchmark results are
kept in dev-notes/foil_results/mizu.
…ntial

loo.function and loo_approximate_posterior call loo_map() directly, so the
two pass-through helpers and their man page are removed. psislw() is
deprecated and no longer uses the parallel package; its cores argument is
ignored.
@VisruthSK

VisruthSK commented Nov 5, 2026 •

Copy link
Copy Markdown
Member Author

I also spun up a quick benchmark trying the recently released mizu but unsurprisingly it didn't improve performance (due to the nature of the computations themselves, mirai + mori suffice). Results are in dev-notes which will be deleted before merging.

@VisruthSK VisruthSK linked an issue Nov 5, 2026 that may be closed by this pull request
loo_map() checks for them only when parallel = TRUE and names the missing package. It keeps the shared arguments in a local variable until the results are back. The test setup starts one daemon without a dispatcher when mirai is installed.
mirai and mori are only suggested, and daemons start with an empty global environment, so a log-likelihood function must be self-contained or get its inputs through data, draws or ...
The loo tests now compare parallel = FALSE with parallel = TRUE, and the test names say so. Comments about forked processes and mclapply are removed.
The migration guide still said options(loo.cores) remained for backward compatibility.
New benchmark runs come from bench/run.sh and land in dev-notes/foil_results/<platform>/w<N>.
Reading a small object from shared memory costs more per access in the worker
than a plain copy, and copying only costs more for large objects or when every
element is its own task. loo_map now shares an argument when it is larger than
share_bytes (4 MiB) or when tasks are sent one per element (batch = FALSE).
options(loo.parallel = FALSE) sets the default of parallel in every function
that has it, and options(loo.share_bytes = ...) sets the size above which
loo_map shares arguments with mori. Document both, test that the sharing
cutoff does not change results, and drop the unused default in
EXX_loo_compute().
Each daemon now returns the normalizing constant and effective sample size of
its columns, so the host no longer makes two more passes over the whole matrix.
Add tests that psis, tis and sis give identical objects sequentially and on two
daemons, and that the stored constants and sample sizes match the weights.
pointwise_loo_calcs() ran several passes over the whole log-likelihood matrix
on the host after psis() returned. It now maps a per-column function through
loo_map(), reading the log-likelihood and log-weight columns from shared
memory, so the columns are computed on the daemons. Add a test that loo.matrix
gives identical results sequentially and on two daemons, and share the
two-daemon helper between test files.
The base branch takes cores and this branch takes parallel with mirai daemons
set up beforehand, so the benchmark script picks the arguments from the
installed loo.matrix. loo_matrix and loo_function stay sequential on both, and
two new benchmarks run each on two workers, with the daemons started and
loaded before timing.
@codecov-commenter

codecov-commenter commented Nov 7, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 97.82609% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 94.71%. Comparing base (af50d00) to head (c0de6b5).

Files with missing lines Patch % Lines
R/importance_sampling.R 88.00% 3 Missing ⚠️
R/loo_subsample.R 95.23% 1 Missing ⚠️
R/psis_approximate_posterior.R 83.33% 1 Missing ⚠️
Additional details and impacted files
@@              Coverage Diff               @@
##           loo-v3.0.0     #419      +/-   ##
==============================================
+ Coverage       92.65%   94.71%   +2.05%     
==============================================
  Files              31       31              
  Lines            3240     3159      -81     
==============================================
- Hits             3002     2992      -10     
+ Misses            238      167      -71     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

… APIs

The base branch takes cores and this branch takes parallel with mirai daemons
started beforehand, so the script picks the arguments from the installed
loo.matrix. Each of loo_matrix and loo_function runs sequentially and on two
workers, with the daemons started and loaded before timing. benchmark_run takes
one expression per call, so each benchmark gets its own call and the shared
setup is built once.
@github-actions

github-actions Bot commented Nov 7, 2026 •

Copy link
Copy Markdown
Contributor

This is how benchmark results would change (along with a 95% confidence interval in relative change) if c0de6b5 is merged into loo-v3.0.0:

Status Name Current Incoming Interval
🚀 loo_matrix 1.14s 963.78ms [-16.67%, -14.44%]
❗🐌 loo_matrix_two_workers 1.45s 2.78s [+91.06%, +93.87%]
❗🐌 loo_function 932.73ms 1.01s [+4.47%, +11.43%]
❗🐌 loo_function_two_workers 573.24ms 897.65ms [+55.12%, +58.07%]

Further explanation regarding interpretation and methodology can be found in the documentation.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Parallelisation of loo using Mirai

2 participants