add unit and integ tests

Author: mollyheamazon

src/sagemaker/modules/train/model_trainer.py

@@ -409,7 +409,9 @@ def _validate_source_code(self, source_code: Optional[SourceCode]):
                     )
                 if not _is_valid_path(source_dir) and not _is_valid_s3_uri(source_dir):
                     raise ValueError(
-                        f"Invalid 'source_dir' path: {source_dir}. " + "Must be a valid directory.",
+                        f"Invalid 'source_dir' path: {source_dir}. "
+                        + "Must be a valid local directory, "
+                        "s3 uri or path to tar.gz file stored locally or in s3.",
                     )
                 if requirements:
                     if not source_dir.endswith(".tar.gz"):
@@ -851,10 +853,7 @@ def _prepare_train_script(
         if source_code.source_dir:
             working_dir = f"cd {SM_CODE_CONTAINER_PATH} \n"
             if source_code.source_dir.endswith(".tar.gz"):
-                if source_code.source_dir.startswith("s3://"):
-                    tarfile_name = os.path.basename(source_code.source_dir)
-                else:
-                    tarfile_name = source_code.source_dir
+                tarfile_name = os.path.basename(source_code.source_dir)
                 working_dir += f"tar --strip-components=1 -xzf {tarfile_name} \n"
 
         if base_command:

tests/data/modules/script_mode/code.tar.gz

@@ -44,6 +44,24 @@
 
 DEFAULT_CPU_IMAGE = "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.0.0-cpu-py310"
 
+TAR_FILE_SOURCE_DIR = f"{DATA_DIR}/modules/script_mode/code.tar.gz"
+TAR_FILE_SOURCE_CODE = SourceCode(
+    source_dir=TAR_FILE_SOURCE_DIR,
+    requirements="requirements.txt",
+    entry_script="custom_script.py",
+)
+
+
+def test_source_dir_local_tar_file(modules_sagemaker_session):
+    model_trainer = ModelTrainer(
+        sagemaker_session=modules_sagemaker_session,
+        training_image=DEFAULT_CPU_IMAGE,
+        source_code=TAR_FILE_SOURCE_CODE,
+        base_job_name="source_dir_local_tar_file",
+    )
+
+    model_trainer.train()
+
 
 def test_hp_contract_basic_py_script(modules_sagemaker_session):
     model_trainer = ModelTrainer(

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

@@ -163,13 +163,46 @@ def model_trainer():
             },
             "should_throw": False,
         },
+        {
+            "init_params": {
+                "training_image": DEFAULT_IMAGE,
+                "source_code": SourceCode(
+                    source_dir=f"{DEFAULT_SOURCE_DIR}/code.tar.gz",
+                    entry_script="custom_script.py",
+                ),
+            },
+            "should_throw": False,
+        },
+        {
+            "init_params": {
+                "training_image": DEFAULT_IMAGE,
+                "source_code": SourceCode(
+                    source_dir="s3://bucket/code",
+                    entry_script="custom_script.py",
+                ),
+            },
+            "should_throw": False,
+        },
+        {
+            "init_params": {
+                "training_image": DEFAULT_IMAGE,
+                "source_code": SourceCode(
+                    source_dir="s3://bucket/code/code.tar.gz",
+                    entry_script="custom_script.py",
+                ),
+            },
+            "should_throw": False,
+        },
     ],
     ids=[
         "no_params",
         "training_image_and_algorithm_name",
         "only_training_image",
         "unsupported_source_code",
-        "supported_source_code",
+        "supported_source_code_local_dir",
+        "supported_source_code_local_tar_file",
+        "supported_source_code_s3_dir",
+        "supported_source_code_s3_tar_file",
     ],
 )
 def test_model_trainer_param_validation(test_case, modules_session):