diff --git a/packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py b/packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py index 5b9f5d0b6..a07eca938 100644 --- a/packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py +++ b/packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py @@ -44,6 +44,7 @@ has_resumable_ingest, validated_index_resume_progress, compute_col_stats_for_parquet, + evaluate_prediction_parquet, resolve_temporal, detect_multicat_delimiter, multicat_delimiter_is_proven, @@ -146,6 +147,7 @@ def __getattr__(name): "has_resumable_ingest", "validated_index_resume_progress", "compute_col_stats_for_parquet", + "evaluate_prediction_parquet", "DiskGraphPlanCompiler", "SplitSpec", "DiskgraphConfigConverter", diff --git a/packages/diskgraph/diskgraph-pybind/src/bp_col_stats.rs b/packages/diskgraph/diskgraph-pybind/src/bp_col_stats.rs index e6703783f..15b81c0a4 100644 --- a/packages/diskgraph/diskgraph-pybind/src/bp_col_stats.rs +++ b/packages/diskgraph/diskgraph-pybind/src/bp_col_stats.rs @@ -1018,7 +1018,7 @@ fn accumulate_by_class( type RowGroupStats = (FileColStats, ByClassStats, Vec); /// The row group's batches of the projected columns and `keys`. -fn row_group_batches( +pub(crate) fn row_group_batches( path: &Path, row_group: usize, parquet_indices: &[usize], diff --git a/packages/diskgraph/diskgraph-pybind/src/bp_evaluate.rs b/packages/diskgraph/diskgraph-pybind/src/bp_evaluate.rs new file mode 100644 index 000000000..6565de8eb --- /dev/null +++ b/packages/diskgraph/diskgraph-pybind/src/bp_evaluate.rs @@ -0,0 +1,932 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +// All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Join batch-prediction outputs with generated labels and compute realized +//! metrics. Both sides are hash-partitioned on their key into Arrow IPC +//! spill files of about `partition_rows` rows each, and at most +//! `MAX_CONCURRENT_JOINS` partitions are joined at once. Metrics follow the +//! fit-time definitions: exact AUROC and AP as in torchmetrics, MAE/RMSE, +//! top-1 accuracy, and MAP@k / precision@k as in PyG. + +use std::collections::hash_map::DefaultHasher; +use std::collections::{BTreeMap, HashMap}; +use std::fs::File; +use std::hash::{Hash, Hasher}; +use std::io::BufWriter; +use std::path::{Path, PathBuf}; +use std::sync::{Arc, Mutex}; + +use anyhow::{bail, Context, Result}; +use arrow::array::TimestampNanosecondArray; +use arrow::array::{Array, AsArray, BooleanArray, FixedSizeBinaryArray, UInt32Array}; +use arrow::compute::{cast, filter, interleave, take_record_batch}; +use arrow::datatypes::Float64Type; +use arrow::ipc::reader::FileReader; +use arrow::ipc::writer::FileWriter; +use arrow::row::{RowConverter, SortField}; +use arrow_array::{ArrayRef, RecordBatch}; +use arrow_schema::{DataType, Field, Schema, SchemaRef, TimeUnit}; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use parquet::arrow::ArrowWriter; +use rayon::prelude::*; +use serde::Deserialize; +use serde_json::{json, Value as JsonValue}; + +use crate::bp_col_stats::row_group_batches; + +#[derive(Clone, Copy, Debug, Deserialize, PartialEq)] +#[serde(rename_all = "snake_case")] +pub enum EvaluateTask { + Binary, + Regression, + Multiclass, + Link, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct EvaluateSpec { + pub task: EvaluateTask, + pub prediction_keys: Vec, + pub label_keys: Vec, + pub label_target: String, + #[serde(default)] + pub actuals_dir: Option, + /// (label column, output column) pairs written to `actuals_dir`. + #[serde(default)] + pub actual_columns: Vec<(String, String)>, + /// (prediction column, output column) pairs written next to them, from + /// the matched prediction (the top-scored row of a grouped task). + #[serde(default)] + pub prediction_columns: Vec<(String, String)>, + /// (output column, UTC nanoseconds) written as a constant timestamp. + #[serde(default)] + pub constant_timestamp: Option<(String, i64)>, + pub spill_dir: PathBuf, + #[serde(default = "default_partition_rows")] + pub partition_rows: usize, +} + +fn default_partition_rows() -> usize { + 1_000_000 +} + +const MAX_CONCURRENT_JOINS: usize = 8; +const KEY: &str = "__key"; + +type Batches = Box> + Send>; + +/// The type a pair of columns is compared as. +fn common_type(a: &DataType, b: &DataType) -> DataType { + let int = |t: &DataType| t.is_integer() || *t == DataType::Boolean; + if a.is_temporal() || b.is_temporal() { + DataType::Timestamp(TimeUnit::Nanosecond, None) + } else if int(a) && int(b) { + DataType::Int64 + } else if [a, b].iter().all(|t| int(t) || t.is_floating()) { + DataType::Float64 + } else { + DataType::Utf8 + } +} + +/// 128-bit hash of each row of `arrays` cast to `types`; `None` on a null. +fn hashes(arrays: &[&ArrayRef], types: &[DataType]) -> Result>> { + let arrays = arrays + .iter() + .zip(types) + .map(|(array, to)| match (to, array.data_type()) { + (DataType::Timestamp(unit, _), DataType::Timestamp(_, tz)) => { + cast(array, &DataType::Timestamp(*unit, tz.clone())) + } + _ => cast(array, to), + }) + .collect::, _>>()?; + let fields = arrays + .iter() + .map(|a| SortField::new(a.data_type().clone())) + .collect(); + let rows = RowConverter::new(fields)?.convert_columns(&arrays)?; + Ok((0..rows.num_rows()) + .map(|row| { + arrays.iter().all(|array| array.is_valid(row)).then(|| { + let half = |seed: u64| { + let mut hasher = DefaultHasher::new(); + seed.hash(&mut hasher); + rows.row(row).hash(&mut hasher); + hasher.finish() + }; + u128::from(half(1)) << 64 | u128::from(half(2)) + }) + }) + .collect()) +} + +fn classes(array: &ArrayRef, class: &DataType) -> Result>> { + let hashes = hashes(&[array], std::slice::from_ref(class))?; + Ok(hashes.into_iter().map(|h| h.map(|h| h as u64)).collect()) +} + +fn column<'a>(batch: &'a RecordBatch, name: &str) -> Result<&'a ArrayRef> { + batch + .column_by_name(name) + .with_context(|| format!("missing column '{name}'")) +} + +fn floats(array: &ArrayRef) -> Result>> { + Ok(cast(array, &DataType::Float64)? + .as_primitive::() + .iter() + .collect()) +} + +fn keys(batch: &RecordBatch) -> Result + '_> { + let keys = column(batch, KEY)?.as_fixed_size_binary(); + Ok((0..keys.len()).map(|row| u128::from_le_bytes(keys.value(row).try_into().unwrap()))) +} + +/// Sorted, distinct class hashes of each label. +fn label_classes(target: &ArrayRef, link: bool, class: &DataType) -> Result>>> { + if !link { + let classes = classes(target, class)?; + return Ok(classes.into_iter().map(|c| c.map(|c| vec![c])).collect()); + } + let item = Field::new_list_field(class.clone(), true); + let list = cast(target, &DataType::List(Arc::new(item)))?; + let list = list.as_list::(); + let items = classes(list.values(), class)?; + let offsets = list.value_offsets(); + Ok((0..list.len()) + .map(|row| { + list.is_valid(row).then(|| { + let range = offsets[row] as usize..offsets[row + 1] as usize; + let mut row: Vec = items[range].iter().flatten().copied().collect(); + row.sort_unstable(); + row.dedup(); + row + }) + }) + .collect()) +} + +type Sink = (PathBuf, Mutex>>>); + +fn spill((path, writer): &Sink, batch: &RecordBatch) -> Result<()> { + let mut writer = writer.lock().unwrap(); + if writer.is_none() { + let file = BufWriter::new(File::create(path)?); + *writer = Some(FileWriter::try_new(file, &batch.schema())?); + } + Ok(writer.as_mut().unwrap().write(batch)?) +} + +fn unspill((path, writer): Sink) -> Result { + let Some(mut writer) = writer.into_inner().unwrap() else { + return Ok(Box::new(std::iter::empty())); + }; + writer.finish()?; + let reader = FileReader::try_new(File::open(path)?, None)?; + Ok(Box::new(reader.map(|batch| Ok(batch?)))) +} + +type RowGroups = Vec<(usize, usize)>; + +fn row_groups(paths: &[&Path]) -> Result<(RowGroups, usize, SchemaRef)> { + let (mut groups, mut rows, mut schema) = (Vec::new(), 0, Arc::new(Schema::empty())); + for (file, path) in paths.iter().enumerate() { + let builder = ParquetRecordBatchReaderBuilder::try_new(File::open(path)?)?; + schema = builder.schema().clone(); + rows += builder.metadata().file_metadata().num_rows() as usize; + groups.extend((0..builder.metadata().num_row_groups()).map(|group| (file, group))); + } + Ok((groups, rows, schema)) +} + +/// Read `keep` columns, hash `keys`, and route rows with a key to `sinks`. +fn scatter( + paths: &[&Path], + groups: &RowGroups, + (keys, types): (&[String], &[DataType]), + keep: &[String], + sinks: &[Sink], +) -> Result<()> { + let columns: Vec = keys.iter().chain(keep).cloned().collect(); + groups.par_iter().try_for_each(|&(file, group)| { + for batch in row_group_batches(paths[file], group, &[], &columns)? { + let batch = batch?; + let key_columns = keys + .iter() + .map(|name| column(&batch, name)) + .collect::>>()?; + let hashes = hashes(&key_columns, types)?; + let mut parts = vec![Vec::new(); sinks.len()]; + for (row, key) in hashes.iter().enumerate() { + if let Some(key) = key { + parts[((key >> 64) as u64 % sinks.len() as u64) as usize].push(row as u32); + } + } + let key = hashes.into_iter().map(|key| key.map(u128::to_le_bytes)); + let key = FixedSizeBinaryArray::try_from_sparse_iter_with_size(key, 16)?; + let mut out = Vec::new(); + for name in keep { + out.push((name.as_str(), column(&batch, name)?.clone(), true)); + } + out.push((KEY, Arc::new(key), true)); + let batch = RecordBatch::try_from_iter_with_nullable(out)?; + for (sink, rows) in sinks.iter().zip(parts) { + if !rows.is_empty() { + spill(sink, &take_record_batch(&batch, &UInt32Array::from(rows))?)?; + } + } + } + Ok(()) + }) +} + +#[derive(Default)] +struct Accumulator { + rows: u64, + labelled: u64, + /// Binary: score order key shifted left, label in bit 0. + pairs: Vec, + count: u64, + sum: f64, + squares: f64, + /// Link: MAP numerators by label count. + map_by_labels: BTreeMap, + k: usize, +} + +impl Accumulator { + fn merge(mut self, mut other: Self) -> Self { + self.rows += other.rows; + self.labelled += other.labelled; + self.pairs.append(&mut other.pairs); + self.count += other.count; + self.sum += other.sum; + self.squares += other.squares; + for (labels, value) in other.map_by_labels { + *self.map_by_labels.entry(labels).or_default() += value; + } + self.k = self.k.max(other.k); + self + } +} + +fn order_key(score: f64) -> u64 { + let bits = (score as f32).to_bits(); + u64::from(if bits & 0x8000_0000 != 0 { + !bits + } else { + bits | 0x8000_0000 + }) +} + +/// A grouped prediction row: score, class and (batch, row). +type Ranked = (f64, Option, (usize, usize)); + +/// Score each label row against the predictions with its key; the first +/// label of a key counts. +fn join( + predictions: Batches, + labels: Batches, + spec: &EvaluateSpec, + (score, class): (&str, &DataType), + actuals: Option, +) -> Result { + let grouped = matches!(spec.task, EvaluateTask::Multiclass | EvaluateTask::Link); + let mut first = HashMap::new(); + let mut groups: HashMap> = HashMap::new(); + let mut kept: Vec> = Vec::new(); + for (index, batch) in predictions.enumerate() { + let batch = batch?; + kept.push( + spec.prediction_columns + .iter() + .map(|(name, _)| column(&batch, name).cloned()) + .collect::>()?, + ); + let scores = floats(column(&batch, score)?)?; + if !grouped { + for (row, (key, score)) in keys(&batch)?.zip(scores).enumerate() { + first.entry(key).or_insert((score, (index, row))); + } + continue; + } + let classes = classes(column(&batch, "CLASS")?, class)?; + for (row, ((key, score), class)) in keys(&batch)?.zip(scores).zip(classes).enumerate() { + let score = score.unwrap_or(f64::NEG_INFINITY); + groups + .entry(key) + .or_default() + .push((score, class, (index, row))); + } + } + let mut acc = Accumulator { + rows: (first.len() + groups.len()) as u64, + k: groups.values().map(Vec::len).max().unwrap_or(0), + ..Default::default() + }; + + let mut writer: Option> = None; + for batch in labels { + let batch = batch?; + let target = column(&batch, &spec.label_target)?; + let mut matched = vec![false; batch.num_rows()]; + let mut picked = Vec::new(); + if !grouped { + for ((row, key), label) in keys(&batch)?.enumerate().zip(floats(target)?) { + let Some(label) = label else { continue }; + let Some((score, at)) = first.remove(&key) else { + continue; + }; + matched[row] = true; + picked.push(at); + match score { + Some(score) if spec.task == EvaluateTask::Binary => { + if score.is_finite() { + acc.pairs + .push(order_key(score) << 1 | u64::from(label != 0.0)); + } + } + Some(prediction) => { + let error = prediction - label; + acc.count += 1; + acc.sum += error.abs(); + acc.squares += error * error; + } + None => {} + } + } + } else { + let link = spec.task == EvaluateTask::Link; + let labels = label_classes(target, link, class)?; + for ((row, key), items) in keys(&batch)?.enumerate().zip(labels) { + let Some(items) = items else { continue }; + let Some(mut ranked) = groups.remove(&key) else { + continue; + }; + matched[row] = true; + // Highest score first; ties keep their order. + ranked.sort_by(|a, b| b.0.total_cmp(&a.0)); + picked.push(ranked[0].2); + let hit = + |class: Option| class.is_some_and(|c| items.binary_search(&c).is_ok()); + if !link { + acc.count += 1; + acc.sum += f64::from(u8::from(hit(ranked[0].1))); + } else if !items.is_empty() { + let (mut hits, mut numerator) = (0usize, 0.0); + for (rank, (_, class, _)) in ranked.iter().enumerate() { + if hit(*class) { + hits += 1; + numerator += hits as f64 / (rank + 1) as f64; + } + } + acc.count += 1; + acc.sum += hits as f64; + *acc.map_by_labels.entry(items.len()).or_default() += numerator; + } + } + } + let matched = BooleanArray::from(matched); + acc.labelled += matched.true_count() as u64; + let Some(path) = actuals.as_ref().filter(|_| matched.true_count() > 0) else { + continue; + }; + let mut out = spec + .actual_columns + .iter() + .map(|(name, out)| Ok((out.as_str(), filter(column(&batch, name)?, &matched)?, true))) + .collect::>>()?; + if let Some((name, ns)) = &spec.constant_timestamp { + let values = vec![*ns; picked.len()]; + let array = TimestampNanosecondArray::from(values).with_timezone("UTC"); + out.push((name.as_str(), Arc::new(array), true)); + } + for (column, (_, name)) in spec.prediction_columns.iter().enumerate() { + let arrays: Vec<&dyn Array> = kept.iter().map(|batch| batch[column].as_ref()).collect(); + out.push((name.as_str(), interleave(&arrays, &picked)?, true)); + } + let out = RecordBatch::try_from_iter_with_nullable(out)?; + if writer.is_none() { + writer = Some(ArrowWriter::try_new( + File::create(path)?, + out.schema(), + None, + )?); + } + writer.as_mut().unwrap().write(&out)?; + } + if let Some(writer) = writer { + writer.close()?; + } + Ok(acc) +} + +/// Exact AUROC and average precision from (score, label) order keys. +fn ranking_metrics(mut pairs: Vec) -> (Option, Option) { + pairs.par_sort_unstable(); + let positives = pairs.iter().filter(|pair| *pair & 1 == 1).count() as f64; + let negatives = pairs.len() as f64 - positives; + if positives == 0.0 || negatives == 0.0 { + return (None, None); + } + let groups: Vec<(f64, f64)> = pairs + .chunk_by(|a, b| a >> 1 == b >> 1) + .map(|group| { + let positive = group.iter().filter(|pair| *pair & 1 == 1).count() as f64; + (positive, group.len() as f64 - positive) + }) + .collect(); + let (mut below, mut area) = (0.0, 0.0); + for &(positive, negative) in &groups { + area += positive * (below + negative / 2.0); + below += negative; + } + let (mut true_positive, mut false_positive, mut previous_recall, mut precision_sum) = + (0.0, 0.0, 0.0, 0.0); + for &(positive, negative) in groups.iter().rev() { + true_positive += positive; + false_positive += negative; + let recall = true_positive / positives; + precision_sum += + (recall - previous_recall) * true_positive / (true_positive + false_positive); + previous_recall = recall; + } + (Some(area / (positives * negatives)), Some(precision_sum)) +} + +fn metrics(acc: Accumulator, task: EvaluateTask) -> JsonValue { + let ratio = |num: f64| (acc.count > 0).then(|| num / acc.count as f64); + match task { + EvaluateTask::Binary => { + let (auroc, ap) = ranking_metrics(acc.pairs); + json!({"auroc": auroc, "ap": ap}) + } + EvaluateTask::Regression => { + json!({"mae": ratio(acc.sum), "rmse": ratio(acc.squares).map(f64::sqrt)}) + } + EvaluateTask::Multiclass => json!({"acc": ratio(acc.sum)}), + EvaluateTask::Link => { + let k = acc.k.max(1); + let map: f64 = acc + .map_by_labels + .iter() + .map(|(labels, numerator)| numerator / (*labels).min(k) as f64) + .sum(); + let precision = ratio(acc.sum / k as f64); + json!({format!("map@{k}"): ratio(map), format!("precision@{k}"): precision}) + } + } +} + +/// Join `prediction_paths` with `label_paths` and compute the realized +/// metrics; returns `{"rows", "labelled", "metrics"}`. +pub fn evaluate_prediction_parquet( + prediction_paths: &[&Path], + label_paths: &[&Path], + spec: &EvaluateSpec, +) -> Result { + if spec.prediction_keys.is_empty() || spec.prediction_keys.len() != spec.label_keys.len() { + bail!("prediction_keys and label_keys must pair up"); + } + let (prediction_groups, prediction_rows, prediction_schema) = row_groups(prediction_paths)?; + let (label_groups, label_rows, label_schema) = row_groups(label_paths)?; + let types = spec + .prediction_keys + .iter() + .zip(&spec.label_keys) + .map(|(p, l)| { + Ok(common_type( + prediction_schema.field_with_name(p)?.data_type(), + label_schema.field_with_name(l)?.data_type(), + )) + }) + .collect::>>()?; + let score = match spec.task { + EvaluateTask::Binary => "True_PROB", + EvaluateTask::Regression => "TARGET_PRED", + EvaluateTask::Multiclass | EvaluateTask::Link => "SCORE", + }; + let mut prediction_columns = vec![score.to_string()]; + prediction_columns.extend(spec.prediction_columns.iter().map(|(name, _)| name.clone())); + let mut class = DataType::Null; + if matches!(spec.task, EvaluateTask::Multiclass | EvaluateTask::Link) { + prediction_columns.push("CLASS".to_string()); + } + prediction_columns.sort(); + prediction_columns.dedup(); + if matches!(spec.task, EvaluateTask::Multiclass | EvaluateTask::Link) { + let mut target = label_schema + .field_with_name(&spec.label_target)? + .data_type(); + if spec.task == EvaluateTask::Link { + let (DataType::List(item) + | DataType::LargeList(item) + | DataType::FixedSizeList(item, _)) = target + else { + bail!("link-prediction labels must be lists, got {target}"); + }; + target = item.data_type(); + } + class = common_type( + prediction_schema.field_with_name("CLASS")?.data_type(), + target, + ); + } + let mut label_columns = vec![spec.label_target.clone()]; + label_columns.extend(spec.actual_columns.iter().map(|(name, _)| name.clone())); + label_columns.sort(); + label_columns.dedup(); + + let partitions = prediction_rows + .max(label_rows) + .div_ceil(spec.partition_rows.max(1)) + .max(1); + std::fs::create_dir_all(&spec.spill_dir)?; + if let Some(dir) = &spec.actuals_dir { + std::fs::create_dir_all(dir)?; + } + let sinks = |side: &str| -> Vec { + (0..partitions) + .map(|p| { + ( + spec.spill_dir.join(format!("{side}-{p:05}.arrow")), + Mutex::default(), + ) + }) + .collect() + }; + let (predictions, labels) = (sinks("predictions"), sinks("labels")); + let prediction_keys = (spec.prediction_keys.as_slice(), types.as_slice()); + scatter( + prediction_paths, + &prediction_groups, + prediction_keys, + &prediction_columns, + &predictions, + )?; + let label_keys = (spec.label_keys.as_slice(), types.as_slice()); + scatter( + label_paths, + &label_groups, + label_keys, + &label_columns, + &labels, + )?; + + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(MAX_CONCURRENT_JOINS) + .build()?; + let acc = pool.install(|| { + predictions + .into_par_iter() + .zip(labels) + .enumerate() + .map(|(p, (predictions, labels))| { + let actuals = + (spec.actuals_dir.as_ref()).map(|dir| dir.join(format!("part-{p:05}.parquet"))); + join( + unspill(predictions)?, + unspill(labels)?, + spec, + (score, &class), + actuals, + ) + }) + .try_reduce(Accumulator::default, |a, b| Ok(a.merge(b))) + })?; + let _ = std::fs::remove_dir_all(&spec.spill_dir); + let (rows, labelled) = (acc.rows, acc.labelled); + Ok(json!({"rows": rows, "labelled": labelled, "metrics": metrics(acc, spec.task)})) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ + BooleanArray, Float64Array, Int64Array, Int64Builder, ListBuilder, StringArray, + TimestampMicrosecondArray, TimestampNanosecondArray, + }; + use parquet::file::properties::WriterProperties; + use tempfile::tempdir; + + fn write(path: &Path, columns: Vec<(&str, ArrayRef)>, rows_per_group: usize) { + let batch = RecordBatch::try_from_iter(columns).unwrap(); + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(rows_per_group)) + .build(); + let mut writer = + ArrowWriter::try_new(File::create(path).unwrap(), batch.schema(), Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + } + + fn spec(task: EvaluateTask, dir: &Path, partition_rows: usize) -> EvaluateSpec { + EvaluateSpec { + task, + prediction_keys: vec!["ENTITY".into(), "ANCHOR_TIMESTAMP".into()], + label_keys: vec!["ENTITY1".into(), "TIMESTAMP".into()], + label_target: "TARGET1".into(), + actuals_dir: Some(dir.join("actuals")), + actual_columns: vec![ + ("ENTITY1".into(), "ENTITY".into()), + ("TIMESTAMP".into(), "ANCHOR_TIMESTAMP".into()), + ("TARGET1".into(), "TARGET".into()), + ], + prediction_columns: vec![], + constant_timestamp: None, + spill_dir: dir.join("spill"), + partition_rows, + } + } + + fn evaluate(dir: &Path, spec: &EvaluateSpec) -> JsonValue { + let predictions = dir.join("predictions.parquet"); + let labels = dir.join("labels.parquet"); + evaluate_prediction_parquet(&[predictions.as_path()], &[labels.as_path()], spec).unwrap() + } + + fn actuals(dir: &Path) -> Vec { + std::fs::read_dir(dir.join("actuals")) + .unwrap() + .flat_map(|entry| { + let file = File::open(entry.unwrap().path()).unwrap(); + ParquetRecordBatchReaderBuilder::try_new(file) + .unwrap() + .build() + .unwrap() + .map(Result::unwrap) + }) + .collect() + } + + fn actual_column( + dir: &Path, + name: &str, + ) -> Vec { + actuals(dir) + .iter() + .flat_map(|batch| { + let array = cast(column(batch, name).unwrap(), &T::DATA_TYPE).unwrap(); + array.as_primitive::().values().to_vec() + }) + .collect() + } + + /// Mann-Whitney AUROC with ties counted half, by brute force. + fn brute_auroc(scores: &[f32], labels: &[bool]) -> f64 { + let (mut wins, mut pairs) = (0.0, 0.0); + for (i, &si) in scores.iter().enumerate() { + for (j, &sj) in scores.iter().enumerate() { + if labels[i] && !labels[j] { + pairs += 1.0; + wins += if si > sj { + 1.0 + } else if si == sj { + 0.5 + } else { + 0.0 + }; + } + } + } + wins / pairs + } + + #[test] + fn binary_join_matches_brute_force_in_one_and_many_partitions() { + let n = 3000i64; + let entity: Vec = (0..n).collect(); + // Several anchors; tz-aware microseconds against naive nanoseconds. + let anchor_us: Vec = (0..n).map(|i| (i % 3) * 86_400_000_000).collect(); + let scores: Vec = (0..n).map(|i| ((i * 7919) % 101) as f32 / 100.0).collect(); + let labels: Vec = (0..n).map(|i| (i * 31 + i / 7) % 5 < 2).collect(); + // Predictions for every entity; labels for all but the last 500. + let m = (n - 500) as usize; + for partition_rows in [1_000_000, 700] { + let dir = tempdir().unwrap(); + write( + &dir.path().join("predictions.parquet"), + vec![ + ( + "ENTITY", + Arc::new(Int64Array::from(entity.clone())) as ArrayRef, + ), + ( + "ANCHOR_TIMESTAMP", + Arc::new( + TimestampMicrosecondArray::from(anchor_us.clone()).with_timezone("UTC"), + ) as ArrayRef, + ), + ( + "True_PROB", + Arc::new(arrow::array::Float32Array::from(scores.clone())) as ArrayRef, + ), + ], + 256, + ); + write( + &dir.path().join("labels.parquet"), + vec![ + ( + "ENTITY1", + Arc::new(Int64Array::from( + entity[..m].iter().rev().copied().collect::>(), + )) as ArrayRef, + ), + ( + "TIMESTAMP", + Arc::new(TimestampNanosecondArray::from( + anchor_us[..m] + .iter() + .rev() + .map(|t| t * 1000) + .collect::>(), + )) as ArrayRef, + ), + ( + "TARGET1", + Arc::new(BooleanArray::from( + labels[..m].iter().rev().copied().collect::>(), + )) as ArrayRef, + ), + ], + 300, + ); + let mut spec_ = spec(EvaluateTask::Binary, dir.path(), partition_rows); + spec_.prediction_columns = vec![("True_PROB".into(), "True_PROB".into())]; + let out = evaluate(dir.path(), &spec_); + assert_eq!(out["rows"], json!(n)); + assert_eq!(out["labelled"], json!(m)); + let auroc = out["metrics"]["auroc"].as_f64().unwrap(); + assert!((auroc - brute_auroc(&scores[..m], &labels[..m])).abs() < 1e-12); + assert!(out["metrics"]["ap"].as_f64().unwrap() > 0.0); + let entities = actual_column::(dir.path(), "ENTITY"); + let emitted = actual_column::(dir.path(), "True_PROB"); + assert_eq!(entities.len(), m); + for (entity, score) in entities.iter().zip(emitted) { + assert_eq!(score, scores[*entity as usize]); + } + assert!(!dir.path().join("spill").exists()); + } + } + + #[test] + fn average_precision_matches_step_definition() { + // Scores 0.9(+), 0.8(-), 0.8(+), 0.1(-): thresholds 0.9, 0.8, 0.1. + let pairs = vec![ + order_key(0.9) << 1 | 1, + order_key(0.8) << 1, + order_key(0.8) << 1 | 1, + order_key(0.1) << 1, + ]; + let (auroc, ap) = ranking_metrics(pairs); + assert!((auroc.unwrap() - 0.875).abs() < 1e-12); + assert!((ap.unwrap() - (0.5 * 1.0 + 0.5 * 2.0 / 3.0)).abs() < 1e-12); + assert_eq!(ranking_metrics(vec![order_key(0.5) << 1 | 1]), (None, None)); + } + + #[test] + fn regression_multiclass_and_link_metrics() { + let dir = tempdir().unwrap(); + let mut spec_ = spec(EvaluateTask::Regression, dir.path(), 1_000_000); + spec_.prediction_keys.truncate(1); + spec_.label_keys.truncate(1); + spec_.actual_columns = vec![ + ("ENTITY1".into(), "ENTITY".into()), + ("TARGET1".into(), "TARGET".into()), + ]; + spec_.constant_timestamp = Some(("ANCHOR_TIMESTAMP".into(), 7)); + write( + &dir.path().join("predictions.parquet"), + vec![ + ( + "ENTITY", + Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef, + ), + ( + "TARGET_PRED", + Arc::new(Float64Array::from(vec![1.0, 3.0, 4.0])) as ArrayRef, + ), + ], + 2, + ); + write( + &dir.path().join("labels.parquet"), + vec![ + ( + "ENTITY1", + Arc::new(StringArray::from(vec!["c", "a", "b"])) as ArrayRef, + ), + ( + "TARGET1", + Arc::new(Float64Array::from(vec![5.0, 1.0, 2.0])) as ArrayRef, + ), + ], + 2, + ); + let out = evaluate(dir.path(), &spec_); + assert_eq!( + out["metrics"], + json!({"mae": 2.0 / 3.0, "rmse": (2.0f64 / 3.0).sqrt()}) + ); + assert_eq!( + actual_column::( + dir.path(), + "ANCHOR_TIMESTAMP" + ), + vec![7, 7, 7] + ); + + // Two anchors of entity 1 (a hit and a miss at rank 1) and entity 2. + spec_ = spec(EvaluateTask::Multiclass, dir.path(), 1_000_000); + spec_.prediction_columns = vec![ + ("CLASS".into(), "CLASS".into()), + ("SCORE".into(), "SCORE".into()), + ]; + let entity = Int64Array::from(vec![1, 1, 1, 1, 2, 2]); + let anchor = TimestampMicrosecondArray::from(vec![0, 0, 1, 1, 0, 0]); + write( + &dir.path().join("predictions.parquet"), + vec![ + ("ENTITY", Arc::new(entity.clone()) as ArrayRef), + ("ANCHOR_TIMESTAMP", Arc::new(anchor.clone()) as ArrayRef), + ( + "CLASS", + Arc::new(Int64Array::from(vec![10, 11, 12, 10, 11, 12])) as ArrayRef, + ), + ( + "SCORE", + Arc::new(Float64Array::from(vec![0.9, 0.5, 0.7, 0.2, 0.8, 0.1])) as ArrayRef, + ), + ], + 3, + ); + write( + &dir.path().join("labels.parquet"), + vec![ + ( + "ENTITY1", + Arc::new(Int64Array::from(vec![1, 1, 2])) as ArrayRef, + ), + ( + "TIMESTAMP", + Arc::new(TimestampNanosecondArray::from(vec![0, 1000, 0])) as ArrayRef, + ), + ( + "TARGET1", + Arc::new(Float64Array::from(vec![10.0, 10.0, 11.0])) as ArrayRef, + ), + ], + 3, + ); + let out = evaluate(dir.path(), &spec_); + assert_eq!(out["metrics"], json!({"acc": 2.0 / 3.0})); + let top: Vec<(i64, f64)> = + actual_column::(dir.path(), "CLASS") + .into_iter() + .zip(actual_column::( + dir.path(), + "SCORE", + )) + .collect(); + assert_eq!(top, vec![(10, 0.9), (12, 0.7), (11, 0.8)]); + assert_eq!( + (out["rows"].clone(), out["labelled"].clone()), + (json!(3), json!(3)) + ); + + spec_ = spec(EvaluateTask::Link, dir.path(), 1_000_000); + let mut targets = ListBuilder::new(Int64Builder::new()); + for items in [vec![11, 13], vec![], vec![11]] { + for item in items { + targets.values().append_value(item); + } + targets.append(true); + } + write( + &dir.path().join("labels.parquet"), + vec![ + ( + "ENTITY1", + Arc::new(Int64Array::from(vec![1, 1, 2])) as ArrayRef, + ), + ( + "TIMESTAMP", + Arc::new(TimestampNanosecondArray::from(vec![0, 1000, 0])) as ArrayRef, + ), + ("TARGET1", Arc::new(targets.finish()) as ArrayRef), + ], + 3, + ); + // (1, 0) ranks 10, 11 and hits 11 at rank 2 of two labels; (1, 1) + // has no labels and is skipped; entity 2 hits 11 at rank 1. + let out = evaluate(dir.path(), &spec_); + assert_eq!( + out["metrics"], + json!({"map@2": (0.5 / 2.0 + 1.0) / 2.0, "precision@2": 0.5}) + ); + } +} diff --git a/packages/diskgraph/diskgraph-pybind/src/lib.rs b/packages/diskgraph/diskgraph-pybind/src/lib.rs index cb68b9035..9929a5996 100644 --- a/packages/diskgraph/diskgraph-pybind/src/lib.rs +++ b/packages/diskgraph/diskgraph-pybind/src/lib.rs @@ -184,6 +184,7 @@ use diskgraph_engine::state::{SamplingState, StatePool}; mod artifact_helpers; mod bp_col_stats; +mod bp_evaluate; mod col_stats_export; mod custom_tt_entry; mod index_metadata; @@ -1556,6 +1557,41 @@ fn compute_col_stats_for_parquet( .map_err(map_anyhow_err) } +/// Join batch-prediction Parquet outputs with generated label Parquet parts +/// and compute realized metrics, streaming and hash-partitioning both sides. +/// +/// Args: +/// prediction_paths: the run's prediction Parquet files. +/// label_paths: the label Parquet parts. +/// spec_json: JSON object with ``task`` (``"binary"``, ``"regression"``, +/// ``"multiclass"`` or ``"link"``), ``prediction_keys`` and +/// ``label_keys`` (paired by position), ``label_target``, +/// ``spill_dir``, and optional ``actuals_dir`` with +/// ``actual_columns`` (``[label name, output name]`` pairs) to write +/// the matched label rows, and ``partition_rows``. +/// +/// Returns a JSON string ``{"rows", "labelled", "metrics"}``: distinct +/// predictions, those with a label, and the metric values (``null`` when +/// undefined). +#[pyfunction] +#[pyo3(signature = (prediction_paths, label_paths, spec_json))] +fn evaluate_prediction_parquet( + py: Python<'_>, + prediction_paths: Vec, + label_paths: Vec, + spec_json: &str, +) -> PyResult { + let spec: bp_evaluate::EvaluateSpec = serde_json::from_str(spec_json) + .map_err(|e| DiskgraphConfigError::new_err(format!("invalid spec_json: {e}")))?; + py.detach(|| -> anyhow::Result { + let predictions: Vec<&std::path::Path> = + prediction_paths.iter().map(std::path::Path::new).collect(); + let labels: Vec<&std::path::Path> = label_paths.iter().map(std::path::Path::new).collect(); + Ok(bp_evaluate::evaluate_prediction_parquet(&predictions, &labels, &spec)?.to_string()) + }) + .map_err(map_anyhow_err) +} + /// Verify a local snapshot directory matches its ``_MANIFEST.json``. /// /// Called by the Python ``DiskgraphService`` after a dataplane→local @@ -1657,6 +1693,7 @@ fn diskgraph_pybind(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(has_resumable_ingest, m)?)?; m.add_function(wrap_pyfunction!(validated_index_resume_progress, m)?)?; m.add_function(wrap_pyfunction!(compute_col_stats_for_parquet, m)?)?; + m.add_function(wrap_pyfunction!(evaluate_prediction_parquet, m)?)?; m.add_function(wrap_pyfunction!(col_stats_export::read_col_stats, m)?)?; m.add_function(wrap_pyfunction!(col_stats_export::read_edge_stats, m)?)?; m.add_class::()?; diff --git a/packages/kumo-core/src/kumo_core/pquery.py b/packages/kumo-core/src/kumo_core/pquery.py index f99177389..27ce5e2f1 100644 --- a/packages/kumo-core/src/kumo_core/pquery.py +++ b/packages/kumo-core/src/kumo_core/pquery.py @@ -20,7 +20,7 @@ from datetime import datetime, timezone from enum import Enum from numbers import Integral -from typing import Any, Mapping, Sequence +from typing import Any, Callable, Iterable, Mapping, Sequence from diskgraph import read_col_stats from kumoapi.common import ValidationResponse @@ -976,8 +976,15 @@ def latest_temporal_timestamp(query: ValidatedNativeQuery, if query.query_type != "temporal": raise PredictiveQueryPlanError( "a source-derived timestamp requires a temporal query") - return max(int(maximum.value) - for _, maximum in _target_time_bounds(query, index_path)) + return _target_maximum(query, index_path, max) + + +def _target_maximum(query: ValidatedNativeQuery, index_path: str, + pick: Callable[[Iterable[int]], int]) -> int: + """``pick`` (``min`` or ``max``) of the target aggregations' source + maxima.""" + return pick(int(maximum.value) + for _, maximum in _target_time_bounds(query, index_path)) def _temporal_timestamps( diff --git a/packages/kumo-core/src/kumo_core/training/table.py b/packages/kumo-core/src/kumo_core/training/table.py index 558e5b273..8c5fef997 100644 --- a/packages/kumo-core/src/kumo_core/training/table.py +++ b/packages/kumo-core/src/kumo_core/training/table.py @@ -453,6 +453,7 @@ def generate_train_table( request: TrainTableRequest, *, runtime: Any | None = None, + validate_targets: bool = True, ) -> GeneratedTrainTableResult: """Compile and execute the real DiskGraph TT-NG node/link producer.""" @@ -531,11 +532,12 @@ def generate_train_table( if not isinstance(table_path, str) or not table_path: raise RuntimeError("execute_tt_gen result missing tt_parquet_path") try: - _validate_classification_targets(request, staging_dir) + if validate_targets: + _validate_classification_targets(request, staging_dir) except ValueError: _discard_invalid_staging_dir(staging_dir, request) raise - top_k = getattr(request.query, "top_k", None) + top_k = getattr(request.query, "top_k", None) if validate_targets else None if (task_type in {"multilabel_classification", "multilabel_ranking"} and top_k is not None): class_count = _observed_target_class_count(staging_dir) diff --git a/packages/kumo-core/tests/test_pquery.py b/packages/kumo-core/tests/test_pquery.py index 08328085b..32fbc1777 100644 --- a/packages/kumo-core/tests/test_pquery.py +++ b/packages/kumo-core/tests/test_pquery.py @@ -15,6 +15,7 @@ _column_stats, _default_split, _split_options, + _target_maximum, _temporal_timestamps, build_generated_training_task_dataset, build_native_pquery_task_plan, @@ -751,7 +752,7 @@ def test_latest_temporal_timestamp_uses_source_maximum(tmp_path, monkeypatch): ("first", "second"), ("second", "first"), ]) -def test_latest_temporal_timestamp_is_order_independent(monkeypatch, order): +def test_target_maxima_are_order_independent(monkeypatch, order): maxima = { "first": 1_700_000_000_000_000_000, "second": 1_800_000_000_000_000_000, @@ -779,6 +780,7 @@ def stats(_index_path, table_name, _scope): assert latest_temporal_timestamp(query, "/index") == maxima["second"] assert calls == list(order) + assert _target_maximum(query, "/index", min) == maxima["first"] def test_generation_plan_keeps_earliest_anchor_needed_for_filter_history( diff --git a/packages/kumo-core/tests/test_train_table.py b/packages/kumo-core/tests/test_train_table.py index 9d156dc1d..3dd497cc1 100644 --- a/packages/kumo-core/tests/test_train_table.py +++ b/packages/kumo-core/tests/test_train_table.py @@ -160,6 +160,10 @@ def test_generated_degenerate_labels_are_rejected_before_publication( }, }, })) + labels = generate_train_table( + request(tmp_path, task_type=task_type, target_stype=target_stype), + runtime=FakeRuntime(str(staging)), validate_targets=False) + assert labels.staging_dir == str(staging) and staging.is_dir() with pytest.raises(GenerationValidationError, match=message): generate_train_table( request(tmp_path, task_type=task_type, target_stype=target_stype), diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index f99f2fe0a..625e99dcc 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -265,3 +265,94 @@ Pick one per registered model, since the two store different formats. Both recor | SageMaker | not available | SageMaker managed MLflow (`MLFLOW_TRACKING_URI` = tracking server ARN) | | Snowflake | `SnowflakeConnector` (Snowflake Model Registry) | not available | | Plain VM | not available (store under the artifact root instead) | any MLflow tracking server | + +## Monitoring prediction runs + +### Realized performance + +`trainer.evaluate(result, snapshot=newer_snapshot)` regenerates the query's labels at the run's anchors on a newer snapshot, joins them with the run's predictions on `ENTITY` (and `ANCHOR_TIMESTAMP`), and returns an `EvaluationResult`: +- `metrics`, with the fit-time `test_*` definitions: binary `auroc`/`ap` (exact, scores compared at float32 precision); regression `mae`/`rmse`; multiclass `acc` of the top-scored `CLASS`; link prediction `map@k`/`precision@k`, with `k` the classes returned per entity and anchor; +- `rows`, `labelled` and `labelled_fraction` (for link prediction, the share of entities with any future link); +- `matured`: `False` while the label window ends after the latest event of the query's target tables; `None` for static queries, which read labels at the snapshot; +- `bundle_matches`, comparing the run's model with the trainer's; +- `actuals`, a Parquet directory of the evaluated rows next to the run: `ENTITY`, `ANCHOR_TIMESTAMP` (the run's prediction time for static queries), the prediction (`True_PROB` and the predicted label `TARGET_PRED`; `TARGET_PRED`; or the top `CLASS` and its `SCORE`) and the label `TARGET`. + +`evaluation.export(PredictionOutputConfig(...))` copies the evaluated rows to a table or directory. The join runs in DiskGraph: both sides stream from Parquet and are hash-partitioned, so memory follows the partition size rather than the run size. `evaluate` never refuses; a failed step leaves `metrics` empty with a `reasons` entry. + +### Score drift + +`trainer.score_drift(result, previous=None, snapshot=None, training_snapshot=None)` returns a `DriftReport` comparing the run's scores with the model's holdout (`psi`, `ks`, `wasserstein`) and with a previous run (`previous_*`), plus per-table row and per-edge count changes between the training and scoring snapshots. The statistics are read off the recorded percentile curves (min, q01..q99, max; linear in between), with no pass over the data: + +| statistic | definition | error from the curves | +|---|---|---| +| `psi` | population stability over the baseline's deciles | — | +| `ks` | largest gap between the two CDFs | ≤ 0.02 + 2× the sketch's rank error | +| `wasserstein` | mean absolute gap between the two quantile functions, in score units | ≤ 0.01 × (sum of both score ranges) + sketch error | + +`wasserstein` is comparable over time for probabilities and scale-dependent for regression. Missing records leave a field `None` with a `reasons` entry. + +### Tracking + +With MLOps tracking on (`kumo.MLOps(tracking=...)`), each `evaluate` and `score_drift` call logs a new run in the tracker's experiment, tagged `kumo.kind=monitoring`, `kumo.prediction_run=` and `kumo.bundle_id`, with the metrics, the drift statistics and the input count changes. MLflow also tags the prediction run with `kumo.monitored=`. + +### Monitoring in a later job + +Export the run to a Parquet directory and load it back: + +```python +from kumo.engine import PredictionResult, Trainer + +trainer = Trainer.load_from_connector(store, path="models/") +run = PredictionResult.load_from_connector(store, path="predictions/2024-06-01/") +report = trainer.evaluate(run, snapshot=newer_snapshot) +drift = trainer.score_drift(run) +``` + +`load_from_connector` requires the export's `_SUCCESS` marker and fetches only its `_provenance.json`; `score_drift` needs nothing else. The Parquet parts are downloaded in parallel when `evaluate` (or `files`) first reads them, into a scratch directory removed with the result. + +### Warehouse model monitors + +Both recipes monitor the evaluated rows, exported to a warehouse table. Neither platform monitors link prediction or ranking outputs. + +**Databricks (Data Quality Monitoring, inference log; databricks-sdk ≥ 0.68).** Export to a Unity Catalog Delta table (change data feed recommended), add the model id, and create the monitor: + +```python +from databricks.sdk import WorkspaceClient +from databricks.sdk.service.dataquality import ( + AggregationGranularity, DataProfilingConfig, InferenceLogConfig, + InferenceProblemType, Monitor, Refresh) + +evaluation.export(PredictionOutputConfig(dbx, PredictionArtifactType.PREDICTIONS, "c.s.churn_evaluated")) +spark.sql("ALTER TABLE c.s.churn_evaluated ADD COLUMN model_id STRING") +spark.sql(f"UPDATE c.s.churn_evaluated SET model_id = '{evaluation.bundle_id}'") + +w = WorkspaceClient() +table_id = w.tables.get("c.s.churn_evaluated").table_id +w.data_quality.create_monitor(monitor=Monitor( + object_type="table", object_id=table_id, + data_profiling_config=DataProfilingConfig( + output_schema_id=w.schemas.get("c.s").schema_id, + assets_dir="/Workspace/Users//monitoring", + inference_log=InferenceLogConfig( + problem_type=InferenceProblemType.INFERENCE_PROBLEM_TYPE_CLASSIFICATION, + prediction_column="TARGET_PRED", label_column="TARGET", + model_id_column="model_id", timestamp_column="ANCHOR_TIMESTAMP", + granularities=[AggregationGranularity.AGGREGATION_GRANULARITY_1_DAY])))) +w.data_quality.create_refresh(object_type="table", object_id=table_id, refresh=Refresh()) +``` + +For regression use `INFERENCE_PROBLEM_TYPE_REGRESSION`; for multiclass, `prediction_column="CLASS"`. The inference log has no probability column, so `True_PROB` is profiled as an ordinary column. + +**Snowflake (Model Monitor).** The model must be a Snowflake Model Registry version, so fit with `kumo.MLOps(registry=ModelOutputConfig(sf, location=...))`. The baseline is the holdout's evaluated rows, `training_result.holdout_df()` (written with `session.write_pandas`); the source is each run's evaluated rows: + +```sql +CREATE MODEL MONITOR churn_monitor WITH + MODEL=churn VERSION=v3 FUNCTION=predict + SOURCE=churn_evaluated BASELINE=churn_holdout_evaluated + WAREHOUSE=wh REFRESH_INTERVAL='1 day' AGGREGATION_WINDOW='1 day' + TIMESTAMP_COLUMN=ANCHOR_TIMESTAMP ID_COLUMNS=(ENTITY) + PREDICTION_SCORE_COLUMNS=(True_PROB) ACTUAL_CLASS_COLUMNS=(TARGET); +``` + +- Binary classification: `PREDICTION_SCORE_COLUMNS` (in [0, 1]) with `ACTUAL_CLASS_COLUMNS`; multiclass: `PREDICTION_CLASS_COLUMNS=(CLASS)` with `ACTUAL_CLASS_COLUMNS`; regression: `PREDICTION_SCORE_COLUMNS=(TARGET_PRED)` with `ACTUAL_SCORE_COLUMNS`. +- The timestamp must be `TIMESTAMP_NTZ`, the columns free of nulls, NaN and infinities, and each column used once. Without `BASELINE` there is no drift. diff --git a/packages/kumo-engine/kumo/engine/__init__.py b/packages/kumo-engine/kumo/engine/__init__.py index 6bc34e37e..f1e8a8d1a 100644 --- a/packages/kumo-engine/kumo/engine/__init__.py +++ b/packages/kumo-engine/kumo/engine/__init__.py @@ -32,6 +32,7 @@ TrainingTableManifest, ) from kumo.engine._mlops import MLOps, mlops +from kumo.engine._monitoring import DriftReport, EvaluationResult from kumo.engine._version import __version__ from kumo.engine.config import configure from kumo.engine.connector import ( @@ -69,6 +70,8 @@ __all__ = [ "__version__", "configure", + "EvaluationResult", + "DriftReport", "EmbeddingModelStore", "GloVeVectors", "TokenEmbeddingModel", diff --git a/packages/kumo-engine/kumo/engine/_monitoring.py b/packages/kumo-engine/kumo/engine/_monitoring.py new file mode 100644 index 000000000..67f62af84 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_monitoring.py @@ -0,0 +1,470 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Realized performance and score drift of prediction runs.""" + +import dataclasses +import json +import math +import shutil +import warnings +from dataclasses import dataclass, field, replace +from pathlib import Path +from typing import Any, Dict, Mapping, Optional, Tuple +from uuid import uuid4 + +import numpy as np +import pandas as pd + +from kumo.engine._manifest import PredictionManifest + + +@dataclass(frozen=True) +class EvaluationResult: + """Metrics of a prediction run against labels from a newer snapshot. + + ``actuals`` is a Parquet directory of the evaluated rows: ``ENTITY``, + ``ANCHOR_TIMESTAMP`` (the prediction time for static queries), the + prediction (``True_PROB`` and ``TARGET_PRED``; ``TARGET_PRED``; or the + top ``CLASS`` and its ``SCORE``) and the label ``TARGET``. + """ + + prediction_id: str + bundle_id: Optional[str] + snapshot_id: str + task_type: Optional[str] + anchors: Tuple[str, ...] + metrics: Mapping[str, Optional[float]] + rows: int + labelled: int + labelled_fraction: Optional[float] + matured: Optional[bool] + bundle_matches: Optional[bool] + reasons: Mapping[str, str] = field(default_factory=dict) + actuals: Optional[str] = None + + def export(self, output_config: Any) -> None: + r"""Copy the evaluated rows to ``output_config``'s predictions + destination (a `PredictionOutputConfig`), e.g. a table to monitor.""" + if self.actuals is None: + raise ValueError(f'No evaluated rows: {dict(self.reasons)}') + sink = output_config.connector.sink_backend + destination = output_config.destinations()['predictions'] + files = sorted(str(path) for path in Path(self.actuals).glob('*.parquet')) + sink.write_files(output_config.connector, destination, files, + output_config.write_mode, + export_id=f'{self.prediction_id}-evaluation-' + f'{Path(self.actuals).name}') + + +@dataclass(frozen=True) +class DriftReport: + """Score and input drift of a prediction run; ``reasons`` explains Nones. + + ``psi``, ``ks`` and ``wasserstein`` compare the run's scores with the + holdout's; ``previous_*`` with ``previous_prediction_id``'s run. All are + computed from the recorded percentile curves (``wasserstein`` is in score + units). + """ + + prediction_id: Optional[str] + bundle_id: Optional[str] + score: Optional[str] + psi: Optional[float] + ks: Optional[float] + wasserstein: Optional[float] + holdout_truncated: Optional[bool] + bundle_matches: Optional[bool] + previous_prediction_id: Optional[str] + previous_psi: Optional[float] + previous_ks: Optional[float] + previous_wasserstein: Optional[float] + input_changes: Optional[Mapping[str, Tuple[Any, Any]]] + reasons: Mapping[str, str] + + +def evaluate(trainer: Any, run: Any, snapshot: Any) -> EvaluationResult: + query = trainer._training_query + if query is None: + raise RuntimeError('Fit or load the Trainer before evaluate()') + reasons: Dict[str, str] = {} + out: Dict[str, Any] = dict(task_type=None, anchors=(), metrics={}, + rows=0, labelled=0, matured=None, + reasons=reasons, actuals=None) + try: + _evaluate(out, trainer, run, snapshot, query) + except Exception as error: + reasons['metrics'] = f"{out['step']} failed: {error}" + out.pop('step') + return EvaluationResult( + prediction_id=run.id, + bundle_id=run.bundle_id, + snapshot_id=snapshot.id, + labelled_fraction=(out['labelled'] / + out['rows'] if out['rows'] else None), + bundle_matches=_same(run.bundle_id, trainer.bundle_id), + **out, + ) + + +_TASKS = { + 'binary_classification': 'binary', + 'regression': 'regression', + 'multiclass_classification': 'multiclass', + 'link_prediction': 'link', + 'temporal_link_prediction': 'link', + 'static_link_prediction': 'link', +} + + +def _evaluate(out: Dict[str, Any], trainer: Any, run: Any, snapshot: Any, + query: str) -> None: + import diskgraph + from kumo_core.pquery import _target_maximum + from kumo_core.training import custom_table + from kumo_core.types import PredictionArtifactType + from pyarrow import parquet + + from kumo.engine.pquery.predictive_query import PredictiveQuery + + reasons = out['reasons'] + out['step'] = 'parsing the query' + pquery = PredictiveQuery.from_snapshot(snapshot, query) + parsed = pquery._parsed + static = parsed.query_type == 'static' + out['task_type'] = parsed.task_type + + out['step'] = 'reading the predictions' + files = [ + str(path) for path in replace( + run, artifact_type=PredictionArtifactType.PREDICTIONS).files + ] + columns = parquet.read_schema(files[0]).names + manifest = run.manifest + anchors = () if static or manifest is None else tuple( + sorted({_ns(value) for value in manifest.prediction_table.anchors})) + keys = ['ENTITY'] + if len(anchors) > 1: + keys.append('ANCHOR_TIMESTAMP') + out['anchors'] = tuple( + pd.Timestamp(anchor, tz='UTC').isoformat() for anchor in anchors) + if anchors: + window_end = parsed.query.get_combined_date_offset_range( + ).end_date_offset + try: + latest = _target_maximum(parsed, str(snapshot.index_path), min) + out['matured'] = bool( + pd.Timestamp(max(anchors)) + window_end <= pd.Timestamp(latest)) + except Exception as error: + reasons['matured'] = str(error) + task = _TASKS.get(parsed.task_type) + if task is None: + reasons['metrics'] = f'{parsed.task_type} metrics are not supported' + return + if not static and not anchors: + reasons['metrics'] = 'the run records no anchor timestamps' + return + if 'ANCHOR_TIMESTAMP' in keys and 'ANCHOR_TIMESTAMP' not in columns: + reasons['metrics'] = ('the run has several anchors but its ' + 'predictions have no ANCHOR_TIMESTAMP column') + return + + out['step'] = 'generating labels' + entity, timestamp, target = (custom_table.entity_column(), + custom_table.timestamp_column(), + custom_table.target_column()) + actuals = _actuals_dir(run, snapshot) + try: + with pquery._label_parts( + snapshot, anchors, + trainer.training_table_generation_plan) as parts: + out['step'] = 'joining labels with predictions' + label_columns = parquet.read_schema(parts[0]).names + shutil.rmtree(actuals, ignore_errors=True) + emitted = [name for name in _EMITTED[task] if name in columns] + spec = { + 'task': task, + 'prediction_keys': keys, + 'label_keys': [entity, timestamp][:len(keys)], + 'label_target': target, + 'actuals_dir': str(actuals), + 'actual_columns': [[entity, 'ENTITY']] + + ([[timestamp, 'ANCHOR_TIMESTAMP']] + if timestamp in label_columns and not static else []) + + [[target, 'TARGET']], + 'prediction_columns': [[name, name] for name in emitted], + 'constant_timestamp': (['ANCHOR_TIMESTAMP', + _prediction_time(manifest)] + if static and manifest is not None + else None), + 'spill_dir': str(actuals.parent / f'.spill-{uuid4().hex}'), + } + joined = json.loads( + diskgraph.evaluate_prediction_parquet(files, parts, + json.dumps(spec))) + except Exception as error: + if (out['step'] != 'generating labels' + or 'empty training table' not in str(error)): + raise + reasons['metrics'] = 'no entity has a label at these anchors' + return + out.update(rows=joined['rows'], labelled=joined['labelled'], + actuals=str(actuals)) + if out['labelled'] == 0: + reasons['metrics'] = 'no prediction has a label' + return + metrics = out['metrics'] = { + name: _finite(value) + for name, value in joined['metrics'].items() + } + if None in metrics.values(): + reasons['metrics'] = 'a metric is undefined for these labels' + + +_EMITTED = { + 'binary': ('True_PROB', 'TARGET_PRED'), + 'regression': ('TARGET_PRED', ), + 'multiclass': ('CLASS', 'SCORE'), + 'link': (), +} + + +def _prediction_time(manifest: PredictionManifest) -> int: + anchors = manifest.prediction_table.anchors + return _ns(anchors[0] if anchors else manifest.timings.started_at) + + +def _actuals_dir(run: Any, snapshot: Any) -> Path: + from kumo.engine._scratch import get_scratch_dir + + selection = getattr(run, 'selection', None) + root = (Path(selection.pointer_path).parent if selection is not None + else get_scratch_dir() / 'evaluations' / str(run.id)) + return root / 'actuals' / str(snapshot.id).replace('/', '_') + + +def score_drift(trainer: Any, run: Any, *, previous: Any, snapshot: Any, + training_snapshot: Any) -> DriftReport: + reasons: Dict[str, str] = {} + score, summary = _score_summary(run) + model = getattr(trainer, 'manifest', None) + holdout = getattr(model, 'holdout', None) + baseline = None + if summary is None: + reasons['holdout'] = 'the run records no output score statistics' + elif holdout is None: + reasons['holdout'] = 'the model records no holdout score statistics' + else: + baseline = holdout.columns.get(score) + if _curves(summary, baseline) is None: + reasons['holdout'] = f'the holdout records no percentiles for {score}' + + earlier = None + if previous is not None: + previous_score, earlier = _score_summary(previous) + if summary is None or earlier is None: + reasons['previous'] = 'a run records no score statistics' + elif previous_score != score: + reasons['previous'] = 'the runs record different scores' + earlier = None + elif _curves(summary, earlier) is None: + reasons['previous'] = 'a run records no percentiles' + + input_changes = None + try: + counts = [s._counts for s in (training_snapshot, snapshot)] + except Exception: + reasons['input_changes'] = 'both snapshots are needed' + else: + input_changes = _count_changes(*counts) + run_bundle = getattr(run, 'bundle_id', None) + stats = (_psi, _ks, _wasserstein) + against = [[None if b is None or summary is None else stat(summary, b) + for stat in stats] for b in (baseline, earlier)] + return DriftReport( + prediction_id=getattr(run, 'id', None), + bundle_id=run_bundle, + score=score, + psi=against[0][0], + ks=against[0][1], + wasserstein=against[0][2], + holdout_truncated=None if holdout is None else holdout.truncated, + bundle_matches=_same(run_bundle, getattr(trainer, 'bundle_id', None)), + previous_prediction_id=getattr(previous, 'id', None), + previous_psi=against[1][0], + previous_ks=against[1][1], + previous_wasserstein=against[1][2], + input_changes=input_changes, + reasons=reasons, + ) + + +def track(report: Any, run: Any, settings: Any) -> None: + r"""Log ``report`` as a new monitoring run of ``run``; best-effort.""" + tracking = None if settings is None else settings.tracking + if tracking is None: + return + from kumo.engine._scratch import new_run_id + from kumo.engine._tracking import RunRecord, log_run + + try: + values = {**report.metrics} if isinstance( + report, EvaluationResult) else {} + for name, value in dataclasses.asdict(report).items(): + if name == 'input_changes': + for key, (before, after) in (value or {}).items(): + values[f'input.{key}.before'] = before + values[f'input.{key}.after'] = after + elif name not in ('metrics', 'reasons', 'actuals'): + values[name] = value + prediction_id = getattr(run, 'id', None) + record = RunRecord( + name=f'monitor-{prediction_id}', + params={ + key: str(value) + for key, value in values.items() + if isinstance(value, str) + }, + metrics={ + key: float(value) + for key, value in values.items() + if isinstance(value, (int, float)) + and not isinstance(value, bool) and math.isfinite(value) + }, + tags={ + key: str(value) + for key, value in ( + ('kumo.kind', 'monitoring'), + ('kumo.run_id', f'{prediction_id}-monitor-{new_run_id()}'), + ('kumo.prediction_run', prediction_id), + ('kumo.bundle_id', getattr(run, 'bundle_id', None)), + ) if value is not None + }, + ) + except Exception as error: + warnings.warn(f'Could not track the monitoring run: {error}', + UserWarning, stacklevel=3) + return + log_run(tracking, 'log_monitoring', record) + + +def _count_changes(before: Mapping[str, Mapping[str, int]], + after: Mapping[str, Mapping[str, int]] + ) -> Dict[str, Tuple[Any, Any]]: + changes = {} + for kind, key in (('rows', 'tables.{}.rows'), ('edges', 'edges.{}')): + old, new = before.get(kind) or {}, after.get(kind) or {} + for name in sorted(set(old) | set(new)): + if not name.endswith('_reverse') and old.get(name) != new.get(name): + changes[key.format(name)] = (old.get(name), new.get(name)) + return changes + + +def _score_summary(run: Any) -> Tuple[Optional[str], Optional[Mapping]]: + try: + columns = run.manifest.stats['predictions'].columns + return next(iter(columns.items())) + except Exception: + return None, None + + +def _psi(actual: Mapping[str, Any], + expected: Optional[Mapping[str, Any]]) -> Optional[float]: + r"""PSI over ``expected``'s deciles; each side's bin shares are read off + its percentile curve (linear between min, q01..q99 and max).""" + curves = _curves(actual, expected) + if curves is None: + return None + edges = np.unique(curves[1][0][10:91:10]) + + def shares(values: np.ndarray, levels: np.ndarray) -> np.ndarray: + upper = np.searchsorted(values, edges, side='right') + lower = np.maximum(upper - 1, 0) + upper = np.minimum(upper, len(values) - 1) + width = values[upper] - values[lower] + step = np.divide(edges - values[lower], width, + out=np.ones_like(edges), where=width > 0) + cdf = levels[lower] + np.clip(step, 0, 1) * (levels[upper] - + levels[lower]) + return np.maximum(np.diff(np.concatenate([[0], cdf, [1]])), 1e-6) + + actual_shares, expected_shares = shares(*curves[0]), shares(*curves[1]) + return _finite( + np.sum((actual_shares - expected_shares) * + np.log(actual_shares / expected_shares))) + + +def _ks(a: Mapping[str, Any], b: Optional[Mapping[str, Any]]) -> Optional[float]: + r"""sup |F_a - F_b| of the percentile curves, at every knot from both + sides (point masses at repeated knots).""" + curves = _curves(a, b) + if curves is None: + return None + points = np.union1d(curves[0][0], curves[1][0]) + + def cdf(values: np.ndarray, levels: np.ndarray, side: str) -> np.ndarray: + index = np.searchsorted(values, points, side=side) + lower = np.clip(index - 1, 0, len(values) - 1) + upper = np.clip(index, 0, len(values) - 1) + width = values[upper] - values[lower] + step = np.divide(points - values[lower], width, + out=np.zeros_like(points), where=width > 0) + inside = levels[lower] + step * (levels[upper] - levels[lower]) + return np.where(index == 0, 0.0, + np.where(index == len(values), 1.0, inside)) + + return _finite( + max( + np.max(np.abs(cdf(*curves[0], side) - cdf(*curves[1], side))) + for side in ('left', 'right'))) + + +def _wasserstein(a: Mapping[str, Any], + b: Optional[Mapping[str, Any]]) -> Optional[float]: + r"""Integral of |Q_a - Q_b| over levels, exact for the linear curves.""" + curves = _curves(a, b) + if curves is None: + return None + d = curves[0][0] - curves[1][0] + left, right = d[:-1], d[1:] + size = np.abs(left) + np.abs(right) + crossing = np.divide(left**2 + right**2, 2 * size, out=np.zeros_like(size), + where=size > 0) + segments = np.where(left * right >= 0, size / 2, crossing) + return _finite(0.01 * np.sum(segments)) + + +def _curves(a: Mapping[str, Any], + b: Optional[Mapping[str, Any]]) -> Optional[tuple]: + try: + curves = (_percentile_curve(a), _percentile_curve(b)) + except Exception: + return None + return None if None in curves else curves + + +def _percentile_curve(summary: Optional[Mapping[str, Any]]) -> Optional[tuple]: + percentiles = (summary or {}).get('percentiles') + if not percentiles or len(percentiles) != 99: + return None + values = np.maximum.accumulate( + np.array([summary['min'], *percentiles, summary['max']], dtype=float)) + return values, np.linspace(0, 1, 101) + + +def _ns(value: Any) -> int: + timestamp = pd.Timestamp(value) + if timestamp.tzinfo is not None: + timestamp = timestamp.tz_convert(None) + return timestamp.value + + +def _same(a: Optional[str], b: Optional[str]) -> Optional[bool]: + return None if a is None or b is None else a == b + + +def _finite(value: Any) -> Optional[float]: + if value is None: + return None + value = float(value) + return value if math.isfinite(value) else None diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py index aee04c13d..ab249a5b0 100644 --- a/packages/kumo-engine/kumo/engine/_tracking.py +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -411,6 +411,35 @@ def _call(self, call: Callable[[Tracking], Any]) -> None: _abandon(tracker) +def log_run(tracking: Any, hook: str, record: RunRecord) -> None: + r"""Log a standalone run with ``tracking.(record)`` on a tracking + thread, waiting up to `FINISH_SECONDS`; never raises.""" + method = getattr(tracking, hook, None) + if method is None or _rank() != 0: + return + + def log() -> None: + try: + method(record) + except Exception: + logger.warning('Experiment tracking failed; the %s run is not ' + 'tracked', record.tags.get('kumo.kind'), + exc_info=True) + + try: + getattr(tracking, 'prepare', _none)() + except Exception: + logger.warning('Experiment tracking failed', exc_info=True) + return + worker = threading.Thread(target=log, name='kumo-tracking', daemon=True) + worker.start() + worker.join(FINISH_SECONDS) + if worker.is_alive(): + logger.warning( + 'Experiment tracking did not finish within %ss; it goes on in ' + 'the background', FINISH_SECONDS) + + def _none(*args: Any) -> None: return None diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 3d4c45901..0b4c6ee5d 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -142,7 +142,7 @@ def start(self, record: RunRecord) -> None: self.experiment_name)) run = mlflow.start_run( experiment_id=experiment_id, - run_name=self.run_name, + run_name=record.name or self.run_name, parent_run_id=self._outer, tags=record.tags, log_system_metrics=self.log_system_metrics, @@ -153,6 +153,18 @@ def start(self, record: RunRecord) -> None: self._log_artifacts(record.artifacts) _log_inputs(mlflow, self._client, self._parent, record.inputs) + def log_monitoring(self, record: RunRecord) -> None: + self.start(record) + experiment_id, run_id = self._experiment_id, self.run_id + self.finish(record, None) + client = self._mlflow.MlflowClient() + for run in client.search_runs( + [experiment_id], filter_string=( + "tags.`kumo.kind` = 'prediction' AND tags.`kumo.run_id` " + f"= '{record.tags['kumo.prediction_run']}'"), + max_results=1): + client.set_tag(run.info.run_id, 'kumo.monitored', run_id) + def on_trial(self, record: RunRecord, event: Event) -> None: key = (event.trial or 0, int(event.params['attempt'])) if event.name == 'trial_started': diff --git a/packages/kumo-engine/kumo/engine/integrations/snowflake.py b/packages/kumo-engine/kumo/engine/integrations/snowflake.py index 4d78d5b56..75ba30e71 100644 --- a/packages/kumo-engine/kumo/engine/integrations/snowflake.py +++ b/packages/kumo-engine/kumo/engine/integrations/snowflake.py @@ -95,7 +95,8 @@ def start(self, record: RunRecord) -> None: elif (not hasattr(tracking, '_run') or tracking._run is not None or getattr(tracking, '_session', None) is not self.session): logger.warning('Snowflake ExperimentTracking is in use in this ' - 'process; this fit is not tracked in Snowflake') + 'process; this %s run is not tracked in Snowflake', + record.tags.get('kumo.kind')) return self._tracking = tracking database = self.database or self.session.get_current_database() @@ -105,8 +106,9 @@ def start(self, record: RunRecord) -> None: database_name=database, schema_name=schema or 'PUBLIC') self._run = tracking.start_run( - self._new_run_name(tracking, experiment.name, - record.tags['kumo.run_id'])) + self._new_run_name( + tracking, experiment.name, record.tags['kumo.run_id'], + f'"{record.name}"' if record.name else self.run_name)) self.run_id = str(self._run.name) if self.native_model_lineage: self._run_context.enter_context(self._run) @@ -114,6 +116,10 @@ def start(self, record: RunRecord) -> None: self._log_params({**record.params, **record.tags}) self._log_artifacts(record.artifacts) + def log_monitoring(self, record: RunRecord) -> None: + self.start(record) + self.finish(record, None) + def on_trial(self, record: RunRecord, event: Event) -> None: if self._run is None: return @@ -161,9 +167,9 @@ def finish(self, record: Optional[RunRecord], self._run_context.close() self._reset() - def _new_run_name(self, tracking: Any, experiment: Any, - run_id: str) -> Optional[str]: - if self.run_name is None: + def _new_run_name(self, tracking: Any, experiment: Any, run_id: str, + run_name: Optional[str]) -> Optional[str]: + if run_name is None: return None # start_run(name) resumes, or raises on, any run matching name as a # SQL LIKE pattern and cannot create an existing run, so pick a name @@ -174,16 +180,16 @@ def _new_run_name(self, tracking: Any, experiment: Any, experiment_name=experiment) ] short = run_id.rsplit('-', 1)[-1] - for name in (self.run_name, _suffixed(self.run_name, short[:8]), - _suffixed(self.run_name, short)): + for name in (run_name, _suffixed(run_name, short[:8]), + _suffixed(run_name, short)): if not _taken(name, existing): break else: name = None - if name != self.run_name: + if name != run_name: logger.warning( 'Snowflake run %s already exists in experiment %s; recording ' - 'this fit as %s', self.run_name, self.experiment_name, + 'this run as %s', run_name, self.experiment_name, name or 'a generated run name') return name diff --git a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py index 7ac4a9d8d..601cd2f59 100644 --- a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py +++ b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py @@ -42,12 +42,13 @@ """ import logging +from contextlib import contextmanager from copy import deepcopy from dataclasses import dataclass, replace from datetime import datetime from functools import cached_property from pathlib import Path -from typing import Any, Optional +from typing import Any, Iterator, Optional from kumo_core.model_plan import ModelPlan as CoreModelPlan from kumo_core.model_plan import ModelPlanValueError @@ -290,25 +291,10 @@ def generate_training_table( if errors: raise ModelPlanValueError(errors) - run_id = new_run_id() - run_dir = train_table_scratch_dir(run_id) + run_dir = train_table_scratch_dir(new_run_id()) run_dir.mkdir(parents=True, exist_ok=True) - - workflow = TrainingWorkflowRequest.from_spec( - graph=snapshot.graph_config, - predictive_query=self.query, - prepared_graph=PreparedGraph( - index_id=snapshot.id, - index_path=str(snapshot.index_path), - ), - run_id=run_id, - scratch_dir=str(run_dir), - artifact_dir=str(run_dir), - training_table_generation_plan=options, - ) - with no_encoder_inference_warnings(): - task_plan = build_training_task_plan(workflow) - generated = generate_train_table(task_plan.train_table) + workflow, task_plan, generated = self._generate( + snapshot, options, run_dir) generation_plan = options if workflow.query.query_type == 'temporal': @@ -341,6 +327,62 @@ def generate_training_table( ) return _finish(table, snapshot, task_plan.train_table, mlops) + def _generate(self, snapshot: GraphSnapshot, options: CoreGenerationPlan, + run_dir: Path, *, timestamps: tuple[int, ...] = (), + validate_targets: bool = True) -> tuple: + workflow = TrainingWorkflowRequest.from_spec( + graph=snapshot.graph_config, + predictive_query=self.query, + prepared_graph=PreparedGraph( + index_id=snapshot.id, + index_path=str(snapshot.index_path), + ), + run_id=run_dir.name, + scratch_dir=str(run_dir), + artifact_dir=str(run_dir), + training_table_generation_plan=options, + timestamps=timestamps, + ) + with no_encoder_inference_warnings(): + task_plan = build_training_task_plan(workflow) + return workflow, task_plan, generate_train_table( + task_plan.train_table, validate_targets=validate_targets) + + @contextmanager + def _label_parts( + self, snapshot: GraphSnapshot, timestamps: tuple[int, ...], + plan: Optional[TrainingTableGenerationPlan]) -> Iterator[list]: + r"""Generate this query's labels at ``timestamps`` on ``snapshot`` + and yield their Parquet parts, removed on exit.""" + import shutil + + from kumo_core.io.json import from_json_dict + from kumo_core.training.export import read_spec, stage_parts + + # The fit's split ranges would not cover later anchors. + options = CoreGenerationPlan() + if plan is not None: + options = replace(from_json_dict(to_json_dict(plan), + CoreGenerationPlan), + split=options.split) + run_dir = train_table_scratch_dir(new_run_id()) + run_dir.mkdir(parents=True, exist_ok=True) + staging = None + try: + _, _, generated = self._generate(snapshot, options, run_dir, + timestamps=timestamps, + validate_targets=False) + staging = Path(generated.staging_dir) + yield [ + str(staging / part) + for parts in stage_parts(read_spec(str(staging))).values() + for part in parts + ] + finally: + if staging is not None: + shutil.rmtree(staging, ignore_errors=True) + shutil.rmtree(run_dir, ignore_errors=True) + def generate_prediction_table( self, snapshot: GraphSnapshot, diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 96d3cf07c..7d9dab289 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -16,6 +16,7 @@ from datetime import datetime, timezone from pathlib import Path from typing import ( + TYPE_CHECKING, Any, Callable, Dict, @@ -88,6 +89,9 @@ PredictionOutputConfig, ) +if TYPE_CHECKING: + from kumo.engine._monitoring import DriftReport, EvaluationResult + logger = logging.getLogger(__name__) @@ -171,6 +175,20 @@ def with_scoring_defaults(self, assert result is not None return result + def evaluate(self, *args: Any, **kwargs: Any) -> 'EvaluationResult': + """Evaluate a run of this model, as :meth:`Trainer.evaluate`.""" + return self._model('evaluate').evaluate(*args, **kwargs) + + def score_drift(self, *args: Any, **kwargs: Any) -> 'DriftReport': + """Score drift of a run of this model, as :meth:`Trainer.score_drift`.""" + return self._model('score_drift').score_drift(*args, **kwargs) + + def _model(self, call: str) -> 'Trainer': + if self._trainer is None: + raise RuntimeError(f'This result has no model to {call} with; ' + 'use Trainer.load on its bundle_dir.') + return self._trainer + def holdout_df(self) -> pd.DataFrame: r"""Read the holdout predictions as a pandas DataFrame. @@ -200,15 +218,75 @@ class PredictionResult: be written), from which ``bundle_id`` and ``snapshot_id`` are read. """ - selection: PredictionSelectionRef + selection: Optional[PredictionSelectionRef] artifact_type: PredictionArtifactType + _exported: Optional['_Export'] = field(default=None, compare=False, + repr=False) + + @classmethod + def load_from_connector(cls, connector: Connector, *, + path: str) -> 'PredictionResult': + r"""Load a Parquet directory written by :meth:`export` for monitoring. + + The export must be complete (its ``_SUCCESS`` marker present). Only + its ``_provenance.json`` is fetched here; the Parquet parts are + downloaded, in parallel, when ``files`` is first read, into a local + directory removed when the result is garbage-collected. + """ + import tempfile + import weakref + + from pyarrow.fs import FileType + + from kumo.engine._manifest import read_manifest + from kumo.engine._scratch import get_scratch_dir + from kumo.engine.connector import file_store + + store = file_store(connector, path) + if store is None: + raise TypeError(f'{path} is not a file destination of this ' + f'connector') + filesystem, root = store + root = root.rstrip('/') + if filesystem.get_file_info( + f'{root}/_SUCCESS').type == FileType.NotFound: + raise FileNotFoundError( + f'{path} has no _SUCCESS marker: the export is incomplete or ' + f'still running') + scratch = get_scratch_dir() + scratch.mkdir(parents=True, exist_ok=True) + local = Path(tempfile.mkdtemp(prefix='fetched-predictions-', + dir=scratch)) + try: + with filesystem.open_input_stream( + f'{root}/{PROVENANCE_FILE}') as reader: + (local / PROVENANCE_FILE).write_bytes(reader.read()) + except (FileNotFoundError, OSError): + pass + manifest = read_manifest(local / PROVENANCE_FILE, PredictionManifest) + here = str(_run_manifest.destination_record(connector, + path).location).rstrip('/') + kind = next((PredictionArtifactType(record.artifact.upper()) + for record in (manifest and manifest.outputs or []) + if str(record.destination.location).rstrip('/') == here), + PredictionArtifactType.EMBEDDINGS + if root.rsplit('/', 1)[-1] == 'embeddings' else + PredictionArtifactType.PREDICTIONS) + exported = _Export(filesystem, root, local, kind) + weakref.finalize(exported, shutil.rmtree, str(local), True) + return cls(None, kind, _exported=exported) @property - def id(self) -> str: - return self.selection.job_id + def id(self) -> Optional[str]: + if self.selection is not None: + return self.selection.job_id + manifest = self.manifest + return None if manifest is None else manifest.id @property def _manifest_path(self) -> Path: + if self._exported is not None: + return self._exported.local / PROVENANCE_FILE return Path(self.selection.pointer_path).parent / PROVENANCE_FILE @property @@ -229,6 +307,11 @@ def snapshot_id(self) -> Optional[str]: @property def num_rows(self) -> int: + if self._exported is not None: + from pyarrow import parquet + return sum( + parquet.ParquetFile(path).metadata.num_rows + for path in self.files) from kumo_core.train import artifact_manifest selected = artifact_manifest.load_prediction_selection( self.selection, validate_objects=False) @@ -238,6 +321,11 @@ def num_rows(self) -> int: @property def files(self) -> tuple[str, ...]: from kumo_core.train import artifact_manifest + if self._exported is not None: + if self._exported.kind is not self.artifact_type: + raise ValueError( + f'This export holds {self._exported.kind.value.lower()}') + return self._exported.files() resolved = artifact_manifest.resolve_prediction_selection_artifact( self.selection, self.artifact_type) if resolved is None: @@ -286,6 +374,36 @@ def _export(self, output_config: PredictionOutputConfig, written.append(records[index]) +@dataclass(frozen=True) +class _Export: + filesystem: Any + root: str + local: Path + kind: PredictionArtifactType + + def files(self) -> tuple[str, ...]: + parts = self.local / 'parts' + if not parts.is_dir(): + from kumo_core.sources.transfer import StagedFile, fetch_all + from pyarrow.fs import FileSelector, FileType + + def fetch(staged: StagedFile, target: str) -> None: + with self.filesystem.open_input_stream(staged.path) as reader, \ + open(target, 'wb') as writer: + shutil.copyfileobj(reader, writer, 1 << 20) + + partial = self.local / 'parts.partial' + shutil.rmtree(partial, ignore_errors=True) + listed = self.filesystem.get_file_info(FileSelector(self.root)) + fetch_all([ + StagedFile(info.path, info.size) for info in listed + if info.type == FileType.File and info.path.endswith( + '.parquet') + ], str(partial), fetch=fetch) + partial.rename(parts) + return tuple(sorted(map(str, parts.glob('*.parquet')))) + + def _sidecar_bytes(manifest: Optional[PredictionManifest]) -> Optional[bytes]: from kumo.engine._manifest import manifest_bytes @@ -1249,6 +1367,49 @@ def _predict( _run_manifest.finish(result, manifest, output_config, settings) return result + def evaluate(self, run: PredictionResult, *, snapshot: GraphSnapshot, + mlops: Optional[Any] = None) -> 'EvaluationResult': + r"""Score a run against labels regenerated on a newer snapshot. + + With MLOps tracking on (``mlops`` or a `kumo.mlops` block), the + result is logged as a new monitoring run of ``run``. + """ + from kumo.engine._mlops import resolve + from kumo.engine._monitoring import evaluate, track + + settings = resolve(mlops) + result = evaluate(self, run, snapshot) + track(result, run, settings) + return result + + def score_drift( + self, + run: PredictionResult, + *, + previous: Optional[PredictionResult] = None, + snapshot: Optional[GraphSnapshot] = None, + training_snapshot: Optional[GraphSnapshot] = None, + mlops: Optional[Any] = None, + ) -> 'DriftReport': + r"""Compare a run's scores with the holdout and a previous run. + + Args: + previous: An earlier run, for run-over-run statistics. + snapshot: The snapshot the run scored, with + ``training_snapshot`` for input row-count changes. + training_snapshot: The snapshot this model trained on. + mlops: Logs the report as a new monitoring run of ``run`` when + tracking is on, like `evaluate`. + """ + from kumo.engine._mlops import resolve + from kumo.engine._monitoring import score_drift, track + + settings = resolve(mlops) + report = score_drift(self, run, previous=previous, snapshot=snapshot, + training_snapshot=training_snapshot) + track(report, run, settings) + return report + @classmethod def load( cls, diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py index 1d4f639d5..1931480d0 100644 --- a/packages/kumo-engine/test/test_mlflow_tracking.py +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -314,3 +314,32 @@ def test_fit_inputs_come_from_the_embedded_parents(): 'training_table': {**entry, 'version': None}} assert fit_inputs(stub('model', 'm')) == {} assert fit_inputs(None) == {} + + +def test_a_monitoring_run_is_new_and_tags_its_prediction_run(tracking_uri): + from kumo.engine._tracking import RunRecord, log_run + from kumo.engine.integrations.mlflow import MlflowTracking + + client = mlflow.MlflowClient() + experiment = client.create_experiment('scores') + predicted = client.create_run(experiment, tags={ + 'kumo.kind': 'prediction', 'kumo.run_id': 'pred-1' + }).info.run_id + client.set_terminated(predicted) + + log_run( + MlflowTracking(experiment_name='scores', log_system_metrics=False), + 'log_monitoring', + RunRecord(name='monitor-pred-1', metrics={'psi': 0.1, 'ks': 0.2}, + tags={'kumo.kind': 'monitoring', 'kumo.run_id': 'm-1', + 'kumo.prediction_run': 'pred-1', + 'kumo.bundle_id': 'b'})) + + monitor, = [run for run in _runs(client) + if run.data.tags['kumo.kind'] == 'monitoring'] + assert monitor.info.run_name == 'monitor-pred-1' + assert monitor.info.status == 'FINISHED' + assert monitor.data.metrics == {'psi': 0.1, 'ks': 0.2} + assert monitor.data.tags['kumo.prediction_run'] == 'pred-1' + assert client.get_run(predicted).data.tags['kumo.monitored'] == ( + monitor.info.run_id) diff --git a/packages/kumo-engine/test/test_monitoring.py b/packages/kumo-engine/test/test_monitoring.py new file mode 100644 index 000000000..55401d113 --- /dev/null +++ b/packages/kumo-engine/test/test_monitoring.py @@ -0,0 +1,426 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import gc +import math +import os +import time +from dataclasses import dataclass, replace +from types import SimpleNamespace +from typing import Any, Optional + +import numpy as np +import pandas as pd +import pytest +from kumoapi.graph import Edge +from kumoapi.model_plan import ModelPlan +from kumoapi.typing import Stype + +import kumo.engine as kumo +from kumo.engine._manifest import ArtifactStats +from kumo.engine._monitoring import _ks, _psi, _wasserstein, evaluate +from kumo.engine._scratch import get_scratch_dir +from kumo.engine.trainer.trainer import PredictionResult, TrainingResult + +CUTOFF = pd.Timestamp('2024-03-31') +ANCHOR = CUTOFF - pd.Timedelta(days=1) +PREDICTED_AT = pd.Timestamp('2024-05-01', tz='UTC') + + +@pytest.fixture(scope='module') +def snapshots(tmp_path_factory): + root = tmp_path_factory.mktemp('monitoring') + previous = get_scratch_dir() + kumo.configure(scratch_dir=root / 'scratch') + rng = np.random.default_rng(0) + data = root / 'data' + data.mkdir() + flag = pd.Series(rng.integers(0, 2, 60).astype(bool), dtype='boolean') + orders = pd.DataFrame({ + 'EID': np.arange(1500), + 'USER_ID': rng.integers(0, 60, 1500), + 'AMT': rng.gamma(2.0, 10.0, 1500), + 'TS': pd.Timestamp('2024-01-01') + + pd.to_timedelta(rng.integers(0, 120, 1500), unit='D'), + }) + + def write(users_flag, rows): + pd.DataFrame({ + 'ID': np.arange(60), + 'FLAG': users_flag + }).to_parquet(data / 'users.parquet') + rows.to_parquet(data / 'orders.parquet') + + hidden = flag.copy() + hidden[:12] = pd.NA + write(hidden, orders[orders.TS < CUTOFF]) + # File versions are modification times in seconds. + for path in data.iterdir(): + os.utime(path, (time.time() - 3600, ) * 2) + connector = kumo.FileConnector(root_dir=str(data)) + users = kumo.Table.from_source_table(connector['users.parquet'], + primary_key='ID').infer_metadata() + events = kumo.Table.from_source_table( + connector['orders.parquet'], primary_key='EID', + time_column='TS').infer_metadata().with_column('USER_ID', + stype=Stype.ID) + graph = kumo.Graph( + tables={ + 'users': users, + 'orders': events + }, edges=[Edge(src_table='orders', fkey='USER_ID', dst_table='users')]) + before = graph.snapshot() + write(flag, orders) + after = graph.snapshot() + yield SimpleNamespace(graph=graph, before=before, after=after, + orders=orders, flag=flag) + kumo.configure(scratch_dir=previous) + + +@pytest.fixture(scope='module') +def classifier(snapshots): + pquery = kumo.PredictiveQuery( + graph=snapshots.graph, query='PREDICT COUNT(orders.*, 0, 7, days) > 0 ' + 'FOR EACH users.ID') + plan = ModelPlan() + plan.training_job.disable_explain = True + plan.optimization.max_epochs = 1 + plan.optimization.max_steps_per_epoch = 2 + trainer = kumo.Trainer(model_plan=plan) + trainer.fit(pquery.generate_training_table(snapshots.before), + snapshot=snapshots.before) + table = pquery.generate_prediction_table(snapshots.before) + return trainer, trainer.predict(table, snapshot=snapshots.before, + binary_classification_threshold=0.5) + + +@dataclass +class _Run: + frame: pd.DataFrame + artifact_type: Any = None + manifest: Any = None + id: str = 'run' + bundle_id: Optional[str] = None + + @property + def files(self): + path = get_scratch_dir() / f'{id(self.frame)}.parquet' + self.frame.to_parquet(path) + return [path] + + +def _evaluate(snapshots, query, frame, anchors=()): + trainer = SimpleNamespace(_training_query=query, bundle_id=None, + training_table_generation_plan=None) + manifest = SimpleNamespace( + prediction_table=SimpleNamespace(anchors=[str(a) for a in anchors]), + timings=SimpleNamespace(started_at=PREDICTED_AT.isoformat())) + return evaluate(trainer, _Run(frame, manifest=manifest), snapshots.after) + + +def _window(orders, anchor): + rows = orders[(orders.TS > anchor) + & (orders.TS <= anchor + pd.Timedelta(days=7))] + return rows.groupby('USER_ID') + + +def test_evaluate_matured_run(snapshots, classifier): + trainer, run = classifier + result = trainer.evaluate(run, snapshot=snapshots.after) + + assert pd.Timestamp(result.anchors[0]).tz_convert(None) == ANCHOR + assert (result.matured, result.bundle_matches, + result.labelled_fraction) == (True, True, 1.0) + assert result.rows == result.labelled == 60 and result.reasons == {} + active = set(_window(snapshots.orders, ANCHOR).groups) + rows = pd.read_parquet(result.actuals).set_index('ENTITY') + labels = rows['TARGET'] + assert labels.to_dict() == {user: user in active for user in range(60)} + expected = run.frame.set_index('ENTITY').loc[rows.index] + assert (rows['True_PROB'] == expected['True_PROB']).all() + assert (rows['TARGET_PRED'] == expected['TARGET_PRED']).all() + + from sklearn.metrics import roc_auc_score + joined = run.frame.set_index('ENTITY').join(labels) + assert result.metrics['auroc'] == pytest.approx( + roc_auc_score(joined['TARGET'], joined['True_PROB']), abs=1e-6) + + immature = trainer.evaluate(run, snapshot=snapshots.before) + assert immature.matured is False + assert immature.metrics == {'auroc': None, 'ap': None} + assert 'undefined' in immature.reasons['metrics'] + + +def test_evaluate_joins_several_anchors(snapshots, classifier, monkeypatch): + trainer, run = classifier + anchors = [ANCHOR - pd.Timedelta(days=14), ANCHOR] + stacked = _Run(pd.concat([ + run.frame.assign(ANCHOR_TIMESTAMP=anchor.tz_localize('UTC')) + for anchor in anchors + ], ignore_index=True)) + recorded = run.manifest + monkeypatch.setattr(PredictionResult, 'files', + property(lambda _: stacked.files)) + monkeypatch.setattr( + PredictionResult, 'manifest', property(lambda _: replace( + recorded, prediction_table=replace( + recorded.prediction_table, anchors=[ + anchor.tz_localize('UTC').isoformat() + for anchor in anchors + ])))) + result = trainer.evaluate(run, snapshot=snapshots.after) + + assert (result.rows, result.labelled, result.matured) == (120, 120, True) + labels = pd.read_parquet(result.actuals).set_index( + ['ANCHOR_TIMESTAMP', 'ENTITY'])['TARGET'] + for anchor in anchors: + active = set(_window(snapshots.orders, anchor).groups) + assert labels[anchor].to_dict() == { + user: user in active + for user in range(60) + } + + monkeypatch.undo() + monkeypatch.setattr(PredictionResult, 'manifest', property(lambda _: None)) + result = trainer.evaluate(run, snapshot=snapshots.after) + assert result.reasons['metrics'] == 'the run records no anchor timestamps' + assert (result.metrics, result.actuals) == ({}, None) + + +def test_evaluate_regression_and_static_queries(snapshots): + frame = pd.DataFrame({'ENTITY': np.arange(60), 'TARGET_PRED': 10.0}) + result = _evaluate(snapshots, + 'PREDICT SUM(orders.AMT, 0, 7, days) FOR EACH users.ID', + frame, [ANCHOR]) + sums = _window(snapshots.orders, ANCHOR)['AMT'].sum() + error = 10.0 - pd.Series({user: sums.get(user, 0.0) for user in range(60)}) + assert result.matured and result.labelled_fraction == 1.0 + assert result.metrics['mae'] == pytest.approx(error.abs().mean(), rel=1e-5) + assert result.metrics['rmse'] == pytest.approx( + math.sqrt((error**2).mean()), rel=1e-5) + + frame = pd.DataFrame({'ENTITY': np.arange(12), 'True_PROB': .5}) + result = _evaluate(snapshots, 'PREDICT users.FLAG FOR EACH users.ID', + frame) + assert (result.matured, result.anchors) == (None, ()) + assert result.rows == 12 and result.labelled_fraction == 1.0 + rows = pd.read_parquet(result.actuals).set_index('ENTITY') + assert all(rows['TARGET'][user] == snapshots.flag[user] + for user in range(12)) + assert (rows['ANCHOR_TIMESTAMP'] == PREDICTED_AT).all() + assert (rows['True_PROB'] == .5).all() + + +def test_evaluation_exports_its_rows(snapshots, classifier, tmp_path): + from kumo.engine.trainer.output_config import PredictionOutputConfig + + trainer, run = classifier + result = trainer.evaluate(run, snapshot=snapshots.after) + result.export(PredictionOutputConfig( + kumo.FileConnector(root_dir=str(tmp_path)), + output_paths={'predictions': 'evaluated/'})) + pd.testing.assert_frame_equal( + pd.read_parquet(tmp_path / 'evaluated'), + pd.read_parquet(result.actuals)) + + +def _join(tmp_path, task, predictions, labels): + import json + + import diskgraph + + predictions.to_parquet(tmp_path / 'predictions.parquet', + row_group_size=97) + labels.to_parquet(tmp_path / 'labels.parquet', row_group_size=89) + spec = {'task': task, 'prediction_keys': ['ENTITY'], + 'label_keys': ['ENTITY1'], 'label_target': 'TARGET1', + 'spill_dir': str(tmp_path / 'spill'), 'partition_rows': 500} + return json.loads(diskgraph.evaluate_prediction_parquet( + [str(tmp_path / 'predictions.parquet')], + [str(tmp_path / 'labels.parquet')], json.dumps(spec))) + + +def test_realized_metrics_match_fit_time_definitions(tmp_path): + import torch + from kumo_core.evaluation.classification import ( + BinaryAUROC, + BinaryAveragePrecision, + ) + from torch_geometric.metrics import LinkPredMAP, LinkPredPrecision + + rng = np.random.default_rng(0) + n = 2000 + score = np.round(rng.random(n), 2).astype(np.float32) + target = rng.random(n) < score + joined = _join(tmp_path, 'binary', + pd.DataFrame({'ENTITY': np.arange(n), 'True_PROB': score}), + pd.DataFrame({'ENTITY1': np.arange(n)[::-1], + 'TARGET1': target[::-1]})) + for name, metric in (('auroc', BinaryAUROC()), + ('ap', BinaryAveragePrecision())): + metric.update(torch.tensor(score), torch.tensor(target).long()) + assert joined['metrics'][name] == pytest.approx( + float(metric.compute()), abs=1e-6) + + entities, k, items = 300, 5, 40 + ranked = -np.sort(-rng.random((entities, k)), axis=1) + classes = np.stack([rng.choice(items, k, replace=False) + for _ in range(entities)]) + truth = [rng.choice(items, rng.integers(0, 4), replace=False) + for _ in range(entities)] + joined = _join( + tmp_path, 'link', + pd.DataFrame({'ENTITY': np.repeat(np.arange(entities), k), + 'CLASS': classes.ravel(), 'SCORE': ranked.ravel()}), + pd.DataFrame({'ENTITY1': np.arange(entities), + 'TARGET1': [list(t) for t in truth]})) + index = torch.tensor(np.stack([ + np.repeat(np.arange(entities), [len(t) for t in truth]), + np.concatenate(truth)])) + for name, metric in ((f'map@{k}', LinkPredMAP(k)), + (f'precision@{k}', LinkPredPrecision(k))): + metric.update(torch.tensor(classes), index) + assert joined['metrics'][name] == pytest.approx( + float(metric.compute()), abs=1e-6) + assert joined['labelled'] == entities + + +def _summary(values): + return {'count': len(values), 'min': float(np.min(values)), + 'max': float(np.max(values)), + 'percentiles': list(np.quantile(values, np.arange(1, 100) / 100))} + + +def test_psi_reads_percentile_curves(): + rng = np.random.default_rng(0) + baseline = _summary(rng.beta(2, 5, 50_000)) + + assert _psi(baseline, baseline) == 0 + assert _psi(_summary(rng.beta(2, 5, 50_000)), baseline) < .01 + assert _psi(_summary(rng.beta(5, 2, 50_000)), baseline) > .2 + tied = _summary(np.r_[np.zeros(40_000), rng.random(10_000)]) + assert _psi(tied, tied) == 0 + assert _psi({'quantiles': {'q50': .5}}, baseline) is None + + +def test_ks_and_wasserstein_from_percentile_curves(): + rng = np.random.default_rng(0) + a, b = rng.normal(0, 1, 100_000), rng.normal(0.5, 2, 100_000) + ks = max(abs(np.searchsorted(np.sort(a), x, 'right') / len(a) - + np.searchsorted(np.sort(b), x, 'right') / len(b)) + for x in np.linspace(-6, 6, 4001)) + levels = (np.arange(10_000) + .5) / 10_000 + w1 = np.mean(np.abs(np.quantile(a, levels) - np.quantile(b, levels))) + + assert _ks(_summary(a), _summary(b)) == pytest.approx(ks, abs=0.02) + assert _wasserstein(_summary(a), _summary(b)) == pytest.approx(w1, rel=.02) + assert _ks(_summary(a), _summary(a)) == _wasserstein( + _summary(a), _summary(a)) == 0 + # Point masses: 30% more ones moves the CDF jump at 0 by 0.3. + zeros = _summary(np.r_[np.zeros(500), np.ones(500)]) + ones = _summary(np.r_[np.zeros(200), np.ones(800)]) + assert _ks(zeros, ones) == pytest.approx(0.3, abs=0.02) + assert _wasserstein(zeros, ones) == pytest.approx(0.3, abs=0.02) + assert _ks(zeros, None) is None and _wasserstein({}, zeros) is None + + +def test_score_drift(snapshots, classifier, monkeypatch): + trainer, run = classifier + report = trainer.score_drift(run, previous=run, snapshot=snapshots.after, + training_snapshot=snapshots.before) + + assert report.score == 'True_PROB' and math.isfinite(report.psi) + assert math.isfinite(report.ks) and math.isfinite(report.wasserstein) + assert report.holdout_truncated is False and report.bundle_matches + assert (report.previous_psi, report.previous_ks, + report.previous_wasserstein) == (0, 0, 0) + assert report.input_changes['tables.orders.rows'] == (len( + snapshots.orders[snapshots.orders.TS < CUTOFF]), len(snapshots.orders)) + assert not any(key.endswith('_reverse') for key in report.input_changes) + assert report.reasons == {} + assert (report.prediction_id, report.previous_prediction_id) == (run.id, + run.id) + + shifted = _summary(np.full(500, 0.999)) + report = trainer.score_drift( + SimpleNamespace(id='r', bundle_id=None, manifest=SimpleNamespace( + stats={'predictions': ArtifactStats( + rows=500, columns={'True_PROB': shifted})}))) + assert report.psi > 0.2 and report.ks > 0.5 + assert report.bundle_matches is None + assert 'input_changes' in report.reasons + + report = trainer.score_drift(SimpleNamespace(manifest=None)) + assert (report.prediction_id, report.psi) == (None, None) + assert report.reasons + monkeypatch.setattr(TrainingResult, 'manifest', property(lambda self: None)) + report = trainer.score_drift(run) + assert (report.psi, report.ks) == (None, None) + assert 'holdout' in report.reasons['holdout'] + + +def test_monitoring_a_run_loaded_from_its_export(snapshots, classifier, + tmp_path): + from kumo.engine.trainer.output_config import PredictionOutputConfig + + trainer, run = classifier + connector = kumo.FileConnector(root_dir=str(tmp_path)) + run.export( + PredictionOutputConfig(connector, output_paths={'predictions': 'p/'})) + + loaded = PredictionResult.load_from_connector(connector, path='p/') + parts = loaded._exported.local / 'parts' + assert loaded.manifest.id == run.id and not parts.exists() + assert trainer.score_drift(loaded).psi is not None + assert not parts.exists() + + assert (loaded.id, loaded.bundle_id, loaded.snapshot_id, + loaded.num_rows) == (run.id, run.bundle_id, run.snapshot_id, + run.num_rows) + assert trainer.evaluate(loaded, snapshot=snapshots.after).metrics == \ + trainer.evaluate(run, snapshot=snapshots.after).metrics + assert trainer.score_drift(loaded).psi == trainer.score_drift(run).psi + + local = loaded._exported.local + assert local.is_dir() + del loaded + gc.collect() + assert not local.exists() + + (tmp_path / 'p' / '_SUCCESS').unlink() + with pytest.raises(FileNotFoundError, match='_SUCCESS'): + PredictionResult.load_from_connector(connector, path='p/') + + +def test_tracked_monitoring_logs_a_new_run(snapshots, classifier, tmp_path): + trainer, run = classifier + + class Tracker: + records: list = [] + + def start(self, record): + pass + + def on_trial(self, record, event): + pass + + def finish(self, record, error): + pass + + def log_monitoring(self, record): + self.records.append(record) + + cfg = kumo.MLOps(artifact_root=str(tmp_path / 'root'), tracking=Tracker()) + report = trainer.score_drift(run, mlops=cfg) + with kumo.mlops(cfg): + trainer.evaluate(run, snapshot=snapshots.after) + + drift, evaluation = Tracker.records + for record in (drift, evaluation): + assert record.tags['kumo.kind'] == 'monitoring' + assert record.tags['kumo.prediction_run'] == run.id + assert record.tags['kumo.bundle_id'] == run.bundle_id + assert drift.metrics['psi'] == report.psi + assert drift.metrics['wasserstein'] == report.wasserstein + assert 'auroc' in evaluation.metrics and 'rows' in evaluation.metrics diff --git a/packages/kumo-engine/test/test_predictive_query.py b/packages/kumo-engine/test/test_predictive_query.py index 2f77622a3..b306a43f7 100644 --- a/packages/kumo-engine/test/test_predictive_query.py +++ b/packages/kumo-engine/test/test_predictive_query.py @@ -309,7 +309,7 @@ def test_generated_weight_rejected_before_artifacts(connector_kind, tmp_path, ) monkeypatch.setattr( 'kumo.engine.pquery.predictive_query.generate_train_table', - lambda request: generated) + lambda request, **kwargs: generated) table = query.generate_training_table(snapshot) @@ -395,7 +395,7 @@ def test_forecast_generation_retains_effective_plan(graph_snapshot, tmp_path, ) monkeypatch.setattr( 'kumo.engine.pquery.predictive_query.generate_train_table', - lambda plan: generated) + lambda plan, **kwargs: generated) table = query.generate_training_table( snapshot, diff --git a/packages/kumo-engine/test/test_snowflake_tracking.py b/packages/kumo-engine/test/test_snowflake_tracking.py index ebf705d57..25a1ce057 100644 --- a/packages/kumo-engine/test/test_snowflake_tracking.py +++ b/packages/kumo-engine/test/test_snowflake_tracking.py @@ -154,7 +154,7 @@ def test_an_existing_run_is_never_resumed(fake, taken, caplog): assert set(user.runs) == {taken, new} assert user.runs[taken]['params'] == {} assert user.runs[new]['status'] == 'FINISHED' - assert 'recording this fit as' in caplog.text + assert 'recording this run as' in caplog.text def test_metric_and_param_limits(fake, monkeypatch):