add the support for dumping best trajectories during training (#153)

* add the option to dump best trajectories during training

* track->dump

* fix test

* typo

Author: yundiqian

compiler_opt/rl/compilation_runner.py

@@ -25,12 +25,13 @@
 
 from absl import flags
 from absl import logging
+import tensorflow as tf
+
 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
-import tensorflow as tf
+from compiler_opt.rl import policy_saver
 
 _COMPILATION_TIMEOUT = flags.DEFINE_integer(
     'compilation_timeout', 60,
@@ -219,11 +220,12 @@ class CompilationResult:
   length: total length of all sequence examples, derived from sequence_examples.
   reward_stats: a dictionary from keys (e.g. function names) to a RewardStat.
   rewards: a list of reward values.
+  policy_rewards: a list of reward values under policy.
   keys: a list of keys.
 
   The object must observe the following invariants:
-  1) The entries of sequence_examples, rewards, and keys correspond to eachoter
-  at the same index.
+  1) The entries of sequence_examples, rewards, policy_rewards and keys
+  correspond to each other at the same index.
 
   2) The keys in reward stats are those in the keys field.
   """
@@ -232,6 +234,7 @@ class CompilationResult:
   length: int = dataclasses.field(init=False)
   reward_stats: Dict[str, RewardStat]
   rewards: List[float]
+  policy_rewards: List[float]
   keys: List[str]
 
   def __post_init__(self, sequence_examples: List[tf.train.SequenceExample]):
@@ -243,8 +246,8 @@ def __post_init__(self, sequence_examples: List[tf.train.SequenceExample]):
     ]
     object.__setattr__(self, 'length', sum(lengths))
 
-    assert (len(self.serialized_sequence_examples) == len(self.rewards) ==
-            (len(self.keys)))
+    assert (len(self.serialized_sequence_examples) == len(self.rewards) == len(
+        self.policy_rewards) == len(self.keys))
     assert set(self.keys) == set(self.reward_stats.keys())
     assert not hasattr(self, 'sequence_examples')
 
@@ -356,6 +359,7 @@ def collect_data(
 
     sequence_example_list = []
     rewards = []
+    policy_rewards = []
     keys = []
     for k, v in policy_result.items():
       sequence_example = v[0]
@@ -376,12 +380,14 @@ def collect_data(
           policy_reward * (1 - self._moving_average_decay_rate))
       rewards.append(
           _calculate_reward(policy=policy_reward, baseline=default_reward))
+      policy_rewards.append(policy_reward)
       keys.append(k)
 
     return CompilationResult(
         sequence_examples=sequence_example_list,
         reward_stats=reward_stat,
         rewards=rewards,
+        policy_rewards=policy_rewards,
         keys=keys)
 
   def compile_fn(

compiler_opt/rl/compilation_runner_test.py

@@ -134,6 +134,7 @@ def test_policy(self, mock_compile_fn):
                     (1 - _MOVING_AVERAGE_DECAY_RATE))
         }, data.reward_stats)
     self.assertAllClose([0.1998002], data.rewards)
+    self.assertAllClose([8], data.policy_rewards)
 
   @mock.patch(constant.BASE_MODULE_DIR +
               '.compilation_runner.CompilationRunner.compile_fn')
@@ -162,6 +163,7 @@ def test_default(self, mock_compile_fn):
                     moving_average_reward=_DEFAULT_REWARD)
         }, data.reward_stats)
     self.assertAllClose([0], data.rewards)
+    self.assertAllClose([10], data.policy_rewards)
 
   @mock.patch(constant.BASE_MODULE_DIR +
               '.compilation_runner.CompilationRunner.compile_fn')
@@ -197,6 +199,7 @@ def test_given_default_size(self, mock_compile_fn):
                     _POLICY_REWARD * (1 - _MOVING_AVERAGE_DECAY_RATE))
         }, data.reward_stats)
     self.assertAllClose([0.199800], data.rewards)
+    self.assertAllClose([8], data.policy_rewards)
 
   @mock.patch(constant.BASE_MODULE_DIR +
               '.compilation_runner.CompilationRunner.compile_fn')

compiler_opt/rl/local_data_collector.py

@@ -17,13 +17,14 @@
 import concurrent.futures
 import itertools
 import time
-from typing import Callable, Dict, Iterator, List, Tuple, Optional
+from typing import Callable, Dict, Iterator, List, Optional, Tuple
 
 from absl import logging
 from tf_agents.trajectories import trajectory
 
 from compiler_opt.distributed import worker
 from compiler_opt.distributed import buffered_scheduler
+from compiler_opt.rl import best_trajectory
 from compiler_opt.rl import compilation_runner
 from compiler_opt.rl import corpus
 from compiler_opt.rl import data_collector
@@ -41,6 +42,7 @@ def __init__(
       parser: Callable[[List[str]], Iterator[trajectory.Trajectory]],
       reward_stat_map: Dict[str, Optional[Dict[str,
                                                compilation_runner.RewardStat]]],
+      best_trajectory_repo: Optional[best_trajectory.BestTrajectoryRepo],
       exit_checker_ctor=data_collector.EarlyExitChecker):
     # TODO(mtrofin): type exit_checker_ctor when we get typing.Protocol support
     super().__init__()
@@ -53,6 +55,7 @@ def __init__(
         compilation_runner
         .CompilationRunnerStub] = self._worker_pool.get_currently_active()
     self._reward_stat_map = reward_stat_map
+    self._best_trajectory_repo = best_trajectory_repo
     self._exit_checker_ctor = exit_checker_ctor
     # _reset_workers is a future that resolves when post-data collection cleanup
     # work completes, i.e. cancelling all work and re-enabling the workers.
@@ -126,7 +129,7 @@ def collect_data(
     """Collect data for a given policy.
 
     Args:
-      policy_path: the path to the policy directory to collect data with.
+      policy: a policy_saver.Policy object to collect data with.
 
     Returns:
       An iterator of batched trajectory.Trajectory that are ready to be fed to
@@ -183,6 +186,15 @@ def wrapup():
     self._reward_stat_map.update(
         {spec.name: res.reward_stats for (spec, res) in successful_work})
 
+    if self._best_trajectory_repo is not None:
+      for spec, res in successful_work:
+        module_name = spec.name
+        for (identifier, reward,
+             sequence_example) in zip(res.keys, res.policy_rewards,
+                                      res.serialized_sequence_examples):
+          self._best_trajectory_repo.update_if_better_trajectory(
+              module_name, identifier, reward, sequence_example)
+
     monitor_dict = {}
     monitor_dict['default'] = {
         'success_modules': len(successful_work),

compiler_opt/rl/local_data_collector_test.py

@@ -16,24 +16,22 @@
 
 # pylint: disable=protected-access
 import collections
-
 import string
 import sys
+from typing import List, Tuple
 
 import tensorflow as tf
 from tf_agents.system import system_multiprocessing as multiprocessing
 
+# This is https://github.com/google/pytype/issues/764
+from google.protobuf import text_format  # pytype: disable=pyi-error
 from compiler_opt.distributed.local.local_worker_manager import LocalWorkerPoolManager
 from compiler_opt.rl import compilation_runner
 from compiler_opt.rl import corpus
 from compiler_opt.rl import data_collector
 from compiler_opt.rl import local_data_collector
 from compiler_opt.rl import policy_saver
 
-# This is https://github.com/google/pytype/issues/764
-from google.protobuf import text_format  # pytype: disable=pyi-error
-from typing import List, Tuple
-
 _policy_str = 'policy'.encode(encoding='utf-8')
 
 _mock_policy = policy_saver.Policy(output_spec=bytes(), policy=_policy_str)
@@ -71,6 +69,7 @@ def mock_collect_data(loaded_module_spec: corpus.LoadedModuleSpec, policy,
                     default_reward=1, moving_average_reward=2)
         },
         rewards=[1.2],
+        policy_rewards=[36],
         keys=['default'])
   else:
     return compilation_runner.CompilationResult(
@@ -81,6 +80,7 @@ def mock_collect_data(loaded_module_spec: corpus.LoadedModuleSpec, policy,
                     default_reward=1, moving_average_reward=3)
         },
         rewards=[3.4],
+        policy_rewards=[18],
         keys=['default'])
 
 
