@epfml/discojs 3.0.1-p20240805130603.0 → 3.0.1-p20240809121820.0

package/dist/aggregator/base.js

@@ -1,5 +1,7 @@
+import createDebug from "debug";
 import { Map, Set } from 'immutable';
 import { EventEmitter } from '../utils/event_emitter.js';
+const debug = createDebug("discojs:aggregator");
 export var AggregationStep;
 (function (AggregationStep) {
     AggregationStep[AggregationStep["ADD"] = 0] = "ADD";
@@ -75,17 +77,16 @@ export class Base extends EventEmitter {
     log(step, from) {
         switch (step) {
             case AggregationStep.ADD:
-                console.log(`Adding contribution from node ${from ?? '"unknown"'} for round (${this.communicationRound}, ${this.round})`);
+                debug(`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})`);
+                debug(`updating contribution from node ${from ?? '"unknown"'} 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`);
+                debug(`buffer full, aggregating weights for round (${this.communicationRound}, ${this.round})`);
                 break;
             default: {
                 const _ = step;

package/dist/aggregator/mean.js

@@ -1,5 +1,7 @@
+import createDebug from "debug";
 import { AggregationStep, Base as Aggregator } from "./base.js";
 import { aggregation } from "../index.js";
+const debug = createDebug("discojs:aggregator:mean");
 /**
  * Mean aggregator whose aggregation step consists in computing the mean of the received weights.
  *
@@ -51,7 +53,8 @@ export class MeanAggregator extends Aggregator {
         else {
             // Print a warning regarding the default behavior when thresholdType is not specified
             if (thresholdType === undefined) {
-                console.warn("[WARN] Setting the aggregator's threshold to 100% of the nodes' contributions by default. " +
+                // TODO enforce validity by splitting features instead of warning
+                debug("[WARN] Setting the aggregator's threshold to 100% of the nodes' contributions by default. " +
                     "To instead wait for a single contribution, set thresholdType = 'absolute'");
                 this.#thresholdType = 'relative';
             }
@@ -72,9 +75,9 @@ export class MeanAggregator extends Aggregator {
             throw new Error("only a single communication round");
         if (!this.nodes.has(nodeId) || !this.isWithinRoundCutoff(round)) {
             if (!this.nodes.has(nodeId))
-                console.warn("Contribution rejected because node id is not registered");
+                debug(`contribution rejected because node ${nodeId} is not registered`);
             if (!this.isWithinRoundCutoff(round))
-                console.warn(`Contribution rejected because round ${round} is not within round cutoff`);
+                debug(`contribution rejected because round ${round} is not within cutoff`);
             return false;
         }
         this.log(this.contributions.hasIn([0, nodeId])

package/dist/client/decentralized/base.js

@@ -1,3 +1,4 @@
+import createDebug from "debug";
 import { Map, Set } from 'immutable';
 import { serialization } from "../../index.js";
 import { Client } from '../index.js';
@@ -6,6 +7,7 @@ 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';
+const debug = createDebug("discojs:client:decentralized");
 /**
  * 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
@@ -49,7 +51,7 @@ export class Base extends Client {
         };
         this.server.send(msg);
         const peerIdMsg = await waitMessage(this.server, type.AssignNodeID);
-        console.log(`[${peerIdMsg.id}] assigned id generated by server`);
+        debug(`[${peerIdMsg.id}] assigned id generated by server`);
         if (this._ownId !== undefined) {
             throw new Error('received id from server but was already received');
         }
@@ -123,11 +125,11 @@ export class Base extends Client {
             // this awaits the peer's weight update and adds it to 
             // our aggregator upon reception
             (conn) => { this.receivePayloads(conn, round); });
-            console.log(`[${this.ownId}] received peers for round ${round}:`, connections.keySeq().toJS());
+            debug(`[${this.ownId}] received peers for round ${round}: %o`, connections.keySeq().toJS());
             this.connections = connections;
         }
         catch (e) {
-            console.error(e);
+            debug(`[${this.ownId}] while beginning round: %o`, e);
             this.aggregator.setNodes(Set(this.ownId));
             this.connections = Map();
         }
@@ -144,20 +146,20 @@ export class Base extends Client {
     receivePayloads(connections, round) {
         connections.forEach(async (connection, peerId) => {
             let currentCommunicationRounds = 0;
-            console.log(`waiting for peer ${peerId}`);
+            debug(`waiting for peer ${peerId}`);
             do {
                 try {
                     const message = await waitMessageWithTimeout(connection, type.Payload, 60_000, "Timeout waiting for a contribution from peer " + peerId);
                     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}`);
+                        debug(`[${this.ownId}] failed to add contribution from peer ${peerId}`);
                     }
                 }
                 catch (e) {
                     if (this.isDisconnected) {
                         return;
                     }
-                    console.error(e instanceof Error ? e.message : e);
+                    debug(`[${this.ownId}] while receiving payloads: %o`, e);
                 }
             } while (++currentCommunicationRounds < this.aggregator.communicationRounds);
         });
@@ -192,13 +194,13 @@ export class Base extends Client {
                                     payload: encoded
                                 };
                                 peer.send(msg);
-                                console.log(`[${this.ownId}] send weight update to peer`, msg.peer, msg);
+                                debug(`[${this.ownId}] send weight update to peer ${msg.peer}: %O`, msg);
                             }
                         }
                     }));
                 }
-                catch {
-                    throw new Error('error while sending weights');
+                catch (cause) {
+                    throw new Error('error while sending weights', { cause });
                 }
             }
             // Wait for aggregation before proceeding to the next communication round.
@@ -213,7 +215,7 @@ export class Base extends Client {
                 if (this.isDisconnected) {
                     return weights;
                 }
-                console.error(e);
+                debug(`[${this.ownId}] while waiting for aggregation: %o`, e);
                 break;
             }
             // There is at least one communication round remaining

package/dist/client/decentralized/peer_pool.js

@@ -1,6 +1,8 @@
+import createDebug from "debug";
 import { Map } from 'immutable';
 import { Peer } from './peer.js';
 import { PeerConnection } from '../event_connection.js';
+const debug = createDebug("discojs:client:decentralized:pool");
 // TODO cleanup old peers
 export class PeerPool {
     id;
@@ -9,7 +11,7 @@ export class PeerPool {
         this.id = id;
     }
     async shutdown() {
-        console.info(`[${this.id}] is shutting down all its connections`);
+        debug(`[${this.id}] is shutting down all its connections`);
         // Add a timeout o.w. the promise hangs forever if the other peer is already disconnected
         await Promise.race([
             Promise.all(this.peers.valueSeq().map((peer) => peer.disconnect())),
@@ -18,7 +20,7 @@ export class PeerPool {
         this.peers = Map();
     }
     signal(peerId, signal) {
-        console.info(`[${this.id}] signals for`, peerId);
+        debug(`[${this.id}] signals for %s`, peerId);
         const peer = this.peers.get(peerId);
         if (peer === undefined) {
             throw new Error(`received signal for unknown peer: ${peerId}`);
@@ -31,17 +33,17 @@ export class PeerPool {
         if (peersToConnect.contains(this.id)) {
             throw new Error('peers to connect contains our id');
         }
-        console.info(`[${this.id}] is connecting peers:`, peersToConnect.toJS());
+        debug(`[${this.id}] is connecting peers: %o`, peersToConnect.toArray());
         const newPeers = Map(peersToConnect
             .filter((id) => !this.peers.has(id))
             .map((id) => [id, new Peer(id, id < this.id)]));
-        console.info(`[${this.id}] asked to connect new peers:`, newPeers.keySeq().toJS());
+        debug(`[${this.id}] asked to connect new peers: %o`, newPeers.keySeq().toArray());
         const newPeersConnections = newPeers.map((peer) => new PeerConnection(this.id, peer, signallingServer));
         // adding peers to pool before connecting them because they must be set to call signal on them
         this.peers = this.peers.merge(newPeersConnections);
         clientHandle(this.peers);
         await Promise.all(newPeersConnections.valueSeq().map((conn) => conn.connect()));
-        console.info(`[${this.id}] knowns connected peers:`, this.peers.keySeq().toJS());
+        debug(`[${this.id}] knowns connected peers: %o`, this.peers.keySeq().toArray());
         return this.peers
             .filter((_, id) => peersToConnect.has(id));
     }

package/dist/client/event_connection.js

@@ -1,9 +1,11 @@
-import WebSocket from 'isomorphic-ws';
-import msgpack from 'msgpack-lite';
+import createDebug from "debug";
+import WebSocket from "isomorphic-ws";
+import msgpack from "msgpack-lite";
 import * as decentralizedMessages from './decentralized/messages.js';
 import { type } from './messages.js';
 import { timeout } from './utils.js';
 import { EventEmitter } from '../utils/event_emitter.js';
+const debug = createDebug("discojs:client:connections");
 export async function waitMessage(connection, type) {
     return await new Promise((resolve) => {
         // "once" is important because we can't resolve the same promise multiple times
@@ -41,7 +43,9 @@ export class PeerConnection extends EventEmitter {
             }
             this.emit(msg.type, msg);
         });
-        this.peer.on('close', () => { console.warn('From', this._ownId, ': peer', this.peer.id, 'closed connection'); });
+        this.peer.on("close", () => {
+            debug(`[${this._ownId}] peer ${this.peer.id} closed connection`);
+        });
         await new Promise((resolve) => {
             this.peer.on('connect', resolve);
         });

package/dist/client/federated/base.js

@@ -1,8 +1,10 @@
+import createDebug from "debug";
 import { serialization, } from "../../index.js";
 import { Base as Client } from "../base.js";
 import { type } from "../messages.js";
 import { waitMessageWithTimeout, WebSocketServer, } from "../event_connection.js";
 import * as messages from "./messages.js";
+const debug = createDebug("discojs:client:federated");
 /**
  * Client class that communicates with a centralized, federated server, when training
  * a specific task in the federated setting.
@@ -53,7 +55,7 @@ export class Base extends Client {
         };
         this.server.send(msg);
         const received = await waitMessageWithTimeout(this.server, type.AssignNodeID);
-        console.info(`[${received.id}] assign id generated by the server`);
+        debug(`[${received.id}] assign id generated by the server`);
         this._ownId = received.id;
     }
     /**
@@ -87,7 +89,7 @@ export class Base extends Client {
         else {
             // Unexpected case: for some reason, the server result is stale.
             // We proceed to the next round without its result.
-            console.info(`[${this.ownId}] Server result is either stale or not received`);
+            debug(`[${this.ownId}] server result is either stale or not received`);
             this.aggregator.nextRound();
         }
         return await this.aggregationResult;
@@ -122,7 +124,7 @@ export class Base extends Client {
             }
         }
         catch (e) {
-            console.error(e);
+            debug(`[${this.ownId}] while receiving results: %o`, e);
         }
     }
 }

package/dist/dataset/data/preprocessing/text_preprocessing.js

@@ -35,7 +35,7 @@ const leftPadding = {
             // to include one more token than maxSequenceLength in order to have the next token's label of the maxSequenceLength'th token
             const maxLength = task.trainingInformation.maxSequenceLength ?? tokenizer.model_max_length;
             const maxLengthPlusLabel = maxLength + 1;
-            let fixedLengthTokens = tf.tensor(tokens, undefined, 'int32'); // cast tokens from float to int for gpt-tfjs
+            let fixedLengthTokens = tf.tensor1d(tokens, 'int32'); // cast tokens from float to int for gpt-tfjs
             if (fixedLengthTokens.size > maxLengthPlusLabel) { // Should never happen because tokenization truncates inputs
                 throw Error("There are more tokens than expected after tokenization and truncation");
             }
@@ -45,6 +45,7 @@ const leftPadding = {
             }
             // if tokens.size == maxLengthPlusLabel we can leave it as it is
             // ys is a one-hot encoding of the next token (i.e. xs shifted by one)
+            // cast because oneHot isn't size-typing its return value
             const ys = tf.oneHot(fixedLengthTokens.slice([1]), tokenizer.model.vocab.length + 1);
             // remove the extra token now that ys is created
             const xs = fixedLengthTokens.slice([0], maxLength);

package/dist/dataset/data/tabular_data.js

@@ -13,9 +13,8 @@ export class TabularData extends Data {
         try {
             await dataset.iterator();
         }
-        catch (e) {
-            console.error('Data input format is not compatible with the chosen task.');
-            throw (e);
+        catch (cause) {
+            throw new Error('data input format not compatible with chosen task', { cause });
         }
         return new TabularData(dataset, task, size);
     }

package/dist/index.d.ts

@@ -6,12 +6,11 @@ export * as client from './client/index.js';
 export * as aggregator from './aggregator/index.js';
 export { WeightsContainer, aggregation } from './weights/index.js';
 export { Logger, ConsoleLogger } from './logging/index.js';
-export { Memory, type ModelInfo, type Path, type ModelSource, Empty as EmptyMemory } from './memory/index.js';
+export { Memory, type ModelInfo, type ModelSource, Empty as EmptyMemory } from './memory/index.js';
 export { Disco, RoundLogs } from './training/index.js';
 export { Validator } from './validation/index.js';
 export { Model, BatchLogs, EpochLogs, ValidationMetrics } from './models/index.js';
 export * as models from './models/index.js';
 export * from './task/index.js';
 export * as defaultTasks from './default_tasks/index.js';
-export * from './types.js';
 export * as async_iterator from "./utils/async_iterator.js";

package/dist/index.js

@@ -13,5 +13,4 @@ export { Model, EpochLogs } from './models/index.js';
 export * as models from './models/index.js';
 export * from './task/index.js';
 export * as defaultTasks from './default_tasks/index.js';
-export * from './types.js';
 export * as async_iterator from "./utils/async_iterator.js";

package/dist/memory/base.d.ts

@@ -1,8 +1,4 @@
 import type { Model, TaskID } from '../index.js';
-/**
- * Model path which uniquely identifies a model in memory.
- */
-export type Path = string;
 /**
  * Type of models stored in memory. Stored models can either be a model currently
  * being trained ("working model") or a regular model saved in memory ("saved model").
@@ -21,10 +17,10 @@ export interface ModelInfo {
 }
 /**
  * A model source uniquely identifies a model stored in memory.
- * It can be in the form of either a model info object or a Path string
+ * It can be in the form of either a model info object or an ID
  * (one-to-one mapping between the two)
  */
-export type ModelSource = ModelInfo | Path;
+export type ModelSource = ModelInfo | string;
 /**
  * Represents a model memory system, providing functions to fetch, save, delete and update models.
  * Stored models can either be a model currently being trained ("working model") or a regular model
@@ -67,7 +63,7 @@ export declare abstract class Memory {
      * @param source The model source
      * @returns The saved model's path
      */
-    abstract saveWorkingModel(source: ModelSource): Promise<Path | undefined>;
+    abstract saveWorkingModel(source: ModelSource): Promise<string | undefined>;
     /**
      * Saves the newly provided model to the given model source.
      * Returns the saved model's path
@@ -75,7 +71,7 @@ export declare abstract class Memory {
      * @param model The new model
      * @returns The saved model's path
      */
-    abstract saveModel(source: ModelSource, model: Model): Promise<Path | undefined>;
+    abstract saveModel(source: ModelSource, model: Model): Promise<string | undefined>;
     /**
      * Moves the model identified by the model source to a file system. This is platform-dependent.
      * @param source The model source
@@ -95,7 +91,7 @@ export declare abstract class Memory {
      * @param source The model source
      * @returns The model path
      */
-    abstract getModelMemoryPath(source: ModelSource): Path | undefined;
+    abstract getModelMemoryPath(source: ModelSource): string | undefined;
     /**
      * Computes the model information corresponding to the given model source, be it a path or model information.
      * This is used to easily switch between model path and information, which are both unique model identifiers

package/dist/memory/empty.d.ts

@@ -1,5 +1,5 @@
 import type { Model } from '../index.js';
-import type { ModelInfo, Path } from './base.js';
+import type { ModelInfo } from './base.js';
 import { Memory } from './base.js';
 /**
  * Represents an empty model memory.
@@ -14,7 +14,7 @@ export declare class Empty extends Memory {
     saveModel(): Promise<undefined>;
     deleteModel(): Promise<void>;
     downloadModel(): Promise<void>;
-    getModelMemoryPath(): Path;
+    getModelMemoryPath(): string;
     getModelInfo(): ModelInfo;
     duplicateSource(): Promise<undefined>;
 }

package/dist/memory/index.d.ts

@@ -1,2 +1,2 @@
 export { Empty } from './empty.js';
-export { Memory, type ModelInfo, type Path, type ModelSource } from './base.js';
+export { Memory, type ModelInfo, type ModelSource } from './base.js';

package/dist/models/gpt/index.d.ts

@@ -4,7 +4,6 @@
 import * as tf from '@tensorflow/tfjs';
 import { PreTrainedTokenizer } from '@xenova/transformers';
 import { WeightsContainer } from '../../index.js';
-import type { Dataset } from '../../dataset/index.js';
 import { BatchLogs, Model, EpochLogs } from "../index.js";
 import type { Prediction, Sample } from '../model.js';
 import { type GPTConfig } from './config.js';
@@ -25,7 +24,13 @@ export declare class GPT extends Model {
      * @param epochs the number of passes of the training dataset
      * @param tracker
      */
-    train(trainingData: Dataset, validationData?: Dataset): AsyncGenerator<BatchLogs, EpochLogs>;
+    train(trainingData: tf.data.Dataset<{
+        xs: tf.Tensor2D;
+        ys: tf.Tensor3D;
+    }>, validationData?: tf.data.Dataset<{
+        xs: tf.Tensor2D;
+        ys: tf.Tensor3D;
+    }>): AsyncGenerator<BatchLogs, EpochLogs>;
     predict(input: Sample): Promise<Prediction>;
     generate(input: string, tokenizer: PreTrainedTokenizer, newTokens?: number): Promise<string>;
     get config(): Required<GPTConfig>;

package/dist/models/gpt/index.js

@@ -1,13 +1,15 @@
 /**
  * this code is taken from gpt-tfjs with modifications from @peacefulotter and @lukemovement
  **/
+import createDebug from "debug";
+import { List } from 'immutable';
 import * as tf from '@tensorflow/tfjs';
 import { WeightsContainer } from '../../index.js';
 import { Model, EpochLogs } from "../index.js";
 import { GPTForCausalLM } from './model.js';
 import { DEFAULT_CONFIG } from './config.js';
 import evaluate from './evaluate.js';
-import { List } from 'immutable';
+const debug = createDebug("discojs:models:gpt");
 export class GPT extends Model {
     model;
     #maxBatchCount;
@@ -126,8 +128,7 @@ export class GPT extends Model {
             this.model.optimizer.dispose();
         }
         const disposeResults = this.model.dispose();
-        if (disposeResults.refCountAfterDispose > 0) {
-            console.error("The GPT model was not disposed correctly (refcount > 0)", disposeResults);
-        }
+        if (disposeResults.refCountAfterDispose > 0)
+            debug("model not disposed correctly: %o", disposeResults);
     }
 }

package/dist/models/gpt/model.js

@@ -1,8 +1,10 @@
+import createDebug from "debug";
 import * as tf from '@tensorflow/tfjs';
 import { getModelSizes, DEFAULT_CONFIG } from './config.js';
 import { getCustomAdam, clipByGlobalNormObj } from './optimizers.js';
 import evaluate from './evaluate.js';
 import { GPTArchitecture } from './layers.js';
+const debug = createDebug("discojs:models:gpt");
 /**
  * GPTModel extends tf.LayersModel and overrides tfjs' default training loop
  *
@@ -87,10 +89,18 @@ class GPTModel extends tf.LayersModel {
                     this.config.evaluateEvery !== undefined &&
                     iteration % this.config.evaluateEvery == 0) {
                     const iterationLogs = await evaluate(this, evalDataset, this.config.maxEvalBatches);
-                    console.log(iterationLogs);
+                    debug('evaluation metrics: %O', iterationLogs);
                 }
                 const memory = tf.memory().numBytes / 1024 / 1024 / 1024;
-                console.log(`Epoch: ${epoch}`, `\tStep: ${iteration} / ${this.config.maxIter}`, `\tLoss: ${loss.toFixed(3)}`, `\tMemory: ${memory.toFixed(2)} GB`, `\tNumber of tensors allocated: ${tf.memory().numTensors}`, `\tPreprocessing time: ${preprocessingTime.toFixed(0)} ms`, `\tWeight update time: ${weightUpdateTime.toFixed(0)} ms`);
+                debug("training metrics: %O", {
+                    epoch,
+                    iteration,
+                    loss,
+                    memory,
+                    allocated: tf.memory().numTensors,
+                    preprocessingTime,
+                    weightUpdateTime,
+                });
                 iteration++;
                 next = await iterator.next();
             }

package/dist/models/gpt/optimizers.js

@@ -38,7 +38,6 @@ class AdamW extends tf.AdamOptimizer {
     excludeFromWeightDecay;
     gradientClipNorm;
     constructor(params) {
-        console.log('Using custom AdamW optimizer');
         const defaultParams = {
             learningRate: 0.1,
             beta1: 0.9,

package/dist/task/task_handler.js

@@ -1,7 +1,9 @@
 import axios from 'axios';
+import createDebug from "debug";
 import { Map } from 'immutable';
 import { serialization } from '../index.js';
 import { isTask } from './task.js';
+const debug = createDebug("discojs:task:handlers");
 const TASK_ENDPOINT = 'tasks';
 export async function pushTask(url, task, model) {
     await axios.post(url.href + TASK_ENDPOINT, {
@@ -19,7 +21,7 @@ export async function fetchTasks(url) {
     else if (!tasks.every(isTask)) {
         for (const task of tasks) {
             if (!isTask(task)) {
-                console.error("task has invalid format:", task);
+                debug("task has invalid format: :O", task);
             }
         }
         throw new Error('invalid tasks response, the task object received is not well formatted');

package/dist/validation/validator.d.ts

@@ -1,4 +1,4 @@
-import type { data, Model, Task, Logger, client as clients, Memory, ModelSource, Features } from '../index.js';
+import type { data, Model, Task, Logger, client as clients, Memory, ModelSource } from '../index.js';
 export declare class Validator {
     readonly task: Task;
     readonly logger: Logger;
@@ -13,10 +13,10 @@ export declare class Validator {
     test(data: data.Data): AsyncGenerator<Array<{
         groundTruth: number;
         pred: number;
-        features: Features;
+        features: number[];
     }>, void>;
     inference(data: data.Data): AsyncGenerator<Array<{
-        features: Features;
+        features: number[];
         pred: number;
     }>, void>;
     getModel(): Promise<Model>;

package/dist/validation/validator.js

@@ -57,7 +57,7 @@ export class Validator {
             hits += List(pred).zip(List(ysLabel)).filter(([p, y]) => p === y).size;
             this.rollingAccuracy = hits / this.size;
             tf.dispose([xs, ys, yPredTensor]);
-            yield List(ysLabel).zip(List(pred), List(currentFeatures))
+            yield (List(ysLabel).zip(List(pred), List(currentFeatures)))
                 .map(([gt, p, f]) => ({ groundTruth: gt, pred: p, features: f }))
                 .toArray();
             next = await iterator.next();

package/package.json

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

package/dist/types.d.ts

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

package/dist/types.js

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