From 4f5280029b0668920891f5e61d8e0c8b8581f7d4 Mon Sep 17 00:00:00 2001 From: Rohan Narayan Date: Wed, 24 Jun 2026 15:57:46 +0000 Subject: [PATCH] fix(tags): use sagemaker-sdk: prefix instead of sagemaker-studio: for JumpStart tags The SDK should tag resources with sagemaker-sdk:jumpstart-model-id (and related keys) to distinguish SDK-initiated operations from Studio UI ones. This fixes incorrect source attribution in usage data. - Rename _get_studio_tags() to _get_jumpstart_tags() in train/utils.py - Change tag key prefix from sagemaker-studio: to sagemaker-sdk: everywhere - Update all trainer call sites and unit tests --- sagemaker-serve/src/sagemaker/serve/model_builder.py | 6 +++--- sagemaker-train/src/sagemaker/train/dpo_trainer.py | 4 ++-- .../src/sagemaker/train/multi_turn_rl_trainer.py | 4 ++-- sagemaker-train/src/sagemaker/train/rlaif_trainer.py | 4 ++-- sagemaker-train/src/sagemaker/train/rlvr_trainer.py | 4 ++-- sagemaker-train/src/sagemaker/train/sft_trainer.py | 4 ++-- sagemaker-train/src/sagemaker/train/utils.py | 6 +++--- sagemaker-train/tests/unit/train/test_dpo_trainer.py | 4 ++-- sagemaker-train/tests/unit/train/test_rlaif_trainer.py | 4 ++-- sagemaker-train/tests/unit/train/test_rlvr_trainer.py | 4 ++-- sagemaker-train/tests/unit/train/test_sft_trainer.py | 4 ++-- 11 files changed, 24 insertions(+), 24 deletions(-) diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 80d067beaa..779a1414ac 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -2508,7 +2508,7 @@ def _build_single_modelbuilder( containers=[container_def], enable_network_isolation=True, tags=[ - {"key": "sagemaker-studio:jumpstart-model-id", + {"key": "sagemaker-sdk:jumpstart-model-id", "value": base_model.hub_content_name}, ], ) @@ -4848,10 +4848,10 @@ def _deploy_nova_model( ) tags = [ - {"key": "sagemaker-studio:jumpstart-model-id", "value": base_model.hub_content_name}, + {"key": "sagemaker-sdk:jumpstart-model-id", "value": base_model.hub_content_name}, ] if base_model.recipe_name: - tags.append({"key": "sagemaker-studio:recipe-name", "value": base_model.recipe_name}) + tags.append({"key": "sagemaker-sdk:recipe-name", "value": base_model.recipe_name}) endpoint = Endpoint.create( endpoint_name=endpoint_name, diff --git a/sagemaker-train/src/sagemaker/train/dpo_trainer.py b/sagemaker-train/src/sagemaker/train/dpo_trainer.py index bd5d9a11bd..7f98146ab9 100644 --- a/sagemaker-train/src/sagemaker/train/dpo_trainer.py +++ b/sagemaker-train/src/sagemaker/train/dpo_trainer.py @@ -6,7 +6,7 @@ from sagemaker.core.resources import TrainingJob, ModelPackageGroup, ModelPackage from sagemaker.core.shapes import VpcConfig from sagemaker.train.defaults import TrainDefaults -from sagemaker.train.utils import _get_unique_name, _get_studio_tags +from sagemaker.train.utils import _get_unique_name, _get_jumpstart_tags from sagemaker.train.configs import StoppingCondition from sagemaker.train.common_utils.finetune_utils import ( _get_fine_tuning_options_and_model_arn, @@ -251,7 +251,7 @@ def train(self, ) vpc_config = self.networking if self.networking else None - tags = _get_studio_tags(self._model_name, get_sagemaker_hub_name()) + tags = _get_jumpstart_tags(self._model_name, get_sagemaker_hub_name()) # Build TrainingJob.create() arguments create_args = { 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 d1f456c371..aad39c5e02 100644 --- a/sagemaker-train/src/sagemaker/train/multi_turn_rl_trainer.py +++ b/sagemaker-train/src/sagemaker/train/multi_turn_rl_trainer.py @@ -44,7 +44,7 @@ from sagemaker.train.common_utils.recipe_utils import _list_hub_models_by_recipe, _is_nova_model from sagemaker.train.constants import get_sagemaker_hub_name from sagemaker.train.defaults import TrainDefaults -from sagemaker.train.utils import _get_unique_name, _get_studio_tags +from sagemaker.train.utils import _get_unique_name, _get_jumpstart_tags logger = logging.getLogger(__name__) @@ -279,7 +279,7 @@ def train( self.training_dataset = training_dataset job_config_doc = self._build_job_config_document() - tags = _get_studio_tags(self._model_name, get_sagemaker_hub_name()) + tags = _get_jumpstart_tags(self._model_name, get_sagemaker_hub_name()) try: job = Job.create( diff --git a/sagemaker-train/src/sagemaker/train/rlaif_trainer.py b/sagemaker-train/src/sagemaker/train/rlaif_trainer.py index f2d8460989..5199c36557 100644 --- a/sagemaker-train/src/sagemaker/train/rlaif_trainer.py +++ b/sagemaker-train/src/sagemaker/train/rlaif_trainer.py @@ -5,7 +5,7 @@ from sagemaker.core.resources import TrainingJob, ModelPackageGroup, MlflowTrackingServer, ModelPackage from sagemaker.core.shapes import VpcConfig from sagemaker.train.defaults import TrainDefaults -from sagemaker.train.utils import _get_unique_name, _get_studio_tags +from sagemaker.train.utils import _get_unique_name, _get_jumpstart_tags from sagemaker.train.common_utils.recipe_utils import _get_hub_content_metadata from sagemaker.ai_registry.dataset import DataSet from sagemaker.ai_registry.evaluator import Evaluator @@ -268,7 +268,7 @@ def train(self, training_dataset: Optional[Union[str, DataSet]] = None, validati ) vpc_config = self.networking if self.networking else None - tags = _get_studio_tags(self._model_name, get_sagemaker_hub_name()) + tags = _get_jumpstart_tags(self._model_name, get_sagemaker_hub_name()) # Build TrainingJob.create() arguments create_args = { diff --git a/sagemaker-train/src/sagemaker/train/rlvr_trainer.py b/sagemaker-train/src/sagemaker/train/rlvr_trainer.py index 333a93fc55..5d9a374039 100644 --- a/sagemaker-train/src/sagemaker/train/rlvr_trainer.py +++ b/sagemaker-train/src/sagemaker/train/rlvr_trainer.py @@ -5,7 +5,7 @@ from sagemaker.core.resources import TrainingJob, ModelPackageGroup, MlflowTrackingServer, ModelPackage from sagemaker.core.shapes import VpcConfig from sagemaker.train.defaults import TrainDefaults -from sagemaker.train.utils import _get_unique_name, _get_studio_tags +from sagemaker.train.utils import _get_unique_name, _get_jumpstart_tags from sagemaker.ai_registry.dataset import DataSet from sagemaker.ai_registry.evaluator import Evaluator from sagemaker.train.configs import StoppingCondition @@ -259,7 +259,7 @@ def train(self, training_dataset: Optional[Union[str, DataSet]] = None, ) vpc_config = self.networking if self.networking else None - tags = _get_studio_tags(self._model_name, get_sagemaker_hub_name()) + tags = _get_jumpstart_tags(self._model_name, get_sagemaker_hub_name()) # Build TrainingJob.create() arguments create_args = { diff --git a/sagemaker-train/src/sagemaker/train/sft_trainer.py b/sagemaker-train/src/sagemaker/train/sft_trainer.py index 233f169d0f..00e24f6f30 100644 --- a/sagemaker-train/src/sagemaker/train/sft_trainer.py +++ b/sagemaker-train/src/sagemaker/train/sft_trainer.py @@ -6,7 +6,7 @@ from sagemaker.core.resources import TrainingJob, ModelPackageGroup, ModelPackage from sagemaker.core.shapes import VpcConfig from sagemaker.train.defaults import TrainDefaults -from sagemaker.train.utils import _get_unique_name, _get_studio_tags +from sagemaker.train.utils import _get_unique_name, _get_jumpstart_tags from sagemaker.ai_registry.dataset import DataSet from sagemaker.train.configs import StoppingCondition from sagemaker.train.common_utils.finetune_utils import ( @@ -250,7 +250,7 @@ def train(self, training_dataset: Optional[Union[str, DataSet]] = None, validati ) vpc_config = self.networking if self.networking else None - tags = _get_studio_tags(self._model_name, get_sagemaker_hub_name()) + tags = _get_jumpstart_tags(self._model_name, get_sagemaker_hub_name()) # Build TrainingJob.create() arguments create_args = { diff --git a/sagemaker-train/src/sagemaker/train/utils.py b/sagemaker-train/src/sagemaker/train/utils.py index 0abd7596b5..9ffba1d716 100644 --- a/sagemaker-train/src/sagemaker/train/utils.py +++ b/sagemaker-train/src/sagemaker/train/utils.py @@ -240,14 +240,14 @@ def _run_clone_command_silent(repo_url, dest_dir): logger.error(f"Error output:\n{e}") raise -def _get_studio_tags(model_id: str, hub_name: str): +def _get_jumpstart_tags(model_id: str, hub_name: str): return [ { - "key": "sagemaker-studio:jumpstart-model-id", + "key": "sagemaker-sdk:jumpstart-model-id", "value": model_id }, { - "key": "sagemaker-studio:jumpstart-hub-name", + "key": "sagemaker-sdk:jumpstart-hub-name", "value": hub_name } ] diff --git a/sagemaker-train/tests/unit/train/test_dpo_trainer.py b/sagemaker-train/tests/unit/train/test_dpo_trainer.py index 1b70e0bf89..7b6eaee7a5 100644 --- a/sagemaker-train/tests/unit/train/test_dpo_trainer.py +++ b/sagemaker-train/tests/unit/train/test_dpo_trainer.py @@ -269,8 +269,8 @@ def test_train_with_tags(self, mock_training_job_create, mock_model_package_conf mock_training_job_create.assert_called_once() call_kwargs = mock_training_job_create.call_args[1] assert call_kwargs["tags"] == [ - {"key": "sagemaker-studio:jumpstart-model-id", "value": "test-model"}, - {"key": "sagemaker-studio:jumpstart-hub-name", "value": "SageMakerPublicHub"} + {"key": "sagemaker-sdk:jumpstart-model-id", "value": "test-model"}, + {"key": "sagemaker-sdk:jumpstart-hub-name", "value": "SageMakerPublicHub"} ] @patch('sagemaker.train.dpo_trainer._validate_and_resolve_model_package_group') diff --git a/sagemaker-train/tests/unit/train/test_rlaif_trainer.py b/sagemaker-train/tests/unit/train/test_rlaif_trainer.py index e5666883e8..a4f16fee42 100644 --- a/sagemaker-train/tests/unit/train/test_rlaif_trainer.py +++ b/sagemaker-train/tests/unit/train/test_rlaif_trainer.py @@ -274,8 +274,8 @@ def test_train_with_tags(self, mock_training_job_create, mock_model_package_conf mock_training_job_create.assert_called_once() call_kwargs = mock_training_job_create.call_args[1] assert call_kwargs["tags"] == [ - {"key": "sagemaker-studio:jumpstart-model-id", "value": "test-model"}, - {"key": "sagemaker-studio:jumpstart-hub-name", "value": "SageMakerPublicHub"} + {"key": "sagemaker-sdk:jumpstart-model-id", "value": "test-model"}, + {"key": "sagemaker-sdk:jumpstart-hub-name", "value": "SageMakerPublicHub"} ] @patch('sagemaker.train.rlaif_trainer._validate_and_resolve_model_package_group') diff --git a/sagemaker-train/tests/unit/train/test_rlvr_trainer.py b/sagemaker-train/tests/unit/train/test_rlvr_trainer.py index 320b81555d..25a5d334a8 100644 --- a/sagemaker-train/tests/unit/train/test_rlvr_trainer.py +++ b/sagemaker-train/tests/unit/train/test_rlvr_trainer.py @@ -272,8 +272,8 @@ def test_train_with_tags(self, mock_training_job_create, mock_model_package_conf mock_training_job_create.assert_called_once() call_kwargs = mock_training_job_create.call_args[1] assert call_kwargs["tags"] == [ - {"key": "sagemaker-studio:jumpstart-model-id", "value": "test-model"}, - {"key": "sagemaker-studio:jumpstart-hub-name", "value": "SageMakerPublicHub"} + {"key": "sagemaker-sdk:jumpstart-model-id", "value": "test-model"}, + {"key": "sagemaker-sdk:jumpstart-hub-name", "value": "SageMakerPublicHub"} ] @patch('sagemaker.train.rlvr_trainer._validate_and_resolve_model_package_group') diff --git a/sagemaker-train/tests/unit/train/test_sft_trainer.py b/sagemaker-train/tests/unit/train/test_sft_trainer.py index 108990f839..b043543e4b 100644 --- a/sagemaker-train/tests/unit/train/test_sft_trainer.py +++ b/sagemaker-train/tests/unit/train/test_sft_trainer.py @@ -284,8 +284,8 @@ def test_train_with_tags(self, mock_training_job_create, mock_model_package_conf mock_training_job_create.assert_called_once() call_kwargs = mock_training_job_create.call_args[1] assert call_kwargs["tags"] == [ - {"key": "sagemaker-studio:jumpstart-model-id", "value": "test-model"}, - {"key": "sagemaker-studio:jumpstart-hub-name", "value": "SageMakerPublicHub"} + {"key": "sagemaker-sdk:jumpstart-model-id", "value": "test-model"}, + {"key": "sagemaker-sdk:jumpstart-hub-name", "value": "SageMakerPublicHub"} ] def test_process_hyperparameters_removes_constructor_handled_keys(self):