AWS Batch環境を作ってAPV-MCTSのセルフプレイによる強化学習を動かしてみた

AWS Batch APV-MCTS AlphaZero

以前の記事で作成したAlphaZeroもどきオセロのAPV-MCTSのセルフプレイによる強化学習をAWS Batchで動かしてみました。

今回のコード

下記、タグv0.0.1になります。

github.com

環境の概要

下記の２つの環境があります。

  • CodeBuild：heta-reversiのdocker imageを作ってECRに登録する
  • AWS Batch：EC2インスタンスを作成しECRのdocker imageを起動しセルフプレイを実行して結果をS3バケットに格納する

全体構成は下記です。

f:id:nihma:20190126205303p:plain

環境構築の手順

setup_batch.shに記載しました。

大まかに下記の３つの構築を行います。

  • VPC環境
  • ECR関連
  • AWS Batch関連

VPC環境

AWS BatchがEC2インスタンスを作成するためのVPC環境を構築します。

  • VPC
  • Internet Gateway ... AWS Batchから利用するためSubnetをパブリックにするために必要
  • Subnet ... AWS Batchから利用するためEC2インスタンスへのIP自動割り当て設定が必要
  • Route Table
  • Security Group

コードです。

setup_batch.sh: L13-L63

## Create VPC
VPC=$(aws ec2 create-vpc --cidr-block 10.0.0.0/16)
VPC_ID=$(echo ${VPC} | jq -r ".Vpc.VpcId")
aws ec2 create-tags \
    --resources ${VPC_ID} \
    --tags Key=Name,Value=test-batch

## Create Internet Gateway
INTERNET_GATEWAY=$(aws ec2 create-internet-gateway)
INTERNET_GATEWAY_ID=$( \
    echo ${INTERNET_GATEWAY} | jq -r ".InternetGateway.InternetGatewayId")
aws ec2 create-tags \
    --resources ${INTERNET_GATEWAY_ID} \
    --tags Key=Name,Value=test-batch
aws ec2 attach-internet-gateway \
    --internet-gateway-id ${INTERNET_GATEWAY_ID} \
    --vpc-id ${VPC_ID}

## Create Subnet
SUBNET=$( \
    aws ec2 create-subnet \
        --vpc-id ${VPC_ID} \
        --cidr-block 10.0.0.0/24 \
        --availability-zone ap-northeast-1a)
SUBNET_ID=$(echo ${SUBNET} | jq -r ".Subnet.SubnetId")
aws ec2 create-tags \
    --resources ${SUBNET_ID} \
    --tags Key=Name,Value=test-batch
aws ec2 modify-subnet-attribute --subnet-id ${SUBNET_ID} --map-public-ip-on-launch

## Create Route Table
ROUTE_TABLE_ID=$( \
    aws ec2 describe-route-tables \
        --filters Name=vpc-id,Values=${VPC_ID} \
             | jq -r ".RouteTables[].RouteTableId")
aws ec2 create-tags \
    --resources ${ROUTE_TABLE_ID} \
    --tags Key=Name,Value=test-batch
aws ec2 create-route \
    --route-table-id ${ROUTE_TABLE_ID} \
    --destination-cidr-block 0.0.0.0/0 \
    --gateway-id ${INTERNET_GATEWAY_ID}
aws ec2 associate-route-table \
    --route-table-id ${ROUTE_TABLE_ID} \
    --subnet-id ${SUBNET_ID}

## Security Group
DEFAULT_SECURITY_GROUP_ID=$( \
    aws ec2 describe-security-groups \
        --filters Name=group-name,Values=default Name=vpc-id,Values=${VPC_ID} \
            | jq -r '.SecurityGroups[].GroupId')

ECR関連

下記を構築しCodeBuildを実行してECRにdocker imageを登録します。

  • S3 ... セルフプレイ結果を格納するバケット
  • ECR repository
  • IAM ecr build role
  • CodeBuild ecr

コードです。

setup_batch.sh: L65-L155

## Create S3
aws s3 mb s3://test-batch-bucket-name --region ap-northeast-1

## Create ECR repository
ECR_REPO_NAME="test-batch"
ECR_REPO=$(aws ecr create-repository --repository-name ${ECR_REPO_NAME})
ECR_REPO_URL=$(echo ${ECR_REPO} | jq -r ".repository.repositoryUri")

