refactor(hf): disable progress reporting and total_bytes optional during progress tracking

kszucs · kszucs · commit 6e36baa851f7 · 2026-02-13T10:14:35.000+01:00
diff --git a/data/src/bin/example.rs b/data/src/bin/example.rs
@@ -91,7 +91,7 @@ async fn clean(mut reader: impl Read, mut writer: impl Write, size: u64) -> Resu
         FileUploadSession::new(TranslatorConfig::local_config(std::env::current_dir()?)?.into(), None).await?;
 
     let mut size_read = 0;
-    let mut handle = translator.start_clean(None, size, None).await;
+    let mut handle = translator.start_clean(None, Some(size), None).await;
 
     loop {
         let bytes = reader.read(&mut read_buf)?;
diff --git a/data/src/data_client.rs b/data/src/data_client.rs
@@ -249,7 +249,7 @@ pub async fn clean_bytes(
     bytes: Vec<u8>,
     name: Option<Arc<str>>,
 ) -> errors::Result<(XetFileInfo, DeduplicationMetrics)> {
-    let mut handle = processor.start_clean(name, bytes.len() as u64, None).await;
+    let mut handle = processor.start_clean(name, Some(bytes.len() as u64), None).await;
     handle.add_data(&bytes).await?;
     handle.finish().await
 }
@@ -270,7 +270,11 @@ pub async fn clean_file(
     let mut buffer = vec![0u8; u64::min(filesize, *xet_config().data.ingestion_block_size) as usize];
 
     let mut handle = processor
-        .start_clean(Some(filename.as_ref().to_string_lossy().into()), filesize, Sha256::from_hex(sha256.as_ref()).ok())
+        .start_clean(
+            Some(filename.as_ref().to_string_lossy().into()),
+            Some(filesize),
+            Sha256::from_hex(sha256.as_ref()).ok(),
+        )
         .await;
 
     loop {
diff --git a/data/src/file_upload_session.rs b/data/src/file_upload_session.rs
@@ -162,7 +162,10 @@ impl FileUploadSession {
 
             // Get a new file id for the completion tracking.  This also registers the size against the total bytes
             // of the file.
-            let file_id = self.completion_tracker.register_new_file(file_name.clone(), file_size).await;
+            let file_id = self
+                .completion_tracker
+                .register_new_file(file_name.clone(), Some(file_size))
+                .await;
 
             // Now, spawn a task
             let ingestion_concurrency_limiter =
@@ -260,7 +263,7 @@ impl FileUploadSession {
     pub async fn start_clean(
         self: &Arc<Self>,
         file_name: Option<Arc<str>>,
-        size: u64,
+        size: Option<u64>,
         sha256: Option<Sha256>,
     ) -> SingleFileCleaner {
         // Get a new file id for the completion tracking
@@ -571,7 +574,7 @@ mod tests {
             .unwrap();
 
         let mut cleaner = upload_session
-            .start_clean(Some("test".into()), read_data.len() as u64, None)
+            .start_clean(Some("test".into()), Some(read_data.len() as u64), None)
             .await;
 
         // Read blocks from the source file and hand them to the cleaning handle
diff --git a/data/src/streaming.rs b/data/src/streaming.rs
@@ -66,13 +66,8 @@ impl XetClient {
     ///
     /// Each call creates a fresh [`FileUploadSession`] because sessions are
     /// consumed on finalization and cannot be reused.
-    pub async fn write(
-        &self,
-        progress_updater: Option<Arc<dyn TrackingProgressUpdater>>,
-        name: Option<Arc<str>>,
-        size: u64,
-    ) -> errors::Result<XetWriter> {
-        XetWriter::new(self.config.clone(), progress_updater, name, size).await
+    pub async fn write(&self, name: Option<Arc<str>>) -> errors::Result<XetWriter> {
+        XetWriter::new(self.config.clone(), name).await
     }
 }
 
@@ -89,17 +84,9 @@ pub struct XetWriter {
 
 impl XetWriter {
     /// Creates a new writer that will upload a single file.
-    ///
-    /// `size` is the total number of bytes that will be written and is used for
-    /// progress tracking. The caller must know the content length before writing.
-    pub async fn new(
-        config: Arc<TranslatorConfig>,
-        progress_updater: Option<Arc<dyn TrackingProgressUpdater>>,
-        name: Option<Arc<str>>,
-        size: u64,
-    ) -> errors::Result<Self> {
-        let session = FileUploadSession::new(config, progress_updater).await?;
-        let handle = session.start_clean(name, size, None).await;
+    pub async fn new(config: Arc<TranslatorConfig>, name: Option<Arc<str>>) -> errors::Result<Self> {
+        let session = FileUploadSession::new(config, None).await?;
+        let handle = session.start_clean(name, None, None).await;
         Ok(Self {
             session: Some(session),
             handle: Some(handle),
@@ -292,7 +279,7 @@ mod tests {
     /// Upload `content` via [`XetWriter`] and download it back via [`XetReader`],
     /// asserting the round-tripped bytes match the original.
     async fn assert_roundtrip(client: &XetClient, content: &[u8]) {
-        let mut writer = client.write(None, None, content.len() as u64).await.unwrap();
+        let mut writer = client.write(None).await.unwrap();
         for chunk in content.chunks(4096) {
             writer.write(Bytes::copy_from_slice(chunk)).await.unwrap();
         }
@@ -341,7 +328,7 @@ mod tests {
         let client = XetClient::new(Some(endpoint), None, None, "test".into()).unwrap();
 
         let content: Vec<u8> = (0..1_000_000).map(|i| (i % 256) as u8).collect();
-        let mut writer = client.write(None, None, content.len() as u64).await.unwrap();
+        let mut writer = client.write(None).await.unwrap();
         for chunk in content.chunks(4096) {
             writer.write(Bytes::copy_from_slice(chunk)).await.unwrap();
         }
@@ -359,7 +346,7 @@ mod tests {
         let endpoint = format!("local://{}", temp_dir.path().display());
         let client = XetClient::new(Some(endpoint), None, None, "test".into()).unwrap();
 
-        let mut writer = client.write(None, None, 5).await.unwrap();
+        let mut writer = client.write(None).await.unwrap();
         writer.write(Bytes::from_static(b"hello")).await.unwrap();
         let file_info = writer.close().await.unwrap();
 
@@ -416,7 +403,7 @@ mod tests {
         let endpoint = format!("local://{}", temp_dir.path().display());
         let client = XetClient::new(Some(endpoint), None, None, "test".into()).unwrap();
 
-        let mut writer = client.write(None, None, 100).await.unwrap();
+        let mut writer = client.write(None).await.unwrap();
         writer.write(Bytes::from_static(b"some data")).await.unwrap();
         writer.abort().await.unwrap();
 
@@ -431,7 +418,7 @@ mod tests {
         let client = XetClient::new(Some(endpoint), None, None, "test".into()).unwrap();
 
         let content = b"Hello, World!";
-        let mut writer = client.write(None, None, Some(content.len() as u64)).await.unwrap();
+        let mut writer = client.write(None).await.unwrap();
         writer.write(Bytes::from_static(content)).await.unwrap();
         let file_info = writer.close().await.unwrap();
 
diff --git a/data/tests/test_session_resume.rs b/data/tests/test_session_resume.rs
@@ -67,7 +67,7 @@ mod tests {
 
             // Feed it half the data, and checkpoint.
             let mut cleaner = file_upload_session
-                .start_clean(Some("data".into()), data.len() as u64, None)
+                .start_clean(Some("data".into()), Some(data.len() as u64), None)
                 .await;
             cleaner.add_data(&data[..half_n]).await.unwrap();
             cleaner.checkpoint().await.unwrap();
@@ -85,7 +85,7 @@ mod tests {
 
             // Feed it half the data, and checkpoint.
             let mut cleaner = file_upload_session
-                .start_clean(Some("data".into()), data.len() as u64, None)
+                .start_clean(Some("data".into()), Some(data.len() as u64), None)
                 .await;
 
             // Add all the data.  Roughly the first half should dedup.
@@ -140,7 +140,7 @@ mod tests {
 
             // Feed it half the data, and checkpoint.
             let mut cleaner = file_upload_session
-                .start_clean(Some("data".into()), data.len() as u64, None)
+                .start_clean(Some("data".into()), Some(data.len() as u64), None)
                 .await;
             cleaner.add_data(&data[..rn]).await.unwrap();
             cleaner.checkpoint().await.unwrap();
@@ -172,7 +172,7 @@ mod tests {
 
             // Feed it half the data, and checkpoint.
             let mut cleaner = file_upload_session
-                .start_clean(Some("data".into()), data.len() as u64, None)
+                .start_clean(Some("data".into()), Some(data.len() as u64), None)
                 .await;
 
             // Add all the data.  Roughly the first half should dedup.
diff --git a/hf_xet/src/progress_update.rs b/hf_xet/src/progress_update.rs
@@ -254,7 +254,7 @@ impl WrappedProgressUpdaterImpl {
                                 py,
                                 PyItemProgressUpdate {
                                     item_name: PyString::new(py, &u.item_name).into(),
-                                    total_bytes: u.total_bytes,
+                                    total_bytes: u.total_bytes.unwrap_or(0),
                                     bytes_completed: u.bytes_completed,
                                     bytes_completion_increment: u.bytes_completion_increment,
                                 },
diff --git a/progress_tracking/src/aggregator.rs b/progress_tracking/src/aggregator.rs
@@ -276,7 +276,7 @@ mod tests {
             .register_updates(ProgressUpdate {
                 item_updates: vec![ItemProgressUpdate {
                     item_name: Arc::from("fileA.txt"),
-                    total_bytes: 100,
+                    total_bytes: Some(100),
                     bytes_completed: 10,
                     bytes_completion_increment: 10,
                 }],
@@ -299,7 +299,7 @@ mod tests {
             .register_updates(ProgressUpdate {
                 item_updates: vec![ItemProgressUpdate {
                     item_name: Arc::from("fileB.txt"),
-                    total_bytes: 200,
+                    total_bytes: Some(200),
                     bytes_completed: 50,
                     bytes_completion_increment: 50,
                 }],
@@ -323,13 +323,13 @@ mod tests {
                 item_updates: vec![
                     ItemProgressUpdate {
                         item_name: Arc::from("fileC.txt"),
-                        total_bytes: 300,
+                        total_bytes: Some(300),
                         bytes_completed: 90,
                         bytes_completion_increment: 90,
                     },
                     ItemProgressUpdate {
                         item_name: Arc::from("fileA.txt"),
-                        total_bytes: 100,
+                        total_bytes: Some(100),
                         bytes_completed: 30,
                         bytes_completion_increment: 20,
                     },
@@ -368,19 +368,19 @@ mod tests {
 
         let a = &flushed.item_updates[0];
         assert_eq!(a.item_name.as_ref(), "fileA.txt");
-        assert_eq!(a.total_bytes, 100);
+        assert_eq!(a.total_bytes, Some(100));
         assert_eq!(a.bytes_completed, 30);
         assert_eq!(a.bytes_completion_increment, 30);
 
         let b = &flushed.item_updates[1];
         assert_eq!(b.item_name.as_ref(), "fileB.txt");
-        assert_eq!(b.total_bytes, 200);
+        assert_eq!(b.total_bytes, Some(200));
         assert_eq!(b.bytes_completed, 50);
         assert_eq!(b.bytes_completion_increment, 50);
 
         let c = &flushed.item_updates[2];
         assert_eq!(c.item_name.as_ref(), "fileC.txt");
-        assert_eq!(c.total_bytes, 300);
+        assert_eq!(c.total_bytes, Some(300));
         assert_eq!(c.bytes_completed, 90);
         assert_eq!(c.bytes_completion_increment, 90);
     }
diff --git a/progress_tracking/src/item_tracking.rs b/progress_tracking/src/item_tracking.rs
@@ -109,7 +109,7 @@ impl SingleItemProgressUpdater {
             .clone()
             .do_item_update(ItemProgressUpdate {
                 item_name: self.item_name.clone(),
-                total_bytes: self.n_bytes.load(Ordering::Relaxed),
+                total_bytes: Some(self.n_bytes.load(Ordering::Relaxed)),
                 bytes_completed: old_completed_count + increment,
                 bytes_completion_increment: increment,
             })
diff --git a/progress_tracking/src/progress_info.rs b/progress_tracking/src/progress_info.rs
@@ -5,7 +5,7 @@ use std::sync::Arc;
 #[derive(Clone, Debug)]
 pub struct ItemProgressUpdate {
     pub item_name: Arc<str>,
-    pub total_bytes: u64,
+    pub total_bytes: Option<u64>,
 
     // Bytes completed are the total bytes completed, either through
     // deduplication, upload/download, loading from cache, etc.
@@ -19,7 +19,10 @@ impl ItemProgressUpdate {
 
         // Just in case the total got updated, as can be the case when we don't know the
         // size ahead of time.
-        self.total_bytes = self.total_bytes.max(other.total_bytes);
+        self.total_bytes = match (self.total_bytes, other.total_bytes) {
+            (Some(a), Some(b)) => Some(a.max(b)),
+            (a, b) => a.or(b),
+        };
         self.bytes_completed = self.bytes_completed.max(other.bytes_completed);
         self.bytes_completion_increment += other.bytes_completion_increment;
     }
diff --git a/progress_tracking/src/upload_tracking.rs b/progress_tracking/src/upload_tracking.rs
diff --git a/progress_tracking/src/verification_wrapper.rs b/progress_tracking/src/verification_wrapper.rs
