Skip to content

Use cmdstan config and decouple stanfit objects from RunSet - #851

Open
amas0 wants to merge 47 commits into
stan-dev:developfrom
amas0:use-save-cmdstan-config
Open

Use cmdstan config and decouple stanfit objects from RunSet#851
amas0 wants to merge 47 commits into
stan-dev:developfrom
amas0:use-save-cmdstan-config

Conversation

@amas0

@amas0 amas0 commented Apr 12, 2026

Copy link
Copy Markdown
Collaborator

PR intending to close #785 and implementing ideas discussed in #848.


Okay. I'll start by saying that the line diff is highly inflated due to needing to replace/generate large numbers of test files to be compatible with the new output files that are being consumed. The actual diff in the core code is:

amas ~/projects/cmdstanpy $ git diff --stat develop HEAD cmdstanpy | tail -n1
 16 files changed, 1597 insertions(+), 1786 deletions(-)

So still a major PR, but not as insane as the top-line numbers look.


This is a major refactor of how stanfit objects are structured and constructed. The major points are:

  1. Config JSON output that was added in Enable save_cmdstan_config #838 is now the primary source of run configuration, replacing CSV comment parsing. (Parsing the variational eta value is the only survivor). Each method now has a typed Pydantic object to validate/parse the config JSON and is stored in fit.config.
  2. Stanfit objects and their construction now decoupled from RunSets. The stanfit objects are all now dataclasses constructed directly from the output files of the cmdstan process. This relegates RunSet's role to constructing the cmdstan argument/output file name patterns.
    • CmdStanMLE takes a converged argument on construction sourced from the retcodes of the process (assumed true when from files).
  3. Stanfit objects now share a generic StanFit base class which collects a good chunk of code that was getting repeated. This breaks down further into SingleFileFit (MLE/VB/Laplace/Pathfinder) and MultiChainFit (Sample/GQ).
  4. Handful of noteworthy public API changes downstream from all of these changes
    • from_csv is now from_output_files
    • save_csvfile is now save_output_files
    • fit.runset is gone
  5. As we now build fit objects from_output_files, we require config JSONs next to the CSV files
  6. Tests updated to include proper output files for cases where they would now be expected (config json/metric files).

Also, since we are on version 2.39 now it seemed a good time to bump the minimum version to 2.37 (this simplified some things).

There are some things that I think can still be improved, but this PR is already close to getting out of hand and I'd like to defer those to try to focus on getting this one "done". Just as note, these are:

  • One consequence of now relying on a greater number of output files is the discovery of those output files has gotten clumsy. The logic in stanfit/__init__.py bothers me. When we pulled everything out of a CSV, we didn't really have to worry about it.
    That being said, I think the solution here might be to promote the config to be the primary file for resolving the others. It includes explicit references to other output files and a proper Stanfit could be built starting from it. Right now we have to crawl around the output dir to find the appropriate output files. The config file, since it is generated by cmdstan, should always have accurate references to the file locations (provided they hadn't been moved).
  • The base StanFit objects that I introduced aren't explicitly ABCs, but they could be. Not sure if there's a reason one way or another? I chose not to make it an ABC for now.

Would really appreciate some thoughts on these changes @WardBrian, especially if any of these changes would cause issues for how this library is used with dev versions of cmdstan.

Disclaimer: AI was used throughout this PR -- critically for helping with regenerating the appropriate test files and applying initial prototype structures across the different methods. All code in this PR was reviewed by me before marking it as ready.

@WardBrian WardBrian left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A few comments/questions, but overall I like the look of the code!

Responding to your comments in the PR:

  • Runset being left out of the result entirely makes sense. We will have to see how this feels when we tackle something with more moving parts like the MCMC class, but for now I agree.
  • Repr changes make sense.
  • Stdout file is good. We may want to also do the same thing with diagnostic files and profiling files, which gets us closer to just including the whole Runset again... maybe we have a small dataclass that is like a cut-down runset and stores the relevant files, can be re-used across all of them?
  • I agree, it makes sense to replace csv with output across the board, both in the save functions and the from_csv function. This would make sense going forward anyway, c.f. stan-dev/design-docs#58
  • Yep, totally agree that InferenceMetadata might want to take a hike when all is said and done. The header could be stored in each class, or in the cut-down-Runset idea ("FileInfo"?)

Comment thread cmdstanpy/stanfit/pathfinder.py Outdated
Comment thread cmdstanpy/stanfit/pathfinder.py Outdated
Comment thread cmdstanpy/stanfit/pathfinder.py Outdated
Comment thread cmdstanpy/stanfit/__init__.py Outdated
f'found {len(csvfiles)}'
)
csv_file = csvfiles[0]
config_file = os.path.splitext(csv_file)[0] + '_config.json'

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I assume this is temporary but it's worth putting a comment saying so, since this is kind of a nasty assumption

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah definitely temporary. I was planning on doing away with this function entirely once the full changes are in.

