def testPredictRequestFromBatchedJson(self):
a_info = meta_graph_pb2.TensorInfo(name="a_unaliased",
dtype=types_pb2.DT_FLOAT)
b_info = meta_graph_pb2.TensorInfo(name="b_unaliased",
dtype=types_pb2.DT_STRING)
input_signature = {"a": a_info, "b": b_info}
batched_strings = {
"a": [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]],
"b": ["a", "c", "e"]
}
expected_predict_request = predict_pb2.PredictRequest(
inputs={
"a":
tensor_util.make_tensor_proto(
[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]],
input_signature["a"].dtype),
"b":
tensor_util.make_tensor_proto(["a", "c", "e"],
input_signature["b"].dtype)
})
expected_predict_request.model_spec.name = "default"
expected_predict_request.model_spec.signature_name = "signature"
predict_request = (tf_prediction_server_lib.create_predict_request(
batched_strings, input_signature, "signature"))
self.assertProtoEquals(predict_request, expected_predict_request)
def testSingleUtf8StringTensor(self):
tensor_info = meta_graph_pb2.TensorInfo(
dtype=tf.string.as_datatype_enum)
outputs_map = {"dummy": tensor_info}
actual = mlprediction.encode_base64([[[u"a", u"b"]], [[u"c", u"d"]]],
outputs_map)
self.assertEqual(actual, [[[u"a", u"b"]], [[u"c", u"d"]]])
def testCreateTFModelFromModelServerClient(self, mock_model_server):
env_updates = {"prediction_engine": "MODEL_SERVER"}
flag_values = {"tensorflow_session_parallelism": 3}
pseudo_flags = type("Flags", (object, ), flag_values)
# Create model's SignatureDef and expected response.
expected_response = get_model_metadata_pb2.GetModelMetadataResponse()
expected_response.model_spec.name = "default"
in_bytes = meta_graph_pb2.TensorInfo(name="x",
dtype=tf.string.as_datatype_enum)
out_bytes = meta_graph_pb2.TensorInfo(name="y",
dtype=tf.string.as_datatype_enum)
inputs = {"in_bytes": in_bytes}
outputs = {"out_bytes": out_bytes}
signatures_def = meta_graph_pb2.SignatureDef(inputs=inputs,
outputs=outputs)
signatures_def_map = get_model_metadata_pb2.SignatureDefMap()
signatures_def_map.signature_def["serving_default"].CopyFrom(
signatures_def)
expected_response.metadata["signature_def"].Pack(signatures_def_map)
mock_model_server.GetModelMetadata.return_value = expected_response
with mock.patch.dict("os.environ", env_updates):
with mock.patch.object(tf_prediction_server_lib,
"_start_model_server"):
tf_prediction_server_lib._start_model_server.return_value = (
None, mock_model_server)
model = tf_prediction_server_lib.create_tf_model(
"/dummy/model/path", pseudo_flags)
# model is a TensorflowModel instance with a ModelServer client.
expected_predict_response = make_response({"out_bytes": ["to encode"]})
mock_model_server.Predict.return_value = expected_predict_response
dummy_instances = []
_, predictions = model.predict(dummy_instances,
stats=mlprediction.Stats())
self.assertEqual(list(predictions), [{
"out_bytes": {
u"b64": base64.b64encode("to encode")
}
}])
def testGetModelSignatureMissingDtype(self, mock_model_server):
expected_response = get_model_metadata_pb2.GetModelMetadataResponse()
expected_response.model_spec.name = "default"
in_bytes = meta_graph_pb2.TensorInfo(name="x")
out_bytes = meta_graph_pb2.TensorInfo(name="y",
dtype=tf.string.as_datatype_enum)
inputs = {"in_bytes": in_bytes}
outputs = {"out_bytes": out_bytes}
signatures_def = meta_graph_pb2.SignatureDef(inputs=inputs,
outputs=outputs)
signatures_def_map = get_model_metadata_pb2.SignatureDefMap()
signatures_def_map.signature_def["serving_default"].CopyFrom(
signatures_def)
expected_response.metadata["signature_def"].Pack(signatures_def_map)
mock_model_server.GetModelMetadata.return_value = expected_response
received_signatures = tf_prediction_server_lib._get_model_signature_map(
mock_model_server)
self.assertTrue("serving_default" not in received_signatures)
def testPredictRequestWithInt64InputTensors(self):
a_info = meta_graph_pb2.TensorInfo(name="a_unaliased",
dtype=types_pb2.DT_INT64)
b_info = meta_graph_pb2.TensorInfo(name="b_unaliased",
dtype=types_pb2.DT_STRING)
input_signature = {"a": a_info, "b": b_info}
batched_strings = {"a": [[1, 2], [3, 4], [5, 6]], "b": ["a", "c", "e"]}
expected_predict_request = predict_pb2.PredictRequest(
inputs={
"a":
tensor_util.make_tensor_proto([[1, 2], [3, 4], [5, 6]],
input_signature["a"].dtype),
"b":
tensor_util.make_tensor_proto(["a", "c", "e"],
input_signature["b"].dtype)
})
expected_predict_request.model_spec.name = "default"
expected_predict_request.model_spec.signature_name = "signature"
predict_request = (tf_prediction_server_lib.create_predict_request(
batched_strings, input_signature, "signature"))
self.assertEqual(predict_request.inputs["a"].dtype, types_pb2.DT_INT64)
self.assertProtoEquals(predict_request, expected_predict_request)
def testMultiTensorWithUtf8Strings(self):
tensor_info_1 = meta_graph_pb2.TensorInfo(
dtype=tf.string.as_datatype_enum)
tensor_info_2 = meta_graph_pb2.TensorInfo(
dtype=tf.string.as_datatype_enum)
tensor_info_3 = meta_graph_pb2.TensorInfo(
dtype=tf.float32.as_datatype_enum)
outputs_map = {
"tensor1": tensor_info_1,
"tensor2": tensor_info_2,
"tensor3": tensor_info_3,
}
actual = mlprediction.encode_base64([{
u"tensor1": [[[u"a", u"b"]], [[u"c", u"d"]]],
u"tensor2": [u"x", u"y", u"z"],
u"tensor3": [1.0, -2.0, 3.14]
}], outputs_map)
self.assertEqual(actual, [{
u"tensor1": [[[u"a", u"b"]], [[u"c", u"d"]]],
u"tensor2": [u"x", u"y", u"z"],
u"tensor3": [1.0, -2.0, 3.14]
}])
def testSingleRank1BytesTensor(self):
tensor_info = meta_graph_pb2.TensorInfo(
dtype=tf.string.as_datatype_enum)
outputs_map = {"dummy_bytes": tensor_info}
actual = mlprediction.encode_base64([u"a", u"b", u"c"], outputs_map)
self.assertEqual(actual, [
{
u"b64": base64.b64encode(u"a")
},
{
u"b64": base64.b64encode(u"b")
},
{
u"b64": base64.b64encode(u"c")
},
])
def testPredictRequestUnexpectedTensor(self):
preprocessed_instances = [
example(1, 3).SerializeToString(),
example(2, -4).SerializeToString(),
example(0, 0).SerializeToString(),
]
input_signature = {
"a":
meta_graph_pb2.TensorInfo(name="a_unaliased",
dtype=types_pb2.DT_STRING)
}
batched_strings = {"a": preprocessed_instances, "b": [[1], [2], [3]]}
with self.assertRaises(mlprediction.PredictionError) as error:
tf_prediction_server_lib.create_predict_request(
batched_strings, input_signature, "signature")
self.assertTrue(
"Unexpected tensor name: b" in error.exception.error_detail)
def testCreatePredictRequestFromBatchedStrings(self, mock_model_server):
def features(x1, x2):
return tf.train.Features(
feature={
"x1":
tf.train.Feature(float_list=tf.train.FloatList(
value=[x1])),
"x2":
tf.train.Feature(float_list=tf.train.FloatList(value=[x2]))
})
batched_instances = {
"examples": [
tf.train.Example(features=features(1, 3)).SerializeToString(),
tf.train.Example(features=features(2, -4)).SerializeToString(),
tf.train.Example(features=features(0, 0)).SerializeToString()
]
}
input_signature = {
"examples":
meta_graph_pb2.TensorInfo(name="examples_unaliased",
dtype=types_pb2.DT_STRING)
}
expected_predict_request = predict_pb2.PredictRequest(
inputs={
"examples":
tensor_util.make_tensor_proto(
batched_instances["examples"],
input_signature["examples"].dtype)
})
expected_predict_request.model_spec.name = "default"
expected_predict_request.model_spec.signature_name = "signature"
predict_request = (tf_prediction_server_lib.create_predict_request(
batched_instances, input_signature, "signature"))
self.assertProtoEquals(expected_predict_request, predict_request)
def testPredictRequestFromBatchedStrings(self):
preprocessed_instances = [
example(1, 3).SerializeToString(),
example(2, -4).SerializeToString(),
example(0, 0).SerializeToString(),
]
input_signature = {
"a":
meta_graph_pb2.TensorInfo(name="a_unaliased",
dtype=types_pb2.DT_STRING)
}
batched_strings = {"a": preprocessed_instances}
expected_predict_request = predict_pb2.PredictRequest(
inputs={
"a":
tensor_util.make_tensor_proto(preprocessed_instances,
input_signature["a"].dtype)
})
expected_predict_request.model_spec.name = "default"
expected_predict_request.model_spec.signature_name = "signature"
predict_request = (tf_prediction_server_lib.create_predict_request(
batched_strings, input_signature, "signature"))
self.assertProtoEquals(predict_request, expected_predict_request)