def testGetInstanceTypeByCPU(self):
"""Test that the instance type matches the requested amount of CPU cores."""
# pylint: disable=protected-access
self.assertEqual('m4.large', common.GetInstanceTypeByCPU(2))
self.assertEqual('m4.16xlarge', common.GetInstanceTypeByCPU(64))
self.assertRaises(ValueError, common.GetInstanceTypeByCPU, 0)
self.assertRaises(ValueError, common.GetInstanceTypeByCPU, 256)
def testGetInstanceTypeByCPU(self):
"""Test that the instance type matches the requested amount of CPU cores."""
self.assertEqual('m4.large', common.GetInstanceTypeByCPU(2))
self.assertEqual('m4.16xlarge', common.GetInstanceTypeByCPU(64))
with self.assertRaises(ValueError):
common.GetInstanceTypeByCPU(0)
with self.assertRaises(ValueError):
common.GetInstanceTypeByCPU(256)
def GetOrCreateAnalysisVm(
self,
vm_name: str,
boot_volume_size: int,
ami: str,
cpu_cores: int,
packages: Optional[List[str]] = None,
ssh_key_name: Optional[str] = None
) -> Tuple[ec2.AWSInstance, bool]:
"""Get or create a new virtual machine for analysis purposes.
Args:
vm_name (str): The instance name tag of the virtual machine.
boot_volume_size (int): The size of the analysis VM boot volume (in GB).
ami (str): The Amazon Machine Image ID to use to create the VM.
cpu_cores (int): Number of CPU cores for the analysis VM.
packages (List[str]): Optional. List of packages to install in the VM.
ssh_key_name (str): Optional. A SSH key pair name linked to the AWS
account to associate with the VM. If none provided, the VM can only
be accessed through in-browser SSH from the AWS management console
with the EC2 client connection package (ec2-instance-connect). Note
that if this package fails to install on the target VM, then the VM
will not be accessible. It is therefore recommended to fill in this
parameter.
Returns:
Tuple[AWSInstance, bool]: A tuple with an AWSInstance object and a
boolean indicating if the virtual machine was created (True) or
reused (False).
Raises:
RuntimeError: If the virtual machine cannot be found or created.
"""
# Re-use instance if it already exists, or create a new one.
try:
instances = self.GetInstancesByName(vm_name)
if instances:
created = False
return instances[0], created
except RuntimeError:
pass
instance_type = common.GetInstanceTypeByCPU(cpu_cores)
startup_script = utils.ReadStartupScript()
if packages:
startup_script = startup_script.replace('${packages[@]}',
' '.join(packages))
# Install ec2-instance-connect to allow SSH connections from the browser.
startup_script = startup_script.replace(
'(exit ${exit_code})',
'apt -y install ec2-instance-connect && (exit ${exit_code})')
client = self.ClientApi(common.EC2_SERVICE)
vm_args = {
'BlockDeviceMappings':
[self._GetBootVolumeConfigByAmi(ami, boot_volume_size)],
'ImageId':
ami,
'MinCount':
1,
'MaxCount':
1,
'InstanceType':
instance_type,
'TagSpecifications':
[common.GetTagForResourceType('instance', vm_name)],
'UserData':
startup_script,
'Placement': {
'AvailabilityZone': self.default_availability_zone
}
}
if ssh_key_name:
vm_args['KeyName'] = ssh_key_name
# Create the instance in AWS
try:
instance = client.run_instances(**vm_args)
# If the call to run_instances was successful, then the API response
# contains the instance ID for the new instance.
instance_id = instance['Instances'][0]['InstanceId']
# Wait for the instance to be running
client.get_waiter('instance_running').wait(
InstanceIds=[instance_id])
# Wait for the status checks to pass
client.get_waiter('instance_status_ok').wait(
InstanceIds=[instance_id])
except (client.exceptions.ClientError,
botocore.exceptions.WaiterError) as exception:
raise RuntimeError('Could not create instance {0:s}: {1:s}'.format(
vm_name, str(exception)))
instance = ec2.AWSInstance(self,
instance_id,
self.default_region,
self.default_availability_zone,
name=vm_name)
created = True
return instance, created
