Merge pull request #2553 from opentensor/feat/roman/hadle-the-maximun-connection-limit

[SDK] Handle server connection limit

Author: basfroman

bittensor/core/subtensor.py

@@ -17,6 +17,7 @@
 from scalecodec.type_registry import load_type_registry_preset
 from scalecodec.types import ScaleType
 from substrateinterface.base import QueryMapResult, SubstrateInterface
+from websockets.exceptions import InvalidStatus
 from websockets.sync import client as ws_client
 
 from bittensor.core import settings
@@ -233,6 +234,7 @@ def _get_substrate(self, force: bool = False):
                     open_timeout=self._connection_timeout,
                     max_size=2**32,
                 )
+
             self.substrate = SubstrateInterface(
                 ss58_format=settings.SS58_FORMAT,
                 use_remote_preset=True,
@@ -244,19 +246,26 @@ def _get_substrate(self, force: bool = False):
                     f"Connected to {self.network} network and {self.chain_endpoint}."
                 )
 
-        except (ConnectionRefusedError, ssl.SSLError) as error:
-            logging.error(
-                f"<red>Could not connect to</red> <blue>{self.network}</blue> <red>network with</red> <blue>{self.chain_endpoint}</blue> <red>chain endpoint.</red>",
+        except ConnectionRefusedError as error:
+            logging.critical(
+                f"[red]Could not connect to[/red] [blue]{self.network}[/blue] [red]network with[/red] [blue]{self.chain_endpoint}[/blue] [red]chain endpoint.[/red]",
             )
             raise ConnectionRefusedError(error.args)
-        except ssl.SSLError as e:
+
+        except ssl.SSLError as error:
             logging.critical(
                 "SSL error occurred. To resolve this issue, run the following command in your terminal:"
             )
             logging.critical("[blue]sudo python -m bittensor certifi[/blue]")
             raise RuntimeError(
                 "SSL configuration issue, please follow the instructions above."
-            ) from e
+            ) from error
+
+        except InvalidStatus as error:
+            logging.critical(
+                f"[red]You have reached the limit of simultaneous connections to the server.[/red]"
+            )
+            raise InvalidStatus(error.response) from error
 
     @staticmethod
     def config() -> "Config":

tests/unit_tests/test_subtensor.py

@@ -2880,3 +2880,27 @@ def test_set_weights_with_commit_reveal_enabled(subtensor, mocker):
         wait_for_finalization=fake_wait_for_finalization,
     )
     assert result == mocked_commit_reveal_v3_extrinsic.return_value
+
+
+def test_connection_limit(mocker):
+    """Test connection limit is not exceeded."""
+    # Technically speaking, this test should exist in integration tests. But to reduce server costs we will leave this
+    # test here.
+
+    # Preps
+    mocker.patch.object(
+        subtensor_module.ws_client,
+        "connect",
+        side_effect=subtensor_module.InvalidStatus(
+            response=mocker.Mock(
+                response=mocker.Mock(
+                    status_code=429, message="test connection limit error"
+                )
+            )
+        ),
+    )
+    # Call with assertions
+
+    with pytest.raises(subtensor_module.InvalidStatus):
+        for i in range(2):
+            Subtensor("test")