Decentralized controller consistency

discojs/src/client/client.ts

@@ -99,7 +99,14 @@ export abstract class Client<N extends Network> extends EventEmitter<{
   /**
    * Handles the disconnection process of the client from any sort of network server.
    */
-  async disconnect(): Promise<void> {}
+  async disconnect(): Promise<void> {
+    // Optional server for the local client
+    // Makes it also safe on multiple calls
+    await this._server?.disconnect();
+    this._server = undefined;
+    this._ownId = undefined;
+    this.promiseForMoreParticipants = undefined;
+  }
 
   /**
    * Emits the round status specified. It also stores the status emitted such that

discojs/src/client/decentralized/decentralized_client.ts

@@ -14,6 +14,7 @@ import {
   waitMessage,
   type PeerConnection,
   waitMessageWithTimeout,
+  sendAndWaitWithRetry,
 } from "#client/event_connection";
 import { PeerPool } from "#client/decentralized/peer_pool";
 import * as messages from "#client/decentralized/messages";
@@ -146,10 +147,15 @@ export class DecentralizedClient extends Client<"decentralized"> {
     const msg: ClientConnected = {
       type: MType.ClientConnected,
     };
-    this.server.send(msg);
-
     const { id, waitForMoreParticipants, nbOfParticipants, joinedMidTraining } =
-      await waitMessage(this.server, MType.NewDecentralizedNodeInfo);
+      await sendAndWaitWithRetry(
+        this.server,
+        msg,
+        MType.NewDecentralizedNodeInfo,
+        {
+          signal: this.crashSignal,
+        },
+      );
 
     this.#modelSyncNeeded = joinedMidTraining;
     this.setNbOfParticipantsUponJoining(nbOfParticipants);
@@ -184,14 +190,13 @@ export class DecentralizedClient extends Client<"decentralized"> {
       const peers = this.#connections.keySeq().toSet();
       this.aggregator.setNodes(this.aggregator.nodes.subtract(peers));
     }
-    // Disconnect from server
-    await this.server?.disconnect();
-    this._server = undefined;
-    this._ownId = undefined;
 
     this.#latestModel?.dispose();
     this.#latestModel = undefined;
 
+    // At the end in case there is a throw while disconnecting the server
+    await super.disconnect();
+
     return Promise.resolve();
   }
 
@@ -206,6 +211,9 @@ export class DecentralizedClient extends Client<"decentralized"> {
    * model syncing happens first to get the latest global model from another peer.
    */
   override async onRoundBeginCommunication(): Promise<void> {
+    // Check if the crash signal has been triggered and throw if so
+    this.crashSignal.throwIfAborted();
+
     if (this.#modelSyncNeeded) {
       // 1. If model sync is needed, send server a request
       this.server.send({ type: MType.ModelSyncRequest });
@@ -289,9 +297,9 @@ export class DecentralizedClient extends Client<"decentralized"> {
       // Therefore, unsolved listeners may accumulate across rounds.
       // We can add listener resolving if this becomes a problem later.
       const msg = await Promise.race([
-        waitMessage(this.server, MType.StartWeightSharing),
-        waitMessage(this.server, MType.RetryPeerConnections),
-        waitMessage(this.server, MType.ConnectionFail),
+        waitMessage(this.server, MType.StartWeightSharing, this.crashSignal),
+        waitMessage(this.server, MType.RetryPeerConnections, this.crashSignal),
+        waitMessage(this.server, MType.ConnectionFail, this.crashSignal),
       ]);
 
       if (msg.type === MType.StartWeightSharing) {
@@ -355,6 +363,7 @@ export class DecentralizedClient extends Client<"decentralized"> {
       const receivedMessage = await waitMessage(
         this.server,
         MType.PeersForRound,
+        this.crashSignal,
       );
       // every peer is ready to share weights, we can now connect to them
       this.saveAndEmit("connecting to peers");
@@ -400,6 +409,7 @@ export class DecentralizedClient extends Client<"decentralized"> {
       // still part of the session: leave the number of participants alone
       this.aggregator.setNodes(Set([this.ownId]));
       this.#connections = Map();
+      this.crashSignal.throwIfAborted(); // Do not swallow if it is a signal error
     }
   }
 

discojs/src/client/decentralized/messages.ts

@@ -8,6 +8,7 @@ import type {
   WaitingForMoreParticipants,
   EnoughParticipants,
   ParticipantsUpdate,
+  CrashClient,
 } from "#client/mtype";
 
 /// Phase 0 communication (between server and peers)
@@ -116,7 +117,8 @@ export type MessageFromServer =
   | RetryPeerConnections
   | ConnectionFail
   | ModelProviderInfo
-  | ProvideModelToPeer;
+  | ProvideModelToPeer
+  | CrashClient;
 
 export type MessageToServer =
   | ClientConnected
@@ -151,6 +153,7 @@ export function isMessageFromServer(o: unknown): o is MessageFromServer {
     case MType.ConnectionFail:
     case MType.ModelProviderInfo:
     case MType.ProvideModelToPeer:
+    case MType.CrashClient:
       return true;
   }
 

discojs/src/client/federated/federated_client.ts

@@ -112,12 +112,8 @@ export class FederatedClient extends Client<"federated"> {
    * Disconnection process when user quits the task.
    */
   override async disconnect(): Promise<void> {
-    await this.server.disconnect();
-    this._server = undefined;
-    this._ownId = undefined;
-    this.promiseForMoreParticipants = undefined;
-
     this.aggregator.setNodes(this.aggregator.nodes.delete(SERVER_NODE_ID));
+    await super.disconnect();
   }
 
   override onRoundBeginCommunication(): Promise<void> {