vzpsi.py

#!/usr/bin/python

from collections import OrderedDict
from glob import glob
from multiprocessing import Process, Queue
from optparse import OptionParser
from os import getuid
from time import sleep
from traceback import format_exc

def parse_opts():
    parser = OptionParser(description='show OpenVZ container processes',usage='usage: vzpsi [options] [ctid]')

    parser.add_option('-s', help='sort by (ctid) pid ppid prog mem cpu rbyt rio wbyt wio', metavar='FIELD', default='ctid')
    parser.add_option('-n', help='display top N processes (all)', metavar='NUMBER', default=0, type='int')
    parser.add_option('-p', help='measurement period in seconds (5)', metavar='SECONDS', default=5, type='float')
    parser.add_option('-q', help='suppress output header', metavar='', action='store_true')

    return parser.parse_args()

def get_mem(pids, queue, period=5):
    mem = {}

    for pid in pids:
        mem[pid] = []

    for i in xrange(period):
        for pid in pids:
            try:
                with open('/proc/%s/status' % pid) as file:
                    status = map(str.split, file.readlines())

                mem[pid].append([float(status[18][1]) , float(status[25][1]) ])
            except:
                pass

        if i < period - 1:
            sleep(1)

    for pid in mem.keys():
        if len(mem[pid]) == period:
            rss = 0.0
            swp = 0.0

            for i in xrange(period):
                rss += mem[pid][i][0]
                swp += mem[pid][i][1]

            mem[pid] = {'mem': int(round(rss / period, 0)), 'swp': int(round(swp / period, 0))}
        else:
            del mem[pid]

    queue.put(mem)

    return True

def get_cpu(pids, queue, period=5):
    sys_start = 0
    sys_end = 0
    pid_start = {}
    pid_end = {}
    cpu = {}

    with open('/proc/stat') as file:
        stat = map(str.split, file.readlines())

    for line in stat:
        if line[0].startswith('cpu'):
            if len(line[0]) > 3:
                sys_start += sum(map(float, line[1:]))
        else:
            break

    for pid in pids:
        try:
            with open('/proc/%s/stat' % pid) as file:
                data = file.read().split()

            pid_start[pid] = [float(data[13]), float(data[14])]
        except:
            pass

    sleep(period)

    with open('/proc/stat') as file:
        stat = map(str.split, file.readlines())

    for line in stat:
        if line[0].startswith('cpu'):
            if len(line[0]) > 3:
                sys_end += sum(map(float, line[1:]))
        else:
            break

    for pid in pid_start:
        try:
            with open('/proc/%s/stat' % pid) as file:
                data = file.read().split()

            pid_end[pid]  = [float(data[13]), float(data[14])]
        except:
            pass

    sys_time = sys_end - sys_start

    for pid in pid_start:
        if pid in pid_end:
            cpu[pid] = {'cpu': round(((pid_end[pid][0] - pid_start[pid][0]) + (pid_end[pid][1] - pid_start[pid][1])) / sys_time * 100, 1)}

    queue.put(cpu)

def get_io(pids, queue, period=5):
    period = float(period)
    start = {}
    end = {}
    io = {}

    for pid in pids:
        try:
            with open('/proc/%s/io' % pid) as file:
                data = map(str.split, file.readlines())

            start[pid] = [float(data[2][1]), float(data[4][1]), float(data[3][1]), float(data[5][1])]
        except:
            pass

    sleep(period)

    for pid in start:
        try:
            with open('/proc/%s/io' % pid) as file:
                data = map(str.split, file.readlines())

            end[pid] = [float(data[2][1]), float(data[4][1]), float(data[3][1]), float(data[5][1])]
            io[pid] = {'rbyt': int(round((end[pid][1] - start[pid][1]) / period / 1024.0, 0)),
                       'rio' : int(round((end[pid][0] - start[pid][0]) / period, 0)),
                       'wbyt': int(round((end[pid][3] - start[pid][3]) / period / 1024.0, 0)),
                       'wio' : int(round((end[pid][2] - start[pid][2]) / period, 0))}
        except:
            pass

    queue.put(io)

