AWS上にマラソンマッチ用のジャッジ環境を作った

※この記事は包括的な解説というよりは、同じようなことをやろうとした人へのインプットになればいいかなと思っています。C++ソースコード用に書きましたが、少し手を入れれば他の言語でも使えると思います。 AWSを触ったことがない人向けには書いていないです。すいません...

<7/23追記> Lambdaのメモリと処理能力について理解があやふやだったので検証した記録を残しました。メモリは1.8GBくらいにするのが良さそうです。

<8/20追記> 実際にコンテストで使ってみたところ、この構成だと不満が多かったです。それに関するレポートを書きました。

<11/7追記> 実際にプログラムを走らせるLambda(図の2番目のLambda)でプログラムをビルドしていたんですが、毎回ビルドするのは無駄が多いことに気がつきました(最適化オプション付きだとそこそこ時間がかかる)。1番目のLambadでビルドしてバイナリをS3にアップロード&2番目のLambdaでS3からバイナリをダウンロードするとコスト面で少し安くなると思います。 また、最近はStepFunctionsいらないかなという気分になっています。StepFunctionsではLambdaの並列実行を担保していてコードはすっきりするのですが、コストが微妙にかかるのと、処理時間がかかってしまうので、ちょっと微妙かなと思い始めました。並列実行が全部終わることを担保するのはクライアント側の方でやればいいかなと思っています(もしくは全部終わったかを確認するLambdaを作るか)。

<11/13追記> tsutajさんに問題を見つけていただいたのですが、GLIBCのバージョンが古いことでエラーが起きているようです。 コンテナのOSのバージョンを上げるように変更をしました。 月報 (2022 年 10 月) - hogecoder

背景

こんにちは。競技プログラミングのコンテストに参加しているyunixです。 マラソンマッチをしていると以下のような悩みにつき当たります:

  • アルゴリズムの変更が良い効果をもたらしたのかを判断するのには大量のテストケースを実行しなければいけない
  • 一つのテストケースが終了するのには2~5秒くらいかかり、手元の環境で全て実行するのを待つのは結構時間がかかる
    • 手元の環境の性能がAtCoderのジャッジ環境より性能が低いので、実行時間を調整してコンテストの制限時間の2倍くらいにすることが多い
    • そうすると一回の評価に4分くらいかかって大変なストレスになる

ラソンマッチでは色々なアイデアを試すので、素早いフィードバックをもらえる環境が欲しいです。 また、マラソンは根気の勝負みたいな面もあるので、アイデアを出して実装する以外の箇所での精神的な負荷を避けたいです(個人の感想)。

この悩みを解決するために、AWSにマラソンのジャッジ環境を作ったので紹介したいと思います。 作る際に気をつけたのは以下の2点です:

  • 素早くテスト実行をして結果をフィードバックできるようにする
    • 例えば、新しいコードをデプロイするのに5分待つ必要がある、みたいなのでは本末転倒になってしまいます
    • コードの変更をすぐに反映してAWS上で実行できるようにする必要があります
  • 複数のコンテストで簡単に使いまわせるようにする
    • AWS上にジャッジ環境を作るのに時間がかかってしまっても本末転倒です
    • 4時間くらいの短期コンテストでも使用できるくらいにしたい

色々な実現方法はあると思うのですが、以下の図のような感じで環境を作りました。 この記事ではこのシステムを作ったときに考えたことやCDKのコードを簡単に紹介したいと思います。

システムの概要

基本的なアイデアは、

  • 一つのテストケースを実行するのに、一つのLambda関数を立ち上げる。
  • テストケースの実行を並列で行う。並列で実行する際の制御(全ての処理が終わったかの確認など)をするためにAWS Step Functionsを利用する
  • ソースコードはLambdaを実行する際にS3から取得して、Lambdaの中でビルドする。こうすることでS3にあるソースコードを変更することでLambdaの中で実行されるコードを変更することができ、デプロイに時間がかからなくなる。
  • visualizer(というかスコア計算用実行可能ファイル)と入力のファイルはS3にあらかじめ入れておき、実行時にLambdaに取得するようにする。コンテストごとにフォルダを分けておき、Step Functionsの起動時にコンテスト名を指定する
    • スコア計算は公式のvisualizerを使う

という感じです。

コンテストごとに使い回すことができるようになっていて、

  • S3バケットの中にコンテストの名前のフォルダを作る
  • そのフォルダの中にテストコード(in/0000.txt, in/0001.txt....)と出力評価用の実行可能ファイル(vis)を設置する

