Parallelize xet uploads (network) (#1715)

See #1704

This parallelizes at least network calls, a complementary PR would put
the wasm module into a worker to parallelize CPU.

cc @assafvayner @jsulz @Kakulukian

Author: coyotte508

packages/hub/src/lib/commit.ts

@@ -23,6 +23,8 @@ import { base64FromBytes } from "../utils/base64FromBytes";
 import { isFrontend } from "../utils/isFrontend";
 import { createBlobs } from "../utils/createBlobs";
 import { uploadShards } from "../utils/uploadShards";
+import { splitAsyncGenerator } from "../utils/splitAsyncGenerator";
+import { mergeAsyncGenerators } from "../utils/mergeAsyncGenerators";
 
 const CONCURRENT_SHAS = 5;
 const CONCURRENT_LFS_UPLOADS = 5;
@@ -323,43 +325,46 @@ export async function* commitIter(params: CommitParams): AsyncGenerator<CommitPr
 						};
 					}
 				}
-				for await (const event of uploadShards(
-					(async function* () {
-						for (const obj of json.objects) {
-							const op = shaToOperation.get(obj.oid);
-							if (!op || !obj.actions?.upload) {
-								continue;
-							}
-							abortSignal?.throwIfAborted();
-
-							yield { content: op.content, path: op.path, sha256: obj.oid };
+				const source = (async function* () {
+					for (const obj of json.objects) {
+						const op = shaToOperation.get(obj.oid);
+						if (!op || !obj.actions?.upload) {
+							continue;
 						}
-					})(),
-					{
-						fetch: params.fetch,
-						accessToken,
-						hubUrl: params.hubUrl ?? HUB_URL,
-						repo: repoId,
-						// todo: maybe leave empty if PR?
-						rev: params.branch ?? "main",
-					}
-				)) {
-					if (event.event === "file") {
-						yield {
-							event: "fileProgress",
-							path: event.path,
-							progress: 1,
-							state: "uploading",
-						};
-					} else if (event.event === "fileProgress") {
-						yield {
-							event: "fileProgress",
-							path: event.path,
-							progress: event.progress,
-							state: "uploading",
-						};
+						abortSignal?.throwIfAborted();
+
+						yield { content: op.content, path: op.path, sha256: obj.oid };
 					}
-				}
+				})();
+				const sources = splitAsyncGenerator(source, 5);
+				yield* mergeAsyncGenerators(
+					sources.map(async function* (source) {
+						for await (const event of uploadShards(source, {
+							fetch: params.fetch,
+							accessToken,
+							hubUrl: params.hubUrl ?? HUB_URL,
+							repo: repoId,
+							// todo: maybe leave empty if PR?
+							rev: params.branch ?? "main",
+						})) {
+							if (event.event === "file") {
+								yield {
+									event: "fileProgress" as const,
+									path: event.path,
+									progress: 1,
+									state: "uploading" as const,
+								};
+							} else if (event.event === "fileProgress") {
+								yield {
+									event: "fileProgress" as const,
+									path: event.path,
+									progress: event.progress,
+									state: "uploading" as const,
+								};
+							}
+						}
+					})
+				);
 			} else {
 				yield* eventToGenerator<CommitProgressEvent, void>((yieldCallback, returnCallback, rejectCallback) => {
 					return promisesQueueStreaming(

packages/hub/src/utils/mergeAsyncGenerators.spec.ts

@@ -0,0 +1,94 @@
+import { describe, expect, it } from "vitest";
+import { mergeAsyncGenerators } from "./mergeAsyncGenerators";
+import { splitAsyncGenerator } from "./splitAsyncGenerator";
+
+describe("mergeAsyncGenerators", () => {
+	const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms));
+	it("should merge multiple async generators", async () => {
+		const generator1 = (async function* () {
+			yield 1;
+			yield 2;
+			await sleep(250);
+			yield 3;
+		})();
+		const generator2 = (async function* () {
+			await sleep(100);
+			yield 4;
+			yield 5;
+			yield 6;
+		})();
+		const generator3 = (async function* () {
+			await sleep(200);
+			yield 7;
+			yield 8;
+			yield 9;
+		})();
+
+		const results: number[] = [];
+
+		for await (const result of mergeAsyncGenerators([generator1, generator2, generator3])) {
+			results.push(result);
+		}
+		expect(results).toEqual([1, 2, 4, 5, 6, 7, 8, 9, 3]);
+	});
+
+	it("should merge multiple async generators from a single source", async () => {
+		const source = (async function* () {
+			yield 1;
+			yield 2;
+			yield 3;
+			yield 4;
+			yield 5;
+			yield 6;
+			yield 7;
+			yield 8;
+			yield 9;
+		})();
+		const sources = splitAsyncGenerator(source, 3);
+
+		const generator1 = (async function* () {
+			for await (const result of sources[0]) {
+				yield { result, gen: 1 };
+				await sleep(100);
+			}
+		})();
+
+		const generator2 = (async function* () {
+			await sleep(50);
+			for await (const result of sources[1]) {
+				yield { result, gen: 2 };
+				await sleep(100);
+			}
+		})();
+
+		const generator3 = (async function* () {
+			await sleep(80);
+			let count = 0;
+			for await (const result of sources[2]) {
+				yield { result, gen: 3 };
+				count++;
+
+				if (count >= 2) {
+					return;
+				}
+			}
+		})();
+
+		const results: { result: number; gen: number }[] = [];
+		for await (const result of mergeAsyncGenerators([generator1, generator2, generator3])) {
+			results.push(result);
+		}
+		expect(results.length).toBe(9);
+		expect(results).toEqual([
+			{ result: 1, gen: 1 },
+			{ result: 2, gen: 2 },
+			{ result: 3, gen: 3 },
+			{ result: 4, gen: 3 },
+			{ result: 5, gen: 1 },
+			{ result: 6, gen: 2 },
+			{ result: 7, gen: 1 },
+			{ result: 8, gen: 2 },
+			{ result: 9, gen: 1 },
+		]);
+	});
+});

packages/hub/src/utils/mergeAsyncGenerators.ts

@@ -0,0 +1,34 @@
+/**
+ * Merge outputs of multiple async generators.
+ */
+export async function* mergeAsyncGenerators<T>(generators: AsyncGenerator<T>[]): AsyncGenerator<T> {
+	const executing: Promise<{ result: IteratorResult<T>; gen: AsyncGenerator<T> }>[] = [];
+
+	const generatorSymbol = Symbol("generator");
+
+	for (const gen of generators) {
+		const p = gen.next().then((result) => ({ result, gen }));
+		Object.defineProperty(p, generatorSymbol, {
+			value: gen,
+		});
+		executing.push(p);
+	}
+
+	while (executing.length > 0) {
+		const next = await Promise.race(executing);
+		const { result, gen } = next;
+
+		const index = executing.findIndex((p) => Object.getOwnPropertyDescriptor(p, generatorSymbol)?.value === gen);
+
+		if (result.done) {
+			executing.splice(index, 1);
+			continue;
+		}
+
+		yield result.value;
+		executing[index] = gen.next().then((result) => ({ result, gen }));
+		Object.defineProperty(executing[index], generatorSymbol, {
+			value: gen,
+		});
+	}
+}

packages/hub/src/utils/splitAsyncGenerator.ts

@@ -0,0 +1,43 @@
+/**
+ * Split an async generator into multiple async generators, all drawing from the same source.
+ */
+export function splitAsyncGenerator<T>(source: AsyncGenerator<T>, n: number): Array<AsyncGenerator<T>> {
+	if (n <= 0) {
+		return [];
+	}
+
+	const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms));
+
+	let takenIndex: number | null = null;
+	const generators: AsyncGenerator<T>[] = [];
+	let remaining = n;
+	for (let i = 0; i < n; i++) {
+		generators.push({
+			next: async () => {
+				while (takenIndex !== null) {
+					await sleep(1);
+				}
+				takenIndex = i;
+				return source.next().then((r) => {
+					takenIndex = null;
+					return r;
+				});
+			},
+			return: async () => {
+				remaining--;
+				if (remaining === 0) {
+					return source.return(undefined);
+				}
+				return {
+					done: true,
+					value: undefined,
+				};
+			},
+			throw: async (error) => {
+				return source.throw(error);
+			},
+			[Symbol.asyncIterator]: () => generators[i],
+		});
+	}
+	return generators;
+}