<commit_before>"""
pbs_bullet

A script to keep an eye on HPC jobs using the PBS system, and
send pushbullet notifications on a few events.

I have not tested this. Use with extreme caution.

<NAME>
<EMAIL>
"""

import sys
import argparse
import json
import urllib2, urllib
import logging
from time import sleep

logging.basicConfig()
logger = logging.getLogger(__name__)

try:
    from subprocess import check_output
except:
    logger.error("pbs_bullet uses the check_output command, added in python 2.7.")
    sys.exit(1)

def arguments():
    parser = argparse.ArgumentParser(
        description="Watch a PBS job. Can optionally send notifications by pushbullet, and kill the job if low on memory.")
    parser.add_argument('jobid', help='Job id to monitor.')
    parser.add_argument('--pushbullet-token', type=str, dest="pb_token", default=None)
    parser.add_argument('--notify-on', type=str, nargs='*', dest="notify_on",
        choices=["start", "finish", "killed", "fail"], default=["start", "finish", "killed", "fail"],
        help="Events to send a notification on.")
    parser.add_argument('--low-mem', type=float, dest="kill_threshold", default=0.,
        help="Kill the job using qdel if free memory on any one node drops below this \%.")
    parser.add_argument('--poll-interval', type=int, dest="poll_interval", default=300,
        help="Number of seconds to wait between each check. Defaults to 5 minutes.")
    parser.add_argument('--log-level', dest='log_level', type=str, choices=['debug',
        'info', 'warning', 'error'], default='info', nargs="?")
    args, extras = parser.parse_known_args()
    return args


def parse_job(jobdetails):
    """
    Turn the output of qstat -f into a dictionary.
    """
    lines = jobdetails.replace("\n\t", "").splitlines()[1:-1]
    return dict(map(lambda line: tuple(line.strip().split(" = ")), lines))

def get_nodes(job_dict):
    """
    Return a list of the nodes in use from a job dictionary.
    """

    nodes = job_dict['exec_host']
    return set(map(lambda x: x.split('/')[0], nodes.split('+')))

def kill_job(jobid):
    """
    Attempt to kill a job using qdel.
    """

    return subprocess.call(['qdel', jobid])

def check_free(node):
    """
    Use rsh and free to get the percentage of free memory on a node.
    """

    return subprocess.check_output(["rsh", node, "free", "|",  "awk",  "'FNR == 3 {print $4/($3+$4)*100}'"])

def start_notify(jobid, jobdetails, nodes, pb_token):
    """
    Send a notification that the job has started.
    """
    title = "%s, id: %s, started." % (jobdetails['Job_Name'], str(jobid))
    body = "Running on nodes %s, and due to finish %s." % (", ".join(nodes), jobdetails['etime']) 
    send_notification(title, body, pb_token)

def finish_notify(jobid, jobdetails, pb_token):
    """
    Send a notification that the job has completed.
    """
    title = "%s, id: %s, finished." % (jobdetails['Job_Name'], str(jobid))
    body = ""
    send_notification(title, body, pb_token)

def kill_notify(jobid, jobdetails, nodes, freemem, pb_token):
    """
    Send a notification that the job is being killed.
    """
    title = "Attempting to kill job %s, id: %s." % (jobdetails['Job_Name'], str(jobid))
    body = ["Free memory on nodes: "]
    body += map(lambda (node, mem): "%s - %f\%", zip(nodes, freemem))
    send_notification(title, "\n".join(body), pb_token)

def send_notification(title, body, token):
    """
    Send a pushbullet notification.
    """
    note = json.dumps({"type":"note", "title": title, "body": body})
    logger.debug("Sending %s to pushbullet." % note)
    request = urllib2.Request('https://api.pushbullet.com/v2/pushes', note, headers={
        'Authorization':"Bearer %s" % token,
        'Content-Type':'application/json',
        'Accept':'*/*'
    })
    try:
        resp = urllib2.urlopen(request)
    except urllib2.HTTPError as e:
        logger.error("Pushbullet notify error.")
        logger.error(e.read())

def main():
    args = arguments()
    numeric_level = getattr(logging, args.log_level.upper(), None)
    if not isinstance(numeric_level, int):
        raise ValueError('Invalid log level: %s' % log_level)

    logger.setLevel(numeric_level)
    jobid = args.jobid
    pb_token = args.pb_token
    sleep_time = args.poll_interval
    notify_on = args.notify_on
    lowmem = args.kill_threshold

    started = False
    finished = False

    while not finished:
        try:
            logger.debug("Checking status for job %s" % jobid)
            jobdetails = parse_job(subprocess.check_output(['qstat', '-f', jobid]))
        except Exception as e:
            logger.error('qstat command failed. Bailing out.')
            logger.error('Error was:')
            logger.error(e)<commit_msg>Added methods for registering as a pb device<commit_after>"""
pbs_bullet

A script to keep an eye on HPC jobs using the PBS system, and
send pushbullet notifications on a few events.

I have not tested this. Use with extreme caution.

<NAME>
<EMAIL>
"""

import sys
import argparse
import json
import urllib2, urllib
import logging
from time import sleep

logging.basicConfig()
logger = logging.getLogger(__name__)

try:
    from subprocess import check_output
except:
    logger.error("pbs_bullet uses the check_output command, added in python 2.7.")
    sys.exit(1)

def arguments():
    parser = argparse.ArgumentParser(
        description="Watch a PBS job. Can optionally send notifications by pushbullet, and kill the job if low on memory.")
    parser.add_argument('jobid', help='Job id to monitor.')
    parser.add_argument('--pushbullet-token', type=str, dest="pb_token", default=None)
    parser.add_argument('--notify-on', type=str, nargs='*', dest="notify_on",
        choices=["start", "finish", "killed", "fail"], default=["start", "finish", "killed", "fail"],
        help="Events to send a notification on.")
    parser.add_argument('--low-mem', type=float, dest="kill_threshold", default=0.,
        help="Kill the job using qdel if free memory on any one node drops below this \%.")
    parser.add_argument('--poll-interval', type=int, dest="poll_interval", default=300,
        help="Number of seconds to wait between each check. Defaults to 5 minutes.")
    parser.add_argument('--log-level', dest='log_level', type=str, choices=['debug',
        'info', 'warning', 'error'], default='info', nargs="?")
    args, extras = parser.parse_known_args()
    return args


def parse_job(jobdetails):
    """
    Turn the output of qstat -f into a dictionary.
    """
    lines = jobdetails.replace("\n\t", "").splitlines()[1:-1]
    return dict(map(lambda line: tuple(line.strip().split(" = ")), lines))

def get_nodes(job_dict):
    """
    Return a list of the nodes in use from a job dictionary.
    """

    nodes = job_dict['exec_host']
    return set(map(lambda x: x.split('/')[0], nodes.split('+')))

def kill_job(jobid):
    """
    Attempt to kill a job using qdel.
    """

    return subprocess.call(['qdel', jobid])

def check_free(node):
    """
    Use rsh and free to get the percentage of free memory on a node.
    """

    return subprocess.check_output(["rsh", node, "free", "|",  "awk",  "'FNR == 3 {print $4/($3+$4)*100}'"])

def start_notify(jobid, jobdetails, nodes, pb_token):
    """
    Send a notification that the job has started.
    """
    title = "%s, id: %s, started." % (jobdetails['Job_Name'], str(jobid))
    body = "Running on nodes %s, and due to finish %s." % (", ".join(nodes), jobdetails['etime']) 
    send_notification(title, body, pb_token)

def finish_notify(jobid, jobdetails, pb_token):
    """
    Send a notification that the job has completed.
    """
    title = "%s, id: %s, finished." % (jobdetails['Job_Name'], str(jobid))
    body = ""
    send_notification(title, body, pb_token)

def kill_notify(jobid, jobdetails, nodes, freemem, pb_token):
    """
    Send a notification that the job is being killed.
    """
    title = "Attempting to kill job %s, id: %s." % (jobdetails['Job_Name'], str(jobid))
    body = ["Free memory on nodes: "]
    body += map(lambda (node, mem): "%s - %f\%", zip(nodes, freemem))
    send_notification(title, "\n".join(body), pb_token)

def send_notification(title, body, token):
    """
    Send a pushbullet notification.
    """
    note = json.dumps({"type":"note", "title": title, "body": body})
    logger.debug("Sending %s to pushbullet." % note)
    request = urllib2.Request('https://api.pushbullet.com/v2/pushes', note, headers={
        'Authorization':"Bearer %s" % token,
        'Content-Type':'application/json',
        'Accept':'*/*'
    })
    try:
        resp = urllib2.urlopen(request)
    except urllib2.HTTPError as e:
        logger.error("Pushbullet notify error.")
        logger.error(e.read())

def create_listener(name, token):
    """
    Register this device to receive pushbullet notifications,
    and if successful return the identifier.
    """
    data = urllib.urlencode({'nickname':name, 'type':'streaming'})
    logger.debug("Sending %s to pushbullet." % data)
    request = urllib2.Request('https://api.pushbullet.com/v2/devices', data, headers={
        'Authorization':"Bearer %s" % token,
        'Accept':'*/*'
    })
    try:
        resp = urllib2.urlopen(request)
        return json.load(resp)['iden']
    except urllib2.HTTPError as e:
        logger.error("Pushbullet register error.")
        logger.error(e.read())

def delete_listener(iden, token):
    """
    Unregister this device as a listener.
    """
    logger.debug("Sending %s to pushbullet." % data)
    request = urllib2.Request('https://api.pushbullet.com/v2/devices/%s' % iden, headers={
        'Authorization':"Bearer %s" % token,
        'Accept':'*/*'
    })
    request.get_method = lambda: 'DELETE'
    try:
        resp = urllib2.urlopen(request)
        return json.load(resp)['iden']
    except urllib2.HTTPError as e:
        logger.error("Pushbullet delete error.")
        logger.error(e.read())

def main():
    args = arguments()
    numeric_level = getattr(logging, args.log_level.upper(), None)
    if not isinstance(numeric_level, int):
        raise ValueError('Invalid log level: %s' % log_level)

    logger.setLevel(numeric_level)
    jobid = args.jobid
    pb_token = args.pb_token
    sleep_time = args.poll_interval
    notify_on = args.notify_on
    lowmem = args.kill_threshold

    started = False
    finished = False

    while not finished:
        try:
            logger.debug("Checking status for job %s" % jobid)
            jobdetails = parse_job(subprocess.check_output(['qstat', '-f', jobid]))
        except Exception as e:
            logger.error('qstat command failed. Bailing out.')
            logger.error('Error was:')
            logger.error(e)
{'fraction_non_alphanumeric': 0.14818595156351488, 'fraction_numerical': 0.005010670873155795, 'mean_word_length': 3.3025948103792415, 'pattern_counts': {'":': 6, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 4, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6238842', 'n_tokens_mistral': 3406, 'n_tokens_neox': 3364, 'n_words': 969}