tenshift_core/sources/
mod.rs

1//! Built-in sources for common data origins.
2//!
3//! - [`GlobSource`]  -  load files matching a glob pattern (e.g., `"data/*.json"`)
4//! - [`MemorySource`]  -  load from an in-memory collection (for testing)
5//! - [`ImageFolderSource`]  -  load files split into subdirectories by class label
6//! - [`DistributedSampler`]  -  shard data across distributed training ranks
7//! - [`JsonlSource`]  -  streaming JSON Lines parser
8//! - `CsvSource`  -  load CSV files (feature: `csv`)
9//! - `RingSource`  -  `io_uring`-accelerated batch file loading (feature: `uring`)
10//!
11//! # Creating Custom Sources
12//!
13//! Implement the [`Source`] trait to add new data origins:
14//!
15//! ```rust
16//! use tenshift_core::source::{Source, SourceIterator};
17//! use tenshift_core::sample::Sample;
18//! use tenshift_core::error::Result;
19//!
20//! struct MySource;
21//!
22//! impl Source for MySource {
23//!     fn open(&self) -> Result<Box<dyn SourceIterator>> {
24//!         // Return an iterator over your data
25//!         Ok(Box::new(std::iter::empty()))
26//!     }
27//!
28//!     fn name(&self) -> &str { "my_source" }
29//! }
30//! ```
31
32use std::path::{Path, PathBuf};
33
34use crate::error::{Error, Result};
35use crate::sample::{Sample, Tensor};
36use crate::source::{Source, SourceIterator};
37
38/// Collect every regular file matching a glob `pattern`, failing loud on any
39/// per-path iteration error.
40///
41/// The former `glob::glob(pattern)?.filter_map(Result::ok)` silently discarded
42/// each [`glob::GlobError`] (an I/O error, e.g. permission denied, while reading
43/// a directory mid-iteration), so an unreadable subtree quietly shrank the
44/// training set with no operator-visible signal - a Law-10 silent input loss in
45/// the data pipeline. Each `GlobError` is now surfaced as [`Error::Io`]. Shared
46/// by [`GlobSource::new`] and `RingSource::from_glob` (ONE-PLACE: previously the
47/// identical filtering was hand-rolled in both).
48pub(crate) fn collect_glob_files(pattern: &str) -> Result<Vec<PathBuf>> {
49    let mut paths = Vec::new();
50    for entry in glob::glob(pattern)? {
51        let path = entry.map_err(|e| Error::Io(std::sync::Arc::new(e.into_error())))?;
52        if path.is_file() {
53            paths.push(path);
54        }
55    }
56    Ok(paths)
57}
58
59#[cfg(feature = "image")]
60pub mod image_folder;
61
62#[cfg(feature = "image")]
63pub use image_folder::ImageFolderSource;
64
65#[cfg(feature = "csv")]
66mod csv_source;
67#[cfg(feature = "uring")]
68mod ring_source;
69
70#[cfg(feature = "csv")]
71pub use csv_source::CsvSource;
72#[cfg(feature = "uring")]
73pub use ring_source::RingSource;
74
75/// A source that reads files matching a glob pattern.
76///
77/// Each file becomes one sample with a `"data"` field containing the raw bytes
78/// and metadata containing the filename.
79///
80/// ```rust
81/// use tenshift_core::sources::GlobSource;
82/// use tenshift_core::source::Source;
83///
84/// # fn example() -> tenshift_core::error::Result<()> {
85/// let source = GlobSource::new("tests/data/*.json")?;
86/// println!("found {} files", match source.len_hint() { Some(v) => v, None => 0 });
87/// # Ok(())
88/// # }
89/// ```
90pub struct GlobSource {
91    paths: Vec<PathBuf>,
92    pattern: String,
93}
94
95impl GlobSource {
96    /// Create a source from a glob pattern.
97    ///
98    /// # Errors
99    ///
100    /// Returns [`Error::InvalidPattern`] if the pattern is invalid, or
101    /// [`Error::EmptySource`] if no files match.
102    pub fn new(pattern: &str) -> Result<Self> {
103        let paths = collect_glob_files(pattern)?;
104
105        if paths.is_empty() {
106            return Err(Error::EmptySource {
107                pattern: pattern.to_string(),
108            });
109        }
110
111        tracing::info!(
112            "GlobSource: found {} files matching '{}'",
113            paths.len(),
114            pattern
115        );
116
117        Ok(Self {
118            paths,
119            pattern: pattern.to_string(),
120        })
121    }
122
123    /// Create a source from an explicit list of file paths.
124    ///
125    /// # Errors
126    ///
127    /// Returns [`Error::EmptySource`] if the path list is empty.
128    pub fn from_paths(paths: Vec<PathBuf>) -> Result<Self> {
129        if paths.is_empty() {
130            return Err(Error::EmptySource {
131                pattern: "<explicit paths>".to_string(),
132            });
133        }
134        Ok(Self {
135            pattern: format!("{} explicit paths", paths.len()),
136            paths,
137        })
138    }
139}
140
141impl Source for GlobSource {
142    fn open(&self) -> Result<Box<dyn SourceIterator>> {
143        Ok(Box::new(GlobIterator {
144            paths: self.paths.clone(),
145            index: 0,
146        }))
147    }
148
149    fn len_hint(&self) -> Option<u64> {
150        Some(self.paths.len() as u64)
151    }
152
153    fn name(&self) -> &str {
154        &self.pattern
155    }
156}
157
158/// Iterator over files in a `GlobSource`.
159struct GlobIterator {
160    paths: Vec<PathBuf>,
161    index: u64,
162}
163
164impl SourceIterator for GlobIterator {
165    fn next_sample(&mut self) -> Option<Result<Sample>> {
166        let Ok(idx) = usize::try_from(self.index) else {
167            return Some(Err(Error::InvalidConfig {
168                reason: "source index exceeded addressable memory on this platform".to_string(),
169            }));
170        };
171        if idx >= self.paths.len() {
172            return None;
173        }
174
175        let path = &self.paths[idx];
176        let index = self.index;
177        self.index += 1;
178
179        Some(load_file_as_sample(path, index))
180    }
181}
182
183/// Maximum JSONL line length (256MB). Prevents OOM from malicious/invalid
184/// input files with unbounded lines.
185const MAX_JSONL_LINE_LENGTH: usize = 256 * 1024 * 1024;
186
187/// Load a single file into a Sample.
188fn load_file_as_sample(path: &Path, index: u64) -> Result<Sample> {
189    // Check size before reading to prevent OOM on huge files
190    let metadata = std::fs::metadata(path).map_err(|e| Error::ReadFailed {
191        path: path.to_path_buf(),
192        reason: e.to_string(),
193    })?;
194    if metadata.len() > crate::pipeline::MAX_LOAD_FILE_SIZE {
195        return Err(Error::ReadFailed {
196            path: path.to_path_buf(),
197            reason: format!(
198                "file size {} exceeds maximum {} bytes",
199                metadata.len(),
200                crate::pipeline::MAX_LOAD_FILE_SIZE
201            ),
202        });
203    }
204
205    let data = std::fs::read(path).map_err(|e| Error::ReadFailed {
206        path: path.to_path_buf(),
207        reason: e.to_string(),
208    })?;
209
210    let filename = match path.file_name() {
211        Some(n) => n.to_string_lossy().to_string(),
212        None => String::new(),
213    };
214
215    Ok(Sample::new()
216        .with("data", Tensor::bytes(data))
217        .with("filename", Tensor::bytes(filename.as_bytes().to_vec()))
218        .with_metadata(path.to_string_lossy(), index))
219}
220
221/// A source that yields samples from an in-memory collection.
222///
223/// Useful for testing and small datasets.
224pub struct MemorySource {
225    samples: Vec<Sample>,
226    name: String,
227}
228
229impl MemorySource {
230    /// Create a source from a vector of samples.
231    pub fn new(name: impl Into<String>, samples: Vec<Sample>) -> Self {
232        Self {
233            samples,
234            name: name.into(),
235        }
236    }
237}
238
239impl Source for MemorySource {
240    fn open(&self) -> Result<Box<dyn SourceIterator>> {
241        let samples: Vec<Result<Sample>> = self.samples.iter().cloned().map(Ok).collect();
242        Ok(Box::new(samples.into_iter()))
243    }
244
245    fn len_hint(&self) -> Option<u64> {
246        Some(self.samples.len() as u64)
247    }
248
249    fn name(&self) -> &str {
250        &self.name
251    }
252}
253
254/// A source wrapper that shards samples across distributed ranks.
255///
256/// Only samples where `index % world_size == rank` are yielded. The source
257/// item's position within the underlying iterator determines the index.
258///
259/// ```rust
260/// use tenshift_core::sample::Sample;
261/// use tenshift_core::source::Source;
262/// use tenshift_core::sources::{DistributedSampler, MemorySource};
263///
264/// let source = MemorySource::new("demo", vec![Sample::new(), Sample::new(), Sample::new()]);
265/// let shard = DistributedSampler::new(source, 1, 2).unwrap();
266/// assert_eq!(shard.len_hint(), Some(1));
267/// ```
268pub struct DistributedSampler<S> {
269    inner: S,
270    rank: usize,
271    world_size: usize,
272    name: String,
273}
274
275impl<S> DistributedSampler<S> {
276    /// Create a distributed sampler over an existing source.
277    ///
278    /// # Errors
279    ///
280    /// Returns [`Error::InvalidConfig`] when `world_size == 0` or `rank >= world_size`.
281    pub fn new(inner: S, rank: usize, world_size: usize) -> Result<Self>
282    where
283        S: Source,
284    {
285        validate_shard_config(rank, world_size)?;
286
287        Ok(Self {
288            name: format!("{}[rank={rank}/{world_size}]", inner.name()),
289            inner,
290            rank,
291            world_size,
292        })
293    }
294}
295
296impl<S> Source for DistributedSampler<S>
297where
298    S: Source,
299{
300    fn open(&self) -> Result<Box<dyn SourceIterator>> {
301        Ok(Box::new(DistributedSamplerIter {
302            inner: self.inner.open()?,
303            index: 0,
304            rank: self.rank,
305            world_size: self.world_size,
306        }))
307    }
308
309    fn len_hint(&self) -> Option<u64> {
310        self.inner
311            .len_hint()
312            .map(|len| distributed_len_hint(len, self.rank, self.world_size))
313    }
314
315    fn name(&self) -> &str {
316        &self.name
317    }
318}
319
320struct DistributedSamplerIter {
321    inner: Box<dyn SourceIterator>,
322    index: u64,
323    rank: usize,
324    world_size: usize,
325}
326
327impl SourceIterator for DistributedSamplerIter {
328    fn next_sample(&mut self) -> Option<Result<Sample>> {
329        loop {
330            let item = self.inner.next_sample()?;
331            let index = self.index;
332            self.index = self.index.saturating_add(1);
333
334            let in_shard = index % self.world_size as u64 == self.rank as u64;
335            // Always forward errors regardless of shard - they represent source-level failures
336            // that should be handled by ErrorPolicy, not silently dropped
337            match item {
338                Ok(sample) if in_shard => return Some(Ok(sample)),
339                Err(error) => return Some(Err(error)),
340                Ok(_) => {}
341            }
342        }
343    }
344}
345
346fn validate_shard_config(rank: usize, world_size: usize) -> Result<()> {
347    if world_size == 0 {
348        return Err(Error::InvalidConfig {
349            reason: "world_size must be greater than zero. Fix: pass a positive shard count."
350                .to_string(),
351        });
352    }
353
354    if rank >= world_size {
355        return Err(Error::InvalidConfig {
356            reason: format!(
357                "rank {rank} is out of range for world_size {world_size}. Fix: use a rank in 0..{world_size}."
358            ),
359        });
360    }
361
362    Ok(())
363}
364
365fn distributed_len_hint(len: u64, rank: usize, world_size: usize) -> u64 {
366    if len <= rank as u64 {
367        0
368    } else {
369        ((len - 1 - rank as u64) / world_size as u64) + 1
370    }
371}
372
373/// A JSONL (JSON Lines) source  -  one JSON object per line.
374///
375/// Each line becomes a sample with a `"json"` field containing the raw JSON bytes.
376pub struct JsonlSource {
377    path: PathBuf,
378}
379
380impl JsonlSource {
381    /// Create a JSONL source from a file path.
382    pub fn new(path: impl Into<PathBuf>) -> Self {
383        Self { path: path.into() }
384    }
385}
386
387impl Source for JsonlSource {
388    fn open(&self) -> Result<Box<dyn SourceIterator>> {
389        let file = std::fs::File::open(&self.path).map_err(|e| Error::ReadFailed {
390            path: self.path.clone(),
391            reason: e.to_string(),
392        })?;
393        let reader = std::io::BufReader::new(file);
394        Ok(Box::new(JsonlIterator {
395            reader,
396            path: self.path.clone(),
397            index: 0,
398        }))
399    }
400
401    fn name(&self) -> &str {
402        match self.path.to_str() {
403            Some(v) => v,
404            None => "jsonl",
405        }
406    }
407}
408
409/// Iterator over lines in a JSONL file.
410struct JsonlIterator {
411    reader: std::io::BufReader<std::fs::File>,
412    path: PathBuf,
413    index: u64,
414}
415
416impl SourceIterator for JsonlIterator {
417    fn next_sample(&mut self) -> Option<Result<Sample>> {
418        use std::io::BufRead;
419
420        let mut line = String::new();
421        match self.reader.read_line(&mut line) {
422            Ok(0) => None, // EOF
423            Ok(n) => {
424                // Enforce maximum line length to prevent OOM from malicious input
425                if n > MAX_JSONL_LINE_LENGTH {
426                    return Some(Err(Error::CorruptData {
427                        path: self.path.clone(),
428                        reason: format!(
429                            "line {} exceeds maximum length of {} bytes (got {} bytes). Fix: check for corrupt data or increase MAX_JSONL_LINE_LENGTH",
430                            self.index + 1,
431                            MAX_JSONL_LINE_LENGTH,
432                            n
433                        ),
434                    }));
435                }
436                let trimmed = line.trim();
437                if trimmed.is_empty() {
438                    // Skip empty lines, recurse
439                    return self.next_sample();
440                }
441                let index = self.index;
442                self.index += 1;
443                // Validate JSON syntax without building a DOM tree.
444                // `RawValue` confirms the line is valid JSON but does not
445                // allocate strings, maps, or vectors  -  just a bounds check.
446                if let Err(error) = serde_json::from_str::<&serde_json::value::RawValue>(trimmed) {
447                    return Some(Err(Error::CorruptData {
448                        path: self.path.clone(),
449                        reason: format!("invalid JSON on line {}: {error}", index + 1),
450                    }));
451                }
452                Some(Ok(Sample::new()
453                    .with("json", Tensor::bytes(trimmed.as_bytes().to_vec()))
454                    .with_metadata(self.path.to_string_lossy(), index)))
455            }
456            Err(e) => Some(Err(Error::ReadFailed {
457                path: self.path.clone(),
458                reason: e.to_string(),
459            })),
460        }
461    }
462}
463
464#[cfg(test)]
465mod glob_tests {
466    use super::collect_glob_files;
467
468    #[test]
469    fn collect_glob_files_returns_only_regular_files() {
470        let dir = tempfile::tempdir().unwrap();
471        std::fs::write(dir.path().join("a.bin"), "a").unwrap();
472        std::fs::write(dir.path().join("b.bin"), "b").unwrap();
473        // A subdirectory matching the glob must be excluded (not a regular file).
474        std::fs::create_dir(dir.path().join("c.bin")).unwrap();
475        // A non-matching file must be excluded by the pattern.
476        std::fs::write(dir.path().join("d.txt"), "d").unwrap();
477
478        let pattern = format!("{}/*.bin", dir.path().display());
479        let mut files = collect_glob_files(&pattern).unwrap();
480        files.sort();
481
482        let names: Vec<_> = files
483            .iter()
484            .map(|p| p.file_name().unwrap().to_str().unwrap().to_string())
485            .collect();
486        assert_eq!(names, vec!["a.bin".to_string(), "b.bin".to_string()]);
487    }
488
489    #[test]
490    fn collect_glob_files_rejects_invalid_pattern() {
491        // An unclosed `[` is an invalid glob pattern; it must surface as an
492        // error, not silently return an empty file list.
493        let err = collect_glob_files("data/[.bin");
494        assert!(err.is_err(), "invalid glob pattern must error, got {err:?}");
495    }
496}
tenshift_core/sources/mod.rs

tenshift_core/sources/
mod.rs
