LocalWorker: Fix nr threads to 1 (#81)

Author: mtrofin

compiler_opt/distributed/local/local_worker_manager.py

@@ -35,7 +35,6 @@
 import queue  # pylint: disable=unused-import
 import threading
 
-import gin
 from absl import logging
 # pylint: disable=unused-import
 from compiler_opt.distributed.worker import Worker
@@ -61,7 +60,7 @@ class TaskResult:
 
 
 def _run_impl(in_q: 'queue.Queue[Task]', out_q: 'queue.Queue[TaskResult]',
-              worker_class: 'type[Worker]', threads: int, *args, **kwargs):
+              worker_class: 'type[Worker]', *args, **kwargs):
   """Worker process entrypoint."""
   # Note: the out_q is typed as taking only TaskResult objects, not
   # Optional[TaskResult], despite that being the type it is used on the Stub
@@ -73,7 +72,7 @@ def _run_impl(in_q: 'queue.Queue[Task]', out_q: 'queue.Queue[TaskResult]',
   # process near-immediately. `threads` only controls how many threads are
   # spawned at a time which execute given tasks. In the typical clang-spawning
   # jobs, this effectively limits the number of clang instances spawned.
-  pool = concurrent.futures.ThreadPoolExecutor(max_workers=threads)
+  pool = concurrent.futures.ThreadPoolExecutor(max_workers=1)
   obj = worker_class(*args, **kwargs)
 
   def make_ondone(msgid):
@@ -109,7 +108,7 @@ def _run(*args, **kwargs):
     raise e
 
 
-def _make_stub(cls: 'type[Worker]', pool_threads: int, *args, **kwargs):
+def _make_stub(cls: 'type[Worker]', *args, **kwargs):
 
   class _Stub():
     """Client stub to a worker hosted by a process."""
@@ -120,13 +119,15 @@ def __init__(self):
         multiprocessing.get_context().Queue()
 
       # this is the process hosting one worker instance.
+      # we set aside 1 thread to coordinate running jobs, and the main thread
+      # to handle high priority requests. The expectation is that the user
+      # achieves concurrency through multiprocessing, not multithreading.
       self._process = multiprocessing.Process(
           target=functools.partial(
               _run,
               worker_class=cls,
               in_q=self._send,
               out_q=self._receive,
-              threads=pool_threads,
               *args,
               **kwargs))
       # lock for the msgid -> reply future map. The map will be set to None
@@ -217,25 +218,16 @@ def __dir__(self):
   return _Stub()
 
 
-@gin.configurable
 class LocalWorkerPool(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,
-               pool_threads: int = 1,
+  def __init__(self, worker_class: 'type[Worker]', count: Optional[int], *args,
                **kwargs):
     if not count:
       count = multiprocessing.cpu_count()
     self._stubs = [
-        _make_stub(worker_class, pool_threads, *args, **kwargs)
-        for _ in range(count // pool_threads)
+        _make_stub(worker_class, *args, **kwargs) for _ in range(count)
     ]
-    # Make sure there's always `count` worker threads, not a rounded `count`
-    if (remainder := count % pool_threads) != 0:
-      self._stubs.append(_make_stub(worker_class, remainder, *args, **kwargs))
 
   def __enter__(self):
     return self._stubs