@epfml/discojs 3.0.1-p20241001093123.0 → 3.0.1-p20241014092014.0

package/dist/client/client.js

@@ -1,6 +1,9 @@
+import createDebug from "debug";
 import axios from 'axios';
 import { serialization } from '../index.js';
 import { EventEmitter } from '../utils/event_emitter.js';
+import { type } from "./messages.js";
+const debug = createDebug("discojs:client");
 /**
  * Main, abstract, class representing a Disco client in a network, which handles
  * communication with other nodes, be it peers or a server.
@@ -9,18 +12,25 @@ export class Client extends EventEmitter {
     url;
     task;
     aggregator;
-    /**
-     * Own ID provided by the network's server.
-     */
+    // Own ID provided by the network's server.
     _ownId;
+    // The network's server.
+    _server;
+    // The aggregator's result produced after aggregation.
+    aggregationResult;
     /**
-     * The network's server.
+     * When the server notifies clients to pause and wait until more
+     * participants join, we rely on this promise to wait
+     * until the server signals that the training can resume
      */
-    _server;
+    promiseForMoreParticipants = undefined;
     /**
-     * The aggregator's result produced after aggregation.
+     * When the server notifies the client that they can resume training
+     * after waiting for more participants, we want to be able to display what
+     * we were doing before waiting (training locally or updating our model).
+     * We use this attribute to store the status to rollback to when we stop waiting
      */
