[feature] external axon flags

bittensor/_axon/__init__.py

@@ -2,6 +2,7 @@
 """
 # The MIT License (MIT)
 # Copyright © 2021 Yuma Rao
+# Copyright © 2022 Opentensor Foundation
 
 # Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated 
 # documentation files (the “Software”), to deal in the Software without restriction, including without limitation 
@@ -60,6 +61,8 @@ def __new__(
             server: 'grpc._Server' = None,
             port: int = None,
             ip: str = None,
+            external_ip: str = None,
+            external_port: int = None,
             max_workers: int = None, 
             maximum_concurrent_rpcs: int = None,
             blacklist: 'Callable' = None,
@@ -96,6 +99,10 @@ def __new__(
                     Binding port.
                 ip (:type:`str`, `optional`):
                     Binding ip.
+                external_ip (:type:`str`, `optional`):
+                    The external ip of the server to broadcast to the network.
+                external_port (:type:`int`, `optional`):
+                    The external port of the server to broadcast to the network.
                 max_workers (:type:`int`, `optional`):
                     Used to create the threadpool if not passed, specifies the number of active threads servicing requests.
                 maximum_concurrent_rpcs (:type:`int`, `optional`):
@@ -115,6 +122,8 @@ def __new__(
         config = copy.deepcopy(config)
         config.axon.port = port if port != None else config.axon.port
         config.axon.ip = ip if ip != None else config.axon.ip
+        config.axon.external_ip = external_ip if external_ip != None else config.axon.external_ip
+        config.axon.external_port = external_port if external_port != None else config.axon.external_port
         config.axon.max_workers = max_workers if max_workers != None else config.axon.max_workers
         config.axon.maximum_concurrent_rpcs = maximum_concurrent_rpcs if maximum_concurrent_rpcs != None else config.axon.maximum_concurrent_rpcs
         config.axon.forward_timeout = forward_timeout if forward_timeout != None else config.axon.forward_timeout
@@ -160,6 +169,8 @@ def __new__(
             server = server,
             ip = config.axon.ip,
             port = config.axon.port,
+            external_ip=config.axon.external_ip, # don't use internal ip if it is None, we will try to find it later
+            external_port=config.axon.external_port or config.axon.port, # default to internal port if external port is not set
             forward = forward_text,
             backward = backward_text,
             synapses = synapses,
@@ -199,9 +210,13 @@ def add_args( cls, parser: argparse.ArgumentParser, prefix: str = None  ):
         prefix_str = '' if prefix == None else prefix + '.'
         try:
             parser.add_argument('--' + prefix_str + 'axon.port', type=int, 
-                    help='''The port this axon endpoint is served on. i.e. 8091''', default = bittensor.defaults.axon.port)
+                    help='''The local port this axon endpoint is binds. i.e. 8091''', default = bittensor.defaults.axon.port)
             parser.add_argument('--' + prefix_str + 'axon.ip', type=str, 
                 help='''The local ip this axon binds to. ie. [::]''', default = bittensor.defaults.axon.ip)
+            parser.add_argument('--' + prefix_str + 'axon.external_port', type=int, required=False,
+                    help='''The port this axon endpoint is served on. i.e. 8091''', default = bittensor.defaults.axon.external_port)
+            parser.add_argument('--' + prefix_str + 'axon.external_ip', type=str, required=False,
+                help='''The external ip this axon broadcasts to the network to. ie. [::]''', default = bittensor.defaults.axon.external_ip)
             parser.add_argument('--' + prefix_str + 'axon.max_workers', type=int, 
                 help='''The maximum number connection handler threads working simultaneously on this endpoint. 
                         The grpc server distributes new worker threads to service requests up to this number.''', default = bittensor.defaults.axon.max_workers)
@@ -230,6 +245,8 @@ def add_defaults(cls, defaults):
         defaults.axon = bittensor.Config()
         defaults.axon.port = os.getenv('BT_AXON_PORT') if os.getenv('BT_AXON_PORT') != None else 8091
         defaults.axon.ip = os.getenv('BT_AXON_IP') if os.getenv('BT_AXON_IP') != None else '[::]'
+        defaults.axon.external_port = os.getenv('BT_AXON_EXTERNAL_PORT') if os.getenv('BT_AXON_EXTERNAL_PORT') != None else None
+        defaults.axon.external_ip = os.getenv('BT_AXON_EXTERNAL_IP') if os.getenv('BT_AXON_EXTERNAL_IP') != None else None
         defaults.axon.max_workers = os.getenv('BT_AXON_MAX_WORERS') if os.getenv('BT_AXON_MAX_WORERS') != None else 10
         defaults.axon.maximum_concurrent_rpcs = os.getenv('BT_AXON_MAXIMUM_CONCURRENT_RPCS') if os.getenv('BT_AXON_MAXIMUM_CONCURRENT_RPCS') != None else 400
 
@@ -244,6 +261,7 @@ def check_config(cls, config: 'bittensor.Config' ):
         """ Check config for axon port and wallet
         """
         assert config.axon.port > 1024 and config.axon.port < 65535, 'port must be in range [1024, 65535]'
+        assert config.axon.external_port is None or (config.axon.external_port > 1024 and config.axon.external_port < 65535), 'external port must be in range [1024, 65535]'
         bittensor.wallet.check_config( config )
 
     @classmethod   

bittensor/_axon/axon_impl.py

@@ -2,6 +2,7 @@
 """
 # The MIT License (MIT)
 # Copyright © 2021 Yuma Rao
