Update number of participants in real-time

discojs/src/client/client.ts

@@ -49,6 +49,11 @@ export abstract class Client<N extends Network> extends EventEmitter<{
 
   // Current number of participants including this client in the training session
   #nbOfParticipants: number = 1;
+  /**
+   * Whether the server sent us a more recent number of participants than the
+   * one carried by the message answering our join request.
+   */
+  #nbOfParticipantsUpdatedSinceJoining = false;
 
   constructor(
     public readonly url: URL, // The network server's URL to connect to
@@ -119,6 +124,8 @@ export abstract class Client<N extends Network> extends EventEmitter<{
    * between a NewNodeInfo message and an EnoughParticipant message.
    */
   protected setupServerCallbacks(setMessageInversionFlag: () => void) {
+    this.#nbOfParticipantsUpdatedSinceJoining = false;
+
     // Setup an event callback if the server signals that we should
     // wait for more participants
     this.server.on(MType.WaitingForMoreParticipants, (event) => {
@@ -131,13 +138,22 @@ export abstract class Client<N extends Network> extends EventEmitter<{
       );
       // Display the waiting status right away
       this.emit("status", "not enough participants");
+      this.#nbOfParticipantsUpdatedSinceJoining = true;
       this.nbOfParticipants = event.nbOfParticipants; // emits the `participants` event
       // Upon receiving a WaitingForMoreParticipants message,
       // the client will await for this promise to resolve before sending its
       // local weight update
       this.promiseForMoreParticipants = this.createPromiseForMoreParticipants();
     });
 
+    // The server tells us whenever a participant joined or left, so that we
+    // don't have to wait for the end of the round to display how many of us
+    // are training together
+    this.server.on(MType.ParticipantsUpdate, (event) => {
+      this.#nbOfParticipantsUpdatedSinceJoining = true;
+      this.nbOfParticipants = event.nbOfParticipants;
+    });
+
     // As an example assume we need at least 2 participants to train,
     // When two participants join almost at the same time, the server
     // sends a NewNodeInfo with waitForMoreParticipants=true to the first participant
@@ -147,10 +163,25 @@ export abstract class Client<N extends Network> extends EventEmitter<{
     this.server.once(MType.EnoughParticipants, (event) => {
       if (this._ownId === undefined) {
         setMessageInversionFlag();
+        this.#nbOfParticipantsUpdatedSinceJoining = true;
         this.nbOfParticipants = event.nbOfParticipants;
       }
     });
   }
+
+  /**
+   * Set the number of participants carried by the message answering our join
+   * request.
+   *
+   * The server computes that count when sending its answer, which we can
+   * process much later as the message also carries the model weights. Any
+   * WaitingForMoreParticipants or EnoughParticipants message processed in the
+   * meantime carries a more recent count, which we keep instead.
+   */
+  protected setNbOfParticipantsUponJoining(nbOfParticipants: number): void {
+    if (this.#nbOfParticipantsUpdatedSinceJoining) return;
+    this.nbOfParticipants = nbOfParticipants;
+  }
   /**
    * Method called when the server notifies the client that there aren't enough
    * participants (anymore) to start/continue training

discojs/src/client/decentralized/decentralized_client.ts

@@ -54,13 +54,6 @@ export class DecentralizedClient extends Client<"decentralized"> {
     return this._server === undefined;
   }
 
-  private setAggregatorNodes(nodes: Set<NodeID>) {
-    this.aggregator.setNodes(nodes);
-    // Emits the `participants` event
-    this.nbOfParticipants =
-      this.aggregator.nodes.size === 0 ? 1 : this.aggregator.nodes.size;
-  }
-
   private cloneWeights(weights: WeightsContainer): WeightsContainer {
     return new WeightsContainer(weights.weights.map((t) => t.clone()));
   }
@@ -156,7 +149,7 @@ export class DecentralizedClient extends Client<"decentralized"> {
       await waitMessage(this.server, MType.NewDecentralizedNodeInfo);
 
     this.#modelSyncNeeded = joinedMidTraining;
-    this.nbOfParticipants = nbOfParticipants;
+    this.setNbOfParticipantsUponJoining(nbOfParticipants);
 
     // This should come right after receiving the message to make sure
     // we don't miss a subsequent message from the server
@@ -186,7 +179,7 @@ export class DecentralizedClient extends Client<"decentralized"> {
 
     if (this.#connections !== undefined) {
       const peers = this.#connections.keySeq().toSet();
-      this.setAggregatorNodes(this.aggregator.nodes.subtract(peers));
+      this.aggregator.setNodes(this.aggregator.nodes.subtract(peers));
     }
     // Disconnect from server
     await this.server?.disconnect();
@@ -311,9 +304,10 @@ export class DecentralizedClient extends Client<"decentralized"> {
         // clear the communication round peer pool
         await this.#pool?.shutdown();
         this.#pool = new PeerPool(this.ownId);
-        // clear the connections
+        // clear the connections. The peers are still part of the session,
+        // so the number of participants is left as the server reported it
         this.#connections = Map();
-        this.setAggregatorNodes(Set(this.ownId));
+        this.aggregator.setNodes(Set([this.ownId]));
         continue;
       } else if (msg.type === MType.ConnectionFail) {
         debug(`[${shortenId(this.ownId)}] disconnect from the server`);
@@ -369,7 +363,10 @@ export class DecentralizedClient extends Client<"decentralized"> {
         throw new Error("received peer list contains our own id");
       }
       // Store the list of peers for the current round including ourselves
-      this.setAggregatorNodes(peers.add(this.ownId));
+      this.aggregator.setNodes(peers.add(this.ownId));
+      // the peers of the round leave out those still syncing their model, so
+      // the server tells us how many participants the session has
+      this.nbOfParticipants = receivedMessage.nbOfParticipants;
       this.aggregator.setRound(receivedMessage.aggregationRound); // the server gives us the round number
 
       // Initiate peer to peer connections with each peer
@@ -394,7 +391,9 @@ export class DecentralizedClient extends Client<"decentralized"> {
         `Error for [${shortenId(this.ownId)}] while beginning round: %o`,
         e,
       );
-      this.setAggregatorNodes(Set(this.ownId));
+      // we can't aggregate with peers we failed to connect to, but they are
+      // still part of the session: leave the number of participants alone
+      this.aggregator.setNodes(Set([this.ownId]));
       this.#connections = Map();
     }
   }

discojs/src/client/decentralized/messages.ts

@@ -7,6 +7,7 @@ import type {
   ClientConnected,
   WaitingForMoreParticipants,
   EnoughParticipants,
+  ParticipantsUpdate,
 } from "#client/mtype";
 
 /// Phase 0 communication (between server and peers)
@@ -40,6 +41,9 @@ export interface PeersForRound {
   type: MType.PeersForRound;
   peers: NodeID[];
   aggregationRound: number;
+  // the peers of a round are only those able to take part in it, so they
+  // aren't the number of participants of the session
+  nbOfParticipants: number;
 }
 
 // peer sends to server to signal all the connections to other peers
@@ -107,6 +111,7 @@ export type MessageFromServer =
   | PeersForRound
   | WaitingForMoreParticipants
   | EnoughParticipants
+  | ParticipantsUpdate
   | StartWeightSharing
   | RetryPeerConnections
   | ConnectionFail
@@ -140,6 +145,7 @@ export function isMessageFromServer(o: unknown): o is MessageFromServer {
       return "peers" in o && Array.isArray(o.peers) && o.peers.every(isNodeID);
     case MType.WaitingForMoreParticipants:
     case MType.EnoughParticipants:
+    case MType.ParticipantsUpdate:
     case MType.StartWeightSharing:
     case MType.RetryPeerConnections:
     case MType.ConnectionFail:

discojs/src/client/federated/federated_client.ts

@@ -85,7 +85,7 @@ export class FederatedClient extends Client<"federated"> {
     this._ownId = id;
     debug(`[${shortenId(id)}] joined session at round ${round} `);
     this.aggregator.setRound(round);
-    this.nbOfParticipants = nbOfParticipants;
+    this.setNbOfParticipantsUponJoining(nbOfParticipants);
     // Upon connecting, the server answers with a boolean
     // which indicates whether there are enough participants or not
     debug(

discojs/src/client/federated/messages.ts

@@ -6,6 +6,7 @@ import type {
   ClientConnected,
   WaitingForMoreParticipants,
   EnoughParticipants,
+  ParticipantsUpdate,
 } from "#client/mtype";
 
 // See ../messages.ts for doc
@@ -15,7 +16,8 @@ export type MessageFederated =
   | SendPayload
   | ReceiveServerPayload
   | WaitingForMoreParticipants
-  | EnoughParticipants;
+  | EnoughParticipants
+  | ParticipantsUpdate;
 
 export interface NewFederatedNodeInfo {
   type: MType.NewFederatedNodeInfo;
@@ -50,6 +52,7 @@ export function isMessageFederated(raw: unknown): raw is MessageFederated {
     case MType.ReceiveServerPayload:
     case MType.WaitingForMoreParticipants:
     case MType.EnoughParticipants:
+    case MType.ParticipantsUpdate:
       return true;
   }
 

discojs/src/client/mtype.ts

@@ -52,6 +52,12 @@ export enum MType {
   EnoughParticipants,
   SendPayload,
   ReceiveServerPayload,
+
+  /* Both schemes */
+  // Message sent by the server when a participant joined or left, so that
+  // clients don't have to wait for the end of the round to learn about it.
+  // Kept last as the enum values are what goes over the wire.
+  ParticipantsUpdate,
 }
 
 export function hasMessageType(
@@ -80,3 +86,8 @@ export interface WaitingForMoreParticipants {
   type: MType.WaitingForMoreParticipants;
   nbOfParticipants: number;
 }
+
+export interface ParticipantsUpdate {
+  type: MType.ParticipantsUpdate;
+  nbOfParticipants: number;
+}

server/src/controllers/decentralized_controller.ts

@@ -78,8 +78,11 @@ export class DecentralizedController<
               joinedMidTraining: joinedMidTraining,
             };
             ws.send(msgpack.encode(msg), { binary: true });
-            // Send an update to participants if we can start/resume training
-            this.sendEnoughParticipantsMsgIfNeeded(peerId);
+            // Send an update to participants if we can start/resume training,
+            // which already carries the number of participants
+            if (!this.sendEnoughParticipantsMsgIfNeeded(peerId))
+              // otherwise just tell them that someone joined
+              this.sendParticipantsUpdateMsg(peerId);
             break;
           }
           // Send by peers at the beginning of each training round to notify
@@ -188,14 +191,20 @@ export class DecentralizedController<
       }
 
       // Check if we are already waiting for new participants to join
-      if (this.waitingForMoreParticipants) return;
+      if (this.waitingForMoreParticipants) {
+        // tell the remaining participants that one of them left
+        this.sendParticipantsUpdateMsg();
+        return;
+      }
       // If no, check if we are still above the minimum number of participant required
       if (this.connections.size >= minNbOfParticipants) {
+        this.sendParticipantsUpdateMsg();
         this.sendPeersForRoundIfNeeded();
         return;
       }
       // If we are below the minimum number of participants
-      // tell remaining participants to wait until more participants join
+      // tell remaining participants to wait until more participants join,
+      // which already carries the number of participants
       this.sendWaitForMoreParticipantsMsg();
     });
   }
