<NME> kill.py
<BEF> #!/usr/bin/python
import os
import sys
import yaml
import psutil
import subprocess
import tempfile
import logging
import getpass

from . import beanstalk
from . import report
from .config import config

log = logging.getLogger(__name__)



def main(args):
    run_name = args['--run']
    job = args['--job']
    jobspec = args['--jobspec']
    archive_base = args['--archive']
    owner = args['--owner']
    machine_type = args['--machine-type']
    preserve_queue = args['--preserve-queue']

    if jobspec:
        split_spec = jobspec.split('/')
        run_name = split_spec[0]
        job = [split_spec[1]]

    if job:
        for job_id in job:
            kill_job(run_name, job_id, archive_base, owner)
    else:
        kill_run(run_name, archive_base, owner, machine_type,
                 preserve_queue=preserve_queue)


def kill_run(run_name, archive_base=None, owner=None, machine_type=None,
             preserve_queue=False):
    run_info = {}
    serializer = report.ResultsSerializer(archive_base)
    if archive_base:
        run_archive_dir = os.path.join(archive_base, run_name)
        if os.path.isdir(run_archive_dir):
            run_info = find_run_info(serializer, run_name)
            if 'machine_type' in run_info:
                machine_type = run_info['machine_type']
                owner = run_info['owner']
            else:
                log.warning("The run info does not have machine type: %s" % run_info)
                log.warning("Run archive used: %s" % run_archive_dir)
                log.info("Using machine type '%s' and owner '%s'" % (machine_type, owner))
        elif machine_type is None:
            # no jobs found in archive and no machine type specified,
            # so we try paddles to see if there is anything scheduled
            run_info = report.ResultsReporter().get_run(run_name)
            machine_type = run_info.get('machine_type', None)
            if machine_type:
                log.info(f"Using machine type '{machine_type}' received from paddles.")
            else:
                raise RuntimeError(f"Cannot find machine type for the run {run_name}; " +
                                    "you must also pass --machine-type")

    if not preserve_queue:
        remove_beanstalk_jobs(run_name, machine_type)
        remove_paddles_jobs(run_name)
    kill_processes(run_name, run_info.get('pids'))
    if owner is not None:
        targets = find_targets(run_name, owner)
        nuke_targets(targets, owner)


def kill_job(run_name, job_id, archive_base=None, owner=None, skip_nuke=False):
    serializer = report.ResultsSerializer(archive_base)
    job_info = serializer.job_info(run_name, job_id)
    if not owner:
        if 'owner' not in job_info:
            raise RuntimeError(
                "I could not figure out the owner of the requested job. "
                "Please pass --owner <owner>.")
        owner = job_info['owner']
    kill_processes(run_name, [job_info.get('pid')])
    # Because targets can be missing for some cases, for example, when all
    # the necessary nodes ain't locked yet, we do not use job_info to get them,
    # but use find_targets():
    targets = find_targets(run_name, owner, job_id)
    if not skip_nuke:
        nuke_targets(targets, owner)


def find_run_info(serializer, run_name):
    log.info("Assembling run information...")
    run_info_fields = [
        'machine_type',
        'owner',
    ]

    pids = []
    run_info = {}
    job_info = {}
    job_num = 0
    jobs = serializer.jobs_for_run(run_name)
    job_total = len(jobs)
    for (job_id, job_dir) in jobs.items():
        if not os.path.isdir(job_dir):
            continue
        job_num += 1
        beanstalk.print_progress(job_num, job_total, 'Reading Job: ')
        job_info = serializer.job_info(run_name, job_id, simple=True)
        for key in job_info.keys():
            if key in run_info_fields and key not in run_info:
                run_info[key] = job_info[key]
        if 'pid' in job_info:
            pids.append(job_info['pid'])
    run_info['pids'] = pids
    return run_info


def remove_paddles_jobs(run_name):
    jobs = report.ResultsReporter().get_jobs(run_name, fields=['status'])
    job_ids = [job['job_id'] for job in jobs if job['status'] == 'queued']
    if job_ids:
        log.info("Deleting jobs from paddles: %s", str(job_ids))
        report.try_delete_jobs(run_name, job_ids)