-    aggregationResult;
+    previousStatus;
     constructor(url, // The network server's URL to connect to
     task, // The client's corresponding task
     aggregator) {
@@ -41,6 +51,94 @@ export class Client extends EventEmitter {
      * Handles the disconnection process of the client from any sort of network server.
      */
     async disconnect() { }
+    /**
+     * Emits the round status specified. It also stores the status emitted such that
+     * if the server tells the client to wait for more participants, it can display
+     * the waiting status and once enough participants join, it can display the previous status again
+     */
+    saveAndEmit(status) {
+        this.previousStatus = status;
+        this.emit("status", status);
+    }
+    /**
+     * For both federated and decentralized clients, we listen to the server to tell
+     * us whether there are enough participants to train. If not, we pause until further notice.
+     * When a client connects to the server, the server answers with the session information (id,
+     * number of participants) and whether there are enough participants.
+     * When there are the server sends a new EnoughParticipant message to update the client.
+     *
+     * `setMessageInversionFlag` is used to address the following scenario:
+     * 1. Client 1 connect to the server
+     * 2. Server answers with message A containing "not enough participants"
+     * 3. Before A arrives a new client joins. There are enough participants now.
+     * 4. Server updates client 1 with message B saying "there are enough participants"
+     * 5. Due to network and message sizes message B can arrive before A.
+     * i.e. "there are enough participants" arrives before "not enough participants"
+     * ending up with client 1 thinking it needs to wait for more participants.
+     *
+     * To keep track of this message inversion, `setMessageInversionFlag`
+     * tells us whether a message inversion occurred (by setting a boolean to true)
+     *
+     * @param setMessageInversionFlag function flagging whether a message inversion occurred
+     * between a NewNodeInfo message and an EnoughParticipant message.
+     */
+    setupServerCallbacks(setMessageInversionFlag) {
+        // Setup an event callback if the server signals that we should 
+        // wait for more participants
+        this.server.on(type.WaitingForMoreParticipants, () => {
+            if (this.promiseForMoreParticipants !== undefined)
+                throw new Error("Server sent multiple WaitingForMoreParticipants messages");
+            debug(`[${shortenId(this.ownId)}] received WaitingForMoreParticipants message from server`);
+            // Display the waiting status right away
+            this.emit("status", "not enough participants");
+            // Upon receiving a WaitingForMoreParticipants message,
+            // the client will await for this promise to resolve before sending its
+            // local weight update
+            this.promiseForMoreParticipants = this.createPromiseForMoreParticipants();
+        });
+        // 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
+        // and directly follows with an EnoughParticipants message when the 2nd participant joins
+        // However, the EnoughParticipants can arrive before the NewNodeInfo (which can be much bigger)
+        // so we check whether we received the EnoughParticipants before being assigned a node ID
+        this.server.once(type.EnoughParticipants, () => {
+            if (this._ownId === undefined) {
+                debug(`Received EnoughParticipants message from server before the NewFederatedNodeInfo message`);
+                setMessageInversionFlag();
+            }
+        });
+    }
+    /**
+     * Method called when the server notifies the client that there aren't enough
+     * participants (anymore) to start/continue training
+     * The method creates a promise that will resolve once the server notifies
+     * the client that the training can resume via a subsequent EnoughParticipants message
+     * @returns a promise which resolves when enough participants joined the session
+     */
+    async createPromiseForMoreParticipants() {
+        return new Promise((resolve) => {
+            // "once" is important because we can't resolve the same promise multiple times
+            this.server.once(type.EnoughParticipants, () => {
+                debug(`[${shortenId(this.ownId)}] received EnoughParticipants message from server`);
+                // Emit the last status emitted before waiting if defined
+                if (this.previousStatus !== undefined)
+                    this.emit("status", this.previousStatus);
+                resolve();
+            });
+        });
+    }
+    async waitForParticipantsIfNeeded() {
+        // we check if we are waiting for more participants before sending our weight update
+        if (this.waitingForMoreParticipants) {
+            // wait for the promise to resolve, which takes as long as it takes for new participants to join
+            debug(`[${shortenId(this.ownId)}] is awaiting the promise for more participants`);
+            this.emit("status", "not enough participants");
+            await this.promiseForMoreParticipants;
+            // Make sure to set the promise back to undefined once resolved
+            this.promiseForMoreParticipants = undefined;
+        }
+    }
     /**
      * Fetches the latest model available on the network's server, for the adequate task.
      * @returns The latest model
@@ -54,13 +152,6 @@ export class Client extends EventEmitter {
         const response = await axios.get(url.href, { responseType: 'arraybuffer' });
         return await serialization.model.decode(new Uint8Array(response.data));
     }
-    // Number of contributors to a collaborative session
-    // If decentralized, it should be the number of peers
-    // If federated, it should the number of participants excluding the server
-    // If local it should be 1
-    get nbOfParticipants() {
-        return this.aggregator.nodes.size; // overriden by the federated client
-    }
     get ownId() {
         if (this._ownId === undefined) {
             throw new Error('the node is not connected');
@@ -73,4 +164,14 @@ export class Client extends EventEmitter {
         }
         return this._server;
     }
+    /**
+     * Whether the client should wait until more
+     * participants join the session, i.e. a promise has been created
+     */
+    get waitingForMoreParticipants() {
+        return this.promiseForMoreParticipants !== undefined;
+    }
+}
+export function shortenId(id) {
+    return id.slice(0, 4);
 }

package/dist/client/decentralized/decentralized_client.d.ts

@@ -1,5 +1,5 @@
 import type { Model, WeightsContainer } from "../../index.js";
-import { Client } from '../index.js';
+import { Client } from '../client.js';
 /**
  * Represents a decentralized client in a network of peers. Peers coordinate each other with the
  * help of the network's server, yet only exchange payloads between each other. Communication
@@ -7,11 +7,8 @@ import { Client } from '../index.js';
  * WebRTC for Node.js.
  */
 export declare class DecentralizedClient extends Client {
-    /**
-     * The pool of peers to communicate with during the current training round.
-     */
-    private pool?;
-    private connections?;
+    #private;
+    getNbOfParticipants(): number;
     private get isDisconnected();
     /**
      * Public method called by disco.ts when starting training. This method sends
@@ -22,12 +19,6 @@ export declare class DecentralizedClient extends Client {
      * peers network information.
      */
     connect(): Promise<Model>;
-    /**
-     * Create a WebSocket connection with the server
-     * The client then waits for the server to forward it other client's network information.
-     * Upon receiving other peer's information, the clients establish a peer-to-peer WebRTC connection.
-     */
-    private connectServer;
     disconnect(): Promise<void>;
     /**
      * At the beginning of a round, each peer tells the server it is ready to proceed
@@ -38,6 +29,13 @@ export declare class DecentralizedClient extends Client {
      *
      */
     onRoundBeginCommunication(): Promise<void>;
+    onRoundEndCommunication(weights: WeightsContainer): Promise<WeightsContainer>;
+    /**
+     * Signal to the server that we are ready to exchange weights.
+     * Once enough peers are ready, the server sends the list of peers for this round
+     * and the peers can establish peer-to-peer connections with each other.
+     */
+    private establishPeerConnections;
     /**
      * At each communication rounds, awaits peers contributions and add them to the client's aggregator.
      * This method is used as callback by getPeers when connecting to the rounds' peers
@@ -45,5 +43,5 @@ export declare class DecentralizedClient extends Client {
      * @param round
      */
     private receivePayloads;
-    onRoundEndCommunication(weights: WeightsContainer): Promise<WeightsContainer>;
+    private exchangeWeightUpdates;
 }

package/dist/client/decentralized/decentralized_client.js

@@ -1,7 +1,7 @@
 import createDebug from "debug";
 import { Map, Set } from 'immutable';
 import { serialization } from "../../index.js";
-import { Client } from '../index.js';
+import { Client, shortenId } from '../client.js';
 import { type } from '../messages.js';
 import { timeout } from '../utils.js';
 import { WebSocketServer, waitMessage, waitMessageWithTimeout } from '../event_connection.js';
@@ -18,8 +18,12 @@ export class DecentralizedClient extends Client {
     /**
      * The pool of peers to communicate with during the current training round.
      */
-    pool;
-    connections;
+    #pool;
+    #connections;
+    getNbOfParticipants() {
+        const nbOfParticipants = this.aggregator.nodes.size;
+        return nbOfParticipants === 0 ? 1 : nbOfParticipants;
+    }
     // Used to handle timeouts and promise resolving after calling disconnect
     get isDisconnected() {
         return this._server === undefined;
@@ -46,42 +50,48 @@ export class DecentralizedClient extends Client {
                 throw new Error(`unknown protocol: ${this.url.protocol}`);
         }
         serverURL.pathname += `decentralized/${this.task.id}`;
-        this._server = await this.connectServer(serverURL);
+        // Create a WebSocket connection with the server
+        // The client then waits for the server to forward it other client's network information.
+        // Upon receiving other peer's information, the clients establish a peer-to-peer WebRTC connection.
+        this._server = await WebSocketServer.connect(serverURL, messages.isMessageFromServer, messages.isMessageToServer);
+        this.server.on(type.SignalForPeer, (event) => {
+            if (this.#pool === undefined)
+                throw new Error('received signal but peer pool is undefined');
+            // Create a WebRTC connection with the peer
+            this.#pool.signal(event.peer, event.signal);
+        });
+        // c.f. setupServerCallbacks doc for explanation
+        let receivedEnoughParticipants = false;
+        this.setupServerCallbacks(() => receivedEnoughParticipants = true);
         const msg = {
             type: type.ClientConnected
         };
         this.server.send(msg);
-        const { id } = await waitMessage(this.server, type.NewDecentralizedNodeInfo);
-        debug(`[${id}] assigned id generated by server`);
+        const { id, waitForMoreParticipants } = await waitMessage(this.server, type.NewDecentralizedNodeInfo);
+        // This should come right after receiving the message to make sure
+        // we don't miss a subsequent message from the server
+        // We check if the server is telling us to wait for more participants
+        // and we also check if a EnoughParticipant message ended up arriving
+        // before the NewNodeInfo
+        if (waitForMoreParticipants && !receivedEnoughParticipants) {
+            // Create a promise that resolves when enough participants join
+            // The client will await this promise before sending its local weight update
+            this.promiseForMoreParticipants = this.createPromiseForMoreParticipants();
+        }
+        debug(`[${shortenId(id)}] assigned id generated by server`);
         if (this._ownId !== undefined) {
             throw new Error('received id from server but was already received');
         }
         this._ownId = id;
-        this.pool = new PeerPool(id);
+        this.#pool = new PeerPool(id);
         return model;
     }
-    /**
-     * Create a WebSocket connection with the server
-     * The client then waits for the server to forward it other client's network information.
-     * Upon receiving other peer's information, the clients establish a peer-to-peer WebRTC connection.
-     */
-    async connectServer(url) {
-        const server = await WebSocketServer.connect(url, messages.isMessageFromServer, messages.isMessageToServer);
-        server.on(type.SignalForPeer, (event) => {
-            if (this.pool === undefined) {
-                throw new Error('received signal but peer pool is undefined');
-            }
-            // Create a WebRTC connection with the peer
-            this.pool.signal(event.peer, event.signal);
-        });
-        return server;
-    }
     async disconnect() {
         // Disconnect from peers
-        await this.pool?.shutdown();
-        this.pool = undefined;
-        if (this.connections !== undefined) {
-            const peers = this.connections.keySeq().toSet();
+        await this.#pool?.shutdown();
+        this.#pool = undefined;
+        if (this.#connections !== undefined) {
+            const peers = this.#connections.keySeq().toSet();
             this.aggregator.setNodes(this.aggregator.nodes.subtract(peers));
         }
         // Disconnect from server
@@ -99,13 +109,41 @@ export class DecentralizedClient extends Client {
      *
      */
     async onRoundBeginCommunication() {
+        // Notify the server we want to join the next round so that the server
+        // waits for us to be ready before sending the list of peers for the round
+        this.server.send({ type: type.JoinRound });
+        // Store the promise for the current round's aggregation result.
+        // We will await for it to resolve at the end of the round when exchanging weight updates.
+        this.aggregationResult = this.aggregator.getPromiseForAggregation();
+        this.saveAndEmit("local training");
+        return Promise.resolve();
+    }
+    async onRoundEndCommunication(weights) {
+        if (this.aggregationResult === undefined) {
+            throw new TypeError('aggregation result promise is undefined');
+        }
+        // Save the status in case participants leave and we switch to waiting for more participants
+        // Once enough new participants join we can display the previous status again
+        this.saveAndEmit("connecting to peers");
+        // First we check if we are waiting for more participants before sending our weight update
+        await this.waitForParticipantsIfNeeded();
+        // Create peer-to-peer connections with all peers for the round
+        await this.establishPeerConnections();
+        // Exchange weight updates with peers and return aggregated weights
+        return await this.exchangeWeightUpdates(weights);
+    }
+    /**
+     * Signal to the server that we are ready to exchange weights.
+     * Once enough peers are ready, the server sends the list of peers for this round
+     * and the peers can establish peer-to-peer connections with each other.
+     */
+    async establishPeerConnections() {
         if (this.server === undefined) {
             throw new Error("peer's server is undefined, make sure to call `client.connect()` first");
         }
-        if (this.pool === undefined) {
+        if (this.#pool === undefined) {
             throw new Error('peer pool is undefined, make sure to call `client.connect()` first');
         }
-        this.emit("status", "Retrieving peers' information");
         // Reset peers list at each round of training to make sure client works with an updated peers
         // list, maintained by the server. Adds any received weights to the aggregator.
         // Tell the server we are ready for the next round
@@ -113,33 +151,29 @@ export class DecentralizedClient extends Client {
         this.server.send(readyMessage);
         // Wait for the server to answer with the list of peers for the round
         try {
-            debug(`[${this.ownId}] is waiting for peer list for round ${this.aggregator.round}`);
-            const receivedMessage = await waitMessageWithTimeout(this.server, type.PeersForRound, undefined, "Timeout waiting for the round's peer list");
+            debug(`[${shortenId(this.ownId)}] is waiting for peer list for round ${this.aggregator.round}`);
+            const receivedMessage = await waitMessage(this.server, type.PeersForRound);
             const peers = Set(receivedMessage.peers);
             if (this.ownId !== undefined && peers.has(this.ownId)) {
                 throw new Error('received peer list contains our own id');
             }
             // Store the list of peers for the current round including ourselves
             this.aggregator.setNodes(peers.add(this.ownId));
+            this.aggregator.setRound(receivedMessage.aggregationRound); // the server gives us the round number
             // Initiate peer to peer connections with each peer
             // When connected, create a promise waiting for each peer's round contribution
-            const connections = await this.pool.getPeers(peers, this.server, 
-            // Init receipt of peers weights
-            // this awaits the peer's weight update and adds it to 
-            // our aggregator upon reception
-            (conn) => { this.receivePayloads(conn, this.aggregator.round); });
-            debug(`[${this.ownId}] received peers for round ${this.aggregator.round}: %o`, connections.keySeq().toJS());
-            this.connections = connections;
+            const connections = await this.#pool.getPeers(peers, this.server, 
+            // Init receipt of peers weights. this awaits the peer's
+            // weight update and adds it to our aggregator upon reception
+            (conn) => this.receivePayloads(conn));
+            debug(`[${shortenId(this.ownId)}] received peers for round ${this.aggregator.round}: %o`, connections.keySeq().toJS());
+            this.#connections = connections;
         }
         catch (e) {
-            debug(`Error for [${this.ownId}] while beginning round: %o`, e);
+            debug(`Error for [${shortenId(this.ownId)}] while beginning round: %o`, e);
             this.aggregator.setNodes(Set(this.ownId));
-            this.connections = Map();
+            this.#connections = Map();
         }
-        // Store the promise for the current round's aggregation result.
-        // We will await for it to resolve at the end of the round when exchanging weight updates.
-        this.aggregationResult = new Promise((resolve) => this.aggregator.once('aggregation', resolve));
-        this.emit("status", "Training the model on the data you connected");
     }
     /**
      * At each communication rounds, awaits peers contributions and add them to the client's aggregator.
@@ -147,66 +181,71 @@ export class DecentralizedClient extends Client {
      * @param connections
      * @param round
      */
-    receivePayloads(connections, round) {
+    receivePayloads(connections) {
         connections.forEach(async (connection, peerId) => {
-            let currentCommunicationRounds = 0;
             debug(`waiting for peer ${peerId}`);
-            do {
+            for (let r = 0; r < this.aggregator.communicationRounds; r++) {
                 try {
                     const message = await waitMessageWithTimeout(connection, type.Payload, 60_000, "Timeout waiting for a contribution from peer " + peerId);
                     const decoded = serialization.weights.decode(message.payload);
-                    if (!this.aggregator.add(peerId, decoded, round, message.round)) {
-                        debug(`[${this.ownId}] failed to add contribution from peer ${peerId}`);
+                    if (!this.aggregator.isValidContribution(peerId, message.aggregationRound)) {
+                        debug(`[${shortenId(this.ownId)}] failed to add contribution from peer ${shortenId(peerId)}`);
+                    }
+                    else {
+                        debug(`[${shortenId(this.ownId)}] received payload from peer ${shortenId(peerId)}` +
+                            ` for round (%d, %d)`, message.aggregationRound, message.communicationRound);
+                        this.aggregator.once("aggregation", () => debug(`[${shortenId(this.ownId)}] aggregated the model` +
+                            ` for round (%d, %d)`, message.aggregationRound, message.communicationRound));
+                        this.aggregator.add(peerId, decoded, message.aggregationRound, message.communicationRound);
                     }
                 }
                 catch (e) {
                     if (this.isDisconnected)
                         return;
-                    debug(`Error for [${this.ownId}] while receiving payloads: %o`, e);
+                    debug(`Error for [${shortenId(this.ownId)}] while receiving payloads: %o`, e);
                 }
-            } while (++currentCommunicationRounds < this.aggregator.communicationRounds);
+            }
         });
     }
-    async onRoundEndCommunication(weights) {
+    async exchangeWeightUpdates(weights) {
         if (this.aggregationResult === undefined) {
             throw new TypeError('aggregation result promise is undefined');
         }
-        this.emit("status", "Updating the model with other participants' models");
+        this.saveAndEmit("updating model");
         // Perform the required communication rounds. Each communication round consists in sending our local payload,
         // followed by an aggregation step triggered by the receipt of other payloads, and handled by the aggregator.
         // A communication round's payload is the aggregation result of the previous communication round. The first
         // communication round simply sends our training result, i.e. model weights updates. This scheme allows for
         // the aggregator to define any complex multi-round aggregation mechanism.
         let result = weights;
-        for (let r = 0; r < this.aggregator.communicationRounds; r++) {
+        for (let communicationRound = 0; communicationRound < this.aggregator.communicationRounds; communicationRound++) {
+            const connections = this.#connections;
+            if (connections === undefined)
+                throw new Error("peer's connections is undefined");
             // Generate our payloads for this communication round and send them to all ready connected peers
-            if (this.connections !== undefined) {
-                const payloads = this.aggregator.makePayloads(result);
-                try {
-                    await Promise.all(payloads.map(async (payload, id) => {
-                        if (id === this.ownId) {
-                            this.aggregator.add(this.ownId, payload, this.aggregator.round, r);
-                        }
-                        else {
-                            const peer = this.connections?.get(id);
-                            if (peer !== undefined) {
-                                const encoded = await serialization.weights.encode(payload);
-                                const msg = {
-                                    type: type.Payload,
-                                    peer: id,
-                                    round: r,
-                                    payload: encoded
-                                };
-                                peer.send(msg);
-                                debug(`[${this.ownId}] send weight update to peer ${msg.peer}: %O`, msg);
-                            }
-                        }
-                    }));
+            const payloads = this.aggregator.makePayloads(result);
+            payloads.forEach(async (payload, id) => {
+                // add our own contribution to the aggregator
+                if (id === this.ownId) {
+                    this.aggregator.add(this.ownId, payload, this.aggregator.round, communicationRound);
+                    return;
                 }
-                catch (cause) {
-                    throw new Error('error while sending weights', { cause });
+                // Send our payload to each peer
+                const peer = connections.get(id);
+                if (peer !== undefined) {
+                    const encoded = await serialization.weights.encode(payload);
+                    const msg = {
+                        type: type.Payload,
+                        peer: id,
+                        aggregationRound: this.aggregator.round,
+                        communicationRound,
+                        payload: encoded
+                    };
+                    peer.send(msg);
+                    debug(`[${shortenId(this.ownId)}] send weight update to peer ${shortenId(msg.peer)}` +
+                        ` for round (%d, %d)`, this.aggregator.round, communicationRound);
                 }
-            }
+            });
             // Wait for aggregation before proceeding to the next communication round.
             // The current result will be used as payload for the eventual next communication round.
             try {
@@ -219,17 +258,15 @@ export class DecentralizedClient extends Client {
                 if (this.isDisconnected) {
                     return weights;
                 }
-                debug(`[${this.ownId}] while waiting for aggregation: %o`, e);
+                debug(`[${shortenId(this.ownId)}] while waiting for aggregation: %o`, e);
                 break;
             }
             // There is at least one communication round remaining
-            if (r < this.aggregator.communicationRounds - 1) {
+            if (communicationRound < this.aggregator.communicationRounds - 1) {
                 // Reuse the aggregation result
-                this.aggregationResult = new Promise((resolve) => this.aggregator.once('aggregation', resolve));
+                this.aggregationResult = this.aggregator.getPromiseForAggregation();
             }
         }
-        // Reset the peers list for the next round
-        this.aggregator.resetNodes();
         return await this.aggregationResult;
     }
 }

package/dist/client/decentralized/messages.d.ts

@@ -1,7 +1,8 @@
 import { weights } from '../../serialization/index.js';
 import { type SignalData } from './peer.js';
 import { type NodeID } from '../types.js';
-import { type, type ClientConnected } from '../messages.js';
+import { type } from '../messages.js';
+import type { ClientConnected, WaitingForMoreParticipants, EnoughParticipants } from '../messages.js';
 export interface NewDecentralizedNodeInfo {
     type: type.NewDecentralizedNodeInfo;
     id: NodeID;
@@ -12,21 +13,26 @@ export interface SignalForPeer {
     peer: NodeID;
     signal: SignalData;
 }
+export interface JoinRound {
+    type: type.JoinRound;
+}
 export interface PeerIsReady {
     type: type.PeerIsReady;
 }
 export interface PeersForRound {
     type: type.PeersForRound;
     peers: NodeID[];
+    aggregationRound: number;
 }
 export interface Payload {
     type: type.Payload;
     peer: NodeID;
-    round: number;
+    aggregationRound: number;
+    communicationRound: number;
     payload: weights.Encoded;
 }
-export type MessageFromServer = NewDecentralizedNodeInfo | SignalForPeer | PeersForRound;
-export type MessageToServer = ClientConnected | SignalForPeer | PeerIsReady;
+export type MessageFromServer = NewDecentralizedNodeInfo | SignalForPeer | PeersForRound | WaitingForMoreParticipants | EnoughParticipants;
+export type MessageToServer = ClientConnected | SignalForPeer | PeerIsReady | JoinRound;
 export type PeerMessage = Payload;
 export declare function isMessageFromServer(o: unknown): o is MessageFromServer;
 export declare function isMessageToServer(o: unknown): o is MessageToServer;

package/dist/client/decentralized/messages.js

@@ -2,9 +2,8 @@ import { weights } from '../../serialization/index.js';
 import { isNodeID } from '../types.js';
 import { type, hasMessageType } from '../messages.js';
 export function isMessageFromServer(o) {
-    if (!hasMessageType(o)) {
+    if (!hasMessageType(o))
         return false;
-    }
     switch (o.type) {
         case type.NewDecentralizedNodeInfo:
             return 'id' in o && isNodeID(o.id) &&
@@ -15,28 +14,30 @@ export function isMessageFromServer(o) {
                 'signal' in o; // TODO check signal content?
         case type.PeersForRound:
             return 'peers' in o && Array.isArray(o.peers) && o.peers.every(isNodeID);
+        case type.WaitingForMoreParticipants:
+        case type.EnoughParticipants:
+            return true;
     }
     return false;
 }
 export function isMessageToServer(o) {
-    if (!hasMessageType(o)) {
+    if (!hasMessageType(o))
         return false;
-    }
     switch (o.type) {
         case type.ClientConnected:
             return true;
         case type.SignalForPeer:
             return 'peer' in o && isNodeID(o.peer) &&
                 'signal' in o; // TODO check signal content?
+        case type.JoinRound:
         case type.PeerIsReady:
             return true;
     }
     return false;
 }
 export function isPeerMessage(o) {
-    if (!hasMessageType(o)) {
+    if (!hasMessageType(o))
         return false;
-    }
     switch (o.type) {
         case type.Payload:
             return ('peer' in o && isNodeID(o.peer) &&

package/dist/client/federated/federated_client.d.ts

@@ -6,25 +6,13 @@ import { Client } from "../client.js";
  */
 export declare class FederatedClient extends Client {
     #private;
-    get nbOfParticipants(): number;
-    /**
-     * Opens a new WebSocket connection with the server and listens to new messages over the channel
-     */
-    private connectServer;
+    getNbOfParticipants(): number;
     /**
      * Initializes the connection to the server, gets our node ID
      * as well as the latest training information: latest global model, current round and
      * whether we are waiting for more participants.
      */
     connect(): Promise<Model>;
-    /**
-     * Method called when the server notifies the client that there aren't enough
-     * participants (anymore) to start/continue training
-     * The method creates a promise that will resolve once the server notifies
-     * the client that the training can resume via a subsequent EnoughParticipants message
-     * @returns a promise which resolves when enough participants joined the session
-     */
-    private waitForMoreParticipants;
     /**
      * Disconnection process when user quits the task.
      */