o41-training-data 0.1.6

//! Turns manifest references into file lists. Consumers checkpoint the
//! resolved config, so a republished manifest never changes a resumed stream.
use crate::{
    config::{Config, Source},
    storage::{Client, Object},
};
use anyhow::{Context, Result, ensure};
use serde::Deserialize;
use std::{collections::HashMap, sync::Arc};

#[derive(Deserialize)]
struct Shard {
    uri: String,
    records: usize,
    size_bytes: u64,
    #[serde(default)]
    etag: String,
}
#[derive(Deserialize)]
struct Stage {
    shards: Vec<Shard>,
}
#[derive(Deserialize)]
struct Split {
    shards: Vec<Shard>,
    #[serde(default)]
    curriculum: HashMap<String, Stage>,
}
#[derive(Deserialize)]
struct Manifest {
    version: u32,
    data_format: String,
    splits: HashMap<String, Split>,
}

fn read(uri: &str, client: &Client) -> Result<Manifest> {
    let source = Source {
        uri: uri.to_owned(),
        etag: None,
        size_bytes: None,
        rows: None,
    };
    let mut object = Object::open(&source, client)?;
    let body = object.range(0, object.size)?;
    let manifest: Manifest = serde_json::from_slice(&body)?;
    ensure!(manifest.version == 1, "expected data manifest version 1");
    Ok(manifest)
}

fn names(keys: impl Iterator<Item = impl AsRef<str>>) -> String {
    let mut names: Vec<_> = keys.map(|k| k.as_ref().to_owned()).collect();
    names.sort();
    names.join(", ")
}

pub fn resolve(mut config: Config) -> Result<Config> {
    let runtime = Arc::new(
        tokio::runtime::Builder::new_multi_thread()
            .worker_threads(1)
            .enable_all()
            .build()?,
    );
    let client = Client::new(config.storage.clone(), runtime);
    let mut manifests = HashMap::<String, Manifest>::new();
    for (stage_index, stage) in config.stages.iter_mut().enumerate() {
        for (source_index, group) in stage.sources.iter_mut().enumerate() {
            let location = format!("stage {stage_index} source {source_index}");
            let Some(uri) = group.manifest_uri.take() else {
                ensure!(
                    group.split.is_none() && group.curriculum_stages.is_empty(),
                    "{location}: split and curriculumStages need manifestUri"
                );
                ensure!(!group.files.is_empty(), "{location}: files or manifestUri");
                continue;
            };
            ensure!(
                group.files.is_empty(),
                "{location}: declares both files and manifestUri"
            );
            let split_name = group
                .split
                .take()
                .ok_or_else(|| anyhow::anyhow!("{location}: manifest source needs split"))?;
            if !manifests.contains_key(&uri) {
                let manifest =
                    read(&uri, &client).with_context(|| format!("{location}: manifest {uri}"))?;
                manifests.insert(uri.clone(), manifest);
            }
            let manifest = &manifests[&uri];
            let split = manifest.splits.get(&split_name).ok_or_else(|| {
                anyhow::anyhow!(
                    "{location}: unknown split {split_name:?}; available splits: {}",
                    names(manifest.splits.keys())
                )
            })?;
            let mut shards = Vec::new();
            if group.curriculum_stages.is_empty() {
                shards.extend(&split.shards);
            } else {
                let mut seen = std::collections::HashSet::new();
                for name in std::mem::take(&mut group.curriculum_stages) {
                    let stage = split.curriculum.get(&name).ok_or_else(|| {
                        anyhow::anyhow!(
                            "{location}: unknown curriculum stage {name:?}; available stages: {}",
                            names(split.curriculum.keys())
                        )
                    })?;
                    shards.extend(stage.shards.iter().filter(|s| seen.insert(&s.uri)));
                }
            }
            ensure!(
                !shards.is_empty(),
                "{location}: manifest selection has no shards"
            );
            match &group.encoding {
                Some(encoding) => ensure!(
                    *encoding == manifest.data_format,
                    "{location}: declared encoding {encoding} differs from manifest data_format {}",
                    manifest.data_format
                ),
                None => group.encoding = Some(manifest.data_format.clone()),
            }
            group.files = shards
                .into_iter()
                .map(|s| Source {
                    uri: s.uri.clone(),
                    etag: Some(s.etag.clone()).filter(|e| !e.is_empty()),
                    size_bytes: Some(s.size_bytes),
                    rows: Some(s.records),
                })
                .collect();
        }
    }
    Ok(config)
}