xet_data/processing/
data_client.rs

1use std::sync::Arc;
2#[cfg(not(target_family = "wasm"))]
3use std::{fs::File, io::Read, path::Path};
4
5#[cfg(not(target_family = "wasm"))]
6use bytes::Bytes;
7use http::header::HeaderMap;
8use tracing::instrument;
9#[cfg(not(target_family = "wasm"))]
10use tracing::{Instrument, Span, info_span};
11use uuid::Uuid;
12use xet_client::cas_client::auth::{AuthConfig, TokenRefresher};
13#[cfg(not(target_family = "wasm"))]
14use xet_core_structures::merklehash::MerkleHash;
15use xet_runtime::core::XetContext;
16#[cfg(not(target_family = "wasm"))]
17use xet_runtime::core::par_utils::run_constrained_with_semaphore;
18
19use super::configurations::{SessionContext, TranslatorConfig};
20use super::file_cleaner::Sha256Policy;
21use super::{FileUploadSession, XetFileInfo};
22#[cfg(not(target_family = "wasm"))]
23use crate::deduplication::Chunker;
24use crate::deduplication::DeduplicationMetrics;
25use crate::error::Result;
26
27pub fn default_config(
28    ctx: &XetContext,
29    endpoint: String,
30    token_info: Option<(String, u64)>,
31    token_refresher: Option<Arc<dyn TokenRefresher>>,
32    custom_headers: Option<Arc<HeaderMap>>,
33) -> Result<TranslatorConfig> {
34    let (token, token_expiration) = token_info.unzip();
35    let auth_cfg = AuthConfig::maybe_new(token, token_expiration, token_refresher);
36
37    let session = SessionContext {
38        endpoint,
39        auth: auth_cfg,
40        custom_headers,
41        repo_paths: vec!["".into()],
42        session_id: Some(Uuid::now_v7().to_string()),
43    };
44
45    TranslatorConfig::new(ctx, session)
46}
47
48#[instrument(skip_all, name = "clean_bytes", fields(bytes.len = bytes.len()))]
49pub async fn clean_bytes(
50    processor: Arc<FileUploadSession>,
51    bytes: Vec<u8>,
52    sha256_policy: Sha256Policy,
53) -> Result<(XetFileInfo, DeduplicationMetrics)> {
54    let (_id, mut handle) = processor.start_clean(None, Some(bytes.len() as u64), sha256_policy)?;
55    handle.add_data(&bytes).await?;
56    let (info, metrics) = handle.finish().await?;
57    Ok((info, metrics))
58}
59
60#[cfg(not(target_family = "wasm"))]
61#[instrument(skip_all, name = "clean_file", fields(file.name = tracing::field::Empty, file.len = tracing::field::Empty))]
62pub async fn clean_file(
63    processor: Arc<FileUploadSession>,
64    filename: impl AsRef<Path>,
65    sha256_policy: Sha256Policy,
66) -> Result<(XetFileInfo, DeduplicationMetrics)> {
67    let mut reader = File::open(&filename)?;
68
69    let filesize = reader.metadata()?.len();
70    let span = Span::current();
71    span.record("file.name", filename.as_ref().to_str());
72    span.record("file.len", filesize);
73    let mut buffer = vec![0u8; u64::min(filesize, *processor.ctx.config.data.ingestion_block_size) as usize];
74
75    let (_id, mut handle) =
76        processor.start_clean(Some(filename.as_ref().to_string_lossy().into()), Some(filesize), sha256_policy)?;
77
78    loop {
79        let bytes = reader.read(&mut buffer)?;
80        if bytes == 0 {
81            break;
82        }
83
84        handle.add_data(&buffer[0..bytes]).await?;
85    }
86
87    let (info, metrics) = handle.finish().await?;
88    Ok((info, metrics))
89}
90
91/// Computes the xet hash for a single file without uploading.
92///
93/// This function performs local-only hash computation by reading the file,
94/// chunking it using content-defined chunking, and computing the aggregated
95/// hash from the chunk hashes. The resulting hash is identical to what would
96/// be returned by upload operations, enabling verification of downloaded files.
97///
98/// # Arguments
99/// * `filename` - Path to the file to hash
100/// * `buffer_size` - Size of the read buffer in bytes
101///
102/// # Returns
103/// * `XetFileInfo` containing the hex-encoded hash and file size
104///
105/// # Errors
106/// * `IoError` if the file cannot be opened or read
107///
108/// # Use Cases
109/// - Verify that downloaded files are correctly reassembled
110/// - Check if a file needs to be uploaded (by comparing hashes)
111/// - Generate cache keys for local file operations
112#[cfg(not(target_family = "wasm"))]
113fn hash_single_file(ctx: XetContext, filename: String, buffer_size: usize) -> Result<XetFileInfo> {
114    let mut reader = File::open(&filename)?;
115    let filesize = reader.metadata()?.len();
116
117    let mut buffer = vec![0u8; buffer_size];
118    let mut chunker = Chunker::default();
119    let mut chunk_hashes: Vec<(MerkleHash, u64)> = Vec::new();
120
121    loop {
122        ctx.check_sigint_shutdown()?;
123
124        let bytes_read = reader.read(&mut buffer)?;
125        if bytes_read == 0 {
126            break;
127        }
128
129        let data = Bytes::copy_from_slice(&buffer[0..bytes_read]);
130        let chunks = chunker.next_block_bytes(&data, false);
131
132        for chunk in chunks {
133            chunk_hashes.push((chunk.hash, chunk.data.len() as u64));
134        }
135    }
136
137    // Get the final chunk if any data remains in the chunker
138    if let Some(final_chunk) = chunker.finish() {
139        chunk_hashes.push((final_chunk.hash, final_chunk.data.len() as u64));
140    }
141
142    let file_hash = xet_core_structures::merklehash::file_hash(&chunk_hashes);
143    Ok(XetFileInfo::new(file_hash.hex(), filesize))
144}
145
146/// Computes xet hashes for multiple files in parallel without uploading.
147///
148/// This function processes multiple files concurrently using a semaphore to limit
149/// parallelism. Each file is hashed independently using `hash_single_file()`.
150/// The resulting hashes are identical to those from upload operations,
151/// enabling validation and verification of file transfers.
152///
153/// # Arguments
154/// * `file_paths` - Vector of file paths to hash
155///
156/// # Returns
157/// * Vector of `XetFileInfo` in the same order as input file paths
158///
159/// # Errors
160/// * Returns error if any file cannot be read or hashed
161///
162/// # Use Cases
163/// - Verify integrity of downloaded files by comparing computed hashes
164/// - Batch validation of multiple files after transfer
165/// - Determine which files need to be uploaded by comparing with server hashes
166///
167/// # Performance
168/// - Uses `file_ingestion_semaphore` to control parallelism
169/// - No authentication or server connection required
170/// - Pure local computation
171#[cfg(not(target_family = "wasm"))]
172#[instrument(skip_all, name = "data_client::hash_files", fields(num_files=file_paths.len()))]
173pub async fn hash_files_async(ctx: &XetContext, file_paths: Vec<String>) -> Result<Vec<XetFileInfo>> {
174    let runtime = ctx.runtime.clone();
175    let semaphore = ctx.common.file_ingestion_semaphore.clone();
176    let buffer_size = *ctx.config.data.ingestion_block_size as usize;
177
178    let hash_futures = file_paths.into_iter().map(|file_path| {
179        let runtime = runtime.clone();
180        let ctx = ctx.clone();
181        async move {
182            runtime
183                .spawn_blocking(move || hash_single_file(ctx, file_path, buffer_size))
184                .await
185                .map_err(|e| std::io::Error::other(e.to_string()))?
186        }
187        .instrument(info_span!("hash_file"))
188    });
189
190    let files = run_constrained_with_semaphore(hash_futures, semaphore).await?;
191
192    Ok(files)
193}
194
195#[cfg(test)]
196mod tests {
197    use dirs::home_dir;
198    use serial_test::serial;
199    use tempfile::tempdir;
200    use xet_runtime::core::XetContext;
201    use xet_runtime::utils::EnvVarGuard;
202
203    use super::*;
204
205    #[test]
206    #[serial(default_config_env)]
207    fn test_default_config_with_hf_home() {
208        let temp_dir = tempdir().unwrap();
209        let _hf_home_guard = EnvVarGuard::set("HF_HOME", temp_dir.path().to_str().unwrap());
210
211        let endpoint = "http://localhost:8080".to_string();
212        let ctx = XetContext::default().unwrap();
213        let result = default_config(&ctx, endpoint, None, None, None);
214
215        assert!(result.is_ok());
216        let config = result.unwrap();
217        assert!(config.shard_cache_directory.starts_with(temp_dir.path()));
218    }
219
220    #[test]
221    #[serial(default_config_env)]
222    fn test_default_config_with_hf_xet_cache_and_hf_home() {
223        let temp_dir_xet_cache = tempdir().unwrap();
224        let temp_dir_hf_home = tempdir().unwrap();
225
226        let hf_xet_cache_guard = EnvVarGuard::set("HF_XET_CACHE", temp_dir_xet_cache.path().to_str().unwrap());
227        let hf_home_guard = EnvVarGuard::set("HF_HOME", temp_dir_hf_home.path().to_str().unwrap());
228
229        let endpoint = "http://localhost:8080".to_string();
230        let ctx = XetContext::default().unwrap();
231        let result = default_config(&ctx, endpoint, None, None, None);
232
233        assert!(result.is_ok());
234        let config = result.unwrap();
235        assert!(config.shard_cache_directory.starts_with(temp_dir_xet_cache.path()));
236
237        drop(hf_xet_cache_guard);
238        drop(hf_home_guard);
239
240        let temp_dir = tempdir().unwrap();
241        let _hf_home_guard = EnvVarGuard::set("HF_HOME", temp_dir.path().to_str().unwrap());
242
243        let endpoint = "http://localhost:8080".to_string();
244        let ctx = XetContext::default().unwrap();
245        let result = default_config(&ctx, endpoint, None, None, None);
246
247        assert!(result.is_ok());
248        let config = result.unwrap();
249        assert!(config.shard_cache_directory.starts_with(temp_dir.path()));
250    }
251
252    #[test]
253    #[serial(default_config_env)]
254    fn test_default_config_with_hf_xet_cache() {
255        let temp_dir = tempdir().unwrap();
256        let _hf_xet_cache_guard = EnvVarGuard::set("HF_XET_CACHE", temp_dir.path().to_str().unwrap());
257
258        let endpoint = "http://localhost:8080".to_string();
259        let ctx = XetContext::default().unwrap();
260        let result = default_config(&ctx, endpoint, None, None, None);
261
262        assert!(result.is_ok());
263        let config = result.unwrap();
264        assert!(config.shard_cache_directory.starts_with(temp_dir.path()));
265    }
266
267    #[test]
268    #[serial(default_config_env)]
269    fn test_default_config_without_env_vars() {
270        let endpoint = "http://localhost:8080".to_string();
271        let ctx = XetContext::default().unwrap();
272        let result = default_config(&ctx, endpoint, None, None, None);
273
274        let expected = home_dir().unwrap().join(".cache").join("huggingface").join("xet");
275
276        assert!(result.is_ok());
277        let config = result.unwrap();
278        let test_cache_dir = &config.shard_cache_directory;
279        assert!(
280            test_cache_dir.starts_with(&expected),
281            "cache dir = {test_cache_dir:?}; does not start with {expected:?}",
282        );
283    }
284
285    #[tokio::test]
286    async fn test_hash_empty_file() {
287        let temp_dir = tempdir().unwrap();
288        let file_path = temp_dir.path().join("empty.txt");
289        std::fs::write(&file_path, b"").unwrap();
290
291        let buffer_size = 8 * 1024 * 1024; // 8MB
292        let ctx = XetContext::default().unwrap();
293        let result = hash_single_file(ctx, file_path.to_str().unwrap().to_string(), buffer_size);
294        assert!(result.is_ok());
295
296        let file_info = result.unwrap();
297        assert_eq!(file_info.file_size(), Some(0));
298        assert!(!file_info.hash().is_empty());
299    }
300
301    #[tokio::test]
302    async fn test_hash_small_file() {
303        let temp_dir = tempdir().unwrap();
304        let file_path = temp_dir.path().join("small.txt");
305        let content = b"Hello, World!";
306        std::fs::write(&file_path, content).unwrap();
307
308        let buffer_size = 8 * 1024 * 1024; // 8MB
309        let ctx = XetContext::default().unwrap();
310        let result = hash_single_file(ctx, file_path.to_str().unwrap().to_string(), buffer_size);
311        assert!(result.is_ok());
312
313        let file_info = result.unwrap();
314        assert_eq!(file_info.file_size(), Some(content.len() as u64));
315        assert!(!file_info.hash().is_empty());
316    }
317
318    #[tokio::test]
319    #[cfg_attr(feature = "smoke-test", ignore)]
320    async fn test_hash_determinism() {
321        let temp_dir = tempdir().unwrap();
322        let file_path = temp_dir.path().join("test.txt");
323
324        // Create a file that is large enough to span multiple buffer reads
325        // Using 20MB to ensure it's larger than typical buffer sizes
326        let file_size = 20 * 1024 * 1024;
327        let content: Vec<u8> = (0..file_size).map(|i| (i % 256) as u8).collect();
328        std::fs::write(&file_path, &content).unwrap();
329
330        let file_path_str = file_path.to_str().unwrap().to_string();
331        let ctx = XetContext::default().unwrap();
332
333        // Hash with 8MB buffer size
334        let result1 = hash_single_file(ctx.clone(), file_path_str.clone(), 8 * 1024 * 1024);
335        assert!(result1.is_ok());
336        let file_info1 = result1.unwrap();
337
338        // Hash with 4MB buffer size
339        let result2 = hash_single_file(ctx.clone(), file_path_str, 4 * 1024 * 1024);
340        assert!(result2.is_ok());
341        let file_info2 = result2.unwrap();
342
343        // Hashes should be identical regardless of buffer size
344        // This verifies that chunker.finish() is called correctly
345        assert_eq!(file_info1.hash(), file_info2.hash());
346        assert_eq!(file_info1.file_size(), file_info2.file_size());
347    }
348
349    #[tokio::test]
350    async fn test_hash_file_not_found() {
351        let buffer_size = 8 * 1024 * 1024; // 8MB
352        let ctx = XetContext::default().unwrap();
353        let result = hash_single_file(ctx, "/nonexistent/file.txt".to_string(), buffer_size);
354        assert!(result.is_err());
355    }
356
357    #[tokio::test]
358    async fn test_hash_files_async() {
359        let temp_dir = tempdir().unwrap();
360
361        let file1_path = temp_dir.path().join("file1.txt");
362        let file2_path = temp_dir.path().join("file2.txt");
363
364        std::fs::write(&file1_path, b"First file content").unwrap();
365        std::fs::write(&file2_path, b"Second file content").unwrap();
366
367        let file_paths = vec![
368            file1_path.to_str().unwrap().to_string(),
369            file2_path.to_str().unwrap().to_string(),
370        ];
371
372        let ctx = XetContext::default().unwrap();
373        let result = hash_files_async(&ctx, file_paths).await;
374        assert!(result.is_ok());
375
376        let file_infos = result.unwrap();
377        assert_eq!(file_infos.len(), 2);
378        assert_eq!(file_infos[0].file_size(), Some(18));
379        assert_eq!(file_infos[1].file_size(), Some(19));
380        assert_ne!(file_infos[0].hash(), file_infos[1].hash());
381    }
382
383    #[tokio::test]
384    #[cfg_attr(feature = "smoke-test", ignore)]
385    async fn test_hash_file_size_multiple_of_buffer() {
386        // Regression test for bug where final chunk wasn't produced when file size
387        // is exactly a multiple of buffer_size. This test verifies that
388        // chunker.finish() is called to flush any remaining data.
389        let temp_dir = tempdir().unwrap();
390        let file_path = temp_dir.path().join("multiple_of_buffer.bin");
391
392        // Create a file that is exactly 16MB
393        let file_size = 16 * 1024 * 1024;
394        let content: Vec<u8> = (0..file_size).map(|i| (i % 256) as u8).collect();
395        std::fs::write(&file_path, &content).unwrap();
396
397        let file_path_str = file_path.to_str().unwrap().to_string();
398        let ctx = XetContext::default().unwrap();
399
400        // Hash with 8MB buffer size - file is exactly 2x buffer size
401        let result1 = hash_single_file(ctx.clone(), file_path_str.clone(), 8 * 1024 * 1024);
402        assert!(result1.is_ok());
403        let file_info1 = result1.unwrap();
404        assert_eq!(file_info1.file_size(), Some(file_size as u64));
405        assert!(!file_info1.hash().is_empty());
406
407        // Hash with 4MB buffer size - file is exactly 4x buffer size
408        let result2 = hash_single_file(ctx.clone(), file_path_str.clone(), 4 * 1024 * 1024);
409        assert!(result2.is_ok());
410        let file_info2 = result2.unwrap();
411
412        // Hash with 2MB buffer size - file is exactly 8x buffer size
413        let result3 = hash_single_file(ctx, file_path_str, 2 * 1024 * 1024);
414        assert!(result3.is_ok());
415        let file_info3 = result3.unwrap();
416
417        // All hashes should be identical regardless of buffer size
418        // This verifies that chunker.finish() is properly called to flush remaining chunks
419        // Without finish(), different buffer sizes would produce different (incomplete) hashes
420        assert_eq!(file_info1.hash(), file_info2.hash(), "Hash mismatch between 8MB and 4MB buffer sizes");
421        assert_eq!(file_info1.hash(), file_info3.hash(), "Hash mismatch between 8MB and 2MB buffer sizes");
422        assert_eq!(file_info1.file_size(), file_info2.file_size());
423        assert_eq!(file_info1.file_size(), file_info3.file_size());
424    }
425}
xet_data/processing/data_client.rs

xet_data/processing/
data_client.rs
