From 3aa05adc627757ff2d621cb18065f5151c875b93 Mon Sep 17 00:00:00 2001 From: Roman Date: Fri, 18 Apr 2025 11:27:42 -0700 Subject: [PATCH 1/3] add `force_register_neuron` method --- bittensor/utils/mock/subtensor_mock.py | 130 +++++++++++++++++++++++++ 1 file changed, 130 insertions(+) diff --git a/bittensor/utils/mock/subtensor_mock.py b/bittensor/utils/mock/subtensor_mock.py index e2aeec8758..9773d525a4 100644 --- a/bittensor/utils/mock/subtensor_mock.py +++ b/bittensor/utils/mock/subtensor_mock.py @@ -1,6 +1,7 @@ from collections.abc import Mapping from dataclasses import dataclass from hashlib import sha256 +from random import randint from types import SimpleNamespace from typing import Any, Optional, Union, TypedDict from unittest.mock import MagicMock, patch @@ -358,6 +359,104 @@ def set_difficulty(self, netuid: int, difficulty: int) -> None: subtensor_state["Difficulty"][netuid][self.block_number] = difficulty + def _register_neuron(self, netuid: int, hotkey: str, coldkey: str) -> int: + subtensor_state = self.chain_state["SubtensorModule"] + if netuid not in subtensor_state["NetworksAdded"]: + raise Exception("Subnet does not exist") + + subnetwork_n = self._get_most_recent_storage( + subtensor_state["SubnetworkN"][netuid] + ) + + if subnetwork_n > 0 and any( + self._get_most_recent_storage(subtensor_state["Keys"][netuid][uid]) + == hotkey + for uid in range(subnetwork_n) + ): + # already_registered + raise Exception("Hotkey already registered") + else: + # Not found + if subnetwork_n >= self._get_most_recent_storage( + subtensor_state["MaxAllowedUids"][netuid] + ): + # Subnet full, replace neuron randomly + uid = randint(0, subnetwork_n - 1) + else: + # Subnet not full, add new neuron + # Append as next uid and increment subnetwork_n + uid = subnetwork_n + subtensor_state["SubnetworkN"][netuid][self.block_number] = ( + subnetwork_n + 1 + ) + + subtensor_state["Stake"][hotkey] = {} + subtensor_state["Stake"][hotkey][coldkey] = {} + subtensor_state["Stake"][hotkey][coldkey][self.block_number] = 0 + + subtensor_state["Uids"][netuid][hotkey] = {} + subtensor_state["Uids"][netuid][hotkey][self.block_number] = uid + + subtensor_state["Keys"][netuid][uid] = {} + subtensor_state["Keys"][netuid][uid][self.block_number] = hotkey + + subtensor_state["Owner"][hotkey] = {} + subtensor_state["Owner"][hotkey][self.block_number] = coldkey + + subtensor_state["Active"][netuid][uid] = {} + subtensor_state["Active"][netuid][uid][self.block_number] = True + + subtensor_state["LastUpdate"][netuid][uid] = {} + subtensor_state["LastUpdate"][netuid][uid][self.block_number] = ( + self.block_number + ) + + subtensor_state["Rank"][netuid][uid] = {} + subtensor_state["Rank"][netuid][uid][self.block_number] = 0.0 + + subtensor_state["Emission"][netuid][uid] = {} + subtensor_state["Emission"][netuid][uid][self.block_number] = 0.0 + + subtensor_state["Incentive"][netuid][uid] = {} + subtensor_state["Incentive"][netuid][uid][self.block_number] = 0.0 + + subtensor_state["Consensus"][netuid][uid] = {} + subtensor_state["Consensus"][netuid][uid][self.block_number] = 0.0 + + subtensor_state["Trust"][netuid][uid] = {} + subtensor_state["Trust"][netuid][uid][self.block_number] = 0.0 + + subtensor_state["ValidatorTrust"][netuid][uid] = {} + subtensor_state["ValidatorTrust"][netuid][uid][self.block_number] = 0.0 + + subtensor_state["Dividends"][netuid][uid] = {} + subtensor_state["Dividends"][netuid][uid][self.block_number] = 0.0 + + subtensor_state["PruningScores"][netuid][uid] = {} + subtensor_state["PruningScores"][netuid][uid][self.block_number] = 0.0 + + subtensor_state["ValidatorPermit"][netuid][uid] = {} + subtensor_state["ValidatorPermit"][netuid][uid][self.block_number] = False + + subtensor_state["Weights"][netuid][uid] = {} + subtensor_state["Weights"][netuid][uid][self.block_number] = [] + + subtensor_state["Bonds"][netuid][uid] = {} + subtensor_state["Bonds"][netuid][uid][self.block_number] = [] + + subtensor_state["Axons"][netuid][hotkey] = {} + subtensor_state["Axons"][netuid][hotkey][self.block_number] = {} + + subtensor_state["Prometheus"][netuid][hotkey] = {} + subtensor_state["Prometheus"][netuid][hotkey][self.block_number] = {} + + if hotkey not in subtensor_state["IsNetworkMember"]: + subtensor_state["IsNetworkMember"][hotkey] = {} + subtensor_state["IsNetworkMember"][hotkey][netuid] = {} + subtensor_state["IsNetworkMember"][hotkey][netuid][self.block_number] = True + + return uid + @staticmethod def _convert_to_balance(balance: Union["Balance", float, int]) -> "Balance": if isinstance(balance, float): @@ -368,6 +467,37 @@ def _convert_to_balance(balance: Union["Balance", float, int]) -> "Balance": return balance + def force_register_neuron( + self, + netuid: int, + hotkey: str, + coldkey: str, + stake: Union["Balance", float, int] = Balance(0), + balance: Union["Balance", float, int] = Balance(0), + ) -> int: + """ + Force register a neuron on the mock chain, returning the UID. + """ + stake = self._convert_to_balance(stake) + balance = self._convert_to_balance(balance) + + subtensor_state = self.chain_state["SubtensorModule"] + if netuid not in subtensor_state["NetworksAdded"]: + raise Exception("Subnet does not exist") + + uid = self._register_neuron(netuid=netuid, hotkey=hotkey, coldkey=coldkey) + + subtensor_state["TotalStake"][self.block_number] = ( + self._get_most_recent_storage(subtensor_state["TotalStake"]) + stake.rao + ) + subtensor_state["Stake"][hotkey][coldkey][self.block_number] = stake.rao + + if balance.rao > 0: + self.force_set_balance(coldkey, balance) + self.force_set_balance(coldkey, balance) + + return uid + def force_set_balance( self, ss58_address: str, balance: Union["Balance", float, int] = Balance(0) ) -> tuple[bool, Optional[str]]: From fd1e2d8b1b8a4cdc719a564e732f9e96d6a155e0 Mon Sep 17 00:00:00 2001 From: Roman Date: Fri, 18 Apr 2025 11:28:06 -0700 Subject: [PATCH 2/3] add test --- .../test_subtensor_integration.py | 28 +++++++++++++++++-- 1 file changed, 26 insertions(+), 2 deletions(-) diff --git a/tests/integration_tests/test_subtensor_integration.py b/tests/integration_tests/test_subtensor_integration.py index 46cb382671..db8bf19453 100644 --- a/tests/integration_tests/test_subtensor_integration.py +++ b/tests/integration_tests/test_subtensor_integration.py @@ -1,10 +1,10 @@ import pytest -from bittensor import NeuronInfo -from bittensor.core.chain_data.axon_info import AxonInfo +from bittensor.core.chain_data import AxonInfo, NeuronInfo from bittensor.core.subtensor import Subtensor from bittensor.utils.balance import Balance from tests.helpers.helpers import FakeWebsocket +from bittensor.utils.mock.subtensor_mock import MockSubtensor @pytest.fixture @@ -126,3 +126,27 @@ async def test_get_neuron_for_pubkey_and_subnet(mocker): assert isinstance(result.total_stake, Balance) assert isinstance(result.axon_info, AxonInfo) assert result.is_null is False + + +def test_mock_subtensor_force_register_neuron(): + """Tests the force_register_neuron method of the MockSubtensor class.""" + # Preps + test_netuid = 1 + subtensor = MockSubtensor() + subtensor.create_subnet(netuid=test_netuid) + + uid1 = subtensor.force_register_neuron(test_netuid, "hk1", "cc1") + uid2 = subtensor.force_register_neuron(test_netuid, "hk2", "cc2") + + # Calls + neurons = subtensor.neurons(test_netuid) + neuron1 = subtensor.neuron_for_uid(uid1, test_netuid) + neuron2 = subtensor.neuron_for_uid(uid2, test_netuid) + + # Assertions + assert len(neurons) == 2 + assert [neuron1, neuron2] == neurons + assert neuron1.hotkey == "hk1" + assert neuron1.coldkey == "cc1" + assert neuron2.hotkey == "hk2" + assert neuron2.coldkey == "cc2" From 05b433a22682bc311018d788f40a2bab40abab45 Mon Sep 17 00:00:00 2001 From: Roman Date: Fri, 18 Apr 2025 11:28:12 -0700 Subject: [PATCH 3/3] imports --- tests/integration_tests/test_metagraph_integration.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/integration_tests/test_metagraph_integration.py b/tests/integration_tests/test_metagraph_integration.py index efcb49d267..005b1ae0c3 100644 --- a/tests/integration_tests/test_metagraph_integration.py +++ b/tests/integration_tests/test_metagraph_integration.py @@ -1,10 +1,11 @@ +import os from unittest import mock -import bittensor import torch -import os -from bittensor.utils.mock import MockSubtensor + +import bittensor from bittensor.core.metagraph import METAGRAPH_STATE_DICT_NDARRAY_KEYS, get_save_dir +from bittensor.utils.mock import MockSubtensor _subtensor_mock: MockSubtensor = MockSubtensor()