Skip to main content

Experiment Tracking with MLflow and Weights & Biases

Flyte-sdk provides robust integrations with popular experiment tracking platforms, MLflow and Weights & Biases (W&B), allowing you to manage, track, and visualize your machine learning experiments directly within your Flyte workflows. These integrations enable automatic logging, contextual configuration, hyperparameter sweeps, and seamless UI link resolution.

MLflow Integration​

Flyte-sdk's MLflow integration streamlines experiment tracking by automatically managing MLflow runs and providing contextual configuration within your Flyte tasks.

Basic Run Tracking​

To track an MLflow run for a Flyte task, apply the @mlflow_run decorator to your task definition. This decorator ensures that an MLflow run is automatically started and managed for the duration of your task's execution. You can then access the current MLflow run object using get_mlflow_run() within your task to log metrics, parameters, and artifacts.

For example, a simple task that performs some computation and logs to MLflow would look like this:

from flyteplugins.mlflow import mlflow_run, get_mlflow_run
from flytekit import task

@mlflow_run
@task
def my_mlflow_task(x: int) -> int:
    run = get_mlflow_run()
    print(f"MLflow Run ID: {run.info.run_id}")
    # Log parameters, metrics, etc. using the 'run' object
    return x * 2

Automatic Logging (Autolog)​

Flyte-sdk extends MLflow's autologging capabilities, allowing you to automatically log metrics, parameters, and models from popular ML frameworks. Enable autologging by setting autolog=True in the @mlflow_run decorator. You can also specify a framework to enable framework-specific autologging.

For instance, to enable autologging for scikit-learn models:

from flyteplugins.mlflow import mlflow_run, Mlflow
from flytekit import task

@mlflow_run(autolog=True, framework="sklearn")
@task(links=(Mlflow(),))  # Automatically generates MLflow UI link
async def train_sklearn_model(n_samples: int = 100) -> None:
    # Your scikit-learn training code here
    pass

For more granular control or to apply autologging across multiple tasks within a workflow, use mlflow_config as a context manager. This allows you to set autologging parameters that propagate to child tasks.

from flyteplugins.mlflow import mlflow_config, Mlflow
from flytekit import task, workflow

@mlflow_run(autolog=True, framework="sklearn")
@task
def train_with_context_autolog(n_samples: int) -> None:
    # This task will use the autologging settings from the context manager
    pass

@workflow
def my_autolog_workflow(n_samples: int = 100):
    with mlflow_config(
        autolog=True,
        run_mode="new",
        # experiment_id=run.info.experiment_id, # Example of reusing an experiment
        autolog_kwargs={"log_input_examples": True},
    ):
        train_with_context_autolog(n_samples=n_samples)

Configuration​

The mlflow_config context manager (from flyteplugins.mlflow._context._MLflowConfig) provides a way to configure MLflow settings such as tracking_uri, experiment_name, run_name, tags, and run_mode. This configuration can be applied globally or within specific parts of your workflow, affecting all nested MLflow runs.

Key parameters for mlflow_config include:

  • tracking_uri: The URI of the MLflow tracking server.
  • experiment_name: The name of the MLflow experiment.
  • run_name: A custom name for the MLflow run.
  • tags: A dictionary of tags to associate with the run.
  • run_mode: Controls how MLflow runs are managed in nested contexts:
    • "auto" (default): Inherits the parent run if available, otherwise creates a new one.
    • "new": Always creates a new, independent MLflow run.
    • "nested": Creates a new run nested under the parent run.

To automatically generate links to the MLflow UI in the Flyte console, include an Mlflow() object in the links argument of your @task decorator. The Mlflow class (defined in flyteplugins.mlflow._link.Mlflow) resolves the link URL based on an explicit link provided or an auto-generated link from link_host and link_template set via mlflow_config().

from flyteplugins.mlflow import mlflow_run, Mlflow
from flytekit import task

@mlflow_run
@task(links=(Mlflow(),)) # Generates a link to the MLflow UI for this task's run
def my_linked_mlflow_task() -> None:
    # Task logic
    pass

# You can also provide an explicit link or override it:
@mlflow_run
@task(links=(Mlflow(link="https://my-custom-mlflow.example.com/run/123"),))
def my_custom_linked_mlflow_task() -> None:
    # Task logic
    pass

When mlflow_config is used with link_host, the Mlflow link will automatically construct the URL. If run_mode="new" is set for a task, the inherited parent link will not be shown, as the task creates its own independent run.

Distributed Training​

In distributed training scenarios, the @mlflow_run decorator is designed to log only from rank 0 to prevent duplicate entries and conflicts in MLflow. It automatically detects the RANK environment variable to determine which process should perform the logging.

Weights & Biases Integration​

Flyte-sdk's Weights & Biases (W&B) integration provides comprehensive tools for experiment tracking, hyperparameter sweeps, and managing distributed training runs.

Basic Run Initialization​

Initialize a W&B run for your Flyte task using the @wandb_init decorator. This decorator handles the wandb.init() call, ensuring proper run lifecycle management. You can access the active W&B run object within your task using get_wandb_run().

from flyteplugins.wandb import wandb_init, get_wandb_run
from flytekit import task

@wandb_init
@task
def my_wandb_task(x: int) -> str:
    run = get_wandb_run()
    if run:
        run.log({"input_x": x})
        print(f"W&B Run ID: {run.id}")
    return f"Processed {x}"

You can control the run creation behavior using the run_mode parameter in @wandb_init:

  • "auto" (default): Uses the parent's run if available, otherwise creates a new one.
  • "new": Always creates a new W&B run with a unique ID.
  • "shared": Always shares the parent's run ID.
from flyteplugins.wandb import wandb_init, get_wandb_run
from flytekit import task

