Add container config to local mode config

Author: trungleduc

src/sagemaker/config/__init__.py

@@ -13,7 +13,11 @@
 """This module configures the default values for SageMaker Python SDK."""
 
 from __future__ import absolute_import
-from sagemaker.config.config import load_sagemaker_config, validate_sagemaker_config  # noqa: F401
+from sagemaker.config.config import (  # noqa: F401
+    load_local_mode_config,
+    load_sagemaker_config,
+    validate_sagemaker_config,
+)
 from sagemaker.config.config_schema import (  # noqa: F401
     KEY,
     TRAINING_JOB,
@@ -161,4 +165,8 @@
     INFERENCE_SPECIFICATION,
     ESTIMATOR,
     DEBUG_HOOK_CONFIG,
+    SAGEMAKER_PYTHON_SDK_LOCAL_MODE_CONFIG_SCHEMA,
+    LOCAL,
+    LOCAL_CODE,
+    CONTAINER_CONFIG,
 )

src/sagemaker/config/config.py

@@ -16,7 +16,7 @@
 The schema of the config file is dictated in config_schema.py in the same module.
 
 """
-from __future__ import absolute_import
+from __future__ import absolute_import, annotations
 
 import pathlib
 import os
@@ -33,12 +33,18 @@
 logger = get_sagemaker_config_logger()
 
 _APP_NAME = "sagemaker"
+# The default name of the config file.
+_CONFIG_FILE_NAME = "config.yaml"
 # The default config file location of the Administrator provided config file. This path can be
 # overridden with `SAGEMAKER_ADMIN_CONFIG_OVERRIDE` environment variable.
-_DEFAULT_ADMIN_CONFIG_FILE_PATH = os.path.join(site_config_dir(_APP_NAME), "config.yaml")
+_DEFAULT_ADMIN_CONFIG_FILE_PATH = os.path.join(site_config_dir(_APP_NAME), _CONFIG_FILE_NAME)
 # The default config file location of the user provided config file. This path can be
 # overridden with `SAGEMAKER_USER_CONFIG_OVERRIDE` environment variable.
-_DEFAULT_USER_CONFIG_FILE_PATH = os.path.join(user_config_dir(_APP_NAME), "config.yaml")
+_DEFAULT_USER_CONFIG_FILE_PATH = os.path.join(user_config_dir(_APP_NAME), _CONFIG_FILE_NAME)
+# The default config file location of the local mode.
+_DEFAULT_LOCAL_MODE_CONFIG_FILE_PATH = os.path.join(
+    os.path.expanduser("~"), ".sagemaker", _CONFIG_FILE_NAME
+)
 
 ENV_VARIABLE_ADMIN_CONFIG_OVERRIDE = "SAGEMAKER_ADMIN_CONFIG_OVERRIDE"
 ENV_VARIABLE_USER_CONFIG_OVERRIDE = "SAGEMAKER_USER_CONFIG_OVERRIDE"
@@ -144,11 +150,21 @@ def validate_sagemaker_config(sagemaker_config: dict = None):
     jsonschema.validate(sagemaker_config, SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA)
 
 
+def load_local_mode_config() -> dict | None:
+    """Loads the local mode config file."""
+    try:
+        content = _load_config_from_file(_DEFAULT_LOCAL_MODE_CONFIG_FILE_PATH)
+    except ValueError:
+        content = None
+
+    return content
+
+
 def _load_config_from_file(file_path: str) -> dict:
     """Placeholder docstring"""
     inferred_file_path = file_path
     if os.path.isdir(file_path):
-        inferred_file_path = os.path.join(file_path, "config.yaml")
+        inferred_file_path = os.path.join(file_path, _CONFIG_FILE_NAME)
     if not os.path.exists(inferred_file_path):
         raise ValueError(
             f"Unable to load the config file from the location: {file_path}"
@@ -194,10 +210,14 @@ def _get_inferred_s3_uri(s3_uri, s3_resource_for_config):
     if len(s3_files_with_same_prefix) > 1:
         # Customer has provided us with a S3 URI which points to a directory
         # search for s3://<bucket>/directory-key-prefix/config.yaml
-        inferred_s3_uri = str(pathlib.PurePosixPath(s3_uri, "config.yaml")).replace("s3:/", "s3://")
+        inferred_s3_uri = str(pathlib.PurePosixPath(s3_uri, _CONFIG_FILE_NAME)).replace(
+            "s3:/", "s3://"
+        )
         if inferred_s3_uri not in s3_files_with_same_prefix:
             # We don't know which file we should be operating with.
-            raise ValueError("Provide an S3 URI of a directory that has a config.yaml file.")
+            raise ValueError(
+                f"Provide an S3 URI of a directory that has a {_CONFIG_FILE_NAME} file."
+            )
         # Customer has a config.yaml present in the directory that was provided as the S3 URI
         return inferred_s3_uri
     return s3_uri

src/sagemaker/config/config_schema.py

@@ -102,7 +102,11 @@
 DISABLE_PROFILER = "DisableProfiler"
 ESTIMATOR = "Estimator"
 DEBUG_HOOK_CONFIG = "DebugHookConfig"
-
+LOCAL = "local"
+LOCAL_CODE = "local_code"
+SERVING_PORT = "serving_port"
+CONTAINER_CONFIG = "container_config"
+REGION_NAME = "region_name"
 
 def _simple_path(*args: str):
     """Appends an arbitrary number of strings to use as path constants"""
@@ -1068,3 +1072,28 @@ def _simple_path(*args: str):
         },
     },
 }
+
+SAGEMAKER_PYTHON_SDK_LOCAL_MODE_CONFIG_SCHEMA = {
+    "$schema": "https://json-schema.org/draft/2020-12/schema",
+    TYPE: OBJECT,
+    ADDITIONAL_PROPERTIES: False,
+    PROPERTIES: {
+        LOCAL: {
+            TYPE: OBJECT,
+            ADDITIONAL_PROPERTIES: False,
+            PROPERTIES: {
+                LOCAL_CODE: {
+                    TYPE: "boolean",
+                },
+                REGION_NAME: {TYPE: "string"},
+                SERVING_PORT: {
+                    TYPE: "integer",
+                },
+                CONTAINER_CONFIG: {
+                    TYPE: OBJECT,
+                },
+            },
+        },
+    },
+    "required": [LOCAL],
+}

src/sagemaker/local/image.py

@@ -36,6 +36,7 @@
 from six.moves.urllib.parse import urlparse
 
 import sagemaker
+from sagemaker.config.config_schema import CONTAINER_CONFIG, LOCAL
 import sagemaker.local.data
 import sagemaker.local.utils
 import sagemaker.utils
@@ -75,7 +76,6 @@ def __init__(
         sagemaker_session=None,
         container_entrypoint=None,
         container_arguments=None,
-        container_default_config=None,
     ):
         """Initialize a SageMakerContainer instance
 
