Flexible runtime validation for configs with variable nodes

[abefrandsen]

Greetings! I'm a big fan of this library and I'm trying to learn the best ways to perform config validation using structured configs and OmegaConf.merge. (I'll try to articulate my question coherently, but I'm at such a basic level where even that is difficult!)

My question is about runtime type validation for a config that can have variable nodes. Here's a toy example of what I'm getting at. In the context of an ML project, I specify a model from a config as follows, where the encoder and prediction_head nodes can be variable but all share a high-level structure:

encoder:
  module_name: MyTransformerEncoder # this is the name of a custom pytorch module
  module_cfg: # specify the init params to instantiate the model, must be consistent with the above-specified module
    d_model: 768
    num_layers: 4
prediction_head:
  module_name: MyFFNPredictor # another custom pytorch module
  module_cfg: # same as above, specify the init params
    num_target_classes: 50
    hidden_dim: 1024

But since I want to support different types of encoders and prediction heads, another equally valid config might be (don't worry about the fact that a transformer encoder is typically not easily swappable with a feedforward network...):

encoder:
  module_name: MyFFNEncoder # use a different encoder this time
  module_cfg: 
    embed_dim: 768
    n_layers: 6
prediction_head: # same as in first example
  module_name: MyFFNPredictor 
  module_cfg: 
    num_target_classes: 50
    hidden_dim: 1024

I want to use a structured config to validate such configs at runtime. A first attempt at this could be:

from dataclasses import dataclass
from omegaconf import OmegaConf, MISSING
from typing import Dict

@dataclass
class Module:
    module_name: str
    module_cfg: Dict

@dataclass
class ModelConfig:
    encoder: Module = MISSING
    prediction_head: Module = MISSING

The two example configs should validate just fine against this ModelConfig schema, but I'd like to further validate that the fields and types given in the module_cfg nodes are consistent with the corresponding value of module_name, and this is where I'm getting stuck. I don't want to write a separate structure config for every possible combination of specific encoders and decoders, but I want to be able to validate all such possible configs. While I found a (probably hackish) halfway solution, I'm very curious about how you all might approach this. Maybe I should just expect the user to provide correct values for these nodes (if they don't the code will almost surely error out as it tries to instantiate the specified module with the given init parameters) and be satisfied with my first attempt above?

Thanks all!