diff --git a/discojs/src/client/client.ts b/discojs/src/client/client.ts index 7333ad363..ae825750b 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 @@ -138,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 @@ -147,10 +163,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..fb9499290 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())); } @@ -156,7 +149,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 @@ -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)); + 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 @@ -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/discojs/src/client/decentralized/messages.ts b/discojs/src/client/decentralized/messages.ts index 1c2b334fb..0bd18011a 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) @@ -40,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 @@ -107,6 +111,7 @@ export type MessageFromServer = | PeersForRound | WaitingForMoreParticipants | EnoughParticipants + | ParticipantsUpdate | StartWeightSharing | RetryPeerConnections | ConnectionFail @@ -140,6 +145,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/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/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..42a78a74b 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(); }); } @@ -274,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/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 4568aad8e..e71d11fe5 100644 --- a/server/tests/client.spec.ts +++ b/server/tests/client.spec.ts @@ -1,14 +1,21 @@ -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, + WeightsContainer, + mtype, defaultTasks, defaultModels, } from "@epfml/discojs"; @@ -116,3 +123,227 @@ describe("federated client", () => { await expect(client.connect()).rejects.toThrow(); }); }); + +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 + * Returns a `close` function tearing down every connections + */ +async function serveStub( + answer: ( + msg: { type: mtype.MType }, + send: (msg: ServerMessage) => void, + ) => void, +): Promise<{ url: URL; close: () => void }> { + const handle = http.createServer(); + 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; + 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 { + 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 + * 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. + * + */ +describe("client joining while the participants change", () => { + // 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, + id: "node-id", + waitForMoreParticipants: true, + payload: undefined, + round: 0, + nbOfParticipants: 1, // outdated by the time we read it + }; + const { url, close } = await serveStub(answerJoin(joinAnswer)); + + 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 { + 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, // outdated by the time we read it + }; + const { url, close } = await serveStub(answerJoin(joinAnswer)); + + 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 { + close(); + } + }); +}); + +/** + * 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 failing to connect to the peers of the round would + const badPeersForRound: decentralizedMessages.PeersForRound = { + type: mtype.MType.PeersForRound, + peers: [ownId], + aggregationRound: 0, + nbOfParticipants: 2, + }; + + 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); + break; + } + }); + + const client = withoutModel( + new DecentralizedClient( + url, + await defaultTasks.cifar10.getTask(), + new MeanAggregator(), + ), + ); + + const participants: number[] = []; + 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]); + + // the round never completes, the peers of every round are unreachable + let failedRoundStart = givesUpOnRound(); + const round = client.onRoundEndCommunication(new WeightsContainer([[1]])); + round.catch(() => undefined); + + // 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]); + } finally { + close(); + } + }); +}); diff --git a/server/tests/e2e/decentralized_controller.spec.ts b/server/tests/e2e/decentralized_controller.spec.ts index 044448b5b..a60e2ff33 100644 --- a/server/tests/e2e/decentralized_controller.spec.ts +++ b/server/tests/e2e/decentralized_controller.spec.ts @@ -212,3 +212,122 @@ 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("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(); + + 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"); 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); + }, +);