fix: address review comments (iteration #1)

Author: aviruthen

sagemaker-core/src/sagemaker/core/training/configs.py

@@ -129,6 +129,7 @@ class SourceCode(BaseConfig):
     ]
     dependencies: Optional[List[str]] = None
 
+
 class OutputDataConfig(shapes.OutputDataConfig):
     """OutputDataConfig.
 

sagemaker-train/src/sagemaker/train/constants.py

@@ -68,4 +68,4 @@
 
 SM_RECIPE = "recipe"
 SM_RECIPE_YAML = "recipe.yaml"
-SM_RECIPE_CONTAINER_PATH = f"/opt/ml/input/data/recipe/{SM_RECIPE_YAML}"
+SM_RECIPE_CONTAINER_PATH = f"/opt/ml/input/data/recipe/{SM_RECIPE_YAML}"

sagemaker-train/src/sagemaker/train/model_trainer.py

@@ -272,6 +272,7 @@ class ModelTrainer(BaseModel):
 
     # Private Attributes for AWS_Batch
     _temp_code_dir: Optional[TemporaryDirectory] = PrivateAttr(default=None)
+    _temp_deps_dir: Optional[TemporaryDirectory] = PrivateAttr(default=None)
 
     CONFIGURABLE_ATTRIBUTES: ClassVar[List[str]] = [
         "role",
@@ -411,6 +412,8 @@ def __del__(self):
                 self._temp_recipe_train_dir.cleanup()
             if self._temp_code_dir is not None:
                 self._temp_code_dir.cleanup()
+            if self._temp_deps_dir is not None:
+                self._temp_deps_dir.cleanup()
 
     def _validate_training_image_and_algorithm_name(
         self, training_image: Optional[str], algorithm_name: Optional[str]
@@ -489,10 +492,10 @@ def _validate_source_code(self, source_code: Optional[SourceCode]):
                             )
             if source_code.dependencies:
                 for dep_path in source_code.dependencies:
-                    if not _is_valid_path(dep_path):
+                    if not _is_valid_path(dep_path, path_type="Directory"):
                         raise ValueError(
                             f"Invalid dependency path: {dep_path}. "
-                            "Each dependency must be a valid local directory or file path."
+                            "Each dependency must be a valid local directory path."
                         )
 
     @staticmethod
@@ -667,17 +670,14 @@ def _create_training_job_args(
             # If dependencies are provided, create a channel for the dependencies
             # The dependencies will be mounted at /opt/ml/input/data/sm_dependencies
             if self.source_code.dependencies:
-                deps_tmp_dir = TemporaryDirectory()
+                self._temp_deps_dir = TemporaryDirectory()
                 for dep_path in self.source_code.dependencies:
                     dep_basename = os.path.basename(os.path.normpath(dep_path))
-                    dest_path = os.path.join(deps_tmp_dir.name, dep_basename)
-                    if os.path.isdir(dep_path):
-                        shutil.copytree(dep_path, dest_path, dirs_exist_ok=True)
-                    else:
-                        shutil.copy2(dep_path, dest_path)
+                    dest_path = os.path.join(self._temp_deps_dir.name, dep_basename)
+                    shutil.copytree(dep_path, dest_path, dirs_exist_ok=True)
                 dependencies_channel = self.create_input_data_channel(
                     channel_name=SM_DEPENDENCIES,
-                    data_source=deps_tmp_dir.name,
+                    data_source=self._temp_deps_dir.name,
                     key_prefix=input_data_key_prefix,
                 )
                 final_input_data_config.append(dependencies_channel)
@@ -841,6 +841,9 @@ def train(
             local_container.train(wait)
         if self._temp_code_dir is not None:
             self._temp_code_dir.cleanup()
+        if self._temp_deps_dir is not None:
+            self._temp_deps_dir.cleanup()
+            self._temp_deps_dir = None
 
 
     def create_input_data_channel(

sagemaker-train/src/sagemaker/train/templates.py

@@ -42,14 +42,23 @@
 INSTALL_DEPENDENCIES = """
 echo "Setting up additional dependencies"
 if [ -d /opt/ml/input/data/sm_dependencies ]; then
