torsh_data/dataloader/
workers.rs

1//! Worker pool functionality for multi-process data loading
2//!
3//! This module provides worker pool implementations for parallel data loading,
4//! including both standard and persistent worker pools.
5
6use crate::{collate::Collate, dataset::Dataset, sampler::BatchSampler};
7use std::collections::HashMap;
8use std::sync::mpsc;
9use std::sync::{Arc, Mutex};
10use std::thread;
11use torsh_core::error::Result;
12
13/// Worker task containing indices to process
14#[derive(Debug, Clone)]
15struct WorkerTask {
16    task_id: usize,
17    indices: Vec<usize>,
18}
19
20/// Worker result containing processed batch
21#[derive(Debug)]
22pub struct WorkerResult<T> {
23    pub task_id: usize,
24    pub result: Result<T>,
25}
26
27/// Handle to a worker thread
28struct WorkerHandle<T> {
29    _thread: thread::JoinHandle<()>,
30    _phantom: std::marker::PhantomData<T>,
31}
32
33/// Worker process manager for multi-process data loading
34///
35/// This struct manages a pool of worker threads that process batches in parallel,
36/// allowing for efficient utilization of multiple CPU cores during data loading.
37///
38/// # Type Parameters
39///
40/// * `D` - Dataset type implementing the Dataset trait
41/// * `C` - Collate function type implementing the Collate trait
42///
43/// # Examples
44///
45/// ```rust,ignore
46/// use std::sync::Arc;
47/// use torsh_data::dataloader::workers::WorkerPool;
48/// use torsh_data::dataset::TensorDataset;
49/// use torsh_data::collate::DefaultCollate;
50///
51/// let dataset = Arc::new(TensorDataset::new(vec![1, 2, 3, 4, 5]));
52/// let collate_fn = Arc::new(DefaultCollate);
53/// let worker_pool = WorkerPool::new(dataset, collate_fn, 4);
54///
55/// // Submit tasks and collect results
56/// worker_pool.submit_task(0, vec![0, 1]).unwrap();
57/// let result = worker_pool.get_result().unwrap();
58/// ```
59pub struct WorkerPool<D, C>
60where
61    D: Dataset + Clone + Send + Sync + 'static,
62    C: Collate<D::Item> + Clone + Send + Sync + 'static,
63    D::Item: Send + 'static,
64    C::Output: Send + 'static,
65{
66    #[allow(dead_code)]
67    dataset: Arc<D>,
68    #[allow(dead_code)]
69    collate_fn: Arc<C>,
70    num_workers: usize,
71    #[allow(dead_code)]
72    workers: Vec<WorkerHandle<C::Output>>,
73    task_sender: mpsc::Sender<WorkerTask>,
74    result_receiver: mpsc::Receiver<WorkerResult<C::Output>>,
75}
76
77impl<D, C> WorkerPool<D, C>
78where
79    D: Dataset + Clone + Send + Sync + 'static,
80    C: Collate<D::Item> + Clone + Send + Sync + 'static,
81    D::Item: Send + 'static,
82    C::Output: Send + 'static,
83{
84    /// Create a new worker pool
85    ///
86    /// # Arguments
87    ///
88    /// * `dataset` - The dataset to process
89    /// * `collate_fn` - The collation function to apply to batches
90    /// * `num_workers` - Number of worker threads to spawn
91    ///
92    /// # Returns
93    ///
94    /// A new WorkerPool ready to process tasks
95    pub fn new(dataset: Arc<D>, collate_fn: Arc<C>, num_workers: usize) -> Self {
96        let (task_sender, task_receiver) = mpsc::channel::<WorkerTask>();
97        let (result_sender, result_receiver) = mpsc::channel::<WorkerResult<C::Output>>();
98
99        let task_receiver = Arc::new(Mutex::new(task_receiver));
100        let mut workers = Vec::with_capacity(num_workers);
101
102        // Spawn worker threads
103        for worker_id in 0..num_workers {
104            let dataset_clone = Arc::clone(&dataset);
105            let collate_fn_clone = Arc::clone(&collate_fn);
106            let task_receiver_clone = Arc::clone(&task_receiver);
107            let result_sender_clone = result_sender.clone();
108
109            let worker_thread = thread::spawn(move || {
110                Self::worker_loop(
111                    worker_id,
112                    dataset_clone,
113                    collate_fn_clone,
114                    task_receiver_clone,
115                    result_sender_clone,
116                );
117            });
118
119            workers.push(WorkerHandle {
120                _thread: worker_thread,
121                _phantom: std::marker::PhantomData,
122            });
123        }
124
125        Self {
126            dataset,
127            collate_fn,
128            num_workers,
129            workers,
130            task_sender,
131            result_receiver,
132        }
133    }
134
135    /// Worker thread loop
136    ///
137    /// This is the main loop that each worker thread runs to process tasks.
138    #[allow(clippy::too_many_arguments)]
139    fn worker_loop(
140        _worker_id: usize,
141        dataset: Arc<D>,
142        collate_fn: Arc<C>,
143        task_receiver: Arc<Mutex<mpsc::Receiver<WorkerTask>>>,
144        result_sender: mpsc::Sender<WorkerResult<C::Output>>,
145    ) {
146        loop {
147            // Try to receive a task
148            let task = {
149                let receiver = match task_receiver.lock() {
150                    Ok(receiver) => receiver,
151                    Err(_) => {
152                        // Mutex is poisoned, worker should exit
153                        break;
154                    }
155                };
156                receiver.recv()
157            };
158
159            match task {
160                Ok(WorkerTask { task_id, indices }) => {
161                    // Process the batch
162                    let batch_result = Self::process_batch(&*dataset, &*collate_fn, indices);
163
164                    let result = WorkerResult {
165                        task_id,
166                        result: batch_result,
167                    };
168
169                    // Send result back
170                    if result_sender.send(result).is_err() {
171                        // Main thread is no longer listening, exit
172                        break;
173                    }
174                }
175                Err(_) => {
176                    // Main thread closed the channel, exit
177                    break;
178                }
179            }
180        }
181    }
182
183    /// Process a batch of indices
184    ///
185    /// # Arguments
186    ///
187    /// * `dataset` - The dataset to load from
188    /// * `collate_fn` - The collation function to apply
189    /// * `indices` - The indices of samples to load
190    ///
191    /// # Returns
192    ///
193    /// The collated batch result
194    fn process_batch(dataset: &D, collate_fn: &C, indices: Vec<usize>) -> Result<C::Output> {
195        let mut samples = Vec::with_capacity(indices.len());
196
197        for idx in indices {
198            match dataset.get(idx) {
199                Ok(sample) => samples.push(sample),
200                Err(e) => return Err(e),
201            }
202        }
203
204        collate_fn.collate(samples)
205    }
206
207    /// Submit a task to the worker pool
208    ///
209    /// # Arguments
210    ///
211    /// * `task_id` - Unique identifier for the task
212    /// * `indices` - Indices of samples to process
213    ///
214    /// # Returns
215    ///
216    /// Result indicating success or failure of task submission
217    pub fn submit_task(&self, task_id: usize, indices: Vec<usize>) -> Result<()> {
218        let task = WorkerTask { task_id, indices };
219
220        self.task_sender.send(task).map_err(|_| {
221            torsh_core::error::TorshError::RuntimeError(
222                "Failed to send task to worker pool".to_string(),
223            )
224        })?;
225
226        Ok(())
227    }
228
229    /// Get a result from the worker pool
230    ///
231    /// This method blocks until a result is available from any worker.
232    ///
233    /// # Returns
234    ///
235    /// The next available worker result
236    pub fn get_result(&self) -> Result<WorkerResult<C::Output>> {
237        self.result_receiver.recv().map_err(|_| {
238            torsh_core::error::TorshError::RuntimeError(
239                "Failed to receive result from worker pool".to_string(),
240            )
241        })
242    }
243
244    /// Try to get a result without blocking
245    ///
246    /// # Returns
247    ///
248    /// Some(result) if a result is available, None otherwise
249    pub fn try_get_result(&self) -> Option<WorkerResult<C::Output>> {
250        self.result_receiver.try_recv().ok()
251    }
252
253    /// Get the number of workers
254    ///
255    /// # Returns
256    ///
257    /// The number of worker threads in the pool
258    pub fn num_workers(&self) -> usize {
259        self.num_workers
260    }
261}
262
263/// Multi-process DataLoader iterator
264///
265/// This iterator coordinates with a WorkerPool to process batches in parallel,
266/// maintaining a pipeline of pending tasks to maximize throughput.
267///
268/// # Type Parameters
269///
270/// * `D` - Dataset type
271/// * `S` - Sampler type
272/// * `C` - Collate function type
273pub struct MultiProcessIterator<'a, D, S, C>
274where
275    D: Dataset + Clone + Send + Sync + 'static,
276    S: BatchSampler,
277    C: Collate<D::Item> + Clone + Send + Sync + 'static,
278    D::Item: Send + 'static,
279    C::Output: Send + 'static,
280{
281    sampler_iter: S::Iter,
282    worker_pool: &'a WorkerPool<D, C>,
283    pending_tasks: HashMap<usize, Vec<usize>>,
284    next_task_id: usize,
285    max_pending: usize,
286}
287
288impl<'a, D, S, C> MultiProcessIterator<'a, D, S, C>
289where
290    D: Dataset + Clone + Send + Sync + 'static,
291    S: BatchSampler,
292    C: Collate<D::Item> + Clone + Send + Sync + 'static,
293    D::Item: Send + 'static,
294    C::Output: Send + 'static,
295{
296    /// Create a new multi-process iterator
297    ///
298    /// # Arguments
299    ///
300    /// * `sampler_iter` - Iterator over batch indices
301    /// * `worker_pool` - Worker pool to use for processing
302    ///
303    /// # Returns
304    ///
305    /// A new MultiProcessIterator
306    pub fn new(sampler_iter: S::Iter, worker_pool: &'a WorkerPool<D, C>) -> Self {
307        let max_pending = worker_pool.num_workers() * 2; // Buffer 2x the number of workers
308
309        Self {
310            sampler_iter,
311            worker_pool,
312            pending_tasks: HashMap::new(),
313            next_task_id: 0,
314            max_pending,
315        }
316    }
317
318    /// Create a new multi-process iterator with custom buffer size
319    ///
320    /// # Arguments
321    ///
322    /// * `sampler_iter` - Iterator over batch indices
323    /// * `worker_pool` - Worker pool to use for processing
324    /// * `max_pending` - Maximum number of pending tasks
325    ///
326    /// # Returns
327    ///
328    /// A new MultiProcessIterator with custom buffering
329    pub fn with_buffer_size(
330        sampler_iter: S::Iter,
331        worker_pool: &'a WorkerPool<D, C>,
332        max_pending: usize,
333    ) -> Self {
334        Self {
335            sampler_iter,
336            worker_pool,
337            pending_tasks: HashMap::new(),
338            next_task_id: 0,
339            max_pending,
340        }
341    }
342
343    /// Submit tasks to keep the pipeline full
344    fn submit_tasks(&mut self) {
345        while self.pending_tasks.len() < self.max_pending {
346            if let Some(indices) = self.sampler_iter.next() {
347                let task_id = self.next_task_id;
348                self.next_task_id += 1;
349
350                if self
351                    .worker_pool
352                    .submit_task(task_id, indices.clone())
353                    .is_ok()
354                {
355                    self.pending_tasks.insert(task_id, indices);
356                } else {
357                    break;
358                }
359            } else {
360                break;
361            }
362        }
363    }
364
365    /// Get the number of pending tasks
366    pub fn pending_count(&self) -> usize {
367        self.pending_tasks.len()
368    }
369
370    /// Check if there are any pending tasks
371    pub fn has_pending_tasks(&self) -> bool {
372        !self.pending_tasks.is_empty()
373    }
374}
375
376impl<D, S, C> Iterator for MultiProcessIterator<'_, D, S, C>
377where
378    D: Dataset + Clone + Send + Sync + 'static,
379    S: BatchSampler,
380    S::Iter: Iterator<Item = Vec<usize>>,
381    C: Collate<D::Item> + Clone + Send + Sync + 'static,
382    D::Item: Send + 'static,
383    C::Output: Send + 'static,
384{
385    type Item = Result<C::Output>;
386
387    fn next(&mut self) -> Option<Self::Item> {
388        // Submit new tasks to keep pipeline full
389        self.submit_tasks();
390
391        // If no pending tasks, we're done
392        if self.pending_tasks.is_empty() {
393            return None;
394        }
395
396        // Wait for a result
397        match self.worker_pool.get_result() {
398            Ok(WorkerResult { task_id, result }) => {
399                // Remove the completed task
400                self.pending_tasks.remove(&task_id);
401                Some(result)
402            }
403            Err(e) => Some(Err(e)),
404        }
405    }
406}
407
408/// Messages for persistent workers
409#[derive(Debug)]
410enum PersistentWorkerMessage {
411    Task(WorkerTask),
412    Shutdown,
413}
414
415/// Handle to a persistent worker
416struct PersistentWorkerHandle {
417    _thread: thread::JoinHandle<()>,
418}
419
420/// Persistent worker pool that keeps workers alive across epochs
421///
422/// Unlike the regular WorkerPool, this implementation keeps worker threads alive
423/// between epochs, reducing the overhead of thread creation and destruction.
424/// This is particularly useful for training scenarios with multiple epochs.
425///
426/// # Examples
427///
428/// ```rust,ignore
429/// use std::sync::Arc;
430/// use torsh_data::dataloader::workers::PersistentWorkerPool;
431/// use torsh_data::dataset::TensorDataset;
432/// use torsh_data::collate::DefaultCollate;
433///
434/// let dataset = Arc::new(TensorDataset::new(vec![1, 2, 3, 4, 5]));
435/// let collate_fn = Arc::new(DefaultCollate);
436/// let worker_pool = PersistentWorkerPool::new(dataset, collate_fn, 4);
437///
438/// // Use across multiple epochs
439/// for epoch in 0..10 {
440///     worker_pool.reset_for_epoch().unwrap();
441///     // Submit tasks for this epoch
442/// }
443///
444/// // Shutdown when done
445/// worker_pool.shutdown().unwrap();
446/// ```
447pub struct PersistentWorkerPool<D, C>
448where
449    D: Dataset + Clone + Send + Sync + 'static,
450    C: Collate<D::Item> + Clone + Send + Sync + 'static,
451    D::Item: Send + 'static,
452    C::Output: Send + 'static,
453{
454    #[allow(dead_code)]
455    dataset: Arc<D>,
456    #[allow(dead_code)]
457    collate_fn: Arc<C>,
458    num_workers: usize,
459    #[allow(dead_code)]
460    workers: Vec<PersistentWorkerHandle>,
461    task_sender: mpsc::Sender<PersistentWorkerMessage>,
462    result_receiver: mpsc::Receiver<WorkerResult<C::Output>>,
463    is_shutdown: Arc<std::sync::atomic::AtomicBool>,
464}
465
466impl<D, C> PersistentWorkerPool<D, C>
467where
468    D: Dataset + Clone + Send + Sync + 'static,
469    C: Collate<D::Item> + Clone + Send + Sync + 'static,
470    D::Item: Send + 'static,
471    C::Output: Send + 'static,
472{
473    /// Create a new persistent worker pool
474    ///
475    /// # Arguments
476    ///
477    /// * `dataset` - The dataset to process
478    /// * `collate_fn` - The collation function to apply to batches
479    /// * `num_workers` - Number of worker threads to spawn
480    ///
481    /// # Returns
482    ///
483    /// A new PersistentWorkerPool ready to process tasks
484    pub fn new(dataset: Arc<D>, collate_fn: Arc<C>, num_workers: usize) -> Self {
485        let (task_sender, task_receiver) = mpsc::channel::<PersistentWorkerMessage>();
486        let (result_sender, result_receiver) = mpsc::channel::<WorkerResult<C::Output>>();
487
488        let task_receiver = Arc::new(Mutex::new(task_receiver));
489        let is_shutdown = Arc::new(std::sync::atomic::AtomicBool::new(false));
490        let mut workers = Vec::with_capacity(num_workers);
491
492        // Spawn persistent worker threads
493        for worker_id in 0..num_workers {
494            let dataset_clone = Arc::clone(&dataset);
495            let collate_fn_clone = Arc::clone(&collate_fn);
496            let task_receiver_clone = Arc::clone(&task_receiver);
497            let result_sender_clone = result_sender.clone();
498            let is_shutdown_clone = Arc::clone(&is_shutdown);
499
500            let worker_thread = thread::spawn(move || {
501                Self::persistent_worker_loop(
502                    worker_id,
503                    dataset_clone,
504                    collate_fn_clone,
505                    task_receiver_clone,
506                    result_sender_clone,
507                    is_shutdown_clone,
508                );
509            });
510
511            workers.push(PersistentWorkerHandle {
512                _thread: worker_thread,
513            });
514        }
515
516        Self {
517            dataset,
518            collate_fn,
519            num_workers,
520            workers,
521            task_sender,
522            result_receiver,
523            is_shutdown,
524        }
525    }
526
527    /// Persistent worker thread loop that stays alive
528    ///
529    /// This loop includes timeout-based shutdown checking to ensure workers
530    /// can be cleanly terminated even if no tasks are pending.
531    #[allow(clippy::too_many_arguments)]
532    fn persistent_worker_loop(
533        _worker_id: usize,
534        dataset: Arc<D>,
535        collate_fn: Arc<C>,
536        task_receiver: Arc<Mutex<mpsc::Receiver<PersistentWorkerMessage>>>,
537        result_sender: mpsc::Sender<WorkerResult<C::Output>>,
538        is_shutdown: Arc<std::sync::atomic::AtomicBool>,
539    ) {
540        loop {
541            // Check for shutdown signal
542            if is_shutdown.load(std::sync::atomic::Ordering::Relaxed) {
543                break;
544            }
545
546            // Try to receive a message with timeout to allow periodic shutdown checks
547            let message = {
548                let receiver = match task_receiver.lock() {
549                    Ok(receiver) => receiver,
550                    Err(_) => {
551                        // Mutex is poisoned, worker should exit
552                        break;
553                    }
554                };
555                receiver.recv_timeout(std::time::Duration::from_millis(100))
556            };
557
558            match message {
559                Ok(PersistentWorkerMessage::Task(WorkerTask { task_id, indices })) => {
560                    // Process the batch
561                    let batch_result = Self::process_batch(&*dataset, &*collate_fn, indices);
562
563                    let result = WorkerResult {
564                        task_id,
565                        result: batch_result,
566                    };
567
568                    // Send result back
569                    if result_sender.send(result).is_err() {
570                        // Main thread is no longer listening, exit
571                        break;
572                    }
573                }
574                Ok(PersistentWorkerMessage::Shutdown) => {
575                    // Explicit shutdown request
576                    break;
577                }
578                Err(mpsc::RecvTimeoutError::Timeout) => {
579                    // Timeout occurred, continue the loop to check for shutdown
580                    continue;
581                }
582                Err(mpsc::RecvTimeoutError::Disconnected) => {
583                    // Main thread closed the channel, exit
584                    break;
585                }
586            }
587        }
588    }
589
590    /// Process a batch of indices (same as WorkerPool)
591    fn process_batch(dataset: &D, collate_fn: &C, indices: Vec<usize>) -> Result<C::Output> {
592        let mut samples = Vec::with_capacity(indices.len());
593
594        for idx in indices {
595            match dataset.get(idx) {
596                Ok(sample) => samples.push(sample),
597                Err(e) => return Err(e),
598            }
599        }
600
601        collate_fn.collate(samples)
602    }
603
604    /// Submit a task to the persistent worker pool
605    ///
606    /// # Arguments
607    ///
608    /// * `task_id` - Unique identifier for the task
609    /// * `indices` - Indices of samples to process
610    ///
611    /// # Returns
612    ///
613    /// Result indicating success or failure of task submission
614    pub fn submit_task(&self, task_id: usize, indices: Vec<usize>) -> Result<()> {
615        let message = PersistentWorkerMessage::Task(WorkerTask { task_id, indices });
616
617        self.task_sender.send(message).map_err(|_| {
618            torsh_core::error::TorshError::RuntimeError(
619                "Failed to send task to persistent worker pool".to_string(),
620            )
621        })?;
622
623        Ok(())
624    }
625
626    /// Get a result from the persistent worker pool
627    ///
628    /// This method blocks until a result is available from any worker.
629    ///
630    /// # Returns
631    ///
632    /// The next available worker result
633    pub fn get_result(&self) -> Result<WorkerResult<C::Output>> {
634        self.result_receiver.recv().map_err(|_| {
635            torsh_core::error::TorshError::RuntimeError(
636                "Failed to receive result from persistent worker pool".to_string(),
637            )
638        })
639    }
640
641    /// Get a result with timeout
642    ///
643    /// # Arguments
644    ///
645    /// * `timeout` - Maximum time to wait for a result
646    ///
647    /// # Returns
648    ///
649    /// The next available worker result or a timeout error
650    pub fn get_result_timeout(
651        &self,
652        timeout: std::time::Duration,
653    ) -> Result<WorkerResult<C::Output>> {
654        self.result_receiver
655            .recv_timeout(timeout)
656            .map_err(|e| match e {
657                mpsc::RecvTimeoutError::Timeout => torsh_core::error::TorshError::RuntimeError(
658                    "Timeout waiting for result from persistent worker pool".to_string(),
659                ),
660                mpsc::RecvTimeoutError::Disconnected => {
661                    torsh_core::error::TorshError::RuntimeError(
662                        "Persistent worker pool disconnected".to_string(),
663                    )
664                }
665            })
666    }
667
668    /// Try to get a result without blocking
669    ///
670    /// # Returns
671    ///
672    /// Some(result) if a result is available, None otherwise
673    pub fn try_get_result(&self) -> Option<WorkerResult<C::Output>> {
674        self.result_receiver.try_recv().ok()
675    }
676
677    /// Get the number of workers
678    ///
679    /// # Returns
680    ///
681    /// The number of worker threads in the pool
682    pub fn num_workers(&self) -> usize {
683        self.num_workers
684    }
685
686    /// Check if the pool is shutdown
687    ///
688    /// # Returns
689    ///
690    /// True if the pool has been shutdown, false otherwise
691    pub fn is_shutdown(&self) -> bool {
692        self.is_shutdown.load(std::sync::atomic::Ordering::Relaxed)
693    }
694
695    /// Shutdown the persistent worker pool
696    ///
697    /// This method gracefully shuts down all worker threads and should be called
698    /// when the pool is no longer needed.
699    ///
700    /// # Returns
701    ///
702    /// Result indicating success or failure of shutdown operation
703    pub fn shutdown(&self) -> Result<()> {
704        // Set shutdown flag
705        self.is_shutdown
706            .store(true, std::sync::atomic::Ordering::Relaxed);
707
708        // Send shutdown messages to all workers
709        for _ in 0..self.num_workers {
710            if self
711                .task_sender
712                .send(PersistentWorkerMessage::Shutdown)
713                .is_err()
714            {
715                // Channel is closed, workers are already shutdown
716                break;
717            }
718        }
719
720        Ok(())
721    }
722
723    /// Reset the pool for a new epoch (clears any pending tasks)
724    ///
725    /// This method can be called between epochs to reset the pool state.
726    /// In the current implementation, this is a placeholder for future
727    /// epoch-specific optimizations.
728    ///
729    /// # Returns
730    ///
731    /// Result indicating success or failure of reset operation
732    pub fn reset_for_epoch(&self) -> Result<()> {
733        // For a full implementation, you might want to:
734        // 1. Clear any pending tasks
735        // 2. Reset worker state
736        // 3. Optionally restart workers if needed
737
738        // For now, this is a placeholder
739        Ok(())
740    }
741}
742
743impl<D, C> Drop for PersistentWorkerPool<D, C>
744where
745    D: Dataset + Clone + Send + Sync + 'static,
746    C: Collate<D::Item> + Clone + Send + Sync + 'static,
747    D::Item: Send + 'static,
748    C::Output: Send + 'static,
749{
750    fn drop(&mut self) {
751        // Ensure workers are shutdown when pool is dropped
752        let _ = self.shutdown();
753    }
754}
755
756/// Configuration for worker pools
757#[derive(Debug, Clone)]
758pub struct WorkerConfig {
759    /// Number of worker threads
760    pub num_workers: usize,
761    /// Maximum number of pending tasks for multi-process iterator
762    pub max_pending_tasks: Option<usize>,
763    /// Whether to use persistent workers
764    pub persistent: bool,
765}
766
767impl Default for WorkerConfig {
768    fn default() -> Self {
769        Self {
770            num_workers: std::thread::available_parallelism()
771                .map(|n| n.get())
772                .unwrap_or(1),
773            max_pending_tasks: None,
774            persistent: false,
775        }
776    }
777}
778
779impl WorkerConfig {
780    /// Create a new worker configuration
781    pub fn new() -> Self {
782        Self::default()
783    }
784
785    /// Set the number of workers
786    pub fn num_workers(mut self, num_workers: usize) -> Self {
787        self.num_workers = num_workers;
788        self
789    }
790
791    /// Set the maximum number of pending tasks
792    pub fn max_pending_tasks(mut self, max_pending: usize) -> Self {
793        self.max_pending_tasks = Some(max_pending);
794        self
795    }
796
797    /// Enable persistent workers
798    pub fn persistent(mut self, persistent: bool) -> Self {
799        self.persistent = persistent;
800        self
801    }
802}
803
804/// Utility functions for worker management
805pub mod utils {
806    use super::*;
807
808    /// Determine optimal number of workers based on system capabilities
809    ///
810    /// # Arguments
811    ///
812    /// * `cpu_intensive` - Whether the workload is CPU-intensive
813    ///
814    /// # Returns
815    ///
816    /// Recommended number of worker threads
817    pub fn optimal_worker_count(cpu_intensive: bool) -> usize {
818        let cpu_count = std::thread::available_parallelism()
819            .map(|n| n.get())
820            .unwrap_or(1);
821
822        if cpu_intensive {
823            // For CPU-intensive tasks, use fewer workers to avoid over-subscription
824            (cpu_count * 3 / 4).max(1)
825        } else {
826            // For I/O-bound tasks, can use more workers
827            cpu_count * 2
828        }
829    }
830
831    /// Create a worker configuration optimized for training
832    ///
833    /// # Returns
834    ///
835    /// A WorkerConfig optimized for training workloads
836    pub fn training_config() -> WorkerConfig {
837        WorkerConfig::new()
838            .num_workers(optimal_worker_count(false))
839            .persistent(true)
840            .max_pending_tasks(optimal_worker_count(false) * 3)
841    }
842
843    /// Create a worker configuration optimized for inference
844    ///
845    /// # Returns
846    ///
847    /// A WorkerConfig optimized for inference workloads
848    pub fn inference_config() -> WorkerConfig {
849        WorkerConfig::new()
850            .num_workers(optimal_worker_count(true))
851            .persistent(false)
852            .max_pending_tasks(optimal_worker_count(true) * 2)
853    }
854}
855
856#[cfg(test)]
857mod tests {
858    use super::*;
859    use crate::{collate::DefaultCollate, dataset::TensorDataset};
860
861    #[test]
862    fn test_worker_pool_creation() {
863        use torsh_core::device::DeviceType;
864        use torsh_tensor::Tensor;
865
866        let tensor = Tensor::from_data(vec![1.0f32, 2.0, 3.0, 4.0, 5.0], vec![5], DeviceType::Cpu)
867            .expect("Tensor should succeed");
868        let dataset = Arc::new(TensorDataset::from_tensor(tensor));
869        let collate_fn = Arc::new(DefaultCollate);
870        let worker_pool = WorkerPool::new(dataset, collate_fn, 2);
871
872        assert_eq!(worker_pool.num_workers(), 2);
873    }
874
875    #[test]
876    fn test_worker_pool_task_submission() {
877        use torsh_core::device::DeviceType;
878        use torsh_tensor::Tensor;
879
880        let tensor = Tensor::from_data(vec![1.0f32, 2.0, 3.0, 4.0, 5.0], vec![5], DeviceType::Cpu)
881            .expect("Tensor should succeed");
882        let dataset = Arc::new(TensorDataset::from_tensor(tensor));
883        let collate_fn = Arc::new(DefaultCollate);
884        let worker_pool = WorkerPool::new(dataset, collate_fn, 2);
885
886        assert!(worker_pool.submit_task(0, vec![0, 1]).is_ok());
887    }
888
889    #[test]
890    fn test_persistent_worker_pool_creation() {
891        use torsh_core::device::DeviceType;
892        use torsh_tensor::Tensor;
893
894        let tensor = Tensor::from_data(vec![1.0f32, 2.0, 3.0, 4.0, 5.0], vec![5], DeviceType::Cpu)
895            .expect("Tensor should succeed");
896        let dataset = Arc::new(TensorDataset::from_tensor(tensor));
897        let collate_fn = Arc::new(DefaultCollate);
898        let worker_pool = PersistentWorkerPool::new(dataset, collate_fn, 2);
899
900        assert_eq!(worker_pool.num_workers(), 2);
901        assert!(!worker_pool.is_shutdown());
902    }
903
904    #[test]
905    fn test_persistent_worker_pool_shutdown() {
906        use torsh_core::device::DeviceType;
907        use torsh_tensor::Tensor;
908
909        let tensor = Tensor::from_data(vec![1.0f32, 2.0, 3.0, 4.0, 5.0], vec![5], DeviceType::Cpu)
910            .expect("Tensor should succeed");
911        let dataset = Arc::new(TensorDataset::from_tensor(tensor));
912        let collate_fn = Arc::new(DefaultCollate);
913        let worker_pool = PersistentWorkerPool::new(dataset, collate_fn, 2);
914
915        assert!(worker_pool.shutdown().is_ok());
916        assert!(worker_pool.is_shutdown());
917    }
918
919    #[test]
920    fn test_worker_config() {
921        let config = WorkerConfig::new()
922            .num_workers(4)
923            .max_pending_tasks(8)
924            .persistent(true);
925
926        assert_eq!(config.num_workers, 4);
927        assert_eq!(config.max_pending_tasks, Some(8));
928        assert!(config.persistent);
929    }
930
931    #[test]
932    fn test_optimal_worker_count() {
933        let cpu_intensive_count = utils::optimal_worker_count(true);
934        let io_bound_count = utils::optimal_worker_count(false);
935
936        assert!(cpu_intensive_count > 0);
937        assert!(io_bound_count > 0);
938        assert!(io_bound_count >= cpu_intensive_count);
939    }
940
941    #[test]
942    fn test_training_config() {
943        let config = utils::training_config();
944
945        assert!(config.num_workers > 0);
946        assert!(config.persistent);
947        assert!(config.max_pending_tasks.is_some());
948    }
949
950    #[test]
951    fn test_inference_config() {
952        let config = utils::inference_config();
953
954        assert!(config.num_workers > 0);
955        assert!(!config.persistent);
956        assert!(config.max_pending_tasks.is_some());
957    }
958}
torsh_data/dataloader/workers.rs

torsh_data/dataloader/
workers.rs
