~kirkland/uec-testing-scripts/euca2ools

#!/usr/bin/python

import datetime
import logging, logging.handlers
import os, os.path
import random
import re
import string
import sys
import yaml
import time

from twisted.internet import reactor, task
from twisted.python import usage
from zipfile import ZipFile, is_zipfile
from shutil import rmtree
from collections import defaultdict, namedtuple

import uec_utils

class Options(usage.Options):
    optFlags = [
        ["cleanup", None, "Clean up" ],
        ["write_log", 'w', "Write log to file" ],
        ["max_instances", 'm', "Run the maximum number of instances" ],
    ]
    optParameters = [
        ["log", "l", "info", 
         "log level wanted (debug, info, warning, error, critical"],
        ["start_stop", "s", None, 
         "start and stop instances over a certain period of time (x minutes)"],
    ]

euca_describe_availability_zones = 'euca-describe-availability-zones'
euca_run_instances               = 'euca-run-instances'
euca_describe_instances          = 'euca-describe-instances'
euca_terminate_instances         = 'euca-terminate-instances'
euca_add_keypair                 = 'euca-add-keypair'
euca_delete_keypair              = 'euca-delete-keypair'
euca_add_group                   = 'euca-add-group'
euca_delete_group                = 'euca-delete-group'
euca_authorize                   = 'euca-authorize'
euca_get_console_output          = 'euca-get-console-output'
euca_conf                        = 'euca_conf'
euca_describe_images             = 'euca-describe-images'

zone       = namedtuple('zone', 'max cpu ram disk')
admin_user = None
opts       = None
overtime   = False

