From 92bb181bae3032b259a615fa59ef1a5fc00eb599 Mon Sep 17 00:00:00 2001 From: Roja Reddy Sareddy Date: Tue, 21 Jul 2026 13:44:45 -0700 Subject: [PATCH 1/4] Fix role issue in mtrl integ tests --- .../tests/integ/train/test_mtrl_evaluator.py | 15 +++------------ .../integ/train/test_mtrl_evaluator_3p_agent.py | 7 ------- .../integ/train/test_mtrl_trainer_integration.py | 8 +------- .../test_multi_turn_rl_trainer_integration.py | 4 ---- 4 files changed, 4 insertions(+), 30 deletions(-) diff --git a/sagemaker-train/tests/integ/train/test_mtrl_evaluator.py b/sagemaker-train/tests/integ/train/test_mtrl_evaluator.py index b79a2763dd..12ea51b36a 100644 --- a/sagemaker-train/tests/integ/train/test_mtrl_evaluator.py +++ b/sagemaker-train/tests/integ/train/test_mtrl_evaluator.py @@ -47,7 +47,6 @@ def _get_test_config(): "s3_output_path": f"s3://sagemaker-{_REGION}-{account_id}/model-evaluation/output-artifacts/", "mlflow_resource_arn": f"arn:aws:sagemaker:{_REGION}:{account_id}:mlflow-app/app-TTAUWUNMUHH6", "model_package_group": f"arn:aws:sagemaker:{_REGION}:{account_id}:model-package-group/openai-reasoning-gpt-oss-20b-mtrl-mpg", - "role": f"arn:aws:iam::{account_id}:role/Admin", "region": _REGION, "account_id": account_id, } @@ -172,7 +171,6 @@ def test_bedrock_agent_config_fields(self, mtrl_trainer, test_config): dataset=test_config["dataset"], s3_output_path=f'{test_config["s3_output_path"]}integ-fields-bedrock/', mlflow_resource_arn=test_config["mlflow_resource_arn"], - role=test_config["role"], region=test_config["region"], agent_config=test_config["agent_arn"], agent_qualifier="PROD", @@ -182,8 +180,7 @@ def test_bedrock_agent_config_fields(self, mtrl_trainer, test_config): evaluator._resolve_agent_arn() ctx = evaluator._build_template_context( - aws_context={"region": test_config["region"], "account_id": test_config["account_id"], - "role_arn": test_config["role"]}, + aws_context={"region": test_config["region"], "account_id": test_config["account_id"]}, artifacts={}, model_package_group_arn=test_config["model_package_group"], ) @@ -208,7 +205,6 @@ def test_lambda_agent_config_fields(self, mtrl_trainer, test_config): dataset=test_config["dataset"], s3_output_path=f'{test_config["s3_output_path"]}integ-fields-lambda/', mlflow_resource_arn=test_config["mlflow_resource_arn"], - role=test_config["role"], region=test_config["region"], agent_config=lambda_arn, ) @@ -217,8 +213,7 @@ def test_lambda_agent_config_fields(self, mtrl_trainer, test_config): evaluator._resolve_agent_arn() ctx = evaluator._build_template_context( - aws_context={"region": test_config["region"], "account_id": test_config["account_id"], - "role_arn": test_config["role"]}, + aws_context={"region": test_config["region"], "account_id": test_config["account_id"]}, artifacts={}, model_package_group_arn=test_config["model_package_group"], ) @@ -239,7 +234,6 @@ def test_model_package_config_fields(self, mtrl_trainer, test_config): dataset=test_config["dataset"], s3_output_path=f'{test_config["s3_output_path"]}integ-fields-mpc/', mlflow_resource_arn=test_config["mlflow_resource_arn"], - role=test_config["role"], region=test_config["region"], agent_config=test_config["agent_arn"], ) @@ -248,8 +242,7 @@ def test_model_package_config_fields(self, mtrl_trainer, test_config): evaluator._resolve_agent_arn() ctx = evaluator._build_template_context( - aws_context={"region": test_config["region"], "account_id": test_config["account_id"], - "role_arn": test_config["role"]}, + aws_context={"region": test_config["region"], "account_id": test_config["account_id"]}, artifacts={}, model_package_group_arn=test_config["model_package_group"], ) @@ -280,7 +273,6 @@ def test_evaluator_construction_with_trainer(self, mtrl_trainer, test_config): dataset=test_config["dataset"], s3_output_path=f'{test_config["s3_output_path"]}integ-construct/', mlflow_resource_arn=test_config["mlflow_resource_arn"], - role=test_config["role"], region=test_config["region"], agent_config=test_config["agent_arn"], ) @@ -298,7 +290,6 @@ def test_evaluator_construction_with_base_model(self, test_config): s3_output_path=f'{test_config["s3_output_path"]}integ-base/', agent_config=test_config["agent_arn"], mlflow_resource_arn=test_config["mlflow_resource_arn"], - role=test_config["role"], region=test_config["region"], ) diff --git a/sagemaker-train/tests/integ/train/test_mtrl_evaluator_3p_agent.py b/sagemaker-train/tests/integ/train/test_mtrl_evaluator_3p_agent.py index 2d16d4318f..1a9d1402c3 100644 --- a/sagemaker-train/tests/integ/train/test_mtrl_evaluator_3p_agent.py +++ b/sagemaker-train/tests/integ/train/test_mtrl_evaluator_3p_agent.py @@ -161,10 +161,6 @@ def _get_3p_test_config(): "MTRL_3P_MLFLOW_ARN", f"arn:aws:sagemaker:{_REGION}:{account_id}:mlflow-app/app-TTAUWUNMUHH6", ), - "role": os.environ.get( - "MTRL_3P_ROLE", - f"arn:aws:iam::{account_id}:role/Admin", - ), "region": os.environ.get("MTRL_3P_REGION", _REGION), "account_id": account_id, } @@ -258,7 +254,6 @@ def test_evaluate_with_lambda_agent_wait_for_completion(self, lambda_agent_arn, agent_config=lambda_agent_arn, s3_output_path=f'{test_config["s3_output_path"]}lambda-e2e/', mlflow_resource_arn=test_config["mlflow_resource_arn"], - role=test_config["role"], region=test_config["region"], accept_eula=True, ) @@ -304,7 +299,6 @@ def test_evaluate_base_model_with_agent_lambda_object(self, lambda_agent_arn, te agent_config=agent, s3_output_path=f'{test_config["s3_output_path"]}lambda-object/', mlflow_resource_arn=test_config["mlflow_resource_arn"], - role=test_config["role"], region=test_config["region"], accept_eula=True, ) @@ -332,7 +326,6 @@ def test_evaluate_with_attached_trainer(self, lambda_agent_arn, test_config): agent_config=lambda_agent_arn, s3_output_path=f'{test_config["s3_output_path"]}attached-trainer/', mlflow_resource_arn=test_config["mlflow_resource_arn"], - role=test_config["role"], region=test_config["region"], accept_eula=True, ) diff --git a/sagemaker-train/tests/integ/train/test_mtrl_trainer_integration.py b/sagemaker-train/tests/integ/train/test_mtrl_trainer_integration.py index 880900f4b8..9f30c3c86b 100644 --- a/sagemaker-train/tests/integ/train/test_mtrl_trainer_integration.py +++ b/sagemaker-train/tests/integ/train/test_mtrl_trainer_integration.py @@ -44,6 +44,7 @@ def _get_account_id(): boto_session = boto3.Session(region_name=_REGION) return boto_session.client("sts").get_caller_identity()["Account"] + # ============================================================ # Per-account resource configuration # ============================================================ @@ -60,7 +61,6 @@ def _get_account_id(): "s3_output_path": "s3://sagemaker-us-west-2-729646638167/mtrl-integ/eval-output/", "mlflow_resource_arn": "arn:aws:sagemaker:us-west-2:729646638167:mlflow-app/app-TTAUWUNMUHH6", "model_package_group": "arn:aws:sagemaker:us-west-2:729646638167:model-package-group/openai-reasoning-gpt-oss-20b-mtrl-mpg", - "role": "arn:aws:iam::729646638167:role/Admin", }, # PREPROD — Staging account (391266019386) "391266019386": { @@ -72,7 +72,6 @@ def _get_account_id(): "s3_output_path": "s3://sagemaker-us-west-2-391266019386/mtrl-integ/eval-output/", "mlflow_resource_arn": "arn:aws:sagemaker:us-west-2:391266019386:mlflow-app/app-P3FRQFRQTNGI", "model_package_group": "arn:aws:sagemaker:us-west-2:391266019386:model-package-group/mtrl-integ-gpt-oss-agentcore", - "role": "arn:aws:iam::391266019386:role/Admin", }, # BETA — Dev/test account (742774200982) "742774200982": { @@ -84,7 +83,6 @@ def _get_account_id(): "s3_output_path": "s3://sagemaker-us-west-2-742774200982/mtrl-integ/eval-output/", "mlflow_resource_arn": "arn:aws:sagemaker:us-west-2:742774200982:mlflow-app/app-6ZU5TXXH2GUX", "model_package_group": "arn:aws:sagemaker:us-west-2:742774200982:model-package-group/openai-reasoning-gpt-oss-20b-mtrl-mpg", - "role": "arn:aws:iam::742774200982:role/Admin", }, } @@ -133,7 +131,6 @@ def attached_trainer(config): output_model_package_group=config["model_package_group"], mlflow_app_arn=config["mlflow_resource_arn"], s3_output_path=config["s3_output_path"], - role=config["role"], accept_eula=True, ) trainer._latest_job = job @@ -165,7 +162,6 @@ def test_evaluate_finetuned_model(self, attached_trainer, config): dataset=config["dataset"], s3_output_path=f'{config["s3_output_path"]}finetuned/', mlflow_resource_arn=config["mlflow_resource_arn"], - role=config["role"], region=_REGION, ) @@ -195,7 +191,6 @@ def test_evaluate_base_model(self, config): agent_config=config["agent_core_arn"], s3_output_path=f'{config["s3_output_path"]}basemodel/', mlflow_resource_arn=config["mlflow_resource_arn"], - role=config["role"], region=_REGION, ) @@ -225,7 +220,6 @@ def test_evaluate_comparison(self, attached_trainer, config): dataset=config["dataset"], s3_output_path=f'{config["s3_output_path"]}comparison/', mlflow_resource_arn=config["mlflow_resource_arn"], - role=config["role"], region=_REGION, evaluate_base_model=True, ) diff --git a/sagemaker-train/tests/integ/train/test_multi_turn_rl_trainer_integration.py b/sagemaker-train/tests/integ/train/test_multi_turn_rl_trainer_integration.py index 71d41128e7..be47203dcb 100644 --- a/sagemaker-train/tests/integ/train/test_multi_turn_rl_trainer_integration.py +++ b/sagemaker-train/tests/integ/train/test_multi_turn_rl_trainer_integration.py @@ -56,7 +56,6 @@ def test_resources(): """Resolve account-specific resource ARNs lazily.""" account_id = _get_account_id() return { - "role_arn": f"arn:aws:iam::{account_id}:role/Admin", "mlflow_arn": f"arn:aws:sagemaker:{_REGION}:{account_id}:mlflow-app/app-TTAUWUNMUHH6", "s3_input_path": f"s3://sagemaker-rft-{account_id}/prompts/gsm8k_small/prompts.parquet", "s3_output_path": f"s3://sagemaker-{_REGION}-{account_id}/model-evaluation/mtrl-trainer-integ/", @@ -77,7 +76,6 @@ def test_train_and_wait(self, sagemaker_session, test_resources): training_dataset=test_resources["s3_input_path"], mlflow_app_arn=test_resources["mlflow_arn"], s3_output_path=test_resources["s3_output_path"], - role=test_resources["role_arn"], accept_eula=True, sagemaker_session=sagemaker_session, ) @@ -101,7 +99,6 @@ def test_train_and_stop(self, sagemaker_session, test_resources): agent_env=AGENT_RUNTIME_ID, training_dataset=test_resources["s3_input_path"], mlflow_app_arn=test_resources["mlflow_arn"], - role=test_resources["role_arn"], accept_eula=True, sagemaker_session=sagemaker_session, ) @@ -132,7 +129,6 @@ def test_train_with_lambda_arn(self, sagemaker_session, test_resources): mlflow_app_arn=test_resources["mlflow_arn"], s3_output_path=test_resources["s3_output_path"], accept_eula=True, - role=test_resources["role_arn"], sagemaker_session=sagemaker_session, ) trainer.hyperparameters.global_batch_size = 32 From e13958da6d4a40ffbfefd8406e4e0c22c7da39e4 Mon Sep 17 00:00:00 2001 From: Roja Reddy Sareddy Date: Wed, 22 Jul 2026 11:24:12 -0700 Subject: [PATCH 2/4] Fix role issue in mtrl integ tests --- .../tests/integ/train/test_mtrl_evaluator.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/sagemaker-train/tests/integ/train/test_mtrl_evaluator.py b/sagemaker-train/tests/integ/train/test_mtrl_evaluator.py index 12ea51b36a..9d493a0508 100644 --- a/sagemaker-train/tests/integ/train/test_mtrl_evaluator.py +++ b/sagemaker-train/tests/integ/train/test_mtrl_evaluator.py @@ -40,6 +40,10 @@ def _get_test_config(): """Build test configuration lazily (only when tests actually run).""" boto_session = boto3.Session(region_name=_REGION) account_id = boto_session.client("sts").get_caller_identity()["Account"] + from sagemaker.core.helper.session_helper import Session + from sagemaker.train.defaults import TrainDefaults + sagemaker_session = Session(boto_session=boto_session) + role_arn = TrainDefaults.get_role(role=None, sagemaker_session=sagemaker_session) return { "base_model": "mock-oss-test", "agent_arn": f"arn:aws:bedrock-agentcore:{_REGION}:{account_id}:runtime/sagemaker_rft_prod_gsm8k_streaming-Yk6O377mUS", @@ -47,6 +51,7 @@ def _get_test_config(): "s3_output_path": f"s3://sagemaker-{_REGION}-{account_id}/model-evaluation/output-artifacts/", "mlflow_resource_arn": f"arn:aws:sagemaker:{_REGION}:{account_id}:mlflow-app/app-TTAUWUNMUHH6", "model_package_group": f"arn:aws:sagemaker:{_REGION}:{account_id}:model-package-group/openai-reasoning-gpt-oss-20b-mtrl-mpg", + "role": role_arn, "region": _REGION, "account_id": account_id, } @@ -180,7 +185,7 @@ def test_bedrock_agent_config_fields(self, mtrl_trainer, test_config): evaluator._resolve_agent_arn() ctx = evaluator._build_template_context( - aws_context={"region": test_config["region"], "account_id": test_config["account_id"]}, + aws_context={"region": test_config["region"], "account_id": test_config["account_id"], "role_arn": test_config["role"]}, artifacts={}, model_package_group_arn=test_config["model_package_group"], ) @@ -213,7 +218,7 @@ def test_lambda_agent_config_fields(self, mtrl_trainer, test_config): evaluator._resolve_agent_arn() ctx = evaluator._build_template_context( - aws_context={"region": test_config["region"], "account_id": test_config["account_id"]}, + aws_context={"region": test_config["region"], "account_id": test_config["account_id"], "role_arn": test_config["role"]}, artifacts={}, model_package_group_arn=test_config["model_package_group"], ) @@ -242,7 +247,7 @@ def test_model_package_config_fields(self, mtrl_trainer, test_config): evaluator._resolve_agent_arn() ctx = evaluator._build_template_context( - aws_context={"region": test_config["region"], "account_id": test_config["account_id"]}, + aws_context={"region": test_config["region"], "account_id": test_config["account_id"], "role_arn": test_config["role"]}, artifacts={}, model_package_group_arn=test_config["model_package_group"], ) From 6147c72ee15f33893a669ca7d1ae3d9b2651f887 Mon Sep 17 00:00:00 2001 From: Roja Reddy Sareddy Date: Wed, 22 Jul 2026 16:17:16 -0700 Subject: [PATCH 3/4] Fix: Telemetry INFO logging to print only once per process --- .../core/telemetry/telemetry_logging.py | 7 ++- .../unit/telemetry/test_telemetry_logging.py | 54 +++++++++++++++++++ 2 files changed, 60 insertions(+), 1 deletion(-) diff --git a/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py b/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py index aa44bac0af..b8f5205e91 100644 --- a/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py +++ b/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py @@ -64,6 +64,7 @@ "For more information, refer to https://sagemaker.readthedocs.io/en/stable/overview.html" "#configuring-and-using-defaults-with-the-sagemaker-python-sdk." ) +_telemetry_msg_shown = False FEATURE_TO_CODE = { str(Feature.SDK_DEFAULTS): 11, @@ -284,8 +285,8 @@ def wrapper(*args, **kwargs): ) if sagemaker_session: + global _telemetry_msg_shown logger.debug("sagemaker_session found, preparing to emit telemetry...") - logger.info(TELEMETRY_OPT_OUT_MESSAGING) response = None caught_ex = None studio_app_type = process_studio_metadata_file() @@ -299,6 +300,10 @@ def wrapper(*args, **kwargs): ) logger.debug("TelemetryOptOut flag is set to: %s", telemetry_opt_out_flag) + if not telemetry_opt_out_flag and not _telemetry_msg_shown: + logger.info(TELEMETRY_OPT_OUT_MESSAGING) + _telemetry_msg_shown = True + # Construct the feature list to track feature combinations feature_list: List[int] = [FEATURE_TO_CODE[str(feature)]] diff --git a/sagemaker-core/tests/unit/telemetry/test_telemetry_logging.py b/sagemaker-core/tests/unit/telemetry/test_telemetry_logging.py index 04154e6ad7..6c7359ceb1 100644 --- a/sagemaker-core/tests/unit/telemetry/test_telemetry_logging.py +++ b/sagemaker-core/tests/unit/telemetry/test_telemetry_logging.py @@ -647,3 +647,57 @@ def train(self): self.assertIn(15, feature_list) # MODEL_CUSTOMIZATION still present self.assertNotIn(19, feature_list) # No NOVA (detection failed gracefully) self.assertNotIn(20, feature_list) # No OSS + + @patch("sagemaker.core.telemetry.telemetry_logging._send_telemetry_request") + @patch("sagemaker.core.telemetry.telemetry_logging.resolve_value_from_config") + def test_telemetry_opt_out_message_shown_only_once( + self, mock_resolve_config, mock_send_telemetry_request + ): + """Test that the telemetry opt-out INFO message is logged only once per process.""" + import sagemaker.core.telemetry.telemetry_logging as telemetry_module + + mock_resolve_config.return_value = False + # Reset the flag to simulate a fresh process + telemetry_module._telemetry_msg_shown = False + + mock_local_client = LocalSagemakerClientMock() + + with patch.object(telemetry_module.logger, "info") as mock_logger_info: + mock_local_client.mock_create_model() + mock_local_client.mock_create_model() + mock_local_client.mock_create_model() + + info_calls = [ + call for call in mock_logger_info.call_args_list + if "telemetry" in str(call).lower() and "opt out" in str(call).lower() + ] + self.assertEqual(len(info_calls), 1, "Telemetry opt-out message should be logged exactly once") + + # Reset the flag for other tests + telemetry_module._telemetry_msg_shown = False + + @patch("sagemaker.core.telemetry.telemetry_logging._send_telemetry_request") + @patch("sagemaker.core.telemetry.telemetry_logging.resolve_value_from_config") + def test_telemetry_opt_out_message_not_shown_when_opted_out( + self, mock_resolve_config, mock_send_telemetry_request + ): + """Test that the telemetry opt-out INFO message is not shown when user has opted out.""" + import sagemaker.core.telemetry.telemetry_logging as telemetry_module + + mock_resolve_config.return_value = True # opted out + # Reset the flag to simulate a fresh process + telemetry_module._telemetry_msg_shown = False + + mock_local_client = LocalSagemakerClientMock() + + with patch.object(telemetry_module.logger, "info") as mock_logger_info: + mock_local_client.mock_create_model() + + info_calls = [ + call for call in mock_logger_info.call_args_list + if "telemetry" in str(call).lower() and "opt out" in str(call).lower() + ] + self.assertEqual(len(info_calls), 0, "Telemetry opt-out message should not appear when opted out") + + # Reset the flag for other tests + telemetry_module._telemetry_msg_shown = False From d27108b1f12b37638a3888740ed2200eb4e80d90 Mon Sep 17 00:00:00 2001 From: Roja Reddy Sareddy Date: Thu, 23 Jul 2026 16:09:34 -0700 Subject: [PATCH 4/4] Fix: Telemetry INFO logging to print only once per process --- .../src/sagemaker/core/telemetry/telemetry_logging.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py b/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py index b8f5205e91..89bbfb6a28 100644 --- a/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py +++ b/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py @@ -301,7 +301,7 @@ def wrapper(*args, **kwargs): logger.debug("TelemetryOptOut flag is set to: %s", telemetry_opt_out_flag) if not telemetry_opt_out_flag and not _telemetry_msg_shown: - logger.info(TELEMETRY_OPT_OUT_MESSAGING) + logger.warning(TELEMETRY_OPT_OUT_MESSAGING) _telemetry_msg_shown = True # Construct the feature list to track feature combinations