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