config: add a simple dataclass based config system

[suvayu]

Features

• Config objects are also validated using Pydantic
• Learning schemes can be provided as a list (default: empty)
• Both TOML and YAML are supported, TOML is preferred

Breaking changes

• train test validation split is now configured together, easier to
  spot mistakes

  train_test_val_split:
    train: 0.7
    validate: 0.15
    test: 0.15

Questions / Remarks

• I'm not sure the input_features_file config should be a config.
  It seems it's part of the dataset, and should be addressed by
  incorporating the information in the data format.
• We can add custom validation like train_test_val_split adds up to 1.
• Is supporting both YAML and TOML worth it? Because of the breaking
  changes mentioned above, it's not 100% backwards compatible.
• Should learning_schemes be mandatory?
• Is time_series_fraction required only when learning_schemes are
  not present? Currently it's ignored.

TODOs

• Add tests
• Include TOML examples
• Generate documentation
  • maybe using JSON schema
• Setup for different models, latent, fno, etc
• composable config files, mostly to split data & model config

Sample session

(I edited both files to fix the breaking changes)

In [1]: from dataclasses import asdict

In [2]: from neuralpdr.config import read_conf

In [3]: conf = read_conf("configs/v1/ml4ps_paper/mlps_model_1.yaml")

In [4]: conf2 = read_conf("configs/v2/test.yaml")

In [5]: asdict(conf)
Out[5]: 
{'start_index': 0,
 'end_index': 300,
 'batch_size': 128,
 'learning_rate': 0.003,
 'weight_scale': 0.2,
 'weight_decay': 0.00014,
 'weight_truncation': 10.0,
 'enc_dec_depth': 2,
 'enc_dec_width': 512,
 'latent_depth': 2,
 'latent_width': 512,
 'latent_bottleneck': 32,
 'latent_final_activation': 'tanh',
 'minimal_timeseries_length': 32,
 'train_test_val_split': {'train': 0.7, 'validate': 0.15, 'test': 0.15},
 'training_batch_subsampling': 1.0,
 'shuffle_every_n_epochs': 1,
 'aux_features': True,
 'save_file_path': PosixPath('mlps_1'),
 'dataset_path': PosixPath('/home/gijsv/projects/emulating-3dpdr/3dpdr_dataset_8192.h5'),
 'input_features_file': PosixPath('configs/v1/features/input_features.yaml'),
 'learning_schemes': [{'timeseries_fraction': 0.33,
   'epochs': 50,
   'lr_scheduler': 'sgdr',
   'warmup_epochs': 2,
   'learning_rate': 0.003},
  {'timeseries_fraction': 0.67,
   'epochs': 50,
   'lr_scheduler': 'sgdr',
   'warmup_epochs': 2,
   'learning_rate': 0.003},
  {'timeseries_fraction': 1.0,
   'epochs': 50,
   'lr_scheduler': 'sgdr',
   'warmup_epochs': 2,
   'learning_rate': 0.003}],
 'double_epochs_last_fraction': False}

In [6]: asdict(conf2)
Out[6]: 
{'start_index': 0,
 'end_index': 498,
 'batch_size': 32,
 'learning_rate': 0.001,
 'weight_scale': 0.001,
 'weight_decay': 0.0001,
 'weight_truncation': 4.0,
 'enc_dec_depth': 4,
 'enc_dec_width': 64,
 'latent_depth': 4,
 'latent_width': 64,
 'latent_bottleneck': 16,
 'latent_final_activation': 'tanh',
 'minimal_timeseries_length': 32,
 'train_test_val_split': {'train': 0.7, 'validate': 0.15, 'test': 0.15},
 'training_batch_subsampling': 1.0,
 'shuffle_every_n_epochs': 1,
 'aux_features': True,
 'save_file_path': PosixPath('results3'),
 'dataset_path': PosixPath('/home/vermarien/data2/emulating-3dpdr/3dpdr_dataset_v2.h5'),
 'input_features_file': PosixPath('configs/v2/features/input_features.yaml'),
 'learning_schemes': [],
 'double_epochs_last_fraction': True}

[GijsVermarien]

Thanks for the PR, I added a small comment on tomli, and have a design question:

Now I see what you mean with toml and yaml support both being enabled. It makes me wonder: if we have two ML models, says one with called latent with parameters width and depth, but now we also have another model called fourier neural operator that has width, depth and spectral_components. Could we facilitate this by creating namespaces where specifyinglatent has a latent pydantic config checker and fno has its own, or could we still get away with a flat configuration namespace still? This would be very useful if we want to spin up more types of models to experiment with. We could even derive this call signature directly from the model.py files.

To answer your questions:

I'm not sure the input_features_file config should be a config. It seems it's part of the dataset, and should be addressed by incorporating the information in the data format.