class Instance(object):
    TEST_STATES      = [ 'not-tested', 'being-tested', 'success', 'failed',
                         'rescheduled', 'lost' ]
    TEST_MAX_RETRIES = 12   
    EMI_USED         = defaultdict(int)
    VMTYPE_USED      = defaultdict(int)
    GROUP_USED       = defaultdict(int)
    KEYPAIR_USED     = defaultdict(int)
    MAX_INSTANCES    = 0
    ZONE_CPU         = []
    ZONE_CPU_DICT    = dict()
    CLUSTERS         = defaultdict(int)

    def __init__(self, emi, vmtype, user, ssh_key, group):
        self.emi           = emi
        self.vmtype        = vmtype
        self.user          = user
        self.ssh_key       = ssh_key
        self.group         = group
        self.id            = None
        self.pub_ip        = None
        self.test_state    = "not-tested"
        self._test_retries = 0
        self._state        = None
        self.req_shutdown  = False
        self.zone          = None
        self.cluster       = None

    def __repr__(self):
        return "<Instance: %s - %s - %s - %s - %s - %s - %s>" % (self.id, \
        self.emi, self.vmtype, self.user, self.ssh_key, self.group, self.state)

    @property
    def logger(self):
        if self.id:
            return logging.getLogger("INSTANCE %s" % (self.id))
        else:
            return logging.getLogger()

    @property
    def private_key(self):
        return self.user.private_keys[self.ssh_key]

    @property
    def state(self):
        return self._state
    @state.setter
    def state(self, value):
        self.logger.debug("State: %s" % (value))
        # The state has changed
        if self._state != value:
            self.logger.debug("New state: %s => %s" % (self._state, value))
            self._state = value
            # Only trigger a test when the state has been changed to running
            if self._state == "running":
                self.logger.debug("Scheduling test")
                reactor.callLater(random.randint(5, 10) + random.random(),
                                 self.test)
            if self._state in ["terminated", "shutting-down"] \
               and not self.req_shutdown:
                self.test_state = 'lost'

    @property
    def test_state(self):
        return self._test_state
    @test_state.setter
    def test_state(self, value):
        self.logger.debug("Test state: %s" % (value))
        if value not in self.TEST_STATES:
            raise ValueError, "Unknow test state: %s" % (value)
        else:
            self._test_state = value

    def getProcessOutput(self, executable, args=(), timeout=60):
        self.logger.debug("Executing: %s - %s" % (executable, args))
        return self.user.getProcessOutput(executable, args, timeout)

    def start(self):
        self.logger.info("Starting instance")
        output = self.getProcessOutput(euca_run_instances,
                                        args= [self.emi,
                                               '-k', self.ssh_key,
                                               '-g', self.group,
                                               '-t', self.vmtype])
        output.addCallbacks(self.started, self.errStarted)
        # Terminate the instance no matter what
        reactor.callLater(45*60, self.terminate)

    def started(self, output):
        self.logger.debug("Instance start output: %s" % (output))
        started = False
        for l in output.split("\n"):
            if l.startswith('INSTANCE'):
                self.id = l.split()[1]
                self.cluster = l.split()[10]
                self.logger.info("Started (%s/%s)" %
                                 (self.emi, self.vmtype))
                started = True
        if not started:
            self.errStarted(output)

    def errStarted(self, output):
        self.logger.warning("Instance failed to start: %s" % (output))

    def test(self):
        if self.state != "running":
            self.logger.debug("Not running - aborting scheduled test.")
            return
        self.logger.info("Testing instance")
        self.test_state = "being-tested"
        self._test_retries += 1
        args = ['-o', 'UserKnownHostsFile=/dev/null',
                '-o', 'StrictHostKeyChecking=no',
                '-o', 'ConnectTimeout=5',
                '-o', 'Batchmode=yes',
                '-i', self.private_key,
                "ubuntu@%s" % (self.pub_ip),
                'echo TEST_SUCCESS']
        self.logger.debug("Ssh args: %s" % (args))
        output = uec_utils.getProcessOutput('ssh', args, errortoo=True)
        output.addCallback(self.tested)

    def tested(self, output):
        if 'TEST_SUCCESS' in output:
            self.logger.info("Test successful [%s]" % self._test_retries)
            self.test_state = "success"
            self.logger.debug("Test output: %s" % (output))
            if overtime:            
                self.terminate()
        elif self._test_retries <= self.TEST_MAX_RETRIES:
            self.logger.info("Rescheduling test (%s/%s)" %
                             (self._test_retries, self.TEST_MAX_RETRIES))
            self.test_state = "rescheduled"
            self.logger.debug("Test output: %s" % (output))
            reactor.callLater(30, self.test)
        else:
            self.logger.warning("Test output: %s" % (output))
            self.logger.info("Test failed")
            self.test_state = "failed"
            output = self.getProcessOutput(euca_get_console_output,
                                           args = [self.id], timeout=300)
            self.req_shutdown  = True
            output.addCallback(self.consoleOutput)


    def consoleOutput(self, output):
        # caution -- only used when a test failed
       self.logger.warning("Console output: START")
       for aLine in re.splitlines(output):
           self.logger.warning(aLine)
       self.logger.warning("Console output: END")
       output = self.getProcessOutput(euca_terminate_instances,
                                      args= [self.id])
       output.addCallback(self.terminated)

    def terminate(self):
        if self.state != "terminated":
            self.logger.info("Terminating")
            if self.id is None:
                self.state = "terminated"
                self.logger.info(
                    "instance has already terminated, last test state=%s" \
                    % self.test_state)
            else:
                if self.test_state not in ( "success", "failed" ):
                    output = self.getProcessOutput(euca_get_console_output,
                                           args = [self.id], timeout=300)
                    output.addCallback(self.consoleOutput)
                else:
                    self.req_shutdown  = True                    
                    output = self.getProcessOutput(euca_terminate_instances,
                                                   args= [self.id])
                    output.addCallback(self.terminated)

    def terminated(self, output):
        self.logger.info("Terminated")
        if self.test_state not in ( "success", "failed" ):
            self.logger.info("terminated state found while test_state = %s" %
                             (self.test_state))
            self.test_state = "failed"
        self.state = "terminated"
        self.logger.debug("terminate output: %s" % (output))

