[Revolution] add associated ip info (#1486)

* add ipinfo to type map and chaindata

* not an option

* add subtensor call for pulling assoc ip

* add do_associate_ips call

* add encode func

* black format

* add state calls api

* add overload for decoding from multiple types

* use new state call to grab without custom rpc

* oops, use correct function

* add validator ip to runtimeapi registry spec

* remove duplicate methods from merge and use new logic

* run black

---------

Co-authored-by: ifrit98 <ifrit98@gmail.com>

bittensor/__init__.py

@@ -164,6 +164,19 @@ def debug(on: bool = True):
                 },
             },
         },
+        "ValidatorIPRuntimeApi": {
+            "methods": {
+                "get_associated_validator_ip_info_for_subnet": {
+                    "params": [
+                        {
+                            "name": "netuid",
+                            "type": "u16",
+                        },
+                    ],
+                    "type": "Vec<u8>",
+                },
+            },
+        },
     },
 }
 

bittensor/chain_data.py

@@ -137,6 +137,13 @@
                 ["ip_type", "u8"],
             ],
         },
+        "IPInfo": {
+            "type": "struct",
+            "type_mapping": [
+                ["ip", "Compact<u128>"],
+                ["ip_type_and_protocol", "Compact<u8>"],
+            ],
+        },
         "StakeInfo": {
             "type": "struct",
             "type_mapping": [
@@ -226,7 +233,8 @@ class ChainDataType(Enum):
     DelegateInfo = 3
     NeuronInfoLite = 4
     DelegatedInfo = 5
-    StakeInfo = 6
+    IPInfo = 6
+    StakeInfo = 7
 
 
 # Constants
@@ -883,6 +891,68 @@ def from_parameter_dict(
         return cls(**dict(parameter_dict))
 
 
+@dataclass
+class IPInfo:
+    r"""
+    Dataclass for associated IP Info.
+    """
+    ip: str
+    ip_type: int
+    protocol: int
+
+    def encode(self) -> Dict[str, Any]:
+        r"""Returns a dictionary of the IPInfo object that can be encoded."""
+        return {
+            "ip": net.ip_to_int(
+                self.ip
+            ),  # IP type and protocol are encoded together as a u8
+            "ip_type_and_protocol": ((self.ip_type << 4) + self.protocol) & 0xFF,
+        }
+
+    @classmethod
+    def from_vec_u8(cls, vec_u8: List[int]) -> Optional["IPInfo"]:
+        r"""Returns a IPInfo object from a vec_u8."""
+        if len(vec_u8) == 0:
+            return None
+
+        decoded = from_scale_encoding(vec_u8, ChainDataType.IPInfo)
+
+        if decoded is None:
+            return None
+
+        return IPInfo.fix_decoded_values(decoded)
+
+    @classmethod
+    def list_from_vec_u8(cls, vec_u8: List[int]) -> List["IPInfo"]:
+        r"""Returns a list of IPInfo objects from a vec_u8."""
+        decoded = from_scale_encoding(vec_u8, ChainDataType.IPInfo, is_vec=True)
+
+        if decoded is None:
+            return []
+
+        decoded = [IPInfo.fix_decoded_values(d) for d in decoded]
+
+        return decoded
+
+    @classmethod
+    def fix_decoded_values(cls, decoded: Dict) -> "IPInfo":
+        r"""Returns a SubnetInfo object from a decoded IPInfo dictionary."""
+        return IPInfo(
+            ip=bittensor.utils.networking.int_to_ip(decoded["ip"]),
+            ip_type=decoded["ip_type_and_protocol"] >> 4,
+            protocol=decoded["ip_type_and_protocol"] & 0xF,
+        )
+
+    def to_parameter_dict(self) -> "torch.nn.ParameterDict":
+        r"""Returns a torch tensor of the subnet info."""
+        return torch.nn.ParameterDict(self.__dict__)
+
+    @classmethod
+    def from_parameter_dict(cls, parameter_dict: "torch.nn.ParameterDict") -> "IPInfo":
+        r"""Returns a IPInfo object from a torch parameter_dict."""
+        return cls(**dict(parameter_dict))
+
+
 # Senate / Proposal data
 
 

bittensor/subtensor.py

@@ -41,6 +41,7 @@
     AxonInfo,
     ProposalVoteData,
     ProposalCallData,
+    IPInfo,
     custom_rpc_type_registry,
 )
 from .errors import *
@@ -812,6 +813,54 @@ def _do_serve_prometheus(
             else:
                 return True, None
 
+    def _do_associate_ips(
+        self,
+        wallet: "bittensor.wallet",
+        ip_info_list: List[IPInfo],
+        netuid: int,
+        wait_for_inclusion: bool = False,
+        wait_for_finalization: bool = True,
+    ) -> Tuple[bool, Optional[str]]:
+        """
+        Sends an associate IPs extrinsic to the chain.
+
+        Args:
+            wallet (:obj:`bittensor.wallet`): Wallet object.
+            ip_info_list (:obj:`List[IPInfo]`): List of IPInfo objects.
+            netuid (:obj:`int`): Netuid to associate IPs to.
+            wait_for_inclusion (:obj:`bool`): If true, waits for inclusion.
+            wait_for_finalization (:obj:`bool`): If true, waits for finalization.
+
+        Returns:
+            success (:obj:`bool`): True if associate IPs was successful.
+            error (:obj:`Optional[str]`): Error message if associate IPs failed, None otherwise.
+        """
+        with self.substrate as substrate:
+            call = substrate.compose_call(
+                call_module="SubtensorModule",
+                call_function="associate_ips",
+                call_params={
+                    "ip_info_list": [ip_info.encode() for ip_info in ip_info_list],
+                    "netuid": netuid,
+                },
+            )
+            extrinsic = substrate.create_signed_extrinsic(
+                call=call, keypair=wallet.hotkey
+            )
+            response = substrate.submit_extrinsic(
+                extrinsic,
+                wait_for_inclusion=wait_for_inclusion,
+                wait_for_finalization=wait_for_finalization,
+            )
+            if wait_for_inclusion or wait_for_finalization:
+                response.process_events()
+                if response.is_success:
+                    return True, None
+                else:
+                    return False, response.error_message
+            else:
+                return True, None
+
     #################
     #### Staking ####
     #################
@@ -2113,6 +2162,40 @@ def bonds(
 
         return b_map
 
+    def associated_validator_ip_info(
+        self, netuid: int, block: Optional[int] = None
+    ) -> Optional[List[IPInfo]]:
+        """Returns the list of all validator IPs associated with this subnet.
+
+        Args:
+            netuid (int):
+                The network uid of the subnet to query.
+            block ( Optional[int] ):
+                block to sync from, or None for latest block.
+
+        Returns:
+            validator_ip_info (Optional[List[IPInfo]]):
+                List of validator IP info objects for subnet.
+                  or None if no validator IPs are associated with this subnet,
+                  e.g. if the subnet does not exist.
+        """
+        hex_bytes_result = self.query_runtime_api(
+            runtime_api="ValidatorIPRuntimeApi",
+            method="get_associated_validator_ip_info_for_subnet",
+            params=[netuid],
+            block=block,
+        )
+
+        if hex_bytes_result == None:
+            return None
+
+        if hex_bytes_result.startswith("0x"):
+            bytes_result = bytes.fromhex(hex_bytes_result[2:])
+        else:
+            bytes_result = bytes.fromhex(hex_bytes_result)
+
+        return IPInfo.list_from_vec_u8(bytes_result)
+
     def get_subnet_burn_cost(self, block: Optional[int] = None) -> int:
         @retry(delay=2, tries=3, backoff=2, max_delay=4)
         def make_substrate_call_with_retry():