olympipe-rs

Parallel data pipeline library for Rust, inspired by olympipe.

Each stage runs in its own thread (or several threads for task(..., count)), communicating via bounded crossbeam channels. Back-pressure is automatic: a slow downstream stage stalls upstream producers.

The default channel holds 4 messages. A message is a single item when the buffer is 1 or 2. With a larger buffer, a stage packs up to 64 items that it already has into one message, and flushes a partial message as soon as the next input would block (a trickle source, the last items, or a slow callback). A full channel therefore holds at most buffer_size items for a one-item packet, and at most buffer_size * 64 items otherwise (256 with the default buffer of 4). Worker threads are named olympipe-<stage>.

.with_max_packet(1) opts out of that packing for every stage added afterwards. The feeder honours it too when you set it before the first stage or terminal call (Pipeline::new(...).with_max_packet(1)), because the feeder does not pull the source until then. After the feeder has started, a later with_max_packet only affects subsequent stages. The feeder channel is sized by new / new_with_buffer and holds at most buffer_size * packet items. with_max_packet(1) restores the old per-item messages: a channel holds at most buffer_size items. new_with_buffer keeps the default packet size (64 when the buffer is larger than 2). Call with_max_packet after with_buffer; with_buffer resets the packet size of later stages to that default and does not change the feeder.

Use case: ML inference pipelines

A common bottleneck in machine learning workloads is keeping the GPU busy. olympipe-rs makes it straightforward to overlap CPU-bound preprocessing and postprocessing with GPU inference in separate threads:

use olympipe_rs::Pipeline;

fn load_and_preprocess(path: &str) -> Vec<u8> {
    path.as_bytes().to_vec()
}
fn run_inference(tensor: Vec<u8>) -> Vec<u8> {
    tensor
}
fn postprocess(output: Vec<u8>) -> Vec<u8> {
    output
}
fn save(_result: Vec<u8>) {}

let image_paths = ["frame.png"];
Pipeline::new(image_paths)
    .with_max_packet(1) // large tensors: one item per message, not 64
    .task(load_and_preprocess, 4) // 4 CPU threads decode & normalise
    .task(run_inference, 1) // 1 GPU thread, always fed
    .task(postprocess, 4) // 4 CPU threads decode results
    .for_each(save);

Because each stage runs independently and channels act as bounded buffers, the GPU thread is never starved waiting for preprocessing, and preprocessed batches don't pile up unboundedly in memory. .with_max_packet(1) keeps that bound at buffer_size items (4 by default) instead of buffer_size * 64, which matters when each item is an image or a tensor.

Quick start

[dependencies]
olympipe-rs = "1"

use olympipe_rs::Pipeline;

let result = Pipeline::new(0..100)
    .filter(|x| x % 2 == 0)
    .task(|x: i32| x * x, 4) // 4 parallel workers
    .limit(10)
    .collect();
assert_eq!(result.len(), 10);

Benchmarks

Comparisons with Python olympipe 2.1.0, and the before/after optimization figures, are in BENCHMARKS.md. cargo bench runs the criterion harness; cargo run --release --example bench_report prints the cross-language report workloads.

API

Constructors

Stages (consume self, return a new Pipeline)

Terminal methods

Split / Gather

split writes each branch to a bounded channel (capacity = buffer size in messages, default 4). A branch that nobody reads applies back-pressure on the router, so both branches must be consumed concurrently (std::thread::scope, or gather). Draining one branch to completion before starting the other can block the router once a buffer fills and stall the pipeline. Dropping a branch disconnects it: the router keeps feeding the other branch and does not block on the dead one.

use olympipe_rs::Pipeline;
use std::thread;

let (evens, odds) = Pipeline::new(0..10i32).split(|x| {
    if x % 2 == 0 { (Some(x), None) } else { (None, Some(x)) }
});

thread::scope(|scope| {
    let odds = scope.spawn(|| odds.collect());
    let mut e = evens.collect();
    let mut o = odds.join().unwrap();
    e.sort();
    o.sort();
    assert_eq!(e, vec![0, 2, 4, 6, 8]);
    assert_eq!(o, vec![1, 3, 5, 7, 9]);
});

split_unbounded keeps the old unbounded outputs (no back-pressure). It is the opt-in for callers that need to drain one branch fully before the other; the router can retain the whole input in memory.

use olympipe_rs::Pipeline;

let (evens, odds) = Pipeline::new(0..10i32).split_unbounded(|x| {
    if x % 2 == 0 {
        (Some(x), None)
    } else {
        (None, Some(x))
    }
});
let e = evens.collect(); // sequential consumption is safe
let o = odds.collect();
assert_eq!(e, vec![0, 2, 4, 6, 8]);
assert_eq!(o, vec![1, 3, 5, 7, 9]);

use olympipe_rs::Pipeline;

let p1 = Pipeline::new(vec![1i32, 2]);
let p2 = Pipeline::new(vec![3i32, 4]);
let mut result = p1.gather(vec![p2]).collect();
result.sort();
assert_eq!(result, vec![1, 2, 3, 4]);

Error handling

A panic in any worker, including the feeder that pulls the source iterator, fails the whole pipeline. The first failure is kept as PipelineError::WorkerPanicked { stage, message } (stage is "task", "filter", "split", "reduce", "checkpoint", "feeder", …). The pipeline is cancelled: every stage checks a relaxed atomic between items and stops, and channel disconnection unblocks anyone waiting on a send or recv. .timeout(duration) records PipelineError::Timeout the same way (it does not only log).

try_collect, try_for_each, and try_wait_for_completion drain, join every worker, and return that error. collect, for_each, and wait_for_completion keep their old signatures and panic with olympipe-rs pipeline failed: {err} — the message includes the original panic text — so a stage failure reaches the caller the way an exception does in Python olympipe. The library does not print the panic itself; the default panic hook already does.

