Skip to content

Commit 20b8859

Browse files
authored
fix: fix update export model sample, and add sample test (#239)
1 parent d24a9c2 commit 20b8859

File tree

3 files changed

+55
-2
lines changed

3 files changed

+55
-2
lines changed

.sample_configs/param_handlers/export_model_sample.py

+3-1
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,9 @@ def make_output_config(gcs_destination_output_uri_prefix: str) -> google.cloud.a
2323
output_config = {
2424
'artifact_destination': {
2525
'output_uri_prefix': gcs_destination_output_uri_prefix
26-
}
26+
},
27+
# For information about export formats: https://cloud.google.com/ai-platform-unified/docs/export/export-edge-model#aiplatform_export_model_sample-drest
28+
'export_format_id': 'tf-saved-model'
2729
}
2830

2931
return output_config

samples/snippets/export_model_sample.py

+5-1
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,11 @@ def export_model_sample(
3030
# This client only needs to be created once, and can be reused for multiple requests.
3131
client = aiplatform.gapic.ModelServiceClient(client_options=client_options)
3232
output_config = {
33-
"artifact_destination": {"output_uri_prefix": gcs_destination_output_uri_prefix}
33+
"artifact_destination": {
34+
"output_uri_prefix": gcs_destination_output_uri_prefix
35+
},
36+
# For information about export formats: https://cloud.google.com/ai-platform-unified/docs/export/export-edge-model#aiplatform_export_model_sample-drest
37+
"export_format_id": "tf-saved-model",
3438
}
3539
name = client.model_path(project=project, location=location, model=model_id)
3640
response = client.export_model(name=name, output_config=output_config)
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
# Copyright 2020 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import os
16+
17+
import pytest
18+
19+
import export_model_sample
20+
21+
PROJECT_ID = os.getenv("BUILD_SPECIFIC_GCLOUD_PROJECT")
22+
MODEL_ID = (
23+
"3422489426196955136" # permanent_swim_run_videos_action_recognition_edge_model
24+
)
25+
GCS_URI = (
26+
"gs://ucaip-samples-test-output/tmp/export_model_sample"
27+
)
28+
29+
30+
@pytest.fixture(scope="function", autouse=True)
31+
def teardown(storage_client):
32+
yield
33+
34+
bucket = storage_client.get_bucket("ucaip-samples-test-output")
35+
blobs = bucket.list_blobs(prefix="tmp/export_model_sample")
36+
for blob in blobs:
37+
blob.delete()
38+
39+
40+
def test_export_model_sample(capsys):
41+
export_model_sample.export_model_sample(
42+
project=PROJECT_ID,
43+
model_id=MODEL_ID,
44+
gcs_destination_output_uri_prefix=GCS_URI
45+
)
46+
out, _ = capsys.readouterr()
47+
assert "output_info" in out

0 commit comments

Comments
 (0)