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