OmegaConf and Hydra Configuration Management

Pass OmegaConf configurations through Flyte tasks

Use DictConfig and ListConfig directly in Flyte task signatures. After installing flyteplugins-omegaconf, importing flyteplugins.omegaconf registers both types with Flyte's TypeEngine:

from omegaconf import DictConfig, OmegaConf

from flyteplugins.omegaconf import register_omegaconf_transformers

register_omegaconf_transformers()

@env.task
async def identity(cfg: DictConfig) -> DictConfig:
    return cfg

cfg = OmegaConf.create({"optimizer": {"lr": 0.001}, "epochs": 10})
run = flyte.run(identity, cfg=cfg)
result = run.outputs()[0]

The explicit call is safe because register_omegaconf_transformers is cached with functools.lru_cache. The package also calls it at import time, and the flyte.plugins.types` entry point calls the same function during plugin discovery. Consequently, normal package import is sufficient when the plugin is installed and discovered.

DictConfigTransformer and ListConfigTransformer are the two concrete registrations. They both inherit OmegaConfTransformerBase; the former supplies the name OmegaConf DictConfig Transformer and container type DictConfig, while the latter supplies the name OmegaConf ListConfig Transformer and ListConfig.

Understand the wire representation

The transformer advertises Flyte SimpleType.STRUCT, but it does not emit an ordinary Flyte map or struct literal. OmegaConfTransformerBase.get_literal_type() adds cache metadata declaring MessagePack serialization. Its to_literal() path is:

1. serialize_omegaconf() recursively converts the configuration into the plugin payload.
2. msgpack.dumps() encodes that payload.
3. The bytes are placed in a Flyte Binary tagged with MESSAGEPACK.
4. The binary is stored in Literal(scalar=Scalar(binary=...)).

On the receiving side, to_python_value() accepts only a literal with scalar.binary. from_binary_idl() rejects a different binary tag, decodes MessagePack with strict_map_key=False, calls deserialize_omegaconf(), and checks that the result matches the transformer's container type. Empty, malformed, or non-binary literal shapes raise TypeError.

The shared codec recursively handles nested dictionaries and lists. It also represents missing values, enums, paths, and tuples with internal payload markers. The marker is data inside the MessagePack payload; it is not the Flyte literal type.

Use DictConfig in tasks

A plain configuration can be created, passed through a task boundary, and returned as an OmegaConf container. The following is the task pattern used by the OmegaConf examples (where env is the configured Flyte task environment):

from omegaconf import DictConfig, OmegaConf

@env.task
async def identity(cfg: DictConfig) -> DictConfig:
    return cfg

cfg = OmegaConf.create({"optimizer": {"lr": 0.001}, "epochs": 10})
run = flyte.run(identity, cfg=cfg)
result = run.outputs()[0]

assert isinstance(result, DictConfig)
assert result.optimizer.lr == 0.001
assert result.epochs == 10

The task examples in plugins/omegaconf/examples/example_dictconfig.py use the same signature to merge and consume configurations:

@env.task(report=True)
async def normalize_config(cfg: DictConfig) -> DictConfig:
    """Receives a plain DictConfig, adds a field, returns it."""
    await log_yaml.aio(cfg, title="Original config in normalize_config")
    return OmegaConf.merge(cfg, {"normalized": True})


@env.task(report=True)
async def merge_with_overrides(base_cfg: DictConfig, override_cfg: DictConfig) -> DictConfig:
    """Merges two DictConfigs inside a task and returns the result."""
    await log_yaml.aio(base_cfg, title="Base config in merge_with_overrides")
    await log_yaml.aio(override_cfg, title="Override config in merge_with_overrides")
    return OmegaConf.merge(base_cfg, override_cfg)

Structured configurations retain their dataclass type when the qualified dataclass is importable in the receiving task environment:

cfg = OmegaConf.structured(TrainConf(optimizer=OptimizerConf(lr=0.01), epochs=20))
run = flyte.run(_identity, cfg=cfg)
result = run.outputs()[0]
assert OmegaConf.get_type(result) == TrainConf

The codec records the qualified schema name. Deserialization imports that class, creates OmegaConf.structured(schema_type), and merges the decoded values into it. If the import fails, it deliberately falls back to a plain DictConfig, so the receiving image must contain the dataclass's importable module if schema preservation is required. Structured configs also remain strict: merging a key that the dataclass does not declare can fail. Use a plain DictConfig for an intentionally open-ended key set.

Account for interpolation and missing values

The codec resolves OmegaConf interpolations while recursively reading values for serialization. A receiving task therefore sees the resolved value rather than the original interpolation expression. The example task documents this behavior explicitly:

@env.task(report=True)
async def use_interpolated_config(cfg: DictConfig) -> float:
    """Receives a config where interpolations are already resolved at serialization time."""
    # cfg.optimizer.lr is the resolved float, not "${base_lr}"
    await log_yaml.aio(cfg, title="Config in use_interpolated_config")
    return float(cfg.optimizer.lr)

MISSING fields are encoded and survive deserialization. They do not become usable values: reading one still raises MissingMandatoryValue.

cfg = OmegaConf.structured(RequiredConf())  # data_path is MISSING
result = await roundtrip(cfg)

with pytest.raises(MissingMandatoryValue):
    _ = result.data_path

If a structured schema cannot be imported, fallback to a plain dictionary-backed configuration also loses the declared field type annotation. Fill required fields before the task boundary when the receiving task needs to read them.

Use ListConfig for sequences

ListConfigTransformer uses the same MessagePack representation but validates that deserialization produces a ListConfig. The repository's list example declares both list inputs and outputs directly:

@env.task
async def scale_values(values: ListConfig, factor: float) -> ListConfig:
    """Multiplies each numeric value in the list by factor."""
    return OmegaConf.create([v * factor for v in values])


@env.task
async def select_best_config(configs: ListConfig) -> DictConfig:
    """Receives a list of DictConfigs, returns the one with the highest lr."""
    best = max(OmegaConf.to_container(configs), key=lambda c: c["optimizer"]["lr"])
    return OmegaConf.create(best)

Nested list and dictionary nodes are preserved, but ListConfig has no structured typed-element reconstruction mode in this plugin. A list comes back as a plain ListConfig, even when elements originated from dataclass-backed values.

OmegaConf dictionary keys used by the documented integration must be strings; integer-keyed dictionaries are not supported. Also use bracket access when a user key collides with a DictConfig attribute. For example, a key named values must be read as cfg["values"], not cfg.values, because the latter resolves to the DictConfig.values() method:

assert cfg.kind == "training-job"
assert cfg["values"].lr == 0.001
assert cfg["values"].epochs == 10

Launch Hydra-composed Flyte workflows

Select the Flyte launcher with hydra/launcher=flyte. hydra_plugins.hydra_flyte_launcher.FlyteLauncher is a thin Hydra-discoverable subclass; its operational implementation is flyteplugins.hydra._launcher.FlyteLauncher. The wrapper exists because Hydra requires launcher targets under the hydra_plugins namespace.

FlyteLauncherConfig is registered in Hydra's ConfigStore as group hydra/launcher, name flyte:

@dataclass
class FlyteLauncherConfig:
    _target_: str = "hydra_plugins.hydra_flyte_launcher.FlyteLauncher"
    mode: str = "remote"
    wait: bool = True
    wait_max_workers: Optional[int] = 32


ConfigStore.instance().store(
    group="hydra/launcher",
    name="flyte",
    node=FlyteLauncherConfig,
)

Use a standard @hydra.main entry point

The real example in plugins/hydra/examples/train.py composes a DictConfig, initializes Flyte, and submits the task:

@hydra.main(version_base=None, config_path="conf", config_name="training")
def main(cfg: DictConfig):
    flyte.init_from_config()
    entry_task = apply_task_env(pipeline, cfg)
    mode = _flyte_launcher_mode()
    runner = flyte.with_runcontext(mode=mode) if mode else flyte
    run = runner.run(
        entry_task,
        cfg=cfg,
        dataset=cfg.data.dataset,
    )
    if mode is None:
        print(run.url)
    return run

Run a local job with python train.py hydra/launcher=flyte hydra.launcher.mode=local. For a remote six-job grid, use Hydra's multirun syntax:

python train.py --multirun \
    hydra/launcher=flyte hydra.launcher.mode=remote \
    optimizer.lr=0.001,0.01,0.1 training.epochs=10,20

The two comma-separated dimensions expand to six combinations. wait=false submits a remote sweep without waiting:

python train.py --multirun \
    hydra/launcher=flyte hydra.launcher.wait=false \
    optimizer.lr=0.001,0.01,0.1

Use the Python SDK helpers

Import the helpers from their implementation module when the application should be driven from Python rather than a @hydra.main script:

from flyteplugins.hydra._run import hydra_run, hydra_sweep

hydra_run(
    pipeline,
    config_path="conf",
    config_name="training",
    overrides=["optimizer.lr=0.01"],
    mode="remote",
)

runs = hydra_sweep(
    pipeline,
    config_path="conf",
    config_name="training",
    overrides=["optimizer.lr=0.001,0.01,0.1"],
    mode="remote",
)

hydra_run() composes one configuration, configures a launcher, launches one override set, and returns that result. hydra_sweep() expands ordinary comma-separated overrides into a Cartesian product and returns one result per job. Both accept config_path, config_name, overrides, mode, wait, wait_max_workers, run_options, task_env_key, and additional non-cfg task arguments.

For example, a repository example passes a dataset and two value overrides:

run = hydra_run(
    pipeline,
    config_path=CONFIG_PATH,
    config_name="training",
    overrides=["optimizer.lr=0.01", "training.epochs=20"],
    dataset="s3://my-bucket/imagenet",
    mode="remote",
)

Overrides beginning with hydra/, such as hydra/sweeper=optuna, take the full Hydra runtime path so the plugin can be discovered. Ordinary application overrides and Hydra namespace value overrides such as hydra.run.dir=... use the basic sweep path.

Use the Flyte CLI

flyte hydra run does not require @hydra.main; it receives the config location, Python file, task name, and override options directly:

flyte hydra run --config-path conf --config-name training \
    train.py pipeline --dataset s3://my-bucket/imagenet

flyte hydra run --local --config-path conf --config-name training \
    train.py pipeline --dataset s3://my-bucket/imagenet

CLI application overrides use --cfg, while Hydra namespace overrides use --hydra-override:

flyte hydra run --config-path conf --config-name training \
    train.py pipeline --dataset s3://my-bucket/imagenet \
    --cfg optimizer.lr=0.01 --cfg training.epochs=20 \
    --hydra-override "hydra.run.dir=./outputs/exp1"

A CLI grid uses repeated --cfg options:

flyte hydra run --multirun --config-path conf --config-name training \
    train.py pipeline --dataset s3://my-bucket/imagenet \
    --cfg "optimizer.lr=0.001,0.01,0.1" --cfg "training.epochs=10,20"

Tune dispatch and result handling

FlyteLauncher.setup() must run before launch(). Setup stores the Hydra context, task function, and composed configuration. launch() initializes Hydra globals and logging, creates hydra.sweep.dir, then dispatches according to the exact mode value:

• mode == "remote" sets Flyte's run mode to remote, submits every job before waiting, and prints each submitted run URL.
• Any other mode follows the local path; the documented configuration values are remote and local.
• Local jobs run sequentially in-process.
• With wait=True, remote jobs are waited on concurrently. Each wait uses quiet=True; wait_max_workers is capped at the run count and forced to at least one. None means one worker per run.
• With wait=False, remote jobs are returned after submission without waiting.

A successful waited remote result is wrapped in FlyteHydraRunResult. It exposes .url, delegates unknown attributes to the underlying Flyte run, converts to float, and iterates over scalar or sequence output. A single Flyte output becomes .value; multiple outputs remain a list. Exceptions during waiting and non-success terminal phases become failed Hydra job results.

Configuration and troubleshooting

Setting	Default	Behavior
hydra.launcher.mode	remote	Remote Flyte submission for exactly "remote"; other values use local dispatch.
hydra.launcher.wait	true	Wait for remote runs after submitting them. Set false for fire-and-forget submission.
hydra.launcher.wait_max_workers	32	Maximum concurrent wait workers. Set to null for one worker per run.
hydra.run.dir	Hydra default	Job output directory used for a single run.
hydra.sweep.dir	Hydra default	Sweep directory; the launcher creates it before dispatch.
hydra.sweep.subdir	Hydra default	Per-job sweep subdirectory passed to Hydra's run_job().

• The task image cannot deserialize a configuration. Install flyteplugins-omegaconf in the environment that executes the task. Import-time registration and the flyte.plugins.types entry point are the plugin's registration paths.
• A structured result is plain DictConfig. Ensure the dataclass's fully qualified module and class are importable in the receiving environment. The codec falls back to OmegaConf.create(values) when schema import fails.
• A required field still raises after a successful run. This is expected for MISSING; serialization preserves the marker, and accessing it raises MissingMandatoryValue until the field is filled.
• A direct transformer conversion fails. The literal must contain scalar.binary, and its binary tag must be Flyte's MESSAGEPACK tag. to_python_value() rejects other literal shapes and from_binary_idl() rejects other tags.
• launch() fails while reading Hydra configuration. Call setup() first; launch() expects config, hydra_context, and task_function to have been populated.
• A large remote sweep creates too many waiting threads. Keep the default cap or set a finite wait_max_workers; None deliberately creates one worker per run.
• An Optuna/custom sweeper is not found. Install the relevant plugin, such as hydra-optuna-sweeper, and use a hydra/ override such as hydra/sweeper=optuna so hydra_sweep() selects the full Hydra runtime.