将 Tensorboard 服务器添加到 Flask 端点

Add Tensorboard server to Flask endpoint

我有一个 Flask 应用程序和一个 Tensorboad 服务器。有没有一种方法可以将 Tensorboard 服务器映射到 Flask 的端点之一,以便在我到达该端点时立即触发 Tensorboard 服务器?

Flask 应用程序

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 服务器代码:

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)

我尝试在 Flask 应用程序中创建一个端点,然后添加 tensorboard_main(host, port, logdir) 希望如果我到达端点,服务器就会启动,但我没有成功。

您可以在此处使用多处理:为 flask 创建一个进程,为 tensorboard 创建另一个进程,然后 运行 在同一台主机上。

代码:

from multiprocessing import Process
from tensorboard.program import TensorBoard
from flask import Flask, jsonify

app = Flask(__name__)


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


def flask_main(app, host, port):
    return app.run(host=host, port=port)


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


if __name__ == "__main__":
    host = "0.0.0.0"
    port_for_tensorboard = "7070"
    port_for_flask = "5000"
    logdir = "/tmp/logdir"
    process_for_tensorboard = Process(target=tensorboard_main, args=(host, port_for_tensorboard, logdir))
    process_for_flask = Process(target=flask_main, args=(app, host, port_for_flask))
    process_for_tensorboard.start()
    process_for_flask.start()
    process_for_tensorboard.join()
    process_for_flask.join()

如果你想让flask端点tensorboard内部显示一些东西,那么你需要查看从一个进程到另一个进程的共享数据(认为它会比这个例子更复杂)

我发现要将 TensorBoard 服务器集成到更大的 Flask 应用程序中,最好的方法是在 WSGI 级别组合应用程序。

tensorboard.program API 允许我们通过自定义服务器 class。这里的签名是 server_class 应该是一个可调用的,它采用 TensorBoard WSGI 应用程序和 returns 一个满足 TensorBoardServer 接口的服务器。

因此代码是:

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