@epfml/discojs 2.0.0 → 2.1.2-p20240506085037.0

package/dist/training/disco.js

@@ -0,0 +1,107 @@
+import { client as clients, EmptyMemory, ConsoleLogger } from '../index.js';
+import { MeanAggregator } from '../aggregator/mean.js';
+import { TrainerBuilder } from './trainer/trainer_builder.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();
+        }
+        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);
+            }
+            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');
+                }
+            }
+        }
+        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');
+    }
+    /**
+     * Starts a training instance for the Disco object's task on the provided data tuple.
+     * @param dataTuple The data tuple
+     */
+    // TODO RoundLogs should contain number of participants but Trainer doesn't need client
+    async *fit(dataTuple) {
+        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 roundLogs of trainer.fitModel(trainData.dataset, validationData.dataset)) {
+            let msg = `Round: ${roundLogs.round}\n`;
+            for (const epochLogs of roundLogs.epochs.values()) {
+                msg += `    Epoch: ${epochLogs.epoch}\n`;
+                msg += `        Training loss: ${epochLogs.training.loss}\n`;
+                if (epochLogs.training.accuracy !== undefined) {
+                    msg += `        Training accuracy: ${epochLogs.training.accuracy}\n`;
+                }
+                if (epochLogs.validation !== undefined) {
+                    msg += `        Validation loss: ${epochLogs.validation.loss}\n`;
+                    msg += `        Validation accuracy: ${epochLogs.validation.accuracy}\n`;
+                }
+            }
+            this.logger.success(msg);
+            yield {
+                ...roundLogs,
+                participants: this.client.nodes.size + 1 // add ourself
+            };
+        }
+        this.logger.success("Training finished.");
+    }
+    /**
+     * Stops the ongoing training instance without disconnecting the client.
+     */
+    async pause() {
+        const trainer = await this.trainer;
+        await trainer.stopTraining();
+    }
+    /**
+     * Completely stops the ongoing training instance.
+     */
+    async close() {
+        await this.pause();
+        await this.client.disconnect();
+    }
+}

package/dist/training/index.d.ts

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

package/dist/training/index.js

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

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

@@ -0,0 +1,20 @@
+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

@@ -0,0 +1,36 @@
+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({ taskID: this.task.id, name: this.task.trainingInformation.modelID }, this.model);
+    }
+}

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

@@ -0,0 +1,12 @@
+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

@@ -0,0 +1,19 @@
+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({ taskID: this.task.id, name: this.task.trainingInformation.modelID }, this.model);
+    }
+}

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

@@ -0,0 +1,33 @@
+import type tf from "@tensorflow/tfjs";
+import { List } from "immutable";
+import type { Model, Task } from "../../index.js";
+import { EpochLogs } from "../../models/model.js";
+export interface RoundLogs {
+    round: number;
+    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<RoundLogs>;
+}

package/dist/training/trainer/trainer.js

@@ -0,0 +1,52 @@
+import { List } from "immutable";
+/** 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;
+    }
+    /**
+     * 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");
+        }
+        await this.onRoundBegin(0);
+        this.training = this.model.train(dataset, valDataset, this.#epochs);
+        for await (const logs of this.training) {
+            // for now, round (sharing on network) == epoch (full pass over local data)
+            yield {
+                round: logs.epoch,
+                epochs: List.of(logs),
+            };
+            if (logs.epoch % this.#roundDuration === 0) {
+                const round = Math.trunc(logs.epoch / this.#roundDuration);
+                await this.onRoundEnd(round);
+                await this.onRoundBegin(round);
+            }
+        }
+        const round = Math.trunc(this.#epochs / this.#roundDuration);
+        await this.onRoundEnd(round);
+        this.training = undefined;
+    }
+}

package/dist/{core/training → training}/trainer/trainer_builder.d.ts

@@ -1,13 +1,12 @@
-import { Client, Task, TrainingInformant, Memory } from '../..';
-import { Trainer } from './trainer';
+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;
-    private readonly trainingInformant;
-    constructor(memory: Memory, task: Task, trainingInformant: TrainingInformant);
+    constructor(memory: Memory, task: Task);
     /**
      * Builds a trainer object.
      *
@@ -15,11 +14,10 @@ export declare class TrainerBuilder {
      * @param distributed whether to build a distributed or local trainer
      * @returns
      */
