Make worker.get_full_worker_args more user friendly (#223)

We can just accept `**kwargs` - this allows usage like

  `worker.get_full_worker_args(some_type, arg1=val1, arg2=val2)`

Author: mtrofin

compiler_opt/distributed/local/local_worker_manager.py

@@ -244,7 +244,7 @@ def create_local_worker_pool(worker_cls: 'type[worker.Worker]',
   """Create a local worker pool for worker_cls."""
   if not count:
     count = multiprocessing.get_context().cpu_count()
-  final_kwargs = worker.get_full_worker_args(worker_cls, kwargs)
+  final_kwargs = worker.get_full_worker_args(worker_cls, **kwargs)
   stubs = [_make_stub(worker_cls, *args, **final_kwargs) for _ in range(count)]
   return worker.FixedWorkerPool(workers=stubs, worker_concurrency=16)
 

compiler_opt/distributed/worker.py

@@ -91,7 +91,7 @@ def get_exception(worker_future: WorkerFuture) -> Optional[Exception]:
     return e
 
 
-def get_full_worker_args(worker_class: 'type[Worker]', current_kwargs):
+def get_full_worker_args(worker_class: 'type[Worker]', **current_kwargs):
   """Get the union of given kwargs and gin config.
 
   This allows the worker hosting process be set up differently from the training

compiler_opt/distributed/worker_test.py

@@ -0,0 +1,42 @@
+# 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.
+"""Test for worker."""
+
+import gin
+
+from absl.testing import absltest
+from compiler_opt.distributed import worker
+
+
+@gin.configurable(module='_test')
+class SomeType:
+
+  def __init__(self, argument):
+    pass
+
+
+class WorkerTest(absltest.TestCase):
+
+  def test_gin_args(self):
+    with gin.unlock_config():
+      gin.bind_parameter('_test.SomeType.argument', 42)
+    real_args = worker.get_full_worker_args(
+        SomeType, more_args=2, even_more_args='hi')
+    self.assertDictEqual(real_args,
+                         dict(argument=42, more_args=2, even_more_args='hi'))
+
+
+if __name__ == '__main__':
+  absltest.main()