@epfml/discojs 2.1.1 → 2.1.2-p20240506085037.0

package/dist/aggregator/base.d.ts

@@ -0,0 +1,180 @@
+import { Map, Set } from 'immutable';
+import type { client, Model, AsyncInformant } from '../index.js';
+export declare enum AggregationStep {
+    ADD = 0,
+    UPDATE = 1,
+    AGGREGATE = 2
+}
+/**
+ * Main, abstract, aggregator class whose role is to buffer contributions and to produce
+ * a result based off their aggregation, whenever some defined condition is met.
+ */
+export declare abstract class Base<T> {
+    /**
+     * The Model whose weights are updated on aggregation.
+     */
+    protected _model?: Model | undefined;
+    /**
+     * The round cut-off for contributions.
+     */
+    protected readonly roundCutoff: number;
+    /**
+     * The number of communication rounds occurring during any given aggregation round.
+     */
+    readonly communicationRounds: number;
+    /**
+     * Contains the ids of all active nodes, i.e. members of the aggregation group at
+     * a given round. It is a subset of all the nodes available in the network.
+     */
+    protected _nodes: Set<client.NodeID>;
+    /**
+     * Contains the contributions received from active nodes, accessible by node id.
+     * It defines the effective aggregation group, which is possibly a subset
+     * of all active nodes, depending on the aggregation scheme.
+     */
+    protected contributions: Map<number, Map<client.NodeID, T>>;
+    /**
+     * Emits the aggregation event whenever an aggregation step is performed.
+     * Triggers the resolve of the result promise and the preparation for the
+     * next aggregation round.
+     */
+    private readonly eventEmitter;
+    protected informant?: AsyncInformant<T>;
+    /**
+     * The result promise which, on resolve, will contain the current aggregation result.
+     * This promise should be fetched by any object making use of an aggregator, in order
+     * to await upon aggregation.
+     */
+    protected result: Promise<T>;
+    /**
+     * The current aggregation round, used for assessing whether a node contribution is recent enough
+     * or not.
+     */
+    protected _round: number;
+    /**
+     * The current communication round. A single aggregation round is made of possibly multiple
+     * communication rounds. This makes the aggregator free to perform intermediate aggregation
+     * steps based off communication with its nodes. Overall, this allows for more complex
+     * aggregation schemes requiring an exchange of information between nodes before aggregating.
+     */
+    protected _communicationRound: number;
+    constructor(
+    /**
+     * The Model whose weights are updated on aggregation.
+     */
+    _model?: Model | undefined, 
+    /**
+     * The round cut-off for contributions.
+     */
+    roundCutoff?: number, 
+    /**
+     * The number of communication rounds occurring during any given aggregation round.
+     */
+    communicationRounds?: number);
+    /**
+     * Adds a node's contribution to the aggregator for the given aggregation and communication rounds.
+     * The contribution will be aggregated during the next aggregation step.
+     * @param nodeId The node's id
+     * @param contribution The node's contribution
+     * @param round For which aggregation round the contribution was made
+     * @param communicationRound For which communication round the contribution was made
+     */
+    abstract add(nodeId: client.NodeID, contribution: T, round: number, communicationRound?: number): boolean;
+    /**
+     * Performs an aggregation step over the received node contributions.
+     * Must store the aggregation's result in the aggregator's result promise.
+     */
+    abstract aggregate(): void;
+    registerObserver(informant: AsyncInformant<T>): void;
+    /**
+     * Returns whether the given round is recent enough, dependent on the
+     * aggregator's round cutoff.
+     * @param round The round
+     * @returns True if the round is recent enough, false otherwise
+     */
+    isWithinRoundCutoff(round: number): boolean;
+    /**
+     * Logs useful messages during the various aggregation steps.
+     * @param step The aggregation step
+     * @param from The node which triggered the logging message
+     */
+    log(step: AggregationStep, from?: client.NodeID): void;
+    /**
+     * Sets the aggregator's TF.js model.
+     * @param model The new TF.js model
+     */
+    setModel(model: Model): void;
+    /**
+     * Adds a node's id to the set of active nodes. A node represents an active neighbor
+     * peer/client within the network, whom we are communicating with during this aggregation
+     * round.
+     * @param nodeId The node to be added
+     */
+    registerNode(nodeId: client.NodeID): boolean;
+    /**
+     * Overwrites the current set of active nodes with the given one. A node represents
+     * an active neighbor peer/client within the network, whom we are communicating with
+     * during this aggregation round.
+     * @param nodeIds The new set of nodes
+     */
+    setNodes(nodeIds: Set<client.NodeID>): void;
+    /**
+     * Empties the current set of "nodes". Usually called at the end of an aggregation round,
+     * if the set of nodes is meant to change or to be actualized.
+     */
+    resetNodes(): void;
+    /**
+     * Sets the aggregator's round number. To be used whenever the aggregator is out of sync
+     * with the network's round.
+     * @param round The new round
+     */
+    setRound(round: number): void;
+    /**
+     * Emits the event containing the aggregation result, which allows the result
+     * promise to resolve and for the next aggregation round to take place.
+     * @param aggregated The aggregation result
+     */
+    protected emit(aggregated: T): void;
+    /**
+     * Updates the aggregator's state to proceed to the next communication round.
+     * If all communication rounds were performed, proceeds to the next aggregation round
+     * and empties the collection of stored contributions.
+     */
+    nextRound(): void;
+    private makeResult;
+    /**
+     * Aggregation steps are performed asynchronously, yet can be awaited upon when required.
+     * This function gives access to the current aggregation result's promise, which will
+     * eventually resolve and contain the result of the very next aggregation step, at the
+     * time of the function call.
+     * @returns The promise containing the aggregation result
+     */
+    receiveResult(): Promise<T>;
+    /**
+     * Constructs the payloads sent to other nodes as contribution.
+     * @param base Object from which the payload is computed
+     */
+    abstract makePayloads(base: T): Map<client.NodeID, T>;
+    abstract isFull(): boolean;
+    /**
+     * The set of node ids, representing our neighbors within the network.
+     */
+    get nodes(): Set<client.NodeID>;
+    /**
+     * The aggregation round.
+     */
+    get round(): number;
+    /**
+     * The aggregator's current size, defined by its number of contributions. The size is bounded by
+     * the amount of all active nodes times the number of communication rounds.
+     */
+    get size(): number;
+    /**
+     * The aggregator's current model.
+     */
+    get model(): Model | undefined;
+    /**
+     * The current communication round.
+     */
+    get communicationRound(): number;
+}

