o41-training-data 0.1.6

//! C callers must supply live handles, valid aligned output pointers and
//! NUL-terminated UTF-8 strings. Calls on one handle must be serialized.
//! Imported Arrow streams and arrays must be released by their Arrow callbacks.
#![allow(clippy::missing_safety_doc)]
mod columns;
mod config;
mod engine;
mod manifest;
mod profile;
mod readers;
mod storage;

pub use config::{Config, Field, Group, Shuffle, Source, Stage, Storage};
pub use engine::{Batch, Coordinate, Position, State};
pub use manifest::resolve;
pub use storage::Client;

use anyhow::{Result, ensure};
use arrow::{ffi_stream::FFI_ArrowArrayStream, record_batch::RecordBatchIterator};
use crossbeam_channel::{Receiver, Sender, bounded, select};
use engine::Engine;
use serde::{Deserialize, Serialize};
use std::{
    cell::RefCell,
    ffi::{CStr, CString, c_char},
    panic::{AssertUnwindSafe, catch_unwind},
    ptr,
    sync::Arc,
    thread::JoinHandle,
};

thread_local! { static OPEN_ERROR: RefCell<CString> = RefCell::new(CString::default()); }
fn error_string(error: impl std::fmt::Display) -> CString {
    CString::new(format!("{error:#}").replace('\0', "\\0")).unwrap()
}
unsafe fn text<'a>(pointer: *const c_char) -> Result<&'a str> {
    ensure!(!pointer.is_null(), "null string");
    Ok(unsafe { CStr::from_ptr(pointer) }.to_str()?)
}
struct Worker {
    batches: Receiver<Result<Option<Batch>>>,
    stop: Sender<()>,
    thread: Option<JoinHandle<()>>,
}
impl Worker {
    fn start(config: Arc<Config>, state: State) -> Result<Self> {
        let mut engine = Engine::new(config.clone(), state)?;
        let (send, batches) = bounded(config.prefetch_batches);
        let (stop, stopped) = bounded(0);
        let thread = std::thread::Builder::new()
            .name("training-data-batches".into())
            .spawn(move || {
                loop {
                    if matches!(
                        stopped.try_recv(),
                        Ok(()) | Err(crossbeam_channel::TryRecvError::Disconnected)
                    ) {
                        return;
                    }
                    let result = catch_unwind(AssertUnwindSafe(|| engine.next()))
                        .unwrap_or_else(|_| Err(anyhow::anyhow!("reader worker panicked")));
                    let terminal = !matches!(&result, Ok(Some(_)));
                    let _wait = profile::Timer::new(9);
                    select! { send(send, result) -> _ => (), recv(stopped) -> _ => return }
                    if terminal {
                        return;
                    }
                }
            })?;
        Ok(Self {
            batches,
            stop,
            thread: Some(thread),
        })
    }
}
impl Drop for Worker {
    fn drop(&mut self) {
        // A rendezvous send cancels either a blocked queue send or the next loop.
        let _ = self.stop.send(());
        if let Some(thread) = self.thread.take() {
            let _ = thread.join();
        }
    }
}
pub struct Loader {
    config: Arc<Config>,
    schema: Arc<arrow::datatypes::Schema>,
    state: State,
    worker: Option<Worker>,
    error: CString,
    checkpoint: CString,
    ended: bool,
}
/// The complete resumable description of a stream: the resolved config
/// without credentials, the loop cursor and the public counters.
#[derive(Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct Checkpoint {
    pub version: u32,
    pub config: Config,
    pub cursor: State,
    pub stage: usize,
    pub batches_yielded: u64,
    pub rows_yielded: u64,
    pub epochs: Vec<u64>,
}
impl Loader {
    pub fn open(config: Config, state: Option<State>) -> Result<Self> {
        let bound = config.validate()?;
        let schema = columns::project_values(&[], &config.fields)?.schema();
        let schema = Arc::new(schema.as_ref().clone().with_metadata(
            std::collections::HashMap::from([
                (
                    "o41_training_data.memory_bound_bytes".into(),
                    bound.to_string(),
                ),
                (
                    "o41_training_data.memory_limit_bytes".into(),
                    config.memory_limit()?.to_string(),
                ),
            ]),
        ));
        let state = state.unwrap_or_else(|| State::new(&config, 0));
        state.validate(&config)?;
        let config = Arc::new(config);
        let worker = Some(Worker::start(config.clone(), state.clone())?);
        Ok(Loader {
            config,
            schema,
            state,
            worker,
            error: CString::default(),
            checkpoint: CString::default(),
            ended: false,
        })
    }
    pub fn resume(checkpoint: Checkpoint, storage: Storage) -> Result<Self> {
        ensure!(checkpoint.version == 1, "unsupported checkpoint version");
        let mut config = checkpoint.config;
        config.storage = storage;
        Self::open(config, Some(checkpoint.cursor))
    }
    /// The next batch, or `None` once a finite stream is exhausted. An error
    /// ends the stream; the loader stays usable only for `checkpoint`.
    // Named after `loader_next` and the Python binding, not after `Iterator`:
    // a loader outlives its stream and is read again through `checkpoint`.
    #[allow(clippy::should_implement_trait)]
    pub fn next(&mut self) -> Result<Option<Batch>> {
        if self.ended {
            return Ok(None);
        }
        let wait = profile::Timer::new(8);
        let result = self
            .worker
            .as_ref()
            .ok_or_else(|| anyhow::anyhow!("loader stopped after a failed advance"))?
            .batches
            .recv()
            .map_err(|_| anyhow::anyhow!("worker stopped"))?;
        drop(wait);
        match result {
            Ok(Some(batch)) => {
                self.state = batch.state.clone();
                Ok(Some(batch))
            }
            Ok(None) => {
                self.ended = true;
                Ok(None)
            }
            Err(error) => {
                self.ended = true;
                Err(error)
            }
        }
    }
    /// Switches to `stage`, keeping the global batch partitioning aligned.
    pub fn advance(&mut self, stage: u64) -> Result<()> {
        ensure!(
            stage < self.config.stages.len() as u64,
            "stage out of range"
        );
        if stage as usize == self.state.stage {
            return Ok(());
        }
        self.worker.take();
        let mut aligned = self.state.clone();
        if !aligned.batch_index.is_multiple_of(self.config.world_size) {
            let mut engine = Engine::new(self.config.clone(), aligned)?;
            engine.align()?;
            aligned = engine.state.clone();
        }
        let mut state = State::new(&self.config, stage as usize);
        state.buffer = aligned.buffer;
        state.partial = aligned.partial;
        state.batch_index = aligned.batch_index;
        state.batches_yielded = aligned.batches_yielded;
        state.rows_yielded = aligned.rows_yielded;
        self.worker = Some(Worker::start(self.config.clone(), state.clone())?);
        self.state = state;
        self.ended = false;
        Ok(())
    }
    pub fn checkpoint(&self) -> Checkpoint {
        let mut config = self.config.as_ref().clone();
        config.storage = config.storage.without_credentials();
        Checkpoint {
            version: 1,
            config,
            cursor: self.state.clone(),
            stage: self.state.stage,
            batches_yielded: self.state.batches_yielded,
            rows_yielded: self.state.rows_yielded,
            epochs: self.state.positions.iter().map(|p| p.epoch).collect(),
        }
    }
}
fn guarded(handle: *mut Loader, operation: impl FnOnce(&mut Loader) -> Result<i32>) -> i32 {
    if handle.is_null() {
        return -1;
    }
    let loader = unsafe { &mut *handle };
    match catch_unwind(AssertUnwindSafe(|| operation(loader))) {
        Ok(Ok(code)) => {
            loader.error = CString::default();
            code
        }
        Ok(Err(error)) => {
            loader.error = error_string(error);
            -1
        }
        Err(_) => {
            loader.error = error_string("panic at loader boundary");
            -1
        }
    }
}
fn opened(operation: impl FnOnce() -> Result<Loader>) -> *mut Loader {
    let result = catch_unwind(AssertUnwindSafe(operation))
        .unwrap_or_else(|_| Err(anyhow::anyhow!("panic opening loader")));
    match result {
        Ok(loader) => Box::into_raw(Box::new(loader)),
        Err(error) => {
            OPEN_ERROR.with(|e| *e.borrow_mut() = error_string(error));
            ptr::null_mut()
        }
    }
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn loader_resolve(config_json: *const c_char) -> *mut c_char {
    let result = catch_unwind(AssertUnwindSafe(|| -> Result<CString> {
        let config: Config = serde_json::from_str(unsafe { text(config_json) }?)?;
        let resolved = resolve(config)?;
        Ok(CString::new(serde_json::to_string(&resolved)?)?)
    }))
    .unwrap_or_else(|_| Err(anyhow::anyhow!("panic resolving config")));
    match result {
        Ok(json) => json.into_raw(),
        Err(error) => {
            OPEN_ERROR.with(|e| *e.borrow_mut() = error_string(error));
            ptr::null_mut()
        }
    }
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn loader_free_string(string: *mut c_char) {
    if !string.is_null() {
        drop(unsafe { CString::from_raw(string) });
    }
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn loader_open(config_json: *const c_char) -> *mut Loader {
    opened(|| {
        let config: Config = serde_json::from_str(unsafe { text(config_json) }?)?;
        Loader::open(config, None)
    })
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn loader_resume(
    checkpoint_json: *const c_char,
    storage_json: *const c_char,
) -> *mut Loader {
    opened(|| {
        let checkpoint: Checkpoint = serde_json::from_str(unsafe { text(checkpoint_json) }?)?;
        let storage: Storage = serde_json::from_str(unsafe { text(storage_json) }?)?;
        Loader::resume(checkpoint, storage)
    })
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn loader_next(
    handle: *mut Loader,
    stream: *mut FFI_ArrowArrayStream,
    stage: *mut u64,
    batch_index: *mut u64,
) -> i32 {
    guarded(handle, |loader| {
        ensure!(
            !stream.is_null() && !stage.is_null() && !batch_index.is_null(),
            "null output"
        );
        let Some(batch) = loader.next()? else {
            return Ok(0);
        };
        let _export = profile::Timer::new(10);
        let columns = batch.columns.with_schema(loader.schema.clone())?;
        let reader = RecordBatchIterator::new(std::iter::once(Ok(columns)), loader.schema.clone());
        unsafe {
            ptr::write(stream, FFI_ArrowArrayStream::new(Box::new(reader)));
            *stage = batch.stage as u64;
            *batch_index = batch.index;
        }
        Ok(1)
    })
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn loader_checkpoint(handle: *mut Loader, out: *mut *const c_char) -> i32 {
    guarded(handle, |loader| {
        ensure!(!out.is_null(), "null output");
        loader.checkpoint = CString::new(serde_json::to_string(&loader.checkpoint())?)?;
        unsafe { ptr::write(out, loader.checkpoint.as_ptr()) };
        Ok(1)
    })
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn loader_advance(handle: *mut Loader, stage: u64) -> i32 {
    guarded(handle, |loader| {
        loader.advance(stage)?;
        Ok(1)
    })
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn loader_last_error(handle: *mut Loader) -> *const c_char {
    if handle.is_null() {
        OPEN_ERROR.with(|e| e.borrow().as_ptr())
    } else {
        unsafe { (*handle).error.as_ptr() }
    }
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn loader_close(handle: *mut Loader) {
    if !handle.is_null() {
        let _ = catch_unwind(AssertUnwindSafe(|| drop(unsafe { Box::from_raw(handle) })));
        profile::report();
    }
}

#[cfg(test)]
mod tests;