Compose multiple _global_ configs

[alexanderswerdlow]

A key feature I've used with Hydra is being able to compose several global level configs. In Hydra, this works out of the box by creating a new subconfig (I often name mine modes) and then I can specify a primary global experiment, along with several "modes". These are very useful for common situations where I need to modify multiple config groups [e.g., change model parameters and trainer params at the same time] without having to define a new experiment [as this is not feasible with multiple of these modes].

When I try to do this in Hydra-zen, I get the following error:

In 'mode/demo_mode': ValidationError raised while composing config:
Merge error: Config is not a subclass of Config. value

With this minimal example:

from dataclasses import dataclass
from hydra_zen import (
    MISSING,
    builds,
    make_config,
    store,
    zen
)

@dataclass
class Trainer:
    batch_size: int = 4
    lr: float = 0.001
    num_workers: int = 4
    seed: int = 32

@dataclass
class BaseConfig:
    trainer: Trainer = MISSING
    
store(Trainer(), group="trainer", name='default')
Config = builds(
    BaseConfig,
    populate_full_signature=True,
    hydra_defaults=["_self_", dict(trainer='default')],
)

store(Config, name="config")

mode_store = store(group="experiment", package="_global_")
mode_store(
    make_config(
        hydra_defaults=["_self_"],
        trainer=dict(lr=0.01),
        bases=(Config,),
    ),
    name="demo_exp",
)

mode_store = store(group="mode", package="_global_")
mode_store(
    make_config(
        hydra_defaults=["_self_"],
        trainer=dict(num_workers=3, seed=2),
        bases=(Config,),
    ),
    name="demo_mode",
)

def main(trainer):
    print(trainer)

if __name__ == "__main__":
    store.add_to_hydra_store()
    zen(main).hydra_main(config_path=None, config_name="config", version_base="1.2")

My goal here is to independently compose these modes as I can do in hydra: python train.py +experiment=demo_exp modes='[mode1,mode2]'.

[jgbos]

Interesting, I've never tried to do that before. The error is on the Hydra side which makes it hard to get around. I'm trying to figure out a work around though, I feel like we can make this work.

btw, I'm using the following command:

$ python train.py +experiment=demo_exp +mode=demo_mode

[jgbos]

Can you provide the Hydra config, I'm forgetting how to do pure Hydra :/

[alexanderswerdlow]

Thanks for getting back! It's not a full standalone example but here's a very simplified overview of what I'll use in Hydra. Let me know if you'd like a fully working example and I can put one together!

The command you sent is exactly the one I tried, although I suppose an even more complete example would be if I had multiple modes I wanted to compose [in addition to the experiment].

config/config.yaml:

defaults:
  - _self_  # Override values within this file with values in selected files.
  - experiment: null
  - modes: null

config/experiment/demo_exp.yaml:

# @package _global_
seed: 0
shuffle: False

config/modes/demo_mode_1.yaml:

# @package _global_
shuffle: True

config/modes/demo_mode_2.yaml:

# @package _global_
shuffle: True
seed: 1

main.py:

@hydra.main(config_path="config", config_name="config")
def run(config):
    pass

Edit: I put together a complete example in the zip below:

To run it use: python run.py experiment=demo_exp 'modes=[fast_run,remote_server]'

hydra.zip

[rsokl]

Thanks for the great minimal example - super helpful.

What is going on here is that Hydra sees that you are using "structured configs" (i.e. dataclasses) and is trying to make things "safer" by enforcing that config-B can only be merged with config-A if B is a subclass of A. In my opinion, this "safety" check does way more harm than good (in fact, I have quite literally only seen it do harm).

Fortunately, we can disable hydra's checks here by depriving it of type information. We can have our config-store convert dataclass types to omegaconf dictionaries upon storing them; these omegaconf dictionaries will have the same fields and hierarchical structure as our original configs, but hydra will no longer fall over itself with subclass checks.

Cutting to the chase, add this to the top of your code and everything works 😄

