Simplify scheduling work on a worker pool (#219)

The user can just specify how to call a worker with some work, instead
of having to go through the trouble of creating a work factory and so
on.

Author: mtrofin

compiler_opt/distributed/buffered_scheduler.py

@@ -19,11 +19,12 @@
 import concurrent.futures
 import threading
 
-from typing import List, Callable, TypeVar
+from typing import Any, Callable, Iterable, List, Optional, Tuple, TypeVar
 
 from compiler_opt.distributed import worker
 
 T = TypeVar('T')
+W = TypeVar('W')
 
 
 def schedule(work: List[Callable[[T], worker.WorkerFuture]],
@@ -80,3 +81,38 @@ def chain_work(wkr: T):
       chain_work(w)
 
   return results
+
+
+def schedule_on_worker_pool(
+    action: Callable[[W, T], Any],
+    jobs: Iterable[T],
+    worker_pool: worker.WorkerPool,
+    buffer_size: Optional[int] = None
+) -> Tuple[List[W], List[worker.WorkerFuture]]:
+  """
+  Schedule the given action on workers from the given worker pool.
+  Args:
+    action: a function that, given a worker and some args, calls that worker
+      with those args.
+    jobs: a list of arguments, each element constituting a unit of work.
+    worker_pool: the worker pool on which to schedule the work.
+    buffer_size: if provided, buffer these many work items, instead of the
+      worker manager's default.
+
+  Returns:
+    a tuple. The first value is the workers that are used to perform the work.
+      The second is a list of futures, one for each work item.
+  """
+
+  def work_factory(args):
+
+    def work(w: worker.Worker):
+      return action(w, args)
+
+    return work
+
+  work = [work_factory(job) for job in jobs]
+  workers: List[W] = worker_pool.get_currently_active()
+  return workers, schedule(work, workers,
+                           (worker_pool.get_worker_concurrency()
+                            if buffer_size is None else buffer_size))

compiler_opt/distributed/buffered_scheduler_test.py

@@ -21,10 +21,42 @@
 from absl.testing import absltest
 from compiler_opt.distributed import worker
 from compiler_opt.distributed import buffered_scheduler
+from compiler_opt.distributed.local import local_worker_manager
 
 
 class BufferedSchedulerTest(absltest.TestCase):
 
+  def test_simple_scheduling(self):
+
+    class TheWorker(worker.Worker):
+
+      def square(self, the_value, extra_factor=1):
+        return the_value * the_value * extra_factor
+
+    with local_worker_manager.LocalWorkerPoolManager(TheWorker, 2) as pool:
+      workers, futures = buffered_scheduler.schedule_on_worker_pool(
+          lambda w, v: w.square(v), range(10), pool)
+      self.assertLen(workers, 2)
+      self.assertLen(futures, 10)
+      worker.wait_for(futures)
+      self.assertListEqual([f.result() for f in futures],
+                           [x * x for x in range(10)])
+
+      _, futures = buffered_scheduler.schedule_on_worker_pool(
+          lambda w, v: w.square(**v), [dict(the_value=v) for v in range(10)],
+          pool)
+      worker.wait_for(futures)
+      self.assertListEqual([f.result() for f in futures],
+                           [x * x for x in range(10)])
+
+      # same idea, but mix some kwargs
+      _, futures = buffered_scheduler.schedule_on_worker_pool(
+          lambda w, v: w.square(v[0], **v[1]),
+          [(v, dict(extra_factor=10)) for v in range(10)], pool)
+      worker.wait_for(futures)
+      self.assertListEqual([f.result() for f in futures],
+                           [x * x * 10 for x in range(10)])
+
   def test_schedules(self):
     call_count = [0] * 4
     locks = [threading.Lock() for _ in range(4)]

compiler_opt/rl/local_data_collector.py

@@ -100,28 +100,24 @@ def _join_pending_jobs(self):
     logging.info('Waiting for pending work from last iteration took %f',
                  time.time() - t1)
 
-  def _schedule_jobs(
-      self, policy: policy_saver.Policy, model_id: int,
-      sampled_modules: List[corpus.LoadedModuleSpec]
-  ) -> List[worker.WorkerFuture[compilation_runner.CompilationResult]]:
+  def _schedule_jobs(self, policy: policy_saver.Policy, model_id: int,
+                     sampled_modules: List[corpus.LoadedModuleSpec]) -> None:
     # by now, all the pending work, which was signaled to cancel, must've
     # finished
     self._join_pending_jobs()
-    jobs = [(loaded_module_spec, policy,
-             self._reward_stat_map[loaded_module_spec.name])
-            for loaded_module_spec in sampled_modules]
-
-    def work_factory(job):
-
-      def work(w: compilation_runner.CompilationRunnerStub):
-        return w.collect_data(*job, model_id=model_id)
-
-      return work
+    jobs = [
+        dict(
+            loaded_module_spec=loaded_module_spec,
+            policy=policy,
+            reward_stat=self._reward_stat_map[loaded_module_spec.name],
+            model_id=model_id) for loaded_module_spec in sampled_modules
+    ]
 
-    work = [work_factory(job) for job in jobs]
-    self._workers = self._worker_pool.get_currently_active()
-    return buffered_scheduler.schedule(
-        work, self._workers, self._worker_pool.get_worker_concurrency())
+    (self._workers,
+     self._current_futures) = buffered_scheduler.schedule_on_worker_pool(
+         action=lambda w, kwargs: w.collect_data(**kwargs),
+         jobs=jobs,
+         worker_pool=self._worker_pool)
 
   def collect_data(
       self, policy: policy_saver.Policy, model_id: int
@@ -145,8 +141,7 @@ def collect_data(
     logging.info('resolving prefetched sample took: %d seconds',
                  time.time() - time1)
     self._next_sample = self._prefetch_next_sample()
-    self._current_futures = self._schedule_jobs(policy, model_id,
-                                                sampled_modules)
+    self._schedule_jobs(policy, model_id, sampled_modules)
 
     def wait_for_termination():
       early_exit = self._exit_checker_ctor(num_modules=self._num_modules)