fix(test): add mlflow_resource_arn fixture that auto-discovers or creates app

Replace hard-coded MLflow app ARN with a conftest fixture that finds an
existing ready app or creates a temporary one (cleaned up after tests).
Prevents failures when the hard-coded app is deleted or quota is full.

X-AI-Prompt: add self-healing mlflow fixture for llm_as_judge integ tests
X-AI-Tool: kiro-cli

Author: lucasjia-aws

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

@@ -48,3 +48,92 @@ def sagemaker_session_us_east_1():
     """Create a SageMaker session in us-east-1 for Nova model tests."""
     boto_session = boto3.Session(region_name=NOVA_REGION)
     return Session(boto_session=boto_session)
+
+
+import time
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+@pytest.fixture(scope="module")
+def mlflow_resource_arn():
+    """Discover or create an MLflow app for integ tests, clean up if created.
+
+    Looks for an existing MLflow app in Created/Updated state. If none exists,
+    creates one and deletes it after the test module finishes.
+    """
+    region = os.environ.get("AWS_DEFAULT_REGION", DEFAULT_REGION)
+    sm_client = boto3.client("sagemaker", region_name=region)
+    created_arn = None
+
+    # Try to find an existing ready app
+    try:
+        paginator = sm_client.get_paginator("list_mlflow_apps")
+        for page in paginator.paginate():
+            for app in page.get("MlflowApps", []):
+                if app.get("Status") in ("Created", "Updated"):
+                    logger.info(f"Using existing MLflow app: {app['Arn']}")
+                    yield app["Arn"]
+                    return
+    except Exception as e:
+        logger.warning(f"Failed to list MLflow apps: {e}")
+
+    # No ready app found — create one
+    logger.info("No ready MLflow app found. Creating one for integ tests...")
+    sts_client = boto3.client("sts", region_name=region)
+    account_id = sts_client.get_caller_identity()["Account"]
+    app_name = f"integ-test-mlflow-{int(time.time())}"
+    artifact_store_uri = f"s3://sagemaker-{region}-{account_id}/mlflow-artifacts"
+
+    # Ensure bucket/prefix exists
+    s3_client = boto3.client("s3", region_name=region)
+    bucket_name = f"sagemaker-{region}-{account_id}"
+    try:
+        s3_client.head_bucket(Bucket=bucket_name)
+    except Exception:
+        if region == "us-east-1":
+            s3_client.create_bucket(Bucket=bucket_name)
+        else:
+            s3_client.create_bucket(
+                Bucket=bucket_name,
+                CreateBucketConfiguration={"LocationConstraint": region},
+            )
+    try:
+        s3_client.put_object(Bucket=bucket_name, Key="mlflow-artifacts/")
+    except Exception:
+        pass
+
+    # Get execution role
+    from sagemaker.train.defaults import TrainDefaults
+    boto_session = boto3.Session(region_name=region)
+    sagemaker_session = Session(boto_session=boto_session)
+    role_arn = TrainDefaults.get_role(role=None, sagemaker_session=sagemaker_session)
+
+    resp = sm_client.create_mlflow_app(
+        Name=app_name,
+        ArtifactStoreUri=artifact_store_uri,
+        RoleArn=role_arn,
+        AccountDefaultStatus="DISABLED",
+    )
+    created_arn = resp["Arn"]
+    logger.info(f"Created MLflow app: {created_arn}")
+
+    # Wait for it to become ready
+    for _ in range(60):
+        desc = sm_client.describe_mlflow_app(Arn=created_arn)
+        status = desc.get("Status")
+        if status in ("Created", "Updated"):
+            break
+        if status in ("Failed", "CreateFailed", "DeleteFailed"):
+            pytest.skip(f"MLflow app creation failed: {desc.get('FailureReason')}")
+        time.sleep(10)
+
+    yield created_arn
+
+    # Cleanup
+    logger.info(f"Cleaning up MLflow app: {created_arn}")
+    try:
+        sm_client.delete_mlflow_app(Arn=created_arn)
+    except Exception as e:
+        logger.warning(f"Failed to delete MLflow app {created_arn}: {e}")

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

@@ -81,7 +81,7 @@
 class TestLLMAsJudgeBaseModelFix:
     """Integration test for base model fix in LLMAsJudgeEvaluator"""
 
-    def test_base_model_evaluation_uses_correct_weights(self):
+    def test_base_model_evaluation_uses_correct_weights(self, mlflow_resource_arn):
         """
         Test that base model evaluation uses original base model weights.
         
@@ -115,6 +115,7 @@ def test_base_model_evaluation_uses_correct_weights(self):
             custom_metrics=TEST_CONFIG["custom_metrics_json"],
             s3_output_path=TEST_CONFIG["s3_output_path"],
             evaluate_base_model=TEST_CONFIG["evaluate_base_model"],
+            mlflow_resource_arn=mlflow_resource_arn,
         )
 
         # Verify evaluator configuration
@@ -250,7 +251,7 @@ def test_base_model_evaluation_uses_correct_weights(self):
             # Re-raise to fail the test
             raise
 
-    def test_base_model_false_still_works(self):
+    def test_base_model_false_still_works(self, mlflow_resource_arn):
         """
         Test that evaluate_base_model=False still works correctly (backward compatibility).
         
@@ -271,6 +272,7 @@ def test_base_model_false_still_works(self):
             builtin_metrics=TEST_CONFIG["builtin_metrics"],
             s3_output_path=TEST_CONFIG["s3_output_path"],
             evaluate_base_model=False,  # Only evaluate custom model
+            mlflow_resource_arn=mlflow_resource_arn,
         )
 
         # Verify evaluator configuration