node.py

import time
from boto.ec2.connection import EC2Connection
from boto.ec2.blockdevicemapping import BlockDeviceType
from boto.ec2.blockdevicemapping import BlockDeviceMapping
from fabric.api import settings, sudo
from parse_client import ParseClient
from job import Job


class Node(object):
 """ Object representing an EC2 node """

 # list of job names
 jobs = []

 def __init__(self, **kwargs):
 """ new Node """
 self.name = kwargs['name']
 self.key_name = kwargs.get('key_name')
 self.zone = kwargs.get('zone')
 self.instance_type = kwargs.get('instance_type')
 self.user = kwargs.get('user')
 self.ami = kwargs.get('ami')
 self.security_group = kwargs.get('security_group')

 # optional
 self.instance_id = kwargs.get('instance_id')
 self.public_dns_name = kwargs.get('public_dns_name')
 self.private_dns_name = kwargs.get('private_dns_name')
 self.ip_address = kwargs.get('ip_address')
 self.private_ip_address = kwargs.get('private_ip_address')

 if kwargs.get('job'):
 self.jobs.append(kwargs['job'])
 elif kwargs.get('jobs'):
 self.jobs = kwargs['jobs']

 # Parse reference
 self.object_id = kwargs.get('objectId')

 # Experimenting
 self.ebs = False


 def to_dict(self):
 """ Return dictoionary reprenstation of object """
 return {
 'name': self.name,
 'instance_id': self.instance_id,
 'ami': self.ami,
 'key_name': self.key_name,
 'zone': self.zone,
 'instance_type': self.instance_type,
 'public_dns_name': self.public_dns_name,
 'private_dns_name': self.private_dns_name,
 'ip_address': self.ip_address,
 'private_ip_address': self.private_ip_address,
 'user': self.user,
 'jobs': self.jobs
 } 


 def launch(self):
 '''Launch a single instance of the provided ami '''
 conn = EC2Connection()
 # Declare the block device mapping for ephemeral disks

 mapping = None
 if self.ebs:
 # TODO - toggle instance storage
 mapping = BlockDeviceMapping()
 eph0 = BlockDeviceType()
 eph1 = BlockDeviceType()
 eph0.ephemeral_name = 'ephemeral0'
 eph1.ephemeral_name = 'ephemeral1'
 mapping['/dev/sdb'] = eph0
 mapping['/dev/sdc'] = eph1
 
 # Now, ask for a reservation
 reservation = conn.run_instances(self.ami, instance_type=self.instance_type, 
 key_name=self.key_name, placement=self.zone, 
 block_device_map=mapping, security_groups=[self.security_group])
 # And assume that the instance we're talking about is the first in the list
 # This is not always a good assumption, and will likely depend on the specifics
 # of your launching situation. For launching an isolated instance while no
 # other actions are taking place, this is sufficient.
 instance = reservation.instances[0]
 print('Waiting for instance to start...')
 # Check up on its status every so often
 try:
 status = instance.update()
 except:
 pass
 while status == 'pending':
 time.sleep(5)
 try:
 status = instance.update()
 except:
 pass
 if status == 'running':
 print('New instance "' + instance.id + '" accessible at ' + instance.public_dns_name)
 # Name the instance
 conn.create_tags([instance.id], {'Name': self.name})

 # update properties
 self.instance_id = instance.id
 self.private_dns_name = instance.private_dns_name
 self.public_dns_name = instance.public_dns_name
 self.ip_address = instance.ip_address
 self.private_ip_address = instance.private_ip_address

 ParseClient.add_node(self.to_dict())

 if self.try_connect():
 self.refresh_jobs()

 return True
 else:
 print('Instance status: ' + status)
 return False


 def try_connect(self):
 """ Test SSH connection in a rety loop """
 tries = 0
 while True:
 try:
 tries += 1
 with settings(host_string='%s@%s' % (self.user, self.ip_address)):
 sudo("hostname")
 return True
 except:
 time.sleep(5)
 if tries >= 10:
 print "Unable to connect after %d tries" % tries
 return False


 def refresh_jobs(self):
 """ Run all jobs """
 with settings(host_string='%s@%s' % (self.user, self.ip_address)):
 # Run init from Base job class, once per run
 job_obj = Job()
 job_obj.update_packages() 
 # Run all nodes jobs
 for job in self.jobs:
 self.run_single_job(job)


 def run_single_job(self, job):
 """ Connect to box and run job """
 template_vars = {
 'nodes': ParseClient.get_all_nodes()
 }
 with settings(host_string='%s@%s' % (self.user, self.ip_address)):
 job_obj = self.get_job_module(job)
 job_obj.run(template_vars) 


 def terminate(self):
 ''' Terminate this instance '''
 conn = EC2Connection()
 try:
 conn.terminate_instances([self.instance_id])
 except:
 pass
 ParseClient.delete_node(self.object_id)


 def mock_launch(self):
 """ For testing """
 mock = {
 'name': self.name,
 'key_name': self.key_name,
 'ami': self.ami,
 'zone': self.zone,
 'instance_type': self.instance_type,
 'user':self.user,
 'jobs':self.jobs,
 'public_dns_name': u'ec2-107-21-159-143.compute-1.amazonaws.com',
 'ip_address': u'107.21.159.143',
 'private_dns_name': u'ip-10-29-6-45.ec2.internal',
 'id': u'i-e2a5559d',
 'dns_name': u'ec2-107-21-159-143.compute-1.amazonaws.com',
 'private_ip_address': u'10.29.6.45'
 }

 if ParseClient.add_node(mock):
 print "Node stored on remote"
 else:
 print "Node storage failed on remote"


 def get_job_module(self, job):
 """ Helper to get hob object dynamically """
 obj = self.get_class("jobs.%s" % job)
 return obj()


 @staticmethod
 def get_class(kls):
 """ Source: http://stackoverflow.com/questions/452969/does-python-have-an-equivalent-to-java-class-forname """
 parts = kls.split('.')
 module = ".".join(parts[:-1])
 mod = __import__( module )
 for comp in parts[1:]:
 mod = getattr(mod, comp)
 return mod


 def add_job(self, job):
 """ Add a job to node, update node on remote """
 if job not in self.jobs:
 self.jobs.append(job)
 ParseClient.update_node(self.object_id, self.to_dict())


 @staticmethod 
 def get_node(name):
 """ Retrive node from remote, return Node object """
 result = ParseClient.get_node(name)
 return result if result is None else Node(**result)


 @staticmethod 
 def get_all_nodes():
 """ Retrieve an array of Node objects """
 result = ParseClient.get_all_nodes()
 return [Node(**n) for n in result]