Create BaseConfig

Author: benieric

src/sagemaker/modules/configs.py

@@ -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,
@@ -94,8 +100,6 @@ 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
@@ -196,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.
@@ -217,7 +221,5 @@ 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,11 +14,12 @@
 from __future__ import absolute_import
 
 from typing import Optional, Dict, Any, List
-from pydantic import BaseModel, PrivateAttr, ConfigDict
+from pydantic import PrivateAttr
 from sagemaker.modules.utils import safe_serialize
+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.
@@ -53,8 +54,6 @@ 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
@@ -74,11 +73,9 @@ def _to_mp_hyperparameters(self) -> Dict[str, Any]:
         return hyperparameters
 
 
-class DistributedConfig(BaseModel):
+class DistributedConfig(BaseConfig):
     """Base class for distributed training configurations."""
 
-    model_config = ConfigDict(validate_assignment=True, extra="forbid")
-
     _type: str = PrivateAttr()
 
     def model_dump(self, *args, **kwargs):