@amas0

amas0 commented Apr 16, 2026

Copy link
Copy Markdown
Collaborator Author

@WardBrian Working through some of the other methods.

For CmdStanMLE, it looks like our typical approach to checking for convergence is looking at the return codes via the runset. Is there a more natural way to do this for reading from the output files directly?

I see we could maybe parse that info out of the stdout file, but I don't love that.

@WardBrian

Copy link
Copy Markdown
Member

Hm, that is indeed tricky. I think we should add the piece of information that currently ends up in stdout (Optimization terminated normally:/with error:) to the output csv, but that won't help with existing, released versions.

Looking at the code, it looks like the only place we actually check the return codes is in model.py, before we call the CmdStanMLE constructor, so maybe it is fine?

@amas0

amas0 commented Apr 16, 2026

Copy link
Copy Markdown
Collaborator Author

Yeah. In my local implementation, I take is_converged as an argument to the constructor. We have it in model.py so it should be fine in practice, but it means we can't get full reconstitution from just output files for MLE (unless we parse stdout).

It does seem like the kind of thing that should be available in the CSV? Like a converged__ column or something? Not a major issue for now, just doesn't make things as clean as I'd like.

@WardBrian

Copy link
Copy Markdown
Member

but it means we can't get full reconstitution from just output files for MLE (unless we parse stdout).

This is true for the existing from_csv function, isn't it?

@amas0

amas0 commented Apr 16, 2026

Copy link
Copy Markdown
Collaborator Author

Yeah good point, this is the snippet from from_csv:

            for i in range(len(runset._retcodes)):
                runset._set_retcode(i, 0)
            return CmdStanMLE(runset)

Looks like in recreating the runset, we just manually set all return codes to 0, which is just assuming convergence I believe.

@WardBrian

Copy link
Copy Markdown
Member

Yeah, so we won't be doing any worse, at least!

@amas0

amas0 commented Apr 29, 2026

Copy link
Copy Markdown
Collaborator Author

@WardBrian Is there an equivalent to saving the metric json for the variational method? Looking at CmdStanVB now and we parse the variational eta value from the stancsv comments, but afaik we don't have an alternative way to source it for now.

@bob-carpenter

Copy link
Copy Markdown
Member

There's no preconditioning in our two variational inference implementations, ADVI and Pathfinder, so there's no equivalent of a metric/mass matrix/preconditioner.

@WardBrian

Copy link
Copy Markdown
Member

@bob-carpenter I believe @amas0 was asking just if we wrote the metadata from the algorithm as a json, similar to the metric in sampling, not if there was actually a preconditioner

Unfortunately to my knowledge the answer to this is no. It's a shame to have to keep all the comment-reading machinery around for something as rarely-needed as the final eta value from ADVI...

@bob-carpenter

Copy link
Copy Markdown
Member

Sorry for the confusion.

Keeping things around for eta is also silly in that estimating eta is one of the major weaknesses of our ADVI implementation (another being restriction to diagonal or dense approximations, a third being the high variance of low-sample ELBO evaluations).

A better approach for eta might be to treat it as another column in the tabular output like step size. Like step size, it will change as the algorithm progresses then hopefully stabilize.

