fix: adjacent slash in s3 key (#3427)

Co-authored-by: Namrata Madan <[email protected]>

Author: nmadan

src/sagemaker/model.py

@@ -569,8 +569,8 @@ def _upload_code(self, key_prefix: str, repack: bool = False) -> None:
                     )
                     return
                 self.sagemaker_session.context.need_runtime_repack.add(id(self))
-                self.sagemaker_session.context.runtime_repack_output_prefix = "s3://{}/{}".format(
-                    bucket, key_prefix
+                self.sagemaker_session.context.runtime_repack_output_prefix = s3.s3_path_join(
+                    "s3://", bucket, key_prefix
                 )
                 # Add the uploaded_code and repacked_model_data to update the container env
                 self.repacked_model_data = self.model_data

src/sagemaker/tensorflow/model.py

@@ -396,8 +396,8 @@ def prepare_container_def(
             # model is not yet there, defer repacking to later during pipeline execution
             if isinstance(self.sagemaker_session, PipelineSession):
                 self.sagemaker_session.context.need_runtime_repack.add(id(self))
-                self.sagemaker_session.context.runtime_repack_output_prefix = "s3://{}/{}".format(
-                    bucket, key_prefix
+                self.sagemaker_session.context.runtime_repack_output_prefix = s3.s3_path_join(
+                    "s3://", bucket, key_prefix
                 )
             else:
                 logging.warning(

tests/unit/sagemaker/workflow/test_model_step.py

@@ -70,6 +70,7 @@
 _TENSORFLOW_PATH = os.path.join(DATA_DIR, "tfs/tfs-test-entrypoint-and-dependencies")
 _REPACK_OUTPUT_KEY_PREFIX = "code-output"
 _MODEL_CODE_LOCATION = f"s3://{_BUCKET}/{_REPACK_OUTPUT_KEY_PREFIX}"
+_MODEL_CODE_LOCATION_TRAILING_SLASH = _MODEL_CODE_LOCATION + "/"
 
 
 @pytest.fixture
@@ -701,7 +702,7 @@ def test_conditional_model_create_and_regis(
                 entry_point=f"{DATA_DIR}/{_SCRIPT_NAME}",
                 role=_ROLE,
                 enable_network_isolation=True,
-                code_location=_MODEL_CODE_LOCATION,
+                code_location=_MODEL_CODE_LOCATION_TRAILING_SLASH,
             ),
             2,
         ),
@@ -725,7 +726,7 @@ def test_conditional_model_create_and_regis(
                 entry_point=f"{DATA_DIR}/{_SCRIPT_NAME}",
                 role=_ROLE,
                 framework_version="1.5.0",
-                code_location=_MODEL_CODE_LOCATION,
+                code_location=_MODEL_CODE_LOCATION_TRAILING_SLASH,
             ),
             2,
         ),
@@ -757,7 +758,7 @@ def test_conditional_model_create_and_regis(
                 image_uri=_IMAGE_URI,
                 entry_point=f"{DATA_DIR}/{_SCRIPT_NAME}",
                 role=_ROLE,
-                code_location=_MODEL_CODE_LOCATION,
+                code_location=_MODEL_CODE_LOCATION_TRAILING_SLASH,
             ),
             2,
         ),
@@ -780,7 +781,9 @@ def assert_test_result(steps: list):
         assert len(steps) == expected_step_num
         if expected_step_num == 2:
             assert steps[0]["Type"] == "Training"
-            if model.key_prefix == _REPACK_OUTPUT_KEY_PREFIX:
+            if model.key_prefix is not None and model.key_prefix.startswith(
+                _REPACK_OUTPUT_KEY_PREFIX
+            ):
                 assert steps[0]["Arguments"]["OutputDataConfig"]["S3OutputPath"] == (
                     f"{_MODEL_CODE_LOCATION}/{model.name}"
                 )