Skip to content

Commit

Permalink
perf(nns): Improve neuron validation batching (#3105)
Browse files Browse the repository at this point in the history
# Why

Currently, the neuron data validation is reading all neuron sections for
a batch of neurons, and the batching is by a fixed batch size. Those can
be improved by: reading specific neuron sections for each task, and
batch by instruction limit.

# What

Implement the above-mentioned optimizations and update benchmarks.
  • Loading branch information
jasonz-dfinity authored Dec 12, 2024
1 parent e3cc69f commit 0e85368
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 31 deletions.
14 changes: 7 additions & 7 deletions rs/nns/governance/canbench/canbench_results.yml
Original file line number Diff line number Diff line change
Expand Up @@ -25,13 +25,13 @@ benches:
scopes: {}
cascading_vote_all_heap:
total:
instructions: 34451954
instructions: 34418252
heap_increase: 0
stable_memory_increase: 128
scopes: {}
cascading_vote_heap_neurons_stable_index:
total:
instructions: 56587595
instructions: 56553893
heap_increase: 0
stable_memory_increase: 128
scopes: {}
Expand Down Expand Up @@ -61,7 +61,7 @@ benches:
scopes: {}
draw_maturity_from_neurons_fund_heap:
total:
instructions: 7233119
instructions: 7242119
heap_increase: 0
stable_memory_increase: 0
scopes: {}
Expand All @@ -85,7 +85,7 @@ benches:
scopes: {}
list_neurons_heap:
total:
instructions: 3900362
instructions: 3900350
heap_increase: 9
stable_memory_increase: 0
scopes: {}
Expand All @@ -103,7 +103,7 @@ benches:
scopes: {}
list_neurons_stable:
total:
instructions: 97878953
instructions: 97878941
heap_increase: 5
stable_memory_increase: 0
scopes: {}
Expand All @@ -121,13 +121,13 @@ benches:
scopes: {}
neuron_data_validation_heap:
total:
instructions: 531768291
instructions: 349808709
heap_increase: 0
stable_memory_increase: 0
scopes: {}
neuron_data_validation_stable:
total:
instructions: 588628069
instructions: 312718406
heap_increase: 0
stable_memory_increase: 0
scopes: {}
Expand Down
55 changes: 31 additions & 24 deletions rs/nns/governance/src/neuron_data_validation.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ use crate::{
neuron::Neuron,
neuron_store::NeuronStore,
pb::v1::Topic,
storage::{with_stable_neuron_indexes, with_stable_neuron_store},
storage::{neurons::NeuronSections, with_stable_neuron_indexes, with_stable_neuron_store},
};

use candid::{CandidType, Deserialize};
Expand Down Expand Up @@ -327,7 +327,7 @@ impl ValidationTask for VecDeque<Box<dyn ValidationTask>> {
/// storage, and therefore past inconsistencies would be preserved indefinitely until found and
/// fixed.
trait CardinalityAndRangeValidator {
const HEAP_NEURON_RANGE_CHUNK_SIZE: usize;
const NEURON_SECTIONS: NeuronSections;

/// Validates the cardinalities of primary data and index to be equal.
fn validate_cardinalities(neuron_store: &NeuronStore) -> Option<ValidationIssue>;
Expand Down Expand Up @@ -396,10 +396,19 @@ impl<Validator: CardinalityAndRangeValidator + Send + Sync> ValidationTask
}

fn validate_next_chunk(&mut self, neuron_store: &NeuronStore) -> Vec<ValidationIssue> {
if let Some(next_neuron_id) = self.heap_next_neuron_id.take() {
// Set a limit on the number of instructions used by this function.
#[cfg(target_arch = "wasm32")]
let instruction_limit = ic_cdk::api::instruction_counter() + 100_000_000;
#[cfg(target_arch = "wasm32")]
let keep_going = || ic_cdk::api::instruction_counter() < instruction_limit;

#[cfg(not(target_arch = "wasm32"))]
let keep_going = || true;

let issues = if let Some(next_neuron_id) = self.heap_next_neuron_id.take() {
neuron_store
.heap_neurons_range(next_neuron_id..)
.take(Validator::HEAP_NEURON_RANGE_CHUNK_SIZE)
.take_while(|_| keep_going())
.flat_map(|neuron| {
self.heap_next_neuron_id = neuron.id().next();
Validator::validate_primary_neuron_has_corresponding_index_entries(neuron)
Expand All @@ -408,29 +417,26 @@ impl<Validator: CardinalityAndRangeValidator + Send + Sync> ValidationTask
} else if let Some(next_neuron_id) = self.stable_next_neuron_id.take() {
with_stable_neuron_store(|stable_neuron_store| {
stable_neuron_store
.range_neurons(next_neuron_id..)
.take(INACTIVE_NEURON_VALIDATION_CHUNK_SIZE)
.range_neurons_sections(next_neuron_id.., Validator::NEURON_SECTIONS)
.take_while(|_| keep_going())
.flat_map(|neuron| {
let neuron_id_to_validate = neuron.id();
self.stable_next_neuron_id = neuron_id_to_validate.next();
self.stable_next_neuron_id = neuron.id().next();
Validator::validate_primary_neuron_has_corresponding_index_entries(&neuron)
})
.collect()
})
} else {
println!("validate_next_chunk should not be called when is_done() is true");
vec![]
}
};
issues
}
}

struct SubaccountIndexValidator;

impl CardinalityAndRangeValidator for SubaccountIndexValidator {
// On average, checking whether an entry exists in a StableBTreeMap takes ~130K instructions,
// and there is exactly one entry for a neuron, so 400 neurons takes 52M instructions (aiming to
// keep it under 60M).
const HEAP_NEURON_RANGE_CHUNK_SIZE: usize = 400;
const NEURON_SECTIONS: NeuronSections = NeuronSections::NONE;

fn validate_cardinalities(neuron_store: &NeuronStore) -> Option<ValidationIssue> {
let cardinality_primary = neuron_store.len() as u64;
Expand Down Expand Up @@ -468,10 +474,10 @@ impl CardinalityAndRangeValidator for SubaccountIndexValidator {
struct PrincipalIndexValidator;

impl CardinalityAndRangeValidator for PrincipalIndexValidator {
// On average, checking whether an entry exists in a StableBTreeMap takes ~130K instructions,
// and there is 1 controler + 1.2 hot keys (=2.2) on average, so 200 neurons takes 57.2M
// instructions (aiming to keep it under 60M).
const HEAP_NEURON_RANGE_CHUNK_SIZE: usize = 200;
const NEURON_SECTIONS: NeuronSections = NeuronSections {
hot_keys: true,
..NeuronSections::NONE
};

fn validate_cardinalities(neuron_store: &NeuronStore) -> Option<ValidationIssue> {
let cardinality_primary_heap: u64 = neuron_store
Expand Down Expand Up @@ -526,10 +532,10 @@ impl CardinalityAndRangeValidator for PrincipalIndexValidator {
struct FollowingIndexValidator;

impl CardinalityAndRangeValidator for FollowingIndexValidator {
// On average, checking whether an entry exists in a StableBTreeMap takes ~130K instructions,
// and there are ~11.3 followee entries on average, so 40 neurons takes 58.76M instructions
// (aiming to keep it under 60M).
const HEAP_NEURON_RANGE_CHUNK_SIZE: usize = 40;
const NEURON_SECTIONS: NeuronSections = NeuronSections {
followees: true,
..NeuronSections::NONE
};

fn validate_cardinalities(neuron_store: &NeuronStore) -> Option<ValidationIssue> {
let cardinality_primary_heap: u64 = neuron_store
Expand Down Expand Up @@ -583,9 +589,10 @@ impl CardinalityAndRangeValidator for FollowingIndexValidator {
struct KnownNeuronIndexValidator;

impl CardinalityAndRangeValidator for KnownNeuronIndexValidator {
// As long as we have <460 known neurons, we will be able to keep the instructions <60M as each
// entry lookup takes ~130K instructions.
const HEAP_NEURON_RANGE_CHUNK_SIZE: usize = 300000;
const NEURON_SECTIONS: NeuronSections = NeuronSections {
known_neuron_data: true,
..NeuronSections::NONE
};
fn validate_cardinalities(neuron_store: &NeuronStore) -> Option<ValidationIssue> {
let cardinality_primary_heap = neuron_store
.heap_neurons_iter()
Expand Down

0 comments on commit 0e85368

Please sign in to comment.