feature: SDK Defaults Config - Support for Session default_bucket

Author: Ruban Hussain

src/sagemaker/config/__init__.py

@@ -93,6 +93,7 @@
     AUTO_ML_VOLUME_KMS_KEY_ID_PATH,
     AUTO_ML_INTER_CONTAINER_ENCRYPTION_PATH,
     ENDPOINT_CONFIG_DATA_CAPTURE_KMS_KEY_ID_PATH,
+    SESSION_S3_BUCKET_PATH,
     MONITORING_SCHEDULE_CONFIG,
     MONITORING_JOB_DEFINITION,
     MONITORING_OUTPUT_CONFIG,
@@ -131,4 +132,8 @@
     EXECUTION_ROLE_ARN,
     ASYNC_INFERENCE_CONFIG,
     SCHEMA_VERSION,
+    PYTHON_SDK,
+    MODULES,
+    S3_BUCKET,
+    SESSION,
 )

src/sagemaker/config/config_schema.py

@@ -89,13 +89,16 @@
 OBJECT = "object"
 ADDITIONAL_PROPERTIES = "additionalProperties"
 ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION = "EnableInterContainerTrafficEncryption"
+SESSION = "Session"
+S3_BUCKET = "S3Bucket"
 
 
 def _simple_path(*args: str):
     """Appends an arbitrary number of strings to use as path constants"""
     return ".".join(args)
 
 
+# Paths for reference elsewhere in the SDK.
 COMPILATION_JOB_VPC_CONFIG_PATH = _simple_path(SAGEMAKER, COMPILATION_JOB, VPC_CONFIG)
 COMPILATION_JOB_KMS_KEY_ID_PATH = _simple_path(
     SAGEMAKER, COMPILATION_JOB, OUTPUT_CONFIG, KMS_KEY_ID
@@ -231,7 +234,6 @@ def _simple_path(*args: str):
 MODEL_PACKAGE_VALIDATION_PROFILES_PATH = _simple_path(
     SAGEMAKER, MODEL_PACKAGE, VALIDATION_SPECIFICATION, VALIDATION_PROFILES
 )
-
 REMOTE_FUNCTION_DEPENDENCIES = _simple_path(
     SAGEMAKER, PYTHON_SDK, MODULES, REMOTE_FUNCTION, DEPENDENCIES
 )
@@ -274,9 +276,6 @@ def _simple_path(*args: str):
 REMOTE_FUNCTION_ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION = _simple_path(
     SAGEMAKER, PYTHON_SDK, MODULES, REMOTE_FUNCTION, ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION
 )
-
-# Paths for reference elsewhere in the SDK.
-# Names include the schema version since the paths could change with other schema versions
 MONITORING_SCHEDULE_INTER_CONTAINER_ENCRYPTION_PATH = _simple_path(
     SAGEMAKER,
     MONITORING_SCHEDULE,
@@ -298,6 +297,7 @@ def _simple_path(*args: str):
 TRAINING_JOB_INTER_CONTAINER_ENCRYPTION_PATH = _simple_path(
     SAGEMAKER, TRAINING_JOB, ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION
 )
+SESSION_S3_BUCKET_PATH = _simple_path(SAGEMAKER, PYTHON_SDK, MODULES, SESSION, S3_BUCKET)
 
 SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA = {
     "$schema": "https://json-schema.org/draft/2020-12/schema",
@@ -447,6 +447,16 @@ def _simple_path(*args: str):
         "s3Uri": {TYPE: "string", "pattern": "^(https|s3)://([^/]+)/?(.*)$", "maxLength": 1024},
         # Regex is taken from https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_AlgorithmSpecification.html#sagemaker-Type-AlgorithmSpecification-ContainerEntrypoint
         "preExecutionCommand": {TYPE: "string", "pattern": r".*"},
+
+        # Regex based on https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_PipelineDefinitionS3Location.html
+        # except with an additional ^ and $ for the beginning and the end to closer align to
+        # https://docs.aws.amazon.com/AmazonS3/latest/userguide/bucketnamingrules.html
+        "s3Bucket": {
+            TYPE: "string",
+            "pattern": r"^[a-z0-9][\.\-a-z0-9]{1,61}[a-z0-9]$",
+            "minLength": 3,
+            "maxLength": 63,
+        },
     },
     PROPERTIES: {
         SCHEMA_VERSION: {
@@ -477,6 +487,16 @@ def _simple_path(*args: str):
                             TYPE: OBJECT,
                             ADDITIONAL_PROPERTIES: False,
                             PROPERTIES: {
+                                SESSION: {
+                                    TYPE: OBJECT,
+                                    ADDITIONAL_PROPERTIES: False,
+                                    PROPERTIES: {
+                                        S3_BUCKET: {
+                                            "description": "Used as `default_bucket` of Session",
+                                            "$ref": "#/definitions/s3Bucket",
+                                        },
+                                    },
+                                },
                                 REMOTE_FUNCTION: {
                                     TYPE: OBJECT,
                                     ADDITIONAL_PROPERTIES: False,
@@ -504,9 +524,9 @@ def _simple_path(*args: str):
                                         VOLUME_KMS_KEY_ID: {"$ref": "#/definitions/kmsKeyId"},
                                         VPC_CONFIG: {"$ref": "#/definitions/vpcConfig"},
                                     },
-                                }
+                                },
                             },
-                        }
+                        },
                     },
                 },
                 # Feature Group