-    build(client: Client, distributed?: boolean): Promise<Trainer>;
+    build(client: clients.Client, distributed?: boolean): Promise<Trainer>;
     /**
-     * If a model exists in memory, laod it, otherwise load model from server
+     * If a model exists in memory, load it, otherwise load model from server
      * @returns
      */
     private getModel;
-    private updateModelInformation;
 }

package/dist/training/trainer/trainer_builder.js

@@ -0,0 +1,43 @@
+import { ModelType } from '../../index.js';
+import { DistributedTrainer } from './distributed_trainer.js';
+import { LocalTrainer } from './local_trainer.js';
+/**
+ * A class that helps build the Trainer and auxiliary classes.
+ */
+export class TrainerBuilder {
+    memory;
+    task;
+    constructor(memory, task) {
+        this.memory = memory;
+        this.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
+     */
+    async build(client, distributed = false) {
+        const model = await this.getModel(client);
+        if (distributed) {
+            return new DistributedTrainer(this.task, this.memory, model, client);
+        }
+        else {
+            return new LocalTrainer(this.task, this.memory, model);
+        }
+    }
+    /**
+     * If a model exists in memory, load it, otherwise load model from server
+     * @returns
+     */
+    async getModel(client) {
+        const modelID = this.task.trainingInformation?.modelID;
+        if (modelID === undefined) {
+            throw new TypeError('model ID is undefined');
+        }
+        const info = { type: ModelType.WORKING, taskID: this.task.id, name: modelID };
+        const model = await (await this.memory.contains(info) ? this.memory.getModel(info) : client.getLatestModel());
+        return model;
+    }
+}

package/dist/types.d.ts

@@ -0,0 +1,8 @@
+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/types.js

@@ -0,0 +1 @@
+export {};

package/dist/utils/event_emitter.d.ts

@@ -0,0 +1,40 @@
+type Listener<T> = (_: T) => void;
+/**
+ * Call handlers on given events
+ *
+ * @typeParam I object/mapping from event name to emitted value type
+ */
+export declare class EventEmitter<I extends Record<string, unknown>> {
+    private listeners;
+    /**
+     * @param initialListeners object/mapping of event name to listener, as if using `on` on created instance
+     */
+    constructor(initialListeners?: {
+        [E in keyof I]?: Listener<I[E]>;
+    });
+    /**
+     * Register listener to call on event
+     *
+     * @param event event name to listen to
+     * @param listener handler to call
+     */
+    on<E extends keyof I>(event: E, listener: Listener<I[E]>): void;
+    /**
+     * Register listener to call once on next event
+     *
+     * @param event event name to listen to
+     * @param listener handler to call next time
+     */
+    once<E extends keyof I>(event: E, listener: Listener<I[E]>): void;
+    /**
+     * Send value to registered listeners of event name
+     *
+     * @param event send to listeners of event name
+     * @param value what to call listeners with
+     */
+    emit<E extends keyof I>(event: E, value: I[E]): void;
+}
+/** `EventEmitter` for all events */
+export declare class Sink extends EventEmitter<Record<string, unknown>> {
+}
+export {};

package/dist/utils/event_emitter.js

@@ -0,0 +1,57 @@
+// inspired by https://danilafe.com/blog/typescript_typesafe_events/
+import { List } from 'immutable';
+/**
+ * Call handlers on given events
+ *
+ * @typeParam I object/mapping from event name to emitted value type
+ */
+export class EventEmitter {
+    listeners = {};
+    /**
+     * @param initialListeners object/mapping of event name to listener, as if using `on` on created instance
+     */
+    constructor(initialListeners = {}) {
+        for (const event in initialListeners) {
+            const listener = initialListeners[event];
+            if (listener !== undefined) {
+                this.on(event, listener);
+            }
+        }
+    }
+    /**
+     * Register listener to call on event
+     *
+     * @param event event name to listen to
+     * @param listener handler to call
+     */
+    on(event, listener) {
+        const eventListeners = this.listeners[event] ?? List();
+        this.listeners[event] = eventListeners.push([false, listener]);
+    }
+    /**
+     * Register listener to call once on next event
+     *
+     * @param event event name to listen to
+     * @param listener handler to call next time
+     */
+    once(event, listener) {
+        const eventListeners = this.listeners[event] ?? List();
+        this.listeners[event] = eventListeners.push([true, listener]);
+    }
+    /**
+     * Send value to registered listeners of event name
+     *
+     * @param event send to listeners of event name
+     * @param value what to call listeners with
+     */
+    emit(event, value) {
+        const eventListeners = this.listeners[event] ?? List();
+        this.listeners[event] = eventListeners.filterNot(([once]) => once);
+        eventListeners.forEach(([_, listener]) => {
+            listener(value);
+        });
+    }
+}
+/** `EventEmitter` for all events */
+export class Sink extends EventEmitter {
+}

package/dist/validation/index.d.ts

@@ -0,0 +1 @@
+export { Validator } from './validator.js';

package/dist/validation/index.js

@@ -0,0 +1 @@
+export { Validator } from './validator.js';

package/dist/validation/validator.d.ts

@@ -0,0 +1,28 @@
+import { List } from 'immutable';
+import type { data, Model, Task, Logger, client as clients, Memory, ModelSource, Features } from '../index.js';
+export declare class Validator {
+    readonly task: Task;
+    readonly logger: Logger;
+    private readonly memory;
+    private readonly source?;
+    private readonly client?;
+    private readonly graphInformant;
+    private size;
+    private _confusionMatrix;
+    constructor(task: Task, logger: Logger, memory: Memory, source?: ModelSource | undefined, client?: clients.Client | undefined);
+    private getLabel;
+    assess(data: data.Data, useConfusionMatrix?: boolean): Promise<Array<{
+        groundTruth: number;
+        pred: number;
+        features: Features;
+    }>>;
+    predict(data: data.Data): Promise<Array<{
+        features: Features;
+        pred: number;
+    }>>;
+    getModel(): Promise<Model>;
+    get accuracyData(): List<number>;
+    get accuracy(): number;
+    get visitedSamples(): number;
+    get confusionMatrix(): number[][] | undefined;
+}

package/dist/validation/validator.js

@@ -0,0 +1,132 @@
+import { List } from 'immutable';
+import * as tf from '@tensorflow/tfjs';
+import { GraphInformant } from '../index.js';
+export class Validator {
+    task;
+    logger;
+    memory;
+    source;
+    client;
+    graphInformant = new GraphInformant();
+    size = 0;
+    _confusionMatrix;
+    constructor(task, logger, memory, source, client) {
+        this.task = task;
+        this.logger = logger;
+        this.memory = memory;
+        this.source = source;
+        this.client = client;
+        if (source === undefined && client === undefined) {
+            throw new Error('To initialize a Validator, either or both a source and client need to be specified');
+        }
+    }
+    async getLabel(ys) {
+        switch (ys.shape[1]) {
+            case 1:
+                return await ys.greaterEqual(tf.scalar(0.5)).data();
+            case 2:
+                return await ys.argMax(1).data();
+            default:
+                throw new Error(`unable to reduce tensor of shape: ${ys.shape.toString()}`);
+        }
+    }
+    async assess(data, useConfusionMatrix = false) {
+        const batchSize = this.task.trainingInformation?.batchSize;
+        if (batchSize === undefined) {
+            throw new TypeError('Batch size is undefined');
+        }
+        const model = await this.getModel();
+        let features = [];
+        const groundTruth = [];
+        let hits = 0;
+        // Get model predictions per batch and flatten the result
+        // Also build the features and ground truth arrays
+        const predictions = (await data.preprocess().dataset.batch(batchSize)
+            .mapAsync(async (e) => {
+            if (typeof e === 'object' && 'xs' in e && 'ys' in e) {
+                const xs = e.xs;
+                const ys = await this.getLabel(e.ys);
+                const pred = await this.getLabel(await model.predict(xs));
+                const currentFeatures = await xs.array();
+                if (Array.isArray(currentFeatures)) {
+                    features = features.concat(currentFeatures);
+                }
+                else {
+                    throw new TypeError('Data format is incorrect');
+                }
+                groundTruth.push(...Array.from(ys));
+                this.size += xs.shape[0];
+                hits += List(pred).zip(List(ys)).filter(([p, y]) => p === y).size;
+                // TODO: Confusion Matrix stats
+                const currentAccuracy = hits / this.size;
+                this.graphInformant.updateAccuracy(currentAccuracy);
+                return Array.from(pred);
+            }
+            else {
+                throw new Error('Input data is missing a feature or the label');
+            }
+        }).toArray()).flat();
+        this.logger.success(`Obtained validation accuracy of ${this.accuracy}`);
+        this.logger.success(`Visited ${this.visitedSamples} samples`);
+        if (useConfusionMatrix) {
+            try {
+                this._confusionMatrix = tf.math.confusionMatrix([], [], 0).arraySync();
+            }
+            catch (e) {
+                console.error(e instanceof Error ? e.message : e);
+                throw new Error('Failed to compute the confusion matrix');
+            }
+        }
+        return List(groundTruth)
+            .zip(List(predictions), List(features))
+            .map(([gt, p, f]) => ({ groundTruth: gt, pred: p, features: f }))
+            .toArray();
+    }
+    async predict(data) {
+        const batchSize = this.task.trainingInformation?.batchSize;
+        if (batchSize === undefined) {
+            throw new TypeError('Batch size is undefined');
+        }
+        const model = await this.getModel();
+        let features = [];
+        // Get model prediction per batch and flatten the result
+        // Also incrementally build the features array
+        const predictions = (await data.preprocess().dataset.batch(batchSize)
+            .mapAsync(async (e) => {
+            const xs = e;
+            const currentFeatures = await xs.array();
+            if (Array.isArray(currentFeatures)) {
+                features = features.concat(currentFeatures);
+            }
+            else {
+                throw new TypeError('Data format is incorrect');
+            }
+            const pred = await this.getLabel(await model.predict(xs));
+            return Array.from(pred);
+        }).toArray()).flat();
+        return List(features).zip(List(predictions))
+            .map(([f, p]) => ({ features: f, pred: p }))
+            .toArray();
+    }
+    async getModel() {
+        if (this.source !== undefined && await this.memory.contains(this.source)) {
+            return await this.memory.getModel(this.source);
+        }
+        if (this.client !== undefined) {
+            return await this.client.getLatestModel();
+        }
+        throw new Error('Could not load the model');
+    }
+    get accuracyData() {
+        return this.graphInformant.data();
+    }
+    get accuracy() {
+        return this.graphInformant.accuracy();
+    }
+    get visitedSamples() {
+        return this.size;
+    }
+    get confusionMatrix() {
+        return this._confusionMatrix;
+    }
+}

package/dist/weights/aggregation.d.ts

@@ -0,0 +1,21 @@
+import type { TensorLike } from './weights_container.js';
+import { WeightsContainer } from './weights_container.js';
+type WeightsLike = Iterable<TensorLike>;
+/**
+ * Sums the given iterable of weights entry-wise.
+ * @param weights The list of weights to sum
+ * @returns The summed weights
+ */
+export declare function sum(weights: Iterable<WeightsLike | WeightsContainer>): WeightsContainer;
+/**
+ * Computes the successive entry-wise difference between the weights of the given iterable.
+ * The operation is not commutative w.r.t. the iterable's ordering.
+ */
+export declare function diff(weights: Iterable<WeightsLike | WeightsContainer>): WeightsContainer;
+/**
+ * Averages the given iterable of weights entry-wise.
+ * @param weights The list of weights to average
+ * @returns The averaged weights
+ */
+export declare function avg(weights: Iterable<WeightsLike | WeightsContainer>): WeightsContainer;
+export {};