From 649a984392360df16fffe4f2d4b36a20940fc2ea Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?L=C3=A9onard=20Amsler?= Date: Wed, 14 Oct 2026 12:39:36 +0200 Subject: [PATCH 1/6] apply and abstract federated controller fixes to the decentralized one --- discojs/src/client/decentralized/messages.ts | 5 +- .../controllers/decentralized_controller.ts | 199 +++++++++++++----- .../src/controllers/federated_controller.ts | 131 +++--------- server/src/controllers/training_controller.ts | 120 ++++++++++- 4 files changed, 286 insertions(+), 169 deletions(-) diff --git a/discojs/src/client/decentralized/messages.ts b/discojs/src/client/decentralized/messages.ts index 0bd18011a..a0bf25138 100644 --- a/discojs/src/client/decentralized/messages.ts +++ b/discojs/src/client/decentralized/messages.ts @@ -8,6 +8,7 @@ import type { WaitingForMoreParticipants, EnoughParticipants, ParticipantsUpdate, + CrashClient, } from "#client/mtype"; /// Phase 0 communication (between server and peers) @@ -116,7 +117,8 @@ export type MessageFromServer = | RetryPeerConnections | ConnectionFail | ModelProviderInfo - | ProvideModelToPeer; + | ProvideModelToPeer + | CrashClient; export type MessageToServer = | ClientConnected @@ -151,6 +153,7 @@ export function isMessageFromServer(o: unknown): o is MessageFromServer { case MType.ConnectionFail: case MType.ModelProviderInfo: case MType.ProvideModelToPeer: + case MType.CrashClient: return true; } diff --git a/server/src/controllers/decentralized_controller.ts b/server/src/controllers/decentralized_controller.ts index 42a78a74b..8a4cb9893 100644 --- a/server/src/controllers/decentralized_controller.ts +++ b/server/src/controllers/decentralized_controller.ts @@ -1,5 +1,4 @@ import createDebug from "debug"; -import { v4 as randomUUID } from "uuid"; import * as msgpack from "@msgpack/msgpack"; import type WebSocket from "ws"; import { Map, Set } from "immutable"; @@ -14,6 +13,21 @@ import MessageTypes = mtype.MType; const debug = createDebug("server:controllers:decentralized"); +/** + * Decentralized training server for a single task. + * + * Six variables track participants at different stages. + * + * - `clientIds` (inherited): connection between an ID and an open WebSocket. + * It is only used to guarantee ids don't collide. + * - `connections` (inherited): map between the ids of clients + * that completed the connection handshake and their socket. + * - `#roundPeers`: the nodes that want to join the current round and their readiness status. + * - `#connectFinishedNodes`: the nodes that have completed the connection handshake and are ready for the current round. + * - `#syncingNodes`: the nodes that are currently syncing with the latest model. + * - `#providerNode`: the client selected to provide the latest model to peers joining in the middle of training. + * + */ export class DecentralizedController< D extends DataType, > extends TrainingController { @@ -40,32 +54,64 @@ export class DecentralizedController< const minNbOfParticipants = this.task.trainingInformation.minNbOfParticipants; - // Peer id of the message sender - let peerId = randomUUID(); - while (this.connections.has(peerId)) { - peerId = randomUUID(); - } + // Generate a new unique client ID for the participant + const peerId = this.newClientId(); const shortId = peerId.slice(0, 4); // How the server responds to messages ws.on("message", (data: Buffer) => { try { const msg: unknown = msgpack.decode(data); - if (!messages.isMessageToServer(msg)) + if (!messages.isMessageToServer(msg)) { + this.crashClient(ws, peerId, "Invalid message received"); return debug("invalid message received: %o", msg); + } + // If the client has not yet established a connection + // and the message is not a ClientConnected message, + // we consider it as coming from an unconnected client + // and respond with a CrashClient message. + if ( + !this.connections.has(peerId) && + msg.type !== MessageTypes.ClientConnected + ) { + debug( + "Received message from an unconnected client [%s], sending CrashClient message", + shortId, + ); + this.crashClient(ws, peerId, "No ClientConnected message received"); + return; + } switch (msg.type) { // A new peer joins the network for a task case MessageTypes.ClientConnected: { - debug(`peer [%s] joined ${this.task.id}`, shortId); - this.connections = this.connections.set(peerId, ws); - - // If the new peer joins in the middle of training, - // it needs to get the latest model from an existing peer - let joinedMidTraining = false; - if (this.#aggregationRound > 0) { - joinedMidTraining = true; - this.#syncingNodes = this.#syncingNodes.add(peerId); + // Verify if this is a new client connection + const isNewClient = !this.connections.has(peerId); + if (!isNewClient) { + debug( + `Duplicate client connection detected for client [%s]`, + shortId, + ); + } else { + debug(`New client connection for client [%s]`, shortId); + // Connect the new client to both the connections map and the aggregator + this.connectClient(peerId, ws); + } + + // Count the client's join attempt + // Send crash message if it exceeded the maximum number of connection attempts + if (!this.countJoinAttempt(peerId)) { + debug( + "Client [%s] exceeded the maximum number of connections (%d), sending CrashClient message", + shortId, + TrainingController.MAX_CLIENT_CONNECTION_RETRY, + ); + this.crashClient( + ws, + peerId, + "Exceeded maximum number of retries to connect", + ); + return; } // Answer with client id in an NewNodeInfo message @@ -75,12 +121,13 @@ export class DecentralizedController< nbOfParticipants: this.connections.size, waitForMoreParticipants: this.connections.size < minNbOfParticipants, - joinedMidTraining: joinedMidTraining, + joinedMidTraining: this.#syncingNodes.has(peerId), }; - ws.send(msgpack.encode(msg), { binary: true }); + this.send(ws, peerId, msgpack.encode(msg)); + // Send an update to participants if we can start/resume training, // which already carries the number of participants - if (!this.sendEnoughParticipantsMsgIfNeeded(peerId)) + if (isNewClient && !this.sendEnoughParticipantsMsgIfNeeded(peerId)) // otherwise just tell them that someone joined this.sendParticipantsUpdateMsg(peerId); break; @@ -97,7 +144,7 @@ export class DecentralizedController< case MessageTypes.PeerIsReady: { this.#roundPeers = this.#roundPeers.set(peerId, true); debug("Received peer ready from: %o", shortId); - this.sendPeersForRoundIfNeeded(); + this.#sendPeersForRoundIfNeeded(); break; } // Forwards a peer's message to another destination peer @@ -109,7 +156,15 @@ export class DecentralizedController< peer: peerId, signal: msg.signal, }; - this.connections.get(msg.peer)?.send(msgpack.encode(forward)); + try { + this.#sendMessageToNode(msg.peer, forward); + } catch (err) { + debug( + "Failed to forward signal to node [%s]: %o", + msg.peer.slice(0, 4), + err, + ); + } break; } case MessageTypes.ConnectionsReady: { @@ -127,7 +182,7 @@ export class DecentralizedController< peerId, true, ); - this.signalWeightSharing(); + this.#signalWeightSharing(); break; } case MessageTypes.ModelSyncRequest: { @@ -143,16 +198,30 @@ export class DecentralizedController< type: MessageTypes.ModelProviderInfo, providerNode: this.#providerNode, }; - this.connections.get(peerId)?.send(msgpack.encode(providerInfo)); + try { + this.#sendMessageToNode(peerId, providerInfo); + } catch (err) { + debug( + "Failed to send provider info to node [%s]: %o", + peerId.slice(0, 4), + err, + ); + } // Signal the provider client with newly joined client's information const newNodeInfo: messages.ProvideModelToPeer = { type: MessageTypes.ProvideModelToPeer, newNode: peerId, }; - this.connections - .get(this.#providerNode) - ?.send(msgpack.encode(newNodeInfo)); + try { + this.#sendMessageToNode(this.#providerNode, newNodeInfo); + } catch (err) { + debug( + "Failed to send new node info to provider node [%s]: %o", + this.#providerNode.slice(0, 4), + err, + ); + } break; } default: { @@ -162,14 +231,13 @@ export class DecentralizedController< } } catch (e) { debug("when processing WebSocket message: %o", e); + this.crashClient(ws, peerId, "Server failed to handle a message"); } }); // Setup callback for client leaving the session ws.on("close", () => { // Remove the participant when the websocket is closed - this.connections = this.connections.delete(peerId); - this.#roundPeers = this.#roundPeers.delete(peerId); - this.#connectFinishedNodes = this.#connectFinishedNodes.delete(peerId); + this.disconnectClient(peerId); // Reset the training session when all participants leave if (this.connections.size === 0) { @@ -199,7 +267,7 @@ export class DecentralizedController< // If no, check if we are still above the minimum number of participant required if (this.connections.size >= minNbOfParticipants) { this.sendParticipantsUpdateMsg(); - this.sendPeersForRoundIfNeeded(); + this.#sendPeersForRoundIfNeeded(); return; } // If we are below the minimum number of participants @@ -207,10 +275,15 @@ export class DecentralizedController< // which already carries the number of participants this.sendWaitForMoreParticipantsMsg(); }); + + // Setup callback for WebSocket errors + ws.on("error", (err) => { + this.disconnect(ws, peerId, err); + }); } - reset(): void { - this.resetConnectionState(); + protected override reset(): void { + super.reset(); this.#roundPeers = Map(); this.#connectFinishedNodes = Map(); @@ -226,25 +299,23 @@ export class DecentralizedController< } } - private sendMessageToNode( + #sendMessageToNode( nodeId: NodeID, message: messages.MessageFromServer, ): void { const conn = this.connections.get(nodeId); if (conn === undefined) { - throw new Error( - `peer ${nodeId} marked as ready but not connection to it`, - ); + throw new Error(`No connection to node ${nodeId}, cannot send message`); } - conn.send(msgpack.encode(message)); + this.send(conn, nodeId, msgpack.encode(message)); } - private broadcastMessageToNodes( + #broadcastMessageToNodes( nodeIds: Iterable, createMessage: (nodeId: NodeID) => messages.MessageFromServer, ): void { for (const nodeId of nodeIds) { - this.sendMessageToNode(nodeId, createMessage(nodeId)); + this.#sendMessageToNode(nodeId, createMessage(nodeId)); } } @@ -253,7 +324,7 @@ export class DecentralizedController< * and if all peers that joined the round are ready to exchange weight updates * If so, send the list of peers for this round to all participants */ - private sendPeersForRoundIfNeeded(): void { + #sendPeersForRoundIfNeeded(): void { const minNbOfParticipants = this.task.trainingInformation.minNbOfParticipants; const nbOfPeersReady = this.#roundPeers.filter((ready) => ready).size; @@ -274,7 +345,7 @@ export class DecentralizedController< ) return; // Once every peer that joined the round is ready, we can start the round - this.broadcastMessageToNodes( + this.#broadcastMessageToNodes( this.#roundPeers.keySeq(), (id): messages.PeersForRound => { debug("Sending peer list to: %o", id.slice(0, 4)); @@ -294,7 +365,7 @@ export class DecentralizedController< this.#roundPeers = this.#roundPeers.map(() => false); // Start timeout to check peer connections are successful - this.startTimeout(); + this.#startTimeout(); } /** @@ -302,15 +373,15 @@ export class DecentralizedController< * with other peers in the round * If so, send StartWeightSharing message to signal peers to proceed */ - private signalWeightSharing(): void { + #signalWeightSharing(): void { // Return if not all participants are ready if (!this.#connectFinishedNodes.every((ready) => ready)) return; // Stop the timeout - this.clearTimeout(); + this.#clearTimeout(); // Send round participants StartWeightSharing messages - this.broadcastMessageToNodes( + this.#broadcastMessageToNodes( this.#roundPeers.keySeq(), (id): messages.StartWeightSharing => { debug("Signaling weight sharing to: %o", id.slice(0, 4)); @@ -330,11 +401,11 @@ export class DecentralizedController< /** * Set a timeout to check peer connections establishment */ - private startTimeout( + #startTimeout( maxTime: number = this.task.trainingInformation.maxPeerConnectionTime, ): void { this.#timeout = setTimeout(() => { - this.handleTimeout(); + this.#handleTimeout(); }, maxTime); } @@ -342,7 +413,7 @@ export class DecentralizedController< * Clear previously set timeout once all peer connections * are established before the timeout */ - private clearTimeout(): void { + #clearTimeout(): void { if (this.#timeout !== undefined) { clearTimeout(this.#timeout); this.#timeout = undefined; @@ -354,8 +425,8 @@ export class DecentralizedController< * Signals peers to discard existing connections and * reestablish connections with the current set of peers */ - private handleTimeout(): void { - this.clearTimeout(); + #handleTimeout(): void { + this.#clearTimeout(); debug( `Connection setup timeout for round ${this.#aggregationRound}, Retrying with same peers`, ); @@ -377,14 +448,14 @@ export class DecentralizedController< // Signal the node that connection is failed for that node debug("Signaling connection failure to: %o", nodeId.slice(0, 4)); - this.sendMessageToNode(nodeId, { + this.#sendMessageToNode(nodeId, { type: MessageTypes.ConnectionFail, }); } }); // Restart the round with remaining clients - this.broadcastMessageToNodes( + this.#broadcastMessageToNodes( this.#roundPeers.keySeq(), (id): messages.RetryPeerConnections => { debug("Signaling connection retry to: %o", id.slice(0, 4)); @@ -402,7 +473,7 @@ export class DecentralizedController< // Reset the connectionRetry since we excluded the failed clients this.#connectionRetry = 0; // Restart the timeout after sending retry messages - this.startTimeout(); + this.#startTimeout(); return; } @@ -411,7 +482,7 @@ export class DecentralizedController< this.#roundPeers = this.#roundPeers.map(() => false); this.#connectFinishedNodes = this.#roundPeers.map(() => false); - this.broadcastMessageToNodes( + this.#broadcastMessageToNodes( this.#roundPeers.keySeq(), (id): messages.RetryPeerConnections => { debug("Signaling connection retry to: %o", id.slice(0, 4)); @@ -423,6 +494,24 @@ export class DecentralizedController< ); // Restart the timeout after sending retry messages - this.startTimeout(); + this.#startTimeout(); + } + + protected override connectClient(nodeId: string, ws: WebSocket): void { + super.connectClient(nodeId, ws); + + // If the new peer joins in the middle of training, + // it needs to get the latest model from an existing peer + if (this.#aggregationRound > 0) { + this.#syncingNodes = this.#syncingNodes.add(nodeId); + } + } + + protected override disconnectClient(nodeId: string): void { + super.disconnectClient(nodeId); + + this.#roundPeers = this.#roundPeers.delete(nodeId); + this.#connectFinishedNodes = this.#connectFinishedNodes.delete(nodeId); + this.#syncingNodes = this.#syncingNodes.delete(nodeId); } } diff --git a/server/src/controllers/federated_controller.ts b/server/src/controllers/federated_controller.ts index 4db12801b..3b29aff8b 100644 --- a/server/src/controllers/federated_controller.ts +++ b/server/src/controllers/federated_controller.ts @@ -1,6 +1,5 @@ import createDebug from "debug"; import type WebSocket from "ws"; -import { v4 as randomUUID } from "uuid"; import * as msgpack from "@msgpack/msgpack"; import type { DataType, Task, Encoded, NodeID } from "@epfml/discojs"; @@ -23,21 +22,16 @@ const debug = createDebug("server:controllers:federated"); * * Four variables track participants at different stages. * - * - `#clientIds`: connection between an ID and an open WebSocket. + * - `clientIds` (inherited): connection between an ID and an open WebSocket. * It is only used to guarantee ids don't collide. * - `connections` (inherited): map between the ids of clients * that completed the connection handshake and their socket. - * This is the set of active participants in the current round. * - `#aggregator.nodes`: the group of nodes that the aggregator expects * contributions from. It decides when a round is completed * (`threshold * nodes.size`) and whether a contribution is valid. * - `#pendingUpdateRecipients`: the clients whose contribution was accepted for * the current round, reset after each aggregation. * It is a subset of `connections` that participated in the current round. - * - * Lifetime asymmetry: `reset()` clears `connections` and rebuilds the - * aggregator, but deliberately leaves `#clientIds` intact. Those sockets may - * still be open, and their ids must stay reserved to avoid collisions. */ export class FederatedController extends TrainingController< D, @@ -59,24 +53,6 @@ export class FederatedController extends TrainingController< * or staled participants */ #latestGlobalWeights: Encoded; - /** - * Complete list of client ids that have connected their websocket - * Make sure two clients don't share the same ID. - */ - #clientIds = new Set(); - /** - * Number of ClientConnected messages each client's socket sent. - * We answer them with the global weights, so a client resending it - * could make the server send them over and over. - * By limiting the number of times a client can send the ClientConnected message, we prevent some abuse. - * Note: It does not prevent a client from connecting with a different websocket and thus clientID. - */ - #clientCounters = new Map(); - - /** - * Maximum number of times a client can try to connect with the same client ID. - */ - static readonly MAX_CLIENT_CONNECTION_RETRY = 5; constructor( task: Task, @@ -128,7 +104,7 @@ export class FederatedController extends TrainingController< aggregator.round, recipientId.slice(0, 4), ); - this.#send(recipientWs, recipientId, encodedMsg); + this.send(recipientWs, recipientId, encodedMsg); }); } catch (err) { debug( @@ -157,16 +133,13 @@ export class FederatedController extends TrainingController< handle(ws: WebSocket): void { const minNbOfParticipants = this.task.trainingInformation.minNbOfParticipants; - // Try generating a new Client id until there no collision with existing ones - let clientId = randomUUID(); - while (this.#clientIds.has(clientId)) { - clientId = randomUUID(); - } - this.#clientIds.add(clientId); + + // Generate a new unique client ID for the participant + const clientId = this.newClientId(); const shortId = clientId.slice(0, 4); ws.on("error", (err) => { - this.#disconnect(ws, clientId, err); + this.disconnect(ws, clientId, err); }); // Setup callbacks triggered upon receiving the different client messages @@ -175,7 +148,7 @@ export class FederatedController extends TrainingController< const msg: unknown = msgpack.decode(data); if (!federatedMessages.isMessageToServer(msg)) { debug("invalid federated message received on WebSocket: %o", msg); - this.#crashClient(ws, clientId, "Invalid federated message"); + this.crashClient(ws, clientId, "Invalid federated message"); return; } @@ -191,11 +164,7 @@ export class FederatedController extends TrainingController< "Received message from an unconnected client [%s], sending CrashClient message", shortId, ); - this.#crashClient( - ws, - clientId, - "No ClientConnected message received", - ); + this.crashClient(ws, clientId, "No ClientConnected message received"); return; } @@ -217,21 +186,18 @@ export class FederatedController extends TrainingController< } else { debug(`New client connection for client [%s]`, shortId); // Connect the new client to both the connections map and the aggregator - this.#connectClient(clientId, ws); + this.connectClient(clientId, ws); } - // Increase the counter for this client ID - const count = (this.#clientCounters.get(clientId) ?? 0) + 1; - this.#clientCounters.set(clientId, count); - - // If the limit is reached for this client ID, we tell the client to crash - if (count > FederatedController.MAX_CLIENT_CONNECTION_RETRY) { + // Count the client's join attempt + // Crash the client if it exceeded the maximum number of connection attempts + if (!this.countJoinAttempt(clientId)) { debug( "Client [%s] exceeded the maximum number of connections (%d), sending CrashClient message", shortId, - FederatedController.MAX_CLIENT_CONNECTION_RETRY, + TrainingController.MAX_CLIENT_CONNECTION_RETRY, ); - this.#crashClient( + this.crashClient( ws, clientId, "Exceeded maximum number of retries to connect", @@ -252,7 +218,7 @@ export class FederatedController extends TrainingController< round: this.#aggregator.round, nbOfParticipants: this.connections.size, }; - this.#send(ws, clientId, msgpack.encode(msg)); + this.send(ws, clientId, msgpack.encode(msg)); // Send an update to participants if we can start/resume training, // which already carries the number of participants @@ -282,7 +248,7 @@ export class FederatedController extends TrainingController< this.#aggregator.round, ); // Send a notification to crash to the client - this.#crashClient( + this.crashClient( ws, clientId, `Received contribution for future round ${round} (current round=${this.#aggregator.round})`, @@ -329,7 +295,7 @@ export class FederatedController extends TrainingController< payload: this.#latestGlobalWeights, nbOfParticipants: this.connections.size, }; - this.#send(ws, clientId, msgpack.encode(msg)); + this.send(ws, clientId, msgpack.encode(msg)); } break; } @@ -337,14 +303,14 @@ export class FederatedController extends TrainingController< } catch (err) { // The client is in an unknown state, drop it rather than waiting for it debug("failed to handle message of client [%s]: %o", shortId, err); - this.#crashClient(ws, clientId, "Server failed to handle a message"); + this.crashClient(ws, clientId, "Server failed to handle a message"); } }); // Setup callback for client leaving the session ws.on("close", () => { // Remove the participant when the websocket is closed - this.#disconnectClient(clientId); + this.disconnectClient(clientId); debug("client [%s] left", shortId); @@ -371,13 +337,12 @@ export class FederatedController extends TrainingController< }); } - reset(): void { - this.resetConnectionState(); + protected override reset(): void { + super.reset(); + this.#pendingUpdateRecipients.clear(); - // Dispose first before generating a new aggregator - this.#aggregator.dispose(); + this.#aggregator.dispose(); // Dispose first before generating a new aggregator this.#aggregator = this.#makeAggregator(); - this.#clientCounters.clear(); this.#latestGlobalWeights = this.initialWeights; } @@ -386,8 +351,8 @@ export class FederatedController extends TrainingController< * Connects a new client to both the connections map and the aggregator. * Ensures consistency between the connections map and the aggregator. */ - #connectClient(clientId: string, ws: WebSocket): void { - this.connections = this.connections.set(clientId, ws); + protected override connectClient(clientId: string, ws: WebSocket): void { + super.connectClient(clientId, ws); this.#aggregator.registerNode(clientId); } @@ -395,51 +360,9 @@ export class FederatedController extends TrainingController< * Disconnects a client from the connections map, the aggregator, and the pending update recipients set. * Ensures consistency between the sets */ - #disconnectClient(clientId: string): void { - this.connections = this.connections.delete(clientId); + protected override disconnectClient(clientId: string): void { + super.disconnectClient(clientId); this.#aggregator.removeNode(clientId); this.#pendingUpdateRecipients.delete(clientId); - this.#clientIds.delete(clientId); - this.#clientCounters.delete(clientId); - } - - /** - * Sends a message to a client, disconnecting it if the message can't be delivered - */ - #send(ws: WebSocket, clientId: NodeID, data: Uint8Array): void { - try { - ws.send(data, (err) => { - if (err !== undefined && err !== null) - this.#disconnect(ws, clientId, err); - }); - } catch (err) { - this.#disconnect(ws, clientId, err); - } - } - - /** - * Tells a client to crash, then disconnects it once the message is sent. - * The server doesn't rely on the client to disconnect itself, which it may - * never do if it is faulty, while it keeps counting as a participant. - */ - #crashClient(ws: WebSocket, clientId: NodeID, reason: string): void { - debug("crashing client [%s]: %s", clientId.slice(0, 4), reason); - const msg: mtype.CrashClient = { type: MessageTypes.CrashClient, reason }; - try { - // disconnect whether the message was delivered or not - ws.send(msgpack.encode(msg), () => ws.terminate()); - } catch (err) { - this.#disconnect(ws, clientId, err); - } - } - - /** - * Forcefully disconnects a client. - * Its websocket "close" handler then removes it from the session, so that - * the remaining participants don't wait for it. - */ - #disconnect(ws: WebSocket, clientId: NodeID, reason: unknown): void { - debug("disconnecting client [%s]: %o", clientId.slice(0, 4), reason); - ws.terminate(); } } diff --git a/server/src/controllers/training_controller.ts b/server/src/controllers/training_controller.ts index fbd0f5ae3..a55aac9ab 100644 --- a/server/src/controllers/training_controller.ts +++ b/server/src/controllers/training_controller.ts @@ -2,6 +2,7 @@ import createDebug from "debug"; import type WebSocket from "ws"; import { Map } from "immutable"; import * as msgpack from "@msgpack/msgpack"; +import { v4 as randomUUID } from "uuid"; import { mtype } from "@epfml/discojs"; import type { DataType, Network, Task, NodeID } from "@epfml/discojs"; @@ -38,19 +39,39 @@ export abstract class TrainingController< * i.e. waiting for more participants or resuming training */ protected connections = Map(); + /** + * Complete list of client ids that have connected their websocket + * Make sure two clients don't share the same ID. + */ + protected clientIds = new Set(); + /** + * Number of ClientConnected messages each client's socket sent. + * We answer them with the global weights, so a client resending it + * could make the server send them over and over. + * By limiting the number of times a client can send the ClientConnected message, we prevent some abuse. + * Note: It does not prevent a client from connecting with a different websocket and thus clientID. + */ + protected clientCounters = Map(); + /** + * Maximum number of times a client can try to connect with the same client ID. + */ + static readonly MAX_CLIENT_CONNECTION_RETRY = 5; constructor(protected readonly task: Task) {} abstract handle(ws: WebSocket): void; - // Reset the controller state - abstract reset(): void; - - // Reset the peer connection state - // Used when the training is finished - protected resetConnectionState(): void { + /** + * Resets the state of the training controller. + * Clears the connections and client counters. + * Typically used when the training is finished or needs to be restarted. + * Deliberately leaves `clientIds` intact. Those sockets may + * still be open, and their ids must stay reserved to avoid future collisions. + */ + protected reset(): void { this.waitingForMoreParticipants = true; this.connections = Map(); + this.clientCounters = Map(); } /** @@ -75,7 +96,7 @@ export abstract class TrainingController< "Sending participants update to client [%s]", participantId.slice(0, 4), ); - participantWs.send(encoded); + this.send(participantWs, participantId, encoded); }); } @@ -106,7 +127,7 @@ export abstract class TrainingController< type: mtype.MType.EnoughParticipants, nbOfParticipants: this.connections.size, }; - participantWs.send(msgpack.encode(msg)); + this.send(participantWs, participantId, msgpack.encode(msg)); }); this.waitingForMoreParticipants = false; // update the attribute return true; @@ -131,7 +152,88 @@ export abstract class TrainingController< type: mtype.MType.WaitingForMoreParticipants, nbOfParticipants: this.connections.size, }; - participantWs.send(msgpack.encode(msg)); + this.send(participantWs, participantId, msgpack.encode(msg)); }); } + + /** + * Sends a message to a client, disconnecting it if the message can't be delivered + */ + protected send(ws: WebSocket, clientId: NodeID, data: Uint8Array): void { + try { + ws.send(data, (err) => { + if (err !== undefined && err !== null) + this.disconnect(ws, clientId, err); + }); + } catch (err) { + this.disconnect(ws, clientId, err); + } + } + + /** + * Tells a client to crash, then disconnects it once the message is sent. + * The server doesn't rely on the client to disconnect itself, which it may + * never do if it is faulty, while it keeps counting as a participant. + */ + protected crashClient(ws: WebSocket, clientId: NodeID, reason: string): void { + debug("crashing client [%s]: %s", clientId.slice(0, 4), reason); + const msg: mtype.CrashClient = { type: mtype.MType.CrashClient, reason }; + try { + // disconnect whether the message was delivered or not + ws.send(msgpack.encode(msg), () => ws.terminate()); + } catch (err) { + this.disconnect(ws, clientId, err); + } + } + + /** + * Forcefully disconnects a client. + * Its websocket "close" handler then removes it from the session, so that + * the remaining participants don't wait for it. + */ + protected disconnect(ws: WebSocket, clientId: NodeID, reason: unknown): void { + debug("disconnecting client [%s]: %o", clientId.slice(0, 4), reason); + ws.terminate(); + } + + /** + * Counts a client's join attempt and determines if it is allowed to try joining again based on the maximum retry limit. + * @param id the client ID for which to count the join attempt + * @returns `true` if the client is allowed to attempt joining again, `false` if the maximum retry limit has been reached + */ + protected countJoinAttempt(id: NodeID): boolean { + const currentCount = this.clientCounters.get(id) ?? 0; + if (currentCount >= TrainingController.MAX_CLIENT_CONNECTION_RETRY) { + return false; + } + this.clientCounters = this.clientCounters.set(id, currentCount + 1); + return true; + } + + protected newClientId(): NodeID { + let nodeId = randomUUID(); + while (this.clientIds.has(nodeId)) { + nodeId = randomUUID(); + } + this.clientIds.add(nodeId); + return nodeId; + } + + /** + * Adds a client that completed the join handshake to the session. + * Subclasses also add it to their own sets, calling super first. + */ + protected connectClient(id: NodeID, ws: WebSocket): void { + this.connections = this.connections.set(id, ws); + } + + /** + * Removes a client whose socket closed from the session and frees its id. + * Subclasses also remove it from their own sets, calling super. + */ + protected disconnectClient(id: NodeID): void { + this.connections = this.connections.delete(id); + this.clientCounters = this.clientCounters.delete(id); + this.clientIds.delete(id); + } } From f7bf9b5efc3473ccca458720ef71776bfe70da0c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?L=C3=A9onard=20Amsler?= Date: Wed, 14 Oct 2026 13:24:41 +0200 Subject: [PATCH 2/6] apply federated client fixes to the decentralized one --- discojs/src/client/client.ts | 9 +++++- .../decentralized/decentralized_client.ts | 30 ++++++++++++------- .../src/client/federated/federated_client.ts | 6 +--- 3 files changed, 29 insertions(+), 16 deletions(-) diff --git a/discojs/src/client/client.ts b/discojs/src/client/client.ts index d8b323232..c34557bb6 100644 --- a/discojs/src/client/client.ts +++ b/discojs/src/client/client.ts @@ -94,7 +94,14 @@ export abstract class Client extends EventEmitter<{ /** * Handles the disconnection process of the client from any sort of network server. */ - async disconnect(): Promise {} + async disconnect(): Promise { + // Optional server for the local client + // Makes it also safe on multiple calls + await this._server?.disconnect(); + this._server = undefined; + this._ownId = undefined; + this.promiseForMoreParticipants = undefined; + } /** * Emits the round status specified. It also stores the status emitted such that diff --git a/discojs/src/client/decentralized/decentralized_client.ts b/discojs/src/client/decentralized/decentralized_client.ts index 4baa9d6a3..cc5700e87 100644 --- a/discojs/src/client/decentralized/decentralized_client.ts +++ b/discojs/src/client/decentralized/decentralized_client.ts @@ -14,6 +14,7 @@ import { waitMessage, type PeerConnection, waitMessageWithTimeout, + sendAndWaitWithRetry, } from "#client/event_connection"; import { PeerPool } from "#client/decentralized/peer_pool"; import * as messages from "#client/decentralized/messages"; @@ -146,10 +147,15 @@ export class DecentralizedClient extends Client<"decentralized"> { const msg: ClientConnected = { type: MType.ClientConnected, }; - this.server.send(msg); - const { id, waitForMoreParticipants, nbOfParticipants, joinedMidTraining } = - await waitMessage(this.server, MType.NewDecentralizedNodeInfo); + await sendAndWaitWithRetry( + this.server, + msg, + MType.NewDecentralizedNodeInfo, + { + signal: this.crashSignal, + }, + ); this.#modelSyncNeeded = joinedMidTraining; this.setNbOfParticipantsUponJoining(nbOfParticipants); @@ -184,14 +190,13 @@ export class DecentralizedClient extends Client<"decentralized"> { const peers = this.#connections.keySeq().toSet(); this.aggregator.setNodes(this.aggregator.nodes.subtract(peers)); } - // Disconnect from server - await this.server?.disconnect(); - this._server = undefined; - this._ownId = undefined; this.#latestModel?.dispose(); this.#latestModel = undefined; + // At the end in case there is a throw while disconnecting the server + await super.disconnect(); + return Promise.resolve(); } @@ -206,6 +211,9 @@ export class DecentralizedClient extends Client<"decentralized"> { * model syncing happens first to get the latest global model from another peer. */ override async onRoundBeginCommunication(): Promise { + // Check if the crash signal has been triggered and throw if so + this.crashSignal.throwIfAborted(); + if (this.#modelSyncNeeded) { // 1. If model sync is needed, send server a request this.server.send({ type: MType.ModelSyncRequest }); @@ -287,9 +295,9 @@ export class DecentralizedClient extends Client<"decentralized"> { // Therefore, unsolved listeners may accumulate across rounds. // We can add listener resolving if this becomes a problem later. const msg = await Promise.race([ - waitMessage(this.server, MType.StartWeightSharing), - waitMessage(this.server, MType.RetryPeerConnections), - waitMessage(this.server, MType.ConnectionFail), + waitMessage(this.server, MType.StartWeightSharing, this.crashSignal), + waitMessage(this.server, MType.RetryPeerConnections, this.crashSignal), + waitMessage(this.server, MType.ConnectionFail, this.crashSignal), ]); if (msg.type === MType.StartWeightSharing) { @@ -353,6 +361,7 @@ export class DecentralizedClient extends Client<"decentralized"> { const receivedMessage = await waitMessage( this.server, MType.PeersForRound, + this.crashSignal, ); // every peer is ready to share weights, we can now connect to them this.saveAndEmit("connecting to peers"); @@ -398,6 +407,7 @@ export class DecentralizedClient extends Client<"decentralized"> { // still part of the session: leave the number of participants alone this.aggregator.setNodes(Set([this.ownId])); this.#connections = Map(); + this.crashSignal.throwIfAborted(); // Do not swallow if it is a signal error } } diff --git a/discojs/src/client/federated/federated_client.ts b/discojs/src/client/federated/federated_client.ts index 7d62c01c0..605e2ed54 100644 --- a/discojs/src/client/federated/federated_client.ts +++ b/discojs/src/client/federated/federated_client.ts @@ -112,12 +112,8 @@ export class FederatedClient extends Client<"federated"> { * Disconnection process when user quits the task. */ override async disconnect(): Promise { - await this.server.disconnect(); - this._server = undefined; - this._ownId = undefined; - this.promiseForMoreParticipants = undefined; - this.aggregator.setNodes(this.aggregator.nodes.delete(SERVER_NODE_ID)); + await super.disconnect(); } override onRoundBeginCommunication(): Promise { From bf5a4d373e37fc9bbd8780e65050178355ce4aca Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?L=C3=A9onard=20Amsler?= Date: Thu, 22 Oct 2026 18:15:30 -0300 Subject: [PATCH 3/6] decentralized controller unit tests --- .../decentralized_controller.spec.ts | 352 ++++++++++++++---- 1 file changed, 283 insertions(+), 69 deletions(-) diff --git a/server/src/controllers/decentralized_controller.spec.ts b/server/src/controllers/decentralized_controller.spec.ts index a16e4986f..fc8ff5364 100644 --- a/server/src/controllers/decentralized_controller.spec.ts +++ b/server/src/controllers/decentralized_controller.spec.ts @@ -4,12 +4,77 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { DecentralizedController } from "../../src/controllers/decentralized_controller.js"; import type { DecentralizedFakeWebSocket } from "../../tests/fake_websocket.js"; import { + expectLastMessageOfType, lastMessageOfType, makeDecentralizedFakeWebSocket, messagesOfType, } from "../../tests/fake_websocket.js"; import MessageTypes = mtype.MType; +import { TrainingController } from "./index.js"; + +async function makeController( + maxConnectionRetry: number, + minNbOfParticipants: number, +): Promise> { + const baseTask = await defaultTasks.cifar10.getTask(); + const task: Task<"image", "decentralized"> = { + ...baseTask, + trainingInformation: { + ...baseTask.trainingInformation, + scheme: "decentralized", + aggregationStrategy: "mean", + roundDuration: 1, + minNbOfParticipants, + maxConnectionRetry, + maxPeerConnectionTime: 60_000, + maxModelSyncTime: 30_000, + }, + }; + + return new DecentralizedController(task); +} + +function connect( + controller: DecentralizedController<"image">, + ws: DecentralizedFakeWebSocket, +): void { + controller.handle(ws); + ws.emitMessage({ type: MessageTypes.ClientConnected }); +} + +function connectAndJoinRound( + controller: DecentralizedController<"image">, + ws: DecentralizedFakeWebSocket, +): void { + connect(controller, ws); + + ws.emitMessage({ + type: MessageTypes.JoinRound, + }); + + ws.emitMessage({ + type: MessageTypes.PeerIsReady, + }); +} + +// runs a whole round (JoinRound, PeerIsReady, ConnectionsReady) +function runRound(peers: DecentralizedFakeWebSocket[]): void { + for (const ws of peers) { + ws.emitMessage({ type: MessageTypes.JoinRound }); + ws.emitMessage({ type: MessageTypes.PeerIsReady }); + } + for (const ws of peers) + ws.emitMessage({ type: MessageTypes.ConnectionsReady }); +} + +/** Checks that the server told a participant to crash, then disconnected it */ +async function expectCrashed(ws: DecentralizedFakeWebSocket): Promise { + expectLastMessageOfType(ws, MessageTypes.CrashClient); + await vi.waitFor(() => { + expect(ws.nbOfTerminations).toBe(1); + }); +} describe("DecentralizedController peer connection retry", () => { beforeEach(() => { @@ -20,46 +85,8 @@ describe("DecentralizedController peer connection retry", () => { vi.useRealTimers(); }); - async function makeController(maxConnectionRetry: number) { - const baseTask = await defaultTasks.cifar10.getTask(); - const task: Task<"image", "decentralized"> = { - ...baseTask, - trainingInformation: { - ...baseTask.trainingInformation, - scheme: "decentralized", - aggregationStrategy: "mean", - roundDuration: 1, - minNbOfParticipants: 2, - maxConnectionRetry, - maxPeerConnectionTime: 60_000, - maxModelSyncTime: 30_000, - }, - }; - - return new DecentralizedController(task); - } - - function connectAndJoinRound( - controller: DecentralizedController<"image">, - ws: DecentralizedFakeWebSocket, - ): void { - controller.handle(ws); - - ws.emitMessage({ - type: MessageTypes.ClientConnected, - }); - - ws.emitMessage({ - type: MessageTypes.JoinRound, - }); - - ws.emitMessage({ - type: MessageTypes.PeerIsReady, - }); - } - it("broadcasts RetryPeerConnections when not all peers finish connecting before timeout", async () => { - const controller = await makeController(3); + const controller = await makeController(3, 2); const ws1 = makeDecentralizedFakeWebSocket(); const ws2 = makeDecentralizedFakeWebSocket(); @@ -94,7 +121,7 @@ describe("DecentralizedController peer connection retry", () => { }); it("excludes failed peers only after maxConnectionRetry retries are exhausted", async () => { - const controller = await makeController(3); + const controller = await makeController(3, 2); const ws1 = makeDecentralizedFakeWebSocket(); const ws2 = makeDecentralizedFakeWebSocket(); @@ -153,33 +180,6 @@ describe("DecentralizedController peer connection retry", () => { }); describe("DecentralizedController participants updates", () => { - async function makeController(): Promise> { - const baseTask = await defaultTasks.cifar10.getTask(); - const task: Task<"image", "decentralized"> = { - ...baseTask, - trainingInformation: { - ...baseTask.trainingInformation, - scheme: "decentralized", - aggregationStrategy: "mean", - roundDuration: 1, - minNbOfParticipants: 2, - maxConnectionRetry: 3, - maxPeerConnectionTime: 60_000, - maxModelSyncTime: 30_000, - }, - }; - - return new DecentralizedController(task); - } - - function connect( - controller: DecentralizedController<"image">, - ws: DecentralizedFakeWebSocket, - ): void { - controller.handle(ws); - ws.emitMessage({ type: MessageTypes.ClientConnected }); - } - /** The counts a peer was told about, after it joined */ const participantsSeen = (ws: DecentralizedFakeWebSocket): number[] => messagesOfType(ws, MessageTypes.ParticipantsUpdate).map( @@ -187,7 +187,7 @@ describe("DecentralizedController participants updates", () => { ); it("counts a peer syncing its model like every other message does", async () => { - const controller = await makeController(); + const controller = await makeController(3, 2); const [ws1, ws2, ws3] = [ makeDecentralizedFakeWebSocket(), @@ -229,7 +229,7 @@ describe("DecentralizedController participants updates", () => { }); it("tells the peers when one joins or leaves", async () => { - const controller = await makeController(); + const controller = await makeController(3, 2); const [ws1, ws2, ws3] = [ makeDecentralizedFakeWebSocket(), @@ -270,3 +270,217 @@ describe("DecentralizedController participants updates", () => { ).to.equal(1); }); }); + +/** + * A peer resends ClientConnected when the server's answer is late or lost. + * The server has to answer again without treating it as a new participant. + */ +describe("Join handshake retry", () => { + it("sends the same NewDecentralizedNodeInfo again", async () => { + const controller = await makeController(3, 2); + const ws = makeDecentralizedFakeWebSocket(); + + connect(controller, ws); + + // send again a ClientConnected message + ws.emitMessage({ type: MessageTypes.ClientConnected }); + + // We should receive the same NewDecentralizedNodeInfo message again + const answers = messagesOfType(ws, MessageTypes.NewDecentralizedNodeInfo); + expect(answers).to.have.length(2); + expect(answers[1]).to.deep.equal(answers[0]); + expect(answers[1]).to.include({ + nbOfParticipants: 1, + waitForMoreParticipants: true, + joinedMidTraining: false, + }); + }); + + it("does not tell the other peers that someone joined again", async () => { + const controller = await makeController(3, 2); + const ws1 = makeDecentralizedFakeWebSocket(); + const ws2 = makeDecentralizedFakeWebSocket(); + + connect(controller, ws1); + connect(controller, ws2); + + // ws1 was told once that training can start, and never that someone joined + expect(messagesOfType(ws1, MessageTypes.EnoughParticipants)).to.have.length( + 1, + ); + expect(messagesOfType(ws1, MessageTypes.ParticipantsUpdate)).to.have.length( + 0, + ); + + // Resend a ClientConnected message from ws2 + ws2.emitMessage({ type: MessageTypes.ClientConnected }); + + // ws1 should not receive any new messages about ws2 joining again + expect(messagesOfType(ws1, MessageTypes.EnoughParticipants)).to.have.length( + 1, + ); + expect(messagesOfType(ws1, MessageTypes.ParticipantsUpdate)).to.have.length( + 0, + ); + + // ws2 should receive the NewDecentralizedNodeInfo message again + expect( + messagesOfType(ws2, MessageTypes.NewDecentralizedNodeInfo), + ).to.have.length(2); + }); + + it("tells a peer joining mid-training again that it has to sync its model", async () => { + const controller = await makeController(3, 2); + const [ws1, ws2, ws3] = [ + makeDecentralizedFakeWebSocket(), + makeDecentralizedFakeWebSocket(), + makeDecentralizedFakeWebSocket(), + ]; + + connect(controller, ws1); + connect(controller, ws2); + + // Run a training round with ws1 and ws2 before ws3 joins + runRound([ws1, ws2]); + + // Connect ws3 and resend its ClientConnected message + connect(controller, ws3); + ws3.emitMessage({ type: MessageTypes.ClientConnected }); + + // Check that ws3 receives two NewDecentralizedNodeInfo messages + const answers = messagesOfType(ws3, MessageTypes.NewDecentralizedNodeInfo); + expect(answers.map((answer) => answer.joinedMidTraining)).to.deep.equal([ + true, + true, + ]); + expect(answers.map((answer) => answer.nbOfParticipants)).to.deep.equal([ + 3, 3, + ]); + }); + + it("keeps a peer that joined before training in the next rounds when it resends", async () => { + const controller = await makeController(3, 2); + const [ws1, ws2, ws3] = [ + makeDecentralizedFakeWebSocket(), + makeDecentralizedFakeWebSocket(), + makeDecentralizedFakeWebSocket(), + ]; + + for (const ws of [ws1, ws2, ws3]) connect(controller, ws); + runRound([ws1, ws2, ws3]); + + ws1.emitMessage({ type: MessageTypes.ClientConnected }); + + expect( + expectLastMessageOfType(ws1, MessageTypes.NewDecentralizedNodeInfo) + .joinedMidTraining, + ).to.be.false; + + // the others can't start the next round without it + for (const ws of [ws2, ws3]) { + ws.emitMessage({ type: MessageTypes.JoinRound }); + ws.emitMessage({ type: MessageTypes.PeerIsReady }); + } + expect(messagesOfType(ws2, MessageTypes.PeersForRound)).to.have.length(1); + + ws1.emitMessage({ type: MessageTypes.JoinRound }); + ws1.emitMessage({ type: MessageTypes.PeerIsReady }); + + const peersForRound = expectLastMessageOfType( + ws2, + MessageTypes.PeersForRound, + ); + expect(peersForRound.aggregationRound).to.equal(1); + expect(peersForRound.peers).to.have.length(2); + }); +}); + +describe("Crashing a peer", () => { + it("because it sends another message before ClientConnected", async () => { + const controller = await makeController(3, 2); + const ws = makeDecentralizedFakeWebSocket(); + controller.handle(ws); + + // Send a message before ClientConnected + ws.emitMessage({ type: MessageTypes.JoinRound }); + + await expectCrashed(ws); + }); + + it("because it sends too many joins", async () => { + const controller = await makeController(3, 2); + const ws = makeDecentralizedFakeWebSocket(); + + connect(controller, ws); + + for (let i = 0; i < TrainingController.MAX_CLIENT_CONNECTION_RETRY; i++) { + ws.emitMessage({ type: MessageTypes.ClientConnected }); + } + + await expectCrashed(ws); + }); + + it("because it sends an invalid message", async () => { + const controller = await makeController(3, 2); + const ws = makeDecentralizedFakeWebSocket(); + connect(controller, ws); + + ws.emitMessage({ type: MessageTypes.ClientConnected }); + ws.emitMessage({ type: MessageTypes.JoinRound }); + ws.emitMessage({ type: MessageTypes.PeerIsReady }); + + // Send an invalid message + ws.emit("message", Buffer.from([0x00])); + await expectCrashed(ws); + }); + + it("because of a socket error, only disconnects the websocket", async () => { + const controller = await makeController(3, 2); + const ws = makeDecentralizedFakeWebSocket(); + connect(controller, ws); + + ws.emit("error", new Error("socket failure")); + + await vi.waitFor(() => { + expect(ws.nbOfTerminations).to.equal(1); + }); + expect(messagesOfType(ws, MessageTypes.CrashClient)).to.have.length(0); + }); + + it("does not block later rounds", async () => { + const controller = await makeController(3, 2); + const [ws1, ws2, ws3] = [ + makeDecentralizedFakeWebSocket(), + makeDecentralizedFakeWebSocket(), + makeDecentralizedFakeWebSocket(), + ]; + + for (const ws of [ws1, ws2, ws3]) connect(controller, ws); + + // Make ws1 join the round + ws1.emitMessage({ type: MessageTypes.JoinRound }); + ws1.emitMessage({ type: MessageTypes.PeerIsReady }); + + // Crash ws1 by sending an invalid message + ws1.emit("message", Buffer.from([0x00])); + await expectCrashed(ws1); + + // The others should learn it + await vi.waitFor(() => { + expect( + expectLastMessageOfType(ws2, MessageTypes.ParticipantsUpdate) + .nbOfParticipants, + ).to.equal(2); + }); + + // The others should be able to keep on training + for (const ws of [ws2, ws3]) { + ws.emitMessage({ type: MessageTypes.JoinRound }); + ws.emitMessage({ type: MessageTypes.PeerIsReady }); + } + // the round starts with the two peers left + expect( + expectLastMessageOfType(ws2, MessageTypes.PeersForRound).peers, + ).to.have.length(1); + }); +}); From 3cd5c9a48820f9e1357834281fcff6d496b2e147 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?L=C3=A9onard=20Amsler?= Date: Thu, 22 Oct 2026 18:45:48 -0300 Subject: [PATCH 4/6] Decentralized client test for disconnection on crash --- server/tests/client.spec.ts | 53 +++++++++++++++++++++++++++++++++++++ 1 file changed, 53 insertions(+) diff --git a/server/tests/client.spec.ts b/server/tests/client.spec.ts index aed196b33..279f791c3 100644 --- a/server/tests/client.spec.ts +++ b/server/tests/client.spec.ts @@ -18,9 +18,11 @@ import { mtype, defaultTasks, defaultModels, + ClientCrashError, } from "@epfml/discojs"; import { afterEach, describe, expect, it, vi } from "vitest"; import { Server } from "../src/index.js"; +import { TrainingController } from "../src/controllers/index.js"; describe("decentralized client", () => { let handle: http.Server; @@ -71,6 +73,57 @@ describe("decentralized client", () => { await expect(client.connect()).rejects.toThrow(); }); + + it("is disconnected by the server crashing it", async () => { + const url = await startServer( + [defaultModels.CIFAR10Classifier], + [defaultTasks.cifar10], + ); + + const client = new DecentralizedClient( + url, + await defaultTasks.cifar10.getTask(), + new MeanAggregator(), + ); + await client.connect(); + const connection = ( + client as unknown as { + _server: { + send: (msg: mtype.ClientConnected) => void; + socket: { readyState: number }; + }; + } + )._server; + + // Listen for the crash signal from the server + const crashed = new Promise((resolve) => + client.crashSignal.addEventListener("abort", resolve, { once: true }), + ); + + // Tries to connect the maximum number of times allowed by the server + for (let i = 0; i < TrainingController.MAX_CLIENT_CONNECTION_RETRY; i++) { + const msg: mtype.ClientConnected = { + type: mtype.MType.ClientConnected, + }; + connection.send(msg); + } + + // Make sure the crash occured and that the reason is correct + await crashed; + expect(client.crashSignal.reason).to.be.instanceOf(ClientCrashError); + expect(client.crashSignal.reason).to.have.property( + "reason", + "Exceeded maximum number of retries to connect", + ); + + // Make sure the connection is closed after the crash signal is received + await vi.waitFor(() => { + expect(connection.socket.readyState).to.equal(WebSocket.CLOSED); + }); + + // Make sure the client can still be disconnected gracefully + await client.disconnect(); + }); }); describe("federated client", () => { From 61666100d6677ff2e31a05b2732363cdde4199fe Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?L=C3=A9onard=20Amsler?= Date: Tue, 27 Oct 2026 08:31:11 -0300 Subject: [PATCH 5/6] Decentralized controller: reset connection retry count and manage timeouts --- server/src/controllers/decentralized_controller.ts | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/server/src/controllers/decentralized_controller.ts b/server/src/controllers/decentralized_controller.ts index 8a4cb9893..b7ed449b2 100644 --- a/server/src/controllers/decentralized_controller.ts +++ b/server/src/controllers/decentralized_controller.ts @@ -379,6 +379,8 @@ export class DecentralizedController< // Stop the timeout this.#clearTimeout(); + // Reset the connection retry count + this.#connectionRetry = 0; // Send round participants StartWeightSharing messages this.#broadcastMessageToNodes( @@ -404,6 +406,9 @@ export class DecentralizedController< #startTimeout( maxTime: number = this.task.trainingInformation.maxPeerConnectionTime, ): void { + // Cancel the previous timeout if it exists + this.#clearTimeout(); + this.#timeout = setTimeout(() => { this.#handleTimeout(); }, maxTime); @@ -472,8 +477,7 @@ export class DecentralizedController< // Reset the connectionRetry since we excluded the failed clients this.#connectionRetry = 0; - // Restart the timeout after sending retry messages - this.#startTimeout(); + return; } @@ -492,9 +496,6 @@ export class DecentralizedController< }; }, ); - - // Restart the timeout after sending retry messages - this.#startTimeout(); } protected override connectClient(nodeId: string, ws: WebSocket): void { From feacc3cb128de15bc75ae8bbb2fcf02bad8eec77 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?L=C3=A9onard=20Amsler?= Date: Sun, 1 Nov 2026 18:37:12 -0300 Subject: [PATCH 6/6] Apply code review changes --- .../decentralized/decentralized_client.ts | 3 +++ discojs/src/client/event_connection.ts | 3 ++- .../controllers/decentralized_controller.ts | 19 ++++++++++--------- server/src/controllers/training_controller.ts | 6 +++--- 4 files changed, 18 insertions(+), 13 deletions(-) diff --git a/discojs/src/client/decentralized/decentralized_client.ts b/discojs/src/client/decentralized/decentralized_client.ts index 269522447..30a2bb4cf 100644 --- a/discojs/src/client/decentralized/decentralized_client.ts +++ b/discojs/src/client/decentralized/decentralized_client.ts @@ -222,6 +222,7 @@ export class DecentralizedClient extends Client<"decentralized"> { const providerInfo = await waitMessageWithTimeout( this.server, MType.ModelProviderInfo, + this.crashSignal, this.task.trainingInformation.maxModelSyncTime, "Timeout while waiting for the latest model provider", ); @@ -427,6 +428,7 @@ export class DecentralizedClient extends Client<"decentralized"> { const message = await waitMessageWithTimeout( connection, MType.Payload, + this.crashSignal, MAX_PAYLOAD_WAIT, "Timeout waiting for a contribution from peer " + peerId, ); @@ -607,6 +609,7 @@ export class DecentralizedClient extends Client<"decentralized"> { const message = await waitMessageWithTimeout( providerConn, MType.SharedModel, + this.crashSignal, this.task.trainingInformation.maxModelSyncTime, "Timeout while waiting for the latest model", ); diff --git a/discojs/src/client/event_connection.ts b/discojs/src/client/event_connection.ts index d5ddd787d..c9bcbf8a5 100644 --- a/discojs/src/client/event_connection.ts +++ b/discojs/src/client/event_connection.ts @@ -65,11 +65,12 @@ export async function waitMessage( export async function waitMessageWithTimeout( connection: EventConnection, type: T, + signal: AbortSignal, timeoutMs?: number, errorMsg: string = "timeout", ): Promise> { return await Promise.race([ - waitMessage(connection, type), + waitMessage(connection, type, signal), timeout(timeoutMs, errorMsg), ]); } diff --git a/server/src/controllers/decentralized_controller.ts b/server/src/controllers/decentralized_controller.ts index b7ed449b2..ff55e3178 100644 --- a/server/src/controllers/decentralized_controller.ts +++ b/server/src/controllers/decentralized_controller.ts @@ -23,7 +23,8 @@ const debug = createDebug("server:controllers:decentralized"); * - `connections` (inherited): map between the ids of clients * that completed the connection handshake and their socket. * - `#roundPeers`: the nodes that want to join the current round and their readiness status. - * - `#connectFinishedNodes`: the nodes that have completed the connection handshake and are ready for the current round. + * - `#connectFinishedNodes`: the nodes that have completed their peer-to-peer webRTC connections + * and are ready for the current round. * - `#syncingNodes`: the nodes that are currently syncing with the latest model. * - `#providerNode`: the client selected to provide the latest model to peers joining in the middle of training. * @@ -291,12 +292,7 @@ export class DecentralizedController< this.#connectionRetry = 0; this.#providerNode = undefined; this.#syncingNodes = Set(); - - // Reset the timeout - if (this.#timeout !== undefined) { - clearTimeout(this.#timeout); - this.#timeout = undefined; - } + this.#clearTimeout(); } #sendMessageToNode( @@ -407,6 +403,7 @@ export class DecentralizedController< maxTime: number = this.task.trainingInformation.maxPeerConnectionTime, ): void { // Cancel the previous timeout if it exists + // A previous timeout exists if we are retrying to establish peer connections this.#clearTimeout(); this.#timeout = setTimeout(() => { @@ -478,6 +475,10 @@ export class DecentralizedController< // Reset the connectionRetry since we excluded the failed clients this.#connectionRetry = 0; + // Start the timeout for the next connection attempt + // Avoids getting stuck waiting for enough PeerIsReady messages + this.#startTimeout(); + return; } @@ -498,7 +499,7 @@ export class DecentralizedController< ); } - protected override connectClient(nodeId: string, ws: WebSocket): void { + protected override connectClient(nodeId: NodeID, ws: WebSocket): void { super.connectClient(nodeId, ws); // If the new peer joins in the middle of training, @@ -508,7 +509,7 @@ export class DecentralizedController< } } - protected override disconnectClient(nodeId: string): void { + protected override disconnectClient(nodeId: NodeID): void { super.disconnectClient(nodeId); this.#roundPeers = this.#roundPeers.delete(nodeId); diff --git a/server/src/controllers/training_controller.ts b/server/src/controllers/training_controller.ts index a55aac9ab..50f05b5e4 100644 --- a/server/src/controllers/training_controller.ts +++ b/server/src/controllers/training_controller.ts @@ -1,6 +1,6 @@ import createDebug from "debug"; import type WebSocket from "ws"; -import { Map } from "immutable"; +import { Map, Set } from "immutable"; import * as msgpack from "@msgpack/msgpack"; import { v4 as randomUUID } from "uuid"; @@ -43,10 +43,10 @@ export abstract class TrainingController< * Complete list of client ids that have connected their websocket * Make sure two clients don't share the same ID. */ - protected clientIds = new Set(); + protected clientIds = Set(); /** * Number of ClientConnected messages each client's socket sent. - * We answer them with the global weights, so a client resending it + * In FL, we answer them with the global weights, so a client resending it * could make the server send them over and over. * By limiting the number of times a client can send the ClientConnected message, we prevent some abuse. * Note: It does not prevent a client from connecting with a different websocket and thus clientID.