Parallelization of loo using mirai and mori

.Rbuildignore

@@ -1,3 +1,4 @@
+^benchmark$
 ^CRAN-RELEASE$
 ^.*\.Rproj$
 ^\.Rproj\.user$

.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/*

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

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)

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: 

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

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)
 }
 

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])
+}