Refactor and test jobsim_client.

dm/tools/jobsim_client/generate_ensure_graph_data_req.py

+#!/usr/bin/env python
+# Copyright 2016 The LUCI Authors. All rights reserved.
+# Use of this source code is governed under the Apache License, Version 2.0
+# that can be found in the LICENSE file.
+
+"""
+This is a dumb script that will:
+  * compile jobsim_client (for linux-amd64)
+  * upload a simple cipd package containing only jobsim_client with the 'latest'
+    ref. The packaeg name is 'infra/experimental/dm/jobsim_client/linux-amd64'.
+  * Print a JSONPB-encoded EnsureGraphDataReq that runs jobsim_client with the
+    provided strings to calculate edit-distance(a, b, transposition?).
+"""
+
+import argparse
+import json
+import os
+import pprint
+import shutil
+import subprocess
+import tempfile
+
+THIS_DIR = os.path.dirname(os.path.abspath(__file__))
+
+def compile_pkg(pkg_dir, os_name, arch_name):
+  print 'building jobsim_client'
+  env = os.environ.copy()
+  env.update(
+    GOOS   = os_name,
+    GOARCH = arch_name,
+  )
+  subprocess.check_call(
+      ['go', 'build', 'github.com/luci/luci-go/dm/tools/jobsim_client'],
+      cwd=pkg_dir, env=env)
+
+def upload_pkg(pkg_dir, pkg_name_prefix, os_name, arch_name):
+  print 'creating jobsim_client package'
+
+  pkg_name = '%s/%s-%s' % (pkg_name_prefix, os_name, arch_name)
+
+  fd, tfile = tempfile.mkstemp('-cipd-output.json')
+  os.close(fd)
+  try:
+    subprocess.check_call(['cipd', 'create', '-name', pkg_name,
+                           '-ref', 'latest', '-in', pkg_dir,
+                           '-install-mode', 'copy', '-json-output', tfile])
+    with open(tfile, 'r') as tfileData:
+      out_json = json.load(tfileData)
+    version = out_json[u'result'][u'instance_id'].encode('utf-8')
+  finally:
+    os.unlink(tfile)
+
+  print 'uploaded %s:%s' % (pkg_name, version)
+
+  return pkg_name, version
+
+def print_req(opts, pkg_name, version):
+  def dumps(obj):
+    return json.dumps(obj, sort_keys=True, separators=(',', ':'))
+
+  cpu = {
+    'amd64': 'x86-64',
+  }[opts.arch]
+
+  os_name = {
+    'linux': 'Linux',
+  }[opts.os]
+
+  command = ['jobsim_client', 'edit-distance', '-dm-host', '${DM.HOST}',
+             '-execution-auth-path', '${DM.EXECUTION.AUTH:PATH}',
+             '-quest-desc-path', '${DM.QUEST.DATA.DESC:PATH}']
+  if opts.use_transposition:
+    command.append('-use-transposition')
+
+  distParams = {
+    'scheduling': {
+      'dimensions': {
+        'cpu': cpu,
+        'os': os_name,
+        'pool': opts.pool,
+      },
+    },
+    'meta': {'name_prefix': 'dm jobsim client'},
+    'job': {
+      'inputs': {
+        'cipd': {
+          'server': 'https://chrome-infra-packages.appspot.com',
+          'by_path': {
+            '.': {
+              'pkg': [
+                {
+                  'name': pkg_name,
+                  'version': version if opts.pin else 'latest',
+                },
+              ]
+            }
+          }
+        }
+      },
+      'command': command,
+    }
+  }
+
+  params = {
+    'a': opts.a,
+    'b': opts.b,
+  }
+
+  desc = {
+    'quest': [
+      {
+        'distributor_config_name': 'swarming',
+        'parameters': dumps(params),
+        'distributor_parameters': dumps(distParams),
+        'meta': {
+          'timeouts': {
+            'start': '600s',
+            'run': '300s',
+            'stop': '300s',
+          }
+        },
+      }
+    ],
+    'quest_attempt': [
+      {'nums': [1]},
+    ]
+  }
+
+  print dumps(desc)
+
+def main():
+  parser = argparse.ArgumentParser(
+      description=__doc__, formatter_class=argparse.RawTextHelpFormatter)
+  parser.add_argument('--use-transposition', action='store_true', default=False,
+                      help=('Use Damerau-Levenshtein distance calculation '
+                            'instead of plain Levenshtein distance.'))
+  parser.add_argument('a', type=str, help='The "a" string to calculate for.')
+  parser.add_argument('b', type=str, help='The "b" string to calculate for.')
+
+  plat_grp = parser.add_argument_group(
+      'platform', 'Options for the target platform of the job.')
+  plat_grp.add_argument('--os', choices=('linux',), default='linux',
+                        help='The OS to compile/run on.')
+  plat_grp.add_argument('--arch', choices=('amd64',), default='amd64',
+                        help='The Arch to compile/run on.')
+  plat_grp.add_argument('--pool', type=str, default='default',
+                        help='The swarming pool to use.')
+  plat_grp.add_argument('--pin', action='store_true', default=False,
+                        help='Emit the request with a pinned package version'
+                        ' instead of "latest".')
+
+  cipd_grp = parser.add_argument_group('cipd', 'cipd packaging options')
+  cipd_grp.add_argument('--cipd-service-url', default=None,
+                        help='The CIPD service to upload to.')
+  cipd_grp.add_argument('--cipd-service-account-json', default=None,
+                        help='The CIPD service account JSON file to use.')
+  cipd_grp.add_argument('--cipd-name',
+                        default='infra/experimental/dm/jobsim_client',
+                        help='The CIPD package name prefix to upload to. This '
+                        'will be appended with the standard os-arch suffix.')
+
+  opts = parser.parse_args()
+
+  # Use local path for determinisim.
+  pkg_dir = os.path.join(THIS_DIR, 'pkg_dir')
+  shutil.rmtree(pkg_dir, ignore_errors=True)
+  os.mkdir(pkg_dir)
+  try:
+    compile_pkg(pkg_dir, opts.os, opts.arch)
+    pkg_name, version = upload_pkg(pkg_dir, opts.cipd_name, opts.os, opts.arch)
+    print_req(opts, pkg_name, version)
+  finally:
+    shutil.rmtree(pkg_dir, ignore_errors=True)
+
+
+if __name__ == '__main__':
+  main()