Skip to content

Commit c0d01f1

Browse files
authored
feat: Add metadata SDK samples for list artifact and list execution (#1514)
* feat: Add metadata SDK samples for list artifact and execution * update default execution display name * Fix lint issues. * update tests based on review feedback to return a list of more than one item * simplify sample file names
1 parent d442248 commit c0d01f1

File tree

6 files changed

+158
-0
lines changed

6 files changed

+158
-0
lines changed

samples/model-builder/conftest.py

+18
Original file line numberDiff line numberDiff line change
@@ -590,6 +590,15 @@ def mock_create_execution(mock_execution):
590590
yield mock_create_execution
591591

592592

593+
@pytest.fixture
594+
def mock_list_execution(mock_execution):
595+
with patch.object(aiplatform.Execution, "list") as mock_list_execution:
596+
# Returning list of 2 executions to avoid confusion with get method
597+
# which returns one unique execution.
598+
mock_list_execution.return_value = [mock_execution, mock_execution]
599+
yield mock_list_execution
600+
601+
593602
@pytest.fixture
594603
def mock_get_artifact(mock_artifact):
595604
with patch.object(aiplatform, "Artifact") as mock_get_artifact:
@@ -625,6 +634,15 @@ def mock_create_artifact(mock_artifact):
625634
yield mock_create_artifact
626635

627636

637+
@pytest.fixture
638+
def mock_list_artifact(mock_artifact):
639+
with patch.object(aiplatform.Artifact, "list") as mock_list_artifact:
640+
# Returning list of 2 artifacts to avoid confusion with get method
641+
# which returns one unique artifact.
642+
mock_list_artifact.return_value = [mock_artifact, mock_artifact]
643+
yield mock_list_artifact
644+
645+
628646
@pytest.fixture
629647
def mock_start_run(mock_experiment_run):
630648
with patch.object(aiplatform, "start_run") as mock_start_run:
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
# Copyright 2022 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+
from typing import Optional
16+
17+
from google.cloud import aiplatform
18+
19+
20+
# [START aiplatform_sdk_create_artifact_with_sdk_sample]
21+
def list_artifact_sample(
22+
project: str,
23+
location: str,
24+
display_name_fitler: Optional[str] = "display_name=\"my_model_*\"",
25+
create_date_filter: Optional[str] = "create_time>\"2022-06-11T12:30:00-08:00\"",
26+
):
27+
aiplatform.init(
28+
project=project,
29+
location=location)
30+
31+
combined_filters = f"{display_name_fitler} AND {create_date_filter}"
32+
33+
return aiplatform.Artifact.list(filter=combined_filters)
34+
35+
36+
# [END aiplatform_sdk_create_artifact_with_sdk_sample]
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
# Copyright 2022 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 list_artifact_sample
16+
17+
import test_constants as constants
18+
19+
20+
def test_list_artifact_with_sdk_sample(mock_artifact, mock_list_artifact):
21+
artifacts = list_artifact_sample.list_artifact_sample(
22+
project=constants.PROJECT,
23+
location=constants.LOCATION,
24+
display_name_fitler=constants.DISPLAY_NAME,
25+
create_date_filter=constants.CREATE_DATE,
26+
)
27+
28+
mock_list_artifact.assert_called_with(
29+
filter=f"{constants.DISPLAY_NAME} AND {constants.CREATE_DATE}"
30+
)
31+
assert len(artifacts) == 2
32+
assert artifacts[0] is mock_artifact
33+
assert artifacts[1] is mock_artifact
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
# Copyright 2022 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+
from typing import Optional
16+
17+
from google.cloud import aiplatform
18+
19+
20+
# [START aiplatform_sdk_create_execution_with_sdk_sample]
21+
def list_execution_sample(
22+
project: str,
23+
location: str,
24+
display_name_fitler: Optional[str] = "display_name=\"my_execution_*\"",
25+
create_date_filter: Optional[str] = "create_time>\"2022-06-11T12:30:00-08:00\"",
26+
):
27+
aiplatform.init(
28+
project=project,
29+
location=location)
30+
31+
combined_filters = f"{display_name_fitler} AND {create_date_filter}"
32+
33+
return aiplatform.Execution.list(filter=combined_filters)
34+
35+
36+
# [END aiplatform_sdk_create_execution_with_sdk_sample]
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
# Copyright 2022 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 list_execution_sample
16+
17+
import test_constants as constants
18+
19+
20+
def test_list_execution_sample(mock_execution, mock_list_execution):
21+
executions = list_execution_sample.list_execution_sample(
22+
project=constants.PROJECT,
23+
location=constants.LOCATION,
24+
display_name_fitler=constants.DISPLAY_NAME,
25+
create_date_filter=constants.CREATE_DATE,
26+
)
27+
28+
mock_list_execution.assert_called_with(
29+
filter=f"{constants.DISPLAY_NAME} AND {constants.CREATE_DATE}"
30+
)
31+
assert len(executions) == 2
32+
assert executions[0] is mock_execution
33+
assert executions[1] is mock_execution

samples/model-builder/test_constants.py

+2
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,8 @@
2828
DISPLAY_NAME = str(uuid4()) # Create random display name
2929
DISPLAY_NAME_2 = str(uuid4())
3030

31+
CREATE_DATE = "2022-06-11T12:30:00-08:00"
32+
3133
STAGING_BUCKET = "gs://my-staging-bucket"
3234
EXPERIMENT_NAME = "fraud-detection-trial-72"
3335
CREDENTIALS = credentials.AnonymousCredentials()

0 commit comments

Comments
 (0)