From 308a886993628ee0737d76bebc4f2715b860c07d Mon Sep 17 00:00:00 2001 From: Andreea Popescu Date: Thu, 4 Apr 2024 13:24:00 +0100 Subject: [PATCH] add neuron certificate discoverability --- bittensor/axon.py | 8 +++-- bittensor/chain_data.py | 24 +++++++++++++ bittensor/extrinsics/serving.py | 9 ++++- bittensor/subtensor.py | 63 +++++++++++++++++++++++++++++++-- bittensor/types.py | 4 ++- bittensor/utils/__init__.py | 2 ++ 6 files changed, 103 insertions(+), 7 deletions(-) diff --git a/bittensor/axon.py b/bittensor/axon.py index 34ce9e51f1..05628c04d1 100644 --- a/bittensor/axon.py +++ b/bittensor/axon.py @@ -42,6 +42,7 @@ from starlette.requests import Request from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint from typing import List, Optional, Tuple, Callable, Any, Dict +from .utils import Certificate from bittensor.errors import ( InvalidRequestNameError, @@ -822,7 +823,10 @@ def stop(self) -> "bittensor.axon": return self def serve( - self, netuid: int, subtensor: Optional[bittensor.subtensor] = None + self, + netuid: int, + subtensor: Optional[bittensor.subtensor] = None, + certificate: Optional[Certificate] = None, ) -> "bittensor.axon": """ Serves the Axon on the specified subtensor connection using the configured wallet. This method @@ -848,7 +852,7 @@ def serve( to start receiving and processing requests from other neurons. """ if subtensor is not None and hasattr(subtensor, "serve_axon"): - subtensor.serve_axon(netuid=netuid, axon=self) + subtensor.serve_axon(netuid=netuid, axon=self, certificate=certificate) return self async def default_verify(self, synapse: bittensor.Synapse): diff --git a/bittensor/chain_data.py b/bittensor/chain_data.py index 05b36a8659..1b54c284c9 100644 --- a/bittensor/chain_data.py +++ b/bittensor/chain_data.py @@ -28,6 +28,7 @@ from .utils import networking as net, U16_MAX, U16_NORMALIZED_FLOAT from .utils.balance import Balance +from .utils import Certificate custom_rpc_type_registry = { "types": { @@ -67,6 +68,12 @@ ["total_daily_return", "Compact"], ], }, + "NeuronCertificate": { + "type": "struct", + "type_mapping": [ + ["certificate", "Vec"], + ], + }, "NeuronInfo": { "type": "struct", "type_mapping": [ @@ -286,6 +293,7 @@ class ChainDataType(Enum): StakeInfo = 6 IPInfo = 7 SubnetHyperparameters = 8 + NeuronCertificate = 9 # Constants @@ -514,6 +522,22 @@ def _neuron_dict_to_namespace(neuron_dict) -> "NeuronInfo": return neuron +# Dataclasses for chain data. +@dataclass +class NeuronCertificate: + r""" + Dataclass for neuron certificate. + """ + + certificate: Certificate + + @classmethod + def from_vec_u8(cls, vec_u8: List[int]) -> "NeuronCertificate": + r"""Returns a NeuronCertificate object from a ``vec_u8``.""" + + return from_scale_encoding(vec_u8, ChainDataType.NeuronCertificate) + + @dataclass class NeuronInfoLite: r""" diff --git a/bittensor/extrinsics/serving.py b/bittensor/extrinsics/serving.py index 4bd7277788..ec91439e58 100644 --- a/bittensor/extrinsics/serving.py +++ b/bittensor/extrinsics/serving.py @@ -20,6 +20,8 @@ import bittensor.utils.networking as net from rich.prompt import Confirm from ..errors import MetadataError +from typing import Optional +from ..utils import Certificate def serve_extrinsic( @@ -34,6 +36,7 @@ def serve_extrinsic( wait_for_inclusion: bool = False, wait_for_finalization=True, prompt: bool = False, + certificate: Optional[Certificate] = None, ) -> bool: r"""Subscribes a Bittensor endpoint to the subtensor chain. @@ -44,6 +47,8 @@ def serve_extrinsic( Endpoint host port i.e., ``192.122.31.4``. port (int): Endpoint port number i.e., ``9221``. + certificate (str): + Certificate. protocol (int): An ``int`` representation of the protocol. netuid (int): @@ -75,6 +80,7 @@ def serve_extrinsic( "protocol": protocol, "placeholder1": placeholder1, "placeholder2": placeholder2, + "certificate": certificate, } bittensor.logging.debug("Checking axon ...") neuron = subtensor.get_neuron_for_pubkey_and_subnet( @@ -144,6 +150,7 @@ def serve_axon_extrinsic( wait_for_inclusion: bool = False, wait_for_finalization: bool = True, prompt: bool = False, + certificate: Optional[Certificate] = None, ) -> bool: r"""Serves the axon to the network. @@ -197,6 +204,7 @@ def serve_axon_extrinsic( wait_for_inclusion=wait_for_inclusion, wait_for_finalization=wait_for_finalization, prompt=prompt, + certificate=certificate, ) return serve_success @@ -264,7 +272,6 @@ def publish_metadata( from retry import retry -from typing import Optional def get_metadata(self, netuid: int, hotkey: str, block: Optional[int] = None) -> str: diff --git a/bittensor/subtensor.py b/bittensor/subtensor.py index afbbf938de..99ce28c01c 100644 --- a/bittensor/subtensor.py +++ b/bittensor/subtensor.py @@ -36,6 +36,7 @@ # Local imports. from .btlogging.defines import BITTENSOR_LOGGER_NAME from .chain_data import ( + NeuronCertificate, NeuronInfo, DelegateInfo, PrometheusInfo, @@ -82,7 +83,12 @@ ) from .extrinsics.root import root_register_extrinsic, set_root_weights_extrinsic from .types import AxonServeCallParams, PrometheusServeCallParams -from .utils import U16_NORMALIZED_FLOAT, ss58_to_vec_u8, U64_NORMALIZED_FLOAT +from .utils import ( + U16_NORMALIZED_FLOAT, + ss58_to_vec_u8, + U64_NORMALIZED_FLOAT, + Certificate, +) from .utils.balance import Balance from .utils.registration import POWSolution @@ -1328,6 +1334,7 @@ def serve( wait_for_inclusion: bool = False, wait_for_finalization=True, prompt: bool = False, + certificate: Optional[Certificate] = None, ) -> bool: """ Registers a neuron's serving endpoint on the Bittensor network. This function announces the @@ -1362,6 +1369,7 @@ def serve( placeholder2, wait_for_inclusion, wait_for_finalization, + certificate=certificate, ) def serve_axon( @@ -1371,6 +1379,7 @@ def serve_axon( wait_for_inclusion: bool = False, wait_for_finalization: bool = True, prompt: bool = False, + certificate: Optional[Certificate] = None, ) -> bool: """ Registers an Axon serving endpoint on the Bittensor network for a specific neuron. This function @@ -1391,7 +1400,12 @@ def serve_axon( computing infrastructure, contributing to the collective intelligence of Bittensor. """ return serve_axon_extrinsic( - self, netuid, axon, wait_for_inclusion, wait_for_finalization + self, + netuid, + axon, + wait_for_inclusion, + wait_for_finalization, + certificate=certificate, ) def _do_serve_axon( @@ -1418,12 +1432,18 @@ def _do_serve_axon( enhancing the decentralized computation capabilities of Bittensor. """ + if call_params["certificate"] is None: + del call_params["certificate"] + call_function = "serve_axon" + else: + call_function = "serve_axon_tls" + @retry(delay=2, tries=3, backoff=2, max_delay=4) def make_substrate_call_with_retry(): with self.substrate as substrate: call = substrate.compose_call( call_module="SubtensorModule", - call_function="serve_axon", + call_function=call_function, call_params=call_params, ) extrinsic = substrate.create_signed_extrinsic( @@ -3952,6 +3972,43 @@ def make_substrate_call_with_retry(): return NeuronInfo.from_vec_u8(result) + def get_neuron_certificate( + self, uid: int, netuid: int, block: Optional[int] = None + ) -> Optional[Certificate]: + """ + Retrieves detailed information about a specific neuron identified by its unique identifier (UID) + within a specified subnet (netuid) of the Bittensor network. + Args: + uid (int): The unique identifier of the neuron. + netuid (int): The unique identifier of the subnet. + block (Optional[int], optional): The blockchain block number for the query. + + Returns: + Optional[Certificate]: the certificate of the neuron if found, ``None`` otherwise. + + This function is used for certificate discovery for setting up mutual tls communication between neurons + """ + + @retry(delay=2, tries=3, backoff=2, max_delay=4) + def make_substrate_call_with_retry(): + with self.substrate as substrate: + block_hash = None if block == None else substrate.get_block_hash(block) + params = [netuid, uid] + if block_hash: + params = params + [block_hash] + return substrate.rpc_request( + method="neuronInfo_getNeuronCertificate", + params=params, # custom rpc method + ) + + json_body = make_substrate_call_with_retry() + result = json_body["result"] + + if result in (None, []): + return None + + return NeuronCertificate.from_vec_u8(result) + def neurons(self, netuid: int, block: Optional[int] = None) -> List[NeuronInfo]: """ Retrieves a list of all neurons within a specified subnet of the Bittensor network. This function diff --git a/bittensor/types.py b/bittensor/types.py index 8aa9b7cde4..97f2c4c401 100644 --- a/bittensor/types.py +++ b/bittensor/types.py @@ -15,7 +15,8 @@ # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER # DEALINGS IN THE SOFTWARE. -from typing import TypedDict +from typing import TypedDict, Optional +from .utils import Certificate class AxonServeCallParams(TypedDict): @@ -28,6 +29,7 @@ class AxonServeCallParams(TypedDict): port: int ip_type: int netuid: int + certificate: Optional[Certificate] class PrometheusServeCallParams(TypedDict): diff --git a/bittensor/utils/__init__.py b/bittensor/utils/__init__.py index fbc3aa7d47..b189c719a3 100644 --- a/bittensor/utils/__init__.py +++ b/bittensor/utils/__init__.py @@ -30,6 +30,8 @@ U16_MAX = 65535 U64_MAX = 18446744073709551615 +Certificate = str + def ss58_to_vec_u8(ss58_address: str) -> List[int]: ss58_bytes: bytes = bittensor.utils.ss58_address_to_bytes(ss58_address)