Add TensorBoard server to Flask endpoint

[lokesh005]

I have a Flask App and a Tensorboad server. Is there a way by which I can map the Tensorboard server to one of the endpoints of Flask so that as soon as I hit that endpoint it triggers the Tensorboard server?

Flask application

from flask import Flask, jsonify, request
app = Flask(__name__)

@app.route('/hello-world', methods=['GET', 'POST'])
def say_hello():
    return jsonify({'result': 'Hello world'})

if __name__ == "__main__":
    app.run(host=host, port=5000)

Tensorboard server code:

from tensorboard.program import TensorBoard, setup_environment
def tensorboard_main(host, port, logdir):
    configuration = list([""])
    configuration.extend(["--host", host])
    configuration.extend(["--port", port])
    configuration.extend(["--logdir", logdir])

    tensorboard = TensorBoard()
    tensorboard.configure(configuration)
    tensorboard.main()

if __name__ == "__main__":
    host = "0.0.0.0"
    port = "7070"
    logdir = '/tmp/logdir'
    tensorboard_main(host, port, logdir)

I tried creating an endpoint in Flask app and then added tensorboard_main(host, port, logdir) in the hope that if I hit the endpoint then the server will start but I got no luck.

I had posted this issue on StackOverflow but didn't get a valid reply hence posting it here.

[wchargin]

Hi @lokesh005! If you’d like to integrate a TensorBoard server into a larger Flask app, the best way to do this is probably to compose the applications at the WSGI level. TensorBoard is a WSGI application internally, so all the necessary pieces of the puzzle are there, but we don’t currently directly expose the WSGI app to clients.

The tensorboard.program API does allow you to pass a custom server class, though, which is essentially a CPS transform of what you want. The signature here is that server_class should be a callable that takes the TensorBoard WSGI app and returns a server that satisfies the TensorBoardServer interface. So you can imagine that you could define something like the following:

import os

import flask
import tensorboard as tb
from werkzeug import serving
from werkzeug.middleware import dispatcher

HOST = "0.0.0.0"
PORT = 7070

flask_app = flask.Flask(__name__)


@flask_app.route("/hello-world", methods=["GET", "POST"])
def say_hello():
    return flask.jsonify({"result": "Hello world"})


class CustomServer(tb.program.TensorBoardServer):
    def __init__(self, tensorboard_app, flags):
        del flags  # unused
        self._app = dispatcher.DispatcherMiddleware(
            flask_app, {"/tensorboard": tensorboard_app}
        )

    def serve_forever(self):
        serving.run_simple(HOST, PORT, self._app)

    def get_url(self):
        return "http://%s:%s" % (HOST, PORT)

    def print_serving_message(self):
        pass  # Werkzeug's `serving.run_simple` handles this


def main():
    program = tb.program.TensorBoard(server_class=CustomServer)
    program.configure(logdir=os.path.expanduser("~/tensorboard_data"))
    program.main()


if __name__ == "__main__":
    main()

Then python main.py starts a server with your Flask app on / and the TensorBoard UI on /tensorboard/ (or /tensorboard; it redirects).

Does this make sense?

IMHO it would be reasonable for tensorboard.program to just directly provide a WSGI application on request; I don’t think that it really breaks the abstraction boundary, and it would give you a bit more flexibility in orchestrating things like this. But in the meantime this workaround may help.

[lokesh005]

Hi @wchargin Thanks for suggesting the possible solution. I tried running the code with correct logdir, but this time it is throwing Failed to load the set of active dashboards. error.

[wchargin]

Could you post the full code that you’re using? The code that I posted works for me with TensorBoard 2.2.2 and Flask 1.1.2:

[lokesh005]

Got it, I was using Tensorboard v2.0.0, due to which I was getting that issue.

[ceevee830]

