Fixes container env generation for S3 URI, update test

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:
@@ -476,9 +476,12 @@ def _script_mode_env_vars(self):
             else:
                 dir_name = self.uploaded_code.s3_prefix
         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(),

tests/unit/sagemaker/model/test_framework_model.py

@@ -25,6 +25,7 @@
 MODEL_IMAGE = "mi"
 ENTRY_POINT = "blah.py"
 ROLE = "some-role"
+S3_SOURCE_DIR = "s3://somebucket/sourcedir.tar.gz"
 
 DATA_DIR = os.path.join(os.path.dirname(__file__), "..", "data")
 SCRIPT_NAME = "dummy_script.py"
@@ -111,6 +112,23 @@ def test_prepare_container_def(time, sagemaker_session):
         "ModelDataUrl": MODEL_DATA,
     }
 
+@patch("shutil.rmtree", MagicMock())
+@patch("tarfile.open", MagicMock())
+@patch("os.listdir", MagicMock(return_value=["blah.py"]))
+@patch("time.strftime", return_value=TIMESTAMP)
+def test_prepare_container_def_s3_src(time, sagemaker_session):
+    model = DummyFrameworkModel(sagemaker_session, source_dir=S3_SOURCE_DIR)
+    assert model.prepare_container_def(INSTANCE_TYPE) == {
+        "Environment": {
+            "SAGEMAKER_PROGRAM": ENTRY_POINT,
+            "SAGEMAKER_SUBMIT_DIRECTORY": "s3://somebucket/sourcedir.tar.gz",
+            "SAGEMAKER_CONTAINER_LOG_LEVEL": "20",
+            "SAGEMAKER_REGION": REGION,
+        },
+        "Image": MODEL_IMAGE,
+        "ModelDataUrl": MODEL_DATA,
+    }
+S3_SOURCE_DIR
 
 @patch("shutil.rmtree", MagicMock())
 @patch("tarfile.open", MagicMock())