fix tests

Author: benieric

src/sagemaker/modules/train/model_trainer.py

@@ -482,7 +482,11 @@ def model_post_init(self, __context: Any):
                     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(f"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(

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

@@ -42,7 +42,7 @@
     entry_script="train.py",
 )
 
-DEFAULT_CPU_IMAGE = "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.0.0-cpu-py31"
+DEFAULT_CPU_IMAGE = "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.0.0-cpu-py310"
 
 
 def test_hp_contract_basic_py_script(modules_sagemaker_session):

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

@@ -1148,8 +1148,21 @@ def test_hyperparameters_not_exist(modules_session):
 @patch("os.path.exists")
 def test_hyperparameters_invalid(mock_exists, modules_session):
     mock_exists.return_value = True
-    # Must be valid YAML or JSON
-    mock_file_open = mock_open(read_data="invalid")
+
+    # YAML contents must be a valid mapping
+    mock_file_open = mock_open(read_data="- item1\n- item2")
+    with patch("builtins.open", mock_file_open):
+        with pytest.raises(ValueError, match="Must be a valid JSON or YAML file."):
+            ModelTrainer(
+                training_image=DEFAULT_IMAGE,
+                role=DEFAULT_ROLE,
+                sagemaker_session=modules_session,
+                compute=DEFAULT_COMPUTE_CONFIG,
+                hyperparameters="hyperparameters.yaml",
+            )
+
+    # Must be valid YAML
+    mock_file_open = mock_open(read_data="* invalid")
     with patch("builtins.open", mock_file_open):
         with pytest.raises(ValueError, match="Must be a valid JSON or YAML file."):
             ModelTrainer(

tox.ini

@@ -83,8 +83,8 @@ passenv =
 commands =
     python -c "import os; os.system('install-custom-pkgs --install-boto-wheels')"
     pip install 'apache-airflow==2.9.3' --constraint "https://raw.githubusercontent.com/apache/airflow/constraints-2.9.3/constraints-3.8.txt"
-    pip install 'torch==2.0.1+cpu' -f 'https://download.pytorch.org/whl/torch_stable.html'
-    pip install 'torchvision==0.15.2+cpu' -f 'https://download.pytorch.org/whl/torch_stable.html'
+    pip install 'torch==2.0.1' -f 'https://download.pytorch.org/whl/torch_stable.html'
+    pip install 'torchvision==0.15.2' -f 'https://download.pytorch.org/whl/torch_stable.html'
     pip install 'dill>=0.3.8'
 
     pytest {posargs}