Efficient Xet download

packages/hub/src/lib/list-files.spec.ts

@@ -60,7 +60,7 @@ describe("listFiles", () => {
 				type: "file",
 			},
 		]);
-	}, 30_000);
+	});
 
 	it("should fetch the list of files from the repo, including last commit", async () => {
 		const cursor = listFiles({
@@ -146,7 +146,7 @@ describe("listFiles", () => {
 				type: "file",
 			},
 		]);
-	}, 30_000);
+	});
 
 	it("should fetch the list of files from the repo, including subfolders", async () => {
 		const cursor = listFiles({
@@ -165,5 +165,5 @@ describe("listFiles", () => {
 		}
 
 		assert(files.some((file) => file.path === "data/XSUM-EMNLP18-Summary-Data-Original.tar.gz"));
-	}, 30_000);
+	});
 });

packages/hub/src/lib/parse-safetensors-metadata.spec.ts

@@ -49,7 +49,7 @@ describe("parseSafetensorsMetadata", () => {
 		assert.deepStrictEqual(parse.parameterCount, { BF16: 176_247_271_424 });
 		assert.deepStrictEqual(sum(Object.values(parse.parameterCount)), 176_247_271_424);
 		// total params = 176B
-	}, 30_000);
+	});
 
 	it("fetch info for single-file with multiple dtypes", async () => {
 		const parse = await parseSafetensorsMetadata({

packages/hub/src/lib/upload-file.spec.ts

@@ -94,5 +94,5 @@ describe("uploadFile", () => {
 				hubUrl: TEST_HUB_URL,
 			});
 		}
-	}, 30_000);
+	});
 });

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

@@ -0,0 +1,96 @@
+import { describe, it, expect } from "vitest";
+import { RangeList } from "./RangeList";
+
+describe("RangeList", () => {
+	it("should add a single range", () => {
+		const rangeList = new RangeList();
+		rangeList.add(1, 100);
+
+		const ranges = rangeList.getAllRanges();
+		expect(ranges).toHaveLength(1);
+		expect(ranges[0]).toEqual({
+			start: 1,
+			end: 100,
+			refCount: 1,
+			data: null,
+		});
+	});
+
+	it("should handle overlapping ranges", () => {
+		const rangeList = new RangeList();
+		rangeList.add(1, 100);
+		rangeList.add(30, 50);
+
+		const ranges = rangeList.getAllRanges();
+		expect(ranges).toHaveLength(3);
+		expect(ranges).toEqual([
+			{ start: 1, end: 30, refCount: 1, data: null },
+			{ start: 30, end: 50, refCount: 2, data: null },
+			{ start: 50, end: 100, refCount: 1, data: null },
+		]);
+	});
+
+	it("should remove a range at existing boundaries", () => {
+		const rangeList = new RangeList();
+		rangeList.add(1, 100);
+		rangeList.add(30, 50);
+		rangeList.remove(30, 50);
+
+		const ranges = rangeList.getAllRanges();
+		expect(ranges).toHaveLength(3);
+		expect(ranges).toEqual([
+			{ start: 1, end: 30, refCount: 1, data: null },
+			{ start: 30, end: 50, refCount: 1, data: null },
+			{ start: 50, end: 100, refCount: 1, data: null },
+		]);
+	});
+
+	it("should throw error when removing range at non-existing boundaries", () => {
+		const rangeList = new RangeList();
+		rangeList.add(1, 100);
+		rangeList.add(30, 50);
+
+		expect(() => rangeList.remove(2, 50)).toThrow("Range boundaries must match existing boundaries");
+	});
+
+	it("should get ranges within boundaries", () => {
+		const rangeList = new RangeList();
+		rangeList.add(1, 100);
+		rangeList.add(30, 50);
+
+		const ranges = rangeList.getRanges(30, 100);
+		expect(ranges).toHaveLength(2);
+		expect(ranges).toEqual([
+			{ start: 30, end: 50, refCount: 2, data: null },
+			{ start: 50, end: 100, refCount: 1, data: null },
+		]);
+	});
+
+	it("should throw error when end is less than or equal to start", () => {
+		const rangeList = new RangeList();
+
+		expect(() => rangeList.add(100, 1)).toThrow("End must be greater than start");
+		expect(() => rangeList.add(1, 1)).toThrow("End must be greater than start");
+		expect(() => rangeList.remove(100, 1)).toThrow("End must be greater than start");
+		expect(() => rangeList.remove(1, 1)).toThrow("End must be greater than start");
+		expect(() => rangeList.getRanges(100, 1)).toThrow("End must be greater than start");
+		expect(() => rangeList.getRanges(1, 1)).toThrow("End must be greater than start");
+	});
+
+	it("should handle multiple overlapping ranges", () => {
+		const rangeList = new RangeList();
+		rangeList.add(1, 100);
+		rangeList.add(30, 50);
+		rangeList.add(40, 60);
+
+		const ranges = rangeList.getAllRanges();
+		expect(ranges).toHaveLength(5);
+		expect(ranges).toEqual([
+			{ start: 1, end: 30, refCount: 1, data: null },
+			{ start: 30, end: 40, refCount: 2, data: null },
+			{ start: 40, end: 50, refCount: 3, data: null },
+			{ start: 50, end: 60, refCount: 2, data: null },
+			{ start: 60, end: 100, refCount: 1, data: null },
+		]);
+	});
+});

