update logic and unit test to raise value error if the file is not .tar.gz

Author: mollyheamazon

src/sagemaker/modules/train/model_trainer.py

@@ -407,7 +407,18 @@ def _validate_source_code(self, source_code: Optional[SourceCode]):
                         "If 'requirements' or 'entry_script' is provided in 'source_code', "
                         + "'source_dir' must also be provided.",
                     )
-                if not _is_valid_path(source_dir) and not _is_valid_s3_uri(source_dir):
+                if not (
+                    _is_valid_path(source_dir, path_type="Directory")
+                    or _is_valid_s3_uri(source_dir, path_type="Directory")
+                    or (
+                        _is_valid_path(source_dir, path_type="File")
+                        and source_dir.endswith(".tar.gz")
+                    )
+                    or (
+                        _is_valid_s3_uri(source_dir, path_type="File")
+                        and source_dir.endswith(".tar.gz")
+                    )
+                ):
                     raise ValueError(
                         f"Invalid 'source_dir' path: {source_dir}. "
                         + "Must be a valid local directory, "

tests/unit/sagemaker/modules/train/test_model_trainer.py

@@ -92,9 +92,6 @@
     source_dir=DEFAULT_SOURCE_DIR,
     entry_script="custom_script.py",
 )
-UNSUPPORTED_SOURCE_CODE = SourceCode(
-    entry_script="train.py",
-)
 DEFAULT_ENTRYPOINT = ["/bin/bash"]
 DEFAULT_ARGUMENTS = [
     "-c",
@@ -152,7 +149,19 @@ def model_trainer():
         {
             "init_params": {
                 "training_image": DEFAULT_IMAGE,
-                "source_code": UNSUPPORTED_SOURCE_CODE,
+                "source_code": SourceCode(
+                    entry_script="train.py",
+                ),
+            },
+            "should_throw": True,
+        },
+        {
+            "init_params": {
+                "training_image": DEFAULT_IMAGE,
+                "source_code": SourceCode(
+                    source_dir="s3://bucket/requirements.txt",
+                    entry_script="custom_script.py",
+                ),
             },
             "should_throw": True,
         },
@@ -177,7 +186,7 @@ def model_trainer():
             "init_params": {
                 "training_image": DEFAULT_IMAGE,
                 "source_code": SourceCode(
-                    source_dir="s3://bucket/code",
+                    source_dir="s3://bucket/code/",
                     entry_script="custom_script.py",
                 ),
             },
@@ -198,7 +207,8 @@ def model_trainer():
         "no_params",
         "training_image_and_algorithm_name",
         "only_training_image",
-        "unsupported_source_code",
+        "unsupported_source_code_missing_source_dir",
+        "unsupported_source_code_s3_other_file",
         "supported_source_code_local_dir",
         "supported_source_code_local_tar_file",
         "supported_source_code_s3_dir",