という作業を行うことで数分程度で新しいコンテストにも対応することができます。

CDKやその他のコード

このシステムを作るのには以下の4つのファイルが必要です

  • AWSリソースを定義したCDKのコード
  • Lambdaの実行環境のDockerfile
  • Lambdaで実行されるPythonのコード
  • 手元からStep Funcrtionsを起動するためのPythonスクリプト

書いたコードをペタペタ貼りつけていこうと思います。 (linter, formatterはめんどくさくて入れていませんでした... ごめんなさい)。

// [lib/cdk-stack.ts]
import { Stack, StackProps, Duration } from 'aws-cdk-lib';
import { Construct } from 'constructs';
import * as s3 from "aws-cdk-lib/aws-s3";
import * as lambda from "aws-cdk-lib/aws-lambda"
import * as sfn from 'aws-cdk-lib/aws-stepfunctions';
import * as tasks from 'aws-cdk-lib/aws-stepfunctions-tasks';

export class MarathonJudgeStack extends Stack {
  constructor(scope: Construct, id: string, props?: StackProps) {
    super(scope, id, props);

    // テストケースやソースコードを保存するためのS3
    const bucket = new s3.Bucket(this, "MarathonBucket", {
      blockPublicAccess: s3.BlockPublicAccess.BLOCK_ALL,
    })
    
    // Step Functionsの1段目になるLambda
    const entryPoint = new lambda.DockerImageFunction(this, "entry-lambda", {
      code: lambda.DockerImageCode.fromImageAsset("lambda", {
        cmd: ["handlers.entry_handler"],
      }),
      timeout: Duration.seconds(30),
      environment: {
        "s3Bucket": bucket.bucketName
      }
    });
    bucket.grantReadWrite(entryPoint)

    // Step Functionsの2段目になるLambda
    const executeTestAndEvaluate = new lambda.DockerImageFunction(this, "evaluation-lambda", {
      code: lambda.DockerImageCode.fromImageAsset("lambda", {
        cmd: ["handlers.exec_handler"],
      }),
      memorySize: 3000,
      timeout: Duration.seconds(30),
      environment: {
        "s3Bucket": bucket.bucketName
      }
    });
    bucket.grantReadWrite(executeTestAndEvaluate)

    // Step Functionsの3段目になるLambda
    const makeSummary = new lambda.DockerImageFunction(this, "summary-lambda", {
      code: lambda.DockerImageCode.fromImageAsset("lambda", {
        cmd: ["handlers.summary_handler"],
      }),
      timeout: Duration.seconds(30),
      environment: {
        "s3Bucket": bucket.bucketName
      }
    });
    bucket.grantReadWrite(makeSummary)

    // AWS Step Functionsの定義
    const stateMachine = new sfn.StateMachine(this, "StateMachine", {
      definition: new tasks.LambdaInvoke(this, "entryLambdaInvokeTask", {
        lambdaFunction: entryPoint,
        resultPath: "$"
      }).next(
        new sfn.Map(this, "mapState", {
          maxConcurrency: 100,
          inputPath: "$.Payload",
          itemsPath: "$.tasks",
          resultPath: "$.Payload.eval_result"
        }).iterator(
          new tasks.LambdaInvoke(this, "exec task", {
            lambdaFunction: executeTestAndEvaluate
          })
        )
      ).next(
        new tasks.LambdaInvoke(this, "summary task", {
          inputPath: "$.Payload.eval_result",
          lambdaFunction: makeSummary,
          resultPath: "$.Payload.summary_result",
        })
      )
    })
  }
}
# Dockerfile
FROM python:3.8-bullseye

RUN apt-get update && \
  apt-get install -y \
  g++ \
  make \
  cmake \
  unzip \
  default-jre \
  default-jdk \
  libcurl4-openssl-dev \
  libopencv-dev && \
  apt-get autoremove -y


COPY handlers.py requirements.txt ./ 

RUN python3.8 -m pip install -r requirements.txt -t .
RUN python3.8 -m pip install boto3
RUN python3.8 -m pip install awslambdaric

ENTRYPOINT ["/usr/local/bin/python", "-m", "awslambdaric"]
CMD ["handlers.entry_handler","handlers.exec_handler", "handlers.summary_handler"]
# handlers.py
import os
import json

import boto3


