def test_main(self):
# Mock out all of utils except parser
hpo._utils = MagicMock()
hpo._utils.add_default_client_arguments = _utils.add_default_client_arguments
# Set some static returns
hpo._utils.create_hyperparameter_tuning_job.return_value = 'job-name'
hpo._utils.get_best_training_job_and_hyperparameters.return_value = 'best_job', {
"key_1": "best_hp_1"
}
hpo._utils.get_image_from_job.return_value = 'training-image'
hpo._utils.get_model_artifacts_from_job.return_value = 'model-artifacts'
hpo.main(required_args)
# Check if correct requests were created and triggered
hpo._utils.create_hyperparameter_tuning_job.assert_called()
hpo._utils.wait_for_hyperparameter_training_job.assert_called()
# Check the file outputs
hpo._utils.write_output.assert_has_calls([
call('/tmp/hpo_job_name_output_path', 'job-name'),
call('/tmp/model_artifact_url_output_path', 'model-artifacts'),
call('/tmp/best_job_name_output_path', 'best_job'),
call('/tmp/best_hyperparameters_output_path',
{"key_1": "best_hp_1"},
json_encode=True),
call('/tmp/training_image_output_path', 'training-image')
])
def test_main_assumes_role(self):
# Mock out all of utils except parser
hpo._utils = MagicMock()
hpo._utils.add_default_client_arguments = _utils.add_default_client_arguments
# Set some static returns
hpo._utils.create_hyperparameter_tuning_job.return_value = 'job-name'
hpo._utils.get_best_training_job_and_hyperparameters.return_value = 'best_job', {"key_1": "best_hp_1"}
assume_role_args = required_args + ['--assume_role', 'my-role']
hpo.main(assume_role_args)
hpo._utils.get_sagemaker_client.assert_called_once_with('us-west-2', None, assume_role_arn='my-role')
def test_main(self):
# Mock out all of utils except parser
hpo._utils = MagicMock()
hpo._utils.add_default_client_arguments = _utils.add_default_client_arguments
# Set some static returns
hpo._utils.create_hyperparameter_tuning_job.return_value = 'job-name'
hpo._utils.get_best_training_job_and_hyperparameters.return_value = 'best_job', 'best_hyperparameters'
hpo._utils.get_image_from_job.return_value = 'training-image'
hpo._utils.get_model_artifacts_from_job.return_value = 'model-artifacts'
with patch('builtins.open', mock_open()) as file_open:
hpo.main(required_args)
# Check if correct requests were created and triggered
hpo._utils.create_hyperparameter_tuning_job.assert_called()
hpo._utils.wait_for_hyperparameter_training_job.assert_called()
# Check the file outputs
file_open.assert_has_calls([
call('/tmp/hpo_job_name.txt', 'w'),
call('/tmp/best_job_name.txt', 'w'),
call('/tmp/best_hyperparameters.txt', 'w'),
call('/tmp/model_artifact_url.txt', 'w'),
call('/tmp/training_image.txt', 'w')
],
any_order=True)
file_open().write.assert_has_calls([
call('job-name'),
call('best_job'),
call('"best_hyperparameters"'),
call('model-artifacts'),
call('training-image'),
],
any_order=False)