Change the return type of buffered_scheduler APIs (#221)

They return canonical futures (i.e. `concurrent.futures.Future`). This is great because we can then use any `Future` APIs, e.g. `concurrent.futures.wait`.

Author: mtrofin

compiler_opt/distributed/buffered_scheduler.py

@@ -29,7 +29,7 @@
 
 def schedule(work: List[Callable[[T], worker.WorkerFuture]],
              workers: List[T],
-             buffer=2) -> List[worker.WorkerFuture]:
+             buffer=2) -> List[concurrent.futures.Future]:
   """
   Assigns work to workers once previous work of the worker are
   completed.
@@ -88,7 +88,7 @@ def schedule_on_worker_pool(
     jobs: Iterable[T],
     worker_pool: worker.WorkerPool,
     buffer_size: Optional[int] = None
-) -> Tuple[List[W], List[worker.WorkerFuture]]:
+) -> Tuple[List[W], List[concurrent.futures.Future]]:
   """
   Schedule the given action on workers from the given worker pool.
   Args:

compiler_opt/distributed/buffered_scheduler_test.py

@@ -26,6 +26,25 @@
 
 class BufferedSchedulerTest(absltest.TestCase):
 
+  def test_canonical_futures(self):
+
+    class WaitingWorker(worker.Worker):
+
+      def wait_seconds(self, n: int):
+        time.sleep(n)
+        return n + 1
+
+    with local_worker_manager.LocalWorkerPoolManager(WaitingWorker, 2) as pool:
+      _, futures = buffered_scheduler.schedule_on_worker_pool(
+          lambda w, v: w.wait_seconds(v), range(4), pool)
+      not_done = futures
+      entered_count = 0
+      while not_done:
+        _, not_done = concurrent.futures.wait(not_done, timeout=0.5)
+        entered_count += 1
+      self.assertGreater(entered_count, 1)
+      self.assertListEqual([r.result() for r in futures], list(range(1, 5)))
+
   def test_simple_scheduling(self):
 
     class TheWorker(worker.Worker):

compiler_opt/rl/local_data_collector.py

@@ -170,7 +170,7 @@ def wrapup():
       # 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)
+      concurrent.futures.wait(self._current_futures)
       worker.wait_for([wkr.enable() for wkr in self._workers])
 
     self._reset_workers = self._pool.submit(wrapup)