find_ngrams_on_ct 1.8 KB

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