Yes and no; It's derived from the dataset since every dataset has different features available. But some might want to use only a subset of the features. I mainly put it it in it's own config since I didn't want to copy it between configurations as they would become very long. Would it be easy to create one canonical parser, but allow for pointing to use include type statements for other configs?

We can add custom validation like train_test_val_split adds up to 1.

Yes, I would add it as a warning. This allows user to easily select subsets for testing the architecture etc.

Is supporting both YAML and TOML worth it? Because of the breaking changes mentioned above, it's not 100% backwards compatible.

Since we're still in 0.2, I would prefer to just make it breaking and move to toml altogether.

Should learning_schemes be mandatory?

If the user only wants to use a constant learning rate and fixed amount of epochs, they could specify the options inside the learning_schemes structure directly in the root namespace. But if this introduces a lot of complexity, I think it's fine to enforce the namespace of learning_schemes with at least epochs and learning_rate specified. We can then assume defaults for timeseries_fraction:float=1.0 andwarmup_epochs:int=0 (only available for sgdr and other periodic schedulers).

Is time_series_fraction required only when learning_schemes are not present? Currently it's ignored.

It really only makes sense in the learning_schemes context.

[suvayu]

one with called latent with parameters width and depth, but now we also have another model called fourier neural operator that has width, depth and spectral_components.

This kind of conditional validation can get complicated. Pydantic has a very specific order of validation. While it's possible to do, I've found it's harder to maintain. What I would prefer is, there is a different signal, say a cli argument, or maybe an attribute in the config that lets us choose the right model. I think discriminated union could work here. Here's how it would look like:

https://github.com/spine-tools/Spine-Database-API/blob/5cd944dc8f114f23981d111aee083048912ca3a1/spinedb_api/models.py#L336-L337

In the above case, the particular array type is determined by the "type" attribute. Are there other such conditional cases? It's best to accumulate them together into one condition, rather than a cascade. Does that make sense?

Yes and no; It's derived from the dataset since every dataset has different features available. But some might want to use only a subset of the features.

Okay, I get it. I'm not happy about it, but I also can't say I have a better solution for this :-p

I mainly put it it in it's own config since I didn't want to copy it between configurations as they would become very long. Would it be easy to create one canonical parser, but allow for pointing to use include type statements for other configs?

Including other config files can be messy IMO. When you allow that, several additional factors start playing a role, like working directory (for relative paths), precedence of values when there is an overlap, etc. What we could do is, treat this path as another type of config file with a separate "model" and validate that as a second step. We could also take into account other info, e.g. we now know the dataset path, so we could detect typos in the config if names don't match. But in that case this second "model" probably should be very simple.

But if this introduces a lot of complexity, I think it's fine to enforce the namespace of learning_schemes with at least epochs and learning_rate specified.

I'll take this route, quite clean.

Is time_series_fraction required only when learning_schemes are not present? Currently it's ignored.

It really only makes sense in the learning_schemes context.

Not sure I follow, do you mean this attribute should be present when learning_schemes are specified? If that's the case, shouldn't it be under the learning_schemes attribute?

[suvayu]

if we have two ML models, says one with called latent with parameters width and depth, but now we also have another model called fourier neural operator that has width, depth and spectral_components.

I guess the configs for latent are the ones latent_*, what are the spectral components for fno? For depth and width, would it be fine to have them as just "depth" and "width" since now we are separating "latent" and "fno"?

[GijsVermarien]

one with called latent with parameters width and depth, but now we also have another model called fourier neural operator that has width, depth and spectral_components.

This kind of conditional validation can get complicated. Pydantic has a very specific order of validation. While it's possible to do, I've found it's harder to maintain. What I would prefer is, there is a different signal, say a cli argument, or maybe an attribute in the config that lets us choose the right model. I think discriminated union could work here. Here's how it would look like:

https://github.com/spine-tools/Spine-Database-API/blob/5cd944dc8f114f23981d111aee083048912ca3a1/spinedb_api/models.py#L336-L337

In the above case, the particular array type is determined by the "type" attribute. Are there other such conditional cases? It's best to accumulate them together into one condition, rather than a cascade. Does that make sense?

I see what you mean! Do we need a better typing system then? I'm mostly thinking how we could allow for several models types, latent is already implemented and FNO might be something we want to implement in the future.

Yes and no; It's derived from the dataset since every dataset has different features available. But some might want to use only a subset of the features.

Okay, I get it. I'm not happy about it, but I also can't say I have a better solution for this :-p

I mainly put it it in it's own config since I didn't want to copy it between configurations as they would become very long. Would it be easy to create one canonical parser, but allow for pointing to use include type statements for other configs?