-    for dep_dir in /opt/ml/input/data/sm_dependencies/*/; do
-        if [ -d "$dep_dir" ]; then
-            echo "Adding $dep_dir to PYTHONPATH"
-            export PYTHONPATH="$dep_dir:$PYTHONPATH"
+    for dep in /opt/ml/input/data/sm_dependencies/*; do
+        if [ -d "$dep" ]; then
+            echo "Adding directory $dep to PYTHONPATH"
+            export PYTHONPATH="$dep:$PYTHONPATH"
+        elif [ -f "$dep" ]; then
+            case "$dep" in
+                *.whl|*.tar.gz)
+                    echo "Installing package $dep via pip"
+                    $SM_PIP_CMD install "$dep"
+                    ;;
+                *)
+                    echo "Adding parent directory of $dep to PYTHONPATH"
+                    export PYTHONPATH="/opt/ml/input/data/sm_dependencies:$PYTHONPATH"
+                    ;;
+            esac
         fi
     done
-    # Also add the root dependencies dir in case of single files
-    export PYTHONPATH="/opt/ml/input/data/sm_dependencies:$PYTHONPATH"
 fi
 """
 

sagemaker-train/tests/unit/train/test_model_trainer_dependencies.py

@@ -11,7 +11,7 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 """Tests for ModelTrainer dependencies feature."""
-from __future__ import absolute_import
+from __future__ import annotations
 
 import os
 import shutil
@@ -53,7 +53,7 @@
 )
 
 
-@pytest.fixture(scope="module", autouse=True)
+@pytest.fixture(autouse=True)
 def modules_session():
     with patch("sagemaker.train.Session", spec=Session) as session_mock:
         session_instance = session_mock.return_value
@@ -107,6 +107,27 @@ def test_validate_source_code_with_invalid_dependency_path_raises_value_error(mo
         )
 
 
+def test_validate_source_code_with_file_dependency_raises_value_error(modules_session):
+    """Verify _validate_source_code raises ValueError when a dependency is a file, not a directory."""
+    dep_file = tempfile.NamedTemporaryFile(suffix=".py", delete=False)
+    dep_file.close()
+    try:
+        with pytest.raises(ValueError, match="Invalid dependency path"):
+            ModelTrainer(
+                training_image=DEFAULT_IMAGE,
+                role=DEFAULT_ROLE,
+                sagemaker_session=modules_session,
+                compute=DEFAULT_COMPUTE_CONFIG,
+                source_code=SourceCode(
+                    source_dir=DEFAULT_SOURCE_DIR,
+                    entry_script="custom_script.py",
+                    dependencies=[dep_file.name],
+                ),
+            )
+    finally:
+        os.unlink(dep_file.name)
+
+
 def test_source_code_dependencies_validation_with_valid_dirs(modules_session):
     """Create SourceCode with dependencies pointing to multiple valid directories."""
     dep_dir1 = tempfile.mkdtemp()
@@ -250,23 +271,39 @@ def test_train_with_dependencies_generates_pythonpath_setup_in_train_script(
                 dependencies=[dep_dir],
             ),
         )
-        trainer.train()
-
-        # Check the generated train script in the temp directory
-        assert trainer._temp_code_dir is not None or True  # temp dir may be cleaned up
-        # The key assertion is that the template was used - check via the training job call
-        mock_training_job.create.assert_called_once()
+        # Call _create_training_job_args to generate the train script without cleanup
+        trainer._create_training_job_args()
+
+        # Read the generated train script and verify it contains PYTHONPATH setup
+        assert trainer._temp_code_dir is not None
+        train_script_path = os.path.join(trainer._temp_code_dir.name, "sm_train.sh")
+        with open(train_script_path) as f:
+            script_content = f.read()
+        assert "sm_dependencies" in script_content
+        assert "PYTHONPATH" in script_content
+        assert "Setting up additional dependencies" in script_content
     finally:
+        if trainer._temp_code_dir is not None:
+            trainer._temp_code_dir.cleanup()
+        if trainer._temp_deps_dir is not None:
+            trainer._temp_deps_dir.cleanup()
         shutil.rmtree(dep_dir)
 
 
-def test_dependencies_copied_to_temp_dir_preserving_basenames(modules_session):
+@patch("sagemaker.train.model_trainer.TrainingJob")
+def test_dependencies_copied_to_temp_dir_preserving_basenames(
+    mock_training_job, modules_session
+):
     """Verify that each dependency directory's basename is preserved when copied."""
     dep_dir = tempfile.mkdtemp(suffix="_mylib")
     sub_file = os.path.join(dep_dir, "module.py")
     with open(sub_file, "w") as f:
         f.write("# test module")
 
+    modules_session.upload_data.return_value = (
+        f"s3://{DEFAULT_BUCKET}/prefix/sm_dependencies"
+    )
+
     try:
         trainer = ModelTrainer(
             training_image=DEFAULT_IMAGE,
@@ -281,9 +318,22 @@ def test_dependencies_copied_to_temp_dir_preserving_basenames(modules_session):
                 dependencies=[dep_dir],
             ),
         )
-        # Verify the source_code has the dependencies set
-        assert trainer.source_code.dependencies == [dep_dir]
+        # Call _create_training_job_args to trigger the copy
+        trainer._create_training_job_args()
+
+        # Verify the dependencies were copied preserving basenames
+        assert trainer._temp_deps_dir is not None
         dep_basename = os.path.basename(os.path.normpath(dep_dir))
-        assert dep_basename.endswith("_mylib")
+        copied_dep_path = os.path.join(trainer._temp_deps_dir.name, dep_basename)
+        assert os.path.isdir(copied_dep_path), (
+            f"Expected dependency directory {copied_dep_path} to exist"
+        )
+        assert os.path.isfile(os.path.join(copied_dep_path, "module.py")), (
+            "Expected module.py to be copied into the dependency directory"
+        )
     finally:
+        if trainer._temp_code_dir is not None:
+            trainer._temp_code_dir.cleanup()
+        if trainer._temp_deps_dir is not None:
+            trainer._temp_deps_dir.cleanup()
         shutil.rmtree(dep_dir)