packages/hub/src/utils/RangeList.ts

@@ -0,0 +1,179 @@
+/**
+ * Code generated with this prompt by Cursor:
+ *
+ * I want to build a class to manage ranges
+ *
+ * I can add ranges to it with a start& an end (both integer, end > start). It should store those ranges efficently.
+ *
+ * When several ranges overlap, eg [1, 100] and [30, 50], I want the class to split the range into non-overlapping ranges, and add a "ref counter" to the ranges. For example,  [1, 30], [30, 50] * 2, [50, 100]
+ *
+ * I also want to be able to remove ranges, it will decrease the ref counter or remove the range altogether. I can only remove ranges at existing boundaries. For example, with the [1, 30], [30, 50] * 2, [50, 100] configuration
+ *
+ * - removing [1, 100] => the only range remaning is [30, 50]
+ * - removing [2, 50] => error, because "2' is not a boundary
+ * - removing [30, 50] => [1, 30], [30, 50], [50, 100] (do not "merge" the ranges back together)
+ *
+ * I want to be able to associate data to each range. And I want to be able to get the ranges inside boundaries. For example , with [1, 30], [30, 50] * 2, [50, 100] configuration
+ *
+ * - getting [30, 100] => I receive [30, 50] * 2, [50, 100], and I can get / modify the data assocaited to each range by accessing their data prop. Note the "*2" is just the ref counter, there is onlly one range object for the interval returned
+ * - getting [2, 50] => I get [30, 50] * 2
+ *
+ * ----
+ *
+ * Could optimize with binary search, but the ranges we want to handle are not that many.
+ */
+interface Range<T> {
+	start: number;
+	end: number;
+	refCount: number;
+	data: T | null;
+}
+
+export class RangeList<T> {
+	private ranges: Range<T>[] = [];
+
+	/**
+	 * Add a range to the list. If it overlaps with existing ranges,
+	 * it will split them and increment reference counts accordingly.
+	 */
+	add(start: number, end: number): void {
+		if (end <= start) {
+			throw new TypeError("End must be greater than start");
+		}
+
+		// Find all ranges that overlap with the new range
+		const overlappingRanges: { index: number; range: Range<T> }[] = [];
+		for (let i = 0; i < this.ranges.length; i++) {
+			const range = this.ranges[i];
+			if (start < range.end && end > range.start) {
+				overlappingRanges.push({ index: i, range });
+			}
+			if (range.data !== null) {
+				throw new Error("Overlapping range already has data");
+			}
+		}
+
+		if (overlappingRanges.length === 0) {
+			// No overlaps, just add the new range
+			this.ranges.push({ start, end, refCount: 1, data: null });
+			this.ranges.sort((a, b) => a.start - b.start);
+			return;
+		}
+
+		// Handle overlaps by splitting ranges
+		const newRanges: Range<T>[] = [];
+		let currentPos = start;
+
+		for (let i = 0; i < overlappingRanges.length; i++) {
+			const { range } = overlappingRanges[i];
+
+			// Add range before overlap if exists
+			if (currentPos < range.start) {
+				newRanges.push({
+					start: currentPos,
+					end: range.start,
+					refCount: 1,
+					data: null,
+				});
+			} else if (range.start < currentPos) {
+				newRanges.push({
+					start: range.start,
+					end: currentPos,
+					refCount: range.refCount,
+					data: null,
+				});
+			}
+
+			// Add overlapping part with increased ref count
+			newRanges.push({
+				start: Math.max(currentPos, range.start),
+				end: Math.min(end, range.end),
+				refCount: range.refCount + 1,
+				data: null,
+			});
+
+			// Add remaining part of existing range if exists
+			if (range.end > end) {
+				newRanges.push({
+					start: end,
+					end: range.end,
+					refCount: range.refCount,
+					data: null,
+				});
+			}
+
+			currentPos = Math.max(currentPos, range.end);
+		}
+
+		// Add remaining part after last overlap if exists
+		if (currentPos < end) {
+			newRanges.push({
+				start: currentPos,
+				end,
+				refCount: 1,
+				data: null,
+			});
+		}
+
+		// Remove old overlapping ranges and insert new ones
+		const firstIndex = overlappingRanges[0].index;
+		const lastIndex = overlappingRanges[overlappingRanges.length - 1].index;
+		this.ranges.splice(firstIndex, lastIndex - firstIndex + 1, ...newRanges);
+		this.ranges.sort((a, b) => a.start - b.start);
+	}
+
+	/**
+	 * Remove a range from the list. The range must start and end at existing boundaries.
+	 */
+	remove(start: number, end: number): void {
+		if (end <= start) {
+			throw new TypeError("End must be greater than start");
+		}
+
+		// Find ranges that need to be modified
+		const affectedRanges: { index: number; range: Range<T> }[] = [];
+		for (let i = 0; i < this.ranges.length; i++) {
+			const range = this.ranges[i];
+			if (start < range.end && end > range.start) {
+				affectedRanges.push({ index: i, range });
+			}
+		}
+
+		if (affectedRanges.length === 0) {
+			throw new Error("No ranges found to remove");
+		}
+
+		// Verify boundaries match
+		if (start !== affectedRanges[0].range.start || end !== affectedRanges[affectedRanges.length - 1].range.end) {
+			throw new Error("Range boundaries must match existing boundaries");
+		}
+
+		// Todo: also check if there's a gap in the middle but it should not happen with our usage
+
+		for (let i = 0; i < affectedRanges.length; i++) {
+			const { range } = affectedRanges[i];
+
+			range.refCount--;
+		}
+
+		this.ranges = this.ranges.filter((range) => range.refCount > 0);
+	}
+
+	/**
+	 * Get all ranges within the specified boundaries.
+	 */
+	getRanges(start: number, end: number): Range<T>[] {
+		if (end <= start) {
+			throw new TypeError("End must be greater than start");
+		}
+
+		return this.ranges.filter((range) => start < range.end && end > range.start);
+	}
+
+	/**
+	 * Get all ranges in the list
+	 */
+	getAllRanges(): Range<T>[] {
+		return [...this.ranges];
+	}
+}

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

@@ -58,7 +58,7 @@ describe("WebBlob", () => {
 		expect(webBlob).toBeInstanceOf(WebBlob);
 		expect(webBlob).toMatchObject({ url });
 		expect(await webBlob.slice(10, 22).text()).toBe("__metadata__");
-	}, 30_000);
+	});
 
 	it("should lazy load a Xet file hosted on Hugging Face", async () => {
 		const stableDiffusionUrl =
@@ -70,7 +70,7 @@ describe("WebBlob", () => {
 		expect(webBlob).toBeInstanceOf(WebBlob);
 		expect(webBlob).toMatchObject({ url });
 		expect(await webBlob.slice(10, 22).text()).toBe("__metadata__");
-	}, 30_000);
+	});
 
 	it("should create a slice on the file", async () => {
 		const expectedText = fullText.slice(10, 20);