Allow ModelTrainer to accept hyperparameter file and create Hyperparameter class

Author: benieric

Diff for: src/sagemaker/modules/hyperparameters.py

@@ -0,0 +1,98 @@
+import os
+import json
+import dataclasses
+from typing import Any, Type, TypeVar
+
+from sagemaker.modules import logger
+
+T = TypeVar("T")
+
+
+class DictConfig:
+    """Class that supports both dict and dot notation access"""
+
+    def __init__(self, **kwargs):
+        # Store the original dict
+        self._data = kwargs
+
+        # Set all items as attributes for dot notation
+        for key, value in kwargs.items():
+            # Recursively convert nested dicts to DictConfig
+            if isinstance(value, dict):
+                value = DictConfig(**value)
+            setattr(self, key, value)
+
+    def __getitem__(self, key: str) -> Any:
+        """Enable dictionary-style access: config['key']"""
+        return self._data[key]
+
+    def __setitem__(self, key: str, value: Any):
+        """Enable dictionary-style assignment: config['key'] = value"""
+        self._data[key] = value
+        setattr(self, key, value)
+
+    def __str__(self) -> str:
+        """String representation"""
+        return str(self._data)
+
+    def __repr__(self) -> str:
+        """Detailed string representation"""
+        return f"DictConfig({self._data})"
+
+
+class Hyperparameters:
+    """Class to load hyperparameters in training container."""
+
+    @staticmethod
+    def load() -> DictConfig:
+        """Loads hyperparameters in training container
+
+        Example:
+
+        .. code:: python
+            from sagemaker.modules.hyperparameters import Hyperparameters
+
+            hps = Hyperparameters.load()
+            print(hps.batch_size)
+
+        Returns:
+            DictConfig: hyperparameters as a DictConfig object
+        """
+        hps = json.loads(os.environ.get("SM_HPS", "{}"))
+        if not hps:
+            logger.warning("No hyperparameters found in SM_HPS environment variable.")
+        return DictConfig(**hps)
+
+    @staticmethod
+    def load_structured(dataclass_type: Type[T]) -> T:
+        """Loads hyperparameters as a structured dataclass
+
+        Example:
+
+        .. code:: python
+            from sagemaker.modules.hyperparameters import Hyperparameters
+
+            @dataclass
+            class TrainingConfig:
+                batch_size: int
+                learning_rate: float
+
+            config = Hyperparameters.load_structured(TrainingConfig)
+            print(config.batch_size) # typed int
+
+        Args:
+            dataclass_type: Dataclass type to structure the config
+
+        Returns:
+            dataclass_type: Instance of provided dataclass type
+        """
+
+        if not dataclasses.is_dataclass(dataclass_type):
+            raise ValueError(f"{dataclass_type} is not a dataclass type.")
+
+        hps = json.loads(os.environ.get("SM_HPS", "{}"))
+        if not hps:
+            logger.warning("No hyperparameters found in SM_HPS environment variable.")
+
+        # Convert hyperparameters to dataclass
+        return dataclass_type(**hps)

Diff for: src/sagemaker/modules/train/model_trainer.py

@@ -17,6 +17,7 @@
 import os
 import json
 import shutil
+import yaml
 from tempfile import TemporaryDirectory
 
 from typing import Optional, List, Union, Dict, Any, ClassVar
@@ -195,8 +196,9 @@ class ModelTrainer(BaseModel):
             Defaults to "File".
         environment (Optional[Dict[str, str]]):
             The environment variables for the training job.
-        hyperparameters (Optional[Dict[str, Any]]):
-            The hyperparameters for the training job.
+        hyperparameters (Optional[Union[Dict[str, Any], str]):
+            The hyperparameters for the training job. Can be a dictionary of hyperparameters
+            or a path to hyperparameters json/yaml file.
         tags (Optional[List[Tag]]):
             An array of key-value pairs. You can use tags to categorize your AWS resources
             in different ways, for example, by purpose, owner, or environment.
@@ -226,7 +228,7 @@ class ModelTrainer(BaseModel):
     checkpoint_config: Optional[CheckpointConfig] = None
     training_input_mode: Optional[str] = "File"
     environment: Optional[Dict[str, str]] = {}
-    hyperparameters: Optional[Dict[str, Any]] = {}
+    hyperparameters: Optional[Union[Dict[str, Any], str]] = {}
     tags: Optional[List[Tag]] = None
     local_container_root: Optional[str] = os.getcwd()
 
@@ -470,6 +472,17 @@ def model_post_init(self, __context: Any):
                 f"StoppingCondition not provided. Using default:\n{self.stopping_condition}"
             )
 
+        if self.hyperparameters and isinstance(self.hyperparameters, str):
+            if not os.path.exists(self.hyperparameters):
+                raise ValueError(f"Hyperparameters file not found: {self.hyperparameters}")
+            logger.info(f"Loading hyperparameters from file: {self.hyperparameters}")
+            if self.hyperparameters.endswith(".json"):
+                with open(self.hyperparameters, "r") as f:
+                    self.hyperparameters = json.load(f)
+            elif self.hyperparameters.endswith(".yaml"):
+                with open(self.hyperparameters, "r") as f:
+                    self.hyperparameters = yaml.safe_load(f)
+
         if self.training_mode == Mode.SAGEMAKER_TRAINING_JOB and self.output_data_config is None:
             session = self.sagemaker_session
             base_job_name = self.base_job_name