feat(train): Add inherited list_supported_models to BaseTrainer

sagemaker-train/src/sagemaker/train/base_trainer.py

@@ -102,6 +102,37 @@ class BaseTrainer(ABC):
     training_image: Optional[str] = None
     latest_training_job: Optional[TrainingJob] = None
 
+    @classmethod
+    @_telemetry_emitter(
+        feature=Feature.MODEL_CUSTOMIZATION,
+        func_name="BaseTrainer.list_supported_models",
+    )
+    def list_supported_models(cls, session=None) -> List[str]:
+        """Return the models that support this trainer's fine-tuning technique.
+
+        Queries SageMakerPublicHub for all models whose ``RecipeCollection``
+        contains a FineTuning recipe for this trainer's customization technique
+        (``cls._customization_technique``, e.g. ``"SFT"``, ``"DPO"``,
+        ``"RLVR"``, ``"RLAIF"``, ``"CPT"``).
+
+        Args:
+            session: Optional boto3 session.
+
+        Returns:
+            Sorted list of hub content model names supporting the technique.
+        """
+        from sagemaker.train.common_utils.recipe_utils import _list_hub_models_by_recipe
+
+        technique = getattr(cls, "_customization_technique", None)
+        if not technique:
+            raise NotImplementedError(
+                f"{cls.__name__} does not define a customization technique and "
+                "cannot list supported models."
+            )
+        return _list_hub_models_by_recipe(
+            recipe_type="FineTuning", technique=technique, session=session
+        )
+
     def __init__(
         self,
         sagemaker_session: Optional[Session] = None,

sagemaker-train/src/sagemaker/train/common_utils/recipe_utils.py

@@ -506,20 +506,25 @@ def resolve_recipe(
 
 
 def _build_recipe_keyword(recipe_type: str, technique: str) -> str:
-    """Build the ``@recipe:`` search keyword for a recipe type and technique.
+    """Build the base ``@recipe:`` search keyword for a recipe type and technique.
 
-    The hub tags recipes as ``@recipe:{type}_{technique}_{strategy}`` (all
-    lowercase).  We match on the ``@recipe:{type}_{technique}_`` prefix so
-    the strategy component (e.g. ``lora``) is ignored.
+    The hub tags recipes as either ``@recipe:{type}_{technique}_{strategy}``
+    (e.g. ``@recipe:finetuning_sft_lora``) or, for techniques with no strategy
+    component, the bare ``@recipe:{type}_{technique}`` (e.g.
+    ``@recipe:finetuning_cpt``) — all lowercase. This returns the base
+    ``@recipe:{type}_{technique}`` form (no trailing underscore); callers match
+    it exactly OR as a ``{base}_`` prefix so the optional strategy component is
+    ignored without matching an unrelated technique that merely shares a prefix
+    (e.g. base ``..._rl`` must not match ``..._rlvr_...``).
 
     Args:
         recipe_type: ``"FineTuning"`` or ``"Evaluation"``.
         technique: Technique value, e.g. ``"MTRL"`` or ``"MTRLEvaluation"``.
 
     Returns:
-        Lowercase keyword prefix string, e.g. ``"@recipe:finetuning_mtrl_"``.
+        Lowercase base keyword string, e.g. ``"@recipe:finetuning_mtrl"``.
     """
-    return f"@recipe:{recipe_type}_{technique}_".lower()
+    return f"@recipe:{recipe_type}_{technique}".lower()
 
 
 def _list_hub_models_by_recipe(
@@ -552,7 +557,7 @@ def _list_hub_models_by_recipe(
             f"recipe_type must be 'FineTuning' or 'Evaluation', got: {recipe_type!r}"
         )
 
-    keyword_prefix = _build_recipe_keyword(recipe_type, technique)
+    keyword_base = _build_recipe_keyword(recipe_type, technique)
 
     region = (getattr(session, "region_name", None) or 
               getattr(getattr(session, "boto_session", None), "region_name", None) or
@@ -577,7 +582,14 @@ def _list_hub_models_by_recipe(
             if not content_name:
                 continue
             keywords = summary.get("HubContentSearchKeywords", [])
-            if any(kw.lower().startswith(keyword_prefix) for kw in keywords):
+            # Match the bare base keyword (techniques with no strategy component,
+            # e.g. "@recipe:finetuning_cpt") OR the "{base}_{strategy}" form
+            # (e.g. "@recipe:finetuning_sft_lora"). The "{base}_" guard prevents
+            # matching an unrelated technique that merely shares a prefix.
+            if any(
+                (kwl := kw.lower()) == keyword_base or kwl.startswith(keyword_base + "_")
+                for kw in keywords
+            ):
                 matched_models.append(content_name)
 
         next_token = response.get("NextToken")

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

@@ -129,6 +129,8 @@ class RLAIFTrainer(BaseTrainer):
             and 'job_name_prefix'. If not specified, no notifications are sent.
     """
 
+    _customization_technique = CustomizationTechnique.RLAIF.value
+
     def __init__(
         self,
         model: Union[str, ModelPackage],

sagemaker-train/tests/integ/train/test_trainer_list_supported_models_integration.py

@@ -0,0 +1,131 @@
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Integration tests for ``BaseTrainer.list_supported_models`` across the
+fine-tuning trainers (SFT / RLVR / RLAIF / DPO / CPT).
+
+These tests run against real SageMaker services in prod us-west-2 and query the
+active SageMaker hub (the integ harness pins a private ``SAGEMAKER_HUB_NAME`` in
+``conftest.py``; falls back to ``SageMakerPublicHub``).
+
+They verify the contract that unit tests (which mock the hub) cannot: that each
+trainer's ``_customization_technique`` string matches how the live hub tags its
+FineTuning recipes. Rather than assert a hard non-empty count -- which is brittle
+because the pinned hub may not carry every technique -- the test independently
+scans the hub once and asserts ``list_supported_models`` returns exactly the set
+of models tagged for that technique (whether that is zero or many). A helper
+regression that (for example) required a ``_{strategy}`` suffix -- and so dropped
+suffix-less techniques like CPT (``@recipe:finetuning_cpt``) -- would surface
+here as a mismatch against the oracle.
+"""
+from __future__ import annotations
+
+import collections
+import os
+
+import boto3
+import pytest
+from sagemaker.core.helper.session_helper import Session
+from sagemaker.train.sft_trainer import SFTTrainer
+from sagemaker.train.rlvr_trainer import RLVRTrainer
+from sagemaker.train.rlaif_trainer import RLAIFTrainer
+from sagemaker.train.dpo_trainer import DPOTrainer
+from sagemaker.train.cpt_trainer import CPTTrainer
+
+_REGION = "us-west-2"
+_FINETUNING_PREFIX = "@recipe:finetuning_"
+
+# (trainer class, expected customization technique string)
+_TRAINER_CASES = [
+    pytest.param(SFTTrainer, "SFT", id="SFT"),
+    pytest.param(RLVRTrainer, "RLVR", id="RLVR"),
+    pytest.param(RLAIFTrainer, "RLAIF", id="RLAIF"),
+    pytest.param(DPOTrainer, "DPO", id="DPO"),
+    pytest.param(CPTTrainer, "CPT", id="CPT"),
+]
+
+
+@pytest.fixture(scope="module")
+def sagemaker_session():
+    boto_session = boto3.Session(region_name=_REGION)
+    yield Session(boto_session=boto_session)
+
+
+@pytest.fixture(scope="module")
+def hub_finetuning_models(sagemaker_session):
+    """Independent oracle: scan the active hub once and map technique token ->
+    sorted list of model names tagged with a matching FineTuning recipe.
+
+    Built independently of the SDK helper (groups by the token immediately after
+    ``@recipe:finetuning_``), so it can catch a regression in that helper rather
+    than merely re-deriving it.
+    """
+    client = sagemaker_session.boto_session.client("sagemaker", region_name=_REGION)
+    hub_name = os.environ.get("SAGEMAKER_HUB_NAME", "SageMakerPublicHub")
+    mapping: dict[str, set] = collections.defaultdict(set)
+    next_token = None
+    while True:
+        kwargs = {"HubName": hub_name, "HubContentType": "Model"}
+        if next_token:
+            kwargs["NextToken"] = next_token
+        response = client.list_hub_contents(**kwargs)
+        for summary in response.get("HubContentSummaries", []):
+            name = summary.get("HubContentName")
+            if not name:
+                continue
+            for keyword in summary.get("HubContentSearchKeywords", []):
+                kwl = keyword.lower()
+                if kwl.startswith(_FINETUNING_PREFIX):
+                    token = kwl[len(_FINETUNING_PREFIX):].split("_")[0]
+                    mapping[token].add(name)
+        next_token = response.get("NextToken")
+        if not next_token:
+            break
+    return {tech: sorted(names) for tech, names in mapping.items()}
+
+
+class TestTrainerListSupportedModels:
+    """List supported models per fine-tuning technique (requires API access)."""
+
+    @pytest.mark.parametrize("trainer_cls,expected_technique", _TRAINER_CASES)
+    def test_list_supported_models(
+        self, trainer_cls, expected_technique, sagemaker_session, hub_finetuning_models
+    ):
+        """Each trainer resolves its technique and returns exactly the hub models
+        tagged for it."""
+        # Sanity: the class attribute the inherited method keys off is set.
+        assert trainer_cls._customization_technique == expected_technique
+
+        result = trainer_cls.list_supported_models(
+            session=sagemaker_session.boto_session
+        )
+
+        # Structural contract: a sorted, de-duplicated list of non-empty strings.
+        assert isinstance(result, list)
+        assert all(isinstance(name, str) and name for name in result)
+        assert result == sorted(result)
+        assert len(set(result)) == len(result)
+
+        # Correctness contract: exactly the models the active hub tags for this
+        # technique (may legitimately be empty if the pinned hub carries none).
+        expected = hub_finetuning_models.get(expected_technique.lower(), [])
+        assert result == expected
+
+    def test_public_hub_has_models_for_core_techniques(self, hub_finetuning_models):
+        """Guard against a silent all-empty hub / broken scan: only meaningful
+        against the public hub, where these techniques are known to be tagged.
+        Skipped when a private test hub is pinned."""
+        if os.environ.get("SAGEMAKER_HUB_NAME", "SageMakerPublicHub") != "SageMakerPublicHub":
+            pytest.skip("private hub pinned; model population is environment-specific")
+        for technique in ("sft", "dpo", "rlvr", "rlaif", "cpt"):
+            assert hub_finetuning_models.get(technique), (
+                f"public hub returned no models for technique '{technique}'"
+            )

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

@@ -300,3 +300,33 @@ def test_model_source_passed_as_override_parameter(
         start_cmd = mock_subprocess.run.call_args_list[-1].args[0]
         overrides = json.loads(start_cmd[start_cmd.index("--override-parameters") + 1])
         assert overrides["recipes.run.model_name_or_path"] == "s3://bucket/checkpoint/step_10"
+
+
+class TestBaseTrainerListSupportedModels:
+    """The inherited ``list_supported_models`` classmethod on ``BaseTrainer``."""
+
+    def test_delegates_with_class_technique(self):
+        class _TechTrainer(BaseTrainer):
+            _customization_technique = "SFT"
+
+            def train(self, *args, **kwargs):  # pragma: no cover - abstract impl
+                return None
+
+        with patch(
+            "sagemaker.train.common_utils.recipe_utils._list_hub_models_by_recipe"
+        ) as mock_list:
+            mock_list.return_value = ["meta-llama/Llama-3"]
+            result = _TechTrainer.list_supported_models()
+
+        assert result == ["meta-llama/Llama-3"]
+        mock_list.assert_called_once_with(
+            recipe_type="FineTuning", technique="SFT", session=None
+        )
+
+    def test_raises_when_technique_missing(self):
+        class _NoTechTrainer(BaseTrainer):
+            def train(self, *args, **kwargs):  # pragma: no cover - abstract impl
+                return None
+
+        with pytest.raises(NotImplementedError, match="customization technique"):
+            _NoTechTrainer.list_supported_models()

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

@@ -738,3 +738,15 @@ def test_dry_run_returns_none_without_submitting(
         mock_create.assert_not_called()
         mock_role.assert_called_once()
         mock_validate_hp.assert_called_once()
+
+
+class TestDPOTrainerListSupportedModels:
+
+    @patch("sagemaker.train.common_utils.recipe_utils._list_hub_models_by_recipe")
+    def test_list_supported_models(self, mock_list):
+        mock_list.return_value = ["meta-llama/Llama-3"]
+        result = DPOTrainer.list_supported_models()
+        assert result == ["meta-llama/Llama-3"]
+        mock_list.assert_called_once_with(
+            recipe_type="FineTuning", technique="DPO", session=None
+        )

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

@@ -666,6 +666,51 @@ def test_invalid_recipe_type_raises(self):
         with pytest.raises(ValueError, match="recipe_type must be"):
             _list_hub_models_by_recipe(recipe_type="Invalid", technique="MTRL")
 
+    @patch("sagemaker.train.common_utils.recipe_utils.boto3.Session")
+    def test_finds_models_with_bare_keyword_no_strategy(self, mock_session_cls):
+        """Techniques whose recipes carry no strategy suffix are tagged with the
+        bare ``@recipe:finetuning_{technique}`` keyword (e.g. CPT). The matcher
+        must find these, not just ``{base}_{strategy}`` forms."""
+        mock_client = MagicMock()
+        mock_session_cls.return_value.client.return_value = mock_client
+
+        mock_client.list_hub_contents.return_value = {
+            "HubContentSummaries": [
+                {
+                    "HubContentName": "model-cpt-bare",
+                    "HubContentSearchKeywords": ["@recipe:finetuning_cpt"],
+                },
+                {
+                    "HubContentName": "model-cpt-suffixed",
+                    "HubContentSearchKeywords": ["@recipe:finetuning_cpt_full"],
+                },
+            ],
+        }
+
+        from sagemaker.train.common_utils.recipe_utils import _list_hub_models_by_recipe
+        result = _list_hub_models_by_recipe(recipe_type="FineTuning", technique="CPT")
+        assert result == ["model-cpt-bare", "model-cpt-suffixed"]
+
+    @patch("sagemaker.train.common_utils.recipe_utils.boto3.Session")
+    def test_does_not_match_technique_sharing_a_prefix(self, mock_session_cls):
+        """A shorter technique must not match a longer one that merely shares its
+        prefix (e.g. ``rl`` must not match ``rlvr``)."""
+        mock_client = MagicMock()
+        mock_session_cls.return_value.client.return_value = mock_client
+
+        mock_client.list_hub_contents.return_value = {
+            "HubContentSummaries": [
+                {
+                    "HubContentName": "model-rlvr",
+                    "HubContentSearchKeywords": ["@recipe:finetuning_rlvr_lora"],
+                },
+            ],
+        }
+
+        from sagemaker.train.common_utils.recipe_utils import _list_hub_models_by_recipe
+        result = _list_hub_models_by_recipe(recipe_type="FineTuning", technique="rl")
+        assert result == []
+
 
 class TestListAgentRuntimes:
     @patch("sagemaker.train.multi_turn_rl_trainer.boto3.Session")

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

@@ -805,3 +805,15 @@ def test_train_passes_sequence_length_to_serverless_config(
         mock_serverless_config.assert_called_once()
         call_kwargs = mock_serverless_config.call_args[1]
         assert call_kwargs["sequence_length"] == "64K"
+
+
+class TestRLAIFTrainerListSupportedModels:
+
+    @patch("sagemaker.train.common_utils.recipe_utils._list_hub_models_by_recipe")
+    def test_list_supported_models(self, mock_list):
+        mock_list.return_value = ["meta-llama/Llama-3"]
+        result = RLAIFTrainer.list_supported_models()
+        assert result == ["meta-llama/Llama-3"]
+        mock_list.assert_called_once_with(
+            recipe_type="FineTuning", technique="RLAIF", session=None
+        )

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

@@ -763,3 +763,15 @@ def test_dry_run_returns_none_without_submitting(
         mock_create.assert_not_called()
         mock_role.assert_called_once()
         mock_validate_hp.assert_called_once()
+
+
+class TestRLVRTrainerListSupportedModels:
+
+    @patch("sagemaker.train.common_utils.recipe_utils._list_hub_models_by_recipe")
+    def test_list_supported_models(self, mock_list):
+        mock_list.return_value = ["meta-llama/Llama-3"]
+        result = RLVRTrainer.list_supported_models()
+        assert result == ["meta-llama/Llama-3"]
+        mock_list.assert_called_once_with(
+            recipe_type="FineTuning", technique="RLVR", session=None
+        )

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

@@ -1556,3 +1556,25 @@ def test_dry_run_raises_on_role_validation_failure(
 
         with pytest.raises(ValueError, match="Missing permissions"):
             trainer.train(dry_run=True)
+
+
+class TestSFTTrainerListSupportedModels:
+
+    @patch("sagemaker.train.common_utils.recipe_utils._list_hub_models_by_recipe")
+    def test_list_supported_models(self, mock_list):
+        mock_list.return_value = ["meta-llama/Llama-3", "Qwen/Qwen3-32B"]
+        result = SFTTrainer.list_supported_models()
+        assert isinstance(result, list)
+        assert "Qwen/Qwen3-32B" in result
+        mock_list.assert_called_once_with(
+            recipe_type="FineTuning", technique="SFT", session=None
+        )
+
+    @patch("sagemaker.train.common_utils.recipe_utils._list_hub_models_by_recipe")
+    def test_list_supported_models_passes_session(self, mock_list):
+        mock_list.return_value = []
+        session = Mock()
+        SFTTrainer.list_supported_models(session=session)
+        mock_list.assert_called_once_with(
+            recipe_type="FineTuning", technique="SFT", session=session
+        )