Script to benchmark time to SSH into an EC2 instance using boto3

boto3_ec2_bench.py

import boto3
import paramiko
import time

# AmzLinux Deep learning AMI: ami-0407d462d1919a19a
# Amazon linux AMI: ami-0c2ab3b8efb09f272

private_key_path = '<PATH_TO_PEM>'   # Set to the path of your private key
aws_keypair = '<KEYPAIR_NAME>'  # Should already exist in your AWS account
ami = 'ami-0407d462d1919a19a'
instance_type = 'm5.xlarge'
region = 'us-west-2'
username = 'ec2-user'

# ============= Launch EC2 instance =============
# Launch an EC2 instance with boto3
ec2 = boto3.resource('ec2', region_name=region)
instance_create_time = time.time()
instance = ec2.create_instances(
    ImageId=ami,
    MinCount=1,
    MaxCount=1,
    InstanceType=instance_type,
    KeyName=aws_keypair
)
instance = instance[0]
instance_id = instance.instance_id


# ======== Wait for instance to spin up ========
print(f'Getting instance {instance_id}')
client = boto3.client('ec2')
ec2 = boto3.resource('ec2', region_name=region)
print('Waiting until instance is running..')
WaiterConfig = {
    'Delay': 2,
    'MaxAttempts': 100
}
waiter = client.get_waiter('instance_running')
waiter.wait(InstanceIds=[instance_id], WaiterConfig=WaiterConfig)
instance_ready_time = time.time()
print('Instance running!')

# ======== Get instance IP address ========
instance.load()
ip_address = instance.public_ip_address
print('Instance public IP: {}'.format(ip_address))

# ======== SSH into instance ========
def ssh_connect_with_retry(ssh, ip_address, username='ec2-user', pkey_path='~/.ssh/id_rsa', retries=0, sleep_duration=5):
    if retries > 100:
        return False
    privkey = paramiko.RSAKey.from_private_key_file(pkey_path)
    try:
        retries += 1
        ssh.connect(hostname=ip_address,
                    username=username, pkey=privkey)
        return True
    except Exception as e:
        print(e)
        time.sleep(sleep_duration)
        print('Retrying SSH connection to {}'.format(ip_address))
        ssh_connect_with_retry(ssh, ip_address, username, pkey_path, retries, sleep_duration)

ssh = paramiko.SSHClient()
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
ssh_connect_with_retry(ssh, ip_address, username=username, pkey_path=private_key_path, retries=0, sleep_duration=2)

# ======== Run commands ========
# run commands
commands = [
    'echo $HOSTNAME'
]
for command in commands:
    print('Running command: {}'.format(command))
    stdin, stdout, stderr = ssh.exec_command(command)
    print(stdout.read().decode('utf-8'))
    print(stderr.read().decode('utf-8'))
ssh_done_time = time.time()

# ========= Cleanup =========
# Close the client connection once the job is done
ssh.close()

# Terminate instance
print('Terminating instance')
instance.terminate()

# ======== Print times ========
print('Total time: {:0.2f}s'.format(ssh_done_time - instance_create_time))
print('Time taken to instance ready: {:0.2f}s'.format(instance_ready_time - instance_create_time))
print('SSH time taken : {:0.2f}s'.format(ssh_done_time - instance_ready_time))