Refactor AsyncSubtensor and Dendrite to offload signing to ProcessPoolExecutor (Fixes #3211)

Author: Dairus01

bittensor/core/async_subtensor.py

@@ -6,6 +6,8 @@
 
 import asyncstdlib as a
 import scalecodec
+from concurrent.futures import ProcessPoolExecutor
+from functools import partial
 from async_substrate_interface import AsyncSubstrateInterface
 from async_substrate_interface.substrate_addons import RetryAsyncSubstrate
 from async_substrate_interface.utils.storage import StorageKey
@@ -163,6 +165,16 @@
     tick_to_price,
 )
 
+
+
+
+def sign_extrinsic_worker(private_key: str, payload: bytes) -> str:
+    from bittensor_wallet import Keypair
+    # We use the private key to recreate the keypair for signing
+    kp = Keypair(private_key=private_key)
+    return f"0x{kp.sign(payload).hex()}"
+
+
 if TYPE_CHECKING:
     from async_substrate_interface import AsyncQueryMapResult
     from async_substrate_interface.types import ScaleObj
@@ -253,6 +265,8 @@ def __init__(
             logging.info(
                 f"Connected to {self.network} network and {self.chain_endpoint}."
             )
+        
+        self._executor = ProcessPoolExecutor(max_workers=1)
 
     async def close(self):
         """Closes the connection to the blockchain.
@@ -279,6 +293,10 @@ async def close(self):
         if self.substrate:
             await self.substrate.close()
 
+        if self._executor:
+            self._executor.shutdown(wait=True)
+            self._executor = None
+
     async def initialize(self):
         """Establishes connection to the blockchain.
 
@@ -328,7 +346,12 @@ async def __aenter__(self):
         return await self.initialize()
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
-        await self.substrate.close()
+        if self.substrate:
+            await self.substrate.close()
+        
+        if self._executor:
+            self._executor.shutdown(wait=True)
+            self._executor = None
 
     # Helpers ==========================================================================================================
 
@@ -5961,6 +5984,51 @@ async def compose_call(
             block_hash=block_hash,
         )
 
+    async def create_signed_extrinsic(
+        self,
+        call: "GenericCall",
+        keypair: "Keypair",
+        nonce: Optional[int] = None,
+        era: Optional[dict] = None,
+        tip: int = 0,
+        tip_asset_id: Optional[int] = None,
+    ):
+        """
+        Creates a signed extrinsic, offloading the signing process to a separate process to avoid blocking the event loop.
+        """
+        if nonce is None:
+            nonce = await self.substrate.get_account_next_index(keypair.ss58_address)
+
+        # Generate the signature payload.
+        # This uses the underlying substrate interface to prepare the bytes to be signed.
+        signature_payload = await self.substrate.generate_signature_payload(
+            call=call,
+            era=era,
+            nonce=nonce,
+            tip=tip,
+            tip_asset_id=tip_asset_id
+        )
+
+        # Offload the signing of the payload to an executor.
+        loop = asyncio.get_running_loop()
+        # signature_payload.data should contain the bytes to sign.
+        signature = await loop.run_in_executor(
+            self._executor,
+            partial(sign_extrinsic_worker, keypair.private_key, signature_payload.data)
+        )
+
+        # Create the extrinsic with the attached signature.
+        extrinsic = await self.substrate.create_extrinsic(
+            call=call,
+            nonce=nonce,
+            era=era,
+            tip=tip,
+            tip_asset_id=tip_asset_id,
+            signature=signature,
+            keypair=keypair,  # Keypair may be needed for address/type info
+        )
+        return extrinsic
+
     async def sign_and_send_extrinsic(
         self,
         call: "GenericCall",
@@ -6029,7 +6097,7 @@ async def sign_and_send_extrinsic(
         if period is not None:
             extrinsic_data["era"] = {"period": period}
 
-        extrinsic_response.extrinsic = await self.substrate.create_signed_extrinsic(
+        extrinsic_response.extrinsic = await self.create_signed_extrinsic(
             **extrinsic_data
         )
         try:

bittensor/core/dendrite.py

@@ -4,6 +4,8 @@
 import time
 import uuid
 import warnings
+from concurrent.futures import ProcessPoolExecutor
+from functools import partial
 from typing import Any, AsyncGenerator, Optional, Union, Type
 
 import aiohttp
@@ -39,6 +41,19 @@ def event_loop_is_running():
         return False
 
 
+def sign_message_worker(private_key_hex: str, message: str) -> str:
+    from bittensor_wallet import Keypair
+    
+    # Re-instantiate the keypair from the private key
+    # ss58_address is required but can be dummy if only signing?
+    # Actually bittensor_wallet Keypair usually takes ss58_address. 
+    # If private_key is provided, it might derive it.
+    # Let's try passing None or empty string if allowed, or check Keypair usage.
+    # Safe bet: pass private key. Keypair usually derives public from private.
+    kp = Keypair(private_key=private_key_hex)
+    return f"0x{kp.sign(message).hex()}"
+
+
 class DendriteMixin:
     """
     The Dendrite class represents the abstracted implementation of a network client module.
@@ -121,6 +136,7 @@ def __init__(self, wallet: Optional[Union["Wallet", "Keypair"]] = None):
         self.synapse_history: list = []
 
         self._session: Optional[aiohttp.ClientSession] = None
+        self._executor = ProcessPoolExecutor(max_workers=1)
 
     @property
     async def session(self) -> aiohttp.ClientSession:
@@ -192,6 +208,10 @@ def close_session(self, using_new_loop: bool = False):
             if using_new_loop:
                 loop.close()
             self._session = None
+        
+        if self._executor:
+            self._executor.shutdown(wait=True)
+            self._executor = None
 
     async def aclose_session(self):
         """
@@ -217,6 +237,10 @@ async def aclose_session(self):
         if self._session:
             await self._session.close()
             self._session = None
+            
+        if self._executor:
+            self._executor.shutdown(wait=True)
+            self._executor = None
 
     def _get_endpoint_url(self, target_axon, request_name):
         """
@@ -568,7 +592,8 @@ async def call(
         url = self._get_endpoint_url(target_axon, request_name=request_name)
 
         # Preprocess synapse for making a request
-        synapse = self.preprocess_synapse_for_request(target_axon, synapse, timeout)
+        synapse = await self.preprocess_synapse_for_request(target_axon, synapse, timeout)
+
 
         try:
             # Log outgoing request
@@ -643,7 +668,8 @@ async def call_stream(
         url = f"http://{endpoint}/{request_name}"
 
         # Preprocess synapse for making a request
-        synapse = self.preprocess_synapse_for_request(target_axon, synapse, timeout)  # type: ignore
+        synapse = await self.preprocess_synapse_for_request(target_axon, synapse, timeout)  # type: ignore
+
 
         try:
             # Log outgoing request
@@ -682,7 +708,7 @@ async def call_stream(
             else:
                 yield synapse
 
-    def preprocess_synapse_for_request(
+    async def preprocess_synapse_for_request(
         self,
         target_axon_info: "AxonInfo",
         synapse: "Synapse",
@@ -719,7 +745,12 @@ def preprocess_synapse_for_request(
 
         # Sign the request using the dendrite, axon info, and the synapse body hash
         message = f"{synapse.dendrite.nonce}.{synapse.dendrite.hotkey}.{synapse.axon.hotkey}.{synapse.dendrite.uuid}.{synapse.body_hash}"
-        synapse.dendrite.signature = f"0x{self.keypair.sign(message).hex()}"
+        
+        loop = asyncio.get_running_loop()
+        synapse.dendrite.signature = await loop.run_in_executor(
+            self._executor,
+            partial(sign_message_worker, self.keypair.private_key, message)
+        )
 
         return synapse