Unit tests federated controller + fixes by leonard-amsler · Pull Request #1227 · epfml/disco

leonard-amsler · 2026-09-22T21:06:38Z

This PR involves:

The creation of a unit test suite for the federated controller
Some fixes applied to the controller to match its correct fonctionality
The modularization of the fake websocket for both unit decentralized and federated testing
Discovery of bugs that need attention and not corrected
Minor: Relocalization of the existing unit/end2end tests for decentralized and federated controllers into folders

Unit tests:

implemented at: server/tests/unit/federated_controller.spec.ts
18 tests in three groups: join handshake (3), aggregation (8, one commented out - see know issue), leaving and reset (7)

Controller discovered bugs:

Contributions for future rounds were accepted. isWithinRoundCutoff was only discarding previous rounds.
Contributions from a socket that never joined (ClientConnected was never sent by the client) were aggregated. We now ensure the connections store has an entry for the client before adding its update to the aggregator.
A stale client was never catching up the current round. A simple offset of 1 was removed to the round id given to stale clients.
The minimum number client needed wasn't reset when all peers disconnected from the protocol and that the controller was reset. The minimum is now set in the #makeAggregator() method, so set by all aggregators.

Aggregator bugs:

The contributions of a node that left the protocol were still stored, possibly validating the minimum of clients required (as the client left, the number of required updates decreases and the update still stored). In other words, the contributions were no longer a subset of the active nodes. We now discard the contributions of unactive clients with disposeContributionsOf in the aggregator.
We weren't checking if the update thresholds were met after the departure of a node. When the last client that was still missing to contribute for a round was leaving, we were in a stale mode forever. Now, we check if the aggregator is full also when nodes quit.

Known issue but not fixed:

A socket that never sends ClientConnected increases the absolute thresholds forever. It blocks the current round even if it sends updates. We could either:
- Notify the client that the server did not receive its ClientConnected message so that it can resend it to be registered in the connections set.
- Skip the need for the ClientConnected message and do its content when receiving an update.
The aggregator is fixed and hardcoded to the mean aggregator with no stale cutoff and a 100% relative aggregation strategy. The usage of the Byzantine robust aggregator is impossible. If we choose, we could use the abstract class aggregator in the controller and instantiate the appropriate one in the constructor (similar to as the getAggregator() for the client). We could add the aggregation threshold in the task description. This would need both server and UI changes. See Hardcoded MeanAggregator for the federated server #1238

Fixes #1220, fixes #1221, fixes #1222

JulienVig · 2026-09-23T11:03:45Z

Regarding known issue not fixed #1: how would you skip the ClientConnected message? In FL for example, it's currently used to let the server know how many contributions it should wait for and send the clients the base model

…t a connection message

leonard-amsler · 2026-09-27T15:18:27Z

New progress on the PR:

Synchronization of the connection list with the aggregator's registered peers. A connected client that does not send the ClientConnected message does not block the training anymore.
The client now sends the ClientConnected message with retries (default 5 with 5 seconds timeout) until it receives the answer from the server. Resolve the problem of a missed ClientConnected message by the server.
If it happens that the server receives model update from a non-registered client, the servers sends a notification to the client and make it crash. This is a robustness mechanism as it should never happen in a correct execution.

…-federated-controller

JulienVig

Thanks for the great work! I left mostly superficial comments and a bit of refactoring but there are two points I would like to discuss:

What should the server answer to a client that sent a future contribution?
Sending the model weights every time we receive a Client Connected sounds a bit risky regarding DDoS attacks, I think we should implement a counter

…ons from the federated specs

…stead of private functions

…oid waiting after reconnection

…ient

… (crash)

…ions

…ebsocket

…-federated-controller

leonard-amsler · 2026-10-07T08:38:24Z

New state

CrashClient: The federated server can now stop a client with a reason. The client aborts any waiting state (waiting for a server message) immediately, or stops before sending its update if it arrives during local training. The caller of Disco.train gets a ClientCrashError, and the webapp shows a dedicated toast.
Browser WebSocket fix: the new @epfml/isomorphic-ws package uses ws (keeping the 1 GiB message limit) on Node and the native WebSocket in browsers. Same structure as @epfml/isomorphic-wrtc. Since GPT-2 training with Goldfish loss #1111, { maxPayload } made browser connections fail. Replaces main's runtime check from Test federated webapp training #1230.
ClientConnected limit: Implementation of a local counter to detect abusive clients sending too many ClientConnected messages, flooding the network with weights. If more than 5 are sent per websocket, the server sends CrashClient instead of resending the global weights.
Smaller fixes and refactors: reconnecting clients no longer stuck waiting for participants, federated specs moved next to their code, fewer and longer join retries.

Tests

Cypress training/server-crash.cy.ts: Webapp test with a scripted WebSocket server (cy.task("scriptServer")).

Notes

The e2e server URL is now http://server:8090. It was previously either http://server:1351 (.env.test) or http://server (test:e2e scripts) with a default port 80, not accessible. To attach a websocket we need a proper accessible hostname and port. The "server" gets translated to 127.0.0.1. Chose port 8090 as it is available, above 1024 (accessible), not 1351 (vite) and not 8080 (real disco server). The scripted server only starts for the server host, so the other tests keep 8080 for their real server.
A crashed client stays registered and the round waits for its contribution until it closes its socket, which is up to the caller of disco.train (the webapp does it).

… them

leonard-amsler · 2026-10-08T14:15:15Z

Description of the mechanism to make a webapp test with a scripted server

Cypress can intercept HTTP requests by calling cy.intercept and send a predefined answer. However, there is no mechanism to intercept websocket request natively in cypress.

In order to make a test that relies on websocket connections of the server (such as cypress/e2e/training/scripted/server-crash.cy.ts), we run a small websocket server scripted in cypress.config.ts, separately from the other e2e tests.

Run	`VITE_SERVER_URL`	HTTP to the server	Websocket
Regular e2e	`http://server`	intercepted by `setupServerWith`, other requests fail (`server` doesn't exist in DNS)	none
Scripted (`DISCO_SCRIPTED_SERVER_E2E=1`)	`http://server:8080`	intercepted by `setupServerWith`, as above	scripted server on `127.0.0.1:8080`
Federated training	`http://localhost:8080`	real server	real server

How the scripted run works

Only with DISCO_SCRIPTED_SERVER_E2E=1: the scripted/ tests are included, Cypress maps server to 127.0.0.1 (hosts), and the scripted server starts on the URL's port.
Tests script its answers by message type with cy.task("scriptServer", { [type]: [messages] }).
The config fails right away if the URL isn't http://server:<port>.
CI job: test-e2e-scripted. Locally: pnpm -F webapp test:e2e:scripted. Make sure you stop a local DISCO server on 8080 first.
There's no conflict with the federated jobs' real server on port 8080, as they never run in the same job.
The URL needs an explicit port, as the scripted server listens on it. .env.test used to say http://server:1351, but 1351 is vite's port for the webapp, which the scripted server can't share. So it now uses 8080, the real server's default port.

The regular e2e tests are unchanged. They keep http://server without a port and without the mapping, so nothing listens there and a forgotten request still fails.

JulienVig

Thanks for the great work, only nitpicking comments about names and documentation. I will create a follow up PR to yours to address a couple things that I realized when reviewing this PR. I'll let you know exactly what it's adressing so that we don't fix the same things in your next branch

JulienVig · 2026-10-08T12:24:05Z

+      messages.isMessageFromServer, // can only receive federated message types from the server
+      messages.isMessageToServer, // idem for messages that the client can send


Suggested change

messages.isMessageFromServer, // can only receive federated message types from the server

messages.isMessageToServer, // idem for messages that the client can send

messages.isMessageFromServer,

messages.isMessageToServer,

The comments are outdated now

JulienVig · 2026-10-08T12:34:22Z

  readonly #preprocessOnce: boolean;
  // Forwarded to compatible models to identify this client in debug output.
  readonly #debugLabel?: string;
+  #closed: boolean;


Can you add a comment describing the attribute?

Fyi, you can also declare it directly in the constructor arguments:

constructor( ... #closed = false, ... )

JulienVig · 2026-10-08T12:37:13Z

   */
  async close(): Promise<void> {
+    if (this.#closed) return;
+    this.#closed = true;


We should set the flag only once the cleanup has been successful in case something interrupts it so I think it's better to move it done to the very bottom (after the finally, not inside)

JulienVig · 2026-10-09T11:02:12Z

      return;
+    } else if (e instanceof ClientCrashError) {
+      toaster.error(
+        "The server stopped your training.<br/>Please rejoin the training by refreshing the page.",


Suggested change

"The server stopped your training.<br/>Please rejoin the training by refreshing the page.",

"The server have disconnected you due to an unexpected event.<br/> Feel free to rejoin the training.",

Let's not try to troubleshoot too much or be too precise in the error message. One future feature is to preserve the state when refreshing so this message would get outdated and we also don't know whether this hypothetical error would actually go away by rejoining.

JulienVig · 2026-10-09T11:09:48Z

+  /**
+   * Maximum number of times a client can try to connect with the same client ID.
+   */
+  static readonly MAX_CLIENT_CONNECTED_PER_SOCKET = 5;


Suggested change

static readonly MAX_CLIENT_CONNECTED_PER_SOCKET = 5;

static readonly MAX_CLIENT_CONNECTION_RETRY = 5;

I know that it's the name of the message type but I find the variable name confusing on its own, it sounds like we could connect up to 5 clients per socket.

JulienVig · 2026-10-09T11:44:43Z

+import { goToTaskOverview } from "../../../support/training";
+
+describe("training page", () => {
+  it("tells the user when the server make the client crash stops the training", () => {


Suggested change

it("tells the user when the server make the client crash stops the training", () => {

it("tells the user when the server crashes the client", () => {

JulienVig · 2026-10-09T11:47:48Z

            payload:
              this.#aggregator.round === 0
-                ? undefined
+                ? undefined // Optimization: no needs to send the initial weights, the client already has them


Suggested change

? undefined // Optimization: no needs to send the initial weights, the client already has them

? undefined // Optimization: no need to send the initial weights, the client already has them

JulienVig · 2026-10-09T11:57:59Z

  // Kept last as the enum values are what goes over the wire.
  ParticipantsUpdate,
+  CrashClient,


Can you document what the CrashClient message is?

leonard-amsler added 4 commits September 22, 2026 19:20

relocation of decentralized tests

18981bc

Modularization of the fake websocket

3baea8f

Test suit for the federated controller & controller fixs

1cd6950

knip & lint

3ce7195

leonard-amsler added 5 commits September 27, 2026 14:38

client's addition to aggregator when receiving ClientConnected + test

697b1de

Helper method to send msg with retry and wait for another one + test

f29ba63

Retry sending the ClientConnected until receiving server's answer

0c33585

Notify client and making it crash if server receives an update withou…

8e5b102

…t a connection message

Merge main into unit-tests-federated-controller

979da3f

leonard-amsler force-pushed the unit-tests-federated-controller branch from e51de03 to 979da3f Compare September 27, 2026 15:11

leonard-amsler marked this pull request as ready for review September 28, 2026 15:12

leonard-amsler requested a review from JulienVig September 28, 2026 15:12

Merge branch 'main' of https://github.com/epfml/disco into unit-tests…

e9892c3

…-federated-controller

JulienVig requested changes Sep 30, 2026

View reviewed changes

This was linked to issues Sep 30, 2026

Federated learning staleness bugs #1219

Open

Fix stale contribution recovery round number #1220

Open

Remove departed contributors from pending federated aggregation #1221

Open

Define first-round stale contribution recovery #1222

Open

leonard-amsler added 5 commits October 2, 2026 20:26

move specs next to code they test, drop timeouts and non-null asserti…

3d07a2a

…ons from the federated specs

split federated messages into from/to groups

e1c3f2f

leave error handling to the callers of Disco.train

09334f9

fewer and longer ClientConnected retries

e196026

document participant sets in the federated controller, # functions in…

e4cc7e9

…stead of private functions

leonard-amsler force-pushed the unit-tests-federated-controller branch from 91e94d2 to e4cc7e9 Compare October 2, 2026 18:35

leonard-amsler added 4 commits October 5, 2026 16:16

clear the participants wait when the federated client disconnects, av…

ae19374

…oid waiting after reconnection

crash the federated client at any point when the server sends CrashCl…

39b1c15

…ient

Remake of isomorphic-ws to use maxPayload only for node, not browsers…

c0e2098

… (crash)

definition of server port in webapp tests to enable websocket connect…

9c774b2

…ions

leonard-amsler added 4 commits October 6, 2026 09:05

webapp error toaster on ClientCrashError

ed6d471

cypress test for the ClientCrashError toast, with a scripted server w…

b3c6881

…ebsocket

Merge branch 'main' of https://github.com/epfml/disco into unit-tests…

c07c136

…-federated-controller

local counter per ws to avoid ClientConnected abuse

0ce589f

leonard-amsler mentioned this pull request Oct 7, 2026

Decentralized and federated controller consistency #1250

Open

run the scripted server e2e tests separately, mapping server only for…

3416c92

… them

leonard-amsler requested a review from JulienVig October 8, 2026 08:20

JulienVig approved these changes Oct 9, 2026

View reviewed changes

JulienVig added this pull request to stack #1252 October 9, 2026 15:46

JulienVig mentioned this pull request Oct 9, 2026

Server forcefully disconnects crashed clients #1253

Draft

		messages.isMessageFromServer, // can only receive federated message types from the server
		messages.isMessageToServer, // idem for messages that the client can send

	"The server stopped your training.<br/>Please rejoin the training by refreshing the page.",
	"The server have disconnected you due to an unexpected event.<br/> Feel free to rejoin the training.",

	static readonly MAX_CLIENT_CONNECTED_PER_SOCKET = 5;
	static readonly MAX_CLIENT_CONNECTION_RETRY = 5;

	it("tells the user when the server make the client crash stops the training", () => {
	it("tells the user when the server crashes the client", () => {

	? undefined // Optimization: no needs to send the initial weights, the client already has them
	? undefined // Optimization: no need to send the initial weights, the client already has them

Conversation

leonard-amsler commented Sep 22, 2026 • edited by JulienVig Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

This PR involves:

Unit tests:

Controller discovered bugs:

Aggregator bugs:

Known issue but not fixed:

Uh oh!

JulienVig commented Sep 23, 2026

Uh oh!

leonard-amsler commented Sep 27, 2026

Uh oh!

JulienVig left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

leonard-amsler commented Oct 7, 2026

New state

Tests

Notes

Uh oh!

leonard-amsler commented Oct 8, 2026

Description of the mechanism to make a webapp test with a scripted server

How the scripted run works

Uh oh!

JulienVig left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

leonard-amsler commented Sep 22, 2026 •

edited by JulienVig

Loading