1 #!/usr/bin/env python 2 3 4 """Run Cluster Telemetry to compute n-grams from SKPs.""" 5 6 7 import argparse 8 import os 9 import re 10 import subprocess 11 import sys 12 import tempfile 13 14 15 NGRAMS_LUA_SUBSTITUTION_STR = '-- CHANGEME\nlocal n = \d+\n-- CHANGEME' 16 SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__)) 17 18 19 def main(): 20 # Parse arguments. 21 parser = argparse.ArgumentParser( 22 description='Run Cluster Telemetry to compute n-grams from SKPs.') 23 parser.add_argument('--n', help='Compute n-grams with this integer as N', 24 required=True) 25 args, extra_args = parser.parse_known_args() 26 27 # Read the n-gram Lua script. 28 script_path = os.path.join(SCRIPT_DIR, 'ngrams.lua') 29 with open(script_path) as f: 30 script_contents = f.read() 31 32 # Replace the default value of n with the value specified by the user. 33 new_contents, subd = re.subn(NGRAMS_LUA_SUBSTITUTION_STR, 34 'local n = %s' % args.n, script_contents, 1) 35 if subd != 1: 36 raise Exception('Unable to replace N in %s; expected to find:\n%s' % ( 37 script_path, sub)) 38 39 # Write the new script contents to a temporary file. 40 tmp_script = tempfile.NamedTemporaryFile(delete=False) 41 try: 42 tmp_script.write(new_contents) 43 tmp_script.close() 44 45 # Run trigger_ct_lua with the new script, forwarding the rest of the 46 # passed-in arguments to the script. 47 trigger_ct_lua = os.path.join(SCRIPT_DIR, 'trigger_ct_lua') 48 cmd = ['python', trigger_ct_lua, 49 '--script', tmp_script.name, 50 '--aggregator', os.path.join(SCRIPT_DIR, 'ngrams_aggregate.lua')] 51 cmd.extend(extra_args) 52 try: 53 subprocess.check_call(cmd) 54 except subprocess.CalledProcessError: 55 exit(1) 56 finally: 57 os.remove(tmp_script.name) 58 59 60 if __name__ == '__main__': 61 main() 62
