def testLocalPredictionTensorflowModelWithStrings(self, mock_create_client): signature_def = meta_graph_pb2.SignatureDef() signature_def.outputs["x_bytes"].dtype = types_pb2.DT_STRING signature_def.inputs["x_bytes"].dtype = types_pb2.DT_STRING mock_client = mock.Mock() mock_client.signature_map = { tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature_def } mock_client.predict.return_value = {"x_bytes": "to_encode"} mock_create_client.return_value = mock_client predictions = mlprediction.local_predict(model_dir=None, instances=[{ "x_bytes": [1, 2, 3] }]) # Validate that the output is correctly base64 encoded (and only once) self.assertEquals(predictions, { "predictions": [{ "x_bytes": { "b64": base64.b64encode("to_encode") } }] })
def testModelWithBytesBasedOutput(self): mock_client = mock.Mock() mock_client.predict.return_value = {"x_bytes": "to_encode"} signature_def = meta_graph_pb2.SignatureDef() signature_def.outputs["x_bytes"].dtype = types_pb2.DT_STRING signature_def.inputs["input_key"].dtype = types_pb2.DT_STRING mock_client.signature_map = { tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature_def } model = mlprediction.create_model(mock_client, "gs://tmp/foo") _, predictions = model.predict({"input_key": "foo"}) self.assertEqual(predictions, [{ "x_bytes": { "b64": base64.b64encode("to_encode") } }])
def setUp(self): self._instances = [{"a": 1, "b": 2}, {"a": 2, "b": 4}] self._model_path = "gs://dummy/model/path" signature_def = meta_graph_pb2.SignatureDef() signature_def.inputs["a"].dtype = types_pb2.DT_INT32 signature_def.inputs["b"].dtype = types_pb2.DT_INT32 signature_def.outputs["c"].dtype = types_pb2.DT_INT32 self._mock_client = mock.Mock() self._mock_client.predict.return_value = {"c": np.array([10, 20])} self._mock_client.signature_map = { tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature_def } self._kwargs = { "signature_name": tf.saved_model.signature_constants. DEFAULT_SERVING_SIGNATURE_DEF_KEY }
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 testModelWithAdditionalOptions(self): # Setup the app. mock_model = mock.Mock() mock_model.predict.return_value = [], [] mock_model.signature_map = { "custom_signature": meta_graph_pb2.SignatureDef() } config = create_app_config(model=mock_model) inference_app = webapp2.WSGIApplication( [("/", server_lib._InferenceHandler)], debug=True, config=config) test_app = webtest.app.TestApp(app=inference_app) serialized_examples = [ example(1, 3).SerializeToString(), example(2, -4).SerializeToString(), example(0, 0).SerializeToString() ] # Act. instances = [ { "b64": base64.b64encode(serialized_examples[0]) }, { "b64": base64.b64encode(serialized_examples[1]) }, { "b64": base64.b64encode(serialized_examples[2]) }, ] body = {"instances": instances, "signature_name": "custom_signature"} test_app.post(url="/", params=json.dumps(body), content_type="application/json") # Assert. mock_model.predict.assert_has_calls([ mock.call(serialized_examples, stats=mock.ANY, signature_name="custom_signature") ])
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 testAppDecodesBytesIn(self): # mock_model = mock.create_autospec(server_lib.Model) mock_model = mock.Mock() signature_def = meta_graph_pb2.SignatureDef() signature_def.outputs["z_bytes"].dtype = types_pb2.DT_STRING mock_model.signature_map = { tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature_def } # Setup the app. config = create_app_config(model=mock_model) inference_app = webapp2.WSGIApplication( [("/", server_lib._InferenceHandler)], debug=True, config=config) test_app = webtest.app.TestApp(app=inference_app) instances = [{ u"x_bytes": { "b64": unicode(base64.b64encode("first")) }, u"y_bytes": { "b64": unicode(base64.b64encode("second")) } }] predictions = [{"z_bytes": "some binary string"}] mock_model.predict.return_value = instances, predictions body = {"instances": instances} response = test_app.post(url="/", params=json.dumps(body), content_type="application/json") # Assert. # Inputs are automatically decoded expected_instances = [{u"x_bytes": "first", u"y_bytes": "second"}] mock_model.predict.assert_has_calls( [mock.call(expected_instances, stats=mock.ANY)]) expected_predictions = {"predictions": predictions} self.assertEqual(response.body, json.dumps(expected_predictions))
def testModelWithOutputCannotJsonEncode(self): # Setup the app. mock_model = mock.Mock() mock_model.predict.return_value = [], [{"x": "\xe1"}] signature_def = meta_graph_pb2.SignatureDef() signature_def.outputs["x"].dtype = types_pb2.DT_STRING mock_model.signature_map = { tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature_def } config = create_app_config(model=mock_model) inference_app = webapp2.WSGIApplication( [("/", server_lib._InferenceHandler)], debug=True, config=config) test_app = webtest.app.TestApp(app=inference_app) # Act. body = {"instances": []} response = test_app.post(url="/", params=json.dumps(body), content_type="application/json") self.assertIn("Failed to json encode the prediction response", response) self.assertIn("suffix the alias of your output tensor with _bytes", response)