class User(object):
    def __init__(self):
        self.user_id     = 'admin'
        self.logger      = logging.getLogger("USER %s" % (self.user_id))
        self.keypairs    = []
        self.groups      = []
        self.cred_dir    = '%s/cred/%s/' %(os.getcwd(), self.user_id)
        self.cred_zip    = 'cred.zip'
        self.nb_groups   = 2
        self.nb_keypairs = 2     

    def __repr__(self):
        return "<User: %s>" % (self.user_id)

    def setup(self):
        rmtree(self.cred_dir, ignore_errors=1)
        os.makedirs(self.cred_dir)    
        output = self.getProcessOutput(euca_conf, 
                                       args = ['--get-credentials', 
                                               self.cred_dir+self.cred_zip,
                                               r'2>&1'],
                                       eucarc = False)
        output.addCallback(self.credentialsAdded)        

    def credentialsAdded(self, output):
        self.logger.debug("Credentials added output: %s" % (output)) 
        if not is_zipfile(self.cred_dir+self.cred_zip):
            self.logger.error("Credential file is not a valid zip")
            reactor.stop()
        ZipFile(self.cred_dir+self.cred_zip).extractall(self.cred_dir)
        self.logger.info("Credentials downloaded")
        for k in self.keypair_names:
            self.logger.debug("Adding keypair %s" % (k))
            output = self.getProcessOutput(euca_add_keypair, args = [k])
            output.addCallback(self.keypairAdded)
        for g in self.group_names:
            self.logger.debug("Adding group %s" % (g))
            output = self.getProcessOutput(euca_add_group,
                                           args = ['-d', 'UEC-test', g])
            output.addCallback(self.groupAdded)
      
    def keypairAdded(self, output):
        self.logger.debug("Keypair added output: %s" % (output))
        key = None
        privkey_fh = None
        for l in output.split("\n"):
            if not key and l.startswith("KEYPAIR"):
                key = l.split()[1]
                self.logger.info("Keypair (%s) added" % (key))
                Instance.KEYPAIR_USED[key]
                self.keypairs.append(key)
                privkey_fh = open(self.private_keys[key], 'w')
                os.chmod(self.private_keys[key], 0600)
                continue
            if privkey_fh:
                privkey_fh.write("%s\n" % (l))
        if privkey_fh:
            privkey_fh.close()

    def groupAdded(self, output):
        self.logger.debug("Group added output: %s" % (output))
        for l in output.split("\n"):
            if l.startswith("GROUP"):
                group = l.split()[1]
                self.logger.info("Group (%s) added" % (group))
                Instance.GROUP_USED[group]
                self.logger.debug("Authorizing group %s" % (group))
                for port in ['22','80']:
                    # FIXME
                    time.sleep(3) # Dummy delay see LP #592641
                    output = self.getProcessOutput(euca_authorize,
                                                   args = [group,
                                                           '-P', 'tcp', 
                                                           '-p', port,
                                                           '-s',  '0.0.0.0/0'])
                    output.addCallback(self.groupAuthorized)

    def groupAuthorized(self, output):
        self.logger.debug("Group authorize output: %s" % (output))
        for l in output.split("\n"):
            if l.startswith("GROUP"):
                group = l.split()[1]
                self.logger.info("Group %s authorized" % (group))
                if not self.groups.count(group): self.groups.append(group)

    @property
    def ready(self):
        return len(self.keypairs) == self.nb_keypairs and \
            len(self.groups) == self.nb_groups

    @property
    def private_keys(self):
        ret = {}
        for k in self.keypairs:
            ret[k] = os.path.join(self.cred_dir, "%s.priv" % (k))
        return ret

    @property
    def keypair_names(self):
        return [ "uectest-k%s" % (k) for k in range(self.nb_keypairs)]

    @property
    def group_names(self):
        return [ "uectest-g%s" % (g) for g in range(self.nb_groups)]

    def getProcessOutput(self, executable, args=(), timeout=60, eucarc = True):
        eucarc = "source %s/eucarc ; " %self.cred_dir if eucarc else ""
        myargs = [ '-c', "%s%s %s" % (eucarc,
                                                       executable,
                                                       string.join(args))]
        self.logger.debug("Running bash command: %s " % (myargs))
        return uec_utils.getProcessOutput('bash', myargs, timeout=timeout,
                                        env=os.environ)

    def cleanup(self):
        for k in self.keypair_names:
            self.logger.debug("Deleting keypair %s" % (k))
            output = self.getProcessOutput(euca_delete_keypair, args = [k])
            output.addCallback(self.keypairDeleted, k)
        for g in self.group_names:
            self.logger.debug("Deleting group %s" % (g))
            output = self.getProcessOutput(euca_delete_group, args = [g])
            output.addCallback(self.groupDeleted, g)

    def keypairDeleted(self, output, key):
        self.logger.info("Keypair (%s) deleted" % (key))
        self.logger.debug("Keypair deleted output: %s" % (output))

    def groupDeleted(self, output, group):
        self.logger.info("Group (%s) deleted" % (group))
        self.logger.debug("Group deleted output: %s" % (output))