Hi @wchargin , is flask supported within plugins inside TB v1.15? I'm getting 404 errors. If it is, can you give an example of an endpoint? I'm calling sAjaxSource: './tables/serverside_table', on the client side. I'm not sure if this is wrong, or if I'm not setting up TB side correctly.

Client code

     $('#selector').DataTable({
        bProcessing: true,
        bServerSide: true,
        sPaginationType: 'full_numbers',
        lengthMenu: [[10, 25, 50, 100], [10, 25, 50, 100]],
        bjQueryUI: true,
        sAjaxSource: './tables/serverside_table',
        columns: [
          {data: 'Column A'},
          {data: 'Column B'},
          {data: 'Column C'},
          {data: 'Column D'},

Server code

from flask import Blueprint, jsonify, request
from tensorboard_plugin_dlprof import table_builder
from tensorboard_plugin_dlprof import dlprof_flask_name

tables = Blueprint('tables', 'dlprof', url_prefix='/tables')

@tables.route("/serverside_table", methods=['GET'])
def serverside_table_content():
    data = table_builder.collect_data_serverside(request)
    return jsonify(data)

Server console

W0807  application.py:401] path /data/plugin/dlprof/tables/serverside_table not found, sending 404

Thanks.

[addadda023]

Hi @wchargin, thanks for the excellent starter code! I am reaching out to ask how to inject login required decorator for /tensorboard/ endpoint. The / endpoint leads to login page and it redirects to /tensorboard after successful login. However, since tensorboard is started in main(), I can access /tensorboard without login. How do I enforce login_required decorator during dispatcher.DispatcherMiddleware call?

from flask import Flask, render_template, request, session, redirect, url_for
import flask
import os
import logging
import functools
import tensorboard as tb
from werkzeug import serving
from werkzeug.middleware import dispatcher


# Log transport
logging.basicConfig(level=logging.DEBUG)

HOST = "0.0.0.0"
PORT = 8080
USERS = {"user": ("user", "xxx")}

app = Flask(__name__)
app.secret_key = "yyy"

def login_required(func):
    @functools.wraps(func)
    def secure_function(*args, **kwargs):
        if "username" not in session:
            return redirect(url_for("login", next=request.url))
        return func(*args, **kwargs)

    return secure_function

def login_required(func):
    @functools.wraps(func)
    def secure_function(*args, **kwargs):
        if "username" not in session:
            return redirect(url_for("login", next=request.url))
        return func(*args, **kwargs)

    return secure_function

@app.route("/", methods=["GET", "POST"])
def login():
    if request.method == "POST":
        username = request.form.get("username")
        password = request.form.get("password")

        if username in USERS and USERS[username][1] == password:
            session["username"] = username
            return redirect("http://%s:%s/tensorboard/" % (HOST, PORT))
    return render_template("login.html")

class CustomServer(tb.program.TensorBoardServer):
    def __init__(self, tensorboard_app, flags):
        del flags  # unused
        self._app = dispatcher.DispatcherMiddleware(
            app, {"/tensorboard": tensorboard_app}
        ) # how to include login_required? 

    def serve_forever(self):
        serving.run_simple(HOST, PORT, self._app)

    def get_url(self):
        return "http://%s:%s" % (HOST, PORT)

    def print_serving_message(self):
        pass  # Werkzeug's `serving.run_simple` handles this

def start_tensorboard():
    program = tb.program.TensorBoard(server_class=CustomServer)
    program.configure(logdir="tensorboard_data/")
    program.main()

def main():
    start_tensorboard()


if __name__ == "__main__":
    main()

[eddieparker]

I also appreciate this starter code, it's fantastic for what I need. Is there any way to get this to work with flask's reloader? I tried adding 'use_reloader=True' to the run_simple but then it complains about signal handlers being installed outside the main thread. I can't find anything that indicates that tensorboard has similar reload facilities on file change detection?

[arif-ahmed-nv]

@addadda023 Did you manage to figure out how to use login for the /tensorboard endpoint?