@epfml/discojs 2.2.1 → 3.0.0

package/dist/training/disco.js

@@ -1,65 +1,73 @@
-import { List } from 'immutable';
-import { async_iterator } from '../index.js';
-import { client as clients, EmptyMemory, ConsoleLogger } from '../index.js';
-import { MeanAggregator } from '../aggregator/mean.js';
-import { enumerate, split } from '../utils/async_iterator.js';
-import { TrainerBuilder } from './trainer/trainer_builder.js';
+import { async_iterator, } from "../index.js";
+import { client as clients, ConsoleLogger, EmptyMemory } from "../index.js";
+import { getAggregator } from "../aggregator/index.js";
+import { enumerate, split } from "../utils/async_iterator.js";
+import { Trainer } from "./trainer.js";
 /**
  * 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 class Disco {
-    task;
-    logger;
-    memory;
-    client;
     trainer;
-    constructor(task, options) {
-        if (options.scheme === undefined) {
-            options.scheme = task.trainingInformation.scheme;
-        }
-        if (options.aggregator === undefined) {
-            options.aggregator = new MeanAggregator();
+    #client;
+    #logger;
+    // small helper to avoid keeping Task & Memory around
+    #updateWorkingModel;
+    constructor(trainer, task, client, memory, logger) {
+        this.trainer = trainer;
+        this.#client = client;
+        this.#logger = logger;
+        this.#updateWorkingModel = () => memory.updateWorkingModel({
+            type: "working",
+            taskID: task.id,
+            name: task.trainingInformation.modelID,
+            tensorBackend: task.trainingInformation.tensorBackend,
+        }, this.trainer.model);
+    }
+    /**
+     * 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 clientConfig client to connect with or parameters on how to create one.
+     **/
+    static async fromTask(task, clientConfig, config) {
+        const { scheme, logger, memory } = {
+            scheme: task.trainingInformation.scheme,
+            logger: new ConsoleLogger(),
+            memory: new EmptyMemory(),
+            ...config,
+        };
+        let client;
+        if (clientConfig instanceof clients.Client) {
+            client = clientConfig;
         }
-        if (options.client === undefined) {
-            if (options.url === undefined) {
-                throw new Error('could not determine client from given parameters');
-            }
-            if (typeof options.url === 'string') {
-                options.url = new URL(options.url);
+        else {
+            let url, aggregator;
+            if (clientConfig instanceof URL) {
+                url = clientConfig;
+                aggregator = getAggregator(task, { scheme });
             }
-            switch (options.scheme) {
-                case 'federated':
-                    options.client = new clients.federated.FederatedClient(options.url, task, options.aggregator);
-                    break;
-                case 'decentralized':
-                    options.client = new clients.decentralized.DecentralizedClient(options.url, task, options.aggregator);
-                    break;
-                case 'local':
-                    options.client = new clients.Local(options.url, task, options.aggregator);
-                    break;
-                default: {
-                    const _ = options.scheme;
-                    throw new Error('should never happen');
-                }
+            else {
+                ({ url, aggregator } = clientConfig);
             }
+            client = clients.getClient(scheme, url, task, aggregator);
         }
-        if (options.logger === undefined) {
-            options.logger = new ConsoleLogger();
-        }
-        if (options.memory === undefined) {
-            options.memory = new EmptyMemory();
-        }
-        if (options.client.task !== task) {
-            throw new Error('client not setup for given task');
-        }
-        this.task = task;
-        this.client = options.client;
-        this.memory = options.memory;
-        this.logger = options.logger;
-        const trainerBuilder = new TrainerBuilder(this.memory, this.task);
-        this.trainer = trainerBuilder.build(this.client, options.scheme !== 'local');
+        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);
     }
     /** Train on dataset, yielding logs of every round. */
     async *trainByRound(dataTuple) {
@@ -96,27 +104,24 @@ export class Disco {
                     ;
     }
     /**
-    * Train on dataset, yield the nested steps.
-    *
-    * Don't forget to await the yielded generator otherwise nothing will progress.
-    * If you don't care about the whole process, use one of the other train methods.
-    **/
-    // TODO RoundLogs should contain number of participants but Trainer doesn't need client
+     * Train on dataset, yield the nested steps.
+     *
+     * Don't forget to await the yielded generator otherwise nothing will progress.
+     * If you don't care about the whole process, use one of the other train methods.
+     **/
     async *train(dataTuple) {
-        this.logger.success("Training started.");
+        this.#logger.success("Training started.");
         const trainData = dataTuple.train.preprocess().batch();
         const validationData = dataTuple.validation?.preprocess().batch() ?? trainData;
-        await this.client.connect();
-        const trainer = await this.trainer;
-        for await (const [round, epochs] of enumerate(trainer.fitModel(trainData.dataset, validationData.dataset))) {
+        await this.#client.connect();
+        for await (const [round, epochs] of enumerate(this.trainer.train(trainData.dataset, validationData.dataset))) {
             yield async function* () {
-                let epochsLogs = List();
-                for await (const [epoch, batches] of enumerate(epochs)) {
+                const [gen, returnedRoundLogs] = split(epochs);
+                for await (const [epoch, batches] of enumerate(gen)) {
                     const [gen, returnedEpochLogs] = split(batches);
                     yield gen;
                     const epochLogs = await returnedEpochLogs;
-                    epochsLogs = epochsLogs.push(epochLogs);
-                    this.logger.success([
+                    this.#logger.success([
                         `Round: ${round}`,
                         `  Epoch: ${epoch}`,
                         `    Training loss: ${epochLogs.training.loss}`,
@@ -129,26 +134,23 @@ export class Disco {
                             : "",
                     ].join("\n"));
                 }
-                return {
-                    epochs: epochsLogs,
-                    participants: this.client.nodes.size + 1, // add ourself
-                };
+                return await returnedRoundLogs;
             }.bind(this)();
+            await this.#updateWorkingModel(this.trainer.model);
         }
-        this.logger.success("Training finished.");
+        this.#logger.success("Training finished.");
     }
     /**
      * Stops the ongoing training instance without disconnecting the client.
      */
     async pause() {
-        const trainer = await this.trainer;
-        await trainer.stopTraining();
+        await this.trainer.stopTraining();
     }
     /**
      * Completely stops the ongoing training instance.
      */
     async close() {
         await this.pause();
-        await this.client.disconnect();
+        await this.#client.disconnect();
     }
 }

package/dist/training/index.d.ts

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

package/dist/training/index.js

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

package/dist/training/trainer.d.ts

@@ -0,0 +1,16 @@
+import * as tf from "@tensorflow/tfjs";
+import { List } from "immutable";
+import type { BatchLogs, EpochLogs, Model, Task } from "../index.js";
+import { Client } from "../client/index.js";
+export interface RoundLogs {
+    epochs: List<EpochLogs>;
+    participants: number;
+}
+/** Train a model and exchange with others **/
+export declare class Trainer {
+    #private;
+    readonly model: Model;
+    constructor(task: Task, model: Model, 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

@@ -0,0 +1,72 @@
+import * as tf from "@tensorflow/tfjs";
+import { List } from "immutable";
+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;
+    #training;
+    constructor(task, model, client) {
+        this.model = model;
+        this.#client = client;
+        this.#roundDuration = task.trainingInformation.roundDuration;
+        this.#epochs = task.trainingInformation.epochs;
+        this.#privacy = task.trainingInformation.privacy;
+        if (!Number.isInteger(this.#epochs / this.#roundDuration))
+            throw new Error(`round duration ${this.#roundDuration} doesn't divide number of epochs ${this.#epochs}`);
+    }
+    async stopTraining() {
+        await this.#training?.return();
+    }
+    async *train(dataset, valDataset) {
+        if (this.#training !== undefined)
+            throw new Error("training already running, stop it before launching a new one");
+        try {
+            this.#training = this.#runRounds(dataset, valDataset);
+            yield* this.#training;
+        }
+        finally {
+            this.#training = undefined;
+        }
+    }
+    async *#runRounds(dataset, valDataset) {
+        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);
+            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);
+            this.model.weights = previousRoundWeights = networkWeights;
+        }
+    }
+    async *#runRound(dataset, valDataset) {
+        let epochsLogs = List();
+        for (let epoch = 0; epoch < this.#roundDuration; epoch++) {
+            const [gen, epochLogs] = async_iterator.split(this.model.train(dataset, valDataset));
+            yield gen;
+            epochsLogs = epochsLogs.push(await epochLogs);
+        }
+        return {
+            epochs: epochsLogs,
+            participants: this.#client.nbOfParticipants,
+        };
+    }
+}
+async function applyPrivacy(previous, current, options) {
+    let ret = current;
+    if (options.clippingRadius !== undefined) {
+        const previousRoundWeights = previous ?? current.map((w) => tf.zerosLike(w));
+        const weightsProgress = current.sub(previousRoundWeights);
+        ret = previousRoundWeights.add(await privacy.clipNorm(weightsProgress, options.clippingRadius));
+    }
+    if (options.noiseScale !== undefined)
+        ret = privacy.addNoise(ret, options.noiseScale);
+    return ret;
+}

package/dist/types.d.ts

@@ -2,7 +2,5 @@ import type { Map } from 'immutable';
 import type { WeightsContainer } from './index.js';
 import type { NodeID } from './client/index.js';
 export type Path = string;
-export type MetadataKey = string;
-export type MetadataValue = string;
 export type Features = number | number[] | number[][] | number[][][] | number[][][][] | number[][][][][];
 export type Contributions = Map<NodeID, WeightsContainer>;

package/dist/weights/weights_container.d.ts

@@ -51,11 +51,6 @@ export declare class WeightsContainer {
      * @returns The tensor located at the index
      */
     get(index: number): tf.Tensor | undefined;
-    /**
-     * Computes the weights container's Frobenius norm
-     * @returns The Frobenius norm
-     */
-    frobeniusNorm(): number;
     concat(other: WeightsContainer): WeightsContainer;
     equals(other: WeightsContainer, margin?: number): boolean;
     /**

package/dist/weights/weights_container.js

@@ -70,13 +70,6 @@ export class WeightsContainer {
     get(index) {
         return this._weights.get(index);
     }
-    /**
-     * Computes the weights container's Frobenius norm
-     * @returns The Frobenius norm
-     */
-    frobeniusNorm() {
-        return Math.sqrt(this.map((w) => w.square().sum()).reduce((a, b) => a.add(b)).dataSync()[0]);
-    }
     concat(other) {
         return WeightsContainer.of(...this.weights, ...other.weights);
     }

package/package.json

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

package/dist/async_informant.d.ts

@@ -1,15 +0,0 @@
-import type { AggregatorBase } from './aggregator/index.js';
-export declare class AsyncInformant<T> {
-    private readonly aggregator;
-    private _round;
-    private _currentNumberOfParticipants;
-    private _totalNumberOfParticipants;
-    private _averageNumberOfParticipants;
-    constructor(aggregator: AggregatorBase<T>);
-    update(): void;
-    get round(): number;
-    get currentNumberOfParticipants(): number;
-    get totalNumberOfParticipants(): number;
-    get averageNumberOfParticipants(): number;
-    getAllStatistics(): Record<'round' | 'currentNumberOfParticipants' | 'totalNumberOfParticipants' | 'averageNumberOfParticipants', number>;
-}

package/dist/async_informant.js

@@ -1,42 +0,0 @@
-export class AsyncInformant {
-    aggregator;
-    _round = 0;
-    _currentNumberOfParticipants = 0;
-    _totalNumberOfParticipants = 0;
-    _averageNumberOfParticipants = 0;
-    constructor(aggregator) {
-        this.aggregator = aggregator;
-    }
-    update() {
-        if (this.round === 0 || this.round < this.aggregator.round) {
-            this._round = this.aggregator.round;
-            this._currentNumberOfParticipants = this.aggregator.size;
-            this._averageNumberOfParticipants = this.totalNumberOfParticipants / this.round;
-            this._totalNumberOfParticipants += this.currentNumberOfParticipants;
-        }
-        else {
-            this._round = this.aggregator.round;
-        }
-    }
-    // Getter functions
-    get round() {
-        return this._round;
-    }
-    get currentNumberOfParticipants() {
-        return this._currentNumberOfParticipants;
-    }
-    get totalNumberOfParticipants() {
-        return this._totalNumberOfParticipants;
-    }
-    get averageNumberOfParticipants() {
-        return this._averageNumberOfParticipants;
-    }
-    getAllStatistics() {
-        return {
-            round: this.round,
-            currentNumberOfParticipants: this.currentNumberOfParticipants,
-            totalNumberOfParticipants: this.totalNumberOfParticipants,
-            averageNumberOfParticipants: this.averageNumberOfParticipants
-        };
-    }
-}

package/dist/training/trainer/distributed_trainer.d.ts

@@ -1,20 +0,0 @@
-import type { Model, Memory, Task, client as clients } from "../../index.js";
-import { Trainer } from "./trainer.js";
-/**
- * Class whose role is to train a model in a distributed way with a given dataset.
- */
-export declare class DistributedTrainer extends Trainer {
-    private readonly task;
-    private readonly memory;
-    private readonly client;
-    private readonly aggregator;
-    /**
-     * DistributedTrainer constructor, accepts same arguments as Trainer and in additional also a client who takes care of communicating weights.
-     */
-    constructor(task: Task, memory: Memory, model: Model, client: clients.Client);
-    onRoundBegin(round: number): Promise<void>;
-    /**
-     * Callback called every time a round is over
-     */
-    onRoundEnd(round: number): Promise<void>;
-}

package/dist/training/trainer/distributed_trainer.js

@@ -1,41 +0,0 @@
-import { Trainer } from "./trainer.js";
-/**
- * Class whose role is to train a model in a distributed way with a given dataset.
- */
-export class DistributedTrainer extends Trainer {
-    task;
-    memory;
-    client;
-    aggregator;
-    /**
-     * DistributedTrainer constructor, accepts same arguments as Trainer and in additional also a client who takes care of communicating weights.
-     */
-    constructor(task, memory, model, client) {
-        super(task, model);
-        this.task = task;
-        this.memory = memory;
-        this.client = client;
-        this.aggregator = this.client.aggregator;
-        this.aggregator.setModel(model);
-    }
-    async onRoundBegin(round) {
-        await this.client.onRoundBeginCommunication(this.model.weights, round);
-    }
-    /**
-     * Callback called every time a round is over
-     */
-    async onRoundEnd(round) {
-        await this.client.onRoundEndCommunication(this.model.weights, round);
-        if (this.aggregator.model !== undefined) {
-            // The aggregator's own aggregation is async. The trainer updates its model to match the aggregator's
-            // after it has completed a round of training.
-            this.model.weights = this.aggregator.model.weights;
-        }
-        await this.memory.updateWorkingModel({
-            type: 'working',
-            taskID: this.task.id,
-            name: this.task.trainingInformation.modelID,
-            tensorBackend: this.task.trainingInformation.tensorBackend
-        }, this.model);
-    }
-}

package/dist/training/trainer/local_trainer.d.ts

@@ -1,12 +0,0 @@
-import type { Memory, Model, Task } from "../../index.js";
-import { Trainer } from "./trainer.js";
-/** Class whose role is to locally (alone) train a model on a given dataset,
- * without any collaborators.
- */
-export declare class LocalTrainer extends Trainer {
-    private readonly task;
-    private readonly memory;
-    constructor(task: Task, memory: Memory, model: Model);
-    onRoundBegin(): Promise<void>;
-    onRoundEnd(): Promise<void>;
-}

package/dist/training/trainer/local_trainer.js

@@ -1,24 +0,0 @@
-import { Trainer } from "./trainer.js";
-/** Class whose role is to locally (alone) train a model on a given dataset,
- * without any collaborators.
- */
-export class LocalTrainer extends Trainer {
-    task;
-    memory;
-    constructor(task, memory, model) {
-        super(task, model);
-        this.task = task;
-        this.memory = memory;
-    }
-    async onRoundBegin() {
-        return await Promise.resolve();
-    }
-    async onRoundEnd() {
-        await this.memory.updateWorkingModel({
-            type: 'working',
-            taskID: this.task.id,
-            name: this.task.trainingInformation.modelID,
-            tensorBackend: this.task.trainingInformation.tensorBackend
-        }, this.model);
-    }
-}

package/dist/training/trainer/trainer.d.ts

@@ -1,32 +0,0 @@
-import type tf from "@tensorflow/tfjs";
-import { List } from "immutable";
-import type { Model, Task } from "../../index.js";
-import { BatchLogs, EpochLogs } from "../../models/index.js";
-export interface RoundLogs {
-    epochs: List<EpochLogs>;
-}
-/** Abstract class whose role is to train a model with a given dataset. This can be either done
- * locally (alone) or in a distributed way with collaborators.
- *
- * 1. Call `fitModel(dataset)` to start training.
- * 2. which will then call onRoundEnd once the round has ended.
- *
- * The onRoundEnd needs to be implemented to specify what actions to do when the round has ended, such as a communication step with collaborators.
- */
-export declare abstract class Trainer {
-    #private;
-    readonly model: Model;
-    private training?;
-    constructor(task: Task, model: Model);
-    protected abstract onRoundBegin(round: number): Promise<void>;
-    protected abstract onRoundEnd(round: number): Promise<void>;
-    /**
-     * Request stop training to be used from the Disco instance or any class that is taking care of the trainer.
-     */
-    stopTraining(): Promise<void>;
-    /**
-     * Start training the model with the given dataset
-     * @param dataset
-     */
-    fitModel(dataset: tf.data.Dataset<tf.TensorContainer>, valDataset: tf.data.Dataset<tf.TensorContainer>): AsyncGenerator<AsyncGenerator<AsyncGenerator<BatchLogs, EpochLogs>, RoundLogs>, void>;
-}

package/dist/training/trainer/trainer.js

@@ -1,61 +0,0 @@
-import { List } from "immutable";
-import * as async_iterator from "../../utils/async_iterator.js";
-/** Abstract class whose role is to train a model with a given dataset. This can be either done
- * locally (alone) or in a distributed way with collaborators.
- *
- * 1. Call `fitModel(dataset)` to start training.
- * 2. which will then call onRoundEnd once the round has ended.
- *
- * The onRoundEnd needs to be implemented to specify what actions to do when the round has ended, such as a communication step with collaborators.
- */
-export class Trainer {
-    model;
-    #roundDuration;
-    #epochs;
-    training;
-    constructor(task, model) {
-        this.model = model;
-        this.#roundDuration = task.trainingInformation.roundDuration;
-        this.#epochs = task.trainingInformation.epochs;
-        if (!Number.isInteger(this.#epochs / this.#roundDuration))
-            throw new Error(`round duration doesn't divide epochs`);
-    }
-    /**
-     * Request stop training to be used from the Disco instance or any class that is taking care of the trainer.
-     */
-    async stopTraining() {
-        await this.training?.return();
-    }
-    /**
-     * Start training the model with the given dataset
-     * @param dataset
-     */
-    async *fitModel(dataset, valDataset) {
-        if (this.training !== undefined)
-            throw new Error("training already running, cancel it before launching a new one");
-        try {
-            this.training = this.#runRounds(dataset, valDataset);
-            yield* this.training;
-        }
-        finally {
-            this.training = undefined;
-        }
-    }
-    async *#runRounds(dataset, valDataset) {
-        const totalRound = Math.trunc(this.#epochs / this.#roundDuration);
-        for (let round = 0; round < totalRound; round++) {
-            await this.onRoundBegin(round);
-            yield this.#runRound(dataset, valDataset);
-            await this.onRoundEnd(round);
-        }
-    }
-    async *#runRound(dataset, valDataset) {
-        let epochsLogs = List();
-        for (let epoch = 0; epoch < this.#roundDuration; epoch++) {
-            const [gen, epochLogs] = async_iterator.split(this.model.train(dataset, valDataset));
-            yield gen;
-            epochsLogs = epochsLogs.push(await epochLogs);
-        }
-        return { epochs: epochsLogs };
-    }
-}

package/dist/training/trainer/trainer_builder.d.ts

@@ -1,23 +0,0 @@
-import type { client as clients, Task, Memory } from '../../index.js';
-import type { Trainer } from './trainer.js';
-/**
- * A class that helps build the Trainer and auxiliary classes.
- */
-export declare class TrainerBuilder {
-    private readonly memory;
-    private readonly task;
-    constructor(memory: Memory, task: Task);
-    /**
-     * Builds a trainer object.
-     *
-     * @param client client to share weights with (either distributed or federated)
-     * @param distributed whether to build a distributed or local trainer
-     * @returns
-     */
-    build(client: clients.Client, distributed?: boolean): Promise<Trainer>;
-    /**
-     * If a model exists in memory, load it, otherwise load model from server
-     * @returns
-     */
-    private getModel;
-}