Make the corpus remote wrt workers (#141)

* Make the corpus remote wrt workers

This change removes the assumption that the corpus is co-located with the workers. The new design relies on the assumption that the workers are reachable via a high-throughput network (datacenter setup). The trainer prefetches the modules for the next iteration and sends them as data blobs to the workers.

The rest of the change trickles out from this:

- the corpus operates internally with a metadata object about modules, (`ModuleSpec`)
- workers are given a `LoadedModuleSpec` which contains the data of the module (and its thinlto index, if needed).
- command line final shape is done on the worker side

This had the side-effect that responsibilities could be reassigned as follows:

- the corpus constructor handles preparing the `ModuleSpec`s, applying all the necessary flag modification handlers (i.e. add/delete/replace) 
- flag modification logic is generic - it doesn't know anything about thinlto, '-cc1', etc. 
- flag contextualization happens on the worker side. Flags reference a context object as a string formatting value, which is then provided by the worker.

In turn, this impacted testing.

This change impacts purely local training similar to how a previous change remoting the policy does - local files are copied by workers redundantly. It doesn not appear to have any impact on training performance, but should we want to, we can optimize the local scenario under the hood.

Author: mtrofin

compiler_opt/rl/compilation_runner.py

@@ -24,7 +24,9 @@
 from typing import Dict, List, Optional, Tuple
 
 from absl import flags
-from compiler_opt.distributed.worker import Worker, WorkerFuture
+from absl import logging
+from compiler_opt.distributed.worker import Worker
+from compiler_opt.distributed.worker import WorkerFuture
 from compiler_opt.rl import constant
 from compiler_opt.rl import policy_saver
 from compiler_opt.rl import corpus
@@ -178,6 +180,8 @@ def start_cancellable_process(
   # Disable tensorflow info messages during data collection
   if _QUIET.value:
     command_env['TF_CPP_MIN_LOG_LEVEL'] = '1'
+  else:
+    logging.info(cmdline)
   with subprocess.Popen(
       cmdline,
       env=command_env,
@@ -251,7 +255,7 @@ class CompilationRunnerStub(metaclass=abc.ABCMeta):
   @abc.abstractmethod
   def collect_data(
       self,
-      module_spec: corpus.ModuleSpec,
+      loaded_module_spec: corpus.LoadedModuleSpec,
       policy: Optional[policy_saver.Policy] = None,
       reward_stat: Optional[Dict[str, RewardStat]] = None
   ) -> WorkerFuture[CompilationResult]:
@@ -310,13 +314,13 @@ def resume_all_work(self):
 
   def collect_data(
       self,
-      module_spec: corpus.ModuleSpec,
+      loaded_module_spec: corpus.LoadedModuleSpec,
       policy: Optional[policy_saver.Policy] = None,
       reward_stat: Optional[Dict[str, RewardStat]] = None) -> CompilationResult:
     """Collect data for the given IR file and policy.
 
     Args:
-      module_spec: a ModuleSpec.
+      loaded_module_spec: a LoadedModuleSpec.
       policy: serialized policy.
       reward_stat: reward stat of this module, None if unknown.
 
@@ -331,21 +335,22 @@ def collect_data(
       ValueError if example under default policy and ml policy does not match.
     """
     with tempfile.TemporaryDirectory() as tempdir:
+      final_cmd_line = loaded_module_spec.build_command_line(tempdir)
       tf_policy_path = ''
       if policy is not None:
         tf_policy_path = os.path.join(tempdir, 'policy')
         policy.to_filesystem(tf_policy_path)
 
       if reward_stat is None:
         default_result = self.compile_fn(
-            module_spec, tf_policy_path='', reward_only=bool(tf_policy_path))
+            final_cmd_line, tf_policy_path='', reward_only=bool(tf_policy_path))
         reward_stat = {
             k: RewardStat(v[1], v[1]) for (k, v) in default_result.items()
         }
 
       if tf_policy_path:
         policy_result = self.compile_fn(
-            module_spec, tf_policy_path, reward_only=False)
+            final_cmd_line, tf_policy_path, reward_only=False)
       else:
         policy_result = default_result
 
@@ -358,7 +363,7 @@ def collect_data(
       if k not in reward_stat:
         raise ValueError(
             (f'Example {k} does not exist under default policy for '
-             f'module {module_spec.name}'))
+             f'cmd line: {final_cmd_line}'))
       default_reward = reward_stat[k].default_reward
       moving_average_reward = reward_stat[k].moving_average_reward
       sequence_example = _overwrite_trajectory_reward(
@@ -380,12 +385,12 @@ def collect_data(
         keys=keys)
 
   def compile_fn(
-      self, module_spec: corpus.ModuleSpec, tf_policy_path: str,
+      self, command_line: corpus.FullyQualifiedCmdLine, tf_policy_path: str,
       reward_only: bool) -> Dict[str, Tuple[tf.train.SequenceExample, float]]:
     """Compiles for the given IR file under the given policy.
 
     Args:
-      module_spec: a ModuleSpec.
+      command_line: the fully qualified command line.
       tf_policy_path: path to TF policy directory on local disk.
       reward_only: whether only return reward.
 

compiler_opt/rl/compilation_runner_test.py

@@ -94,6 +94,9 @@ def _mock_compile_fn(file_paths, tf_policy_path, reward_only):  # pylint: disabl
 
 _mock_policy = policy_saver.Policy(bytes(), bytes())
 
+_mock_loaded_module_spec = corpus.LoadedModuleSpec(
+    name='dummy', loaded_ir=bytes())
+
 
 class CompilationRunnerTest(tf.test.TestCase):
 
@@ -111,7 +114,7 @@ def test_policy(self, mock_compile_fn):
     runner = compilation_runner.CompilationRunner(
         moving_average_decay_rate=_MOVING_AVERAGE_DECAY_RATE)
     data = runner.collect_data(
-        module_spec=corpus.ModuleSpec(name='dummy'), policy=_mock_policy)
+        loaded_module_spec=_mock_loaded_module_spec, policy=_mock_policy)
     self.assertEqual(2, mock_compile_fn.call_count)
 
     expected_example = _get_sequence_example_with_reward(
@@ -139,7 +142,7 @@ def test_default(self, mock_compile_fn):
     runner = compilation_runner.CompilationRunner(
         moving_average_decay_rate=_MOVING_AVERAGE_DECAY_RATE)
 
-    data = runner.collect_data(module_spec=corpus.ModuleSpec(name='dummy'))
+    data = runner.collect_data(loaded_module_spec=_mock_loaded_module_spec)
     # One call when we ask for the default policy, because it can provide both
     # trace and default size.
     self.assertEqual(1, mock_compile_fn.call_count)
@@ -168,7 +171,7 @@ def test_given_default_size(self, mock_compile_fn):
         moving_average_decay_rate=_MOVING_AVERAGE_DECAY_RATE)
 
     data = runner.collect_data(
-        module_spec=corpus.ModuleSpec(name='dummy'),
+        loaded_module_spec=_mock_loaded_module_spec,
         policy=_mock_policy,
         reward_stat={
             'default':
@@ -205,7 +208,7 @@ def test_exception_handling(self, mock_compile_fn):
 
     with self.assertRaisesRegex(subprocess.CalledProcessError, 'error'):
       _ = runner.collect_data(
-          module_spec=corpus.ModuleSpec(name='dummy'),
+          loaded_module_spec=_mock_loaded_module_spec,
           policy=_mock_policy,
           reward_stat=None)
     self.assertEqual(1, mock_compile_fn.call_count)