Skip to content

Commit

Permalink
Update aws image
Browse files Browse the repository at this point in the history
  • Loading branch information
@asonnino
asonnino committed Feb 8, 2024
1 parent ec74756 commit 7e705f6
Showing 1 changed file with 129 additions and 116 deletions.
245 changes: 129 additions & 116 deletions benchmark/benchmark/instance.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,8 @@
class AWSError(Exception):
def __init__(self, error):
assert isinstance(error, ClientError)
self.message = error.response['Error']['Message']
self.code = error.response['Error']['Code']
self.message = error.response["Error"]["Message"]
self.code = error.response["Error"]["Code"]
super().__init__(self.message)


Expand All @@ -21,14 +21,14 @@ def __init__(self, settings):
self.settings = settings
self.clients = OrderedDict()
for region in settings.aws_regions:
self.clients[region] = boto3.client('ec2', region_name=region)
self.clients[region] = boto3.client("ec2", region_name=region)

@classmethod
def make(cls, settings_file='settings.json'):
def make(cls, settings_file="settings.json"):
try:
return cls(Settings.load(settings_file))
except SettingsError as e:
raise BenchError('Failed to load settings', e)
raise BenchError("Failed to load settings", e)

def _get(self, state):
# Possible states are: 'pending', 'running', 'shutting-down',
Expand All @@ -37,21 +37,15 @@ def _get(self, state):
for region, client in self.clients.items():
r = client.describe_instances(
Filters=[
{
'Name': 'tag:Name',
'Values': [self.settings.testbed]
},
{
'Name': 'instance-state-name',
'Values': state
}
{"Name": "tag:Name", "Values": [self.settings.testbed]},
{"Name": "instance-state-name", "Values": state},
]
)
instances = [y for x in r['Reservations'] for y in x['Instances']]
instances = [y for x in r["Reservations"] for y in x["Instances"]]
for x in instances:
ids[region] += [x['InstanceId']]
if 'PublicIpAddress' in x:
ips[region] += [x['PublicIpAddress']]
ids[region] += [x["InstanceId"]]
if "PublicIpAddress" in x:
ips[region] += [x["PublicIpAddress"]]
return ids, ips

def _wait(self, state):
Expand All @@ -65,77 +59,97 @@ def _wait(self, state):

def _create_security_group(self, client):
client.create_security_group(
Description='HotStuff node',
Description="HotStuff node",
GroupName=self.settings.testbed,
)

