Default to using self.cancellation_manager

compiler_opt/rl/compilation_runner.py

@@ -334,20 +334,14 @@ def collect_data(
     """
     if reward_stat is None:
       default_result = self.compile_fn(
-          module_spec,
-          tf_policy_path='',
-          reward_only=bool(tf_policy_path),
-          cancellation_manager=self._cancellation_manager)
+          module_spec, 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,
-          cancellation_manager=self._cancellation_manager)
+          module_spec, tf_policy_path, reward_only=False)
     else:
       policy_result = default_result
 
@@ -383,17 +377,13 @@ def collect_data(
 
   def compile_fn(
       self, module_spec: corpus.ModuleSpec, tf_policy_path: str,
-      reward_only: bool,
-      cancellation_manager: Optional[WorkerCancellationManager]
-  ) -> Dict[str, Tuple[tf.train.SequenceExample, float]]:
+      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.
       tf_policy_path: path to TF policy directory on local disk.
       reward_only: whether only return reward.
-      cancellation_manager: a WorkerCancellationManager to handle early
-        termination
 
     Returns:
       A dict mapping from example identifier to tuple containing:

compiler_opt/rl/compilation_runner_test.py

@@ -76,10 +76,8 @@ def _get_sequence_example(feature_value):
   return text_format.Parse(sequence_example_text, tf.train.SequenceExample())
 
 
-def _mock_compile_fn(file_paths, tf_policy_path, reward_only,
-                     cancellation_manager):  # pylint: disable=unused-argument
+def _mock_compile_fn(file_paths, tf_policy_path, reward_only):  # pylint: disable=unused-argument
   del file_paths
-  del cancellation_manager
   if tf_policy_path:
     sequence_example = _get_sequence_example(_POLICY_FEATURE_VALUE)
     native_size = _POLICY_REWARD

compiler_opt/rl/inlining/inlining_runner.py

@@ -17,7 +17,7 @@
 import io
 import os
 import tempfile
-from typing import Dict, Optional, Tuple
+from typing import Dict, Tuple
 
 import gin
 import tensorflow as tf
@@ -47,17 +47,13 @@ def __init__(self, llvm_size_path: str, *args, **kwargs):
 
   def compile_fn(
       self, module_spec: corpus.ModuleSpec, tf_policy_path: str,
-      reward_only: bool, cancellation_manager: Optional[
-          compilation_runner.WorkerCancellationManager]
-  ) -> Dict[str, Tuple[tf.train.SequenceExample, float]]:
+      reward_only: bool) -> Dict[str, Tuple[tf.train.SequenceExample, float]]:
     """Run inlining for the given IR file under the given policy.
 
     Args:
       module_spec: a ModuleSpec.
       tf_policy_path: path to TF policy direcoty on local disk.
       reward_only: whether only return native size.
-      cancellation_manager: handler for early termination by killing any running
-      processes
 
     Returns:
       A dict mapping from example identifier to tuple containing:
@@ -71,6 +67,7 @@ def compile_fn(
       cancelled work.
       RuntimeError: if llvm-size produces unexpected output.
     """
+
     working_dir = tempfile.mkdtemp()
 
     log_path = os.path.join(working_dir, 'log')
@@ -91,12 +88,12 @@ def compile_fn(
             ['-mllvm', '-ml-inliner-model-under-training=' + tf_policy_path])
       compilation_runner.start_cancellable_process(command_line,
                                                    self._compilation_timeout,
-                                                   cancellation_manager)
+                                                   self._cancellation_manager)
       command_line = [self._llvm_size_path, output_native_path]
       output_bytes = compilation_runner.start_cancellable_process(
           command_line,
           timeout=self._compilation_timeout,
-          cancellation_manager=cancellation_manager,
+          cancellation_manager=self._cancellation_manager,
           want_output=True)
       if not output_bytes:
         raise RuntimeError(f'Empty llvm-size output: {" ".join(command_line)}')

compiler_opt/rl/regalloc/regalloc_runner.py

@@ -18,7 +18,7 @@
 import io
 import os
 import tempfile
-from typing import Dict, Optional, Tuple
+from typing import Dict, Tuple
 
 import gin
 import tensorflow as tf
@@ -45,17 +45,13 @@ class RegAllocRunner(compilation_runner.CompilationRunner):
   # construction
   def compile_fn(
       self, module_spec: corpus.ModuleSpec, tf_policy_path: str,
-      reward_only: bool, cancellation_manager: Optional[
-          compilation_runner.WorkerCancellationManager]
-  ) -> Dict[str, Tuple[tf.train.SequenceExample, float]]:
+      reward_only: bool) -> Dict[str, Tuple[tf.train.SequenceExample, float]]:
     """Run inlining for the given IR file under the given policy.
 
     Args:
       module_spec: a ModuleSpec.
       tf_policy_path: path to TF policy direcoty on local disk.
       reward_only: whether only return reward.
-      cancellation_manager: handler for early termination by killing any running
-        processes
 
     Returns:
       A dict mapping from example identifier to tuple containing:
@@ -69,6 +65,7 @@ def compile_fn(
         cancelled work.
       RuntimeError: if llvm-size produces unexpected output.
     """
+
     working_dir = tempfile.mkdtemp()
 
     log_path = os.path.join(working_dir, 'log')
@@ -88,7 +85,7 @@ def compile_fn(
         command_line.extend(['-mllvm', '-regalloc-model=' + tf_policy_path])
       compilation_runner.start_cancellable_process(command_line,
                                                    self._compilation_timeout,
-                                                   cancellation_manager)
+                                                   self._cancellation_manager)
 
       sequence_example = struct_pb2.Struct()