@epfml/discojs 2.0.0 → 2.1.2-p20240506085037.0

package/dist/aggregator/secure.js

@@ -0,0 +1,91 @@
+import { Map, List, Range } from 'immutable';
+import * as tf from '@tensorflow/tfjs';
+import { AggregationStep, Base as Aggregator } from './base.js';
+import { aggregation } 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 class SecureAggregator extends Aggregator {
+    maxShareValue;
+    constructor(model, maxShareValue = 100) {
+        super(model, 0, 2);
+        this.maxShareValue = maxShareValue;
+    }
+    aggregate() {
+        this.log(AggregationStep.AGGREGATE);
+        if (this.communicationRound === 0) {
+            // Sum the received shares
+            const result = aggregation.sum(this.contributions.get(0)?.values());
+            this.emit(result);
+        }
+        else if (this.communicationRound === 1) {
+            // Average the received partial sums
+            const result = aggregation.avg(this.contributions.get(1)?.values());
+            if (this.model !== undefined) {
+                this.model.weights = result;
+            }
+            this.emit(result);
+        }
+        else {
+            throw new Error('communication round is out of bounds');
+        }
+    }
+    add(nodeId, contribution, round, communicationRound) {
+        if (this.nodes.has(nodeId) && this.isWithinRoundCutoff(round)) {
+            this.log(this.contributions.hasIn([communicationRound, nodeId]) ? AggregationStep.UPDATE : AggregationStep.ADD, nodeId);
+            this.contributions = this.contributions.setIn([communicationRound, nodeId], contribution);
+            this.informant?.update();
+            if (this.isFull()) {
+                this.aggregate();
+            }
+            return true;
+        }
+        return false;
+    }
+    isFull() {
+        const contribs = this.contributions.get(this.communicationRound);
+        if (contribs === undefined) {
+            return false;
+        }
+        return contribs.size === this.nodes.size;
+    }
+    makePayloads(weights) {
+        if (this.communicationRound === 0) {
+            const shares = this.generateAllShares(weights);
+            // Abitrarily assign our shares to the available nodes
+            return Map(List(this.nodes).zip(shares));
+        }
+        else {
+            // Send our partial sum to every other nodes
+            return this.nodes.toMap().map(() => weights);
+        }
+    }
+    /**
+     * Generate N additive shares that aggregate to the secret weights array, where N is the number of peers.
+     */
+    generateAllShares(secret) {
+        if (this.nodes.size === 0) {
+            throw new Error('too few participants to generate shares');
+        }
+        // Generate N-1 shares
+        const shares = Range(0, this.nodes.size - 1)
+            .map(() => this.generateRandomShare(secret))
+            .toList();
+        // The last share completes the sum
+        return shares.push(secret.sub(aggregation.sum(shares)));
+    }
+    /**
+     * 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) {
+        const MAX_SEED_BITS = 47;
+        const random = crypto.getRandomValues(new BigInt64Array(1))[0];
+        const seed = Number(BigInt.asUintN(MAX_SEED_BITS, random));
+        return secret.map((t) => tf.randomUniform(t.shape, -this.maxShareValue, this.maxShareValue, 'float32', seed));
+    }
+}

package/dist/async_informant.d.ts

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

@@ -0,0 +1,42 @@
+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/client/base.d.ts

@@ -0,0 +1,76 @@
+import type { Set } from 'immutable';
+import type { Model, Task, WeightsContainer } from '../index.js';
+import type { NodeID } from './types.js';
+import type { EventConnection } from './event_connection.js';
+import type { Aggregator } from '../aggregator/index.js';
+/**
+ * Main, abstract, class representing a Disco client in a network, which handles
+ * communication with other nodes, be it peers or a server.
+ */
+export declare abstract class Base {
+    /**
+     * The network server's URL to connect to.
+     */
+    readonly url: URL;
+    /**
+     * The client's corresponding task.
+     */
+    readonly task: Task;
+    /**
+     * The client's aggregator.
+     */
+    readonly aggregator: Aggregator;
+    /**
+     * Own ID provided by the network's server.
+     */
+    protected _ownId?: NodeID;
+    /**
+     * The network's server.
+     */
+    protected _server?: EventConnection;
+    /**
+     * The aggregator's result produced after aggregation.
+     */
+    protected aggregationResult?: Promise<WeightsContainer>;
+    constructor(
+    /**
+     * The network server's URL to connect to.
+     */
+    url: URL, 
+    /**
+     * The client's corresponding task.
+     */
+    task: Task, 
+    /**
+     * The client's aggregator.
+     */
+    aggregator: Aggregator);
+    /**
+     * Handles the connection process from the client to any sort of network server.
+     */
+    connect(): Promise<void>;
+    /**
+     * Handles the disconnection process of the client from any sort of network server.
+     */
+    disconnect(): Promise<void>;
+    /**
+     * Fetches the latest model available on the network's server, for the adequate task.
+     * @returns The latest model
+     */
+    getLatestModel(): Promise<Model>;
+    /**
+     * Communication callback called at the beginning of every training round.
+     * @param _weights The most recent local weight updates
+     * @param _round The current training round
+     */
+    onRoundBeginCommunication(_weights: WeightsContainer, _round: number): Promise<void>;
+    /**
+     * Communication callback called the end of every training round.
+     * @param _weights The most recent local weight updates
+     * @param _round The current training round
+     */
+    onRoundEndCommunication(_weights: WeightsContainer, _round: number): Promise<void>;
+    get nodes(): Set<NodeID>;
+    get ownId(): NodeID;
+    get server(): EventConnection;
+}

package/dist/client/base.js

@@ -0,0 +1,88 @@
+import axios from 'axios';
+import { serialization } from '../index.js';
+/**
+ * Main, abstract, class representing a Disco client in a network, which handles
+ * communication with other nodes, be it peers or a server.
+ */
+export class Base {
+    url;
+    task;
+    aggregator;
+    /**
+     * Own ID provided by the network's server.
+     */
+    _ownId;
+    /**
+     * The network's server.
+     */
+    _server;
+    /**
+     * The aggregator's result produced after aggregation.
+     */
+    aggregationResult;
+    constructor(
+    /**
+     * The network server's URL to connect to.
+     */
+    url, 
+    /**
+     * The client's corresponding task.
+     */
+    task, 
+    /**
+     * The client's aggregator.
+     */
+    aggregator) {
+        this.url = url;
+        this.task = task;
+        this.aggregator = aggregator;
+    }
+    /**
+     * Handles the connection process from the client to any sort of network server.
+     */
+    async connect() { }
+    /**
+     * Handles the disconnection process of the client from any sort of network server.
+     */
+    async disconnect() { }
+    /**
+     * Fetches the latest model available on the network's server, for the adequate task.
+     * @returns The latest model
+     */
+    async getLatestModel() {
+        const url = new URL('', this.url.href);
+        if (!url.pathname.endsWith('/')) {
+            url.pathname += '/';
+        }
+        url.pathname += `tasks/${this.task.id}/model.json`;
+        const response = await axios.get(url.href, { responseType: 'arraybuffer' });
+        return await serialization.model.decode(new Uint8Array(response.data));
+    }
+    /**
+     * Communication callback called at the beginning of every training round.
+     * @param _weights The most recent local weight updates
+     * @param _round The current training round
+     */
+    async onRoundBeginCommunication(_weights, _round) { }
+    /**
+     * Communication callback called the end of every training round.
+     * @param _weights The most recent local weight updates
+     * @param _round The current training round
+     */
+    async onRoundEndCommunication(_weights, _round) { }
+    get nodes() {
+        return this.aggregator.nodes;
+    }
+    get ownId() {
+        if (this._ownId === undefined) {
+            throw new Error('the node is not connected');
+        }
+        return this._ownId;
+    }
+    get server() {
+        if (this._server === undefined) {
+            throw new Error('server undefined, not connected');
+        }
+        return this._server;
+    }
+}

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

@@ -0,0 +1,32 @@
+import { type WeightsContainer } from '../../index.js';
+import { Client } from '../index.js';
+import { type PeerConnection } from '../event_connection.js';
+import * as messages from './messages.js';
+/**
+ * Represents a decentralized client in a network of peers. Peers coordinate each other with the
+ * help of the network's server, yet only exchange payloads between each other. Communication
+ * with the server is based off regular WebSockets, whereas peer-to-peer communication uses
+ * WebRTC for Node.js.
+ */
+export declare class Base extends Client {
+    /**
+     * The pool of peers to communicate with during the current training round.
+     */
+    private pool?;
+    private connections?;
+    /**
+     * Send message to server that this client is ready for the next training round.
+     */
+    private waitForPeers;
+    protected sendMessagetoPeer(peer: PeerConnection, msg: messages.PeerMessage): void;
+    /**
+     * Creation of the WebSocket for the server, connection of client to that WebSocket,
+     * deals with message reception from the decentralized client's perspective (messages received by client).
+     */
+    private connectServer;
+    connect(): Promise<void>;
+    disconnect(): Promise<void>;
+    onRoundBeginCommunication(_: WeightsContainer, round: number): Promise<void>;
+    onRoundEndCommunication(weights: WeightsContainer, round: number): Promise<void>;
+    private receivePayloads;
+}

package/dist/client/decentralized/base.js

@@ -0,0 +1,192 @@
+import { Map, Set } from 'immutable';
+import { serialization } from '../../index.js';
+import { Client } from '../index.js';
+import { type } from '../messages.js';
+import { timeout } from '../utils.js';
+import { WebSocketServer, waitMessage, waitMessageWithTimeout } from '../event_connection.js';
+import { PeerPool } from './peer_pool.js';
+import * as messages from './messages.js';
+/**
+ * Represents a decentralized client in a network of peers. Peers coordinate each other with the
+ * help of the network's server, yet only exchange payloads between each other. Communication
+ * with the server is based off regular WebSockets, whereas peer-to-peer communication uses
+ * WebRTC for Node.js.
+ */
+export class Base extends Client {
+    /**
+     * The pool of peers to communicate with during the current training round.
+     */
+    pool;
+    connections;
+    /**
+     * Send message to server that this client is ready for the next training round.
+     */
+    async waitForPeers(round) {
+        console.info(`[${this.ownId}] is ready for round`, round);
+        // Broadcast our readiness
+        const readyMessage = { type: type.PeerIsReady };
+        if (this.server === undefined) {
+            throw new Error('server undefined, could not connect peers');
+        }
+        this.server.send(readyMessage);
+        // Wait for peers to be connected before sending any update information
+        try {
+            const receivedMessage = await waitMessageWithTimeout(this.server, type.PeersForRound);
+            if (this.nodes.size > 0) {
+                throw new Error('got new peer list from server but was already received for this round');
+            }
+            const peers = Set(receivedMessage.peers);
+            console.info(`[${this.ownId}] received peers for round:`, peers.toJS());
+            if (this.ownId !== undefined && peers.has(this.ownId)) {
+                throw new Error('received peer list contains our own id');
+            }
+            this.aggregator.setNodes(peers.add(this.ownId));
+            if (this.pool === undefined) {
+                throw new Error('waiting for peers but peer pool is undefined');
+            }
+            const connections = await this.pool.getPeers(peers, this.server, 
+            // Init receipt of peers weights
+            (conn) => { this.receivePayloads(conn, round); });
+            console.info(`[${this.ownId}] received peers for round ${round}:`, connections.keySeq().toJS());
+            return connections;
+        }
+        catch (e) {
+            console.error(e);
+            this.aggregator.setNodes(Set(this.ownId));
+            return Map();
+        }
+    }
+    sendMessagetoPeer(peer, msg) {
+        console.info(`[${this.ownId}] send message to peer`, msg.peer, msg);
+        peer.send(msg);
+    }
+    /**
+     * Creation of the WebSocket for the server, connection of client to that WebSocket,
+     * deals with message reception from the decentralized client's perspective (messages received by client).
+     */
+    async connectServer(url) {
+        const server = await WebSocketServer.connect(url, messages.isMessageFromServer, messages.isMessageToServer);
+        server.on(type.SignalForPeer, (event) => {
+            console.info(`[${this.ownId}] received signal from`, event.peer);
+            if (this.pool === undefined) {
+                throw new Error('received signal but peer pool is undefined');
+            }
+            this.pool.signal(event.peer, event.signal);
+        });
+        return server;
+    }
+    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 += `deai/${this.task.id}`;
+        this._server = await this.connectServer(serverURL);
+        const msg = {
+            type: type.ClientConnected
+        };
+        this.server.send(msg);
+        const peerIdMsg = await waitMessage(this.server, type.AssignNodeID);
+        console.info(`[${peerIdMsg.id}] assigned id generated by server`);
+        if (this._ownId !== undefined) {
+            throw new Error('received id from server but was already received');
+        }
+        this._ownId = peerIdMsg.id;
+        this.pool = new PeerPool(peerIdMsg.id);
+    }
+    async disconnect() {
+        // Disconnect from peers
+        await this.pool?.shutdown();
+        this.pool = undefined;
+        if (this.connections !== undefined) {
+            const peers = this.connections.keySeq().toSet();
+            this.aggregator.setNodes(this.aggregator.nodes.subtract(peers));
+        }
+        // Disconnect from server
+        await this.server?.disconnect();
+        this._server = undefined;
+        this._ownId = undefined;
+        return Promise.resolve();
+    }
+    async onRoundBeginCommunication(_, round) {
+        // Reset peers list at each round of training to make sure client works with an updated peers
+        // list, maintained by the server. Adds any received weights to the aggregator.
+        this.connections = await this.waitForPeers(round);
+        // Store the promise for the current round's aggregation result.
+        this.aggregationResult = this.aggregator.receiveResult();
+    }
+    async onRoundEndCommunication(weights, round) {
+        let result = weights;
+        // Perform the required communication rounds. Each communication round consists in sending our local payload,
+        // followed by an aggregation step triggered by the receipt of other payloads, and handled by the aggregator.
+        // A communication round's payload is the aggregation result of the previous communication round. The first
+        // communication round simply sends our training result, i.e. model weights updates. This scheme allows for
+        // the aggregator to define any complex multi-round aggregation mechanism.
+        for (let r = 0; r < this.aggregator.communicationRounds; r++) {
+            // Generate our payloads for this communication round and send them to all ready connected peers
+            if (this.connections !== undefined) {
+                const payloads = this.aggregator.makePayloads(result);
+                try {
+                    await Promise.all(payloads.map(async (payload, id) => {
+                        if (id === this.ownId) {
+                            this.aggregator.add(this.ownId, payload, round, r);
+                        }
+                        else {
+                            const connection = this.connections?.get(id);
+                            if (connection !== undefined) {
+                                const encoded = await serialization.weights.encode(payload);
+                                this.sendMessagetoPeer(connection, {
+                                    type: type.Payload,
+                                    peer: id,
+                                    round: r,
+                                    payload: encoded
+                                });
+                            }
+                        }
+                    }));
+                }
+                catch {
+                    throw new Error('error while sending weights');
+                }
+            }
+            if (this.aggregationResult === undefined) {
+                throw new TypeError('aggregation result promise is undefined');
+            }
+            // Wait for aggregation before proceeding to the next communication round.
+            // The current result will be used as payload for the eventual next communication round.
+            result = await Promise.race([this.aggregationResult, timeout()]);
+            // There is at least one communication round remaining
+            if (r < this.aggregator.communicationRounds - 1) {
+                // Reuse the aggregation result
+                this.aggregationResult = this.aggregator.receiveResult();
+            }
+        }
+        // Reset the peers list for the next round
+        this.aggregator.resetNodes();
+    }
+    receivePayloads(connections, round) {
+        console.info(`[${this.ownId}] Accepting new contributions for round ${round}`);
+        connections.forEach(async (connection, peerId) => {
+            let receivedPayloads = 0;
+            do {
+                try {
+                    const message = await waitMessageWithTimeout(connection, type.Payload);
+                    const decoded = serialization.weights.decode(message.payload);
+                    if (!this.aggregator.add(peerId, decoded, round, message.round)) {
+                        console.warn(`[${this.ownId}] Failed to add contribution from peer ${peerId}`);
+                    }
+                }
+                catch (e) {
+                    console.warn(e instanceof Error ? e.message : e);
+                }
+            } while (++receivedPayloads < this.aggregator.communicationRounds);
+        });
+    }
+}

