|
| 1 | +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"). You |
| 4 | +# may not use this file except in compliance with the License. A copy of |
| 5 | +# the License is located at |
| 6 | +# |
| 7 | +# http://aws.amazon.com/apache2.0/ |
| 8 | +# |
| 9 | +# or in the "license" file accompanying this file. This file is |
| 10 | +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF |
| 11 | +# ANY KIND, either express or implied. See the License for the specific |
| 12 | +# language governing permissions and limitations under the License. |
| 13 | +"""Integration test for private hub artifact resolution. |
| 14 | +
|
| 15 | +Verifies that ModelBuilder.from_jumpstart_config with hub_name resolves |
| 16 | +model artifacts through the private hub content reference, rather than |
| 17 | +falling back to the public JumpStart S3 cache. |
| 18 | +
|
| 19 | +This test creates its own private hub, adds a ModelReference, runs the |
| 20 | +build flow, and tears everything down afterward. |
| 21 | +""" |
| 22 | +from __future__ import absolute_import |
| 23 | + |
| 24 | +import uuid |
| 25 | +import time |
| 26 | +import logging |
| 27 | + |
| 28 | +import boto3 |
| 29 | +import pytest |
| 30 | + |
| 31 | +from sagemaker.serve.model_builder import ModelBuilder |
| 32 | +from sagemaker.core.jumpstart.configs import JumpStartConfig |
| 33 | +from sagemaker.train.configs import Compute |
| 34 | + |
| 35 | +logger = logging.getLogger(__name__) |
| 36 | + |
| 37 | +# Use a small/fast JumpStart model for the test |
| 38 | +TEST_MODEL_ID = "huggingface-llm-gemma-2b" |
| 39 | +TEST_MODEL_VERSION = "*" |
| 40 | +TEST_INSTANCE_TYPE = "ml.g5.xlarge" |
| 41 | +TEST_REGION = "us-east-1" |
| 42 | + |
| 43 | +HUB_NAME_PREFIX = "sdk-integ-test-private-hub" |
| 44 | + |
| 45 | + |
| 46 | +@pytest.fixture(scope="module") |
| 47 | +def private_hub(): |
| 48 | + """Create a private hub with a ModelReference, tear down after tests.""" |
| 49 | + sm = boto3.client("sagemaker", region_name=TEST_REGION) |
| 50 | + hub_name = f"{HUB_NAME_PREFIX}-{uuid.uuid4().hex[:8]}" |
| 51 | + |
| 52 | + # --- Setup --- |
| 53 | + logger.info("Creating private hub: %s", hub_name) |
| 54 | + sm.create_hub( |
| 55 | + HubName=hub_name, |
| 56 | + HubDescription="SDK integration test for private hub artifact resolution", |
| 57 | + Tags=[ |
| 58 | + {"Key": "Purpose", "Value": "sdk-integ-test"}, |
| 59 | + {"Key": "AutoCleanup", "Value": "true"}, |
| 60 | + ], |
| 61 | + ) |
| 62 | + |
| 63 | + # Wait for hub to be ready |
| 64 | + for _ in range(30): |
| 65 | + resp = sm.describe_hub(HubName=hub_name) |
| 66 | + if resp["HubStatus"] == "InService": |
| 67 | + break |
| 68 | + time.sleep(2) |
| 69 | + else: |
| 70 | + pytest.fail(f"Hub {hub_name} did not reach InService state") |
| 71 | + |
| 72 | + # Add a ModelReference to the public JumpStart content |
| 73 | + public_arn = ( |
| 74 | + f"arn:aws:sagemaker:{TEST_REGION}:aws:hub-content/" |
| 75 | + f"SageMakerPublicHub/Model/{TEST_MODEL_ID}" |
| 76 | + ) |
| 77 | + logger.info("Creating hub content reference to: %s", public_arn) |
| 78 | + sm.create_hub_content_reference( |
| 79 | + HubName=hub_name, |
| 80 | + SageMakerPublicHubContentArn=public_arn, |
| 81 | + ) |
| 82 | + |
| 83 | + # Wait for content reference to be available |
| 84 | + for _ in range(60): |
| 85 | + try: |
| 86 | + contents = sm.list_hub_contents( |
| 87 | + HubName=hub_name, |
| 88 | + HubContentType="ModelReference", |
| 89 | + ) |
| 90 | + summaries = contents.get("HubContentSummaries", []) |
| 91 | + if any( |
| 92 | + s["HubContentName"] == TEST_MODEL_ID |
| 93 | + and s.get("HubContentStatus") == "Available" |
| 94 | + for s in summaries |
| 95 | + ): |
| 96 | + break |
| 97 | + except Exception: |
| 98 | + pass |
| 99 | + time.sleep(3) |
| 100 | + else: |
| 101 | + pytest.fail( |
| 102 | + f"ModelReference for {TEST_MODEL_ID} not available in " |
| 103 | + f"hub {hub_name} after 3 minutes" |
| 104 | + ) |
| 105 | + |
| 106 | + yield hub_name |
| 107 | + |
| 108 | + # --- Teardown --- |
| 109 | + logger.info("Cleaning up hub: %s", hub_name) |
| 110 | + try: |
| 111 | + # Delete hub content references first |
| 112 | + contents = sm.list_hub_contents( |
| 113 | + HubName=hub_name, |
| 114 | + HubContentType="ModelReference", |
| 115 | + ) |
| 116 | + for content in contents.get("HubContentSummaries", []): |
| 117 | + try: |
| 118 | + sm.delete_hub_content_reference( |
| 119 | + HubName=hub_name, |
| 120 | + HubContentName=content["HubContentName"], |
| 121 | + HubContentType="ModelReference", |
| 122 | + ) |
| 123 | + except Exception as e: |
| 124 | + logger.warning("Failed to delete content ref: %s", e) |
| 125 | + |
| 126 | + # Delete the hub |
| 127 | + sm.delete_hub(HubName=hub_name) |
| 128 | + logger.info("Hub %s deleted", hub_name) |
| 129 | + except Exception as e: |
| 130 | + logger.warning("Hub cleanup failed: %s", e) |
| 131 | + |
| 132 | + |
| 133 | +@pytest.fixture(scope="module") |
| 134 | +def sagemaker_session(): |
| 135 | + """Create a SageMaker session using default credentials.""" |
| 136 | + from sagemaker.core.helper.session_helper import Session |
| 137 | + |
| 138 | + return Session(boto_session=boto3.Session(region_name=TEST_REGION)) |
| 139 | + |
| 140 | + |
| 141 | +@pytest.fixture(scope="module") |
| 142 | +def execution_role(): |
| 143 | + """Get a SageMaker execution role from the caller's identity.""" |
| 144 | + sts = boto3.client("sts", region_name=TEST_REGION) |
| 145 | + identity = sts.get_caller_identity() |
| 146 | + account_id = identity["Account"] |
| 147 | + # Use the standard SageMaker execution role naming convention |
| 148 | + return f"arn:aws:iam::{account_id}:role/Admin" |
| 149 | + |
| 150 | + |
| 151 | +@pytest.mark.slow_test |
| 152 | +def test_from_jumpstart_config_derives_hub_arn(private_hub, sagemaker_session): |
| 153 | + """Verify from_jumpstart_config correctly derives hub_arn from hub_name.""" |
| 154 | + js_config = JumpStartConfig( |
| 155 | + model_id=TEST_MODEL_ID, |
| 156 | + model_version=TEST_MODEL_VERSION, |
| 157 | + hub_name=private_hub, |
| 158 | + ) |
| 159 | + |
| 160 | + mb = ModelBuilder.from_jumpstart_config( |
| 161 | + jumpstart_config=js_config, |
| 162 | + compute=Compute(instance_type=TEST_INSTANCE_TYPE), |
| 163 | + sagemaker_session=sagemaker_session, |
| 164 | + ) |
| 165 | + |
| 166 | + assert mb.hub_arn is not None, ( |
| 167 | + f"hub_arn is None after from_jumpstart_config with " |
| 168 | + f"hub_name={private_hub}" |
| 169 | + ) |
| 170 | + assert private_hub in mb.hub_arn |
| 171 | + logger.info("hub_arn correctly derived: %s", mb.hub_arn) |
| 172 | + |
| 173 | + |
| 174 | +@pytest.mark.slow_test |
| 175 | +def test_build_resolves_artifacts_via_private_hub( |
| 176 | + private_hub, execution_role, sagemaker_session |
| 177 | +): |
| 178 | + """Verify build() resolves model data through the private hub.""" |
| 179 | + js_config = JumpStartConfig( |
| 180 | + model_id=TEST_MODEL_ID, |
| 181 | + model_version=TEST_MODEL_VERSION, |
| 182 | + hub_name=private_hub, |
| 183 | + accept_eula=True, |
| 184 | + ) |
| 185 | + |
| 186 | + mb = ModelBuilder.from_jumpstart_config( |
| 187 | + jumpstart_config=js_config, |
| 188 | + role_arn=execution_role, |
| 189 | + compute=Compute(instance_type=TEST_INSTANCE_TYPE), |
| 190 | + sagemaker_session=sagemaker_session, |
| 191 | + ) |
| 192 | + |
| 193 | + mb.build() |
| 194 | + |
| 195 | + # After build, model data URI must NOT reference the public cache |
| 196 | + model_data = getattr(mb, "s3_model_data_url", None) |
| 197 | + assert model_data is not None, "No model_data found after build()" |
| 198 | + |
| 199 | + model_data_str = str(model_data) |
| 200 | + assert "jumpstart-cache-prod" not in model_data_str, ( |
| 201 | + f"Model data still points to public JumpStart cache: " |
| 202 | + f"{model_data_str}. Expected private hub artifact resolution." |
| 203 | + ) |
| 204 | + logger.info("Model data resolved to: %s", model_data_str) |
0 commit comments