Enabling users to add auto-config support for custom types

[rsokl]

Summary

It would be nice for users to be able to register custom conversion strategies so that hydra-zen's various config-creation functions can know how to convert values of those types into corresponding structured config representations

Motivation

hydra-zen supports values of types that are not natively supported by Hydra.

E.g. we provide custom support for values of pathlib.Path and complex. This means that values of these types can be provided, natively, to all of hydra-zen's config-creation functions (just, builds, and make_config) and utilities (to_yaml, save_yaml) (as of #259 )

>>> from pathlib import Path
>>> from hydra_zen import just, builds, make_config, to_yaml
>>> Conf = make_config(value=2.0 + 3.0j)
>>> print(to_yaml(Conf))
 value:
  real: 2.0
  imag: 3.0
  _target_: builtins.complex

>>> just([Path.cwd(), 1+2j])
[ConfigPath(_args_=('C:\\Users\\Ryan Soklaski',), _target_='pathlib.Path'),
 ConfigComplex(real=1.0, imag=2.0, _target_='builtins.complex')]

Under the hood , we store a dictionary ZEN_VALUE_CONVERSION that maps type -> conversion_fn, where conversion_fn describes how to take a value of that type and represent that value as a structured config (such that instantiating the structured config returns that value).

Think of it as pickling, but in the form of structured configs. Then instantiate does the unpickling.

It would be nice if users could leverage their own conversion functions!

Implementation Details

There would be a two ways that we can expose this functionality: via a registration API and via an object-oriented API. These are not exclusive; indeed, we would probably enable all of these:

hydra_zen.register_converter (happy to change the name)

We would closely mirror Hypothesis' register_type_strategy implementation and enable users to manually register conversion functions:

from hydra_zen import register_converter, builds, just, to_yaml

class WeirdClass:
    def __init__(self, x: int):
        self.x = x + 2

register_converter(WeirdClass, lambda obj: builds(WeirdClass, x=obj.x - 2)

now hydra-zen's config creation functions will understand how to convert WeirdClass instances to corresponding configs

>>> Conf = make_config(a=1, b=WeirdClass(x=10))
>>> print(to_yaml(Conf))
a: 1
b:
  _target_: __main__.WeirdClass
  x: 10

These type conversion functions can be registered en masse and automatically using entrypoints. This is again following in Hypothesis' footsteps.

hydra-zen's various utility functions will automatically "understand" these values, so one will be able to write:

from hydra_zen import save_yaml

save_yaml(WeirdClass(x=100), "weird.yml")

and this will simply work!

__as_config__ (also happy to change this name)

As an alternative, library authors can also implement a __as_config__ method that hydra-zen will look for:

from typing import Self, Type

from hydra_zen import register_converter, builds, just, to_yaml
from hydra_zen import Builds


class WeirdClass:
    def __init__(self, x: int):
        self.x = x + 2
    
    def __as_config__(self: Self) -> Type[Builds[Type[Self]]]:
        # tells `hydra_zen.just` how to convert instances of 
        # `WeirdClass` to structured configs
        return builds(type(self), x=x - 2)

Now, make_config (and other config-creation functions) will see WeirdClass's __as_config__ method, and defer to that:

>>> Conf = make_config(a='hi', b=WeirdClass(x=22))
>>> print(to_yaml(Conf))
a: 'hi'
b:
  _target_: __main__.WeirdClass
  x: 2

Some additional considerations

Presently, hydra-zen will not check for subclass relationships. E.g. we can create configs of type complex, but we don't handle values of subclasses of complex. We may want register_converter to be able to specify whether or not the strategy is valid for subclasses. Fortunately __as_config__ can handle this elegantly.

Another thing to consider is that this approach to config-creation should be used with restraint. Even though we could register strategies for all torch types, we should not encourage users to instantiate an entire resnet, just to pass it to just to be made into a config... Well, maybe this is okay, but it seems awfully slow and "wasteful".

Ultimately, this functionality is meant to make it painless to take Python objects that you are working with and to create their structured config representations. This does not accommodate common configuration use cases, like specifying interpolated fields:

>>> WeirdClass(x="${x}")
TypeError: can't add 2 + "${x}"

Feedback

Does this make sense? Does it seem useful? Are there any pitfalls that I am missing?

[addisonklinke]

I like the idea of an __as_config__ interface per class because it keeps the conversion code local to the rest of the source. That should make its usage more apparent for newcomers to a given code base

I think I'm running into a similar issue where a custom class cannot be converted to Hydra-compatible primitives, but I'm not 100% sure this is what you're getting at. Specifically PyTorch Lightning Flash has InputTransform whose running_stage init parameter (from the super class) is typed as Optional[SomeCustomClass]

import flash
from hydra_zen import builds
from omegaconf import OmegaConf

InputTransformConf = builds(flash.core.data.io.input_transform.InputTransform, populate_full_signature=True)
print(OmegaConf.to_yaml(InputTransformConf))

This causes an error

hydra_zen.errors.HydraZenUnsupportedPrimitiveError: Building: InputTransform ..
The configured value <property object at 0x7f9f589a02c0>, for field `running_stage`, 
is not supported by Hydra -- serializing or instantiating this config would ultimately result in an error.
Consider using `hydra_zen.builds(<class 'property'>, ...)` to

It sounds like the desired/proposed solution would be to handle the typing for running_stage with either

• hydra_zen.register_converter
• InputTransform.__as_config__

The second option isn't really available in this case since I'm importing from a 3rd party library, so it is nice to have both approaches

[rsokl]

Thanks for sharing this @addisonklinke. It is very helpful for me to see how you are using hydra-zen.

I will need to look into the specific example that you posted (glancing at the source code you linked, I don't see how this error popped up. It looks like the default value for running_stage should be None), but this does roughly look like the case where hydra_zen.register_converter would be a remedy for this.

Now that I have more time to work on hydra-zen, I hope to push forward with hydra-zen's auto-config capabilities

[addisonklinke]

@rsokl Another library class where I ran into issues

from hydra_zen import builds
import torch

cfg = builds(torch.optim.SGD, populate_full_signature=True)

This raises the unsupported primitive error for lr=<required parameter>. Is there a recommended workaround until this feature is implemented?

[rsokl]

Their <required parameter> thing is so weird. They have it so that the can validate param groups.

Anyway, we can add auto-config support for this by replacing it with MISSING. In the meantime, I recommend:

from omegaconf import MISSING

cfg = builds(torch.optim.SGD, lr=MISSING, populate_full_signature=True)

[rsokl]

@addisonklinke 8a8179b adds support for <required parameter>. builds(SGD, populate_full_signature=True, zen_partial=True) now works automatically in hydra-zen 0.9.0

[rsokl]

@addisonklinke I know this issue is oh-so very old, but I have officially introduced a means for adding auto-config support to all config-creation functions!

https://mit-ll-responsible-ai.github.io/hydra-zen/changes.html#rc1-2023-10-23