Skip to content

Commit c65ad99

Browse files
mohamedzeidan2021Mohamed Zeidan
andauthored
fix: normalize ModelTrainer Tag objects inherited by HyperparameterTuner (#6482)
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 Co-authored-by: Mohamed Zeidan <zeidmo@amazon.com>
1 parent ce3ac25 commit c65ad99

2 files changed

Lines changed: 137 additions & 4 deletions

File tree

‎sagemaker-train/src/sagemaker/train/tuner.py‎

Lines changed: 31 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,25 @@
7373
logger = logging.getLogger(__name__)
7474

7575

76+
def _to_boto_tag(tag):
77+
"""Normalize a single tag into boto-style ``{"Key": ..., "Value": ...}`` form.
78+
79+
``ModelTrainer`` stores tags as ``sagemaker.core.shapes.Tag`` objects, while the
80+
tuner carries boto-style dictionaries, so inherited tags must be converted before
81+
they are merged.
82+
"""
83+
if isinstance(tag, dict):
84+
if "key" in tag and "value" in tag and "Key" not in tag:
85+
return {"Key": tag["key"], "Value": tag["value"]}
86+
return tag
87+
88+
key = getattr(tag, "key", None)
89+
if key is not None:
90+
return {"Key": key, "Value": getattr(tag, "value", None)}
91+
92+
return tag
93+
94+
7695
class WarmStartTypes(Enum):
7796
"""Types of warm start supported for hyperparameter tuning jobs."""
7897

@@ -315,10 +334,15 @@ def _prepare_tags_for_tuning(self):
315334
self.model_trainer or self.model_trainer_dict[sorted(self.model_trainer_dict.keys())[0]]
316335
)
317336

318-
model_trainer_tags = getattr(model_trainer, "tags", []) or []
337+
model_trainer_tags = [
338+
_to_boto_tag(tag) for tag in (getattr(model_trainer, "tags", None) or [])
339+
]
319340

320-
if self.tags is None and len(model_trainer_tags) > 0:
321-
self.tags = []
341+
if self.tags is None:
342+
if model_trainer_tags:
343+
self.tags = []
344+
else:
345+
self.tags = [_to_boto_tag(tag) for tag in self.tags]
322346

323347
for tag in model_trainer_tags:
324348
if tag not in self.tags:
@@ -1277,7 +1301,10 @@ def _start_tuning_job(self, inputs):
12771301
if self.tags:
12781302
from sagemaker.core.shapes import Tag
12791303

1280-
tag_objects = [Tag(key=tag["Key"], value=tag["Value"]) for tag in self.tags]
1304+
tag_objects = [
1305+
Tag(key=boto_tag["Key"], value=boto_tag["Value"])
1306+
for boto_tag in (_to_boto_tag(tag) for tag in self.tags)
1307+
]
12811308

12821309
# Build tuning request
12831310
tuning_request = {

‎sagemaker-train/tests/unit/train/test_tuner.py‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
HyperparameterTuner,
2424
WarmStartTypes,
2525
GRID_SEARCH,
26+
_to_boto_tag,
2627
)
2728
from sagemaker.core.parameter import (
2829
CategoricalParameter,
@@ -37,6 +38,7 @@
3738
Channel,
3839
DataSource,
3940
S3DataSource,
41+
Tag,
4042
)
4143