package/dist/aggregator/base.js

@@ -0,0 +1,236 @@
+import { Map, Set } from 'immutable';
+import { EventEmitter } from '../utils/event_emitter.js';
+export var AggregationStep;
+(function (AggregationStep) {
+    AggregationStep[AggregationStep["ADD"] = 0] = "ADD";
+    AggregationStep[AggregationStep["UPDATE"] = 1] = "UPDATE";
+    AggregationStep[AggregationStep["AGGREGATE"] = 2] = "AGGREGATE";
+})(AggregationStep || (AggregationStep = {}));
+/**
+ * Main, abstract, aggregator class whose role is to buffer contributions and to produce
+ * a result based off their aggregation, whenever some defined condition is met.
+ */
+export class Base {
+    _model;
+    roundCutoff;
+    communicationRounds;
+    /**
+     * Contains the ids of all active nodes, i.e. members of the aggregation group at
+     * a given round. It is a subset of all the nodes available in the network.
+     */
+    _nodes;
+    /**
+     * Contains the contributions received from active nodes, accessible by node id.
+     * It defines the effective aggregation group, which is possibly a subset
+     * of all active nodes, depending on the aggregation scheme.
+     */
+    contributions;
+    /**
+     * Emits the aggregation event whenever an aggregation step is performed.
+     * Triggers the resolve of the result promise and the preparation for the
+     * next aggregation round.
+     */
+    eventEmitter = new EventEmitter();
+    informant;
+    /**
+     * The result promise which, on resolve, will contain the current aggregation result.
+     * This promise should be fetched by any object making use of an aggregator, in order
+     * to await upon aggregation.
+     */
+    result;
+    /**
+     * The current aggregation round, used for assessing whether a node contribution is recent enough
+     * or not.
+     */
+    _round = 0;
+    /**
+     * The current communication round. A single aggregation round is made of possibly multiple
+     * communication rounds. This makes the aggregator free to perform intermediate aggregation
+     * steps based off communication with its nodes. Overall, this allows for more complex
+     * aggregation schemes requiring an exchange of information between nodes before aggregating.
+     */
+    _communicationRound = 0;
+    constructor(
+    /**
+     * The Model whose weights are updated on aggregation.
+     */
+    _model, 
+    /**
+     * The round cut-off for contributions.
+     */
+    roundCutoff = 0, 
+    /**
+     * The number of communication rounds occurring during any given aggregation round.
+     */
+    communicationRounds = 1) {
+        this._model = _model;
+        this.roundCutoff = roundCutoff;
+        this.communicationRounds = communicationRounds;
+        this.contributions = Map();
+        this._nodes = Set();
+        // Make the initial result promise
+        this.result = this.makeResult();
+        // On every aggregation, update the object's state to match the current aggregation
+        // and communication rounds.
+        this.eventEmitter.on('aggregation', () => {
+            this.nextRound();
+        });
+    }
+    registerObserver(informant) {
+        this.informant = informant;
+    }
+    /**
+     * Returns whether the given round is recent enough, dependent on the
+     * aggregator's round cutoff.
+     * @param round The round
+     * @returns True if the round is recent enough, false otherwise
+     */
+    isWithinRoundCutoff(round) {
+        return this.round - round <= this.roundCutoff;
+    }
+    /**
+     * Logs useful messages during the various aggregation steps.
+     * @param step The aggregation step
+     * @param from The node which triggered the logging message
+     */
+    log(step, from) {
+        switch (step) {
+            case AggregationStep.ADD:
+                console.log(`> Adding contribution from node ${from ?? '"unknown"'} for round (${this.communicationRound}, ${this.round})`);
+                break;
+            case AggregationStep.UPDATE:
+                if (from === undefined) {
+                    return;
+                }
+                console.log(`> Updating contribution from node ${from} for round (${this.communicationRound}, ${this.round})`);
+                break;
+            case AggregationStep.AGGREGATE:
+                console.log('*'.repeat(80));
+                console.log(`Buffer is full. Aggregating weights for round (${this.communicationRound}, ${this.round})\n`);
+                break;
+            default: {
+                const _ = step;
+                throw new Error('should never happen');
+            }
+        }
+    }
+    /**
+     * Sets the aggregator's TF.js model.
+     * @param model The new TF.js model
+     */
+    setModel(model) {
+        this._model = model;
+    }
+    /**
+     * Adds a node's id to the set of active nodes. A node represents an active neighbor
+     * peer/client within the network, whom we are communicating with during this aggregation
+     * round.
+     * @param nodeId The node to be added
+     */
+    registerNode(nodeId) {
+        if (!this.nodes.has(nodeId)) {
+            this._nodes = this._nodes.add(nodeId);
+            return true;
+        }
+        return false;
+    }
+    /**
+     * Overwrites the current set of active nodes with the given one. A node represents
+     * an active neighbor peer/client within the network, whom we are communicating with
+     * during this aggregation round.
+     * @param nodeIds The new set of nodes
+     */
+    setNodes(nodeIds) {
+        this._nodes = nodeIds;
+    }
+    /**
+     * Empties the current set of "nodes". Usually called at the end of an aggregation round,
+     * if the set of nodes is meant to change or to be actualized.
+     */
+    resetNodes() {
+        this._nodes = Set();
+    }
+    /**
+     * Sets the aggregator's round number. To be used whenever the aggregator is out of sync
+     * with the network's round.
+     * @param round The new round
+     */
+    setRound(round) {
+        if (round > this.round) {
+            this._round = round;
+        }
+    }
+    /**
+     * Emits the event containing the aggregation result, which allows the result
+     * promise to resolve and for the next aggregation round to take place.
+     * @param aggregated The aggregation result
+     */
+    emit(aggregated) {
+        this.eventEmitter.emit('aggregation', aggregated);
+    }
+    /**
+     * Updates the aggregator's state to proceed to the next communication round.
+     * If all communication rounds were performed, proceeds to the next aggregation round
+     * and empties the collection of stored contributions.
+     */
+    nextRound() {
+        if (++this._communicationRound === this.communicationRounds) {
+            this._communicationRound = 0;
+            this._round++;
+            this.contributions = Map();
+        }
+        this.result = this.makeResult();
+        this.informant?.update();
+    }
+    async makeResult() {
+        return await new Promise((resolve) => {
+            this.eventEmitter.once('aggregation', (w) => {
+                resolve(w);
+            });
+        });
+    }
+    /**
+     * Aggregation steps are performed asynchronously, yet can be awaited upon when required.
+     * This function gives access to the current aggregation result's promise, which will
+     * eventually resolve and contain the result of the very next aggregation step, at the
+     * time of the function call.
+     * @returns The promise containing the aggregation result
+     */
+    async receiveResult() {
+        return await this.result;
+    }
+    /**
+     * The set of node ids, representing our neighbors within the network.
+     */
+    get nodes() {
+        return this._nodes;
+    }
+    /**
+     * The aggregation round.
+     */
+    get round() {
+        return this._round;
+    }
+    /**
+     * The aggregator's current size, defined by its number of contributions. The size is bounded by
+     * the amount of all active nodes times the number of communication rounds.
+     */
+    get size() {
+        return this.contributions
+            .valueSeq()
+            .map((m) => m.size)
+            .reduce((totalSize, size) => totalSize + size) ?? 0;
+    }
+    /**
+     * The aggregator's current model.
+     */
+    get model() {
+        return this._model;
+    }
+    /**
+     * The current communication round.
+     */
+    get communicationRound() {
+        return this._communicationRound;
+    }
+}

