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