Struct WorkerPool 

pub struct WorkerPool<D, C>
where
    D: Dataset + Clone + Send + Sync + 'static,
    C: Collate<D::Item> + Clone + Send + Sync + 'static,
    D::Item: Send + 'static,
    C::Output: Send + 'static,
{ /* private fields */ }

Worker process manager for multi-process data loading

This struct manages a pool of worker threads that process batches in parallel, allowing for efficient utilization of multiple CPU cores during data loading.

Type Parameters

• D - Dataset type implementing the Dataset trait
• C - Collate function type implementing the Collate trait

Examples

use std::sync::Arc;
use torsh_data::dataloader::workers::WorkerPool;
use torsh_data::dataset::TensorDataset;
use torsh_data::collate::DefaultCollate;

let dataset = Arc::new(TensorDataset::new(vec![1, 2, 3, 4, 5]));
let collate_fn = Arc::new(DefaultCollate);
let worker_pool = WorkerPool::new(dataset, collate_fn, 4);

// Submit tasks and collect results
worker_pool.submit_task(0, vec![0, 1]).unwrap();
let result = worker_pool.get_result().unwrap();

Implementations

impl<D, C> WorkerPool<D, C>

where D: Dataset + Clone + Send + Sync + 'static, C: Collate<D::Item> + Clone + Send + Sync + 'static, D::Item: Send + 'static, C::Output: Send + 'static,

pub fn new(dataset: Arc<D>, collate_fn: Arc<C>, num_workers: usize) -> Self

Create a new worker pool

Arguments

• dataset - The dataset to process
• collate_fn - The collation function to apply to batches
• num_workers - Number of worker threads to spawn

Returns

A new WorkerPool ready to process tasks

pub fn submit_task(&self, task_id: usize, indices: Vec<usize>) -> Result<()>

Submit a task to the worker pool

Arguments

• task_id - Unique identifier for the task
• indices - Indices of samples to process

Returns

Result indicating success or failure of task submission

pub fn get_result(&self) -> Result<WorkerResult<C::Output>>

Get a result from the worker pool

This method blocks until a result is available from any worker.

Returns

The next available worker result

pub fn try_get_result(&self) -> Option<WorkerResult<C::Output>>

Try to get a result without blocking

Returns

Some(result) if a result is available, None otherwise

pub fn num_workers(&self) -> usize

Get the number of workers

Returns

The number of worker threads in the pool