feat: Extract reward_lambda_arn from Nova recipes to training job hyperparameters

Author: Malav Shastri

src/sagemaker/modules/train/sm_recipes/utils.py

@@ -312,6 +312,12 @@ def _get_args_from_nova_recipe(
         if lambda_arn:
             args["hyperparameters"]["eval_lambda_arn"] = lambda_arn
 
+    # Handle reward lambda configuration
+    run_config = recipe.get("run", {})
+    reward_lambda_arn = run_config.get("reward_lambda_arn", "")
+    if reward_lambda_arn:
+        args["hyperparameters"]["reward_lambda_arn"] = reward_lambda_arn
+
     _register_custom_resolvers()
 
     # Resolve Final Recipe

src/sagemaker/pytorch/estimator.py

@@ -1251,6 +1251,12 @@ def _setup_for_nova_recipe(
             if lambda_arn:
                 args["hyperparameters"]["eval_lambda_arn"] = lambda_arn
 
+        # Handle reward lambda configuration
+        run_config = recipe.get("run", {})
+        reward_lambda_arn = run_config.get("reward_lambda_arn", "")
+        if reward_lambda_arn:
+            args["hyperparameters"]["reward_lambda_arn"] = reward_lambda_arn
+
         # Resolve and save the final recipe
         self._recipe_resolve_and_save(recipe, recipe_name, args["source_dir"])
 

tests/unit/sagemaker/modules/train/sm_recipes/test_utils.py

@@ -478,3 +478,57 @@ def test_get_args_from_nova_recipe_with_evaluation(test_case):
         recipe=recipe, compute=test_case["compute"], role=test_case["role"]
     )
     assert args == test_case["expected_args"]
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        {
+            "recipe": {
+                "run": {
+                    "model_type": "amazon.nova",
+                    "model_name_or_path": "dummy-test",
+                    "reward_lambda_arn": "arn:aws:lambda:us-east-1:123456789012:function:MyRewardLambdaFunction",
+                },
+            },
+            "compute": Compute(instance_type="ml.m5.xlarge", instance_count=2),
+            "role": "arn:aws:iam::123456789012:role/SageMakerRole",
+            "expected_args": {
+                "compute": Compute(instance_type="ml.m5.xlarge", instance_count=2),
+                "hyperparameters": {
+                    "base_model": "dummy-test",
+                    "reward_lambda_arn": "arn:aws:lambda:us-east-1:123456789012:function:MyRewardLambdaFunction",
+                },
+                "training_image": None,
+                "source_code": None,
+                "distributed": None,
+            },
+        },
+        {
+            "recipe": {
+                "run": {
+                    "model_type": "amazon.nova",
+                    "model_name_or_path": "dummy-test",
+                    # No reward_lambda_arn - should not be in hyperparameters
+                },
+            },
+            "compute": Compute(instance_type="ml.m5.xlarge", instance_count=2),
+            "role": "arn:aws:iam::123456789012:role/SageMakerRole",
+            "expected_args": {
+                "compute": Compute(instance_type="ml.m5.xlarge", instance_count=2),
+                "hyperparameters": {
+                    "base_model": "dummy-test",
+                },
+                "training_image": None,
+                "source_code": None,
+                "distributed": None,
+            },
+        },
+    ],
+)
+def test_get_args_from_nova_recipe_with_reward_lambda(test_case):
+    recipe = OmegaConf.create(test_case["recipe"])
+    args, _ = _get_args_from_nova_recipe(
+        recipe=recipe, compute=test_case["compute"], role=test_case["role"]
+    )
+    assert args == test_case["expected_args"]