Repository navigation
Implemention discussion: using Stan config JSON throughout library and other changes #848
Description
As part of ongoing work to revamp the IO of the library (#785), we would like to transition from sourcing Stan config info out of comment lines in the Stan CSV to the structured configuration json files (#714). Since #838, the config files are now generated as part of running any Stan method and saved to the output dir along with other files. This includes storing the location of the config in the RunSet object.
In line with work in #844 which introduced pydantic as a dependency for parsing and validating data sourced from json files, we'd like to do the same here. A underlying challenge here is that the config information is used across a number of different methods and not always in a consistent way. There is also a desire to maintain a flexible interface that can accommodate new or in-development outputs/methods. An ideal solution would standardize the IO configuration across methods while not constraining future development.
It's not obvious to me the right path, so I'd like to use this thread to get some thoughts down and discuss ideas before moving forward with a particular implementation.
Where do we actually use the configuration in the library?
To my understanding, the primary use of parsing the Stan configuration is the from_csv() function to reconstruct Fit object from the Stan CSV files. When constructing a Fit object directly from a method like model.sample(), the configuration stored is sourced from the arguments passed by the user. (I believe the only time there is an explicit comparison between the output configuration and the user args is in CmdStanMCMC._validate_csv_files()).
So, under the current approach, we really only need to parse and validate configuration insofar as it is used to create Fit objects from CmdStan output files (this would need a change from a from_csv function to a from_output_files function). The good news is that the required fields to do this is relatively small.
As a side note, in theory I think there may be some value in validating user-passed arguments against the arguments parsed out of config files. Could be the kind of thing that would catch bugs that are otherwise hiding.
What should a configuration object look like? How should validation be done?
When it comes to using the configuration, one way I see that this can be done is a single pydantic model for all the different methods, such as:
from typing import Literal
from pydantic import BaseModel, ConfigDict
class ConfigInfo(BaseModel):
"""Structured representation of a config JSON file"""
model_config = ConfigDict(extra="allow")
model_name: str
stan_major_version: int
stan_minor_version: int
stan_patch_version: int
method: Literal["sample", "optimize", "variational", "laplace", "pathfinder"]
# Sample
algorithm: str | None = None
num_samples: int | None = None
num_warmup: int | None = None
save_warmup: bool = False
thin: int = 1
max_depth: int | None = None
# Optimize
save_iterations: bool = False
jacobian: bool = False
# Variational
iter: int | None = None
grad_samples: int | None = None
elbo_samples: int | None = None
eta: float | None = None
tol_rel_obj: float | None = None
eval_elbo: int | None = None
output_samples: int | None = None
# Laplace
mode: str | None = None
draws: int | None = None
# Pathfinder
num_draws: int | None = None
num_paths: int | None = None
psis_resample: bool | None = None
calculate_lp: bool | None = None This keeps everything together, but co-mingles the various config properties and any validation process would need to do so in a method-aware manner. The model_config = ConfigDict(extra="allow") allows the pydantic model to allow fields that are not explicitly specified and handles there being many more fields in the output than are needed by the library (but still gives access to them if desired).
Something like this would probably be the easiest to integrate within the current library in the sense that it would require the fewest changes.
We could alternatively create separate models for each method, e.g.:
class BaseConfig(BaseModel):
"""Common fields for all methods."""
model_config = ConfigDict(extra="allow")
model_name: str
stan_major_version: int
stan_minor_version: int
stan_patch_version: int
class SampleConfig(BaseConfig):
method: Literal["sample"] = "sample"
algorithm: str
num_samples: int
num_warmup: int
save_warmup: bool = False
thin: int = 1
max_depth: int | None = None
class OptimizeConfig(BaseConfig):
method: Literal["optimize"] = "optimize"
algorithm: str
save_iterations: bool = False
jacobian: bool = False
... which, in combination with pydantic's discriminated union feature would give us method discriminated config objects that I think would be nicer to work with. This would allow us to keep validation method specific.
This would probably require more changes throughout the library, but I think it would be easier to reason about within the library.
Thinking about further changes
Since these IO changes are aligned with trying to do work for a 2.0 release, I wonder if there are other organizational changes in the library that could be made that would make it a bit easier to maintain/develop against. I find that not having a consistent interface/protocol for the different methods to output (like adhering to a StanFit protocol/ABC) is a pain point. Primarily, I think this could clean up the IO pipeline by making things more consistent among the other methods.
In any case, I would appreciate some thoughts (@WardBrian ) on what a reasonable approach for this would be? Or if you think a better approach is something that I had not suggested, I'm all ears.
In some sense I'm trying to close the gap between the state that this library is in today and an ideal version that would be created from scratch today. Not so easy a task.
Activity
WardBrian commentedon Jan 27, 2026
Thank you for these thoughts! I was meaning to write out a similar post about pre-2.0 work.
My instinct is to say we could also go entirely the other way and have the individual fit objects constructed entirely from the output files in all cases, no additional information. This makes basically
from_output_files(or at least the individual branches of it) into the constructors, and the CmdStanModel methods would basically beThe top-level from_output_files just minimally determines which class to call the from_files method on and then calls it. Nice and simple, no need to reconstruct a Runset or try to divine any arguments specifically inside this function.
I am hesitant to get too into checking the user's arguments are mirrored in the output json. Not only because any potential bug would likely be caught elsewhere, but also because of things like
adapt engaged=1andadapt engaged=trueboth meaning the same thingRe: shared vs discriminated config. I can see cases for both.
The shared one with all the fields being optional essentially feels like a just-slightly-more-typed version of
json.loading the config and reading out whatever keys we need in a given Fit object. This is nice in the sense that it is 'minimal': if the config schema changed in some way, but your code path doesn't need any of the bit that changed, your code will still run.The discriminated union is how I'd be inclined to go, but it does have the drawback of any upfront validation could stop something that was going to work due to an 'incidental' incompatibility. Maybe that's fine, and we haven't seen the config change all that much (when it does, it's usually additive only).
Certainly not easy, though admirable! Part of the difficulty with something like a StanFit ABC is that Stan itself is not terribly consistent with it's outputs. For example, ADVI gives you means and approximate draws, a protocol would only expose one or the other, or have to include some extra methods (or arguments to methods) that aren't used by the others. For an example of how this can go wrong, see how scipy.stats calls all the arguments to it's distributions
locandscale, for 'consistency', even when the distribution is not a location-scale family, and the mess that createsThat being said there are some 'obvious' improvements one could make in this direction, like whether things are accessed by properties or methods, etc.
amas0 commentedon Jan 27, 2026
I'm glad to see you suggest this, because it occurred to me while thinking this through that having the two separate paths to create the same object was unnecessary (I could easily imagine confusing situations where the
CmdStanMCMCobject from the user args is inconsistent with the one created from the output files).This would also cleanly separate out the CmdStanModel and the methods to run the model from the actual output of the model. With the underlying Stan process sitting in between.
I don't think this would be too much of a lift! Mostly I'd wonder what the user experience for this should be.
Do you have any anticipated examples of this in mind? I think there are defensive ways this can be managed, e.g. a flag to fall back to a BaseConfig that doesn't validate. I imagine the biggest concern would be trying to ensure that an in-development or new version of Stan doesn't always need corresponding changes to CmdStanPy to work.
That's fair. I suppose a question for this is what does the intersection of the output of the various methods look like? Is there something there that can serve as a foundation to build off of. Like do all methods now minimally output a CSV and (with
save_cmdstan_config) a config file?WardBrian commentedon Jan 28, 2026
Precisely!
I like how nutpie does this (obviously I don't think
pauseorresumecould be implemented, butwaitwould be natural andabortcould return the runset before passing to from_files)Yes, that's the big one. In particular, it's nice if new cmdstan versions work with the currently released cmdstanpy (even if the new features may be inaccessible). The example I'm thinking of concretely is a few versions ago cmdstan updated from reporting boolean flags as 0/1 to
true/false(though, this did break cmdstanpy already at the time, so perhaps I'm being over-worried that a pydantic validator would make it more likely or not)Yes, though some produce multiple CSVs (sampling, pathfinder with certain flags that I don't think we currently expose)
I almost made a base class because of how similar the code between CmdStanLaplace and CmdStanPathfinder is, so maybe that's a good place to look for commonalities as a start.
WardBrian commentedon Apr 9, 2026
@amas0 is this still something you're interested in working on?
amas0 commentedon Apr 9, 2026
Yes. Apologies for being slow on it!
I hadn't touched this in a while, but just last week I started looking at it again. Revisiting this, I settled on this being the key idea:
To this end, I started working on a refactor of the
CmdStanPathfinderobject (as a simple example) to be reconstructed entirely from output files. My plan was to share what that looked like as a template to be applied to the rest of the objects. The config JSON would, of course, be part of those files used in reconstruction.If that seems reasonable, I'll treat this nudge as enough motivation to focus on it in the near term.
WardBrian commentedon Apr 9, 2026
No worries! I appreciate you continuing to contribute to cmdstanpy :)
That sounds like a great place to start, IMO
WardBrian commentedon Sep 22, 2026
#851