diff --git a/iris-mpc-cpu/Cargo.toml b/iris-mpc-cpu/Cargo.toml index 6389a7c54..c0e862807 100644 --- a/iris-mpc-cpu/Cargo.toml +++ b/iris-mpc-cpu/Cargo.toml @@ -26,6 +26,7 @@ itertools.workspace = true num-traits.workspace = true prost = "0.13" rand.workspace = true +rand_distr = "0.4.3" rstest = "0.23.0" serde.workspace = true serde_json.workspace = true @@ -37,7 +38,6 @@ tracing.workspace = true tracing-subscriber.workspace = true tracing-test = "0.2.5" uuid.workspace = true -rand_distr = "0.4.3" [dev-dependencies] criterion = { version = "0.5.1", features = ["async_tokio"] } diff --git a/iris-mpc-cpu/bin/hnsw_algorithm_metrics.rs b/iris-mpc-cpu/bin/hnsw_algorithm_metrics.rs index dd13ff310..8e43a9c4f 100644 --- a/iris-mpc-cpu/bin/hnsw_algorithm_metrics.rs +++ b/iris-mpc-cpu/bin/hnsw_algorithm_metrics.rs @@ -6,25 +6,31 @@ use iris_mpc_cpu::{ hawkers::plaintext_store::PlaintextStore, hnsw::{ metrics::{ - EventCounter, HnswEventCounterLayer, VertexOpeningsLayer, COMPARE_DIST_EVENT, EVAL_DIST_EVENT, LAYER_SEARCH_EVENT, OPEN_NODE_EVENT - }, searcher::{HnswParams, HnswSearcher} - } + EventCounter, HnswEventCounterLayer, VertexOpeningsLayer, COMPARE_DIST_EVENT, + EVAL_DIST_EVENT, LAYER_SEARCH_EVENT, OPEN_NODE_EVENT, + }, + searcher::{HnswParams, HnswSearcher}, + }, }; use rand::SeedableRng; -use std::{collections::HashMap, error::Error, sync::{Arc, Mutex}}; +use std::{ + collections::HashMap, + error::Error, + sync::{Arc, Mutex}, +}; use tracing_subscriber::prelude::*; #[derive(Parser)] #[allow(non_snake_case)] struct Args { #[clap(default_value = "64")] - M: usize, + M: usize, #[clap(default_value = "128")] - ef_constr: usize, + ef_constr: usize, #[clap(default_value = "64")] - ef_search: usize, + ef_search: usize, #[clap(default_value = "10000")] - database_size: usize, + database_size: usize, layer_probability: Option, } @@ -94,18 +100,21 @@ fn print_stats(counters: &Arc, verbose: bool) { } } -fn configure_tracing() -> (Arc, Arc>>) { +fn configure_tracing() -> ( + Arc, + Arc>>, +) { let counters = Arc::new(EventCounter::default()); let counting_layer = HnswEventCounterLayer { counters: counters.clone(), }; - let counter_map: Arc>> - = Arc::new(Mutex::new(HashMap::new())); + let counter_map: Arc>> = + Arc::new(Mutex::new(HashMap::new())); let vertex_openings_layer = VertexOpeningsLayer { - counter_map: counter_map.clone() + counter_map: counter_map.clone(), }; tracing_subscriber::registry() diff --git a/iris-mpc-cpu/src/hawkers/aby3_store.rs b/iris-mpc-cpu/src/hawkers/aby3_store.rs index 20c82d800..b859a4f91 100644 --- a/iris-mpc-cpu/src/hawkers/aby3_store.rs +++ b/iris-mpc-cpu/src/hawkers/aby3_store.rs @@ -603,9 +603,7 @@ impl Aby3Store { #[cfg(test)] mod tests { use super::*; - use crate::{ - database_generators::generate_galois_iris_shares, hnsw::HnswSearcher, - }; + use crate::{database_generators::generate_galois_iris_shares, hnsw::HnswSearcher}; use aes_prng::AesRng; use hawk_pack::graph_store::GraphMem; use itertools::Itertools; diff --git a/iris-mpc-cpu/src/hawkers/plaintext_store.rs b/iris-mpc-cpu/src/hawkers/plaintext_store.rs index fa04b12f2..e03dc3e1b 100644 --- a/iris-mpc-cpu/src/hawkers/plaintext_store.rs +++ b/iris-mpc-cpu/src/hawkers/plaintext_store.rs @@ -1,6 +1,6 @@ use crate::hnsw::{ - HnswSearcher, metrics::{COMPARE_DIST_EVENT, EVAL_DIST_EVENT}, + HnswSearcher, }; use aes_prng::AesRng; use hawk_pack::{graph_store::GraphMem, VectorStore}; diff --git a/iris-mpc-cpu/src/hnsw/metrics.rs b/iris-mpc-cpu/src/hnsw/metrics.rs index cd829cefd..2c8cb4f86 100644 --- a/iris-mpc-cpu/src/hnsw/metrics.rs +++ b/iris-mpc-cpu/src/hnsw/metrics.rs @@ -1,15 +1,20 @@ -use std::{collections::HashMap, sync::{ - atomic::{AtomicUsize, Ordering}, - Arc, Mutex, -}}; -use tracing::{field::{Field, Visit}, Event, Id, Subscriber}; +use std::{ + collections::HashMap, + fmt::Debug, + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, Mutex, + }, +}; +use tracing::{ + field::{Field, Visit}, + span::Attributes, + Event, Id, Subscriber, +}; use tracing_subscriber::{ layer::{Context, Layer}, + registry::LookupSpan, }; -use std::fmt::Debug; -use tracing_subscriber::registry::LookupSpan; - -use tracing::span::Attributes; pub const LAYER_SEARCH_EVENT: u64 = 0; pub const OPEN_NODE_EVENT: u64 = 1; @@ -28,7 +33,8 @@ pub struct HnswEventCounterLayer { } impl Layer for HnswEventCounterLayer - where S: Subscriber + for <'a> LookupSpan<'a> +where + S: Subscriber + for<'a> LookupSpan<'a>, { fn on_event(&self, event: &Event<'_>, _ctx: Context<'_, S>) { let mut visitor = EventVisitor::default(); @@ -70,7 +76,6 @@ impl Visit for EventVisitor { fn record_debug(&mut self, _field: &Field, _value: &dyn Debug) {} } - /// Tracing library Layer for counting detailed HNSW layer search operations pub struct VertexOpeningsLayer { // Measure number of vertex openings for different lc and ef values @@ -78,15 +83,16 @@ pub struct VertexOpeningsLayer { } impl Layer for VertexOpeningsLayer - where S: Subscriber + for <'a> LookupSpan<'a> +where + S: Subscriber + for<'a> LookupSpan<'a>, { - // fn register_callsite(&self, _metadata: &'static Metadata<'static>) -> Interest { - // Interest::sometimes() + // fn register_callsite(&self, _metadata: &'static Metadata<'static>) -> + // Interest { Interest::sometimes() // } // fn enabled(&self, metadata: &Metadata<'_>, _ctx: Context<'_, S>) -> bool { - // let is_search_layer_span = metadata.is_span() && metadata.name() == "search_layer"; - // is_search_layer_span || metadata.is_event() + // let is_search_layer_span = metadata.is_span() && metadata.name() == + // "search_layer"; is_search_layer_span || metadata.is_event() // } fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) { @@ -104,13 +110,20 @@ impl Layer for VertexOpeningsLayer // open node event must have parent span representing open node function let current_span = ctx.current_span(); let span_id = current_span.id().unwrap(); - if let Some(LayerSearchFields { lc: Some(lc), ef: Some(ef) }) - = ctx.span(span_id).unwrap().extensions().get::() + if let Some(LayerSearchFields { + lc: Some(lc), + ef: Some(ef), + }) = ctx + .span(span_id) + .unwrap() + .extensions() + .get::() { let mut counter_map = self.counter_map.lock().unwrap(); let increment_amount = visitor.amount.unwrap_or(1); - *counter_map.entry((*lc as usize, *ef as usize)).or_insert(0usize) - += increment_amount as usize; + *counter_map + .entry((*lc as usize, *ef as usize)) + .or_insert(0usize) += increment_amount as usize; } else { panic!("Open node event is missing associated span fields"); } diff --git a/iris-mpc-cpu/src/hnsw/mod.rs b/iris-mpc-cpu/src/hnsw/mod.rs index 9a6c320ef..edcae0c1c 100644 --- a/iris-mpc-cpu/src/hnsw/mod.rs +++ b/iris-mpc-cpu/src/hnsw/mod.rs @@ -1,4 +1,4 @@ -pub mod searcher; pub mod metrics; +pub mod searcher; pub use searcher::HnswSearcher; diff --git a/iris-mpc-cpu/src/hnsw/searcher.rs b/iris-mpc-cpu/src/hnsw/searcher.rs index 25626ecdc..2101604b9 100644 --- a/iris-mpc-cpu/src/hnsw/searcher.rs +++ b/iris-mpc-cpu/src/hnsw/searcher.rs @@ -4,7 +4,7 @@ //* //* https://github.com/Inversed-Tech/hawk-pack/ -use tracing::{info, instrument}; +use super::metrics; pub use hawk_pack::data_structures::queue::{ FurthestQueue, FurthestQueueV, NearestQueue, NearestQueueV, }; @@ -13,8 +13,7 @@ use rand::RngCore; use rand_distr::{Distribution, Geometric}; use serde::{Deserialize, Serialize}; use std::collections::HashSet; - -use super::metrics; +use tracing::{info, instrument}; // specify construction and search parameters by layer up to this value minus 1 // any higher layers will use the last set of parameters @@ -70,8 +69,8 @@ impl HnswParams { ef_construction: usize, ef_search: usize, M: usize, - layer_probability: f64) -> Self - { + layer_probability: f64, + ) -> Self { let M_arr = [M; N_PARAM_LAYERS]; let mut M_max_arr = [M; N_PARAM_LAYERS]; M_max_arr[0] = 2 * M; @@ -470,7 +469,6 @@ impl HnswSearcher { } } - #[cfg(test)] mod tests { use super::*; diff --git a/iris-mpc-cpu/src/network/grpc.rs b/iris-mpc-cpu/src/network/grpc.rs index 74db712b4..1e1a0542a 100644 --- a/iris-mpc-cpu/src/network/grpc.rs +++ b/iris-mpc-cpu/src/network/grpc.rs @@ -340,8 +340,8 @@ mod tests { use super::*; use crate::{ execution::{local::generate_local_identities, player::Role}, - hnsw::HnswSearcher, hawkers::aby3_store::Aby3Store, + hnsw::HnswSearcher, }; use aes_prng::AesRng; use rand::SeedableRng;