Separate worker and worker manager args (#453)

Without this, it's hard to extend the argument list for the local worker manager.

Author: mtrofin

compiler_opt/distributed/buffered_scheduler_test.py

@@ -34,7 +34,8 @@ def wait_seconds(self, n: int):
         time.sleep(n)
         return n + 1
 
-    with local_worker_manager.LocalWorkerPoolManager(WaitingWorker, 2) as pool:
+    with local_worker_manager.LocalWorkerPoolManager(
+        WaitingWorker, count=2) as pool:
       _, futures = buffered_scheduler.schedule_on_worker_pool(
           lambda w, v: w.wait_seconds(v), range(4), pool)
       not_done = futures
@@ -52,7 +53,8 @@ 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:
+    with local_worker_manager.LocalWorkerPoolManager(
+        TheWorker, count=2) as pool:
       workers, futures = buffered_scheduler.schedule_on_worker_pool(
           lambda w, v: w.square(v), range(10), pool)
       self.assertLen(workers, 2)

compiler_opt/distributed/local/local_worker_manager.py

@@ -253,9 +253,9 @@ def __dir__(self):
   return _Stub()
 
 
-def create_local_worker_pool(worker_cls: 'type[worker.Worker]',
-                             count: int | None, parse_argv: bool, *args,
-                             **kwargs) -> worker.FixedWorkerPool:
+def _create_local_worker_pool(worker_cls: 'type[worker.Worker]',
+                              count: int | None, parse_argv: bool, *args,
+                              **kwargs) -> worker.FixedWorkerPool:
   """Create a local worker pool for worker_cls."""
   if not count:
     count = _get_context().cpu_count()
@@ -267,7 +267,7 @@ def create_local_worker_pool(worker_cls: 'type[worker.Worker]',
   return worker.FixedWorkerPool(workers=stubs, worker_concurrency=16)
 
 
-def close_local_worker_pool(pool: worker.FixedWorkerPool):
+def _close_local_worker_pool(pool: worker.FixedWorkerPool):
   """Close the given LocalWorkerPool."""
   # first, trigger killing the worker process and exiting of the msg pump,
   # which will also clear out any pending futures.
@@ -281,16 +281,21 @@ def close_local_worker_pool(pool: worker.FixedWorkerPool):
 class LocalWorkerPoolManager(AbstractContextManager):
   """A pool of workers hosted on the local machines, each in its own process."""
 
-  def __init__(self, worker_class: 'type[worker.Worker]', count: int | None,
-               *args, **kwargs):
-    self._pool = create_local_worker_pool(worker_class, count, True, *args,
-                                          **kwargs)
+  def __init__(self,
+               worker_class: 'type[worker.Worker]',
+               *,
+               count: int | None,
+               worker_args: tuple = (),
+               worker_kwargs: dict | None = None):
+    worker_kwargs = {} if worker_kwargs is None else worker_kwargs
+    self._pool = _create_local_worker_pool(worker_class, count, True,
+                                           *worker_args, **worker_kwargs)
 
   def __enter__(self) -> worker.FixedWorkerPool:
     return self._pool
 
   def __exit__(self, *args):
-    close_local_worker_pool(self._pool)
+    _close_local_worker_pool(self._pool)
 
   def __del__(self):
     self.__exit__()

compiler_opt/distributed/local/local_worker_manager_test.py

@@ -85,7 +85,8 @@ def test_pool(self):
     kwarg = 'bar'
 
     with local_worker_manager.LocalWorkerPoolManager(
-        JobNormal, 2, arg, kwarg=kwarg) as pool:
+        JobNormal, count=2, worker_args=(arg,),
+        worker_kwargs=dict(kwarg=kwarg)) as pool:
       p1 = pool.get_currently_active()[0]
       p2 = pool.get_currently_active()[1]
       set_futures = [p1.set_token(1), p2.set_token(2)]
@@ -108,15 +109,15 @@ def test_pool(self):
 
   def test_failure(self):
 
-    with local_worker_manager.LocalWorkerPoolManager(JobFail, 2) as pool:
+    with local_worker_manager.LocalWorkerPoolManager(JobFail, count=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.get_currently_active()[0].method().result()
 
   def test_worker_crash_while_waiting(self):
 
-    with local_worker_manager.LocalWorkerPoolManager(JobSlow, 2) as pool:
+    with local_worker_manager.LocalWorkerPoolManager(JobSlow, count=2) as pool:
       p = pool.get_currently_active()[0]
       f = p.method()
       self.assertFalse(f.done())
@@ -127,12 +128,14 @@ def test_worker_crash_while_waiting(self):
           _ = f.result()
 
   def test_flag_parsing(self):
-    with local_worker_manager.LocalWorkerPoolManager(JobGetFlags, 1) as pool:
+    with local_worker_manager.LocalWorkerPoolManager(
+        JobGetFlags, count=1) as pool:
       result = pool.get_currently_active()[0].method().result()
       self.assertEqual(result['the_flag'], 1)
 
     with mock.patch('sys.argv', sys.argv + ['--test_only_flag=42']):
-      with local_worker_manager.LocalWorkerPoolManager(JobGetFlags, 1) as pool:
+      with local_worker_manager.LocalWorkerPoolManager(
+          JobGetFlags, count=1) as pool:
         result = pool.get_currently_active()[0].method().result()
         self.assertEqual(result['the_flag'], 42)
 

compiler_opt/es/blackbox_evaluator_test.py

@@ -28,7 +28,10 @@ class BlackboxEvaluatorTests(absltest.TestCase):
 
   def test_sampling_get_results(self):
     with local_worker_manager.LocalWorkerPoolManager(
-        blackbox_test_utils.ESWorker, count=3, arg='', kwarg='') as pool:
+        blackbox_test_utils.ESWorker,
+        count=3,
+        worker_args=('',),
+        worker_kwargs=dict(kwarg='')) as pool:
       perturbations = [b'00', b'01', b'10']
       evaluator = blackbox_evaluator.SamplingBlackboxEvaluator(None, 5, 5, None)
       # pylint: disable=protected-access
@@ -52,7 +55,10 @@ def test_get_rewards(self):
 
   def test_trace_get_results(self):
     with local_worker_manager.LocalWorkerPoolManager(
-        blackbox_test_utils.ESTraceWorker, count=3, arg='', kwarg='') as pool:
+        blackbox_test_utils.ESTraceWorker,
+        count=3,
+        worker_args=('',),
+        worker_kwargs=dict(kwarg='')) as pool:
       perturbations = [b'00', b'01', b'10']
       test_corpus = corpus.create_corpus_for_testing(
           location=self.create_tempdir().full_path,
@@ -65,7 +71,10 @@ def test_trace_get_results(self):
 
   def test_trace_set_baseline(self):
     with local_worker_manager.LocalWorkerPoolManager(
-        blackbox_test_utils.ESTraceWorker, count=1, arg='', kwarg='') as pool:
+        blackbox_test_utils.ESTraceWorker,
+        count=1,
+        worker_args=('',),
+        worker_kwargs=dict(kwarg='')) as pool:
       test_corpus = corpus.create_corpus_for_testing(
           location=self.create_tempdir().full_path,
           elements=[corpus.ModuleSpec(name='name1', size=1)])

