diff --git a/compiler_opt/__init__.py b/compiler_opt/__init__.py
index 6078d4f..7c35352 100644
--- a/compiler_opt/__init__.py
+++ b/compiler_opt/__init__.py
@@ -15,6 +15,8 @@
 """Ensure flags are initialized for e.g. pytest harness case."""
 
 import sys
+#import compiler_opt.core.dask.worker_manager as default_factory
+import compiler_opt.core.multiprocessing.worker_manager as default_factory
 
 from absl import flags
 
@@ -28,3 +30,5 @@
 if not flags.FLAGS.is_parsed():
   flags.FLAGS(sys.argv, known_only=True)
 assert flags.FLAGS.is_parsed()
+
+worker_generator = default_factory.get_compilation_jobs
diff --git a/compiler_opt/core/abstract_worker.py b/compiler_opt/core/abstract_worker.py
new file mode 100644
index 0000000..a1ed868
--- /dev/null
+++ b/compiler_opt/core/abstract_worker.py
@@ -0,0 +1,24 @@
+# coding=utf-8
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Common abstraction for a worker contract."""
+
+from abc import abstractmethod
+
+
+class AbstractWorker:
+
+  @abstractmethod
+  def cancel_all_work(self):
+    ...
diff --git a/compiler_opt/core/dask/worker_manager.py b/compiler_opt/core/dask/worker_manager.py
new file mode 100644
index 0000000..d4f5f47
--- /dev/null
+++ b/compiler_opt/core/dask/worker_manager.py
@@ -0,0 +1,88 @@
+# coding=utf-8
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Dask - based middleware implementation."""
+from absl import logging
+import gin
+import dask.config
+import dask.utils
+import multiprocessing
+import tempfile
+
+from compiler_opt.core.abstract_worker import AbstractWorker
+from concurrent.futures import ThreadPoolExecutor
+from dask.distributed import Client, Worker, LocalCluster
+from typing import Callable, Optional, Tuple
+
+
+class MTWorker(Worker):
+  """Multi-threaded worker."""
+
+  def __init__(self, *args, **kwargs):
+    super().__init__(*args, **kwargs)
+    if self.executors['actor']:
+      self.executors['actor'].shutdown()
+    self.executors['actor'] = ThreadPoolExecutor(
+        max_workers=None, thread_name_prefix='Dask-Actor-MT')
+
+
+class LocalManager:
+  """Local, dask-based worker manager."""
+
+  def __init__(self):
+    self._tmpdir = tempfile.TemporaryDirectory()
+    dask.config.set({
+        'temporary-directory': self._tmpdir.name,
+        'distributed.worker.daemon': False,
+        'work-stealing': False
+    })
+
+    self._client = Client(
+        dashboard_address=None,
+        processes=True,
+        n_workers=1,
+        worker_class=MTWorker)
+    print(self._client)
+
+  def shutdown(self):
+    self._client.close()
+    self._tmpdir.cleanup()
+
+  def get_client(self):
+    return self._client
+
+
+def get_local_compilation_jobs(ctor: Callable[[], AbstractWorker],
+                               count: Optional[int]) -> Tuple[Callable, list]:
+
+  class DaskStubWrapper(AbstractWorker):
+
+    def __init__(self, stub: AbstractWorker):
+      self._stub = stub
+
+    def cancel_all_work(self):
+      return self._stub.cancel_all_work(separate_thread=False)
+
+    def __getattr__(self, name):
+      return self._stub.__getattr__(name)
+
+  if not count:
+    count = multiprocessing.cpu_count()
+  instance = LocalManager()
+  workers = [
+      instance.get_client().submit(ctor, actor=True) for _ in range(count)
+  ]
+  return instance.shutdown, [
+      DaskStubWrapper(worker.result()) for worker in workers
+  ]
diff --git a/compiler_opt/core/multiprocessing/worker_manager.py b/compiler_opt/core/multiprocessing/worker_manager.py
new file mode 100644
index 0000000..46fd911
--- /dev/null
+++ b/compiler_opt/core/multiprocessing/worker_manager.py
@@ -0,0 +1,111 @@
+# coding=utf-8
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Dask - based middleware implementation."""
+from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
+import concurrent.futures
+import functools
+import multiprocessing
+import multiprocessing.connection
+import queue
+import threading
+
+from compiler_opt.core.abstract_worker import AbstractWorker
+from typing import Callable, Optional, Tuple
+
+def worker(ctor, in_q:queue.Queue, out_q:queue.Queue):
+  pool = ThreadPoolExecutor()
+  obj = ctor()
+
+  def make_ondone(msgid):
+
+    def on_done(f: concurrent.futures.Future):
+      if f.exception():
+        out_q.put((msgid, False, f.exception()))
+      else:
+        out_q.put((msgid, True, f.result()))
+    return on_done
+
+  while True:
+    msgid, fname, args, kwargs, urgent = in_q.get()
+    the_func = getattr(obj, fname)
+    if urgent:
+      try:
+        res = the_func(*args, **kwargs)
+        out_q.put((msgid, True, res))
+      except Exception as e:
+        out_q.put((msgid, False, e))
+    else:
+      pool.submit(the_func, *args,
+                  **kwargs).add_done_callback(make_ondone(msgid))
+
+
+class Stub:
+  def __init__(self, ctor):
+    self._send = multiprocessing.get_context().Queue()
+    self._receive = multiprocessing.get_context().Queue()
+
+    self._process = multiprocessing.Process(
+        target=functools.partial(worker, ctor, self._send, self._receive))
+    self._lock = threading.Lock()
+    self._map:dict[object, concurrent.futures.Future] = {}
+    self._pump = threading.Thread(target=self._msg_pump)
+    self._done = threading.Event()
+    self._msgidlock = threading.Lock()
+    self._msgid = 0
+    self._process.start()
+    self._pump.start()
+
+  def _msg_pump(self):
+    while not self._done.is_set():
+      msgid, succeeded, value = self._receive.get()
+      with self._lock:
+        future = self._map[msgid]
+        del self._map[msgid]
+        if succeeded:
+          future.set_result(value)
+        else:
+          future.set_exception(value)
+
+  def __getattr__(self, name):
+    with self._msgidlock:
+      msgid = self._msgid
+      self._msgid += 1
+    def remote_call(*args, **kwargs):
+      self._send.put(
+          (msgid, name, args, kwargs, name == 'cancel_all_work'))
+      future = concurrent.futures.Future()
+      with self._lock:
+        self._map[msgid] = future
+      return future
+
+    return remote_call
+  
+  def kill(self):
+    try:
+      self._process.kill()
+    except:
+      pass
+    self._done.set()
+
+def get_compilation_jobs(ctor: Callable[[], AbstractWorker],
+                         count: Optional[int]) -> Tuple[Callable, list]:
+
+  if not count:
+    count = multiprocessing.cpu_count()
+  stubs = [Stub(ctor) for _ in range(count)]
+  def shutdown():
+    for s in stubs:
+      s.kill()
+  return shutdown, [Stub(ctor) for _ in range(count)]
diff --git a/compiler_opt/rl/compilation_runner.py b/compiler_opt/rl/compilation_runner.py
index 464d908..4722d1c 100644
--- a/compiler_opt/rl/compilation_runner.py
+++ b/compiler_opt/rl/compilation_runner.py
@@ -14,10 +14,10 @@
 # limitations under the License.
 """Module for running compilation and collect training data."""
 
-import concurrent
+from abc import abstractclassmethod, abstractmethod
+import concurrent.futures
 import dataclasses
 import json
-import multiprocessing
 import subprocess
 import threading
 from typing import Dict, List, Optional, Tuple
@@ -119,18 +119,6 @@
     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
@@ -157,6 +145,9 @@
     self._done = False
     self._lock = threading.Lock()
 
+  def enable(self):
+    self._done = False
+
   def register_process(self, p: 'subprocess.Popen[bytes]'):
     """Register a process for potential cancellation."""
     with self._lock:
@@ -165,12 +156,13 @@
         return
     kill_process_ignore_exceptions(p)
 
-  def signal(self):
+  def kill_all_processes(self):
     """Cancel any pending work."""
     with self._lock:
       self._done = True
     for p in self._processes:
       kill_process_ignore_exceptions(p)
+    return len(self._processes)
 
   def unregister_process(self, p: 'subprocess.Popen[bytes]'):
     with self._lock:
@@ -262,22 +254,24 @@
     assert not hasattr(self, 'sequence_examples')
 
 
+class CompilationRunnerStub:
+  """The interface of a stub to CompilationRunner."""
+
+  @abstractmethod
+  def collect_data(
+      self, file_paths: Tuple[str, ...], tf_policy_path: str,
+      reward_stat: Optional[Dict[str, RewardStat]]
+  ) -> concurrent.futures.Future[CompilationResult]:
+    ...
+
+  @abstractmethod
+  def cancel_all_work(self) -> concurrent.futures.Future:
+    ...
+
+
 class CompilationRunner:
   """Base class for collecting compilation data."""
 
-  _POOL: concurrent.futures.ThreadPoolExecutor = None
-
-  @staticmethod
-  def init_pool():
-    """Worker process initialization."""
-    CompilationRunner._POOL = concurrent.futures.ThreadPoolExecutor()
-
-  @staticmethod
-  def _get_pool():
-    """Internal API for fetching the cancellation token waiting pool."""
-    assert CompilationRunner._POOL
-    return CompilationRunner._POOL
-
   def __init__(self,
                clang_path: Optional[str] = None,
                launcher_path: Optional[str] = None,
@@ -293,40 +287,17 @@
     self._launcher_path = launcher_path
     self._moving_average_decay_rate = moving_average_decay_rate
     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.
+  def enable(self):
+    self._cancellation_manager.enable()
 
-    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()
-
-    CompilationRunner._get_pool().submit(signaler)
-    return ret
+  def cancel_all_work(self):
+    return 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:
@@ -346,14 +317,12 @@
       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()
       }
@@ -363,7 +332,7 @@
           file_paths,
           tf_policy_path,
           reward_only=False,
-          cancellation_manager=cancellation_manager)
+          cancellation_manager=self._cancellation_manager)
     else:
       policy_result = default_result
 
diff --git a/compiler_opt/rl/inlining/__init__.py b/compiler_opt/rl/inlining/__init__.py
index 6487ac8..bce5988 100644
--- a/compiler_opt/rl/inlining/__init__.py
+++ b/compiler_opt/rl/inlining/__init__.py
@@ -28,6 +28,9 @@
   def get_runner(self, *args, **kwargs):
     return inlining_runner.InliningRunner(*args, **kwargs)
 
+  def get_runner_ctor(self):
+    return inlining_runner.InliningRunner
+
   def get_signature_spec(self):
     return config.get_inlining_signature_spec()
 
diff --git a/compiler_opt/rl/local_data_collector.py b/compiler_opt/rl/local_data_collector.py
index 3b80bbc..fea34e4 100644
--- a/compiler_opt/rl/local_data_collector.py
+++ b/compiler_opt/rl/local_data_collector.py
@@ -14,19 +14,31 @@
 # limitations under the License.
 """Module for collecting data locally."""
 
+import concurrent.futures
 import itertools
 import random
 import time
-from typing import Callable, Dict, Iterator, List, Tuple, Optional
+from typing import Callable, Dict, Iterator, Iterable, List, Tuple, Optional
 
 from absl import logging
-import multiprocessing  # for Pool
+import gin
 from tf_agents.trajectories import trajectory
 
+from compiler_opt import worker_generator
 from compiler_opt.rl import compilation_runner
 from compiler_opt.rl import data_collector
 
 
+def wait_for(futures:Iterable[concurrent.futures.Future]):
+  """Dask futures don't support more than result()."""
+  for f in futures:
+    try:
+      _ = f.result()
+    except:
+      pass
+
+
+@gin.configurable()
 class LocalDataCollector(data_collector.DataCollector):
   """class for local data collection."""
 
