diff --git a/server/tests/e2e/decentralized.spec.ts b/server/tests/e2e/decentralized.spec.ts index 2e164a719..a6293902d 100644 --- a/server/tests/e2e/decentralized.spec.ts +++ b/server/tests/e2e/decentralized.spec.ts @@ -1,12 +1,13 @@ import type * as http from "node:http"; import type { + DataFormat, DataType, + Dataset, RoundStatus, Client, Task, TaskProvider, ModelCard, - Network, EpochLogs, } from "@epfml/discojs"; import { @@ -19,106 +20,19 @@ import { WeightsContainer, } from "@epfml/discojs"; import { List } from "immutable"; -import { afterEach, describe, expect, it } from "vitest"; +import { afterEach, beforeAll, beforeEach, describe, expect, it } from "vitest"; import { Server } from "../../src/index.js"; -import { datasets, Queue } from "../utils.js"; +import { datasets } from "../utils.js"; +import { + Participant, + arrayFromAsync, + expectAllWSToBeClose, + expectPeersToAgreeOnModel, + expectWSToBeClose, + recordModelsAtRoundBoundary, +} from "./helpers.js"; import * as tf from "@tensorflow/tfjs-node"; -async function WSIntoList(ws: WeightsContainer): Promise>> { - return List( - (await Promise.all(ws.weights.map(async (w) => await w.data()))).map( - (arr) => List(arr), - ), - ); -} - -async function expectWSToBeClose( - left: WeightsContainer, - right: WeightsContainer, -): Promise { - for (const tensors of (await WSIntoList(left)).zip(await WSIntoList(right))) - for (const [l, r] of tensors[0].zip(tensors[1])) - expect(l).to.be.closeTo(r, 1e-4); -} - -// function from federated.spec.ts -async function arrayFromAsync(iter: AsyncIterable): Promise { - const ret: T[] = []; - for await (const e of iter) { - // TODO trick to allow other Promises to run - // else one client might progress alone without communicating with others - // will be fixed when client orchestrations in the server is correctly done - await new Promise((resolve) => setTimeout(resolve, 10)); - - ret.push(e); - } - return ret; -} - -// function to check if weights across all participants are close to each other -async function expectAllWSToBeClose( - weights: WeightsContainer[], -): Promise { - const reference = weights[0]; - - await Promise.all( - weights.map(async (current) => { - await expectWSToBeClose(reference, current); - }), - ); -} - -/** - * Records the model a peer holds at each round boundary: once - * onRoundEndCommunication has returned and before the next local round trains - * on it. Peers hold the very same model at those points, whereas - * `trainer.model.weights` read afterwards also contains each peer's own local - * training, which is not reproducible across peers. - * - * Pass `keep: "latest"` in the tests measuring tensor memory so that the - * recorded models don't grow with the number of rounds. - */ -function recordModelsAtRoundBoundary( - disco: Disco, - { keep = "all" }: { keep?: "all" | "latest" } = {}, -): { - all: () => readonly WeightsContainer[]; - latest: () => WeightsContainer; - dispose: () => void; -} { - const models: WeightsContainer[] = []; - - disco.on("status", (status) => { - if (status !== "local training") return; - if (keep === "latest") models.splice(0).forEach((m) => m.dispose()); - models.push( - new WeightsContainer( - disco.trainer.model.weights.weights.map((w) => w.clone()), - ), - ); - }); - - return { - all: () => models, - latest: () => { - const model = models.at(-1); - if (model === undefined) - throw new Error("the peer hasn't reached a round boundary yet"); - return model; - }, - dispose: () => models.splice(0).forEach((m) => m.dispose()), - }; -} - -/** The peers should hold the same model at their latest round boundary */ -async function expectPeersToAgreeOnModel( - ...recordings: { latest: () => WeightsContainer }[] -): Promise { - const [first, ...others] = recordings; - for (const other of others) - await expectWSToBeClose(first.latest(), other.latest()); -} - const expectWeightsToEqual = (a: WeightsContainer, b: WeightsContainer) => { expect(a.weights.length).to.equal(b.weights.length); @@ -317,7 +231,7 @@ describe("end-to-end decentralized", { timeout: 50_000 }, () => { }), ); - await expectAllWSToBeClose(results.map(([weights]) => weights)); + await expectAllWSToBeClose(...results.map(([weights]) => weights)); } finally { await Promise.all(discos.map((disco) => disco.close())); } @@ -376,7 +290,7 @@ describe("end-to-end decentralized", { timeout: 50_000 }, () => { }), ); - await expectAllWSToBeClose(results.map(([weights]) => weights)); + await expectAllWSToBeClose(...results.map(([weights]) => weights)); } finally { await Promise.all(discos.map((disco) => disco.close())); } @@ -429,7 +343,7 @@ describe("end-to-end decentralized", { timeout: 50_000 }, () => { }), ); - await expectAllWSToBeClose(results.map(([weights]) => weights)); + await expectAllWSToBeClose(...results.map(([weights]) => weights)); } finally { await Promise.all(discos.map((disco) => disco.close())); } @@ -467,269 +381,293 @@ describe("end-to-end decentralized", { timeout: 50_000 }, () => { }; } - // syncs model after participants drop below minNbOfParticipants and newcomers join with mean aggregator - it("emit expected events", { timeout: 150_000 }, async () => { - const { task, taskProvider } = await lusCovidDecentralized(); - const url = await startServer(defaultModels.LUSClassifier, taskProvider); - const dataset = await datasets.loadLusCOVID(); + /** + * At each round (each call to `disco.trainByRound().next()`) the event cycle is: + * a) During onRoundBeginCommunication, + * 1. a peer that joined mid-training syncs its model with the latest one + * 2. the peer notifies the server that they want to join the next round + * 3. the peer waits until there are enough participants, setting the status + * to "not enough participants" while it does + * 4. finishes by updating the status to "local training" + * b) local training (the status remains "local training") + * c) During onRoundEndCommunication + * 1. the peer sets its status to "waiting for peers to share weights" + * and notifies the server that they are ready to share weights + * 2. wait for the server to answer with the current round's peers list + * this is where the nb of participants is updated + * 3. set status to "connecting to peers" and establish the connections + * 4. set status to "updating model" and exchange weight updates + * + * A single call to next() performs a full round: a), b) and c). It only + * resolves once the peers exchanged their weight updates, so a call made + * while the round can't complete stays pending. `Participant` therefore + * splits a round in `startRound()` and `completeRound()`, so that the tests + * can choreograph through the status and participants events instead of + * awaiting a round right away. + * Note that RoundLogs.participants is the count seen at the end of local + * training, before the weight exchange. + * + * Every step of that choreography is a named function below asserting the + * events it expects. Each test then plays the steps leading to the one it + * covers and ends with that step, which keeps a failure pointing at a single + * step of the timeline. Put end to end, the steps make up the timeline: + * - User 1 joins the task + * - User 2 joins + * - User 2 leaves (Since minNbOfParticipants condition is not satisfied, the training stops) + * - User 3 joins (User 3 gets the latest model from User 1 and start local training from that model) + */ + describe("emit expected events", { timeout: 150_000 }, () => { + type Peer = Participant<"image", "decentralized">; + + /** Statuses a peer goes through in c), once it is done training */ + const SHARES_WEIGHTS: readonly RoundStatus[] = [ + "waiting for peers to share weights", + "connecting to peers", + "updating model", + ]; + /** Statuses a peer goes through during a round it can complete */ + const FULL_ROUND: readonly RoundStatus[] = [ + "local training", + ...SHARES_WEIGHTS, + ]; + + let task: Task<"image", "decentralized">; + let url: URL; + let dataset: Dataset; + const joined: Peer[] = []; + + beforeAll(async () => { + dataset = await datasets.loadLusCOVID(); + }); + + beforeEach(async () => { + const decentralized = await lusCovidDecentralized(); + task = decentralized.task; + url = await startServer( + defaultModels.LUSClassifier, + decentralized.taskProvider, + ); + }); + + afterEach(async () => { + await Promise.all(joined.splice(0).map(async (p) => await p.leave())); + }); + + /** Have a new peer join the task, closed at the end of the test */ + function join(name: string): Peer { + const peer = new Participant(name, task, url, dataset); + joined.push(peer); + return peer; + } /** - * At each round (each call to `disco.trainByRound().next()`) the event cycle is: - * a) During onRoundBeginCommunication, - * 1. a peer that joined mid-training syncs its model with the latest one - * 2. the peer notifies the server that they want to join the next round - * 3. the peer waits until there are enough participants, setting the status - * to "not enough participants" while it does - * 4. finishes by updating the status to "local training" - * b) local training (the status remains "local training") - * c) During onRoundEndCommunication - * 1. the peer sets its status to "waiting for peers to share weights" - * and notifies the server that they are ready to share weights - * 2. wait for the server to answer with the current round's peers list - * this is where the nb of participants is updated - * 3. set status to "connecting to peers" and establish the connections - * 4. set status to "updating model" and exchange weight updates - * - * A single call to next() performs a full round: a), b) and c). It only - * resolves once the peers exchanged their weight updates, so a call made - * while the round can't complete stays pending. The test therefore holds - * the pending next() promises and choreographs through the status and - * participants events instead of awaiting next() right away. - * Note that RoundLogs.participants is the count seen at the end of local - * training, before the weight exchange. - * - * Test timeline looks like this: - * - User 1 joins the task - * - User 2 joins - * - User 2 leaves (Since minNbOfParticipants condition is not satisfied, the training stops) - * - User 3 joins (User 3 gets the latest model from User 1 and start local training from that model) - * - User 1 & 3 leave + * A peer joining a task nobody else is on: it waits in a) and doesn't even + * start training, so its round can't complete. */ + async function joinsAlone(name: string): Promise { + const peer = join(name).startRound(); - const discoUser1 = new Disco(task, url, { preprocessOnce: true }); - const discoUser2 = new Disco(task, url, { preprocessOnce: true }); - - // Register listeners for user1 and user2 events - const statusUser1 = new Queue(); - const nbParticipantsUser1 = new Queue(); - const statusUser2 = new Queue(); - const nbParticipantsUser2 = new Queue(); - discoUser1.on("status", (status) => statusUser1.put(status)); - discoUser1.on("participants", (participants) => - nbParticipantsUser1.put(participants), - ); - discoUser2.on("status", (status) => statusUser2.put(status)); - discoUser2.on("participants", (participants) => - nbParticipantsUser2.put(participants), - ); + await peer.expectStatuses("not enough participants"); + await peer.expectParticipants(1); - const modelsUser1 = recordModelsAtRoundBoundary(discoUser1); - const modelsUser2 = recordModelsAtRoundBoundary(discoUser2); - - let user2Closed = false; - - const generatorUser1 = discoUser1.trainByRound(dataset); - const generatorUser2 = discoUser2.trainByRound(dataset); - - /* ROUND 1 */ - /* USER 1 JOINS */ - - // User 1 is alone so it stays in a): it doesn't train yet and the round - // can't complete, so the promise stays pending - const round1User1Promise = generatorUser1.next(); - expect(await statusUser1.next()).equal("not enough participants"); - // We expect only one participant - expect(await nbParticipantsUser1.next()).equal(1); - - /* USER 2 JOINS */ - /* minNbOfParticipants condition satisfied, local training starts */ - - const round1User2Promise = generatorUser2.next(); - // User 2 has enough participants right away and goes to b) - expect(await statusUser2.next()).equal("local training"); - // User 1 is released from a) and trains too - expect(await nbParticipantsUser1.next()).equal(2); - expect(await statusUser1.next()).equal("local training"); - expect(await nbParticipantsUser2.next()).equal(2); - - /* ROUND 1 COMPLETES */ - - // Both peers reach c), the server answers with the round's peers list, - // they exchange their updates and both pending next() calls resolve - const [round1User1, round1User2] = await Promise.all([ - round1User1Promise, - round1User2Promise, - ]); - expect(round1User1.done).to.be.false; - expect(round1User2.done).to.be.false; - if (round1User1.done || round1User2.done) - throw new Error("User 1 or 2 finished training at the 1st round"); - expect(round1User1.value.participants).equal(2); - expect(round1User2.value.participants).equal(2); - - expect(await statusUser1.next()).equal( - "waiting for peers to share weights", - ); - expect(await statusUser1.next()).equal("connecting to peers"); - expect(await statusUser1.next()).equal("updating model"); - expect(await statusUser2.next()).equal( - "waiting for peers to share weights", - ); - expect(await statusUser2.next()).equal("connecting to peers"); - expect(await statusUser2.next()).equal("updating model"); - // Receiving the peers list updates the participants - expect(await nbParticipantsUser1.next()).equal(2); - expect(await nbParticipantsUser2.next()).equal(2); - - /* ROUND 2 */ - - // Both users are present so the round runs a), b) and c) to completion - const [round2User1, round2User2] = await Promise.all([ - generatorUser1.next(), - generatorUser2.next(), - ]); - expect(round2User1.done).to.be.false; - expect(round2User2.done).to.be.false; - if (round2User1.done || round2User2.done) - throw new Error("User 1 or 2 finished training at the 2nd round"); - expect(round2User1.value.participants).equal(2); - expect(round2User2.value.participants).equal(2); - - // Both users did a), b) and c) - expect(await statusUser1.next()).equal("local training"); - expect(await statusUser1.next()).equal( - "waiting for peers to share weights", - ); - expect(await statusUser1.next()).equal("connecting to peers"); - expect(await statusUser1.next()).equal("updating model"); - expect(await statusUser2.next()).equal("local training"); - expect(await statusUser2.next()).equal( - "waiting for peers to share weights", - ); - expect(await statusUser2.next()).equal("connecting to peers"); - expect(await statusUser2.next()).equal("updating model"); - expect(await nbParticipantsUser1.next()).equal(2); - expect(await nbParticipantsUser2.next()).equal(2); - - // Weights should have converged after exchanging updates - await expectPeersToAgreeOnModel(modelsUser1, modelsUser2); - - /* USER 2 LEAVES */ - - // Round 3 starts for User 1 before closing User 2, so User 1 trains and - // then enters c) where it emits "waiting for peers to share weights". - // It cannot reach "connecting to peers" yet: that only happens once the - // server answers with the round's peer list. - const round3User1Promise = generatorUser1.next(); - expect(await statusUser1.next()).equal("local training"); - expect(await statusUser1.next()).equal( - "waiting for peers to share weights", - ); + return peer; + } - await discoUser2.close(); - user2Closed = true; + /** A peer joining a waiting one, releasing it from a) so that both train */ + async function joinsWaitingPeer( + waiting: Peer, + name: string, + ): Promise { + const peer = join(name).startRound(); + + // the newcomer has enough participants right away and goes to b) + await peer.expectStatuses("local training"); + // the waiting peer is released from a) and trains too + await waiting.expectParticipants(2); + await waiting.expectStatuses("local training"); + await peer.expectParticipants(2); + + return peer; + } - // Check if User 1 got a signal that there is not enough participants - expect(await nbParticipantsUser1.next()).equal(1); - expect(await statusUser1.next()).equal("not enough participants"); + /** + * Complete the round every peer started, expecting each of them to go + * through the given statuses and to have seen every other peer. + */ + async function completesRound( + peers: readonly Peer[], + statuses: readonly RoundStatus[], + ): Promise { + const logs = await Promise.all( + peers.map(async (peer) => await peer.completeRound()), + ); + for (const log of logs) expect(log.participants).equal(peers.length); - /* USER 3 JOINS */ + for (const peer of peers) await peer.expectStatuses(...statuses); + // receiving the round's peers list updates the participants + for (const peer of peers) await peer.expectParticipants(peers.length); + } - // Create User 3 and register event listeners - const discoUser3 = new Disco(task, url, { preprocessOnce: true }); - const statusUser3 = new Queue(); - const nbParticipantsUser3 = new Queue(); - discoUser3.on("status", (status) => statusUser3.put(status)); - discoUser3.on("participants", (participants) => - nbParticipantsUser3.put(participants), - ); + /** Peers done with their local training exchange their weight updates */ + const exchangesWeights = (...peers: readonly Peer[]) => + completesRound(peers, SHARES_WEIGHTS); - const waitForUser3ModelSynced = new Promise((resolve) => { - discoUser3.on("modelSynced", (weights) => { - if (weights !== undefined) resolve(weights); - }); + /** A round during which every peer is present, so a), b) and c) run */ + const runsFullRound = (...peers: readonly Peer[]) => + completesRound(peers, FULL_ROUND); + + /** A peer starting a round which it can't complete on its own */ + async function startsRoundAndWaitsForWeights(peer: Peer): Promise { + peer.startRound(); + + // the peer trains and then enters c). It cannot reach "connecting to + // peers" yet: that only happens once the server answers with the round's + // peers list + await peer.expectStatuses( + "local training", + "waiting for peers to share weights", + ); + } + + /** A peer leaving, dropping the remaining one below minNbOfParticipants */ + async function leavesTask(leaving: Peer, remaining: Peer): Promise { + await leaving.leave(); + + await remaining.expectParticipants(1); + await remaining.expectStatuses("not enough participants"); + } + + /** A newcomer joining a peer which is stuck in c) waiting for weights */ + async function joinsPeerWaitingForWeights( + waiting: Peer, + name: string, + ): Promise { + const peer = join(name).startRound(); + + await peer.expectParticipants(2); + // the newcomer syncs its model in a) then trains + await peer.expectStatuses("local training"); + // the waiting peer learns the newcomer joined and is still in c) waiting + // for it to be ready, so it rolls back to the status it had before + // waiting for more participants + await waiting.expectParticipants(2); + await waiting.expectStatuses("waiting for peers to share weights"); + + return peer; + } + + /** The waiting peer and the newcomer exchange their weight updates */ + async function exchangesWeightsWithNewcomer( + waiting: Peer, + newcomer: Peer, + ): Promise { + const logs = await Promise.all([ + waiting.completeRound(), + newcomer.completeRound(), + ]); + // both trained while two participants were around: the waiting peer + // before the other one left, the newcomer after it joined + for (const log of logs) expect(log.participants).equal(2); + + // the waiting peer already announced it was waiting for weights + await waiting.expectStatuses("connecting to peers", "updating model"); + await waiting.expectParticipants(2); + await newcomer.expectStatuses(...SHARES_WEIGHTS); + await newcomer.expectParticipants(2); + } + + it("a peer joining alone waits for a second participant", async () => { + await joinsAlone("user 1"); }); - const modelsUser3 = recordModelsAtRoundBoundary(discoUser3); - const generatorUser3 = discoUser3.trainByRound(dataset); - - /* ROUND 3 COMPLETES */ - /* User 3's first round completes User 1's round 3 */ - - const round1User3Promise = generatorUser3.next(); - expect(await nbParticipantsUser3.next()).equal(2); - // User 3 syncs its model in a) then trains - expect(await statusUser3.next()).equal("local training"); - // User 1 learns User 3 joined and is still in c) waiting for User 3 to be - // ready, so it rolls back to the status it had before waiting for more - // participants - expect(await nbParticipantsUser1.next()).equal(2); - expect(await statusUser1.next()).equal( - "waiting for peers to share weights", - ); + it("a joining peer lets both of them train", async () => { + const user1 = await joinsAlone("user 1"); - // User 3's model should have been synced to the latest global model, i.e. - // the result of User 1 and User 2's last aggregation. User 1 hasn't - // started a new round, so that is still its latest round boundary. - const user3SyncedWeights = await waitForUser3ModelSynced; - await expectWSToBeClose(user3SyncedWeights, modelsUser1.latest()); - - // User 1 and User 3 exchange their updates and both rounds resolve - const [round3User1, round1User3] = await Promise.all([ - round3User1Promise, - round1User3Promise, - ]); - expect(round3User1.done).to.be.false; - expect(round1User3.done).to.be.false; - if (round3User1.done || round1User3.done) - throw new Error("User 1 or 3 finished training at the 3rd round"); - // User 1 trained while User 2 was still there - expect(round3User1.value.participants).equal(2); - expect(round1User3.value.participants).equal(2); - - expect(await statusUser1.next()).equal("connecting to peers"); - expect(await statusUser1.next()).equal("updating model"); - expect(await nbParticipantsUser1.next()).equal(2); - expect(await statusUser3.next()).equal( - "waiting for peers to share weights", - ); - expect(await statusUser3.next()).equal("connecting to peers"); - expect(await statusUser3.next()).equal("updating model"); - expect(await nbParticipantsUser3.next()).equal(2); - - /* ROUND 4 */ - /* first full round shared by User 1 and User 3 */ - - const [round4User1, round2User3] = await Promise.all([ - generatorUser1.next(), - generatorUser3.next(), - ]); - expect(round4User1.done).to.be.false; - expect(round2User3.done).to.be.false; - if (round4User1.done || round2User3.done) - throw new Error("User 1 or 3 finished training at the 4th round"); - expect(round4User1.value.participants).equal(2); - expect(round2User3.value.participants).equal(2); - - expect(await statusUser1.next()).equal("local training"); - expect(await statusUser1.next()).equal( - "waiting for peers to share weights", - ); - expect(await statusUser1.next()).equal("connecting to peers"); - expect(await statusUser1.next()).equal("updating model"); - expect(await statusUser3.next()).equal("local training"); - expect(await statusUser3.next()).equal( - "waiting for peers to share weights", - ); - expect(await statusUser3.next()).equal("connecting to peers"); - expect(await statusUser3.next()).equal("updating model"); + await joinsWaitingPeer(user1, "user 2"); + }); - // Weights should have converged between User 1 and User 3 after the exchange - await expectPeersToAgreeOnModel(modelsUser1, modelsUser3); + it("peers exchange their weights at the end of the round", async () => { + const user1 = await joinsAlone("user 1"); + const user2 = await joinsWaitingPeer(user1, "user 2"); - await discoUser3.close(); - await discoUser1.close().catch(() => {}); - if (!user2Closed) await discoUser2.close().catch(() => {}); + await exchangesWeights(user1, user2); + }); + + // regression test, peer used to display missing participants when + // it was not the case + it("doesn't report missing participants while a peer trains", async () => { + const user1 = await joinsAlone("user 1"); + await joinsWaitingPeer(user1, "user 2"); + + // user 1 is done training first and has to wait for user 2 to be ready, + // but shouldn't be told that participants are missing: user 2 is here, + // only still training. The status following "local training" must go + // straight to the weight exchange. + await user1.expectStatuses("waiting for peers to share weights"); + }); + + it("a round runs the whole cycle when both peers are present", async () => { + const user1 = await joinsAlone("user 1"); + const user2 = await joinsWaitingPeer(user1, "user 2"); + await exchangesWeights(user1, user2); + + await runsFullRound(user1, user2); + + // weights should have converged after exchanging updates + await expectPeersToAgreeOnModel( + user1.modelsAtRoundBoundary, + user2.modelsAtRoundBoundary, + ); + }); + + it("a peer waiting for weights is told when the other leaves", async () => { + const user1 = await joinsAlone("user 1"); + const user2 = await joinsWaitingPeer(user1, "user 2"); + await exchangesWeights(user1, user2); + await startsRoundAndWaitsForWeights(user1); + + await leavesTask(user2, user1); + }); + + it("a peer joining mid-training syncs the latest model", async () => { + const user1 = await joinsAlone("user 1"); + const user2 = await joinsWaitingPeer(user1, "user 2"); + await exchangesWeights(user1, user2); + await startsRoundAndWaitsForWeights(user1); + await leavesTask(user2, user1); + + const user3 = await joinsPeerWaitingForWeights(user1, "user 3"); + + // user 3's model should have been synced to the latest global model, + // i.e. the result of user 1 and user 2's last aggregation. User 1 hasn't + // completed a new round, so that is still its latest round boundary. + await expectWSToBeClose( + await user3.syncedModel(), + user1.modelsAtRoundBoundary.latest(), + ); + + // the server should accept user 3's weights (they should not be + // outdated) and let both peers complete their round + await exchangesWeightsWithNewcomer(user1, user3); + }); + + it("a peer which joined mid-training then runs full rounds", async () => { + const user1 = await joinsAlone("user 1"); + const user2 = await joinsWaitingPeer(user1, "user 2"); + await exchangesWeights(user1, user2); + await startsRoundAndWaitsForWeights(user1); + await leavesTask(user2, user1); + const user3 = await joinsPeerWaitingForWeights(user1, "user 3"); + await exchangesWeightsWithNewcomer(user1, user3); + + await runsFullRound(user1, user3); + + // weights should have converged between user 1 and user 3 too + await expectPeersToAgreeOnModel( + user1.modelsAtRoundBoundary, + user3.modelsAtRoundBoundary, + ); + }); }); /** @@ -1154,82 +1092,4 @@ describe("end-to-end decentralized", { timeout: 50_000 }, () => { } }, ); - - // regression test, peer used to display missing participants when - // it was not the case - it( - "doesn't report missing participants when peer is sharing its weights", - { timeout: 100_000 }, - async () => { - const { task, taskProvider } = await lusCovidDecentralized(); - const url = await startServer(defaultModels.LUSClassifier, taskProvider); - const dataset = await datasets.loadLusCOVID(); - - /** - * A call to trainByRound().next() runs a full round: a), b) and c), and - * the wait for more participants happens in a). The timeline is: - * - User 1 joins the task by themselves and waits in a) for a second - * participant ("not enough participants" is expected there, User 1 - * really is alone) - * - User 2 joins, both train locally - * - User 1 is done training and waits in c) for User 2 to share its - * weights - * - * User 1 has to wait for User 2 to be ready but, once User 2 joined, - * shouldn't be told that participants are missing: User 2 is here, only - * still training. The statuses following "local training" must go - * straight to the weight exchange. - */ - - /* USER 1 JOINS */ - - const discoUser1 = new Disco(task, url, { preprocessOnce: true }); - const statusUser1 = new Queue(); - discoUser1.on("status", (status) => { - statusUser1.put(status); - }); - const generatorUser1 = discoUser1.trainByRound(dataset); - - // a) blocks until a second participant joins and the round can only - // complete afterwards, so don't await it yet - const logUser1Round1 = generatorUser1.next(); - expect(await statusUser1.next()).equal("not enough participants"); - - /* USER 2 JOINS, BOTH CAN TRAIN */ - - const discoUser2 = new Disco(task, url, { preprocessOnce: true }); - const statusUser2 = new Queue(); - discoUser2.on("status", (status) => { - statusUser2.put(status); - }); - const generatorUser2 = discoUser2.trainByRound(dataset); - const logUser2Round1 = generatorUser2.next(); - - // there are enough participants now, User 1 trains locally - expect(await statusUser1.next()).equal("local training"); - expect(await statusUser2.next()).equal("local training"); - - /* USER 1 IS DONE TRAINING, USER 2 HASN'T SHARED ITS WEIGHTS YET */ - - // User 1 waits in c) for User 2 to be ready but should NOT report - // missing participants: the next statuses must be the weight exchange - expect(await statusUser1.next()).equal( - "waiting for peers to share weights", - ); - - /* USER 2 IS DONE TRAINING TOO */ - - await Promise.all([logUser1Round1, logUser2Round1]); - expect(await statusUser1.next()).equal("connecting to peers"); - expect(await statusUser1.next()).equal("updating model"); - expect(await statusUser2.next()).equal( - "waiting for peers to share weights", - ); - expect(await statusUser2.next()).equal("connecting to peers"); - expect(await statusUser2.next()).equal("updating model"); - - await discoUser1.close(); - await discoUser2.close(); - }, - ); }); diff --git a/server/tests/e2e/federated.spec.ts b/server/tests/e2e/federated.spec.ts index 5953c433e..6c3c7c6e9 100644 --- a/server/tests/e2e/federated.spec.ts +++ b/server/tests/e2e/federated.spec.ts @@ -12,42 +12,20 @@ import type { } from "@epfml/discojs"; import { Disco, defaultTasks, defaultModels, GPT } from "@epfml/discojs"; import { List } from "immutable"; -import { assert, afterEach, describe, expect, it } from "vitest"; +import { + assert, + afterEach, + beforeAll, + beforeEach, + describe, + expect, + it, +} from "vitest"; import { Server } from "../../src/index.js"; -import { Queue, datasets } from "../utils.js"; +import { datasets } from "../utils.js"; +import { Participant, arrayFromAsync, expectWSToBeClose } from "./helpers.js"; import * as tf from "@tensorflow/tfjs-node"; -// Array.fromAsync not yet widely used (2024) -async function arrayFromAsync(iter: AsyncIterable): Promise { - const ret: T[] = []; - for await (const e of iter) { - // TODO trick to allow other Promises to run - // else one client might progress alone without communicating with others - // will be fixed when client orchestrations in the server is correctly done - await new Promise((resolve) => setTimeout(resolve, 10)); - - ret.push(e); - } - return ret; -} - -async function WSIntoList(ws: WeightsContainer): Promise>> { - return List( - (await Promise.all(ws.weights.map(async (w) => await w.data()))).map( - (arr) => List(arr), - ), - ); -} - -async function expectWSToBeClose( - left: WeightsContainer, - right: WeightsContainer, -): Promise { - for (const tensors of (await WSIntoList(left)).zip(await WSIntoList(right))) - for (const [l, r] of tensors[0].zip(tensors[1])) - expect(l).to.be.closeTo(r, 1e-4); -} - describe("end-to-end federated", () => { let handle: http.Server | undefined; async function startServer( @@ -227,150 +205,193 @@ describe("end-to-end federated", () => { assert.isTrue(r1[0].equals(r2[0])); }); - it("clients emit expected events", { timeout: 100_000 }, async () => { - const task = await defaultTasks.lusCovid.getTask(); - task.trainingInformation = { - ...task.trainingInformation, - roundDuration: 1, - minNbOfParticipants: 2, - }; - const taskProvider = { - ...defaultTasks.lusCovid, - getTask: () => Promise.resolve(task), - }; - const url = await startServer(defaultModels.LUSClassifier, taskProvider); - const dataset = await datasets.loadLusCOVID(); + /** + * When disco.trainByRound is called for the first time, the client connects + * to the server which returns the latest model, current round and nb of + * participants. Then at each round the event cycle is: + * a) onRoundBeginCommunication which updates the status to "local training" + * b) local training (the status remains "local training") + * c) onRoundEndCommunication which sends the local update and + * receives the global weights while emitting the status UPDATE + * + * Given this, it is important to note that a single call to + * disco.trainByRound().next() performs a full round: a), b) and c). + * It only resolves once the server aggregated the round, so when a client + * is alone (minNbOfParticipants isn't met) the call stays pending until + * another participant joins and the round completes. `Participant` therefore + * splits a round in `startRound()` and `completeRound()`, so that the tests + * can choreograph through the status and participants events instead of + * awaiting a round right away. + * + * Every step of that choreography is a named function below asserting the + * events it expects. Each test then plays the steps leading to the one it + * covers and ends with that step, which keeps a failure pointing at a single + * step of the timeline. + */ + describe("clients emit expected events", { timeout: 100_000 }, () => { + type Client = Participant<"image", "federated">; + + /** Statuses a client goes through during a round it can complete */ + const FULL_ROUND: readonly RoundStatus[] = [ + "local training", + "updating model", + ]; + + let task: Task<"image", "federated">; + let url: URL; + let dataset: Dataset; + const joined: Client[] = []; + + beforeAll(async () => { + dataset = await datasets.loadLusCOVID(); + }); + + beforeEach(async () => { + const baseTask = await defaultTasks.lusCovid.getTask(); + task = { + ...baseTask, + trainingInformation: { + ...baseTask.trainingInformation, + roundDuration: 1, + minNbOfParticipants: 2, + }, + }; + + url = await startServer(defaultModels.LUSClassifier, { + ...defaultTasks.lusCovid, + getTask: () => Promise.resolve(task), + }); + }); + + afterEach(async () => { + await Promise.all(joined.splice(0).map(async (c) => await c.leave())); + }); + + /** Have a new client join the task, closed at the end of the test */ + function join(name: string): Client { + const client = new Participant(name, task, url, dataset); + joined.push(client); + return client; + } /** - * When disco.trainByRound is called for the first time, the client connects to the server - * which returns the latest model, current round and nb of participants. - * Then at each round the event cycle is: - * a) onRoundBeingCommunication which updates the status to "local training" - * b) local training (the status remains "local training") - * c) onRoundEndCommunication which sends the local update and - * receives the global weights while emitting the status UPDATE - * - * Given this, it is important to note that a single call to - * disco.trainByRound().next() performs a full round: a), b) and c). - * It only resolves once the server aggregated the round, so when a client - * is alone (minNbOfParticipants isn't met) the call stays pending until - * another participant joins and the round completes. Tests therefore hold - * the pending next() promise and choreograph through the status and - * participants events instead of awaiting next() right away. - * - * In this test the timeline is: - * - User 1 joins the task by themselves - * - User 2 joins - * - User 1 leaves - * - User 3 joins - * - User 2 & 3 leave + * A client joining a task nobody else is on: it trains locally right away + * but can't share its update, so it stays pending in c). */ + async function joinsAlone(name: string): Promise { + const client = join(name).startRound(); - // Create User 1 - const discoUser1 = new Disco(task, url, { preprocessOnce: true }); - const statusUser1 = new Queue(); - const nbParticipantsUser1 = new Queue(); - discoUser1.on("status", (status) => statusUser1.put(status)); - discoUser1.on("participants", (participants) => - nbParticipantsUser1.put(participants), - ); - const generatorUser1 = discoUser1.trainByRound(dataset); - - // Have User 1 join the task and train locally. The round can't complete - // while User 1 is alone so the promise stays pending in c) - const logUser1Round1Promise = generatorUser1.next(); - expect(await statusUser1.next()).equal("local training"); - expect(await nbParticipantsUser1.next()).equal(1); - expect(await statusUser1.next()).equal("not enough participants"); - - // Create User 2 - const discoUser2 = new Disco(task, url, { preprocessOnce: true }); - const statusUser2 = new Queue(); - const nbParticipantsUser2 = new Queue(); - discoUser2.on("status", (status) => statusUser2.put(status)); - discoUser2.on("participants", (participants) => - nbParticipantsUser2.put(participants), - ); - const generatorUser2 = discoUser2.trainByRound(dataset); - - // Have User 2 join the task and train for one round - const logUser2Round1Promise = generatorUser2.next(); - // User 2 connects to the server which triggers the participant event - expect(await nbParticipantsUser2.next()).equal(2); - expect(await statusUser2.next()).equal("local training"); - // User 1 receives the EnoughParticipants message with the participants, - // its previous status is restored and it proceeds to share its update - expect(await nbParticipantsUser1.next()).equal(2); - expect(await statusUser1.next()).equal("local training"); - expect(await statusUser1.next()).equal("updating model"); - // User 2 finishes training and shares its update too - expect(await statusUser2.next()).equal("updating model"); - - // The server aggregates the round and answers with the new global weights - // along with the participants, resolving both pending next() calls - await Promise.all([logUser1Round1Promise, logUser2Round1Promise]); - expect(await nbParticipantsUser1.next()).equal(2); - expect(await nbParticipantsUser2.next()).equal(2); - - // Proceed with round 2, both users are present so the round completes - await Promise.all([generatorUser1.next(), generatorUser2.next()]); - // User 1 and 2 did a), b) and c) - expect(await statusUser1.next()).equal("local training"); - expect(await statusUser1.next()).equal("updating model"); - expect(await statusUser2.next()).equal("local training"); - expect(await statusUser2.next()).equal("updating model"); - // Receive the server payload during c) along with the participants - expect(await nbParticipantsUser1.next()).equal(2); - expect(await nbParticipantsUser2.next()).equal(2); - - // Have user 1 quit the session - await discoUser1.close(); - // User 2 receives the WaitingForMoreParticipants message - expect(await statusUser2.next()).equal("not enough participants"); - expect(await nbParticipantsUser2.next()).equal(1); - - // Make User 2 start round 3, it trains and then waits in c) for - // another participant - const logUser2Round3Promise = generatorUser2.next(); - expect(await statusUser2.next()).equal("local training"); - expect(await statusUser2.next()).equal("not enough participants"); - - // Create User 3 - const discoUser3 = new Disco(task, url, { preprocessOnce: true }); - const statusUser3 = new Queue(); - const nbParticipantsUser3 = new Queue(); - discoUser3.on("status", (status) => statusUser3.put(status)); - discoUser3.on("participants", (participants) => - nbParticipantsUser3.put(participants), - ); - const generatorUser3 = discoUser3.trainByRound(dataset); - - // User 3 joins mid-training and trains one local round - const logUser3Round1Promise = generatorUser3.next(); - expect(await nbParticipantsUser3.next()).equal(2); - expect(await statusUser3.next()).equal("local training"); - - // User 2 receives the EnoughParticipants message, its previous status - // is restored and it proceeds to share its update - expect(await nbParticipantsUser2.next()).equal(2); - expect(await statusUser2.next()).equal("local training"); - expect(await statusUser2.next()).equal("updating model"); - // User 3 finishes training and sends their weights to the server - expect(await statusUser3.next()).equal("updating model"); - - // the server should accept user 3's weights (should not be outdated) - // and aggregate the global weights, resolving both rounds - await Promise.all([logUser2Round3Promise, logUser3Round1Promise]); - // User 2 and 3 finish c) - expect(await nbParticipantsUser2.next()).equal(2); - expect(await nbParticipantsUser3.next()).equal(2); - - await discoUser2.close(); - expect(await statusUser3.next()).equal("not enough participants"); - // WaitForMoreParticipants message - expect(await nbParticipantsUser3.next()).equal(1); - - await discoUser3.close(); + // a) and b), the client trains without waiting for anyone + await client.expectStatuses("local training"); + await client.expectParticipants(1); + // c), sharing the update needs a second participant + await client.expectStatuses("not enough participants"); + + return client; + } + + /** + * A client joining a waiting one: both share their update, the server + * aggregates them and answers with the new global weights. + */ + async function joinsWaitingClient( + waiting: Client, + name: string, + ): Promise { + const client = join(name).startRound(); + + // the new client connects to the server, which triggers the participant + // event, and trains + await client.expectParticipants(2); + await client.expectStatuses("local training"); + // the waiting client receives the EnoughParticipants message with the + // participants, its previous status is restored and it shares its update + await waiting.expectParticipants(2); + await waiting.expectStatuses("local training", "updating model"); + // the new client finishes training and shares its update too + await client.expectStatuses("updating model"); + + // the server aggregates the round and answers with the new global + // weights along with the participants, resolving both pending rounds + await Promise.all([waiting.completeRound(), client.completeRound()]); + await waiting.expectParticipants(2); + await client.expectParticipants(2); + + return client; + } + + /** A round during which every client is present, so a), b) and c) run */ + async function runFullRound(...clients: readonly Client[]): Promise { + await Promise.all(clients.map(async (c) => await c.completeRound())); + + for (const client of clients) await client.expectStatuses(...FULL_ROUND); + // the server payload received during c) carries the participants + for (const client of clients) + await client.expectParticipants(clients.length); + } + + /** A client leaving, the remaining one is left without enough participants */ + async function leavesTask( + leaving: Client, + remaining: Client, + ): Promise { + await leaving.leave(); + + // the remaining client receives the WaitingForMoreParticipants message + await remaining.expectStatuses("not enough participants"); + await remaining.expectParticipants(1); + } + + /** A client starting a round while it knows it is the only participant */ + async function startsRoundAlone(client: Client): Promise { + client.startRound(); + + // it trains, then waits in c) for another participant + await client.expectStatuses("local training", "not enough participants"); + } + + it("a client joining alone trains then waits for a participant", async () => { + await joinsAlone("user 1"); + }); + + it("a joining client completes the round of the waiting one", async () => { + const user1 = await joinsAlone("user 1"); + + await joinsWaitingClient(user1, "user 2"); + }); + + it("a round runs the whole cycle when both clients are present", async () => { + const user1 = await joinsAlone("user 1"); + const user2 = await joinsWaitingClient(user1, "user 2"); + + await runFullRound(user1, user2); + }); + + it("a client is notified when a participant leaves", async () => { + const user1 = await joinsAlone("user 1"); + const user2 = await joinsWaitingClient(user1, "user 2"); + + await leavesTask(user1, user2); + }); + + it("a client left alone trains but waits to share its update", async () => { + const user1 = await joinsAlone("user 1"); + const user2 = await joinsWaitingClient(user1, "user 2"); + await leavesTask(user1, user2); + + await startsRoundAlone(user2); + }); + + it("a client joining mid-training completes the pending round", async () => { + const user1 = await joinsAlone("user 1"); + const user2 = await joinsWaitingClient(user1, "user 2"); + await leavesTask(user1, user2); + await startsRoundAlone(user2); + + // the server should accept user 3's weights (they should not be + // outdated) and aggregate them with user 2's pending round + await joinsWaitingClient(user2, "user 3"); + }); }); /** diff --git a/server/tests/e2e/helpers.ts b/server/tests/e2e/helpers.ts new file mode 100644 index 000000000..4d79981a1 --- /dev/null +++ b/server/tests/e2e/helpers.ts @@ -0,0 +1,237 @@ +import type { + DataFormat, + DataType, + Dataset, + Network, + RoundLogs, + RoundStatus, + Task, +} from "@epfml/discojs"; +import { Disco, WeightsContainer } from "@epfml/discojs"; +import { List } from "immutable"; +import { expect } from "vitest"; + +import { Queue } from "../utils.js"; + +// Array.fromAsync not yet widely used (2024) +export async function arrayFromAsync(iter: AsyncIterable): Promise { + const ret: T[] = []; + for await (const e of iter) { + // TODO trick to allow other Promises to run + // else one client might progress alone without communicating with others + // will be fixed when client orchestrations in the server is correctly done + await new Promise((resolve) => setTimeout(resolve, 10)); + + ret.push(e); + } + return ret; +} + +async function WSIntoList(ws: WeightsContainer): Promise>> { + return List( + (await Promise.all(ws.weights.map(async (w) => await w.data()))).map( + (arr) => List(arr), + ), + ); +} + +export async function expectWSToBeClose( + left: WeightsContainer, + right: WeightsContainer, +): Promise { + for (const tensors of (await WSIntoList(left)).zip(await WSIntoList(right))) + for (const [l, r] of tensors[0].zip(tensors[1])) + expect(l).to.be.closeTo(r, 1e-4); +} + +/** Every given model should be close to the first one */ +export async function expectAllWSToBeClose( + ...weights: readonly WeightsContainer[] +): Promise { + const [reference, ...others] = weights; + + await Promise.all( + others.map(async (current) => await expectWSToBeClose(reference, current)), + ); +} + +/** The models a participant held at round boundaries */ +export interface ModelRecording { + /** Every recorded model, oldest first. Only the last one with `keep: "latest"`. */ + all: () => readonly WeightsContainer[]; + /** The most recently recorded model, throws if there is none yet. */ + latest: () => WeightsContainer; + /** Release every recorded model */ + dispose: () => void; +} + +/** + * Records the model a peer holds at each round boundary: once + * onRoundEndCommunication has returned and before the next local round trains + * on it. Peers hold the very same model at those points, whereas + * `trainer.model.weights` read afterwards also contains each peer's own local + * training, which is not reproducible across peers. + * + * Pass `keep: "latest"` in the tests measuring tensor memory so that the + * recorded models don't grow with the number of rounds. + */ +export function recordModelsAtRoundBoundary< + D extends DataType, + N extends Network, +>( + disco: Disco, + { keep = "all" }: { keep?: "all" | "latest" } = {}, +): ModelRecording { + const models: WeightsContainer[] = []; + + disco.on("status", (status) => { + if (status !== "local training") return; + if (keep === "latest") models.splice(0).forEach((m) => m.dispose()); + models.push( + new WeightsContainer( + disco.trainer.model.weights.weights.map((w) => w.clone()), + ), + ); + }); + + return { + all: () => models, + latest: () => { + const model = models.at(-1); + if (model === undefined) + throw new Error("the peer hasn't reached a round boundary yet"); + return model; + }, + dispose: () => models.splice(0).forEach((m) => m.dispose()), + }; +} + +/** The peers should hold the same model at their latest round boundary */ +export async function expectPeersToAgreeOnModel( + ...recordings: readonly Pick[] +): Promise { + const [first, ...others] = recordings; + for (const other of others) + await expectWSToBeClose(first.latest(), other.latest()); +} + +/** + * A client taking part in a task, with the events it emits recorded so that a + * test can assert them in the order they were emitted. + * + * Rounds are driven in two steps because a single call to + * `trainByRound().next()` performs a whole round and only resolves once the + * other participants completed theirs: {@link startRound} kicks a round off and + * {@link completeRound} awaits it. A test can therefore assert the events + * emitted while a round is still in flight, which is how a participant waiting + * for others is observed. + */ +export class Participant { + readonly disco: Disco; + /** The models held at round boundaries, only the latest one is kept */ + readonly modelsAtRoundBoundary: ModelRecording; + + readonly #name: string; + readonly #rounds: AsyncGenerator; + readonly #statuses = new Queue(); + readonly #participants = new Queue(); + readonly #syncedModel: Promise; + + #pendingRound: Promise> | undefined; + #left = false; + + constructor( + name: string, + task: Task, + url: URL, + dataset: Dataset, + ) { + this.#name = name; + this.disco = new Disco(task, url, { + preprocessOnce: true, + debugLabel: name, + }); + + this.disco.on("status", (status) => this.#statuses.put(status)); + this.disco.on("participants", (participants) => + this.#participants.put(participants), + ); + this.#syncedModel = new Promise((resolve) => + this.disco.on("modelSynced", (weights) => { + if (weights !== undefined) resolve(weights); + }), + ); + + this.modelsAtRoundBoundary = recordModelsAtRoundBoundary(this.disco, { + keep: "latest", + }); + this.#rounds = this.disco.trainByRound(dataset); + } + + /** Start a round without waiting for it to complete */ + startRound(): this { + if (this.#pendingRound !== undefined) + throw new Error(`${this.#name} is already running a round`); + + const round = this.#rounds.next(); + // a round can stay pending forever, e.g. when the participant ends up + // alone, and leaving the task then rejects it: never let that surface as + // an unhandled rejection + round.catch(() => undefined); + this.#pendingRound = round; + + return this; + } + + /** Wait for the current round, starting one if none is running */ + async completeRound(): Promise { + if (this.#pendingRound === undefined) this.startRound(); + const pending = this.#pendingRound; + if (pending === undefined) throw new Error("unreachable"); + + const round = await pending; + this.#pendingRound = undefined; + + if (round.done) + throw new Error(`${this.#name} stopped training earlier than expected`); + + return round.value; + } + + /** Assert the next emitted statuses, in order */ + async expectStatuses(...expected: readonly RoundStatus[]): Promise { + for (const status of expected) + expect(await this.#statuses.next(), `${this.#name} status`).to.equal( + status, + ); + } + + /** Assert the next emitted participant counts, in order */ + async expectParticipants(...expected: readonly number[]): Promise { + for (const count of expected) + expect( + await this.#participants.next(), + `${this.#name} participants`, + ).to.equal(count); + } + + /** + * The model the participant synced with when joining mid-training, + * only resolves for a participant which joined an ongoing session. + */ + syncedModel(): Promise { + return this.#syncedModel; + } + + /** Leave the task, releasing everything the participant holds */ + async leave(): Promise { + if (this.#left) return; + this.#left = true; + + try { + await this.disco.close(); + } finally { + this.modelsAtRoundBoundary.dispose(); + } + } +}