o41-training-data 0.1.9

use anyhow::{Result, ensure};
use serde::{Deserialize, Serialize};

#[derive(Clone, Debug, Deserialize, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct Source {
    pub uri: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub etag: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub size_bytes: Option<u64>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub rows: Option<usize>,
}
/// Explicit object-store settings. Nothing is read from the environment.
#[derive(Clone, Debug, Deserialize, Serialize, Default)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct Storage {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub base_url: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub region: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub access_key_id: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub secret_access_key: Option<String>,
    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
    pub virtual_hosted_style_request: bool,
}
impl Storage {
    pub fn without_credentials(&self) -> Self {
        Self {
            access_key_id: None,
            secret_access_key: None,
            ..self.clone()
        }
    }
}
#[derive(Clone, Debug, PartialEq, Deserialize, Serialize)]
#[serde(deny_unknown_fields)]
pub struct Field {
    pub name: String,
    pub from: String,
    pub dtype: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub length: Option<usize>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub shape: Option<Vec<usize>>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub pad: Option<serde_json::Value>,
}
impl Field {
    pub fn dimensions(&self) -> Vec<usize> {
        self.shape
            .clone()
            .unwrap_or_else(|| self.length.into_iter().collect())
    }
    pub fn width(&self) -> usize {
        self.dimensions().iter().product()
    }
}
/// One weighted source of a stage: either resolved `files` or a manifest
/// reference that `manifest::resolve` turns into files.
#[derive(Clone, Debug, Deserialize, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct Group {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub encoding: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub fields: Option<Vec<Field>>,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub files: Vec<Source>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub manifest_uri: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub split: Option<String>,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub curriculum_stages: Vec<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub weight: Option<f64>,
}
#[derive(Clone, Debug, Deserialize, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct Shuffle {
    pub buffer_rows: usize,
}
#[derive(Clone, Debug, Deserialize, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct Stage {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub encoding: Option<String>,
    pub sources: Vec<Group>,
    pub batch_size: usize,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub shuffle: Option<Shuffle>,
    pub infinite: bool,
    #[serde(default)]
    pub drop_last: bool,
}
#[derive(Clone, Debug, Deserialize, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct Config {
    #[serde(default)]
    pub storage: Storage,
    pub stages: Vec<Stage>,
    pub fields: Vec<Field>,
    pub seed: u64,
    pub rank: u64,
    pub world_size: u64,
    /// Keep a rank's batch when the finite stream cannot fill every rank.
    /// Synchronized training keeps the default, false.
    #[serde(default, skip_serializing_if = "is_false")]
    pub allow_incomplete_group: bool,
    pub threads: usize,
    pub prefetch_batches: usize,
    #[serde(default = "lookahead")]
    pub lookahead: usize,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub memory_limit_bytes: Option<u64>,
    #[serde(default = "record_limit")]
    pub max_record_bytes: usize,
    #[serde(default = "row_group_limit")]
    pub max_row_group_bytes: usize,
    #[serde(default = "row_limit")]
    pub max_shard_rows: usize,
}
/// Decoded bytes a Parquet reader's batch of wide rows holds: rows of up to
/// 64 KiB decode 128 at a time.
const DECODE_BYTES: u128 = 8 * 1024 * 1024;
fn is_false(value: &bool) -> bool {
    !*value
}
fn lookahead() -> usize {
    2
}
fn record_limit() -> usize {
    4 * 1024 * 1024
}
fn row_group_limit() -> usize {
    64 * 1024 * 1024
}
fn row_limit() -> usize {
    65536
}
impl Config {
    pub fn memory_limit(&self) -> Result<u64> {
        if let Some(limit) = self.memory_limit_bytes {
            return Ok(limit);
        }
        let pages = unsafe { libc::sysconf(libc::_SC_PHYS_PAGES) };
        let page_size = unsafe { libc::sysconf(libc::_SC_PAGESIZE) };
        ensure!(
            pages > 0 && page_size > 0,
            "cannot read host physical memory"
        );
        Ok(((pages as u128 * page_size as u128) / 2).try_into()?)
    }
    /// Bytes the memory bound counts for a row: eight for each `i64` or `f64`
    /// element, four for any other.
    fn row_bytes(&self) -> u128 {
        self.fields
            .iter()
            .map(|f| {
                f.width() as u128
                    * if matches!(f.dtype.as_str(), "i64" | "f64") {
                        8
                    } else {
                        4
                    }
            })
            .sum()
    }
    /// Rows a Parquet reader decodes at once: 128, or as many as fit
    /// [`DECODE_BYTES`], at least one.
    pub(crate) fn decode_rows(&self) -> usize {
        (DECODE_BYTES / self.row_bytes().max(1)).clamp(1, 128) as usize
    }
    /// Decoded shards held at once, and the largest's rows: each reader slot
    /// holds one, and a file takes at most one slot.
    fn shards(&self) -> (u128, u128) {
        let files = self
            .stages
            .iter()
            .flat_map(|s| &s.sources)
            .flat_map(|g| &g.files);
        let slots =
            (self.threads as u128 * self.lookahead as u128).min(files.clone().count() as u128);
        let rows = files
            .map(|f| f.rows.unwrap_or(self.max_shard_rows))
            .max()
            .unwrap_or(0) as u128;
        (slots, rows)
    }
    pub fn memory_bound(&self) -> Result<u64> {
        let groups: Vec<_> = self.stages.iter().flat_map(|s| &s.sources).collect();
        let row = self.row_bytes();
        let metadata = std::iter::once(&self.fields)
            .chain(groups.iter().filter_map(|g| g.fields.as_ref()))
            .map(|fields| {
                fields
                    .iter()
                    .map(|f| (4096 + f.name.len() + f.from.len()) as u128)
                    .sum::<u128>()
            })
            .max()
            .unwrap_or(0);
        let (slots, rows) = self.shards();
        let batch_rows = self.stages.iter().map(|s| s.batch_size).max().unwrap_or(0) as u128;
        let batch = (row + 128) * batch_rows + metadata;
        let shuffle = self
            .stages
            .iter()
            .filter_map(|s| s.shuffle.as_ref())
            .map(|s| s.buffer_rows)
            .max()
            .unwrap_or(0) as u128;
        let sources = self
            .stages
            .iter()
            .map(|s| s.sources.len())
            .max()
            .unwrap_or(0) as u128;
        let encodings: Vec<_> = self
            .stages
            .iter()
            .flat_map(|s| {
                s.sources
                    .iter()
                    .map(|g| g.encoding.as_ref().or(s.encoding.as_ref()))
            })
            .collect();
        let parquet = encodings.iter().any(|e| e.is_some_and(|e| e == "parquet"));
        // A Parquet reader writes its rows into the shard it returns; any
        // other projects them into chunks, then copies them into one.
        let streamed = encodings.iter().any(|e| e.is_none_or(|e| e != "parquet"));
        let decoded = if streamed {
            128
        } else {
            self.decode_rows() as u128
        };
        let shard = row * rows;
        let fetched = crate::storage::HELD as u128
            * if parquet {
                (self.max_row_group_bytes as u64).max(crate::storage::WINDOW) as u128
            } else {
                crate::storage::WINDOW as u128
            }
            + crate::storage::PARTS as u128;
        let workspace = (64 * self.max_record_bytes as u128).max(if parquet {
            8 * self.max_row_group_bytes as u128
        } else {
            0
        }) + 16 * row * rows.min(decoded)
            + if streamed { 2 * shard } else { 0 }
            + metadata * rows.div_ceil(128)
            + fetched
            + 32 * 1024 * 1024;
        let buffered = (shuffle + batch_rows + 1) * (row + 1024);
        let snapshots = 64 * (shuffle + sources);
        let bound = slots
            .checked_mul(shard + metadata)
            .and_then(|n| {
                (self.threads as u128)
                    .min(slots)
                    .checked_mul(workspace)
                    .and_then(|w| n.checked_add(w))
            })
            .and_then(|n| n.checked_add(buffered))
            .and_then(|n| {
                (self.prefetch_batches as u128 + 3)
                    .checked_mul(batch * 3 + snapshots)
                    .and_then(|q| n.checked_add(q))
            })
            .ok_or_else(|| anyhow::anyhow!("memory bound exceeds supported range"))?;
        Ok(bound.try_into()?)
    }

