fixes to GQ time handling and documentation

Author: jgabry

NEWS.md

@@ -11,6 +11,7 @@ set to `TRUE` for an entire R session via new global options. (#1159)
 * `save_metric_files()` now gives an informative error when metric files were not created and keeps saved metric files after the fitted model is garbage-collected. (#1021)
 * `cmdstan_model()` no longer fails when `MAKEFLAGS` enables directory-printing 
 output while reading `STANCFLAGS` from `make`. (#1163)
+* `$generate_quantities()` now reports per-process execution times with CmdStan 2.39 or newer, and `read_cmdstan_csv()` returns these times from standalone generated quantities CSV files. (#1168)
 * `laplace()` no longer overwrites the internally generated optimizer CSV when
 `mode = NULL` and `output_basename` is supplied. The internally generated
 optimizer CSV now uses the filename `<output_basename>-mode-1.csv`.

R/csv.R

@@ -78,6 +78,12 @@
 #' For [standalone generated quantities][model-method-generate-quantities] the
 #' returned list also includes the following components:
 #'
+#' * `time`: Run time information for the individual processes, with one row in
+#' the `chains` data frame per CSV file. The returned object is the same as for
+#' the [$time()][fit-method-time] method except the total run time can't be
+#' inferred directly from the CSV files (they may have been generated in
+#' parallel) and is therefore `NA`. For CmdStan versions before 2.39 the
+#' individual process times are reported as zero.
 #' * `generated_quantities`: A [`draws_array`][posterior::draws_array] of
 #' the generated quantities.
 #'
@@ -654,6 +660,21 @@ for (method in unavailable_methods_CmdStanFit_CSV) {
 
 # csv reading internals ---------------------------------------------------
 
+parse_generated_quantities_time <- function(line) {
+  suffix <- "seconds (Generated Quantities)"
+  line <- trimws(line)
+  if (!startsWith(line, "Elapsed Time:") || !endsWith(line, suffix)) {
+    return(NULL)
+  }
+  time <- trimws(sub(suffix, "", line, fixed = TRUE))
+  time <- trimws(sub("Elapsed Time:", "", time, fixed = TRUE))
+  time <- suppressWarnings(as.double(time))
+  if (is.na(time)) {
+    return(NULL)
+  }
+  time
+}
+
 #' Reads the sampling arguments and the diagonal of the
 #' inverse mass matrix from the comments in a CSV file.
 #'
@@ -787,10 +808,11 @@ read_csv_metadata <- function(csv_file) {
           tmp <- gsub("seconds (Total)", "", tmp, fixed = TRUE)
           tmp <- trimws(gsub(" Elapsed Time: ", "", tmp, fixed = TRUE))
           total_time <- as.numeric(tmp)
-        } else if (grepl("(Generated Quantities)", tmp, fixed = TRUE)) {
-          tmp <- gsub("seconds (Generated Quantities)", "", tmp, fixed = TRUE)
-          tmp <- trimws(gsub("Elapsed Time:", "", tmp, fixed = TRUE))
-          total_time <- as.numeric(tmp)
+        } else {
+          generated_quantities_time <- parse_generated_quantities_time(tmp)
+          if (!is.null(generated_quantities_time)) {
+            total_time <- generated_quantities_time
+          }
         }
         if (!is.null(csv_file_info$method) &&
             csv_file_info$method == "diagnose" &&

R/fit.R

@@ -1090,19 +1090,26 @@ CmdStanFit$set("public", name = "metric_files", value = metric_files)
 #'
 #' @name fit-method-time
 #' @aliases time
-#' @description Report the run time in seconds. For MCMC additional information
-#'   is provided about the run times of individual chains and the warmup and
-#'   sampling phases. For Laplace approximation the time only include the time
-#'   for drawing the approximate sample and does not include the time
-#'   taken to run the `$optimize()` method.
+#' @description Report the run time in seconds. For MCMC and standalone
+#'   generated quantities additional information is provided about the run
+#'   times of individual chains or processes. For MCMC, timing information is
+#'   also provided for the warmup and sampling phases. For Laplace approximation
+#'   the time only includes the time for drawing the approximate sample and does
+#'   not include the time taken to run the `$optimize()` method.
 #'
 #' @return
 #' A list with elements
-#' * `total`: (scalar) The total run time. For MCMC this may be different than
-#' the sum of the chain run times if parallelization was used.
-#' * `chains`: (data frame) For MCMC only, timing info for the individual
-#' chains. The data frame has columns `"chain_id"`, `"warmup"`, `"sampling"`,
-#' and `"total"`.
+#' * `total`: (scalar) The total run time. For MCMC and standalone generated
+#' quantities this may differ from the sum of the individual run times if
+#' parallelization was used.
+#' * `chains`: (data frame) For MCMC and standalone generated quantities,
+#' timing information for the individual chains. For MCMC the data frame has
+#' columns `"chain_id"`, `"warmup"`, `"sampling"`, and `"total"`. For standalone
+#' generated quantities, each row corresponds to one fitted-parameter CSV file
+#' and one CmdStan process, and the data frame has columns `"chain_id"` and
+#' `"total"`. Variational or optimization input therefore produces one row.
+#' With CmdStan versions before 2.39, standalone generated quantities process
+#' times are reported as zero.
 #'
 #' @seealso [`CmdStanMCMC`], [`CmdStanMLE`], [`CmdStanVB`], [`CmdStanGQ`]
 #'
@@ -2250,7 +2257,7 @@ CmdStanPathfinder$set("public", name = "lp_approx", value = lp_approx)
 #'
 #'  |**Method**|**Description**|
 #'  |:----------|:---------------|
-#'  [`$time()`][fit-method-time] | Report the total run time. |
+#'  [`$time()`][fit-method-time] | Report total and process-specific run times. |
 #'  [`$output()`][fit-method-output] | Return the stdout and stderr of all chains or pretty print the output for a single chain. |
 #'  [`$return_codes()`][fit-method-return_codes]  |  Return the return codes from the CmdStan runs. |
 #'

R/run.R

@@ -357,7 +357,7 @@ CmdStanRun <- R6::R6Class(
       } else if (self$method() == "generate_quantities") {
         chain_time <- data.frame(
           chain_id = self$procs$proc_ids()[self$procs$is_finished()],
-          total = self$procs$proc_total_time()[self$procs$is_finished()]
+          total = as.vector(self$procs$proc_total_time())
         )
 
         time <- list(total = self$procs$total_time(), chains = chain_time)
@@ -1166,10 +1166,13 @@ CmdStanGQProcs <- R6::R6Class(
         if (nzchar(line)) {
           if (self$proc_state(id) == 1 && grepl("refresh = ", line, perl = TRUE)) {
             self$set_proc_state(id, new_state = 1.5)
-          } else if (grepl("Elapsed Time:", line, fixed = TRUE)) {
-            private$proc_total_time_[[id]] <- as.double(trimws(sub("Elapsed Time:", "", sub("seconds (Generated Quantities)", "", line, fixed = TRUE), fixed = TRUE)))
-          } else if (self$proc_state(id) >= 2 && private$show_stdout_messages_) {
-            cat("Chain", id, line, "\n")
+          } else {
+            generated_quantities_time <- parse_generated_quantities_time(line)
+            if (!is.null(generated_quantities_time)) {
+              private$proc_total_time_[[id]] <- generated_quantities_time
+            } else if (self$proc_state(id) >= 2 && private$show_stdout_messages_) {
+              cat("Chain", id, line, "\n")
+            }
           }
         } else {
           # after the metadata is printed and we found a blank line

man/CmdStanGQ.Rd

@@ -558,21 +558,22 @@ test_that("returning time works for gq read_cmdstan_csv from fit object", {
   expect_named(gq_csv$time$chains, c("chain_id", "total"))
   if (cmdstan_version() >= "2.39.0") {
     # per-chain times should be non-zero (parsed from CmdStan timing output)
-    expect_true(all(gq_csv$time$chains$total > 0))
+    expect_gt(min(gq_csv$time$chains$total), 0)
   } else {
     # for version < 2.39 per-chain times are reported as 0
-    expect_true(all(gq_csv$time$chains$total == 0))
+    expect_equal(gq_csv$time$chains$total, rep(0, fit_gq$num_chains()))
   }
 })
 
 test_that("gq time from read_cmdstan_csv matches time from fit_gq$time()", {
-  expect_equivalent(
+  expect_equal(
     read_cmdstan_csv(fit_gq$output_files())$time$chains,
-    fit_gq$time()$chains
+    fit_gq$time()$chains,
+    ignore_attr = TRUE
   )
 })
 
-test_that("returning time is NULL for gq CSV without timing", {
+test_that("gq CSV without timing returns zero chain time", {
   csv_files <- test_path("resources", "csv", "bernoulli_ppc-1-gq.csv")
   csv_data <- read_cmdstan_csv(csv_files)
   expect_equal(csv_data$time$total, NA_integer_)

man/fit-method-time.Rd

@@ -136,17 +136,46 @@ test_that("time() works after gq", {
     ncols = 2
   )
   expect_named(run_times$chains, c("chain_id", "total"))
-  # total wall-clock time is always positive
-  expect_true(run_times$total > 0)
+  # CmdStanR measures total wall-clock time for every CmdStan version.
+  expect_gt(run_times$total, 0)
   if (cmdstan_version() >= "2.39.0") {
     # CmdStan >= 2.39 reports per-chain timing for generated quantities
-    expect_true(all(run_times$chains$total > 0))
+    expect_gt(min(run_times$chains$total), 0)
   } else {
     # for CmdStan < 2.39 per-chain times are reported as 0
-    expect_true(all(run_times$chains$total == 0))
+    expect_equal(run_times$chains$total, rep(0, fit_gq$num_chains()))
   }
 })
 
+test_that("time() aligns successful gq process ids and times", {
+  # Reproduce a partial failure without running noisy CmdStan subprocesses.
+  # The first process fails so the successful process ids and times can be
+  # checked for the misalignment caused by filtering the times twice.
+  MockGQProcs <- R6::R6Class(
+    classname = "MockGQProcs",
+    inherit = CmdStanGQProcs,
+    public = list(
+      initialize = function() {
+        super$initialize(num_procs = 3)
+        private$proc_state_[] <- c(7, 6, 6)
+        private$proc_total_time_[] <- c(0, 0.1, 0.2)
+        private$total_time_ <- 0.3
+      }
+    )
+  )
+  procs <- fit_gq$runset$procs
+  on.exit(fit_gq$runset$procs <- procs)
+  fit_gq$runset$procs <- MockGQProcs$new()
+
+  expect_equal(
+    fit_gq$runset$time(),
+    list(
+      total = 0.3,
+      chains = data.frame(chain_id = c(2L, 3L), total = c(0.1, 0.2))
+    )
+  )
+})
+
 test_that("fitted_params_files() works", {
   expect_equal(
     fit_gq$fitted_params_files(),

man/read_cmdstan_csv.Rd

@@ -77,8 +77,11 @@ test_that("generate_quantities works with draws_array", {
 test_that("generate_quantities works with VB and draws_matrix", {
   fit <- testing_fit("bernoulli", method = "variational", seed = 123)
   expect_gq_output(
-    mod_gq$generate_quantities(data = data_list, fitted_params = fit)
+    fit_gq <- mod_gq$generate_quantities(data = data_list, fitted_params = fit)
   )
+  run_times <- fit_gq$time()
+  expect_equal(run_times$chains$chain_id, 1)
+  expect_gte(run_times$chains$total, 0)
   expect_gq_output(
     mod_gq$generate_quantities(data = data_list, fitted_params = fit$draws())
   )