EFS as a persistence layer for AWS Batch

Using EFS with Step Functions orchestrated pipelines

Cloud platforms like AWS, GCP, and Azure have certainly made large scale computation easier. The deployment and maintenance of complex compute ecosystems like Spark, Hadoop, and Tensorflow libraries has become standard practice.

However, many legacy systems were not developed in line with the cloud revolution, and single-box deployment is still a common compute architecture. When the compute time for a job starts to spill over to 12+ hours, and the memory requirements for a single task hit the 100+ Gb range, this starts to become a pain-point. Some very powerful software libraries such as Stata (Economics), LS-Dyna (Engineering), and MATSim (Transport Modelling) follow these set of constraints.

In Arup’s City Modelling Lab, we use orchestrate several heavy single-box computational tasks like MATSim, and we need to run workflows around them. We also need to reliably run and persist outputs from these simulations.

We examined several solutions to tackle this problem. Ultimately, we chose AWS Batch for execution, AWS Step Functions for orchestration, and AWS Elastic File System (EFS) for persistence. This article explores why we made these choices and provides a template for how to deploy this solution.

Working with file system constraints

AWS Batch is a tool for resource management and task submission in the cloud. The service is comprised of three essential features

• Autoscaling ECS clusters (compute environments) as infrastructure
• Docker container executions as a computational unit (job)
• Managed queues for task scheduling

The combination of these features makes AWS Batch a strong alternative to traditional resource and task management strategies such as long-running servers, deploying EMR clusters, and other orchestration-as-a-service platforms.

First and foremost, managed job queues and autoscaling compute environments make life easier for an engineer. However, having docker containers as a computational unit is not a complete out-of-the box solution for the workflow management problem.

The main issue with using containers is that they are killed after execution. This means that a persistence layer needs to be defined so that containers can be used in conjunction with each other in workflows.

The basic approach to persistence layers for docker containers is volume mounting. Essentially, one or more directories on the computational environment is mounted when the container is launched. Any results from a container’s execution can then be written to the mounted volume, and thus the host file system.

This approach works well for long-running servers, but a server’s file system can no longer be considered a permanent fixture in the realm of autoscaling and elastic compute. For this reason, we need to consider more involved alternatives.

Cloud persistence solutions

Some potential solutions to the cloud persistence layer problem are

• Use S3 as an intermediary for results
• Use a standalone or cloud-provider database
• Use mounted volumes on long-running servers
• Use mounted EFS volumes to persist results and during execution

S3 as an intermediary

S3 is a reliable persistence layer that can scale to basically infinity. However, repeated S3 copy/puts introduce a lot of latency, especially if done for each task in a workflow. As the number of orchestrated jobs grows, there is also the problem of managing concurrency.

Using a cloud-provider or standalone database

Using a database as your persistence layer is the typical best practice for long running services. This solution will take care of the problems of concurrency and storage. However, in the context of task orchestration, it is not entirely a clear winner. The database solution will still introduce some non-trivial latency for large read/writes to the database when starting and finishing tasks.

Mounted volumes on long-running servers

This solution should be seen as the most sketchy of the bunch. The idea under-pinning elastic computing and autoscaling is that computers fail, and as you scale, computers fail more often. Furthermore, with fixed disk space, a cluster is even more liable to hit its limit and fall over.

Mounted EFS volumes as a persistence layer

AWS Elastic File System (EFS) is an answer to the problem of file system persistence across EC2 instances. EFS volumes are preserved by AWS even after EC2 servers are shut down. So, EFS solves the problem of file-system preservation on instance failure. EFS volumes are mounted to EC2 instances at launch based on a launch template, rather than copied to the container during each execution. This avoids the latency problem of repeated calls that database/S3 solutions introduce.

AWS Batch with mounted EFS Volumes

Now that I have argued for why EFS is a good a persistence layer, the rest of this article will cover how to set up EFS mounting for AWS Batch compute environments.

Set up EFS

Amazon has some introductory material on setting up an EFS volume here. The most important part of this process, for our purposes, is to make sure that security groups are configured properly.

We need to make sure that the EFS volume allows connections from the compute environment’s subnet, and simultaneously ensure the volume and server are sufficiently protected from the outside world. One approach to setting up VPC, subnets and security groups in order to achieve this is covered well in this article.

Build a Launch Template

We will use basic python scripting in order to encode and submit our launch template.

The launch template we will use executes the following MIME file, where region (e.g. eu-west-1) and file system_id (e.g. fs-1234567) have been set accordingly. Save the following as something like efs-mime.yml.

MIME-Version: 1.0
Content-Type: multipart/mixed; boundary="==MYBOUNDARY=="
--==MYBOUNDARY==
Content-Type: text/cloud-config; charset="us-ascii"
packages:
- amazon-efs-utils
runcmd:
- region=my_region
- file_system_id=my_filesystem
- efs_directory=/mnt/efs
- mkdir -p ${efs_directory}
- mount -t nfs4 -o nfsvers=4.1,rsize=1048576,wsize=1048576,hard,timeo=600,retrans=2,noresvport ${file_system_id}.efs.${region}.amazonaws.com:/ ${efs_directory}
--==MYBOUNDARY==--

We can then create and submit a launch template (referenced via template_name) that will be available for EC2 instances with the following script. Note, for some reason, the MIME file needs to be base64 encoded.

