def GetModelPath(self, input_dict: Dict[str, List[types.Artifact]]) -> str:
"""Get input model path to push.
Pusher can push various types of artifacts if it contains the model. This
method decides which artifact type is given to the Pusher and extracts the
real model path. Subclass of Pusher Executor should use this method to
acquire the source model path.
Args:
input_dict: A dictionary of artifacts that is given as the fisrt argument
to the Executor.Do() method.
Returns:
A resolved input model path.
Raises:
RuntimeError: If no model path is found from input_dict.
"""
# Check input_dict['model'] first.
models = input_dict.get(standard_component_specs.MODEL_KEY)
if models:
model = artifact_utils.get_single_instance(models)
return path_utils.serving_model_path(
model.uri, path_utils.is_old_model_artifact(model))
# Falls back to input_dict['infra_blessing']
blessed_models = input_dict.get(
standard_component_specs.INFRA_BLESSING_KEY)
if not blessed_models:
# Should not reach here; Pusher.__init__ prohibits creating a component
# without having any of model or infra_blessing inputs.
raise RuntimeError('Pusher has no model input.')
model = artifact_utils.get_single_instance(blessed_models)
if not model.get_int_custom_property(_INFRA_BLESSING_MODEL_FLAG_KEY):
raise RuntimeError('InfraBlessing does not contain a model. Check '
'request_spec.make_warmup is set to True.')
return path_utils.stamped_model_path(model.uri)
def _CreateWarmupModel(self, blessing: types.Artifact, model_path: str,
warmup_requests: List[iv_types.Request]):
output_model_path = path_utils.stamped_model_path(blessing.uri)
io_utils.copy_dir(src=model_path, dst=output_model_path)
io_utils.write_tfrecord_file(
path_utils.warmup_file_path(output_model_path),
*[_convert_to_prediction_log(r) for r in warmup_requests])
blessing.set_int_custom_property(_MODEL_FLAG_KEY, 1)
def testDo_WarmupNotCreatedWithoutRequests(self):
infra_validator = executor.Executor(self._context)
del self._exec_properties[REQUEST_SPEC_KEY] # No request
with mock.patch.object(infra_validator, '_ValidateOnce'):
infra_validator.Do(self._input_dict, self._output_dict,
self._exec_properties)
warmup_file = path_utils.warmup_file_path(
path_utils.stamped_model_path(self._blessing.uri))
self.assertFileDoesNotExist(warmup_file)
def testDo_MakeSavedModelWarmup(self):
infra_validator = executor.Executor(self._context)
self._request_spec.make_warmup = True
self._exec_properties[REQUEST_SPEC_KEY] = (
proto_utils.proto_to_json(self._request_spec))
with mock.patch.object(infra_validator, '_ValidateOnce'):
infra_validator.Do(self._input_dict, self._output_dict,
self._exec_properties)
warmup_file = path_utils.warmup_file_path(
path_utils.stamped_model_path(self._blessing.uri))
self.assertFileExists(warmup_file)
self.assertEqual(self._blessing.get_int_custom_property('has_model'), 1)
def testDo_InfraBlessingAsModel(self):
infra_blessing = standard_artifacts.InfraBlessing()
infra_blessing.uri = os.path.join(self._output_data_dir,
'infra_blessing')
infra_blessing.set_int_custom_property('blessed', True)
infra_blessing.set_int_custom_property('has_model', 1)
# Create dummy model
blessed_model_path = path_utils.stamped_model_path(infra_blessing.uri)
fileio.makedirs(blessed_model_path)
io_utils.write_string_file(
os.path.join(blessed_model_path, 'my-model'), '')
self._executor.Do({INFRA_BLESSING_KEY: [infra_blessing]},
self._output_dict, self._exec_properties)
self.assertPushed()
self.assertTrue(
fileio.exists(os.path.join(self._model_push.uri, 'my-model')))
def testStampedModelPath(self):
self.assertEqual(path_utils.stamped_model_path('/my-artifact'),
'/my-artifact/stamped_model')