Skip to content

Commit

Permalink
fix an issue with merging axon.py from release v1.8.0
Browse files Browse the repository at this point in the history
  • Loading branch information
@mo0haned
mo0haned committed Feb 5, 2025
1 parent 2ec32dc commit 5d53013
Showing 1 changed file with 0 additions and 94 deletions.
94 changes: 0 additions & 94 deletions compute/axon.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,18 +18,13 @@
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.

# Standard library
# Standard library
import copy
import time
import time
import uuid
from inspect import Signature
from typing import TYPE_CHECKING, Callable, Optional
from inspect import Signature
from typing import TYPE_CHECKING, Callable, Optional

# Third-party
# Third-party
import uvicorn
from fastapi import FastAPI, APIRouter
Expand All @@ -47,19 +42,6 @@
from bittensor.utils import format_error_message, networking as net
from bittensor.utils import ( format_error_message, networking as net, unlock_key, Certificate )

# Local
# Bittensor
import bittensor
from bittensor.core.axon import Axon as axon
from bittensor.core.axon import FastAPIThreadedServer, AxonMiddleware
from bittensor.core.subtensor import Subtensor as subtensor
from bittensor.core.config import Config
from bittensor.core.threadpool import PriorityThreadPoolExecutor
from bittensor.core.extrinsics.serving import do_serve_axon
from bittensor.utils.btlogging import logging
from bittensor.utils import format_error_message, networking as net
from bittensor.utils import ( format_error_message, networking as net, unlock_key, Certificate )

# Local
from compute import __version_as_int__
from compute.utils.version import get_local_version
Expand All @@ -83,16 +65,9 @@ def custom_serve_extrinsic(
wait_for_inclusion: bool = False,
wait_for_finalization: bool = True,
certificate: Certificate | None = None,
<<<<<<< HEAD
wait_for_finalization: bool = True,
certificate: Certificate | None = None,
=======
>>>>>>> 1df75b2 (fix: update type hint for certificate parameter in custom_serve_extrinsic to support union type)
) -> bool:
"""Subscribes a Bittensor endpoint to the subtensor chain.
"""Subscribes a Bittensor endpoint to the subtensor chain.

