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', } cmle_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': cmle_runner._get_tf_runtime_version(), 'python_version': cmle_runner._get_python_version(), }, body) mock_get.assert_called_with(name='op_name')
def Do(self, input_dict: Dict[Text, List[types.TfxArtifact]], output_dict: Dict[Text, List[types.TfxArtifact]], 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. RuntimeError: if the Google Cloud AI Platform training job failed. """ self._log_startup(input_dict, output_dict, exec_properties) if not self.CheckBlessing(input_dict, output_dict): return model_export = types.get_single_instance(input_dict['model_export']) model_export_uri = model_export.uri model_blessing_uri = types.get_single_uri(input_dict['model_blessing']) model_push = types.get_single_instance(output_dict['model_push']) # TODO(jyzhao): should this be in driver or executor. if not tf.gfile.Exists(os.path.join(model_blessing_uri, 'BLESSED')): model_push.set_int_custom_property('pushed', 0) tf.logging.info('Model on %s was not blessed', ) return exec_properties_copy = exec_properties.copy() custom_config = exec_properties_copy.pop('custom_config', {}) ai_platform_serving_args = custom_config['ai_platform_serving_args'] # 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) if ai_platform_serving_args is not None: cmle_runner.deploy_model_for_cmle_serving( model_path, model_version, ai_platform_serving_args) # Make sure artifacts are populated in a standard way by calling # tfx.pusher.executor.Executor.Do(). exec_properties_copy['push_destination'] = exec_properties.get( 'push_destination', self._make_local_temp_destination()) super(Executor, self).Do(input_dict, output_dict, exec_properties_copy)
def Do(self, input_dict: Dict[Text, List[types.TfxType]], output_dict: Dict[Text, List[types.TfxType]], exec_properties: Dict[Text, Any]) -> None: """Push model to target directory if blessed. 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. A push action delivers the model exports produced by Trainer to the destination defined in component config. 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: A dict of execution properties, including: - push_destination: JSON string of pusher_pb2.PushDestination instance, providing instruction of destination to push model. Returns: None """ self._log_startup(input_dict, output_dict, exec_properties) if not self.CheckBlessing(input_dict, output_dict): return model_push = types.get_single_instance(output_dict['model_push']) model_push_uri = model_push.uri model_export = types.get_single_instance(input_dict['model_export']) model_export_uri = model_export.uri tf.logging.info('Model pushing.') # Copy the model we are pushing into 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) tf.logging.info('Model version is %s', model_version) io_utils.copy_dir(model_path, os.path.join(model_push_uri, model_version)) tf.logging.info('Model written to %s.', model_push_uri) # Copied to a fixed outside path, which can be listened by model server. # # If model is already successfully copied to outside before, stop copying. # This is because model validator might blessed same model twice (check # mv driver) with different blessing output, we still want Pusher to # handle the mv output again to keep metadata tracking, but no need to # copy to outside path again.. # TODO(jyzhao): support rpc push and verification. push_destination = pusher_pb2.PushDestination() json_format.Parse(exec_properties['push_destination'], push_destination) serving_path = os.path.join(push_destination.filesystem.base_directory, model_version) if tf.gfile.Exists(serving_path): tf.logging.info( 'Destination directory %s already exists, skipping current push.', serving_path) else: # tf.serving won't load partial model, it will retry until fully copied. io_utils.copy_dir(model_path, serving_path) tf.logging.info('Model written to serving path %s.', serving_path) model_push.set_int_custom_property('pushed', 1) model_push.set_string_custom_property('pushed_model', model_export_uri) model_push.set_int_custom_property('pushed_model_id', model_export.id) tf.logging.info('Model pushed to %s.', serving_path) if exec_properties.get('custom_config'): cmle_serving_args = exec_properties.get( 'custom_config', {}).get('cmle_serving_args') if cmle_serving_args is not None: tf.logging.warn( '\'cmle_serving_args\' is deprecated, please use custom executor ' 'in tfx.extensions.google_cloud_ai_platform.pusher instead' ) return cmle_runner.deploy_model_for_cmle_serving( serving_path, model_version, cmle_serving_args)