tests: Implement integration tests covering JumpStart PrivateHub workflows

src/sagemaker/jumpstart/model.py

@@ -1036,13 +1036,15 @@ def _get_deployment_configs(
                 image_uri=image_uri,
                 region=self.region,
                 model_version=self.model_version,
+                hub_arn=self.hub_arn,
             )
             deploy_kwargs = get_deploy_kwargs(
                 model_id=self.model_id,
                 instance_type=instance_type_to_use,
                 sagemaker_session=self.sagemaker_session,
                 region=self.region,
                 model_version=self.model_version,
+                hub_arn=self.hub_arn,
             )
 
             deployment_config_metadata = DeploymentConfigMetadata(

tests/integ/sagemaker/jumpstart/conftest.py

@@ -16,24 +16,44 @@
 import boto3
 import pytest
 from botocore.config import Config
+from sagemaker.jumpstart.hub.hub import Hub
 from sagemaker.session import Session
 from tests.integ.sagemaker.jumpstart.constants import (
     ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID,
+    ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME,
+    HUB_NAME_PREFIX,
     JUMPSTART_TAG,
+    SM_JUMPSTART_PUBLIC_HUB_NAME,
+)
+
+from sagemaker.jumpstart.types import (
+    HubContentType,
 )
 
 
 from tests.integ.sagemaker.jumpstart.utils import (
     get_test_artifact_bucket,
     get_test_suite_id,
+    get_sm_session,
 )
 
 from sagemaker.jumpstart.constants import JUMPSTART_DEFAULT_REGION_NAME
 
 
 def _setup():
     print("Setting up...")
-    os.environ.update({ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID: get_test_suite_id()})
+    test_suit_id = get_test_suite_id()
+    test_hub_name = f"{HUB_NAME_PREFIX}{test_suit_id}"
+    test_hub_description = "PySDK Integ Test Private Hub"
+
+    os.environ.update({ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID: test_suit_id})
+    os.environ.update({ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME: test_hub_name})
+
+    # Create a private hub to use for the test session
+    hub = Hub(
+        hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME], sagemaker_session=get_sm_session()
+    )
+    hub.create(description=test_hub_description)
 
 
 def _teardown():
@@ -113,6 +133,37 @@ def _teardown():
     bucket = s3_resource.Bucket(test_cache_bucket)
     bucket.objects.filter(Prefix=test_suite_id + "/").delete()
 
+    # delete private hubs
+    _delete_hubs(sagemaker_session)
+
+
+def _delete_hubs(sagemaker_session):
+    # list Hubs created by PySDK integration tests
+    list_hub_response = sagemaker_session.list_hubs(
+        name_contains=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME]
+    )
+
+    for hub in list_hub_response["HubSummaries"]:
+        if hub["HubName"] != SM_JUMPSTART_PUBLIC_HUB_NAME:
+            # delete all hub contents first
+            _delete_hub_contents(sagemaker_session, hub["HubName"])
+            sagemaker_session.delete_hub(hub["HubName"])
+
+
+def _delete_hub_contents(sagemaker_session, test_hub_name):
+    # list hub_contents for the given hub
+    list_hub_content_response = sagemaker_session.list_hub_contents(
+        hub_name=test_hub_name, hub_content_type=HubContentType.MODEL_REFERENCE.value
+    )
+
+    # delete hub_contents for the given hub
+    for models in list_hub_content_response["HubContentSummaries"]:
+        sagemaker_session.delete_hub_content_reference(
+            hub_name=test_hub_name,
+            hub_content_type=HubContentType.MODEL_REFERENCE.value,
+            hub_content_name=models["HubContentName"],
+        )
+
 
 @pytest.fixture(scope="session", autouse=True)
 def setup(request):

tests/integ/sagemaker/jumpstart/constants.py

@@ -37,8 +37,13 @@ def _to_s3_path(filename: str, s3_prefix: Optional[str]) -> str:
 
 ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID = "JUMPSTART_SDK_TEST_SUITE_ID"
 
+ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME = "JUMPSTART_SDK_TEST_HUB_NAME"
+
 JUMPSTART_TAG = "JumpStart-SDK-Integ-Test-Suite-Id"
 
