Default to using self.cancellation_manager

Part of google#96
Northbadge · Aug 19, 2022 · 13c223b · 13c223b
1 parent 390faa2
commit 13c223b
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/compiler_opt/rl/inlining/inlining_runner.py b/compiler_opt/rl/inlining/inlining_runner.py
@@ -71,6 +71,9 @@ def compile_fn(
       cancelled work.
       RuntimeError: if llvm-size produces unexpected output.
     """
+    if cancellation_manager is None:
+      cancellation_manager = self._cancellation_manager
+
     working_dir = tempfile.mkdtemp()
 
     log_path = os.path.join(working_dir, 'log')

diff --git a/compiler_opt/rl/regalloc/regalloc_runner.py b/compiler_opt/rl/regalloc/regalloc_runner.py
@@ -69,6 +69,9 @@ def compile_fn(
         cancelled work.
       RuntimeError: if llvm-size produces unexpected output.
     """
+    if cancellation_manager is None:
+      cancellation_manager = self._cancellation_manager
+
     working_dir = tempfile.mkdtemp()
 
     log_path = os.path.join(working_dir, 'log')