MPI sampling

.github/workflows/R-CMD-check.yaml

@@ -53,6 +53,12 @@ jobs:
           mingw32-make --version
           Get-Command mingw32-make | Select-Object -ExpandProperty Definition
         shell: powershell
+
+      - name: Install MPI
+        if: runner.os == 'Linux'
+        run: |
+          sudo apt-get install -y openmpi-bin
+          echo "CMDSTANR_RUN_MPI_TESTS=TRUE" >> $GITHUB_ENV
 
       - uses: r-lib/actions/setup-r@master
         with:

.github/workflows/Test-coverage.yaml

@@ -27,7 +27,11 @@ jobs:
       - uses: r-lib/actions/setup-pandoc@master
 
       - name: Install Ubuntu dependencies
-        run: sudo apt-get install libcurl4-openssl-dev
+        run: |
+          sudo apt-get install libcurl4-openssl-dev
+          sudo apt-get install -y openmpi-bin
+          echo "CMDSTANR_RUN_MPI_TESTS=TRUE" >> $GITHUB_ENV
+
       - name: Query dependencies
         run: |
           install.packages('remotes')

R/model.R

@@ -914,6 +914,180 @@ sample_method <- function(data = NULL,
 }
 CmdStanModel$set("public", name = "sample", value = sample_method)
 
