@epfml/discojs 3.0.1-p20240902094132.0 → 3.0.1-p20240902162912.0

package/dist/training/disco.d.ts

@@ -1,32 +1,35 @@
 import { client as clients, BatchLogs, EpochLogs, Logger, Memory, Task, TrainingInformation } from "../index.js";
 import type { TypedLabeledDataset } from "../index.js";
 import type { Aggregator } from "../aggregator/index.js";
+import { EventEmitter } from "../utils/event_emitter.js";
 import { RoundLogs, Trainer } from "./trainer.js";
-interface Config {
+interface DiscoConfig {
     scheme: TrainingInformation["scheme"];
     logger: Logger;
     memory: Memory;
 }
+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";
 /**
  * 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,
  * communication with nodes, logs and model memory.
  */
-export declare class Disco {
+export declare class Disco extends EventEmitter<{
+    'status': RoundStatus;
+}> {
     #private;
     readonly trainer: Trainer;
-    private constructor();
     /**
      * Connect to the given task and get ready to train.
      *
-     * Will load the model from memory if available or fetch it from the server.
-     *
+     * @param task
      * @param clientConfig client to connect with or parameters on how to create one.
-     **/
-    static fromTask(task: Task, clientConfig: clients.Client | URL | {
+     * @param config the DiscoConfig
+     */
+    constructor(task: Task, clientConfig: clients.Client | URL | {
         aggregator: Aggregator;
         url: URL;
-    }, config: Partial<Config>): Promise<Disco>;
+    }, config: Partial<DiscoConfig>);
     /** Train on dataset, yielding logs of every round. */
     trainByRound(dataset: TypedLabeledDataset): AsyncGenerator<RoundLogs>;
     /** Train on dataset, yielding logs of every epoch. */
@@ -42,10 +45,6 @@ export declare class Disco {
      * If you don't care about the whole process, use one of the other train methods.
      **/
     train(dataset: TypedLabeledDataset): AsyncGenerator<AsyncGenerator<AsyncGenerator<BatchLogs, EpochLogs>, RoundLogs>>;
-    /**
-     * Stops the ongoing training instance without disconnecting the client.
-     */
-    pause(): Promise<void>;
     /**
      * Completely stops the ongoing training instance.
      */

package/dist/training/disco.js

@@ -1,6 +1,7 @@
 import { async_iterator, client as clients, ConsoleLogger, EmptyMemory, } from "../index.js";
 import { getAggregator } from "../aggregator/index.js";
 import { enumerate, split } from "../utils/async_iterator.js";
+import { EventEmitter } from "../utils/event_emitter.js";
 import { Trainer } from "./trainer.js";
 import { labeledDatasetToDataSplit } from "../dataset/data/helpers.js";
 /**
@@ -8,27 +9,21 @@ import { labeledDatasetToDataSplit } from "../dataset/data/helpers.js";
  * a convenient object providing a reduced yet complete API that wraps model training,
  * communication with nodes, logs and model memory.
  */
-export class Disco {
+export class Disco extends EventEmitter {
     trainer;
     #client;
     #logger;
     #memory;
     #task;
-    constructor(trainer, task, client, memory, logger) {
-        this.trainer = trainer;
-        this.#client = client;
-        this.#logger = logger;
-        this.#memory = memory;
-        this.#task = task;
-    }
     /**
      * Connect to the given task and get ready to train.
      *
-     * Will load the model from memory if available or fetch it from the server.
-     *
+     * @param task
      * @param clientConfig client to connect with or parameters on how to create one.
-     **/
-    static async fromTask(task, clientConfig, config) {
+     * @param config the DiscoConfig
+     */
+    constructor(task, clientConfig, config) {
+        super();
         const { scheme, logger, memory } = {
             scheme: task.trainingInformation.scheme,
             logger: new ConsoleLogger(),
@@ -52,18 +47,13 @@ export class Disco {
         }
         if (client.task !== task)
             throw new Error("client not setup for given task");
-        let model;
-        const memoryInfo = {
-            type: "working",
-            taskID: task.id,
-            name: task.trainingInformation.modelID,
-            tensorBackend: task.trainingInformation.tensorBackend,
-        };
-        if (await memory.contains(memoryInfo))
-            model = await memory.getModel(memoryInfo);
-        else
-            model = await client.getLatestModel();
-        return new Disco(new Trainer(task, model, client), task, client, memory, logger);
+        this.#logger = logger;
+        this.#client = client;
+        this.#memory = memory;
+        this.#task = task;
+        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));
     }
     /** Train on dataset, yielding logs of every round. */
     async *trainByRound(dataset) {
@@ -106,11 +96,12 @@ export class Disco {
      * If you don't care about the whole process, use one of the other train methods.
      **/
     async *train(dataset) {
-        this.#logger.success("Training started.");
+        this.#logger.success("Training started");
         const data = await labeledDatasetToDataSplit(this.#task, dataset);
         const trainData = data.train.preprocess().batch().dataset;
         const validationData = data.validation?.preprocess().batch().dataset ?? trainData;
-        await this.#client.connect();
+        // the client fetches the latest weights upon connection
+        this.trainer.model = await this.#client.connect();
         for await (const [round, epochs] of enumerate(this.trainer.train(trainData, validationData))) {
             yield async function* () {
                 const [gen, returnedRoundLogs] = split(epochs);
@@ -123,6 +114,7 @@ export class Disco {
                         `  Epoch: ${epoch}`,
                         `    Training loss: ${epochLogs.training.loss}`,
                         `    Training accuracy: ${epochLogs.training.accuracy}`,
+                        `    Peak memory: ${epochLogs.peakMemory}`,
                         epochLogs.validation !== undefined
                             ? `    Validation loss: ${epochLogs.validation.loss}`
                             : "",
@@ -140,19 +132,12 @@ export class Disco {
                 tensorBackend: this.#task.trainingInformation.tensorBackend,
             }, this.trainer.model);
         }
-        this.#logger.success("Training finished.");
-    }
-    /**
-     * Stops the ongoing training instance without disconnecting the client.
-     */
-    async pause() {
-        await this.trainer.stopTraining();
+        this.#logger.success("Training finished");
     }
     /**
      * Completely stops the ongoing training instance.
      */
     async close() {
-        await this.pause();
         await this.#client.disconnect();
     }
 }

package/dist/training/index.d.ts

@@ -1,2 +1,2 @@
-export { Disco } from './disco.js';
+export { Disco, RoundStatus } from './disco.js';
 export { RoundLogs, Trainer } from './trainer.js';

package/dist/training/trainer.d.ts

@@ -9,8 +9,9 @@ export interface RoundLogs {
 /** Train a model and exchange with others **/
 export declare class Trainer {
     #private;
-    readonly model: Model;
-    constructor(task: Task, model: Model, client: Client);
+    get model(): Model;
+    set model(model: Model);
+    constructor(task: Task, client: Client);
     stopTraining(): Promise<void>;
     train(dataset: tf.data.Dataset<tf.TensorContainer>, valDataset: tf.data.Dataset<tf.TensorContainer>): AsyncGenerator<AsyncGenerator<AsyncGenerator<BatchLogs, EpochLogs>, RoundLogs>, void>;
 }

package/dist/training/trainer.js

@@ -4,14 +4,21 @@ import { privacy } from "../index.js";
 import * as async_iterator from "../utils/async_iterator.js";
 /** Train a model and exchange with others **/
 export class Trainer {
-    model;
     #client;
     #roundDuration;
     #epochs;
     #privacy;
+    #model;
     #training;
-    constructor(task, model, client) {
-        this.model = model;
+    get model() {
+        if (this.#model === undefined)
+            throw new Error("trainer's model has not been set");
+        return this.#model;
+    }
+    set model(model) {
+        this.#model = model;
+    }
+    constructor(task, client) {
         this.#client = client;
         this.#roundDuration = task.trainingInformation.roundDuration;
         this.#epochs = task.trainingInformation.epochs;
@@ -37,12 +44,12 @@ export class Trainer {
         const totalRound = Math.trunc(this.#epochs / this.#roundDuration);
         let previousRoundWeights;
         for (let round = 0; round < totalRound; round++) {
-            await this.#client.onRoundBeginCommunication(this.model.weights, round);
+            await this.#client.onRoundBeginCommunication();
             yield this.#runRound(dataset, valDataset);
             let localWeights = this.model.weights;
             if (this.#privacy !== undefined)
                 localWeights = await applyPrivacy(previousRoundWeights, localWeights, this.#privacy);
-            const networkWeights = await this.#client.onRoundEndCommunication(localWeights, round);
+            const networkWeights = await this.#client.onRoundEndCommunication(localWeights);
             this.model.weights = previousRoundWeights = networkWeights;
         }
     }

package/dist/utils/event_emitter.js

@@ -47,9 +47,7 @@ export class EventEmitter {
     emit(event, value) {
         const eventListeners = this.listeners[event] ?? List();
         this.listeners[event] = eventListeners.filterNot(([once]) => once);
-        eventListeners.forEach(([_, listener]) => {
-            listener(value);
-        });
+        eventListeners.forEach(([_, listener]) => { listener(value); });
     }
 }
 /** `EventEmitter` for all events */

package/package.json

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

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

@@ -1,38 +0,0 @@
-import { type WeightsContainer } from "../../index.js";
-import { Base as Client } from "../base.js";
-/**
- * Client class that communicates with a centralized, federated server, when training
- * a specific task in the federated setting.
- */
-export declare class Base extends Client {
-    #private;
-    /**
-     * Arbitrary node id assigned to the federated server which we are communicating with.
-     * Indeed, the server acts as a node within the network. In the federated setting described
-     * by this client class, the server is the only node which we are communicating with.
-     */
-    static readonly SERVER_NODE_ID = "federated-server-node-id";
-    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 and get our own node id.
-     * TODO: In the federated setting, should return the current server-side round
-     * for the task.
-     */
-    connect(): Promise<void>;
-    /**
-     * Disconnection process when user quits the task.
-     */
-    disconnect(): Promise<void>;
-    onRoundBeginCommunication(): Promise<void>;
-    onRoundEndCommunication(weights: WeightsContainer, round: number): Promise<WeightsContainer>;
-    /**
-     * Send a message containing our local weight updates to the federated server.
-     * And waits for the server to reply with the most recent aggregated weights
-     * @param payload The weight updates to send
-     */
-    private sendPayloadAndReceiveResult;
-}

package/dist/client/federated/base.js

@@ -1,130 +0,0 @@
-import createDebug from "debug";
-import { serialization, } from "../../index.js";
-import { Base as Client } from "../base.js";
-import { type } from "../messages.js";
-import { waitMessageWithTimeout, WebSocketServer, } from "../event_connection.js";
-import * as messages from "./messages.js";
-const debug = createDebug("discojs:client:federated");
-/**
- * Client class that communicates with a centralized, federated server, when training
- * a specific task in the federated setting.
- */
-export class Base extends Client {
-    /**
-     * Arbitrary node id assigned to the federated server which we are communicating with.
-     * Indeed, the server acts as a node within the network. In the federated setting described
-     * by this client class, the server is the only node which we are communicating with.
-     */
-    static SERVER_NODE_ID = "federated-server-node-id";
-    // 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;
-    // the number of participants excluding the server
-    get nbOfParticipants() {
-        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, messages.isMessageFederated);
-        return server;
-    }
-    /**
-     * Initializes the connection to the server and get our own node id.
-     * TODO: In the federated setting, should return the current server-side round
-     * for the task.
-     */
-    async connect() {
-        const serverURL = new URL("", this.url.href);
-        switch (this.url.protocol) {
-            case "http:":
-                serverURL.protocol = "ws:";
-                break;
-            case "https:":
-                serverURL.protocol = "wss:";
-                break;
-            default:
-                throw new Error(`unknown protocol: ${this.url.protocol}`);
-        }
-        serverURL.pathname += `feai/${this.task.id}`;
-        this._server = await this.connectServer(serverURL);
-        this.aggregator.registerNode(Base.SERVER_NODE_ID);
-        const msg = {
-            type: type.ClientConnected,
-        };
-        this.server.send(msg);
-        const received = await waitMessageWithTimeout(this.server, type.AssignNodeID);
-        debug(`[${received.id}] assign id generated by the server`);
-        this._ownId = received.id;
-    }
-    /**
-     * Disconnection process when user quits the task.
-     */
-    async disconnect() {
-        await this.server.disconnect();
-        this._server = undefined;
-        this._ownId = undefined;
-        this.aggregator.setNodes(this.aggregator.nodes.delete(Base.SERVER_NODE_ID));
-        return Promise.resolve();
-    }
-    onRoundBeginCommunication() {
-        // Prepare the result promise for the incoming round
-        this.aggregationResult = new Promise((resolve) => this.aggregator.once('aggregation', resolve));
-        return Promise.resolve();
-    }
-    async onRoundEndCommunication(weights, round) {
-        // NB: For now, we suppose a fully-federated setting.
-        if (this.aggregationResult === undefined) {
-            throw new Error("local aggregation result was not set");
-        }
-        // Send our local contribution to the server
-        // and receive the most recent weights as an answer to our contribution
-        const serverResult = await this.sendPayloadAndReceiveResult(this.aggregator.makePayloads(weights).first());
-        if (serverResult !== undefined &&
-            this.aggregator.add(Base.SERVER_NODE_ID, serverResult, round, 0)) {
-            // Regular case: the server sends us its aggregation result which will serve our
-            // own aggregation result.
-        }
-        else {
-            // Unexpected case: for some reason, the server result is stale.
-            // We proceed to the next round without its result.
-            debug(`[${this.ownId}] server result is either stale or not received`);
-            this.aggregator.nextRound();
-        }
-        return await this.aggregationResult;
-    }
-    /**
-     * Send a message containing our local weight updates to the federated server.
-     * And waits for the server to reply with the most recent aggregated weights
-     * @param payload The weight updates to send
-     */
-    async sendPayloadAndReceiveResult(payload) {
-        const msg = {
-            type: type.SendPayload,
-            payload: await serialization.weights.encode(payload),
-            round: this.aggregator.round,
-        };
-        this.server.send(msg);
-        // Waits for the server's result for its current (most recent) round and add it to our aggregator.
-        // Updates the aggregator's round if it's behind the server's.
-        try {
-            // It is important than the client immediately awaits the server result or it may miss it
-            const { payload, round, nbOfParticipants } = await waitMessageWithTimeout(this.server, type.ReceiveServerPayload);
-            const serverRound = round;
-            this.#nbOfParticipants = nbOfParticipants; // Save the current participants
-            // Store the server result only if it is not stale
-            if (this.aggregator.round <= round) {
-                const serverResult = serialization.weights.decode(payload);
-                // Update the local round to match the server's
-                if (this.aggregator.round < serverRound) {
-                    this.aggregator.setRound(serverRound);
-                }
-                return serverResult;
-            }
-        }
-        catch (e) {
-            debug(`[${this.ownId}] while receiving results: %o`, e);
-        }
-    }
-}

package/dist/client/local.d.ts

@@ -1,5 +0,0 @@
-import { WeightsContainer } from "../weights/weights_container.js";
-import { Base } from "./base.js";
-export declare class Local extends Base {
-    onRoundEndCommunication(weights: WeightsContainer): Promise<WeightsContainer>;
-}

package/dist/client/local.js

@@ -1,6 +0,0 @@
-import { Base } from "./base.js";
-export class Local extends Base {
-    onRoundEndCommunication(weights) {
-        return Promise.resolve(weights);
-    }
-}

package/dist/task/digest.d.ts

@@ -1,5 +0,0 @@
-export interface Digest {
-    algorithm: string;
-    value: string;
-}
-export declare function isDigest(raw: unknown): raw is Digest;

package/dist/task/digest.js

@@ -1,14 +0,0 @@
-export function isDigest(raw) {
-    if (typeof raw !== 'object' || raw === null) {
-        return false;
-    }
-    const { algorithm, value } = raw;
-    if (!(typeof algorithm === 'string' &&
-        typeof value === 'string')) {
-        return false;
-    }
-    const repack = { algorithm, value };
-    const _correct = repack;
-    const _total = repack;
-    return true;
-}