Skip to content

Commit 1d15f82

Browse files
Ark-kuncopybara-github
authored andcommitted
chore: Made the base.Logger class inherit from logging.Logger
This removes the need to dynamically dispatch all `Logger` methods and attributes via `__getattr__`. This change improves the static type checking and the IDE experience. PiperOrigin-RevId: 566442238
1 parent 41a3a83 commit 1d15f82

File tree

2 files changed

+25
-28
lines changed

2 files changed

+25
-28
lines changed

google/cloud/aiplatform/base.py

Lines changed: 24 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@
5757
_DEFAULT_RETRY = retry.Retry()
5858

5959

60-
class Logger:
60+
class VertexLogger(logging.getLoggerClass()):
6161
"""Logging wrapper class with high level helper methods."""
6262

6363
def __init__(self, name: str):
@@ -66,12 +66,12 @@ def __init__(self, name: str):
6666
Args:
6767
name (str): Name to associate with logger.
6868
"""
69-
self._logger = logging.getLogger(name)
70-
self._logger.setLevel(logging.INFO)
69+
super().__init__(name)
70+
self.setLevel(logging.INFO)
7171

7272
# To avoid writing duplicate logs, skip adding the new handler if
7373
# StreamHandler already exists in logger hierarchy.
74-
logger = self._logger
74+
logger = self
7575
while logger:
7676
for handler in logger.handlers:
7777
if isinstance(handler, logging.StreamHandler):
@@ -81,7 +81,7 @@ def __init__(self, name: str):
8181
handler = logging.StreamHandler(sys.stdout)
8282
handler.setLevel(logging.INFO)
8383

84-
self._logger.addHandler(handler)
84+
self.addHandler(handler)
8585

8686
def log_create_with_lro(
8787
self,
@@ -96,12 +96,10 @@ def log_create_with_lro(
9696
lro (operation.Operation):
9797
Optional. Backing LRO for creation.
9898
"""
99-
self._logger.info(f"Creating {cls.__name__}")
99+
self.info(f"Creating {cls.__name__}")
100100

101101
if lro:
102-
self._logger.info(
103-
f"Create {cls.__name__} backing LRO: {lro.operation.name}"
104-
)
102+
self.info(f"Create {cls.__name__} backing LRO: {lro.operation.name}")
105103

106104
def log_create_complete(
107105
self,
@@ -120,11 +118,9 @@ def log_create_complete(
120118
Vertex AI Resource proto.Message
121119
variable_name (str): Name of variable to use for code snippet
122120
"""
123-
self._logger.info(f"{cls.__name__} created. Resource name: {resource.name}")
124-
self._logger.info(f"To use this {cls.__name__} in another session:")
125-
self._logger.info(
126-
f"{variable_name} = aiplatform.{cls.__name__}('{resource.name}')"
127-
)
121+
self.info(f"{cls.__name__} created. Resource name: {resource.name}")
122+
self.info(f"To use this {cls.__name__} in another session:")
123+
self.info(f"{variable_name} = aiplatform.{cls.__name__}('{resource.name}')")
128124

129125
def log_create_complete_with_getter(
130126
self,
@@ -143,11 +139,9 @@ def log_create_complete_with_getter(
143139
Vertex AI Resource proto.Message
144140
variable_name (str): Name of variable to use for code snippet
145141
"""
146-
self._logger.info(f"{cls.__name__} created. Resource name: {resource.name}")
147-
self._logger.info(f"To use this {cls.__name__} in another session:")
148-
self._logger.info(
149-
f"{variable_name} = aiplatform.{cls.__name__}.get('{resource.name}')"
150-
)
142+
self.info(f"{cls.__name__} created. Resource name: {resource.name}")
143+
self.info(f"To use this {cls.__name__} in another session:")
144+
self.info(f"{variable_name} = aiplatform.{cls.__name__}.get('{resource.name}')")
151145

152146
def log_action_start_against_resource(
153147
self, action: str, noun: str, resource_noun_obj: "VertexAiResourceNoun"
@@ -160,7 +154,7 @@ def log_action_start_against_resource(
160154
resource_noun_obj (VertexAiResourceNoun):
161155
Resource noun object the action is acting against.
162156
"""
163-
self._logger.info(
157+
self.info(
164158
f"{action} {resource_noun_obj.__class__.__name__} {noun}: {resource_noun_obj.resource_name}"
165159
)
166160

@@ -180,9 +174,7 @@ def log_action_started_against_resource_with_lro(
180174
Resource noun object the action is acting against.
181175
lro (operation.Operation): Backing LRO for action.
182176
"""
183-
self._logger.info(
184-
f"{action} {cls.__name__} {noun} backing LRO: {lro.operation.name}"
185-
)
177+
self.info(f"{action} {cls.__name__} {noun} backing LRO: {lro.operation.name}")
186178

187179
def log_action_completed_against_resource(
188180
self, noun: str, action: str, resource_noun_obj: "VertexAiResourceNoun"
@@ -195,13 +187,18 @@ def log_action_completed_against_resource(
195187
resource_noun_obj (VertexAiResourceNoun):
196188
Resource noun object the action is acting against
197189
"""
198-
self._logger.info(
190+
self.info(
199191
f"{resource_noun_obj.__class__.__name__} {noun} {action}. Resource name: {resource_noun_obj.resource_name}"
200192
)
201193

202-
def __getattr__(self, attr: str):
203-
"""Forward remainder of logging to underlying logger."""
204-
return getattr(self._logger, attr)
194+
195+
def Logger(name: str) -> VertexLogger: # pylint: disable=invalid-name
196+
old_class = logging.getLoggerClass()
197+
try:
198+
logging.setLoggerClass(VertexLogger)
199+
return logging.getLogger(name)
200+
finally:
201+
logging.setLoggerClass(old_class)
205202

206203

207204
_LOGGER = Logger(__name__)

google/cloud/aiplatform/pipeline_jobs.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -409,7 +409,7 @@ def submit(
409409
credentials=self.credentials,
410410
)
411411
except: # noqa: E722
412-
_LOGGER._logger.exception(
412+
_LOGGER.exception(
413413
"Error when trying to get or create a GCS bucket for the pipeline output artifacts"
414414
)
415415

0 commit comments

Comments
 (0)