fix: container env generation for S3 URI and add test for the same (#2971)

Author: shreyapandit

src/sagemaker/model.py

@@ -466,7 +466,7 @@ def _upload_code(self, key_prefix: str, repack: bool = False) -> None:
             )
 
     def _script_mode_env_vars(self):
-        """Placeholder docstring"""
+        """Returns a mapping of environment variables for script mode execution"""
         script_name = None
         dir_name = None
         if self.uploaded_code:
@@ -478,8 +478,11 @@ def _script_mode_env_vars(self):
         elif self.entry_point is not None:
             script_name = self.entry_point
             if self.source_dir is not None:
-                dir_name = "file://" + self.source_dir
-
+                dir_name = (
+                    self.source_dir
+                    if self.source_dir.startswith("s3://")
+                    else "file://" + self.source_dir
+                )
         return {
             SCRIPT_PARAM_NAME.upper(): script_name or str(),
             DIR_PARAM_NAME.upper(): dir_name or str(),

tests/unit/sagemaker/model/test_model.py

@@ -26,6 +26,8 @@
 from sagemaker.sklearn.model import SKLearnModel
 from sagemaker.tensorflow.model import TensorFlowModel
 from sagemaker.xgboost.model import XGBoostModel
+from sagemaker.workflow.properties import Properties
+
 
 MODEL_DATA = "s3://bucket/model.tar.gz"
 MODEL_IMAGE = "mi"
@@ -42,7 +44,6 @@
 BRANCH = "test-branch-git-config"
 COMMIT = "ae15c9d7d5b97ea95ea451e4662ee43da3401d73"
 ENTRY_POINT_INFERENCE = "inference.py"
-
 SCRIPT_URI = "s3://codebucket/someprefix/sourcedir.tar.gz"
 IMAGE_URI = "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-inference:1.9.0-gpu-py38"
 
@@ -71,6 +72,23 @@ def sagemaker_session():
     return sms
 
 
+@patch("shutil.rmtree", MagicMock())
+@patch("tarfile.open", MagicMock())
+@patch("os.listdir", MagicMock(return_value=[ENTRY_POINT_INFERENCE]))
+def test_prepare_container_def_with_model_src_s3_returns_correct_url(sagemaker_session):
+    model = Model(
+        entry_point=ENTRY_POINT_INFERENCE,
+        role=ROLE,
+        sagemaker_session=sagemaker_session,
+        source_dir=SCRIPT_URI,
+        image_uri=MODEL_IMAGE,
+        model_data=Properties("Steps.MyStep"),
+    )
+    container_def = model.prepare_container_def(INSTANCE_TYPE, "ml.eia.medium")
+
+    assert container_def["Environment"]["SAGEMAKER_SUBMIT_DIRECTORY"] == SCRIPT_URI
+
+
 def test_prepare_container_def_with_model_data():
     model = Model(MODEL_IMAGE)
     container_def = model.prepare_container_def(INSTANCE_TYPE, "ml.eia.medium")