def entry_handler(event, context):
    timestamp = event["timestamp"]
    contest_name = event["contest_name"]
    N = 50
    result = [
       {
            "test_path": f"{contest_name}/in/{i:04}.txt", 
            "timestamp": timestamp, 
            "contest": contest_name, 
            "id": f"{i:04}"
       } 
       for i in range(0, N, 1)
    ]
    return {"tasks": result}


def exec_handler(event, context):
    print(event, os.environ["s3Bucket"])
    # 前の段からテストファイルのパス・コンテスト名を取得
    test_path = event["test_path"]
    contest_name = event["contest"]
    test_id = event["id"]
    timestamp = event["timestamp"]

    # ソースコードを取得してビルドする
    bucket = boto3.resource("s3").Bucket(os.environ["s3Bucket"])
    bucket.download_file(f"{contest_name}/main.cpp", "/tmp/main.cpp")
    os.system("g++ -O2 /tmp/main.cpp -o /tmp/a.out")

    # スコア計算のスクリプトを取得
    bucket.download_file(f"{contest_name}/vis", "/tmp/vis")
    os.system("chmod 777 /tmp/vis")

    # テストファイルを取得
    bucket.download_file(test_path, "/tmp/test.dat")

    # テストの実行・スコアの評価
    os.system("/tmp/a.out < /tmp/test.dat > /tmp/out.dat")
    os.system("cd /tmp && /tmp/vis /tmp/test.dat /tmp/out.dat > /tmp/score.dat")


    bucket.upload_file("/tmp/out.svg", f"{contest_name}/{timestamp}/out_{test_id}.svg")
    with open("/tmp/score.dat", "r") as f:
        score = f.read()
    

    return {"score": score, "test_file": test_path}


def summary_handler(event, context):
    summary = [result["Payload"] for result in event]
    print(summary)

    return json.dumps(summary)
# 手元から実行するスクリプト
import boto3
import json
import time
# from analysis import * 

session = boto3.Session()

bucket = session.resource("s3").Bucket("バケット名")
bucket.upload_file("./main.cpp", "AHC001/main.cpp")

sfn = session.client("stepfunctions")
response = sfn.start_execution(
    input=json.dumps({"timestamp": time.time(), "contest_name": "AHC001"}),
    stateMachineArn="state machineのARN"
)

while True:
    res = sfn.describe_execution(executionArn=response["executionArn"])
    status = res["status"]
    if status == "RUNNING":
        time.sleep(5)
        continue
    else:
        break

scores = json.loads(json.loads(res["output"])["Payload"]["summary_result"]["Payload"])
print("scores:", scores)
for score in scores:
    print(score)
# compare_with_base_score(scores)

テスト実行結果

上記のPythonスクリプトを実行すると以下のような出力が10秒くらいで得られます(一つWAを出してますね...):

{'score': '938843176\n', 'test_file': 'AHC001/in/0000.txt'}
{'score': '941416888\n', 'test_file': 'AHC001/in/0001.txt'}
{'score': '929691207\n', 'test_file': 'AHC001/in/0002.txt'}
{'score': '935296655\n', 'test_file': 'AHC001/in/0003.txt'}
{'score': '943402555\n', 'test_file': 'AHC001/in/0004.txt'}
{'score': '929585418\n', 'test_file': 'AHC001/in/0005.txt'}
{'score': '904585256\n', 'test_file': 'AHC001/in/0006.txt'}
{'score': '962340546\n', 'test_file': 'AHC001/in/0007.txt'}
{'score': '923769345\n', 'test_file': 'AHC001/in/0008.txt'}
{'score': '0\n', 'test_file': 'AHC001/in/0009.txt'}
{'score': '941221190\n', 'test_file': 'AHC001/in/0010.txt'}
{'score': '944463901\n', 'test_file': 'AHC001/in/0011.txt'}
...
{'score': '947283108\n', 'test_file': 'AHC001/in/0049.txt'}

ここまで来たらあとはスコアの統計を取ったり、別の方法で可視化したり色々やることができます。 4時間のコンテストだったら和を取ったり平均を取ったりするだけかもしれませんが、長期間のコンテストだったらテストケースの傾向ごとに集計するとかの分析をしてもいいかもしれませんね。

コスト

(7/23追記) Lambdaのメモリに関する実験も参照してください。おそらくLambdaのメモリは2GBにするのが最適で、以下の計算のコストの2/3くらいになると思います。

この環境のクラウド料金のほとんどはLambdaの計算時間によるものです。なのでLambdaの実行時間だけを考えればおおむね問題ないと思われます。

