🐛[BUG]: OmegaConf's ListConfig and DictConfig are not JSON serializable

[mnabian]

Version

0.3.0

On which installation method(s) does this occur?

No response

Describe the issue

With the recent checkpoint refactor, there is a constraint that the model arguments should be JSON serializable. The dominant approach for handling configs in Modulus is hydra configs, which relies on OmegaConf. Some of OmegaConf's data types are not serializable, including ListConfig and DictConfig, which limits the use of the new checkpointing feature with hydra configs. We need to custom JSON encoder to handle these types.

Minimum reproducible example

No response

Relevant log output

File "/usr/local/lib/python3.10/dist-packages/modulus/models/module.py", line 181, in save
    json.dump(self._args, f)
  File "/usr/lib/python3.10/json/__init__.py", line 179, in dump
    for chunk in iterable:
  File "/usr/lib/python3.10/json/encoder.py", line 431, in _iterencode
    yield from _iterencode_dict(o, _current_indent_level)
  File "/usr/lib/python3.10/json/encoder.py", line 405, in _iterencode_dict
    yield from chunks
  File "/usr/lib/python3.10/json/encoder.py", line 405, in _iterencode_dict
    yield from chunks
  File "/usr/lib/python3.10/json/encoder.py", line 438, in _iterencode
    o = _default(o)
  File "/usr/lib/python3.10/json/encoder.py", line 179, in default
    raise TypeError(f'Object of type {o.__class__.__name__} '
TypeError: Object of type ListConfig is not JSON serializable

Environment details

No response

[NickGeneva]

Hi @mnabian

Can you elaborate on what your running in this case? I'm imaging that the JSON saved in the model checkpoints is written by the checkpoint util, then read by the same util. Thus so long as the checkpoint functions can parse the JSON file in isolation, its fine.

How does this JSON get processed by hydra?

[mnabian]

@NickGeneva this error occurs when saving the checkpoint and input arguments of a model instantiated using Hydra configs for the input arguments. In this case, the model arguments might contain some of OmegaConf's data types that aren't JSON-serializable. As a result, it gives an error when attempting to save the arguments in a JSON file

[akshaysubr]

@mnabian Which model is accepting direct hydra configs into its constructor? And is that a strict constraint?

I'm assuming it is something like this?

@hydra.main(...)
def main(cfg: DictConfig):
    model = MyModel(cfg.model_params)

[mnabian]

@akshaysubr currently only the models in the unified recipe for weather prediction accept direct hydra configs. But I would also like to change other examples and have them all using hydra configs as a unified approach for handling configs in Modulus-Launch: NVIDIA/modulus-launch#91
No this is not a strict constraint as long as we have a proper way of instantiating models with args from hydra configs.

[mnabian]

This will be fixed at the call site. See the discussions here: #134