fix: make configs safer

Author: benieric

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
 
@@ -94,6 +94,8 @@ class SourceCode(BaseModel):
             If not specified, entry_script must be provided.
     """
 
+    model_config = ConfigDict(validate_assignment=True, extra="forbid")
+
     source_dir: Optional[str] = None
     requirements: Optional[str] = None
     entry_script: Optional[str] = None
@@ -215,5 +217,7 @@ class InputData(BaseModel):
             S3DataSource object, or FileSystemDataSource object.
     """
 
+    model_config = ConfigDict(validate_assignment=True, extra="forbid")
+
     channel_name: str = None
     data_source: Union[str, FileSystemDataSource, S3DataSource] = None

src/sagemaker/modules/distributed.py

@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 
 from typing import Optional, Dict, Any, List
-from pydantic import BaseModel, PrivateAttr
+from pydantic import BaseModel, PrivateAttr, ConfigDict
 from sagemaker.modules.utils import safe_serialize
 
 
@@ -53,6 +53,8 @@ class SMP(BaseModel):
             parallelism or expert parallelism.
     """
 
+    model_config = ConfigDict(validate_assignment=True, extra="forbid")
+
     hybrid_shard_degree: Optional[int] = None
     sm_activation_offloading: Optional[bool] = None
     activation_loading_horizon: Optional[int] = None
@@ -75,6 +77,8 @@ def _to_mp_hyperparameters(self) -> Dict[str, Any]:
 class DistributedConfig(BaseModel):
     """Base class for distributed training configurations."""
 
+    model_config = ConfigDict(validate_assignment=True, extra="forbid")
+
     _type: str = PrivateAttr()
 
     def model_dump(self, *args, **kwargs):

src/sagemaker/modules/train/model_trainer.py

@@ -792,14 +792,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()
@@ -831,6 +831,10 @@ def _prepare_train_script(
                     + "Only .py and .sh scripts are supported."
                 )
             execute_driver = EXECUTE_BASIC_SCRIPT_DRIVER
+        else:
+            raise ValueError(
+                f"Invalid configuration, please provide a valid SourceCode: {source_code}"
+            )
 
         train_script = TRAIN_SCRIPT_TEMPLATE.format(
             working_dir=working_dir,

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

@@ -438,7 +438,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": EXECUTE_MPI_DRIVER,
             "expected_hyperparameters": {},
@@ -1059,3 +1059,12 @@ 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)