From 8a69883d28435080559a3a24e20a6d5ed20f4a09 Mon Sep 17 00:00:00 2001 From: Julien Vignoud Date: Thu, 10 Sep 2026 17:35:09 +0200 Subject: [PATCH 1/5] fix: reset participants and hide participant counter when training locally --- .../components/training/TrainerDashboard.vue | 13 +++++-- .../training/TrainingInformation.vue | 2 ++ .../__tests__/TrainerDashboard.spec.ts | 34 +++++++++++++++++++ 3 files changed, 47 insertions(+), 2 deletions(-) diff --git a/webapp/src/components/training/TrainerDashboard.vue b/webapp/src/components/training/TrainerDashboard.vue index 8daa208c5..520c687fb 100644 --- a/webapp/src/components/training/TrainerDashboard.vue +++ b/webapp/src/components/training/TrainerDashboard.vue @@ -98,7 +98,7 @@ :batches-of-epoch="batchesOfEpochLogs" :has-validation-data="hasValidationData" :is-training="isTraining" - :is-training-alone="isTrainingAlone" + :is-training-alone="isTrainingLocally" :nb-participants="nbParticipants" /> @@ -167,6 +167,12 @@ const hasValidationData = computed( const isTraining = computed(() => trainingGenerator.value !== undefined); const isTrainingAlone = ref(false); +// Whether the training involves no collaborator at all: either the user chose +// to train alone or the task itself is a local one +const isTrainingLocally = computed( + () => + isTrainingAlone.value || props.task.trainingInformation.scheme === "local", +); // number of participants in the training session const nbParticipants = ref(1); @@ -180,6 +186,9 @@ async function startTraining(): Promise { roundsLogs.value = List(); epochsOfRoundLogs.value = List(); batchesOfEpochLogs.value = List(); + // the client tells us how many participants there are once connected, + // until then we only know about ourselves + nbParticipants.value = 1; // Vue proxy doesn't work with Dataset's private fields const dataset = toRaw(props.dataset); @@ -192,7 +201,7 @@ async function startTraining(): Promise { console.log("server URL:", CONFIG.serverUrl.toString()); const disco = new Disco(props.task, CONFIG.serverUrl, { - scheme: isTrainingAlone.value + scheme: isTrainingLocally.value ? "local" : props.task.trainingInformation.scheme, }); diff --git a/webapp/src/components/training/TrainingInformation.vue b/webapp/src/components/training/TrainingInformation.vue index e7b28426f..b0b3e009e 100644 --- a/webapp/src/components/training/TrainingInformation.vue +++ b/webapp/src/components/training/TrainingInformation.vue @@ -41,7 +41,9 @@ + { infos.props("rounds").last()?.epochs.last()?.training.accuracy, ).toBeGreaterThan(0); }); + +it("hides the participants when training alone", async () => { + const wrapper = await setupForTask(); + + expect(wrapper.text()).toContain("number of participants"); + + // nobody to collaborate with, the count would be stuck at one + await wrapper.get("#train-locally-bttn").trigger("click"); + + expect(wrapper.text()).not.toContain("number of participants"); +}); + +it( + "resets the participants when a training starts", + { timeout: 20_000 }, + async () => { + const wrapper = await setupForTask(); + const infos = wrapper.getComponent(TrainingInformation); + + // stand for a previous session having left a count behind, which only a + // connected client can otherwise produce + (wrapper.vm as unknown as { nbParticipants: number }).nbParticipants = 4; + await wrapper.vm.$nextTick(); + expect(infos.props("nbParticipants")).to.equal(4); + + // train alone so that the round completes without a server to connect to + await wrapper.get("#train-locally-bttn").trigger("click"); + await wrapper.get("#start-training-bttn").trigger("click"); + while (infos.props("rounds").isEmpty()) + await new Promise((resolve) => setTimeout(resolve, 100)); + + expect(infos.props("nbParticipants")).to.equal(1); + }, +); From b9cccf9c9e9dc5abd920f61ea9faff58ad3b53c3 Mon Sep 17 00:00:00 2001 From: Julien Vignoud Date: Fri, 11 Sep 2026 12:12:56 +0200 Subject: [PATCH 2/5] fix: handle message inversion overwriting participant count --- discojs/src/client/client.ts | 23 ++++ .../decentralized/decentralized_client.ts | 2 +- .../src/client/federated/federated_client.ts | 2 +- server/tests/client.spec.ts | 126 +++++++++++++++++- 4 files changed, 150 insertions(+), 3 deletions(-) diff --git a/discojs/src/client/client.ts b/discojs/src/client/client.ts index 7333ad363..3184519d1 100644 --- a/discojs/src/client/client.ts +++ b/discojs/src/client/client.ts @@ -49,6 +49,11 @@ export abstract class Client extends EventEmitter<{ // Current number of participants including this client in the training session #nbOfParticipants: number = 1; + /** + * Whether the server sent us a more recent number of participants than the + * one carried by the message answering our join request. + */ + #nbOfParticipantsUpdatedSinceJoining = false; constructor( public readonly url: URL, // The network server's URL to connect to @@ -119,6 +124,8 @@ export abstract class Client extends EventEmitter<{ * between a NewNodeInfo message and an EnoughParticipant message. */ protected setupServerCallbacks(setMessageInversionFlag: () => void) { + this.#nbOfParticipantsUpdatedSinceJoining = false; + // Setup an event callback if the server signals that we should // wait for more participants this.server.on(MType.WaitingForMoreParticipants, (event) => { @@ -131,6 +138,7 @@ export abstract class Client extends EventEmitter<{ ); // Display the waiting status right away this.emit("status", "not enough participants"); + this.#nbOfParticipantsUpdatedSinceJoining = true; this.nbOfParticipants = event.nbOfParticipants; // emits the `participants` event // Upon receiving a WaitingForMoreParticipants message, // the client will await for this promise to resolve before sending its @@ -147,10 +155,25 @@ export abstract class Client extends EventEmitter<{ this.server.once(MType.EnoughParticipants, (event) => { if (this._ownId === undefined) { setMessageInversionFlag(); + this.#nbOfParticipantsUpdatedSinceJoining = true; this.nbOfParticipants = event.nbOfParticipants; } }); } + + /** + * Set the number of participants carried by the message answering our join + * request. + * + * The server computes that count when sending its answer, which we can + * process much later as the message also carries the model weights. Any + * WaitingForMoreParticipants or EnoughParticipants message processed in the + * meantime carries a more recent count, which we keep instead. + */ + protected setNbOfParticipantsUponJoining(nbOfParticipants: number): void { + if (this.#nbOfParticipantsUpdatedSinceJoining) return; + this.nbOfParticipants = nbOfParticipants; + } /** * Method called when the server notifies the client that there aren't enough * participants (anymore) to start/continue training diff --git a/discojs/src/client/decentralized/decentralized_client.ts b/discojs/src/client/decentralized/decentralized_client.ts index 64428c9fd..d9a0432f6 100644 --- a/discojs/src/client/decentralized/decentralized_client.ts +++ b/discojs/src/client/decentralized/decentralized_client.ts @@ -156,7 +156,7 @@ export class DecentralizedClient extends Client<"decentralized"> { await waitMessage(this.server, MType.NewDecentralizedNodeInfo); this.#modelSyncNeeded = joinedMidTraining; - this.nbOfParticipants = nbOfParticipants; + this.setNbOfParticipantsUponJoining(nbOfParticipants); // This should come right after receiving the message to make sure // we don't miss a subsequent message from the server diff --git a/discojs/src/client/federated/federated_client.ts b/discojs/src/client/federated/federated_client.ts index 179e40da6..dca53bac6 100644 --- a/discojs/src/client/federated/federated_client.ts +++ b/discojs/src/client/federated/federated_client.ts @@ -85,7 +85,7 @@ export class FederatedClient extends Client<"federated"> { this._ownId = id; debug(`[${shortenId(id)}] joined session at round ${round} `); this.aggregator.setRound(round); - this.nbOfParticipants = nbOfParticipants; + this.setNbOfParticipantsUponJoining(nbOfParticipants); // Upon connecting, the server answers with a boolean // which indicates whether there are enough participants or not debug( diff --git a/server/tests/client.spec.ts b/server/tests/client.spec.ts index 4568aad8e..bb91ead72 100644 --- a/server/tests/client.spec.ts +++ b/server/tests/client.spec.ts @@ -1,14 +1,20 @@ -import type * as http from "node:http"; +import * as http from "node:http"; +import * as msgpack from "@msgpack/msgpack"; +import { WebSocketServer } from "ws"; import type { DataType, + Model, Network, TaskProvider, ModelCard, + decentralizedMessages, + federatedMessages, } from "@epfml/discojs"; import { MeanAggregator, DecentralizedClient, FederatedClient, + mtype, defaultTasks, defaultModels, } from "@epfml/discojs"; @@ -116,3 +122,121 @@ describe("federated client", () => { await expect(client.connect()).rejects.toThrow(); }); }); + +/** + * A client learns how many participants there are from the message answering + * its join request, but the server computes that count when sending it. As that + * message also carries the model weights it can be big and slow, so a smaller + * EnoughParticipants sent right after can be processed first, carrying a more + * recent count which the join answer must not overwrite. + * + * Reproducing that ordering needs both messages to be sent back to back, which + * the real server doesn't do, hence the stubbed one below. + */ +describe("client joining while the participants change", () => { + type JoinAnswer = + | federatedMessages.MessageFederated + | decentralizedMessages.MessageFromServer; + + /** Have the client skip fetching the base model, only messages matter here */ + function withoutModel< + C extends { getLatestModel: () => Promise> }, + >(client: C): C { + client.getLatestModel = () => Promise.resolve({} as Model); + return client; + } + + /** Answer a join request with the given messages, in the order given */ + async function serveJoinAnswer( + ...answers: readonly JoinAnswer[] + ): Promise<[http.Server, URL]> { + const handle = http.createServer(); + new WebSocketServer({ server: handle }).on("connection", (ws) => + ws.on("message", (data: Buffer) => { + const msg: unknown = msgpack.decode(data); + if ( + !mtype.hasMessageType(msg) || + msg.type !== mtype.MType.ClientConnected + ) + return; + for (const answer of answers) ws.send(msgpack.encode(answer)); + }), + ); + + const url = await new Promise((resolve) => + handle.listen(0, "127.0.0.1", () => { + const address = handle.address(); + if (address === null || typeof address === "string") + throw new Error("server didn't listen on a port"); + resolve(new URL(`http://127.0.0.1:${address.port}/`)); + }), + ); + + return [handle, url]; + } + + // sent when another participant joined between the two messages + const enoughParticipants: mtype.EnoughParticipants = { + type: mtype.MType.EnoughParticipants, + nbOfParticipants: 2, + }; + + it("keeps the newer count when federated", async () => { + const joinAnswer: federatedMessages.NewFederatedNodeInfo = { + type: mtype.MType.NewFederatedNodeInfo, + id: "node-id", + waitForMoreParticipants: true, + payload: undefined, + round: 0, + nbOfParticipants: 1, + }; + const [handle, url] = await serveJoinAnswer(joinAnswer, enoughParticipants); + + const client = withoutModel( + new FederatedClient( + url, + await defaultTasks.titanic.getTask(), + new MeanAggregator(), + ), + ); + + try { + await client.connect(); + + expect(client.nbOfParticipants).to.equal(2); + expect(client.waitingForMoreParticipants).to.be.false; + } finally { + await client.disconnect(); + handle.close(); + } + }); + + it("keeps the newer count when decentralized", async () => { + const joinAnswer: decentralizedMessages.NewDecentralizedNodeInfo = { + type: mtype.MType.NewDecentralizedNodeInfo, + id: "node-id", + waitForMoreParticipants: true, + joinedMidTraining: false, + nbOfParticipants: 1, + }; + const [handle, url] = await serveJoinAnswer(joinAnswer, enoughParticipants); + + const client = withoutModel( + new DecentralizedClient( + url, + await defaultTasks.cifar10.getTask(), + new MeanAggregator(), + ), + ); + + try { + await client.connect(); + + expect(client.nbOfParticipants).to.equal(2); + expect(client.waitingForMoreParticipants).to.be.false; + } finally { + await client.disconnect(); + handle.close(); + } + }); +}); From 0eb64dc4e61e385a140bd8ebaf1727d4e28a7325 Mon Sep 17 00:00:00 2001 From: Julien Vignoud Date: Fri, 11 Sep 2026 13:17:35 +0200 Subject: [PATCH 3/5] fix: decouple nodes from number of participant --- .../decentralized/decentralized_client.ts | 23 ++- server/tests/client.spec.ts | 173 +++++++++++++----- 2 files changed, 141 insertions(+), 55 deletions(-) diff --git a/discojs/src/client/decentralized/decentralized_client.ts b/discojs/src/client/decentralized/decentralized_client.ts index d9a0432f6..2028a9555 100644 --- a/discojs/src/client/decentralized/decentralized_client.ts +++ b/discojs/src/client/decentralized/decentralized_client.ts @@ -54,13 +54,6 @@ export class DecentralizedClient extends Client<"decentralized"> { return this._server === undefined; } - private setAggregatorNodes(nodes: Set) { - this.aggregator.setNodes(nodes); - // Emits the `participants` event - this.nbOfParticipants = - this.aggregator.nodes.size === 0 ? 1 : this.aggregator.nodes.size; - } - private cloneWeights(weights: WeightsContainer): WeightsContainer { return new WeightsContainer(weights.weights.map((t) => t.clone())); } @@ -186,7 +179,7 @@ export class DecentralizedClient extends Client<"decentralized"> { if (this.#connections !== undefined) { const peers = this.#connections.keySeq().toSet(); - this.setAggregatorNodes(this.aggregator.nodes.subtract(peers)); + this.aggregator.setNodes(this.aggregator.nodes.subtract(peers)); } // Disconnect from server await this.server?.disconnect(); @@ -311,9 +304,10 @@ export class DecentralizedClient extends Client<"decentralized"> { // clear the communication round peer pool await this.#pool?.shutdown(); this.#pool = new PeerPool(this.ownId); - // clear the connections + // clear the connections. The peers are still part of the session, + // so the number of participants is left as the server reported it this.#connections = Map(); - this.setAggregatorNodes(Set(this.ownId)); + this.aggregator.setNodes(Set([this.ownId])); continue; } else if (msg.type === MType.ConnectionFail) { debug(`[${shortenId(this.ownId)}] disconnect from the server`); @@ -369,7 +363,10 @@ export class DecentralizedClient extends Client<"decentralized"> { throw new Error("received peer list contains our own id"); } // Store the list of peers for the current round including ourselves - this.setAggregatorNodes(peers.add(this.ownId)); + const roundNodes = peers.add(this.ownId); + this.aggregator.setNodes(roundNodes); + // the server told us who takes part in the round, emits `participants` + this.nbOfParticipants = roundNodes.size; this.aggregator.setRound(receivedMessage.aggregationRound); // the server gives us the round number // Initiate peer to peer connections with each peer @@ -394,7 +391,9 @@ export class DecentralizedClient extends Client<"decentralized"> { `Error for [${shortenId(this.ownId)}] while beginning round: %o`, e, ); - this.setAggregatorNodes(Set(this.ownId)); + // we can't aggregate with peers we failed to connect to, but they are + // still part of the session: leave the number of participants alone + this.aggregator.setNodes(Set([this.ownId])); this.#connections = Map(); } } diff --git a/server/tests/client.spec.ts b/server/tests/client.spec.ts index bb91ead72..ed3b0c489 100644 --- a/server/tests/client.spec.ts +++ b/server/tests/client.spec.ts @@ -14,6 +14,7 @@ import { MeanAggregator, DecentralizedClient, FederatedClient, + WeightsContainer, mtype, defaultTasks, defaultModels, @@ -123,6 +124,49 @@ describe("federated client", () => { }); }); +type ServerMessage = + | federatedMessages.MessageFederated + | decentralizedMessages.MessageFromServer; + +/** Have a client skip fetching the base model, only messages matter here */ +function withoutModel< + C extends { getLatestModel: () => Promise> }, +>(client: C): C { + client.getLatestModel = () => Promise.resolve({} as Model); + return client; +} + +/** + * A server answering a client's messages with whatever `answer` sends, so that + * a test can produce orderings the real server doesn't. + */ +async function serveStub( + answer: ( + msg: { type: mtype.MType }, + send: (msg: ServerMessage) => void, + ) => void, +): Promise<[http.Server, URL]> { + const handle = http.createServer(); + new WebSocketServer({ server: handle }).on("connection", (ws) => + ws.on("message", (data: Buffer) => { + const msg: unknown = msgpack.decode(data); + if (!mtype.hasMessageType(msg)) return; + answer(msg, (answer) => ws.send(msgpack.encode(answer))); + }), + ); + + const url = await new Promise((resolve) => + handle.listen(0, "127.0.0.1", () => { + const address = handle.address(); + if (address === null || typeof address === "string") + throw new Error("server didn't listen on a port"); + resolve(new URL(`http://127.0.0.1:${address.port}/`)); + }), + ); + + return [handle, url]; +} + /** * A client learns how many participants there are from the message answering * its join request, but the server computes that count when sending it. As that @@ -134,47 +178,6 @@ describe("federated client", () => { * the real server doesn't do, hence the stubbed one below. */ describe("client joining while the participants change", () => { - type JoinAnswer = - | federatedMessages.MessageFederated - | decentralizedMessages.MessageFromServer; - - /** Have the client skip fetching the base model, only messages matter here */ - function withoutModel< - C extends { getLatestModel: () => Promise> }, - >(client: C): C { - client.getLatestModel = () => Promise.resolve({} as Model); - return client; - } - - /** Answer a join request with the given messages, in the order given */ - async function serveJoinAnswer( - ...answers: readonly JoinAnswer[] - ): Promise<[http.Server, URL]> { - const handle = http.createServer(); - new WebSocketServer({ server: handle }).on("connection", (ws) => - ws.on("message", (data: Buffer) => { - const msg: unknown = msgpack.decode(data); - if ( - !mtype.hasMessageType(msg) || - msg.type !== mtype.MType.ClientConnected - ) - return; - for (const answer of answers) ws.send(msgpack.encode(answer)); - }), - ); - - const url = await new Promise((resolve) => - handle.listen(0, "127.0.0.1", () => { - const address = handle.address(); - if (address === null || typeof address === "string") - throw new Error("server didn't listen on a port"); - resolve(new URL(`http://127.0.0.1:${address.port}/`)); - }), - ); - - return [handle, url]; - } - // sent when another participant joined between the two messages const enoughParticipants: mtype.EnoughParticipants = { type: mtype.MType.EnoughParticipants, @@ -190,7 +193,11 @@ describe("client joining while the participants change", () => { round: 0, nbOfParticipants: 1, }; - const [handle, url] = await serveJoinAnswer(joinAnswer, enoughParticipants); + const [handle, url] = await serveStub((msg, send) => { + if (msg.type !== mtype.MType.ClientConnected) return; + send(joinAnswer); + send(enoughParticipants); + }); const client = withoutModel( new FederatedClient( @@ -219,7 +226,11 @@ describe("client joining while the participants change", () => { joinedMidTraining: false, nbOfParticipants: 1, }; - const [handle, url] = await serveJoinAnswer(joinAnswer, enoughParticipants); + const [handle, url] = await serveStub((msg, send) => { + if (msg.type !== mtype.MType.ClientConnected) return; + send(joinAnswer); + send(enoughParticipants); + }); const client = withoutModel( new DecentralizedClient( @@ -240,3 +251,79 @@ describe("client joining while the participants change", () => { } }); }); + +/** + * Failing to connect to the round's peers leaves the aggregator with nobody to + * expect a contribution from, but the peers are still part of the session: the + * server is the one telling the client how many participants there are, so a + * failed round start must not have the client report being alone. + */ +describe("peer failing to begin a round", () => { + it("keeps reporting the participants the server gave", async () => { + const ownId = "node-id"; + + const joinAnswer: decentralizedMessages.NewDecentralizedNodeInfo = { + type: mtype.MType.NewDecentralizedNodeInfo, + id: ownId, + waitForMoreParticipants: false, + joinedMidTraining: false, + nbOfParticipants: 2, + }; + // a peer list containing our own id makes the client fail to begin the + // round, as connecting to the peers of the round would + const badPeersForRound: decentralizedMessages.PeersForRound = { + type: mtype.MType.PeersForRound, + peers: [ownId], + aggregationRound: 0, + }; + + let rounds = 0; + const [handle, url] = await serveStub((msg, send) => { + switch (msg.type) { + case mtype.MType.ClientConnected: + send(joinAnswer); + break; + case mtype.MType.PeerIsReady: + send(badPeersForRound); + // let the client handle the failure and listen again before telling + // it to retry, then to give up + rounds++; + setTimeout(() => + send( + rounds === 1 + ? { type: mtype.MType.RetryPeerConnections } + : { type: mtype.MType.ConnectionFail }, + ), + ); + break; + } + }); + + const client = withoutModel( + new DecentralizedClient( + url, + await defaultTasks.cifar10.getTask(), + new MeanAggregator(), + ), + ); + const participants: number[] = []; + client.on("participants", (nbOfParticipants) => + participants.push(nbOfParticipants), + ); + + try { + await client.connect(); + await client.onRoundBeginCommunication(); + + await expect( + client.onRoundEndCommunication(new WeightsContainer([[1]])), + ).rejects.toThrow("Client disconnected after connection failure"); + + // the server only ever said there were two of us + expect(participants).to.deep.equal([2]); + expect(rounds).to.equal(2); // the retry did happen + } finally { + handle.close(); + } + }); +}); From cf0636880c12e0ba63e67fbb3344b99ea6f56109 Mon Sep 17 00:00:00 2001 From: Julien Vignoud Date: Fri, 11 Sep 2026 15:08:01 +0200 Subject: [PATCH 4/5] feat: server updates clients with number of participants upon join & leave --- discojs/src/client/client.ts | 8 ++ discojs/src/client/decentralized/messages.ts | 3 + discojs/src/client/federated/messages.ts | 5 +- discojs/src/client/mtype.ts | 11 ++ .../controllers/decentralized_controller.ts | 17 ++- .../src/controllers/federated_controller.ts | 15 ++- server/src/controllers/training_controller.ts | 33 +++++- server/tests/client.spec.ts | 111 ++++++++++-------- .../e2e/decentralized_controller.spec.ts | 77 ++++++++++++ server/tests/e2e/federated.spec.ts | 52 ++++++++ 10 files changed, 276 insertions(+), 56 deletions(-) diff --git a/discojs/src/client/client.ts b/discojs/src/client/client.ts index 3184519d1..ae825750b 100644 --- a/discojs/src/client/client.ts +++ b/discojs/src/client/client.ts @@ -146,6 +146,14 @@ export abstract class Client extends EventEmitter<{ this.promiseForMoreParticipants = this.createPromiseForMoreParticipants(); }); + // The server tells us whenever a participant joined or left, so that we + // don't have to wait for the end of the round to display how many of us + // are training together + this.server.on(MType.ParticipantsUpdate, (event) => { + this.#nbOfParticipantsUpdatedSinceJoining = true; + this.nbOfParticipants = event.nbOfParticipants; + }); + // As an example assume we need at least 2 participants to train, // When two participants join almost at the same time, the server // sends a NewNodeInfo with waitForMoreParticipants=true to the first participant diff --git a/discojs/src/client/decentralized/messages.ts b/discojs/src/client/decentralized/messages.ts index 1c2b334fb..309dd7b71 100644 --- a/discojs/src/client/decentralized/messages.ts +++ b/discojs/src/client/decentralized/messages.ts @@ -7,6 +7,7 @@ import type { ClientConnected, WaitingForMoreParticipants, EnoughParticipants, + ParticipantsUpdate, } from "#client/mtype"; /// Phase 0 communication (between server and peers) @@ -107,6 +108,7 @@ export type MessageFromServer = | PeersForRound | WaitingForMoreParticipants | EnoughParticipants + | ParticipantsUpdate | StartWeightSharing | RetryPeerConnections | ConnectionFail @@ -140,6 +142,7 @@ export function isMessageFromServer(o: unknown): o is MessageFromServer { return "peers" in o && Array.isArray(o.peers) && o.peers.every(isNodeID); case MType.WaitingForMoreParticipants: case MType.EnoughParticipants: + case MType.ParticipantsUpdate: case MType.StartWeightSharing: case MType.RetryPeerConnections: case MType.ConnectionFail: diff --git a/discojs/src/client/federated/messages.ts b/discojs/src/client/federated/messages.ts index 9172efb56..20e932bef 100644 --- a/discojs/src/client/federated/messages.ts +++ b/discojs/src/client/federated/messages.ts @@ -6,6 +6,7 @@ import type { ClientConnected, WaitingForMoreParticipants, EnoughParticipants, + ParticipantsUpdate, } from "#client/mtype"; // See ../messages.ts for doc @@ -15,7 +16,8 @@ export type MessageFederated = | SendPayload | ReceiveServerPayload | WaitingForMoreParticipants - | EnoughParticipants; + | EnoughParticipants + | ParticipantsUpdate; export interface NewFederatedNodeInfo { type: MType.NewFederatedNodeInfo; @@ -50,6 +52,7 @@ export function isMessageFederated(raw: unknown): raw is MessageFederated { case MType.ReceiveServerPayload: case MType.WaitingForMoreParticipants: case MType.EnoughParticipants: + case MType.ParticipantsUpdate: return true; } diff --git a/discojs/src/client/mtype.ts b/discojs/src/client/mtype.ts index f1b964aa5..30bbd7f1a 100644 --- a/discojs/src/client/mtype.ts +++ b/discojs/src/client/mtype.ts @@ -52,6 +52,12 @@ export enum MType { EnoughParticipants, SendPayload, ReceiveServerPayload, + + /* Both schemes */ + // Message sent by the server when a participant joined or left, so that + // clients don't have to wait for the end of the round to learn about it. + // Kept last as the enum values are what goes over the wire. + ParticipantsUpdate, } export function hasMessageType( @@ -80,3 +86,8 @@ export interface WaitingForMoreParticipants { type: MType.WaitingForMoreParticipants; nbOfParticipants: number; } + +export interface ParticipantsUpdate { + type: MType.ParticipantsUpdate; + nbOfParticipants: number; +} diff --git a/server/src/controllers/decentralized_controller.ts b/server/src/controllers/decentralized_controller.ts index 08108c5b3..5cbfb7cdf 100644 --- a/server/src/controllers/decentralized_controller.ts +++ b/server/src/controllers/decentralized_controller.ts @@ -78,8 +78,11 @@ export class DecentralizedController< joinedMidTraining: joinedMidTraining, }; ws.send(msgpack.encode(msg), { binary: true }); - // Send an update to participants if we can start/resume training - this.sendEnoughParticipantsMsgIfNeeded(peerId); + // Send an update to participants if we can start/resume training, + // which already carries the number of participants + if (!this.sendEnoughParticipantsMsgIfNeeded(peerId)) + // otherwise just tell them that someone joined + this.sendParticipantsUpdateMsg(peerId); break; } // Send by peers at the beginning of each training round to notify @@ -188,14 +191,20 @@ export class DecentralizedController< } // Check if we are already waiting for new participants to join - if (this.waitingForMoreParticipants) return; + if (this.waitingForMoreParticipants) { + // tell the remaining participants that one of them left + this.sendParticipantsUpdateMsg(); + return; + } // If no, check if we are still above the minimum number of participant required if (this.connections.size >= minNbOfParticipants) { + this.sendParticipantsUpdateMsg(); this.sendPeersForRoundIfNeeded(); return; } // If we are below the minimum number of participants - // tell remaining participants to wait until more participants join + // tell remaining participants to wait until more participants join, + // which already carries the number of participants this.sendWaitForMoreParticipantsMsg(); }); } diff --git a/server/src/controllers/federated_controller.ts b/server/src/controllers/federated_controller.ts index 9e1d6485f..348aee498 100644 --- a/server/src/controllers/federated_controller.ts +++ b/server/src/controllers/federated_controller.ts @@ -171,8 +171,11 @@ export class FederatedController extends TrainingController< nbOfParticipants: this.connections.size, }; ws.send(msgpack.encode(msg)); - // Send an update to participants if we can start/resume training - this.sendEnoughParticipantsMsgIfNeeded(clientId); + // Send an update to participants if we can start/resume training, + // which already carries the number of participants + if (!this.sendEnoughParticipantsMsgIfNeeded(clientId)) + // otherwise just tell them that someone joined + this.sendParticipantsUpdateMsg(clientId); break; } /* @@ -250,10 +253,14 @@ export class FederatedController extends TrainingController< if ( this.connections.size >= minNbOfParticipants || this.waitingForMoreParticipants - ) + ) { + // tell the remaining participants that one of them left + this.sendParticipantsUpdateMsg(); return; + } - // tell remaining participants to wait until more participants join + // tell remaining participants to wait until more participants join, + // which already carries the number of participants this.sendWaitForMoreParticipantsMsg(); }); } diff --git a/server/src/controllers/training_controller.ts b/server/src/controllers/training_controller.ts index 96c0ddf38..fbd0f5ae3 100644 --- a/server/src/controllers/training_controller.ts +++ b/server/src/controllers/training_controller.ts @@ -53,12 +53,40 @@ export abstract class TrainingController< this.connections = Map(); } + /** + * Notifies participants of how many of them there are. + * + * @param exclude a participant to leave out, typically one which just joined + * and already learned the count from the answer to its join request + */ + protected sendParticipantsUpdateMsg(exclude?: NodeID): void { + const msg: mtype.ParticipantsUpdate = { + type: mtype.MType.ParticipantsUpdate, + nbOfParticipants: this.connections.size, + }; + const encoded = msgpack.encode(msg); + + const recipients = + exclude !== undefined + ? this.connections.delete(exclude) + : this.connections; + recipients.forEach((participantWs, participantId) => { + debug( + "Sending participants update to client [%s]", + participantId.slice(0, 4), + ); + participantWs.send(encoded); + }); + } + /** * If enough participants joined, notifies them that the training can start/resume * * @param currentId the id of the participant that just joined + * @returns whether the participants were notified, which also tells them the + * number of participants */ - protected sendEnoughParticipantsMsgIfNeeded(currentId: NodeID) { + protected sendEnoughParticipantsMsgIfNeeded(currentId: NodeID): boolean { // If we are currently waiting for more participants to join and we now have enough, // broadcast to previously waiting participants that the training can start if ( @@ -81,7 +109,10 @@ export abstract class TrainingController< participantWs.send(msgpack.encode(msg)); }); this.waitingForMoreParticipants = false; // update the attribute + return true; } + + return false; } /** diff --git a/server/tests/client.spec.ts b/server/tests/client.spec.ts index ed3b0c489..2a4f773d2 100644 --- a/server/tests/client.spec.ts +++ b/server/tests/client.spec.ts @@ -137,17 +137,19 @@ function withoutModel< } /** - * A server answering a client's messages with whatever `answer` sends, so that - * a test can produce orderings the real server doesn't. + * A server answering a client's messages with whatever `answer` sends + * Returns a `close` function tearing down every connections */ async function serveStub( answer: ( msg: { type: mtype.MType }, send: (msg: ServerMessage) => void, ) => void, -): Promise<[http.Server, URL]> { +): Promise<{ url: URL; close: () => void }> { const handle = http.createServer(); - new WebSocketServer({ server: handle }).on("connection", (ws) => + const socket = new WebSocketServer({ server: handle }); + + socket.on("connection", (ws) => ws.on("message", (data: Buffer) => { const msg: unknown = msgpack.decode(data); if (!mtype.hasMessageType(msg)) return; @@ -164,9 +166,18 @@ async function serveStub( }), ); - return [handle, url]; + return { + url, + close: () => { + socket.clients.forEach((ws) => ws.terminate()); + handle.close(); + }, + }; } +/** Resolve once every already queued microtask ran */ +const settled = () => new Promise((resolve) => setImmediate(resolve)); + /** * A client learns how many participants there are from the message answering * its join request, but the server computes that count when sending it. As that @@ -174,16 +185,23 @@ async function serveStub( * EnoughParticipants sent right after can be processed first, carrying a more * recent count which the join answer must not overwrite. * - * Reproducing that ordering needs both messages to be sent back to back, which - * the real server doesn't do, hence the stubbed one below. */ describe("client joining while the participants change", () => { - // sent when another participant joined between the two messages + // sent because another participant joined right after this one asked to join const enoughParticipants: mtype.EnoughParticipants = { type: mtype.MType.EnoughParticipants, nbOfParticipants: 2, }; + /** Answer a join request with the newer count first, the join answer second */ + const answerJoin = + (joinAnswer: ServerMessage) => + (msg: { type: mtype.MType }, send: (msg: ServerMessage) => void) => { + if (msg.type !== mtype.MType.ClientConnected) return; + send(enoughParticipants); + send(joinAnswer); + }; + it("keeps the newer count when federated", async () => { const joinAnswer: federatedMessages.NewFederatedNodeInfo = { type: mtype.MType.NewFederatedNodeInfo, @@ -191,13 +209,9 @@ describe("client joining while the participants change", () => { waitForMoreParticipants: true, payload: undefined, round: 0, - nbOfParticipants: 1, + nbOfParticipants: 1, // outdated by the time we read it }; - const [handle, url] = await serveStub((msg, send) => { - if (msg.type !== mtype.MType.ClientConnected) return; - send(joinAnswer); - send(enoughParticipants); - }); + const { url, close } = await serveStub(answerJoin(joinAnswer)); const client = withoutModel( new FederatedClient( @@ -213,8 +227,7 @@ describe("client joining while the participants change", () => { expect(client.nbOfParticipants).to.equal(2); expect(client.waitingForMoreParticipants).to.be.false; } finally { - await client.disconnect(); - handle.close(); + close(); } }); @@ -224,13 +237,9 @@ describe("client joining while the participants change", () => { id: "node-id", waitForMoreParticipants: true, joinedMidTraining: false, - nbOfParticipants: 1, + nbOfParticipants: 1, // outdated by the time we read it }; - const [handle, url] = await serveStub((msg, send) => { - if (msg.type !== mtype.MType.ClientConnected) return; - send(joinAnswer); - send(enoughParticipants); - }); + const { url, close } = await serveStub(answerJoin(joinAnswer)); const client = withoutModel( new DecentralizedClient( @@ -246,8 +255,7 @@ describe("client joining while the participants change", () => { expect(client.nbOfParticipants).to.equal(2); expect(client.waitingForMoreParticipants).to.be.false; } finally { - await client.disconnect(); - handle.close(); + close(); } }); }); @@ -270,31 +278,22 @@ describe("peer failing to begin a round", () => { nbOfParticipants: 2, }; // a peer list containing our own id makes the client fail to begin the - // round, as connecting to the peers of the round would + // round, as failing to connect to the peers of the round would const badPeersForRound: decentralizedMessages.PeersForRound = { type: mtype.MType.PeersForRound, peers: [ownId], aggregationRound: 0, }; - let rounds = 0; - const [handle, url] = await serveStub((msg, send) => { + let push: ((msg: ServerMessage) => void) | undefined; + const { url, close } = await serveStub((msg, send) => { + push = send; switch (msg.type) { case mtype.MType.ClientConnected: send(joinAnswer); break; case mtype.MType.PeerIsReady: send(badPeersForRound); - // let the client handle the failure and listen again before telling - // it to retry, then to give up - rounds++; - setTimeout(() => - send( - rounds === 1 - ? { type: mtype.MType.RetryPeerConnections } - : { type: mtype.MType.ConnectionFail }, - ), - ); break; } }); @@ -306,24 +305,44 @@ describe("peer failing to begin a round", () => { new MeanAggregator(), ), ); + const participants: number[] = []; - client.on("participants", (nbOfParticipants) => - participants.push(nbOfParticipants), - ); + client.on("participants", (nbOfParticipants) => { + participants.push(nbOfParticipants); + }); + // the client gives up on the round right after that status, so waiting for + // it then letting the pending microtasks run is enough to observe it + let giveUp: (() => void) | undefined; + client.on("status", (status) => { + if (status === "connecting to peers") giveUp?.(); + }); + /** Resolves when the client gives up on the round it is beginning */ + const givesUpOnRound = () => + new Promise((resolve) => (giveUp = resolve)); try { await client.connect(); await client.onRoundBeginCommunication(); + expect(participants).to.deep.equal([2]); - await expect( - client.onRoundEndCommunication(new WeightsContainer([[1]])), - ).rejects.toThrow("Client disconnected after connection failure"); + // the round never completes, the peers of every round are unreachable + let failedRoundStart = givesUpOnRound(); + const round = client.onRoundEndCommunication(new WeightsContainer([[1]])); + round.catch(() => undefined); - // the server only ever said there were two of us + // failing to connect to the round's peers + await failedRoundStart; + await settled(); + expect(participants).to.deep.equal([2]); + + // being told to start the round over, then failing again + failedRoundStart = givesUpOnRound(); + push?.({ type: mtype.MType.RetryPeerConnections }); + await failedRoundStart; + await settled(); expect(participants).to.deep.equal([2]); - expect(rounds).to.equal(2); // the retry did happen } finally { - handle.close(); + close(); } }); }); diff --git a/server/tests/e2e/decentralized_controller.spec.ts b/server/tests/e2e/decentralized_controller.spec.ts index 044448b5b..ff9bb8a4e 100644 --- a/server/tests/e2e/decentralized_controller.spec.ts +++ b/server/tests/e2e/decentralized_controller.spec.ts @@ -212,3 +212,80 @@ describe("DecentralizedController peer connection retry", () => { ).to.have.length(4); }); }); + +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: FakeWebSocket, + ): void { + controller.handle(ws); + ws.emitMessage({ type: MessageTypes.ClientConnected }); + } + + /** The counts a peer was told about, after it joined */ + const participantsSeen = (ws: FakeWebSocket): number[] => + messagesOfType(ws, MessageTypes.ParticipantsUpdate).map( + (msg) => msg.nbOfParticipants, + ); + + it("tells the peers when one joins or leaves", async () => { + const controller = await makeController(); + + const [ws1, ws2, ws3] = [ + makeFakeWebSocket(), + makeFakeWebSocket(), + makeFakeWebSocket(), + ]; + + connect(controller, ws1); + // the second peer meets the minimum, which EnoughParticipants carries + connect(controller, ws2); + expect(participantsSeen(ws1)).to.deep.equal([]); + expect( + lastMessageOfType(ws1, MessageTypes.EnoughParticipants)?.nbOfParticipants, + ).to.equal(2); + + // the third one changes nothing but the count + connect(controller, ws3); + expect(participantsSeen(ws1)).to.deep.equal([3]); + expect(participantsSeen(ws2)).to.deep.equal([3]); + // it learned the count from the answer to its own join request + expect(participantsSeen(ws3)).to.deep.equal([]); + expect( + lastMessageOfType(ws3, MessageTypes.NewDecentralizedNodeInfo) + ?.nbOfParticipants, + ).to.equal(3); + + // leaving while the minimum is still met + ws3.emitClose(); + expect(participantsSeen(ws1)).to.deep.equal([3, 2]); + expect(participantsSeen(ws2)).to.deep.equal([3, 2]); + + // dropping below it is carried by WaitingForMoreParticipants instead + ws2.emitClose(); + expect(participantsSeen(ws1)).to.deep.equal([3, 2]); + expect( + lastMessageOfType(ws1, MessageTypes.WaitingForMoreParticipants) + ?.nbOfParticipants, + ).to.equal(1); + }); +}); diff --git a/server/tests/e2e/federated.spec.ts b/server/tests/e2e/federated.spec.ts index 6c3c7c6e9..708b8ca96 100644 --- a/server/tests/e2e/federated.spec.ts +++ b/server/tests/e2e/federated.spec.ts @@ -330,6 +330,39 @@ describe("end-to-end federated", () => { await client.expectParticipants(clients.length); } + /** + * A client joining a session which already has enough participants. It + * releases nobody and completes nobody's round, it only makes the others + * one more: the server is the only one able to tell them. + */ + async function joinsSession( + present: readonly Client[], + name: string, + ): Promise { + const client = join(name).startRound(); + + const participants = present.length + 1; + await client.expectParticipants(participants); + await client.expectStatuses("local training"); + for (const other of present) await other.expectParticipants(participants); + + return client; + } + + /** + * A client leaving a session which keeps enough participants: the others + * carry on, one fewer. + */ + async function leavesSession( + leaving: Client, + remaining: readonly Client[], + ): Promise { + await leaving.leave(); + + for (const other of remaining) + await other.expectParticipants(remaining.length); + } + /** A client leaving, the remaining one is left without enough participants */ async function leavesTask( leaving: Client, @@ -382,6 +415,25 @@ describe("end-to-end federated", () => { await startsRoundAlone(user2); }); + it("clients are notified when a participant joins", async () => { + const user1 = await joinsAlone("user 1"); + const user2 = await joinsWaitingClient(user1, "user 2"); + + // the minimum is already met so user 3 changes nothing but the count, + // which the others would otherwise only learn when a round of theirs + // completes, if one ever does + await joinsSession([user1, user2], "user 3"); + }); + + it("clients are notified when a participant leaves", async () => { + const user1 = await joinsAlone("user 1"); + const user2 = await joinsWaitingClient(user1, "user 2"); + const user3 = await joinsSession([user1, user2], "user 3"); + + // the remaining two still have enough participants to carry on + await leavesSession(user3, [user1, user2]); + }); + it("a client joining mid-training completes the pending round", async () => { const user1 = await joinsAlone("user 1"); const user2 = await joinsWaitingClient(user1, "user 2"); From 4347094d05b7fb0b82edfcb70f827a4debbf18db Mon Sep 17 00:00:00 2001 From: Julien Vignoud Date: Fri, 11 Sep 2026 15:24:56 +0200 Subject: [PATCH 5/5] fix: exclude model syncing peers from round participant numbers --- .../decentralized/decentralized_client.ts | 8 ++-- discojs/src/client/decentralized/messages.ts | 3 ++ .../controllers/decentralized_controller.ts | 1 + server/tests/client.spec.ts | 1 + .../e2e/decentralized_controller.spec.ts | 42 +++++++++++++++++++ 5 files changed, 51 insertions(+), 4 deletions(-) diff --git a/discojs/src/client/decentralized/decentralized_client.ts b/discojs/src/client/decentralized/decentralized_client.ts index 2028a9555..fb9499290 100644 --- a/discojs/src/client/decentralized/decentralized_client.ts +++ b/discojs/src/client/decentralized/decentralized_client.ts @@ -363,10 +363,10 @@ export class DecentralizedClient extends Client<"decentralized"> { throw new Error("received peer list contains our own id"); } // Store the list of peers for the current round including ourselves - const roundNodes = peers.add(this.ownId); - this.aggregator.setNodes(roundNodes); - // the server told us who takes part in the round, emits `participants` - this.nbOfParticipants = roundNodes.size; + this.aggregator.setNodes(peers.add(this.ownId)); + // the peers of the round leave out those still syncing their model, so + // the server tells us how many participants the session has + this.nbOfParticipants = receivedMessage.nbOfParticipants; this.aggregator.setRound(receivedMessage.aggregationRound); // the server gives us the round number // Initiate peer to peer connections with each peer diff --git a/discojs/src/client/decentralized/messages.ts b/discojs/src/client/decentralized/messages.ts index 309dd7b71..0bd18011a 100644 --- a/discojs/src/client/decentralized/messages.ts +++ b/discojs/src/client/decentralized/messages.ts @@ -41,6 +41,9 @@ export interface PeersForRound { type: MType.PeersForRound; peers: NodeID[]; aggregationRound: number; + // the peers of a round are only those able to take part in it, so they + // aren't the number of participants of the session + nbOfParticipants: number; } // peer sends to server to signal all the connections to other peers diff --git a/server/src/controllers/decentralized_controller.ts b/server/src/controllers/decentralized_controller.ts index 5cbfb7cdf..42a78a74b 100644 --- a/server/src/controllers/decentralized_controller.ts +++ b/server/src/controllers/decentralized_controller.ts @@ -283,6 +283,7 @@ export class DecentralizedController< type: MessageTypes.PeersForRound, peers: this.#roundPeers.delete(id).keySeq().toArray(), aggregationRound: this.#aggregationRound, + nbOfParticipants: this.connections.size, }; }, ); diff --git a/server/tests/client.spec.ts b/server/tests/client.spec.ts index 2a4f773d2..e71d11fe5 100644 --- a/server/tests/client.spec.ts +++ b/server/tests/client.spec.ts @@ -283,6 +283,7 @@ describe("peer failing to begin a round", () => { type: mtype.MType.PeersForRound, peers: [ownId], aggregationRound: 0, + nbOfParticipants: 2, }; let push: ((msg: ServerMessage) => void) | undefined; diff --git a/server/tests/e2e/decentralized_controller.spec.ts b/server/tests/e2e/decentralized_controller.spec.ts index ff9bb8a4e..a60e2ff33 100644 --- a/server/tests/e2e/decentralized_controller.spec.ts +++ b/server/tests/e2e/decentralized_controller.spec.ts @@ -247,6 +247,48 @@ describe("DecentralizedController participants updates", () => { (msg) => msg.nbOfParticipants, ); + it("counts a peer syncing its model like every other message does", async () => { + const controller = await makeController(); + + const [ws1, ws2, ws3] = [ + makeFakeWebSocket(), + makeFakeWebSocket(), + makeFakeWebSocket(), + ]; + + // two peers run a first round together + connect(controller, ws1); + connect(controller, ws2); + for (const ws of [ws1, ws2]) { + ws.emitMessage({ type: MessageTypes.JoinRound }); + ws.emitMessage({ type: MessageTypes.PeerIsReady }); + } + for (const ws of [ws1, ws2]) + ws.emitMessage({ type: MessageTypes.ConnectionsReady }); + + // a third one joins, which has to sync its model before taking part + connect(controller, ws3); + expect( + lastMessageOfType(ws3, MessageTypes.NewDecentralizedNodeInfo) + ?.joinedMidTraining, + ).to.be.true; + expect( + lastMessageOfType(ws3, MessageTypes.NewDecentralizedNodeInfo) + ?.nbOfParticipants, + ).to.equal(3); + expect(participantsSeen(ws1)).to.deep.equal([3]); + + // the other two start a round without it, as it is still syncing + for (const ws of [ws1, ws2]) { + ws.emitMessage({ type: MessageTypes.JoinRound }); + ws.emitMessage({ type: MessageTypes.PeerIsReady }); + } + const peersForRound = lastMessageOfType(ws1, MessageTypes.PeersForRound); + // the round leaves the syncing peer out, but it is still a participant + expect(peersForRound?.peers).to.have.length(1); + expect(peersForRound?.nbOfParticipants).to.equal(3); + }); + it("tells the peers when one joins or leaves", async () => { const controller = await makeController();