|
23 | 23 | HyperparameterTuner, |
24 | 24 | WarmStartTypes, |
25 | 25 | GRID_SEARCH, |
| 26 | + _to_boto_tag, |
26 | 27 | ) |
27 | 28 | from sagemaker.core.parameter import ( |
28 | 29 | CategoricalParameter, |
|
37 | 38 | Channel, |
38 | 39 | DataSource, |
39 | 40 | S3DataSource, |
| 41 | + Tag, |
40 | 42 | ) |
41 | 43 |
|
42 | 44 | # --------------------------------------------------------------------------- |
@@ -803,3 +805,107 @@ def test_build_training_job_definition_passes_through_output_data_config(self): |
803 | 805 | assert ( |
804 | 806 | definition.output_data_config.s3_output_path == "s3://bucket/output" |
805 | 807 | ), "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