[Bug Fix] Wrapping args of LocalWorkerPoolManager with worker_kwargs (#466)

Co-authored-by: svkeerthy <[email protected]>

Author: svkeerthy

compiler_opt/rl/distributed/ppo_collect_lib.py

@@ -167,8 +167,9 @@ def sequence_example_iterator_fn(seq_ex: list[str]):
   with worker_manager_class(
       worker_class=problem_config.get_runner_type(),
       count=num_workers,
-      moving_average_decay_rate=1,
-      create_observer_fns=create_observer_fns) as worker_pool:
+      worker_kwargs=dict(
+          moving_average_decay_rate=1,
+          create_observer_fns=create_observer_fns)) as worker_pool:
 
     data_collector = local_data_collector.LocalDataCollector(
         cps=cps,

compiler_opt/rl/distributed/ppo_eval_lib.py

@@ -105,7 +105,7 @@ def sequence_example_iterator_fn(seq_ex: list[str]):
   with worker_manager_class(
       worker_class=problem_config.get_runner_type(),
       count=num_workers,
-      moving_average_decay_rate=1) as worker_pool:
+      worker_kwargs=dict(moving_average_decay_rate=1)) as worker_pool:
     logging.info('constructed pool')
     collector = local_data_collector.LocalDataCollector(
         cps=cps,

compiler_opt/rl/train_locally.py

@@ -146,7 +146,8 @@ def sequence_example_iterator_fn(seq_ex: list[str]):
   with worker_manager_class(
       worker_class=problem_config.get_runner_type(),
       count=FLAGS.num_workers,
-      moving_average_decay_rate=moving_average_decay_rate) as worker_pool:
+      worker_kwargs=dict(
+          moving_average_decay_rate=moving_average_decay_rate)) as worker_pool:
     data_collector = local_data_collector.LocalDataCollector(
         cps=cps,
         num_modules=num_modules,