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)