o41_training_data/
lib.rs

1//! C callers must supply live handles, valid aligned output pointers and
2//! NUL-terminated UTF-8 strings. Calls on one handle must be serialized.
3//! Imported Arrow streams and arrays must be released by their Arrow callbacks.
4#![allow(clippy::missing_safety_doc)]
5mod columns;
6mod config;
7mod engine;
8mod manifest;
9mod profile;
10mod readers;
11mod storage;
12
13pub use config::{Config, Field, Group, Shuffle, Source, Stage, Storage};
14pub use engine::{Batch, Coordinate, Position, State};
15pub use manifest::resolve;
16pub use storage::Client;
17
18use anyhow::{Result, ensure};
19use arrow::{ffi_stream::FFI_ArrowArrayStream, record_batch::RecordBatchIterator};
20use crossbeam_channel::{Receiver, Sender, bounded, select};
21use engine::Engine;
22use serde::{Deserialize, Serialize};
23use std::{
24    cell::RefCell,
25    ffi::{CStr, CString, c_char},
26    panic::{AssertUnwindSafe, catch_unwind},
27    ptr,
28    sync::Arc,
29    thread::JoinHandle,
30};
31
32thread_local! { static OPEN_ERROR: RefCell<CString> = RefCell::new(CString::default()); }
33fn error_string(error: impl std::fmt::Display) -> CString {
34    CString::new(format!("{error:#}").replace('\0', "\\0")).unwrap()
35}
36unsafe fn text<'a>(pointer: *const c_char) -> Result<&'a str> {
37    ensure!(!pointer.is_null(), "null string");
38    Ok(unsafe { CStr::from_ptr(pointer) }.to_str()?)
39}
40struct Worker {
41    batches: Receiver<Result<(Option<Batch>, u64)>>,
42    stop: Sender<()>,
43    thread: Option<JoinHandle<()>>,
44}
45impl Worker {
46    fn start(config: Arc<Config>, state: State) -> Result<Self> {
47        let mut engine = Engine::new(config.clone(), state)?;
48        let (send, batches) = bounded(config.prefetch_batches);
49        let (stop, stopped) = bounded(0);
50        let thread = std::thread::Builder::new()
51            .name("training-data-batches".into())
52            .spawn(move || {
53                loop {
54                    if matches!(
55                        stopped.try_recv(),
56                        Ok(()) | Err(crossbeam_channel::TryRecvError::Disconnected)
57                    ) {
58                        return;
59                    }
60                    let result = catch_unwind(AssertUnwindSafe(|| engine.next()))
61                        .unwrap_or_else(|_| Err(anyhow::anyhow!("reader worker panicked")));
62                    let terminal = !matches!(&result, Ok(Some(_)));
63                    let result = result.map(|batch| (batch, engine.dropped_rows));
64                    let _wait = profile::Timer::new(9);
65                    select! { send(send, result) -> _ => (), recv(stopped) -> _ => return }
66                    if terminal {
67                        return;
68                    }
69                }
70            })?;
71        Ok(Self {
72            batches,
73            stop,
74            thread: Some(thread),
75        })
76    }
77}
78impl Drop for Worker {
79    fn drop(&mut self) {
80        let _ = self.stop.send(());
81        if let Some(thread) = self.thread.take() {
82            let _ = thread.join();
83        }
84    }
85}
86pub struct Loader {
87    config: Arc<Config>,
88    schema: Arc<arrow::datatypes::Schema>,
89    state: State,
90    worker: Option<Worker>,
91    error: CString,
92    checkpoint: CString,
93    ended: bool,
94    dropped_rows: u64,
95}
96/// The complete resumable description of a stream: the resolved config
97/// without credentials, the loop cursor and the public counters.
98#[derive(Clone, Serialize, Deserialize)]
99#[serde(rename_all = "camelCase", deny_unknown_fields)]
100pub struct Checkpoint {
101    pub version: u32,
102    pub config: Config,
103    pub cursor: State,
104    pub stage: usize,
105    pub batches_yielded: u64,
106    pub rows_yielded: u64,
107    pub epochs: Vec<u64>,
108}
109impl Loader {
110    pub fn open(config: Config, state: Option<State>) -> Result<Self> {
111        let bound = config.validate()?;
112        let schema = columns::project_values(&[], &config.fields)?.schema();
113        let schema = Arc::new(schema.as_ref().clone().with_metadata(
114            std::collections::HashMap::from([
115                (
116                    "o41_training_data.memory_bound_bytes".into(),
117                    bound.to_string(),
118                ),
119                (
120                    "o41_training_data.memory_limit_bytes".into(),
121                    config.memory_limit()?.to_string(),
122                ),
123            ]),
124        ));
125        let state = state.unwrap_or_else(|| State::new(&config, 0));
126        state.validate(&config)?;
127        let config = Arc::new(config);
128        let worker = Some(Worker::start(config.clone(), state.clone())?);
129        Ok(Loader {
130            config,
131            schema,
132            state,
133            worker,
134            error: CString::default(),
135            checkpoint: CString::default(),
136            ended: false,
137            dropped_rows: 0,
138        })
139    }
140    pub fn resume(checkpoint: Checkpoint, storage: Storage) -> Result<Self> {
141        ensure!(checkpoint.version == 1, "unsupported checkpoint version");
142        let mut config = checkpoint.config;
143        config.storage = storage;
144        Self::open(config, Some(checkpoint.cursor))
145    }
146    /// The next batch, or `None` once a finite stream is exhausted. An error
147    /// ends the stream; the loader stays usable only for `checkpoint`.
148    #[allow(clippy::should_implement_trait)]
149    pub fn next(&mut self) -> Result<Option<Batch>> {
150        if self.ended {
151            return Ok(None);
152        }
153        let wait = profile::Timer::new(8);
154        let result = self
155            .worker
156            .as_ref()
157            .ok_or_else(|| anyhow::anyhow!("loader stopped after a failed advance"))?
158            .batches
159            .recv()
160            .map_err(|_| anyhow::anyhow!("worker stopped"))?;
161        drop(wait);
162        match result {
163            Ok((Some(batch), _)) => {
164                self.state = batch.state.clone();
165                Ok(Some(batch))
166            }
167            Ok((None, dropped_rows)) => {
168                self.ended = true;
169                self.dropped_rows = dropped_rows;
170                Ok(None)
171            }
172            Err(error) => {
173                self.ended = true;
174                Err(error)
175            }
176        }
177    }
178    /// Switches to `stage`, keeping the global batch partitioning aligned.
179    pub fn advance(&mut self, stage: u64) -> Result<()> {
180        ensure!(
181            stage < self.config.stages.len() as u64,
182            "stage out of range"
183        );
184        if stage as usize == self.state.stage {
185            return Ok(());
186        }
187        self.worker.take();
188        let mut aligned = self.state.clone();
189        if !aligned.batch_index.is_multiple_of(self.config.world_size) {
190            let mut engine = Engine::new(self.config.clone(), aligned)?;
191            engine.align()?;
192            aligned = engine.state.clone();
193        }
194        let mut state = State::new(&self.config, stage as usize);
195        state.buffer = aligned.buffer;
196        state.partial = aligned.partial;
197        state.batch_index = aligned.batch_index;
198        state.batches_yielded = aligned.batches_yielded;
199        state.rows_yielded = aligned.rows_yielded;
200        self.worker = Some(Worker::start(self.config.clone(), state.clone())?);
201        self.state = state;
202        self.ended = false;
203        self.dropped_rows = 0;
204        Ok(())
205    }
206    /// Rows discarded by batching on this rank since opening or advancing the
207    /// loader. Available after `next` returns `None`; a prefetched tail does
208    /// not count before the consumer reaches the end. Summing across ranks
209    /// includes a short final batch and any incomplete final group of batches.
210    pub fn dropped_rows(&self) -> u64 {
211        self.dropped_rows
212    }
213    pub fn checkpoint(&self) -> Checkpoint {
214        let mut config = self.config.as_ref().clone();
215        config.storage = config.storage.without_credentials();
216        Checkpoint {
217            version: 1,
218            config,
219            cursor: self.state.clone(),
220            stage: self.state.stage,
221            batches_yielded: self.state.batches_yielded,
222            rows_yielded: self.state.rows_yielded,
223            epochs: self.state.positions.iter().map(|p| p.epoch).collect(),
224        }
225    }
226}
227fn guarded(handle: *mut Loader, operation: impl FnOnce(&mut Loader) -> Result<i32>) -> i32 {
228    if handle.is_null() {
229        return -1;
230    }
231    let loader = unsafe { &mut *handle };
232    match catch_unwind(AssertUnwindSafe(|| operation(loader))) {
233        Ok(Ok(code)) => {
234            loader.error = CString::default();
235            code
236        }
237        Ok(Err(error)) => {
238            loader.error = error_string(error);
239            -1
240        }
241        Err(_) => {
242            loader.error = error_string("panic at loader boundary");
243            -1
244        }
245    }
246}
247fn opened(operation: impl FnOnce() -> Result<Loader>) -> *mut Loader {
248    let result = catch_unwind(AssertUnwindSafe(operation))
249        .unwrap_or_else(|_| Err(anyhow::anyhow!("panic opening loader")));
250    match result {
251        Ok(loader) => Box::into_raw(Box::new(loader)),
252        Err(error) => {
253            OPEN_ERROR.with(|e| *e.borrow_mut() = error_string(error));
254            ptr::null_mut()
255        }
256    }
257}
258#[unsafe(no_mangle)]
259pub unsafe extern "C" fn loader_resolve(config_json: *const c_char) -> *mut c_char {
260    let result = catch_unwind(AssertUnwindSafe(|| -> Result<CString> {
261        let config: Config = serde_json::from_str(unsafe { text(config_json) }?)?;
262        let resolved = resolve(config)?;
263        Ok(CString::new(serde_json::to_string(&resolved)?)?)
264    }))
265    .unwrap_or_else(|_| Err(anyhow::anyhow!("panic resolving config")));
266    match result {
267        Ok(json) => json.into_raw(),
268        Err(error) => {
269            OPEN_ERROR.with(|e| *e.borrow_mut() = error_string(error));
270            ptr::null_mut()
271        }
272    }
273}
274#[unsafe(no_mangle)]
275pub unsafe extern "C" fn loader_free_string(string: *mut c_char) {
276    if !string.is_null() {
277        drop(unsafe { CString::from_raw(string) });
278    }
279}
280#[unsafe(no_mangle)]
281pub unsafe extern "C" fn loader_open(config_json: *const c_char) -> *mut Loader {
282    opened(|| {
283        let config: Config = serde_json::from_str(unsafe { text(config_json) }?)?;
284        Loader::open(config, None)
285    })
286}
287#[unsafe(no_mangle)]
288pub unsafe extern "C" fn loader_resume(
289    checkpoint_json: *const c_char,
290    storage_json: *const c_char,
291) -> *mut Loader {
292    opened(|| {
293        let checkpoint: Checkpoint = serde_json::from_str(unsafe { text(checkpoint_json) }?)?;
294        let storage: Storage = serde_json::from_str(unsafe { text(storage_json) }?)?;
295        Loader::resume(checkpoint, storage)
296    })
297}
298#[unsafe(no_mangle)]
299pub unsafe extern "C" fn loader_next(
300    handle: *mut Loader,
301    stream: *mut FFI_ArrowArrayStream,
302    stage: *mut u64,
303    batch_index: *mut u64,
304) -> i32 {
305    guarded(handle, |loader| {
306        ensure!(
307            !stream.is_null() && !stage.is_null() && !batch_index.is_null(),
308            "null output"
309        );
310        let Some(batch) = loader.next()? else {
311            return Ok(0);
312        };
313        let _export = profile::Timer::new(10);
314        let columns = batch.columns.with_schema(loader.schema.clone())?;
315        let reader = RecordBatchIterator::new(std::iter::once(Ok(columns)), loader.schema.clone());
316        unsafe {
317            ptr::write(stream, FFI_ArrowArrayStream::new(Box::new(reader)));
318            *stage = batch.stage as u64;
319            *batch_index = batch.index;
320        }
321        Ok(1)
322    })
323}
324#[unsafe(no_mangle)]
325pub unsafe extern "C" fn loader_checkpoint(handle: *mut Loader, out: *mut *const c_char) -> i32 {
326    guarded(handle, |loader| {
327        ensure!(!out.is_null(), "null output");
328        loader.checkpoint = CString::new(serde_json::to_string(&loader.checkpoint())?)?;
329        unsafe { ptr::write(out, loader.checkpoint.as_ptr()) };
330        Ok(1)
331    })
332}
333#[unsafe(no_mangle)]
334pub unsafe extern "C" fn loader_advance(handle: *mut Loader, stage: u64) -> i32 {
335    guarded(handle, |loader| {
336        loader.advance(stage)?;
337        Ok(1)
338    })
339}
340#[unsafe(no_mangle)]
341pub unsafe extern "C" fn loader_last_error(handle: *mut Loader) -> *const c_char {
342    if handle.is_null() {
343        OPEN_ERROR.with(|e| e.borrow().as_ptr())
344    } else {
345        unsafe { (*handle).error.as_ptr() }
346    }
347}
348#[unsafe(no_mangle)]
349pub unsafe extern "C" fn loader_close(handle: *mut Loader) {
350    if !handle.is_null() {
351        let _ = catch_unwind(AssertUnwindSafe(|| drop(unsafe { Box::from_raw(handle) })));
352        profile::report();
353    }
354}
355
356#[cfg(test)]
357mod tests;
o41_training_data/lib.rs

o41_training_data/
lib.rs
