Spaces:
Paused
Paused
| # Copyright 2017 The TensorFlow Authors. All Rights Reserved. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| # ============================================================================== | |
| """Utilities for TensorBoard command line program. | |
| This is a lightweight module for bringing up a TensorBoard HTTP server | |
| or emulating the `tensorboard` shell command. | |
| Those wishing to create custom builds of TensorBoard can use this module | |
| by swapping out `tensorboard.main` with the custom definition that | |
| modifies the set of plugins and static assets. | |
| This module does not depend on first-party plugins or the default web | |
| server assets. Those are defined in `tensorboard.default`. | |
| """ | |
| from abc import ABCMeta | |
| from abc import abstractmethod | |
| import argparse | |
| import atexit | |
| from collections import defaultdict | |
| import errno | |
| import logging | |
| import mimetypes | |
| import os | |
| import shlex | |
| import signal | |
| import socket | |
| import sys | |
| import threading | |
| import time | |
| import urllib.parse | |
| from absl import flags as absl_flags | |
| from absl.flags import argparse_flags | |
| from werkzeug import serving | |
| from tensorboard import manager | |
| from tensorboard import version | |
| from tensorboard.backend import application | |
| from tensorboard.backend.event_processing import data_ingester as local_ingester | |
| from tensorboard.backend.event_processing import event_file_inspector as efi | |
| from tensorboard.data import server_ingester | |
| from tensorboard.plugins.core import core_plugin | |
| from tensorboard.util import tb_logging | |
| logger = tb_logging.get_logger() | |
| # Default subcommand name. This is a user-facing CLI and should not change. | |
| _SERVE_SUBCOMMAND_NAME = "serve" | |
| # Internal flag name used to store which subcommand was invoked. | |
| _SUBCOMMAND_FLAG = "__tensorboard_subcommand" | |
| # Message printed when we actually use the data server, so that users are not | |
| # caught terribly by surprise. | |
| _DATA_SERVER_ADVISORY_MESSAGE = """ | |
| NOTE: Using experimental fast data loading logic. To disable, pass | |
| "--load_fast=false" and report issues on GitHub. More details: | |
| https://github.com/tensorflow/tensorboard/issues/4784 | |
| """ | |
| # Message printed with `--load_fast=true` if the data server could not start up. | |
| # To be formatted with one `DataServerStartupError` interpoland. | |
| _DATA_SERVER_STARTUP_ERROR_MESSAGE_TEMPLATE = """\ | |
| Could not start data server: %s. | |
| Try with --load_fast=false and report issues on GitHub. Details: | |
| https://github.com/tensorflow/tensorboard/issues/4784 | |
| """ | |
| class TensorBoard: | |
| """Class for running TensorBoard. | |
| Fields: | |
| plugin_loaders: Set from plugins passed to constructor. | |
| assets_zip_provider: Set by constructor. | |
| server_class: Set by constructor. | |
| flags: An argparse.Namespace set by the configure() method. | |
| cache_key: As `manager.cache_key`; set by the configure() method. | |
| """ | |
| def __init__( | |
| self, | |
| plugins=None, | |
| assets_zip_provider=None, | |
| server_class=None, | |
| subcommands=None, | |
| ): | |
| """Creates new instance. | |
| Args: | |
| plugins: A list of TensorBoard plugins to load, as TBPlugin classes or | |
| TBLoader instances or classes. If not specified, defaults to first-party | |
| plugins. | |
| assets_zip_provider: A function that provides a zip file containing | |
| assets to the application. If `None`, the default TensorBoard web | |
| assets will be used. (If building from source, your binary must | |
| explicitly depend on `//tensorboard:assets_lib` if you pass `None`.) | |
| server_class: An optional factory for a `TensorBoardServer` to use | |
| for serving the TensorBoard WSGI app. If provided, its callable | |
| signature should match that of `TensorBoardServer.__init__`. | |
| subcommands: An optional list of TensorBoardSubcommand objects, which | |
| extend the functionality of the CLI. | |
| :type plugins: | |
| list[ | |
| base_plugin.TBLoader | Type[base_plugin.TBLoader] | | |
| Type[base_plugin.TBPlugin] | |
| ] | |
| """ | |
| if plugins is None: | |
| from tensorboard import default | |
| plugins = default.get_plugins() | |
| if assets_zip_provider is None: | |
| try: | |
| from tensorboard import assets | |
| except ImportError as e: | |
| # `tensorboard.assets` is not a strict Bazel dep; clients are | |
| # required to either depend on `//tensorboard:assets_lib` or | |
| # pass a valid assets provider. | |
| raise ImportError( | |
| "No `assets_zip_provider` given, but `tensorboard.assets` " | |
| "could not be imported to resolve defaults" | |
| ) from e | |
| assets_zip_provider = assets.get_default_assets_zip_provider() | |
| if server_class is None: | |
| server_class = create_port_scanning_werkzeug_server | |
| if subcommands is None: | |
| subcommands = [] | |
| self.plugin_loaders = [ | |
| application.make_plugin_loader(p) for p in plugins | |
| ] | |
| self.assets_zip_provider = assets_zip_provider | |
| self.server_class = server_class | |
| self.subcommands = {} | |
| for subcommand in subcommands: | |
| name = subcommand.name() | |
| if name in self.subcommands or name == _SERVE_SUBCOMMAND_NAME: | |
| raise ValueError("Duplicate subcommand name: %r" % name) | |
| self.subcommands[name] = subcommand | |
| self.flags = None | |
| def configure(self, argv=("",), **kwargs): | |
| """Configures TensorBoard behavior via flags. | |
| This method will populate the "flags" property with an argparse.Namespace | |
| representing flag values parsed from the provided argv list, overridden by | |
| explicit flags from remaining keyword arguments. | |
| Args: | |
| argv: Can be set to CLI args equivalent to sys.argv; the first arg is | |
| taken to be the name of the path being executed. | |
| kwargs: Additional arguments will override what was parsed from | |
| argv. They must be passed as Python data structures, e.g. | |
| `foo=1` rather than `foo="1"`. | |
| Returns: | |
| Either argv[:1] if argv was non-empty, or [''] otherwise, as a mechanism | |
| for absl.app.run() compatibility. | |
| Raises: | |
| ValueError: If flag values are invalid. | |
| """ | |
| base_parser = argparse_flags.ArgumentParser( | |
| prog="tensorboard", | |
| description=( | |
| "TensorBoard is a suite of web applications for " | |
| "inspecting and understanding your TensorFlow runs " | |
| "and graphs. https://github.com/tensorflow/tensorboard " | |
| ), | |
| ) | |
| subparsers = base_parser.add_subparsers( | |
| help="TensorBoard subcommand (defaults to %r)" | |
| % _SERVE_SUBCOMMAND_NAME | |
| ) | |
| serve_subparser = subparsers.add_parser( | |
| _SERVE_SUBCOMMAND_NAME, | |
| help="start local TensorBoard server (default subcommand)", | |
| ) | |
| serve_subparser.set_defaults( | |
| **{_SUBCOMMAND_FLAG: _SERVE_SUBCOMMAND_NAME} | |
| ) | |
| if len(argv) < 2 or argv[1].startswith("-"): | |
| # This invocation, if valid, must not use any subcommands: we | |
| # don't permit flags before the subcommand name. | |
| serve_parser = base_parser | |
| else: | |
| # This invocation, if valid, must use a subcommand: we don't take | |
| # any positional arguments to `serve`. | |
| serve_parser = serve_subparser | |
| for name, subcommand in self.subcommands.items(): | |
| subparser = subparsers.add_parser( | |
| name, | |
| help=subcommand.help(), | |
| description=subcommand.description(), | |
| ) | |
| subparser.set_defaults(**{_SUBCOMMAND_FLAG: name}) | |
| subcommand.define_flags(subparser) | |
| for loader in self.plugin_loaders: | |
| loader.define_flags(serve_parser) | |
| arg0 = argv[0] if argv else "" | |
| flags = base_parser.parse_args(argv[1:]) # Strip binary name from argv. | |
| if getattr(flags, _SUBCOMMAND_FLAG, None) is None: | |
| # Manually assign default value rather than using `set_defaults` | |
| # on the base parser to work around Python bug #9351 on old | |
| # versions of `argparse`: <https://bugs.python.org/issue9351> | |
| setattr(flags, _SUBCOMMAND_FLAG, _SERVE_SUBCOMMAND_NAME) | |
| self.cache_key = manager.cache_key( | |
| working_directory=os.getcwd(), | |
| arguments=argv[1:], | |
| configure_kwargs=kwargs, | |
| ) | |
| if arg0: | |
| # Only expose main module Abseil flags as TensorBoard native flags. | |
| # This is the same logic Abseil's ArgumentParser uses for determining | |
| # which Abseil flags to include in the short helpstring. | |
| for flag in set(absl_flags.FLAGS.get_key_flags_for_module(arg0)): | |
| if hasattr(flags, flag.name): | |
| raise ValueError("Conflicting Abseil flag: %s" % flag.name) | |
| setattr(flags, flag.name, flag.value) | |
| for k, v in kwargs.items(): | |
| if not hasattr(flags, k): | |
| raise ValueError("Unknown TensorBoard flag: %s" % k) | |
| setattr(flags, k, v) | |
| if getattr(flags, _SUBCOMMAND_FLAG) == _SERVE_SUBCOMMAND_NAME: | |
| for loader in self.plugin_loaders: | |
| loader.fix_flags(flags) | |
| self.flags = flags | |
| return [arg0] | |
| def main(self, ignored_argv=("",)): | |
| """Blocking main function for TensorBoard. | |
| This method is called by `tensorboard.main.run_main`, which is the | |
| standard entrypoint for the tensorboard command line program. The | |
| configure() method must be called first. | |
| Args: | |
| ignored_argv: Do not pass. Required for Abseil compatibility. | |
| Returns: | |
| Process exit code, i.e. 0 if successful or non-zero on failure. In | |
| practice, an exception will most likely be raised instead of | |
| returning non-zero. | |
| :rtype: int | |
| """ | |
| self._install_signal_handler(signal.SIGTERM, "SIGTERM") | |
| self._fix_mime_types() | |
| subcommand_name = getattr(self.flags, _SUBCOMMAND_FLAG) | |
| if subcommand_name == _SERVE_SUBCOMMAND_NAME: | |
| runner = self._run_serve_subcommand | |
| else: | |
| runner = self.subcommands[subcommand_name].run | |
| return runner(self.flags) or 0 | |
| def _run_serve_subcommand(self, flags): | |
| # TODO(#2801): Make `--version` a flag on only the base parser, not `serve`. | |
| if flags.version_tb: | |
| print(version.VERSION) | |
| return 0 | |
| if flags.inspect: | |
| # TODO(@wchargin): Convert `inspect` to a normal subcommand? | |
| logger.info( | |
| "Not bringing up TensorBoard, but inspecting event files." | |
| ) | |
| event_file = os.path.expanduser(flags.event_file) | |
| efi.inspect(flags.logdir, event_file, flags.tag) | |
| return 0 | |
| try: | |
| server = self._make_server() | |
| server.print_serving_message() | |
| self._register_info(server) | |
| server.serve_forever() | |
| return 0 | |
| except TensorBoardServerException as e: | |
| logger.error(e.msg) | |
| sys.stderr.write("ERROR: %s\n" % e.msg) | |
| sys.stderr.flush() | |
| return -1 | |
| def launch(self): | |
| """Python API for launching TensorBoard. | |
| This method is the same as main() except it launches TensorBoard in | |
| a separate permanent thread. The configure() method must be called | |
| first. | |
| Returns: | |
| The URL of the TensorBoard web server. | |
| :rtype: str | |
| """ | |
| # Make it easy to run TensorBoard inside other programs, e.g. Colab. | |
| server = self._make_server() | |
| thread = threading.Thread( | |
| target=server.serve_forever, name="TensorBoard" | |
| ) | |
| thread.daemon = True | |
| thread.start() | |
| return server.get_url() | |
| def _register_info(self, server): | |
| """Write a TensorBoardInfo file and arrange for its cleanup. | |
| Args: | |
| server: The result of `self._make_server()`. | |
| """ | |
| server_url = urllib.parse.urlparse(server.get_url()) | |
| info = manager.TensorBoardInfo( | |
| version=version.VERSION, | |
| start_time=int(time.time()), | |
| port=server_url.port, | |
| pid=os.getpid(), | |
| path_prefix=self.flags.path_prefix, | |
| logdir=self.flags.logdir or self.flags.logdir_spec, | |
| db=self.flags.db, | |
| cache_key=self.cache_key, | |
| ) | |
| atexit.register(manager.remove_info_file) | |
| manager.write_info_file(info) | |
| def _install_signal_handler(self, signal_number, signal_name): | |
| """Set a signal handler to gracefully exit on the given signal. | |
| When this process receives the given signal, it will run `atexit` | |
| handlers and then exit with `0`. | |
| Args: | |
| signal_number: The numeric code for the signal to handle, like | |
| `signal.SIGTERM`. | |
| signal_name: The human-readable signal name. | |
| """ | |
| # Note to maintainers: Google-internal code overrides this | |
| # method (cf. cl/334534610). Double-check changes before | |
| # modifying API. | |
| old_signal_handler = None # set below | |
| def handler(handled_signal_number, frame): | |
| # In case we catch this signal again while running atexit | |
| # handlers, take the hint and actually die. | |
| signal.signal(signal_number, signal.SIG_DFL) | |
| sys.stderr.write( | |
| "TensorBoard caught %s; exiting...\n" % signal_name | |
| ) | |
| # The main thread is the only non-daemon thread, so it suffices to | |
| # exit hence. | |
| if old_signal_handler not in (signal.SIG_IGN, signal.SIG_DFL): | |
| old_signal_handler(handled_signal_number, frame) | |
| sys.exit(0) | |
| old_signal_handler = signal.signal(signal_number, handler) | |
| def _fix_mime_types(self): | |
| """Fix incorrect entries in the `mimetypes` registry. | |
| On Windows, the Python standard library's `mimetypes` reads in | |
| mappings from file extension to MIME type from the Windows | |
| registry. Other applications can and do write incorrect values | |
| to this registry, which causes `mimetypes.guess_type` to return | |
| incorrect values, which causes TensorBoard to fail to render on | |
| the frontend. | |
| This method hard-codes the correct mappings for certain MIME | |
| types that are known to be either used by TensorBoard or | |
| problematic in general. | |
| """ | |
| # Known to be problematic when Visual Studio is installed: | |
| # <https://github.com/tensorflow/tensorboard/issues/3120> | |
| mimetypes.add_type("text/javascript", ".js") | |
| # Not known to be problematic, but used by TensorBoard: | |
| mimetypes.add_type("font/woff2", ".woff2") | |
| mimetypes.add_type("text/html", ".html") | |
| def _start_subprocess_data_ingester(self): | |
| """Creates, starts, and returns a `SubprocessServerDataIngester`.""" | |
| flags = self.flags | |
| server_binary = server_ingester.get_server_binary() | |
| ingester = server_ingester.SubprocessServerDataIngester( | |
| server_binary=server_binary, | |
| logdir=flags.logdir, | |
| reload_interval=flags.reload_interval, | |
| channel_creds_type=flags.grpc_creds_type, | |
| samples_per_plugin=flags.samples_per_plugin, | |
| extra_flags=shlex.split(flags.extra_data_server_flags), | |
| ) | |
| ingester.start() | |
| return ingester | |
| def _make_data_ingester(self): | |
| """Determines the right data ingester, starts it, and returns it.""" | |
| flags = self.flags | |
| if flags.grpc_data_provider: | |
| ingester = server_ingester.ExistingServerDataIngester( | |
| flags.grpc_data_provider, | |
| channel_creds_type=flags.grpc_creds_type, | |
| ) | |
| ingester.start() | |
| return ingester | |
| if flags.load_fast == "true": | |
| try: | |
| return self._start_subprocess_data_ingester() | |
| except server_ingester.NoDataServerError as e: | |
| msg = "Option --load_fast=true not available: %s\n" % e | |
| sys.stderr.write(msg) | |
| sys.exit(1) | |
| except server_ingester.DataServerStartupError as e: | |
| msg = _DATA_SERVER_STARTUP_ERROR_MESSAGE_TEMPLATE % e | |
| sys.stderr.write(msg) | |
| sys.exit(1) | |
| if flags.load_fast == "auto" and _should_use_data_server(flags): | |
| try: | |
| ingester = self._start_subprocess_data_ingester() | |
| sys.stderr.write(_DATA_SERVER_ADVISORY_MESSAGE) | |
| sys.stderr.flush() | |
| return ingester | |
| except server_ingester.NoDataServerError as e: | |
| logger.info("No data server: %s", e) | |
| except server_ingester.DataServerStartupError as e: | |
| logger.info( | |
| "Data server error: %s; falling back to multiplexer", e | |
| ) | |
| ingester = local_ingester.LocalDataIngester(flags) | |
| ingester.start() | |
| return ingester | |
| def _make_data_provider(self): | |
| """Returns `(data_provider, deprecated_multiplexer)`.""" | |
| ingester = self._make_data_ingester() | |
| # Stash ingester so that it can avoid GCing Windows file handles. | |
| # (See comment in `SubprocessServerDataIngester.start` for details.) | |
| self._ingester = ingester | |
| deprecated_multiplexer = None | |
| if isinstance(ingester, local_ingester.LocalDataIngester): | |
| deprecated_multiplexer = ingester.deprecated_multiplexer | |
| return (ingester.data_provider, deprecated_multiplexer) | |
| def _make_server(self): | |
| """Constructs the TensorBoard WSGI app and instantiates the server.""" | |
| (data_provider, deprecated_multiplexer) = self._make_data_provider() | |
| app = application.TensorBoardWSGIApp( | |
| self.flags, | |
| self.plugin_loaders, | |
| data_provider, | |
| self.assets_zip_provider, | |
| deprecated_multiplexer, | |
| ) | |
| return self.server_class(app, self.flags) | |
| def _should_use_data_server(flags): | |
| if flags.logdir_spec and not flags.logdir: | |
| logger.info( | |
| "Note: --logdir_spec is not supported with --load_fast behavior; " | |
| "falling back to slower Python-only load path. To use the data " | |
| "server, replace --logdir_spec with --logdir." | |
| ) | |
| return False | |
| if not flags.logdir: | |
| # Using some other legacy mode; not supported. | |
| return False | |
| if "://" in flags.logdir and not flags.logdir.startswith("gs://"): | |
| logger.info( | |
| "Note: --load_fast behavior only supports local and GCS (gs://) " | |
| "paths; falling back to slower Python-only load path." | |
| ) | |
| return False | |
| if flags.detect_file_replacement is True: | |
| logger.info( | |
| "Note: --detect_file_replacement=true is not supported with " | |
| "--load_fast behavior; falling back to slower Python-only load " | |
| "path." | |
| ) | |
| return False | |
| return True | |
| class TensorBoardSubcommand(metaclass=ABCMeta): | |
| """Experimental private API for defining subcommands for tensorboard. | |
| The intended use is something like: | |
| `tensorboard <sub_cmd_name> <additional_args...>` | |
| Since our hosted service at http://tensorboard.dev has been shut down, this | |
| functionality is no longer used, but the support for subcommands remains, | |
| in case it is ever useful again. | |
| """ | |
| def name(self): | |
| """Name of this subcommand, as specified on the command line. | |
| This must be unique across all subcommands. | |
| Returns: | |
| A string. | |
| """ | |
| pass | |
| def define_flags(self, parser): | |
| """Configure an argument parser for this subcommand. | |
| Flags whose names start with two underscores (e.g., `__foo`) are | |
| reserved for use by the runtime and must not be defined by | |
| subcommands. | |
| Args: | |
| parser: An `argparse.ArgumentParser` scoped to this subcommand, | |
| which this function should mutate. | |
| """ | |
| pass | |
| def run(self, flags): | |
| """Execute this subcommand with user-provided flags. | |
| Args: | |
| flags: An `argparse.Namespace` object with all defined flags. | |
| Returns: | |
| An `int` exit code, or `None` as an alias for `0`. | |
| """ | |
| pass | |
| def help(self): | |
| """Short, one-line help text to display on `tensorboard --help`.""" | |
| return None | |
| def description(self): | |
| """Description to display on `tensorboard SUBCOMMAND --help`.""" | |
| return None | |
| class TensorBoardServer(metaclass=ABCMeta): | |
| """Class for customizing TensorBoard WSGI app serving.""" | |
| def __init__(self, wsgi_app, flags): | |
| """Create a flag-configured HTTP server for TensorBoard's WSGI app. | |
| Args: | |
| wsgi_app: The TensorBoard WSGI application to create a server for. | |
| flags: argparse.Namespace instance of TensorBoard flags. | |
| """ | |
| raise NotImplementedError() | |
| def serve_forever(self): | |
| """Blocking call to start serving the TensorBoard server.""" | |
| raise NotImplementedError() | |
| def get_url(self): | |
| """Returns a URL at which this server should be reachable.""" | |
| raise NotImplementedError() | |
| def print_serving_message(self): | |
| """Prints a user-friendly message prior to server start. | |
| This will be called just before `serve_forever`. | |
| """ | |
| sys.stderr.write( | |
| "TensorBoard %s at %s (Press CTRL+C to quit)\n" | |
| % (version.VERSION, self.get_url()) | |
| ) | |
| sys.stderr.flush() | |
| class TensorBoardServerException(Exception): | |
| """Exception raised by TensorBoardServer for user-friendly errors. | |
| Subclasses of TensorBoardServer can raise this exception in order to | |
| generate a clean error message for the user rather than a | |
| stacktrace. | |
| """ | |
| def __init__(self, msg): | |
| self.msg = msg | |
| class TensorBoardPortInUseError(TensorBoardServerException): | |
| """Error raised when attempting to bind to a port that is in use. | |
| This should be raised when it is expected that binding to another | |
| similar port would succeed. It is used as a signal to indicate that | |
| automatic port searching should continue rather than abort. | |
| """ | |
| pass | |
| def with_port_scanning(cls): | |
| """Create a server factory that performs port scanning. | |
| This function returns a callable whose signature matches the | |
| specification of `TensorBoardServer.__init__`, using `cls` as an | |
| underlying implementation. It passes through `flags` unchanged except | |
| in the case that `flags.port is None`, in which case it repeatedly | |
| instantiates the underlying server with new port suggestions. | |
| Args: | |
| cls: A valid implementation of `TensorBoardServer`. This class's | |
| initializer should raise a `TensorBoardPortInUseError` upon | |
| failing to bind to a port when it is expected that binding to | |
| another nearby port might succeed. | |
| The initializer for `cls` will only ever be invoked with `flags` | |
| such that `flags.port is not None`. | |
| Returns: | |
| A function that implements the `__init__` contract of | |
| `TensorBoardServer`. | |
| """ | |
| def init(wsgi_app, flags): | |
| # base_port: what's the first port to which we should try to bind? | |
| # should_scan: if that fails, shall we try additional ports? | |
| # max_attempts: how many ports shall we try? | |
| should_scan = flags.port is None | |
| base_port = ( | |
| core_plugin.DEFAULT_PORT if flags.port is None else flags.port | |
| ) | |
| if base_port > 0xFFFF: | |
| raise TensorBoardServerException( | |
| "TensorBoard cannot bind to port %d > %d" % (base_port, 0xFFFF) | |
| ) | |
| max_attempts = 100 if should_scan else 1 | |
| base_port = min(base_port + max_attempts, 0x10000) - max_attempts | |
| for port in range(base_port, base_port + max_attempts): | |
| subflags = argparse.Namespace(**vars(flags)) | |
| subflags.port = port | |
| try: | |
| return cls(wsgi_app=wsgi_app, flags=subflags) | |
| except TensorBoardPortInUseError: | |
| if not should_scan: | |
| raise | |
| # All attempts failed to bind. | |
| raise TensorBoardServerException( | |
| "TensorBoard could not bind to any port around %s " | |
| "(tried %d times)" % (base_port, max_attempts) | |
| ) | |
| return init | |
| class _WSGIRequestHandler(serving.WSGIRequestHandler): | |
| """Custom subclass of Werkzeug request handler to use HTTP/1.1.""" | |
| # The default on the http.server is HTTP/1.0 for legacy reasons: | |
| # https://docs.python.org/3/library/http.server.html#http.server.BaseHTTPRequestHandler.protocol_version | |
| # Override here to use HTTP/1.1 to avoid needing a new TCP socket and Python | |
| # thread for each HTTP request. The tradeoff is we must always specify the | |
| # Content-Length header, or do chunked encoding for streaming. | |
| protocol_version = "HTTP/1.1" | |
| class WerkzeugServer(serving.ThreadedWSGIServer, TensorBoardServer): | |
| """Implementation of TensorBoardServer using the Werkzeug dev server.""" | |
| # ThreadedWSGIServer handles this in werkzeug 0.12+ but we allow 0.11.x. | |
| daemon_threads = True | |
| def __init__(self, wsgi_app, flags): | |
| self._flags = flags | |
| host = flags.host | |
| port = flags.port | |
| self._auto_wildcard = flags.bind_all | |
| if self._auto_wildcard: | |
| # Serve on all interfaces, and attempt to serve both IPv4 and IPv6 | |
| # traffic through one socket. | |
| host = self._get_wildcard_address(port) | |
| elif host is None: | |
| host = "localhost" | |
| self._host = host | |
| self._url = None # Will be set by get_url() below | |
| self._fix_werkzeug_logging() | |
| def is_port_in_use(port): | |
| with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: | |
| return s.connect_ex(("localhost", port)) == 0 | |
| try: | |
| if is_port_in_use(port): | |
| raise TensorBoardPortInUseError( | |
| "TensorBoard could not bind to port %d, it was already in use" | |
| % port | |
| ) | |
| super().__init__(host, port, wsgi_app, _WSGIRequestHandler) | |
| except socket.error as e: | |
| if hasattr(errno, "EACCES") and e.errno == errno.EACCES: | |
| raise TensorBoardServerException( | |
| "TensorBoard must be run as superuser to bind to port %d" | |
| % port | |
| ) | |
| elif hasattr(errno, "EADDRINUSE") and e.errno == errno.EADDRINUSE: | |
| if port == 0: | |
| raise TensorBoardServerException( | |
| "TensorBoard unable to find any open port" | |
| ) | |
| else: | |
| raise TensorBoardPortInUseError( | |
| "TensorBoard could not bind to port %d, it was already in use" | |
| % port | |
| ) | |
| elif ( | |
| hasattr(errno, "EADDRNOTAVAIL") | |
| and e.errno == errno.EADDRNOTAVAIL | |
| ): | |
| raise TensorBoardServerException( | |
| "TensorBoard could not bind to unavailable address %s" | |
| % host | |
| ) | |
| elif ( | |
| hasattr(errno, "EAFNOSUPPORT") and e.errno == errno.EAFNOSUPPORT | |
| ): | |
| raise TensorBoardServerException( | |
| "Tensorboard could not bind to unsupported address family %s" | |
| % host | |
| ) | |
| # Raise the raw exception if it wasn't identifiable as a user error. | |
| raise | |
| def _get_wildcard_address(self, port): | |
| """Returns a wildcard address for the port in question. | |
| This will attempt to follow the best practice of calling | |
| getaddrinfo() with a null host and AI_PASSIVE to request a | |
| server-side socket wildcard address. If that succeeds, this | |
| returns the first IPv6 address found, or if none, then returns | |
| the first IPv4 address. If that fails, then this returns the | |
| hardcoded address "::" if socket.has_ipv6 is True, else | |
| "0.0.0.0". | |
| """ | |
| fallback_address = "::" if socket.has_ipv6 else "0.0.0.0" | |
| if hasattr(socket, "AI_PASSIVE"): | |
| try: | |
| addrinfos = socket.getaddrinfo( | |
| None, | |
| port, | |
| socket.AF_UNSPEC, | |
| socket.SOCK_STREAM, | |
| socket.IPPROTO_TCP, | |
| socket.AI_PASSIVE, | |
| ) | |
| except socket.gaierror as e: | |
| logger.warning( | |
| "Failed to auto-detect wildcard address, assuming %s: %s", | |
| fallback_address, | |
| str(e), | |
| ) | |
| return fallback_address | |
| addrs_by_family = defaultdict(list) | |
| for family, _, _, _, sockaddr in addrinfos: | |
| # Format of the "sockaddr" socket address varies by address family, | |
| # but [0] is always the IP address portion. | |
| addrs_by_family[family].append(sockaddr[0]) | |
| if hasattr(socket, "AF_INET6") and addrs_by_family[socket.AF_INET6]: | |
| return addrs_by_family[socket.AF_INET6][0] | |
| if hasattr(socket, "AF_INET") and addrs_by_family[socket.AF_INET]: | |
| return addrs_by_family[socket.AF_INET][0] | |
| logger.warning( | |
| "Failed to auto-detect wildcard address, assuming %s", | |
| fallback_address, | |
| ) | |
| return fallback_address | |
| def server_bind(self): | |
| """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) | |
| # 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 | |
| ) | |
| has_v6only_option = hasattr(socket, "IPPROTO_IPV6") and hasattr( | |
| socket, "IPV6_V6ONLY" | |
| ) | |
| if self._auto_wildcard and socket_is_v6 and has_v6only_option: | |
| try: | |
| self.socket.setsockopt( | |
| socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 0 | |
| ) | |
| except socket.error as e: | |
| # Log a warning on failure to dual-bind, except for EAFNOSUPPORT | |
| # since that's expected if IPv4 isn't supported at all (IPv6-only). | |
| if ( | |
| hasattr(errno, "EAFNOSUPPORT") | |
| and e.errno != errno.EAFNOSUPPORT | |
| ): | |
| logger.warning( | |
| "Failed to dual-bind to IPv4 wildcard: %s", str(e) | |
| ) | |
| super().server_bind() | |
| def handle_error(self, request, client_address): | |
| """Override to get rid of noisy EPIPE errors.""" | |
| del request # unused | |
| # Kludge to override a SocketServer.py method so we can get rid of noisy | |
| # EPIPE errors. They're kind of a red herring as far as errors go. For | |
| # example, `curl -N http://localhost:6006/ | head` will cause an EPIPE. | |
| exc_info = sys.exc_info() | |
| e = exc_info[1] | |
| if isinstance(e, IOError) and e.errno == errno.EPIPE: | |
| logger.warning( | |
| "EPIPE caused by %s in HTTP serving" % str(client_address) | |
| ) | |
| else: | |
| logger.error("HTTP serving error", exc_info=exc_info) | |
| def get_url(self): | |
| if not self._url: | |
| if self._auto_wildcard: | |
| display_host = socket.getfqdn() | |
| # Confirm that the connection is open, otherwise change to `localhost` | |
| try: | |
| socket.create_connection( | |
| (display_host, self.server_port), timeout=1 | |
| ) | |
| except socket.error as e: | |
| display_host = "localhost" | |
| else: | |
| host = self._host | |
| display_host = ( | |
| "[%s]" % host | |
| if ":" in host and not host.startswith("[") | |
| else host | |
| ) | |
| self._url = "http://%s:%d%s/" % ( | |
| display_host, | |
| self.server_port, | |
| self._flags.path_prefix.rstrip("/"), | |
| ) | |
| return self._url | |
| def print_serving_message(self): | |
| if self._flags.host is None and not self._flags.bind_all: | |
| sys.stderr.write( | |
| "Serving TensorBoard on localhost; to expose to the network, " | |
| "use a proxy or pass --bind_all\n" | |
| ) | |
| sys.stderr.flush() | |
| super().print_serving_message() | |
| def _fix_werkzeug_logging(self): | |
| """Fix werkzeug logging setup so it inherits TensorBoard's log level. | |
| This addresses a change in werkzeug 0.15.0+ [1] that causes it set its own | |
| log level to INFO regardless of the root logger configuration. We instead | |
| want werkzeug to inherit TensorBoard's root logger log level (set via absl | |
| to WARNING by default). | |
| [1]: https://github.com/pallets/werkzeug/commit/4cf77d25858ff46ac7e9d64ade054bf05b41ce12 | |
| """ | |
| # Log once at DEBUG to force werkzeug to initialize its singleton logger, | |
| # which sets the logger level to INFO it if is unset, and then access that | |
| # object via logging.getLogger('werkzeug') to durably revert the level to | |
| # unset (and thus make messages logged to it inherit the root logger level). | |
| self.log( | |
| "debug", "Fixing werkzeug logger to inherit TensorBoard log level" | |
| ) | |
| logging.getLogger("werkzeug").setLevel(logging.NOTSET) | |
| create_port_scanning_werkzeug_server = with_port_scanning(WerkzeugServer) | |