@@ -92,8 +92,6 @@ def __init__(
                 to use when interacting with SageMaker.
             container_entrypoint (str): the container entrypoint to execute
             container_arguments (str): the container entrypoint arguments
-            container_default_config (Dict | None): the dict of user-defined docker
-                configuration. Defaults to ``None``
         """
         from sagemaker.local.local_session import LocalSession
 
@@ -106,7 +104,6 @@ def __init__(
         self.image = image
         self.container_entrypoint = container_entrypoint
         self.container_arguments = container_arguments
-        self.container_default_config = container_default_config or {}
         # Since we are using a single docker network, Generate a random suffix to attach to the
         # container names. This way multiple jobs can run in parallel.
         suffix = "".join(random.choice(string.ascii_lowercase + string.digits) for _ in range(5))
@@ -773,7 +770,7 @@ def _compose(self, detached=False):
 
         logger.info("docker command: %s", " ".join(compose_cmd))
         return compose_cmd
-    
+
     def _create_docker_host(
         self,
         host: str,
@@ -785,8 +782,8 @@ def _create_docker_host(
         """Creates the docker host configuration.
 
         Args:
-            host (str): The host address 
-            environment (List[str]): List of environment variables 
+            host (str): The host address
+            environment (List[str]): List of environment variables
             optml_subdirs (Set[str]): Set of subdirs
             command (str): Either 'train' or 'serve'
             volumes (list): List of volumes that will be mapped to the containers
@@ -797,9 +794,15 @@ def _create_docker_host(
         container_name_prefix = "".join(
             random.choice(string.ascii_lowercase + string.digits) for _ in range(10)
         )
+        container_default_config = (
+            sagemaker.utils.get_config_value(
+                f"{LOCAL}.{CONTAINER_CONFIG}", self.sagemaker_session.config
+            )
+            or {}
+        )
 
         host_config = {
-            **self.container_default_config,
+            **container_default_config,
             "image": self.image,
             "container_name": f"{container_name_prefix}-{host}",
             "stdin_open": True,

src/sagemaker/local/local_session.py

@@ -11,21 +11,24 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 """Placeholder docstring"""
-from __future__ import absolute_import
+from __future__ import absolute_import, annotations
 
 import logging
-import os
 import platform
 from datetime import datetime
+from typing import Dict
 
 import boto3
 from botocore.exceptions import ClientError
+import jsonschema
 
 from sagemaker.config import (
-    load_sagemaker_config,
-    validate_sagemaker_config,
+    SAGEMAKER_PYTHON_SDK_LOCAL_MODE_CONFIG_SCHEMA,
     SESSION_DEFAULT_S3_BUCKET_PATH,
     SESSION_DEFAULT_S3_OBJECT_KEY_PREFIX_PATH,
+    load_local_mode_config,
+    load_sagemaker_config,
+    validate_sagemaker_config,
 )
 from sagemaker.local.image import _SageMakerContainer
 from sagemaker.local.utils import get_docker_host
@@ -83,7 +86,7 @@ def create_processing_job(
         Environment=None,
         ProcessingInputs=None,
         ProcessingOutputConfig=None,
-        **kwargs
+        **kwargs,
     ):
         """Creates a processing job in Local Mode
 
@@ -128,7 +131,6 @@ def create_processing_job(
             sagemaker_session=self.sagemaker_session,
             container_entrypoint=container_entrypoint,
             container_arguments=container_arguments,
-            container_default_config=self._container_default_config
         )
         processing_job = _LocalProcessingJob(container)
         logger.info("Starting processing job")
@@ -166,7 +168,7 @@ def create_training_job(
         ResourceConfig,
         InputDataConfig=None,
         Environment=None,
-        **kwargs
+        **kwargs,
     ):
         """Create a training job in Local Mode.
 
@@ -231,7 +233,7 @@ def create_transform_job(
         TransformInput,
         TransformOutput,
         TransformResources,
-        **kwargs
+        **kwargs,
     ):
         """Create the transform job.
 
@@ -727,19 +729,25 @@ def _initialize(
             sagemaker_session=self,
         )
 
-        local_mode_config_file = os.path.join(os.path.expanduser("~"), ".sagemaker", "config.yaml")
-        if os.path.exists(local_mode_config_file):
+        self.config = load_local_mode_config()
+        if self._disable_local_code and self.config and "local" in self.config:
+            self.config["local"]["local_code"] = False
+
+    @Session.config.setter
+    def config(self, value: Dict | None):
+        """Setter of the local mode config"""
+        if value is not None:
             try:
-                import yaml
-            except ImportError as e:
-                logger.error(_module_import_error("yaml", "Local mode", "local"))
+                jsonschema.validate(value, SAGEMAKER_PYTHON_SDK_LOCAL_MODE_CONFIG_SCHEMA)
+            except jsonschema.ValidationError as e:
+                logger.error("Failed to validate the local mode config")
                 raise e
+            self._config = value
+        else:
+            self._config = value
 
-            self.config = yaml.safe_load(open(local_mode_config_file, "r"))
-            if self._disable_local_code and "local" in self.config:
-                self.config["local"]["local_code"] = False
-
-            self.sagemaker_runtime_client = LocalSagemakerRuntimeClient(self.config)
+        # update the runtime client on config changed
+        self.sagemaker_runtime_client = LocalSagemakerRuntimeClient(self._config)
 
     def logs_for_job(self, job_name, wait=False, poll=5, log_type="All"):
         """A no-op method meant to override the sagemaker client.

src/sagemaker/session.py

@@ -11,7 +11,7 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 """Placeholder docstring"""
-from __future__ import absolute_import, print_function
+from __future__ import absolute_import, annotations, print_function
 
 import json
 import logging
@@ -239,7 +239,7 @@ def __init__(
         self.s3_client = None
         self.resource_groups_client = None
         self.resource_group_tagging_client = None
-        self.config = None
+        self._config = None
         self.lambda_client = None
         self.settings = settings
 
@@ -326,6 +326,16 @@ def _initialize(
             sagemaker_session=self,
         )
 
+    @property
+    def config(self) -> Dict | None:
+        """The config for the local mode, unused in a normal session"""
+        return self._config
+
+    @config.setter
+    def config(self, value: Dict | None):
+        """The config for the local mode, unused in a normal session"""
+        self._config = value
+
     @property
     def boto_region_name(self):
         """Placeholder docstring"""

tox.ini

@@ -42,6 +42,7 @@ ignore =
     FI55,
     FI56,
     FI57,
+    FI58,
     W503
 
 require-code = True