1#!/usr/bin/env python 2# 3# 4# Copyright 2016, NICTA 5# 6# This software may be distributed and modified according to the terms of 7# the BSD 2-Clause license. Note that NO WARRANTY is provided. 8# See "LICENSE_BSD2.txt" for details. 9# 10# @TAG(NICTA_BSD) 11# 12 13''' 14Monitors the total CPU usage of a process and its children. Usage is similar 15to the UNIX `time` utility. 16 17NB: In order to get up-to-date information, we don't use interfaces such 18 as os.times, which only includes terminated and waited-for children. 19 Instead, we poll the process tree regularly. This means that when a 20 child process dies, its CPU time since the last poll is lost and not 21 included in the total time. 22 23 Hence the total reported time will be an underestimate of the true 24 CPU usage, especially for short-lived child processes. 25''' 26 27from __future__ import print_function 28import os 29import psutil 30import signal 31import subprocess 32import sys 33import threading 34import time 35import warnings 36 37try: 38 import psutil 39 if not hasattr(psutil.Process, "children") and hasattr(psutil.Process, "get_children"): 40 psutil.Process.children = psutil.Process.get_children 41 if not hasattr(psutil.Process, "memory_maps") and hasattr(psutil.Process, "get_memory_maps"): 42 psutil.Process.memory_maps = psutil.Process.get_memory_maps 43 44except ImportError: 45 print("Error: 'psutil' module not available. Run\n" 46 "\n" 47 " pip install --user psutil\n" 48 "\n" 49 "to install.", file=sys.stderr) 50 sys.exit(1) 51 52# The psutil.Process.cpu_times() API changed at version 4.1.0. 53# Earlier versions give the user and system times for the queried 54# process only. Later versions return two additional tuple members 55# for the total user and system times of its child processes. 56# For compatibility with both versions, we ignore the additional 57# values. 58def cpu_time_of(process): 59 cpu_times = process.cpu_times() 60 return cpu_times[0] + cpu_times[1] 61 62class Poller(threading.Thread): 63 '''Subclass of threading.Thread that monitors CPU usage of another process. 64 Use run() to start the process. 65 Use cpu_usage() to retrieve the latest estimate of CPU usage.''' 66 def __init__(self, pid): 67 super(Poller, self).__init__() 68 # Daemonise ourselves to avoid delaying exit of the process of our 69 # calling thread. 70 self.daemon = True 71 self.pid = pid 72 self.finished = False 73 self.started = threading.Semaphore(0) 74 self.proc = None 75 76 # Reported stat. 77 self.cpu = 0.0 78 79 # Remember CPU times of recently seen children. 80 # This is to prevent double-counting for child processes. 81 self.current_children = {} # {(pid, create_time): CPU time} 82 # CPU time of dead children is recorded here. 83 self.old_children_cpu = 0.0 84 85 def run(self): 86 def update(): 87 total = 0.0 88 89 # Fetch process's usage. 90 try: 91 if self.proc is None: 92 self.proc = psutil.Process(self.pid) 93 total += cpu_time_of(self.proc) 94 95 # Fetch children's usage. 96 new_current_children = {} 97 for c in self.proc.children(recursive=True): 98 try: 99 t = cpu_time_of(c) 100 new_current_children[(c.pid, c.create_time())] = t 101 total += t 102 except psutil.NoSuchProcess: 103 pass 104 except psutil.AccessDenied: 105 pass 106 107 # For children that are no longer running, remember their 108 # most recently recorded CPU time. 109 reaped_cpu = 0.0 110 for c_id, c_t in self.current_children.items(): 111 if c_id not in new_current_children: 112 reaped_cpu += c_t 113 self.old_children_cpu += reaped_cpu 114 self.current_children = new_current_children 115 total += self.old_children_cpu 116 117 except psutil.AccessDenied as err: 118 warnings.warn("access denied: pid=%d" % err.pid, RuntimeWarning) 119 120 # Add 1 ns allowance for floating-point rounding, which occurs when we 121 # accumulate current_children times for dead processes into reaped_cpu. 122 # (Floating point epsilon is about 1e-15.) 123 if total + 1e-9 < self.cpu: 124 try: 125 cmd = repr(' '.join(self.proc.cmdline())) 126 except Exception: 127 cmd = '??' 128 warnings.warn("cpu non-monotonic: %.15f -> %.15f, pid=%d, cmd=%s" % 129 (self.cpu, total, self.pid, cmd), 130 RuntimeWarning) 131 return total 132 133 # Fetch a sample, and notify others that we have started. 134 self.cpu = update() 135 self.started.release() 136 137 # Poll the process periodically. 138 # 139 # We poll quickly at the beginning and use exponential backout 140 # to try and get better stats on short-lived processes. 141 # 142 polling_interval = 0.01 143 max_interval = 0.5 144 while not self.finished: 145 time.sleep(polling_interval) 146 try: 147 self.cpu = update() 148 except psutil.NoSuchProcess: 149 break 150 if polling_interval < max_interval: 151 polling_interval = min(polling_interval * 1.5, max_interval) 152 153 def cpu_usage(self): 154 return self.cpu 155 156 def __enter__(self): 157 return self 158 159 def __exit__(self, *_): 160 self.finished = True 161 162def process_poller(pid): 163 '''Initiate polling of a subprocess. This is intended to be used in a 164 `with` block.''' 165 # Create a new thread and start it up. 166 p = Poller(pid) 167 p.start() 168 169 # Wait for the thread to record at least one sample before continuing. 170 p.started.acquire() 171 172 return p 173 174def main(): 175 if len(sys.argv) <= 1 or sys.argv[1] in ['-?', '--help']: 176 print('Usage: %s command args...\n Measure total CPU ' \ 177 'usage of a command' % sys.argv[0], file=sys.stderr) 178 return -1 179 180 # Run the command requested. 181 try: 182 p = subprocess.Popen(sys.argv[1:]) 183 except OSError: 184 print('command not found', file=sys.stderr) 185 return -1 186 187 cpu = 0 188 m = process_poller(p.pid) 189 while True: 190 try: 191 p.returncode = p.wait() 192 break 193 except KeyboardInterrupt: 194 # The user Ctrl-C-ed us. The child should have received SIGINT; 195 # continue waiting for it to finish 196 pass 197 198 print('Total cpu %f seconds' % m.cpu_usage(), file=sys.stderr) 199 200 return p.returncode 201 202if __name__ == '__main__': 203 sys.exit(main()) 204