diff --git a/pallets/subtensor/src/macros/hooks.rs b/pallets/subtensor/src/macros/hooks.rs index 1b7d5fd77..84d4f99e9 100644 --- a/pallets/subtensor/src/macros/hooks.rs +++ b/pallets/subtensor/src/macros/hooks.rs @@ -162,7 +162,9 @@ mod hooks { // Migrate pending emissions .saturating_add(migrations::migrate_pending_emissions::migrate_pending_emissions::()) // Reset unactive subnets - .saturating_add(migrations::migrate_reset_unactive_sn::migrate_reset_unactive_sn::()); + .saturating_add(migrations::migrate_reset_unactive_sn::migrate_reset_unactive_sn::()) + // Remove unknown neuron axon, certificate prom + .saturating_add(migrations::migrate_remove_unknown_neuron_axon_cert_prom::migrate_remove_unknown_neuron_axon_cert_prom::()); weight } diff --git a/pallets/subtensor/src/migrations/migrate_remove_unknown_neuron_axon_cert_prom.rs b/pallets/subtensor/src/migrations/migrate_remove_unknown_neuron_axon_cert_prom.rs new file mode 100644 index 000000000..4553333e6 --- /dev/null +++ b/pallets/subtensor/src/migrations/migrate_remove_unknown_neuron_axon_cert_prom.rs @@ -0,0 +1,83 @@ +use super::*; +use crate::HasMigrationRun; +use frame_support::{traits::Get, weights::Weight}; +use scale_info::prelude::string::String; +use sp_std::collections::btree_set::BTreeSet; + +pub fn migrate_remove_unknown_neuron_axon_cert_prom() -> Weight { + let migration_name = b"migrate_remove_neuron_axon_cert_prom".to_vec(); + let mut weight: Weight = T::DbWeight::get().reads(1); + + // Skip if already executed + if HasMigrationRun::::get(&migration_name) { + log::info!( + target: "runtime", + "Migration '{}' already run - skipping.", + String::from_utf8_lossy(&migration_name) + ); + return weight; + } + log::info!( + "Running migration '{}'", + String::from_utf8_lossy(&migration_name) + ); + + for network in NetworksAdded::::iter_keys() { + weight.saturating_accrue(T::DbWeight::get().reads(1)); + + let hotkeys = BTreeSet::from_iter(Uids::::iter_key_prefix(network)); + weight.saturating_accrue(T::DbWeight::get().reads(hotkeys.len() as u64)); + + // Axons + let axons = Axons::::iter_key_prefix(network).collect::>(); + weight.saturating_accrue(T::DbWeight::get().reads(axons.len() as u64)); + let mut cleaned_axons: u32 = 0; + for axon_hotkey in axons { + if !hotkeys.contains(&axon_hotkey) { + Axons::::remove(network, &axon_hotkey); + cleaned_axons = cleaned_axons.saturating_add(1); + } + } + weight.saturating_accrue(T::DbWeight::get().writes(cleaned_axons as u64)); + + // Neuron Certificates + let certificates = NeuronCertificates::::iter_key_prefix(network).collect::>(); + weight.saturating_accrue(T::DbWeight::get().reads(certificates.len() as u64)); + let mut cleaned_certificates: u32 = 0; + for certificate_hotkey in certificates { + if !hotkeys.contains(&certificate_hotkey) { + NeuronCertificates::::remove(network, &certificate_hotkey); + cleaned_certificates = cleaned_certificates.saturating_add(1); + } + } + weight.saturating_accrue(T::DbWeight::get().writes(cleaned_certificates as u64)); + + // Prometheus + let prometheus = Prometheus::::iter_key_prefix(network).collect::>(); + weight.saturating_accrue(T::DbWeight::get().reads(prometheus.len() as u64)); + let mut cleaned_prometheus: u32 = 0; + for prometheus_hotkey in prometheus { + if !hotkeys.contains(&prometheus_hotkey) { + Prometheus::::remove(network, &prometheus_hotkey); + cleaned_prometheus = cleaned_prometheus.saturating_add(1); + } + } + weight.saturating_accrue(T::DbWeight::get().writes(cleaned_prometheus as u64)); + + log::info!( + "Cleaned {cleaned_axons} axons, {cleaned_certificates} neuron certificates, {cleaned_prometheus} prometheus for network {network}" + ); + } + + HasMigrationRun::::insert(&migration_name, true); + weight = weight.saturating_add(T::DbWeight::get().writes(1)); + + log::info!( + "Migration '{:?}' completed successfully.", + String::from_utf8_lossy(&migration_name) + ); + + log::info!("{weight:#?} weight"); + + weight +} diff --git a/pallets/subtensor/src/migrations/mod.rs b/pallets/subtensor/src/migrations/mod.rs index 41c1333a8..2715be278 100644 --- a/pallets/subtensor/src/migrations/mod.rs +++ b/pallets/subtensor/src/migrations/mod.rs @@ -37,6 +37,7 @@ pub mod migrate_remove_network_modality; pub mod migrate_remove_stake_map; pub mod migrate_remove_tao_dividends; pub mod migrate_remove_total_hotkey_coldkey_stakes_this_interval; +pub mod migrate_remove_unknown_neuron_axon_cert_prom; pub mod migrate_remove_unused_maps_and_values; pub mod migrate_remove_zero_total_hotkey_alpha; pub mod migrate_reset_bonds_moving_average; diff --git a/pallets/subtensor/src/subnets/uids.rs b/pallets/subtensor/src/subnets/uids.rs index 669f74bcc..fda21bc33 100644 --- a/pallets/subtensor/src/subnets/uids.rs +++ b/pallets/subtensor/src/subnets/uids.rs @@ -88,14 +88,13 @@ impl Pallet { BlockAtRegistration::::insert(netuid, uid_to_replace, block_number); // Fill block at registration. IsNetworkMember::::insert(new_hotkey.clone(), netuid, true); // Fill network is member. - // 4. Clear neuron certificates - NeuronCertificates::::remove(netuid, old_hotkey.clone()); + // 4. Clear neuron axons, certificates and prometheus info + Axons::::remove(netuid, &old_hotkey); + NeuronCertificates::::remove(netuid, &old_hotkey); + Prometheus::::remove(netuid, &old_hotkey); // 5. Reset new neuron's values. Self::clear_neuron(netuid, uid_to_replace); - - // 5a. reset axon info for the new uid. - Axons::::remove(netuid, old_hotkey); } /// Appends the uid to the network. diff --git a/pallets/subtensor/src/tests/migration.rs b/pallets/subtensor/src/tests/migration.rs index b694459ea..7e266487e 100644 --- a/pallets/subtensor/src/tests/migration.rs +++ b/pallets/subtensor/src/tests/migration.rs @@ -2724,3 +2724,69 @@ fn test_migrate_reset_unactive_sn_idempotence() { assert_eq!(TotalIssuance::::get(), total_issuance_before); }); } + +#[test] +fn test_migrate_remove_unknown_neuron_axon_cert_prom() { + use crate::migrations::migrate_remove_unknown_neuron_axon_cert_prom::*; + const MIGRATION_NAME: &[u8] = b"migrate_remove_neuron_axon_cert_prom"; + + new_test_ext(1).execute_with(|| { + setup_for(NetUid::from(2), 64, 1231); + setup_for(NetUid::from(42), 256, 15151); + setup_for(NetUid::from(99), 1024, 32323); + assert!(!HasMigrationRun::::get(MIGRATION_NAME)); + + let w = migrate_remove_unknown_neuron_axon_cert_prom::(); + assert!(!w.is_zero(), "Weight must be non-zero"); + + assert!(HasMigrationRun::::get(MIGRATION_NAME)); + assert_for(NetUid::from(2), 64, 1231); + assert_for(NetUid::from(42), 256, 15151); + assert_for(NetUid::from(99), 1024, 32323); + }); + + fn setup_for(netuid: NetUid, uids: u32, items: u32) { + NetworksAdded::::insert(netuid, true); + + for i in 1u32..=uids { + let hk = U256::from(netuid.inner() as u32 * 1000 + i); + Uids::::insert(netuid, hk, i as u16); + } + + for i in 1u32..=items { + let hk = U256::from(netuid.inner() as u32 * 1000 + i); + Axons::::insert(netuid, hk, AxonInfo::default()); + NeuronCertificates::::insert(netuid, hk, NeuronCertificate::default()); + Prometheus::::insert(netuid, hk, PrometheusInfo::default()); + } + } + + fn assert_for(netuid: NetUid, uids: u32, items: u32) { + assert_eq!( + Axons::::iter_key_prefix(netuid).count(), + uids as usize + ); + assert_eq!( + NeuronCertificates::::iter_key_prefix(netuid).count(), + uids as usize + ); + assert_eq!( + Prometheus::::iter_key_prefix(netuid).count(), + uids as usize + ); + + for i in 1u32..=uids { + let hk = U256::from(netuid.inner() as u32 * 1000 + i); + assert!(Axons::::contains_key(netuid, hk)); + assert!(NeuronCertificates::::contains_key(netuid, hk)); + assert!(Prometheus::::contains_key(netuid, hk)); + } + + for i in uids + 1u32..=items { + let hk = U256::from(netuid.inner() as u32 * 1000 + i); + assert!(!Axons::::contains_key(netuid, hk)); + assert!(!NeuronCertificates::::contains_key(netuid, hk)); + assert!(!Prometheus::::contains_key(netuid, hk)); + } + } +} diff --git a/pallets/subtensor/src/tests/uids.rs b/pallets/subtensor/src/tests/uids.rs index 8fee5f750..16a480bec 100644 --- a/pallets/subtensor/src/tests/uids.rs +++ b/pallets/subtensor/src/tests/uids.rs @@ -32,7 +32,6 @@ fn test_replace_neuron() { let new_hotkey_account_id = U256::from(2); let _new_colkey_account_id = U256::from(12345); - let certificate = NeuronCertificate::try_from(vec![1, 2, 3]).unwrap(); let evm_address = H160::from_slice(&[1_u8; 20]); //add network add_network(netuid, tempo, 0); @@ -71,28 +70,11 @@ fn test_replace_neuron() { }); Bonds::::insert(NetUidStorageIndex::from(netuid), neuron_uid, vec![(0, 1)]); - // serve axon mock address - let ip: u128 = 1676056785; - let port: u16 = 9999; - let ip_type: u8 = 4; - assert!( - SubtensorModule::serve_axon( - <::RuntimeOrigin>::signed(hotkey_account_id), - netuid, - 0, - ip, - port, - ip_type, - 0, - 0, - 0 - ) - .is_ok() - ); - - // Set a neuron certificate for it - NeuronCertificates::::insert(netuid, hotkey_account_id, certificate); + Axons::::insert(netuid, hotkey_account_id, AxonInfoOf::default()); + NeuronCertificates::::insert(netuid, hotkey_account_id, NeuronCertificateOf::default()); + Prometheus::::insert(netuid, hotkey_account_id, PrometheusInfoOf::default()); AssociatedEvmAddress::::insert(netuid, neuron_uid, (evm_address, 1)); + // Replace the neuron. SubtensorModule::replace_neuron(netuid, neuron_uid, &new_hotkey_account_id, block_number); @@ -139,10 +121,9 @@ fn test_replace_neuron() { ); // Check axon info is reset. - let axon_info = SubtensorModule::get_axon_info(netuid, &curr_hotkey.unwrap()); - assert_eq!(axon_info.ip, 0); - assert_eq!(axon_info.port, 0); - assert_eq!(axon_info.ip_type, 0); + assert!(!Axons::::contains_key(netuid, curr_hotkey.unwrap())); + assert!(!NeuronCertificates::::contains_key(netuid, curr_hotkey.unwrap())); + assert!(!Prometheus::::contains_key(netuid, curr_hotkey.unwrap())); // Check bonds are cleared. assert_eq!(