package/dist/aggregator/get.d.ts

@@ -0,0 +1,16 @@
+import type { Task } from '../index.js';
+import { aggregator } from '../index.js';
+/**
+ * Enumeration of the available types of aggregator.
+ */
+export declare enum AggregatorChoice {
+    MEAN = 0,
+    SECURE = 1,
+    BANDIT = 2
+}
+/**
+ * Provides the aggregator object adequate to the given task.
+ * @param task The task
+ * @returns The aggregator
+ */
+export declare function getAggregator(task: Task): aggregator.Aggregator;

package/dist/aggregator/get.js

@@ -0,0 +1,31 @@
+import { aggregator } from '../index.js';
+/**
+ * Enumeration of the available types of aggregator.
+ */
+export var AggregatorChoice;
+(function (AggregatorChoice) {
+    AggregatorChoice[AggregatorChoice["MEAN"] = 0] = "MEAN";
+    AggregatorChoice[AggregatorChoice["SECURE"] = 1] = "SECURE";
+    AggregatorChoice[AggregatorChoice["BANDIT"] = 2] = "BANDIT";
+})(AggregatorChoice || (AggregatorChoice = {}));
+/**
+ * Provides the aggregator object adequate to the given task.
+ * @param task The task
+ * @returns The aggregator
+ */
+export function getAggregator(task) {
+    const error = new Error('not implemented');
+    switch (task.trainingInformation.aggregator) {
+        case AggregatorChoice.MEAN:
+            return new aggregator.MeanAggregator();
+        case AggregatorChoice.BANDIT:
+            throw error;
+        case AggregatorChoice.SECURE:
+            if (task.trainingInformation.scheme !== 'decentralized') {
+                throw new Error('secure aggregation is currently supported for decentralized only');
+            }
+            return new aggregator.SecureAggregator();
+        default:
+            return new aggregator.MeanAggregator();
+    }
+}