## Create IAM ecr build role
cat <<EOF > Trust-Policy.json
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Principal": {
                "Service": "codebuild.amazonaws.com"
            },
            "Action": "sts:AssumeRole"
        }
    ]
}
EOF
ROLE_ECR_BUILD=$(aws iam create-role --role-name test-batch-build-ecr \
                                     --assume-role-policy-document file://Trust-Policy.json)
aws iam attach-role-policy --policy-arn arn:aws:iam::aws:policy/AdministratorAccess \
                           --role-name test-batch-build-ecr
ROLE_ECR_BUILD_ARN=$(echo ${ROLE_ECR_BUILD} | jq -r ".Role.Arn")

## Create CodeBuild ecr
cat <<EOF > Source.json
{
  "type": "CODECOMMIT",
  "location": "https://git-codecommit.ap-northeast-1.amazonaws.com/v1/repos/heta-reversi",
  "buildspec": "buildspec_ecr.yml"
}
EOF
cat <<EOF > Artifacts.json
{
  "type": "NO_ARTIFACTS"
}
EOF
cat <<EOF > Environment.json
{
  "type": "LINUX_CONTAINER",
  "image": "aws/codebuild/docker:18.09.0",
  "computeType": "BUILD_GENERAL1_SMALL",
  "environmentVariables": [
    {
      "name": "REGION",
      "value": "ap-northeast-1",
      "type": "PLAINTEXT"
    },
    {
      "name": "IMAGE_REPO_NAME",
      "value": "example",
      "type": "PLAINTEXT"
    },
    {
      "name": "IMAGE_TAG",
      "value": "latest",
      "type": "PLAINTEXT"
    },
    {
      "name": "ECR_REPO_URL",
      "value": "${ECR_REPO_URL}",
      "type": "PLAINTEXT"
    }
  ]
}
EOF
aws codebuild create-project --name test-batch-ecr \
                               --source file://Source.json \
                               --artifacts file://Artifacts.json \
                               --environment file://Environment.json \
                               --service-role ${ROLE_ECR_BUILD_ARN}
CODEBUILD_ID=$(aws codebuild start-build --project-name test-batch-ecr --source-version master | tr -d "\n" | jq -r '.build.id')
echo "started.. id is ${CODEBUILD_ID}"
while true
do
  sleep 10s
  STATUS=$(aws codebuild batch-get-builds --ids "${CODEBUILD_ID}" | tr -d "\n" | jq -r '.builds[].buildStatus')
  echo "..status is ${STATUS}."
  if [ "${STATUS}" != "IN_PROGRESS" ]; then
    if [ "${STATUS}" != "SUCCEEDED" ]; then
      echo "faild."
    fi
    echo "done."
    break
  fi
done

AWS Batch関連

AWS Batch環境です。

  • Instance role
  • IAM ecr service role
  • Batch compute environment ... 今回はcpuを0〜4にしてみました
  • Batch job queue
  • IAM job role
  • job definition ... 今回はcpuを4、memoryを2000にしてみました

setup_batch.sh: L157-L264

## Create Instance role
cat <<EOF > Trust-Policy.json
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Principal": {
                "Service": "ec2.amazonaws.com"
            },
            "Action": "sts:AssumeRole"
        }
    ]
}
EOF
aws iam create-role --role-name test-batch-instance \
                    --assume-role-policy-document file://Trust-Policy.json
aws iam attach-role-policy --policy-arn arn:aws:iam::aws:policy/service-role/AmazonEC2ContainerServiceforEC2Role \
                           --role-name test-batch-instance
INSTANCE_ROLE=$(aws iam create-instance-profile --instance-profile-name test-batch-instance)
INSTANCE_ROLE_ARN=$(echo ${INSTANCE_ROLE} | jq -r ".InstanceProfile.Arn")
aws iam add-role-to-instance-profile --role-name test-batch-instance --instance-profile-name test-batch-instance

