RAM usage script

ram_usage.py

#!/usr/bin/env python
# -*- coding: utf-8 vi:noet
# Measure RAM usage by a process (and its children) on Linux

import sys, io, os, logging, subprocess, time

def ram_usage(pid, depth=0):
	pss = 0

	logging.debug("%s Checking %d", " " * depth, pid)

	with io.open("/proc/{pid}/smaps".format(**locals()), "rb") as f:
		for line in f:
			line = line.rstrip().decode()
			if line.startswith("Pss:"):
				h, q, u = line.split()
				mul = {
				 "kB": (1<<10),
				}
				q = int(q) * mul[u]
				pss += q

	logging.debug("%s Self usage %d B", " " * depth, pss)

	with io.open("/proc/{pid}/task/{pid}/children".format(**locals()), "rb") as f:
		pids = [ int(x) for x in f.read().rstrip().decode().split() ]
		for spid in pids:
			pss += ram_usage(spid, depth=depth+1)

	logging.debug("%s Total usage %d B", " " * depth, pss)

	return pss


if __name__ == "__main__":
	import argparse

	logger = logging.getLogger()

	parser = argparse.ArgumentParser(
	 description="RAM usage",
	)

	parser.add_argument("--log-level",
	 default="INFO",
	 help="Logging level (eg. INFO, see Python logging docs)",
	)

	parser.add_argument("--period",
	 type=float,
	 default=1.0,
	)

	grp = parser.add_mutually_exclusive_group(required=True)

	grp.add_argument("-p", "--pid",
	 type=int,
	)

	grp.add_argument("-e", "--execute",
	 nargs=argparse.REMAINDER,
	)

	try:
		import argcomplete
		argcomplete.autocomplete(parser)
	except:
		pass

	args = parser.parse_args()

	logger.setLevel(getattr(logging, args.log_level))

	if args.pid:
		while True:
			pss = ram_usage(args.pid)
			print(pss)
			time.sleep(args.period)

	else:
		cmd = args.execute

		p = subprocess.Popen(cmd)

		while p.poll() is None:
			pss = ram_usage(p.pid)
			print(pss)
			time.sleep(args.period)

		p.wait()