@amas0

amas0 commented Apr 30, 2026

Copy link
Copy Markdown
Collaborator Author

Yeah, in particular, I was referring to this section that appears in the output CSV if you run the variational method:

lp__,log_p__,log_g__,theta
# Stepsize adaptation complete.
# eta = 0.1
0,0,0,0.5838444

I just meant that we output the algorithm stepsize info to the CSV in the same way that we write the sampler adaptation output (as comment lines after the column headers). But in the sample case, we can output equivalent info to a metric JSON file, but that doesn't seem to be the case for the variational method here.

So, I guess the question is what do we want to do with this in the near term? I think it makes sense to have it be an additional column as Bob suggested? I suppose the question is whether it is important enough to maintain the comment parsing in the near term or should we drop keeping track of it for now until an alternative source is available?

@WardBrian

WardBrian commented May 1, 2026

Copy link
Copy Markdown
Member

I'm not sure it makes sense as a column, since the advi output is only the one row of the varational parameters, in which eta has already been adapted, and then forward sampling iterations

Github code search doesn't seem to find any public code which actually calls .eta, but I'm not sure this is really a good measure.

Maybe we preserve a dumber version of comment parsing just for that, i.e. a function that just does

for line in file:
  if line.startswith("# eta ="):
    return float(line.split('=')[1])

Rather than all the present machinery? We can also try to add a structured/json output like for the metric, but that would take longer to thread through the cmdstan code and we'd still need something in the meantime

Chains were always labelled 1..N, so a run with explicit chain_ids came
back with the wrong ids and selecting a chain from the rebuilt fit gave a
different chain than in the original. CmdStan records the id in each
config; a single-process run has one config holding the first chain's id,
with the remaining chains following on from it.
amas0 added 13 commits August 24, 2026 19:05
StanConfig[X] parametrized generics have no importable module-level
name, so instances could not be pickled. Four fit classes worked
around this with __getstate__/__setstate__ hacks that serialized the
config as a dict; CmdStanPathfinder and CmdStanLaplace were missing
the workaround and could not be pickled at all.

Introduce named subclasses (SampleRunConfig, OptimizeRunConfig, etc.)
with a from_json constructor, use them throughout the fit classes, and
delete the per-class pickle workarounds. parse_config now only handles
the auto-detect case used by from_output_files.

Adds pickle round-trip tests for pathfinder and laplace fits.
Introduce stanfit/base.py with StanFit, holding the members every fit
container duplicates (stan_variables, attribute access, pickling,
column_names, save_output_files), and SingleFileFit, holding the
common behavior of fits whose draws live in one Stan CSV file (lazy
draw loading, draws, stan_variable, method_variables, create_inits,
and the from_files parsing boilerplate).

save_output_files gains the writability check and missing-file
handling previously only present in CmdStanMCMC, applied uniformly
via the _FILE_ATTRS class attribute.

CmdStanPathfinder is the first class ported; it now contains only its
method-specific members (from_files, is_resampled, repr). The other
fit classes will follow in subsequent commits.

Adds a save_output_files regression test for pathfinder.
Keeps only the Laplace-specific members: the mode attribute and its
resolution in from_files, draws_pd, draws_xr, and repr. Draw loading,
variable access, create_inits, pickling, and save_output_files are
inherited. laplace.py shrinks from 358 to 161 lines.
The variational mean and sample are now views of the first and
remaining rows of the inherited draws array rather than separate
fields, and the eta comment-parse happens in the _assemble override.
Keeps the VB-specific members: eta, the variational_params and
variational_sample accessors, the mean/sample toggle on
stan_variable(s), create_inits over the sample, and repr.
vb.py shrinks from 322 to 213 lines.
The MLE and saved iterations are now views of the last row and the
whole of the inherited draws array rather than separate fields; the
base _assemble is used unchanged. The repeated convergence and
save_iterations warnings are factored into small helpers, and the
pd/dict accessors are derived from the np ones. Keeps the
MLE-specific members: converged, the optimized_* accessors, the
iterations toggle on stan_variable(s), create_inits from the
estimate, and repr. mle.py shrinks from 361 to 232 lines.
Move draws() from SingleFileFit down to CmdStanPathfinder and
CmdStanLaplace. Inheriting it on CmdStanVB and CmdStanMLE would have
exposed the raw CSV rows, where the first row is the variational mean
and the rows are optimizer iterations respectively; too easy to
mistake for a posterior sample. method_variables remains shared.
MultiChainFit holds the shared behavior of fits with one Stan CSV
file per chain: the chains property, list-normalizing from_files
parsing, and cross-chain config validation, with a _comparable_config
hook that subclasses extend with the settings that affect their draw
layout.