task_or is unchanged and is the item-aware callback form: a Result::Err from the closure does not fail the pipeline. A panic inside task_or still does. try_task / task_with_mode are the mode-based form (ErrorMode); see below.

For recoverable errors in a stage, use task_or:

use olympipe_rs::Pipeline;

let values = Pipeline::new(vec!["1", "two", "3"])
    .task_or(
        |s| s.parse::<i32>().map_err(|e| e),
        |s, _err| {
            eprintln!("skipping {s}");
            None
        },
    )
    .collect();
assert_eq!(values, vec![1, 3]);

Stateful workers

class_task builds one value per worker thread, inside that thread, and lends it to the method as &mut. count == 0 is one worker and keeps input order, same as task. class_task_with_close receives the state when the worker stops: input exhausted, downstream disconnected, or the pipeline cancelled. It is not called if init or the method panics (Drop of the state still runs on unwind once init has returned).

use olympipe_rs::Pipeline;

let values = Pipeline::new(1..=3)
    .class_task(|| 0, |state, x: i32| { *state += x; *state }, 1)
    .collect();
assert_eq!(values, vec![1, 3, 6]);

N-way split

split stays the two-branch path. split_by sends each item to branch f(&item) (an index >= n fails the pipeline, stage "split"). split_array is the optional-tuple form: one item can be copied to several branches or to none. Bounded variants use buffer_size and must be consumed concurrently (gather, gather_all, wait_for_results, or one thread per branch). The _unbounded variants can be drained one after another.

use olympipe_rs::Pipeline;

let mut values = Pipeline::gather_all(
    Pipeline::new(0..6i32).split_by(2, |x| if x % 2 == 0 { 0 } else { 1 }),
)
.collect();
values.sort();
assert_eq!(values, vec![0, 1, 2, 3, 4, 5]);

Error modes

ErrorMode defaults to FailFast. try_task maps a Result: FailFast cancels the pipeline with PipelineError::TaskFailed { stage: "try_task", .. } and does not catch panics. Skip drops the item, SkipAndLog also writes olympipe-rs: [try_task] skipped item: {err} to stderr, and Replace(v) emits a clone of v. In those three modes a panic inside the closure is a per-item failure too. task_with_mode is the same policy for an infallible closure; FailFast is exactly task (panic stage "task").

use olympipe_rs::{ErrorMode, Pipeline};

let values = Pipeline::new(vec!["1", "x", "3"])
    .try_task(|s: &str| s.parse::<i32>(), 1, ErrorMode::Replace(0))
    .collect();
assert_eq!(values, vec![1, 0, 3]);

Waiting on several pipelines

wait_for_results / try_wait_for_results drain this pipeline and the others at the same time and return one Vec per pipeline, self first. Sequential drains deadlock when the pipelines are bounded branches of one split. wait_for_all_completions discards output. wait_and_reduce folds with (accumulator, item). try_collect_pair is the two-pipeline form for different item types. A failure cancels pipelines that share state, and every terminal is still joined so the others return.

use olympipe_rs::Pipeline;

let sum = Pipeline::new(1..=4).wait_and_reduce(0, |acc, x| acc + x);
assert_eq!(sum, 10);

use olympipe_rs::{try_collect_pair, Pipeline};

let (numbers, labels) = try_collect_pair(
    Pipeline::new(vec![1, 2]),
    Pipeline::new(vec!["a", "b"]),
)
.expect("pair");
assert_eq!(numbers, vec![1, 2]);
assert_eq!(labels, vec!["a", "b"]);

Design notes

Checkpoint

checkpoint keeps one worker. reducer(state, item) folds every input into initial. When the input is exhausted and the pipeline has not failed, emitter(state) is streamed downstream (same packetised channel and back-pressure as other stages). A failure or cancellation emits nothing. A panic in either closure is stage "checkpoint".

use olympipe_rs::Pipeline;

let sorted = Pipeline::new(vec![3, 1, 2]).checkpoint(
    Vec::new(),
    |mut state, item| {
        state.push(item);
        state
    },
    |mut state| {
        state.sort();
        state
    },
);
assert_eq!(sorted.collect(), vec![1, 2, 3]);

Cached task

cached_task is behind the cache feature (serde_json + sha2). The default build does not depend on them.

[dependencies]
olympipe-rs = { version = "1", features = ["cache"] }

Rust cannot hash a closure, so the first argument is a step name you choose. Change it when the computation changes ("resize-v2"). The key is hex SHA-256 of the raw SHA-256 of that name concatenated with the serde_json bytes of the input. Files live at {cache_dir}/{sha256(step)[:16]}/{key}.json. A hit does not call f. A miss calls f and writes the JSON by renaming a temporary file in the same directory. A corrupt file is recomputed and overwritten. An error while writing the cache still forwards the value. There is no uncache: the stage yields U directly. Keys follow serde_json, so prefer BTreeMap over HashMap when the input contains a map.

# #[cfg(feature = "cache")]
# fn demo() {
use olympipe_rs::Pipeline;

let dir = std::env::temp_dir().join(format!(
    "olympipe-rs-readme-cache-{}",
    std::process::id()
));
let _ = std::fs::remove_dir_all(&dir);
let values = Pipeline::new(vec![1, 2, 1])
    .cached_task("double-v1", &dir, |x: i32| x * 2, 1)
    .collect();
assert_eq!(values, vec![2, 4, 2]);
let _ = std::fs::remove_dir_all(&dir);
# }
# #[cfg(not(feature = "cache"))]
# fn demo() {}
# demo();

Python olympipe parity