tests: Implement integration tests covering JumpStart PrivateHub workflows

Author: malavhs

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,32 +16,50 @@
 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 sagemaker.jumpstart.constants import JUMPSTART_LOGGER
 
 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})
+    hub = Hub(hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME], sagemaker_session=get_sm_session())
+    hub.create(description=test_hub_description)
+    describe_hub_response = hub.describe()
+    JUMPSTART_LOGGER.info(f"Describe Hub {describe_hub_response}")
 
 def _teardown():
     print("Tearing down...")
 
     test_cache_bucket = get_test_artifact_bucket()
 
     test_suite_id = os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]
+    test_hub_name = os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME]
 
     boto3_session = boto3.Session(region_name=JUMPSTART_DEFAULT_REGION_NAME)
 
@@ -113,6 +131,40 @@ 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=HUB_NAME_PREFIX)
+
+    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'])
+            JUMPSTART_LOGGER.info(f"Deleting {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
+    )
+    JUMPSTART_LOGGER.info(f"Listing HubContents {list_hub_content_response}")
+
+    #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/model/test_jumpstart_model.py

@@ -56,6 +56,9 @@
 
 def test_non_prepacked_jumpstart_model(setup):
 
+    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(

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

@@ -0,0 +1,175 @@
+# 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 io
+import os
+import sys
+import time
+from unittest import mock
+import logging
+
+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="module")
+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/model/__init__.py

@@ -0,0 +1,40 @@
+import pytest
+import os
+from unittest.mock import MagicMock, patch
+from sagemaker.jumpstart.hub.hub import Hub
+from sagemaker.jumpstart.constants import JUMPSTART_LOGGER
+
+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="malavhs Test 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/model/test_jumpstart_private_hub_model.py

@@ -0,0 +1,38 @@
+import pytest
+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,
+)
+import tests
+
+
+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