fix:unskip and add mock test setup for mtrl integ tests (#5937)

* fix:unskip and mark MTRL tests as gpu intensive

* Add Mock setup for mtrl integ tests

* Add Mock setup for mtrl integ tests

* Add Mock setup for mtrl integ tests

---------

Co-authored-by: Roja Reddy Sareddy <rsareddy@amazon.com>

Author: rsareddy0329

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

@@ -22,6 +22,13 @@
 DEFAULT_REGION = "us-west-2"
 
 
+@pytest.fixture(autouse=True, scope="session")
+def use_private_hub():
+    os.environ["SAGEMAKER_HUB_NAME"] = "sdktest"
+    yield
+    del os.environ["SAGEMAKER_HUB_NAME"]
+
+
 @pytest.fixture(scope="module")
 def sagemaker_session():
     region = os.environ.get("AWS_DEFAULT_REGION")

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

@@ -41,7 +41,7 @@ def _get_test_config():
     boto_session = boto3.Session(region_name=_REGION)
     account_id = boto_session.client("sts").get_caller_identity()["Account"]
     return {
-        "base_model": "openai-reasoning-gpt-oss-20b",
+        "base_model": "mock-oss-test",
         "agent_arn": f"arn:aws:bedrock-agentcore:{_REGION}:{account_id}:runtime/sagemaker_rft_prod_gsm8k_streaming-Yk6O377mUS",
         "dataset": f"s3://sagemaker-rft-{account_id}/prompts/gsm8k_small/prompts.parquet",
         "s3_output_path": f"s3://sagemaker-{_REGION}-{account_id}/model-evaluation/output-artifacts/",
@@ -142,7 +142,7 @@ def mtrl_trainer(sagemaker_session_mtrl, test_config):
 
     trainer = object.__new__(MultiTurnRLTrainer)
     trainer._model_name = test_config["base_model"]
-    trainer._model_arn = f"arn:aws:sagemaker:{_REGION}:aws:hub-content/SageMakerPublicHub/Model/{test_config['base_model']}/1.0.0"
+    trainer._model_arn = f"arn:aws:sagemaker:{_REGION}:{test_config['account_id']}:hub-content/sdktest/Model/{test_config['base_model']}/0.0.1"
     trainer.agent_env = test_config["agent_arn"]
     trainer.bedrock_agentcore_qualifier = "DEFAULT"
     trainer.output_model_package_group = test_config["model_package_group"]

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

@@ -148,7 +148,7 @@ def _get_3p_test_config():
     boto_session = boto3.Session(region_name=_REGION)
     account_id = boto_session.client("sts").get_caller_identity()["Account"]
     return {
-        "base_model": "openai-reasoning-gpt-oss-20b",
+        "base_model": "mock-oss-test",
         "dataset": os.environ.get(
             "MTRL_3P_DATASET",
             f"s3://sagemaker-rft-{account_id}/prompts/gsm8k_small/prompts.parquet",
@@ -240,21 +240,23 @@ def lambda_agent_arn(test_config):
     return _ensure_lambda_exists(test_config["account_id"])
 
 
+@pytest.mark.gpu_intensive
+@pytest.mark.serial
 class TestMTRLEvaluator3PAgentIntegration:
     """Integration tests for MultiTurnRLEvaluator with Lambda-based 3P agent."""
 
-    def test_evaluate_base_model_with_lambda_agent(self, lambda_agent_arn, test_config):
-        """Test evaluating a base model using a Lambda ARN as agent_config.
+    def test_evaluate_with_lambda_agent_wait_for_completion(self, lambda_agent_arn, test_config):
+        """Test full end-to-end: start evaluation, wait for completion, and verify discoverability.
 
-        This is the primary 3P integration pattern: customer provides a
-        Lambda function that wraps their agent (LangChain, Strands, etc.)
-        and the evaluator runs rollouts against it.
+        This test validates the complete lifecycle including wait() using
+        the standard sagemaker-core pipeline execution path, and verifies
+        the evaluation is discoverable via get_all().
         """
         evaluator = MultiTurnRLEvaluator(
             model=test_config["base_model"],
             dataset=test_config["dataset"],
             agent_config=lambda_agent_arn,
-            s3_output_path=f'{test_config["s3_output_path"]}lambda-base-model/',
+            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"],
@@ -267,9 +269,27 @@ def test_evaluate_base_model_with_lambda_agent(self, lambda_agent_arn, test_conf
         assert execution.arn is not None
         assert "pipeline" in execution.arn.lower()
         logger.info(f"Started 3P agent base model evaluation: {execution.arn}")
-        logger.info(f"Status: {execution.status.overall_status}")
 
-    @pytest.mark.skip(reason="Quota limited (1 concurrent eval job) - run manually")
+        execution.wait(timeout=EVALUATION_TIMEOUT_SECONDS)
+        assert execution.status.overall_status in ("Succeeded", "Failed", "Stopped")
+        logger.info(f"Execution completed: {execution.status.overall_status}")
+
+        if execution.status.overall_status == "Failed":
+            logger.error(f"Failure reason: {execution.status.failure_reason}")
+
+        # Verify it's discoverable via get_all
+        found = False
+        for ex in MultiTurnRLEvaluator.get_all(region=test_config["region"]):
+            if ex.arn == execution.arn:
+                found = True
+                break
+
+        assert found, (
+            f"Evaluation {execution.arn} not found via get_all(). "
+            "Pipeline tagging may not be working correctly."
+        )
+        logger.info(f"Successfully discovered evaluation via get_all: {execution.arn}")
+
     def test_evaluate_base_model_with_agent_lambda_object(self, lambda_agent_arn, test_config):
         """Test evaluating using an CustomAgentLambda object as agent_config.
 
@@ -295,83 +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}")
 
-    @pytest.mark.skip(reason="Quota limited (1 concurrent eval job) - run manually")
-    def test_evaluate_with_lambda_agent_wait_for_completion(self, lambda_agent_arn, test_config):
-        """Test full end-to-end: start evaluation and wait for completion.
-
-        This test validates the complete lifecycle including wait() using
-        the standard sagemaker-core pipeline execution path.
-        """
-        evaluator = MultiTurnRLEvaluator(
-            model=test_config["base_model"],
-            dataset=test_config["dataset"],
-            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,
-        )
-
-        execution = evaluator.evaluate()
-        assert execution is not None
-
-        logger.info(f"Waiting for execution: {execution.arn}")
-        execution.wait()
-
-        assert execution.status.overall_status in ("Succeeded", "Failed", "Stopped")
-        logger.info(f"Execution completed: {execution.status.overall_status}")
-
-        if execution.status.overall_status == "Failed":
-            logger.error(f"Failure reason: {execution.status.failure_reason}")
-
-    @pytest.mark.skip(reason="Quota limited (1 concurrent eval job) - run manually")
-    def test_evaluate_lambda_agent_discoverable_via_get_all(self, lambda_agent_arn, test_config):
-        """Test that 3P agent evaluations are discoverable via get_all.
-
-        Validates that evaluations started with Lambda agents show up in
-        the standard get_all() discovery path (pipeline tagging works).
-        """
-        evaluator = MultiTurnRLEvaluator(
-            model=test_config["base_model"],
-            dataset=test_config["dataset"],
-            agent_config=lambda_agent_arn,
-            s3_output_path=f'{test_config["s3_output_path"]}lambda-discovery/',
-            mlflow_resource_arn=test_config["mlflow_resource_arn"],
-            role=test_config["role"],
-            region=test_config["region"],
-            accept_eula=True,
-        )
-
-        execution = evaluator.evaluate()
-        assert execution is not None
-        started_arn = execution.arn
-
-        # Give pipeline time to register
-        time.sleep(10)
-
-        # Verify it's discoverable via get_all
-        found = False
-        for ex in MultiTurnRLEvaluator.get_all(region=test_config["region"]):
-            if ex.arn == started_arn:
-                found = True
-                break
-
-        assert found, (
-            f"Evaluation {started_arn} not found via get_all(). "
-            "Pipeline tagging may not be working correctly."
-        )
-        logger.info(f"Successfully discovered evaluation via get_all: {started_arn}")
-
+        execution.wait(timeout=EVALUATION_TIMEOUT_SECONDS)
+        assert execution.status.overall_status == "Succeeded"
 
-
-    @pytest.mark.skip(reason="Quota limited (1 concurrent eval job) - run manually")
     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(
@@ -390,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(timeout=EVALUATION_TIMEOUT_SECONDS)
+        assert execution.status.overall_status == "Succeeded"

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

@@ -52,8 +52,9 @@ def _get_account_id():
     # PROD — Main account (729646638167)
     "729646638167": {
         "env_name": "PROD",
-        "existing_job_name": "openai-reasoning-gpt-oss-20b-mtrl-20260602215955",
-        "base_model": "openai-reasoning-gpt-oss-20b",
+        #"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",
         "s3_output_path": "s3://sagemaker-us-west-2-729646638167/mtrl-integ/eval-output/",
@@ -65,7 +66,7 @@ def _get_account_id():
     "391266019386": {
         "env_name": "PREPROD",
         "existing_job_name": "mtrl-integ-gpt-oss-agentcore-1779143704358",
-        "base_model": "openai-reasoning-gpt-oss-20b",
+        "base_model": "mock-oss-test",
         "agent_core_arn": "arn:aws:bedrock-agentcore:us-west-2:391266019386:runtime/mtrl_integ_gsm8k_streaming-bIz4H5Echk",
         "dataset": "s3://sagemaker-rft-beta-391266019386/prompts/gsm8k_small/prompts.parquet",
         "s3_output_path": "s3://sagemaker-us-west-2-391266019386/mtrl-integ/eval-output/",
@@ -77,7 +78,7 @@ def _get_account_id():
     "742774200982": {
         "env_name": "BETA",
         "existing_job_name": "openai-reasoning-gpt-oss-20b-mtrl-20260601114439",
-        "base_model": "openai-reasoning-gpt-oss-20b",
+        "base_model": "mock-oss-test",
         "agent_core_arn": "arn:aws:bedrock-agentcore:us-west-2:742774200982:runtime/sagemaker_rft_prod_gsm8k_streaming-UwSB6LEfEq",
         "dataset": "s3://sagemaker-rft-beta-742774200982/prompts/gsm8k_small/prompts.parquet",
         "s3_output_path": "s3://sagemaker-us-west-2-742774200982/mtrl-integ/eval-output/",
@@ -139,6 +140,8 @@ def attached_trainer(config):
     return trainer
 
 
+@pytest.mark.gpu_intensive
+@pytest.mark.serial
 class TestMTRLEvalIntegration:
     """Integration tests for MTRL evaluation: attach → evaluate → wait for success."""
 
@@ -184,7 +187,6 @@ def test_evaluate_finetuned_model(self, attached_trainer, config):
             f"reason: {execution.status.failure_reason}"
         )
 
-    @pytest.mark.skip(reason="Quota limited (1 concurrent eval job) - run manually")
     def test_evaluate_base_model(self, config):
         """Evaluate the base model only — submit and wait for completion."""
         evaluator = MultiTurnRLEvaluator(
@@ -245,35 +247,3 @@ def test_evaluate_comparison(self, attached_trainer, config):
             f"[{config['env_name']}] Comparison eval failed with status: {status}, "
             f"reason: {execution.status.failure_reason}"
         )
-
-    @pytest.mark.skip(reason="Quota limited (1 concurrent eval job) - run manually")
-    def test_evaluate_with_hyperparam_override(self, attached_trainer, config):
-        """Test that hyperparameter overrides are passed through to the eval job."""
-        evaluator = MultiTurnRLEvaluator(
-            model=attached_trainer,
-            dataset=config["dataset"],
-            s3_output_path=f'{config["s3_output_path"]}hyperparam-override/',
-            mlflow_resource_arn=config["mlflow_resource_arn"],
-            role=config["role"],
-            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
-        assert execution.arn is not None
-        logger.info(f"[{config['env_name']}] Started hyperparam override eval: {execution.arn}")
-
-        execution.wait(timeout=EVAL_TIMEOUT)
-
-        status = execution.status.overall_status
-        logger.info(f"[{config['env_name']}] Hyperparam override eval completed: {status}")
-
-        assert status == "Succeeded", (
-            f"[{config['env_name']}] Hyperparam override eval failed with status: {status}, "
-            f"reason: {execution.status.failure_reason}"
-        )

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

@@ -39,8 +39,9 @@ def _get_account_id():
     return _ACCOUNT_ID
 
 AGENT_RUNTIME_ID = "sagemaker_rft_prod_gsm8k_streaming-Yk6O377mUS"
-BASE_MODEL = "openai-reasoning-gpt-oss-20b"
-EXISTING_JOB_NAME = "openai-reasoning-gpt-oss-20b-mtrl-20260602005937"
+#BASE_MODEL = "openai-reasoning-gpt-oss-20b"
+BASE_MODEL = "mock-oss-test"
+EXISTING_JOB_NAME = "openai-reasoning-gpt-oss-20b-mtrl-20260602215955"
 
 
 @pytest.fixture(scope="module")
@@ -63,7 +64,8 @@ def test_resources():
     }
 
 
-@pytest.mark.skip(reason="GPU resource intensive — run manually")
+@pytest.mark.gpu_intensive
+@pytest.mark.serial
 class TestMultiTurnRLTrainerBedrockAgent:
     """Test MTRL training with Bedrock AgentCore runtime."""
 
@@ -116,7 +118,8 @@ def test_train_and_stop(self, sagemaker_session, test_resources):
         assert job.job_status in ("Stopping", "Stopped")
 
 
-@pytest.mark.skip(reason="GPU resource intensive — run manually")
+@pytest.mark.gpu_intensive
+@pytest.mark.serial
 class TestMultiTurnRLTrainerLambdaAgent:
     """Test MTRL training with Lambda agent."""
 
@@ -145,7 +148,7 @@ def test_train_with_lambda_arn(self, sagemaker_session, test_resources):
         assert job.output_model_package_arn is not None
 
 
-@pytest.mark.skip(reason="GPU resource intensive — run manually")
+
 class TestMultiTurnRLTrainerAttach:
     """Test attaching to existing MTRL jobs."""
 
@@ -162,6 +165,7 @@ def test_attach_and_get_properties(self, sagemaker_session):
         assert attached_job.output_model_package_arn is not None
         assert attached_job.s3_output_path is not None
 
+    @pytest.mark.skip(reason="GPU resource intensive — run manually")
     def test_get_all_jobs(self, sagemaker_session):
         """Test listing all MTRL jobs."""
         jobs = list(AgentRFTJob.get_all(
@@ -172,7 +176,6 @@ def test_get_all_jobs(self, sagemaker_session):
         assert all(j.job_status == "Completed" for j in jobs)
 
 
-@pytest.mark.skip(reason="GPU resource intensive — run manually")
 class TestMultiTurnRLTrainerListModels:
     """Test listing supported models (requires API access)."""
 
@@ -190,3 +193,6 @@ def test_list_bedrock_agentcore_runtimes(self, sagemaker_session):
             session=sagemaker_session.boto_session
         )
         assert isinstance(runtimes, list)
+
+
+