Merge pull request #1168 from stan-dev/gq-timing

jgabry · web-flow · commit e27a88fee335 · 2026-07-17T18:24:18.000-06:00
add generate_quantities timing support
diff --git a/NEWS.md b/NEWS.md
@@ -12,6 +12,7 @@ set to `TRUE` for an entire R session via new global options. (#1159)
 * `cmdstan_model()` no longer fails when `MAKEFLAGS` enables directory-printing 
 output while reading `STANCFLAGS` from `make`. (#1163)
 * `cmdstan_model()` now retains include paths when initialized with both a Stan file and a precompiled executable, fixing model introspection for programs that use `#include` (#1094).
+* `$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`.
diff --git a/R/csv.R b/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.
 #'
@@ -456,6 +462,7 @@ read_cmdstan_csv <- function(files,
     }
     list(
       metadata = metadata,
+      time = list(total = NA_integer_, chains = metadata$time),
       generated_quantities = draws
     )
   } else if (metadata$method == "pathfinder") {
@@ -653,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.
 #'
@@ -786,6 +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 {
+          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" &&
@@ -827,6 +854,11 @@ read_csv_metadata <- function(csv_file) {
       sampling = sampling_time,
       total = total_time
     )
+  } else if (csv_file_info$method == "generate_quantities") {
+    csv_file_info$time <- data.frame(
+      chain_id = csv_file_info$id,
+      total = total_time
+    )
   }
   csv_file_info$model <- NULL
   csv_file_info$engaged <- NULL
diff --git a/R/fit.R b/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. |
 #'
diff --git a/R/run.R b/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)
@@ -1143,6 +1143,8 @@ CmdStanGQProcs <- R6::R6Class(
         if (self$is_still_working(id) && !self$is_queued(id) && !self$is_alive(id)) {
           # if the process just finished make sure we process all
           # input and mark the process finished
+          self$process_output(id)
+          self$process_error_output(id)
           if (self$get_proc(id)$get_exit_status() == 0) {
             self$set_proc_state(id = id, new_state = 5) # mark_proc_stop will mark this process successful
           } else {
@@ -1164,8 +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 (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
diff --git a/man/CmdStanGQ.Rd b/man/CmdStanGQ.Rd
diff --git a/man/fit-method-time.Rd b/man/fit-method-time.Rd
diff --git a/man/read_cmdstan_csv.Rd b/man/read_cmdstan_csv.Rd
diff --git a/tests/testthat/resources/csv/bernoulli_ppc-1-gq-with-timing.csv b/tests/testthat/resources/csv/bernoulli_ppc-1-gq-with-timing.csv
@@ -0,0 +1,26 @@
+# stan_version_major = 2
+# stan_version_minor = 39
+# stan_version_patch = 0
+# model = bernoulli_ppc_model
+# method = generate_quantities
+#   generate_quantities
+#     fitted_params = /tmp/RtmpCvOIQ1/fittedParams-202006271227-1-b85b52.csv
+# id = 1
+# data
+#   file = /home/rok/.cmdstanr/cmdstan-2.23.0/examples/bernoulli/bernoulli.data.json
+# init = 2 (Default)
+# random
+#   seed = 123
+# output
+#   file = /tmp/RtmpCvOIQ1/bernoulli_ppc-202006271227-1-986540.csv
+#   diagnostic_file =  (Default)
+#   refresh = 250
+y_rep.1,y_rep.2,y_rep.3,y_rep.4,y_rep.5,y_rep.6,y_rep.7,y_rep.8,y_rep.9,y_rep.10
+0,0,0,0,0,0,0,0,0,0
+0,1,1,0,1,0,1,1,0,0
+0,0,0,0,0,0,0,0,0,0
+0,0,0,0,0,1,0,0,0,1
+1,0,0,0,0,0,0,0,0,0
+# 
+#  Elapsed Time: 0.123 seconds (Generated Quantities)
+# 
diff --git a/tests/testthat/resources/csv/bernoulli_ppc-2-gq-with-timing.csv b/tests/testthat/resources/csv/bernoulli_ppc-2-gq-with-timing.csv
@@ -0,0 +1,26 @@
+# stan_version_major = 2
+# stan_version_minor = 39
+# stan_version_patch = 0
+# model = bernoulli_ppc_model
+# method = generate_quantities
+#   generate_quantities
+#     fitted_params = /tmp/RtmpCvOIQ1/fittedParams-202006271227-2-b85b52.csv
+# id = 2
+# data
+#   file = /home/rok/.cmdstanr/cmdstan-2.23.0/examples/bernoulli/bernoulli.data.json
+# init = 2 (Default)
+# random
+#   seed = 456
+# output
+#   file = /tmp/RtmpCvOIQ1/bernoulli_ppc-202006271227-2-986540.csv
+#   diagnostic_file =  (Default)
+#   refresh = 250
+y_rep.1,y_rep.2,y_rep.3,y_rep.4,y_rep.5,y_rep.6,y_rep.7,y_rep.8,y_rep.9,y_rep.10
+1,0,1,0,0,0,1,0,0,1
+0,1,0,0,1,0,0,1,0,0
+1,0,0,0,0,1,0,0,0,0
+0,0,1,0,0,0,0,1,0,0
+0,1,0,0,0,0,0,0,1,0
+# 
+#  Elapsed Time: 0.456 seconds (Generated Quantities)
+# 
diff --git a/tests/testthat/test-csv.R b/tests/testthat/test-csv.R
@@ -524,6 +524,65 @@ test_that("time from read_cmdstan_csv matches time from fit$time()", {
   )
 })
 
+test_that("returning time works for gq read_cmdstan_csv from static CSV", {
+  csv_files <- test_path("resources", "csv", "bernoulli_ppc-1-gq-with-timing.csv")
+  csv_data <- read_cmdstan_csv(csv_files)
+  expect_equal(csv_data$time$total, NA_integer_)
+  expect_equal(csv_data$time$chains, data.frame(
+    chain_id = 1,
+    total = 0.123
+  ))
+
+  csv_files <- c(
+    test_path("resources", "csv", "bernoulli_ppc-1-gq-with-timing.csv"),
+    test_path("resources", "csv", "bernoulli_ppc-2-gq-with-timing.csv")
+  )
+  csv_data <- read_cmdstan_csv(csv_files)
+  expect_equal(csv_data$time$total, NA_integer_)
+  expect_equal(csv_data$time$chains, data.frame(
+    chain_id = c(1, 2),
+    total = c(0.123, 0.456)
+  ))
+})
+
+test_that("returning time works for gq read_cmdstan_csv from fit object", {
+  gq_csv <- read_cmdstan_csv(fit_gq$output_files())
+  expect_equal(gq_csv$time$total, NA_integer_)
+  checkmate::expect_data_frame(
+    gq_csv$time$chains,
+    any.missing = FALSE,
+    types = c("numeric", "numeric"),
+    nrows = fit_gq$num_chains(),
+    ncols = 2
+  )
+  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_gt(min(gq_csv$time$chains$total), 0)
+  } else {
+    # for version < 2.39 per-chain times are reported as 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_equal(
+    read_cmdstan_csv(fit_gq$output_files())$time$chains,
+    fit_gq$time()$chains,
+    ignore_attr = TRUE
+  )
+})
+
+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_)
+  expect_equal(csv_data$time$chains, data.frame(
+    chain_id = 1,
+    total = 0
+  ))
+})
+
 test_that("read_cmdstan_csv reads seed correctly", {
   opt <- read_cmdstan_csv(fit_bernoulli_optimize$output_files())
   vi <- read_cmdstan_csv(fit_bernoulli_variational$output_files())
diff --git a/tests/testthat/test-fit-gq.R b/tests/testthat/test-fit-gq.R
@@ -135,6 +135,45 @@ test_that("time() works after gq", {
     nrows = fit_gq$runset$num_procs(),
     ncols = 2
   )
+  expect_named(run_times$chains, c("chain_id", "total"))
+  # 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_gt(min(run_times$chains$total), 0)
+  } else {
+    # for CmdStan < 2.39 per-chain times are reported as 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", {
diff --git a/tests/testthat/test-model-generate_quantities.R b/tests/testthat/test-model-generate_quantities.R
@@ -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())
   )
