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
|
#!/usr/bin/env python
"""Run Cluster Telemetry to compute n-grams from SKPs."""
import argparse
import os
import re
import subprocess
import sys
import tempfile
NGRAMS_LUA_SUBSTITUTION_STR = '-- CHANGEME\nlocal n = \d+\n-- CHANGEME'
SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))
def main():
# Parse arguments.
parser = argparse.ArgumentParser(
description='Run Cluster Telemetry to compute n-grams from SKPs.')
parser.add_argument('--n', help='Compute n-grams with this integer as N',
required=True)
args, extra_args = parser.parse_known_args()
# Read the n-gram Lua script.
script_path = os.path.join(SCRIPT_DIR, 'ngrams.lua')
with open(script_path) as f:
script_contents = f.read()
# Replace the default value of n with the value specified by the user.
new_contents, subd = re.subn(NGRAMS_LUA_SUBSTITUTION_STR,
'local n = %s' % args.n, script_contents, 1)
if subd != 1:
raise Exception('Unable to replace N in %s; expected to find:\n%s' % (
script_path, sub))
# Write the new script contents to a temporary file.
tmp_script = tempfile.NamedTemporaryFile(delete=False)
try:
tmp_script.write(new_contents)
tmp_script.close()
# Run trigger_ct_lua with the new script, forwarding the rest of the
# passed-in arguments to the script.
trigger_ct_lua = os.path.join(SCRIPT_DIR, 'trigger_ct_lua')
cmd = ['python', trigger_ct_lua,
'--script', tmp_script.name,
'--aggregator', os.path.join(SCRIPT_DIR, 'ngrams_aggregate.lua')]
cmd.extend(extra_args)
try:
subprocess.check_call(cmd)
except subprocess.CalledProcessError:
exit(1)
finally:
os.remove(tmp_script.name)
if __name__ == '__main__':
main()
|
