Hydra Launcher for Flyte

The flyte-sdk Hydra launcher allows you to run Hydra-configured applications on remote Flyte clusters. It supports single runs and multi-run sweeps (including grid search and Bayesian optimization via Optuna) by dispatching each Hydra job as a separate Flyte execution.

Running with the Hydra CLI

If your script uses the standard @hydra.main decorator, you can use the Flyte launcher by selecting it in your command line or configuration.

# Run a single job on Flyte
python train.py hydra/launcher=flyte hydra.launcher.mode=remote

# Run a grid sweep on Flyte
python train.py --multirun \
    hydra/launcher=flyte \
    optimizer.lr=0.001,0.01,0.1 \
    training.epochs=10,20

The launcher defaults to remote mode when used via the Flyte CLI or SDK, but you can explicitly set hydra.launcher.mode=local to run jobs sequentially in your local process.

Running with the Flyte CLI

The flyte hydra run command allows you to run Flyte tasks with Hydra configuration without requiring the @hydra.main decorator in your script. It inherits standard flyte run flags like --project, --domain, and --image.

# Single remote run with value overrides
flyte hydra run --config-path conf --config-name training \
    train.py pipeline \
    --cfg optimizer.lr=0.01 --cfg training.epochs=20

# Grid sweep (6 parallel remote executions)
flyte hydra run --multirun --config-path conf --config-name training \
    train.py pipeline \
    --cfg "optimizer.lr=0.001,0.01,0.1" --cfg "training.epochs=10,20"

Application-level overrides are passed using the name of the DictConfig parameter in your task (e.g., --cfg if your task signature is def pipeline(cfg: DictConfig):).

Programmatic Usage via SDK

The flyteplugins.hydra module provides hydra_run and hydra_sweep for launching jobs from Python. These functions wrap the FlyteLauncher and handle configuration composition.

Single Run

Use hydra_run to execute a single task with a composed configuration.

from flyteplugins.hydra import hydra_run
from my_script import pipeline

# Run remotely and wait for the result
result = hydra_run(
    pipeline,
    config_path="conf",
    config_name="training",
    overrides=["optimizer.lr=0.01"],
    dataset="s3://my-bucket/imagenet",
    mode="remote"
)

Sweeps and Hyperparameter Optimization

Use hydra_sweep to run multiple jobs. Comma-separated values in overrides trigger a grid search.

from flyteplugins.hydra import hydra_sweep

# Grid sweep: 2 models x 3 learning rates = 6 parallel Flyte runs
runs = hydra_sweep(
    pipeline,
    config_path="conf",
    config_name="training",
    overrides=["model=resnet,vit", "optimizer.lr=0.001,0.01,0.1"],
    mode="remote",
)

To use custom sweepers like Optuna, include the sweeper selection in the overrides list. This triggers the full Hydra runtime for plugin discovery.

# Optuna TPE/Bayesian sweep
runs = hydra_sweep(
    pipeline,
    config_path="conf",
    config_name="training",
    overrides=[
        "hydra/sweeper=optuna",
        "hydra.sweeper.n_trials=20",
        "hydra.sweeper.n_jobs=4",
        "optimizer.lr=interval(1e-4,1e-1)",
    ],
    mode="remote",
)

Mapping Hydra Configs to Flyte Resources

The task_env mechanism allows you to map Hydra configuration groups to Flyte task overrides such as resources (CPU/GPU) and container images. The FlyteLauncher automatically applies these overrides to the entry task.

Configuration Schema

In your Hydra config directory (e.g., conf/task_env/a100.yaml), define the resources:

# conf/task_env/a100.yaml
pipeline:
  resources:
    gpu: "1"
    mem: "32Gi"
    cpu: "8"
  image: "cr.flyte.org/my-project/training:latest"

Applying Overrides

You can select these presets using Hydra overrides:

flyte hydra run ... --cfg "+task_env=a100"

For child tasks called within your workflow, use apply_task_env from flyteplugins.hydra._run to manually apply the same configuration-driven overrides:

from flyteplugins.hydra import apply_task_env

@flyte.workflow
def my_workflow(cfg: DictConfig):
    # Apply overrides from cfg.task_env.train_model to the train_model task
    configured_task = apply_task_env(train_model, cfg)
    return configured_task(cfg=cfg)

Launcher Configuration

The FlyteLauncher can be configured via Hydra's hydra.launcher namespace.

Option	Default	Description
mode	remote	Execution mode: remote (submits to Flyte) or local (runs in-process).
wait	true	Whether to wait for remote Flyte runs to reach a terminal phase.
wait_max_workers	32	Maximum worker threads used to wait for remote runs concurrently.
task_env_key	task_env	The top-level config key used for task.override lookups.

Example override via CLI:

python train.py --multirun hydra/launcher=flyte hydra.launcher.wait=false