feature: add model_type hyperparameter support for Nova recipes (#5291)

Co-authored-by: xibei chen <[email protected]>

Author: sylvie7788

src/sagemaker/pytorch/estimator.py

@@ -1180,6 +1180,9 @@ def _setup_for_nova_recipe(
         # Set up Nova-specific configuration
         run_config = recipe.get("run", {})
         model_name_or_path = run_config.get("model_name_or_path")
+        # Set hyperparameters model_type
+        model_type = run_config.get("model_type")
+        args["hyperparameters"]["model_type"] = model_type
 
         # Set hyperparameters based on model_name_or_path
         if model_name_or_path:

tests/unit/test_pytorch_nova.py

@@ -795,3 +795,40 @@ def test_setup_for_nova_recipe_with_distillation(mock_resolve_save, sagemaker_se
                 pytorch._hyperparameters.get("role_arn")
                 == "arn:aws:iam::123456789012:role/SageMakerRole"
             )
+
+
+@patch("sagemaker.pytorch.estimator.PyTorch._recipe_resolve_and_save")
+def test_setup_for_nova_recipe_sets_model_type(mock_resolve_save, sagemaker_session):
+    """Test that _setup_for_nova_recipe correctly sets model_type hyperparameter."""
+    # Create a mock nova recipe with model_type
+    recipe = OmegaConf.create(
+        {
+            "run": {
+                "model_type": "amazon.nova.llama-2-7b",
+                "model_name_or_path": "llama/llama-2-7b",
+                "replicas": 1,
+            }
+        }
+    )
+
+    with patch(
+        "sagemaker.pytorch.estimator.PyTorch._recipe_load", return_value=("nova_recipe", recipe)
+    ):
+        mock_resolve_save.return_value = recipe
+
+        pytorch = PyTorch(
+            training_recipe="nova_recipe",
+            role=ROLE,
+            sagemaker_session=sagemaker_session,
+            instance_count=INSTANCE_COUNT,
+            instance_type=INSTANCE_TYPE_GPU,
+            image_uri=IMAGE_URI,
+            framework_version="1.13.1",
+            py_version="py3",
+        )
+
+        # Check that the Nova recipe was correctly identified
+        assert pytorch.is_nova_recipe is True
+
+        # Verify that model_type hyperparameter was set correctly
+        assert pytorch._hyperparameters.get("model_type") == "amazon.nova.llama-2-7b"