fix(rlaif): accept preset reward_prompt template names (#6192)

* fix: update timeout of RLVR sequence_length test

* fix: add xdist_group to Nova Bedrock deployment tests

* Address rlaif trainer fix

* Address rlaif trainer fix

---------

Co-authored-by: Roja Reddy Sareddy <rsareddy@amazon.com>

Author: rsareddy0329

sagemaker-train/src/sagemaker/train/rlaif_trainer.py

@@ -49,7 +49,7 @@ class RLAIFTrainer(BaseTrainer):
             training_type=TrainingType.LORA,
             model_package_group="my-model-group",
             reward_model_id="reward-model-id",
-            reward_prompt="Rate the helpfulness of this response on a scale of 1-10",
+            reward_prompt="summarize",
             training_dataset="s3://bucket/rlaif_data.jsonl"
         )
 
@@ -60,7 +60,7 @@ class RLAIFTrainer(BaseTrainer):
             model="meta-llama/Llama-2-7b-hf",
             model_package_group="my-rlaif-models",
             reward_model_id="reward-model-id",
-            reward_prompt="Rate the helpfulness of this response on a scale of 1-10"
+            reward_prompt="summarize"
         )
         
         # Create training job (non-blocking)
@@ -396,8 +396,14 @@ def _process_hyperparameters(self):
         # Process reward_prompt parameter
         if hasattr(self, 'reward_prompt') and self.reward_prompt:
             if isinstance(self.reward_prompt, str):
-                if self.reward_prompt.startswith("Builtin"):
-                    # Handle builtin reward prompts
+                # Resolution order:
+                #   1. Preset template name -> resolved locally against the recipe's
+                #      judge_prompt_template enum (no API call). Accepts "Builtin.Summarize",
+                #      "summarize", or "summarize.jinja".
+                #   2. Evaluator ARN -> validated/assigned as-is.
+                #   3. Otherwise -> HubContent name lookup (custom registered prompt),
+                #      which raises a clear error if not found.
+                if self._is_preset_reward_prompt(self.reward_prompt):
                     self._update_judge_prompt_template_direct(self.reward_prompt)
                 else:
                     # Handle evaluator ARN or hub content name
@@ -411,9 +417,44 @@ def _process_hyperparameters(self):
                 evaluator_arn = _extract_evaluator_arn(self.reward_prompt, "reward_prompt")
                 self._evaluator_arn = evaluator_arn
 
+    @staticmethod
+    def _normalize_template_name(value: str) -> str:
+        """Normalize a preset name or enum path to a comparable key.
+
+        Handles an optional "Builtin." prefix, any path prefix, and an optional
+        ".jinja" suffix, case-insensitively. For example "Builtin.Summarize",
+        "summarize", "summarize.jinja", and "/opt/ml/code/verl/summarize.jinja"
+        all normalize to "summarize".
+        """
+        name = (value or "").strip()
+        if name.lower().startswith("builtin."):
+            name = name.split(".", 1)[1]
+        name = name.split("/")[-1]  # basename
+        if name.lower().endswith(".jinja"):
+            name = name[: -len(".jinja")]
+        return name.lower()
+
+    def _get_judge_prompt_template_enum(self):
+        """Return the recipe's judge_prompt_template enum values (already in memory)."""
+        if not self.hyperparameters or not getattr(self.hyperparameters, "_specs", None):
+            return []
+        judge_prompt_spec = self.hyperparameters._specs.get("judge_prompt_template", {})
+        return judge_prompt_spec.get("enum", []) or []
+
+    def _is_preset_reward_prompt(self, reward_prompt: str) -> bool:
+        """True if reward_prompt matches a recipe preset template (local, no API call).
+
+        An explicit "Builtin." prefix always routes to preset resolution so the
+        user gets a clear "not available" error instead of a HubContent lookup.
+        """
+        if reward_prompt.startswith("Builtin"):
+            return True
+        enum_keys = {self._normalize_template_name(e) for e in self._get_judge_prompt_template_enum()}
+        return self._normalize_template_name(reward_prompt) in enum_keys
+
     def _process_non_builtin_reward_prompt(self):
