Skip to content

Commit c01df66

Browse files
yeesiancopybara-github
authored andcommitted
feat: Allow module_name to be specified with the base logger.
PiperOrigin-RevId: 612860933
1 parent 0b33bdd commit c01df66

File tree

1 file changed

+17
-4
lines changed

1 file changed

+17
-4
lines changed

google/cloud/aiplatform/base.py

+17-4
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,8 @@ def log_create_complete(
9292
cls: Type["VertexAiResourceNoun"],
9393
resource: proto.Message,
9494
variable_name: str,
95+
*,
96+
module_name: str = "aiplatform",
9597
):
9698
"""Logs create event is complete.
9799
@@ -102,17 +104,23 @@ def log_create_complete(
102104
Vertex AI Resource Noun class that is being created.
103105
resource (proto.Message):
104106
Vertex AI Resource proto.Message
105-
variable_name (str): Name of variable to use for code snippet
107+
variable_name (str):
108+
Name of variable to use for code snippet.
109+
module_name (str):
110+
The module namespace under which the Vertex AI Resource Noun
111+
is available. Defaults to `aiplatform`.
106112
"""
107113
self.info(f"{cls.__name__} created. Resource name: {resource.name}")
108114
self.info(f"To use this {cls.__name__} in another session:")
109-
self.info(f"{variable_name} = aiplatform.{cls.__name__}('{resource.name}')")
115+
self.info(f"{variable_name} = {module_name}.{cls.__name__}('{resource.name}')")
110116

111117
def log_create_complete_with_getter(
112118
self,
113119
cls: Type["VertexAiResourceNoun"],
114120
resource: proto.Message,
115121
variable_name: str,
122+
*,
123+
module_name: str = "aiplatform",
116124
):
117125
"""Logs create event is complete.
118126
@@ -123,11 +131,16 @@ def log_create_complete_with_getter(
123131
Vertex AI Resource Noun class that is being created.
124132
resource (proto.Message):
125133
Vertex AI Resource proto.Message
126-
variable_name (str): Name of variable to use for code snippet
134+
variable_name (str):
135+
Name of variable to use for code snippet.
136+
module_name (str):
137+
The module namespace under which the Vertex AI Resource Noun
138+
is available. Defaults to `aiplatform`.
127139
"""
128140
self.info(f"{cls.__name__} created. Resource name: {resource.name}")
129141
self.info(f"To use this {cls.__name__} in another session:")
130-
self.info(f"{variable_name} = aiplatform.{cls.__name__}.get('{resource.name}')")
142+
usage_message = f"{module_name}.{cls.__name__}.get('{resource.name}')"
143+
self.info(f"{variable_name} = {usage_message}")
131144

132145
def log_delete_with_lro(
133146
self,

0 commit comments

Comments
 (0)