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

package/dist/client/federated/federated_client.js

@@ -1,6 +1,6 @@
 import createDebug from "debug";
 import { serialization } from "../../index.js";
-import { Client } from "../client.js";
+import { Client, shortenId } from "../client.js";
 import { type } from "../messages.js";
 import { waitMessage, waitMessageWithTimeout, WebSocketServer, } from "../event_connection.js";
 import * as messages from "./messages.js";
@@ -19,38 +19,10 @@ export class FederatedClient extends Client {
     // Total number of other federated contributors, including this client, excluding the server
     // E.g., if 3 users are training a federated model, nbOfParticipants is 3
     #nbOfParticipants = 1;
-    /**
-     * 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
-     */
-    #promiseForMoreParticipants = undefined;
-    /**
-     * 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
-     */
-    #previousStatus = undefined;
-    /**
-    * Whether the client should wait until more
-    * participants join the session, i.e. a promise has been created
-    */
-    get #waitingForMoreParticipants() {
-        return this.#promiseForMoreParticipants !== undefined;
-    }
     // the number of participants excluding the server
-    get nbOfParticipants() {
+    getNbOfParticipants() {
         return this.#nbOfParticipants;
     }
-    /**
-     * Opens a new WebSocket connection with the server and listens to new messages over the channel
-     */
-    async connectServer(url) {
-        const server = await WebSocketServer.connect(url, messages.isMessageFederated, // can only receive federated message types from the server
-        messages.isMessageFederated);
-        return server;
-    }
     /**
      * Initializes the connection to the server, gets our node ID
      * as well as the latest training information: latest global model, current round and
@@ -70,31 +42,12 @@ export class FederatedClient extends Client {
                 throw new Error(`unknown protocol: ${this.url.protocol}`);
         }
         serverURL.pathname += `federated/${this.task.id}`;
-        this._server = await this.connectServer(serverURL);
-        // Setup an event callback if the server signals that we should 
-        // wait for more participants
-        this.server.on(type.WaitingForMoreParticipants, () => {
-            debug(`[${id.slice(0, 4)}] received WaitingForMoreParticipants message from server`);
-            // Display the waiting status right away
-            this.emit("status", "Waiting for more participants");
-            // Upon receiving a WaitingForMoreParticipants message,
-            // the client will await for this promise to resolve before sending its
-            // local weight update
-            this.#promiseForMoreParticipants = this.waitForMoreParticipants();
-        });
-        // 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 NewFederatedNodeInfo 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 NewFederatedNodeInfo (which is much bigger)
-        // so we check whether we received the EnoughParticipants before being assigned a node ID
+        // Opens a new WebSocket connection with the server and listens to new messages over the channel
+        this._server = await WebSocketServer.connect(serverURL, messages.isMessageFederated, // can only receive federated message types from the server
+        messages.isMessageFederated);
+        // c.f. setupServerCallbacks doc for explanation
         let receivedEnoughParticipants = false;
-        this.server.once(type.EnoughParticipants, () => {
-            if (this._ownId === undefined) {
-                debug(`Received EnoughParticipants message from server before the NewFederatedNodeInfo message`);
-                receivedEnoughParticipants = true;
-            }
-        });
+        this.setupServerCallbacks(() => receivedEnoughParticipants = true);
         this.aggregator.registerNode(SERVER_NODE_ID);
         const msg = {
             type: type.ClientConnected,
@@ -109,40 +62,21 @@ export class FederatedClient extends Client {
         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.waitForMoreParticipants();
+            this.promiseForMoreParticipants = this.createPromiseForMoreParticipants();
         }
         if (this._ownId !== undefined) {
             throw new Error('received id from server but was already received');
         }
         this._ownId = id;
-        debug(`[${id.slice(0, 4)}] joined session at round ${round} `);
+        debug(`[${shortenId(id)}] joined session at round ${round} `);
         this.aggregator.setRound(round);
         this.#nbOfParticipants = nbOfParticipants;
         // Upon connecting, the server answers with a boolean
         // which indicates whether there are enough participants or not
-        debug(`[${this.ownId.slice(0, 4)}] upon connecting, wait for participant flag %o`, this.#waitingForMoreParticipants);
+        debug(`[${shortenId(this.ownId)}] upon connecting, wait for participant flag %o`, this.waitingForMoreParticipants);
         model.weights = serialization.weights.decode(payload);
         return 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
-     */
-    async waitForMoreParticipants() {
-        return new Promise((resolve) => {
-            // "once" is important because we can't resolve the same promise multiple times
-            this.server.once(type.EnoughParticipants, () => {
-                debug(`[${this.ownId.slice(0, 4)}] received EnoughParticipants message from server`);
-                // Emit the last status emitted before waiting if defined
-                if (this.#previousStatus !== undefined)
-                    this.emit("status", this.#previousStatus);
-                resolve();
-            });
-        });
-    }
     /**
      * Disconnection process when user quits the task.
      */
@@ -155,10 +89,7 @@ export class FederatedClient extends Client {
     onRoundBeginCommunication() {
         // Prepare the result promise for the incoming round
         this.aggregationResult = new Promise((resolve) => this.aggregator.once('aggregation', resolve));
-        // 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.#previousStatus = "Training the model on the data you connected";
-        this.emit("status", this.#previousStatus);
+        this.saveAndEmit("local training");
         return Promise.resolve();
     }
     /**
@@ -176,18 +107,8 @@ export class FederatedClient extends Client {
             throw new Error("local aggregation result was not set");
         }
         // First 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(`[${this.ownId.slice(0, 4)}] is awaiting the promise for more participants`);
-            this.emit("status", "Waiting for more participants");
-            await this.#promiseForMoreParticipants;
-            // Make sure to set the promise back to undefined once resolved
-            this.#promiseForMoreParticipants = 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.#previousStatus = "Updating the model with other participants' models";
-        this.emit("status", this.#previousStatus);
+        await this.waitForParticipantsIfNeeded();
+        this.saveAndEmit("updating model");
         // Send our local contribution to the server
         // and receive the server global update for this round as an answer to our contribution
         const payloadToServer = this.aggregator.makePayloads(weights).first();
@@ -198,9 +119,9 @@ export class FederatedClient extends Client {
         };
         // Need to await the resulting global model right after sending our local contribution
         // to make sure we don't miss it
-        debug(`[${this.ownId.slice(0, 4)}] sent its local update to the server for round ${this.aggregator.round}`);
+        debug(`[${shortenId(this.ownId)}] sent its local update to the server for round ${this.aggregator.round}`);
         this.server.send(msg);
-        debug(`[${this.ownId.slice(0, 4)}] is waiting for server update for round ${this.aggregator.round + 1}`);
+        debug(`[${shortenId(this.ownId)}] is waiting for server update for round ${this.aggregator.round + 1}`);
         const { payload: payloadFromServer, round: serverRound, nbOfParticipants } = await waitMessage(this.server, type.ReceiveServerPayload); // Wait indefinitely for the server update
         this.#nbOfParticipants = nbOfParticipants; // Save the current participants
         const serverResult = serialization.weights.decode(payloadFromServer);

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

@@ -1,6 +1,7 @@
 import { type weights } from '../../serialization/index.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 type MessageFederated = ClientConnected | NewFederatedNodeInfo | SendPayload | ReceiveServerPayload | WaitingForMoreParticipants | EnoughParticipants;
 export interface NewFederatedNodeInfo {
     type: type.NewFederatedNodeInfo;
@@ -21,10 +22,4 @@ export interface ReceiveServerPayload {
     round: number;
     nbOfParticipants: number;
 }
-export interface EnoughParticipants {
-    type: type.EnoughParticipants;
-}
-export interface WaitingForMoreParticipants {
-    type: type.WaitingForMoreParticipants;
-}
 export declare function isMessageFederated(raw: unknown): raw is MessageFederated;

package/dist/client/local_client.d.ts

@@ -5,6 +5,7 @@ import { Client } from "./client.js";
  * with anyone. Thus LocalClient doesn't do anything during communication
  */
 export declare class LocalClient extends Client {
+    getNbOfParticipants(): number;
     onRoundBeginCommunication(): Promise<void>;
     onRoundEndCommunication(weights: WeightsContainer): Promise<WeightsContainer>;
 }

package/dist/client/local_client.js

@@ -4,6 +4,9 @@ import { Client } from "./client.js";
  * with anyone. Thus LocalClient doesn't do anything during communication
  */
 export class LocalClient extends Client {
+    getNbOfParticipants() {
+        return 1;
+    }
     onRoundBeginCommunication() {
         return Promise.resolve();
     }

package/dist/client/messages.d.ts

@@ -3,19 +3,26 @@ import type * as federated from './federated/messages.js';
 export declare enum type {
     ClientConnected = 0,
     NewDecentralizedNodeInfo = 1,
-    SignalForPeer = 2,
+    JoinRound = 2,
     PeerIsReady = 3,
     PeersForRound = 4,
-    Payload = 5,
-    NewFederatedNodeInfo = 6,
-    WaitingForMoreParticipants = 7,
-    EnoughParticipants = 8,
-    SendPayload = 9,
-    ReceiveServerPayload = 10
+    SignalForPeer = 5,
+    Payload = 6,
+    NewFederatedNodeInfo = 7,
+    WaitingForMoreParticipants = 8,
+    EnoughParticipants = 9,
+    SendPayload = 10,
+    ReceiveServerPayload = 11
 }
 export interface ClientConnected {
     type: type.ClientConnected;
 }
+export interface EnoughParticipants {
+    type: type.EnoughParticipants;
+}
+export interface WaitingForMoreParticipants {
+    type: type.WaitingForMoreParticipants;
+}
 export type Message = decentralized.MessageFromServer | decentralized.MessageToServer | decentralized.PeerMessage | federated.MessageFederated;
 export type NarrowMessage<D> = Extract<Message, {
     type: D;

package/dist/client/messages.js

@@ -9,34 +9,36 @@ export var type;
     // answers with  its peer id and also tells the client whether we are waiting
     // for more participants before starting training
     type[type["NewDecentralizedNodeInfo"] = 1] = "NewDecentralizedNodeInfo";
-    // Message forwarded by the server from a client to another client
-    // to establish a peer-to-peer (WebRTC) connection
-    type[type["SignalForPeer"] = 2] = "SignalForPeer";
+    // Message sent by peers to the server to signal they want to 
+    // join the next round
+    type[type["JoinRound"] = 2] = "JoinRound";
     // Message sent by nodes to server signaling they are ready to 
     // start the next round
     type[type["PeerIsReady"] = 3] = "PeerIsReady";
     // Sent by the server to participating peers containing the list
     // of peers for the round
     type[type["PeersForRound"] = 4] = "PeersForRound";
+    // Message forwarded by the server from a client to another client
+    // to establish a peer-to-peer (WebRTC) connection
+    type[type["SignalForPeer"] = 5] = "SignalForPeer";
     // The weight update
-    type[type["Payload"] = 5] = "Payload";
+    type[type["Payload"] = 6] = "Payload";
     /* Federated */
     // The server answers the ClientConnected message with the necessary information
     // to start training: node id, latest model global weights, current round etc
-    type[type["NewFederatedNodeInfo"] = 6] = "NewFederatedNodeInfo";
+    type[type["NewFederatedNodeInfo"] = 7] = "NewFederatedNodeInfo";
     // Message sent by server to notify clients that there are not enough
     // participants to continue training
-    type[type["WaitingForMoreParticipants"] = 7] = "WaitingForMoreParticipants";
+    type[type["WaitingForMoreParticipants"] = 8] = "WaitingForMoreParticipants";
     // Message sent by server to notify clients that there are now enough
     // participants to start training collaboratively
-    type[type["EnoughParticipants"] = 8] = "EnoughParticipants";
-    type[type["SendPayload"] = 9] = "SendPayload";
-    type[type["ReceiveServerPayload"] = 10] = "ReceiveServerPayload";
+    type[type["EnoughParticipants"] = 9] = "EnoughParticipants";
+    type[type["SendPayload"] = 10] = "SendPayload";
+    type[type["ReceiveServerPayload"] = 11] = "ReceiveServerPayload";
 })(type || (type = {}));
 export function hasMessageType(raw) {
-    if (typeof raw !== 'object' || raw === null) {
+    if (typeof raw !== 'object' || raw === null)
         return false;
-    }
     const o = raw;
     if (!('type' in o && typeof o.type === 'number' && o.type in type)) {
         return false;

package/dist/default_tasks/cifar10.js

@@ -29,8 +29,8 @@ export const cifar10 = {
                 IMAGE_W: 224,
                 LABEL_LIST: ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck'],
                 scheme: 'decentralized',
+                aggregationStrategy: 'mean',
                 privacy: { clippingRadius: 20, noiseScale: 1 },
-                decentralizedSecure: true,
                 minNbOfParticipants: 3,
                 maxShareValue: 100,
                 tensorBackend: 'tfjs'

package/dist/default_tasks/lus_covid.js

@@ -28,6 +28,7 @@ export const lusCovid = {
                 LABEL_LIST: ['COVID-Positive', 'COVID-Negative'],
                 dataType: 'image',
                 scheme: 'federated',
+                aggregationStrategy: 'mean',
                 minNbOfParticipants: 2,
                 tensorBackend: 'tfjs'
             }

package/dist/default_tasks/mnist.js

@@ -29,7 +29,7 @@ export const mnist = {
                 preprocessingFunctions: [data.ImagePreprocessing.Resize, data.ImagePreprocessing.Normalize],
                 LABEL_LIST: ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'],
                 scheme: 'decentralized',
-                decentralizedSecure: true,
+                aggregationStrategy: 'secure',
                 minNbOfParticipants: 3,
                 maxShareValue: 100,
                 tensorBackend: 'tfjs'

package/dist/default_tasks/simple_face.js

@@ -28,6 +28,7 @@ export const simpleFace = {
                 IMAGE_W: 200,
                 LABEL_LIST: ['child', 'adult'],
                 scheme: 'federated',
+                aggregationStrategy: 'mean',
                 minNbOfParticipants: 2,
                 tensorBackend: 'tfjs'
             }

package/dist/default_tasks/titanic.js

@@ -62,6 +62,7 @@ export const titanic = {
                     'Survived'
                 ],
                 scheme: 'federated',
+                aggregationStrategy: 'mean',
                 minNbOfParticipants: 2,
                 tensorBackend: 'tfjs'
             }

package/dist/default_tasks/wikitext.js

@@ -25,6 +25,7 @@ export const wikitext = {
                 dataType: 'text',
                 preprocessingFunctions: [data.TextPreprocessing.Tokenize, data.TextPreprocessing.LeftPadding],
                 scheme: 'federated',
+                aggregationStrategy: 'mean',
                 minNbOfParticipants: 2,
                 epochs: 6,
                 // Unused by wikitext because data already comes split

package/dist/task/training_information.d.ts

@@ -18,10 +18,9 @@ export interface TrainingInformation {
     LABEL_LIST?: string[];
     scheme: 'decentralized' | 'federated' | 'local';
     privacy?: Privacy;
-    decentralizedSecure?: boolean;
     maxShareValue?: number;
     minNbOfParticipants: number;
-    aggregator?: 'mean' | 'secure';
+    aggregationStrategy?: 'mean' | 'secure';
     tokenizer?: string | PreTrainedTokenizer;
     maxSequenceLength?: number;
     tensorBackend: 'tfjs' | 'gpt';

package/dist/task/training_information.js

@@ -24,7 +24,7 @@ export function isTrainingInformation(raw) {
     if (typeof raw !== 'object' || raw === null) {
         return false;
     }
-    const { IMAGE_H, IMAGE_W, LABEL_LIST, aggregator, batchSize, dataType, decentralizedSecure, privacy, epochs, inputColumns, maxShareValue, minNbOfParticipants, outputColumns, preprocessingFunctions, roundDuration, scheme, validationSplit, tokenizer, maxSequenceLength, tensorBackend } = raw;
+    const { IMAGE_H, IMAGE_W, LABEL_LIST, aggregationStrategy, batchSize, dataType, privacy, epochs, inputColumns, maxShareValue, minNbOfParticipants, outputColumns, preprocessingFunctions, roundDuration, scheme, validationSplit, tokenizer, maxSequenceLength, tensorBackend } = raw;
     if (typeof dataType !== 'string' ||
         typeof epochs !== 'number' ||
         typeof batchSize !== 'number' ||
@@ -33,8 +33,7 @@ export function isTrainingInformation(raw) {
         typeof minNbOfParticipants !== 'number' ||
         (tokenizer !== undefined && typeof tokenizer !== 'string' && !(tokenizer instanceof PreTrainedTokenizer)) ||
         (maxSequenceLength !== undefined && typeof maxSequenceLength !== 'number') ||
-        (aggregator !== undefined && typeof aggregator !== 'string') ||
-        (decentralizedSecure !== undefined && typeof decentralizedSecure !== 'boolean') ||
+        (aggregationStrategy !== undefined && typeof aggregationStrategy !== 'string') ||
         (privacy !== undefined && !isPrivacy(privacy)) ||
         (maxShareValue !== undefined && typeof maxShareValue !== 'number') ||
         (IMAGE_H !== undefined && typeof IMAGE_H !== 'number') ||
@@ -45,8 +44,8 @@ export function isTrainingInformation(raw) {
         (preprocessingFunctions !== undefined && !Array.isArray(preprocessingFunctions))) {
         return false;
     }
-    if (aggregator !== undefined) {
-        switch (aggregator) {
+    if (aggregationStrategy !== undefined) {
+        switch (aggregationStrategy) {
             case 'mean': break;
             case 'secure': break;
             default: return false;
@@ -58,7 +57,7 @@ export function isTrainingInformation(raw) {
         case 'text': break;
         default: return false;
     }
-    // interdepences on data type
+    // interdependencies on data type
     if (dataType === 'image') {
         if (typeof IMAGE_H !== 'number' || typeof IMAGE_W !== 'number') {
             return false;
@@ -87,10 +86,9 @@ export function isTrainingInformation(raw) {
         IMAGE_W,
         IMAGE_H,
         LABEL_LIST,
-        aggregator,
+        aggregationStrategy,
         batchSize,
         dataType,
-        decentralizedSecure,
         privacy,
         epochs,
         inputColumns,

package/dist/training/disco.d.ts

@@ -7,7 +7,10 @@ interface DiscoConfig {
     scheme: TrainingInformation["scheme"];
     logger: Logger;
 }
-export type RoundStatus = "Waiting for more participants" | "Retrieving peers' information" | "Updating the model with other participants' models" | "Training the model on the data you connected";
+export type RoundStatus = 'not enough participants' | // Server notification to wait for more participants
+'updating model' | // fetching/aggregating local updates into a global model
+'local training' | // Training the model locally
+'connecting to peers';
 /**
  * Top-level class handling distributed training from a client's perspective. It is meant to be
  * a convenient object providing a reduced yet complete API that wraps model training and

package/dist/training/trainer.js

@@ -62,7 +62,7 @@ export class Trainer {
         }
         return {
             epochs: epochsLogs,
-            participants: this.#client.nbOfParticipants,
+            participants: this.#client.getNbOfParticipants(),
         };
     }
 }

package/dist/utils/event_emitter.d.ts

@@ -1,11 +1,11 @@
-type Listener<T> = (_: T) => void;
+type Listener<T> = (_: T) => void | Promise<void>;
 /**
  * Call handlers on given events
  *
  * @typeParam I object/mapping from event name to emitted value type
  */
 export declare class EventEmitter<I extends Record<string, unknown>> {
-    private listeners;
+    #private;
     /**
      * @param initialListeners object/mapping of event name to listener, as if using `on` on created instance
      */
@@ -13,7 +13,7 @@ export declare class EventEmitter<I extends Record<string, unknown>> {
         [E in keyof I]?: Listener<I[E]>;
     });
     /**
-     * Register listener to call on event
+     * Register listener to call on event.
      *
      * @param event event name to listen to
      * @param listener handler to call

package/dist/utils/event_emitter.js

@@ -6,7 +6,8 @@ import { List } from 'immutable';
  * @typeParam I object/mapping from event name to emitted value type
  */
 export class EventEmitter {
-    listeners = {};
+    // List of callbacks to run per event
+    #listeners = {};
     /**
      * @param initialListeners object/mapping of event name to listener, as if using `on` on created instance
      */
@@ -19,14 +20,14 @@ export class EventEmitter {
         }
     }
     /**
-     * Register listener to call on event
+     * Register listener to call on event.
      *
      * @param event event name to listen to
      * @param listener handler to call
      */
     on(event, listener) {
-        const eventListeners = this.listeners[event] ?? List();
-        this.listeners[event] = eventListeners.push([false, listener]);
+        const eventListeners = this.#listeners[event] ?? List();
+        this.#listeners[event] = eventListeners.push([false, listener]);
     }
     /**
      * Register listener to call once on next event
@@ -35,8 +36,8 @@ export class EventEmitter {
      * @param listener handler to call next time
      */
     once(event, listener) {
-        const eventListeners = this.listeners[event] ?? List();
-        this.listeners[event] = eventListeners.push([true, listener]);
+        const eventListeners = this.#listeners[event] ?? List();
+        this.#listeners[event] = eventListeners.push([true, listener]);
     }
     /**
      * Send value to registered listeners of event name
@@ -45,9 +46,9 @@ export class EventEmitter {
      * @param value what to call listeners with
      */
     emit(event, value) {
-        const eventListeners = this.listeners[event] ?? List();
-        this.listeners[event] = eventListeners.filterNot(([once]) => once);
-        eventListeners.forEach(([_, listener]) => { listener(value); });
+        const eventListeners = this.#listeners[event] ?? List();
+        this.#listeners[event] = eventListeners.filterNot(([once]) => once);
+        eventListeners.forEach(async ([_, listener]) => { await listener(value); });
     }
 }
 /** `EventEmitter` for all events */

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@epfml/discojs",
-  "version": "3.0.1-p20241001093123.0",
+  "version": "3.0.1-p20241014092014.0",
   "type": "module",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",