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 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 testGetTensorflowRuntime(self):
self.assertEqual('1.14', runner._get_tf_runtime_version('1.14'))
self.assertEqual('1.15', runner._get_tf_runtime_version('1.15.0'))
self.assertEqual('1.15', runner._get_tf_runtime_version('1.15.1'))
self.assertEqual('1.15', runner._get_tf_runtime_version('2.0.0'))
self.assertEqual('1.15', runner._get_tf_runtime_version('2.0.1'))
self.assertEqual('2.1', runner._get_tf_runtime_version('2.1.0'))
self.assertEqual('2.2', runner._get_tf_runtime_version('2.2.0'))
self.assertEqual('2.2', runner._get_tf_runtime_version('2.3.0'))
def testGetTensorflowRuntime(self):
self.assertEqual('1.14', runner._get_tf_runtime_version('1.14'))
self.assertEqual('1.15', runner._get_tf_runtime_version('1.15.0'))
self.assertEqual('1.15', runner._get_tf_runtime_version('1.15.1'))
self.assertEqual('1.15', runner._get_tf_runtime_version('2.0.0'))
self.assertEqual('1.15', runner._get_tf_runtime_version('2.0.1'))
self.assertEqual('2.1', runner._get_tf_runtime_version('2.1.0'))
# TODO(b/157039850) Remove this once CAIP model support TF 2.2 runtime.
self.assertEqual('2.1', runner._get_tf_runtime_version('2.2.0'))
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()
