report-gpu-utilization

#!/usr/bin/python

import subprocess
import sys
import re
import syslog
import time

#  Globals
nvidia_smi_headers_regexp = '.*?Timestamp *: ([^\n]*)(.*)'
nvidia_smi_stanzas_regexp = '.*?Process ID *: (\d+).*?GPU Utilization *: (\d+).*?Memory Utilization *: (\d+)(.*)'
proc_pid_environ_regexp = '(?:^|\0)JOB_ID=(\d+)'
verboselevel = 2
syslog_opened_flag = False

def main():
    global nvidia_smi_regexp

    #  Run nvidia-smi
    nvidia_smi_output = subprocess.Popen("nvidia-smi -q -d ACCOUNTING", stdout=subprocess.PIPE, shell=True).communicate(None)[0]

    #  Extract time and convert to timestamp and preserve rest of output.
    m = re.search(nvidia_smi_headers_regexp, nvidia_smi_output, flags=re.MULTILINE|re.DOTALL)
    timestamp = int(time.mktime(time.strptime(m.group(1))))
    nvidia_smi_output = m.group(2)

    #  Chop up GPU stanzas, deriving SGE job ID and extracting utilisation values and
    #  displaying values.
    while True:
        m = re.search(nvidia_smi_stanzas_regexp, nvidia_smi_output, flags=re.MULTILINE|re.DOTALL)
        if m is None:
            break
        pid               = int(m.group(1))
        gpu_utilisation   = int(m.group(2))
        mem_utilisation   = int(m.group(3))
        nvidia_smi_output = m.group(4)
        jobid             = pid2jobid(pid)
        sys.stdout.write('%d,%d,%d,%d\n' % (timestamp, jobid, gpu_utilisation, mem_utilisation))

#  Support functions
def pid2jobid(pid):
    global proc_pid_environ_regexp
    f = open('/proc/%d/environ' % (pid)).read()
    m = re.search(proc_pid_environ_regexp, f, flags=re.MULTILINE|re.DOTALL)
    jobid = int(m.group(1))
    debug(10, 'pid2jobid: pid=%d, jobid=%d' % (pid, jobid))
    return(jobid)

#  Messaging functions
def debug(level, text):
    global verboselevel
    if verboselevel < level:
        return
    msg_writer("DEBUG[%s]: %s" % (level, text), syslog.LOG_ERR)

def info(text):
    global verboselevel
    if verboselevel < 3:
        return
    msg_writer("INFO: %s" % (text), syslog.LOG_ERR)

def warning(level, text):
    global verboselevel
    if verboselevel < 2:
        return
    msg_writer("WARNING: %s" % (text), syslog.LOG_ERR)

def error(text):
    msg_writer("ERROR: %s" % (text), syslog.LOG_ERR)
    sys.exit(1)

def msg_writer(text, syslog_level):
    msg_writer_stderr(text)
    msg_writer_syslog(text, syslog_level)

def msg_writer_stderr(text):
    global progname
    sys.stderr.write('%s: %s\n' % (progname, text))

def msg_writer_syslog(text, syslog_level):
    global syslog_opened_flag, progname
    if not syslog_opened_flag:
        if sys.version_info < (2, 7):
            syslog.openlog(progname, syslog.LOG_PID)
        else:
            syslog.openlog(progname, logoption=syslog.LOG_PID)
        syslog_opened_flag = True
    syslog.syslog((syslog.LOG_LOCAL0|syslog_level), text)

#  This code is for debugging and unit testing
if __name__ == '__main__':
    progname = sys.argv[0].rpartition("/")[2]
    main()