Move logic to runners

Author: Northbadge

compiler_opt/distributed/local/local_worker_manager.py

@@ -32,9 +32,7 @@
 import functools
 import multiprocessing
 import threading
-import os
 import psutil
-import signal
 
 from absl import logging
 # pylint: disable=unused-import
@@ -134,7 +132,6 @@ def __init__(self):
       # when we stop.
       self._lock = threading.Lock()
       self._map: Dict[int, concurrent.futures.Future] = {}
-      self.is_paused = False
 
       # thread draining the pipe
       self._pump = threading.Thread(target=self._msg_pump)
@@ -209,25 +206,10 @@ def shutdown(self):
       try:
         # Killing the process triggers observer exit, which triggers msg_pump
         # exit
-        self.resume()
         self._process.kill()
       except:  # pylint: disable=bare-except
         pass
 
-    def pause(self):
-      if self.is_paused:
-        return
-      self.is_paused = True
-      # used to send the STOP signal; does not actually kill the process
-      os.kill(self._process.pid, signal.SIGSTOP)
-
-    def resume(self):
-      if not self.is_paused:
-        return
-      self.is_paused = False
-      # used to send the CONTINUE signal; does not actually kill the process
-      os.kill(self._process.pid, signal.SIGCONT)
-
     def set_nice(self, val: int):
       """Sets the nice-ness of the process, this modifies how the OS
       schedules it. Only works on Unix, since val is presumed to be an int.

compiler_opt/distributed/local/local_worker_manager_test.py

@@ -58,25 +58,6 @@ def method(self):
     time.sleep(3600)
 
 
-class JobCounter(Worker):
-  """Test worker."""
-
-  def __init__(self):
-    self.times = []
-
-  @classmethod
-  def is_priority_method(cls, method_name: str) -> bool:
-    return method_name == 'get_times'
-
-  def start(self):
-    while True:
-      self.times.append(time.time())
-      time.sleep(0.05)
-
-  def get_times(self):
-    return self.times
-
-
 class LocalWorkerManagerTest(absltest.TestCase):
 
   def test_pool(self):
@@ -118,34 +99,6 @@ def test_worker_crash_while_waiting(self):
         with self.assertRaises(concurrent.futures.CancelledError):
           _ = f.result()
 
-  def test_pause_resume(self):
-
-    with local_worker_manager.LocalWorkerPool(JobCounter, 1) as pool:
-      p = pool[0]
-
-      # Fill the q for 1 second
-      p.start()
-      time.sleep(1)
-
-      # Then pause the process for 1 second
-      p.pause()
-      time.sleep(1)
-
-      # Then resume the process and wait 1 more second
-      p.resume()
-      time.sleep(1)
-
-      times = p.get_times().result()
-
-      # If pause/resume worked, there should be a gap of at least 0.5 seconds.
-      # Otherwise, this will throw an exception.
-      self.assertNotEmpty(times)
-      last_time = times[0]
-      for cur_time in times:
-        if cur_time - last_time > 0.5:
-          return
-      raise ValueError('Failed to find a 2 second gap in times.')
-
 
 if __name__ == '__main__':
   multiprocessing.handle_test_main(absltest.main)

compiler_opt/rl/compilation_runner.py

@@ -18,9 +18,11 @@
 import dataclasses
 import json
 import os
+import signal
 import subprocess
 import threading
-from typing import Dict, List, Optional, Tuple
+import time
+from typing import Dict, List, Optional, Tuple, Union
 
 from absl import flags
 from compiler_opt.distributed.worker import Worker, WorkerFuture
@@ -102,13 +104,22 @@ class WorkerCancellationManager:
   managing resources.
   """
 
-  def __init__(self):
+  @dataclasses.dataclass
+  class ProcData:
+    process: 'subprocess.Popen[bytes]'
+    timeout: threading.Timer
+    time_left: float
+    start_time: float
+
+  def __init__(self, timeout: float = _COMPILATION_TIMEOUT.value):
     # the queue is filled only by workers, and drained only by the single
     # consumer. we use _done to manage access to the queue. We can then assume
     # empty() is accurate and get() never blocks.
-    self._processes = set()
+    self._processes: Dict[int, WorkerCancellationManager.ProcData] = {}
     self._done = False
+    self._paused = False
     self._lock = threading.Lock()
+    self._timeout = timeout
 
   def enable(self):
     with self._lock:
@@ -118,26 +129,70 @@ def register_process(self, p: 'subprocess.Popen[bytes]'):
     """Register a process for potential cancellation."""
     with self._lock:
       if not self._done:
-        self._processes.add(p)
+        self._processes[p.pid] = self.ProcData(
+            process=p,
+            timeout=threading.Timer(self._timeout,
+                                    kill_process_ignore_exceptions, (p,)),
+            time_left=self._timeout,
+            start_time=time.time())
+        self._processes[p.pid].timeout.start()
         return
     kill_process_ignore_exceptions(p)
 
   def kill_all_processes(self):
     """Cancel any pending work."""
     with self._lock:
       self._done = True
-    for p in self._processes:
-      kill_process_ignore_exceptions(p)
+    for pdata in self._processes.values():
+      kill_process_ignore_exceptions(pdata.process)
+
+  def pause_all_processes(self):
+    with self._lock:
+      if self._paused:
+        return
+      self._paused = True
+
+      cur_time = time.time()
+      for pid, pdata in self._processes.items():
+        pdata.timeout.cancel()
+        pdata.time_left -= cur_time - pdata.start_time
+        if pdata.time_left > 0:
+          # used to send the STOP signal; does not actually kill the process
+          os.kill(pid, signal.SIGSTOP)
+        else:
+          #  In case we cancelled right after the timeout expired,
+          #  but before actually killing the process.
+          kill_process_ignore_exceptions(pdata.process)
+
+  def resume_all_processes(self):
+    with self._lock:
+      if not self._paused:
+        return
+      self._paused = False
+
+      cur_time = time.time()
+      for pid, pdata in self._processes.items():
+        pdata.timeout = threading.Timer(pdata.time_left,
+                                        kill_process_ignore_exceptions,
+                                        (pdata.process,))
+        pdata.timeout.start()
+        pdata.start_time = cur_time
+        # used to send the CONTINUE signal; does not actually kill the process
+        os.kill(pid, signal.SIGCONT)
 
   def unregister_process(self, p: 'subprocess.Popen[bytes]'):
     with self._lock:
-      if not self._done:
-        self._processes.remove(p)
+      if p.pid in self._processes:
+        self._processes[p.pid].timeout.cancel()
+        del self._processes[p.pid]
+
+  def __del__(self):
+    if len(self._processes) > 0:
+      raise RuntimeError('Cancellation manager deleted while containing items.')
 
 
 def start_cancellable_process(
     cmdline: List[str],
-    timeout: float,
     cancellation_manager: Optional[WorkerCancellationManager],
     want_output: bool = False) -> Optional[bytes]:
   """Start a cancellable process.
@@ -166,14 +221,10 @@ def start_cancellable_process(
     if cancellation_manager:
       cancellation_manager.register_process(p)
 
-    try:
-      retcode = p.wait(timeout=timeout)
-    except subprocess.TimeoutExpired as e:
-      kill_process_ignore_exceptions(p)
-      raise e
-    finally:
-      if cancellation_manager:
-        cancellation_manager.unregister_process(p)
+    retcode = p.wait()
+
+    if cancellation_manager:
+      cancellation_manager.unregister_process(p)
     if retcode != 0:
       raise ProcessKilledError(
       ) if retcode == -9 else subprocess.CalledProcessError(retcode, cmdline)
@@ -249,12 +300,16 @@ class CompilationRunner(Worker):
 
   @classmethod
   def is_priority_method(cls, method_name: str) -> bool:
-    return method_name == 'cancel_all_work'
-
-  def __init__(self,
-               clang_path: Optional[str] = None,
-               launcher_path: Optional[str] = None,
-               moving_average_decay_rate: float = 1):
+    return method_name in {
+        'cancel_all_work', 'pause_all_work', 'resume_all_work'
+    }
+
+  def __init__(
+      self,
+      clang_path: Optional[str] = None,
+      launcher_path: Optional[str] = None,
+      moving_average_decay_rate: float = 1,
+      cancellation_manager: Optional[WorkerCancellationManager] = None):
     """Initialization of CompilationRunner class.
 
     Args:
@@ -265,8 +320,8 @@ def __init__(self,
     self._clang_path = clang_path
     self._launcher_path = launcher_path
     self._moving_average_decay_rate = moving_average_decay_rate
-    self._compilation_timeout = _COMPILATION_TIMEOUT.value
-    self._cancellation_manager = WorkerCancellationManager()
+    self._cancellation_manager = (
+        cancellation_manager or WorkerCancellationManager())
 
   # re-allow the cancellation manager accept work.
   def enable(self):
@@ -275,6 +330,12 @@ def enable(self):
   def cancel_all_work(self):
     self._cancellation_manager.kill_all_processes()
 
+  def pause_all_work(self):
+    self._cancellation_manager.pause_all_processes()
+
+  def resume_all_work(self):
+    self._cancellation_manager.resume_all_processes()
+
   def collect_data(
       self, module_spec: corpus.ModuleSpec, tf_policy_path: str,
       reward_stat: Optional[Dict[str, RewardStat]]) -> CompilationResult:

compiler_opt/rl/compilation_runner_test.py

@@ -17,6 +17,7 @@
 import os
 import string
 import subprocess
+import threading
 import time
 from unittest import mock
 
@@ -213,9 +214,9 @@ def test_exception_handling(self, mock_compile_fn):
     self.assertEqual(1, mock_compile_fn.call_count)
 
   def test_start_subprocess_output(self):
-    ct = compilation_runner.WorkerCancellationManager()
+    cm = compilation_runner.WorkerCancellationManager(100)
     output = compilation_runner.start_cancellable_process(
-        ['ls', '-l'], timeout=100, cancellation_manager=ct, want_output=True)
+        ['ls', '-l'], cancellation_manager=cm, want_output=True)
     if output:
       output_str = output.decode('utf-8')
     else:
@@ -227,14 +228,31 @@ def test_timeout_kills_process(self):
                                  'test_timeout_kills_test_file')
     if os.path.exists(sentinel_file):
       os.remove(sentinel_file)
-    with self.assertRaises(subprocess.TimeoutExpired):
+    with self.assertRaises(compilation_runner.ProcessKilledError):
+      cm = compilation_runner.WorkerCancellationManager(0.5)
       compilation_runner.start_cancellable_process(
           ['bash', '-c', 'sleep 1s ; touch ' + sentinel_file],
-          timeout=0.5,
-          cancellation_manager=None)
+          cancellation_manager=cm)
     time.sleep(2)
     self.assertFalse(os.path.exists(sentinel_file))
 
+  def test_pause_resume(self):
+    # This also makes sure timeouts are restored properly.
+    cm = compilation_runner.WorkerCancellationManager(1)
+    start_time = time.time()
+
+    def stop_and_start():
+      time.sleep(0.25)
+      cm.pause_all_processes()
+      time.sleep(2)
+      cm.resume_all_processes()
+
+    threading.Thread(target=stop_and_start).start()
+    compilation_runner.start_cancellable_process(['sleep', '0.5'],
+                                                 cancellation_manager=cm)
+    # should be at least 2 seconds due to the pause.
+    self.assertGreater(time.time() - start_time, 2)
+
 
 if __name__ == '__main__':
   tf.test.main()

compiler_opt/rl/inlining/inlining_runner.py

@@ -90,12 +90,10 @@ def _compile_fn(
         command_line.extend(
             ['-mllvm', '-ml-inliner-model-under-training=' + tf_policy_path])
       compilation_runner.start_cancellable_process(command_line,
-                                                   self._compilation_timeout,
                                                    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,
           want_output=True)
       if not output_bytes:

compiler_opt/rl/local_data_collector_test.py

@@ -80,9 +80,13 @@ def mock_collect_data(module_spec, tf_policy_dir, reward_stat):
 class Sleeper(compilation_runner.CompilationRunner):
   """Test CompilationRunner that just sleeps."""
 
+  def __init__(self):
+    super().__init__(
+        cancellation_manager=compilation_runner.WorkerCancellationManager(3600))
+
   def collect_data(self, module_spec, tf_policy_path, reward_stat):
     _ = module_spec, tf_policy_path, reward_stat
-    compilation_runner.start_cancellable_process(['sleep', '3600s'], 3600,
+    compilation_runner.start_cancellable_process(['sleep', '3600s'],
                                                  self._cancellation_manager)
 
     return compilation_runner.CompilationResult(

compiler_opt/rl/regalloc/regalloc_runner.py

@@ -87,7 +87,6 @@ def _compile_fn(
       if tf_policy_path:
         command_line.extend(['-mllvm', '-regalloc-model=' + tf_policy_path])
       compilation_runner.start_cancellable_process(command_line,
-                                                   self._compilation_timeout,
                                                    cancellation_manager)
 
       sequence_example = struct_pb2.Struct()