Skip to main content

OmegaConf Configuration Support

flyte-sdk provides first-class support for OmegaConf through a dedicated plugin. This integration allows developers to pass DictConfig and ListConfig objects between Flyte tasks while preserving hierarchical structures, type validation, and rich Python types like Enums and Paths.

Type Transformers and Registration​

The core of the OmegaConf support lies in two custom type transformers defined in plugins/omegaconf/src/flyteplugins/omegaconf/base_transformer.py:

  • DictConfigTransformer: Handles omegaconf.DictConfig objects.
  • ListConfigTransformer: Handles omegaconf.ListConfig objects.

These transformers are automatically registered with the Flyte TypeEngine via the register_omegaconf_transformers function in plugins/omegaconf/src/flyteplugins/omegaconf/__init__.py. This registration typically happens through entry points when the flyteplugins-omegaconf package is installed, making the types available for use in task signatures without manual setup.

Serialization and Data Handling​

flyte-sdk serializes OmegaConf objects into MessagePack binaries to ensure efficient transport and compatibility with Flyte's internal data model. The logic for this conversion is implemented in plugins/omegaconf/src/flyteplugins/omegaconf/codec.py.

Interpolation Resolution​

A critical behavior of the flyte-sdk implementation is that interpolations are resolved at serialization time. When a DictConfig is passed from one task to another, the receiving task sees the concrete values, not the interpolation strings (e.g., ${base_lr}). This is achieved in _serialize_child by accessing the parent node with the key, which triggers OmegaConf's internal resolution mechanism.

Rich Type Support​

The serialize_omegaconf and deserialize_omegaconf functions handle several types that are common in configuration but not natively supported by standard JSON/MessagePack:

  • Enums: Preserved by storing the fully qualified class name and the member name.
  • Paths: pathlib.Path and pathlib.PurePath objects are serialized as strings and reconstructed as Path objects.
  • Tuples: Converted to lists during serialization but reconstructed as tuples.
  • MISSING: Fields marked as MISSING in OmegaConf are preserved. Accessing them in a downstream task will still raise a MissingMandatoryValue error.

Structured Configs​

One of the most powerful features of the OmegaConf plugin is its support for Structured Configs (configs backed by Python dataclasses).

When you pass a DictConfig created via OmegaConf.structured(MyDataclass), flyte-sdk attempts to preserve the underlying schema. The _schema_name function extracts the fully qualified name of the dataclass, which is then stored in the serialized payload.

# plugins/omegaconf/examples/example_structured_config.py

@dataclass
class OptimizerConf:
    lr: float = 0.001

@dataclass
class TrainConf:
    optimizer: OptimizerConf = field(default_factory=OptimizerConf)

@env.task
async def train(cfg: DictConfig):
    # cfg is a TrainConf-backed DictConfig with type validation
    assert OmegaConf.get_type(cfg) == TrainConf
    print(cfg.optimizer.lr)

cfg = OmegaConf.structured(TrainConf())
# When passed to 'train', the schema is preserved

Schema Resolution Fallback​

During deserialization, the plugin attempts to re-import the dataclass using the stored qualified name.

  • If the dataclass is importable in the receiving task's environment, the config is reconstructed using OmegaConf.structured(), maintaining full type validation.
  • If the dataclass cannot be imported (e.g., the code is missing in the worker image), the plugin falls back to creating a plain DictConfig from the raw values.

Visualization and Reporting​

The plugin includes a reporting utility in plugins/omegaconf/src/flyteplugins/omegaconf/report.py that allows developers to visualize their configurations directly in the Flyte Console.

The log_yaml function renders a DictConfig or ListConfig as a human-readable YAML block and appends it to a Flyte report tab (defaulting to a tab named "OmegaConf").

from flyteplugins.omegaconf import log_yaml

@env.task(report=True)
async def train_model(cfg: DictConfig):
    # This renders the config in the Flyte UI
    await log_yaml.aio(cfg, title="Training Hyperparameters")
    ...

Constraints and Limitations​

  • String Keys: OmegaConf requires keys in a DictConfig to be strings. flyte-sdk does not support integer-keyed dictionaries within OmegaConf containers.
  • ListConfig Typing: While DictConfig supports structured schemas via dataclasses, ListConfig objects always round-trip as plain ListConfig containers; there is no support for "Structured ListConfigs" with element-level type enforcement.
  • Importability: For Structured Configs to maintain their type-validation benefits, the underlying dataclass must be available in the Python path of both the producing and consuming tasks.