@@ -35,71 +47,67 @@
       file_paths: Tuple[Tuple[str, ...], ...],
       num_workers: int,
       num_modules: int,
-      runner: compilation_runner.CompilationRunner,
       parser: Callable[[List[str]], Iterator[trajectory.Trajectory]],
       reward_stat_map: Dict[str, Optional[Dict[str,
                                                compilation_runner.RewardStat]]],
+      worker_ctor: Callable[[], compilation_runner.CompilationRunnerStub],
       exit_checker_ctor=data_collector.EarlyExitChecker):
     # TODO(mtrofin): type exit_checker_ctor when we get typing.Protocol support
     super().__init__()
 
     self._file_paths = file_paths
     self._num_modules = num_modules
-    self._runner = runner
     self._parser = parser
-    self._pool = multiprocessing.get_context().Pool(
-        num_workers, initializer=compilation_runner.CompilationRunner.init_pool)
+    self._shutdown, self._worker_pool = worker_generator(
+        worker_ctor, count=num_workers)
     self._reward_stat_map = reward_stat_map
 
     self._exit_checker_ctor = exit_checker_ctor
-    self._pending_work = None
-    # hold on to the token so it won't get GCed before all its wait()
-    # complete
-    self._last_token = None
+    self._pending_work: concurrent.futures.Future = None
+    self._current_work: list[Tuple[str, concurrent.futures.Future]] = []
+    self._pool = concurrent.futures.ThreadPoolExecutor()
+
+  def get_last_work(self):
+    return self._current_work
 
   def close_pool(self):
