diff --git a/sagemaker-train/src/sagemaker/train/rlaif_trainer.py b/sagemaker-train/src/sagemaker/train/rlaif_trainer.py index b18f9f5751..a3060cec0f 100644 --- a/sagemaker-train/src/sagemaker/train/rlaif_trainer.py +++ b/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,11 +483,15 @@ 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) @@ -454,25 +499,26 @@ def _update_judge_prompt_template_direct(self, reward_prompt): 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." ) diff --git a/sagemaker-train/tests/unit/train/test_rlaif_trainer.py b/sagemaker-train/tests/unit/train/test_rlaif_trainer.py index 8128a280ea..9851f87c6f 100644 --- a/sagemaker-train/tests/unit/train/test_rlaif_trainer.py +++ b/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()