def get_pids():
    dirs = glob('/proc/[0-9]*')
    pids = {}

    for dir in dirs:
        try:
            with open(dir + '/status') as file:
                status = map(str.split, file.readlines())

            if len(status) == 46:
                with open(dir + '/cmdline') as file:
                    pids[int(status[3][1])] = {'prog': status[0][1], 'ppid': int(status[4][1]), 'ctid': int(status[11][1]), 'cmd': file.read().replace('\x00', ' ')}
        except:
            pass

    return pids

def display(procs, top=0, quiet=False):
    i = 0

    if top is 0:
        top = len(procs)

    if quiet is not True:
        print 'ctid  pid      ppid     prog              cpu      mem      rbyt        rio        wbyt        wio'
        print '----  -------  -------  ----------------  -------  -------  ----------  ---------  ----------  ---------'

    for pid, data in procs.items():
        if i < top:
            if data['mem'] < 1024:
                data['mem'] = '%s KB' % int(data['mem'])
            else:
                data['mem'] = '%s MB' % round(data['mem'] / 1024.0, 1)

            if data['rbyt'] < 1024:
                data['rbyt'] = '%s KB/s' % data['rbyt']
            else:
                data['rbyt'] = '%s MB/s' % int(round(data['rbyt'] / 1024.0, 0))

            if data['wbyt'] < 1024:
                data['wbyt'] = '%s KB/s' % data['wbyt']
            else:
                data['wbyt'] = '%s MB/s' % int(round(data['wbyt'] / 1024.0, 0))

            print ('%s' % data['ctid']).ljust(6) + \
                  ('%s' % pid).ljust(9) + \
                  ('%s' % data['ppid']).ljust(9) +  \
                  data['prog'][:16].ljust(18) + \
                  ('%s %%' % data['cpu']).ljust(9) + \
                  ('%s' % data['mem']).ljust(9) + \
                  ('%s' % data['rbyt']).ljust(12) + \
                  ('%s /s' % data['rio']).ljust(11) + \
                  ('%s' % data['wbyt']).ljust(12) + \
                  ('%s /s' % data['wio']).ljust(11)

            i += 1
        else:
            break

if __name__ == '__main__':
    if getuid() != 0:
        print 'vzspi must be run as root'
        exit(1)
    try:
        data = []
        procs = {}
        pids = get_pids()
        queue = Queue()

        opts, args = parse_opts()

        workers = [Process(target=get_mem, args=(pids, queue, opts.p)),
                   Process(target=get_cpu, args=(pids, queue, opts.p)),
                   Process(target=get_io, args=(pids, queue, opts.p))]

        for worker in workers:
            worker.start()

        for wokrer in workers:
            data.append(queue.get())

        for worker in workers:
            worker.join()

        for pid in pids:
            if pid in data[0] and pid in data[1] and pid in data[2]:
                procs[pid] = dict(data[0][pid].items() + data[1][pid].items() + data[2][pid].items() + pids[pid].items())

        try:
            if args[0]:
                for pid, data in procs.items():
                    if data['ctid'] != int(args[0]):
                        del procs[pid]
        except:
            pass

        if opts.s == 'pid':
            procs = OrderedDict(sorted(procs.items()))

        elif opts.s == 'ctid' or opts.s == 'name' or opts.s == 'ppid':
            try:
                procs = OrderedDict(sorted(procs.items(), key=lambda item: item[1][opts.s]))
            except KeyError:
                print 'invalid sort field'
                exit(2)
        else:
            try:
                procs = OrderedDict(sorted(procs.items(), key=lambda item: item[1][opts.s], reverse=True))
            except KeyError:
                print 'invalid sort field'
                exit(3)

        if len(procs) < opts.n:
            opts.n = len(procs)

        display(procs, opts.n, opts.q)

        exit(0)
    except Exception:
        print format_exc()
        exit(4)