Merge branch 'staging' into tests/zyzniewski/more_e2e_tests

Author: zyzniewski-reef

bittensor/core/async_subtensor.py

@@ -1479,6 +1479,41 @@ async def get_neuron_certificate(
             return None
         return None
 
+    async def get_all_neuron_certificates(
+        self,
+        netuid: int,
+        block: Optional[int] = None,
+        block_hash: Optional[str] = None,
+        reuse_block: bool = False,
+    ) -> dict[str, Certificate]:
+        """
+        Retrieves the TLS certificates for neurons within a specified subnet (netuid) of the Bittensor network.
+
+        Arguments:
+            netuid: The unique identifier of the subnet.
+            block: The blockchain block number for the query.
+            block_hash: The hash of the block to retrieve the parameter from. Do not specify if using block or
+                reuse_block.
+            reuse_block: Whether to use the last-used block. Do not set if using block_hash or block.
+
+        Returns:
+            {ss58: Certificate} for the key/Certificate pairs on the subnet
+
+        This function is used for certificate discovery for setting up mutual tls communication between neurons.
+        """
+        query_certificates = await self.query_map(
+            module="SubtensorModule",
+            name="NeuronCertificates",
+            params=[netuid],
+            block=block,
+            block_hash=block_hash,
+            reuse_block=reuse_block,
+        )
+        output = {}
+        async for key, item in query_certificates:
+            output[decode_account_id(key)] = Certificate(item.value)
+        return output
+
     async def get_neuron_for_pubkey_and_subnet(
         self,
         hotkey_ss58: str,

bittensor/core/chain_data/metagraph_info.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass
 from typing import Optional, Union
 
+from bittensor.core import settings
 from bittensor.core.chain_data.axon_info import AxonInfo
 from bittensor.core.chain_data.chain_identity import ChainIdentity
 from bittensor.core.chain_data.info_base import InfoBase
@@ -234,7 +235,10 @@ def _from_dict(cls, decoded: dict) -> "MetagraphInfo":
             rank=[u16tf(rk) for rk in decoded.get("rank", [])],
             block_at_registration=decoded["block_at_registration"],
             alpha_stake=[_tbwu(ast, _netuid) for ast in decoded["alpha_stake"]],
-            tao_stake=[_tbwu(ts) for ts in decoded["tao_stake"]],
+            tao_stake=[
+                _tbwu(ts) * settings.ROOT_TAO_STAKE_WEIGHT
+                for ts in decoded["tao_stake"]
+            ],
             total_stake=[_tbwu(ts, _netuid) for ts in decoded["total_stake"]],
             # Dividend break down
             tao_dividends_per_hotkey=[

bittensor/core/metagraph.py

@@ -42,7 +42,6 @@
 
 
 Tensor = Union["torch.nn.Parameter", NDArray]
-ROOT_TAO_STAKES_WEIGHT = 0.18
 
 
 METAGRAPH_STATE_DICT_NDARRAY_KEYS = [
@@ -1598,7 +1597,10 @@ async def _get_all_stakes_from_chain(self):
                 dtype=self._dtype_registry["float32"],
             )
             self.tao_stake = self._create_tensor(
-                [b.tao * ROOT_TAO_STAKES_WEIGHT for b in subnet_state.tao_stake],
+                [
+                    b.tao * settings.ROOT_TAO_STAKE_WEIGHT
+                    for b in subnet_state.tao_stake
+                ],
                 dtype=self._dtype_registry["float32"],
             )
             self.total_stake = self.stake = self._create_tensor(
@@ -1902,7 +1904,10 @@ def _get_all_stakes_from_chain(self):
                 dtype=self._dtype_registry["float32"],
             )
             self.tao_stake = self._create_tensor(
-                [b.tao * ROOT_TAO_STAKES_WEIGHT for b in subnet_state.tao_stake],
+                [
+                    b.tao * settings.ROOT_TAO_STAKE_WEIGHT
+                    for b in subnet_state.tao_stake
+                ],
                 dtype=self._dtype_registry["float32"],
             )
             self.total_stake = self.stake = self._create_tensor(

bittensor/core/settings.py

@@ -6,6 +6,7 @@
 
 from munch import munchify
 
+ROOT_TAO_STAKE_WEIGHT = 0.18
 
 READ_ONLY = os.getenv("READ_ONLY") == "1"
 

bittensor/core/subtensor.py

@@ -1123,6 +1123,32 @@ def get_neuron_certificate(
             return None
         return None
 
+    def get_all_neuron_certificates(
+        self, netuid: int, block: Optional[int] = None
+    ) -> dict[str, Certificate]:
+        """
+        Retrieves the TLS certificates for neurons within a specified subnet (netuid) of the Bittensor network.
+
+        Arguments:
+            netuid: The unique identifier of the subnet.
+            block: The blockchain block number for the query.
+
+        Returns:
+            {ss58: Certificate} for the key/Certificate pairs on the subnet
+
+        This function is used for certificate discovery for setting up mutual tls communication between neurons.
+        """
+        query_certificates = self.query_map(
+            module="SubtensorModule",
+            name="NeuronCertificates",
+            params=[netuid],
+            block=block,
+        )
+        output = {}
+        for key, item in query_certificates:
+            output[decode_account_id(key)] = Certificate(item.value)
+        return output
+
     def get_neuron_for_pubkey_and_subnet(
         self, hotkey_ss58: str, netuid: int, block: Optional[int] = None
     ) -> Optional["NeuronInfo"]:

bittensor/utils/balance.py

@@ -53,8 +53,13 @@ def __float__(self):
         return self.tao
 
     def __str__(self):
-        """Returns the Balance object as a string in the format "symbolvalue", where the value is in tao."""
-        return f"{self.unit}{float(self.tao):,.9f}"
+        """
+        Returns the Balance object as a string in the format "symbolvalue", where the value is in tao.
+        """
+        if self.unit == units[0]:
+            return f"{self.unit}{float(self.tao):,.9f}"
+        else:
+            return f"\u200e{float(self.tao):,.9f}{self.unit}\u200e"
 
     def __rich__(self):
         int_tao, fract_tao = format(float(self.tao), "f").split(".")

tests/e2e_tests/test_neuron_certificate.py

@@ -48,5 +48,8 @@ async def test_neuron_certificate(subtensor, alice_wallet):
         )
         == encoded_certificate
     )
+    all_certs_query = subtensor.get_all_neuron_certificates(netuid=netuid)
+    assert alice_wallet.hotkey.ss58_address in all_certs_query.keys()
+    assert all_certs_query[alice_wallet.hotkey.ss58_address] == encoded_certificate
 
     logging.info("✅ Passed test_neuron_certificate")

tests/unit_tests/test_async_subtensor.py

@@ -2676,3 +2676,18 @@ async def test_set_subnet_identity(mocker, subtensor):
         wait_for_inclusion=False,
     )
     assert result == mocked_extrinsic.return_value
+
+
+@pytest.mark.asyncio
+async def test_get_all_neuron_certificates(mocker, subtensor):
+    fake_netuid = 12
+    mocked_query_map_subtensor = mocker.AsyncMock()
+    mocker.patch.object(subtensor.substrate, "query_map", mocked_query_map_subtensor)
+    await subtensor.get_all_neuron_certificates(fake_netuid)
+    mocked_query_map_subtensor.assert_awaited_once_with(
+        module="SubtensorModule",
+        storage_function="NeuronCertificates",
+        params=[fake_netuid],
+        block_hash=None,
+        reuse_block_hash=False,
+    )

tests/unit_tests/test_subtensor.py

@@ -3060,3 +3060,16 @@ def test_set_subnet_identity(mocker, subtensor):
         wait_for_inclusion=False,
     )
     assert result == mocked_extrinsic.return_value
+
+
+def test_get_all_neuron_certificates(mocker, subtensor):
+    fake_netuid = 12
+    mocked_query_map_subtensor = mocker.MagicMock()
+    mocker.patch.object(subtensor.substrate, "query_map", mocked_query_map_subtensor)
+    subtensor.get_all_neuron_certificates(fake_netuid)
+    mocked_query_map_subtensor.assert_called_once_with(
+        module="SubtensorModule",
+        storage_function="NeuronCertificates",
+        params=[fake_netuid],
+        block_hash=None,
+    )