bug fixes; profiling setup

Author: seanses

hf_xet_wasm/Cargo.lock

@@ -36,7 +36,7 @@ web-sys = { version = "0.3.72", features = [
 ] }
 js-sys = "0.3.72"
 futures = "0.3.31"
-sha2 = "0.10.8"
+sha2 = { version = "0.10.8", features = ["asm"] }
 blake3 = "1.7.0"
 getrandom = { version = "0.3", features = ["wasm_js"] }
 wasm_thread = "0.3"

hf_xet_wasm/Cargo.toml

@@ -9,7 +9,9 @@ cargo install wasm-bindgen-cli
 - Build with `./build_wasm.sh` (bash) 
 
 #### Run Instructions
-Serve the web directory using a local http server, for example, https://crates.io/crates/sfz.
+First fill up the four `[FILL_ME]` fields in examples/index.html with a desired testing target.
+
+Then serve the web directory using a local http server, for example, https://crates.io/crates/sfz.
 
 - Install sfz:
 ```bash

hf_xet_wasm/README.md

@@ -1,6 +1,7 @@
 /**
  * Commits a file to a Hugging Face dataset.
  *
+ * @param {string} hf_endpoint - The HF Hub endpoint.
  * @param {string} file_name The name of the file to commit.
  * @param {string} sha256 The SHA256 hash of the file (as a string).
  * @param {number} file_size The size of the file in bytes.
@@ -10,7 +11,7 @@
  * @param {string} hf_token The HF token for auth
  * @returns {Promise<string>} A promise that resolves if the commit is successful, or rejects if an error occurs.
  */
-async function commit(endpoint, file_name, sha256, file_size, repo_type, repo_id, revision, hf_token) {
+async function commit(hf_endpoint, file_name, sha256, file_size, repo_type, repo_id, revision, hf_token) {
     const obj1 = {
         key: "header",
         value: {
@@ -32,7 +33,7 @@ async function commit(endpoint, file_name, sha256, file_size, repo_type, repo_id
     // Serialize to JSON string and concatenate with a newline.
     const body = `${JSON.stringify(obj1)}\n${JSON.stringify(obj2)}`;
 
-    const url = `${endpoint}/api/${repo_type}s/${repo_id}/commit/${revision}`;
+    const url = `${hf_endpoint}/api/${repo_type}s/${repo_id}/commit/${revision}`;
 
     try {
         const response = await fetch(url, {

hf_xet_wasm/examples/commit.js

@@ -19,17 +19,17 @@
         document.getElementById("file_picker").addEventListener(
             "change",
             async function () {
-                const endpoint = "http://localhost:5564"
-                const repo_type = "<your repo type>";
-                const repo_id = "<your repo>";
-                const hf_token = "<your token>";
+                const hf_endpoint = "[FILL_ME]";
+                const repo_type = "[FILL_ME]";
+                const repo_id = "[FILL_ME]";
+                const hf_token = "[FILL_ME]";
 
                 const file = this.files[0];
 
                 console.log("getting auth info...");
 
                 const xetmetadata = await fetchXetMetadataFromRepoInfo({
-                    endpoint,
+                    hfEndpoint: hf_endpoint,
                     tokenType: "write",
                     repoId: repo_id,
                     repoType: repo_type,
@@ -76,21 +76,21 @@
 
                 const casEndpoint = xetmetadata.endpoint;
                 console.log("casEndpoint:", casEndpoint);
-								const tokenRefresher = new TokenRefresher();
-								const tokenInfo = new TokenInfo(xetmetadata.accessToken, xetmetadata.expirationUnixEpoch);
-								const xetSession = new XetSession(casEndpoint, tokenInfo, tokenRefresher);
+                const tokenRefresher = new TokenRefresher();
+                const tokenInfo = new TokenInfo(xetmetadata.accessToken, xetmetadata.expirationUnixEpoch);
+                const xetSession = new XetSession(casEndpoint, tokenInfo, tokenRefresher);
 
-								const pf = await xetSession.uploadFileFromBlob("file_being_tracked", file);
-								console.log("file cleaned to Xet session", pf);
-								await xetSession.finalize();
+                const pf = await xetSession.uploadFileFromBlob("file_being_tracked", file);
+                console.log("file cleaned to Xet session", pf);
+                await xetSession.finalize();
 
                 const file_name = file.name;
                 const file_size = file.size;
 
                 const sha256 = pf.sha256;
                 console.log("result: " + sha256);
 
-                let response = await commit(endpoint, file_name, sha256, file_size, repo_type, repo_id, "main", hf_token);
+                let response = await commit(hf_endpoint, file_name, sha256, file_size, repo_type, repo_id, "main", hf_token);
                 document.getElementById('result').textContent = response;
             },
             false

hf_xet_wasm/examples/index.html

@@ -5,6 +5,8 @@ use futures::AsyncReadExt;
 use hf_xet_wasm::blob_reader::BlobReader;
 use hf_xet_wasm::configurations::{DataConfig, RepoSalt, ShardConfig, TranslatorConfig};
 use hf_xet_wasm::wasm_file_upload_session::FileUploadSession;
+use hf_xet_wasm::wasm_timer::Timer;
+use log::Level;
 use tokio::sync::mpsc;
 use utils::auth::AuthConfig;
 use wasm_bindgen::prelude::*;
@@ -13,7 +15,7 @@ use wasm_thread as thread;
 fn main() {
     #[cfg(target_arch = "wasm32")]
     {
-        console_log::init().unwrap();
+        console_log::init_with_level(Level::Info).unwrap();
         console_error_panic_hook::set_once();
     }
 
@@ -142,7 +144,10 @@ pub async fn test_async_blob_reader(file: web_sys::File) -> String {
 pub async fn clean_file(file: web_sys::File, endpoint: String, jwt_token: String, expiration: u64) -> String {
     log::debug!("clean_file called with {file:?}, {endpoint}, {jwt_token}, {expiration}");
 
+    let _timer = Timer::new_enforce_report("clean file main");
+
     let filename = file.name();
+    let filesize = file.size();
 
     let Ok(blob) = file.slice() else {
         log::error!("failed to convert a file to blob");
@@ -169,16 +174,19 @@ pub async fn clean_file(file: web_sys::File, endpoint: String, jwt_token: String
 
     let upload_session = Arc::new(FileUploadSession::new(Arc::new(config)));
 
-    let mut handle = upload_session.start_clean(filename);
+    let mut handle = upload_session.start_clean(filename, None);
 
     const READ_BUF_SIZE: usize = 8 * 1024 * 1024;
     let mut buf = vec![0u8; READ_BUF_SIZE];
     let mut total_read = 0;
+    let mut last_report = 0.;
     loop {
+        let _timer = Timer::new(format!("read file at {total_read}"));
         let Ok(bytes) = reader.read(&mut buf).await else {
             log::error!("failed to read from reader");
             return "".to_owned();
         };
+        drop(_timer);
         if bytes == 0 {
             break;
         }
@@ -192,7 +200,13 @@ pub async fn clean_file(file: web_sys::File, endpoint: String, jwt_token: String
             return "".to_owned();
         };
 
-        log::debug!("processed {total_read} bytes");
+        log::debug!("read {total_read} bytes");
+
+        let percentage = total_read as f64 / filesize * 100.;
+        if (percentage - last_report) > 10. {
+            log::info!("processing {percentage:.2}% of file");
+            last_report = percentage;
+        }
     }
     let Ok((file_hash, sha256, _metrics)) = handle.finish().await else {
         log::error!("failed to finish cleaner");

hf_xet_wasm/examples/simple.rs

@@ -27,16 +27,17 @@ function xetMetadataOrNone(jsonData) {
 }
 
 async function fetchXetMetadataFromRepoInfo({
-																							endpoint,
-																							tokenType,
-																							repoId,
-																							repoType,
-																							headers,
-																							params = null,
-																						}) {
+	hfEndpoint,
+	tokenType,
+	repoId,
+	repoType,
+	headers,
+	params = null
+}) {
 	/**
 	 * Uses the repo info to request a XET access token from Hub.
 	 *
+	 * @param {string} hfEndpoint - The HF Hub endpoint.
 	 * @param {string} tokenType - Type of the token to request: "read" or "write".
 	 * @param {string} repoId - A namespace (user or an organization) and a repo name separated by a `/`.
 	 * @param {string} repoType - Type of the repo to upload to: "model", "dataset", or "space".
@@ -46,7 +47,7 @@ async function fetchXetMetadataFromRepoInfo({
 	 * @throws {Error} If the Hub API returned an error or the response is improperly formatted.
 	 */
 
-	const url = `${endpoint}/api/${repoType}s/${repoId}/xet-${tokenType}-token/main`;
+	const url = `${hfEndpoint}/api/${repoType}s/${repoId}/xet-${tokenType}-token/main`;
 	console.log(`${url}`);
 
 	return fetchXetMetadataWithUrl(url, headers, params);

hf_xet_wasm/examples/xet_meta.js

@@ -10,7 +10,7 @@ mod sha256;
 mod wasm_deduplication_interface;
 mod wasm_file_cleaner;
 pub mod wasm_file_upload_session;
-
+pub mod wasm_timer;
 mod xorb_uploader;
 
 pub use session::XetSession;

hf_xet_wasm/src/lib.rs

@@ -67,7 +67,7 @@ impl XetSession {
     #[wasm_bindgen(js_name = "uploadFileFromBlob")]
     pub async fn upload_file_from_blob(&mut self, tracker_id: String, blob: Blob) -> Result<JsValue, JsValue> {
         // read from blob async
-        let mut cleaner = self.upload.start_clean(tracker_id);
+        let mut cleaner = self.upload.start_clean(tracker_id, None);
 
         let mut reader = BlobReader::new(blob)?;
 

hf_xet_wasm/src/session.rs

@@ -1,11 +1,44 @@
-use std::sync::Arc;
-
 use deduplication::Chunk;
 use merklehash::MerkleHash;
 use sha2::{Digest, Sha256};
 
 use super::errors::*;
 
+pub enum ShaGeneration {
+    Value(MerkleHash),
+    Action(ShaGenerator),
+}
+
+impl ShaGeneration {
+    pub fn new(hash: Option<MerkleHash>) -> Self {
+        match hash {
+            Some(h) => Self::Value(h),
+            None => Self::Action(ShaGenerator::new()),
+        }
+    }
+
+    pub fn update(&mut self, new_chunks: &[Chunk]) {
+        match self {
+            ShaGeneration::Value(_) => {},
+            ShaGeneration::Action(sha_generator) => sha_generator.update(new_chunks),
+        }
+    }
+
+    pub fn update_with_bytes(&mut self, new_bytes: &[u8]) {
+        match self {
+            ShaGeneration::Value(_) => {},
+            ShaGeneration::Action(sha_generator) => sha_generator.update_with_bytes(new_bytes),
+        }
+    }
+
+    pub fn finalize(self) -> Result<MerkleHash> {
+        match self {
+            ShaGeneration::Value(hash) => Ok(hash),
+            ShaGeneration::Action(sha_generator) => sha_generator.finalize(),
+        }
+    }
+}
+
 pub struct ShaGenerator {
     hasher: Sha256,
 }
@@ -17,13 +50,17 @@ impl ShaGenerator {
         }
     }
 
-    pub async fn update(&mut self, new_chunks: Arc<[Chunk]>) {
+    pub fn update(&mut self, new_chunks: &[Chunk]) {
         for chunk in new_chunks.iter() {
             self.hasher.update(&chunk.data);
         }
     }
 
-    pub async fn finalize(self) -> Result<MerkleHash> {
+    pub fn update_with_bytes(&mut self, new_bytes: &[u8]) {
+        self.hasher.update(new_bytes);
+    }
+
+    pub fn finalize(self) -> Result<MerkleHash> {
         let sha256 = self.hasher.finalize();
         let hex_str = format!("{sha256:x}");
         Ok(MerkleHash::from_hex(&hex_str)?)