Add Tensorboard server to Flask endpoint

Question

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.

Answer 1

I found out that to integrate a TensorBoard server into a larger Flask app, the best way would be to compose the applications at the WSGI level.

The tensorboard.program API allows us to pass a custom server class. The signature here is that server_class should be a callable which takes the TensorBoard WSGI app and returns a server that satisfies the TensorBoardServer interface.

Hence the code is:

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()