src/sagemaker/local/local_session.py

@@ -21,7 +21,11 @@
 import boto3
 from botocore.exceptions import ClientError
 
-from sagemaker.config import load_sagemaker_config, validate_sagemaker_config
+from sagemaker.config import (
+    load_sagemaker_config,
+    validate_sagemaker_config,
+    SESSION_S3_BUCKET_PATH,
+)
 from sagemaker.local.image import _SageMakerContainer
 from sagemaker.local.utils import get_docker_host
 from sagemaker.local.entities import (
@@ -34,7 +38,7 @@
     _LocalPipeline,
 )
 from sagemaker.session import Session
-from sagemaker.utils import get_config_value, _module_import_error
+from sagemaker.utils import get_config_value, _module_import_error, resolve_value_from_config
 
 logger = logging.getLogger(__name__)
 
@@ -700,15 +704,22 @@ def _initialize(
             # create a default S3 resource, but only if it needs to fetch from S3
             self.sagemaker_config = load_sagemaker_config(s3_resource=self.s3_resource)
 
-        sagemaker_config_file = os.path.join(os.path.expanduser("~"), ".sagemaker", "config.yaml")
-        if os.path.exists(sagemaker_config_file):
+        # after sagemaker_config initialization, update self._default_bucket_name_override if needed
+        self._default_bucket_name_override = resolve_value_from_config(
+            direct_input=self._default_bucket_name_override,
+            config_path=SESSION_S3_BUCKET_PATH,
+            sagemaker_session=self,
+        )
+
+        local_mode_config_file = os.path.join(os.path.expanduser("~"), ".sagemaker", "config.yaml")
+        if os.path.exists(local_mode_config_file):
             try:
                 import yaml
             except ImportError as e:
                 logger.error(_module_import_error("yaml", "Local mode", "local"))
                 raise e
 
-            self.config = yaml.safe_load(open(sagemaker_config_file, "r"))
+            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
 

src/sagemaker/session.py

@@ -92,6 +92,7 @@
     FEATURE_GROUP_ROLE_ARN_PATH,
     FEATURE_GROUP_ONLINE_STORE_CONFIG_PATH,
     FEATURE_GROUP_OFFLINE_STORE_CONFIG_PATH,
+    SESSION_S3_BUCKET_PATH,
 )
 from sagemaker.deprecations import deprecated_class
 from sagemaker.inputs import ShuffleConfig, TrainingInput, BatchDataCaptureConfig
