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 dcc624cfde..ebb8de2474 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 @@ -316,8 +316,10 @@ def test_evaluate_with_attached_trainer(self, lambda_agent_arn, test_config): """Test evaluating a fine-tuned model by attaching to an existing training job.""" from sagemaker.train.multi_turn_rl_trainer import MultiTurnRLTrainer + # Fixture job whose output package is pinned via the `pysdk-ci-protected=true` + # tag. Keep in sync with test_mtrl_trainer_integration.py. attached_job = MultiTurnRLTrainer.attach( - "mock-oss-test-mtrl-20260910094327", session=boto3.Session(region_name=_REGION) + "mock-oss-test-mtrl-20260929124814", session=boto3.Session(region_name=_REGION) ) evaluator = MultiTurnRLEvaluator( 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 6617f4c123..0ddb73bb31 100644 --- a/sagemaker-train/tests/integ/train/test_mtrl_trainer_integration.py +++ b/sagemaker-train/tests/integ/train/test_mtrl_trainer_integration.py @@ -53,8 +53,10 @@ def _get_account_id(): # PROD — Main account (729646638167) "729646638167": { "env_name": "PROD", - # "existing_job_name": "mock-oss-test-mtrl-20260611170946", - "existing_job_name": "mock-oss-test-mtrl-20260910094327", + # Fixture job; its output package (mock-oss-test-mtrl-mpg/219) is pinned via the + # `pysdk-ci-protected=true` tag so the CI resource cleaner keeps it. + # Keep in sync with test_mtrl_evaluator_3p_agent.py. + "existing_job_name": "mock-oss-test-mtrl-20260929124814", "base_model": "mock-oss-test", "agent_core_arn": "arn:aws:bedrock-agentcore:us-west-2:729646638167:runtime/" "sagemaker_rft_prod_gsm8k_streaming-Yk6O377mUS", 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 cf65c049bc..b10510b0db 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 @@ -21,6 +21,7 @@ import time import boto3 +import botocore.exceptions import pytest from sagemaker.core.helper.session_helper import Session from sagemaker.train.multi_turn_rl_trainer import MultiTurnRLTrainer @@ -44,6 +45,13 @@ def _get_account_id(): BASE_MODEL = "mock-oss-test" EXISTING_JOB_NAME = "mock-oss-test-mtrl-20260616153024" +# Training jobs in this module register their output into a dedicated scratch group. +# Without it they fall back to the SDK default ``{BASE_MODEL}-mtrl-mpg`` +# (``mock-oss-test-mtrl-mpg``), which also holds the fixed fixture package that the +# MTRL evaluator integ tests attach to. Every new version there made the resource +# cleaner trim (delete) that fixture package. +TRAIN_OUTPUT_MODEL_PACKAGE_GROUP = f"{BASE_MODEL}-mtrl-train-mpg" + @pytest.fixture(scope="module") def sagemaker_session(): @@ -52,6 +60,33 @@ def sagemaker_session(): yield session +@pytest.fixture(scope="module") +def train_output_model_package_group(sagemaker_session): + """Get-or-create the scratch output ModelPackageGroup and return its name. + + ``MultiTurnRLTrainer`` only validates an explicitly passed group, so it must + exist before the trainer is constructed. Race-safe for concurrent runs: losing + the create race ("already exists") is treated as success. + """ + sm_client = sagemaker_session.boto_session.client("sagemaker") + try: + sm_client.describe_model_package_group( + ModelPackageGroupName=TRAIN_OUTPUT_MODEL_PACKAGE_GROUP + ) + except botocore.exceptions.ClientError: + try: + sm_client.create_model_package_group( + ModelPackageGroupName=TRAIN_OUTPUT_MODEL_PACKAGE_GROUP, + ModelPackageGroupDescription=( + "Scratch output group for MTRL trainer integ tests (sagemaker-train)" + ), + ) + except botocore.exceptions.ClientError as e: + if "already exists" not in str(e): + raise + return TRAIN_OUTPUT_MODEL_PACKAGE_GROUP + + @pytest.fixture(scope="module") def test_resources(): """Resolve account-specific resource ARNs lazily.""" @@ -69,12 +104,15 @@ def test_resources(): class TestMultiTurnRLTrainerBedrockAgent: """Test MTRL training with Bedrock AgentCore runtime.""" - def test_train_and_wait(self, sagemaker_session, test_resources): + def test_train_and_wait( + self, sagemaker_session, test_resources, train_output_model_package_group + ): """Test complete MTRL workflow with Bedrock AgentCore agent.""" trainer = MultiTurnRLTrainer( model=BASE_MODEL, agent_env=AGENT_RUNTIME_ID, training_dataset=test_resources["s3_input_path"], + output_model_package_group=train_output_model_package_group, mlflow_app_arn=test_resources["mlflow_arn"], s3_output_path=test_resources["s3_output_path"], accept_eula=True, @@ -93,12 +131,15 @@ def test_train_and_wait(self, sagemaker_session, test_resources): assert job.output_model_package_arn is not None assert job.s3_output_path is not None - def test_train_and_stop(self, sagemaker_session, test_resources): + def test_train_and_stop( + self, sagemaker_session, test_resources, train_output_model_package_group + ): """Test creating and stopping an MTRL job.""" trainer = MultiTurnRLTrainer( model=BASE_MODEL, agent_env=AGENT_RUNTIME_ID, training_dataset=test_resources["s3_input_path"], + output_model_package_group=train_output_model_package_group, mlflow_app_arn=test_resources["mlflow_arn"], accept_eula=True, sagemaker_session=sagemaker_session, @@ -121,12 +162,15 @@ def test_train_and_stop(self, sagemaker_session, test_resources): class TestMultiTurnRLTrainerLambdaAgent: """Test MTRL training with Lambda agent.""" - def test_train_with_lambda_arn(self, sagemaker_session, test_resources): + def test_train_with_lambda_arn( + self, sagemaker_session, test_resources, train_output_model_package_group + ): """Test MTRL workflow using an existing Lambda ARN as agent.""" trainer = MultiTurnRLTrainer( model=BASE_MODEL, agent_env=test_resources["lambda_arn"], training_dataset=test_resources["s3_input_path"], + output_model_package_group=train_output_model_package_group, mlflow_app_arn=test_resources["mlflow_arn"], s3_output_path=test_resources["s3_output_path"], accept_eula=True,