package/dist/aggregator/index.d.ts

@@ -0,0 +1,7 @@
+import type { WeightsContainer } from '../weights/index.js';
+import type { Base } from './base.js';
+export { Base as AggregatorBase, AggregationStep } from './base.js';
+export { MeanAggregator } from './mean.js';
+export { SecureAggregator } from './secure.js';
+export { getAggregator, AggregatorChoice } from './get.js';
+export type Aggregator = Base<WeightsContainer>;

package/dist/aggregator/index.js

@@ -0,0 +1,4 @@
+export { Base as AggregatorBase, AggregationStep } from './base.js';
+export { MeanAggregator } from './mean.js';
+export { SecureAggregator } from './secure.js';
+export { getAggregator, AggregatorChoice } from './get.js';

package/dist/aggregator/mean.d.ts

@@ -0,0 +1,23 @@
+import type { Map } from 'immutable';
+import { Base as Aggregator } from './base.js';
+import type { Model, WeightsContainer, client } from '../index.js';
+/**
+ * Mean aggregator whose aggregation step consists in computing the mean of the received weights.
+ */
+export declare class MeanAggregator extends Aggregator<WeightsContainer> {
+    /**
+     * The threshold t to fulfill to trigger an aggregation step. It can either be:
+     * - relative: 0 < t <= 1, thus requiring t * |nodes| contributions
+     * - absolute: t > 1, thus requiring t contributions
+     */
+    readonly threshold: number;
+    constructor(model?: Model, roundCutoff?: number, threshold?: number);
+    /**
+     * Checks whether the contributions buffer is full, according to the set threshold.
+     * @returns Whether the contributions buffer is full
+     */
+    isFull(): boolean;
+    add(nodeId: client.NodeID, contribution: WeightsContainer, round: number): boolean;
+    aggregate(): void;
+    makePayloads(weights: WeightsContainer): Map<client.NodeID, WeightsContainer>;
+}