Args:
subtensor (bittensor.core.subtensor.Subtensor): Subtensor instance object.
wallet (bittensor_wallet.Wallet): Bittensor wallet object.
Expand All @@ -105,30 +80,14 @@ def custom_serve_extrinsic(
wait_for_inclusion (bool): If set, waits for the extrinsic to enter a block before returning ``true``, or returns ``false`` if the extrinsic fails to enter the block within the timeout.
wait_for_finalization (bool): If set, waits for the extrinsic to be finalized on the chain before returning ``true``, or returns ``false`` if the extrinsic fails to be finalized within the timeout.
certificate (Certificate | None): An optional certificate object that can be used for secure communication.
subtensor (bittensor.core.subtensor.Subtensor): Subtensor instance object.
wallet (bittensor_wallet.Wallet): Bittensor wallet object.
ip (str): Endpoint host port i.e., ``192.122.31.4``.
port (int): Endpoint port number i.e., ``9221``.
protocol (int): An ``int`` representation of the protocol.
netuid (int): The network uid to serve on.
placeholder1 (int): A placeholder for future use.
placeholder2 (int): A placeholder for future use.
wait_for_inclusion (bool): If set, waits for the extrinsic to enter a block before returning ``true``, or returns ``false`` if the extrinsic fails to enter the block within the timeout.
wait_for_finalization (bool): If set, waits for the extrinsic to be finalized on the chain before returning ``true``, or returns ``false`` if the extrinsic fails to be finalized within the timeout.
certificate (Certificate | None): An optional certificate object that can be used for secure communication.
Returns:
success (bool): Flag is ``true`` if extrinsic was finalized or uncluded in the block. If we did not wait for finalization / inclusion, the response is ``true``.
success (bool): Flag is ``true`` if extrinsic was finalized or uncluded in the block. If we did not wait for finalization / inclusion, the response is ``true``.
"""
# Decrypt hotkey
if not (unlock := unlock_key(wallet, "hotkey")).success:
logging.error(unlock.message)
return False
if not (unlock := unlock_key(wallet, "hotkey")).success:
logging.error(unlock.message)
return False
params: "AxonServeCallParams" = {
"version": __version_as_int__,
"version": __version_as_int__,
"ip": net.ip_to_int(ip),
"port": port,
Expand All @@ -140,15 +99,12 @@ def custom_serve_extrinsic(
"placeholder1": placeholder1,
"placeholder2": placeholder2,
"certificate": certificate,
"certificate": certificate,
}
logging.debug("Checking axon ...")
logging.debug("Checking axon ...")
neuron = subtensor.get_neuron_for_pubkey_and_subnet(
wallet.hotkey.ss58_address, netuid=netuid
)
neuron_up_to_date = not neuron.is_null and params == {
"version": __version_as_int__,
"version": __version_as_int__,
"ip": net.ip_to_int(neuron.axon_info.ip),
"port": neuron.axon_info.port,
Expand All @@ -164,19 +120,14 @@ def custom_serve_extrinsic(
output["coldkey"] = wallet.coldkeypub.ss58_address
output["hotkey"] = wallet.hotkey.ss58_address
if neuron_up_to_date:
logging.debug(
logging.debug(
f"Axon already served on: AxonInfo({wallet.hotkey.ss58_address},{ip}:{port}) "
)
return True

logging.debug(
f"Serving axon with: AxonInfo({wallet.hotkey.ss58_address},{ip}:{port}) -> {subtensor.network}:{netuid}"
logging.debug(
f"Serving axon with: AxonInfo({wallet.hotkey.ss58_address},{ip}:{port}) -> {subtensor.network}:{netuid}"
)
success, error_message = do_serve_axon(
self=subtensor,
success, error_message = do_serve_axon(
self=subtensor,
wallet=wallet,
Expand All @@ -186,24 +137,18 @@ def custom_serve_extrinsic(
)

if wait_for_inclusion or wait_for_finalization:
if success is True:
logging.debug(
f"Axon served with: AxonInfo({wallet.hotkey.ss58_address},{ip}:{port}) on {subtensor.network}:{netuid} "
)
if success is True:
logging.debug(
f"Axon served with: AxonInfo({wallet.hotkey.ss58_address},{ip}:{port}) on {subtensor.network}:{netuid} "
)
return True
else:
logging.error(f"Failed: {format_error_message(error_message)}")
logging.error(f"Failed: {format_error_message(error_message)}")
return False
else:
return True

bittensor.core.extrinsics.serving.serve_extrinsic = custom_serve_extrinsic
bittensor.core.extrinsics.serving.serve_extrinsic = custom_serve_extrinsic

class ComputeSubnetSubtensor(subtensor):
def __init__(
Expand All @@ -226,16 +171,11 @@ def __init__(
self,
wallet: Optional["Wallet"] = None,
config: Optional["Config"] = None,
wallet: Optional["Wallet"] = None,
config: Optional["Config"] = None,
port: Optional[int] = None,
ip: Optional[str] = None,
external_ip: Optional[str] = None,
external_port: Optional[int] = None,
max_workers: Optional[int] = None,
):
"""Creates a new bittensor.Axon object from passed arguments.

):
"""Creates a new bittensor.Axon object from passed arguments.
Expand All @@ -247,13 +187,6 @@ def __init__(
external_ip (:type:`Optional[str]`): The external ip of the server to broadcast to the network.
external_port (:type:`Optional[int]`): The external port of the server to broadcast to the network.
max_workers (:type:`Optional[int]`): Used to create the threadpool if not passed, specifies the number of active threads servicing requests.
config (:obj:`Optional[bittensor.core.config.Config]`): bittensor.Axon.config()
wallet (:obj:`Optional[bittensor_wallet.Wallet]`): bittensor wallet with hotkey and coldkeypub.
port (:type:`Optional[int]`): Binding port.
ip (:type:`Optional[str]`): Binding ip.
external_ip (:type:`Optional[str]`): The external ip of the server to broadcast to the network.
external_port (:type:`Optional[int]`): The external port of the server to broadcast to the network.
max_workers (:type:`Optional[int]`): Used to create the threadpool if not passed, specifies the number of active threads servicing requests.
"""

