add mirai/mori backend for memory-efficient parallelization

.Rbuildignore

@@ -11,3 +11,4 @@
 ^arf.*\.tgz$
 ^doc$
 ^Meta$
+^bench$

.gitignore

@@ -14,3 +14,5 @@ attic/
 branch/
 /doc/
 /Meta/
+bench/logs/
+bench/results/

DESCRIPTION

@@ -1,6 +1,6 @@
 Package: arf
 Title: Adversarial Random Forests
-Version: 0.2.5
+Version: 0.3.0
 Authors@R: c(
     person("Marvin N.", "Wright", , "cran@wrig.de", role = c("aut", "cre"),
            comment = c(ORCID = "0000-0002-8542-6291")),
@@ -9,7 +9,9 @@ Authors@R: c(
     person("Kristin", "Blesch", role = "aut",
            comment = c(ORCID = "0000-0001-6241-3079")),
     person("Jan", "Kapar", role = "aut",
-           comment = c(ORCID = "0009-0000-6408-2840"))
+           comment = c(ORCID = "0009-0000-6408-2840")),
+    person("Lukas", "Burk", , "cran@lukasburk.de", role = "ctb",
+           comment = c(ORCID = "0000-0001-7528-3795"))
   )
 Description: Adversarial random forests (ARFs) recursively partition data
     into fully factorized leaves, where features are jointly independent.
@@ -33,19 +35,22 @@ Imports:
     ranger,
     stringr,
     truncnorm
-Suggests: 
+Suggests:
     doFuture,
     doParallel,
     ggplot2,
     knitr,
+    mirai,
     mlbench,
+    mori,
     palmerpenguins,
+    pkgload,
     rmarkdown,
     testthat (>= 3.0.0),
     tibble
-VignetteBuilder: 
+VignetteBuilder:
     knitr
+Config/roxygen2/version: 8.0.0
 Config/testthat/edition: 3
 Encoding: UTF-8
 Roxygen: list(markdown = TRUE)
-RoxygenNote: 7.3.2

NEWS.md

@@ -1,3 +1,10 @@
+# arf 0.3.0
+* Add mirai/mori parallel backend as an alternative to foreach/doParallel (#62)
+	* Shares large read-only inputs (training data, forest, learned parameters) across workers via mori, lowering memory use in `adversarial_rf()`, `forde()`, `forge()`, `expct()`, and `lik()`
+	* Enable with active mirai daemons or `options(arf.backend)`, see `?arf-options`
+* Reduce peak memory of `expct()` by processing conditions in bounded blocks (up to ~16x lower in internal benchmarks on large forests with many conditions, tune via `options(arf.block_rows)`, see `?arf-options`)
+* Reduce memory and dispatch overhead in `forde()` (fused per-tree parameter pass, per-tree coverage) and `lik()` (per-batch reduction)
+
 # arf 0.2.5
 * Export sample_from_leaves() for intra-leaf marginal sampling
 
@@ -27,4 +34,4 @@
 * Speed boost for the adversarial resampling step 
 * Early stopping option for adversarial training
 * alpha parameter for regularizing multinomial distributions in forde
-* Unified treatment of colnames with internal semantics (y, obs, tree, leaf)
+* Unified treatment of colnames with internal semantics (y, obs, tree, leaf)

R/adversarial_rf.R

@@ -16,8 +16,10 @@
 #' @param prune Impose \code{min_node_size} by pruning? 
 #' @param verbose Print discriminator accuracy after each round? Will also show 
 #'   additional warnings.
-#' @param parallel Compute in parallel? Must register backend beforehand, e.g. 
-#'   via \code{doParallel} or \code{doFuture}; see examples.
+#' @param parallel Compute in parallel? Enables multithreaded ranger training
+#'   (no backend needed) and parallelizes the pruning step, which requires a
+#'   registered \code{foreach} backend (\code{doParallel}, \code{doFuture}) or
+#'   active \code{mirai} daemons. See \code{\link{arf-options}}.
 #' @param ... Extra parameters to be passed to \code{ranger}.
 #' 
 #' @details 
@@ -92,6 +94,9 @@
 #' # ... or with doFuture
 #' doFuture::registerDoFuture()
 #' future::plan("multisession", workers = 4)
+#'
+#' # ... or with mirai (shares large read-only inputs across workers via mori)
+#' mirai::daemons(4)
 #' }
 #' 
 #' @seealso
@@ -183,41 +188,52 @@ adversarial_rf <- function(
     }
   }
 
