batch.sh

#!/bin/bash

set -e

REPOSITORY=$(basename $(realpath $(dirname "$0")))
VCPUS=16
MEMORY=$(expr 1024 \* 32)
COMMAND='[]'
JOB_QUEUE=default

export AWS_PROFILE=al2
export AWS_REGION=us-east-1

SUBMIT=0
while :; do
    case $1 in
        -s|--submit)
            SUBMIT=1
            ;;

        --)
            shift
            break
            ;;

        -?*)
            error "Unknown option ${1}"
            echo ""
            usage
            ;;

        *)
            break
    esac
    shift
done

function get-or-create-repository-uri() {
	REPOSITORY="$1"
	# TODO: make this stop outputing an error on first run
	REPOSITORY_URI=$(aws ecr describe-repositories --repository-names $REPOSITORY --query "(repositories[?repositoryName==\`$REPOSITORY\`])[0].repositoryUri" --output text || echo "error")
	if [ $REPOSITORY_URI = "error" ]; then
		aws ecr create-repository --repository-name $REPOSITORY &> /dev/null
		REPOSITORY_URI=$(aws ecr describe-repositories --repository-names $REPOSITORY --query "(repositories[?repositoryName==\`$REPOSITORY\`])[0].repositoryUri" --output text 2> /dev/null)
	fi
	echo "$REPOSITORY_URI"
}

function submit-job() {
	REPOSITORY=$1
	JOB_QUEUE=$2
	aws batch submit-job --job-name $REPOSITORY --job-definition $REPOSITORY --job-queue $JOB_QUEUE --query 'jobId' --output text
}

function get-job-status() {
	JOB_ID="$1"
	aws batch describe-jobs --jobs $JOB_ID --query "(jobs[?jobId==\`$JOB_ID\`])[0].status" --output text
}

function get-job-log-stream() {
	JOB_ID="$1"
	aws batch describe-jobs --jobs $JOB_ID --query "(jobs[?jobId==\`$JOB_ID\`])[0].container.logStreamName" --output text
}

echo "==> Logging Docker into ECR"
$(aws ecr get-login --no-include-email)

echo "==> Getting or Creating ECR Repository"
REPOSITORY_URI=$(get-or-create-repository-uri "$REPOSITORY")
echo "$REPOSITORY_URI"

echo "==> Building Docker Image"
docker build . -t $REPOSITORY

echo "==> Tagging Docker Image"
docker tag $REPOSITORY:latest $REPOSITORY_URI

echo "==> Pushing Docker Image"
docker push $REPOSITORY_URI

CONTAINER_PROPERTIES=$(jq -n \
	--arg image "$REPOSITORY_URI" \
	--argjson memory $MEMORY \
	--argjson vcpus $VCPUS \
	--argjson command "$COMMAND" \
	'{"image": $image, "vcpus": $vcpus, "command": $command, "memory": $memory}')

echo "Registering Job Definition ($REPOSITORY) with Container Properties:"
echo "$CONTAINER_PROPERTIES"

aws batch register-job-definition \
		--job-definition-name $REPOSITORY \
		--type container \
		--container-properties "$CONTAINER_PROPERTIES"

if [ $SUBMIT -eq 0 ]; then

	echo -e "\n\nSubmit the job by running:"
	echo -e "\naws batch submit-job --job-name $REPOSITORY --job-definition $REPOSITORY --job-queue $JOB_QUEUE"

else

	# create job queue

	JOB_ID=$(submit-job $REPOSITORY $JOB_QUEUE)
	echo "Started job with id: $JOB_ID"

	echo "Waiting until job has begun"
	STATUS=$(get-job-status $JOB_ID)
	while [ $STATUS == "SUBMITTED" ] || [ $STATUS == "RUNNABLE" ]; do
	    printf '.'
	    STATUS=$(get-job-status $JOB_ID)
	    sleep 5
	done
	printf '\n'

	STATUS=$(get-job-status $JOB_ID)
	echo "Job Status: $STATUS"

	LOG_GROUP=/aws/batch/job
	LOG_STREAM=$(get-job-log-stream $JOB_ID)

	echo "Tailing Cloudwatch Logs: $LOG_GROUP $LOG_STREAM"

	awslogs get $LOG_GROUP $LOG_STREAM -s 24hr --watch --no-group --no-stream --timestamp

fi

setup-aws-batch.sh

#!/bin/sh

set -e

export AWS_PROFILE=al2
export AWS_REGION=us-east-1

aws emr create-default-roles

COMPUTE_ENVIRONMENT_ARN=$(aws batch describe-compute-environments --compute-environments default --query 'computeEnvironments[0].computeEnvironmentArn' --output text)

if [ $COMPUTE_ENVIRONMENT_ARN = "None" ]; then

	ACCOUNT_ID=$(aws sts get-caller-identity --query 'Account' --output text)
	VPC_ID=$(aws ec2 describe-vpcs --query '(Vpcs[?IsDefault==`true`])[0].VpcId' --output text)
	SECURITY_GROUP=$(aws ec2 describe-security-groups --group-names default --query 'SecurityGroups[0].GroupId')
	SUBNET_IDS=$(aws ec2 describe-subnets --query "(Subnets[?VpcId==\`$VPC_ID\`]).SubnetId" | jq -Mc)

AssumeRolePolicyDocument=$(cat <<END_HEREDOC
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Principal": {
                "Service": "batch.amazonaws.com"
            },
            "Action": "sts:AssumeRole"
        }
    ]
}
END_HEREDOC
)

	SERVICE_ROLE=$(aws iam list-roles --query '(Roles[?RoleName==`AWSBatchServiceRole`])[0].Arn' --output text)
	if [ $SERVICE_ROLE = "None" ]; then
		echo "Creating AWSBatchServiceRole..."
		aws iam create-role --role-name AWSBatchServiceRole --assume-role-policy-document "$AssumeRolePolicyDocument"
		aws iam attach-role-policy --role-name AWSBatchServiceRole --policy-arn arn:aws:iam::aws:policy/service-role/AWSBatchServiceRole
	fi

config=$(cat <<END_HEREDOC
{
	"computeEnvironmentName": "default",
	"type": "MANAGED",
	"state": "ENABLED",
	"computeResources": {
		"type": "EC2",
		"minvCpus": 0,
		"maxvCpus": 256,
		"desiredvCpus": 2,
		"instanceTypes": ["optimal"],
		"subnets": $SUBNET_IDS,
		"securityGroupIds": [
			$SECURITY_GROUP
		],
		"ec2KeyPair": "default",
		"instanceRole": "arn:aws:iam::$ACCOUNT_ID:instance-profile/EMR_EC2_DefaultRole"
	},
	"serviceRole": "$SERVICE_ROLE"
}
END_HEREDOC
)

	echo "Creating Compute Environment with Config:"
	echo "$config"

	aws batch create-compute-environment --cli-input-json "$config"

fi

JOB_QUEUE_ARN=$(aws batch describe-job-queues --job-queues default --query 'jobQueues[0].jobQueueArn' --output text)

if [ $JOB_QUEUE_ARN = "None" ]; then

config=$(cat <<'END_HEREDOC'
{
    "jobQueueName": "default",
    "state": "ENABLED",
    "priority": 10,
    "computeEnvironmentOrder": [
        {
            "order": 10,
            "computeEnvironment": "default"
        }
    ]
}
END_HEREDOC
)

	echo "Creating Job Queue with Config:"
	echo "$config"

	aws batch create-job-queue --cli-input-json "$config"

fi