def test_build_tf_config_with_multiple_hosts(trainer):
hosts = ['algo-1', 'algo-2', 'algo-3', 'algo-4']
current_host = 'algo-3'
test_wrapper = trainer.Trainer(customer_script=mock_script,
current_host=current_host,
hosts=hosts,
model_path=model_path)
tf_config = test_wrapper.build_tf_config()
expected_tf_config = {
'environment': 'cloud',
'cluster': {
'master': ['algo-1:2222'],
'ps': ['algo-1:2223', 'algo-2:2223', 'algo-3:2223', 'algo-4:2223'],
'worker': ['algo-2:2222', 'algo-3:2222', 'algo-4:2222']
},
'task': {
'index': 1,
'type': 'worker'
}
}
assert tf_config == expected_tf_config
assert test_wrapper.task_type == 'worker'
def test_train_estimator_fn(os_environ, botocore, boto3, inspect_args, trainer, modules):
'''
this test ensures that customers functions estimator_fn, train_input_fn, eval_input_fn, and serving_input_fn are
being invoked with the right params
'''
customer_script = MagicMock(spec=['estimator_fn', 'train_input_fn', 'eval_input_fn', 'serving_input_fn'])
training_dir_path = 'mytrainingpath'
_trainer = trainer.Trainer(customer_script=customer_script,
current_host=current_host,
hosts=hosts,
model_path='s3://my/s3/path',
customer_params={'training_steps': 10, 'num_gpu': 20},
input_channels={'training': training_dir_path})
modules.learn_runner.run.side_effect = lambda experiment_fn, run_config, hparams: experiment_fn(run_config, hparams)
modules.Experiment.side_effect = execute_input_functions
modules.saved_model_export_utils.make_export_strategy.side_effect = make_export_strategy_fn
_trainer.train()
modules.learn_runner.run.assert_called()
modules.Experiment.assert_called()
expected_params = {'num_gpu': 20, 'min_eval_frequency': 1000, 'training_steps': 10, 'save_checkpoints_secs': 300}
customer_script.estimator_fn.assert_called_with(modules.RunConfig(), modules.HParams().values())
customer_script.train_input_fn.assert_called_with(training_dir=training_dir_path, hyperparameters=expected_params)
customer_script.eval_input_fn.assert_called_with(training_dir=training_dir_path, hyperparameters=expected_params)
customer_script.serving_input_fn.assert_called_with(expected_params)
def test_trainer_run_config_params(os_environ, botocore, boto3, inspect_args, trainer, modules):
'''
this test ensures that customers functions model_fn, train_input_fn, eval_input_fn, and serving_input_fn are
being invoked with the right params
'''
customer_script = MagicMock(spec=['model_fn', 'train_input_fn', 'eval_input_fn', 'serving_input_fn'])
_trainer = trainer.Trainer(customer_script=customer_script,
current_host=current_host,
hosts=hosts,
model_path='s3://my/s3/path',
eval_steps=23,
customer_params={'save_summary_steps': 1,
'save_checkpoints_secs': 2,
'save_checkpoints_steps': 3,
'keep_checkpoint_max': 4,
'keep_checkpoint_every_n_hours': 5,
'log_step_count_steps': 6},
input_channels={'training': 'mytrainingpath'})
modules.learn_runner.run.side_effect = lambda experiment_fn, run_config, hparams: experiment_fn(run_config, hparams)
_trainer.train()
modules.RunConfig.assert_called_with(
save_summary_steps=1,
save_checkpoints_secs=2,
save_checkpoints_steps=3,
keep_checkpoint_max=4,
keep_checkpoint_every_n_hours=5,
log_step_count_steps=6,
model_dir='s3://my/s3/path'
)
def test_configure_s3_file_system(os_env, botocore, boto_client, trainer):
trainer.Trainer(customer_script=mock_script,
current_host=current_host,
hosts=hosts,
model_path='s3://my/s3/path')
boto_client('s3').get_bucket_location.assert_called_once_with(Bucket='my')
calls = [
call('S3_USE_HTTPS', '1'),
call('S3_REGION',
boto_client('s3').get_bucket_location()['LocationConstraint'])
]
os_env.__setitem__.assert_has_calls(calls, any_order=True)
def test_trainer_experiment_params(os_environ, botocore, boto3, inspect_args,
trainer, modules):
'''
this test ensures that customers functions model_fn, train_input_fn, eval_input_fn, and serving_input_fn are
being invoked with the right params
'''
customer_script = MagicMock(spec=[
'model_fn', 'train_input_fn', 'eval_input_fn', 'serving_input_fn'
])
_trainer = trainer.Trainer(customer_script=customer_script,
current_host=current_host,
hosts=hosts,
model_path='s3://my/s3/path',
eval_steps=23,
customer_params={
'min_eval_frequency': 2,
'local_eval_frequency': 3,
'eval_delay_secs': 7,
'continuous_eval_throttle_secs': 25,
'train_steps_per_iteration': 13
},
input_channels={'training': 'mytrainingpath'})
modules.learn_runner.run.side_effect = lambda experiment_fn, run_config, hparams: experiment_fn(
run_config, hparams)
_trainer.train()
modules.Experiment.assert_called_with(
estimator=modules.estimator.Estimator(),
train_input_fn=ANY,
eval_input_fn=ANY,
export_strategies=ANY,
train_steps=ANY,
eval_steps=23,
min_eval_frequency=2,
local_eval_frequency=3,
eval_delay_secs=7,
continuous_eval_throttle_secs=25,
train_steps_per_iteration=13)
def test_build_tf_config_with_one_host(trainer):
hosts = ['algo-1']
current_host = 'algo-1'
test_wrapper = trainer.Trainer(customer_script=mock_script,
current_host=current_host,
hosts=hosts,
model_path=model_path)
tf_config = test_wrapper.build_tf_config()
expected_tf_config = {
'environment': 'cloud',
'cluster': {
'master': ['algo-1:2222']
},
'task': {'index': 0, 'type': 'master'}
}
assert tf_config == expected_tf_config
assert test_wrapper.task_type == 'master'
