@epfml/discojs 3.0.1-p20240902100041.0 → 3.0.1-p20240904094219.0

package/dist/aggregator/base.d.ts

@@ -58,13 +58,22 @@ export declare abstract class Base<T> extends EventEmitter<{
     communicationRounds?: number);
     /**
      * Adds a node's contribution to the aggregator for the given aggregation and communication rounds.
+     * The aggregation round is increased whenever a new global model is obtained and local models are updated.
+     * Within one aggregation round there may be multiple communication rounds (such as for the decentralized secure aggregation
+     *  which requires multiple steps to obtain a global model)
      * The contribution will be aggregated during the next aggregation step.
      * @param nodeId The node's id
      * @param contribution The node's contribution
-     * @param round For which aggregation round the contribution was made
-     * @param communicationRound For which communication round the contribution was made
+     * @param round aggregation round of the contribution was made
+     * @param communicationRound communication round the contribution was made within the aggregation round
+     * @returns boolean, true if the contribution has been successfully taken into account or False if it has been rejected
      */
     abstract add(nodeId: client.NodeID, contribution: T, round: number, communicationRound?: number): boolean;
+    /**
+     * Evaluates whether a given participant contribution can be used in the current aggregation round
+     * the boolean returned by `this.add` is obtained via `this.isValidContribution`
+     */
+    isValidContribution(nodeId: client.NodeID, round: number): boolean;
     /**
      * Performs an aggregation step over the received node contributions.
      * Must store the aggregation's result in the aggregator's result promise.
@@ -91,6 +100,11 @@ export declare abstract class Base<T> extends EventEmitter<{
      * @returns True is the node wasn't already in the list of nodes, False if already included
      */
     registerNode(nodeId: client.NodeID): boolean;
+    /**
+     * Remove a node's id from the set of active nodes.
+     * @param nodeId The node to be removed
+     */
+    removeNode(nodeId: client.NodeID): void;
     /**
      * Overwrites the current set of active nodes with the given one. A node represents
      * an active neighbor peer/client within the network, whom we are communicating with

package/dist/aggregator/base.js

@@ -60,6 +60,21 @@ export class Base extends EventEmitter {
         // and communication rounds.
         this.on('aggregation', () => this.nextRound());
     }
+    /**
+     * Evaluates whether a given participant contribution can be used in the current aggregation round
+     * the boolean returned by `this.add` is obtained via `this.isValidContribution`
+     */
+    isValidContribution(nodeId, round) {
+        if (!this.nodes.has(nodeId)) {
+            debug("Contribution rejected because node id is not registered");
+            return false;
+        }
+        if (!this.isWithinRoundCutoff(round)) {
+            debug(`Contribution rejected because round ${round} is not within round cutoff`);
+            return false;
+        }
+        return true;
+    }
     /**
      * Returns whether the given round is recent enough, dependent on the
      * aggregator's round cutoff.
@@ -77,16 +92,16 @@ export class Base extends EventEmitter {
     log(step, from) {
         switch (step) {
             case AggregationStep.ADD:
-                debug(`adding contribution from node ${from ?? '"unknown"'} for round (${this.communicationRound}, ${this.round})`);
+                debug(`Adding contribution from node ${from ?? '"unknown"'} for aggregation round ${this.round} and communication round ${this.communicationRound}`);
                 break;
             case AggregationStep.UPDATE:
                 if (from === undefined) {
                     return;
                 }
-                debug(`updating contribution from node ${from ?? '"unknown"'} for round (${this.communicationRound}, ${this.round})`);
+                debug(`Updating contribution from node ${from} for aggregation round ${this.round} and communication round ${this.communicationRound}`);
                 break;
             case AggregationStep.AGGREGATE:
-                debug(`buffer full, aggregating weights for round (${this.communicationRound}, ${this.round})`);
+                debug(`Buffer is full. Aggregating weights for round aggregation round ${this.round} and communication round ${this.communicationRound}`);
                 break;
             default: {
                 const _ = step;
@@ -108,6 +123,13 @@ export class Base extends EventEmitter {
         }
         return false;
     }
+    /**
+     * Remove a node's id from the set of active nodes.
+     * @param nodeId The node to be removed
+     */
+    removeNode(nodeId) {
+        this._nodes = this._nodes.delete(nodeId);
+    }
     /**
      * Overwrites the current set of active nodes with the given one. A node represents
      * an active neighbor peer/client within the network, whom we are communicating with

package/dist/aggregator/mean.d.ts

@@ -30,6 +30,7 @@ export declare class MeanAggregator extends Aggregator<WeightsContainer> {
     constructor(roundCutoff?: number, threshold?: number, thresholdType?: ThresholdType);
     /** Checks whether the contributions buffer is full. */
     isFull(): boolean;
