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

benieric · web-flow · commit f668f3bd3fdf · 2025-03-04T18:30:04.000-08:00
diff --git a/src/sagemaker/model.py b/src/sagemaker/model.py
@@ -745,6 +745,8 @@ def is_repack(self) -> bool:
         Returns:
             bool: if the source need to be repacked or not
         """
+        if self.source_dir is None or self.entry_point is None:
+            return False
         return self.source_dir and self.entry_point and not self.git_config
 
     def _upload_code(self, key_prefix: str, repack: bool = False) -> None:
@@ -2143,6 +2145,8 @@ def is_repack(self) -> bool:
         Returns:
             bool: if the source need to be repacked or not
         """
+        if self.source_dir is None or self.entry_point is None:
+            return False
         return self.source_dir and self.entry_point and not (self.key_prefix or self.git_config)
 
 
diff --git a/src/sagemaker/modules/train/model_trainer.py b/src/sagemaker/modules/train/model_trainer.py
@@ -18,8 +18,8 @@
 import json
 import shutil
 from tempfile import TemporaryDirectory
-
 from typing import Optional, List, Union, Dict, Any, ClassVar
+import yaml
 
 from graphene.utils.str_converters import to_camel_case, to_snake_case
 
@@ -194,8 +194,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.
@@ -225,7 +226,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()
 
@@ -469,6 +470,29 @@ 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}")
+            with open(self.hyperparameters, "r") as f:
+                contents = f.read()
+                try:
+                    self.hyperparameters = json.loads(contents)
+                    logger.debug("Hyperparameters loaded as JSON")
+                except json.JSONDecodeError:
+                    try:
+                        logger.info(f"contents: {contents}")
+                        self.hyperparameters = yaml.safe_load(contents)
+                        if not isinstance(self.hyperparameters, dict):
+                            raise ValueError("YAML contents must be a valid mapping")
+                        logger.info(f"hyperparameters: {self.hyperparameters}")
+                        logger.debug("Hyperparameters loaded as YAML")
+                    except (yaml.YAMLError, ValueError):
+                        raise ValueError(
+                            f"Invalid hyperparameters file: {self.hyperparameters}. "
+                            "Must be a valid JSON or YAML file."
+                        )
+
         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
diff --git a/src/sagemaker/pipeline.py b/src/sagemaker/pipeline.py
@@ -17,6 +17,8 @@
 
 import sagemaker
 from sagemaker import ModelMetrics, Model