def checkRessources(instances, admin_user, first_call = False):
    if not admin_user.ready:
        reactor.callLater(10, checkRessources, instances, admin_user, 
                          first_call = True)  
        return
    logging.debug("Checking ressource availability with %s" % (admin_user))
    logging.debug("Known instances: %s" % (instances))
    output = admin_user.getProcessOutput(euca_describe_availability_zones,
                                         args= ['verbose'])
    output.addCallback(availableEMI, instances, admin_user, first_call)

def availableEMI(output, instances, admin_user, first_call):   
    logging.debug("describe-availability-zones: %s" % (output))
    # Get a list of types which have available ressource
    available_types = []
    zone_cpu        = []
    global zone
    global opts
    for l in output.split("\n"):
        m = re.search(r"(\S+)\s+(\d+)\s/\s(\d+)\s+(\d+)\s+(\d+)\s+(\d+)", l)
        if m and int(m.group(2)) != 0:
            available_types.append(m.group(1))
        if m and first_call:
            zone_cpu.append(m.group(1))            
            Instance.ZONE_CPU_DICT[m.group(1)] = zone._make([int(m.group(3)), 
                                                             int(m.group(4)),
                                                             int(m.group(5)),
                                                             int(m.group(6))])
    if first_call:
        Instance.ZONE_CPU = zone_cpu
        max_i = Instance.ZONE_CPU_DICT[zone_cpu[0]].max  
        if opts["max_instances"]:
            Instance.MAX_INSTANCES = max_i
        while max_i > 0:
            for zone in zone_cpu:
                if max_i >= Instance.ZONE_CPU_DICT[zone].cpu:
                    max_i -= Instance.ZONE_CPU_DICT[zone].cpu
                    Instance.MAX_INSTANCES += 1
        if overtime:
            reactor.callLater(int(overtime)*60, TerminaTestedInstances, instances) 
    output = admin_user.getProcessOutput(euca_describe_images, args= [])
    output.addCallback(availableRessource, instances, 
                       admin_user, available_types)
      
