fix: propagate user-supplied tags to created resources

sagemaker-serve/src/sagemaker/serve/model_builder.py

@@ -114,6 +114,7 @@
 from sagemaker.core.enums import EndpointType
 from sagemaker.core.common_utils import (
     Tags,
+    TagsDict,
     ModelApprovalStatusEnum,
     _resolve_routing_config,
     format_tags,
@@ -168,6 +169,19 @@
 SAGEMAKER_OUTPUT_LOCATION = "sagemaker_s3_output"
 
 
+def _tags_as_list(tags: Optional[Tags]) -> Optional[List[TagsDict]]:
+    """Expand a single ``{key: value}`` tag mapping into the list form.
+
+    ``deploy()`` accepts tags either as a list of tag dicts or as one ``{key: value}``
+    mapping (see the ``Tags`` alias), but the resource ``create()`` calls are typed
+    ``List[Tag]`` and reject a bare mapping. A list is returned unchanged, for pydantic
+    to validate and coerce.
+    """
+    if isinstance(tags, dict):
+        return [{"key": k, "value": v} for k, v in tags.items()]
+    return tags
+
+
 @dataclass
 class ModelBuilder(_InferenceRecommenderMixin, _ModelBuilderServers, _ModelBuilderUtils):
     """Unified interface for building and deploying machine learning models.
@@ -5678,8 +5692,12 @@ def _deploy_model_customization(
                 endpoint_name=endpoint_name,
                 initial_instance_count=initial_instance_count,
                 wait=kwargs.get("wait", True),
+                tags=kwargs.get("tags"),
             )
 
+        # Passed to whichever endpoint-creation path runs below.
+        endpoint_tags = _tags_as_list(kwargs.get("tags"))
+
         # Fetch model package
         model_package = self._fetch_model_package()
 
@@ -5700,7 +5718,8 @@ def _deploy_model_customization(
                 ],
             )
             endpoint = Endpoint.create(
-                endpoint_name=endpoint_name, endpoint_config_name=endpoint_name
+                endpoint_name=endpoint_name, endpoint_config_name=endpoint_name,
+                tags=endpoint_tags,
             )
             if kwargs.get("wait", True):
                 endpoint.wait_for_status("InService")
@@ -5723,7 +5742,8 @@ def _deploy_model_customization(
             )
             logger.info("Endpoint core call starting")
             endpoint = Endpoint.create(
-                endpoint_name=endpoint_name, endpoint_config_name=endpoint_name
+                endpoint_name=endpoint_name, endpoint_config_name=endpoint_name,
+                tags=endpoint_tags,
             )
             endpoint.wait_for_status("InService")
         else:
@@ -5953,6 +5973,7 @@ def _deploy_nova_model(
         endpoint_name: str,
         initial_instance_count: int = 1,
         wait: bool = True,
+        tags: Optional[Tags] = None,
     ) -> Endpoint:
         """Deploy a Nova model directly to an endpoint without inference components.
 
@@ -5981,16 +6002,21 @@ def _deploy_nova_model(
             ],
         )
 