-    self._join_pending_jobs()
-    if self._pool:
-      # Stop accepting new work
-      self._pool.close()
-      self._pool.join()
-      self._pool = None
+    self.join_pending_jobs()
+    # Stop accepting new work
+    for p in self._worker_pool:
+      p.cancel_all_work()
+    self._worker_pool = None
+    self._shutdown()
+    self._pool.shutdown()
 
-  def _join_pending_jobs(self):
+  def join_pending_jobs(self):
+    t1 = time.time()
     if self._pending_work:
-      t1 = time.time()
-      for w in self._pending_work:
-        w.wait()
+      concurrent.futures.wait([self._pending_work])
 
-      self._pending_work = None
-      # this should have taken negligible time, normally, since all the work
-      # has been cancelled and the workers had time to process the cancellation
-      # while training was unfolding.
-      logging.info('Waiting for pending work from last iteration took %f',
-                   time.time() - t1)
-    self._last_token = None
+    self._pending_work = None
+    self._current_work = []
+    # this should have taken negligible time, normally, since all the work
+    # has been cancelled and the workers had time to process the cancellation
+    # while training was unfolding.
+    logging.info('Waiting for pending work from last iteration took %f',
+                  time.time() - t1)
 
-  def _schedule_jobs(self, policy_path, sampled_file_paths):
+  def _schedule_jobs(
+      self, policy_path, sampled_file_paths
+  ) -> List[concurrent.futures.Future[compilation_runner.CompilationResult]]:
     # by now, all the pending work, which was signaled to cancel, must've
     # finished
