Let users join federated training mid session

.github/workflows/lint-test-build.yml

@@ -288,7 +288,7 @@ jobs:
  cache: npm
  - run: npm ci
  - run: npm --workspace={discojs,discojs-node} run build
- - run: npm --workspace=server test
+ - run: DEBUG='server:controllers:federated|discojs:client:federated' npm --workspace=server test
 
  test-webapp:
  needs: [build-lib, build-lib-web, download-datasets]

cli/src/cli.ts

@@ -26,7 +26,7 @@ async function runUser(
  const trainingScheme = task.trainingInformation.scheme
  const aggregator = aggregators.getAggregator(task)
  const client = clients.getClient(trainingScheme, url, task, aggregator) 
- const disco = await Disco.fromTask(task, client, { scheme: trainingScheme });
+ const disco = new Disco(task, client, { scheme: trainingScheme });
 
  const logs = List(await arrayFromAsync(disco.trainByRound(data)));
  await new Promise((res, _) => setTimeout(() => res('timeout'), 1000)) // Wait for other peers to finish

discojs-web/src/memory/memory.ts

@@ -7,11 +7,14 @@
  * The working model is loaded from IndexedDB for training (model.fit) only.
  */
 import { Map } from 'immutable'
+import createDebug from "debug"
 import * as tf from '@tensorflow/tfjs'
 
 import type { Model, ModelInfo, ModelSource } from '@epfml/discojs'
 import { Memory, models } from '@epfml/discojs'
 
+const debug = createDebug('discojs-web:memory')
+
 export class IndexedDB extends Memory {
  override getModelMemoryPath (source: ModelSource): string {
  if (typeof source === 'string') {
@@ -100,7 +103,8 @@ export class IndexedDB extends Memory {
  modelInfo['tensorBackend'] = 'gpt'
  includeOptimizer = false // true raises an error
  } else {
- throw new Error('unknown model type')
+ debug('unknown working model type %o', model)
+ throw new Error(`unknown model type while updating working model`)
  }
  const indexedDBURL = this.getModelMemoryPath(modelInfo)
  await model.extract().save(indexedDBURL, { includeOptimizer })
@@ -137,8 +141,9 @@ export class IndexedDB extends Memory {
  } else if (model instanceof models.GPT) { 
  modelInfo['tensorBackend'] = 'gpt'
  includeOptimizer = false // true raises an error
- } else {
- throw new Error('unknown model type')
+ } else {
+ debug('unknown saved model type %o', model)
+ throw new Error('unknown model type while saving model')
  }
  const indexedDBURL = this.getModelMemoryPath(modelInfo)
  await model.extract().save(indexedDBURL, { includeOptimizer })

discojs/src/aggregator/base.ts

@@ -69,14 +69,34 @@ export abstract class Base<T> extends EventEmitter<{'aggregation': T }> {
 
  /**
  * Adds a node's contribution to the aggregator for the given aggregation and communication rounds.
+ * The aggregation round is increased whenever a new global model is obtained and local models are updated.
+ * Within one aggregation round there may be multiple communication rounds (such as for the decentralized secure aggregation
+ * which requires multiple steps to obtain a global model)
  * 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
+ * @param round aggregation round of the contribution was made
+ * @param communicationRound communication round the contribution was made within the aggregation round
+ * @returns boolean, true if the contribution has been successfully taken into account or False if it has been rejected
  */
  abstract add (nodeId: client.NodeID, contribution: T, round: number, communicationRound?: number): boolean
 
+ /**
+ * Evaluates whether a given participant contribution can be used in the current aggregation round
+ * the boolean returned by `this.add` is obtained via `this.isValidContribution`
+ */
+ isValidContribution(nodeId: client.NodeID, round: number): boolean {
+ if (!this.nodes.has(nodeId)) {
+ debug("Contribution rejected because node id is not registered")
+ return false;
+ }
+ if (!this.isWithinRoundCutoff(round)) {
+ debug(`Contribution rejected because round ${round} is not within round cutoff`)
+ return false;
+ }
+ return true
+ }
+
  /**
  * Performs an aggregation step over the received node contributions.
  * Must store the aggregation's result in the aggregator's result promise.
@@ -101,16 +121,16 @@ export abstract class Base<T> extends EventEmitter<{'aggregation': T }> {
  log (step: AggregationStep, from?: client.NodeID): void {
  switch (step) {
  case AggregationStep.ADD:
- debug(`adding contribution from node ${from ?? '"unknown"'} for round (${this.communicationRound}, ${this.round})`);
+ debug(`Adding contribution from node ${from ?? '"unknown"'} for aggregation round ${this.round} and communication round ${this.communicationRound}`);
  break
  case AggregationStep.UPDATE:
  if (from === undefined) {
  return
  }
- debug(`updating contribution from node ${from ?? '"unknown"'} for round (${this.communicationRound}, ${this.round})`)
+ debug(`Updating contribution from node ${from} for aggregation round ${this.round} and communication round ${this.communicationRound}`)
  break
  case AggregationStep.AGGREGATE:
- debug(`buffer full, aggregating weights for round (${this.communicationRound}, ${this.round})`)
+ debug(`Buffer is full. Aggregating weights for round aggregation round ${this.round} and communication round ${this.communicationRound}`)
  break
  default: {
  const _: never = step
@@ -134,6 +154,14 @@ export abstract class Base<T> extends EventEmitter<{'aggregation': T }> {
  return false
  }
 
+ /**
+ * Remove a node's id from the set of active nodes.
+ * @param nodeId The node to be removed
+ */
+ removeNode (nodeId: client.NodeID): void{
+ this._nodes = this._nodes.delete(nodeId)
+ }
+
  /**
  * 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

discojs/src/aggregator/mean.ts

@@ -16,6 +16,7 @@ type ThresholdType = 'relative' | 'absolute'
 export class MeanAggregator extends Aggregator<WeightsContainer> {
  readonly #threshold: number;
  readonly #thresholdType: ThresholdType;
+ #minNbOfParticipants: number | undefined;
 
  /**
  * Create a mean aggregator that averages all weight updates received when a specified threshold is met.
@@ -77,6 +78,11 @@ export class MeanAggregator extends Aggregator<WeightsContainer> {
 
  /** Checks whether the contributions buffer is full. */
  override isFull(): boolean {
+ // Make sure that we are over the minimum number of participants
+ // if specified
+ if (this.#minNbOfParticipants !== undefined &&
+ this.nodes.size < this.#minNbOfParticipants) return false
+
  const thresholdValue =
  this.#thresholdType == 'relative'
  ? this.#threshold * this.nodes.size
@@ -85,6 +91,10 @@ export class MeanAggregator extends Aggregator<WeightsContainer> {
  return (this.contributions.get(0)?.size ?? 0) >= thresholdValue;
  }
 
+ set minNbOfParticipants(minNbOfParticipants: number) {
+ this.#minNbOfParticipants = minNbOfParticipants
+ }
+
  override add(
  nodeId: client.NodeID,
  contribution: WeightsContainer,
@@ -94,11 +104,7 @@ export class MeanAggregator extends Aggregator<WeightsContainer> {
  if (currentContributions !== 0)
  throw new Error("only a single communication round");
 
- if (!this.nodes.has(nodeId) || !this.isWithinRoundCutoff(round)) {
- if (!this.nodes.has(nodeId)) debug(`contribution rejected because node ${nodeId} is not registered`);
- if (!this.isWithinRoundCutoff(round)) debug(`contribution rejected because round ${round} is not within cutoff`);
- return false; 
- }
+ if (!this.isValidContribution(nodeId, round)) return false
 
  this.log(
  this.contributions.hasIn([0, nodeId])
@@ -122,6 +128,7 @@ export class MeanAggregator extends Aggregator<WeightsContainer> {
  this.log(AggregationStep.AGGREGATE);
 
  const result = aggregation.avg(currentContributions.values());
+ this._round = this._round++
  this.emit('aggregation', result);
  }
 

discojs/src/aggregator/secure.ts

@@ -60,8 +60,7 @@ export class SecureAggregator extends Aggregator<WeightsContainer> {
  throw new Error("requires communication round to be 0 or 1");
  }
 
- if (!this.nodes.has(nodeId) || !this.isWithinRoundCutoff(round))
- return false;
+ if (!this.isValidContribution(nodeId, round)) return false
 
  this.log(
  this.contributions.hasIn([communicationRound, nodeId])

discojs/src/client/base.ts → discojs/src/client/client.ts

@@ -1,16 +1,17 @@
 import axios from 'axios'
 
-import type { Model, Task, WeightsContainer } from '../index.js'
+import type { Model, Task, WeightsContainer, RoundStatus } from '../index.js'
 import { serialization } from '../index.js'
 import type { NodeID } from './types.js'
 import type { EventConnection } from './event_connection.js'
 import type { Aggregator } from '../aggregator/index.js'
+import { EventEmitter } from '../utils/event_emitter.js'
 
 /**
  * Main, abstract, class representing a Disco client in a network, which handles
  * communication with other nodes, be it peers or a server.
  */
-export abstract class Base {
+export abstract class Client extends EventEmitter<{'status': RoundStatus}>{
  /**
  * Own ID provided by the network's server.
  */
@@ -25,24 +26,21 @@ export abstract class Base {
  protected aggregationResult?: Promise<WeightsContainer>
 
  constructor (
- /**
- * The network server's URL to connect to.
- */
- public readonly url: URL,
- /**
- * The client's corresponding task.
- */
- public readonly task: Task,
- /**
- * The client's aggregator.
- */
- public readonly aggregator: Aggregator
- ) {}
+ public readonly url: URL, // The network server's URL to connect to
+ public readonly task: Task, // The client's corresponding task
+ public readonly aggregator: Aggregator,
+ ) {
+ super()
+ }
 
  /**
  * Handles the connection process from the client to any sort of network server.
+ * This method is overriden by the federated and decentralized clients
+ * By default, it fetches and returns the server's base model
  */
- async connect (): Promise<void> {}
+ async connect(): Promise<Model> {
+ return this.getLatestModel()
+ }
 
  /**
  * Handles the disconnection process of the client from any sort of network server.
@@ -66,24 +64,15 @@ export abstract class Base {
 
  /**
  * 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: WeightsContainer,
- _round: number,
- ): Promise<void> {}
+ abstract onRoundBeginCommunication(): 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
+ * @param weights The local weight update resulting for the current local training round
  * @returns aggregated weights or the local weights upon error
  */
- abstract onRoundEndCommunication(
- _weights: WeightsContainer,
- _round: number,
- ): Promise<WeightsContainer>;
+ abstract onRoundEndCommunication(weights: WeightsContainer): Promise<WeightsContainer>;
 
  // Number of contributors to a collaborative session
  // If decentralized, it should be the number of peers