現在はテスト実行をするためのLambdaに3GBのメモリを積んでいます。 メモリはこんなにいらないのですが、LambdaはCPU性能とメモリの大きさが比例しているためこうしています。 (実はメモリをどの辺りにするのが本番のジャッジ環境に近くなるのかちゃんと検証していないです。ちゃんと計測しないと...)

AWSの料金計算サイトによると(https://aws.amazon.com/jp/lambda/pricing/)、3GBのメモリを積んでいたとき、1msあたり0.0000000500USDかかるそうです。 おおむね一回の計算に5秒かかって、50個のテストケースが並列で実行されるので、

0.0000000500USD/ms * 5000 ms * 50 = 0.0125ドル

になります。一回の実行で1.5円から2円くらいですね。

調子に乗って実行しまくっていると痛い目を見そうですが、まあこのくらいなら許容範囲かな... 少なくとも手元で実行して待つストレスよりは安いと思います。

他のアイデア

テスト実行は手元から行う必要は必ずしもありません。 Githubに変更をpushしたときにGithub Actionsからテストが実行されて、指定されたissueのところにスコアの統計が記録されるなどしたら面白いかもしれません。

Lambdaのメモリに関する実験など

結論: Lambdaのメモリは1800MBくらいにするのが良さそう。この時AtCoderのコードテストの8割弱くらいの計算回数が得られました(コンパイルオプションとかは見直したほうがいいかも...)。

LambdaのメモリとCPUに関する仕様

上記のコストのところで、Lambdaに割り当てられるCPUパワーはメモリに比例すると書きました。 この記述自体は正しいのですが、ある一定の性能に達するとコアを増やす方向にCPUパワーを上げるようです。 最大の10GBのメモリで6vcpu割り当てられ、並列にスレッドやプロセスを動かすアプリは早くなるよ、と書いてあります。 aws.amazon.com

これを裏付ける情報として、以下のリンクなどを見るとシングルスレッドの処理は1vCPUが割り当てられる1769MBで頭打ちになっています。 www.sentiatechblog.com

Lambdaに割り当てるメモリ量を変えながら実験

以下のことを確かめたいです:

  • Lambdaのメモリはどのくらいがよい?上記のシングルスレッドの計算が頭打ちになるという理解は正しいか?
  • メモリが最適なところで、計算回数的にはAtCoderのジャッジ環境と比べての性能はどうか?

以下のような検証を行いました:

  • AHC012(ケーキを切るやつ)のコードを使って、焼きなましの近傍計算(カット位置の平行移動)を3秒間に何回繰り返せるかを、Lambdaのメモリを変化させながら計算させる
  • 実験に使ったコードはこれです: https://atcoder.jp/contests/ahc012/submissions/33416019
  • また、入力は配布されたテストツールのin/0000.txtを使用しました。

実験結果は以下の図のようになりました:

Lambdaのメモリと焼きなまし回数の関係

Lambdaの仕様から予想されるように、1769MBくらいまではメモリに対して概ね線形で処理回数が増えていそうで、それ以降は頭打ちになっています。 頭打ちになったところではおおよそ640k回近傍計算を行っています。AtCoderのコードテストで同じコードを実行させたら、820k回くらいだったので、概ね8割くらいの性能になっていそうです(コンパイルオプションが一緒ではないのでもしかしたらもうちょっと変わるかも)。

このような事情なので、メモリは1769MBに設定するのが一番コスパが良さそうです。性能的にもAtCoderのジャッジ環境と大きく変わらなくて良さそうです。

不満ポイント

テスト数50~100ケースくらいだったら上記のやり方でも良いのですが、それよりもテストケース数が増えてくるとStep Functionsが以下の点でボトルネックになってしまうようです:

  • step functionsが管理できるpayloadのサイズが256 kBが最大であり、標準エラーなどを返そうとするとテストケースが増えてくると結構辛くなってくる
    • テスト数を増やすとこのlimitに引っかかって実行が失敗してしまう
  • 実行テスト数が増えてくるとstep functionsの完了に時間がかかるようになってしまう

これに対する対応ですが、少し安直ですがクライアント側(python)でなんとかすることにしました。 1つのstate machineの中で実行するLambdaの数がボトルネックになっているので、state machineを並列のプロセスで起動して、それぞれ異なるseedの範囲のテストケース(50から100くらい)を実行させるようにしました。 結果をクライアント側で受け取ってから結合すれば良いという感じです。