Collect telemetry for upload path

Get the telemetry data to the py layer. This will be sent over the
telemetry endpoint later.

Related to STO-7

Author: Ajit Banerjee

.gitignore

@@ -17,3 +17,4 @@ debug/
 .vscode
 venv
 **/*.env
+**/uv.lock

data/src/clean.rs

@@ -25,6 +25,7 @@ use std::mem::take;
 use std::ops::DerefMut;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
+use std::time::Instant;
 use tokio::sync::mpsc::error::TryRecvError;
 use tokio::sync::mpsc::{channel, Receiver, Sender};
 use tokio::sync::Mutex;
@@ -80,6 +81,9 @@ pub struct Cleaner {
 
     // Auxiliary info
     file_name: Option<PathBuf>,
+
+    // Telemetry
+    start: Instant,
 }
 
 impl Cleaner {
@@ -117,6 +121,7 @@ impl Cleaner {
             tracking_info: Mutex::new(Default::default()),
             small_file_buffer: Mutex::new(Some(Vec::with_capacity(small_file_threshold))),
             file_name: file_name.map(|f| f.to_owned()),
+            start: Instant::now(),
         });
 
         Self::run(cleaner.clone(), chunk_c).await;
@@ -242,8 +247,9 @@ impl Cleaner {
         Ok(false)
     }
 
