def __init__(self, model_dir, skip_preprocessing):
self.model_dir = model_dir
session, signature = mlprediction.load_model(model_dir)
client = mlprediction.SessionClient(session, signature)
self.model = mlprediction.create_model(
client, model_dir, skip_preprocessing=skip_preprocessing)
def testUserPostprocessOnly(self):
create_version_json = """
{
"version": {
"processor_class": "google.cloud.ml.prediction.testdata.user_custom_python.user_model.TfUserPostprocessor"
}
}
"""
env_map = {"create_version_request": create_version_json}
with mock.patch.dict("os.environ", env_map):
model = mlprediction.create_model(self._mock_client,
self._model_path)
# Verify the default TensorFlowModel is instantiated.
self.assertIsInstance(model, mlprediction.TensorFlowModel)
# Verify postprocessing (which divides values by 2) is applied to the
# predicted results.
self.assertEqual(model.predict(self._instances, **self._kwargs),
(self._instances, [{
"c": 5
}, {
"c": 10
}]))
# Verify no preprocessing performed.
self._mock_client.predict.assert_has_calls([
mock.call({
"a": [1, 2],
"b": [2, 4]
},
stats=mock.ANY,
signature_name=mock.ANY)
])
def __init__(self,
model_dir,
tags,
framework=mlprediction.TENSORFLOW_FRAMEWORK_NAME):
self.model_dir = model_dir
client = mlprediction.create_client(framework, model_dir, tags)
self.model = mlprediction.create_model(client, model_dir, framework)
def testMissingPredictAndProcessMethod(self):
create_version_json = """
{
"version": {
"model_class": "google.cloud.ml.prediction.testdata.user_custom_python.user_model.MissingPredict"
}
}
"""
env_map = {"create_version_request": create_version_json}
with mock.patch.dict("os.environ", env_map):
with self.assertRaises(mlprediction.PredictionError) as error:
client = None
dummy_model_path = "gs://dummy/model/path"
mlprediction.create_model(client, dummy_model_path)
self.assertEqual(
error.exception.error_detail,
("The provided model class, MissingPredict, is missing "
"the required predict method."))
def testUserModelTooFewPredictArgs(self):
create_version_json = """
{
"version": {
"model_class": "google.cloud.ml.prediction.testdata.user_custom_python.user_model.PredictMethodTooFewArgs"
}
}
"""
env_map = {"create_version_request": create_version_json}
with mock.patch.dict("os.environ", env_map):
with self.assertRaises(mlprediction.PredictionError) as error:
client = None
dummy_model_path = "gs://dummy/model/path"
mlprediction.create_model(client, dummy_model_path)
self.assertEqual(
error.exception.error_detail,
("The provided model class, PredictMethodTooFewArgs, has "
"a predict method with an invalid signature. "
"Expected signature: ['self', 'instances'] "
"User signature: ['self']"))
def testMissingUserModelClassModule(self):
create_version_json = """
{
"version": {
"model_class": "wrong_module.UserModel",
"package_uris": ["gs://test_package"]
}
}
"""
env_map = {"create_version_request": create_version_json}
with mock.patch.dict("os.environ", env_map):
with self.assertRaises(mlprediction.PredictionError) as error:
client = None
dummy_model_path = "gs://dummy/model/path"
mlprediction.create_model(client, dummy_model_path)
self.assertEqual(
error.exception.error_detail,
"wrong_module.UserModel cannot be found. "
"Please make sure (1) model_class is the fully qualified "
"function name, and (2) model_class uses the correct "
"package name as provided by the package_uris: "
"['gs://test_package']")
def testUserModelMissingCustomMethod(self):
create_version_json = """
{
"version": {
"model_class": "google.cloud.ml.prediction.testdata.user_custom_python.wrong_user_model",
"package_uris": ["gs://test_package"]
}
}
"""
env_map = {"create_version_request": create_version_json}
with mock.patch.dict("os.environ", env_map):
with self.assertRaises(mlprediction.PredictionError) as error:
client = None
dummy_model_path = "gs://dummy/model/path"
mlprediction.create_model(client, dummy_model_path)
self.assertEqual(
error.exception.error_detail,
("google.cloud.ml.prediction.testdata.user_custom_python."
"wrong_user_model cannot be found. Please make "
"sure (1) model_class is the fully qualified function "
"name, and (2) model_class uses the correct package name "
"as provided by the package_uris: ['gs://test_package']"))
def testModelCreationNoCustomCode(self):
create_version_json = """
{
"version": {}
}
"""
env_map = {"create_version_request": create_version_json}
with mock.patch.dict("os.environ", env_map):
client = mock.Mock()
client.signature_map = {"serving_default": None}
dummy_model_path = "gs://dummy/model/path"
model = mlprediction.create_model(client, dummy_model_path)
self.assertIsInstance(model, mlprediction.TensorFlowModel)
def testUserModelCreationFromClassMethod(self):
create_version_json = """
{
"version": {
"model_class": "google.cloud.ml.prediction.testdata.user_custom_python.user_model.UserModel"
}
}
"""
env_map = {"create_version_request": create_version_json}
with mock.patch.dict("os.environ", env_map):
client = None
dummy_model_path = "gs://dummy/model/path"
model = mlprediction.create_model(client, dummy_model_path)
self.assertIsInstance(model, user_model.UserModel)
def testNoUserProcessor(self):
model = mlprediction.create_model(self._mock_client, self._model_path)
self.assertIsInstance(model, mlprediction.TensorFlowModel)
self.assertEqual(model.predict(self._instances, **self._kwargs),
(self._instances, [{
"c": 10
}, {
"c": 20
}]))
self._mock_client.predict.assert_has_calls([
mock.call({
"a": [1, 2],
"b": [2, 4]
},
stats=mock.ANY,
signature_name=mock.ANY)
])
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 testTensorFlowModelCreationNoCreateVersionRequest(self):
client = mock.Mock()
client.signature_map = {"serving_default": None}
dummy_model_path = "gs://dummy/model/path"
model = mlprediction.create_model(client, dummy_model_path)
self.assertIsInstance(model, mlprediction.TensorFlowModel)
