Implemention discussion: using Stan config JSON throughout library and other changes

[amas0]

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.

[WardBrian]

Thank you for these thoughts! I was meaning to write out a similar post about pre-2.0 work.

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.

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 be

• validate arguments
• subprocess out to cmdstan
  • at this point, we can basically forget the arguments the user supplied. we could even go crazy and give a non-blocking option, etc.
• When the subprocess returns, call CmdStanMCMC#from_files on the runset's files.

The 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=1 and adapt engaged=true both meaning the same thing

... pydantic thoughts ...

Re: 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).

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.

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 loc and scale, for 'consistency', even when the distribution is not a location-scale family, and the mess that creates

That being said there are some 'obvious' improvements one could make in this direction, like whether things are accessed by properties or methods, etc.

[amas0]

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...

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 CmdStanMCMC object 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.

we could even go crazy and give a non-blocking option

I don't think this would be too much of a lift! Mostly I'd wonder what the user experience for this should be.

it does have the drawback of any upfront validation could stop something that was going to work due to an 'incidental' incompatibility

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.

Part of the difficulty with something like a StanFit ABC is that Stan itself is not terribly consistent with it's outputs

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]

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.

Precisely!

I don't think this would be too much of a lift! Mostly I'd wonder what the user experience for this should be.

I like how nutpie does this (obviously I don't think pause or resume could be implemented, but wait would be natural and abort could return the runset before passing to from_files)

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.

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)

Like do all methods now minimally output a CSV and (with save_cmdstan_config) a config file?

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.