client.authorize_security_group_ingress(
GroupName=self.settings.testbed,
IpPermissions=[
{
'IpProtocol': 'tcp',
'FromPort': 22,
'ToPort': 22,
'IpRanges': [{
'CidrIp': '0.0.0.0/0',
'Description': 'Debug SSH access',
}],
'Ipv6Ranges': [{
'CidrIpv6': '::/0',
'Description': 'Debug SSH access',
}],
"IpProtocol": "tcp",
"FromPort": 22,
"ToPort": 22,
"IpRanges": [
{
"CidrIp": "0.0.0.0/0",
"Description": "Debug SSH access",
}
],
"Ipv6Ranges": [
{
"CidrIpv6": "::/0",
"Description": "Debug SSH access",
}
],
},
{
'IpProtocol': 'tcp',
'FromPort': self.settings.consensus_port,
'ToPort': self.settings.consensus_port,
'IpRanges': [{
'CidrIp': '0.0.0.0/0',
'Description': 'Consensus port',
}],
'Ipv6Ranges': [{
'CidrIpv6': '::/0',
'Description': 'Consensus port',
}],
"IpProtocol": "tcp",
"FromPort": self.settings.consensus_port,
"ToPort": self.settings.consensus_port,
"IpRanges": [
{
"CidrIp": "0.0.0.0/0",
"Description": "Consensus port",
}
],
"Ipv6Ranges": [
{
"CidrIpv6": "::/0",
"Description": "Consensus port",
}
],
},
{
'IpProtocol': 'tcp',
'FromPort': self.settings.mempool_port,
'ToPort': self.settings.mempool_port,
'IpRanges': [{
'CidrIp': '0.0.0.0/0',
'Description': 'Mempool port',
}],
'Ipv6Ranges': [{
'CidrIpv6': '::/0',
'Description': 'Mempool port',
}],
"IpProtocol": "tcp",
"FromPort": self.settings.mempool_port,
"ToPort": self.settings.mempool_port,
"IpRanges": [
{
"CidrIp": "0.0.0.0/0",
"Description": "Mempool port",
}
],
"Ipv6Ranges": [
{
"CidrIpv6": "::/0",
"Description": "Mempool port",
}
],
},
{
'IpProtocol': 'tcp',
'FromPort': self.settings.front_port,
'ToPort': self.settings.front_port,
'IpRanges': [{
'CidrIp': '0.0.0.0/0',
'Description': 'Front end to accept clients transactions',
}],
'Ipv6Ranges': [{
'CidrIpv6': '::/0',
'Description': 'Front end to accept clients transactions',
}],
"IpProtocol": "tcp",
"FromPort": self.settings.front_port,
"ToPort": self.settings.front_port,
"IpRanges": [
{
"CidrIp": "0.0.0.0/0",
"Description": "Front end to accept clients transactions",
}
],
"Ipv6Ranges": [
{
"CidrIpv6": "::/0",
"Description": "Front end to accept clients transactions",
}
],
},
]
],
)

def _get_ami(self, client):
# The AMI changes with regions.
response = client.describe_images(
Filters=[{
'Name': 'description',
'Values': ['Canonical, Ubuntu, 20.04 LTS, amd64 focal image build on 2020-10-26']
}]
Filters=[
{
"Name": "description",
"Values": [
"Canonical, Ubuntu, 22.04 LTS, amd64 jammy image build on 2023-09-19"
],
}
]
)
return response['Images'][0]['ImageId']
return response["Images"][0]["ImageId"]

def create_instances(self, instances):
assert isinstance(instances, int) and instances > 0
Expand All @@ -146,14 +160,14 @@ def create_instances(self, instances):
self._create_security_group(client)
except ClientError as e:
error = AWSError(e)
if error.code != 'InvalidGroup.Duplicate':
raise BenchError('Failed to create security group', error)
if error.code != "InvalidGroup.Duplicate":
raise BenchError("Failed to create security group", error)

