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>>>,
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 _wait = profile::Timer::new(9);
64                    select! { send(send, result) -> _ => (), recv(stopped) -> _ => return }
65                    if terminal {
66                        return;
67                    }
68                }
69            })?;
70        Ok(Self {
71            batches,
72            stop,
73            thread: Some(thread),
74        })
75    }
76}
77impl Drop for Worker {
78    fn drop(&mut self) {
79        let _ = self.stop.send(());
80        if let Some(thread) = self.thread.take() {
81            let _ = thread.join();
82        }
83    }
84}
85pub struct Loader {
86    config: Arc<Config>,
87    schema: Arc<arrow::datatypes::Schema>,
88    state: State,
89    worker: Option<Worker>,
90    error: CString,
91    checkpoint: CString,
92    ended: bool,
93}
94/// The complete resumable description of a stream: the resolved config
95/// without credentials, the loop cursor and the public counters.
96#[derive(Clone, Serialize, Deserialize)]
97#[serde(rename_all = "camelCase", deny_unknown_fields)]
98pub struct Checkpoint {
99    pub version: u32,
100    pub config: Config,
101    pub cursor: State,
102    pub stage: usize,
103    pub batches_yielded: u64,
104    pub rows_yielded: u64,
105    pub epochs: Vec<u64>,
106}
107impl Loader {
108    pub fn open(config: Config, state: Option<State>) -> Result<Self> {
109        let bound = config.validate()?;
110        let schema = columns::project_values(&[], &config.fields)?.schema();
111        let schema = Arc::new(schema.as_ref().clone().with_metadata(
112            std::collections::HashMap::from([
113                (
114                    "o41_training_data.memory_bound_bytes".into(),
115                    bound.to_string(),
116                ),
117                (
118                    "o41_training_data.memory_limit_bytes".into(),
119                    config.memory_limit()?.to_string(),
120                ),
121            ]),
122        ));
123        let state = state.unwrap_or_else(|| State::new(&config, 0));
124        state.validate(&config)?;
125        let config = Arc::new(config);
126        let worker = Some(Worker::start(config.clone(), state.clone())?);
127        Ok(Loader {
128            config,
129            schema,
130            state,
131            worker,
132            error: CString::default(),
133            checkpoint: CString::default(),
134            ended: false,
135        })
136    }
137    pub fn resume(checkpoint: Checkpoint, storage: Storage) -> Result<Self> {
138        ensure!(checkpoint.version == 1, "unsupported checkpoint version");
139        let mut config = checkpoint.config;
140        config.storage = storage;
141        Self::open(config, Some(checkpoint.cursor))
142    }
143    /// The next batch, or `None` once a finite stream is exhausted. An error
144    /// ends the stream; the loader stays usable only for `checkpoint`.
145    #[allow(clippy::should_implement_trait)]
146    pub fn next(&mut self) -> Result<Option<Batch>> {
147        if self.ended {
148            return Ok(None);
149        }
150        let wait = profile::Timer::new(8);
151        let result = self
152            .worker
153            .as_ref()
154            .ok_or_else(|| anyhow::anyhow!("loader stopped after a failed advance"))?
155            .batches
156            .recv()
157            .map_err(|_| anyhow::anyhow!("worker stopped"))?;
158        drop(wait);
159        match result {
160            Ok(Some(batch)) => {
161                self.state = batch.state.clone();
162                Ok(Some(batch))
163            }
164            Ok(None) => {
165                self.ended = true;
166                Ok(None)
167            }
168            Err(error) => {
169                self.ended = true;
170                Err(error)
171            }
172        }
173    }
174    /// Switches to `stage`, keeping the global batch partitioning aligned.
175    pub fn advance(&mut self, stage: u64) -> Result<()> {
176        ensure!(
177            stage < self.config.stages.len() as u64,
178            "stage out of range"
179        );
180        if stage as usize == self.state.stage {
181            return Ok(());
182        }
183        self.worker.take();
184        let mut aligned = self.state.clone();
185        if !aligned.batch_index.is_multiple_of(self.config.world_size) {
186            let mut engine = Engine::new(self.config.clone(), aligned)?;
187            engine.align()?;
188            aligned = engine.state.clone();
189        }
190        let mut state = State::new(&self.config, stage as usize);
191        state.buffer = aligned.buffer;
192        state.partial = aligned.partial;
193        state.batch_index = aligned.batch_index;
194        state.batches_yielded = aligned.batches_yielded;
195        state.rows_yielded = aligned.rows_yielded;
196        self.worker = Some(Worker::start(self.config.clone(), state.clone())?);
197        self.state = state;
198        self.ended = false;
199        Ok(())
200    }
201    pub fn checkpoint(&self) -> Checkpoint {
202        let mut config = self.config.as_ref().clone();
203        config.storage = config.storage.without_credentials();
204        Checkpoint {
205            version: 1,
206            config,
207            cursor: self.state.clone(),
208            stage: self.state.stage,
209            batches_yielded: self.state.batches_yielded,
210            rows_yielded: self.state.rows_yielded,
211            epochs: self.state.positions.iter().map(|p| p.epoch).collect(),
212        }
213    }
214}
215fn guarded(handle: *mut Loader, operation: impl FnOnce(&mut Loader) -> Result<i32>) -> i32 {
216    if handle.is_null() {
217        return -1;
218    }
219    let loader = unsafe { &mut *handle };
220    match catch_unwind(AssertUnwindSafe(|| operation(loader))) {
221        Ok(Ok(code)) => {
222            loader.error = CString::default();
223            code
224        }
225        Ok(Err(error)) => {
226            loader.error = error_string(error);
227            -1
228        }
229        Err(_) => {
230            loader.error = error_string("panic at loader boundary");
231            -1
232        }
233    }
234}
235fn opened(operation: impl FnOnce() -> Result<Loader>) -> *mut Loader {
236    let result = catch_unwind(AssertUnwindSafe(operation))
237        .unwrap_or_else(|_| Err(anyhow::anyhow!("panic opening loader")));
238    match result {
239        Ok(loader) => Box::into_raw(Box::new(loader)),
240        Err(error) => {
241            OPEN_ERROR.with(|e| *e.borrow_mut() = error_string(error));
242            ptr::null_mut()
243        }
244    }
245}
246#[unsafe(no_mangle)]
247pub unsafe extern "C" fn loader_resolve(config_json: *const c_char) -> *mut c_char {
248    let result = catch_unwind(AssertUnwindSafe(|| -> Result<CString> {
249        let config: Config = serde_json::from_str(unsafe { text(config_json) }?)?;
250        let resolved = resolve(config)?;
251        Ok(CString::new(serde_json::to_string(&resolved)?)?)
252    }))
253    .unwrap_or_else(|_| Err(anyhow::anyhow!("panic resolving config")));
254    match result {
255        Ok(json) => json.into_raw(),
256        Err(error) => {
257            OPEN_ERROR.with(|e| *e.borrow_mut() = error_string(error));
258            ptr::null_mut()
259        }
260    }
261}
262#[unsafe(no_mangle)]
263pub unsafe extern "C" fn loader_free_string(string: *mut c_char) {
264    if !string.is_null() {
265        drop(unsafe { CString::from_raw(string) });
266    }
267}
268#[unsafe(no_mangle)]
269pub unsafe extern "C" fn loader_open(config_json: *const c_char) -> *mut Loader {
270    opened(|| {
271        let config: Config = serde_json::from_str(unsafe { text(config_json) }?)?;
272        Loader::open(config, None)
273    })
274}
275#[unsafe(no_mangle)]
276pub unsafe extern "C" fn loader_resume(
277    checkpoint_json: *const c_char,
278    storage_json: *const c_char,
279) -> *mut Loader {
280    opened(|| {
281        let checkpoint: Checkpoint = serde_json::from_str(unsafe { text(checkpoint_json) }?)?;
282        let storage: Storage = serde_json::from_str(unsafe { text(storage_json) }?)?;
283        Loader::resume(checkpoint, storage)
284    })
285}
286#[unsafe(no_mangle)]
287pub unsafe extern "C" fn loader_next(
288    handle: *mut Loader,
289    stream: *mut FFI_ArrowArrayStream,
290    stage: *mut u64,
291    batch_index: *mut u64,
292) -> i32 {
293    guarded(handle, |loader| {
294        ensure!(
295            !stream.is_null() && !stage.is_null() && !batch_index.is_null(),
296            "null output"
297        );
298        let Some(batch) = loader.next()? else {
299            return Ok(0);
300        };
301        let _export = profile::Timer::new(10);
302        let columns = batch.columns.with_schema(loader.schema.clone())?;
303        let reader = RecordBatchIterator::new(std::iter::once(Ok(columns)), loader.schema.clone());
304        unsafe {
305            ptr::write(stream, FFI_ArrowArrayStream::new(Box::new(reader)));
306            *stage = batch.stage as u64;
307            *batch_index = batch.index;
308        }
309        Ok(1)
310    })
311}
312#[unsafe(no_mangle)]
313pub unsafe extern "C" fn loader_checkpoint(handle: *mut Loader, out: *mut *const c_char) -> i32 {
314    guarded(handle, |loader| {
315        ensure!(!out.is_null(), "null output");
316        loader.checkpoint = CString::new(serde_json::to_string(&loader.checkpoint())?)?;
317        unsafe { ptr::write(out, loader.checkpoint.as_ptr()) };
318        Ok(1)
319    })
320}
321#[unsafe(no_mangle)]
322pub unsafe extern "C" fn loader_advance(handle: *mut Loader, stage: u64) -> i32 {
323    guarded(handle, |loader| {
324        loader.advance(stage)?;
325        Ok(1)
326    })
327}
328#[unsafe(no_mangle)]
329pub unsafe extern "C" fn loader_last_error(handle: *mut Loader) -> *const c_char {
330    if handle.is_null() {
331        OPEN_ERROR.with(|e| e.borrow().as_ptr())
332    } else {
333        unsafe { (*handle).error.as_ptr() }
334    }
335}
336#[unsafe(no_mangle)]
337pub unsafe extern "C" fn loader_close(handle: *mut Loader) {
338    if !handle.is_null() {
339        let _ = catch_unwind(AssertUnwindSafe(|| drop(unsafe { Box::from_raw(handle) })));
340        profile::report();
341    }
342}
343
344#[cfg(test)]
345mod tests;
o41_training_data/lib.rs

o41_training_data/
lib.rs
