tensorbored 2.21.0rc1769983804__py3-none-any.whl

tensorbored/manager.py

@@ -0,0 +1,487 @@
+# Copyright 2019 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.
+# ==============================================================================
+"""Private utilities for managing multiple TensorBoard processes."""
+
+import base64
+import dataclasses
+import datetime
+import errno
+import json
+import os
+import subprocess
+import tempfile
+import time
+import typing
+
+from typing import Optional
+
+from tensorbored import version
+from tensorbored.util import tb_logging
+
+
+@dataclasses.dataclass(frozen=True)
+class TensorBoardInfo:
+    """Holds the information about a running TensorBoard instance.
+
+    Attributes:
+      version: Version of the running TensorBoard.
+      start_time: Seconds since epoch.
+      pid: ID of the process running TensorBoard.
+      port: Port on which TensorBoard is running.
+      path_prefix: Relative prefix to the path, may be empty.
+      logdir: Data location used by the TensorBoard server, may be empty.
+      db: Database connection used by the TensorBoard server, may be empty.
+      cache_key: Opaque, as given by `cache_key` below.
+    """
+
+    version: str
+    start_time: int
+    pid: int
+    port: int
+    path_prefix: str
+    logdir: str
+    db: str
+    cache_key: str
+
+
+def data_source_from_info(info):
+    """Format the data location for the given TensorBoardInfo.
+
+    Args:
+      info: A TensorBoardInfo value.
+
+    Returns:
+      A human-readable string describing the logdir or database connection
+      used by the server: e.g., "logdir /tmp/logs".
+    """
+    if info.db:
+        return "db %s" % info.db
+    else:
+        return "logdir %s" % info.logdir
+
+
+def _info_to_string(info):
+    """Convert a `TensorBoardInfo` to string form to be stored on disk.
+
+    The format returned by this function is opaque and should only be
+    interpreted by `_info_from_string`.
+
+    Args:
+      info: A valid `TensorBoardInfo` object.
+
+    Raises:
+      ValueError: If any field on `info` is not of the correct type.
+
+    Returns:
+      A string representation of the provided `TensorBoardInfo`.
+    """
+    field_name_to_type = typing.get_type_hints(TensorBoardInfo)
+    for key, field_type in field_name_to_type.items():
+        if not isinstance(getattr(info, key), field_type):
+            raise ValueError(
+                "expected %r of type %s, but found: %r"
+                % (key, field_type, getattr(info, key))
+            )
+    if info.version != version.VERSION:
+        raise ValueError(
+            "expected 'version' to be %r, but found: %r"
+            % (version.VERSION, info.version)
+        )
+    json_value = dataclasses.asdict(info)
+    return json.dumps(json_value, sort_keys=True, indent=4)
+
+
+def _info_from_string(info_string):
+    """Parse a `TensorBoardInfo` object from its string representation.
+
+    Args:
+      info_string: A string representation of a `TensorBoardInfo`, as
+        produced by a previous call to `_info_to_string`.
+
+    Returns:
+      A `TensorBoardInfo` value.
+
+    Raises:
+      ValueError: If the provided string is not valid JSON, or if it is
+        missing any required fields, or if any field is of incorrect type.
+    """
+    field_name_to_type = typing.get_type_hints(TensorBoardInfo)
+    try:
+        json_value = json.loads(info_string)
+    except ValueError:
+        raise ValueError("invalid JSON: %r" % (info_string,))
+    if not isinstance(json_value, dict):
+        raise ValueError("not a JSON object: %r" % (json_value,))
+    expected_keys = frozenset(field_name_to_type.keys())
+    actual_keys = frozenset(json_value)
+    missing_keys = expected_keys - actual_keys
+    if missing_keys:
+        raise ValueError(
+            "TensorBoardInfo missing keys: %r" % (sorted(missing_keys),)
+        )
+    # For forward compatibility, silently ignore unknown keys.
+
+    # Validate and deserialize fields.
+    fields = {}
+    for key, field_type in field_name_to_type.items():
+        if not isinstance(json_value[key], field_type):
+            raise ValueError(
+                "expected %r of type %s, but found: %r"
+                % (key, field_type, json_value[key])
+            )
+        fields[key] = json_value[key]
+
+    return TensorBoardInfo(**fields)
+
+
+def cache_key(working_directory, arguments, configure_kwargs):
+    """Compute a `TensorBoardInfo.cache_key` field.
+
+    The format returned by this function is opaque. Clients may only
+    inspect it by comparing it for equality with other results from this
+    function.
+
+    Args:
+      working_directory: The directory from which TensorBoard was launched
+        and relative to which paths like `--logdir` and `--db` are
+        resolved.
+      arguments: The command-line args to TensorBoard, as `sys.argv[1:]`.
+        Should be a list (or tuple), not an unparsed string. If you have a
+        raw shell command, use `shlex.split` before passing it to this
+        function.
+      configure_kwargs: A dictionary of additional argument values to
+        override the textual `arguments`, with the same semantics as in
+        `tensorboard.program.TensorBoard.configure`. May be an empty
+        dictionary.
+
+    Returns:
+      A string such that if two (prospective or actual) TensorBoard
+      invocations have the same cache key then it is safe to use one in
+      place of the other. The converse is not guaranteed: it is often safe
+      to change the order of TensorBoard arguments, or to explicitly set
+      them to their default values, or to move them between `arguments`
+      and `configure_kwargs`, but such invocations may yield distinct
+      cache keys.
+    """
+    if not isinstance(arguments, (list, tuple)):
+        raise TypeError(
+            "'arguments' should be a list of arguments, but found: %r "
+            "(use `shlex.split` if given a string)" % (arguments,)
+        )
+    datum = {
+        "working_directory": working_directory,
+        "arguments": arguments,
+        "configure_kwargs": configure_kwargs,
+    }
+    raw = base64.b64encode(
+        json.dumps(datum, sort_keys=True, separators=(",", ":")).encode("utf-8")
+    )
+    # `raw` is of type `bytes`, even though it only contains ASCII
+    # characters; we want it to be `str` in both Python 2 and 3.
+    return str(raw.decode("ascii"))
+
+
+def _get_info_dir():
+    """Get path to directory in which to store info files.
+
+    The directory returned by this function is "owned" by this module. If
+    the contents of the directory are modified other than via the public
+    functions of this module, subsequent behavior is undefined.
+
+    The directory will be created if it does not exist.
+    """
+    path = os.path.join(tempfile.gettempdir(), ".tensorboard-info")
+    try:
+        os.makedirs(path)
+    except OSError as e:
+        if e.errno == errno.EEXIST and os.path.isdir(path):
+            pass
+        else:
+            raise
+    else:
+        os.chmod(path, 0o777)
+    return path
+
+
+def _get_info_file_path():
+    """Get path to info file for the current process.
+
+    As with `_get_info_dir`, the info directory will be created if it
+    does not exist.
+    """
+    return os.path.join(_get_info_dir(), "pid-%d.info" % os.getpid())
+
+
+def write_info_file(tensorboard_info):
+    """Write TensorBoardInfo to the current process's info file.
+
+    This should be called by `main` once the server is ready. When the
+    server shuts down, `remove_info_file` should be called.
+
+    Args:
+      tensorboard_info: A valid `TensorBoardInfo` object.
+
+    Raises:
+      ValueError: If any field on `info` is not of the correct type.
+    """
+    payload = "%s\n" % _info_to_string(tensorboard_info)
+    with open(_get_info_file_path(), "w") as outfile:
+        outfile.write(payload)
+
+
+def remove_info_file():
+    """Remove the current process's TensorBoardInfo file, if it exists.
+
+    If the file does not exist, no action is taken and no error is
+    raised.
+    """
+    try:
+        os.unlink(_get_info_file_path())
+    except OSError as e:
+        if e.errno == errno.ENOENT:
+            # The user may have wiped their temporary directory or something.
+            # Not a problem: we're already in the state that we want to be in.
+            pass
+        else:
+            raise
+
+
+def get_all():
+    """Return TensorBoardInfo values for running TensorBoard processes.
+
+    This function may not provide a perfect snapshot of the set of running
+    processes. Its result set may be incomplete if the user has cleaned
+    their /tmp/ directory while TensorBoard processes are running. It may
+    contain extraneous entries if TensorBoard processes exited uncleanly
+    (e.g., with SIGKILL or SIGQUIT).
+
+    Entries in the info directory that do not represent valid
+    `TensorBoardInfo` values will be silently ignored.
+
+    Returns:
+      A fresh list of `TensorBoardInfo` objects.
+    """
+    info_dir = _get_info_dir()
+    results = []
+    for filename in os.listdir(info_dir):
+        filepath = os.path.join(info_dir, filename)
+        try:
+            with open(filepath) as infile:
+                contents = infile.read()
+        except IOError as e:
+            if e.errno == errno.EACCES:
+                # May have been written by this module in a process whose
+                # `umask` includes some bits of 0o444.
+                continue
+            else:
+                raise
+        try:
+            info = _info_from_string(contents)
+        except ValueError:
+            # Ignore unrecognized files, logging at debug only.
+            tb_logging.get_logger().debug(
+                "invalid info file: %r",
+                filepath,
+                exc_info=True,
+            )
+        else:
+            results.append(info)
+    return results
+
+
+@dataclasses.dataclass(frozen=True)
+class StartReused:
+    """Possible return value of the `start` function.
+
+    Indicates that a call to `start` was compatible with an existing
+    TensorBoard process, which can be reused according to the provided
+    info.
+
+    Attributes:
+      info: A `TensorBoardInfo` object.
+    """
+
+    info: TensorBoardInfo
+
+
+@dataclasses.dataclass(frozen=True)
+class StartLaunched:
+    """Possible return value of the `start` function.
+
+    Indicates that a call to `start` successfully launched a new
+    TensorBoard process, which is available with the provided info.
+
+    Attributes:
+      info: A `TensorBoardInfo` object.
+    """
+
+    info: TensorBoardInfo
+
+
+@dataclasses.dataclass(frozen=True)
+class StartFailed:
+    """Possible return value of the `start` function.
+
+    Indicates that a call to `start` tried to launch a new TensorBoard
+    instance, but the subprocess exited with the given exit code and
+    output streams. (If the contents of the output streams are no longer
+    available---e.g., because the user has emptied /tmp/---then the
+    corresponding values will be `None`.)
+
+    Attributes:
+      exit_code: As `Popen.returncode` (negative for signal).
+      stdout: Error message to stdout if the stream could not be read.
+      stderr: Error message to stderr if the stream could not be read.
+    """
+
+    exit_code: int
+    stdout: Optional[str]
+    stderr: Optional[str]
+
+
+@dataclasses.dataclass(frozen=True)
+class StartExecFailed:
+    """Possible return value of the `start` function.
+
+    Indicates that a call to `start` failed to invoke the subprocess.
+
+    Attributes:
+      os_error: `OSError` due to `Popen` invocation.
+      explicit_binary: If the TensorBoard executable was chosen via the
+        `TENSORBOARD_BINARY` environment variable, then this field contains
+        the path to that binary; otherwise `None`.
+    """
+
+    os_error: OSError
+    explicit_binary: Optional[str]
+
+
+@dataclasses.dataclass(frozen=True)
+class StartTimedOut:
+    """Possible return value of the `start` function.
+
+    Indicates that a call to `start` launched a TensorBoard process, but
+    that process neither exited nor wrote its info file within the allowed
+    timeout period. The process may still be running under the included
+    PID.
+
+    Attributes:
+      pid: ID of the process running TensorBoard.
+    """
+
+    pid: int
+
+
+def start(arguments, timeout=datetime.timedelta(seconds=60)):
+    """Start a new TensorBoard instance, or reuse a compatible one.
+
+    If the cache key determined by the provided arguments and the current
+    working directory (see `cache_key`) matches the cache key of a running
+    TensorBoard process (see `get_all`), that process will be reused.
+
+    Otherwise, a new TensorBoard process will be spawned with the provided
+    arguments, using the `tensorbored` binary from the system path.
+
+    Args:
+      arguments: List of strings to be passed as arguments to
+        `tensorboard`. (If you have a raw command-line string, see
+        `shlex.split`.)
+      timeout: `datetime.timedelta` object describing how long to wait for
+        the subprocess to initialize a TensorBoard server and write its
+        `TensorBoardInfo` file. If the info file is not written within
+        this time period, `start` will assume that the subprocess is stuck
+        in a bad state, and will give up on waiting for it and return a
+        `StartTimedOut` result. Note that in such a case the subprocess
+        will not be killed. Default value is 60 seconds.
+
+    Returns:
+      A `StartReused`, `StartLaunched`, `StartFailed`, or `StartTimedOut`
+      object.
+    """
+    this_cache_key = cache_key(
+        working_directory=os.getcwd(),
+        arguments=arguments,
+        configure_kwargs={},
+    )
+    match = _find_matching_instance(this_cache_key)
+    if match:
+        return StartReused(info=match)
+
+    stdout_fd, stdout_path = tempfile.mkstemp(prefix=".tensorboard-stdout-")
+    stderr_fd, stderr_path = tempfile.mkstemp(prefix=".tensorboard-stderr-")
+    start_time_seconds = time.time()
+    explicit_tb = os.environ.get("TENSORBOARD_BINARY", None)
+    try:
+        p = subprocess.Popen(
+            ["tensorbored" if explicit_tb is None else explicit_tb] + arguments,
+            stdout=stdout_fd,
+            stderr=stderr_fd,
+        )
+    except OSError as e:
+        return StartExecFailed(os_error=e, explicit_binary=explicit_tb)
+    finally:
+        os.close(stdout_fd)
+        os.close(stderr_fd)
+
+    poll_interval_seconds = 0.5
+    end_time_seconds = start_time_seconds + timeout.total_seconds()
+    while time.time() < end_time_seconds:
+        time.sleep(poll_interval_seconds)
+        subprocess_result = p.poll()
+        if subprocess_result is not None:
+            return StartFailed(
+                exit_code=subprocess_result,
+                stdout=_maybe_read_file(stdout_path),
+                stderr=_maybe_read_file(stderr_path),
+            )
+        info = _find_matching_instance(this_cache_key)
+        if info:
+            # Don't check that `info.pid == p.pid`, since on Windows that may
+            # not be the case: see #4300.
+            return StartLaunched(info=info)
+    else:
+        return StartTimedOut(pid=p.pid)
+
+
+def _find_matching_instance(cache_key):
+    """Find a running TensorBoard instance compatible with the cache key.
+
+    Returns:
+      A `TensorBoardInfo` object, or `None` if none matches the cache key.
+    """
+    infos = get_all()
+    candidates = [info for info in infos if info.cache_key == cache_key]
+    for candidate in sorted(candidates, key=lambda x: x.port):
+        # TODO(@wchargin): Check here that the provided port is still live.
+        return candidate
+    return None
+
+
+def _maybe_read_file(filename):
+    """Read the given file, if it exists.
+
+    Args:
+      filename: A path to a file.
+
+    Returns:
+      A string containing the file contents, or `None` if the file does
+      not exist.
+    """
+    try:
+        with open(filename) as infile:
+            return infile.read()
+    except IOError as e:
+        if e.errno == errno.ENOENT:
+            return None