# metaflow-coordinator

[![CI](https://github.com/npow/metaflow-coordinator/actions/workflows/ci.yml/badge.svg)](https://github.com/npow/metaflow-coordinator/actions/workflows/ci.yml) [![PyPI](https://img.shields.io/pypi/v/metaflow-coordinator)](https://pypi.org/project/metaflow-coordinator/) [![License: Apache-2.0](https://img.shields.io/badge/License-Apache--2.0-blue.svg)](LICENSE) [![Python 3.9+](https://img.shields.io/badge/python-3.9+-blue.svg)](https://www.python.org/downloads/)

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.

```mermaid
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
```

## 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

```bash
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.

```python
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`](examples/echo_service.py) for a full runnable version.

## Usage

### Work Queue

```python
@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

```python
@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.

```python
# @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`](examples/redis_cache.py) and [`examples/postgres_results.py`](examples/postgres_results.py) for full working examples.

### Built-in services

```python
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

```python
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.

```python
# 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

```python
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:

```bash
pip install "metaflow-coordinator[s3]"
```

## Development

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

## License

Apache 2.0 — see [LICENSE](LICENSE).