+from sagemaker import local
+from sagemaker import session
 from sagemaker.config import (
     ENDPOINT_CONFIG_KMS_KEY_ID_PATH,
     MODEL_VPC_CONFIG_PATH,
@@ -560,3 +562,16 @@ def delete_model(self):
             raise ValueError("The SageMaker model must be created before attempting to delete.")
 
         self.sagemaker_session.delete_model(self.name)
+
+    def _init_sagemaker_session_if_does_not_exist(self, instance_type=None):
+        """Set ``self.sagemaker_session`` to ``LocalSession`` or ``Session`` if it's not already.
+
+        The type of session object is determined by the instance type.
+        """
+        if self.sagemaker_session:
+            return
+
+        if instance_type in ("local", "local_gpu"):
+            self.sagemaker_session = local.LocalSession(sagemaker_config=self._sagemaker_config)
+        else:
+            self.sagemaker_session = session.Session(sagemaker_config=self._sagemaker_config)
diff --git a/src/sagemaker/workflow/steps.py b/src/sagemaker/workflow/steps.py
@@ -645,6 +645,7 @@ def arguments(self) -> RequestType:
             request_dict = self.step_args
         else:
             if isinstance(self.model, PipelineModel):
+                self.model._init_sagemaker_session_if_does_not_exist()
                 request_dict = self.model.sagemaker_session._create_model_request(
                     name="",
                     role=self.model.role,
@@ -653,6 +654,7 @@ def arguments(self) -> RequestType:
                     enable_network_isolation=self.model.enable_network_isolation,
                 )
             else:
+                self.model._init_sagemaker_session_if_does_not_exist()
                 request_dict = self.model.sagemaker_session._create_model_request(
                     name="",
                     role=self.model.role,
diff --git a/tests/data/modules/params_script/hyperparameters.json b/tests/data/modules/params_script/hyperparameters.json
@@ -0,0 +1,15 @@
+{
+    "integer": 1,
+    "boolean": true,
+    "float": 3.14,
+    "string": "Hello World",
+    "list": [1, 2, 3],
+    "dict": {
+        "string": "value",
+        "integer": 3,
+        "float": 3.14,
+        "list": [1, 2, 3],
+        "dict": {"key": "value"},
+        "boolean": true
+    }
+}
diff --git a/tests/data/modules/params_script/hyperparameters.yaml b/tests/data/modules/params_script/hyperparameters.yaml
@@ -0,0 +1,19 @@
+integer: 1
+boolean: true
+float: 3.14
+string: "Hello World"
+list:
+  - 1
+  - 2
+  - 3
+dict:
+  string: value
+  integer: 3
+  float: 3.14
+  list:
+    - 1
+    - 2
+    - 3
+  dict:
+    key: value
+  boolean: true
diff --git a/tests/data/modules/params_script/requirements.txt b/tests/data/modules/params_script/requirements.txt
@@ -0,0 +1 @@
+omegaconf
diff --git a/tests/data/modules/params_script/train.py b/tests/data/modules/params_script/train.py
@@ -16,6 +16,9 @@
 import argparse
 import json
 import os
+from typing import List, Dict, Any
+from dataclasses import dataclass
+from omegaconf import OmegaConf
 
 EXPECTED_HYPERPARAMETERS = {
     "integer": 1,
@@ -26,6 +29,7 @@
     "dict": {
         "string": "value",
         "integer": 3,
+        "float": 3.14,
         "list": [1, 2, 3],
         "dict": {"key": "value"},
         "boolean": True,
@@ -117,7 +121,7 @@ def main():
     assert isinstance(params["dict"], dict)
 
     params = json.loads(os.environ["SM_TRAINING_ENV"])["hyperparameters"]
-    print(params)
+    print(f"SM_TRAINING_ENV -> hyperparameters: {params}")
     assert params["string"] == EXPECTED_HYPERPARAMETERS["string"]
     assert params["integer"] == EXPECTED_HYPERPARAMETERS["integer"]
     assert params["boolean"] == EXPECTED_HYPERPARAMETERS["boolean"]
@@ -132,9 +136,96 @@ def main():
     assert isinstance(params["float"], float)
     assert isinstance(params["list"], list)
     assert isinstance(params["dict"], dict)
-    print(f"SM_TRAINING_ENV -> hyperparameters: {params}")
 
-    print("Test passed.")
+    # Local JSON - DictConfig OmegaConf
+    params = OmegaConf.load("hyperparameters.json")
+
+    print(f"Local hyperparameters.json: {params}")
+    assert params.string == EXPECTED_HYPERPARAMETERS["string"]
+    assert params.integer == EXPECTED_HYPERPARAMETERS["integer"]
+    assert params.boolean == EXPECTED_HYPERPARAMETERS["boolean"]
+    assert params.float == EXPECTED_HYPERPARAMETERS["float"]
+    assert params.list == EXPECTED_HYPERPARAMETERS["list"]
+    assert params.dict == EXPECTED_HYPERPARAMETERS["dict"]
+    assert params.dict.string == EXPECTED_HYPERPARAMETERS["dict"]["string"]
+    assert params.dict.integer == EXPECTED_HYPERPARAMETERS["dict"]["integer"]
+    assert params.dict.boolean == EXPECTED_HYPERPARAMETERS["dict"]["boolean"]
+    assert params.dict.float == EXPECTED_HYPERPARAMETERS["dict"]["float"]
+    assert params.dict.list == EXPECTED_HYPERPARAMETERS["dict"]["list"]
+    assert params.dict.dict == EXPECTED_HYPERPARAMETERS["dict"]["dict"]
+
+    @dataclass
+    class DictConfig:
+        string: str
+        integer: int
+        boolean: bool
+        float: float
+        list: List[int]
+        dict: Dict[str, Any]
+
+    @dataclass
+    class HPConfig:
+        string: str
+        integer: int
+        boolean: bool
+        float: float
+        list: List[int]
+        dict: DictConfig
+
+    # Local JSON - Structured OmegaConf
+    hp_config: HPConfig = OmegaConf.merge(
+        OmegaConf.structured(HPConfig), OmegaConf.load("hyperparameters.json")
+    )
+    print(f"Local hyperparameters.json - Structured: {hp_config}")
+    assert hp_config.string == EXPECTED_HYPERPARAMETERS["string"]
+    assert hp_config.integer == EXPECTED_HYPERPARAMETERS["integer"]
+    assert hp_config.boolean == EXPECTED_HYPERPARAMETERS["boolean"]
+    assert hp_config.float == EXPECTED_HYPERPARAMETERS["float"]
+    assert hp_config.list == EXPECTED_HYPERPARAMETERS["list"]
+    assert hp_config.dict == EXPECTED_HYPERPARAMETERS["dict"]
+    assert hp_config.dict.string == EXPECTED_HYPERPARAMETERS["dict"]["string"]
+    assert hp_config.dict.integer == EXPECTED_HYPERPARAMETERS["dict"]["integer"]
+    assert hp_config.dict.boolean == EXPECTED_HYPERPARAMETERS["dict"]["boolean"]
+    assert hp_config.dict.float == EXPECTED_HYPERPARAMETERS["dict"]["float"]
+    assert hp_config.dict.list == EXPECTED_HYPERPARAMETERS["dict"]["list"]
+    assert hp_config.dict.dict == EXPECTED_HYPERPARAMETERS["dict"]["dict"]
+
+    # Local YAML - Structured OmegaConf
+    hp_config: HPConfig = OmegaConf.merge(
+        OmegaConf.structured(HPConfig), OmegaConf.load("hyperparameters.yaml")
+    )
+    print(f"Local hyperparameters.yaml - Structured: {hp_config}")
+    assert hp_config.string == EXPECTED_HYPERPARAMETERS["string"]
+    assert hp_config.integer == EXPECTED_HYPERPARAMETERS["integer"]
+    assert hp_config.boolean == EXPECTED_HYPERPARAMETERS["boolean"]
+    assert hp_config.float == EXPECTED_HYPERPARAMETERS["float"]
+    assert hp_config.list == EXPECTED_HYPERPARAMETERS["list"]
+    assert hp_config.dict == EXPECTED_HYPERPARAMETERS["dict"]
+    assert hp_config.dict.string == EXPECTED_HYPERPARAMETERS["dict"]["string"]
+    assert hp_config.dict.integer == EXPECTED_HYPERPARAMETERS["dict"]["integer"]
+    assert hp_config.dict.boolean == EXPECTED_HYPERPARAMETERS["dict"]["boolean"]
+    assert hp_config.dict.float == EXPECTED_HYPERPARAMETERS["dict"]["float"]
+    assert hp_config.dict.list == EXPECTED_HYPERPARAMETERS["dict"]["list"]
+    assert hp_config.dict.dict == EXPECTED_HYPERPARAMETERS["dict"]["dict"]
+    print(f"hyperparameters.yaml -> hyperparameters: {hp_config}")
+
+    # HP Dict - Structured OmegaConf
+    hp_dict = json.loads(os.environ["SM_HPS"])
+    hp_config: HPConfig = OmegaConf.merge(OmegaConf.structured(HPConfig), OmegaConf.create(hp_dict))
+    print(f"SM_HPS - Structured: {hp_config}")
+    assert hp_config.string == EXPECTED_HYPERPARAMETERS["string"]
+    assert hp_config.integer == EXPECTED_HYPERPARAMETERS["integer"]
+    assert hp_config.boolean == EXPECTED_HYPERPARAMETERS["boolean"]
+    assert hp_config.float == EXPECTED_HYPERPARAMETERS["float"]
+    assert hp_config.list == EXPECTED_HYPERPARAMETERS["list"]
+    assert hp_config.dict == EXPECTED_HYPERPARAMETERS["dict"]
+    assert hp_config.dict.string == EXPECTED_HYPERPARAMETERS["dict"]["string"]
+    assert hp_config.dict.integer == EXPECTED_HYPERPARAMETERS["dict"]["integer"]
+    assert hp_config.dict.boolean == EXPECTED_HYPERPARAMETERS["dict"]["boolean"]
+    assert hp_config.dict.float == EXPECTED_HYPERPARAMETERS["dict"]["float"]
+    assert hp_config.dict.list == EXPECTED_HYPERPARAMETERS["dict"]["list"]
+    assert hp_config.dict.dict == EXPECTED_HYPERPARAMETERS["dict"]["dict"]
+    print(f"SM_HPS -> hyperparameters: {hp_config}")
 
 
 if __name__ == "__main__":
diff --git a/tests/integ/sagemaker/modules/train/test_model_trainer.py b/tests/integ/sagemaker/modules/train/test_model_trainer.py
diff --git a/tests/unit/sagemaker/model/test_framework_model.py b/tests/unit/sagemaker/model/test_framework_model.py
diff --git a/tests/unit/sagemaker/model/test_model.py b/tests/unit/sagemaker/model/test_model.py
diff --git a/tests/unit/sagemaker/modules/train/test_model_trainer.py b/tests/unit/sagemaker/modules/train/test_model_trainer.py
