fix: Replace SageMakerConfig class with just method invocations

Author: Balaji Sankar

src/sagemaker/config/__init__.py

@@ -13,7 +13,7 @@
 """This module configures the default values for SageMaker Python SDK."""
 
 from __future__ import absolute_import
-from sagemaker.config.config import SageMakerConfig  # noqa: F401
+from sagemaker.config.config import fetch_sagemaker_config, validate_sagemaker_config  # noqa: F401
 from sagemaker.config.config_schema import (  # noqa: F401
     KEY,
     TRAINING_JOB,
@@ -130,4 +130,5 @@
     RESOURCE_CONFIG,
     EXECUTION_ROLE_ARN,
     ASYNC_INFERENCE_CONFIG,
+    SCHEMA_VERSION,
 )

src/sagemaker/config/config.py

@@ -24,7 +24,7 @@
 from typing import List
 import boto3
 import yaml
-from jsonschema import validate
+import jsonschema
 from platformdirs import site_config_dir, user_config_dir
 from botocore.utils import merge_dicts
 from six.moves.urllib.parse import urlparse
@@ -51,87 +51,65 @@
 S3_PREFIX = "s3://"
 
 
-class SageMakerConfig(object):
-    """A class that encapsulates the configuration for the SageMaker Python SDK.
+def fetch_sagemaker_config(
+    additional_config_paths: List[str] = None, s3_resource=_DEFAULT_S3_RESOURCE
+) -> dict:
+    """Helper method that loads config files and merges them.
+
+    By default, this method first searches for config files in the default locations
+    defined by the SDK.
+
+    Users can override the default admin and user config file paths using the
+    SAGEMAKER_ADMIN_CONFIG_OVERRIDE and SAGEMAKER_USER_CONFIG_OVERRIDE environment variables,
+    respectively.
+
+    Additional config file paths can also be provided as a parameter.
+
+    This method then:
+    * Loads each config file, whether it is Amazon S3 or the local file system.
+    * Validates the schema of the config files.
+    * Merges the files in the same order.
+
+    This method throws exceptions in the following cases:
+    * jsonschema.exceptions.ValidationError: Schema validation fails for one or more config
+    files.
+    * RuntimeError: The method is unable to retrieve the list of all S3 files with the
+    same prefix or is unable to retrieve the file.
+    * ValueError: There are no S3 files with the prefix when an S3 URI is provided.
+    * ValueError: There is no config.yaml file in the S3 bucket when an S3 URI is provided.
+    * ValueError: A file doesn't exist in a path that was specified by the user as part of an
+    environment variable or additional configuration file path. This doesn't include the default
+    config file locations.
+
+    Args:
+        additional_config_paths: List of config file paths.
+            These paths can be one of the following. In the case of a directory, this method
+            searches for a config.yaml file in that directory. This method does not perform a
+            recursive search of folders in that directory.
+            * Local file path
+            * Local directory path
+            * S3 URI of the config file
+            * S3 URI of the directory containing the config file
+            Note: S3 URI follows the format s3://<bucket>/<Key prefix>
+        s3_resource (boto3.resource("s3")): The Boto3 S3 resource. This is used to fetch
+        config files from S3. If it is not provided, this method creates a default S3 resource
+        See :py:meth:boto3.session.Session.resource. This argument is not needed if the
+        config files are present in the local file system.
 