package/dist/client/decentralized/index.d.ts

@@ -0,0 +1,2 @@
+export { Base as DecentralizedClient } from './base.js';
+export * as messages from './messages.js';

package/dist/client/decentralized/index.js

@@ -0,0 +1,2 @@
+export { Base as DecentralizedClient } from './base.js';
+export * as messages from './messages.js';

package/dist/client/decentralized/messages.d.ts

@@ -0,0 +1,28 @@
+import { weights } from '../../serialization/index.js';
+import { type SignalData } from './peer.js';
+import { type NodeID } from '../types.js';
+import { type, type ClientConnected, type AssignNodeID } from '../messages.js';
+export interface SignalForPeer {
+    type: type.SignalForPeer;
+    peer: NodeID;
+    signal: SignalData;
+}
+export interface PeerIsReady {
+    type: type.PeerIsReady;
+}
+export interface PeersForRound {
+    type: type.PeersForRound;
+    peers: NodeID[];
+}
+export interface Payload {
+    type: type.Payload;
+    peer: NodeID;
+    round: number;
+    payload: weights.Encoded;
+}
+export type MessageFromServer = AssignNodeID | SignalForPeer | PeersForRound;
+export type MessageToServer = ClientConnected | SignalForPeer | PeerIsReady;
+export type PeerMessage = Payload;
+export declare function isMessageFromServer(o: unknown): o is MessageFromServer;
+export declare function isMessageToServer(o: unknown): o is MessageToServer;
+export declare function isPeerMessage(o: unknown): o is PeerMessage;

