Add flytekit-omegaconf plugin

[mg515]

Why are the changes needed?

Flytekit natively supports serialization of many data types for exchanging information between tasks.

The Flytekit OmegaConf Plugin extends these by the DictConfig type from the OmegaConf package as well as related types that are being used by the hydra package for configuration management.

Task example

from dataclasses import dataclass
from flytekit import task, workflow
from omegaconf import DictConfig
import flytekitplugins.omegaconf  # noqa F401


@dataclass
class MySimpleConf:
    _target_: str = "lightning_module.MyEncoderModule"
    learning_rate: float = 0.0001


@task
def my_task(cfg: DictConfig) -> None:
    print(f"Doing things with {cfg.learning_rate=}")


@workflow
def pipeline(cfg: DictConfig) -> None:
    my_task(cfg=cfg)


if __name__ == "__main__":
    from omegaconf import OmegaConf

    cfg = OmegaConf.structured(MySimpleConf)
    pipeline(cfg=cfg)

How was this patch tested?

Unit tests are added within the same PR.

Additionally, the plugin has been in private use for a while and is tested within multiple projects.

Check all the applicable boxes

• I updated the documentation accordingly.
• All new and existing tests passed.
• All commits are signed-off.

Acknowledgements

The first version of the plugin was developed roughly at the end of year 2021 by some of the current and former members of Merantix Momentum. Shoutout to @SebS94 @fg91 @AdemFr @YmirKhang @axkoenig @adrianloy @ThomasWollmann

[welcome]

Thank you for opening this pull request! 🙌

These tips will help get your PR across the finish line:

• Most of the repos have a PR template; if not, fill it out to the best of your knowledge.
• Sign off your commits (Reference: DCO Guide).

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 76.15%. Comparing base (fa2aa0b) to head (0a745d0).
Report is 227 commits behind head on master.

❗ There is a different number of reports uploaded between BASE (fa2aa0b) and HEAD (0a745d0). Click for more details.

HEAD has 148 uploads less than BASE

Flag	BASE (fa2aa0b)	HEAD (0a745d0)
	151	3

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2299      +/-   ##
==========================================
- Coverage   83.89%   76.15%   -7.75%     
==========================================
  Files         342      185     -157     
  Lines       25483    18883    -6600     
  Branches     3725     3715      -10     
==========================================
- Hits        21380    14380    -7000     
- Misses       3472     3868     +396     
- Partials      631      635       +4     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[SebS94]

Thanks for taking care of this Miha!

One small consideration - would it be more fitting to rename this plugin into flytekit-omegaconf as this is what it actually serialises? Obviously, our main consideration is the combination with hydra, but there is no actual direct dependency here.

[mg515]

@SebS94 was having the same thoughts myself as well, could as well do it yeah.

[eapolinario]

I did one pass. It's looking pretty good. I left a few comments, but no showstopper. I worry about conflicts with other plugins, like pydantic. We don't necessarily have to fix everything in this PR.

[fg91]

As discussed here, shell we remove the mode here and only provide a global mode?

[eapolinario]

bump.

[mg515]

Addressed in e6e1896

We provide global mode and a context manager for ability to control it locally still.
Please resolve the conversation if you think this is a suitable solution.

[kumare3]

@fg91 / @mg515 I would love for this to go in if you folks are into pushing it?

[cameronraysmith]

Note

The thread from this comment to here is migrated to flyteorg/flyte#5576

I think it's a good idea to support omegaconf's dictconfig directly because it is a foundational layer, but would just note that you get something like this, in terms of the more general ability to configure flyte workflows with hydra, almost for free if you use hydra-zen because it already provides a bridge from omegaconf's types to dataclasses, which are of course already supported by Flyte.

[fg91]

I think it's a good idea to support omegaconf's dictconfig directly because it is a foundational layer, but would just note that you get something like this almost for free if you use hydra-zen because it already provides a bridge from omegaconf's types to dataclasses.

Do you have a minimal working example that shows how a DictConfig can be passed to a Flyte task using hydra-zen? 🙏

[cameronraysmith]

@fg91 let's redirect this discussion to flyteorg/flyte#5576 to avoid getting into details that aren't directly relevant to this particular PR.

I probably should have done this instead of posting #2299 (comment) here, but I'll leave it for the record.

[eapolinario]

All the string manipulation to load modules in runtime seems a bit brittle to me, but won't block the PR on it. Besides that, we really should get on the same page on some of the outstanding @fg91 's comments before proceeding with this change.

[eapolinario]

bump.

[eapolinario]

Thank you for the persistence. This is looking great!

[eapolinario]

Actually before we get this in, can you add flytekit-omegaconf to the list of plugins in CI?

[welcome]

Congrats on merging your first pull request! 🎉

[cameronraysmith]

Many thanks @mg515!