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)
Exemplo n.º 2
0
 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)
Exemplo n.º 3
0
    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