Use LocalWorkerManager (#48)

This change switches to using LocalWorkerManager. The main change to the compilation runner code is that cancellation is now handled 'all-out' - we just cancel all jobs on the worker.

Issue #31

Author: mtrofin

compiler_opt/distributed/worker.py

@@ -14,10 +14,46 @@
 # limitations under the License.
 """Common abstraction for a worker contract."""
 
+import abc
+from typing import Generic, Iterable, TypeVar
+
 
 class Worker:
 
   @classmethod
   def is_priority_method(cls, method_name: str) -> bool:
     _ = method_name
     return False
+
+
+T = TypeVar('T')
+
+
+# Dask's Futures are limited. This captures that.
+class WorkerFuture(Generic[T], metaclass=abc.ABCMeta):
+
+  @abc.abstractmethod
+  def result(self) -> T:
+    raise NotImplementedError()
+
+  @abc.abstractmethod
+  def done(self) -> bool:
+    raise NotImplementedError()
+
+
+def wait_for(futures: Iterable[WorkerFuture]):
+  """Dask futures don't support more than result() and done()."""
+  for f in futures:
+    try:
+      _ = f.result()
+    except:  # pylint: disable=bare-except
+      pass
+
+
+def get_exception(worker_future: WorkerFuture):
+  assert worker_future.done()
+  try:
+    _ = worker_future.result()
+    return None
+  except Exception as e:  # pylint: disable=broad-except
+    return e

compiler_opt/rl/compilation_runner.py

@@ -14,18 +14,17 @@
 # limitations under the License.
 """Module for running compilation and collect training data."""
 
-import concurrent
+import abc
 import dataclasses
 import json
-import multiprocessing
 import subprocess
 import threading
 from typing import Dict, List, Optional, Tuple
 
 from absl import flags
-import tensorflow as tf
-
+from compiler_opt.distributed.worker import Worker, WorkerFuture
 from compiler_opt.rl import constant
+import tensorflow as tf
 
 _COMPILATION_TIMEOUT = flags.DEFINE_integer(
     'compilation_timeout', 60,
@@ -122,18 +121,6 @@ def __init__(self):
     Exception.__init__(self)
 
 
-class ProcessCancellationToken:
-
-  def __init__(self):
-    self._event = multiprocessing.Manager().Event()
-
-  def signal(self):
-    self._event.set()
-
-  def wait(self):
-    self._event.wait()
-
-
 def kill_process_ignore_exceptions(p: 'subprocess.Popen[bytes]'):
   # kill the process and ignore exceptions. Exceptions would be thrown if the
   # process has already been killed/finished (which is inherently in a race
@@ -160,6 +147,10 @@ def __init__(self):
     self._done = False
     self._lock = threading.Lock()
 
+  def enable(self):
+    with self._lock:
+      self._done = False
+
   def register_process(self, p: 'subprocess.Popen[bytes]'):
     """Register a process for potential cancellation."""
     with self._lock:
@@ -168,7 +159,7 @@ def register_process(self, p: 'subprocess.Popen[bytes]'):
         return
     kill_process_ignore_exceptions(p)
 
-  def signal(self):
+  def kill_all_processes(self):
     """Cancel any pending work."""
     with self._lock:
       self._done = True
@@ -265,21 +256,31 @@ def __post_init__(self, sequence_examples: List[tf.train.SequenceExample]):
     assert not hasattr(self, 'sequence_examples')
 
 
-class CompilationRunner:
-  """Base class for collecting compilation data."""
+class CompilationRunnerStub(metaclass=abc.ABCMeta):
+  """The interface of a stub to CompilationRunner, for type checkers."""
 
-  _POOL: concurrent.futures.ThreadPoolExecutor = None
+  @abc.abstractmethod
+  def collect_data(
+      self, file_paths: Tuple[str, ...], tf_policy_path: str,
+      reward_stat: Optional[Dict[str, RewardStat]]
+  ) -> WorkerFuture[CompilationResult]:
+    raise NotImplementedError()
 
-  @staticmethod
-  def init_pool():
-    """Worker process initialization."""
-    CompilationRunner._POOL = concurrent.futures.ThreadPoolExecutor()
+  @abc.abstractmethod
+  def cancel_all_work(self) -> WorkerFuture:
+    raise NotImplementedError()
 
-  @staticmethod
-  def _get_pool():
-    """Internal API for fetching the cancellation token waiting pool."""
-    assert CompilationRunner._POOL
-    return CompilationRunner._POOL
+  @abc.abstractmethod
+  def enable(self) -> WorkerFuture:
+    raise NotImplementedError()
+
+
+class CompilationRunner(Worker):
+  """Base class for collecting compilation data."""
+
+  @classmethod
+  def is_priority_method(cls, method_name: str) -> bool:
+    return method_name == 'cancel_all_work'
 
   def __init__(self,
                clang_path: Optional[str] = None,
@@ -302,40 +303,18 @@ def __init__(self,
     self._additional_flags = additional_flags
     self._delete_flags = delete_flags
     self._compilation_timeout = _COMPILATION_TIMEOUT.value
+    self._cancellation_manager = WorkerCancellationManager()
 
-  def _get_cancellation_manager(
-      self, cancellation_token: Optional[ProcessCancellationToken]
-  ) -> Optional[WorkerCancellationManager]:
-    """Convert the ProcessCancellationToken into a WorkerCancellationManager.
-
-    The conversion also registers the ProcessCancellationToken wait() on a
-    thread which will call the WorkerCancellationManager upon completion.
-    Since the token is always signaled, the thread always completes its work.
-
-    Args:
-      cancellation_token: the ProcessCancellationToken to convert.
-
-    Returns:
-      a WorkerCancellationManager, if a ProcessCancellationToken was given.
-    """
-    if not cancellation_token:
-      return None
-    ret = WorkerCancellationManager()
-
-    def signaler():
-      cancellation_token.wait()
-      ret.signal()
+  # re-allow the cancellation manager accept work.
+  def enable(self):
+    self._cancellation_manager.enable()
 
-    CompilationRunner._get_pool().submit(signaler)
-    return ret
+  def cancel_all_work(self):
+    self._cancellation_manager.kill_all_processes()
 
   def collect_data(
-      self,
-      file_paths: Tuple[str, ...],
-      tf_policy_path: str,
-      reward_stat: Optional[Dict[str, RewardStat]],
-      cancellation_token: Optional[ProcessCancellationToken] = None
-  ) -> CompilationResult:
+      self, file_paths: Tuple[str, ...], tf_policy_path: str,
+      reward_stat: Optional[Dict[str, RewardStat]]) -> CompilationResult:
     """Collect data for the given IR file and policy.
 
     Args:
@@ -355,14 +334,12 @@ def collect_data(
       compilation_runner.ProcessKilledException is passed through.
       ValueError if example under default policy and ml policy does not match.
     """
-    cancellation_manager = self._get_cancellation_manager(cancellation_token)
-
     if reward_stat is None:
       default_result = self._compile_fn(
           file_paths,
           tf_policy_path='',
           reward_only=bool(tf_policy_path),
-          cancellation_manager=cancellation_manager)
+          cancellation_manager=self._cancellation_manager)
       reward_stat = {
           k: RewardStat(v[1], v[1]) for (k, v) in default_result.items()
       }
@@ -372,7 +349,7 @@ def collect_data(
           file_paths,
           tf_policy_path,
           reward_only=False,
-          cancellation_manager=cancellation_manager)
+          cancellation_manager=self._cancellation_manager)
     else:
       policy_result = default_result