fix: duplicate logs in Colab

jaycee-li · copybara-github · commit 9b7525987e05 · 2023-10-05T17:04:57.000-07:00
PiperOrigin-RevId: 571173374
diff --git a/google/cloud/aiplatform/base.py b/google/cloud/aiplatform/base.py
@@ -69,20 +69,6 @@ def __init__(self, name: str):
         super().__init__(name)
         self.setLevel(logging.INFO)
 
-        # To avoid writing duplicate logs, skip adding the new handler if
-        # StreamHandler already exists in logger hierarchy.
-        logger = self
-        while logger:
-            for handler in logger.handlers:
-                if isinstance(handler, logging.StreamHandler):
-                    return
-            logger = logger.parent
-
-        handler = logging.StreamHandler(sys.stdout)
-        handler.setLevel(logging.INFO)
-
-        self.addHandler(handler)
-
     def log_create_with_lro(
         self,
         cls: Type["VertexAiResourceNoun"],
@@ -196,7 +182,22 @@ def Logger(name: str) -> VertexLogger:  # pylint: disable=invalid-name
     old_class = logging.getLoggerClass()
     try:
         logging.setLoggerClass(VertexLogger)
-        return logging.getLogger(name)
+        logger = logging.getLogger(name)
+
+        # To avoid writing duplicate logs, skip adding the new handler if
+        # StreamHandler already exists in logger hierarchy.
+        parent_logger = logger
+        while parent_logger:
+            for handler in parent_logger.handlers:
+                if isinstance(handler, logging.StreamHandler):
+                    return logger
+            parent_logger = parent_logger.parent
+
+        handler = logging.StreamHandler(sys.stdout)
+        handler.setLevel(logging.INFO)
+        logger.addHandler(handler)
+
+        return logger
     finally:
         logging.setLoggerClass(old_class)
 
diff --git a/tests/unit/aiplatform/test_base.py b/tests/unit/aiplatform/test_base.py
@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 
-# Copyright 2020 Google LLC
+# Copyright 2023 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -16,6 +16,7 @@
 #
 
 from importlib import reload
+import logging
 import pytest
 import time
 from typing import Optional
@@ -24,6 +25,9 @@
 from google.cloud.aiplatform import initializer
 
 
+_TEST_LOGGER_NAME = "test_logger"
+
+
 class _TestClass(base.FutureManager):
     def __init__(self, x):
         self.x = x
@@ -199,3 +203,12 @@ def test_create_and_add_return_arg(self, sync):
         assert isinstance(a, _TestClass)
         assert isinstance(b, _TestClassDownStream)
         assert isinstance(c, _TestClass)
+
+
+class TestLogger:
+    def test_logger_handler(self):
+        logger = base.Logger(_TEST_LOGGER_NAME)
+
+        assert logger.level == logging.INFO
+        # the logger won't have a StreamHandler because root logger already has one
+        assert not logger.handlers
