123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240 |
- #!/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()
- #
|