diff --git a/tensorboard/plugins/core/core_plugin.py b/tensorboard/plugins/core/core_plugin.py index cfc815e134..c43a4636fc 100644 --- a/tensorboard/plugins/core/core_plugin.py +++ b/tensorboard/plugins/core/core_plugin.py @@ -341,6 +341,21 @@ def define_flags(self, parser): % DEFAULT_PORT, ) + parser.add_argument( + "--reuse_port", + metavar="BOOL", + # Custom str-to-bool converter since regular bool() doesn't work. + type=lambda v: {"true": True, "false": False}.get(v.lower(), v), + choices=[True, False], + default=False, + help="""\ +Enables the SO_REUSEPORT option on the socket opened by TensorBoard's HTTP +server, for platforms that support it. This is useful in cases when a parent +process has obtained the port already and wants to delegate access to the +port to TensorBoard as a subprocess.(default: %(default)s).\ +""", + ) + parser.add_argument( "--load_fast", action="store_true", diff --git a/tensorboard/plugins/core/core_plugin_test.py b/tensorboard/plugins/core/core_plugin_test.py index 963cdd08c5..269356fb3e 100644 --- a/tensorboard/plugins/core/core_plugin_test.py +++ b/tensorboard/plugins/core/core_plugin_test.py @@ -57,6 +57,7 @@ def __init__( path_prefix="", generic_data="true", grpc_data_provider="", + reuse_port=False, ): self.bind_all = bind_all self.host = host @@ -69,6 +70,7 @@ def __init__( self.path_prefix = path_prefix self.generic_data = generic_data self.grpc_data_provider = grpc_data_provider + self.reuse_port = reuse_port class CorePluginFlagsTest(tf.test.TestCase): diff --git a/tensorboard/program.py b/tensorboard/program.py index 5e77cec796..19c9cdc67e 100644 --- a/tensorboard/program.py +++ b/tensorboard/program.py @@ -682,13 +682,21 @@ def _get_wildcard_address(self, port): return fallback_address def server_bind(self): - """Override to enable IPV4 mapping for IPV6 sockets when desired. + """Override to set custom options on the socket.""" + if self._flags.reuse_port: + try: + socket.SO_REUSEPORT + except AttributeError: + raise TensorBoardServerException( + "TensorBoard --reuse_port option is not supported on this platform" + ) + self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1) - The main use case for this is so that when no host is specified, - TensorBoard can listen on all interfaces for both IPv4 and IPv6 - connections, rather than having to choose v4 or v6 and hope the - browser didn't choose the other one. - """ + # Enable IPV4 mapping for IPV6 sockets when desired. + # The main use case for this is so that when no host is specified, + # TensorBoard can listen on all interfaces for both IPv4 and IPv6 + # connections, rather than having to choose v4 or v6 and hope the + # browser didn't choose the other one. socket_is_v6 = ( hasattr(socket, "AF_INET6") and self.socket.family == socket.AF_INET6 diff --git a/tensorboard/program_test.py b/tensorboard/program_test.py index b346c32b42..d878cbf4cf 100644 --- a/tensorboard/program_test.py +++ b/tensorboard/program_test.py @@ -89,6 +89,7 @@ def make_flags(self, **kwargs): flags = argparse.Namespace() kwargs.setdefault("host", None) kwargs.setdefault("bind_all", kwargs["host"] is None) + kwargs.setdefault("reuse_port", False) for k, v in kwargs.items(): setattr(flags, k, v) return flags