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