Skip to content

Instantly share code, notes, and snippets.

@romilbhardwaj
Created September 6, 2022 17:53
Show Gist options
  • Save romilbhardwaj/e94ce5b5006e00d4f62d81b6df553058 to your computer and use it in GitHub Desktop.
Save romilbhardwaj/e94ce5b5006e00d4f62d81b6df553058 to your computer and use it in GitHub Desktop.
Script to benchmark time to SSH into an EC2 instance using boto3
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))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment