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 testDeployModelForAIPPredictionError(self):
self._setUpPredictionMocks()
self._mock_get.return_value.execute.return_value = {
'done': True,
'error': {
'code': 999,
'message': 'it was an error.'
},
}
with self.assertRaises(RuntimeError):
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_models_create_body = {
'name': self._model_name,
'regions': [],
'labels': self._job_labels
}
self._assertDeployModelMockCalls(
expected_models_create_body=expected_models_create_body,
expect_set_default=False)
def testDeployModelForVertexPredictionError(self):
self._setUpVertexPredictionMocks()
self._mock_endpoint_list.side_effect = [[], [self._mock_endpoint]]
self._mock_model_deploy.side_effect = errors.HttpError(
httplib2.Response(info={'status': 429}), b'')
with self.assertRaises(RuntimeError):
runner.deploy_model_for_aip_prediction(
serving_path=self._serving_path,
model_version_name=self._model_name,
ai_platform_serving_args=self._ai_platform_serving_args_vertex,
labels=self._job_labels,
serving_container_image_uri=self._serving_container_image_uri,
endpoint_region=self._endpoint_region,
enable_vertex=True)
expected_endpoint_create_body = {
'display_name': self._endpoint_name,
'labels': self._job_labels,
}
expected_model_upload_body = {
'display_name': self._model_name,
'artifact_uri': self._serving_path,
'serving_container_image_uri': self._serving_container_image_uri,
}
expected_model_deploy_body = {
'endpoint': self._mock_endpoint,
'traffic_percentage': 100,
}
self._assertDeployModelMockCallsVertex(
expected_endpoint_create_body=expected_endpoint_create_body,
expected_model_upload_body=expected_model_upload_body,
expected_model_deploy_body=expected_model_deploy_body)
def testDeployModelForVertexPrediction(self):
self._setUpVertexPredictionMocks()
self._mock_endpoint_list.side_effect = [[], [self._mock_endpoint]]
runner.deploy_model_for_aip_prediction(
serving_path=self._serving_path,
model_version_name=self._model_name,
ai_platform_serving_args=self._ai_platform_serving_args_vertex,
labels=self._job_labels,
serving_container_image_uri=self._serving_container_image_uri,
endpoint_region=self._endpoint_region,
enable_vertex=True)
expected_endpoint_create_body = {
'display_name': self._endpoint_name,
'labels': self._job_labels,
}
expected_model_upload_body = {
'display_name': self._model_name,
'artifact_uri': self._serving_path,
'serving_container_image_uri': self._serving_container_image_uri,
}
expected_model_deploy_body = {
'endpoint': self._mock_endpoint,
'traffic_percentage': 100,
}
self._assertDeployModelMockCallsVertex(
expected_endpoint_create_body=expected_endpoint_create_body,
expected_model_upload_body=expected_model_upload_body,
expected_model_deploy_body=expected_model_deploy_body)
def Do(self, input_dict: Dict[Text, List[types.Artifact]],
output_dict: Dict[Text, List[types.Artifact]],
exec_properties: Dict[Text, Any]):
"""Overrides the tfx_pusher_executor.
Args:
input_dict: Input dict from input key to a list of artifacts, including:
- model_export: exported model from trainer.
- model_blessing: model blessing path from model_validator.
output_dict: Output dict from key to a list of artifacts, including:
- model_push: A list of 'ModelPushPath' artifact of size one. It will
include the model in this push execution if the model was pushed.
exec_properties: Mostly a passthrough input dict for
tfx.components.Pusher.executor. custom_config.ai_platform_serving_args
is consumed by this class. For the full set of parameters supported by
Google Cloud AI Platform, refer to
https://cloud.google.com/ml-engine/docs/tensorflow/deploying-models#creating_a_model_version.
Returns:
None
Raises:
ValueError:
If ai_platform_serving_args is not in exec_properties.custom_config.
If Serving model path does not start with gs://.
RuntimeError: if the Google Cloud AI Platform training job failed.
"""
self._log_startup(input_dict, output_dict, exec_properties)
model_push = artifact_utils.get_single_instance(
output_dict[tfx_pusher_executor.PUSHED_MODEL_KEY])
if not self.CheckBlessing(input_dict):
model_push.set_int_custom_property('pushed', 0)
return
model_export = artifact_utils.get_single_instance(
input_dict[tfx_pusher_executor.MODEL_KEY])
model_export_uri = model_export.uri
exec_properties_copy = exec_properties.copy()
custom_config = exec_properties_copy.pop('custom_config', {})
ai_platform_serving_args = custom_config[SERVING_ARGS_KEY]
if not ai_platform_serving_args:
raise ValueError(
'\'ai_platform_serving_args\' is missing in \'custom_config\'')
# Deploy the model.
model_path = path_utils.serving_model_path(model_export_uri)
# Note: we do not have a logical model version right now. This
# model_version is a timestamp mapped to trainer's exporter.
model_version = os.path.basename(model_path)
executor_class_path = '%s.%s' % (self.__class__.__module__,
self.__class__.__name__)
runner.deploy_model_for_aip_prediction(
model_path,
model_version,
ai_platform_serving_args,
executor_class_path,
)
model_push.set_int_custom_property('pushed', 1)
model_push.set_string_custom_property('pushed_model', model_path)
def testDeployModelForAIPPrediction(self, mock_discovery):
mock_discovery.build.return_value = self._mock_api_client
self._setUpPredictionMocks()
runner.deploy_model_for_aip_prediction(self._serving_path,
self._model_version,
self._ai_platform_serving_args,
self._executor_class_path)
expected_models_create_body = {'name': self._model_name, 'regions': []}
self._assertDeployModelMockCalls(
expected_models_create_body=expected_models_create_body)
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._assertDeployModelMockCalls(
{
'name': 'model_name',
'regions': ['custom-region']
}, False)
def testDeployModelForAIPPrediction(self):
self._setUpPredictionMocks()
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_models_create_body = {
'name': self._model_name,
'regions': [],
'labels': self._job_labels
}
self._assertDeployModelMockCalls(
expected_models_create_body=expected_models_create_body)
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 testDeployModelForAIPPredictionWithCustomRegion(self):
self._setUpPredictionMocks()
self._ai_platform_serving_args['regions'] = ['custom-region']
runner.deploy_model_for_aip_prediction(
serving_path=self._serving_path,
model_version_name=self._model_version,
ai_platform_serving_args=self._ai_platform_serving_args,
labels=self._job_labels,
api=self._mock_api_client)
expected_models_create_body = {
'name': self._model_name,
'regions': ['custom-region'],
'labels': self._job_labels
}
self._assertDeployModelMockCalls(
expected_models_create_body=expected_models_create_body)
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 testDeployModelForAIPPredictionError(self, mock_discovery):
mock_discovery.build.return_value = self._mock_api_client
self._setUpPredictionMocks()
self._mock_get.return_value.execute.return_value = {
'done': True,
'error': {
'code': 999,
'message': 'it was an error.'
},
}
with self.assertRaises(RuntimeError):
runner.deploy_model_for_aip_prediction(self._serving_path,
self._model_version,
self._ai_platform_serving_args,
self._executor_class_path)
self._assertDeployModelMockCalls({
'name': 'model_name',
'regions': []
}, False)
def testDeployModelForVertexPredictionWithCustomMachineType(self):
self._setUpVertexPredictionMocks()
self._mock_endpoint_list.side_effect = [[], [self._mock_endpoint]]
self._ai_platform_serving_args_vertex[
'machine_type'] = 'custom_machine_type'
runner.deploy_model_for_aip_prediction(
serving_path=self._serving_path,
model_version_name=self._model_name,
ai_platform_serving_args=self._ai_platform_serving_args_vertex,
labels=self._job_labels,
serving_container_image_uri=self._serving_container_image_uri,
endpoint_region=self._endpoint_region,
enable_vertex=True)
expected_model_deploy_body = {
'endpoint': self._mock_endpoint,
'traffic_percentage': 100,
'machine_type': 'custom_machine_type',
}
self._assertDeployModelMockCallsVertex(
expected_model_deploy_body=expected_model_deploy_body)
def testDeployModelForVertexPredictionWithCustomRegion(self):
self._setUpVertexPredictionMocks()
self._mock_endpoint_list.side_effect = [[], [self._mock_endpoint]]
self._mock_init = mock.Mock()
aiplatform.init = self._mock_init
self._endpoint_region = 'custom-region'
runner.deploy_model_for_aip_prediction(
serving_path=self._serving_path,
model_version_name=self._model_name,
ai_platform_serving_args=self._ai_platform_serving_args_vertex,
labels=self._job_labels,
serving_container_image_uri=self._serving_container_image_uri,
endpoint_region=self._endpoint_region,
enable_vertex=True)
expected_init_body = {
'project': self._project_id,
'location': 'custom-region',
}
self._mock_init.assert_called_with(**expected_init_body)
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)
def Do(self, input_dict: Dict[Text, List[types.Artifact]],
output_dict: Dict[Text, List[types.Artifact]],
exec_properties: Dict[Text, Any]):
"""Overrides the tfx_pusher_executor.
Args:
input_dict: Input dict from input key to a list of artifacts, including:
- model_export: exported model from trainer.
- model_blessing: model blessing path from model_validator.
output_dict: Output dict from key to a list of artifacts, including:
- model_push: A list of 'ModelPushPath' artifact of size one. It will
include the model in this push execution if the model was pushed.
exec_properties: Mostly a passthrough input dict for
tfx.components.Pusher.executor. custom_config.ai_platform_serving_args
is consumed by this class. For the full set of parameters supported by
Google Cloud AI Platform, refer to
https://cloud.google.com/ml-engine/docs/tensorflow/deploying-models#creating_a_model_version.
Raises:
ValueError:
If ai_platform_serving_args is not in exec_properties.custom_config.
If Serving model path does not start with gs://.
RuntimeError: if the Google Cloud AI Platform training job failed.
"""
self._log_startup(input_dict, output_dict, exec_properties)
model_push = artifact_utils.get_single_instance(
output_dict[tfx_pusher_executor.PUSHED_MODEL_KEY])
if not self.CheckBlessing(input_dict):
self._MarkNotPushed(model_push)
return
model_export = artifact_utils.get_single_instance(
input_dict[tfx_pusher_executor.MODEL_KEY])
exec_properties_copy = exec_properties.copy()
custom_config = exec_properties_copy.pop(_CUSTOM_CONFIG_KEY, {})
ai_platform_serving_args = custom_config.get(SERVING_ARGS_KEY)
if not ai_platform_serving_args:
raise ValueError(
'\'ai_platform_serving_args\' is missing in \'custom_config\'')
# Deploy the model.
io_utils.copy_dir(
src=path_utils.serving_model_path(model_export.uri),
dst=model_push.uri)
model_path = model_push.uri
# TODO(jjong): Introduce Versioning.
# Note that we're adding "v" prefix as Cloud AI Prediction only allows the
# version name that starts with letters, and contains letters, digits,
# underscore only.
model_version = 'v{}'.format(int(time.time()))
executor_class_path = '%s.%s' % (self.__class__.__module__,
self.__class__.__name__)
runner.deploy_model_for_aip_prediction(
model_path,
model_version,
ai_platform_serving_args,
executor_class_path,
)
self._MarkPushed(
model_push,
pushed_destination=_CAIP_MODEL_VERSION_PATH_FORMAT.format(
project_id=ai_platform_serving_args['project_id'],
model=ai_platform_serving_args['model_name'],
version=model_version),
pushed_version=model_version)
def Do(self, input_dict: Dict[Text, List[types.Artifact]],
output_dict: Dict[Text, List[types.Artifact]],
exec_properties: Dict[Text, Any]):
"""Overrides the tfx_pusher_executor.
Args:
input_dict: Input dict from input key to a list of artifacts, including:
- model_export: exported model from trainer.
- model_blessing: model blessing path from evaluator.
output_dict: Output dict from key to a list of artifacts, including:
- model_push: A list of 'ModelPushPath' artifact of size one. It will
include the model in this push execution if the model was pushed.
exec_properties: Mostly a passthrough input dict for
tfx.components.Pusher.executor. The following keys in `custom_config`
are consumed by this class:
- ai_platform_serving_args: For the full set of parameters supported
by Google Cloud AI Platform, refer to
https://cloud.google.com/ml-engine/reference/rest/v1/projects.models.versions#Version.
- endpoint: Optional endpoint override. Should be in format of
`https://[region]-ml.googleapis.com`. Default to global endpoint if
not set. Using regional endpoint is recommended by Cloud AI Platform.
When set, 'regions' key in ai_platform_serving_args cannot be set.
For more details, please see
https://cloud.google.com/ai-platform/prediction/docs/regional-endpoints#using_regional_endpoints
Raises:
ValueError:
If ai_platform_serving_args is not in exec_properties.custom_config.
If Serving model path does not start with gs://.
If 'endpoint' and 'regions' are set simultanuously.
RuntimeError: if the Google Cloud AI Platform training job failed.
"""
self._log_startup(input_dict, output_dict, exec_properties)
custom_config = json_utils.loads(
exec_properties.get(_CUSTOM_CONFIG_KEY, 'null'))
if custom_config is not None and not isinstance(custom_config, Dict):
raise ValueError(
'custom_config in execution properties needs to be a '
'dict.')
ai_platform_serving_args = custom_config.get(SERVING_ARGS_KEY)
if not ai_platform_serving_args:
raise ValueError(
'\'ai_platform_serving_args\' is missing in \'custom_config\'')
endpoint = custom_config.get(ENDPOINT_ARGS_KEY)
if endpoint and 'regions' in ai_platform_serving_args:
raise ValueError(
'\'endpoint\' and \'ai_platform_serving_args.regions\' cannot be set simultanuously'
)
model_push = artifact_utils.get_single_instance(
output_dict[standard_component_specs.PUSHED_MODEL_KEY])
if not self.CheckBlessing(input_dict):
self._MarkNotPushed(model_push)
return
service_name, api_version = runner.get_service_name_and_api_version(
ai_platform_serving_args)
# Deploy the model.
io_utils.copy_dir(src=self.GetModelPath(input_dict),
dst=model_push.uri)
model_path = model_push.uri
# TODO(jjong): Introduce Versioning.
# Note that we're adding "v" prefix as Cloud AI Prediction only allows the
# version name that starts with letters, and contains letters, digits,
# underscore only.
model_version = 'v{}'.format(int(time.time()))
executor_class_path = '%s.%s' % (self.__class__.__module__,
self.__class__.__name__)
with telemetry_utils.scoped_labels(
{telemetry_utils.LABEL_TFX_EXECUTOR: executor_class_path}):
job_labels = telemetry_utils.get_labels_dict()
endpoint = endpoint or runner.DEFAULT_ENDPOINT
api = discovery.build(
service_name,
api_version,
client_options=client_options.ClientOptions(api_endpoint=endpoint),
)
runner.deploy_model_for_aip_prediction(
api,
model_path,
model_version,
ai_platform_serving_args,
job_labels,
)
self._MarkPushed(
model_push,
pushed_destination=_CAIP_MODEL_VERSION_PATH_FORMAT.format(
project_id=ai_platform_serving_args['project_id'],
model=ai_platform_serving_args['model_name'],
version=model_version),
pushed_version=model_version)
def Do(self, input_dict: Dict[Text, List[types.Artifact]],
output_dict: Dict[Text, List[types.Artifact]],
exec_properties: Dict[Text, Any]) -> None:
"""Runs batch inference on a given model with given input examples.
This function creates a new model (if necessary) and a new model version
before inference, and cleans up resources after inference. It provides
re-executability as it cleans up (only) the model resources that are created
during the process even inference job failed.
Args:
input_dict: Input dict from input key to a list of Artifacts.
- examples: examples for inference.
- model: exported model.
- model_blessing: model blessing result
output_dict: Output dict from output key to a list of Artifacts.
- output: bulk inference results.
exec_properties: A dict of execution properties.
- data_spec: JSON string of bulk_inferrer_pb2.DataSpec instance.
- custom_config: custom_config.ai_platform_serving_args need to contain
the serving job parameters sent to Google Cloud AI Platform. For the
full set of parameters, refer to
https://cloud.google.com/ml-engine/reference/rest/v1/projects.models
Returns:
None
"""
self._log_startup(input_dict, output_dict, exec_properties)
if 'examples' not in input_dict:
raise ValueError('\'examples\' is missing in input dict.')
if 'inference_result' not in output_dict:
raise ValueError('\'inference_result\' is missing in output dict.')
output = artifact_utils.get_single_instance(output_dict['inference_result'])
if 'model' not in input_dict:
raise ValueError('Input models are not valid, model '
'need to be specified.')
if 'model_blessing' in input_dict:
model_blessing = artifact_utils.get_single_instance(
input_dict['model_blessing'])
if not model_utils.is_model_blessed(model_blessing):
output.set_int_custom_property('inferred', 0)
logging.info('Model on %s was not blessed', model_blessing.uri)
return
else:
logging.info('Model blessing is not provided, exported model will be '
'used.')
if _CUSTOM_CONFIG_KEY not in exec_properties:
raise ValueError('Input exec properties are not valid, {} '
'need to be specified.'.format(_CUSTOM_CONFIG_KEY))
custom_config = json_utils.loads(
exec_properties.get(_CUSTOM_CONFIG_KEY, 'null'))
if custom_config is not None and not isinstance(custom_config, Dict):
raise ValueError('custom_config in execution properties needs to be a '
'dict.')
ai_platform_serving_args = custom_config.get(SERVING_ARGS_KEY)
if not ai_platform_serving_args:
raise ValueError(
'\'ai_platform_serving_args\' is missing in \'custom_config\'')
service_name, api_version = runner.get_service_name_and_api_version(
ai_platform_serving_args)
executor_class_path = '%s.%s' % (self.__class__.__module__,
self.__class__.__name__)
with telemetry_utils.scoped_labels(
{telemetry_utils.LABEL_TFX_EXECUTOR: executor_class_path}):
job_labels = telemetry_utils.get_labels_dict()
model = artifact_utils.get_single_instance(input_dict['model'])
model_path = path_utils.serving_model_path(model.uri)
logging.info('Use exported model from %s.', model_path)
# Use model artifact uri to generate model version to guarantee the
# 1:1 mapping from model version to model.
model_version = 'version_' + hashlib.sha256(model.uri.encode()).hexdigest()
inference_spec = self._get_inference_spec(model_path, model_version,
ai_platform_serving_args)
data_spec = bulk_inferrer_pb2.DataSpec()
json_format.Parse(exec_properties['data_spec'], data_spec)
api = discovery.build(service_name, api_version)
new_model_created = False
try:
new_model_created = runner.create_model_for_aip_prediction_if_not_exist(
api, job_labels, ai_platform_serving_args)
runner.deploy_model_for_aip_prediction(
api,
model_path,
model_version,
ai_platform_serving_args,
job_labels,
skip_model_creation=True,
set_default_version=False,
)
self._run_model_inference(data_spec, input_dict['examples'], output.uri,
inference_spec)
except Exception as e:
logging.error('Error in executing CloudAIBulkInferrerComponent: %s',
str(e))
output.set_int_custom_property('inferred', 0)
raise
finally:
# Guarantee newly created resources are cleaned up even if theinference
# job failed.
# Clean up the newly deployed model.
runner.delete_model_version_from_aip_if_exists(api, model_version,
ai_platform_serving_args)
if new_model_created:
runner.delete_model_from_aip_if_exists(api, ai_platform_serving_args)
# Mark the inferenence as successful after resources are cleaned up.
output.set_int_custom_property('inferred', 1)
def Do(self, input_dict: Dict[str, List[types.Artifact]],
output_dict: Dict[str, List[types.Artifact]],
exec_properties: Dict[str, Any]):
"""Overrides the tfx_pusher_executor.
Args:
input_dict: Input dict from input key to a list of artifacts, including:
- model_export: exported model from trainer.
- model_blessing: model blessing path from evaluator.
output_dict: Output dict from key to a list of artifacts, including:
- model_push: A list of 'ModelPushPath' artifact of size one. It will
include the model in this push execution if the model was pushed.
exec_properties: Mostly a passthrough input dict for
tfx.components.Pusher.executor. The following keys in `custom_config`
are consumed by this class:
- ai_platform_serving_args: For the full set of parameters supported
by
- Google Cloud AI Platform, refer to
https://cloud.google.com/ml-engine/reference/rest/v1/projects.models.versions#Version.
- Google Cloud Vertex AI, refer to
https://googleapis.dev/python/aiplatform/latest/aiplatform.html?highlight=deploy#google.cloud.aiplatform.Model.deploy
- endpoint: Optional endpoint override.
- For Google Cloud AI Platform, this should be in format of
`https://[region]-ml.googleapis.com`. Default to global endpoint if
not set. Using regional endpoint is recommended by Cloud AI
Platform. When set, 'regions' key in ai_platform_serving_args cannot
be set. For more details, please see
https://cloud.google.com/ai-platform/prediction/docs/regional-endpoints#using_regional_endpoints
- For Google Cloud Vertex AI, this should be just be `region` (e.g.
'us-central1'). For available regions, please see
https://cloud.google.com/vertex-ai/docs/general/locations
Raises:
ValueError:
If ai_platform_serving_args is not in exec_properties.custom_config.
If Serving model path does not start with gs://.
If 'endpoint' and 'regions' are set simultaneously.
RuntimeError: if the Google Cloud AI Platform training job failed.
"""
self._log_startup(input_dict, output_dict, exec_properties)
custom_config = json_utils.loads(
exec_properties.get(_CUSTOM_CONFIG_KEY, 'null'))
if custom_config is not None and not isinstance(custom_config, Dict):
raise ValueError('custom_config in execution properties needs to be a '
'dict.')
ai_platform_serving_args = custom_config.get(constants.SERVING_ARGS_KEY)
if not ai_platform_serving_args:
raise ValueError(
'\'ai_platform_serving_args\' is missing in \'custom_config\'')
model_push = artifact_utils.get_single_instance(
output_dict[standard_component_specs.PUSHED_MODEL_KEY])
if not self.CheckBlessing(input_dict):
self._MarkNotPushed(model_push)
return
# Deploy the model.
io_utils.copy_dir(src=self.GetModelPath(input_dict), dst=model_push.uri)
model_path = model_push.uri
executor_class_path = '%s.%s' % (self.__class__.__module__,
self.__class__.__name__)
with telemetry_utils.scoped_labels(
{telemetry_utils.LABEL_TFX_EXECUTOR: executor_class_path}):
job_labels = telemetry_utils.make_labels_dict()
enable_vertex = custom_config.get(constants.ENABLE_VERTEX_KEY)
if enable_vertex:
if custom_config.get(constants.ENDPOINT_ARGS_KEY):
deprecation_utils.warn_deprecated(
'\'endpoint\' is deprecated. Please use'
'\'ai_platform_vertex_region\' instead.'
)
if 'regions' in ai_platform_serving_args:
deprecation_utils.warn_deprecated(
'\'ai_platform_serving_args.regions\' is deprecated. Please use'
'\'ai_platform_vertex_region\' instead.'
)
endpoint_region = custom_config.get(constants.VERTEX_REGION_KEY)
# TODO(jjong): Introduce Versioning.
# Note that we're adding "v" prefix as Cloud AI Prediction only allows the
# version name that starts with letters, and contains letters, digits,
# underscore only.
model_name = 'v{}'.format(int(time.time()))
container_image_uri = custom_config.get(
constants.VERTEX_CONTAINER_IMAGE_URI_KEY)
pushed_model_path = runner.deploy_model_for_aip_prediction(
serving_container_image_uri=container_image_uri,
model_version_name=model_name,
ai_platform_serving_args=ai_platform_serving_args,
endpoint_region=endpoint_region,
labels=job_labels,
serving_path=model_path,
enable_vertex=True,
)
self._MarkPushed(
model_push,
pushed_destination=pushed_model_path)
else:
endpoint = custom_config.get(constants.ENDPOINT_ARGS_KEY)
if endpoint and 'regions' in ai_platform_serving_args:
raise ValueError(
'\'endpoint\' and \'ai_platform_serving_args.regions\' cannot be set simultaneously'
)
# TODO(jjong): Introduce Versioning.
# Note that we're adding "v" prefix as Cloud AI Prediction only allows the
# version name that starts with letters, and contains letters, digits,
# underscore only.
model_version = 'v{}'.format(int(time.time()))
endpoint = endpoint or runner.DEFAULT_ENDPOINT
service_name, api_version = runner.get_service_name_and_api_version(
ai_platform_serving_args)
api = discovery.build(
service_name,
api_version,
requestBuilder=telemetry_utils.TFXHttpRequest,
client_options=client_options.ClientOptions(api_endpoint=endpoint),
)
pushed_model_version_path = runner.deploy_model_for_aip_prediction(
serving_path=model_path,
model_version_name=model_version,
ai_platform_serving_args=ai_platform_serving_args,
api=api,
labels=job_labels,
)
self._MarkPushed(
model_push,
pushed_destination=pushed_model_version_path,
pushed_version=model_version)