    /// Structural checks that do not need storage. Every source must already be
    /// resolved to files; manifest references are resolved by `manifest::resolve`.
    pub fn validate(&self) -> Result<u64> {
        ensure!(!self.stages.is_empty(), "no stages");
        ensure!(
            !self.fields.is_empty(),
            "fields must be a nonempty allowlist"
        );
        ensure!(
            self.lookahead > 0
                && self.max_record_bytes > 0
                && self.max_shard_rows > 0
                && self.max_row_group_bytes > 0,
            "invalid memory limits"
        );
        ensure!(
            self.world_size > 0 && self.rank < self.world_size,
            "invalid rank layout"
        );
        ensure!(self.threads > 0, "threads must be positive");
        for (stage_index, stage) in self.stages.iter().enumerate() {
            ensure!(
                stage.sources.len() <= self.threads.saturating_mul(self.lookahead),
                "threads * lookahead must hold one current shard per source"
            );
            ensure!(
                stage.batch_size > 0 && !stage.sources.is_empty(),
                "empty stage"
            );
            ensure!(
                stage.shuffle.as_ref().is_none_or(|s| s.buffer_rows > 0),
                "empty shuffle buffer"
            );
            let mut sum = 0.0;
            for (source_index, group) in stage.sources.iter().enumerate() {
                ensure!(
                    group.manifest_uri.is_none() && group.split.is_none(),
                    "stage {stage_index} source {source_index}: manifest source is not resolved"
                );
                ensure!(
                    matches!(
                        group
                            .encoding
                            .as_ref()
                            .or(stage.encoding.as_ref())
                            .map(String::as_str),
                        Some("msgpack.zstf" | "jsonl.gz" | "parquet")
                    ),
                    "stage {stage_index} source {source_index}: unknown encoding"
                );
                let w = group.weight.unwrap_or(1.0);
                ensure!(
                    w.is_finite() && w >= 0.0 && !group.files.is_empty(),
                    "stage {stage_index} source {source_index}: invalid source group"
                );
                ensure!(
                    group
                        .files
                        .iter()
                        .all(|f| f.rows.is_none_or(|n| n <= self.max_shard_rows)),
                    "manifest rows exceed maxShardRows"
                );
                sum += w;
            }
            ensure!(sum.is_finite() && sum > 0.0, "invalid total weight");
        }
        for fields in std::iter::once(&self.fields).chain(
            self.stages
                .iter()
                .flat_map(|s| &s.sources)
                .filter_map(|g| g.fields.as_ref()),
        ) {
            let mut names = std::collections::HashSet::new();
            for field in fields {
                ensure!(names.insert(&field.name), "duplicate projected name");
                ensure!(
                    matches!(field.dtype.as_str(), "i32" | "i64" | "f32" | "f64" | "bool"),
                    "unknown dtype"
                );
                ensure!(
                    field.length.is_none_or(|n| n > 0 && n <= i32::MAX as usize),
                    "invalid field length"
                );
                ensure!(
                    field.shape.is_none() || field.length.is_none(),
                    "use length or shape, not both"
                );
                if let Some(shape) = &field.shape {
                    ensure!(
                        !shape.is_empty() && shape.iter().all(|&n| n > 0),
                        "invalid field shape"
                    );
                    ensure!(
                        shape
                            .iter()
                            .try_fold(1usize, |a, &b| a.checked_mul(b))
                            .is_some_and(|n| n <= i32::MAX as usize),
                        "field shape is too large"
                    );
                }
            }
        }
        for (stage_index, stage) in self.stages.iter().enumerate() {
            for (source_index, source) in stage.sources.iter().enumerate() {
                let fields = source.fields.as_ref().unwrap_or(&self.fields);
                ensure!(
                    fields.len() == self.fields.len(),
                    "stage {stage_index} source {source_index} {}: field count differs",
                    source.files[0].uri
                );
                for (a, b) in fields.iter().zip(&self.fields) {
                    ensure!(
                        a.name == b.name && a.dtype == b.dtype && a.dimensions() == b.dimensions(),
                        "stage {stage_index} source {source_index} {}: incompatible field {}",
                        source.files[0].uri,
                        b.name
                    );
                }
            }
        }
        let bound = self.memory_bound()?;
        let limit = self.memory_limit()?;
        let (slots, rows) = self.shards();
        ensure!(
            bound <= limit,
            "memory bound {bound} bytes exceeds memoryLimitBytes {limit}: it holds {slots} \
             decoded shards of up to {rows} rows of {} bytes; shards of fewer rows lower it",
            self.row_bytes()
        );
        Ok(bound)
    }
}