Skip to content

Commit 29527dc

Browse files
authored
test(train): Stop MTRL training integ tests churning fixture MPG (#6362)
MTRL training integ tests registered their output into the SDK default group mock-oss-test-mtrl-mpg, which also holds the fixed fixture package that the MTRL evaluator integ tests attach to. The CI resource cleaner keeps only the oldest and newest versions in that group, so every new training run caused the fixture package (mpg/158) to be deleted, and test_evaluate_finetuned_model / test_evaluate_with_attached_trainer have failed with "ModelPackage ... does not exist" since ~2026-09-12. Send training output to a dedicated scratch group (mock-oss-test-mtrl-train-mpg, get-or-created by a fixture), and refresh the fixture job to mock-oss-test-mtrl-20260929124814 (mpg/219), which is pinned with the pysdk-ci-protected=true tag honored by the cleaner.
1 parent eba9a87 commit 29527dc

3 files changed

Lines changed: 54 additions & 6 deletions

File tree

‎sagemaker-train/tests/integ/train/test_mtrl_evaluator_3p_agent.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -316,8 +316,10 @@ def test_evaluate_with_attached_trainer(self, lambda_agent_arn, test_config):
316316
"""Test evaluating a fine-tuned model by attaching to an existing training job."""
317317
from sagemaker.train.multi_turn_rl_trainer import MultiTurnRLTrainer
318318

319+
# Fixture job whose output package is pinned via the `pysdk-ci-protected=true`
320+
# tag. Keep in sync with test_mtrl_trainer_integration.py.
319321
attached_job = MultiTurnRLTrainer.attach(
320-
"mock-oss-test-mtrl-20260910094327", session=boto3.Session(region_name=_REGION)
322+
"mock-oss-test-mtrl-20260929124814", session=boto3.Session(region_name=_REGION)
321323
)
322324

323325
evaluator = MultiTurnRLEvaluator(

‎sagemaker-train/tests/integ/train/test_mtrl_trainer_integration.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -53,8 +53,10 @@ def _get_account_id():
5353
# PROD — Main account (729646638167)
5454
"729646638167": {
5555
"env_name": "PROD",
56-
# "existing_job_name": "mock-oss-test-mtrl-20260611170946",
57-
"existing_job_name": "mock-oss-test-mtrl-20260910094327",
56+
# Fixture job; its output package (mock-oss-test-mtrl-mpg/219) is pinned via the
57+
# `pysdk-ci-protected=true` tag so the CI resource cleaner keeps it.
58+
# Keep in sync with test_mtrl_evaluator_3p_agent.py.
59+
"existing_job_name": "mock-oss-test-mtrl-20260929124814",
5860
"base_model": "mock-oss-test",
5961
"agent_core_arn": "arn:aws:bedrock-agentcore:us-west-2:729646638167:runtime/"
6062
"sagemaker_rft_prod_gsm8k_streaming-Yk6O377mUS",

‎sagemaker-train/tests/integ/train/test_multi_turn_rl_trainer_integration.py‎

Lines changed: 47 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import time
2222

2323
import boto3
24+
import botocore.exceptions
2425
import pytest
2526
from sagemaker.core.helper.session_helper import Session
2627
from sagemaker.train.multi_turn_rl_trainer import MultiTurnRLTrainer
@@ -44,6 +45,13 @@ def _get_account_id():
4445
BASE_MODEL = "mock-oss-test"
4546
EXISTING_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")
4957
def 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")
5691
def test_resources():
5792
"""Resolve account-specific resource ARNs lazily."""
@@ -69,12 +104,15 @@ def test_resources():
69104
class 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):
121162
class 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

Comments
 (0)