Distributed TensorFlow Training on AWS EFS

Question

Pulumi · Accepted Answer

To accomplish distributed TensorFlow training on AWS EFS (Elastic File System), we'll need to set up an environment where multiple machines can access a shared file system to read data and write checkpoints while training a machine learning model. AWS EFS is a fully-managed file storage service for applications that require a file system interface and file system semantics.

Here's how we can achieve this with Pulumi:

1. **Create an EFS File System**: This will serve as the shared storage for all the machines running the TensorFlow training jobs.
2. **Create EFS Mount Targets**: For each Availability Zone, create a mount target so the EC2 instances in those zones can access the EFS File System.
3. **Provision EC2 Instances**: Launch multiple EC2 instances that will perform the distributed training. Ensure these instances are in the same VPC (Virtual Private Cloud) and subnets as the EFS Mount Targets.
4. **Configure Security Groups**: Update the security group settings to allow traffic between the EC2 instances and the mount targets on the necessary ports.
5. **Set up TensorFlow**: This isn't directly related to Pulumi, but you'll need to set up TensorFlow for distributed training on each EC2 instance.

We'll use the `aws.efs.FileSystem` and `aws.efs.MountTarget` resources from the AWS package to create the EFS system and mount targets. Then, we'll create EC2 instances using `aws.ec2.Instance`.

Below is a Pulumi program in Python that sets up the EFS and EC2 instances. Note that the program does not include setting up TensorFlow itself:

```python
import pulumi
import pulumi_aws as aws

# Create an EFS File System for shared storage
efs_file_system = aws.efs.FileSystem("distributed-training-fs",
                                     performance_mode="generalPurpose",
                                     throughput_mode="bursting")

# Create EFS Mount Targets for all necessary subnets.
# Replace `subnet_ids` with the actual subnet IDs you're using.
subnet_ids = ["subnet-0bb1c79de3EXAMPLE", "subnet-0bb2c79de4EXAMPLE"]  # Example subnet IDs
security_group = aws.ec2.SecurityGroup("efs-sg",
                                       description="Allow EFS access",
                                       ingress=[{
                                           'protocol': 'tcp',
                                           'from_port': 2049,
                                           'to_port': 2049,
                                           'cidr_blocks': ['0.0.0.0/0']  # Update to restrict access
                                       }])

mount_targets = []
for subnet_id in subnet_ids:
    mount_target = aws.efs.MountTarget(f"efs-mt-{subnet_id}",
                                       file_system_id=efs_file_system.id,
                                       subnet_id=subnet_id,
                                       security_groups=[security_group.id])
    mount_targets.append(mount_target)

# Create multiple EC2 Instances to perform the distributed training
# Replace `ami` with the actual AMI ID you're using and ensure the instance has the necessary specs.
num_instances = 3
training_instances = []
for i in range(num_instances):
    instance = aws.ec2.Instance(f"training-instance-{i}",
                                instance_type="t3.large",  # Placeholder instance type
                                ami="ami-0c55b159cbfafe1f0",  # Placeholder Amazon Machine Image
                                subnet_id=subnet_ids[i % len(subnet_ids)],  # Round-robin over available subnets
                                security_groups=[security_group.id])
    training_instances.append(instance)

# Export the DNS names of the EFS Mount Targets and the public IPs of the EC2 Instances
pulumi.export('efs_mount_target_dns_names', pulumi.Output.all(*mount_targets).apply(
    lambda targets: [target.dns_name for target in targets]
))
pulumi.export('ec2_instance_public_ips', pulumi.Output.all(*training_instances).apply(
    lambda instances: [instance.public_ip for instance in instances]
))
```

This program sets up an EFS file system and mount targets so that it's accessible across multiple availability zones for high availability and fault tolerance. Security groups are configured to allow traffic on TCP port 2049, which is the NFS port used by EFS. It's suggested that you replace `'0.0.0.0/0'` with more restrictive CIDR blocks for security purposes. EC2 instances are provisioned and associated with the security groups to allow them to access the mount targets.

Remember that you need to install the Pulumi AWS plugin and configure AWS access credentials in your environment before running this program. The EC2 instances require manual setup for TensorFlow after the infrastructure is provisioned by this Pulumi program. Ensure the IAM roles and policies attached to the EC2 instances allow them to access the EFS file system.