Add Mock setup for mtrl integ tests

Author: Roja Reddy Sareddy

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

@@ -270,7 +270,7 @@ def test_evaluate_with_lambda_agent_wait_for_completion(self, lambda_agent_arn,
         assert "pipeline" in execution.arn.lower()
         logger.info(f"Started 3P agent base model evaluation: {execution.arn}")
 
-        execution.wait()
+        execution.wait(EVALUATION_TIMEOUT_SECONDS)
         assert execution.status.overall_status in ("Succeeded", "Failed", "Stopped")
         logger.info(f"Execution completed: {execution.status.overall_status}")
 
@@ -315,12 +315,15 @@ def test_evaluate_base_model_with_agent_lambda_object(self, lambda_agent_arn, te
         assert execution.arn is not None
         logger.info(f"Started CustomAgentLambda object evaluation: {execution.arn}")
 
+        execution.wait(EVALUATION_TIMEOUT_SECONDS)
+        assert execution.status.overall_status == "Succeeded"
+
     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
 
         attached_job = MultiTurnRLTrainer.attach(
-            "openai-reasoning-gpt-oss-20b-mtrl-20260602164546", session=boto3.Session(region_name=_REGION)
+            "mock-oss-test-mtrl-20260615143910", session=boto3.Session(region_name=_REGION)
         )
 
         evaluator = MultiTurnRLEvaluator(
@@ -339,3 +342,6 @@ def test_evaluate_with_attached_trainer(self, lambda_agent_arn, test_config):
         assert execution is not None
         assert execution.arn is not None
         logger.info(f"Started attached trainer evaluation: {execution.arn}")
+
+        execution.wait(EVALUATION_TIMEOUT_SECONDS)
+        assert execution.status.overall_status == "Succeeded"

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

@@ -52,7 +52,8 @@ def _get_account_id():
     # PROD — Main account (729646638167)
     "729646638167": {
         "env_name": "PROD",
-        "existing_job_name": "openai-reasoning-gpt-oss-20b-mtrl-20260602215955",
+        #"existing_job_name": "mock-oss-test-mtrl-20260611170946",
+        "existing_job_name": "mock-oss-test-mtrl-20260615143910",
         "base_model": "mock-oss-test",
         "agent_core_arn": "arn:aws:bedrock-agentcore:us-west-2:729646638167:runtime/sagemaker_rft_prod_gsm8k_streaming-Yk6O377mUS",
         "dataset": "s3://sagemaker-rft-729646638167/prompts/gsm8k_small/prompts.parquet",
@@ -158,10 +159,7 @@ def test_attach_to_existing_job(self, config):
         logger.info(f"[{config['env_name']}] Output model package: {job.output_model_package_arn}")
 
     def test_evaluate_finetuned_model(self, attached_trainer, config):
-        """Evaluate a fine-tuned model from attached trainer — submit and wait for completion.
-
-        Also validates hyperparameter overrides are passed through to the eval job.
-        """
+        """Evaluate a fine-tuned model from attached trainer — submit and wait for completion."""
         evaluator = MultiTurnRLEvaluator(
             model=attached_trainer,
             dataset=config["dataset"],
@@ -171,10 +169,6 @@ def test_evaluate_finetuned_model(self, attached_trainer, config):
             region=_REGION,
         )
 
-        # Override MTRL-specific hyperparams
-        evaluator.hyperparameters.sampling_max_tokens = 1024
-        evaluator.hyperparameters.eval_group_size = 4
-
         execution = evaluator.evaluate()
 
         assert execution is not None