Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- #!/usr/bin/python
- """Get the set of instance id's (exclude spot instances)!"""
- import boto3
- import datetime
- import json
- import redis
- import sys
- import base64
- VERSION = 0.2
- import base64
- import struct
- def encode(n):
- n = int(n[2:], 16)
- data = struct.pack('<Q', n).rstrip('\x00')
- if len(data) == 0:
- data = '\x00'
- s = base64.urlsafe_b64encode(data).rstrip('=')
- return "{:>6}".format(s) # right just with spaces?
- def decode(s):
- data = base64.urlsafe_b64decode(s + '==')
- n = struct.unpack('<Q', data + '\x00'* (8-len(data)) )
- return "i-{:>08}".format(hex(n[0])[2:])
- def cache(innerfunc):
- redis_conn = redis.StrictRedis()
- key = (innerfunc.__name__, VERSION)
- def wrapped():
- cached_val = redis_conn.get(key)
- if cached_val:
- val = json.loads(cached_val)
- else:
- val = innerfunc()
- redis_conn.setex(key, 60*30, json.dumps(val))
- return val
- return wrapped
- def datetime_to_timestamp(idatetime):
- if isinstance(idatetime, datetime.datetime):
- epoch = datetime.datetime.utcfromtimestamp(0)
- return (idatetime.replace(tzinfo=None) - epoch).total_seconds()
- else:
- return idatetime
- def recursive_date_to_epoch(obj):
- """find any date time objects and transform them into time since epoch"""
- if isinstance(obj, dict):
- for key, val in obj.iteritems():
- obj[key] = datetime_to_timestamp(val)
- recursive_date_to_epoch(val)
- elif isinstance(obj, list):
- for i, j in enumerate(obj):
- obj[i] = datetime_to_timestamp(j)
- recursive_date_to_epoch(j)
- @cache
- def get_running_instances():
- conn = boto3.client('ec2', 'us-east-1')
- response = conn.describe_instances(
- Filters=[
- {
- 'Name': 'instance-state-name',
- 'Values': [
- 'running',
- ]
- },
- ],
- )['Reservations']
- recursive_date_to_epoch(response)
- all_instances = list()
- for i in response:
- all_instances.extend(i['Instances'])
- non_spot_instances = [x for x in all_instances if x.get("InstanceLifecycle") != "spot"]
- return non_spot_instances
- def tag_reservation(obj_id, tag_key, tag_val):
- conn = boto3.client('ec2', 'us-east-1')
- response = conn.create_tags(
- Resources=[
- obj_id,
- ],
- Tags=[
- {
- 'Key': tag_key,
- 'Value': tag_val
- },
- ]
- )
- #@cache
- def get_reservations():
- conn = boto3.client('ec2', 'us-east-1')
- response = conn.describe_reserved_instances(
- Filters=[
- {
- 'Name': 'state',
- 'Values': [
- 'active'
- ]
- }
- ]
- )['ReservedInstances']
- recursive_date_to_epoch(response)
- return response
- def partition(iterable, func):
- ans = dict()
- for i in iterable:
- key = func(i)
- ans.setdefault(key, list())
- ans[key].append(i)
- return ans
- def unpack(packed_string):
- ans = list()
- try:
- for startpos in xrange(0, len(packed_string), 6):
- packed_id = packed_string[startpos: startpos+6]
- ans.append(decode(packed_id))
- return sorted(ans)
- except Exception as oops:
- return list()
- class ReservationManager(object):
- instance_group_function = lambda x: (x['InstanceType'], x['Placement']['AvailabilityZone'])
- res_group_function = lambda x: (x['InstanceType'], x['AvailabilityZone'])
- def __init__(self):
- self.running_instances = get_running_instances()
- self.partitioned_instances = partition(self.running_instances, self.instance_group_function)
- self.clean_instances()
- self.reservations = get_reservations()
- self.partitioned_reservations = partition(self.reservations, self.res_group_function)
- self.clean_reservations()
- self.instance_to_reservation = dict()
- # filter out reserved instances? and remove used reservations
- self.reservation_to_instances = dict()
- for res in self.reservations:
- res_id = res['ReservedInstancesId']
- self.reservation_to_instances[res_id] = set()
- self.iid_to_instance = {instance['InstanceId']: instance for instance in self.running_instances}
- self.sync_with_saved()
- def sync_with_saved(self):
- self.clean_instances()
- for instance in self.running_instances:
- res_id = instance.get('Tags', {}).get('reservation_id')
- instance_id = instance['InstanceId']
- try:
- self.reservation_to_instances[res_id].add(instance_id)
- except:
- pass # you land here when the reservation has expired
- for res in self.reservations:
- res_id = res['ReservedInstancesId']
- res.setdefault('Tags', dict())
- tags = res.get('Tags', {})
- current_instances = tags.get('instances', '')
- unpacked = unpack(current_instances)
- res['Tags']['instances'] = set(unpacked)
- for iid in unpacked:
- if iid in self.iid_to_instance: # i.e. it's a running instance
- self.instance_to_reservation[iid] = res_id
- def clean_instances(self):
- for instance in self.running_instances:
- for key in ["NetworkInterfaces", "BlockDeviceMappings", "SecurityGroups", "IamInstanceProfile", "Monitoring", "State"]:
- instance.pop(key, None)
- tag_list = instance.get('Tags')
- if type(tag_list) == type(list()):
- tag_dict = dict([(x['Key'], x['Value']) for x in tag_list])
- instance['Tags'] = tag_dict
- def clean_reservations(self):
- for res in self.reservations:
- tag_list = res.get('Tags')
- if type(tag_list) == type(list()):
- tag_dict = dict([(x['Key'], x['Value']) for x in tag_list])
- res['Tags'] = tag_dict
- def update_for_partition(self, partition):
- (instance_type, az) = partition
- instances = self.partitioned_instances[partition]
- instances.sort(key=lambda x: x['LaunchTime']) # reserve oldest instance first
- reservations = self.partitioned_reservations[partition]
- reservations.sort(key=lambda x: x['End'], reverse=True) # sort from ending furthest in the future to soonest
- first_unreserved_index = 0
- total_slots = sum([x['InstanceCount'] for x in reservations])
- for reservation in reservations:
- res_id = reservation['ReservedInstancesId']
- for slot in xrange(reservation['InstanceCount']):
- try:
- instance = instances[first_unreserved_index]
- except IndexError:
- return
- instance_id = instance['InstanceId']
- first_unreserved_index += 1
- self.reservation_to_instances[res_id].add(instance_id)
- def tag_reservations(self):
- for res in self.reservations:
- res_id = res['ReservedInstancesId']
- old_instances = sorted(res['Tags']['instances'])
- new_instances = sorted(self.reservation_to_instances[res_id])
- if old_instances != new_instances:
- new_tag = "".join([encode(instance) for instance in new_instances])
- tag_reservation(res_id, 'instances', new_tag)
- else:
- pass
- def adjust_single_reservation(self, res_id):
- print >> sys.stderr, "Would update {}".format(res_id)
- pass
- def adjust_reservations(self):
- for res in self.reservations:
- res_id = res['ReservedInstancesId']
- slots = res['InstanceCount']
- instance_type = res['InstanceType']
- used = len(self.reservation_to_instances[res_id])
- if slots != used:
- print "{} of {} {} reserved for {}".format(used, instance_type, slots, res_id)
- self.adjust_single_reservation(res_id)
- # figure out how to use the excess here...
- def update_reserved_instances(self):
- for partition in sorted(set(self.partitioned_instances) & set(self.partitioned_reservations)):
- self.update_for_partition(partition)
- for key, val in self.reservation_to_instances.iteritems():
- self.reservation_to_instances[key] = sorted(val)
- # see if there are reservations which are not used up...
- self.adjust_reservations()
- # update reservation configurations... (ONLY FOR THOSE WHICH CHANGED)
- # tag both ways...
- self.tag_reservations()
- def main():
- reservation_manager = ReservationManager()
- reservation_manager.update_reserved_instances()
- if "__main__" == __name__:
- main()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement