diff --git a/bittensor/axon.py b/bittensor/axon.py index 8cefadfe61..e042220ede 100644 --- a/bittensor/axon.py +++ b/bittensor/axon.py @@ -34,6 +34,7 @@ import warnings from inspect import signature, Signature, Parameter from typing import List, Optional, Tuple, Callable, Any, Dict, Awaitable +from .utils import Certificate import uvicorn from fastapi import APIRouter, Depends, FastAPI @@ -832,7 +833,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 @@ -858,7 +862,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 029cb29829..234fa771d9 100644 --- a/bittensor/chain_data.py +++ b/bittensor/chain_data.py @@ -34,6 +34,7 @@ from .utils import networking as net, RAOPERTAO, U16_NORMALIZED_FLOAT from .utils.balance import Balance from .utils.registration import torch, use_torch +from .utils import Certificate custom_rpc_type_registry = { "types": { @@ -73,6 +74,12 @@ ["total_daily_return", "Compact"], ], }, + "NeuronCertificate": { + "type": "struct", + "type_mapping": [ + ["certificate", "Vec"], + ], + }, "NeuronInfo": { "type": "struct", "type_mapping": [ @@ -333,6 +340,7 @@ class ChainDataType(Enum): IPInfo = 7 SubnetHyperparameters = 8 ScheduledColdkeySwapInfo = 9 + NeuronCertificate = 10 def from_scale_encoding( @@ -540,6 +548,21 @@ def from_weights_bonds_and_neuron_lite( return cls(**n_dict) +# 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: """Dataclass for neuron metadata, but without the weights and bonds.""" diff --git a/bittensor/extrinsics/serving.py b/bittensor/extrinsics/serving.py index bba5367de1..02cc384bcd 100644 --- a/bittensor/extrinsics/serving.py +++ b/bittensor/extrinsics/serving.py @@ -26,6 +26,8 @@ import bittensor.utils.networking as net from bittensor.utils import format_error_message from ..errors import MetadataError +from typing import Optional +from ..utils import Certificate def serve_extrinsic( @@ -40,6 +42,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. @@ -50,6 +53,8 @@ def serve_extrinsic( Endpoint host port i.e., ``192.122.31.4``. port (int): Endpoint port number i.e., ``9221``. + certificate (str): + TLS Certificate protocol (int): An ``int`` representation of the protocol. netuid (int): @@ -81,6 +86,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( @@ -148,6 +154,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. @@ -200,6 +207,7 @@ def serve_axon_extrinsic( protocol=4, wait_for_inclusion=wait_for_inclusion, wait_for_finalization=wait_for_finalization, + certificate=certificate, ) return serve_success diff --git a/bittensor/subtensor.py b/bittensor/subtensor.py index 05ee9bb2c8..9b3b08ee6e 100644 --- a/bittensor/subtensor.py +++ b/bittensor/subtensor.py @@ -45,6 +45,7 @@ from bittensor.utils import torch, weight_utils, format_error_message from .chain_data import ( DelegateInfoLite, + NeuronCertificate, NeuronInfo, DelegateInfo, PrometheusInfo, @@ -114,6 +115,7 @@ ss58_to_vec_u8, U64_NORMALIZED_FLOAT, networking, + Certificate, ) from .utils.balance import Balance from .utils.registration import POWSolution @@ -1813,6 +1815,7 @@ def serve( placeholder2: int = 0, wait_for_inclusion: bool = False, wait_for_finalization=True, + certificate: Optional[Certificate] = None, ) -> bool: """ Registers a neuron's serving endpoint on the Bittensor network. This function announces the @@ -1849,6 +1852,7 @@ def serve( placeholder2, wait_for_inclusion, wait_for_finalization, + certificate=certificate, ) def serve_axon( @@ -1857,6 +1861,7 @@ def serve_axon( axon: "bittensor.axon", wait_for_inclusion: bool = False, wait_for_finalization: bool = True, + certificate: Optional[Certificate] = None, ) -> bool: """ Registers an Axon serving endpoint on the Bittensor network for a specific neuron. This function @@ -1876,7 +1881,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( @@ -1903,11 +1913,17 @@ 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=1, tries=3, backoff=2, max_delay=4, logger=_logger) def make_substrate_call_with_retry(): call = self.substrate.compose_call( call_module="SubtensorModule", - call_function="serve_axon", + call_function=call_function, call_params=call_params, ) extrinsic = self.substrate.create_signed_extrinsic( @@ -5022,6 +5038,41 @@ 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 the TLS certificate for 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, logger=_logger) + def make_substrate_call_with_retry(): + block_hash = None if block == None else self.substrate.get_block_hash(block) + params = [netuid, uid] + if block_hash: + params = params + [block_hash] + return self.substrate.rpc_request( + method="neuronInfo_getNeuronCertificate", + params=params, # custom rpc method + ) + + json_body = make_substrate_call_with_retry() + + if not (result := json_body.get("result", 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 700a656131..e48b49a857 100644 --- a/bittensor/utils/__init__.py +++ b/bittensor/utils/__init__.py @@ -31,6 +31,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)