Running a Parameter Server with Flame App

By Klaus Ma
Flame ML

In AI-related distributed workloads, large amounts of data often need to move between many worker nodes. Flame provides an object cache to help pass data objects across nodes. This walkthrough uses a simple Parameter Server (ps–worker) pattern to show how Flame App fits a training-style workflow. Later posts in the series will cover internal mechanics and optimizations for different scenarios.

Overview

The ps–worker pattern is a classic distributed training architecture:

  • The Parameter Server holds global model weights and applies gradient updates.
  • Multiple workers compute gradients on different data batches in parallel.
  • Each worker fetches the latest weights, computes gradients, and sends them back to the Parameter Server.

This example uses flamepy.app to declare distributed services and handle communication between them.

Architecture

   ┌─────────────────────┐
   │  Parameter Server   │  - Stores model weights
   │                     │  - Applies aggregated gradients
   └─────────┬───────────┘
             │
      ┌──────┴──────┐
      │             │
┌─────▼─────┐   ┌─────▼─────┐  - Fetch weights
│ Worker 1  │   │ Worker 2  │  - Compute gradients on data batches
└───────────┘   └───────────┘

Components

  • ConvNet — Small convolutional network for MNIST digit classification.
  • ParameterServer — Keeps model state and applies gradient updates with SGD.
  • DataWorker — Computes gradients on mini-batches from the training set.
  • Main training loop — Coordinates synchronous training iterations.

Files

  • main.py — Entry point: initializes the app, declares services, and runs the training loop.
  • ps.py — Model, parameter server, and data workers.
  • pyproject.toml — Dependencies and project metadata.

Requirements

  • Python >= 3.12
  • PyTorch and torchvision
  • Flame Python SDK (flamepy)
  • NumPy
  • filelock

How to run

Build the Flame cluster (if it is not already running):

docker compose build
docker compose up -d

Go to the example directory:

cd examples/ps

Run the example:

python main.py

The script downloads the MNIST dataset on first run and starts training.

Expected output

100.0%
100.0%
100.0%
100.0%
Running synchronous parameter server training.
Iter 0:         accuracy is 16.5
Iter 10:        accuracy is 32.9
Final accuracy is 32.9.

In this simplified run, accuracy usually rises from about 10% (random guessing) to roughly 30–35% after 20 training iterations, as shown above. Higher accuracy may require more iterations, a larger model, or evaluation on the full dataset.

Key concepts

1. Declaring services with Flame App

import flamepy.app as app
from ps import DataWorker, ParameterServer

app.init("ps-example")

@app.service(autoscale=False, warmup=1)
class ParameterServerService(ParameterServer):
    def __init__(self):
        super().__init__(1e-2)

@app.service(warmup=2)
class DataWorkerService(DataWorker):
    pass

ps_svc = ParameterServerService()
worker_svc = DataWorkerService()
workers_svc = [worker_svc, worker_svc]

The decorators register the two service types. The parameter server uses one fixed instance; worker calls can run on warm worker instances. The full example calls app.destroy() in a finally block.

2. Asynchronous remote calls

gradients = [worker.compute_gradients(current_weights) for worker in workers_svc]
current_weights = ps_svc.apply_gradients(*gradients).get()

Method calls on services return futures. Use .get() to block and retrieve the result.

3. Synchronous training

The training loop waits for every worker to finish its gradients before the parameter server applies an update:

for i in range(20):
    # Start all gradient computations in parallel
    gradients = [worker.compute_gradients(current_weights) for worker in workers_svc]
    # Wait for all gradients, then apply the update
    current_weights = ps_svc.apply_gradients(*gradients).get()

This is a synchronous parameter server: each iteration waits for all workers.

For the App lifecycle, see Running Monte Carlo Pi with Flame App (Python).

Customization

Change the number of workers in main.py:

@app.service(warmup=4)
class DataWorkerService(DataWorker):
    pass

Change the learning rate when constructing the parameter server:

@app.service(autoscale=False, warmup=1)
class ParameterServerService(ParameterServer):
    def __init__(self):
        super().__init__(1e-3)

Train for more iterations by changing the loop bound in main.py from range(20) to e.g. range(50).

Notes

  • The example uses filelock so MNIST can be downloaded safely when several workers start at the same time.
  • Evaluation is capped at 1024 samples so iteration during development stays fast.
  • The model is deliberately small: the goal is to demonstrate the pattern, not state-of-the-art accuracy.
  • Reinforcement learning: see the examples/rl/ directory in the Flame repo.
  • Other distributed patterns: browse under examples/.

Troubleshooting

If something goes wrong:

  1. Services do not start — Confirm the Flame cluster is up (docker compose ps).
  2. Import errors — After changing sdk/python, rebuild containers: docker compose build.
  3. Test timeouts — Inspect logs: docker logs flame-executor-manager and docker logs flame-session-manager.

References

Related Posts