Explicitly pass the worker's gin config (#166)

If the local worker manager's processes aren't forked (which depends on
what the default for `multiprocessing.get_context` is) the gin bindings
won't be available. So we just pass them explicitly.

Author: mtrofin

compiler_opt/distributed/local/local_worker_manager.py

@@ -36,7 +36,7 @@
 
 from absl import logging
 # pylint: disable=unused-import
-from compiler_opt.distributed.worker import Worker, FixedWorkerPool
+from compiler_opt.distributed import worker
 
 from contextlib import AbstractContextManager
 from multiprocessing import connection
@@ -59,8 +59,8 @@ class TaskResult:
   value: Any
 
 
-def _run_impl(pipe: connection.Connection, worker_class: 'type[Worker]', *args,
-              **kwargs):
+def _run_impl(pipe: connection.Connection, worker_class: 'type[worker.Worker]',
+              *args, **kwargs):
   """Worker process entrypoint."""
 
   # A setting of 1 does not inhibit the while loop below from running since
@@ -111,7 +111,7 @@ def _run(*args, **kwargs):
     raise e
 
 
-def _make_stub(cls: 'type[Worker]', *args, **kwargs):
+def _make_stub(cls: 'type[worker.Worker]', *args, **kwargs):
 
   class _Stub:
     """Client stub to a worker hosted by a process."""
@@ -241,16 +241,17 @@ def __dir__(self):
 class LocalWorkerPoolManager(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,
-               **kwargs):
+  def __init__(self, worker_class: 'type[worker.Worker]', count: Optional[int],
+               *args, **kwargs):
     if not count:
       count = multiprocessing.get_context().cpu_count()
+    final_kwargs = worker.get_full_worker_args(worker_class, kwargs)
     self._stubs = [
-        _make_stub(worker_class, *args, **kwargs) for _ in range(count)
+        _make_stub(worker_class, *args, **final_kwargs) for _ in range(count)
     ]
 
-  def __enter__(self) -> FixedWorkerPool:
-    return FixedWorkerPool(workers=self._stubs, worker_concurrency=10)
+  def __enter__(self) -> worker.FixedWorkerPool:
+    return worker.FixedWorkerPool(workers=self._stubs, worker_concurrency=10)
 
   def __exit__(self, *args):
     # first, trigger killing the worker process and exiting of the msg pump,

compiler_opt/distributed/worker.py

@@ -15,8 +15,11 @@
 """Common abstraction for a worker contract."""
 
 import abc
+import sys
 from typing import Any, List, Iterable, Optional, Protocol, TypeVar
 
+import gin
+
 
 class Worker(Protocol):
 
@@ -86,3 +89,23 @@ def get_exception(worker_future: WorkerFuture) -> Optional[Exception]:
     return None
   except Exception as e:  # pylint: disable=broad-except
     return e
+
+
+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
+  process - e.g. no need to initialize gin variables there, for example.
+  """
+  gin_config = {}
+  try:
+    gin_config = gin.get_bindings(worker_class)
+  except ValueError:
+    # we don't have a way to check if `worker_class` is even known to gin, and
+    # it's not a requirement that it were. Tests, for instance, don't use gin.
+    pass
+  # Issue #38
+  if sys.version_info.minor >= 9:
+    return current_kwargs | gin_config
+  else:
+    return {**current_kwargs, **gin_config}