def remove_beanstalk_jobs(run_name, tube_name):
    qhost = config.queue_host
    qport = config.queue_port
    if qhost is None or qport is None:
        raise RuntimeError(
            'Beanstalk queue information not found in {conf_path}'.format(
                conf_path=config.yaml_path))
    log.info("Checking Beanstalk Queue...")
    beanstalk_conn = beanstalk.connect()
    real_tube_name = beanstalk.watch_tube(beanstalk_conn, tube_name)

    curjobs = beanstalk_conn.stats_tube(real_tube_name)['current-jobs-ready']
    if curjobs != 0:
        x = 1
        while x != curjobs:
            x += 1
            job = beanstalk_conn.reserve(timeout=20)
            if job is None:
                continue
            job_config = yaml.safe_load(job.body)
            if run_name == job_config['name']:
                job_id = job.stats()['id']
                msg = "Deleting job from queue. ID: " + \
                    "{id} Name: {name} Desc: {desc}".format(
                        id=str(job_id),
                        name=job_config['name'],
                        desc=job_config['description'],
                    )
                log.info(msg)
                job.delete()
    else:
        print("No jobs in Beanstalk Queue")
    beanstalk_conn.close()


def kill_processes(run_name, pids=None):
    if pids:
        to_kill = set(pids).intersection(psutil.pids())
    else:
        to_kill = find_pids(run_name)

    # Remove processes that don't match run-name from the set
    to_check = set(to_kill)
    for pid in to_check:
        if not process_matches_run(pid, run_name):
            to_kill.remove(pid)

    if len(to_kill) == 0:
        log.info("No teuthology processes running")
    else:
        log.info("Killing Pids: " + str(to_kill))
        may_need_sudo = \
            psutil.Process(int(pid)).username() != getpass.getuser()
        if may_need_sudo:
            sudo_works = subprocess.Popen(['sudo', '-n', 'true']).wait() == 0
            if not sudo_works:
                log.debug("Passwordless sudo not configured; not using sudo")
        use_sudo = may_need_sudo and sudo_works
        for pid in to_kill:
            args = ['kill', str(pid)]
            # Don't attempt to use sudo if it's not necessary
            if use_sudo:
                args = ['sudo', '-n'] + args
            subprocess.call(args)


def process_matches_run(pid, run_name):
    try:
        p = psutil.Process(pid)
        cmd = p.cmdline()
        if run_name in cmd and sys.argv[0] not in cmd:
            return True
    except psutil.NoSuchProcess:
        pass
    return False

    to_nuke = []
    for target in targets:
        to_nuke.append(target.split('@')[1].split('.')[0])

    target_file = tempfile.NamedTemporaryFile(delete=False)
    target_file.write(yaml.safe_dump(targets_dict))
    return run_pids


def find_targets(run_name, owner, job_id=None):
    lock_args = [
        'teuthology-lock',
        '--list-targets',
        '--desc-pattern',
        '/' + run_name + '/' + str(job_id or ''),
        '--status',
        'up',
        '--owner',
        owner
    ]
    proc = subprocess.Popen(lock_args, stdout=subprocess.PIPE)
    stdout, stderr = proc.communicate()
    out_obj = yaml.safe_load(stdout)
    if not out_obj or 'targets' not in out_obj:
        return {}

    return out_obj


def nuke_targets(targets_dict, owner):
    targets = targets_dict.get('targets')
    if not targets:
        log.info("No locked machines. Not nuking anything")
        return

    to_nuke = []
    for target in targets:
        to_nuke.append(misc.decanonicalize_hostname(target))

    target_file = tempfile.NamedTemporaryFile(delete=False, mode='w+t')
    target_file.write(yaml.safe_dump(targets_dict))
    target_file.close()

    log.info("Nuking machines: " + str(to_nuke))
    nuke_args = [
        'teuthology-nuke',
        '-t',
        target_file.name,
        '--owner',
        owner
    ]
    nuke_args.extend(['--reboot-all', '--unlock'])

    proc = subprocess.Popen(
        nuke_args,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT)
    for line in proc.stdout:
        line = line.replace(b'\r', b'').replace(b'\n', b'')
        log.info(line.decode())
        sys.stdout.flush()

    os.unlink(target_file.name)

<MSG> More hostname handling fixes

Signed-off-by: Zack Cerza <f801c831581d4150a2793939287636221d62131e@inktank.com>

<DFF> @@ -11,6 +11,7 @@ import getpass
 from . import beanstalk
 from . import report
 from .config import config
+from . import misc
 
 log = logging.getLogger(__name__)
 
@@ -205,7 +206,7 @@ def nuke_targets(targets_dict, owner):
 
     to_nuke = []
     for target in targets:
-        to_nuke.append(target.split('@')[1].split('.')[0])
+        to_nuke.append(misc.decanonicalize_hostname(target))
 
     target_file = tempfile.NamedTemporaryFile(delete=False)
     target_file.write(yaml.safe_dump(targets_dict))

{'addition_count': 2, 'commit_subject': 'More hostname handling fixes', 'deletion_count': 1, 'file_extension': '.py', 'lang': 'py', 'license': 'mit', 'repo_name': 'ceph/teuthology'}