-    This class is used to define default values provided by the user.
-
-    This class is integrated with sagemaker.session.Session. Users of the SageMaker Python SDK
-    have the ability to pass a SageMakerConfig object to sagemaker.session.Session. If a
-    SageMakerConfig object is not provided by the user, then sagemaker.session.Session
-    creates its own SageMakerConfig object.
-
-    Note: After sagemaker.session.Session is initialized, it operates with the configuration
-    values defined at that instant. If you modify the configuration files or file paths after
-    sagemaker.session.Session is initialized, those changes are not reflected in
-    sagemaker.session.Session. To incorporate the changes in the configuration files,
-    initialize sagemaker.session.Session again.
     """
-
-    def __init__(self, additional_config_paths: List[str] = None, s3_resource=_DEFAULT_S3_RESOURCE):
-        """Initializes the SageMakerConfig object.
-
-        By default, this method first searches for config files in the default locations
-        defined by the SDK.
-
-        Users can override the default admin and user config file paths using the
-        SAGEMAKER_ADMIN_CONFIG_OVERRIDE and SAGEMAKER_USER_CONFIG_OVERRIDE environment variables,
-        respectively.
-
-        Additional config file paths can also be provided as a constructor parameter.
-
-        This method then:
-        * Loads each config file, whether it is Amazon S3 or the local file system.
-        * Validates the schema of the config files.
-        * Merges the files in the same order.
-
-        This method throws exceptions in the following cases:
-        * jsonschema.exceptions.ValidationError: Schema validation fails for one or more config
-        files.
-        * RuntimeError: The method is unable to retrieve the list of all S3 files with the
-        same prefix or is unable to retrieve the file.
-        * ValueError: There are no S3 files with the prefix when an S3 URI is provided.
-        * ValueError: There is no config.yaml file in the S3 bucket when an S3 URI is provided.
-        * ValueError: A file doesn't exist in a path that was specified by the user as part of an
-        environment variable or additional configuration file path. This doesn't include the default
-        config file locations.
-
-        Args:
-            additional_config_paths: List of config file paths.
-                These paths can be one of the following. In the case of a directory, this method
-                searches for a config.yaml file in that directory. This method does not perform a
-                recursive search of folders in that directory.
-                * Local file path
-                * Local directory path
-                * S3 URI of the config file
-                * S3 URI of the directory containing the config file
-                Note: S3 URI follows the format s3://<bucket>/<Key prefix>
-            s3_resource (boto3.resource("s3")): The Boto3 S3 resource. This is used to fetch
-            config files from S3. If it is not provided, this method creates a default S3 resource
-            See :py:meth:boto3.session.Session.resource. This argument is not needed if the
-            config files are present in the local file system.
-
-        """
-        default_config_path = os.getenv(
-            ENV_VARIABLE_ADMIN_CONFIG_OVERRIDE, _DEFAULT_ADMIN_CONFIG_FILE_PATH
-        )
-        user_config_path = os.getenv(
-            ENV_VARIABLE_USER_CONFIG_OVERRIDE, _DEFAULT_USER_CONFIG_FILE_PATH
-        )
-        self.config_paths = [default_config_path, user_config_path]
-        if additional_config_paths:
-            self.config_paths += additional_config_paths
-        self.config_paths = list(filter(lambda item: item is not None, self.config_paths))
-        self.config = _load_config_files(self.config_paths, s3_resource)
-
-
-def _load_config_files(file_paths: List[str], s3_resource_for_config) -> dict:
-    """Placeholder docstring"""
+    default_config_path = os.getenv(
+        ENV_VARIABLE_ADMIN_CONFIG_OVERRIDE, _DEFAULT_ADMIN_CONFIG_FILE_PATH
+    )
+    user_config_path = os.getenv(ENV_VARIABLE_USER_CONFIG_OVERRIDE, _DEFAULT_USER_CONFIG_FILE_PATH)
+    config_paths = [default_config_path, user_config_path]
+    if additional_config_paths:
+        config_paths += additional_config_paths
+    config_paths = list(filter(lambda item: item is not None, config_paths))
     merged_config = {}
-    for file_path in file_paths:
+    for file_path in config_paths:
         config_from_file = {}
         if file_path.startswith(S3_PREFIX):
-            config_from_file = _load_config_from_s3(file_path, s3_resource_for_config)
+            config_from_file = _load_config_from_s3(file_path, s3_resource)
         else:
             try:
                 config_from_file = _load_config_from_file(file_path)
@@ -145,11 +123,24 @@ def _load_config_files(file_paths: List[str], s3_resource_for_config) -> dict:
                     # Exceptions.
                     raise
         if config_from_file:
-            validate(config_from_file, SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA)
+            validate_sagemaker_config(config_from_file)
             merge_dicts(merged_config, config_from_file)
     return merged_config
 
 
+def validate_sagemaker_config(sagemaker_config: dict = None):
+    """Helper method that validates whether the schema of a given dictionary.
+
+    This method will validate whether the dictionary adheres to the schema
+    defined at `~sagemaker.config.config_schema.SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA`
+
+    Args:
+        sagemaker_config: A dictionary containing default values for the
+                SageMaker Python SDK. (default: None).
+    """
+    jsonschema.validate(sagemaker_config, SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA)
+
+
 def _load_config_from_file(file_path: str) -> dict:
     """Placeholder docstring"""
     inferred_file_path = file_path

src/sagemaker/local/local_session.py

@@ -21,7 +21,7 @@
 import boto3
 from botocore.exceptions import ClientError
 
-from sagemaker.config import SageMakerConfig
+from sagemaker.config import fetch_sagemaker_config, validate_sagemaker_config
 from sagemaker.local.image import _SageMakerContainer
 from sagemaker.local.utils import get_docker_host
 from sagemaker.local.entities import (
@@ -605,7 +605,7 @@ def __init__(
         default_bucket=None,
         s3_endpoint_url=None,
         disable_local_code=False,
-        sagemaker_config: SageMakerConfig = None,
+        sagemaker_config: dict = None,
     ):
         """Create a Local SageMaker Session.
 