-        tags = [
+        endpoint_tags = [
             {"key": "sagemaker-sdk:jumpstart-model-id", "value": base_model.hub_content_name},
         ]
         if base_model.recipe_name:
-            tags.append({"key": "sagemaker-sdk:recipe-name", "value": base_model.recipe_name})
+            endpoint_tags.append(
+                {"key": "sagemaker-sdk:recipe-name", "value": base_model.recipe_name}
+            )
+
+        # Merge user-provided tags
+        endpoint_tags.extend(_tags_as_list(tags) or [])
 
         endpoint = Endpoint.create(
             endpoint_name=endpoint_name,
             endpoint_config_name=endpoint_name,
-            tags=tags,
+            tags=endpoint_tags,
         )
 
         if wait:

sagemaker-serve/tests/unit/test_model_builder_deploy.py

@@ -602,3 +602,75 @@ def test_reset_build_state_clears_upload_state(self):
 
 if __name__ == "__main__":
     unittest.main()
+
+
+class TestDeployNovaModelTags(unittest.TestCase):
+    """Test tag propagation through the Nova model-customization deploy path.
+
+    ``deploy(tags=...)`` previously dropped tags on this path, so they never reached the
+    created Endpoint.
+    """
+
+    PROJECT_TAG = {"key": "sagemaker:project-id", "value": "p-12345"}
+
+    def _make_builder(self):
+        """Build a ModelBuilder without running __init__, for private-method testing."""
+        builder = ModelBuilder.__new__(ModelBuilder)
+        builder.instance_type = "ml.g5.xlarge"
+        builder.role_arn = "arn:aws:iam::123456789012:role/TestRole"
+        builder.built_model = MagicMock(model_name="test-model")
+
+        base_model = MagicMock(hub_content_name="nova-lite", recipe_name="test-recipe")
+        model_package = MagicMock()
+        model_package.inference_specification.containers = [MagicMock(base_model=base_model)]
+        builder._fetch_model_package = MagicMock(return_value=model_package)
+        return builder
+
+    @patch("sagemaker.serve.model_builder.Endpoint")
+    @patch("sagemaker.serve.model_builder.EndpointConfig")
+    def test_user_tags_merged_with_jumpstart_tags(self, mock_endpoint_config, mock_endpoint):
+        """User tags must be merged in without displacing the JumpStart tags."""
+        builder = self._make_builder()
+
+        builder._deploy_nova_model(
+            endpoint_name="test-endpoint", wait=False, tags=[self.PROJECT_TAG]
+        )
+
+        created_tags = mock_endpoint.create.call_args.kwargs["tags"]
+        self.assertIn(self.PROJECT_TAG, created_tags)
+        self.assertIn(
+            {"key": "sagemaker-sdk:jumpstart-model-id", "value": "nova-lite"}, created_tags
+        )
+        self.assertIn(
+            {"key": "sagemaker-sdk:recipe-name", "value": "test-recipe"}, created_tags
+        )
+
+    @patch("sagemaker.serve.model_builder.Endpoint")
+    @patch("sagemaker.serve.model_builder.EndpointConfig")
+    def test_dict_form_tags_accepted(self, mock_endpoint_config, mock_endpoint):
+        """deploy() accepts a single {key: value} mapping, which must be expanded."""
+        builder = self._make_builder()
+
+        builder._deploy_nova_model(
+            endpoint_name="test-endpoint",
+            wait=False,
+            tags={"sagemaker:project-id": "p-12345"},
+        )
+
+        self.assertIn(self.PROJECT_TAG, mock_endpoint.create.call_args.kwargs["tags"])
+
+    @patch("sagemaker.serve.model_builder.Endpoint")
+    @patch("sagemaker.serve.model_builder.EndpointConfig")
+    def test_no_tags_leaves_jumpstart_tags_only(self, mock_endpoint_config, mock_endpoint):
+        """Omitting tags must not change existing behaviour."""
+        builder = self._make_builder()
+
+        builder._deploy_nova_model(endpoint_name="test-endpoint", wait=False)
+
+        self.assertEqual(
+            mock_endpoint.create.call_args.kwargs["tags"],
+            [
+                {"key": "sagemaker-sdk:jumpstart-model-id", "value": "nova-lite"},
+                {"key": "sagemaker-sdk:recipe-name", "value": "test-recipe"},
+            ],
+        )

sagemaker-train/src/sagemaker/train/base_trainer.py

@@ -720,6 +720,7 @@ def _yaml_safe_default(value):
             sagemaker_session=sagemaker_session,
             role=role,
             base_job_name=base_job_name,
+            tags=self.tags,
         )
 
         # Execute training

sagemaker-train/src/sagemaker/train/dpo_trainer.py

