Merge pull request #2732 from opentensor/fix/zyzniewski/burned_register_supports_root_subnet

Fix: burned_register supports root subnet (netuid=0 param)

Author: zyzniewski-reef

bittensor/core/async_subtensor.py

@@ -2915,6 +2915,14 @@ async def burned_register(
             bool: ``True`` if the registration is successful, False otherwise.
         """
         async with self:
+            if netuid == 0:
+                return await root_register_extrinsic(
+                    subtensor=self,
+                    wallet=wallet,
+                    wait_for_inclusion=wait_for_inclusion,
+                    wait_for_finalization=wait_for_finalization,
+                )
+
             return await burned_register_extrinsic(
                 subtensor=self,
                 wallet=wallet,
@@ -3205,35 +3213,6 @@ async def root_register(
         Returns:
             `True` if registration was successful, otherwise `False`.
         """
-        netuid = 0
-        logging.info(
-            f"Registering on netuid [blue]0[/blue] on network: [blue]{self.network}[/blue]"
-        )
-
-        # Check current recycle amount
-        logging.info("Fetching recycle amount & balance.")
-        block_hash = block_hash if block_hash else await self.get_block_hash()
-
-        try:
-            recycle_call, balance = await asyncio.gather(
-                self.get_hyperparameter(
-                    param_name="Burn", netuid=netuid, block_hash=block_hash
-                ),
-                self.get_balance(wallet.coldkeypub.ss58_address, block_hash=block_hash),
-            )
-        except TypeError as e:
-            logging.error(f"Unable to retrieve current recycle. {e}")
-            return False
-
-        current_recycle = Balance.from_rao(int(recycle_call))
-
-        # Check balance is sufficient
-        if balance < current_recycle:
-            logging.error(
-                f"[red]Insufficient balance {balance} to register neuron. "
-                f"Current recycle is {current_recycle} TAO[/red]."
-            )
-            return False
 
         return await root_register_extrinsic(
             subtensor=self,

bittensor/core/extrinsics/asyncex/root.py

@@ -7,6 +7,7 @@
 
 from bittensor.core.errors import SubstrateRequestException
 from bittensor.utils import u16_normalized_float, format_error_message, unlock_key
+from bittensor.utils.balance import Balance
 from bittensor.utils.btlogging import logging
 from bittensor.utils.weight_utils import (
     normalize_max_weight,
@@ -62,6 +63,33 @@ async def root_register_extrinsic(
             the response is `True`.
     """
     netuid = 0
+    logging.info(
+        f"Registering on netuid [blue]{netuid}[/blue] on network: [blue]{subtensor.network}[/blue]"
+    )
+
+    logging.info("Fetching recycle amount & balance.")
+    block_hash = await subtensor.get_block_hash()
+    recycle_call, balance = await asyncio.gather(
+        subtensor.get_hyperparameter(
+            param_name="Burn",
+            netuid=netuid,
+            block_hash=block_hash,
+        ),
+        subtensor.get_balance(
+            wallet.coldkeypub.ss58_address,
+            block_hash=block_hash,
+        ),
+    )
+
+    current_recycle = Balance.from_rao(int(recycle_call))
+
+    if balance < current_recycle:
+        logging.error(
+            f"[red]Insufficient balance {balance} to register neuron. "
+            f"Current recycle is {current_recycle} TAO[/red]."
+        )
+        return False
+
     if not (unlock := unlock_key(wallet)).success:
         logging.error(unlock.message)
         return False

bittensor/core/extrinsics/root.py

@@ -11,6 +11,7 @@
     unlock_key,
     torch,
 )
+from bittensor.utils.balance import Balance
 from bittensor.utils.btlogging import logging
 from bittensor.utils.weight_utils import (
     normalize_max_weight,
@@ -65,6 +66,30 @@ def root_register_extrinsic(
             response is `True`.
     """
     netuid = 0
+    logging.info(
+        f"Registering on netuid [blue]{netuid}[/blue] on network: [blue]{subtensor.network}[/blue]"
+    )
+
+    logging.info("Fetching recycle amount & balance.")
+    block = subtensor.get_current_block()
+    recycle_call = subtensor.get_hyperparameter(
+        param_name="Burn",
+        netuid=netuid,
+        block=block,
+    )
+    balance = subtensor.get_balance(
+        wallet.coldkeypub.ss58_address,
+        block=block,
+    )
+
+    current_recycle = Balance.from_rao(int(recycle_call))
+
+    if balance < current_recycle:
+        logging.error(
+            f"[red]Insufficient balance {balance} to register neuron. "
+            f"Current recycle is {current_recycle} TAO[/red]."
+        )
+        return False
 
     if not (unlock := unlock_key(wallet)).success:
         logging.error(unlock.message)

bittensor/core/subtensor.py

@@ -2234,6 +2234,15 @@ def burned_register(
         Returns:
             bool: ``True`` if the registration is successful, False otherwise.
         """
+
+        if netuid == 0:
+            return root_register_extrinsic(
+                subtensor=self,
+                wallet=wallet,
+                wait_for_inclusion=wait_for_inclusion,
+                wait_for_finalization=wait_for_finalization,
+            )
+
         return burned_register_extrinsic(
             subtensor=self,
             wallet=wallet,
@@ -2521,32 +2530,6 @@ def root_register(
         Returns:
             `True` if registration was successful, otherwise `False`.
         """
-        logging.info(
-            f"Registering on netuid [blue]0[/blue] on network: [blue]{self.network}[/blue]"
-        )
-
-        # Check current recycle amount
-        logging.info("Fetching recycle amount & balance.")
-        block = self.get_current_block()
-
-        try:
-            recycle_call = cast(
-                int, self.get_hyperparameter(param_name="Burn", netuid=0, block=block)
-            )
-            balance = self.get_balance(wallet.coldkeypub.ss58_address, block=block)
-        except TypeError as e:
-            logging.error(f"Unable to retrieve current recycle. {e}")
-            return False
-
-        current_recycle = Balance.from_rao(int(recycle_call))
-
-        # Check balance is sufficient
-        if balance < current_recycle:
-            logging.error(
-                f"[red]Insufficient balance {balance} to register neuron. "
-                f"Current recycle is {current_recycle} TAO[/red]."
-            )
-            return False
 
         return root_register_extrinsic(
             subtensor=self,

tests/unit_tests/extrinsics/asyncex/test_root.py

@@ -3,6 +3,8 @@
 from bittensor.core.errors import SubstrateRequestException
 from bittensor.core.extrinsics.asyncex import root as async_root
 
+from bittensor.utils.balance import Balance
+
 
 @pytest.mark.asyncio
 async def test_get_limits_success(subtensor, mocker):
@@ -62,6 +64,16 @@ async def test_root_register_extrinsic_success(subtensor, fake_wallet, mocker):
         "query",
         return_value=fake_uid,
     )
+    mocker.patch.object(
+        subtensor,
+        "get_hyperparameter",
+        return_value=Balance(0),
+    )
+    mocker.patch.object(
+        subtensor,
+        "get_balance",
+        return_value=Balance(1),
+    )
 
     # Call
     result = await async_root.root_register_extrinsic(
@@ -86,10 +98,53 @@ async def test_root_register_extrinsic_success(subtensor, fake_wallet, mocker):
     assert result is True
 
 
+@pytest.mark.asyncio
+async def test_root_register_extrinsic_insufficient_balance(
+    subtensor,
+    fake_wallet,
+    mocker,
+):
+    mocker.patch.object(
+        subtensor,
+        "get_hyperparameter",
+        return_value=Balance(1),
+    )
+    mocker.patch.object(
+        subtensor,
+        "get_balance",
+        return_value=Balance(0),
+    )
+
+    result = await async_root.root_register_extrinsic(
+        subtensor=subtensor,
+        wallet=fake_wallet,
+        wait_for_inclusion=True,
+        wait_for_finalization=True,
+    )
+
+    assert result is False
+
+    subtensor.get_balance.assert_called_once_with(
+        fake_wallet.coldkeypub.ss58_address,
+        block_hash=subtensor.substrate.get_chain_head.return_value,
+    )
+    subtensor.substrate.submit_extrinsic.assert_not_called()
+
+
 @pytest.mark.asyncio
 async def test_root_register_extrinsic_unlock_failed(subtensor, fake_wallet, mocker):
     """Tests registration fails due to unlock failure."""
     # Preps
+    mocker.patch.object(
+        subtensor,
+        "get_hyperparameter",
+        return_value=Balance(0),
+    )
+    mocker.patch.object(
+        subtensor,
+        "get_balance",
+        return_value=Balance(1),
+    )
     mocked_unlock_key = mocker.patch.object(
         async_root,
         "unlock_key",
@@ -117,6 +172,16 @@ async def test_root_register_extrinsic_already_registered(
     # Preps
     fake_wallet.hotkey.ss58_address = "fake_hotkey_address"
 
+    mocker.patch.object(
+        subtensor,
+        "get_hyperparameter",
+        return_value=Balance(0),
+    )
+    mocker.patch.object(
+        subtensor,
+        "get_balance",
+        return_value=Balance(1),
+    )
     mocked_unlock_key = mocker.patch.object(
         async_root,
         "unlock_key",
@@ -152,6 +217,16 @@ async def test_root_register_extrinsic_transaction_failed(
     # Preps
     fake_wallet.hotkey.ss58_address = "fake_hotkey_address"
 
+    mocker.patch.object(
+        subtensor,
+        "get_hyperparameter",
+        return_value=Balance(0),
+    )
+    mocker.patch.object(
+        subtensor,
+        "get_balance",
+        return_value=Balance(1),
+    )
     mocked_unlock_key = mocker.patch.object(
         async_root,
         "unlock_key",
@@ -193,6 +268,16 @@ async def test_root_register_extrinsic_uid_not_found(subtensor, fake_wallet, moc
     # Preps
     fake_wallet.hotkey.ss58_address = "fake_hotkey_address"
 
+    mocker.patch.object(
+        subtensor,
+        "get_hyperparameter",
+        return_value=Balance(0),
+    )
+    mocker.patch.object(
+        subtensor,
+        "get_balance",
+        return_value=Balance(1),
+    )
     mocked_unlock_key = mocker.patch.object(
         async_root,
         "unlock_key",

tests/unit_tests/extrinsics/test_root.py

@@ -1,6 +1,7 @@
 import pytest
 from bittensor.core.subtensor import Subtensor
 from bittensor.core.extrinsics import root
+from bittensor.utils.balance import Balance
 
 
 @pytest.fixture
@@ -75,6 +76,11 @@ def test_root_register_extrinsic(
         "query",
         return_value=hotkey_registered[1],
     )
+    mocker.patch.object(
+        mock_subtensor,
+        "get_balance",
+        return_value=Balance(1),
+    )
 
     # Act
     result = root.root_register_extrinsic(
@@ -100,6 +106,31 @@ def test_root_register_extrinsic(
         )
 
 
+def test_root_register_extrinsic_insufficient_balance(
+    mock_subtensor,
+    mock_wallet,
+    mocker,
+):
+    mocker.patch.object(
+        mock_subtensor,
+        "get_balance",
+        return_value=Balance(0),
+    )
+
+    success = root.root_register_extrinsic(
+        subtensor=mock_subtensor,
+        wallet=mock_wallet,
+    )
+
+    assert success is False
+
+    mock_subtensor.get_balance.assert_called_once_with(
+        mock_wallet.coldkeypub.ss58_address,
+        block=mock_subtensor.get_current_block.return_value,
+    )
+    mock_subtensor.substrate.submit_extrinsic.assert_not_called()
+
+
 @pytest.mark.parametrize(
     "wait_for_inclusion, wait_for_finalization, netuids, weights, expected_success",
     [