Skip to content

Commit f0570cb

Browse files
vinnysenthilvinnysenthil
authored
fix: XAI Metadata compatibility with Model.upload (#705)
* Add tests to check XAI metadata builder and Model.upload compatibility * Use get_metadata_protobuf in test * Fix version mismatch bug, tests passing
1 parent f0b58b0 commit f0570cb

4 files changed

Lines changed: 64 additions & 13 deletions

File tree

‎google/cloud/aiplatform/explain/metadata/tf/v1/saved_model_metadata_builder.py‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,7 @@
1717
from google.protobuf import json_format
1818
from typing import Any, Dict, List, Optional
1919

20-
from google.cloud.aiplatform.compat.types import (
21-
explanation_metadata_v1beta1 as explanation_metadata,
22-
)
20+
from google.cloud.aiplatform.compat.types import explanation_metadata
2321
from google.cloud.aiplatform.explain.metadata import metadata_builder
2422

2523

‎google/cloud/aiplatform/explain/metadata/tf/v2/saved_model_metadata_builder.py‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,7 @@
1818
from typing import Optional, List, Dict, Any, Tuple
1919

2020
from google.cloud.aiplatform.explain.metadata import metadata_builder
21-
from google.cloud.aiplatform.compat.types import (
22-
explanation_metadata_v1beta1 as explanation_metadata,
23-
)
21+
from google.cloud.aiplatform.compat.types import explanation_metadata
2422

2523

2624
class SavedModelMetadataBuilder(metadata_builder.MetadataBuilder):

‎tests/unit/aiplatform/test_explain_saved_model_metadata_builder_tf1_test.py‎

Lines changed: 31 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,12 +15,15 @@
1515
# limitations under the License.
1616
#
1717

18+
import pytest
1819
import tensorflow.compat.v1 as tf
1920

21+
from google.cloud.aiplatform import models
2022
from google.cloud.aiplatform.explain.metadata.tf.v1 import saved_model_metadata_builder
21-
from google.cloud.aiplatform.compat.types import (
22-
explanation_metadata_v1beta1 as explanation_metadata,
23-
)
23+
from google.cloud.aiplatform.compat.types import explanation_metadata
24+
25+
import test_models
26+
from test_models import upload_model_mock, get_model_mock # noqa: F401
2427

2528

2629
class SavedModelMetadataBuilderTF1Test(tf.test.TestCase):
@@ -108,3 +111,28 @@ def test_get_metadata_protobuf_double_output(self):
108111
)
109112

110113
assert md_builder.get_metadata_protobuf() == expected_object
114+
115+
@pytest.mark.usefixtures("upload_model_mock", "get_model_mock")
116+
def test_model_upload_compatibility(self):
117+
self._set_up()
118+
md_builder = saved_model_metadata_builder.SavedModelMetadataBuilder(
119+
self.model_path, tags=[tf.saved_model.tag_constants.SERVING]
120+
)
121+
122+
generated_md = md_builder.get_metadata_protobuf()
123+
124+
try:
125+
models.Model.upload(
126+
display_name=test_models._TEST_MODEL_NAME,
127+
serving_container_image_uri=test_models._TEST_SERVING_CONTAINER_IMAGE,
128+
explanation_parameters=test_models._TEST_EXPLANATION_PARAMETERS,
129+
explanation_metadata=generated_md, # Test metadata from builder
130+
labels=test_models._TEST_LABEL,
131+
)
132+
except TypeError as e:
133+
if "Parameter to MergeFrom() must be instance of same class" in str(e):
134+
pytest.fail(
135+
f"Model.upload() expects different proto version, more info: {e}"
136+
)
137+
else:
138+
raise e

‎tests/unit/aiplatform/test_explain_saved_model_metadata_builder_tf2_test.py‎

Lines changed: 31 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,14 +15,16 @@
1515
# limitations under the License.
1616
#
1717

18-
18+
import pytest
1919
import tensorflow as tf
2020
import numpy as np
2121

22+
from google.cloud.aiplatform import models
2223
from google.cloud.aiplatform.explain.metadata.tf.v2 import saved_model_metadata_builder
23-
from google.cloud.aiplatform.compat.types import (
24-
explanation_metadata_v1beta1 as explanation_metadata,
25-
)
24+
from google.cloud.aiplatform.compat.types import explanation_metadata
25+
26+
import test_models
27+
from test_models import upload_model_mock, get_model_mock # noqa: F401
2628

2729

2830
class SavedModelMetadataBuilderTF2Test(tf.test.TestCase):
@@ -184,3 +186,28 @@ def test_model_with_feature_column(self):
184186
"outputs": {"output_1": {"outputTensorName": "output_1"}},
185187
}
186188
assert expected_md == generated_md
189+
190+
@pytest.mark.usefixtures("upload_model_mock", "get_model_mock")
191+
def test_model_upload_compatibility(self):
192+
self._set_up_sequential()
193+
194+
builder = saved_model_metadata_builder.SavedModelMetadataBuilder(
195+
self.saved_model_path
196+
)
197+
generated_md = builder.get_metadata_protobuf()
198+
199+
try:
200+
models.Model.upload(
201+
display_name=test_models._TEST_MODEL_NAME,
202+
serving_container_image_uri=test_models._TEST_SERVING_CONTAINER_IMAGE,
203+
explanation_parameters=test_models._TEST_EXPLANATION_PARAMETERS,
204+
explanation_metadata=generated_md, # Test metadata from builder
205+
labels=test_models._TEST_LABEL,
206+
)
207+
except TypeError as e:
208+
if "Parameter to MergeFrom() must be instance of same class" in str(e):
209+
pytest.fail(
210+
f"Model.upload() expects different proto version, more info: {e}"
211+
)
212+
else:
213+
raise e

0 commit comments

Comments
 (0)