compiler_opt/es/blackbox_learner_test.py

@@ -146,7 +146,10 @@ def test_prune_skipped_perturbations(self):
 
   def test_run_step(self):
     with local_worker_manager.LocalWorkerPoolManager(
-        blackbox_test_utils.ESWorker, count=3, arg='', kwarg='') as pool:
+        blackbox_test_utils.ESWorker,
+        count=3,
+        worker_args=('',),
+        worker_kwargs=dict(kwarg='')) as pool:
       self._learner.run_step(pool)  # pylint: disable=protected-access
       # expected length calculated from expected shapes of variables
       self.assertEqual(len(self._learner.get_model_weights()), 17154)

compiler_opt/rl/imitation_learning/generate_bc_trajectories_lib.py

@@ -955,15 +955,15 @@ def gen_trajectories(
       min(os.cpu_count(), num_workers) if num_workers else os.cpu_count())
   with worker_manager_class(
       worker_class_type,
-      worker_count,
-      obs_action_specs=obs_action_spec,
-      mlgo_task_type=mlgo_task_type,
-      callable_policies=callable_policies,
-      explore_on_features=explore_on_features,
-      persistent_objects_path=persistent_objects_path,
-      explicit_temps_dir=explicit_temps_dir,
-      gin_config_str=gin.config_str(),
-  ) as lwm:
+      count=worker_count,
+      worker_kwargs=dict(
+          obs_action_specs=obs_action_spec,
+          mlgo_task_type=mlgo_task_type,
+          callable_policies=callable_policies,
+          explore_on_features=explore_on_features,
+          persistent_objects_path=persistent_objects_path,
+          explicit_temps_dir=explicit_temps_dir,
+          gin_config_str=gin.config_str())) as lwm:
 
     _, result_futures = buffered_scheduler.schedule_on_worker_pool(
         action=lambda w, j: w.select_best_exploration(loaded_module_spec=j),

compiler_opt/tools/generate_default_trace.py

@@ -151,12 +151,13 @@ def generate_trace(
     with performance_context as performance_writer:
       with worker_manager_class(
           FilteringWorker,
-          _NUM_WORKERS.value,
-          policy_path=_POLICY_PATH.value,
-          key_filter=_KEY_FILTER.value,
-          runner_type=runner_type,
-          runner_kwargs=worker.get_full_worker_args(
-              runner_type, moving_average_decay_rate=0)) as lwm:
+          count=_NUM_WORKERS.value,
+          worker_kwargs=dict(
+              policy_path=_POLICY_PATH.value,
+              key_filter=_KEY_FILTER.value,
+              runner_type=runner_type,
+              runner_kwargs=worker.get_full_worker_args(
+                  runner_type, moving_average_decay_rate=0))) as lwm:
 
         _, result_futures = buffered_scheduler.schedule_on_worker_pool(
             action=lambda w, j: w.compile_and_filter(j),