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.

Author: Xichao Wang

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")

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",
             ],
         },

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"])