444 lines
14 KiB
Python
444 lines
14 KiB
Python
# 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.
|
|
# ==============================================================================
|
|
"""Utilities for using TensorBoard in notebook contexts, like Colab.
|
|
|
|
These APIs are experimental and subject to change.
|
|
"""
|
|
|
|
|
|
import datetime
|
|
import errno
|
|
import html
|
|
import json
|
|
import os
|
|
import random
|
|
import shlex
|
|
import textwrap
|
|
import time
|
|
|
|
from tensorboard import manager
|
|
|
|
|
|
# Return values for `_get_context` (see that function's docs for
|
|
# details).
|
|
_CONTEXT_COLAB = "_CONTEXT_COLAB"
|
|
_CONTEXT_IPYTHON = "_CONTEXT_IPYTHON"
|
|
_CONTEXT_NONE = "_CONTEXT_NONE"
|
|
|
|
|
|
def _get_context():
|
|
"""Determine the most specific context that we're in.
|
|
|
|
Returns:
|
|
_CONTEXT_COLAB: If in Colab with an IPython notebook context.
|
|
_CONTEXT_IPYTHON: If not in Colab, but we are in an IPython notebook
|
|
context (e.g., from running `jupyter notebook` at the command
|
|
line).
|
|
_CONTEXT_NONE: Otherwise (e.g., by running a Python script at the
|
|
command-line or using the `ipython` interactive shell).
|
|
"""
|
|
# In Colab, the `google.colab` module is available, but the shell
|
|
# returned by `IPython.get_ipython` does not have a `get_trait`
|
|
# method.
|
|
try:
|
|
import google.colab # noqa: F401
|
|
import IPython
|
|
except ImportError:
|
|
pass
|
|
else:
|
|
if IPython.get_ipython() is not None:
|
|
# We'll assume that we're in a Colab notebook context.
|
|
return _CONTEXT_COLAB
|
|
|
|
# In an IPython command line shell or Jupyter notebook, we can
|
|
# directly query whether we're in a notebook context.
|
|
try:
|
|
import IPython
|
|
except ImportError:
|
|
pass
|
|
else:
|
|
ipython = IPython.get_ipython()
|
|
if ipython is not None and ipython.has_trait("kernel"):
|
|
return _CONTEXT_IPYTHON
|
|
|
|
# Otherwise, we're not in a known notebook context.
|
|
return _CONTEXT_NONE
|
|
|
|
|
|
def load_ipython_extension(ipython):
|
|
"""Deprecated: use `%load_ext tensorboard` instead.
|
|
|
|
Raises:
|
|
RuntimeError: Always.
|
|
"""
|
|
raise RuntimeError(
|
|
"Use '%load_ext tensorboard' instead of '%load_ext tensorboard.notebook'."
|
|
)
|
|
|
|
|
|
def _load_ipython_extension(ipython):
|
|
"""Load the TensorBoard notebook extension.
|
|
|
|
Intended to be called from `%load_ext tensorboard`. Do not invoke this
|
|
directly.
|
|
|
|
Args:
|
|
ipython: An `IPython.InteractiveShell` instance.
|
|
"""
|
|
_register_magics(ipython)
|
|
|
|
|
|
def _register_magics(ipython):
|
|
"""Register IPython line/cell magics.
|
|
|
|
Args:
|
|
ipython: An `InteractiveShell` instance.
|
|
"""
|
|
ipython.register_magic_function(
|
|
_start_magic,
|
|
magic_kind="line",
|
|
magic_name="tensorboard",
|
|
)
|
|
|
|
|
|
def _start_magic(line):
|
|
"""Implementation of the `%tensorboard` line magic."""
|
|
return start(line)
|
|
|
|
|
|
def start(args_string):
|
|
"""Launch and display a TensorBoard instance as if at the command line.
|
|
|
|
Args:
|
|
args_string: Command-line arguments to TensorBoard, to be
|
|
interpreted by `shlex.split`: e.g., "--logdir ./logs --port 0".
|
|
Shell metacharacters are not supported: e.g., "--logdir 2>&1" will
|
|
point the logdir at the literal directory named "2>&1".
|
|
"""
|
|
context = _get_context()
|
|
try:
|
|
import IPython
|
|
import IPython.display
|
|
except ImportError:
|
|
IPython = None
|
|
|
|
if context == _CONTEXT_NONE:
|
|
handle = None
|
|
print("Launching TensorBoard...")
|
|
else:
|
|
handle = IPython.display.display(
|
|
IPython.display.Pretty("Launching TensorBoard..."),
|
|
display_id=True,
|
|
)
|
|
|
|
def print_or_update(message):
|
|
if handle is None:
|
|
print(message)
|
|
else:
|
|
handle.update(IPython.display.Pretty(message))
|
|
|
|
parsed_args = shlex.split(args_string, comments=True, posix=True)
|
|
start_result = manager.start(parsed_args)
|
|
|
|
if isinstance(start_result, manager.StartLaunched):
|
|
_display(
|
|
port=start_result.info.port,
|
|
print_message=False,
|
|
display_handle=handle,
|
|
)
|
|
|
|
elif isinstance(start_result, manager.StartReused):
|
|
template = (
|
|
"Reusing TensorBoard on port {port} (pid {pid}), started {delta} ago. "
|
|
"(Use '!kill {pid}' to kill it.)"
|
|
)
|
|
message = template.format(
|
|
port=start_result.info.port,
|
|
pid=start_result.info.pid,
|
|
delta=_time_delta_from_info(start_result.info),
|
|
)
|
|
print_or_update(message)
|
|
_display(
|
|
port=start_result.info.port,
|
|
print_message=False,
|
|
display_handle=None,
|
|
)
|
|
|
|
elif isinstance(start_result, manager.StartFailed):
|
|
|
|
def format_stream(name, value):
|
|
if value == "":
|
|
return ""
|
|
elif value is None:
|
|
return "\n<could not read %s>" % name
|
|
else:
|
|
return "\nContents of %s:\n%s" % (name, value.strip())
|
|
|
|
message = (
|
|
"ERROR: Failed to launch TensorBoard (exited with %d).%s%s"
|
|
% (
|
|
start_result.exit_code,
|
|
format_stream("stderr", start_result.stderr),
|
|
format_stream("stdout", start_result.stdout),
|
|
)
|
|
)
|
|
print_or_update(message)
|
|
|
|
elif isinstance(start_result, manager.StartExecFailed):
|
|
the_tensorboard_binary = (
|
|
"%r (set by the `TENSORBOARD_BINARY` environment variable)"
|
|
% (start_result.explicit_binary,)
|
|
if start_result.explicit_binary is not None
|
|
else "`tensorboard`"
|
|
)
|
|
if start_result.os_error.errno == errno.ENOENT:
|
|
message = (
|
|
"ERROR: Could not find %s. Please ensure that your PATH contains "
|
|
"an executable `tensorboard` program, or explicitly specify the path "
|
|
"to a TensorBoard binary by setting the `TENSORBOARD_BINARY` "
|
|
"environment variable." % (the_tensorboard_binary,)
|
|
)
|
|
else:
|
|
message = "ERROR: Failed to start %s: %s" % (
|
|
the_tensorboard_binary,
|
|
start_result.os_error,
|
|
)
|
|
print_or_update(textwrap.fill(message))
|
|
|
|
elif isinstance(start_result, manager.StartTimedOut):
|
|
message = (
|
|
"ERROR: Timed out waiting for TensorBoard to start. "
|
|
"It may still be running as pid %d." % start_result.pid
|
|
)
|
|
print_or_update(message)
|
|
|
|
else:
|
|
raise TypeError(
|
|
"Unexpected result from `manager.start`: %r.\n"
|
|
"This is a TensorBoard bug; please report it." % start_result
|
|
)
|
|
|
|
|
|
def _time_delta_from_info(info):
|
|
"""Format the elapsed time for the given TensorBoardInfo.
|
|
|
|
Args:
|
|
info: A TensorBoardInfo value.
|
|
|
|
Returns:
|
|
A human-readable string describing the time since the server
|
|
described by `info` started: e.g., "2 days, 0:48:58".
|
|
"""
|
|
delta_seconds = int(time.time()) - info.start_time
|
|
return str(datetime.timedelta(seconds=delta_seconds))
|
|
|
|
|
|
def display(port=None, height=None):
|
|
"""Display a TensorBoard instance already running on this machine.
|
|
|
|
Args:
|
|
port: The port on which the TensorBoard server is listening, as an
|
|
`int`, or `None` to automatically select the most recently
|
|
launched TensorBoard.
|
|
height: The height of the frame into which to render the TensorBoard
|
|
UI, as an `int` number of pixels, or `None` to use a default value
|
|
(currently 800).
|
|
"""
|
|
_display(port=port, height=height, print_message=True, display_handle=None)
|
|
|
|
|
|
def _display(port=None, height=None, print_message=False, display_handle=None):
|
|
"""Internal version of `display`.
|
|
|
|
Args:
|
|
port: As with `display`.
|
|
height: As with `display`.
|
|
print_message: True to print which TensorBoard instance was selected
|
|
for display (if applicable), or False otherwise.
|
|
display_handle: If not None, an IPython display handle into which to
|
|
render TensorBoard.
|
|
"""
|
|
if height is None:
|
|
height = 800
|
|
|
|
if port is None:
|
|
infos = manager.get_all()
|
|
if not infos:
|
|
raise ValueError(
|
|
"Can't display TensorBoard: no known instances running."
|
|
)
|
|
else:
|
|
info = max(manager.get_all(), key=lambda x: x.start_time)
|
|
port = info.port
|
|
else:
|
|
infos = [i for i in manager.get_all() if i.port == port]
|
|
info = max(infos, key=lambda x: x.start_time) if infos else None
|
|
|
|
if print_message:
|
|
if info is not None:
|
|
message = (
|
|
"Selecting TensorBoard with {data_source} "
|
|
"(started {delta} ago; port {port}, pid {pid})."
|
|
).format(
|
|
data_source=manager.data_source_from_info(info),
|
|
delta=_time_delta_from_info(info),
|
|
port=info.port,
|
|
pid=info.pid,
|
|
)
|
|
print(message)
|
|
else:
|
|
# The user explicitly provided a port, and we don't have any
|
|
# additional information. There's nothing useful to say.
|
|
pass
|
|
|
|
fn = {
|
|
_CONTEXT_COLAB: _display_colab,
|
|
_CONTEXT_IPYTHON: _display_ipython,
|
|
_CONTEXT_NONE: _display_cli,
|
|
}[_get_context()]
|
|
return fn(port=port, height=height, display_handle=display_handle)
|
|
|
|
|
|
def _display_colab(port, height, display_handle):
|
|
"""Display a TensorBoard instance in a Colab output frame.
|
|
|
|
The Colab VM is not directly exposed to the network, so the Colab
|
|
runtime provides a service worker tunnel to proxy requests from the
|
|
end user's browser through to servers running on the Colab VM: the
|
|
output frame may issue requests to https://localhost:<port> (HTTPS
|
|
only), which will be forwarded to the specified port on the VM.
|
|
|
|
It does not suffice to create an `iframe` and let the service worker
|
|
redirect its traffic (`<iframe src="https://localhost:6006">`),
|
|
because for security reasons service workers cannot intercept iframe
|
|
traffic. Instead, we manually fetch the TensorBoard index page with an
|
|
XHR in the output frame, and inject the raw HTML into `document.body`.
|
|
|
|
By default, the TensorBoard web app requests resources against
|
|
relative paths, like `./data/logdir`. Within the output frame, these
|
|
requests must instead hit `https://localhost:<port>/data/logdir`. To
|
|
redirect them, we change the document base URI, which transparently
|
|
affects all requests (XHRs and resources alike).
|
|
"""
|
|
import IPython.display
|
|
|
|
shell = """
|
|
(async () => {
|
|
const url = new URL(await google.colab.kernel.proxyPort(%PORT%, {'cache': true}));
|
|
url.searchParams.set('tensorboardColab', 'true');
|
|
const iframe = document.createElement('iframe');
|
|
iframe.src = url;
|
|
iframe.setAttribute('width', '100%');
|
|
iframe.setAttribute('height', '%HEIGHT%');
|
|
iframe.setAttribute('frameborder', 0);
|
|
document.body.appendChild(iframe);
|
|
})();
|
|
"""
|
|
replacements = [
|
|
("%PORT%", "%d" % port),
|
|
("%HEIGHT%", "%d" % height),
|
|
]
|
|
for (k, v) in replacements:
|
|
shell = shell.replace(k, v)
|
|
script = IPython.display.Javascript(shell)
|
|
|
|
if display_handle:
|
|
display_handle.update(script)
|
|
else:
|
|
IPython.display.display(script)
|
|
|
|
|
|
def _display_ipython(port, height, display_handle):
|
|
import IPython.display
|
|
|
|
frame_id = "tensorboard-frame-{:08x}".format(random.getrandbits(64))
|
|
shell = """
|
|
<iframe id="%HTML_ID%" width="100%" height="%HEIGHT%" frameborder="0">
|
|
</iframe>
|
|
<script>
|
|
(function() {
|
|
const frame = document.getElementById(%JSON_ID%);
|
|
const url = new URL(%URL%, window.location);
|
|
const port = %PORT%;
|
|
if (port) {
|
|
url.port = port;
|
|
}
|
|
frame.src = url;
|
|
})();
|
|
</script>
|
|
"""
|
|
proxy_url = os.environ.get("TENSORBOARD_PROXY_URL")
|
|
if proxy_url is not None:
|
|
# Allow %PORT% in $TENSORBOARD_PROXY_URL
|
|
proxy_url = proxy_url.replace("%PORT%", "%d" % port)
|
|
replacements = [
|
|
("%HTML_ID%", html.escape(frame_id, quote=True)),
|
|
("%JSON_ID%", json.dumps(frame_id)),
|
|
("%HEIGHT%", "%d" % height),
|
|
("%PORT%", "0"),
|
|
("%URL%", json.dumps(proxy_url)),
|
|
]
|
|
else:
|
|
replacements = [
|
|
("%HTML_ID%", html.escape(frame_id, quote=True)),
|
|
("%JSON_ID%", json.dumps(frame_id)),
|
|
("%HEIGHT%", "%d" % height),
|
|
("%PORT%", "%d" % port),
|
|
("%URL%", json.dumps("/")),
|
|
]
|
|
|
|
for (k, v) in replacements:
|
|
shell = shell.replace(k, v)
|
|
iframe = IPython.display.HTML(shell)
|
|
if display_handle:
|
|
display_handle.update(iframe)
|
|
else:
|
|
IPython.display.display(iframe)
|
|
|
|
|
|
def _display_cli(port, height, display_handle):
|
|
del height # unused
|
|
del display_handle # unused
|
|
message = "Please visit http://localhost:%d in a web browser." % port
|
|
print(message)
|
|
|
|
|
|
def list():
|
|
"""Print a listing of known running TensorBoard instances.
|
|
|
|
TensorBoard instances that were killed uncleanly (e.g., with SIGKILL
|
|
or SIGQUIT) may appear in this list even if they are no longer
|
|
running. Conversely, this list may be missing some entries if your
|
|
operating system's temporary directory has been cleared since a
|
|
still-running TensorBoard instance started.
|
|
"""
|
|
infos = manager.get_all()
|
|
if not infos:
|
|
print("No known TensorBoard instances running.")
|
|
return
|
|
|
|
print("Known TensorBoard instances:")
|
|
for info in infos:
|
|
template = (
|
|
" - port {port}: {data_source} (started {delta} ago; pid {pid})"
|
|
)
|
|
print(
|
|
template.format(
|
|
port=info.port,
|
|
data_source=manager.data_source_from_info(info),
|
|
delta=_time_delta_from_info(info),
|
|
pid=info.pid,
|
|
)
|
|
)
|