Struct PersistentWorkerPool 

pub struct PersistentWorkerPool<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 */ }

Persistent worker pool that keeps workers alive across epochs

Unlike the regular WorkerPool, this implementation keeps worker threads alive between epochs, reducing the overhead of thread creation and destruction. This is particularly useful for training scenarios with multiple epochs.

Examples

use std::sync::Arc;
use torsh_data::dataloader::workers::PersistentWorkerPool;
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 = PersistentWorkerPool::new(dataset, collate_fn, 4);

// Use across multiple epochs
for epoch in 0..10 {
    worker_pool.reset_for_epoch().unwrap();
    // Submit tasks for this epoch
}

// Shutdown when done
worker_pool.shutdown().unwrap();

Implementations

impl<D, C> PersistentWorkerPool<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 persistent 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 PersistentWorkerPool ready to process tasks

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

Submit a task to the persistent 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 persistent worker pool

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

Returns

The next available worker result

pub fn get_result_timeout( &self, timeout: Duration, ) -> Result<WorkerResult<C::Output>>

Get a result with timeout

Arguments

• timeout - Maximum time to wait for a result

Returns

The next available worker result or a timeout error

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

pub fn is_shutdown(&self) -> bool

Check if the pool is shutdown

Returns

True if the pool has been shutdown, false otherwise

pub fn shutdown(&self) -> Result<()>

Shutdown the persistent worker pool

This method gracefully shuts down all worker threads and should be called when the pool is no longer needed.

Returns

Result indicating success or failure of shutdown operation

pub fn reset_for_epoch(&self) -> Result<()>

Reset the pool for a new epoch (clears any pending tasks)

This method can be called between epochs to reset the pool state. In the current implementation, this is a placeholder for future epoch-specific optimizations.

Returns

Result indicating success or failure of reset operation

Trait Implementations

impl<D, C> Drop for PersistentWorkerPool<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,

fn drop(&mut self)

Executes the destructor for this type.

fn pin_drop(self: Pin<&mut Self>)

🔬This is a nightly-only experimental API. (pin_ergonomics)

Execute the destructor for this type, but different to Drop::drop, it requires self to be pinned.