Merge pull request #2611 from opentensor/feat/roman/update-metagraph-class

Update metagraph class with `rao` stuff

Author: basfroman

bittensor/core/chain_data/subnet_state.py

@@ -4,7 +4,7 @@
 """
 
 from dataclasses import dataclass
-from typing import Optional
+from typing import Optional, Union
 
 from scalecodec.utils.ss58 import ss58_encode
 
@@ -39,7 +39,7 @@ class SubnetState:
     emission_history: list[list[int]]
 
     @classmethod
-    def from_vec_u8(cls, vec_u8: list[int]) -> Optional["SubnetState"]:
+    def from_vec_u8(cls, vec_u8: Union[list[int], bytes]) -> Optional["SubnetState"]:
         if len(vec_u8) == 0:
             return None
         decoded = from_scale_encoding(vec_u8, ChainDataType.SubnetState)

bittensor/core/metagraph.py

@@ -9,23 +9,25 @@
 from typing import Optional, Union
 
 import numpy as np
+from async_substrate_interface.errors import SubstrateRequestException
 from numpy.typing import NDArray
 
+from bittensor.core import settings
+from bittensor.core.chain_data import AxonInfo, SubnetState
 from bittensor.utils.btlogging import logging
 from bittensor.utils.registration import torch, use_torch
 from bittensor.utils.weight_utils import (
     convert_weight_uids_and_vals_to_tensor,
     convert_bond_uids_and_vals_to_tensor,
     convert_root_weight_uids_and_vals_to_tensor,
 )
-from bittensor.core import settings
-from bittensor.core.chain_data import AxonInfo
 
 # For annotation purposes
 if typing.TYPE_CHECKING:
     from bittensor.core.subtensor import Subtensor
     from bittensor.core.async_subtensor import AsyncSubtensor
     from bittensor.core.chain_data import NeuronInfo, NeuronInfoLite
+    from bittensor.utils.balance import Balance
 
 
 Tensor = Union["torch.nn.Parameter", NDArray]
@@ -219,8 +221,6 @@ class MetagraphMixin(ABC):
     n: Tensor
     neurons: list[Union["NeuronInfo", "NeuronInfoLite"]]
     block: Tensor
-    stake: Tensor
-    total_stake: Tensor
     ranks: Tensor
     trust: Tensor
     consensus: Tensor
@@ -234,11 +234,34 @@ class MetagraphMixin(ABC):
     weights: Tensor
     bonds: Tensor
     uids: Tensor
+    alpha_stake: Tensor
+    tao_stake: Tensor
+    stake: Tensor
     axons: list[AxonInfo]
     chain_endpoint: Optional[str]
     subtensor: Optional["AsyncSubtensor"]
     _dtype_registry = {"int64": np.int64, "float32": np.float32, "bool": bool}
 
+    @property
+    def TS(self) -> list["Balance"]:
+        """
+        Represents the tao stake of each neuron in the Bittensor network.
+
+        Returns:
+            list["Balance"]: The list of tao stake of each neuron in the network.
+        """
+        return self.tao_stake
+
+    @property
+    def AS(self) -> list["Balance"]:
+        """
+        Represents the alpha stake of each neuron in the Bittensor network.
+
+        Returns:
+            list["Balance"]: The list of alpha stake of each neuron in the network.
+        """
+        return self.alpha_stake
+
     @property
     def S(self) -> Union[NDArray, "torch.nn.Parameter"]:
         """
