def testGetAiPlatformTrainingPythonVersion(self):
if sys.version_info.major == 2:
self.assertEqual('2.7', runner._get_caip_python_version('1.14'))
self.assertEqual('2.7', runner._get_caip_python_version('1.15'))
else: # 3.x
self.assertEqual('3.5', runner._get_caip_python_version('1.14'))
self.assertEqual('3.7', runner._get_caip_python_version('1.15'))
self.assertEqual('3.7', runner._get_caip_python_version('2.1'))
def testDeployModelForCMLEServing(self, mock_discovery):
serving_path = os.path.join(self._output_data_dir, 'serving_path')
model_version = 'model_version'
mock_discovery.build.return_value = self._mock_api_client
mock_create = mock.Mock()
mock_create.return_value.execute.return_value = {'name': 'op_name'}
self._mock_api_client.projects().models().versions(
).create = mock_create
mock_get = mock.Mock()
self._mock_api_client.projects().operations().get = mock_get
mock_get.return_value.execute.return_value = {
'done': 'Done',
}
runner.deploy_model_for_cmle_serving(serving_path, model_version,
self._cmle_serving_args)
mock_create.assert_called_with(body=mock.ANY,
parent='projects/{}/models/{}'.format(
self._project_id, 'model_name'))
(_, kwargs) = mock_create.call_args
body = kwargs['body']
self.assertDictEqual(
{
'name': 'v{}'.format(model_version),
'deployment_uri': serving_path,
'runtime_version': runner._get_tf_runtime_version(),
'python_version': runner._get_caip_python_version(),
}, body)
mock_get.assert_called_with(name='op_name')
def testDeployModelForAIPPredictionWithCustomMachineType(self):
self._setUpPredictionMocks()
self._ai_platform_serving_args['machine_type'] = 'custom_machine_type'
runner.deploy_model_for_aip_prediction(self._mock_api_client,
self._serving_path,
self._model_version,
self._ai_platform_serving_args,
self._job_labels)
expected_versions_create_body = {
'name':
self._model_version,
'deployment_uri':
self._serving_path,
'machine_type':
'custom_machine_type',
'runtime_version':
runner._get_tf_runtime_version(tf.__version__),
'python_version':
runner._get_caip_python_version(
runner._get_tf_runtime_version(tf.__version__)),
'labels':
self._job_labels,
}
self._assertDeployModelMockCalls(
expected_versions_create_body=expected_versions_create_body)
def _assertDeployModelMockCalls(self,
expected_models_create_body=None,
expected_versions_create_body=None,
expect_set_default=True):
if not expected_models_create_body:
expected_models_create_body = {
'name': self._model_name,
'regions': [],
}
if not expected_versions_create_body:
with telemetry_utils.scoped_labels({
telemetry_utils.LABEL_TFX_EXECUTOR:
self._executor_class_path
}):
labels = telemetry_utils.get_labels_dict()
expected_versions_create_body = {
'name':
self._model_version,
'deployment_uri':
self._serving_path,
'runtime_version':
runner._get_tf_runtime_version(tf.__version__),
'python_version':
runner._get_caip_python_version(
runner._get_tf_runtime_version(tf.__version__)),
'labels':
labels
}
self._mock_models_create.assert_called_with(
body=mock.ANY,
parent='projects/{}'.format(self._project_id),
)
(_, models_create_kwargs) = self._mock_models_create.call_args
self.assertDictEqual(expected_models_create_body,
models_create_kwargs['body'])
self._mock_versions_create.assert_called_with(
body=mock.ANY,
parent='projects/{}/models/{}'.format(self._project_id,
self._model_name))
(_, versions_create_kwargs) = self._mock_versions_create.call_args
self.assertDictEqual(expected_versions_create_body,
versions_create_kwargs['body'])
if not expect_set_default:
return
self._mock_set_default.assert_called_with(
name='projects/{}/models/{}/versions/{}'.format(
self._project_id, self._model_name, self._model_version))
self._mock_set_default_execute.assert_called_with()
def _assertDeployModelMockCalls(self,
expected_models_create_body=None,
expected_versions_create_body=None,
expect_set_default=True):
if not expected_models_create_body:
expected_models_create_body = {
'name': self._model_name,
'regions': [],
'labels': self._job_labels
}
if not expected_versions_create_body:
expected_versions_create_body = {
'name':
self._model_version,
'deployment_uri':
self._serving_path,
'runtime_version':
runner._get_tf_runtime_version(tf.__version__),
'python_version':
runner._get_caip_python_version(
runner._get_tf_runtime_version(tf.__version__)),
'labels':
self._job_labels
}
self._mock_models_create.assert_called_with(
body=mock.ANY,
parent='projects/{}'.format(self._project_id),
)
models_create_kwargs = self._mock_models_create.call_args[1]
self.assertDictEqual(expected_models_create_body,
models_create_kwargs['body'])
self._mock_versions_create.assert_called_with(
body=mock.ANY,
parent='projects/{}/models/{}'.format(self._project_id,
self._model_name))
versions_create_kwargs = self._mock_versions_create.call_args[1]
self.assertDictEqual(expected_versions_create_body,
versions_create_kwargs['body'])
if not expect_set_default:
return
self._mock_set_default.assert_called_with(
name='projects/{}/models/{}/versions/{}'.format(
self._project_id, self._model_name, self._model_version))
self._mock_set_default_execute.assert_called_with()
def testDeployModelForAIPPredictionWithCustomRegion(self, mock_discovery):
serving_path = os.path.join(self._output_data_dir, 'serving_path')
model_version = 'model_version'
mock_discovery.build.return_value = self._mock_api_client
mock_create = mock.Mock()
mock_create.return_value.execute.return_value = {'name': 'op_name'}
self._mock_api_client.projects().models().versions(
).create = mock_create
mock_get = mock.Mock()
self._mock_api_client.projects().operations().get = mock_get
mock_set_default = mock.Mock()
self._mock_api_client.projects().models().versions(
).setDefault = mock_set_default
mock_set_default_execute = mock.Mock()
self._mock_api_client.projects().models().versions().setDefault(
).execute = mock_set_default_execute
mock_get.return_value.execute.return_value = {
'done': 'Done',
'response': {
'name': model_version
},
}
self._ai_platform_serving_args['regions'] = ['custom-region']
runner.deploy_model_for_aip_prediction(serving_path, model_version,
self._ai_platform_serving_args)
mock_create.assert_called_with(body=mock.ANY,
parent='projects/{}/models/{}'.format(
self._project_id, 'model_name'))
(_, kwargs) = mock_create.call_args
body = kwargs['body']
self.assertDictEqual(
{
'name': 'v{}'.format(model_version),
'regions': ['custom-region'],
'deployment_uri': serving_path,
'runtime_version': runner._get_tf_runtime_version(
tf.__version__),
'python_version': runner._get_caip_python_version(),
}, body)
mock_get.assert_called_with(name='op_name')
mock_set_default.assert_called_with(
name='projects/{}/models/{}/versions/{}'.format(
self._project_id, 'model_name', model_version))
mock_set_default_execute.assert_called_with()
def testDeployModelForAIPPredictionWithCustomRegion(self, mock_discovery):
mock_discovery.build.return_value = self._mock_api_client
self._setUpPredictionMocks()
self._ai_platform_serving_args['regions'] = ['custom-region']
runner.deploy_model_for_aip_prediction(self._serving_path,
self._model_version,
self._ai_platform_serving_args,
self._executor_class_path)
self._mock_models_create.assert_called_with(
body=mock.ANY,
parent='projects/{}'.format(self._project_id),
)
(_, models_create_kwargs) = self._mock_models_create.call_args
models_create_body = models_create_kwargs['body']
self.assertDictEqual(
{
'name': 'model_name',
'regions': ['custom-region']
}, models_create_body)
self._mock_versions_create.assert_called_with(
body=mock.ANY,
parent='projects/{}/models/{}'.format(self._project_id,
'model_name'))
(_, versions_create_kwargs) = self._mock_versions_create.call_args
versions_create_body = versions_create_kwargs['body']
with telemetry_utils.scoped_labels(
{telemetry_utils.TFX_EXECUTOR: self._executor_class_path}):
labels = telemetry_utils.get_labels_dict()
runtime_version = runner._get_tf_runtime_version(tf.__version__)
self.assertDictEqual(
{
'name': self._model_version,
'deployment_uri': self._serving_path,
'runtime_version': runtime_version,
'python_version':
runner._get_caip_python_version(runtime_version),
'labels': labels,
}, versions_create_body)
self._mock_get.assert_called_with(name='op_name')
self._mock_set_default.assert_called_with(
name='projects/{}/models/{}/versions/{}'.format(
self._project_id, 'model_name', self._model_version))
self._mock_set_default_execute.assert_called_with()
def testDeployModelForAIPPredictionWithCustomRuntime(self):
self._setUpPredictionMocks()
self._ai_platform_serving_args['runtime_version'] = '1.23.45'
runner.deploy_model_for_aip_prediction(self._mock_api_client,
self._serving_path,
self._model_version,
self._ai_platform_serving_args,
self._job_labels)
expected_versions_create_body = {
'name': self._model_version,
'deployment_uri': self._serving_path,
'runtime_version': '1.23.45',
'python_version': runner._get_caip_python_version('1.23.45'),
'labels': self._job_labels,
}
self._assertDeployModelMockCalls(
expected_versions_create_body=expected_versions_create_body)
def testStartCMLETraining(self, mock_discovery, mock_dependency_utils):
mock_discovery.build.return_value = self._mock_api_client
mock_create = mock.Mock()
self._mock_api_client.projects().jobs().create = mock_create
mock_get = mock.Mock()
self._mock_api_client.projects().jobs().get.return_value = mock_get
mock_get.execute.return_value = {
'state': 'SUCCEEDED',
}
mock_dependency_utils.build_ephemeral_package.return_value = self._fake_package
class_path = 'foo.bar.class'
runner.start_cmle_training(self._inputs, self._outputs,
self._exec_properties, class_path,
self._training_inputs)
mock_dependency_utils.build_ephemeral_package.assert_called_with()
mock_create.assert_called_with(body=mock.ANY,
parent='projects/{}'.format(
self._project_id))
(_, kwargs) = mock_create.call_args
body = kwargs['body']
self.assertDictContainsSubset(
{
'pythonVersion':
runner._get_caip_python_version(),
'runtimeVersion':
'.'.join(tf.__version__.split('.')[0:2]),
'jobDir':
self._job_dir,
'args': [
'--executor_class_path', class_path, '--inputs', '{}',
'--outputs', '{}', '--exec-properties',
'{"custom_config": {}}'
],
'pythonModule':
'tfx.scripts.run_executor',
'packageUris': [os.path.join(self._job_dir, 'fake_package')],
}, body['trainingInput'])
self.assertStartsWith(body['jobId'], 'tfx_')
mock_get.execute.assert_called_with()
def testDeployModelForAIPPredictionWithCustomRuntime(self, mock_discovery):
mock_discovery.build.return_value = self._mock_api_client
self._setUpPredictionMocks()
self._ai_platform_serving_args['runtime_version'] = '1.23.45'
runner.deploy_model_for_aip_prediction(self._serving_path,
self._model_version,
self._ai_platform_serving_args,
self._executor_class_path)
with telemetry_utils.scoped_labels(
{telemetry_utils.LABEL_TFX_EXECUTOR: self._executor_class_path}):
labels = telemetry_utils.get_labels_dict()
expected_versions_create_body = {
'name': self._model_version,
'deployment_uri': self._serving_path,
'runtime_version': '1.23.45',
'python_version': runner._get_caip_python_version('1.23.45'),
'labels': labels,
}
self._assertDeployModelMockCalls(
expected_versions_create_body=expected_versions_create_body)
