relay-throughput-testing/src/log_system_usage.py

#!/usr/bin/python3
#
import time
import threading
import subprocess
import re
import sys
import os
import pickle
import gzip
#
PROC_STAT_HEADERS = ('user', 'nice', 'system', 'idle', 'iowait', 'irq', 'softirq', 'steal', 'guest', 'guest_nice')
PROC_PID_STAT_HEADERS = ('pid', 'comm', 'state', 'ppid', 'pgrp', 'session', 'tty_nr', 'tpgid', 'flags', 'minflt',
                         'cminflt', 'majflt', 'cmajflt', 'utime', 'stime', 'cutime', 'cstime', 'priority', 'nice',
                         'num_threads', 'itrealvalue', 'starttime', 'vsize', 'rss', 'rsslim', 'startcode',
                         'endcode', 'startstack', 'kstkesp', 'kstkeip', 'signal', 'blocked', 'sigignore',
                         'sigcatch', 'wchan', 'nswap', 'cnswap', 'exit_signal', 'processor', 'rt_priority',
                         'policy', 'delayacct_blkio_ticks', 'guest_time', 'cguest_time', 'start_data', 'end_data',
                         'start_brk', 'arg_start', 'arg_end', 'env_start', 'env_end', 'exit_code')
#
def get_cpu_stats(path='/proc/stat'):
	"""
	Get CPU statistics from /proc/stat. Output is returned for the system and each CPU.

	Ex:

	{'system': {'user': 8696397, 'nice': 22431, ...},
	 'cpus': {
	     0: {'user': 4199206, 'nice': 11615, ...},
	     1: {'user': 4199308, 'nice': 10642, ...}
	 }
	}
	"""
	#
	with open(path, 'r') as f:
		lines = f.readlines()
	#
	cpu_lines = [l for l in lines if l.startswith('cpu')]
	stats = {'system':None, 'cpus':{}}
	#
	for l in cpu_lines:
		l_split = l.split()
		cpu_index = l_split[0][3:]
		cpu_stats = {x[0]: int(x[1]) for x in zip(PROC_STAT_HEADERS, l_split[1:])}
		#
		if cpu_index == '':
			stats['system'] = cpu_stats
		else:
			stats['cpus'][int(cpu_index)] = cpu_stats
		#
	#
	return stats
#
def parse_stat_file(path):
	with open(path, 'r') as f:
		contents = f.read()
	#
	raw_stats = re.findall("(\(.*\)|\S+)", contents, flags=re.DOTALL)
	proc_stats = {x[0]: x[1] for x in zip(PROC_PID_STAT_HEADERS, raw_stats)}
	for k in proc_stats:
		if k != 'comm' and k != 'state':
			proc_stats[k] = int(proc_stats[k])
		#
	#
	return proc_stats
#
def get_proc_stats(pid):
	pid = int(pid)
	path = os.path.join('/proc', str(pid), 'stat')
	#
	return parse_stat_file(path)
#
def get_thread_stats(tid):
	tid = int(tid)
	path = os.path.join('/proc', str(tid), 'task', str(tid), 'stat')
	#
	return parse_stat_file(path)
#
def calculate_cpu_stats(stats):
	idle = stats['idle'] + stats['iowait']
	non_idle = stats['user'] + stats['nice'] + stats['system'] + stats['irq'] + stats['softirq'] + stats['steal']
	#
	return (idle, non_idle)
#
def calculate_cpu_usage(initial, current):
	"""
	Calculation adapted from: https://stackoverflow.com/questions/23367857/accurate-calculation-of-cpu-usage-given-in-percentage-in-linux/
	"""
	#
	(initial_idle, initial_non_idle) = calculate_cpu_stats(initial)
	initial_total = initial_idle + initial_non_idle
	#
	(current_idle, current_non_idle) = calculate_cpu_stats(current)
	current_total = current_idle + current_non_idle
	#
	return (current_non_idle-initial_non_idle)/(current_total-initial_total)
#
def calculate_process_cpu_usage(process_initial, process_current, cpu_initial, cpu_current):
	(initial_idle, initial_non_idle) = calculate_cpu_stats(cpu_initial)
	initial_total = initial_idle + initial_non_idle
	#
	(current_idle, current_non_idle) = calculate_cpu_stats(cpu_current)
	current_total = current_idle + current_non_idle
	#
	process_initial_non_idle = process_initial['utime'] + process_initial['stime']
	process_current_non_idle = process_current['utime'] + process_current['stime']
	#
	return (process_current_non_idle-process_initial_non_idle)/(current_total-initial_total)
#
def calculate_cpu_usage_continuous(stats):
	cpu_usages = []
	for i in range(len(stats)-1):
		cpu_usages.append(calculate_cpu_usage(stats[i], stats[i+1]))
	#
	return cpu_usages
