From c8515a6054b5971b522bb5c9616b3782bdf14cb6 Mon Sep 17 00:00:00 2001 From: Benjamin Himes Date: Wed, 27 Nov 2024 08:47:57 +0200 Subject: [PATCH 1/5] Updates `ensure_connected` for the change to websockets. --- bittensor/utils/networking.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/bittensor/utils/networking.py b/bittensor/utils/networking.py index aca8ee2b04..d4ff98af46 100644 --- a/bittensor/utils/networking.py +++ b/bittensor/utils/networking.py @@ -10,7 +10,7 @@ import netaddr import requests from retry import retry -from websocket import WebSocketConnectionClosedException +from websockets.exceptions import ConnectionClosed, ConnectionClosedError from bittensor.utils.btlogging import logging @@ -193,7 +193,7 @@ def wrapper(self, *args, **kwargs): try: return func(self, *args, **kwargs) - except WebSocketConnectionClosedException: + except (ConnectionClosed, ConnectionClosedError): logging.console.warning( "WebSocket connection closed. Attempting to reconnect 5 times..." ) From ff367694f0794670c98e4008844b9931222b3184 Mon Sep 17 00:00:00 2001 From: Benjamin Himes Date: Wed, 27 Nov 2024 08:51:57 +0200 Subject: [PATCH 2/5] Only need to catch ConnectionClosed, as ConnectionClosedError subclasses ConnectionClosed. --- bittensor/utils/networking.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/bittensor/utils/networking.py b/bittensor/utils/networking.py index d4ff98af46..f47c88512e 100644 --- a/bittensor/utils/networking.py +++ b/bittensor/utils/networking.py @@ -10,7 +10,7 @@ import netaddr import requests from retry import retry -from websockets.exceptions import ConnectionClosed, ConnectionClosedError +from websockets.exceptions import ConnectionClosed from bittensor.utils.btlogging import logging @@ -193,7 +193,7 @@ def wrapper(self, *args, **kwargs): try: return func(self, *args, **kwargs) - except (ConnectionClosed, ConnectionClosedError): + except ConnectionClosed: logging.console.warning( "WebSocket connection closed. Attempting to reconnect 5 times..." ) From e16aa17c5417e1e48beba99a11c044f37dc18525 Mon Sep 17 00:00:00 2001 From: Benjamin Himes Date: Wed, 27 Nov 2024 09:32:48 +0200 Subject: [PATCH 3/5] Regenerates the websocket connection if it is closed. --- bittensor/core/subtensor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bittensor/core/subtensor.py b/bittensor/core/subtensor.py index 6c17cc87b0..fa46242277 100644 --- a/bittensor/core/subtensor.py +++ b/bittensor/core/subtensor.py @@ -216,7 +216,7 @@ def _get_substrate(self): """Establishes a connection to the Substrate node using configured parameters.""" try: # Set up params. - if not self.websocket: + if 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, From 918242f89da622c1387d832282c6734284c0f98f Mon Sep 17 00:00:00 2001 From: Benjamin Himes Date: Wed, 27 Nov 2024 09:56:42 +0200 Subject: [PATCH 4/5] Add property to FakeWebsocket --- tests/helpers/helpers.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/helpers/helpers.py b/tests/helpers/helpers.py index 1a7581ccd5..4b8253e0ef 100644 --- a/tests/helpers/helpers.py +++ b/tests/helpers/helpers.py @@ -125,6 +125,8 @@ def get_mock_neuron_by_uid(uid: int, **kwargs) -> NeuronInfo: class FakeWebsocket(ClientConnection): + close_code = None + def __init__(self, *args, seed, **kwargs): protocol = ClientProtocol(parse_uri("ws://127.0.0.1:9945")) super().__init__(socket=None, protocol=protocol, **kwargs) From f050db7da8c087298be03c22830cadecd9e5e1dc Mon Sep 17 00:00:00 2001 From: Benjamin Himes Date: Wed, 27 Nov 2024 10:36:48 +0200 Subject: [PATCH 5/5] Ruff --- tests/helpers/helpers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/helpers/helpers.py b/tests/helpers/helpers.py index 4b8253e0ef..a125e15329 100644 --- a/tests/helpers/helpers.py +++ b/tests/helpers/helpers.py @@ -126,7 +126,7 @@ def get_mock_neuron_by_uid(uid: int, **kwargs) -> NeuronInfo: class FakeWebsocket(ClientConnection): close_code = None - + def __init__(self, *args, seed, **kwargs): protocol = ClientProtocol(parse_uri("ws://127.0.0.1:9945")) super().__init__(socket=None, protocol=protocol, **kwargs)