+#' Run Stan's MCMC algorithms with MPI
+#'
+#' @name model-method-mpi-sample
+#' @aliases mpi_sample
+#' @family CmdStanModel methods
+#'
+#' @description The `$mpi_sample()` method of a [`CmdStanModel`] object runs the
+#'   default MCMC algorithm in CmdStan (`algorithm=hmc engine=nuts`) with MPI 
+#'   (STAN_MPI makefile flag), to produce a set of draws from the posterior
+#'   distribution of a model conditioned on some data.
+#' 
+#'   In order to use MPI with Stan, an MPI implementation must be installed.
+#'   For Unix systems the most commonly used implementations are MPICH and OpenMPI.
+#'   The implementations provide an MPI C++ compiler wrapper (for example mpicxx), 
+#'   which is required to compile the model.
+#'
+#'   An example of compiling with STAN_MPI:
+#'   ```
+#'   cpp_options = list(STAN_MPI = TRUE, CXX="mpicxx", TBB_CXX_TYPE="gcc")
+#'   mod <- cmdstan_model("model.stan", cpp_options = cpp_options)
+#'   ```
+#'   The C++ options that need supplied to the compile call are: 
+#'   - `STAN_MPI`: Enables the use of MPI with Stan
+#'   - `CXX`: The name of the MPI C++ compiler wrapper (typicall mpicxx)
+#'   - `TBB_CXX_TYPE`: The C++ compiler the MPI wrapper wraps. Typically gcc on
+#'   Linux and clang on macOS.
+#' 
+#'   In the call to the `$mpi_sample()` method, we can additionally provide
+#'   the name of the MPI launcher (`mpi_cmd`), which defaults to "mpiexec",
+#'   and any other MPI launch arguments. In most cases, it is enough to
+#'   only define the number of processes with `mpi_args = list("n" = 4)`.
+#' 
+#'   An example of a call of `$mpi_sample()`:
+#'   ```
+#'   cpp_options = list(STAN_MPI = TRUE, CXX="mpicxx", TBB_CXX_TYPE="gcc")
+#'   fit <- mod$mpi_sample(data_list, mpi_args = c("-n", 4))
+#'   ```
+#'
+#' @section Usage:
+#'   ```
+#'   $mpi_sample(
+#'     data = NULL,
+#'     mpi_cmd = "mpiexec",
+#'     mpi_args = NULL,
+#'     seed = NULL,
+#'     refresh = NULL,
+#'     init = NULL,
+#'     save_latent_dynamics = FALSE,
+#'     output_dir = NULL,
+#'     chains = 4,
+#'     parallel_chains = getOption("mc.cores", 1),
+#'     chain_ids = seq_len(chains),
+#'     iter_warmup = NULL,
+#'     iter_sampling = NULL,
+#'     save_warmup = FALSE,
+#'     thin = NULL,
+#'     max_treedepth = NULL,
+#'     adapt_engaged = TRUE,
+#'     adapt_delta = NULL,
+#'     step_size = NULL,
+#'     metric = NULL,
+#'     metric_file = NULL,
+#'     inv_metric = NULL,
+#'     init_buffer = NULL,
+#'     term_buffer = NULL,
+#'     window = NULL,
+#'     fixed_param = FALSE,
+#'     sig_figs = NULL,
+#'     validate_csv = TRUE,
+#'     show_messages = TRUE
+#'   )
+#'   ```
+#'
+#' @section Arguments:
+#'   * `mpi_cmd`: (character vector) The MPI launcher used for launching MPI processes.
+#'     The default launcher is `mpiexec`.
+#'   * `mpi_args`: (list) A list of arguments to use when launching MPI processes.
+#'     For example, mpi_args = list("n" = 4) launches the executable as
+#'     `mpiexec -n 4 model_executable`, followed by CmdStan arguments
+#'     for the model executable.
+#'   * `data`, `seed`, `refresh`, `init`, `save_latent_dynamics`, `output_dir`,
+#'     `chains`, `parallel_chains`, `chain_ids`, `iter_warmup`, `iter_sampling`,
+#'     `save_warmup`, `thin`, `max_treedepth`, `adapt_engaged`, `adapt_delta`,
+#'     `step_size`, `metric`, `metric_file`, `inv_metric`, `init_buffer`,
+#'     `term_buffer`, `window`, `fixed_param`, `sig_figs`, `validate_csv`,
+#'     `show_messages`:
+#'      Same as for the [`$sample()`][model-method-sample] method.
+#'
+#' @section Value: The `$mpi_sample()` method returns a [`CmdStanMCMC`] object.
+#'
+#' @template seealso-docs
+#' @inherit cmdstan_model examples
+#'
+NULL
+mpi_sample_method <- function(data = NULL,
+                          mpi_cmd = "mpiexec",
+                          mpi_args = NULL,
+                          seed = NULL,
+                          refresh = NULL,
+                          init = NULL,
+                          save_latent_dynamics = FALSE,
+                          output_dir = NULL,
+                          chains = 1,
+                          parallel_chains = getOption("mc.cores", 1),
+                          chain_ids = seq_len(chains),
+                          iter_warmup = NULL,
+                          iter_sampling = NULL,
+                          save_warmup = FALSE,
+                          thin = NULL,
+                          max_treedepth = NULL,
+                          adapt_engaged = TRUE,
+                          adapt_delta = NULL,
+                          step_size = NULL,
+                          metric = NULL,
+                          metric_file = NULL,
+                          inv_metric = NULL,
+                          init_buffer = NULL,
+                          term_buffer = NULL,
+                          window = NULL,
+                          fixed_param = FALSE,
+                          sig_figs = NULL,
+                          validate_csv = TRUE,
+                          show_messages = TRUE) {
+
+  if (fixed_param) {
+    chains <- 1
+    parallel_chains <- 1
+    save_warmup <- FALSE
+  }
+
+  checkmate::assert_integerish(chains, lower = 1, len = 1)
+  checkmate::assert_integerish(parallel_chains, lower = 1, null.ok = TRUE)
+  checkmate::assert_integerish(chain_ids, lower = 1, len = chains, unique = TRUE, null.ok = FALSE)
+  sample_args <- SampleArgs$new(
+    iter_warmup = iter_warmup,
+    iter_sampling = iter_sampling,
+    save_warmup = save_warmup,
+    thin = thin,
+    max_treedepth = max_treedepth,
+    adapt_engaged = adapt_engaged,
+    adapt_delta = adapt_delta,
+    step_size = step_size,
+    metric = metric,
+    metric_file = metric_file,
+    inv_metric = inv_metric,
+    init_buffer = init_buffer,
+    term_buffer = term_buffer,
+    window = window,
+    fixed_param = fixed_param
+  )
+  cmdstan_args <- CmdStanArgs$new(
+    method_args = sample_args,
+    model_name = strip_ext(basename(self$exe_file())),
+    exe_file = self$exe_file(),
+    proc_ids = chain_ids,
+    data_file = process_data(data),
+    save_latent_dynamics = save_latent_dynamics,
+    seed = seed,
+    init = init,
+    refresh = refresh,
+    output_dir = output_dir,
+    validate_csv = validate_csv,
+    sig_figs = sig_figs
+  )
+  cmdstan_procs <- CmdStanMCMCProcs$new(
+    num_procs = chains,
+    parallel_procs = parallel_chains,
+    show_stderr_messages = show_messages
+  )
+  runset <- CmdStanRun$new(args = cmdstan_args, procs = cmdstan_procs)
+  runset$run_cmdstan_mpi(mpi_cmd, mpi_args)
+  CmdStanMCMC$new(runset)
+}
+CmdStanModel$set("public", name = "mpi_sample", value = mpi_sample_method)
 
 #' Run Stan's optimization algorithms
 #'