CmdStanGQ keeps only its GQ-specific members: previous_fit and the
draw accessors that merge it with the generated quantities. Its
config validation needs were exactly the base ones, so it vanishes
entirely. gq.py shrinks from 815 to 668 lines.
Keeps the sampler-specific members: metric handling, warmup/draw
count validation, divergence and treedepth tabulation, summary and
diagnose, the draws accessors, and chain-aware create_inits. The
cross-chain config validation reduces to a _comparable_config
override extending the base checks with num_samples, num_warmup,
save_warmup, and thin. The metric, diagnostic, and profile file
lists join the shared save_output_files via _FILE_ATTRS.
mcmc.py shrinks from 1014 to 842 lines.
The refactor onto shared StanFit base classes moved many public
methods (stan_variables, save_output_files, create_inits, etc.) into
the bases, so the api docs need :inherited-members: to keep rendering
them on each fit class. The base classes themselves stay out of the
documented API.

Also clarifies why cross-chain config validation parses sibling
configs with the class of the held config.
InferenceMetadata.from_csv read the entire Stan CSV, draws included,
just to take its header line; a scan_header helper stops at the first
non-comment line instead.

from_output_files now raises for generate_quantities outputs rather
than warning and returning None, which removes None from its return
type
@amas0
amas0 marked this pull request as ready for review August 25, 2026 04:04
@amas0

amas0 commented Aug 25, 2026

Copy link
Copy Markdown
Collaborator Author

Okay, I think I'm happy enough with where this stands to have it ready for review. At the very least to have someone else other than me look at this.

@WardBrian

Copy link
Copy Markdown
Member

@amas0 it is on my list for this week!

@WardBrian WardBrian left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Starting small, really only reviewed a few pieces for now and decided to stop with my questions

)


def _sidecar(csv_file: Path, kind: str) -> Path:

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[nit] don't love the name, maybe "_accompanying_json" or something?

Comment on lines +132 to +134
named for its first chain. Outputs written by running CmdStan directly
with a single output name and ``num_chains`` follow a different naming
scheme and are not currently supported.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this comment correct? Because that will definitely be a blocker, we don't want to drop support for the built-in parallelism

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

To clarify, this refers specifically to if you run a model with built-in parallelism manually in cmdstan without cmdstanpy at all and then try to load it in to cmdstanpy. When you run it via cmdstanpy, we set the output filenames in the cmdstan call to conform to an expected naming convention.

As long as you are trying to load in a cmdstan output that was run via cmdstanpy, you will be fine.

Instantiate a CmdStan object from a the Stan CSV files from a CmdStan run.
CSV files are specified from either a list of Stan CSV files or a single
Instantiate a CmdStan object from the output files of a CmdStan run.
Output files are specified as either a list of Stan CSV files or a single

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I feel like we should allow the _config.json to be in the list, and use that if so rather than going looking for it?

thin=thin,
save_warmup=save_warmup,
if method_name == 'sample':
# Collect the sidecar JSONs that exist next to each CSV. A

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the word sidecar is the new em-dash. I almost never see it outside llm comments/thought traces (unless you count by bartenders)

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah...I had a similar reaction. I satisfied myself when I saw it was apparently a thing, but I can clean that bit up.

Comment on lines +206 to +207
builder = _SINGLE_FILE_METHODS.get(method_name)
if builder is None:

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't know if we really have enough of these to justify this over just an if/else chain or a match statement

