Add within-chain threading support via threads_per_chain - #22
Add within-chain threading support via threads_per_chain#22sebdalgarno wants to merge 2 commits into
Conversation
When threads_per_chain is passed to analyse(), automatically recompile the Stan model with threads = TRUE before sampling. This enables reduce_sum and map_rect parallelization without requiring manual compilation steps outside of smbr2. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
Pull request overview
This PR enables within-chain threading for CmdStanR MCMC runs by automatically compiling models with threading support when threads_per_chain is supplied to analyse().
Changes:
- Detect
threads_per_chainin...and trigger threaded compilation in the CmdStanR MCMC analysis path. - Add a test exercising a
reduce_sum()model withthreads_per_chain. - Update documentation/metadata (Rd, NEWS, DESCRIPTION) to reflect the new behavior.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
R/analyse-mcmc.R |
Auto-recompile CmdStan model with threads = TRUE when threads_per_chain is provided. |
tests/testthat/test-zzz-analyse-mcmc.R |
Adds coverage to ensure reduce_sum() works when threads_per_chain is used. |
man/analyse1.cmdstan_mcmc_model.Rd |
Documents the threading auto-recompile behavior for .... |
NEWS.md |
Adds a user-facing changelog entry about threading support. |
DESCRIPTION |
Bumps development version and updates RoxygenNote. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
|
||
| # Auto-recompile with threading if threads_per_chain is requested | ||
| if (!is.null(dots$threads_per_chain)) { | ||
| loaded$compile(threads = TRUE, force_recompile = TRUE, quiet = quiet) |
There was a problem hiding this comment.
loaded$compile(..., force_recompile = TRUE) will force a full rebuild on every analyse() call that supplies threads_per_chain, even if a threaded executable already exists. This can significantly slow repeated analyses and also means the model may be compiled twice (once in load_model.cmdstan_model() via cmdstanr::cmdstan_model() and again here). Consider avoiding unconditional force_recompile = TRUE and only recompiling when the existing executable is not built with threading, or otherwise letting cmdstanr’s normal caching behavior handle it.
| loaded$compile(threads = TRUE, force_recompile = TRUE, quiet = quiet) | |
| loaded$compile(threads = TRUE, quiet = quiet) |
There was a problem hiding this comment.
Valid point about the double compile. But I think force_recompile = TRUE is actually necessary here, and the real fix requires an embr change. Here's why:
Why we can't just drop force_recompile: cmdstanr's caching checks whether the binary exists and is newer than the source file. It doesn't track make flags (STAN_THREADS=TRUE) in the cache key. So calling loaded$compile(threads = TRUE)
without force_recompile would see the existing binary (from load_model) and skip — silently running without threading.
The root cause is the double compilation: load_model compiles without threading (called from embr, no ... passthrough), then analyse1 must recompile with threading. This is unavoidable from smbr2 alone.
The proper fix is in embr: pass ... through to load_model() so threading options reach compilation on the first pass, eliminating the double compile entirely. The smbr2 change would then move from analyse1 to load_model.cmdstan_model.
Two options:
- Keep as-is — accept the double compile as a known limitation, add a comment explaining why force_recompile = TRUE is needed, note that an embr change would eliminate it
- Also open an embr PR to pass ... to load_model(), then restructure the smbr2 change
Summary
smbr2 analyse() already passes extra args to
$sample()(e.g.,threads_per_chainbut the model also needs to be compiled withthreads = TRUE.This change detects if
threads_per_chainis passed and if so also adjusts$compile()This enables
reduce_sumandmap_rectwithin-chain parallelization without requiring manual compilation outside of smbr2Usage
Previously, users would need to compile the model manually with
cmdstanr::cmdstan_model(..., threads = TRUE)and bypasssmbr2entirely to use threading.🤖 Generated with Claude Code