Switch generate_default_trace to worker infrastructure. (#222)

Author: mtrofin

compiler_opt/tools/generate_default_trace.py

@@ -14,26 +14,30 @@
 # limitations under the License.
 """Generate initial training data from the behavior of the current heuristic."""
 
+import concurrent.futures
 import contextlib
 import functools
-import os
-import queue
 import re
-import subprocess
 from typing import Dict, List, Optional, Union, Tuple  # pylint:disable=unused-import
 
 from absl import app
 from absl import flags
 from absl import logging
 import gin
-import multiprocessing
+
 import tensorflow as tf
 
+from compiler_opt.distributed import worker
+from compiler_opt.distributed import buffered_scheduler
+from compiler_opt.distributed.local import local_worker_manager
+
 from compiler_opt.rl import compilation_runner
 from compiler_opt.rl import corpus
 from compiler_opt.rl import policy_saver
 from compiler_opt.rl import registry
 
+from tf_agents.system import system_multiprocessing as multiprocessing
+
 # see https://bugs.python.org/issue33315 - we do need these types, but must
 # currently use them as string annotations
 
@@ -76,64 +80,45 @@ def get_runner() -> compilation_runner.CompilationRunner:
   return problem_config.get_runner_type()(moving_average_decay_rate=0)
 
 
-def worker(policy_path: Optional[str],
-           work_queue: 'queue.Queue[corpus.LoadedModuleSpec]',
-           results_queue: 'queue.Queue[ResultsQueueEntry]',
-           key_filter: Optional[str]):
-  """Describes the job each paralleled worker process does.
+class FilteringWorker(worker.Worker):
+  """Worker that performs a computation and optionally filters the result.
 
-  The worker picks a workitem from the work_queue, process it, and deposits
-  a result on the results_queue, in either success or failure cases.
-  The results_queue items are tuples (workitem, result). On failure, the result
-  is None.
 
   Args:
-    runner: the data collector.
     policy_path: the policy_path to generate trace with.
-    work_queue: the queue of unprocessed work items.
-    results_queue: the queue where results are deposited.
     key_filter: regex filter for key names to include, or None to include all.
   """
-  try:
-    runner = get_runner()
-    m = re.compile(key_filter) if key_filter else None
-    policy = policy_saver.Policy.from_filesystem(
+
+  def __init__(self, policy_path: Optional[str], key_filter: Optional[str]):
+    self._policy_path = policy_path
+    self._key_filter = re.compile(key_filter) if key_filter else None
+    self._runner = get_runner()
+    self._policy = policy_saver.Policy.from_filesystem(
         policy_path) if policy_path else None
-    while True:
-      try:
-        loaded_module_spec = work_queue.get_nowait()
-      except queue.Empty:
-        return
-      try:
-        data = runner.collect_data(
-            loaded_module_spec=loaded_module_spec,
-            policy=policy,
-            reward_stat=None,
-            model_id=0)
-        if not m:
-          results_queue.put(
-              (loaded_module_spec.name, data.serialized_sequence_examples,
-               data.reward_stats))
-          continue
-        new_reward_stats = {}
-        new_sequence_examples = []
-        for k, sequence_example in zip(data.keys,
-                                       data.serialized_sequence_examples):
-          if not m.match(k):
-            continue
-          new_reward_stats[k] = data.reward_stats[k]
-          new_sequence_examples.append(sequence_example)
-        results_queue.put(
-            (loaded_module_spec.name, new_sequence_examples, new_reward_stats))
-      except (subprocess.CalledProcessError, subprocess.TimeoutExpired,
-              RuntimeError):
-        logging.error('Failed to compile %s.', loaded_module_spec.name)
-        results_queue.put(None)
-  except BaseException as e:  # pylint: disable=broad-except
-    results_queue.put(e)
-
-
-def main(_):
+
+  def compile_and_filter(
+      self, loaded_module_spec: corpus.LoadedModuleSpec
+  ) -> Tuple[str, List[str], Dict[str, compilation_runner.RewardStat]]:
+    data = self._runner.collect_data(
+        loaded_module_spec=loaded_module_spec,
+        policy=self._policy,
+        reward_stat=None,
+        model_id=0)
+    if self._key_filter is None:
+      return (loaded_module_spec.name, data.serialized_sequence_examples,
+              data.reward_stats)
+    new_reward_stats = {}
+    new_sequence_examples = []
+    for k, sequence_example in zip(data.keys,
+                                   data.serialized_sequence_examples):
+      if not self._key_filter.match(k):
+        continue
+      new_reward_stats[k] = data.reward_stats[k]
+      new_sequence_examples.append(sequence_example)
+    return (loaded_module_spec.name, new_sequence_examples, new_reward_stats)
+
+
+def main(worker_manager_class=local_worker_manager.LocalWorkerPoolManager):
 
   gin.parse_config_files_and_bindings(
       _GIN_FILES.value, bindings=_GIN_BINDINGS.value, skip_unknown=False)
@@ -160,74 +145,60 @@ def main(_):
   # other smaller files are processed in parallel
   corpus_elements = cps.sample(k=sampled_modules, sort=True)
 
-  worker_count = (
-      min(os.cpu_count(), _NUM_WORKERS.value)
-      if _NUM_WORKERS.value else os.cpu_count())
-
   tfrecord_context = (
       tf.io.TFRecordWriter(_OUTPUT_PATH.value)
       if _OUTPUT_PATH.value else contextlib.nullcontext())
   performance_context = (
       tf.io.gfile.GFile(_OUTPUT_PERFORMANCE_PATH.value, 'w')
       if _OUTPUT_PERFORMANCE_PATH.value else contextlib.nullcontext())
+  work = [
+      cps.load_module_spec(corpus_element) for corpus_element in corpus_elements
+  ]
 
   with tfrecord_context as tfrecord_writer:
     with performance_context as performance_writer:
-      ctx = multiprocessing.get_context()
-      m = ctx.Manager()
-      results_queue: 'queue.Queue[ResultsQueueEntry]' = m.Queue()
-      work_queue: 'queue.Queue[corpus.LoadedModuleSpec]' = m.Queue()
-      for corpus_element in corpus_elements:
-        work_queue.put(cps.load_module_spec(corpus_element))
-
-      # pylint:disable=g-complex-comprehension
-      processes = [
-          ctx.Process(
-              target=functools.partial(worker, _POLICY_PATH.value, work_queue,
-                                       results_queue, _KEY_FILTER.value))
-          for _ in range(0, worker_count)
-      ]
-      # pylint:enable=g-complex-comprehension
-
-      for p in processes:
-        p.start()
-
-      total_successful_examples = 0
-      total_work = len(corpus_elements)
-      total_failed_examples = 0
-      total_training_examples = 0
-      for _ in range(total_work):
-        logging.log_every_n_seconds(logging.INFO,
-                                    '%d success, %d failed out of %d', 10,
-                                    total_successful_examples,
-                                    total_failed_examples, total_work)
-
-        results = results_queue.get()
-        if isinstance(results, BaseException):
-          logging.fatal(results)
-        if not results:
-          total_failed_examples += 1
-          continue
-
-        total_successful_examples += 1
-        module_name, records, reward_stat = results
-        if tfrecord_writer:
-          total_training_examples += len(records)
-          for r in records:
-            tfrecord_writer.write(r)
-        if performance_writer:
-          for key, value in reward_stat.items():
-            performance_writer.write(
-                (f'{module_name},{key},{value.default_reward},'
-                 f'{value.moving_average_reward}\n'))
-
-      print((f'{total_successful_examples} of {len(corpus_elements)} modules '
-             f'succeeded, and {total_training_examples} trainining examples '
-             'written'))
-      for p in processes:
-        p.join()
+      with worker_manager_class(
+          FilteringWorker,
+          _NUM_WORKERS.value,
+          policy_path=_POLICY_PATH.value,
+          key_filter=_KEY_FILTER.value) as lwm:
+
+        _, result_futures = buffered_scheduler.schedule_on_worker_pool(
+            action=lambda w, j: w.compile_and_filter(j),
+            jobs=work,
+            worker_pool=lwm)
+        total_successful_examples = 0
+        total_work = len(corpus_elements)
+        total_failed_examples = 0
+        total_training_examples = 0
+        not_done = result_futures
+        while not_done:
+          (done, not_done) = concurrent.futures.wait(not_done, 10)
+          succeeded = [
+              r for r in done if not r.cancelled() and r.exception() is None
+          ]
+          total_successful_examples += len(succeeded)
+          total_failed_examples += (len(done) - len(succeeded))
+          for r in succeeded:
+            module_name, records, reward_stat = r.result()
+            if tfrecord_writer:
+              total_training_examples += len(records)
+              for r in records:
+                tfrecord_writer.write(r)
+            if performance_writer:
+              for key, value in reward_stat.items():
+                performance_writer.write(
+                    (f'{module_name},{key},{value.default_reward},'
+                     f'{value.moving_average_reward}\n'))
+          logging.info('%d success, %d failed out of %d',
+                       total_successful_examples, total_failed_examples,
+                       total_work)
+
+  print((f'{total_successful_examples} of {len(corpus_elements)} modules '
+         f'succeeded, and {total_training_examples} trainining examples '
+         'written'))
 
 
 if __name__ == '__main__':
   flags.mark_flag_as_required('data_path')
-  app.run(main)
+  multiprocessing.handle_main(functools.partial(app.run, main))

compiler_opt/tools/generate_default_trace_test.py

@@ -28,6 +28,8 @@
 from compiler_opt.rl import compilation_runner
 from compiler_opt.tools import generate_default_trace
 
+from tf_agents.system import system_multiprocessing as multiprocessing
+
 flags.FLAGS['num_workers'].allow_override = True
 flags.FLAGS['gin_files'].allow_override = True
 flags.FLAGS['gin_bindings'].allow_override = True
@@ -105,12 +107,12 @@ def test_api(self, mock_get_runner):
         output_performance_path=os.path.join(tmp_dir.full_path,
                                              'output_performance'),
     ):
-      generate_default_trace.main(None)
+      generate_default_trace.main()
 
   def test_get_runner(self):
     runner = generate_default_trace.get_runner()
     self.assertIsInstance(runner, compilation_runner.CompilationRunner)
 
 
 if __name__ == '__main__':
-  absltest.main()
+  multiprocessing.handle_main(absltest.main)