From 784f6bb059c89d911c855fa825265b304db679d0 Mon Sep 17 00:00:00 2001 From: Xichao Wang Date: Sun, 18 Oct 2026 10:53:35 +0000 Subject: [PATCH] fix(training-session): discover base models by Training Sessions capability list_base_models() and base_model resolution keep a Hub model when it carries @capability:training_session, or, until the Public Hub publishes that keyword, when it is in the client Region's static list of models served by Training Sessions. Replaces the MTRL recipe keyword filter, which listed only models that also have a managed MTRL recipe. --- .../train/training_session/_utils.py | 54 +++++++++++++++---- .../training_session/test_service_client.py | 20 +++---- .../unit/train/training_session/test_utils.py | 44 ++++++++++++--- 3 files changed, 90 insertions(+), 28 deletions(-) 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"])