@@ -93,7 +93,11 @@ def collect_data(self, loaded_module_spec, policy, reward_stat):
                                                  self._cancellation_manager)
 
     return compilation_runner.CompilationResult(
-        sequence_examples=[], reward_stats={}, rewards=[], keys=[])
+        sequence_examples=[],
+        reward_stats={},
+        rewards=[],
+        policy_rewards=[],
+        keys=[])
 
 
 class MyRunner(compilation_runner.CompilationRunner):
@@ -154,7 +158,8 @@ def _test_iterator_fn(data_list):
           num_modules=9,
           worker_pool=lwp,
           parser=create_test_iterator_fn(),
-          reward_stat_map=collections.defaultdict(lambda: None))
+          reward_stat_map=collections.defaultdict(lambda: None),
+          best_trajectory_repo=None)
 
       # reset the sampler, so the next time we collect, we collect the same
       # modules. We do it before the collect_data call, because that's when
@@ -226,6 +231,7 @@ def wait(self, _):
           worker_pool=lwp,
           parser=parser,
           reward_stat_map=collections.defaultdict(lambda: None),
+          best_trajectory_repo=None,
           exit_checker_ctor=QuickExiter)
       collector.collect_data(policy=_mock_policy)
       collector._join_pending_jobs()

compiler_opt/rl/train_locally.py

@@ -19,6 +19,7 @@
 import json
 import os
 import time
+from typing import List
 
 from absl import app
 from absl import flags
@@ -27,10 +28,10 @@
 import tensorflow as tf
 from tf_agents.agents import tf_agent
 from tf_agents.system import system_multiprocessing as multiprocessing
-from typing import List
 
 from compiler_opt.distributed.local.local_worker_manager import LocalWorkerPoolManager
 from compiler_opt.rl import agent_creators
+from compiler_opt.rl import best_trajectory
 from compiler_opt.rl import compilation_runner
 from compiler_opt.rl import constant
 from compiler_opt.rl import corpus
@@ -69,6 +70,7 @@ def train_eval(worker_manager_class=LocalWorkerPoolManager,
                train_sequence_length=1,
                deploy_policy_name='saved_policy',
                use_random_network_distillation=False,
+               dump_best_trajectory=False,
                moving_average_decay_rate=1):
   """Train for LLVM inliner."""
   root_dir = FLAGS.root_dir
@@ -134,6 +136,15 @@ def sequence_example_iterator_fn(seq_ex: List[str]):
     logging.info('Loaded Reward Stat Map from disk, containing %d modules',
                  len(reward_stat_map))
 
+  best_trajectory_repo = None
+  best_trajecroty_repo_path = os.path.join(root_dir,
+                                           'best_trajectory_repo.json')
+  if dump_best_trajectory:
+    best_trajectory_repo = best_trajectory.BestTrajectoryRepo(
+        action_name=action_spec.name)
+    if tf.io.gfile.exists(best_trajecroty_repo_path):
+      best_trajectory_repo.load_from_json_file(best_trajecroty_repo_path)
+
   with worker_manager_class(
       worker_class=problem_config.get_runner_type(),
       count=FLAGS.num_workers,
@@ -143,7 +154,8 @@ def sequence_example_iterator_fn(seq_ex: List[str]):
         num_modules=num_modules,
         worker_pool=worker_pool,
         parser=sequence_example_iterator_fn,
-        reward_stat_map=reward_stat_map)
+        reward_stat_map=reward_stat_map,
+        best_trajectory_repo=best_trajectory_repo)
 
     # Repeat for num_policy_iterations iterations.
     t1 = time.time()
@@ -155,6 +167,9 @@ def sequence_example_iterator_fn(seq_ex: List[str]):
       with tf.io.gfile.GFile(reward_stat_map_path, 'w') as f:
         json.dump(reward_stat_map, f, cls=constant.DataClassJSONEncoder)
 
+      if best_trajectory_repo is not None:
+        best_trajectory_repo.sink_to_json_file(best_trajecroty_repo_path)
+
       policy_path = os.path.join(root_dir, 'policy',
                                  str(llvm_trainer.global_step_numpy()))
       saver.save(policy_path)

compiler_opt/tools/generate_default_trace_test.py

@@ -58,6 +58,7 @@ def collect_data(self, loaded_module_spec, policy, reward_stat):
                     default_reward=1, moving_average_reward=2)
         },
         rewards=[1.2],
+        policy_rewards=[18],
         keys=['default'])