-    async fn dedup(&self, chunks: &[ChunkYieldType]) -> Result<()> {
+    async fn dedup(&self, chunks: &[ChunkYieldType]) -> Result<u64> {
         info!("Dedup {} chunks", chunks.len());
+        let mut total_compressed_bytes = 0;
         let mut tracking_info = self.tracking_info.lock().await;
 
         let enable_global_dedup = self.enable_global_dedup_queries;
@@ -472,13 +478,14 @@ impl Cleaner {
                     tracking_info.cas_data.data.extend(bytes);
 
                     if tracking_info.cas_data.data.len() > TARGET_CAS_BLOCK_SIZE {
-                        let cas_hash = register_new_cas_block(
+                        let (cas_hash, compressed_bytes) = register_new_cas_block(
                             &mut tracking_info.cas_data,
                             &self.shard_manager,
                             &self.cas,
                             &self.cas_prefix,
                         )
                         .await?;
+                        total_compressed_bytes += compressed_bytes;
 
                         for i in take(&mut tracking_info.current_cas_file_info_indices) {
                             tracking_info.file_info[i].cas_hash = cas_hash;
@@ -492,7 +499,7 @@ impl Cleaner {
             }
         }
 
-        Ok(())
+        Ok(total_compressed_bytes)
     }
 
     async fn finish(&self) -> Result<()> {
@@ -525,7 +532,8 @@ impl Cleaner {
         Ok(())
     }
 
-    async fn summarize_dedup_info(&self) -> Result<(MerkleHash, u64)> {
+    async fn summarize_dedup_info(&self) -> Result<(MerkleHash, u64, u64)> {
+        let mut total_compressed_bytes = 0;
         let mut tracking_info = self.tracking_info.lock().await;
 
         let file_hash = file_node_hash(
@@ -585,13 +593,14 @@ impl Cleaner {
             if cas_data_accumulator.data.len() >= TARGET_CAS_BLOCK_SIZE {
                 let mut new_cas_data = take(cas_data_accumulator.deref_mut());
                 drop(cas_data_accumulator); // Release the lock.
-                register_new_cas_block(
+                let (_cas_hash, compressed_bytes) = register_new_cas_block(
                     &mut new_cas_data,
                     &self.shard_manager,
                     &self.cas,
                     &self.cas_prefix,
                 )
                 .await?;
+                total_compressed_bytes += compressed_bytes;
             } else {
                 drop(cas_data_accumulator);
             }
@@ -601,11 +610,11 @@ impl Cleaner {
 
         *tracking_info = Default::default();
 
-        Ok((file_hash, file_size))
+        Ok((file_hash, file_size, total_compressed_bytes))
     }
 
     async fn to_pointer_file(&self) -> Result<String> {
-        let (hash, filesize) = self.summarize_dedup_info().await?;
+        let (hash, filesize, compressed_size) = self.summarize_dedup_info().await?;
         let pointer_file = PointerFile::init_from_info(
             &self
                 .file_name
@@ -614,6 +623,8 @@ impl Cleaner {
                 .unwrap_or_default(),
             &hash.hex(),
             filesize,
+            compressed_size,
+            self.start.elapsed(),
         );
         Ok(pointer_file.to_string())
     }

data/src/data_processing.rs

@@ -19,6 +19,7 @@ use std::mem::take;
 use std::ops::DerefMut;
 use std::path::Path;
 use std::sync::Arc;
+use std::time::Instant;
 use tokio::sync::Mutex;
 use tracing::error;
 
@@ -59,6 +60,9 @@ pub struct PointerFileTranslator {
 
     /* ----- Deduped data shared across files ----- */
     global_cas_data: Arc<Mutex<CASDataAggregator>>,
+    // Telemetry
+    /* ----- Telemetry ----- */
+    pub start: Instant,
 }
 
 // Constructors
@@ -93,6 +97,7 @@ impl PointerFileTranslator {
             remote_shards,
             cas: cas_client,
             global_cas_data: Default::default(),
+            start: Instant::now(),
         })
     }
 }
@@ -207,7 +212,7 @@ pub(crate) async fn register_new_cas_block(
     shard_manager: &Arc<ShardFileManager>,
     cas: &Arc<dyn Staging + Send + Sync>,
     cas_prefix: &str,
-) -> Result<MerkleHash> {
+) -> Result<(MerkleHash, u64)> {
     let cas_hash = cas_node_hash(&cas_data.chunks[..]);
 
     let raw_bytes_len = cas_data.data.len();
@@ -282,7 +287,7 @@ pub(crate) async fn register_new_cas_block(
     cas_data.chunks.clear();
     cas_data.pending_file_info.clear();
 
-    Ok(cas_hash)
+    Ok((cas_hash, compressed_bytes_len as u64))
 }
 
 /// Smudge operations

data/src/pointer_file.rs

@@ -3,6 +3,7 @@ use crate::constants::POINTER_FILE_LIMIT;
 use merklehash::{DataHashHexParseError, MerkleHash};
 use static_assertions::const_assert;
 use std::{collections::BTreeMap, fs, path::Path};
+use std::time::Duration;
 use toml::Value;
 use tracing::{debug, error, warn};
 
@@ -33,6 +34,12 @@ pub struct PointerFile {
 
     /// The size of the file pointed to by this pointer file
     filesize: u64,
+
+    /// The addition to CAS bytes due to this file
+    compressed_size: Option<u64>,
+
+    /// The duration for cleaning or smudging this file
+    pub latency: Option<Duration>,
 }
 
 impl PointerFile {
@@ -59,6 +66,8 @@ impl PointerFile {
                 is_valid,
                 hash,
                 filesize,
+                compressed_size: None,
+                latency: None,
             };
         }
 
@@ -73,6 +82,8 @@ impl PointerFile {
                 is_valid,
                 hash,
                 filesize,
+                compressed_size: None,
+                latency: None,
             };
         }
 
@@ -117,6 +128,8 @@ impl PointerFile {
             is_valid,
             hash,
             filesize,
+            compressed_size: None,
+            latency: None,
         }
     }
 
@@ -134,6 +147,8 @@ impl PointerFile {
             is_valid: false,
             hash: empty_string,
             filesize: 0,
+            compressed_size: None,
+            latency: None,
         };
 
         let Ok(file_meta) = fs::metadata(path).map_err(|e| {
@@ -156,13 +171,15 @@ impl PointerFile {
         PointerFile::init_from_string(&contents, path)
     }
 
-    pub fn init_from_info(path: &str, hash: &str, filesize: u64) -> Self {
+    pub fn init_from_info(path: &str, hash: &str, filesize: u64, compressed_size: u64, latency: Duration) -> Self {
         Self {
             version_string: CURRENT_VERSION.to_string(),
             path: path.to_string(),
             is_valid: true,
             hash: hash.to_string(),
             filesize,
+            compressed_size: Some(compressed_size),
+            latency: Some(latency),
         }
     }
 
@@ -194,6 +211,10 @@ impl PointerFile {
     pub fn filesize(&self) -> u64 {
         self.filesize
     }
+
+    pub fn compressed_size(&self) -> u64 {self.compressed_size.unwrap_or(0)}
+
+    pub fn latency(&self) -> f64 {self.latency.map(|dur| dur.as_secs_f64()).unwrap_or(0f64)}
 }
 
 pub fn is_xet_pointer_file(data: &[u8]) -> bool {
@@ -322,4 +343,6 @@ mod tests {
         let test = PointerFile::init_from_string(&test_contents, &empty_string);
         assert!(!test.is_valid()); // new version is not valid
     }
+
+    // todo add init_from_info test
 }

hf_xet/pyproject.toml

@@ -11,6 +11,9 @@ classifiers = [
     "Programming Language :: Python :: Implementation :: PyPy",
 ]
 dynamic = ["version"]
+dependencies = [
+    "maturin>=1.7.4",
+]
 [project.optional-dependencies]
 tests = [
     "pytest",

hf_xet/src/data_client.rs

@@ -55,15 +55,15 @@ pub async fn download_async(
     endpoint: Option<String>,
     token_info: Option<(String, u64)>,
     token_refresher: Option<Arc<dyn TokenRefresher>>,
-) -> errors::Result<Vec<String>> {
+) -> errors::Result<Vec<PointerFile>> {
     let config = default_config(
         endpoint.unwrap_or(DEFAULT_CAS_ENDPOINT.to_string()),
         token_info,
         token_refresher,
     )?;
     let processor = Arc::new(PointerFileTranslator::new(config).await?);
     let processor = &processor;
-    let paths = tokio_par_for_each(
+    let pfs = tokio_par_for_each(
         pointer_files,
         MAX_CONCURRENT_DOWNLOADS,
         |pointer_file, _| async move {
@@ -77,7 +77,7 @@ pub async fn download_async(
         ParallelError::TaskError(e) => e,
     })?;
 
-    Ok(paths)
+    Ok(pfs)
 }
 
 async fn clean_file(processor: &PointerFileTranslator, f: String) -> errors::Result<PointerFile> {
@@ -95,7 +95,6 @@ async fn clean_file(processor: &PointerFileTranslator, f: String) -> errors::Res
 
         handle.add_bytes(read_buf[0..bytes].to_vec()).await?;
     }
-
     let pf_str = handle.result().await?;
     let pf = PointerFile::init_from_string(&pf_str, path.to_str().unwrap());
     Ok(pf)
@@ -104,15 +103,17 @@ async fn clean_file(processor: &PointerFileTranslator, f: String) -> errors::Res
 async fn smudge_file(
     proc: &PointerFileTranslator,
     pointer_file: &PointerFile,
-) -> errors::Result<String> {
+) -> errors::Result<PointerFile> {
     let path = PathBuf::from(pointer_file.path());
     if let Some(parent_dir) = path.parent() {
         fs::create_dir_all(parent_dir)?;
     }
     let mut f = File::create(&path)?;
     proc.smudge_file_from_pointer(&pointer_file, &mut f, None)
         .await?;
-    Ok(pointer_file.path().to_string())
+    let mut pointer_file_clone = pointer_file.clone();
+    pointer_file_clone.latency = Some(proc.start.elapsed());
+    Ok(pointer_file_clone)
 }
 
 #[cfg(test)]

hf_xet/src/lib.rs

@@ -10,8 +10,32 @@ use pyo3::prelude::*;
 use pyo3::pyfunction;
 use std::fmt::Debug;
 use std::sync::Arc;
+use std::time::Duration;
 use token_refresh::WrappedTokenRefresher;
 
+// This will be the information that will finally be sent to the telemetry endpoint
+//
+// #[pyclass]
+// #[derive(Clone, Debug)]
+// pub struct UploadTelemetry {
+//     #[pyo3(get)]
+//     total_time_ms: u64,
+//     #[pyo3(get)]
+//     uploaded_bytes: u64,
+// }
+//
+// #[pyclass]
+// #[derive(Clone, Debug)]
+// pub struct DownloadTelemetry {
+//     #[pyo3(get)]
+//     total_time_ms: u64,
+//     #[pyo3(get)]
+//     downloaded_bytes: u64,
+//     #[pyo3(get)]
+//     cached_bytes: u32,
+// }
+
+
 #[pyfunction]
 #[pyo3(signature = (file_paths, endpoint, token_info, token_refresher), text_signature = "(file_paths: List[str], endpoint: Optional[str], token_info: Optional[(str, int)], token_refresher: Optional[Callable[[], (str, int)]]) -> List[PyPointerFile]")]
 pub fn upload_files(
@@ -63,8 +87,10 @@ pub fn download_files(
             .block_on(async move {
                 data_client::download_async(pfs, endpoint, token_info, refresher).await
             })
+            .map(|pfs| pfs.into_iter().map(|pointer_file| pointer_file.path().to_string()).collect())
             .map_err(|e| PyException::new_err(format!("{e:?}")))
     })
+
 }
 
 // helper to convert the implemented WrappedTokenRefresher into an Arc<dyn TokenRefresher>
@@ -82,21 +108,28 @@ pub struct PyPointerFile {
     hash: String,
     #[pyo3(get)]
     filesize: u64,
+    #[pyo3(get)]
+    compressed_size: u64,
+    #[pyo3(get)]
+    latency: f64,
 }
 
+
 impl From<PointerFile> for PyPointerFile {
     fn from(pf: PointerFile) -> Self {
         Self {
             path: pf.path().to_string(),
             hash: pf.hash_string().to_string(),
             filesize: pf.filesize(),
+            compressed_size: pf.compressed_size(),
+            latency: pf.latency(),
         }
     }
 }
 
 impl From<PyPointerFile> for PointerFile {
     fn from(pf: PyPointerFile) -> Self {
-        PointerFile::init_from_info(&pf.path, &pf.hash, pf.filesize)
+        PointerFile::init_from_info(&pf.path, &pf.hash, pf.filesize, pf.compressed_size, Duration::from_secs_f64(pf.latency))
     }
 }
 
@@ -108,6 +141,8 @@ impl PyPointerFile {
             path,
             hash,
             filesize,
+            compressed_size: 0,
+            latency: 0.0,
         }
     }