Merge pull request #2917 from opentensor/feat/thewhaleking/new-websockets-option

New websockets option

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

@@ -31,6 +31,8 @@ class SubtensorApi:
         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
@@ -78,11 +80,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
@@ -124,6 +128,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