hf-xet 1.1.6.dev0__tar.gz → 1.1.6.dev2__tar.gz

{hf_xet-1.1.6.dev0 → hf_xet-1.1.6.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hf-xet
-Version: 1.1.6.dev0
+Version: 1.1.6.dev2
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy

{hf_xet-1.1.6.dev0 → hf_xet-1.1.6.dev2}/cas_client/Cargo.toml

@@ -34,6 +34,7 @@ url = { workspace = true }
 lazy_static = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
+rand = { workspace = true }
 
 reqwest-middleware = "0.4"
 reqwest-retry = "0.7"

{hf_xet-1.1.6.dev0 → hf_xet-1.1.6.dev2}/cas_client/src/remote_client.rs

@@ -291,7 +291,7 @@ impl RemoteClient {
     // at the beginning of the download, but queried in segments. Range downloads are executed with
     // a certain degree of parallelism, but writing out to storage is sequential. Ideal when the external
     // storage uses HDDs.
-    #[instrument(skip_all, name = "RemoteClient::reconstruct_file_segmented", fields(file.hash = file_hash.hex()
+    #[instrument(skip_all, name = "RemoteClient::reconstruct_file_segmented", fields(file.hash = file_hash.hex(), file.range = format!("{byte_range:?}")
     ))]
     async fn reconstruct_file_to_writer_segmented(
         &self,
@@ -300,6 +300,7 @@ impl RemoteClient {
         writer: &OutputProvider,
         progress_updater: Option<Arc<SingleItemProgressUpdater>>,
     ) -> Result<u64> {
+        eprintln!("Attempting to download file={file_hash:?} with range {byte_range:?}.");
         // queue size is inherently bounded by degree of concurrency.
         let (task_tx, mut task_rx) = mpsc::unbounded_channel::<DownloadQueueItem<SequentialTermDownload>>();
         let (running_downloads_tx, mut running_downloads_rx) =
@@ -336,18 +337,23 @@ impl RemoteClient {
                 match item {
                     DownloadQueueItem::End => {
                         // everything processed
-                        debug!("download queue emptied");
+                        eprintln!("download queue emptied");
                         drop(running_downloads_tx);
                         break;
                     },
                     DownloadQueueItem::DownloadTask(term_download) => {
                         // acquire the permit before spawning the task, so that there's limited
                         // number of active downloads.
+
+                        use rand::Rng;
+                        let x: u64 = rand::rng().random();
+                        eprintln!("{x:x}: Acquiring download permit.");
                         let permit = DOWNLOAD_CONNECTION_CONCURRENCY_LIMITER.clone().acquire_owned().await?;
-                        debug!("spawning 1 download task");
+                        eprintln!("{x:x}: Permit acquried: spawning 1 download task");
                         let future: JoinHandle<Result<(TermDownloadResult<Vec<u8>>, OwnedSemaphorePermit)>> =
                             tokio::spawn(async move {
                                 let data = term_download.run().await?;
+                                eprintln!("{x:x}: Download task completed.");
                                 Ok((data, permit))
                             });
                         running_downloads_tx.send(future)?;
@@ -355,7 +361,7 @@ impl RemoteClient {
                     DownloadQueueItem::Metadata(fetch_info) => {
                         // query for the file info of the first segment
                         let segment_size = download_scheduler_clone.next_segment_size()?;
-                        debug!("querying file info of size {segment_size}");
+                        eprintln!("querying file info of size {segment_size}");
                         let (segment, maybe_remainder) = fetch_info.take_segment(segment_size);
 
                         let Some((offset_into_first_range, terms)) = segment.query().await? else {
@@ -367,7 +373,7 @@ impl RemoteClient {
                         let segment = Arc::new(segment);
                         // define the term download tasks
                         let mut remaining_segment_len = segment_size;
-                        debug!("enqueueing {} download tasks", terms.len());
+                        eprintln!("enqueueing {} download tasks", terms.len());
                         for (i, term) in terms.into_iter().enumerate() {
                             let skip_bytes = if i == 0 { offset_into_first_range } else { 0 };
                             let take = remaining_total_len
@@ -391,7 +397,7 @@ impl RemoteClient {
 
                             remaining_total_len -= take;
                             remaining_segment_len -= take;
-                            debug!("enqueueing {download_task:?}");
+                            eprintln!("enqueueing {download_task:?}");
                             task_tx.send(DownloadQueueItem::DownloadTask(download_task))?;
                         }
 
@@ -409,31 +415,49 @@ impl RemoteClient {
         });
 
         let mut writer = writer.get_writer_at(0)?;
-        let mut total_written = 0;
-        while let Some(result) = running_downloads_rx.recv().await {
-            match result.await {
-                Ok(Ok((mut download_result, permit))) => {
-                    let data = take(&mut download_result.payload);
-                    writer.write_all(&data)?;
-                    // drop permit after data written out so they don't accumulate in memory unbounded
-                    drop(permit);
-
-                    if let Some(updater) = progress_updater.as_ref() {
-                        updater.update(data.len() as u64).await;
-                    }
 
-                    total_written += data.len() as u64;
+        let write_task = tokio::spawn(async move {
+            let mut total_written = 0;
+            while let Some(result) = running_downloads_rx.recv().await {
+                match result.await {
+                    Ok(Ok((mut download_result, permit))) => {
+                        eprintln!(
+                            "Task receive loop: Received ok result with data of size {}",
+                            download_result.payload.len()
+                        );
+                        let data = take(&mut download_result.payload);
+                        writer.write_all(&data)?;
+                        // drop permit after data written out so they don't accumulate in memory unbounded
+                        drop(permit);
+
+                        if let Some(updater) = progress_updater.as_ref() {
+                            updater.update(data.len() as u64).await;
+                        }
+
+                        total_written += data.len() as u64;
 
-                    // Now inspect the download metrics and tune the download degree of concurrency
-                    download_scheduler.tune_on(download_result)?;
-                },
-                Ok(Err(e)) => Err(e)?,
-                Err(e) => Err(anyhow!("{e:?}"))?,
+                        // Now inspect the download metrics and tune the download degree of concurrency
+                        download_scheduler.tune_on(download_result)?;
+                    },
+                    Ok(Err(e)) => {
+                        eprintln!("Task receive loop: Received error result: {e:?}");
+                        Err(e)?
+                    },
+                    Err(e) => Err(anyhow!("{e:?}"))?,
+                }
             }
-        }
-        writer.flush()?;
+            eprintln!("Task receive loop exited.");
+            writer.flush()?;
+
+            Result::Ok(total_written)
+        });
+
+        eprintln!("Both tasks spawned; waiting.");
+        let (write_res, queue_res) = tokio::join! {write_task, queue_dispatcher};
 
-        queue_dispatcher.await??;
+        queue_res??;
+        let total_written = write_res??;
+        eprintln!("Finished; total_written=total_written.");
 
         Ok(total_written)
     }
@@ -724,7 +748,7 @@ impl Client for RemoteClient {
     ) -> Result<u64> {
         // If the user has set the `HF_XET_RECONSTRUCT_WRITE_SEQUENTIALLY=true` env variable, then we
         // should write the file to the output sequentially instead of in parallel.
-        if *RECONSTRUCT_WRITE_SEQUENTIALLY {
+        if true || *RECONSTRUCT_WRITE_SEQUENTIALLY {
             info!("reconstruct terms sequentially");
             self.reconstruct_file_to_writer_segmented(hash, byte_range, output_provider, progress_updater)
                 .await

{hf_xet-1.1.6.dev0 → hf_xet-1.1.6.dev2}/data/src/file_downloader.rs

@@ -37,7 +37,7 @@ impl FileDownloader {
         Ok(Self { config, client })
     }
 
-    #[instrument(skip_all, name = "FileDownloader::smudge_file_from_hash", fields(hash=file_id.hex()))]
+    #[instrument(skip_all, name = "FileDownloader::smudge_file_from_hash", fields(hash=file_id.hex(), file_name = file_name.to_string(), range = format!("{:?}", range)))]
     pub async fn smudge_file_from_hash(
         &self,
         file_id: &MerkleHash,
@@ -46,6 +46,7 @@ impl FileDownloader {
         range: Option<FileRange>,
         progress_updater: Option<Arc<ItemProgressUpdater>>,
     ) -> Result<u64> {
+        eprintln!("Attempting to download {file_name:?}, range={range:?},file_id={file_id:?}");
         let file_progress_tracker = progress_updater.map(|p| ItemProgressUpdater::item_tracker(&p, file_name, None));
 
         // Currently, this works by always directly querying the remote server.

{hf_xet-1.1.6.dev0 → hf_xet-1.1.6.dev2}/hf_xet/Cargo.lock

@@ -277,6 +277,7 @@ dependencies = [
  "merklehash",
  "more-asserts",
  "progress_tracking",
+ "rand 0.9.2",
  "reqwest",
  "reqwest-middleware",
  "reqwest-retry",
@@ -311,7 +312,7 @@ dependencies = [
  "mdb_shard",
  "merklehash",
  "more-asserts",
- "rand 0.9.1",
+ "rand 0.9.2",
  "serde",
  "thiserror 2.0.12",
  "tokio",
@@ -384,7 +385,7 @@ dependencies = [
  "merklehash",
  "mockall",
  "once_cell",
- "rand 0.9.1",
+ "rand 0.9.2",
  "thiserror 2.0.12",
  "tokio",
  "tracing",
@@ -625,7 +626,7 @@ dependencies = [
  "parutils",
  "progress_tracking",
  "prometheus",
- "rand 0.9.1",
+ "rand 0.9.2",
  "rand_chacha 0.9.0",
  "regex",
  "serde",
@@ -809,7 +810,7 @@ dependencies = [
  "colored",
  "lazy_static",
  "libc",
- "rand 0.9.1",
+ "rand 0.9.2",
  "tracing",
  "whoami",
  "winapi",
@@ -1079,7 +1080,7 @@ checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c"
 
 [[package]]
 name = "hf_xet"
-version = "1.1.6-dev0"
+version = "1.1.6-dev2"
 dependencies = [
  "async-trait",
  "bipbuffer",
@@ -1093,6 +1094,7 @@ dependencies = [
  "pprof",
  "progress_tracking",
  "pyo3",
+ "rand 0.9.2",
  "serde",
  "serde_json",
  "signal-hook",
@@ -1581,7 +1583,7 @@ dependencies = [
  "itertools 0.14.0",
  "lazy_static",
  "merklehash",
- "rand 0.9.1",
+ "rand 0.9.2",
  "regex",
  "serde",
  "static_assertions",
@@ -1625,7 +1627,7 @@ dependencies = [
  "blake3",
  "getrandom 0.3.3",
  "heed",
- "rand 0.9.1",
+ "rand 0.9.2",
  "safe-transmute",
  "serde",
 ]
@@ -2269,7 +2271,7 @@ dependencies = [
  "bytes",
  "getrandom 0.3.3",
  "lru-slab",
- "rand 0.9.1",
+ "rand 0.9.2",
  "ring",
  "rustc-hash",
  "rustls",
@@ -2323,9 +2325,9 @@ dependencies = [
 
 [[package]]
 name = "rand"
-version = "0.9.1"
+version = "0.9.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9fbfd9d094a40bf3ae768db9361049ace4c0e04a4fd6b359518bd7b73a73dd97"
+checksum = "6db2770f06117d490610c7488547d543617b21bfa07796d7a12f6f1bd53850d1"
 dependencies = [
  "rand_chacha 0.9.0",
  "rand_core 0.9.3",
@@ -3341,7 +3343,7 @@ version = "1.2.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "470dbf6591da1b39d43c14523b2b469c86879a53e8b758c8e090a470fe7b1fbe"
 dependencies = [
- "rand 0.9.1",
+ "rand 0.9.2",
  "web-time",
 ]
 

{hf_xet-1.1.6.dev0 → hf_xet-1.1.6.dev2}/hf_xet/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "hf_xet"
-version = "1.1.6-dev0"
+version = "1.1.6-dev2"
 edition = "2021"
 license = "Apache-2.0"
 
@@ -44,6 +44,7 @@ tracing-subscriber = { version = "0.3", features = [
     "env-filter",
 ] }
 tracing-appender = "0.2"
+rand = "0.9.2"
 
 # Unix-specific dependencies
 [target.'cfg(unix)'.dependencies]

{hf_xet-1.1.6.dev0 → hf_xet-1.1.6.dev2}/hf_xet/src/lib.rs

@@ -107,14 +107,23 @@ pub fn download_files(
     let refresher = token_refresher.map(WrappedTokenRefresher::from_func).transpose()?.map(Arc::new);
     let updaters = progress_updater.map(try_parse_progress_updaters).transpose()?;
 
-    async_run(py, async move {
+    use rand::Rng;
+
+    let x: u64 = rand::rng().random();
+    eprintln!("{x:x}: Downloading files {file_infos:?}");
+
+    let res = async_run(py, async move {
         let out: Vec<String> =
             data_client::download_async(file_infos, endpoint, token_info, refresher.map(|v| v as Arc<_>), updaters)
                 .await
                 .map_err(convert_data_processing_error)?;
 
         PyResult::Ok(out)
-    })
+    });
+
+    eprintln!("{x:x}: Completed.");
+
+    res
 }
 
 fn try_parse_progress_updaters(funcs: Vec<Py<PyAny>>) -> PyResult<Vec<Arc<dyn TrackingProgressUpdater>>> {