Improve submit_extrinsic util

bittensor/core/extrinsics/commit_weights.py

@@ -71,7 +71,7 @@ def do_commit_weights(
         keypair=wallet.hotkey,
     )
     response = submit_extrinsic(
-        substrate=self.substrate,
+        subtensor=self,
         extrinsic=extrinsic,
         wait_for_inclusion=wait_for_inclusion,
         wait_for_finalization=wait_for_finalization,
@@ -184,7 +184,7 @@ def do_reveal_weights(
         keypair=wallet.hotkey,
     )
     response = submit_extrinsic(
-        substrate=self.substrate,
+        subtensor=self,
         extrinsic=extrinsic,
         wait_for_inclusion=wait_for_inclusion,
         wait_for_finalization=wait_for_finalization,

bittensor/core/extrinsics/registration.py

@@ -66,7 +66,7 @@ def _do_pow_register(
     )
     extrinsic = self.substrate.create_signed_extrinsic(call=call, keypair=wallet.hotkey)
     response = submit_extrinsic(
-        substrate=self.substrate,
+        self,
         extrinsic=extrinsic,
         wait_for_inclusion=wait_for_inclusion,
         wait_for_finalization=wait_for_finalization,
@@ -297,7 +297,7 @@ def _do_burned_register(
         call=call, keypair=wallet.coldkey
     )
     response = submit_extrinsic(
-        self.substrate,
+        self,
         extrinsic,
         wait_for_inclusion=wait_for_inclusion,
         wait_for_finalization=wait_for_finalization,

bittensor/core/extrinsics/root.py

@@ -33,7 +33,7 @@ def _do_root_register(
         call=call, keypair=wallet.coldkey
     )
     response = submit_extrinsic(
-        self.substrate,
+        self,
         extrinsic,
         wait_for_inclusion=wait_for_inclusion,
         wait_for_finalization=wait_for_finalization,
@@ -159,7 +159,7 @@ def _do_set_root_weights(
         era={"period": period},
     )
     response = submit_extrinsic(
-        self.substrate,
+        self,
         extrinsic,
         wait_for_inclusion=wait_for_inclusion,
         wait_for_finalization=wait_for_finalization,

bittensor/core/extrinsics/serving.py

@@ -75,7 +75,7 @@ def do_serve_axon(
     )
     extrinsic = self.substrate.create_signed_extrinsic(call=call, keypair=wallet.hotkey)
     response = submit_extrinsic(
-        substrate=self.substrate,
+        self,
         extrinsic=extrinsic,
         wait_for_inclusion=wait_for_inclusion,
         wait_for_finalization=wait_for_finalization,
@@ -289,7 +289,7 @@ def publish_metadata(
 
         extrinsic = substrate.create_signed_extrinsic(call=call, keypair=wallet.hotkey)
         response = submit_extrinsic(
-            substrate,
+            self,
             extrinsic,
             wait_for_inclusion=wait_for_inclusion,
             wait_for_finalization=wait_for_finalization,

bittensor/core/extrinsics/set_weights.py

@@ -83,7 +83,7 @@ def do_set_weights(
         era={"period": period},
     )
     response = submit_extrinsic(
-        substrate=self.substrate,
+        self,
         extrinsic=extrinsic,
         wait_for_inclusion=wait_for_inclusion,
         wait_for_finalization=wait_for_finalization,

bittensor/core/extrinsics/transfer.py

@@ -70,7 +70,7 @@ def do_transfer(
         call=call, keypair=wallet.coldkey
     )
     response = submit_extrinsic(
-        substrate=self.substrate,
+        self,
         extrinsic=extrinsic,
         wait_for_inclusion=wait_for_inclusion,
         wait_for_finalization=wait_for_finalization,

bittensor/core/extrinsics/utils.py

@@ -11,7 +11,8 @@
 from bittensor.utils import format_error_message
 
 if TYPE_CHECKING:
-    from substrateinterface import SubstrateInterface, ExtrinsicReceipt
+    from bittensor.core.subtensor import Subtensor
+    from substrateinterface import ExtrinsicReceipt
     from scalecodec.types import GenericExtrinsic
 
 try:
@@ -26,7 +27,7 @@
 
 
 def submit_extrinsic(
-    substrate: "SubstrateInterface",
+    subtensor: "Subtensor",
     extrinsic: "GenericExtrinsic",
     wait_for_inclusion: bool,
     wait_for_finalization: bool,
@@ -39,7 +40,7 @@ def submit_extrinsic(
     it logs the error and re-raises the exception.
 
     Args:
-        substrate (substrateinterface.SubstrateInterface): The substrate interface instance used to interact with the blockchain.
+        subtensor: The Subtensor instance used to interact with the blockchain.
         extrinsic (scalecodec.types.GenericExtrinsic): The extrinsic to be submitted to the blockchain.
         wait_for_inclusion (bool): Whether to wait for the extrinsic to be included in a block.
         wait_for_finalization (bool): Whether to wait for the extrinsic to be finalized on the blockchain.
@@ -51,20 +52,22 @@ def submit_extrinsic(
         SubstrateRequestException: If the submission of the extrinsic fails, the error is logged and re-raised.
     """
     extrinsic_hash = extrinsic.extrinsic_hash
-    starting_block = substrate.get_block()
+    starting_block = subtensor.substrate.get_block()
 
     timeout = EXTRINSIC_SUBMISSION_TIMEOUT
     event = threading.Event()
 
     def submit():
         try:
-            response_ = substrate.submit_extrinsic(
+            response_ = subtensor.substrate.submit_extrinsic(
                 extrinsic,
                 wait_for_inclusion=wait_for_inclusion,
                 wait_for_finalization=wait_for_finalization,
             )
         except SubstrateRequestException as e:
-            logging.error(format_error_message(e.args[0], substrate=substrate))
+            logging.error(
+                format_error_message(e.args[0], substrate=subtensor.substrate)
+            )
             # Re-raise the exception for retrying of the extrinsic call. If we remove the retry logic,
             # the raise will need to be removed.
             raise
@@ -76,21 +79,24 @@ def submit():
         response = None
         future = executor.submit(submit)
         if not event.wait(timeout):
-            logging.error("Timed out waiting for extrinsic submission.")
-            after_timeout_block = substrate.get_block()
+            logging.error("Timed out waiting for extrinsic submission. Reconnecting.")
+            # force reconnection of the websocket
+            subtensor._get_substrate(force=True)
+            after_timeout_block = subtensor.substrate.get_block()
 
             for block_num in range(
                 starting_block["header"]["number"],
                 after_timeout_block["header"]["number"] + 1,
             ):
-                block_hash = substrate.get_block_hash(block_num)
+                block_hash = subtensor.substrate.get_block_hash(block_num)
                 try:
-                    response = substrate.retrieve_extrinsic_by_hash(
+                    response = subtensor.substrate.retrieve_extrinsic_by_hash(
                         block_hash, f"0x{extrinsic_hash.hex()}"
                     )
-                except ExtrinsicNotFound:
+                except (ExtrinsicNotFound, SubstrateRequestException):
                     continue
                 if response:
+                    logging.debug(f"Recovered extrinsic: {extrinsic}")
                     break
             if response is None:
                 logging.error(

bittensor/core/subtensor.py

@@ -212,11 +212,21 @@ def close(self):
         if self.substrate:
             self.substrate.close()
 
-    def _get_substrate(self):
-        """Establishes a connection to the Substrate node using configured parameters."""
+    def _get_substrate(self, force: bool = False):
+        """
+        Establishes a connection to the Substrate node using configured parameters.
+
+        Args:
+            force: forces a reconnection if this flag is set
+
+        """
         try:
             # Set up params.
-            if self.websocket is None or self.websocket.close_code is not None:
+            if force and self.websocket:
+                logging.debug("Closing websocket connection")
+                self.websocket.close()
+
+            if force or self.websocket is None or self.websocket.close_code is not None:
                 self.websocket = ws_client.connect(
                     self.chain_endpoint,
                     open_timeout=self._connection_timeout,
@@ -229,7 +239,7 @@ def _get_substrate(self):
                 websocket=self.websocket,
             )
             if self.log_verbose:
-                logging.debug(
+                logging.info(
                     f"Connected to {self.network} network and {self.chain_endpoint}."
                 )
 

tests/e2e_tests/test_commit_weights.py

@@ -6,6 +6,7 @@
 from bittensor.core.subtensor import Subtensor
 from bittensor.utils.balance import Balance
 from bittensor.utils.weight_utils import convert_weights_and_uids_for_emit
+from bittensor.core.extrinsics import utils
 from tests.e2e_tests.utils.chain_interactions import (
     add_stake,
     register_subnet,
@@ -32,6 +33,7 @@ async def test_commit_and_reveal_weights(local_chain):
         AssertionError: If any of the checks or verifications fail
     """
     netuid = 1
+    utils.EXTRINSIC_SUBMISSION_TIMEOUT = 12  # handle fast blocks
     print("Testing test_commit_and_reveal_weights")
     # Register root as Alice
     keypair, alice_wallet = setup_wallet("//Alice")

tests/unit_tests/extrinsics/test_utils.py

@@ -9,37 +9,44 @@
 )
 
 from bittensor.core.extrinsics import utils
+from bittensor.core.subtensor import Subtensor
 
 
 @pytest.fixture
 def set_extrinsics_timeout_env(monkeypatch):
     monkeypatch.setenv("EXTRINSIC_SUBMISSION_TIMEOUT", "1")
 
 
-def test_submit_extrinsic_timeout():
+@pytest.fixture
+def mock_subtensor():
+    mock_subtensor = MagicMock(autospec=Subtensor)
+    mock_substrate = MagicMock(autospec=SubstrateInterface)
+    mock_subtensor.substrate = mock_substrate
+    yield mock_subtensor
+
+
+def test_submit_extrinsic_timeout(mock_subtensor):
     timeout = 1
 
     def wait(extrinsic, wait_for_inclusion, wait_for_finalization):
         time.sleep(timeout + 0.01)
         return True
 
-    mock_substrate = MagicMock(autospec=SubstrateInterface)
-    mock_substrate.submit_extrinsic = wait
+    mock_subtensor.substrate.submit_extrinsic = wait
     mock_extrinsic = MagicMock(autospec=GenericExtrinsic)
     with patch.object(utils, "EXTRINSIC_SUBMISSION_TIMEOUT", timeout):
         with pytest.raises(SubstrateRequestException):
-            utils.submit_extrinsic(mock_substrate, mock_extrinsic, True, True)
+            utils.submit_extrinsic(mock_subtensor, mock_extrinsic, True, True)
 
 
-def test_submit_extrinsic_success():
-    mock_substrate = MagicMock(autospec=SubstrateInterface)
-    mock_substrate.submit_extrinsic.return_value = True
+def test_submit_extrinsic_success(mock_subtensor):
+    mock_subtensor.substrate.submit_extrinsic.return_value = True
     mock_extrinsic = MagicMock(autospec=GenericExtrinsic)
-    result = utils.submit_extrinsic(mock_substrate, mock_extrinsic, True, True)
+    result = utils.submit_extrinsic(mock_subtensor, mock_extrinsic, True, True)
     assert result is True
 
 
-def test_submit_extrinsic_timeout_env(set_extrinsics_timeout_env):
+def test_submit_extrinsic_timeout_env(set_extrinsics_timeout_env, mock_subtensor):
     importlib.reload(utils)
     timeout = utils.EXTRINSIC_SUBMISSION_TIMEOUT
     assert timeout < 5  # should be less than 5 seconds as taken from test env var
@@ -48,23 +55,21 @@ def wait(extrinsic, wait_for_inclusion, wait_for_finalization):
         time.sleep(timeout + 1)
         return True
 
-    mock_substrate = MagicMock(autospec=SubstrateInterface)
-    mock_substrate.submit_extrinsic = wait
+    mock_subtensor.substrate.submit_extrinsic = wait
     mock_extrinsic = MagicMock(autospec=GenericExtrinsic)
     with pytest.raises(SubstrateRequestException):
-        utils.submit_extrinsic(mock_substrate, mock_extrinsic, True, True)
+        utils.submit_extrinsic(mock_subtensor, mock_extrinsic, True, True)
 
 
-def test_submit_extrinsic_success_env(set_extrinsics_timeout_env):
+def test_submit_extrinsic_success_env(set_extrinsics_timeout_env, mock_subtensor):
     importlib.reload(utils)
-    mock_substrate = MagicMock(autospec=SubstrateInterface)
-    mock_substrate.submit_extrinsic.return_value = True
+    mock_subtensor.substrate.submit_extrinsic.return_value = True
     mock_extrinsic = MagicMock(autospec=GenericExtrinsic)
-    result = utils.submit_extrinsic(mock_substrate, mock_extrinsic, True, True)
+    result = utils.submit_extrinsic(mock_subtensor, mock_extrinsic, True, True)
     assert result is True
 
 
-def test_submit_extrinsic_timeout_env_float(monkeypatch):
+def test_submit_extrinsic_timeout_env_float(monkeypatch, mock_subtensor):
     monkeypatch.setenv("EXTRINSIC_SUBMISSION_TIMEOUT", "1.45")  # use float
 
     importlib.reload(utils)
@@ -76,11 +81,10 @@ def wait(extrinsic, wait_for_inclusion, wait_for_finalization):
         time.sleep(timeout + 0.3)  # sleep longer by float
         return True
 
-    mock_substrate = MagicMock(autospec=SubstrateInterface)
-    mock_substrate.submit_extrinsic = wait
+    mock_subtensor.substrate.submit_extrinsic = wait
     mock_extrinsic = MagicMock(autospec=GenericExtrinsic)
     with pytest.raises(SubstrateRequestException):
-        utils.submit_extrinsic(mock_substrate, mock_extrinsic, True, True)
+        utils.submit_extrinsic(mock_subtensor, mock_extrinsic, True, True)
 
 
 def test_import_timeout_env_parse(monkeypatch):