def availableRessource(output, instances, admin_user, available_types):
    # Get a list of available EMI
    available_emi = []
    #logging.debug("describe-images: %s" % (output))
    for l in output.split("\n"):
        m = re.search(r"IMAGE\s+(emi.*?)\s+", l)
        if m : available_emi.append(m.group(1))
    logging.debug("Available VM types: %s" % (available_types))
    logging.debug("Available EMI: %s" % (available_emi))
    [ Instance.VMTYPE_USED[t] for t in available_types ]
    [ Instance.EMI_USED[i] for i in available_emi ]
    # Start one instance of an emi that has ressource available
    emi           = None
    vmtype        = None
    ressources_ok = True
    try:
        emi    = [ i for i in available_emi if Instance.EMI_USED[i] \
                   == min(Instance.EMI_USED.values())][0]
        vmtype = [ t for t in available_types if Instance.VMTYPE_USED[t] \
                   == min(Instance.VMTYPE_USED.values())][0]
        key    = [ k for k in admin_user.keypairs if Instance.KEYPAIR_USED[k] \
                   == min(Instance.KEYPAIR_USED.values())][0]
        group  = [ g for g in admin_user.groups if Instance.GROUP_USED[g] \
                   == min(Instance.GROUP_USED.values())][0]
    except IndexError:
        logging.info("Not enough ressource available to start a new EMI")
        ressources_ok = False
    
    if opts["max_instances"] and Instance.ZONE_CPU[0] in available_types:
        vmtype = Instance.ZONE_CPU[0]
        
    if emi and vmtype:
        logging.info("Creating new instance")
        i = Instance(emi = emi, vmtype = vmtype, 
                     user = admin_user, ssh_key = key, group = group)
        instances.append(i)
        Instance.EMI_USED[emi]       += 1
        Instance.VMTYPE_USED[vmtype] += 1
        Instance.KEYPAIR_USED[key]   += 1
        Instance.GROUP_USED[group]   += 1
        i.start()
        
    if not opts['start_stop']:
        if (ressources_ok and len(instances) < Instance.MAX_INSTANCES):
            reactor.callLater(random.randint(20, 30) + random.random(),
                          checkRessources, instances, admin_user)
        else:
            logging.info("Max instances reached. Stop checking for %s" %
                         "available ressource.")
            reactor.callLater(30, TerminaTestedInstances, instances)
    elif overtime:
        reactor.callLater(random.randint(20, 30) + random.random(),
              checkRessources, instances, admin_user)
        
def TerminaTestedInstances(instances):   
    global overtime 
    overtime = False
    logging.debug("Check if all the instances have been tested")    
    if len(instances) == len(filter(lambda i: i.test_state in \
                                    ( "success", "failed", "lost" ), 
                                    instances)):      
        logging.info("Terminate all tested instances")        
        for i in instances:
            if i.state not in 'terminated': i.terminate()
    if len(instances) == len(filter(lambda i: i.state == 'terminated',
                                    instances)):
        logging.info("Test run completed")
        reactor.stop()
    else:
        reactor.callLater(random.randint(10, 15) + random.random(),
                  TerminaTestedInstances, instances)
            
def checkInstancesState(instances, admin_user):
    if not admin_user.ready:
        reactor.callLater(15, checkInstancesState, instances, admin_user)
        return
    logging.debug("Checking instances state")
    output = admin_user.getProcessOutput(euca_describe_instances)
    output.addCallback(instancesState, instances, admin_user)

def instancesState(output, instances, admin_user):
    logging.debug("describe instance: %s", output)
    for l in output.split("\n"):
        if l.startswith('INSTANCE'):
            try:
                i_id = l.split()[1]
                i = filter(lambda i: i.id == i_id, instances)[0]
            except IndexError:
                #logging.debug("Unknown instance %s - skipping" % (i_id))
                continue
            state =  l.split()[5]
            i.state = state
            pub_ip = l.split()[3]
            if pub_ip != '0.0.0.0':
                i.logger.debug("Setting pub ip (%s)" % (pub_ip))
                i.pub_ip = pub_ip
            m = re.search(r"\s+(\w+)\s+eki-", l)      
            if m: i.zone = m.group(1)
    # Stop when the number of instances have been run and tested
    #if config['max_instances_to_start'] != 0 \
    #   and len(instances) >= config['max_instances_to_start'] \
    #   and len(filter(lambda i: i.state != 'terminated', instances)) == 0:
    #    logging.info("Test run completed")
    #    reactor.stop()
    #else:
    reactor.callLater(random.randint(10, 20) + random.random(),
                          checkInstancesState, instances, admin_user)        
        
