federated-learning-nextflow

Turning a single-process FedAvg simulation into a real distributed workflow — every client a separate SLURM job on its own GPU, orchestrated by Nextflow, with model weights passed between rounds as files.

Algorithm	FedAvg — clients train locally, server averages weights, broadcasts, repeat
Orchestration	Nextflow DSL2, SLURM executor, one GPU per client task
State exchange	PyTorch state_dict on disk — the workflow manager only understands files
Verified run	4 clients × 5 rounds = 26 tasks, 6m 41s, exit 0 on Leonardo Booster
Result	MNIST 93.2% → 97.5% test accuracy over 10 rounds

What the exercise showed

• Nextflow DSL2 will not let you call a process twice. Federated learning is a loop. That collision is the whole engineering problem here, and the workaround is ugly but honest.
• File-based state exchange is what makes FL fit a workflow manager at all. Weights become artifacts, so Nextflow can track provenance, cache, and resume.
• Orchestration overhead is real but small — 26 SLURM jobs for 20 client-trainings, 0.5 CPU-hours total.

---

The interesting problem: a DAG engine running a loop

Nextflow builds a directed acyclic graph. A process may be invoked exactly once per workflow. Federated learning is inherently iterative — round n consumes round n−1's global model.

The workaround: import the same process under a different alias for every round, so each round gets an independent process instance.

include { ClientTrain as CT0; ServerAggregate as SA0 } from './modules/fl_processes'
include { ClientTrain as CT1; ServerAggregate as SA1 } from './modules/fl_processes'
...                                                   // up to CT9 / SA9

def CT = [CT0, CT1, CT2, CT3, CT4, CT5, CT6, CT7, CT8, CT9]
def SA = [SA0, SA1, SA2, SA3, SA4, SA5, SA6, SA7, SA8, SA9]

model = ServerInit()
for (int i = 0; i < r; i++) {
    model = SA[i]( CT[i]( model.combine(Channel.of(0..<n)), i ).collect(), i )
}

Piece	What it does
model.combine(Channel.of(0..<n))	Pairs the global model with each client id — one task per client
CT[i](...)	n client trainings launched in parallel, each its own SLURM job
.collect()	Barrier — waits for every client in the round to finish
SA[i](...)	FedAvg aggregation, emits the new global model for round i+1

→ Round count is capped at 10, with an explicit error if you ask for more. Adding rounds means adding aliases.

→ This is a genuine limitation of DAG-oriented workflow managers for iterative ML. Nextflow is built for bioinformatics pipelines — run each stage once, over many samples. A recursion-capable engine, or Nextflow's own preview recursion feature, would express this directly.

The three-script decomposition

The centralised simulation (fl/base_fed_avg.py) keeps clients and server as live objects in one process. A workflow can't do that — tasks are separate processes on separate nodes. So the loop body is split into scripts that communicate only through files:

Script	Reads	Writes
server_init.py	—	global_model.pt — untrained weights, runs once
client_train.py	global_model.pt, client id	client_<id>_model.pt after local epochs
server_aggregate.py	all client_*_model.pt	global_model.pt for the next round

→ Every message in the FL protocol becomes a file on disk. That's what buys you Nextflow's caching, -resume, and provenance for free — and it's the standard shape for making any iterative distributed algorithm workflow-native.

FedAvg itself is a uniform average, not weighted by dataset size:

for key in client_state_dicts[0].keys():
    stacked = torch.stack([client[key].float() for client in client_state_dicts], dim=0)
    averaged_state_dict[key] = torch.mean(stacked, dim=0)

Running it

nextflow run workflow/workflow.nf -profile hpc \
    --clients 4 --rounds 5 --epochs 2 --gpu

Parameter	Default	Meaning
--clients	2	Parallel client tasks per round
--rounds	3	FL rounds (max 10)
--epochs	1	Local epochs per client per round
--batch-size	32	Mini-batch size
--learning-rate	0.01	SGD learning rate
--one-class	false	Non-IID — one MNIST class per client
--gpu	false	CUDA inside client tasks

On a cluster, workflow/slurm/job_nextflow.sh submits the Nextflow driver as a job; Nextflow then submits each task itself:

sbatch workflow/slurm/job_nextflow.sh

Resource requests are per-process — only ClientTrain asks for a GPU:

process {
    executor = 'slurm'
    queue    = 'boost_usr_prod'
    cpus     = 2
    memory   = '8 GB'

    withName: 'ClientTrain' {
        cpus           = 4
        memory         = '32 GB'
        maxForks       = 4
        clusterOptions = '--gres=gpu:1'
    }
}

Results

4 clients, 10 rounds, 2 local epochs, batch 32, lr 0.01, IID split, one A100 per client.

Round	Train loss	Test loss	Test accuracy
1	0.25	0.22	93.2%
5	0.078	0.093	96.9%
10	0.035	0.083	97.5%

All four clients track each other almost exactly — expected for an IID split, where every client sees the same distribution and FedAvg has nothing to reconcile. --one-class makes it non-IID and the curves diverge.

Verified workflow run — 4 clients × 5 rounds on Leonardo Booster:

executor >  slurm (26)
[6c/f0413f] process > ServerInit                        [100%] 1 of 1 ✔
[e4/ef3702] process > CT0 (round=0 | client=2)          [100%] 4 of 4 ✔
[ee/930e1f] process > SA0 (round=0)                     [100%] 1 of 1 ✔
...
[63/d7398b] process > SA4 (round=4)                     [100%] 1 of 1 ✔
Duration    : 6m 41s
CPU hours   : 0.5
Succeeded   : 26

26 = 1 init + (4 clients × 5 rounds) + 5 aggregations. Per-client metrics and GPU telemetry — power, utilisation, temperature, network — are tracked in Weights & Biases.

Caveats

Caveat	Detail
10-round ceiling	An artefact of the alias workaround, not of FedAvg. More rounds = more aliases
Uniform averaging	FedAvg proper weights by client dataset size (n_k/n); this averages uniformly. Equivalent for equal splits, wrong for unbalanced ones
SLURM queue polling is flaky	WARN: [SLURM] queue status cannot be fetched ... exit status 143 appeared on long runs. Nextflow recovers, but it's noisy
MNIST, not a real workload	The point is the orchestration; the model is deliberately trivial
No secure aggregation	Weights are written to a shared filesystem in the clear. Real FL needs encryption or differential privacy — out of scope here

Layout

fl/
  base_fed_avg.py       centralised simulation — the starting point
  server_init.py        \
  client_train.py        }  the same loop, decomposed for the workflow
  server_aggregate.py   /
  common/               client, server, model, parser, utils
workflow/
  workflow.nf           the round loop and alias workaround
  modules/fl_processes.nf   ClientTrain and ServerAggregate definitions
  nextflow.config       SLURM executor and per-process resources
  slurm/job_nextflow.sh submits the Nextflow driver itself
results/                metrics and the DAG diagram
docs/                   presentation

Where this came from

Course	P2.12 — Hybrid Workflows for Federated AI, MHPC, ICTP / SISSA Trieste, 2025–26
Cluster	Leonardo Booster, CINECA
Skeleton	The abstract BaseClient/BaseServer interfaces, MNIST loader, model and CLI are course-provided, GPL-3.0 — see LICENSE
My work	The concrete Client and Server implementations, FedAvg aggregation, the three-script decomposition, the entire Nextflow layer, SLURM integration and W&B tracking
Write-up	docs/dynamic-workflows-presentation.pdf