Have the worker pool context produce a pool object (#154)

The motivation is:
- allowing worker managers specify the level of concurrency for a
  worker. This is set to '10' right now by the data collector, but it's
  really specific to the worker set up. For the local case, where we max
  out the hardware threads with workers, 10 is overly-generous. For a
  distributed case, a number approaching the hardware threads would be
  more appropriate
- allowing distributed worker managers update the set of available
  workers over time. Workers could get preempted, or new ones become
  available. Having a way to periodically check and update what's
  available - albeit there are never guarantees a worker, once
  discovered, stays alive - helps avoid artificial starvation.

The rest is renames that fall out of this refactoring.

Author: mtrofin

compiler_opt/distributed/local/local_worker_manager.py

@@ -36,7 +36,7 @@
 
 from absl import logging
 # pylint: disable=unused-import
-from compiler_opt.distributed.worker import Worker
+from compiler_opt.distributed.worker import Worker, FixedWorkerPool
 
 from contextlib import AbstractContextManager
 from multiprocessing import connection
@@ -238,7 +238,7 @@ def __dir__(self):
   return _Stub()
 
 
-class LocalWorkerPool(AbstractContextManager):
+class LocalWorkerPoolManager(AbstractContextManager):
   """A pool of workers hosted on the local machines, each in its own process."""
 
   def __init__(self, worker_class: 'type[Worker]', count: Optional[int], *args,
@@ -249,8 +249,8 @@ def __init__(self, worker_class: 'type[Worker]', count: Optional[int], *args,
         _make_stub(worker_class, *args, **kwargs) for _ in range(count)
     ]
 
-  def __enter__(self):
-    return self._stubs
+  def __enter__(self) -> FixedWorkerPool:
+    return FixedWorkerPool(workers=self._stubs, worker_concurrency=10)
 
   def __exit__(self, *args):
     # first, trigger killing the worker process and exiting of the msg pump,

compiler_opt/distributed/local/local_worker_manager_test.py

@@ -62,9 +62,9 @@ class LocalWorkerManagerTest(absltest.TestCase):
 
   def test_pool(self):
 
-    with local_worker_manager.LocalWorkerPool(JobNormal, 2) as pool:
-      p1 = pool[0]
-      p2 = pool[1]
+    with local_worker_manager.LocalWorkerPoolManager(JobNormal, 2) as pool:
+      p1 = pool.get_currently_active()[0]
+      p2 = pool.get_currently_active()[1]
       set_futures = [p1.set_token(1), p2.set_token(2)]
       done, not_done = concurrent.futures.wait(set_futures)
       self.assertLen(done, 2)
@@ -81,16 +81,16 @@ def test_pool(self):
 
   def test_failure(self):
 
-    with local_worker_manager.LocalWorkerPool(JobFail, 2) as pool:
+    with local_worker_manager.LocalWorkerPoolManager(JobFail, 2) as pool:
       with self.assertRaises(concurrent.futures.CancelledError):
         # this will fail because we didn't pass the arg to the ctor, so the
         # worker hosting process will crash.
-        pool[0].method().result()
+        pool.get_currently_active()[0].method().result()
 
   def test_worker_crash_while_waiting(self):
 
-    with local_worker_manager.LocalWorkerPool(JobSlow, 2) as pool:
-      p = pool[0]
+    with local_worker_manager.LocalWorkerPoolManager(JobSlow, 2) as pool:
+      p = pool.get_currently_active()[0]
       f = p.method()
       self.assertFalse(f.done())
       try:

compiler_opt/distributed/worker.py

@@ -14,7 +14,8 @@
 # limitations under the License.
 """Common abstraction for a worker contract."""
 
-from typing import Iterable, Optional, Protocol, TypeVar
+import abc
+from typing import Any, List, Iterable, Optional, Protocol, TypeVar
 
 
 class Worker(Protocol):
@@ -28,6 +29,34 @@ def is_priority_method(cls, method_name: str) -> bool:
 T = TypeVar('T')
 
 
+class WorkerPool(metaclass=abc.ABCMeta):
+  """Abstraction of a pool of workers that may be refreshed."""
+
+  # Issue #155 would strongly-type the return type.
+  @abc.abstractmethod
+  def get_currently_active(self) -> List[Any]:
+    raise NotImplementedError()
+
+  @abc.abstractmethod
+  def get_worker_concurrency(self) -> int:
+    raise NotImplementedError()
+
+
+class FixedWorkerPool(WorkerPool):
+  """A WorkerPool built from a fixed list of workers."""
+
+  # Issue #155 would strongly-type `workers`
+  def __init__(self, workers: List[Any], worker_concurrency: int = 2):
+    self._workers = workers
+    self._worker_concurrency = worker_concurrency
+
+  def get_currently_active(self):
+    return self._workers
+
+  def get_worker_concurrency(self):
+    return self._worker_concurrency
+
+
 # Dask's Futures are limited. This captures that.
 class WorkerFuture(Protocol[T]):
 

compiler_opt/rl/local_data_collector.py

@@ -37,7 +37,7 @@ def __init__(
       self,
       cps: corpus.Corpus,
       num_modules: int,
-      worker_pool: List[compilation_runner.CompilationRunnerStub],
+      worker_pool: worker.WorkerPool,
       parser: Callable[[List[str]], Iterator[trajectory.Trajectory]],
       reward_stat_map: Dict[str, Optional[Dict[str,
                                                compilation_runner.RewardStat]]],
@@ -49,6 +49,9 @@ def __init__(
     self._num_modules = num_modules
     self._parser = parser
     self._worker_pool = worker_pool
+    self._workers: List[
+        compilation_runner
+        .CompilationRunnerStub] = self._worker_pool.get_currently_active()
     self._reward_stat_map = reward_stat_map
     self._exit_checker_ctor = exit_checker_ctor
     # _reset_workers is a future that resolves when post-data collection cleanup
@@ -75,8 +78,11 @@ def _prefetch_next_sample(self):
 
   def close_pool(self):
     self._join_pending_jobs()
-    for p in self._worker_pool:
+    # if the pool lost some workers, that's fine - we don't need to tell them
+    # anything anymore. To the new ones, the call is redudant (fine).
+    for p in self._workers:
       p.cancel_all_work()
+    self._workers = None
     self._worker_pool = None
 
   def _join_pending_jobs(self):
@@ -110,7 +116,9 @@ def work(w: compilation_runner.CompilationRunnerStub):
       return work
 
     work = [work_factory(job) for job in jobs]
-    return buffered_scheduler.schedule(work, self._worker_pool, buffer=10)
+    self._workers = self._worker_pool.get_currently_active()
+    return buffered_scheduler.schedule(
+        work, self._workers, self._worker_pool.get_worker_concurrency())
 
   def collect_data(
       self, policy: policy_saver.Policy
@@ -158,13 +166,13 @@ def get_num_finished_work():
 
     # signal whatever work is left to finish, and re-enable workers.
     def wrapup():
-      cancel_futures = [wkr.cancel_all_work() for wkr in self._worker_pool]
+      cancel_futures = [wkr.cancel_all_work() for wkr in self._workers]
       worker.wait_for(cancel_futures)
       # now that the workers killed pending compilations, make sure the workers
       # drained their working queues first - they should all complete quickly
       # since the cancellation manager is killing immediately any process starts
       worker.wait_for(self._current_futures)
-      worker.wait_for([wkr.enable() for wkr in self._worker_pool])
+      worker.wait_for([wkr.enable() for wkr in self._workers])
 
     self._reset_workers = self._pool.submit(wrapup)
 

compiler_opt/rl/local_data_collector_test.py

@@ -23,7 +23,7 @@
 import tensorflow as tf
 from tf_agents.system import system_multiprocessing as multiprocessing
 
-from compiler_opt.distributed.local.local_worker_manager import LocalWorkerPool
+from compiler_opt.distributed.local.local_worker_manager import LocalWorkerPoolManager
 from compiler_opt.rl import compilation_runner
 from compiler_opt.rl import corpus
 from compiler_opt.rl import data_collector
@@ -142,7 +142,7 @@ def _test_iterator_fn(data_list):
       return _test_iterator_fn
 
     sampler = DeterministicSampler()
-    with LocalWorkerPool(worker_class=MyRunner, count=4) as lwp:
+    with LocalWorkerPoolManager(worker_class=MyRunner, count=4) as lwp:
       collector = local_data_collector.LocalDataCollector(
           cps=corpus.create_corpus_for_testing(
               location=self.create_tempdir(),
@@ -214,7 +214,7 @@ def __init__(self, num_modules):
       def wait(self, _):
         return False
 
-    with LocalWorkerPool(worker_class=Sleeper, count=4) as lwp:
+    with LocalWorkerPoolManager(worker_class=Sleeper, count=4) as lwp:
       collector = local_data_collector.LocalDataCollector(
           cps=corpus.create_corpus_for_testing(
               location=self.create_tempdir(),

compiler_opt/rl/train_locally.py

@@ -29,7 +29,7 @@
 from tf_agents.system import system_multiprocessing as multiprocessing
 from typing import List
 
-from compiler_opt.distributed.local.local_worker_manager import LocalWorkerPool
+from compiler_opt.distributed.local.local_worker_manager import LocalWorkerPoolManager
 from compiler_opt.rl import agent_creators
 from compiler_opt.rl import compilation_runner
 from compiler_opt.rl import constant
@@ -59,7 +59,7 @@
 
 
 @gin.configurable
-def train_eval(worker_manager_class=LocalWorkerPool,
+def train_eval(worker_manager_class=LocalWorkerPoolManager,
                agent_name=constant.AgentName.PPO,
                warmstart_policy_dir=None,
                num_policy_iterations=0,