Xet upload: Fixup file progress (#1728)

Follow up #1727

Tested the wrong file that was already deduplicated

Author: coyotte508

packages/hub/scripts/bench.ts

@@ -2,14 +2,16 @@ import { uploadShards } from "../src/utils/uploadShards.js";
 import { sha256 } from "../src/utils/sha256.js";
 import { parseArgs } from "node:util";
 import { tmpdir } from "node:os";
-import { join } from "node:path";
+import path, { join } from "node:path";
 import { writeFile, readFile, stat, mkdir } from "node:fs/promises";
 import type { RepoId } from "../src/types/public.js";
 import { toRepoId } from "../src/utils/toRepoId.js";
 import type { CommitOperation } from "../src/index.js";
 import { commitIter, downloadFile } from "../src/index.js";
 import { SplicedBlob } from "../src/utils/SplicedBlob.js";
 import { pathToFileURL } from "node:url";
+import { existsSync } from "node:fs";
+import { FileBlob } from "../src/utils/FileBlob.js";
 
 /**
  * This script downloads the files from openai-community/gpt2 and simulates an upload to a xet repo.
@@ -37,6 +39,11 @@ const FILES_TO_DOWNLOAD = [
 		filename: "64.tflite",
 		sha256: "cfcd510b239d90b71ee87d4e57a5a8c2d55b2a941e5d9fe5852298268ddbe61b",
 	},
+	{
+		url: "https://huggingface.co/openai-community/gpt2/resolve/main/model.safetensors?download=true",
+		filename: "model.safetensors",
+		sha256: "248dfc3911869ec493c76e65bf2fcf7f615828b0254c12b473182f0f81d3a707",
+	},
 ];
 
 const FILES_TO_EDIT = [
@@ -84,8 +91,7 @@ async function* createFileSource(files: Array<{ filepath: string; filename: stri
 }> {
 	for (const file of files) {
 		console.log(`Processing ${file.filename}...`);
-		const buffer = await readFile(file.filepath);
-		let blob = new Blob([buffer]);
+		let blob: Blob = await FileBlob.create(file.filepath);
 
 		if (file.filename.endsWith(".edited")) {
 			const edits = FILES_TO_EDIT.find((f) => f.filename === file.filename)?.edits;
@@ -110,7 +116,7 @@ async function* createFileSource(files: Array<{ filepath: string; filename: stri
 		const sha256ToCheck =
 			FILES_TO_DOWNLOAD.find((f) => f.filename === file.filename)?.sha256 ||
 			FILES_TO_EDIT.find((f) => f.filename === file.filename)?.sha256;
-		if (sha256Hash !== sha256ToCheck) {
+		if (sha256ToCheck !== undefined && sha256Hash !== sha256ToCheck) {
 			throw new Error(`SHA256 mismatch for ${file.filename}: ${sha256Hash} !== ${sha256ToCheck}`);
 		}
 
@@ -214,6 +220,10 @@ async function main() {
 				short: "c",
 				default: false,
 			},
+			localFilePath: {
+				type: "string",
+				short: "f",
+			},
 			write: {
 				type: "boolean",
 				short: "w",
@@ -250,6 +260,13 @@ async function main() {
 		files.push({ filepath, filename: fileInfo.filename });
 	}
 
+	if (args.localFilePath) {
+		if (!existsSync(args.localFilePath)) {
+			throw new Error(`Local file ${args.localFilePath} does not exist`);
+		}
+		files.push({ filepath: args.localFilePath, filename: path.basename(args.localFilePath) });
+	}
+
 	// Parse repo
 	const repoName = args.repo;
 
@@ -279,7 +296,20 @@ async function main() {
 	// Process files through uploadShards
 	const fileSource = createFileSource(files);
 
-	for await (const event of uploadShards(fileSource, uploadParams)) {
+	const fileProgress: Record<string, number> = {};
+
+	for await (const event of uploadShards(fileSource, {
+		...uploadParams,
+		yieldCallback: (event) => {
+			if (!fileProgress[event.path]) {
+				fileProgress[event.path] = event.progress;
+			}
+			if (event.progress < fileProgress[event.path]) {
+				throw new Error(`Progress for ${event.path} went down from ${fileProgress[event.path]} to ${event.progress}`);
+			}
+			fileProgress[event.path] = event.progress;
+		},
+	})) {
 		switch (event.event) {
 			case "file": {
 				console.log(`\n📁 Processed file: ${event.path}`);
@@ -303,6 +333,14 @@ async function main() {
 			case "fileProgress": {
 				const progress = (event.progress * 100).toFixed(1);
 				console.log(`   📈 Progress for ${event.path}: ${progress}%`);
+
+				if (!fileProgress[event.path]) {
+					fileProgress[event.path] = event.progress;
+				}
+				if (event.progress < fileProgress[event.path]) {
+					throw new Error(`Progress for ${event.path} went down from ${fileProgress[event.path]} to ${event.progress}`);
+				}
+				fileProgress[event.path] = event.progress;
 				break;
 			}
 		}

packages/hub/src/utils/FileBlob.ts

@@ -96,6 +96,10 @@ export class FileBlob extends Blob {
 	 * Returns a stream around the part of the file delimited by the FileBlob.
 	 */
 	override stream(): ReturnType<Blob["stream"]> {
+		if (this.start === this.end) {
+			return new Blob([]).stream();
+		}
+
 		return Readable.toWeb(createReadStream(this.path, { start: this.start, end: this.end - 1 })) as ReturnType<
 			Blob["stream"]
 		>;

packages/hub/src/utils/createXorbs.ts

@@ -40,6 +40,7 @@ class CurrentXorbInfo {
 
 	fileProcessedBytes: Record<string, number>;
 	fileUploadedBytes: Record<string, number>;
+	fileSize: Record<string, number>;
 	data: Uint8Array;
 	immutableData: {
 		chunkIndex: number;
@@ -52,6 +53,7 @@ class CurrentXorbInfo {
 		this.chunks = [];
 		this.fileProcessedBytes = {};
 		this.fileUploadedBytes = {};
+		this.fileSize = {};
 		this.data = new Uint8Array(XORB_SIZE);
 		this.immutableData = null;
 	}
@@ -68,10 +70,13 @@ class CurrentXorbInfo {
 			hash: computeXorbHash(xorbChunksCleaned),
 			chunks: xorbChunksCleaned,
 			id: this.id,
-			files: Object.entries(this.fileProcessedBytes).map(([path, progress]) => ({
+			files: Object.entries(this.fileProcessedBytes).map(([path, processedBytes]) => ({
 				path,
-				progress,
-				lastSentProgress: this.fileUploadedBytes[path] ?? 0,
+				progress: processedBytes / this.fileSize[path],
+				lastSentProgress:
+					((this.fileUploadedBytes[path] ?? 0) +
+						(processedBytes - (this.fileUploadedBytes[path] ?? 0)) * PROCESSING_PROGRESS_RATIO) /
+					this.fileSize[path],
 			})),
 		};
 	}
@@ -111,7 +116,7 @@ export async function* createXorbs(
 	const chunkCache = new ChunkCache();
 	let xorb = new CurrentXorbInfo();
 
-	const nextXorb = (currentFile: { path: string; uploadedBytes: number }): XorbEvent => {
+	const nextXorb = (currentFile: { path: string; uploadedBytes: number; size: number }): XorbEvent => {
 		const event = xorb.event(chunkModule.compute_xorb_hash.bind(chunkModule));
 
 		xorbId++;
@@ -120,6 +125,7 @@ export async function* createXorbs(
 		xorb.fileUploadedBytes = {
 			[currentFile.path]: currentFile.uploadedBytes,
 		};
+		xorb.fileSize[currentFile.path] = currentFile.size;
 
 		return event;
 	};
@@ -143,10 +149,12 @@ export async function* createXorbs(
 
 	try {
 		for await (const fileSource of fileSources) {
+			xorb.fileSize[fileSource.path] = fileSource.content.size;
+
 			// Load dedup info for the first chunk of the file, if it's potentially modified by the splice
 			if (fileSource.content instanceof SplicedBlob && fileSource.content.firstSpliceIndex < MAX_CHUNK_SIZE) {
 				await loadDedupInfoToCache(
-					fileSource.content.originalBlob.slice(0, fileSource.content.firstSpliceIndex),
+					fileSource.content.originalBlob.slice(0, MAX_CHUNK_SIZE),
 					remoteXorbHashes,
 					params,
 					chunkCache,
@@ -235,7 +243,7 @@ export async function* createXorbs(
 					if (cacheData === undefined) {
 						if (!writeChunk(xorb, chunkToCopy, chunk.hash)) {
 							// Failure to write chunk, maybe because it went over xorb size limit
-							yield nextXorb({ path: fileSource.path, uploadedBytes: processedBytes / fileSource.content.size });
+							yield nextXorb({ path: fileSource.path, uploadedBytes: processedBytes, size: fileSource.content.size });
 
 							chunkIndex = 0;
 							chunkXorbId = xorbId;
@@ -290,7 +298,7 @@ export async function* createXorbs(
 					}
 
 					if (xorb.chunks.length >= MAX_XORB_CHUNKS) {
-						yield nextXorb({ path: fileSource.path, uploadedBytes: processedBytes });
+						yield nextXorb({ path: fileSource.path, uploadedBytes: processedBytes, size: fileSource.content.size });
 
 						for (const event of pendingFileEvents) {
 							event.representation = event.representation.map((rep) => ({