parutils makeover remove async_scoped (#454)

removing async_scoped dep and creating new parallel util to replace
tokio_par_for_each that relied on async_scoped.

Any usage of tokio_par_for_each which is the only fn used out of
parutils has been replaced with the new
`tokio_run_max_concurrency_fold_result_with_semaphore`

TODO:
- [x] add more tests
- [x] use semaphore acquired from the global semaphore provider where/if
relevant.

Author: assafvayner

Cargo.lock

@@ -12,7 +12,6 @@ members = [
     "file_utils",
     "mdb_shard",
     "merklehash",
-    "parutils",
     "progress_tracking",
     "utils",
     "xet_threadpool",
@@ -32,7 +31,6 @@ debug = 1
 
 [workspace.dependencies]
 anyhow = "1"
-async-scoped = { version = "0.7", features = ["use-tokio"] }
 async-trait = "0.1"
 base64 = "0.22"
 bincode = "1.3"

Cargo.toml

@@ -52,7 +52,6 @@ xet-core enables huggingface_hub to utilize xet storage for uploading and downlo
 * [hf_xet](./hf_xet): Python integration with Rust code, uses maturin to build hfxet Python package. Main integration with HF Hub Python package.
 * [mdb_shard](./mdb_shard): Shard operations, including Shard format, dedupe probing, benchmarks, and utilities.
 * [merklehash](./merklehash): MerkleHash type, 256-bit hash, widely used across many crates.
-* [parutils](./parutils): Provides parallel execution utilities relying on Tokio (ex. parallel foreach).
 * [progress_reporting](./progress_reporting): offers ReportedWriter so progress for Writer operations can be displayed.
 * [utils](./utils): general utilities, including singleflight, progress, serialization_utils and threadpool.
 

README.md

@@ -54,7 +54,7 @@ utils::configurable_constants! {
 }
 
 lazy_static! {
-    static ref DOWNLOAD_CONCURRENCY_LIMITER: GlobalSemaphoreHandle =
+    static ref DOWNLOAD_CHUNK_RANGE_CONCURRENCY_LIMITER: GlobalSemaphoreHandle =
         global_semaphore_handle!(*NUM_CONCURRENT_RANGE_GETS);
 }
 
@@ -329,7 +329,8 @@ impl RemoteClient {
         let download_scheduler = DownloadSegmentLengthTuner::from_configurable_constants();
         let download_scheduler_clone = download_scheduler.clone();
 
-        let download_concurrency_limiter = ThreadPool::current().global_semaphore(*DOWNLOAD_CONCURRENCY_LIMITER);
+        let download_concurrency_limiter =
+            ThreadPool::current().global_semaphore(*DOWNLOAD_CHUNK_RANGE_CONCURRENCY_LIMITER);
 
         let queue_dispatcher: JoinHandle<Result<()>> = tokio::spawn(async move {
             let mut remaining_total_len = total_len;
@@ -479,7 +480,8 @@ impl RemoteClient {
         let term_download_client = self.http_client_with_retry.clone();
         let download_scheduler = DownloadSegmentLengthTuner::from_configurable_constants();
 
-        let download_concurrency_limiter = ThreadPool::current().global_semaphore(*DOWNLOAD_CONCURRENCY_LIMITER);
+        let download_concurrency_limiter =
+            ThreadPool::current().global_semaphore(*DOWNLOAD_CHUNK_RANGE_CONCURRENCY_LIMITER);
 
         let process_result = move |result: TermDownloadResult<u64>,
                                    total_written: &mut u64,

cas_client/src/remote_client.rs

@@ -22,7 +22,6 @@ deduplication = { path = "../deduplication" }
 error_printer = { path = "../error_printer" }
 mdb_shard = { path = "../mdb_shard" }
 merklehash = { path = "../merklehash" }
-parutils = { path = "../parutils" }
 progress_tracking = { path = "../progress_tracking" }
 utils = { path = "../utils" }
 xet_threadpool = { path = "../xet_threadpool" }

data/Cargo.toml

@@ -10,17 +10,19 @@ use cas_client::{CacheConfig, FileProvider, OutputProvider, CHUNK_CACHE_SIZE_BYT
 use cas_object::CompressionScheme;
 use deduplication::DeduplicationMetrics;
 use dirs::home_dir;
-use parutils::{tokio_par_for_each, ParallelError};
 use progress_tracking::item_tracking::ItemProgressUpdater;
 use progress_tracking::TrackingProgressUpdater;
 use tracing::{info, info_span, instrument, Instrument, Span};
 use ulid::Ulid;
 use utils::auth::{AuthConfig, TokenRefresher};
 use utils::normalized_path_from_user_string;
+use xet_threadpool::utils::run_constrained_with_semaphore;
+use xet_threadpool::{global_semaphore_handle, GlobalSemaphoreHandle, ThreadPool};
 
 use crate::configurations::*;
-use crate::constants::{INGESTION_BLOCK_SIZE, MAX_CONCURRENT_DOWNLOADS, MAX_CONCURRENT_FILE_INGESTION};
+use crate::constants::{INGESTION_BLOCK_SIZE, MAX_CONCURRENT_DOWNLOADS};
 use crate::errors::DataProcessingError;
+use crate::file_upload_session::CONCURRENT_FILE_INGESTION_LIMITER;
 use crate::{errors, FileDownloader, FileUploadSession, XetFileInfo};
 
 utils::configurable_constants! {
@@ -125,30 +127,22 @@ pub async fn upload_bytes_async(
     let config = default_config(endpoint.unwrap_or(DEFAULT_CAS_ENDPOINT.clone()), None, token_info, token_refresher)?;
     Span::current().record("session_id", &config.session_id);
 
+    let semaphore = ThreadPool::current().global_semaphore(*CONCURRENT_FILE_INGESTION_LIMITER);
     let upload_session = FileUploadSession::new(config, progress_updater).await?;
-    let blobs_with_spans = add_spans(file_contents, || info_span!("clean_task"));
-
-    // clean the bytes
-    let files = tokio_par_for_each(blobs_with_spans, *MAX_CONCURRENT_FILE_INGESTION, |(blob, span), _| {
-        async {
-            let (xf, _metrics) = clean_bytes(upload_session.clone(), blob).await?;
-            Ok(xf)
-        }
-        .instrument(span.unwrap_or_else(|| info_span!("unexpected_span")))
-    })
-    .await
-    .map_err(|e| match e {
-        ParallelError::JoinError => DataProcessingError::InternalError("Join error".to_string()),
-        ParallelError::TaskError(e) => e,
-    })?;
+    let clean_futures = file_contents.into_iter().map(|blob| {
+        let upload_session = upload_session.clone();
+        async move { clean_bytes(upload_session, blob).await.map(|(xf, _metrics)| xf) }
+            .instrument(info_span!("clean_task"))
+    });
+    let files = run_constrained_with_semaphore(clean_futures, semaphore).await?;
 
     // Push the CAS blocks and flush the mdb to disk
     let _metrics = upload_session.finalize().await?;
 
     Ok(files)
 }
 
-#[instrument(skip_all, name = "data_client::upload_files", 
+#[instrument(skip_all, name = "data_client::upload_files",
     fields(session_id = tracing::field::Empty,
     num_files=file_paths.len(),
     new_bytes = tracing::field::Empty,
@@ -157,7 +151,7 @@ pub async fn upload_bytes_async(
     new_chunks = tracing::field::Empty,
     deduped_chunks = tracing::field::Empty,
     defrag_prevented_dedup_chunks = tracing::field::Empty
-))]
+    ))]
 pub async fn upload_async(
     file_paths: Vec<String>,
     endpoint: Option<String>,
@@ -201,6 +195,11 @@ pub async fn download_async(
     token_refresher: Option<Arc<dyn TokenRefresher>>,
     progress_updaters: Option<Vec<Arc<dyn TrackingProgressUpdater>>>,
 ) -> errors::Result<Vec<String>> {
+    lazy_static! {
+        static ref CONCURRENT_FILE_DOWNLOAD_LIMITER: GlobalSemaphoreHandle =
+            global_semaphore_handle!(*MAX_CONCURRENT_DOWNLOADS);
+    }
+
     if let Some(updaters) = &progress_updaters {
         if updaters.len() != file_infos.len() {
             return Err(DataProcessingError::ParameterError(
@@ -212,30 +211,19 @@ pub async fn download_async(
         default_config(endpoint.unwrap_or(DEFAULT_CAS_ENDPOINT.to_string()), None, token_info, token_refresher)?;
     Span::current().record("session_id", &config.session_id);
 
+    let processor = Arc::new(FileDownloader::new(config).await?);
     let updaters = match progress_updaters {
         None => vec![None; file_infos.len()],
         Some(updaters) => updaters.into_iter().map(Some).collect(),
     };
-    let file_with_progress = file_infos.into_iter().zip(updaters).collect::<Vec<_>>();
-    let extended_file_info_list = add_spans(file_with_progress, || info_span!("download_file"));
-
-    let processor = &Arc::new(FileDownloader::new(config).await?);
-    let paths = tokio_par_for_each(
-        extended_file_info_list,
-        *MAX_CONCURRENT_DOWNLOADS,
-        |(((file_info, file_path), updater), span), _| {
-            async move {
-                let proc = processor.clone();
-                smudge_file(&proc, &file_info, &file_path, updater).await
-            }
-            .instrument(span.unwrap_or_else(|| info_span!("unexpected_span")))
-        },
-    )
-    .await
-    .map_err(|e| match e {
-        ParallelError::JoinError => DataProcessingError::InternalError("Join error".to_string()),
-        ParallelError::TaskError(e) => e,
-    })?;
+    let smudge_file_futures = file_infos.into_iter().zip(updaters).map(|((file_info, file_path), updater)| {
+        let proc = processor.clone();
+        async move { smudge_file(&proc, &file_info, &file_path, updater).await }.instrument(info_span!("download_file"))
+    });
+
+    let semaphore = ThreadPool::current().global_semaphore(*CONCURRENT_FILE_DOWNLOAD_LIMITER);
+
+    let paths = run_constrained_with_semaphore(smudge_file_futures, semaphore).await?;
 
     Ok(paths)
 }
@@ -298,23 +286,12 @@ async fn smudge_file(
     Ok(file_path.to_string())
 }
 
-/// Adds spans to the indicated list for each element.
-///
-/// Although a span will be added for each element, we need an Option<Span> since
-/// tokio_par_for_each requires the input list be Default, which Span isn't.
-pub fn add_spans<I, F: Fn() -> Span>(v: Vec<I>, create_span: F) -> Vec<(I, Option<Span>)> {
-    let spans: Vec<Option<Span>> = v.iter().map(|_| Some(create_span())).collect();
-    v.into_iter().zip(spans).collect()
-}
-
 #[cfg(test)]
 mod tests {
     use std::env;
 
     use serial_test::serial;
     use tempfile::tempdir;
-    use tracing::info;
-    use tracing_test::traced_test;
 
     use super::*;
 
@@ -402,41 +379,4 @@ mod tests {
             "cache dir = {test_cache_dir:?}; does not start with {expected:?}",
         );
     }
-
-    #[tokio::test(flavor = "multi_thread")]
-    #[traced_test]
-    async fn test_add_spans() {
-        let outer_span = info_span!("outer_span");
-        async {
-            let v = vec!["a", "b", "c"];
-            let expected_len = v.len();
-            let v_plus = add_spans(v, || info_span!("task_span"));
-            assert_eq!(v_plus.len(), expected_len);
-            tokio_par_for_each(v_plus, expected_len, |(s, span), i| {
-                async move {
-                    info!("inside: {s},{i}");
-                    Ok::<(), ()>(())
-                }
-                .instrument(span.unwrap())
-            })
-            .await
-            .unwrap();
-        }
-        .instrument(outer_span)
-        .await;
-
-        assert!(logs_contain("inside: a,0"));
-        assert!(logs_contain("inside: b,1"));
-        assert!(logs_contain("inside: c,2"));
-        logs_assert(|lines: &[&str]| {
-            match lines
-                .iter()
-                .filter(|line| line.contains("task_span") && line.contains("outer_span"))
-                .count()
-            {
-                3 => Ok(()),
-                n => Err(format!("Expected 3 lines, got {n}")),
-            }
-        });
-    }
 }

data/src/data_client.rs

@@ -8,6 +8,7 @@ use thiserror::Error;
 use tokio::sync::AcquireError;
 use tracing::error;
 use utils::errors::{AuthError, SingleflightError};
+use xet_threadpool::utils::ParutilsError;
 
 #[derive(Error, Debug)]
 pub enum DataProcessingError {
@@ -91,3 +92,14 @@ impl From<SingleflightError<DataProcessingError>> for DataProcessingError {
         }
     }
 }
+
+impl From<ParutilsError<DataProcessingError>> for DataProcessingError {
+    fn from(value: ParutilsError<DataProcessingError>) -> Self {
+        match value {
+            ParutilsError::Join(e) => DataProcessingError::JoinError(e),
+            ParutilsError::Acquire(e) => DataProcessingError::PermitAcquisitionError(e),
+            ParutilsError::Task(e) => e,
+            e => DataProcessingError::InternalError(e.to_string()),
+        }
+    }
+}