-        """Process non-builtin reward prompt (ARN or hub content name)."""
-        # Remove judge_prompt_template for non-builtin prompts
+        """Process non-preset reward prompt (ARN or hub content name)."""
+        # Remove judge_prompt_template for non-preset prompts
         if hasattr(self.hyperparameters, 'judge_prompt_template'):
             delattr(self.hyperparameters, 'judge_prompt_template')
             self.hyperparameters._specs.pop('judge_prompt_template', None)
@@ -442,37 +483,42 @@ def _process_non_builtin_reward_prompt(self):
 
 
     def _update_judge_prompt_template_direct(self, reward_prompt):
-        """Update judge_prompt_template based on Builtin reward function."""
+        """Resolve a preset reward prompt name to the recipe's judge_prompt_template value.
+
+        Accepts "Builtin.Summarize", "summarize", or "summarize.jinja" and matches
+        it against the recipe's judge_prompt_template enum (normalized by basename,
+        with an optional ".jinja" suffix). No API call is made.
+        """
         # Get available templates from hyperparameters specs
-        judge_prompt_spec = self.hyperparameters._specs.get('judge_prompt_template', {})
-        available_templates = judge_prompt_spec.get('enum', [])
-        
+        available_templates = self._get_judge_prompt_template_enum()
+
         if not available_templates:
             # If no enum found, use the current value as the only available option
             current_value = getattr(self.hyperparameters, 'judge_prompt_template', None)
             if current_value:
                 available_templates = [current_value]
             else:
                 return
-        
-        # Extract template name after "Builtin." and convert to lowercase
-        template_name = reward_prompt.split(".", 1)[1].lower()
-        
-        # Find matching template by extracting filename without extension
+
+        # Normalize the requested name (strips optional "Builtin." prefix and ".jinja")
+        template_name = self._normalize_template_name(reward_prompt)
+
+        # Find matching template by normalized basename
         matching_template = None
         for template in available_templates:
-            template_filename = template.split("/")[-1].replace(".jinja", "").lower()
-            if template_filename == template_name:
+            if self._normalize_template_name(template) == template_name:
                 matching_template = template
                 break
 
         if matching_template:
             self.hyperparameters.judge_prompt_template = matching_template
         else:
-            available_options = [f"Builtin.{t.split('/')[-1].replace('.jinja', '')}" for t in available_templates]
+            available_options = [self._normalize_template_name(t) for t in available_templates]
             raise ValueError(
-                f"Selected reward function option '{reward_prompt}' is not available. "
-                f"Choose one from the available options: {available_options}. "
-                f"Example: reward_prompt='Builtin.summarize'"
+                f"Selected reward prompt '{reward_prompt}' is not an available preset. "
+                f"Choose one from the available options: {available_options} "
+                f"(pass the name directly, e.g. reward_prompt='{available_options[0]}', "
+                f"or with the 'Builtin.' prefix). "
+                f"Alternatively pass an evaluator ARN or a registered HubContent prompt name."
             )
 

sagemaker-train/tests/unit/train/test_rlaif_trainer.py

@@ -394,20 +394,23 @@ def test_process_hyperparameters_early_return_on_none(self):
         # No exception should be raised
 
     def test_update_judge_prompt_template_direct_with_matching_template(self):