from hydra_zen import store
from hydra_zen.wrapper import default_to_config
from dataclasses import is_dataclass
from omegaconf import OmegaConf

def destructure(x):
    x = default_to_config(x)  # apply the default auto-config logic of `store`
    if is_dataclass(x):
        # Recursively converts: 
        # dataclass -> omegaconf-dict (backed by dataclass types) 
        #           -> dict -> omegaconf dict (no types)
        return OmegaConf.create(OmegaConf.to_container(OmegaConf.create(x)))  # type: ignore
    return x

store = store(to_config=destructure)

You do lose type-checking support for all configs that you store in this way. I.e. Hydra will no longer be able to see that a field is annotated as int but then raise when you passed it a str at the CLI. That being said, I think this is a reasonable tradeoff in order to enable your workflow, which is indeed very powerful¹.

Footnotes

1. Also, Hydra's runtime type checking is severely limited and flawed, and there are much better ways to add runtime type checking to your hydra-zen app. ↩

[rsokl]

cc @jgbos I wish I had thought of this like 2 years ago 😭

[jgbos]

Did you just come up with that @rsokl ? That is awesome.

@alexanderswerdlow I think below matches closely with the code you provided:

from dataclasses import dataclass
from hydra_zen import MISSING, builds, make_config, store, zen
from hydra_zen import store
from hydra_zen.wrapper import default_to_config
from dataclasses import is_dataclass
from omegaconf import OmegaConf


def destructure(x):
    x = default_to_config(x)  # apply the default auto-config logic of `store`
    if is_dataclass(x):
        # Recursively converts:
        # dataclass -> omegaconf-dict (backed by dataclass types)
        #           -> dict -> omegaconf dict (no types)
        return OmegaConf.create(OmegaConf.to_container(OmegaConf.create(x)))  # type: ignore
    return x


store = store(to_config=destructure)

exp_store = store(group="experiment", package="_global_")
mode_store = store(group="modes", package="_global_")


@dataclass
class Config:
    batch_size: int = 4
    lr: float = 0.001
    num_workers: int = 4
    seed: int = 32
    shuffle: bool = True


store(Config, name="config")

exp_store(
    make_config(
        hydra_defaults=["_self_"],
        seed=0,
        shuffle=False,
        bases=(Config,),
    ),
    name="demo_exp",
)

mode_store(
    make_config(
        hydra_defaults=["_self_"],
        num_workers=3,
        bases=(Config,),
    ),
    name="demo_mode",
)

mode_store(
    make_config(
        hydra_defaults=["_self_"],
        seed=220,
        bases=(Config,),
    ),
    name="demo_mode2",
)


def main(trainer):
    print(trainer)


if __name__ == "__main__":
    store.add_to_hydra_store()
    zen(main).hydra_main(config_path=None, config_name="config", version_base="1.2")

Execute with python run.py +experiment=demo_exp +modes=[demo_mode,demo_mode2]

[alexanderswerdlow]

@jgbos @rsokl This is great, thank you so much!

