diff --git a/sagemaker-train/src/sagemaker/train/training_session/_utils.py b/sagemaker-train/src/sagemaker/train/training_session/_utils.py index a69ad271f2..05e452463c 100644 --- a/sagemaker-train/src/sagemaker/train/training_session/_utils.py +++ b/sagemaker-train/src/sagemaker/train/training_session/_utils.py @@ -20,7 +20,7 @@ import logging import warnings import weakref -from typing import Any, Dict, List, Optional +from typing import Any, Dict, FrozenSet, List, Optional from sagemaker.train.training_session.types import ( Datum, @@ -34,12 +34,34 @@ _PUBLIC_HUB_NAME = "SageMakerPublicHub" -#: Hub content search-keyword prefix identifying a base model this client can -#: fine-tune. Each model's recipes are tagged -#: ``@recipe:{type}_{technique}_{strategy}`` (lowercase) on the Public Hub; -#: matching on the prefix selects models carrying a compatible fine-tuning -#: recipe regardless of the parameter-efficient strategy it uses. -_TRAINABLE_RECIPE_KEYWORD_PREFIX = "@recipe:finetuning_mtrl_" +#: Hub content search keyword identifying a base model this client can train. +#: The Public Hub publishes each model's ``Capabilities`` as +#: ``@capability:{capability}`` (lowercase) keywords; ``TRAINING_SESSION`` is +#: set per region, so the keyword is present only where the model is served. +_TRAINABLE_CAPABILITY_KEYWORD = "@capability:training_session" + +#: Hub content names served by Training Sessions, by Region, for the interval +#: before the Public Hub publishes the capability keyword above. Mirrors the +#: service's model configuration; a listed model that is absent from a +#: Region's Hub is simply not reported there. Remove when the keyword is live. +_TRAINABLE_MODELS_BY_REGION: Dict[str, FrozenSet[str]] = { + "us-west-2": frozenset( + { + "huggingface-vlm-gemma-4-31b-it", + "huggingface-vlm-qwen3-6-27b", + "huggingface-vlm-qwen3-8-27b", + "openai-reasoning-gpt-oss-20b", + } + ), + "us-east-1": frozenset({"openai-reasoning-gpt-oss-20b"}), + "ap-south-1": frozenset( + { + "huggingface-vlm-gemma-4-31b-it", + "huggingface-vlm-qwen3-6-27b", + "openai-reasoning-gpt-oss-20b", + } + ), +} #: Hub content search-keyword prefix carrying a model's published id, spelled #: ``@huggingface-id:{provider}/{model}``. Free on the list summary, but keyword @@ -206,17 +228,29 @@ def _public_model_id(hub_content_name: str, keywords: List[str], hub_client: Any return model_id +def _is_trainable(hub_content_name: str, keywords: List[str], region: Optional[str]) -> bool: + """Return True if the Hub entry is served by Training Sessions in this Region. + + A model qualifies when its keywords carry the training-session capability + or when it is in the Region's static list (:data:`_TRAINABLE_MODELS_BY_REGION`). + """ + if any(kw.lower() == _TRAINABLE_CAPABILITY_KEYWORD for kw in keywords): + return True + return hub_content_name in _TRAINABLE_MODELS_BY_REGION.get(region or "", frozenset()) + + def _scan_trainable_hub_models(hub_client: Any) -> Dict[str, str]: """Return ``{model id: Hub content name}`` for the region's fine-tunable models. - Queries the SageMaker Public Hub and keeps the models that advertise a - compatible fine-tuning recipe (via their ``HubContentSearchKeywords``), so + Queries the SageMaker Public Hub and keeps the models that advertise the + training-session capability (via their ``HubContentSearchKeywords``), so the set is discovered live rather than hardcoded and a newly enabled model appears with no SDK change. Args: hub_client: A boto3 ``sagemaker`` client (its region determines results). """ + region = getattr(getattr(hub_client, "meta", None), "region_name", None) trainable: List["tuple[str, List[str]]"] = [] next_token = None while True: @@ -231,7 +265,7 @@ def _scan_trainable_hub_models(hub_client: Any) -> Dict[str, str]: if _is_excluded_base_model(hub_content_name): continue keywords = item.get("HubContentSearchKeywords", []) or [] - if not any(kw.lower().startswith(_TRAINABLE_RECIPE_KEYWORD_PREFIX) for kw in keywords): + if not _is_trainable(hub_content_name, keywords, region): continue trainable.append((hub_content_name, keywords)) next_token = resp.get("NextToken") diff --git a/sagemaker-train/tests/unit/train/training_session/test_service_client.py b/sagemaker-train/tests/unit/train/training_session/test_service_client.py index 7abd6f5aa9..df7e52bf29 100644 --- a/sagemaker-train/tests/unit/train/training_session/test_service_client.py +++ b/sagemaker-train/tests/unit/train/training_session/test_service_client.py @@ -530,8 +530,8 @@ def test_list_keeps_a_status_this_version_does_not_know(patch_clients): class FakeHubClient: """Stubs DescribeHubContent + ListHubContents for base-model resolution.""" - #: Recipe keyword marking a model as fine-tunable (what the SDK filters on). - _TRAINABLE_KEYWORD = "@recipe:finetuning_mtrl_lora" + #: Capability keyword marking a model as trainable (what the SDK filters on). + _TRAINABLE_KEYWORD = "@capability:training_session" def __init__( self, @@ -547,7 +547,7 @@ def __init__( else ["hubcontent-vlm-model-b", "hubcontent-text-model-a"] ) # Search keywords per content name. Default: every listed model carries - # the fine-tuning recipe keyword, i.e. is trainable. + # the training-session capability keyword, i.e. is trainable. self._keywords_by_name = ( keywords_by_name if keywords_by_name is not None @@ -1092,8 +1092,8 @@ def test_create_rejects_unresolvable_base_model(patch_clients): svc.create_lora_training_client("not-a-real-model", **kwargs) -def test_list_base_models_filters_by_recipe_keyword(patch_clients): - # Only models carrying a fine-tuning recipe keyword are trainable; a model +def test_list_base_models_filters_by_capability_keyword(patch_clients): + # Only models carrying the training-session capability keyword are trainable; a model # present in the hub without one must be excluded. Each trainable model is # reported by the model id its Hub entry advertises. install, _ = patch_clients @@ -1108,11 +1108,11 @@ def test_list_base_models_filters_by_recipe_keyword(patch_clients): keywords_by_name={ "hubcontent-vlm-model-b": [ "@recipe:finetuning_sft_lora", - "@recipe:finetuning_mtrl_lora", + "@capability:training_session", "@huggingface-id:provider-b/model-b", ], "hubcontent-vlm-model-c": [ - "@recipe:finetuning_mtrl_lora", + "@capability:training_session", "@huggingface-id:provider-c/model-c", ], "some-eval-only-model": ["@recipe:evaluation_mtrlevaluation"], @@ -1126,7 +1126,7 @@ def test_list_base_models_filters_by_recipe_keyword(patch_clients): def test_list_base_models_excludes_unsupported_names(patch_clients): - # The recipe keyword alone does not make a model usable through this client. + # The capability keyword alone does not make a model usable through this client. # A name the client does not support end to end is excluded even when the # Hub tags it as fine-tunable. install, _ = patch_clients @@ -1140,11 +1140,11 @@ def test_list_base_models_excludes_unsupported_names(patch_clients): ], keywords_by_name={ "hubcontent-vlm-model-b": [ - "@recipe:finetuning_mtrl_lora", + "@capability:training_session", "@huggingface-id:provider-b/model-b", ], "hubcontent-nova-model-x": [ - "@recipe:finetuning_mtrl_lora", + "@capability:training_session", "@huggingface-id:provider-x/model-x", ], }, diff --git a/sagemaker-train/tests/unit/train/training_session/test_utils.py b/sagemaker-train/tests/unit/train/training_session/test_utils.py index bd4871aab2..dc18b3af71 100644 --- a/sagemaker-train/tests/unit/train/training_session/test_utils.py +++ b/sagemaker-train/tests/unit/train/training_session/test_utils.py @@ -89,7 +89,7 @@ class _FakeCatalogHub: Args: models: ``{hub_content_name: (arn, hf_id_or_None)}``. advertise_hf_id: Put ``@huggingface-id:`` on the list summary, as the - real Hub does. False leaves only the recipe keyword, so the + real Hub does. False leaves only the capability keyword, so the Describe fallback is what has to supply the id. region: Reported as ``meta.region_name``, as a real client does. The catalogue is cached per client rather than per Region, so this @@ -98,7 +98,7 @@ class _FakeCatalogHub: that are not fine-tunable at all. """ - _RECIPE_KW = "@recipe:finetuning_mtrl_lora" + _CAPABILITY_KW = "@capability:training_session" def __init__(self, models, *, advertise_hf_id=False, region=None, keywords_by_name=None): self._models = models @@ -112,7 +112,7 @@ def __init__(self, models, *, advertise_hf_id=False, region=None, keywords_by_na def _keywords(self, name): if name in self._keywords_by_name: return self._keywords_by_name[name] - keywords = [self._RECIPE_KW] + keywords = [self._CAPABILITY_KW] hf_id = self._models[name][1] if self._advertise_hf_id and hf_id: keywords.append(f"@huggingface-id:{hf_id.lower()}") @@ -157,7 +157,7 @@ def list_hub_contents(self, **kwargs): { "HubContentName": "hubcontent-text-model-a", "HubContentSearchKeywords": [ - "@recipe:finetuning_mtrl_lora", + "@capability:training_session", "@huggingface-id:provider-a/model-a", ], } @@ -246,6 +246,34 @@ def test_excludes_models_without_a_finetuning_recipe(self): ) assert list_base_models_in_region(hub) == ["provider-a/model-a"] + def test_static_region_list_surfaces_a_model_without_the_keyword(self): + # Until the Hub publishes the capability, a Region's static list stands in. + hub = _FakeCatalogHub( + {"huggingface-vlm-qwen3-8-27b": (_MODEL_A_ARN, "Qwen/Qwen3.8-27B")}, + advertise_hf_id=True, + region="us-west-2", + keywords_by_name={"huggingface-vlm-qwen3-8-27b": ["@huggingface-id:qwen/qwen3.8-27b"]}, + ) + assert list_base_models_in_region(hub) == ["Qwen/Qwen3.8-27B"] + assert resolve_base_model("Qwen/Qwen3.8-27B", hub) == (_MODEL_A_ARN, "Qwen/Qwen3.8-27B") + + def test_static_region_list_is_per_region(self): + hub = _FakeCatalogHub( + {"huggingface-vlm-qwen3-6-27b": (_MODEL_A_ARN, "Qwen/Qwen3.6-27B")}, + advertise_hf_id=True, + region="us-east-1", + keywords_by_name={"huggingface-vlm-qwen3-6-27b": ["@huggingface-id:qwen/qwen3.6-27b"]}, + ) + assert list_base_models_in_region(hub) == [] + + def test_static_region_list_needs_a_region(self): + hub = _FakeCatalogHub( + {"huggingface-vlm-qwen3-8-27b": (_MODEL_A_ARN, "Qwen/Qwen3.8-27B")}, + advertise_hf_id=True, + keywords_by_name={"huggingface-vlm-qwen3-8-27b": ["@huggingface-id:qwen/qwen3.8-27b"]}, + ) + assert list_base_models_in_region(hub) == [] + def test_paginates(self): hub = _FakeCatalogHub(self._MODELS, advertise_hf_id=True) pages = [ @@ -254,7 +282,7 @@ def test_paginates(self): { "HubContentName": "hubcontent-text-model-a", "HubContentSearchKeywords": [ - "@recipe:finetuning_mtrl_lora", + "@capability:training_session", "@huggingface-id:provider-a/model-a", ], } @@ -266,7 +294,7 @@ def test_paginates(self): { "HubContentName": "hubcontent-vlm-model-b", "HubContentSearchKeywords": [ - "@recipe:finetuning_mtrl_lora", + "@capability:training_session", "@huggingface-id:provider-b/model-b", ], } @@ -303,7 +331,7 @@ def test_ignores_a_keyword_that_is_not_a_model_id(self): # in the very error that rejects it; the content name is listed instead. hub = _FakeCatalogHub({"hubcontent-single": (_MODEL_A_ARN, None)}) hub._keywords_by_name = { - "hubcontent-single": ["@recipe:finetuning_mtrl_lora", "@huggingface-id:singleseg"] + "hubcontent-single": ["@capability:training_session", "@huggingface-id:singleseg"] } listed = list_base_models_in_region(hub) @@ -410,7 +438,7 @@ class _KeywordOnlyHub(_FakeCatalogHub): """Advertises the canonical id; its Hub document carries none.""" def _keywords(self, name): - return ["@recipe:finetuning_mtrl_lora", "@huggingface-id:Provider-A/Model-A"] + return ["@capability:training_session", "@huggingface-id:Provider-A/Model-A"] def describe_hub_content(self, **kwargs): self.describe_calls.append(kwargs["HubContentName"])