-        """Test _update_judge_prompt_template_direct with matching template."""
-        mock_hyperparams = Mock()
-        mock_hyperparams._specs = {
-            'judge_prompt_template': {
-                'enum': ['templates/summarize.jinja', 'templates/helpfulness.jinja']
+        """Test _update_judge_prompt_template_direct resolves Builtin, plain, and .jinja names."""
+        for reward_prompt in ("Builtin.summarize", "summarize", "summarize.jinja", "Builtin.Summarize"):
+            mock_hyperparams = Mock()
+            mock_hyperparams._specs = {
+                'judge_prompt_template': {
+                    'enum': ['templates/summarize.jinja', 'templates/helpfulness.jinja']
+                }
             }
-        }
-        
-        trainer = RLAIFTrainer.__new__(RLAIFTrainer)
-        trainer.hyperparameters = mock_hyperparams
-        
-        trainer._update_judge_prompt_template_direct("Builtin.summarize")
-        
-        assert mock_hyperparams.judge_prompt_template == 'templates/summarize.jinja'
+
+            trainer = RLAIFTrainer.__new__(RLAIFTrainer)
+            trainer.hyperparameters = mock_hyperparams
+
+            trainer._update_judge_prompt_template_direct(reward_prompt)
+
+            assert mock_hyperparams.judge_prompt_template == 'templates/summarize.jinja', (
+                f"failed for input {reward_prompt!r}"
+            )
 
     def test_update_judge_prompt_template_direct_with_no_enum(self):
         """Test _update_judge_prompt_template_direct when no enum is available."""
@@ -434,7 +437,7 @@ def test_update_judge_prompt_template_direct_no_matching_template(self):
         trainer = RLAIFTrainer.__new__(RLAIFTrainer)
         trainer.hyperparameters = mock_hyperparams
 
-        with pytest.raises(ValueError, match="Selected reward function option 'Builtin.nonexistent' is not available"):
+        with pytest.raises(ValueError, match="Selected reward prompt 'Builtin.nonexistent' is not an available preset"):
             trainer._update_judge_prompt_template_direct("Builtin.nonexistent")
 
     def test_update_judge_prompt_template_direct_early_return(self):
@@ -449,6 +452,58 @@ def test_update_judge_prompt_template_direct_early_return(self):
         # Should return early without error
         trainer._update_judge_prompt_template_direct("Builtin.anything")
 
+    def test_normalize_template_name(self):
+        """Normalization strips Builtin. prefix, path, and optional .jinja; lowercases."""
+        cases = {
+            "summarize": "summarize",
+            "summarize.jinja": "summarize",
+            "Builtin.Summarize": "summarize",
+            "Builtin.summarize.jinja": "summarize",
+            "/opt/ml/code/verl/summarize.jinja": "summarize",
+            "bedrock/RLAIF/PandaLM/prompts/grader.jinja": "grader",
+            "  Summarize  ": "summarize",
+        }
+        for raw, expected in cases.items():
+            assert RLAIFTrainer._normalize_template_name(raw) == expected, f"failed for {raw!r}"
+
+    def test_is_preset_reward_prompt_matches_enum_without_prefix(self):
+        """Plain names that match the enum are presets (no API call)."""
+        mock_hyperparams = Mock()
+        mock_hyperparams._specs = {
+            'judge_prompt_template': {
+                'enum': ['/opt/ml/code/verl/summarize.jinja', 'bedrock/RLAIF/PandaLM/prompts/grader.jinja']
+            }
+        }
+        trainer = RLAIFTrainer.__new__(RLAIFTrainer)
+        trainer.hyperparameters = mock_hyperparams
+
+        assert trainer._is_preset_reward_prompt("summarize") is True
+        assert trainer._is_preset_reward_prompt("summarize.jinja") is True
+        assert trainer._is_preset_reward_prompt("Builtin.Summarize") is True
+        assert trainer._is_preset_reward_prompt("grader") is True
+        # Builtin.* always routes to preset resolution (for a clear error later)
+        assert trainer._is_preset_reward_prompt("Builtin.anything") is True
+        # A raw prompt / unknown name is not a preset -> falls through to ARN/Hub
+        assert trainer._is_preset_reward_prompt("Rate the helpfulness 1-10") is False
+        assert trainer._is_preset_reward_prompt("arn:aws:sagemaker:us-east-1:1:evaluator/x") is False
+
+    def test_process_hyperparameters_routes_plain_preset_to_template(self):
+        """A plain preset name sets judge_prompt_template and never calls Hub."""
+        mock_hyperparams = Mock()
+        mock_hyperparams._specs = {
+            'judge_prompt_template': {'enum': ['/opt/ml/code/verl/summarize.jinja']}
+        }
+        trainer = RLAIFTrainer.__new__(RLAIFTrainer)
+        trainer.hyperparameters = mock_hyperparams
+        trainer.reward_prompt = "summarize"
+        trainer.reward_model_id = None
+
+        with patch('sagemaker.train.rlaif_trainer._get_hub_content_metadata') as mock_hub:
+            trainer._process_hyperparameters()
+
+        mock_hub.assert_not_called()
+        assert mock_hyperparams.judge_prompt_template == '/opt/ml/code/verl/summarize.jinja'
+
     def test_process_non_builtin_reward_prompt_removes_judge_template(self):
         """Test _process_non_builtin_reward_prompt removes judge_prompt_template."""
         mock_hyperparams = Mock()