feature: SDK Defaults Config - Support for Session default_bucket_prefix & ability to set an SDK-wide S3 Key Prefix through Session

Author: Ruban Hussain

src/sagemaker/amazon/amazon_estimator.py

@@ -20,7 +20,7 @@
 
 from six.moves.urllib.parse import urlparse
 
-from sagemaker import image_uris
+from sagemaker import image_uris, s3_utils
 from sagemaker.amazon import validation
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.common import write_numpy_to_dense_tensor
@@ -93,8 +93,15 @@ def __init__(
             enable_network_isolation=enable_network_isolation,
             **kwargs
         )
-        data_location = data_location or "s3://{}/sagemaker-record-sets/".format(
-            self.sagemaker_session.default_bucket()
+
+        data_location = data_location or (
+            s3_utils.s3_path_join(
+                "s3://",
+                self.sagemaker_session.default_bucket(),
+                self.sagemaker_session.default_bucket_prefix,
+                "sagemaker-record-sets",
+            )
+            + "/"
         )
         self._data_location = data_location
 

src/sagemaker/automl/automl.py

@@ -17,7 +17,7 @@
 from typing import Optional, List, Dict
 from six import string_types
 
-from sagemaker import Model, PipelineModel
+from sagemaker import Model, PipelineModel, s3
 from sagemaker.automl.candidate_estimator import CandidateEstimator
 from sagemaker.config import (
     AUTO_ML_ROLE_ARN_PATH,
@@ -663,7 +663,14 @@ def _prepare_for_auto_ml_job(self, job_name=None):
             self.current_job_name = name_from_base(base_name, max_length=32)
 
         if self.output_path is None:
-            self.output_path = "s3://{}/".format(self.sagemaker_session.default_bucket())
+            self.output_path = (
+                s3.s3_path_join(
+                    "s3://",
+                    self.sagemaker_session.default_bucket(),
+                    self.sagemaker_session.default_bucket_prefix,
+                )
+                + "/"
+            )
 
     @classmethod
     def _get_supported_inference_keys(cls, container, default=None):

src/sagemaker/config/__init__.py

@@ -94,6 +94,7 @@
     AUTO_ML_INTER_CONTAINER_ENCRYPTION_PATH,
     ENDPOINT_CONFIG_DATA_CAPTURE_KMS_KEY_ID_PATH,
     SESSION_S3_BUCKET_PATH,
+    SESSION_S3_OBJECT_KEY_PREFIX_PATH,
     MONITORING_SCHEDULE_CONFIG,
     MONITORING_JOB_DEFINITION,
     MONITORING_OUTPUT_CONFIG,

src/sagemaker/config/config_schema.py

@@ -91,6 +91,7 @@
 ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION = "EnableInterContainerTrafficEncryption"
 SESSION = "Session"
 S3_BUCKET = "S3Bucket"
+S3_OBJECT_KEY_PREFIX = "S3ObjectKeyPrefix"
 
 
 def _simple_path(*args: str):
@@ -298,6 +299,10 @@ def _simple_path(*args: str):
     SAGEMAKER, TRAINING_JOB, ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION
 )
 SESSION_S3_BUCKET_PATH = _simple_path(SAGEMAKER, PYTHON_SDK, MODULES, SESSION, S3_BUCKET)
+SESSION_S3_OBJECT_KEY_PREFIX_PATH = _simple_path(
+    SAGEMAKER, PYTHON_SDK, MODULES, SESSION, S3_OBJECT_KEY_PREFIX
+)
+
 
 SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA = {
     "$schema": "https://json-schema.org/draft/2020-12/schema",
@@ -447,7 +452,6 @@ 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
@@ -492,9 +496,26 @@ def _simple_path(*args: str):
                                     ADDITIONAL_PROPERTIES: False,
                                     PROPERTIES: {
                                         S3_BUCKET: {
-                                            "description": "Used as `default_bucket` of Session",
+                                            "description": "sets `default_bucket` of Session",
                                             "$ref": "#/definitions/s3Bucket",
                                         },
+                                        S3_OBJECT_KEY_PREFIX: {
+                                            "description": (
+                                                "sets `default_s3_object_key_prefix` of Session"
+                                            ),
+                                            TYPE: "string",
+                                            # Regex based on
+                                            # https://docs.aws.amazon.com/AmazonS3/latest/userguide/object-keys.html
+                                            # For now, the regex only allows the "safe characters"
+                                            # specified by S3. If needed, the regex can be loosened
+                                            # (but not tightened) in the future without
+                                            # introducing backward incompatibility.
+                                            "pattern": (
+                                                r"^[a-zA-Z0-9!_.*'()-]+(/[a-zA-Z0-9!_.*'()-]+)*$"
+                                            ),
+                                            "minLength": 0,
+                                            "maxLength": 1024,
+                                        },
                                     },
                                 },
                                 REMOTE_FUNCTION: {

src/sagemaker/djl_inference/model.py

@@ -23,7 +23,7 @@
 from typing import Optional, Union, Dict, Any
 
 import sagemaker
-from sagemaker import s3, Predictor, image_uris, fw_utils
+from sagemaker import s3, Predictor, image_uris, fw_utils, s3_utils
 from sagemaker.deserializers import JSONDeserializer, BaseDeserializer
 from sagemaker.djl_inference import defaults
 from sagemaker.model import FrameworkModel
@@ -527,7 +527,9 @@ def prepare_container_def(
             deploy_key_prefix = fw_utils.model_code_key_prefix(
                 self.key_prefix, self.name, self.image_uri
             )
-            bucket = self.bucket or self.sagemaker_session.default_bucket()
+            bucket, deploy_key_prefix = s3_utils.calculate_bucket_and_prefix(
+                self.bucket, deploy_key_prefix, self.sagemaker_session
+            )
             uploaded_code = fw_utils.tar_and_upload_dir(
                 self.sagemaker_session.boto_session,
                 bucket,

src/sagemaker/estimator.py

@@ -27,7 +27,7 @@
 from six.moves.urllib.parse import urlparse
 
 import sagemaker
-from sagemaker import git_utils, image_uris, vpc_utils
+from sagemaker import git_utils, image_uris, vpc_utils, s3
 from sagemaker.analytics import TrainingJobAnalytics
 from sagemaker.config import (
     TRAINING_JOB_VOLUME_KMS_KEY_ID_PATH,
@@ -672,6 +672,9 @@ def __init__(
             enable_network_isolation=self._enable_network_isolation,
         )
 
+        # Internal flag
+        self._is_output_path_set_from_default_bucket_and_prefix = False
+
     @abstractmethod
     def training_image_uri(self):
         """Return the Docker image to use for training.
@@ -772,7 +775,12 @@ def _prepare_for_training(self, job_name=None):
             if self.sagemaker_session.local_mode and local_code:
                 self.output_path = ""
             else:
-                self.output_path = "s3://{}/".format(self.sagemaker_session.default_bucket())
+                self.output_path = s3.s3_path_join(
+                    "s3://",
+                    self.sagemaker_session.default_bucket(),
+                    self.sagemaker_session.default_bucket_prefix,
+                )
+                self._is_output_path_set_from_default_bucket_and_prefix = True
 
         if self.git_config:
             updated_paths = git_utils.git_clone_repo(
@@ -847,7 +855,8 @@ def _stage_user_code_in_s3(self) -> str:
         if is_pipeline_variable(self.output_path):
             if self.code_location is None:
                 code_bucket = self.sagemaker_session.default_bucket()
-                code_s3_prefix = self._assign_s3_prefix()
+                key_prefix = self.sagemaker_session.default_bucket_prefix
+                code_s3_prefix = self._assign_s3_prefix(key_prefix)
                 kms_key = None
             else:
                 code_bucket, key_prefix = parse_s3_url(self.code_location)
@@ -860,16 +869,33 @@ def _stage_user_code_in_s3(self) -> str:
             if local_mode:
                 if self.code_location is None:
                     code_bucket = self.sagemaker_session.default_bucket()
-                    code_s3_prefix = self._assign_s3_prefix()
+                    key_prefix = self.sagemaker_session.default_bucket_prefix
+                    code_s3_prefix = self._assign_s3_prefix(key_prefix)
                     kms_key = None
                 else:
                     code_bucket, key_prefix = parse_s3_url(self.code_location)
                     code_s3_prefix = self._assign_s3_prefix(key_prefix)
                     kms_key = None
             else:
                 if self.code_location is None:
-                    code_bucket, _ = parse_s3_url(self.output_path)
-                    code_s3_prefix = self._assign_s3_prefix()
+                    # TODO: if output_path was set from Session, include prefix. If it was set by
+                    #  the user, do not set the prefix (which would change behavior of existing
+                    #  notebooks)
+                    code_bucket, possible_key_prefix = parse_s3_url(self.output_path)
+
+                    if self._is_output_path_set_from_default_bucket_and_prefix:
+                        # Only include possible_key_prefix if the output_path was created from the
+                        # Session's default bucket and prefix. In that scenario, possible_key_prefix
+                        # will either be "" or Session.default_bucket_prefix.
+                        # Note: We cannot do `if (code_bucket == session.default_bucket() and
+                        # key_prefix == session.default_bucket_prefix)` instead because the user
+                        # could have passed in equivalent values themselves to output_path. And
+                        # including the prefix in that case could result in a potentially backwards
+                        # incompatible behavior change for the end user.
+                        code_s3_prefix = self._assign_s3_prefix(possible_key_prefix)
+                    else:
+                        code_s3_prefix = self._assign_s3_prefix()
+
                     kms_key = self.output_kms_key
                 else:
                     code_bucket, key_prefix = parse_s3_url(self.code_location)
@@ -1060,8 +1086,12 @@ def _set_source_s3_uri(self, rule):
         if "source_s3_uri" in (rule.rule_parameters or {}):
             parse_result = urlparse(rule.rule_parameters["source_s3_uri"])
             if parse_result.scheme != "s3":
-                desired_s3_uri = os.path.join(
-                    "s3://", self.sagemaker_session.default_bucket(), rule.name, str(uuid.uuid4())
+                desired_s3_uri = s3.s3_path_join(
+                    "s3://",
+                    self.sagemaker_session.default_bucket(),
+                    self.sagemaker_session.default_bucket_prefix,
+                    rule.name,
+                    str(uuid.uuid4()),
                 )
                 s3_uri = S3Uploader.upload(
                     local_path=rule.rule_parameters["source_s3_uri"],

src/sagemaker/experiments/_helper.py

@@ -19,6 +19,7 @@
 
 import botocore
 
+from sagemaker import s3_utils
 from sagemaker.experiments._utils import is_already_exist_error
 
 logger = logging.getLogger(__name__)
@@ -75,8 +76,15 @@ def upload_artifact(self, file_path):
             raise ValueError(
                 "{} does not exist or is not a file. Please supply a file path.".format(file_path)
             )
-        if not self.artifact_bucket:
-            self.artifact_bucket = self.sagemaker_session.default_bucket()
+
+        # If self.artifact_bucket is falsy, it will be set to sagemaker_session.default_bucket.
+        # In that case, and if sagemaker_session.default_bucket_prefix exists, self.artifact_prefix
+        # needs to be updated too (because not updating self.artifact_prefix would result in
+        # different behavior the 1st time this method is called vs the 2nd).
+        self.artifact_bucket, self.artifact_prefix = s3_utils.calculate_bucket_and_prefix(
+            self.artifact_bucket, self.artifact_prefix, self.sagemaker_session
+        )
+
         artifact_name = os.path.basename(file_path)
         artifact_s3_key = "{}/{}/{}".format(
             self.artifact_prefix, self.trial_component_name, artifact_name
@@ -96,8 +104,15 @@ def upload_object_artifact(self, artifact_name, artifact_object, file_extension=
         Returns:
             str: The s3 URI of the uploaded file and the version of the file.
         """
-        if not self.artifact_bucket:
-            self.artifact_bucket = self.sagemaker_session.default_bucket()
+
+        # If self.artifact_bucket is falsy, it will be set to sagemaker_session.default_bucket.
+        # In that case, and if sagemaker_session.default_bucket_prefix exists, self.artifact_prefix
+        # needs to be updated too (because not updating self.artifact_prefix would result in
+        # different behavior the 1st time this method is called vs the 2nd).
+        self.artifact_bucket, self.artifact_prefix = s3_utils.calculate_bucket_and_prefix(
+            self.artifact_bucket, self.artifact_prefix, self.sagemaker_session
+        )
+
         if file_extension:
             artifact_name = (
                 artifact_name + ("" if file_extension.startswith(".") else ".") + file_extension

src/sagemaker/lambda_helper.py

@@ -17,6 +17,8 @@
 import zipfile
 import time
 from botocore.exceptions import ClientError
+
+from sagemaker import s3, s3_utils
 from sagemaker.session import Session
 
 
@@ -109,12 +111,15 @@ def create(self):
         if self.script is not None:
             code = {"ZipFile": _zip_lambda_code(self.script)}
         else:
-            bucket = self.s3_bucket or self.session.default_bucket()
+            bucket, key_prefix = s3_utils.calculate_bucket_and_prefix(
+                self.s3_bucket, None, self.session
+            )
             key = _upload_to_s3(
                 s3_client=_get_s3_client(self.session),
                 function_name=self.function_name,
                 zipped_code_dir=self.zipped_code_dir,
                 s3_bucket=bucket,
+                s3_key_prefix=key_prefix,
             )
             code = {"S3Bucket": bucket, "S3Key": key}
 
@@ -148,7 +153,10 @@ def update(self):
                         ZipFile=_zip_lambda_code(self.script),
                     )
                 else:
-                    bucket = self.s3_bucket or self.session.default_bucket()
+                    bucket, key_prefix = s3_utils.calculate_bucket_and_prefix(
+                        self.s3_bucket, None, self.session
+                    )
+
                     # get function name to be used in S3 upload path
                     if self.function_arn:
                         versioned_function_name = self.function_arn.split("funtion:")[-1]
@@ -167,6 +175,7 @@ def update(self):
                             function_name=function_name_for_s3,
                             zipped_code_dir=self.zipped_code_dir,
                             s3_bucket=bucket,
+                            s3_key_prefix=key_prefix,
                         ),
                     )
                 return response
@@ -255,7 +264,7 @@ def _get_lambda_client(session):
     return lambda_client
 
 
-def _upload_to_s3(s3_client, function_name, zipped_code_dir, s3_bucket):
+def _upload_to_s3(s3_client, function_name, zipped_code_dir, s3_bucket, s3_key_prefix=None):
     """Upload the zipped code to S3 bucket provided in the Lambda instance.
 
     Lambda instance must have a path to the zipped code folder and a S3 bucket to upload
@@ -264,7 +273,13 @@ def _upload_to_s3(s3_client, function_name, zipped_code_dir, s3_bucket):
 
     Returns: the S3 key where the code is uploaded.
     """
-    key = "{}/{}/{}".format("lambda", function_name, "code")
+
+    key = s3.s3_path_join(
+        s3_key_prefix,
+        "lambda",
+        function_name,
+        "code",
+    )
     s3_client.upload_file(zipped_code_dir, s3_bucket, key)
     return key
 

src/sagemaker/local/local_session.py

@@ -25,6 +25,7 @@
     load_sagemaker_config,
     validate_sagemaker_config,
     SESSION_S3_BUCKET_PATH,
+    SESSION_S3_OBJECT_KEY_PREFIX_PATH,
 )
 from sagemaker.local.image import _SageMakerContainer
 from sagemaker.local.utils import get_docker_host
@@ -610,6 +611,7 @@ def __init__(
         s3_endpoint_url=None,
         disable_local_code=False,
         sagemaker_config: dict = None,
+        default_bucket_prefix=None,
     ):
         """Create a Local SageMaker Session.
 
@@ -632,6 +634,10 @@ def __init__(
                 this dictionary can be generated by calling
                 :func:`~sagemaker.config.load_sagemaker_config` and then be provided to the
                 Session.
+            default_bucket_prefix (str): The default prefix to use for S3 Object Keys. When
+                objects are saved to the Session's default_bucket, the Object Key used will
+                start with the default_bucket_prefix. If not provided here or within
+                sagemaker_config, no additional prefix will be added.
         """
         self.s3_endpoint_url = s3_endpoint_url
         # We use this local variable to avoid disrupting the __init__->_initialize API of the
@@ -643,6 +649,7 @@ def __init__(
             boto_session=boto_session,
             default_bucket=default_bucket,
             sagemaker_config=sagemaker_config,
+            default_bucket_prefix=default_bucket_prefix,
         )
 
         if platform.system() == "Windows":
@@ -710,6 +717,12 @@ def _initialize(
             config_path=SESSION_S3_BUCKET_PATH,
             sagemaker_session=self,
         )
+        # after sagemaker_config initialization, update self.default_bucket_prefix if needed
+        self.default_bucket_prefix = resolve_value_from_config(
+            direct_input=self.default_bucket_prefix,
+            config_path=SESSION_S3_OBJECT_KEY_PREFIX_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):