+SM_JUMPSTART_PUBLIC_HUB_NAME = "SageMakerPublicHub"
+
+HUB_NAME_PREFIX = "PySDK-HubTest-"
 
 TRAINING_DATASET_MODEL_DICT = {
     ("huggingface-spc-bert-base-cased", "1.0.0"): ("training-datasets/QNLI-tiny/"),

tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py

@@ -0,0 +1,174 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
+import os
+import time
+
+import pytest
+from sagemaker.enums import EndpointType
+from sagemaker.jumpstart.hub.hub import Hub
+from sagemaker.jumpstart.hub.utils import generate_hub_arn_for_init_kwargs
+from sagemaker.predictor import retrieve_default
+
+from sagemaker.jumpstart.constants import JUMPSTART_LOGGER
+
+import tests.integ
+
+from sagemaker.jumpstart.model import JumpStartModel
+from tests.integ.sagemaker.jumpstart.constants import (
+    ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME,
+    ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID,
+    JUMPSTART_TAG,
+)
+from tests.integ.sagemaker.jumpstart.utils import (
+    get_public_hub_model_arn,
+    get_sm_session,
+)
+
+MAX_INIT_TIME_SECONDS = 5
+
+TEST_MODEL_IDS = {
+    "catboost-classification-model",
+    "huggingface-txt2img-conflictx-complex-lineart",
+    "meta-textgeneration-llama-2-7b",
+    "meta-textgeneration-llama-3-2-1b",
+    "catboost-regression-model",
+}
+
+
+@pytest.fixture(scope="session")
+def add_models():
+    # Create Model References to test in Hub
+    hub_instance = Hub(
+        hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME], sagemaker_session=get_sm_session()
+    )
+    for model in TEST_MODEL_IDS:
+        hub_instance.create_model_reference(model_arn=get_public_hub_model_arn(hub_instance, model))
+
+
+def test_jumpstart_hub_model(setup, add_models):
+
+    JUMPSTART_LOGGER.info("starting test")
+    JUMPSTART_LOGGER.info(f"get identity {get_sm_session().get_caller_identity_arn()}")
+
+    model_id = "catboost-classification-model"
+
+    model = JumpStartModel(
+        model_id=model_id,
+        role=get_sm_session().get_caller_identity_arn(),
+        sagemaker_session=get_sm_session(),
+        hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME],
+    )
+
+    # uses ml.m5.4xlarge instance
+    model.deploy(
+        tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
+    )
+
+
+def test_jumpstart_hub_gated_model(setup, add_models):
+
+    model_id = "meta-textgeneration-llama-3-2-1b"
+
+    model = JumpStartModel(
+        model_id=model_id,
+        role=get_sm_session().get_caller_identity_arn(),
+        sagemaker_session=get_sm_session(),
+        hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME],
+    )
+
+    # uses ml.g6.xlarge instance
+    predictor = model.deploy(
+        accept_eula=True,
+        tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
+    )
+
+    payload = {
+        "inputs": "some-payload",
+        "parameters": {"max_new_tokens": 256, "top_p": 0.9, "temperature": 0.6},
+    }
+
+    response = predictor.predict(payload, custom_attributes="accept_eula=true")
+
+    assert response is not None
+
+
+def test_jumpstart_gated_model_inference_component_enabled(setup, add_models):
+
+    model_id = "meta-textgeneration-llama-2-7b"
+
+    hub_name = os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME]
+
+    region = tests.integ.test_region()
+
+    sagemaker_session = get_sm_session()
+
+    hub_arn = generate_hub_arn_for_init_kwargs(
+        hub_name=hub_name, region=region, session=sagemaker_session
+    )
+
+    model = JumpStartModel(
+        model_id=model_id,
+        role=get_sm_session().get_caller_identity_arn(),
+        sagemaker_session=sagemaker_session,
+        hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME],
+    )
+
+    # uses ml.g5.2xlarge instance
+    model.deploy(
+        tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
+        accept_eula=True,
+        endpoint_type=EndpointType.INFERENCE_COMPONENT_BASED,
+    )
+
+    predictor = retrieve_default(
+        endpoint_name=model.endpoint_name,
+        sagemaker_session=sagemaker_session,
+        tolerate_vulnerable_model=True,
+        hub_arn=hub_arn,
+    )
+
+    payload = {
+        "inputs": "some-payload",
+        "parameters": {"max_new_tokens": 256, "top_p": 0.9, "temperature": 0.6},
+    }
+
+    response = predictor.predict(payload)
+
+    assert response is not None
+
+    model = JumpStartModel.attach(
+        predictor.endpoint_name, sagemaker_session=sagemaker_session, hub_name=hub_name
+    )
+    assert model.model_id == model_id
+    assert model.endpoint_name == predictor.endpoint_name
+    assert model.inference_component_name == predictor.component_name
+
+
+def test_instatiating_model(setup, add_models):
+
+    model_id = "catboost-regression-model"
+
+    start_time = time.perf_counter()
+
+    JumpStartModel(
+        model_id=model_id,
+        role=get_sm_session().get_caller_identity_arn(),
+        sagemaker_session=get_sm_session(),
+        hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME],
+    )
+
+    elapsed_time = time.perf_counter() - start_time
+
+    assert elapsed_time <= MAX_INIT_TIME_SECONDS

