SkyPilot extension for Metaflow

This extension adds support for executing steps in Metaflow flows on any cloud provider via SkyPilot.

Installation

pip install metaflow-skypilot

SkyPilot also requires cloud credentials to be configured. Follow the SkyPilot setup guide for your cloud provider.

Basic Usage

Add the @skypilot decorator to any step you want to run on the cloud:

from metaflow import FlowSpec, step, skypilot

class MyFlow(FlowSpec):

    @step
    def start(self):
        self.next(self.train)

    @skypilot(cpus='2+', memory='8+')
    @step
    def train(self):
        # This step runs on cloud via SkyPilot
        print("Training on the cloud!")
        self.next(self.end)

    @step
    def end(self):
        print("Done!")

if __name__ == '__main__':
    MyFlow()

Resource Configuration

The @skypilot decorator accepts all sky.Resources parameters directly, an example with a few of them is given below:

@skypilot(
    infra='aws',
    cpus='4+',
    memory='16+',
    accelerators='A100:1',
)
@step
def gpu_step(self):
    ...

Reusing a Named Cluster

By default, each run provisions a fresh ephemeral cluster that is torn down after the job finishes. To reuse a persistent cluster across runs, pass a cluster_name:

@skypilot(
    cpus='2+',
    cluster_name='my-persistent-cluster',
)
@step
def my_step(self):
    ...

• Without cluster_name: a new cluster is provisioned, runs the job, and is terminated after 10 idle minutes.
• With cluster_name: the cluster is reused across runs (auto-started if stopped). It stops automatically after 10 idle minutes but is not terminated — it will be restarted on the next run.

Each task always runs in an isolated working directory (~/metaflow/assets/<job_name>/) regardless of cluster type, so there are no filesystem clashes when reusing a cluster.

Using with @pypi

Use @pypi to install Python dependencies on the remote VM:

@skypilot(cpus='2+')
@pypi(python='3.9', packages={'numpy': '1.24.0', 'pandas': '2.0.0'})
@step
def my_step(self):
    import numpy as np
    ...

Supplying Credentials

Cloud credentials for accessing the Metaflow datastore (e.g. S3) can be supplied in three ways:

• Instance IAM role / cloud identity: if the provisioned resource has access to the datastore via its cloud identity, no extra configuration is needed.
• Environment variables via the @environment decorator:

@environment(vars={
    "AWS_ACCESS_KEY_ID": "XXXX",
    "AWS_SECRET_ACCESS_KEY": "YYYY"
})
@skypilot(cpus='2+')
@step
def my_step(self):
    ...

• Secrets manager via the @secrets decorator.

If you are on the Outerbounds platform, authentication is handled automatically.

Things to Note

• @skypilot cannot be combined with @kubernetes, @batch, or @slurm on the same step.
• @parallel is not supported with @skypilot.
• The minimum step timeout is 60 seconds.

Fin.