pathway: add logging

Author: Zhankuil

src/sagemaker/remote_function/client.py

@@ -20,7 +20,6 @@
 from typing import Dict, List, Tuple, Any
 import functools
 import inspect
-import logging
 
 from botocore.exceptions import ClientError
 
@@ -29,6 +28,7 @@
 from sagemaker.session import Session
 from sagemaker.s3 import s3_path_join
 from sagemaker.remote_function.job import _JobSettings, _Job
+from sagemaker.remote_function import logging_config
 
 
 _API_CALL_LIMIT = {
@@ -44,8 +44,9 @@
 _CANCELLED = "CANCELLED"
 _FINISHED = "FINISHED"
 
-LOGGER = logging.getLogger(__name__)
-LOGGER.setLevel(logging.INFO)
+
+logging_config.basic_config()
+logger = logging_config.get_logger()
 
 
 def remote(
@@ -191,7 +192,7 @@ def _submit_worker(executor):
             else:
                 executor._running_jobs[job.job_name] = job
     except Exception:  # pylint: disable=broad-except
-        LOGGER.exception("Error occurred while submitting CreateTrainingJob requests.")
+        logger.exception("Error occurred while submitting CreateTrainingJob requests.")
 
 
 def _polling_worker(executor):
@@ -227,13 +228,13 @@ def _polling_worker(executor):
                         != "LimitExceededException"
                     ):
                         # Couldn't check the job status, move on
-                        LOGGER.exception(
+                        logger.exception(
                             "Error occurred while checking the status of job %s", job_name
                         )
                         del executor._running_jobs[job_name]
                         executor._semaphore.release()
     except Exception:  # pylint: disable=broad-except
-        LOGGER.exception("Error occurred while monitoring the job statuses.")
+        logger.exception("Error occurred while monitoring the job statuses.")
 
 
 # TODO: 1) add map method.

src/sagemaker/remote_function/core/runtime_environment.py

@@ -14,17 +14,16 @@
 
 from __future__ import absolute_import
 
-import logging
 import sys
 import os
 import shlex
 import subprocess
 from sagemaker.s3 import s3_path_join, S3Uploader, S3Downloader
 from sagemaker.session import Session
 from sagemaker.remote_function.errors import RuntimeEnvironmentError
+from sagemaker.remote_function.logging_config import get_logger
 
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+logger = get_logger()
 
 
 class RuntimeEnvironmentManager:
@@ -182,10 +181,10 @@ def _log_error(process):
         for line in iter(pipe.readline, b""):
             error_str = str(line, "UTF-8")
             if "ERROR:" in error_str:
-                logging.error(error_str)
+                logger.error(error_str)
                 error_logs = error_logs + error_str
             else:
-                logging.warn(error_str)
+                logger.warn(error_str)
 
     return error_logs
 

src/sagemaker/remote_function/core/stored_function.py

@@ -16,12 +16,17 @@
 import os
 import pathlib
 import shutil
+
 from sagemaker.utils import _tmpdir
 from sagemaker.s3 import s3_path_join, S3Uploader
+from sagemaker.remote_function import logging_config
 
 import sagemaker.remote_function.core.serialization as serialization
 
 
+logger = logging_config.get_logger()
+
+
 class StoredFunction:
     """Class representing a remote function stored in S3."""
 
@@ -81,15 +86,27 @@ def _zip_and_upload_source_dir(self, source_dir):
 
     def load_and_invoke(self) -> None:
         """Load and deserialize the function and the arguments and then execute it."""
+
+        logger.info(
+            f"Deserializing function code from {s3_path_join(self.s3_base_uri, 'function.pkl')}"
+        )
         func = serialization.deserialize_func_from_s3(
             self.sagemaker_session, s3_path_join(self.s3_base_uri, "function.pkl")
         )
+
+        logger.info(
+            f"Deserializing function arguments from {s3_path_join(self.s3_base_uri, 'arguments.pkl')}"
+        )
         args, kwargs = serialization.deserialize_obj_from_s3(
             self.sagemaker_session, s3_path_join(self.s3_base_uri, "arguments.pkl")
         )
 
+        logger.info("Invoking the function")
         result = func(*args, **kwargs)
 
+        logger.info(
+            f"Serializing the function return and uploading to {s3_path_join(self.s3_base_uri, 'results.pkl')}"
+        )
         serialization.serialize_obj_to_s3(
             result,
             self.sagemaker_session,

src/sagemaker/remote_function/job.py

@@ -27,8 +27,7 @@
 JOBS_CONTAINER_ENTRYPOINT = ["invoke-remote-function"]
 
 
-# TODO: 1) provide reasonable default values, e.g. image_uri
-#  2) extend this class to load job settings from the configuration files.
+# TODO: extend this class to load job settings from the configuration files.
 class _JobSettings:
     """Helper class that processes the job settings.
 
@@ -62,7 +61,7 @@ def __init__(
 
         self.environment_variables = environment_variables
 
-        self.image_uri = image_uri or "fake_image_uri"
+        self.image_uri = image_uri or _JobSettings._get_default_image_uri()
         self.dependencies = dependencies
 
         self.instance_type = instance_type
@@ -92,6 +91,13 @@ def __init__(
 
         self.tags = [] if tags is None else [{"Key": k, "Value": v} for k, v in tags]
 
+    @staticmethod
+    def _get_default_image_uri():
+        """Get the default image uri"""
+
+        # TODO: provide default image uri if not set
+        raise ValueError("image_uri must be set")
+
 
 class _Job:
     """Helper class that interacts with the SageMaker training service."""

src/sagemaker/remote_function/job_driver.py

@@ -15,11 +15,13 @@
 
 import argparse
 import sys
+
 import boto3
 
 from sagemaker.session import Session
 from sagemaker.remote_function.errors import handle_error
 from sagemaker.remote_function.core.runtime_environment import RuntimeEnvironmentManager
+from sagemaker.remote_function import logging_config
 
 
 SUCCESS_EXIT_CODE = 0
@@ -79,6 +81,10 @@ def _execute_remote_function(sagemaker_session, s3_base_uri, s3_kms_key):
 
 def main():
     """Entry point for job driver script"""
+
+    logging_config.basic_config()
+    logger = logging_config.get_logger()
+
     exit_code = SUCCESS_EXIT_CODE
     try:
         args = _parse_agrs()
@@ -105,6 +111,7 @@ def main():
         _execute_remote_function(sagemaker_session, s3_base_uri, s3_kms_key)
 
     except Exception as e:  # pylint: disable=broad-except
+        logger.exception("Error encountered when invoking the remote function.")
         exit_code = handle_error(e, sagemaker_session, s3_base_uri, s3_kms_key)
     finally:
         sys.exit(exit_code)

src/sagemaker/remote_function/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.
+"""Utilities related to logging."""
+from __future__ import absolute_import
+
+import logging
+
+
+def get_logger():
+    """Return a logger with the name 'sagemaker'"""
+    logger = logging.getLogger("sagemaker")
+    logger.setLevel(logging.INFO)
+    return logger
+
+
+def basic_config():
+    """Set logger configuration."""
+    logging.basicConfig(format="%(asctime)s %(name)-12s %(levelname)-8s %(message)s")

tests/integ/sagemaker/remote_function/test_decorator.py

@@ -80,6 +80,7 @@ def divide(x, y):
         divide(10, 2)
 
 
+# TODO: add VPC settings, update SageMakerRole with KMS permissions
 @pytest.mark.skip
 def test_advanced_job_setting(
     sagemaker_session, dummy_container_without_error, cpu_instance_type, s3_kms_key
@@ -88,7 +89,6 @@ def test_advanced_job_setting(
         role=ROLE,
         image_uri=dummy_container_without_error,
         instance_type=cpu_instance_type,
-        # TODO: add VPC settings
         s3_kms_key=s3_kms_key,
         sagemaker_session=sagemaker_session,
     )

tests/unit/sagemaker/remote_function/core/test_serialization.py

@@ -92,3 +92,15 @@ def test_serialize_deserialize_data():
     deserialized = deserialize_obj_from_s3(sagemaker_session=Mock(), s3_uri=s3_uri)
 
     assert deserialized == [10]
+
+
+@patch("sagemaker.s3.S3Uploader.upload_bytes", new=upload)
+@patch("sagemaker.s3.S3Downloader.read_bytes", new=read)
+def test_serialize_deserialize_none():
+
+    s3_uri = random_s3_uri()
+    serialize_obj_to_s3(None, sagemaker_session=Mock(), s3_uri=s3_uri, s3_kms_key=KMS_KEY)
+
+    deserialized = deserialize_obj_from_s3(sagemaker_session=Mock(), s3_uri=s3_uri)
+
+    assert deserialized is None

tests/unit/sagemaker/remote_function/test_client.py

@@ -171,16 +171,18 @@ def test_executor_invalid_arguments():
             e.submit(job_function, 1, 2, c=3, d=4)
 
 
-def test_executor_submit_after_shutdown():
+@patch("sagemaker.remote_function.client._JobSettings")
+def test_executor_submit_after_shutdown(*args):
     with pytest.raises(RuntimeError):
         with RemoteExecutor(max_parallel_job=1, s3_root_uri="s3://bucket/") as e:
             pass
         e.submit(job_function, 1, 2, c=3, d=4)
 
 
 @patch("sagemaker.remote_function.client._API_CALL_LIMIT", new=API_CALL_LIMIT)
+@patch("sagemaker.remote_function.client._JobSettings")
 @patch("sagemaker.remote_function.client._Job.start")
-def test_executor_submit_happy_case(mock_start):
+def test_executor_submit_happy_case(mock_start, *args):
     mock_job = Mock()
     mock_job.describe.return_value = COMPLETED_TRAINING_JOB
     mock_job.job_name = TRAINING_JOB_NAME
@@ -202,8 +204,9 @@ def test_executor_submit_happy_case(mock_start):
 
 
 @patch("sagemaker.remote_function.client._API_CALL_LIMIT", new=API_CALL_LIMIT)
+@patch("sagemaker.remote_function.client._JobSettings")
 @patch("sagemaker.remote_function.client._Job.start")
-def test_executor_submit_enforcing_max_parallel_jobs(mock_start):
+def test_executor_submit_enforcing_max_parallel_jobs(mock_start, *args):
     mock_job = Mock()
     mock_job.describe.return_value = INPROGRESS_TRAINING_JOB
     mock_start.return_value = mock_job
@@ -227,8 +230,9 @@ def test_executor_submit_enforcing_max_parallel_jobs(mock_start):
 
 
 @patch("sagemaker.remote_function.client._API_CALL_LIMIT", new=API_CALL_LIMIT)
+@patch("sagemaker.remote_function.client._JobSettings")
 @patch("sagemaker.remote_function.client._Job.start")
-def test_executor_fails_to_start_job(mock_start):
+def test_executor_fails_to_start_job(mock_start, *args):
     mock_job = Mock()
     mock_job.describe.return_value = COMPLETED_TRAINING_JOB
 
@@ -245,8 +249,9 @@ def test_executor_fails_to_start_job(mock_start):
 
 
 @patch("sagemaker.remote_function.client._API_CALL_LIMIT", new=API_CALL_LIMIT)
+@patch("sagemaker.remote_function.client._JobSettings")
 @patch("sagemaker.remote_function.client._Job.start")
-def test_executor_submit_and_cancel(mock_start):
+def test_executor_submit_and_cancel(mock_start, *args):
     mock_job = Mock()
     mock_job.describe.return_value = INPROGRESS_TRAINING_JOB
     mock_start.return_value = mock_job
@@ -270,8 +275,9 @@ def test_executor_submit_and_cancel(mock_start):
 
 
 @patch("sagemaker.remote_function.client._API_CALL_LIMIT", new=API_CALL_LIMIT)
+@patch("sagemaker.remote_function.client._JobSettings")
 @patch("sagemaker.remote_function.client._Job.start")
-def test_executor_describe_job_throttled_temporarily(mock_start):
+def test_executor_describe_job_throttled_temporarily(mock_start, *args):
     throttling_error = ClientError(
         error_response={"Error": {"Code": "LimitExceededException"}},
         operation_name="SomeOperation",
@@ -298,8 +304,9 @@ def test_executor_describe_job_throttled_temporarily(mock_start):
 
 
 @patch("sagemaker.remote_function.client._API_CALL_LIMIT", new=API_CALL_LIMIT)
+@patch("sagemaker.remote_function.client._JobSettings")
 @patch("sagemaker.remote_function.client._Job.start")
-def test_executor_describe_job_failed_permanently(mock_start):
+def test_executor_describe_job_failed_permanently(mock_start, *args):
     mock_job = Mock()
     mock_job.describe.side_effect = RuntimeError()
     mock_start.return_value = mock_job
@@ -333,7 +340,8 @@ def test_executor_describe_job_failed_permanently(mock_start):
         ),
     ],
 )
-def test_executor_submit_invalid_function_args(args, kwargs, error_message):
+@patch("sagemaker.remote_function.client._JobSettings")
+def test_executor_submit_invalid_function_args(MockJobSettings, args, kwargs, error_message):
     with pytest.raises(TypeError) as e:
         with RemoteExecutor(max_parallel_job=1, s3_root_uri="s3://bucket/") as executor:
             executor.submit(job_function, *args, **kwargs)

tests/unit/sagemaker/remote_function/test_job.py

@@ -12,6 +12,7 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
+import pytest
 from mock import patch, Mock, ANY
 from sagemaker.remote_function.job import _JobSettings, _Job
 
@@ -80,13 +81,20 @@ def job_function(a, b=1, *, c, d=3):
 @patch("sagemaker.remote_function.job.Session", return_value=mock_session())
 @patch("sagemaker.remote_function.job.get_execution_role", return_value=DEFAULT_ROLE_ARN)
 def test_default_settings(*args):
-    job_settings = _JobSettings()
-    assert job_settings.image_uri == "fake_image_uri"
-    assert job_settings.s3_root_uri == f"s3://{BUCKET}/fake_image_uri"
+    job_settings = _JobSettings(image_uri="image_uri")
+    assert job_settings.image_uri == "image_uri"
+    assert job_settings.s3_root_uri == f"s3://{BUCKET}/image_uri"
     assert job_settings.role == DEFAULT_ROLE_ARN
     assert job_settings.tags == []
 
 
+@patch("sagemaker.remote_function.job.Session", return_value=mock_session())
+@patch("sagemaker.remote_function.job.get_execution_role", return_value=DEFAULT_ROLE_ARN)
+def test_fails_on_missing_image_uri(*args):
+    with pytest.raises(ValueError):
+        _JobSettings(image_uri=None)
+
+
 @patch("sagemaker.remote_function.job.RuntimeEnvironmentManager")
 @patch("sagemaker.remote_function.job.StoredFunction")
 @patch("sagemaker.remote_function.job.Session", return_value=mock_session())