Merge pull request #2212 from opentensor/clear-axon-prom

Clear unknown neuron axon/prom/cert

Author: sam0x17

pallets/subtensor/src/macros/hooks.rs

@@ -162,7 +162,9 @@ mod hooks {
                 // Migrate pending emissions
                 .saturating_add(migrations::migrate_pending_emissions::migrate_pending_emissions::<T>())
                 // Reset unactive subnets
-                .saturating_add(migrations::migrate_reset_unactive_sn::migrate_reset_unactive_sn::<T>());
+                .saturating_add(migrations::migrate_reset_unactive_sn::migrate_reset_unactive_sn::<T>())
+                // Remove unknown neuron axon, certificate prom
+                .saturating_add(migrations::migrate_remove_unknown_neuron_axon_cert_prom::migrate_remove_unknown_neuron_axon_cert_prom::<T>());
             weight
         }
 

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<T: Config>() -> 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::<T>::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::<T>::iter_keys() {
+        weight.saturating_accrue(T::DbWeight::get().reads(1));
+
+        let hotkeys = BTreeSet::from_iter(Uids::<T>::iter_key_prefix(network));
+        weight.saturating_accrue(T::DbWeight::get().reads(hotkeys.len() as u64));
+
+        // Axons
+        let axons = Axons::<T>::iter_key_prefix(network).collect::<Vec<_>>();
+        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::<T>::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::<T>::iter_key_prefix(network).collect::<Vec<_>>();
+        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::<T>::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::<T>::iter_key_prefix(network).collect::<Vec<_>>();
+        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::<T>::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::<T>::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
+}

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;

pallets/subtensor/src/subnets/uids.rs

@@ -88,14 +88,13 @@ impl<T: Config> Pallet<T> {
         BlockAtRegistration::<T>::insert(netuid, uid_to_replace, block_number); // Fill block at registration.
         IsNetworkMember::<T>::insert(new_hotkey.clone(), netuid, true); // Fill network is member.
 
-        // 4. Clear neuron certificates
-        NeuronCertificates::<T>::remove(netuid, old_hotkey.clone());
+        // 4. Clear neuron axons, certificates and prometheus info
+        Axons::<T>::remove(netuid, &old_hotkey);
+        NeuronCertificates::<T>::remove(netuid, &old_hotkey);
+        Prometheus::<T>::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::<T>::remove(netuid, old_hotkey);
     }
 
     /// Appends the uid to the network.

pallets/subtensor/src/tests/migration.rs

@@ -2724,3 +2724,69 @@ fn test_migrate_reset_unactive_sn_idempotence() {
         assert_eq!(TotalIssuance::<Test>::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::<Test>::get(MIGRATION_NAME));
+
+        let w = migrate_remove_unknown_neuron_axon_cert_prom::<Test>();
+        assert!(!w.is_zero(), "Weight must be non-zero");
+
+        assert!(HasMigrationRun::<Test>::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::<Test>::insert(netuid, true);
+
+        for i in 1u32..=uids {
+            let hk = U256::from(netuid.inner() as u32 * 1000 + i);
+            Uids::<Test>::insert(netuid, hk, i as u16);
+        }
+
+        for i in 1u32..=items {
+            let hk = U256::from(netuid.inner() as u32 * 1000 + i);
+            Axons::<Test>::insert(netuid, hk, AxonInfo::default());
+            NeuronCertificates::<Test>::insert(netuid, hk, NeuronCertificate::default());
+            Prometheus::<Test>::insert(netuid, hk, PrometheusInfo::default());
+        }
+    }
+
+    fn assert_for(netuid: NetUid, uids: u32, items: u32) {
+        assert_eq!(
+            Axons::<Test>::iter_key_prefix(netuid).count(),
+            uids as usize
+        );
+        assert_eq!(
+            NeuronCertificates::<Test>::iter_key_prefix(netuid).count(),
+            uids as usize
+        );
+        assert_eq!(
+            Prometheus::<Test>::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::<Test>::contains_key(netuid, hk));
+            assert!(NeuronCertificates::<Test>::contains_key(netuid, hk));
+            assert!(Prometheus::<Test>::contains_key(netuid, hk));
+        }
+
+        for i in uids + 1u32..=items {
+            let hk = U256::from(netuid.inner() as u32 * 1000 + i);
+            assert!(!Axons::<Test>::contains_key(netuid, hk));
+            assert!(!NeuronCertificates::<Test>::contains_key(netuid, hk));
+            assert!(!Prometheus::<Test>::contains_key(netuid, hk));
+        }
+    }
+}

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,15 @@ fn test_replace_neuron() {
         });
         Bonds::<Test>::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(
-                <<Test as Config>::RuntimeOrigin>::signed(hotkey_account_id),
-                netuid,
-                0,
-                ip,
-                port,
-                ip_type,
-                0,
-                0,
-                0
-            )
-            .is_ok()
+        Axons::<Test>::insert(netuid, hotkey_account_id, AxonInfoOf::default());
+        NeuronCertificates::<Test>::insert(
+            netuid,
+            hotkey_account_id,
+            NeuronCertificateOf::default(),
         );
-
-        // Set a neuron certificate for it
-        NeuronCertificates::<Test>::insert(netuid, hotkey_account_id, certificate);
+        Prometheus::<Test>::insert(netuid, hotkey_account_id, PrometheusInfoOf::default());
         AssociatedEvmAddress::<Test>::insert(netuid, neuron_uid, (evm_address, 1));
+
         // Replace the neuron.
         SubtensorModule::replace_neuron(netuid, neuron_uid, &new_hotkey_account_id, block_number);
 
@@ -139,10 +125,15 @@ 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::<Test>::contains_key(netuid, curr_hotkey.unwrap()));
+        assert!(!NeuronCertificates::<Test>::contains_key(
+            netuid,
+            curr_hotkey.unwrap()
+        ));
+        assert!(!Prometheus::<Test>::contains_key(
+            netuid,
+            curr_hotkey.unwrap()
+        ));
 
         // Check bonds are cleared.
         assert_eq!(