improve subtensors

Author: Roman

bittensor/core/async_subtensor.py

@@ -9,6 +9,7 @@
 import numpy as np
 import scalecodec
 from async_substrate_interface import AsyncSubstrateInterface
+from async_substrate_interface.substrate_addons import RetryAsyncSubstrate
 from bittensor_drand import get_encrypted_commitment
 from bittensor_wallet.utils import SS58_FORMAT
 from numpy.typing import NDArray
@@ -114,17 +115,21 @@ def __init__(
         self,
         network: Optional[str] = None,
         config: Optional["Config"] = None,
-        _mock: bool = False,
         log_verbose: bool = False,
+        fallback_chains: Optional[list[str]] = None,
+        retry_forever: bool = False,
+        _mock: bool = False,
     ):
         """
         Initializes an instance of the AsyncSubtensor class.
 
         Arguments:
             network (str): The network name or type to connect to.
             config (Optional[Config]): Configuration object for the AsyncSubtensor instance.
-            _mock: Whether this is a mock instance. Mainly just for use in testing.
             log_verbose (bool): Enables or disables verbose logging.
+            fallback_chains (list): List of fallback chains to use if no network is specified. Defaults to `None`.
+            retry_forever (bool): Whether to retry forever on connection errors. Defaults to `False`.
+            _mock: Whether this is a mock instance. Mainly just for use in testing.
 
         Raises:
             Any exceptions raised during the setup, configuration, or connection process.
@@ -151,6 +156,9 @@ def __init__(
             chain_name="Bittensor",
             _mock=_mock,
         )
+        self.substrate = self._get_substrate(
+            fallback_chains=fallback_chains, retry_forever=retry_forever, _mock=_mock
+        )
         if self.log_verbose:
             logging.info(
                 f"Connected to {self.network} network and {self.chain_endpoint}."
@@ -283,6 +291,38 @@ async def get_hyperparameter(
 
         return getattr(result, "value", result)
 
+    def _get_substrate(
+        self,
+        fallback_chains: Optional[list[str]] = None,
+        retry_forever: bool = False,
+        _mock: bool = False,
+    ) -> Union[AsyncSubstrateInterface, RetryAsyncSubstrate]:
+        """Creates the Substrate instance based on provided arguments.
+
+        Arguments:
+            fallback_chains (list): List of fallback chains to use if no network is specified. Defaults to `None`.
+            retry_forever (bool): Whether to retry forever on connection errors. Defaults to `False`.
+            _mock: Whether this is a mock instance. Mainly just for use in testing.
+
+        Returns:
+            the instance of the SubstrateInterface or RetrySyncSubstrate class.
+        """
+        if fallback_chains or retry_forever:
+            return RetryAsyncSubstrate(
+                url=self.chain_endpoint,
+                fallback_chains=fallback_chains,
+                retry_forever=retry_forever,
+                _mock=_mock,
+            )
+        return AsyncSubstrateInterface(
+            url=self.chain_endpoint,
+            ss58_format=SS58_FORMAT,
+            type_registry=TYPE_REGISTRY,
+            use_remote_preset=True,
+            chain_name="Bittensor",
+            _mock=_mock,
+        )
+
     # Subtensor queries ===========================================================================================
 
     async def query_constant(

bittensor/core/subtensor.py

@@ -6,6 +6,7 @@
 import numpy as np
 import scalecodec
 from async_substrate_interface.errors import SubstrateRequestException
+from async_substrate_interface.substrate_addons import RetrySyncSubstrate
 from async_substrate_interface.sync_substrate import SubstrateInterface
 from async_substrate_interface.types import ScaleObj
 from bittensor_drand import get_encrypted_commitment
@@ -116,17 +117,21 @@ def __init__(
         self,
         network: Optional[str] = None,
         config: Optional["Config"] = None,
-        _mock: bool = False,
         log_verbose: bool = False,
+        fallback_chains: Optional[list[str]] = None,
+        retry_forever: bool = False,
+        _mock: bool = False,
     ):
         """
         Initializes an instance of the Subtensor class.
 
         Arguments:
             network (str): The network name or type to connect to.
             config (Optional[Config]): Configuration object for the AsyncSubtensor instance.
-            _mock: Whether this is a mock instance. Mainly just for use in testing.
             log_verbose (bool): Enables or disables verbose logging.
+            fallback_chains (list): List of fallback chains to use if no network is specified. Defaults to `None`.
+            retry_forever (bool): Whether to retry forever on connection errors. Defaults to `False`.
+            _mock: Whether this is a mock instance. Mainly just for use in testing.
 
         Raises:
             Any exceptions raised during the setup, configuration, or connection process.
@@ -143,13 +148,8 @@ def __init__(
             f"Connecting to network: [blue]{self.network}[/blue], "
             f"chain_endpoint: [blue]{self.chain_endpoint}[/blue]> ..."
         )
-        self.substrate = SubstrateInterface(
-            url=self.chain_endpoint,
-            ss58_format=SS58_FORMAT,
-            type_registry=TYPE_REGISTRY,
-            use_remote_preset=True,
-            chain_name="Bittensor",
-            _mock=_mock,
+        self.substrate = self._get_substrate(
+            fallback_chains=fallback_chains, retry_forever=retry_forever, _mock=_mock
         )
         if self.log_verbose:
             logging.info(
@@ -163,11 +163,41 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         self.close()
 
     def close(self):
-        """
-        Closes the websocket connection
-        """
+        """Closes the websocket connection."""
         self.substrate.close()
 
+    def _get_substrate(
+        self,
+        fallback_chains: Optional[list[str]] = None,
+        retry_forever: bool = False,
+        _mock: bool = False,
+    ) -> Union[SubstrateInterface, RetrySyncSubstrate]:
+        """Creates the Substrate instance based on provided arguments.
+
+        Arguments:
+            fallback_chains (list): List of fallback chains to use if no network is specified. Defaults to `None`.
+            retry_forever (bool): Whether to retry forever on connection errors. Defaults to `False`.
+            _mock: Whether this is a mock instance. Mainly just for use in testing.
+
+        Returns:
+            the instance of the SubstrateInterface or RetrySyncSubstrate class.
+        """
+        if fallback_chains or retry_forever:
+            return RetrySyncSubstrate(
+                url=self.chain_endpoint,
+                fallback_chains=fallback_chains,
+                retry_forever=retry_forever,
+                _mock=_mock,
+            )
+        return SubstrateInterface(
+            url=self.chain_endpoint,
+            ss58_format=SS58_FORMAT,
+            type_registry=TYPE_REGISTRY,
+            use_remote_preset=True,
+            chain_name="Bittensor",
+            _mock=_mock,
+        )
+
     # Subtensor queries ===========================================================================================
 
     def query_constant(