Merge pull request #895 from epfml/894-participant-julien

Fix number of participant inconsistency

Author: JulienVig

discojs/src/client/client.ts

@@ -20,7 +20,10 @@ 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.
  */
-export abstract class Client extends EventEmitter<{'status': RoundStatus}>{
+export abstract class Client extends EventEmitter<{
+  'status': RoundStatus,
+  'participants': number
+}>{
   // Own ID provided by the network's server.
   protected _ownId?: NodeID
   // The network's server.
@@ -40,7 +43,10 @@ export abstract class Client extends EventEmitter<{'status': RoundStatus}>{
    * 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
    */
-  private previousStatus: RoundStatus | undefined;
+  #previousStatus: RoundStatus | undefined;
+
+  // Current number of participants including this client in the training session
+  #nbOfParticipants: number = 1;
 
   constructor (
     public readonly url: URL, // The network server's URL to connect to
@@ -82,7 +88,7 @@ export abstract class Client extends EventEmitter<{'status': RoundStatus}>{
    * the waiting status and once enough participants join, it can display the previous status again
    */ 
   protected saveAndEmit(status: RoundStatus) {
-    this.previousStatus = status
+    this.#previousStatus = status
     this.emit("status", status)
   }
 
@@ -111,12 +117,13 @@ export abstract class Client extends EventEmitter<{'status': RoundStatus}>{
   protected setupServerCallbacks(setMessageInversionFlag: () => void) {
     // Setup an event callback if the server signals that we should 
     // wait for more participants
-    this.server.on(type.WaitingForMoreParticipants, () => {
+    this.server.on(type.WaitingForMoreParticipants, (event) => {
       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")
+      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
       // local weight update
@@ -129,10 +136,10 @@ export abstract class Client extends EventEmitter<{'status': RoundStatus}>{
     // 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, () => {
+    this.server.once(type.EnoughParticipants, (event) => {
       if (this._ownId === undefined) {
-        debug(`Received EnoughParticipants message from server before the NewFederatedNodeInfo message`)
         setMessageInversionFlag()
+        this.nbOfParticipants = event.nbOfParticipants
       }
     })
   }
@@ -146,10 +153,11 @@ export abstract class Client extends EventEmitter<{'status': RoundStatus}>{
   protected async createPromiseForMoreParticipants(): Promise<void> {
     return new Promise<void>((resolve) => {
       // "once" is important because we can't resolve the same promise multiple times
-      this.server.once(type.EnoughParticipants, () => {
+      this.server.once(type.EnoughParticipants, (event) => {
         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)
+        if (this.#previousStatus !== undefined) this.emit("status", this.#previousStatus)
+        this.nbOfParticipants = event.nbOfParticipants
         resolve()
       })
     })
@@ -190,7 +198,18 @@ export abstract class Client extends EventEmitter<{'status': RoundStatus}>{
   * If federated, it should the number of participants excluding the server
   * If local it should be 1
   */
-  abstract getNbOfParticipants(): number;
+  public get nbOfParticipants(): number {
+    return this.#nbOfParticipants
+  }
+
+  /**
+   * Setter for the number of participants
+   * It emits the number of participants to the client
+   */
+  public set nbOfParticipants(nbOfParticipants: number) {
+    this.#nbOfParticipants = nbOfParticipants
+    this.emit("participants", nbOfParticipants)
+  }
 
   get ownId(): NodeID {
     if (this._ownId === undefined) {

discojs/src/client/decentralized/decentralized_client.ts

@@ -26,15 +26,16 @@ export class DecentralizedClient extends Client {
   #pool?: PeerPool
   #connections?: Map<NodeID, PeerConnection>
 
-  override getNbOfParticipants(): number {
-    const nbOfParticipants = this.aggregator.nodes.size
-    return nbOfParticipants === 0 ? 1 : nbOfParticipants
-  }
-
   // Used to handle timeouts and promise resolving after calling disconnect
   private get isDisconnected() : boolean {
     return this._server === undefined
   }
+
+  private setAggregatorNodes(nodes: Set<NodeID>) {
+    this.aggregator.setNodes(nodes)
+    // Emits the `participants` event
+    this.nbOfParticipants = this.aggregator.nodes.size === 0 ? 1 : this.aggregator.nodes.size
+  }
 
   /**
    * Public method called by disco.ts when starting training. This method sends
@@ -77,7 +78,11 @@ export class DecentralizedClient extends Client {
     }
     this.server.send(msg)
 
-    const { id, waitForMoreParticipants } = await waitMessage(this.server, type.NewDecentralizedNodeInfo)
+    const { id, waitForMoreParticipants,
+      nbOfParticipants } = await waitMessage(this.server, type.NewDecentralizedNodeInfo)
+    
+    this.nbOfParticipants = nbOfParticipants
+    
 
     // This should come right after receiving the message to make sure
     // we don't miss a subsequent message from the server
@@ -107,7 +112,7 @@ export class DecentralizedClient extends Client {
 
     if (this.#connections !== undefined) {
       const peers = this.#connections.keySeq().toSet()
-      this.aggregator.setNodes(this.aggregator.nodes.subtract(peers))
+      this.setAggregatorNodes(this.aggregator.nodes.subtract(peers))
     }
     // Disconnect from server
     await this.server?.disconnect()
@@ -180,7 +185,7 @@ export class DecentralizedClient extends Client {
         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.setAggregatorNodes(peers.add(this.ownId))
       this.aggregator.setRound(receivedMessage.aggregationRound) // the server gives us the round number
 
       // Initiate peer to peer connections with each peer
@@ -197,7 +202,7 @@ export class DecentralizedClient extends Client {
       this.#connections = connections
     } catch (e) {
       debug(`Error for [${shortenId(this.ownId)}] while beginning round: %o`, e);
-      this.aggregator.setNodes(Set(this.ownId))
+      this.setAggregatorNodes(Set(this.ownId))
       this.#connections = Map()
     }
   }

discojs/src/client/decentralized/messages.ts

@@ -11,6 +11,7 @@ export interface NewDecentralizedNodeInfo {
   type: type.NewDecentralizedNodeInfo
   id: NodeID
   waitForMoreParticipants: boolean
+  nbOfParticipants: number
 }
 
 // WebRTC signal to forward to other node

discojs/src/client/federated/federated_client.ts

@@ -24,14 +24,6 @@ const SERVER_NODE_ID = "federated-server-node-id";
  * a specific task in the federated setting.
  */
 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: number = 1;
-
-  // the number of participants excluding the server
-  override getNbOfParticipants(): number {
-    return this.#nbOfParticipants
-  }
 
   /**
    * Initializes the connection to the server, gets our node ID
@@ -92,7 +84,7 @@ export class FederatedClient extends Client {
     this._ownId = id;
     debug(`[${shortenId(id)}] joined session at round ${round} `);
     this.aggregator.setRound(round)
-    this.#nbOfParticipants = nbOfParticipants
+    this.nbOfParticipants = nbOfParticipants
     // Upon connecting, the server answers with a boolean
     // which indicates whether there are enough participants or not
     debug(`[${shortenId(this.ownId)}] upon connecting, wait for participant flag %o`, this.waitingForMoreParticipants)
@@ -161,8 +153,7 @@ export class FederatedClient extends Client {
       round: serverRound,
       nbOfParticipants
     } = await waitMessage( this.server, type.ReceiveServerPayload); // Wait indefinitely for the server update
-    
-    this.#nbOfParticipants = nbOfParticipants // Save the current participants
+    this.nbOfParticipants = nbOfParticipants // Save the current participants
     const serverResult = serialization.weights.decode(payloadFromServer);
     this.aggregator.setRound(serverRound);
 

discojs/src/client/local_client.ts

@@ -7,10 +7,6 @@ import { Client } from "./client.js";
  */
 export class LocalClient extends Client {
 
-  override getNbOfParticipants(): number {
-    return 1;
-  }
-
   override onRoundBeginCommunication(): Promise<void> {
     return Promise.resolve();
   }

discojs/src/client/messages.ts

@@ -47,10 +47,12 @@ export interface ClientConnected {
 
 export interface EnoughParticipants {
   type: type.EnoughParticipants
+  nbOfParticipants: number
 }
 
 export interface WaitingForMoreParticipants {
   type: type.WaitingForMoreParticipants
+  nbOfParticipants: number
 }
 
 export type Message =

discojs/src/processing/text.spec.ts

@@ -67,7 +67,7 @@ describe("text processing", () => {
     ]
     const tokens = tokenize(tokenizer, text);
     expect(tokens.toArray()).to.be.deep.equal(expectedTokens);
-  });
+  }).timeout(3000);
 
   it("can tokenize text with the GPT2 tokenizer", async () => {
     const tokenizer = await AutoTokenizer.from_pretrained("Xenova/gpt2");

discojs/src/training/disco.ts

@@ -49,6 +49,7 @@ export type RoundStatus = 'not enough participants' | // Server notification to
  */
 export class Disco<D extends DataType> extends EventEmitter<{
   status: RoundStatus;
+  participants: number
 }> {
   public readonly trainer: Trainer<D>;
   readonly #client: clients.Client;
@@ -99,6 +100,7 @@ export class Disco<D extends DataType> extends EventEmitter<{
     this.trainer = new Trainer(task, client);
     // Simply propagate the training status events emitted by the client
     this.#client.on("status", (status) => this.emit("status", status));
+    this.#client.on("participants", (nbParticipants) => this.emit("participants", nbParticipants));
   }
 
   /** Train on dataset, yielding logs of every round. */

discojs/src/training/trainer.ts

@@ -124,7 +124,7 @@ export class Trainer<D extends DataType> {
 
     return {
       epochs: epochsLogs,
-      participants: this.#client.getNbOfParticipants(),
+      participants: this.#client.nbOfParticipants,
     };
   }
 }

server/src/controllers/decentralized_controller.ts

@@ -50,6 +50,7 @@ export class DecentralizedController<
             const msg: messages.NewDecentralizedNodeInfo = {
               type: MessageTypes.NewDecentralizedNodeInfo,
               id: peerId,
+              nbOfParticipants: this.connections.size,
               waitForMoreParticipants: this.connections.size < minNbOfParticipants
             }
             ws.send(msgpack.encode(msg), { binary: true })