@@ -251,7 +274,7 @@ def S(self) -> Union[NDArray, "torch.nn.Parameter"]:
             NDArray: A tensor representing the stake of each neuron in the network. Higher values signify a greater
                 stake held by the respective neuron.
         """
-        return self.total_stake
+        return self.stake
 
     @property
     def R(self) -> Union[NDArray, "torch.nn.Parameter"]:
@@ -554,8 +577,6 @@ def state_dict(self):
             "version": self.version,
             "n": self.n,
             "block": self.block,
-            "stake": self.stake,
-            "total_stake": self.total_stake,
             "ranks": self.ranks,
             "trust": self.trust,
             "consensus": self.consensus,
@@ -571,6 +592,9 @@ def state_dict(self):
             "uids": self.uids,
             "axons": self.axons,
             "neurons": self.neurons,
+            "alpha_stake": self.alpha_stake,
+            "tao_stake": self.tao_stake,
+            "stake": self.stake,
         }
 
     @staticmethod
@@ -1284,6 +1308,9 @@ async def sync(
         if not lite:
             await self._set_weights_and_bonds(subtensor=subtensor)
 
+        # Fills in the stake associated attributes of a class instance from a chain response.
+        await self._get_all_stakes_from_chain(subtensor=subtensor)
+
     async def _initialize_subtensor(
         self, subtensor: "AsyncSubtensor"
     ) -> "AsyncSubtensor":
@@ -1448,6 +1475,46 @@ async def _process_root_weights(
             )
         return tensor_param
 
+    async def _get_all_stakes_from_chain(
+        self, subtensor: Optional["AsyncSubtensor"] = None
+    ):
+        """Fills in the stake associated attributes of a class instance from a chain response."""
+        try:
+            if not subtensor:
+                subtensor = self._initialize_subtensor(subtensor=subtensor)
+
+            hex_bytes_result = await subtensor.query_runtime_api(
+                runtime_api="SubnetInfoRuntimeApi",
+                method="get_subnet_state",
+                params=[self.netuid],
+            )
+
+            if hex_bytes_result is None:
+                logging.debug(
+                    f"Unable to retrieve subnet state for netuid `{self.netuid}`."
+                )
+                return []
+
+            if hex_bytes_result.startswith("0x"):
+                bytes_result = bytes.fromhex(hex_bytes_result[2:])
+            else:
+                bytes_result = bytes.fromhex(hex_bytes_result)
+
+            subnet_state: "SubnetState" = SubnetState.from_vec_u8(bytes_result)
+            if self.netuid == 0:
+                self.total_stake = self.stake = self.tao_stake = self.alpha_stake = (
+                    subnet_state.tao_stake
+                )
+                return subnet_state
+
+            self.alpha_stake = subnet_state.alpha_stake
+            self.tao_stake = [b * 0.018 for b in subnet_state.tao_stake]
+            self.total_stake = self.stake = subnet_state.total_stake
+            return subnet_state
+
+        except (SubstrateRequestException, AttributeError) as e:
+            logging.debug(e)
+
 
 class Metagraph(NumpyOrTorch):
     def __init__(
@@ -1512,6 +1579,8 @@ def sync(
 
                 metagraph.sync(block=history_block, lite=False, subtensor=subtensor)
         """
+
+        # Initialize subtensor
         subtensor = self._initialize_subtensor(subtensor)
 
         if (
@@ -1538,6 +1607,9 @@ def sync(
         if not lite:
             self._set_weights_and_bonds(subtensor=subtensor)
 
+        # Fills in the stake associated attributes of a class instance from a chain response.
+        self._get_all_stakes_from_chain(subtensor=subtensor)
+
     def _initialize_subtensor(self, subtensor: "Subtensor") -> "Subtensor":
         """
         Initializes the subtensor to be used for syncing the metagraph.
@@ -1694,6 +1766,44 @@ def _process_root_weights(
             )
         return tensor_param
 
+    def _get_all_stakes_from_chain(self, subtensor: Optional["Subtensor"] = None):
+        """Fills in the stake associated attributes of a class instance from a chain response."""
+        try:
+            if not subtensor:
+                subtensor = self._initialize_subtensor()
+
+            hex_bytes_result = subtensor.query_runtime_api(
+                runtime_api="SubnetInfoRuntimeApi",
+                method="get_subnet_state",
+                params=[self.netuid],
+            )
+
+            if hex_bytes_result is None:
+                logging.debug(
+                    f"Unable to retrieve subnet state for netuid `{self.netuid}`."
+                )
+                return []
+
+            if hex_bytes_result.startswith("0x"):
+                bytes_result = bytes.fromhex(hex_bytes_result[2:])
+            else:
+                bytes_result = bytes.fromhex(hex_bytes_result)
+
+            subnet_state: "SubnetState" = SubnetState.from_vec_u8(bytes_result)
+            if self.netuid == 0:
+                self.total_stake = self.stake = self.tao_stake = self.alpha_stake = (
+                    subnet_state.tao_stake
+                )
+                return subnet_state
+
+            self.alpha_stake = subnet_state.alpha_stake
+            self.tao_stake = [b * 0.018 for b in subnet_state.tao_stake]
+            self.total_stake = self.stake = subnet_state.total_stake
+            return subnet_state
+
+        except (SubstrateRequestException, AttributeError) as e:
+            logging.debug(e)
+
 
 async def async_metagraph(
     netuid: int,