4244
# ---------------------------------------------------------------------------
@@ -803,3 +805,107 @@ def test_build_training_job_definition_passes_through_output_data_config(self):
803805
assert (
804806
definition.output_data_config.s3_output_path == "s3://bucket/output"
805807
), "s3_output_path should be preserved"
808+
809+
810+
class TestHyperparameterTunerTagInheritance:
811+
"""Test tag inheritance from ModelTrainer (issue #6472)."""
812+
813+
@pytest.fixture
814+
def trainer_with_tag_objects(self):
815+
"""Create a mock ModelTrainer whose tags are v3 ``Tag`` objects."""
816+
trainer = _create_mock_model_trainer()
817+
trainer.tags = [Tag(key="project", value="ml-project")]
818+
return trainer
819+
820+
@staticmethod
821+
def _create_tuner(model_trainer, tags=None, include_jumpstart_tags=False):
822+
"""Create a tuner with JumpStart tagging disabled by default."""
823+
tuner = HyperparameterTuner(
824+
model_trainer=model_trainer,
825+
objective_metric_name="accuracy",
826+
hyperparameter_ranges=_create_single_hp_range(),
827+
tags=tags,
828+
)
829+
tuner.sagemaker_session.settings.include_jumpstart_tags = include_jumpstart_tags
830+
return tuner
831+
832+
def test_inherits_tag_objects_from_model_trainer(self, trainer_with_tag_objects):
833+
"""Test v3 ``Tag`` objects are normalized to boto-style dicts when inherited."""
834+
tuner = self._create_tuner(trainer_with_tag_objects)
835+
836+
tuner._prepare_tags_for_tuning()
837+
838+
assert tuner.tags == [{"Key": "project", "Value": "ml-project"}]
839+
840+
def test_merges_tag_objects_with_boto_style_tuner_tags(self, trainer_with_tag_objects):
841+
"""Test inherited ``Tag`` objects merge with the tuner's own boto-style tags."""
842+
tuner = self._create_tuner(
843+
trainer_with_tag_objects, tags=[{"Key": "owner", "Value": "team"}]
844+
)
845+
846+
tuner._prepare_tags_for_tuning()
847+
848+
assert tuner.tags == [
849+
{"Key": "owner", "Value": "team"},
850+
{"Key": "project", "Value": "ml-project"},
851+
]
852+
853+
def test_does_not_duplicate_equivalent_tags(self, trainer_with_tag_objects):
854+
"""Test a tuner tag equal to an inherited ``Tag`` object is not added twice."""
855+
tuner = self._create_tuner(
856+
trainer_with_tag_objects, tags=[{"Key": "project", "Value": "ml-project"}]
857+
)
858+
859+
tuner._prepare_tags_for_tuning()
860+
861+
assert tuner.tags == [{"Key": "project", "Value": "ml-project"}]
862+
863+
def test_tags_stay_none_when_no_tags_are_set(self):
864+
"""Test tags remain None when neither the tuner nor the trainer has tags."""
865+
trainer = _create_mock_model_trainer()
866+
trainer.tags = None
867+
tuner = self._create_tuner(trainer)
868+
869+
tuner._prepare_tags_for_tuning()
870+
871+
assert tuner.tags is None
872+
873+
def test_jumpstart_tag_helper_receives_boto_style_tags(self, trainer_with_tag_objects):
874+
"""Test tags are normalized before being passed to the JumpStart tag helper.
875+
876+
``add_jumpstart_uri_tags`` subscripts each tag with ``tag["Key"]``, so it fails
877+
on ``Tag`` objects too.
878+
"""
879+
tuner = self._create_tuner(trainer_with_tag_objects, include_jumpstart_tags=True)
880+
881+
with patch("sagemaker.train.tuner.add_jumpstart_uri_tags") as mock_add_tags:
882+
mock_add_tags.side_effect = lambda tags, **kwargs: tags
883+
tuner._prepare_tags_for_tuning()
884+
885+
passed_tags = mock_add_tags.call_args[1]["tags"]
886+
assert passed_tags == [{"Key": "project", "Value": "ml-project"}]
887+
888+
def test_start_tuning_job_converts_tag_objects(self, trainer_with_tag_objects):
889+
"""Test ``_start_tuning_job`` accepts ``Tag`` objects left on ``tuner.tags``."""
890+
tuner = self._create_tuner(trainer_with_tag_objects)
891+
tuner._current_job_name = "test-tuning-job"
892+
tuner.tags = [Tag(key="project", value="ml-project")]
893+
894+
with patch("sagemaker.train.tuner.HyperParameterTuningJob") as mock_tuning_job:
895+
tuner._start_tuning_job(None)
896+
897+
tags = mock_tuning_job.create.call_args[1]["tags"]
898+
assert tags == [Tag(key="project", value="ml-project")]
899+
900+
def test_to_boto_tag_normalizes_lowercase_dicts(self):
901+
"""Test lowercase tag dicts are normalized to boto-style dicts."""
902+
assert _to_boto_tag({"key": "project", "value": "ml-project"}) == {
903+
"Key": "project",
904+
"Value": "ml-project",
905+
}
906+
907+
def test_to_boto_tag_passes_through_unknown_tags(self):
908+
"""Test tags that are neither ``Tag`` objects nor dicts are left untouched."""
909+
parameter = ParameterString(name="TagValue")
910+
911+
assert _to_boto_tag(parameter) is parameter

0 commit comments

Comments
 (0)