Skip to content

Repository files navigation

metaflow-coordinator

CI PyPI License: Apache-2.0 Python 3.9+

Coordinate shared services across your parallel Metaflow workers — work queues, rate limiters, barriers, Redis, Postgres, custom APIs — scoped to one run.

How it works

The coordinator step runs a service that workers discover and call while both run concurrently. When workers finish, the service tears itself down automatically.

graph LR
    start([start]) --> C[run_coordinator]
    start --> L[launch_workers]
    L --> W["run_worker × N"]
    W --> JW[join_workers]
    C --> J([join])
    JW --> J
    J --> E([end])
    C <-.-> W
Loading

The problem

When you fan out to N parallel worker steps, those workers share nothing. There's no built-in way for them to pull from a shared queue, throttle concurrent API calls, or stream results back to the coordinator. The usual answer is Redis or a database — real infrastructure for a temporary need that exists only for the duration of one run.

Coordination patterns

Each pattern has a runnable example in examples/.

Task distribution

Pattern Use when… Example
Work Queue Tasks are heterogeneous; workers self-assign the next available item work_queue.py
Priority Queue Same, but high-priority items jump the queue as workers make discoveries priority_queue.py
Staged Pipeline Work flows through ordered stages; any worker can handle any stage staged_pipeline.py

Throttling & synchronization

Pattern Use when… Example
Rate Limiter Workers call a rate-limited API; enforce ≤ N concurrent requests fleet-wide rate_limiter.py
Synchronization Barrier Workers run in rounds and must all finish a round before the next begins gradient_aggregator.py

Shared state

Pattern Use when… Example
Shared Cache Workers have overlapping sub-problems; skip redundant work via a shared cache shared_cache.py (in-process) · redis_cache.py (Redis)
Broadcast + Consensus Every worker sees the same input; aggregate with majority vote or best-of-N agent_ensemble.py
Database Results Workers write structured results to a shared table; query with full SQL at the end postgres_results.py

Advanced

Pattern Use when… Example
Adaptive Search Prune weak candidates between rounds (successive halving / tournament) tournament.py
External Process The coordinator needs to run a binary (Redis, nginx, DuckDB) instead of a Python service shard_server.py
Nested Coordinators Sub-groups of workers each need their own independent coordinator nested_coordinators.py

Quick start

pip install metaflow-coordinator
# For AWS Batch / Kubernetes (URL exchange via S3):
pip install "metaflow-coordinator[s3]"

A coordinator step runs a service; worker steps discover and call it. Both run concurrently inside the same flow.

import httpx
from metaflow import FlowSpec, current, step
from metaflow_coordinator import (
    ResultCollector, await_service,
    coordinator_join, worker_join,
)

class MyFlow(FlowSpec):

    @step
    def start(self):
        self.coordinator_id = current.run_id
        self.worker_ids = list(range(4))
        self.next(self.run_coordinator, self.launch_workers)

    @step
    def run_coordinator(self):
        collector = ResultCollector(n_workers=4)
        collector.run(service_id=self.coordinator_id, namespace="my-flow")
        self.results = collector.results_by_worker   # {worker_id: value}
        self.next(self.join)

    @step
    def launch_workers(self):
        self.next(self.run_worker, foreach="worker_ids")

    @step
    def run_worker(self):
        url = await_service(self.coordinator_id, namespace="my-flow", timeout=60)
        httpx.post(f"{url}/submit", json={"worker_id": self.input, "result": self.input * 2})
        self.next(self.join_workers)

    @step
    @worker_join
    def join_workers(self, inputs):
        self.next(self.join)

    @step
    @coordinator_join
    def join(self, inputs):
        self.next(self.end)

    @step
    def end(self):
        print(self.results)

ResultCollector.run() blocks until all 4 workers post to POST /submit, then tears down. No threads, no custom FastAPI app, no /complete calls needed.

See examples/echo_service.py for a full runnable version.

Usage

Work Queue

@step
def run_coordinator(self):
    wq = WorkQueue(items=my_records, drain_delay=2.0)
    wq.run(service_id=self.coordinator_id, namespace="my-flow")
    self.results = wq.results
    self.next(self.join)

@step
def run_worker(self):
    url = await_service(self.coordinator_id, namespace="my-flow", timeout=60)
    while True:
        item = httpx.post(f"{url}/pull", json={}).json()
        if item["done"]:
            break
        httpx.post(f"{url}/submit", json={"item_id": item["item_id"], "result": process(item["payload"])})
    self.next(self.join_workers)

