29603096773

Committed 17 Jul 2026 06:14PM UTC coverage: 95.82% (-0.005%) from 95.825%

Build # 29603096773

Build Type

push

github

Committed by

web-flow

Commit Message

v0.25.0-alpha (#132)

* Bump parquet from 58.3.0 to 59.1.0 (#126)

Bumps [parquet](https://github.com/apache/arrow-rs) from 58.3.0 to 59.1.0.
- [Release notes](https://github.com/apache/arrow-rs/releases)
- [Changelog](https://github.com/apache/arrow-rs/blob/main/CHANGELOG.md)
- [Commits](https://github.com/apache/arrow-rs/compare/58.3.0...59.1.0)

---
updated-dependencies:
- dependency-name: parquet
  dependency-version: 59.1.0
  dependency-type: direct:production
  update-type: version-update:semver-major
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>

* Add `Sampler` mixed strategy tests (#127)

* Remove `hf-hub` dependency and legacy `datasets-server` fallback paths  (#130)

Why: Full control over download and cache handling.

Also, fixes new 503 errors which were starting to originate from deprecated
`datasets-server` fallbacks.

Replace the datasets-server /parquet, /size, and /info endpoints with
the Hub API tree endpoint (/api/datasets/{dataset}/tree/main). Remove
the hf-hub crate dependency entirely, eliminating the three legacy
fallback paths: sibling-based candidate resolution, ClassLabel /info
resolution, and global row count /size queries.
Key changes:
- remote_url_for_candidate handles bare relative paths, full HTTP URLs,
  and url::-prefixed candidates
- candidate_target_path resolves CDN /resolve/ URLs and bare relative
  paths
- Tree endpoint scoped to config subdirectory; falls back to root for
  "default"/empty config
- Path traversal guard rejects ParentDir/Prefix components
- HEAD fallback in download_and_materialize_shard_with_runtime uses
  fetch_remote_size_with_runtime
- target_matches_expected_size rejects 0-byte files when expected_bytes
  is None
- Mock servers return full CDN URLs; HfMockServer uses org/dataset
  dataset name
- Removed: label_maps, size_endpoint, info_endpoint, known_total_rows(),
  collect_candidates_from_sibling... (continued)

Coverage Stats

5137 of 5662 new or added lines in 35 files covered. (90.73%)

11 existing lines in 4 files now uncovered.

21068 of 21987 relevant lines covered (95.82%)

5318.83 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

91.95

/crates/triplets-hf-source/src/disk_cache/disk_cache.rs

use crate::config::HuggingFaceRowsConfig;
use crate::constants::{HF_ALL_SPLITS_DIR, HF_GROUP};
use cache_manager::CacheRoot;
use simd_r_drive::DataStore;
use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex, MutexGuard};
use tracing::warn;
use triplets_core::SamplerError;

#[cfg(test)]
use std::sync::OnceLock;
#[cfg(test)]
use tempfile::TempDir;

/// Shared handle to the open-store cache.  Stored on `HuggingFaceRowsConfig`
/// so all methods have access without passing it separately.
#[derive(Clone)]
pub struct StoreCache(pub(crate) Arc<Mutex<HashMap<PathBuf, Arc<DataStore>>>>);

impl std::fmt::Debug for StoreCache {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("StoreCache").finish_non_exhaustive()
    }
}

impl StoreCache {
    pub(crate) fn new() -> Self {
        StoreCache(Arc::new(Mutex::new(HashMap::new())))
    }

    pub(crate) fn lock(
        &self,
    ) -> Result<MutexGuard<'_, HashMap<PathBuf, Arc<DataStore>>>, SamplerError> {
        self.0.lock().map_err(|_| SamplerError::SourceUnavailable {
            source_id: "store_cache".to_string(),
            reason: "row-store cache lock poisoned".to_string(),
        })
    }

    pub(crate) fn lock_ok(&self) -> Option<MutexGuard<'_, HashMap<PathBuf, Arc<DataStore>>>> {
        self.0.lock().ok()
    }
}

pub(crate) fn managed_cache_root() -> Result<CacheRoot, String> {
    #[cfg(test)]
    {
        static TEST_CACHE_ROOT: OnceLock<TempDir> = OnceLock::new();
        let root = TEST_CACHE_ROOT
            .get_or_init(|| TempDir::new().expect("failed to create test HF cache root"));
        Ok(CacheRoot::from_root(root.path()))
    }

    #[cfg(not(test))]
    {
        CacheRoot::from_discovery()
            .map_err(|err| format!("failed discovering managed cache root: {err}"))
    }
}

pub(crate) fn ensure_cache_group(relative_group: PathBuf) -> Result<PathBuf, String> {
    let cache_root = managed_cache_root()?;
    cache_root.ensure_group(&relative_group).map_err(|err| {
        format!(
            "failed creating managed cache group '{}': {err}",
            relative_group.display()
        )
    })
}

/// Evict a stale store from the cache and unlink the file so the shard
/// gets re-downloaded on the next cycle.
pub(crate) fn remove_stale_store(config: &HuggingFaceRowsConfig, path: &Path) {
    let _ = config
        .store_cache
        .lock_ok()
        .map(|mut cache| cache.remove(path));
    if let Err(err) = fs::remove_file(path) {
        warn!(
            "[triplets:hf] failed to remove stale store {}: {}",
            path.display(),
            err
        );
    }
}

/// Get or open a store through the shared cache.  Never opens a duplicate
/// handle — if the path is already in `store_cache`, returns the cached
/// `Arc`; otherwise opens and inserts it.
pub(crate) fn open_store_via_cache(
    config: &HuggingFaceRowsConfig,
    path: &Path,
) -> Result<Arc<DataStore>, SamplerError> {
    // Fast path: check the cache while holding the lock briefly.
    {
        let cache = config.store_cache.lock()?;
        if let Some(store) = cache.get(path).cloned() {
            return Ok(store);
        }
    }
    // Open the store outside the lock so that concurrent calls (e.g. from
    // rayon's parallel iteration in build_shard_index) can proceed in
    // parallel instead of being serialized on the mutex.
    let store = Arc::new(crate::shard_indexer::open_shard_store(config, path)?);
    // Re-acquire the lock and insert into the cache.  If another thread
    // already inserted the same path, our duplicate handle is harmless
    // (the cache retains the first one).  We return our handle either way
    // — both point to the same underlying file.
    let mut cache = config.store_cache.lock()?;
    cache
        .entry(path.to_path_buf())
        .or_insert_with(|| store.clone());
    Ok(store)
}

/// Resolve a managed snapshot directory for a list-based Hugging Face source.
pub fn managed_hf_list_snapshot_dir(
    dataset: &str,
    config: &str,
    split: &str,
    replica_idx: usize,
) -> Result<PathBuf, String> {
    // Empty split (all-splits mode) uses HF_ALL_SPLITS_DIR so the path hierarchy stays valid
    // and won't collide with a split literally named "" on any filesystem.
    let split_dir = if split.is_empty() {
        HF_ALL_SPLITS_DIR
    } else {
        split
    };
    ensure_cache_group(
        PathBuf::from(HF_GROUP)
            .join("source-list")
            .join(dataset.replace('/', "__"))
            .join(config)
            .join(split_dir)
            .join(format!("replica_{replica_idx}")),
    )
}

/// Resolve a managed snapshot directory for a single Hugging Face source.
pub fn managed_hf_snapshot_dir(
    dataset: &str,
    config: &str,
    split: &str,
) -> Result<PathBuf, String> {
    let split_dir = if split.is_empty() {
        HF_ALL_SPLITS_DIR
    } else {
        split
    };
    ensure_cache_group(
        PathBuf::from(HF_GROUP)
            .join(dataset.replace('/', "__"))
            .join(config)
            .join(split_dir),
    )
}

1	use crate::config::HuggingFaceRowsConfig;
2	use crate::constants::{HF_ALL_SPLITS_DIR, HF_GROUP};
3	use cache_manager::CacheRoot;
4	use simd_r_drive::DataStore;
5	use std::collections::HashMap;
6	use std::fs;
7	use std::path::{Path, PathBuf};
8	use std::sync::{Arc, Mutex, MutexGuard};
9	use tracing::warn;
10	use triplets_core::SamplerError;
11
12	#[cfg(test)]
13	use std::sync::OnceLock;
14	#[cfg(test)]
15	use tempfile::TempDir;
16
17	/// Shared handle to the open-store cache. Stored on `HuggingFaceRowsConfig`
18	/// so all methods have access without passing it separately.
19	#[derive(Clone)]
20	pub struct StoreCache(pub(crate) Arc<Mutex<HashMap<PathBuf, Arc<DataStore>>>>);
21
22	impl std::fmt::Debug for StoreCache {
NEW 23	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {	×
NEW 24	f.debug_struct("StoreCache").finish_non_exhaustive()	×
NEW 25	}	×
26	}
27
28	impl StoreCache {
29	pub(crate) fn new() -> Self {	295✔
30	StoreCache(Arc::new(Mutex::new(HashMap::new())))	295✔
31	}	295✔
32
33	pub(crate) fn lock(	470✔
34	&self,	470✔
35	) -> Result<MutexGuard<'_, HashMap<PathBuf, Arc<DataStore>>>, SamplerError> {	470✔
36	self.0.lock().map_err(\|_\| SamplerError::SourceUnavailable {	470✔
NEW 37	source_id: "store_cache".to_string(),	×
NEW 38	reason: "row-store cache lock poisoned".to_string(),	×
NEW 39	})	×
40	}	470✔
41
42	pub(crate) fn lock_ok(&self) -> Option<MutexGuard<'_, HashMap<PathBuf, Arc<DataStore>>>> {	195✔
43	self.0.lock().ok()	195✔
44	}	195✔
45	}
46
47	pub(crate) fn managed_cache_root() -> Result<CacheRoot, String> {	25✔
48	#[cfg(test)]
49	{
50	static TEST_CACHE_ROOT: OnceLock<TempDir> = OnceLock::new();
51	let root = TEST_CACHE_ROOT	17✔
52	.get_or_init(\|\| TempDir::new().expect("failed to create test HF cache root"));	17✔
53	Ok(CacheRoot::from_root(root.path()))	17✔
54	}
55
56	#[cfg(not(test))]
57	{
58	CacheRoot::from_discovery()	8✔
59	.map_err(\|err\| format!("failed discovering managed cache root: {err}"))	8✔
60	}
61	}	25✔
62
63	pub(crate) fn ensure_cache_group(relative_group: PathBuf) -> Result<PathBuf, String> {	25✔
64	let cache_root = managed_cache_root()?;	25✔
65	cache_root.ensure_group(&relative_group).map_err(\|err\| {	25✔
66	format!(	1✔
67	"failed creating managed cache group '{}': {err}",
68	relative_group.display()	1✔
69	)
70	})	1✔
71	}	25✔
72
73	/// Evict a stale store from the cache and unlink the file so the shard
74	/// gets re-downloaded on the next cycle.
75	pub(crate) fn remove_stale_store(config: &HuggingFaceRowsConfig, path: &Path) {	5✔
76	let _ = config	5✔
77	.store_cache	5✔
78	.lock_ok()	5✔
79	.map(\|mut cache\| cache.remove(path));	5✔
80	if let Err(err) = fs::remove_file(path) {	5✔
81	warn!(	1✔
82	"[triplets:hf] failed to remove stale store {}: {}",
NEW 83	path.display(),	×
84	err
85	);
86	}	4✔
87	}	5✔
88
89	/// Get or open a store through the shared cache. Never opens a duplicate
90	/// handle — if the path is already in `store_cache`, returns the cached
91	/// `Arc`; otherwise opens and inserts it.
92	pub(crate) fn open_store_via_cache(	23✔
93	config: &HuggingFaceRowsConfig,	23✔
94	path: &Path,	23✔
95	) -> Result<Arc<DataStore>, SamplerError> {	23✔
96	// Fast path: check the cache while holding the lock briefly.
97	{
98	let cache = config.store_cache.lock()?;	23✔
99	if let Some(store) = cache.get(path).cloned() {	23✔
100	return Ok(store);	5✔
101	}	18✔
102	}
103	// Open the store outside the lock so that concurrent calls (e.g. from
104	// rayon's parallel iteration in build_shard_index) can proceed in
105	// parallel instead of being serialized on the mutex.
106	let store = Arc::new(crate::shard_indexer::open_shard_store(config, path)?);	18✔
107	// Re-acquire the lock and insert into the cache. If another thread
108	// already inserted the same path, our duplicate handle is harmless
109	// (the cache retains the first one). We return our handle either way
110	// — both point to the same underlying file.
111	let mut cache = config.store_cache.lock()?;	18✔
112	cache	18✔
113	.entry(path.to_path_buf())	18✔
114	.or_insert_with(\|\| store.clone());	18✔
115	Ok(store)	18✔
116	}	23✔
117
118	/// Resolve a managed snapshot directory for a list-based Hugging Face source.
119	pub fn managed_hf_list_snapshot_dir(	19✔
120	dataset: &str,	19✔
121	config: &str,	19✔
122	split: &str,	19✔
123	replica_idx: usize,	19✔
124	) -> Result<PathBuf, String> {	19✔
125	// Empty split (all-splits mode) uses HF_ALL_SPLITS_DIR so the path hierarchy stays valid
126	// and won't collide with a split literally named "" on any filesystem.
127	let split_dir = if split.is_empty() {	19✔
128	HF_ALL_SPLITS_DIR	1✔
129	} else {
130	split	18✔
131	};
132	ensure_cache_group(	19✔
133	PathBuf::from(HF_GROUP)	19✔
134	.join("source-list")	19✔
135	.join(dataset.replace('/', "__"))	19✔
136	.join(config)	19✔
137	.join(split_dir)	19✔
138	.join(format!("replica_{replica_idx}")),	19✔
139	)
140	}	19✔
141
142	/// Resolve a managed snapshot directory for a single Hugging Face source.
143	pub fn managed_hf_snapshot_dir(	5✔
144	dataset: &str,	5✔
145	config: &str,	5✔
146	split: &str,	5✔
147	) -> Result<PathBuf, String> {	5✔
148	let split_dir = if split.is_empty() {	5✔
149	HF_ALL_SPLITS_DIR	2✔
150	} else {
151	split	3✔
152	};
153	ensure_cache_group(	5✔
154	PathBuf::from(HF_GROUP)	5✔
155	.join(dataset.replace('/', "__"))	5✔
156	.join(config)	5✔
157	.join(split_dir),	5✔
158	)
159	}	5✔

jzombie / rust-triplets / 29603096773

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous