tenshift-core 0.1.2

Thread-safe, backpressure-aware data loading pipeline for iterative processing
Documentation
//! Runtime configuration for tenshift pipeline execution.
//!
//! These settings control how the source thread, worker pool, and collector
//! cooperate. The module sits at the boundary between the public builder API
//! and the executor's concrete threading behavior.
//!
//! # Auto-Scaling Heuristics
//!
//! The default configuration automatically adapts to the execution environment:
//!
//! ## Worker Count
//! **Default:** `num_cpus().min(8)`
//!
//! Worker threads execute stateless transforms in parallel. The default caps
//! at 8 because:
//! - Most ML preprocessing is memory-bandwidth-bound, not CPU-bound
//! - Beyond 8 workers, contention on memory channels often reduces throughput
//! - Users with truly CPU-bound transforms can override with `.workers(n)`
//!
//! ## Prefetch Size
//! **Default:** `num_workers * 2`, minimum 8
//!
//! The prefetch buffer sits between workers and collector. Its size balances:
//! - **Worker starvation prevention:** Workers produce chunks faster than the
//!   collector consumes them during bursts. Without sufficient prefetch,
//!   workers block on full channels and stall.
//! - **Memory bounds:** Each prefetch slot holds one full chunk. A fixed large
//!   buffer would OOM on high-core-count systems.
//!
//! The `* 2` multiplier provides one full chunk of headroom per worker,
//! ensuring continuous work even with slight consumer jitter.
//!
//! ## Channel Chunk Size
//! **Default:** 64 samples
//!
//! Samples flow through internal channels in chunks of 64. This amortizes:
//! - **Channel synchronization overhead:** Sending 64 samples costs ~same as 1
//! - **Cache locality:** Sequential processing within chunks improves hit rates
//! - **Backpressure timing:** 64-sample granularity keeps OOM prevention responsive

#![allow(clippy::module_name_repetitions)]

use super::num_cpus;
use std::time::Duration;

/// The default random seed used for deterministic operations (e.g. shuffling).
pub const DEFAULT_SHUFFLE_SEED: u64 = 0x517c_c1b7_2722_0a95;

/// Maximum file size to load (256MB).
pub const MAX_LOAD_FILE_SIZE: u64 = 256 * 1024 * 1024;

/// Configuration for the pipeline.
///
/// Created via [`PipelineConfig::default()`](Default) or modified through
/// the [`Pipeline`](crate::Pipeline) builder methods.
#[derive(Debug, Clone)]
pub struct PipelineConfig {
    /// Number of worker threads for parallel data loading and transforms.
    pub num_workers: usize,
    /// Size of the prefetch buffer (number of items held in memory).
    ///
    /// Default: `num_workers * 2` with a minimum of 8. See module-level docs
    /// for the rationale behind this heuristic.
    pub prefetch_size: usize,
    /// What to do when a source item or transform fails.
    pub on_error: ErrorPolicy,
    /// Optional seed for deterministic shuffling.
    pub seed: Option<u64>,
    /// Number of epochs (passes over the data). 0 = infinite.
    pub epochs: usize,
    /// Number of samples to chunk together in internal channels.
    ///
    /// Default: 64. Larger chunks amortize synchronization overhead but use
    /// more memory. See module-level docs for the rationale.
    pub channel_chunk_size: usize,
    /// Maximum number of out-of-order processed chunks buffered while waiting for a missing sequence.
    pub pending_sequence_limit: usize,
    /// Maximum time to wait for a missing processed sequence before skipping it.
    pub sequence_gap_timeout: Duration,
    /// Maximum time to wait for the source to produce a new item before shutting down.
    pub source_timeout: Option<Duration>,
    /// Whether to drop the final incomplete batch instead of flushing it.
    pub drop_last: bool,
    /// Pin worker threads to specific physical CPU cores to eliminate OS scheduler migration.
    pub pin_threads: bool,
    pub(crate) shard: Option<(usize, usize)>,
    #[doc(hidden)]
    pub(crate) test_start_sequence: u64,
}

/// What to do when a data loading error occurs.
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ErrorPolicy {
    /// Skip the bad item, log a warning, continue.
    Skip,
    /// Stop the entire pipeline on the first error.
    Fail,
}

impl Default for PipelineConfig {
    fn default() -> Self {
        const FALLBACK_WORKERS: usize = 4;
        let workers = num_cpus().unwrap_or_else(|error| {
            tracing::warn!(
                "tenshift: could not detect available parallelism ({error});                  defaulting to {FALLBACK_WORKERS} workers. Override with PipelineConfig::workers()."
            );
            FALLBACK_WORKERS
        });
        Self {
            num_workers: workers,
            // Scale prefetch with worker count to keep all workers fed.
            // Old default of 8 caused stalls on 64+ core machines.
            // See module-level docs for the full rationale.
            prefetch_size: workers.saturating_mul(2).max(8),
            on_error: ErrorPolicy::Skip,
            seed: None,
            epochs: 1,
            // Larger chunks amortize channel synchronization overhead.
            // 64 is the sweet spot for ML workloads (see module docs).
            channel_chunk_size: 64,
            pending_sequence_limit: 1000,
            sequence_gap_timeout: Duration::from_secs(30),
            source_timeout: None,
            drop_last: false,
            pin_threads: false,
            shard: None,
            test_start_sequence: 0,
        }
    }
}

impl PipelineConfig {
    /// Automatically set prefetch size to workers * 2.
    ///
    /// Callers can use `.prefetch_auto()` instead of `.prefetch(2)` for CPU-bound pipelines.
    ///
    /// # Rationale
    ///
    /// This recalculates the prefetch size based on the current worker count,
    /// ensuring the buffer scales appropriately when workers are customized.
    #[must_use]
    pub fn prefetch_auto(mut self) -> Self {
        self.prefetch_size = self.num_workers.saturating_mul(2).max(2);
        self
    }
}