Prefect2で、リソース確保・クリーンアップのタスクをwith構文で実行する

Prefectを使って、「EC2インスタンスを立ち上げて、シェルスクリプトを実行し、処理が終わったらシャットダウンする」ようなflowを実行しています。

こちらの記事を書いたPrefectの導入時点では、次のようにtry ~ finally で実行し、「シェルコマンドでエラーが起きた場合でもEC2インスタンスをシャットダウンする」ような処理を実現していました。

from tasks import run_ec2_instance, execute_commands, terminate_ec2_instance

@flow
def sample_flow():
    instance_id = run_ec2_instance(
        image_id="ami-******", # 自作のAMI
        instance_type="t2.micro",
        instance_profile='SSMInstanceRole', # ssmで接続するため、最低限この権限が必要
    )

    try:
        execute_commands(
            instance_id=instance_id,
            commands=["echo 'hello task1'"],
            terminate_after=30,
        )
        ...

    finally:
        terminate_ec2_instance(instance_id)

if __name__ == "__main__":
    sample_flow()

ただ、「How to clean up resources used in a flow?」にある通り、こういうリソースの確保、解放の処理は、Pythonの通常の関数（open等）のようにコンテキストマネージャ（with文）で大幅に簡略化できるようです。（ドキュメントから見つけられずにできないと思いこんでいただけでした）

次のようなコンテキストマネージャーを用意すると、

import contextlib

@contextlib.contextmanager
def setup_ec2_instance(image_id: str, instance_type: str, instance_profile: str):
    """EC2インスタンスのリソースを準備する。

    Args:
        image_id (str): AMIのID
        instance_type (str): インスタンスタイプ
        instance_profile (str): インスタンスプロファイル

    Yields:
        str: 起動したEC2インスタンスのID
    """
    instance_id = run_ec2_instance(image_id=image_id, instance_type=instance_type, instance_profile=instance_profile)
    try:
        yield instance_id
    finally:
        terminate_ec2_instance(instance_id=instance_id)

次のように、いつものPythonで見かける感じにスッキリできます。これで毎回try~finallyを書く必要が無くなりました🙌

from tasks import run_ec2_instance, execute_commands, terminate_ec2_instance

@flow
def sample_flow():
    with setup_ec2_instance(
        image_id="ami-******", # 自作のAMI
        instance_type="t2.micro",
        instance_profile='SSMInstanceRole', # ssmで接続するため、最低限この権限が必要
    ) as instance_id:
        execute_commands(
            instance_id=instance_id,
            commands=["echo 'hello task1'"],
            terminate_after=30,
        )

if __name__ == "__main__":
    sample_flow()

「How to clean up resources used in a flow?」を読むと、Prefect 1の時点ではコンテキストマネージャ用の専用の機能が用意されていたのですが、Pythonの標準的な方法(contextlib)で実装できるように変更されたようですね。