Implement fuzzer weight setting. (#4392)

This allows manipulating the weight of individual FuzzerJob entries in
the database.

Complements the `fuzz-target set` command.

Author: letitz

butler.py

@@ -132,6 +132,25 @@ def _add_weights_fuzzer_subparser(weights_subparsers):
   aggregate_parser.add_argument(
       '-j', '--jobs', help='Which jobs to aggregate.', nargs='+')
 
+  set_parser = subparsers.add_parser(
+      'set', help='Set the weight of a FuzzerJob entry.')
+  set_parser.add_argument(
+      '-f',
+      '--fuzzer',
+      help='The fuzzer field of the entry to modify.',
+      required=True)
+  set_parser.add_argument(
+      '-j',
+      '--job',
+      help='The job field of the entry to modify.',
+      required=True)
+  set_parser.add_argument(
+      '-w',
+      '--weight',
+      help='The new weight to set.',
+      type=float,
+      required=True)
+
 
 def _add_weights_batches_subparser(weights_subparsers):
   """Adds a parser for the `weights fuzzer-batch` command."""

src/local/butler/weights.py

@@ -39,8 +39,8 @@
 
 def _iter_weights(
     fuzzer_jobs: Sequence[data_types.FuzzerJob]) -> Sequence[float]:
-  for fj in fuzzer_jobs:
-    yield fj.actual_weight
+  for fuzzer_job in fuzzer_jobs:
+    yield fuzzer_job.actual_weight
 
 
 def _sum_weights(fuzzer_jobs: Sequence[data_types.FuzzerJob]) -> float:
@@ -123,7 +123,8 @@ def _display_fuzzer_jobs(fuzzer_jobs: Sequence[data_types.FuzzerJob],
   printer = _print_with_prefix(prefix)
 
   fuzzer_jobs = list(fuzzer_jobs)
-  fuzzer_jobs.sort(key=lambda fj: fj.actual_weight, reverse=True)
+  fuzzer_jobs.sort(
+      key=lambda fuzzer_job: fuzzer_job.actual_weight, reverse=True)
 
   total_weight = _sum_weights(fuzzer_jobs)
 
@@ -316,6 +317,45 @@ def _aggregate_fuzzer_jobs(
   _print_stats(others, total_weight)
 
 
+def _set_fuzzer_job_weight(
+    fuzzer: str,
+    job: str,
+    weight: float,
+) -> None:
+  """Sets the matching FuzzerJob's weight to the given value."""
+  fuzzer_jobs = list(
+      data_types.FuzzerJob.query(data_types.FuzzerJob.fuzzer == fuzzer,
+                                 data_types.FuzzerJob.job == job))
+
+  if not fuzzer_jobs:
+    print('No matching FuzzerJob entries found for ' +
+          f'fuzzer {fuzzer} and job {job}')
+    return
+
+  if len(fuzzer_jobs) > 1:
+    print('Bailing out! Multiple FuzzerJob entries found for ' +
+          f'fuzzer {fuzzer} and job {job}: {fuzzer_jobs}')
+    return
+
+  fuzzer_job = fuzzer_jobs[0]
+
+  print(f'Fuzzer: {fuzzer_job.fuzzer}')
+  print(f'Job: {fuzzer_job.job}')
+  print(f'Platform: {fuzzer_job.platform}')
+  print(f'Multiplier: {fuzzer_job.multiplier}')
+  print(f'Old weight: {fuzzer_job.weight}')
+  print(f'-> New weight: {weight}')
+
+  answer = input('Do you want to apply this mutation? [y,n] ')
+  if answer.lower() != 'y':
+    print('Not applying mutation.')
+    return
+
+  fuzzer_job.weight = weight
+  fuzzer_job.put()
+  print('Mutation applied.')
+
+
 def _set_fuzz_target_job_weight(
     fuzz_target_name: str,
     job: str,
@@ -362,6 +402,8 @@ def _execute_fuzzer_command(args) -> None:
       raise TypeError(f'--format {repr(args.format)} unrecognized')
   elif cmd == 'aggregate':
     _aggregate_fuzzer_jobs(args.platform, fuzzers=args.fuzzers, jobs=args.jobs)
+  elif cmd == 'set':
+    _set_fuzzer_job_weight(args.fuzzer, args.job, args.weight)
   else:
     raise TypeError(f'weights fuzzer command {repr(cmd)} unrecognized')