federer/dist/src/client/Client.js

Experiments in asynchronous federated learning and decentralized learning

"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.Client = void 0;
const tslib_1 = require("tslib");
const worker_threads_1 = require("worker_threads");
const path = tslib_1.__importStar(require("path"));
const socket_io_client_1 = require("socket.io-client");
const Comlink = tslib_1.__importStar(require("comlink"));
const node_adapter_1 = tslib_1.__importDefault(require("comlink/dist/umd/node-adapter"));
class Client {
    constructor(serverUrl, numberDatapoints, worker, workerThread, options, logger) {
        this.worker = worker;
        this.workerThread = workerThread;
        this.options = options;
        this.logger = logger;
        this.socket = socket_io_client_1.io(serverUrl);
        this.socket.on("download", async (message) => {
            this.logger.info("started training");
            const tensors = message.weights;
            this.logger.debug(`Transfer to worker: ${new Date().getTime()}`);
            const newWeights = await this.worker.train(Comlink.transfer(tensors, [
                tensors instanceof ArrayBuffer ? tensors : tensors.buffer,
            ]), message.round);
            this.logger.info("finished training");
            this.upload(message.round, newWeights);
        });
        this.socket.on("connect", () => {
            this.logger.info("connected to the server");
            this.socket.emit("ready", {
                deltaUpdates: this.options.deltaUpdates,
                clientId: this.options.id,
                numberDatapoints: numberDatapoints,
                numberEpochs: this.options.trainOptions.epochs,
            });
        });
        this.socket.on("disconnect", (reason) => {
            this.logger.warn(`detected a disconnect from the server for the following reason: '${reason}'`);
        });
    }
    static async create(serverUrl, options, logger) {
        const workerThread = new worker_threads_1.Worker(path.join(__dirname, "ClientWorker.js"));
        const workerBuilder = Comlink.wrap(node_adapter_1.default(workerThread));
        const worker = await workerBuilder.build(options);
        const numberDatapoints = await worker.countNumberDataPoints();
        return new Client(serverUrl, numberDatapoints, worker, workerThread, options, logger);
    }
    upload(round, weights) {
        // If the delay is 0, we don't use `setTimeout`, but emit directly.
        // This avoids delaying the response by a tiny amount by going through the
        // event loop (which may have other jobs to do, which could in theory
        // block the response for a bit).
        const upload = () => this.socket.emit("upload", { round, weights });
        if (this.options.replyDelay > 0) {
            setTimeout(upload, this.options.replyDelay);
        }
        else {
            upload();
        }
    }
    async terminate() {
        this.socket.close();
        await this.workerThread.terminate();
    }
}
exports.Client = Client;
//# sourceMappingURL=Client.js.map