Use cmdstan config and decouple stanfit objects from RunSet

[amas0]

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]

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"?)

[amas0]

@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]

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]

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]

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]

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]

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

[amas0]

@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]

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]

@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]

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]

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]

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

[WardBrian]

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

[amas0]

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]

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]

I'm good with that.

[WardBrian]

Ok, I'm up to 96/103 'viewed' on this, which means I mostly have only the big files like GQ, MCMC, and base left. A few questions in the mean time

[amas0]

Okay I implemented a refactor of the from_output_files flow that I think clarifies the approach (although it is still admittedly complex).

To start, I finally the loading from file logic into its own stanfit/from_outputs.py file and cleaned out __init__.py. Within this new file we have all the logic of how we go from a set of output files -> StanFit object. The two pathways are (from your suggestion):

• Fit file discovery based on expected naming convention/config output names. This can be done by passing a directory, a config json path, or a stan csv path. Each of these try to get to a config and then uses the file names in the config to collect the appropriate other files (which are assumed to be next to the config).
• Explicit file passing. This expects a list of file paths to be passed and will attempt to classify the files into known kinds and build an appropriate fit from there. This does no file discovery, so if you, for example take a multi-chain sample run and only pass in the output files for a single chain, you will get a single chain fit back.

This diagram shows visually the valid paths of how one can build a fit via this function:

flowchart TD
    A["from_output_files"]
    B["Single path provided"]
    C["All files explicitly passed"]
    A --> B
    A --> C
    D["Directory"]
    E["Config JSON"]
    F["Stan CSV"]
    G["Find sibling config JSON"]
    B --> D
    B --> E
    B --> F
    F --> G
    H["Build FitFiles manifest"]
    E --> H
    G --> H
    I["Find all valid Configs"]
    D --> I
    J["One unique valid config (multiple configs from the same sample fit produce the same config)"]
    K["Two unique valid configs (Laplace edge case with mode fit)"]
    I --> J
    I --> K
    J --> H
    K --> H
    L["Classify each input file by kind"]
    C --> L
    M["Verify minimally valid set of files included (no discovery)"]
    L --> M
    M --> H
    N["Construct StanFit from FitFiles"]
    H --> N

Loading

If one wishes to be more explicit, of course, one can always call a particular StanFit's from_files method directly.

[amas0]

Okay I think I addressed the comments you had.

There's a section of the docs that needs to be rewritten if we decide to go through with these changes, so I'll get to that in a bit.

[WardBrian]

Sorry for my delay! I'm quite happy with how things are working

[WardBrian]

It feels like this could be simplified even more using the new base class. For example, should _draws_start be in the base classes instead? (I think it's more or less equivalent to the new _draws_for_inits?)

[amas0]

I gave a shot at looking at what _draws_start in the base class would look like, but I didn't feel it was much of an improvement.

That being said, I did make a handful of changes that I think are cleanups of some of the longstanding code in here (also fixed up some bugs that AI flagged).

Let me know if you think there are any other pathways you wanna go here.

[WardBrian]

This file is a lot of code but it all looks pretty reasonable. I do wonder if it should be split up or if the parts that are e.g. super-Laplace specific be moved to the laplace.py file

[amas0]

I did another pass on this file -- some of the laplace stuff made sense to move out to laplace.py. The file probably could be broken up further, but I wasn't sure if splitting this logic across multiple files was better than just a ~700 line file? If you have a particular suggestion, I can see how that might play out.

[amas0]

As a side note, I made an update to the docsrc conf.py to remove a check on a nearly 10 year old sphinx version? I think we can probably assume that we'll have newer versions than that?

[WardBrian]

That seems like a safe bet…

I haven’t forgotten about this PR, just had a bunch of other stuff to do with the 2.40 release. Hopefully next week I will have enough time to get this review done and the PR merged

[amas0]

No worries! I figure if it took me the better part of a year to get this together, what's a few more weeks. Appreciate all the comments.