update: logging improvements (aws#901)

Co-authored-by: Zhankui Lu <[email protected]>

Author: 2 people

src/sagemaker/remote_function/client.py

@@ -52,7 +52,6 @@
 _CANCELLED = "CANCELLED"
 _FINISHED = "FINISHED"
 
-logging_config.basic_config()
 logger = logging_config.get_logger()
 
 

src/sagemaker/remote_function/invoke_function.py

@@ -78,7 +78,6 @@ def _execute_remote_function(sagemaker_session, s3_base_uri, s3_kms_key, run_in_
 def main():
     """Entry point for invoke function script"""
 
-    logging_config.basic_config()
     logger = logging_config.get_logger()
 
     exit_code = SUCCESS_EXIT_CODE

src/sagemaker/remote_function/logging_config.py

@@ -14,15 +14,25 @@
 from __future__ import absolute_import
 
 import logging
+import time
+
+
+class _UTCFormatter(logging.Formatter):
+    """Class that overrides the default local time provider in log formatter."""
+
+    converter = time.gmtime
 
 
 def get_logger():
     """Return a logger with the name 'sagemaker'"""
-    logger = logging.getLogger("sagemaker")
-    logger.setLevel(logging.INFO)
-    return logger
-
+    sagemaker_logger = logging.getLogger("sagemaker.remote_function")
+    if len(sagemaker_logger.handlers) == 0:
+        sagemaker_logger.setLevel(logging.INFO)
+        handler = logging.StreamHandler()
+        formatter = _UTCFormatter("%(asctime)s %(name)s %(levelname)-8s %(message)s")
+        handler.setFormatter(formatter)
+        sagemaker_logger.addHandler(handler)
+        # don't stream logs with the root logger handler
+        sagemaker_logger.propagate = 0
 
-def basic_config():
-    """Set logger configuration."""
-    logging.basicConfig(format="%(asctime)s %(name)-12s %(levelname)-8s %(message)s")
+    return sagemaker_logger

src/sagemaker/remote_function/runtime_environment/bootstrap_runtime_environment.py

@@ -14,17 +14,17 @@
 from __future__ import absolute_import
 
 import argparse
-import logging
 import sys
 import os
 import shutil
 import pathlib
 
 if __package__ is None or __package__ == "":
-    from runtime_environment_manager import RuntimeEnvironmentManager
+    from runtime_environment_manager import RuntimeEnvironmentManager, get_logger
 else:
     from sagemaker.remote_function.runtime_environment.runtime_environment_manager import (
         RuntimeEnvironmentManager,
+        get_logger,
     )
 
 SUCCESS_EXIT_CODE = 0
@@ -37,10 +37,7 @@
 SAGEMAKER_WHL_CHANNEL = "sagemaker_whl_file"
 
 
-logging.basicConfig(
-    format="%(asctime)s %(name)-12s %(levelname)-8s %(message)s", level=logging.INFO
-)
-logger = logging.getLogger("sagemaker")
+logger = get_logger()
 
 
 def main():

src/sagemaker/remote_function/runtime_environment/runtime_environment_manager.py

@@ -14,17 +14,38 @@
 
 from __future__ import absolute_import
 
+
 import logging
 import sys
 import shlex
 import os
 import subprocess
 import tempfile
+import time
+
+
+class _UTCFormatter(logging.Formatter):
+    """Class that overrides the default local time provider in log formatter."""
+
+    converter = time.gmtime
+
+
+def get_logger():
+    """Return a logger with the name 'sagemaker'"""
+    sagemaker_logger = logging.getLogger("sagemaker.remote_function")
+    if len(sagemaker_logger.handlers) == 0:
+        sagemaker_logger.setLevel(logging.INFO)
+        handler = logging.StreamHandler()
+        formatter = _UTCFormatter("%(asctime)s %(name)s %(levelname)-8s %(message)s")
+        handler.setFormatter(formatter)
+        sagemaker_logger.addHandler(handler)
+        # don't stream logs with the root logger handler
+        sagemaker_logger.propagate = 0
+
+    return sagemaker_logger
+
 
-logging.basicConfig(
-    format="%(asctime)s %(name)-12s %(levelname)-8s %(message)s", level=logging.INFO
-)
-logger = logging.getLogger("sagemaker")
+logger = get_logger()
 
 
 class RuntimeEnvironmentManager:
@@ -143,7 +164,7 @@ def bootstrap(
         """Bootstraps the runtime environment by installing the additional dependencies if any.
 
         Args:
-            dependencies_s3_uri (str): S3 URI where dependencies file exists.
+            local_dependencies_file (str): path where dependencies file exists.
             conda_env (str): conda environment to be activated. Default is None.
 
         Returns: None

tests/integ/sagemaker/remote_function/test_decorator.py

@@ -16,6 +16,7 @@
 
 import pytest
 import os
+import logging
 import pandas as pd
 from sagemaker.experiments.run import Run, load_run
 from tests.integ.sagemaker.experiments.helpers import cleanup_exp_resources
@@ -68,6 +69,7 @@ def test_decorated_function_raises_exception(
         sagemaker_session=sagemaker_session,
     )
     def divide(x, y):
+        logging.warning(f"{x}/{y}")
         return x / y
 
     with pytest.raises(ZeroDivisionError):

tests/unit/sagemaker/remote_function/test_logging_config.py

@@ -0,0 +1,28 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
+import logging
+
+from sagemaker.remote_function.logging_config import get_logger
+
+
+def test_logger_config():
+    logging.basicConfig(level=logging.INFO)
+
+    logger_1 = get_logger()
+    assert len(logger_1.handlers) == 1
+
+    logger_2 = get_logger()
+    assert logger_2 is logger_1
+    assert len(logger_2.handlers) == 1