package/dist/client/decentralized/messages.js

@@ -0,0 +1,44 @@
+import { weights } from '../../serialization/index.js';
+import { isNodeID } from '../types.js';
+import { type, hasMessageType } from '../messages.js';
+export function isMessageFromServer(o) {
+    if (!hasMessageType(o)) {
+        return false;
+    }
+    switch (o.type) {
+        case type.AssignNodeID:
+            return 'id' in o && isNodeID(o.id);
+        case type.SignalForPeer:
+            return 'peer' in o && isNodeID(o.peer) &&
+                'signal' in o; // TODO check signal content?
+        case type.PeersForRound:
+            return 'peers' in o && Array.isArray(o.peers) && o.peers.every(isNodeID);
+    }
+    return false;
+}
+export function isMessageToServer(o) {
+    if (!hasMessageType(o)) {
+        return false;
+    }
+    switch (o.type) {
+        case type.ClientConnected:
+            return true;
+        case type.SignalForPeer:
+            return 'peer' in o && isNodeID(o.peer) &&
+                'signal' in o; // TODO check signal content?
+        case type.PeerIsReady:
+            return true;
+    }
+    return false;
+}
+export function isPeerMessage(o) {
+    if (!hasMessageType(o)) {
+        return false;
+    }
+    switch (o.type) {
+        case type.Payload:
+            return ('peer' in o && isNodeID(o.peer) &&
+                'payload' in o && weights.isEncoded(o.payload));
+    }
+    return false;
+}

package/dist/client/decentralized/peer.d.ts

@@ -0,0 +1,40 @@
+/// <reference types="node" resolution-mode="require"/>
+import type { NodeID } from '../types.js';
+export type SignalData = {
+    type: 'answer' | 'offer' | 'pranswer' | 'rollback';
+    sdp?: string;
+} | {
+    type: 'transceiverRequest';
+    transceiverRequest: {
+        kind: string;
+    };
+} | {
+    type: 'renegotiate';
+    renegotiate: true;
+} | {
+    type: 'candidate';
+    candidate: RTCIceCandidate;
+};
+interface Events {
+    'close': () => void;
+    'connect': () => void;
+    'signal': (signal: SignalData) => void;
+    'data': (data: Buffer) => void;
+}
+export declare class Peer {
+    readonly id: NodeID;
+    private readonly peer;
+    private bufferSize?;
+    private sendCounter;
+    private sendQueue;
+    private receiving;
+    constructor(id: NodeID, initiator?: boolean);
+    send(msg: Buffer): void;
+    private flush;
+    get maxChunkSize(): number;
+    private chunk;
+    destroy(): Promise<void>;
+    signal(signal: SignalData): void;
+    on<K extends keyof Events>(event: K, listener: Events[K]): void;
+}
+export {};