+# Copyright © 2022 Opentensor Foundation
 
 # Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated 
 # documentation files (the “Software”), to deal in the Software without restriction, including without limitation 
@@ -43,6 +44,8 @@ def __init__(
         wallet: 'bittensor.wallet',
         ip: str,
         port: int,
+        external_ip: str,
+        external_port: int,
         server: 'grpc._Server',
         forward: 'Callable',
         backward: 'Callable',
@@ -73,6 +76,8 @@ def __init__(
         """
         self.ip = ip
         self.port = port
+        self.external_ip = external_ip
+        self.external_port = external_port
         self.wallet = wallet
         self.server = server
         self.forward_callback = forward if forward != None else self.default_forward_callback

bittensor/_config/config_impl.py

@@ -3,6 +3,7 @@
 """
 # The MIT License (MIT)
 # Copyright © 2021 Yuma Rao
+# Copyright © 2022 Opentensor Foundation
 
 # Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated 
 # documentation files (the “Software”), to deal in the Software without restriction, including without limitation 
@@ -53,6 +54,8 @@ def to_defaults(self):
             if 'axon' in self.keys():
                 bittensor.defaults.axon.port = self.axon.port
                 bittensor.defaults.axon.ip = self.axon.ip
+                bittensor.defaults.axon.external_port = self.axon.external_port
+                bittensor.defaults.axon.external_ip = self.axon.external_ip
                 bittensor.defaults.axon.max_workers = self.axon.max_workers
                 bittensor.defaults.axon.maximum_concurrent_rpcs = self.axon.maximum_concurrent_rpcs
 

bittensor/_subtensor/subtensor_impl.py

@@ -412,15 +412,18 @@ def serve_axon (
             except net.UPNPCException as upnpc_exception:
                 raise RuntimeError('Failed to hole-punch with upnpc with exception {}'.format( upnpc_exception )) from upnpc_exception
         else:
-            external_port = axon.port
+            external_port = axon.external_port
 
         # ---- Get external ip ----
-        try:
-            external_ip = net.get_external_ip()
-            bittensor.__console__.print(":white_heavy_check_mark: [green]Found external ip: {}[/green]".format( external_ip ))
-            bittensor.logging.success(prefix = 'External IP', sufix = '<blue>{}</blue>'.format( external_ip ))
-        except Exception as E:
-            raise RuntimeError('Unable to attain your external ip. Check your internet connection. error: {}'.format(E)) from E
+        if axon.external_ip == None:
+            try:
+                external_ip = net.get_external_ip()
+                bittensor.__console__.print(":white_heavy_check_mark: [green]Found external ip: {}[/green]".format( external_ip ))
+                bittensor.logging.success(prefix = 'External IP', sufix = '<blue>{}</blue>'.format( external_ip ))
+            except Exception as E:
+                raise RuntimeError('Unable to attain your external ip. Check your internet connection. error: {}'.format(E)) from E
+        else:
+            external_ip = axon.external_ip
 
         # ---- Subscribe to chain ----
         serve_success = self.serve(

tests/unit_tests/bittensor_tests/test_axon.py

@@ -1,5 +1,6 @@
 # The MIT License (MIT)
 # Copyright © 2021 Yuma Rao
+# Copyright © 2022 Opentensor Foundation
 
 # Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated 
 # documentation files (the “Software”), to deal in the Software without restriction, including without limitation 
@@ -16,6 +17,7 @@
 # DEALINGS IN THE SOFTWARE.
 
 import time
+import unittest
 import unittest.mock as mock
 import uuid
 
@@ -1041,6 +1043,149 @@ def test_axon_is_destroyed():
     axonB.__del__()
     assert is_port_in_use( port ) == False
 
+# test external axon args
+class TestExternalAxon(unittest.TestCase):
+    """
+    Tests the external axon config flags 
+    `--axon.external_port` and `--axon.external_ip`
+    Need to verify the external config is used when broadcasting to the network
+    and the internal config is used when creating the grpc server
+
+    Also test the default behaviour when no external axon config is provided
+    (should use the internal axon config, like usual)
+    """
+
+    def test_external_ip_not_set_dont_use_internal_ip(self):
+        # Verify that not setting the external ip arg will NOT default to the internal axon ip
+        mock_add_insecure_port = mock.MagicMock(return_value=None)
+        mock_server = mock.MagicMock(
+            add_insecure_port=mock_add_insecure_port
+        )
+
+        mock_config = bittensor.axon.config()
+
+        axon = bittensor.axon ( ip = 'fake_ip', server=mock_server, config=mock_config )
+        assert axon.external_ip != axon.ip # should be different
+        assert axon.external_ip is None # should be None
+
+    def test_external_port_not_set_use_internal_port(self):
+        # Verify that not setting the external port arg will default to the internal axon port
+        mock_config = bittensor.axon.config()
+
+        axon = bittensor.axon ( port = 1234, config=mock_config )
+        assert axon.external_port == axon.port
+
+    def test_external_port_set_full_address_internal(self):
+        internal_port = 1234
+        external_port = 5678
+
+        mock_add_insecure_port = mock.MagicMock(return_value=None)
+        mock_server = mock.MagicMock(
+            add_insecure_port=mock_add_insecure_port
+        )
+
+        mock_config = bittensor.axon.config()
+
+        _ = bittensor.axon( port=internal_port, external_port=external_port, server=mock_server, config=mock_config )
+
+        mock_add_insecure_port.assert_called_once()
+        args, _ = mock_add_insecure_port.call_args
+        full_address0 = args[0]
+
+        assert f'{internal_port}' in full_address0 and f':{external_port}' not in full_address0
+
+        mock_add_insecure_port.reset_mock()
+
+        # Test using config
+        mock_config = bittensor.axon.config()
+
+        mock_config.axon.port = internal_port
+        mock_config.axon.external_port = external_port
+
+        _ = bittensor.axon( config=mock_config, server=mock_server )
+
+        mock_add_insecure_port.assert_called_once()
+        args, _ = mock_add_insecure_port.call_args
+        full_address0 = args[0]
+
+        assert f'{internal_port}' in full_address0, f'{internal_port} was not found in {full_address0}'
+        assert f':{external_port}' not in full_address0, f':{external_port} was found in {full_address0}'
+
+    def test_external_ip_set_full_address_internal(self):
+        internal_ip = 'fake_ip_internal'
+        external_ip = 'fake_ip_external'
+
+        mock_add_insecure_port = mock.MagicMock(return_value=None)
+        mock_server = mock.MagicMock(
+            add_insecure_port=mock_add_insecure_port
+        )
+
+        mock_config = bittensor.axon.config()
+
+        _ = bittensor.axon( ip=internal_ip, external_ip=external_ip, server=mock_server, config=mock_config )
+
+        mock_add_insecure_port.assert_called_once()
+        args, _ = mock_add_insecure_port.call_args
+        full_address0 = args[0]
+
+        assert f'{internal_ip}' in full_address0 and f'{external_ip}' not in full_address0
+
+        mock_add_insecure_port.reset_mock()
+
+        # Test using config
+        mock_config = bittensor.axon.config()
+        mock_config.axon.external_ip = external_ip
+        mock_config.axon.ip = internal_ip
+
+        _ = bittensor.axon( config=mock_config, server=mock_server )
+
+        mock_add_insecure_port.assert_called_once()
+        args, _ = mock_add_insecure_port.call_args
+        full_address0 = args[0]
+
+        assert f'{internal_ip}' in full_address0, f'{internal_ip} was not found in {full_address0}'
+        assert f'{external_ip}' not in full_address0, f'{external_ip} was found in {full_address0}'
+
+    def test_external_ip_port_set_full_address_internal(self):
+        internal_ip = 'fake_ip_internal'
+        external_ip = 'fake_ip_external'
+        internal_port = 1234
+        external_port = 5678
+
+        mock_add_insecure_port = mock.MagicMock(return_value=None)
+        mock_server = mock.MagicMock(
+            add_insecure_port=mock_add_insecure_port
+        )
+
+        mock_config = bittensor.axon.config()
+
+        _ = bittensor.axon( ip=internal_ip, external_ip=external_ip, port=internal_port, external_port=external_port, server=mock_server, config=mock_config )
+
+        mock_add_insecure_port.assert_called_once()
+        args, _ = mock_add_insecure_port.call_args
+        full_address0 = args[0]
+
+        assert f'{internal_ip}:{internal_port}' == full_address0 and f'{external_ip}:{external_port}' != full_address0
+
+        mock_add_insecure_port.reset_mock()
+
+        # Test using config
+        mock_config = bittensor.axon.config()
+
+        mock_config.axon.ip = internal_ip
+        mock_config.axon.external_ip = external_ip
+        mock_config.axon.port = internal_port
+        mock_config.axon.external_port = external_port
+
+        _ = bittensor.axon( config=mock_config, server=mock_server )
+
+        mock_add_insecure_port.assert_called_once()
+        args, _ = mock_add_insecure_port.call_args
+        full_address1 = args[0]
+
+        assert f'{internal_ip}:{internal_port}' == full_address1, f'{internal_ip}:{internal_port} is not eq to {full_address1}'
+        assert f'{external_ip}:{external_port}' != full_address1, f'{external_ip}:{external_port} is eq to {full_address1}'
+
 
 if __name__ == "__main__":
     # test_forward_joint_success()