tenshift-core 0.1.2

Thread-safe, backpressure-aware data loading pipeline for iterative processing
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109

//! Threaded executor for running a built tenshift pipeline.
//!
//! Once the builder has described the architecture, this module materializes
//! the source thread, worker pool, and collector thread and wires their bounded
//! channels together into a running iterator.

#![allow(clippy::module_name_repetitions)]
use super::{
    collector, source_thread, tuned_channel_capacities, validate_stage_order, worker, wrap_source,
    Pipeline, PipelineIterator, SampleChunk, Stage,
};
use crate::error::Result;
use crate::sample::Sample;
use crate::transform::Transform;
use crossbeam_channel::bounded;
use std::sync::atomic::{AtomicBool, AtomicU64};
use std::sync::Arc;
use std::time::Instant;

#[allow(clippy::too_many_lines)]
pub(crate) fn start_pipeline(pipeline: Pipeline) -> Result<PipelineIterator> {
    validate_stage_order(&pipeline.stages, &pipeline.collate_mode)?;
    let source = wrap_source(pipeline.source, pipeline.config.shard)?;

    let shutdown = Arc::new(AtomicBool::new(false));
    let errors_skipped = Arc::new(AtomicU64::new(0));

    let mut stateless: Vec<Box<dyn Transform>> = Vec::new();
    let mut collector_stages = Vec::new();
    for stage in pipeline.stages {
        match stage {
            Stage::Stateless(transform) => stateless.push(transform),
            Stage::Shuffle(buffer_size) => collector_stages.push(Stage::Shuffle(buffer_size)),
            Stage::Batch(batch_size) => collector_stages.push(Stage::Batch(batch_size)),
        }
    }

    let stateless = Arc::new(stateless);
    let chunk = pipeline.config.channel_chunk_size;
    let capacities =
        tuned_channel_capacities(pipeline.config.prefetch_size, pipeline.config.num_workers);
    let output_paused = Arc::new(AtomicBool::new(false));
    let (raw_tx, raw_rx) = bounded::<SampleChunk>(capacities.raw_chunks);
    let (proc_tx, proc_rx) = bounded::<SampleChunk>(capacities.processed_chunks);
    let (out_tx, out_rx) = bounded::<Vec<Sample>>(capacities.output_batches);

    let mut all_handles = Vec::new();

    let source = Arc::from(source);
    let fatal_error = Arc::new(std::sync::Mutex::new(None));

    let source_handle = source_thread::spawn_source_thread(
        source,
        raw_tx.clone(),
        chunk,
        pipeline.config.epochs,
        pipeline.config.source_timeout,
        pipeline.config.on_error,
        pipeline.config.test_start_sequence,
        Arc::clone(&output_paused),
        Arc::clone(&shutdown),
        Arc::clone(&errors_skipped),
        Arc::clone(&fatal_error),
    )?;
    all_handles.push(source_handle);

    let mut worker_handles = worker::spawn_worker_threads(
        pipeline.config.num_workers,
        pipeline.config.pin_threads,
        pipeline.config.on_error,
        raw_rx.clone(),
        proc_tx.clone(),
        stateless,
        Arc::clone(&shutdown),
        Arc::clone(&errors_skipped),
        Arc::clone(&fatal_error),
    )?;
    all_handles.append(&mut worker_handles);

    drop(raw_rx);
    drop(proc_tx);

    let collector_handle = collector::spawn_collector_thread(
        collector_stages,
        pipeline.collate_mode,
        pipeline.config.on_error,
        pipeline.config.seed,
        pipeline.config.pending_sequence_limit,
        pipeline.config.sequence_gap_timeout,
        pipeline.config.drop_last,
        proc_rx.clone(),
        out_tx.clone(),
        Arc::clone(&output_paused),
        Arc::clone(&shutdown),
        Arc::clone(&errors_skipped),
        Arc::clone(&fatal_error),
    )?;
    all_handles.push(collector_handle);

    Ok(PipelineIterator {
        receiver: out_rx,
        shutdown,
        workers: all_handles,
        items_yielded: 0,
        errors_skipped,
        started_at: Instant::now(),
        fatal_error,
    })
}