Make generate_default_trace_test closer to actual use (#225)

- making `main` take care of gin bindings, and any additional
  parameters, so the worker pool type doesn't get overridden by `main`'s
  args

- making sure gin configuration for the underlying worker makes it
  through.

This also fixes the actual tool, since we weren't passing the gin config
of the underlying task correctly.

Author: mtrofin

compiler_opt/tools/generate_default_trace.py

@@ -75,11 +75,6 @@
                           BaseException]
 
 
-def get_runner() -> compilation_runner.CompilationRunner:
-  problem_config = registry.get_configuration()
-  return problem_config.get_runner_type()(moving_average_decay_rate=0)
-
-
 class FilteringWorker(worker.Worker):
   """Worker that performs a computation and optionally filters the result.
 
@@ -89,10 +84,12 @@ class FilteringWorker(worker.Worker):
     key_filter: regex filter for key names to include, or None to include all.
   """
 
-  def __init__(self, policy_path: Optional[str], key_filter: Optional[str]):
+  def __init__(self, policy_path: Optional[str], key_filter: Optional[str],
+               runner_type: 'type[compilation_runner.CompilationRunner]',
+               runner_kwargs):
     self._policy_path = policy_path
     self._key_filter = re.compile(key_filter) if key_filter else None
-    self._runner = get_runner()
+    self._runner = runner_type(**runner_kwargs)
     self._policy = policy_saver.Policy.from_filesystem(
         policy_path) if policy_path else None
 
@@ -118,11 +115,15 @@ def compile_and_filter(
     return (loaded_module_spec.name, new_sequence_examples, new_reward_stats)
 
 
-def main(worker_manager_class=local_worker_manager.LocalWorkerPoolManager):
-
+def main(_):
   gin.parse_config_files_and_bindings(
       _GIN_FILES.value, bindings=_GIN_BINDINGS.value, skip_unknown=False)
   logging.info(gin.config_str())
+  generate_trace()
+
+
+def generate_trace(
+    worker_manager_class=local_worker_manager.LocalWorkerPoolManager):
 
   config = registry.get_configuration()
 
@@ -155,13 +156,17 @@ def main(worker_manager_class=local_worker_manager.LocalWorkerPoolManager):
       cps.load_module_spec(corpus_element) for corpus_element in corpus_elements
   ]
 
+  runner_type = config.get_runner_type()
   with tfrecord_context as tfrecord_writer:
     with performance_context as performance_writer:
       with worker_manager_class(
           FilteringWorker,
           _NUM_WORKERS.value,
           policy_path=_POLICY_PATH.value,
-          key_filter=_KEY_FILTER.value) as lwm:
+          key_filter=_KEY_FILTER.value,
+          runner_type=runner_type,
+          runner_kwargs=worker.get_full_worker_args(
+              runner_type, moving_average_decay_rate=0)) as lwm:
 
         _, result_futures = buffered_scheduler.schedule_on_worker_pool(
             action=lambda w, j: w.compile_and_filter(j),

compiler_opt/tools/generate_default_trace_test.py

@@ -35,9 +35,14 @@
 flags.FLAGS['gin_bindings'].allow_override = True
 
 
+@gin.configurable(module='runners')
 class MockCompilationRunner(compilation_runner.CompilationRunner):
   """A compilation runner just for test."""
 
+  def __init__(self, sentinel=None):
+    assert sentinel == 42
+    super().__init__()
+
   def collect_data(self,
                    loaded_module_spec,
                    policy=None,
@@ -75,11 +80,11 @@ def setUp(self):
     with gin.unlock_config():
       gin.parse_config_files_and_bindings(
           config_files=['compiler_opt/rl/inlining/gin_configs/common.gin'],
-          bindings=None)
+          bindings=['runners.MockCompilationRunner.sentinel=42'])
     return super().setUp()
 
-  @mock.patch('compiler_opt.tools.generate_default_trace.get_runner')
-  def test_api(self, mock_get_runner):
+  @mock.patch('compiler_opt.rl.inlining.InliningConfig.get_runner_type')
+  def test_generate_trace(self, mock_get_runner):
 
     tmp_dir = self.create_tempdir()
     module_names = ['a', 'b', 'c', 'd']
@@ -97,7 +102,7 @@ def test_api(self, mock_get_runner):
           os.path.join(tmp_dir.full_path, module_name + '.cmd'), 'w') as f:
         f.write('-cc1')
 
-    mock_compilation_runner = MockCompilationRunner()
+    mock_compilation_runner = MockCompilationRunner
     mock_get_runner.return_value = mock_compilation_runner
 
     with flagsaver.flagsaver(
@@ -107,11 +112,7 @@ def test_api(self, mock_get_runner):
         output_performance_path=os.path.join(tmp_dir.full_path,
                                              'output_performance'),
     ):
-      generate_default_trace.main()
-
-  def test_get_runner(self):
-    runner = generate_default_trace.get_runner()
-    self.assertIsInstance(runner, compilation_runner.CompilationRunner)
+      generate_default_trace.generate_trace()
 
 
 if __name__ == '__main__':