Merge pull request #2659 from opentensor/fix/thewhaleking/fix-async-torch-metagraph

Metagraph Improvements

Author: basfroman

bittensor/core/metagraph.py

@@ -42,6 +42,7 @@
 
 
 Tensor = Union["torch.nn.Parameter", NDArray]
+ROOT_TAO_STAKES_WEIGHT = 0.018
 
 
 METAGRAPH_STATE_DICT_NDARRAY_KEYS = [
@@ -542,6 +543,15 @@ def __init__(
 
                 metagraph = Metagraph(netuid=123, network="finney", lite=True, sync=True)
         """
+        self.lite = lite
+        self.subtensor = subtensor
+        self.should_sync = sync
+        self.netuid = netuid
+        self.network, self.chain_endpoint = determine_chain_endpoint_and_network(
+            network
+        )
+        self.neurons = []
+        self.axons: list[AxonInfo] = []
 
     def __str__(self) -> str:
         """
@@ -1053,10 +1063,6 @@ def __init__(
         """
         BaseClass.__init__(self)
         MetagraphMixin.__init__(self, netuid, network, lite, sync, subtensor)
-        self.netuid = netuid
-        self.network, self.chain_endpoint = determine_chain_endpoint_and_network(
-            network
-        )
         self._dtype_registry = {
             "int64": torch.int64,
             "float32": torch.float32,
@@ -1117,10 +1123,6 @@ def __init__(
         self.uids = torch.nn.Parameter(
             torch.tensor([], dtype=torch.int64), requires_grad=False
         )
-        self.axons: list[AxonInfo] = []
-        self.neurons = []
-        self.subtensor = subtensor
-        self.should_sync = sync
         self.alpha_stake = torch.nn.Parameter(
             torch.tensor([], dtype=torch.float32), requires_grad=False
         )
@@ -1249,9 +1251,6 @@ def __init__(
         self.tao_stake: Tensor = np.array([], dtype=np.int64)
         self.stake: Tensor = np.array([], dtype=np.int64)
         self.total_stake: Tensor = np.array([], dtype=np.int64)
-
-        self.axons: list[AxonInfo] = []
-        self.neurons = []
         self.subtensor = subtensor
         self.should_sync = sync
 
@@ -1353,7 +1352,7 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
     async def sync(
         self,
         block: Optional[int] = None,
-        lite: bool = True,
+        lite: Optional[bool] = None,
         subtensor: Optional["AsyncSubtensor"] = None,
     ):
         """
@@ -1364,8 +1363,9 @@ async def sync(
         Args:
             block (Optional[int]): A specific block number to synchronize with. If None, the metagraph syncs with the
                 latest block. This allows for historical analysis or specific state examination of the network.
-            lite (bool): If True, a lite version of the metagraph is used for quicker synchronization. This is
+            lite (Optional[bool]): If True, a lite version of the metagraph is used for quicker synchronization. This is
                 beneficial when full detail is not necessary, allowing for reduced computational and time overhead.
+                Defaults to `True`.
             subtensor (Optional[bittensor.core.subtensor.Subtensor]): An instance of the subtensor class from Bittensor,
                 providing an interface to the underlying blockchain data. If provided, this instance is used for data
                 retrieval during synchronization.
@@ -1400,6 +1400,9 @@ async def sync(
 
                 metagraph.sync(block=history_block, lite=False, subtensor=subtensor)
         """
+        if lite is None:
+            lite = self.lite
+
         subtensor = await self._initialize_subtensor(subtensor)
 
         if (
@@ -1618,8 +1621,14 @@ async def _get_all_stakes_from_chain(self):
                 )
                 return subnet_state
 
-            self.alpha_stake = subnet_state.alpha_stake
-            self.tao_stake = [b * 0.018 for b in subnet_state.tao_stake]
+            self.alpha_stake = self._create_tensor(
+                [b.tao for b in subnet_state.alpha_stake],
+                dtype=self._dtype_registry["float32"],
+            )
+            self.tao_stake = self._create_tensor(
+                [b.tao * ROOT_TAO_STAKES_WEIGHT for b in subnet_state.tao_stake],
+                dtype=self._dtype_registry["float32"],
+            )
             self.total_stake = self.stake = self._create_tensor(
                 [stake.tao for stake in subnet_state.total_stake],
                 dtype=self._dtype_registry["float32"],
@@ -1651,7 +1660,7 @@ def __init__(
     def sync(
         self,
         block: Optional[int] = None,
-        lite: bool = True,
+        lite: Optional[bool] = None,
         subtensor: Optional["Subtensor"] = None,
     ):
         """
@@ -1662,8 +1671,9 @@ def sync(
         Args:
             block (Optional[int]): A specific block number to synchronize with. If None, the metagraph syncs with the
                 latest block. This allows for historical analysis or specific state examination of the network.
-            lite (bool): If True, a lite version of the metagraph is used for quicker synchronization. This is
+            lite (Optional[bool]): If True, a lite version of the metagraph is used for quicker synchronization. This is
                 beneficial when full detail is not necessary, allowing for reduced computational and time overhead.
+                Defaults to `True`.
             subtensor (Optional[bittensor.core.subtensor.Subtensor]): An instance of the subtensor class from Bittensor,
                 providing an interface to the underlying blockchain data. If provided, this instance is used for data
                 retrieval during synchronization.
@@ -1698,6 +1708,8 @@ def sync(
 
                 metagraph.sync(block=history_block, lite=False, subtensor=subtensor)
         """
+        if lite is None:
+            lite = self.lite
 
         # Initialize subtensor
         subtensor = self._initialize_subtensor(subtensor=subtensor)
@@ -1918,7 +1930,7 @@ def _get_all_stakes_from_chain(self):
                 dtype=self._dtype_registry["float32"],
             )
             self.tao_stake = self._create_tensor(
-                [b.tao * 0.018 for b in subnet_state.tao_stake],
+                [b.tao * ROOT_TAO_STAKES_WEIGHT for b in subnet_state.tao_stake],
                 dtype=self._dtype_registry["float32"],
             )
             self.total_stake = self.stake = self._create_tensor(