|
import random |
|
|
|
|
|
class NaiveResourceManager(object): |
|
r""" |
|
Overview: |
|
the naive resource manager |
|
Interface: |
|
__init__, assign_collector, assign_learner, update |
|
""" |
|
|
|
def __init__(self) -> None: |
|
r""" |
|
Overview: |
|
init the resouce manager |
|
""" |
|
self._worker_type = ['collector', 'learner'] |
|
self._resource_info = {k: {} for k in self._worker_type} |
|
|
|
def assign_collector(self, collector_task: dict) -> dict: |
|
r""" |
|
Overview: |
|
assign the collector_task randomly and return the resouce info |
|
Arguments: |
|
- collector_task (:obj:`dict`): the collector task to assign |
|
""" |
|
available_collector_list = list(self._resource_info['collector'].keys()) |
|
if len(available_collector_list) > 0: |
|
selected_collector = random.sample(available_collector_list, 1)[0] |
|
info = self._resource_info['collector'].pop(selected_collector) |
|
return {'collector_id': selected_collector, 'resource_info': info} |
|
else: |
|
return None |
|
|
|
def assign_learner(self, learner_task: dict) -> dict: |
|
r""" |
|
Overview: |
|
assign the learner_task randomly and return the resouce info |
|
Arguments: |
|
- learner_task (:obj:`dict`): the learner task to assign |
|
""" |
|
available_learner_list = list(self._resource_info['learner'].keys()) |
|
if len(available_learner_list) > 0: |
|
selected_learner = random.sample(available_learner_list, 1)[0] |
|
info = self._resource_info['learner'].pop(selected_learner) |
|
return {'learner_id': selected_learner, 'resource_info': info} |
|
else: |
|
return None |
|
|
|
def have_assigned(self, name: id, worker_id: str) -> bool: |
|
assert name in self._worker_type, "invalid worker_type: {}".format(name) |
|
if name == 'collector': |
|
return worker_id in self._resource_info['collector'] |
|
elif name == 'learner': |
|
return worker_id in self._resource_info['learner'] |
|
|
|
def delete(self, name: id, worker_id: str) -> bool: |
|
assert name in self._worker_type, "invalid worker_type: {}".format(name) |
|
if worker_id in self._resource_info[name]: |
|
self._resource_info.pop(worker_id) |
|
return True |
|
else: |
|
return False |
|
|
|
def update(self, name: str, worker_id: str, resource_info: dict) -> None: |
|
r""" |
|
Overview: |
|
update the reource info |
|
""" |
|
assert name in self._worker_type, "invalid worker_type: {}".format(name) |
|
self._resource_info[name][worker_id] = resource_info |
|
|