For the demo I had to add the following so that I could access the . properties instead of just getting a dict [although for some reason my main project didn't require this]. I still haven't fully tested it, but the quick ones I did with my larger project seemed to work just fine!

def main(trainer):
    cfg: Trainer = OmegaConf.create(trainer)
    print(cfg.lr)

I'll also note that my use-case probably isn't ideal—with better code design, I shouldn't need to globally modify multiple config groups—but I've found it very useful for the type of research code I write where things are constantly changing.

[jgbos]

@alexanderswerdlow I'll try to find some time to play with the "." method issue. Ideally you don't have to import and use omegaconf for your main functions.

@rsokl should we move this to discussions as this is a very interesting use case?

[rsokl]

@alexanderswerdlow - just to close the loop on this specific issue - I figured out a fix so that you don't need that additional OmegaConf call. Just change

store(Trainer(), group="trainer", name="default")

# The config stored here does not have a `_target_` field, so
# `zen` is unable to instantiate it back into an instance of `Trainer`,
# so it just gets converted to a `dict` by `zen`, which is why
# you had to do an additional `OmegaConf.create`

to

store(Trainer, group="trainer", name="default")

# Here, `store` automatically applies `builds(Trainer, populate_full_signature=True)` and thus 
# the config has `Trainer` as its target. So `trainer`
# will be a genuine instance of `Trainer` - this is preferable
# because now you can use things like methods on `Trainer`!

(This might be an indicator that store's auto-config abilities can a little bit too magical/subtle, but 🤷‍♂️)

Now the full working example looks like

from dataclasses import dataclass
from hydra_zen import MISSING, builds, make_config, store, zen

from hydra_zen import store, instantiate
from hydra_zen.wrapper import default_to_config
from dataclasses import is_dataclass
from omegaconf import OmegaConf



def remove_types(x):
    x = default_to_config(x)
    if is_dataclass(x):
        # recursively converts: 
        # dataclass -> omegaconf-dict (backed by dataclass types) 
        #           -> dict -> omegaconf dict (no types)
        return OmegaConf.create(OmegaConf.to_container(OmegaConf.create(x)))  # type: ignore
    return x


store = store(to_config=remove_types)


@dataclass
class Trainer:
    batch_size: int = 4
    lr: float = 0.001
    num_workers: int = 4
    seed: int = 32


@dataclass
class BaseConfig:
    trainer: Trainer = MISSING


store(Trainer, group="trainer", name="default")

Config = builds(
    BaseConfig,
    populate_full_signature=True,
    hydra_defaults=["_self_", dict(trainer="default")],
)

store(Config, name="config")

mode_store = store(group="experiment", package="_global_")
mode_store(
    make_config(
        hydra_defaults=["_self_"],
        trainer=dict(lr=0.01),
        bases=(Config,),
    ),
    name="demo_exp",
)

mode_store = store(group="mode", package="_global_")

mode_store(
    make_config(
        hydra_defaults=["_self_"],
        trainer=dict(num_workers=3, seed=2),
        bases=(Config,),
    ),
    name="demo_mode",
)


def main(trainer: Trainer):
    assert isinstance(trainer, Trainer)
    print(trainer)


if __name__ == "__main__":
    store.add_to_hydra_store()
    zen(main).hydra_main(config_path=None, config_name="config", version_base="1.2")

$ python ppp.py +experiment=demo_exp
Trainer(batch_size=4, lr=0.01, num_workers=4, seed=32)

@jgbos we should definitely document this - either informally in a Discussion, or more formally in a How To. I don't know exactly what parts to distill (e.g. the _global_ aspect is somewhat documented in Configure Multiple Experiments already), but stripping out the type info from a config is probably the most crucial part.

I actually think we might want to make remove_types part of the ZenStore API. E.g. something like

store = ZenStore(remove_types=True)
store(SomeDataClass, name="foo")  # type info is *not* stripped upon storing
store.add_to_hydra_store()  # types are stripped from configs only as they are added to global store

This way you can retrieve configs from your store that still are proper dataclasses with type info, but Hydra is still deprived of the type info such that it stops with its aggressive subclass checks. Thoughts?

[jgbos]

Wouldn't that name be a bit confusing if it doesn't remove the types for hydra-zen but only for hydra? Maybe the option should only be on add_to_hydra_store?

[rsokl]

Yeah it might be a little confusing, but for the most part I do want people to have access to type-info. It is only hydra that is misbehaving, so it should only be hydra who is deprived 😝 .

Maybe the option should only be on add_to_hydra_store?

Yeah, this seems like a good idea. Although one nice thing about exposing the option in ZenStore.__init__ as well is that you can set it in your specific store and not have to modify your code where add_to_hydra_store gets called. E.g. you might have 5 apps that use the same store - it would be nice to not have to modify all those call sites.

🎄