+    set minNbOfParticipants(minNbOfParticipants: number);
     add(nodeId: client.NodeID, contribution: WeightsContainer, round: number, currentContributions?: number): boolean;
     aggregate(): void;
     makePayloads(weights: WeightsContainer): Map<client.NodeID, WeightsContainer>;

package/dist/aggregator/mean.js

@@ -9,6 +9,7 @@ const debug = createDebug("discojs:aggregator:mean");
 export class MeanAggregator extends Aggregator {
     #threshold;
     #thresholdType;
+    #minNbOfParticipants;
     /**
      * Create a mean aggregator that averages all weight updates received when a specified threshold is met.
      * By default, initializes an aggregator that waits for 100% of the nodes' contributions and that
@@ -65,21 +66,24 @@ export class MeanAggregator extends Aggregator {
     }
     /** Checks whether the contributions buffer is full. */
     isFull() {
+        // Make sure that we are over the minimum number of participants
+        // if specified
+        if (this.#minNbOfParticipants !== undefined &&
+            this.nodes.size < this.#minNbOfParticipants)
+            return false;
         const thresholdValue = this.#thresholdType == 'relative'
             ? this.#threshold * this.nodes.size
             : this.#threshold;
         return (this.contributions.get(0)?.size ?? 0) >= thresholdValue;
     }
+    set minNbOfParticipants(minNbOfParticipants) {
+        this.#minNbOfParticipants = minNbOfParticipants;
+    }
     add(nodeId, contribution, round, currentContributions = 0) {
         if (currentContributions !== 0)
             throw new Error("only a single communication round");
-        if (!this.nodes.has(nodeId) || !this.isWithinRoundCutoff(round)) {
-            if (!this.nodes.has(nodeId))
-                debug(`contribution rejected because node ${nodeId} is not registered`);
-            if (!this.isWithinRoundCutoff(round))
-                debug(`contribution rejected because round ${round} is not within cutoff`);
+        if (!this.isValidContribution(nodeId, round))
             return false;
-        }
         this.log(this.contributions.hasIn([0, nodeId])
             ? AggregationStep.UPDATE
             : AggregationStep.ADD, nodeId);
@@ -94,6 +98,7 @@ export class MeanAggregator extends Aggregator {
             throw new Error("aggregating without any contribution");
         this.log(AggregationStep.AGGREGATE);
         const result = aggregation.avg(currentContributions.values());
+        // Emitting the event runs the superclass' callback to increment the round
         this.emit('aggregation', result);
     }
     makePayloads(weights) {

package/dist/aggregator/secure.js

@@ -48,7 +48,7 @@ export class SecureAggregator extends Aggregator {
             default:
                 throw new Error("requires communication round to be 0 or 1");
         }
-        if (!this.nodes.has(nodeId) || !this.isWithinRoundCutoff(round))
+        if (!this.isValidContribution(nodeId, round))
             return false;
         this.log(this.contributions.hasIn([communicationRound, nodeId])
             ? AggregationStep.UPDATE

package/dist/client/{base.d.ts → client.d.ts}

@@ -1,23 +1,17 @@
-import type { Model, Task, WeightsContainer } from '../index.js';
+import type { Model, Task, WeightsContainer, RoundStatus } from '../index.js';
 import type { NodeID } from './types.js';
 import type { EventConnection } from './event_connection.js';
 import type { Aggregator } from '../aggregator/index.js';
+import { EventEmitter } from '../utils/event_emitter.js';
 /**
  * Main, abstract, class representing a Disco client in a network, which handles
  * communication with other nodes, be it peers or a server.
  */
-export declare abstract class Base {
-    /**
-     * The network server's URL to connect to.
-     */
+export declare abstract class Client extends EventEmitter<{
+    'status': RoundStatus;
+}> {
     readonly url: URL;
-    /**
-     * The client's corresponding task.
-     */
     readonly task: Task;
-    /**
-     * The client's aggregator.
-     */
     readonly aggregator: Aggregator;
     /**
      * Own ID provided by the network's server.
@@ -31,23 +25,15 @@ export declare abstract class Base {
      * The aggregator's result produced after aggregation.
      */
     protected aggregationResult?: Promise<WeightsContainer>;
-    constructor(
-    /**
-     * The network server's URL to connect to.
-     */
-    url: URL, 
-    /**
-     * The client's corresponding task.
-     */
-    task: Task, 
-    /**
-     * The client's aggregator.
-     */
+    constructor(url: URL, // The network server's URL to connect to
+    task: Task, // The client's corresponding task
     aggregator: Aggregator);
     /**
      * Handles the connection process from the client to any sort of network server.
+     * This method is overriden by the federated and decentralized clients
+     * By default, it fetches and returns the server's base model
      */
-    connect(): Promise<void>;
+    connect(): Promise<Model>;
     /**
      * Handles the disconnection process of the client from any sort of network server.
      */
@@ -59,17 +45,14 @@ export declare abstract class Base {
     getLatestModel(): Promise<Model>;
     /**
      * Communication callback called at the beginning of every training round.
-     * @param _weights The most recent local weight updates
-     * @param _round The current training round
      */
-    onRoundBeginCommunication(_weights: WeightsContainer, _round: number): Promise<void>;
+    abstract onRoundBeginCommunication(): Promise<void>;
     /**
      * Communication callback called the end of every training round.
-     * @param _weights The most recent local weight updates
-     * @param _round The current training round
+     * @param weights The local weight update resulting for the current local training round
      * @returns aggregated weights or the local weights upon error
      */
-    abstract onRoundEndCommunication(_weights: WeightsContainer, _round: number): Promise<WeightsContainer>;
+    abstract onRoundEndCommunication(weights: WeightsContainer): Promise<WeightsContainer>;
     get nbOfParticipants(): number;
     get ownId(): NodeID;
     get server(): EventConnection;

package/dist/client/{base.js → client.js}

@@ -1,10 +1,11 @@
 import axios from 'axios';
 import { serialization } from '../index.js';
+import { EventEmitter } from '../utils/event_emitter.js';
 /**
  * Main, abstract, class representing a Disco client in a network, which handles
  * communication with other nodes, be it peers or a server.
  */
-export class Base {
+export class Client extends EventEmitter {
     url;
     task;
     aggregator;
@@ -20,27 +21,22 @@ export class Base {
      * The aggregator's result produced after aggregation.
      */
     aggregationResult;
-    constructor(
-    /**
-     * The network server's URL to connect to.
-     */
-    url, 
-    /**
-     * The client's corresponding task.
-     */
-    task, 
-    /**
-     * The client's aggregator.
-     */
+    constructor(url, // The network server's URL to connect to
+    task, // The client's corresponding task
     aggregator) {
+        super();
         this.url = url;
         this.task = task;
         this.aggregator = aggregator;
     }
     /**
      * Handles the connection process from the client to any sort of network server.
+     * This method is overriden by the federated and decentralized clients
+     * By default, it fetches and returns the server's base model
      */
-    async connect() { }
+    async connect() {
+        return this.getLatestModel();
+    }
     /**
      * Handles the disconnection process of the client from any sort of network server.
      */
@@ -58,12 +54,6 @@ export class Base {
         const response = await axios.get(url.href, { responseType: 'arraybuffer' });
         return await serialization.model.decode(new Uint8Array(response.data));
     }
-    /**
-     * Communication callback called at the beginning of every training round.
-     * @param _weights The most recent local weight updates
-     * @param _round The current training round
-     */
-    async onRoundBeginCommunication(_weights, _round) { }
     // 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

package/dist/client/decentralized/{base.d.ts → decentralized_client.d.ts}

@@ -1,4 +1,4 @@
-import type { WeightsContainer } from "../../index.js";
+import type { Model, WeightsContainer } from "../../index.js";
 import { Client } from '../index.js';
 /**
  * Represents a decentralized client in a network of peers. Peers coordinate each other with the
@@ -6,7 +6,7 @@ import { Client } from '../index.js';
  * with the server is based off regular WebSockets, whereas peer-to-peer communication uses
  * WebRTC for Node.js.
  */
-export declare class Base extends Client {
+export declare class DecentralizedClient extends Client {
     /**
      * The pool of peers to communicate with during the current training round.
      */
@@ -21,7 +21,7 @@ export declare class Base extends Client {
      * create peer-to-peer WebRTC connections with peers. The server is used to exchange
      * peers network information.
      */
-    connect(): Promise<void>;
+    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.
@@ -37,7 +37,7 @@ export declare class Base extends Client {
      * and waits for it to resolve.
      *
      */
-    onRoundBeginCommunication(_: WeightsContainer, round: number): Promise<void>;
+    onRoundBeginCommunication(): Promise<void>;
     /**
      * 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 +45,5 @@ export declare class Base extends Client {
      * @param round
      */
     private receivePayloads;
-    onRoundEndCommunication(weights: WeightsContainer, round: number): Promise<WeightsContainer>;
+    onRoundEndCommunication(weights: WeightsContainer): Promise<WeightsContainer>;
 }

package/dist/client/decentralized/{base.js → decentralized_client.js}

@@ -14,7 +14,7 @@ const debug = createDebug("discojs:client:decentralized");
  * with the server is based off regular WebSockets, whereas peer-to-peer communication uses
  * WebRTC for Node.js.
  */
-export class Base extends Client {
+export class DecentralizedClient extends Client {
     /**
      * The pool of peers to communicate with during the current training round.
      */
@@ -33,6 +33,7 @@ export class Base extends Client {
      * peers network information.
      */
     async connect() {
+        const model = await super.connect(); // Get the server base model
         const serverURL = new URL('', this.url.href);
         switch (this.url.protocol) {
             case 'http:':
@@ -44,19 +45,20 @@ export class Base extends Client {
             default:
                 throw new Error(`unknown protocol: ${this.url.protocol}`);
         }
-        serverURL.pathname += `deai/${this.task.id}`;
+        serverURL.pathname += `decentralized/${this.task.id}`;
         this._server = await this.connectServer(serverURL);
         const msg = {
             type: type.ClientConnected
         };
         this.server.send(msg);
-        const peerIdMsg = await waitMessage(this.server, type.AssignNodeID);
-        debug(`[${peerIdMsg.id}] assigned id generated by server`);
+        const { id } = await waitMessage(this.server, type.NewDecentralizedNodeInfo);
+        debug(`[${id}] assigned id generated by server`);
         if (this._ownId !== undefined) {
             throw new Error('received id from server but was already received');
         }
-        this._ownId = peerIdMsg.id;
-        this.pool = new PeerPool(peerIdMsg.id);
+        this._ownId = id;
+        this.pool = new PeerPool(id);
+        return model;
     }
     /**
      * Create a WebSocket connection with the server
@@ -96,21 +98,22 @@ export class Base extends Client {
      * and waits for it to resolve.
      *
      */
-    async onRoundBeginCommunication(_, round) {
+    async onRoundBeginCommunication() {
         if (this.server === undefined) {
             throw new Error("peer's server is undefined, make sure to call `client.connect()` first");
         }
         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.
-        // this.connections = await this.waitForPeers(round)
         // Tell the server we are ready for the next round
         const readyMessage = { type: type.PeerIsReady };
         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");
             const peers = Set(receivedMessage.peers);
             if (this.ownId !== undefined && peers.has(this.ownId)) {
@@ -124,18 +127,19 @@ export class Base extends Client {
             // Init receipt of peers weights
             // this awaits the peer's weight update and adds it to 
             // our aggregator upon reception
-            (conn) => { this.receivePayloads(conn, round); });
-            debug(`[${this.ownId}] received peers for round ${round}: %o`, connections.keySeq().toJS());
+            (conn) => { this.receivePayloads(conn, this.aggregator.round); });
+            debug(`[${this.ownId}] received peers for round ${this.aggregator.round}: %o`, connections.keySeq().toJS());
             this.connections = connections;
         }
         catch (e) {
-            debug(`[${this.ownId}] while beginning round: %o`, e);
+            debug(`Error for [${this.ownId}] while beginning round: %o`, e);
             this.aggregator.setNodes(Set(this.ownId));
             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.
@@ -156,18 +160,18 @@ export class Base extends Client {
                     }
                 }
                 catch (e) {
-                    if (this.isDisconnected) {
+                    if (this.isDisconnected)
                         return;
-                    }
-                    debug(`[${this.ownId}] while receiving payloads: %o`, e);
+                    debug(`Error for [${this.ownId}] while receiving payloads: %o`, e);
                 }
             } while (++currentCommunicationRounds < this.aggregator.communicationRounds);
         });
     }
-    async onRoundEndCommunication(weights, round) {
+    async onRoundEndCommunication(weights) {
         if (this.aggregationResult === undefined) {
             throw new TypeError('aggregation result promise is undefined');
         }
+        this.emit("status", "Updating the model with other participants' models");
         // 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
@@ -181,7 +185,7 @@ export class Base extends Client {
                 try {
                     await Promise.all(payloads.map(async (payload, id) => {
                         if (id === this.ownId) {
-                            this.aggregator.add(this.ownId, payload, round, r);
+                            this.aggregator.add(this.ownId, payload, this.aggregator.round, r);
                         }
                         else {
                             const peer = this.connections?.get(id);

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

@@ -1,2 +1,2 @@
-export { Base as DecentralizedClient } from './base.js';
+export { DecentralizedClient } from './decentralized_client.js';
 export * as messages from './messages.js';

package/dist/client/decentralized/index.js

@@ -1,2 +1,2 @@
-export { Base as DecentralizedClient } from './base.js';
+export { DecentralizedClient } from './decentralized_client.js';
 export * as messages from './messages.js';

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

@@ -1,7 +1,12 @@
 import { weights } from '../../serialization/index.js';
 import { type SignalData } from './peer.js';
 import { type NodeID } from '../types.js';
-import { type, type ClientConnected, type AssignNodeID } from '../messages.js';
+import { type, type ClientConnected } from '../messages.js';
+export interface NewDecentralizedNodeInfo {
+    type: type.NewDecentralizedNodeInfo;
+    id: NodeID;
+    waitForMoreParticipants: boolean;
+}
 export interface SignalForPeer {
     type: type.SignalForPeer;
     peer: NodeID;
@@ -20,7 +25,7 @@ export interface Payload {
     round: number;
     payload: weights.Encoded;
 }
-export type MessageFromServer = AssignNodeID | SignalForPeer | PeersForRound;
+export type MessageFromServer = NewDecentralizedNodeInfo | SignalForPeer | PeersForRound;
 export type MessageToServer = ClientConnected | SignalForPeer | PeerIsReady;
 export type PeerMessage = Payload;
 export declare function isMessageFromServer(o: unknown): o is MessageFromServer;

package/dist/client/decentralized/messages.js

@@ -6,8 +6,10 @@ export function isMessageFromServer(o) {
         return false;
     }
     switch (o.type) {
-        case type.AssignNodeID:
-            return 'id' in o && isNodeID(o.id);
+        case type.NewDecentralizedNodeInfo:
+            return 'id' in o && isNodeID(o.id) &&
+                'waitForMoreParticipants' in o &&
+                typeof o.waitForMoreParticipants === 'boolean';
         case type.SignalForPeer:
             return 'peer' in o && isNodeID(o.peer) &&
                 'signal' in o; // TODO check signal content?

package/dist/client/event_connection.js

@@ -95,8 +95,8 @@ export class WebSocketServer extends EventEmitter {
     }
     disconnect() {
         return new Promise((resolve, reject) => {
-            this.socket.once('close', resolve);
-            this.socket.once('error', reject);
+            this.socket.onclose = () => resolve();
+            this.socket.onerror = (e) => reject(e.message);
             this.socket.close();
         });
     }

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

@@ -0,0 +1,44 @@
+import type { Model, WeightsContainer } from "../../index.js";
+import { Client } from "../client.js";
+/**
+ * Client class that communicates with a centralized, federated server, when training
+ * a specific task in the federated setting.
+ */
+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;
+    /**
+     * 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.
+     */
+    disconnect(): Promise<void>;
+    onRoundBeginCommunication(): Promise<void>;
+    /**
+     * Send the local weight update to the server and waits (indefinitely) for the server global update
+     *
+     * If the waitingForMoreParticipants flag is set, we first wait (also indefinitely) until the
+     * server notifies us that the training can resume.
+     *
+    // NB: For now, we suppose a fully-federated setting.
+     * @param weights Local weights sent to the server at the end of the local training round
+     * @returns the new global weights sent by the server
+     */
+    onRoundEndCommunication(weights: WeightsContainer): Promise<WeightsContainer>;
+}