def test_report_task_result(self):
self.master.task_d = _TaskDispatcher(
{
"shard_1": (0, 10),
"shard_2": (0, 9)
},
{},
{},
records_per_task=3,
num_epochs=2,
)
master = MasterServicer(3, evaluation_service=None, master=self.master)
# task to number of runs.
tasks = defaultdict(int)
while True:
req = elasticdl_pb2.GetTaskRequest()
req.worker_id = random.randint(1, 10)
task = master.get_task(req, None)
if not task.shard_name:
break
self.assertEqual(self.master.task_d._doing[task.task_id][0],
req.worker_id)
task_key = (task.shard_name, task.start, task.end)
tasks[task_key] += 1
report = elasticdl_pb2.ReportTaskResultRequest()
report.task_id = task.task_id
if task.start == 0 and tasks[task_key] == 1:
# Simulate error reports.
report.err_message = "Worker error"
master.report_task_result(report, None)
self.assertDictEqual(
{
("shard_1", 0, 3): 3,
("shard_1", 3, 6): 2,
("shard_1", 6, 9): 2,
("shard_1", 9, 10): 2,
("shard_2", 0, 3): 3,
("shard_2", 3, 6): 2,
("shard_2", 6, 9): 2,
},
tasks,
)
def get_task(self, task_type=None):
"""
get task from master
"""
req = elasticdl_pb2.GetTaskRequest()
req.worker_id = self._worker_id
if task_type is not None:
req.task_type = task_type
try:
res = self._stub.get_task(req)
except Exception:
# Master may have stopped GRPC service when there are no more
# tasks. This will result in a GRPC call exception.
self.logger.info(
"Cannot connect to master, assuming no more tasks"
)
res = elasticdl_pb2.Task()
return res
def testGetEmptyTask(self):
master = MasterServicer(
3,
_TaskDispatcher({}, {}, {}, records_per_task=3, num_epochs=2),
evaluation_service=None,
)
req = elasticdl_pb2.GetTaskRequest()
# No task yet, make sure the returned versions are as expected.
req.worker_id = 1
task = master.get_task(req, None)
self.assertEqual("", task.shard_name)
self.assertEqual(0, task.model_version)
master._version = 1
task = master.get_task(req, None)
self.assertEqual("", task.shard_name)
self.assertEqual(1, task.model_version)
def testGetEmptyTask(self):
master = MasterServicer(
2,
3,
None,
_TaskDispatcher({}, {}, {}, records_per_task=3, num_epochs=2),
init_var=[],
checkpoint_filename_for_init="",
checkpoint_service=CheckpointService("", 0, 0, False),
evaluation_service=None,
)
req = elasticdl_pb2.GetTaskRequest()
# No task yet, make sure the returned versions are as expected.
req.worker_id = 1
task = master.GetTask(req, None)
self.assertEqual("", task.shard_file_name)
self.assertEqual(0, task.model_version)
master._version = 1
task = master.GetTask(req, None)
self.assertEqual("", task.shard_file_name)
self.assertEqual(1, task.model_version)
def get_task(self, task_type=None):
"""Get a task from master.
Args:
task_type: elasticdl_pb.TaskType
the training phase, c.f. /elasticdl/proto/elasticdl.proto
Returns:
the task unit assigned by master,
c.f. /elasticdl/proto/elasticdl.proto
"""
req = elasticdl_pb2.GetTaskRequest()
req.worker_id = self._worker_id
if task_type is not None:
req.task_type = task_type
try:
res = self._stub.get_task(req)
except Exception:
# the master node would stop the gRPC service if no more tasks.
# And this will result a gRPC call exception.
res = elasticdl_pb2.Task()
return res
def distributed_train_and_evaluate(
feature_shape,
model_zoo_path,
model_def,
model_params="",
eval_metrics_fn="eval_metrics_fn",
loss="loss",
training=True,
dataset_name=DatasetName.IMAGE_DEFAULT,
callback_classes=[],
use_async=False,
get_model_steps=1,
ps_channels=None,
pservers=None,
distribution_strategy=DistributionStrategy.PARAMETER_SERVER,
):
"""Runs distributed training and evaluation with a local master. Grpc
calls are mocked by local master call.
Args:
feature_shape: The shape of model input.
model_zoo_path: The directory that contains user-defined model files
or a specific model file.
model_def: The import path to the model definition function/class in
the model zoo, e.g. "cifar10_subclass.CustomModel".
model_params: The dictionary of model parameters in a string that will
be used to instantiate the model, e.g. "param1=1,param2=2".
eval_metrics_fn: The name of the evaluation metrics function defined
in the model file.
loss: The name of the loss function defined in the model file.
training: True for job type `TRAIN_WITH_EVALUATION`, False for
job type `EVALUATION`.
dataset_name: A dataset name from `DatasetName`.
callback_classes: A List of callbacks that will be called at given
stages of the training procedure.
use_async: A bool. True if using asynchronous updates.
get_model_steps: Worker will perform `get_model` from the parameter
server every this many steps.
ps_channels: A channel list to all parameter server pods.
pservers: A list of parameter server pods.
distribution_strategy: The distribution startegy used by workers, e.g.
DistributionStrategy.PARAMETER_SERVER or
DistributionStrategy.AllreduceStrategy.
Returns:
An integer indicating the model version after the distributed training
and evaluation.
"""
job_type = (JobType.TRAINING_WITH_EVALUATION
if training else JobType.EVALUATION_ONLY)
evaluation_steps = 1 if job_type == JobType.TRAINING_WITH_EVALUATION else 0
batch_size = 8 if dataset_name == DatasetName.IMAGENET else 16
pservers = pservers or []
ps_channels = ps_channels or []
model_module = load_module(get_module_file_path(model_zoo_path,
model_def)).__dict__
for channel in ps_channels:
grpc.channel_ready_future(channel).result()
worker_arguments = [
"--worker_id",
"1",
"--job_type",
job_type,
"--minibatch_size",
batch_size,
"--model_zoo",
model_zoo_path,
"--model_def",
model_def,
"--model_params",
model_params,
"--loss",
loss,
"--get_model_steps",
get_model_steps,
"--distribution_strategy",
distribution_strategy,
]
args = parse_worker_args(worker_arguments)
worker = Worker(args, ps_channels=ps_channels)
if dataset_name in [DatasetName.IMAGENET, DatasetName.FRAPPE]:
record_num = batch_size
else:
record_num = 128
shards = {
create_recordio_file(record_num, dataset_name, feature_shape): (
0,
record_num,
)
}
if training:
training_shards = shards
evaluation_shards = shards
else:
training_shards = {}
evaluation_shards = shards
task_d = _TaskDispatcher(
training_shards,
evaluation_shards,
{},
records_per_task=64,
num_epochs=1,
)
if training:
evaluation_service = EvaluationService(
None,
task_d,
0,
0,
evaluation_steps,
False,
model_module[eval_metrics_fn],
)
else:
evaluation_service = EvaluationService(
None,
task_d,
0,
0,
evaluation_steps,
True,
model_module[eval_metrics_fn],
)
task_d.set_evaluation_service(evaluation_service)
master = MasterServicer(
batch_size,
task_d,
evaluation_service=evaluation_service,
)
callbacks = [
callback_class(master, worker) for callback_class in callback_classes
]
in_process_master = InProcessMaster(master, callbacks)
worker._stub = in_process_master
for pservicer in pservers:
pservicer._master_stub = in_process_master
worker.run()
req = elasticdl_pb2.GetTaskRequest()
req.worker_id = 1
task = master.get_task(req, None)
# No more task.
if task.shard_name:
raise RuntimeError(
"There are some tasks unfinished after worker exits.")
return master._version
def distributed_train_and_evaluate(
feature_shape,
model_zoo_path,
model_def,
model_params="",
eval_metrics_fn="eval_metrics_fn",
training=True,
dataset_name=DatasetName.IMAGE_DEFAULT,
callback_classes=[],
use_async=False,
get_model_steps=1,
):
"""Runs distributed training and evaluation with a local master. Grpc
calls are mocked by local master call.
Args:
feature_shape: The shape of model input.
model_zoo_path: The directory that contains user-defined model files
or a specific model file.
model_def: The import path to the model definition function/class in
the model zoo, e.g. "cifar10_subclass.CustomModel".
model_params: The dictionary of model parameters in a string that will
be used to instantiate the model, e.g. "param1=1,param2=2".
training: True for job type `TRAIN_WITH_EVALUATION`, False for
job type `EVALUATION`.
dataset_name: A dataset name from `DatasetName`.
callback_classes: A List of callbacks that will be called at given
stages of the training procedure.
use_async: A python bool. True if using asynchronous updates.
get_model_steps: Worker will perform `get_model` from the parameter
server every this many steps.
Returns:
An integer indicating the model version after the distributed training
and evaluation.
"""
job_type = (JobType.TRAINING_WITH_EVALUATION
if training else JobType.EVALUATION_ONLY)
batch_size = 8 if dataset_name == DatasetName.IMAGENET else 16
arguments = [
"--worker_id",
"1",
"--job_type",
job_type,
"--minibatch_size",
batch_size,
"--model_zoo",
model_zoo_path,
"--model_def",
model_def,
"--model_params",
model_params,
"--get_model_steps",
get_model_steps,
]
args = parse_worker_args(arguments)
worker = Worker(args)
if dataset_name in [DatasetName.IMAGENET, DatasetName.FRAPPE]:
record_num = batch_size
else:
record_num = 128
shards = {
create_recordio_file(record_num, dataset_name, feature_shape): (
0,
record_num,
)
}
if training:
training_shards = shards
evaluation_shards = shards
else:
training_shards = {}
evaluation_shards = shards
task_d = _TaskDispatcher(
training_shards,
evaluation_shards,
{},
records_per_task=64,
num_epochs=1,
)
model_module = load_module(get_module_file_path(model_zoo_path,
model_def)).__dict__
checkpoint_service = CheckpointService("", 0, 0, True)
if training:
evaluation_service = EvaluationService(
checkpoint_service,
None,
task_d,
0,
0,
1,
False,
model_module[eval_metrics_fn],
)
else:
evaluation_service = EvaluationService(
checkpoint_service,
None,
task_d,
0,
0,
0,
True,
model_module[eval_metrics_fn],
)
task_d.set_evaluation_service(evaluation_service)
grads_to_wait = 1 if use_async else 2
master = MasterServicer(
grads_to_wait,
batch_size,
worker._opt_fn(),
task_d,
init_var=[],
checkpoint_filename_for_init="",
checkpoint_service=checkpoint_service,
evaluation_service=evaluation_service,
use_async=use_async,
)
callbacks = [
callback_class(master, worker) for callback_class in callback_classes
]
worker._stub = InProcessMaster(master, callbacks)
for var in worker._model.trainable_variables:
master.set_model_var(var.name, var.numpy())
worker.run()
req = elasticdl_pb2.GetTaskRequest()
req.worker_id = 1
task = master.GetTask(req, None)
# No more task.
if task.shard_name:
raise RuntimeError(
"There are some tasks unfinished after worker exits.")
return master._version
def distributed_train_and_evaluate(
self,
feature_shape,
model_def,
model_params="",
training=True,
dataset="",
):
"""
Run distributed training and evaluation with a local master.
grpc calls are mocked by local master call.
"""
job_type = (JobType.TRAINING_ONLY
if training else JobType.EVALUATION_ONLY)
batch_size = 16
worker = Worker(
1,
job_type,
batch_size,
_model_zoo_path,
model_def=model_def,
model_params=model_params,
channel=None,
)
if dataset == "imagenet":
batch_size = 8
shards = {create_imagenet_recordio_file(8, feature_shape): (0, 8)}
elif dataset == "frappe":
shards = {
create_frappe_recordio_file(16, feature_shape, 5383): (0, 16)
}
else:
shards = {create_recordio_file(128, feature_shape): (0, 128)}
if training:
training_shards = shards
evaluation_shards = shards
else:
training_shards = {}
evaluation_shards = shards
task_d = _TaskDispatcher(
training_shards,
evaluation_shards,
{},
records_per_task=64,
num_epochs=1,
)
# Initialize checkpoint service
checkpoint_service = CheckpointService("", 0, 0, True)
if training:
evaluation_service = EvaluationService(checkpoint_service, None,
task_d, 0, 0, 1, False)
else:
evaluation_service = EvaluationService(checkpoint_service, None,
task_d, 0, 0, 0, True)
task_d.set_evaluation_service(evaluation_service)
# The master service
master = MasterServicer(
2,
batch_size,
worker._opt_fn(),
task_d,
init_var=[],
checkpoint_filename_for_init="",
checkpoint_service=checkpoint_service,
evaluation_service=evaluation_service,
)
worker._stub = InProcessMaster(master)
for var in worker._model.trainable_variables:
master.set_model_var(var.name, var.numpy())
worker.run()
req = elasticdl_pb2.GetTaskRequest()
req.worker_id = 1
task = master.GetTask(req, None)
# No more task.
self.assertTrue(not task.shard_name)
def distributed_train_and_evaluate(self, training=True):
"""
Run distributed training and evaluation with a local master.
grpc calls are mocked by local master call.
"""
class _Master(InProcessMaster):
def ReportGradient(self, req):
if 2 < self._m._version < 80:
# For testing of retrain when gradient not accepted.
# Increase master version to reject the gradient.
self._m._version += 1
return self._m.ReportGradient(req, None)
def ReportEvaluationMetrics(self, req):
if 2 < self._m._version < 80:
# Testing of evaluation retries. Increase the master
# version so the evaluation metrics will not be accepted.
self._m._version += 1
return self._m.ReportEvaluationMetrics(req, None)
job_type = (JobType.TRAINING_ONLY
if training else JobType.EVALUATION_ONLY)
batch_size = 16
worker = Worker(
1,
job_type,
batch_size,
_model_zoo_path,
model_def="test_module.custom_model",
channel=None,
)
shards = {create_recordio_file(128): 128}
if training:
training_shards = shards
evaluation_shards = {}
else:
training_shards = {}
evaluation_shards = shards
task_d = _TaskDispatcher(
training_shards,
evaluation_shards,
{},
records_per_task=64,
num_epochs=1,
)
if not training:
evaluation_service = EvaluationService(None, None, task_d, 0, 0, 0,
True)
task_d.set_evaluation_service(evaluation_service)
else:
evaluation_service = None
master = MasterServicer(
2,
batch_size,
worker._opt_fn(),
task_d,
init_var=[],
checkpoint_filename_for_init="",
checkpoint_service=None,
evaluation_service=evaluation_service,
)
worker._stub = _Master(master)
for var in worker._model.trainable_variables:
master.set_model_var(var.name, var.numpy())
worker.run()
req = elasticdl_pb2.GetTaskRequest()
req.worker_id = 1
task = master.GetTask(req, None)
# No more task.
self.assertTrue(not task.shard_name)
def distributed_train_and_evaluate(
self,
training=True,
callback_classes=[],
use_async=False,
grads_to_wait=2,
get_model_steps=1,
):
"""
Run distributed training and evaluation with a local master.
grpc calls are mocked by local master call.
"""
if use_async and grads_to_wait > 1:
raise ValueError(
"grads_to_wait should be 1 when using asynchronous SGD."
)
job_type = (
JobType.TRAINING_ONLY if training else JobType.EVALUATION_ONLY
)
batch_size = 16
worker = Worker(
1,
job_type,
batch_size,
_model_zoo_path,
model_def="test_module.custom_model",
channel=None,
get_model_steps=get_model_steps,
)
shards = {create_recordio_file(128): (0, 128)}
if training:
training_shards = shards
evaluation_shards = {}
else:
training_shards = {}
evaluation_shards = shards
task_d = _TaskDispatcher(
training_shards,
evaluation_shards,
{},
records_per_task=64,
num_epochs=1,
)
if not training:
evaluation_service = EvaluationService(
None, None, task_d, 0, 0, 0, True
)
task_d.set_evaluation_service(evaluation_service)
else:
evaluation_service = None
master = MasterServicer(
grads_to_wait,
batch_size,
worker._opt_fn(),
task_d,
init_var=[],
checkpoint_filename_for_init="",
checkpoint_service=None,
evaluation_service=evaluation_service,
use_async=use_async,
)
callbacks = [
callback_class(master, worker, self)
for callback_class in callback_classes
]
worker._stub = InProcessMaster(master, callbacks)
for var in worker._model.trainable_variables:
master.set_model_var(var.name, var.numpy())
worker.run()
req = elasticdl_pb2.GetTaskRequest()
req.worker_id = 1
task = master.GetTask(req, None)
# No more task.
self.assertTrue(not task.shard_name)
