From dd2fb22510d28b3b289e4d0449c77819def51d70 Mon Sep 17 00:00:00 2001 From: Mohamed Zeidan Date: Thu, 15 Oct 2026 13:38:14 -0700 Subject: [PATCH] fix: normalize ModelTrainer Tag objects inherited by HyperparameterTuner ModelTrainer stores tags as sagemaker.core.shapes.Tag objects, while HyperparameterTuner carries boto-style {"Key": ..., "Value": ...} dicts. _prepare_tags_for_tuning() merged the trainer's Tag objects into that list without converting them, so the resulting list held both shapes. Both consumers of that list subscript each tag: * add_jumpstart_uri_tags() -> tag_key_in_array() reads tag["Key"] * _start_tuning_job() builds Tag(key=tag["Key"], value=tag["Value"]) Either path raised TypeError: 'Tag' object is not subscriptable before the tuning job was created, so any tagged ModelTrainer broke tune(). Normalize each tag to boto-style form before merging, which fixes both call sites and lets the dedup check compare equivalent tags. Tags are left as None when neither the tuner nor the trainer defines any. Fixes #6472 --- sagemaker-train/src/sagemaker/train/tuner.py | 35 +++++- .../tests/unit/train/test_tuner.py | 106 ++++++++++++++++++ 2 files changed, 137 insertions(+), 4 deletions(-) diff --git a/sagemaker-train/src/sagemaker/train/tuner.py b/sagemaker-train/src/sagemaker/train/tuner.py index a3b2098591..6b0a2bcb4e 100644 --- a/sagemaker-train/src/sagemaker/train/tuner.py +++ b/sagemaker-train/src/sagemaker/train/tuner.py @@ -73,6 +73,25 @@ logger = logging.getLogger(__name__) +def _to_boto_tag(tag): + """Normalize a single tag into boto-style ``{"Key": ..., "Value": ...}`` form. + + ``ModelTrainer`` stores tags as ``sagemaker.core.shapes.Tag`` objects, while the + tuner carries boto-style dictionaries, so inherited tags must be converted before + they are merged. + """ + if isinstance(tag, dict): + if "key" in tag and "value" in tag and "Key" not in tag: + return {"Key": tag["key"], "Value": tag["value"]} + return tag + + key = getattr(tag, "key", None) + if key is not None: + return {"Key": key, "Value": getattr(tag, "value", None)} + + return tag + + class WarmStartTypes(Enum): """Types of warm start supported for hyperparameter tuning jobs.""" @@ -315,10 +334,15 @@ def _prepare_tags_for_tuning(self): self.model_trainer or self.model_trainer_dict[sorted(self.model_trainer_dict.keys())[0]] ) - model_trainer_tags = getattr(model_trainer, "tags", []) or [] + model_trainer_tags = [ + _to_boto_tag(tag) for tag in (getattr(model_trainer, "tags", None) or []) + ] - if self.tags is None and len(model_trainer_tags) > 0: - self.tags = [] + if self.tags is None: + if model_trainer_tags: + self.tags = [] + else: + self.tags = [_to_boto_tag(tag) for tag in self.tags] for tag in model_trainer_tags: if tag not in self.tags: @@ -1277,7 +1301,10 @@ def _start_tuning_job(self, inputs): if self.tags: from sagemaker.core.shapes import Tag - tag_objects = [Tag(key=tag["Key"], value=tag["Value"]) for tag in self.tags] + tag_objects = [ + Tag(key=boto_tag["Key"], value=boto_tag["Value"]) + for boto_tag in (_to_boto_tag(tag) for tag in self.tags) + ] # Build tuning request tuning_request = { diff --git a/sagemaker-train/tests/unit/train/test_tuner.py b/sagemaker-train/tests/unit/train/test_tuner.py index 0def7315f5..40b89c9869 100644 --- a/sagemaker-train/tests/unit/train/test_tuner.py +++ b/sagemaker-train/tests/unit/train/test_tuner.py @@ -23,6 +23,7 @@ HyperparameterTuner, WarmStartTypes, GRID_SEARCH, + _to_boto_tag, ) from sagemaker.core.parameter import ( CategoricalParameter, @@ -37,6 +38,7 @@ Channel, DataSource, S3DataSource, + Tag, ) # --------------------------------------------------------------------------- @@ -803,3 +805,107 @@ def test_build_training_job_definition_passes_through_output_data_config(self): assert ( definition.output_data_config.s3_output_path == "s3://bucket/output" ), "s3_output_path should be preserved" + + +class TestHyperparameterTunerTagInheritance: + """Test tag inheritance from ModelTrainer (issue #6472).""" + + @pytest.fixture + def trainer_with_tag_objects(self): + """Create a mock ModelTrainer whose tags are v3 ``Tag`` objects.""" + trainer = _create_mock_model_trainer() + trainer.tags = [Tag(key="project", value="ml-project")] + return trainer + + @staticmethod + def _create_tuner(model_trainer, tags=None, include_jumpstart_tags=False): + """Create a tuner with JumpStart tagging disabled by default.""" + tuner = HyperparameterTuner( + model_trainer=model_trainer, + objective_metric_name="accuracy", + hyperparameter_ranges=_create_single_hp_range(), + tags=tags, + ) + tuner.sagemaker_session.settings.include_jumpstart_tags = include_jumpstart_tags + return tuner + + def test_inherits_tag_objects_from_model_trainer(self, trainer_with_tag_objects): + """Test v3 ``Tag`` objects are normalized to boto-style dicts when inherited.""" + tuner = self._create_tuner(trainer_with_tag_objects) + + tuner._prepare_tags_for_tuning() + + assert tuner.tags == [{"Key": "project", "Value": "ml-project"}] + + def test_merges_tag_objects_with_boto_style_tuner_tags(self, trainer_with_tag_objects): + """Test inherited ``Tag`` objects merge with the tuner's own boto-style tags.""" + tuner = self._create_tuner( + trainer_with_tag_objects, tags=[{"Key": "owner", "Value": "team"}] + ) + + tuner._prepare_tags_for_tuning() + + assert tuner.tags == [ + {"Key": "owner", "Value": "team"}, + {"Key": "project", "Value": "ml-project"}, + ] + + def test_does_not_duplicate_equivalent_tags(self, trainer_with_tag_objects): + """Test a tuner tag equal to an inherited ``Tag`` object is not added twice.""" + tuner = self._create_tuner( + trainer_with_tag_objects, tags=[{"Key": "project", "Value": "ml-project"}] + ) + + tuner._prepare_tags_for_tuning() + + assert tuner.tags == [{"Key": "project", "Value": "ml-project"}] + + def test_tags_stay_none_when_no_tags_are_set(self): + """Test tags remain None when neither the tuner nor the trainer has tags.""" + trainer = _create_mock_model_trainer() + trainer.tags = None + tuner = self._create_tuner(trainer) + + tuner._prepare_tags_for_tuning() + + assert tuner.tags is None + + def test_jumpstart_tag_helper_receives_boto_style_tags(self, trainer_with_tag_objects): + """Test tags are normalized before being passed to the JumpStart tag helper. + + ``add_jumpstart_uri_tags`` subscripts each tag with ``tag["Key"]``, so it fails + on ``Tag`` objects too. + """ + tuner = self._create_tuner(trainer_with_tag_objects, include_jumpstart_tags=True) + + with patch("sagemaker.train.tuner.add_jumpstart_uri_tags") as mock_add_tags: + mock_add_tags.side_effect = lambda tags, **kwargs: tags + tuner._prepare_tags_for_tuning() + + passed_tags = mock_add_tags.call_args[1]["tags"] + assert passed_tags == [{"Key": "project", "Value": "ml-project"}] + + def test_start_tuning_job_converts_tag_objects(self, trainer_with_tag_objects): + """Test ``_start_tuning_job`` accepts ``Tag`` objects left on ``tuner.tags``.""" + tuner = self._create_tuner(trainer_with_tag_objects) + tuner._current_job_name = "test-tuning-job" + tuner.tags = [Tag(key="project", value="ml-project")] + + with patch("sagemaker.train.tuner.HyperParameterTuningJob") as mock_tuning_job: + tuner._start_tuning_job(None) + + tags = mock_tuning_job.create.call_args[1]["tags"] + assert tags == [Tag(key="project", value="ml-project")] + + def test_to_boto_tag_normalizes_lowercase_dicts(self): + """Test lowercase tag dicts are normalized to boto-style dicts.""" + assert _to_boto_tag({"key": "project", "value": "ml-project"}) == { + "Key": "project", + "Value": "ml-project", + } + + def test_to_boto_tag_passes_through_unknown_tags(self): + """Test tags that are neither ``Tag`` objects nor dicts are left untouched.""" + parameter = ParameterString(name="TagValue") + + assert _to_boto_tag(parameter) is parameter