def test_hptuning_job_remote_runner_on_payload_deserializes_correctly(
self, mock_path_exists, mock_job_service_client):
job_client = mock.Mock()
mock_job_service_client.return_value = job_client
create_hptuning_job_response = mock.Mock()
job_client.create_hyperparameter_tuning_job.return_value = create_hptuning_job_response
create_hptuning_job_response.name = self._hptuning_job_name
get_hptuning_job_response = mock.Mock()
job_client.get_hyperparameter_tuning_job.return_value = get_hptuning_job_response
get_hptuning_job_response.state = gca_job_state.JobState.JOB_STATE_SUCCEEDED
mock_path_exists.return_value = False
hyperparameter_tuning_job_remote_runner.create_hyperparameter_tuning_job(
self._type, self._project, self._location, self._payload,
self._gcp_resources)
expected_parent = f"projects/{self._project}/locations/{self._location}"
expected_job_spec = json_util.recursive_remove_empty(
json.loads(self._payload, strict=False))
job_client.create_hyperparameter_tuning_job.assert_called_once_with(
parent=expected_parent,
hyperparameter_tuning_job=expected_job_spec)
def test_hptuning_job_remote_runner_cancel(self, mock_execution_context,
mock_post_requests, _,
mock_auth,
mock_job_service_client):
creds = mock.Mock()
creds.token = 'fake_token'
mock_auth.return_value = [creds, "project"]
job_client = mock.Mock()
mock_job_service_client.return_value = job_client
create_hptuning_job_response = mock.Mock()
job_client.create_hyperparameter_tuning_job.return_value = create_hptuning_job_response
create_hptuning_job_response.name = self._hptuning_job_name
get_hptuning_job_response = mock.Mock()
job_client.get_hyperparameter_tuning_job.return_value = get_hptuning_job_response
get_hptuning_job_response.state = gca_job_state.JobState.JOB_STATE_SUCCEEDED
mock_execution_context.return_value = None
hyperparameter_tuning_job_remote_runner.create_hyperparameter_tuning_job(
self._type, self._project, self._location, self._payload,
self._gcp_resources)
# Call cancellation handler
mock_execution_context.call_args[1]['on_cancel']()
mock_post_requests.assert_called_once_with(
url=
f'{self._hptuning_job_uri_prefix}{self._hptuning_job_name}:cancel',
data='',
headers={
'Content-type': 'application/json',
'Authorization': 'Bearer fake_token',
})
def test_hptuning_job_remote_runner_raises_exception_empty_uri_in_gcp_resources(
self, _, mock_job_service_client):
job_client = mock.Mock()
mock_job_service_client.return_value = job_client
create_hptuning_job_response = mock.Mock()
job_client.create_hyperparameter_tuning_job.return_value = create_hptuning_job_response
create_hptuning_job_response.name = self._hptuning_job_name
get_hptuning_job_response_success = mock.Mock()
get_hptuning_job_response_success.state = gca_job_state.JobState.JOB_STATE_SUCCEEDED
job_client.get_hyperparameter_tuning_job.side_effect = [
get_hptuning_job_response_success
]
# Write the job proto to output
hptuning_job_resources = GcpResources()
hptuning_job_resource_1 = hptuning_job_resources.resources.add()
hptuning_job_resource_1.resource_type = "HyperparameterTuningJob"
hptuning_job_resource_1.resource_uri = ""
with open(self._gcp_resources, "w") as f:
f.write(json_format.MessageToJson(hptuning_job_resources))
with self.assertRaisesRegex(
ValueError,
"Job Name in gcp_resource is not formatted correctly or is empty."
):
hyperparameter_tuning_job_remote_runner.create_hyperparameter_tuning_job(
self._type, self._project, self._location, self._payload,
self._gcp_resources)
def test_hptuning_job_remote_runner_raises_exception_with_more_than_one_resource_in_gcp_resources(
self, _, mock_job_service_client):
job_client = mock.Mock()
mock_job_service_client.return_value = job_client
create_hptuning_job_response = mock.Mock()
job_client.create_hyperparameter_tuning_job.return_value = create_hptuning_job_response
create_hptuning_job_response.name = self._hptuning_job_name
get_hptuning_job_response_success = mock.Mock()
get_hptuning_job_response_success.state = gca_job_state.JobState.JOB_STATE_SUCCEEDED
job_client.get_hyperparameter_tuning_job.side_effect = [
get_hptuning_job_response_success
]
# Write the job proto to output
hptuning_job_resources = GcpResources()
hptuning_job_resource_1 = hptuning_job_resources.resources.add()
hptuning_job_resource_1.resource_type = "HyperparameterTuningJob"
hptuning_job_resource_1.resource_uri = f"{self._hptuning_job_uri_prefix}{self._hptuning_job_name}"
hptuning_job_resource_2 = hptuning_job_resources.resources.add()
hptuning_job_resource_2.resource_type = "HyperparameterTuningJob"
hptuning_job_resource_2.resource_uri = f"{self._hptuning_job_uri_prefix}{self._hptuning_job_name}"
with open(self._gcp_resources, "w") as f:
f.write(json_format.MessageToJson(hptuning_job_resources))
with self.assertRaisesRegex(
ValueError,
"gcp_resources should contain one resource, found 2"):
hyperparameter_tuning_job_remote_runner.create_hyperparameter_tuning_job(
self._type, self._project, self._location, self._payload,
self._gcp_resources)
def test_hptuning_job_remote_runner_returns_gcp_resources(
self, mock_time_sleep, mock_path_exists, mock_job_service_client):
job_client = mock.Mock()
mock_job_service_client.return_value = job_client
create_hptuning_job_response = mock.Mock()
job_client.create_hyperparameter_tuning_job.return_value = create_hptuning_job_response
create_hptuning_job_response.name = self._hptuning_job_name
get_hptuning_job_response_success = mock.Mock()
get_hptuning_job_response_success.state = gca_job_state.JobState.JOB_STATE_SUCCEEDED
job_client.get_hyperparameter_tuning_job.side_effect = [
get_hptuning_job_response_success
]
mock_path_exists.return_value = False
hyperparameter_tuning_job_remote_runner.create_hyperparameter_tuning_job(
self._type, self._project, self._location, self._payload,
self._gcp_resources)
with open(self._gcp_resources) as f:
serialized_gcp_resources = f.read()
# Instantiate GCPResources Proto
hptuning_job_resources = json_format.Parse(
serialized_gcp_resources, GcpResources())
self.assertLen(hptuning_job_resources.resources, 1)
hptuning_job_name = hptuning_job_resources.resources[
0].resource_uri[len(self._hptuning_job_uri_prefix):]
self.assertEqual(hptuning_job_name, self._hptuning_job_name)
def test_hptuning_job_remote_runner_retries_to_get_status_on_non_completed_job(
self, mock_time_sleep, mock_path_exists, mock_job_service_client):
job_client = mock.Mock()
mock_job_service_client.return_value = job_client
create_hptuning_job_response = mock.Mock()
job_client.create_hyperparameter_tuning_job.return_value = create_hptuning_job_response
create_hptuning_job_response.name = self._hptuning_job_name
get_hptuning_job_response_success = mock.Mock()
get_hptuning_job_response_success.state = gca_job_state.JobState.JOB_STATE_SUCCEEDED
get_hptuning_job_response_running = mock.Mock()
get_hptuning_job_response_running.state = gca_job_state.JobState.JOB_STATE_RUNNING
job_client.get_hyperparameter_tuning_job.side_effect = [
get_hptuning_job_response_running,
get_hptuning_job_response_success
]
mock_path_exists.return_value = False
hyperparameter_tuning_job_remote_runner.create_hyperparameter_tuning_job(
self._type, self._project, self._location, self._payload,
self._gcp_resources)
mock_time_sleep.assert_called_once_with(
job_remote_runner._POLLING_INTERVAL_IN_SECONDS)
self.assertEqual(job_client.get_hyperparameter_tuning_job.call_count,
2)
def test_hptuning_job_remote_runner_raises_exception_on_internal_error(
self, mock_path_exists, mock_job_service_client):
job_client = mock.Mock()
mock_job_service_client.return_value = job_client
create_hptuning_job_response = mock.Mock()
job_client.create_hyperparameter_tuning_job.return_value = create_hptuning_job_response
create_hptuning_job_response.name = self._hptuning_job_name
get_hptuning_job_response = mock.Mock()
job_client.get_hyperparameter_tuning_job.return_value = get_hptuning_job_response
get_hptuning_job_response.state = gca_job_state.JobState.JOB_STATE_FAILED
mock_path_exists.return_value = False
with self.assertRaises(SystemExit):
hyperparameter_tuning_job_remote_runner.create_hyperparameter_tuning_job(
self._type, self._project, self._location, self._payload,
self._gcp_resources)
def test_hptuning_job_remote_runner_on_region_is_set_correctly_in_client_options(
self, mock_job_service_client):
job_client = mock.Mock()
mock_job_service_client.return_value = job_client
create_hptuning_job_response = mock.Mock()
job_client.create_hyperparameter_tuning_job.return_value = create_hptuning_job_response
create_hptuning_job_response.name = self._hptuning_job_name
get_hptuning_job_response = mock.Mock()
job_client.get_hyperparameter_tuning_job.return_value = get_hptuning_job_response
get_hptuning_job_response.state = gca_job_state.JobState.JOB_STATE_SUCCEEDED
hyperparameter_tuning_job_remote_runner.create_hyperparameter_tuning_job(
self._type, self._project, self._location, self._payload,
self._gcp_resources)
mock_job_service_client.assert_called_once_with(client_options={
"api_endpoint":
"test_region-aiplatform.googleapis.com"
},
client_info=mock.ANY)
