fix(train): use latest model package in LLM-as-judge base model integ tests

The hardcoded model package version 1 predates the backend's
SageMakerPublicHub requirement for serverless training jobs, causing
consistent failures across all PRs.

Author: mollyheamazon

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

@@ -18,6 +18,7 @@
 """
 from __future__ import absolute_import
 
+import boto3
 import json
 import time
 import pytest
@@ -64,51 +65,78 @@
 }
 
 # Test configuration
+MODEL_PACKAGE_GROUP = "sdk-test-finetuned-models"
+REGION = "us-west-2"
+ACCOUNT_ID = "729646638167"
+
 TEST_CONFIG = {
-    "model_package_arn": "arn:aws:sagemaker:us-west-2:729646638167:model-package/sdk-test-finetuned-models/1",
     "evaluator_model": "anthropic.claude-3-5-haiku-20241022-v1:0",
-    "dataset_s3_uri": "s3://sagemaker-us-west-2-729646638167/model-customization/eval/gen_qa.jsonl",
+    "dataset_s3_uri": f"s3://sagemaker-{REGION}-{ACCOUNT_ID}/model-customization/eval/gen_qa.jsonl",
     "builtin_metrics": ["Completeness", "Faithfulness"],
     "custom_metrics_json": json.dumps([CUSTOM_METRIC_DICT]),
-    "s3_output_path": "s3://sagemaker-us-west-2-729646638167/model-customization/eval/base-model-fix-test/",
-    "mlflow_tracking_server_arn": "arn:aws:sagemaker:us-west-2:729646638167:mlflow-app/app-TTAUWUNMUHH6",
+    "s3_output_path": f"s3://sagemaker-{REGION}-{ACCOUNT_ID}/model-customization/eval/base-model-fix-test/",
+    "mlflow_tracking_server_arn": f"arn:aws:sagemaker:{REGION}:{ACCOUNT_ID}:mlflow-app/app-TTAUWUNMUHH6",
     "evaluate_base_model": True,  # This is the key difference - testing base model evaluation
-    "region": "us-west-2",
+    "region": REGION,
 }
 
 
+def _get_latest_model_package_arn():
+    """Return the ARN of the latest approved model package, or None."""
+    sm_client = boto3.client("sagemaker", region_name=REGION)
+    packages = sm_client.list_model_packages(
+        ModelPackageGroupName=MODEL_PACKAGE_GROUP,
+        ModelApprovalStatus="Approved",
+        SortBy="CreationTime",
+        SortOrder="Descending",
+        MaxResults=1,
+    )
+    summaries = packages.get("ModelPackageSummaryList", [])
+    if not summaries:
+        return None
+    return summaries[0]["ModelPackageArn"]
+
+
 @pytest.mark.serial
 class TestLLMAsJudgeBaseModelFix:
     """Integration test for base model fix in LLMAsJudgeEvaluator"""
 
     def test_base_model_evaluation_uses_correct_weights(self, mlflow_resource_arn):
         """
         Test that base model evaluation uses original base model weights.
-        
+
         This test verifies the fix for the bug where base model evaluation
         incorrectly used fine-tuned model weights. The test:
-        
+
         1. Creates an evaluator with evaluate_base_model=True
         2. Starts the evaluation pipeline
-        3. Verifies the pipeline has both EvaluateBaseInferenceModel and 
+        3. Verifies the pipeline has both EvaluateBaseInferenceModel and
            EvaluateCustomInferenceModel steps
         4. Waits for completion
         5. Compares results to ensure base and custom models produce different outputs
-        
+
         Expected behavior:
         - EvaluateBaseInferenceModel should use only BaseModelArn (no ModelPackageConfig)
         - EvaluateCustomInferenceModel should use ModelPackageConfig with SourceModelPackageArn
         - Results should show different performance between base and custom models
         """
+        model_package_arn = _get_latest_model_package_arn()
+        if not model_package_arn:
+            pytest.skip(
+                f"No approved model packages in group '{MODEL_PACKAGE_GROUP}'. "
+                "Run SFT/RLVR training first."
+            )
+
         logger.info("=" * 80)
         logger.info("Testing Base Model Fix: evaluate_base_model=True")
         logger.info("=" * 80)
-        
+
         # Step 1: Create evaluator with evaluate_base_model=True
         logger.info("Creating LLMAsJudgeEvaluator with evaluate_base_model=True")
-        
+        logger.info(f"Using model package: {model_package_arn}")
+
         evaluator = LLMAsJudgeEvaluator(
-            model=TEST_CONFIG["model_package_arn"],
+            model=model_package_arn,
             evaluator_model=TEST_CONFIG["evaluator_model"],
             dataset=TEST_CONFIG["dataset_s3_uri"],
             builtin_metrics=TEST_CONFIG["builtin_metrics"],
@@ -254,19 +282,27 @@ def test_base_model_evaluation_uses_correct_weights(self, mlflow_resource_arn):
     def test_base_model_false_still_works(self, mlflow_resource_arn):
         """
         Test that evaluate_base_model=False still works correctly (backward compatibility).
-        
+
         This test ensures the fix doesn't break existing functionality when
         evaluate_base_model=False (the default behavior).
         """
+        model_package_arn = _get_latest_model_package_arn()
+        if not model_package_arn:
+            pytest.skip(
+                f"No approved model packages in group '{MODEL_PACKAGE_GROUP}'. "
+                "Run SFT/RLVR training first."
+            )
+
         logger.info("=" * 80)
         logger.info("Testing Backward Compatibility: evaluate_base_model=False")
         logger.info("=" * 80)
-        
+
         # Create evaluator with evaluate_base_model=False
         logger.info("Creating LLMAsJudgeEvaluator with evaluate_base_model=False")
-        
+        logger.info(f"Using model package: {model_package_arn}")
+
         evaluator = LLMAsJudgeEvaluator(
-            model=TEST_CONFIG["model_package_arn"],
+            model=model_package_arn,
             evaluator_model=TEST_CONFIG["evaluator_model"],
             dataset=TEST_CONFIG["dataset_s3_uri"],
             builtin_metrics=TEST_CONFIG["builtin_metrics"],