Rate Limiter

@step
def run_coordinator(self):
    SemaphoreService(max_concurrent=5, n_workers=self.n_workers).run(
        service_id=self.coordinator_id, namespace="my-flow"
    )
    self.next(self.join)

@step
def run_worker(self):
    url = await_service(self.coordinator_id, namespace="my-flow", timeout=60)
    httpx.post(f"{url}/acquire")     # blocks until a slot is free
    result = call_rate_limited_api()
    httpx.post(f"{url}/release")
    httpx.post(f"{url}/done")
    self.next(self.join_workers)

External process (Redis, Postgres, nginx, DuckDB, …)

ProcessService wraps any binary. The url_scheme parameter controls the URL workers receive — "redis" gives redis://host:port, "postgresql" gives postgresql://host:port, etc. Workers connect with their native client.

# @conda installs the redis-server binary
@conda(packages={"redis": "7.2"})
@step
def run_coordinator(self):
    tracker   = CompletionTracker(n_workers=self.n_workers)
    redis_svc = ProcessService(
        command=["redis-server", "--port", "{port}"],
        done=tracker.done,
        url_scheme="redis",
    )
    SessionServiceGroup({"redis": redis_svc, "tracker": tracker}).run(
        service_id=self.coordinator_id, namespace="my-flow"
    )
    self.next(self.join)

@pypi(packages={"redis": "5.0"})
@step
def run_worker(self):
    import redis
    urls = discover_services(self.coordinator_id, names=["redis", "tracker"],
                             namespace="my-flow", timeout=120)
    r = redis.Redis.from_url(urls["redis"])
    r.set("key", "value")
    httpx.post(f"{urls['tracker']}/complete")
    self.next(self.join_workers)

See examples/redis_cache.py and examples/postgres_results.py for full working examples.

Built-in services

from metaflow_coordinator import (
    CompletionTracker,   # fires done when N workers call POST /complete
    ResultCollector,     # fires done when N workers call POST /submit; stores results
    BarrierService,      # synchronization barrier across R rounds
    WorkQueue,           # pull-based task distribution
    SemaphoreService,    # concurrency limiter
)

Core API

Service types

from metaflow_coordinator import FastAPIService, ProcessService, SessionServiceGroup

FastAPIService(app=my_app, done=done_event, drain_delay=2.0).run(
    service_id=self.coordinator_id, namespace="my-flow"
)

ProcessService(
    command=["redis-server", "--port", "{port}"],
    done=tracker.done,
    url_scheme="redis",           # workers receive redis://host:port
    ready=HttpReady(path="/health"),
)

SessionServiceGroup({"redis": redis_svc, "tracker": tracker}).run(
    service_id=self.coordinator_id, namespace="my-flow"
)

Service discovery

Use await_service when the coordinator runs a single service, and discover_services when it runs a SessionServiceGroup. The names you pass must match the keys of the group dict.

# Single service (FastAPIService / ProcessService)
url = await_service(self.coordinator_id, namespace="my-flow", timeout=120)

# Multiple services — names must match SessionServiceGroup keys
# coordinator: SessionServiceGroup({"redis": redis_svc, "tracker": tracker}).run(...)
urls = discover_services(self.coordinator_id, names=["redis", "tracker"],
                         namespace="my-flow", timeout=120)
urls["redis"]    # → redis://host:port
urls["tracker"]  # → http://host:port

Join decorators

from metaflow_coordinator import coordinator_join, worker_join

@step
@coordinator_join          # handles merge_artifacts for the coordinator/workers join
def join(self, inputs):
    self.next(self.end)

@step
@worker_join               # handles merge_artifacts for foreach-reduce steps
def join_workers(self, inputs):
    self.results = [inp.result for inp in inputs]
    self.next(self.join)

Remote execution

Add @batch or @kubernetes to coordinator and worker steps. Service URLs are exchanged via S3 automatically — no VPC changes needed as long as coordinator and workers share a VPC. Install the S3 extra to enable this:

pip install "metaflow-coordinator[s3]"

Development

git clone https://github.com/npow/metaflow-coordinator.git
cd metaflow-coordinator
pip install -e ".[dev]"
pytest -v

License

Apache 2.0 — see LICENSE.

About

Coordinate shared services across your parallel Metaflow workers — work queues, rate limiters, barriers, Redis, Postgres, custom APIs — scoped to one run.

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages