diff --git a/sagemaker-train/tests/integ/train/test_mtrl_evaluator.py b/sagemaker-train/tests/integ/train/test_mtrl_evaluator.py index b79a2763dd..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,7 +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": f"arn:aws:iam::{account_id}:role/Admin", + "role": role_arn, "region": _REGION, "account_id": account_id, } @@ -172,7 +176,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 +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"], - "role_arn": test_config["role"]}, + 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"], ) @@ -208,7 +210,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 +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"], - "role_arn": test_config["role"]}, + 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"], ) @@ -239,7 +239,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 +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"], - "role_arn": test_config["role"]}, + 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"], ) @@ -280,7 +278,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 +295,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