fix handling of cmdstan in tests for remote CI runs

Author: shmercer

R/adaptive_btl_refit.R

@@ -1944,9 +1944,9 @@
                                                  cmdstan,
                                                  seed,
                                                  model_fn = NULL) {
-  .btl_mcmc_require_cmdstanr()
   resolved_cmdstan <- .btl_mcmc_resolve_cmdstan_config(cmdstan %||% list())
   if (is.null(model_fn)) {
+    .btl_mcmc_require_cmdstanr()
     model_fn <- cmdstanr::cmdstan_model
   }
   if (!is.function(model_fn)) {
@@ -2245,6 +2245,10 @@
   diagnostics <- NULL
   mcmc_config_used <- NULL
   cmdstan_schedule_used <- NULL
+  cmdstan_fit_fn <- refit_contract_ctx[["cmdstan_fit_fn"]] %||% .adaptive_link_fit_transform_cmdstan
+  if (!is.function(cmdstan_fit_fn)) {
+    rlang::abort("`refit_contract$cmdstan_fit_fn` must be a function when provided.")
+  }
   repair_attempts <- 0L
   max_attempts <- 3L
   for (attempt in seq_len(max_attempts)) {
@@ -2259,19 +2263,19 @@
       ),
       joint_used = joint_used
     )
-    cmdstan_fit <- .adaptive_link_fit_transform_cmdstan(
+    cmdstan_fit <- cmdstan_fit_fn(
       stan_data = stan_data_base,
       variable_names = variable_names,
       cmdstan = utils::modifyList(
-        refit_contract_ctx$cmdstan %||% list(),
+        refit_contract_ctx[["cmdstan"]] %||% list(),
         list(
           chains = as.integer(cmdstan_schedule_used$chains),
           iter_warmup = as.integer(cmdstan_schedule_used$iter_warmup),
           iter_sampling = as.integer(cmdstan_schedule_used$iter_sampling)
         )
       ),
       seed = as.integer((seed + attempt * 1009L) %% .Machine$integer.max),
-      model_fn = refit_contract_ctx$cmdstan_model_fn %||% NULL
+      model_fn = refit_contract_ctx[["cmdstan_model_fn"]] %||% NULL
     )
     draws_matrix <- as.matrix(cmdstan_fit$draws_matrix)
     diagnostics <- .adaptive_link_cmdstan_validate_diagnostics(
@@ -3286,6 +3290,7 @@
       }
     }
 
+    btl_config <- out$config$btl_config %||% list()
     cross_all <- .adaptive_link_cross_edges(out, spoke_id = spoke_id, last_refit_step = NULL)
     cross_since <- .adaptive_link_cross_edges(out, spoke_id = spoke_id, last_refit_step = last_step)
     startup_gap <- .adaptive_link_phase_b_startup_gap_for_spoke(out, spoke_id = spoke_id)
@@ -3306,11 +3311,13 @@
       hub_lock_mode = lock_mode,
       hub_lock_kappa = kappa,
       shift_only_theta_treatment = theta_treatment,
-      cmdstan = out$config$btl_config$cmdstan %||% list(),
+      cmdstan = btl_config[["cmdstan"]] %||% list(),
+      cmdstan_fit_fn = btl_config[["cmdstan_fit_fn"]] %||% NULL,
+      cmdstan_model_fn = btl_config[["cmdstan_model_fn"]] %||% NULL,
       link_diagnostics_thresholds = list(
-        divergences_max = as.integer(out$config$btl_config$divergences_max %||% 0L),
-        max_rhat = as.double(out$config$btl_config$max_rhat %||% 1.01),
-        min_ess_bulk = as.double(out$config$btl_config$ess_bulk_min %||% 400)
+        divergences_max = as.integer(btl_config$divergences_max %||% 0L),
+        max_rhat = as.double(btl_config$max_rhat %||% 1.01),
+        min_ess_bulk = as.double(btl_config$ess_bulk_min %||% 400)
       )
     )
     hub_theta_init <- if (identical(refit_mode, "joint_refit") && length(hub_current) > 0L) {
@@ -5345,7 +5352,7 @@ default_btl_fit_fn <- function(state, config) {
     results = results,
     ids = ids_fit,
     model_variant = config$model_variant %||% "btl_e_b",
-    cmdstan = config$cmdstan %||% list()
+    cmdstan = config[["cmdstan"]] %||% list()
   )
 
   fit_contract <- .adaptive_btl_extract_fit_contract(fit_out)

R/bayes_btl_mcmc_adaptive.R

@@ -535,7 +535,7 @@ as_btl_fit_contract_from_mcmc <- function(mcmc_fit, ids) {
     Y = as.integer(bt_data$Y)
   )
 
-  cmdstan <- config$cmdstan %||% list()
+  cmdstan <- config[["cmdstan"]] %||% list()
   if (!is.list(cmdstan)) {
     rlang::abort("`config$cmdstan` must be a list when provided.")
   }

R/btl_mcmc_contracts.R

@@ -259,10 +259,11 @@ validate_btl_mcmc_config <- function(config) {
     "`keep_draws` must be logical.")
   .btl_mcmc_check(.btl_mcmc_intish(config$thin_draws) && config$thin_draws >= 1L,
     "`thin_draws` must be >= 1.")
-  if (!is.list(config$cmdstan)) {
+  cmdstan <- config[["cmdstan"]]
+  if (!is.list(cmdstan)) {
     rlang::abort("`config$cmdstan` must be a list when provided.")
   }
-  cmdstan_output_dir <- config$cmdstan$output_dir %||% NULL
+  cmdstan_output_dir <- cmdstan[["output_dir"]] %||% NULL
   if (!is.null(cmdstan_output_dir)) {
     .btl_mcmc_check(is.character(cmdstan_output_dir) && length(cmdstan_output_dir) == 1L,
       "`config$cmdstan$output_dir` must be a length-1 character path.")
@@ -804,7 +805,7 @@ build_round_log_row <- function(state,
   row$mcmc_cores_detected_physical <- as.integer(mcmc_config_used$cores_detected_physical %||% NA_integer_)
   row$mcmc_cores_detected_logical <- as.integer(mcmc_config_used$cores_detected_logical %||% NA_integer_)
   threads_per_chain <- mcmc_config_used$threads_per_chain %||%
-    config$cmdstan$threads_per_chain %||% 1L
+    config[["cmdstan"]][["threads_per_chain"]] %||% 1L
   row$mcmc_threads_per_chain <- as.integer(threads_per_chain %||% NA_integer_)
   row$mcmc_cmdstanr_version <- as.character(mcmc_config_used$cmdstanr_version %||% NA_character_)
   row

tests/testthat/helper-fixtures.R

@@ -135,3 +135,94 @@ make_deterministic_fit_fn <- function(ids, fit = NULL) {
     get_calls = function() env$calls
   )
 }
+
+make_test_link_cmdstan_fit_fn <- function() {
+  function(stan_data, variable_names, cmdstan, seed, model_fn = NULL) {
+    n_draws <- 4L
+    draw_offsets <- c(-0.03, -0.01, 0.01, 0.03)
+    delta_center <- if (is.numeric(stan_data$hub_ref_cross) && is.numeric(stan_data$spoke_ref_cross)) {
+      mean(as.double(stan_data$hub_ref_cross) - as.double(stan_data$spoke_ref_cross), na.rm = TRUE)
+    } else {
+      0
+    }
+    if (!is.finite(delta_center)) {
+      delta_center <- 0
+    }
+    hub_prior_signal <- mean(as.double(stan_data$hub_prior_sd %||% numeric()), na.rm = TRUE)
+    if (is.finite(hub_prior_signal)) {
+      delta_center <- delta_center + (hub_prior_signal * 0.01)
+    }
+
+    build_theta_draws <- function(base_vals, prefix) {
+      base_vals <- as.double(base_vals %||% numeric())
+      if (length(base_vals) < 1L) {
+        return(NULL)
+      }
+      out <- vapply(
+        seq_along(base_vals),
+        function(idx) base_vals[[idx]] + draw_offsets + ((idx - 1L) * 0.005),
+        numeric(n_draws)
+      )
+      colnames(out) <- paste0(prefix, "[", seq_along(base_vals), "]")
+      out
+    }
+
+    draws <- matrix(nrow = n_draws, ncol = 0L)
+    if ("delta" %in% variable_names) {
+      draws <- cbind(draws, delta = delta_center + draw_offsets)
+    }
+    if ("log_alpha" %in% variable_names) {
+      draws <- cbind(draws, log_alpha = c(-0.04, -0.01, 0.01, 0.04))
+    }
+
+    theta_hub_draws <- build_theta_draws(stan_data$hub_ref, "theta_hub")
+    if (!is.null(theta_hub_draws) &&
+      ("theta_hub" %in% variable_names || any(grepl("^theta_hub\\[", variable_names)))) {
+      keep <- if ("theta_hub" %in% variable_names) {
+        rep(TRUE, ncol(theta_hub_draws))
+      } else {
+        colnames(theta_hub_draws) %in% variable_names
+      }
+      draws <- cbind(draws, theta_hub_draws[, keep, drop = FALSE])
+    }
+
+    theta_spoke_draws <- build_theta_draws(stan_data$spoke_ref, "theta_spoke")
+    if (!is.null(theta_spoke_draws) &&
+      ("theta_spoke" %in% variable_names || any(grepl("^theta_spoke\\[", variable_names)))) {
+      keep <- if ("theta_spoke" %in% variable_names) {
+        rep(TRUE, ncol(theta_spoke_draws))
+      } else {
+        colnames(theta_spoke_draws) %in% variable_names
+      }
+      draws <- cbind(draws, theta_spoke_draws[, keep, drop = FALSE])
+    }
+
+    if (ncol(draws) < 1L) {
+      draws <- matrix(delta_center + draw_offsets, ncol = 1L)
+      colnames(draws) <- "delta"
+    }
+
+    list(
+      fit = NULL,
+      draws_matrix = draws,
+      diagnostics = list(
+        divergences = 0L,
+        max_rhat = 1.0,
+        min_ess_bulk = 1000
+      ),
+      mcmc_config_used = list(
+        chains = as.integer(cmdstan$chains %||% 4L),
+        parallel_chains = as.integer(cmdstan$parallel_chains %||% cmdstan$chains %||% 4L),
+        threads_per_chain = as.integer(cmdstan$threads_per_chain %||% 1L),
+        cmdstanr_version = "test"
+      )
+    )
+  }
+}
+
+test_link_btl_config <- function(x = list()) {
+  utils::modifyList(
+    list(cmdstan_fit_fn = make_test_link_cmdstan_fit_fn()),
+    x %||% list()
+  )
+}

tests/testthat/test-5019-console-progress-and-refit-block.R

@@ -81,7 +81,7 @@ test_that("adaptive_rank_run_live prints linking-specific refit summary lines",
             phase_a_mode = "import",
             phase_a_artifacts = artifacts
           ),
-          btl_config = list(refit_pairs_target = 2L, stability_lag = 1L),
+          btl_config = test_link_btl_config(list(refit_pairs_target = 2L, stability_lag = 1L)),
           progress = "all",
           progress_redraw_every = 1L
         )

tests/testthat/test-5026-adaptive-rank-wrapper.R

@@ -599,7 +599,7 @@ test_that("adaptive_rank wrapper supports link_one_spoke import flow", {
       phase_a_mode = "import",
       phase_a_artifacts = artifacts[c("1", "2")]
     ),
-    btl_config = list(refit_pairs_target = 2L),
+    btl_config = test_link_btl_config(list(refit_pairs_target = 2L)),
     progress = "none",
     seed = 13L
   )
@@ -643,7 +643,7 @@ test_that("adaptive_rank wrapper supports link_multi_spoke concurrent flow", {
       phase_a_mode = "import",
       phase_a_artifacts = artifacts
     ),
-    btl_config = list(refit_pairs_target = 2L),
+    btl_config = test_link_btl_config(list(refit_pairs_target = 2L)),
     progress = "none",
     seed = 17L
   )
@@ -681,7 +681,7 @@ test_that("adaptive_rank wrapper falls back to rank_raw when linked ranks are un
       hub_id = 1L,
       phase_a_mode = "run"
     ),
-    btl_config = list(refit_pairs_target = 1L),
+    btl_config = test_link_btl_config(list(refit_pairs_target = 1L)),
     progress = "none",
     seed = 31L
   )

tests/testthat/test-5049-linking-candidates-round-routing.R

@@ -714,6 +714,7 @@ test_that("coverage source propagates through selection and linking stage rows",
   state$warm_start_done <- TRUE
   state <- mark_link_phase_b_ready(state)
   state$round$staged_active <- TRUE
+  state$config$btl_config <- test_link_btl_config(state$config$btl_config %||% list())
   draws <- matrix(
     seq_along(state$item_ids),
     nrow = 4L,

tests/testthat/test-5050-linking-refit-transforms.R

@@ -10,6 +10,7 @@ make_linking_refit_state <- function(adaptive_config = list()) {
     seed = 123L,
     adaptive_config = utils::modifyList(base_cfg, adaptive_config)
   )
+  state$config$btl_config <- test_link_btl_config(state$config$btl_config %||% list())
 
   draws <- matrix(
     c(
@@ -3786,6 +3787,7 @@ test_that("linking CmdStan schedule and refit seed are stable under fixed inputs
 
 test_that("linking refit retries CmdStan effort until diagnostics pass", {
   state <- make_linking_refit_state(list(link_refit_mode = "shift_only"))
+  state$config$btl_config$cmdstan_fit_fn <- NULL
   state <- append_cross_step(state, 1L, "s21", "h1", 1L, spoke_id = 2L)
   state <- append_cross_step(state, 2L, "h2", "s22", 0L, spoke_id = 2L)
 

tests/testthat/test-5052-linking-calibration-harness.R

@@ -10,6 +10,7 @@ test_that("offline linking calibration is deterministic and writes required arti
     judge_b = 0.05,
     judge_eps = 0.03,
     n_steps = 30L,
+    btl_config = test_link_btl_config(),
     output_dir = out_dir,
     progress = "none"
   )
@@ -22,6 +23,7 @@ test_that("offline linking calibration is deterministic and writes required arti
     judge_b = 0.05,
     judge_eps = 0.03,
     n_steps = 30L,
+    btl_config = test_link_btl_config(),
     output_dir = withr::local_tempdir(),
     progress = "none"
   )
@@ -81,6 +83,7 @@ test_that("offline calibration reuses canonical production selection utilities",
     seed = 17L,
     set_sizes = c(3L, 3L),
     n_steps = 50L,
+    btl_config = test_link_btl_config(),
     progress = "none"
   )