metaflow-prebuilt

A Metaflow extension that pre-bakes conda/PyPI dependencies into Docker images at deploy time. Tasks start immediately — no conda bootstrap overhead at runtime.

How it works

When you deploy a flow with --environment=prebuilt, the extension:

1. Resolves the conda/PyPI environment spec for each step.
2. Builds a Docker image with the environment already installed.
3. Pushes the image to a registry.
4. Configures the remote runner (Batch, Kubernetes, etc.) to pull that image instead of bootstrapping conda at task start.

Install

pip install metaflow-prebuilt

Optional extras for specific backends:

pip install 'metaflow-prebuilt[ecr]'       # ECR registry + CodeBuild service
pip install 'metaflow-prebuilt[gcr]'       # GCR registry
pip install 'metaflow-prebuilt[kaniko]'    # Kaniko build service (K8s)

Quick start

from metaflow import FlowSpec, step, conda

class MyFlow(FlowSpec):

    @conda(packages={"numpy": "1.26.4"})
    @step
    def start(self):
        import numpy
        print(numpy.__version__)
        self.next(self.end)

    @step
    def end(self):
        pass

if __name__ == "__main__":
    MyFlow()

Deploy with the prebuilt environment (requires a configured registry and build service — see below):

# Deploy: builds and pushes Docker images, then submits the flow
python my_flow.py --environment=prebuilt run

# Or on Batch/Kubernetes:
python my_flow.py --environment=prebuilt batch run

Configuration

Select backends via environment variables:

Extension points

DockerBuildService

Controls how Docker images are built and pushed. Select with METAFLOW_PREBUILT_BUILD_SERVICE=<name>.

ImageRegistry

Controls where images are stored. Select with METAFLOW_PREBUILT_IMAGE_REGISTRY=<name>.

Contributing a custom backend

Custom build service

Subclass DockerBuildService and implement build_and_push:

from metaflow_extensions.prebuilt.plugins.conda.build_service import DockerBuildService

class MyBuildService(DockerBuildService):
    def build_and_push(self, dockerfile, context_files, image_tag,
                       push_credentials, echo, target_platform=None) -> bool:
        # build and push; return True on success, False on failure.
        # target_platform (e.g. "linux/amd64") is the resolved step's arch for
        # cross-arch builds; None => builder default. Local builders should honor
        # it; remote builders with a fixed build arch may ignore it.
        ...

Register it in your pyproject.toml:

[project.entry-points."metaflow_prebuilt.build_services"]
mybuild = "mypackage.my_module:MyBuildService"

Then use it with METAFLOW_PREBUILT_BUILD_SERVICE=mybuild.

Custom image registry

Subclass ImageRegistry and implement image_tag, push_credentials, and pull_config:

from metaflow_extensions.prebuilt.plugins.conda.image_registry import ImageRegistry

class MyRegistry(ImageRegistry):
    def image_tag(self, env_id) -> str: ...
    def push_credentials(self) -> dict: ...
    def pull_config(self, pull_tag) -> dict: ...

Register it:

[project.entry-points."metaflow_prebuilt.image_registries"]
myregistry = "mypackage.my_module:MyRegistry"

Then use it with METAFLOW_PREBUILT_IMAGE_REGISTRY=myregistry.

License

Apache 2.0