Skip to content

Commit 6c7e1fa

Browse files
sasha-gitgsasha-gitg
authored
chore: upgrade GAPIC Metadata from v1beta1 to v1 and add experiment logging integration test (#778)
1 parent 141c008 commit 6c7e1fa

File tree

13 files changed

+174
-53
lines changed

13 files changed

+174
-53
lines changed

‎google/cloud/aiplatform/compat/__init__.py‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,8 +40,10 @@
4040
types.accelerator_type = types.accelerator_type_v1beta1
4141
types.annotation = types.annotation_v1beta1
4242
types.annotation_spec = types.annotation_spec_v1beta1
43+
types.artifact = types.artifact_v1beta1
4344
types.batch_prediction_job = types.batch_prediction_job_v1beta1
4445
types.completion_stats = types.completion_stats_v1beta1
46+
types.context = types.context_v1beta1
4547
types.custom_job = types.custom_job_v1beta1
4648
types.data_item = types.data_item_v1beta1
4749
types.data_labeling_job = types.data_labeling_job_v1beta1
@@ -52,6 +54,8 @@
5254
types.endpoint = types.endpoint_v1beta1
5355
types.endpoint_service = types.endpoint_service_v1beta1
5456
types.env_var = types.env_var_v1beta1
57+
types.event = types.event_v1beta1
58+
types.execution = types.execution_v1beta1
5559
types.explanation = types.explanation_v1beta1
5660
types.explanation_metadata = types.explanation_metadata_v1beta1
5761
types.hyperparameter_tuning_job = types.hyperparameter_tuning_job_v1beta1
@@ -60,6 +64,8 @@
6064
types.job_state = types.job_state_v1beta1
6165
types.machine_resources = types.machine_resources_v1beta1
6266
types.manual_batch_tuning_parameters = types.manual_batch_tuning_parameters_v1beta1
67+
types.metadata_service = types.metadata_service_v1beta1
68+
types.metadata_store = types.metadata_store_v1beta1
6369
types.model = types.model_v1beta1
6470
types.model_evaluation = types.model_evaluation_v1beta1
6571
types.model_evaluation_slice = types.model_evaluation_slice_v1beta1
@@ -72,7 +78,6 @@
7278
types.specialist_pool_service = types.specialist_pool_service_v1beta1
7379
types.study = types.study_v1beta1
7480
types.training_pipeline = types.training_pipeline_v1beta1
75-
types.metadata_service = types.metadata_service_v1beta1
7681
types.tensorboard_service = types.tensorboard_service_v1beta1
7782
types.tensorboard_data = types.tensorboard_data_v1beta1
7883
types.tensorboard_experiment = types.tensorboard_experiment_v1beta1
@@ -93,8 +98,10 @@
9398
types.accelerator_type = types.accelerator_type_v1
9499
types.annotation = types.annotation_v1
95100
types.annotation_spec = types.annotation_spec_v1
101+
types.artifact = types.artifact_v1
96102
types.batch_prediction_job = types.batch_prediction_job_v1
97103
types.completion_stats = types.completion_stats_v1
104+
types.context = types.context_v1
98105
types.custom_job = types.custom_job_v1
99106
types.data_item = types.data_item_v1
100107
types.data_labeling_job = types.data_labeling_job_v1
@@ -105,6 +112,8 @@
105112
types.endpoint = types.endpoint_v1
106113
types.endpoint_service = types.endpoint_service_v1
107114
types.env_var = types.env_var_v1
115+
types.event = types.event_v1
116+
types.execution = types.execution_v1
108117
types.explanation = types.explanation_v1
109118
types.explanation_metadata = types.explanation_metadata_v1
110119
types.hyperparameter_tuning_job = types.hyperparameter_tuning_job_v1
@@ -113,6 +122,8 @@
113122
types.job_state = types.job_state_v1
114123
types.machine_resources = types.machine_resources_v1
115124
types.manual_batch_tuning_parameters = types.manual_batch_tuning_parameters_v1
125+
types.metadata_service = types.metadata_service_v1
126+
types.metadata_store = types.metadata_store_v1
116127
types.model = types.model_v1
117128
types.model_evaluation = types.model_evaluation_v1
118129
types.model_evaluation_slice = types.model_evaluation_slice_v1

‎google/cloud/aiplatform/compat/services/__init__.py‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,9 @@
2424
from google.cloud.aiplatform_v1beta1.services.job_service import (
2525
client as job_service_client_v1beta1,
2626
)
27+
from google.cloud.aiplatform_v1beta1.services.metadata_service import (
28+
client as metadata_service_client_v1beta1,
29+
)
2730
from google.cloud.aiplatform_v1beta1.services.model_service import (
2831
client as model_service_client_v1beta1,
2932
)
@@ -36,9 +39,6 @@
3639
from google.cloud.aiplatform_v1beta1.services.specialist_pool_service import (
3740
client as specialist_pool_service_client_v1beta1,
3841
)
39-
from google.cloud.aiplatform_v1beta1.services.metadata_service import (
40-
client as metadata_service_client_v1beta1,
41-
)
4242
from google.cloud.aiplatform_v1beta1.services.tensorboard_service import (
4343
client as tensorboard_service_client_v1beta1,
4444
)
@@ -52,6 +52,9 @@
5252
from google.cloud.aiplatform_v1.services.job_service import (
5353
client as job_service_client_v1,
5454
)
55+
from google.cloud.aiplatform_v1.services.metadata_service import (
56+
client as metadata_service_client_v1,
57+
)
5558
from google.cloud.aiplatform_v1.services.model_service import (
5659
client as model_service_client_v1,
5760
)
@@ -70,6 +73,7 @@
7073
dataset_service_client_v1,
7174
endpoint_service_client_v1,
7275
job_service_client_v1,
76+
metadata_service_client_v1,
7377
model_service_client_v1,
7478
pipeline_service_client_v1,
7579
prediction_service_client_v1,

‎google/cloud/aiplatform/compat/types/__init__.py‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,10 @@
1919
accelerator_type as accelerator_type_v1beta1,
2020
annotation as annotation_v1beta1,
2121
annotation_spec as annotation_spec_v1beta1,
22+
artifact as artifact_v1beta1,
2223
batch_prediction_job as batch_prediction_job_v1beta1,
2324
completion_stats as completion_stats_v1beta1,
25+
context as context_v1beta1,
2426
custom_job as custom_job_v1beta1,
2527
data_item as data_item_v1beta1,
2628
data_labeling_job as data_labeling_job_v1beta1,
@@ -31,6 +33,8 @@
3133
endpoint as endpoint_v1beta1,
3234
endpoint_service as endpoint_service_v1beta1,
3335
env_var as env_var_v1beta1,
36+
event as event_v1beta1,
37+
execution as execution_v1beta1,
3438
explanation as explanation_v1beta1,
3539
explanation_metadata as explanation_metadata_v1beta1,
3640
hyperparameter_tuning_job as hyperparameter_tuning_job_v1beta1,
@@ -39,6 +43,8 @@
3943
job_state as job_state_v1beta1,
4044
machine_resources as machine_resources_v1beta1,
4145
manual_batch_tuning_parameters as manual_batch_tuning_parameters_v1beta1,
46+
metadata_service as metadata_service_v1beta1,
47+
metadata_store as metadata_store_v1beta1,
4248
model as model_v1beta1,
4349
model_evaluation as model_evaluation_v1beta1,
4450
model_evaluation_slice as model_evaluation_slice_v1beta1,
@@ -52,7 +58,6 @@
5258
specialist_pool_service as specialist_pool_service_v1beta1,
5359
study as study_v1beta1,
5460
training_pipeline as training_pipeline_v1beta1,
55-
metadata_service as metadata_service_v1beta1,
5661
tensorboard as tensorboard_v1beta1,
5762
tensorboard_data as tensorboard_data_v1beta1,
5863
tensorboard_experiment as tensorboard_experiment_v1beta1,
@@ -64,8 +69,10 @@
6469
accelerator_type as accelerator_type_v1,
6570
annotation as annotation_v1,
6671
annotation_spec as annotation_spec_v1,
72+
artifact as artifact_v1,
6773
batch_prediction_job as batch_prediction_job_v1,
6874
completion_stats as completion_stats_v1,
75+
context as context_v1,
6976
custom_job as custom_job_v1,
7077
data_item as data_item_v1,
7178
data_labeling_job as data_labeling_job_v1,
@@ -76,6 +83,8 @@
7683
endpoint as endpoint_v1,
7784
endpoint_service as endpoint_service_v1,
7885
env_var as env_var_v1,
86+
event as event_v1,
87+
execution as execution_v1,
7988
explanation as explanation_v1,
8089
explanation_metadata as explanation_metadata_v1,
8190
hyperparameter_tuning_job as hyperparameter_tuning_job_v1,
@@ -84,6 +93,8 @@
8493
job_state as job_state_v1,
8594
machine_resources as machine_resources_v1,
8695
manual_batch_tuning_parameters as manual_batch_tuning_parameters_v1,
96+
metadata_service as metadata_service_v1,
97+
metadata_store as metadata_store_v1,
8798
model as model_v1,
8899
model_evaluation as model_evaluation_v1,
89100
model_evaluation_slice as model_evaluation_slice_v1,
@@ -103,8 +114,10 @@
103114
accelerator_type_v1,
104115
annotation_v1,
105116
annotation_spec_v1,
117+
artifact_v1,
106118
batch_prediction_job_v1,
107119
completion_stats_v1,
120+
context_v1,
108121
custom_job_v1,
109122
data_item_v1,
110123
data_labeling_job_v1,
@@ -115,6 +128,8 @@
115128
endpoint_v1,
116129
endpoint_service_v1,
117130
env_var_v1,
131+
event_v1,
132+
execution_v1,
118133
explanation_v1,
119134
explanation_metadata_v1,
120135
hyperparameter_tuning_job_v1,
@@ -123,6 +138,8 @@
123138
job_state_v1,
124139
machine_resources_v1,
125140
manual_batch_tuning_parameters_v1,
141+
metadata_service_v1,
142+
metadata_store_v1,
126143
model_v1,
127144
model_evaluation_v1,
128145
model_evaluation_slice_v1,
@@ -138,8 +155,10 @@
138155
accelerator_type_v1beta1,
139156
annotation_v1beta1,
140157
annotation_spec_v1beta1,
158+
artifact_v1beta1,
141159
batch_prediction_job_v1beta1,
142160
completion_stats_v1beta1,
161+
context_v1beta1,
143162
custom_job_v1beta1,
144163
data_item_v1beta1,
145164
data_labeling_job_v1beta1,
@@ -150,6 +169,8 @@
150169
endpoint_v1beta1,
151170
endpoint_service_v1beta1,
152171
env_var_v1beta1,
172+
event_v1beta1,
173+
execution_v1beta1,
153174
explanation_v1beta1,
154175
explanation_metadata_v1beta1,
155176
hyperparameter_tuning_job_v1beta1,
@@ -158,6 +179,8 @@
158179
job_state_v1beta1,
159180
machine_resources_v1beta1,
160181
manual_batch_tuning_parameters_v1beta1,
182+
metadata_service_v1beta1,
183+
metadata_store_v1beta1,
161184
model_v1beta1,
162185
model_evaluation_v1beta1,
163186
model_evaluation_slice_v1beta1,

‎google/cloud/aiplatform/metadata/artifact.py‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -20,16 +20,17 @@
2020
import proto
2121

2222
from google.cloud.aiplatform import utils
23-
from google.cloud.aiplatform.metadata.resource import _Resource
24-
from google.cloud.aiplatform_v1beta1 import ListArtifactsRequest
25-
from google.cloud.aiplatform_v1beta1.types import artifact as gca_artifact
23+
from google.cloud.aiplatform.compat.types import artifact as gca_artifact
24+
from google.cloud.aiplatform.compat.types import metadata_service
25+
from google.cloud.aiplatform.metadata import resource
2626

2727

28-
class _Artifact(_Resource):
28+
class _Artifact(resource._Resource):
2929
"""Metadata Artifact resource for Vertex AI"""
3030

3131
_resource_noun = "artifacts"
3232
_getter_method = "get_artifact"
33+
_delete_method = "delete_artifact"
3334

3435
@classmethod
3536
def _create_resource(
@@ -86,5 +87,7 @@ def _list_resources(
8687
filter (str):
8788
Optional. filter string to restrict the list result
8889
"""
89-
list_request = ListArtifactsRequest(parent=parent, filter=filter,)
90+
list_request = metadata_service.ListArtifactsRequest(
91+
parent=parent, filter=filter,
92+
)
9093
return client.list_artifacts(request=list_request)

‎google/cloud/aiplatform/metadata/context.py‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -20,16 +20,17 @@
2020
import proto
2121

2222
from google.cloud.aiplatform import utils
23-
from google.cloud.aiplatform.metadata.resource import _Resource
24-
from google.cloud.aiplatform_v1beta1 import ListContextsRequest
25-
from google.cloud.aiplatform_v1beta1.types import context as gca_context
23+
from google.cloud.aiplatform.compat.types import metadata_service
24+
from google.cloud.aiplatform.compat.types import context as gca_context
25+
from google.cloud.aiplatform.metadata import resource
2626

2727

28-
class _Context(_Resource):
28+
class _Context(resource._Resource):
2929
"""Metadata Context resource for Vertex AI"""
3030

3131
_resource_noun = "contexts"
3232
_getter_method = "get_context"
33+
_delete_method = "delete_context"
3334

3435
def add_artifacts_and_executions(
3536
self,
@@ -106,5 +107,7 @@ def _list_resources(
106107
Optional. filter string to restrict the list result
107108
"""
108109

109-
list_request = ListContextsRequest(parent=parent, filter=filter,)
110+
list_request = metadata_service.ListContextsRequest(
111+
parent=parent, filter=filter,
112+
)
110113
return client.list_contexts(request=list_request)

‎google/cloud/aiplatform/metadata/execution.py‎

Lines changed: 16 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -21,18 +21,19 @@
2121
from google.api_core import exceptions
2222

2323
from google.cloud.aiplatform import utils
24-
from google.cloud.aiplatform.metadata.artifact import _Artifact
25-
from google.cloud.aiplatform.metadata.resource import _Resource
26-
from google.cloud.aiplatform_v1beta1 import Event
27-
from google.cloud.aiplatform_v1beta1.types import execution as gca_execution
28-
from google.cloud.aiplatform_v1beta1.types.metadata_service import ListExecutionsRequest
24+
from google.cloud.aiplatform.compat.types import event as gca_event
25+
from google.cloud.aiplatform.compat.types import execution as gca_execution
26+
from google.cloud.aiplatform.compat.types import metadata_service
27+
from google.cloud.aiplatform.metadata import artifact
28+
from google.cloud.aiplatform.metadata import resource
2929

3030

31-
class _Execution(_Resource):
31+
class _Execution(resource._Resource):
3232
"""Metadata Execution resource for Vertex AI"""
3333

3434
_resource_noun = "executions"
3535
_getter_method = "get_execution"
36+
_delete_method = "delete_execution"
3637

3738
def add_artifact(
3839
self, artifact_resource_name: str, input: bool,
@@ -46,16 +47,16 @@ def add_artifact(
4647
Required. Whether Artifact is an input event to the Execution or not.
4748
"""
4849

49-
event = Event(
50+
event = gca_event.Event(
5051
artifact=artifact_resource_name,
51-
type_=Event.Type.INPUT if input else Event.Type.OUTPUT,
52+
type_=gca_event.Event.Type.INPUT if input else gca_event.Event.Type.OUTPUT,
5253
)
5354

5455
self.api_client.add_execution_events(
5556
execution=self.resource_name, events=[event],
5657
)
5758

58-
def query_input_and_output_artifacts(self) -> Sequence[_Artifact]:
59+
def query_input_and_output_artifacts(self) -> Sequence[artifact._Artifact]:
5960
"""query the input and output artifacts connected to the execution.
6061
6162
Returns:
@@ -70,13 +71,13 @@ def query_input_and_output_artifacts(self) -> Sequence[_Artifact]:
7071
return []
7172

7273
return [
73-
_Artifact(
74-
resource=artifact,
74+
artifact._Artifact(
75+
resource=metadata_artifact,
7576
project=self.project,
7677
location=self.location,
7778
credentials=self.credentials,
7879
)
79-
for artifact in artifacts
80+
for metadata_artifact in artifacts
8081
]
8182

8283
@classmethod
@@ -120,7 +121,9 @@ def _list_resources(
120121
Optional. filter string to restrict the list result
121122
"""
122123

123-
list_request = ListExecutionsRequest(parent=parent, filter=filter,)
124+
list_request = metadata_service.ListExecutionsRequest(
125+
parent=parent, filter=filter,
126+
)
124127
return client.list_executions(request=list_request)
125128

126129
@classmethod

‎google/cloud/aiplatform/metadata/metadata_store.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@
2424
from google.cloud.aiplatform import base, initializer
2525
from google.cloud.aiplatform import compat
2626
from google.cloud.aiplatform import utils
27-
from google.cloud.aiplatform_v1beta1.types import metadata_store as gca_metadata_store
27+
from google.cloud.aiplatform.compat.types import metadata_store as gca_metadata_store
2828

2929

3030
class _MetadataStore(base.VertexAiResourceNounWithFutureManager):
@@ -176,7 +176,7 @@ def _create(
176176
gapic_metadata_store = gca_metadata_store.MetadataStore(
177177
encryption_spec=initializer.global_config.get_encryption_spec(
178178
encryption_spec_key_name=encryption_spec_key_name,
179-
select_version=compat.V1BETA1,
179+
select_version=compat.DEFAULT_VERSION,
180180
)
181181
)
182182

0 commit comments

Comments
 (0)