-    self._join_pending_jobs()
-    cancellation_token = compilation_runner.ProcessCancellationToken()
+    self.join_pending_jobs()
     jobs = [(file_paths, policy_path,
-             self._reward_stat_map['-'.join(file_paths)], cancellation_token)
+             self._reward_stat_map['-'.join(file_paths)])
             for file_paths in sampled_file_paths]
 
-    # Make sure we're not missing failures in workers. All but
-    # ProcessKilledError, which we want to ignore.
-    def error_callback(e):
-      if isinstance(e, compilation_runner.ProcessKilledError):
-        return
-      logging.exception('Error in worker: %r', e)
-
-    return (cancellation_token, [
-        self._pool.apply_async(
-            self._runner.collect_data, job, error_callback=error_callback)
-        for job in jobs
-    ])
+    # Naive load balancing
+    ret = []
+    for i in range(len(jobs)):
+      ret.append(self._worker_pool[i % len(self._worker_pool)].collect_data(
+          *(jobs[i])))
+    return ret
 
   def collect_data(
       self, policy_path: str
@@ -117,39 +125,50 @@
       information is viewable in TensorBoard.
     """
     sampled_file_paths = random.sample(self._file_paths, k=self._num_modules)
-    ct, results = self._schedule_jobs(policy_path, sampled_file_paths)
+    results = self._schedule_jobs(policy_path, sampled_file_paths)
 
     def wait_for_termination():
       early_exit = self._exit_checker_ctor(num_modules=self._num_modules)
 
       def get_num_finished_work():
-        finished_work = sum(res.ready() for res in results)
+        finished_work = sum(res.done() for res in results)
         return finished_work
 
       return early_exit.wait(get_num_finished_work)
 
     wait_seconds = wait_for_termination()
-    # signal whatever work is left to finish
-    ct.signal()
-    current_work = zip(sampled_file_paths, results)
-    finished_work = [(paths, res) for paths, res in current_work if res.ready()]
-    successful_work = [
-        (paths, res) for paths, res in finished_work if res.successful()
+    self._current_work = list(zip(sampled_file_paths, results))
+    finished_work = [
+        (paths, res) for paths, res in self._current_work if res.done()
     ]
-    failures = len(finished_work) - len(successful_work)
 
+    def is_successful(f):
+      try:
+        _ = f.result()
+        return True
+      except:  # pylint: disable=bare-except
+        return False
+
+    successful_work = [(paths, res.result())
+                       for paths, res in finished_work
+                       if is_successful(res)]
+    failures = len(finished_work) - len(successful_work)
     logging.info(('%d of %d modules finished in %d seconds (%d failures).'),
                  len(finished_work), self._num_modules, wait_seconds, failures)
+    # signal whatever work is left to finish
+    def wrapup():
+      cancel_futures = [wkr.cancel_all_work() for wkr in self._worker_pool]
+      wait_for(cancel_futures)
+      wait_for(results)
+      wait_for([wkr.enable() for wkr in self._worker_pool])
+    self._pending_work = self._pool.submit(wrapup)
 
     sequence_examples = list(
-        itertools.chain.from_iterable([
-            res.get().serialized_sequence_examples
-            for (_, res) in successful_work
-        ]))
-    total_trajectory_length = sum(
-        res.get().length for (_, res) in successful_work)
+        itertools.chain.from_iterable(
+            [res.serialized_sequence_examples for (_, res) in successful_work]))
+    total_trajectory_length = sum(res.length for (_, res) in successful_work)
     self._reward_stat_map.update({
-        '-'.join(file_paths): res.get().reward_stats
+        '-'.join(file_paths): res.reward_stats
         for (file_paths, res) in successful_work
     })
 
@@ -160,19 +179,13 @@
     }
     rewards = list(
         itertools.chain.from_iterable(
-            [res.get().rewards for (_, res) in successful_work]))
+            [res.rewards for (_, res) in successful_work]))
     monitor_dict[
         'reward_distribution'] = data_collector.build_distribution_monitor(
             rewards)
 
     parsed = self._parser(sequence_examples)
 
-    self._pending_work = [res for res in results if not res.ready()]
-    # if some of the cancelled work hasn't yet processed the signal, let's let
-    # it do that while we process the data. We also need to hold on to the
-    # current token, so its Event object doesn't get GC-ed here.
-    if self._pending_work:
-      self._last_token = ct
     return parsed, monitor_dict
 
   def on_dataset_consumed(self,
diff --git a/compiler_opt/rl/local_data_collector_test.py b/compiler_opt/rl/local_data_collector_test.py
index 51a35b4..84f0e2e 100644
--- a/compiler_opt/rl/local_data_collector_test.py
+++ b/compiler_opt/rl/local_data_collector_test.py
@@ -44,7 +44,7 @@
   return text_format.Parse(sequence_example_text, tf.train.SequenceExample())
 
 
-def mock_collect_data(file_paths, tf_policy_dir, reward_stat, _):
+def mock_collect_data(file_paths, tf_policy_dir, reward_stat):
   assert file_paths == ('a', 'b')
   assert tf_policy_dir == 'policy'
   assert reward_stat is None or reward_stat == {
@@ -77,30 +77,10 @@
 class Sleeper(compilation_runner.CompilationRunner):
   """Test CompilationRunner that just sleeps."""
 
-  # pylint: disable=super-init-not-called
-  def __init__(self, killed, timedout, living):
-    self._killed = killed
-    self._timedout = timedout
-    self._living = living
-    self._lock = mp.Manager().Lock()
-
-  def collect_data(self, file_paths, tf_policy_path, reward_stat,
-                   cancellation_token):
-    _ = reward_stat
-    cancellation_manager = self._get_cancellation_manager(cancellation_token)
-    try:
-      compilation_runner.start_cancellable_process(['sleep', '3600s'], 3600,
-                                                   cancellation_manager)
-    except compilation_runner.ProcessKilledError as e:
-      with self._lock:
-        self._killed.value += 1
-      raise e
-    except subprocess.TimeoutExpired as e:
-      with self._lock:
-        self._timedout.value += 1
-      raise e
-    with self._lock:
-      self._living.value += 1
+  def collect_data(self, file_paths, tf_policy_path, reward_stat):
+    _ = file_paths, tf_policy_path, reward_stat
+    compilation_runner.start_cancellable_process(['sleep', '3600s'], 3600,
+                                                 self._cancellation_manager)
     return compilation_runner.CompilationResult(
         sequence_examples=[], reward_stats={}, rewards=[], keys=[])
 
@@ -108,10 +88,15 @@
 class LocalDataCollectorTest(tf.test.TestCase):
 
   def test_local_data_collector(self):
-    mock_compilation_runner = mock.create_autospec(
-        compilation_runner.CompilationRunner)
 
-    mock_compilation_runner.collect_data = mock_collect_data
+    def make_runner():
+
+      class MyRunner(compilation_runner.CompilationRunner):
+
+        def collect_data(self, *args, **kwargs):
+          return mock_collect_data(*args, **kwargs)
+
+      return MyRunner()
 
     def create_test_iterator_fn():
 
@@ -132,8 +117,8 @@
         file_paths=tuple([('a', 'b')] * 100),
         num_workers=4,
         num_modules=9,
-        runner=mock_compilation_runner,
         parser=create_test_iterator_fn(),
+        worker_ctor=make_runner,
         reward_stat_map=collections.defaultdict(lambda: None))
 
     data_iterator, monitor_dict = collector.collect_data(policy_path='policy')
@@ -161,11 +146,6 @@
     collector.close_pool()
 
   def test_local_data_collector_task_management(self):
-    killed = mp.Manager().Value('i', value=0)
-    timedout = mp.Manager().Value('i', value=0)
-    living = mp.Manager().Value('i', value=0)
-
-    mock_compilation_runner = Sleeper(killed, timedout, living)
 
     def parser(_):
       pass
@@ -182,17 +162,18 @@
         file_paths=tuple([('a', 'b')] * 200),
         num_workers=4,
         num_modules=4,
-        runner=mock_compilation_runner,
+        worker_ctor=Sleeper,
         parser=parser,
         reward_stat_map=collections.defaultdict(lambda: None),
         exit_checker_ctor=QuickExiter)
     collector.collect_data(policy_path='policy')
-    # close the pool first, so we are forced to wait for the workers to process
-    # their cancellation.
+    collector.join_pending_jobs()
+    killed = 0
+    for _, w in collector.get_last_work():
+      self.assertRaises(compilation_runner.ProcessKilledError, w.result)
+      killed += 1
+    self.assertEquals(killed, 4)
     collector.close_pool()
-    self.assertEqual(killed.value, 4)
-    self.assertEqual(living.value, 0)
-    self.assertEqual(timedout.value, 0)
 
 
 if __name__ == '__main__':
diff --git a/compiler_opt/rl/problem_configuration.py b/compiler_opt/rl/problem_configuration.py
index ff6d089..28f169d 100644
--- a/compiler_opt/rl/problem_configuration.py
+++ b/compiler_opt/rl/problem_configuration.py
@@ -94,7 +94,10 @@
   @abc.abstractmethod
   def get_runner(self, *args, **kwargs) -> compilation_runner.CompilationRunner:
     raise NotImplementedError
-
+  
+  @abc.abstractmethod
+  def get_runner_ctor(self):
+    raise NotImplementedError
 
 def is_thinlto(module_paths: Iterable[str]) -> bool:
   return tf.io.gfile.exists(next(iter(module_paths)) + '.thinlto.bc')
diff --git a/compiler_opt/rl/train_locally.py b/compiler_opt/rl/train_locally.py
index 5a80a9b..f4576c8 100644
--- a/compiler_opt/rl/train_locally.py
+++ b/compiler_opt/rl/train_locally.py
@@ -25,7 +25,6 @@
 from absl import logging
 import gin
 import tensorflow as tf
-from tf_agents.system import system_multiprocessing as multiprocessing
 from typing import List
 
 from compiler_opt.rl import agent_creators
@@ -111,9 +110,6 @@
       file_paths = [(path + '.bc', path + '.cmd', path + '.thinlto.bc')
                     for path in module_paths]
 
-  runner = problem_config.get_runner(
-      moving_average_decay_rate=moving_average_decay_rate)
-
   dataset_fn = data_reader.create_sequence_example_dataset_fn(
       agent_name=agent_name,
       time_step_spec=time_step_spec,
@@ -145,7 +141,9 @@
       file_paths=file_paths,
       num_workers=FLAGS.num_workers,
       num_modules=FLAGS.num_modules,
-      runner=runner,
+      worker_ctor=functools.partial(
+          problem_config.get_runner_ctor(),
+          moving_average_decay_rate=moving_average_decay_rate),
       parser=sequence_example_iterator_fn,
       reward_stat_map=reward_stat_map)
 
@@ -185,4 +183,4 @@
 
 if __name__ == '__main__':
   flags.mark_flag_as_required('data_path')
-  multiprocessing.handle_main(functools.partial(app.run, main))
+  app.run(main)
diff --git a/requirements.txt b/requirements.txt
index 45a37e8..c61d9d6 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -3,10 +3,14 @@
 cachetools==4.2.2
 certifi==2021.5.30
 charset_normalizer==2.0.4
+click==8.1.3
 cloudpickle==1.6.0
+dask[distributed]==2022.6.0
 decorator==5.0.9
+distributed==2022.6.0
 dm-tree==0.1.6
 flatbuffers==2.0
+fsspec==2022.5.0
 future==0.18.2
 gast==0.4.0
 gin==0.1.006
@@ -16,15 +20,21 @@
 google-pasta==0.2.0
 grpcio==1.39.0
 gym==0.19.0
+HeapDict==1.0.1
 h5py==3.6.0
 idna==3.2
+Jinja2==3.1.2
 keras==2.7.0
 keras-preprocessing==1.1.2
 libclang==12.0.0
+locket==1.0.0
 markdown==3.3.4
+MarkupSafe==2.1.1
+msgpack==1.0.4
 numpy==1.22.1
 oauthlib==3.1.1
 opt-einsum==3.3.0
+partd==1.2.0
 pillow==8.3.1
 protobuf==3.17.3
 pyasn1==0.4.8
@@ -36,6 +46,8 @@
 scipy==1.7.1
 setuptools==57.4.0
 six==1.16.0
+sortedcontainers==2.4.0
+tblib==1.7.0
 tensorboard==2.7.0
 tensorboard-data-server==0.6.1
 tensorboard-plugin-wit==1.8.0
@@ -45,10 +57,12 @@
 tensorflow-probability==0.15.0
 termcolor==1.1.0
 tf-agents==0.11.0
+toolz==0.11.2
+tornado==6.1
 typing-extensions==4.0.1
 urllib3==1.26.6
 werkzeug==2.0.1
 wheel==0.37.0
 wrapt==1.13.3
+zict==2.2.0
 zipp==3.7.0
-