fix(train): discover Training Sessions models by capability keyword

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