diff --git a/bittensor/core/axon.py b/bittensor/core/axon.py index ce967acb80..378315f77d 100644 --- a/bittensor/core/axon.py +++ b/bittensor/core/axon.py @@ -59,7 +59,7 @@ from bittensor.core.stream import StreamingSynapse from bittensor.core.synapse import Synapse, TerminalInfo from bittensor.core.threadpool import PriorityThreadPoolExecutor -from bittensor.utils import networking +from bittensor.utils import networking, Certificate from bittensor.utils.axon_utils import allowed_nonce_window_ns, calculate_diff_seconds from bittensor.utils.btlogging import logging @@ -807,7 +807,12 @@ def stop(self) -> "Axon": self.started = False return self - def serve(self, netuid: int, subtensor: Optional["Subtensor"] = None) -> "Axon": + def serve( + self, + netuid: int, + subtensor: Optional["Subtensor"] = None, + certificate: Optional[Certificate] = None, + ) -> "Axon": """ Serves the Axon on the specified subtensor connection using the configured wallet. This method registers the Axon with a specific subnet within the Bittensor network, identified by the ``netuid``. @@ -832,7 +837,7 @@ def serve(self, netuid: int, subtensor: Optional["Subtensor"] = None) -> "Axon": 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: "Synapse"): diff --git a/bittensor/core/chain_data/neuron_certificate.py b/bittensor/core/chain_data/neuron_certificate.py new file mode 100644 index 0000000000..13068b3661 --- /dev/null +++ b/bittensor/core/chain_data/neuron_certificate.py @@ -0,0 +1,19 @@ +from dataclasses import dataclass + +from bittensor.core.chain_data.utils import from_scale_encoding, ChainDataType +from bittensor.utils import Certificate + + +# 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) diff --git a/bittensor/core/chain_data/utils.py b/bittensor/core/chain_data/utils.py index 9c21c9d22e..d4f86208ab 100644 --- a/bittensor/core/chain_data/utils.py +++ b/bittensor/core/chain_data/utils.py @@ -22,6 +22,7 @@ class ChainDataType(Enum): SubnetHyperparameters = 8 ScheduledColdkeySwapInfo = 9 AccountId = 10 + NeuronCertificate = 11 def from_scale_encoding( @@ -178,6 +179,12 @@ def from_scale_encoding_using_type_string( ["pruning_score", "Compact"], ], }, + "NeuronCertificate": { + "type": "struct", + "type_mapping": [ + ["certificate", "Vec"], + ], + }, "axon_info": { "type": "struct", "type_mapping": [ @@ -256,6 +263,7 @@ def from_scale_encoding_using_type_string( ["arbitration_block", "Compact"], ], }, + } } diff --git a/bittensor/core/extrinsics/serving.py b/bittensor/core/extrinsics/serving.py index 02630ff91c..f161fffdb4 100644 --- a/bittensor/core/extrinsics/serving.py +++ b/bittensor/core/extrinsics/serving.py @@ -20,7 +20,12 @@ from bittensor.core.errors import MetadataError from bittensor.core.extrinsics.utils import submit_extrinsic from bittensor.core.settings import version_as_int -from bittensor.utils import format_error_message, networking as net, unlock_key +from bittensor.utils import ( + format_error_message, + networking as net, + unlock_key, + Certificate, +) from bittensor.utils.btlogging import logging from bittensor.utils.networking import ensure_connected @@ -40,6 +45,7 @@ def do_serve_axon( call_params: "AxonServeCallParams", wait_for_inclusion: bool = False, wait_for_finalization: bool = True, + certificate: Optional[Certificate] = None, ) -> tuple[bool, Optional[dict]]: """ Internal method to submit a serve axon transaction to the Bittensor blockchain. This method creates and submits a transaction, enabling a neuron's ``Axon`` to serve requests on the network. @@ -57,9 +63,15 @@ def do_serve_axon( This function is crucial for initializing and announcing a neuron's ``Axon`` service on the network, 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" + 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(call=call, keypair=wallet.hotkey) @@ -68,6 +80,7 @@ def do_serve_axon( extrinsic=extrinsic, wait_for_inclusion=wait_for_inclusion, wait_for_finalization=wait_for_finalization, + certificate=certificate, ) if wait_for_inclusion or wait_for_finalization: response.process_events() @@ -90,6 +103,7 @@ def serve_extrinsic( placeholder2: int = 0, wait_for_inclusion: bool = False, wait_for_finalization=True, + certificate: Optional[Certificate] = None, ) -> bool: """Subscribes a Bittensor endpoint to the subtensor chain. @@ -124,6 +138,7 @@ def serve_extrinsic( "protocol": protocol, "placeholder1": placeholder1, "placeholder2": placeholder2, + "certificate": certificate, } logging.debug("Checking axon ...") neuron = subtensor.get_neuron_for_pubkey_and_subnet( @@ -159,6 +174,7 @@ def serve_extrinsic( call_params=params, wait_for_finalization=wait_for_finalization, wait_for_inclusion=wait_for_inclusion, + certificate=certificate, ) if wait_for_inclusion or wait_for_finalization: @@ -182,6 +198,7 @@ def serve_axon_extrinsic( axon: "Axon", wait_for_inclusion: bool = False, wait_for_finalization: bool = True, + certificate: Optional[Certificate] = None, ) -> bool: """Serves the axon to the network. @@ -224,6 +241,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/core/subtensor.py b/bittensor/core/subtensor.py index efaff369a4..5f3f09984f 100644 --- a/bittensor/core/subtensor.py +++ b/bittensor/core/subtensor.py @@ -43,6 +43,7 @@ DelegateInfo, NeuronInfo, NeuronInfoLite, + NeuronCertificate, PrometheusInfo, SubnetHyperparameters, SubnetInfo, @@ -77,6 +78,7 @@ ss58_to_vec_u8, u16_normalized_float, hex_to_bytes, + Certificate, ) from bittensor.utils.balance import Balance from bittensor.utils.btlogging import logging @@ -1013,6 +1015,7 @@ def serve_axon( axon: "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 is used to set up the Axon, a key component of a neuron that handles incoming queries and data processing tasks. @@ -1029,7 +1032,7 @@ def serve_axon( By registering an Axon, the neuron becomes an active part of the network's distributed 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 ) # metagraph @@ -1170,6 +1173,37 @@ def get_neuron_for_pubkey_and_subnet( block=block, ) + def get_neuron_certificate( + self, hotkey: str, 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: + hotkey (str): The hotkey to query. + 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 + """ + + certificate = self.query_module( + module="SubtensorModule", + name="NeuronCertificates", + block=block, + params=[netuid, hotkey], + ) + if not hasattr(certificate, "serialize"): + return None + certificate = certificate.serialize() + if not certificate: + return None + return certificate.get("certificate", None) + @networking.ensure_connected def neuron_for_uid( self, uid: Optional[int], netuid: int, block: Optional[int] = None diff --git a/bittensor/core/types.py b/bittensor/core/types.py index 9fd2b4d052..577df5b6ba 100644 --- a/bittensor/core/types.py +++ b/bittensor/core/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 bittensor.utils import Certificate class AxonServeCallParams(TypedDict): @@ -26,6 +27,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 7a42dff0cf..cdcadf465f 100644 --- a/bittensor/utils/__init__.py +++ b/bittensor/utils/__init__.py @@ -40,6 +40,8 @@ U16_MAX = 65535 U64_MAX = 18446744073709551615 +Certificate = str + UnlockStatus = namedtuple("UnlockStatus", ["success", "message"]) diff --git a/bittensor/utils/mock/subtensor_mock.py b/bittensor/utils/mock/subtensor_mock.py index 817be08434..4b67b862d2 100644 --- a/bittensor/utils/mock/subtensor_mock.py +++ b/bittensor/utils/mock/subtensor_mock.py @@ -47,6 +47,7 @@ class AxonServeCallParams(TypedDict): port: int ip_type: int netuid: int + certificate: Optional[str] class PrometheusServeCallParams(TypedDict):