Repository navigation
Conversation
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.
Member
Author
|
I also spun up a quick benchmark trying the recently released |
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 Report❌ Patch coverage is
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. 🚀 New features to boost your workflow:
|
… 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.
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:
Further explanation regarding interpretation and methodology can be found in the documentation. |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
coreswith mirai-basedparallelloo's parallel loops now run on persistent mirai daemons instead of
parallel::mclapply()(macOS, Linux) andparLapply()(Windows). Thecoresargument is replaced byparallel, which defaults toTRUE.Old and new approach
cores = non each call, oroptions(mc.cores)/options(loo.cores)parallel = TRUE/FALSEon each call,mirai::daemons(n)once per sessioncores == 1/mclapply/parLapplybranch at every call siteloo_map()Breaking changes
coresis removed from every function that had it. Useparallel = TRUEorFALSE. The deprecatedpsislw()keeps the argument, ignores it and runs sequentially.parallel = TRUE(the default) you need daemons first:mirai::daemons(n). Without them the call errors.options(mc.cores)andoptions(loo.cores)are no longer read, andloo_cores()is gone.miraiandmoriare added to Imports andparallelis dropped.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. Withparallel = FALSEit islapply(). Otherwise:Xis split into blocks that run on the daemons, and results come back in order.batch = TRUEsends each daemon one block of consecutive elements.batch = FALSEsends 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.mori::share(). Remote daemons receive copies. The dispatcher URL decides which applies.parallel = FALSE.Results
loo-v3.0.0withcores = N(PSOCK cluster per call)mirai::daemons(N)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):diamonds_functionloo_functionloo_subsamplerelative_effwaic_functiondiamonds_functionloo_functionloo_subsamplerelative_effwaic_functionThe function-based paths do not ship large closures: the log-likelihood function serializes to 412 bytes.
loo_subsample()andrelative_eff()stay a few percent slower without start-up, which I have not looked into. Scaling plots, relative to sequential and in plain seconds, aredev-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.
ap_matrixbrms_topleveldiamonds_functiondiamonds_matrixloo_five_callsloo_functionloo_function_heavy_drawsloo_subsamplemetricsmoment_matchingpsis_matrix_largepsis_matrix_smallradon_loo_apradon_loo_functionrelative_effsar_loosubsample_methodsvig_large_datavig_milk_klinevig_moment_matchingvig_roachesvig_wells_rstanwaic_functionwells_loo_subsamplewells_relative_eff_loowells4_weightsExact numbers
ap_matrixap_matrixap_matrixbrms_toplevelbrms_toplevelbrms_topleveldiamonds_functiondiamonds_functiondiamonds_functiondiamonds_matrixdiamonds_matrixdiamonds_matrixloo_five_callsloo_five_callsloo_five_callsloo_functionloo_functionloo_functionloo_function_heavy_drawsloo_function_heavy_drawsloo_function_heavy_drawsloo_subsampleloo_subsampleloo_subsamplemetricsmetricsmetricsmoment_matchingmoment_matchingmoment_matchingpsis_matrix_largepsis_matrix_largepsis_matrix_largepsis_matrix_smallpsis_matrix_smallpsis_matrix_smallradon_loo_apradon_loo_apradon_loo_apradon_loo_functionradon_loo_functionradon_loo_functionrelative_effrelative_effrelative_effsar_loosar_loosar_loosubsample_methodssubsample_methodssubsample_methodsvig_large_datavig_large_datavig_large_datavig_milk_klinevig_milk_klinevig_milk_klinevig_moment_matchingvig_moment_matchingvig_moment_matchingvig_roachesvig_roachesvig_roachesvig_wells_rstanvig_wells_rstanvig_wells_rstanwaic_functionwaic_functionwaic_functionwells_loo_subsamplewells_loo_subsamplewells_loo_subsamplewells_relative_eff_loowells_relative_eff_loowells_relative_eff_loowells4_weightswells4_weightswells4_weightsLinux
Change in time against the baseline, 80% CrI. Positive is slower. Raw reports are in
dev-notes/foil_results/linux.ap_matrixbrms_topleveldiamonds_functiondiamonds_matrixloo_five_callsloo_function_heavy_drawsloo_functionloo_subsamplemetricsmoment_matchingpsis_matrix_largepsis_matrix_smallradon_loo_apradon_loo_functionrelative_effsar_loosubsample_methodsvig_large_datavig_milk_klinevig_moment_matchingvig_roachesvig_wells_rstanwaic_functionwells4_weightswells_loo_subsamplewells_relative_eff_looTypical usage
The loo calls (
relative_eff,loo,waicwhere the source uses it,loo_compare) from two Bayesian Workflow scripts (roaches, arsenic) and theloo2-example,loo2-with-rstanandloo2-weightsvignettes, with the fits done beforehand. Median seconds of 7 repetitions on WSL (20 threads), daemons already running. 1 worker iscores = 1againstparallel = FALSE. Code and data setup are indev-notes/vignette-benchmarks.bayesian-workflow/arsenic_loobayesian-workflow/arsenic_loobayesian-workflow/arsenic_loobayesian-workflow/arsenic_loobayesian-workflow/arsenic_looloo2-with-rstanloo2-with-rstanloo2-with-rstanloo2-with-rstanloo2-with-rstanbayesian-workflow/roachesbayesian-workflow/roachesbayesian-workflow/roachesbayesian-workflow/roachesbayesian-workflow/roachesloo2-exampleloo2-exampleloo2-exampleloo2-exampleloo2-exampleloo2-weightsloo2-weightsloo2-weightsloo2-weightsloo2-weightsOn the small workloads the old code gets slower than serial as workers are added (
loo2-weightsgoes from 0.14 s to 1.47 s at 16); the new code rises less (0.12 s to 0.50 s). Onloo2-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 whyarsenic_loois faster at 1 worker.loo2-elpd,loo2-mixisandloo2-non-factorizedare not covered.One daemon against
lapplyloo_map()on one mirai daemon againstparallel = FALSE, real job sizes, Linux, time change with 80% CrI. Positive is slower; this is the fixed cost of the daemon round trip.ap_matrixbrms_topleveldiamonds_functiondiamonds_matrixloo_five_callsloo_function_heavy_drawsloo_functionloo_subsamplemetricsmoment_matchingpsis_matrix_largepsis_matrix_smallradon_loo_apradon_loo_functionrelative_effsar_loosubsample_methodsvig_large_datavig_milk_klinevig_moment_matchingvig_roachesvig_wells_rstanwaic_functionwells4_weightswells_loo_subsamplewells_relative_eff_loomori against plain copies
The same six benchmarks at 4 workers on Windows with
loo_map()as written, againstloo_map()withmori::share()replaced by the identity (arguments copied to each daemon).diamonds_matrixloo_functionpsis_matrix_largepsis_matrix_smallrelative_effwaic_functionFootnotes
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 ↩