R/run.R

@@ -20,7 +20,6 @@ CmdStanRun <- R6::R6Class(
       }
       invisible(self)
     },
-
     num_procs = function() self$procs$num_procs(),
     proc_ids = function() self$procs$proc_ids(),
     exe_file = function() self$args$exe_file,
@@ -150,6 +149,10 @@ CmdStanRun <- R6::R6Class(
       }
     },
 
+    run_cmdstan_mpi = function(mpi_cmd, mpi_args) {
+      private$run_sample_(mpi_cmd, mpi_args)
+    },
+
     # run bin/stansummary or bin/diagnose
     # @param tool The name of the tool in `bin/` to run.
     # @param flags An optional character vector of flags (e.g. c("--sig_figs=1")).
@@ -222,10 +225,15 @@ CmdStanRun <- R6::R6Class(
 
 
 # run helpers -------------------------------------------------
-.run_sample <- function() {
+.run_sample <- function(mpi_cmd = NULL, mpi_args = NULL) {
   procs <- self$procs
   on.exit(procs$cleanup(), add = TRUE)
-
+  if (!is.null(mpi_cmd)) {
+    if (is.null(mpi_args)) {
+      mpi_args = list()
+    }
+    mpi_args[["exe"]] <- self$exe_file()
+  }
   # add path to the TBB library to the PATH variable
   if (cmdstan_version() >= "2.21" && os_is_windows()) {
     path_to_TBB <- file.path(cmdstan_path(), "stan", "lib", "stan_math", "lib", "tbb")
@@ -261,7 +269,9 @@ CmdStanRun <- R6::R6Class(
         id = chain_id,
         command = self$command(),
         args = self$command_args()[[chain_id]],
-        wd = dirname(self$exe_file())
+        wd = dirname(self$exe_file()),
+        mpi_cmd = mpi_cmd,
+        mpi_args = mpi_args
       )
       procs$mark_proc_start(chain_id)
       procs$set_active_procs(procs$active_procs() + 1)
@@ -475,12 +485,21 @@ CmdStanProcs <- R6::R6Class(
     get_proc = function(id) {
       private$processes_[[id]]
     },
-    new_proc = function(id, command, args, wd) {
+    new_proc = function(id, command, args, wd, mpi_cmd = NULL, mpi_args = NULL) {
+      if (!is.null(mpi_cmd)) {
+        exe_name <- mpi_args[["exe"]]
+        mpi_args[["exe"]] <- NULL
+        mpi_args_vector <- c()
+        for (i in names(mpi_args)) {
+          mpi_args_vector <- c(paste0("-", i), mpi_args[[i]], mpi_args_vector)
+        }
+        args = c(mpi_args_vector, exe_name, args)
+        command <- mpi_cmd
+      }
       private$processes_[[id]] <- processx::process$new(
         command = command,
         args = args,
         wd = wd,
-        echo_cmd = FALSE,
         stdout = "|",
         stderr = "|"
       )