# Build and check config.
Expand All @@ -265,48 +198,30 @@ def __init__(
config.axon.external_ip = external_ip or config.axon.external_ip
config.axon.external_port = external_port or config.axon.external_port
config.axon.max_workers = max_workers or config.axon.max_workers
config.axon.ip = ip or config.axon.ip
config.axon.port = port or config.axon.port
config.axon.external_ip = external_ip or config.axon.external_ip
config.axon.external_port = external_port or config.axon.external_port
config.axon.max_workers = max_workers or config.axon.max_workers
axon.check_config(config)
self.config = config # type: ignore
self.config = config # type: ignore

# Get wallet or use default.
self.wallet = wallet or Wallet(config=self.config)
self.wallet = wallet or Wallet(config=self.config)

# Build axon objects.
self.uuid = str(uuid.uuid1())
self.ip = self.config.axon.ip # type: ignore
self.port = self.config.axon.port # type: ignore
self.ip = self.config.axon.ip # type: ignore
self.port = self.config.axon.port # type: ignore
self.external_ip = (
self.config.axon.external_ip # type: ignore
if self.config.axon.external_ip is not None # type: ignore
else net.get_external_ip()
self.config.axon.external_ip # type: ignore
if self.config.axon.external_ip is not None # type: ignore
else net.get_external_ip()
)
self.external_port = (
self.config.axon.external_port # type: ignore
if self.config.axon.external_port is not None # type: ignore
else self.config.axon.port # type: ignore
self.config.axon.external_port # type: ignore
if self.config.axon.external_port is not None # type: ignore
else self.config.axon.port # type: ignore
)
self.full_address = str(self.config.axon.ip) + ":" + str(self.config.axon.port) # type: ignore
self.full_address = str(self.config.axon.ip) + ":" + str(self.config.axon.port) # type: ignore
self.started = False

# Build middleware
self.thread_pool = PriorityThreadPoolExecutor(
max_workers=self.config.axon.max_workers # type: ignore
self.thread_pool = PriorityThreadPoolExecutor(
max_workers=self.config.axon.max_workers # type: ignore
)
Expand All @@ -319,17 +234,10 @@ def __init__(
self.forward_fns: dict[str, Callable | None] = {}
self.verify_fns: dict[str, Callable | None] = {}

self.forward_class_types: dict[str, list[Signature]] = {}
self.blacklist_fns: dict[str, Callable | None] = {}
self.priority_fns: dict[str, Callable | None] = {}
self.forward_fns: dict[str, Callable | None] = {}
self.verify_fns: dict[str, Callable | None] = {}


# Instantiate FastAPI
self.app = FastAPI()
log_level = "trace" if logging.__trace_on__ else "critical"
log_level = "trace" if logging.__trace_on__ else "critical"
self.fast_config = uvicorn.Config(
self.app, host="0.0.0.0", port=self.config.axon.port, log_level=log_level
)
Expand All @@ -340,8 +248,6 @@ def __init__(
# Build ourselves as the middleware.
self.middleware_cls = ComputeSubnetAxonMiddleware
self.app.add_middleware(self.middleware_cls, axon=self)
self.middleware_cls = ComputeSubnetAxonMiddleware
self.app.add_middleware(self.middleware_cls, axon=self)

# Attach default forward.
def ping(r: bittensor.Synapse) -> bittensor.Synapse:
Expand Down

0 comments on commit 5d53013

Please sign in to comment.