Metaflow Torchrun Integration

0.2.1 · active · verified Fri Apr 17

Metaflow-torchrun is a Python library that provides a `@torchrun` decorator to enable distributed PyTorch training within Metaflow steps. It abstracts away the complexities of launching and managing `torchrun` processes, allowing users to integrate distributed ML workflows seamlessly into their Metaflow flows. The current version is 0.2.1, with relatively frequent updates since its initial release.

Common errors

```
ModuleNotFoundError: No module named 'torchrun'
```
cause `torch` (which bundles `torchrun`) is not installed or not accessible in the Python environment where the Metaflow step runs.

fix Install the PyTorch library: `pip install torch` (or the appropriate command for your system/CUDA version, e.g., `pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118`).
```
TypeError: torchrun() got an unexpected keyword argument 'my_arg'
```
cause An unknown or misspelled argument was passed directly to the `@torchrun` decorator, or arguments intended for the PyTorch script were passed incorrectly.

fix Ensure only valid decorator arguments like `nproc_per_node`, `args`, `module`, etc., are used. For arguments to your PyTorch script, use the `args` parameter: `@torchrun(args=['--my-script-arg', 'value'])`.
```
MetaflowException: The @torchrun decorator is only supported when running a flow from a Python file, not interactively.
```
cause Attempting to run a Metaflow flow with the `@torchrun` decorator from an interactive environment (e.g., Jupyter notebook, IPython console), which is not supported.

fix Save your Metaflow flow as a Python file (e.g., `my_flow.py`) and execute it from the command line: `python my_flow.py run`.

Warnings

gotcha The `@torchrun` decorator is incompatible with interactive notebooks (e.g., Jupyter), requiring Metaflow flows to be run from a standalone Python file.
Fix: Run Metaflow flows decorated with `@torchrun` as a standalone Python script from your terminal: `python your_flow.py run`.
gotcha Ensure `torch` is installed and `torchrun` is available in the Python environment where the Metaflow step executes.
Fix: Install `torch` (which includes `torchrun`) in your environment: `pip install torch` (or `pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118` for CUDA-enabled versions).
gotcha Arguments intended for your PyTorch script, when passed to the `args` parameter of the `@torchrun` decorator, must be a list of strings.
Fix: Pass arguments as `args=['--epochs', '10', '--lr', '0.01']` instead of a single string, dictionary, or other types. For example: `@torchrun(args=['--model-name', 'resnet50'])`.
gotcha The current implementation of `@torchrun` distributes processes across multiple GPUs/CPUs *on a single node*, not across multiple machines.
Fix: Understand that `metaflow-torchrun` simplifies single-node distributed training. For multi-node setups, more advanced Metaflow features or other orchestrators would be required.

Install

pip install metaflow-torchrun Install metaflow-torchrun
pip install metaflow torch Install core dependencies

Imports

torchrun
```
from metaflow_torchrun import torchrun
```

Quickstart

This example demonstrates a basic Metaflow flow using the `@torchrun` decorator. Save this code as a Python file (e.g., `my_flow.py`) and run it from your terminal using `python my_flow.py run`. It will launch 2 parallel processes on the local machine within the `train_distributed` step, each printing its rank information. Ensure `metaflow-torchrun`, `metaflow`, and `torch` are installed in your environment.

import os
import logging
from metaflow import FlowSpec, step, current
from metaflow_torchrun import torchrun

# Set up basic logging to see output from torchrun
logging.basicConfig(level=logging.INFO)

class MyDistributedFlow(FlowSpec):

    @step
    def start(self):
        print(f"Starting flow {current.flow_name}...")
        self.next(self.train_distributed)

    # Decorate a step with @torchrun to enable distributed execution
    @torchrun(nproc_per_node=2) # Use 2 processes for illustration on a single node
    @step
    def train_distributed(self):
        # Inside a torchrun decorated step, the script is executed by multiple processes.
        # Each process will have environment variables like LOCAL_RANK, RANK, WORLD_SIZE.
        local_rank = int(os.environ.get("LOCAL_RANK", -1))
        global_rank = int(os.environ.get("RANK", -1)) # RANK is the global rank
        world_size = int(os.environ.get("WORLD_SIZE", -1))

        print(f"Hello from process {global_rank}/{world_size} (local rank {local_rank})\n" +
              f" in Metaflow run {current.run_id} step {current.step_name}.")

        # In a real scenario, you'd initialize a process group and run DDP here.
        # import torch.distributed as dist
        # if local_rank != -1:
        #     dist.init_process_group("nccl" if torch.cuda.is_available() else "gloo", rank=global_rank, world_size=world_size)
        # Your PyTorch model training code...
        # if local_rank != -1:
        #     dist.destroy_process_group()

        self.next(self.end)

    @step
    def end(self):
        print("Distributed training flow completed!")

if __name__ == "__main__":
    # To run this flow:
    # 1. Save it as a Python file, e.g., `my_flow.py`
    # 2. Run from your terminal: `python my_flow.py run`
    # Ensure `metaflow-torchrun`, `metaflow`, and `torch` are installed.
    MyDistributedFlow()

view raw JSON →