Aggregate memory usage of a process and its children.

mem_usage.py

"""Aggregate memory usage of a process and its children.

This utility measures the resident set size of a process and its children, and
returns the result either in bytes or formatted in a more convenient
format. Note: this "total resident size" is often an inflated measure of how
much memory a process group uses, as it will double-count e.g. memory used by
dynamic libraries and shared memory. It serves well as (a) a conservative
estimate, and (b) when sampled over time, to estimate changes in memory usage.

Requires psutil.

"""

from __future__ import division, print_function
import argparse
import sys

import psutil


if hasattr(sys, 'getwindowsversion'):
    _attr = 'wset'
else:
    _attr = 'rss'


def mem_usage(process):
    return getattr(process.memory_info_ex(), _attr)


def humanize(nbytes):
    suffixes = tuple('bKMGT')
    running = nbytes
    for suffix in suffixes:
        if running < 1024:
            return '{0:.01f}{1}'.format(running, suffix)
        running /= 1024
    return '{0:.01f}{1}'.format(1024 * running, suffixes[-1])


def main():
    parser = argparse.ArgumentParser(
        description='Aggregate memory usage for a process and its children'
    )
    parser.add_argument('rootpid', type=int)
    parser.add_argument('-v', '--verbose', action='store_true')
    parser.add_argument('--human', action='store_true')
    args = parser.parse_args()

    root = psutil.Process(args.rootpid)
    total_rss = sum(mem_usage(child)
                    for child in root.children(recursive=True))

    if args.human:
        formatted_rss = humanize(total_rss)
    else:
        formatted_rss = str(total_rss)

    if args.verbose:
        print('Total RSS for process {0} and children: {1}'.format(
            args.rootpid, formatted_rss
        ))
    else:
        print(formatted_rss)


if __name__ == '__main__':
    main()