def _create_session(self, payload):
"""
Creates a new session and optimizer based on the parameters sent by a
model developer. It then asks the optimizer what's the first thing to do
and the service starts working on that session from that point on.
"""
# NOTE: We removed the 'optimizer_type' argument since for the MVP we're
# only considering the 'FederatedAveragingOptimizer' for now.
# TODO: We need to incorporate session id's when we're ready.
if not self.dataset_manager:
raise Exception("Dataset Manager has not been set. Communication \
Manager needs to be configured first!")
logging.info("New optimizer session is being set up...")
initialization_payload = payload.get(TxEnum.CONTENT.name)
self.optimizer = FederatedAveragingOptimizer(
initialization_payload,
self.dataset_manager
)
logging.info("Optimizer session is set! Now doing kickoff...")
event_type, payload = self.optimizer.kickoff()
callback = callback_handler_no_default(
event_type,
self.EVENT_TYPE_2_CALLBACK,
)
callback(payload)
logging.info("Kickoff complete.")
def test_optimizer_fails_on_wrong_event_type(initialization_payload):
optimizer = FederatedAveragingOptimizer(initialization_payload)
try:
dummy_payload = {}
_, _ = optimizer.ask("WRONG_EVENT_TYPE", dummy_payload)
raise Exception("Optimizer should have failed!")
except Exception as e:
assert str(e) == "Invalid callback passed!"
def test_optimizer_can_kickoff(runner, initialization_payload,
dataset_manager):
optimizer = FederatedAveragingOptimizer(initialization_payload,
dataset_manager)
event_type, job_arr = optimizer.kickoff()
assert event_type == ActionableEventTypes.SCHEDULE_JOBS.name
for job in job_arr:
assert job.job_type == JobTypes.JOB_INIT.name or job.job_type == JobTypes.JOB_SPLIT.name
def test_optimizer_schedules_communication_after_training(initialization_payload, dataset_manager, train_dmlresult_obj):
optimizer = FederatedAveragingOptimizer(
initialization_payload,
dataset_manager
)
trained_weights = train_dmlresult_obj.results['weights']
train_dmlresult_obj.job = train_dmlresult_obj.job.copy_constructor()
event_type, job_arr = optimizer.ask(RawEventTypes.JOB_DONE.name, train_dmlresult_obj)
assert all(np.allclose(arr1, arr2) for arr1,arr2 in zip(optimizer.job.weights, trained_weights))
assert event_type == ActionableEventTypes.SCHEDULE_JOBS.name
for job in job_arr:
assert job.job_type == JobTypes.JOB_COMM.name
def test_optimizer_schedules_training_after_initialization(initialization_payload, dataset_manager, init_dmlresult_obj, split_dmlresult_obj):
optimizer = FederatedAveragingOptimizer(
initialization_payload,
dataset_manager
)
print(optimizer.job.raw_filepath)
initial_weights = init_dmlresult_obj.results['weights']
init_dmlresult_obj.job = init_dmlresult_obj.job.copy_constructor()
split_dmlresult_obj.job = split_dmlresult_obj.job.copy_constructor()
event_type, job_arr = optimizer.ask(RawEventTypes.JOB_DONE.name, init_dmlresult_obj)
assert all(np.allclose(arr1, arr2) for arr1,arr2 in zip(optimizer.job.weights, initial_weights))
assert event_type == ActionableEventTypes.NOTHING.name
event_type, job_arr = optimizer.ask(RawEventTypes.JOB_DONE.name, split_dmlresult_obj)
assert event_type == ActionableEventTypes.SCHEDULE_JOBS.name
for job in job_arr:
assert job.job_type == JobTypes.JOB_TRAIN.name
def bootstrap(repo_id="testRepo", api_key="demo-api-key", test=False):
"""
Bootstraps the data provider unix service.
It instantiates the Configuration Manager, Dataset Manager, Communication
Manager and the Execution Pipeline.
"""
# 1. Set up Configuration Manager.
config_manager = ConfigurationManager()
config_manager.bootstrap()
runner = DMLRunner(config_manager)
optimizer = FederatedAveragingOptimizer(runner, repo_id)
loop = asyncio.get_event_loop()
websocket_client = WebSocketClient(optimizer, config_manager, repo_id,
api_key, test)
# mappings = dataset_manager.get_mappings()
# 7. Wait for the threads to end.
# TODO: Need to make it work as a daemon.
loop.run_until_complete(websocket_client.prepare_dml())
class CommunicationManager(object):
"""
Communication Manager
Manages the communication within the modules of the service throughout all
active DML Sessions.
Right now, the service can only have one session at a time. This will be
reflected in the code by just having one optimizer property and no
information about sessions throughout the code.
"""
def __init__(self):
"""
Initializes the class.
"""
logging.info("Setting up Communication Manager...")
self.optimizer = None # NOTE: Since right now we're only dealing with
# one session at a time, the class is only
# concerned with one optimizer.
self.scheduler = None
# NOTE: This should be updated when Gateway PR is merged and we make
# the Communication Manager error-handle out of order/missed messages.
self.EVENT_TYPE_2_CALLBACK = {
MessageEventTypes.NEW_SESSION.name: self._create_session,
ActionableEventTypes.SCHEDULE_JOBS.name: self._schedule_jobs,
ActionableEventTypes.TERMINATE.name: self._terminate_session,
ActionableEventTypes.NOTHING.name: self._do_nothing,
}
logging.info("Communication Manager is set up!")
# Setup Methods
def configure(self, scheduler):
"""
Configures the Communication Manager so that it can submit jobs to the
Scheduler.
"""
logging.info("Configuring the Communication Manager...")
self.scheduler = scheduler
logging.info("Communication Manager is configured!")
# Public methods
def inform(self, event_type, payload):
"""
Method called by other modules to inform the Communication Manager about
events that are going on in the service.
If there isn't an active optimizer, the Communication Manager will handle it
Else, payloads are relayed to the active optimizer which decides how
and to whom the Communication Manager should communicate the event/message.
For example: A runner could inform the Communication Manager that the
node is done training a particular model, to which the Optimizer could
decide it's time to communicate the new weights to the network.
"""
logging.info("Information has been received: {}".format(event_type))
if self.optimizer:
# We have an active session so we ask the optimizer what to do.
event_type, payload = self.optimizer.ask(event_type, payload)
callback = callback_handler_no_default(
event_type,
self.EVENT_TYPE_2_CALLBACK
)
callback(payload)
# Callbacks
def _create_session(self, payload):
"""
Creates a new session and optimizer based on the parameters sent by a
model developer. It then asks the optimizer what's the first thing to do
and the service starts working on that session from that point on.
"""
# NOTE: We removed the 'optimizer_type' argument since for the MVP we're
# only considering the 'FederatedAveragingOptimizer' for now.
# TODO: We need to incorporate session id's when we're ready.
logging.info("New optimizer session is being set up...")
initialization_payload = payload.get(TxEnum.CONTENT.name)
self.optimizer = FederatedAveragingOptimizer(initialization_payload)
logging.info("Optimizer session is set! Now doing kickoff...")
event_type, payload = self.optimizer.kickoff()
callback = callback_handler_no_default(
event_type,
self.EVENT_TYPE_2_CALLBACK,
)
callback(payload)
logging.info("Kickoff complete.")
def _schedule_jobs(self, dmljob_obj_arr):
"""
Helper function that schedules a DML Job through the Scheduler.
"""
if not self.scheduler:
raise Exception("Communication Manager needs to be configured first!")
for i, dmljob_obj in enumerate(dmljob_obj_arr):
assert isinstance(dmljob_obj, DMLJob), \
"Obj {} is not DMLJob, is {}".format(i, type(dmljob_obj))
self.scheduler.add_job(dmljob_obj)
def _terminate_session(self, payload):
"""
Helper function that removes the current session (and optimizer) from
the Communication Manager.
Note that it ignores the payload for now (since we're only dealing with
one session).
"""
self.optimizer = None
logging.info("Session terminated.")
def _do_nothing(self, payload):
"""
Do nothing
"""
logging.info("Doing nothing")