try:
# Create all instances.
size = instances * len(self.clients)
progress = progress_bar(
self.clients.values(), prefix=f'Creating {size} instances'
self.clients.values(), prefix=f"Creating {size} instances"
)
for client in progress:
client.run_instances(
Expand All @@ -163,37 +177,38 @@ def create_instances(self, instances):
MaxCount=instances,
MinCount=instances,
SecurityGroups=[self.settings.testbed],
TagSpecifications=[{
'ResourceType': 'instance',
'Tags': [{
'Key': 'Name',
'Value': self.settings.testbed
}]
}],
TagSpecifications=[
{
"ResourceType": "instance",
"Tags": [{"Key": "Name", "Value": self.settings.testbed}],
}
],
EbsOptimized=True,
BlockDeviceMappings=[{
'DeviceName': '/dev/sda1',
'Ebs': {
'VolumeType': 'gp2',
'VolumeSize': 200,
'DeleteOnTermination': True
BlockDeviceMappings=[
{
"DeviceName": "/dev/sda1",
"Ebs": {
"VolumeType": "gp2",
"VolumeSize": 200,
"DeleteOnTermination": True,
},
}
}],
],
)

# Wait for the instances to boot.
Print.info('Waiting for all instances to boot...')
self._wait(['pending'])
Print.heading(f'Successfully created {size} new instances')
Print.info("Waiting for all instances to boot...")
self._wait(["pending"])
Print.heading(f"Successfully created {size} new instances")
except ClientError as e:
raise BenchError('Failed to create AWS instances', AWSError(e))
raise BenchError("Failed to create AWS instances", AWSError(e))

def terminate_instances(self):
try:
ids, _ = self._get(['pending', 'running', 'stopping', 'stopped'])
ids, _ = self._get(["pending", "running", "stopping", "stopped"])
size = sum(len(x) for x in ids.values())
if size == 0:
Print.heading(f'All instances are shut down')
Print.heading(f"All instances are shut down")
return

# Terminate instances.
Expand All @@ -202,64 +217,62 @@ def terminate_instances(self):
client.terminate_instances(InstanceIds=ids[region])

# Wait for all instances to properly shut down.
Print.info('Waiting for all instances to shut down...')
self._wait(['shutting-down'])
Print.info("Waiting for all instances to shut down...")
self._wait(["shutting-down"])
for client in self.clients.values():
client.delete_security_group(
GroupName=self.settings.testbed
)
client.delete_security_group(GroupName=self.settings.testbed)

Print.heading(f'Testbed of {size} instances destroyed')
Print.heading(f"Testbed of {size} instances destroyed")
except ClientError as e:
raise BenchError('Failed to terminate instances', AWSError(e))
raise BenchError("Failed to terminate instances", AWSError(e))

def start_instances(self, max):
size = 0
try:
ids, _ = self._get(['stopping', 'stopped'])
ids, _ = self._get(["stopping", "stopped"])
for region, client in self.clients.items():
if ids[region]:
target = ids[region]
target = target if len(target) < max else target[:max]
size += len(target)
client.start_instances(InstanceIds=target)
Print.heading(f'Starting {size} instances')
Print.heading(f"Starting {size} instances")
except ClientError as e:
raise BenchError('Failed to start instances', AWSError(e))
raise BenchError("Failed to start instances", AWSError(e))

def stop_instances(self):
try:
ids, _ = self._get(['pending', 'running'])
ids, _ = self._get(["pending", "running"])
for region, client in self.clients.items():
if ids[region]:
client.stop_instances(InstanceIds=ids[region])
size = sum(len(x) for x in ids.values())
Print.heading(f'Stopping {size} instances')
Print.heading(f"Stopping {size} instances")
except ClientError as e:
raise BenchError(AWSError(e))

def hosts(self, flat=False):
try:
_, ips = self._get(['pending', 'running'])
_, ips = self._get(["pending", "running"])
return [x for y in ips.values() for x in y] if flat else ips
except ClientError as e:
raise BenchError('Failed to gather instances IPs', AWSError(e))
raise BenchError("Failed to gather instances IPs", AWSError(e))

def print_info(self):
hosts = self.hosts()
key = self.settings.key_path
text = ''
text = ""
for region, ips in hosts.items():
text += f'\n Region: {region.upper()}\n'
text += f"\n Region: {region.upper()}\n"
for i, ip in enumerate(ips):
new_line = '\n' if (i+1) % 6 == 0 else ''
text += f'{new_line} {i}\tssh -i {key} ubuntu@{ip}\n'
new_line = "\n" if (i + 1) % 6 == 0 else ""
text += f"{new_line} {i}\tssh -i {key} ubuntu@{ip}\n"
print(
'\n'
'----------------------------------------------------------------\n'
' INFO:\n'
'----------------------------------------------------------------\n'
f' Available machines: {sum(len(x) for x in hosts.values())}\n'
f'{text}'
'----------------------------------------------------------------\n'
"\n"
"----------------------------------------------------------------\n"
" INFO:\n"
"----------------------------------------------------------------\n"
f" Available machines: {sum(len(x) for x in hosts.values())}\n"
f"{text}"
"----------------------------------------------------------------\n"
)

0 comments on commit 7e705f6

Please sign in to comment.