2121import time
2222
2323import boto3
24+ import botocore .exceptions
2425import pytest
2526from sagemaker .core .helper .session_helper import Session
2627from sagemaker .train .multi_turn_rl_trainer import MultiTurnRLTrainer
@@ -44,6 +45,13 @@ def _get_account_id():
4445BASE_MODEL = "mock-oss-test"
4546EXISTING_JOB_NAME = "mock-oss-test-mtrl-20260616153024"
4647
48+ # Training jobs in this module register their output into a dedicated scratch group.
49+ # Without it they fall back to the SDK default ``{BASE_MODEL}-mtrl-mpg``
50+ # (``mock-oss-test-mtrl-mpg``), which also holds the fixed fixture package that the
51+ # MTRL evaluator integ tests attach to. Every new version there made the resource
52+ # cleaner trim (delete) that fixture package.
53+ TRAIN_OUTPUT_MODEL_PACKAGE_GROUP = f"{ BASE_MODEL } -mtrl-train-mpg"
54+
4755
4856@pytest .fixture (scope = "module" )
4957def sagemaker_session ():
@@ -52,6 +60,33 @@ def sagemaker_session():
5260 yield session
5361
5462
63+ @pytest .fixture (scope = "module" )
64+ def train_output_model_package_group (sagemaker_session ):
65+ """Get-or-create the scratch output ModelPackageGroup and return its name.
66+
67+ ``MultiTurnRLTrainer`` only validates an explicitly passed group, so it must
68+ exist before the trainer is constructed. Race-safe for concurrent runs: losing
69+ the create race ("already exists") is treated as success.
70+ """
71+ sm_client = sagemaker_session .boto_session .client ("sagemaker" )
72+ try :
73+ sm_client .describe_model_package_group (
74+ ModelPackageGroupName = TRAIN_OUTPUT_MODEL_PACKAGE_GROUP
75+ )
76+ except botocore .exceptions .ClientError :
77+ try :
78+ sm_client .create_model_package_group (
79+ ModelPackageGroupName = TRAIN_OUTPUT_MODEL_PACKAGE_GROUP ,
80+ ModelPackageGroupDescription = (
81+ "Scratch output group for MTRL trainer integ tests (sagemaker-train)"
82+ ),
83+ )
84+ except botocore .exceptions .ClientError as e :
85+ if "already exists" not in str (e ):
86+ raise
87+ return TRAIN_OUTPUT_MODEL_PACKAGE_GROUP
88+
89+
5590@pytest .fixture (scope = "module" )
5691def test_resources ():
5792 """Resolve account-specific resource ARNs lazily."""
@@ -69,12 +104,15 @@ def test_resources():
69104class TestMultiTurnRLTrainerBedrockAgent :
70105 """Test MTRL training with Bedrock AgentCore runtime."""
71106
72- def test_train_and_wait (self , sagemaker_session , test_resources ):
107+ def test_train_and_wait (
108+ self , sagemaker_session , test_resources , train_output_model_package_group
109+ ):
73110 """Test complete MTRL workflow with Bedrock AgentCore agent."""
74111 trainer = MultiTurnRLTrainer (
75112 model = BASE_MODEL ,
76113 agent_env = AGENT_RUNTIME_ID ,
77114 training_dataset = test_resources ["s3_input_path" ],
115+ output_model_package_group = train_output_model_package_group ,
78116 mlflow_app_arn = test_resources ["mlflow_arn" ],
79117 s3_output_path = test_resources ["s3_output_path" ],
80118 accept_eula = True ,
@@ -93,12 +131,15 @@ def test_train_and_wait(self, sagemaker_session, test_resources):
93131 assert job .output_model_package_arn is not None
94132 assert job .s3_output_path is not None
95133
96- def test_train_and_stop (self , sagemaker_session , test_resources ):
134+ def test_train_and_stop (
135+ self , sagemaker_session , test_resources , train_output_model_package_group
136+ ):
97137 """Test creating and stopping an MTRL job."""
98138 trainer = MultiTurnRLTrainer (
99139 model = BASE_MODEL ,
100140 agent_env = AGENT_RUNTIME_ID ,
101141 training_dataset = test_resources ["s3_input_path" ],
142+ output_model_package_group = train_output_model_package_group ,
102143 mlflow_app_arn = test_resources ["mlflow_arn" ],
103144 accept_eula = True ,
104145 sagemaker_session = sagemaker_session ,
@@ -121,12 +162,15 @@ def test_train_and_stop(self, sagemaker_session, test_resources):
121162class TestMultiTurnRLTrainerLambdaAgent :
122163 """Test MTRL training with Lambda agent."""
123164
124- def test_train_with_lambda_arn (self , sagemaker_session , test_resources ):
165+ def test_train_with_lambda_arn (
166+ self , sagemaker_session , test_resources , train_output_model_package_group
167+ ):
125168 """Test MTRL workflow using an existing Lambda ARN as agent."""
126169 trainer = MultiTurnRLTrainer (
127170 model = BASE_MODEL ,
128171 agent_env = test_resources ["lambda_arn" ],
129172 training_dataset = test_resources ["s3_input_path" ],
173+ output_model_package_group = train_output_model_package_group ,
130174 mlflow_app_arn = test_resources ["mlflow_arn" ],
131175 s3_output_path = test_resources ["s3_output_path" ],
132176 accept_eula = True ,
0 commit comments