diff --git a/.Rbuildignore b/.Rbuildignore index 63463781..f23c768d 100644 --- a/.Rbuildignore +++ b/.Rbuildignore @@ -1,3 +1,4 @@ +^benchmark$ ^CRAN-RELEASE$ ^.*\.Rproj$ ^\.Rproj\.user$ diff --git a/.gitignore b/.gitignore index 9c69942e..37d0b018 100644 --- a/.gitignore +++ b/.gitignore @@ -16,6 +16,7 @@ loo.Rproj # produced vignettes vignettes/loo2-lfo_cache/* vignettes/loo2-non-factorizable_cache/* +vignettes/articles-online-only/*.html vignettes/*.html vignettes/*.pdf inst/doc @@ -26,3 +27,7 @@ tests/testthat/Rplots.pdf cran-comments.md CRAN-RELEASE release-prep.R + +agent/* +data/* +scratch-files/* \ No newline at end of file diff --git a/DESCRIPTION b/DESCRIPTION index 4b15bd82..3b8c3234 100644 --- a/DESCRIPTION +++ b/DESCRIPTION @@ -37,6 +37,8 @@ Depends: Imports: checkmate, matrixStats (>= 0.52), + mirai, + mori, parallel, posterior (>= 1.7.0), stats @@ -62,3 +64,4 @@ LazyData: TRUE Roxygen: list(markdown = TRUE) SystemRequirements: pandoc (>= 1.12.3), pandoc-citeproc Config/roxygen2/version: 8.0.0 +RoxygenNote: 7.3.3 diff --git a/NAMESPACE b/NAMESPACE index 3405d737..562c6a24 100644 --- a/NAMESPACE +++ b/NAMESPACE @@ -130,6 +130,7 @@ export(loo_approximate_posterior.matrix) export(loo_compare) export(loo_crps) export(loo_i) +export(loo_mirai) export(loo_model_weights) export(loo_model_weights.default) export(loo_moment_match) diff --git a/NEWS.md b/NEWS.md index d998539c..c8137a14 100644 --- a/NEWS.md +++ b/NEWS.md @@ -1,5 +1,17 @@ # loo (development version) +* Parallelization with **mirai**: user-managed `mirai::daemons()` pools, + new `loo_mirai()` convenience wrapper for map-style workflows, and deprecated + per-call `cores` / `options(mc.cores)` bridges. +* The `cores` argument and `options(mc.cores)` on loo functions are deprecated + and will be removed in a future release; they still start a per-call local + daemon pool when the resolved value is `> 1` and no pool is connected. A + once-per-session warning is emitted when `cores` is passed explicitly (even + `cores = 1`), when `options(mc.cores)` is set, or when the resolved value is + `> 1`. +* Removed support for `options(loo.cores)`, which had been deprecated since + loo 2.x in favor of `options(mc.cores)`. + # loo 2.10.0 * Updates to `loo_compare` output by @jgabry, @avehtari, @florence-bockting in #300: diff --git a/R/effective_sample_sizes.R b/R/effective_sample_sizes.R index 360a3098..6ee4f2a2 100644 --- a/R/effective_sample_sizes.R +++ b/R/effective_sample_sizes.R @@ -50,6 +50,7 @@ relative_eff.default <- function(x, chain_id, ...) { #' @template matrix #' relative_eff.matrix <- function(x, chain_id, ..., cores = getOption("mc.cores", 1)) { + cores <- loo_cores(cores, call = match.call()) x <- llmatrix_to_array(x, chain_id) relative_eff.array(x, cores = cores) } @@ -59,35 +60,33 @@ relative_eff.matrix <- function(x, chain_id, ..., cores = getOption("mc.cores", #' @template array #' relative_eff.array <- function(x, ..., cores = getOption("mc.cores", 1)) { + cores <- loo_cores(cores, call = match.call()) stopifnot(length(dim(x)) == 3) S <- prod(dim(x)[1:2]) # posterior sample size = iter * chains - if (cores == 1) { - n_eff_vec <- apply(x, 3, posterior::ess_mean) - } else { - if (!os_is_windows()) { - n_eff_list <- - parallel::mclapply( - mc.cores = cores, - X = seq_len(dim(x)[3]), - FUN = function(i) posterior::ess_mean(x[, , i, drop = TRUE]) - ) - } else { - cl <- parallel::makePSOCKcluster(cores) - on.exit(parallel::stopCluster(cl)) - n_eff_list <- - parallel::parLapply( - cl = cl, - X = seq_len(dim(x)[3]), - fun = function(i) posterior::ess_mean(x[, , i, drop = TRUE]) - ) - } - n_eff_vec <- unlist(n_eff_list, use.names = FALSE) - } + # The full draws array is reused across observations, so it is broadcast via + # shared memory on a local pool. Each worker reads only its slice `x[, , i]`. + n_eff_list <- with_loo_daemons( + cores, + loo_map( + seq_len(dim(x)[3]), + relative_eff_i_array, + cores = cores, + broadcast = list(x = x) + ) + ) + n_eff_vec <- unlist(n_eff_list, use.names = FALSE) return(n_eff_vec / S) } +#' Worker computing `ess_mean()` for a single slice of a draws array +#' @noRd +#' @keywords internal +relative_eff_i_array <- function(i, x) { + posterior::ess_mean(x[, , i, drop = TRUE]) +} + #' @export #' @templateVar fn relative_eff #' @template function @@ -100,50 +99,41 @@ relative_eff.function <- cores = getOption("mc.cores", 1), data = NULL, draws = NULL) { + cores <- loo_cores(cores, call = match.call()) f_i <- validate_llfun(x) # not really an llfun, should return exp(ll) or exp(-ll) N <- dim(data)[1] - if (cores == 1) { - n_eff_list <- - lapply( - X = seq_len(N), - FUN = function(i) { - val_i <- f_i(data_i = data[i, , drop = FALSE], draws = draws, ...) - relative_eff.default(as.vector(val_i), chain_id = chain_id, cores = 1) - } - ) - } else { - if (!os_is_windows()) { - n_eff_list <- - parallel::mclapply( - X = seq_len(N), - FUN = function(i) { - val_i <- f_i(data_i = data[i, , drop = FALSE], draws = draws, ...) - relative_eff.default(as.vector(val_i), chain_id = chain_id, cores = 1) - }, - mc.cores = cores - ) - } else { - cl <- parallel::makePSOCKcluster(cores) - parallel::clusterExport(cl=cl, varlist=c("draws", "chain_id", "data"), envir=environment()) - on.exit(parallel::stopCluster(cl)) - n_eff_list <- - parallel::parLapply( - cl = cl, - X = seq_len(N), - fun = function(i) { - val_i <- f_i(data_i = data[i, , drop = FALSE], draws = draws, ...) - relative_eff.default(as.vector(val_i), chain_id = chain_id, cores = 1) - } - ) - } - } + # `data` and `draws` are reused for every observation, so they are broadcast + # via shared memory on a local pool and serialized on a remote pool. + n_eff_list <- with_loo_daemons( + cores, + loo_map( + seq_len(N), + relative_eff_i_function, + f_i = f_i, + chain_id = chain_id, + re_dots = list(...), + cores = cores, + broadcast = list(data = data, draws = draws) + ) + ) n_eff_vec <- unlist(n_eff_list, use.names = FALSE) return(n_eff_vec) } +#' Worker computing the relative effective sample size for observation `i` +#' @noRd +#' @keywords internal +relative_eff_i_function <- function(i, f_i, data, draws, chain_id, re_dots) { + val_i <- do.call( + f_i, + c(list(data_i = data[i, , drop = FALSE], draws = draws), re_dots) + ) + relative_eff.default(as.vector(val_i), chain_id = chain_id, cores = 1) +} + #' @export #' @describeIn relative_eff #' If `x` is an object of class `"psis"`, `relative_eff()` simply returns diff --git a/R/helpers.R b/R/helpers.R index 38b401dd..8d2fb382 100644 --- a/R/helpers.R +++ b/R/helpers.R @@ -171,14 +171,17 @@ nlist <- function(...) { } -# Check how many cores to use and throw deprecation warning if loo.cores is used -loo_cores <- function(cores) { - loo_cores_op <- getOption("loo.cores", NA) - if (!is.na(loo_cores_op) && (loo_cores_op != cores)) { - cores <- loo_cores_op - warning("'loo.cores' is deprecated, please use 'mc.cores' or pass 'cores' explicitly.", - call. = FALSE) - } +# Check how many cores to use and throw deprecation warnings for legacy options +loo_cores <- function(cores, call = NULL) { + explicit <- !is.null(call) && "cores" %in% names(call) + mc_cores_op <- getOption("mc.cores") + mc_cores_set <- !explicit && !is.null(mc_cores_op) + + loo_deprecate_cores( + cores = cores, + explicit = explicit, + mc_cores_set = mc_cores_set + ) return(cores) } diff --git a/R/importance_sampling.R b/R/importance_sampling.R index f8dfd4d3..0028a4d1 100644 --- a/R/importance_sampling.R +++ b/R/importance_sampling.R @@ -21,7 +21,7 @@ importance_sampling.array <- ..., r_eff = 1, cores = getOption("mc.cores", 1)) { - cores <- loo_cores(cores) + cores <- loo_cores(cores, call = match.call()) stopifnot(length(dim(log_ratios)) == 3) assert_importance_sampling_method_is_implemented(method) log_ratios <- validate_ll(log_ratios) @@ -38,7 +38,7 @@ importance_sampling.matrix <- ..., r_eff = 1, cores = getOption("mc.cores", 1)) { - cores <- loo_cores(cores) + cores <- loo_cores(cores, call = match.call()) assert_importance_sampling_method_is_implemented(method) log_ratios <- validate_ll(log_ratios) r_eff <- prepare_psis_r_eff(r_eff, len = ncol(log_ratios)) @@ -198,29 +198,23 @@ do_importance_sampling <- function(log_ratios, r_eff, cores, method) { stop("Incorrect IS method.") } - if (cores == 1) { - lw_list <- lapply(seq_len(N), function(i) - is_fun(log_ratios_i = log_ratios[, i], tail_len_i = tail_len[i])) - } else { - if (!os_is_windows()) { - lw_list <- parallel::mclapply( - X = seq_len(N), - mc.cores = cores, - FUN = function(i) - is_fun(log_ratios_i = log_ratios[, i], tail_len_i = tail_len[i]) - ) - } else { - cl <- parallel::makePSOCKcluster(cores) - on.exit(parallel::stopCluster(cl)) - lw_list <- - parallel::parLapply( - cl = cl, - X = seq_len(N), - fun = function(i) - is_fun(log_ratios_i = log_ratios[, i], tail_len_i = tail_len[i]) - ) - } - } + # Each observation needs a different column of `log_ratios`, but the whole + # matrix is reused across the map, so it is a broadcast object: `loo_map()` + # shares it via shared memory on a local pool (zero-copy column access) and + # falls back to serialization on a remote pool. Serial work runs as a plain + # lapply(). `with_loo_daemons()` provides a pool when this is the top-level + # call (e.g. psis()) and reuses an outer pool when called from loo(). + lw_list <- with_loo_daemons( + cores, + loo_map( + seq_len(N), + do_is_i, + is_fun = is_fun, + tail_len = tail_len, + cores = cores, + broadcast = list(log_ratios = log_ratios) + ) + ) log_weights <- psis_apply(lw_list, "log_weights", fun_val = numeric(S)) pareto_k <- psis_apply(lw_list, "pareto_k") @@ -234,3 +228,24 @@ do_importance_sampling <- function(log_ratios, r_eff, cores, method) { method = rep(method, length(pareto_k)) # Conform to other attr that exist per obs. ) } + +#' Apply an importance sampling method to a single observation +#' +#' @noRd +#' @keywords internal +#' @description +#' Worker function mapped over observations (matrix columns) by +#' `do_importance_sampling()`, either serially via [lapply()] or in parallel +#' via [mirai::mirai_map()]. +#' @param i Integer column index of the observation to process. +#' @param is_fun The per-observation importance sampling function to apply, one +#' of `do_psis_i()`, `do_tis_i()`, or `do_sis_i()`. +#' @param log_ratios Matrix of log ratios (`-loglik`). May be a shared-memory +#' object created by [mori::share()] to avoid copying to each worker. +#' @param tail_len Vector of tail lengths used to fit the GPD, one per +#' observation. +#' @return The result of `is_fun` for observation `i` (a list with elements +#' such as `log_weights` and `pareto_k`). +do_is_i <- function(i, is_fun, log_ratios, tail_len) { + is_fun(log_ratios_i = log_ratios[, i], tail_len_i = tail_len[i]) +} \ No newline at end of file diff --git a/R/loo.R b/R/loo.R index 10b1bdc7..2d1a4ed0 100644 --- a/R/loo.R +++ b/R/loo.R @@ -257,7 +257,7 @@ loo.function <- cores = getOption("mc.cores", 1), is_method = c("psis", "tis", "sis")) { is_method <- match.arg(is_method) - cores <- loo_cores(cores) + cores <- loo_cores(cores, call = match.call()) stopifnot(is.data.frame(data) || is.matrix(data), !is.null(draws)) assert_importance_sampling_method_is_implemented(is_method) .llfun <- validate_llfun(x) @@ -665,52 +665,23 @@ parallel_importance_sampling_list <- function(N, .loo_i, .llfun, data, draws, r_eff, save_psis, cores, method, ...){ - if (cores == 1) { - psis_list <- - lapply( - X = seq_len(N), - FUN = .loo_i, - llfun = .llfun, - data = data, - draws = draws, - r_eff = r_eff, - save_psis = save_psis, - is_method = method, - ... - ) - } else { - if (!os_is_windows()) { - # On Mac or Linux use mclapply() for multiple cores - psis_list <- - parallel::mclapply( - mc.cores = cores, - X = seq_len(N), - FUN = .loo_i, - llfun = .llfun, - data = data, - draws = draws, - r_eff = r_eff, - save_psis = save_psis, - is_method = method, - ... - ) - } else { - # On Windows use makePSOCKcluster() and parLapply() for multiple cores - cl <- parallel::makePSOCKcluster(cores) - on.exit(parallel::stopCluster(cl)) - psis_list <- - parallel::parLapply( - cl = cl, - X = seq_len(N), - fun = .loo_i, - llfun = .llfun, - data = data, - draws = draws, - r_eff = r_eff, - save_psis = save_psis, - is_method = method, - ... - ) - } - } + # `draws` (and `data`) are reused identically for every observation, so they + # are broadcast objects: shared once via shared memory on a local pool + # (recovering the copy-on-write benefit fork gave the old mclapply() path) + # and serialized on a remote pool. A single cross-platform code path replaces + # the previous mclapply()/parLapply() branching. + with_loo_daemons( + cores, + loo_map( + seq_len(N), + .loo_i, + llfun = .llfun, + r_eff = r_eff, + save_psis = save_psis, + is_method = method, + ..., + cores = cores, + broadcast = list(data = data, draws = draws) + ) + ) } diff --git a/R/loo_approximate_posterior.R b/R/loo_approximate_posterior.R index 3a3bce09..6d24e8fb 100644 --- a/R/loo_approximate_posterior.R +++ b/R/loo_approximate_posterior.R @@ -47,6 +47,7 @@ loo_approximate_posterior.array <- ..., save_psis = FALSE, cores = getOption("mc.cores", 1)) { + cores <- loo_cores(cores, call = match.call()) checkmate::assert_flag(save_psis) checkmate::assert_int(cores) checkmate::assert_matrix(log_p, mode = "numeric", nrows = dim(x)[1], ncols = dim(x)[2]) @@ -75,6 +76,7 @@ loo_approximate_posterior.matrix <- ..., save_psis = FALSE, cores = getOption("mc.cores", 1)) { + cores <- loo_cores(cores, call = match.call()) checkmate::assert_flag(save_psis) checkmate::assert_int(cores) checkmate::assert_numeric(log_p, len = nrow(x)) @@ -117,7 +119,7 @@ loo_approximate_posterior.function <- checkmate::assert_numeric(log_p, len = length(log_g)) checkmate::assert_numeric(log_g, len = length(log_p)) - cores <- loo_cores(cores) + cores <- loo_cores(cores, call = match.call()) stopifnot(is.data.frame(data) || is.matrix(data), !is.null(draws)) .llfun <- validate_llfun(x) N <- dim(data)[1] diff --git a/R/loo_model_weights.R b/R/loo_model_weights.R index 946dc7c3..4abc26c0 100644 --- a/R/loo_model_weights.R +++ b/R/loo_model_weights.R @@ -180,7 +180,7 @@ loo_model_weights.default <- r_eff_list = NULL, cores = getOption("mc.cores", 1)) { - cores <- loo_cores(cores) + cores <- loo_cores(cores, call = match.call()) method <- match.arg(method) K <- length(x) # number of models @@ -188,15 +188,24 @@ loo_model_weights.default <- N <- ncol(x[[1]]) # number of data points validate_log_lik_list(x) validate_r_eff_list(r_eff_list, K, N) - lpd_point <- matrix(NA, N, K) - elpd_loo <- rep(NA, K) - for (k in 1:K) { - r_eff_k <- r_eff_list[[k]] # possibly NULL - log_likelihood <- x[[k]] - loo_object <- loo(log_likelihood, r_eff = r_eff_k, cores = cores) - lpd_point[, k] <- loo_object$pointwise[, "elpd_loo"] #calculate log(p_k (y_i | y_-i)) - elpd_loo[k] <- loo_object$estimates["elpd_loo", "Estimate"] - } + # Establish a single daemon pool for all K models so each inner loo() + # reuses it instead of spinning a pool up and down K times. + loo_objects <- with_loo_daemons( + cores, + lapply(seq_len(K), function(k) { + loo(x[[k]], r_eff = r_eff_list[[k]], cores = cores) + }) + ) + lpd_point <- vapply( + loo_objects, + function(o) o$pointwise[, "elpd_loo"], #calculate log(p_k (y_i | y_-i)) + FUN.VALUE = numeric(N) + ) + elpd_loo <- vapply( + loo_objects, + function(o) o$estimates["elpd_loo", "Estimate"], + FUN.VALUE = numeric(1) + ) } else if (is.psis_loo(x[[1]])) { validate_psis_loo_list(x) lpd_point <- do.call(cbind, lapply(x, function(obj) obj$pointwise[, "elpd_loo"])) diff --git a/R/loo_moment_matching.R b/R/loo_moment_matching.R index 110eff93..67b9625d 100644 --- a/R/loo_moment_matching.R +++ b/R/loo_moment_matching.R @@ -81,6 +81,7 @@ loo_moment_match.default <- function(x, loo, post_draws, log_lik_i, checkmate::assertNumber(k_threshold, null.ok=TRUE) checkmate::assertLogical(split) checkmate::assertLogical(cov) + cores <- loo_cores(cores, call = match.call()) checkmate::assertNumber(cores) @@ -111,32 +112,27 @@ loo_moment_match.default <- function(x, loo, post_draws, log_lik_i, kfs <- rep(0,N) I <- which(ks > k_threshold) - loo_moment_match_i_fun <- function(i) { - loo_moment_match_i(i = i, x = x, log_lik_i = log_lik_i, - unconstrain_pars = unconstrain_pars, - log_prob_upars = log_prob_upars, - log_lik_i_upars = log_lik_i_upars, - max_iters = max_iters, k_threshold = k_threshold, - split = split, cov = cov, N = N, S = S, upars = upars, - orig_log_prob = orig_log_prob, k = ks[i], - is_method = is_method, npars = npars, ...) - } - - if (cores == 1) { - mm_list <- lapply(X = I, FUN = function(i) loo_moment_match_i_fun(i)) - } - else { - if (!os_is_windows()) { - mm_list <- parallel::mclapply(X = I, mc.cores = cores, - FUN = function(i) loo_moment_match_i_fun(i)) - } - else { - cl <- parallel::makePSOCKcluster(cores) - on.exit(parallel::stopCluster(cl)) - mm_list <- parallel::parLapply(cl = cl, X = I, - fun = function(i) loo_moment_match_i_fun(i)) - } - } + # The large unconstrained-draws matrix `upars` and the `orig_log_prob` vector + # are reused for every high-Pareto-k observation, so they are broadcast via + # shared memory on a local pool. The worker is the namespace-level + # `loo_moment_match_i_worker()` (rather than a closure over this frame) so the + # broadcast objects are not also dragged along inside a captured environment. + mm_list <- with_loo_daemons( + cores, + loo_map( + I, + loo_moment_match_i_worker, + x = x, ks = ks, log_lik_i = log_lik_i, + unconstrain_pars = unconstrain_pars, + log_prob_upars = log_prob_upars, + log_lik_i_upars = log_lik_i_upars, + max_iters = max_iters, k_threshold = k_threshold, + split = split, cov = cov, N = N, S = S, + is_method = is_method, npars = npars, mm_dots = list(...), + cores = cores, + broadcast = list(upars = upars, orig_log_prob = orig_log_prob) + ) + ) # update results for (ii in seq_along(I)) { @@ -230,6 +226,46 @@ loo_moment_match.default <- function(x, loo, post_draws, log_lik_i, #' @param ... Further arguments passed to the custom functions documented above. #' @return List with the updated elpd values and diagnostics #' +#' Worker wrapper around `loo_moment_match_i()`` for parallel mapping +#' +#' @noRd +#' @keywords internal +#' @description +#' A namespace-level (non-closure) adapter mapped over high-Pareto-k +#' observation indices by `loo_map()`. Keeping it at namespace scope means it +#' does not capture the calling frame, so large objects shared via +#' [mori::share()] (`upars`, `orig_log_prob`) are not duplicated inside a +#' serialized closure environment. The per-observation Pareto k is selected +#' here from the full `ks` vector, and any extra arguments are forwarded +#' through `mm_dots`. +#' @param i Integer observation index. +#' @param ks Full vector of Pareto k estimates; `ks[i]` is used for this fold. +#' @param mm_dots A list of additional arguments forwarded to +#' `loo_moment_match_i()` (the `...` from [loo_moment_match()]). +#' @return The result of `loo_moment_match_i()` for observation `i`. +loo_moment_match_i_worker <- function(i, x, ks, log_lik_i, unconstrain_pars, + log_prob_upars, log_lik_i_upars, + max_iters, k_threshold, split, cov, + N, S, upars, orig_log_prob, is_method, + npars, mm_dots) { + do.call( + loo_moment_match_i, + c( + list( + i = i, x = x, log_lik_i = log_lik_i, + unconstrain_pars = unconstrain_pars, + log_prob_upars = log_prob_upars, + log_lik_i_upars = log_lik_i_upars, + max_iters = max_iters, k_threshold = k_threshold, + split = split, cov = cov, N = N, S = S, upars = upars, + orig_log_prob = orig_log_prob, k = ks[i], + is_method = is_method, npars = npars + ), + mm_dots + ) + ) +} + loo_moment_match_i <- function(i, x, log_lik_i, diff --git a/R/loo_subsample.R b/R/loo_subsample.R index bcac4b17..d5c17e70 100644 --- a/R/loo_subsample.R +++ b/R/loo_subsample.R @@ -117,7 +117,7 @@ loo_subsample.function <- estimator = "diff_srs", llgrad = NULL, llhess = NULL) { - cores <- loo_cores(cores) + cores <- loo_cores(cores, call = match.call()) # Asserting inputs .llfun <- validate_llfun(x) stopifnot(is.data.frame(data) || is.matrix(data), !is.null(draws)) @@ -133,7 +133,6 @@ loo_subsample.function <- r_eff <- prepare_psis_r_eff(r_eff, len = dim(data)[1]) } checkmate::assert_flag(save_psis) - cores <- loo_cores(cores) checkmate::assert_choice(loo_approximation, choices = loo_approximation_choices(), null.ok = FALSE) checkmate::assert_int(loo_approximation_draws, lower = 1, upper = .ndraws(draws), null.ok = TRUE) @@ -276,7 +275,7 @@ update.psis_loo_ss <- function(object, ..., if (!is.null(draws)) { # No current checks } - cores <- loo_cores(cores) + cores <- loo_cores(cores, call = match.call()) # Update elpd approximations if (!is.null(loo_approximation) | !is.null(loo_approximation_draws)) { @@ -494,17 +493,18 @@ lpd_i <- function(i, llfun, data, draws) { #' @noRd #' @return a vector of computed log probability densities compute_lpds <- function(N, data, draws, llfun, cores) { - if (cores == 1) { - lpds <- lapply(X = seq_len(N), FUN = lpd_i, llfun, data, draws) - } else { - if (.Platform$OS.type != "windows") { - lpds <- mclapply(X = seq_len(N), mc.cores = cores, FUN = lpd_i, llfun, data, draws) - } else { - cl <- makePSOCKcluster(cores) - on.exit(stopCluster(cl)) - lpds <- parLapply(cl, X = seq_len(N), fun = lpd_i, llfun, data, draws) - } - } + # `draws` (and `data`) are reused for every observation, so they are shared + # once via shared memory on a local pool and serialized on a remote pool. + lpds <- with_loo_daemons( + cores, + loo_map( + seq_len(N), + lpd_i, + llfun = llfun, + cores = cores, + broadcast = list(data = data, draws = draws) + ) + ) unlist(lpds) } diff --git a/R/parallel.R b/R/parallel.R new file mode 100644 index 00000000..93bb4b6b --- /dev/null +++ b/R/parallel.R @@ -0,0 +1,382 @@ +#' Package-internal state for the parallel backend +#' +#' @noRd +#' @keywords internal +#' @description +#' Holds small bits of session-scoped state used by the parallel helpers: +#' +#' * `informed_cores_ignored`: guards the "a pool is connected so `cores` is +#' ignored" message in `with_loo_daemons()` so it is only emitted once per +#' session. +#' * `warned_cores_deprecated`: guards the `cores` argument deprecation warning +#' in `loo_deprecate_cores()` so it is only emitted once per session. +#' * `warned_cores_in_mirai`: guards the `cores`-in-`args_list` warning in +#' `loo_mirai()` so it is only emitted once per session. +#' +#' It also serves as the object attached to session-scoped parallel state. +.loo_internal <- new.env(parent = emptyenv()) + +#' Inform (once per session) that `cores` is ignored while a pool is connected +#' +#' @noRd +#' @keywords internal +#' @description +#' Emitted by `with_loo_daemons()` when a mirai daemon pool is connected (e.g. +#' the user called [mirai::daemons()] themselves) and the current call passed +#' `cores <= 1`. When a pool is connected loo always uses it, so the `cores` +#' argument is ignored and the work runs in parallel regardless of its value. +#' This message makes that behaviour visible the +#' first time a call looks like it asked for serial execution; the usual cause +#' is relying on the default `cores = getOption("mc.cores", 1)` after setting +#' up daemons. +loo_inform_cores_ignored <- function() { + if (isTRUE(.loo_internal$informed_cores_ignored)) { + return(invisible(NULL)) + } + .loo_internal$informed_cores_ignored <- TRUE + message( + "A mirai daemon pool is connected, so 'cores' is ignored and this call ", + "runs in parallel on the existing pool. Call mirai::daemons(0) to stop ", + "the pool if you want serial execution." + ) + invisible(NULL) +} + +#' Warn (once per session) that the `cores` argument is deprecated +#' +#' @noRd +#' @keywords internal +#' @description +#' Emitted from `loo_cores()` when the user passes `cores` explicitly (even +#' `cores = 1`), when `options(mc.cores)` is set, or when the resolved value +#' is `cores > 1`. The argument/option still works for now by starting a scoped +#' local [mirai::daemons()] pool when `cores > 1`, but users should migrate to +#' [mirai::daemons()] or [loo_mirai()]. +#' @param cores Integer number of cores requested by the user. +#' @param explicit Whether `cores` was passed explicitly by the user. +#' @param mc_cores_set Whether `options(mc.cores)` is set and `cores` was not +#' passed explicitly. +loo_deprecate_cores <- function(cores, explicit = FALSE, mc_cores_set = FALSE) { + if (!explicit && !mc_cores_set && cores <= 1L) { + return(invisible(NULL)) + } + if (isTRUE(.loo_internal$warned_cores_deprecated)) { + return(invisible(NULL)) + } + .loo_internal$warned_cores_deprecated <- TRUE + warning( + "The 'cores' argument and options('mc.cores') are deprecated and will be ", + "removed in a future release. For one-off parallel calls, 'cores' / ", + "'mc.cores' currently start a per-call local daemon pool when > 1; use ", + "mirai::daemons() before calling loo functions, or loo_mirai(..., ", + "n_daemons = k) for map-style workflows over many models.", + call. = FALSE + ) + invisible(NULL) +} + +#' Warn (once per session) about `cores` in `loo_mirai()` `args_list` +#' +#' @noRd +#' @keywords internal +loo_deprecate_cores_in_mirai <- function() { + if (isTRUE(.loo_internal$warned_cores_in_mirai)) { + return(invisible(NULL)) + } + .loo_internal$warned_cores_in_mirai <- TRUE + warning( + "Passing 'cores' in args_list to loo_mirai() is deprecated and ignored. ", + "Parallelism is controlled by loo_mirai() via n_daemons or an existing ", + "mirai::daemons() pool.", + call. = FALSE + ) + invisible(NULL) +} + +#' Prepare argument lists for [loo_mirai()] map jobs +#' +#' @noRd +#' @keywords internal +#' @description +#' Strips any user-supplied `cores` element (with a deprecation warning) and +#' forces `cores = 1` so each mapped job runs serially on its worker. Outer +#' parallelism comes only from `loo_mirai()` itself. +loo_mirai_prepare_args <- function(args) { + if ("cores" %in% names(args)) { + loo_deprecate_cores_in_mirai() + args$cores <- NULL + } + args$cores <- 1L + args +} + +#' Evaluate parallel work with an appropriate mirai daemon pool +#' +#' @noRd +#' @keywords internal +#' @description +#' Central entry point used by loo's parallel code paths to ensure a +#' [mirai::daemons()] pool exists for the duration of a computation. It is +#' deliberately a good citizen of the user's session: +#' +#' * A daemon pool is already connected (e.g. the user called +#' [mirai::daemons()] themselves, possibly with remote/HPC daemons): `code` +#' runs on the existing pool, which is left untouched. **A connected pool +#' always wins**, so the `cores` argument is ignored entirely in this case. +#' If `cores <= 1` here -- i.e. the call looked like it requested serial +#' execution -- a one-time-per-session message notes that `cores` is being +#' ignored (via `loo_inform_cores_ignored()`). +#' * `cores <= 1` with no pool connected: runs `code` serially without touching +#' daemons. +#' * Otherwise: a pool of `cores` local daemons is created for the duration of +#' `code` and automatically reset afterwards (via the scoped +#' `with(mirai::daemons(), ...)` method), so no daemon processes are left +#' running once the call returns. +#' +#' This keeps a single pool alive across the whole top-level computation +#' (rather than spinning daemons up and down for each unit of work) while +#' respecting any pool the user has already declared. Because it reuses an +#' existing pool, it is safe to nest: an inner call made while an outer call +#' already established a pool simply reuses it instead of creating another. +#' +#' @param cores Integer number of cores requested by the user. When no pool is +#' connected this is the per-call "enable parallel" switch (and the size of +#' the ephemeral pool). When a pool is already connected it is ignored -- +#' loo always uses the connected pool. +#' @param code Expression to evaluate. Lazily evaluated in the calling +#' environment, after any daemon pool has been set up. +#' @return The value of `code`. +with_loo_daemons <- function(cores, code) { + if (loo_has_pool()) { + if (cores <= 1) { + loo_inform_cores_ignored() + } + return(code) + } + if (cores <= 1) { + return(code) + } + with(mirai::daemons(cores), code) +} + +#' Run a \pkg{loo} function over argument lists in parallel with mirai +#' +#' Intended for map-style workflows such as computing [loo()], [psis()], or +#' other \pkg{loo} functions for many models in one session. +#' +#' @param fun A \pkg{loo} function (e.g. [loo()], [psis()]) to apply to each +#' element of `args_list`. +#' @param args_list A list of lists. Each element is passed to `fun` via +#' `do.call(fun, element)`. +#' @param n_daemons Number of local daemons to start for this call. When +#' `NULL` (the default), an existing [mirai::daemons()] pool is reused if +#' connected; otherwise the work runs serially. When an integer `>= 2`, a +#' local pool of that size is created for the duration of the call and +#' torn down automatically when it returns (unless a pool was already +#' connected, in which case it is reused and left untouched). +#' @return A list of results in the same order as `args_list`. +#' @details Parallelism is controlled by `loo_mirai()` (via `n_daemons` or an +#' existing daemon pool), not by a `cores` argument on the mapped calls. +#' Any `cores` element in `args_list` is deprecated, ignored, and replaced +#' with `cores = 1` so each job runs serially on its worker. +#' @export +#' @examples +#' r_eff <- relative_eff(exp(example_loglik_array())) +#' ll_list <- list(example_loglik_matrix(), example_loglik_matrix()) +#' args_list <- lapply(ll_list, function(LL) { +#' list(log_ratios = -LL, r_eff = r_eff) +#' }) +#' +#' # Serial when no pool is connected and n_daemons is NULL +#' loo_mirai(psis, args_list) +#' +#' \dontrun{ +#' # Convenience wrapper: start workers, map, tear down +#' loo_mirai(psis, args_list, n_daemons = 2) +#' +#' # Or manage the pool yourself and call loo functions directly +#' mirai::daemons(4) +#' loo_mirai(loo, lapply(ll_list, function(LL) { +#' list(x = LL, r_eff = r_eff) +#' })) +#' mirai::daemons(0) +#' } +loo_mirai <- function(fun, args_list, n_daemons = NULL) { + if (!is.function(fun)) { + stop("'fun' must be a function.", call. = FALSE) + } + if (!identical(environment(fun), asNamespace("loo"))) { + stop( + "'fun' must be a loo function (e.g. loo(), psis()).", + call. = FALSE + ) + } + if (!is.list(args_list)) { + stop("'args_list' must be a list.", call. = FALSE) + } + + run_map <- function() { + prepared <- lapply(args_list, loo_mirai_prepare_args) + if (!loo_has_pool()) { + return(lapply(prepared, function(a) do.call(fun, a))) + } + mirai::mirai_map( + prepared, + function(.args, .fun) do.call(.fun, .args), + .args = list(.fun = fun) + )[mirai::.stop] + } + + if (!is.null(n_daemons)) { + n_daemons <- as.integer(n_daemons) + if (length(n_daemons) != 1L || is.na(n_daemons) || n_daemons < 1L) { + stop( + "'n_daemons' must be a single positive integer or NULL.", + call. = FALSE + ) + } + if (n_daemons == 1L || loo_has_pool()) { + return(run_map()) + } + return(with(mirai::daemons(n_daemons), run_map())) + } + + run_map() +} + +#' Is a mirai daemon pool currently connected? +#' +#' @noRd +#' @keywords internal +#' @return `TRUE` if at least one daemon connection exists for the active +#' compute profile, otherwise `FALSE`. +loo_has_pool <- function() { + conns <- tryCatch(mirai::status()$connections, error = function(e) 0L) + isTRUE(as.integer(conns) > 0L) +} + +#' Number of workers available for chunking decisions +#' +#' @noRd +#' @keywords internal +#' @description +#' Returns the number of connected daemons when a pool exists (so chunking +#' matches the actual worker count, including user-supplied or remote pools), +#' otherwise falls back to the requested `cores`. +#' @param cores Integer number of cores requested by the user. +loo_n_workers <- function(cores) { + conns <- tryCatch(mirai::status()$connections, error = function(e) 0L) + conns <- as.integer(conns) + if (length(conns) != 1L || is.na(conns) || conns < 1L) { + return(as.integer(cores)) + } + conns +} + +#' Is the active daemon pool on the local machine? +#' +#' @noRd +#' @keywords internal +#' @description +#' Determines whether shared memory ([mori::share()]) can be used safely with +#' the active pool. Shared memory only works when workers run on the same +#' physical machine, so we only treat same-host transports as local: +#' +#' * `abstract://` and `ipc://` are same-machine inter-process transports used +#' by local [mirai::daemons()] pools, so these are treated as local. +#' * `tcp://` (and anything else) may be a remote pool, or the host URL that +#' remote SSH/HPC daemons dial back to, so it is treated as **not** local. +#' loo then falls back to ordinary serialization instead of shared memory. +#' +#' This is intentionally conservative: an incorrect "local" classification +#' would produce wrong results on a remote pool, whereas an incorrect "remote" +#' classification merely forgoes the zero-copy optimisation. +#' @return `TRUE` if the pool is confirmed local, otherwise `FALSE`. +loo_pool_is_local <- function() { + urls <- tryCatch(mirai::status()$daemons, error = function(e) NULL) + if (!is.character(urls) || length(urls) == 0L) { + return(FALSE) + } + all(grepl("^(abstract|ipc)://", urls)) +} + +#' Map a worker over elements, serially or across a mirai daemon pool +#' +#' @noRd +#' @keywords internal +#' @description +#' Single cross-platform entry point for loo's per-observation parallelism. +#' Replaces the previous platform-branching +#' [parallel::mclapply()] / [parallel::parLapply()] code paths with a single +#' [mirai::mirai_map()] path, while preserving the serial [lapply()] behaviour +#' when no daemon pool is connected. +#' +#' Object transport is chosen automatically: +#' +#' * `broadcast` objects are reused identically by every element (e.g. the +#' posterior `draws` matrix). On a local pool they are written once into +#' shared memory with [mori::share()] so each daemon maps the same physical +#' pages (zero-copy). On a remote pool, where shared memory is unavailable, +#' they are serialized instead; chunking bounds the number of copies sent to +#' roughly one per worker. +#' * Small per-call arguments are passed through `...`. +#' +#' @param X A vector or list to iterate over. Each element is passed as the +#' first argument to `FUN`. +#' @param FUN Worker function. Called as `FUN(x, , <...>)`; the +#' names in `broadcast` and `...` must match `FUN`'s formals. +#' @param ... Small constant arguments forwarded to `FUN` for every element. +#' @param cores Integer number of cores requested by the user. Used only as a +#' fallback worker count for chunking when (unexpectedly) no pool is +#' connected; it does not gate execution. Parallelism is used whenever a +#' daemon pool is connected (the connected pool always wins over `cores`). +#' @param broadcast Named list of large objects reused by every element. See +#' Description for how these are transported. +#' @param chunk Chunking strategy. `"auto"` (default) splits `X` into roughly +#' one chunk per worker to amortise per-task overhead -- best for cheap +#' per-element work over many elements. `"never"` dispatches one task per +#' element for finer load balancing -- best for expensive, uneven per-element +#' work. `"never"` is automatically promoted to `"auto"` on a remote pool +#' that carries `broadcast` objects, to avoid re-sending them per task. +#' @return A list of `FUN` results in the same order as `X`. +loo_map <- function(X, FUN, ..., cores = 1L, broadcast = list(), + chunk = c("auto", "never")) { + chunk <- match.arg(chunk) + dots <- list(...) + + if (!loo_has_pool()) { + return(do.call(lapply, c(list(X, FUN), broadcast, dots))) + } + + local_pool <- loo_pool_is_local() + if (length(broadcast) > 0L) { + if (local_pool) { + broadcast <- lapply(broadcast, mori::share) + } else if (chunk == "never") { + chunk <- "auto" + } + } + const_args <- c(broadcast, dots) + + if (chunk == "never") { + return( + mirai::mirai_map( + X, + function(.x, .FUN, .const) do.call(.FUN, c(list(.x), .const)), + .args = list(.FUN = FUN, .const = const_args) + )[mirai::.stop] + ) + } + + n_chunks <- min(loo_n_workers(cores), length(X)) + positions <- parallel::splitIndices(length(X), n_chunks) + chunks <- lapply(positions, function(p) X[p]) + chunk_results <- mirai::mirai_map( + chunks, + function(.chunk, .FUN, .const) { + lapply(.chunk, function(.x) do.call(.FUN, c(list(.x), .const))) + }, + .args = list(.FUN = FUN, .const = const_args) + )[mirai::.stop] + do.call(c, chunk_results) +} diff --git a/R/psis_approximate_posterior.R b/R/psis_approximate_posterior.R index ebd9b4a4..7c372f1a 100644 --- a/R/psis_approximate_posterior.R +++ b/R/psis_approximate_posterior.R @@ -92,7 +92,7 @@ ap_psis <- function(log_ratios, log_p, log_g, ...) { ap_psis.array <- function(log_ratios, log_p, log_g, ..., cores = getOption("mc.cores", 1)) { - cores <- loo_cores(cores) + cores <- loo_cores(cores, call = match.call()) stopifnot(length(dim(log_ratios)) == 3) log_ratios <- validate_ll(log_ratios) log_ratios <- llarray_to_matrix(log_ratios) @@ -113,7 +113,7 @@ ap_psis.matrix <- function(log_ratios, log_p, log_g, cores = getOption("mc.cores", 1)) { checkmate::assert_numeric(log_p, len = nrow(log_ratios)) checkmate::assert_numeric(log_g, len = nrow(log_ratios)) - cores <- loo_cores(cores) + cores <- loo_cores(cores, call = match.call()) log_ratios <- validate_ll(log_ratios) log_ratios <- correct_log_ratios(log_ratios, log_p = log_p, log_g = log_g) diff --git a/R/psislw.R b/R/psislw.R index 808e657f..6463c50e 100644 --- a/R/psislw.R +++ b/R/psislw.R @@ -14,12 +14,10 @@ #' @param wtrunc For truncating very large weights to \eqn{S}^`wtrunc`. Set #' to zero for no truncation. #' @param cores The number of cores to use for parallelization. This defaults to -#' the option `mc.cores` which can be set for an entire R session by -#' `options(mc.cores = NUMBER)`, the old option `loo.cores` is now -#' deprecated but will be given precedence over `mc.cores` until it is -#' removed. **As of version 2.0.0, the default is now 1 core if -#' `mc.cores` is not set, but we recommend using as many (or close to as -#' many) cores as possible.** +#' the option `mc.cores`, which can be set for an entire R session by +#' `options(mc.cores = NUMBER)`. **As of version 2.0.0, the default is now 1 +#' core if `mc.cores` is not set, but we recommend using as many (or close to +#' as many) cores as possible.** #' @param llfun,llargs See [loo.function()]. #' @param ... Ignored when `psislw()` is called directly. The `...` is #' only used internally when `psislw()` is called by the [loo()] @@ -40,7 +38,7 @@ psislw <- function(lw, wcp = 0.2, wtrunc = 3/4, ...) { .Deprecated("psis") - cores <- loo_cores(cores) + cores <- loo_cores(cores, call = match.call()) .psis <- function(lw_i) { x <- lw_i - max(lw_i) diff --git a/R/zzz.R b/R/zzz.R index 87dc9e20..41549ca5 100644 --- a/R/zzz.R +++ b/R/zzz.R @@ -5,10 +5,9 @@ "- Online documentation and vignettes at mc-stan.org/loo" ) packageStartupMessage( - "- As of v2.0.0 loo defaults to 1 core ", - "but we recommend using as many as possible. ", - "Use the 'cores' argument or set options(mc.cores = NUM_CORES) ", - "for an entire session. " + "- Parallelization uses mirai. The 'cores' argument and options('mc.cores') ", + "are deprecated; see vignette('loo2-parallel') for mirai::daemons() and ", + "loo_mirai()." ) if (os_is_windows()) { packageStartupMessage( diff --git a/_pkgdown.yml b/_pkgdown.yml index 0a216a02..43739552 100644 --- a/_pkgdown.yml +++ b/_pkgdown.yml @@ -52,6 +52,7 @@ articles: - loo2-non-factorized - loo2-lfo - loo2-large-data + - loo2-parallel - loo2-moment-matching - loo2-mixis - title: Frequently asked questions @@ -104,6 +105,7 @@ reference: - elpd - title: Other functions contents: + - loo_mirai - loo_predictive_metric - crps - elpd diff --git a/benchmark/README.md b/benchmark/README.md new file mode 100644 index 00000000..195f383e --- /dev/null +++ b/benchmark/README.md @@ -0,0 +1,154 @@ +# loo parallel benchmarks + +These scripts measure the performance of loo's parallel code paths and compare +two installed versions of the package side by side: + +- **`baseline`** — a pre-`mirai` version (the old `mclapply`/`parLapply` + backend), e.g. the released version from CRAN. +- **`new`** — the current working tree (the `mirai` + `mori` backend, including + user-managed `mirai::daemons()` session pools). + +The same user-facing calls (`psis()`, `loo()`) are timed for every version; only +the internal parallel backend differs. For the `new` version we additionally +time a **persist** mode that starts `mirai::daemons()` once and reuses the +pool across timed iterations, so we can separate per-call daemon +spawn/teardown overhead from the steady-state cost. + +## Files + +| File | Purpose | +|---|---| +| `benchmark-parallel.R` | Times `psis()`/`loo()` across cores for one installed version; writes `/tmp/bench-