@@ -180,7 +181,8 @@ def __init__(
             default_bucket (str): The default Amazon S3 bucket to be used by this session.
                 This will be created the next time an Amazon S3 bucket is needed (by calling
                 :func:`default_bucket`).
-                If not provided, a default bucket will be created based on the following format:
+                If not provided, it will be fetched from the sagemaker_config. If not configured
+                there either, a default bucket will be created based on the following format:
                 "sagemaker-{region}-{aws-account-id}".
                 Example: "sagemaker-my-custom-bucket".
             settings (sagemaker.session_settings.SessionSettings): Optional. Set of optional
@@ -200,8 +202,13 @@ def __init__(
                 :func:`~sagemaker.config.load_sagemaker_config` and then be provided to the
                 Session.
         """
+
+        # sagemaker_config is validated and initialized inside :func:`_initialize`,
+        # so if default_bucket is None and the sagemaker_config has a default S3 bucket configured,
+        # _default_bucket_name_override will be set again inside :func:`_initialize`.
         self._default_bucket = None
         self._default_bucket_name_override = default_bucket
+
         self.s3_resource = None
         self.s3_client = None
         self.resource_groups_client = None
@@ -280,6 +287,13 @@ def _initialize(
             # create a default S3 resource, but only if it needs to fetch from S3
             self.sagemaker_config = load_sagemaker_config(s3_resource=self.s3_resource)
 
+        # after sagemaker_config initialization, update self._default_bucket_name_override if needed
+        self._default_bucket_name_override = resolve_value_from_config(
+            direct_input=self._default_bucket_name_override,
+            config_path=SESSION_S3_BUCKET_PATH,
+            sagemaker_session=self,
+        )
+
     @property
     def boto_region_name(self):
         """Placeholder docstring"""
@@ -484,7 +498,8 @@ def default_bucket(self):
         This function will create the s3 bucket if it does not exist.
 
         Returns:
-            str: The name of the default bucket, which is of the form:
+            str: The name of the default bucket. If the name was not explicitly specified through
+            the Session or sagemaker_config, the bucket will take the form:
                 ``sagemaker-{region}-{AWS account ID}``.
         """
 

tests/data/config/config.yaml

@@ -1,5 +1,9 @@
 SchemaVersion: '1.0'
 SageMaker:
+  PythonSDK:
+    Modules:
+      Session:
+        S3Bucket: 'sagemaker-python-sdk-test-bucket'
   FeatureGroup:
     OnlineStoreConfig:
       SecurityConfig:

tests/integ/test_local_mode.py

@@ -23,6 +23,8 @@
 import stopit
 
 import tests.integ.lock as lock
+from sagemaker.config import SESSION_S3_BUCKET_PATH
+from sagemaker.utils import resolve_value_from_config
 from tests.integ import DATA_DIR
 from mock import Mock, ANY
 
@@ -70,6 +72,13 @@ def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client,
 
         self.sagemaker_config = kwargs.get("sagemaker_config", None)
 
+        # after sagemaker_config initialization, update self._default_bucket_name_override if needed
+        self._default_bucket_name_override = resolve_value_from_config(
+            direct_input=self._default_bucket_name_override,
+            config_path=SESSION_S3_BUCKET_PATH,
+            sagemaker_session=self,
+        )
+
 
 class LocalPipelineNoS3Session(LocalPipelineSession):
     """
@@ -91,6 +100,13 @@ def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client,
 
         self.sagemaker_config = kwargs.get("sagemaker_config", None)
 
+        # after sagemaker_config initialization, update self._default_bucket_name_override if needed
+        self._default_bucket_name_override = resolve_value_from_config(
+            direct_input=self._default_bucket_name_override,
+            config_path=SESSION_S3_BUCKET_PATH,
+            sagemaker_session=self,
+        )
+
 
 @pytest.fixture(scope="module")
 def sagemaker_local_session_no_local_code(boto_session):

tests/unit/__init__.py

@@ -70,11 +70,27 @@
     MODEL,
     ASYNC_INFERENCE_CONFIG,
     SCHEMA_VERSION,
+    PYTHON_SDK,
+    MODULES,
+    S3_BUCKET,
+    SESSION,
 )
 
 DATA_DIR = os.path.join(os.path.dirname(__file__), "..", "data")
 PY_VERSION = "py3"
 
+SAGEMAKER_CONFIG_SESSION = {
+    SCHEMA_VERSION: "1.0",
+    SAGEMAKER: {
+        PYTHON_SDK: {
+            MODULES: {
+                SESSION: {
+                    S3_BUCKET: "sagemaker-config-session-s3-bucket",
+                },
+            },
+        },
+    },
+}
 
 SAGEMAKER_CONFIG_MONITORING_SCHEDULE = {
     SCHEMA_VERSION: "1.0",

tests/unit/sagemaker/config/conftest.py

@@ -37,6 +37,13 @@ def valid_tags():
     return [{"Key": "tag1", "Value": "tagValue1"}]
 
 
+@pytest.fixture()
+def valid_session_config():
+    return {
+        "S3Bucket": "sagemaker-python-sdk-test-bucket",
+    }
+
+
 @pytest.fixture()
 def valid_feature_group_config(valid_iam_role_arn):
     security_storage_config = {"KmsKeyId": "kmskeyid1"}
@@ -191,6 +198,7 @@ def valid_remote_function_config(valid_iam_role_arn, valid_tags, valid_vpc_confi
 
 @pytest.fixture()
 def valid_config_with_all_the_scopes(
+    valid_session_config,
     valid_feature_group_config,
     valid_monitoring_schedule_config,
     valid_endpointconfig_config,
@@ -206,6 +214,11 @@ def valid_config_with_all_the_scopes(
     valid_remote_function_config,
 ):
     return {
+        "PythonSDK": {
+            "Modules": {
+                "Session": valid_session_config,
+            }
+        },
         "FeatureGroup": valid_feature_group_config,
         "MonitoringSchedule": valid_monitoring_schedule_config,
         "EndpointConfig": valid_endpointconfig_config,

tests/unit/sagemaker/config/test_config_schema.py

@@ -199,3 +199,37 @@ def test_invalid_s3uri_schema(base_config_with_schema):
     config["SageMaker"] = {"PythonSDK": {"Modules": {"RemoteFunction": {"S3RootUri": "bad_regex"}}}}
     with pytest.raises(exceptions.ValidationError):
         validate(config, SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA)
+
+@pytest.mark.parametrize(
+    "bucket_name",
+    [
+        "docexamplebucket1",
+        "log-delivery-march-2020",
+        "my-hosted-content",
+        "docexamplewebsite.com",
+        "www.docexamplewebsite.com",
+        "my.example.s3.bucket",
+    ],
+)
+def test_session_s3_bucket_schema(base_config_with_schema, bucket_name):
+    config = {"PythonSDK": {"Modules": {"Session": {"S3Bucket": bucket_name}}}}
+    _validate_config(base_config_with_schema, config)
+
+
+@pytest.mark.parametrize(
+    "invalid_bucket_name",
+    [
+        "ab",
+        "this-is-sixty-four-characters-total-which-is-one-above-the-limit",
+        "UPPERCASE-LETTERS",
+        "special_characters",
+        "special-characters@",
+        ".dot-at-the-beginning",
+        "-dash-at-the-beginning",
+        "dot-at-the-end.",
+        "dash-at-the-end-",
+    ],
+)
+def test_invalid_session_s3_bucket_schema(base_config_with_schema, invalid_bucket_name):
+    with pytest.raises(exceptions.ValidationError):
+        test_session_s3_bucket_schema(base_config_with_schema, invalid_bucket_name)