Skip to content

Commit 28e56ef

Browse files
SinaChavoshiSinaChavoshi
authored
feat: Add support for SDK Method metrics tracking via _USER_AGENT_SDK… (#1591)
* feat: Add support for SDK Method metrics tracking via _USER_AGENT_SDK_COMMAND * add unit tests * refactor for readability * correct spelling
1 parent 8e92871 commit 28e56ef

File tree

10 files changed

+502
-15
lines changed

10 files changed

+502
-15
lines changed

‎google/cloud/aiplatform/base.py‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
from google.cloud.aiplatform import initializer
4747
from google.cloud.aiplatform import utils
4848
from google.cloud.aiplatform.compat.types import encryption_spec as gca_encryption_spec
49+
from google.cloud.aiplatform.constants import base as base_constants
4950
from google.protobuf import json_format
5051

5152
# This is the default retry callback to be used with get methods.
@@ -499,7 +500,19 @@ def __init__(
499500
self.location = location or initializer.global_config.location
500501
self.credentials = credentials or initializer.global_config.credentials
501502

502-
self.api_client = self._instantiate_client(self.location, self.credentials)
503+
appended_user_agent = None
504+
if base_constants.USER_AGENT_SDK_COMMAND:
505+
appended_user_agent = [
506+
f"sdk_command/{base_constants.USER_AGENT_SDK_COMMAND}"
507+
]
508+
# Reset the value for the USER_AGENT_SDK_COMMAND to avoid counting future unrelated api calls.
509+
base_constants.USER_AGENT_SDK_COMMAND = ""
510+
511+
self.api_client = self._instantiate_client(
512+
location=self.location,
513+
credentials=self.credentials,
514+
appended_user_agent=appended_user_agent,
515+
)
503516

504517
@classmethod
505518
def _instantiate_client(

‎google/cloud/aiplatform/constants/base.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,3 +88,7 @@
8888

8989
# Used in constructing the requests user_agent header for metrics reporting.
9090
USER_AGENT_PRODUCT = "model-builder"
91+
# This field is used to pass the name of the specific SDK method
92+
# that is being used for usage metrics tracking purposes.
93+
# For more details on go/oneplatform-api-analytics
94+
USER_AGENT_SDK_COMMAND = ""

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

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
from google.cloud.aiplatform.compat.types import (
2929
metadata_service as gca_metadata_service,
3030
)
31+
from google.cloud.aiplatform.constants import base as base_constants
3132
from google.cloud.aiplatform.metadata import metadata_store
3233
from google.cloud.aiplatform.metadata import resource
3334
from google.cloud.aiplatform.metadata import utils as metadata_utils
@@ -114,6 +115,7 @@ def _create_resource(
114115
artifact_id=resource_id,
115116
)
116117

118+
# TODO() refactor code to move _create to _Resource class.
117119
@classmethod
118120
def _create(
119121
cls,
@@ -175,7 +177,19 @@ def _create(
175177
Instantiated representation of the managed Metadata resource.
176178
177179
"""
178-
api_client = cls._instantiate_client(location=location, credentials=credentials)
180+
appended_user_agent = []
181+
if base_constants.USER_AGENT_SDK_COMMAND:
182+
appended_user_agent = [
183+
f"sdk_command/{base_constants.USER_AGENT_SDK_COMMAND}"
184+
]
185+
# Reset the value for the USER_AGENT_SDK_COMMAND to avoid counting future unrelated api calls.
186+
base_constants.USER_AGENT_SDK_COMMAND = ""
187+
188+
api_client = cls._instantiate_client(
189+
location=location,
190+
credentials=credentials,
191+
appended_user_agent=appended_user_agent,
192+
)
179193

180194
parent = utils.full_resource_name(
181195
resource_name=metadata_store_id,
@@ -311,6 +325,13 @@ def create(
311325
Returns:
312326
Artifact: Instantiated representation of the managed Metadata Artifact.
313327
"""
328+
# Add User Agent Header for metrics tracking if one is not specified
329+
# If one is already specified this call was initiated by a sub class.
330+
if not base_constants.USER_AGENT_SDK_COMMAND:
331+
base_constants.USER_AGENT_SDK_COMMAND = (
332+
"aiplatform.metadata.artifact.Artifact.create"
333+
)
334+
314335
return cls._create(
315336
resource_id=resource_id,
316337
schema_title=schema_title,

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

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323

2424
from google.cloud.aiplatform import base
2525
from google.cloud.aiplatform import utils
26+
from google.cloud.aiplatform.constants import base as base_constants
2627
from google.cloud.aiplatform.metadata import utils as metadata_utils
2728
from google.cloud.aiplatform.compat.types import context as gca_context
2829
from google.cloud.aiplatform.compat.types import (
@@ -136,6 +137,13 @@ def create(
136137
Returns:
137138
Context: Instantiated representation of the managed Metadata Context.
138139
"""
140+
# Add User Agent Header for metrics tracking if one is not specified
141+
# If one is already specified this call was initiated by a sub class.
142+
if not base_constants.USER_AGENT_SDK_COMMAND:
143+
base_constants.USER_AGENT_SDK_COMMAND = (
144+
"aiplatform.metadata.context.Context.create"
145+
)
146+
139147
return cls._create(
140148
resource_id=resource_id,
141149
schema_title=schema_title,
@@ -202,7 +210,19 @@ def _create(
202210
Instantiated representation of the managed Metadata resource.
203211
204212
"""
205-
api_client = cls._instantiate_client(location=location, credentials=credentials)
213+
appended_user_agent = []
214+
if base_constants.USER_AGENT_SDK_COMMAND:
215+
appended_user_agent = [
216+
f"sdk_command/{base_constants.USER_AGENT_SDK_COMMAND}"
217+
]
218+
# Reset the value for the USER_AGENT_SDK_COMMAND to avoid counting future unrelated api calls.
219+
base_constants.USER_AGENT_SDK_COMMAND = ""
220+
221+
api_client = cls._instantiate_client(
222+
location=location,
223+
credentials=credentials,
224+
appended_user_agent=appended_user_agent,
225+
)
206226

207227
parent = utils.full_resource_name(
208228
resource_name=metadata_store_id,

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

Lines changed: 105 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -20,14 +20,14 @@
2020
import proto
2121
from google.auth import credentials as auth_credentials
2222

23-
from google.cloud.aiplatform import base
2423
from google.cloud.aiplatform import models
2524
from google.cloud.aiplatform import utils
2625
from google.cloud.aiplatform.compat.types import event as gca_event
2726
from google.cloud.aiplatform.compat.types import execution as gca_execution
2827
from google.cloud.aiplatform.compat.types import (
2928
metadata_service as gca_metadata_service,
3029
)
30+
from google.cloud.aiplatform.constants import base as base_constants
3131
from google.cloud.aiplatform.metadata import artifact
3232
from google.cloud.aiplatform.metadata import metadata_store
3333
from google.cloud.aiplatform.metadata import resource
@@ -142,18 +142,110 @@ def create(
142142
Execution: Instantiated representation of the managed Metadata Execution.
143143
144144
"""
145-
self = cls._empty_constructor(
146-
project=project, location=location, credentials=credentials
145+
# Add User Agent Header for metrics tracking if one is not specified
146+
# If one is already specified this call was initiated by a sub class.
147+
if not base_constants.USER_AGENT_SDK_COMMAND:
148+
base_constants.USER_AGENT_SDK_COMMAND = (
149+
"aiplatform.metadata.execution.Execution.create"
150+
)
151+
152+
return cls._create(
153+
resource_id=resource_id,
154+
schema_title=schema_title,
155+
display_name=display_name,
156+
schema_version=schema_version,
157+
description=description,
158+
metadata=metadata,
159+
state=state,
160+
metadata_store_id=metadata_store_id,
161+
project=project,
162+
location=location,
163+
credentials=credentials,
164+
)
165+
166+
# TODO() refactor code to move _create to _Resource class.
167+
@classmethod
168+
def _create(
169+
cls,
170+
schema_title: str,
171+
*,
172+
state: gca_execution.Execution.State = gca_execution.Execution.State.RUNNING,
173+
resource_id: Optional[str] = None,
174+
display_name: Optional[str] = None,
175+
schema_version: Optional[str] = None,
176+
metadata: Optional[Dict[str, Any]] = None,
177+
description: Optional[str] = None,
178+
metadata_store_id: str = "default",
179+
project: Optional[str] = None,
180+
location: Optional[str] = None,
181+
credentials=Optional[auth_credentials.Credentials],
182+
) -> "Execution":
183+
"""
184+
Creates a new Metadata Execution.
185+
186+
Args:
187+
schema_title (str):
188+
Required. schema_title identifies the schema title used by the Execution.
189+
state (gca_execution.Execution.State.RUNNING):
190+
Optional. State of this Execution. Defaults to RUNNING.
191+
resource_id (str):
192+
Optional. The <resource_id> portion of the Execution name with
193+
the format. This is globally unique in a metadataStore:
194+
projects/123/locations/us-central1/metadataStores/<metadata_store_id>/executions/<resource_id>.
195+
display_name (str):
196+
Optional. The user-defined name of the Execution.
197+
schema_version (str):
198+
Optional. schema_version specifies the version used by the Execution.
199+
If not set, defaults to use the latest version.
200+
metadata (Dict):
201+
Optional. Contains the metadata information that will be stored in the Execution.
202+
description (str):
203+
Optional. Describes the purpose of the Execution to be created.
204+
metadata_store_id (str):
205+
Optional. The <metadata_store_id> portion of the resource name with
206+
the format:
207+
projects/123/locations/us-central1/metadataStores/<metadata_store_id>/artifacts/<resource_id>
208+
If not provided, the MetadataStore's ID will be set to "default".
209+
project (str):
210+
Optional. Project used to create this Execution. Overrides project set in
211+
aiplatform.init.
212+
location (str):
213+
Optional. Location used to create this Execution. Overrides location set in
214+
aiplatform.init.
215+
credentials (auth_credentials.Credentials):
216+
Optional. Custom credentials used to create this Execution. Overrides
217+
credentials set in aiplatform.init.
218+
219+
Returns:
220+
Execution: Instantiated representation of the managed Metadata Execution.
221+
222+
"""
223+
appended_user_agent = []
224+
if base_constants.USER_AGENT_SDK_COMMAND:
225+
appended_user_agent = [
226+
f"sdk_command/{base_constants.USER_AGENT_SDK_COMMAND}"
227+
]
228+
# Reset the value for the USER_AGENT_SDK_COMMAND to avoid counting future unrelated api calls.
229+
base_constants.USER_AGENT_SDK_COMMAND = ""
230+
231+
api_client = cls._instantiate_client(
232+
location=location,
233+
credentials=credentials,
234+
appended_user_agent=appended_user_agent,
235+
)
236+
237+
parent = utils.full_resource_name(
238+
resource_name=metadata_store_id,
239+
resource_noun=metadata_store._MetadataStore._resource_noun,
240+
parse_resource_name_method=metadata_store._MetadataStore._parse_resource_name,
241+
format_resource_name_method=metadata_store._MetadataStore._format_resource_name,
242+
project=project,
243+
location=location,
147244
)
148-
super(base.VertexAiResourceNounWithFutureManager, self).__init__()
149245

150246
resource = Execution._create_resource(
151-
client=self.api_client,
152-
parent=metadata_store._MetadataStore._format_resource_name(
153-
project=self.project,
154-
location=self.location,
155-
metadata_store=metadata_store_id,
156-
),
247+
client=api_client,
248+
parent=parent,
157249
schema_title=schema_title,
158250
resource_id=resource_id,
159251
metadata=metadata,
@@ -162,6 +254,9 @@ def create(
162254
schema_version=schema_version,
163255
state=state,
164256
)
257+
self = cls._empty_constructor(
258+
project=project, location=location, credentials=credentials
259+
)
165260
self._gca_resource = resource
166261

167262
return self

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

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
from google.cloud.aiplatform import compat
2626
from google.cloud.aiplatform import utils
2727
from google.cloud.aiplatform.compat.types import metadata_store as gca_metadata_store
28+
from google.cloud.aiplatform.constants import base as base_constants
2829

2930

3031
class _MetadataStore(base.VertexAiResourceNounWithFutureManager):
@@ -115,7 +116,6 @@ def get_or_create(
115116
Instantiated representation of the managed metadata store resource.
116117
117118
"""
118-
119119
store = cls._get(
120120
metadata_store_name=metadata_store_id,
121121
project=project,
@@ -176,7 +176,20 @@ def _create(
176176
Instantiated representation of the managed metadata store resource.
177177
178178
"""
179-
api_client = cls._instantiate_client(location=location, credentials=credentials)
179+
appended_user_agent = []
180+
if base_constants.USER_AGENT_SDK_COMMAND:
181+
appended_user_agent = [
182+
f"sdk_command/{base_constants.USER_AGENT_SDK_COMMAND}"
183+
]
184+
# Reset the value for the USER_AGENT_SDK_COMMAND to avoid counting future unrelated api calls.
185+
base_constants.USER_AGENT_SDK_COMMAND = ""
186+
187+
api_client = cls._instantiate_client(
188+
location=location,
189+
credentials=credentials,
190+
appended_user_agent=appended_user_agent,
191+
)
192+
180193
gapic_metadata_store = gca_metadata_store.MetadataStore(
181194
encryption_spec=initializer.global_config.get_encryption_spec(
182195
encryption_spec_key_name=encryption_spec_key_name,

‎google/cloud/aiplatform/metadata/schema/base_artifact.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323

2424
from google.cloud.aiplatform.compat.types import artifact as gca_artifact
2525
from google.cloud.aiplatform.metadata import artifact
26+
from google.cloud.aiplatform.constants import base as base_constants
2627
from google.cloud.aiplatform.metadata import constants
2728

2829

@@ -114,6 +115,11 @@ def _init_with_resource_name(
114115
Artifact name with the following format, this is globally unique in a metadataStore:
115116
projects/123/locations/us-central1/metadataStores/<metadata_store_id>/artifacts/<resource_id>.
116117
"""
118+
# Add User Agent Header for metrics tracking if one is not specified
119+
# If one is already specified this call was initiated by a sub class.
120+
if not base_constants.USER_AGENT_SDK_COMMAND:
121+
base_constants.USER_AGENT_SDK_COMMAND = "aiplatform.metadata.schema.base_artifact.BaseArtifactSchema._init_with_resource_name"
122+
117123
super(BaseArtifactSchema, self).__init__(artifact_name=artifact_name)
118124

119125
def create(
@@ -144,6 +150,10 @@ def create(
144150
Returns:
145151
Artifact: Instantiated representation of the managed Metadata Artifact.
146152
"""
153+
# Add User Agent Header for metrics tracking.
154+
base_constants.USER_AGENT_SDK_COMMAND = (
155+
"aiplatform.metadata.schema.base_artifact.BaseArtifactSchema.create"
156+
)
147157

148158
# Check if metadata exists to avoid proto read error
149159
metadata = None

‎google/cloud/aiplatform/metadata/schema/base_context.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
from google.auth import credentials as auth_credentials
2323

2424
from google.cloud.aiplatform.compat.types import context as gca_context
25+
from google.cloud.aiplatform.constants import base as base_constants
2526
from google.cloud.aiplatform.metadata import constants
2627
from google.cloud.aiplatform.metadata import context
2728

@@ -91,6 +92,11 @@ def _init_with_resource_name(
9192
Context name with the following format, this is globally unique in a metadataStore:
9293
projects/123/locations/us-central1/metadataStores/<metadata_store_id>/contexts/<resource_id>.
9394
"""
95+
# Add User Agent Header for metrics tracking if one is not specified
96+
# If one is already specified this call was initiated by a sub class.
97+
if not base_constants.USER_AGENT_SDK_COMMAND:
98+
base_constants.USER_AGENT_SDK_COMMAND = "aiplatform.metadata.schema.base_context.BaseContextSchema._init_with_resource_name"
99+
94100
super(BaseContextSchema, self).__init__(resource_name=context_name)
95101

96102
def create(
@@ -122,6 +128,11 @@ def create(
122128
Context: Instantiated representation of the managed Metadata Context.
123129
124130
"""
131+
# Add User Agent Header for metrics tracking.
132+
base_constants.USER_AGENT_SDK_COMMAND = (
133+
"aiplatform.metadata.schema.base_context.BaseContextSchema.create"
134+
)
135+
125136
# Check if metadata exists to avoid proto read error
126137
metadata = None
127138
if self._gca_resource.metadata:

0 commit comments

Comments
 (0)