Compositional Configuration in Python with jsonargparse
On how to mange complex command line configurations.
Hello everybody.
I want to show a cool design pattern our machine learning team has been using recently. It’s based on the jsonargparse library my collegue Mauricio developed. Jsonargparse, as the name implies, is a library that gives your argparse jsonic superpowers:
- Read parameters from the environment, CLI and json files
- Dump your config directly into a json file
- Compose hierarchical parsers
- and much much more …
But today I want to focus on the topic of how to construct and pass complex configurations. The ideas and problems I’m using come mainly from machine learning - as this is what I’m working on - but I’ll try to make it more general. If you are not familiar with machine learning, you can just skip some elaborations like the following one:
The prime example where we need a huge configuration would be the training of some neural network. Those have usually dozen of parameters at different logical levels.
I want to show you where and how the configuration problem occurs. I’ll also show how jsonargparse solves this by enabling the composition of configurations.
The problem of configuration
A stack of verbose functions
Let’s say we are building a system with a big stack of function calls and a lot of different parameters at different levels:
At some point this is going to break. Imagine we add a parameter to function_with_less_params(). Now we have to change all of our function calls between the initial setting of the parameter and this function, just to carry that parameter to where it needs to be. Additionally, the definitions become harder and harder to read. Maintaining the system becomes tedious.
Passing configuration objects
Probably we are missing some abstractions and our parameters have some structure.
For example, in machine learning we might group batch size and epochs into a single hyper-parameter config.
We can refactor our existing approach, by passing a configuration object. In this example, I’m just going to use a dictionary, but that could be any type of object.
We still pass the configuration around though, which again clutters our calls, especially if we have a lot of other values we want to pass. We could make the configuration global, but a better approach would be to create another object, which holds the configuration as a member variable. All functions become methods of this object and have thereby access to it. Then, only the functionality, which needs this configuration, has access to it.
Now it is super easy for us to add new parameters to the existing code.
The journey so far
We saw that, if we have a lot of them, passing around single parameters is not a good idea. Our code becomes eye cancer. Changing things in our stack of calls is like a game of Jenga.
Collecting them in a single config object we make available via self is a far better solution.
Exposing ourselves in public
At some point, we reach the stage of actually executing our code. Usually there are two options:
- Somebody imports this code and runs it there. The importer uses our documentation to run it accordingly.
- We run the code ourselves. In machine learning, during the experimentation phase, that usually means via the command line.
To run our program via the command line, we need to make our configuration available there. The standard tool would be to use Python’s inbuilt argparse to construct it:
Configurations of Configurations, Breaker of Sanity
So far so fine. But once we have multiple classes, similar to ConfigHolder, that we want to parameterize, we run into another issue, namely we need to manage the coherence between our CLI and our classes.
A real life example would be the encoder and decoder of a sequence-to-sequence neural network: They are logically separate and need different configurations.
We can use multiple parameter groups in argparse to avoid namespace collisions and represent the different levels of calls, but maintaining this is still going to be hard again. Every time we change the parameter of one of our classes, we need to remember to update it in the parser accordingly and without typos. Here jsonargparse comes into play. I’ll switch gears and provide an executable example here. It’s still a dummy one, the structure is just more realistic. Our example will have three parts:
- A producer, producing some number
- A transformer, transforming those numbers
- A manager, which chains the previous two together
The core idea
Every class will be responsible for its own configuration and will expose it via a method. Higher level classes build their configuration from the configuration of the lower level classes they use.
The first step is not strictly necessary, but it removes some boilerplate code. You could easily implement this pattern without it. We create a super class, from which all the other 3 will inherit from. This class will force its children to implement get_config_parser(). In the case, when we pass no configuration, the default will be used.
Our next class, Producer, is going to inherit from it and will implement the get_config_parser() method. The call function will implement the logic we want to execute.
As you can see, we are creating an ArgumentParser from jsonargparse, populate it and return it. We do exactly the same thing for the Transformer, which I will omit here. (It only multiplies each input with a factor). We then combine both of those, to implement our Manager class:
The Manager class builds its parser from the Producer and Transformer classes’ via the ActionParser class. We could also add here parameters specific only to the Manager. Manager instantiates Producer and Transformer in its init, by passing the relevant part of its config to the respective object’s init.
CLI examples
So let’s see some examples, of how we would use this:
Conclusion
I hopefully showed you, how you can build hierarchical, maintainable configurations via jsonargparse. Since we manage our configurations close to the actual usage inside the code base, we can easily change things and don’t forget to update our CLI. Jsonargparse has even more cool features I did not go into here, like dumping the config into a json and reading it again. Check it out.
Bonus
We can also directly use the config namespaces, when we implement our code. Remember, when we set cfg=None. Now we don’t even have to configure anything when we import Manager. If we want to configure it, we just get the parsed default and change the values we want:
