jobinfo

#!/usr/bin/env python2

# jobinfo


import math
import os
import re
import subprocess
import sys
import time
from collections import namedtuple as NT

# Define utilities

def parse_time(t):
    # Format: [DD-[hh:]]mm:ss
    time_parts = re.compile(r'(((?P<days>\d+)-)?(?P<hours>\d\d):)?' +
                            r'(?P<minutes>\d\d):(?P<seconds>\d\d(\.\d+)?)')
    m = time_parts.match(t)
    if m is None:
        return 0.0, 0, 0, 0
    ss = float(m.group('seconds'))
    mm = int(m.group('minutes'))
    hh = int(m.group('hours') or '0')
    dd = int(m.group('days') or '0')
    return ss, mm, hh, dd

def elapsed_to_seconds(elapsed):
    ss, mm, hh, dd = parse_time(elapsed)
    return dd*24*60*60 + hh*60*60 + mm*60 + ss

def format_bs(x):
    postfix = ' KMGTPE'
    e = int(math.log(x+1, 2)/10)
    return "%.2f%s" % (x / 2**(10 * e), postfix[e])

def mem_to_kb(x):
    if x.endswith('G'):
        req = float(x[:-1])
        return req * 1024 * 1024
    elif x.endswith('M'):
        req = float(x[:-1])
        return req * 1024
    elif x.endswith('K'):
        req = float(x[:-1])
        return req
    else:
        return "??"

def totalmem(x, meta):
    if x.endswith('c'):
        req = "%s" % x[:-1]
        creq = mem_to_kb(req)
        return creq * int(meta.ncpus)
    elif x.endswith('n'):
        req = "%s" % x[:-1]
        creq = mem_to_kb(req)
        return creq * int(meta.NNodes)
    else:
        creq = mem_to_kb(x)
        return creq

def whoami():
    import pwd
    return pwd.getpwuid(os.getuid()).pw_name

def without_fraction(s):
    return "." in s and s[:s.rfind(".")] or s

# Define constructors

def str_set(x=None):
    if x in [None, '']:
        return set()
    return set([x])

def gpu_str(s=None):
    # Pending job
    if s is '':
        return str("--")
    # No GPUs requested
    if s is None or "gres/gpu" not in s.lower():
        return 0
    # GPUs requested
    # Number of GPUs requested
    g = re.match(r".*gres/gpu=(?P<num_gpus>\d+)", s.lower())
    g_str = g.group('num_gpus')
    # Type of GPU requested
    t = re.match(r".*gres/gpu:(?P<gpu_type>\w\d+)", s.lower())
    t_str = t.group('gpu_type')
    return g_str + " (" + t_str + ")"

def byte_size(s=None):
    if s in [None, "", "16?"]:
        return -1.0
    m = {'K': 10, 'M': 20, 'G': 30, 'T': 40, 'P': 50, 'E': 60}
    scale = 2 ** m.get(s[-1], 0)
    if scale != 1:
        s = s[:-1]
    return scale * float(s)

def date_str(s=None):
    if s is None or s.strip() == "":
        return "9999-01-01T00:00:00"
    return s

# Define combinators

def keep_first(a, b):
    return a == '' and b or a

def time_max(a, b):
    if 'UNLIMITED' in [a, b]:
        return 'UNLIMITED'
    if a in ['', 'INVALID']:
        return b
    if b in ['', 'INVALID']:
        return a
    return max(a, b)

# Define formatters

def f_tasks(x, meta):
    if not x:
        return "--"
    return str(x)

def f_drw(x, meta):
    if x < 0:
        return "--"
    return "%s" % format_bs(x)

def f_cpu(x, meta):
    total = elapsed_to_seconds(meta.TotalCPU)
    if total == 0:
        return "--"
    xs = elapsed_to_seconds(x)
    return "%5.2f%%" % (xs / total * 100)

def f_cpueff(x, meta):
    total = elapsed_to_seconds(meta.TotalCPU)
    xint = int(x)
    if xint == 0 or total == 0:
        return "--"
    return "%5.2f%%" % (total / xint * 100)

def f_mem(x, meta):
    if x.endswith('c'):
        return "%s/core" % (x[:-1])
    elif x.endswith('n'):
        return "%s/node" % (x[:-1])
    else:
        return x

def f_rss(x, meta):
    if x < 0:
        return "--"
    totalx = x * int(meta.NTasks)
    return "%s (%s)" % (format_bs(totalx), "estimate")

def f_memeff(x, meta):
    if x < 0:
        return "--"
    max = x / 1024 * int(meta.NTasks)
    req = str(meta.ReqMem)
    creq = totalmem(req, meta)
    if str(creq).endswith('?'):
        return "??"
    return "%5.2f%%" % (max / creq * 100)

def f_time(x, meta):
    all_times = [meta.timelimit, meta.elapsed, meta.TotalCPU, meta.planned]
    max_time_len = max(len(without_fraction(y)) for y in all_times)
    ss, mm, hh, dd = parse_time(x)
    if dd > 0:
        dd = ("%i-" % dd)
    else:
        dd = ""
    res = "%s%02i:%02i:%02i" % (dd, hh, mm, ss)
    if res.strip() == "00:00:00" and meta.start.lower() == "unknown":
        return "--"
    return res.rjust(max_time_len)

def f_str(x, meta):
    return str(x)

def f_date(x, meta):
    if str(x).lower() == "unknown":
        return "--"
    return str(x)

def f_exit(x, meta):
    if meta.end.lower() == "unknown":
        return "--"
    return x

def f_state(states, meta):
    if len(states) > 1:
        states = states - set(["COMPLETED", ""])
    reason = meta.reason
    if reason != '':
        reason = ' ' + reason
    deps = meta.dependencies
    if deps != '':
        deps = " (%s)" % deps
    return ','.join(states) + reason + deps

# Define fields and how they will be processed

hide, show = False, True
Field = NT('Field', 'name ctor combinator shown prefer_sstat formatter desc')
FIELDS = [
        Field("JobID",         str,        keep_first,   show,  False,  f_str,     "Job ID"),
        Field("JobName",       str,        keep_first,   show,  False,  f_str,     "Job name"),
        Field("User",          str,        keep_first,   show,  False,  f_str,     "User"),
        Field("Account",       str,        keep_first,   show,  False,  f_str,     "Account"),
        Field("WorkDir",       str,        keep_first,   show,  False,  f_str,     "Working directory"),
        Field("Cluster",       str,        keep_first,   show,  False,  f_str,     "Cluster"),
        Field("Partition",     str,        keep_first,   show,  False,  f_str,     "Partition"),
        Field("NNodes",        str,        keep_first,   show,  False,  f_str,     "Nodes"),
        Field("NodeList",      str,        keep_first,   show,  False,  f_str,     "Nodelist"),
        Field("NTasks",        str,        max,          show,  False,  f_tasks,   "Tasks"),
        Field("ncpus",         int,        max,          show,  False,  f_str,     "CPUs"),
        Field("AllocTRES",     gpu_str,    keep_first,   show,  False,  f_str,     "GPUs"),
        Field("State",         str_set,    set.union,    show,  False,  f_state,   "State"),
        Field("ExitCode",      str,        keep_first,   show,  False,  f_exit,    "Exit code"),
        Field("Submit",        str,        keep_first,   show,  False,  f_str,     "Submit time"),
        Field("start",         date_str,   min,          show,  False,  f_date,    "Start time"),
        Field("end",           str,        time_max,     show,  False,  f_date,    "End time"),
        Field("planned",       date_str,   min,          show,  False,  f_time,    "Wait time"),
        Field("timelimit",     str,        time_max,     show,  False,  f_time,    "Reserved walltime"),
        Field("elapsed",       str,        time_max,     show,  False,  f_time,    "Used walltime"),
        Field("CPUTime",       str,        keep_first,   show,  False,  f_time,    "Used CPU walltime"),
        Field("TotalCPU",      str,        keep_first,   show,  False,  f_time,    "Used CPU time"),
        Field("UserCPU",       str,        keep_first,   show,  False,  f_cpu,     "% User (computation)"),
        Field("SystemCPU",     str,        keep_first,   show,  False,  f_cpu,     "% System (I/O)"),
        Field("CPUTimeRaw",    str,        keep_first,   show,  False,  f_cpueff,  "CPU efficiency"),
        Field("ReqMem",        str,        keep_first,   show,  False,  f_mem,     "Reserved memory"),
        Field("MaxRSS",        byte_size,  max,          show,  True,   f_rss,     "Max memory used"),
        Field("maxrss",        byte_size,  max,          show,  True,   f_memeff,  "Memory efficiency"),
        Field("MaxDiskWrite",  byte_size,  max,          show,  True,   f_drw,     "Max disk write"),
        Field("MaxDiskRead",   byte_size,  max,          show,  True,   f_drw,     "Max disk read")
        ]

FIELD_NAMES = [f.name for f in FIELDS]
FIELD_NAMES_SSTAT = [f.name for f in FIELDS if f.prefer_sstat]
FIELD_CTORS = [f.ctor for f in FIELDS]
FIELD_COMB = [f.combinator for f in FIELDS]
FORMAT_STR = "--format=%s" % (",".join(FIELD_NAMES))
FORMAT_SSTAT_STR = "--format=%s" % (",".join(FIELD_NAMES_SSTAT))
Meta = NT('Meta', FIELD_NAMES + ['dependencies', 'reason'])

# Get, process, and combine values

def combine(xs):
    r = xs[0]
    for x in xs[1:]:
        for i, comb in enumerate(FIELD_COMB):
            r[i] = comb(r[i], x[i])
    return r

def get_values(jobid):
    info = subprocess.Popen(['sacct', FORMAT_STR, '--parsable', '--noheader', '-j', jobid], stdout=subprocess.PIPE)
    xs = []
    for line in info.stdout:
        fields = line.strip().split('|')
        # Skip <job_id>.extern values (not needed)
        if fields[0].endswith(".extern"):
            continue
        xs.append([ctor(s) for ctor, s in zip(FIELD_CTORS, fields)])
    if len(xs) == 0:
        print >>sys.stderr, "Error: No such job"
        sys.exit(1)
    return xs

def get_sstat_values(jobid):
    info = subprocess.Popen(['sstat', FORMAT_SSTAT_STR, '--parsable', '--noheader', '-a', '-j', jobid], stdout=subprocess.PIPE)
    xs = []
    for line in info.stdout:
        j = 0
        vals = line.strip().split('|')
        x = []
        for f in FIELDS:
            if f.prefer_sstat:
                x.append(f.ctor(vals[j]))
                j += 1
            else:
                x.append(f.ctor())
            xs.append(x)
    return xs

def main(jobid):
    y = combine(get_values(jobid))
    meta = Meta._make(y + ['', ''])
    ys = [y]
    if "RUNNING" in meta.State and (os.getuid() == 0 or meta.User == whoami()):
        tmp = get_sstat_values("%s,%s.batch" % (jobid, jobid))
        if len(tmp) != 0:
            ys.append(combine(tmp))
    if "PENDING" in meta.State:
        info = subprocess.Popen(['squeue', '--format=%E;%R', '--noheader', '-a', '-j', jobid], stdout=subprocess.PIPE)
        dependencies, reason = info.stdout.readline().strip().split(";")
        if dependencies == "(null)":
            dependencies = ""
    else:
        dependencies = ""
        reason = ""
    y = combine(ys)
    meta = Meta._make(y + [dependencies, reason])

    for i,(name,parse,comb,show,prefer_sstat,format,desc) in enumerate(FIELDS):
        val = y[i]
        if show:
            print("%-20s : %s" % (desc, format(val, meta)))

def usage(pipe):
    print >>pipe, \
"""jobinfo / Display detailed job information
Usage:   jobinfo <job_id>
Example: jobinfo 111111"""

if __name__ == "__main__":
    if "-h" in sys.argv or "--help" in sys.argv:
        usage(sys.stdout)
        sys.exit(0)
    if len(sys.argv) != 2:
        print >>sys.stderr, "Error: There should be one argument (job ID)"
        usage(sys.stderr)
        sys.exit(1)
    jobid = sys.argv[1]
    if len(set(jobid) - set("0123456789_+.")) > 0:
        print >>sys.stderr, "Error: The argument does not look like a valid job ID"
        usage(sys.stderr)
        sys.exit(1)
    main(jobid)