feature: Added Config parser for SageMaker Python SDK (aws#817)

Co-authored-by: Balaji Sankar <[email protected]>

Author: 2 people

src/sagemaker/config/config_factory.py

@@ -0,0 +1,369 @@
+# 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 athis 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.
+"""This module configures the default values for SageMaker Python SDK.
+
+It supports loading Config files from local file system/S3.
+The schema of the Config file is dictated in config_schema.py in the same module.
+
+"""
+from __future__ import absolute_import
+
+import abc
+import logging
+import os
+import yaml
+
+from botocore.utils import merge_dicts
+from sagemaker import s3
+from sagemaker.config.config_schema import SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA_V1_0
+from sagemaker.utils import get_config_value
+
+logger = logging.getLogger("sagemaker")
+
+
+class SageMakerConfig(abc.ABC):
+    """SageMakerConfig class encapsulates the Config for SageMaker Python SDK.
+
+    This class also exposes methods to retrieve the Config.
+    Note: This class shouldn't be directly instantiated.
+
+    .. tip::
+        Use SageMakerConfigFactory to initialize the SageMakerConfig class.
+        Subclasses which override ``__init__`` should invoke ``super()``.
+    """
+
+    def __init__(self):
+        """Initializes a SageMakerConfig object.
+
+        Note: This constructor invokes _load_config() method which should be implemented by the
+        subclasses.
+
+        Once the Config is loaded, this constructor will validate the schema of the config file.
+
+        .. tip::
+            Subclasses which override ``__init__`` should invoke ``super()``.
+
+        """
+        self.config_values = self._load_config()
+        self.validate()
+
+    @abc.abstractmethod
+    def _load_config(self) -> dict:
+        """Should be implemented by the Subclass.
+
+        Subclass is responsible for fetching the config file.
+
+        Returns:
+            dict: Config as a python dictionary which obeys the schema mentioned in config_schema.py
+        """
+
+    def validate(self):
+        """Validates the schema of the Config.
+
+        SchemaError exception is thrown if the schema is not obeyed.
+        """
+        if len(self.config_values) != 0:
+            # If Config file is not found, then Config object will be empty.
+            # We will NOT be doing schema validations for an empty config
+            SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA_V1_0.validate(self.config_values)
+
+    def merge(self, config_to_be_merged):
+        """Merges the given SageMakerConfig object with another SageMakerConfig object.
+
+        This method will merge the configuration values. If a same key is present in both the
+        config files, then the value present in the 'config_to_be_merged' will overwrite the
+        current value.
+
+        Note: This method re-uses the merge_dicts utility method implemented in botocore.
+
+        Args:
+            config_to_be_merged (sagemaker.config_management.config_factory.SageMakerConfig):
+                The additional SageMakerConfig object that needs to be merged with the current
+                object.
+
+        """
+        merge_dicts(self.config_values, config_to_be_merged.config_values)
+
+    def get_config_value(self, key_path):
+        """Given a key path, retrieve the corresponding value in the Config.
+
+        Args:
+            key_path (str): Key path of the config entry.
+            Nested entries are represented using '.' (Dot).
+
+        Returns:
+            object: Represents the object that is present in that key path of the Config.
+                Returns None if the key path doesn't exist.
+        """
+        return get_config_value(key_path, self.config_values)
+
+    def get_config(self):
+        """Returns the entire configuration as a dictionary.
+
+        Returns: The entire configuration as a dictionary.
+        """
+        return self.config_values
+
+
+class _SageMakerConfigFromFile(SageMakerConfig):
+    """An implementation of the SageMakerConfig which loads the Config from a local file system."""
+
+    def __init__(self, file_path):
+        """Constructor for _SageMakerConfigFromFile.
+
+        Note: This internally invokes the super class constructor.
+
+        Args:
+            file_path (str): The local file system path of the YAML file from which the
+            config needs to be loaded. Note: This needs to be an absolute path.
+        """
+        self.file_path = file_path
+        super(_SageMakerConfigFromFile, self).__init__()
+
+    def _load_config(self) -> dict:
+        """Overridden implementation of the _load_config method.
+
+        This method loads the config file from the path that was specified as
+        part of the constructor.
+        If the path that was provided, corresponds to a directory then this method
+        will try to search for 'config.yaml' in that directory.
+        Note: We will not be doing any recursive search.
+
+        Returns:
+            dict: A dictionary representing the configurations.
+            Note: This dictionary will be empty if an invalid file path is provided
+        """
+        config = {}
+        if self.file_path:
+            inferred_file_path = self.file_path
+            if os.path.isdir(self.file_path):
+                inferred_file_path = os.path.join(self.file_path, "config.yaml")
+            if _validate_file_exists(inferred_file_path):
+                config = yaml.safe_load(open(inferred_file_path, "r"))
+        return config
+
+
+class _SageMakerConfigFromS3(SageMakerConfig):
+    """An implementation of the SageMakerConfig which loads the Config from a S3 location."""
+
+    def __init__(self, s3_uri, sagemaker_session=None):
+        """Constructor for _SageMakerConfigFromFile.
+
+        This internally invokes the super class constructor.
+
+        Args:
+            sagemaker_session (sagemaker.session.Session): Session object which
+                manages interactions with Amazon SageMaker APIs and any other
+                AWS services needed. If not specified, a default session object will be created
+                using the default AWS configuration chain.
+            s3_uri (str): An S3 uri that refers to a location in which config file is present.
+                s3_uri must start with 's3://'.
+                An example s3_uri: 's3://example-bucket/config.yaml'.
+                If there are multiple S3 objects with the same key prefix,
+                then we will infer the path to be s3://example-bucket/somekeyprefix/config.yaml
+
+        """
+        self.sagemaker_session = sagemaker_session
+        self.s3_uri = _get_inferred_s3_path(s3_uri, self.sagemaker_session)
+        super(_SageMakerConfigFromS3, self).__init__()
+
+    def _load_config(self) -> dict:
+        """Overridden implementation of the _load_config method.
+
+        This method loads the config file from the S3 URI that was
+        specified as part of the constructor.
+
+        Note: Currently we don't support Client side KMS. This capability might be added in the
+        future
+
+        Returns:
+            dict: A dictionary representing the configurations.
+            Note: This dictionary will be empty if an invalid S3 URI is provided
+
+        """
+        config = {}
+        s3_file_content = None
+        if self.s3_uri:
+            try:
+                s3_file_content = s3.S3Downloader.read_file(self.s3_uri, self.sagemaker_session)
+            except Exception as e:  # pylint: disable=W0703
+                # if customers didn't provide us with a valid S3 File/insufficient read permission,
+                # we DO want to silently ignore that and operate with an empty config.
+                logger.warning(
+                    "Unable to fetch Config file from S3 with URI: %s due to %s", self.s3_uri, e
+                )
+            if s3_file_content:
+                config = yaml.safe_load(s3_file_content)
+        return config
+
+
+class SageMakerConfigFactory(object):
+    """Factory Class to create SageMakerConfig object.
+
+    Also supports specifying an additional config for override.
+    """
+
+    @staticmethod
+    def build_sagemaker_config(
+        default_config_location: str = None,
+        additional_override_config_location: str = None,
+        sagemaker_session=None,
+    ) -> SageMakerConfig:
+        """Factory method to build a SageMakerConfig object.
+
+        Note: This method also supports building an additional config override
+        (in case if there are multiple configs).
+        This method will then merge the additional config override to the base (default) config.
+
+        Args:
+            default_config_location (str): File path of the Config.
+                This can even be a S3 URI/Local File path.
+            additional_override_config_location (str):
+                File path of the Config override. This can even be a S3 URI/Local File path.
+            sagemaker_session (sagemaker.session.Session): Session object which
+                manages interactions with Amazon SageMaker APIs and any other
+                AWS services needed. If not specified, a default session object will be created
+                using the default AWS configuration chain.
+
+        Returns:
+            SageMakerConfig: A SageMakerConfig object which contains the merged Config values.
+        """
+        inferred_default_config_path = _get_default_config_path(default_config_location)
+        inferred_additional_config_override_path = _get_additional_config_override_path(
+            additional_override_config_location
+        )
+        default_config = _get_sagemaker_config(inferred_default_config_path, sagemaker_session)
+        additional_override_config = _get_sagemaker_config(
+            inferred_additional_config_override_path, sagemaker_session
+        )
+        default_config.merge(additional_override_config)
+        return default_config
+
+
+# Default path in which the SageMaker Python SDK looks for Config objects.
+DEFAULT_CONFIG_FILE_PATH = os.path.join(
+    os.path.expanduser("~"), ".sagemaker", "defaults", "sdk-default-config.yaml"
+)
+ENV_VARIABLE_DEFAULT_CONFIG_OVERRIDE = "SAGEMAKER_DEFAULT_CONFIG_OVERRIDE"
+ENV_VARIABLE_ADDITIONAL_CONFIG_OVERRIDE = "SAGEMAKER_ADDITIONAL_CONFIG_OVERRIDE"
+
+
+def _get_sagemaker_config(file_path: str, sagemaker_session) -> SageMakerConfig:
+    """Constructs SageMakerConfig object based on the file path that was provided.
+
+    Args:
+        file_path (str): File path of the Config. This can even be a S3 URI/Local File path.
+        sagemaker_session (sagemaker.session.Session): Session object which
+            manages interactions with Amazon SageMaker APIs and any other
+            AWS services needed. If not specified, a default session object will be created
+            using the default AWS configuration chain.
+
+    Returns:
+        SageMakerConfig: The SageMakerConfig object based on the file path that was provided.
+    """
+    if file_path and file_path.startswith("s3://"):
+        return _SageMakerConfigFromS3(file_path, sagemaker_session)
+    return _SageMakerConfigFromFile(file_path)
+
+
+def _validate_file_exists(file_path) -> bool:
+    """Validates whether a file/directory corresponding to that path exists
+
+    Args:
+        file_path (str): The file path for which we need to verify its existence,
+
+    Returns:
+        bool: Boolean indicating whether a file/directory corresponding to that path exists.
+    """
+    if not os.path.exists(file_path):
+        logger.warning("The specified file path: %s for the Config file doesn't exist.", file_path)
+        return False
+    return True
+
+
+def _get_default_config_path(default_config_directory_method_parameter: str = None) -> str:
+    """Returns the default config path.
+
+    Args:
+        default_config_directory_method_parameter (str):
+            Corresponds to the default_config_location parameter of SageMakerConfigFactory
+
+    Returns:
+        str: If default_config_directory_method_parameter is passed, then we will use that value.
+            Else, this method will try to fetch the value from SAGEMAKER_CONFIG_FILE environment
+            variable.If SAGEMAKER_CONFIG_FILE environment variable is not set,
+            we will infer the path to be DEFAULT_CONFIG_FILE_PATH
+    """
+    if default_config_directory_method_parameter:
+        return default_config_directory_method_parameter
+    return os.getenv(ENV_VARIABLE_DEFAULT_CONFIG_OVERRIDE, DEFAULT_CONFIG_FILE_PATH)
+
+
+def _get_additional_config_override_path(
+    additional_config_override_method_parameter: str = None,
+) -> str:
+    """Returns the additional config override path.
+
+    Args:
+        additional_config_override_method_parameter (str): Corresponds to the
+            additional_override_config_location parameter of SageMakerConfigFactory
+
+    Returns:
+        str: If additional_config_override_method_parameter is passed, then we will use that value.
+            Else, this method will try to fetch the value from SAGEMAKER_ADDITIONAL_CONFIG_OVERRIDE
+            environment variable
+    """
+    if additional_config_override_method_parameter:
+        return additional_config_override_method_parameter
+    return os.getenv(ENV_VARIABLE_ADDITIONAL_CONFIG_OVERRIDE, None)
+
+
+def _get_inferred_s3_path(s3_uri, sagemaker_session):
+    """Verifies whether the given S3 URI exists and returns the URI.
+
+    If there are multiple S3 objects with the same key prefix,
+    then this method will verify whether S3 URI + /config.yaml exists.
+    s3://example-bucket/somekeyprefix/config.yaml
+
+    Args:
+        s3_uri (str) : An S3 uri that refers to a location in which config file is present.
+            s3_uri must start with 's3://'.
+            An example s3_uri: 's3://example-bucket/config.yaml'.
+        sagemaker_session (sagemaker.session.Session): Session object which
+            manages interactions with Amazon SageMaker APIs and any other
+            AWS services needed. If not specified, a default session object will be created
+            using the default AWS configuration chain.
+
+    Returns:
+        str: Valid S3 URI of the Config file. None if it doesn't exist.
+    """
+    try:
+        s3_files_with_same_prefix = s3.S3Downloader.list(s3_uri, sagemaker_session)
+    except Exception as e:  # pylint: disable=W0703
+        # if customers didn't provide us with a valid S3 File/insufficient read permission,
+        # we DO want to silently ignore that and operate with an empty config.
+        logger.warning("Unable to read from S3 with URI: %s due to %s", s3_uri, e)
+        return None
+    if len(s3_files_with_same_prefix) == 0:
+        # Provided S3 URI is invalid.
+        # we DO want to silently ignore that and operate with an empty config.
+        return None
+    if len(s3_files_with_same_prefix) > 1:
+        # Customer has provided us with a S3 URI which points to a directory
+        inferred_s3_path = s3.s3_path_join(s3_uri, "config.yaml")
+        if inferred_s3_path not in s3_files_with_same_prefix:
+            # We don't know which file we should be operating with.
+            return None
+        # Customer has a config.yaml present in the directory that was provided as the S3 URI
+        return inferred_s3_path
+    return s3_uri