def cleanUp(instances):
    logging.info("Cleaning up")
    logging.debug("Cleaning instances: %s" % (instances))
    [ i.terminate() for i in instances ]
    logging.debug("Cleaning user: %s" % (admin_user.user_id))
    admin_user.cleanup()

def printStats(instances, once=False):
    stats = { 'started' : len(instances) }
    for s in Instance.TEST_STATES:
        stats[s] = len(filter(lambda i: i.test_state == s, instances))
        try:
            stats["%s_rate" % (s)] = float("%.2f" \
                                     % (float(stats[s])/stats['started']))
        except ZeroDivisionError:
            stats["%s_rate" % (s)] = 0.00
        stats["%s_instances" % (s)] = \
             [ i.id for i in instances if i.test_state == s ]
    logging.getLogger('STATS').info("Stats: %s" % (yaml.dump(stats)))
    if not once:
        reactor.callLater(60, printStats, instances)        
   
def printFinalStats(instances):
    stats = { 'started' : len(instances) }
    for s in Instance.TEST_STATES:
        stats[s] = len(filter(lambda i: i.test_state == s, instances))
        try:
            stats["%s_rate" % (s)] = float("%.2f" \
                                     % (float(stats[s])/stats['started']))
        except ZeroDivisionError:
            stats["%s_rate" % (s)] = 0.00
    for s in sorted(stats.keys()):
        print "%s : %s" % (s, stats[s])      
        
def main():
    LEVELS = {'debug': logging.DEBUG,
               'info': logging.INFO,
               'warning': logging.WARNING,
               'error': logging.ERROR,
               'critical': logging.CRITICAL}
    global opts 
    global admin_user
    global overtime
    opts = Options()
   
    try:
        opts.parseOptions() # When given no argument, parses sys.argv[1:]
    except usage.UsageError, errortext:
        print '%s: %s' % (sys.argv[0], errortext)
        print '%s: Try --help for usage details.' % (sys.argv[0])
        return 1 

    # Verify that script is run as root
    if os.getuid():
        sys.stderr.write(
            "This script needs superuser permissions to run correctly\n")
        return 1        
        
    level = LEVELS.get(opts['log'], logging.NOTSET)
   
    if opts['write_log']:
        suffix = "all_vmtypes"
        if opts["max_instances"]:
            suffix = "max_instances"
        elif opts["start_stop"]:
            suffix = "start_stop"
        LOG_FILENAME = "/var/log/uec_instances.%s.log" %suffix
        handler = logging.handlers.RotatingFileHandler(LOG_FILENAME,
                                                       backupCount=2)
        logging.getLogger().addHandler(handler)
        handler.setFormatter(
            logging.Formatter("[%(asctime)s] %(levelname)s %(message)s"))
        logging.getLogger().setLevel(level)
        handler.doRollover() 
    else:
        logging.basicConfig(level=level)              
    logging.info("log level set to " + opts['log'])
    instances = []   
    admin_user = User() # Create admin user
    if opts["start_stop"]:
        overtime = opts["start_stop"]
    if opts["cleanup"]:
        cleanUp(instances)
    else:
        reactor.callWhenRunning(admin_user.setup) # setup admin user
        reactor.callWhenRunning(checkRessources, instances, admin_user, 
                                first_call = True)
        reactor.callWhenRunning(checkInstancesState, instances, admin_user)
        printStats(instances)
        reactor.run()
        cleanUp(instances)
        printStats(instances, once=True)
        printFinalStats(instances)
        if len(filter(lambda i: i.test_state == 'success', instances)) != \
           len(instances):
            return 1
        else:
            return 0

if __name__ == "__main__":
    sys.exit(main())