tests/integ/sagemaker/jumpstart/private_hub/test_hub.py

@@ -0,0 +1,39 @@
+import pytest
+from sagemaker.jumpstart.hub.hub import Hub
+
+from tests.integ.sagemaker.jumpstart.utils import (
+    get_sm_session,
+)
+from tests.integ.sagemaker.jumpstart.utils import (
+    get_test_suite_id,
+)
+from tests.integ.sagemaker.jumpstart.constants import (
+    HUB_NAME_PREFIX,
+)
+
+
+@pytest.fixture
+def hub_instance():
+    HUB_NAME = f"{HUB_NAME_PREFIX}-{get_test_suite_id()}"
+    hub = Hub(HUB_NAME, sagemaker_session=get_sm_session())
+    yield hub
+
+
+def test_private_hub(setup, hub_instance):
+    # Createhub
+    create_hub_response = hub_instance.create(
+        description="This is a Test Private Hub.",
+        display_name="PySDK integration tests Hub",
+        search_keywords=["jumpstart-sdk-integ-test"],
+    )
+
+    # Create Hub Verifications
+    assert create_hub_response is not None
+
+    # Describe Hub
+    hub_description = hub_instance.describe()
+    assert hub_description is not None
+
+    # Delete Hub
+    delete_hub_response = hub_instance.delete()
+    assert delete_hub_response is not None

tests/integ/sagemaker/jumpstart/private_hub/test_hub_content.py

@@ -0,0 +1,34 @@
+import os
+from sagemaker.jumpstart.hub.hub import Hub
+
+from sagemaker.jumpstart.hub.interfaces import DescribeHubContentResponse
+from tests.integ.sagemaker.jumpstart.utils import (
+    get_sm_session,
+)
+from tests.integ.sagemaker.jumpstart.utils import get_public_hub_model_arn
+from tests.integ.sagemaker.jumpstart.constants import (
+    ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME,
+)
+
+
+def test_hub_model_reference(setup):
+    model_id = "meta-textgenerationneuron-llama-3-2-1b-instruct"
+
+    hub_instance = Hub(
+        hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME], sagemaker_session=get_sm_session()
+    )
+
+    # Create Model Reference
+    create_model_response = hub_instance.create_model_reference(
+        model_arn=get_public_hub_model_arn(hub_instance, model_id)
+    )
+    assert create_model_response is not None
+
+    # Describe Model
+    describe_model_response = hub_instance.describe_model(model_name=model_id)
+    assert describe_model_response is not None
+    assert type(describe_model_response) == DescribeHubContentResponse
+
+    # Delete Model Reference
+    delete_model_response = hub_instance.delete_model_reference(model_name=model_id)
+    assert delete_model_response is not None

tests/integ/sagemaker/jumpstart/utils.py

@@ -32,6 +32,7 @@
 )
 from sagemaker.jumpstart.constants import JUMPSTART_DEFAULT_REGION_NAME
 from sagemaker.jumpstart.utils import get_jumpstart_content_bucket
+from sagemaker.jumpstart.hub.hub import Hub
 
 from sagemaker.session import Session
 
@@ -115,6 +116,20 @@ def download_file(local_download_path, s3_bucket, s3_key, s3_client) -> None:
     s3_client.download_file(s3_bucket, s3_key, local_download_path)
 
 
+def get_public_hub_model_arn(hub: Hub, model_id: str) -> str:
+    filter_value = f"model_id == {model_id}"
+    response = hub.list_sagemaker_public_hub_models(filter=filter_value)
+
+    models = response["hub_content_summaries"]
+    while response["next_token"]:
+        response = hub.list_sagemaker_public_hub_models(
+            filter=filter_value, next_token=response["next_token"]
+        )
+        models.extend(response["hub_content_summaries"])
+
+    return models[0]["hub_content_arn"]
+
+
 class EndpointInvoker:
     def __init__(
         self,