Add pause/resume/context to workers

- Allows a user to start/stop processes at will, via OS signals SIGSTOP and SIGCONT.
- Allows a user to bind processes to specific CPUs.
- Allows local_worker_pool to be used outside of a context manager
- Switch workers to be Protocol based, so Workers are effectively duck-typed (i.e. anything that has the required methods passes as a Worker)

Part of google#96

Author: Northbadge

compiler_opt/distributed/local/local_worker_manager.py

@@ -32,14 +32,17 @@
 import functools
 import multiprocessing
 import threading
+import os
+import psutil
+import signal
 
 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
+from typing import Any, Callable, Dict, Optional, List
 
 
 @dataclasses.dataclass(frozen=True)
@@ -131,6 +134,7 @@ 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)
@@ -210,10 +214,37 @@ 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.
+      """
+      psutil.Process(self._process.pid).nice(val)
+
+    def set_affinity(self, val: List[int]):
+      """Sets the CPU affinity of the process, this modifies which cores the OS
+      schedules it on.
+      """
+      psutil.Process(self._process.pid).cpu_affinity(val)
+
     def join(self):
       self._observer.join()
       self._pump.join()
@@ -247,3 +278,11 @@ def __exit__(self, *args):
     # now wait for the message pumps to indicate they exit.
     for s in self._stubs:
       s.join()
+
+  def __del__(self):
+    self.__exit__()
+
+  @property
+  def stubs(self):
+    # Return a shallow copy, to avoid something messing the internal list up
+    return list(self._stubs)

compiler_opt/distributed/worker.py

@@ -14,17 +14,28 @@
 # limitations under the License.
 """Common abstraction for a worker contract."""
 
-from typing import Iterable, Optional, TypeVar, Protocol
+from typing import Iterable, Optional, TypeVar, Protocol, runtime_checkable
 
 
-class Worker:
+class Worker(Protocol):
 
   @classmethod
   def is_priority_method(cls, method_name: str) -> bool:
     _ = method_name
     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')
 
 

requirements.txt

@@ -27,6 +27,7 @@ oauthlib==3.1.1
 opt-einsum==3.3.0
 pillow==8.3.1
 protobuf==3.17.3
+psutil==5.9.0
 pyasn1==0.4.8
 pyasn1_modules==0.2.8
 pyglet==1.5.0