Remove custom timeout

Author: Northbadge

compiler_opt/distributed/local/local_worker_manager.py

@@ -31,16 +31,16 @@
 import dataclasses
 import functools
 import multiprocessing
-import threading
 import psutil
+import threading
 
 from absl import logging
 # pylint: disable=unused-import
 from compiler_opt.distributed.worker import Worker
 
 from contextlib import AbstractContextManager
 from multiprocessing import connection
-from typing import Any, Callable, Dict, Optional, List
+from typing import Any, Callable, Dict, List, Optional
 
 
 @dataclasses.dataclass(frozen=True)

compiler_opt/distributed/worker.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 """Common abstraction for a worker contract."""
 
-from typing import Iterable, Optional, TypeVar, Protocol, runtime_checkable
+from typing import Iterable, Optional, Protocol, TypeVar
 
 
 class Worker(Protocol):
@@ -25,17 +25,6 @@ def is_priority_method(cls, method_name: str) -> bool:
     return False
 
 
-@runtime_checkable
-class ContextAwareWorker(Worker, Protocol):
-  """ContextAwareWorkers use set_context to modify internal state, this allows
-  it to behave differently when run remotely vs locally. The user of a
-  ContextAwareWorker can check for this with isinstance(obj, ContextAwareWorker)
-  """
-
-  def set_context(self, local: bool) -> None:
-    return
-
-
 T = TypeVar('T')
 
 

compiler_opt/rl/compilation_runner.py

@@ -21,7 +21,6 @@
 import signal
 import subprocess
 import threading
-import time
 from typing import Dict, List, Optional, Tuple
 
 from absl import flags
@@ -104,22 +103,14 @@ class WorkerCancellationManager:
   managing resources.
   """
 
-  @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):
+  def __init__(self):
     # 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: Dict[int, WorkerCancellationManager.ProcData] = {}
+    self._processes = set()
     self._done = False
     self._paused = False
     self._lock = threading.Lock()
-    self._timeout = timeout
 
   def enable(self):
     with self._lock:
@@ -129,65 +120,40 @@ def register_process(self, p: 'subprocess.Popen[bytes]'):
     """Register a process for potential cancellation."""
     with self._lock:
       if not self._done:
-        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())
-        if self._paused:
-          os.kill(p.pid, signal.SIGSTOP)
-        else:
-          self._processes[p.pid].timeout.start()
+        self._processes.add(p)
         return
     kill_process_ignore_exceptions(p)
 
   def kill_all_processes(self):
     """Cancel any pending work."""
     with self._lock:
       self._done = True
-    for pdata in self._processes.values():
-      kill_process_ignore_exceptions(pdata.process)
+    for p in self._processes:
+      kill_process_ignore_exceptions(p)
 
   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)
+      for p in self._processes:
+        # used to send the STOP signal; does not actually kill the process
+        os.kill(p.pid, signal.SIGSTOP)
 
   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
+      for p in self._processes:
         # used to send the CONTINUE signal; does not actually kill the process
-        os.kill(pid, signal.SIGCONT)
+        os.kill(p.pid, signal.SIGCONT)
 
   def unregister_process(self, p: 'subprocess.Popen[bytes]'):
     with self._lock:
-      if p.pid in self._processes:
-        self._processes[p.pid].timeout.cancel()
-        del self._processes[p.pid]
+      self._processes.remove(p)
 
   def __del__(self):
     if len(self._processes) > 0:
@@ -196,6 +162,7 @@ def __del__(self):
 
 def start_cancellable_process(
     cmdline: List[str],
+    timeout: float,
     cancellation_manager: Optional[WorkerCancellationManager],
     want_output: bool = False) -> Optional[bytes]:
   """Start a cancellable process.
@@ -224,10 +191,15 @@ def start_cancellable_process(
     if cancellation_manager:
       cancellation_manager.register_process(p)
 
-    retcode = p.wait()
+    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)
 
-    if cancellation_manager:
-      cancellation_manager.unregister_process(p)
     if retcode != 0:
       raise ProcessKilledError(
       ) if retcode == -9 else subprocess.CalledProcessError(retcode, cmdline)
@@ -307,12 +279,10 @@ def is_priority_method(cls, method_name: str) -> bool:
         '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):
+  def __init__(self,
+               clang_path: Optional[str] = None,
+               launcher_path: Optional[str] = None,
+               moving_average_decay_rate: float = 1):
     """Initialization of CompilationRunner class.
 
     Args:
@@ -323,8 +293,8 @@ def __init__(
     self._clang_path = clang_path
     self._launcher_path = launcher_path
     self._moving_average_decay_rate = moving_average_decay_rate
-    self._cancellation_manager = (
-        cancellation_manager or WorkerCancellationManager())
+    self._compilation_timeout = _COMPILATION_TIMEOUT.value
+    self._cancellation_manager = WorkerCancellationManager()
 
   # re-allow the cancellation manager accept work.
   def enable(self):

compiler_opt/rl/compilation_runner_test.py

@@ -214,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):
-    cm = compilation_runner.WorkerCancellationManager(100)
+    cm = compilation_runner.WorkerCancellationManager()
     output = compilation_runner.start_cancellable_process(
-        ['ls', '-l'], cancellation_manager=cm, want_output=True)
+        ['ls', '-l'], timeout=100, cancellation_manager=cm, want_output=True)
     if output:
       output_str = output.decode('utf-8')
     else:
@@ -228,30 +228,30 @@ def test_timeout_kills_process(self):
                                  'test_timeout_kills_test_file')
     if os.path.exists(sentinel_file):
       os.remove(sentinel_file)
-    with self.assertRaises(compilation_runner.ProcessKilledError):
-      cm = compilation_runner.WorkerCancellationManager(0.5)
+    with self.assertRaises(subprocess.TimeoutExpired):
       compilation_runner.start_cancellable_process(
           ['bash', '-c', 'sleep 1s ; touch ' + sentinel_file],
-          cancellation_manager=cm)
+          timeout=0.5,
+          cancellation_manager=None)
     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)
+    cm = compilation_runner.WorkerCancellationManager()
     start_time = time.time()
 
     def stop_and_start():
       time.sleep(0.25)
       cm.pause_all_processes()
-      time.sleep(2)
+      time.sleep(1)
       cm.resume_all_processes()
 
     threading.Thread(target=stop_and_start).start()
     compilation_runner.start_cancellable_process(['sleep', '0.5'],
+                                                 30,
                                                  cancellation_manager=cm)
-    # should be at least 2 seconds due to the pause.
-    self.assertGreater(time.time() - start_time, 2)
+    # should be at least 1 second due to the pause.
+    self.assertGreater(time.time() - start_time, 1)
 
 
 if __name__ == '__main__':

compiler_opt/rl/inlining/inlining_runner.py

@@ -90,10 +90,12 @@ 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,13 +80,9 @@ 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'],
+    compilation_runner.start_cancellable_process(['sleep', '3600s'], 3600,
                                                  self._cancellation_manager)
 
     return compilation_runner.CompilationResult(

compiler_opt/rl/regalloc/regalloc_runner.py

@@ -87,6 +87,7 @@ 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()