Jupyter・Colabでバックエンドフレームワークを実行するメモ

概要

Jupyterのコードセルでバックエンドフレームワークのコードを書いて、できるだけJupyter内で動作確認できるように試行錯誤したメモ。

成果物

Colabで実行できるようにしました

技術メモ

ColabでサーバーにアクセスできるようにURLを取得

ローカルのJupyterの場合はlocalhostで起動したサーバーに接続できますが、Colabではそうはいかないのでサーバーに接続するためのURLを取得する必要があります。

Colabではgoogle.colab.kernel.proxyPortというJavaScript APIでポートを指定すればそのポートで立てたサーバーにアクセスするためのURLが取得できます。

# ColabのサーバーURLを取得(Colab以外の環境ではローカルホストを返す)

server_port = 5000  # ポート番号

try:
    from google.colab.output import eval_js  # type: ignore
    server_url = eval_js(f"google.colab.kernel.proxyPort({server_port})").strip("/")
except ImportError:
    server_url = f"http://localhost:{server_port}"

サーバーの開始、再起動

URLさえ取得できればサーバーを起動するのは簡単なのですが(とはいえ普通に起動すると強制終了するまでブロックして他のコードセルが実行できなくなるので工夫は必要ですが)、FlaskやFastAPIではそのサーバーを終了するための関数が用意されていません。どうもWSGIの仕様？らしいですが、Ctrl+Cを送るなりPythonプロセス自体を終わらせるなりしろという感じのようです。

JupyterやColabでPythonプロセス自体を終わらせるとIPythonカーネルも終了してしまうのでおいそれとPythonプロセスをKillすることはできません。

multiprocessingモジュールでプロセスを作ってterminateで終わらせる方法もあるようですが、Jupyterではなぜかうまく動きませんでした…

というわけで、仕方なく%%writefileマジックコマンドでコードセルに書いたFlaskのコードをPythonファイルに出力し、Pythonコマンドで実行してサーバーを立てるという最後の手段を使うしかありませんでした…

%%writefile flask_app.py
from flask import Flask

app = Flask(__name__)

@app.route("/")
def hello():
    return "Hello World!"

app.run()

とりあえずJupyterとは別のPythonプロセスでサーバーを立てたのでPythonプロセスを終了させてサーバーを終了させることはできるようになりましたが、プロセスの終わらせ方がWindowsとColab(Linux)で違ってたりするのでさらなる試行錯誤を強いられる羽目になりましたが、とりあえず以下のような感じで終了させることができました。もっとスマートな方法があればいいのですがなんとも…

import time
import subprocess
import sys
import signal

# IPythonのユーザー名前空間を取得
from IPython import get_ipython  # type: ignore
ipython = get_ipython()
userns = ipython.user_ns

# プラットフォームによってシグナルの種類とPythonコマンドを変更
if sys.platform == "win32":
    termsignal = signal.CTRL_C_EVENT
    py_name = "python"
else:
    termsignal = signal.SIGTERM
    py_name = "python3"

# 以前に起動したサーバーを停止
if "server_process" in userns:
    server_process = userns["server_process"]
    if server_process.poll() is None:
        print("Stopping the previous server...")
        try:
            server_process.send_signal(termsignal)
            try:
                server_process.wait(10)
            except subprocess.TimeoutExpired:
                server_process.kill()
                time.sleep(5)
        except:
            pass

# 前のセルで書き出したflask_app.pyを実行

# Colab環境ではpopenの引数にshell=Trueを指定するとうまくいかない？
try:
    import google.colab  # type: ignore
    shell_flag = False
except ImportError:
    shell_flag = True

# サーバーを起動
print("Starting the server...")
server_process = subprocess.Popen([py_name, "flask_app.py"], shell=shell_flag, text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

time.sleep(5)

if server_process.poll() is not None:
    print("Failed to start the server.")
    print(server_process.stdout.read())
    print(server_process.stderr.read())
else:
    print("The server is running.")
    print(f"URL: {server_url}")

これでFlaskコードのセルを修正して実行したあとにこのセルを実行すれば今動いているサーバーを終わらせて修正したコードでサーバーを再起動してくれるようになりました。

モジュール作ってみた

以上のコードをもう少し改良してモジュール化しました、%pipでインストールしてマジックコマンドを先頭に設定するだけでコードセルのコードをサーバーとして起動できるようになっています。

以下で動作の確認ができます。