From ef9999eb1d2ee288b35b31749dd40f25f099204b Mon Sep 17 00:00:00 2001 From: Rohan Gujarathi Date: Fri, 7 Aug 2026 06:35:57 +0000 Subject: [PATCH] fix: propagate user-supplied tags to created resources Tags passed to the trainer, evaluator and ModelBuilder classes were accepted but never reached the resources they create, so they could not be used for tag-based resource association. - Trainers (SFT, DPO, RLVR, RLAIF, MultiTurnRL): the serverless train() path built its tag list from the JumpStart tags only and discarded self.tags. MultiTurnRLTrainer additionally computed tags and then never passed them to Job.create(). - BaseTrainer._train_serverful_smtj() did not forward tags to ModelTrainer.from_recipe(). - BaseEvaluator had no tags field at all; add one and merge it into the pipeline tags. - MultiTurnRLEvaluator passed only the evaluation discovery tag to CreatePipeline and dropped customer tags. - ModelBuilder.deploy(tags=...) dropped tags on the model-customization and Nova paths, which bypass the normal _deploy_core_endpoint() handling. The create APIs are typed List[Tag] and validated by pydantic, which coerces plain key/value dicts and accepts Tag objects, so the merged tag lists are passed through as-is. ModelBuilder additionally accepts a single {key: value} mapping on deploy(), which is expanded into list form for those APIs. --- .../src/sagemaker/serve/model_builder.py | 36 +++++- .../tests/unit/test_model_builder_deploy.py | 72 +++++++++++ .../src/sagemaker/train/base_trainer.py | 1 + .../src/sagemaker/train/dpo_trainer.py | 3 + .../train/evaluate/base_evaluator.py | 6 + .../train/evaluate/multi_turn_rl_evaluator.py | 38 +++++- .../sagemaker/train/multi_turn_rl_trainer.py | 4 + .../src/sagemaker/train/rlaif_trainer.py | 3 + .../src/sagemaker/train/rlvr_trainer.py | 3 + .../src/sagemaker/train/sft_trainer.py | 3 + .../unit/train/evaluate/test_execution.py | 45 +++++++ .../train/evaluate/test_mtrl_evaluator.py | 116 ++++++++++++++++- .../tests/unit/train/test_sft_trainer.py | 119 ++++++++++++++++++ 13 files changed, 441 insertions(+), 8 deletions(-) diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 10db5a7d4a..e28cb9d15d 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/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: diff --git a/sagemaker-serve/tests/unit/test_model_builder_deploy.py b/sagemaker-serve/tests/unit/test_model_builder_deploy.py index 3a0fca3d8e..d9c82d96fa 100644 --- a/sagemaker-serve/tests/unit/test_model_builder_deploy.py +++ b/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"}, + ], + ) diff --git a/sagemaker-train/src/sagemaker/train/base_trainer.py b/sagemaker-train/src/sagemaker/train/base_trainer.py index c17199a611..35ed38340e 100644 --- a/sagemaker-train/src/sagemaker/train/base_trainer.py +++ b/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 diff --git a/sagemaker-train/src/sagemaker/train/dpo_trainer.py b/sagemaker-train/src/sagemaker/train/dpo_trainer.py index a83d1ed2fe..8d7da42322 100644 --- a/sagemaker-train/src/sagemaker/train/dpo_trainer.py +++ b/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, diff --git a/sagemaker-train/src/sagemaker/train/evaluate/base_evaluator.py b/sagemaker-train/src/sagemaker/train/evaluate/base_evaluator.py index 29ab44133a..99e1b33f82 100644 --- a/sagemaker-train/src/sagemaker/train/evaluate/base_evaluator.py +++ b/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, diff --git a/sagemaker-train/src/sagemaker/train/evaluate/multi_turn_rl_evaluator.py b/sagemaker-train/src/sagemaker/train/evaluate/multi_turn_rl_evaluator.py index 754c0944e4..9bdfd0ec1d 100644 --- a/sagemaker-train/src/sagemaker/train/evaluate/multi_turn_rl_evaluator.py +++ b/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}") diff --git a/sagemaker-train/src/sagemaker/train/multi_turn_rl_trainer.py b/sagemaker-train/src/sagemaker/train/multi_turn_rl_trainer.py index daf3c8f5d9..cbb5f75895 100644 --- a/sagemaker-train/src/sagemaker/train/multi_turn_rl_trainer.py +++ b/sagemaker-train/src/sagemaker/train/multi_turn_rl_trainer.py @@ -300,6 +300,9 @@ 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, @@ -307,6 +310,7 @@ def train( 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, ) diff --git a/sagemaker-train/src/sagemaker/train/rlaif_trainer.py b/sagemaker-train/src/sagemaker/train/rlaif_trainer.py index 5077c288d5..f432f6958e 100644 --- a/sagemaker-train/src/sagemaker/train/rlaif_trainer.py +++ b/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, diff --git a/sagemaker-train/src/sagemaker/train/rlvr_trainer.py b/sagemaker-train/src/sagemaker/train/rlvr_trainer.py index f20b54e7b1..2bcab9b55d 100644 --- a/sagemaker-train/src/sagemaker/train/rlvr_trainer.py +++ b/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, diff --git a/sagemaker-train/src/sagemaker/train/sft_trainer.py b/sagemaker-train/src/sagemaker/train/sft_trainer.py index 3f14c348f0..3d62d844d8 100644 --- a/sagemaker-train/src/sagemaker/train/sft_trainer.py +++ b/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, diff --git a/sagemaker-train/tests/unit/train/evaluate/test_execution.py b/sagemaker-train/tests/unit/train/evaluate/test_execution.py index c0ba1c6cad..8337823912 100644 --- a/sagemaker-train/tests/unit/train/evaluate/test_execution.py +++ b/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.""" diff --git a/sagemaker-train/tests/unit/train/evaluate/test_mtrl_evaluator.py b/sagemaker-train/tests/unit/train/evaluate/test_mtrl_evaluator.py index fd77922292..35ebc5d845 100644 --- a/sagemaker-train/tests/unit/train/evaluate/test_mtrl_evaluator.py +++ b/sagemaker-train/tests/unit/train/evaluate/test_mtrl_evaluator.py @@ -53,7 +53,9 @@ def test_creates_pipeline_when_not_found(self, mock_boto3_client, mock_pe_cls): mock_boto3_client.return_value = mock_client mock_client.list_pipelines.return_value = {"PipelineSummaries": []} mock_client.start_pipeline_execution.return_value = { - "PipelineExecutionArn": f"arn:aws:sagemaker:us-west-2:123:pipeline/{PIPELINE_PREFIX}/execution/exec-1" + "PipelineExecutionArn": ( + f"arn:aws:sagemaker:us-west-2:123:pipeline/{PIPELINE_PREFIX}/execution/exec-1" + ) } result = evaluator._start_mtrl_execution( @@ -139,6 +141,118 @@ def test_find_existing_pipeline_search_fails(self, mock_boto3_client, mock_pe_cl mock_client.create_pipeline.assert_called_once() assert result.arn.endswith("exec-1") + @patch("sagemaker.core.resources.PipelineExecution") + @patch("boto3.client") + def test_lowercase_user_tags_sent_capitalized(self, mock_boto3_client, mock_pe_cls): + """Lowercase tags must be converted before reaching boto3 CreatePipeline. + + This path uses raw boto3, which rejects the lowercase key/value form that the other + evaluators pass to the pydantic-validated Pipeline.create. + """ + evaluator = self._make_evaluator() + evaluator.tags = [{"key": "sagemaker:project-id", "value": "p-12345"}] + mock_client = MagicMock() + mock_boto3_client.return_value = mock_client + mock_client.list_pipelines.return_value = {"PipelineSummaries": []} + mock_client.start_pipeline_execution.return_value = { + "PipelineExecutionArn": ( + f"arn:aws:sagemaker:us-west-2:123:pipeline/{PIPELINE_PREFIX}/execution/exec-1" + ) + } + + evaluator._start_mtrl_execution( + pipeline_definition='{"Steps": []}', + name="test-eval", + role_arn=ROLE, + region=REGION, + ) + + sent_tags = mock_client.create_pipeline.call_args.kwargs["Tags"] + assert {"Key": "sagemaker:project-id", "Value": "p-12345"} in sent_tags + # Every entry must use the capitalized form boto3 requires. + for tag in sent_tags: + assert set(tag.keys()) == {"Key", "Value"} + + @patch("sagemaker.core.resources.PipelineExecution") + @patch("boto3.client") + def test_capitalized_user_tags_passed_through(self, mock_boto3_client, mock_pe_cls): + """Tags already in the capitalized form must be preserved unchanged.""" + evaluator = self._make_evaluator() + evaluator.tags = [{"Key": "sagemaker:project-id", "Value": "p-12345"}] + mock_client = MagicMock() + mock_boto3_client.return_value = mock_client + mock_client.list_pipelines.return_value = {"PipelineSummaries": []} + mock_client.start_pipeline_execution.return_value = { + "PipelineExecutionArn": ( + f"arn:aws:sagemaker:us-west-2:123:pipeline/{PIPELINE_PREFIX}/execution/exec-1" + ) + } + + evaluator._start_mtrl_execution( + pipeline_definition='{"Steps": []}', + name="test-eval", + role_arn=ROLE, + region=REGION, + ) + + sent_tags = mock_client.create_pipeline.call_args.kwargs["Tags"] + assert {"Key": "sagemaker:project-id", "Value": "p-12345"} in sent_tags + + @patch("sagemaker.core.resources.PipelineExecution") + @patch("boto3.client") + def test_no_user_tags_keeps_only_discovery_tag(self, mock_boto3_client, mock_pe_cls): + """Omitting tags must leave the evaluation discovery tag as the only tag.""" + evaluator = self._make_evaluator() + evaluator.tags = None + mock_client = MagicMock() + mock_boto3_client.return_value = mock_client + mock_client.list_pipelines.return_value = {"PipelineSummaries": []} + mock_client.start_pipeline_execution.return_value = { + "PipelineExecutionArn": ( + f"arn:aws:sagemaker:us-west-2:123:pipeline/{PIPELINE_PREFIX}/execution/exec-1" + ) + } + + evaluator._start_mtrl_execution( + pipeline_definition='{"Steps": []}', + name="test-eval", + role_arn=ROLE, + region=REGION, + ) + + sent_tags = mock_client.create_pipeline.call_args.kwargs["Tags"] + assert sent_tags == [{"Key": "SagemakerModelEvaluation", "Value": "true"}] + + +class TestTagsWithCapitalizedKeys: + """Tests for the tag-casing conversion used by the MTRL evaluator paths.""" + + def test_converts_lowercase(self): + from sagemaker.train.evaluate.multi_turn_rl_evaluator import _tags_with_capitalized_keys + + assert _tags_with_capitalized_keys([{"key": "a", "value": "b"}]) == [ + {"Key": "a", "Value": "b"} + ] + + def test_preserves_capitalized(self): + from sagemaker.train.evaluate.multi_turn_rl_evaluator import _tags_with_capitalized_keys + + assert _tags_with_capitalized_keys([{"Key": "a", "Value": "b"}]) == [ + {"Key": "a", "Value": "b"} + ] + + def test_returns_empty_for_none(self): + from sagemaker.train.evaluate.multi_turn_rl_evaluator import _tags_with_capitalized_keys + + assert _tags_with_capitalized_keys(None) == [] + + def test_skips_entries_missing_key_or_value(self): + from sagemaker.train.evaluate.multi_turn_rl_evaluator import _tags_with_capitalized_keys + + assert _tags_with_capitalized_keys([{"nope": "x"}, {"key": "ok", "value": "1"}]) == [ + {"Key": "ok", "Value": "1"} + ] + class TestModelResolutionWithLatestJob: """Tests for model resolution handling _latest_job (AgentRFT flow).""" diff --git a/sagemaker-train/tests/unit/train/test_sft_trainer.py b/sagemaker-train/tests/unit/train/test_sft_trainer.py index 7586a8d7df..0ccd34eb13 100644 --- a/sagemaker-train/tests/unit/train/test_sft_trainer.py +++ b/sagemaker-train/tests/unit/train/test_sft_trainer.py @@ -1,5 +1,6 @@ import pytest from unittest.mock import Mock, patch, MagicMock +from sagemaker.core.shapes import Tag from sagemaker.train.sft_trainer import SFTTrainer from sagemaker.train.common import TrainingType from sagemaker.core.resources import ModelPackage @@ -288,6 +289,124 @@ def test_train_with_tags(self, mock_training_job_create, mock_model_package_conf {"key": "sagemaker-sdk:jumpstart-hub-name", "value": "SageMakerPublicHub"} ] + @patch('sagemaker.train.sft_trainer._resolve_model_and_name') + @patch('sagemaker.train.sft_trainer._get_fine_tuning_options_and_model_arn') + @patch('sagemaker.train.sft_trainer.TrainDefaults.get_role') + @patch('sagemaker.train.sft_trainer.TrainDefaults.get_sagemaker_session') + @patch('sagemaker.train.sft_trainer._get_unique_name') + @patch('sagemaker.train.sft_trainer._validate_and_resolve_model_package_group') + @patch('sagemaker.train.sft_trainer._create_input_data_config') + @patch('sagemaker.train.sft_trainer._convert_input_data_to_channels') + @patch('sagemaker.train.sft_trainer._create_output_config') + @patch('sagemaker.train.sft_trainer._create_mlflow_config') + @patch('sagemaker.train.sft_trainer._create_model_package_config') + @patch('sagemaker.core.resources.TrainingJob.create') + def test_train_merges_user_tags_with_jumpstart_tags( + self, + mock_training_job_create, + mock_model_package_config, + mock_mlflow_config, + mock_output_config, + mock_convert_channels, + mock_input_config, + mock_validate_group, + mock_unique_name, + mock_get_sagemaker_session, + mock_get_role, + mock_get_options, + mock_resolve_model, + ): + """User-supplied tags must be propagated to the TrainingJob, not dropped.""" + mock_validate_group.return_value = "test-group" + mock_resolve_model.return_value = ("test-model", "test-model") + mock_get_sagemaker_session.return_value = Mock() + mock_fine_tuning_options = Mock() + mock_fine_tuning_options.to_dict.return_value = {"learning_rate": "0.001"} + mock_get_options.return_value = (mock_fine_tuning_options, "model-arn", False) + mock_get_role.return_value = "test-role" + mock_unique_name.return_value = "test-job-name" + mock_input_config.return_value = [Mock()] + mock_convert_channels.return_value = [Mock()] + mock_output_config.return_value = Mock() + mock_mlflow_config.return_value = Mock() + mock_model_package_config.return_value = Mock() + mock_training_job = Mock() + mock_training_job.arn = "arn:aws:sagemaker:us-east-1:123456789012:training-job/test-job" + mock_training_job.wait = Mock() + mock_training_job_create.return_value = mock_training_job + + trainer = SFTTrainer( + model="test-model", + model_package_group="test-group", + training_dataset="s3://bucket/train", + tags=[{"key": "sagemaker:project-id", "value": "p-12345"}], + ) + trainer.train(wait=False) + + call_kwargs = mock_training_job_create.call_args[1] + assert call_kwargs["tags"] == [ + {"key": "sagemaker-sdk:jumpstart-model-id", "value": "test-model"}, + {"key": "sagemaker-sdk:jumpstart-hub-name", "value": "SageMakerPublicHub"}, + {"key": "sagemaker:project-id", "value": "p-12345"}, + ] + + @patch('sagemaker.train.sft_trainer._resolve_model_and_name') + @patch('sagemaker.train.sft_trainer._get_fine_tuning_options_and_model_arn') + @patch('sagemaker.train.sft_trainer.TrainDefaults.get_role') + @patch('sagemaker.train.sft_trainer.TrainDefaults.get_sagemaker_session') + @patch('sagemaker.train.sft_trainer._get_unique_name') + @patch('sagemaker.train.sft_trainer._validate_and_resolve_model_package_group') + @patch('sagemaker.train.sft_trainer._create_input_data_config') + @patch('sagemaker.train.sft_trainer._convert_input_data_to_channels') + @patch('sagemaker.train.sft_trainer._create_output_config') + @patch('sagemaker.train.sft_trainer._create_mlflow_config') + @patch('sagemaker.train.sft_trainer._create_model_package_config') + @patch('sagemaker.core.resources.TrainingJob.create') + def test_train_accepts_tag_objects( + self, + mock_training_job_create, + mock_model_package_config, + mock_mlflow_config, + mock_output_config, + mock_convert_channels, + mock_input_config, + mock_validate_group, + mock_unique_name, + mock_get_sagemaker_session, + mock_get_role, + mock_get_options, + mock_resolve_model, + ): + """Tag objects must be accepted alongside the plain JumpStart tag dicts.""" + mock_validate_group.return_value = "test-group" + mock_resolve_model.return_value = ("test-model", "test-model") + mock_get_sagemaker_session.return_value = Mock() + mock_fine_tuning_options = Mock() + mock_fine_tuning_options.to_dict.return_value = {"learning_rate": "0.001"} + mock_get_options.return_value = (mock_fine_tuning_options, "model-arn", False) + mock_get_role.return_value = "test-role" + mock_unique_name.return_value = "test-job-name" + mock_input_config.return_value = [Mock()] + mock_convert_channels.return_value = [Mock()] + mock_output_config.return_value = Mock() + mock_mlflow_config.return_value = Mock() + mock_model_package_config.return_value = Mock() + mock_training_job = Mock() + mock_training_job.arn = "arn:aws:sagemaker:us-east-1:123456789012:training-job/test-job" + mock_training_job.wait = Mock() + mock_training_job_create.return_value = mock_training_job + + trainer = SFTTrainer( + model="test-model", + model_package_group="test-group", + training_dataset="s3://bucket/train", + tags=[Tag(key="sagemaker:project-id", value="p-12345")], + ) + trainer.train(wait=False) + + call_kwargs = mock_training_job_create.call_args[1] + assert Tag(key="sagemaker:project-id", value="p-12345") in call_kwargs["tags"] + def test_process_hyperparameters_removes_constructor_handled_keys(self): """Test that _process_hyperparameters removes keys handled by constructor inputs.""" # Create mock hyperparameters with all possible keys