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

* 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.

* fix(train): always use SageMakerPublicHub for base model ARN in evaluations

The backend now enforces that serverless training jobs only accept
BaseModelArn values pointing to SageMakerPublicHub. The SAGEMAKER_HUB_NAME
env var is for training recipe lookups only and should not affect the
base model ARN passed to evaluation pipelines.

Also remove ModelApprovalStatus filter from test helper since training
tests never set approval status on output packages.

* fix(train): fall back to SageMakerPublicHub when model not found in private hub

When _get_hub_content_metadata fails to find a model in the configured
private hub (e.g. sdktest), retry with SageMakerPublicHub. This handles
models like meta-textgeneration-llama-3-2-1b-instruct that only exist
in the public hub.

Author: mollyheamazon

sagemaker-train/src/sagemaker/train/common_utils/model_resolution.py

@@ -6,7 +6,6 @@
 - ModelPackage objects or ARNs (fine-tuned models)
 """
 
-import os
 import json
 import boto3
 from typing import Union, Optional, Dict, Any
@@ -239,11 +238,8 @@ def _resolve_model_package_object(self, model_package: 'ModelPackage') -> _Model
                     arn_parts = model_pkg_arn.split(':')
                     if len(arn_parts) >= 4:
                         region = arn_parts[3]
-                        # Use SAGEMAKER_HUB_NAME if set (private hub), otherwise fall back to public hub
-                        hub_name = os.environ.get("SAGEMAKER_HUB_NAME", "SageMakerPublicHub")
-                        # Private hubs are account-scoped; public hub uses 'aws' as account
-                        hub_account = "aws" if hub_name == "SageMakerPublicHub" else arn_parts[4]
-                        base_model_arn = f"arn:aws:sagemaker:{region}:{hub_account}:hub-content/{hub_name}/Model/{hub_content_name}/{hub_content_version}"
+                        # Base model always lives in SageMakerPublicHub (SAGEMAKER_HUB_NAME is for training recipes only)
+                        base_model_arn = f"arn:aws:sagemaker:{region}:aws:hub-content/SageMakerPublicHub/Model/{hub_content_name}/{hub_content_version}"
 
         # If we couldn't extract or construct base model ARN, this is not a supported model package
         if not base_model_arn:

sagemaker-train/src/sagemaker/train/common_utils/recipe_utils.py

@@ -63,14 +63,30 @@ def _get_hub_content_metadata(
         ... )
         >>> print(metadata['HubContentName'])
     """
-    hub_content = HubContent.get(
-        hub_name=hub_name,
-        hub_content_type=hub_content_type,
-        hub_content_name=hub_content_name,
-        region=region,
-        session=session
-    )
-    
+    try:
+        hub_content = HubContent.get(
+            hub_name=hub_name,
+            hub_content_type=hub_content_type,
+            hub_content_name=hub_content_name,
+            region=region,
+            session=session
+        )
+    except Exception:
+        if hub_name != "SageMakerPublicHub":
+            logger.info(
+                f"Hub content '{hub_content_name}' not found in '{hub_name}', "
+                f"falling back to SageMakerPublicHub"
+            )
+            hub_content = HubContent.get(
+                hub_name="SageMakerPublicHub",
+                hub_content_type=hub_content_type,
+                hub_content_name=hub_content_name,
+                region=region,
+                session=session
+            )
+        else:
+            raise
+
     # Convert to dict for easier access
     hub_content_dict = hub_content.__dict__
 

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,77 @@
 }
 
 # 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 model package, or None."""
+    sm_client = boto3.client("sagemaker", region_name=REGION)
+    packages = sm_client.list_model_packages(
+        ModelPackageGroupName=MODEL_PACKAGE_GROUP,
+        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 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 +281,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 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"],