@@ -274,6 +283,7 @@ export class DecentralizedController<
           type: MessageTypes.PeersForRound,
           peers: this.#roundPeers.delete(id).keySeq().toArray(),
           aggregationRound: this.#aggregationRound,
+          nbOfParticipants: this.connections.size,
         };
       },
     );

server/src/controllers/federated_controller.ts

@@ -171,8 +171,11 @@ export class FederatedController<D extends DataType> extends TrainingController<
             nbOfParticipants: this.connections.size,
           };
           ws.send(msgpack.encode(msg));
-          // Send an update to participants if we can start/resume training
-          this.sendEnoughParticipantsMsgIfNeeded(clientId);
+          // Send an update to participants if we can start/resume training,
+          // which already carries the number of participants
+          if (!this.sendEnoughParticipantsMsgIfNeeded(clientId))
+            // otherwise just tell them that someone joined
+            this.sendParticipantsUpdateMsg(clientId);
           break;
         }
         /*
@@ -250,10 +253,14 @@ export class FederatedController<D extends DataType> extends TrainingController<
       if (
         this.connections.size >= minNbOfParticipants ||
         this.waitingForMoreParticipants
-      )
+      ) {
+        // tell the remaining participants that one of them left
+        this.sendParticipantsUpdateMsg();
         return;
+      }
 
-      // tell remaining participants to wait until more participants join
+      // tell remaining participants to wait until more participants join,
+      // which already carries the number of participants
       this.sendWaitForMoreParticipantsMsg();
     });
   }

server/src/controllers/training_controller.ts

@@ -53,12 +53,40 @@ export abstract class TrainingController<
     this.connections = Map<NodeID, WebSocket>();
   }
 
+  /**
+   * Notifies participants of how many of them there are.
+   *
+   * @param exclude a participant to leave out, typically one which just joined
+   * and already learned the count from the answer to its join request
+   */
+  protected sendParticipantsUpdateMsg(exclude?: NodeID): void {
+    const msg: mtype.ParticipantsUpdate = {
+      type: mtype.MType.ParticipantsUpdate,
+      nbOfParticipants: this.connections.size,
+    };
+    const encoded = msgpack.encode(msg);
+
+    const recipients =
+      exclude !== undefined
+        ? this.connections.delete(exclude)
+        : this.connections;
+    recipients.forEach((participantWs, participantId) => {
+      debug(
+        "Sending participants update to client [%s]",
+        participantId.slice(0, 4),
+      );
+      participantWs.send(encoded);
+    });
+  }
+
   /**
    * If enough participants joined, notifies them that the training can start/resume
    *
    * @param currentId the id of the participant that just joined
+   * @returns whether the participants were notified, which also tells them the
+   * number of participants
    */
-  protected sendEnoughParticipantsMsgIfNeeded(currentId: NodeID) {
+  protected sendEnoughParticipantsMsgIfNeeded(currentId: NodeID): boolean {
     // If we are currently waiting for more participants to join and we now have enough,
     // broadcast to previously waiting participants that the training can start
     if (
@@ -81,7 +109,10 @@ export abstract class TrainingController<
           participantWs.send(msgpack.encode(msg));
         });
       this.waitingForMoreParticipants = false; // update the attribute
+      return true;
     }
+
+    return false;
   }
 
   /**