diff --git a/sagemaker-train/src/sagemaker/train/base_trainer.py b/sagemaker-train/src/sagemaker/train/base_trainer.py index 2f1ab48117..35801cb787 100644 --- a/sagemaker-train/src/sagemaker/train/base_trainer.py +++ b/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, diff --git a/sagemaker-train/src/sagemaker/train/common_utils/recipe_utils.py b/sagemaker-train/src/sagemaker/train/common_utils/recipe_utils.py index 4fd314068f..c4abf30ff0 100644 --- a/sagemaker-train/src/sagemaker/train/common_utils/recipe_utils.py +++ b/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") diff --git a/sagemaker-train/src/sagemaker/train/rlaif_trainer.py b/sagemaker-train/src/sagemaker/train/rlaif_trainer.py index b93aeb010d..b18f9f5751 100644 --- a/sagemaker-train/src/sagemaker/train/rlaif_trainer.py +++ b/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], diff --git a/sagemaker-train/tests/integ/train/test_trainer_list_supported_models_integration.py b/sagemaker-train/tests/integ/train/test_trainer_list_supported_models_integration.py new file mode 100644 index 0000000000..e3ee65aa5f --- /dev/null +++ b/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}'" + ) diff --git a/sagemaker-train/tests/unit/train/test_base_trainer_compute.py b/sagemaker-train/tests/unit/train/test_base_trainer_compute.py index 2922c326c1..af42e22802 100644 --- a/sagemaker-train/tests/unit/train/test_base_trainer_compute.py +++ b/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() diff --git a/sagemaker-train/tests/unit/train/test_dpo_trainer.py b/sagemaker-train/tests/unit/train/test_dpo_trainer.py index 5dfae85bfd..8ef48041a3 100644 --- a/sagemaker-train/tests/unit/train/test_dpo_trainer.py +++ b/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 + ) diff --git a/sagemaker-train/tests/unit/train/test_multi_turn_rl_trainer.py b/sagemaker-train/tests/unit/train/test_multi_turn_rl_trainer.py index 291c7cc79e..360deebcb2 100644 --- a/sagemaker-train/tests/unit/train/test_multi_turn_rl_trainer.py +++ b/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") diff --git a/sagemaker-train/tests/unit/train/test_rlaif_trainer.py b/sagemaker-train/tests/unit/train/test_rlaif_trainer.py index 9667c53202..8128a280ea 100644 --- a/sagemaker-train/tests/unit/train/test_rlaif_trainer.py +++ b/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 + ) diff --git a/sagemaker-train/tests/unit/train/test_rlvr_trainer.py b/sagemaker-train/tests/unit/train/test_rlvr_trainer.py index 75dd95c607..25839df202 100644 --- a/sagemaker-train/tests/unit/train/test_rlvr_trainer.py +++ b/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 + ) diff --git a/sagemaker-train/tests/unit/train/test_sft_trainer.py b/sagemaker-train/tests/unit/train/test_sft_trainer.py index e4e3803eba..1239bbd6d8 100644 --- a/sagemaker-train/tests/unit/train/test_sft_trainer.py +++ b/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 + )