@@ -327,6 +327,9 @@ def train(self,
         vpc_config = self.networking if self.networking else None
         tags = _get_jumpstart_tags(self._model_name, get_sagemaker_hub_name())
 
+        # Merge user-provided tags with the JumpStart tags
+        tags.extend(self.tags or [])
+
         # Build TrainingJob.create() arguments
         create_args = {
             "training_job_name": current_training_job_name,

sagemaker-train/src/sagemaker/train/evaluate/base_evaluator.py

@@ -123,6 +123,8 @@ class BaseEvaluator(BaseModel):
             3. Model package group name string (will fetch the object and extract ARN)
             Required when model is a JumpStart model ID. Optional when model is a ModelPackage
             ARN/object (will be inferred automatically).
+        tags (Optional[List[TagsDict]]): Tags applied to the evaluation pipeline when it is
+            created, which cascade to the pipeline's step jobs.
     """
 
     region: Optional[str] = None
@@ -138,6 +140,7 @@ class BaseEvaluator(BaseModel):
     networking: Optional[VpcConfig] = None
     kms_key_id: Optional[str] = None
     model_package_group: Optional[Union[str, ModelPackageGroup]] = None
+    tags: Optional[List[TagsDict]] = None
     compute: Optional[Union[Compute, HyperPodCompute]] = None
     training_image: Optional[str] = None
     recipe: Optional[str] = None
@@ -955,6 +958,9 @@ def _start_execution(
         if self._is_jumpstart_model:
             from sagemaker.core.jumpstart.utils import add_jumpstart_model_info_tags
             tags = add_jumpstart_model_info_tags(tags, self.model, "*")
+
+        # Merge user-provided tags
+        tags.extend(self.tags or [])
 
         execution = EvaluationPipelineExecution.start(
             eval_type=eval_type,

sagemaker-train/src/sagemaker/train/evaluate/multi_turn_rl_evaluator.py

@@ -44,6 +44,35 @@
 _MAX_STOPPING_CONDITION_SECONDS = 72 * 60 * 60
 
 
+def _tags_with_capitalized_keys(tags: Optional[List[Dict[str, str]]]) -> List[Dict[str, str]]:
+    """Convert tags to the capitalized ``Key``/``Value`` form this evaluator's paths require.
+
+    Both the raw boto3 ``CreatePipeline`` call and the ``Tags`` block rendered into the
+    pipeline definition use the API's capitalized form, whereas the other evaluators hand
+    tags to the pydantic-validated ``Pipeline.create``, which takes the lowercase form.
+    Either form is accepted here so that the inherited ``tags`` field behaves the same way
+    across evaluator subclasses.
+
+    Args:
+        tags: Tags using either ``key``/``value`` or ``Key``/``Value``, or None.
+
+    Returns:
+        The tags in capitalized form; empty when none were supplied. Entries missing a key
+        or value are skipped, since SageMaker rejects them.
+    """
+    normalized = []
+    for tag in tags or []:
+        if isinstance(tag, dict):
+            key = tag.get("Key", tag.get("key"))
+            value = tag.get("Value", tag.get("value"))
+        else:
+            key = getattr(tag, "key", None)
+            value = getattr(tag, "value", None)
+        if key is not None and value is not None:
+            normalized.append({"Key": key, "Value": value})
+    return normalized
+
+
 class MultiTurnRLEvaluator(BaseEvaluator):
     """Evaluate a multi-turn RL agent model against a held-out prompt dataset.
 
@@ -532,7 +561,7 @@ def _build_job_config_doc(include_mpc: bool, mlflow_run_name: str) -> str:
             "vpc_config": bool(networking),
             "vpc_security_group_ids": vpc_security_group_ids,
             "vpc_subnets": vpc_subnets,
-            "tags": self.tags,
+            "tags": _tags_with_capitalized_keys(self.tags) or None,
             # Pre-stringified JobConfigDocument for the templates.
             "job_config_document_str": job_config_doc_str,
             "job_config_document_ft_str": job_config_doc_ft_str,
@@ -684,6 +713,11 @@ def _start_mtrl_execution(self, pipeline_definition, name, role_arn, region):
         pipeline_prefix = _get_pipeline_name_prefix(EvalType.MTRL)
         pipeline_name = pipeline_prefix
 
+        # Customer tags are merged into the pipeline tags. This path uses raw boto3, which
+        # requires the API's capitalized Key/Value form.
+        pipeline_tags = [{"Key": _TAG_SAGEMAKER_MODEL_EVALUATION, "Value": "true"}]
+        pipeline_tags.extend(_tags_with_capitalized_keys(self.tags))
+
         # Search for existing MTRL pipeline
         existing_pipeline_name = None
         try:
@@ -710,7 +744,7 @@ def _start_mtrl_execution(self, pipeline_definition, name, role_arn, region):
                 PipelineDisplayName=pipeline_name,
                 PipelineDescription="MTRL evaluation pipeline",
                 ClientRequestToken=str(uuid.uuid4()),
-                Tags=[{"Key": _TAG_SAGEMAKER_MODEL_EVALUATION, "Value": "true"}],
+                Tags=pipeline_tags,
             )
             _logger.info(f"Created pipeline: {pipeline_name}")
 

sagemaker-train/src/sagemaker/train/multi_turn_rl_trainer.py

@@ -300,13 +300,17 @@ def train(
 
         tags = _get_jumpstart_tags(self._model_name, get_sagemaker_hub_name())
 
+        # Merge user-provided tags with the JumpStart tags
+        tags.extend(self.tags or [])
+
         try:
             job = Job.create(
                 job_name=current_job_name,
                 job_category=JOB_CATEGORY,
                 role_arn=role,
                 job_config_schema_version=JOB_CONFIG_SCHEMA_VERSION,
                 job_config_document=job_config_doc,
+                tags=tags,
                 session=sagemaker_session.boto_session,
                 region=sagemaker_session.boto_session.region_name,
             )

sagemaker-train/src/sagemaker/train/rlaif_trainer.py

@@ -290,6 +290,9 @@ def train(self, training_dataset: Optional[Union[str, DataSet]] = None, validati
         vpc_config = self.networking if self.networking else None
         tags = _get_jumpstart_tags(self._model_name, get_sagemaker_hub_name())
 
+        # Merge user-provided tags with the JumpStart tags
+        tags.extend(self.tags or [])
+
         # Build TrainingJob.create() arguments
         create_args = {
             "training_job_name": current_training_job_name,

sagemaker-train/src/sagemaker/train/rlvr_trainer.py

@@ -498,6 +498,9 @@ def train(self, training_dataset: Optional[Union[str, DataSet]] = None,
         vpc_config = self.networking if self.networking else None
         tags = _get_jumpstart_tags(self._model_name, get_sagemaker_hub_name())
 
+        # Merge user-provided tags with the JumpStart tags
+        tags.extend(self.tags or [])
+
         # Build TrainingJob.create() arguments
         create_args = {
             "training_job_name": current_training_job_name,

sagemaker-train/src/sagemaker/train/sft_trainer.py

@@ -394,6 +394,9 @@ def train(self, training_dataset: Optional[Union[str, DataSet]] = None, validati
         vpc_config = self.networking if self.networking else None
         tags = _get_jumpstart_tags(self._model_name, get_sagemaker_hub_name())
 
+        # Merge user-provided tags with the JumpStart tags
+        tags.extend(self.tags or [])
+
         # Build TrainingJob.create() arguments
         create_args = {
             "training_job_name": current_training_job_name,

sagemaker-train/tests/unit/train/evaluate/test_execution.py

@@ -129,6 +129,51 @@ def test_create_pipeline_success(self, mock_pipeline_class, mock_get_name, mock_
         )
         assert result == mock_pipeline
 
+    @patch("sagemaker.train.evaluate.execution._get_pipeline_name")
+    @patch("sagemaker.train.evaluate.execution.Pipeline")
+    def test_create_pipeline_propagates_user_tags(
+        self, mock_pipeline_class, mock_get_name, mock_session
+    ):
+        """User tags must reach Pipeline.create alongside the evaluation discovery tag."""
+        mock_get_name.return_value = DEFAULT_PIPELINE_NAME
+        mock_pipeline_class.create.return_value = MagicMock()
+
+        _create_evaluation_pipeline(
+            eval_type=EvalType.BENCHMARK,
+            role_arn=DEFAULT_ROLE,
+            pipeline_definition=DEFAULT_PIPELINE_DEFINITION,
+            session=mock_session,
+            region=DEFAULT_REGION,
+            tags=[{"key": "sagemaker:project-id", "value": "p-12345"}],
+        )
+
+        created_tags = mock_pipeline_class.create.call_args.kwargs["tags"]
+        pairs = [(t.key, t.value) for t in created_tags]
+        assert ("sagemaker:project-id", "p-12345") in pairs
+        # The evaluation discovery tag must still be present.
+        assert any(key == "SagemakerModelEvaluation" for key, _ in pairs)
+
+    @patch("sagemaker.train.evaluate.execution._get_pipeline_name")
+    @patch("sagemaker.train.evaluate.execution.Pipeline")
+    def test_create_pipeline_accepts_capitalized_user_tags(
+        self, mock_pipeline_class, mock_get_name, mock_session
+    ):
+        """Capitalized user tags must also be converted into Tag objects."""
+        mock_get_name.return_value = DEFAULT_PIPELINE_NAME
+        mock_pipeline_class.create.return_value = MagicMock()
+
+        _create_evaluation_pipeline(
+            eval_type=EvalType.BENCHMARK,
+            role_arn=DEFAULT_ROLE,
+            pipeline_definition=DEFAULT_PIPELINE_DEFINITION,
+            session=mock_session,
+            region=DEFAULT_REGION,
+            tags=[{"Key": "sagemaker:project-id", "Value": "p-12345"}],
+        )
+
+        pairs = [(t.key, t.value) for t in mock_pipeline_class.create.call_args.kwargs["tags"]]
+        assert ("sagemaker:project-id", "p-12345") in pairs
+
     @patch("sagemaker.train.evaluate.execution.Pipeline")
     def test_create_pipeline_waits_for_status(self, mock_pipeline_class, mock_session):
         """Test that pipeline waits for active status."""