## Create IAM ecr service role
cat <<EOF > Trust-Policy.json
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Principal": {
                "Service": "batch.amazonaws.com"
            },
            "Action": "sts:AssumeRole"
        }
    ]
}
EOF
ROLE_SERVICE=$(aws iam create-role --role-name test-batch-service \
                                   --assume-role-policy-document file://Trust-Policy.json)
aws iam attach-role-policy --policy-arn arn:aws:iam::aws:policy/service-role/AWSBatchServiceRole \
                           --role-name test-batch-service
ROLE_SERVICE_ARN=$(echo ${ROLE_SERVICE} |jq -r ".Role.Arn")

## Create Batch compute environment
cat << EOF > compute-environment.spec.json
{
    "computeEnvironmentName": "test-compute-environment",
    "type": "MANAGED",
    "state": "ENABLED",
    "computeResources": {
        "type": "EC2",
        "minvCpus": 0,
        "maxvCpus": 4,
        "desiredvCpus": 0,
        "instanceTypes": ["optimal"],
        "subnets": ["${SUBNET_ID}"],
        "securityGroupIds": ["${DEFAULT_SECURITY_GROUP_ID}"],
        "instanceRole": "${INSTANCE_ROLE_ARN}"
    },
    "serviceRole": "${ROLE_SERVICE_ARN}"
}
EOF
COMPUTE_ENV=$(aws batch create-compute-environment --cli-input-json file://compute-environment.spec.json)
COMPUTE_ENV_ARN=$(echo ${COMPUTE_ENV} | jq -r '.computeEnvironmentArn')

## Create Batch job queue
JOB_QUEUE=$(aws batch create-job-queue \
  --job-queue-name test-job-queue \
  --priority 1 \
  --compute-environment-order order=1,computeEnvironment=${COMPUTE_ENV_ARN})
JOB_QUEUE_ARN=$(echo ${JOB_QUEUE} | jq -r '.jobQueueArn')

## Create IAM job role
cat <<EOF > Trust-Policy.json
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Principal": {
                "Service": "ecs-tasks.amazonaws.com"
            },
            "Action": "sts:AssumeRole"
        }
    ]
}
EOF
ROLE_JOB=$(aws iam create-role --role-name test-batch-job \
                               --assume-role-policy-document file://Trust-Policy.json)
aws iam attach-role-policy --policy-arn arn:aws:iam::aws:policy/AmazonS3FullAccess \
                           --role-name test-batch-job
ROLE_JOB_ARN=$(echo ${ROLE_JOB} | jq -r ".Role.Arn")

## Create job definition
cat << EOF > job-definition.spec.json
{
  "image": "${ECR_REPO_URL}",
  "vcpus": 4,
  "memory": 2000,
  "jobRoleArn": "${ROLE_JOB_ARN}"
}
EOF
JOB_DEF=$(aws batch register-job-definition \
  --job-definition-name test-job-definition \
  --type container \
  --container-properties file://job-definition.spec.json)
JOB_DEF_ARN=$(echo ${JOB_DEF} | jq -r '.jobDefinitionArn')

セルフプレイ方法

こちらもsetup_batch.shに記載しました。

JOBを登録してステータスがSUCCEEDEDになったらS3バケットから結果ファイルを取り出します。

setup_batch.sh: L266-L278

## Submit job
JOB=$(aws batch submit-job \
    --job-name "test-job" \
    --job-queue "${JOB_QUEUE_ARN}" \
    --job-definition "${JOB_DEF_ARN}")
JOB_ID=$(echo ${JOB} | jq -r ".jobId")

## Show job status
aws batch describe-jobs --jobs ${JOB_ID} | jq -r ".jobs[].status"

## Get model file
aws s3 ls test-batch-bucket-name/data/
aws s3 sync s3://test-batch-bucket-name/data .

まとめ

今回はAPV-MCTSのセルフプレイによる強化学習を動かすAWS Batch環境を作ってみました。

ただ、マシンパワーが足りず学習が終わりそうにないため今回は下記のように暫定でパラメータをかなり小さな値にしました。*1これでもc4.xlargeで４時間かかりました。

main.py: L16-L24

# FIXME: Stop global variables by modularizing
default_params = {
    'choice_asynchronous_policy_and_value_monte_carlo_tree_search_try_num': 1500,
    'dual_net_trainer_self_play_try_num':                                   2500,
    'dual_net_trainer_create_new_model_epoch_num':                          100,
    'dual_net_trainer_evaluation_try_num':                                  400,
    'dual_net_trainer_evaluation_win_num':                                  220,
    'dual_net_trainer_try_num':                                             100
}

main.py: L713-L722

    elif len(args) > 2 and args[1] == 'create-model-batch':
        bucket_name = args[2]

        # FIXME: Stop global variables by modularizing
        default_params['choice_asynchronous_policy_and_value_monte_carlo_tree_search_try_num'] = 150
        default_params['dual_net_trainer_self_play_try_num']                                   = 25
        default_params['dual_net_trainer_create_new_model_epoch_num']                          = 10
        default_params['dual_net_trainer_evaluation_try_num']                                  = 40
        default_params['dual_net_trainer_evaluation_win_num']                                  = 22
        default_params['dual_net_trainer_try_num']                                             = 1

元のパラメータにて現実的な時間で学習できるようにするためには、nvidia-docker2とGPUインスタンスでGPU計算化したり、複数インスタンスを使った分散化したり、など改善の余地があると思います。

*1:リファクタリングしないと...