tenshift-core 0.1.3

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
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161

import os
import re

with open("src/pipeline.rs", "r") as f:
    text = f.read()

# Helper to capture everything from a pattern up to the end of the block or struct
def extract(pattern, end_marker=None):
    match = re.search(pattern, text)
    if not match:
        return ""
    start = match.start()
    
    if end_marker:
        end_match = re.search(end_marker, text[start:])
        if end_match:
            return text[start:start+end_match.end()]
        
    # fallback brace matching
    brace_level = 0
    in_block = False
    for i in range(start, len(text)):
        if text[i] == '{':
            if not in_block:
                in_block = True
            brace_level += 1
        elif text[i] == '}':
            brace_level -= 1
            if in_block and brace_level == 0:
                return text[start:i+1]
    return text[start:]

def write_mod(filename, content):
    with open("src/pipeline/" + filename, "w") as f:
        f.write("#![allow(clippy::module_name_repetitions)]\n\n")
        f.write("use crate::error::{Error, Result};\n")
        f.write("use crate::sample::{DType, Sample, Tensor};\n")
        f.write("use crate::source::Source;\n")
        f.write("use crate::transform::*;\n")
        f.write("use super::*;\n")
        f.write("use std::collections::BTreeMap;\n")
        f.write("use std::panic::{catch_unwind, AssertUnwindSafe};\n")
        f.write("use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};\n")
        f.write("use std::sync::Arc;\n")
        f.write("use std::time::{Duration, Instant};\n")
        f.write("use crossbeam_channel::{bounded, Receiver, RecvTimeoutError, Sender};\n")
        f.write("use std::any::Any;\n\n")
        f.write(content)

os.makedirs("src/pipeline", exist_ok=True)

# 1. config.rs
config = extract(r"/// Configuration for the pipeline\.\n#\[derive\(Debug, Clone\)\]\npub struct PipelineConfig")
config += "\n\n" + extract(r"/// What to do when a data loading error occurs\.\n#\[non_exhaustive\]\n#\[derive\(Debug, Clone, Copy, PartialEq, Eq\)\]\npub enum ErrorPolicy")
config += "\n\n" + extract(r"impl Default for PipelineConfig")
write_mod("config.rs", config)

# 2. core_types.rs
types = "pub(crate) type CollateFn = dyn Fn(Vec<Sample>) -> Result<Sample> + Send + Sync + 'static;\n\n"
types += extract(r"enum Stage") + "\n\n"
types += extract(r"enum CollectorStage") + "\n\n"
types += extract(r"enum CollectorItem") + "\n\n"
types += extract(r"struct SampleChunk") + "\n\n"
types += extract(r"enum SourceMessage") + "\n\n"
types += extract(r"enum CollateMode") + "\n\n"
types = types.replace("enum ", "pub(crate) enum ").replace("struct SampleChunk", "pub(crate) struct SampleChunk")
# fix internal fields
types = types.replace("CollectorStage::", "crate::pipeline::core_types::CollectorStage::")
write_mod("core_types.rs", types)

# 3. iterator.rs
iterator = extract(r"/// A running pipeline that yields batches of samples\.\npub struct PipelineIterator")
iterator += "\n\n" + extract(r"/// Error returned when waiting for the next batch with a timeout\.\n#\[derive\(Debug, Clone, Copy, PartialEq, Eq\)\]\npub enum NextTimeoutError")
iterator += "\n\n" + extract(r"impl PipelineIterator")
iterator += "\n\n" + extract(r"impl Iterator for PipelineIterator")
iterator += "\n\n" + extract(r"impl Drop for PipelineIterator")
iterator += "\n\n" + extract(r"/// Observability snapshot from a running or completed pipeline\.\n#\[derive\(Debug, Clone\)\]\npub struct PipelineStats")
iterator += "\n\n" + extract(r"impl std::fmt::Display for PipelineStats")
write_mod("iterator.rs", iterator)

# 4. collate.rs
collate = extract(r"fn collate_batch")
collate += "\n\n" + extract(r"fn default_collate")
collate = collate.replace("fn collate_batch", "pub(crate) fn collate_batch").replace("fn default_collate", "pub(crate) fn default_collate")
write_mod("collate.rs", collate)

# 5. worker.rs
worker = extract(r"fn apply_stateless")
worker = worker.replace("fn apply_stateless", "pub(crate) fn apply_stateless")
write_mod("worker.rs", worker)

# 6. collector.rs
collector = extract(r"fn push_collector_stage")
collector += "\n\n" + extract(r"fn finish_collector_stage")
collector += "\n\n" + extract(r"fn emit_items")
collector += "\n\n" + extract(r"fn handle_collector_error")
collector = collector.replace("fn push_", "pub(crate) fn push_") \
    .replace("fn finish_", "pub(crate) fn finish_") \
    .replace("fn emit_items", "pub(crate) fn emit_items") \
    .replace("fn handle_collector_error", "pub(crate) fn handle_collector_error")
write_mod("collector.rs", collector)

# 7. source_thread.rs
thrd = extract(r"fn recv_source_message")
thrd += "\n\n" + extract(r"fn run_source_epoch")
thrd += "\n\n" + extract(r"fn run_source_iterator")
thrd = thrd.replace("fn recv_source_message", "pub(crate) fn recv_source_message") \
    .replace("fn run_source_epoch", "pub(crate) fn run_source_epoch") \
    .replace("fn run_source_iterator", "pub(crate) fn run_source_iterator")
write_mod("source_thread.rs", thrd)

# 8. utils.rs
utils = extract(r"fn wrap_source")
utils += "\n\n" + extract(r"fn validate_stage_order")
utils += "\n\n" + extract(r"fn panic_message")
utils += "\n\n" + extract(r"fn num_cpus")
utils = utils.replace("fn wrap_source", "pub(crate) fn wrap_source") \
    .replace("fn validate_stage_order", "pub(crate) fn validate_stage_order") \
    .replace("fn panic_message", "pub(crate) fn panic_message") \
    .replace("fn num_cpus", "pub(crate) fn num_cpus")
write_mod("utils.rs", utils)

# 9. builder.rs
builder = extract(r"/// A composable data loading pipeline\.\n///\n/// Build it with \[`Pipeline::from_source`\], chain transforms, then iterate\.\npub struct Pipeline")
builder += "\n\n" + extract(r"impl Pipeline")
write_mod("builder.rs", builder)

# Make the mod.rs file
mod_rs = """//! Pipeline  -  the composable data loading engine.
//!
//! A pipeline is a chain: `Source -> [Workers] -> [Collector] -> Consumer`.
//! Multiple worker threads load data and apply stateless transforms in parallel.
//! A single collector thread applies ordering-sensitive stages such as shuffle,
//! batching, and collation.

pub mod config;
pub mod core_types;
pub mod collate;
pub mod collector;
pub mod iterator;
pub mod source_thread;
pub mod utils;
pub mod worker;
pub mod builder;

pub use config::{PipelineConfig, ErrorPolicy};
pub use builder::Pipeline;
pub use iterator::{PipelineIterator, NextTimeoutError, PipelineStats};
pub(crate) use core_types::*;
pub(crate) use collate::*;
pub(crate) use collector::*;
pub(crate) use source_thread::*;
pub(crate) use utils::*;
pub(crate) use worker::*;
"""

with open("src/pipeline/mod.rs", "w") as f:
    f.write(mod_rs)

with open("src/pipeline.rs", "w") as f:
    f.write("pub mod pipeline;\npub use pipeline::*;\n")