metaflow-torchrun

What it is and what it does

metaflow-torchrun is a plugin that bridges Metaflow workflows and PyTorch's torchrun distributed training framework. It provides a @torchrun decorator that lets you run existing PyTorch distributed programs (like DDP training) as parallel steps in a Metaflow DAG without modifying your training code. The decorator automatically handles torchrun argument selection based on Metaflow compute resource requests (CPU, GPU, memory), network discovery, and subprocess orchestration.

You use it by stacking the @torchrun decorator on a Metaflow step, then calling current.torch.run() with your training entrypoint and arguments. This pattern is useful for multi-node training jobs submitted to AWS Batch or Kubernetes, where you want Metaflow to manage task orchestration and torchrun to manage the distributed torch processes within each task.

Use it for:

• Run multi-node PyTorch DDP training as a step in a Metaflow workflow without subprocess boilerplate
• Orchestrate distributed GPT or transformer training across parallel Metaflow tasks on Kubernetes or AWS Batch
• Execute existing torchrun scripts inside Metaflow without rewriting training code for the workflow framework
• Automate network discovery and torchrun argument configuration based on Metaflow resource decorators

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

Provides a Metaflow decorator that integrates PyTorch distributed training via torchrun, allowing you to run multi-node torch programs as parallel Metaflow tasks without modifying your training code.

Yes, if you use Metaflow and PyTorch distributed training together. The package has no runtime dependencies, low install friction, active maintenance, and solves a real integration gap. License treatment is unclear—verify Apache License compliance before production use. No known vulnerabilities.

Install

metaflow-torchrun on PyPI

pip install metaflow-torchrun

uv add metaflow-torchrun

poetry add metaflow-torchrun

Installing metaflow-torchrun

Before you install

Installation is straightforward with no runtime dependencies. The package is actively maintained with a recent release, though Python version support is unspecified.

License in practice

License treatment is unclear; the description mentions Apache License but SPDX metadata is not present in the package record, so verify the actual license terms before use.

Quickstart

pip install metaflow-torchrun

from metaflow import FlowSpec, step, torchrun

class MyFlow(FlowSpec):
    @kubernetes(gpu=1)
    @torchrun
    @step
    def train(self):
        current.torch.run(
            entrypoint="main.py",
            entrypoint_args={"main-arg-1": "123"},
            nproc_per_node=1
        )

Requires Metaflow to be installed and configured; GPU/compute environment must support the resource decorators used (e.g., @kubernetes or @batch).

Verify before relying

• Exact Python version requirements and compatibility range
• Whether Metaflow itself is a required dependency or assumed to be pre-installed
• GPU/CUDA version requirements for torchrun integration

Package facts

Evidence: metaflow_torchrun-0.2.2-py3-none-any.whl

Tags

More Distributed Computing packages