Including other config files can be messy IMO. When you allow that, several additional factors start playing a role, like working directory (for relative paths), precedence of values when there is an overlap, etc. What we could do is, treat this path as another type of config file with a separate "model" and validate that as a second step. We could also take into account other info, e.g. we now know the dataset path, so we could detect typos in the config if names don't match. But in that case this second "model" probably should be very simple.

I think this works, with precendce I think it's fine to throw errors just to be safe. It's mostly that things are often used in a composable fashion. One example would be that I want to run experiments on Snellius with one common dataset, actually it would make a lot of sense to specify the dataset just once in a common config and then models separately. I'm not sure if there is a good solution for such "composable" configurations? In the end it's fine if it's all mashed together at runtime and run through just one parser...

But if this introduces a lot of complexity, I think it's fine to enforce the namespace of learning_schemes with at least epochs and learning_rate specified.

I'll take this route, quite clean.

Is time_series_fraction required only when learning_schemes are not present? Currently it's ignored.

It really only makes sense in the learning_schemes context.

Not sure I follow, do you mean this attribute should be present when learning_schemes are specified? If that's the case, shouldn't it be under the learning_schemes attribute?

If there is only one "learning scheme", it would be fine to put learning rate and epochs in the root namespace. But once you have more than one, it has to be put in a learning_schemes list. The flat namespace for one learning scheme is just to make things easier when onboarding people. But we can enforce the learning_schemes namespace if its easier to verify that way.

[GijsVermarien]

Do you need more input from my side to move forward?

[suvayu]

I see what you mean! Do we need a better typing system then? I'm mostly thinking how we could allow for several models types, latent is already implemented and FNO might be something we want to implement in the future.

No need for a separate typing system. I already mocked up something using discriminated unions. Can you give a possible set of options that I can use a placeholder for a future FNO model? Based on your earlier comment, I have used something like this: depth: int, width: int, and spectrum: list[float]. Is this okay for now as a placeholder?

It's mostly that things are often used in a composable fashion. One example would be that I want to run experiments on Snellius with one common dataset [...]

Okay, that use case makes sense. I'll think how best to do this. Basically the issue is having one config file refer to another. If instead the files are provided independently, and are combined based on a clear precedence rule, e.g. order on the command line, then it is pretty straightforward.

... But we can enforce the learning_schemes namespace if its easier to verify that way.

It's difficult to validate if certain attributes are present/absent for different cases. I'll enforce it as a list, and for the simple case there will be only a single entry.

For now I think I have everything. The new conceptual work is small, I think the big part is making the existing code & tests compatible with the new config. I think I'll work on this again next week.

[GijsVermarien]

I see what you mean! Do we need a better typing system then? I'm mostly thinking how we could allow for several models types, latent is already implemented and FNO might be something we want to implement in the future.

No need for a separate typing system. I already mocked up something using discriminated unions. Can you give a possible set of options that I can use a placeholder for a future FNO model? Based on your earlier comment, I have used something like this: depth: int, width: int, and spectrum: list[float]. Is this okay for now as a placeholder?

This is perfect for now, I find that if you have only N=2 you are way further ahead than with N=1.

It's mostly that things are often used in a composable fashion. One example would be that I want to run experiments on Snellius with one common dataset [...]

Okay, that use case makes sense. I'll think how best to do this. Basically the issue is having one config file refer to another. If instead the files are provided independently, and are combined based on a clear precedence rule, e.g. order on the command line, then it is pretty straightforward.

Super

... But we can enforce the learning_schemes namespace if its easier to verify that way.

It's difficult to validate if certain attributes are present/absent for different cases. I'll enforce it as a list, and for the simple case there will be only a single entry.

Sounds good as well

For now I think I have everything. The new conceptual work is small, I think the big part is making the existing code & tests compatible with the new config. I think I'll work on this again next week.

Excellent!

[suvayu]

Hi @GijsVermarien, thanks for updating the dependencies. I was finding it difficult :-p

What do you think of the PR otherwise? Can you review it?

[GijsVermarien]

Hi Suvayu, great to hear it's ready for review. Could you quickly double check the original TODOs to clarify which you addressed and which you didn't yet. Cheers, Gijs

[suvayu]

Sure, this addresses:

• Makes the package setup modern using pyproject.toml
• Overhauls the config system
• Disabled Neptune related things (but didn't build a replacement)
• Fixed some function signatures guided by static analysis.

I have started another branch on top of this (linters) with more static analysis guided fixes. I'll submit that end of this week.

I'll tackle the data loader after this in a separate PR.

[GijsVermarien]

I went through the changes and it looks fine. Let's accept and see if anything pops up whilst I retrain the model on a new dataset I received recently.

[GijsVermarien]

Why do we need Python<3.11 for tomli at this time?