Adds ws_shutdown_timer arg to AsyncSubtensor and SubtensorAPI

Author: thewhaleking

bittensor/core/async_subtensor.py

@@ -119,6 +119,7 @@ def __init__(
         retry_forever: bool = False,
         _mock: bool = False,
         archive_endpoints: Optional[list[str]] = None,
+        websocket_shutdown_timer: float = 5.0,
     ):
         """
         Initializes an instance of the AsyncSubtensor class.
@@ -157,6 +158,7 @@ def __init__(
             retry_forever=retry_forever,
             _mock=_mock,
             archive_endpoints=archive_endpoints,
+            ws_shutdown_timer=websocket_shutdown_timer
         )
         if self.log_verbose:
             logging.info(
@@ -296,6 +298,7 @@ def _get_substrate(
         retry_forever: bool = False,
         _mock: bool = False,
         archive_endpoints: Optional[list[str]] = None,
+        ws_shutdown_timer: float = 5.0
     ) -> Union[AsyncSubstrateInterface, RetryAsyncSubstrate]:
         """Creates the Substrate instance based on provided arguments.
 
@@ -307,6 +310,8 @@ def _get_substrate(
             archive_endpoints: Similar to fallback_endpoints, but specifically only archive nodes. Will be used in cases
                 where you are requesting a block that is too old for your current (presumably lite) node. Defaults to
                 `None`
+            ws_shutdown_timer: Amount of time, in seconds, to wait after the last response from the chain to close the
+                connection.
 
         Returns:
             the instance of the SubstrateInterface or RetrySyncSubstrate class.
@@ -322,6 +327,7 @@ def _get_substrate(
                 chain_name="Bittensor",
                 _mock=_mock,
                 archive_nodes=archive_endpoints,
+                ws_shutdown_timer=ws_shutdown_timer
             )
         return AsyncSubstrateInterface(
             url=self.chain_endpoint,
@@ -330,6 +336,7 @@ def _get_substrate(
             use_remote_preset=True,
             chain_name="Bittensor",
             _mock=_mock,
+            ws_shutdown_timer=ws_shutdown_timer
         )
 
     # Subtensor queries ===========================================================================================

bittensor/core/subtensor_api/__init__.py

@@ -29,6 +29,11 @@ class SubtensorApi:
         retry_forever: Whether to retry forever on connection errors. Defaults to `False`.
         log_verbose: Enables or disables verbose logging.
         mock: Whether this is a mock instance. Mainly just for use in testing.
+        archive_endpoints: Similar to fallback_endpoints, but specifically only archive nodes. Will be used in cases
+                where you are requesting a block that is too old for your current (presumably lite) node. Defaults to
+                `None`
+        websocket_shutdown_timer: Amount of time, in seconds, to wait after the last response from the chain to close
+            the connection. Only applicable to AsyncSubtensor.
 
     Example:
         # sync version
@@ -76,11 +81,13 @@ def __init__(
         log_verbose: bool = False,
         mock: bool = False,
         archive_endpoints: Optional[list[str]] = None,
+        websocket_shutdown_timer: float = 5.0
     ):
         self.network = network
         self._fallback_endpoints = fallback_endpoints
         self._archive_endpoints = archive_endpoints
         self._retry_forever = retry_forever
+        self._ws_shutdown_timer = websocket_shutdown_timer
         self._mock = mock
         self.log_verbose = log_verbose
         self.is_async = async_subtensor
@@ -122,6 +129,7 @@ def _get_subtensor(self) -> Union["_Subtensor", "_AsyncSubtensor"]:
                 retry_forever=self._retry_forever,
                 _mock=self._mock,
                 archive_endpoints=self._archive_endpoints,
+                websocket_shutdown_timer=self._ws_shutdown_timer
             )
             self.initialize = _subtensor.initialize
             return _subtensor

pyproject.toml

@@ -36,7 +36,7 @@ dependencies = [
     "uvicorn",
     "bittensor-drand>=0.5.0",
     "bittensor-wallet>=3.0.8",
-    "async-substrate-interface>=1.3.0"
+    "async-substrate-interface>=1.3.1"
 ]
 
 [project.optional-dependencies]