1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
|
#!/usr/bin/env python3
# utils/rusage.py - Utility to measure resource usage -*- python -*-
#
# This source file is part of the Swift.org open source project
#
# Copyright (c) 2017 Apple Inc. and the Swift project authors
# Licensed under Apache License v2.0 with Runtime Library Exception
#
# See https://swift.org/LICENSE.txt for license information
# See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
#
#
# Simple informative/supervisory wrapper around getrusage() and, optionally,
# setrlimit()
#
# By default this script lets its subprocess run to completion whatever its
# consumed limits, and only prints a limit-violation if it occurs after the
# fact. If you want to be a bit more abrupt you can run with --enforce and
# it will also run subprocesses under setrlimit().
#
# Note: setrlimit(RLIMIT_RSS) does nothing on macOS. It's unenforced.
#
# If there's a limit-violation observed (by getrusage()), it's printed to
# stderr and the wrapper exits with an error. Whether or not the violation
# was caused by enforcement or mere observation.
#
# If there's no limit-violation, the subprocess' own exit code is
# propagated, either silently or, if --verbose is passed, after printing
# the actual getrusage() memory and time values (such that they can be used
# as a limit in future runs).
#
import argparse
import csv
import datetime
import resource
import subprocess
import sys
import time
class MemAction(argparse.Action):
def __init__(self, *args, **kwargs):
super(MemAction, self).__init__(*args, **kwargs)
def __call__(self, parser, namespace, v, option_string=None):
r = None
if v.endswith('K'):
r = int(v[:-1]) * 1024
elif v.endswith('M'):
r = int(v[:-1]) * 1024 * 1024
elif v.endswith('G'):
r = int(v[:-1]) * 1024 * 1024 * 1024
else:
r = int(v)
setattr(namespace, self.dest, r)
class TimeAction(argparse.Action):
def __init__(self, *args, **kwargs):
super(TimeAction, self).__init__(*args, **kwargs)
def __call__(self, parser, namespace, v, option_string=None):
r = None
if v.endswith('ms'):
r = float(v[:-2]) / 1000.0
elif v.endswith('us'):
r = float(v[:-2]) / 1000000.0
else:
r = float(v)
setattr(namespace, self.dest, r)
parser = argparse.ArgumentParser()
parser.add_argument("--mem",
metavar="M",
help="memory (in bytes, or ..'K', 'M', 'G')",
action=MemAction)
parser.add_argument("--time",
metavar="T",
help="time (in secs, or ..'ms', 'us')",
action=TimeAction)
parser.add_argument("--wall-time",
help="wall time (in secs, or ..'ms', 'us')",
action='store_true')
parser.add_argument("--enforce",
action='store_true',
default=False,
help="call setrlimit() before running subprocess")
parser.add_argument("--verbose",
action='store_true',
default=False,
help="always report status and usage")
parser.add_argument("--csv",
action='store_true',
default=False,
help="write results as CSV")
parser.add_argument("--csv-header",
action='store_true',
default=False,
help="Emit CSV header")
parser.add_argument("--csv-output", default="-",
type=argparse.FileType('wb', 0),
help="Write CSV output to file")
parser.add_argument("--csv-name", type=str,
default=str(datetime.datetime.now()),
help="Label row in CSV with name")
parser.add_argument('remainder', nargs=argparse.REMAINDER,
help="subcommand to run under supervision")
args = parser.parse_args()
if len(args.remainder) == 0:
parser.print_help()
sys.exit(1)
if args.enforce:
if args.time is not None:
secs = max(1, int(args.time))
(soft, hard) = resource.getrlimit(resource.RLIMIT_CPU)
secs = min(soft, hard, secs)
if args.verbose:
sys.stderr.write("rusage: setrlimit(RLIMIT_CPU, %d)\n"
% secs)
resource.setrlimit(resource.RLIMIT_CPU, (secs, secs))
if args.mem is not None:
(soft, hard) = resource.getrlimit(resource.RLIMIT_RSS)
mem = min(soft, hard, args.mem)
if args.verbose:
sys.stderr.write("rusage: setrlimit(RLIMIT_RSS, %d)\n"
% mem)
resource.setrlimit(resource.RLIMIT_RSS, (mem, mem))
start = time.time()
ret = subprocess.call(args.remainder)
end = time.time()
wall_time = end - start
used = resource.getrusage(resource.RUSAGE_CHILDREN)
if args.verbose:
sys.stderr.write("rusage: subprocess exited %d\n" % ret)
over_mem = args.mem is not None and used.ru_maxrss > args.mem
over_time = args.time is not None and used.ru_utime > args.time
if args.verbose or over_mem:
sys.stderr.write("rusage: subprocess mem: %d bytes\n"
% used.ru_maxrss)
if over_mem:
sys.stderr.write("rusage: exceeded limit: %d bytes\n"
% args.mem)
if args.verbose or over_time:
sys.stderr.write("rusage: subprocess time: %.6f secs\n"
% used.ru_utime)
if args.wall_time:
sys.stderr.write("rusage: subprocess wall time: %.6f secs\n"
% wall_time)
if over_time:
sys.stderr.write("rusage: exceeded limit: %.6f secs\n"
% args.time)
if args.csv:
fieldnames = ["time", "mem", "run"]
row = {
'time': used.ru_utime,
'mem': used.ru_maxrss,
'run': args.csv_name
}
if args.wall_time:
row['wall'] = wall_time
fieldnames.insert(1, 'wall')
out = csv.DictWriter(args.csv_output, fieldnames, dialect='excel-tab')
if args.csv_header:
out.writeheader()
out.writerow(row)
if over_mem or over_time:
sys.exit(-1)
sys.exit(ret)
|
