Skip to content

Instantly share code, notes, and snippets.

@soumith
Created January 13, 2021 04:09
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save soumith/3b5c65443cbd078b3dcd9393f1af6b1e to your computer and use it in GitHub Desktop.
Save soumith/3b5c65443cbd078b3dcd9393f1af6b1e to your computer and use it in GitHub Desktop.
import os
import threading
from flask import Flask
import socketio
from flask_cors import cross_origin, CORS
import dlevent
import logging
import json
import random
try:
import html
html_escape = html.escape
del html
except ImportError:
import cgi
html_escape = cgi.escape
del cgi
app = None
def _dashboard_thread(web_root, ip, port, quiet=False):
global app
root_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "../"))
static_folder = os.path.join(root_dir, web_root, "build")
if not quiet:
print("static_folder:", static_folder)
app = Flask(__name__, static_folder=static_folder, static_url_path="")
sio = socketio.Server(async_mode="threading", cors_allowed_origins="*")
app.wsgi_app = socketio.WSGIApp(sio, app.wsgi_app)
dlevent.sio = sio
CORS(app, resources={r"*": {"origins": "*"}})
if quiet:
log = logging.getLogger('werkzeug')
log.setLevel(logging.ERROR)
log.disabled = True
log = logging.getLogger('socketio.server')
log.setLevel(logging.ERROR)
log.disabled = True
app.logger.disabled = True
@app.route("/")
@cross_origin(origin="*")
def index():
return app.send_static_file("index.html")
if os.getenv("MCDASHBOARD_PORT"):
port = os.getenv("MCDASHBOARD_PORT")
print("setting MC dashboard port from env variable MCDASHBOARD_PORT={}".format(port))
if os.getenv("MCDASHBOARD_IP"):
ip = os.getenv("MCDASHBOARD_IP")
print("setting MC dashboard ip from env variable MCDASHBOARD_IP={}".format(ip))
app.run(ip, threaded=True, port=port, debug=False)
def start(web_root="tools/dashboard_web", ip="0.0.0.0", port=8000):
t = threading.Thread(target=_dashboard_thread, args=(web_root, ip, port, True))
t.start()
# Return values for `_get_context` (see that function's docs for
# details).
_CONTEXT_COLAB = "_CONTEXT_COLAB"
_CONTEXT_IPYTHON = "_CONTEXT_IPYTHON"
_CONTEXT_NONE = "_CONTEXT_NONE"
# this function is generously pulled from https://github.com/tensorflow/tensorboard/blob/d7fc1eab51ea6d0fb60e22eac7b271b4abc076e8/tensorboard/notebook.py#L53
# it is licensed under Apache 2.0 and Copyright 2017 The TensorFlow Authors. All rights reserved.
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 display(port=8000, height=None, print_message=False, display_handle=None):
if height is None:
height = 800
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_ipython(port, height, display_handle):
import IPython.display
frame_id = "dldashboard-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>
"""
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)
# 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.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import datetime
import errno
import json
import os
import random
import shlex
import textwrap
import time
try:
import html
html_escape = html.escape
del html
except ImportError:
import cgi
html_escape = cgi.escape
del cgi
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,
)
)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment