Add a model_id parameter to collect_data. (#178)

The model_id parameter allows joining compilation results with the model
used during the compilation. This is useful for debugging errors in
training as well as for ensuring experience freshness in the distributed
training setting.

Author: jacob-hegna

compiler_opt/rl/compilation_runner.py

@@ -237,6 +237,9 @@ class CompilationResult:
   policy_rewards: List[float]
   keys: List[str]
 
+  # The id of the model used to generate this compilation result
+  model_id: Optional[int]
+
   def __post_init__(self, sequence_examples: List[tf.train.SequenceExample]):
     object.__setattr__(self, 'serialized_sequence_examples',
                        [x.SerializeToString() for x in sequence_examples])
@@ -260,8 +263,8 @@ def collect_data(
       self,
       loaded_module_spec: corpus.LoadedModuleSpec,
       policy: Optional[policy_saver.Policy] = None,
-      reward_stat: Optional[Dict[str, RewardStat]] = None
-  ) -> WorkerFuture[CompilationResult]:
+      reward_stat: Optional[Dict[str, RewardStat]] = None,
+      model_id: Optional[int] = None) -> WorkerFuture[CompilationResult]:
     raise NotImplementedError()
 
   @abc.abstractmethod
@@ -315,17 +318,18 @@ def pause_all_work(self):
   def resume_all_work(self):
     self._cancellation_manager.resume_all_processes()
 
-  def collect_data(
-      self,
-      loaded_module_spec: corpus.LoadedModuleSpec,
-      policy: Optional[policy_saver.Policy] = None,
-      reward_stat: Optional[Dict[str, RewardStat]] = None) -> CompilationResult:
+  def collect_data(self,
+                   loaded_module_spec: corpus.LoadedModuleSpec,
+                   policy: Optional[policy_saver.Policy] = None,
+                   reward_stat: Optional[Dict[str, RewardStat]] = None,
+                   model_id: Optional[int] = None) -> CompilationResult:
     """Collect data for the given IR file and policy.
 
     Args:
       loaded_module_spec: a LoadedModuleSpec.
       policy: serialized policy.
       reward_stat: reward stat of this module, None if unknown.
+      model_id: id for the model used to collect data.
 
     Returns:
       A CompilationResult. In particular:
@@ -341,7 +345,8 @@ def collect_data(
       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')
+        model_id_suffix = f'-{model_id}' if model_id is not None else ''
+        tf_policy_path = os.path.join(tempdir, 'policy' + model_id_suffix)
         policy.to_filesystem(tf_policy_path)
 
       if reward_stat is None:
@@ -388,7 +393,8 @@ def collect_data(
         reward_stats=reward_stat,
         rewards=rewards,
         policy_rewards=policy_rewards,
-        keys=keys)
+        keys=keys,
+        model_id=model_id)
 
   def compile_fn(
       self, command_line: corpus.FullyQualifiedCmdLine, tf_policy_path: str,

compiler_opt/rl/data_collector.py

@@ -52,12 +52,13 @@ class DataCollector(metaclass=abc.ABCMeta):
 
   @abc.abstractmethod
   def collect_data(
-      self, policy: policy_saver.Policy
+      self, policy: policy_saver.Policy, model_id: int
   ) -> Tuple[Iterator[trajectory.Trajectory], Dict[str, Dict[str, float]]]:
     """Collect data for a given policy.
 
     Args:
       policy_path: the path to the policy directory to collect data with.
+      model_id: the id of the model used to collect data.
 
     Returns:
       An iterator of batched trajectory.Trajectory that are ready to be fed to
@@ -126,7 +127,7 @@ def wait(self, get_num_finished_work):
 
     Args:
       get_num_finished_work: a callable object which returns the amount of
-      finished work.
+        finished work.
 
     Returns:
       The amount of time waited.

compiler_opt/rl/local_data_collector.py

@@ -101,7 +101,7 @@ def _join_pending_jobs(self):
                  time.time() - t1)
 
   def _schedule_jobs(
-      self, policy: policy_saver.Policy,
+      self, policy: policy_saver.Policy, model_id: int,
       sampled_modules: List[corpus.LoadedModuleSpec]
   ) -> List[worker.WorkerFuture[compilation_runner.CompilationResult]]:
     # by now, all the pending work, which was signaled to cancel, must've
@@ -114,7 +114,7 @@ def _schedule_jobs(
     def work_factory(job):
 
       def work(w: compilation_runner.CompilationRunnerStub):
-        return w.collect_data(*job)
+        return w.collect_data(*job, model_id=model_id)
 
       return work
 
@@ -124,7 +124,7 @@ def work(w: compilation_runner.CompilationRunnerStub):
         work, self._workers, self._worker_pool.get_worker_concurrency())
 
   def collect_data(
-      self, policy: policy_saver.Policy
+      self, policy: policy_saver.Policy, model_id: int
   ) -> Tuple[Iterator[trajectory.Trajectory], Dict[str, Dict[str, float]]]:
     """Collect data for a given policy.
 
@@ -145,7 +145,8 @@ def collect_data(
     logging.info('resolving prefetched sample took: %d seconds',
                  time.time() - time1)
     self._next_sample = self._prefetch_next_sample()
-    self._current_futures = self._schedule_jobs(policy, sampled_modules)
+    self._current_futures = self._schedule_jobs(policy, model_id,
+                                                sampled_modules)
 
     def wait_for_termination():
       early_exit = self._exit_checker_ctor(num_modules=self._num_modules)

compiler_opt/rl/local_data_collector_test.py

@@ -52,7 +52,7 @@ def _get_sequence_example(feature_value):
 
 
 def mock_collect_data(loaded_module_spec: corpus.LoadedModuleSpec, policy,
-                      reward_stat):
+                      reward_stat, model_id):
   assert loaded_module_spec.name.startswith('dummy')
   assert policy.policy == _policy_str
   assert reward_stat is None or reward_stat == {
@@ -70,7 +70,8 @@ def mock_collect_data(loaded_module_spec: corpus.LoadedModuleSpec, policy,
         },
         rewards=[1.2],
         policy_rewards=[36],
-        keys=['default'])
+        keys=['default'],
+        model_id=model_id)
   else:
     return compilation_runner.CompilationResult(
         sequence_examples=[_get_sequence_example(feature_value=2)],
@@ -81,13 +82,14 @@ def mock_collect_data(loaded_module_spec: corpus.LoadedModuleSpec, policy,
         },
         rewards=[3.4],
         policy_rewards=[18],
-        keys=['default'])
+        keys=['default'],
+        model_id=model_id)
 
 
 class Sleeper(compilation_runner.CompilationRunner):
   """Test CompilationRunner that just sleeps."""
 
-  def collect_data(self, loaded_module_spec, policy, reward_stat):
+  def collect_data(self, loaded_module_spec, policy, reward_stat, model_id):
     _ = loaded_module_spec, policy, reward_stat
     compilation_runner.start_cancellable_process(['sleep', '3600s'], 3600,
                                                  self._cancellation_manager)
@@ -97,7 +99,8 @@ def collect_data(self, loaded_module_spec, policy, reward_stat):
         reward_stats={},
         rewards=[],
         policy_rewards=[],
-        keys=[])
+        keys=[],
+        model_id=model_id)
 
 
 class MyRunner(compilation_runner.CompilationRunner):
@@ -166,7 +169,8 @@ def _test_iterator_fn(data_list):
       # we'll re-sample to prefetch the next batch.
       sampler.reset()
 
-      data_iterator, monitor_dict = collector.collect_data(policy=_mock_policy)
+      data_iterator, monitor_dict = collector.collect_data(
+          policy=_mock_policy, model_id=0)
       data = list(data_iterator)
       self.assertEqual([1, 2, 3], data)
       expected_monitor_dict_subset = {
@@ -184,7 +188,8 @@ def _test_iterator_fn(data_list):
             **monitor_dict,
             **expected_monitor_dict_subset
         })
-      data_iterator, monitor_dict = collector.collect_data(policy=_mock_policy)
+      data_iterator, monitor_dict = collector.collect_data(
+          policy=_mock_policy, model_id=0)
       data = list(data_iterator)
       # because we reset the sampler, these are the same modules
       self.assertEqual([4, 5, 6], data)
@@ -233,7 +238,7 @@ def wait(self, _):
           reward_stat_map=collections.defaultdict(lambda: None),
           best_trajectory_repo=None,
           exit_checker_ctor=QuickExiter)
-      collector.collect_data(policy=_mock_policy)
+      collector.collect_data(policy=_mock_policy, model_id=0)
       collector._join_pending_jobs()
       killed = 0
       for w in collector._current_futures:

compiler_opt/rl/train_locally.py

@@ -176,7 +176,8 @@ def sequence_example_iterator_fn(seq_ex: List[str]):
 
       dataset_iter, monitor_dict = data_collector.collect_data(
           policy=policy_saver.Policy.from_filesystem(
-              os.path.join(policy_path, deploy_policy_name)))
+              os.path.join(policy_path, deploy_policy_name)),
+          model_id=llvm_trainer.global_step_numpy())
       llvm_trainer.train(dataset_iter, monitor_dict, num_iterations)
 
       data_collector.on_dataset_consumed(dataset_iter)

compiler_opt/tools/generate_default_trace.py

@@ -108,7 +108,8 @@ def worker(policy_path: Optional[str],
         data = runner.collect_data(
             loaded_module_spec=loaded_module_spec,
             policy=policy,
-            reward_stat=None)
+            reward_stat=None,
+            model_id=0)
         if not m:
           results_queue.put(
               (loaded_module_spec.name, data.serialized_sequence_examples,

compiler_opt/tools/generate_default_trace_test.py

@@ -36,7 +36,7 @@
 class MockCompilationRunner(compilation_runner.CompilationRunner):
   """A compilation runner just for test."""
 
-  def collect_data(self, loaded_module_spec, policy, reward_stat):
+  def collect_data(self, loaded_module_spec, policy, reward_stat, model_id):
     sequence_example_text = """
       feature_lists {
         feature_list {
@@ -59,7 +59,8 @@ def collect_data(self, loaded_module_spec, policy, reward_stat):
         },
         rewards=[1.2],
         policy_rewards=[18],
-        keys=['default'])
+        keys=['default'],
+        model_id=model_id)
 
 
 class GenerateDefaultTraceTest(absltest.TestCase):