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.

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/tests/integ/train/test_llm_as_judge_base_model_fix.py

@@ -82,11 +82,10 @@
 
 
 def _get_latest_model_package_arn():
-    """Return the ARN of the latest approved model package, or None."""
+    """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,
-        ModelApprovalStatus="Approved",
         SortBy="CreationTime",
         SortOrder="Descending",
         MaxResults=1,
@@ -123,7 +122,7 @@ def test_base_model_evaluation_uses_correct_weights(self, mlflow_resource_arn):
         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}'. "
+                f"No model packages in group '{MODEL_PACKAGE_GROUP}'. "
                 "Run SFT/RLVR training first."
             )
 
@@ -289,7 +288,7 @@ def test_base_model_false_still_works(self, mlflow_resource_arn):
         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}'. "
+                f"No model packages in group '{MODEL_PACKAGE_GROUP}'. "
                 "Run SFT/RLVR training first."
             )