Respawn and patch for elastigroup (#530)

Adds patch and respawn functionality to elastigroup stacks

senza/cli.py

@@ -19,13 +19,14 @@
 import senza.stups.taupage as taupage
 import requests
 import yaml
+import senza.respawn as respawn
 
 from botocore.exceptions import ClientError
 from clickclick import (Action, FloatRange, OutputFormat, choice, error,
                         fatal_error, info, ok)
 from clickclick.console import print_table
 
-from spotinst.components import elastigroup_api
+from .spotinst.components import elastigroup_api
 from .arguments import (GLOBAL_OPTIONS, json_output_option, output_option,
                         parameter_file_option, region_option,
                         stacktrace_visible_option, watch_option,
@@ -42,8 +43,7 @@
 from .manaus.exceptions import VPCError
 from .manaus.route53 import Route53, Route53Record
 from .manaus.utils import extract_client_error_code
-from .patch import patch_auto_scaling_group
-from .respawn import get_auto_scaling_group, respawn_auto_scaling_group
+from .patch import patch_auto_scaling_group, patch_elastigroup
 from .stups.piu import Piu
 from .subcommands.config import cmd_config
 from .subcommands.root import cli
@@ -1476,7 +1476,7 @@ def get_auto_scaling_groups(stack_refs, region):
 def patch(stack_ref, region, image, instance_type, user_data):
     '''Patch specific properties of existing stack.
 
-    Currently only supports patching ASG launch configurations.'''
+    Currently supports patching ASG launch configurations and ElastiGroup groups.'''
 
     stack_refs = get_stack_refs(stack_ref)
     region = get_region(region)
@@ -1499,25 +1499,55 @@ def patch(stack_ref, region, image, instance_type, user_data):
 
     asg = BotoClientProxy('autoscaling', region)
 
-    for asg_name in get_auto_scaling_groups(stack_refs, region):
-        with Action('Patching Auto Scaling Group {}..'.format(asg_name)) as act:
-            result = asg.describe_auto_scaling_groups(AutoScalingGroupNames=[asg_name])
-            groups = result['AutoScalingGroups']
-            for group in groups:
-                if not patch_auto_scaling_group(group, region, properties):
-                    act.ok('NO CHANGES')
+    stacks = get_stacks(stack_refs, region)
+    for group in get_auto_scaling_groups_and_elasti_groups(stacks, region):
+        if group['type'] == ELASTIGROUP_TYPE:
+            patch_spotinst_elastigroup(properties, group['resource_id'], region, group['stack_name'])
+        elif group['type'] == AUTO_SCALING_GROUP_TYPE:
+            patch_aws_asg(properties, region, asg, group['resource_id'])
+
+
+def patch_aws_asg(properties, region, asg, asg_name):
+    '''
+    Patch an AWS Auto Scaling Group
+    '''
+    with Action('Patching Auto Scaling Group {}..'.format(asg_name)) as act:
+        result = asg.describe_auto_scaling_groups(AutoScalingGroupNames=[asg_name])
+        groups = result['AutoScalingGroups']
+        for group in groups:
+            if not patch_auto_scaling_group(group, region, properties):
+                act.ok('NO CHANGES')
+
+
+def patch_spotinst_elastigroup(properties, elastigroup_id, region, stack_name):
+    '''
+    Patch specific properties of an existing ElastiGroup
+    '''
+
+    spotinst_account_data = elastigroup_api.get_spotinst_account_data(region, stack_name)
+
+    with Action('Patching ElastiGroup {} (ID: {})..'.format(stack_name, elastigroup_id)) as act:
+        groups = elastigroup_api.get_elastigroup(elastigroup_id, spotinst_account_data)
+
+        for group in groups:
+            if not patch_elastigroup(group, properties, elastigroup_id, spotinst_account_data):
+                act.ok('NO CHANGES')
 
 
 @cli.command('respawn-instances')
 @click.argument('stack_ref', nargs=-1)
 @click.option('--inplace',
-              is_flag=True, help='Perform inplace update, do not scale out')
+              is_flag=True, help='Perform inplace update, do not scale out. Ignored for ElastiGroups.')
 @click.option('-f', '--force',
               is_flag=True,
-              help='Force respawn even if Launch Configuration is unchanged')
+              help='Force respawn even if Launch Configuration is unchanged. Ignored for ElastiGroups.')
+@click.option('--batch_size_percentage',
+              metavar='PERCENTAGE',
+              help='Percentage (int value) of the ElastiGroup cluster that is respawned in each step.'
+                   ' Valid only for ElastiGroups. The default value for this of 20.')
 @region_option
 @stacktrace_visible_option
-def respawn_instances(stack_ref, inplace, force, region):
+def respawn_instances(stack_ref, inplace, force, batch_size_percentage, region):
     '''Replace all EC2 instances in Auto Scaling Group(s)
 
     Performs a rolling update to prevent downtimes.'''
@@ -1526,8 +1556,12 @@ def respawn_instances(stack_ref, inplace, force, region):
     region = get_region(region)
     check_credentials(region)
 
-    for asg_name in get_auto_scaling_groups(stack_refs, region):
-        respawn_auto_scaling_group(asg_name, region, inplace=inplace, force=force)
+    stacks = get_stacks(stack_refs, region)
+    for group in get_auto_scaling_groups_and_elasti_groups(stacks, region):
+        if group['type'] == AUTO_SCALING_GROUP_TYPE:
+            respawn.respawn_auto_scaling_group(group['resource_id'], region, inplace=inplace, force=force)
+        elif group['type'] == ELASTIGROUP_TYPE:
+            respawn.respawn_elastigroup(group['resource_id'], group['stack_name'], region, batch_size_percentage)
 
 
 @cli.command()
@@ -1545,14 +1579,13 @@ def scale(stack_ref, region, desired_capacity, force):
     region = get_region(region)
     check_credentials(region)
 
-    asg = BotoClientProxy('autoscaling', region)
-
     stacks = get_stacks(stack_refs, region)
     stack_count = len(stacks)
     if not force and stack_count > 1:
         confirm_str = 'Number of stacks to be scaled - {}. Do you want to continue?'.format(stack_count)
         click.confirm(confirm_str, abort=True)
 
+    asg = BotoClientProxy('autoscaling', region)
     for group in get_auto_scaling_groups_and_elasti_groups(stacks, region):
         if group['type'] == AUTO_SCALING_GROUP_TYPE:
             scale_auto_scaling_group(asg, group['resource_id'], desired_capacity)
@@ -1564,36 +1597,30 @@ def scale_elastigroup(elastigroup_id, stack_name, desired_capacity, region):
     '''
     Commands to scale an ElastiGroup
     '''
-    cf = boto3.client('cloudformation', region)
-    template = cf.get_template(StackName=stack_name)['TemplateBody']
+    spotinst_account_data = elastigroup_api.get_spotinst_account_data(region, stack_name)
 
-    spotinst_token = template['Mappings']['Senza']['Info']['SpotinstAccessToken']
-    spotinst_account_id = template['Resources']['AppServerConfig']['Properties']['accountId']
+    groups = elastigroup_api.get_elastigroup(elastigroup_id, spotinst_account_data)
 
-    group = elastigroup_api.get_elastigroup(elastigroup_id, spotinst_account_id, spotinst_token)
-    capacity = group['capacity']
+    for group in groups:
+        capacity = group['capacity']
 
-    with Action('Scaling ElastiGroup {} (ID: {}) from {} to {} instances..'.format(
-            stack_name, elastigroup_id, capacity['target'], desired_capacity)) as act:
-        if capacity['target'] == desired_capacity:
-            act.ok('NO CHANGES')
-        else:
-            minimum = desired_capacity if capacity['minimum'] > desired_capacity else capacity['minimum']
-            maximum = desired_capacity if capacity['maximum'] < desired_capacity else capacity['maximum']
+        with Action('Scaling ElastiGroup {} (ID: {}) from {} to {} instances..'.format(
+                stack_name, elastigroup_id, capacity['target'], desired_capacity)) as act:
+            if capacity['target'] == desired_capacity:
+                act.ok('NO CHANGES')
+            else:
+                minimum = desired_capacity if capacity['minimum'] > desired_capacity else capacity['minimum']
+                maximum = desired_capacity if capacity['maximum'] < desired_capacity else capacity['maximum']
 
-            elastigroup_api.update_capacity(minimum,
-                                            maximum,
-                                            desired_capacity,
-                                            elastigroup_id,
-                                            spotinst_account_id,
-                                            spotinst_token)
+                elastigroup_api.update_capacity(minimum, maximum, desired_capacity, elastigroup_id,
+                                                spotinst_account_data)
 
 
 def scale_auto_scaling_group(asg, asg_name, desired_capacity):
     '''
     Commands to scale an AWS Auto Scaling Group
     '''
-    group = get_auto_scaling_group(asg, asg_name)
+    group = respawn.get_auto_scaling_group(asg, asg_name)
     current_capacity = group['DesiredCapacity']
     with Action('Scaling {} from {} to {} instances..'.format(
             asg_name, current_capacity, desired_capacity)) as act:

senza/patch.py

@@ -1,9 +1,11 @@
 
 import codecs
+import base64
 import datetime
 
 import yaml
 
+from .spotinst.components import elastigroup_api
 from .exceptions import InvalidUserDataType
 from .manaus.boto_proxy import BotoClientProxy
 
@@ -29,6 +31,19 @@
 ])
 
 
+def should_patch_user_data(new_val, old_val):
+    '''
+    Validate if User Data should be patched.
+    '''
+    current_user_data = yaml.safe_load(old_val)
+    if isinstance(new_val, dict):
+        return True
+    elif isinstance(current_user_data, dict):
+        raise InvalidUserDataType(type(current_user_data),
+                                  type(new_val))
+    return False
+
+
 def patch_user_data(old: str, new: dict):
     first_line, sep, data = old.partition('\n')
     data = yaml.safe_load(data)
@@ -60,16 +75,46 @@ def patch_auto_scaling_group(group: dict, region: str, properties: dict):
             for key, val in properties.items():
 
                 if key == 'UserData':
-                    current_user_data = yaml.safe_load(kwargs['UserData'])
-                    if isinstance(val, dict):
+                    if should_patch_user_data(val, kwargs['UserData']):
                         kwargs[key] = patch_user_data(kwargs[key], val)
-                    elif isinstance(current_user_data, dict):
-                        raise InvalidUserDataType(type(current_user_data),
-                                                  type(val))
                 else:
                     kwargs[key] = val
             asg.create_launch_configuration(**kwargs)
             asg.update_auto_scaling_group(AutoScalingGroupName=group['AutoScalingGroupName'],
                                           LaunchConfigurationName=kwargs['LaunchConfigurationName'])
             changed = True
     return changed
+
+
+def patch_elastigroup(group, properties, elastigroup_id, spotinst_account_data):
+    '''
+    Patch specific properties of an existing ElastiGroup
+    '''
+    changed = False
+    properties_to_patch = {}
+
+    group_user_data = group['compute']['launchSpecification']['userData']
+    current_user_data = codecs.decode(group_user_data.encode('utf-8'), 'base64').decode('utf-8')
+
+    current_properties = {
+        'ImageId': group['compute']['launchSpecification']['imageId'],
+        'InstanceType': group['compute']['instanceTypes']['ondemand'],
+        'UserData': current_user_data
+    }
+
+    for key, val in properties.items():
+        if key in current_properties:
+            if key == 'UserData':
+                if should_patch_user_data(val, current_properties[key]):
+                    patched_user_data = patch_user_data(current_properties[key], val)
+                    encoded_user_data = base64.urlsafe_b64encode(patched_user_data.encode('utf-8')).decode('utf-8')
+                    properties_to_patch[key] = encoded_user_data
+            else:
+                if current_properties[key] != val:
+                    properties_to_patch[key] = val
+
+    if len(properties_to_patch) > 0:
+        elastigroup_api.patch_elastigroup(properties_to_patch, elastigroup_id, spotinst_account_data)
+        changed = True
+
+    return changed

senza/respawn.py

@@ -6,10 +6,15 @@
 from clickclick import Action, info
 
 from .manaus.boto_proxy import BotoClientProxy
+from .spotinst.components import elastigroup_api
 
 SCALING_PROCESSES_TO_SUSPEND = ['AZRebalance', 'AlarmNotification', 'ScheduledActions']
 RUNNING_LIFECYCLE_STATES = set(['Pending', 'InService', 'Rebooting'])
 
+ELASTIGROUP_TERMINATED_DEPLOY_STATUS = ['stopped', 'failed']
+
+DEFAULT_BATCH_SIZE = 20
+
 
 def get_auto_scaling_group(asg, asg_name: str):
     '''Get boto3 Auto Scaling Group by name or raise exception'''
@@ -151,3 +156,37 @@ def respawn_auto_scaling_group(asg_name: str, region: str, inplace: bool=False,
                                       inplace)
     else:
         info('Nothing to do')
+
+
+def respawn_elastigroup(elastigroup_id: str, stack_name: str, region: str, batch_size: int):
+    '''
+    Respawn all instances in the ElastiGroup.
+    '''
+
+    if batch_size is None or batch_size < 1:
+        batch_size = DEFAULT_BATCH_SIZE
+
+    spotinst_account = elastigroup_api.get_spotinst_account_data(region, stack_name)
+
+    info('Redeploying the cluster for ElastiGroup {} (ID {})'.format(stack_name, elastigroup_id))
+
+    deploy_output = elastigroup_api.deploy(batch_size=batch_size, grace_period=600, elastigroup_id=elastigroup_id,
+                                           spotinst_account_data=spotinst_account)
+
+    deploy_count = len(deploy_output)
+    deploys_finished = 0
+    with Action('Waiting for deploy to complete. Total of {} deploys'.format(deploy_count)) as act:
+        while True:
+            for deploy in deploy_output:
+                deploy_status = elastigroup_api.deploy_status(deploy['id'], elastigroup_id, spotinst_account)
+                for ds in deploy_status:
+                    if ds['id'] == deploy['id']:
+                        if ds['progress']['value'] >= 100\
+                                or ds['status'].lower() in ELASTIGROUP_TERMINATED_DEPLOY_STATUS:
+                            deploys_finished += 1
+                            info('Deploy {} finished with status {}'.format(ds['id'], ds['status']))
+
+            if deploys_finished == deploy_count:
+                break
+            time.sleep(2)
+            act.progress()

spotinst/components/elastigroup.py → senza/spotinst/components/elastigroup.py

@@ -15,7 +15,7 @@
 from senza.components.taupage_auto_scaling_group import check_application_id, check_application_version, \
     check_docker_image_exists, generate_user_data
 from senza.utils import ensure_keys
-from spotinst import MissingSpotinstAccount
+from senza.spotinst import MissingSpotinstAccount
 
 SPOTINST_LAMBDA_FORMATION_ARN = 'arn:aws:lambda:{}:178579023202:function:spotinst-cloudformation'
 SPOTINST_API_URL = 'https://api.spotinst.io'