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

o41_training_data/
lib.rs
