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 testWarmupFilePath(self):
self.assertEqual(path_utils.warmup_file_path('/my-model'),
'/my-model/assets.extra/tf_serving_warmup_requests')
def _assertPushedModelHasWarmup(self):
pushed_model = self._get_latest_output_artifact('Pusher', 'pushed_model')
self._assertFileExists(path_utils.warmup_file_path(pushed_model.uri))