#!python
import boto3
'''
Method to create a launch template for ec2 deployments in AWS
'''
def create_launch_template(aws_region, template_name, mime_file):
    with open(mime_file, "rb") as f:
        read_data = f.read()
        user_data = base64.b64encode(read_data)
    launch_template_json = {
        "LaunchTemplateName": template_name,
        "LaunchTemplateData": {
            "UserData": str(user_data)[2:-1]
        }
    }    
    ec2_client = boto3.client('ec2', region_name=aws_region)
    ec2_client.create_launch_template(**launch_template_json)

After running this script for your region and template_name, a launch template should be visible in the launch templates EC2 control panel tab.

Set up a Compute Environment with the launch template

We will launch our compute environment using the boto3 python library. First, we need to do some setup to make sure that our environment will be configured properly. You will need to fill in MY_* fields based on the instance types you want, the security groups you set up for the EFS volume, etc.

Once updated, the following script provides an example of how you can launch a compute environment with EFS mounted and a linked job queue.

#!python
import boto3
import time
environment_name = MY_ENVIRONMENT_NAME
queue_name = MY_QUEUE_NAME
# security info
region = MY_REGION
subnets = MY_SUBNETS
service_role = MY_SERVICE_ROLE
instance_role = MY_INSTANCE_ROLE
security_group_ids = MY_SECURITY_GROUPS
# instance info
instance_types = MY_INSTANCE_TYPES
min_cpu = MIN_CPU
max_cpu = MAX_CPU
desired_cpu = DESIRED_CPU
ec2_key_pair = MY_EC2_KEY_PAIR
launch_template = {
  "launchTemplateName": str(template_name), # template from above
  "version": "$Latest"
}
compute_resources_def = {
    'type': "EC2",
    'minvCpus': min_cpu,
    'maxvCpus': max_cpu,
    'desiredvCpus': desired_cpu,
    'instanceTypes': instance_types,
    'subnets': subnets,
    'instanceRole': instance_role,
    'securityGroupIds': security_group_ids,
    'ec2KeyPair': ec2_key_pair
    'launchTemplate'=launch_template
}
batch_client = boto3.client('batch', region_name=region)
batch_client.create_compute_environment(
    computeEnvironmentName=environment_name,
    type="EC2",
    state="ENABLED",
    computeResources=compute_resources_def,
    serviceRole=service_role
)
# wait for creation
time.sleep(15)
batch_client.create_job_queue(
    jobQueueName=queue_name,
    state="ENABLED",
    priority=1,
    computeEnvironmentOrder=[
        {
            'order': 1,
            'computeEnvironment': environment_name
        },
    ]
)

After running the above script, and assuming that your security groups have been configured properly. A new compute environment and job queue should be visible in the AWS Batch panel. All EC2 instances launched in your new compute environment should have the EFS volume mounted at /mnt/efs .

Try to ssh into a launched instance (if no instances are available, try setting min_vcpu > 0 and at least the number of cores in your smallest instance type in MY_INSTANCE_TYPES). In the future, if you ls /mnt/efs, you should be able to see your EFS files, although this will be empty right now.

Submit a Job with access to your EFS volume

EFS mounting happens at the job registration phase. So, let's assume we have information like the following configured in a JSON file somewhere

{
    "jobDefinition": "test",
    "jobName": "test",
    "image": "test",
    "memory": 6000,
    "vcpus": 2,
    "mount_container_path": "/efs",
    "mount_source_volume": "efs",
    "mount_source_path": "/mnt/efs"
}

The information can then be parsed and passed to set up a job definition that knows to mount the EFS file system volume to the docker image.

#!python
import boto3
job_definition = MY_JOB_DEFINITION
job_name = MY_JOB_NAME
container_memory = MY_MEMORY_REQUIRED
container_vcpu = MY_VCPU_REQUIRED
mount_container_path = "/efs"
mount_source_volume = "efs"
mount_source_path = "/mnt/efs"
region = MY_REGION
command = MY_COMMAND
docker_image = MY_DOCKER_IMAGE
request_kwargs = {
    'jobDefinitionName': job_definition,
    'type': "container",
    'containerProperties': {
        'command': command,
        'image': docker_image,
        'memory': container_memory,
        'vcpus': container_vcpu,
    },
    'retryStrategy': {'attempts': 1}
}
if mount_container_path and mount_source_volume:
    request_kwargs['containerProperties']['volumes'] = [
        {
            'name': mount_source_volume,
            'host': {
                'sourcePath': mount_source_path
            }
        },
    ]
    request_kwargs['containerProperties']['mountPoints'] = [
        {
            'containerPath': mount_container_path,
            'readOnly': False,
            'sourceVolume': mount_source_volume
        }
    ]
batch_client = boto3.client('batch', region_name=region)
batch_client.register_job_definition(**request_kwargs)

After registering the job definition, we can submit jobs based on that definition, which will behave essentially as any other AWS Batch jobs. The EFS volume will be available during execution of these jobs, and we can easily persist data and task results by reading/writing to the mounted volume.

We find this particularly useful for our MATSim simulations. We can now effectively deploy simulations with arbitrary memory requirements, on arbitrarily large networks, and with arbitrarily large populations.

All we need to do is pop the network and populations into EFS, reference the mounted folder in the container’s command, and launch the MATSim container job. After the job finishes and the container is killed, we will still be able to view the results in EFS. Voila!