fix:unskip and mark MTRL tests as gpu intensive

Author: Roja Reddy Sareddy

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

@@ -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()
+        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,77 +315,6 @@ 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}")
-
-
-
-    @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

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

@@ -139,6 +139,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."""
 
@@ -156,7 +158,10 @@ 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."""
+        """Evaluate a fine-tuned model from attached trainer — submit and wait for completion.
+
+        Also validates hyperparameter overrides are passed through to the eval job.
+        """
         evaluator = MultiTurnRLEvaluator(
             model=attached_trainer,
             dataset=config["dataset"],
@@ -166,6 +171,10 @@ 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
@@ -184,7 +193,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 +253,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

@@ -40,7 +40,7 @@ def _get_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"
+EXISTING_JOB_NAME = "openai-reasoning-gpt-oss-20b-mtrl-20260602215955"
 
 
 @pytest.fixture(scope="module")
@@ -63,7 +63,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 +117,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 +147,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 +164,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 +175,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)."""