#
def calculate_process_cpu_usage_continuous(process_stats, cpu_stats):
	process_usages = []
	assert all([len(process_stats) == len(cpu_stats[i]) for i in cpu_stats])
	for i in range(len(process_stats)-1):
		using_core_0 = process_stats[i]['processor']
		using_core_1 = process_stats[i+1]['processor']
		usage_0 = calculate_process_cpu_usage(process_stats[i], process_stats[i+1], cpu_stats[using_core_0][i], cpu_stats[using_core_0][i+1])
		usage_1 = calculate_process_cpu_usage(process_stats[i], process_stats[i+1], cpu_stats[using_core_1][i], cpu_stats[using_core_1][i+1])
		process_usages.append((usage_0+usage_1)/2)
	#
	return process_usages
#
#def get_running_processes():
#	lines = subprocess.check_output(['ps', '-a', '-x', '-o', 'pid,state,args', '--no-headers']).decode('utf-8').split('\n')
#	lines = [line.strip() for line in lines]
#	lines = [line.split(' ', 2) for line in lines if len(line) != 0]
#	#
#	data = []
#	for line in lines:
#		data.append({'pid':int(line[0]), 'state':line[1], 'args':line[2]})
#	#
#	return data
#
def log_cpu_stats(path, interval, pids, stop_event):
	"""
	Log the cpu stats to a gz compressed JSON file. Storing JSON
	seems to only use about 10% more disk space than storing
	bytes directly (4 bytes per value), so JSON is used for
	simplicity.

	path: file to save to
	interval: how many seconds to wait before getting more data
	stop_event: a threading.Event which stops the function
	"""
	#
	pids = [int(pid) for pid in pids]
	threads = {pid: [int(tid) for tid in os.listdir('/proc/{}/task'.format(pid))] for pid in pids}
	stats = {'timestamps':[],
	         'cpu':{'system':[],
	                'id':{x: [] for x in get_cpu_stats()['cpus'].keys()}},
	         'process':{x: {'pid': [],
	                        'tid': {y: [] for y in threads[x]}} for x in pids}}
	#
	while not stop_event.is_set():
		current_time = time.time()
		stats['timestamps'].append(current_time)
		#
		current_stats = get_cpu_stats()
		stats['cpu']['system'].append(current_stats['system'])
		for cpu in current_stats['cpus']:
			stats['cpu']['id'][cpu].append(current_stats['cpus'][cpu])
		#
		for pid in pids:
			stats['process'][pid]['pid'].append(get_proc_stats(pid))
			for tid in threads[pid]:
				stats['process'][pid]['tid'][tid].append(get_thread_stats(tid))
			#
		#
		wait_time = max(0, interval-(time.time()-current_time))
		stop_event.wait(wait_time)
	#
	with gzip.GzipFile(path, 'wb') as f:
		pickle.dump(stats, f, protocol=4)
	#
#
def load_cpu_stats(path):
	with gzip.GzipFile(path, 'rb') as f:
		return pickle.load(f)
	#
#
'''
def log_cpu_stats(path, interval, stop_event):
	path: file to save to
	interval: how many seconds to wait before getting more data
	stop_event: a threading.Event which stops the function
	#
	with gzip.GzipFile(path+'.2.gz', 'w') as f:
		f.write(' '.join(PROC_STAT_HEADERS).encode('utf-8'))
		f.write('\n\n'.encode('utf-8'))
		#
		while not stop_event.is_set():
			f.write(str(time.time()).encode('utf-8'))
			f.write('\n'.encode('utf-8'))
			stats = get_cpu_stats()
			f.write('cpu '.encode('utf-8'))
			#f.write(' '.join([str(stats['system'][x]) for x in PROC_STAT_HEADERS]).encode('utf-8'))
			f.write(b''.join([stats['system'][x].to_bytes(4, byteorder='big') for x in PROC_STAT_HEADERS]))
			f.write('\n'.encode('utf-8'))
			for cpu in stats['cpus']:
				f.write('cpu{} '.format(cpu).encode('utf-8'))
				#f.write(' '.join([str(stats['cpus'][cpu][x]) for x in PROC_STAT_HEADERS]).encode('utf-8'))
				f.write(b''.join([stats['cpus'][cpu][x].to_bytes(4, byteorder='big') for x in PROC_STAT_HEADERS]))
				f.write('\n'.encode('utf-8'))
			#
			f.write('\n'.encode('utf-8'))
			time.sleep(interval)
		#
	#
#
'''
if __name__ == '__main__':
	stop_event = threading.Event()
	#
	assert len(sys.argv) >= 3
	interval = float(sys.argv[1])
	file_name = sys.argv[2]
	if len(sys.argv) > 3:
		pids = sys.argv[3].split(',')
	else:
		pids = []
	#
	t = threading.Thread(target=log_cpu_stats, args=(file_name, interval, pids, stop_event))
	t.start()
	#
	try:
		while True:
			time.sleep(100)
		#
	except KeyboardInterrupt:
		stop_event.set()
		print()
	#
	t.join()
#