diff --git a/cli/Cargo.lock b/cli/Cargo.lock index c34372b..a0a5f54 100644 --- a/cli/Cargo.lock +++ b/cli/Cargo.lock @@ -3894,6 +3894,16 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "papaya" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "997ee03cd38c01469a7046643714f0ad28880bcb9e6679ff0666e24817ca19b7" +dependencies = [ + "equivalent", + "seize", +] + [[package]] name = "parking" version = "2.2.1" @@ -4291,11 +4301,14 @@ dependencies = [ "axum", "clap", "comfy-table", + "crossbeam-queue", "datafusion", "dirs", "fastembed", "futures", "genai", + "itertools 0.14.0", + "papaya", "parquet", "predicates", "rand 0.8.6", @@ -5159,6 +5172,16 @@ dependencies = [ "libc", ] +[[package]] +name = "seize" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b55fb86dfd3a2f5f76ea78310a88f96c4ea21a3031f8d212443d56123fd0521" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + [[package]] name = "semver" version = "1.0.28" diff --git a/cli/Cargo.toml b/cli/Cargo.toml index efd9c27..a2ad207 100644 --- a/cli/Cargo.toml +++ b/cli/Cargo.toml @@ -10,11 +10,14 @@ async-trait = "0.1" axum = "0.8.9" clap = { version = "4.6.1", features = ["derive"] } comfy-table = "7.2.1" +crossbeam-queue = "0.3.12" datafusion = "46.0.0" dirs = "6" fastembed = "5" futures = "0.3" genai = "0.6.5" +itertools = "0.14.0" +papaya = "0.2.4" parquet = { version = "54", features = ["arrow"] } rand = "0.8" reqwest = { version = "0.13.3", features = ["json", "rustls"] } diff --git a/cli/src/metrics_store.rs b/cli/src/metrics_store.rs index b5a187e..2f52af1 100644 --- a/cli/src/metrics_store.rs +++ b/cli/src/metrics_store.rs @@ -1,4 +1,3 @@ -use std::collections::HashMap; use std::path::PathBuf; use std::sync::Arc; @@ -6,9 +5,11 @@ use anyhow::{Context, Result}; use arrow::array::{Array, Float64Array, Int64Array, StringArray, StringViewArray}; use arrow::datatypes::{DataType, Field, Schema}; use arrow::record_batch::RecordBatch; +use crossbeam_queue::SegQueue; use datafusion::prelude::*; +use itertools::Itertools; +use papaya::HashMap; use parquet::arrow::ArrowWriter; -use tokio::sync::Mutex; use crate::db::MetricPointSummary; use crate::time::{format_utc, now_utc}; @@ -31,10 +32,8 @@ struct BufferedMetric { /// SQL over the per-run Parquet files. #[derive(Debug, Clone)] pub struct MetricsStore { - // TODO: use a concurrent queue rather than sequential hashmap, - // - // https://docs.rs/crossbeam-queue/latest/crossbeam_queue/struct.SegQueue.html - buffer: Arc>>>, + /// Mapping from `run_id` to queue of buffered metrics. + buffer: Arc>>, dir: PathBuf, } @@ -51,7 +50,7 @@ impl MetricsStore { std::fs::create_dir_all(&dir) .with_context(|| format!("failed to create metrics dir {}", dir.display()))?; Ok(Self { - buffer: Arc::new(Mutex::new(HashMap::new())), + buffer: Arc::new(HashMap::new()), dir, }) } @@ -65,8 +64,10 @@ impl MetricsStore { metric_value: f64, unit: Option<&str>, ) { - let mut buffer = self.buffer.lock().await; - buffer.entry(run_id).or_default().push(BufferedMetric { + let buffer = &self.buffer; + let buffer_guard = buffer.guard(); + let run_queue = buffer.get_or_insert_with(run_id, || SegQueue::new(), &buffer_guard); + run_queue.push(BufferedMetric { run_id, step_id, metric_name: metric_name.to_owned(), @@ -84,11 +85,17 @@ impl MetricsStore { /// /// Returns an error if the flush failed. pub async fn flush(&self, run_id: i64) -> Result<()> { - let metrics = { - let mut buffer = self.buffer.lock().await; - buffer.remove(&run_id).unwrap_or_default() + let buffer = &self.buffer; + let buffer_guard = buffer.guard(); + let Some(metrics_q) = buffer.remove(&run_id, &buffer_guard) else { + return Ok(()); }; + let mut metrics = Vec::new(); + while let Some(metric) = metrics_q.pop() { + metrics.push(metric); + } + if metrics.is_empty() { return Ok(()); } @@ -108,29 +115,43 @@ impl MetricsStore { // so that timestamps are stored natively in Parquet rather than as strings. ])); - let run_ids = Int64Array::from_iter_values(metrics.iter().map(|m| m.run_id)); - let step_ids = Int64Array::from(metrics.iter().map(|m| m.step_id).collect::>()); - let metric_names = StringArray::from( - metrics - .iter() - .map(|m| m.metric_name.as_str()) - .collect::>(), - ); - let metric_values = Float64Array::from_iter_values(metrics.iter().map(|m| m.metric_value)); - let units = StringArray::from( - metrics - .iter() - .map(|m| m.unit.as_deref()) - .collect::>(), - ); - let created_at_strings: Vec = - metrics.iter().map(|m| format_utc(m.created_at)).collect(); - let created_ats = StringArray::from( - created_at_strings - .iter() - .map(std::string::String::as_str) - .collect::>(), - ); + let (run_ids, step_ids, metric_names, metric_values, units, created_ats) = { + let metrics_iter = metrics.into_iter(); + + let vals_vec = metrics_iter.map(|m| { + ( + m.run_id, + m.step_id, + m.metric_name, + m.metric_value, + m.unit, + format_utc(m.created_at), + ) + }); + let unzipped: ( + Vec, + Vec>, + Vec, + Vec, + Vec>, + Vec, + ) = vals_vec.into_iter().multiunzip(); + + let run_ids = Int64Array::from(unzipped.0); + let step_ids = Int64Array::from(unzipped.1); + let metric_names = StringArray::from(unzipped.2); + let metric_values = Float64Array::from(unzipped.3); + let units = StringArray::from(unzipped.4); + let created_ats = StringArray::from(unzipped.5); + ( + run_ids, + step_ids, + metric_names, + metric_values, + units, + created_ats, + ) + }; let batch = RecordBatch::try_new( schema.clone(),