-  # Prune leaves to ensure min_node_size w.r.t. real data
+  # Prune leaves to ensure min_node_size w.r.t. real data. Per-tree work lives in
+  # arf_prune_tree() (prune_workers.R). This is a meaningful share of runtime once
+  # ranger training is threaded (the prune loop is pure R, unaffected by
+  # ranger's num.threads), so it gets the same backend treatment as the rest.
   if (isTRUE(prune)) {
     pred <- stats::predict(rf0, x_real, type = 'terminalNodes')$predictions + 1L
-    prune <- function(tree) {
-      # Nodes to prune are leaves which contain fewer than min_node_size real samples
-      out <- rf0$forest$child.nodeIDs[[tree]]
-      leaves <- which(out[[1]] == 0L)
-      to_prune <- leaves[!(leaves %in% which(tabulate(pred[, tree]) >= min_node_size))]
-      while(length(to_prune) > 0) {
-        if (1 %in% to_prune) {
-          # Never prune the root
-          break
-        }
-        for (tp in to_prune) {
-          # Find parent
-          parent <- which((out[[1]] + 1L) == tp)
-          if (length(parent) > 0) {
-            # If node to prune (tp) is the left child of parent, replace left child with right child
-            out[[1]][parent] <- out[[2]][parent]
-          } else {
-            # If node to prune (tp) is the right child of parent, replace right child with left child
-            parent <- which((out[[2]] + 1L) == tp)
-            out[[2]][parent] <- out[[1]][parent]
-          }
-        }
-        # If both children of a parent are to be pruned, prune the parent in the next round
-        # This happens if both children have been pruned
-        to_prune <- which((out[[1]] + 1L) %in% to_prune)
-      }
-      return(out)
+    prune_one <- function(b) {
+      arf_prune_tree(b, rf0$forest$child.nodeIDs, pred, min_node_size)
     }
-    if (isTRUE(parallel)) {
-      rf0$forest$child.nodeIDs <- foreach(b = seq_len(num_trees)) %dopar% prune(b)
+    use_mirai <- FALSE
+    if (num_trees > 1) {
+      backend <- arf_select_backend(parallel)
+      use_mirai <- identical(backend, "mirai")
+    }
+    if (use_mirai) {
+      # arf_prune_tree's body is base-R, so pass it as an object (daemons need no
+      # arf loaded). Share the two big read-only objects (pred is n x num_trees;
+      # child.nodeIDs is the forest) once via mori. Chunk contiguously and c() so
+      # the flat result keeps tree order 1..num_trees (unname: mirai_map names
+      # chunks, but child.nodeIDs must stay an unnamed list).
+      pred_shared <- mori::share(pred)
+      child_shared <- mori::share(rf0$forest$child.nodeIDs)
+      chunks <- arf_tree_chunks(num_trees, mirai::status()$connections)
+      chunk_fn <- function(trees, worker, child_nodeIDs, pred, min_node_size) {
+        lapply(trees, worker, child_nodeIDs = child_nodeIDs,
+               pred = pred, min_node_size = min_node_size)
+      }
+      # Both functions are base-R with explicit args: strip their environments
+      # before shipping. chunk_fn's would otherwise be THIS frame (rf0, dat,
+      # x_real: hundreds of MB serialized into every task); arf_prune_tree's
+      # namespace env would force daemons to load arf.
+      # See the closure note above arf_mirai_tree_map() in mirai_helpers.R.
+      environment(chunk_fn) <- globalenv()
+      prune_worker <- arf_prune_tree
+      environment(prune_worker) <- globalenv()
+      res <- mirai::mirai_map(chunks, chunk_fn,
+                              .args = list(worker = prune_worker,
+                                           child_nodeIDs = child_shared,
+                                           pred = pred_shared,
+                                           min_node_size = min_node_size))[]
+      arf_stop_on_mirai_error(res)
+      rf0$forest$child.nodeIDs <- unname(do.call(c, res))
+    } else if (isTRUE(parallel) && num_trees > 1) {
+      rf0$forest$child.nodeIDs <- foreach(b = seq_len(num_trees)) %dopar% prune_one(b)
     } else {
-      rf0$forest$child.nodeIDs <- foreach(b = seq_len(num_trees)) %do% prune(b)
+      rf0$forest$child.nodeIDs <- foreach(b = seq_len(num_trees)) %do% prune_one(b)
     }
   }
 

R/arf-options.R

@@ -0,0 +1,103 @@
+#' arf package options
+#'
+#' Options controlling the parallel backend and its messaging, set via
+#' \code{\link{options}}.
+#'
+#' @details
+#' \describe{
+#'   \item{\code{arf.backend}}{Parallel backend used when \code{parallel = TRUE}:
+#'     \code{"foreach"} or \code{"mirai"}. If unset, arf uses \code{"mirai"} when
+#'     mirai daemons are running and \code{"foreach"} otherwise.}
+#'   \item{\code{arf.verbose}}{Report the selected backend once per backend
+#'     configuration per session? Default \code{TRUE}; set \code{FALSE} to
+#'     silence.}
+#'   \item{\code{arf.block_rows}}{Cap on rows materialized per block of
+#'     conditions in \code{\link{expct}}. Default \code{5e6}. Lower it to
+#'     trade speed for memory on large forests with many conditions.}
+#' }
+#'
+#' \code{arf.block_rows} does not affect results, only peak memory and speed.
+#' For memory-constrained hardware, combine a small daemon count with a lower
+#' \code{arf.block_rows} and the \code{batch}/\code{stepsize} arguments of
+#' \code{\link{lik}}, \code{\link{forge}} and \code{\link{expct}}.
+#'
+#' The \code{"foreach"} backend uses whatever adapter is registered (e.g.
+#' \code{doParallel}, \code{doFuture}). The \code{"mirai"} backend uses
+#' \code{mirai} daemons and shares large read-only inputs (training data,
+#' forest, learned parameters) across workers via \code{mori}, so workers do
+#' not each copy them. Speed is comparable between the backends. The memory
+#' benefit is largest for the tree-parallel operations (\code{\link{forde}},
+#' \code{\link{adversarial_rf}}) on large forests with many workers, where
+#' \code{foreach} memory grows with the worker count and \code{mirai} stays
+#' much flatter (roughly half to a third at 16 workers in internal benchmarks).
+#' For small workloads the daemon pool adds a fixed overhead that can outweigh
+#' the sharing, so prefer \code{"mirai"} at scale and either backend otherwise.
+#' Daemons started via \code{future.mirai} (e.g.
+#' \code{plan(future.mirai::mirai_multisession)}) are detected like any other
+#' mirai daemons, so futureverse users get the \code{"mirai"} backend
+#' automatically. All backend packages are in Suggests; install the ones you
+#' use.
+#'
+#' Reproducibility of stochastic operations (\code{\link{forge}}, categorical
+#' \code{\link{expct}}) under parallel execution: \code{\link{set.seed}} only
+#' governs the calling process, not the workers. With the \code{"mirai"}
+#' backend, seed the daemons instead: \code{mirai::daemons(n, seed = 42)}
+#' gives reproducible results provided the daemon count, the seed and the
+#' sequence of calls on a fresh daemon pool are kept fixed (changing the
+#' daemon count changes how work is chunked and therefore the random stream
+#' assignment). This also applies to pools started via \code{future.mirai}:
+#' \code{future}'s own seed machinery covers only work dispatched through its
+#' API, which arf's backend bypasses, and \code{plan()} does not accept a
+#' \code{seed} argument, so seed the pool with \code{mirai::daemons()}
+#' directly. For the \code{"foreach"} backend, register \code{doRNG} on top of
+#' the adapter (\code{doRNG::registerDoRNG(42)} after
+#' \code{registerDoParallel()} or \code{registerDoFuture()}): results are then
+#' reproducible, identical across adapters, and independent of the worker
+#' count, provided \code{stepsize} is set explicitly (its default depends on
+#' the worker count). Without \code{doRNG}, \code{doFuture} flags the
+#' stochastic operations with "UNRELIABLE VALUE" warnings. Sequential
+#' execution (\code{parallel = FALSE}) with \code{set.seed} is exact as
+#' always.
+#'
+#' @examples
+#' \dontrun{
+#' arf <- adversarial_rf(iris)
+#'
+#' # foreach backend
+#' doParallel::registerDoParallel(cores = 4)
+#' psi <- forde(arf, iris)
+#'
+#' # mirai backend: start daemons, then call as usual
+#' mirai::daemons(4)
+#' psi <- forde(arf, iris)
+#' mirai::daemons(0)  # shut down when done
+#'
+#' # futureverse: future.mirai daemons are detected automatically
+#' future::plan(future.mirai::mirai_multisession, workers = 4)
+#' psi <- forde(arf, iris)
+#'
+#' # force a backend regardless of what is registered
+#' options(arf.backend = "mirai")
+#'
+#' # silence the backend message
+#' options(arf.verbose = FALSE)
+#'
+#' # reproducible parallel sampling with mirai: seeded daemons
+#' # (fixed daemon count, fresh pool)
+#' evi <- data.frame(Species = sample(levels(iris$Species), 100, replace = TRUE))
+#' mirai::daemons(4, seed = 42)
+#' # stepsize = evidence rows per step; 25 = 100 conditions / 4 workers,
+#' # i.e. the default sizing, made explicit
+#' x <- forge(psi, n_synth = 1, evidence = evi, stepsize = 25)
+#' mirai::daemons(0)
+#'
+#' # reproducible parallel sampling with foreach: doRNG on top of the
+#' # adapter; set stepsize explicitly (its default depends on worker count)
+#' doParallel::registerDoParallel(cores = 4)
+#' doRNG::registerDoRNG(42)
+#' x <- forge(psi, n_synth = 1, evidence = evi, stepsize = 25)
+#' }
+#'
+#' @name arf-options
+#' @aliases arf.backend arf.verbose arf.block_rows
+NULL

R/arf-package.R

@@ -36,6 +36,9 @@
 #' # ... or with doFuture
 #' doFuture::registerDoFuture()
 #' future::plan("multisession", workers = 4)
+#'
+#' # ... or with mirai (shares large read-only inputs across workers via mori)
+#' mirai::daemons(4)
 #' }
 "_PACKAGE"