package/dist/aggregator/mean.js

@@ -0,0 +1,69 @@
+import { AggregationStep, Base as Aggregator } from './base.js';
+import { aggregation } from '../index.js';
+/**
+ * Mean aggregator whose aggregation step consists in computing the mean of the received weights.
+ */
+export class MeanAggregator extends Aggregator {
+    /**
+     * The threshold t to fulfill to trigger an aggregation step. It can either be:
+     * - relative: 0 < t <= 1, thus requiring t * |nodes| contributions
+     * - absolute: t > 1, thus requiring t contributions
+     */
+    threshold;
+    constructor(model, roundCutoff = 0, threshold = 1) {
+        super(model, roundCutoff, 1);
+        // Default threshold is 100% of node participation
+        if (threshold === undefined) {
+            this.threshold = 1;
+            // Threshold must be positive
+        }
+        else if (threshold <= 0) {
+            throw new Error('threshold must be positive');
+            // Thresholds greater than 1 are considered absolute instead of relative to the number of nodes
+        }
+        else if (threshold > 1 && Math.round(threshold) !== threshold) {
+            throw new Error('absolute thresholds must integers');
+        }
+        else {
+            this.threshold = threshold;
+        }
+    }
+    /**
+     * Checks whether the contributions buffer is full, according to the set threshold.
+     * @returns Whether the contributions buffer is full
+     */
+    isFull() {
+        if (this.threshold <= 1) {
+            const contribs = this.contributions.get(this.communicationRound);
+            if (contribs === undefined) {
+                return false;
+            }
+            return contribs.size >= this.threshold * this.nodes.size;
+        }
+        return this.contributions.size >= this.threshold;
+    }
+    add(nodeId, contribution, round) {
+        if (this.nodes.has(nodeId) && this.isWithinRoundCutoff(round)) {
+            this.log(this.contributions.hasIn([0, nodeId]) ? AggregationStep.UPDATE : AggregationStep.ADD, nodeId);
+            this.contributions = this.contributions.setIn([0, nodeId], contribution);
+            this.informant?.update();
+            if (this.isFull()) {
+                this.aggregate();
+            }
+            return true;
+        }
+        return false;
+    }
+    aggregate() {
+        this.log(AggregationStep.AGGREGATE);
+        const result = aggregation.avg(this.contributions.get(0)?.values());
+        if (this.model !== undefined) {
+            this.model.weights = result;
+        }
+        this.emit(result);
+    }
+    makePayloads(weights) {
+        // Communicate our local weights to every other node, be it a peer or a server
+        return this.nodes.toMap().map(() => weights);
+    }
+}

package/dist/aggregator/secure.d.ts

@@ -0,0 +1,27 @@
+import { Map, List } from 'immutable';
+import { Base as Aggregator } from './base.js';
+import type { Model, WeightsContainer, client } from '../index.js';
+/**
+ * Aggregator implementing secure multi-party computation for decentralized learning.
+ * An aggregation consists of two communication rounds:
+ * - first, nodes communicate their secret shares to each other;
+ * - then, they sum their received shares and communicate the result.
+ * Finally, nodes are able to average the received partial sums to establish the aggregation result.
+ */
+export declare class SecureAggregator extends Aggregator<WeightsContainer> {
+    private readonly maxShareValue;
+    constructor(model?: Model, maxShareValue?: number);
+    aggregate(): void;
+    add(nodeId: client.NodeID, contribution: WeightsContainer, round: number, communicationRound: number): boolean;
+    isFull(): boolean;
+    makePayloads(weights: WeightsContainer): Map<client.NodeID, WeightsContainer>;
+    /**
+     * Generate N additive shares that aggregate to the secret weights array, where N is the number of peers.
+     */
+    generateAllShares(secret: WeightsContainer): List<WeightsContainer>;
+    /**
+     * Generates one share in the same shape as the secret that is populated with values randomly chosen from
+     * a uniform distribution between (-maxShareValue, maxShareValue).
+     */
+    generateRandomShare(secret: WeightsContainer): WeightsContainer;
+}