diff --git a/bittensor/core/async_subtensor.py b/bittensor/core/async_subtensor.py index f91df67952..11c01f5059 100644 --- a/bittensor/core/async_subtensor.py +++ b/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 =========================================================================================== diff --git a/bittensor/core/subtensor_api/__init__.py b/bittensor/core/subtensor_api/__init__.py index 464b325b27..5ae0bf7134 100644 --- a/bittensor/core/subtensor_api/__init__.py +++ b/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