@@ -618,6 +618,16 @@ def __init__(
             disable_local_code (bool): Set ``True`` to override the default AWS configuration
                 chain to disable the ``local.local_code`` setting, which may not be supported for
                 some SDK features (default: False).
+            sagemaker_config: A dictionary containing default values for the
+                SageMaker Python SDK. (default: None). The dictionary must adhere to the schema
+                defined at `~sagemaker.config.config_schema.SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA`.
+                If sagemaker_config is not provided and configuration files exist (at the default
+                paths for admins and users, or paths set through the environment variables
+                SAGEMAKER_ADMIN_CONFIG_OVERRIDE and SAGEMAKER_USER_CONFIG_OVERRIDE),
+                a new dictionary will be generated from those configuration files. Alternatively,
+                this dictionary can be generated by calling
+                :func:`~sagemaker.config.fetch_sagemaker_config` and then be provided to the
+                Session.
         """
         self.s3_endpoint_url = s3_endpoint_url
         # We use this local variable to avoid disrupting the __init__->_initialize API of the
@@ -635,12 +645,7 @@ def __init__(
             logger.warning("Windows Support for Local Mode is Experimental")
 
     def _initialize(
-        self,
-        boto_session,
-        sagemaker_client,
-        sagemaker_runtime_client,
-        sagemaker_config: SageMakerConfig = None,
-        **kwargs
+        self, boto_session, sagemaker_client, sagemaker_runtime_client, **kwargs
     ):  # pylint: disable=unused-argument
         """Initialize this Local SageMaker Session.
 
@@ -669,20 +674,20 @@ def _initialize(
         self.sagemaker_client = LocalSagemakerClient(self)
         self.sagemaker_runtime_client = LocalSagemakerRuntimeClient(self.config)
         self.local_mode = True
+        sagemaker_config = kwargs.get("sagemaker_config", None)
+        if sagemaker_config:
+            validate_sagemaker_config(sagemaker_config)
 
         if self.s3_endpoint_url is not None:
             self.s3_resource = boto_session.resource("s3", endpoint_url=self.s3_endpoint_url)
             self.s3_client = boto_session.client("s3", endpoint_url=self.s3_endpoint_url)
-            self.sagemaker_config = sagemaker_config or (
-                SageMakerConfig(s3_resource=self.s3_resource)
-                if "sagemaker_config" not in kwargs
-                else kwargs.get("sagemaker_config")
+            self.sagemaker_config = (
+                sagemaker_config if sagemaker_config else fetch_sagemaker_config(
+                    s3_resource=self.s3_resource)
             )
         else:
-            self.sagemaker_config = sagemaker_config or (
-                SageMakerConfig()
-                if "sagemaker_config" not in kwargs
-                else kwargs.get("sagemaker_config")
+            self.sagemaker_config = (
+                sagemaker_config if sagemaker_config else fetch_sagemaker_config()
             )
 
         sagemaker_config_file = os.path.join(os.path.expanduser("~"), ".sagemaker", "config.yaml")

src/sagemaker/processing.py

@@ -1246,20 +1246,10 @@ def __init__(
         self.s3_data_distribution_type = s3_data_distribution_type
         self.s3_compression_type = s3_compression_type
         self.s3_input = s3_input
-        self._dataset_definition = dataset_definition
+        self.dataset_definition = dataset_definition
         self.app_managed = app_managed
         self._create_s3_input()
 
-    @property
-    def dataset_definition(self):
-        """Getter for DataSetDefinition
-
-        Returns:
-            DatasetDefinition: The DatasetDefinition Object.
-
-        """
-        return self._dataset_definition
-
     def _to_request_dict(self):
         """Generates a request dictionary using the parameters provided to the class."""
 

src/sagemaker/session.py

@@ -33,7 +33,7 @@
 import sagemaker.logs
 from sagemaker import vpc_utils
 from sagemaker._studio import _append_project_tags
-from sagemaker.config import SageMakerConfig  # noqa: F401
+from sagemaker.config import fetch_sagemaker_config, validate_sagemaker_config  # noqa: F401
 from sagemaker.config import (
     KEY,
     TRAINING_JOB,
@@ -157,7 +157,7 @@ def __init__(
         default_bucket=None,
         settings=SessionSettings(),
         sagemaker_metrics_client=None,
-        sagemaker_config: SageMakerConfig = None,
+        sagemaker_config: dict = None,
     ):
         """Initialize a SageMaker ``Session``.
 
@@ -189,9 +189,16 @@ def __init__(
                 Client which makes SageMaker Metrics related calls to Amazon SageMaker
                 (default: None). If not provided, one will be created using
                 this instance's ``boto_session``.
-            sagemaker_config (sagemaker.config.SageMakerConfig): The SageMakerConfig object which
-                holds the default values for the SageMaker Python SDK. (default: None). If not
-                provided, one will be created.
+            sagemaker_config (dict): A dictionary containing default values for the
+                SageMaker Python SDK. (default: None). The dictionary must adhere to the schema
+                defined at `~sagemaker.config.config_schema.SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA`.
+                If sagemaker_config is not provided and configuration files exist (at the default
+                paths for admins and users, or paths set through the environment variables
+                SAGEMAKER_ADMIN_CONFIG_OVERRIDE and SAGEMAKER_USER_CONFIG_OVERRIDE),
+                a new dictionary will be generated from those configuration files. Alternatively,
+                this dictionary can be generated by calling
+                :func:`~sagemaker.config.fetch_sagemaker_config` and then be provided to the
+                Session.
         """
         self._default_bucket = None
         self._default_bucket_name_override = default_bucket
@@ -217,7 +224,7 @@ def _initialize(
         sagemaker_runtime_client,
         sagemaker_featurestore_runtime_client,
         sagemaker_metrics_client,
-        sagemaker_config: SageMakerConfig = None,
+        sagemaker_config: dict = None,
     ):
         """Initialize this SageMaker Session.
 
@@ -260,13 +267,13 @@ def _initialize(
 
         self.local_mode = False
         if sagemaker_config:
-            self.sagemaker_config = sagemaker_config
+            validate_sagemaker_config(sagemaker_config)
         else:
             if self.s3_resource is None:
                 s3 = self.boto_session.resource("s3", region_name=self.boto_region_name)
             else:
                 s3 = self.s3_resource
-            self.sagemaker_config = SageMakerConfig(s3_resource=s3)
+            self.sagemaker_config = fetch_sagemaker_config(s3_resource=s3)
 
     @property
     def boto_region_name(self):

src/sagemaker/utils.py

@@ -37,6 +37,7 @@
 from six.moves.urllib import parse
 
 from sagemaker import deprecations
+from sagemaker.config import validate_sagemaker_config
 from sagemaker.session_settings import SessionSettings
 from sagemaker.workflow import is_pipeline_variable, is_pipeline_parameter_string
 
@@ -1090,7 +1091,9 @@ def get_sagemaker_config_value(sagemaker_session, key):
     """
     if not sagemaker_session:
         return None
-    config_value = get_config_value(key, sagemaker_session.sagemaker_config.config)
+    if sagemaker_session.sagemaker_config:
+        validate_sagemaker_config(sagemaker_session.sagemaker_config)
+    config_value = get_config_value(key, sagemaker_session.sagemaker_config)
     # Copy the value so any modifications to the output will not modify the source config
     return copy.deepcopy(config_value)