Comment thread cmdstanpy/__init__.py
CmdStanPathfinder,
CmdStanVB,
from_csv,
from_output_files,

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's unfortunate we don't have a release in the 1.x cycle that formally deprecated from_csv, we might want to keep it around as an alias as a result. Thoughts?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

My approach with this PR was to specifically target a 2.0 release, which is why I opted for a clean compatibility break here. That being said an alias that replicates the previous experience of using from_csv probably wouldn't be an issue?

Comment thread cmdstanpy/model.py

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

General comment: I really like how the changes look in this file, to me it really motivates the rest of the changes to be able to have this all be nicer!

@amas0

amas0 commented Aug 26, 2026

Copy link
Copy Markdown
Collaborator Author

@WardBrian you've rightfully picked at the code around file discovery so far. I wanted to raise this point I made in PR description (in case you missed it) to get your thoughts:

One consequence of now relying on a greater number of output files is the discovery of those output files has gotten clumsy. The logic in stanfit/__init__.py bothers me. When we pulled everything out of a CSV, we didn't really have to worry about it.
That being said, I think the solution here might be to promote the config to be the primary file for resolving the others. It includes explicit references to other output files and a proper Stanfit could be built starting from it. Right now we have to crawl around the output dir to find the appropriate output files. The config file, since it is generated by cmdstan, should always have accurate references to the file locations (provided they hadn't been moved).

I had a working version of a change where we essentially start from the config.json, which has enough information in it to reconstruct the output file names without needing to crawl around the directory (I believe). The current structure we have goes through a few heuristics to identify which CSVs are draws (vs diagnostics/profile) and to resolve which chain id each draw csv corresponds to. Starting from a config(s) would remove that step -- it would also enable easier support for the case where you ran cmdstan outside of cmdstanpy and are trying to load fit in (related to your comment above). It would also handle the case where there are multiple fits in the same directory (which is a current problem area if you provide a directory).

The big drawback is the potential brittleness of the path references in the config. If a fit is generated and the output files are moved, then those references are no longer correct. One mitigation there would be to have the stanfit's save_output_files(...) method also update the path references in the config when moved. So managing fit "bundles" from within cmdstanpy would be safe, but not necessarily outside of that.

We do now require a config to be present to construct a stanfit, so I don't think it's too crazy to consider this route, but it would certainly be a shift from past behavior.

@WardBrian

Copy link
Copy Markdown
Member

If I recall correctly, in the multi threaded case the output field in the config.json is still not a literal valid path, but instead the “template” without the chain id. Is that not the case?

@amas0

amas0 commented Aug 26, 2026

Copy link
Copy Markdown
Collaborator Author

Ah yeah, I see - that's how it is for cmdstan directly. I had confused myself with how we run it through cmdstanpy.

In these changes I take advantage of the 2.37 output file csv to have full paths for each chain csv in the config.

@WardBrian

Copy link
Copy Markdown
Member

Ah, because we now pass the names we’d like as a comma separated list? That’s nice

@amas0

amas0 commented Aug 26, 2026

Copy link
Copy Markdown
Collaborator Author

Yeah, it makes it so that if we run a model via cmdstanpy then we get a pretty clean output set. I think it was also necessary to maintain the naming convention discussed in #835 now that we have the config in the mix?

The underlying tension here I think is between whether the output files were generated by cmdstan directly or indirectly via cmdstanpy. In the latter case, we get a good amount of control to have a clean output to work with. In the former case, it's trickier.

@WardBrian

Copy link
Copy Markdown
Member

How about:

  • if the config files come from cmdstanpy or follow our desired format, you can pass a subset (maybe just the config) or a directory and we figure out the rest
  • if they didn’t, you need to list every single one of them in the call to the function

Reasonable?

@amas0

amas0 commented Aug 26, 2026

Copy link
Copy Markdown
Collaborator Author

I'm good with that.

Comment thread test/data/runset-good/bern-1_config.json
Comment thread test/test_laplace.py Outdated
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Path forward on IO operations

3 participants