@wandb_init(run_mode="new")
@task
async def child_task_new_run(x: int) -> str:
    run = get_wandb_run()
    if run:
        print(f"New W&B Run ID: {run.id}")
    return f"New run for {x}"

Configuration​

The wandb_config context manager (from flyteplugins.wandb._context._WandBConfig) allows you to configure various W&B settings, including project, entity, name, group, tags, and config dictionaries. It also supports Flyte-specific parameters like run_mode and rank_scope for distributed training.

Use wandb_config as a context manager or pass its output to task overrides to apply configuration:

from flyteplugins.wandb import wandb_config, wandb_init
from flytekit import task, workflow

@wandb_init
@task
def configured_wandb_task(value: float):
    run = get_wandb_run()
    if run:
        run.log({"metric": value})

@workflow
def my_configured_workflow():
    with wandb_config(project="my-flyte-project", entity="my-team", tags=["experiment-A"]):
        configured_wandb_task(value=0.5)

Sweeps (Hyperparameter Optimization)​

Flyte-sdk integrates with W&B Sweeps for hyperparameter optimization. Apply the @wandb_sweep decorator to a task that will act as a sweep agent. This task will pull configurations from the W&B Sweep controller and execute runs.

Define your sweep configuration using wandb_sweep_config (from flyteplugins.wandb._context._WandBSweepConfig), specifying the method, metric, and parameters. You can then pass this configuration to your workflow.

from flyteplugins.wandb import wandb_sweep, wandb_sweep_config, wandb_config
from flytekit import task, workflow

@wandb_sweep
@task
async def sweep_agent(agent_id: int, sweep_id: str, count: int = 5) -> int:
    # This task will run as a W&B sweep agent
    # It will automatically initialize wandb runs based on sweep config
    print(f"Running W&B sweep agent {agent_id} for sweep {sweep_id}")
    return agent_id

@workflow
def my_sweep_workflow():
    sweep_config = wandb_sweep_config(
        method="random",
        metric={"name": "loss", "goal": "minimize"},
        parameters={
            "learning_rate": {"min": 0.0001, "max": 0.1},
            "epochs": {"values": [5, 10, 15]},
        },
    )
    # Pass wandb_config and wandb_sweep_config to the workflow or task overrides
    # Example of passing via workflow context (simplified for illustration):
    # sweep_agent.with_overrides(wandb_config=wandb_config(...), wandb_sweep_config=sweep_config)
    # In a real scenario, these would be passed to the task that creates the sweep, and then inherited.
    pass

Within a @wandb_sweep task, you can retrieve the current sweep ID using get_wandb_sweep_id().

Similar to MLflow, you can generate links to the W&B UI for runs and sweeps. Use Wandb() for individual runs and WandbSweep() for sweep pages within the links argument of @task.

from flyteplugins.wandb import wandb_init, Wandb, wandb_sweep, WandbSweep
from flytekit import task

@wandb_init
@task(links=(Wandb(),))
def my_wandb_run_task() -> None:
    # Task that generates a link to its W&B run page
    pass

@wandb_sweep
@task(links=(WandbSweep(),))
def my_wandb_sweep_task() -> None:
    # Task that generates a link to its W&B sweep page
    pass

The Wandb class (from flyteplugins.wandb._link.Wandb) and WandbSweep class (from flyteplugins.wandb._link.WandbSweep) automatically construct the appropriate URLs based on the context, including host, project, entity, and run/sweep IDs.

Artifacts and Run Directories​

Flyte-sdk allows you to manage W&B artifacts and run directories. While not directly shown in the provided code snippets, the research findings indicate the existence of download_wandb_run_dir() and download_wandb_sweep_dirs() functions. Additionally, setting download_logs=True in @wandb_init or @wandb_sweep can automate the downloading of W&B logs after task completion.

Distributed Training​

For distributed training, W&B integration offers fine-grained control over how runs are logged across multiple processes or nodes. The run_mode and rank_scope parameters in wandb_config and @wandb_init govern this behavior:

  • run_mode:
    • "auto": In single-node distributed training, only rank 0 logs. In multi-node, behavior depends on rank_scope.
    • "shared": All ranks log to a single shared W&B run.
    • "new": Each rank gets its own W&B run, which can be grouped in the W&B UI.
  • rank_scope (Flyte-specific, for run_mode="auto" or "shared" in multi-node):
    • "global" (default): Only global rank 0 logs (or all ranks log to one shared run globally).
    • "worker": Local rank 0 of each worker logs (or ranks per worker log to one shared run per worker).

These parameters are crucial for correctly aggregating or separating W&B runs in complex distributed setups.

Common Patterns and Gotchas​

When integrating MLflow and Weights & Biases with Flyte-sdk, keep the following in mind:

  • Outermost Decorator: The @mlflow_run, @wandb_init, and @wandb_sweep decorators must always be the outermost decorators applied to a Flyte task. This ensures they correctly manage the lifecycle of the respective experiment tracking runs.
  • Traces and Context: Traces (sub-workflows or nested function calls within a task) cannot directly use these decorators. However, they can access the parent task's active MLflow or W&B run context using get_mlflow_run() or get_wandb_run() respectively.
  • Run Mode Management: Pay close attention to the run_mode parameter for both MLflow (_MLflowConfig.run_mode) and W&B (_WandBConfig.run_mode). This parameter dictates whether a new run is created, an existing parent run is inherited, or runs are shared, which is critical for organizing your experiments.
  • W&B Sweep Configuration: For W&B Sweeps, ensure that project and entity are correctly set via wandb_config() or directly in the WandbSweep link definition. Without these, W&B may not be able to query or display your sweeps correctly.