Merge branch 'master' into master-distributed-config-extensible

Author: benieric

doc/frameworks/pytorch/using_pytorch.rst

@@ -28,8 +28,6 @@ To train a PyTorch model by using the SageMaker Python SDK:
 Prepare a PyTorch Training Script
 =================================
 
-Your PyTorch training script must be a Python 3.6 compatible source file.
-
 Prepare your script in a separate source file than the notebook, terminal session, or source file you're
 using to submit the script to SageMaker via a ``PyTorch`` Estimator. This will be discussed in further detail below.
 

src/sagemaker/modules/configs.py

@@ -22,7 +22,7 @@
 from __future__ import absolute_import
 
 from typing import Optional, Union
-from pydantic import BaseModel, model_validator
+from pydantic import BaseModel, model_validator, ConfigDict
 
 import sagemaker_core.shapes as shapes
 
@@ -74,7 +74,13 @@
 ]
 
 
-class SourceCode(BaseModel):
+class BaseConfig(BaseModel):
+    """BaseConfig"""
+
+    model_config = ConfigDict(validate_assignment=True, extra="forbid")
+
+
+class SourceCode(BaseConfig):
     """SourceCode.
 
     The SourceCode class allows the user to specify the source code location, dependencies,
@@ -194,7 +200,7 @@ def _to_vpc_config(self) -> shapes.VpcConfig:
         return shapes.VpcConfig(**filtered_dict)
 
 
-class InputData(BaseModel):
+class InputData(BaseConfig):
     """InputData.
 
     This config allows the user to specify an input data source for the training job.

src/sagemaker/modules/distributed.py

@@ -17,12 +17,14 @@
 
 from abc import ABC, abstractmethod
 from typing import Optional, Dict, Any, List
-from pydantic import BaseModel
+from pydantic import PrivateAttr
+
 from sagemaker.modules.utils import safe_serialize
 from sagemaker.modules.constants import SM_DRIVERS_LOCAL_PATH
+from sagemaker.modules.configs import BaseConfig
 
 
-class SMP(BaseModel):
+class SMP(BaseConfig):
     """SMP.
 
     This class is used for configuring the SageMaker Model Parallelism v2 parameters.
@@ -76,7 +78,7 @@ def _to_mp_hyperparameters(self) -> Dict[str, Any]:
         return hyperparameters
 
 
-class DistributedConfig(BaseModel, ABC):
+class DistributedConfig(BaseConfig, ABC):
     """Abstract base class for distributed training configurations.
 
     This class defines the interface that all distributed training configurations

src/sagemaker/modules/train/model_trainer.py

@@ -204,7 +204,9 @@ class ModelTrainer(BaseModel):
             "LOCAL_CONTAINER" mode.
     """
 
-    model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True, validate_assignment=True, extra="forbid"
+    )
 
     training_mode: Mode = Mode.SAGEMAKER_TRAINING_JOB
     sagemaker_session: Optional[Session] = None
@@ -362,9 +364,10 @@ def _populate_intelligent_defaults_from_model_trainer_space(self):
 
     def __del__(self):
         """Destructor method to clean up the temporary directory."""
-        # Clean up the temporary directory if it exists
-        if self._temp_recipe_train_dir is not None:
-            self._temp_recipe_train_dir.cleanup()
+        # Clean up the temporary directory if it exists and class was initialized
+        if hasattr(self, "__pydantic_fields_set__"):
+            if self._temp_recipe_train_dir is not None:
+                self._temp_recipe_train_dir.cleanup()
 
     def _validate_training_image_and_algorithm_name(
         self, training_image: Optional[str], algorithm_name: Optional[str]
@@ -796,14 +799,14 @@ def _prepare_train_script(
         """Prepare the training script to be executed in the training job container.
 
         Args:
-            source_code (SourceCodeConfig): The source code configuration.
+            source_code (SourceCode): The source code configuration.
         """
 
         base_command = ""
         if source_code.command:
             if source_code.entry_script:
                 logger.warning(
-                    "Both 'command' and 'entry_script' are provided in the SourceCodeConfig. "
+                    "Both 'command' and 'entry_script' are provided in the SourceCode. "
                     + "Defaulting to 'command'."
                 )
             base_command = source_code.command.split()
@@ -832,6 +835,13 @@ def _prepare_train_script(
                     + "Only .py and .sh scripts are supported."
                 )
             execute_driver = EXECUTE_BASIC_SCRIPT_DRIVER
+        else:
+            # This should never be reached, as the source_code should have been validated.
+            raise ValueError(
+                f"Unsupported SourceCode or DistributedConfig: {source_code}, {distributed}."
+                + "Please provide a valid configuration with atleast one of 'command'"
+                + " or entry_script'."
+            )
 
         train_script = TRAIN_SCRIPT_TEMPLATE.format(
             working_dir=working_dir,

src/sagemaker/serve/detector/dependency_manager.py

@@ -34,22 +34,34 @@ def capture_dependencies(dependencies: dict, work_dir: Path, capture_all: bool =
     """Placeholder docstring"""
     path = work_dir.joinpath("requirements.txt")
     if "auto" in dependencies and dependencies["auto"]:
+        import site
+
+        pkl_path = work_dir.joinpath(PKL_FILE_NAME)
+        dest_path = path
+        site_packages_dir = site.getsitepackages()[0]
+        pickle_command_dir = "/sagemaker/serve/detector"
+
         command = [
             sys.executable,
-            Path(__file__).parent.joinpath("pickle_dependencies.py"),
-            "--pkl_path",
-            work_dir.joinpath(PKL_FILE_NAME),
-            "--dest",
-            path,
+            "-c",
         ]
 
         if capture_all:
-            command.append("--capture_all")
+            command.append(
+                f"from pickle_dependencies import get_all_requirements;"
+                f'get_all_requirements("{dest_path}")'
+            )
+        else:
+            command.append(
+                f"from pickle_dependencies import get_requirements_for_pkl_file;"
+                f'get_requirements_for_pkl_file("{pkl_path}", "{dest_path}")'
+            )
 
         subprocess.run(
             command,
             env={"SETUPTOOLS_USE_DISTUTILS": "stdlib"},
             check=True,
+            cwd=site_packages_dir + pickle_command_dir,
         )
 
         with open(path, "r") as f:

src/sagemaker/serve/detector/pickle_dependencies.py

@@ -3,7 +3,6 @@
 from __future__ import absolute_import
 from pathlib import Path
 from typing import List
-import argparse
 import email.parser
 import email.policy
 import json
@@ -129,32 +128,3 @@ def get_all_requirements(dest: Path):
             version = package_info.get("version")
 
             out.write(f"{name}=={version}\n")
-
-
-def parse_args():
-    """Placeholder docstring"""
-    parser = argparse.ArgumentParser(
-        prog="pkl_requirements", description="Generates a requirements.txt for a cloudpickle file"
-    )
-    parser.add_argument("--pkl_path", required=True, help="path of the pkl file")
-    parser.add_argument("--dest", required=True, help="path of the destination requirements.txt")
-    parser.add_argument(
-        "--capture_all",
-        action="store_true",
-        help="capture all dependencies in current environment",
-    )
-    args = parser.parse_args()
-    return (Path(args.pkl_path), Path(args.dest), args.capture_all)
-
-
-def main():
-    """Placeholder docstring"""
-    pkl_path, dest, capture_all = parse_args()
-    if capture_all:
-        get_all_requirements(dest)
-    else:
-        get_requirements_for_pkl_file(pkl_path, dest)
-
-
-if __name__ == "__main__":
-    main()

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

@@ -18,6 +18,7 @@
 import json
 import os
 import pytest
+from pydantic import ValidationError
 from unittest.mock import patch, MagicMock, ANY
 
 from sagemaker import image_uris
@@ -442,7 +443,7 @@ def test_create_input_data_channel(mock_default_bucket, mock_upload_data, model_
         {
             "source_code": DEFAULT_SOURCE_CODE,
             "distributed": MPI(
-                custom_mpi_options=["-x", "VAR1", "-x", "VAR2"],
+                mpi_additional_options=["-x", "VAR1", "-x", "VAR2"],
             ),
             "expected_template": EXEUCTE_DISTRIBUTED_DRIVER.format(
                 driver_name="MPI", driver_script="mpi_driver.py"
@@ -1059,3 +1060,36 @@ def mock_upload_data(path, bucket, key_prefix):
         hyper_parameters=hyperparameters,
         environment=environment,
     )
+
+
+def test_safe_configs():
+    # Test extra fails
+    with pytest.raises(ValueError):
+        SourceCode(entry_point="train.py")
+    # Test invalid type fails
+    with pytest.raises(ValueError):
+        SourceCode(entry_script=1)
+
+
+@patch("sagemaker.modules.train.model_trainer.TemporaryDirectory")
+def test_destructor_cleanup(mock_tmp_dir, modules_session):
+
+    with pytest.raises(ValidationError):
+        model_trainer = ModelTrainer(
+            training_image=DEFAULT_IMAGE,
+            role=DEFAULT_ROLE,
+            sagemaker_session=modules_session,
+            compute="test",
+        )
+    mock_tmp_dir.cleanup.assert_not_called()
+
+    model_trainer = ModelTrainer(
+        training_image=DEFAULT_IMAGE,
+        role=DEFAULT_ROLE,
+        sagemaker_session=modules_session,
+        compute=DEFAULT_COMPUTE_CONFIG,
+    )
+    model_trainer._temp_recipe_train_dir = mock_tmp_dir
+    mock_tmp_dir.assert_not_called()
+    del model_trainer
+    mock_tmp_dir.cleanup.assert_called_once()