Introduce a Nesting Worker In some distributed object middleware, there's a limitation that one remote process host one object instance only. During training, we have 'lull' periods when we don't collect data for an experiment (because we run the training algo on the data). If we run a bunch of related experiments, we can reuse the remote resource during that time. This can be advantageous if the total resource utilization of 2N remote tasks used at (say) 50% is higher than that of N tasks used at 100% (which is usually the case - there's some book keeping cost) The design uses a special kind of worker which manages a bunch of underlying workers to which it delegates all calls. The current patch just introduces the bare minimum support. Subsequent patches will add better error handing (e.g. worker id doesn't exist), and potentially some load balancing support, if needed.
diff --git a/compiler_opt/distributed/nesting_worker.py b/compiler_opt/distributed/nesting_worker.py new file mode 100644 index 0000000..9d5b8bf --- /dev/null +++ b/compiler_opt/distributed/nesting_worker.py
@@ -0,0 +1,97 @@ +# 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. +"""A worker hosting multiple workers.""" + +from typing import Optional + +import threading +from compiler_opt.distributed import worker +from contextlib import AbstractContextManager + +WorkerID = int + + +class NestingWorker(worker.Worker): + """A worker which hosts a number of other worker objects.""" + + def __init__(self): + self._lock = threading.Lock() + self._workers = {} + self._next_id = 0 + + def create(self, cls: 'type[worker.Worker]', *args, **kwargs) -> WorkerID: + worker_instance = cls(*args, **kwargs) + with self._lock: + worker_id = self._next_id + self._workers[worker_id] = worker_instance + self._next_id += 1 + return worker_id + + def release(self, worker_id: WorkerID): + with self._lock: + self._workers.pop(worker_id) + + def call(self, worker_id: WorkerID, method: str, *args, **kwargs): + with self._lock: + worker_instance = self._workers[worker_id] + return getattr(worker_instance, method)(*args, **kwargs) + + def _get_registered_workers(self): + return self._workers + + +def create_nested_worker_manager(underlying_pool: worker.WorkerPool): + """Create a worker manager class on an underlying pool of NestingWorkers.""" + class _Stub: + """Stub to worker hosted by a NestingWorker""" + def __init__(self, nesting_worker: NestingWorker, worker_id: WorkerID): + self._nesting_worker = nesting_worker + self._id = worker_id + + def __getattr__(self, method: str): + + def func(*args, **kwargs): + return self._nesting_worker.call(self._id, method, *args, **kwargs) + + return func + + def release(self): + self._nesting_worker.release(self._id) + + class _Nester(AbstractContextManager): + """The worker manager class.""" + def __init__(self, worker_class: 'type[worker.Worker]', + count: Optional[int], *args, **kwargs): + self._underlying_pool = underlying_pool + self._pool = [] + current_workers = self._underlying_pool.get_currently_active() + for i in range(count): + nesting_worker = current_workers[i % len(current_workers)] + self._pool.append( + _Stub( + nesting_worker=nesting_worker, + worker_id=nesting_worker.create(worker_class, *args, + **kwargs).result())) + + def __enter__(self): + return worker.FixedWorkerPool( + workers=self._pool, + worker_concurrency=self._underlying_pool.get_worker_concurrency()) + + def __exit__(self, *args): + for w in self._pool: + w.release() + + return _Nester
diff --git a/compiler_opt/distributed/nesting_worker_test.py b/compiler_opt/distributed/nesting_worker_test.py new file mode 100644 index 0000000..ddeefc4 --- /dev/null +++ b/compiler_opt/distributed/nesting_worker_test.py
@@ -0,0 +1,49 @@ +# 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 buffered_scheduler.""" + +# pylint: disable=protected-access +from absl.testing import absltest +from compiler_opt.distributed import nesting_worker +from compiler_opt.distributed import worker +from compiler_opt.distributed.local import local_worker_manager + + +class ToyWorker(worker.Worker): + + def say_hi(self): + return 'hi!' + + +class NestingWorkerTest(absltest.TestCase): + + def test_setup(self): + with local_worker_manager.LocalWorkerPoolManager( + worker_class=nesting_worker.NestingWorker, count=2) as lwpm: + nesting_manager = nesting_worker.create_nested_worker_manager(lwpm) + lwpm_active = lwpm.get_currently_active() + self.assertLen(lwpm_active, 2) + with nesting_manager(worker_class=ToyWorker, count=10) as mgr: + self.assertLen(mgr.get_currently_active(), 10) + for nw in lwpm_active: + self.assertLen(nw._get_registered_workers().result(), 5) + for wkr in mgr.get_currently_active(): + self.assertEqual(wkr.say_hi().result(), 'hi!') + for nw in lwpm_active: + self.assertLen(nw._get_registered_workers().result(), 0) + + +if __name__ == '__main__': + absltest.main()