use super::*;
use arrow::{array::Int64Array, ffi_stream::ArrowArrayStreamReader};
use std::{
io::Write,
time::{Duration, Instant},
};
fn write_jsonl(path: &std::path::Path, ids: std::ops::Range<i64>) {
let mut gzip = flate2::write::GzEncoder::new(
std::fs::File::create(path).unwrap(),
flate2::Compression::default(),
);
for id in ids {
writeln!(gzip, "{{\"id\":{id}}}").unwrap();
}
gzip.finish().unwrap();
}
unsafe fn pull(handle: *mut Loader) -> (u64, u64, Vec<i64>) {
let mut stream = FFI_ArrowArrayStream::empty();
let mut stage = 0;
let mut index = 0;
assert_eq!(
unsafe { loader_next(handle, &mut stream, &mut stage, &mut index) },
1
);
let mut reader = ArrowArrayStreamReader::try_new(stream).unwrap();
let batch = reader.next().unwrap().unwrap();
let ids = batch
.column(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
.to_vec();
(stage, index, ids)
}
#[test]
fn advance_with_exactly_eight_queued_batches() {
let root = std::env::temp_dir().join("o41-training-data-prefetch");
std::fs::create_dir_all(&root).unwrap();
let path = root.join("prefetch.jsonl.gz");
write_jsonl(&path, 0..9);
let stage = serde_json::json!({"encoding":"jsonl.gz", "sources":[{"files":[{"uri":path}]}], "batchSize":2,"shuffle":{"bufferRows":5},"infinite":true});
let mut answers = Vec::new();
for prefetch in [0, 8] {
let config = serde_json::json!({"storage":{},"fields":[{"name":"id","from":"id","dtype":"i64"}],"stages":[stage,stage],"seed":7,"rank":0,"worldSize":1,"threads":2,"prefetchBatches":prefetch});
let config = CString::new(config.to_string()).unwrap();
unsafe {
let handle = loader_open(config.as_ptr());
assert!(!handle.is_null());
for _ in 0..3 {
pull(handle);
}
let deadline = Instant::now() + Duration::from_secs(5);
while (*handle).worker.as_ref().unwrap().batches.len() < prefetch
&& Instant::now() < deadline
{
std::thread::sleep(Duration::from_millis(5));
}
assert_eq!((*handle).worker.as_ref().unwrap().batches.len(), prefetch);
assert_eq!(loader_advance(handle, 1), 1);
let result: Vec<_> = (0..12).map(|_| pull(handle)).collect();
assert!(result.iter().all(|(stage, _, _)| *stage == 1));
answers.push(result);
loader_close(handle);
}
}
assert_eq!(answers[0], answers[1]);
}
#[test]
fn resolve_checkpoint_and_resume_are_exact() {
let root = std::env::temp_dir().join("o41-training-data-resume");
std::fs::create_dir_all(&root).unwrap();
let a = root.join("a.jsonl.gz");
let b = root.join("b.jsonl.gz");
write_jsonl(&a, 0..7);
write_jsonl(&b, 7..13);
let shard = |path: &std::path::Path, records: usize| serde_json::json!({"uri": path, "records": records, "size_bytes": std::fs::metadata(path).unwrap().len(), "etag": ""});
let manifest = root.join("data-manifest-v1.json");
std::fs::write(
&manifest,
serde_json::json!({"version": 1, "data_format": "jsonl.gz", "splits": {"train": {
"records": 13, "shards": [shard(&a, 7), shard(&b, 6)],
"curriculum": {"easy": {"records": 7, "shards": [shard(&a, 7)]},
"all": {"records": 13, "shards": [shard(&a, 7), shard(&b, 6)]}}}}})
.to_string(),
)
.unwrap();
let config = serde_json::json!({
"storage": {},
"fields": [{"name":"id","from":"id","dtype":"i64"}],
"stages": [{"sources": [{"manifestUri": manifest, "split": "train", "curriculumStages": ["easy", "all"]}],
"batchSize": 4, "infinite": true}],
"seed": 1, "rank": 0, "worldSize": 1, "threads": 2, "prefetchBatches": 2
});
let config = CString::new(config.to_string()).unwrap();
unsafe {
let resolved = loader_resolve(config.as_ptr());
assert!(
!resolved.is_null(),
"{:?}",
CStr::from_ptr(loader_last_error(ptr::null_mut()))
);
let parsed: Config =
serde_json::from_str(CStr::from_ptr(resolved).to_str().unwrap()).unwrap();
let files: Vec<_> = parsed.stages[0].sources[0]
.files
.iter()
.map(|f| (f.uri.clone(), f.rows))
.collect();
assert_eq!(
files,
vec![
(a.to_string_lossy().into_owned(), Some(7)),
(b.to_string_lossy().into_owned(), Some(6))
]
);
assert_eq!(
parsed.stages[0].sources[0].encoding.as_deref(),
Some("jsonl.gz")
);
let handle = loader_open(resolved);
assert!(
!handle.is_null(),
"{:?}",
CStr::from_ptr(loader_last_error(ptr::null_mut()))
);
loader_free_string(resolved);
let first: Vec<_> = (0..3).map(|_| pull(handle)).collect();
assert_eq!(first[0].2, vec![0, 1, 2, 3]);
let mut out: *const c_char = ptr::null();
assert_eq!(loader_checkpoint(handle, &mut out), 1);
let checkpoint = CStr::from_ptr(out).to_str().unwrap().to_owned();
let expected: Vec<_> = (0..5).map(|_| pull(handle)).collect();
loader_close(handle);
let document: serde_json::Value = serde_json::from_str(&checkpoint).unwrap();
assert_eq!(document["batchesYielded"], 3);
assert_eq!(document["rowsYielded"], 12);
assert!(document["config"]["storage"].get("accessKeyId").is_none());
let checkpoint = CString::new(checkpoint).unwrap();
let storage = CString::new("{}").unwrap();
let resumed = loader_resume(checkpoint.as_ptr(), storage.as_ptr());
assert!(
!resumed.is_null(),
"{:?}",
CStr::from_ptr(loader_last_error(ptr::null_mut()))
);
let actual: Vec<_> = (0..5).map(|_| pull(resumed)).collect();
assert_eq!(actual, expected);
loader_close(resumed);
}
}
#[test]
fn vector_rows_keep_bits_across_eviction_shuffle_ranks_and_resume() {
use arrow::array::{FixedSizeListArray, Float32Array};
let root = std::env::temp_dir().join(format!("o41-row-assembly-{}", std::process::id()));
std::fs::create_dir_all(&root).unwrap();
let mut files = Vec::new();
for file in 0..3 {
let path = root.join(format!("{file}.jsonl.gz"));
let mut gzip = flate2::write::GzEncoder::new(
std::fs::File::create(&path).unwrap(),
flate2::Compression::default(),
);
for id in file * 7..(file + 1) * 7 {
writeln!(
gzip,
"{}",
serde_json::json!({"id":id,"x":[-0.0f32,id as f32/4.0,-(id as f32)/8.0]})
)
.unwrap();
}
gzip.finish().unwrap();
files.push(serde_json::json!({"uri":path}));
}
let check = |batch: &arrow::record_batch::RecordBatch| {
let ids = batch
.column(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap();
let lists = batch
.column(1)
.as_any()
.downcast_ref::<FixedSizeListArray>()
.unwrap();
let values = lists
.values()
.as_any()
.downcast_ref::<Float32Array>()
.unwrap();
let bits: Vec<_> = values.values().iter().map(|v| v.to_bits()).collect();
for (&id, actual) in ids.values().iter().zip(bits.as_chunks::<3>().0.iter()) {
let expected = [-0.0f32, id as f32 / 4.0, -(id as f32) / 8.0].map(f32::to_bits);
assert_eq!(*actual, expected);
}
(ids.values().to_vec(), bits)
};
let read = |loader: &mut Loader| {
let batch = loader.next().unwrap().unwrap();
(batch.index, check(&batch.columns))
};
for shuffle in [
serde_json::Value::Null,
serde_json::json!({"bufferRows":17}),
] {
for rank in 0..2 {
let config: Config = serde_json::from_value(serde_json::json!({
"storage":{}, "fields":[{"name":"id","from":"id","dtype":"i64"},
{"name":"x","from":"x","dtype":"f32","shape":[3]}],
"stages":[{"encoding":"jsonl.gz","sources":[{"files":files}],
"batchSize":5,"shuffle":shuffle,"infinite":true}],
"seed":19,"rank":rank,"worldSize":2,"threads":2,"prefetchBatches":2
}))
.unwrap();
let mut loader = Loader::open(config, None).unwrap();
let retained = loader.next().unwrap().unwrap().columns;
let saved = check(&retained);
for _ in 0..7 {
read(&mut loader);
}
let checkpoint = loader.checkpoint();
let expected: Vec<_> = (0..12).map(|_| read(&mut loader)).collect();
let mut resumed = Loader::resume(checkpoint, Storage::default()).unwrap();
let actual: Vec<_> = (0..12).map(|_| read(&mut resumed)).collect();
assert_eq!(actual, expected);
assert_eq!(check(&retained), saved);
}
}
std::fs::remove_dir_all(root).unwrap();
}